From 9f890a921805a06058819a02f4a3ebc006514f3b Mon Sep 17 00:00:00 2001 From: Michael Suo Date: Mon, 28 Oct 2019 13:45:42 -0700 Subject: [PATCH 01/64] make sure clang-tidy is diffing against the right thing (#28788) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28788 Okay, my last fix was wrong because it turns out that the base SHA is computed at PR time using the actual repo's view of the base ref, not the user's. So if the user doesn't rebase on top of the latest master before putting up the PR, the diff thing is wrong anyway. This PR fixes the issue by not relying on any of these API details and just getting the merge-base of the base and head refs, which should guarantee we are diffing against the right thing. This solution taken from https://github.com/github/VisualStudio/pull/1008 Test Plan: Imported from OSS Differential Revision: D18172391 Pulled By: suo fbshipit-source-id: 491a50119194508b2eefa5bd39fe813ca85f27b1 --- .github/workflows/lint.yml | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 581ad371f8262..1f2515fe6d3e9 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -209,7 +209,10 @@ jobs: run: | set -eux git remote add upstream https://github.com/pytorch/pytorch - git fetch upstream "${{ github.base_ref}}" + git fetch upstream "$GITHUB_BASE_REF" + BASE_SHA=${{ github.event.pull_request.base.sha }} + HEAD_SHA=${{ github.event.pull_request.head.sha }} + MERGE_BASE=$(git merge-base $BASE_SHA $HEAD_SHA) if [[ ! -d build ]]; then git submodule update --init --recursive @@ -238,8 +241,9 @@ jobs: # The negative filters below are to exclude files that include onnx_pb.h or # caffe2_pb.h, otherwise we'd have to build protos as part of this CI job. python tools/clang_tidy.py \ + --verbose \ --paths torch/csrc/ \ - --diff "${{ github.event.pull_request.base.sha}}" \ + --diff "$MERGE_BASE" \ -g"-torch/csrc/jit/export.cpp" \ -g"-torch/csrc/jit/import.cpp" \ -g"-torch/csrc/jit/netdef_converter.cpp" \ @@ -252,6 +256,6 @@ jobs: check_name: 'clang-tidy' linter_output_path: 'clang-tidy-output.txt' commit_sha: ${{ steps.get_pr_tip.outputs.commit_sha }} - regex: '^(?.*?):(?\d+):(?\d+): (?.*?) (?\[.*\])' + regex: '^(?.*?):(?\d+):(?\d+): (?.*?) \[(?.*)\]' env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} From 02d318461e5c7bded304c42ed7075de84f71dac6 Mon Sep 17 00:00:00 2001 From: Jianyu Huang Date: Mon, 28 Oct 2019 13:48:22 -0700 Subject: [PATCH 02/64] Temporarily disable test_numerical_consistency_per_channel due to failure (#28807) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28807 `FAIL: test_numerical_consistency_per_channel (_main_.TestFakeQuantizePerChannel)` This test is failing consistently on master, we can't find a clean blame. ghstack-source-id: 92763176 Test Plan: CI Differential Revision: D18181496 fbshipit-source-id: 5948af06c4cb7dea9a8db1366deb7c12f6ec1c72 --- test/test_fake_quant.py | 1 + test/test_quantizer.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_fake_quant.py b/test/test_fake_quant.py index ee82db730fbf3..bcd413d69f4a2 100644 --- a/test/test_fake_quant.py +++ b/test/test_fake_quant.py @@ -256,6 +256,7 @@ def test_backward_per_channel(self, device, X): @given(device=st.sampled_from(['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu']), X=hu.per_channel_tensor(shapes=hu.array_shapes(1, 5,), qparams=hu.qparams(dtypes=torch.quint8))) + @unittest.skip("temporarily disable the test") def test_numerical_consistency_per_channel(self, device, X): r"""Comparing numerical consistency between CPU quantize/dequantize op and the CPU fake quantize op """ diff --git a/test/test_quantizer.py b/test/test_quantizer.py index 361d7d0146651..3df63f017deab 100644 --- a/test/test_quantizer.py +++ b/test/test_quantizer.py @@ -39,7 +39,7 @@ def __init__(self): @unittest.skipUnless('fbgemm' in torch.backends.quantized.supported_engines, " Quantized operations require FBGEMM. FBGEMM is only optimized for CPUs" " with instruction set support avx2 or newer.") -@unittest.skip("temoprarily disable the test") +@unittest.skip("temporarily disable the test") class QuantizerTestCase(TestCase): @_tmp_donotuse_dont_inline_everything def test_default(self): From 9e64c54c019e08c1e0a6d4074bde0feb218953e2 Mon Sep 17 00:00:00 2001 From: Jianyu Huang Date: Mon, 28 Oct 2019 14:20:35 -0700 Subject: [PATCH 03/64] Add the warning message for API with linear modules (#28766) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28766 Add the warning message to explicitly ask the users to upgrade the deprecated `torch.jit.quantized` API to the new `torch.quantization.quantize_dynamic` API. ghstack-source-id: 92711620 Test Plan: CI Differential Revision: D18164903 fbshipit-source-id: e6aff2527f335c2d9f362e6856ce8597edb52aaa --- torch/jit/quantized.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/torch/jit/quantized.py b/torch/jit/quantized.py index 0cffc40522d8a..4fe3a738d0388 100644 --- a/torch/jit/quantized.py +++ b/torch/jit/quantized.py @@ -7,6 +7,8 @@ from torch.nn.utils.rnn import PackedSequence +import warnings + class QuantizedLinear(torch.jit.ScriptModule): __constants__ = ['scale', 'zero_point'] @@ -602,6 +604,9 @@ def quantize_rnn_cell_modules(module): def quantize_linear_modules(module, dtype=torch.int8): + warnings.warn("quantize_linear_modules function has been deprecated. " + "Please use torch.quantization.quantize_dynamic API instead.") + reassign = {} for name, mod in module.named_modules(): if mod is module: From f33813d5890396ac241d7a7dbbc3f0e9a304ee3c Mon Sep 17 00:00:00 2001 From: Peter Bell Date: Mon, 28 Oct 2019 14:21:34 -0700 Subject: [PATCH 04/64] Return NotImplemented from all binary math ops (#27423) Summary: Fixes https://github.com/pytorch/pytorch/issues/26333 Fixes the operators missed in https://github.com/pytorch/pytorch/issues/26507 and includes a test for all operators. Pull Request resolved: https://github.com/pytorch/pytorch/pull/27423 Differential Revision: D17835390 Pulled By: ezyang fbshipit-source-id: 7a1351c7ccc8ad11454dbaa00d3701dcee4f06a8 --- .../expect/TestOperators.test_equal.expect | 8 ++-- test/onnx/expect/TestOperators.test_ge.expect | 8 ++-- test/onnx/expect/TestOperators.test_gt.expect | 8 ++-- test/onnx/expect/TestOperators.test_le.expect | 8 ++-- test/onnx/expect/TestOperators.test_lt.expect | 8 ++-- test/test_torch.py | 45 +++++++++++++++++++ tools/autograd/gen_python_functions.py | 30 ++++++++++++- .../templates/python_torch_functions.cpp | 13 ++++++ torch/tensor.py | 36 ++++++++++++--- 9 files changed, 136 insertions(+), 28 deletions(-) diff --git a/test/onnx/expect/TestOperators.test_equal.expect b/test/onnx/expect/TestOperators.test_equal.expect index d53ae0ded1c96..2c5fdc0dc3668 100644 --- a/test/onnx/expect/TestOperators.test_equal.expect +++ b/test/onnx/expect/TestOperators.test_equal.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Equal" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_ge.expect b/test/onnx/expect/TestOperators.test_ge.expect index abb3a62ef1a11..df2af85135dad 100644 --- a/test/onnx/expect/TestOperators.test_ge.expect +++ b/test/onnx/expect/TestOperators.test_ge.expect @@ -3,8 +3,8 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Less" } @@ -15,7 +15,7 @@ graph { } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -31,7 +31,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_gt.expect b/test/onnx/expect/TestOperators.test_gt.expect index 0dbf73c132439..01cd7d05f7619 100644 --- a/test/onnx/expect/TestOperators.test_gt.expect +++ b/test/onnx/expect/TestOperators.test_gt.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Greater" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_le.expect b/test/onnx/expect/TestOperators.test_le.expect index 82674fb9038e9..b362b7699bd01 100644 --- a/test/onnx/expect/TestOperators.test_le.expect +++ b/test/onnx/expect/TestOperators.test_le.expect @@ -3,8 +3,8 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Greater" } @@ -15,7 +15,7 @@ graph { } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -31,7 +31,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/onnx/expect/TestOperators.test_lt.expect b/test/onnx/expect/TestOperators.test_lt.expect index b4688b4342c21..a18842e06592d 100644 --- a/test/onnx/expect/TestOperators.test_lt.expect +++ b/test/onnx/expect/TestOperators.test_lt.expect @@ -3,14 +3,14 @@ producer_name: "pytorch" producer_version: "1.3" graph { node { - input: "0" - input: "1" + input: "x" + input: "y" output: "2" op_type: "Less" } name: "torch-jit-export" input { - name: "0" + name: "x" type { tensor_type { elem_type: 6 @@ -32,7 +32,7 @@ graph { } } input { - name: "1" + name: "y" type { tensor_type { elem_type: 6 diff --git a/test/test_torch.py b/test/test_torch.py index 1ce1e9ed3752f..723f0197f4350 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -14149,6 +14149,50 @@ def caller(cls, caller(cls, *test) +tensor_binary_ops = [ + '__lt__', '__le__', + '__gt__', '__ge__', + '__eq__', '__ne__', + + '__add__', '__radd__', '__iadd__', + '__sub__', '__rsub__', '__isub__', + '__mul__', '__rmul__', '__imul__', + '__matmul__', '__rmatmul__', '__imatmul__', + '__truediv__', '__rtruediv__', '__itruediv__', + '__floordiv__', '__rfloordiv__', '__ifloordiv__', + '__mod__', '__rmod__', '__imod__', + '__divmod__', '__rdivmod__', '__idivmod__', + '__pow__', '__rpow__', '__ipow__', + '__lshift__', '__rlshift__', '__ilshift__', + '__rshift__', '__rrshift__', '__irshift__', + '__and__', '__rand__', '__iand__', + '__xor__', '__rxor__', '__ixor__', + '__or__', '__ror__', '__ior__', +] + + +# Test that binary math operations return NotImplemented for unknown types. +def generate_not_implemented_tests(cls): + class UnknownType: + pass + + for op in tensor_binary_ops: + @dtypes(*_types) + def test(self, device, dtype): + # Generate the inputs + tensor = _small_2d(dtype, device) + + # Runs the tensor op on the device + result = getattr(tensor, op)(UnknownType()) + self.assertEqual(result, NotImplemented) + + test_name = "test_{}_not_implemented".format(op) + assert not hasattr(cls, test_name), "{0} already in {1}".format( + test_name, cls.__name__) + + setattr(cls, test_name, test) + + class TestTensorDeviceOps(TestCase): pass @@ -14162,6 +14206,7 @@ class TestTorch(TestCase, _TestTorchMixin): # pytest will fail. add_neg_dim_tests() generate_tensor_op_tests(TestTensorDeviceOps) +generate_not_implemented_tests(TestTorchDeviceType) instantiate_device_type_tests(TestTorchDeviceType, globals()) instantiate_device_type_tests(TestDevicePrecision, globals(), except_for='cpu') instantiate_device_type_tests(TestTensorDeviceOps, globals(), except_for='cpu') diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index e0c8c600464d4..2fb93baf1d1ee 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -48,6 +48,28 @@ 'div(Tensor, Scalar)', 'div_(Tensor, Scalar)', ] +# Python binary operator dunder methods +BINARY_OP_NAMES = [ + '__lt__', '__le__', + '__gt__', '__ge__', + '__eq__', '__ne__', + + '__add__', '__radd__', '__iadd__', + '__sub__', '__rsub__', '__isub__', + '__mul__', '__rmul__', '__imul__', + '__matmul__', '__rmatmul__', '__imatmul__', + '__truediv__', '__rtruediv__', '__itruediv__', + '__floordiv__', '__rfloordiv__', '__ifloordiv__', + '__mod__', '__rmod__', '__imod__', + '__divmod__', '__rdivmod__', '__idivmod__', + '__pow__', '__rpow__', '__ipow__', + '__lshift__', '__rlshift__', '__ilshift__', + '__rshift__', '__rrshift__', '__irshift__', + '__and__', '__rand__', '__iand__', + '__xor__', '__rxor__', '__ixor__', + '__or__', '__ror__', '__ior__', +] + PY_VARIABLE_METHOD_VARARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { @@ -120,6 +142,9 @@ PY_VARIABLE_METHOD_DEF = CodeTemplate("""\ {"${name}", (PyCFunction)${pycfunc_voidcast}${pycname}, ${flags}, NULL},""") +PY_VARIABLE_METHOD_BINOP_DEF = CodeTemplate("""\ +{"${name}", (PyCFunction)${pycfunc_voidcast}TypeError_to_NotImplemented_<${pycname}>, ${flags}, NULL},""") + PY_RETURN_NAMEDTUPLE_DEF = CodeTemplate("""\ static PyStructSequence_Field fields${namedtuple_type_index}[] = { ${namedtuple_fields} {nullptr} @@ -745,7 +770,10 @@ def process_function(name, declarations): env['flags'] += ' | METH_STATIC' py_methods.append(tmpl.substitute(env)) - py_method_defs.append(PY_VARIABLE_METHOD_DEF.substitute(env)) + if name in BINARY_OP_NAMES: + py_method_defs.append(PY_VARIABLE_METHOD_BINOP_DEF.substitute(env)) + else: + py_method_defs.append(PY_VARIABLE_METHOD_DEF.substitute(env)) for name in sorted(python_functions.keys()): process_function(name, python_functions[name]) diff --git a/tools/autograd/templates/python_torch_functions.cpp b/tools/autograd/templates/python_torch_functions.cpp index 299e82fe42e6e..935a01f6cc10d 100644 --- a/tools/autograd/templates/python_torch_functions.cpp +++ b/tools/autograd/templates/python_torch_functions.cpp @@ -447,6 +447,19 @@ static PyObject * THPVariable_numel(PyObject* self_, PyObject* args, PyObject* k END_HANDLE_TH_ERRORS } +// Wrapper converts a raised TypeError into returning NotImplemented +// Used to implement binary arithmetic operators +template +static PyObject * TypeError_to_NotImplemented_(PyObject* self, PyObject* args, PyObject* kwargs) { + PyObject* ret = Func(self, args, kwargs); + if (!ret && PyErr_ExceptionMatches(PyExc_TypeError)) { + PyErr_Clear(); + Py_INCREF(Py_NotImplemented); + ret = Py_NotImplemented; + } + return ret; +} + // generated methods start here ${py_methods} diff --git a/torch/tensor.py b/torch/tensor.py index 618efeb3b875c..4f88d1ba054cf 100644 --- a/torch/tensor.py +++ b/torch/tensor.py @@ -10,6 +10,25 @@ from torch._six import imap from torch._C import _add_docstr from numbers import Number +import functools + + +def _wrap_type_error_to_not_implemented(f): + from torch import _six + import inspect + + # functools.wraps doesn't work well with methods in python 2 + method_assignments = ('__name__', '__doc__') + assigned = (method_assignments if _six.PY2 and inspect.ismethoddescriptor(f) + else functools.WRAPPER_ASSIGNMENTS) + + @functools.wraps(f, assigned=assigned) + def wrapped(*args, **kwargs): + try: + return f(*args, **kwargs) + except TypeError: + return NotImplemented + return wrapped # NB: If you subclass Tensor, and want to share the subclassed class @@ -372,17 +391,20 @@ def __format__(self, format_spec): return object.__format__(self, format_spec) def __ipow__(self, other): - raise NotImplementedError("in-place pow not implemented") + return NotImplemented + @_wrap_type_error_to_not_implemented def __rpow__(self, other): return self.new_tensor(other) ** self + @_wrap_type_error_to_not_implemented def __floordiv__(self, other): result = self / other if result.dtype.is_floating_point: result = result.trunc() return result + @_wrap_type_error_to_not_implemented def __rfloordiv__(self, other): result = other / self if result.dtype.is_floating_point: @@ -391,12 +413,12 @@ def __rfloordiv__(self, other): __neg__ = _C._TensorBase.neg - __eq__ = _C._TensorBase.eq - __ne__ = _C._TensorBase.ne - __lt__ = _C._TensorBase.lt - __le__ = _C._TensorBase.le - __gt__ = _C._TensorBase.gt - __ge__ = _C._TensorBase.ge + __eq__ = _wrap_type_error_to_not_implemented(_C._TensorBase.eq) + __ne__ = _wrap_type_error_to_not_implemented(_C._TensorBase.ne) + __lt__ = _wrap_type_error_to_not_implemented(_C._TensorBase.lt) + __le__ = _wrap_type_error_to_not_implemented(_C._TensorBase.le) + __gt__ = _wrap_type_error_to_not_implemented(_C._TensorBase.gt) + __ge__ = _wrap_type_error_to_not_implemented(_C._TensorBase.ge) __abs__ = _C._TensorBase.abs def __len__(self): From 688a9dbe3cf090a7c10efac5e69be2d41dc5f77e Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Mon, 28 Oct 2019 14:43:51 -0700 Subject: [PATCH 05/64] Move type casting to c10/util/TypeCast.h (#28426) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28426 Type casting is used in copy, and will be used also in tensor iterator in the next stacked diff. I move it to c10 to make it serve as an common util for different things. I also add two dynamic casting functions - fetch_and_cast - cast_and_store fetch_and_cast fetch a value with dynamic type specified by a ScalarType from a void pointer and cast it to a static type. cast_and_store casts a static typed value into dynamic type specified by a ScalarType, and store it into a void pointer. Test Plan: Imported from OSS Differential Revision: D18170996 Pulled By: ezyang fbshipit-source-id: 41658afd5c0ab58c6b6c510424893d9a2a0c059e --- aten/src/ATen/cpu/vec256/vec256_base.h | 4 +- aten/src/ATen/native/Copy.h | 38 ------ aten/src/ATen/native/cpu/CopyKernel.cpp | 7 +- aten/src/ATen/native/cuda/Copy.cu | 3 +- c10/core/ScalarType.h | 4 + c10/util/TypeCast.h | 172 ++++++++++++++++++++++++ 6 files changed, 183 insertions(+), 45 deletions(-) create mode 100644 c10/util/TypeCast.h diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h index 64063b3da2b92..d7c4aab6ce209 100644 --- a/aten/src/ATen/cpu/vec256/vec256_base.h +++ b/aten/src/ATen/cpu/vec256/vec256_base.h @@ -13,6 +13,7 @@ #include #include #include +#include #if defined(__GNUC__) #define __at_align32__ __attribute__((aligned(32))) @@ -681,8 +682,7 @@ inline void convert(const src_T *src, dst_T *dst, int64_t n) { # pragma unroll #endif for (int64_t i = 0; i < n; i++) { - *dst = static_cast( - static_cast>(*src)); + *dst = c10::static_cast_with_inter_type(*src); src++; dst++; } diff --git a/aten/src/ATen/native/Copy.h b/aten/src/ATen/native/Copy.h index a8d16f6f7b871..2dfd9e9f4922b 100644 --- a/aten/src/ATen/native/Copy.h +++ b/aten/src/ATen/native/Copy.h @@ -9,44 +9,6 @@ struct TensorIterator; namespace native { -// Note [Implicit conversion between signed and unsigned] -// C and C++ have a lovely set of implicit conversion rules, where casting -// signed integral values to unsigned integral values is always valid -// (it basically treats the value as if using modulo arithmetic), however -// converting negative floating point values to unsigned integral types -// is UB! This means that: (double)-1 -> (int64_t)-1 -> (uint8_t)255 is -// guaranteed to look like this, but we have (double)-1 -> (uint8_t) -// because it's UB. This also makes UBSan really angry. -// -// I think those rules are stupid and we really shouldn't conform to them. -// The structs below ensure that for all unsigned types we use (currently -// only uint8_t), we will do an intermediate convertion via int64_t, -// to ensure that any negative values are wrapped around correctly. -// -// Note that conversions from doubles to signed integral types that can't -// represent a particular value after truncating the fracitonal part are UB as well, -// but fixing them is not as simple as adding an int64_t intermediate, beacuse the -// int64_t -> conversion is UB for those large values anyway. -// I guess in that case we just have to live with that, but it's definitely less -// surprising than the thing above. -// -// For the curious: -// https://en.cppreference.com/w/cpp/language/implicit_conversion -// The relevant paragraph is "Floating-integral conversions". - -template -struct inter_copy_type { - using type = T; -}; - -template <> -struct inter_copy_type { - using type = int64_t; -}; - -template -using inter_copy_type_t = typename inter_copy_type::type; - using copy_fn = void (*)(TensorIterator&, bool non_blocking); DECLARE_DISPATCH(copy_fn, copy_stub); diff --git a/aten/src/ATen/native/cpu/CopyKernel.cpp b/aten/src/ATen/native/cpu/CopyKernel.cpp index 2362b8f20a9c6..a48c7afc07855 100644 --- a/aten/src/ATen/native/cpu/CopyKernel.cpp +++ b/aten/src/ATen/native/cpu/CopyKernel.cpp @@ -4,6 +4,7 @@ #include #include #include +#include namespace at { namespace native { @@ -14,8 +15,7 @@ void copy_kernel_cast(TensorIterator& iter) { if (isComplexType(iter.dtype(1))) { AT_DISPATCH_COMPLEX_TYPES(iter.dtype(1), "copy_kernel_cast", [&] { cpu_kernel(iter, [=](scalar_t a) -> self_T { - return static_cast( - static_cast>(std::real(a))); + return c10::static_cast_with_inter_type(std::real(a)); }); }); } @@ -28,8 +28,7 @@ void copy_kernel_cast(TensorIterator& iter) { "copy_kernel_cast", [&] { cpu_kernel(iter, [=](scalar_t a) -> self_T { - return static_cast( - static_cast>(a)); + return c10::static_cast_with_inter_type(a); }); }); } diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu index 4e201a6576af3..407d5df90ca5f 100644 --- a/aten/src/ATen/native/cuda/Copy.cu +++ b/aten/src/ATen/native/cuda/Copy.cu @@ -8,6 +8,7 @@ #include #include #include +#include namespace at { namespace native { @@ -17,7 +18,7 @@ using namespace at::cuda; template void copy_kernel_impl(TensorIterator& iter) { gpu_kernel(iter, []GPU_LAMBDA(src_t x) -> dst_t { - return static_cast(static_cast>(x)); + return c10::static_cast_with_inter_type(x); }); } diff --git a/c10/core/ScalarType.h b/c10/core/ScalarType.h index de5a0d2f20de2..b9be182946b4c 100644 --- a/c10/core/ScalarType.h +++ b/c10/core/ScalarType.h @@ -166,6 +166,10 @@ struct ScalarTypeToCPPType { _(c10::quint8, QUInt8) \ _(c10::qint32, QInt32) +#define AT_FORALL_COMPLEX_TYPES(_) \ + _(std::complex, ComplexFloat) \ + _(std::complex, ComplexDouble) + static inline caffe2::TypeMeta scalarTypeToTypeMeta(ScalarType scalar_type) { #define DEFINE_CASE(ctype, name) \ case ScalarType::name: \ diff --git a/c10/util/TypeCast.h b/c10/util/TypeCast.h new file mode 100644 index 0000000000000..1a00a0abb7315 --- /dev/null +++ b/c10/util/TypeCast.h @@ -0,0 +1,172 @@ +#pragma once + +#include +#include +#include + + +namespace c10 { + +// Note [Implicit conversion between signed and unsigned] +// C and C++ have a lovely set of implicit conversion rules, where casting +// signed integral values to unsigned integral values is always valid +// (it basically treats the value as if using modulo arithmetic), however +// converting negative floating point values to unsigned integral types +// is UB! This means that: (double)-1 -> (int64_t)-1 -> (uint8_t)255 is +// guaranteed to look like this, but we have (double)-1 -> (uint8_t) +// because it's UB. This also makes UBSan really angry. +// +// I think those rules are stupid and we really shouldn't conform to them. +// The structs below ensure that for all unsigned types we use (currently +// only uint8_t), we will do an intermediate convertion via int64_t, +// to ensure that any negative values are wrapped around correctly. +// +// Note that conversions from doubles to signed integral types that can't +// represent a particular value after truncating the fracitonal part are UB as well, +// but fixing them is not as simple as adding an int64_t intermediate, beacuse the +// int64_t -> conversion is UB for those large values anyway. +// I guess in that case we just have to live with that, but it's definitely less +// surprising than the thing above. +// +// For the curious: +// https://en.cppreference.com/w/cpp/language/implicit_conversion +// The relevant paragraph is "Floating-integral conversions". + +template +struct inter_copy_type { + using type = T; +}; + +template <> +struct inter_copy_type { + using type = int64_t; +}; + +template +using inter_copy_type_t = typename inter_copy_type::type; + +template +C10_HOST_DEVICE inline dest_t static_cast_with_inter_type(src_t src) { + return static_cast( + static_cast>(src)); +} + +// Dynamic type casting utils: +// - fetch_and_cast +// - cast_and_store +// +// fetch_and_cast fetch a value with dynamic type specified by a ScalarType +// from a void pointer and cast it to a static type. +// +// cast_and_store casts a static typed value into dynamic type specified +// by a ScalarType, and store it into a void pointer. +// +// NOTE: +// +// Dynamic casting allows us to support type promotion without blowing up +// the combination space: For example, without dynamic cast, in order to +// implement `add_` with type promotion, we would need something like +// +// AT_DISPATCH_ALL_TYPES(output.dtype(), +// AT_DISPATCH_ALL_TYPES(input1.dtype(), +// AT_DISPATCH_ALL_TYPES(input2.dtype(), +// [](arg0_t a, arg1_t b) -> out_t { return a + b; } +// ) +// ) +// ) +// +// If we support N dtypes, the above code would generate the a+b kernel for +// all the N * N * N different supported types, the compilation time and +// binary size would become horrible. +// +// Dynamic casting might sounds like a bad idea in terms of performance. +// Especially if you ever do it in a loop, you are going to do a billion tests. +// But in practice it is not as bad as it might look: +// +// - on CPU, this is a branch that always has the same outcome, therefore +// hopefully the branch predictor could do the job pretty well +// - on GPU, these branches will not diverge, so we could still have the same +// warp executing the same line of code +// - Most kernels, like `add`, are bandwidth bound, adding a few clock cycles to +// check an integer does not hurt the performance much because the ALUs would +// wait for load instructions anyway. +// +// For the discussion and benchmark, refer to: +// - https://github.com/pytorch/pytorch/pull/28343 +// - https://github.com/pytorch/pytorch/pull/28344 +// - https://github.com/pytorch/pytorch/pull/28345 +// + +#ifdef C10_HOST_DEVICE +#define ERROR_UNSUPPORTED_CAST assert(false); +#else +#define ERROR_UNSUPPORTED_CAST TORCH_CHECK(false, "Unexpected scalar type"); +#endif + +// Fetch a value with dynamic type src_type from ptr, and cast it to static type dest_t. +#define FETCH_AND_CAST_CASE(type, scalartype) case ScalarType::scalartype: return static_cast_with_inter_type(*(const type *)ptr); +#define FETCH_AND_CAST_COMPLEX_CASE(type, scalartype) case ScalarType::scalartype: return static_cast_with_inter_type(std::real(*(const type *)ptr)); +template +C10_HOST_DEVICE inline dest_t fetch_and_cast(const ScalarType src_type, const void *ptr) { + switch (src_type) { + AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, FETCH_AND_CAST_CASE) +#ifndef C10_HOST_DEVICE + AT_FORALL_COMPLEX_TYPES(FETCH_AND_CAST_COMPLEX_CASE) +#endif + default:; + } + ERROR_UNSUPPORTED_CAST + return dest_t(0); // just to avoid compiler warning +} + +// Cast a value with static type src_t into dynamic dest_type, and store it to ptr. +#define CAST_AND_STORE_CASE(type, scalartype) case ScalarType::scalartype: *(type *)ptr = static_cast_with_inter_type(value); return; +template +C10_HOST_DEVICE inline void cast_and_store(const ScalarType dest_type, void *ptr, src_t value) { + switch (dest_type) { + AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) + default:; + } + ERROR_UNSUPPORTED_CAST +} + +template<> +inline void cast_and_store>(const ScalarType dest_type, void *ptr, std::complex value_) { + auto value = std::real(value_); + switch (dest_type) { + AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) + default:; + } + ERROR_UNSUPPORTED_CAST +} +template<> +inline void cast_and_store>(const ScalarType dest_type, void *ptr, std::complex value_) { + auto value = std::real(value_); + switch (dest_type) { + AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) + default:; + } + ERROR_UNSUPPORTED_CAST +} + +#define DEFINE_UNCASTABLE(T, scalartype_) \ +template<> \ +inline T fetch_and_cast(const ScalarType src_type, const void *ptr) { \ + assert(ScalarType::scalartype_ == src_type); \ + return *(const T *)ptr; \ +} \ +template<> \ +inline void cast_and_store(const ScalarType dest_type, void *ptr, T value) { \ + assert(ScalarType::scalartype_ == dest_type); \ + *(T *)ptr = value; \ +} + +AT_FORALL_QINT_TYPES(DEFINE_UNCASTABLE) + +#undef FETCH_AND_CAST_CASE +#undef FETCH_AND_CAST_COMPLEX_CASE +#undef CAST_AND_STORE_CASE +#undef DEFINE_UNCASTABLE +#undef ERROR_UNSUPPORTED_CAST + +} // namespace c10 From b9f099ed936a116ffd1c0ebb0d065b50fa753332 Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Mon, 28 Oct 2019 14:43:51 -0700 Subject: [PATCH 06/64] Make TensorIterator stop promoting types by copying (#28427) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28427 Fixes: https://github.com/pytorch/pytorch/issues/26401 This PR fixes the issue by using the newly added dynamic cast inside `TensorIterator` so that instead of converting the type at the beginning (which generates extra kernel launches), the `TensorIterator` do a load-cast-compute-store for each element while looping. So there is only one read and one write of memory. **nvprof:** ```python import torch _100M = 100 * 1024 ** 2 r = torch.randn(_100M, dtype=torch.float32, device='cuda') d = torch.randn(_100M, dtype=torch.float64, device='cuda') torch.cuda.synchronize() torch.cuda.profiler.start() r.add_(d) torch.cuda.profiler.stop() torch.cuda.synchronize() ``` ``` ==11407== NVPROF is profiling process 11407, command: /home/xgao/anaconda3/bin/python simple.py ==11407== Profiling application: /home/xgao/anaconda3/bin/python simple.py ==11407== Profiling result: Type Time(%) Time Calls Avg Min Max Name GPU activities: 100.00% 2.0611ms 1 2.0611ms 2.0611ms 2.0611ms _ZN2at6native18elementwise_kernelILi512ELi1EZNS0_15gpu_kernel_implIZZZNS0_15add_kernel_cudaERNS_14TensorIteratorEN3c106ScalarEENKUlvE_clEvENKUlvE1_clEvEUlddE_EEvS4_RKT_EUliE_EEviT1_ API calls: 100.00% 1.05006s 1 1.05006s 1.05006s 1.05006s cudaLaunchKernel 0.00% 2.7740us 2 1.3870us 673ns 2.1010us cudaGetDevice 0.00% 2.3730us 1 2.3730us 2.3730us 2.3730us cudaSetDevice 0.00% 830ns 1 830ns 830ns 830ns cudaGetLastError ``` **benchmark** ```python import torch print(torch.__version__) print(torch.version.git_version) _100M = 100 * 1024 ** 2 r = torch.randn(_100M, dtype=torch.float32, device='cuda') d = torch.randn(_100M, dtype=torch.float64, device='cuda') torch.cuda.synchronize() %timeit r.add_(d); torch.cuda.synchronize() ``` original ``` 1.4.0a0+7d277b0 7d277b0670eb1f9098a7e098e93b20453e8b5c9f 6.83 ms ± 1.12 ms per loop (mean ± std. dev. of 7 runs, 100 loops each) ``` after ``` 1.4.0a0+f0f2f65 f0f2f654cba9b8c569f0bcd583732bbc891f80b2 2.08 ms ± 139 ns per loop (mean ± std. dev. of 7 runs, 100 loops each) ``` For more benchmark, see: https://github.com/pytorch/pytorch/pull/28344 Test Plan: Imported from OSS Differential Revision: D18170997 Pulled By: ezyang fbshipit-source-id: 9c82c1c89583f3e6202c5d790b9b73ad9f960fad --- aten/src/ATen/native/TensorIterator.cpp | 31 ++++++----- aten/src/ATen/native/TensorIterator.h | 7 +++ aten/src/ATen/native/cuda/BinaryOpsKernel.cu | 42 +++++++-------- aten/src/ATen/native/cuda/Loops.cuh | 56 ++++++++++++++++---- aten/src/ATen/test/tensor_iterator_test.cpp | 1 + 5 files changed, 93 insertions(+), 44 deletions(-) diff --git a/aten/src/ATen/native/TensorIterator.cpp b/aten/src/ATen/native/TensorIterator.cpp index 5d8b94f395e1c..cbf9cdbc83fed 100644 --- a/aten/src/ATen/native/TensorIterator.cpp +++ b/aten/src/ATen/native/TensorIterator.cpp @@ -148,7 +148,7 @@ static void validate_dtype(OperandInfo& op, ScalarType common_dtype, CommonDType } } -static void maybe_promote_common_dtype(OperandInfo& op, ScalarType common_dtype) { +static void maybe_copy_casting_to_common_dtype(OperandInfo& op, ScalarType common_dtype) { if (op.tensor.defined() && op.tensor.scalar_type() != common_dtype) { op.dtype = common_dtype; @@ -165,7 +165,7 @@ static void maybe_promote_common_dtype(OperandInfo& op, ScalarType common_dtype) void TensorIterator::compute_types() { bool missing_dtypes = false; bool missing_output_dtypes = false; - ScalarType common_dtype = dtype(); + common_dtype_ = dtype(); for (auto& op : operands_) { if (!op.tensor.defined() && !op.is_type_defined()) { missing_dtypes = true; @@ -183,22 +183,24 @@ void TensorIterator::compute_types() { bool compute_common_dtype_only_for_inputs = (common_dtype_strategy_ == CommonDTypeStrategy::PROMOTE_INPUTS); bool may_have_differing_types = true; + bool common_device_is_cuda = false; if (missing_dtypes || compute_common_dtype) { auto operands = compute_common_dtype_only_for_inputs ? at::ArrayRef(operands_).slice(noutputs()) : operands_; auto common_type = compute_common_type_(operands); auto common_device = std::get<0>(common_type); - common_dtype = std::get<1>(common_type); + common_device_is_cuda = common_device.is_cuda(); + common_dtype_ = std::get<1>(common_type); may_have_differing_types = !std::get<2>(common_type); bool has_cpu_scalar = false; for (auto& op : operands_) { if (!op.is_type_defined()) { op.device = common_device; - op.dtype = common_dtype; + op.dtype = common_dtype_; } else if (compute_common_dtype && - (op.device != common_device || op.dtype != common_dtype)) { + (op.device != common_device || op.dtype != common_dtype_)) { if (allow_cpu_scalars_ && op.tensor.defined() && op.tensor.dim() == 0 && - common_device.is_cuda() && op.tensor.device().is_cpu() && + common_device_is_cuda && op.tensor.device().is_cpu() && !has_cpu_scalar) { // don't cast CPU scalars in CUDA ops that directly support them. op.device = op.tensor.device(); @@ -206,8 +208,8 @@ void TensorIterator::compute_types() { has_cpu_scalar = true; } else if (promote_gpu_output_dtypes_ && op.tensor.defined() && !op.is_output && - op.tensor.scalar_type() == kHalf && common_dtype == kFloat && - op.tensor.device().is_cuda() && common_device.is_cuda()) { + op.tensor.scalar_type() == kHalf && common_dtype_ == kFloat && + op.tensor.device().is_cuda() && common_device_is_cuda) { // allow input tensor type upcasting for fp16 to fp32 in fused kernel // on GPU op.device = op.tensor.device(); @@ -217,7 +219,7 @@ void TensorIterator::compute_types() { if (compute_common_dtype_only_for_inputs && op.is_output) { op.dtype = op.tensor.scalar_type(); } else { - op.dtype = common_dtype; + op.dtype = common_dtype_; } } } @@ -226,12 +228,17 @@ void TensorIterator::compute_types() { for (auto &op : operands_) { if (may_have_differing_types) { - validate_dtype(op, common_dtype, common_dtype_strategy_); - if (compute_common_dtype && (!compute_common_dtype_only_for_inputs || !op.is_output)) { - maybe_promote_common_dtype(op, common_dtype); + validate_dtype(op, common_dtype_, common_dtype_strategy_); + bool cast_by_copy = compute_common_dtype && !common_device_is_cuda && (!compute_common_dtype_only_for_inputs || !op.is_output); + if (cast_by_copy) { + maybe_copy_casting_to_common_dtype(op, common_dtype_); } } + if (op.tensor.defined() && op.tensor.scalar_type() != common_dtype_) { + have_differing_types_ = true; + } + if (op.tensor.defined() && op.device != op.tensor.device()) { if (op.is_output) { AT_ERROR("output with device ", op.tensor.device(), diff --git a/aten/src/ATen/native/TensorIterator.h b/aten/src/ATen/native/TensorIterator.h index 69ed861c888c7..eef28b710cd4a 100644 --- a/aten/src/ATen/native/TensorIterator.h +++ b/aten/src/ATen/native/TensorIterator.h @@ -191,6 +191,7 @@ struct CAFFE2_API TensorIterator { IntArrayRef strides(int arg) const { return operands_[arg].stride_bytes; } void* data_ptr(int arg) const; ScalarType dtype(int arg=0) const { return operands_[arg].tensor.scalar_type(); } + ScalarType common_dtype() const { return common_dtype_; } ScalarType input_dtype(int arg=0) const { return operands_[num_outputs_ + arg].dtype; } Device device(int arg=0) const { return operands_[arg].device; } DeviceType device_type(int arg=0) const { return device(arg).type(); } @@ -286,6 +287,10 @@ struct CAFFE2_API TensorIterator { /// CUDA reductions. bool is_final_output() const { return final_output_; } + bool needs_dynamic_casting() const { + return (common_dtype_strategy_ != CommonDTypeStrategy::NONE) && have_differing_types_; + } + void set_check_mem_overlap(bool check_mem_overlap) { check_mem_overlap_ = check_mem_overlap; } @@ -352,6 +357,7 @@ struct CAFFE2_API TensorIterator { SmallVector operands_; int num_outputs_ = 0; CommonDTypeStrategy common_dtype_strategy_ = CommonDTypeStrategy::CHECK; + ScalarType common_dtype_ = ScalarType::Undefined; bool has_coalesced_dimensions_ = false; bool accumulate_ = false; bool resize_outputs_ = true; @@ -360,6 +366,7 @@ struct CAFFE2_API TensorIterator { bool promote_gpu_output_dtypes_ = false; bool final_output_ = true; bool check_mem_overlap_ = false; + bool have_differing_types_ = false; }; /// A container-like struct that acts as if it contains splits of a /// TensorIterator that can use 32-bit indexing. Taken together the splits cover diff --git a/aten/src/ATen/native/cuda/BinaryOpsKernel.cu b/aten/src/ATen/native/cuda/BinaryOpsKernel.cu index 002539717de44..95f1a75883b56 100644 --- a/aten/src/ATen/native/cuda/BinaryOpsKernel.cu +++ b/aten/src/ATen/native/cuda/BinaryOpsKernel.cu @@ -13,7 +13,7 @@ namespace at { namespace native { void add_kernel_cuda(TensorIterator& iter, Scalar alpha_scalar) { - AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(), "add_cuda/sub_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "add_cuda/sub_cuda", [&]() { auto alpha = alpha_scalar.to(); gpu_kernel_with_scalars(iter, [alpha]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a + alpha * b; @@ -26,11 +26,11 @@ static void sub_kernel_cuda(TensorIterator& iter, Scalar alpha_scalar) { } void div_kernel_cuda(TensorIterator& iter) { - if (!isIntegralType(iter.dtype(), /*includeBool*/ false) && iter.is_cpu_scalar(2)) { + if (!isIntegralType(iter.common_dtype(), /*includeBool*/ false) && iter.is_cpu_scalar(2)) { // optimization for floating-point types: if the second operand is a CPU // scalar, compute a * reciprocal(b). Note that this may lose one bit of // precision compared to computing the division. - AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.dtype(), "div_cuda", [&]() { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "div_cuda", [&]() { auto inv_b = scalar_t(1.0 / iter.scalar_value(2)); iter.remove_operand(2); gpu_kernel(iter, [inv_b]GPU_LAMBDA(scalar_t a) -> scalar_t { @@ -38,7 +38,7 @@ void div_kernel_cuda(TensorIterator& iter) { }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "div_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "div_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a / b; }); @@ -47,13 +47,13 @@ void div_kernel_cuda(TensorIterator& iter) { } void mul_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { // Workaround for the error: '*' in boolean context, suggest '&&' instead [-Werror=int-in-bool-context] gpu_kernel_with_scalars(iter, []GPU_LAMBDA(bool a, bool b) -> bool { return a && b; }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "mul_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "mul_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a * b; }); @@ -62,7 +62,7 @@ void mul_kernel_cuda(TensorIterator& iter) { } void atan2_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.dtype(), "atan2_cuda", [&]() { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "atan2_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return ::atan2(a, b); }); @@ -70,14 +70,14 @@ void atan2_kernel_cuda(TensorIterator& iter) { } void logical_xor_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "logical_xor_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return bool(a) != bool(b); }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "logical_xor_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "logical_xor_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return static_cast(bool(a) != bool(b)); }); @@ -86,14 +86,14 @@ void logical_xor_kernel_cuda(TensorIterator& iter) { } void lt_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "lt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a < b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "lt_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "lt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a < b; }); @@ -102,14 +102,14 @@ void lt_kernel_cuda(TensorIterator& iter) { } void le_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "le_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a <= b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "le_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "le_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a <= b; }); @@ -118,14 +118,14 @@ void le_kernel_cuda(TensorIterator& iter) { } void gt_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "gt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a > b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "gt_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "gt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a > b; }); @@ -134,14 +134,14 @@ void gt_kernel_cuda(TensorIterator& iter) { } void ge_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "ge_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a >= b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "ge_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "ge_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a >= b; }); @@ -150,14 +150,14 @@ void ge_kernel_cuda(TensorIterator& iter) { } void eq_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "eq_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a == b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "eq_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a == b; }); @@ -166,14 +166,14 @@ void eq_kernel_cuda(TensorIterator& iter) { } void ne_kernel_cuda(TensorIterator& iter) { - if (iter.dtype() == ScalarType::Bool) { + if (iter.common_dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "ne_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a != b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "ne_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a != b; }); diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh index 38d25982fe0b5..e5b52accbb253 100644 --- a/aten/src/ATen/native/cuda/Loops.cuh +++ b/aten/src/ATen/native/cuda/Loops.cuh @@ -35,6 +35,7 @@ #include #include #include +#include // Marks a lambda as executable on both the host and device. The __host__ // attribute is important so that we can access static type information from @@ -116,6 +117,20 @@ invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[] return invoke_impl(f, data, strides, i, Indices{}); } +template +C10_HOST_DEVICE typename traits::result_type +invoke_impl(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i, + c10::guts::index_sequence) { + return f(c10::fetch_and_cast::type>(dtypes[I], data[I] + i * strides[I])...); +} + +template > +C10_HOST_DEVICE typename traits::result_type +invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i) { + using Indices = c10::guts::make_index_sequence; + return invoke_impl(f, data, strides, dtypes, i, Indices{}); +} + template void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { using traits = function_traits; @@ -130,6 +145,10 @@ void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { data[i] = (char*)iter.data_ptr(i); } + at::detail::Array dtypes; + for (int i = 0; i < ntensors; i++) { + dtypes[i] = iter.tensor(i).scalar_type(); + } int64_t numel = iter.numel(); if (iter.is_trivial_1d()) { @@ -138,19 +157,35 @@ void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { for (int i = 0; i < ntensors; i++) { strides[i] = inner_strides[i]; } - - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - arg0_t* out = (arg0_t*)(data[0] + strides[0] * idx); - *out = invoke(f, &data.data[1], &strides.data[1], idx); - }); + if (iter.needs_dynamic_casting()) { + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + void* out = data[0] + strides[0] * idx; + arg0_t result = invoke(f, &data.data[1], &strides.data[1], &dtypes.data[1], idx); + c10::cast_and_store(dtypes[0], out, result); + }); + } else { + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + arg0_t* out = (arg0_t*)(data[0] + strides[0] * idx); + *out = invoke(f, &data.data[1], &strides.data[1], idx); + }); + } } else { auto offset_calc = make_offset_calculator(iter); - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - auto offsets = offset_calc.get(idx); - arg0_t* out = (arg0_t*)(data[0] + offsets[0]); - *out = invoke(f, &data.data[1], &offsets.data[1], 1); - }); + if (iter.needs_dynamic_casting()) { + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + auto offsets = offset_calc.get(idx); + void* out = data[0] + offsets[0]; + arg0_t result = invoke(f, &data.data[1], &offsets.data[1], &dtypes.data[1], 1); + c10::cast_and_store(dtypes[0], out, result); + }); + } else { + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + auto offsets = offset_calc.get(idx); + arg0_t* out = (arg0_t*)(data[0] + offsets[0]); + *out = invoke(f, &data.data[1], &offsets.data[1], 1); + }); + } } } @@ -174,7 +209,6 @@ void gpu_kernel(TensorIterator& iter, const func_t& f) { } gpu_kernel_impl(iter, f); - iter.cast_outputs(); } template diff --git a/aten/src/ATen/test/tensor_iterator_test.cpp b/aten/src/ATen/test/tensor_iterator_test.cpp index 200091de260eb..a2bbc653298fa 100644 --- a/aten/src/ATen/test/tensor_iterator_test.cpp +++ b/aten/src/ATen/test/tensor_iterator_test.cpp @@ -190,6 +190,7 @@ TEST(TensorIteratorTest, ComputeCommonDTypeInputOnly) { EXPECT_TRUE(iter.dtype(0) == at::kBool); EXPECT_TRUE(iter.dtype(1) == at::kDouble); EXPECT_TRUE(iter.dtype(2) == at::kDouble); + EXPECT_TRUE(iter.common_dtype() == at::kDouble); } TEST(TensorIteratorTest, DoNotComputeCommonDTypeInputOnly) { From 5c5b2c68dba8577e584736ad6240d4e7a29c4d2a Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Mon, 28 Oct 2019 14:43:51 -0700 Subject: [PATCH 07/64] Simplify copy kernel (#28428) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28428 Using the new type promotion and dynamic casting added to `TensorIterator`, the copy kernels could be greatly simplified. Benchmark on CUDA: ```python import torch import timeit import pandas import itertools from tqdm.notebook import tqdm import math print(torch.__version__) print() _10M = 10 * 1024 ** 2 d = {} for from_, to in tqdm(itertools.product(torch.testing.get_all_dtypes(), repeat=2)): if from_ not in d: d[from_] = {} a = torch.empty(_10M, dtype=from_, device='cuda') min_ = math.inf for i in range(100): torch.cuda.synchronize() start = timeit.default_timer() a.to(to) torch.cuda.synchronize() end = timeit.default_timer() elapsed = end - start if elapsed < min_: min_ = elapsed d[from_][to] = int(min_ * 1000 * 1000) pandas.DataFrame(d) ``` original: ![image](https://user-images.githubusercontent.com/1032377/67623519-e3e6dd80-f7da-11e9-86ea-9cc9f237123b.png) new: ![image](https://user-images.githubusercontent.com/1032377/67623527-fc56f800-f7da-11e9-82bd-dc1ff9821b68.png) Test Plan: Imported from OSS Differential Revision: D18170995 Pulled By: ezyang fbshipit-source-id: 461b53641813dc6cfa872a094ae917e750c60759 --- aten/src/ATen/native/cuda/Copy.cu | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu index 407d5df90ca5f..b9c90ec156cae 100644 --- a/aten/src/ATen/native/cuda/Copy.cu +++ b/aten/src/ATen/native/cuda/Copy.cu @@ -8,20 +8,12 @@ #include #include #include -#include namespace at { namespace native { using namespace at::cuda; -template -void copy_kernel_impl(TensorIterator& iter) { - gpu_kernel(iter, []GPU_LAMBDA(src_t x) -> dst_t { - return c10::static_cast_with_inter_type(x); - }); -} - // device-to-device copy, does type conversion static void copy_device_to_device(TensorIterator& iter, bool non_blocking) { int64_t numel = iter.numel(); @@ -66,11 +58,11 @@ static void copy_device_to_device(TensorIterator& iter, bool non_blocking) { cudaMemcpyDeviceToDevice, copy_stream)); } else { - AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(0), "copy_", [&] { - using dst_t = scalar_t; - AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(1), "copy_", [&] { - copy_kernel_impl(iter); - }); + // this is done intentionally done after build because copy has a "promotion" + // rule that always "promote" to target dtype. + iter.promote_common_dtype(); + AT_DISPATCH_ALL_TYPES_AND3(kHalf, kBool, kBFloat16, iter.dtype(0), "copy_", [&] { + gpu_kernel(iter, []GPU_LAMBDA(scalar_t x) { return x; }); }); } From f63cf96c4d2922f37abbbac65cb55a5a26fb8358 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Mon, 28 Oct 2019 14:53:53 -0700 Subject: [PATCH 08/64] Update C++ parity table for torch::nn::Linear (#28804) Summary: Since we have merged https://github.com/pytorch/pytorch/pull/27382 (thanks pbelevich!) Pull Request resolved: https://github.com/pytorch/pytorch/pull/28804 Differential Revision: D18185714 Pulled By: yf225 fbshipit-source-id: 1148f5837fbf578843b989fc53fd334519943cdd --- test/cpp_api_parity/parity-tracker.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index a326b54c90072..ba9eb78a75458 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -91,7 +91,7 @@ torch.nn.TransformerDecoder|No|No torch.nn.TransformerEncoderLayer|No|No torch.nn.TransformerDecoderLayer|No|No torch.nn.Identity|Yes|No -torch.nn.Linear|No|No +torch.nn.Linear|Yes|No torch.nn.Bilinear|Yes|No torch.nn.Flatten|No|No torch.nn.Dropout|No|No From 949678bd9e82c12efe77fb0fb5471a17ded128f4 Mon Sep 17 00:00:00 2001 From: James Reed Date: Mon, 28 Oct 2019 16:43:20 -0700 Subject: [PATCH 09/64] Small fixes for torchbind (#28800) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28800 Fix up namespaces and make friendly error message when registered class doesn't inherit from the right base Test Plan: Imported from OSS Differential Revision: D18175067 Pulled By: jamesr66a fbshipit-source-id: 5c7cf3a49fb45db502d84eb3f9a69be126ee59fb --- torch/custom_class.h | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/torch/custom_class.h b/torch/custom_class.h index c28f6b14d94ea..3b9c73c9da7ef 100644 --- a/torch/custom_class.h +++ b/torch/custom_class.h @@ -60,6 +60,9 @@ detail::types init() { return detail::types{}; } template class class_ { + static_assert(std::is_base_of::value, + "torch::jit::class_ requires T to inherit from CustomClassHolder"); + std::string className; std::string qualClassName; c10::optional> pyClass = c10::nullopt; @@ -70,7 +73,7 @@ class class_ { const std::string topModule = "__torch__.torch"; public: - class_(string className_) : className(std::move(className_)) { + class_(std::string className_) : className(std::move(className_)) { // Currently we register everything as a python class just for convenience. // We'll want to remove this at some point to get rid of the python // dependency. It would require significant changes to class registration, @@ -90,7 +93,7 @@ class class_ { PyObject* rawPyObj = py_object.release().ptr(); return rawPyObj; }; - getClassConverter()[qualClassName] = castToPython; + at::getClassConverter()[qualClassName] = castToPython; // We currently represent custom classes as torchscript classes with a // capsule attribute @@ -100,9 +103,9 @@ class class_ { classTypePtr->addAttribute("capsule", CapsuleType::get()); c10::getCustomClassTypeMap().insert({typeid(c10::intrusive_ptr).name(), - StrongTypePtr(classCu, classTypePtr)}); + c10::StrongTypePtr(classCu, classTypePtr)}); c10::getCustomClassTypeMap().insert({typeid(c10::tagged_capsule).name(), - StrongTypePtr(classCu, classTypePtr)}); + c10::StrongTypePtr(classCu, classTypePtr)}); classCu->register_type(classTypePtr); } @@ -124,7 +127,7 @@ class class_ { return *this; } template - class_& def(string name, Func f) { + class_& def(std::string name, Func f) { auto res = def_(name, f, detail::args_t{}); return *this; } @@ -140,19 +143,19 @@ class class_ { std::vector addInputs_( Func f, std::shared_ptr graph, - guts::index_sequence) { + at::guts::index_sequence) { using argTypes = - typename guts::infer_function_traits_t::parameter_types; + typename at::guts::infer_function_traits_t::parameter_types; std::vector res = { - addInput>::call( + addInput>::call( graph)...}; return res; } template std::vector addInputs(Func f, std::shared_ptr graph) { constexpr auto numArgs = - guts::infer_function_traits_t::number_of_parameters; - return addInputs_(f, graph, guts::make_index_sequence()); + at::guts::infer_function_traits_t::number_of_parameters; + return addInputs_(f, graph, at::guts::make_index_sequence()); } template @@ -192,11 +195,11 @@ class class_ { classTypePtr->addMethod(method); } template - class_& def_(string name, Func f, detail::types funcInfo) { + class_& def_(std::string name, Func f, detail::types funcInfo) { pyClass->def(name.c_str(), f); auto func = [f](c10::intrusive_ptr cur, Types... args) { - return guts::invoke(f, *cur, args...); + return at::guts::invoke(f, *cur, args...); }; defineMethod(name, std::move(func), funcInfo.hasRet); return *this; From 6f90567e0c73b2436a8ae5525f017c134aaf54db Mon Sep 17 00:00:00 2001 From: Jianyu Huang Date: Mon, 28 Oct 2019 17:35:51 -0700 Subject: [PATCH 10/64] Add the unittest import for test_fake_quant.py (#28815) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28815 Add the unittest import ghstack-source-id: 92789329 Test Plan: CI Differential Revision: D18191989 fbshipit-source-id: c54e0309e21156c33e4fec01bfba17a1c30463c9 --- test/test_fake_quant.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_fake_quant.py b/test/test_fake_quant.py index bcd413d69f4a2..0a4443717a270 100644 --- a/test/test_fake_quant.py +++ b/test/test_fake_quant.py @@ -10,12 +10,13 @@ from torch.quantization import FakeQuantize from torch.quantization import default_observer, default_per_channel_weight_observer import io +import unittest + # Reference method for fake quantize def _fake_quantize_per_tensor_affine_reference(X, scale, zero_point, quant_min, quant_max): res = (torch.clamp(torch.round(X.cpu() * (1.0 / scale) + zero_point), quant_min, quant_max) - zero_point) * scale return res - # Reference method for the gradient of the fake quantize operator def _fake_quantize_per_tensor_affine_grad_reference(dY, X, scale, zero_point, quant_min, quant_max): Xq = torch.round(X.cpu() * (1.0 / scale) + zero_point) From c6ad68cf10cb35d6b9f53cc3060ad98d8cd1f82a Mon Sep 17 00:00:00 2001 From: svcscm Date: Mon, 28 Oct 2019 17:39:43 -0700 Subject: [PATCH 11/64] Updating submodules Summary: GitHub commits: https://github.com/facebook/fbthrift/commit/724e939772a49a7d18466aad787437c6c76184cc https://github.com/facebook/folly/commit/f4fb4266c044a959f8f01bee88c2f502939f527f https://github.com/facebook/proxygen/commit/95d4b19724ebb58b9c3faea0ab8ff5045efd7928 https://github.com/facebookincubator/mvfst/commit/8b8131450ee24587687f9284485aa8b5413e5a23 https://github.com/facebookincubator/profilo/commit/ac8faa6528a5784c683904b422c5843c2b5a0bf6 https://github.com/pytorch/fbgemm/commit/5487e2b1a2184ee41f49a2935922dc8ed92f6464 Test Plan: n/a Reviewed By: zpao fbshipit-source-id: 9b9f4cccd869638215c17111361a6f6c480c73af --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index 98141ffe1b165..5487e2b1a2184 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 98141ffe1b1657459512b621d622c2cf3a1537a3 +Subproject commit 5487e2b1a2184ee41f49a2935922dc8ed92f6464 From 1e3e1f5bf9745d3c82dc0834064e443d6d142785 Mon Sep 17 00:00:00 2001 From: Linbin Yu Date: Mon, 28 Oct 2019 17:43:49 -0700 Subject: [PATCH 12/64] Fix build error in VariableTypeManual Summary: build error in internal pt mobile build ``` xplat/caffe2/torch/csrc/autograd/VariableTypeManual.cpp:118:49: error: address of function 'requires_grad' will always evaluate to 'true' [-Werror,-Wpointer-bool-conversion] autograd::utils::requires_grad_leaf_error(requires_grad) ~~~~~~~~ ^~~~~~~~~~~~~ xplat/caffe2/torch/csrc/autograd/VariableTypeManual.cpp:118:49: note: prefix with the address-of operator to silence this warning ``` I think the variable name in requires_grad_leaf_error is wrong. Test Plan: mobile build works Reviewed By: pbelevich Differential Revision: D18192663 fbshipit-source-id: a3d3ebb9039022eb228c1d183a1076f65f9e84e0 --- torch/csrc/autograd/VariableTypeManual.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp index b141003256807..be1fa2b076e95 100644 --- a/torch/csrc/autograd/VariableTypeManual.cpp +++ b/torch/csrc/autograd/VariableTypeManual.cpp @@ -115,7 +115,7 @@ int64_t _version(const Tensor & self) { Tensor& requires_grad_(Tensor& self, bool _requires_grad) { if (!self.is_leaf() && !_requires_grad) { throw std::runtime_error( - autograd::utils::requires_grad_leaf_error(requires_grad) + autograd::utils::requires_grad_leaf_error(_requires_grad) ); } return self.set_requires_grad(_requires_grad); From 52dd587123a8290bcc55b1fd00522cfc39719260 Mon Sep 17 00:00:00 2001 From: jon-tow Date: Mon, 28 Oct 2019 21:32:30 -0700 Subject: [PATCH 13/64] C++ API parity: Upsample (#28413) Summary: Adds `interpolate` functional and `Upsample` module support for the C++ API. **Issue**: https://github.com/pytorch/pytorch/issues/25883 **Reviewer**: yf225 Pull Request resolved: https://github.com/pytorch/pytorch/pull/28413 Differential Revision: D18165014 Pulled By: yf225 fbshipit-source-id: ecae2f432a301b1f4afa7c038b2d104cbad139f2 --- caffe2/CMakeLists.txt | 1 + test/cpp/api/enum.cpp | 10 ++ test/cpp/api/functional.cpp | 86 ++++++++++ test/cpp/api/modules.cpp | 156 ++++++++++++++++++ test/cpp_api_parity/parity-tracker.md | 2 +- tools/build_variables.py | 1 + torch/csrc/api/include/torch/enum.h | 10 ++ torch/csrc/api/include/torch/nn/functional.h | 1 + .../include/torch/nn/functional/upsampling.h | 107 ++++++++++++ torch/csrc/api/include/torch/nn/modules.h | 1 + .../api/include/torch/nn/modules/upsampling.h | 44 +++++ torch/csrc/api/include/torch/nn/options.h | 1 + .../api/include/torch/nn/options/upsampling.h | 73 ++++++++ torch/csrc/api/src/enum.cpp | 5 + torch/csrc/api/src/nn/modules/upsampling.cpp | 49 ++++++ 15 files changed, 546 insertions(+), 1 deletion(-) create mode 100644 torch/csrc/api/include/torch/nn/functional/upsampling.h create mode 100644 torch/csrc/api/include/torch/nn/modules/upsampling.h create mode 100644 torch/csrc/api/include/torch/nn/options/upsampling.h create mode 100644 torch/csrc/api/src/nn/modules/upsampling.cpp diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index d33a8c532a864..328e2d9a73b34 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -575,6 +575,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/pixelshuffle.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/pooling.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/rnn.cpp + ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/upsampling.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/container/functional.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/modules/container/named_any.cpp ${TORCH_SRC_DIR}/csrc/api/src/nn/options/activation.cpp diff --git a/test/cpp/api/enum.cpp b/test/cpp/api/enum.cpp index 3872875cfc9c2..e490ef9b22dfa 100644 --- a/test/cpp/api/enum.cpp +++ b/test/cpp/api/enum.cpp @@ -32,6 +32,11 @@ TEST(EnumTest, AllEnums) { torch::enumtype::kReflect, torch::enumtype::kReplicate, torch::enumtype::kCircular, + torch::enumtype::kNearest, + torch::enumtype::kBilinear, + torch::enumtype::kBicubic, + torch::enumtype::kTrilinear, + torch::enumtype::kArea, torch::enumtype::kSum, torch::enumtype::kMean, torch::enumtype::kMax @@ -54,6 +59,11 @@ TEST(EnumTest, AllEnums) { TORCH_ENUM_PRETTY_PRINT_TEST(Reflect) TORCH_ENUM_PRETTY_PRINT_TEST(Replicate) TORCH_ENUM_PRETTY_PRINT_TEST(Circular) + TORCH_ENUM_PRETTY_PRINT_TEST(Nearest) + TORCH_ENUM_PRETTY_PRINT_TEST(Bilinear) + TORCH_ENUM_PRETTY_PRINT_TEST(Bicubic) + TORCH_ENUM_PRETTY_PRINT_TEST(Trilinear) + TORCH_ENUM_PRETTY_PRINT_TEST(Area) TORCH_ENUM_PRETTY_PRINT_TEST(Sum) TORCH_ENUM_PRETTY_PRINT_TEST(Mean) TORCH_ENUM_PRETTY_PRINT_TEST(Max) diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp index 2273775ce1820..b9de3651ee54d 100644 --- a/test/cpp/api/functional.cpp +++ b/test/cpp/api/functional.cpp @@ -1266,6 +1266,92 @@ TEST_F(FunctionalTest, Threshold) { } } +TEST_F(FunctionalTest, Interpolate) { + { + // 1D interpolation + auto input = torch::ones({1, 1, 2}); + auto options = InterpolateOptions() + .size({4}) + .mode(torch::kNearest); + auto output = F::interpolate(input, options); + auto expected = torch::ones({1, 1, 4}); + + ASSERT_TRUE(output.allclose(expected)); + } + { + // 2D interpolation + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + auto input = torch::ones({1, 1, 2, 2}); + auto options = InterpolateOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBilinear) + .align_corners(align_corners); + auto output = F::interpolate(input, options); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + // 3D interpolation + for (const auto align_corners : {true, false}) { + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + auto input = torch::ones({1, 1, 2, 2, 2}); + auto options = + InterpolateOptions() + .scale_factor({scale_factor, scale_factor, scale_factor}) + .mode(torch::kTrilinear) + .align_corners(align_corners); + auto output = F::interpolate(input, options); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = + torch::ones({1, 1, expected_size, expected_size, expected_size}); + + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + auto input = torch::randn({3, 2, 2}); + ASSERT_THROWS_WITH( + F::interpolate(input[0], InterpolateOptions().size({4, 4})), + "Input Error: Only 3D, 4D and 5D input Tensors supported (got 2D) " + "for the modes: nearest | linear | bilinear | bicubic | trilinear (got kNearest)"); + ASSERT_THROWS_WITH( + F::interpolate( + torch::reshape(input, {1, 1, 1, 3, 2, 2}), + InterpolateOptions().size({1, 1, 1, 3, 4, 4})), + "Input Error: Only 3D, 4D and 5D input Tensors supported (got 6D) " + "for the modes: nearest | linear | bilinear | bicubic | trilinear (got kNearest)"); + ASSERT_THROWS_WITH( + F::interpolate(input, InterpolateOptions()), + "either size or scale_factor should be defined"); + ASSERT_THROWS_WITH( + F::interpolate( + input, + InterpolateOptions().size({3, 4, 4}).scale_factor({0.5})), + "only one of size or scale_factor should be defined"); + ASSERT_THROWS_WITH( + F::interpolate(input, InterpolateOptions().scale_factor({3, 2})), + "scale_factor shape must match input shape. " + "Input is 1D, scale_factor size is 2"); + ASSERT_THROWS_WITH( + F::interpolate( + input, + InterpolateOptions() + .mode(torch::kNearest) + .align_corners(true)), + "align_corners option can only be set with the " + "interpolating modes: linear | bilinear | bicubic | trilinear"); + } +} + TEST_F(FunctionalTest, Pad) { { auto input = torch::arange(6, torch::kDouble).reshape({1, 2, 3}); diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 8403977edfdcb..20f692b4e7d25 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -1688,6 +1688,153 @@ TEST_F(ModulesTest, Threshold) { } } +TEST_F(ModulesTest, Upsampling1D) { + { + Upsample model(UpsampleOptions() + .size({4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor}) + .mode(torch::kLinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + // linear (1D) upsampling spatial invariance + Upsample model(UpsampleOptions() + .scale_factor({3}) + .mode(torch::kLinear) + .align_corners(false)); + auto input = torch::zeros({1, 1, 9}); + input.narrow(2, 0, 4).normal_(); + auto output = model->forward(input); + auto expected = model->forward(input.narrow(2, 0, 5)); + + ASSERT_TRUE(torch::allclose(output.narrow(2, 0, 15), expected)); + } +} + +TEST_F(ModulesTest, Upsampling2D) { + { + Upsample model(UpsampleOptions() + .size({4, 4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBilinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model(UpsampleOptions() + .scale_factor({scale_factor, scale_factor}) + .mode(torch::kBicubic) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = torch::ones({1, 1, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } +} + +TEST_F(ModulesTest, Upsampling3D) { + { + Upsample model(UpsampleOptions() + .size({4, 4, 4}) + .mode(torch::kNearest)); + auto input = torch::ones({1, 1, 2, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected = torch::ones({1, 1, 4, 4, 4}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + { + for (const auto align_corners : {true, false}) { + // test float scale factor up & down sampling + for (const auto scale_factor : {0.5, 1.5, 2.0}) { + Upsample model( + UpsampleOptions() + .scale_factor({scale_factor, scale_factor, scale_factor}) + .mode(torch::kTrilinear) + .align_corners(align_corners)); + auto input = torch::ones({1, 1, 2, 2, 2}, torch::requires_grad()); + auto output = model->forward(input); + auto expected_size = + static_cast(std::floor(input.size(-1) * scale_factor)); + auto expected = + torch::ones({1, 1, expected_size, expected_size, expected_size}); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(s.ndimension(), 0); + ASSERT_TRUE(output.allclose(expected)); + } + } + } +} + TEST_F(ModulesTest, PrettyPrintIdentity) { ASSERT_EQ(c10::str(Identity()), "torch::nn::Identity()"); } @@ -1998,6 +2145,15 @@ TEST_F(ModulesTest, PrettyPrintConv) { "torch::nn::Conv2d(input_channels=3, output_channels=4, kernel_size=[5, 6], stride=[1, 2])"); } +TEST_F(ModulesTest, PrettyPrintUpsample) { + ASSERT_EQ( + c10::str(Upsample(UpsampleOptions().size({2, 4, 4}))), + "torch::nn::Upsample(size=[2, 4, 4], mode=kNearest)"); + ASSERT_EQ( + c10::str(Upsample(UpsampleOptions().scale_factor({0.5, 1.5}).mode(torch::kBilinear))), + "torch::nn::Upsample(scale_factor=[0.5, 1.5], mode=kBilinear)"); +} + TEST_F(ModulesTest, PrettyPrintUnfold) { ASSERT_EQ( c10::str(Unfold(torch::IntArrayRef({2, 4}))), diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index ba9eb78a75458..8a5d5f45082c1 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -121,7 +121,7 @@ torch.nn.CosineEmbeddingLoss|Yes|No torch.nn.MultiMarginLoss|Yes|No torch.nn.TripletMarginLoss|Yes|No torch.nn.PixelShuffle|Yes|No -torch.nn.Upsample|No|No +torch.nn.Upsample|Yes|No torch.nn.DataParallel|No|No torch.nn.parallel.DistributedDataParallel|No|No torch.nn.utils.clip_grad_norm_|Yes|No diff --git a/tools/build_variables.py b/tools/build_variables.py index 3e52cb4a9aecc..3f980938b3229 100644 --- a/tools/build_variables.py +++ b/tools/build_variables.py @@ -227,6 +227,7 @@ def add_torch_libs(): "torch/csrc/api/src/nn/modules/pixelshuffle.cpp", "torch/csrc/api/src/nn/modules/pooling.cpp", "torch/csrc/api/src/nn/modules/rnn.cpp", + "torch/csrc/api/src/nn/modules/upsampling.cpp", "torch/csrc/api/src/nn/modules/container/functional.cpp", "torch/csrc/api/src/nn/modules/container/named_any.cpp", "torch/csrc/api/src/nn/options/activation.cpp", diff --git a/torch/csrc/api/include/torch/enum.h b/torch/csrc/api/include/torch/enum.h index 111f5b5a6ba5f..6ba65befc1522 100644 --- a/torch/csrc/api/include/torch/enum.h +++ b/torch/csrc/api/include/torch/enum.h @@ -49,6 +49,11 @@ TORCH_ENUM_DECLARE(Constant) TORCH_ENUM_DECLARE(Reflect) TORCH_ENUM_DECLARE(Replicate) TORCH_ENUM_DECLARE(Circular) +TORCH_ENUM_DECLARE(Nearest) +TORCH_ENUM_DECLARE(Bilinear) +TORCH_ENUM_DECLARE(Bicubic) +TORCH_ENUM_DECLARE(Trilinear) +TORCH_ENUM_DECLARE(Area) TORCH_ENUM_DECLARE(Sum) TORCH_ENUM_DECLARE(Mean) TORCH_ENUM_DECLARE(Max) @@ -73,6 +78,11 @@ struct enum_name { TORCH_ENUM_PRETTY_PRINT(Reflect) TORCH_ENUM_PRETTY_PRINT(Replicate) TORCH_ENUM_PRETTY_PRINT(Circular) + TORCH_ENUM_PRETTY_PRINT(Nearest) + TORCH_ENUM_PRETTY_PRINT(Bilinear) + TORCH_ENUM_PRETTY_PRINT(Bicubic) + TORCH_ENUM_PRETTY_PRINT(Trilinear) + TORCH_ENUM_PRETTY_PRINT(Area) TORCH_ENUM_PRETTY_PRINT(Sum) TORCH_ENUM_PRETTY_PRINT(Mean) TORCH_ENUM_PRETTY_PRINT(Max) diff --git a/torch/csrc/api/include/torch/nn/functional.h b/torch/csrc/api/include/torch/nn/functional.h index 7187eff904aad..90ce395950145 100644 --- a/torch/csrc/api/include/torch/nn/functional.h +++ b/torch/csrc/api/include/torch/nn/functional.h @@ -9,4 +9,5 @@ #include #include #include +#include #include diff --git a/torch/csrc/api/include/torch/nn/functional/upsampling.h b/torch/csrc/api/include/torch/nn/functional/upsampling.h new file mode 100644 index 0000000000000..165ce9593041b --- /dev/null +++ b/torch/csrc/api/include/torch/nn/functional/upsampling.h @@ -0,0 +1,107 @@ +#pragma once + +#include +#include + +#include + +namespace torch { +namespace nn { +namespace functional { + +inline Tensor interpolate(const Tensor& input, InterpolateOptions options) { + auto _check_size_scale_factor = [options](size_t dim) { + if (options.size().empty() && options.scale_factor().empty()) { + TORCH_CHECK(false, "either size or scale_factor should be defined"); + } + if (!options.size().empty() && !options.scale_factor().empty()) { + TORCH_CHECK(false, "only one of size or scale_factor should be defined"); + } + if (!options.scale_factor().empty() && + options.scale_factor().size() != dim) { + TORCH_CHECK( + false, + "scale_factor shape must match input shape. " + "Input is ", dim, "D, scale_factor size is ", + options.scale_factor().size()); + } + }; + + auto _output_size = [input, options, _check_size_scale_factor](size_t dim) { + _check_size_scale_factor(dim); + if (!options.size().empty()) { + return options.size(); + } + auto scale_factors = options.scale_factor(); + + std::vector sizes; + for (size_t i = 0; i < dim; ++i) { + sizes.push_back(static_cast(std::floor( + static_cast(input.size(i + 2)) * scale_factors[i]))); + } + return sizes; + }; + + if (c10::get_if(&options.mode()) || + c10::get_if(&options.mode())) { + if (options.align_corners() != c10::nullopt) { + TORCH_CHECK( + false, + "align_corners option can only be set with the " + "interpolating modes: linear | bilinear | bicubic | trilinear"); + } + } else { + if (options.align_corners() == c10::nullopt) { + TORCH_WARN( + "Default upsampling behavior when mode is linear, bilinear, bicubic, " + "or trilinear, has changed to align_corners=False since 0.4.0. " + "Please specify align_corners=True if the old behavior is desired. " + "See the documentation of nn.Upsample for details."); + options.align_corners(false); + } + } + + if (input.dim() == 3 && c10::get_if(&options.mode())) { + return torch::upsample_nearest1d(input, _output_size(1)); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_nearest2d(input, _output_size(2)); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return torch::upsample_nearest3d(input, _output_size(3)); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + return adaptive_avg_pool1d(input, _output_size(1)); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return adaptive_avg_pool2d(input, _output_size(2)); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return adaptive_avg_pool3d(input, _output_size(3)); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + return torch::upsample_linear1d(input, _output_size(1), *options.align_corners()); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 3D input, but bilinear mode needs 4D input"); + } else if (input.dim() == 3 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 3D input, but trilinear mode needs 5D input"); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 4D input, but linear mode needs 3D input"); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_bilinear2d(input, _output_size(2), *options.align_corners()); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 4D input, but trilinear mode needs 5D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 5D input, but linear mode needs 3D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + TORCH_CHECK(false, "Got 5D input, but bilinear mode needs 4D input"); + } else if (input.dim() == 5 && c10::get_if(&options.mode())) { + return torch::upsample_trilinear3d(input, _output_size(3), *options.align_corners()); + } else if (input.dim() == 4 && c10::get_if(&options.mode())) { + return torch::upsample_bicubic2d(input, _output_size(2), *options.align_corners()); + } else { + TORCH_CHECK( + false, + "Input Error: Only 3D, 4D and 5D input Tensors supported " + "(got ", input.dim(), "D) for the modes: nearest | linear | bilinear | bicubic | trilinear " + "(got ", c10::visit(enumtype::enum_name{}, options.mode()), ")"); + } +} + +} // namespace functional +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/modules.h b/torch/csrc/api/include/torch/nn/modules.h index 375066214c0a4..2239e02161000 100644 --- a/torch/csrc/api/include/torch/nn/modules.h +++ b/torch/csrc/api/include/torch/nn/modules.h @@ -20,5 +20,6 @@ #include #include #include +#include #include #include diff --git a/torch/csrc/api/include/torch/nn/modules/upsampling.h b/torch/csrc/api/include/torch/nn/modules/upsampling.h new file mode 100644 index 0000000000000..f661ba346eecd --- /dev/null +++ b/torch/csrc/api/include/torch/nn/modules/upsampling.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace torch { +namespace nn { + +/// Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D +/// (volumetric) data. +/// +/// See https://pytorch.org/docs/stable/nn.html#Upsample to learn more +/// about the exact semantics of this module. +class TORCH_API UpsampleImpl : public Cloneable { + public: + explicit UpsampleImpl(const UpsampleOptions& options_ = {}); + + void reset() override; + + /// Pretty prints the `Upsample` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + Tensor forward(const Tensor& input); + + /// The options with which this `Module` was constructed. + UpsampleOptions options; +}; + +/// A `ModuleHolder` subclass for `UpsampleImpl`. +/// See the documentation for `UpsampleImpl` class to learn what +/// methods it provides, or the documentation for `ModuleHolder` to learn about +/// PyTorch's module storage semantics. +TORCH_MODULE(Upsample); + +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/options.h b/torch/csrc/api/include/torch/nn/options.h index 0cc7d338640e8..893f92141ffdb 100644 --- a/torch/csrc/api/include/torch/nn/options.h +++ b/torch/csrc/api/include/torch/nn/options.h @@ -12,3 +12,4 @@ #include #include #include +#include diff --git a/torch/csrc/api/include/torch/nn/options/upsampling.h b/torch/csrc/api/include/torch/nn/options/upsampling.h new file mode 100644 index 0000000000000..67f62fad47470 --- /dev/null +++ b/torch/csrc/api/include/torch/nn/options/upsampling.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +namespace torch { +namespace nn { + +/// Options for a `D`-dimensional interpolate functional. +struct TORCH_API InterpolateOptions { + typedef c10::variant< + enumtype::kNearest, + enumtype::kLinear, + enumtype::kBilinear, + enumtype::kBicubic, + enumtype::kTrilinear, + enumtype::kArea> mode_t; + + /// output spatial sizes. + TORCH_ARG(std::vector, size) = {}; + + /// multiplier for spatial size. + TORCH_ARG(std::vector, scale_factor) = {}; + + /// the upsampling algorithm: one of "nearest", "linear", "bilinear", + /// "bicubic", "trilinear", and "area". Default: "nearest" + TORCH_ARG(mode_t, mode) = torch::kNearest; + + /// Geometrically, we consider the pixels of the input and output as squares + /// rather than points. If set to "True", the input and output tensors are + /// aligned by the center points of their corner pixels, preserving the values + /// at the corner pixels. If set to "False", the input and output tensors + /// are aligned by the corner points of their corner pixels, and the + /// interpolation uses edge value padding for out-of-boundary values, making + /// this operation *independent* of input size when :attr:`scale_factor` is + /// kept the same. This only has an effect when :attr:`mode` is "linear", + /// "bilinear", "bicubic" or "trilinear". Default: "False" + TORCH_ARG(c10::optional, align_corners) = c10::nullopt; +}; + +/// Options for a `D`-dimensional Upsample module. +struct TORCH_API UpsampleOptions { + /// output spatial sizes. + TORCH_ARG(std::vector, size) = {}; + + /// multiplier for spatial size. + TORCH_ARG(std::vector, scale_factor) = {}; + + /// the upsampling algorithm: one of "nearest", "linear", "bilinear", + /// "bicubic" and "trilinear". Default: "nearest" + typedef c10::variant< + enumtype::kNearest, + enumtype::kLinear, + enumtype::kBilinear, + enumtype::kBicubic, + enumtype::kTrilinear> mode_t; + TORCH_ARG(mode_t, mode) = torch::kNearest; + + /// if "True", the corner pixels of the input and output tensors are + /// aligned, and thus preserving the values at those pixels. This only has + /// effect when :attr:`mode` is "linear", "bilinear", or + /// "trilinear". Default: "False" + TORCH_ARG(c10::optional, align_corners) = c10::nullopt; +}; + +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/src/enum.cpp b/torch/csrc/api/src/enum.cpp index b59eb955e9635..374d2839af8c7 100644 --- a/torch/csrc/api/src/enum.cpp +++ b/torch/csrc/api/src/enum.cpp @@ -17,6 +17,11 @@ TORCH_ENUM_DEFINE(Constant) TORCH_ENUM_DEFINE(Reflect) TORCH_ENUM_DEFINE(Replicate) TORCH_ENUM_DEFINE(Circular) +TORCH_ENUM_DEFINE(Nearest) +TORCH_ENUM_DEFINE(Bilinear) +TORCH_ENUM_DEFINE(Bicubic) +TORCH_ENUM_DEFINE(Trilinear) +TORCH_ENUM_DEFINE(Area) TORCH_ENUM_DEFINE(Sum) TORCH_ENUM_DEFINE(Mean) TORCH_ENUM_DEFINE(Max) diff --git a/torch/csrc/api/src/nn/modules/upsampling.cpp b/torch/csrc/api/src/nn/modules/upsampling.cpp new file mode 100644 index 0000000000000..8699f73ce600a --- /dev/null +++ b/torch/csrc/api/src/nn/modules/upsampling.cpp @@ -0,0 +1,49 @@ +#include + +#include + +namespace F = torch::nn::functional; + +namespace torch { +namespace nn { + +UpsampleImpl::UpsampleImpl(const UpsampleOptions& options_) // NOLINT(modernize-pass-by-value) + : options(options_) {} + +void UpsampleImpl::reset() {} + +void UpsampleImpl::pretty_print(std::ostream& stream) const { + stream << "torch::nn::Upsample("; + if (!options.scale_factor().empty()) { + stream << "scale_factor=" << at::ArrayRef(options.scale_factor()); + } else { + stream << "size=" << at::ArrayRef(options.size()); + } + stream << ", mode=" << c10::visit(enumtype::enum_name{}, options.mode()) << ")"; +} + +Tensor UpsampleImpl::forward(const Tensor& input) { + InterpolateOptions::mode_t mode; + if (c10::get_if(&options.mode())) { + mode = torch::kNearest; + } else if (c10::get_if(&options.mode())) { + mode = torch::kLinear; + } else if (c10::get_if(&options.mode())) { + mode = torch::kBilinear; + } else if (c10::get_if(&options.mode())) { + mode = torch::kBicubic; + } else if (c10::get_if(&options.mode())) { + mode = torch::kTrilinear; + } + + return F::interpolate( + input, + InterpolateOptions() + .size(options.size()) + .scale_factor(options.scale_factor()) + .mode(mode) + .align_corners(options.align_corners())); +} + +} // namespace nn +} // namespace torch From 097da5524955c42433926ba604cbd32df915c6b0 Mon Sep 17 00:00:00 2001 From: Lu Fang Date: Mon, 28 Oct 2019 21:45:40 -0700 Subject: [PATCH 14/64] Fix BC check CI (#28816) Summary: Skip the functions which were reverted. Pull Request resolved: https://github.com/pytorch/pytorch/pull/28816 Reviewed By: hl475 Differential Revision: D18196628 Pulled By: houseroad fbshipit-source-id: 30d43fcd57efb21b870c6a630b7ee305604dc603 --- test/backward_compatibility/check_backward_compatibility.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index bec599e906a84..6a22579d51d4c 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -23,6 +23,9 @@ ('thnn_conv_depthwise2d_backward', datetime.date(2019, 10, 30)), ('thnn_conv3d_backward', datetime.date(2019, 10, 30)), ('empty_like', datetime.date(2019, 10, 30)), + ('rand_like', datetime.date(2019, 11, 11)), + ('ones_like', datetime.date(2019, 11, 11)), + ('full_like', datetime.date(2019, 11, 11)), ] From a0339c8d8f31e1e80ebecd7f3c002a836ecbf371 Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Mon, 28 Oct 2019 22:18:49 -0700 Subject: [PATCH 15/64] `bootstrap.sh` refactor (#28809) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28809 ### Summary This PR adds the interactive mode to `bootstrap.sh`. Instead of passing the credential information from command parameters(`-t`,`-p`), we're going to ask the user enter that information and save it to a config file, such that next time you don't have to enter again. So all you need now, is one line command ```shell ./bootstrap ``` ### Test Plan - TestApp.ipa can be installed on any devices - Don't break CI jobs Test Plan: Imported from OSS Differential Revision: D18194032 Pulled By: xta0 fbshipit-source-id: a416ef7f13fa565e2c10bb55f94a8ce994b4e869 --- ios/TestApp/.gitignore | 1 + ios/TestApp/README.md | 13 +++++---- .../TestApp/Base.lproj/Main.storyboard | 8 +++++- ios/TestApp/TestApp/Benchmark.mm | 1 - ios/TestApp/TestApp/ViewController.mm | 28 +++++++++++++------ ios/TestApp/bootstrap.sh | 27 +++++++++++++----- 6 files changed, 54 insertions(+), 24 deletions(-) diff --git a/ios/TestApp/.gitignore b/ios/TestApp/.gitignore index 1038da3961c8f..eaf0cd27b306e 100644 --- a/ios/TestApp/.gitignore +++ b/ios/TestApp/.gitignore @@ -1 +1,2 @@ model.pt +.config diff --git a/ios/TestApp/README.md b/ios/TestApp/README.md index 28aa88b7a9f77..ce2842d28016a 100644 --- a/ios/TestApp/README.md +++ b/ios/TestApp/README.md @@ -21,9 +21,10 @@ The TestApp is currently being used as a dummy app by Circle CI for nightly jobs The benchmark folder contains two scripts that help you setup the benchmark project. The `setup.rb` does the heavy-lifting jobs of setting up the XCode project, whereas the `trace_model.py` is a Python script that you can tweak to generate your model for benchmarking. Simply follow the steps below to setup the project 1. In the PyTorch root directory, run `BUILD_PYTORCH_MOBILE=1 IOS_ARCH=arm64 ./scripts/build_ios.sh` to generate the custom build from **Master** branch -2. Navigate to the `benchmark` folder, run `python trace_model.py` to get your model generated. -3. In the same directory, run `ruby setup.rb` to setup the XCode project. -4. Open the `TestApp.xcodeproj`, you're ready to go. +2. Navigate to the `benchmark` folder, run `python trace_model.py` to generate your model. +3. In the same directory, open `config.json`. Those are the input parameters you can tweak. +4. Again, in the same directory, run `ruby setup.rb` to setup the XCode project. +5. Open the `TestApp.xcodeproj`, you're ready to go. The benchmark code is written in C++, see `benchmark.mm` for more details. @@ -35,13 +36,13 @@ For those who want to do perf testing but don't want touch XCode, `bootstrap.sh` 2. A valid provisioning profile for code signing 3. A valid team identifier -To run the script, simply type the command below and make sure your phone is unlocked and connected via USB. +To run the script, simply type the command below and make sure your phone is connected via USB. ```shell -./bootstrap -t ${TEAM_ID} -p ${PROVISIONING_PROFILE} +./bootstrap ``` -The benchmark log will be displayed on the screen. +Open the app on your device, the benchmark result will be displayed on the screen. > Note This requires ios-deploy to be installed. Please have a look at [ios-deploy](https://github.com/ios-control/ios-deploy). To quickly install it, use `npm -g i ios-deploy` diff --git a/ios/TestApp/TestApp/Base.lproj/Main.storyboard b/ios/TestApp/TestApp/Base.lproj/Main.storyboard index c78975d739870..31dd0f40e5673 100644 --- a/ios/TestApp/TestApp/Base.lproj/Main.storyboard +++ b/ios/TestApp/TestApp/Base.lproj/Main.storyboard @@ -33,7 +33,13 @@ - + + + + + + + diff --git a/ios/TestApp/TestApp/Benchmark.mm b/ios/TestApp/TestApp/Benchmark.mm index cadcc421a2891..b8cabc0326596 100644 --- a/ios/TestApp/TestApp/Benchmark.mm +++ b/ios/TestApp/TestApp/Benchmark.mm @@ -77,7 +77,6 @@ + (NSString*)run { if (print_output) { std::cout << module.forward(inputs) << std::endl; } - UI_LOG(@"Start benchmarking...", nil); UI_LOG(@"Running warmup runs", nil); CAFFE_ENFORCE(warmup >= 0, "Number of warm up runs should be non negative, provided ", warmup, "."); diff --git a/ios/TestApp/TestApp/ViewController.mm b/ios/TestApp/TestApp/ViewController.mm index 4e4f0dcb884f4..97f5e69b439fc 100644 --- a/ios/TestApp/TestApp/ViewController.mm +++ b/ios/TestApp/TestApp/ViewController.mm @@ -7,7 +7,8 @@ @interface ViewController () @end -@implementation ViewController +@implementation ViewController { +} - (void)viewDidLoad { [super viewDidLoad]; @@ -18,20 +19,29 @@ - (void)viewDidLoad { NSDictionary* config = [NSJSONSerialization JSONObjectWithData:configData options:NSJSONReadingAllowFragments error:&err]; + if (err) { NSLog(@"Parse config.json failed!"); return; } + [Benchmark setup:config]; + [self runBenchmark]; +} +- (void)runBenchmark { + self.textView.text = @"Start benchmarking...\n"; dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - if ([Benchmark setup:config]) { - NSString* text = [Benchmark run]; - dispatch_async(dispatch_get_main_queue(), ^{ - self.textView.text = text; - }); - } else { - NSLog(@"Setup benchmark config failed!"); - } + NSString* text = [Benchmark run]; + dispatch_async(dispatch_get_main_queue(), ^{ + self.textView.text = [self.textView.text stringByAppendingString:text]; + }); + }); +} + +- (IBAction)reRun:(id)sender { + self.textView.text = @""; + dispatch_async(dispatch_get_main_queue(), ^{ + [self runBenchmark]; }); } diff --git a/ios/TestApp/bootstrap.sh b/ios/TestApp/bootstrap.sh index 3fdb9721e867c..9785bc18fe923 100755 --- a/ios/TestApp/bootstrap.sh +++ b/ios/TestApp/bootstrap.sh @@ -22,14 +22,30 @@ bootstrap() { XCODE_PROJ_PATH="./TestApp.xcodeproj" XCODE_TARGET="TestApp" XCODE_BUILD="./build" - if [ -d ${XCODE_BUILD} ]; then + if [ ! -f "./.config" ]; then + touch .config + echo "" >> .config + else + source .config + fi + if [ -z "${TEAM_ID}" ]; then + reply=$(bash -c 'read -r -p "Team Id:" tmp; echo $tmp') + TEAM_ID="${reply}" + echo "TEAM_ID=${TEAM_ID}" >> .config + fi + if [ -z "${PROFILE}" ]; then + reply=$(bash -c 'read -r -p "Provisioning Profile:" tmp; echo $tmp') + PROFILE="${reply}" + echo "PROFILE=${PROFILE}" >> .config + fi + if [ -d "${XCODE_BUILD}" ]; then echo "found the old XCode build, remove it" - rm -rf ${XCODE_BUILD} + rm -rf "${XCODE_BUILD}" fi - cd ${BENCHMARK_DIR} + cd "${BENCHMARK_DIR}" echo "Generating model" python trace_model.py - ruby setup.rb -t ${TEAM_ID} + ruby setup.rb -t "${TEAM_ID}" cd .. #run xcodebuild if ! [ -x "$(command -v xcodebuild)" ]; then @@ -77,7 +93,4 @@ esac shift done -echo TEAM_ID = "${TEAM_ID}" -echo PROFILE = "${PROFILE}" - bootstrap From 295401f04c88892b953ead2662f40a3e879f8ee7 Mon Sep 17 00:00:00 2001 From: svcscm Date: Mon, 28 Oct 2019 22:22:18 -0700 Subject: [PATCH 16/64] Updating submodules Summary: GitHub commits: https://github.com/pytorch/fbgemm/commit/edee4921c47162b1d8a1e96cc4ab0e49ae8816d2 Test Plan: n/a Reviewed By: zpao fbshipit-source-id: b69770ac1a801b372fba0e112124b25ad1572821 --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index 5487e2b1a2184..edee4921c4716 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 5487e2b1a2184ee41f49a2935922dc8ed92f6464 +Subproject commit edee4921c47162b1d8a1e96cc4ab0e49ae8816d2 From f6692146e771bd37d7e3483e2af58c236ffc96ca Mon Sep 17 00:00:00 2001 From: Xiaomeng Yang Date: Mon, 28 Oct 2019 23:24:43 -0700 Subject: [PATCH 17/64] Add Conv3dInt8 (#28768) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28768 Add Conv3dInt8 Test Plan: buck test mode/dev-nosan caffe2/test:quantized -- "Conv" Reviewed By: jianyuh Differential Revision: D18023661 fbshipit-source-id: 8fc7a4350baf29271dfd6fa3c1c4b10e60e2fdbf --- aten/src/ATen/native/quantized/cpu/qconv.cpp | 713 ++++++++++++------- test/test_quantized.py | 309 +++++--- 2 files changed, 663 insertions(+), 359 deletions(-) diff --git a/aten/src/ATen/native/quantized/cpu/qconv.cpp b/aten/src/ATen/native/quantized/cpu/qconv.cpp index c2993717f7552..cc6efcd084565 100644 --- a/aten/src/ATen/native/quantized/cpu/qconv.cpp +++ b/aten/src/ATen/native/quantized/cpu/qconv.cpp @@ -1,41 +1,140 @@ +#include +#include +#include + #include -#include #include +#include #include #include #include #include #include -#include namespace at { namespace native { namespace { -SmallVector convOutputShape( +template +bool ConvDimChecks( + int64_t act_dims, + int64_t stride_dims, + int64_t padding_dims, + int64_t dilation_dims) { + TORCH_CHECK( + act_dims == kSpatialDim + 2, + "quantized::conv", + kSpatialDim, + "d(): Expected activation tensor to have ", + kSpatialDim + 2, + " dimensions."); + TORCH_CHECK( + stride_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected stride tensor to have ", + kSpatialDim, + " dimensions."); + TORCH_CHECK( + padding_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected padding tensor to have ", + kSpatialDim, + " dimensions."); + TORCH_CHECK( + dilation_dims == kSpatialDim, + "quantized::conv", + kSpatialDim, + "d(): Expected dilation tensor to have ", + kSpatialDim, + " dimensions."); + return true; +} + +#ifdef USE_FBGEMM + +template +SmallVector MakeConvOutputShape( + int N, + int M, + const std::array& output_image_shape); + +template <> +SmallVector MakeConvOutputShape<2>( + int N, + int M, + const std::array& output_image_shape) { + return {N, M, output_image_shape[0], output_image_shape[1]}; +} + +template <> +SmallVector MakeConvOutputShape<3>( + int N, + int M, + const std::array& output_image_shape) { + return {N, + M, + output_image_shape[0], + output_image_shape[1], + output_image_shape[2]}; +} + +#endif // USE_FBGEMM + +#ifdef USE_PYTORCH_QNNPACK + +template +SmallVector MakeConvOutputShape( int N, // mini-batch - int K, // output channels - int H, // input height - int W, // input width + int M, // output channels + const std::vector& input_image_shape, + const std::vector& kernel, + const torch::List& stride, + const torch::List& padding, + const torch::List& dilation); + +template <> +SmallVector MakeConvOutputShape<2>( + int N, // mini-batch + int M, // output channels + const std::vector& input_image_shape, const std::vector& kernel, const torch::List& stride, const torch::List& padding, const torch::List& dilation) { - SmallVector out_shape; - out_shape.push_back(N); - - int H_out = std::floor( - (H + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1); - int W_out = std::floor( - (W + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1); - out_shape.push_back(H_out); - out_shape.push_back(W_out); - // TODO: reorder it to NCHW order once the memory format regression is fixed - out_shape.push_back(K); - - return out_shape; + const int H = input_image_shape[0]; + const int W = input_image_shape[1]; + const int64_t Y_H = + (H + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1; + const int64_t Y_W = + (W + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1; + return {N, M, Y_H, Y_W}; } +template <> +SmallVector MakeConvOutputShape<3>( + int N, // mini-batch + int M, // output channels + const std::vector& input_image_shape, + const std::vector& kernel, + const torch::List& stride, + const torch::List& padding, + const torch::List& dilation) { + const int D = input_image_shape[0]; + const int H = input_image_shape[1]; + const int W = input_image_shape[2]; + const int64_t Y_D = + (D + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0] + 1; + const int64_t Y_H = + (H + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1] + 1; + const int64_t Y_W = + (W + 2 * padding[2] - dilation[2] * (kernel[2] - 1) - 1) / stride[2] + 1; + return {N, M, Y_D, Y_H, Y_W}; +} + +#endif // USE_PYTORCH_QNNPACK + /* * FBGEMM uses vpmaddubsw instruction to multiply activations (uint8_t) and * weights (int8_t). @@ -65,27 +164,99 @@ SmallVector convOutputShape( * is 32767. * */ -template -class QConv2dInt8 final : public c10::OperatorKernel { +template +class QConvInt8 final : public c10::OperatorKernel { public: - void conv_checks( - int64_t act_dims, - int64_t stride_dims, - int64_t padding_dims, - int64_t dilation_dims) { - TORCH_CHECK( - act_dims == 4, - "quantized::conv2d(): Expected activation tensor to have 4 dimensions."); - TORCH_CHECK( - stride_dims == 2, "quantized::conv2d(): Supports 2D convolution only"); - TORCH_CHECK( - padding_dims == 2, "quantized::conv2d(): Supports 2D convolution only"); + Tensor operator()( + Tensor act, + Tensor packed_weight, + torch::List stride, + torch::List padding, + torch::List dilation, + int64_t groups, + double output_scale, + int64_t output_zero_point) { + auto& ctx = at::globalContext(); + +#ifdef USE_FBGEMM + if (ctx.qEngine() == at::QEngine::FBGEMM) { + return FbgemmConv( + act, + packed_weight, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point); + } +#endif // USE_FBGEMM + +#ifdef USE_PYTORCH_QNNPACK + if (ctx.qEngine() == at::QEngine::QNNPACK) { + TORCH_CHECK(kSpatialDim == 2, "QNNPACK only suuports Conv2d now."); + return QnnpackConv( + act, + packed_weight, + stride, + padding, + dilation, + groups, + output_scale, + output_zero_point); + } +#endif + TORCH_CHECK( - dilation_dims == 2, - "quantized::conv2d(): Supports 2D convolution only"); + false, + "Didn't find engine for operation quantized::conv ", + toString(ctx.qEngine())); } + + private: #ifdef USE_FBGEMM - at::Tensor fbgemm_conv( + static const float* GetBiasData( + const PackedConvWeight& pack_data, + Tensor* bias) { + const float* bias_data = nullptr; + if (pack_data.bias.has_value()) { + *bias = pack_data.bias.value(); + TORCH_CHECK( + bias->dtype() == at::kFloat, + "[QConv3D] The 'bias' tensor must have 'torch.float' dtype"); + *bias = bias->contiguous(); + TORCH_CHECK(bias->dim() == 1, "bias should be a vector (1D Tensor)"); + const int M = pack_data.w->outputChannels(); + TORCH_CHECK(bias->size(0) == M, "bias should have ", M, " elements."); + bias_data = bias->data_ptr(); + } + return bias_data; + } + + static void GetQuantizationParams( + const PackedConvWeight& pack_data, + float act_scale, + float out_scale, + std::vector* output_multiplier_float, + std::vector* act_times_w_scale) { + if (pack_data.q_scheme == kPerTensorAffine) { + *act_times_w_scale = {(act_scale * pack_data.w_scale[0])}; + *output_multiplier_float = {act_times_w_scale->front() / out_scale}; + } else if (pack_data.q_scheme == kPerChannelAffine) { + const int M = pack_data.w->outputChannels(); + output_multiplier_float->resize(M); + act_times_w_scale->resize(M); + for (int i = 0; i < M; ++i) { + act_times_w_scale->at(i) = (act_scale * pack_data.w_scale[i]); + output_multiplier_float->at(i) = act_times_w_scale->at(i) / out_scale; + } + } else { + TORCH_CHECK( + false, "[QConv", kSpatialDim, "D] Unknown quantization scheme"); + } + } + + at::Tensor FbgemmConv( Tensor act, Tensor packed_weight, torch::List stride, @@ -105,178 +276,230 @@ class QConv2dInt8 final : public c10::OperatorKernel { // See https://github.com/pytorch/pytorch/issues/23403 TORCH_CHECK( fbgemm::fbgemmSupportedCPU(), "Your CPU does not support FBGEMM."); - conv_checks( + ConvDimChecks( act.ndimension(), stride.size(), padding.size(), dilation.size()); - int N = act.size(0); - int C = act.size(1); - int H = act.size(2); - int W = act.size(3); - - // FBGEMM requires NHWC - // TODO: change it to contiguous(MemoryFormat::ChannelsLast) once a perf - // regression of it is fixed. Today it's equivalent because `act` sizes - // are not used below - Tensor act_contig = act.permute({0, 2, 3, 1}).contiguous(); - const uint8_t* act_ptr = - reinterpret_cast(act_contig.data_ptr()); - - PackedConvWeight<2>& pack_ptr = - cpp_custom_type_hack::cast>(packed_weight); - auto packB = pack_ptr.w.get(); - auto& col_offsets = pack_ptr.col_offsets; - auto& kernel = pack_ptr.kernel; - - int K = packB->outputChannels(); - - int pad_l = padding[0]; - int pad_t = padding[1]; - int stride_h = stride[0]; - int stride_w = stride[1]; - int kernel_h = kernel[0]; - int kernel_w = kernel[1]; - // clang-format off - TORCH_CHECK(C == (packB->inputChannels()), - "[QConv2D] Given groups=", groups, ", weight of size ", - K, ", ", kernel_h, ", ", kernel_w, ", ", packB->inputChannels(), - ", expected input (NCHW) ", N, ", ", C, ", ", H, ", ", W, - " to have ", (packB->inputChannels() * groups), - " channels, but got ", C, " channels instead"); - // clang-format on - fbgemm::conv_param_t<> conv_p( - N, // Batch size - C, // Number of input channels - K, // Number of output channels - {H, W}, - groups, - {kernel_h, kernel_w}, - {stride_h, stride_w}, - {pad_l, pad_t, pad_l, pad_t}, - {static_cast(dilation[0]), static_cast(dilation[1])}); - - float act_scale = act.q_scale(); - int32_t act_zero_point = act.q_zero_point(); - - const float* bias_ptr = nullptr; - at::Tensor bias; - if (pack_ptr.bias.has_value()) { - bias = pack_ptr.bias.value(); + const int N = act.size(0); + const int C = act.size(1); + const int D = kSpatialDim == 2 ? 1 : act.size(2); + const int H = act.size(kSpatialDim); + const int W = act.size(kSpatialDim + 1); + + const Tensor act_nhwc = kSpatialDim == 2 + ? act.contiguous(MemoryFormat::ChannelsLast) + : fbgemm_utils::ConvertToChannelsLast3dTensor(act); + const uint8_t* act_data = + reinterpret_cast(act_nhwc.data_ptr()); + PackedConvWeight& pack_data = + cpp_custom_type_hack::cast>( + packed_weight); + auto* pack_w = pack_data.w.get(); + const auto& col_offsets = pack_data.col_offsets; + const auto& kernel = pack_data.kernel; + + const int M = pack_w->outputChannels(); + const int kernel_d = kSpatialDim == 2 ? 1 : kernel[0]; + const int kernel_h = kernel[kSpatialDim - 2]; + const int kernel_w = kernel[kSpatialDim - 1]; + const int pad_d = kSpatialDim == 2 ? 0 : padding[0]; + const int pad_h = padding[kSpatialDim - 2]; + const int pad_w = padding[kSpatialDim - 1]; + const int stride_d = kSpatialDim == 2 ? 1 : stride[0]; + const int stride_h = stride[kSpatialDim - 2]; + const int stride_w = stride[kSpatialDim - 1]; + const int dilation_d = kSpatialDim == 2 ? 1 : dilation[0]; + const int dilation_h = dilation[kSpatialDim - 2]; + const int dilation_w = dilation[kSpatialDim - 1]; + + if (kSpatialDim == 2) { TORCH_CHECK( - bias.dtype() == at::kFloat, - "[QConv2D] The 'bias' tensor must have 'torch.float' dtype"); - bias = bias.contiguous(); - TORCH_CHECK(bias.dim() == 1, "bias should be a vector (1D Tensor)"); + C == pack_w->inputChannels(), + "[QConv2D] Given groups=", + groups, + ", weight of size ", + M, + ", ", + kernel_h, + ", ", + kernel_w, + ", ", + pack_w->inputChannels(), + ", expected input (NCHW) ", + N, + ", ", + C, + ", ", + H, + ", ", + W, + " to have ", + pack_w->inputChannels(), + " channels, but got ", + C, + " channels instead"); + } else { TORCH_CHECK( - bias.size(0) == K, - "bias should have K elements: " + c10::to_string(K)); - bias_ptr = bias.data_ptr(); + C == pack_w->inputChannels(), + "[QConv3D] Given groups=", + groups, + ", weight of size ", + M, + ", ", + kernel_d, + ", ", + kernel_h, + ", ", + kernel_w, + ", ", + pack_w->inputChannels(), + ", expected input (NCDHW) ", + N, + ", ", + C, + ", ", + D, + ", ", + H, + ", ", + W, + " to have ", + pack_w->inputChannels(), + " channels, but got ", + C, + " channels instead"); } - std::vector output_multiplier_float(1, 0.0); - std::vector act_times_w_scale(1, 1.0); + fbgemm::conv_param_t conv_p = + fbgemm_utils::MakeFbgemmConvParam( + N, // Batch size + C, // Number of input channels + M, // Number of output channels + kSpatialDim == 2 ? std::vector{H, W} + : std::vector{D, H, W}, + groups, + kSpatialDim == 2 ? std::vector{kernel_h, kernel_w} + : std::vector{kernel_d, kernel_h, kernel_w}, + kSpatialDim == 2 ? std::vector{stride_h, stride_w} + : std::vector{stride_d, stride_h, stride_w}, + kSpatialDim == 2 ? std::vector{pad_h, pad_w} + : std::vector{pad_d, pad_h, pad_w}, + kSpatialDim == 2 + ? std::vector{dilation_h, dilation_w} + : std::vector{dilation_d, dilation_h, dilation_w}); + + const float act_scale = act.q_scale(); + const int32_t act_zero_point = act.q_zero_point(); + + Tensor bias; + const float* bias_data = GetBiasData(pack_data, &bias); + TORCH_CHECK( - pack_ptr.w_scale.size() == pack_ptr.w_zp.size(), + pack_data.w_scale.size() == pack_data.w_zp.size(), "Weight scales and zero points vectors should have the same size."); + std::vector output_multiplier_float; + std::vector act_times_w_scale; + GetQuantizationParams( + pack_data, + act_scale, + output_scale, + &output_multiplier_float, + &act_times_w_scale); - if (pack_ptr.q_scheme == kPerTensorAffine) { - act_times_w_scale[0] = (act_scale * pack_ptr.w_scale[0]); - output_multiplier_float[0] = - act_times_w_scale[0] / static_cast(output_scale); - } else if (pack_ptr.q_scheme == kPerChannelAffine) { - output_multiplier_float.resize(K, 0.0); - act_times_w_scale.resize(K, 1.0); - for (int i = 0; i < K; ++i) { - act_times_w_scale[i] = (act_scale * pack_ptr.w_scale[i]); - output_multiplier_float[i] = - act_times_w_scale[i] / static_cast(output_scale); - } - } else { - TORCH_CHECK(false, "[QConv2D] Unknown quantization scheme"); - } - - // TODO: change the following to NCHW sizes once perf is fixed - SmallVector outShape{ - N, conv_p.OUT_DIM[0], conv_p.OUT_DIM[1], K}; + const SmallVector output_shape = + MakeConvOutputShape(N, M, conv_p.OUT_DIM); TORCH_CHECK( std::all_of( - outShape.begin(), outShape.end(), [](int64_t i) { return i > 0; }), - "[QConv2D] each dimension of output tensor should be greater than 0") - - // Force output format to be NHWC - // TODO: consider preserving input format - // TODO: add MemoryFormat::ChannelsLast here once perf is fixed - Tensor output = _empty_affine_quantized( - outShape, device(kCPU).dtype(kQUInt8), output_scale, output_zero_point); - auto buffer = at::empty(output.sizes(), output.options().dtype(at::kInt)); - - int num_tasks = at::get_num_threads(); + output_shape.begin(), + output_shape.end(), + [](int64_t i) { return i > 0; }), + "[QConv", + kSpatialDim, + "D] each dimension of output tensor should be greater than 0"); + + Tensor output = kSpatialDim == 2 + ? _empty_affine_quantized( + output_shape, + device(kCPU).dtype(kQUInt8), + output_scale, + output_zero_point, + MemoryFormat::ChannelsLast) + : fbgemm_utils::MakeEmptyAffineQuantizedChannelsLast3dTensor( + output_shape[0], + output_shape[1], + output_shape[2], + output_shape[3], + output_shape[4], + device(kCPU).dtype(kQUInt8), + output_scale, + output_zero_point); + Tensor buffer = at::empty(output.sizes(), output.options().dtype(at::kInt)); + const int num_tasks = at::get_num_threads(); at::parallel_for(0, num_tasks, 1, [&](int64_t begin, int64_t end) { - fbgemm::DoNothing<> NoOpObj{}; + fbgemm::DoNothing<> kNoOpObj{}; for (int task_id = begin; task_id < end; ++task_id) { - if (pack_ptr.q_scheme == kPerTensorAffine) { + if (pack_data.q_scheme == kPerTensorAffine) { fbgemm::ReQuantizeOutput< - ReluFused, + kReluFused, fbgemm::QuantizationGranularity::TENSOR, float> - outputProcObj( - NoOpObj, + output_proc_obj( + kNoOpObj, output_multiplier_float.data(), output_zero_point, act_zero_point, - pack_ptr.w_zp.data(), + pack_data.w_zp.data(), nullptr, /* row offset buffer */ col_offsets.data(), - bias_ptr, - K, + bias_data, + M, groups, act_times_w_scale.data()); - fbgemm::fbgemmConv( + fbgemm::fbgemmConv( conv_p, - act_ptr, - *packB, + act_data, + *pack_w, reinterpret_cast(output.data_ptr()), buffer.data_ptr(), - outputProcObj, + output_proc_obj, task_id /* thread_id*/, num_tasks /* num_threads */); - - } else if (pack_ptr.q_scheme == kPerChannelAffine) { + } else if (pack_data.q_scheme == kPerChannelAffine) { fbgemm::ReQuantizeOutput< - ReluFused, + kReluFused, fbgemm::QuantizationGranularity::OUT_CHANNEL, float> - outputProcObj( - NoOpObj, + output_proc_obj( + kNoOpObj, output_multiplier_float.data(), output_zero_point, act_zero_point, - pack_ptr.w_zp.data(), + pack_data.w_zp.data(), nullptr, /* row offset buffer */ col_offsets.data(), - bias_ptr, - K, + bias_data, + M, groups, act_times_w_scale.data()); - fbgemm::fbgemmConv( + fbgemm::fbgemmConv( conv_p, - act_ptr, - *packB, + act_data, + *pack_w, reinterpret_cast(output.data_ptr()), buffer.data_ptr(), - outputProcObj, + output_proc_obj, task_id /* thread_id*/, num_tasks /* num_threads */); } } }); - // TODO: remove permute once MemoryLayout is added above - return output.permute({0, 3, 1, 2}); + return output; } #endif + #ifdef USE_PYTORCH_QNNPACK - at::Tensor qnnpack_conv( + at::Tensor QnnpackConv( Tensor act, Tensor packed_weight, torch::List stride, @@ -285,44 +508,42 @@ class QConv2dInt8 final : public c10::OperatorKernel { int64_t groups, double output_scale, int64_t output_zero_point) { - conv_checks( + ConvDimChecks( act.ndimension(), stride.size(), padding.size(), dilation.size()); - PackedConvWeightsQnnp& pack_ptr = + PackedConvWeightsQnnp& pack_data = cpp_custom_type_hack::cast(packed_weight); - auto packB = pack_ptr.w.get(); - auto kernel = pack_ptr.kernel; - auto kernel_zp = pack_ptr.w_zp; - auto kernel_scale = pack_ptr.w_scale; + auto* pack_w = pack_data.w.get(); + const auto& kernel = pack_data.kernel; + const auto& kernel_zp = pack_data.w_zp; + const auto& kernel_scale = pack_data.w_scale; const uint32_t kernel_h = kernel[0]; const uint32_t kernel_w = kernel[1]; // TODO Can be replaced with packB->getOutputChannels() when update pre-pack // to actually do the packing. - const auto out_ch = pack_ptr.bias.size(0); + const auto out_ch = pack_data.bias.size(0); // inputs are in semantic NCHW format - int N = act.size(0); - int in_ch = act.size(1); - int H = act.size(2); - int W = act.size(3); - int K = out_ch; // output channels - // TODO: change it to contiguous(MemoryFormat::ChannelsLast) once a perf - // regression of it is fixed. Today it's equivalent because `act` sizes - // are not used below - Tensor input_contig = act.permute({0, 2, 3, 1}).contiguous(); - - uint32_t stride_h = stride[0]; - uint32_t stride_w = stride[1]; - uint32_t pad_t = padding[0]; - uint32_t pad_l = padding[1]; - uint32_t dilation_h = dilation[0]; - uint32_t dilation_w = dilation[1]; - - auto output_min = ReluFused + const int N = act.size(0); + const int C = act.size(1); + const int H = act.size(2); + const int W = act.size(3); + const int M = out_ch; // output channels + + const Tensor act_nhwc = act.contiguous(MemoryFormat::ChannelsLast); + + const uint32_t stride_h = stride[0]; + const uint32_t stride_w = stride[1]; + const uint32_t pad_h = padding[0]; + const uint32_t pad_w = padding[1]; + const uint32_t dilation_h = dilation[0]; + const uint32_t dilation_w = dilation[1]; + + auto output_min = kReluFused ? activationLimits(output_scale, output_zero_point, Activation::RELU) .first : std::numeric_limits::min(); - auto output_max = ReluFused + auto output_max = kReluFused ? activationLimits(output_scale, output_zero_point, Activation::RELU) .second : std::numeric_limits::max(); @@ -330,29 +551,26 @@ class QConv2dInt8 final : public c10::OperatorKernel { {kernel_w, kernel_h}, {stride_w, stride_h}, {dilation_w, dilation_h}, - {pad_t, pad_l, pad_t, pad_l}, + {pad_h, pad_w, pad_h, pad_w}, groups, - in_ch, - out_ch, + C, + M, kernel_zp, kernel_scale, output_min, output_max); - // TODO: change convOutputShape to return NCHW sizes once perf is fixed - // Force output format to be NHWC - // TODO: consider preserving input format - // TODO: add MemoryFormat::ChannelsLast here once perf is fixed - auto input_scale = input_contig.q_scale(); + auto input_scale = act_nhwc.q_scale(); // Re-quantizing the bias based on input scale and weight scale. - if (!pack_ptr.input_scale.has_value() || - pack_ptr.input_scale.value() != input_scale) { + if (!pack_data.input_scale.has_value() || + pack_data.input_scale.value() != input_scale) { // Get the original weight and adjust it to uint8 from int8 auto weight_contig = - pack_ptr.orig_weight.contiguous(MemoryFormat::ChannelsLast); - auto bias_fp32 = pack_ptr.bias; - int8_t* w_data = (int8_t*)weight_contig.data_ptr(); + pack_data.orig_weight.contiguous(MemoryFormat::ChannelsLast); + auto bias_fp32 = pack_data.bias; + int8_t* w_data = + reinterpret_cast(weight_contig.data_ptr()); Tensor qnnp_weight = at::_empty_affine_quantized( weight_contig.sizes(), at::device(kCPU).dtype(kQUInt8), @@ -368,106 +586,69 @@ class QConv2dInt8 final : public c10::OperatorKernel { auto bias = at::quantize_per_tensor( bias_fp32, kernel_scale * input_scale, 0, kQInt32); // Update the input scale to not pack again. - pack_ptr.input_scale = input_scale; - pack_ptr.w.reset(); - pack_ptr.w = guts::make_unique( + pack_data.input_scale = input_scale; + pack_data.w.reset(); + pack_data.w = guts::make_unique( conv_p, - (uint8_t*)qnnp_w_data, - (int32_t*)bias.data_ptr()); - packB = pack_ptr.w.get(); + reinterpret_cast(qnnp_w_data), + reinterpret_cast(bias.data_ptr())); + pack_w = pack_data.w.get(); } - TORCH_INTERNAL_ASSERT(packB != nullptr, "Packed Weights are NULL"); - auto outShape = - convOutputShape(N, K, H, W, kernel, stride, padding, dilation); + TORCH_INTERNAL_ASSERT(pack_w != nullptr, "Packed Weights are NULL"); + const auto output_shape = MakeConvOutputShape( + N, M, {H, W}, kernel, stride, padding, dilation); TORCH_CHECK( std::all_of( - outShape.begin(), outShape.end(), [](int64_t i) { return i > 0; }), - "quantized::conv2d (qnnpack): each dimension of output tensor should be greater " - "than 0") - TORCH_CHECK( - (outShape[3] == out_ch), - "quantized::conv2d (qnnpack): Number of filters must be equal to number of " - "output channels") + output_shape.begin(), + output_shape.end(), + [](int64_t i) { return i > 0; }), + "quantized::conv2d (qnnpack): each dimension of output tensor should " + "be greater than 0.") // Allocate output Tensor and a buffer for QNNPACK to use Tensor output = at::_empty_affine_quantized( - outShape, + output_shape, at::device(kCPU).dtype(kQUInt8), output_scale, - output_zero_point); + output_zero_point, + MemoryFormat::ChannelsLast); - const pytorch_qnnp_status runStatus = qnnpack::qnnpackConv( + const pytorch_qnnp_status run_status = qnnpack::qnnpackConv( conv_p, - packB->getPackedWeights(), + pack_w->getPackedWeights(), N, H, W, - input_contig.q_scale(), - input_contig.q_zero_point(), - (uint8_t*)input_contig.data_ptr(), + act_nhwc.q_scale(), + act_nhwc.q_zero_point(), + reinterpret_cast(act_nhwc.data_ptr()), output.q_scale(), output.q_zero_point(), - (uint8_t*)output.data_ptr(), + reinterpret_cast(output.data_ptr()), caffe2::mobile_pthreadpool()); TORCH_INTERNAL_ASSERT( - runStatus == pytorch_qnnp_status_success, + run_status == pytorch_qnnp_status_success, "failed to run quantized::conv2d (qnnpack) operator"); - // TODO: remove permute once MemoryLayout is added above - return output.permute({0, 3, 1, 2}); + return output; } #endif - Tensor operator()( - Tensor act, - Tensor packed_weight, - torch::List stride, - torch::List padding, - torch::List dilation, - int64_t groups, - double output_scale, - int64_t output_zero_point) { - auto& ctx = at::globalContext(); -#ifdef USE_FBGEMM - if (ctx.qEngine() == at::QEngine::FBGEMM) { - return fbgemm_conv( - act, - packed_weight, - stride, - padding, - dilation, - groups, - output_scale, - output_zero_point); - } -#endif -#ifdef USE_PYTORCH_QNNPACK - if (ctx.qEngine() == at::QEngine::QNNPACK) { - return qnnpack_conv( - act, - packed_weight, - stride, - padding, - dilation, - groups, - output_scale, - output_zero_point); - } -#endif - TORCH_CHECK( - false, - "Didn't find engine for operation quantized::conv ", - toString(ctx.qEngine())); - } }; static auto registry = c10::RegisterOperators() .op("quantized::conv2d", - c10::RegisterOperators::options().kernel>( + c10::RegisterOperators::options().kernel>( TensorTypeId::QuantizedCPUTensorId)) .op("quantized::conv2d_relu", - c10::RegisterOperators::options().kernel>( + c10::RegisterOperators::options().kernel>( + TensorTypeId::QuantizedCPUTensorId)) + .op("quantized::conv3d", + c10::RegisterOperators::options().kernel>( + TensorTypeId::QuantizedCPUTensorId)) + .op("quantized::conv3d_relu", + c10::RegisterOperators::options().kernel>( TensorTypeId::QuantizedCPUTensorId)); } // namespace diff --git a/test/test_quantized.py b/test/test_quantized.py index 1b88fe0f67ec7..6d27b78e0ad7c 100644 --- a/test/test_quantized.py +++ b/test/test_quantized.py @@ -1319,6 +1319,114 @@ def _test_qconv_unpack_impl( np.testing.assert_equal( W_q.q_zero_point(), W_unpacked.q_zero_point()) + def _test_qconv_impl( + self, qconv_fn, qconv_prepack_fn, conv_op, batch_size, + input_channels_per_group, input_feature_map_shape, + output_channels_per_group, groups, kernels, strides, pads, dilations, + X_scale, X_zero_point, W_scale, W_zero_point, Y_scale, Y_zero_point, + use_bias, use_relu, use_channelwise + ): + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + # Padded input size should be at least as big as dilated kernel + for i in range(len(kernels)): + assume(input_feature_map_shape[i] + 2 * pads[i] + >= dilations[i] * (kernels[i] - 1) + 1) + W_scale = W_scale * output_channels + W_zero_point = W_zero_point * output_channels + # Resize W_scale and W_zero_points arrays equal to output_channels + W_scale = W_scale[:output_channels] + W_zero_point = W_zero_point[:output_channels] + # For testing, we use small values for weights and for activations + # so that no overflow occurs in vpmaddubsw instruction. If the + # overflow occurs in qconv implementation and if there is no + # overflow + # In reference we can't exactly match the results with reference. + # Please see the comment in qconv implementation file + # aten/src/ATen/native/quantized/cpu/qconv.cpp for more details. + (W_value_min, W_value_max) = (-5, 5) + # the operator expects them in the format + # (output_channels, input_channels/groups, + # kernel_d, kernel_h, kernel_w) + W_init = torch.randint( + W_value_min, + W_value_max, + (output_channels, input_channels_per_group,) + kernels, + ) + b_init = torch.randint(0, 10, (output_channels,)) + + (X_value_min, X_value_max) = (0, 4) + X_init = torch.randint( + X_value_min, + X_value_max, + (batch_size, input_channels,) + input_feature_map_shape, + ) + X = X_scale * (X_init - X_zero_point).float() + + if use_channelwise: + W_shape = (-1, 1) + (1,) * len(kernels) + W_scales_tensor = torch.tensor(W_scale, dtype=torch.float) + W_zero_points_tensor = torch.tensor(W_zero_point, dtype=torch.float) + W = W_scales_tensor.reshape(*W_shape) * ( + W_init.float() - W_zero_points_tensor.reshape(*W_shape)).float() + b = X_scale * W_scales_tensor * b_init.float() + else: + W = W_scale[0] * (W_init - W_zero_point[0]).float() + b = X_scale * W_scale[0] * b_init.float() + + # Assign weights + conv_op.weight = torch.nn.Parameter(W, requires_grad=False) + conv_op.bias = torch.nn.Parameter( + b, requires_grad=False) if use_bias else None + result_ref = conv_op(X) + if use_relu: + relu = torch.nn.ReLU() + result_ref = relu(result_ref) + + # Quantize reference results for comparision + result_ref_q = torch.quantize_per_tensor( + result_ref, scale=Y_scale, zero_point=Y_zero_point, + dtype=torch.quint8) + X_q = torch.quantize_per_tensor( + X, scale=X_scale, zero_point=X_zero_point, dtype=torch.quint8) + if use_channelwise: + W_q = torch.quantize_per_channel( + W, W_scales_tensor, W_zero_points_tensor.long(), 0, + dtype=torch.qint8) + else: + W_q = torch.quantize_per_tensor( + W, scale=W_scale[0], zero_point=W_zero_point[0], + dtype=torch.qint8) + + bias_float = b if use_bias else None + W_prepack = qconv_prepack_fn( + W_q, bias_float, strides, pads, dilations, groups) + Y_q = qconv_fn( + X_q, + W_prepack, + strides, + pads, + dilations, + groups, + Y_scale, + Y_zero_point, + ) + + # Make sure the results match + # assert_array_almost_equal compares using the following formula: + # abs(desired-actual) < 1.5 * 10**(-decimal) + # (https://docs.scipy.org/doc/numpy/reference/generated/numpy.testing.assert_almost_equal.html) + # We use decimal = 0 to ignore off-by-1 differences between + # reference and test. Off-by-1 differences arise due to the order of + # round and zero_point addition operation, i.e., if addition + # followed by round is used by reference and round followed by + # addition is used by test, the results may differ by 1. + # For example, the result of round(2.5) + 1 is 3 while + # round(2.5 + 1) is 4 assuming the rounding mode is + # round-to-nearest, ties-to-even. + np.testing.assert_array_almost_equal( + result_ref_q.int_repr().numpy(), Y_q.int_repr().numpy(), decimal=0) + """Tests the correctness of quantized convolution op.""" @given(batch_size=st.integers(1, 3), input_channels_per_group=st.sampled_from([2, 4, 5, 8, 16, 32]), @@ -1376,106 +1484,33 @@ def test_qconv( return use_channelwise = False + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + kernels = (kernel_h, kernel_w) + strides = (stride_h, stride_w) + pads = (pad_h, pad_w) + dilations = (dilation, dilation) + with override_quantized_engine(qengine): qconv = torch.ops.quantized.conv2d if use_relu: qconv = torch.ops.quantized.conv2d_relu qconv_prepack = torch.ops.quantized.conv_prepack - # C - input_channels = input_channels_per_group * groups - # K - output_channels = output_channels_per_group * groups - dilation_h = dilation_w = dilation - # Padded input size should be at least as big as dilated kernel - assume(height + 2 * pad_h >= dilation_h * (kernel_h - 1) + 1) - assume(width + 2 * pad_w >= dilation_w * (kernel_w - 1) + 1) - W_scale = W_scale * output_channels - W_zero_point = W_zero_point * output_channels - # Resize W_scale and W_zero_points arrays equal to output_channels - W_scale = W_scale[:output_channels] - W_zero_point = W_zero_point[:output_channels] - # For testing, we use small values for weights and for activations so that no overflow occurs - # in vpmaddubsw instruction. If the overflow occurs in qconv implementation and if there is no overflow - # in reference we can't exactly match the results with reference. - # Please see the comment in qconv implementation file (aten/src/ATen/native/quantized/cpu/qconv.cpp) - # for more details. - W_value_min = -5 - W_value_max = 5 - # the operator expects them in the format (output_channels, input_channels/groups, kernel_h, kernel_w) - W_init = torch.from_numpy( - np.random.randint( - W_value_min, - W_value_max, - (output_channels, int(input_channels / groups), kernel_h, kernel_w)), - ) - b_init = torch.from_numpy(np.random.randint(0, 10, (output_channels,))) - stride = [stride_h, stride_w] - pad = [pad_h, pad_w] - dilation = [dilation_h, dilation_w] - X_value_min = 0 - X_value_max = 4 - X_init = torch.from_numpy(np.random.randint( - X_value_min, X_value_max, (batch_size, input_channels, height, width))) - X = X_scale * (X_init - X_zero_point).to(dtype=torch.float) - if use_channelwise: - W_scales_tensor = torch.tensor(W_scale, dtype=torch.float) - W_zero_points_tensor = torch.tensor(W_zero_point, dtype=torch.float) - W = W_scales_tensor.reshape(-1, 1, 1, 1) * (W_init.to(dtype=torch.float) - - W_zero_points_tensor.reshape(-1, 1, 1, 1)).to(dtype=torch.float) - b = X_scale * W_scales_tensor * (b_init - 0).to(dtype=torch.float) - else: - W = W_scale[0] * (W_init - W_zero_point[0]).to(dtype=torch.float) - b = X_scale * W_scale[0] * (b_init - 0).to(dtype=torch.float) - # Existing floating point conv operator - conv_op = torch.nn.Conv2d(input_channels, - output_channels, - (kernel_h, kernel_w), - (stride_h, stride_w), - (pad_h, pad_w), - (dilation_h, dilation_w), - groups) - # assign weights - conv_op.weight = torch.nn.Parameter(W, requires_grad=False) - conv_op.bias = torch.nn.Parameter(b, requires_grad=False) if use_bias else None - result_ref = conv_op(X) - if use_relu: - relu = torch.nn.ReLU() - result_ref = relu(result_ref) - # quantize reference results for comparision - result_ref_q = torch.quantize_per_tensor(result_ref, scale=Y_scale, zero_point=Y_zero_point, dtype=torch.quint8) - X_q = torch.quantize_per_tensor(X, scale=X_scale, zero_point=X_zero_point, dtype=torch.quint8) - if use_channelwise: - W_q = torch.quantize_per_channel(W, - W_scales_tensor, - W_zero_points_tensor.to(dtype=torch.long), - 0, - dtype=torch.qint8) - else: - W_q = torch.quantize_per_tensor(W, scale=W_scale[0], zero_point=W_zero_point[0], dtype=torch.qint8) - bias_float = b if use_bias else None - W_prepack = qconv_prepack(W_q, bias_float, stride, pad, dilation, groups) - Y_q = qconv( - X_q, - W_prepack, - stride, - pad, - dilation, + conv_op = torch.nn.Conv2d( + input_channels, + output_channels, + kernels, + strides, + pads, + dilations, groups, - Y_scale, - Y_zero_point, ) - # Make sure the results match - # assert_array_almost_equal compares using the following formula: - # abs(desired-actual) < 1.5 * 10**(-decimal) - # (https://docs.scipy.org/doc/numpy/reference/generated/numpy.testing.assert_almost_equal.html) - # We use decimal = 0 to ignore off-by-1 differences between reference and - # test. Off-by-1 differences arise due to the order of round and - # zero_point addition operation, i.e., if addition followed by round is - # used by reference and round followed by addition is used by test, the - # results may differ by 1. - # For example, the result of round(2.5) + 1 is 3 while round(2.5 + 1) is 4 - # assuming the rounding mode is round-to-nearest, ties-to-even. - np.testing.assert_array_almost_equal(result_ref_q.int_repr().numpy(), Y_q.int_repr().numpy(), decimal=0) + self._test_qconv_impl( + qconv, qconv_prepack, conv_op, batch_size, + input_channels_per_group, (height, width), + output_channels_per_group, groups, kernels, strides, pads, + dilations, X_scale, X_zero_point, W_scale, W_zero_point, + Y_scale, Y_zero_point, use_bias, use_relu, use_channelwise) """Tests the correctness of the quantized::qconv_unpack op.""" @given( @@ -1514,6 +1549,94 @@ def test_qconv_unpack( qconv_prepack, qconv_unpack, inputs, (stride_h, stride_w), (pad_h, pad_w), channelwise) + @given(batch_size=st.integers(1, 4), + input_channels_per_group=st.sampled_from([2, 4, 5, 8, 16]), + D=st.integers(4, 8), + H=st.integers(4, 8), + W=st.integers(4, 8), + output_channels_per_group=st.sampled_from([2, 4, 5, 8, 16]), + groups=st.integers(1, 3), + kernel_d=st.integers(1, 4), + kernel_h=st.integers(1, 4), + kernel_w=st.integers(1, 4), + stride_d=st.integers(1, 2), + stride_h=st.integers(1, 2), + stride_w=st.integers(1, 2), + pad_d=st.integers(0, 2), + pad_h=st.integers(0, 2), + pad_w=st.integers(0, 2), + dilation=st.integers(1, 2), + X_scale=st.floats(1.2, 1.6), + X_zero_point=st.integers(0, 4), + W_scale=st.lists(st.floats(0.2, 1.6), min_size=1, max_size=2), + W_zero_point=st.lists(st.integers(-5, 5), min_size=1, max_size=2), + Y_scale=st.floats(4.2, 5.6), + Y_zero_point=st.integers(0, 4), + use_bias=st.booleans(), + use_relu=st.booleans(), + use_channelwise=st.booleans(), + qengine=st.sampled_from(("fbgemm",))) + def test_qconv3d( + self, + batch_size, + input_channels_per_group, + D, + H, + W, + output_channels_per_group, + groups, + kernel_d, + kernel_h, + kernel_w, + stride_d, + stride_h, + stride_w, + pad_d, + pad_h, + pad_w, + dilation, + X_scale, + X_zero_point, + W_scale, + W_zero_point, + Y_scale, + Y_zero_point, + use_bias, + use_relu, + use_channelwise, + qengine + ): + if qengine not in torch.backends.quantized.supported_engines: + return + + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + kernels = (kernel_d, kernel_h, kernel_w) + strides = (stride_d, stride_h, stride_w) + pads = (pad_d, pad_h, pad_w) + dilations = (dilation, dilation, dilation) + + with override_quantized_engine(qengine): + qconv = torch.ops.quantized.conv3d + if use_relu: + qconv = torch.ops.quantized.conv3d_relu + qconv_prepack = torch.ops.quantized.conv3d_prepack + conv_op = torch.nn.Conv3d( + input_channels, + output_channels, + kernels, + strides, + pads, + dilations, + groups, + ) + self._test_qconv_impl( + qconv, qconv_prepack, conv_op, batch_size, + input_channels_per_group, (D, H, W), output_channels_per_group, + groups, kernels, strides, pads, dilations, X_scale, + X_zero_point, W_scale, W_zero_point, Y_scale, Y_zero_point, + use_bias, use_relu, use_channelwise) + """Tests the correctness of the quantized::qconv3d_unpack op.""" @given( inputs=hu.tensor_conv( From dff159804f92d9f6caad64c08bf21a770f5004f7 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Tue, 29 Oct 2019 07:40:31 -0700 Subject: [PATCH 18/64] Revert D18170995: Simplify copy kernel Test Plan: revert-hammer Differential Revision: D18170995 Original commit changeset: 461b53641813 fbshipit-source-id: 1ebb119325d746a153982ac3209d3570a7e18d88 --- aten/src/ATen/native/cuda/Copy.cu | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu index b9c90ec156cae..407d5df90ca5f 100644 --- a/aten/src/ATen/native/cuda/Copy.cu +++ b/aten/src/ATen/native/cuda/Copy.cu @@ -8,12 +8,20 @@ #include #include #include +#include namespace at { namespace native { using namespace at::cuda; +template +void copy_kernel_impl(TensorIterator& iter) { + gpu_kernel(iter, []GPU_LAMBDA(src_t x) -> dst_t { + return c10::static_cast_with_inter_type(x); + }); +} + // device-to-device copy, does type conversion static void copy_device_to_device(TensorIterator& iter, bool non_blocking) { int64_t numel = iter.numel(); @@ -58,11 +66,11 @@ static void copy_device_to_device(TensorIterator& iter, bool non_blocking) { cudaMemcpyDeviceToDevice, copy_stream)); } else { - // this is done intentionally done after build because copy has a "promotion" - // rule that always "promote" to target dtype. - iter.promote_common_dtype(); - AT_DISPATCH_ALL_TYPES_AND3(kHalf, kBool, kBFloat16, iter.dtype(0), "copy_", [&] { - gpu_kernel(iter, []GPU_LAMBDA(scalar_t x) { return x; }); + AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(0), "copy_", [&] { + using dst_t = scalar_t; + AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(1), "copy_", [&] { + copy_kernel_impl(iter); + }); }); } From 0301f5f30be7a0c89427a8494d32c600143bd4a4 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Tue, 29 Oct 2019 07:40:31 -0700 Subject: [PATCH 19/64] Revert D18170997: Make TensorIterator stop promoting types by copying Test Plan: revert-hammer Differential Revision: D18170997 Original commit changeset: 9c82c1c89583 fbshipit-source-id: 8862d9628864d23a087f2895870386772a634e45 --- aten/src/ATen/native/TensorIterator.cpp | 31 +++++------ aten/src/ATen/native/TensorIterator.h | 7 --- aten/src/ATen/native/cuda/BinaryOpsKernel.cu | 42 +++++++-------- aten/src/ATen/native/cuda/Loops.cuh | 56 ++++---------------- aten/src/ATen/test/tensor_iterator_test.cpp | 1 - 5 files changed, 44 insertions(+), 93 deletions(-) diff --git a/aten/src/ATen/native/TensorIterator.cpp b/aten/src/ATen/native/TensorIterator.cpp index cbf9cdbc83fed..5d8b94f395e1c 100644 --- a/aten/src/ATen/native/TensorIterator.cpp +++ b/aten/src/ATen/native/TensorIterator.cpp @@ -148,7 +148,7 @@ static void validate_dtype(OperandInfo& op, ScalarType common_dtype, CommonDType } } -static void maybe_copy_casting_to_common_dtype(OperandInfo& op, ScalarType common_dtype) { +static void maybe_promote_common_dtype(OperandInfo& op, ScalarType common_dtype) { if (op.tensor.defined() && op.tensor.scalar_type() != common_dtype) { op.dtype = common_dtype; @@ -165,7 +165,7 @@ static void maybe_copy_casting_to_common_dtype(OperandInfo& op, ScalarType commo void TensorIterator::compute_types() { bool missing_dtypes = false; bool missing_output_dtypes = false; - common_dtype_ = dtype(); + ScalarType common_dtype = dtype(); for (auto& op : operands_) { if (!op.tensor.defined() && !op.is_type_defined()) { missing_dtypes = true; @@ -183,24 +183,22 @@ void TensorIterator::compute_types() { bool compute_common_dtype_only_for_inputs = (common_dtype_strategy_ == CommonDTypeStrategy::PROMOTE_INPUTS); bool may_have_differing_types = true; - bool common_device_is_cuda = false; if (missing_dtypes || compute_common_dtype) { auto operands = compute_common_dtype_only_for_inputs ? at::ArrayRef(operands_).slice(noutputs()) : operands_; auto common_type = compute_common_type_(operands); auto common_device = std::get<0>(common_type); - common_device_is_cuda = common_device.is_cuda(); - common_dtype_ = std::get<1>(common_type); + common_dtype = std::get<1>(common_type); may_have_differing_types = !std::get<2>(common_type); bool has_cpu_scalar = false; for (auto& op : operands_) { if (!op.is_type_defined()) { op.device = common_device; - op.dtype = common_dtype_; + op.dtype = common_dtype; } else if (compute_common_dtype && - (op.device != common_device || op.dtype != common_dtype_)) { + (op.device != common_device || op.dtype != common_dtype)) { if (allow_cpu_scalars_ && op.tensor.defined() && op.tensor.dim() == 0 && - common_device_is_cuda && op.tensor.device().is_cpu() && + common_device.is_cuda() && op.tensor.device().is_cpu() && !has_cpu_scalar) { // don't cast CPU scalars in CUDA ops that directly support them. op.device = op.tensor.device(); @@ -208,8 +206,8 @@ void TensorIterator::compute_types() { has_cpu_scalar = true; } else if (promote_gpu_output_dtypes_ && op.tensor.defined() && !op.is_output && - op.tensor.scalar_type() == kHalf && common_dtype_ == kFloat && - op.tensor.device().is_cuda() && common_device_is_cuda) { + op.tensor.scalar_type() == kHalf && common_dtype == kFloat && + op.tensor.device().is_cuda() && common_device.is_cuda()) { // allow input tensor type upcasting for fp16 to fp32 in fused kernel // on GPU op.device = op.tensor.device(); @@ -219,7 +217,7 @@ void TensorIterator::compute_types() { if (compute_common_dtype_only_for_inputs && op.is_output) { op.dtype = op.tensor.scalar_type(); } else { - op.dtype = common_dtype_; + op.dtype = common_dtype; } } } @@ -228,17 +226,12 @@ void TensorIterator::compute_types() { for (auto &op : operands_) { if (may_have_differing_types) { - validate_dtype(op, common_dtype_, common_dtype_strategy_); - bool cast_by_copy = compute_common_dtype && !common_device_is_cuda && (!compute_common_dtype_only_for_inputs || !op.is_output); - if (cast_by_copy) { - maybe_copy_casting_to_common_dtype(op, common_dtype_); + validate_dtype(op, common_dtype, common_dtype_strategy_); + if (compute_common_dtype && (!compute_common_dtype_only_for_inputs || !op.is_output)) { + maybe_promote_common_dtype(op, common_dtype); } } - if (op.tensor.defined() && op.tensor.scalar_type() != common_dtype_) { - have_differing_types_ = true; - } - if (op.tensor.defined() && op.device != op.tensor.device()) { if (op.is_output) { AT_ERROR("output with device ", op.tensor.device(), diff --git a/aten/src/ATen/native/TensorIterator.h b/aten/src/ATen/native/TensorIterator.h index eef28b710cd4a..69ed861c888c7 100644 --- a/aten/src/ATen/native/TensorIterator.h +++ b/aten/src/ATen/native/TensorIterator.h @@ -191,7 +191,6 @@ struct CAFFE2_API TensorIterator { IntArrayRef strides(int arg) const { return operands_[arg].stride_bytes; } void* data_ptr(int arg) const; ScalarType dtype(int arg=0) const { return operands_[arg].tensor.scalar_type(); } - ScalarType common_dtype() const { return common_dtype_; } ScalarType input_dtype(int arg=0) const { return operands_[num_outputs_ + arg].dtype; } Device device(int arg=0) const { return operands_[arg].device; } DeviceType device_type(int arg=0) const { return device(arg).type(); } @@ -287,10 +286,6 @@ struct CAFFE2_API TensorIterator { /// CUDA reductions. bool is_final_output() const { return final_output_; } - bool needs_dynamic_casting() const { - return (common_dtype_strategy_ != CommonDTypeStrategy::NONE) && have_differing_types_; - } - void set_check_mem_overlap(bool check_mem_overlap) { check_mem_overlap_ = check_mem_overlap; } @@ -357,7 +352,6 @@ struct CAFFE2_API TensorIterator { SmallVector operands_; int num_outputs_ = 0; CommonDTypeStrategy common_dtype_strategy_ = CommonDTypeStrategy::CHECK; - ScalarType common_dtype_ = ScalarType::Undefined; bool has_coalesced_dimensions_ = false; bool accumulate_ = false; bool resize_outputs_ = true; @@ -366,7 +360,6 @@ struct CAFFE2_API TensorIterator { bool promote_gpu_output_dtypes_ = false; bool final_output_ = true; bool check_mem_overlap_ = false; - bool have_differing_types_ = false; }; /// A container-like struct that acts as if it contains splits of a /// TensorIterator that can use 32-bit indexing. Taken together the splits cover diff --git a/aten/src/ATen/native/cuda/BinaryOpsKernel.cu b/aten/src/ATen/native/cuda/BinaryOpsKernel.cu index 95f1a75883b56..002539717de44 100644 --- a/aten/src/ATen/native/cuda/BinaryOpsKernel.cu +++ b/aten/src/ATen/native/cuda/BinaryOpsKernel.cu @@ -13,7 +13,7 @@ namespace at { namespace native { void add_kernel_cuda(TensorIterator& iter, Scalar alpha_scalar) { - AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "add_cuda/sub_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(), "add_cuda/sub_cuda", [&]() { auto alpha = alpha_scalar.to(); gpu_kernel_with_scalars(iter, [alpha]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a + alpha * b; @@ -26,11 +26,11 @@ static void sub_kernel_cuda(TensorIterator& iter, Scalar alpha_scalar) { } void div_kernel_cuda(TensorIterator& iter) { - if (!isIntegralType(iter.common_dtype(), /*includeBool*/ false) && iter.is_cpu_scalar(2)) { + if (!isIntegralType(iter.dtype(), /*includeBool*/ false) && iter.is_cpu_scalar(2)) { // optimization for floating-point types: if the second operand is a CPU // scalar, compute a * reciprocal(b). Note that this may lose one bit of // precision compared to computing the division. - AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "div_cuda", [&]() { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.dtype(), "div_cuda", [&]() { auto inv_b = scalar_t(1.0 / iter.scalar_value(2)); iter.remove_operand(2); gpu_kernel(iter, [inv_b]GPU_LAMBDA(scalar_t a) -> scalar_t { @@ -38,7 +38,7 @@ void div_kernel_cuda(TensorIterator& iter) { }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "div_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "div_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a / b; }); @@ -47,13 +47,13 @@ void div_kernel_cuda(TensorIterator& iter) { } void mul_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { // Workaround for the error: '*' in boolean context, suggest '&&' instead [-Werror=int-in-bool-context] gpu_kernel_with_scalars(iter, []GPU_LAMBDA(bool a, bool b) -> bool { return a && b; }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "mul_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "mul_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a * b; }); @@ -62,7 +62,7 @@ void mul_kernel_cuda(TensorIterator& iter) { } void atan2_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.common_dtype(), "atan2_cuda", [&]() { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(iter.dtype(), "atan2_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return ::atan2(a, b); }); @@ -70,14 +70,14 @@ void atan2_kernel_cuda(TensorIterator& iter) { } void logical_xor_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "logical_xor_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return bool(a) != bool(b); }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "logical_xor_cuda", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "logical_xor_cuda", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return static_cast(bool(a) != bool(b)); }); @@ -86,14 +86,14 @@ void logical_xor_kernel_cuda(TensorIterator& iter) { } void lt_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "lt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a < b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "lt_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "lt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a < b; }); @@ -102,14 +102,14 @@ void lt_kernel_cuda(TensorIterator& iter) { } void le_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "le_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a <= b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "le_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "le_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a <= b; }); @@ -118,14 +118,14 @@ void le_kernel_cuda(TensorIterator& iter) { } void gt_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "gt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a > b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "gt_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "gt_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a > b; }); @@ -134,14 +134,14 @@ void gt_kernel_cuda(TensorIterator& iter) { } void ge_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "ge_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a >= b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "ge_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "ge_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a >= b; }); @@ -150,14 +150,14 @@ void ge_kernel_cuda(TensorIterator& iter) { } void eq_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "eq_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a == b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "eq_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a == b; }); @@ -166,14 +166,14 @@ void eq_kernel_cuda(TensorIterator& iter) { } void ne_kernel_cuda(TensorIterator& iter) { - if (iter.common_dtype() == ScalarType::Bool) { + if (iter.dtype() == ScalarType::Bool) { AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.input_dtype(), "ne_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool { return a != b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.common_dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND(kHalf, iter.dtype(), "ne_cpu", [&]() { gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { return a != b; }); diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh index e5b52accbb253..38d25982fe0b5 100644 --- a/aten/src/ATen/native/cuda/Loops.cuh +++ b/aten/src/ATen/native/cuda/Loops.cuh @@ -35,7 +35,6 @@ #include #include #include -#include // Marks a lambda as executable on both the host and device. The __host__ // attribute is important so that we can access static type information from @@ -117,20 +116,6 @@ invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[] return invoke_impl(f, data, strides, i, Indices{}); } -template -C10_HOST_DEVICE typename traits::result_type -invoke_impl(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i, - c10::guts::index_sequence) { - return f(c10::fetch_and_cast::type>(dtypes[I], data[I] + i * strides[I])...); -} - -template > -C10_HOST_DEVICE typename traits::result_type -invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i) { - using Indices = c10::guts::make_index_sequence; - return invoke_impl(f, data, strides, dtypes, i, Indices{}); -} - template void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { using traits = function_traits; @@ -145,10 +130,6 @@ void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { data[i] = (char*)iter.data_ptr(i); } - at::detail::Array dtypes; - for (int i = 0; i < ntensors; i++) { - dtypes[i] = iter.tensor(i).scalar_type(); - } int64_t numel = iter.numel(); if (iter.is_trivial_1d()) { @@ -157,35 +138,19 @@ void gpu_kernel_impl(TensorIterator& iter, const func_t& f) { for (int i = 0; i < ntensors; i++) { strides[i] = inner_strides[i]; } + - if (iter.needs_dynamic_casting()) { - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - void* out = data[0] + strides[0] * idx; - arg0_t result = invoke(f, &data.data[1], &strides.data[1], &dtypes.data[1], idx); - c10::cast_and_store(dtypes[0], out, result); - }); - } else { - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - arg0_t* out = (arg0_t*)(data[0] + strides[0] * idx); - *out = invoke(f, &data.data[1], &strides.data[1], idx); - }); - } + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + arg0_t* out = (arg0_t*)(data[0] + strides[0] * idx); + *out = invoke(f, &data.data[1], &strides.data[1], idx); + }); } else { auto offset_calc = make_offset_calculator(iter); - if (iter.needs_dynamic_casting()) { - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - auto offsets = offset_calc.get(idx); - void* out = data[0] + offsets[0]; - arg0_t result = invoke(f, &data.data[1], &offsets.data[1], &dtypes.data[1], 1); - c10::cast_and_store(dtypes[0], out, result); - }); - } else { - launch_kernel(numel, [=]GPU_LAMBDA(int idx) { - auto offsets = offset_calc.get(idx); - arg0_t* out = (arg0_t*)(data[0] + offsets[0]); - *out = invoke(f, &data.data[1], &offsets.data[1], 1); - }); - } + launch_kernel(numel, [=]GPU_LAMBDA(int idx) { + auto offsets = offset_calc.get(idx); + arg0_t* out = (arg0_t*)(data[0] + offsets[0]); + *out = invoke(f, &data.data[1], &offsets.data[1], 1); + }); } } @@ -209,6 +174,7 @@ void gpu_kernel(TensorIterator& iter, const func_t& f) { } gpu_kernel_impl(iter, f); + iter.cast_outputs(); } template diff --git a/aten/src/ATen/test/tensor_iterator_test.cpp b/aten/src/ATen/test/tensor_iterator_test.cpp index a2bbc653298fa..200091de260eb 100644 --- a/aten/src/ATen/test/tensor_iterator_test.cpp +++ b/aten/src/ATen/test/tensor_iterator_test.cpp @@ -190,7 +190,6 @@ TEST(TensorIteratorTest, ComputeCommonDTypeInputOnly) { EXPECT_TRUE(iter.dtype(0) == at::kBool); EXPECT_TRUE(iter.dtype(1) == at::kDouble); EXPECT_TRUE(iter.dtype(2) == at::kDouble); - EXPECT_TRUE(iter.common_dtype() == at::kDouble); } TEST(TensorIteratorTest, DoNotComputeCommonDTypeInputOnly) { From 5fbec1f55df59156edff4084023086823227fbb0 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Tue, 29 Oct 2019 07:40:31 -0700 Subject: [PATCH 20/64] Revert D18170996: Move type casting to c10/util/TypeCast.h Test Plan: revert-hammer Differential Revision: D18170996 Original commit changeset: 41658afd5c0a fbshipit-source-id: 394e84bbc52bdd708609304261ffa1513a771d57 --- aten/src/ATen/cpu/vec256/vec256_base.h | 4 +- aten/src/ATen/native/Copy.h | 38 ++++++ aten/src/ATen/native/cpu/CopyKernel.cpp | 7 +- aten/src/ATen/native/cuda/Copy.cu | 3 +- c10/core/ScalarType.h | 4 - c10/util/TypeCast.h | 172 ------------------------ 6 files changed, 45 insertions(+), 183 deletions(-) delete mode 100644 c10/util/TypeCast.h diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h index d7c4aab6ce209..64063b3da2b92 100644 --- a/aten/src/ATen/cpu/vec256/vec256_base.h +++ b/aten/src/ATen/cpu/vec256/vec256_base.h @@ -13,7 +13,6 @@ #include #include #include -#include #if defined(__GNUC__) #define __at_align32__ __attribute__((aligned(32))) @@ -682,7 +681,8 @@ inline void convert(const src_T *src, dst_T *dst, int64_t n) { # pragma unroll #endif for (int64_t i = 0; i < n; i++) { - *dst = c10::static_cast_with_inter_type(*src); + *dst = static_cast( + static_cast>(*src)); src++; dst++; } diff --git a/aten/src/ATen/native/Copy.h b/aten/src/ATen/native/Copy.h index 2dfd9e9f4922b..a8d16f6f7b871 100644 --- a/aten/src/ATen/native/Copy.h +++ b/aten/src/ATen/native/Copy.h @@ -9,6 +9,44 @@ struct TensorIterator; namespace native { +// Note [Implicit conversion between signed and unsigned] +// C and C++ have a lovely set of implicit conversion rules, where casting +// signed integral values to unsigned integral values is always valid +// (it basically treats the value as if using modulo arithmetic), however +// converting negative floating point values to unsigned integral types +// is UB! This means that: (double)-1 -> (int64_t)-1 -> (uint8_t)255 is +// guaranteed to look like this, but we have (double)-1 -> (uint8_t) +// because it's UB. This also makes UBSan really angry. +// +// I think those rules are stupid and we really shouldn't conform to them. +// The structs below ensure that for all unsigned types we use (currently +// only uint8_t), we will do an intermediate convertion via int64_t, +// to ensure that any negative values are wrapped around correctly. +// +// Note that conversions from doubles to signed integral types that can't +// represent a particular value after truncating the fracitonal part are UB as well, +// but fixing them is not as simple as adding an int64_t intermediate, beacuse the +// int64_t -> conversion is UB for those large values anyway. +// I guess in that case we just have to live with that, but it's definitely less +// surprising than the thing above. +// +// For the curious: +// https://en.cppreference.com/w/cpp/language/implicit_conversion +// The relevant paragraph is "Floating-integral conversions". + +template +struct inter_copy_type { + using type = T; +}; + +template <> +struct inter_copy_type { + using type = int64_t; +}; + +template +using inter_copy_type_t = typename inter_copy_type::type; + using copy_fn = void (*)(TensorIterator&, bool non_blocking); DECLARE_DISPATCH(copy_fn, copy_stub); diff --git a/aten/src/ATen/native/cpu/CopyKernel.cpp b/aten/src/ATen/native/cpu/CopyKernel.cpp index a48c7afc07855..2362b8f20a9c6 100644 --- a/aten/src/ATen/native/cpu/CopyKernel.cpp +++ b/aten/src/ATen/native/cpu/CopyKernel.cpp @@ -4,7 +4,6 @@ #include #include #include -#include namespace at { namespace native { @@ -15,7 +14,8 @@ void copy_kernel_cast(TensorIterator& iter) { if (isComplexType(iter.dtype(1))) { AT_DISPATCH_COMPLEX_TYPES(iter.dtype(1), "copy_kernel_cast", [&] { cpu_kernel(iter, [=](scalar_t a) -> self_T { - return c10::static_cast_with_inter_type(std::real(a)); + return static_cast( + static_cast>(std::real(a))); }); }); } @@ -28,7 +28,8 @@ void copy_kernel_cast(TensorIterator& iter) { "copy_kernel_cast", [&] { cpu_kernel(iter, [=](scalar_t a) -> self_T { - return c10::static_cast_with_inter_type(a); + return static_cast( + static_cast>(a)); }); }); } diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu index 407d5df90ca5f..4e201a6576af3 100644 --- a/aten/src/ATen/native/cuda/Copy.cu +++ b/aten/src/ATen/native/cuda/Copy.cu @@ -8,7 +8,6 @@ #include #include #include -#include namespace at { namespace native { @@ -18,7 +17,7 @@ using namespace at::cuda; template void copy_kernel_impl(TensorIterator& iter) { gpu_kernel(iter, []GPU_LAMBDA(src_t x) -> dst_t { - return c10::static_cast_with_inter_type(x); + return static_cast(static_cast>(x)); }); } diff --git a/c10/core/ScalarType.h b/c10/core/ScalarType.h index b9be182946b4c..de5a0d2f20de2 100644 --- a/c10/core/ScalarType.h +++ b/c10/core/ScalarType.h @@ -166,10 +166,6 @@ struct ScalarTypeToCPPType { _(c10::quint8, QUInt8) \ _(c10::qint32, QInt32) -#define AT_FORALL_COMPLEX_TYPES(_) \ - _(std::complex, ComplexFloat) \ - _(std::complex, ComplexDouble) - static inline caffe2::TypeMeta scalarTypeToTypeMeta(ScalarType scalar_type) { #define DEFINE_CASE(ctype, name) \ case ScalarType::name: \ diff --git a/c10/util/TypeCast.h b/c10/util/TypeCast.h deleted file mode 100644 index 1a00a0abb7315..0000000000000 --- a/c10/util/TypeCast.h +++ /dev/null @@ -1,172 +0,0 @@ -#pragma once - -#include -#include -#include - - -namespace c10 { - -// Note [Implicit conversion between signed and unsigned] -// C and C++ have a lovely set of implicit conversion rules, where casting -// signed integral values to unsigned integral values is always valid -// (it basically treats the value as if using modulo arithmetic), however -// converting negative floating point values to unsigned integral types -// is UB! This means that: (double)-1 -> (int64_t)-1 -> (uint8_t)255 is -// guaranteed to look like this, but we have (double)-1 -> (uint8_t) -// because it's UB. This also makes UBSan really angry. -// -// I think those rules are stupid and we really shouldn't conform to them. -// The structs below ensure that for all unsigned types we use (currently -// only uint8_t), we will do an intermediate convertion via int64_t, -// to ensure that any negative values are wrapped around correctly. -// -// Note that conversions from doubles to signed integral types that can't -// represent a particular value after truncating the fracitonal part are UB as well, -// but fixing them is not as simple as adding an int64_t intermediate, beacuse the -// int64_t -> conversion is UB for those large values anyway. -// I guess in that case we just have to live with that, but it's definitely less -// surprising than the thing above. -// -// For the curious: -// https://en.cppreference.com/w/cpp/language/implicit_conversion -// The relevant paragraph is "Floating-integral conversions". - -template -struct inter_copy_type { - using type = T; -}; - -template <> -struct inter_copy_type { - using type = int64_t; -}; - -template -using inter_copy_type_t = typename inter_copy_type::type; - -template -C10_HOST_DEVICE inline dest_t static_cast_with_inter_type(src_t src) { - return static_cast( - static_cast>(src)); -} - -// Dynamic type casting utils: -// - fetch_and_cast -// - cast_and_store -// -// fetch_and_cast fetch a value with dynamic type specified by a ScalarType -// from a void pointer and cast it to a static type. -// -// cast_and_store casts a static typed value into dynamic type specified -// by a ScalarType, and store it into a void pointer. -// -// NOTE: -// -// Dynamic casting allows us to support type promotion without blowing up -// the combination space: For example, without dynamic cast, in order to -// implement `add_` with type promotion, we would need something like -// -// AT_DISPATCH_ALL_TYPES(output.dtype(), -// AT_DISPATCH_ALL_TYPES(input1.dtype(), -// AT_DISPATCH_ALL_TYPES(input2.dtype(), -// [](arg0_t a, arg1_t b) -> out_t { return a + b; } -// ) -// ) -// ) -// -// If we support N dtypes, the above code would generate the a+b kernel for -// all the N * N * N different supported types, the compilation time and -// binary size would become horrible. -// -// Dynamic casting might sounds like a bad idea in terms of performance. -// Especially if you ever do it in a loop, you are going to do a billion tests. -// But in practice it is not as bad as it might look: -// -// - on CPU, this is a branch that always has the same outcome, therefore -// hopefully the branch predictor could do the job pretty well -// - on GPU, these branches will not diverge, so we could still have the same -// warp executing the same line of code -// - Most kernels, like `add`, are bandwidth bound, adding a few clock cycles to -// check an integer does not hurt the performance much because the ALUs would -// wait for load instructions anyway. -// -// For the discussion and benchmark, refer to: -// - https://github.com/pytorch/pytorch/pull/28343 -// - https://github.com/pytorch/pytorch/pull/28344 -// - https://github.com/pytorch/pytorch/pull/28345 -// - -#ifdef C10_HOST_DEVICE -#define ERROR_UNSUPPORTED_CAST assert(false); -#else -#define ERROR_UNSUPPORTED_CAST TORCH_CHECK(false, "Unexpected scalar type"); -#endif - -// Fetch a value with dynamic type src_type from ptr, and cast it to static type dest_t. -#define FETCH_AND_CAST_CASE(type, scalartype) case ScalarType::scalartype: return static_cast_with_inter_type(*(const type *)ptr); -#define FETCH_AND_CAST_COMPLEX_CASE(type, scalartype) case ScalarType::scalartype: return static_cast_with_inter_type(std::real(*(const type *)ptr)); -template -C10_HOST_DEVICE inline dest_t fetch_and_cast(const ScalarType src_type, const void *ptr) { - switch (src_type) { - AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, FETCH_AND_CAST_CASE) -#ifndef C10_HOST_DEVICE - AT_FORALL_COMPLEX_TYPES(FETCH_AND_CAST_COMPLEX_CASE) -#endif - default:; - } - ERROR_UNSUPPORTED_CAST - return dest_t(0); // just to avoid compiler warning -} - -// Cast a value with static type src_t into dynamic dest_type, and store it to ptr. -#define CAST_AND_STORE_CASE(type, scalartype) case ScalarType::scalartype: *(type *)ptr = static_cast_with_inter_type(value); return; -template -C10_HOST_DEVICE inline void cast_and_store(const ScalarType dest_type, void *ptr, src_t value) { - switch (dest_type) { - AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) - default:; - } - ERROR_UNSUPPORTED_CAST -} - -template<> -inline void cast_and_store>(const ScalarType dest_type, void *ptr, std::complex value_) { - auto value = std::real(value_); - switch (dest_type) { - AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) - default:; - } - ERROR_UNSUPPORTED_CAST -} -template<> -inline void cast_and_store>(const ScalarType dest_type, void *ptr, std::complex value_) { - auto value = std::real(value_); - switch (dest_type) { - AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, CAST_AND_STORE_CASE) - default:; - } - ERROR_UNSUPPORTED_CAST -} - -#define DEFINE_UNCASTABLE(T, scalartype_) \ -template<> \ -inline T fetch_and_cast(const ScalarType src_type, const void *ptr) { \ - assert(ScalarType::scalartype_ == src_type); \ - return *(const T *)ptr; \ -} \ -template<> \ -inline void cast_and_store(const ScalarType dest_type, void *ptr, T value) { \ - assert(ScalarType::scalartype_ == dest_type); \ - *(T *)ptr = value; \ -} - -AT_FORALL_QINT_TYPES(DEFINE_UNCASTABLE) - -#undef FETCH_AND_CAST_CASE -#undef FETCH_AND_CAST_COMPLEX_CASE -#undef CAST_AND_STORE_CASE -#undef DEFINE_UNCASTABLE -#undef ERROR_UNSUPPORTED_CAST - -} // namespace c10 From eb551041850c59419f254d0f31711940f22671d2 Mon Sep 17 00:00:00 2001 From: svcscm Date: Tue, 29 Oct 2019 10:11:16 -0700 Subject: [PATCH 21/64] Updating submodules Summary: GitHub commits: https://github.com/pytorch/fbgemm/commit/214b370edb34437d7cd4c861bcd09775c5e330cc Test Plan: n/a Reviewed By: zpao fbshipit-source-id: aa03f9a37d316c232fdf2e4289c32ec68a22b469 --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index edee4921c4716..214b370edb344 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit edee4921c47162b1d8a1e96cc4ab0e49ae8816d2 +Subproject commit 214b370edb34437d7cd4c861bcd09775c5e330cc From 47faee2faee367896617d5e379ac57d28cc380a4 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 29 Oct 2019 11:40:04 -0700 Subject: [PATCH 22/64] Switching tests to ProfilingExecutor (rebased) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28535 Differential Revision: D18197932 Pulled By: Krovatkin fbshipit-source-id: 2639b205e899f800787ee57c157447d54e4669c3 --- aten/src/ATen/core/jit_type.h | 60 ++-- caffe2/CMakeLists.txt | 1 + test/cpp/jit/test_utils.cpp | 4 +- test/jit_utils.py | 177 ++++++---- test/test_jit.py | 310 ++++++++++++------ test/test_jit_fuser.py | 114 +++++-- tools/build_variables.py | 1 + torch/csrc/jit/autodiff.cpp | 26 +- torch/csrc/jit/fuser/executor.cpp | 4 + torch/csrc/jit/graph_executor.cpp | 40 ++- torch/csrc/jit/graph_executor.h | 1 + torch/csrc/jit/init.cpp | 13 +- torch/csrc/jit/ir.cpp | 2 + torch/csrc/jit/operator.cpp | 5 +- torch/csrc/jit/passes/alias_analysis.cpp | 1 - torch/csrc/jit/passes/bailout_graph.cpp | 33 +- torch/csrc/jit/passes/clear_undefinedness.cpp | 38 +++ torch/csrc/jit/passes/clear_undefinedness.h | 24 ++ torch/csrc/jit/passes/guard_elimination.cpp | 64 +++- .../jit/passes/specialize_autogradzero.cpp | 115 ++++--- .../jit/profiling_graph_executor_impl.cpp | 119 ++++--- torch/csrc/jit/register_prim_ops.cpp | 38 ++- 22 files changed, 821 insertions(+), 369 deletions(-) create mode 100644 torch/csrc/jit/passes/clear_undefinedness.cpp create mode 100644 torch/csrc/jit/passes/clear_undefinedness.h diff --git a/aten/src/ATen/core/jit_type.h b/aten/src/ATen/core/jit_type.h index ac14ff62f02a0..9947d6c0b178e 100644 --- a/aten/src/ATen/core/jit_type.h +++ b/aten/src/ATen/core/jit_type.h @@ -556,6 +556,12 @@ struct CAFFE2_API TensorType : public Type { return r; } + TensorTypePtr withPossiblyUndefined() { + auto r = clone(); + r->undefined_ = c10::nullopt; + return r; + } + c10::optional undefined() const { return undefined_; } static TensorTypePtr get(); @@ -563,29 +569,37 @@ struct CAFFE2_API TensorType : public Type { static const TypeKind Kind = TypeKind::TensorType; private: - TensorType(const at::Tensor &tensor) - : Type(TypeKind::TensorType), scalar_type_(tensor.scalar_type()), - device_(tensor.device()), sizes_(tensor.sizes().size()), - strides_(tensor.sizes().size()), - requires_grad_(tensor.requires_grad()), undefined_(false) { - if (!tensor.is_mkldnn() && !tensor.is_sparse()) { - sizes_ = tensor.sizes().vec(); - strides_ = tensor.strides().vec(); - } - } - TensorType(c10::optional scalar_type, - c10::optional device, const VaryingShape &sizes, - const VaryingStrides &strides, - c10::optional requires_grad, - c10::optional undefined = false) - : Type(TypeKind::TensorType), scalar_type_(scalar_type), - device_(device), sizes_(sizes), strides_(strides), - requires_grad_(requires_grad), undefined_(undefined) {} - - TensorTypePtr clone() const { - return TensorTypePtr(new TensorType(scalar_type_, device_, sizes_, - strides_, requires_grad_, - undefined_)); + TensorType(const at::Tensor& tensor) + : Type(TypeKind::TensorType), + scalar_type_(tensor.scalar_type()), + device_(tensor.device()), + sizes_(tensor.sizes().size()), + strides_(tensor.sizes().size()), + requires_grad_(tensor.requires_grad()), + undefined_(!tensor.defined()) { + if (!tensor.is_mkldnn() && !tensor.is_sparse()) { + sizes_ = tensor.sizes().vec(); + strides_ = tensor.strides().vec(); + } + } + TensorType( + c10::optional scalar_type, + c10::optional device, + const VaryingShape& sizes, + const VaryingStrides& strides, + c10::optional requires_grad, + c10::optional undefined = false) + : Type(TypeKind::TensorType), + scalar_type_(scalar_type), + device_(device), + sizes_(sizes), + strides_(strides), + requires_grad_(requires_grad), + undefined_(undefined) {} + + TensorTypePtr clone() const { + return TensorTypePtr(new TensorType( + scalar_type_, device_, sizes_, strides_, requires_grad_, undefined_)); } static std::vector contiguousStridesOf(at::IntArrayRef sizes) { diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 328e2d9a73b34..032902e78d422 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -386,6 +386,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) ${TORCH_SRC_DIR}/csrc/jit/passes/batch_mm.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/bailout_graph.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/canonicalize.cpp + ${TORCH_SRC_DIR}/csrc/jit/passes/clear_undefinedness.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/constant_propagation.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/constant_pooling.cpp ${TORCH_SRC_DIR}/csrc/jit/passes/common_subexpression_elimination.cpp diff --git a/test/cpp/jit/test_utils.cpp b/test/cpp/jit/test_utils.cpp index d2eef8a43d3e8..f86026e90b413 100644 --- a/test/cpp/jit/test_utils.cpp +++ b/test/cpp/jit/test_utils.cpp @@ -1,4 +1,6 @@ #include +#include +#include namespace torch { namespace jit { @@ -37,7 +39,7 @@ std::pair runGradient( static const auto as_tensorlist = [](const Stack& stack) { return fmap(stack, [](const IValue& i) { return i.toTensor(); }); }; - + ClearUndefinedness(grad_spec.df); Code f_code{grad_spec.f}, df_code{grad_spec.df}; InterpreterState f_interpreter{f_code}, df_interpreter{df_code}; diff --git a/test/jit_utils.py b/test/jit_utils.py index 138beee6a40a6..7f5ed32bcb7b2 100644 --- a/test/jit_utils.py +++ b/test/jit_utils.py @@ -12,6 +12,7 @@ import torch.jit.quantized import zipfile import functools +from enum import Enum # Testing utils from common_utils import TestCase, IS_WINDOWS, \ @@ -32,6 +33,24 @@ import tempfile import textwrap +IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR = False + +class ProfilingMode(Enum): + OFF = 1 + EXECUTOR = 2 + FULL = 3 + +@contextmanager +def enable_profiling_mode(flag): + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + old_prof_exec_state = torch._C._jit_set_profiling_executor(flag != ProfilingMode.OFF) + old_prof_mode_state = torch._C._jit_set_profiling_mode(flag == ProfilingMode.FULL) + try: + yield + finally: + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + torch._C._jit_set_profiling_executor(old_prof_exec_state) + torch._C._jit_set_profiling_mode(old_prof_mode_state) def execWrapper(code, glob, loc): if PY2: @@ -39,7 +58,6 @@ def execWrapper(code, glob, loc): else: exec(code, glob, loc) - def do_input_map(fn, input): return _nested_map(lambda t: isinstance(t, torch.Tensor), fn)(input) @@ -304,24 +322,34 @@ def get_frame_vars(self, frames_up): return defined_vars def checkScriptRaisesRegex(self, script, inputs, exception, regex, - outputs=None, capture_output=False): + outputs=None, capture_output=False, profiling=ProfilingMode.FULL): """ Checks that a given function will throw the correct exception, when executed with normal python, the string frontend, and the AST frontend """ - # normal python - with self.assertRaisesRegex(exception, regex): - script(*inputs) - # string frontend - with self.assertRaisesRegex(exception, regex): - source = textwrap.dedent(inspect.getsource(script)) - cu = torch.jit.CompilationUnit(source) - ge = getattr(cu, script.__name__) - ge(*inputs) - # python AST frontend - with self.assertRaisesRegex(exception, regex): - ge = torch.jit.script(script) - ge(*inputs) + + with enable_profiling_mode(profiling): + # normal python + with self.assertRaisesRegex(exception, regex): + script(*inputs) + # string frontend + with self.assertRaisesRegex(exception, regex): + source = textwrap.dedent(inspect.getsource(script)) + cu = torch.jit.CompilationUnit(source) + ge = getattr(cu, script.__name__) + # profiling run + with self.assertRaisesRegex(exception, regex): + ge(*inputs) + # optimized run + ge(*inputs) + # python AST frontend + with self.assertRaisesRegex(exception, regex): + ge = torch.jit.script(script) + # profiling run + with self.assertRaisesRegex(exception, regex): + ge(*inputs) + # optimized run + ge(*inputs) def checkScript(self, script, @@ -330,59 +358,71 @@ def checkScript(self, optimize=True, inputs_requires_grad=False, capture_output=False, - frames_up=1): + frames_up=1, + profiling=ProfilingMode.FULL): with torch.jit.optimized_execution(optimize): - if isinstance(script, str): - # Compile the string to a Script function - cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) - - # Execute the Python function so we can run it later and get its - # outputs - frame = self.get_frame_vars(frames_up) - the_locals = {} - execWrapper(script, glob=frame, loc=the_locals) - frame.update(the_locals) - - python_fn = frame[name] - scripted_fn = getattr(cu, name) - else: - - # Check the string frontend first - source = textwrap.dedent(inspect.getsource(script)) - self.checkScript( - source, - inputs, - script.__name__, - capture_output, - frames_up=2) - - # Continue checking the Python frontend - scripted_fn = torch.jit.script(script, _frames_up=1) - python_fn = script - - if inputs_requires_grad: - recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs) - else: - recording_inputs = inputs + with enable_profiling_mode(profiling): + if isinstance(script, str): + # Compile the string to a Script function + # with enable_profiling_mode(profiling): + cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) + + # Execute the Python function so we can run it later and get its + # outputs + + frame = self.get_frame_vars(frames_up) + the_locals = {} + execWrapper(script, glob=frame, loc=the_locals) + frame.update(the_locals) + + python_fn = frame[name] + scripted_fn = getattr(cu, name) + else: - if capture_output: - with self.capture_stdout() as script_stdout: + # Check the string frontend first + source = textwrap.dedent(inspect.getsource(script)) + self.checkScript( + source, + inputs, + script.__name__, + capture_output, + profiling=profiling, + frames_up=2) + + # Continue checking the Python frontend + scripted_fn = torch.jit.script(script, _frames_up=1) + python_fn = script + + if inputs_requires_grad: + recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs) + else: + recording_inputs = inputs + + if capture_output: + with self.capture_stdout() as script_stdout: + script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout() as opt_script_stdout: + opt_script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout() as _python_stdout: + python_outputs = python_fn(*inputs) + if not IS_WINDOWS: + self.assertExpected(script_stdout[0], subname='stdout') + self.assertEqual(python_outputs, opt_script_outputs) + else: + # profiling run script_outputs = scripted_fn(*recording_inputs) - with self.capture_stdout() as _python_stdout: + # optimized run + opt_script_outputs = scripted_fn(*recording_inputs) python_outputs = python_fn(*inputs) - if not IS_WINDOWS: - self.assertExpected(script_stdout[0], subname='stdout') - else: - script_outputs = scripted_fn(*recording_inputs) - python_outputs = python_fn(*inputs) - self.assertEqual(python_outputs, script_outputs) - - return scripted_fn + self.assertEqual(python_outputs, script_outputs) + self.assertEqual(script_outputs, opt_script_outputs) + return scripted_fn def checkTrace(self, func, reference_tensors, input_tensors=None, drop=None, allow_unused=False, verbose=False, inputs_require_grads=True, check_tolerance=1e-5, export_import=True, _force_outplace=False): + # TODO: check gradients for parameters, not just inputs def allSum(vs): # drop allows us to remove some values from ever being used @@ -413,8 +453,11 @@ def input_reduce(input, fn, acc): else: recording_inputs = reference_tensors + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `checkTrace` already does all the checks + # against python function ge = torch.jit.trace(func, input_tensors, check_tolerance=check_tolerance, - _force_outplace=_force_outplace) + _force_outplace=_force_outplace, check_trace=False) if export_import: ge = self.getExportImportCopy(ge) @@ -427,7 +470,6 @@ def input_reduce(input, fn, acc): outputs_ge = ge(*nograd_inputs) self.assertEqual(outputs, outputs_ge) - # test single grad case outputs = func(*recording_inputs) if inputs_require_grads: grads = torch.autograd.grad(allSum(outputs), flattened_recording_inputs, @@ -441,8 +483,11 @@ def input_reduce(input, fn, acc): if inputs_require_grads: self.assertEqual(grads, grads_ge) - # test the grad grad case + self.assertEqual(outputs, outputs_ge) + if inputs_require_grads: + self.assertEqual(grads, grads_ge) + # test the grad grad case outputs = func(*recording_inputs) l1 = allSum(outputs) if inputs_require_grads: @@ -514,14 +559,6 @@ def checkModule(self, nn_module, args): return sm -@contextmanager -def enable_profiling_mode(): - torch._C._jit_set_profiling_mode(True) - try: - yield - finally: - torch._C._jit_set_profiling_mode(False) - @contextmanager def inline_everything_mode(should_inline): old = torch._C._jit_get_inline_everything_mode() diff --git a/test/test_jit.py b/test/test_jit.py index b555f1877b78e..b5d1e5fbcdbc3 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -41,9 +41,9 @@ skipIfRocm, skipIfNoLapack, suppress_warnings, IS_SANDCASTLE, \ freeze_rng_state, set_rng_seed, slowTest, TemporaryFileName, skipIfCompiledWithoutNumpy from jit_utils import JitTestCase, enable_cpu_fuser, disable_autodiff_subgraph_inlining, \ - _trace, enable_cpu_fuser_if, enable_profiling_mode, do_input_map, \ + _trace, enable_cpu_fuser_if, enable_profiling_mode, ProfilingMode, do_input_map, \ execWrapper, _inline_everything, _tmp_donotuse_dont_inline_everything, \ - get_forward, get_forward_graph, get_module_method + get_forward, get_forward_graph, get_module_method, IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR from common_nn import module_tests, new_module_tests, criterion_tests from common_methods_invocations import method_tests as autograd_method_tests from common_methods_invocations import create_input, unpack_variables, \ @@ -102,6 +102,63 @@ def LSTMCellF(input, hx, cx, *params): return LSTMCell(input, (hx, cx), *params) +def doAutodiffCheck(testname): + + if not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + return True + + + # these tests are disabled because BailOut nodes + # inserted by ProfilingExecutor interfere with + # subgraph slicing of Differentiable Graphs + test_exceptions = [ + # functional + 'test_nn_dropout', + 'test_nn_log_softmax', + 'test_nn_relu', + 'test_nn_softmax', + 'test_nn_threshold', + 'test_nn_lp_pool2d', + 'test_nn_lp_pool1d', + 'test_nn_gumbel_softmax_hard', + 'test_nn_gumbel_softmax', + 'test_nn_multilabel_soft_margin_loss', + 'test_nn_batch_norm', + # AutogradJitGenerated + 'test___rdiv___constant', + 'test___rdiv___scalar_constant', + ] + + if testname in test_exceptions: + return False + return True + +func_call = torch._C.ScriptFunction.__call__ +meth_call = torch._C.ScriptMethod.__call__ + +def prof_callable(callable, *args, **kwargs): + if 'profile_and_replay' in kwargs: + del kwargs['profile_and_replay'] + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + with enable_profiling_mode(ProfilingMode.FULL): + callable(*args, **kwargs) + return callable(*args, **kwargs) + + return callable(*args, **kwargs) + +def prof_func_call(*args, **kwargs): + return prof_callable(func_call, *args, **kwargs) + +def prof_meth_call(*args, **kwargs): + return prof_callable(meth_call, *args, **kwargs) + +torch._C.ScriptFunction.__call__ = prof_func_call +torch._C.ScriptMethod.__call__ = prof_meth_call + +if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + # enable profiling graph executor for all tests in this file by default + torch._C._jit_set_profiling_executor(True) + def LSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None): hx, cx = hidden @@ -204,13 +261,16 @@ def get_execution_plan(graph_executor_state): return execution_plans[0] -def get_grad_executor(plan_state, diff_graph_idx=None): +def get_grad_executor(plan_state, diff_graph_idx=None, skip_check=False): if diff_graph_idx is None: nodes = list(plan_state.graph.nodes()) - if len(nodes) == 1 or (len(nodes) == 2 and nodes[1].kind() == "prim::TupleConstruct"): - pass - else: - raise RuntimeError("Can't get a grad_executor for a non-differentiable graph") + + if not skip_check: + nodes = list(filter(lambda n : n.kind() != "prim::BailOut" and n.kind() != "prim::BailoutTemplate", nodes)) + if len(nodes) == 1 or (len(nodes) == 2 and nodes[1].kind() == "prim::TupleConstruct"): + pass + else: + raise RuntimeError("Can't get a grad_executor for a non-differentiable graph") grad_executors = list(plan_state.code.grad_executor_states()) return grad_executors[diff_graph_idx or 0] @@ -224,10 +284,10 @@ def all_backward_graphs(script_module, diff_graph_idx=None): return [p.graph.copy() for p in bwd_plans] -def backward_graph(script_module, diff_graph_idx=None): +def backward_graph(script_module, diff_graph_idx=None, skip_check=False): ge_state = script_module.get_debug_state() fwd_plan = get_execution_plan(ge_state) - grad_executor_state = get_grad_executor(fwd_plan, diff_graph_idx=diff_graph_idx) + grad_executor_state = get_grad_executor(fwd_plan, diff_graph_idx=diff_graph_idx, skip_check=skip_check) bwd_plan = get_execution_plan(grad_executor_state) # Running JIT passes requires that we own the graph (with a shared_ptr). # The debug state struct does not own its graph so we make a copy of it. @@ -502,9 +562,14 @@ def f(x, y): def test_peephole_optimize_shape_ops(self): def test_input(func, input, result): - self.assertEqual(func(input), result) + # if result == 2 we will trigger a bailout and + # the unprofiled graph should return the correct result + self.assertEqual(func(input, profile_and_replay=True), result) gre = func.graph_for(input) - FileCheck().check_not("prim::If").run(gre) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + FileCheck().check("prim::Constant").check_next("prim::BailoutTemplate").run(gre) + else: + FileCheck().check_not("prim::If").run(gre) def test_dim(): @torch.jit.script @@ -1970,6 +2035,7 @@ def forward(self, d): inputs = {'x': torch.rand(3, 4), 'y': torch.rand(3, 4)} module = torch.jit.trace(Test(), inputs) + FileCheck().check('aten::values').check('prim::ListUnpack').run(str(module.graph)) def test_input_dict_flattens_recursive(self): @@ -2154,20 +2220,21 @@ def test_dropout_cuda(self): # which is not included in TestJitGeneratedFunctional x = torch.ones(4, 4).cuda().requires_grad_() - @torch.jit.script - def func(x): - return torch.nn.functional.dropout(x) + with enable_profiling_mode(ProfilingMode.FULL): + @torch.jit.script + def func(x): + return torch.nn.functional.dropout(x) - with freeze_rng_state(): - out_ref = torch.nn.functional.dropout(x) - grad_ref = torch.autograd.grad(out_ref.sum(), x) + with freeze_rng_state(): + out_ref = torch.nn.functional.dropout(x) + grad_ref = torch.autograd.grad(out_ref.sum(), x) - with freeze_rng_state(): - out = func(x) - grad = torch.autograd.grad(out.sum(), x) + with freeze_rng_state(): + out = func(x) + grad = torch.autograd.grad(out.sum(), x) - self.assertEqual(out, out_ref) - self.assertEqual(grad, grad_ref) + self.assertEqual(out, out_ref) + self.assertEqual(grad, grad_ref) def test_conv(self): x = torch.ones(20, 16, 50, 40) @@ -2297,8 +2364,7 @@ def rand(*args): self.checkTrace(lambda a, b: a * b + b, [rand(1), rand(1)], [rand(2, 3), rand(2, 3)]) # trivial identity - self.checkTrace(lambda a, b: ( - b, a), [rand(1), rand(1)]) + self.checkTrace(lambda a, b: (b, a), [rand(1), rand(1)]) def foo(a): t = a * a @@ -2319,7 +2385,8 @@ def test_ge_unoptimized(self): @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser support for Sandcastle") @enable_cpu_fuser def test_ge_optimized(self): - self.run_ge_tests(True, False) + with enable_profiling_mode(ProfilingMode.FULL): + self.run_ge_tests(True, False) @unittest.skipIf(not RUN_CUDA, "requires CUDA") def test_ge_cuda(self): @@ -5269,6 +5336,7 @@ def func(alpha, beta, x, y): # NOTE: cannot optimize yet because broadcasts are not inserted before the fuser runs self.checkScript(func, [alpha, beta, x, y], optimize=False) + @unittest.skipIf(not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "skip if profiling isn't enabled") def test_profiling_graph_executor(self): @torch.jit.script def def_in_one_branch(x, z): @@ -5281,12 +5349,15 @@ def def_in_one_branch(x, z): a = torch.rand(2, 3) - with enable_profiling_mode(): - # the first call is profiled - profiled_graph_str = str(def_in_one_branch.graph_for(a, False)) + with enable_profiling_mode(ProfilingMode.FULL): + # check prim::profile are inserted + profiled_graph_str = str(def_in_one_branch.graph_for(a, True)) FileCheck().check_count("prim::profile", 4).run(profiled_graph_str) - # the second call is optimized + # this call is optimized for + # the given shape of (2, 3) def_in_one_branch(a, False) + # change shape to (3) + # so we go down a bailout path a = torch.ones(3) # check prim::BailOuts are inserted bailout_graph_str = str(def_in_one_branch.graph_for(a, True)) @@ -5296,7 +5367,6 @@ def def_in_one_branch(x, z): # this triggers 2 bailouts self.assertEqual(def_in_one_branch(a, True), 3.0) - def test_resize_input_ops(self): # resize_ and resize_as resize the input tensor. because our shape analysis # is flow invariant, we set any Tensor that can alias a resized Tensor @@ -5389,7 +5459,7 @@ def test(x, y, z): # and the output of the node conservatively setting grad to true inps = (torch.tensor(1.0, requires_grad=True), torch.tensor(1), 10) - test(*inps) + test(*inps, profile_and_replay=True) graph = test.graph_for(*inps) loop = graph.findNode("prim::Loop") @@ -5398,8 +5468,13 @@ def test(x, y, z): loop_outputs = list(loop_body.outputs()) self.assertTrue(loop_inputs[1].requires_grad()) - self.assertFalse(loop_outputs[1].requires_grad()) - self.assertTrue(loop.output().requires_grad()) + + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + bailouts_in_outer_block = graph.findAllNodes("prim::BailOut", False) + self.assertFalse(bailouts_in_outer_block[1].output().requires_grad()) + else: + self.assertTrue(loop.output().requires_grad()) + self.assertFalse(loop_outputs[1].requires_grad()) def test_view_shape_prop(self): cu = torch.jit.CompilationUnit(''' @@ -6041,6 +6116,7 @@ def test(x, y): with self.assertRaisesRegex(Exception, ""): test(1, None) + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the current version of Profiler doesn't profile/specialize Optionals") def test_optional_tensor(self): @torch.jit.script def fn(x, y): @@ -6081,6 +6157,7 @@ def fn(x, y, b): g = torch.jit.last_executed_optimized_graph() self.assertEqual(next(g.outputs()).type().str(), "Tensor") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the current version of Profiler doesn't profile/specialize Optionals") def test_optional_list(self): @torch.jit.script def fn(x, y): @@ -6784,17 +6861,23 @@ def func(): ''') ops = ['tensor', 'as_tensor'] inputs = ['[1]', '[False]', '[2.5]', '0.5', '1', 'False', '[[1]]'] - expected_shape = ["Long(*)", ("Bool(*)"), "Double(*)", "Double()", "Long()", "Bool()", "Long(*, *)"] + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + expected_shape = ["Long(1)", "Bool(1)", "Double(1)", "Double()", "Long()", "Bool()", "Long(1, 1)"] + else: + expected_shape = ["Long(*)", ("Bool(*)"), "Double(*)", "Double()", "Long()", "Bool()", "Long(*, *)"] for op in ops: for inp, expect in zip(inputs, expected_shape): code = tensor_template.format(tensor_op=op, input=inp) scope = {} exec(code, globals(), scope) - self.checkScript(code, ()) - cu = torch.jit.CompilationUnit(code) - torch._C._jit_pass_complete_shape_analysis(cu.func.graph, (), False) - FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(cu.func.graph) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + fn = self.checkScript(code, ()) + FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(fn.graph_for()) + else: + cu = torch.jit.CompilationUnit(code) + torch._C._jit_pass_complete_shape_analysis(cu.func.graph, (), False) + FileCheck().check(expect).check("aten::{tensor_op}".format(tensor_op=op)).run(cu.func.graph) @torch.jit.script def test_dtype(inp_dtype): @@ -6802,17 +6885,26 @@ def test_dtype(inp_dtype): a = torch.tensor(1.0, dtype=torch.float, requires_grad=True) return a, torch.tensor(1.0, dtype=inp_dtype) # noqa T484 - g = test_dtype.graph_for(5) - # first should have type set second should not - FileCheck().check("Float() = aten::tensor").check("Tensor = aten::tensor").run(g) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + g = test_dtype.graph_for(5, profile_and_replay=True) + # both should have completed shapes + FileCheck().check("Tensor = aten::tensor").check("Float() = prim::BailOut").check("Tensor = aten::tensor").check("Half() = prim::BailOut").run(g) + else: + g = test_dtype.graph_for(5) + # first should have type set second should not + FileCheck().check("Float() = aten::tensor").check("Tensor = aten::tensor").run(g) @torch.jit.script def test_as_tensor_tensor_input(input): a = torch.as_tensor(input, dtype=input.dtype) return a, torch.as_tensor(input, dtype=torch.float) - g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4)) - FileCheck().check("Tensor = aten::as_tensor").check("Float(*, *) = aten::as_tensor").run(g) + if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4), profile_and_replay=True) + FileCheck().check("Tensor = aten::as_tensor").check("Float(3, 4) = prim::BailOut").check("Tensor = aten::as_tensor").check("Float(3, 4) = prim::BailOut").run(g) + else: + g = test_as_tensor_tensor_input.graph_for(torch.ones(3, 4)) + FileCheck().check("Tensor = aten::as_tensor").check("Float(*, *) = aten::as_tensor").run(g) def test_tensor_requires_grad(self): @@ -6909,7 +7001,7 @@ def s(t, to_str, non_blocking=None, device=None, cuda=None): code = template.format(to_str=to_str, device=device, non_blocking=non_blocking, cuda=cuda) scope = {} cu = torch.jit.CompilationUnit(code) - return cu.func(t) + return cu.func(t, profile_and_replay=True) def test_copy_behavior(t, non_blocking=False): self.assertIs(t, s(t, 't.to(t, non_blocking=non_blocking)', non_blocking)) @@ -10797,6 +10889,7 @@ def foo(): self.checkScript(foo, ()) + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the original version of test_rand") def test_rand(self): def test_rand(): a = torch.rand([3, 4]) @@ -10819,6 +10912,28 @@ def randint(): # and shape analysis dtype is the same. FileCheck().check("Double(*, *)").check_not("Float(*, *)").run(randint.graph_for()) + @unittest.skipIf(not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "the original version of test_rand") + def test_rand_profiling(self): + def test_rand(): + a = torch.rand([3, 4]) + return a + 1.0 - a + + fn = self.checkScript(test_rand, ()) + out = fn() + self.assertEqual(out.dtype, torch.double) + # Testing shape analysis correctly setting type + FileCheck().check("Double(3, 4)").check_not("Float(3, 4)").run(fn.graph_for()) + + @torch.jit.script + def randint(): + return torch.randint(0, 5, [1, 2]) + + out = randint(profile_and_replay=True) + self.assertEqual(out.dtype, torch.double) + # although the type should be int here, testing that the runtime dtype + # and shape analysis dtype is the same. + FileCheck().check("Double(1, 2)").check_not("Float(1, 2)").run(randint.graph_for()) + def test_erase_number_types(self): def func(a): b = 7 + 1 + 3 @@ -10847,8 +10962,9 @@ def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh): fw_graph = slstm.graph_for(*inputs) bw_graph = backward_graph(slstm, diff_graph_idx=0) - self.assertTrue('prim::MMBatchSide' in str(fw_graph)) - self.assertTrue('prim::MMTreeReduce' in str(bw_graph)) + if not IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + self.assertTrue('prim::MMBatchSide' in str(fw_graph)) + self.assertTrue('prim::MMTreeReduce' in str(bw_graph)) sout = slstm(*inputs) out = lstm(*inputs) @@ -13693,7 +13809,7 @@ def fn(): x.add_(torch.ones(2, 3)) return x_view - self.checkScript(fn, ()) + self.checkScript(fn, (), profiling=ProfilingMode.EXECUTOR) def test_cpp_function_tensor_str(self): x = torch.randn(2, 2) @@ -16322,7 +16438,10 @@ def new_fn(*tensors_): def create_traced_fn(self, fn): def traced_fn(*inputs, **kwargs): fn_tensors, inputs_tensors = partial_apply_nontensors(fn, inputs, **kwargs) - traced = torch.jit.trace(fn_tensors, inputs_tensors) + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `check_against_reference` already does all the checks + # against python function + traced = torch.jit.trace(fn_tensors, inputs_tensors, check_trace=False) self.assertExportImport(traced.graph, inputs_tensors) output = traced(*inputs_tensors) traced_fn.last_graph = traced.graph_for(*inputs_tensors) @@ -16393,7 +16512,6 @@ def create_script_fn(self, method_name, func_type, output_process_fn): def script_fn(*args, **kwargs): formals, tensors, actuals = get_script_args(args) call = get_call(method_name, func_type, actuals, kwargs) - script = script_template.format(', '.join(formals), call) CU = torch.jit.CompilationUnit(script) @@ -16443,7 +16561,8 @@ def clone_inputs(requires_grad): # test no gradients case outputs = self.runAndSaveRNG(reference_func, nograd_inputs, kwargs) - outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs) + with enable_profiling_mode(ProfilingMode.FULL): + outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs) self.assertEqual(outputs, outputs_test) if check_types: @@ -16453,43 +16572,42 @@ def clone_inputs(requires_grad): # skip grad tests return - # test single grad case - outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) - grads = torch.autograd.grad(allSum(outputs), recording_tensors, - allow_unused=allow_unused) - - outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) - grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors, - allow_unused=allow_unused) - self.assertEqual(outputs, outputs_test) - self.assertEqual(grads, grads_test) - - # test the grad grad case - if self._testMethodName in nn_functional_single_grad: - return - - outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) - l1 = allSum(outputs) - grads = torch.autograd.grad(l1, recording_tensors, create_graph=True, - allow_unused=allow_unused) - l2 = (allSum(grads) * l1) - grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused) - - recording_inputs, recording_tensors = clone_inputs(True) - - outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) - l1_test = allSum(outputs_test) - grads_test = torch.autograd.grad( - l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused) - l2_test = (allSum(grads_test) * l1_test) - grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused) + with enable_profiling_mode(ProfilingMode.FULL): + # test single grad case + outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) + grads = torch.autograd.grad(allSum(outputs), recording_tensors, + allow_unused=allow_unused) + outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) + grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors, + allow_unused=allow_unused) + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + # test the grad grad case + if self._testMethodName in nn_functional_single_grad: + return - self.assertEqual(outputs, outputs_test) - self.assertEqual(grads, grads_test) - for g2, g2_test in zip(grads2, grads2_test): - if g2 is None and g2_test is None: - continue - self.assertTrue(torch.allclose(g2, g2_test, atol=5e-4, rtol=1e-4)) + outputs = self.runAndSaveRNG(reference_func, recording_inputs, kwargs) + l1 = allSum(outputs) + grads = torch.autograd.grad(l1, recording_tensors, create_graph=True, + allow_unused=allow_unused) + + l2 = (allSum(grads) * l1) + grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused) + recording_inputs, recording_tensors = clone_inputs(True) + outputs_test = self.runAndSaveRNG(func, recording_inputs, kwargs) + l1_test = allSum(outputs_test) + grads_test = torch.autograd.grad( + l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused) + + l2_test = (allSum(grads_test) * l1_test) + grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused) + + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + for g2, g2_test in zip(grads2, grads2_test): + if g2 is None and g2_test is None: + continue + self.assertTrue(torch.allclose(g2, g2_test, atol=5e-4, rtol=1e-4)) # NB: torch.jit.script, when used as a function, uses the current scope @@ -17116,7 +17234,9 @@ def fn(*inputs, **kwargs): if IS_SANDCASTLE: autodiff_nodes = autodiff_nodes + fusible_nodes fusible_nodes = [] - self.assertAutodiffNode(traced_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) + + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(traced_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) if not is_magic_method and test_name not in EXCLUDE_SCRIPT: script_fn = create_script_fn(self, name, 'method', output_process_fn) @@ -17127,10 +17247,11 @@ def fn(*inputs, **kwargs): if IS_SANDCASTLE: autodiff_nodes = autodiff_nodes + fusible_nodes fusible_nodes = [] - self.assertAutodiffNode(script_fn.last_graph, - should_autodiff_node and test_name not in EXCLUDE_SCRIPT_AD_CHECK, - autodiff_nodes, - fusible_nodes) + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(script_fn.last_graph, + should_autodiff_node and test_name not in EXCLUDE_SCRIPT_AD_CHECK, + autodiff_nodes, + fusible_nodes) # functional interface tests if hasattr(torch, name) and name not in EXCLUDE_FUNCTIONAL: @@ -17204,8 +17325,8 @@ def fn(*inputs, **kwargs): f_args_variable = (self_variable,) + args_variable f_args_tensor = (self_tensor,) + args_tensor - should_autodiff_node, autodiff_nodes, fusible_nodes = normalize_check_ad(check_ad, name) + if test_name not in EXCLUDE_SCRIPT: def run_test(): # XXX: this test should always run with disable_autodiff_subgraph_inlining(True), @@ -17214,7 +17335,8 @@ def run_test(): script_fn = create_script_fn(self, name, 'nn_functional', output_process_fn) check_against_reference(self, script_fn, fn, f_args_variable, kwargs_variable, no_grad=no_grad) # For tests we disabled AD subgraph inlining, make sure it's not falling back to autograd - self.assertAutodiffNode(script_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) + if (doAutodiffCheck(test_name)): + self.assertAutodiffNode(script_fn.last_graph, should_autodiff_node, autodiff_nodes, fusible_nodes) if test_name in EXCLUDE_PYTHON_PRINT: with torch.jit._disable_emit_hooks(): @@ -18221,7 +18343,7 @@ def use_foo(foo, foo2, tup): input = (f, f2, (f, f3)) sfoo = self.checkScript(use_foo, input) graphstr = str(sfoo.graph_for(*input)) - FileCheck().check_count("Double(*, *) = prim::GetAttr", 4).run(graphstr) + FileCheck().check_count("prim::GetAttr", 4).run(graphstr) def test_class_sorting(self): global Foo # see [local resolution in python] diff --git a/test/test_jit_fuser.py b/test/test_jit_fuser.py index 6529509059f6b..6c9db9b5b736b 100644 --- a/test/test_jit_fuser.py +++ b/test/test_jit_fuser.py @@ -16,13 +16,49 @@ from test_jit import JitTestCase, enable_cpu_fuser, RUN_CUDA, RUN_CUDA_HALF, RUN_CUDA_MULTI_GPU, \ backward_graph, all_backward_graphs, get_lstm_inputs, get_milstm_inputs, \ LSTMCellC, LSTMCellF, LSTMCellS, MiLSTMCell, _inline_everything +from jit_utils import enable_profiling_mode, ProfilingMode, IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR + +if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: + torch._C._jit_set_profiling_executor(True) + torch._C._jit_set_profiling_mode(True) + + +def strip_profiling_nodes(nodes): + profiling_opcodes = set(['prim::BailoutTemplate', 'prim::BailOut']) + return [n for n in nodes if n.kind() not in profiling_opcodes] + + +def warmup_backward(f, *args): + profiling_count = 2 + results = [] + for i in range(profiling_count): + if len(args) > 0: + r = torch.autograd.grad(f, *args) + results.append(r) + else: + f.backward(retain_graph=True) + + return results + + +def warmup_forward(f, *args): + profiling_count = 2 + for i in range(profiling_count): + results = f(*args) + + return results class TestFuser(JitTestCase): def assertAllFused(self, graph, except_for=()): - if [n.kind() for n in graph.nodes()] == ['prim::DifferentiableGraph']: - graph = next(graph.nodes()).g('Subgraph') - allowed_nodes = {'prim::Constant', 'prim::FusionGroup', 'prim::TupleConstruct'} | set(except_for) + + diff_graphs = [n for n in graph.nodes() if n.kind() == 'prim::DifferentiableGraph'] + if len(diff_graphs) > 0: + self.assertEqual(len(diff_graphs), 1) + graph = diff_graphs[0].g('Subgraph') + + allowed_nodes = {'prim::Constant', 'prim::FusionGroup', 'prim::BailoutTemplate', + 'prim::BailOut', 'prim::TupleConstruct'} | set(except_for) self.assertTrue(all(node.kind() in allowed_nodes for node in graph.nodes()), 'got {}'.format(graph)) self.assertTrue([node.kind() for node in graph.nodes()].count('prim::FusionGroup') == 1) @@ -87,6 +123,7 @@ def scaleshift(x, scale, shift): @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @unittest.skipIf(not RUN_CUDA_HALF, "no half support") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "no half support with profiling on") def test_cuda_half(self): x = torch.randn(4, 4, dtype=torch.half, device='cuda') y = torch.randn(4, 4, dtype=torch.half, device='cuda') @@ -200,7 +237,7 @@ def f(x, y): ge = self.checkTrace(f, (x, y)) graph = ge.graph_for(x, y) - FileCheck().check("broadcast_tensors").check('with prim::FusionGroup_0') \ + FileCheck().check("broadcast_tensors").check('with prim::FusionGroup_') \ .check_count('ConstantChunk', 2, exactly=True).run(str(graph)) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @@ -266,11 +303,11 @@ def funcOptMax(a, b): funcs = (func2, funcInf, funcOptMin, funcOptMax) for f, inputs in product(funcs, [[a, b], [a, nan]]): inp1, inp2 = inputs - s = self.checkScript(f, (inp1, inp2)) + s = self.checkScript(f, (inp1, inp2), profiling=ProfilingMode.FULL) self.assertAllFused(s.graph_for(inp1, inp2), except_for={'aten::size', 'aten::_size_if_not_equal'}) - c = s(inp1, inp2) - c.sum().backward() + with enable_profiling_mode(ProfilingMode.FULL): + warmup_backward(c.sum()) graph = backward_graph(s) self.assertAllFused(graph, except_for={'aten::Float'}) @@ -283,8 +320,10 @@ def func(x): a = torch.randn(4, 4, dtype=torch.float, device='cuda', requires_grad=True) s = torch.jit.script(func, (a,)) c = s(a) - c.sum().backward() - graph = backward_graph(s) + c = s(a) + warmup_backward(c.sum()) + # skip_check to skip extra bailout nodes in between + graph = backward_graph(s, skip_check=True) self.assertAllFused(graph, except_for={'aten::div', 'prim::Constant'}) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @@ -422,6 +461,7 @@ def test_exp_cuda(self): self.assertAllFused(ge.graph_for(x, y)) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "broken with profiling on") @_inline_everything def test_fuse_decompose_normalization(self): class ResLike(torch.jit.ScriptModule): @@ -497,12 +537,22 @@ def fn_test_scalar_arg(x, p): scripted = torch.jit.script(fn_test_scalar_arg, (x, p)) self.assertEqual(fn_test_scalar_arg(x, p), scripted(x, p)) self.assertAllFused(scripted.graph_for(x, p)) + x.requires_grad_(True) + + # use another function otherwise we will bailout + # and won't be able to do fused checks + def fn_test_scalar_arg_requires_grad(x, p): + # type: (Tensor, float) -> Tensor + return p * (x * x + x) + + scripted = torch.jit.script(fn_test_scalar_arg_requires_grad, (x, p)) out = scripted(x, p) self.assertAllFused(scripted.graph_for(x, p), except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") + @unittest.skipIf(IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR, "broken with profiling on") @enable_cpu_fuser def test_fuser_deduplication(self): # See that fusion kernel outputs are deduplicated when removing _grad_sum_to_size in the fuser's compilation @@ -513,14 +563,16 @@ def f(x, y): b = torch.randn(5, 5, requires_grad=True) a = torch.randn(5, 5, requires_grad=True) s = self.checkScript(f, (a, b)) - self.assertAllFused(s.graph_for(a, b), except_for={'aten::size', 'aten::_size_if_not_equal', 'prim::BroadcastSizes'}) + self.assertAllFused(s.graph_for(a, b), except_for={ + 'aten::size', 'aten::_size_if_not_equal', 'prim::BroadcastSizes'}) c = s(a, b) - ga, gb = torch.autograd.grad(c.sum(), [a, b]) + results = warmup_backward(c.sum(), [a, b]) + ga2, gb2 = results.pop() graph = backward_graph(s) self.assertAllFused(graph) # check that a, b share storage, i.e. were generated as a single output in the fuser - self.assertEqual(ga.data_ptr(), gb.data_ptr()) + self.assertEqual(ga2.data_ptr(), gb2.data_ptr()) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") @enable_cpu_fuser @@ -559,10 +611,11 @@ def iou(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2): self.assertAllFused(s.graph_for(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2), except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) - c = s(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2) - torch.autograd.grad(c.sum(), [b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2]) - graph = backward_graph(s) - self.assertAllFused(graph, except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) + with enable_profiling_mode(True): + c = s(b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2) + warmup_backward(c.sum(), [b1x1, b1y1, b1x2, b1y2, b2x1, b2y1, b2x2, b2y2]) + graph = backward_graph(s) + self.assertAllFused(graph, except_for={'aten::size', 'prim::BroadcastSizes', 'aten::_size_if_not_equal'}) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") @unittest.skipIf(not RUN_CUDA_MULTI_GPU, "needs non-zero device") @@ -630,17 +683,19 @@ def doit(x, y): def test_lstm_cuda(self): inputs = get_lstm_inputs('cuda', training=True) module = self.checkScript(LSTMCellS, inputs) + return forward_graph = module.graph_for(*inputs) self.assertGraphContainsExactly( forward_graph, 'prim::FusionGroup', 1, consider_subgraphs=True) - self.assertTrue(len(list(forward_graph.nodes())) == 2) + self.assertTrue(len(strip_profiling_nodes(forward_graph.nodes())) == 2) # Everything is differentiable but TupleConstruct return FileCheck().check("DifferentiableGraph").check_next("TupleConstruct") \ .check_next("return").run(str(forward_graph)) - hy, cy = module(*inputs) - (hy + cy).sum().backward() - backward = backward_graph(module) + with enable_profiling_mode(True): + hy, cy = module(*inputs) + warmup_backward((hy + cy).sum()) + backward = backward_graph(module) self.assertAllFused(backward, except_for=("aten::t", "aten::mm", "aten::_grad_sum_to_size")) @@ -679,9 +734,10 @@ def test_lstm_traced_cuda(self): inputs = get_lstm_inputs('cuda') ge = self.checkTrace(LSTMCellF, inputs) graph = ge.graph_for(*inputs) - FileCheck().check_not("Chunk").check_not("aten::add").check_not("aten::sigmoid") \ + # .check_not("aten::add") don't get pulled into FusionGroup because of BailOuts + FileCheck().check_not("Chunk").check_not("aten::sigmoid") \ .check_not("aten::tanh").check("FusionGroup").check_next("TupleConstruct") \ - .check_next("return").check_not("FusionGroup_1").run(str(graph)) + .check_next("return").check_not("FusionGroup_2").run(str(graph)) @unittest.skipIf(IS_SANDCASTLE, "NYI: fuser CPU support for Sandcastle") @unittest.skip("Test is flaky, see https://github.com/pytorch/pytorch/issues/8746") @@ -711,7 +767,7 @@ def test_milstm_cuda(self): FileCheck().check("DifferentiableGraph").check_next("TupleConstruct") \ .check_next("return").check("FusionGroup").run(str(forward_graph)) hy, cy = module(*inputs) - (hy + cy).sum().backward() + warmup_backward((hy + cy).sum()) @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA") def test_rand_cuda(self): @@ -758,6 +814,7 @@ def fn_test_erf(x): ge = self.checkTrace(fn_test_erf, (x,)) self.assertAllFused(ge.graph_for(x)) x.requires_grad_(True) + ge = self.checkTrace(fn_test_erf, (x,)) self.assertAllFused(ge.graph_for(x), except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @@ -856,7 +913,7 @@ def my_broadcasted_cell(a, b, c): s1 = torch.randn(5, 1, requires_grad=True, device='cuda') s2 = torch.randn(5, 5, requires_grad=True, device='cuda') - module = self.checkScript(my_broadcasted_cell, (s1, s1, s1)) + module = self.checkScript(my_broadcasted_cell, (s1, s1, s1), profiling=ProfilingMode.FULL) forward_graph = module.graph_for(s1, s1, s1) self.assertAllFused(forward_graph, except_for=("aten::size", "prim::BroadcastSizes", "aten::_size_if_not_equal")) @@ -864,9 +921,13 @@ def my_broadcasted_cell(a, b, c): old_plans = set() for i in range(3): # if we have s2, then the s1 are _grad_sum_to_size'd + args = s2 if i < 1 else s1, s2 if i < 2 else s1, s2 args = [a.detach_().requires_grad_() for a in args] + # recompile, so we don't trigger bailouts + module = self.checkScript(my_broadcasted_cell, args, profiling=ProfilingMode.FULL) res = module(s2 if i < 1 else s1, s2 if i < 2 else s1, s2) + warmup_backward(res.sum(), args) grads = torch.autograd.grad(res.sum(), args) for inp, gr in zip(args, grads): self.assertEqual(inp.shape, gr.shape) @@ -878,9 +939,8 @@ def my_broadcasted_cell(a, b, c): assert backward is None backward = g old_plans.add(str(backward)) - self.assertEqual(len([1 for o in next(backward.outputs()).node().inputs() - if o.node().kind() == "aten::_grad_sum_to_size"]), i) - self.assertEqual(len([1 for o in next(backward.outputs()).node().inputs() if o.node().kind() == "prim::Param"]), 3 - i) + num_grads = 1 if i > 0 else 0 + self.assertEqual(len([n for n in backward.nodes() if n.kind() == 'aten::_grad_sum_to_size']), num_grads) if __name__ == '__main__': diff --git a/tools/build_variables.py b/tools/build_variables.py index 3f980938b3229..32ae67cbb4ec1 100644 --- a/tools/build_variables.py +++ b/tools/build_variables.py @@ -108,6 +108,7 @@ "torch/csrc/jit/passes/canonicalize_ops.cpp", "torch/csrc/jit/passes/decompose_ops.cpp", "torch/csrc/jit/passes/canonicalize.cpp", + "torch/csrc/jit/passes/clear_undefinedness.cpp", "torch/csrc/jit/passes/common_subexpression_elimination.cpp", "torch/csrc/jit/passes/constant_propagation.cpp", "torch/csrc/jit/passes/constant_pooling.cpp", diff --git a/torch/csrc/jit/autodiff.cpp b/torch/csrc/jit/autodiff.cpp index 6e26375ad660b..af13353414308 100644 --- a/torch/csrc/jit/autodiff.cpp +++ b/torch/csrc/jit/autodiff.cpp @@ -788,12 +788,31 @@ static void lambdaLiftReverse(Gradient& grad_desc, ReverseDetails& rev_info) { for (auto& offset : grad_desc.df_input_captured_outputs) add_capture(graph.outputs()[offset]); - GRAPH_DUMP(" forward graph: ", &graph); - GRAPH_DEBUG(" backward graph: ", *(reverse_block->owningNode())); grad_desc.df = std::make_shared(); grad_desc.df->block()->cloneFrom(reverse_block, [&](Value* v) { return grad_desc.df->inputs()[capture_to_formal_index.at(v)]; }); + + // if we actually profile we can rely on profiling information + // so we don't have to mark every gradient as possibly undefined + if (!getProfilingMode() && getExecutorMode()) { + for (size_t i = 0; i < grad_desc.df_input_vjps.size(); i++) { + auto tt = grad_desc.df->block()->inputs().at(i); + if (auto ttt = tt->type()->cast()) { + tt->setType(ttt->withPossiblyUndefined()); + } else if (auto lt = tt->type()->cast()) { + auto undef_type = + lt->getElementType()->expect()->withPossiblyUndefined(); + tt->setType(ListType::create(undef_type)); + } else { + // unexpected type + TORCH_INTERNAL_ASSERT(false); + } + } + } + + GRAPH_DUMP(" forward graph: ", &graph); + GRAPH_DEBUG(" backward graph: ", *(reverse_block->owningNode())); // reverse_node was just to hold onto reverse_block in a debuggable way // we can remove it now. reverse_block->owningNode()->destroy(); @@ -829,9 +848,6 @@ Gradient differentiate(std::shared_ptr& graph) { // Fills in f, df, f_real_outputs, df_input_captures, // modifies df_input_vjps (new vjps are added for temporaries) lambdaLiftReverse(grad_desc, rev_info); - // It's possible the we've cloned the same constants many times, so - // de-duplicate them - ConstantPooling(grad_desc.df); packReturnValuesIntoTuple(grad_desc.df); return grad_desc; } diff --git a/torch/csrc/jit/fuser/executor.cpp b/torch/csrc/jit/fuser/executor.cpp index 69ddad91f016a..d485a6c09a0e2 100644 --- a/torch/csrc/jit/fuser/executor.cpp +++ b/torch/csrc/jit/fuser/executor.cpp @@ -338,6 +338,10 @@ bool runFusion(const int64_t key, Stack& stack, std::string* code_out) { inputs.emplace_back(all_inputs[i].toTensor()); } + if (!inputs.at(0).defined()) { + return false; + } + // Determines device to dispatch to. at::Device device = inputs.at(0).device(); // If there's a device mismatch in the inputs or if one of the input is a diff --git a/torch/csrc/jit/graph_executor.cpp b/torch/csrc/jit/graph_executor.cpp index 2129aa0f897cd..4abb884d31c27 100644 --- a/torch/csrc/jit/graph_executor.cpp +++ b/torch/csrc/jit/graph_executor.cpp @@ -221,6 +221,7 @@ struct DifferentiableGraphBackward : public autograd::Node { input_instructions_.unpack(std::move(inputs), stack); captures_.unpack(stack, shared_from_this()); + GRAPH_DEBUG("Running DifferentiableGraphBackward for ", &executor); executor.run(stack); unpackReturnTuple(stack); @@ -552,6 +553,7 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { // Phase 0. Inline functions, then clean up any artifacts that the inliner // left in that may inhibit optimization Inline(*opt_graph); + LowerGradOf(*opt_graph); specializeAutogradZero(*opt_graph); LowerSimpleTuples(opt_graph); ConstantPooling(opt_graph); @@ -621,10 +623,10 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { GraphExecutor::GraphExecutor(std::shared_ptr graph) : pImpl( - getProfilingMode() ? dynamic_cast( - new ProfilingGraphExecutorImpl(graph)) - : dynamic_cast( - new GraphExecutorImpl(graph))) {} + getExecutorMode() ? dynamic_cast( + new ProfilingGraphExecutorImpl(graph)) + : dynamic_cast( + new GraphExecutorImpl(graph))) {} void GraphExecutor::run(Stack& inputs) { return pImpl->run(inputs); @@ -643,7 +645,6 @@ GraphExecutorState GraphExecutor::getDebugState() { } void runRequiredPasses(const std::shared_ptr& g) { - LowerGradOf(*g); // implicit inserted expand nodes are not necessarily always valid // when used inside script methods that might have unstable shapes // we remove the implicitly created ones, and have shape analysis @@ -681,14 +682,33 @@ static bool mayIntroduceGradient(const Block* b) { } bool needsGradient(const std::shared_ptr& graph) { - if (!autograd::GradMode::is_enabled()) + if (!autograd::GradMode::is_enabled()) { return false; - if (mayIntroduceGradient(graph->block())) + } + + if (mayIntroduceGradient(graph->block())) { return true; - for (const Value* input : graph->inputs()) { - if (input->type()->requires_grad()) - return true; } + + if (getProfilingMode()) { + for (const Value* input : graph->inputs()) { + for (const auto& use : input->uses()) { + if (use.user->kind() == prim::BailOut) { + auto ptt = use.user->output()->type()->expect(); + if (ptt->requiresGrad() && *ptt->requiresGrad()) { + return true; + } + } + } + } + } else { + for (const Value* input : graph->inputs()) { + if (input->type()->requires_grad()) { + return true; + } + } + } + return false; } diff --git a/torch/csrc/jit/graph_executor.h b/torch/csrc/jit/graph_executor.h index 14cc8376b3f16..8c8dd3697253e 100644 --- a/torch/csrc/jit/graph_executor.h +++ b/torch/csrc/jit/graph_executor.h @@ -61,6 +61,7 @@ TORCH_API void debugSetAutodiffSubgraphInlining(bool state); TORCH_API std::shared_ptr lastExecutedOptimizedGraph(); TORCH_API std::atomic &getProfilingMode(); +TORCH_API std::atomic& getExecutorMode(); struct TORCH_API GraphOptimizerEnabledGuard { GraphOptimizerEnabledGuard(bool state) diff --git a/torch/csrc/jit/init.cpp b/torch/csrc/jit/init.cpp index 8523822923a62..1192b7a616d56 100644 --- a/torch/csrc/jit/init.cpp +++ b/torch/csrc/jit/init.cpp @@ -325,7 +325,18 @@ void initJITBindings(PyObject* module) { }) .def( "_jit_set_profiling_mode", - [](bool profiling_flag) { getProfilingMode() = profiling_flag; }) + [](bool profiling_flag) { + bool oldState = getProfilingMode(); + getProfilingMode() = profiling_flag; + return oldState; + }) + .def( + "_jit_set_profiling_executor", + [](bool profiling_flag) { + bool oldState = getExecutorMode(); + getExecutorMode() = profiling_flag; + return oldState; + }) .def( "_jit_set_inline_everything_mode", [](bool enabled) { script::getInlineEverythingMode() = enabled; }) diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp index 2241e31e6837e..9f90f660c8e2f 100644 --- a/torch/csrc/jit/ir.cpp +++ b/torch/csrc/jit/ir.cpp @@ -940,6 +940,8 @@ bool Node::hasSideEffects() const { case prim::CallMethod: case prim::BailoutTemplate: case prim::profile: + case prim::BailOut: + case prim::Guard: return true; } diff --git a/torch/csrc/jit/operator.cpp b/torch/csrc/jit/operator.cpp index 02b6a19ae77fd..cc8aaf2dbf8ed 100644 --- a/torch/csrc/jit/operator.cpp +++ b/torch/csrc/jit/operator.cpp @@ -220,8 +220,9 @@ bool Operator::matches(const Node* node) const { const auto& formals = schema().arguments(); // not enough inputs - if (actuals.size() < formals.size()) + if (actuals.size() < formals.size()) { return false; + } TypeEnv type_env; for (size_t i = 0; i < formals.size(); ++i) { @@ -231,6 +232,7 @@ bool Operator::matches(const Node* node) const { if (!matched_type.success()) { return false; } + TypePtr resolved = tryEvalTypeVariables(formal, type_env); if (resolved) { formal = resolved; @@ -239,6 +241,7 @@ bool Operator::matches(const Node* node) const { // not resolved all type variables, e.g. if None was matched to Optional[T] // we will not succeed at matching T. However None <: Optional[T] so this // check can still succeed. + if (!actuals[i]->type()->isSubtypeOf(formal)) { return false; } diff --git a/torch/csrc/jit/passes/alias_analysis.cpp b/torch/csrc/jit/passes/alias_analysis.cpp index 75c3ef2e8bd10..2290a57796107 100644 --- a/torch/csrc/jit/passes/alias_analysis.cpp +++ b/torch/csrc/jit/passes/alias_analysis.cpp @@ -1264,7 +1264,6 @@ bool aliasAnalysisHasSpecialCaseFor(Symbol symbol) { prim::Drop, at::onnx::Reshape, at::onnx::Shape, - prim::AutogradAnyNonZero, prim::AutogradAdd, }; diff --git a/torch/csrc/jit/passes/bailout_graph.cpp b/torch/csrc/jit/passes/bailout_graph.cpp index 95eb6d5e66fa6..e02b4a7765134 100644 --- a/torch/csrc/jit/passes/bailout_graph.cpp +++ b/torch/csrc/jit/passes/bailout_graph.cpp @@ -1,7 +1,8 @@ -#include #include #include +#include #include +#include #include #include #include @@ -23,7 +24,7 @@ static std::unordered_set collectLoopCounts(Node *n) { it = outerNode->owningBlock(); } - return loopCounts; + return std::move(loopCounts); } struct BailOutGraphBuilderForNode { @@ -142,10 +143,14 @@ struct BailOutGraphBuilderForNode { } std::shared_ptr buildBailOutGraphFrom(Node* n) { - + // add graph inputs for guard's input + // and loop counts for loops `n` is contained in + // to make sure we can line bailout grap's inputs up properly + // with arguments to this BailOut node. for (auto bi : n->inputs()) { getOrAddInputForValue(bi); } + buildBailOutBlockFrom(n); // add graph outputs for (auto ov : graph_->outputs()) { @@ -239,11 +244,9 @@ struct BailOutInserter { // currently, there's always one guaded input bailout_node->addInput(it->input()); - // we need to collect loop counts to - // record the number of iterations already run - // however, liveness doesn't capture loop counts - // if they aren't used explicitly in a loop - // so we collect them manually here + // collect loop counts since liveness won't collect them + // if they aren't used explicitly, but they are used + // by BailOut graphs if we trigger a bailout inside a loop auto loopCounts = collectLoopCounts(*it); for (auto lc : loopCounts) { bailout_node->addInput(lc); @@ -296,8 +299,8 @@ void InsertBailOuts(std::shared_ptr graph) { // index matches the given `index` static Node* locateBailOutNodeInUnoptimizedGraph(Block* b, int64_t index) { for (auto n : b->nodes()) { - if (n->kind() == prim::BailOut && n->hasAttribute(attr::index) && - n->i(attr::index) == index) { + if ((n->kind() == prim::BailOut || n->kind() == prim::Guard) && + n->hasAttribute(attr::index) && n->i(attr::index) == index) { return n; } for (auto ib : n->blocks()) { @@ -313,7 +316,7 @@ static Node* locateBailOutNodeInUnoptimizedGraph(Block* b, int64_t index) { // to its users static void removeBailouts(Block* b) { for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) { - if (it->kind() == prim::BailOut) { + if (it->kind() == prim::BailOut || it->kind() == prim::Guard) { // clear profiling information it->inputs().at(0)->setType(TensorType::get()); it->output()->replaceAllUsesWith(it->inputs().at(0)); @@ -333,16 +336,20 @@ TORCH_API std::shared_ptr BuildBailOutGraphFrom( const std::shared_ptr& target) { auto orig_bailout_node = locateBailOutNodeInUnoptimizedGraph(orig->block(), bailout_index); + + GRAPH_DEBUG("bailout triggered for ", *orig_bailout_node); + GRAPH_DUMP("original bailout graph ", orig); TORCH_INTERNAL_ASSERT( orig_bailout_node->inputs().at(0)->type()->cast() == nullptr); TORCH_INTERNAL_ASSERT( - orig_bailout_node && orig_bailout_node->kind() == prim::BailOut && + orig_bailout_node && + (orig_bailout_node->kind() == prim::BailOut || + orig_bailout_node->kind() == prim::Guard) && bailout_index == orig_bailout_node->i(attr::index)); BailOutGraphBuilderForNode bg(orig, target); auto bailout_graph = bg.buildBailOutGraphFrom(orig_bailout_node); removeBailouts(bailout_graph->block()); - ConstantPooling(bailout_graph); return bailout_graph; } diff --git a/torch/csrc/jit/passes/clear_undefinedness.cpp b/torch/csrc/jit/passes/clear_undefinedness.cpp new file mode 100644 index 0000000000000..972f7e62aa516 --- /dev/null +++ b/torch/csrc/jit/passes/clear_undefinedness.cpp @@ -0,0 +1,38 @@ +#include +#include + +namespace torch { +namespace jit { + +void clearUndefinedness(Value* o) { + if (o->type()->kind() == TensorType::Kind) { + o->setType(TensorType::get()); + } else if ( + o->type()->kind() == ListType::Kind && + o->type()->expect()->getElementType()->kind() == + TensorType::Kind) { + o->setType(ListType::create(TensorType::get())); + } +} + +void clearUndefinedness(Block* block) { + for (auto n : block->nodes()) { + for (auto o : n->outputs()) { + clearUndefinedness(o); + } + for (auto ib : n->blocks()) { + clearUndefinedness(ib); + } + } +} + +void ClearUndefinedness(const std::shared_ptr& graph) { + for (auto i : graph->inputs()) { + clearUndefinedness(i); + } + clearUndefinedness(graph->block()); + GRAPH_DUMP("After removeUndefinedness: ", graph); +} + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/passes/clear_undefinedness.h b/torch/csrc/jit/passes/clear_undefinedness.h new file mode 100644 index 0000000000000..17f78bfda3cba --- /dev/null +++ b/torch/csrc/jit/passes/clear_undefinedness.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace torch { +namespace jit { + +// Undefinedness makes argument matching fail for regular tensor operations +// if 1+ arguments are undefined or possibly undefined tensors. +// Technically, undefined tensors are **not** tensors as the regular tensor +// operations do not know how to handle them. +// However, in practice, there are guards and conversion operators that +// **always** gate regular operations if undefined tensors may be present +// Eventually, we would love to move to the world where we use optionals +// in lieu of undefined tensors. +// When this happens, this pass will be removed +TORCH_API void ClearUndefinedness(const std::shared_ptr& graph); + +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/passes/guard_elimination.cpp b/torch/csrc/jit/passes/guard_elimination.cpp index 0fefdf5a0a50e..4f160acea99fb 100644 --- a/torch/csrc/jit/passes/guard_elimination.cpp +++ b/torch/csrc/jit/passes/guard_elimination.cpp @@ -16,7 +16,10 @@ struct GuardElimination { aliasDb_(c10::guts::make_unique(graph_)) {} void run() { - moveGuardsToDefs(graph_->block()); + const size_t MAX_ATTEMPTS = 5; + size_t attempts = MAX_ATTEMPTS; + while (attempts-- && moveGuardsToDefs(graph_->block())) { + } GRAPH_DUMP("After moveGuardsToDefs", graph_); coalesceGuards(graph_->block()); GRAPH_DUMP("After coalesceGuards", graph_); @@ -24,7 +27,16 @@ struct GuardElimination { GRAPH_DUMP("After eliminateRedundantGuards", graph_); } - void moveGuardsToDefs(Block* b) { + static bool isLoweredGradOf(Node* n) { + if (n->kind() != prim::If) { + return false; + } + + return n->input(0)->node()->kind() == prim::AutogradAnyNonZero; + } + + bool moveGuardsToDefs(Block* b) { + bool changed = false; for (auto it = b->nodes().begin(); it != b->nodes().end();) { auto n = *it; if (n->kind() == prim::Guard) { @@ -39,6 +51,7 @@ struct GuardElimination { guardee = *n->owningBlock()->nodes().begin(); } bool moved = aliasDb_->moveAfterTopologicallyValid(n, guardee); + changed |= moved; if (moved) { GRAPH_UPDATE( "Moved ", @@ -53,6 +66,21 @@ struct GuardElimination { } } } + + if (b->owningNode() && + isLoweredGradOf( + b->owningNode()) /*b->owningNode()->kind() == prim::If*/) { + for (auto it = b->nodes().begin(); it != b->nodes().end();) { + auto block_node = *it++; + if (block_node->kind() != prim::Guard) { + break; + } + block_node->moveBefore(b->owningNode()); + changed = true; + } + } + + return changed; } void coalesceGuards(Block* b) { @@ -183,7 +211,6 @@ struct GuardElimination { // Guards can be removed if all inputs are guarded and `isSummarized()` // returns // false or inputs are `prim::Constant` - // bool removableGuard(Node *n) { const static auto no_exceptions = std::unordered_set{}; @@ -208,6 +235,17 @@ struct GuardElimination { case aten::neg: case prim::ConstantChunk: case aten::size: + case aten::abs: + case aten::sign: + case aten::pow: + case aten::relu: + case aten::threshold: + case aten::avg_pool2d: + case prim::AutogradAdd: + case prim::AutogradZero: + case aten::rand_like: + case aten::erf: + case aten::erfc: return checkInputs(n, no_exceptions); case aten::cat: // check that the dimension argument is constant @@ -240,7 +278,27 @@ struct GuardElimination { } } return false; + + // this is checked by one of the tests in test_jit_fuser.py + case prim::ListUnpack: { + // check if the input is a constant chunk + // used for LSTM fusions + auto chunk = n->input(0)->node(); + if (chunk->kind() != aten::chunk) { + return false; + } + return checkInputs(chunk, no_exceptions); + } + // this is checked by one of the tests in test_jit_fuser.py + case aten::broadcast_tensors: { + auto list_construct = n->input(0)->node(); + if (list_construct->kind() != prim::ListConstruct) { + return false; + } + return checkInputs(list_construct, no_exceptions); + } case prim::Guard: + case prim::GradOf: return true; default: GRAPH_DEBUG("cannot remove ", n->kind().toQualString()); diff --git a/torch/csrc/jit/passes/specialize_autogradzero.cpp b/torch/csrc/jit/passes/specialize_autogradzero.cpp index f90fde68dd94c..bc0bd59ad59a3 100644 --- a/torch/csrc/jit/passes/specialize_autogradzero.cpp +++ b/torch/csrc/jit/passes/specialize_autogradzero.cpp @@ -1,5 +1,4 @@ #include -#include #include namespace torch { @@ -18,10 +17,14 @@ void specializeAutogradZero(Graph &g) { for (Value* input : g.inputs()) { const auto& tp = input->type(); if (auto tt = tp->cast()) { - if (tt->undefined() && *tt->undefined()) { - state[input] = State::Zero; + if (tt->undefined()) { + if (*tt->undefined()) { + state[input] = State::Zero; + } else { + state[input] = State::Nonzero; + } } else { - state[input] = State::Nonzero; + state[input] = State::Unknown; } } else if ( tp->isSubtypeOf(TensorType::get()) || @@ -34,75 +37,23 @@ void specializeAutogradZero(Graph &g) { for (auto it = g.nodes().begin(); it != g.nodes().end(); ++it) { auto n = *it; - switch (n->kind()) { - case prim::GradOf: { - auto all_zeros = - std::all_of(n->inputs().begin(), n->inputs().end(), [&](Value* v) { - return state[v] == State::Zero; - }); - // Property 1: if all the gradInputs to the GradOf are Zero - // then the gradOutputs are also zero and will be represented as - // AutogradZero nodes - if (all_zeros) { - auto zero = g.createAutogradZero()->insertAfter(n)->output(); - for (auto o : n->outputs()) { - GRAPH_UPDATE("Replacing output %", o->debugName(), - " with AutogradZero %", zero->debugName()); - o->replaceAllUsesWith(zero); - } - } else { - // Property 2: GradOfs are required to correctly handle combinations - // of Nonzero and zero inputs. They are expected to produce - // Nonzero output tensors in this case. - // Remove the GradOf, splicing its body back into the surrounding - // block - auto body = n->blocks().at(0); - for (auto input : n->inputs()) { - // we should never get into a situation when specializing a GradOf - // where we do not know if a value is Nonzero since at the top level - // a gradient graph is composed of Linear nodes and AutogradAdds - // and LinearNodes only appear in these graphs - AT_ASSERT(state[input] != State::Unknown); - } - // hoist the nodes in the GradOf body to be before the linear block - GRAPH_UPDATE("Hoisting out ", getHeader(*it)); - for (auto it = body->nodes().begin(); it != body->nodes().end();) { - auto block_node = *it++; - block_node->moveBefore(n); - } - - for (size_t i = 0; i < n->outputs().size(); ++i) { - GRAPH_UPDATE("Replacing prim::GradOf's use %", - n->outputs().at(i)->debugName(), - " with hoisted value %", - body->outputs().at(i)->debugName()); - n->outputs().at(i)->replaceAllUsesWith(body->outputs().at(i)); - } - } - GRAPH_UPDATE("Destroying ", getHeader(*it)); - it.destroyCurrent(); - } break; + switch (n->kind()) { case prim::AutogradAdd: { auto a = n->input(0); auto b = n->input(1); // if one is Autograd zero, we can just drop the add if (state[a] == State::Zero) { // Zero + b == b - GRAPH_UPDATE("Simplifying ", getHeader(n), " where %", a->debugName(), - " is AutogradZero to %", b->debugName()); n->output()->replaceAllUsesWith(b); it.destroyCurrent(); } else if (state[b] == State::Zero) { // a + Zero == a - GRAPH_UPDATE("Simplifying ", getHeader(n), " where %", b->debugName(), - " is AutogradZero to %", a->debugName()); n->output()->replaceAllUsesWith(a); it.destroyCurrent(); } else if (state[a] == State::Nonzero && state[b] == State::Nonzero) { // when both are Nonzero, we can use a normal, optimizable add // instruction - WithInsertPoint guard(n); auto* g = n->owningGraph(); auto* cOne = g->insertConstant(1); @@ -113,8 +64,6 @@ void specializeAutogradZero(Graph &g) { auto* add_output = add_node->output(); state[add_output] = State::Nonzero; n->output()->replaceAllUsesWith(add_output); - GRAPH_UPDATE("Simplifying ", getHeader(n), " to ", - getHeader(add_node)); it.destroyCurrent(); } else { // otherwise we have conditionally-Nonzero things, and we need @@ -152,6 +101,54 @@ void specializeAutogradZero(Graph &g) { : State::Unknown; } } break; + // Lowered GradOf block + case prim::If: { + auto if_input = n->input(0)->node(); + if (if_input->kind() == prim::AutogradAnyNonZero) { + auto all_zeros = std::all_of( + if_input->inputs().begin(), + if_input->inputs().end(), + [&](Value* v) { return state[v] == State::Zero; }); + + auto all_nonzeros = std::all_of( + if_input->inputs().begin(), + if_input->inputs().end(), + [&](Value* v) { return state[v] == State::Nonzero; }); + // Property 1: if all the gradInputs to the GradOf are Zero + // then the gradOutputs are also zero and will be represented as + // AutogradZero nodes + if (all_zeros) { + auto zero = g.createAutogradZero()->insertAfter(n)->output(); + state[zero] = State::Zero; + for (auto o : n->outputs()) { + o->replaceAllUsesWith(zero); + } + it.destroyCurrent(); + break; + } + + if (all_nonzeros) { + auto body = n->blocks().at(0); + // hoist the nodes in the GradOf body to be before the linear block + for (auto it = body->nodes().begin(); it != body->nodes().end();) { + auto block_node = *it++; + block_node->moveBefore(n); + } + + for (size_t i = 0; i < n->outputs().size(); ++i) { + n->outputs().at(i)->replaceAllUsesWith(body->outputs().at(i)); + state[body->outputs().at(i)] = State::Nonzero; + } + it.destroyCurrent(); + break; + } + } + + for (auto o : n->outputs()) { + state[o] = State::Unknown; + } + break; + } default: for (auto o : n->outputs()) { state[o] = State::Unknown; diff --git a/torch/csrc/jit/profiling_graph_executor_impl.cpp b/torch/csrc/jit/profiling_graph_executor_impl.cpp index 4b3c9779b610c..b6fa448bcc45d 100644 --- a/torch/csrc/jit/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/profiling_graph_executor_impl.cpp @@ -1,9 +1,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -18,34 +20,37 @@ namespace torch { namespace jit { static std::atomic profiling_mode{false}; +static std::atomic executor_mode{false}; + std::atomic& getProfilingMode() { return profiling_mode; } +std::atomic& getExecutorMode() { + return executor_mode; +} + +static bool needsGradientInProfilingMode(Block* b) { + for (auto n : b->nodes()) { + if (n->kind() == prim::BailOut) { + auto ptt = n->output()->type()->expect(); + if (ptt->requiresGrad() && *ptt->requiresGrad()) { + return true; + } + } + + for (auto ib : n->blocks()) { + if (needsGradientInProfilingMode(ib)) { + return true; + } + } + } + return false; +} std::shared_ptr ProfilingGraphExecutorImpl::prepareGraph( const std::shared_ptr& graph, Stack& stack) { auto g = graph->copy(); - ArgumentSpec spec = - arg_spec_creator_.create(autograd::GradMode::is_enabled(), stack); - arg_spec_creator_.specializeTypes(*g, spec); - runRequiredPasses(g); - PropagateRequiresGrad(g); - ConstantPropagation(g); - if (needsGradient(g)) { - auto diff_nodes = CreateAutodiffSubgraphs( - g, getAutodiffSubgraphInlining() ? autodiffSubgraphNodeThreshold : 1); - for (Node* dnode : diff_nodes) { - auto diff_graph = std::move(dnode->g(attr::Subgraph)); - Gradient gradient = differentiate(diff_graph); - // do not optimize DifferentiableGraphs, since - // ideally they will be profiled and then optimized separetely - // when their corresponding DifferentiableGraphOp is called - packGradient(gradient, dnode); - } - InlineAutodiffSubgraphs( - g, getAutodiffSubgraphInlining() ? autodiffSubgraphInlineThreshold : 1); - } return g; } @@ -54,54 +59,61 @@ ProfilingGraphExecutorImpl::ProfilingGraphExecutorImpl( : GraphExecutorImplBase(graph), arg_spec_creator_(*this->graph) {} ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { + GRAPH_DEBUG("Running ProfilingGraphExecutorImpl ", this); if (optimized_plan_) { return *optimized_plan_; } - if (!pr_) { - pr_ = ProfilingRecord::instrumentGraph(prepareGraph(graph, stack)); - auto copy = pr_->graph()->copy(); - LowerGradOf(*copy); - RemoveExpands(copy); - CanonicalizeOps(copy); - EliminateDeadCode(copy); - profiling_plan_ = ExecutionPlan(copy); - // fall-through - } + std::shared_ptr copy; + if (getProfilingMode()) { + if (!pr_) { + pr_ = ProfilingRecord::instrumentGraph(prepareGraph(graph, stack)); + auto copy = pr_->graph()->copy(); + LowerGradOf(*copy); + specializeAutogradZero(*copy); + runRequiredPasses(copy); + GRAPH_DUMP("Profiled Graph: ", copy); + profiling_plan_ = ExecutionPlan(copy); + // fall-through + } - if (!pr_->ready()) { - return *profiling_plan_; + if (!pr_->ready()) { + return *profiling_plan_; + } + copy = pr_->graph()->copy(); + + } else { + copy = graph->copy(); } - // copy already has differentiableGraphs - auto copy = pr_->graph()->copy(); if (!getGraphExecutorOptimize()) { runRequiredPasses(copy); optimized_plan_ = ExecutionPlan(copy); return *optimized_plan_; } - // insert bailouts InsertGuards(copy); - // get rid of autograd specific ops - // we can probably make guard_elimination.cpp - // to handle these ops - specializeAutogradZero(*copy); - // hoist out GradOf blocks - // otherwise we will need to teach - // liveness and buildBailOut graphs - // about them LowerGradOf(*copy); - // constant fold into ConstantChunk - CanonicalizeOps(copy); - EliminateRedundantGuards(copy); - InsertBailOuts(copy); - // TODO: this runs specializeAutogradZero ?? - GRAPH_DUMP("After InsertBailOuts: ", copy); + if (getProfilingMode()) { + EliminateRedundantGuards(copy); + InsertBailOuts(copy); + GRAPH_DUMP("After InsertBailOuts: ", copy); + } + + specializeAutogradZero(*copy); + if (!getProfilingMode()) { + ClearUndefinedness(copy); + } + runRequiredPasses(copy); ConstantPropagation(copy); runOptimization(copy); - if (needsGradient(copy)) { + + // TODO: insert grad propagation + bool needs_gradient = getProfilingMode() + ? needsGradientInProfilingMode(copy->block()) + : needsGradient(copy); + if (needs_gradient) { auto diff_nodes = CreateAutodiffSubgraphs( copy, getAutodiffSubgraphInlining() ? autodiffSubgraphNodeThreshold : 1); @@ -120,6 +132,7 @@ ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { runNondiffOptimization(copy); } EliminateDeadCode(copy); + GRAPH_DUMP("Optimized Graph : ", copy); // cache optimized_plan_ = ExecutionPlan(copy); return *optimized_plan_; @@ -127,7 +140,11 @@ ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { GraphExecutorState ProfilingGraphExecutorImpl::getDebugState() { - AT_ERROR("not supported"); + GraphExecutorState state; + TORCH_INTERNAL_ASSERT(optimized_plan_); + auto opt_plan = *optimized_plan_; + state.execution_plans.emplace(ArgumentSpec{0, 0}, opt_plan); + return state; } } // namespace jit diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp index d35c558a4c090..032641c05546b 100644 --- a/torch/csrc/jit/register_prim_ops.cpp +++ b/torch/csrc/jit/register_prim_ops.cpp @@ -931,15 +931,26 @@ RegisterOperators reg( }, aliasAnalysisSpecialCase()), Operator( - prim::AutogradAnyNonZero, + "prim::AutogradAnyNonZero(...) -> int", [](const Node* node) -> Operation { size_t num_inputs = node->inputs().size(); - return [=](Stack& stack) { + return [num_inputs](Stack& stack) { bool result = false; - for (const IValue& t : last(stack, num_inputs)) { - if (t.toTensor().defined()) { - result = true; - break; + for (const IValue& v : last(stack, num_inputs)) { + if (v.isTensor()) { + if (v.toTensor().defined()) { + result = true; + break; + } + } else if (v.isTensorList()) { + for (const at::Tensor& t : v.toTensorListRef()) { + result = true; + } + if (result) { + break; + } + } else { + TORCH_INTERNAL_ASSERT(false); } } drop(stack, num_inputs); @@ -947,18 +958,25 @@ RegisterOperators reg( return 0; }; }, - aliasAnalysisSpecialCase()), + aliasAnalysisFromSchema()), Operator( prim::AutogradAdd, [](Stack& stack) { at::Tensor a, b; pop(stack, a, b); - if (!a.defined()) + if (!a.defined() && !b.defined()) { + // undef + undef == undef + stack.emplace_back(a); + } + else if (!a.defined()) { stack.emplace_back(b); - else if (!b.defined()) + } + else if (!b.defined()) { stack.emplace_back(a); - else + } + else { stack.emplace_back(a + b); + } return 0; }, aliasAnalysisSpecialCase()), From ef5a6b22623c8b8abec2bb4672fc34f2ea7960a4 Mon Sep 17 00:00:00 2001 From: Jianyu Huang Date: Tue, 29 Oct 2019 11:50:25 -0700 Subject: [PATCH 23/64] Avoid the misleading zero_point and scale [2/2] (#28827) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28827 When we print the `DynamicLinear` module, we don't want to print the scale and zero points as they are not needed for the dynamic quantization. Let's take the output of RoBERTa model as an example: Before this PR: ``` (19): TransformerEncoderLayer( (dropout): Dropout(p=0.1, inplace=False) (attention): MultiheadAttention( (dropout): Dropout(p=0.1, inplace=False) (input_projection): DynamicQuantizedLinear(in_features=1024, out_features=3072, scale=1.0, zero_point=0) (output_projection): DynamicQuantizedLinear(in_features=1024, out_features=1024, scale=1.0, zero_point=0) ) (residual_mlp): ResidualMLP( (mlp): Sequential( (0): DynamicQuantizedLinear(in_features=1024, out_features=4096, scale=1.0, zero_point=0) (1): GeLU() (2): Dropout(p=0.1, inplace=False) (3): DynamicQuantizedLinear(in_features=4096, out_features=1024, scale=1.0, zero_point=0) (4): Dropout(p=0.1, inplace=False) ) ) (attention_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) ) (20): TransformerEncoderLayer( (dropout): Dropout(p=0.1, inplace=False) (attention): MultiheadAttention( (dropout): Dropout(p=0.1, inplace=False) (input_projection): DynamicQuantizedLinear(in_features=1024, out_features=3072, scale=1.0, zero_point=0) (output_projection): DynamicQuantizedLinear(in_features=1024, out_features=1024, scale=1.0, zero_point=0) ) (residual_mlp): ResidualMLP( (mlp): Sequential( (0): DynamicQuantizedLinear(in_features=1024, out_features=4096, scale=1.0, zero_point=0) (1): GeLU() (2): Dropout(p=0.1, inplace=False) (3): DynamicQuantizedLinear(in_features=4096, out_features=1024, scale=1.0, zero_point=0) (4): Dropout(p=0.1, inplace=False) ) ) (attention_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) ) ``` After this PR: ``` (19): TransformerEncoderLayer( (dropout): Dropout(p=0.1, inplace=False) (attention): MultiheadAttention( (dropout): Dropout(p=0.1, inplace=False) (input_projection): DynamicQuantizedLinear(in_features=1024, out_features=3072) (output_projection): DynamicQuantizedLinear(in_features=1024, out_features=1024) ) (residual_mlp): ResidualMLP( (mlp): Sequential( (0): DynamicQuantizedLinear(in_features=1024, out_features=4096) (1): GeLU() (2): Dropout(p=0.1, inplace=False) (3): DynamicQuantizedLinear(in_features=4096, out_features=1024) (4): Dropout(p=0.1, inplace=False) ) ) (attention_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) ) (20): TransformerEncoderLayer( (dropout): Dropout(p=0.1, inplace=False) (attention): MultiheadAttention( (dropout): Dropout(p=0.1, inplace=False) (input_projection): DynamicQuantizedLinear(in_features=1024, out_features=3072) (output_projection): DynamicQuantizedLinear(in_features=1024, out_features=1024) ) (residual_mlp): ResidualMLP( (mlp): Sequential( (0): DynamicQuantizedLinear(in_features=1024, out_features=4096) (1): GeLU() (2): Dropout(p=0.1, inplace=False) (3): DynamicQuantizedLinear(in_features=4096, out_features=1024) (4): Dropout(p=0.1, inplace=False) ) ) (attention_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) ) ``` ghstack-source-id: 92807317 Test Plan: CI Differential Revision: D18197022 fbshipit-source-id: e41635330cfdfb008a0468d6a8ff67a06f7e1c59 --- torch/nn/quantized/dynamic/modules/linear.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/torch/nn/quantized/dynamic/modules/linear.py b/torch/nn/quantized/dynamic/modules/linear.py index 04b942f6dacd2..360eb233ba725 100644 --- a/torch/nn/quantized/dynamic/modules/linear.py +++ b/torch/nn/quantized/dynamic/modules/linear.py @@ -45,6 +45,11 @@ def forward(self, x): def _get_name(self): return 'DynamicQuantizedLinear' + def extra_repr(self): + return 'in_features={}, out_features={}'.format( + self.in_features, self.out_features + ) + @classmethod def from_float(cls, mod): r"""Create a dynamic quantized module from a float module or qparams_dict From 4703854321bcbcddc3956641cd0f6fd98910057e Mon Sep 17 00:00:00 2001 From: Mingzhe Li Date: Tue, 29 Oct 2019 11:51:33 -0700 Subject: [PATCH 24/64] change softmax input shape (#28836) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28836 as title Test Plan: ``` buck run mode/opt //caffe2/benchmarks/operator_benchmark/pt:softmax_test Invalidating internal cached state: Buck configuration options changed between invocations. This may cause slower builds. Changed value project.buck_out='buck-out/opt' (was 'buck-out/dev') ... and 56 more. See logs for all changes Parsing buck files: finished in 6.2 sec Creating action graph: finished in 8.8 sec Building: finished in 05:42.6 min (100%) 28336/28336 jobs, 23707 updated Total time: 05:57.7 min # ---------------------------------------- # PyTorch/Caffe2 Operator Micro-benchmarks # ---------------------------------------- # Tag : short # Benchmarking PyTorch: Softmax /proc/self/fd/4/softmax_test.py:57: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument. """ # Mode: Eager # Name: Softmax_N4_C3_H256_W256 # Input: N: 4, C: 3, H: 256, W: 256 Forward Execution Time (us) : 18422.487 Reviewed By: hl475 Differential Revision: D18202335 fbshipit-source-id: 0bb376cb465d998a49196e148d48d436126ae334 --- benchmarks/operator_benchmark/pt/softmax_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/operator_benchmark/pt/softmax_test.py b/benchmarks/operator_benchmark/pt/softmax_test.py index 8b4f4e6c4ac27..9d20455d49245 100644 --- a/benchmarks/operator_benchmark/pt/softmax_test.py +++ b/benchmarks/operator_benchmark/pt/softmax_test.py @@ -17,8 +17,8 @@ # Configs for softmax ops softmax_configs_short = op_bench.config_list( attrs=[ - [4, 3, 128, 128], - [8, 3, 256, 256], + [4, 3, 256, 256], + [8, 3, 512, 512], ], attr_names=[ 'N', 'C', 'H', 'W' From f42768f8c0f4f74efa55530af8230c1a2e5753f4 Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Tue, 29 Oct 2019 11:52:31 -0700 Subject: [PATCH 25/64] Add scripts to run cuda-memcheck (#28127) Summary: This PR adds scripts that could be used for https://github.com/pytorch/pytorch/issues/26052 Example output: ``` Success: TestTorchDeviceTypeCPU.test_advancedindex_big_cpu Success: TestTorchDeviceTypeCPU.test_addcmul_cpu Success: TestTorchDeviceTypeCPU.test_addbmm_cpu_float32 Success: TestTorchDeviceTypeCPU.test_advancedindex_cpu_float16 Success: TestTorchDeviceTypeCPU.test_addmv_cpu Success: TestTorchDeviceTypeCPU.test_addcdiv_cpu Success: TestTorchDeviceTypeCPU.test_all_any_empty_cpu Success: TestTorchDeviceTypeCPU.test_atan2_cpu Success: TestTorchDeviceTypeCPU.test_advancedindex_cpu_float64 Success: TestTorchDeviceTypeCPU.test_baddbmm_cpu_float32 Success: TestTorchDeviceTypeCPU.test_atan2_edgecases_cpu Success: TestTorchDeviceTypeCPU.test_add_cpu Success: TestTorchDeviceTypeCPU.test_addr_cpu_bfloat16 Success: TestTorchDeviceTypeCPU.test_addr_cpu_float32 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/28127 Differential Revision: D18184255 Pulled By: mruberry fbshipit-source-id: 7fd4bd9faf9f8b37b369f631c63f26eb965b16e7 --- test/scripts/cuda_memcheck_common.py | 98 +++++++++++++++++++++ test/scripts/run_cuda_memcheck.py | 123 +++++++++++++++++++++++++++ 2 files changed, 221 insertions(+) create mode 100644 test/scripts/cuda_memcheck_common.py create mode 100755 test/scripts/run_cuda_memcheck.py diff --git a/test/scripts/cuda_memcheck_common.py b/test/scripts/cuda_memcheck_common.py new file mode 100644 index 0000000000000..7f7dc8253393f --- /dev/null +++ b/test/scripts/cuda_memcheck_common.py @@ -0,0 +1,98 @@ +# this file contains a simple parser that parses report +# from cuda-memcheck + +class ParseError(Exception): + """Whenever the simple parser is unable to parse the report, this exception will be raised""" + pass + + +class Report: + """A report is a container of errors, and a summary on how many errors are found""" + + HEAD = 'ERROR SUMMARY: ' + TAIL = ' errors' + + def __init__(self, text, errors): + self.text = text + self.num_errors = int(text[len(self.HEAD):len(text) - len(self.TAIL)]) + self.errors = errors + if len(errors) != self.num_errors: + raise ParseError("Number of errors does not match") + + +class Error: + """Each error is a section in the output of cuda-memcheck. + Each error in the report has an error message and a backtrace. It looks like: + + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaGetLastError + 0x163) [0x4c493] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x5b77a05] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x39d6d1d] + ========= ..... + """ + + def __init__(self, lines): + self.message = lines[0] + lines = lines[2:] + self.stack = [l.strip() for l in lines] + + +def parse(message): + """A simple parser that parses the report of cuda-memcheck. This parser is meant to be simple + and it only split the report into separate errors and a summary. Where each error is further + splitted into error message and backtrace. No further details are parsed. + + A report contains multiple errors and a summary on how many errors are detected. It looks like: + + ========= CUDA-MEMCHECK + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaPointerGetAttributes. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaPointerGetAttributes + 0x1a9) [0x428b9] + ========= Host Frame:/home/xgao/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch.so [0x5b778a9] + ========= ..... + ========= + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= Host Frame:/usr/local/cuda/lib64/libcudart.so.10.1 (cudaGetLastError + 0x163) [0x4c493] + ========= ..... + ========= + ========= ..... + ========= + ========= Program hit cudaErrorInvalidValue (error 1) due to "invalid argument" on CUDA API call to cudaGetLastError. + ========= Saved host backtrace up to driver entry point at error + ========= Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so.1 [0x38c7b3] + ========= ..... + ========= Host Frame:python (_PyEval_EvalFrameDefault + 0x6a0) [0x1d0ad0] + ========= Host Frame:python (_PyEval_EvalCodeWithName + 0xbb9) [0x116db9] + ========= + ========= ERROR SUMMARY: 4 errors + """ + errors = [] + HEAD = '=========' + headlen = len(HEAD) + started = False + in_message = False + message_lines = [] + lines = message.splitlines() + for l in lines: + if l == HEAD + ' CUDA-MEMCHECK': + started = True + continue + if not started or not l.startswith(HEAD): + continue + l = l[headlen + 1:] + if l.startswith('ERROR SUMMARY:'): + return Report(l, errors) + if not in_message: + in_message = True + message_lines = [l] + elif l == '': + errors.append(Error(message_lines)) + in_message = False + else: + message_lines.append(l) + raise ParseError("No error summary found") diff --git a/test/scripts/run_cuda_memcheck.py b/test/scripts/run_cuda_memcheck.py new file mode 100755 index 0000000000000..f80fa84350c74 --- /dev/null +++ b/test/scripts/run_cuda_memcheck.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python + +"""This script runs cuda-memcheck on the specified unit test. Each test case +is run in its isolated process with a timeout so that: +1) different test cases won't influence each other, and +2) in case of hang, the script would still finish in a finite amount of time. +The output will be written to a log file result.log + +Example usage: + python run_cuda_memcheck.py ../test_torch.py 600 + +Note that running cuda-memcheck could be very slow. +""" + +import asyncio +import torch +import multiprocessing +import argparse +import subprocess +import tqdm +import re +import cuda_memcheck_common as cmc + +ALL_TESTS = [] +GPUS = torch.cuda.device_count() + +# parse arguments +parser = argparse.ArgumentParser(description="Run isolated cuda-memcheck on unit tests") +parser.add_argument('filename', help="the python file for a test, such as test_torch.py") +parser.add_argument('timeout', type=int, help='kill the test if it does not terminate in a certain amount of seconds') +parser.add_argument('--strict', action='store_true', + help='Whether to show cublas/cudnn errors. These errors are ignored by default because' + 'cublas/cudnn does not run error-free under cuda-memcheck, and ignoring these errors') +parser.add_argument('--nproc', type=int, default=multiprocessing.cpu_count(), + help='Number of processes running tests, default to number of cores in the system') +parser.add_argument('--gpus', default='all', + help='GPU assignments for each process, it could be "all", or : separated list like "1,2:3,4:5,6"') +args = parser.parse_args() + +# Filters that ignores cublas/cudnn errors +# TODO (@zasdfgbnm): When can we remove this? Will cublas/cudnn run error-free under cuda-memcheck? +def is_ignored_only(output): + try: + report = cmc.parse(output) + except cmc.ParseError: + # in case the simple parser fails parsing the output of cuda memcheck + # then this error is never ignored. + return False + count_ignored_errors = 0 + for e in report.errors: + if 'libcublas' in ''.join(e.stack) or 'libcudnn' in ''.join(e.stack): + count_ignored_errors += 1 + return count_ignored_errors == report.num_errors + +# Discover tests: +# To get a list of tests, run: +# pytest --setup-only test/test_torch.py +# and then parse the output +proc = subprocess.Popen(['pytest', '--setup-only', args.filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE) +stdout, stderr = proc.communicate() +lines = stdout.decode().strip().splitlines() +for line in lines: + if '(fixtures used:' in line: + line = line.strip().split()[0] + line = line[line.find('::') + 2:] + line = line.replace('::', '.') + ALL_TESTS.append(line) + +# Run tests: +# Since running cuda-memcheck on PyTorch unit tests is very slow, these tests must be run in parallel. +# This is done by using the coroutine feature in new Python versions. A number of coroutines are created; +# they create subprocesses and awaiting them to finish. The number of running subprocesses could be +# specified by the user and by default is the same as the number of CPUs in the machine. +# These subprocesses are balanced across different GPUs on the system by assigning one devices per process, +# or as specified by the user +progress = 0 +logfile = open('result.log', 'w') +progressbar = tqdm.tqdm(total=len(ALL_TESTS)) + +async def run1(coroutine_id): + global progress + + if args.gpus == 'all': + gpuid = coroutine_id % GPUS + else: + gpu_assignments = args.gpus.split(':') + assert args.nproc == len(gpu_assignments), 'Please specify GPU assignmnent for each process, separated by :' + gpuid = gpu_assignments[coroutine_id] + + while progress < len(ALL_TESTS): + test = ALL_TESTS[progress] + progress += 1 + cmd = f'CUDA_VISIBLE_DEVICES={gpuid} cuda-memcheck --error-exitcode 1 python {args.filename} {test}' + proc = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) + try: + stdout, stderr = await asyncio.wait_for(proc.communicate(), args.timeout) + except asyncio.TimeoutError: + print('Timeout:', test, file=logfile) + proc.kill() + else: + if proc.returncode == 0: + print('Success:', test, file=logfile) + else: + stdout = stdout.decode() + stderr = stderr.decode() + should_display = args.strict or not is_ignored_only(stdout) + if should_display: + print('Fail:', test, file=logfile) + print(stdout, file=logfile) + print(stderr, file=logfile) + else: + print('Ignored:', test, file=logfile) + del proc + progressbar.update(1) + +async def main(): + tasks = [asyncio.create_task(run1(i)) for i in range(args.nproc)] + for t in tasks: + await t + +if __name__ == '__main__': + loop = asyncio.get_event_loop() + loop.run_until_complete(main()) From ac4c72db3b607ba3b99e1c7e4f611fbd26e93558 Mon Sep 17 00:00:00 2001 From: Amy Yang Date: Tue, 29 Oct 2019 11:52:38 -0700 Subject: [PATCH 26/64] add DNNLOWP static qparam choosing to pybind Summary: as title Test Plan: test in stacked diff Reviewed By: csummersea Differential Revision: D18123726 fbshipit-source-id: ce75db1e6f314a822a94ebdfc11988fab50ee836 --- caffe2/quantization/server/dnnlowp.h | 7 +++ caffe2/quantization/server/p99.cc | 6 +-- caffe2/quantization/server/pybind.cc | 69 ++++++++++++++++++++++++++++ 3 files changed, 79 insertions(+), 3 deletions(-) diff --git a/caffe2/quantization/server/dnnlowp.h b/caffe2/quantization/server/dnnlowp.h index ae6fd7e81170c..0d7414f962412 100644 --- a/caffe2/quantization/server/dnnlowp.h +++ b/caffe2/quantization/server/dnnlowp.h @@ -148,6 +148,13 @@ class QuantizationFactory { return weight_kind_; } + void SetWeightP99Threshold(float threshold) { + weight_p99_threshold_ = threshold; + } + void SetActivationP99Threshold(float threshold) { + activation_p99_threshold_ = threshold; + } + explicit QuantizationFactory( int activation_precision = 8, // precision used for activations in main operations like matmul diff --git a/caffe2/quantization/server/p99.cc b/caffe2/quantization/server/p99.cc index 9c3aac7b48a58..e4d5eb48e6cfc 100644 --- a/caffe2/quantization/server/p99.cc +++ b/caffe2/quantization/server/p99.cc @@ -14,8 +14,8 @@ TensorQuantizationParams P99::ChooseQuantizationParams( std::vector bins_f( dnnlowp::adjust_hist_to_include_zero(hist, &min, &max)); int nbins = bins_f.size(); - assert(min <= 0.f); - assert(max >= 0.f); + CAFFE_ENFORCE(min <= 0.f); + CAFFE_ENFORCE(max >= 0.f); float org_max = max; float org_min = min; float bin_width = (max - min) / nbins; @@ -32,7 +32,7 @@ TensorQuantizationParams P99::ChooseQuantizationParams( sum += bins_f[i]; CDF[i] = (double)sum / total_sum; } - assert(threshold_ > 0.5 && threshold_ < 1); + CAFFE_ENFORCE(threshold_ > 0.5 && threshold_ < 1); double left_quantile = (1.0f - threshold_) / 2.0f; double right_quantile = 1.0f - left_quantile; int i_begin = 0; diff --git a/caffe2/quantization/server/pybind.cc b/caffe2/quantization/server/pybind.cc index 4709be4c0e2fa..3c2cad0c4991e 100644 --- a/caffe2/quantization/server/pybind.cc +++ b/caffe2/quantization/server/pybind.cc @@ -1,6 +1,8 @@ #include +#include #include "activation_distribution_observer.h" #include "caffe2_dnnlowp_utils.h" +#include "quantization_error_minimization.h" namespace caffe2 { namespace python { @@ -172,4 +174,71 @@ PYBIND11_MODULE(dnnlowp_pybind11, m) { CAFFE_ENFORCE(transformed_net.SerializeToString(&protob)); return pybind11::bytes(protob); }); + + pybind11::class_(m, "QueryTensorQparam") + .def_property_readonly( + "scale", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.scale; + }) + .def_property_readonly( + "zero_point", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.zero_point; + }) + .def_property_readonly( + "min", + [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.Min(); + }) + .def_property_readonly( + "max", [](dnnlowp::TensorQuantizationParams& qparam) { + return qparam.Max(); + }); + + m.def( + "ChooseStaticQuantizationParams", + [](float min, + float max, + const std::vector& bins, + bool preserve_sparsity, + int precision, + const std::string& quant_scheme, + float p99_threshold, + bool is_weight) { + dnnlowp::Histogram hist = dnnlowp::Histogram(min, max, bins); + + dnnlowp::QuantizationFactory::QuantizationKind quant_kind = + dnnlowp::QuantizationFactory::MIN_MAX_QUANTIZATION; + if (quant_scheme.compare("L2_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L2_MIN_QUANTIZATION; + } else if (quant_scheme.compare("L2_MIN_QUANTIZATION_APPROX") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L2_MIN_QUANTIZATION_APPROX; + } else if (quant_scheme.compare("KL_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::KL_MIN_QUANTIZATION; + } else if (quant_scheme.compare("P99_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::P99_QUANTIZATION; + } else if (quant_scheme.compare("L1_MIN_QUANTIZATION") == 0) { + quant_kind = dnnlowp::QuantizationFactory::L1_MIN_QUANTIZATION; + } else { + LOG(INFO) << "Using DNNLOWP default MIN_MAX_QUANTIZATION"; + } + dnnlowp::QuantizationFactory* qfactory = + dnnlowp::QuantizationFactory::GetDefaultInstance(); + if (is_weight) { + qfactory->SetWeightP99Threshold(p99_threshold); + } else { + qfactory->SetActivationP99Threshold(p99_threshold); + } + return qfactory->ChooseQuantizationParams( + hist, quant_kind, precision, preserve_sparsity, is_weight); + }, + pybind11::arg("min"), + pybind11::arg("max"), + pybind11::arg("bins"), + pybind11::arg("preserve_sparsity") = true, + pybind11::arg("precision") = 8, + pybind11::arg("quant_scheme") = "min_max", + pybind11::arg("p99_threshold") = 0.99, + pybind11::arg("is_weight") = false); } From 0a68e8bab09937439dcd8af66a7421e3f32c7cb4 Mon Sep 17 00:00:00 2001 From: Mingzhe Li Date: Tue, 29 Oct 2019 11:53:03 -0700 Subject: [PATCH 27/64] fix op bench runtime error when use_jit is enabled (#28837) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28837 The JIT code used in op bench is not compatibility with latest JIT code path. This diff aims to resolve that issue. Test Plan: ```buck run mode/opt //caffe2/benchmarks/operator_benchmark/pt:add_test -- --use_jit Building: finished in 02:29.8 min (100%) 7055/7055 jobs, 1 updated Total time: 02:30.3 min # ---------------------------------------- # PyTorch/Caffe2 Operator Micro-benchmarks # ---------------------------------------- # Tag : short # Benchmarking PyTorch: add # Mode: JIT # Name: add_M64_N64_K64_cpu # Input: M: 64, N: 64, K: 64, device: cpu Forward Execution Time (us) : 118.052 Reviewed By: hl475 Differential Revision: D18197057 fbshipit-source-id: 92edae8a48abc4115a558a91ba46cc9c3edb2eb8 --- benchmarks/operator_benchmark/benchmark_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/operator_benchmark/benchmark_pytorch.py b/benchmarks/operator_benchmark/benchmark_pytorch.py index 8edbd1090c2ff..7ec3914decac4 100644 --- a/benchmarks/operator_benchmark/benchmark_pytorch.py +++ b/benchmarks/operator_benchmark/benchmark_pytorch.py @@ -76,7 +76,7 @@ def _generate_jit_forward_graph(self): @torch.jit.script def _jit_forward_graph(iters, place_holder): # type: (int, Tensor) - result = torch.jit.annotate(torch.Tensor, None) + result = torch.jit.annotate(torch.Tensor, place_holder) for _ in range(iters): result = func(place_holder) return result From 607defa8a932147989adfeb2f37768ff30eaf011 Mon Sep 17 00:00:00 2001 From: Mingzhe Li Date: Tue, 29 Oct 2019 11:53:15 -0700 Subject: [PATCH 28/64] print per block avg time when running on AI-PEP machines (#28838) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28838 as title Test Plan: ``` buck run mode/opt //caffe2/benchmarks/operator_benchmark/pt:softmax_test -- --ai_pep_format true Total time: 02:36.7 min # ---------------------------------------- # PyTorch/Caffe2 Operator Micro-benchmarks # ---------------------------------------- # Tag : short # Benchmarking PyTorch: Softmax /proc/self/fd/4/softmax_test.py:57: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument. """ PyTorchObserver {"type": "PyTorch_Softmax_N4_C3_H128_W128", "metric": "latency", "unit": "ms", "value": "4.83197245048359"} PyTorchObserver {"type": "PyTorch_Softmax_N4_C3_H128_W128", "metric": "latency", "unit": "ms", "value": "4.839232977246866"} PyTorchObserver {"type": "PyTorch_Softmax_N4_C3_H128_W128", "metric": "latency", "unit": "ms", "value": "4.7970924858236685"} PyTorchObserver {"type": "PyTorch_Softmax_N4_C3_H128_W128", "metric": "latency", "unit": "ms", "value": "4.708389271399938"} # Benchmarking PyTorch: Softmax ... Reviewed By: hl475 Differential Revision: D18202504 fbshipit-source-id: 4a332763432b3b5886f241bb2ce49d4df481a6f3 --- benchmarks/operator_benchmark/benchmark_core.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/benchmarks/operator_benchmark/benchmark_core.py b/benchmarks/operator_benchmark/benchmark_core.py index 7a3d15237d55a..14bd5a2094c30 100644 --- a/benchmarks/operator_benchmark/benchmark_core.py +++ b/benchmarks/operator_benchmark/benchmark_core.py @@ -84,9 +84,6 @@ def __init__(self, args): if self.args.test_name is not None: self.args.tag_filter = None - if self.args.ai_pep_format: - self.print_per_iter = True - def _print_header(self): DASH_LINE = '-' * 40 @@ -199,10 +196,18 @@ def _measure_time(self, launch_test, test_case, iters, print_per_iter): report_run_time = 1e6 * run_time_sec / iters time_trace.append(report_run_time) + # Print out the time spent in each epoch in ms + if self.args.ai_pep_format: + test_name = '_'.join([test_case.framework, test_case.test_config.test_name]) + print("PyTorchObserver " + json.dumps( + { + "type": test_name, + "metric": "latency", + "unit": "ms", + "value": str(report_run_time / 1e3), + } + )) if results_are_significant: - # Print out the last 50 values when running with AI PEP - if self.args.ai_pep_format: - test_case._print_per_iter() break # Re-estimate the hopefully-sufficient From efbaa8a563e0229689c9e5599d932a4865c406cc Mon Sep 17 00:00:00 2001 From: Anjali Chourdia Date: Tue, 29 Oct 2019 11:56:14 -0700 Subject: [PATCH 29/64] added a check for zero stride Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28784 Differential Revision: D18178889 Pulled By: anjali411 fbshipit-source-id: 976810bf3f9def3a8f5ca6885b1e049b831f06f3 --- aten/src/ATen/native/Convolution.cpp | 13 ++++++------- test/test_nn.py | 12 +++++++++--- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp index fabeecf3c76b2..630301ba6e019 100644 --- a/aten/src/ATen/native/Convolution.cpp +++ b/aten/src/ATen/native/Convolution.cpp @@ -31,7 +31,7 @@ struct ConvParams { bool is_output_padding_neg() const; bool is_output_padding_big() const; bool is_padding_neg() const; - bool is_stride_neg() const; + bool is_stride_nonpos() const; void view1d_as_2d(); bool use_cpu_depthwise3x3_winograd(const at::Tensor& input, const at::Tensor& weight) const; bool use_cudnn(const at::Tensor& input) const; @@ -105,15 +105,14 @@ auto ConvParams::is_padding_neg() const -> bool { return is_non_neg; } -auto ConvParams::is_stride_neg() const -> bool { - bool is_non_neg = false; +auto ConvParams::is_stride_nonpos() const -> bool { + bool is_nonpos = false; for (int s : stride) { - is_non_neg |= (s < 0); + is_nonpos |= (s <= 0); } - return is_non_neg; + return is_nonpos; } - auto ConvParams::view1d_as_2d() -> void { if (stride.size() == 1) { stride.insert(stride.begin(), 1); @@ -381,7 +380,7 @@ static void check_shape_forward(const at::Tensor& input, TORCH_CHECK(!params.is_padding_neg(), "negative padding is not supported"); TORCH_CHECK(!params.is_output_padding_neg(), "negative output_padding is not supported"); - TORCH_CHECK(!params.is_stride_neg(), "negative stride is not supported"); + TORCH_CHECK(!params.is_stride_nonpos(), "non-positive stride is not supported"); TORCH_CHECK(weight_dim == k, "Expected ", weight_dim, "-dimensional input for ", weight_dim, diff --git a/test/test_nn.py b/test/test_nn.py index 13fa15a8ae419..44b01e840f3da 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -889,7 +889,7 @@ def test_invalid_conv1d(self): # Negative stride check module = nn.Conv1d(in_channels=3, out_channels=6, kernel_size=3, stride=-1, bias=True).to(dtype) input = torch.randn(1, 3, 4).to(dtype) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def test_mismatch_shape_conv2d(self): @@ -918,7 +918,13 @@ def test_invalid_conv2d(self): # Negative stride check module = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4, stride=-1, bias=True).to(dtype) input = torch.randn(1, 3, 4, 4).to(dtype) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): + module(input) + + # Zero stride check + module = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4, stride=0, bias=True).to(dtype) + input = torch.randn(1, 3, 4, 4).to(dtype) + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def test_invalid_conv3d(self): @@ -930,7 +936,7 @@ def test_invalid_conv3d(self): # Negative stride check module = torch.nn.Conv3d(1, 1, kernel_size=3, stride=-2) input = torch.empty(1, 1, 4, 4, 4) - with self.assertRaisesRegex(RuntimeError, 'negative stride is not supported'): + with self.assertRaisesRegex(RuntimeError, 'non-positive stride is not supported'): module(input) def _test_alpha_dropout(self, cls, input): From b7d472a109f3b6c2bca8920a2991e04c47ee0afd Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 11:58:10 -0700 Subject: [PATCH 30/64] Some fixes for jit overview doc (#28112) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28112 att Test Plan: reading Imported from OSS Differential Revision: D18173102 fbshipit-source-id: d8574758288bfce08eaf0f4f6163284defb56d6e --- torch/csrc/jit/docs/OVERVIEW.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/torch/csrc/jit/docs/OVERVIEW.md b/torch/csrc/jit/docs/OVERVIEW.md index daa5999b3f6ec..8212468343afb 100644 --- a/torch/csrc/jit/docs/OVERVIEW.md +++ b/torch/csrc/jit/docs/OVERVIEW.md @@ -352,7 +352,7 @@ JIT programs are created using either the tracing frontend (`torch.jit.trace`) o The tracer produces graphs by recording what actual operations are done on tensors. The entry point from Python into C++ for tracing using `torch.jit.trace` is `_create_method_from_trace`. -A thread local instance of the TracingState object maintains a mapping between actual data being computing during the trace (e.g. Tensors) stored in IValues, and the abstract `Value*` in the Graph that would compute that value. The functions `void setValueTrace(const IValue&, Value*)` and `Value* getValueTrace(const IValue&)` are used by the tracer to maintain this mapping. +A thread local instance of the TracingState object maintains a mapping between actual data being computed during the trace (e.g. Tensors) stored in IValues, and the abstract `Value*` in the Graph that would compute that value. The functions `void setValueTrace(const IValue&, Value*)` and `Value* getValueTrace(const IValue&)` are used by the tracer to maintain this mapping. An initial IValue to Value mapping is setup up between the inputs to the function being traced and symbolic Value inputs to the Graph being constructed. If we are tracing a `torch.nn.Module`, the tracer also adds Parameters and sub-Modules to the Module being constructed that correspond to the Python `torch.nn.Module` being traced. These values are also added as mapping so that uses of the Parameters in the trace will create uses of the Parameters in the Graph. @@ -384,7 +384,7 @@ The functions `addInputs` and `addOutput` are overloaded to handle the different Currently set/getValueTrace only works on Tensors and Futures. Other types are not natively traced. Instead aggregates like tuples or lists are often flattened into tensors at the end of a trace and explicitly constructed from individual tensors at the beginning of this trace. -The tracer has special behavior when tracing calls to other TorchScript functions. This behavior is implemented in the GraphExecutor right before a Graph is about to be run. If tracing is enabled while running the graph, the GraphExecutor will disable tracing, run the graph as normal, and then inline the Graph into the trace. It then hooks up the IValues computed by running the Graph to inlined Graph's out Values in the inlined graph. +The tracer has special behavior when tracing calls to other TorchScript functions. This behavior is implemented in the GraphExecutor right before a Graph is about to be run. If tracing is enabled while running the graph, the GraphExecutor will disable tracing, run the graph as normal, and then inline the Graph into the trace. It then hooks up the IValues computed by running the Graph to out Values in the inlined graph. > *When a trace calls a TorchScript function, that function is preserved as is, meaning that control-flow is preserved.* This makes it possible to "fix" tracing issues by writing the subset of the program that cannot be traced in script and having the trace invoke it. @@ -392,7 +392,7 @@ The resulting Graph created by tracing is installed as the 'forward' method of t ## Script ## -The script frontend directly converts Python syntax into Modules. Like many compilers this happens in two phases. First, we generate an abstract syntax tree (AST), which is constructed out of Tree objects. The compiler (misnamed, but that is the name of the file) then does semantic analysis on the Tree and lowers it into a Module. We can generate Trees in two ways: (1) using frontend.py, which takes the Python AST and transliterates it into Tree objects, or (2) via the Lexer and Parser which parse python syntax directly. The Lexer/Parser path may seem redundant but it is crucially important. We need to define builtin functions ([script/builtin_functions.cpp](../script/builtin_functions.cpp)) when Python is not linked. We allow users to load TorchScript programs directly from strings without Python ([api/include/torch/jit.h](../../../api/include/torch/jit.h)). We also use this Python syntax as the serialization format for TorchScript, since it allows us to make changes to our IR without breaking backward compatibility. Furthermore, the Lexer is reused to implement the FunctionSchema parser, which turns FunctionSchema declarations from strings into FunctionSchema objects. +The script frontend directly converts Python syntax into Modules. Like many compilers this happens in two phases. First, we generate an abstract syntax tree (AST), which is constructed out of Tree objects. The compiler (misnamed, but that is the name of the file) then does semantic analysis on the Tree and lowers it into a Module. We can generate Trees in two ways: (1) using frontend.py, which takes the Python AST and transliterates it into Tree objects, or (2) via the Lexer and Parser which parse python syntax directly. The Lexer/Parser path may seem redundant but it is crucially important. We need to define builtin functions ([script/builtin_functions.cpp](../script/builtin_functions.cpp)) when Python is not linked. We allow users to load TorchScript programs directly from strings without Python ([api/include/torch/jit.h](../../api/include/torch/jit.h)). We also use this Python syntax as the serialization format for TorchScript, since it allows us to make changes to our IR without breaking backward compatibility. Furthermore, the Lexer is reused to implement the FunctionSchema parser, which turns FunctionSchema declarations from strings into FunctionSchema objects. The following sections look into each the stages in the script frontend in detail. @@ -531,7 +531,7 @@ SugaredValues are how the compiler represents non-first class values during Grap SugaredValues are also how we interact with Python runtime during the compilation process. For instance, `math.pi` is resolved to 3.1415... by first resolving `math` to a SugaredValue representing accesses to Python modules (PythonModuleValue) whose `attr` function turns python numbers into `prim::Constant` Nodes in the graph. -Finally, normal Values are also represented by the SimpleValue SugaredValue in places where it is valid either a SugaredValue or a normal Value to appear. +Finally, normal Values are also represented by the SimpleValue SugaredValue in places where it is valid that either a SugaredValue or a normal Value will appear. ## Resolver ## @@ -601,7 +601,7 @@ In addition to being mutable, tensors also have a set of dynamically determined * size - the precise size of the tensor * requires_grad - whether the tensor is recording its gradient with autograd -Changes in these properties change how operators on tensor will evaluate and would make certain optimization invalid. For instance, if we have fuser capable of generating new cuda kernels but not cpu kernels, it is only valid to fuse operations where the inputs are known to run only on CUDA devices. The GraphExecutor's job is to still enable optimization even when certains combinations of properties prevent optimizations for occurring. +Changes in these properties change how operators on tensor will evaluate and would make certain optimization invalid. For instance, if we have fuser capable of generating new cuda kernels but not cpu kernels, it is only valid to fuse operations where the inputs are known to run only on CUDA devices. The GraphExecutor's job is to still enable optimization even when certains combinations of properties prevent optimizations from occurring. Nodes in a graph are executed *serially* in the order they appear in a block. Nodes may be reordered either during optimization or by the interpreter itself if it can be proven that it is not distinguishable from the original serial execution order. These semantics are necessary since the combination of mutable tensors and potential alias between tensors makes it unsafe to perform arbitrary reordering otherwise. However, the AliasInfo object can accurately track how alias propagate through builtin operators so optimization passes can query when certain reorders or optimizations are safe. @@ -797,7 +797,7 @@ graph(%x : Tensor, return (%30) ``` -Execution starts in `GraphExecutor::run`, which takes takes a Stack of inputs. +Execution starts in `GraphExecutor::run`, which takes a Stack of inputs. *Specialization* The executor *specializes* the Graph for the particular set of inputs. Specialization is handled by the `ArgumentSpec` object which extracts a "signature" composed of all the properties being specialized. We only specialize to the properties of Tensors. The ArgumentSpec only records properties for Tensors that either appear directly in the inputs to the graph or inside Tuples that are inputs to the Graph. The properties recorded are currently: @@ -1050,7 +1050,7 @@ one specifies a file(s) in `PYTORCH_JIT_LOG_LEVEL`. `GRAPH_DEBUG` can be enabled by prefixing a file name with an `>` as in `>alias_analysis`. `>>` and `>>>` are also valid and **currently** are equivalent to `GRAPH_DEBUG` as there is no logging level that is -higher than `GRAPH_DEBUG`. +higher than `GRAPH_DEBUG`. ## DifferentiableGraphOp ## From 83331bf12373c54eb7f3af7ca4e50b91a22d23ba Mon Sep 17 00:00:00 2001 From: Michael Suo Date: Tue, 29 Oct 2019 11:58:51 -0700 Subject: [PATCH 31/64] don't overspecify required python version (#28842) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28842 We don't care which python version, and github actions has changed the versions available, breaking our CI. So just pin it to 3-something to make it more future proof Test Plan: Imported from OSS Differential Revision: D18205349 Pulled By: suo fbshipit-source-id: bf260dc29a138dd8bf8c85081a182aae298fe86d --- .github/workflows/lint.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 1f2515fe6d3e9..67f0cf8dbbe09 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -13,7 +13,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -28,7 +28,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -69,7 +69,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -84,7 +84,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master @@ -99,7 +99,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Fetch PyTorch uses: actions/checkout@master @@ -174,7 +174,7 @@ jobs: - name: Setup Python uses: actions/setup-python@v1 with: - python-version: 3.7.4 + python-version: 3.x architecture: x64 - name: Checkout PyTorch uses: actions/checkout@master From e57a11977314e56e69a16fe192a380b536205b1f Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Tue, 29 Oct 2019 11:59:59 -0700 Subject: [PATCH 32/64] Remove autograd copy_ specific isFloatingPoint (#28279) Summary: Remove autograd copy_ specific isFloatingPoint and use c10's isFloatingType (and isComplexType). Before this, .to or .copy_ would drop requires_grad for bfloat16 as the floating types were only considered to be double, float, and half. Pull Request resolved: https://github.com/pytorch/pytorch/pull/28279 Differential Revision: D18176084 Pulled By: izdeby fbshipit-source-id: 8a005a6105e4a827be5c8163135e693a7daae4f4 --- test/test_autograd.py | 15 ++++++++++++++- torch/csrc/autograd/VariableTypeManual.cpp | 5 ++++- torch/csrc/autograd/VariableTypeUtils.h | 4 ---- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/test/test_autograd.py b/test/test_autograd.py index 72b6ae588197b..a5ba360f58909 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -39,7 +39,7 @@ mask_not_all_zeros, S) from common_device_type import (instantiate_device_type_tests, skipCUDAIfRocm, - onlyCUDA, dtypes, dtypesIfCUDA, + onlyCPU, onlyCUDA, dtypes, dtypesIfCUDA, deviceCountAtLeast, skipCUDAIfCudnnVersionLessThan) # load_tests from common_utils is used to automatically filter tests for @@ -3958,6 +3958,19 @@ def test_inputbuffer_add_multidevice(self, devices): output = input.to(device=devices[1]) + input.to(device=devices[1]) output.backward() + @onlyCPU + def test_copy_(self, device): + # At the time of writing this test, copy_ is not generated from native_functions.yaml + # there was a bug that bfloat16 was not recognized as floating. + x = torch.randn(10, device=device, requires_grad=True) + floating_dt = [dt for dt in torch.testing.get_all_dtypes() if dt.is_floating_point] + for dt in floating_dt: + y = torch.empty(10, device=device, dtype=dt) + y.copy_(x) + self.assertTrue(y.requires_grad) + z = x.to(torch.bfloat16) + self.assertTrue(z.requires_grad) + @onlyCUDA def test_cross_device_reentrant_autograd(self, device): # Output on gpu so that this task will be associated with the gpu thread diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp index be1fa2b076e95..165c0a06898bd 100644 --- a/torch/csrc/autograd/VariableTypeManual.cpp +++ b/torch/csrc/autograd/VariableTypeManual.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -149,7 +150,9 @@ Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking) { check_inplace(self); std::shared_ptr grad_fn; auto requires_grad = compute_requires_grad(self, src); - requires_grad &= isFloatingPoint(self.scalar_type()); + // currently, isFloatingType will return false for (floating) complex types, + // so this might have to be amended when they should be differentiable + requires_grad &= isFloatingType(self.scalar_type()); if (requires_grad) { grad_fn = std::make_shared(); grad_fn->set_next_edges(collect_next_edges(self, src)); diff --git a/torch/csrc/autograd/VariableTypeUtils.h b/torch/csrc/autograd/VariableTypeUtils.h index e42d80aedad13..fa503d47eaa87 100644 --- a/torch/csrc/autograd/VariableTypeUtils.h +++ b/torch/csrc/autograd/VariableTypeUtils.h @@ -80,10 +80,6 @@ inline void increment_version(Tensor & t) { as_variable_ref(t).bump_version(); } -inline bool isFloatingPoint(ScalarType s) { - return s == kFloat || s == kDouble || s == kHalf; -} - struct Flatten : IterArgs { Flatten(variable_list& out) : out(out) {} variable_list& out; From 7e8c48bff5cd2357fdbe5ed360e816dd5afe4ad6 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Tue, 29 Oct 2019 12:13:09 -0700 Subject: [PATCH 33/64] argmax for half datatype fix (#28787) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28787 Stack from [ghstack](https://github.com/ezyang/ghstack): * **#28787 argmax for half datatype fix** Test Plan: Imported from OSS Differential Revision: D18194420 Pulled By: pbelevich fbshipit-source-id: d2abec1ea8a9ce3a93aec5a2c5bba57d163197e6 --- aten/src/ATen/native/cuda/ReduceOpsKernel.cu | 82 +++++++++++++------- aten/src/ATen/test/CMakeLists.txt | 3 +- aten/src/ATen/test/reduce_ops_test.cpp | 24 ++++++ test/test_torch.py | 2 + 4 files changed, 84 insertions(+), 27 deletions(-) create mode 100644 aten/src/ATen/test/reduce_ops_test.cpp diff --git a/aten/src/ATen/native/cuda/ReduceOpsKernel.cu b/aten/src/ATen/native/cuda/ReduceOpsKernel.cu index 57b1adec8f9b0..ae53440bd5e0f 100644 --- a/aten/src/ATen/native/cuda/ReduceOpsKernel.cu +++ b/aten/src/ATen/native/cuda/ReduceOpsKernel.cu @@ -141,50 +141,80 @@ void or_kernel_cuda(TensorIterator& iter) { }), false); } -template +template void max_values_kernel_cuda_impl(TensorIterator& iter) { gpu_reduce_kernel( - iter, func_wrapper ([]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { - return (THCNumerics::isnan(a) || a > b) ? a : b; - }), at::numeric_limits::lower_bound()); + iter, func_wrapper ([]GPU_LAMBDA(acc_t a, acc_t b) -> acc_t { + return (THCNumerics::isnan(a) || a > b) ? a : b; + }), at::numeric_limits::lower_bound()); } -template +template void min_values_kernel_cuda_impl(TensorIterator& iter) { gpu_reduce_kernel( - iter, func_wrapper ([]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t { - return (THCNumerics::isnan(a) || a < b) ? a : b; - }), at::numeric_limits::upper_bound()); + iter, func_wrapper ([]GPU_LAMBDA(acc_t a, acc_t b) -> acc_t { + return (THCNumerics::isnan(a) || a < b) ? a : b; + }), at::numeric_limits::upper_bound()); } void max_values_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(), "max_values_cuda", [&]() { - max_values_kernel_cuda_impl(iter); - }); + if (iter.dtype(1) == kHalf) { + max_values_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(), "max_values_cuda", [&]() { + max_values_kernel_cuda_impl(iter); + }); + } } void min_values_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(), "min_values_cuda", [&]() { - min_values_kernel_cuda_impl(iter); - }); + if (iter.dtype(1) == kHalf) { + min_values_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(), "min_values_cuda", [&]() { + min_values_kernel_cuda_impl(iter); + }); + } } +template +void argmax_kernel_cuda_impl(TensorIterator& iter) { + gpu_reduce_kernel( + iter, + ArgMaxOps{}, + thrust::pair(at::numeric_limits::lower_bound(), 0)); +}; + +template +void argmin_kernel_cuda_impl(TensorIterator& iter) { + gpu_reduce_kernel( + iter, + ArgMinOps{}, + thrust::pair(at::numeric_limits::upper_bound(), 0)); +}; + void argmax_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmax_cuda", [&]() { - gpu_reduce_kernel( - iter, - ArgMaxOps{}, - thrust::pair(at::numeric_limits::lower_bound(), 0)); - }); + if (iter.dtype(1) == kHalf) { + // Instead of implementing is_nan and warp_shfl_down + // we can convert halves to float and do all the operations in float + argmax_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmax_cuda", [&]() { + argmax_kernel_cuda_impl(iter); + }); + } } void argmin_kernel_cuda(TensorIterator& iter) { - AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmin_cuda", [&]() { - gpu_reduce_kernel( - iter, - ArgMinOps{}, - thrust::pair(at::numeric_limits::upper_bound(), 0)); - }); + if (iter.dtype(1) == kHalf) { + // Instead of implementing is_nan and warp_shfl_down + // we can convert halves to float and do all the operations in float + argmin_kernel_cuda_impl(iter); + } else { + AT_DISPATCH_ALL_TYPES(iter.dtype(1), "argmin_cuda", [&]() { + argmin_kernel_cuda_impl(iter); + }); + } } REGISTER_DISPATCH(std_var_stub, &std_var_kernel_cuda); diff --git a/aten/src/ATen/test/CMakeLists.txt b/aten/src/ATen/test/CMakeLists.txt index 373c3f9205308..e9647f15a4ef8 100644 --- a/aten/src/ATen/test/CMakeLists.txt +++ b/aten/src/ATen/test/CMakeLists.txt @@ -31,7 +31,8 @@ list(APPEND ATen_CPU_TEST_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/memory_overlapping_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpu_generator_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/pow_test.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/variant_test.cpp) + ${CMAKE_CURRENT_SOURCE_DIR}/variant_test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/reduce_ops_test.cpp) list(APPEND ATen_CUDA_TEST_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/cuda_integer_divider_test.cu diff --git a/aten/src/ATen/test/reduce_ops_test.cpp b/aten/src/ATen/test/reduce_ops_test.cpp new file mode 100644 index 0000000000000..429761b563a07 --- /dev/null +++ b/aten/src/ATen/test/reduce_ops_test.cpp @@ -0,0 +1,24 @@ +#include + +#include +#include + +using namespace at; + +TEST(ReduceOpsTest, MaxValuesAndMinValues) { + const int W = 10; + const int H = 10; + if (hasCUDA()) { + for (const auto dtype : {kHalf, kFloat, kDouble, kShort, kInt, kLong}) { + auto a = at::rand({H, W}, TensorOptions(kCUDA).dtype(at::kHalf)); + ASSERT_FLOAT_EQ( + a.max_values(c10::IntArrayRef{0, 1}).item(), + a.max().item() + ); + ASSERT_FLOAT_EQ( + a.min_values(c10::IntArrayRef{0, 1}).item(), + a.min().item() + ); + } + } +} diff --git a/test/test_torch.py b/test/test_torch.py index 723f0197f4350..af10e97dae7c0 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -8272,6 +8272,8 @@ def test_dim_reduction(self, device): torch.int64, torch.int32, torch.int16] + if self.device_type == 'cuda': # 'cpu' and 'xla' do not support half + types.append(torch.half) # This won't test for 256bit instructions, since we usually # only work on 1 cacheline (1024bit) at a time and these From 6b5bfd4cfc3c281e61f63e0b62d919cc600041e4 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 12:28:06 -0700 Subject: [PATCH 34/64] Make inserted child module names unique (#27237) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27237 Making inserted observer module and wrapper module names unique Test Plan: test_jit.py Imported from OSS Differential Revision: D18182917 fbshipit-source-id: 77aa5997fbf024c73085866550372b5e68ad9ae1 --- test/test_jit.py | 49 +++++++++++++------------- torch/csrc/jit/passes/quantization.cpp | 23 ++++++++---- 2 files changed, 42 insertions(+), 30 deletions(-) diff --git a/test/test_jit.py b/test/test_jit.py index b5d1e5fbcdbc3..0195d3b7bcdf9 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -1005,23 +1005,23 @@ def forward(self, x): } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) assert len([x for x, _ in m._c._get_modules() - if x.startswith('observer_for_')]) == 0, \ + if x.startswith('_observer_')]) == 0, \ 'Expected to have 0 observer submodules' - FileCheck().check_not('ClassType = prim::GetAttr[name="observer_for_') \ + FileCheck().check_not('ClassType = prim::GetAttr[name="_observer_') \ .check('ClassType = prim::GetAttr[name="conv"](%self)') \ .check_next('Tensor = prim::CallMethod[name="forward"]') \ - .check_not('ClassType = prim::GetAttr[name="observer_for_') \ + .check_not('ClassType = prim::GetAttr[name="_observer_') \ .run(str(get_forward_graph(m._c))) assert len([x for x, _ in m._c._get_module('conv')._get_modules() - if x.startswith('observer_for_')]) == 3, \ + if x.startswith('_observer_')]) == 3, \ 'Expected to have 3 observer submodules' - FileCheck().check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .check('Tensor = aten::conv2d') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .run(str(m._c._get_module("conv")._get_method('conv2d_forward').graph)) @_tmp_donotuse_dont_inline_everything @@ -1044,17 +1044,17 @@ def forward(self, x): return self.sub(self.conv(x)) def check_observed(s): - FileCheck().check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ - .check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ + .check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') \ .run(str(s)) def check_not_observed(s): - FileCheck().check_not('ClassType = prim::GetAttr[name="observer_for_') \ - .check_not('prim::CallMethod[name="forward"](%observer_for_') \ + FileCheck().check_not('ClassType = prim::GetAttr[name="_observer_') \ + .check_not('prim::CallMethod[name="forward"](%_observer_') \ .run(str(s)) m = torch.jit.script(M()) @@ -1118,15 +1118,15 @@ def test_module(module, relu_call, num_observers): } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) assert len([x for x, _ in m._c._get_modules() - if x.startswith('observer_for_')]) == num_observers, \ + if x.startswith('_observer_')]) == num_observers, \ 'Expected to have ' + str(num_observers) + ' observer submodules' c = FileCheck().check('ClassType = prim::GetAttr[name="conv"]') \ .check_next('prim::CallMethod[name="forward"]') \ - .check_not('ClassType = prim::GetAttr[name="observer_for_') \ + .check_not('ClassType = prim::GetAttr[name="_observer_') \ .check(relu_call) if num_observers == 1: - c = c.check('ClassType = prim::GetAttr[name="observer_for_') \ - .check_next('prim::CallMethod[name="forward"](%observer_for_') + c = c.check('ClassType = prim::GetAttr[name="_observer_') \ + .check_next('prim::CallMethod[name="forward"](%_observer_') c.run(str(get_forward_graph(m._c))) # TODO: add checks for conv and relu later, graph looks correct but this pr # has too many changes already @@ -1153,8 +1153,9 @@ def forward(self, x): weight=weight_observer._c) } torch._C._jit_pass_insert_observers(m._c, "forward", qconfig_dict, True) - assert m._c._get_module('conv')._get_module('observer_for_input.1')._get_attribute('dtype') != \ - m._c._get_module('conv')._get_module('observer_for_weight.1')._get_attribute('dtype') + dtypes = set([obs._get_attribute('dtype') for x, obs in m._c._get_module('conv')._get_modules() + if x.startswith('_observer_')]) + assert len(dtypes) == 2, 'Expected to have 2 different types of dtype' @_tmp_donotuse_dont_inline_everything def test_insert_quant_dequant(self): diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp index 3fdd3d8519e82..55cce3168f71e 100644 --- a/torch/csrc/jit/passes/quantization.cpp +++ b/torch/csrc/jit/passes/quantization.cpp @@ -145,6 +145,11 @@ class InsertObserversHelper { // Values that are the output of GetAttr[name="bias"] and they // will be propagated through the function call hierarchy std::unordered_set bias_values_; + // Unique id generator for observer module, used for generating + // unique observer names when we insert observer module, we + // record the current unique id used to avoid incrementing from 0 + // every time to find a unique id. + int uid_ = 0; }; // Clone observer module and add it to the original module, @@ -165,12 +170,15 @@ Node* InsertObserversHelper::insertObserverFor( } else { observer_module = std::get<0>(qconfig); } - std::string observer_name = "observer_for_" + v->debugName(); script::Module observer = observer_module.clone(); + std::string observer_name = "_observer_" + std::to_string(uid_++); + while (module.find_module(observer_name)) { + observer_name = "_observer_" + std::to_string(uid_++); + } module.register_module(observer_name, observer); // Get handle of observer module Node* observer_instance = g->create(c10::prim::GetAttr); - // self.observer_for_v + // self._observer_v observer_instance->addInput(g->inputs()[0]); observer_instance->s_(c10::attr::name, observer_name); observer_instance->output()->setDebugName(observer_name); @@ -426,7 +434,7 @@ c10::optional findObserverName(Value* v) { u.user->s(attr::name) == "forward") { auto module_instance = u.user->inputs().at(0); if (module_instance->node()->kind() == prim::GetAttr && - module_instance->node()->s(attr::name).find("observer_for_") != + module_instance->node()->s(attr::name).find("_observer_") != std::string::npos) { return module_instance->node()->s(attr::name); } @@ -544,7 +552,7 @@ c10::optional QuantizeHelper::findChildModuleToQuantize( child_instance->node()->kind() == prim::GetAttr, "Child instance should come from GetAttr."); auto child_module_name = child_instance->node()->s(attr::name); - if (child_module_name.find("observer_for_") == std::string::npos) { + if (child_module_name.find("_observer_") == std::string::npos) { auto child_module = module_.find_module(child_module_name); TORCH_INTERNAL_ASSERT( child_module, @@ -1029,8 +1037,11 @@ graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, } auto w_quant_val = match_vmap.at(vmap.at("w_quant")); // unique name for the module based on %w_quant - auto module_name = - module_name_prefix + std::to_string(w_quant_val->unique()); + int uid = 0; + auto module_name = module_name_prefix + std::to_string(uid++); + while (module.find_module(module_name)) { + module_name_prefix + std::to_string(uid++); + } module.register_module(module_name, wrapper_module); // Add GetAttr of the packed module From 22d70bc1ecfafc9fc867da9c21bb081e4c02ab45 Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Tue, 29 Oct 2019 13:00:21 -0700 Subject: [PATCH 35/64] Add OfflineTensor Summary: OfflineTensor will be a shell to just carry the shape and dtype. No data will be stored. This should help us plumb through the onnxifi process. Test Plan: ``` buck test caffe2/caffe2/fb/opt:onnxifi_with_offline_tensor_test ``` Reviewed By: ChunliF, zrphercule Differential Revision: D18187208 fbshipit-source-id: 57c70f6f9897a5fc66580c81295db108acd03862 --- caffe2/onnx/offline_tensor.cc | 86 +++++++++++++++++++++++++++++++++++ caffe2/onnx/offline_tensor.h | 51 +++++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 caffe2/onnx/offline_tensor.cc create mode 100644 caffe2/onnx/offline_tensor.h diff --git a/caffe2/onnx/offline_tensor.cc b/caffe2/onnx/offline_tensor.cc new file mode 100644 index 0000000000000..73327cbf8eafa --- /dev/null +++ b/caffe2/onnx/offline_tensor.cc @@ -0,0 +1,86 @@ +#include "caffe2/onnx/offline_tensor.h" + +namespace caffe2 { + +namespace { +// These constants need to be aligned with onnxifi.h +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT16 = 10; +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT32 = 1; +constexpr uint64_t kONNXIFI_DATATYPE_UINT8 = 2; +constexpr uint64_t kONNXIFI_DATATYPE_INT32 = 6; +constexpr uint64_t kONNXIFI_DATATYPE_INT8 = 3; +constexpr uint64_t kONNXIFI_DATATYPE_INT64 = 7; +constexpr uint64_t kONNXIFI_DATATYPE_INT16 = 5; +constexpr uint64_t kONNXIFI_DATATYPE_UINT16 = 4; +} // namespace + +CAFFE_KNOWN_TYPE(OfflineTensor); + +bool OfflineTensorShapeFunctions::IsSameMetaType(TypeIdentifier id) { + return id == TypeMeta::Id(); +} + +TypeIdentifier OfflineTensorShapeFunctions::GetTypeMetaId() { + return TypeMeta::Id(); +} + +TypeMeta OfflineTensorShapeFunctions::GetExternalTensorType(const void* c) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + + return offline_tensor->shape_tensor.dtype(); +} + +vector OfflineTensorShapeFunctions::GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + return GetTensorInfo(&(offline_tensor->shape_tensor), capacity, device); +} + +void OfflineTensorShapeFunctions::SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* /* unused */, + std::vector>* /* unused */, + ExternalTensorDescriptor* desc) { + const auto& offline_tensor = blob->template Get(); + const Tensor& shape_tensor = offline_tensor.shape_tensor; + + if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT64; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT16; + } else { + CAFFE_THROW("Unsupported tensor type: ", shape_tensor.dtype().name()); + } + desc->buffer = 0; + + desc->quantizationParams = 0; + desc->quantizationAxis = 0; + + // Set up dim and shape + const auto shape = shape_tensor.sizes(); + desc->dimensions = shape.size(); + shapes->emplace_back(shape.cbegin(), shape.cend()); + desc->shape = shapes->back().data(); +} + +REGISTER_EXTERNAL_TENSOR_FUNCTIONS( + (TypeMeta::Id()), + OfflineTensorShapeFunctions); +} // namespace caffe2 diff --git a/caffe2/onnx/offline_tensor.h b/caffe2/onnx/offline_tensor.h new file mode 100644 index 0000000000000..a1707055068aa --- /dev/null +++ b/caffe2/onnx/offline_tensor.h @@ -0,0 +1,51 @@ +#pragma once + +#include +#include +#include + +namespace caffe2 { + +struct OfflineTensor { + // A shell tensor to record shape and dtype + Tensor shape_tensor{CPU}; + + void setShapeAndType( + const std::vector& sizes, + at::Device device, + caffe2::TypeMeta data_type) { + shape_tensor.unsafeGetTensorImpl()->set_storage( + at::Storage::create_legacy(device, data_type)); + shape_tensor.Resize(sizes); + CHECK(!shape_tensor.storage_initialized()); + CHECK(shape_tensor.dtype_initialized()); + } +}; + +class OfflineTensorShapeFunctions : public ExternalTensorFunctionsBase { + public: + explicit OfflineTensorShapeFunctions() : ExternalTensorFunctionsBase() {} + ~OfflineTensorShapeFunctions() override {} + bool isQuantized() const override { + return false; + } + bool IsSameMetaType(TypeIdentifier id) override; + void SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* all_scales, + std::vector>* all_offsets, + ExternalTensorDescriptor* desc) override; + void LoadInfoOfBlob( + const Blob* /* unused */, + std::vector* /* unused */, + std::vector* /* unused */, + uint32_t* /* unused */) override {} + TypeIdentifier GetTypeMetaId() override; + TypeMeta GetExternalTensorType(const void* c) override; + vector GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) override; +}; +} // namespace caffe2 From d8c368bd62455f2c9d0bf036640a87281db0face Mon Sep 17 00:00:00 2001 From: Dylan Bespalko Date: Tue, 29 Oct 2019 13:35:29 -0700 Subject: [PATCH 36/64] CPU-strided-complex support for compare and pointwise ops (#28735) Summary: In-tree changes to pytorch to support complex numbers are being submitted here. Out-of-tree support for complex numbers is here: [pytorch-cpu-strided-complex extension](https://gitlab.com/pytorch-complex/pytorch-cpu-strided-complex) These changes optimize complex Vec256 math kernels so that are within 2X real number performance on average. [Benchmarks are here](https://docs.google.com/spreadsheets/d/17pObcrSTpV4BOOX9FYf1vIX3QUlEgQhLvL1IBEyJyzs/edit#gid=0) Changes so far: - [x] Added complex support for eq, neq, max, and min ops. - max/min ops need to compare the absolute value for complex numbers (using zabs). - [x] Added complex support for is_nonzero and where. - where op compares the absolute value for complex numbers (using zabs). - [x] Added complex support for linear interp and and pointwise ops. - [x] Added complex support for check_convert and Linspace/Logspace. - std::complex does not support ++operator. - All compilers from clang, g++, c++ on aarch64, x86 produce the same assembly code when using `+=1' instead of `++`. [example for loop](https://godbolt.org/z/O6NW_p) - [x] Added complex support for log, log2, log10. - [x] Optimized Vec256 operators using various logarithmic identities. - `asin()`, `acos()`, `atan()` is optimized using a `ln()` identity. - `sqrt()` is optimized by splitting the computation into real and imag parts. - several `_mm256_mul_pd` are avoided by using `_mm256_xor_pd` ops instead. - [x] Added complex support for pow. - exp is cast to `std::complex`. - no special optimization is added when the `exp` is real because the `std::pow()` operator expects a std::complex number. Pull Request resolved: https://github.com/pytorch/pytorch/pull/28735 Differential Revision: D18170691 Pulled By: ezyang fbshipit-source-id: 6f167398e112cdeab02fcfde8b543cb6629c865a --- aten/src/ATen/cpu/vec256/vec256_base.h | 12 ++ .../ATen/cpu/vec256/vec256_complex_double.h | 109 +++++++++++++----- .../ATen/cpu/vec256/vec256_complex_float.h | 109 +++++++++++++----- aten/src/ATen/native/BinaryOps.cpp | 2 +- aten/src/ATen/native/RangeFactories.cpp | 8 +- aten/src/ATen/native/TensorCompare.cpp | 4 +- aten/src/ATen/native/TensorFactories.cpp | 2 +- aten/src/ATen/native/cpu/BinaryOpsKernel.cpp | 8 +- aten/src/ATen/native/cpu/CrossKernel.cpp | 2 +- aten/src/ATen/native/cpu/LerpKernel.cpp | 14 ++- .../ATen/native/cpu/PointwiseOpsKernel.cpp | 4 +- aten/src/ATen/native/cpu/PowKernel.cpp | 60 +++++++++- .../ATen/native/cpu/TensorCompareKernel.cpp | 11 +- aten/src/ATen/native/cpu/UnaryOpsKernel.cpp | 2 +- 14 files changed, 260 insertions(+), 87 deletions(-) diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h index 64063b3da2b92..01481783e7048 100644 --- a/aten/src/ATen/cpu/vec256/vec256_base.h +++ b/aten/src/ATen/cpu/vec256/vec256_base.h @@ -256,9 +256,21 @@ struct Vec256 { Vec256 log1p() const { return map(std::log1p); } + template ::value, int>::type = 0> Vec256 log2() const { + // other_t_log2 is for SFINAE and clarity. Make sure it is not changed. + static_assert(std::is_same::value, "other_t_log2 must be T"); return map(std::log2); } + template ::value, int>::type = 0> + Vec256 log2() const { + // complex_t_log2 is for SFINAE and clarity. Make sure it is not changed. + static_assert(std::is_same::value, "complex_t_log2 must be T"); + const T log_2 = T(std::log(2.0)); + return Vec256(map(std::log))/Vec256(log_2); + } Vec256 ceil() const { return map(at::native::ceil_impl); } diff --git a/aten/src/ATen/cpu/vec256/vec256_complex_double.h b/aten/src/ATen/cpu/vec256/vec256_complex_double.h index db56dfd408480..4ac48edac0575 100644 --- a/aten/src/ATen/cpu/vec256/vec256_complex_double.h +++ b/aten/src/ATen/cpu/vec256/vec256_complex_double.h @@ -13,7 +13,7 @@ namespace { #if defined(__AVX__) && !defined(_MSC_VER) -template <> class Vec256> { +template <> class Vec256> { private: __m256d values; public: @@ -135,26 +135,58 @@ template <> class Vec256> { const __m256d imag_mask = _mm256_castsi256_pd(_mm256_setr_epi64x(0x0000000000000000, 0xFFFFFFFFFFFFFFFF, 0x0000000000000000, 0xFFFFFFFFFFFFFFFF)); return _mm256_and_pd(values, imag_mask); - } + } Vec256> imag() const { return _mm256_permute_pd(imag_(), 0x05); //b a } __m256d conj_() const { - const __m256d conj_mask = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); - return _mm256_mul_pd(values, conj_mask); //a -b + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + return _mm256_xor_pd(values, sign_mask); // a -b } Vec256> conj() const { return conj_(); } - Vec256> acos() const { - return map(std::acos); + Vec256> log() const { + // Most trigonomic ops use the log() op to improve complex number performance. + return map(std::log); + } + Vec256> log2() const { + const __m256d log2_ = _mm256_set1_pd(std::log(2)); + return _mm256_div_pd(log(), log2_); + } + Vec256> log10() const { + const __m256d log10_ = _mm256_set1_pd(std::log(10)); + return _mm256_div_pd(log(), log10_); + } + Vec256> log1p() const { + AT_ERROR("not supported for complex numbers"); } Vec256> asin() const { - return map(std::asin); + // asin(x) + // = -i*ln(iz + sqrt(1 -z^2)) + // = -i*ln((ai - b) + sqrt(1 - (a + bi)*(a + bi))) + // = -i*ln((-b + ai) + sqrt(1 - (a**2 - b**2) - 2*abi)) + const __m256d one = _mm256_set1_pd(1); + + auto conj = conj_(); + auto b_a = _mm256_permute_pd(conj, 0x05); //-b a + auto ab = _mm256_mul_pd(conj, b_a); //-ab -ab + auto im = _mm256_add_pd(ab, ab); //-2ab -2ab + + auto val_2 = _mm256_mul_pd(values, values); // a*a b*b + auto re = _mm256_hsub_pd(val_2, _mm256_permute_pd(val_2, 0x05)); // a*a-b*b b*b-a*a + re = _mm256_sub_pd(one, re); + + auto root = Vec256(_mm256_blend_pd(re, im, 0x0A)).sqrt(); //sqrt(re + i*im) + auto ln = Vec256(_mm256_add_pd(b_a, root)).log(); //ln(iz + sqrt()) + return Vec256(_mm256_permute_pd(ln.values, 0x05)).conj(); //-i*ln() } - Vec256> atan() const { - return map(std::atan); + Vec256> acos() const { + // acos(x) = pi/2 - asin(x) + const __m256d pi_2 = _mm256_setr_pd(M_PI/2, 0.0, M_PI/2, 0.0); + return _mm256_sub_pd(pi_2, asin()); } + Vec256> atan() const; Vec256> atan2(const Vec256> &b) const { AT_ERROR("not supported for complex numbers"); } @@ -170,18 +202,6 @@ template <> class Vec256> { Vec256> expm1() const { AT_ERROR("not supported for complex numbers"); } - Vec256> log() const { - return map(std::log); - } - Vec256> log2() const { - AT_ERROR("not supported for complex numbers"); - } - Vec256> log10() const { - return map(std::log10); - } - Vec256> log1p() const { - AT_ERROR("not supported for complex numbers"); - } Vec256> sin() const { return map(std::sin); } @@ -217,14 +237,32 @@ template <> class Vec256> { return _mm256_round_pd(values, (_MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC)); } Vec256> sqrt() const { - return map(std::sqrt); + // sqrt(a + bi) + // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 + b**2) - a)i] + // = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i] + + const __m256d scalar = _mm256_set1_pd(std::sqrt(2)/2); //sqrt(2)/2 sqrt(2)/2 + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + auto sign = _mm256_and_pd(values, sign_mask); + auto factor = _mm256_or_pd(scalar, sign); + + auto a_a = _mm256_xor_pd(_mm256_movedup_pd(values), sign_mask); // a -a + auto res_re_im = _mm256_sqrt_pd(_mm256_add_pd(abs_(), a_a)); // sqrt(abs + a) sqrt(abs - a) + return _mm256_mul_pd(factor, res_re_im); } Vec256> reciprocal() const; Vec256> rsqrt() const { - return map([](const std::complex &x) { return (std::complex)(1)/std::sqrt(x); }); + return sqrt().reciprocal(); } Vec256> pow(const Vec256> &exp) const { - AT_ERROR("not supported for complex numbers"); + __at_align32__ std::complex x_tmp[size()]; + __at_align32__ std::complex y_tmp[size()]; + store(x_tmp); + exp.store(y_tmp); + for (int i = 0; i < size(); i++) { + x_tmp[i] = std::pow(x_tmp[i], y_tmp[i]); + } + return loadu(x_tmp); } // Comparison using the _CMP_**_OQ predicate. // `O`: get false if an operand is NaN @@ -259,11 +297,11 @@ template <> Vec256> inline operator-(const Vec256 Vec256> inline operator*(const Vec256> &a, const Vec256> &b) { //(a + bi) * (c + di) = (ac - bd) + (ad + bc)i - const __m256d neg = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); auto ac_bd = _mm256_mul_pd(a, b); //ac bd auto d_c = _mm256_permute_pd(b, 0x05); //d c - d_c = _mm256_mul_pd(neg, d_c); //d -c + d_c = _mm256_xor_pd(sign_mask, d_c); //d -c auto ad_bc = _mm256_mul_pd(a, d_c); //ad -bc auto ret = _mm256_hsub_pd(ac_bd, ad_bc); //ac - bd ad + bc @@ -274,11 +312,11 @@ template <> Vec256> inline operator/(const Vec256> Vec256>::reciprocal() const{ //re + im*i = (a + bi) / (c + di) //re = (ac + bd)/abs_2() = c/abs_2() //im = (bc - ad)/abs_2() = d/abs_2() - const __m256d neg = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0); - auto c_d = _mm256_mul_pd(neg, values); //c -d + const __m256d sign_mask = _mm256_setr_pd(0.0, -0.0, 0.0, -0.0); + auto c_d = _mm256_xor_pd(sign_mask, values); //c -d return _mm256_div_pd(c_d, abs_2_()); } +Vec256> Vec256>::atan() const { + // atan(x) = i/2 * ln((i + z)/(i - z)) + const __m256d i = _mm256_setr_pd(0.0, 1.0, 0.0, 1.0); + const Vec256 i_half = _mm256_setr_pd(0.0, 0.5, 0.0, 0.5); + + auto sum = Vec256(_mm256_add_pd(i, values)); // a 1+b + auto sub = Vec256(_mm256_sub_pd(i, values)); // -a 1-b + auto ln = (sum/sub).log(); // ln((i + z)/(i - z)) + return i_half*ln; // i/2*ln() +} + template <> Vec256> inline maximum(const Vec256>& a, const Vec256>& b) { auto abs_a = a.abs_2_(); diff --git a/aten/src/ATen/cpu/vec256/vec256_complex_float.h b/aten/src/ATen/cpu/vec256/vec256_complex_float.h index 2e388914f2b93..2149e79099ecd 100644 --- a/aten/src/ATen/cpu/vec256/vec256_complex_float.h +++ b/aten/src/ATen/cpu/vec256/vec256_complex_float.h @@ -171,26 +171,58 @@ template <> class Vec256> { const __m256 imag_mask = _mm256_castsi256_ps(_mm256_setr_epi32(0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF, 0x00000000, 0xFFFFFFFF)); return _mm256_and_ps(values, imag_mask); - } + } Vec256> imag() const { return _mm256_permute_ps(imag_(), 0x55); //b a } __m256 conj_() const { - const __m256 conj_mask = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); - return _mm256_mul_ps(values, conj_mask); //a -b + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + return _mm256_xor_ps(values, sign_mask); // a -b } Vec256> conj() const { return conj_(); } - Vec256> acos() const { - return map(std::acos); + Vec256> log() const { + // Most trigonomic ops use the log() op to improve complex number performance. + return map(std::log); + } + Vec256> log2() const { + const __m256 log2_ = _mm256_set1_ps(std::log(2)); + return _mm256_div_ps(log(), log2_); + } + Vec256> log10() const { + const __m256 log10_ = _mm256_set1_ps(std::log(10)); + return _mm256_div_ps(log(), log10_); + } + Vec256> log1p() const { + AT_ERROR("not supported for complex numbers"); } Vec256> asin() const { - return map(std::asin); + // asin(x) + // = -i*ln(iz + sqrt(1 -z^2)) + // = -i*ln((ai - b) + sqrt(1 - (a + bi)*(a + bi))) + // = -i*ln((-b + ai) + sqrt(1 - (a**2 - b**2) - 2*abi)) + const __m256 one = _mm256_set1_ps(1); + + auto conj = conj_(); + auto b_a = _mm256_permute_ps(conj, 0x55); //-b a + auto ab = _mm256_mul_ps(conj, b_a); //-ab -ab + auto im = _mm256_add_ps(ab, ab); //-2ab -2ab + + auto val_2 = _mm256_mul_ps(values, values); // a*a b*b + auto re = _mm256_hsub_ps(val_2, _mm256_permute_ps(val_2, 0x55)); // a*a-b*b b*b-a*a + re = _mm256_sub_ps(one, re); + + auto root = Vec256(_mm256_blend_ps(re, im, 0xAA)).sqrt(); //sqrt(re + i*im) + auto ln = Vec256(_mm256_add_ps(b_a, root)).log(); //ln(iz + sqrt()) + return Vec256(_mm256_permute_ps(ln.values, 0x55)).conj(); //-i*ln() } - Vec256> atan() const { - return map(std::atan); + Vec256> acos() const { + // acos(x) = pi/2 - asin(x) + const __m256 pi_2 = _mm256_setr_ps(M_PI/2, 0.0, M_PI/2, 0.0, M_PI/2, 0.0, M_PI/2, 0.0); + return _mm256_sub_ps(pi_2, asin()); } + Vec256> atan() const; Vec256> atan2(const Vec256> &b) const { AT_ERROR("not supported for complex numbers"); } @@ -206,18 +238,6 @@ template <> class Vec256> { Vec256> expm1() const { AT_ERROR("not supported for complex numbers"); } - Vec256> log() const { - return map(std::log); - } - Vec256> log2() const { - AT_ERROR("not supported for complex numbers"); - } - Vec256> log10() const { - return map(std::log10); - } - Vec256> log1p() const { - AT_ERROR("not supported for complex numbers"); - } Vec256> sin() const { return map(std::sin); } @@ -253,14 +273,32 @@ template <> class Vec256> { return _mm256_round_ps(values, (_MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC)); } Vec256> sqrt() const { - return map(std::sqrt); + // sqrt(a + bi) + // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 + b**2) - a)i] + // = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i] + + const __m256 scalar = _mm256_set1_ps(std::sqrt(2)/2); //sqrt(2)/2 sqrt(2)/2 + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + auto sign = _mm256_and_ps(values, sign_mask); + auto factor = _mm256_or_ps(scalar, sign); + + auto a_a = _mm256_xor_ps(_mm256_moveldup_ps(values), sign_mask); // a -a + auto res_re_im = _mm256_sqrt_ps(_mm256_add_ps(abs_(), a_a)); // sqrt(abs + a) sqrt(abs - a) + return _mm256_mul_ps(factor, res_re_im); } Vec256> reciprocal() const; Vec256> rsqrt() const { - return map([](const std::complex &x) { return (std::complex)(1)/std::sqrt(x); }); + return sqrt().reciprocal(); } Vec256> pow(const Vec256> &exp) const { - AT_ERROR("not supported for complex numbers"); + __at_align32__ std::complex x_tmp[size()]; + __at_align32__ std::complex y_tmp[size()]; + store(x_tmp); + exp.store(y_tmp); + for (int i = 0; i < size(); i++) { + x_tmp[i] = std::pow(x_tmp[i], y_tmp[i]); + } + return loadu(x_tmp); } // Comparison using the _CMP_**_OQ predicate. // `O`: get false if an operand is NaN @@ -295,11 +333,11 @@ template <> Vec256> inline operator-(const Vec256 Vec256> inline operator*(const Vec256> &a, const Vec256> &b) { //(a + bi) * (c + di) = (ac - bd) + (ad + bc)i - const __m256 neg = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); auto ac_bd = _mm256_mul_ps(a, b); //ac bd auto d_c = _mm256_permute_ps(b, 0x55); //d c - d_c = _mm256_mul_ps(neg, d_c); //d -c + d_c = _mm256_xor_ps(sign_mask, d_c); //d -c auto ad_bc = _mm256_mul_ps(a, d_c); //ad -bc auto ret = _mm256_hsub_ps(ac_bd, ad_bc); //ac - bd ad + bc @@ -310,11 +348,11 @@ template <> Vec256> inline operator/(const Vec256> Vec256>::reciprocal() const { //re + im*i = (a + bi) / (c + di) //re = (ac + bd)/abs_2() = c/abs_2() //im = (bc - ad)/abs_2() = d/abs_2() - const __m256 neg = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0); - auto c_d = _mm256_mul_ps(neg, values); //c -d + const __m256 sign_mask = _mm256_setr_ps(0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0); + auto c_d = _mm256_xor_ps(sign_mask, values); //c -d return _mm256_div_ps(c_d, abs_2_()); } +Vec256> Vec256>::atan() const { + // atan(x) = i/2 * ln((i + z)/(i - z)) + const __m256 i = _mm256_setr_ps(0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0); + const Vec256 i_half = _mm256_setr_ps(0.0, 0.5, 0.0, 0.5, 0.0, 0.5, 0.0, 0.5); + + auto sum = Vec256(_mm256_add_ps(i, values)); // a 1+b + auto sub = Vec256(_mm256_sub_ps(i, values)); // -a 1-b + auto ln = (sum/sub).log(); // ln((i + z)/(i - z)) + return i_half*ln; // i/2*ln() +} + template <> Vec256> inline maximum(const Vec256>& a, const Vec256>& b) { auto abs_a = a.abs_2_(); diff --git a/aten/src/ATen/native/BinaryOps.cpp b/aten/src/ATen/native/BinaryOps.cpp index 23c15ca45dd36..d3bb70d59e8fc 100644 --- a/aten/src/ATen/native/BinaryOps.cpp +++ b/aten/src/ATen/native/BinaryOps.cpp @@ -150,7 +150,7 @@ static Tensor wrapped_scalar_tensor(Scalar scalar) { static void check_convert(Scalar scalar, ScalarType scalarType) { // Validate that is possible to convert scalar to tensor dtype without overflow - AT_DISPATCH_ALL_TYPES_AND3(at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, scalarType, "check_convert", [&]{ + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, scalarType, "check_convert", [&]{ scalar.to(); }); } diff --git a/aten/src/ATen/native/RangeFactories.cpp b/aten/src/ATen/native/RangeFactories.cpp index 0210e998d9264..02262240a9382 100644 --- a/aten/src/ATen/native/RangeFactories.cpp +++ b/aten/src/ATen/native/RangeFactories.cpp @@ -21,14 +21,14 @@ Tensor& linspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps } else if (steps == 1) { r.fill_(start); } else { - AT_DISPATCH_FLOATING_TYPES(r.scalar_type(), "linspace_cpu", [&]() { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(r.scalar_type(), "linspace_cpu", [&]() { scalar_t scalar_start = start.to(); scalar_t scalar_end = end.to(); scalar_t *data_ptr = r.data_ptr(); scalar_t step = (scalar_end - scalar_start) / static_cast(steps - 1); at::parallel_for(0, steps, internal::GRAIN_SIZE, [&](int64_t p_begin, int64_t p_end) { scalar_t is = static_cast(p_begin); - for (int64_t i = p_begin; i < p_end; ++i, ++is) { + for (int64_t i = p_begin; i < p_end; ++i, is+=1) { //std::complex does not support ++operator data_ptr[i] = scalar_start + step*is; } }); @@ -54,7 +54,7 @@ Tensor& logspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps } else if (steps == 1) { r.fill_(std::pow(base, start.to())); } else { - AT_DISPATCH_FLOATING_TYPES(r.scalar_type(), "logspace_cpu", [&]() { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(r.scalar_type(), "logspace_cpu", [&]() { scalar_t scalar_base = static_cast(base); scalar_t scalar_start = start.to(); scalar_t scalar_end = end.to(); @@ -62,7 +62,7 @@ Tensor& logspace_cpu_out(Tensor& result, Scalar start, Scalar end, int64_t steps scalar_t step = (scalar_end - scalar_start) / static_cast(steps - 1); at::parallel_for(0, steps, internal::GRAIN_SIZE, [&](int64_t p_begin, int64_t p_end) { scalar_t is = static_cast(p_begin); - for (int64_t i = p_begin; i < p_end; ++i, ++is) { + for (int64_t i = p_begin; i < p_end; ++i, is+=1) { //std::complex does not support ++operator data_ptr[i]= std::pow(scalar_base, scalar_start + step*is); } }); diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp index 3916a82f00b0e..0e2132c9eca92 100644 --- a/aten/src/ATen/native/TensorCompare.cpp +++ b/aten/src/ATen/native/TensorCompare.cpp @@ -99,6 +99,8 @@ bool is_nonzero(const Tensor& self) { Scalar localScalar = self.item(); if (localScalar.isFloatingPoint()) { return localScalar.to() != 0; + } else if (localScalar.isComplex()) { + return localScalar.to>() != std::complex(0.0, 0.0); } else if (localScalar.isIntegral(false)){ return localScalar.to() != 0; } else if (localScalar.isBoolean()) { @@ -123,7 +125,7 @@ std::vector where(const Tensor& condition) { Tensor _s_where_cpu(const Tensor& condition, const Tensor& self, const Tensor& other) { Tensor ret = at::empty(self.sizes(), self.options()); - AT_DISPATCH_ALL_TYPES(ret.scalar_type(), "where_cpu", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(ret.scalar_type(), "where_cpu", [&] { where_cpu(ret, condition, self, other); }); return ret; diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp index d27285d244b43..07cbca3e57363 100644 --- a/aten/src/ATen/native/TensorFactories.cpp +++ b/aten/src/ATen/native/TensorFactories.cpp @@ -845,7 +845,7 @@ template Tensor tensor_cpu(ArrayRef values, const TensorOptions& options) { auto result = at::empty(values.size(), options); AT_ASSERT(result.is_contiguous()); - AT_DISPATCH_ALL_TYPES(result.scalar_type(), "tensor_cpu", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "tensor_cpu", [&] { std::copy(values.begin(), values.end(), result.template data_ptr()); }); return result; diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp index 7f20f8a5f9ff7..71c4fb71bc5c9 100644 --- a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp @@ -185,14 +185,14 @@ void ge_kernel(TensorIterator& iter) { void eq_kernel(TensorIterator& iter) { if (iter.dtype() == ScalarType::Bool) { - AT_DISPATCH_ALL_TYPES_AND2(kBool, kBFloat16, iter.input_dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kBool, kBFloat16, iter.input_dtype(), "eq_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> bool { return a == b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kBFloat16, iter.dtype(), "eq_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(kBFloat16, iter.dtype(), "eq_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> scalar_t { return a == b; @@ -203,14 +203,14 @@ void eq_kernel(TensorIterator& iter) { void ne_kernel(TensorIterator& iter) { if (iter.dtype() == ScalarType::Bool) { - AT_DISPATCH_ALL_TYPES_AND2(kBool, kBFloat16, iter.input_dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(kBool, kBFloat16, iter.input_dtype(), "ne_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> bool { return a != b; }); }); } else { - AT_DISPATCH_ALL_TYPES_AND(kBFloat16, iter.dtype(), "ne_cpu", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(kBFloat16, iter.dtype(), "ne_cpu", [&]() { cpu_kernel(iter, [=](scalar_t a, scalar_t b) -> scalar_t { return a != b; diff --git a/aten/src/ATen/native/cpu/CrossKernel.cpp b/aten/src/ATen/native/cpu/CrossKernel.cpp index 3e243403b5f49..0d405447bc1a9 100644 --- a/aten/src/ATen/native/cpu/CrossKernel.cpp +++ b/aten/src/ATen/native/cpu/CrossKernel.cpp @@ -65,7 +65,7 @@ static void apply_cross(Tensor& result, const Tensor& a, const Tensor& b, const } static void cross_kernel_impl(Tensor& result, const Tensor& a, const Tensor& b, const int64_t dim) { - AT_DISPATCH_ALL_TYPES(result.scalar_type(), "cross", [&]() { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "cross", [&]() { apply_cross(result, a, b, dim); }); } diff --git a/aten/src/ATen/native/cpu/LerpKernel.cpp b/aten/src/ATen/native/cpu/LerpKernel.cpp index 5e6d044030626..313f3e3bf9367 100644 --- a/aten/src/ATen/native/cpu/LerpKernel.cpp +++ b/aten/src/ATen/native/cpu/LerpKernel.cpp @@ -20,14 +20,15 @@ static void lerp_kernel_scalar( TORCH_CHECK(self.dtype() == end.dtype(), "expected dtype ", self.dtype(), " for `end` but got dtype ", end.dtype()); auto iter = TensorIterator::binary_op(ret, self, end, /*check_mem_overlap=*/true); - AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_scalar", [&] { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(ret.scalar_type(), "lerp_kernel_scalar", [&] { + using value_t = typename ztype::value_t; scalar_t weight_val = weight.to(); at::native::cpu_kernel( iter, [weight_val](scalar_t self_val, scalar_t end_val) { - return (weight_val < 0.5) + return (zabs(weight_val) < 0.5) ? self_val + weight_val * (end_val - self_val) - : end_val - (end_val - self_val) * (1 - weight_val); + : end_val - (end_val - self_val) * (scalar_t(1) - weight_val); }); }); } @@ -49,13 +50,14 @@ static void lerp_kernel_tensor( iter.add_input(end); iter.add_input(weights); iter.build(); - AT_DISPATCH_FLOATING_TYPES(ret.scalar_type(), "lerp_kernel_tensor", [&] { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(ret.scalar_type(), "lerp_kernel_tensor", [&] { + using value_t = typename ztype::value_t; at::native::cpu_kernel( iter, [](scalar_t self_val, scalar_t end_val, scalar_t weight_val) { - return (weight_val < 0.5) + return (zabs(weight_val) < 0.5) ? self_val + weight_val * (end_val - self_val) - : end_val - (end_val - self_val) * (1 - weight_val); + : end_val - (end_val - self_val) * (scalar_t(1) - weight_val); }); }); } diff --git a/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp b/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp index e2cc77e1ffb4a..92d15953372ad 100644 --- a/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp @@ -12,7 +12,7 @@ namespace { static void addcmul_cpu_kernel(TensorIterator& iter, Scalar value) { ScalarType dtype = iter.dtype(0); - AT_DISPATCH_ALL_TYPES(dtype, "addcmul_cpu_out", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcmul_cpu_out", [&] { scalar_t scalar_val = value.to(); auto scalar_vec = Vec256(scalar_val); cpu_kernel_vec( @@ -30,7 +30,7 @@ static void addcmul_cpu_kernel(TensorIterator& iter, Scalar value) { static void addcdiv_cpu_kernel(TensorIterator& iter, Scalar value) { ScalarType dtype = iter.dtype(0); - AT_DISPATCH_ALL_TYPES(dtype, "addcdiv_cpu_out", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(dtype, "addcdiv_cpu_out", [&] { scalar_t scalar_val = value.to(); auto scalar_vec = Vec256(scalar_val); cpu_kernel_vec( diff --git a/aten/src/ATen/native/cpu/PowKernel.cpp b/aten/src/ATen/native/cpu/PowKernel.cpp index 1109063966dac..085d57bac7732 100644 --- a/aten/src/ATen/native/cpu/PowKernel.cpp +++ b/aten/src/ATen/native/cpu/PowKernel.cpp @@ -11,8 +11,8 @@ namespace at { namespace native { namespace { void pow_tensor_tensor_kernel(TensorIterator& iter) { - if (isFloatingType(iter.dtype())) { - AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "pow", [&]() { + if (isFloatingType(iter.dtype()) || isComplexType(iter.dtype())) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(iter.dtype(), "pow", [&]() { using Vec = Vec256; cpu_kernel_vec(iter, [=](scalar_t base, scalar_t exp) -> scalar_t { @@ -91,6 +91,62 @@ void pow_tensor_scalar_kernel(TensorIterator& iter, Scalar exp_scalar) { ); } }); + } else if (isComplexType(iter.dtype())) { + const auto exp = exp_scalar.to>(); + // Floating types allow AVX2 vector optimizations for pow/sqrt/rsqrt: + AT_DISPATCH_COMPLEX_TYPES(iter.dtype(), "pow", [&]() { + using Vec = Vec256; + if (exp == 0.5) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return std::sqrt(base); + }, + [](Vec base) -> Vec { return base.sqrt(); } + ); + } else if (exp == 2.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return base * base; + }, + [](Vec base) -> Vec { return base * base; } + ); + } else if (exp == 3.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return base * base * base; + }, + [](Vec base) -> Vec { return base * base * base; } + ); + } else if (exp == -0.5) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / std::sqrt(base); + }, + [](Vec base) -> Vec { return base.rsqrt(); } + ); + } else if (exp == -1.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / base; + }, + [](Vec base) -> Vec { return base.reciprocal(); } + ); + } else if (exp == -2.0) { + cpu_kernel_vec(iter, + [](scalar_t base) -> scalar_t { + return scalar_t(1.0) / (base * base); + }, + [](Vec base) -> Vec { return (base * base).reciprocal(); } + ); + } else { + cpu_kernel_vec(iter, + [=](scalar_t base) -> scalar_t { + return std::pow(base, scalar_t(exp)); + }, + [=](Vec base) -> Vec { return base.pow(scalar_t(exp)); } // std::pow cannot accept mixed complex data types. + ); + } + }); } else { // Integral types do not allow AVX2 vector optimizations for pow/sqrt/rsqrt. // Trying to implement pow/sqrt/rsqrt as loop in vec256_int.h does not allow diff --git a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp index aa7a5ebeaabc6..fb7d93cd8237b 100644 --- a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp +++ b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp @@ -8,6 +8,7 @@ #include #include #include +#include namespace at { namespace native { namespace { @@ -34,6 +35,8 @@ struct Reduction { } } int64_t batch = numel / (n * stride); + using value_t = typename ztype::value_t; + value_t (*zabs_)(scalar_t) = zabs; if (stride == 1) { parallel_for(0, batch, 1, [=](int64_t begin, int64_t end) { for (int64_t b = begin; b < end; b++) { @@ -42,7 +45,7 @@ struct Reduction { index_t result_index = 0; for (int64_t k = 0; k < n; k++) { scalar_t value = data[k]; - bool cmp = greater ? (result > value) : (result < value); + bool cmp = greater ? (zabs_(result) > zabs_(value)) : (zabs_(result) < zabs_(value)); result = cmp ? result : value; result_index = cmp ? result_index : k; if (_isnan(result)) { @@ -63,7 +66,7 @@ struct Reduction { index_t result_index = 0; for (int64_t k = 0; k < n; k++) { scalar_t value = data[k * stride]; - bool cmp = greater ? (result > value) : (result < value); + bool cmp = greater ? (zabs_(result) > zabs_(value)) : (zabs_(result) < zabs_(value)); result = cmp ? result : value; result_index = cmp ? result_index : k; if (_isnan(result)) { @@ -83,7 +86,7 @@ static void max_kernel_impl( Tensor& max_indices, const Tensor& self, c10::optional dim) { - AT_DISPATCH_ALL_TYPES_AND(ScalarType::Bool, self.scalar_type(), "max", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(ScalarType::Bool, self.scalar_type(), "max", [&] { Reduction::apply(max, max_indices, self, dim, true); }); } @@ -93,7 +96,7 @@ static void min_kernel_impl( Tensor& min_indices, const Tensor& self, c10::optional dim) { - AT_DISPATCH_ALL_TYPES_AND(ScalarType::Bool, self.scalar_type(), "min", [&] { + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(ScalarType::Bool, self.scalar_type(), "min", [&] { Reduction::apply(min, min_indices, self, dim, false); }); } diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp index 33e1b5574b31b..a87fa83df4e2d 100644 --- a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp @@ -428,7 +428,7 @@ IMPLEMENT_COMPLEX_KERNEL(FLOATING, floor) IMPLEMENT_COMPLEX_KERNEL(FLOATING, log) IMPLEMENT_COMPLEX_KERNEL(FLOATING, log10) IMPLEMENT_FLOAT_KERNEL(FLOATING, log1p) -IMPLEMENT_FLOAT_KERNEL(FLOATING, log2) +IMPLEMENT_COMPLEX_KERNEL(FLOATING, log2) IMPLEMENT_COMPLEX_KERNEL(FLOATING, round) IMPLEMENT_COMPLEX_KERNEL(FLOATING, sin) // IMPLEMENT_FLOAT_KERNEL(FLOATING, sinh) From e33b4b6761d2cf9745e94af1308d98ebd168a3ab Mon Sep 17 00:00:00 2001 From: Will Feng Date: Tue, 29 Oct 2019 14:13:37 -0700 Subject: [PATCH 37/64] Use c10::variant-based enums for Reduction Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27942 Test Plan: Imported from OSS Differential Revision: D18202857 Pulled By: yf225 fbshipit-source-id: 0303ce2508e3b7665c6a91ae270a7d0ef0e45900 --- test/cpp/api/enum.cpp | 9 +- test/cpp/api/functional.cpp | 6 +- test/cpp/api/modules.cpp | 6 +- torch/csrc/api/include/torch/enum.h | 89 ++++++++++++++++++- .../api/include/torch/nn/functional/loss.h | 64 +++++++++---- .../api/include/torch/nn/functional/padding.h | 2 +- .../include/torch/nn/functional/upsampling.h | 2 +- torch/csrc/api/include/torch/nn/init.h | 1 - .../csrc/api/include/torch/nn/options/loss.h | 63 ++++++++----- torch/csrc/api/src/enum.cpp | 2 + torch/csrc/api/src/nn/modules/embedding.cpp | 4 +- torch/csrc/api/src/nn/modules/loss.cpp | 2 +- torch/csrc/api/src/nn/modules/upsampling.cpp | 2 +- 13 files changed, 196 insertions(+), 56 deletions(-) diff --git a/test/cpp/api/enum.cpp b/test/cpp/api/enum.cpp index e490ef9b22dfa..8f6c36c9f543c 100644 --- a/test/cpp/api/enum.cpp +++ b/test/cpp/api/enum.cpp @@ -1,6 +1,5 @@ #include -#include #include #include @@ -10,7 +9,7 @@ v = torch::k##name; \ std::string pretty_print_name("k"); \ pretty_print_name.append(#name); \ - ASSERT_EQ(c10::visit(torch::enumtype::enum_name{}, v), pretty_print_name); \ + ASSERT_EQ(torch::enumtype::get_enum_name(v), pretty_print_name); \ } TEST(EnumTest, AllEnums) { @@ -39,7 +38,9 @@ TEST(EnumTest, AllEnums) { torch::enumtype::kArea, torch::enumtype::kSum, torch::enumtype::kMean, - torch::enumtype::kMax + torch::enumtype::kMax, + torch::enumtype::kNone, + torch::enumtype::kBatchMean > v; TORCH_ENUM_PRETTY_PRINT_TEST(Linear) @@ -67,4 +68,6 @@ TEST(EnumTest, AllEnums) { TORCH_ENUM_PRETTY_PRINT_TEST(Sum) TORCH_ENUM_PRETTY_PRINT_TEST(Mean) TORCH_ENUM_PRETTY_PRINT_TEST(Max) + TORCH_ENUM_PRETTY_PRINT_TEST(None) + TORCH_ENUM_PRETTY_PRINT_TEST(BatchMean) } diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp index b9de3651ee54d..bd1a7623ff6e7 100644 --- a/test/cpp/api/functional.cpp +++ b/test/cpp/api/functional.cpp @@ -131,7 +131,7 @@ TEST_F(FunctionalTest, SoftMarginLossNoReduction) { auto input = torch::tensor({2., 4., 1., 3.}, torch::requires_grad()); auto target = torch::tensor({-1., 1., 1., -1.}, torch::kFloat); auto output = - F::soft_margin_loss(input, target, torch::Reduction::None); + F::soft_margin_loss(input, target, torch::kNone); auto expected = torch::tensor({2.1269281, 0.01814993, 0.3132617, 3.0485873}, torch::kFloat); auto s = output.sum(); s.backward(); @@ -144,7 +144,7 @@ TEST_F(FunctionalTest, MultiLabelSoftMarginLossWeightedNoReduction) { auto input = torch::tensor({{0., 2., 2., 0.}, {2., 1., 0., 1.}}, torch::requires_grad()); auto target = torch::tensor({{0., 0., 1., 0.}, {1., 0., 1., 1.}}, torch::kFloat); auto weight = torch::tensor({0.1, 0.6, 0.4, 0.8}, torch::kFloat); - auto options = MultiLabelSoftMarginLossOptions().reduction(torch::Reduction::None).weight(weight); + auto options = MultiLabelSoftMarginLossOptions().reduction(torch::kNone).weight(weight); auto output = F::multilabel_soft_margin_loss(input, target, options); auto expected = torch::tensor({0.4876902, 0.3321295}, torch::kFloat); @@ -491,7 +491,7 @@ TEST_F(FunctionalTest, MultiLabelMarginLossNoReduction) { auto input = torch::tensor({{0.1, 0.2, 0.4, 0.8}}, torch::requires_grad()); auto target = torch::tensor({{3, 0, -1, 1}}, torch::kLong); auto output = F::multilabel_margin_loss( - input, target, torch::Reduction::None); + input, target, torch::kNone); auto expected = torch::tensor({0.8500}, torch::kFloat); auto s = output.sum(); s.backward(); diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 20f692b4e7d25..89b277f9c3822 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -1189,7 +1189,7 @@ TEST_F(ModulesTest, MultiLabelMarginLossDefaultOptions) { } TEST_F(ModulesTest, MultiLabelMarginLossNoReduction) { - MultiLabelMarginLoss loss(torch::Reduction::None); + MultiLabelMarginLoss loss(torch::kNone); auto input = torch::tensor({{0.1, 0.2, 0.4, 0.8}}, torch::requires_grad()); auto target = torch::tensor({{3, 0, -1, 1}}, torch::kLong); auto output = loss->forward(input, target); @@ -1255,7 +1255,7 @@ TEST_F(ModulesTest, MultiLabelSoftMarginLossDefaultOptions) { } TEST_F(ModulesTest, SoftMarginLossNoReduction) { - SoftMarginLoss loss(torch::Reduction::None); + SoftMarginLoss loss(torch::kNone); auto input = torch::tensor({2., 4., 1., 3.}, torch::requires_grad()); auto target = torch::tensor({-1., 1., 1., -1.}, torch::kFloat); auto output = loss->forward(input, target); @@ -1271,7 +1271,7 @@ TEST_F(ModulesTest, MultiLabelSoftMarginLossWeightedNoReduction) { auto input = torch::tensor({{0., 2., 2., 0.}, {2., 1., 0., 1.}}, torch::requires_grad()); auto target = torch::tensor({{0., 0., 1., 0.}, {1., 0., 1., 1.}}, torch::kFloat); auto weight = torch::tensor({0.1, 0.6, 0.4, 0.8}, torch::kFloat); - auto options = MultiLabelSoftMarginLossOptions().reduction(torch::Reduction::None).weight(weight); + auto options = MultiLabelSoftMarginLossOptions().reduction(torch::kNone).weight(weight); MultiLabelSoftMarginLoss loss = MultiLabelSoftMarginLoss(options); auto output = loss->forward(input, target); auto expected = torch::tensor({0.4876902, 0.3321295}, torch::kFloat); diff --git a/torch/csrc/api/include/torch/enum.h b/torch/csrc/api/include/torch/enum.h index 6ba65befc1522..e9534fc1dfb9e 100644 --- a/torch/csrc/api/include/torch/enum.h +++ b/torch/csrc/api/include/torch/enum.h @@ -2,6 +2,8 @@ #include +#include +#include #include #include @@ -32,6 +34,64 @@ std::string operator()(const enumtype::k##name& v) const { \ return k + #name; \ } +// NOTE: Backstory on why we need the following two macros: +// +// Consider the following options class: +// +// ``` +// struct TORCH_API SomeOptions { +// typedef c10::variant reduction_t; +// SomeOptions(reduction_t reduction = torch::kMean) : reduction_(reduction) {} +// +// TORCH_ARG(reduction_t, reduction); +// }; +// ``` +// +// and the functional that uses it: +// +// ``` +// Tensor some_functional( +// const Tensor& input, +// SomeOptions options = {}) { +// ... +// } +// ``` +// +// Normally, we would expect this to work: +// +// `F::some_functional(input, torch::kNone)` +// +// However, it throws the following error instead: +// +// ``` +// error: could not convert ‘torch::kNone’ from ‘const torch::enumtype::kNone’ to ‘torch::nn::SomeOptions’ +// ``` +// +// To get around this problem, we explicitly provide the following constructors for `SomeOptions`: +// +// ``` +// SomeOptions(torch::enumtype::kNone reduction) : reduction_(torch::kNone) {} +// SomeOptions(torch::enumtype::kMean reduction) : reduction_(torch::kMean) {} +// SomeOptions(torch::enumtype::kSum reduction) : reduction_(torch::kSum) {} +// ``` +// +// so that the conversion from `torch::kNone` to `SomeOptions` would work. +// +// Note that we also provide the default constructor `SomeOptions() {}`, so that +// `SomeOptions options = {}` can work. +#define TORCH_OPTIONS_CTOR_VARIANT_ARG3(OPTIONS_NAME, ARG_NAME, TYPE1, TYPE2, TYPE3) \ +OPTIONS_NAME() {} \ +OPTIONS_NAME(torch::enumtype::TYPE1 ARG_NAME) : ARG_NAME##_(torch::TYPE1) {} \ +OPTIONS_NAME(torch::enumtype::TYPE2 ARG_NAME) : ARG_NAME##_(torch::TYPE2) {} \ +OPTIONS_NAME(torch::enumtype::TYPE3 ARG_NAME) : ARG_NAME##_(torch::TYPE3) {} + +#define TORCH_OPTIONS_CTOR_VARIANT_ARG4(OPTIONS_NAME, ARG_NAME, TYPE1, TYPE2, TYPE3, TYPE4) \ +OPTIONS_NAME() {} \ +OPTIONS_NAME(torch::enumtype::TYPE1 ARG_NAME) : ARG_NAME##_(torch::TYPE1) {} \ +OPTIONS_NAME(torch::enumtype::TYPE2 ARG_NAME) : ARG_NAME##_(torch::TYPE2) {} \ +OPTIONS_NAME(torch::enumtype::TYPE3 ARG_NAME) : ARG_NAME##_(torch::TYPE3) {} \ +OPTIONS_NAME(torch::enumtype::TYPE4 ARG_NAME) : ARG_NAME##_(torch::TYPE4) {} + TORCH_ENUM_DECLARE(Linear) TORCH_ENUM_DECLARE(Conv1D) TORCH_ENUM_DECLARE(Conv2D) @@ -57,10 +117,13 @@ TORCH_ENUM_DECLARE(Area) TORCH_ENUM_DECLARE(Sum) TORCH_ENUM_DECLARE(Mean) TORCH_ENUM_DECLARE(Max) +TORCH_ENUM_DECLARE(None) +TORCH_ENUM_DECLARE(BatchMean) namespace torch { namespace enumtype { -struct enum_name { + +struct _compute_enum_name { TORCH_ENUM_PRETTY_PRINT(Linear) TORCH_ENUM_PRETTY_PRINT(Conv1D) TORCH_ENUM_PRETTY_PRINT(Conv2D) @@ -86,6 +149,30 @@ struct enum_name { TORCH_ENUM_PRETTY_PRINT(Sum) TORCH_ENUM_PRETTY_PRINT(Mean) TORCH_ENUM_PRETTY_PRINT(Max) + TORCH_ENUM_PRETTY_PRINT(None) + TORCH_ENUM_PRETTY_PRINT(BatchMean) }; + +template +std::string get_enum_name(V variant_enum) { + return c10::visit(enumtype::_compute_enum_name{}, variant_enum); +} + +template +at::Reduction::Reduction reduction_get_enum(V variant_enum) { + if (c10::get_if(&variant_enum)) { + return at::Reduction::None; + } else if (c10::get_if(&variant_enum)) { + return at::Reduction::Mean; + } else if (c10::get_if(&variant_enum)) { + return at::Reduction::Sum; + } else { + TORCH_CHECK( + false, + get_enum_name(variant_enum), " is not a valid value for reduction"); + return at::Reduction::END; + } +} + } // namespace enumtype } // namespace torch diff --git a/torch/csrc/api/include/torch/nn/functional/loss.h b/torch/csrc/api/include/torch/nn/functional/loss.h index 7700fd18ef526..19068e54871bd 100644 --- a/torch/csrc/api/include/torch/nn/functional/loss.h +++ b/torch/csrc/api/include/torch/nn/functional/loss.h @@ -10,21 +10,30 @@ inline Tensor l1_loss( const Tensor& input, const Tensor& target, const L1LossOptions& options = {}) { - return torch::l1_loss(input, target, options.reduction()); + return torch::l1_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor kl_div( const Tensor& input, const Tensor& target, const KLDivLossOptions& options = {}) { - return torch::kl_div(input, target, options.reduction()); + return torch::kl_div( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor mse_loss( const Tensor& input, const Tensor& target, const MSELossOptions& options = {}) { - return torch::mse_loss(input, target, options.reduction()); + return torch::mse_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor binary_cross_entropy( @@ -32,7 +41,10 @@ inline Tensor binary_cross_entropy( const Tensor& target, const BCELossOptions& options = {}) { return torch::binary_cross_entropy( - input, target, options.weight(), options.reduction()); + input, + target, + options.weight(), + enumtype::reduction_get_enum(options.reduction())); } inline Tensor hinge_embedding_loss( @@ -40,7 +52,10 @@ inline Tensor hinge_embedding_loss( const Tensor& target, const HingeEmbeddingLossOptions& options = {}) { return torch::hinge_embedding_loss( - input, target, options.margin(), options.reduction()); + input, + target, + options.margin(), + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multi_margin_loss( @@ -58,7 +73,7 @@ inline Tensor multi_margin_loss( options.p(), options.margin(), options.weight(), - options.reduction() + enumtype::reduction_get_enum(options.reduction()) ); } @@ -68,21 +83,31 @@ inline Tensor cosine_embedding_loss( const Tensor& target, const CosineEmbeddingLossOptions& options) { return torch::cosine_embedding_loss( - input1, input2, target, options.margin(), options.reduction()); + input1, + input2, + target, + options.margin(), + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multilabel_margin_loss( const Tensor& input, const Tensor& target, const MultiLabelMarginLossOptions& options = {}) { - return torch::multilabel_margin_loss(input, target, options.reduction()); + return torch::multilabel_margin_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor soft_margin_loss( const Tensor& input, const Tensor& target, const SoftMarginLossOptions& options = {}) { - return torch::soft_margin_loss(input, target, options.reduction()); + return torch::soft_margin_loss( + input, + target, + enumtype::reduction_get_enum(options.reduction())); } inline Tensor multilabel_soft_margin_loss( @@ -98,15 +123,18 @@ inline Tensor multilabel_soft_margin_loss( Tensor ret; - if (options.reduction() == torch::Reduction::None) { - ret = loss; - } else if (options.reduction() == torch::Reduction::Mean) { - ret = loss.mean(); - } else if (options.reduction() == torch::Reduction::Sum) { - ret = loss.sum(); + if (c10::get_if(&options.reduction())) { + ret = loss; + } else if (c10::get_if(&options.reduction())) { + ret = loss.mean(); + } else if (c10::get_if(&options.reduction())) { + ret = loss.sum(); } else { - ret = input; - TORCH_INTERNAL_ASSERT(true, options.reduction(), " is not valid"); + ret = input; + TORCH_INTERNAL_ASSERT( + false, + enumtype::get_enum_name(options.reduction()), + " is not valid"); } return ret; } @@ -124,7 +152,7 @@ inline Tensor triplet_margin_loss( options.p(), options.eps(), options.swap(), - options.reduction()); + enumtype::reduction_get_enum(options.reduction())); } } // namespace functional diff --git a/torch/csrc/api/include/torch/nn/functional/padding.h b/torch/csrc/api/include/torch/nn/functional/padding.h index 039055d23edc0..48715fd95e93c 100644 --- a/torch/csrc/api/include/torch/nn/functional/padding.h +++ b/torch/csrc/api/include/torch/nn/functional/padding.h @@ -36,7 +36,7 @@ inline Tensor pad(const Tensor& input, const PadOptions& options) { TORCH_CHECK( options.value() == 0, "Padding mode \"", - c10::visit(torch::enumtype::enum_name{}, options.mode()), + torch::enumtype::get_enum_name(options.mode()), "\" doesn't take in value argument"); if (input.dim() == 3) { TORCH_CHECK(options.pad().size() == 2, "3D tensors expect 2 values for padding"); diff --git a/torch/csrc/api/include/torch/nn/functional/upsampling.h b/torch/csrc/api/include/torch/nn/functional/upsampling.h index 165ce9593041b..2cf66c8699289 100644 --- a/torch/csrc/api/include/torch/nn/functional/upsampling.h +++ b/torch/csrc/api/include/torch/nn/functional/upsampling.h @@ -98,7 +98,7 @@ inline Tensor interpolate(const Tensor& input, InterpolateOptions options) { false, "Input Error: Only 3D, 4D and 5D input Tensors supported " "(got ", input.dim(), "D) for the modes: nearest | linear | bilinear | bicubic | trilinear " - "(got ", c10::visit(enumtype::enum_name{}, options.mode()), ")"); + "(got ", enumtype::get_enum_name(options.mode()), ")"); } } diff --git a/torch/csrc/api/include/torch/nn/init.h b/torch/csrc/api/include/torch/nn/init.h index eeb0627592565..954cbc8470def 100644 --- a/torch/csrc/api/include/torch/nn/init.h +++ b/torch/csrc/api/include/torch/nn/init.h @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include diff --git a/torch/csrc/api/include/torch/nn/options/loss.h b/torch/csrc/api/include/torch/nn/options/loss.h index 2d3437b862e70..077b640b4cd74 100644 --- a/torch/csrc/api/include/torch/nn/options/loss.h +++ b/torch/csrc/api/include/torch/nn/options/loss.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -9,60 +10,72 @@ namespace nn { /// Options for a L1 loss module. struct TORCH_API L1LossOptions { - L1LossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(L1LossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output. - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a KLDiv loss module. struct TORCH_API KLDivLossOptions { - KLDivLossOptions(Reduction::Reduction reduction = Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG4(KLDivLossOptions, reduction, kNone, kBatchMean, kSum, kMean) /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction); + /// ``'none'`` | ``'batchmean'`` | ``'sum'`` | ``'mean'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a MSE loss module. struct TORCH_API MSELossOptions { - MSELossOptions(Reduction::Reduction reduction = Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(MSELossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction); + /// ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a BCE loss module. struct TORCH_API BCELossOptions { + typedef c10::variant reduction_t; + /// A manual rescaling weight given to the loss of each batch element. TORCH_ARG(Tensor, weight) = {}; /// Specifies the reduction to apply to the output. - TORCH_ARG(Reduction::Reduction, reduction) = Reduction::Mean; + /// ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'`` + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a Hinge Embedding loss functional and module. struct TORCH_API HingeEmbeddingLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Default: 1 TORCH_ARG(double, margin) = 1.0; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-margin loss functional and module. struct TORCH_API MultiMarginLossOptions { + typedef c10::variant reduction_t; + /// Has a default value of :math:`1`. :math:`1` and :math:`2` /// are the only supported values. TORCH_ARG(int64_t, p) = 1; @@ -76,53 +89,59 @@ struct TORCH_API MultiMarginLossOptions { /// ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied, /// ``'mean'``: the sum of the output will be divided by the number of /// elements in the output, ``'sum'``: the output will be summed. Default: ``'mean'`` - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a Hinge Embedding loss functional and module. struct TORCH_API CosineEmbeddingLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Should be a number from -1 to 1, 0 /// to 0.5 is suggested. Default: 0.0 TORCH_ARG(double, margin) = 0.0; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-label margin loss functional and module. struct TORCH_API MultiLabelMarginLossOptions { - MultiLabelMarginLossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(MultiLabelMarginLossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a soft margin loss functional and module. struct TORCH_API SoftMarginLossOptions { - SoftMarginLossOptions(torch::Reduction::Reduction reduction = torch::Reduction::Mean) - : reduction_(reduction) {} + typedef c10::variant reduction_t; + + TORCH_OPTIONS_CTOR_VARIANT_ARG3(SoftMarginLossOptions, reduction, kNone, kMean, kSum) /// Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction); + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a multi-label soft margin loss functional and module. struct TORCH_API MultiLabelSoftMarginLossOptions { + typedef c10::variant reduction_t; + /// A manual rescaling weight given to each /// class. If given, it has to be a Tensor of size `C`. Otherwise, it is /// treated as if having all ones. @@ -132,13 +151,15 @@ struct TORCH_API MultiLabelSoftMarginLossOptions { /// 'none': no reduction will be applied, 'mean': the sum of the output will /// be divided by the number of elements in the output, 'sum': the output will /// be summed. Default: 'mean' - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; // ============================================================================ /// Options for a triplet-margin-Loss functional and module. struct TORCH_API TripletMarginLossOptions { + typedef c10::variant reduction_t; + /// Specifies the threshold for which the distance of a negative sample must /// reach in order to incur zero loss. Default: 1 TORCH_ARG(double, margin) = 1.0; @@ -150,7 +171,7 @@ struct TORCH_API TripletMarginLossOptions { /// E. Riba et al. Default: False TORCH_ARG(bool, swap) = false; /// Specifies the reduction to apply to the output. Default: Mean - TORCH_ARG(torch::Reduction::Reduction, reduction) = torch::Reduction::Mean; + TORCH_ARG(reduction_t, reduction) = torch::kMean; }; } // namespace nn diff --git a/torch/csrc/api/src/enum.cpp b/torch/csrc/api/src/enum.cpp index 374d2839af8c7..9cb76e0303a92 100644 --- a/torch/csrc/api/src/enum.cpp +++ b/torch/csrc/api/src/enum.cpp @@ -25,3 +25,5 @@ TORCH_ENUM_DEFINE(Area) TORCH_ENUM_DEFINE(Sum) TORCH_ENUM_DEFINE(Mean) TORCH_ENUM_DEFINE(Max) +TORCH_ENUM_DEFINE(None) +TORCH_ENUM_DEFINE(BatchMean) diff --git a/torch/csrc/api/src/nn/modules/embedding.cpp b/torch/csrc/api/src/nn/modules/embedding.cpp index a4a583f6ae6f6..81d2b34205b2a 100644 --- a/torch/csrc/api/src/nn/modules/embedding.cpp +++ b/torch/csrc/api/src/nn/modules/embedding.cpp @@ -150,7 +150,7 @@ torch::Tensor EmbeddingBagImpl::forward( !per_sample_weights_.defined() || c10::get_if(&options.mode()), "embedding_bag: per_sample_weights was not null. ", "per_sample_weights is only supported for mode='kSum' (got mode='", - c10::visit(torch::enumtype::enum_name{}, options.mode()), "').Please open a feature request on GitHub."); + torch::enumtype::get_enum_name(options.mode()), "').Please open a feature request on GitHub."); return std::get<0>( torch::embedding_bag( @@ -179,7 +179,7 @@ void EmbeddingBagImpl::pretty_print(std::ostream& stream) const { stream << ", sparse=" << std::boolalpha << options.sparse(); } if (!c10::get_if(&options.mode())) { - stream << ", mode=" << c10::visit(torch::enumtype::enum_name{}, options.mode()); + stream << ", mode=" << torch::enumtype::get_enum_name(options.mode()); } stream << ")"; } diff --git a/torch/csrc/api/src/nn/modules/loss.cpp b/torch/csrc/api/src/nn/modules/loss.cpp index 3eaf2a1ba0b5e..9ca998c616079 100644 --- a/torch/csrc/api/src/nn/modules/loss.cpp +++ b/torch/csrc/api/src/nn/modules/loss.cpp @@ -102,7 +102,7 @@ void MultiMarginLossImpl::reset() { void MultiMarginLossImpl::pretty_print(std::ostream& stream) const { stream << "torch::nn::MultiMarginLoss(p=" << options.p() << ", margin=" << options.margin() << ", weight=" << options.weight() - << ", reduction=" << options.reduction() << ")"; + << ", reduction=" << enumtype::get_enum_name(options.reduction()) << ")"; } Tensor MultiMarginLossImpl::forward(const Tensor& input, const Tensor& target) { diff --git a/torch/csrc/api/src/nn/modules/upsampling.cpp b/torch/csrc/api/src/nn/modules/upsampling.cpp index 8699f73ce600a..401cb16ce78b3 100644 --- a/torch/csrc/api/src/nn/modules/upsampling.cpp +++ b/torch/csrc/api/src/nn/modules/upsampling.cpp @@ -19,7 +19,7 @@ void UpsampleImpl::pretty_print(std::ostream& stream) const { } else { stream << "size=" << at::ArrayRef(options.size()); } - stream << ", mode=" << c10::visit(enumtype::enum_name{}, options.mode()) << ")"; + stream << ", mode=" << enumtype::get_enum_name(options.mode()) << ")"; } Tensor UpsampleImpl::forward(const Tensor& input) { From 4045d6c3fa89b73d36af9a8857242093d875f429 Mon Sep 17 00:00:00 2001 From: Michael Suo Date: Tue, 29 Oct 2019 14:14:41 -0700 Subject: [PATCH 38/64] Revert D18187208: Add OfflineTensor Test Plan: revert-hammer Differential Revision: D18187208 Original commit changeset: 57c70f6f9897 fbshipit-source-id: d13b089ceb645b2a9852923cd21a752a2f45a15b --- caffe2/onnx/offline_tensor.cc | 86 ----------------------------------- caffe2/onnx/offline_tensor.h | 51 --------------------- 2 files changed, 137 deletions(-) delete mode 100644 caffe2/onnx/offline_tensor.cc delete mode 100644 caffe2/onnx/offline_tensor.h diff --git a/caffe2/onnx/offline_tensor.cc b/caffe2/onnx/offline_tensor.cc deleted file mode 100644 index 73327cbf8eafa..0000000000000 --- a/caffe2/onnx/offline_tensor.cc +++ /dev/null @@ -1,86 +0,0 @@ -#include "caffe2/onnx/offline_tensor.h" - -namespace caffe2 { - -namespace { -// These constants need to be aligned with onnxifi.h -constexpr uint64_t kONNXIFI_DATATYPE_FLOAT16 = 10; -constexpr uint64_t kONNXIFI_DATATYPE_FLOAT32 = 1; -constexpr uint64_t kONNXIFI_DATATYPE_UINT8 = 2; -constexpr uint64_t kONNXIFI_DATATYPE_INT32 = 6; -constexpr uint64_t kONNXIFI_DATATYPE_INT8 = 3; -constexpr uint64_t kONNXIFI_DATATYPE_INT64 = 7; -constexpr uint64_t kONNXIFI_DATATYPE_INT16 = 5; -constexpr uint64_t kONNXIFI_DATATYPE_UINT16 = 4; -} // namespace - -CAFFE_KNOWN_TYPE(OfflineTensor); - -bool OfflineTensorShapeFunctions::IsSameMetaType(TypeIdentifier id) { - return id == TypeMeta::Id(); -} - -TypeIdentifier OfflineTensorShapeFunctions::GetTypeMetaId() { - return TypeMeta::Id(); -} - -TypeMeta OfflineTensorShapeFunctions::GetExternalTensorType(const void* c) { - const OfflineTensor* offline_tensor = - reinterpret_cast(c); - - return offline_tensor->shape_tensor.dtype(); -} - -vector OfflineTensorShapeFunctions::GetExternalTensorInfo( - const void* c, - size_t* capacity, - DeviceOption* device) { - const OfflineTensor* offline_tensor = - reinterpret_cast(c); - return GetTensorInfo(&(offline_tensor->shape_tensor), capacity, device); -} - -void OfflineTensorShapeFunctions::SetupExternalTensorDescriptor( - const Blob* blob, - std::vector>* shapes, - std::vector>* /* unused */, - std::vector>* /* unused */, - ExternalTensorDescriptor* desc) { - const auto& offline_tensor = blob->template Get(); - const Tensor& shape_tensor = offline_tensor.shape_tensor; - - if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_FLOAT32; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_INT32; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_INT8; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_UINT8; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_INT64; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_INT16; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_FLOAT16; - } else if (shape_tensor.template IsType()) { - desc->dataType = kONNXIFI_DATATYPE_UINT16; - } else { - CAFFE_THROW("Unsupported tensor type: ", shape_tensor.dtype().name()); - } - desc->buffer = 0; - - desc->quantizationParams = 0; - desc->quantizationAxis = 0; - - // Set up dim and shape - const auto shape = shape_tensor.sizes(); - desc->dimensions = shape.size(); - shapes->emplace_back(shape.cbegin(), shape.cend()); - desc->shape = shapes->back().data(); -} - -REGISTER_EXTERNAL_TENSOR_FUNCTIONS( - (TypeMeta::Id()), - OfflineTensorShapeFunctions); -} // namespace caffe2 diff --git a/caffe2/onnx/offline_tensor.h b/caffe2/onnx/offline_tensor.h deleted file mode 100644 index a1707055068aa..0000000000000 --- a/caffe2/onnx/offline_tensor.h +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace caffe2 { - -struct OfflineTensor { - // A shell tensor to record shape and dtype - Tensor shape_tensor{CPU}; - - void setShapeAndType( - const std::vector& sizes, - at::Device device, - caffe2::TypeMeta data_type) { - shape_tensor.unsafeGetTensorImpl()->set_storage( - at::Storage::create_legacy(device, data_type)); - shape_tensor.Resize(sizes); - CHECK(!shape_tensor.storage_initialized()); - CHECK(shape_tensor.dtype_initialized()); - } -}; - -class OfflineTensorShapeFunctions : public ExternalTensorFunctionsBase { - public: - explicit OfflineTensorShapeFunctions() : ExternalTensorFunctionsBase() {} - ~OfflineTensorShapeFunctions() override {} - bool isQuantized() const override { - return false; - } - bool IsSameMetaType(TypeIdentifier id) override; - void SetupExternalTensorDescriptor( - const Blob* blob, - std::vector>* shapes, - std::vector>* all_scales, - std::vector>* all_offsets, - ExternalTensorDescriptor* desc) override; - void LoadInfoOfBlob( - const Blob* /* unused */, - std::vector* /* unused */, - std::vector* /* unused */, - uint32_t* /* unused */) override {} - TypeIdentifier GetTypeMetaId() override; - TypeMeta GetExternalTensorType(const void* c) override; - vector GetExternalTensorInfo( - const void* c, - size_t* capacity, - DeviceOption* device) override; -}; -} // namespace caffe2 From aa949b12b330a49394fe85bdba4bf7063e673b89 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 14:16:10 -0700 Subject: [PATCH 39/64] InsertObservers (#27238) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27238 att Test Plan: test_jit.py insert_observers Imported from OSS Differential Revision: D18182914 fbshipit-source-id: 718300f259a2e38e730d3e7cc6308813fd1112af --- torch/csrc/jit/passes/quantization.cpp | 57 +++++++++++++------------- 1 file changed, 28 insertions(+), 29 deletions(-) diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp index 55cce3168f71e..cb2766757a0f6 100644 --- a/torch/csrc/jit/passes/quantization.cpp +++ b/torch/csrc/jit/passes/quantization.cpp @@ -329,36 +329,35 @@ void InsertObserversHelper::insertObservers( bias_values_.emplace(v); } } - if (v->node()->kind() == prim::CallMethod) { - // If we find a call to a method of a child module, - // we'll recursively insert observers for the forward function to - // the child module. - auto module_instance = v->node()->inputs()[0]; - auto module_method_name = v->node()->s(attr::name); - // TODO: looks like this block is not related to v? maybe we should - // move this outside - script::Module callee_module; - if (module_instance->node()->kind() == prim::GetAttr) { - auto child_module_name = module_instance->node()->s(attr::name); - auto child_module = module.find_module(child_module_name); - TORCH_INTERNAL_ASSERT( - child_module, - "Child module " + child_module_name + " does not exist"); - callee_module = child_module.value(); - } else { - TORCH_INTERNAL_ASSERT( - module_instance == graph->inputs()[0], - "We only support call method either on %self" - "or child instance in insert_observers_pass right now"); - callee_module = module; - } - auto method_graph = - callee_module.get_method(module_method_name).graph(); - propagateValues(v->node(), method_graph); - // Recursively insert observer for the forward function of child - // module - insertObservers(callee_module, module_method_name); + } + + if (n->kind() == prim::CallMethod) { + // If we find a call to a method of a child module, + // we'll recursively insert observers for the forward function to + // the child module. + auto module_instance = n->inputs()[0]; + auto module_method_name = n->s(attr::name); + script::Module callee_module; + if (module_instance->node()->kind() == prim::GetAttr) { + auto child_module_name = module_instance->node()->s(attr::name); + auto child_module = module.find_module(child_module_name); + TORCH_INTERNAL_ASSERT( + child_module, + "Child module " + child_module_name + " does not exist"); + callee_module = child_module.value(); + } else { + TORCH_INTERNAL_ASSERT( + module_instance == graph->inputs()[0], + "We only support call method either on %self" + "or child instance in insert_observers_pass right now"); + callee_module = module; } + auto method_graph = + callee_module.get_method(module_method_name).graph(); + propagateValues(n, method_graph); + // Recursively insert observer for the forward function of child + // module + insertObservers(callee_module, module_method_name); } for (Block* subblock : n->blocks()) { From 793e2914e4792086abbcf605c7a45f6c2608cb4c Mon Sep 17 00:00:00 2001 From: Huayu Li Date: Tue, 29 Oct 2019 14:53:42 -0700 Subject: [PATCH 40/64] Support full id interations (#28769) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28769 Support full id interaction. Test Plan: * unit-tests * buck test caffe2/caffe2/python/operator_test:pack_ops_test -- * buck test caffe2/caffe2/fb/dper/layer_models/tests:sparse_nn_attention_test -- test_sparse_nn_full_id * canary * apply SUM + full id with max_length as 20 on SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID: f147253340 (v1: f146340704) # of embeddings for this features is 20: {F219139816} The corresponding ops: two lookups, which is as expected. ``` op { input: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_0/Repeat_0/sparse_lookup/w" input: "feature_preproc/output_features:SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM:values" input: "feature_preproc/output_features:SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM:lengths" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_0/Repeat_0/sparse_lookup/output" name: "" type: "SparseLengthsSum" } op { input: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/sparse_lookup/w" input: "feature_preproc/output_features:SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM:values" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/sparse_lookup/output" name: "" type: "Gather" } op { input: "feature_preproc/output_features:SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM:lengths" input: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/sparse_lookup/output" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/PackSegments/embedding_packed" name: "" type: "PackSegments" arg { name: "max_length" i: 20 } arg { name: "pad_minf" i: 0 } } op { input: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/PackSegments/embedding_packed" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/Reshape/reshaped_record" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/Reshape/old_shape" name: "" type: "Reshape" arg { name: "shape" ints: -1 ints: 1280 } } op { input: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/Reshape/reshaped_record" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_0" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_1" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_2" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_3" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_4" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_5" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_6" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_7" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_8" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_9" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_10" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_11" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_12" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_13" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_14" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_15" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_16" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_17" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_18" output: "nested/dot/SPARSE_AD_MEDIA_XRAY_V11_TOPIC_ID_AUTO_FIRST_X_AUTO_UNIGRAM/Pool_Option_1/Repeat_0/full_id/split/output_19" name: "" type: "Split" arg { name: "axis" i: 1 } } ``` Reviewed By: chonglinsun Differential Revision: D18083520 fbshipit-source-id: f592fb7734dd4e3e712ba42dc0afcd0b32a4afa0 --- caffe2/operators/pack_segments.cc | 27 ++++++++------ caffe2/python/operator_test/pack_ops_test.py | 38 +++++++++++++++----- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/caffe2/operators/pack_segments.cc b/caffe2/operators/pack_segments.cc index 8e82ed0a6a86f..ee65d5f7373d6 100644 --- a/caffe2/operators/pack_segments.cc +++ b/caffe2/operators/pack_segments.cc @@ -33,11 +33,6 @@ bool PackSegmentsOp::DoRunWithType2() { total_length += l[i]; } if (max_length_ != -1) { - // Final dim must be greater than the max_length - CAFFE_ENFORCE_GE( - max_length_, - max_length, - "Pre-defined max_length should be greater than the real max_length"); max_length = max_length_; } @@ -89,13 +84,14 @@ bool PackSegmentsOp::DoRunWithType2() { const auto* d = static_cast(data.raw_data()); int64_t start = 0; for (int64_t i = 0; i < lengths.size(0); ++i) { + auto len = l[i] <= max_length ? l[i] : max_length; context_.CopyItemsSameDevice( data.dtype(), - l[i] * block_size, + len * block_size, d + block_bytesize * start, out + block_bytesize * max_length * i); if (return_presence_mask_) { - memset(presence_mask_data + max_length * i, (int)true, l[i]); + memset(presence_mask_data + max_length * i, (int)true, len); } start += l[i]; } @@ -128,7 +124,14 @@ bool UnpackSegmentsOp::DoRunWithType2() { } const T* l = lengths.template data(); - int64_t total_l = std::accumulate(l, l + lengths.size(0), (int64_t)0); + int64_t total_l = 0; + if (max_length_ != -1) { + for (int64_t i = 0; i < lengths.size(0); ++i) { + total_l += (int64_t)(l[i] <= max_length_ ? l[i] : max_length_); + } + } else { + total_l = std::accumulate(l, l + lengths.size(0), (int64_t)0); + } auto shape = data.sizes().vec(); CAFFE_ENFORCE_EQ( @@ -146,12 +149,16 @@ bool UnpackSegmentsOp::DoRunWithType2() { const auto* d = static_cast(data.raw_data()); int64_t start = 0; for (int64_t i = 0; i < lengths.size(0); ++i) { + auto len = l[i]; + if (max_length_ != -1 && l[i] > max_length_) { + len = max_length_; + } context_.CopyItemsSameDevice( data.dtype(), - l[i] * block_size, + len * block_size, d + block_bytesize * data.size(1) * i, out + block_bytesize * start); - start += l[i]; + start += len; } return true; } diff --git a/caffe2/python/operator_test/pack_ops_test.py b/caffe2/python/operator_test/pack_ops_test.py index b5ee13740e45d..5c83566ed07a5 100644 --- a/caffe2/python/operator_test/pack_ops_test.py +++ b/caffe2/python/operator_test/pack_ops_test.py @@ -21,13 +21,13 @@ def pack_segments_ref(lengths, data, max_length=max_length): constant_values = 0 if data.dtype.char == 'S': constant_values = '' - if max_length is not None: - assert(max_length > np.max(lengths)) - else: + if max_length is None: max_length = np.max(lengths) + start = 0 for idx in range(np.size(lengths)): - chunk = data[np.sum(lengths[:idx]):np.sum(lengths[:idx + 1])] - pad_length = max_length - lengths[idx] + len = lengths[idx] if max_length >= lengths[idx] else max_length + chunk = data[start : start + len] + pad_length = max_length - len # ((0, pad_length), (0, 0)) says add pad_length rows of padding # below chunk and 0 rows of padding elsewhere @@ -38,10 +38,12 @@ def pack_segments_ref(lengths, data, max_length=max_length): constant_values=constant_values ) ) + start += lengths[idx] result = [arr] if return_presence_mask: presence_arr = [] for length in lengths: + length = length if max_length >= length else max_length pad_length = max_length - length presence_arr.append( np.pad( @@ -57,9 +59,12 @@ def pack_segments_ref(lengths, data, max_length=max_length): @serial.given( num_seq=st.integers(10, 100), cell_size=st.integers(1, 10), + max_length_buffer=st.integers(-5, 5), **hu.gcs ) - def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): + def test_pack_with_max_length_ops( + self, num_seq, cell_size, max_length_buffer, gc, dc + ): # create data lengths = np.arange(num_seq, dtype=np.int32) + 1 num_cell = np.sum(lengths) @@ -74,7 +79,7 @@ def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): + "=" * 60 ) # run test - max_length = num_seq + 1 + max_length = num_seq + max_length_buffer op = core.CreateOperator( 'PackSegments', ['l', 'd'], ['t'], max_length=max_length) workspace.FeedBlob('l', lengths) @@ -105,7 +110,24 @@ def test_pack_with_max_length_ops(self, num_seq, cell_size, gc, dc): max_length=max_length, device_option=gc)) assert(workspace.FetchBlob('t').shape[1] == max_length) - assert((workspace.FetchBlob('newd') == workspace.FetchBlob('d')).all()) + + def _cal_unpacked_data(data): + if max_length >= num_seq: + return data + output = None + start = 0 + for i, length in enumerate(lengths): + new_len = max_length if length > max_length else length + chunk = data[start: start + new_len] + if output is None: + output = chunk + else: + output = np.concatenate((output, chunk), axis=0) + start += length + return output + + true_newd = _cal_unpacked_data(workspace.FetchBlob('d')) + assert((workspace.FetchBlob('newd') == true_newd).all()) @given( num_seq=st.integers(10, 500), From c9423c30b39f42f20ad753de5abef6bafb5ba2ac Mon Sep 17 00:00:00 2001 From: David Reiss Date: Tue, 29 Oct 2019 16:01:28 -0700 Subject: [PATCH 41/64] Add host build for pytorch_android (#27662) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27662 This adds a new gradle subproject at pytorch_android/host and tweaks the top-level build.gradle to only run some Android bits on the other projects. Referencing Java sources from inside the host directory feels a bit hacky, but getting host and Android Gradle builds to coexist in the same directory hit several roadblocks. We can try a bigger refactor to separate the Android-specific and non-Android-specific parts of the code, but that seems overkill at this point for 4 Java files. This doesn't actually run without some local changes to fbjni, but I want to get the files landed to avoid unnecessary merge conflicts. Test Plan: Imported from OSS Differential Revision: D18210317 Pulled By: dreiss fbshipit-source-id: dafb54dde06a5a9a48fc7b7065d9359c5c480795 --- android/build.gradle | 54 ++++++++++++----------- android/pytorch_android/host/build.gradle | 33 ++++++++++++++ android/settings.gradle | 4 +- 3 files changed, 65 insertions(+), 26 deletions(-) create mode 100644 android/pytorch_android/host/build.gradle diff --git a/android/build.gradle b/android/build.gradle index b29299d7a4a3b..70f33c4e0769c 100644 --- a/android/build.gradle +++ b/android/build.gradle @@ -1,33 +1,37 @@ -buildscript { - ext { - minSdkVersion = 21 - targetSdkVersion = 28 - compileSdkVersion = 28 - buildToolsVersion = '28.0.3' - - coreVersion = "1.2.0" - extJUnitVersion = "1.1.1" - runnerVersion = "1.2.0" - rulesVersion = "1.2.0" - junitVersion = "4.12" +allprojects { + if (name == "pytorch_host") { + return } - repositories { - google() - mavenLocal() - mavenCentral() - jcenter() - } + buildscript { + ext { + minSdkVersion = 21 + targetSdkVersion = 28 + compileSdkVersion = 28 + buildToolsVersion = '28.0.3' + + coreVersion = "1.2.0" + extJUnitVersion = "1.1.1" + runnerVersion = "1.2.0" + rulesVersion = "1.2.0" + junitVersion = "4.12" + } + + repositories { + google() + mavenLocal() + mavenCentral() + jcenter() + } - dependencies { - classpath 'com.android.tools.build:gradle:3.3.2' - classpath "com.jfrog.bintray.gradle:gradle-bintray-plugin:${GRADLE_BINTRAY_PLUGIN_VERSION}" - classpath "com.github.dcendents:android-maven-gradle-plugin:${ANDROID_MAVEN_GRADLE_PLUGIN_VERSION}" - classpath "org.jfrog.buildinfo:build-info-extractor-gradle:4.9.8" + dependencies { + classpath 'com.android.tools.build:gradle:3.3.2' + classpath "com.jfrog.bintray.gradle:gradle-bintray-plugin:${GRADLE_BINTRAY_PLUGIN_VERSION}" + classpath "com.github.dcendents:android-maven-gradle-plugin:${ANDROID_MAVEN_GRADLE_PLUGIN_VERSION}" + classpath "org.jfrog.buildinfo:build-info-extractor-gradle:4.9.8" + } } -} -allprojects { repositories { google() jcenter() diff --git a/android/pytorch_android/host/build.gradle b/android/pytorch_android/host/build.gradle new file mode 100644 index 0000000000000..7a096ad60d260 --- /dev/null +++ b/android/pytorch_android/host/build.gradle @@ -0,0 +1,33 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the Apache-2 license found in the +// LICENSE file in the root directory of this source tree. + +plugins { + id 'java-library' +} + +repositories { + mavenLocal() + jcenter() +} + +sourceSets { + main { + java.srcDir '../src/main/java' + } + test { + java { + srcDir '../src/androidTest/java' + exclude '**/PytorchInstrumented*' + } + resources.srcDirs = ["../src/androidTest/assets"] + } +} + +dependencies { + compileOnly 'com.google.code.findbugs:jsr305:3.0.1' + implementation 'com.facebook.soloader:nativeloader:0.8.0' + implementation 'com.facebook.fbjni:fbjni:0.0.3-SNAPSHOT' + testImplementation 'junit:junit:4.12' +} diff --git a/android/settings.gradle b/android/settings.gradle index 277b15c0afc99..99e442b2ae2f5 100644 --- a/android/settings.gradle +++ b/android/settings.gradle @@ -1,4 +1,6 @@ -include ':app', ':pytorch_android', ':fbjni', ':pytorch_android_torchvision' +include ':app', ':pytorch_android', ':fbjni', ':pytorch_android_torchvision', ':pytorch_host' project(':fbjni').projectDir = file('libs/fbjni_local') project(':pytorch_android_torchvision').projectDir = file('pytorch_android_torchvision') + +project(':pytorch_host').projectDir = file('pytorch_android/host') From 34455c68b5e43afeb398808dfc6806170fba014f Mon Sep 17 00:00:00 2001 From: David Reiss Date: Tue, 29 Oct 2019 16:01:28 -0700 Subject: [PATCH 42/64] Remove unnecessary BUILD_DIR variable in Android CMake build (#27663) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27663 CMake sets CMAKE_BINARY_DIR and creates it automatically. Using this allows us to use the -B command-line flag to CMake to specify an alternate output directory. Test Plan: Imported from OSS Differential Revision: D18210316 Pulled By: dreiss fbshipit-source-id: ba2f6bd4b881ddd00de73fe9c33d82645ad5495d --- android/pytorch_android/CMakeLists.txt | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt index 5ebae28741906..18665eaf16bf6 100644 --- a/android/pytorch_android/CMakeLists.txt +++ b/android/pytorch_android/CMakeLists.txt @@ -24,11 +24,8 @@ target_include_directories(pytorch PUBLIC ${libtorch_include_DIR} ) -set(BUILD_DIR ${CMAKE_SOURCE_DIR}/build) -file(MAKE_DIRECTORY ${BUILD_DIR}) - set(fbjni_DIR ${CMAKE_CURRENT_LIST_DIR}/../libs/fbjni/) -set(fbjni_BUILD_DIR ${BUILD_DIR}/fbjni/${ANDROID_ABI}) +set(fbjni_BUILD_DIR ${CMAKE_BINARY_DIR}/fbjni/${ANDROID_ABI}) add_subdirectory(${fbjni_DIR} ${fbjni_BUILD_DIR}) From 80e270a76cf4d182be95991d6af814e3555eb1bd Mon Sep 17 00:00:00 2001 From: David Reiss Date: Tue, 29 Oct 2019 16:01:28 -0700 Subject: [PATCH 43/64] Add support for host build to pytorch_android native code (#27664) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27664 When ANDROID_ABI is not set, find libtorch headers and libraries from the LIBTORCH_HOME build variable (which must be set by hand), place output under a "host" directory, and use dynamic linking instead of static. This doesn't actually work without some local changes to fbjni, but I want to get the changes landed to avoid unnecessary merge conflicts. Test Plan: Imported from OSS Differential Revision: D18210315 Pulled By: dreiss fbshipit-source-id: 685a62de3c2a0a52bec7fd6fb95113058456bac8 --- android/pytorch_android/CMakeLists.txt | 95 +++++++++++++++++--------- 1 file changed, 64 insertions(+), 31 deletions(-) diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt index 18665eaf16bf6..69b3f97aa2997 100644 --- a/android/pytorch_android/CMakeLists.txt +++ b/android/pytorch_android/CMakeLists.txt @@ -4,7 +4,19 @@ set(CMAKE_CXX_STANDARD 11) set(CMAKE_VERBOSE_MAKEFILE ON) set(pytorch_android_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp) -set(libtorch_include_DIR ${pytorch_android_DIR}/libtorch_include/${ANDROID_ABI}) + +if (ANDROID_ABI) + set(libtorch_include_DIR ${pytorch_android_DIR}/libtorch_include/${ANDROID_ABI}) + set(BUILD_SUBDIR ${ANDROID_ABI}) +else() + if (NOT LIBTORCH_HOME) + message(FATAL_ERROR + "pytorch_android requires LIBTORCH_HOME to be defined for non-Android builds.") + endif() + set(libtorch_include_DIR ${LIBTORCH_HOME}/include) + link_directories(${LIBTORCH_HOME}/lib) + set(BUILD_SUBDIR host) +endif() message(STATUS "libtorch dir:${libtorch_DIR}") @@ -25,36 +37,57 @@ target_include_directories(pytorch PUBLIC ) set(fbjni_DIR ${CMAKE_CURRENT_LIST_DIR}/../libs/fbjni/) -set(fbjni_BUILD_DIR ${CMAKE_BINARY_DIR}/fbjni/${ANDROID_ABI}) +set(fbjni_BUILD_DIR ${CMAKE_BINARY_DIR}/fbjni/${BUILD_SUBDIR}) add_subdirectory(${fbjni_DIR} ${fbjni_BUILD_DIR}) -function(import_static_lib name) - add_library(${name} STATIC IMPORTED) - set_property( - TARGET ${name} - PROPERTY IMPORTED_LOCATION - ${CMAKE_CURRENT_LIST_DIR}/src/main/jniLibs/${ANDROID_ABI}/${name}.a) -endfunction(import_static_lib) - -import_static_lib(libtorch) -import_static_lib(libc10) -import_static_lib(libnnpack) -import_static_lib(libpytorch_qnnpack) -import_static_lib(libeigen_blas) -import_static_lib(libcpuinfo) -import_static_lib(libclog) - -target_link_libraries(pytorch - fbjni - -Wl,--gc-sections - -Wl,--whole-archive - libtorch - -Wl,--no-whole-archive - libc10 - libnnpack - libpytorch_qnnpack - libeigen_blas - libcpuinfo - libclog -) +if (ANDROID_ABI) + + function(import_static_lib name) + add_library(${name} STATIC IMPORTED) + set_property( + TARGET ${name} + PROPERTY IMPORTED_LOCATION + ${CMAKE_CURRENT_LIST_DIR}/src/main/jniLibs/${ANDROID_ABI}/${name}.a) + endfunction(import_static_lib) + + import_static_lib(libtorch) + import_static_lib(libc10) + import_static_lib(libnnpack) + import_static_lib(libpytorch_qnnpack) + import_static_lib(libeigen_blas) + import_static_lib(libcpuinfo) + import_static_lib(libclog) + + # Link most things statically on Android. + target_link_libraries(pytorch + fbjni + -Wl,--gc-sections + -Wl,--whole-archive + libtorch + -Wl,--no-whole-archive + libc10 + libnnpack + libpytorch_qnnpack + libeigen_blas + libcpuinfo + libclog + ) + +else() + + # Prefer dynamic linking on the host + target_link_libraries(pytorch + fbjni + -Wl,--gc-sections + -Wl,--whole-archive + torch + -Wl,--no-whole-archive + c10 + nnpack + pytorch_qnnpack + cpuinfo + clog + ) + +endif() From d201ff89259d5afa799c58fd707f74655ab6dff9 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 16:01:41 -0700 Subject: [PATCH 44/64] Factor out insertPrepackUnpackForLinear (#27239) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27239 att Test Plan: python test/test_jit.py 'TestJit.test_insert_prepack_unpack' Imported from OSS Differential Revision: D18182913 fbshipit-source-id: 7cbaac9159520d9e873079d10bf80764f2ec27ae --- torch/csrc/jit/passes/quantization.cpp | 71 ++++++++++++++------------ 1 file changed, 37 insertions(+), 34 deletions(-) diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp index cb2766757a0f6..51190dd1005e4 100644 --- a/torch/csrc/jit/passes/quantization.cpp +++ b/torch/csrc/jit/passes/quantization.cpp @@ -625,6 +625,42 @@ void InsertQuantDeQuantImpl( qh.destroyNodes(); } +void insertPrepackUnpackForLinear(std::shared_ptr& graph) { + std::string linear_with_quant = R"( +graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): + %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) + %w_dequant = aten::dequantize(%w_quant) + %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) + return (%r) )"; + + std::string linear_with_quant_prepack = R"( +graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): + %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) + %packed_params = quantized::linear_prepack(%w_quant, %b) + %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack(%packed_params) + %w_dequant = aten::dequantize(%w_quant_unpacked) + %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) + return (%r) )"; + + // Filter to match linear CallFunction + auto filter = [](const Match& match, + const std::unordered_map& vmap) { + const auto& match_vmap = match.values_map; + auto linear_node = match_vmap.at(vmap.at("linear"))->node(); + auto func = + linear_node->output()->type()->expect()->function(); + auto func_name = getFuncName(func->qualname()); + if (func_name == "linear") { + return true; + } + return false; + }; + + SubgraphRewriter rewriter; + rewriter.RegisterRewritePattern(linear_with_quant, linear_with_quant_prepack); + rewriter.runOnGraph(graph, filter); +} + void insertPrepackUnpackForConv2d(std::shared_ptr& graph) { std::string conv_with_quant = R"( graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, %dilation, %groups): @@ -917,40 +953,7 @@ graph(%self, %scale, %zero_point, %dtype): } void InsertPrepackUnpack(std::shared_ptr& graph) { - std::string linear_with_quant = R"( -graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): - %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) - %w_dequant = aten::dequantize(%w_quant) - %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) - return (%r) )"; - - std::string linear_with_quant_prepack = R"( -graph(%linear, %a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype): - %w_quant = aten::quantize_per_tensor(%w, %w_scale, %w_zero_point, %w_dtype) - %packed_params = quantized::linear_prepack(%w_quant, %b) - %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::linear_unpack(%packed_params) - %w_dequant = aten::dequantize(%w_quant_unpacked) - %r = prim::CallFunction(%linear, %a_dequant, %w_dequant, %b) - return (%r) )"; - - // Filter to match linear CallFunction - auto filter = [](const Match& match, - const std::unordered_map& vmap) { - const auto& match_vmap = match.values_map; - auto linear_node = match_vmap.at(vmap.at("linear"))->node(); - auto func = - linear_node->output()->type()->expect()->function(); - auto func_name = getFuncName(func->qualname()); - if (func_name == "linear") { - return true; - } - return false; - }; - - SubgraphRewriter rewriter; - rewriter.RegisterRewritePattern(linear_with_quant, linear_with_quant_prepack); - rewriter.runOnGraph(graph, filter); - + insertPrepackUnpackForLinear(graph); insertPrepackUnpackForConv2d(graph); } From f1f86994bc9c8451110aae40fe69eb13f50f8681 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Tue, 29 Oct 2019 16:52:31 -0700 Subject: [PATCH 45/64] Fix implementation of F::kl_div / F::mse_loss / F::binary_cross_entropy Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28806 Test Plan: Imported from OSS Differential Revision: D18202859 Pulled By: yf225 fbshipit-source-id: 1aa19111cd5111dd5f2779f7f00f07f2f2e16d4d --- .../api/include/torch/nn/functional/loss.h | 79 ++++++++++++++++--- 1 file changed, 66 insertions(+), 13 deletions(-) diff --git a/torch/csrc/api/include/torch/nn/functional/loss.h b/torch/csrc/api/include/torch/nn/functional/loss.h index 19068e54871bd..1405647c90f00 100644 --- a/torch/csrc/api/include/torch/nn/functional/loss.h +++ b/torch/csrc/api/include/torch/nn/functional/loss.h @@ -1,5 +1,6 @@ #pragma once +#include #include namespace torch { @@ -20,31 +21,83 @@ inline Tensor kl_div( const Tensor& input, const Tensor& target, const KLDivLossOptions& options = {}) { - return torch::kl_div( - input, - target, - enumtype::reduction_get_enum(options.reduction())); + torch::Reduction::Reduction reduction_enum; + + if (c10::get_if(&options.reduction())) { + TORCH_WARN("reduction: 'mean' divides the total loss by both the batch size and the support size." + "'batchmean' divides only by the batch size, and aligns with the KL div math definition." + "'mean' will be changed to behave the same as 'batchmean' in the next major release."); + } + + // special case for batchmean + if (c10::get_if(&options.reduction())) { + reduction_enum = torch::Reduction::Sum; + } else { + reduction_enum = enumtype::reduction_get_enum(options.reduction()); + } + + auto reduced = torch::kl_div(input, target, reduction_enum); + + if (c10::get_if(&options.reduction()) && input.dim() != 0) { + reduced = reduced / input.sizes()[0]; + } + + return reduced; } inline Tensor mse_loss( const Tensor& input, const Tensor& target, const MSELossOptions& options = {}) { - return torch::mse_loss( - input, - target, - enumtype::reduction_get_enum(options.reduction())); + if (!(target.sizes() == input.sizes())) { + TORCH_WARN("Using a target size (", target.sizes(), + ") that is different to the input size (", input.sizes(), "). ", + "This will likely lead to incorrect results due to broadcasting. ", + "Please ensure they have the same size."); + } + torch::Tensor ret; + if (target.requires_grad()) { + ret = torch::pow(input - target, 2); + if (!c10::get_if(&options.reduction())) { + ret = (c10::get_if(&options.reduction())) ? torch::mean(ret) : torch::sum(ret); + } + } else { + std::vector broadcast_tensors = torch::broadcast_tensors({input, target}); + auto expanded_input = broadcast_tensors[0]; + auto expanded_target = broadcast_tensors[1]; + ret = torch::mse_loss( + expanded_input, + expanded_target, + enumtype::reduction_get_enum(options.reduction())); + } + return ret; } inline Tensor binary_cross_entropy( const Tensor& input, const Tensor& target, const BCELossOptions& options = {}) { - return torch::binary_cross_entropy( - input, - target, - options.weight(), - enumtype::reduction_get_enum(options.reduction())); + auto reduction_enum = enumtype::reduction_get_enum(options.reduction()); + + if (target.sizes() != input.sizes()) { + TORCH_WARN("Using a target size (", target.sizes(), ") ", + "that is different to the input size (", input.sizes(), ") is deprecated. ", + "Please ensure they have the same size."); + } + if (input.numel() != target.numel()) { + TORCH_CHECK( + false, + "Target and input must have the same number of elements. target nelement (", target.numel(), ") " + "!= input nelement (", input.numel(), ")"); + } + + auto weight = options.weight(); + if (weight.defined()) { + auto new_size = at::infer_size(target.sizes(), weight.sizes()); + weight = weight.expand(new_size); + } + + return torch::binary_cross_entropy(input, target, weight, reduction_enum); } inline Tensor hinge_embedding_loss( From 4e56455b091befa7d07654d31c251ad83104905b Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 29 Oct 2019 16:58:49 -0700 Subject: [PATCH 46/64] whitelist autogradanynonzero (#28852) Summary: prim::AutogradAnyNonZero is optimized away under normal circumstances (a graph executor specializes tensor arguments and runs `specializeAutogradZero`), so the change should be backward compatible for as long as we are running the original executor. Pull Request resolved: https://github.com/pytorch/pytorch/pull/28852 Differential Revision: D18213118 Pulled By: Krovatkin fbshipit-source-id: 223f172c59e5f2b05460db7de98edbadc45dd73d --- test/backward_compatibility/check_backward_compatibility.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index 6a22579d51d4c..a3742c316895b 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -26,6 +26,7 @@ ('rand_like', datetime.date(2019, 11, 11)), ('ones_like', datetime.date(2019, 11, 11)), ('full_like', datetime.date(2019, 11, 11)), + ('AutogradAnyNonZero', datetime.date(2019, 11, 11)), ] From b1ea19ca1785c7a1ae19ae60512b7ec7fe58cc55 Mon Sep 17 00:00:00 2001 From: Jianyu Huang Date: Tue, 29 Oct 2019 17:18:46 -0700 Subject: [PATCH 47/64] Update the misleading comments for zero_points and scale in dynamic quant linear module [1/2] (#28767) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28767 The scale and zero_point are for the output activation tensor, not for the weight tensor. We removed them here because we don't need the zero points and scales for the output tensors in dynamic quantization. ghstack-source-id: 92807318 Test Plan: CI Differential Revision: D18164949 fbshipit-source-id: 0f9172bfef615c30dc28e1dd4448a9f3cc897c2e --- torch/nn/quantized/dynamic/modules/linear.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/torch/nn/quantized/dynamic/modules/linear.py b/torch/nn/quantized/dynamic/modules/linear.py index 360eb233ba725..7574dd53eb761 100644 --- a/torch/nn/quantized/dynamic/modules/linear.py +++ b/torch/nn/quantized/dynamic/modules/linear.py @@ -17,8 +17,6 @@ class Linear(nnq.Linear): shape :math:`(\text{out\_features}, \text{in\_features})`. bias (Tensor): the non-learnable bias of the module of shape :math:`(\text{out\_features})`. If :attr:`bias` is ``True``, the values are initialized to zero. - scale: `scale` parameter of weight Quantized Tensor, type: double - zero_point: `zero_point` parameter for weight Quantized Tensor, type: long Examples:: From 8f1564b8ab38f6a4e7188c5aee17c8c571433279 Mon Sep 17 00:00:00 2001 From: Shihao Xu Date: Tue, 29 Oct 2019 17:24:42 -0700 Subject: [PATCH 48/64] Add enum type to rpc registry for consolidating RPC initialization code path (#28628) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28628 Consolidate code paths of ProcessGroupAgent construction and other RPC Backend construction. ghstack-source-id: 92845348 Differential Revision: D5516188 fbshipit-source-id: 151d9b7b74f68631d6673fecc74dec525949b8f0 --- test/dist_autograd_test.py | 64 ++++++++++++- test/dist_utils.py | 99 ++++++++++---------- test/rpc_test.py | 109 +++++++++++----------- torch/distributed/rpc/__init__.py | 14 ++- torch/distributed/rpc/api.py | 46 +++------ torch/distributed/rpc/backend_registry.py | 80 +++++++++++++--- 6 files changed, 256 insertions(+), 156 deletions(-) diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py index c283d6f537b25..16e8266597598 100644 --- a/test/dist_autograd_test.py +++ b/test/dist_autograd_test.py @@ -8,7 +8,6 @@ import torch.distributed.autograd as dist_autograd import torch.distributed.rpc as rpc from dist_utils import INIT_METHOD_TEMPLATE, dist_init, TEST_CONFIG -from torch.distributed.rpc import RpcBackend import threading @@ -258,6 +257,18 @@ def _verify_graph_for_nested_rpc_call(self, ctx): def _test_graph(self, fn): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) @@ -307,6 +318,18 @@ def test_graph_for_python_call(self): @dist_init(setup_model_parallel=True) def test_graph_for_py_nested_call(self): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) @@ -361,6 +384,18 @@ def test_graph_for_py_nested_call(self): @dist_init(setup_model_parallel=True) def test_graph_for_py_nested_call_itself(self): dst_rank = (self.rank + 1) % self.world_size + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) @@ -469,6 +504,17 @@ def test_context_cleanup_many_workers(self): @dist_init(setup_model_parallel=True) def test_context_cleanup_nested_rpc(self): + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + dst_rank = (self.rank + 1) % self.world_size nested_dst_rank = (dst_rank + 1) % self.world_size with dist_autograd.context() as context_id: @@ -685,10 +731,22 @@ def test_backward_autograd_engine_error(self): # Run backwards, and validate we receive an error. dist_autograd.backward([val.sum()]) - @unittest.skipIf(TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP, + @unittest.skip("Using sleep to simulate syncronization is flaky") + @unittest.skipIf(TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", "Skipping this test temporarily since ProcessGroupAgent does not report errors on node failures") @dist_init(clean_shutdown=False) def test_backward_node_failure(self): + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + with dist_autograd.context() as context_id: t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) @@ -702,7 +760,7 @@ def test_backward_node_failure(self): # Kill all odd rank nodes. if self.rank % 2 == 0: # Wait a bit for all other nodes to die. - time.sleep(5) + time.sleep(5) # This is flaky. with self.assertRaisesRegex(RuntimeError, "Request aborted during client shutdown"): # Run backwards, and validate we receive an error since all # other nodes are dead. diff --git a/test/dist_utils.py b/test/dist_utils.py index 4231e7cf36661..0d0add5afc358 100644 --- a/test/dist_utils.py +++ b/test/dist_utils.py @@ -1,12 +1,11 @@ from __future__ import absolute_import, division, print_function, unicode_literals -from functools import wraps, partial import threading +from functools import partial, wraps from os import getenv import torch.distributed as dist import torch.distributed.rpc as rpc -from torch.distributed.rpc.api import RpcBackend if not dist.is_available(): @@ -15,7 +14,7 @@ class TestConfig: - __slots__ = ["rpc_backend"] + __slots__ = ["rpc_backend_name"] def __init__(self, *args, **kwargs): assert len(args) == 0, "TestConfig only takes kwargs." @@ -23,7 +22,7 @@ def __init__(self, *args, **kwargs): setattr(self, k, v) -TEST_CONFIG = TestConfig(rpc_backend=getenv("RPC_BACKEND", RpcBackend.PROCESS_GROUP)) +TEST_CONFIG = TestConfig(rpc_backend_name=getenv("RPC_BACKEND_NAME", "PROCESS_GROUP")) INIT_METHOD_TEMPLATE = "file://{file_name}" @@ -51,7 +50,7 @@ def set_termination_signal(): _TERMINATION_SIGNAL.set() -def dist_init(test_method=None, setup_model_parallel=True, clean_shutdown=True): +def dist_init(old_test_method=None, setup_model_parallel=True, clean_shutdown=True): """ We use this decorator for setting up and tearing down state since MultiProcessTestCase runs each `test*` method in a separate process and @@ -59,17 +58,21 @@ def dist_init(test_method=None, setup_model_parallel=True, clean_shutdown=True): 'setUp' and 'tearDown' methods of unittest. """ - # If we use dist_init without arguments (ex: @dist_init), test_method is + # If we use dist_init without arguments (ex: @dist_init), old_test_method is # appropriately set and we return the wrapper appropriately. On the other # hand if dist_init has arguments (ex: @dist_init(clean_shutdown=False)), - # test_method is None and we return a functools.partial which is the real + # old_test_method is None and we return a functools.partial which is the real # decorator that is used and as a result we recursively call dist_init with - # test_method and the rest of the arguments appropriately set. - if test_method is None: - return partial(dist_init, setup_model_parallel=setup_model_parallel, clean_shutdown=clean_shutdown) - - @wraps(test_method) - def wrapper(self, *arg, **kwargs): + # old_test_method and the rest of the arguments appropriately set. + if old_test_method is None: + return partial( + dist_init, + setup_model_parallel=setup_model_parallel, + clean_shutdown=clean_shutdown, + ) + + @wraps(old_test_method) + def new_test_method(self, *arg, **kwargs): self.worker_id = self.rank self.worker_name_to_id = { "worker{}".format(rank): rank for rank in range(self.world_size) @@ -79,52 +82,48 @@ def wrapper(self, *arg, **kwargs): global _ALL_NODE_NAMES _ALL_NODE_NAMES = self.worker_name_to_id.keys() - dist.init_process_group( - backend="gloo", - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) # Use enough 'num_send_recv_threads' until we fix https://github.com/pytorch/pytorch/issues/26359 rpc.init_model_parallel( self_name="worker%d" % self.rank, - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) - test_method(self, *arg, **kwargs) - - if setup_model_parallel: - if clean_shutdown: - # Follower reports done. - if self.rank == MASTER_RANK: - on_master_follower_report_done("worker{}".format(MASTER_RANK)) - else: - rpc.rpc_async( - "worker{}".format(MASTER_RANK), - on_master_follower_report_done, - args=("worker{}".format(self.rank),), - ) - - # Master waits for followers to report done. - # Follower waits for master's termination command. - _TERMINATION_SIGNAL.wait() - if self.rank == MASTER_RANK: - # Master sends termination command. - futs = [] - for dst_rank in range(self.world_size): - # torch.distributed.rpc module does not support sending to self. - if dst_rank == MASTER_RANK: - continue - dst_name = "worker{}".format(dst_rank) - fut = rpc.rpc_async(dst_name, set_termination_signal, args=()) - futs.append(fut) - for fut in futs: - assert fut.wait() is None, "Sending termination signal failed." + return_value = old_test_method(self, *arg, **kwargs) + + if setup_model_parallel and clean_shutdown: + # Follower reports done. + if self.rank == MASTER_RANK: + on_master_follower_report_done("worker{}".format(MASTER_RANK)) + else: + rpc.rpc_async( + "worker{}".format(MASTER_RANK), + on_master_follower_report_done, + args=("worker{}".format(self.rank),), + ) + + # Master waits for followers to report done. + # Follower waits for master's termination command. + _TERMINATION_SIGNAL.wait() + if self.rank == MASTER_RANK: + # Master sends termination command. + futs = [] + for dst_rank in range(self.world_size): + # torch.distributed.rpc module does not support sending to self. + if dst_rank == MASTER_RANK: + continue + dst_name = "worker{}".format(dst_rank) + fut = rpc.rpc_async(dst_name, set_termination_signal, args=()) + futs.append(fut) + for fut in futs: + assert fut.wait() is None, "Sending termination signal failed." # Close RPC. rpc.join_rpc() - return wrapper + + return return_value + + return new_test_method diff --git a/test/rpc_test.py b/test/rpc_test.py index 775a83e686e08..f8d8f71c8276b 100644 --- a/test/rpc_test.py +++ b/test/rpc_test.py @@ -11,23 +11,24 @@ import torch.distributed.rpc as rpc from common_utils import load_tests from dist_utils import INIT_METHOD_TEMPLATE, TEST_CONFIG, dist_init -from torch.distributed.rpc import RpcBackend from torch.distributed.rpc.internal import PythonUDF, _internal_rpc_pickler def requires_process_group_agent(message=""): def decorator(old_func): return unittest.skipUnless( - TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP, - message, + TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP", message )(old_func) + return decorator VALUE_FUTURE = concurrent.futures.Future() -def stub_start_rpc_backend_handler(store, self_name, self_rank, worker_name_to_id): +def stub_start_rpc_backend_handler( + store, self_name, self_rank, worker_name_to_id, *args, **kwargs +): return mock.Mock() # RpcAgent. @@ -229,12 +230,21 @@ def test_register_rpc_backend_and_start_rpc_backend( self, mock_rpc_agent, mock_dist_autograd_init ): backend_name = "stub_backend" - rpc.register_backend( + + backend = rpc.backend_registry.register_backend( backend_name, stub_start_rpc_backend_handler ) + + with self.assertRaisesRegex( + RuntimeError, "^RPC backend .+: already registered$" + ): + rpc.backend_registry.register_backend( + backend_name, stub_start_rpc_backend_handler + ) + rpc.init_model_parallel( self_name="worker1", - backend=backend_name, + backend=backend, init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -243,16 +253,10 @@ def test_register_rpc_backend_and_start_rpc_backend( @requires_process_group_agent("PROCESS_GROUP rpc backend specific test, skip") @dist_init(setup_model_parallel=False) def test_duplicate_name(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) with self.assertRaisesRegex(RuntimeError, "is not unique"): rpc.init_model_parallel( self_name="duplicate_name", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -261,87 +265,82 @@ def test_duplicate_name(self): @dist_init(setup_model_parallel=False) def test_reinit(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) rpc.init_model_parallel( self_name="worker{}".format(self.rank), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, ) + + # This is for the below `dist.barrier`. + # For `RpcAgent` other than `ProcessGroupAgent`, + # no `_default_pg` is initialized. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) # Wait for all init to complete. dist.barrier() + with self.assertRaisesRegex(RuntimeError, "is already initialized"): rpc.init_model_parallel( self_name="worker{}".format(self.rank), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, ) rpc.join_rpc() - @dist_init(setup_model_parallel=False) - def test_init_invalid_backend(self): - with self.assertRaisesRegex(RuntimeError, "Unrecognized RPC backend"): - rpc.init_model_parallel( - self_name="worker{}".format(self.rank), - backend="invalid", - init_method=self.init_method, - self_rank=self.rank, - worker_name_to_id=self.worker_name_to_id, - ) - @dist_init(setup_model_parallel=False) def test_invalid_names(self): - dist.init_process_group( - backend=dist.Backend.GLOO, - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) - with self.assertRaisesRegex(RuntimeError, "Worker name must match"): rpc.init_model_parallel( self_name="abc*", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + base_file_name = self.file_name + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "1" with self.assertRaisesRegex(RuntimeError, "Worker name must match"): rpc.init_model_parallel( self_name=" ", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "2" with self.assertRaisesRegex(RuntimeError, "must be non-empty"): rpc.init_model_parallel( self_name="", - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, num_send_recv_threads=16, ) + # Use a different file path for FileStore to avoid rendezvous mismatch. + self.file_name = base_file_name + "3" # If the number in the message does not match, it is likely that the # value of MAX_NAME_LEN in RPC WorkerInfo has changed. with self.assertRaisesRegex(RuntimeError, "shorter than 128"): rpc.init_model_parallel( self_name="".join(["a" for _ in range(500)]), - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -363,7 +362,6 @@ def test_invalid_names(self): # would add extra overhead to the call, and normal use cases won't # create a progress group and exit without doing anything. Hence, it is # not worthy to introduce the overhead just for this test case. - dist.barrier() @dist_init def test_add(self): @@ -450,15 +448,9 @@ def test_sync_rpc(self): @dist_init(setup_model_parallel=False) def test_join_rpc(self): # Initialize RPC. - dist.init_process_group( - backend="gloo", - init_method=self.init_method, - rank=self.rank, - world_size=self.world_size, - ) rpc.init_model_parallel( self_name="worker%d" % self.rank, - backend=TEST_CONFIG.rpc_backend, + backend=rpc.backend_registry.BackendType[TEST_CONFIG.rpc_backend_name], init_method=self.init_method, self_rank=self.rank, worker_name_to_id=self.worker_name_to_id, @@ -932,5 +924,18 @@ def test_requires_process_group_agent_decorator(self): def test_func(): return "expected result" - if TEST_CONFIG.rpc_backend == RpcBackend.PROCESS_GROUP: + if TEST_CONFIG.rpc_backend_name == "PROCESS_GROUP": self.assertEqual(test_func(), "expected result") + + def test_dist_init_decorator(self): + @dist_init(setup_model_parallel=False) + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + @dist_init + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") diff --git a/torch/distributed/rpc/__init__.py b/torch/distributed/rpc/__init__.py index b1e0e17049503..39b76c423a6a2 100644 --- a/torch/distributed/rpc/__init__.py +++ b/torch/distributed/rpc/__init__.py @@ -2,7 +2,7 @@ import sys -from .backend_registry import * # noqa: F401 +from . import backend_registry if sys.version_info >= (3, 0): @@ -13,7 +13,7 @@ def init_model_parallel( self_name, - backend=RpcBackend.PROCESS_GROUP, + backend=backend_registry.BackendType.PROCESS_GROUP, init_method=None, self_rank=-1, worker_name_to_id=None, @@ -43,14 +43,22 @@ def init_model_parallel( init_method(str): backend specific init arguments. num_send_recv_threads(int): Number of threads for send/recv work. """ + # Rendezvous. + world_size = len(worker_name_to_id) + rendezvous_iterator = torch.distributed.rendezvous( + init_method, rank=self_rank, world_size=world_size + ) + store, _, _ = next(rendezvous_iterator) + # Initialize RPC. _init_rpc( backend, - init_method, + store, self_name, self_rank, worker_name_to_id, num_send_recv_threads, ) + # Initialize Autograd. torch.distributed.autograd._init(api._agent.get_worker_info().id) diff --git a/torch/distributed/rpc/api.py b/torch/distributed/rpc/api.py index 4d9534873f6e3..84540654de159 100644 --- a/torch/distributed/rpc/api.py +++ b/torch/distributed/rpc/api.py @@ -2,15 +2,13 @@ from torch.distributed import invoke_remote_builtin, invoke_remote_python_udf from torch.distributed import _start_rpc_agent from torch.distributed import _destroy_rref_context, _cleanup_python_rpc_handler -from torch.distributed import ProcessGroupAgent from torch.distributed import WorkerInfo -from .backend_registry import is_backend_registered, init_backend +from . import backend_registry from .internal import _internal_rpc_pickler, PythonUDF import functools import sys import torch -from enum import Enum _agent = None @@ -58,14 +56,12 @@ def sync_rpc(): _agent.sync() -class RpcBackend(Enum): - PROCESS_GROUP = 1 # TODO: add a context manager to wrap _init_rpc and join_rpc def _init_rpc( - backend=RpcBackend.PROCESS_GROUP, - init_method=None, + backend=backend_registry.BackendType.PROCESS_GROUP, + store=None, self_name=None, self_rank=-1, worker_name_to_id=None, @@ -79,33 +75,15 @@ def _init_rpc( if _agent: raise RuntimeError("RPC is already initialized") - if backend == RpcBackend.PROCESS_GROUP: - from torch.distributed.distributed_c10d import _get_default_group - - group = _get_default_group() - if (self_rank != -1) and (self_rank != group.rank()): - raise RuntimeError("self_rank argument {} doesn't match pg rank {}".format( - self_rank, group.rank())) - if (worker_name_to_id is not None) and (len(worker_name_to_id) != group.size()): - raise RuntimeError("worker_name_to_id argument {} doesn't match pg size {}".format( - worker_name_to_id, group.size())) - # TODO: add try-except and destroy _agent in all processes if any fails. - _agent = ProcessGroupAgent(self_name, group, num_send_recv_threads) - elif is_backend_registered(backend): - # Rendezvous. - world_size = len(worker_name_to_id) - rendezvous_iterator = torch.distributed.rendezvous(init_method, self_rank, world_size) - store, self_rank, world_size = next(rendezvous_iterator) - # Initialize RPC. - _agent = init_backend( - backend, - store=store, - self_name=self_name, - self_rank=self_rank, - worker_name_to_id=worker_name_to_id, - ) - else: - raise RuntimeError("Unrecognized RPC backend ", backend) + # Initialize RPC. + _agent = backend_registry.init_backend( + backend, + store=store, + self_name=self_name, + self_rank=self_rank, + worker_name_to_id=worker_name_to_id, + num_send_recv_threads=num_send_recv_threads, + ) _start_rpc_agent(_agent) diff --git a/torch/distributed/rpc/backend_registry.py b/torch/distributed/rpc/backend_registry.py index e20a9e5d3b681..4fabdf20a934d 100644 --- a/torch/distributed/rpc/backend_registry.py +++ b/torch/distributed/rpc/backend_registry.py @@ -1,15 +1,16 @@ from __future__ import absolute_import, division, print_function, unicode_literals +import collections +import enum -_BACKEND_REGISTRY = {} +import torch.distributed as dist +import torch.distributed.distributed_c10d as dc10d -def _get_backend_registry(): - return _BACKEND_REGISTRY +BackendValue = collections.namedtuple("BackendValue", ["init_backend_handler"]) - -def is_backend_registered(backend_name): - return backend_name in _get_backend_registry() +# Create an enum type, `BackendType`, with empty members. +BackendType = enum.Enum(value="BackendType", names={}) def register_backend(backend_name, init_backend_handler): @@ -21,14 +22,65 @@ def register_backend(backend_name, init_backend_handler): `_init_rpc()` function is called with a backend. This returns the agent. """ - backend_registry = _get_backend_registry() - if backend_name in backend_registry: + global BackendType + if backend_name in BackendType.__members__.keys(): raise RuntimeError("RPC backend {}: already registered".format(backend_name)) - backend_registry[backend_name] = init_backend_handler + # Create a new enum type, `BackendType`, with extended members. + existing_enum_dict = {member.name: member.value for member in BackendType} + extended_enum_dict = dict( + {backend_name: BackendValue(init_backend_handler=init_backend_handler)}, + **existing_enum_dict + ) + BackendType = enum.Enum(value="BackendType", names=extended_enum_dict) + return BackendType[backend_name] + + +def init_backend(backend, *args, **kwargs): + return backend.value.init_backend_handler(*args, **kwargs) + + +def process_group_init_backend_handler( + store, + self_name, + self_rank, + worker_name_to_id, + num_send_recv_threads, + *args, + **kwargs +): + # Initialize ProcessGroup. + if dist.is_initialized(): + raise RuntimeError( + "Default process group must not be initialized before `init_model_parallel`." + ) + + world_size = len(worker_name_to_id) + dist.init_process_group( + backend="gloo", store=store, rank=self_rank, world_size=world_size + ) + + try: + group = dc10d._get_default_group() + assert group is not None, "Failed to initialize default ProcessGroup." + + if (self_rank != -1) and (self_rank != group.rank()): + raise RuntimeError( + "self_rank argument {} doesn't match pg rank {}".format( + self_rank, group.rank() + ) + ) + if (worker_name_to_id is not None) and (len(worker_name_to_id) != group.size()): + raise RuntimeError( + "worker_name_to_id argument {} doesn't match pg size {}".format( + worker_name_to_id, group.size() + ) + ) + # TODO: add try-except and destroy _agent in all processes if any fails. + return dist.ProcessGroupAgent(self_name, group, num_send_recv_threads) + except Exception as ex: + dist.destroy_process_group() + raise ex + -def init_backend(backend_name, *args, **kwargs): - backend_registry = _get_backend_registry() - if backend_name not in backend_registry: - raise RuntimeError("No rpc_init handler for {}.".format(backend_name)) - return backend_registry[backend_name](*args, **kwargs) +register_backend("PROCESS_GROUP", process_group_init_backend_handler) From cbc234bcebe3a155a1cbe7e02282d64c52ffd0d4 Mon Sep 17 00:00:00 2001 From: nuka137 Date: Tue, 29 Oct 2019 17:25:56 -0700 Subject: [PATCH 49/64] C++ API: torch::nn::BatchNorm1d (#28176) Summary: Add torch::nn::BatchNorm1d function/module support for the C++ API. torch::nn::BatchNorm{2,3}d will be added after this PR is merged. Related Issue: https://github.com/pytorch/pytorch/issues/25883 Reviewer: yf225 I would like to discuss about below items. * Necessity of `num_batches_tracked` in `BatchNormImplBase` * `num_batches_tracked` is needed to calculate `momentum` when we do not feed `momentum` argument in Python API. But in C++ API, `momentum` argument has a default value. * `num_batches_tracked` is only used for counting up `BatchNorm1d::foward()` call. I think it is no necessary for user anymore. * The design of `BatchNorm{1,2,3}dOptions` * We have already `BatchNormOptions` used for deprecated `BatchNorm` module. However, it is hard to use it for `BatchNorm{1,2,3}dOptions` because of the arguments disagreement of each modules. * In this PR, I introduce `BatchNormOptionsv2` template class for the `BatchNorm{1,2,3}dOptions`. But I'm not sure this design is good or not. Pull Request resolved: https://github.com/pytorch/pytorch/pull/28176 Differential Revision: D18196843 Pulled By: yf225 fbshipit-source-id: 667e2b5de4150d5776c41b9088c9e6c2ead24cd4 --- test/cpp/api/functional.cpp | 27 ++++ test/cpp/api/modulelist.cpp | 2 +- test/cpp/api/modules.cpp | 83 +++++++++++- test/cpp/api/sequential.cpp | 4 +- test/cpp_api_parity/parity-tracker.md | 2 +- torch/csrc/api/include/torch/nn/functional.h | 1 + .../include/torch/nn/functional/batchnorm.h | 36 ++++++ .../api/include/torch/nn/modules/batchnorm.h | 67 +++++++++- .../api/include/torch/nn/options/batchnorm.h | 36 ++++-- torch/csrc/api/src/nn/modules/batchnorm.cpp | 118 ++++++++++++++++-- torch/csrc/api/src/nn/options/batchnorm.cpp | 2 +- 11 files changed, 341 insertions(+), 37 deletions(-) create mode 100644 torch/csrc/api/include/torch/nn/functional/batchnorm.h diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp index bd1a7623ff6e7..c058fbd0e645a 100644 --- a/test/cpp/api/functional.cpp +++ b/test/cpp/api/functional.cpp @@ -1266,6 +1266,33 @@ TEST_F(FunctionalTest, Threshold) { } } +TEST_F(FunctionalTest, BatchNorm1d) { + int num_features = 5; + double eps = 1e-05; + double momentum = 0.1; + + auto input = torch::randn({2, 5}); + auto mean = torch::randn(5); + auto variance = torch::rand(5); + auto weight = torch::ones({num_features}); + auto bias = torch::zeros({num_features}); + auto output = F::batch_norm( + input, mean, variance, + BatchNormOptions().weight(weight).bias(bias).momentum(momentum).eps(eps), + /*training=*/false); + auto expected = (input - mean) / torch::sqrt(variance + eps); + ASSERT_TRUE(output.allclose(expected)); +} + +TEST_F(FunctionalTest, BatchNorm1dDefaultOptions) { + auto input = torch::randn({2, 5}); + auto mean = torch::randn(5); + auto variance = torch::rand(5); + auto output = F::batch_norm(input, mean, variance); + auto expected = (input - mean) / torch::sqrt(variance + 1e-5); + ASSERT_TRUE(output.allclose(expected)); +} + TEST_F(FunctionalTest, Interpolate) { { // 1D interpolation diff --git a/test/cpp/api/modulelist.cpp b/test/cpp/api/modulelist.cpp index e4620ea9d4add..68ff0f6d795ef 100644 --- a/test/cpp/api/modulelist.cpp +++ b/test/cpp/api/modulelist.cpp @@ -281,7 +281,7 @@ TEST_F(ModuleListTest, PrettyPrintModuleList) { " (0): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (1): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (2): torch::nn::Dropout(rate=0.5)\n" - " (3): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (3): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (4): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (5): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 89b277f9c3822..4d1ac87b15ccf 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -1001,7 +1001,7 @@ TEST_F(ModulesTest, BatchNormStateful) { BatchNorm bn(5); // Is stateful by default. - ASSERT_TRUE(bn->options.stateful()); + ASSERT_TRUE(bn->options.track_running_stats()); ASSERT_TRUE(bn->running_mean.defined()); ASSERT_EQ(bn->running_mean.dim(), 1); @@ -1023,7 +1023,7 @@ TEST_F(ModulesTest, BatchNormStateful) { ASSERT_EQ(bn->bias.size(0), 5); } TEST_F(ModulesTest, BatchNormStateless) { - BatchNorm bn(BatchNormOptions(5).stateful(false).affine(false)); + BatchNorm bn(BatchNormOptions(5).track_running_stats(false).affine(false)); ASSERT_FALSE(bn->running_mean.defined()); ASSERT_FALSE(bn->running_var.defined()); @@ -1033,7 +1033,7 @@ TEST_F(ModulesTest, BatchNormStateless) { ASSERT_THROWS_WITH( bn(torch::ones({2, 5})), "Calling BatchNorm::forward is only permitted " - "when the 'stateful' option is true (was false). " + "when the 'track_running_stats' option is true (was false). " "Use BatchNorm::pure_forward instead."); } @@ -1051,6 +1051,71 @@ TEST_F(ModulesTest, BatchNormPureForward) { ASSERT_TRUE(output.allclose(expected)); } +TEST_F(ModulesTest, BatchNormLegacyWarning) { + std::stringstream buffer; + torch::test::CerrRedirect cerr_redirect(buffer.rdbuf()); + + BatchNorm bn(5); + + ASSERT_EQ( + count_substr_occurrences( + buffer.str(), + "torch::nn::BatchNorm module is deprecated" + ), + 1); +} + +TEST_F(ModulesTest, BatchNorm1dStateful) { + BatchNorm1d bn(BatchNorm1dOptions(5)); + + ASSERT_TRUE(bn->options.track_running_stats()); + + ASSERT_TRUE(bn->running_mean.defined()); + ASSERT_EQ(bn->running_mean.dim(), 1); + ASSERT_EQ(bn->running_mean.size(0), 5); + + ASSERT_TRUE(bn->running_var.defined()); + ASSERT_EQ(bn->running_var.dim(), 1); + ASSERT_EQ(bn->running_var.size(0), 5); + + ASSERT_TRUE(bn->num_batches_tracked.defined()); + ASSERT_EQ(bn->num_batches_tracked.dim(), 1); + ASSERT_EQ(bn->num_batches_tracked.size(0), 1); + + ASSERT_TRUE(bn->options.affine()); + + ASSERT_TRUE(bn->weight.defined()); + ASSERT_EQ(bn->weight.dim(), 1); + ASSERT_EQ(bn->weight.size(0), 5); + + ASSERT_TRUE(bn->bias.defined()); + ASSERT_EQ(bn->bias.dim(), 1); + ASSERT_EQ(bn->bias.size(0), 5); +} + +TEST_F(ModulesTest, BatchNorm1dStateless) { + BatchNorm1d bn(BatchNorm1dOptions(5).track_running_stats(false).affine(false)); + + ASSERT_FALSE(bn->running_mean.defined()); + ASSERT_FALSE(bn->running_var.defined()); + ASSERT_FALSE(bn->num_batches_tracked.defined()); + ASSERT_FALSE(bn->weight.defined()); + ASSERT_FALSE(bn->bias.defined()); +} + +TEST_F(ModulesTest, BatchNorm1d) { + BatchNorm1d bn(BatchNorm1dOptions(5)); + bn->eval(); + + auto input = torch::randn({2, 5}, torch::requires_grad()); + auto output = bn->forward(input); + auto s = output.sum(); + s.backward(); + + ASSERT_EQ(input.sizes(), input.grad().sizes()); + ASSERT_TRUE(input.grad().allclose(torch::ones({2, 5}))); +} + TEST_F(ModulesTest, Linear_CUDA) { Linear model(5, 2); model->to(torch::kCUDA); @@ -2303,9 +2368,17 @@ TEST_F(ModulesTest, PrettyPrintFunctional) { TEST_F(ModulesTest, PrettyPrintBatchNorm) { ASSERT_EQ( c10::str(BatchNorm( - BatchNormOptions(4).eps(0.5).momentum(0.1).affine(false).stateful( + BatchNormOptions(4).eps(0.5).momentum(0.1).affine(false).track_running_stats( true))), - "torch::nn::BatchNorm(features=4, eps=0.5, momentum=0.1, affine=false, stateful=true)"); + "torch::nn::BatchNorm(num_features=4, eps=0.5, momentum=0.1, affine=false, track_running_stats=true)"); +} + +TEST_F(ModulesTest, PrettyPrintBatchNorm1d) { + ASSERT_EQ( + c10::str(BatchNorm1d( + BatchNorm1dOptions(4).eps(0.5).momentum(0.1).affine(false) + .track_running_stats(true))), + "torch::nn::BatchNorm1d(4, eps=0.5, momentum=0.1, affine=false, track_running_stats=true)"); } TEST_F(ModulesTest, PrettyPrintLayerNorm) { diff --git a/test/cpp/api/sequential.cpp b/test/cpp/api/sequential.cpp index 543df6606589b..c33932186f17e 100644 --- a/test/cpp/api/sequential.cpp +++ b/test/cpp/api/sequential.cpp @@ -412,7 +412,7 @@ TEST_F(SequentialTest, PrettyPrintSequential) { " (0): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (1): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (2): torch::nn::Dropout(rate=0.5)\n" - " (3): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (3): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (4): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (5): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); @@ -431,7 +431,7 @@ TEST_F(SequentialTest, PrettyPrintSequential) { " (linear): torch::nn::Linear(in_features=10, out_features=3, bias=true)\n" " (conv2d): torch::nn::Conv2d(input_channels=1, output_channels=2, kernel_size=[3, 3], stride=[1, 1])\n" " (dropout): torch::nn::Dropout(rate=0.5)\n" - " (batchnorm): torch::nn::BatchNorm(features=5, eps=1e-05, momentum=0.1, affine=true, stateful=true)\n" + " (batchnorm): torch::nn::BatchNorm(num_features=5, eps=1e-05, momentum=0.1, affine=true, track_running_stats=true)\n" " (embedding): torch::nn::Embedding(num_embeddings=4, embedding_dim=10)\n" " (lstm): torch::nn::LSTM(input_size=4, hidden_size=5, layers=1, dropout=0)\n" ")"); diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index 8a5d5f45082c1..5ad56fe027a57 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -69,7 +69,7 @@ torch.nn.Softmax|Yes|No torch.nn.Softmax2d|Yes|No torch.nn.LogSoftmax|Yes|No torch.nn.AdaptiveLogSoftmaxWithLoss|No|No -torch.nn.BatchNorm1d|No|No +torch.nn.BatchNorm1d|Yes|No torch.nn.BatchNorm2d|No|No torch.nn.BatchNorm3d|No|No torch.nn.GroupNorm|No|No diff --git a/torch/csrc/api/include/torch/nn/functional.h b/torch/csrc/api/include/torch/nn/functional.h index 90ce395950145..6531b7fb8ade9 100644 --- a/torch/csrc/api/include/torch/nn/functional.h +++ b/torch/csrc/api/include/torch/nn/functional.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include diff --git a/torch/csrc/api/include/torch/nn/functional/batchnorm.h b/torch/csrc/api/include/torch/nn/functional/batchnorm.h new file mode 100644 index 0000000000000..a180f3fb43f95 --- /dev/null +++ b/torch/csrc/api/include/torch/nn/functional/batchnorm.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +namespace torch { +namespace nn { +namespace functional { + +inline Tensor batch_norm(const Tensor& input, const Tensor& running_mean, + const Tensor& running_var, const BatchNormOptions& options = {}, bool training = false) { + if (training) { + auto size = input.sizes(); + int64_t size_prods = size[0]; + for (size_t i = 0; i < size.size() - 2; i++) { + size_prods *= size[i + 2]; + } + TORCH_CHECK(size_prods != 1, + "Expected more than 1 value per channel when training, got input size ", size); + } + + return torch::batch_norm( + input, + options.weight(), + options.bias(), + running_mean, + running_var, + training, + options.momentum().value(), + options.eps(), + at::globalContext().userEnabledCuDNN()); +} + +} // namespace functional +} // namespace nn +} // namespace torch diff --git a/torch/csrc/api/include/torch/nn/modules/batchnorm.h b/torch/csrc/api/include/torch/nn/modules/batchnorm.h index 210effd35a17c..eb3e26700fe19 100644 --- a/torch/csrc/api/include/torch/nn/modules/batchnorm.h +++ b/torch/csrc/api/include/torch/nn/modules/batchnorm.h @@ -25,8 +25,8 @@ namespace nn { /// \endrst class TORCH_API BatchNormImpl : public torch::nn::Cloneable { public: - explicit BatchNormImpl(int64_t features) - : BatchNormImpl(BatchNormOptions(features)) {} + explicit BatchNormImpl(int64_t num_features) + : BatchNormImpl(BatchNormOptions(num_features)) {} explicit BatchNormImpl(const BatchNormOptions& options_); void reset() override; @@ -37,7 +37,7 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { /// Applies batch normalization on the `input` using the stored mean and /// variance. /// - /// The module must be constructed with `stateful = true` when calling this + /// The module must be constructed with `track_running_stats = true` when calling this /// method, as the module will otherwise not store running statistics. If you /// want to supply the mean and variance yourself, use `pure_forward`. Tensor forward(const Tensor& input); @@ -61,11 +61,11 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { Tensor bias; /// The running mean. - /// Only defined if the `stateful` option was `true` upon construction. + /// Only defined if the `track_running_stats` option was `true` upon construction. Tensor running_mean; /// The running variance. - /// Only defined if the `stateful` option was `true` upon construction. + /// Only defined if the `track_running_stats` option was `true` upon construction. Tensor running_var; }; @@ -75,5 +75,62 @@ class TORCH_API BatchNormImpl : public torch::nn::Cloneable { /// module storage semantics. TORCH_MODULE(BatchNorm); +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BatchNorm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +/// Base class for all (dimension-specialized) batchnorm modules. +template +class TORCH_API BatchNormImplBase : public torch::nn::Cloneable { + protected: + virtual void _check_input_dim(const Tensor& input) = 0; + + public: + explicit BatchNormImplBase(const BatchNormOptions& options_); + + Tensor forward(const Tensor& input); + + void reset_running_stats(); + + void reset() override; + + /// Pretty prints the `BatchNorm{1,2,3}d` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + /// The options with which this module was constructed. + BatchNormOptions options; + + /// The learned weight. + /// Only defined if the `affine` option was `true` upon construction. + Tensor weight; + + /// The learned bias. + /// Only defined if the `affine` option was `true` upon construction. + Tensor bias; + + /// The running mean. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor running_mean; + + /// The running variance. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor running_var; + + /// The number of the forward call. + /// Only defined if the `track_running_stats` option was `true` upon construction. + Tensor num_batches_tracked; +}; + +/// Applies the BatchNorm1d function. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.BatchNorm1d to learn +/// about the exact behavior of this module. +class TORCH_API BatchNorm1dImpl : public BatchNormImplBase<1, BatchNorm1dImpl> { + protected: + virtual void _check_input_dim(const Tensor& input) override; + + public: + using BatchNormImplBase<1, BatchNorm1dImpl>::BatchNormImplBase; +}; + +TORCH_MODULE(BatchNorm1d); + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/include/torch/nn/options/batchnorm.h b/torch/csrc/api/include/torch/nn/options/batchnorm.h index ca6a952603d97..2aa175cd09144 100644 --- a/torch/csrc/api/include/torch/nn/options/batchnorm.h +++ b/torch/csrc/api/include/torch/nn/options/batchnorm.h @@ -9,26 +9,40 @@ namespace nn { /// Options for the `BatchNorm` module. struct TORCH_API BatchNormOptions { - /* implicit */ BatchNormOptions(int64_t features); + BatchNormOptions() {} + + /* implicit */ BatchNormOptions(int64_t num_features); + /// The number of features of the input tensor. /// Changing this parameter after construction __has no effect__. - TORCH_ARG(int64_t, features); + TORCH_ARG(int64_t, num_features); + + /// The epsilon value added for numerical stability. + /// Changing this parameter after construction __is effective__. + TORCH_ARG(double, eps) = 1e-5; + + /// A momentum multiplier for the mean and variance. + /// Changing this parameter after construction __is effective__. + TORCH_ARG(c10::optional, momentum) = 0.1; + /// Whether to learn a scale and bias that are applied in an affine /// transformation on the input. /// Changing this parameter after construction __has no effect__. TORCH_ARG(bool, affine) = true; + /// Whether to store and update batch statistics (mean and variance) in the - /// module. If `false`, you should call `pure_forward` and supply those batch - /// statistics yourself. + /// module. /// Changing this parameter after construction __has no effect__. - TORCH_ARG(bool, stateful) = true; - /// The epsilon value added for numerical stability. - /// Changing this parameter after construction __is effective__. - TORCH_ARG(double, eps) = 1e-5; - /// A momentum multiplier for the mean and variance. - /// Changing this parameter after construction __is effective__. - TORCH_ARG(double, momentum) = 0.1; + TORCH_ARG(bool, track_running_stats) = true; + + /// This parameter is only used in `F::batch_norm`. + TORCH_ARG(Tensor, weight) = Tensor(); + + /// This parameter is only used in `F::batch_norm`. + TORCH_ARG(Tensor, bias) = Tensor(); }; +using BatchNorm1dOptions = BatchNormOptions; + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/src/nn/modules/batchnorm.cpp b/torch/csrc/api/src/nn/modules/batchnorm.cpp index 806d77bdb2c92..816d07dc31323 100644 --- a/torch/csrc/api/src/nn/modules/batchnorm.cpp +++ b/torch/csrc/api/src/nn/modules/batchnorm.cpp @@ -1,7 +1,9 @@ +#include #include #include #include +#include #include @@ -10,41 +12,45 @@ #include #include +namespace F = torch::nn::functional; + namespace torch { namespace nn { BatchNormImpl::BatchNormImpl(const BatchNormOptions& options_) : options(options_) { + TORCH_WARN("torch::nn::BatchNorm module is deprecated." + "Use BatchNorm{1,2,3}d instead."); reset(); } void BatchNormImpl::reset() { if (options.affine()) { weight = register_parameter( - "weight", torch::empty({options.features()}).uniform_()); - bias = register_parameter("bias", torch::zeros({options.features()})); + "weight", torch::empty({options.num_features()}).uniform_()); + bias = register_parameter("bias", torch::zeros({options.num_features()})); } - if (options.stateful()) { + if (options.track_running_stats()) { running_mean = - register_buffer("running_mean", torch::zeros({options.features()})); + register_buffer("running_mean", torch::zeros({options.num_features()})); running_var = - register_buffer("running_var", torch::ones({options.features()})); + register_buffer("running_var", torch::ones({options.num_features()})); } } void BatchNormImpl::pretty_print(std::ostream& stream) const { stream << std::boolalpha - << "torch::nn::BatchNorm(features=" << options.features() - << ", eps=" << options.eps() << ", momentum=" << options.momentum() - << ", affine=" << options.affine() << ", stateful=" << options.stateful() + << "torch::nn::BatchNorm(num_features=" << options.num_features() + << ", eps=" << options.eps() << ", momentum=" << options.momentum().value() + << ", affine=" << options.affine() << ", track_running_stats=" << options.track_running_stats() << ")"; } Tensor BatchNormImpl::forward(const Tensor& input) { TORCH_CHECK( - options.stateful(), + options.track_running_stats(), "Calling BatchNorm::forward is only permitted when " - "the 'stateful' option is true (was false). " + "the 'track_running_stats' option is true (was false). " "Use BatchNorm::pure_forward instead."); return pure_forward(input, running_mean, running_var); } @@ -67,10 +73,100 @@ Tensor BatchNormImpl::pure_forward( mean, variance, is_training(), - options.momentum(), + options.momentum().value(), options.eps(), torch::cuda::cudnn_is_available()); } +template +BatchNormImplBase::BatchNormImplBase(const BatchNormOptions& options_) + : options(options_) { + reset(); +} + +template +void BatchNormImplBase::reset_running_stats() { + if (options.track_running_stats()) { + running_mean.zero_(); + running_var.fill_(1); + num_batches_tracked.zero_(); + } +} + +template +void BatchNormImplBase::reset() { + if (options.affine()) { + weight = this->register_parameter("weight", torch::empty({options.num_features()})); + bias = this->register_parameter("bias", torch::empty({options.num_features()})); + } else { + weight = this->register_parameter("weight", Tensor()); + bias = this->register_parameter("bias", Tensor()); + } + if (options.track_running_stats()) { + running_mean = this->register_buffer("running_mean", torch::zeros({options.num_features()})); + running_var = this->register_buffer("running_var", torch::ones({options.num_features()})); + num_batches_tracked = this->register_buffer("num_batches_tracked", torch::tensor(0, torch::dtype(torch::kLong))); + } else { + running_mean = this->register_buffer("running_mean", Tensor()); + running_var = this->register_buffer("running_var", Tensor()); + num_batches_tracked = this->register_buffer("num_batches_tracked", Tensor()); + } + + reset_running_stats(); + if (options.affine()) { + torch::nn::init::ones_(weight); + torch::nn::init::zeros_(bias); + } +} + +template +void BatchNormImplBase::pretty_print(std::ostream& stream) const { + stream << std::boolalpha + << "torch::nn::BatchNorm" << D << "d(" + << options.num_features() << ", " + << "eps=" << options.eps() << ", " + << "momentum=" << options.momentum().value() << ", " + << "affine=" << options.affine() << ", " + << "track_running_stats=" << options.track_running_stats() << ")"; +} + +template +Tensor BatchNormImplBase::forward(const Tensor& input) { + _check_input_dim(input); + + double exponential_average_factor; + if (options.momentum() == c10::nullopt) { + exponential_average_factor = 0.0; + } else { + exponential_average_factor = options.momentum().value(); + } + + if (this->is_training() && options.track_running_stats()) { + if (num_batches_tracked.defined()) { + num_batches_tracked += 1; + if (options.momentum() == c10::nullopt) { // use cumulative moving average + exponential_average_factor = 1.0 / num_batches_tracked.item(); + } else { // use exponential moving average + exponential_average_factor = options.momentum().value(); + } + } + } + + return F::batch_norm( + input, + running_mean, + running_var, + BatchNormOptions().weight(weight).bias(bias).momentum(exponential_average_factor).eps(options.eps()), + this->is_training() || !options.track_running_stats()); +} + +void BatchNorm1dImpl::_check_input_dim(const Tensor& input) { + TORCH_CHECK( + input.dim() == 2 || input.dim() == 3, + "expected 2D or 3D input (got ", input.dim(), "D input)"); +} + +template class BatchNormImplBase<1, BatchNorm1dImpl>; + } // namespace nn } // namespace torch diff --git a/torch/csrc/api/src/nn/options/batchnorm.cpp b/torch/csrc/api/src/nn/options/batchnorm.cpp index 60c363b8d9e5c..2144443913749 100644 --- a/torch/csrc/api/src/nn/options/batchnorm.cpp +++ b/torch/csrc/api/src/nn/options/batchnorm.cpp @@ -3,7 +3,7 @@ namespace torch { namespace nn { -BatchNormOptions::BatchNormOptions(int64_t features) : features_(features) {} +BatchNormOptions::BatchNormOptions(int64_t num_features) : num_features_(num_features) {} } // namespace nn } // namespace torch From 57c9b1cefc8fed09ffb251040b9db2481ec34f81 Mon Sep 17 00:00:00 2001 From: Zafar Takhirov Date: Tue, 29 Oct 2019 17:32:29 -0700 Subject: [PATCH 50/64] Enabling inplace relu Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28710 Test Plan: Imported from OSS Differential Revision: D18146120 Pulled By: z-a-f fbshipit-source-id: d8f0982f5a2ae35f7deb34e67cdb64be700a9d6c --- torch/nn/quantized/modules/activation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torch/nn/quantized/modules/activation.py b/torch/nn/quantized/modules/activation.py index d906b0a5c47a8..422e7897bbd91 100644 --- a/torch/nn/quantized/modules/activation.py +++ b/torch/nn/quantized/modules/activation.py @@ -31,10 +31,10 @@ class ReLU(torch.nn.ReLU): """ def __init__(self, inplace=False): super(ReLU, self).__init__(inplace) - assert not inplace, 'torch.nn.quantized.ReLU does not support inplace' + self.inplace = inplace def forward(self, input): - return torch.nn.quantized.functional.relu(input) + return torch.nn.quantized.functional.relu(input, inplace=self.inplace) def _get_name(self): return 'QuantizedReLU' From dfe7b25eafd88130306991299574145f3333106e Mon Sep 17 00:00:00 2001 From: mrsalehi Date: Tue, 29 Oct 2019 17:51:06 -0700 Subject: [PATCH 51/64] Add nn::Flatten to C++ Frontend (#28072) Summary: Adds torch::nn::Flatten module support for the C++ API. Issue: https://github.com/pytorch/pytorch/issues/25883 Reviewer: yf225 Pull Request resolved: https://github.com/pytorch/pytorch/pull/28072 Differential Revision: D18202778 Pulled By: yf225 fbshipit-source-id: 43345dcbdf2f50d75746bf9a0ba293b84df275ab --- test/cpp/api/modules.cpp | 35 +++++++++++++++++++ test/cpp_api_parity/parity-tracker.md | 2 +- .../api/include/torch/nn/modules/linear.h | 25 +++++++++++++ .../api/include/torch/nn/options/linear.h | 10 ++++++ torch/csrc/api/src/nn/modules/linear.cpp | 15 ++++++++ 5 files changed, 86 insertions(+), 1 deletion(-) diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 4d1ac87b15ccf..e32d1ebbe5e97 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -296,6 +296,36 @@ TEST_F(ModulesTest, Identity) { ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); } +TEST_F(ModulesTest, Flatten) { + Flatten flatten; + auto input = torch::tensor({{1, 3, 4}, {2, 5, 6}}, torch::requires_grad()); + auto output = flatten->forward(input); + auto expected = torch::tensor({{1, 3, 4}, {2, 5, 6}}, torch::kFloat); + auto s = output.sum(); + + s.backward(); + ASSERT_TRUE(torch::equal(output, expected)); + ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); + + // Testing with optional arguments start_dim and end_dim + Flatten flatten_optional_dims(FlattenOptions().start_dim(2).end_dim(3)); + input = torch::tensor({ + {{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}, + {{{9, 10}, {11, 12}}, {{13, 14}, {15, 16}}} + }, torch::requires_grad()); // Tensor with sizes (2, 2, 2, 2) + + output = flatten_optional_dims->forward(input); + expected = torch::tensor({ + {{1, 2, 3, 4}, {5, 6, 7, 8}}, + {{9, 10, 11, 12}, {13, 14, 15, 16}} + }, torch::kFloat); // Tensor with sizes (2, 2, 4) + + s = output.sum(); + s.backward(); + ASSERT_TRUE(torch::equal(output, expected)); + ASSERT_TRUE(torch::equal(input.grad(), torch::ones_like(input))); +} + TEST_F(ModulesTest, AdaptiveMaxPool1d) { AdaptiveMaxPool1d model(3); auto x = torch::tensor({{{1, 2, 3, 4, 5}}}, torch::requires_grad()); @@ -1904,6 +1934,11 @@ TEST_F(ModulesTest, PrettyPrintIdentity) { ASSERT_EQ(c10::str(Identity()), "torch::nn::Identity()"); } +TEST_F(ModulesTest, PrettyPrintFlatten) { + ASSERT_EQ(c10::str(Flatten()), "torch::nn::Flatten()"); + ASSERT_EQ(c10::str(Flatten(FlattenOptions().start_dim(2).end_dim(4))), "torch::nn::Flatten()"); +} + TEST_F(ModulesTest, ReflectionPad1d) { { ReflectionPad1d m(ReflectionPad1dOptions(2)); diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index 5ad56fe027a57..0fbda1e04e728 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -93,7 +93,7 @@ torch.nn.TransformerDecoderLayer|No|No torch.nn.Identity|Yes|No torch.nn.Linear|Yes|No torch.nn.Bilinear|Yes|No -torch.nn.Flatten|No|No +torch.nn.Flatten|Yes|No torch.nn.Dropout|No|No torch.nn.Dropout2d|No|No torch.nn.Dropout3d|No|No diff --git a/torch/csrc/api/include/torch/nn/modules/linear.h b/torch/csrc/api/include/torch/nn/modules/linear.h index bb8db0d1b6bff..22d2138db2f9e 100644 --- a/torch/csrc/api/include/torch/nn/modules/linear.h +++ b/torch/csrc/api/include/torch/nn/modules/linear.h @@ -67,6 +67,31 @@ TORCH_MODULE(Linear); // ============================================================================ +/// A placeholder for Flatten operator +class TORCH_API FlattenImpl : public Cloneable { + public: + explicit FlattenImpl(const FlattenOptions& options_ = {}); + + void reset() override; + + /// Pretty prints the `Flatten` module into the given `stream`. + void pretty_print(std::ostream& stream) const override; + + /// Applies a flatten transform on the `input`. + Tensor forward(const Tensor& input); + + /// The options used to configure this module. + FlattenOptions options; +}; + +/// A `ModuleHolder` subclass for `FlattenImpl`. +/// See the documentation for `FlattenImpl` class to learn what methods it +/// provides, or the documentation for `ModuleHolder` to learn about PyTorch's +/// module storage semantics. +TORCH_MODULE(Flatten); + +// ============================================================================ + /// Applies a billinear transformation with optional bias. class TORCH_API BilinearImpl : public Cloneable { public: diff --git a/torch/csrc/api/include/torch/nn/options/linear.h b/torch/csrc/api/include/torch/nn/options/linear.h index f84fb48822f47..f99e815b55ed5 100644 --- a/torch/csrc/api/include/torch/nn/options/linear.h +++ b/torch/csrc/api/include/torch/nn/options/linear.h @@ -22,6 +22,16 @@ struct TORCH_API LinearOptions { // ============================================================================ +/// Options for the `Flatten` module. +struct TORCH_API FlattenOptions { + /// first dim to flatten + TORCH_ARG(int64_t, start_dim) = 1; + /// last dim to flatten + TORCH_ARG(int64_t, end_dim) = -1; +}; + +// ============================================================================ + /// Options for the `Bilinear` module. struct TORCH_API BilinearOptions { BilinearOptions(int64_t in1_features, int64_t in2_features, int64_t out_features); diff --git a/torch/csrc/api/src/nn/modules/linear.cpp b/torch/csrc/api/src/nn/modules/linear.cpp index c457cffa52dea..0fc134c1af67a 100644 --- a/torch/csrc/api/src/nn/modules/linear.cpp +++ b/torch/csrc/api/src/nn/modules/linear.cpp @@ -61,6 +61,20 @@ Tensor LinearImpl::forward(const Tensor& input) { // ============================================================================ +FlattenImpl::FlattenImpl(const FlattenOptions& options_) : options(options_) {} + +void FlattenImpl::reset() {} + +void FlattenImpl::pretty_print(std::ostream& stream) const { + stream << "torch::nn::Flatten()"; +} + +Tensor FlattenImpl::forward(const Tensor& input) { + return input.flatten(options.start_dim(), options.end_dim()); +} + +// ============================================================================ + BilinearImpl::BilinearImpl(const BilinearOptions& options_) : options(options_) { reset(); } @@ -90,5 +104,6 @@ void BilinearImpl::pretty_print(std::ostream& stream) const { Tensor BilinearImpl::forward(const Tensor& input1, const Tensor& input2) { return F::bilinear(input1, input2, weight, bias); } + } // namespace nn } // namespace torch From 1c436ded44b58cdc24c558151d15fb019f31cf99 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 19:02:35 -0700 Subject: [PATCH 52/64] Remove `test_quantizer.py` and reuse one of its test in `test_quantization.py` (#27269) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27269 Remove `test_quantizer.py`, add and rewrite one of the tests in `test_quantizer` in `test_quantization.py` The conv test is removed for now since conv pattern is still broken, we'll add another test later ghstack-source-id: 92869823 Test Plan: python test/test_quantization.py Imported from OSS Differential Revision: D18182916 fbshipit-source-id: 325b5d8e877228d6a513e3ddf52c974479250d42 --- .jenkins/pytorch/test.sh | 2 +- test/run_test.py | 1 - test/test_quantization.py | 44 +++++++++- test/test_quantizer.py | 164 -------------------------------------- 4 files changed, 43 insertions(+), 168 deletions(-) delete mode 100644 test/test_quantizer.py diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 6741d7d6bf4fa..0b41c519587af 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -110,7 +110,7 @@ test_python_nn() { } test_python_all_except_nn() { - time python test/run_test.py --exclude nn --verbose --bring-to-front quantization quantized quantized_tensor quantized_nn_mods quantizer + time python test/run_test.py --exclude nn --verbose --bring-to-front quantization quantized quantized_tensor quantized_nn_mods assert_git_not_dirty } diff --git a/test/run_test.py b/test/run_test.py index 631fa28a33cb2..52bfff8df0511 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -47,7 +47,6 @@ 'quantized', 'quantized_tensor', 'quantized_nn_mods', - 'quantizer', 'sparse', 'torch', 'type_info', diff --git a/test/test_quantization.py b/test/test_quantization.py index 8b172d6b29d6d..64f4303b7efe4 100644 --- a/test/test_quantization.py +++ b/test/test_quantization.py @@ -31,6 +31,7 @@ AnnotatedSubNestedModel, AnnotatedCustomConfigNestedModel from jit_utils import _tmp_donotuse_dont_inline_everything +from jit_utils import get_forward from hypothesis import given from hypothesis import strategies as st @@ -695,10 +696,49 @@ def test_single_layer(self): [self.calib_data], inplace=False) result_eager = model_eager(self.calib_data[0][0]) - torch._C._jit_pass_quant_fusion(model_script._c._get_module('fc1')._get_method('forward').graph) - result_script = model_script._c._get_method('forward')(self.calib_data[0][0]) + result_script = get_forward(model_script._c)(self.calib_data[0][0]) self.assertEqual(result_eager, result_script) + @unittest.skip("quantization for inlined linear is not working right now") + def test_nested(self): + # Eager mode + eager_model = AnnotatedNestedModel() + # default_per_channel_qconfig is not scriptable right now, + # temporarily change to default_qconfig until default_per_channel_qconfig is fixed + eager_model.sub2.fc1.qconfig = default_qconfig + + # Graph mode + script_model = NestedModel() + # Copy weights for eager_model + script_model.sub1.fc.weight = torch.nn.Parameter(eager_model.sub1.fc.weight.detach()) + script_model.sub1.fc.bias = torch.nn.Parameter(eager_model.sub1.fc.bias.detach()) + script_model.sub2.fc1.weight = torch.nn.Parameter(eager_model.sub2.fc1.module.weight.detach()) + script_model.sub2.fc1.bias = torch.nn.Parameter(eager_model.sub2.fc1.module.bias.detach()) + script_model.sub2.fc2.weight = torch.nn.Parameter(eager_model.sub2.fc2.weight.detach()) + script_model.sub2.fc2.bias = torch.nn.Parameter(eager_model.sub2.fc2.bias.detach()) + script_model.fc3.weight = torch.nn.Parameter(eager_model.fc3.module.weight.detach()) + script_model.fc3.bias = torch.nn.Parameter(eager_model.fc3.module.bias.detach()) + print(eager_model(self.calib_data[0][0])) + # Quantize eager module + quantized_eager_model = quantize(eager_model, test_only_eval_fn, self.calib_data) + + qconfig_dict = { + 'sub2.fc1': default_qconfig, + 'fc3': default_qconfig + } + quantized_script_model = quantize_script( + torch.jit.script(script_model), + qconfig_dict, + test_only_eval_fn, + [self.calib_data], + inplace=False) + + eager_result = quantized_eager_model(self.calib_data[0][0]) + print(get_forward(quantized_script_model._c._get_module('fc3')).graph) + script_result = get_forward(quantized_script_model._c)(self.calib_data[0][0]) + print(eager_result, script_result) + self.assertEqual(eager_result, script_result) + class FunctionalModuleTest(QuantizationTestCase): # Histogram Observers are slow, so have no-deadline to ensure test doesn't time out diff --git a/test/test_quantizer.py b/test/test_quantizer.py deleted file mode 100644 index 3df63f017deab..0000000000000 --- a/test/test_quantizer.py +++ /dev/null @@ -1,164 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -from __future__ import unicode_literals - -import unittest -import torch.jit -from jit_utils import _tmp_donotuse_dont_inline_everything -from torch._jit_internal import Optional -import torch.nn as nn -from common_utils import TestCase, run_tests -from common_quantization import NestedModel, AnnotatedNestedModel -from torch.quantization import QuantStub, DeQuantStub, \ - quantize, default_eval_fn, QConfig - -class Observer(torch.nn.Module): - __annotations__ = {'scale' : Optional[torch.Tensor], 'zero_point': Optional[torch.Tensor]} - - def __init__(self): - super(Observer, self).__init__() - self.dtype = torch.quint8 - self.qscheme = torch.per_tensor_affine - self.scale, self.zero_point = None, None - - def forward(self, x): - self.scale = torch.tensor([2.0]) - self.zero_point = torch.tensor([3]) - return x - - @torch.jit.export - def calculate_qparams(self): - return self.scale, self.zero_point - -class WeightObserver(Observer): - def __init__(self): - super(WeightObserver, self).__init__() - self.dtype = torch.qint8 - -@unittest.skipUnless('fbgemm' in torch.backends.quantized.supported_engines, - " Quantized operations require FBGEMM. FBGEMM is only optimized for CPUs" - " with instruction set support avx2 or newer.") -@unittest.skip("temporarily disable the test") -class QuantizerTestCase(TestCase): - @_tmp_donotuse_dont_inline_everything - def test_default(self): - class TestM(nn.Module): - def __init__(self, qconfig): - super(TestM, self).__init__() - self.conv = nn.Conv2d(3, 1, 3).float() - self.conv.weight.data.fill_(1.0) - self.conv.bias.data.fill_(0.01) - self.qconfig = qconfig - self.quant = QuantStub() - self.dequant = DeQuantStub() - - def forward(self, x): - return self.dequant(self.conv(self.quant(x))) - - class TestScriptM(torch.jit.ScriptModule): - def __init__(self): - super(TestScriptM, self).__init__() - self.conv = nn.Conv2d(3, 1, 3).float() - self.conv.bias.data.fill_(0.01) - - @torch.jit.script_method - def forward(self, x): - y = self.conv(x) - return y - - # Test Data - data = [(torch.randn(10, 3, 10, 10, dtype=torch.float), 1)] - - # Eager mode - fake_qconfig = QConfig(activation=Observer, weight=WeightObserver) - eager_module = TestM(fake_qconfig) - # Script mode - script_module = TestScriptM() - script_module.conv.weight = torch.nn.Parameter(eager_module.conv.weight.detach()) - quantized_eager_module = quantize(eager_module, default_eval_fn, data) - - def get_forward(m): - return m._c._get_method('forward') - # TODO: test jit.script as well - ScriptedObserver = torch.jit.script(Observer()) - ScriptedWeightObserver = torch.jit.script(WeightObserver()) - qconfig_dict = { - '': - QConfig( - activation=ScriptedObserver._c, - weight=ScriptedWeightObserver._c) - } - torch._C._jit_pass_insert_observers(script_module._c, - "forward", - qconfig_dict) - # Run ScriptM Model and Collect statistics - get_forward(script_module)(data[0][0]) - - # Insert quantize and dequantize calls - script_module._c = torch._C._jit_pass_insert_quant_dequant(script_module._c, "forward") - # Note that observer modules are not removed right now - torch._C._jit_pass_quant_fusion(script_module._c._get_method('forward').graph) - get_forward(script_module)(data[0][0]) - eager_result = quantized_eager_module(data[0][0]) - script_result = get_forward(script_module)(data[0][0]) - self.assertEqual(eager_result, script_result) - - @_tmp_donotuse_dont_inline_everything - def test_qconfig_dict(self): - data = [(torch.randn(10, 5, dtype=torch.float) * 20, 1)] - - # Eager mode - qconfig = QConfig(activation=Observer, weight=WeightObserver) - eager_module = AnnotatedNestedModel() - eager_module.fc3.qconfig = qconfig - eager_module.sub2.fc1.qconfig = qconfig - # Assign weights - eager_module.sub1.fc.weight.data.fill_(1.0) - eager_module.sub2.fc1.module.weight.data.fill_(1.0) - eager_module.sub2.fc2.weight.data.fill_(1.0) - eager_module.fc3.module.weight.data.fill_(1.0) - - script_module = torch.jit.script(NestedModel()) - # Copy weights for eager_module - script_module.sub1.fc.weight = eager_module.sub1.fc.weight - script_module.sub2.fc1.weight = eager_module.sub2.fc1.module.weight - script_module.sub2.fc2.weight = eager_module.sub2.fc2.weight - script_module.fc3.weight = eager_module.fc3.module.weight - - # Quantize eager module - quantized_eager_module = quantize(eager_module, default_eval_fn, data) - - def get_forward(m): - return m._c._get_method('forward') - - # Quantize script_module - torch._C._jit_pass_constant_propagation(get_forward(script_module).graph) - - ScriptedObserver = torch.jit.script(Observer()) - ScriptedWeightObserver = torch.jit.script(WeightObserver()) - scripted_qconfig = QConfig( - activation=ScriptedObserver._c, - weight=ScriptedWeightObserver._c) - qconfig_dict = { - 'sub2.fc1': scripted_qconfig, - 'fc3': scripted_qconfig - } - torch._C._jit_pass_insert_observers(script_module._c, - "forward", - qconfig_dict) - - # Run script_module and Collect statistics - get_forward(script_module)(data[0][0]) - - # Insert quantize and dequantize calls - script_module._c = torch._C._jit_pass_insert_quant_dequant(script_module._c, "forward") - # Note that observer modules are not removed right now - torch._C._jit_pass_quant_fusion(script_module._c._get_method('forward').graph) - get_forward(script_module)(data[0][0]) - eager_result = quantized_eager_module(data[0][0]) - script_result = get_forward(script_module)(data[0][0]) - self.assertEqual(eager_result, script_result) - -if __name__ == '__main__': - run_tests() From ec81cd55fccd217f5f5906b5f01d15825b37afc8 Mon Sep 17 00:00:00 2001 From: vishwakftw Date: Tue, 29 Oct 2019 19:22:34 -0700 Subject: [PATCH 53/64] Migrate implementations of triu and tril to a separate file (#28750) Summary: Having them in BatchLinearAlgebra.cpp/.cu seemed out of place, since they are more general purpose and this code was interspersed between LAPACK and MAGMA wrappers as well. Changelog: - Move tril* / triu* to TriangularOps.cpp/.cu Pull Request resolved: https://github.com/pytorch/pytorch/pull/28750 Test Plan: - Builds should complete successfully to ensure that the migration is error-free - Tests should pass to ensure the methods that the front-end is unaffected. Differential Revision: D18205456 Pulled By: soumith fbshipit-source-id: 41966b9ddfe9f196f4d7c6a5e466782c1985d3d9 --- aten/src/ATen/native/BatchLinearAlgebra.cpp | 145 ---------------- aten/src/ATen/native/LinearAlgebraUtils.h | 53 +----- aten/src/ATen/native/TriangularOps.cpp | 158 ++++++++++++++++++ aten/src/ATen/native/TriangularOpsUtils.h | 59 +++++++ .../ATen/native/cuda/BatchLinearAlgebra.cu | 97 ----------- aten/src/ATen/native/cuda/TriangularOps.cu | 110 ++++++++++++ 6 files changed, 328 insertions(+), 294 deletions(-) create mode 100644 aten/src/ATen/native/TriangularOps.cpp create mode 100644 aten/src/ATen/native/TriangularOpsUtils.h create mode 100644 aten/src/ATen/native/cuda/TriangularOps.cu diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp index ed18cc3790b6e..8096848855c63 100644 --- a/aten/src/ATen/native/BatchLinearAlgebra.cpp +++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp @@ -624,151 +624,6 @@ std::tuple _lu_with_info_cpu(const Tensor& self, bool pi return std::make_tuple(self_working_copy, pivots_tensor, infos_tensor); } -// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -template -static void apply_triu_tril_single( - scalar_t* result, scalar_t* self, bool inplace, - int64_t k, int64_t n, int64_t m, - int64_t res_row_stride, int64_t res_col_stride, - int64_t self_row_stride, int64_t self_col_stride) { - - constexpr int64_t zero = 0; - - if (upper) { - at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { - for (auto i = start; i < end; i++) { - for (int64_t j = 0; j < std::min(m, i + k); j++) { - result[i * res_row_stride + j * res_col_stride] = 0; - } - if (!inplace) { // copy the rest of the self if not inplace - for (int64_t j = std::max(zero, i + k); j < m; j++) { - result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; - } - } - } - }); - } else { - at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { - for (auto i = start; i < end; i++) { - for (int64_t j = std::max(zero, i + k + 1); j < m; j++) { - result[i * res_row_stride + j * res_col_stride] = 0; - } - if (!inplace) { // copy the rest of the self if not inplace - for (int64_t j = zero; j < std::min(m, i + k + 1); j++) { - result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; - } - } - } - }); - } -} - -template -void apply_triu_tril(Tensor& result, const Tensor& self, bool inplace, int64_t k) { - auto n = self.size(-2); - auto m = self.size(-1); - auto self_data = self.data_ptr(); - auto self_stride = (self.dim() > 2 && self.stride(-3) > 0) ? self.stride(-3) : 1; - auto batchsize = batchCountTrilTriu(result); - auto self_row_stride = self.stride(-2); - auto self_column_stride = self.stride(-1); - - auto result_data = result.data_ptr(); - int64_t result_stride, result_row_stride, result_column_stride; - if (result_data != self_data) { - result_stride = (result.dim() > 2 && result.stride(-3) > 0) ? result.stride(-3) : 1; - result_row_stride = result.stride(-2); - result_column_stride = result.stride(-1); - } else { - result_stride = self_stride; - result_row_stride = self_row_stride; - result_column_stride = self_column_stride; - } - - at::parallel_for(0, batchsize, 0, [&](int64_t start, int64_t end) { - for (auto b = start; b < end; b++) { - scalar_t* self_batch = &self_data[b * self_stride]; - scalar_t* result_batch = &result_data[b * result_stride]; - apply_triu_tril_single( - result_batch, self_batch, inplace, k, n, m, - result_row_stride, result_column_stride, self_row_stride, self_column_stride); - } - }); -} - -Tensor tril(const Tensor& self, int64_t k) { - Tensor result = at::empty({0}, self.options()); - at::tril_out(result, self, k); - return result; -} - -Tensor& tril_cpu_(Tensor &self, int64_t k) { - if (self.numel() == 0) { - return self; - } - bool inplace; - Tensor self_c; - std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); - Tensor result = inplace ? self : at::empty_like(self); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ - apply_triu_tril(result, self_c, inplace, k); - }); - if (!inplace) self.copy_(result); - return self; -} - -Tensor& tril_cpu_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - Tensor self_c; - std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ - apply_triu_tril(result, self_c, false, k); - }); - return result; -} - -Tensor triu(const Tensor& self, int64_t k) { - Tensor result = at::empty({0}, self.options()); - at::triu_out(result, self, k); - return result; -} - -Tensor& triu_cpu_(Tensor &self, int64_t k) { - if (self.numel() == 0) { - return self; - } - bool inplace; - Tensor self_c; - std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); - Tensor result = inplace ? self : at::empty_like(self); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ - apply_triu_tril(result, self_c, inplace, k); - }); - if (!inplace) self.copy_(result); - return self; -} - -Tensor& triu_cpu_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - Tensor self_c; - std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ - apply_triu_tril(result, self_c, false, k); - }); - return result; -} - // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triangular_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ template diff --git a/aten/src/ATen/native/LinearAlgebraUtils.h b/aten/src/ATen/native/LinearAlgebraUtils.h index e3d7ee52bfe3f..795f403d94f81 100644 --- a/aten/src/ATen/native/LinearAlgebraUtils.h +++ b/aten/src/ATen/native/LinearAlgebraUtils.h @@ -1,4 +1,4 @@ -#include +#include #include #include #include @@ -45,57 +45,6 @@ static inline int64_t matrixStride(const Tensor& batched_matrices) { return batched_matrices.size(-1) * batched_matrices.size(-2); } -/* - * Given batches of matrices with arbitrary batch dim, - * computes the number of batches for Triu and Tril. This ignores stride 0 dimension - */ -static inline int64_t batchCountTrilTriu(const Tensor& batched_matrices) { - int64_t result = 1; - for (int64_t i = 0; i < batched_matrices.ndimension() - 2; i++) { - if (batched_matrices.stride(i) != 0) { - result *= batched_matrices.size(i); - } - } - return result; -} - -/* Checks a necessary property for the triu and tril implementations, hence the name. - * Here batch contiguity is checked for tensors with greater than 4 dimensions. - * Contiguous tensors and tensors with less than 3 dimensions pass this check - */ -static inline std::tuple checkTrilTriuBatchContiguous(const Tensor& tensor, bool allow_zero_stride) { - // Complete contiguity is the most desired property, which is why - // we return true if the tensor is contiguous - if (tensor.is_contiguous()) { - auto default_strides_for_size = at::detail::defaultStrides(tensor.sizes()); - if (tensor.strides() == default_strides_for_size) { - return std::make_tuple(true, tensor); - } else { - return std::make_tuple(false, tensor.as_strided(tensor.sizes(), default_strides_for_size)); - } - } - - int64_t dims = tensor.dim(); - - // Tensors with dimension less than 4 are handled by default - if (allow_zero_stride && dims <= 3) { - return std::make_tuple(true, tensor); - } - - int64_t expected_stride = tensor.size(-1) * tensor.size(-2); - for (int64_t i = dims - 3; i >= 0; i--) { - // Skip trivial dimension; - if (allow_zero_stride && i == 0 && (tensor.stride(i) == 0 || tensor.size(i) == 1)) { - continue; - } - if (expected_stride != tensor.stride(i)) { - return std::make_tuple(false, tensor.contiguous()); - } - expected_stride *= tensor.size(i); - } - return std::make_tuple(true, tensor); -} - // Returns the epsilon value for floating types except half static inline double _get_epsilon(const ScalarType& sc_type) { switch (sc_type) { diff --git a/aten/src/ATen/native/TriangularOps.cpp b/aten/src/ATen/native/TriangularOps.cpp new file mode 100644 index 0000000000000..e712b5082102b --- /dev/null +++ b/aten/src/ATen/native/TriangularOps.cpp @@ -0,0 +1,158 @@ +#include +#include +#include +#include + +#include +#include + +namespace at { +namespace native { + +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +template +static void apply_triu_tril_single( + scalar_t* result, scalar_t* self, bool inplace, + int64_t k, int64_t n, int64_t m, + int64_t res_row_stride, int64_t res_col_stride, + int64_t self_row_stride, int64_t self_col_stride) { + + constexpr int64_t zero = 0; + + if (upper) { + at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { + for (auto i = start; i < end; i++) { + for (int64_t j = 0; j < std::min(m, i + k); j++) { + result[i * res_row_stride + j * res_col_stride] = 0; + } + if (!inplace) { // copy the rest of the self if not inplace + for (int64_t j = std::max(zero, i + k); j < m; j++) { + result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; + } + } + } + }); + } else { + at::parallel_for(0, n, 0, [&](int64_t start, int64_t end) { + for (auto i = start; i < end; i++) { + for (int64_t j = std::max(zero, i + k + 1); j < m; j++) { + result[i * res_row_stride + j * res_col_stride] = 0; + } + if (!inplace) { // copy the rest of the self if not inplace + for (int64_t j = zero; j < std::min(m, i + k + 1); j++) { + result[i * res_row_stride + j * res_col_stride] = self[i * self_row_stride + j * self_col_stride]; + } + } + } + }); + } +} + +template +void apply_triu_tril(Tensor& result, const Tensor& self, bool inplace, int64_t k) { + auto n = self.size(-2); + auto m = self.size(-1); + auto self_data = self.data_ptr(); + auto self_stride = (self.dim() > 2 && self.stride(-3) > 0) ? self.stride(-3) : 1; + auto batchsize = batchCountTrilTriu(result); + auto self_row_stride = self.stride(-2); + auto self_column_stride = self.stride(-1); + + auto result_data = result.data_ptr(); + int64_t result_stride, result_row_stride, result_column_stride; + if (result_data != self_data) { + result_stride = (result.dim() > 2 && result.stride(-3) > 0) ? result.stride(-3) : 1; + result_row_stride = result.stride(-2); + result_column_stride = result.stride(-1); + } else { + result_stride = self_stride; + result_row_stride = self_row_stride; + result_column_stride = self_column_stride; + } + + at::parallel_for(0, batchsize, 0, [&](int64_t start, int64_t end) { + for (auto b = start; b < end; b++) { + scalar_t* self_batch = &self_data[b * self_stride]; + scalar_t* result_batch = &result_data[b * result_stride]; + apply_triu_tril_single( + result_batch, self_batch, inplace, k, n, m, + result_row_stride, result_column_stride, self_row_stride, self_column_stride); + } + }); +} + +Tensor tril(const Tensor& self, int64_t k) { + Tensor result = at::empty({0}, self.options()); + at::tril_out(result, self, k); + return result; +} + +Tensor& tril_cpu_(Tensor &self, int64_t k) { + if (self.numel() == 0) { + return self; + } + bool inplace; + Tensor self_c; + std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); + Tensor result = inplace ? self : at::empty_like(self); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ + apply_triu_tril(result, self_c, inplace, k); + }); + if (!inplace) self.copy_(result); + return self; +} + +Tensor& tril_cpu_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + Tensor self_c; + std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "tril", [&]{ + apply_triu_tril(result, self_c, false, k); + }); + return result; +} + +Tensor triu(const Tensor& self, int64_t k) { + Tensor result = at::empty({0}, self.options()); + at::triu_out(result, self, k); + return result; +} + +Tensor& triu_cpu_(Tensor &self, int64_t k) { + if (self.numel() == 0) { + return self; + } + bool inplace; + Tensor self_c; + std::tie(inplace, self_c) = checkTrilTriuBatchContiguous(self, true); + Tensor result = inplace ? self : at::empty_like(self); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ + apply_triu_tril(result, self_c, inplace, k); + }); + if (!inplace) self.copy_(result); + return self; +} + +Tensor& triu_cpu_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + Tensor self_c; + std::tie(std::ignore, self_c) = checkTrilTriuBatchContiguous(self, false); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), "triu", [&]{ + apply_triu_tril(result, self_c, false, k); + }); + return result; +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/TriangularOpsUtils.h b/aten/src/ATen/native/TriangularOpsUtils.h new file mode 100644 index 0000000000000..33fb38bd8ca0b --- /dev/null +++ b/aten/src/ATen/native/TriangularOpsUtils.h @@ -0,0 +1,59 @@ +#include +#include + +namespace at { +namespace native { + +/* + * Given batches of matrices with arbitrary batch dim, + * computes the number of batches for Triu and Tril. This ignores stride 0 dimension + */ +static inline int64_t batchCountTrilTriu(const Tensor& batched_matrices) { + int64_t result = 1; + for (int64_t i = 0; i < batched_matrices.ndimension() - 2; i++) { + if (batched_matrices.stride(i) != 0) { + result *= batched_matrices.size(i); + } + } + return result; +} + +/* Checks a necessary property for the triu and tril implementations, hence the name. + * Here batch contiguity is checked for tensors with greater than 4 dimensions. + * Contiguous tensors and tensors with less than 3 dimensions pass this check + */ +static inline std::tuple checkTrilTriuBatchContiguous(const Tensor& tensor, bool allow_zero_stride) { + // Complete contiguity is the most desired property, which is why + // we return true if the tensor is contiguous + if (tensor.is_contiguous()) { + auto default_strides_for_size = at::detail::defaultStrides(tensor.sizes()); + if (tensor.strides() == default_strides_for_size) { + return std::make_tuple(true, tensor); + } else { + return std::make_tuple(false, tensor.as_strided(tensor.sizes(), default_strides_for_size)); + } + } + + int64_t dims = tensor.dim(); + + // Tensors with dimension less than 4 are handled by default + if (allow_zero_stride && dims <= 3) { + return std::make_tuple(true, tensor); + } + + int64_t expected_stride = tensor.size(-1) * tensor.size(-2); + for (int64_t i = dims - 3; i >= 0; i--) { + // Skip trivial dimension; + if (allow_zero_stride && i == 0 && (tensor.stride(i) == 0 || tensor.size(i) == 1)) { + continue; + } + if (expected_stride != tensor.stride(i)) { + return std::make_tuple(false, tensor.contiguous()); + } + expected_stride *= tensor.size(i); + } + return std::make_tuple(true, tensor); +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu index f67d78056040a..7ac2beb6211fe 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu @@ -929,103 +929,6 @@ std::tuple _lu_with_info_cuda(const Tensor& self, bool p return std::make_tuple(self_working_copy, pivots_tensor, infos_tensor); } -// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -template -#ifdef __HIP_PLATFORM_HCC__ -C10_LAUNCH_BOUNDS_1(512) -#endif -__global__ -void triu_tril_kernel( - cuda::detail::TensorInfo result_info, - const cuda::detail::TensorInfo self_info, - const int64_t k, const int64_t N) { - int64_t linear_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (linear_idx >= N) { - return; - } - - auto dims = self_info.dims; - - IndexType self_offset = 0, result_offset = 0; - // Compute column index and corresponding offset - IndexType col = linear_idx % self_info.sizes[dims - 1]; - linear_idx /= self_info.sizes[dims - 1]; - self_offset += self_info.strides[dims - 1] * col; - result_offset += result_info.strides[dims - 1] * col; - - // Compute row index and corresponding offset - IndexType row = linear_idx % self_info.sizes[dims - 2]; - linear_idx /= self_info.sizes[dims - 2]; - self_offset += self_info.strides[dims - 2] * row; - result_offset += result_info.strides[dims - 2] * row; - - // Compute remaining offsets - IndexType running_index; - #pragma unroll - for (IndexType i = dims - 3; i >= 0; --i) { - running_index = linear_idx % self_info.sizes[i]; - linear_idx /= self_info.sizes[i]; - self_offset += running_index * self_info.strides[i]; - result_offset += running_index * result_info.strides[i]; - } - - bool mask = upper ? (col - row >= k) : (col - row <= k); - result_info.data[result_offset] = mask ? self_info.data[self_offset] : scalar_t(0); -} - -template -Tensor& triu_tril_cuda_template(Tensor& result, const Tensor& self, int64_t k, const char* name) { - int64_t N = self.numel(); - dim3 dim_block = cuda::getApplyBlock(); - dim3 dim_grid((N + dim_block.x - 1) / dim_block.x); - AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), name, [&]{ - if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) { - auto result_info = cuda::detail::getTensorInfo(result); - auto self_info = cuda::detail::getTensorInfo(self); - triu_tril_kernel - <<>>( - result_info, self_info, k, N); - } else { - auto result_info = cuda::detail::getTensorInfo(result); - auto self_info = cuda::detail::getTensorInfo(self); - triu_tril_kernel - <<>>( - result_info, self_info, k, N); - } - }); - AT_CUDA_CHECK(cudaGetLastError()); - return result; -} - -Tensor& tril_cuda_(Tensor &self, int64_t k) { - return tril_cuda_out(self, self, k); -} - -Tensor& tril_cuda_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - return triu_tril_cuda_template(result, self, k, "tril"); -} - -Tensor& triu_cuda_(Tensor &self, int64_t k) { - return triu_cuda_out(self, self, k); -} - -Tensor& triu_cuda_out(Tensor &result, const Tensor& self, int64_t k) { - if (result.sizes() != self.sizes()) { - result.resize_as_(self); - } - if (self.numel() == 0) { - return result; - } - return triu_tril_cuda_template(result, self, k, "triu"); -} - // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triangular_solve ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ template diff --git a/aten/src/ATen/native/cuda/TriangularOps.cu b/aten/src/ATen/native/cuda/TriangularOps.cu new file mode 100644 index 0000000000000..4a5e1e64aadb7 --- /dev/null +++ b/aten/src/ATen/native/cuda/TriangularOps.cu @@ -0,0 +1,110 @@ +#include +#include +#include +#include + +#include +#include + +namespace at { +namespace native { + +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ triu/tril ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +template +#ifdef __HIP_PLATFORM_HCC__ +C10_LAUNCH_BOUNDS_1(512) +#endif +__global__ +void triu_tril_kernel( + cuda::detail::TensorInfo result_info, + const cuda::detail::TensorInfo self_info, + const int64_t k, const int64_t N) { + int64_t linear_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (linear_idx >= N) { + return; + } + + auto dims = self_info.dims; + + IndexType self_offset = 0, result_offset = 0; + // Compute column index and corresponding offset + IndexType col = linear_idx % self_info.sizes[dims - 1]; + linear_idx /= self_info.sizes[dims - 1]; + self_offset += self_info.strides[dims - 1] * col; + result_offset += result_info.strides[dims - 1] * col; + + // Compute row index and corresponding offset + IndexType row = linear_idx % self_info.sizes[dims - 2]; + linear_idx /= self_info.sizes[dims - 2]; + self_offset += self_info.strides[dims - 2] * row; + result_offset += result_info.strides[dims - 2] * row; + + // Compute remaining offsets + IndexType running_index; + #pragma unroll + for (IndexType i = dims - 3; i >= 0; --i) { + running_index = linear_idx % self_info.sizes[i]; + linear_idx /= self_info.sizes[i]; + self_offset += running_index * self_info.strides[i]; + result_offset += running_index * result_info.strides[i]; + } + + bool mask = upper ? (col - row >= k) : (col - row <= k); + result_info.data[result_offset] = mask ? self_info.data[self_offset] : scalar_t(0); +} + +template +Tensor& triu_tril_cuda_template(Tensor& result, const Tensor& self, int64_t k, const char* name) { + int64_t N = self.numel(); + dim3 dim_block = cuda::getApplyBlock(); + dim3 dim_grid((N + dim_block.x - 1) / dim_block.x); + AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::Bool, self.scalar_type(), name, [&]{ + if (cuda::detail::canUse32BitIndexMath(result) && cuda::detail::canUse32BitIndexMath(self)) { + auto result_info = cuda::detail::getTensorInfo(result); + auto self_info = cuda::detail::getTensorInfo(self); + triu_tril_kernel + <<>>( + result_info, self_info, k, N); + } else { + auto result_info = cuda::detail::getTensorInfo(result); + auto self_info = cuda::detail::getTensorInfo(self); + triu_tril_kernel + <<>>( + result_info, self_info, k, N); + } + }); + AT_CUDA_CHECK(cudaGetLastError()); + return result; +} + +Tensor& tril_cuda_(Tensor &self, int64_t k) { + return tril_cuda_out(self, self, k); +} + +Tensor& tril_cuda_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + return triu_tril_cuda_template(result, self, k, "tril"); +} + +Tensor& triu_cuda_(Tensor &self, int64_t k) { + return triu_cuda_out(self, self, k); +} + +Tensor& triu_cuda_out(Tensor &result, const Tensor& self, int64_t k) { + if (result.sizes() != self.sizes()) { + result.resize_as_(self); + } + if (self.numel() == 0) { + return result; + } + return triu_tril_cuda_template(result, self, k, "triu"); +} + +} // namespace native +} // namespace at From 400293fcc63c24ec330e71066144c6911dd345f9 Mon Sep 17 00:00:00 2001 From: Shen Li Date: Tue, 29 Oct 2019 19:37:14 -0700 Subject: [PATCH 54/64] Support remote for builtin operators in distributed autograd (#28630) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28630 This includes: 1. Respect autograd context in rpc.remote for builtin ops 2. Force setting autograd context in RRef.to_here() even if the message for to_here() does not contain any tensor. Test Plan: Imported from OSS Differential Revision: D18138562 Pulled By: mrshenli fbshipit-source-id: a39ec83e556d19130f22eb317927241a017000ba --- test/dist_autograd_test.py | 138 ++++++++++++------ torch/csrc/distributed/autograd/utils.cpp | 48 +++--- torch/csrc/distributed/autograd/utils.h | 12 +- .../csrc/distributed/rpc/python_functions.cpp | 11 +- .../distributed/rpc/request_callback_impl.cpp | 10 +- torch/csrc/distributed/rpc/rref.cpp | 33 ++++- torch/csrc/distributed/rpc/rref_proto.cpp | 19 ++- torch/csrc/distributed/rpc/rref_proto.h | 13 +- torch/csrc/distributed/rpc/utils.cpp | 3 + 9 files changed, 201 insertions(+), 86 deletions(-) diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py index 16e8266597598..acfdb07679b34 100644 --- a/test/dist_autograd_test.py +++ b/test/dist_autograd_test.py @@ -85,7 +85,8 @@ def backward(ctx, input): class ExecMode(Enum): LOCAL = 1 # Run the operation locally. - REMOTE = 2 # Run the operation using RPC. + RPC_SYNC = 2 # Run the operation using rpc_sync + REMOTE = 3 # Run the operation using remote. @unittest.skipIf( @@ -98,9 +99,15 @@ def _exec_func(self, exec_mode, method, *args): if len(args) == 1 and isinstance(args[0], list): return method(*args[0]) return method(*args) - else: + elif ExecMode.RPC_SYNC == exec_mode: return rpc.rpc_sync('worker{}'.format(self._next_rank()), method, args=(args)) + elif ExecMode.REMOTE == exec_mode: + rref = rpc.remote('worker{}'.format(self._next_rank()), method, + args=(args)) + return rref.to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) def _next_rank(self): if hasattr(self, 'dst_rank'): @@ -124,7 +131,7 @@ def world_size(self): def init_method(self): return INIT_METHOD_TEMPLATE.format(file_name=self.file_name) - @dist_init(setup_model_parallel=True) + @dist_init def test_autograd_context(self): # Verify max possible id. max_auto_increment = 281474976710655 @@ -150,7 +157,7 @@ def test_autograd_context(self): ): dist_autograd._retrieve_context(context_id) - @dist_init(setup_model_parallel=True) + @dist_init def test_nested_context(self): with dist_autograd.context() as context_id: # Nested contexts not supported. @@ -255,7 +262,7 @@ def _verify_graph_for_nested_rpc_call(self, ctx): "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() ) - def _test_graph(self, fn): + def _test_graph(self, fn, exec_mode): dst_rank = (self.rank + 1) % self.world_size # This is for the below `dist.barrier`. @@ -272,7 +279,16 @@ def _test_graph(self, fn): with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) - ret = rpc.rpc_sync("worker{}".format(dst_rank), fn, args=(t1, t2)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), fn, args=(t1, t2)) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), fn, args=(t1, t2) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format(dst_rank), _set_rpc_done, args=(context_id, 1)) @@ -306,16 +322,20 @@ def _test_graph(self, fn): with self.assertRaises(RuntimeError): ctx = dist_autograd._current_context() - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_builtin_call(self): - self._test_graph(torch.add) + self._test_graph(torch.add, ExecMode.RPC_SYNC) - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_python_call(self): - self._test_graph(my_py_add) + self._test_graph(my_py_add, ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_builtin_remote_call(self): + self._test_graph(torch.add, ExecMode.REMOTE) # 3-layer nested calls - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_py_nested_call(self): dst_rank = (self.rank + 1) % self.world_size @@ -381,7 +401,7 @@ def test_graph_for_py_nested_call(self): dist.barrier() # Rank0->Rank1->Rank0 - @dist_init(setup_model_parallel=True) + @dist_init def test_graph_for_py_nested_call_itself(self): dst_rank = (self.rank + 1) % self.world_size @@ -430,7 +450,7 @@ def test_graph_for_py_nested_call_itself(self): # autograd context before another worker tries to access it. dist.barrier() - @dist_init(setup_model_parallel=True) + @dist_init def test_no_graph_with_tensors_not_require_grad(self): dst_rank = (self.rank + 1) % self.world_size with dist_autograd.context() as context_id: @@ -452,16 +472,28 @@ def test_no_graph_with_tensors_not_require_grad(self): with self.assertRaises(RuntimeError): ctx = dist_autograd._retrieve_context(ctx_ids[1]) - @dist_init(setup_model_parallel=True) - def test_rpc_complex_args(self): + def _test_rpc_complex_args(self, exec_mode): with dist_autograd.context() as context_id: num_tensors = 10 tensors = [] for i in range(num_tensors): tensors.append(torch.ones(3, 3, requires_grad=(i % 2 == 0))) - ret = rpc.rpc_sync( - "worker{}".format(self._next_rank()), torch.stack, args=(tensors,) - ) + + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(self._next_rank()), + torch.stack, + args=(tensors,) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(self._next_rank()), + torch.stack, + args=(tensors,) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + self.assertEqual(torch.stack(tensors), ret) # Verify appropriate tensors have been attached the autograd graph. @@ -485,8 +517,15 @@ def test_rpc_complex_args(self): dst_rank = (self.rank + 1) % self.world_size self.assertEqual(worker_ids[0], dst_rank) + @dist_init + def test_rpc_complex_args(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC) + + @dist_init + def test_remote_complex_args(self): + self._test_rpc_complex_args(ExecMode.REMOTE) - @dist_init(setup_model_parallel=True) + @dist_init def test_context_cleanup_many_workers(self): dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} with dist_autograd.context() as context_id: @@ -502,7 +541,7 @@ def test_context_cleanup_many_workers(self): success = _all_contexts_cleaned_up() self.assertTrue(success) - @dist_init(setup_model_parallel=True) + @dist_init def test_context_cleanup_nested_rpc(self): # This is for the below `dist.barrier`. # For `RpcAgent` other than `ProcessGroupAgent`, @@ -530,7 +569,7 @@ def test_context_cleanup_nested_rpc(self): success = _all_contexts_cleaned_up() self.assertTrue(success) - @dist_init(setup_model_parallel=True) + @dist_init def test_worker_ids_recorded(self): dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} with dist_autograd.context() as context_id: @@ -561,7 +600,7 @@ def test_worker_ids_recorded(self): self.assertEqual(len(worker_ids), len(dst_ranks)) self.assertEqual(set(worker_ids), dst_ranks) - @dist_init(setup_model_parallel=True) + @dist_init def test_error_in_context(self): with dist_autograd.context() as context_id: t1 = torch.rand(3, 3, requires_grad=True) @@ -574,11 +613,11 @@ def test_error_in_context(self): args=(t1, t2)) def _verify_backwards(self, exec_mode, tensors, context_id, local_grads, *args): - if exec_mode == ExecMode.REMOTE: - self._verify_backwards_remote(tensors, context_id, local_grads, *args) - else: + if exec_mode == ExecMode.LOCAL: torch.autograd.backward(tensors) return [arg.grad for arg in args] + else: + self._verify_backwards_remote(tensors, context_id, local_grads, *args) def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): dist_autograd.backward(tensors) @@ -598,20 +637,21 @@ def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): self.assertEqual(ngrads, len(grads)) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_simple(self): # Run the same code locally and with dist autograd and verify gradients # are same. local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: ret = self._exec_func(exec_mode, torch.add, t1, t2) loss = ret.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_multiple_round_trips(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) @@ -620,7 +660,7 @@ def test_backward_multiple_round_trips(self): t4 = torch.rand((3, 3)) t5 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: # Multiple RPCs between different nodes. val = self._exec_func(exec_mode, torch.add, t1, t2) @@ -631,9 +671,10 @@ def test_backward_multiple_round_trips(self): val = self._exec_func(exec_mode, torch.matmul, val, val) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_different_tensor_dims(self): local_grads = None t1 = torch.rand((4, 6), requires_grad=True) @@ -641,33 +682,35 @@ def test_backward_different_tensor_dims(self): t3 = torch.rand((5, 7), requires_grad=True) t4 = torch.rand((7, 9)) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: val = self._exec_func(exec_mode, torch.matmul, t1, t2) val = self._exec_func(exec_mode, torch.chain_matmul, [val, t3, t4]) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_unused_tensors(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) t3 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: s = self._exec_func(exec_mode, torch.stack, (t1, t2, t3)) val = self._exec_func(exec_mode, torch.matmul, torch.narrow(s, 0, 0, 1), torch.narrow(s, 0, 2, 1)) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2, t3) + local_grads = ret if ret else local_grads - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_multiple_output_tensors(self): local_grads = None t = torch.rand((10, 2), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: with dist_autograd.context() as context_id: tensor_list = self._exec_func(exec_mode, torch.split, t, 2) t1 = tensor_list[0] @@ -677,7 +720,8 @@ def test_backward_multiple_output_tensors(self): val = self._exec_func(exec_mode, torch.chain_matmul, [t1, t2, t3]) loss = val.sum() - local_grads = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t) + ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t) + local_grads = ret if ret else local_grads def _run_test_backward_unused_send_function_in_thread(self): with dist_autograd.context() as context_id: @@ -695,7 +739,7 @@ def _run_test_backward_unused_send_function_in_thread(self): dist_autograd.backward([val.sum()]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_unused_send_function(self): # Run the test in a thread which would never finish. t = threading.Thread(target=self._run_test_backward_unused_send_function_in_thread) @@ -706,7 +750,7 @@ def test_backward_unused_send_function(self): # Verify thread is still alive (indicating backward hasn't completed yet). self.assertTrue(t.is_alive()) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_autograd_engine_error(self): with dist_autograd.context() as context_id: t1 = torch.rand((3, 3), requires_grad=True) @@ -769,7 +813,7 @@ def test_backward_node_failure(self): # Exit all other nodes. pass - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_without_context(self): t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) @@ -779,7 +823,7 @@ def test_backward_without_context(self): args=(t1, t2)) dist_autograd.backward([res.sum()]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_without_rpc(self): dst_rank = self.rank with dist_autograd.context() as context_id: @@ -795,7 +839,7 @@ def test_backward_without_rpc(self): self.assertEqual(torch.ones(3, 3), grads[t1]) self.assertEqual(torch.ones(3, 3), grads[t2]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_invalid_args(self): with dist_autograd.context() as context_id: @@ -817,12 +861,12 @@ def test_backward_invalid_args(self): t = torch.rand(1, requires_grad=True) dist_autograd.backward([t]) - @dist_init(setup_model_parallel=True) + @dist_init def test_backward_multiple_roots(self): local_grads = None t1 = torch.rand((3, 3), requires_grad=True) t2 = torch.rand((3, 3), requires_grad=True) - for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: with dist_autograd.context() as context_id: r1 = self._exec_func(exec_mode, torch.add, t1, t2).sum() r2 = self._exec_func(exec_mode, torch.mul, t1, t2).sum() diff --git a/torch/csrc/distributed/autograd/utils.cpp b/torch/csrc/distributed/autograd/utils.cpp index 8c0b9602cbc2d..da85418aa0e05 100644 --- a/torch/csrc/distributed/autograd/utils.cpp +++ b/torch/csrc/distributed/autograd/utils.cpp @@ -42,37 +42,45 @@ DistAutogradContext* addRecvRpcBackward( const AutogradMetadata& autogradMetadata, std::vector& tensors, rpc::worker_id_t fromWorkerId) { - TORCH_INTERNAL_ASSERT( - torch::autograd::compute_requires_grad(tensors), - "Received tensors do not require grad, addRecvRpcBackward should not be called"); // Initialize autograd context if necessary. auto& autogradContainer = DistAutogradContainer::getInstance(); DistAutogradContext& autogradContext = autogradContainer.getOrCreateContext(autogradMetadata.autogradContextId); - // Attach the tensors as inputs to the autograd function. - auto grad_fn = std::make_shared( - autogradMetadata, autogradContext, fromWorkerId); - for (auto& tensor : tensors) { - torch::autograd::set_history(tensor, grad_fn); + if (!tensors.empty()) { + TORCH_INTERNAL_ASSERT( + torch::autograd::compute_requires_grad(tensors), + "Received tensors do not require grad, addRecvRpcBackward should not be called"); + + // Attach the tensors as inputs to the autograd function. + auto grad_fn = std::make_shared( + autogradMetadata, autogradContext, fromWorkerId); + for (auto& tensor : tensors) { + torch::autograd::set_history(tensor, grad_fn); + } + + // Now update the autograd context with the necessary information. + autogradContext.addRecvFunction( + grad_fn, autogradMetadata.autogradMessageId); } - // Now update the autograd context with the necessary information. - autogradContext.addRecvFunction(grad_fn, autogradMetadata.autogradMessageId); return &autogradContext; } Message getMessageWithAutograd( const rpc::worker_id_t dstId, torch::distributed::rpc::Message&& wrappedRpcMsg, - MessageType msgType) { + MessageType msgType, + bool forceGradRecording) { auto& autogradContainer = DistAutogradContainer::getInstance(); // If there is no valid context and no tensor requires grads, send original // rpc message. otherwise, attach grad info and grad functions and send // rpcWithAutograd message. + auto tensorsRequireGrad = + torch::autograd::compute_requires_grad(wrappedRpcMsg.tensors()); if (!autogradContainer.hasValidContext() || - !torch::autograd::compute_requires_grad(wrappedRpcMsg.tensors())) { + (!forceGradRecording && !tensorsRequireGrad)) { return std::move(wrappedRpcMsg); } @@ -88,9 +96,11 @@ Message getMessageWithAutograd( autogradMetadata, std::move(wrappedRpcMsg)); - // Record autograd information for 'send'. - addSendRpcBackward( - autogradContext, autogradMetadata, rpcWithAutograd->tensors(), dstId); + if (tensorsRequireGrad) { + // Record autograd information for 'send'. + addSendRpcBackward( + autogradContext, autogradMetadata, rpcWithAutograd->tensors(), dstId); + } return std::move(*rpcWithAutograd).toMessage(); } @@ -98,9 +108,13 @@ Message getMessageWithAutograd( std::shared_ptr sendMessageWithAutograd( RpcAgent& agent, const WorkerInfo& dst, - torch::distributed::rpc::Message&& wrappedRpcMsg) { + torch::distributed::rpc::Message&& wrappedRpcMsg, + bool forceGradRecording) { auto msg = getMessageWithAutograd( - dst.id_, std::move(wrappedRpcMsg), MessageType::FORWARD_AUTOGRAD_REQ); + dst.id_, + std::move(wrappedRpcMsg), + MessageType::FORWARD_AUTOGRAD_REQ, + forceGradRecording); return agent.send(dst, std::move(msg)); } diff --git a/torch/csrc/distributed/autograd/utils.h b/torch/csrc/distributed/autograd/utils.h index 2f368cdf17fa2..cc147c059f1d7 100644 --- a/torch/csrc/distributed/autograd/utils.h +++ b/torch/csrc/distributed/autograd/utils.h @@ -34,19 +34,23 @@ TORCH_API DistAutogradContext* addRecvRpcBackward( // This method is a wrapper utility used internally to wrap autograd info // and attach autograd function for each type of rpc call if it has valid -// context and tensors require grads, in this case, return RpcWithAutograd -// message; otherwise return original rpc message. +// context and tensors require grads or forceGradRecording is true, in this +// case, return RpcWithAutograd message; otherwise return original rpc message. +// NB: forceGradRecording is useful when the request does not contain any tensor +// but the corresponding response does. TORCH_API rpc::Message getMessageWithAutograd( const rpc::worker_id_t dstId, rpc::Message&& wrappedRpcMsg, - rpc::MessageType msgType); + rpc::MessageType msgType, + bool forceGradRecording = false); // Send message after autograd checking TORCH_API std::shared_ptr sendMessageWithAutograd( rpc::RpcAgent& agent, const rpc::WorkerInfo& dst, - rpc::Message&& wrappedRpcMsg); + rpc::Message&& wrappedRpcMsg, + bool forceGradRecording = false); } // namespace autograd } // namespace distributed diff --git a/torch/csrc/distributed/rpc/python_functions.cpp b/torch/csrc/distributed/rpc/python_functions.cpp index 6a13eedf25ad5..418a3be2f8faa 100644 --- a/torch/csrc/distributed/rpc/python_functions.cpp +++ b/torch/csrc/distributed/rpc/python_functions.cpp @@ -143,11 +143,12 @@ PyRRef pyRemoteBuiltin( ctx.getWorkerId() != dst.id_, "Does not support creating RRef on self yet."); auto userRRef = ctx.createUserRRef(dst.id_); - auto fm = agent.send( - dst, - ScriptRemoteCall( - op, std::move(stack), userRRef->rrefId(), userRRef->forkId()) - .toMessage()); + + auto scriptRemoteCall = c10::guts::make_unique( + op, std::move(stack), userRRef->rrefId(), userRRef->forkId()); + + auto fm = sendMessageWithAutograd( + agent, dst, std::move(*scriptRemoteCall).toMessage()); ctx.addPendingUser(userRRef->forkId(), userRRef); fm->addCallback(finishAcceptUserRRef); diff --git a/torch/csrc/distributed/rpc/request_callback_impl.cpp b/torch/csrc/distributed/rpc/request_callback_impl.cpp index 2cc7a8cc81df6..a56d5b4463c74 100644 --- a/torch/csrc/distributed/rpc/request_callback_impl.cpp +++ b/torch/csrc/distributed/rpc/request_callback_impl.cpp @@ -82,7 +82,7 @@ Message RequestCallbackImpl::processRpc( ownerRRef->setValue(std::move(stack.front())); ctx.addForkOfOwner(src.retRRefId(), src.retForkId()); - return std::move(RemoteRet(src.retRRefId(), src.retForkId())).toMessage(); + return RemoteRet(src.retRRefId(), src.retForkId()).toMessage(); } case MessageType::PYTHON_REMOTE_CALL: { auto& prc = static_cast(rpc); @@ -96,7 +96,7 @@ Message RequestCallbackImpl::processRpc( ownerRRef->setValue( PythonRpcHandler::getInstance().runPythonUDF(prc.serializedPyObj())); ctx.addForkOfOwner(rrefId, forkId); - return std::move(RemoteRet(rrefId, forkId)).toMessage(); + return RemoteRet(rrefId, forkId).toMessage(); } case MessageType::SCRIPT_RREF_FETCH_CALL: { auto& srf = static_cast(rpc); @@ -104,7 +104,7 @@ Message RequestCallbackImpl::processRpc( // TODO: make this asynchronous std::shared_ptr> rref = ctx.getOrCreateOwnerRRef(srf.rrefId()); - return std::move(ScriptRRefFetchRet({rref->getValue()})).toMessage(); + return ScriptRRefFetchRet({rref->getValue()}).toMessage(); } case MessageType::PYTHON_RREF_FETCH_CALL: { auto& prf = static_cast(rpc); @@ -114,7 +114,7 @@ Message RequestCallbackImpl::processRpc( ctx.getOrCreateOwnerRRef(prf.rrefId()); SerializedPyObj result = PythonRpcHandler::getInstance().serialize(rref->getValue()); - return std::move(PythonRRefFetchRet(result.toIValues())).toMessage(); + return PythonRRefFetchRet(result.toIValues()).toMessage(); } case MessageType::RREF_USER_DELETE: { auto& rud = static_cast(rpc); @@ -132,7 +132,7 @@ Message RequestCallbackImpl::processRpc( auto& rfr = static_cast(rpc); auto& ctx = RRefContext::getInstance(); ctx.addForkOfOwner(rfr.rrefId(), rfr.forkId()); - return std::move(RRefAck()).toMessage(); + return RRefAck().toMessage(); } case MessageType::FORWARD_AUTOGRAD_REQ: { auto& rpcWithAutograd = static_cast(rpc); diff --git a/torch/csrc/distributed/rpc/rref.cpp b/torch/csrc/distributed/rpc/rref.cpp index 1d08f97ca7b0b..60b84b71195ca 100644 --- a/torch/csrc/distributed/rpc/rref.cpp +++ b/torch/csrc/distributed/rpc/rref.cpp @@ -1,7 +1,10 @@ #include +#include +#include #include #include +#include namespace torch { namespace distributed { @@ -133,14 +136,36 @@ template <> std::shared_ptr UserRRef::toHere() { auto future = std::make_shared(nullptr); auto agent = RpcAgent::getDefaultRpcAgent(); - auto futureResponse = agent->send( + + // ScriptRRefFetchCall message always carries autograd context id even if + // the message itself does not contain any tensor, because the response would + // potentially contain tensors. + auto futureResponse = autograd::sendMessageWithAutograd( + *agent, agent->getWorkerInfo(ownerId_), - ScriptRRefFetchCall(rrefId()).toMessage()); + ScriptRRefFetchCall(ownerId_, rrefId()).toMessage(), + true /* forceGradRecording */); futureResponse->addCallback([future](const Message& message) { RRefContext::handleException(message); - auto rfr = ScriptRRefFetchRet::fromMessage(message); - future->markCompleted(rfr->values().front()); + auto response = deserializeResponse(message); + if (message.type() == MessageType::FORWARD_AUTOGRAD_RESP) { + auto& rpcWithAutograd = + static_cast(*response); + + // Attach 'recv' autograd function. + addRecvRpcBackward( + rpcWithAutograd.autogradMetadata(), + rpcWithAutograd.tensors(), + rpcWithAutograd.fromWorkerId()); + + auto& wrappedRpc = rpcWithAutograd.wrappedRpc(); + auto& rfr = static_cast(wrappedRpc); + future->markCompleted(rfr.values().front()); + } else { + auto& rfr = static_cast(*response); + future->markCompleted(rfr.values().front()); + } }); return future; } diff --git a/torch/csrc/distributed/rpc/rref_proto.cpp b/torch/csrc/distributed/rpc/rref_proto.cpp index 740a5b470647f..9d3cb54b676d6 100644 --- a/torch/csrc/distributed/rpc/rref_proto.cpp +++ b/torch/csrc/distributed/rpc/rref_proto.cpp @@ -77,11 +77,26 @@ std::pair ForkMessageBase::fromMessage( /////////////////////////// RRef Protocol ////////////////////////////////// +Message ScriptRRefFetchCall::toMessage() && { + std::vector ivalues; + ivalues.reserve(2); + ivalues.emplace_back(rrefId_.toIValue()); + ivalues.emplace_back(fromWorkerId_); + return fromIValues(std::move(ivalues), MessageType::SCRIPT_RREF_FETCH_CALL); +} + std::unique_ptr ScriptRRefFetchCall::fromMessage( const Message& message) { + auto values = toIValues(message, MessageType::SCRIPT_RREF_FETCH_CALL); + TORCH_INTERNAL_ASSERT( + values.size() == 2, "ScriptRRefFetchCall expects 2 IValues from message"); + auto id = values[1].toInt(); + TORCH_INTERNAL_ASSERT( + id >= std::numeric_limits::min() && + id <= std::numeric_limits::max(), + "ScriptRRefFetchCall fromWorkerId exceeds worker_id_t limit.") return c10::guts::make_unique( - RRefId::fromIValue(RRefMessageBase::fromMessage( - message, MessageType::SCRIPT_RREF_FETCH_CALL))); + worker_id_t(id), RRefId::fromIValue(values[0])); } std::unique_ptr PythonRRefFetchCall::fromMessage( diff --git a/torch/csrc/distributed/rpc/rref_proto.h b/torch/csrc/distributed/rpc/rref_proto.h index 55eeb97b183b0..dc52b9b4509cc 100644 --- a/torch/csrc/distributed/rpc/rref_proto.h +++ b/torch/csrc/distributed/rpc/rref_proto.h @@ -51,11 +51,20 @@ class TORCH_API ForkMessageBase : public RRefMessageBase { // UserRRef uses this message to fetch the remote RRef value from the owner. class TORCH_API ScriptRRefFetchCall final : public RRefMessageBase { public: - explicit ScriptRRefFetchCall(const RRefId& rrefId) - : RRefMessageBase(rrefId, MessageType::SCRIPT_RREF_FETCH_CALL) {} + ScriptRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId) + : RRefMessageBase(rrefId, MessageType::SCRIPT_RREF_FETCH_CALL), + fromWorkerId_(fromWorkerId) {} + inline worker_id_t fromWorkerId() const { + return fromWorkerId_; + } + + Message toMessage() && override; static std::unique_ptr fromMessage( const Message& message); + + private: + const worker_id_t fromWorkerId_; }; class TORCH_API PythonRRefFetchCall final : public RRefMessageBase { diff --git a/torch/csrc/distributed/rpc/utils.cpp b/torch/csrc/distributed/rpc/utils.cpp index df619dcc15829..b4dd2cd606d03 100644 --- a/torch/csrc/distributed/rpc/utils.cpp +++ b/torch/csrc/distributed/rpc/utils.cpp @@ -71,6 +71,9 @@ std::unique_ptr deserializeResponse(const Message& response) { case MessageType::REMOTE_RET: { return RemoteRet::fromMessage(response); } + case MessageType::SCRIPT_RREF_FETCH_RET: { + return ScriptRRefFetchRet::fromMessage(response); + } case MessageType::RREF_ACK: { return RRefAck::fromMessage(response); } From 043530a9b90f26da0fd0469dcc8ab330f252ad65 Mon Sep 17 00:00:00 2001 From: Shen Li Date: Tue, 29 Oct 2019 19:37:14 -0700 Subject: [PATCH 55/64] Support remote for Python UDF in distributed autograd Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28656 Test Plan: Imported from OSS Differential Revision: D18138561 Pulled By: mrshenli fbshipit-source-id: 798e7c00465b5a299f7b4642683bc407895bc7da --- test/dist_autograd_test.py | 286 ++++++++++++++++-- .../csrc/distributed/rpc/python_functions.cpp | 20 +- torch/csrc/distributed/rpc/rref.cpp | 54 ++-- torch/csrc/distributed/rpc/rref_proto.cpp | 19 +- torch/csrc/distributed/rpc/rref_proto.h | 9 +- torch/csrc/distributed/rpc/utils.cpp | 3 + 6 files changed, 341 insertions(+), 50 deletions(-) diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py index acfdb07679b34..45ca4dddba6e6 100644 --- a/test/dist_autograd_test.py +++ b/test/dist_autograd_test.py @@ -36,10 +36,46 @@ def _set_rpc_done(ctx_id, rank_distance): known_context_ids.append(ctx_id) +def _check_rpc_done(rank_distance): + while not rpc_done[rank_distance]: + time.sleep(0.1) + + +def _torch_ones(sizes, requires_grad=False): + return torch.ones(sizes, requires_grad=requires_grad) + + +# creates an owner rref on the given dst, and the rref holds a torch.ones tensor +# of the given size. +def _create_ones_rref_on(dst, sizes): + return rpc.remote( + dst, + _torch_ones, + args=(sizes,), + kwargs={"requires_grad": True} + ) + + +# This method must be called on the rref owner, and verifies that the grad of +# rref tensor equals to the given grad. +def _compare_owner_value(context_id, rref, grad): + grads = dist_autograd.get_gradients(context_id) + return torch.equal(grads[rref.local_value().wait()], grad) + + def my_py_add(t1, t2): return torch.add(t1, t2) +def my_rref_add(rref_t1, t2): + ret = torch.add(rref_t1.local_value().wait(), t2) + return ret + + +def my_nested_rref_add(dst, rref_t1, t2): + return rpc.rpc_sync(dst, my_rref_add, args=(rref_t1, t2)) + + def my_py_nested_call(t1, t2, dst, world_size, hops): next_dst = (dst + 1) % world_size if hops > 0: @@ -68,6 +104,17 @@ def _all_contexts_cleaned_up(timeout_seconds=10): return success +# This function creates a dis atugorad context, run rpc_sync on the given ps, +# and then blocks until the ps has verified the grads are correctly accumulated. +def _run_trainer(rref_t1, t2, ps, rank_diff): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_rref_add, args=(rref_t1, t2)) + dist_autograd.backward([ret.sum()]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0, )) + + from torch.autograd import Function from torch.autograd.function import once_differentiable @@ -119,9 +166,7 @@ def _next_rank(self): return self.dst_rank def _check_rpc_done(self, rank_distance): - while not rpc_done[rank_distance]: - time.sleep(0.1) - pass + _check_rpc_done(rank_distance) @property def world_size(self): @@ -334,9 +379,12 @@ def test_graph_for_python_call(self): def test_graph_for_builtin_remote_call(self): self._test_graph(torch.add, ExecMode.REMOTE) - # 3-layer nested calls @dist_init - def test_graph_for_py_nested_call(self): + def test_graph_for_python_remote_call(self): + self._test_graph(my_py_add, ExecMode.REMOTE) + + # 3-layer nested calls + def _test_graph_for_py_nested_call(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size # This is for the below `dist.barrier`. @@ -354,8 +402,21 @@ def test_graph_for_py_nested_call(self): t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) nest_dst_rank = (dst_rank + 1) % self.world_size - ret = rpc.rpc_sync("worker{}".format(dst_rank), - my_py_nested_call, args=(t1, t2, dst_rank, self.world_size, 1)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + for rd in [1, 2, 3]: rpc.rpc_sync("worker{}".format((self.rank + rd) % self.world_size), _set_rpc_done, args=(context_id, rd)) @@ -400,9 +461,16 @@ def test_graph_for_py_nested_call(self): # autograd context before another worker tries to access it. dist.barrier() - # Rank0->Rank1->Rank0 @dist_init - def test_graph_for_py_nested_call_itself(self): + def test_graph_for_py_nested_call(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_py_nested_remote_call(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE) + + # Rank0->Rank1->Rank0 + def _test_graph_for_py_nested_call_itself(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size # This is for the below `dist.barrier`. @@ -419,9 +487,33 @@ def test_graph_for_py_nested_call_itself(self): with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=True) t2 = torch.zeros(3, 3, requires_grad=True) - ret = rpc.rpc_sync("worker{}".format(dst_rank), - my_py_nested_call, - args=(t1, t2, (self.rank - 1 + self.world_size) % self.world_size, self.world_size, 0)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0 + ) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0 + ) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format((self.rank + 1) % self.world_size), _set_rpc_done, args=(context_id, 1)) @@ -451,12 +543,33 @@ def test_graph_for_py_nested_call_itself(self): dist.barrier() @dist_init - def test_no_graph_with_tensors_not_require_grad(self): + def test_graph_for_py_nested_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC) + + @dist_init + def test_graph_for_py_nested_remote_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE) + + def _test_no_graph_with_tensors_not_require_grad(self, exec_mode): dst_rank = (self.rank + 1) % self.world_size with dist_autograd.context() as context_id: t1 = torch.ones(3, 3, requires_grad=False) t2 = torch.zeros(3, 3, requires_grad=False) - ret = rpc.rpc_sync("worker{}".format(dst_rank), torch.add, args=(t1, t2)) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + "worker{}".format(dst_rank), + torch.add, + args=(t1, t2) + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + "worker{}".format(dst_rank), + torch.add, + args=(t1, t2) + ).to_here().wait() + else: + raise ValueError("Unrecognized ExecMode {}".format(exec_mode)) + rpc.rpc_sync("worker{}".format(dst_rank), _set_rpc_done, args=(context_id, 1)) @@ -468,9 +581,24 @@ def test_no_graph_with_tensors_not_require_grad(self): # Wait for the prev rank to be done with rpc. self._check_rpc_done(1) - # prev context id is not passed over as tensors do not require grads - with self.assertRaises(RuntimeError): - ctx = dist_autograd._retrieve_context(ctx_ids[1]) + if ExecMode.RPC_SYNC == exec_mode: + # prev context id is not passed over as tensors do not require + # grads + with self.assertRaises(RuntimeError): + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + elif ExecMode.REMOTE == exec_mode: + # NB: RRef.to_here() always passes the autograd context to the + # the callee, as the caller does not know whether the return + # value would contain a requires_grad tensor or not. + pass + + @dist_init + def test_no_graph_with_tensors_not_require_grad(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE) def _test_rpc_complex_args(self, exec_mode): with dist_autograd.context() as context_id: @@ -636,7 +764,6 @@ def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): self.assertEqual(ngrads, len(grads)) - @dist_init def test_backward_simple(self): # Run the same code locally and with dist autograd and verify gradients @@ -651,6 +778,129 @@ def test_backward_simple(self): ret = self._verify_backwards(exec_mode, [loss], context_id, local_grads, t1, t2) local_grads = ret if ret else local_grads + # The current rank first creates a tensor on the rref_owner, and then passes + # the rref with another tensor to the callee to run either my_rref_add or + # my_nested_rref_add, depending on whether the callee is the rref owner. + # The grad of tensor lives on the current rank, and the grad of the rref + # tensor lives on the rref owner. + def _test_backward_rref(self, callee, rref_owner): + local_grads = None + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + local_ret.sum().backward() + with dist_autograd.context() as context_id: + rref_t1 = rpc.remote( + rref_owner, + _torch_ones, + args=((3, 3),), + kwargs={"requires_grad": True} + ) + + if callee == rref_owner: + rref = rpc.remote(callee, my_rref_add, args=(rref_t1, t2)) + else: + rref = rpc.remote( + callee, + my_nested_rref_add, + args=(rref_owner, rref_t1, t2) + ) + ret = rref.to_here().wait() + dist_autograd.backward([ret.sum()]) + + # verify grads on caller + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t2, grads) + self.assertEqual(grads[t2], t2.grad) + + # verify grads on rref owner + self.assertTrue( + rpc.rpc_sync( + rref_owner, + _compare_owner_value, + args=(context_id, rref_t1, t1.grad) + ) + ) + + @dist_init + def test_backward_rref(self): + callee = "worker{}".format(self._next_rank()) + rref_owner = callee + self._test_backward_rref(callee, rref_owner) + + @dist_init + def test_backward_rref_multi(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._test_backward_rref(callee, rref_owner) + + @dist_init + def test_backward_rref_nested(self): + callee = "worker{}".format((self.rank + 1) % self.world_size) + rref_owner = "worker{}".format((self.rank + 2) % self.world_size) + self._test_backward_rref(callee, rref_owner) + + # In this test, every rank will serve as a parameter server (ps) and a + # driver, and then kicks off trainers on the other three ranks. So, we have: + # ps = rank0 with trainers = rank1/2/3 + # ps = rank2 with trainers = rank2/3/0 + # ps = rank3 with trainers = rank3/0/1 + # ps = rank4 with trainers = rank0/1/2 + # + # These four test ps-trainer groups run on completely separate autograd + # graphs, but they share the same set of underlying RpcAgents. + @dist_init + def test_trainer_ps(self): + local_grads = None + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + local_ret.sum().backward() + + # create rref on self + # TODO: simplify this once we support rpc to self + self_name = "worker{}".format(self.rank) + rref_t1 = rpc.rpc_sync( + "worker{}".format(self._next_rank()), + _create_ones_rref_on, + args=(self_name, (3, 3)) + ) + + # kick off forward and backward pass on three other workers (trainers) + rank_diffs = [1, 2, 3] + futures = [] + for rank_diff in rank_diffs: + futures.append(rpc.rpc_async( + "worker{}".format((self.rank + rank_diff) % self.world_size), + _run_trainer, + args=(rref_t1, t2, self_name, rank_diff) + )) + + # check if the trainers have done with their backward pass + for rank_diff in rank_diffs: + self._check_rpc_done(rank_diff) + + # trainers are done and holding the context for verification + accumulate_grad_func = None + for rank_diff in rank_diffs: + # make sure grads are accumulated for the same tensors and values + # are all correct + ctx_id = ctx_ids[rank_diff] + grads = dist_autograd.get_gradients(ctx_id) + local_t1 = rref_t1.local_value().wait() + self.assertIn(local_t1, grads) + self.assertEqual(grads[local_t1], t1.grad) + + # unblock trainers + _set_rpc_done(None, 0) + + # wait until all trainers are done + for fut in futures: + fut.wait() + @dist_init def test_backward_multiple_round_trips(self): local_grads = None diff --git a/torch/csrc/distributed/rpc/python_functions.cpp b/torch/csrc/distributed/rpc/python_functions.cpp index 418a3be2f8faa..50d81fe5c4473 100644 --- a/torch/csrc/distributed/rpc/python_functions.cpp +++ b/torch/csrc/distributed/rpc/python_functions.cpp @@ -37,7 +37,6 @@ std::shared_ptr matchBuiltinOp( // ``createStackForSchema`` to avoid throwing an error. stack = torch::jit::createStackForSchema( op->schema(), args, kwargs, c10::nullopt); - } catch (std::runtime_error& e) { VLOG(1) << "Couldn't match schema: " << op->schema() << " to args: " << args << " and kwargs: " << kwargs @@ -178,13 +177,20 @@ PyRRef pyRemotePythonUdf( ctx.getWorkerId() != dst.id_, "Does not support creating RRef on self yet."); auto userRRef = ctx.createUserRRef(dst.id_); - auto fm = agent.send( + + auto pythonRemoteCall = c10::guts::make_unique( + SerializedPyObj(std::move(pickledPythonUDF), std::move(tensors)), + userRRef->rrefId().toIValue(), + userRRef->forkId().toIValue()); + + // set forceGradRecording to true as even if the args does not contain any + // tensor, the return value might still contain tensors. + auto fm = sendMessageWithAutograd( + agent, dst, - PythonRemoteCall( - SerializedPyObj(std::move(pickledPythonUDF), std::move(tensors)), - userRRef->rrefId().toIValue(), - userRRef->forkId().toIValue()) - .toMessage()); + std::move(*pythonRemoteCall).toMessage(), + true /*forceGradRecording*/ + ); ctx.addPendingUser(userRRef->forkId(), userRRef); fm->addCallback(finishAcceptUserRRef); diff --git a/torch/csrc/distributed/rpc/rref.cpp b/torch/csrc/distributed/rpc/rref.cpp index 60b84b71195ca..78e884e572397 100644 --- a/torch/csrc/distributed/rpc/rref.cpp +++ b/torch/csrc/distributed/rpc/rref.cpp @@ -22,6 +22,26 @@ constexpr int PARENT_IDX = 5; // index of parent in the tuple // NB: if more fields are added, make sure this field is also bumped constexpr int RFD_TUPLE_SIZE = 6; // number of RRefForkData fields in py::tuple +template +T& unwrapAutogradMessage( + const Message& message, + std::unique_ptr& response) { + if (message.type() == MessageType::FORWARD_AUTOGRAD_RESP) { + auto& rpcWithAutograd = static_cast(*response); + + // Attach 'recv' autograd function. + addRecvRpcBackward( + rpcWithAutograd.autogradMetadata(), + rpcWithAutograd.tensors(), + rpcWithAutograd.fromWorkerId()); + + auto& wrappedRpc = rpcWithAutograd.wrappedRpc(); + return static_cast(wrappedRpc); + } else { + return static_cast(*response); + } +} + } // namespace std::atomic RRefContext::nextLocalId_{0}; @@ -149,23 +169,8 @@ std::shared_ptr UserRRef::toHere() { futureResponse->addCallback([future](const Message& message) { RRefContext::handleException(message); auto response = deserializeResponse(message); - if (message.type() == MessageType::FORWARD_AUTOGRAD_RESP) { - auto& rpcWithAutograd = - static_cast(*response); - - // Attach 'recv' autograd function. - addRecvRpcBackward( - rpcWithAutograd.autogradMetadata(), - rpcWithAutograd.tensors(), - rpcWithAutograd.fromWorkerId()); - - auto& wrappedRpc = rpcWithAutograd.wrappedRpc(); - auto& rfr = static_cast(wrappedRpc); - future->markCompleted(rfr.values().front()); - } else { - auto& rfr = static_cast(*response); - future->markCompleted(rfr.values().front()); - } + auto& rfr = unwrapAutogradMessage(message, response); + future->markCompleted(rfr.values().front()); }); return future; } @@ -174,14 +179,21 @@ template <> std::shared_ptr UserRRef::toHere() { auto future = std::make_shared(nullptr); auto agent = RpcAgent::getDefaultRpcAgent(); - auto futureResponse = agent->send( + + // PythonRRefFetchCall message always carries autograd context id even if + // the message itself does not contain any tensor, because the response would + // potentially contain tensors. + auto futureResponse = autograd::sendMessageWithAutograd( + *agent, agent->getWorkerInfo(ownerId_), - PythonRRefFetchCall(rrefId()).toMessage()); + PythonRRefFetchCall(ownerId_, rrefId()).toMessage(), + true /* forceGradRecording */); futureResponse->addCallback([future](const Message& message) { RRefContext::handleException(message); - auto rfr = PythonRRefFetchRet::fromMessage(message); - future->markCompleted(c10::ivalue::Tuple::create(rfr->values())); + auto response = deserializeResponse(message); + auto& rfr = unwrapAutogradMessage(message, response); + future->markCompleted(c10::ivalue::Tuple::create(rfr.values())); }); return future; } diff --git a/torch/csrc/distributed/rpc/rref_proto.cpp b/torch/csrc/distributed/rpc/rref_proto.cpp index 9d3cb54b676d6..4c503338d6e56 100644 --- a/torch/csrc/distributed/rpc/rref_proto.cpp +++ b/torch/csrc/distributed/rpc/rref_proto.cpp @@ -99,11 +99,26 @@ std::unique_ptr ScriptRRefFetchCall::fromMessage( worker_id_t(id), RRefId::fromIValue(values[0])); } +Message PythonRRefFetchCall::toMessage() && { + std::vector ivalues; + ivalues.reserve(2); + ivalues.emplace_back(rrefId_.toIValue()); + ivalues.emplace_back(fromWorkerId_); + return fromIValues(std::move(ivalues), MessageType::PYTHON_RREF_FETCH_CALL); +} + std::unique_ptr PythonRRefFetchCall::fromMessage( const Message& message) { + auto values = toIValues(message, MessageType::PYTHON_RREF_FETCH_CALL); + TORCH_INTERNAL_ASSERT( + values.size() == 2, "PythonRRefFetchCall expects 2 IValues from message"); + auto id = values[1].toInt(); + TORCH_INTERNAL_ASSERT( + id >= std::numeric_limits::min() && + id <= std::numeric_limits::max(), + "PythonRRefFetchCall fromWorkerId exceeds worker_id_t limit.") return c10::guts::make_unique( - RRefId::fromIValue(RRefMessageBase::fromMessage( - message, MessageType::PYTHON_RREF_FETCH_CALL))); + worker_id_t(id), RRefId::fromIValue(values[0])); } const std::vector& RRefFetchRet::values() { diff --git a/torch/csrc/distributed/rpc/rref_proto.h b/torch/csrc/distributed/rpc/rref_proto.h index dc52b9b4509cc..d49daa2300f51 100644 --- a/torch/csrc/distributed/rpc/rref_proto.h +++ b/torch/csrc/distributed/rpc/rref_proto.h @@ -69,11 +69,16 @@ class TORCH_API ScriptRRefFetchCall final : public RRefMessageBase { class TORCH_API PythonRRefFetchCall final : public RRefMessageBase { public: - explicit PythonRRefFetchCall(const RRefId& rrefId) - : RRefMessageBase(rrefId, MessageType::PYTHON_RREF_FETCH_CALL) {} + PythonRRefFetchCall(worker_id_t fromWorkerId, const RRefId& rrefId) + : RRefMessageBase(rrefId, MessageType::PYTHON_RREF_FETCH_CALL), + fromWorkerId_(fromWorkerId) {} + Message toMessage() && override; static std::unique_ptr fromMessage( const Message& message); + + private: + const worker_id_t fromWorkerId_; }; // OwnerRRef uses this message to send the RRef value to a remote UserRRef diff --git a/torch/csrc/distributed/rpc/utils.cpp b/torch/csrc/distributed/rpc/utils.cpp index b4dd2cd606d03..e40a2ebe51e74 100644 --- a/torch/csrc/distributed/rpc/utils.cpp +++ b/torch/csrc/distributed/rpc/utils.cpp @@ -74,6 +74,9 @@ std::unique_ptr deserializeResponse(const Message& response) { case MessageType::SCRIPT_RREF_FETCH_RET: { return ScriptRRefFetchRet::fromMessage(response); } + case MessageType::PYTHON_RREF_FETCH_RET: { + return PythonRRefFetchRet::fromMessage(response); + } case MessageType::RREF_ACK: { return RRefAck::fromMessage(response); } From a8b63cacbca5ef83bc105488c6b5cbc7cfa10537 Mon Sep 17 00:00:00 2001 From: svcscm Date: Tue, 29 Oct 2019 19:52:58 -0700 Subject: [PATCH 56/64] Updating submodules Summary: GitHub commits: https://github.com/facebook/fbthrift/commit/4b2da87ee665419f4bea4f1948af54a52e0e6adf https://github.com/facebook/fbzmq/commit/b997eec151aa8fd8f16ace83a84f743cf7e60e49 https://github.com/facebook/mcrouter/commit/9f34d1f643ba017ac8a0dd62c8349123d96148c9 https://github.com/facebook/rocksdb/commit/a3960fc875233308976351f185b672e8f01296ec https://github.com/facebook/wangle/commit/541c404784c95dd63a6b1acce2bfb2ac2d2696a3 https://github.com/facebookincubator/katran/commit/b2438faaf0e92e9ec5de011f7e9b988bda73a243 https://github.com/facebookincubator/mvfst/commit/06335bac7c7c6585b26b91df8399f77fa8ccf15f https://github.com/pytorch/fbgemm/commit/2ac6f45e20b207840f74645b757ada13a94eb8c3 Test Plan: n/a Reviewed By: zpao fbshipit-source-id: a6fb756d2d210d0505c889ba6c0e207e6a2d074d --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index 214b370edb344..2ac6f45e20b20 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 214b370edb34437d7cd4c861bcd09775c5e330cc +Subproject commit 2ac6f45e20b207840f74645b757ada13a94eb8c3 From 790563b37482b9b98077a67623ddd1b5db0105f6 Mon Sep 17 00:00:00 2001 From: Yinghai Lu Date: Tue, 29 Oct 2019 21:57:47 -0700 Subject: [PATCH 57/64] Add OfflineTensor (#28855) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28855 Resubmit: OfflineTensor will be a shell to just carry the shape and dtype. No data will be stored. This should help us plumb through the onnxifi process. Test Plan: ``` buck test caffe2/caffe2/fb/opt:onnxifi_with_offline_tensor_test ``` Reviewed By: ipiszy, ChunliF Differential Revision: D18212824 fbshipit-source-id: 5c8aaed2ef11d719dfa2a2901875efd66806ea56 --- caffe2/onnx/offline_tensor.cc | 88 +++++++++++++++++++++++++++++++++++ caffe2/onnx/offline_tensor.h | 53 +++++++++++++++++++++ 2 files changed, 141 insertions(+) create mode 100644 caffe2/onnx/offline_tensor.cc create mode 100644 caffe2/onnx/offline_tensor.h diff --git a/caffe2/onnx/offline_tensor.cc b/caffe2/onnx/offline_tensor.cc new file mode 100644 index 0000000000000..a80f2bbd49b62 --- /dev/null +++ b/caffe2/onnx/offline_tensor.cc @@ -0,0 +1,88 @@ +#include "caffe2/onnx/offline_tensor.h" + +namespace caffe2 { + +#ifndef C10_MOBILE +namespace { +// These constants need to be aligned with onnxifi.h +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT16 = 10; +constexpr uint64_t kONNXIFI_DATATYPE_FLOAT32 = 1; +constexpr uint64_t kONNXIFI_DATATYPE_UINT8 = 2; +constexpr uint64_t kONNXIFI_DATATYPE_INT32 = 6; +constexpr uint64_t kONNXIFI_DATATYPE_INT8 = 3; +constexpr uint64_t kONNXIFI_DATATYPE_INT64 = 7; +constexpr uint64_t kONNXIFI_DATATYPE_INT16 = 5; +constexpr uint64_t kONNXIFI_DATATYPE_UINT16 = 4; +} // namespace + +CAFFE_KNOWN_TYPE(OfflineTensor); + +bool OfflineTensorShapeFunctions::IsSameMetaType(TypeIdentifier id) { + return id == TypeMeta::Id(); +} + +TypeIdentifier OfflineTensorShapeFunctions::GetTypeMetaId() { + return TypeMeta::Id(); +} + +TypeMeta OfflineTensorShapeFunctions::GetExternalTensorType(const void* c) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + + return offline_tensor->shape_tensor.dtype(); +} + +vector OfflineTensorShapeFunctions::GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) { + const OfflineTensor* offline_tensor = + reinterpret_cast(c); + return GetTensorInfo(&(offline_tensor->shape_tensor), capacity, device); +} + +void OfflineTensorShapeFunctions::SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* /* unused */, + std::vector>* /* unused */, + ExternalTensorDescriptor* desc) { + const auto& offline_tensor = blob->template Get(); + const Tensor& shape_tensor = offline_tensor.shape_tensor; + + if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT32; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT8; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT64; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_INT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_FLOAT16; + } else if (shape_tensor.template IsType()) { + desc->dataType = kONNXIFI_DATATYPE_UINT16; + } else { + CAFFE_THROW("Unsupported tensor type: ", shape_tensor.dtype().name()); + } + desc->buffer = 0; + + desc->quantizationParams = 0; + desc->quantizationAxis = 0; + + // Set up dim and shape + const auto shape = shape_tensor.sizes(); + desc->dimensions = shape.size(); + shapes->emplace_back(shape.cbegin(), shape.cend()); + desc->shape = shapes->back().data(); +} + +REGISTER_EXTERNAL_TENSOR_FUNCTIONS( + (TypeMeta::Id()), + OfflineTensorShapeFunctions); +#endif +} // namespace caffe2 diff --git a/caffe2/onnx/offline_tensor.h b/caffe2/onnx/offline_tensor.h new file mode 100644 index 0000000000000..73287630ee6a3 --- /dev/null +++ b/caffe2/onnx/offline_tensor.h @@ -0,0 +1,53 @@ +#pragma once + +#include +#include "caffe2/core/operator.h" +#include "caffe2/core/tensor.h" + +namespace caffe2 { + +#ifndef C10_MOBILE +struct OfflineTensor { + // A shell tensor to record shape and dtype + Tensor shape_tensor{CPU}; + + void setShapeAndType( + const std::vector& sizes, + at::Device device, + caffe2::TypeMeta data_type) { + shape_tensor.unsafeGetTensorImpl()->set_storage( + at::Storage::create_legacy(device, data_type)); + shape_tensor.Resize(sizes); + CHECK(!shape_tensor.storage_initialized()); + CHECK(shape_tensor.dtype_initialized()); + } +}; + +class OfflineTensorShapeFunctions : public ExternalTensorFunctionsBase { + public: + explicit OfflineTensorShapeFunctions() : ExternalTensorFunctionsBase() {} + ~OfflineTensorShapeFunctions() override {} + bool isQuantized() const override { + return false; + } + bool IsSameMetaType(TypeIdentifier id) override; + void SetupExternalTensorDescriptor( + const Blob* blob, + std::vector>* shapes, + std::vector>* all_scales, + std::vector>* all_offsets, + ExternalTensorDescriptor* desc) override; + void LoadInfoOfBlob( + const Blob* /* unused */, + std::vector* /* unused */, + std::vector* /* unused */, + uint32_t* /* unused */) override {} + TypeIdentifier GetTypeMetaId() override; + TypeMeta GetExternalTensorType(const void* c) override; + vector GetExternalTensorInfo( + const void* c, + size_t* capacity, + DeviceOption* device) override; +}; +#endif +} // namespace caffe2 From eb00af37bd57685e4f1b7102ea423cae3a7004b1 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 29 Oct 2019 21:57:49 -0700 Subject: [PATCH 58/64] insert_prepack_unpack for conv (#27346) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27346 att Test Plan: test_jit.py Imported from OSS Differential Revision: D18182915 fbshipit-source-id: d646ae76ce44f5d12e974c776a3e92e5e163493c --- torch/csrc/jit/passes/quantization.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp index 51190dd1005e4..bd6f7543b937b 100644 --- a/torch/csrc/jit/passes/quantization.cpp +++ b/torch/csrc/jit/passes/quantization.cpp @@ -675,7 +675,7 @@ graph(%a_dequant, %w, %b, %w_scale, %w_zero_point, %w_dtype, %stride, %padding, %packed_params = quantized::conv_prepack(%w_quant, %b, %stride, %padding, %dilation, %groups) %w_quant_unpacked : Tensor, %b_unpacked : Tensor? = quantized::conv_unpack(%packed_params) %w_dequant = aten::dequantize(%w_quant_unpacked) - %r = aten::conv2d(%a_dequant, %w_dequant, %b, %stride, %padding, %dilation, %groups) + %r = aten::conv2d(%a_dequant, %w_dequant, %b_unpacked, %stride, %padding, %dilation, %groups) return (%r) )"; SubgraphRewriter rewriter; From 496f740824fbae657927788c35442eef0a1d8486 Mon Sep 17 00:00:00 2001 From: Benny Chen Date: Tue, 29 Oct 2019 23:30:30 -0700 Subject: [PATCH 59/64] Connect with clip range gather operator (#28866) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28866 When we are working on the fix for int32 instead of int64, we also need to take care of the ClipRangesGatherSigridHash since this is the operator that actually gets used during inference. Test Plan: Added unittest to cover for the new case Reviewed By: ipiszy Differential Revision: D17147237 fbshipit-source-id: 2b562b72a6ae8f7282e54d822467b8204fb1055e --- caffe2/opt/custom/converter.cc | 6 ++++++ caffe2/opt/custom/converter_test.cc | 33 +++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100644 caffe2/opt/custom/converter_test.cc diff --git a/caffe2/opt/custom/converter.cc b/caffe2/opt/custom/converter.cc index 37afc2dc731ba..f52bf18880829 100644 --- a/caffe2/opt/custom/converter.cc +++ b/caffe2/opt/custom/converter.cc @@ -214,6 +214,10 @@ class ClipRangesGatherSigridHashConverter : public Converter { if (args.HasArgument("max_values")) { c->setMaxValues(args.GetRepeatedArgument("max_values")); } + if (args.HasArgument("hash_into_int32")) { + c->setHashIntoInt32( + args.GetSingleArgument("hash_into_int32", false)); + } return nnOp; } @@ -230,6 +234,8 @@ class ClipRangesGatherSigridHashConverter : public Converter { caffe2::MakeArgument>("salts", fuse->getSalts())); op.add_arg()->CopyFrom(caffe2::MakeArgument>( "max_values", fuse->getMaxValues())); + op.add_arg()->CopyFrom(caffe2::MakeArgument( + "hash_into_int32", fuse->getHashIntoInt32())); return op; } diff --git a/caffe2/opt/custom/converter_test.cc b/caffe2/opt/custom/converter_test.cc new file mode 100644 index 0000000000000..4313c3a14087f --- /dev/null +++ b/caffe2/opt/custom/converter_test.cc @@ -0,0 +1,33 @@ +#include "caffe2/core/common.h" +#include "caffe2/core/test_utils.h" +#include "caffe2/opt/converter.h" +#include "caffe2/opt/custom/concat_elim.h" +#include "caffe2/predictor/emulator/data_filler.h" +#include "caffe2/utils/proto_utils.h" + +#include + +using namespace caffe2::testing; +using namespace caffe2::emulator; +using caffe2::OperatorDef; +using std::vector; + +TEST(Converter, ClipRangesGatherSigridHashConverter) { + OperatorDef op; + op.set_type("ClipRangesGatherSigridHash"); + op.add_arg()->CopyFrom(caffe2::MakeArgument("hash_into_int32", true)); + auto nnDef = convertToNeuralNetOperator(op); + auto* pNNDef = + static_cast(nnDef.get()); + EXPECT_TRUE(pNNDef); + EXPECT_TRUE(pNNDef->getHashIntoInt32()); + + OperatorDef op2; + op2.set_type("ClipRangesGatherSigridHash"); + op2.add_arg()->CopyFrom(caffe2::MakeArgument("hash_into_int32", false)); + auto nnDef2 = convertToNeuralNetOperator(op2); + auto* pNNDef2 = + static_cast(nnDef2.get()); + EXPECT_TRUE(pNNDef2); + EXPECT_FALSE(pNNDef2->getHashIntoInt32()); +} From 726f0ce946ed385d82b1d93b940bcc6fb8c61c0e Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Wed, 30 Oct 2019 08:25:23 -0700 Subject: [PATCH 60/64] Increase verbosity of Hypothesis on CI. (#28799) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/28799 When the verbosity is quiet, hypothesis no longer prints the real error when it finds multiple falsifying examples: it just says that there are two failures. This is supremely unuseful. Make it print more. Signed-off-by: Edward Z. Yang Test Plan: Imported from OSS Differential Revision: D18206936 Pulled By: ezyang fbshipit-source-id: 03bb60ba24cee28706bb3d1f0858c32b6743a109 --- test/common_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/common_utils.py b/test/common_utils.py index 8c04fbcefd5dd..de04bdb542b5b 100644 --- a/test/common_utils.py +++ b/test/common_utils.py @@ -462,7 +462,7 @@ def __exit__(self, exec_type, exec_value, traceback): suppress_health_check=[hypothesis.HealthCheck.too_slow], database=None, max_examples=100, - verbosity=hypothesis.Verbosity.quiet)) + verbosity=hypothesis.Verbosity.normal)) hypothesis.settings.register_profile( "dev", hypothesis.settings( @@ -486,7 +486,7 @@ def __exit__(self, exec_type, exec_value, traceback): database=None, max_examples=100, min_satisfying_examples=1, - verbosity=hypothesis.Verbosity.quiet)) + verbosity=hypothesis.Verbosity.normal)) hypothesis.settings.register_profile( "dev", hypothesis.settings( From 2526f97464a24244d07b312931bdb1c54d8295c2 Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Wed, 30 Oct 2019 08:31:12 -0700 Subject: [PATCH 61/64] Include hierarchy information in C++ API loading error messages (#28499) Summary: Before, we would only give the key we are looking for (i.e. typically just "No such serialized tensor 'weight'", no matter for which submodule we were looking for a weight. Now we error with "No such serialized tensor '0.conv1.weight'" or similar. The analogous information is added to missing module error messages. I threw in a test, and it saved me already... Pull Request resolved: https://github.com/pytorch/pytorch/pull/28499 Differential Revision: D18122442 Pulled By: yf225 fbshipit-source-id: a134b6d06ca33de984a11d6fea923244bcd9fb95 --- test/cpp/api/serialize.cpp | 33 +++++++++++++++++++ .../include/torch/serialize/input-archive.h | 1 + .../csrc/api/src/serialize/input-archive.cpp | 17 ++++++---- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/test/cpp/api/serialize.cpp b/test/cpp/api/serialize.cpp index 053d534a1719f..c6f9322f2f999 100644 --- a/test/cpp/api/serialize.cpp +++ b/test/cpp/api/serialize.cpp @@ -127,6 +127,39 @@ TEST(SerializeTest, NonContiguous) { ASSERT_TRUE(x.allclose(y)); } +TEST(SerializeTest, ErrorOnMissingKey) { + struct B : torch::nn::Module { + B(const std::string& name_c) { + register_buffer(name_c, torch::ones(5, torch::kFloat)); + } + }; + struct A : torch::nn::Module { + A(const std::string& name_b, const std::string& name_c) { + register_module(name_b, std::make_shared(name_c)); + } + }; + struct M : torch::nn::Module { + M(const std::string& name_a, + const std::string& name_b, + const std::string& name_c) { + register_module(name_a, std::make_shared(name_b, name_c)); + } + }; + + // create a hierarchy of models with names differing below the top level + auto model1 = std::make_shared("a", "b", "c"); + auto model2 = std::make_shared("a", "b", "x"); + auto model3 = std::make_shared("a", "x", "c"); + + std::stringstream stream; + torch::save(model1, stream); + // We want the errors to contain hierarchy information, too. + ASSERT_THROWS_WITH( + torch::load(model2, stream), "No such serialized tensor 'a.b.x'"); + ASSERT_THROWS_WITH( + torch::load(model3, stream), "No such serialized submodule: 'a.x'"); +} + TEST(SerializeTest, XOR) { // We better be able to save and load an XOR model! auto getLoss = [](Sequential model, uint32_t batch_size) { diff --git a/torch/csrc/api/include/torch/serialize/input-archive.h b/torch/csrc/api/include/torch/serialize/input-archive.h index 9965a442f6a7d..d4a4318e7537a 100644 --- a/torch/csrc/api/include/torch/serialize/input-archive.h +++ b/torch/csrc/api/include/torch/serialize/input-archive.h @@ -102,6 +102,7 @@ class TORCH_API InputArchive final { private: jit::script::Module module_; + std::string hierarchy_prefix_; }; } // namespace serialize } // namespace torch diff --git a/torch/csrc/api/src/serialize/input-archive.cpp b/torch/csrc/api/src/serialize/input-archive.cpp index 856ddcb7a4c28..5ebbe1fe7a249 100644 --- a/torch/csrc/api/src/serialize/input-archive.cpp +++ b/torch/csrc/api/src/serialize/input-archive.cpp @@ -63,15 +63,17 @@ void InputArchive::read( Tensor& tensor, bool is_buffer) { TORCH_CHECK( - try_read(key, tensor, is_buffer), - "No such serialized tensor '", - key, - "'"); + try_read(key, tensor, is_buffer), + "No such serialized tensor '", + hierarchy_prefix_, + key, + "'"); } bool InputArchive::try_read(const std::string& key, InputArchive& archive) { if (auto named_module = module_.find_module(key)) { archive.module_ = std::move(*named_module); + archive.hierarchy_prefix_ = hierarchy_prefix_ + key + "."; return true; } else { return false; @@ -80,8 +82,11 @@ bool InputArchive::try_read(const std::string& key, InputArchive& archive) { void InputArchive::read(const std::string& key, InputArchive& archive) { TORCH_CHECK( - try_read(key, archive), - "No such serialized submodule: '", key, "'"); + try_read(key, archive), + "No such serialized submodule: '", + hierarchy_prefix_, + key, + "'"); } void InputArchive::load_from(const std::string& filename, From 0a5c573bddb3a6ca8116aa3906759d386b24d5d5 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Mon, 28 Oct 2019 16:17:27 -0700 Subject: [PATCH 62/64] specialize undefinedness --- test/jit_utils.py | 5 +- test/test_jit_fuser.py | 2 +- torch/csrc/jit/graph_executor.cpp | 95 +++++++++++++++++++++++++++---- 3 files changed, 88 insertions(+), 14 deletions(-) diff --git a/test/jit_utils.py b/test/jit_utils.py index 7f5ed32bcb7b2..54148c090c742 100644 --- a/test/jit_utils.py +++ b/test/jit_utils.py @@ -33,7 +33,7 @@ import tempfile import textwrap -IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR = False +IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR = True class ProfilingMode(Enum): OFF = 1 @@ -44,7 +44,8 @@ class ProfilingMode(Enum): def enable_profiling_mode(flag): if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: old_prof_exec_state = torch._C._jit_set_profiling_executor(flag != ProfilingMode.OFF) - old_prof_mode_state = torch._C._jit_set_profiling_mode(flag == ProfilingMode.FULL) + #old_prof_mode_state = torch._C._jit_set_profiling_mode(flag == ProfilingMode.FULL) + old_prof_mode_state = torch._C._jit_set_profiling_mode(False) try: yield finally: diff --git a/test/test_jit_fuser.py b/test/test_jit_fuser.py index 6c9db9b5b736b..9429fb27b14ea 100644 --- a/test/test_jit_fuser.py +++ b/test/test_jit_fuser.py @@ -20,7 +20,7 @@ if IN_TRANSITION_TO_PROFILING_GRAPH_EXECUTOR: torch._C._jit_set_profiling_executor(True) - torch._C._jit_set_profiling_mode(True) + torch._C._jit_set_profiling_mode(False) def strip_profiling_nodes(nodes): diff --git a/torch/csrc/jit/graph_executor.cpp b/torch/csrc/jit/graph_executor.cpp index 4abb884d31c27..1635e5d9c2b1e 100644 --- a/torch/csrc/jit/graph_executor.cpp +++ b/torch/csrc/jit/graph_executor.cpp @@ -206,21 +206,27 @@ static void unpackReturnTuple(Stack &stack) { stack.insert(stack.end(), tuple->elements().begin(), tuple->elements().end()); } +struct DifferentiableGraphOp; + struct DifferentiableGraphBackward : public autograd::Node { DifferentiableGraphBackward( - GraphExecutor executor, + const std::shared_ptr& unspec_graph, size_t input_size, - size_t capture_size) - : executor(std::move(executor)), - captures_(capture_size), - input_instructions_(input_size) {} + size_t capture_size, + c10::optional& grad_executor) + : captures_(capture_size), + input_instructions_(input_size), + unspecialized_graph_(unspec_graph), + grad_executor_(grad_executor) {} variable_list apply(variable_list&& inputs) override { Stack stack; - stack.reserve(captures_.size() + inputs.size()); + size_t num_args = captures_.size() + inputs.size(); + stack.reserve(num_args); input_instructions_.unpack(std::move(inputs), stack); captures_.unpack(stack, shared_from_this()); + GraphExecutor& executor = getExecutor(stack); GRAPH_DEBUG("Running DifferentiableGraphBackward for ", &executor); executor.run(stack); unpackReturnTuple(stack); @@ -259,6 +265,66 @@ struct DifferentiableGraphBackward : public autograd::Node { captures_.capture(val, is_output); } + + static c10::TensorTypePtr getTensorType(bool defined) { + auto tensor_type = TensorType::get(); + + if (defined) { + return tensor_type; + } + + return tensor_type->withUndefined(); + } + + GraphExecutor& getExecutor(Stack& stack) { + + // tensor lists are hashed as a single boolean value + // since all tensors will be either defined or undefined + + std::vector hash; + + for (IValue& v : stack) { + if (v.isTensorList()) { + auto list = v.toTensorListRef(); + hash.push_back(list.size() > 0 ? list[0].defined() : true); + } else if (v.isTensor()) { + hash.push_back(v.toTensor().defined()); + } else { + // assume that every other type is defined + hash.push_back(true); + } + } + + + TORCH_INTERNAL_ASSERT(unspecialized_graph_->inputs().size() == hash.size()); + if (grad_executors_.count(hash) == 0) { + + std::shared_ptr spec_copy = unspecialized_graph_->copy(); + + for (auto i = 0; i < hash.size(); i++) { + auto input_type = spec_copy->inputs().at(i); + bool defined = hash[i]; + if (input_type->type()->kind() == TensorType::Kind) { + input_type->setType(getTensorType(defined)); + } else if ( + input_type->type()->kind() == ListType::Kind && + input_type->type()->expect()->getElementType()->kind() == + TensorType::Kind) { + input_type->setType(ListType::create(getTensorType(defined))); + } + } + grad_executors_[hash] = GraphExecutor(spec_copy); + } + + + // set last optimized graph + // make a copy because DifferentiableBackward might disappear + // by the time we get to use diff_op_.grad_executor + grad_executor_ = GraphExecutor(grad_executors_[hash].graph()) ; + return grad_executors_[hash]; + + } + void addOutputForTensor(const at::Tensor& tensor) { auto v = Variable(tensor); add_next_edge(v.defined() ? v.gradient_edge() : autograd::Edge{}); @@ -319,6 +385,9 @@ struct DifferentiableGraphBackward : public autograd::Node { GraphExecutor executor; CaptureList captures_; UnpackInstructions input_instructions_; + std::unordered_map, GraphExecutor> grad_executors_; + std::shared_ptr unspecialized_graph_; + c10::optional& grad_executor_; }; // an optimized way of executing the subgraph computed directly on @@ -330,17 +399,18 @@ struct DifferentiableGraphOp { DifferentiableGraphOp(Gradient grad) : f(grad.f), grad(std::move(grad)), - grad_executor(this->grad.df), + grad_executor(), num_inputs(this->grad.f->inputs().size()), num_outputs(this->grad.f->outputs().size()) {} // XXX: keep in mind that stack can be larger than the inputs we need! int operator()(Stack& stack) const { auto grad_fn = std::make_shared( - grad_executor, + this->grad.df, grad.df_input_vjps.size(), grad.df_input_captured_inputs.size() + - grad.df_input_captured_outputs.size()); + grad.df_input_captured_outputs.size(), + grad_executor); { auto inputs = last(stack, num_inputs); @@ -378,6 +448,7 @@ struct DifferentiableGraphOp { private: friend GraphExecutor* detail::getGradExecutor(Operation& op); + friend struct DifferentiableGraphBackward; at::Tensor detach(at::Tensor t) const { if (!t.defined()) { @@ -426,7 +497,7 @@ struct DifferentiableGraphOp { Code f; Gradient grad; - GraphExecutor grad_executor; + mutable c10::optional grad_executor; const size_t num_inputs; const size_t num_outputs; @@ -459,7 +530,9 @@ namespace detail { GraphExecutor* getGradExecutor(Operation& op) { if (auto diff_op = op.target()) { - return &diff_op->grad_executor; + + TORCH_INTERNAL_ASSERT(diff_op->grad_executor.has_value()) + return &(*diff_op->grad_executor); } return nullptr; } From 91c3cdc650d63052b1ba270104524057348b5af6 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Wed, 30 Oct 2019 10:01:11 -0700 Subject: [PATCH 63/64] investigating test failures --- test/test_jit_fuser.py | 2 +- torch/csrc/jit/graph_executor.cpp | 2 +- torch/csrc/jit/ir.cpp | 9 +++---- torch/csrc/jit/passes/graph_fuser.cpp | 26 ++++++++++++++++--- .../jit/profiling_graph_executor_impl.cpp | 2 ++ 5 files changed, 29 insertions(+), 12 deletions(-) diff --git a/test/test_jit_fuser.py b/test/test_jit_fuser.py index 9429fb27b14ea..c35d0a8e6ca7c 100644 --- a/test/test_jit_fuser.py +++ b/test/test_jit_fuser.py @@ -58,7 +58,7 @@ def assertAllFused(self, graph, except_for=()): graph = diff_graphs[0].g('Subgraph') allowed_nodes = {'prim::Constant', 'prim::FusionGroup', 'prim::BailoutTemplate', - 'prim::BailOut', 'prim::TupleConstruct'} | set(except_for) + 'prim::BailOut', 'prim::TupleConstruct', 'aten::size', 'aten::_size_if_not_equal', "prim::BroadcastSizes"} | set(except_for) self.assertTrue(all(node.kind() in allowed_nodes for node in graph.nodes()), 'got {}'.format(graph)) self.assertTrue([node.kind() for node in graph.nodes()].count('prim::FusionGroup') == 1) diff --git a/torch/csrc/jit/graph_executor.cpp b/torch/csrc/jit/graph_executor.cpp index 1635e5d9c2b1e..e3f4a09e224cc 100644 --- a/torch/csrc/jit/graph_executor.cpp +++ b/torch/csrc/jit/graph_executor.cpp @@ -299,8 +299,8 @@ struct DifferentiableGraphBackward : public autograd::Node { TORCH_INTERNAL_ASSERT(unspecialized_graph_->inputs().size() == hash.size()); if (grad_executors_.count(hash) == 0) { + GRAPH_DEBUG("creating a specialized copy for ", this); std::shared_ptr spec_copy = unspecialized_graph_->copy(); - for (auto i = 0; i < hash.size(); i++) { auto input_type = spec_copy->inputs().at(i); bool defined = hash[i]; diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp index 9f90f660c8e2f..dee4d1e94e929 100644 --- a/torch/csrc/jit/ir.cpp +++ b/torch/csrc/jit/ir.cpp @@ -244,22 +244,19 @@ std::ostream &Node::print(std::ostream &out, size_t level, auto* pyOp = static_cast(this); out << "^" << pyOp->name(); pyOp->writeScalars(out); - } else if (print_attributes) { - if (hasAttribute(attr::Subgraph) && groups) { + } else if (hasAttribute(attr::Subgraph) && groups) { out << kind().toQualString() << "_" << groups->size(); - if (numAttributes() > 1 && kind() != prim::DifferentiableGraph) { + if (print_attributes && numAttributes() > 1 && kind() != prim::DifferentiableGraph) { printAttributes(out, /*ignore_subgraph=*/true); } groups->push_back(this); } else { out << kind().toQualString(); - if (hasAttributes()) { + if (print_attributes && hasAttributes()) { printAttributes(out); } } - } - out << "(" << inputs() << ")"; if (print_scopes) { diff --git a/torch/csrc/jit/passes/graph_fuser.cpp b/torch/csrc/jit/passes/graph_fuser.cpp index b07464e4bc8f9..3c1010bf58c0c 100644 --- a/torch/csrc/jit/passes/graph_fuser.cpp +++ b/torch/csrc/jit/passes/graph_fuser.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -199,7 +200,10 @@ struct GraphFuser { fusableDevice &= isFusableDevice(output); } } - return fusableDevice && isFusableMap(node); + + bool is_fusable_map = isFusableMap(node); + GRAPH_DEBUG("isFusableDefault for ", getHeader(node), " fusableDevice = ", fusableDevice, " is_fusable_map = ", is_fusable_map); + return fusableDevice && is_fusable_map; } bool isFusableMap(Node* node) { @@ -211,25 +215,33 @@ struct GraphFuser { } bool isFusableCatNode(Node* node) { - if (node->kind() != aten::cat) + if (node->kind() != aten::cat) { return false; - if (!node->is_constant(attr::dim)) + } + if (!node->is_constant(attr::dim)) { + GRAPH_DEBUG("attr::dim of ", getHeader(node), " isn't a prim::Constant"); return false; + } auto tensors_node = node->namedInput(attr::tensors)->node(); if ((tensors_node->inputs().size() + node->outputs().size()) > subgraph_arg_limit_) { + + GRAPH_DEBUG("too many arguments for ", getHeader(node)); return false; } if (tensors_node->kind() != prim::ListConstruct) + GRAPH_DEBUG("inputs tensors don't come from prim::ListConstruct for ", getHeader(node)); return false; // NB: Note that technically other uses of the list aren't a big problem for // us. It would be enough to place the prim::FusedConcat before the // prim::ListConstruct, and allUsersAreThisConsumerOrOccurAfterIt would // still be satisfied. However, I don't expect this to be necessary any time // soon, and so we're simply assuming that we don't have to deal with it. - if (tensors_node->output()->uses().size() > 1) + if (tensors_node->output()->uses().size() > 1) { + GRAPH_DEBUG("there's more than one user of the input tensors of ", getHeader(node)); return false; + } return true; } @@ -254,6 +266,8 @@ struct GraphFuser { } void mergeFusionGroups(Node* consumer_group, Node* producer_group) { + + GRAPH_UPDATE("Merging a producer group ", getHeader(producer_group), " into ",getHeader(consumer_group)); // Now we have two fusion groups! // Revert the fusion - place all inner nodes of producer back in the outer // graph. @@ -387,6 +401,8 @@ struct GraphFuser { subgraph.eraseInput(p); } } + + GRAPH_UPDATE("Merging ", getHeader(n), " into ",getHeader(group)); return subgraph.insertNode(in_graph); } @@ -402,6 +418,7 @@ struct GraphFuser { auto sel = group->addOutput(); sel->copyMetadata(n->output()); n->replaceAllUsesWith(group); + GRAPH_UPDATE("Replacing ", getHeader(n), " with ",getHeader(group)); n->destroy(); return group; } @@ -1207,6 +1224,7 @@ void PeepholeOptimizeShapeExpressions(Block* block) { void FuseGraph(std::shared_ptr& graph) { GraphFuser(graph->block(), graph).run(); + GRAPH_DUMP("After GraphFuser: ", graph); // After FuseGraph some common subexpressions may come back EliminateCommonSubexpression(graph); // We might have emitted a fair amount of useless shape propagating code, so diff --git a/torch/csrc/jit/profiling_graph_executor_impl.cpp b/torch/csrc/jit/profiling_graph_executor_impl.cpp index b6fa448bcc45d..be7291cbb4865 100644 --- a/torch/csrc/jit/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/profiling_graph_executor_impl.cpp @@ -114,9 +114,11 @@ ExecutionPlan ProfilingGraphExecutorImpl::getPlanFor(Stack& stack) { ? needsGradientInProfilingMode(copy->block()) : needsGradient(copy); if (needs_gradient) { + GRAPH_DEBUG(this, " needs gradients"); auto diff_nodes = CreateAutodiffSubgraphs( copy, getAutodiffSubgraphInlining() ? autodiffSubgraphNodeThreshold : 1); + GRAPH_DEBUG(" diff_nodes' size is ", diff_nodes.size()); for (Node *dnode : diff_nodes) { auto diff_graph = std::move(dnode->g(attr::Subgraph)); Gradient gradient = differentiate(diff_graph); From 6195436d8c829bda7a4f251750e1290a1b0ec9d0 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Thu, 31 Oct 2019 10:11:29 -0700 Subject: [PATCH 64/64] static graph_executor --- test/test_jit_fuser.py | 1 + torch/csrc/jit/graph_executor.cpp | 39 +++++++++++++++++++++---------- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/test/test_jit_fuser.py b/test/test_jit_fuser.py index c35d0a8e6ca7c..4b1aac24c2462 100644 --- a/test/test_jit_fuser.py +++ b/test/test_jit_fuser.py @@ -445,6 +445,7 @@ def fn(x, y, z): z = torch.randn(4, 2, dtype=torch.float, device='cuda') ge = self.checkTrace(fn, (x, y, z)) graph = ge.graph_for(x, y, z) + print(str(graph)) self.assertAllFused(graph, except_for={'aten::add'}) FileCheck().check("FusedConcat").check_next("return").run(str(graph)) diff --git a/torch/csrc/jit/graph_executor.cpp b/torch/csrc/jit/graph_executor.cpp index e3f4a09e224cc..a047a65ebe89e 100644 --- a/torch/csrc/jit/graph_executor.cpp +++ b/torch/csrc/jit/graph_executor.cpp @@ -49,6 +49,8 @@ #include #include +static std::unique_ptr ge; + namespace torch { namespace jit { @@ -212,12 +214,14 @@ struct DifferentiableGraphBackward : public autograd::Node { DifferentiableGraphBackward( const std::shared_ptr& unspec_graph, size_t input_size, - size_t capture_size, - c10::optional& grad_executor) + size_t capture_size + //c10::optional& grad_executor + ) : captures_(capture_size), input_instructions_(input_size), - unspecialized_graph_(unspec_graph), - grad_executor_(grad_executor) {} + unspecialized_graph_(unspec_graph->copy()) + // grad_executor_(grad_executor) + {} variable_list apply(variable_list&& inputs) override { Stack stack; @@ -313,14 +317,19 @@ struct DifferentiableGraphBackward : public autograd::Node { input_type->setType(ListType::create(getTensorType(defined))); } } + grad_executors_[hash] = GraphExecutor(spec_copy); + + //grad_executors_.insert({hash, GraphExecutor(spec_copy)}); } // set last optimized graph // make a copy because DifferentiableBackward might disappear // by the time we get to use diff_op_.grad_executor - grad_executor_ = GraphExecutor(grad_executors_[hash].graph()) ; + + //grad_executor_ = GraphExecutor(grad_executors_[hash].graph()) ; + ge.reset(new GraphExecutor(grad_executors_[hash].graph())); return grad_executors_[hash]; } @@ -387,7 +396,7 @@ struct DifferentiableGraphBackward : public autograd::Node { UnpackInstructions input_instructions_; std::unordered_map, GraphExecutor> grad_executors_; std::shared_ptr unspecialized_graph_; - c10::optional& grad_executor_; + //c10::optional& grad_executor_; }; // an optimized way of executing the subgraph computed directly on @@ -399,7 +408,7 @@ struct DifferentiableGraphOp { DifferentiableGraphOp(Gradient grad) : f(grad.f), grad(std::move(grad)), - grad_executor(), + grad_executor(this->grad.df), num_inputs(this->grad.f->inputs().size()), num_outputs(this->grad.f->outputs().size()) {} @@ -409,8 +418,8 @@ struct DifferentiableGraphOp { this->grad.df, grad.df_input_vjps.size(), grad.df_input_captured_inputs.size() + - grad.df_input_captured_outputs.size(), - grad_executor); + grad.df_input_captured_outputs.size()//, + /*grad_executor*/); { auto inputs = last(stack, num_inputs); @@ -497,7 +506,8 @@ struct DifferentiableGraphOp { Code f; Gradient grad; - mutable c10::optional grad_executor; + //mutable c10::optional grad_executor; + GraphExecutor grad_executor; const size_t num_inputs; const size_t num_outputs; @@ -528,11 +538,16 @@ RegisterOperators reg_graph_executor_ops({Operator( namespace detail { + + GraphExecutor* getGradExecutor(Operation& op) { if (auto diff_op = op.target()) { - TORCH_INTERNAL_ASSERT(diff_op->grad_executor.has_value()) - return &(*diff_op->grad_executor); + //TORCH_INTERNAL_ASSERT(diff_op->grad_executor.has_value()) + //return &(*diff_op->grad_executor); + //@#$ to do try this next + return &diff_op->grad_executor; + return ge.get(); } return nullptr; }