From e7cccc23b916d832978d77103f9910b1cd2a01c2 Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Thu, 10 Jun 2021 01:41:25 -0700 Subject: [PATCH 001/305] Add query and synchronize to c10::Stream (#59560) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59560 `at::cuda::CUDAStream` has the `query` and `synchronize` methods, but `c10::Stream` does not, and I couldn't find any generic way to accomplish this. Hence I added helpers to do this to the DeviceGuardImpl interface, and then defined these methods on `c10::Stream`. (I had to do it out-of-line to circumvent a circular dependency). ghstack-source-id: 130932249 Test Plan: CI Reviewed By: ezyang Differential Revision: D28931377 fbshipit-source-id: cd0c19cf021e305d0c0cf9af364afb445d010248 --- .../hip/impl/HIPGuardImplMasqueradingAsCUDA.h | 11 +++++++++ c10/core/Stream.cpp | 15 ++++++++++++ c10/core/Stream.h | 10 +++++++- c10/core/impl/DeviceGuardImplInterface.h | 24 +++++++++++++++++++ c10/core/impl/VirtualGuardImpl.h | 7 ++++++ c10/cuda/impl/CUDAGuardImpl.h | 12 ++++++++++ 6 files changed, 78 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h b/aten/src/ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h index 2113993b718a4..d7c1580c78709 100644 --- a/aten/src/ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h +++ b/aten/src/ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h @@ -190,6 +190,17 @@ struct HIPGuardImplMasqueradingAsCUDA final : public c10::impl::DeviceGuardImplI return (err == hipSuccess); } + // Stream-related functions + bool queryStream(const Stream& stream) const override { + HIPStreamMasqueradingAsCUDA hip_stream{stream}; + return hip_stream.query(); + } + + void synchronizeStream(const Stream& stream) const override { + HIPStreamMasqueradingAsCUDA hip_stream{stream}; + hip_stream.synchronize(); + } + void recordDataPtrOnStream( const c10::DataPtr& data_ptr, const Stream& stream) const override { diff --git a/c10/core/Stream.cpp b/c10/core/Stream.cpp index 1a56c9d685671..43819c1cc9f1c 100644 --- a/c10/core/Stream.cpp +++ b/c10/core/Stream.cpp @@ -1,7 +1,22 @@ #include +#include namespace c10 { +// Return whether all asynchronous work previously enqueued on this stream +// has completed running on the device. +bool Stream::query() const { + impl::VirtualGuardImpl impl{device_.type()}; + return impl.queryStream(*this); +} + +// Wait (by blocking the calling thread) until all asynchronous work enqueued +// on this stream has completed running on the device. +void Stream::synchronize() const { + impl::VirtualGuardImpl impl{device_.type()}; + impl.synchronizeStream(*this); +} + // Not very parsable, but I don't know a good compact syntax for streams. // Feel free to change this into something more compact if needed. std::ostream& operator<<(std::ostream& stream, const Stream& s) { diff --git a/c10/core/Stream.h b/c10/core/Stream.h index c149dc260e0bf..446d590b588c0 100644 --- a/c10/core/Stream.h +++ b/c10/core/Stream.h @@ -54,7 +54,7 @@ using StreamId = int32_t; * functionality (e.g., get the cudaStream_t of a CUDA stream.) There are * wrapper classes which provide this functionality, e.g., CUDAStream. */ -class Stream final { +class C10_API Stream final { private: Device device_; StreamId id_; @@ -107,6 +107,14 @@ class Stream final { event.block(*this); } + // Return whether all asynchronous work previously enqueued on this stream + // has completed running on the device. + bool query() const; + + // Wait (by blocking the calling thread) until all asynchronous work enqueued + // on this stream has completed running on the device. + void synchronize() const; + // The purpose of this function is to more conveniently permit binding // of Stream to and from Python. Without packing, I have to setup a whole // class with two fields (device and stream id); with packing I can just diff --git a/c10/core/impl/DeviceGuardImplInterface.h b/c10/core/impl/DeviceGuardImplInterface.h index 28699e2818810..a87f25b60eed1 100644 --- a/c10/core/impl/DeviceGuardImplInterface.h +++ b/c10/core/impl/DeviceGuardImplInterface.h @@ -176,6 +176,22 @@ struct C10_API DeviceGuardImplInterface { */ virtual DeviceIndex deviceCount() const noexcept = 0; + /** + * Return true if all the work previously enqueued on the stream for + * asynchronous execution has completed running on the device. + */ + virtual bool queryStream(const Stream& stream) const { + TORCH_CHECK(false, "Backend doesn't support querying streams."); + } + + /** + * Wait (by blocking the calling thread) until all the work previously + * enqueued on the stream has completed running on the device. + */ + virtual void synchronizeStream(const Stream& stream) const { + TORCH_CHECK(false, "Backend doesn't support synchronizing streams."); + } + /** * Ensure the caching allocator (if any) is aware that the given DataPtr is * being used on the given stream, and that it should thus avoid recycling the @@ -241,6 +257,14 @@ struct NoOpDeviceGuardImpl final : public DeviceGuardImplInterface { } void destroyEvent(void* event, const DeviceIndex device_index) const noexcept override {} + + // Stream-related functions + bool queryStream(const Stream& stream) const override { + return true; + } + void synchronizeStream(const Stream& stream) const override { + // Don't wait for anything. + } }; // The registry is NON-owning. Each stored pointer is std::atomic so diff --git a/c10/core/impl/VirtualGuardImpl.h b/c10/core/impl/VirtualGuardImpl.h index 112a9dc253bd8..1b8da6a1a23e3 100644 --- a/c10/core/impl/VirtualGuardImpl.h +++ b/c10/core/impl/VirtualGuardImpl.h @@ -69,6 +69,13 @@ class VirtualGuardImpl final : public DeviceGuardImplInterface { impl_->destroyEvent(event, device_index); } + bool queryStream(const Stream& stream) const override { + return impl_->queryStream(stream); + } + void synchronizeStream(const Stream& stream) const override { + impl_->synchronizeStream(stream); + } + void recordDataPtrOnStream(const c10::DataPtr& data_ptr, const Stream& stream) const override { impl_->recordDataPtrOnStream(data_ptr, stream); diff --git a/c10/cuda/impl/CUDAGuardImpl.h b/c10/cuda/impl/CUDAGuardImpl.h index 86d6799a6d1e8..9a173a1f1657b 100644 --- a/c10/cuda/impl/CUDAGuardImpl.h +++ b/c10/cuda/impl/CUDAGuardImpl.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -170,6 +171,17 @@ struct CUDAGuardImpl final : public c10::impl::DeviceGuardImplInterface { return (err == cudaSuccess); } + // Stream-related functions + bool queryStream(const Stream& stream) const override { + CUDAStream cuda_stream{stream}; + return cuda_stream.query(); + } + + void synchronizeStream(const Stream& stream) const override { + CUDAStream cuda_stream{stream}; + cuda_stream.synchronize(); + } + void recordDataPtrOnStream(const c10::DataPtr& data_ptr, const Stream& stream) const override { CUDAStream cuda_stream{stream}; From 0d7d316dc1a889937227d29978ee86315f33bc24 Mon Sep 17 00:00:00 2001 From: Oleg Khabinov Date: Thu, 10 Jun 2021 01:54:31 -0700 Subject: [PATCH 002/305] [fx ir] Support lists and dicts in FX IR GraphDrawer (#58775) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58775 Reviewed By: RoshanPAN Differential Revision: D28613939 fbshipit-source-id: 4164e2dd772b59240ea3907001fe4ebddb003060 --- torch/fx/passes/graph_drawer.py | 54 +++++++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 12 deletions(-) diff --git a/torch/fx/passes/graph_drawer.py b/torch/fx/passes/graph_drawer.py index 014da33d89e3f..ae72aaab86868 100644 --- a/torch/fx/passes/graph_drawer.py +++ b/torch/fx/passes/graph_drawer.py @@ -6,6 +6,7 @@ import pydot from typing import Dict, Any from torch.fx.node import _get_qualified_name +from torch.fx.passes.shape_prop import TensorMetadata _COLOR_MAP = { "placeholder": '"AliceBlue"', @@ -130,21 +131,50 @@ def _get_node_label(self, module: torch.fx.GraphModule, node: torch.fx.Node) -> label += "|" + self._typename(node.target) + r"\l" tensor_meta = node.meta.get('tensor_meta') - if tensor_meta is not None: - label += "|" + "dtype" + "=" + str(tensor_meta.dtype) + r"\l" - label += "|" + "shape" + "=" + str(tuple(tensor_meta.shape)) + r"\l" - label += "|" + "stride" + "=" + str(tensor_meta.stride) + r"\l" - if tensor_meta.is_quantized: - if tensor_meta.qscheme in { - torch.per_tensor_affine, - torch.per_tensor_symmetric, - }: - label += "|" + "q_scale" + "=" + str(tensor_meta.q_scale) + r"\l" - label += "|" + "q_zero_point" + "=" + str(tensor_meta.q_zero_point) + r"\l" - label += "|" + "qscheme" + "=" + str(tensor_meta.qscheme) + r"\l" + label += self._tensor_meta_to_label(tensor_meta) return label + "}" + def _tensor_meta_to_label(self, tm) -> str: + if tm is None: + return "" + elif isinstance(tm, TensorMetadata): + return self._stringify_tensor_meta(tm) + elif isinstance(tm, list): + result = "" + for item in tm: + result += self._tensor_meta_to_label(item) + return result + elif isinstance(tm, dict): + result = "" + for k, v in tm.items(): + result += self._tensor_meta_to_label(v) + return result + elif isinstance(tm, tuple): + result = "" + for item in tm: + result += self._tensor_meta_to_label(item) + return result + else: + raise RuntimeError(f"Unsupported tensor meta type {type(tm)}") + + def _stringify_tensor_meta(self, tm: TensorMetadata) -> str: + result = "" + if not hasattr(tm, "dtype"): + print("tm", tm) + result += "|" + "dtype" + "=" + str(tm.dtype) + r"\l" + result += "|" + "shape" + "=" + str(tuple(tm.shape)) + r"\l" + result += "|" + "stride" + "=" + str(tm.stride) + r"\l" + if tm.is_quantized: + if tm.qscheme in { + torch.per_tensor_affine, + torch.per_tensor_symmetric, + }: + result += "|" + "q_scale" + "=" + str(tm.q_scale) + r"\l" + result += "|" + "q_zero_point" + "=" + str(tm.q_zero_point) + r"\l" + result += "|" + "qscheme" + "=" + str(tm.qscheme) + r"\l" + return result + def _get_tensor_label(self, t: torch.Tensor) -> str: return str(t.dtype) + str(list(t.shape)) + r"\l" From 96651458eb97a1898a307a0f7c03fa963b4cc7d3 Mon Sep 17 00:00:00 2001 From: Facebook Community Bot Date: Thu, 10 Jun 2021 04:39:49 -0700 Subject: [PATCH 003/305] Automated submodule update: tensorpipe (#59374) Summary: This is an automated pull request to update the first-party submodule for [pytorch/tensorpipe](https://github.com/pytorch/tensorpipe). New submodule commit: https://github.com/pytorch/tensorpipe/commit/e942ea15138d1ca5eac08b194858b54fa1e4ab2f Pull Request resolved: https://github.com/pytorch/pytorch/pull/59374 Test Plan: Ensure that CI jobs succeed on GitHub before landing. Reviewed By: lw Differential Revision: D28867855 fbshipit-source-id: e1325046003f5c546f02024ff4c427c91721cd7e --- third_party/tensorpipe | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/tensorpipe b/third_party/tensorpipe index a0c6aa14224d7..e942ea15138d1 160000 --- a/third_party/tensorpipe +++ b/third_party/tensorpipe @@ -1 +1 @@ -Subproject commit a0c6aa14224d76f64952ffbca371e761493d4049 +Subproject commit e942ea15138d1ca5eac08b194858b54fa1e4ab2f From 5e3e50472864fa5bd8a7357935c6fcb2c1f785da Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Thu, 10 Jun 2021 06:35:37 -0700 Subject: [PATCH 004/305] Update TensorPipe submodule (#59789) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59789 The bot messed up in D28867855 (https://github.com/pytorch/pytorch/commit/96651458eb97a1898a307a0f7c03fa963b4cc7d3) so I've got to do it manually. Test Plan: CI Reviewed By: beauby Differential Revision: D29027901 fbshipit-source-id: 9438e0cfbe932fbbd1e252ab57e2b1b23f9e44cf --- third_party/tensorpipe | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/tensorpipe b/third_party/tensorpipe index e942ea15138d1..42a67277c1882 160000 --- a/third_party/tensorpipe +++ b/third_party/tensorpipe @@ -1 +1 @@ -Subproject commit e942ea15138d1ca5eac08b194858b54fa1e4ab2f +Subproject commit 42a67277c1882c90cec0da6e57afb20247424994 From 58412740ae433c89320ec45eb96a28bb28b23fb6 Mon Sep 17 00:00:00 2001 From: Heitor Schueroff Date: Thu, 10 Jun 2021 08:00:43 -0700 Subject: [PATCH 005/305] Added doc for torch.einsum sublist format (#57038) Summary: Adds documentation for the new sublist format for `torch.einsum` closes https://github.com/pytorch/pytorch/issues/21412 Pull Request resolved: https://github.com/pytorch/pytorch/pull/57038 Reviewed By: mruberry Differential Revision: D28994431 Pulled By: heitorschueroff fbshipit-source-id: 3dfb154fe6e4c440ac67c2dd92727bb5ecfe289e --- torch/functional.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/torch/functional.py b/torch/functional.py index 0ca423e717b4c..8f50ae80ccad0 100644 --- a/torch/functional.py +++ b/torch/functional.py @@ -209,6 +209,14 @@ def einsum(*args): run faster or consume less memory. Projects like opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/) can optimize the formula for you. + .. note:: + + As of PyTorch 1.10 :func:`torch.einsum` also supports the sublist format (see examples below). In this format, + subscripts for each operand are specified by sublists, list of integers in the range [0, 52). These sublists + follow their operands, and an extra sublist can appear at the end of the input to specify the output's + subscripts., e.g.`torch.einsum(op1, sublist1, op2, sublist2, ..., [subslist_out])`. Python's `Ellipsis` object + may be provided in a sublist to enable broadcasting as described in the Equation section above. + Args: equation (string): The subscripts for the Einstein summation. operands (List[Tensor]): The tensors to compute the Einstein summation of. @@ -246,6 +254,17 @@ def einsum(*args): [[ 2.8153, 1.8787, -4.3839, -1.2112], [ 0.3728, -2.1131, 0.0921, 0.8305]]]) + # with sublist format and ellipsis + >>> torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2]) + tensor([[[-1.0564, -1.5904, 3.2023, 3.1271], + [-1.6706, -0.8097, -0.8025, -2.1183]], + + [[ 4.2239, 0.3107, -0.5756, -0.2354], + [-1.4558, -0.3460, 1.5087, -0.8530]], + + [[ 2.8153, 1.8787, -4.3839, -1.2112], + [ 0.3728, -2.1131, 0.0921, 0.8305]]]) + # batch permute >>> A = torch.randn(2, 3, 4, 5) >>> torch.einsum('...ij->...ji', A).shape From 3176f1669154d8d161de88c07a57972d60e7c3ae Mon Sep 17 00:00:00 2001 From: Kimish Patel Date: Thu, 10 Jun 2021 08:23:10 -0700 Subject: [PATCH 006/305] [Pytorch benchmark] Add BMM benchmark (#59595) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59595 ghstack-source-id: 130946743 Test Plan: bmm_test Reviewed By: mingzhe09088 Differential Revision: D28873228 fbshipit-source-id: 6e4cb04bb6c63f5f68d8f23c13738e2d58ab499c --- benchmarks/operator_benchmark/pt/bmm_test.py | 29 ++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 benchmarks/operator_benchmark/pt/bmm_test.py diff --git a/benchmarks/operator_benchmark/pt/bmm_test.py b/benchmarks/operator_benchmark/pt/bmm_test.py new file mode 100644 index 0000000000000..aa6441fe5764e --- /dev/null +++ b/benchmarks/operator_benchmark/pt/bmm_test.py @@ -0,0 +1,29 @@ +import operator_benchmark as op_bench +import torch + +"""Microbenchmarks for add_ operator. Supports both Caffe2/PyTorch.""" + +class BmmBenchmark(op_bench.TorchBenchmarkBase): + def init(self, B, M, N, K, device): + self.inputs = { + "batch1": torch.rand((B, M, K), device=device, requires_grad=self.auto_set()), + "batch2": torch.rand((B, K, N,), device=device, requires_grad=self.auto_set()) + } + self.set_module_name("bmm") + + def forward(self, batch1, batch2): + return torch.bmm(batch1, batch2) + +bmm_configs = op_bench.cross_product_configs( + B=[2, 100], + M=[8, 256], + N=[256, 16], + K=[16, 32], + device=['cpu'], + tags=["short"], +) + +op_bench.generate_pt_test(bmm_configs, BmmBenchmark) + +if __name__ == "__main__": + op_bench.benchmark_runner.main() From 4f79270b8933a1e580343917137f0131754cda21 Mon Sep 17 00:00:00 2001 From: Kimish Patel Date: Thu, 10 Jun 2021 08:23:10 -0700 Subject: [PATCH 007/305] [PyTorch ] Thread parallel bmm across batch dim (#59596) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59596 Parallelize batch matmul across batch dim. This was found to improve perf for some usecases on mobile. ghstack-source-id: 130989569 Test Plan: CI unit tests Reviewed By: albanD Differential Revision: D26833417 fbshipit-source-id: 9b84d89d29883a6c9d992d993844dd31a25f76b1 --- aten/src/ATen/native/LinearAlgebra.cpp | 60 +++++++++++++++++++++++--- test/test_torch.py | 53 +++++++++++++++++++++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp index cf8189d3a10e0..6bce8120e9e39 100644 --- a/aten/src/ATen/native/LinearAlgebra.cpp +++ b/aten/src/ATen/native/LinearAlgebra.cpp @@ -1256,14 +1256,64 @@ static inline Tensor& bmm_out_or_baddbmm_(Tensor& self_or_result, const Tensor& && self_or_result.is_contiguous()) { at::native::_baddbmm_mkl_(self_or_result, batch1, batch2, beta, alpha); } else { // split along batch dimension +#ifdef C10_MOBILE + /* + * We only do multithreading when Inference mode is enabled because various + * thread local state is not appropriately propagated through + * at::parallel_for. e.g. RecordFunction related state, dispatchKeySet Big + * concern with this is that if we use at::parallel_for where state is not + * propagated then dispatch machinery may work differently on main thread + * vs. other threads, leading to undefined behavior. + * Thus it is recommended to not use at::parallel_for where lambdas do + * ops that go through dispatcher. + * For now we circument this by InferenceMode guard in order to unlock + * performance. + * Longer term we probably want a separate API that explicitly calls out + * the TLS that it propagates. + * Also note that this is enabled for mobile only because blas + * implementation for non-mobile build is already multithreaded. + */ + // Benchmarking was done as follows: + // bmm_test: operator benchmark under + // benchmarks/operator_benchmarks/pt/bmm_test.py Ran this benchmark for + // various matrix sizes on Samsung S8U + const bool enable_multithreaded_bmm = c10::InferenceMode::is_enabled() && + bs >= 4 && res_rows >= 4 && res_cols >= 16 && contraction_size >= 16; +#else + const bool enable_multithreaded_bmm{false}; +#endif if (is_bmm_out) { - for (int64_t b = 0; b < bs; b++) { - auto r = self_or_result.select(0, b); - addmm_impl_cpu_(r, r, batch1.select(0, b), batch2.select(0, b), 0, 1); + if (enable_multithreaded_bmm) { + auto bmm_out_fn = [&](uint64_t start, uint64_t end) { + c10::InferenceMode guard; + for (int64_t b = start; b < end; b++) { + auto r = self_or_result.select(0, b); + addmm_impl_cpu_( + r, r, batch1.select(0, b), batch2.select(0, b), 0, 1); + } + }; + at::parallel_for(0, bs, 1, bmm_out_fn); + } else { + for (int64_t b = 0; b < bs; b++) { + auto r = self_or_result.select(0, b); + addmm_impl_cpu_(r, r, batch1.select(0, b), batch2.select(0, b), 0, 1); + } } } else { - for (int64_t b = 0; b < bs; b++) { - self_or_result.select(0, b).addmm_(batch1.select(0, b), batch2.select(0, b), beta, alpha); + if (enable_multithreaded_bmm) { + auto bmm_fn = [&](uint64_t start, uint64_t end) { + c10::InferenceMode guard; + for (int64_t b = start; b < end; b++) { + self_or_result.select(0, b).addmm_( + batch1.select(0, b), batch2.select(0, b), beta, alpha); + } + }; + at::parallel_for(0, bs, 1, bmm_fn); + } else { + for (int64_t b = 0; b < bs; b++) { + self_or_result.select(0, b).addmm_( + batch1.select(0, b), batch2.select(0, b), beta, alpha); + } } } } diff --git a/test/test_torch.py b/test/test_torch.py index 561ddd7e9feac..b30697e099443 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -6,6 +6,7 @@ import gc import io import inspect +import itertools import math import random import re @@ -8015,6 +8016,58 @@ def test_resurrected_weak_ref(self): del y x.sigmoid() + @torch.inference_mode() + def test_bmm_multithreaded(self): + device = 'cpu' + num_threads = torch.get_num_threads() + + torch.set_num_threads(4) + batch_sizes = [1, 10] + M, N, O = 23, 8, 12 + dtype = torch.float32 + numpy_dtype = dtype + + def invert_perm(p): + d = {x: i for i, x in enumerate(p)} + return (d[0], d[1], d[2]) + + def generate_inputs(num_batches): + # transposed tensors + for perm1, perm2 in itertools.product(itertools.permutations((0, 1, 2)), repeat=2): + b1 = make_tensor((num_batches, M, N), device, dtype, low=-1, high=1) + b2 = make_tensor((num_batches, N, O), device, dtype, low=-1, high=1) + b1 = b1.permute(perm1).contiguous().permute(invert_perm(perm1)) + b2 = b2.permute(perm2).contiguous().permute(invert_perm(perm2)) + yield b1, b2 + # broadcasting tensors + for b1, b2, b3, b4, b5, b6 in itertools.product((True, False), repeat=6): + shape1 = (num_batches if b1 else 1, M if b2 else 1, N if b3 else 1) + shape2 = (num_batches if b4 else 1, N if b5 else 1, O if b6 else 1) + b1 = make_tensor(shape1, device, dtype, low=-1, high=1).expand(num_batches, M, N) + b2 = make_tensor(shape2, device, dtype, low=-1, high=1).expand(num_batches, N, O) + yield b1, b2 + # zero-sized tensors + for z1, z2, z3, z4 in itertools.product((True, False), repeat=4): + shape1 = (num_batches if z1 else 0, M if z2 else 0, N if z3 else 0) + shape2 = (num_batches if z1 else 0, N if z3 else 0, O if z4 else 0) + b1 = torch.randn(shape1, dtype=dtype, device=device) + b2 = torch.randn(shape2, dtype=dtype, device=device) + yield b1, b2 + + try: + for num_batches in batch_sizes: + for (b1, b2), perm3 in itertools.product(generate_inputs(num_batches), itertools.permutations((0, 1, 2))): + res1 = torch.bmm(b1, b2) + res2 = torch.full((num_batches, M, O), math.nan, dtype=dtype, device=device) \ + .permute(perm3).contiguous().permute(invert_perm(perm3)) + torch.bmm(b1, b2, out=res2) + expect = torch.from_numpy( + b1.to(numpy_dtype).cpu().numpy() @ b2.to(numpy_dtype).cpu().numpy()).to(device=device, dtype=dtype) + self.assertEqual(expect, res1) + self.assertEqual(expect, res2) + finally: + torch.set_num_threads(num_threads) + # TODO: these empy classes are temporarily instantiated for XLA compatibility # once XLA updates their test suite it should be removed From 51d954e8e42d2992dd455f922ce21b22983245fa Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Thu, 10 Jun 2021 09:40:20 -0700 Subject: [PATCH 008/305] Link ATEN tests with OpenMP runtime (#59733) Summary: Even if OpenMP extensions are supported by compiler, OpenMP runtime library is not always implicitly added as dependency by linker Above fixes linker problems on Apple M1, when libomp.dylib is installed via conda, when tests that directly use OpenMP pragams fail to link with following errors: ``` /Library/Developer/CommandLineTools/usr/bin/c++ -Wno-deprecated -fvisibility-inlines-hidden -Wno-deprecated-declarations -DUSE_PTHREADPOOL -Xpreprocessor -fopenmp -I/Users/nshulga/miniforge3/include -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wno-narrowing -Wall -Wextra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -Wno-invalid-partial-specialization -Wno-typedef-redefinition -Wno-unknown-warning-option -Wno-unused-private-field -Wno-inconsistent-missing-override -Wno-aligned-allocation-unavailable -Wno-c++14-extensions -Wno-constexpr-not-const -Wno-missing-braces -Qunused-arguments -fcolor-diagnostics -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-unused-private-field -Wno-missing-braces -Wno-c++14-extensions -Wno-constexpr-not-const -O3 -DNDEBUG -DNDEBUG -arch arm64 -isysroot /Library/Developer/CommandLineTools/SDKs/MacOSX11.3.sdk -Wl,-search_paths_first -Wl,-headerpad_max_install_names -rdynamic caffe2/CMakeFiles/test_parallel.dir/__/aten/src/ATen/test/test_parallel.cpp.o -o bin/test_parallel -Wl,-rpath,/Users/nshulga/git/pytorch/build/lib lib/libgtest_main.a lib/libtorch.dylib lib/libtorch_cpu.dylib lib/libprotobuf.a lib/libc10.dylib lib/libgtest.a && : Undefined symbols for architecture arm64: "___kmpc_fork_call", referenced from: TestParallel_NestedParallel_Test::TestBody() in test_parallel.cpp.o TestParallel_Exceptions_Test::TestBody() in test_parallel.cpp.o "_omp_get_max_threads", referenced from: TestParallel_NestedParallel_Test::TestBody() in test_parallel.cpp.o TestParallel_Exceptions_Test::TestBody() in test_parallel.cpp.o "_omp_get_num_threads", referenced from: _.omp_outlined. in test_parallel.cpp.o _.omp_outlined..31 in test_parallel.cpp.o "_omp_get_thread_num", referenced from: _.omp_outlined. in test_parallel.cpp.o _.omp_outlined..31 in test_parallel.cpp.o "_omp_in_parallel", referenced from: TestParallel_NestedParallel_Test::TestBody() in test_parallel.cpp.o TestParallel_Exceptions_Test::TestBody() in test_parallel.cpp.o ld: symbol(s) not found for architecture arm64 clang: error: linker command failed with exit code 1 (use -v to see invocation) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/59733 Reviewed By: walterddr, seemethere Differential Revision: D29005511 Pulled By: malfet fbshipit-source-id: daab5e1b0a58d9b60a8992ef40c743e4b619dac7 --- caffe2/CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index e9de60729ce6f..7b34d8e19b108 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -1664,6 +1664,11 @@ if(BUILD_TEST) get_filename_component(test_name ${test_src} NAME_WE) add_executable(${test_name} "${test_src}") target_link_libraries(${test_name} torch_library gtest_main) + if(USE_OPENMP) + # -fopenmp is a compile time flag and as result not guaranteed + # to link executable against OpenMP runtime library + target_link_libraries(${test_name} ${OpenMP_CXX_LIBRARIES}) + endif() target_include_directories(${test_name} PRIVATE $) target_include_directories(${test_name} PRIVATE $) target_include_directories(${test_name} PRIVATE ${Caffe2_CPU_INCLUDE}) From 7af9252ed789f72194aa522f7b168c88e00ecc44 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 10 Jun 2021 09:41:15 -0700 Subject: [PATCH 009/305] [skip ci] export_slow_tests.py - Add option to ignore small differences (#59759) Summary: This would lower the number of unnecessary commits to pytorch/test-infra by only exporting a different stats file when the stats are varying enough. This way, if the slow test cases we gather from S3 are the same and their times are trivially different, then we do not bother exporting a different stats file when the --ignore-small-diffs option is enabled. We instead export the stats already in test-infra, so that when it tries to commit, it realizes it would be an empty commit and not add to the git history. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59759 Test Plan: Run `python tools/export_slow_tests.py --ignore-small-diffs `. Reviewed By: walterddr Differential Revision: D29032712 Pulled By: janeyx99 fbshipit-source-id: 41d522a4c5f710e776acd1512d41be9791d0cf63 --- tools/export_slow_tests.py | 43 ++++++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/tools/export_slow_tests.py b/tools/export_slow_tests.py index 3aad29cd0b7ad..c0cb43a4a6144 100644 --- a/tools/export_slow_tests.py +++ b/tools/export_slow_tests.py @@ -6,11 +6,12 @@ import statistics from collections import defaultdict from tools.stats_utils.s3_stat_parser import get_previous_reports_for_branch, Report, Version2Report -from typing import cast, DefaultDict, Dict, List +from typing import cast, DefaultDict, Dict, List, Any +from urllib.request import urlopen SLOW_TESTS_FILE = '.pytorch-slow-tests.json' SLOW_TEST_CASE_THRESHOLD_SEC = 60.0 - +RELATIVE_DIFFERENCE_THRESHOLD = 0.1 def get_test_case_times() -> Dict[str, float]: reports: List[Report] = get_previous_reports_for_branch('origin/viable/strict', "") @@ -37,11 +38,35 @@ def filter_slow_tests(test_cases_dict: Dict[str, float]) -> Dict[str, float]: return {test_case: time for test_case, time in test_cases_dict.items() if time >= SLOW_TEST_CASE_THRESHOLD_SEC} -def export_slow_tests(filename: str) -> None: +def get_test_infra_slow_tests() -> Dict[str, float]: + url = "https://raw.githubusercontent.com/pytorch/test-infra/master/stats/slow-tests.json" + contents = urlopen(url, timeout=1).read().decode('utf-8') + return cast(Dict[str, float], json.loads(contents)) + + +def too_similar(calculated_times: Dict[str, float], other_times: Dict[str, float], threshold: float) -> bool: + # check that their keys are the same + if calculated_times.keys() != other_times.keys(): + return False + + for test_case, test_time in calculated_times.items(): + other_test_time = other_times[test_case] + relative_difference = abs((other_test_time - test_time) / max(other_test_time, test_time)) + if relative_difference > threshold: + return False + return True + + +def export_slow_tests(options: Any) -> None: + filename = options.filename if os.path.exists(filename): print(f'Overwriting existent file: {filename}') with open(filename, 'w+') as file: slow_test_times: Dict[str, float] = filter_slow_tests(get_test_case_times()) + if options.ignore_small_diffs: + test_infra_slow_tests_dict = get_test_infra_slow_tests() + if too_similar(slow_test_times, test_infra_slow_tests_dict, options.ignore_small_diffs): + slow_test_times = test_infra_slow_tests_dict json.dump(slow_test_times, file, indent=' ', separators=(',', ': '), sort_keys=True) file.write('\n') @@ -58,12 +83,22 @@ def parse_args() -> argparse.Namespace: const=SLOW_TESTS_FILE, help='Specify a file path to dump slow test times from previous S3 stats. Default file path: .pytorch-slow-tests.json', ) + parser.add_argument( + '--ignore-small-diffs', + nargs='?', + type=float, + const=RELATIVE_DIFFERENCE_THRESHOLD, + help='Compares generated results with stats/slow-tests.json in pytorch/test-infra. If the relative differences ' + 'between test times for each test are smaller than the threshold and the set of test cases have not ' + 'changed, we will export the stats already in stats/slow-tests.json. Else, we will export the calculated ' + 'results. The default threshold is 10%.', + ) return parser.parse_args() def main() -> None: options = parse_args() - export_slow_tests(options.filename) + export_slow_tests(options) if __name__ == '__main__': From f2406240800e7cb3b806a4de94e9ef3d4838c547 Mon Sep 17 00:00:00 2001 From: Edvard Ghazaryan Date: Thu, 10 Jun 2021 10:31:51 -0700 Subject: [PATCH 010/305] displays graph node's info (#59679) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59679 Displays info about graph's nodes Test Plan: Expected view: %wide_offset.1 : Tensor = aten::add(%wide.1, %self._mu, %4) i0: Tensor CPUFloatType {32, 50} i1: Tensor CPUFloatType {1, 50} i2: int {1} o0: Tensor CPUFloatType {32, 50} %wide_normalized.1 : Tensor = aten::mul(%wide_offset.1, %self._sigma) i0: Tensor CPUFloatType {32, 50} i1: Tensor CPUFloatType {1, 50} o0: Tensor CPUFloatType {32, 50} %wide_preproc.1 : Tensor = aten::clamp(%wide_normalized.1, %5, %6) i0: Tensor CPUFloatType {32, 50} i1: double {0} i2: double {10} o0: Tensor CPUFloatType {32, 50} %user_emb_t.1 : Tensor = aten::transpose(%user_emb.1, %4, %7) i0: Tensor CPUFloatType {32, 1, 32} i1: int {1} i2: int {2} o0: Tensor CPUFloatType {32, 32, 1} %dp_unflatten.1 : Tensor = aten::bmm(%ad_emb_packed.1, %user_emb_t.1) i0: Tensor CPUFloatType {32, 1, 32} i1: Tensor CPUFloatType {32, 32, 1} o0: Tensor CPUFloatType {32, 1, 1} %31 : Tensor = static_runtime::flatten_copy(%dp_unflatten.1, %4, %8) i0: Tensor CPUFloatType {32, 1, 1} i1: int {1} i2: int {-1} o0: Tensor CPUFloatType {32, 1} %19 : Tensor[] = prim::ListConstruct(%31, %wide_preproc.1) i0: Tensor CPUFloatType {32, 1} i1: Tensor CPUFloatType {32, 50} o0: TensorList {2} %input.1 : Tensor = aten::cat(%19, %4) i0: TensorList {2} i1: int {1} o0: Tensor CPUFloatType {32, 51} %fc1.1 : Tensor = aten::addmm(%self._fc_b, %input.1, %29, %4, %4) i0: Tensor CPUFloatType {1} i1: Tensor CPUFloatType {32, 51} i2: Tensor CPUFloatType {51, 1} i3: int {1} i4: int {1} o0: Tensor CPUFloatType {32, 1} %23 : Tensor = aten::sigmoid(%fc1.1) i0: Tensor CPUFloatType {32, 1} o0: Tensor CPUFloatType {32, 1} %24 : (Tensor) = prim::TupleConstruct(%23) i0: Tensor CPUFloatType {32, 1} o0: Tuple {1} Reviewed By: hlu1 Differential Revision: D28592852 fbshipit-source-id: 09174014f7d0ce25c511025d2b376f14e16c8a4a --- torch/csrc/jit/runtime/static/impl.cpp | 63 ++++++++++++++++++++++++++ torch/csrc/jit/runtime/static/impl.h | 2 + 2 files changed, 65 insertions(+) diff --git a/torch/csrc/jit/runtime/static/impl.cpp b/torch/csrc/jit/runtime/static/impl.cpp index 7ab60de813228..9b23c0ea3bacc 100644 --- a/torch/csrc/jit/runtime/static/impl.cpp +++ b/torch/csrc/jit/runtime/static/impl.cpp @@ -872,6 +872,69 @@ float StaticRuntime::benchmark_model( return millis / static_cast(main_runs); } +bool display_ivalue(const IValue& iv) { + if (iv.isTensor()) { + std::cout << "Tensor " << iv.toTensor().toString() << " {"; + for (auto i = 0; i < iv.toTensor().sizes().size(); ++i) { + std::cout << iv.toTensor().sizes()[i]; + if (iv.toTensor().sizes().size() > i + 1) { + std::cout << ", "; + } + } + std::cout << "}\n"; + return true; + } else if (iv.isTensorList()) { + std::cout << "TensorList {" << iv.toTensorList().size() << "}\n"; + return true; + } else if (iv.isGenericDict()) { + std::cout << "Dict {" << iv.toGenericDict().size() << "}\n"; + return true; + } else if (iv.isTuple()) { + std::cout << "Tuple {" << iv.toTuple()->elements().size() << "}\n"; + return true; + } else if (iv.isInt()) { + std::cout << "int {" << iv.toInt() << "}\n"; + return true; + } else if (iv.isBool()) { + std::cout << "bool {" << iv.toBool() << "}\n"; + return true; + } else if (iv.isDouble()) { + std::cout << "double {" << iv.toDouble() << "}\n"; + return true; + } + return false; +} + +void display_pnode_info(const ProcessedNode& pnode) { + pnode.node()->print(std::cout, 0, nullptr, false); + const std::vector& inputs = pnode.inputs(); + for (auto i = 0; i < inputs.size(); ++i) { + std::cout << "\ti" << i << ": "; + if (!display_ivalue(*inputs[i])) { + std::cout << *(pnode.node()->inputs()[i]->type()) << '\n'; + } + } + const std::vector& outputs = pnode.outputs(); + for (auto i = 0; i < outputs.size(); ++i) { + std::cout << "\to" << i << ": "; + if (!display_ivalue(outputs[i])) { + std::cout << *(pnode.node()->outputs()[i]->type()) << '\n'; + } + } +} + +void StaticRuntime::display_nodes(const std::vector& args) { + c10::InferenceMode mode; + std::vector stack(args); + for (size_t i = 0; i < stack.size(); i++) { + Input(i) = stack[i]; + } + for (auto& node : nodes_) { + node.run(); + display_pnode_info(node); + } +} + StaticRuntime::IndividualMetrics StaticRuntime::benchmark_individual_ops( const std::vector& args, const std::unordered_map& kwargs, diff --git a/torch/csrc/jit/runtime/static/impl.h b/torch/csrc/jit/runtime/static/impl.h index 50a84dade7ec4..01cbcb7a83bf6 100644 --- a/torch/csrc/jit/runtime/static/impl.h +++ b/torch/csrc/jit/runtime/static/impl.h @@ -188,6 +188,8 @@ class TORCH_API StaticRuntime { const std::vector& args, const std::unordered_map& kwargs); + void display_nodes(const std::vector& args); + void benchmark( const std::vector& args, const std::unordered_map& kwargs, From 9bcef86d1831ac0af8104e57b367b0020e38e02c Mon Sep 17 00:00:00 2001 From: Jeffrey Wan Date: Thu, 10 Jun 2021 11:31:27 -0700 Subject: [PATCH 011/305] Split slow gradcheck periodic CI job so that it does not time out (#59736) Summary: Fixes #{issue number} Pull Request resolved: https://github.com/pytorch/pytorch/pull/59736 Reviewed By: albanD Differential Revision: D29008100 Pulled By: soulitzer fbshipit-source-id: 76da971356fd985dfbfa56d3573f31ef04701773 --- .circleci/config.yml | 12 ++++++++++-- .../workflows/workflows-scheduled-ci.yml | 12 ++++++++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 62b00b0792204..273ea971706ed 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -9537,10 +9537,18 @@ workflows: build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-build" docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" - pytorch_linux_test: - name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_tests + name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_test1 requires: - periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_build - build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-tests" + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test1" + docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + use_cuda_docker_runtime: "1" + resource_class: gpu.medium + - pytorch_linux_test: + name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_test2 + requires: + - periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_build + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test2" docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" use_cuda_docker_runtime: "1" resource_class: gpu.medium diff --git a/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml b/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml index 01a416558736c..8a2115a796807 100644 --- a/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml +++ b/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml @@ -186,10 +186,18 @@ build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-build" docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" - pytorch_linux_test: - name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_tests + name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_test1 requires: - periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_build - build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-tests" + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test1" + docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + use_cuda_docker_runtime: "1" + resource_class: gpu.medium + - pytorch_linux_test: + name: periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_test2 + requires: + - periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_build + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test2" docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" use_cuda_docker_runtime: "1" resource_class: gpu.medium From c90260905f38026b9b47c792b9d1ed9b39a2cd5b Mon Sep 17 00:00:00 2001 From: Kshiteej K Date: Thu, 10 Jun 2021 11:58:23 -0700 Subject: [PATCH 012/305] [fix] torch.{lin, log}space(): properly examine passed dtype (#53685) Summary: Fixes https://github.com/pytorch/pytorch/issues/53171 Pull Request resolved: https://github.com/pytorch/pytorch/pull/53685 Reviewed By: jbschlosser Differential Revision: D28331863 Pulled By: anjali411 fbshipit-source-id: e89359b607d058158cfa1c9a82389d9a4a71185b --- aten/src/ATen/native/TensorFactories.cpp | 29 +++--- test/test_tensor_creation_ops.py | 27 +++-- tools/autograd/gen_python_functions.py | 2 +- .../templates/python_torch_functions.cpp | 98 +++++++++++++++++++ tools/pyi/gen_pyi.py | 4 + 5 files changed, 134 insertions(+), 26 deletions(-) diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp index 7c8dc389b8973..5dee5f309138a 100644 --- a/aten/src/ATen/native/TensorFactories.cpp +++ b/aten/src/ATen/native/TensorFactories.cpp @@ -517,25 +517,20 @@ namespace { TensorOptions linspace_logspace_infer_options( const Scalar& start, const Scalar& end, - const TensorOptions& options) { - auto result_options = options; + const TensorOptions& options, + const char* fn_name) { if (start.isComplex() || end.isComplex()) { - // Since result_options.has_dtype() returns true (dtype is default type), - // even if the user hasn't specified the dtype. - // We just check to see if either `start` or `end` is complex, - // and if the `result_dtype` is not complex (be it default float type or - // user provided), we cast it to default complex dtype with a Warning!. - auto result_dtype = c10::typeMetaToScalarType(options.dtype()); - if (!at::isComplexType(result_dtype)) { - TORCH_WARN( - "As either `start` or `stop` is complex, return type will be the complex dtype corresponding to default dtype.", - "In future, this may throw an error when a non-complex dtype arg is passed as input along ", - "with complex valued start or end value."); - result_options = result_options.dtype(c10::get_default_complex_dtype()); + const auto default_complex_dtype = c10::get_default_complex_dtype(); + if (options.has_dtype()) { + auto dtype = c10::typeMetaToScalarType(options.dtype()); + TORCH_CHECK(at::isComplexType(dtype), + fn_name, ": inferred dtype ", default_complex_dtype, " can't be safely cast to passed dtype ", dtype); + } else { + return options.dtype(default_complex_dtype); } } - return result_options; + return options.has_dtype() ? options : options.dtype(c10::get_default_dtype()); } } // anonymous namespace @@ -554,7 +549,7 @@ Tensor linspace( const auto steps_ = steps.value_or(100); TORCH_CHECK(steps_ >= 0, "number of steps must be non-negative"); - auto result_options = linspace_logspace_infer_options(start, end, options); + auto result_options = linspace_logspace_infer_options(start, end, options, "torch.linspace()"); Tensor result = at::empty({steps_}, result_options); return at::linspace_out(result, start, end, steps); } @@ -575,7 +570,7 @@ Tensor logspace( const auto steps_ = steps.value_or(100); TORCH_CHECK(steps_ >= 0, "number of steps must be non-negative"); - auto result_options = linspace_logspace_infer_options(start, end, options); + auto result_options = linspace_logspace_infer_options(start, end, options, "torch.logspace()"); Tensor result = at::empty({steps_}, result_options); return at::logspace_out(result, start, end, steps, base); } diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py index 69ce474379376..e56bd658d9d1d 100644 --- a/test/test_tensor_creation_ops.py +++ b/test/test_tensor_creation_ops.py @@ -2754,6 +2754,15 @@ def test_linspace(self, device, dtype): # steps = 0 self.assertEqual(torch.linspace(0, 1, 0, device=device, dtype=dtype).numel(), 0, atol=0, rtol=0) + if dtype == torch.float: + # passed dtype can't be safely casted to inferred dtype + with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"): + torch.linspace(0, 1j, 5, device=device, dtype=dtype) + with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"): + torch.linspace(0j, 1, 5, device=device, dtype=dtype) + with self.assertRaisesRegex(RuntimeError, r"torch.linspace\(\): inferred dtype"): + torch.linspace(0j, 1j, 5, device=device, dtype=dtype) + # Check linspace for generating the correct output for each dtype. start = 0 if dtype == torch.uint8 else -100 expected_lin = torch.tensor([start + .5 * i for i in range(401)], device=device, dtype=torch.double) @@ -2797,14 +2806,7 @@ def _test_linspace_logspace_deduction_helper(self, fn, device): if isinstance(start, complex) or isinstance(end, complex): dtype = torch.cfloat - if dtype == torch.cfloat: - # TODO(kshitij12345): Fix unnecessary warning - # Reference: https://github.com/pytorch/pytorch/issues/53171 - with self.assertWarnsRegex(UserWarning, - "As either `start` or `stop` is complex"): - self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype) - else: - self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype) + self.assertEqual(fn(start, end, steps=100, device=device).dtype, dtype) def test_linspace_deduction(self, device): # Test deduction from input parameters. @@ -2896,6 +2898,15 @@ def test_logspace(self, device, dtype): self.assertEqual(torch.logspace(0, 1, 1, device=device, dtype=dtype), torch.ones(1, device=device, dtype=dtype), atol=0, rtol=0) + if dtype == torch.float: + # passed dtype can't be safely casted to inferred dtype + with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"): + torch.logspace(0, 1j, 5, device=device, dtype=dtype) + with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"): + torch.logspace(0j, 1, 5, device=device, dtype=dtype) + with self.assertRaisesRegex(RuntimeError, r"torch.logspace\(\): inferred dtype"): + torch.logspace(0j, 1j, 5, device=device, dtype=dtype) + # Check precision - start, stop and base are chosen to avoid overflow # steps is chosen so that step size is not subject to rounding error # a tolerance is needed for gpu tests due to differences in computation diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index 4f54906f7634b..a59f67bc244d8 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -80,7 +80,7 @@ '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward', '.*_forward_out', '_unsafe_view', 'tensor', '_?sparse_coo_tensor.*', '_?sparse_csr_tensor.*', - '_arange.*', '_range.*', '_linspace.*', '_logspace.*', + '_arange.*', '_range.*', 'linspace.*', 'logspace.*', '_sparse_add_out', '_sparse_div.*', '_sparse_mul.*', '_sparse_sub.*', '_sparse_dense_add_out', 'index', 'unique_dim_consecutive', '_cumsum.*', '_cumprod.*', '_sum.*', '_prod.*', diff --git a/tools/autograd/templates/python_torch_functions.cpp b/tools/autograd/templates/python_torch_functions.cpp index 5c0fd35416118..a9c5e3ae0d66e 100644 --- a/tools/autograd/templates/python_torch_functions.cpp +++ b/tools/autograd/templates/python_torch_functions.cpp @@ -467,6 +467,102 @@ static PyObject * THPVariable_get_device(PyObject* self_, PyObject* args, PyObje static PyObject * THPVariable_numel(PyObject* self_, PyObject* args, PyObject* kwargs); +// linspace +static PyObject * THPVariable_linspace(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "linspace(Scalar start, Scalar end, int64_t? steps=None, *, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)", + }, /*traceable=*/true); + + ParsedArgs<9> parsed_args; + auto _r = parser.parse(nullptr, args, kwargs, parsed_args); + if(_r.has_torch_function()) { + return handle_torch_function(_r, nullptr, args, kwargs, THPVariableFunctionsModule, "torch"); + } + if (_r.isNone(3)) { + // aten::linspace(Scalar start, Scalar end, int? steps=None, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + + // NOTE: r.scalartype(X) gives the default dtype if r.isNone(X) + // This leads to problem in the operator argument checks, + // when either `start` or `end` is complex and dtype is None + const auto options = TensorOptions() + .dtype(_r.scalartypeOptional(4)) + .device(_r.device(6)) + .layout(_r.layoutOptional(5)) + .requires_grad(_r.toBool(8)) + .pinned_memory(_r.toBool(7)); + torch::utils::maybe_initialize_cuda(options); + + auto dispatch_linspace = [](Scalar start, Scalar end, c10::optional steps, TensorOptions options) -> Tensor { + pybind11::gil_scoped_release no_gil; + return torch::linspace(start, end, steps, options); + }; + return wrap(dispatch_linspace(_r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), options)); + } else { + // aten::linspace.out(Scalar start, Scalar end, int? steps=None, *, Tensor(a!) out) -> Tensor(a!) + check_out_type_matches(_r.tensor(3), _r.scalartype(4), + _r.isNone(4), _r.layoutOptional(5), + _r.device(6), _r.isNone(6)); + + auto dispatch_linspace_out = [](Tensor out, Scalar start, Scalar end, c10::optional steps) -> Tensor { + pybind11::gil_scoped_release no_gil; + return at::linspace_out(out, start, end, steps); + }; + return wrap(dispatch_linspace_out(_r.tensor(3), _r.scalar(0), _r.scalar(1), _r.toInt64Optional(2)).set_requires_grad(_r.toBool(8))); + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +// logspace +static PyObject * THPVariable_logspace(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "logspace(Scalar start, Scalar end, int64_t? steps=None, double base=10.0, *, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)", + }, /*traceable=*/true); + + ParsedArgs<10> parsed_args; + auto _r = parser.parse(nullptr, args, kwargs, parsed_args); + if(_r.has_torch_function()) { + return handle_torch_function(_r, nullptr, args, kwargs, THPVariableFunctionsModule, "torch"); + } + if (_r.isNone(4)) { + // aten::logspace(Scalar start, Scalar end, int? steps=None, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + + // NOTE: r.scalartype(X) gives the default dtype if r.isNone(X) + // This leads to problem in the operator argument checks, + // when either `start` or `end` is complex and dtype is None + const auto options = TensorOptions() + .dtype(_r.scalartypeOptional(5)) + .device(_r.device(7)) + .layout(_r.layoutOptional(6)) + .requires_grad(_r.toBool(9)) + .pinned_memory(_r.toBool(8)); + torch::utils::maybe_initialize_cuda(options); + + auto dispatch_logspace = [](Scalar start, Scalar end, c10::optional steps, double base, TensorOptions options) -> Tensor { + pybind11::gil_scoped_release no_gil; + return torch::logspace(start, end, steps, base, options); + }; + return wrap(dispatch_logspace(_r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), _r.toDouble(3), options)); + } else { + // aten::logspace.out(Scalar start, Scalar end, int? steps=None, float base=10.0, *, Tensor(a!) out) -> Tensor(a!) + check_out_type_matches(_r.tensor(4), _r.scalartype(5), + _r.isNone(5), _r.layoutOptional(6), + _r.device(7), _r.isNone(7)); + + auto dispatch_logspace_out = [](Tensor out, Scalar start, Scalar end, c10::optional steps, double base) -> Tensor { + pybind11::gil_scoped_release no_gil; + return at::logspace_out(out, start, end, steps, base); + }; + return wrap(dispatch_logspace_out(_r.tensor(4), _r.scalar(0), _r.scalar(1), _r.toInt64Optional(2), _r.toDouble(3)).set_requires_grad(_r.toBool(9))); + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + // generated forward declarations start here ${py_forwards} @@ -496,6 +592,8 @@ static PyMethodDef torch_functions[] = { {"from_numpy", THPVariable_from_numpy, METH_STATIC | METH_O, NULL}, {"full", castPyCFunctionWithKeywords(THPVariable_full), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, {"hsmm", castPyCFunctionWithKeywords(THPVariable_hspmm), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, + {"linspace", castPyCFunctionWithKeywords(THPVariable_linspace), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, + {"logspace", castPyCFunctionWithKeywords(THPVariable_logspace), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, {"nonzero", castPyCFunctionWithKeywords(THPVariable_nonzero), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, {"randint", castPyCFunctionWithKeywords(THPVariable_randint), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, {"range", castPyCFunctionWithKeywords(THPVariable_range), METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}, diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py index 430ff6dfef819..0b68b4c5fdcbd 100644 --- a/tools/pyi/gen_pyi.py +++ b/tools/pyi/gen_pyi.py @@ -314,6 +314,10 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) - .format(FACTORY_PARAMS), 'def arange(end: Number, *, out: Optional[Tensor]=None, {}) -> Tensor: ...' .format(FACTORY_PARAMS)], + 'linspace': ['def linspace(start: Number, end: Number, steps: Optional[_int]=None, *,' + ' out: Optional[Tensor]=None, {}) -> Tensor: ...'.format(FACTORY_PARAMS)], + 'logspace': ['def logspace(start: Number, end: Number, steps: Optional[_int]=None, base: _float=10.0, *,' + ' out: Optional[Tensor]=None, {}) -> Tensor: ...'.format(FACTORY_PARAMS)], 'randint': ['def randint(low: _int, high: _int, size: _size, *,' ' generator: Optional[Generator]=None, {}) -> Tensor: ...' .format(FACTORY_PARAMS), From 1fc3576d97b5c551ea1add6e81da530070f2e03f Mon Sep 17 00:00:00 2001 From: Charles David Hernandez Date: Thu, 10 Jun 2021 12:15:31 -0700 Subject: [PATCH 013/305] Fixing and enabling tests that check fake_quant matches quant+dequant (#59095) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59095 These tests were disabled, I'm unsure as to why. I've re-enabled them and remade them to expand testing to different devices and dtypes Test Plan: python test/test_quantization.py TestFakeQuantizeOps.test_numerical_consistency Imported from OSS Reviewed By: bdhirsh Differential Revision: D29018745 fbshipit-source-id: 28188f32bafd1f1704c00ba49d09ed719dd1aeb2 --- test/quantization/core/test_workflow_ops.py | 80 +++++++++++---------- 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/test/quantization/core/test_workflow_ops.py b/test/quantization/core/test_workflow_ops.py index 21266031c7131..2b275fde5b982 100644 --- a/test/quantization/core/test_workflow_ops.py +++ b/test/quantization/core/test_workflow_ops.py @@ -408,26 +408,6 @@ def test_learnable_backward_per_tensor_cuda(self, X): self._test_learnable_backward_per_tensor( X, 'cuda', scale_base, zero_point_base) - @given(device=st.sampled_from(['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu']), - X=hu.tensor(shapes=hu.array_shapes(1, 5,), - qparams=hu.qparams(dtypes=torch.quint8))) - # https://github.com/pytorch/pytorch/issues/30604 - @unittest.skip("temporarily disable the test") - def test_numerical_consistency_per_tensor(self, device, X): - r"""Comparing numerical consistency between CPU quantize/dequantize op and the CPU fake quantize op - """ - np.random.seed(NP_RANDOM_SEED) - X, (scale, zero_point, torch_type) = X - quant_min = torch.iinfo(torch_type).min - quant_max = torch.iinfo(torch_type).max - - X = to_tensor(X, device) - # quantize_per_tensor and dequantize are only implemented in CPU - Y = torch.dequantize(torch.quantize_per_tensor(X.cpu(), scale, zero_point, torch_type)) - Y_prime = torch.fake_quantize_per_tensor_affine( - X, scale, zero_point, quant_min, quant_max) - np.testing.assert_allclose(Y, Y_prime.cpu(), rtol=tolerance, atol=tolerance) - @given(device=st.sampled_from(['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu']), X=hu.tensor(shapes=hu.array_shapes(1, 5,), qparams=hu.qparams(dtypes=[torch.quint8])), @@ -840,26 +820,48 @@ def test_learnable_backward_per_channel_cuda(self, X): self._test_learnable_backward_per_channel( X_base, 'cuda', scale_base, zero_point_base, axis) - @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 - """ - np.random.seed(NP_RANDOM_SEED) - X, (scale, zero_point, axis, torch_type) = X - quant_min = torch.iinfo(torch_type).min - quant_max = torch.iinfo(torch_type).max + def test_numerical_consistency_per_tensor(self): + self._test_numerical_consistency('per_tensor') - X = to_tensor(X, device) - scale = to_tensor(scale, device) - zero_point = torch.tensor(zero_point).to(dtype=torch.int64, device=device) - # quantize_linear and dequantize are only implemented in CPU - Y = torch.dequantize(torch.quantize_per_channel(X.cpu(), scale.cpu(), zero_point.cpu(), axis, torch_type)) - Y_prime = torch.fake_quantize_per_channel_affine( - X, scale, zero_point, axis, quant_min, quant_max) - np.testing.assert_allclose(Y, Y_prime.cpu(), rtol=tolerance, atol=tolerance) + def test_numerical_consistency_per_channel(self): + self._test_numerical_consistency('per_channel') + + def _test_numerical_consistency(self, test_type): + r"""Comparing numerical consistency between quantize/dequantize op and the fake quantize op across devices and dtypes + """ + torch.random.manual_seed(NP_RANDOM_SEED) + torch_types = [torch.qint8, torch.quint8] + float_types = [torch.float, torch.float16, torch.float64] + zero_types = [torch.long] + devices = [torch.device('cpu'), torch.device('cuda')] if torch.cuda.is_available() else [torch.device('cpu')] + axis = 1 + for i in range(20): + for torch_type, float_type, device, zero_type in itertools.product(torch_types, float_types, devices, zero_types): + X = torch.randn(3, 3, device=device).to(float_type) + scales = (10 * torch.randn(3, device=device)).abs() + scale = scales.mean().to(float).item() + zeros = (10 * torch.randn(3, device=device)).abs().to(dtype=zero_type) + zero = zeros.max().view(1).item() + quant_min = torch.iinfo(torch_type).min + quant_max = torch.iinfo(torch_type).max + + test_was_run = False + if test_type == "per_tensor": + test_was_run = True + Y = torch.dequantize(torch.quantize_per_tensor(X.to('cpu').to(torch.float), + scale, zero, torch_type)).to(device).to(float_type) + Y_prime = torch.fake_quantize_per_tensor_affine(X, scale, zero, quant_min, quant_max) + self.assertEqual( + Y, Y_prime, "Difference found between dequant+quant_per_tensor and fake_quantize_per_tensor") + + if test_type == "per_channel": + test_was_run = True + Y = torch.dequantize(torch.quantize_per_channel(X.to('cpu').to(torch.float), scales.to( + 'cpu'), zeros.to('cpu'), axis, torch_type)).to(device).to(float_type) + Y_prime = torch.fake_quantize_per_channel_affine(X, scales, zeros, axis, quant_min, quant_max) + self.assertEqual( + Y, Y_prime, "Difference found between dequant+quant_per_channel and fake_quantize_per_channel") + self.assertTrue(test_was_run) if __name__ == '__main__': raise RuntimeError("This test file is not meant to be run directly, use:\n\n" From cc32dcadd97bfd121030adddabdddc543d480a09 Mon Sep 17 00:00:00 2001 From: Yi Zhang Date: Thu, 10 Jun 2021 12:19:13 -0700 Subject: [PATCH 014/305] Fix Error when run python setup.py install again on Windows (#59689) Summary: Fix https://github.com/pytorch/pytorch/issues/59688 So far, .build.ninja should be removed before building the source code on Windows at any time Pull Request resolved: https://github.com/pytorch/pytorch/pull/59689 Reviewed By: bdhirsh Differential Revision: D29032960 Pulled By: walterddr fbshipit-source-id: 2b8162cd119820d3b6d8715745ec29b9c381e01f --- tools/setup_helpers/cmake.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py index d60dc36f13d0d..1a7c2ab0a2ac0 100644 --- a/tools/setup_helpers/cmake.py +++ b/tools/setup_helpers/cmake.py @@ -178,17 +178,19 @@ def generate( if rerun and os.path.isfile(self._cmake_cache_file): os.remove(self._cmake_cache_file) - ninja_build_file = os.path.join(self.build_dir, 'build.ninja') - if os.path.exists(self._cmake_cache_file) and not ( - USE_NINJA and not os.path.exists(ninja_build_file)): - # Everything's in place. Do not rerun. - return + ninja_deps_file = os.path.join(self.build_dir, '.ninja_deps') if IS_WINDOWS and USE_NINJA and os.path.exists(ninja_deps_file): # Cannot rerun ninja on Windows due to a ninja bug. # The workaround is to remove `.ninja_deps`. os.remove(ninja_deps_file) + ninja_build_file = os.path.join(self.build_dir, 'build.ninja') + if os.path.exists(self._cmake_cache_file) and not ( + USE_NINJA and not os.path.exists(ninja_build_file)): + # Everything's in place. Do not rerun. + return + args = [] if USE_NINJA: # Avoid conflicts in '-G' and the `CMAKE_GENERATOR` From fb620a27d08fc5ad00b386505e23e2a51f02366b Mon Sep 17 00:00:00 2001 From: albanD Date: Thu, 10 Jun 2021 12:25:58 -0700 Subject: [PATCH 015/305] [WIP] Add slow gradcheck build for the ci/slow-gradcheck label (#59020) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59020 Reviewed By: bdhirsh Differential Revision: D29036891 Pulled By: albanD fbshipit-source-id: b1f87b2cb38642097ad4079d1e818fa5997bedb4 --- .circleci/cimodel/data/pytorch_build_data.py | 13 +++++++++ .../cimodel/data/pytorch_build_definitions.py | 16 ++++++++--- .circleci/config.yml | 28 +++++++++++++++++++ .circleci/generate_config_yml.py | 11 ++++++-- .circleci/verbatim-sources/header-section.yml | 3 ++ .github/pytorch-circleci-labels.yml | 2 ++ 6 files changed, 67 insertions(+), 6 deletions(-) diff --git a/.circleci/cimodel/data/pytorch_build_data.py b/.circleci/cimodel/data/pytorch_build_data.py index 98b17f4328609..2ff3516a532e2 100644 --- a/.circleci/cimodel/data/pytorch_build_data.py +++ b/.circleci/cimodel/data/pytorch_build_data.py @@ -35,6 +35,11 @@ ("10.2", [ ("3.6", [ ("shard_test", [X(True)]), + ("slow_gradcheck", [ + (True, [ + ('shard_test', [XImportant(True)]), + ]), + ]), ("libtorch", [ (True, [ ('build_only', [X(True)]), @@ -176,10 +181,18 @@ def child_constructor(self): "cuda_gcc_override": CudaGccOverrideConfigNode, "coverage": CoverageConfigNode, "pure_torch": PureTorchConfigNode, + "slow_gradcheck": SlowGradcheckConfigNode, } return next_nodes[experimental_feature] +class SlowGradcheckConfigNode(TreeConfigNode): + def init2(self, node_name): + self.props["is_slow_gradcheck"] = True + + def child_constructor(self): + return ExperimentalFeatureConfigNode + class PureTorchConfigNode(TreeConfigNode): def modify_label(self, label): return "PURE_TORCH=" + str(label) diff --git a/.circleci/cimodel/data/pytorch_build_definitions.py b/.circleci/cimodel/data/pytorch_build_definitions.py index e87c65cc086e0..a044d716cc3d2 100644 --- a/.circleci/cimodel/data/pytorch_build_definitions.py +++ b/.circleci/cimodel/data/pytorch_build_definitions.py @@ -258,7 +258,7 @@ def gen_tree(): return configs_list -def instantiate_configs(): +def instantiate_configs(only_slow_gradcheck): config_list = [] @@ -277,8 +277,12 @@ def instantiate_configs(): is_onnx = fc.find_prop("is_onnx") or False is_pure_torch = fc.find_prop("is_pure_torch") or False is_vulkan = fc.find_prop("is_vulkan") or False + is_slow_gradcheck = fc.find_prop("is_slow_gradcheck") or False parms_list_ignored_for_docker_image = [] + if only_slow_gradcheck ^ is_slow_gradcheck: + continue + python_version = None if compiler_name == "cuda" or compiler_name == "android": python_version = fc.find_prop("pyver") @@ -342,6 +346,10 @@ def instantiate_configs(): if build_only or is_pure_torch: restrict_phases = ["build"] + if is_slow_gradcheck: + parms_list_ignored_for_docker_image.append("old") + parms_list_ignored_for_docker_image.append("gradcheck") + gpu_resource = None if cuda_version and cuda_version != "10": gpu_resource = "medium" @@ -381,7 +389,7 @@ def instantiate_configs(): tags_list=RC_PATTERN) c.dependent_tests = gen_docs_configs(c) - if cuda_version == "10.2" and python_version == "3.6" and not is_libtorch: + if cuda_version == "10.2" and python_version == "3.6" and not is_libtorch and not is_slow_gradcheck: c.dependent_tests = gen_dependent_configs(c) if ( @@ -408,9 +416,9 @@ def instantiate_configs(): return config_list -def get_workflow_jobs(): +def get_workflow_jobs(only_slow_gradcheck=False): - config_list = instantiate_configs() + config_list = instantiate_configs(only_slow_gradcheck) x = [] for conf_options in config_list: diff --git a/.circleci/config.yml b/.circleci/config.yml index 273ea971706ed..72bcc82787474 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -17,6 +17,9 @@ parameters: run_master_build: type: boolean default: false + run_slow_gradcheck_build: + type: boolean + default: false executors: windows-with-nvidia-gpu: @@ -9349,6 +9352,31 @@ workflows: vc_version: "" vc_year: "2019" when: << pipeline.parameters.run_master_build >> + slow_gradcheck_build: + jobs: + - pytorch_linux_build: + name: pytorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_old_gradcheck_build + requires: + - "docker-pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-build" + docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + - pytorch_linux_test: + name: pytorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_old_gradcheck_test1 + requires: + - pytorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_old_gradcheck_build + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test1" + docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + use_cuda_docker_runtime: "1" + resource_class: gpu.medium + - pytorch_linux_test: + name: pytorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_old_gradcheck_test2 + requires: + - pytorch_linux_xenial_cuda10_2_cudnn7_py3_gcc7_old_gradcheck_build + build_environment: "pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7-old-gradcheck-test2" + docker_image: "308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7" + use_cuda_docker_runtime: "1" + resource_class: gpu.medium + when: << pipeline.parameters.run_slow_gradcheck_build >> scheduled-ci: triggers: - schedule: diff --git a/.circleci/generate_config_yml.py b/.circleci/generate_config_yml.py index b51c29180f512..69634e960e4a2 100755 --- a/.circleci/generate_config_yml.py +++ b/.circleci/generate_config_yml.py @@ -145,14 +145,17 @@ def gen_build_workflows_tree(): binary_build_definitions.get_post_upload_jobs, binary_build_definitions.get_binary_smoke_test_jobs, ] + build_jobs = [f() for f in build_workflows_functions] + master_build_jobs = filter_master_only_jobs(build_jobs) binary_build_functions = [ binary_build_definitions.get_binary_build_jobs, binary_build_definitions.get_nightly_tests, binary_build_definitions.get_nightly_uploads, ] - build_jobs = [f() for f in build_workflows_functions] - master_build_jobs = filter_master_only_jobs(build_jobs) + + slow_gradcheck_jobs = pytorch_build_definitions.get_workflow_jobs(only_slow_gradcheck=True) + return { "workflows": { "binary_builds": { @@ -167,6 +170,10 @@ def gen_build_workflows_tree(): "when": r"<< pipeline.parameters.run_master_build >>", "jobs": master_build_jobs, }, + "slow_gradcheck_build": { + "when": r"<< pipeline.parameters.run_slow_gradcheck_build >>", + "jobs": slow_gradcheck_jobs, + }, } } diff --git a/.circleci/verbatim-sources/header-section.yml b/.circleci/verbatim-sources/header-section.yml index 3cac363731ffe..527340d542907 100644 --- a/.circleci/verbatim-sources/header-section.yml +++ b/.circleci/verbatim-sources/header-section.yml @@ -17,6 +17,9 @@ parameters: run_master_build: type: boolean default: false + run_slow_gradcheck_build: + type: boolean + default: false executors: windows-with-nvidia-gpu: diff --git a/.github/pytorch-circleci-labels.yml b/.github/pytorch-circleci-labels.yml index 7f98bd686dbd9..0c030a6154057 100644 --- a/.github/pytorch-circleci-labels.yml +++ b/.github/pytorch-circleci-labels.yml @@ -22,3 +22,5 @@ labels_to_circle_params: - run_build ci/master: parameter: run_master_build + ci/slow-gradcheck: + parameter: run_slow_gradcheck_build From 0099c25b854938343a0eb4f2210b4be839f1c33c Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 10 Jun 2021 12:56:04 -0700 Subject: [PATCH 016/305] fx quant: remove some dead code in observer insertion (redo) (#59799) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59799 This is a redo of #58574, easier to create a new PR than to fix rebase conflicts, as there have been a large number of refactors to the underlying code. Removes some code which was incorrectly added by #57519 but never actually used for anything. Test Plan: ``` python test/test_quantization.py TestQuantizeFx python test/test_quantization.py TestQuantizeFxOps ``` Imported from OSS Reviewed By: jerryzh168 Differential Revision: D29031955 fbshipit-source-id: f407d181070cb283382965952821e3647c705544 --- torch/quantization/fx/prepare.py | 83 +++++++++++++------------------- 1 file changed, 34 insertions(+), 49 deletions(-) diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index 5245d0ae1079d..8cc3c268f5928 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -432,58 +432,43 @@ def maybe_insert_output_observer_for_node( root_node, matched_nodes, pattern, qhandler, qconfig = matches.get( node.name, (None, None, None, None, None)) - if qhandler is not None: - assert qconfig is not None + if qhandler is None: + return None - is_standalone_module = qhandler is not None and \ - isinstance(qhandler, StandaloneModuleQuantizeHandler) - - should_insert_observer = \ - qhandler.should_insert_observer_for_output( - qconfig, model.training) - # TODO(future PR): move the following logic to - # should_insert_observer_for_output - should_insert_observer = should_insert_observer and \ - activation_is_statically_quantized(qconfig) - - # we never insert observers to output of standalone module, we assume - # if needed, they are inserted inside the standalone module - should_insert_observer = should_insert_observer and \ - (not is_standalone_module) - - if should_insert_observer: - act_post_process_ctr = qconfig.activation - if activation_is_int8_quantized(qconfig): - act_post_process_ctr = \ - get_default_output_activation_post_process_map().get( - matched_pattern, - act_post_process_ctr) - observer = act_post_process_ctr() - new_obs = insert_observer(node, observer, model, modules, graph) - # set the type, so the next node can read it - node_name_to_target_dtype[new_obs.name] = \ - node_name_to_target_dtype[node.name] - return new_obs + assert qconfig is not None + assert node.op != 'output', 'observer insertion for outputs is handled elsewhere' - elif node.op == 'output': - prev_node = node.args[0] - assert isinstance(prev_node, Node) - prev_node_dtype = node_name_to_target_dtype[prev_node.name] - node_dtype = node_name_to_target_dtype[node.name] - should_insert_observer = ( - prev_node_dtype == torch.float and - node_dtype != torch.float - ) - if should_insert_observer: - assert qconfig is not None - observer = qconfig.activation() - new_obs = insert_observer( - prev_node, observer, model, modules, graph) - # set the type, so the next node can read it - node_name_to_target_dtype[new_obs.name] = node_dtype - return new_obs + is_standalone_module = qhandler is not None and \ + isinstance(qhandler, StandaloneModuleQuantizeHandler) - return None + should_insert_observer = \ + qhandler.should_insert_observer_for_output( + qconfig, model.training) + # TODO(future PR): move the following logic to + # should_insert_observer_for_output + should_insert_observer = should_insert_observer and \ + activation_is_statically_quantized(qconfig) + + # we never insert observers to output of standalone module, we assume + # if needed, they are inserted inside the standalone module + should_insert_observer = should_insert_observer and \ + (not is_standalone_module) + + if should_insert_observer: + act_post_process_ctr = qconfig.activation + if activation_is_int8_quantized(qconfig): + act_post_process_ctr = \ + get_default_output_activation_post_process_map().get( + matched_pattern, + act_post_process_ctr) + observer = act_post_process_ctr() + new_obs = insert_observer(node, observer, model, modules, graph) + # set the type, so the next node can read it + node_name_to_target_dtype[new_obs.name] = \ + node_name_to_target_dtype[node.name] + return new_obs + else: + return None def maybe_insert_observers_before_graph_output( graph_output_node: Node, From 54cc477ea393ebbfd0adedcac7991a1cf8a666ba Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Thu, 10 Jun 2021 13:34:02 -0700 Subject: [PATCH 017/305] .github: Ensure cleaner windows workspace (#59742) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59742 It looks like Windows workers were failing out due to some leftovers from previous builds, this should hopefully remedy some of those errors Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: janeyx99 Differential Revision: D29009076 Pulled By: seemethere fbshipit-source-id: 426d54df14ec580cb24b818c48e2f4bd36159181 --- .circleci/scripts/windows_cudnn_install.sh | 12 ++++++++---- .github/templates/windows_ci_workflow.yml.j2 | 8 ++++++++ .github/workflows/pytorch-win-vs2019-cpu-py3.yml | 8 ++++++++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/.circleci/scripts/windows_cudnn_install.sh b/.circleci/scripts/windows_cudnn_install.sh index 15b382ca412ca..a389fc4f284e9 100644 --- a/.circleci/scripts/windows_cudnn_install.sh +++ b/.circleci/scripts/windows_cudnn_install.sh @@ -20,9 +20,13 @@ else fi cudnn_installer_link="https://ossci-windows.s3.amazonaws.com/${cudnn_installer_name}.zip" +cudnn_install_folder="C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${CUDA_VERSION}/" -curl --retry 3 -O $cudnn_installer_link -7z x ${cudnn_installer_name}.zip -ocudnn -cp -r cudnn/cuda/* "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${CUDA_VERSION}/" +curl --retry 3 -O "$cudnn_installer_link" +7z x "${cudnn_installer_name}.zip" -ocudnn +# shellcheck recommends to use '${var:?}/*' to avoid potentially expanding to '/*' +# Remove all of the directories before attempting to copy files +rm -rf "${cudnn_install_folder:?}/*" +cp -rf cudnn/cuda/* "${cudnn_install_folder}" rm -rf cudnn -rm -f ${cudnn_installer_name}.zip +rm -f "${cudnn_installer_name}.zip" diff --git a/.github/templates/windows_ci_workflow.yml.j2 b/.github/templates/windows_ci_workflow.yml.j2 index 530f430e326c4..594423fee3604 100644 --- a/.github/templates/windows_ci_workflow.yml.j2 +++ b/.github/templates/windows_ci_workflow.yml.j2 @@ -38,6 +38,10 @@ jobs: uses: actions/checkout@v2 with: submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf - name: Install Visual Studio 2019 toolchain shell: powershell run: | @@ -77,6 +81,10 @@ jobs: uses: actions/checkout@v2 with: submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf - name: Install Visual Studio 2019 toolchain shell: powershell run: | diff --git a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml index 659cc1e5f8904..f96181f4f8219 100644 --- a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml @@ -37,6 +37,10 @@ jobs: uses: actions/checkout@v2 with: submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf - name: Install Visual Studio 2019 toolchain shell: powershell run: | @@ -76,6 +80,10 @@ jobs: uses: actions/checkout@v2 with: submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf - name: Install Visual Studio 2019 toolchain shell: powershell run: | From e2c784d940983a1581e449356affc9d7de5bdf56 Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Thu, 10 Jun 2021 13:34:02 -0700 Subject: [PATCH 018/305] [reland] .github: Add Windows GPU workflow (#58782) (#59752) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59752 Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29009775 Pulled By: seemethere fbshipit-source-id: 5be1b818b5653a4fdbfe4a79731317068dc1a5d1 --- .github/scripts/generate_ci_workflows.py | 8 + .github/templates/windows_ci_workflow.yml.j2 | 27 ++- .../workflows/pytorch-win-vs2019-cpu-py3.yml | 1 - .../pytorch-win-vs2019-cuda10-cudnn7-py3.yml | 188 ++++++++++++++++++ .jenkins/pytorch/win-test.sh | 6 +- 5 files changed, 224 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index 239b82cd743a1..d27f593d67137 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -23,11 +23,13 @@ def PyTorchWindowsWorkflow( *, build_environment: str, test_runner_type: str, + cuda_version: str, on_pull_request: bool = False ) -> PyTorchWorkflow: return { "build_environment": build_environment, "test_runner_type": test_runner_type, + "cuda_version": cuda_version, "on_pull_request": on_pull_request, } @@ -70,8 +72,14 @@ def generate_workflow_file( WINDOWS_WORKFLOWS = [ PyTorchWindowsWorkflow( build_environment="pytorch-win-vs2019-cpu-py3", + cuda_version="cpu", test_runner_type=WINDOWS_CPU_TEST_RUNNER, on_pull_request=True + ), + PyTorchWindowsWorkflow( + build_environment="pytorch-win-vs2019-cuda10-cudnn7-py3", + cuda_version="10.1", + test_runner_type=WINDOWS_CUDA_TEST_RUNNER, ) ] diff --git a/.github/templates/windows_ci_workflow.yml.j2 b/.github/templates/windows_ci_workflow.yml.j2 index 594423fee3604..2d792b1de5fde 100644 --- a/.github/templates/windows_ci_workflow.yml.j2 +++ b/.github/templates/windows_ci_workflow.yml.j2 @@ -15,16 +15,19 @@ on: env: BUILD_ENVIRONMENT: !{{ build_environment }} BUILD_WHEEL: 1 - CUDA_VERSION: "cpu" + CUDA_VERSION: "!{{ cuda_version }}" IN_CI: 1 INSTALL_WINDOWS_SDK: 1 JOB_BASE_NAME: test PYTHON_VERSION: "3.6" SCCACHE_BUCKET: "ossci-compiler-cache" - TORCH_CUDA_ARCH_LIST: "5.2;7.5" VC_PRODUCT: "BuildTools" VC_VERSION: "" VC_YEAR: "2019" +{%- if cuda_version != "cpu" %} + TORCH_CUDA_ARCH_LIST: "7.0" + USE_CUDA: 1 +{%- endif %} concurrency: group: !{{ build_environment }}-${{ github.event.pull_request.number || github.sha }} @@ -46,6 +49,16 @@ jobs: shell: powershell run: | .\.circleci\scripts\vs_install.ps1 +{%- if cuda_version != "cpu" %} + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh +{%- endif %} - name: Build shell: bash run: | @@ -89,6 +102,16 @@ jobs: shell: powershell run: | .\.circleci\scripts\vs_install.ps1 +{%- if cuda_version != "cpu" %} + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh +{%- endif %} - uses: seemethere/download-artifact-s3@0504774707cbc8603d7dca922e8026eb8bf3b47b name: Download PyTorch Build Artifacts with: diff --git a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml index f96181f4f8219..8cfeea07e597c 100644 --- a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml @@ -20,7 +20,6 @@ env: JOB_BASE_NAME: test PYTHON_VERSION: "3.6" SCCACHE_BUCKET: "ossci-compiler-cache" - TORCH_CUDA_ARCH_LIST: "5.2;7.5" VC_PRODUCT: "BuildTools" VC_VERSION: "" VC_YEAR: "2019" diff --git a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml new file mode 100644 index 0000000000000..9d3b9ed3aff53 --- /dev/null +++ b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml @@ -0,0 +1,188 @@ +# @generated DO NOT EDIT MANUALLY +# Template is at: .github/templates/windows_ci_workflow.yml.j2 +# Generation script: .github/scripts/generate_ci_workflows.py +name: Windows CI (pytorch-win-vs2019-cuda10-cudnn7-py3) + +on: + push: + branches: + - master + - release/* + workflow_dispatch: + +env: + BUILD_ENVIRONMENT: pytorch-win-vs2019-cuda10-cudnn7-py3 + BUILD_WHEEL: 1 + CUDA_VERSION: "10.1" + IN_CI: 1 + INSTALL_WINDOWS_SDK: 1 + JOB_BASE_NAME: test + PYTHON_VERSION: "3.6" + SCCACHE_BUCKET: "ossci-compiler-cache" + VC_PRODUCT: "BuildTools" + VC_VERSION: "" + VC_YEAR: "2019" + TORCH_CUDA_ARCH_LIST: "7.0" + USE_CUDA: 1 + +concurrency: + group: pytorch-win-vs2019-cuda10-cudnn7-py3-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: true + +jobs: + build: + runs-on: "windows.4xlarge" + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf + - name: Install Visual Studio 2019 toolchain + shell: powershell + run: | + .\.circleci\scripts\vs_install.ps1 + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh + - name: Build + shell: bash + run: | + .jenkins/pytorch/win-build.sh + # Upload to github so that people can click and download artifacts + - name: Upload artifacts to Github + if: always() + uses: actions/upload-artifact@v2 + # Don't fail on upload to GH since it's only for user convenience + continue-on-error: true + with: + retention-days: 14 + if-no-files-found: error + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\w\build-results + - name: Upload artifacts to s3 + if: always() + uses: seemethere/upload-artifact-s3@9d7ceb0ab39c2c88d93ef7792b27425b27d59162 + with: + retention-days: 14 + if-no-files-found: error + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\w\build-results + + test: + runs-on: windows.8xlarge.nvidia.gpu + env: + JOB_BASE_NAME: pytorch-win-vs2019-cuda10-cudnn7-py3-test + needs: + - build + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf + - name: Install Visual Studio 2019 toolchain + shell: powershell + run: | + .\.circleci\scripts\vs_install.ps1 + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh + - uses: seemethere/download-artifact-s3@0504774707cbc8603d7dca922e8026eb8bf3b47b + name: Download PyTorch Build Artifacts + with: + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\${{ github.run_id }}\build-results + - name: Check build-results folder + shell: powershell + run: | + tree /F C:\$Env:GITHUB_RUN_ID\build-results + # Needed for coverage in win-test.sh + - uses: actions/setup-python@v2 + name: Setup Python3 + with: + python-version: '3.x' + - name: Run test scripts + shell: bash + env: + PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ + run: | + .jenkins/pytorch/win-test.sh + - uses: actions/upload-artifact@v2 + name: Store PyTorch Test Reports + if: always() + with: + name: test-reports + retention-days: 14 + if-no-files-found: error + path: + test/**/*.xml + + # this is a separate step from test because the log files from test are too + # long: basically, GitHub tries to render all of the log files when you click + # through an action causing extreme slowdown on actions that contain too many + # logs (like test); we can always move it back to the other one, but it + # doesn't create the best experience + render_test_results: + if: always() + needs: + - test + runs-on: ubuntu-18.04 + # TODO: Make this into a composite step + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + # deep clone, to allow tools/print_test_stats.py to use Git commands + fetch-depth: 0 + - uses: actions/download-artifact@v2 + name: Download PyTorch Test Reports + with: + name: test-reports + path: test/test-reports + - uses: actions/setup-python@v2 + with: + python-version: 3.9 + - name: Install dependencies + # boto3 version copied from .circleci/docker/common/install_conda.sh + run: | + pip install -r requirements.txt + pip install boto3==1.16.34 junitparser rich + - name: Output Test Results (Click Me) + run: | + python tools/render_junit.py test + - name: Parse ref + id: parse-ref + run: .github/scripts/parse_ref.py + - name: Display and upload test statistics (Click Me) + # temporary hack: set CIRCLE_* vars, until we update + # tools/print_test_stats.py to natively support GitHub Actions + env: + SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_OSSCI_METRICS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_OSSCI_METRICS_SECRET_ACCESS_KEY }} + CIRCLE_BRANCH: ${{ steps.parse-ref.outputs.branch }} + CIRCLE_JOB: pytorch-win-vs2019-cuda10-cudnn7-py3 + CIRCLE_PR_NUMBER: ${{ github.event.pull_request.number }} + CIRCLE_SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + CIRCLE_TAG: ${{ steps.parse-ref.outputs.tag }} + CIRCLE_WORKFLOW_ID: ${{ github.run_id }} # dunno if this corresponds + run: | + export PYTHONPATH=$PWD + python tools/print_test_stats.py --upload-to-s3 --compare-with-s3 test diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index 6dabab6cf5c7d..aa6f89e164778 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -55,9 +55,9 @@ fi run_tests() { # Run nvidia-smi if available - for path in /c/Program Files/NVIDIA Corporation/NVSMI/nvidia-smi.exe /c/Windows/System32/nvidia-smi.exe; do - if [ -x $path ]; then - $path; + for path in '/c/Program Files/NVIDIA Corporation/NVSMI/nvidia-smi.exe' /c/Windows/System32/nvidia-smi.exe; do + if [[ -x "$path" ]]; then + "$path"; break fi done From e7ad82eb2f9856682d95a69114a04655a32c5ad2 Mon Sep 17 00:00:00 2001 From: Erjia Guan Date: Thu, 10 Jun 2021 14:03:19 -0700 Subject: [PATCH 019/305] [DataLoader] Add option to refine type during runtime validation for DP instance (#56066) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56066 Test Plan: Imported from OSS Reviewed By: VitalyFedyunin Differential Revision: D27776646 Pulled By: ejguan fbshipit-source-id: 695ff7775177653d809c5917d938c706281e1298 --- test/test_datapipe.py | 52 +++++- torch/utils/data/_decorator.py | 6 +- torch/utils/data/_typing.py | 22 ++- torch/utils/data/typing.ipynb | 288 ++++++++++++++++++++++++++++++--- 4 files changed, 338 insertions(+), 30 deletions(-) diff --git a/test/test_datapipe.py b/test/test_datapipe.py index ac5006e116ce6..afdfc41a4f07f 100644 --- a/test/test_datapipe.py +++ b/test/test_datapipe.py @@ -385,7 +385,10 @@ def fn(item, dtype=torch.float, *, sum=False): data = torch.tensor(item, dtype=dtype) return data if not sum else data.sum() - map_dp = input_dp.map(lambda ls: ls * 2, nesting_level=0) + with warnings.catch_warnings(record=True) as wa: + map_dp = input_dp.map(lambda ls: ls * 2, nesting_level=0) + self.assertEqual(len(wa), 1) + self.assertRegex(str(wa[0].message), r"^Lambda function is not supported for pickle") self.assertEqual(len(input_dp), len(map_dp)) for x, y in zip(map_dp, input_dp): self.assertEqual(x, y * 2) @@ -406,11 +409,7 @@ def fn(item, dtype=torch.float, *, sum=False): map_dp = input_dp.map(fn, nesting_level=4) with self.assertRaises(IndexError): - for x, y in zip(map_dp, input_dp): - self.assertEqual(len(x), len(y)) - for a, b in zip(x, y): - print(a, b) - self.assertEqual(a, torch.tensor(b, dtype=torch.float)) + list(map_dp) with self.assertRaises(ValueError): input_dp.map(fn, nesting_level=-2) @@ -1052,7 +1051,7 @@ def __iter__(self) -> Iterator[Tuple[int, T_co]]: [1, '1', 2, '2']) for ds in dss: dp = DP(ds) - with self.assertRaisesRegex(RuntimeError, r"Expected an instance of subtype"): + with self.assertRaisesRegex(RuntimeError, r"Expected an instance as subtype"): list(dp) with runtime_validation_disabled(): @@ -1060,9 +1059,46 @@ def __iter__(self) -> Iterator[Tuple[int, T_co]]: with runtime_validation_disabled(): self.assertEqual(list(dp), ds) - with self.assertRaisesRegex(RuntimeError, r"Expected an instance of subtype"): + with self.assertRaisesRegex(RuntimeError, r"Expected an instance as subtype"): list(dp) + def test_reinforce(self): + T = TypeVar('T', int, str) + + + class DP(IterDataPipe[T]): + def __init__(self, ds): + self.ds = ds + + @runtime_validation + def __iter__(self) -> Iterator[T]: + for d in self.ds: + yield d + + ds = list(range(10)) + # Valid type reinforcement + dp = DP(ds).reinforce_type(int) + self.assertTrue(dp.type, int) + self.assertEqual(list(dp), ds) + + # Invalid type + with self.assertRaisesRegex(TypeError, r"'expected_type' must be a type"): + dp = DP(ds).reinforce_type(1) + + # Type is not subtype + with self.assertRaisesRegex(TypeError, r"Expected 'expected_type' as subtype of"): + dp = DP(ds).reinforce_type(float) + + # Invalid data at runtime + dp = DP(ds).reinforce_type(str) + with self.assertRaisesRegex(RuntimeError, r"Expected an instance as subtype"): + list(dp) + + # Context Manager to disable the runtime validation + with runtime_validation_disabled(): + self.assertEqual(list(d for d in dp), ds) + + if __name__ == '__main__': run_tests() diff --git a/torch/utils/data/_decorator.py b/torch/utils/data/_decorator.py index 1f9143820cd7a..6e64fd5b7fef4 100644 --- a/torch/utils/data/_decorator.py +++ b/torch/utils/data/_decorator.py @@ -109,7 +109,7 @@ def deterministic_wrapper_fn(self, *args, **kwargs) -> IterDataPipe: ###################################################### -# typing +# Type validation ###################################################### # Validate each argument of DataPipe with hint as a subtype of the hint. def argument_validation(f): @@ -173,8 +173,8 @@ def wrapper(self): it = f(self) for d in it: if not self.type.issubtype_of_instance(d): - raise RuntimeError("Expected an instance of subtype {}, but found {}" - .format(self.type, d)) + raise RuntimeError("Expected an instance as subtype of {}, but found {}({})" + .format(self.type, d, type(d))) yield d return wrapper diff --git a/torch/utils/data/_typing.py b/torch/utils/data/_typing.py index 7f8f51eb22e7b..1b2c581b208d7 100644 --- a/torch/utils/data/_typing.py +++ b/torch/utils/data/_typing.py @@ -339,6 +339,9 @@ def __hash__(self): def _dp_init_subclass(sub_cls, *args, **kwargs): + # Add function for datapipe instance to reinforce the type + sub_cls.reinforce_type = reinforce_type + # TODO: # - add global switch for type checking at compile-time @@ -371,5 +374,22 @@ def _dp_init_subclass(sub_cls, *args, **kwargs): ", but found {}".format(sub_cls.__name__, _type_repr(hints['return']))) data_type = return_hint.__args__[0] if not issubtype(data_type, sub_cls.type.param): - raise TypeError("Expected return type of '__iter__' is a subtype of {}, but found {}" + raise TypeError("Expected return type of '__iter__' as a subtype of {}, but found {}" " for {}".format(sub_cls.type, _type_repr(data_type), sub_cls.__name__)) + +def reinforce_type(self, expected_type): + r""" + Reinforce the type for DataPipe instance. And the 'expected_type' is required + to be a subtype of the original type hint to restrict the type requirement + of DataPipe instance. + """ + if isinstance(expected_type, tuple): + expected_type = Tuple[expected_type] + _type_check(expected_type, msg="'expected_type' must be a type") + + if not issubtype(expected_type, self.type.param): + raise TypeError("Expected 'expected_type' as subtype of {}, but found {}" + .format(self.type, _type_repr(expected_type))) + + self.type = _DataPipeType(expected_type) + return self diff --git a/torch/utils/data/typing.ipynb b/torch/utils/data/typing.ipynb index 30fea38512cc3..17c0b78b060b1 100644 --- a/torch/utils/data/typing.ipynb +++ b/torch/utils/data/typing.ipynb @@ -30,7 +30,23 @@ "# Hide traceback of Error\n", "import functools\n", "ipython = get_ipython()\n", - "ipython.showtraceback = functools.partial(ipython.showtraceback, exception_only=True)" + "def showtraceback(self, exc_tuple=None, filename=None, tb_offset=None,\n", + " exception_only=False, running_compiled_code=False):\n", + " try:\n", + " try:\n", + " etype, value, tb = self._get_exc_info(exc_tuple)\n", + " except ValueError:\n", + " print('No traceback available to show.', file=sys.stderr)\n", + " return\n", + "\n", + " # Hide traceback\n", + " stb = self.InteractiveTB.get_exception_only(etype, value)\n", + "\n", + " self._showtraceback(etype, value, stb)\n", + "\n", + " except KeyboardInterrupt:\n", + " print('\\n' + self.get_exception_only(), file=sys.stderr)\n", + "ipython.showtraceback = functools.partial(showtraceback, ipython)" ] }, { @@ -54,7 +70,6 @@ "evalue": "Expected 'Iterator' as the return annotation for `__iter__` of InvalidDP1, but found str", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected 'Iterator' as the return annotation for `__iter__` of InvalidDP1, but found str\n" ] } @@ -79,11 +94,10 @@ "outputs": [ { "ename": "TypeError", - "evalue": "Expected return type of '__iter__' is a subtype of int, but found str for InvalidDP2", + "evalue": "Expected return type of '__iter__' as a subtype of int, but found str for InvalidDP2", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", - "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected return type of '__iter__' is a subtype of int, but found str for InvalidDP2\n" + "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected return type of '__iter__' as a subtype of int, but found str for InvalidDP2\n" ] } ], @@ -316,7 +330,6 @@ "evalue": "Expected argument 'dp' as a IterDataPipe, but found ", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected argument 'dp' as a IterDataPipe, but found \n" ] } @@ -342,7 +355,6 @@ "evalue": "Expected type of argument 'dp' as a subtype of hint typing.Union[int, tuple], but found str", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected type of argument 'dp' as a subtype of hint typing.Union[int, tuple], but found str\n" ] } @@ -379,6 +391,8 @@ "source": [ "## Runtime\n", "\n", + "\n", + "### Decorator\n", "Runtime type checking is enabled by a decorator `runtime_validation`. Users can opt in by attaching the decorator to `__iter__` to check the output data is an instance of subtype of `type` attribute of the DataPipe.\n", "\n", "Note: This decorator is only allowed to be attached to `__iter__` for now. It can be extended into `__getitem__` and further `nonblocking` functions.\n", @@ -428,11 +442,10 @@ }, { "ename": "RuntimeError", - "evalue": "Expected an instance of subtype typing.Tuple[int, +T_co], but found ('3', 3)", + "evalue": "Expected an instance as subtype of typing.Tuple[int, +T_co], but found ('3', 3)()", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", - "\u001b[0;31mRuntimeError\u001b[0m\u001b[0;31m:\u001b[0m Expected an instance of subtype typing.Tuple[int, +T_co], but found ('3', 3)\n" + "\u001b[0;31mRuntimeError\u001b[0m\u001b[0;31m:\u001b[0m Expected an instance as subtype of typing.Tuple[int, +T_co], but found ('3', 3)()\n" ] } ], @@ -489,11 +502,10 @@ }, { "ename": "RuntimeError", - "evalue": "Expected an instance of subtype typing.Tuple[int, +T_co], but found [3, 3]", + "evalue": "Expected an instance as subtype of typing.Tuple[int, +T_co], but found [3, 3]()", "output_type": "error", "traceback": [ - "An exception has occurred, use %tb to see the full traceback.\n", - "\u001b[0;31mRuntimeError\u001b[0m\u001b[0;31m:\u001b[0m Expected an instance of subtype typing.Tuple[int, +T_co], but found [3, 3]\n" + "\u001b[0;31mRuntimeError\u001b[0m\u001b[0;31m:\u001b[0m Expected an instance as subtype of typing.Tuple[int, +T_co], but found [3, 3]()\n" ] } ], @@ -544,16 +556,256 @@ "name": "stdout", "output_type": "stream", "text": [ - "(1, 1)\n", - "(2, '2')\n", - "(3, 3.0)\n" + "[(1, 1), (2, '2'), (3, 3.0)]\n" ] } ], "source": [ "dp = DP([(1, 1), (2, '2'), (3, 3.)])\n", - "for d in dp:\n", - " print(d)" + "print(list(dp))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Reinforce type for DataPipe instance" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [], + "source": [ + "T = TypeVar('T', int, str)\n", + "ds = list(range(10))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- If the DataPipe class is not decorated with `runtime_validation` and the DataPipe instance calls `reinforce_type`, a warning will be raised." + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/erjia/workspace/pytorch-dev/typing/torch/utils/data/_typing.py:346: UserWarning: The type of data generated from `DataPipe` instance won't be validated at runtime. Decorator of `runtime_validation` is required to be attached to `__iter__` method of for runtime type validation\n", + " warnings.warn(\"The type of data generated from `DataPipe` instance won't be validated \"\n" + ] + } + ], + "source": [ + "class DP(IterDataPipe[T]):\n", + " def __init__(self, ds):\n", + " self.ds = ds\n", + " \n", + " def __iter__(self):\n", + " for d in self.ds:\n", + " yield d\n", + "dp = DP(ds).reinforce_type(int)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [], + "source": [ + "class DP(IterDataPipe[T]):\n", + " def __init__(self, ds):\n", + " self.ds = ds\n", + " \n", + " @runtime_validation\n", + " def __iter__(self):\n", + " for d in self.ds:\n", + " yield d" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- expected type must be a subtype of the original type hint" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "ename": "TypeError", + "evalue": "Expected 'expected_type' as a subtype of ~T, but found float", + "output_type": "error", + "traceback": [ + "\u001b[0;31mTypeError\u001b[0m\u001b[0;31m:\u001b[0m Expected 'expected_type' as a subtype of ~T, but found float\n" + ] + } + ], + "source": [ + "dp = DP(ds).reinforce_type(float)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- Integer data is not subtype of str" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": {}, + "outputs": [ + { + "ename": "RuntimeError", + "evalue": "Expected an instance as subtype of str, but found 0()", + "output_type": "error", + "traceback": [ + "\u001b[0;31mRuntimeError\u001b[0m\u001b[0;31m:\u001b[0m Expected an instance as subtype of str, but found 0()\n" + ] + } + ], + "source": [ + "dp = DP(ds).reinforce_type(str)\n", + "list(dp)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- Compatible with context mangager to disable validation" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n" + ] + } + ], + "source": [ + "with runtime_validation_disabled():\n", + " print(list(dp))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- Valid type enforcement" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n" + ] + } + ], + "source": [ + "dp = DP(ds).reinforce_type(int)\n", + "print(list(dp))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "- Different type based on the logic of class initialization" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [], + "source": [ + "class DP(IterDataPipe[Union[int, str]]):\n", + " def __init__(self, label):\n", + " if label == 'int':\n", + " self.reinforce_type(int)\n", + " elif label == 'str':\n", + " self.reinforce_type(str)" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": { + "scrolled": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "int\n" + ] + } + ], + "source": [ + "dp = DP('int')\n", + "print(dp.type)" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "str\n" + ] + } + ], + "source": [ + "dp = DP('str')\n", + "print(dp.type)" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "typing.Union[int, str]\n" + ] + } + ], + "source": [ + "dp = DP('')\n", + "print(dp.type)" ] } ], From e71db0bb82b1cafb01da6541791c8907bb468a5d Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Thu, 10 Jun 2021 14:14:55 -0700 Subject: [PATCH 020/305] .jenkins: Ignore exit code of nvidia-smi (#59826) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59826 It's only informational and will run on Windows CPU executors as well Fixes issues found in https://github.com/pytorch/pytorch/runs/2797531966 Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: janeyx99 Differential Revision: D29042951 Pulled By: seemethere fbshipit-source-id: 862094e53417c0a59d7728bf680be37b806b5a6f --- .jenkins/pytorch/win-test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index aa6f89e164778..96fe2e6225be4 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -57,7 +57,7 @@ run_tests() { # Run nvidia-smi if available for path in '/c/Program Files/NVIDIA Corporation/NVSMI/nvidia-smi.exe' /c/Windows/System32/nvidia-smi.exe; do if [[ -x "$path" ]]; then - "$path"; + "$path" || echo "true"; break fi done From ac6b5beade03d72c549466a9c195dd39497ef3a2 Mon Sep 17 00:00:00 2001 From: Serhat Yilmaz Date: Thu, 10 Jun 2021 14:19:54 -0700 Subject: [PATCH 021/305] [torch][segment_reduce] Add support for mean reduction (cpu) (#59521) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59521 This diff is adding support for mean reduction for CPU (fwd + bckwd). Will add cuda implementation in subsequent PR. We are using "cub::DeviceSegmentedReduce" for other aggregation, trying to see how to support mean or will write custom kernel for it. Next Steps: - cuda support for mean - 2d data input support - more testing - benchmarking Test Plan: updated unit test. Still relying on manual data for ease of debugging. Will add more tests that covers edge cases once major features are complete. Reviewed By: ngimel Differential Revision: D28922547 fbshipit-source-id: 2fad53bbad2cce714808ff95759cbdbd45bb4ce6 --- aten/src/ATen/native/SegmentReduce.cpp | 109 ++++++++++++------ aten/src/ATen/native/SegmentReduce.h | 11 +- aten/src/ATen/native/cuda/SegmentReduce.cu | 1 + aten/src/ATen/native/native_functions.yaml | 4 +- .../check_backward_compatibility.py | 3 + test/test_segment_reductions.py | 44 ++++--- tools/autograd/derivatives.yaml | 2 +- 7 files changed, 120 insertions(+), 54 deletions(-) diff --git a/aten/src/ATen/native/SegmentReduce.cpp b/aten/src/ATen/native/SegmentReduce.cpp index b1bb8a0cb040c..c17cb67b5b60e 100644 --- a/aten/src/ATen/native/SegmentReduce.cpp +++ b/aten/src/ATen/native/SegmentReduce.cpp @@ -13,7 +13,18 @@ DEFINE_DISPATCH(_segment_reduce_backward_stub); namespace { +SegmentReductionType get_reduction_enum(const c10::string_view& reduce) { + if (reduce == "max") { + return SegmentReductionType::MAX; + } else if (reduce == "mean") { + return SegmentReductionType::MEAN; + } else { + TORCH_CHECK(false, "unsopported reduction given! ", reduce); + } +} + Tensor _segment_reduce_cpu_kernel( + SegmentReductionType reduction, const Tensor& data, const Tensor& lengths, int64_t axis, @@ -29,20 +40,41 @@ Tensor _segment_reduce_cpu_kernel( const auto* values_data = data.data_ptr(); int64_t k = 0; for (int64_t i = 0; i < batch_size; ++i) { - scalar_t initial_value = initial.has_value() - ? initial.value().to() - : std::numeric_limits::lowest(); + // ===== step1: initialize starting value + scalar_t initial_value; + if (initial.has_value()) { + initial_value = initial.value().to(); + } else if (reduction == SegmentReductionType::MAX) { + initial_value = std::numeric_limits::lowest(); + } else if (reduction == SegmentReductionType::MEAN) { + initial_value = 0; + } + + // ===== step2: apply reduction for (int64_t j = 0; j < lengths_data[i]; ++j) { const auto data = values_data[k]; - initial_value = at::_isnan(data) - ? data - : std::max(initial_value, data); + // TODO: There is no need to branch with every element + if (reduction == SegmentReductionType::MAX) { + initial_value = at::_isnan(data) + ? data + : std::max(initial_value, data); + } else if (reduction == SegmentReductionType::MEAN) { + initial_value = at::_isnan(data) ? data : (initial_value + data); + } k++; } - // If unsafe is false, check on lengths or indices should cover cases - // where lengths for a particular segment is negative. If unsafe - // is true, simply set to initial_value for particular reduction + + // ===== step3: finalize reduction + TORCH_CHECK(lengths_data[i] >= 0); + if (lengths_data[i] == 0 && !initial.has_value()) { + output_data[i] = static_cast(NAN); + continue; + } output_data[i] = initial_value; + if (reduction == SegmentReductionType::MEAN && lengths_data[i] > 0 && + !at::_isnan(output_data[i])) { + output_data[i] = output_data[i] / lengths_data[i]; + } } })); @@ -53,6 +85,7 @@ Tensor _segment_reduce_cpu_backward_kernel( const Tensor& grad_contig, const Tensor& output_contig, const Tensor& data_contig, + SegmentReductionType reduction, const Tensor& lengths_contig) { auto grad_input = at::zeros({data_contig.sizes()}, grad_contig.options()); @@ -71,23 +104,35 @@ Tensor _segment_reduce_cpu_backward_kernel( const auto* values_data = data_contig.data_ptr(); int64_t k = 0; for (int64_t i = 0; i < batch_size; ++i) { - int64_t counter = 0; - for (int64_t j = 0; j < lengths_data[i]; ++j) { - if (at::_isnan(values_data[k]) || - values_data[k] == output_data[i]) { - grad_input_data[k] = grad_data[i]; - counter++; - } - k++; - } - // Average gradient based on number of maximum elements in the segment - if (counter < 2) { + if (lengths_data[i] == 0) { continue; } - for (int64_t j = 0; j < lengths_data[i]; ++j) { - int64_t index = k - j - 1; - if (grad_input_data[index] > 0) { - grad_input_data[index] = grad_input_data[index] / counter; + if (reduction == SegmentReductionType::MAX) { + int64_t counter = 0; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + if (at::_isnan(values_data[k]) || + values_data[k] == output_data[i]) { + grad_input_data[k] = grad_data[i]; + counter++; + } + k++; + } + // Average gradient based on number of maximum elements in the + // segment + if (counter < 2) { + continue; + } + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t index = k - j - 1; + if (grad_input_data[index] > 0) { + grad_input_data[index] = grad_input_data[index] / counter; + } + } + } else if (reduction == SegmentReductionType::MEAN) { + auto grad_val = grad_data[i] / lengths_data[i]; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + grad_input_data[k] = grad_val; + k++; } } } @@ -98,12 +143,6 @@ Tensor _segment_reduce_cpu_backward_kernel( } // namespace -enum SegmentReductionType { MAX }; -static const std::map segmentReduce2REDUCE = - { - {"max", MAX}, -}; - Tensor segment_reduce_kernel( const Tensor& data, c10::string_view reduce, @@ -116,9 +155,6 @@ Tensor segment_reduce_kernel( TORCH_CHECK(axis == 0, "Currently only dim=0 is supported!"); TORCH_CHECK(data.dim() == 1); TORCH_CHECK(data.numel() > 0); - TORCH_CHECK( - at::native::segmentReduce2REDUCE.at(reduce) == MAX, - "Currently only 'max' reduction is supported!"); // length related checks TORCH_CHECK( @@ -136,11 +172,13 @@ Tensor segment_reduce_kernel( TORCH_CHECK(lengths_value.sum().item() == data.numel()); } + auto reduction = get_reduction_enum(reduce); const auto data_contig = data.contiguous(); const auto lengths_contig = lengths_value.contiguous(); return _segment_reduce_stub( data_contig.device().type(), + reduction, data_contig, lengths_contig, axis, @@ -160,10 +198,11 @@ REGISTER_VSX_DISPATCH(_segment_reduce_stub, &_segment_reduce_cpu_kernel); // Currently some computation is beind duplicated across forward and backward. // TODO: Cache indices in forward pass to re-use in backward -Tensor segment_reduce_backward_kernel( +Tensor _segment_reduce_backward_kernel( const Tensor& grad, const Tensor& output, const Tensor& data, + c10::string_view reduce, const c10::optional& lengths) { TORCH_CHECK( lengths.has_value(), @@ -175,11 +214,13 @@ Tensor segment_reduce_backward_kernel( const auto data_contig = data.contiguous(); const auto lengths_contig = lengths_value.contiguous(); + auto reduction = get_reduction_enum(reduce); return _segment_reduce_backward_stub( grad_contig.device().type(), grad_contig, output_contig, data_contig, + reduction, lengths_contig); } diff --git a/aten/src/ATen/native/SegmentReduce.h b/aten/src/ATen/native/SegmentReduce.h index fb4ac8a82ea10..8bb7ece4e7d07 100644 --- a/aten/src/ATen/native/SegmentReduce.h +++ b/aten/src/ATen/native/SegmentReduce.h @@ -7,15 +7,22 @@ namespace at { namespace native { +enum SegmentReductionType { MAX, MEAN }; + using segment_reduce_fn = Tensor (*)( + SegmentReductionType, const Tensor&, const Tensor&, int64_t, const c10::optional&); DECLARE_DISPATCH(segment_reduce_fn, _segment_reduce_stub); -using segment_reduce_backward_fn = - Tensor (*)(const Tensor&, const Tensor&, const Tensor&, const Tensor&); +using segment_reduce_backward_fn = Tensor (*)( + const Tensor&, + const Tensor&, + const Tensor&, + SegmentReductionType, + const Tensor&); DECLARE_DISPATCH(segment_reduce_backward_fn, _segment_reduce_backward_stub); } // namespace native diff --git a/aten/src/ATen/native/cuda/SegmentReduce.cu b/aten/src/ATen/native/cuda/SegmentReduce.cu index cb4f005267ba9..1728961960573 100644 --- a/aten/src/ATen/native/cuda/SegmentReduce.cu +++ b/aten/src/ATen/native/cuda/SegmentReduce.cu @@ -42,6 +42,7 @@ Tensor _get_complete_sum(const Tensor& lengths) { } Tensor _segment_reduce_cuda_kernel( + SegmentReductionType reduction, const Tensor& data, const Tensor& lengths, int64_t axis, diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 8bdd0a6bc13f8..0fa16189054c0 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -10078,10 +10078,10 @@ dispatch: CPU, CUDA: segment_reduce_kernel -- func: segment_reduce_backward(Tensor grad, Tensor output, Tensor data, *, Tensor? lengths=None) -> Tensor +- func: _segment_reduce_backward(Tensor grad, Tensor output, Tensor data, str reduce, *, Tensor? lengths=None) -> Tensor variants: function dispatch: - CPU, CUDA: segment_reduce_backward_kernel + CPU, CUDA: _segment_reduce_backward_kernel - func: pad_sequence(Tensor[] sequences, bool batch_first=False, float padding_value=0.0) -> Tensor python_module: nn diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index ee7b55c312f2f..caa79a9c41823 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -94,6 +94,9 @@ ("aten::conj", datetime.date(2021, 8, 1)), ("aten::_conj", datetime.date(2021, 8, 1)), ("aten::conj.out", datetime.date(2021, 8, 1)), + ("aten::segment_reduce_backward", datetime.date(2021, 6, 15)), + ("aten::segment_reduce", datetime.date(2021, 8, 26)), + ("aten::_segment_reduce_backward", datetime.date(2021, 8, 26)), ] def allow_listed(schema, allow_list): diff --git a/test/test_segment_reductions.py b/test/test_segment_reductions.py index 8526083f8fd28..7338d462fcd84 100644 --- a/test/test_segment_reductions.py +++ b/test/test_segment_reductions.py @@ -12,7 +12,7 @@ class TestSegmentReductions(TestCase): - def _test_max_simple_1d(self, device, dtype, unsafe, axis): + def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): lengths = torch.tensor([1, 2, 3, 0], device=device) data = torch.tensor( [1, float("nan"), 3, 4, 5, 5], @@ -21,30 +21,40 @@ def _test_max_simple_1d(self, device, dtype, unsafe, axis): requires_grad=True, ) initial_value = 0 - expected_result = torch.tensor( - [1, float("nan"), 5, initial_value], device=device, dtype=dtype - ) + if reduction == "max": + expected_result = torch.tensor( + [1, float("nan"), 5, initial_value], device=device, dtype=dtype + ) + expected_grad = torch.tensor( + [1, 1, 0, 0, 0.5, 0.5], device=device, dtype=dtype + ) + elif reduction == "mean": + expected_result = torch.tensor( + [1, float("nan"), 4.666, initial_value], device=device, dtype=dtype + ) + expected_grad = torch.tensor( + [1.0, 0.5, 0.5, 0.333, 0.333, 0.333], device=device, dtype=dtype + ) actual_result = torch.segment_reduce( data=data, - reduce="max", + reduce=reduction, lengths=lengths, axis=axis, unsafe=unsafe, initial=initial_value, ) self.assertEqual( - expected_result, actual_result, rtol=1e-03, atol=1e-05, equal_nan=True + expected_result, actual_result, rtol=1e-02, atol=1e-05, equal_nan=True ) - # Backward is only supported for cpu tensors for now. Return early if cuda + # TODO: Remove this check once cuda backward support is implemented if data.is_cuda: return # Test backward - expected_grad = torch.tensor([1, 1, 0, 0, 0.5, 0.5], device=device, dtype=dtype) actual_result.sum().backward() self.assertEqual( - expected_grad, data.grad, rtol=1e-03, atol=1e-05, equal_nan=True + expected_grad, data.grad, rtol=1e-02, atol=1e-05, equal_nan=True ) # gradcheck does not work well with bfloat16 or fp16 cpu types @@ -61,7 +71,7 @@ def _test_max_simple_1d(self, device, dtype, unsafe, axis): gradcheck( lambda x: torch.segment_reduce( data=x, - reduce="max", + reduce=reduction, lengths=lengths, axis=axis, unsafe=unsafe, @@ -73,11 +83,15 @@ def _test_max_simple_1d(self, device, dtype, unsafe, axis): @dtypesIfCUDA(torch.half, torch.bfloat16, torch.float, torch.double) @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) - def test_max_simple_1d(self, device, dtype): - self._test_max_simple_1d(device, dtype, False, 0) - self._test_max_simple_1d(device, dtype, False, -1) - self._test_max_simple_1d(device, dtype, True, 0) - self._test_max_simple_1d(device, dtype, True, -1) + def test_simple_1d(self, device, dtype): + for reduction in ("max", "mean"): + # TODO: Remove if once mean reduction for cuda is implemented + if reduction == "mean" and device != "cpu": + continue + self._test_simple_1d(reduction, device, dtype, False, 0) + self._test_simple_1d(reduction, device, dtype, False, -1) + self._test_simple_1d(reduction, device, dtype, True, 0) + self._test_simple_1d(reduction, device, dtype, True, -1) instantiate_device_type_tests(TestSegmentReductions, globals()) diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index 0619cebfd4761..ba248da1773ef 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -2118,4 +2118,4 @@ output_differentiability: [False] - name: segment_reduce(Tensor data, str reduce, *, Tensor? lengths=None, Tensor? indices=None, int axis=0, bool unsafe=False, Scalar? initial=None) -> Tensor - data: segment_reduce_backward(grad, result, data, lengths) + data: _segment_reduce_backward(grad, result, data, reduce, lengths) From e9e9291dc1dc34f102fb846edbcaf9e2c24ba0e0 Mon Sep 17 00:00:00 2001 From: Chen Lai Date: Thu, 10 Jun 2021 15:01:56 -0700 Subject: [PATCH 022/305] [After fix] Reuse constant and bump bytecode to v5 (#59722) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59722 Reintroduce sharing constant between bytecode and torchscript (same as #58629) after the fix #59642 Test Plan: Imported from OSS Reviewed By: iseeyuan Differential Revision: D29002345 Pulled By: cccclai fbshipit-source-id: d9c8e474ff57d0509580183206df038a24ad27e3 --- caffe2/serialize/versions.h | 7 +- test/mobile/test_bytecode.py | 94 ++++++++++--------- .../csrc/jit/serialization/export_module.cpp | 6 +- 3 files changed, 56 insertions(+), 51 deletions(-) diff --git a/caffe2/serialize/versions.h b/caffe2/serialize/versions.h index ad17f58a8e9f3..eee98e4fa9cd3 100644 --- a/caffe2/serialize/versions.h +++ b/caffe2/serialize/versions.h @@ -65,13 +65,16 @@ constexpr uint64_t kProducedFileFormatVersion = 0x3L; // 0x1L: Initial version // 0x2L: (Comment missing) // 0x3L: (Comment missing) -// 0x4L: (Comment missing) // 0x4L: (update) Added schema to function tuple. Forward-compatible change. // 0x5L: (update) Update bytecode is sharing constant tensor files from torchscript, and only serialize +// extra tensors that are not in the torchscript constant table. Also update tensor storage schema adapting +// to the unify format, the root key of tensor storage is updated from {index} to +// {the_pointer_value_the_tensor.storage}, for example: `140245072983168.storage` +// Forward-compatibility change. // 0x6L: Implicit opereator versioning using number of specified argument. // Refer to the summary of https://github.com/pytorch/pytorch/pull/56845 // for details. -constexpr uint64_t kProducedBytecodeVersion = 0x4L; +constexpr uint64_t kProducedBytecodeVersion = 0x5L; static_assert(kProducedBytecodeVersion >= kProducedFileFormatVersion, "kProducedBytecodeVersion must be higher or equal to kProducedFileFormatVersion."); diff --git a/test/mobile/test_bytecode.py b/test/mobile/test_bytecode.py index b6f1683c94b87..5511e6a63b085 100644 --- a/test/mobile/test_bytecode.py +++ b/test/mobile/test_bytecode.py @@ -4,7 +4,7 @@ import tempfile import torch import torch.utils.show_pickle -from torch.utils.mobile_optimizer import optimize_for_mobile +# from torch.utils.mobile_optimizer import optimize_for_mobile from torch.jit.mobile import ( _load_for_lite_interpreter, _get_model_bytecode_version, @@ -189,51 +189,53 @@ def test_bytecode_values_for_all_backport_functions(self): current_from_version -= 1 shutil.rmtree(tmpdirname) - def test_all_backport_functions(self): - # Backport from the latest bytecode version to the minimum support version - # Load, run the backport model, and check version - class TestModule(torch.nn.Module): - def __init__(self, v): - super().__init__() - self.x = v - - def forward(self, y: int): - increment = torch.ones([2, 4], dtype=torch.float64) - return self.x + y + increment - - module_input = 1 - expected_mobile_module_result = 3 * torch.ones([2, 4], dtype=torch.float64) - - # temporary input model file and output model file will be exported in the temporary folder - with tempfile.TemporaryDirectory() as tmpdirname: - tmp_input_model_path = Path(tmpdirname, "tmp_script_module.ptl") - script_module = torch.jit.script(TestModule(1)) - optimized_scripted_module = optimize_for_mobile(script_module) - exported_optimized_scripted_module = optimized_scripted_module._save_for_lite_interpreter(str(tmp_input_model_path)) - - current_from_version = _get_model_bytecode_version(tmp_input_model_path) - current_to_version = current_from_version - 1 - tmp_output_model_path = Path(tmpdirname, "tmp_script_module_backport.ptl") - - while current_to_version >= MINIMUM_TO_VERSION: - # Backport the latest model to `to_version` to a tmp file "tmp_script_module_backport" - backport_success = _backport_for_mobile(tmp_input_model_path, tmp_output_model_path, current_to_version) - assert(backport_success) - - backport_version = _get_model_bytecode_version(tmp_output_model_path) - assert(backport_version == current_to_version) - - # Load model and run forward method - mobile_module = _load_for_lite_interpreter(str(tmp_input_model_path)) - mobile_module_result = mobile_module(module_input) - torch.testing.assert_allclose(mobile_module_result, expected_mobile_module_result) - current_to_version -= 1 - - # Check backport failure case - backport_success = _backport_for_mobile(tmp_input_model_path, tmp_output_model_path, MINIMUM_TO_VERSION - 1) - assert(not backport_success) - # need to clean the folder before it closes, otherwise will run into git not clean error - shutil.rmtree(tmpdirname) + # Please run this test manually when working on backport. + # This test passes in OSS, but fails internally, likely due to missing step in build + # def test_all_backport_functions(self): + # # Backport from the latest bytecode version to the minimum support version + # # Load, run the backport model, and check version + # class TestModule(torch.nn.Module): + # def __init__(self, v): + # super().__init__() + # self.x = v + + # def forward(self, y: int): + # increment = torch.ones([2, 4], dtype=torch.float64) + # return self.x + y + increment + + # module_input = 1 + # expected_mobile_module_result = 3 * torch.ones([2, 4], dtype=torch.float64) + + # # temporary input model file and output model file will be exported in the temporary folder + # with tempfile.TemporaryDirectory() as tmpdirname: + # tmp_input_model_path = Path(tmpdirname, "tmp_script_module.ptl") + # script_module = torch.jit.script(TestModule(1)) + # optimized_scripted_module = optimize_for_mobile(script_module) + # exported_optimized_scripted_module = optimized_scripted_module._save_for_lite_interpreter(str(tmp_input_model_path)) + + # current_from_version = _get_model_bytecode_version(tmp_input_model_path) + # current_to_version = current_from_version - 1 + # tmp_output_model_path = Path(tmpdirname, "tmp_script_module_backport.ptl") + + # while current_to_version >= MINIMUM_TO_VERSION: + # # Backport the latest model to `to_version` to a tmp file "tmp_script_module_backport" + # backport_success = _backport_for_mobile(tmp_input_model_path, tmp_output_model_path, current_to_version) + # assert(backport_success) + + # backport_version = _get_model_bytecode_version(tmp_output_model_path) + # assert(backport_version == current_to_version) + + # # Load model and run forward method + # mobile_module = _load_for_lite_interpreter(str(tmp_input_model_path)) + # mobile_module_result = mobile_module(module_input) + # torch.testing.assert_allclose(mobile_module_result, expected_mobile_module_result) + # current_to_version -= 1 + + # # Check backport failure case + # backport_success = _backport_for_mobile(tmp_input_model_path, tmp_output_model_path, MINIMUM_TO_VERSION - 1) + # assert(not backport_success) + # # need to clean the folder before it closes, otherwise will run into git not clean error + # shutil.rmtree(tmpdirname) # Check just the test_backport_bytecode_from_file_to_file mechanism but not the function implementations def test_backport_bytecode_from_file_to_file(self): diff --git a/torch/csrc/jit/serialization/export_module.cpp b/torch/csrc/jit/serialization/export_module.cpp index 856897be534cd..68b7f20af4047 100644 --- a/torch/csrc/jit/serialization/export_module.cpp +++ b/torch/csrc/jit/serialization/export_module.cpp @@ -436,7 +436,7 @@ void ScriptModuleSerializer::serialize( /*archive_name=*/"constants", /*archive_dir=*/"", /*tensor_dir=*/"constants/", - /*tensor_cdata_naming_scheme=*/false); + /*tensor_cdata_naming_scheme=*/true); writeByteCode(module, save_mobile_debug_info); writeMobileMetadata(module, extra_files); @@ -644,8 +644,8 @@ void ScriptModuleSerializer::writeByteCode( telements, /*archive_name=*/"bytecode", /*archive_dir=*/"", - /*tensor_dir=*/"bytecode/", - /*tensor_cdata_naming_scheme=*/false); + /*tensor_dir=*/"constants/", + /*tensor_cdata_naming_scheme=*/true); auto debug_info_telements = Tup(std::move(debug_info_elements)); From 60ba451731acaebea11b01be4ec5967c5a413a0c Mon Sep 17 00:00:00 2001 From: Nils Plath Date: Thu, 10 Jun 2021 15:11:49 -0700 Subject: [PATCH 023/305] [torch] Remove using directive from header (#59728) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59728 I noticed Sandcastle jobs failing with: ``` fbcode/caffe2/torch/csrc/api/include/torch/nn/modules/rnn.h:19:35: error: using namespace directive in global context in header [-Werror,-Wheader-hygiene] using namespace torch::nn::utils::rnn; ``` (cf. V3 of D28939167 or https://www.internalfb.com/intern/sandcastle/job/36028797455955174/). Removing `using namespace ...` fixes the problem. ~~... also applied code formatting ...~~ Test Plan: Sandcastle Reviewed By: jbschlosser Differential Revision: D29000888 fbshipit-source-id: 10917426828fc0c82b982da435ce891dc2bb6eec --- torch/csrc/api/include/torch/nn/modules/rnn.h | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/torch/csrc/api/include/torch/nn/modules/rnn.h b/torch/csrc/api/include/torch/nn/modules/rnn.h index 7ea669944bc8b..adf8500d194e9 100644 --- a/torch/csrc/api/include/torch/nn/modules/rnn.h +++ b/torch/csrc/api/include/torch/nn/modules/rnn.h @@ -16,8 +16,6 @@ #include #include -using namespace torch::nn::utils::rnn; - namespace torch { namespace nn { @@ -113,7 +111,7 @@ class TORCH_API RNNImpl : public detail::RNNImplBase { FORWARD_HAS_DEFAULT_ARGS({1, AnyValue(Tensor())}) public: - std::tuple forward_with_packed_input(const PackedSequence& packed_input, Tensor hx = {}); + std::tuple forward_with_packed_input(const torch::nn::utils::rnn::PackedSequence& packed_input, Tensor hx = {}); RNNOptions options; @@ -159,8 +157,8 @@ class TORCH_API LSTMImpl : public detail::RNNImplBase { FORWARD_HAS_DEFAULT_ARGS({1, AnyValue(torch::optional>())}) public: - std::tuple> forward_with_packed_input( - const PackedSequence& packed_input, torch::optional> hx_opt = {}); + std::tuple> forward_with_packed_input( + const torch::nn::utils::rnn::PackedSequence& packed_input, torch::optional> hx_opt = {}); LSTMOptions options; @@ -211,7 +209,7 @@ class TORCH_API GRUImpl : public detail::RNNImplBase { FORWARD_HAS_DEFAULT_ARGS({1, AnyValue(torch::Tensor())}) public: - std::tuple forward_with_packed_input(const PackedSequence& packed_input, Tensor hx = {}); + std::tuple forward_with_packed_input(const torch::nn::utils::rnn::PackedSequence& packed_input, Tensor hx = {}); GRUOptions options; From c2c35c017045d6d0f9f8e551b5ff0af1f58c1acf Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Thu, 10 Jun 2021 16:51:37 -0700 Subject: [PATCH 024/305] [Binary] Link whole CuDNN for CUDA-11.1 (#59802) Summary: Fixes https://github.com/pytorch/pytorch/issues/50153 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59802 Reviewed By: driazati, seemethere Differential Revision: D29033537 Pulled By: malfet fbshipit-source-id: e816fc71f273ae0b4ba8a0621d5368a2078561a1 --- .circleci/scripts/binary_populate_env.sh | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.circleci/scripts/binary_populate_env.sh b/.circleci/scripts/binary_populate_env.sh index 3e581f0969574..329bf09a3ae10 100755 --- a/.circleci/scripts/binary_populate_env.sh +++ b/.circleci/scripts/binary_populate_env.sh @@ -74,6 +74,12 @@ if [[ ${DESIRED_CUDA} == "cpu" ]]; then USE_GOLD_LINKER="ON" fi +USE_WHOLE_CUDNN="OFF" +# Link whole cuDNN for CUDA-11.1 to include fp16 fast kernels +if [[ "$(uname)" == "Linux" && "${DESIRED_CUDA}" == "cu111" ]]; then + USE_WHOLE_CUDNN="ON" +fi + # Default to nightly, since that's where this normally uploads to PIP_UPLOAD_FOLDER='nightly/' # We put this here so that OVERRIDE_PACKAGE_VERSION below can read from it @@ -176,7 +182,8 @@ export CIRCLE_BRANCH="$CIRCLE_BRANCH" export CIRCLE_WORKFLOW_ID="$CIRCLE_WORKFLOW_ID" export USE_GOLD_LINKER="${USE_GOLD_LINKER}" -export USE_GLOO_WITH_OPENSSL=1 +export USE_GLOO_WITH_OPENSSL="ON" +export USE_WHOLE_CUDNN="${USE_WHOLE_CUDNN}" # =================== The above code will be executed inside Docker container =================== EOL From 8a7c0d082f8581d24a35f66988ed1b26a8e0c001 Mon Sep 17 00:00:00 2001 From: albanD Date: Thu, 10 Jun 2021 19:26:51 -0700 Subject: [PATCH 025/305] ger is an alias to outer, not the other way around (#59710) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59710 This is the exact same PR as before. The version that landed was actually outdated compared to the github PR and that's why it failed on master... Sorry for the noise. Test Plan: Imported from OSS Reviewed By: zou3519 Differential Revision: D28995764 Pulled By: albanD fbshipit-source-id: 8f7ae3356a886d45787c5e6ca53a4e7b033e306e --- aten/src/ATen/native/native_functions.yaml | 6 +----- tools/autograd/derivatives.yaml | 4 ---- tools/autograd/gen_variable_type.py | 2 +- torch/overrides.py | 4 ++-- 4 files changed, 4 insertions(+), 12 deletions(-) diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 0fa16189054c0..5e0dd9917dd9f 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -9858,20 +9858,16 @@ - func: inner.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) -# torch.outer, alias for torch.ger - func: outer(Tensor self, Tensor vec2) -> Tensor variants: function, method - func: outer.out(Tensor self, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) +# torch.ger, alias for torch.outer - func: ger(Tensor self, Tensor vec2) -> Tensor variants: function, method - dispatch: - CompositeExplicitAutograd: ger - func: ger.out(Tensor self, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) - dispatch: - CompositeExplicitAutograd: ger_out - func: linalg_norm(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor python_module: linalg diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index ba248da1773ef..d8af17ffc737a 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -578,10 +578,6 @@ - name: geqrf(Tensor self) -> (Tensor a, Tensor tau) self: not_implemented("geqrf") -- name: ger(Tensor self, Tensor vec2) -> Tensor - self: grad.mv(vec2.conj()) - vec2: grad.t().mv(self.conj()) - - name: indices(Tensor(a) self) -> Tensor(a) output_differentiability: [False] diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py index 3c7136aa0d49a..5b667a6cf9692 100644 --- a/tools/autograd/gen_variable_type.py +++ b/tools/autograd/gen_variable_type.py @@ -87,7 +87,7 @@ 'sgn', 'asin', 'acos', 'sub', 'div', 'cat', 'view_as_complex', 'neg', 'complex', 'select', '_s_where', 'as_strided', 'slice', 'constant_pad_nd', 'unbind', 'split', 'split_with_sizes', 'unsafe_split', 'split_with_sizes_backward', - 'dot', 'vdot', 'cholesky', 'triangular_solve', 'mm', '_unsafe_view', 'mv', 'ger', + 'dot', 'vdot', 'cholesky', 'triangular_solve', 'mm', '_unsafe_view', 'mv', 'outer', 'bmm', 'diagonal', 'alias', 'atan', 'log', 'log10', 'log1p', 'log2', 'reciprocal', 'tan', 'pow', 'rsqrt', 'tanh', 'tanh_backward', 'asinh', 'acosh', 'atanh', 'take', 'fill_', 'exp', 'nonzero', 'mean', 'inverse', 'solve', 'linalg_cholesky', 'addcmul', 'addcdiv', diff --git a/torch/overrides.py b/torch/overrides.py index 2f931d53de635..75bde5decb787 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -464,8 +464,8 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.geqrf: lambda input, out=None: -1, torch.i0: lambda input, out=None: -1, torch.inner: lambda input, other, out=None: -1, - torch.outer: lambda input, vec2, out=None: -1, # alias for torch.ger - torch.ger: lambda input, vec2, out=None: -1, + torch.outer: lambda input, vec2, out=None: -1, + torch.ger: lambda input, vec2, out=None: -1, # alias for torch.outer torch.gradient: lambda input, spacing=None, dim=None, edge_order=1: -1, torch.grid_sampler: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, torch.grid_sampler_2d: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, From a524ee00ca8474e8d9b8bb83814316f54100acf8 Mon Sep 17 00:00:00 2001 From: albanD Date: Thu, 10 Jun 2021 19:26:51 -0700 Subject: [PATCH 026/305] Forward AD formulas batch 3 (#59711) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59711 This is the exact same PR as before. This was reverted before the PR below was faulty. Test Plan: Imported from OSS Reviewed By: zou3519 Differential Revision: D28995762 Pulled By: albanD fbshipit-source-id: 65940ad93bced9b5f97106709d603d1cd7260812 --- tools/autograd/derivatives.yaml | 62 +++++++++++++++++++ torch/autograd/gradcheck.py | 44 ++++++++----- torch/csrc/autograd/FunctionsManual.cpp | 19 ++++++ torch/csrc/autograd/FunctionsManual.h | 2 + .../_internal/common_methods_invocations.py | 59 ++++++++++++++++++ 5 files changed, 172 insertions(+), 14 deletions(-) diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index d8af17ffc737a..c67c4e6812622 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -438,9 +438,11 @@ - name: cummax(Tensor self, int dim) -> (Tensor values, Tensor indices) self: cummaxmin_backward(grad, self, indices, dim) + values: self_t.gather(dim, indices) - name: cummin(Tensor self, int dim) -> (Tensor values, Tensor indices) self: cummaxmin_backward(grad, self, indices, dim) + values: self_t.gather(dim, indices) - name: conv_tbc(Tensor self, Tensor weight, Tensor bias, int pad=0) -> Tensor self, weight, bias: "grad.defined() ? conv_tbc_backward(grad, self, weight, bias, pad) : std::tuple()" @@ -450,6 +452,7 @@ - name: deg2rad(Tensor self) -> Tensor self: deg2rad_backward(grad) + result: auto_element_wise - name: linalg_det(Tensor self) -> Tensor self: linalg_det_backward(grad, self, result) @@ -470,24 +473,30 @@ - name: div.Tensor(Tensor self, Tensor other) -> Tensor self: div_tensor_self_backward(grad, other, self.scalar_type()) other: div_tensor_other_backward(grad, self, other) + result: self_t / other_p - other_t * (self_p / other_p) / other_p - name: div.Scalar(Tensor self, Scalar other) -> Tensor self: div_tensor_self_backward(grad, at::scalar_to_tensor(other), self.scalar_type()) + result: self_t / other - name: div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor self: div_tensor_self_backward(grad, other, self.scalar_type(), rounding_mode) other: div_tensor_other_backward(grad, self, other, rounding_mode) + result: "rounding_mode.has_value() ? result.new_zeros(result.sizes()) : self_t / other_p - other_t * (self_p / other_p) / other_p" - name: div.Scalar_mode(Tensor self, Scalar other, *, str? rounding_mode) -> Tensor self: div_tensor_self_backward(grad, at::scalar_to_tensor(other), self.scalar_type(), rounding_mode) + result: "rounding_mode.has_value() ? result.new_zeros(result.sizes()) : self_t / other" - name: dot(Tensor self, Tensor tensor) -> Tensor self: grad * tensor.conj() tensor: grad * self.conj() + result: at::dot(self_t, tensor_p) + at::dot(self_p, tensor_t) - name: vdot(Tensor self, Tensor other) -> Tensor self: grad.conj() * other other: grad * self + result: at::vdot(self_t, other_p) + at::vdot(self_p, other_t) - name: _fused_dropout(Tensor self, float p, Generator? generator=None) -> (Tensor, Tensor) self: _fused_dropout_backward(grad, result1, p) @@ -513,15 +522,19 @@ - name: exp(Tensor self) -> Tensor self: grad * result.conj() + result: auto_element_wise - name: exp2(Tensor self) -> Tensor self: grad * result * M_LN2 + result: auto_element_wise - name: expm1(Tensor self) -> Tensor self: grad * (result + 1) + result: auto_element_wise - name: expand(Tensor(a) self, int[] size, *, bool implicit=False) -> Tensor(a) self: at::sum_to(grad, self.sizes()) + result: auto_linear - name: exponential_(Tensor(a!) self, float lambd=1, *, Generator? generator=None) -> Tensor(a!) self: zeros_like(grad) @@ -547,6 +560,7 @@ - name: floor(Tensor self) -> Tensor self: zeros_like(grad) + result: auto_element_wise - name: fmod.Scalar(Tensor self, Scalar other) -> Tensor self: grad @@ -557,13 +571,16 @@ - name: frac(Tensor self) -> Tensor self: grad + result: self_t - name: frexp.Tensor(Tensor self) -> (Tensor mantissa, Tensor exponent) self: grad / exponent.exp2() + mantissa: self_t / exponent.exp2() - name: gather(Tensor self, int dim, Tensor index, *, bool sparse_grad=False) -> Tensor self: gather_backward(grad, self, dim, index, sparse_grad) index: non_differentiable + result: auto_linear - name: ge_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) self: zeros_like(self) @@ -609,10 +626,12 @@ - name: hardswish(Tensor self) -> Tensor self: hardswish_backward(grad, self) + result: auto_element_wise - name: hypot(Tensor self, Tensor other) -> Tensor self: grad * self / result other: grad * other / result + result: self_t * self_p / result + other_t * other_p / result - name: i0(Tensor self) -> Tensor self: grad * at::special_i1(self) @@ -644,6 +663,7 @@ # This is because source is not broadcastable to index, as source.dim() < index.dim() source: "maybe_multiply(source.dim() > 0 ? grad.index_select(dim, index).expand_as(source) : grad.index_select(dim, index.squeeze(0)), alpha)" index: non_differentiable + result: self_t.index_add_(dim, index, maybe_multiply(source_t, alpha)) - name: index_copy_(Tensor(a!) self, int dim, Tensor index, Tensor source) -> Tensor(a!) self: grad.clone().index_fill_(dim, index, 0) @@ -652,27 +672,33 @@ # This is because source is not broadcastable to index, as source.dim() < index.dim() source: "source.dim() > 0 ? grad.index_select(dim, index).expand_as(source) : grad.index_select(dim, index.squeeze(0))" index: non_differentiable + result: self_t.index_copy_(dim, index, source_t) - name: index_fill_.int_Scalar(Tensor(a!) self, int dim, Tensor index, Scalar value) -> Tensor(a!) self: grad.clone().index_fill_(dim, index, 0) index: non_differentiable + result: self_t.index_fill_(dim, index, 0) - name: index_fill_.int_Tensor(Tensor(a!) self, int dim, Tensor index, Tensor value) -> Tensor(a!) self: grad.clone().index_fill_(dim, index, 0) value: grad.index_select(dim, std::get<0>(at::_unique(index, /*sorted=*/false))).sum() index: non_differentiable + result: self_t.index_fill_(dim, index, value_t) - name: index_put_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor(a!) self: "accumulate ? grad : grad.clone().index_put_(indices, zeros_like(values), false)" values: grad.index(indices) + result: self_t.index_put_(indices, values_t, accumulate) - name: _index_put_impl_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False, bool unsafe=False) -> Tensor(a!) self: "accumulate ? grad : grad.clone().index_put_(indices, zeros_like(values), false)" values: grad.index(indices) + result: at::_index_put_impl_(self_t, indices, values_t, accumulate, unsafe) - name: index_select(Tensor self, int dim, Tensor index) -> Tensor self: index_select_backward(grad, self.sizes(), dim, index) index: non_differentiable + result: auto_linear - name: inverse(Tensor self) -> Tensor self: -at::matmul(result.conj().transpose(-2, -1), at::matmul(grad, result.conj().transpose(-2, -1))) @@ -685,6 +711,7 @@ - name: kthvalue(Tensor self, int k, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) self: value_selecting_reduction_backward(grad, dim, indices, self.sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) - name: le_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) self: zeros_like(self) @@ -696,63 +723,81 @@ - name: lerp.Scalar(Tensor self, Tensor end, Scalar weight) -> Tensor self: "weight.isComplex() ? grad * (1 - weight.conj().toComplexDouble()) : grad * (1 - weight.toDouble())" end: grad * weight.conj() + result: at::lerp(self_t, end_t, weight) - name: lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor self: grad * (1 - weight).conj() end: grad * weight.conj() weight: grad * (end - self).conj() + result: at::lerp(self_t, end_t, weight_p) + weight_t * (end_p - self_p) - name: lgamma(Tensor self) -> Tensor self: grad * digamma(self) + result: auto_element_wise - name: digamma(Tensor self) -> Tensor self: grad * polygamma(1, self) + result: auto_element_wise - name: polygamma(int n, Tensor self) -> Tensor self: grad * polygamma(n + 1, self) + result: auto_element_wise - name: polygamma_(Tensor(a!) self, int n) -> Tensor(a!) self: grad * polygamma(n + 1, self) + result: auto_element_wise - name: log(Tensor self) -> Tensor self: grad.div(self.conj()) + result: auto_element_wise - name: log10(Tensor self) -> Tensor self: grad / (self.conj() * 2.3025850929940456) + result: auto_element_wise - name: log1p(Tensor self) -> Tensor self: log1p_backward(grad, self) + result: auto_element_wise - name: log2(Tensor self) -> Tensor self: grad / (self.conj() * 0.6931471805599453) + result: auto_element_wise - name: logaddexp(Tensor self, Tensor other) -> Tensor self: grad / (1 + exp(other - self)) other: grad / (1 + exp(self - other)) + result: self_t / (1 + exp(other_p - self_p)) + other_t / (1 + exp(self_p - other_p)) - name: logaddexp2(Tensor self, Tensor other) -> Tensor self: grad / (1 + pow(2, other - self)) other: grad / (1 + pow(2, self - other)) + result: self_t / (1 + pow(2, other_p - self_p)) + other_t / (1 + pow(2, self_p - other_p)) - name: xlogy.Tensor(Tensor self, Tensor other) -> Tensor self: grad * at::xlogy((self != 0), other) other: grad * self / other + result: self_t * at::xlogy((self_p != 0), other_p) + other_t * self_p / other_p - name: xlogy.Scalar_Self(Scalar self, Tensor other) -> Tensor other: grad * self / other + result: auto_element_wise - name: xlogy.Scalar_Other(Tensor self, Scalar other) -> Tensor self: grad * at::xlogy((self != 0), other) + result: auto_element_wise - name: special_xlog1py(Tensor self, Tensor other) -> Tensor self: grad * other.log1p() other: grad * self / (other + 1) + result: self_t * other_p.log1p() + other_t * self_p / (other_p + 1) - name: special_xlog1py.self_scalar(Scalar self, Tensor other) -> Tensor other: grad * self / (other + 1) + result: auto_element_wise - name: special_xlog1py.other_scalar(Tensor self, Scalar other) -> Tensor self: grad * log1p(other.toDouble()) + result: auto_element_wise - name: logdet(Tensor self) -> Tensor self: logdet_backward(grad, self, result) @@ -792,43 +837,53 @@ - name: masked_fill_.Scalar(Tensor(a!) self, Tensor mask, Scalar value) -> Tensor(a!) self: grad.clone().masked_fill_(mask, 0) mask: non_differentiable + result: self_t.masked_fill_(mask, 0) - name: masked_fill_.Tensor(Tensor(a!) self, Tensor mask, Tensor value) -> Tensor(a!) self: grad.clone().masked_fill_(mask, 0) value: at::where(mask, grad, zeros_like(grad)).sum() mask: non_differentiable + result: self_t.masked_fill_(mask, value_t) - name: masked_scatter_(Tensor(a!) self, Tensor mask, Tensor source) -> Tensor(a!) self: grad.clone().masked_fill_(mask, 0) source: masked_scatter_backward(grad, mask, source.sizes()) mask: non_differentiable + result: self_t.masked_scatter_(mask, source_t) - name: masked_select(Tensor self, Tensor mask) -> Tensor self: masked_select_backward(grad, self, mask) mask: non_differentiable + result: auto_linear - name: matrix_exp(Tensor self) -> Tensor self: matrix_exp_backward(self, grad) - name: max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) self: value_selecting_reduction_backward(grad, dim, indices, self.sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) - name: max(Tensor self) -> Tensor self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) - name: maximum(Tensor self, Tensor other) -> Tensor self: at::where(self == other, grad / 2, grad).masked_fill_(self < other, 0) other: at::where(self == other, grad / 2, grad).masked_fill_(self > other, 0) + result: other_t + at::where(self_p == other_p, 0.5, (self_p > other_p).to(result.scalar_type())) * (self_t - other_t) - name: fmax(Tensor self, Tensor other) -> Tensor self: grad.clone().masked_fill_((self >= other).logical_or_(other.isnan()).logical_not_(), 0) other: grad.clone().masked_fill_((self >= other).logical_or_(other.isnan()), 0) + result: other_t + (self_p > other_p).logical_or_(other_p.isnan()) * (self_t - other_t) - name: mean(Tensor self, *, ScalarType? dtype=None) -> Tensor self: grad.expand(self.sizes()).to(self.scalar_type()) / self.numel() + result: auto_linear - name: mean.dim(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor self: sum_backward(grad, self.sizes(), dim, keepdim).to(self.scalar_type()) / _safe_size(self.sizes(), dim) + result: auto_linear - name: median(Tensor self) -> Tensor self: evenly_distribute_backward(grad, self, result) @@ -858,17 +913,21 @@ - name: min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) self: value_selecting_reduction_backward(grad, dim, indices, self.sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) - name: min(Tensor self) -> Tensor self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) - name: minimum(Tensor self, Tensor other) -> Tensor self: at::where(self == other, grad / 2, grad).masked_fill_(self > other, 0) other: at::where(self == other, grad / 2, grad).masked_fill_(self < other, 0) + result: other_t + at::where(self_p == other_p, 0.5, (self_p < other_p).to(result.scalar_type())) * (self_t - other_t) - name: fmin(Tensor self, Tensor other) -> Tensor self: grad.clone().masked_fill_((self <= other).logical_or_(other.isnan()).logical_not_(), 0) other: grad.clone().masked_fill_((self <= other).logical_or_(other.isnan()), 0) + result: other_t + (self_p <= other_p).logical_or_(other_p.isnan()) * (self_t - other_t) - name: amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor self: scale_grad_by_count(restore_reduced_dims(grad, dim, keepdim), restore_reduced_dims(result, dim, keepdim) == self, dim) @@ -879,9 +938,11 @@ - name: mm(Tensor self, Tensor mat2) -> Tensor self: mm_mat1_backward(grad, mat2, self.sizes(), self.strides(), 1) mat2: mm_mat2_backward(grad, self, mat2.sizes(), mat2.strides(), 1) + result: at::mm(self_t, mat2_p) + at::mm(self_p, mat2_t) - name: mode(Tensor self, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) self: value_selecting_reduction_backward(grad, dim, indices, self.sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) - name: mul.Tensor(Tensor self, Tensor other) -> Tensor self: mul_tensor_backward(grad, other, self.scalar_type()) @@ -890,6 +951,7 @@ - name: mul.Scalar(Tensor self, Scalar other) -> Tensor self: mul_tensor_backward(grad, at::scalar_to_tensor(other), self.scalar_type()) + result: self_t * other - name: mv(Tensor self, Tensor vec) -> Tensor self: grad.ger(vec.conj()) diff --git a/torch/autograd/gradcheck.py b/torch/autograd/gradcheck.py index db6318fdceb8e..5e7f56736ee69 100644 --- a/torch/autograd/gradcheck.py +++ b/torch/autograd/gradcheck.py @@ -340,7 +340,8 @@ def _get_analytical_jacobian_forward_ad(fn, inputs, outputs, *, check_grad_dtype # To be consistent with numerical evaluation, we actually compute one reduction per input for i, (fw_grad, u) in enumerate(zip(fw_grads, all_u)): fw_grad.copy_(u.view_as(fw_grad)) - dual_outputs = _as_tuple(fn(*dual_inputs)) + raw_outputs = _as_tuple(fn(*dual_inputs)) + dual_outputs = filter(_is_float_or_complex_tensor, raw_outputs) for index_o, d_o in enumerate(dual_outputs): val, res = fwAD.unpack_dual(d_o) if check_grad_dtypes and val.is_complex() != res.is_complex(): @@ -358,9 +359,13 @@ def _get_analytical_jacobian_forward_ad(fn, inputs, outputs, *, check_grad_dtype for i, fw_grad in enumerate(fw_grads): for lin_idx, grad_idx in enumerate(product(*[range(m) for m in fw_grad.size()])): fw_grad[grad_idx] = 1. - dual_outputs = _as_tuple(fn(*dual_inputs)) + raw_outputs = _as_tuple(fn(*dual_inputs)) + dual_outputs = filter(_is_float_or_complex_tensor, raw_outputs) for index_o, d_o in enumerate(dual_outputs): - _, res = fwAD.unpack_dual(d_o) + val, res = fwAD.unpack_dual(d_o) + if check_grad_dtypes and val.is_complex() != res.is_complex(): + raise GradcheckError('Forward AD gradient has dtype mismatch.') + if res is None: jacobians[i][index_o][lin_idx].zero_() else: @@ -419,7 +424,7 @@ def _mul_tensor_or_tuple(u, k): return k * u -def _get_numerical_jvp_wrt_specific_input(fn, input_idx, inputs, outputs, u, eps, is_forward_ad=False) -> List[torch.Tensor]: +def _get_numerical_jvp_wrt_specific_input(fn, input_idx, inputs, u, eps, is_forward_ad=False) -> List[torch.Tensor]: input = inputs[input_idx] input_to_perturb = _get_input_to_perturb(input) wrapped_fn = _with_prepare_inputs(fn, inputs, input_idx, input_to_perturb, True) @@ -430,18 +435,28 @@ def _get_numerical_jvp_wrt_specific_input(fn, input_idx, inputs, outputs, u, eps return _compute_numerical_jvps_wrt_specific_input(jvp_fn, u, input.is_complex(), is_forward_ad) -def _get_numerical_vJu(fn, inputs, inp_indices, outputs, all_u, all_v, eps, is_forward_ad): +def _get_numerical_vJu(fn, inputs, inp_indices, func_out, all_u, all_v, eps, is_forward_ad): # Note that all_v can also be None, in that case, this function only computes Ju. reduced_jacobians: List[List[torch.Tensor]] = [] for i, (inp_idx, u) in enumerate(zip(inp_indices, all_u)): - all_Ju = _get_numerical_jvp_wrt_specific_input(fn, inp_idx, inputs, outputs, u, eps, is_forward_ad) + all_Ju = _get_numerical_jvp_wrt_specific_input(fn, inp_idx, inputs, u, eps, is_forward_ad) + # Filter out the Ju for non floating point outputs + filtered_Ju = [] + func_out = _as_tuple(func_out) + assert len(all_Ju) == len(func_out) + for Ju, output in zip(all_Ju, func_out): + if _is_float_or_complex_tensor(output): + filtered_Ju.append(Ju) + else: + # TODO: handle the other Ju + pass if all_v is not None: jacobian_scalars: List[torch.Tensor] = [] - for v, Ju in zip(all_v, all_Ju): + for v, Ju in zip(all_v, filtered_Ju): jacobian_scalars.append(_dot_with_type_promotion(v, Ju)) reduced_jacobians.append(jacobian_scalars) else: - reduced_jacobians.append(all_Ju) + reduced_jacobians.append(filtered_Ju) return reduced_jacobians @@ -678,7 +693,7 @@ def _check_no_differentiable_outputs(func, inputs, func_out, eps) -> bool: def _check_no_differentiable_outputs_fast(func, func_out, all_inputs, inputs_indices, all_u, eps, nondet_tol): for inp_idx, u in zip(inputs_indices, all_u): - jvps = _get_numerical_jvp_wrt_specific_input(func, inp_idx, all_inputs, _as_tuple(func_out), u, eps) + jvps = _get_numerical_jvp_wrt_specific_input(func, inp_idx, all_inputs, u, eps) for jvp in jvps: if jvp.numel() == 0: continue @@ -959,13 +974,14 @@ def _gradcheck_real_imag(gradcheck_fn, func, func_out, tupled_inputs, outputs, e def _slow_gradcheck(func, func_out, tupled_inputs, outputs, eps, rtol, atol, check_grad_dtypes, nondet_tol, *, use_forward_ad=False, complex_indices=None, test_imag=False): + func_out = _as_tuple(func_out) if not outputs: - return _check_no_differentiable_outputs(func, tupled_inputs, _as_tuple(func_out), eps) + return _check_no_differentiable_outputs(func, tupled_inputs, func_out, eps) numerical = _transpose(_get_numerical_jacobian(func, tupled_inputs, outputs, eps=eps, is_forward_ad=use_forward_ad)) if use_forward_ad: - analytical_forward = _get_analytical_jacobian_forward_ad(func, tupled_inputs, outputs, check_grad_dtypes=check_grad_dtypes) + analytical_forward = _get_analytical_jacobian_forward_ad(func, tupled_inputs, func_out, check_grad_dtypes=check_grad_dtypes) for i, n_per_out in enumerate(numerical): for j, n in enumerate(n_per_out): @@ -1140,11 +1156,11 @@ def _fast_gradcheck(func, func_out, inputs, outputs, eps, rtol, inp_tensors_idx, inp_tensors = _get_inp_tensors(inputs) all_v, all_u, all_u_dense = _make_vectors(inp_tensors, outputs, use_forward_ad=use_forward_ad) - numerical_vJu = _get_numerical_vJu(func, inputs, inp_tensors_idx, outputs, all_u, all_v, eps, is_forward_ad=use_forward_ad) + numerical_vJu = _get_numerical_vJu(func, inputs, inp_tensors_idx, func_out, all_u, all_v, eps, is_forward_ad=use_forward_ad) if use_forward_ad: assert all_v is None - analytical_vJu = _get_analytical_jacobian_forward_ad(func, inputs, outputs, all_u=all_u, - check_grad_dtypes=check_grad_dtypes) + analytical_vJu = _get_analytical_jacobian_forward_ad(func, inputs, _as_tuple(func_out), + all_u=all_u, check_grad_dtypes=check_grad_dtypes) else: if not outputs: _check_no_differentiable_outputs_fast(func, func_out, inputs, inp_tensors_idx, all_u, eps, nondet_tol) diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp index 8ad4f459f7f07..f4b57ccf2c337 100644 --- a/torch/csrc/autograd/FunctionsManual.cpp +++ b/torch/csrc/autograd/FunctionsManual.cpp @@ -917,6 +917,13 @@ Tensor evenly_distribute_backward(Tensor grad, const Tensor & input, const Tenso } } +Tensor evenly_read_jvp(const Tensor& fw_grad, const Tensor & input, const Tensor & value) { + auto mask = (input == value); + auto count = mask.sum(); + auto grad_output = fw_grad / count; + return at::sum(mask * grad_output); +} + static Tensor var_backward(const Tensor & grad, const Tensor & self, int64_t correction) { // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-avoid-magic-numbers,cppcoreguidelines-narrowing-conversions) return (2.0 / (self.numel() - correction)) * grad * (self - self.mean()); @@ -3564,6 +3571,18 @@ Tensor cumprod_jvp(Tensor self_t, Tensor self_p, Tensor result, int dim) { } } +Tensor gather_with_keepdimed_indices(const Tensor& input, int64_t dim, const Tensor& indices, bool keepdim) { + auto full_indices = indices; + if (!keepdim) { + full_indices = indices.unsqueeze(dim); + } + auto out_fw_grad = at::gather(input, dim, full_indices); + if (!keepdim) { + out_fw_grad = out_fw_grad.squeeze(dim); + } + + return out_fw_grad; +} } // namespace details } // namespace generated diff --git a/torch/csrc/autograd/FunctionsManual.h b/torch/csrc/autograd/FunctionsManual.h index 7f365c9ce5db0..94d86bbd55cfb 100644 --- a/torch/csrc/autograd/FunctionsManual.h +++ b/torch/csrc/autograd/FunctionsManual.h @@ -255,6 +255,8 @@ Tensor lu_unpack_backward( Tensor cat_jvp(at::TensorList tensors, int64_t dim); Tensor cumprod_jvp(Tensor self_t, Tensor self_p, Tensor result, int dim); +Tensor gather_with_keepdimed_indices(const Tensor& input, int64_t dim, const Tensor& indices, bool keepdim); +Tensor evenly_read_jvp(const Tensor& fw_grad, const Tensor & input, const Tensor & value); } // namespace details } // namespace generated diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index ecac834845764..1cda7f822db50 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -4527,6 +4527,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), assert_autodiffed=True, sample_inputs_func=sample_inputs_dot_vdot, + supports_forward_ad=True, skips=( # some test samples works for ROCM backward but not all SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', @@ -4536,6 +4537,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.float16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), sample_inputs_func=sample_inputs_dot_vdot, + supports_forward_ad=True, skips=( # some test samples works for ROCM backward but not all SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', @@ -4728,6 +4730,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('broadcast_to', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_broadcast_to), UnaryUfuncInfo('bitwise_not', ref=np.bitwise_not, @@ -4946,17 +4949,20 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCPU=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), sample_inputs_func=partial(sample_inputs_cumulative_ops, supports_dtype_kwargs=False), + supports_forward_ad=True, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), OpInfo('cummin', dtypesIfCPU=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), sample_inputs_func=partial(sample_inputs_cumulative_ops, supports_dtype_kwargs=False), + supports_forward_ad=True, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), UnaryUfuncInfo('deg2rad', ref=np.radians, decorators=(precisionOverride({torch.bfloat16: 7e-1, torch.float16: 7e-1}),), dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/pull/51283#issuecomment-770614273 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', @@ -4972,6 +4978,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): variant_test_name='no_rounding_mode', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), sample_inputs_func=partial(sample_inputs_binary_pwise, rhs_exclude_zero=True), + supports_forward_ad=True, assert_autodiffed=True), OpInfo('div', aliases=('divide',), @@ -4979,6 +4986,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.half, torch.bfloat16), sample_inputs_func=partial(sample_inputs_binary_pwise, extra_kwargs={ "rounding_mode": 'trunc'}, rhs_exclude_zero=True), + supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/issues/59174 SkipInfo('TestCommon', 'test_variant_consistency_jit'), @@ -4990,6 +4998,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.half, torch.bfloat16), sample_inputs_func=partial(sample_inputs_binary_pwise, extra_kwargs={ "rounding_mode": 'floor'}, rhs_exclude_zero=True), + supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/issues/59174 SkipInfo('TestCommon', 'test_variant_consistency_jit'), @@ -4997,6 +5006,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True), OpInfo('true_divide', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=partial(sample_inputs_binary_pwise, rhs_exclude_zero=True)), UnaryUfuncInfo('exp', ref=np_unary_ufunc_integer_promotion_wrapper(np.exp), @@ -5017,6 +5027,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), ), assert_autodiffed=True, + supports_forward_ad=True, safe_casts_outputs=True), OpInfo('expand', op=lambda self, shape: self.expand(shape), @@ -5025,10 +5036,12 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Because expand does not have a function variant. SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + supports_forward_ad=True, supports_out=False), OpInfo('expand_as', op=lambda self, other: self.expand_as(other), dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=sample_inputs_expand_as, skips=( # Because expand_as does not have a function variant. @@ -5055,10 +5068,12 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('fmax', op=torch.fmax, dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_binary,), OpInfo('fmin', op=torch.fmin, dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_binary,), OpInfo('fmod', dtypes=all_types_and(torch.float16), @@ -5083,6 +5098,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=floating_types_and(torch.bfloat16, torch.float16), dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), assert_autodiffed=True, + supports_forward_ad=True, # Reference for disabling extremals # https://github.com/pytorch/pytorch/issues/51948 handles_extremals=False), @@ -5171,6 +5187,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): ref=np.floor, dtypes=floating_types_and(torch.bfloat16), dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, assert_autodiffed=True), OpInfo('flip', op=torch.flip, @@ -5253,6 +5270,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): # skip testing torch.frexp as it is not supported by ROCm platform yet decorators=[skipCUDAIfRocm], supports_out=False, + supports_forward_ad=True, skips=( # skips below tests as torch.frexp returns tuple-like (mantissa, exponent) as outputs, # while theses tests currently requires output to a single tensor. @@ -5321,6 +5339,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('kthvalue', dtypes=all_types(), dtypesIfCUDA=all_types_and(torch.float16), + supports_forward_ad=True, sample_inputs_func=sample_inputs_kthvalue), OpInfo('le', aliases=('less_equal',), @@ -5503,6 +5522,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), assert_autodiffed=True, safe_casts_outputs=True, + supports_forward_ad=True, decorators=(precisionOverride({torch.bfloat16: 5e-2}),), skips=( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal', @@ -5517,6 +5537,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True, dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), safe_casts_outputs=True, + supports_forward_ad=True, skips=( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], @@ -5529,6 +5550,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), decorators=(precisionOverride({torch.bfloat16: 1e-1}),), safe_casts_outputs=True, + supports_forward_ad=True, assert_autodiffed=True), UnaryUfuncInfo('log2', ref=np.log2, @@ -5537,6 +5559,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), assert_autodiffed=True, safe_casts_outputs=True, + supports_forward_ad=True, decorators=(precisionOverride({torch.bfloat16: 1e-1}),), skips=( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal', @@ -5548,6 +5571,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=floating_types(), dtypesIfCUDA=floating_types_and(torch.bfloat16), dtypesIfROCM=floating_types_and(torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=lambda op_info, device, dtype, requires_grad=False, **kwargs: (SampleInput(make_tensor((S, S), device, dtype, requires_grad=requires_grad), args=(make_tensor((S, S), device, dtype, requires_grad=requires_grad),)),)), @@ -5555,6 +5579,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=floating_types(), dtypesIfCUDA=floating_types_and(torch.bfloat16), dtypesIfROCM=floating_types_and(torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=lambda op_info, device, dtype, requires_grad=False, **kwargs: (SampleInput(make_tensor((S, S), device, dtype, requires_grad=requires_grad), args=(make_tensor((S, S), device, dtype, requires_grad=requires_grad),)),)), @@ -5618,13 +5643,16 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('masked_fill', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), sample_inputs_func=sample_inputs_masked_fill, + supports_forward_ad=True, supports_out=False), OpInfo('masked_scatter', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), sample_inputs_func=sample_inputs_masked_scatter, + supports_forward_ad=True, supports_out=False), OpInfo('masked_select', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=sample_inputs_masked_select), OpInfo('matrix_exp', dtypesIfCPU=floating_and_complex_types_and(torch.bfloat16), @@ -5667,12 +5695,14 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): variant_test_name='binary', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_max_min_binary, + supports_forward_ad=True, assert_autodiffed=True,), OpInfo('max', op=torch.max, variant_test_name='reduction_with_dim', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_max_min_reduction_with_dim, + supports_forward_ad=True, skips=( # max does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'),)), @@ -5681,6 +5711,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): variant_test_name='reduction_no_dim', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_reduction_no_dim,), OpInfo('median', dtypes=all_types(), @@ -5729,12 +5760,14 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): variant_test_name='binary', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_max_min_binary, + supports_forward_ad=True, assert_autodiffed=True,), OpInfo('min', op=torch.min, variant_test_name='reduction_with_dim', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_max_min_reduction_with_dim, + supports_forward_ad=True, skips=( # min does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), @@ -5744,6 +5777,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): variant_test_name='reduction_no_dim', dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_reduction_no_dim,), OpInfo('sum', dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), @@ -5758,6 +5792,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('mean', dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), assert_autodiffed=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_reduction_wrapper(supports_multiple_dims=True), # Need to skip out test because one of the overload for mean does not support it # TODO(@heitorschueroff) fix this when implementing ReductionInfo @@ -5771,10 +5806,12 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('maximum', op=torch.maximum, dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_binary,), OpInfo('minimum', op=torch.minimum, dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, sample_inputs_func=sample_inputs_max_min_binary,), OpInfo('nn.functional.hardswish', aten_name="hardswish", @@ -5783,6 +5820,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): sample_inputs_func=sample_inputs_hardswish, dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), supports_gradgrad=False, + supports_forward_ad=True, supports_out=False, autodiff_nonfusible_nodes=["aten::hardswish"]), OpInfo('nn.functional.leaky_relu', @@ -5854,6 +5892,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCPU=all_types_and_complex_and(torch.float16, torch.bfloat16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), assert_autodiffed=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_mm, skips=( # mm does not correctly warn when resizing out= inputs @@ -5865,6 +5904,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('mode', op=torch.mode, dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, sample_inputs_func=sample_inputs_mode,), MvlGammaInfo(variant_test_name='mvlgamma_p_1', domain=(1e-4, float('inf')), @@ -6289,12 +6329,14 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): ref=np_unary_ufunc_integer_promotion_wrapper(np.exp2), dtypes=all_types_and(torch.bool, torch.half), dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, safe_casts_outputs=True), UnaryUfuncInfo('expm1', aliases=('special.expm1', ), ref=np_unary_ufunc_integer_promotion_wrapper(np.expm1), dtypes=all_types_and(torch.bool, torch.bfloat16), dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, safe_casts_outputs=True, assert_autodiffed=True, skips=( @@ -6387,6 +6429,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), dtypesIfROCM=floating_and_complex_types_and(torch.half, torch.bfloat16), sample_inputs_func=sample_inputs_lerp, + supports_forward_ad=True, assert_autodiffed=True), OpInfo('linalg.inv', aten_name='linalg_inv', @@ -6523,6 +6566,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), safe_casts_outputs=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_polygamma, skips=( # Probably related to the way the function is @@ -6548,6 +6592,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), safe_casts_outputs=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_polygamma, skips=( # Redundant tests @@ -6567,6 +6612,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), safe_casts_outputs=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_polygamma, skips=( # Redundant tests @@ -6587,6 +6633,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), safe_casts_outputs=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_polygamma, skips=( # Redundant tests @@ -6608,6 +6655,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), safe_casts_outputs=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_polygamma, skips=( # Redundant tests @@ -6672,26 +6720,31 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), sample_inputs_func=sample_inputs_gather, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_forward_ad=True, ), OpInfo('index_fill', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_inplace_autograd=False, skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),), supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_index_fill), OpInfo('index_copy', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_inplace_autograd=False, supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_index_copy, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), OpInfo('index_select', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), sample_inputs_func=sample_inputs_index_select, + supports_forward_ad=True, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), OpInfo('index_add', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_out=False, + supports_forward_ad=True, sample_inputs_func=sample_inputs_index_add, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), OpInfo('__getitem__', @@ -6705,6 +6758,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_out=False, supports_inplace_autograd=True, + supports_forward_ad=True, sample_inputs_func=sample_inputs_index_put, skips=( SkipInfo('TestCommon', 'test_variant_consistency_jit'), @@ -6752,6 +6806,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=floating_types(), dtypesIfCPU=floating_types_and(torch.bfloat16), dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, sample_inputs_func=sample_inputs_hypot, ), OpInfo('vstack', @@ -6891,6 +6946,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCPU=all_types_and(torch.bool, torch.half, torch.bfloat16), dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), supports_inplace_autograd=True, + supports_forward_ad=True, safe_casts_outputs=True, sample_inputs_func=sample_inputs_xlogy), OpInfo('zero_', @@ -6908,6 +6964,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): aten_name='special_xlog1py', dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), safe_casts_outputs=True, + supports_forward_ad=True, skips=( SkipInfo('TestOpInfo', 'test_supported_backward', device_type='cpu', dtypes=[torch.float16]), @@ -7017,6 +7074,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): decorators=(precisionOverride({torch.float16: 5e-1}),), dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), + supports_forward_ad=True, safe_casts_outputs=True), UnaryUfuncInfo('special.entr', ref=scipy.special.entr if TEST_SCIPY else _NOTHING, @@ -7077,6 +7135,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and(torch.bool, torch.half), # "digamma" not implemented for 'BFloat16' backward_dtypesIfCPU=all_types_and(torch.bool), + supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/pull/50140#discussion_r552615345 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal', From a3db8e0a26f9834d3d4a2e829d632c902fe95428 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Thu, 10 Jun 2021 19:50:22 -0700 Subject: [PATCH 027/305] [docs] Add torch.package documentation preamble (#59491) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59491 **Summary** This commit adds a preamble to the `torch.package` documentation page that explains briefly what `torch.package` is. **Test Plan** Continous integration. Captura de Pantalla 2021-06-04 a la(s) 3 57 01 p m Test Plan: Imported from OSS Reviewed By: Lilyjjo Differential Revision: D29050630 Pulled By: SplitInfinity fbshipit-source-id: 70a3fd43f076751c6ea83be3ead291686c641158 --- docs/source/package.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/source/package.rst b/docs/source/package.rst index 38fb5fecc1c99..6582806db97c7 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -7,6 +7,14 @@ torch.package This module is experimental and has not yet been publicly released. +``torch.package`` adds support for creating hermetic packages containing arbitrary +PyTorch code. These packages can be saved, shared, used to load and execute models +at a later date or on a different machine, and can even be deployed to production using +``torch::deploy``. + +This document contains tutorials, how-to guides, explanations, and an API reference that +will help you learn more about ``torch.package`` and how to use it. + API Reference ------------- .. autoclass:: torch.package.PackagingError From 6a18ca7a07724ccae4b8c96426becdb046ab2817 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Thu, 10 Jun 2021 19:50:22 -0700 Subject: [PATCH 028/305] [docs] Add tutorials section to torch.package docs (#59499) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59499 **Summary** This commit adds a tutorials section to the torch.package docs. **Test Plan** Continuous integration. Captura de Pantalla 2021-06-04 a la(s) 5 10 31 p m Test Plan: Imported from OSS Reviewed By: Lilyjjo Differential Revision: D29050628 Pulled By: SplitInfinity fbshipit-source-id: c17ab0100a9d63e7af8da7a618143cedbd0a5872 --- docs/source/package.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/source/package.rst b/docs/source/package.rst index 6582806db97c7..b80798e514b87 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -15,6 +15,15 @@ at a later date or on a different machine, and can even be deployed to productio This document contains tutorials, how-to guides, explanations, and an API reference that will help you learn more about ``torch.package`` and how to use it. +Tutorials +--------- +Packaging your first model +^^^^^^^^^^^^^^^^^^^^^^^^^^ +A tutorial that guides you through packaging and unpackaging a simple model is available +`on Colab `_. +After completing this exercise, you will be familiar with the basic API for creating and using +Torch packages. + API Reference ------------- .. autoclass:: torch.package.PackagingError From 062dde7285f9bbcc028e063b11c3e724580a3df6 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Thu, 10 Jun 2021 19:50:22 -0700 Subject: [PATCH 029/305] [docs] Add "how do I" section to torch.package docs (#59503) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59503 **Summary** This commit adds a "how do I..." section to the `torch.package` documentation. This section contains short guides about how to solve real-world problems that frequently recur while using `torch.package`. **Test Plan** Continuous integration. Captura de Pantalla 2021-06-04 a la(s) 9 19 54 p m Test Plan: Imported from OSS Reviewed By: Lilyjjo Differential Revision: D29050629 Pulled By: SplitInfinity fbshipit-source-id: 2b7800732e0a3c1c947f110c05562aed5174a87f --- docs/source/package.rst | 427 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 427 insertions(+) diff --git a/docs/source/package.rst b/docs/source/package.rst index b80798e514b87..531dfed44c672 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -24,6 +24,433 @@ A tutorial that guides you through packaging and unpackaging a simple model is a After completing this exercise, you will be familiar with the basic API for creating and using Torch packages. +How do I... +----------- +See what is inside a package? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Treat the package like a ZIP archive +"""""""""""""""""""""""""""""""""""" +The container format for a ``torch.package`` is ZIP, so any tools that work with standard ZIP files should +work for exploring the contents. Some common ways to interact with ZIP files: + +* ``unzip my_package.pt`` will unzip the ``torch.package`` archive to disk, where you can freely inspect its contents. + + +:: + + $ unzip my_package.pt && tree my_package + my_package + ├── .data + │ ├── 94304870911616.storage + │ ├── 94304900784016.storage + │ ├── extern_modules + │ └── version + ├── models + │ └── model_1.pkl + └── torchvision + └── models + ├── resnet.py + └── utils.py + ~ cd my_package && cat torchvision/models/resnet.py + ... + + +* The Python ``zipfile`` module provides a standard way to read and write ZIP archive contents. + + +:: + + from zipfile import ZipFile + with ZipFile("my_package.pt") as myzip: + file_bytes = myzip.read("torchvision/models/resnet.py") + # edit file_bytes in some way + myzip.writestr("torchvision/models/resnet.py", new_file_bytes) + + +* vim has the ability to natively read ZIP archives. You can even edit files and :``write`` them back into the archive! + + +:: + + # add this to your .vimrc to treat `*.pt` files as zip files + au BufReadCmd *.pt call zip#Browse(expand("")) + + ~ vi my_package.pt + + +Use the ``file_structure()`` API +"""""""""""""""""""""""""""""""" +:class:`PackageImporter` and :class:`PackageExporter` provide a ``file_structure()`` method, which will return a printable +and queryable ``Folder`` object. The ``Folder`` object is a simple directory structure that you can use to explore the +current contents of a ``torch.package``. + +The ``Folder`` object itself is directly printable and will print out a file tree representation. To filter what is returned, +use the glob-style ``include`` and ``exclude`` filtering arguments. + + +:: + + with PackageExporter('my_package.pt', verbose=False) as pe: + pe.save_pickle('models', 'model_1.pkl', mod) + # can limit printed items with include/exclude args + print(pe.file_structure(include=["**/utils.py", "**/*.pkl"], exclude="**/*.storages")) + + importer = PackageImporter('my_package.pt') + print(importer.file_structure()) # will print out all files + + +Output: + + +:: + + # filtered with glob pattern: + # include=["**/utils.py", "**/*.pkl"], exclude="**/*.storages" + ─── my_package.pt + ├── models + │ └── model_1.pkl + └── torchvision + └── models + └── utils.py + + # all files + ─── my_package.pt + ├── .data + │ ├── 94304870911616.storage + │ ├── 94304900784016.storage + │ ├── extern_modules + │ └── version + ├── models + │ └── model_1.pkl + └── torchvision + └── models + ├── resnet.py + └── utils.py + + +You can also query ``Folder`` objects with the ``has_file()`` method. + + +:: + + exporter_file_structure = exporter.file_structure() + found: bool = exporter_file_structure.has_file("package_a/subpackage.py") + + +Include arbitrary resources with my package and access them later? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:class:`PackageExporter` exposes three methods, ``save_pickle``, ``save_text`` and ``save_binary`` that allow you to save +Python objects, text, and binary data to a package. + + +:: + + with torch.PackageExporter("package.pt") as exporter: + # Pickles the object and saves to `my_resources/tens.pkl` in the archive. + exporter.save_pickle("my_resources", "tensor.pkl", torch.randn(4)) + exporter.save_text("config_stuff", "words.txt", "a sample string") + exporter.save_binary("raw_data", "binary", my_bytes) + + +:class:`PackageImporter` exposes complementary methods named ``load_pickle``, ``load_text`` and ``load_binary`` that allow you to load +Python objects, text and binary data from a package. + + +:: + + importer = torch.PackageImporter("package.pt") + my_tensor = importer.load_pickle("my_resources", "tensor.pkl") + text = importer.load_text("config_stuff", "words.txt") + binary = importer.load_binary("raw_data", "binary") + + +Customize how a class is packaged? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +``torch.package`` allows for the customization of how classes are packaged. This behavior is accessed through defining the method +``__reduce_package__`` on a class and by defining a corresponding de-packaging function. This is similar to defining ``__reduce__`` for +Python’s normal pickling process. + +Steps: + +1. Define the method ``__reduce_package__(self, exporter: PackageExporter)`` on the target class. This method should do the work to save the class instance inside of the package, and should return a tuple of the corresponding de-packaging function with the arguments needed to invoke the de-packaging function. This method is called by the ``PackageExporter`` when it encounters an instance of the target class. +2. Define a de-packaging function for the class. This de-packaging function should do the work to reconstruct and return an instance of the class. The function signature’s first parameter should be a ``PackageImporter`` instance, and the rest of the parameters are user defined. + + +:: + + # foo.py [Example of customizing how class Foo is packaged] + from torch.package import PackageExporter, PackageImporter + import time + + + class Foo: + def __init__(self, my_string: str): + super().__init__() + self.my_string = my_string + self.time_imported = 0 + self.time_exported = 0 + + def __reduce_package__(self, exporter: PackageExporter): + """ + Called by ``torch.package.PackageExporter``'s Pickler's ``persistent_id`` when + saving an instance of this object. This method should do the work to save this + object inside of the ``torch.package`` archive. + + Returns function w/ arguments to load the object from a + ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function. + """ + + # use this pattern to ensure no naming conflicts with normal dependencies, + # anything saved under this module name shouldn't conflict with other + # items in the package + generated_module_name = f"foo-generated._{exporter.get_unique_id()}" + exporter.save_text( + generated_module_name, + "foo.txt", + self.my_string + ", with exporter modification!", + ) + time_exported = time.clock_gettime(1) + + # returns de-packaging function w/ arguments to invoke with + return (unpackage_foo, (generated_module_name, time_exported,)) + + + def unpackage_foo( + importer: PackageImporter, generated_module_name: str, time_exported: float + ) -> Foo: + """ + Called by ``torch.package.PackageImporter``'s Pickler's ``persistent_load`` function + when depickling a Foo object. + Performs work of loading and returning a Foo instance from a ``torch.package`` archive. + """ + time_imported = time.clock_gettime(1) + foo = Foo(importer.load_text(generated_module_name, "foo.txt")) + foo.time_imported = time_imported + foo.time_exported = time_exported + return foo + + +:: + + # example of saving instances of class Foo + + import torch + from torch.package import PackageImporter, PackageExporter + import foo + + foo_1 = foo.Foo("foo_1 initial string") + foo_2 = foo.Foo("foo_2 initial string") + with PackageExporter('foo_package.pt', verbose=False) as pe: + # save as normal, no extra work necessary + pe.save_pickle('foo_collection', 'foo1.pkl', foo_1) + pe.save_pickle('foo_collection', 'foo2.pkl', foo_2) + print(pe.file_structure()) + + pi = PackageImporter('foo_package.pt') + imported_foo = pi.load_pickle('foo_collection', 'foo1.pkl') + print(f"foo_1 string: '{imported_foo.my_string}'") + print(f"foo_1 export time: {imported_foo.time_exported}") + print(f"foo_1 import time: {imported_foo.time_imported}") + + +:: + + # output of running above script + ─── foo_package + ├── foo-generated + │ ├── _0 + │ │ └── foo.txt + │ └── _1 + │ └── foo.txt + ├── foo_collection + │ ├── foo1.pkl + │ └── foo2.pkl + └── foo.py + + foo_1 string: 'foo_1 initial string, with reduction modification!' + foo_1 export time: 9857706.650140837 + foo_1 import time: 9857706.652698385 + + +Test in my source code whether or not it is executing inside a package? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +A :class:`PackageImporter` will add the attribute ``__torch_package__`` to every module that it initializes. Your code can check for the +presence of this attribute to determine whether it is executing in a packaged context or not. + + +:: + + # In foo/bar.py: + + if "__torch_package__" in dir(): # true if the code is being loaded from a package + def is_in_package(): + return True + + UserException = Exception + else: + def is_in_package(): + return False + + UserException = UnpackageableException + + +Now, the code will behave differently depending on whether it’s imported normally through your Python environment or imported from a +``torch.package``. + + +:: + + from foo.bar import is_in_package + + print(is_in_package()) # False + + loaded_module = PackageImporter(my_pacakge).import_module("foo.bar") + loaded_module.is_in_package() # True + + +*Warning*: in general, it’s bad practice to have code that behaves differently depending on whether it’s packaged or not. This can lead to +hard-to-debug issues that are sensitive to how you imported your code. If your package is intended to be heavily used, consider restructuring +your code so that it behaves the same way no matter how it was loaded. + + +Patch code into a package? +^^^^^^^^^^^^^^^^^^^^^^^^^^ +:class:`PackageExporter` offers a ``save_source_string()`` method that allows one to save arbitrary Python source code to a module of your choosing. + + +:: + + with PackageExporter(f) as exporter: + # Save the my_module.foo available in your current Python environment. + exporter.save_module("my_module.foo") + + # This saves the provided string to my_module/foo.py in the package archive. + # It will override the my_module.foo that was previously saved. + exporter.save_source_string("my_module.foo", textwrap.dedent( + """\ + def my_function(): + print('hello world') + """ + )) + + # If you want to treat my_module.bar as a package + # (e.g. save to `my_module/bar/__init__.py` instead of `my_module/bar.py) + # pass is_package=True, + exporter.save_source_string("my_module.bar", + "def foo(): print('hello')\n", + is_package=True) + + importer = PackageImporter(f) + importer.import_module("my_module.foo").my_function() # prints 'hello world' + + +Access package contents from packaged code? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:class:`PackageImporter` implements the +`importlib.resources `_ +API for accessing resources from inside a package. + + +:: + + with PackageExporter(f) as exporter: + # saves text to one/a.txt in the archive + exporter.save_text("my_resource", "a.txt", "hello world!") + # saves the tensor to my_pickle/obj.pkl + exporter.save_pickle("my_pickle", "obj.pkl", torch.ones(2, 2)) + + # see below for module contents + exporter.save_module("foo") + exporter.save_module("bar") + + +The ``importlib.resources`` API allows access to resources from within packaged code. + + +:: + + # foo.py: + import importlib.resources + import my_resource + + # returns "hello world!" + def get_my_resource(): + return importlib.resources.read_text(my_resource, "a.txt") + + +Using ``importlib.resources`` is the recommended way to access package contents from within packaged code, since it complies +with the Python standard. However, it is also possible to access the parent :class:`PackageImporter` instance itself from within +packaged code. + + +:: + + # bar.py: + import torch_package_importer # this is the PackageImporter that imported this module. + + # Prints "hello world!", equivalient to importlib.resources.read_text + def get_my_resource(): + return torch_package_importer.load_text("my_resource", "a.txt") + + # You also do things that the importlib.resources API does not support, like loading + # a pickled object from the package. + def get_my_pickle(): + return torch_package_importer.load_pickle("my_pickle", "obj.pkl") + + +Distinguish between packaged code and non-packaged code? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +To tell if an object’s code is from a ``torch.package``, use the ``torch.package.is_from_package()`` function. +Note: if an object is from a package but its definition is from a module marked ``extern`` or from ``stdlib``, +this check will return ``False``. + + +:: + + importer = PackageImporter(f) + mod = importer.import_module('foo') + obj = importer.load_pickle('model', 'model.pkl') + txt = importer.load_text('text', 'my_test.txt') + + assert is_from_package(mod) + assert is_from_package(obj) + assert not is_from_package(txt) # str is from stdlib, so this will return False + + +Re-export an imported object? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +To re-export an object that was previously imported by a :class:`PackageImporter`, you must make the new :class:`PackageExporter` +aware of the original :class:`PackageImporter` so that it can find source code for your object’s dependencies. + + +:: + + importer = PackageImporter(f) + obj = importer.load_pickle("model", "model.pkl") + + # re-export obj in a new package + with PackageExporter(f2, importer=(importer, sys_importer)) as exporter: + exporter.save_pickle("model", "model.pkl", obj) + + +Package a TorchScript module? +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +To package a TorchScript model, use the same ``save_pickle`` and ``load_pickle`` APIs as you would with any other object. +Saving TorchScript objects that are attributes or submodules is supported as well with no extra work. + + +:: + + # save TorchScript just like any other object + with PackageExporter(file_name, verbose=True) as e: + e.save_pickle("res", "script_model.pkl", scripted_model) + e.save_pickle("res", "mixed_model.pkl", python_model_with_scripted_submodule) + # load as normal + importer = PackageImporter(file_name) + loaded_script = importer.load_pickle("res", "script_model.pkl") + loaded_mixed = importer.load_pickle("res", "mixed_model.pkl" + + API Reference ------------- .. autoclass:: torch.package.PackagingError From 0e222db0874948a3d0af240a3db7a3e4163a1b19 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Thu, 10 Jun 2021 19:50:22 -0700 Subject: [PATCH 030/305] [docs] Add explanation section to torch.package docs (#59833) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59833 **Summary** This commit adds an explanation section to the `torch.package` documentation. This section clarifies and illuminates various aspects of the internals of `torch.package` that might be of interest to users. **Test Plan** Continuous integration. Test Plan: Imported from OSS Reviewed By: Lilyjjo Differential Revision: D29050626 Pulled By: SplitInfinity fbshipit-source-id: 78e0cda00f69506ef2dfc52d6df63694b502269e --- docs/source/package.rst | 325 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 324 insertions(+), 1 deletion(-) diff --git a/docs/source/package.rst b/docs/source/package.rst index 531dfed44c672..99e6e75b531d1 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -308,7 +308,7 @@ Now, the code will behave differently depending on whether it’s imported norma loaded_module.is_in_package() # True -*Warning*: in general, it’s bad practice to have code that behaves differently depending on whether it’s packaged or not. This can lead to +**Warning**: in general, it’s bad practice to have code that behaves differently depending on whether it’s packaged or not. This can lead to hard-to-debug issues that are sensitive to how you imported your code. If your package is intended to be heavily used, consider restructuring your code so that it behaves the same way no matter how it was loaded. @@ -451,6 +451,329 @@ Saving TorchScript objects that are attributes or submodules is supported as wel loaded_mixed = importer.load_pickle("res", "mixed_model.pkl" +Explanation +----------- +``torch.package`` Format Overview +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +A ``torch.package`` file is a ZIP archive which conventionally uses the ``.pt`` extension. Inside the ZIP archive, there are two kinds of files: + +* Framework files, which are placed in the ``.data/``. +* User files, which is everything else. + +As an example, this is what a fully packaged ResNet model from ``torchvision`` looks like: + + +:: + + resnet + ├── .data # All framework-specific data is stored here. + │ │ # It's named to avoid conflicts with user-serialized code. + │ ├── 94286146172688.storage # tensor data + │ ├── 94286146172784.storage + │ ├── extern_modules # text file with names of extern modules (e.g. 'torch') + │ ├── version # version metadata + │ ├── ... + ├── model # the pickled model + │ └── model.pkl + └── torchvision # all code dependencies are captured as source files + └── models + ├── resnet.py + └── utils.py + + +Framework files +""""""""""""""" +The ``.data/`` directory is owned by torch.package, and its contents are considered to be a private implementation detail. +The ``torch.package`` format makes no guarantees about the contents of ``.data/``, but any changes made will be backward compatible +(that is, newer version of PyTorch will always be able to load older ``torch.packages``). + +Currently, the ``.data/`` directory contains the following items: + +* ``version``: a version number for the serialized format, so that the ``torch.package`` import infrastructures knows how to load this package. +* ``extern_modules``: a list of modules that are considered ``extern:class:`PackageImporter`. ``extern`` modules will be imported using the loading environment’s system importer. +* ``*.storage``: serialized tensor data. + + +:: + + .data + ├── 94286146172688.storage + ├── 94286146172784.storage + ├── extern_modules + ├── version + ├── ... + + +User files +"""""""""" +All other files in the archive were put there by a user. The layout is identical to a Python +`regular package `_. For a deeper dive in how Python packaging works, +please consult `this essay `_ (it’s slightly out of date, so double-check implementation details +with the `Python reference documentation `_). + + +:: + + + ├── model # the pickled model + │ └── model.pkl + ├── another_package + │ ├── __init__.py + │ ├── foo.txt # a resource file , see importlib.resources + │ └── ... + └── torchvision + └── models + ├── resnet.py # torchvision.models.resnet + └── utils.py # torchvision.models.utils + + +How ``torch.package`` finds your code's dependencies +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Analyzing an object's dependencies +"""""""""""""""""""""""""""""""""" +When you issue a ``save_pickle(obj, ...)`` call, :class:`PackageExporter` will pickle the object normally. Then, it uses the +``pickletools`` standard library module to parse the pickle bytecode. + +In a pickle, an object is saved along with a ``GLOBAL`` opcode that describes where to find the implementation of the object’s type, like: + + +:: + + GLOBAL 'torchvision.models.resnet Resnet` + + +The dependency resolver will gather up all ``GLOBAL`` ops and mark them as dependencies of your pickled object. +For more information about pickling and the pickle format, please consult `the Python docs `_. + +Analyzing a module's dependencies +""""""""""""""""""""""""""""""""" +When a Python module is identified as a dependency, ``torch.package`` walks the module’s python AST representation and looks for import statements with +full support for the standard forms: ``from x import y``, ``import z``, ``from w import v as u``, etc. When one of these import statements are +encountered, ``torch.package`` registers the imported modules as dependencies that are then themselves parsed in the same AST walking way. + +**Note**: AST parsing has limited support for the ``__import__(...)`` syntax and does not support ``importlib.import_module`` calls. In general, you should +not expect dynamic imports to be detected by ``torch.package``. + + +Dependency Management +^^^^^^^^^^^^^^^^^^^^^ +``torch.package`` automatically finds the Python modules that your code and objects depend on. This process is called dependency resolution. +For each module that the dependency resolver finds, you must specify an *action* to take. + +The allowed actions are: + +* ``intern``: put this module into the package. +* ``extern``: declare this module as an external dependency of the package. +* ``mock``: stub out this module. +* ``deny``: depending on this module will raise an error during package export. + +Finally, there is one more important action that is not technically part of ``torch.package``: + +* Refactoring: remove or change the dependencies in your code. + +Note that actions are only defined on entire Python modules. There is no way to package “just” a function or class from module and leave the rest out. +This is by design. Python does not offer clean boundaries between objects defined in a module. The only defined unit of dependency organization is a +module, so that’s what ``torch.package`` uses. + +Actions are applied to modules using patterns. Patterns can either be module names (``"foo.bar"``) or globs (like ``"foo.**"``). You associate a pattern +with an action using methods on :class:`PackageImporter`, e.g. + + +:: + + my_exporter.intern("torchvision.**") + my_exporter.extern("numpy") + + +If a module matches a pattern, the corresponding action is applied to it. For a given module, patterns will be checked in the order that they were defined, +and the first action will be taken. + + +``intern`` +"""""""""" +If a module is ``intern``-ed, it will be placed into the package. + +This action is your model code, or any related code you want to package. For example, if you are trying to package a ResNet from ``torchvision``, +you will need to ``intern`` the module torchvision.models.resnet. + +On package import, when your packaged code tries to import an ``intern``-ed module, PackageImporter will look inside your package for that module. +If it can’t find that module, an error will be raised. This ensures that each :class:`PackageImporter` is isolated from the loading environment—even +if you have ``my_interned_module`` available in both your package and the loading environment, :class:`PackageImporter` will only use the version in your +package. + +**Note**: Only Python source modules can be ``intern``-ed. Other kinds of modules, like C extension modules and bytecode modules, will raise an error if +you attempt to ``intern`` them. These kinds of modules need to be ``mock``-ed or ``extern``-ed. + + +``extern`` +"""""""""" +If a module is ``extern``-ed, it will not be packaged. Instead, it will be added to a list of external dependencies for this package. You can find this +list on ``package_exporter.extern_modules``. + +On package import, when time packaged code tries to import an ``extern``-ed module, :class:`PackageImporter` will use the default Python importer to find +that module, as if you did ``importlib.import_module("my_externed_module")``. If it can’t find that module, an error will be raised. + +In this way, you can depend on third-party libraries like ``numpy`` and ``scipy`` from within your package without having to package them too. + +**Warning**: If any external library changes in a backwards-incompatible way, your package may fail to load. If you need long-term reproducibility +for your package, try to limit your use of ``extern``. + + +``mock`` +"""""""" +If a module is ``mock``-ed, it will not be packaged. Instead a stub module will be packaged in its place. The stub module will allow you to retrieve +objects from it (so that ``from my_mocked_module import foo`` will not error), but any use of that object will raise a ``NotImplementedError``. + +``mock`` should be used for code that you “know” will not be needed in the loaded package, but you still want to available for use in non-packaged contents. +For example, initialization/configuration code, or code only used for debugging/training. + +**Warning**: In general, ``mock`` should be used as a last resort. It introduces behavioral differences between packaged code and non-packaged code, +which may lead to later confusion. Prefer instead to refactor your code to remove unwanted dependencies. + + +Refactoring +""""""""""" +The best way to manage dependencies is to not have dependencies at all! Often, code can be refactored to remove unnecessary dependencies. Here are some +guidelines for writing code with clean dependencies (which are also generally good practices!): + +**Include only what you use**. Do not leave unused imports in our code. The dependency resolver is not smart enough to tell that they are indeed unused, +and will try to process them. + +**Qualify your imports**. For example, instead of writing import foo and later using ``foo.bar.baz``, prefer to write ``from foo.bar import baz``. This more +precisely specifies your real dependency (``foo.bar``) and lets the dependency resolver know you don’t need all of ``foo``. + +**Split up large files with unrelated functionality into smaller ones**. If your ``utils`` module contains a hodge-podge of unrelated functionality, any module +that depends on ``utils`` will need to pull in lots of unrelated dependencies, even if you only needed a small part of it. Prefer instead to define +single-purpose modules that can be packaged independently of one another. + + +Patterns +"""""""" +Patterns allow you to specify groups of modules with a convenient syntax. The syntax and behavior of patterns follows the Bazel/Buck +`glob() `_. + +A module that we are trying to match against a pattern is called a candidate. A candidate is composed of a list of segments separated by a +separator string, e.g. ``foo.bar.baz``. + +A pattern contains one or more segments. Segments can be: + +* A literal string (e.g. ``foo``), which matches exactly. +* A string containing a wildcard (e.g. ``torch``, or ``foo*baz*``). The wildcard matches any string, including the empty string. +* A double wildcard (``**``). This matches against zero or more complete segments. + +Examples: + +* ``torch.**``: matches ``torch`` and all its submodules, e.g. ``torch.nn`` and ``torch.nn.functional``. +* ``torch.*``: matches ``torch.nn`` or ``torch.functional``, but not ``torch.nn.functional`` or ``torch`` +* ``torch*.**``: matches ``torch``, ``torchvision``, and all of their submodules + +When specifying actions, you can pass multiple patterns, e.g. + + +:: + + exporter.intern(["torchvision.models.**", "torchvision.utils.**"]) + + +A module will match against this action if it matches any of the patterns. + +You can also specify patterns to exlcude, e.g. + + +:: + + exporter.mock("**", exclude=["torchvision.**"]) + + +A module will not match against this action if it matches any of the exclude patterns. In this example, we are mocking all modules except +``torchvision`` and its submodules. + +When a module could potentially match against multiple actions, the first action defined will be taken. + + +``torch.package`` sharp edges +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Avoid global state in your modules +"""""""""""""""""""""""""""""""""" +Python makes it really easy to bind objects and run code at module-level scope. This is generally fine—after all, functions and classes are bound to +names this way. However, things become more complicated when you define an object at module scope with the intention of mutating it, introducing mutable +global state. + +Mutable global state is quite useful—it can reduce boilerplate, allow for open registration into tables, etc. But unless employed very carefully, it can +cause complications when used with ``torch.package``. + +Every :class:`PackageImporter` creates an independent environment for its contents. This is nice because it means we load multiple packages and ensure +they are isolated from each other, but when modules are written in a way that assumes shared mutable global state, this behavior can create hard-to-debug +errors. + +Types are not shared between packages and the loading environment +""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""" +Any class that you import from a :class:`PackageImporter` will be a version of the class specific to that importer. For example: + + +:: + + from foo import MyClass + + my_class_instance = MyClass() + + with PackageExporter(f) as exporter: + exporter.save_module("foo") + + importer = PackageImporter(f) + imported_MyClass = importer.import_module("foo").MyClass + + assert isinstance(my_class_instance, MyClass) # works + assert isinstance(my_class_instance, imported_MyClass) # ERROR! + + +In this example, ``MyClass`` and ``import_MyClass`` are *not the same type*. In this specific example, ``MyClass`` and ``import_MyClass`` have exactly the +same implementation, so you might thing it’s okay to consider them the same class. But consider the situation where ``import_MyClass`` is coming from an +older package with an entirely different implementation of ``MyClass`` — in that case, it’s unsafe to consider them the same class. + +Under the hood, each importer has a prefix that allows it to uniquely identify classes: + + +:: + + print(MyClass.__name__) # prints "foo.MyClass" + print(imported_MyClass.__name__) # prints .foo.MyClass + + +That means you should not expect ``isinstance`` checks to work when one of the arguments if from a package and the other is not. If you need this +functionality, consider the following options: + +* Doing duck typing (just using the class instead of explicitly checking that it is of a given type). +* Make the typing relationship an explicit part of the class contract. For example, you can add an attribute tag ``self.handler = "handle_me_this_way"`` and have client code check for the value of ``handler`` instead of checking the type directly. + + +How ``torch.package`` keeps packages isolated from each other +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Each :class:`PackageImporter` instance creates an independent, isolated environment for its modules and objects. Modules in a package can only import +other packaged modules, or modules marked ``extern``. If you use multiple :class:`PackageImporter` instances to load a single package, you will get +multiple independent environments that do not interact. + +This is achieved by extending Python’s import infrastructure with a custom importer. :class:`PackageImporter` provides the same core API as the +``importlib`` importer; namely, it implements the ``import_module`` and ``__import__`` methods. + +When you invoke :meth:`PackageImporter.import_module`, :class:`PackageImporter` will construct and return a new module, much as the system importer does. +However, :class:`PackageImporter` patches the returned module to use ``self`` (i.e. that :class:`PackageImporter` instance) to fulfill future import +requests by looking in the package rather than searching the user’s Python environment. + +Mangling +"""""""" +To avoid confusion (“is this ``foo.bar`` object the one from my package, or the one from my Python environment?”), :class:`PackageImporter` mangles the +``__name__`` and ``__file__`` of all imported modules, by adding a *mangle prefix* to them. + +For ``__name__``, a name like ``torchvision.models.resnet18`` becomes ``.torchvision.models.resnet18``. + +For ``__file__``, a name like ``torchvision/models/resnet18.py`` becomes ``.torchvision/modules/resnet18.py``. + +Name mangling helps avoid inadvertent punning of module names between different packages, and helps you debug by making stack traces and print +statements more clearly show whether they are referring to packaged code or not. For developer-facing details about mangling, consult +``mangling.md`` in ``torch/package/``. + + API Reference ------------- .. autoclass:: torch.package.PackagingError From 4025f95a20141d245c406a3f1f8b8ddf6c1c2643 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Thu, 10 Jun 2021 19:50:22 -0700 Subject: [PATCH 031/305] [docs] Add table of contents to torch.package docs (#59842) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59842 Test Plan: Continuous integration. Captura de Pantalla 2021-06-10 a la(s) 5 13 07 p m Reviewed By: Lilyjjo Differential Revision: D29050627 Pulled By: SplitInfinity fbshipit-source-id: 76c25ed4002cbaf072036e2e14e7857c15077df7 --- docs/source/package.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/source/package.rst b/docs/source/package.rst index 99e6e75b531d1..ed4d6457846bf 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -15,6 +15,11 @@ at a later date or on a different machine, and can even be deployed to productio This document contains tutorials, how-to guides, explanations, and an API reference that will help you learn more about ``torch.package`` and how to use it. + +.. contents:: :local: + :depth: 2 + + Tutorials --------- Packaging your first model From 54bfd41a2ee96f5b142a03bb009742398c1e9545 Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Thu, 10 Jun 2021 20:47:09 -0700 Subject: [PATCH 032/305] Fix torch.angle on aarch64 (#59832) Summary: angle should return 0 for positive values, pi for negative and keep nans in place, which can be accomplished using two blendv functions. Fixes number of unary test failures on M1/aarch64 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59832 Reviewed By: kimishpatel Differential Revision: D29046402 Pulled By: malfet fbshipit-source-id: cb93ad2de140f7a54796387fc11053c507a1d4e9 --- aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h index 628a19143d1ce..08638bba66aa4 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h @@ -308,7 +308,10 @@ template <> class Vectorized { return Vectorized(vabsq_f32(values.val[0]), vabsq_f32(values.val[1])); } Vectorized angle() const { - return Vectorized(0.f); + auto zero = Vectorized(0); + auto pi = Vectorized(c10::pi); + auto tmp = blendv(zero, pi, *this < zero); + return blendv(tmp, *this, isnan()); } Vectorized real() const { return *this; From c7890b4a8e076f1540b2db351470b7a22fd8a0e3 Mon Sep 17 00:00:00 2001 From: Lily Johnson Date: Thu, 10 Jun 2021 21:20:42 -0700 Subject: [PATCH 033/305] [package] doc string cleanup extravaganza (#59843) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59843 Test Plan: Imported from OSS Reviewed By: SplitInfinity Differential Revision: D29049342 Pulled By: Lilyjjo fbshipit-source-id: 3330fb439f28dda0cafef5797ff61311f4afbf76 --- .../package/file_structure_representation.py | 6 +- torch/package/package_exporter.py | 84 ++++++++++--------- torch/package/package_importer.py | 38 ++++----- 3 files changed, 68 insertions(+), 60 deletions(-) diff --git a/torch/package/file_structure_representation.py b/torch/package/file_structure_representation.py index bc0cfe386ca4f..ad2e721b7ca04 100644 --- a/torch/package/file_structure_representation.py +++ b/torch/package/file_structure_representation.py @@ -7,7 +7,7 @@ class Directory: """A file structure representation. Organized as Directory nodes that have lists of their Directory children. Directories for a package are created by calling - :meth:`PackageExporter.file_structure` or :meth:`PackageImporter.file_structure`.""" + :meth:`PackageImporter.file_structure`.""" def __init__(self, name: str, is_dir: bool): self.name = name @@ -43,12 +43,12 @@ def _add_file(self, file_path: str): dir.children[file] = Directory(file, False) def has_file(self, filename: str) -> bool: - """Checks if a file is present in a Directory. + """Checks if a file is present in a :class:`Directory`. Args: filename (str): Path of file to search for. Returns: - bool: if a Directory contains the specified file. + bool: If a :class:`Directory` contains the specified file. """ lineage = filename.split("/", maxsplit=1) child = lineage[0] diff --git a/torch/package/package_exporter.py b/torch/package/package_exporter.py index b42de6a8ae063..9444fae168fd6 100644 --- a/torch/package/package_exporter.py +++ b/torch/package/package_exporter.py @@ -40,7 +40,7 @@ class _ModuleProviderAction(Enum): - """Represents one of the actions that exporter can take on a module. + """Represents one of the actions that :class:`PackageExporter` can take on a module. See :meth:`PackageExporter.extern` and friends for a description of what the actions do. """ @@ -151,11 +151,11 @@ class PackageExporter: The importer for packages ensures that code in the module can only be loaded from within the package, except for modules explicitly listed as external using :meth:`extern`. - The file `extern_modules` in the zip archive lists all the modules that a package externally depends on. + The file ``extern_modules`` in the zip archive lists all the modules that a package externally depends on. This prevents "implicit" dependencies where the package runs locally because it is importing a locally-installed package, but then fails when the package is copied to another machine. - When source code is added to the package, the exporter optionally can scan it + When source code is added to the package, the exporter can optionally scan it for further code dependencies (``dependencies=True``). It looks for import statements, resolves relative references to qualified module names, and performs an action specified by the user (See: :meth:`extern`, :meth:`mock`, and :meth:`intern`). @@ -176,7 +176,7 @@ def __init__( Create an exporter. Args: - f: The location to export to. Can be a ``string``/``Path`` object containing a filename, + f: The location to export to. Can be a ``string``/``Path`` object containing a filename or a binary I/O object. importer: If a single Importer is passed, use that to search for modules. If a sequence of importers are passsed, an ``OrderedImporter`` will be constructed out of them. @@ -289,14 +289,14 @@ def save_source_string( is_package: bool = False, dependencies: bool = True, ): - """Adds `src` as the source code for `module_name` in the exported package. + """Adds ``src`` as the source code for ``module_name`` in the exported package. Args: - module_name (str): e.g. `my_package.my_subpackage`, code will be saved to provide code for this package. + module_name (str): e.g. ``my_package.my_subpackage``, code will be saved to provide code for this package. src (str): The Python source code to save for this package. - is_package (bool, optional): If True, this module is treated as a package. Packages are allowed to have submodules - (e.g. my_package.my_subpackage.my_subsubpackage), and resources can be saved inside them. Defaults to ``False``. - dependencies (bool, optional): If True, we scan the source for dependencies. + is_package (bool, optional): If ``True``, this module is treated as a package. Packages are allowed to have submodules + (e.g. ``my_package.my_subpackage.my_subsubpackage``), and resources can be saved inside them. Defaults to ``False``. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. """ self.dependency_graph.add_node( module_name, @@ -421,7 +421,7 @@ def save_module(self, module_name: str, dependencies=True): module object, and then using its ``__file__`` attribute to find the source code. Args: - module_name (str): e.g. `my_package.my_subpackage`, code will be saved to provide code + module_name (str): e.g. ``my_package.my_subpackage``, code will be saved to provide code for this package. dependencies (bool, optional): If ``True``, we scan the source for dependencies. """ @@ -480,7 +480,7 @@ def save_pickle( ): """Save a python object to the archive using pickle. Equivalent to :func:`torch.save` but saving into the archive rather than a stand-alone file. Stanard pickle does not save the code, only the objects. - If `dependencies` is true, this method will also scan the pickled objects for which modules are required + If ``dependencies`` is true, this method will also scan the pickled objects for which modules are required to reconstruct them and save the relevant code. To be able to save an object where ``type(obj).__name__`` is ``my_module.MyObject``, @@ -489,7 +489,7 @@ def save_pickle( for this to work. Args: - package (str): The name of module package this resource should go in (e.g. "my_package.my_subpackage") + package (str): The name of module package this resource should go in (e.g. ``"my_package.my_subpackage"``). resource (str): A unique name for the resource, used to identify it to load. obj (Any): The object to save, must be picklable. dependencies (bool, optional): If ``True``, we scan the source for dependencies. @@ -533,7 +533,7 @@ def save_text(self, package: str, resource: str, text: str): """Save text data to the package. Args: - package (str): The name of module package this resource should go it (e.g. "my_package.my_subpackage") + package (str): The name of module package this resource should go it (e.g. ``"my_package.my_subpackage"``). resource (str): A unique name for the resource, used to identify it to load. text (str): The contents to save. """ @@ -543,7 +543,7 @@ def save_binary(self, package, resource, binary: bytes): """Save raw bytes to the package. Args: - package (str): The name of module package this resource should go it (e.g. "my_package.my_subpackage") + package (str): The name of module package this resource should go it (e.g. ``"my_package.my_subpackage"``). resource (str): A unique name for the resource, used to identify it to load. binary (str): The data to save. """ @@ -562,8 +562,8 @@ def register_extern_hook(self, hook: ActionHook) -> RemovableHandle: Returns: :class:`torch.utils.hooks.RemovableHandle`: - a handle that can be used to remove the added hook by calling - ``handle.remove()`` + A handle that can be used to remove the added hook by calling + ``handle.remove()``. """ handle = RemovableHandle(self._extern_hooks) self._extern_hooks[handle.id] = hook @@ -581,8 +581,8 @@ def register_mock_hook(self, hook: ActionHook) -> RemovableHandle: Returns: :class:`torch.utils.hooks.RemovableHandle`: - a handle that can be used to remove the added hook by calling - ``handle.remove()`` + A handle that can be used to remove the added hook by calling + ``handle.remove()``. """ handle = RemovableHandle(self._mock_hooks) self._mock_hooks[handle.id] = hook @@ -600,8 +600,8 @@ def register_intern_hook(self, hook: ActionHook) -> RemovableHandle: Returns: :class:`torch.utils.hooks.RemovableHandle`: - a handle that can be used to remove the added hook by calling - ``handle.remove()`` + A handle that can be used to remove the added hook by calling + ``handle.remove()``. """ handle = RemovableHandle(self._intern_hooks) self._intern_hooks[handle.id] = hook @@ -625,7 +625,7 @@ def intern( allow_empty (bool): An optional flag that specifies whether the intern modules specified by this call to the ``intern`` method must be matched to some module during packaging. If an ``intern`` module glob - pattern is added with ``allow_empty=False``, and ``close`` is called (either explicitly or via ``__exit__``) + pattern is added with ``allow_empty=False``, and :meth:`close` is called (either explicitly or via ``__exit__``) before any modules match that pattern, an exception is thrown. If ``allow_empty=True``, no such exception is thrown. """ @@ -647,22 +647,27 @@ def mock( Use this function to mock this functionality out without having to modify the original code. Args: - include (Union[List[str], str]): A string e.g. "my_package.my_subpackage", or list of strings + include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings for the names of the modules to be mocked out. Strings can also be a glob-style pattern string that may match multiple modules. Any required dependencies that match this pattern string will be mocked out automatically. - Examples: - 'torch.**' -- matches torch and all submodules of torch, e.g. 'torch.nn' and torch.nn.functional' - 'torch.*' -- matches 'torch.nn' or 'torch.functional', but not 'torch.nn.functional' + Examples : + ``'torch.**'`` -- matches ``torch`` and all submodules of torch, e.g. ``'torch.nn'`` + and ``'torch.nn.functional'`` + + ``'torch.*'`` -- matches ``'torch.nn'`` or ``'torch.functional'``, but not + ``'torch.nn.functional'`` exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string. - e.g. include='torch.**', exclude='torch.foo' will mock all torch packages except 'torch.foo' Default: [] + e.g. ``include='torch.**', exclude='torch.foo'`` will mock all torch packages except ``'torch.foo'``, + Default: is ``[]``. allow_empty (bool): An optional flag that specifies whether the mock implementation(s) specified by this call - to the `mock` method must be matched to some module during packaging. If a mock is added with allow_empty=False, - and `close` is called (either explicitly or via `__exit__`) and the mock has not been matched to a module - used by the package being exported, an exception is thrown. If allow_empty=True, no such exception is thrown. + to the :meth:`mock` method must be matched to some module during packaging. If a mock is added with + ``allow_empty=False``, and :meth:`close` is called (either explicitly or via ``__exit__``) and the mock has + not been matched to a module used by the package being exported, an exception is thrown. + If ``allow_empty=True``, no such exception is thrown. """ self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( @@ -676,21 +681,24 @@ def extern( exclude: "GlobPattern" = (), allow_empty: bool = True, ): - """Include `module` in the list of external modules the package can import. + """Include ``module`` in the list of external modules the package can import. This will prevent dependency discovery from saving it in the package. The importer will load an external module directly from the standard import system. Code for extern modules must also exist in the process loading the package. Args: - include (Union[List[str], str]): A string e.g. "my_package.my_subpackage", or list of strings - for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock` + include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings + for the names of the modules to be externed. This can also be a glob-style pattern, as + described in :meth:`mock`. - exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string. + exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the + include string. allow_empty (bool): An optional flag that specifies whether the extern modules specified by this call - to the `extern` method must be matched to some module during packaging. If an extern module glob pattern is added - with allow_empty=False, and `close` is called (either explicitly or via `__exit__`) before any modules match that - pattern, an exception is thrown. If allow_empty=True, no such exception is thrown. + to the ``extern`` method must be matched to some module during packaging. If an extern module glob + pattern is added with ``allow_empty=False``, and :meth:`close` is called (either explicitly or via + ``__exit__``) before any modules match that pattern, an exception is thrown. If ``allow_empty=True``, + no such exception is thrown. """ self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( @@ -702,8 +710,8 @@ def deny(self, include: "GlobPattern", *, exclude: "GlobPattern" = ()): If a dependency on any matching packages is found, a :class:`PackagingError` is raised. Args: - include (Union[List[str], str]): A string e.g. "my_package.my_subpackage", or list of strings - for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock` + include (Union[List[str], str]): A string e.g. ``"my_package.my_subpackage"``, or list of strings + for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock`. exclude (Union[List[str], str]): An optional pattern that excludes some patterns that match the include string. """ diff --git a/torch/package/package_importer.py b/torch/package/package_importer.py index 6d673edf30b73..f901d46a3e96c 100644 --- a/torch/package/package_importer.py +++ b/torch/package/package_importer.py @@ -29,7 +29,7 @@ class PackageImporter(Importer): - """Importers allow you to load code written to packages by PackageExporter. + """Importers allow you to load code written to packages by :class:`PackageExporter`. Code is loaded in a hermetic way, using files from the package rather than the normal python import system. This allows for the packaging of PyTorch model code and data so that it can be run @@ -37,12 +37,12 @@ class PackageImporter(Importer): The importer for packages ensures that code in the module can only be loaded from within the package, except for modules explicitly listed as external during export. - The file `extern_modules` in the zip archive lists all the modules that a package externally depends on. + The file ``extern_modules`` in the zip archive lists all the modules that a package externally depends on. This prevents "implicit" dependencies where the package runs locally because it is importing a locally-installed package, but then fails when the package is copied to another machine. """ - """The dictionary of already loaded modules from this package, equivalent to `sys.modules` but + """The dictionary of already loaded modules from this package, equivalent to ``sys.modules`` but local to this importer. """ modules: Dict[str, types.ModuleType] @@ -52,12 +52,12 @@ def __init__( file_or_buffer: Union[str, torch._C.PyTorchFileReader, Path, BinaryIO], module_allowed: Callable[[str], bool] = lambda module_name: True, ): - """Open `file_or_buffer` for importing. This checks that the imported package only requires modules - allowed by `module_allowed` + """Open ``file_or_buffer`` for importing. This checks that the imported package only requires modules + allowed by ``module_allowed`` Args: file_or_buffer: a file-like object (has to implement :meth:`read`, :meth:`readline`, :meth:`tell`, and :meth:`seek`), - or a string or os.PathLike object containing a file name. + a string, or an ``os.PathLike`` object containing a filename. module_allowed (Callable[[str], bool], optional): A method to determine if a externally provided module should be allowed. Can be used to ensure packages loaded do not depend on modules that the server does not support. Defaults to allowing anything. @@ -111,14 +111,14 @@ def __init__( def import_module(self, name: str, package=None): """Load a module from the package if it hasn't already been loaded, and then return the module. Modules are loaded locally - to the importer and will appear in `self.modules` rather than `sys.modules` + to the importer and will appear in ``self.modules`` rather than ``sys.modules``. Args: name (str): Fully qualified name of the module to load. - package ([type], optional): Unused, but present to match the signature of importlib.import_module. Defaults to None. + package ([type], optional): Unused, but present to match the signature of importlib.import_module. Defaults to ``None``. Returns: - types.ModuleType: the (possibly already) loaded module. + types.ModuleType: The (possibly already) loaded module. """ return self._gcd_import(name) @@ -126,7 +126,7 @@ def load_binary(self, package: str, resource: str) -> bytes: """Load raw bytes. Args: - package (str): The name of module package (e.g. "my_package.my_subpackage") + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). resource (str): The unique name for the resource. Returns: @@ -146,10 +146,10 @@ def load_text( """Load a string. Args: - package (str): The name of module package (e.g. "my_package.my_subpackage") + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). resource (str): The unique name for the resource. - encoding (str, optional): Passed to `decode`. Defaults to 'utf-8'. - errors (str, optional): Passed to `decode`. Defaults to 'strict'. + encoding (str, optional): Passed to ``decode``. Defaults to ``'utf-8'``. + errors (str, optional): Passed to ``decode``. Defaults to ``'strict'``. Returns: str: The loaded text. @@ -159,15 +159,15 @@ def load_text( def load_pickle(self, package: str, resource: str, map_location=None) -> Any: """Unpickles the resource from the package, loading any modules that are needed to construct the objects - using :meth:`import_module` + using :meth:`import_module`. Args: - package (str): The name of module package (e.g. "my_package.my_subpackage") + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). resource (str): The unique name for the resource. - map_location: Passed to `torch.load` to determine how tensors are mapped to devices. Defaults to None. + map_location: Passed to `torch.load` to determine how tensors are mapped to devices. Defaults to ``None``. Returns: - Any: the unpickled object. + Any: The unpickled object. """ pickle_file = self._zipfile_path(package, resource) restore_location = _get_restore_location(map_location) @@ -249,7 +249,7 @@ def set_deserialization_context(): def id(self): """ - Returns internal identifier that torch.package uses to distinguish PackageImporter instances. + Returns internal identifier that torch.package uses to distinguish :class:`PackageImporter` instances. Looks like:: @@ -262,7 +262,7 @@ def file_structure( """Returns a file structure representation of package's zipfile. Args: - include (Union[List[str], str]): An optional string e.g. "my_package.my_subpackage", or optional list of strings + include (Union[List[str], str]): An optional string e.g. ``"my_package.my_subpackage"``, or optional list of strings for the names of the files to be inluded in the zipfile representation. This can also be a glob-style pattern, as described in :meth:`PackageExporter.mock` From 455afdf974fc2546cb79f4cc27d307d4b24437ba Mon Sep 17 00:00:00 2001 From: Facebook Community Bot Date: Thu, 10 Jun 2021 21:52:31 -0700 Subject: [PATCH 034/305] Automated submodule update: FBGEMM (#59715) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59715 This is an automated pull request to update the first-party submodule for [pytorch/FBGEMM](https://github.com/pytorch/FBGEMM). New submodule commit: https://github.com/pytorch/FBGEMM/commit/0520ad5f95db754fbc0ccfb7b563986b6d77bb20 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59687 Test Plan: Ensure that CI jobs succeed on GitHub before landing. Reviewed By: jianyuh Differential Revision: D28986238 Pulled By: jspark1105 fbshipit-source-id: 12f68830b5b7a858fbc301af50593281852af51f --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index 5885ee1387ab3..0520ad5f95db7 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 5885ee1387ab363a6cc7d28fa9d4dc3ef1cd5f23 +Subproject commit 0520ad5f95db754fbc0ccfb7b563986b6d77bb20 From 6eabbea47ced39111766a043a92a22a11cd45975 Mon Sep 17 00:00:00 2001 From: Eddie Yan Date: Thu, 10 Jun 2021 22:06:11 -0700 Subject: [PATCH 035/305] Disable cuDNN persistent RNN on A30 (#59830) Summary: https://github.com/pytorch/pytorch/issues/59829 cherry-picked from ptrblck 's change CC ngimel xwang233 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59830 Reviewed By: bdhirsh Differential Revision: D29046145 Pulled By: ngimel fbshipit-source-id: 270ab3bb6c1c7c759497a15eb38b20a177c94adb --- aten/src/ATen/native/cudnn/RNN.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/native/cudnn/RNN.cpp b/aten/src/ATen/native/cudnn/RNN.cpp index da830992b7e44..f81de80cf5ec2 100644 --- a/aten/src/ATen/native/cudnn/RNN.cpp +++ b/aten/src/ATen/native/cudnn/RNN.cpp @@ -726,7 +726,8 @@ namespace { (tensors.seq_length >=20 && bsize <=96) || (tensors.seq_length >=10 && bsize <=32)); } - } else if (prop->major >= 8) { + } else if (prop->major >= 8 && prop->multiProcessorCount >= 98) { + // SM count check excludes A30 (similar issue to A40) if (prop->minor == 6) { // Excludes sm_86 GPU devices from using persistent rnn. // This is because there are some edge cases that will throw exceptions with cudnn 8.0.5 on Nvidia A40 GPU. From 2112074f25d79af634b3d44fa3bb808f18e2923c Mon Sep 17 00:00:00 2001 From: Hao Lu Date: Thu, 10 Jun 2021 23:37:54 -0700 Subject: [PATCH 036/305] [Static Runtime] Add schema check to several aten ops (#59603) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59603 D28698997 (https://github.com/pytorch/pytorch/commit/10345010f761e961b3f6785a874374f7dc2be1c2) was reverted because I forgot to replace the ``` VLOG(1) << "Found schema mismatch"; n->schema().dump(); ``` block in `aten::clamp_min` with `LogAndDumpSchema(n)` and that led to the bazel build to fail. I don't know why it makes the bazel build though. Test Plan: OSS CI. Reviewed By: ajyu Differential Revision: D28950177 fbshipit-source-id: 9bb1c6619e6b68415a3349f04933c2fcd24cc9a2 --- torch/csrc/jit/runtime/static/ops.cpp | 710 ++++++++++++++------------ torch/csrc/jit/runtime/static/ops.h | 5 + 2 files changed, 386 insertions(+), 329 deletions(-) diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index d25849c0fc8fa..7e22b8a6b397a 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -274,9 +274,12 @@ REGISTER_OPERATOR_FUNCTOR( // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::mul, aten_mul, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::mul.Tensor(Tensor self, Tensor other) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } + return [](ProcessedNode* p_node) { const auto& in0_t = p_node->Input(0).toTensor(); const auto& in1_t = p_node->Input(1).toTensor(); @@ -291,7 +294,9 @@ REGISTER_OPERATOR_FUNCTOR(aten::mul, aten_mul, [](Node* n) -> SROperator { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::addmm, aten_addmm, [](Node* n) -> SROperator { - if (n->inputs().size() != 5) { + if (!n->matches(torch::schema( + "aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -309,13 +314,16 @@ REGISTER_OPERATOR_FUNCTOR(aten::addmm, aten_addmm, [](Node* n) -> SROperator { }; }); -// clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor -// clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::clamp, aten_clamp, [](Node* n) -> SROperator { - if (n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor")) && + !n->matches(torch::schema( + "aten::clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } + return [](ProcessedNode* p_node) { const auto& in0_t = p_node->Input(0).toTensor(); @@ -339,7 +347,9 @@ REGISTER_OPERATOR_FUNCTOR(aten::clamp, aten_clamp, [](Node* n) -> SROperator { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::bmm, aten_bmm, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches( + torch::schema("aten::bmm(Tensor self, Tensor mat2) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -355,29 +365,30 @@ REGISTER_OPERATOR_FUNCTOR(aten::bmm, aten_bmm, [](Node* n) -> SROperator { }); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - aten::nan_to_num, - aten_nan_to_num, - [](Node* n) -> SROperator { - if (n->inputs().size() != 4) { - return nullptr; - } - return [](ProcessedNode* p_node) { - const auto& in0_t = p_node->Input(0).toTensor(); - const auto in1_d = p_node->Input(1).toOptional(); - const auto in2_d = p_node->Input(2).toOptional(); - const auto in3_d = p_node->Input(3).toOptional(); - if (p_node->Output(0).isNone()) { - p_node->Output(0) = create_empty_from(in0_t); - } - auto& out_t = p_node->Output(0).toTensor(); - fastResizeToZero(out_t); - at::native::nan_to_num_out(in0_t, in1_d, in2_d, in3_d, out_t); - }; - }); +REGISTER_OPERATOR_FUNCTOR(aten::nan_to_num, aten_nan_to_num, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "aten::nan_to_num(Tensor self, float? nan=None, float? posinf=None, float? neginf=None) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + const auto& in0_t = p_node->Input(0).toTensor(); + const auto in1_d = p_node->Input(1).toOptional(); + const auto in2_d = p_node->Input(2).toOptional(); + const auto in3_d = p_node->Input(3).toOptional(); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = create_empty_from(in0_t); + } + auto& out_t = p_node->Output(0).toTensor(); + fastResizeToZero(out_t); + at::native::nan_to_num_out(in0_t, in1_d, in2_d, in3_d, out_t); + }; +}); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::cat, aten_cat, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches( + torch::schema("aten::cat(Tensor[] tensors, int dim=0) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -394,7 +405,9 @@ REGISTER_OPERATOR_FUNCTOR(aten::cat, aten_cat, [](Node* n) -> SROperator { // Split out into a function to appease MSVC's pre-processor SROperator aten_stack(Node* n) { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::stack(Tensor[] tensors, int dim=0) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -413,23 +426,22 @@ SROperator aten_stack(Node* n) { REGISTER_OPERATOR_FUNCTOR(aten::stack, aten_stack, aten_stack); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - aten::leaky_relu, - aten_leaky_relu, - [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { - return nullptr; - } - return [](ProcessedNode* p_node) { - const auto& in0_t = p_node->Input(0).toTensor(); - const auto in1_s = p_node->Input(1).toScalar(); - if (p_node->Output(0).isNone()) { - p_node->Output(0) = create_empty_from(in0_t); - } - auto& out_t = p_node->Output(0).toTensor(); - at::cpu::leaky_relu_out(out_t, in0_t, in1_s); - }; - }); +REGISTER_OPERATOR_FUNCTOR(aten::leaky_relu, aten_leaky_relu, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "aten::leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + const auto& in0_t = p_node->Input(0).toTensor(); + const auto in1_s = p_node->Input(1).toScalar(); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = create_empty_from(in0_t); + } + auto& out_t = p_node->Output(0).toTensor(); + at::cpu::leaky_relu_out(out_t, in0_t, in1_s); + }; +}); namespace { @@ -584,7 +596,8 @@ std::shared_ptr createSigmoid() { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::relu, aten_relu, [](Node* n) -> SROperator { - if (n->inputs().size() != 1) { + if (!n->matches(torch::schema("aten::relu(Tensor self) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } auto te = createRelu(); @@ -607,7 +620,8 @@ REGISTER_OPERATOR_FUNCTOR(aten::relu, aten_relu, [](Node* n) -> SROperator { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::tanh, aten_tanh, [](Node* n) -> SROperator { - if (n->inputs().size() != 1) { + if (!n->matches(torch::schema("aten::tanh(Tensor self) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } auto te = createTanh(); @@ -633,7 +647,8 @@ REGISTER_OPERATOR_FUNCTOR( aten::sigmoid, aten_sigmoid, [](Node* n) -> SROperator { - if (n->inputs().size() != 1) { + if (!n->matches(torch::schema("aten::sigmoid(Tensor self) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } auto te = createSigmoid(); @@ -656,7 +671,9 @@ REGISTER_OPERATOR_FUNCTOR( // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::logit, aten_logit, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::logit(Tensor self, float? eps=None) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } c10::optional clamp = c10::nullopt; @@ -686,10 +703,11 @@ REGISTER_OPERATOR_FUNCTOR(aten::logit, aten_logit, [](Node* n) -> SROperator { }; }); -// clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::clone, aten_clone, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -721,7 +739,9 @@ REGISTER_OPERATOR_FUNCTOR( quantized::embedding_bag_byte_rowwise_offsets, quantized_embedding_bag_byte_rowwise_offsets, [](Node* n) -> SROperator { - if (n->inputs().size() != 9) { + if (!n->matches(torch::schema( + "quantized::embedding_bag_byte_rowwise_offsets(Tensor weight, Tensor indices, Tensor? offsets=None, bool scale_grad_by_freq=False, int mode=0, bool pruned_weights=False, Tensor? per_sample_weights=None, Tensor? compressed_indices_mapping=None, bool include_last_offset=False) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -757,7 +777,9 @@ REGISTER_OPERATOR_FUNCTOR( quantized::embedding_bag_4bit_rowwise_offsets, embedding_bag_4bit_rowwise_offsets, [](Node* n) -> SROperator { - if (n->inputs().size() != 9) { + if (!n->matches(torch::schema( + "quantized::embedding_bag_4bit_rowwise_offsets(Tensor weight, Tensor indices, Tensor? offsets=None, bool scale_grad_by_freq=False, int mode=0, bool pruned_weights=False, Tensor? per_sample_weights=None, Tensor? compressed_indices_mapping=None, bool include_last_offset=False) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -791,36 +813,37 @@ REGISTER_OPERATOR_FUNCTOR( // The out variant takes precedence over native // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - aten::narrow_copy, - aten_narrow_copy, - [](Node* n) -> SROperator { - if (n->inputs().size() != 4) { - return nullptr; - } - return [](ProcessedNode* p_node) { - const auto& self = p_node->Input(0).toTensor(); // self - const auto dim = p_node->Input(1).toInt(); // dim - int64_t start = 0; - if (p_node->Input(2).isScalar()) { - start = p_node->Input(2).toInt(); - } else { - auto& t = p_node->Input(2).toTensor(); - start = t.item(); - } - auto length = p_node->Input(3).toInt(); // length +REGISTER_OPERATOR_FUNCTOR(aten::narrow_copy, aten_narrow_copy, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "aten::narrow_copy(Tensor self, int dim, int start, int length) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + const auto& self = p_node->Input(0).toTensor(); // self + const auto dim = p_node->Input(1).toInt(); // dim + int64_t start = 0; + if (p_node->Input(2).isScalar()) { + start = p_node->Input(2).toInt(); + } else { + auto& t = p_node->Input(2).toTensor(); + start = t.item(); + } + auto length = p_node->Input(3).toInt(); // length - if (p_node->Output(0).isNone()) { - p_node->Output(0) = create_empty_from(self); - } - auto& output = p_node->Output(0).toTensor(); - fastResizeToZero(output); - at::native::narrow_copy_dense_cpu_out(self, dim, start, length, output); - }; - }); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = create_empty_from(self); + } + auto& output = p_node->Output(0).toTensor(); + fastResizeToZero(output); + at::native::narrow_copy_dense_cpu_out(self, dim, start, length, output); + }; +}); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::index, aten_index, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::index.Tensor(Tensor self, Tensor?[] indices) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -837,7 +860,13 @@ REGISTER_OPERATOR_FUNCTOR(aten::index, aten_index, [](Node* n) -> SROperator { }); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::pow, aten_pow, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> Tensor")) && + !n->matches(torch::schema( + "aten::pow.Scalar(Scalar self, Tensor exponent) -> Tensor")) && + !n->matches(torch::schema( + "aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -980,6 +1009,13 @@ REGISTER_OPERATOR_FUNCTOR(aten::sum, aten_sum, [](Node* n) -> SROperator { if (n->inputs().size() != 2 && n->inputs().size() != 4) { return nullptr; } + if (!n->matches(torch::schema( + "aten::sum(Tensor self, *, ScalarType? dtype=None) -> Tensor")) && + !n->matches(torch::schema( + "aten::sum.dim_IntList(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } return [](ProcessedNode* p_node) { const at::Tensor& self = p_node->Input(0).toTensor(); @@ -1010,7 +1046,9 @@ REGISTER_OPERATOR_FUNCTOR(aten::sum, aten_sum, [](Node* n) -> SROperator { std::function getNativeOperation(Node* n) { if (n->kind() == c10::Symbol::fromQualString("aten::transpose")) { - if (n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1020,7 +1058,9 @@ std::function getNativeOperation(Node* n) { p_node->Output(0) = at::native::transpose(in0_t, in1_i, in2_i); }; } else if (n->kind() == c10::Symbol::fromQualString("aten::flatten")) { - if (n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::flatten.using_ints(Tensor(a) self, int start_dim=0, int end_dim=-1) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1119,7 +1159,9 @@ std::function getNativeOperation(Node* n) { } }; } else if (n->kind() == c10::Symbol::fromQualString("aten::permute")) { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::permute(Tensor(a) self, int[] dims) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1128,7 +1170,9 @@ std::function getNativeOperation(Node* n) { p_node->Output(0) = at::native::permute(in0_t, in1_iv); }; } else if (n->kind() == c10::Symbol::fromQualString("aten::reshape")) { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::reshape(Tensor(a) self, int[] shape) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1137,7 +1181,9 @@ std::function getNativeOperation(Node* n) { p_node->Output(0) = at::native::reshape(in0_t, in1_iv); }; } else if (n->kind() == c10::Symbol::fromQualString("aten::slice")) { - if (n->inputs().size() != 5) { + if (!n->matches(torch::schema( + "aten::slice.Tensor(Tensor(a) self, int dim=0, int? start=0, int? end=9223372036854775807, int step=1) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1149,7 +1195,11 @@ std::function getNativeOperation(Node* n) { p_node->Output(0) = at::native::slice(in0_t, in1_i, in2_i, in3_i, in4_i); }; } else if (n->kind() == c10::Symbol::fromQualString("aten::narrow")) { - if (n->inputs().size() != 4) { + if (!n->matches(torch::schema( + "aten::narrow(Tensor(a) self, int dim, int start, int length) -> Tensor(a)")) && + !n->matches(torch::schema( + "aten::narrow.Tensor(Tensor(a) self, int dim, Tensor start, int length) -> Tensor(a)"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1183,7 +1233,11 @@ std::function getNativeOperation(Node* n) { at::native::slice(self, dim, start, start + length, 1); }; } else if (n->kind() == c10::Symbol::fromQualString("aten::to")) { - if (n->inputs().size() != 5) { + if (!n->matches(torch::schema( + "aten::to.other(Tensor self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor")) && + !n->matches(torch::schema( + "aten::to.dtype(Tensor self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1208,111 +1262,112 @@ std::function getNativeOperation(Node* n) { } // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - aten::embedding_bag, - aten_embedding_bag, - [](Node* n) -> SROperator { - // TODO: Support only 9 args once the old signature has been removed. - if (n->inputs().size() != 8 && n->inputs().size() != 9) { - return nullptr; +REGISTER_OPERATOR_FUNCTOR(aten::embedding_bag, aten_embedding_bag, [](Node* n) -> SROperator { + // TODO: Support only 9 args once the old signature has been removed. + if (!n->matches(torch::schema( + "aten::embedding_bag(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq=False, int mode=0, bool sparse=False, Tensor? per_sample_weights=None, bool include_last_offset=False) -> (Tensor, Tensor, Tensor, Tensor)")) && + !n->matches(torch::schema( + "aten::embedding_bag.padding_idx(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq, int mode, bool sparse, Tensor? per_sample_weights, bool include_last_offset, int? padding_idx) -> (Tensor, Tensor, Tensor, Tensor)"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + const auto& weight = p_node->Input(0).toTensor(); + const auto& indices = p_node->Input(1).toTensor(); + const auto& offsets = p_node->Input(2).toTensor(); + auto scale_grad_by_freq = p_node->Input(3).toBool(); + auto mode = p_node->Input(4).to(); + auto sparse = p_node->Input(5).toBool(); + auto per_sample_weights = p_node->Input(6).toOptional(); + auto include_last_offset = p_node->Input(7).toBool(); + c10::optional padding_idx; + if (p_node->inputs().size() == 9) { + if (p_node->Input(8).isNone()) { + padding_idx = c10::nullopt; + } else { + padding_idx = p_node->Input(8).toInt(); } - return [](ProcessedNode* p_node) { - const auto& weight = p_node->Input(0).toTensor(); - const auto& indices = p_node->Input(1).toTensor(); - const auto& offsets = p_node->Input(2).toTensor(); - auto scale_grad_by_freq = p_node->Input(3).toBool(); - auto mode = p_node->Input(4).to(); - auto sparse = p_node->Input(5).toBool(); - auto per_sample_weights = p_node->Input(6).toOptional(); - auto include_last_offset = p_node->Input(7).toBool(); - c10::optional padding_idx; - if (p_node->inputs().size() == 9) { - if (p_node->Input(8).isNone()) { - padding_idx = c10::nullopt; - } else { - padding_idx = p_node->Input(8).toInt(); - } - } - - at::native::check_arguments( - weight, - indices, - offsets, - mode, - per_sample_weights, - include_last_offset); - - std::ignore = scale_grad_by_freq; - std::ignore = sparse; + } - if (p_node->Output(0).isNone()) { - p_node->Output(0) = at::empty( - {include_last_offset ? offsets.sizes()[0] - 1 - : offsets.sizes()[0], - weight.sizes()[1]}, - weight.options()); - } else { - at::native::resize_( - p_node->Output(0).toTensor(), - {include_last_offset ? offsets.sizes()[0] - 1 - : offsets.sizes()[0], - weight.sizes()[1]}, - c10::nullopt); - } - at::Tensor& output = p_node->Output(0).toTensor(); + at::native::check_arguments( + weight, + indices, + offsets, + mode, + per_sample_weights, + include_last_offset); - if (p_node->Output(1).isNone()) { - p_node->Output(1) = at::empty({0}, offsets.options()); - } - at::Tensor& offset2bag = p_node->Output(1).toTensor(); - at::native::make_offset2bag_out( - offset2bag, - output, - weight, - indices, - offsets, - mode, - per_sample_weights, - padding_idx.value_or(-1)); + std::ignore = scale_grad_by_freq; + std::ignore = sparse; - if (p_node->Output(2).isNone()) { - p_node->Output(2) = at::empty(offsets.sizes(), offsets.options()); - } - at::Tensor& bag_size = p_node->Output(2).toTensor(); - at::native::make_bag_size_out( - bag_size, offsets, indices, mode, include_last_offset, false); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = at::empty( + {include_last_offset ? offsets.sizes()[0] - 1 : offsets.sizes()[0], + weight.sizes()[1]}, + weight.options()); + } else { + at::native::resize_( + p_node->Output(0).toTensor(), + {include_last_offset ? offsets.sizes()[0] - 1 : offsets.sizes()[0], + weight.sizes()[1]}, + c10::nullopt); + } + at::Tensor& output = p_node->Output(0).toTensor(); - if (p_node->Output(3).isNone()) { - p_node->Output(3) = at::empty(bag_size.sizes(), offsets.options()); - } - at::Tensor& max_indices = p_node->Output(3).toTensor(); - at::native::make_max_indices_out( - max_indices, - weight, - indices, - offsets, - bag_size, - mode, - include_last_offset); + if (p_node->Output(1).isNone()) { + p_node->Output(1) = at::empty({0}, offsets.options()); + } + at::Tensor& offset2bag = p_node->Output(1).toTensor(); + at::native::make_offset2bag_out( + offset2bag, + output, + weight, + indices, + offsets, + mode, + per_sample_weights, + padding_idx.value_or(-1)); + + if (p_node->Output(2).isNone()) { + p_node->Output(2) = at::empty(offsets.sizes(), offsets.options()); + } + at::Tensor& bag_size = p_node->Output(2).toTensor(); + at::native::make_bag_size_out( + bag_size, offsets, indices, mode, include_last_offset, false); - at::native::_embedding_bag_cpu_impl_out( - output, - offset2bag, - bag_size, - max_indices, - weight, - indices, - offsets, - mode, - per_sample_weights, - include_last_offset, - padding_idx.value_or(-1)); - }; - }); + if (p_node->Output(3).isNone()) { + p_node->Output(3) = at::empty(bag_size.sizes(), offsets.options()); + } + at::Tensor& max_indices = p_node->Output(3).toTensor(); + at::native::make_max_indices_out( + max_indices, + weight, + indices, + offsets, + bag_size, + mode, + include_last_offset); + + at::native::_embedding_bag_cpu_impl_out( + output, + offset2bag, + bag_size, + max_indices, + weight, + indices, + offsets, + mode, + per_sample_weights, + include_last_offset, + padding_idx.value_or(-1)); + }; +}); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::repeat, aten_repeat, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::repeat(Tensor self, int[] repeats) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1328,7 +1383,15 @@ REGISTER_OPERATOR_FUNCTOR(aten::repeat, aten_repeat, [](Node* n) -> SROperator { }); REGISTER_OPERATOR_FUNCTOR(aten::div, aten_div, [](Node* n) -> SROperator { - if (n->inputs().size() != 2 && n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::div.Tensor(Tensor self, Tensor other) -> Tensor")) && + !n->matches(torch::schema( + "aten::div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor")) && + !n->matches(torch::schema( + "aten::div.Scalar(Tensor self, Scalar other) -> Tensor")) && + !n->matches(torch::schema( + "aten::div.Scalar_mode(Tensor self, Scalar other, *, str? rounding_mode) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1353,7 +1416,11 @@ REGISTER_OPERATOR_FUNCTOR(aten::div, aten_div, [](Node* n) -> SROperator { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::sub, aten_sub, [](Node* n) -> SROperator { - if (n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor")) && + !n->matches(torch::schema( + "aten::sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1373,12 +1440,15 @@ REGISTER_OPERATOR_FUNCTOR(aten::sub, aten_sub, [](Node* n) -> SROperator { }; }); +// TODO: support clamp_min.Tensor(Tensor self, Tensor min) -> Tensor // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR( aten::clamp_min, aten_clamp_min, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches(torch::schema( + "aten::clamp_min(Tensor self, Scalar min) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1395,7 +1465,9 @@ REGISTER_OPERATOR_FUNCTOR( // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::argmin, aten_argmin, [](Node* n) -> SROperator { - if (n->inputs().size() != 3) { + if (!n->matches(torch::schema( + "aten::argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1411,83 +1483,64 @@ REGISTER_OPERATOR_FUNCTOR(aten::argmin, aten_argmin, [](Node* n) -> SROperator { }; }); -REGISTER_OPERATOR_FUNCTOR( - aten::layer_norm, - aten_layer_norm, - [](Node* n) -> SROperator { - if (n->inputs().size() != 6) { - return nullptr; - } - return [](ProcessedNode* p_node) { - // ignore Input(5): `bool cudnn_enable=True` - const auto& input = p_node->Input(0).toTensor(); - const auto normalized_shape = p_node->Input(1).toIntVector(); - auto weight_opt = p_node->Input(2).toOptional(); - auto bias_opt = p_node->Input(3).toOptional(); - float eps = p_node->Input(4).toDouble(); - - c10::MaybeOwned weight_maybe_owned = - at::borrow_from_optional_tensor(weight_opt); - const at::Tensor& weight = *weight_maybe_owned; - c10::MaybeOwned bias_maybe_owned = - at::borrow_from_optional_tensor(bias_opt); - const at::Tensor& bias = *bias_maybe_owned; - - auto M_N = at::native::_check_layer_norm_inputs( - input, normalized_shape, weight, bias); - auto M = M_N.first; - auto N = M_N.second; - auto X = input.expect_contiguous(); - auto gamma = weight.expect_contiguous(); - auto beta = bias.expect_contiguous(); +REGISTER_OPERATOR_FUNCTOR(aten::layer_norm, aten_layer_norm, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "aten::layer_norm(Tensor input, int[] normalized_shape, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enable=True) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + // ignore Input(5): `bool cudnn_enable=True` + const auto& input = p_node->Input(0).toTensor(); + const auto normalized_shape = p_node->Input(1).toIntVector(); + auto weight_opt = p_node->Input(2).toOptional(); + auto bias_opt = p_node->Input(3).toOptional(); + float eps = p_node->Input(4).toDouble(); + + c10::MaybeOwned weight_maybe_owned = + at::borrow_from_optional_tensor(weight_opt); + const at::Tensor& weight = *weight_maybe_owned; + c10::MaybeOwned bias_maybe_owned = + at::borrow_from_optional_tensor(bias_opt); + const at::Tensor& bias = *bias_maybe_owned; + + auto M_N = at::native::_check_layer_norm_inputs( + input, normalized_shape, weight, bias); + auto M = M_N.first; + auto N = M_N.second; + auto X = input.expect_contiguous(); + auto gamma = weight.expect_contiguous(); + auto beta = bias.expect_contiguous(); - if (p_node->Output(0).isNone()) { - p_node->Output(0) = at::native::empty_like( - *X, - c10::nullopt /* dtype */, - c10::nullopt /* layout */, - c10::nullopt /* device */, - c10::nullopt /* pin_memory */, - at::MemoryFormat::Contiguous); - } else { - at::native::resize_( - p_node->Output(0).toTensor(), X->sizes(), c10::nullopt); - } - at::Tensor& output = p_node->Output(0).toTensor(); - at::Tensor mean = create_empty_from({M}, *X); - at::Tensor rstd = create_empty_from({M}, *X); - - at::native::layer_norm_cpu_out( - output, - mean, - rstd, - input, - normalized_shape, - *gamma, - *beta, - eps, - M, - N); - }; - }); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = at::native::empty_like( + *X, + c10::nullopt /* dtype */, + c10::nullopt /* layout */, + c10::nullopt /* device */, + c10::nullopt /* pin_memory */, + at::MemoryFormat::Contiguous); + } else { + at::native::resize_( + p_node->Output(0).toTensor(), X->sizes(), c10::nullopt); + } + at::Tensor& output = p_node->Output(0).toTensor(); + at::Tensor mean = create_empty_from({M}, *X); + at::Tensor rstd = create_empty_from({M}, *X); + + at::native::layer_norm_cpu_out( + output, mean, rstd, input, normalized_shape, *gamma, *beta, eps, M, N); + }; +}); -/* Support the following signatures of norm: - * norm.ScalarOpt_dtype(Tensor self, Scalar? p, *, ScalarType dtype) - * norm.ScalarOpt_dim_dtype(Tensor self, Scalar? p, int[1] dim, bool keepdim, *, - * ScalarType dtype) - * norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) - */ REGISTER_OPERATOR_FUNCTOR(aten::norm, aten_norm, [](Node* n) -> SROperator { - if (n->inputs().size() <= 2) { - LOG(ERROR) - << "Please implement static runtime support for aten::norm 2-arg version"; - return nullptr; - } - // check that the third arg is scalar or int[] - auto val_2 = toIValue(n->inputs()[2]); - if (val_2 && !(val_2->isIntList() || val_2->isInt())) { - LOG(ERROR) - << "Please implement static runtime support for aten::norm w/ DimnameList"; + if (!n->matches(torch::schema( + "aten::norm.ScalarOpt_dtype(Tensor self, Scalar? p, *, ScalarType dtype) -> Tensor")) && + !n->matches(torch::schema( + "aten::norm.ScalarOpt_dim_dtype(Tensor self, Scalar? p, int[1] dim, bool keepdim, *, ScalarType dtype) -> Tensor")) && + !n->matches(torch::schema( + "aten::norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } @@ -1534,7 +1587,9 @@ REGISTER_OPERATOR_FUNCTOR(aten::norm, aten_norm, [](Node* n) -> SROperator { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_OPERATOR_FUNCTOR(aten::matmul, aten_matmul, [](Node* n) -> SROperator { - if (n->inputs().size() != 2) { + if (!n->matches( + torch::schema("aten::matmul(Tensor self, Tensor other) -> Tensor"))) { + LogAndDumpSchema(n); return nullptr; } return [](ProcessedNode* p_node) { @@ -1551,76 +1606,73 @@ REGISTER_OPERATOR_FUNCTOR(aten::matmul, aten_matmul, [](Node* n) -> SROperator { }); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - quantized::linear, - quantized_linear, - [](Node* n) -> SROperator { - if (n->inputs().size() != 4) { - return nullptr; - } - const auto w = toIValue(n->inputs()[1]); - c10::intrusive_ptr packed_weight; - if (w) { - packed_weight = w->toCustomClass(); - } - return [packed_weight](ProcessedNode* p_node) { - const auto& input = p_node->Input(0).toTensor(); - const auto output_scale = p_node->Input(2).toDouble(); - const auto output_zero_point = p_node->Input(3).toInt(); +REGISTER_OPERATOR_FUNCTOR(quantized::linear, quantized_linear, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "quantized::linear(Tensor X, __torch__.torch.classes.quantized.LinearPackedParamsBase W_prepack, float Y_scale_i, int Y_zero_point_i) -> Tensor Y"))) { + LogAndDumpSchema(n); + return nullptr; + } + const auto w = toIValue(n->inputs()[1]); + c10::intrusive_ptr packed_weight; + if (w) { + packed_weight = w->toCustomClass(); + } + return [packed_weight](ProcessedNode* p_node) { + const auto& input = p_node->Input(0).toTensor(); + const auto output_scale = p_node->Input(2).toDouble(); + const auto output_zero_point = p_node->Input(3).toInt(); - if (p_node->Output(0).isNone()) { - p_node->Output(0) = at::native::empty_affine_quantized( - {0}, - c10::kQUInt8, - c10::nullopt, - c10::kCPU, - false, - output_scale, - output_zero_point, - c10::nullopt); - } - auto& out_t = p_node->Output(0).toTensor(); - fastResizeToZero(out_t); + if (p_node->Output(0).isNone()) { + p_node->Output(0) = at::native::empty_affine_quantized( + {0}, + c10::kQUInt8, + c10::nullopt, + c10::kCPU, + false, + output_scale, + output_zero_point, + c10::nullopt); + } + auto& out_t = p_node->Output(0).toTensor(); + fastResizeToZero(out_t); - if (packed_weight) { - packed_weight->apply_out( - input, output_scale, output_zero_point, out_t); - } else { - // Weights could be quantized on the fly - auto packed_weight_tmp = - p_node->Input(1).toCustomClass(); - packed_weight_tmp->apply_out( - input, output_scale, output_zero_point, out_t); - } - }; - }); + if (packed_weight) { + packed_weight->apply_out(input, output_scale, output_zero_point, out_t); + } else { + // Weights could be quantized on the fly + auto packed_weight_tmp = + p_node->Input(1).toCustomClass(); + packed_weight_tmp->apply_out( + input, output_scale, output_zero_point, out_t); + } + }; +}); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_OPERATOR_FUNCTOR( - aten::full_like, - aten_full_like, - [](Node* n) -> SROperator { - if (n->inputs().size() != 7) { - return nullptr; - } - return [](ProcessedNode* p_node) { - const auto in1_s = p_node->Input(1).toScalar(); - if (p_node->Output(0).isNone()) { - const auto& in0_t = p_node->Input(0).toTensor(); - const auto dtype = p_node->Input(2).toOptional(); - const auto layout = p_node->Input(3).toOptional(); - const auto device = p_node->Input(4).toOptional(); - const auto pin_memory = p_node->Input(5).toOptional(); - const auto memory_format = - p_node->Input(6).toOptional(); - - p_node->Output(0) = at::native::empty_like( - in0_t, dtype, layout, device, pin_memory, memory_format); - } - auto& out_t = p_node->Output(0).toTensor(); - at::native::fill_out(out_t, in1_s); - }; - }); +REGISTER_OPERATOR_FUNCTOR(aten::full_like, aten_full_like, [](Node* n) -> SROperator { + if (!n->matches(torch::schema( + "aten::full_like(Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor"))) { + LogAndDumpSchema(n); + return nullptr; + } + return [](ProcessedNode* p_node) { + const auto in1_s = p_node->Input(1).toScalar(); + if (p_node->Output(0).isNone()) { + const auto& in0_t = p_node->Input(0).toTensor(); + const auto dtype = p_node->Input(2).toOptional(); + const auto layout = p_node->Input(3).toOptional(); + const auto device = p_node->Input(4).toOptional(); + const auto pin_memory = p_node->Input(5).toOptional(); + const auto memory_format = + p_node->Input(6).toOptional(); + + p_node->Output(0) = at::native::empty_like( + in0_t, dtype, layout, device, pin_memory, memory_format); + } + auto& out_t = p_node->Output(0).toTensor(); + at::native::fill_out(out_t, in1_s); + }; +}); } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/runtime/static/ops.h b/torch/csrc/jit/runtime/static/ops.h index 7fb7e1a51df51..d35df5f806587 100644 --- a/torch/csrc/jit/runtime/static/ops.h +++ b/torch/csrc/jit/runtime/static/ops.h @@ -133,5 +133,10 @@ inline std::string PrintNode(const Node* node) { return ss.str(); } +inline void LogAndDumpSchema(const Node* node) { + VLOG(1) << "Found schema mismatch"; + node->schema().dump(); +} + } // namespace jit } // namespace torch From 4b9135523229c9d78ba20c696b86670bd765a16d Mon Sep 17 00:00:00 2001 From: Gary Miguel Date: Fri, 11 Jun 2021 00:07:01 -0700 Subject: [PATCH 037/305] [ONNX] remove raw export type (#59160) Summary: [ONNX] remove raw export type Pull Request resolved: https://github.com/pytorch/pytorch/pull/59160 Reviewed By: tugsbayasgalan Differential Revision: D28937039 Pulled By: SplitInfinity fbshipit-source-id: 79bf91605526aa32a7304e75f50fe55d872bd4e8 --- test/jit/test_onnx_export.py | 17 ---- torch/_C/_onnx.pyi | 1 - torch/csrc/jit/serialization/export.cpp | 23 +---- torch/csrc/onnx/init.cpp | 1 - torch/csrc/onnx/onnx.h | 1 - torch/onnx/__init__.py | 7 +- torch/onnx/utils.py | 122 ++++++++++++------------ 7 files changed, 66 insertions(+), 106 deletions(-) diff --git a/test/jit/test_onnx_export.py b/test/jit/test_onnx_export.py index fdc30c22063a2..21ee19d262aab 100644 --- a/test/jit/test_onnx_export.py +++ b/test/jit/test_onnx_export.py @@ -212,23 +212,6 @@ def forward(self, x): mte, (torch.zeros(1, 2, 3, dtype=torch.float),), None, verbose=False, example_outputs=outputs) - def test_onnx_raw_export_script_truediv(self): - class ModuleToExport(torch.jit.ScriptModule): - def __init__(self): - super(ModuleToExport, self).__init__() - - @torch.jit.script_method - def forward(self, x): - z = x.size(0) / 2 - return x + z - - mte = ModuleToExport() - outputs = mte(torch.zeros(1, 2, 3)) - torch.onnx.export_to_pretty_string( - mte, (torch.zeros(1, 2, 3),), None, verbose=False, - add_node_names=False, do_constant_folding=False, - example_outputs=outputs, export_raw_ir=True) - def test_onnx_export_script_non_alpha_add_sub(self): class ModuleToExport(torch.jit.ScriptModule): def __init__(self): diff --git a/torch/_C/_onnx.pyi b/torch/_C/_onnx.pyi index 7ab3cd9c567d2..9489c9cf907b7 100644 --- a/torch/_C/_onnx.pyi +++ b/torch/_C/_onnx.pyi @@ -28,7 +28,6 @@ class OperatorExportTypes(Enum): ONNX = ... ONNX_ATEN = ... ONNX_ATEN_FALLBACK = ... - RAW = ... ONNX_FALLTHROUGH = ... class TrainingMode(Enum): diff --git a/torch/csrc/jit/serialization/export.cpp b/torch/csrc/jit/serialization/export.cpp index d53c400acfd7e..a0397015ffa72 100644 --- a/torch/csrc/jit/serialization/export.cpp +++ b/torch/csrc/jit/serialization/export.cpp @@ -472,9 +472,7 @@ void EncoderBase::EncodeBlock( EncodeValueInfo(graph_proto, v, output, dynamic_axes); } for (auto node : block->nodes()) { - bool is_raw_export = - operator_export_type_ == onnx_torch::OperatorExportTypes::RAW; - if (node->mustBeNone() && !is_raw_export) { + if (node->mustBeNone()) { // None nodes are used to implement optional inputs. One // way to "not provide" an optional input is to create an // Undefined node, and pass its output as that input. @@ -485,7 +483,7 @@ void EncoderBase::EncodeBlock( p_n->set_doc_string(node->sourceRange().str()); } for (auto input : node->inputs()) { - if (input->node()->mustBeNone() && !is_raw_export) { + if (input->node()->mustBeNone()) { p_n->add_input(""); } else { p_n->add_input(input->debugName()); @@ -505,9 +503,7 @@ void EncoderBase::EncodeBlock( domains_.insert(domain); p_n->set_domain(domain); } - if (is_raw_export) { - AT_ASSERT(!node->kind().is_onnx()); - } else if (operator_export_type_ == onnx_torch::OperatorExportTypes::ONNX) { + if (operator_export_type_ == onnx_torch::OperatorExportTypes::ONNX) { AT_ASSERT( !node->kind().is_aten() && !node->kind().is_prim() && !node->kind().is_attr()); @@ -521,15 +517,6 @@ void EncoderBase::EncodeBlock( AddAttribute( p_n, node, attr_name, use_external_data_format, onnx_file_path); } - if (is_raw_export && node->blocks().size() > 0) { - auto blocks = p_n->add_attribute(); - blocks->set_name("_blocks"); - blocks->set_type(onnx::AttributeProto_AttributeType_GRAPHS); - for (auto block : node->blocks()) { - auto graph = blocks->add_graphs(); - EncodeBlock(graph, block, initializers); - } - } if (node->kind() == ::c10::onnx::Loop) { AT_ASSERT(node->blocks().size() == 1); @@ -765,9 +752,7 @@ GraphEncoder::GraphEncoder( const std::string& onnx_file_path) : EncoderBase(operator_export_type, strip_doc), defer_weight_export_(defer_weight_export) { - if (operator_export_type != onnx_torch::OperatorExportTypes::RAW) { - validateGraph(graph, operator_export_type); - } + validateGraph(graph, operator_export_type); if (use_external_data_format) { TORCH_CHECK( diff --git a/torch/csrc/onnx/init.cpp b/torch/csrc/onnx/init.cpp index 1f988af1fb788..2223ebc2f12e4 100644 --- a/torch/csrc/onnx/init.cpp +++ b/torch/csrc/onnx/init.cpp @@ -28,7 +28,6 @@ void initONNXBindings(PyObject* module) { .value("ONNX", OperatorExportTypes::ONNX) .value("ONNX_ATEN", OperatorExportTypes::ONNX_ATEN) .value("ONNX_ATEN_FALLBACK", OperatorExportTypes::ONNX_ATEN_FALLBACK) - .value("RAW", OperatorExportTypes::RAW) .value("ONNX_FALLTHROUGH", OperatorExportTypes::ONNX_FALLTHROUGH); py::enum_(onnx, "TrainingMode") diff --git a/torch/csrc/onnx/onnx.h b/torch/csrc/onnx/onnx.h index 6779fd0282be9..2dbe40e8ad602 100644 --- a/torch/csrc/onnx/onnx.h +++ b/torch/csrc/onnx/onnx.h @@ -6,7 +6,6 @@ enum class OperatorExportTypes { ONNX, // Strict ONNX export ONNX_ATEN, // ONNX With ATen op everywhere ONNX_ATEN_FALLBACK, // ONNX export with ATen fallback - RAW, // Raw export (no ONNX) ONNX_FALLTHROUGH, // Export supported ONNX ops. Pass through unsupported ops. }; diff --git a/torch/onnx/__init__.py b/torch/onnx/__init__.py index f16da3ac89776..a52f329041863 100644 --- a/torch/onnx/__init__.py +++ b/torch/onnx/__init__.py @@ -30,7 +30,7 @@ def _export(*args, **kwargs): def export(model, args, f, export_params=True, verbose=False, training=TrainingMode.EVAL, - input_names=None, output_names=None, aten=False, export_raw_ir=False, + input_names=None, output_names=None, aten=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=True, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, @@ -118,8 +118,6 @@ def forward(self, k, x): aten (bool, default False): [DEPRECATED. use operator_export_type] export the model in aten mode. If using aten mode, all the ops original exported by the functions in symbolic_opset.py are exported as ATen ops. - export_raw_ir (bool, default False): [DEPRECATED. use operator_export_type] - export the internal IR directly instead of converting it to ONNX ops. operator_export_type (enum, default OperatorExportTypes.ONNX): OperatorExportTypes.ONNX: All ops are exported as regular ONNX ops (with ONNX namespace). @@ -146,7 +144,6 @@ def forward(self, k, x): In the above example, aten::triu is not supported in ONNX, hence exporter falls back on this op. - OperatorExportTypes.RAW: Export raw ir. OperatorExportTypes.ONNX_FALLTHROUGH: If an op is not supported in ONNX, fall through and export the operator as is, as a custom ONNX op. Using this mode, the op can be exported and implemented by @@ -273,7 +270,7 @@ def forward(self, k, x): from torch.onnx import utils return utils.export(model, args, f, export_params, verbose, training, - input_names, output_names, aten, export_raw_ir, + input_names, output_names, aten, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py index 973122513d662..b7eb09f3fdd2f 100644 --- a/torch/onnx/utils.py +++ b/torch/onnx/utils.py @@ -71,15 +71,14 @@ def select_model_mode_for_export(model, mode): def export(model, args, f, export_params=True, verbose=False, training=None, - input_names=None, output_names=None, aten=False, export_raw_ir=False, + input_names=None, output_names=None, aten=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=True, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, enable_onnx_checker=True, use_external_data_format=False): - if aten or export_raw_ir: + if aten: assert operator_export_type is None - assert aten ^ export_raw_ir - operator_export_type = OperatorExportTypes.ONNX_ATEN if aten else OperatorExportTypes.RAW + operator_export_type = OperatorExportTypes.ONNX_ATEN elif operator_export_type is None: if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE: operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK @@ -153,60 +152,59 @@ def _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop=Fa torch._C._jit_pass_lint(graph) from torch.onnx.symbolic_helper import _onnx_shape_inference, _export_onnx_opset_version - if operator_export_type != OperatorExportTypes.RAW: - torch._C._jit_pass_peephole(graph, True) - torch._C._jit_pass_lower_all_tuples(graph) - # in _jit_pass_onnx, symbolic functions are called for each node for conversion. - # However, there are nodes that cannot be converted without additional context. - # For example, the number of outputs from split (and whether it is static or dynamic) is unknown - # until the point where it is unpacked by listUnpack node. - # This pass does a preprocess, and prepares the nodes such that enough context can be received - # by the symbolic function. - torch._C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module) - torch._C._jit_pass_onnx_preprocess(graph) - - # onnx does not support tuples, so try to remove them - torch._C._jit_pass_lint(graph) - - # onnx only supports tensors, but 1 / 2 = 0.5 and tensor(1) / tensor(2) = 0 - torch._C._jit_pass_prepare_division_for_onnx(graph) - - torch._C._jit_pass_onnx_remove_print(graph) - torch._C._jit_pass_onnx_preprocess_caffe2(graph) - - if operator_export_type == OperatorExportTypes.ONNX_ATEN_FALLBACK: - torch.onnx.symbolic_helper._quantized_ops.clear() - # Unpack quantized weights for conv and linear ops and insert into graph. - torch._C._jit_pass_onnx_unpack_quantized_weights(graph, params_dict) - # Insert permutes before and after each conv op to ensure correct order. - torch._C._jit_pass_onnx_quantization_insert_permutes(graph, params_dict) - - # Find consecutive permutes that are no-ops and remove them. - torch._C._jit_pass_custom_pattern_based_rewrite_graph(""" - graph(%Pi): - %Pq = quantized::nhwc2nchw(%Pi) - %Pr = quantized::nchw2nhwc(%Pq) - return (%Pr)""", """ - graph(%Ri): - return (%Ri)""", graph) - - # onnx only supports tensors, so we turn all out number types into tensors - torch._C._jit_pass_erase_number_types(graph) - - if _onnx_shape_inference: - input_names = [] if input_names is None else input_names - dynamic_axes = {} if dynamic_axes is None else dynamic_axes - torch._C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names) - graph = torch._C._jit_pass_onnx(graph, operator_export_type) - torch._C._jit_pass_lint(graph) - - torch._C._jit_pass_onnx_scalar_type_analysis(graph, True, _export_onnx_opset_version) - torch._C._jit_pass_lint(graph) - - torch._C._jit_pass_onnx_fold_if(graph) - - torch._C._jit_pass_onnx_peephole(graph, _export_onnx_opset_version, fixed_batch_size) - torch._C._jit_pass_lint(graph) + torch._C._jit_pass_peephole(graph, True) + torch._C._jit_pass_lower_all_tuples(graph) + # in _jit_pass_onnx, symbolic functions are called for each node for conversion. + # However, there are nodes that cannot be converted without additional context. + # For example, the number of outputs from split (and whether it is static or dynamic) is unknown + # until the point where it is unpacked by listUnpack node. + # This pass does a preprocess, and prepares the nodes such that enough context can be received + # by the symbolic function. + torch._C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module) + torch._C._jit_pass_onnx_preprocess(graph) + + # onnx does not support tuples, so try to remove them + torch._C._jit_pass_lint(graph) + + # onnx only supports tensors, but 1 / 2 = 0.5 and tensor(1) / tensor(2) = 0 + torch._C._jit_pass_prepare_division_for_onnx(graph) + + torch._C._jit_pass_onnx_remove_print(graph) + torch._C._jit_pass_onnx_preprocess_caffe2(graph) + + if operator_export_type == OperatorExportTypes.ONNX_ATEN_FALLBACK: + torch.onnx.symbolic_helper._quantized_ops.clear() + # Unpack quantized weights for conv and linear ops and insert into graph. + torch._C._jit_pass_onnx_unpack_quantized_weights(graph, params_dict) + # Insert permutes before and after each conv op to ensure correct order. + torch._C._jit_pass_onnx_quantization_insert_permutes(graph, params_dict) + + # Find consecutive permutes that are no-ops and remove them. + torch._C._jit_pass_custom_pattern_based_rewrite_graph(""" + graph(%Pi): + %Pq = quantized::nhwc2nchw(%Pi) + %Pr = quantized::nchw2nhwc(%Pq) + return (%Pr)""", """ + graph(%Ri): + return (%Ri)""", graph) + + # onnx only supports tensors, so we turn all out number types into tensors + torch._C._jit_pass_erase_number_types(graph) + + if _onnx_shape_inference: + input_names = [] if input_names is None else input_names + dynamic_axes = {} if dynamic_axes is None else dynamic_axes + torch._C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names) + graph = torch._C._jit_pass_onnx(graph, operator_export_type) + torch._C._jit_pass_lint(graph) + + torch._C._jit_pass_onnx_scalar_type_analysis(graph, True, _export_onnx_opset_version) + torch._C._jit_pass_lint(graph) + + torch._C._jit_pass_onnx_fold_if(graph) + + torch._C._jit_pass_onnx_peephole(graph, _export_onnx_opset_version, fixed_batch_size) + torch._C._jit_pass_lint(graph) # graph is not a valid jit graph anymore because types have been replaced # (e.g. int with Tensor), so it now contains operators that don't actually @@ -520,16 +518,16 @@ def _model_to_graph(model, args, verbose=False, def export_to_pretty_string(model, args, f, export_params=True, verbose=False, training=None, - input_names=None, output_names=None, aten=False, export_raw_ir=False, + input_names=None, output_names=None, aten=False, operator_export_type=None, export_type=ExportTypes.PROTOBUF_FILE, example_outputs=None, google_printer=False, opset_version=None, _retain_param_name=True, keep_initializers_as_inputs=None, custom_opsets=None, add_node_names=True, do_constant_folding=True): - if aten or export_raw_ir: + if aten: assert operator_export_type is None - assert aten ^ export_raw_ir - operator_export_type = OperatorExportTypes.ONNX_ATEN if aten else OperatorExportTypes.RAW + assert aten + operator_export_type = OperatorExportTypes.ONNX_ATEN elif operator_export_type is None: operator_export_type = OperatorExportTypes.ONNX return _export_to_pretty_string(model, args, f, export_params, verbose, training, From df759a3d9ecd71f0abeecefc75a26b6888aaf7ec Mon Sep 17 00:00:00 2001 From: Bert Maher Date: Fri, 11 Jun 2021 02:22:37 -0700 Subject: [PATCH 038/305] [nnc] Do not fuse matmul/conv2d if inputs are discontiguous. (#59754) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59754 Also, if inputs are contiguous, use their Placeholders directly rather than generating contiguous Tensors from them. The rationale for this change is that aten::matmul and aten::conv2d support transposed inputs; if NNC generates a physical transpose to perform an external call, performance will be strictly worse than not fusing (sometimes dramatically so, as in the attached benchmark). Test Plan: benchmark Reviewed By: ZolotukhinM Differential Revision: D29010209 fbshipit-source-id: da6d71b155c83e8d6e306089042b6b0af8f80900 --- torch/csrc/jit/tensorexpr/kernel.cpp | 50 ++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index deafb9d6609cd..315f93062bbe6 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -230,6 +230,23 @@ bool conv2dIsSupported( } return true; } + +static bool isContiguous(const torch::jit::Value* v) { + auto const& tt = v->type()->cast(); + if (!tt) { + return false; + } + if (!tt->isComplete()) { + return false; + } + auto const& sizes = tt->sizes().concrete_sizes(); + auto const& strides = tt->strides().concrete_sizes(); + if (!sizes || !strides) { + return false; + } + return *strides == TensorType::contiguousStridesOf(*sizes); +} + // The fuser only supports conv2d with very specific properties: // - Static shapes: 4-d input and filter, 1-d bias. // - Constant strides/padding/dilation/groups @@ -250,6 +267,14 @@ bool conv2dIsSupportedJit(const torch::jit::Node* node) { GRAPH_DEBUG("some params aren't static"); return false; } + + // All inputs should be contiguous so no transposition is required. + if (!isContiguous(node->input(0)) || !isContiguous(node->input(1)) || + !isContiguous(node->input(2))) { + GRAPH_DEBUG("conv2dIsSupported: some inputs are not contiguous"); + return false; + } + return conv2dIsSupported( *input, *weight, @@ -277,6 +302,12 @@ bool matmulIsSupported(const torch::jit::Node* node) { return false; } + // Inputs should be contiguous, or the TE will needlessly transpose them. + if (!isContiguous(node->input(0)) || !isContiguous(node->input(1))) { + GRAPH_DEBUG("matmulIsSupported: Input shapes are not contiguous"); + return false; + } + return true; } @@ -3093,6 +3124,16 @@ void TensorExprKernel::genInputDebugNames() { input_name_map_ = std::move(value_to_name); } +template +static std::vector toExprHandles(const std::vector& sizes) { + std::vector dims; + dims.reserve(sizes.size()); + for (auto const& size : sizes) { + dims.emplace_back(IntImm::make(size)); + } + return dims; +} + Tensor* TensorExprKernel::bindInput(const torch::jit::Value* input) { auto const& t = input->type(); Tensor* result = nullptr; @@ -3104,6 +3145,15 @@ Tensor* TensorExprKernel::bindInput(const torch::jit::Value* input) { input->debugName() + "' are unknown"; throw malformed_input(msg); } + if (isContiguous(input)) { + Placeholder inBuffer( + "t" + input_name_map_[input], + ToDtype(static_cast(*tt->scalarType())), + toExprHandles(*tt->sizes().concrete_sizes())); + bufs_.emplace(input, inBuffer.data()); + bufferArgs_.emplace_back(inBuffer); + break; + } Placeholder inBuffer( "t" + input_name_map_[input], ToDtype(static_cast(*tt->scalarType())), From b5e832111e5e4bb3dd66d716d398b81fe70c6af0 Mon Sep 17 00:00:00 2001 From: Bert Maher Date: Fri, 11 Jun 2021 02:22:37 -0700 Subject: [PATCH 039/305] [nnc] Limit the number of inputs to a fusion group. Summary: nvrtc has a hard limit to the size of kernel parameters, and llvm has a tendency to OOM with huge parameter lists, so let's limit the number of inputs to something sensible. Test Plan: tested on pyper OOM test case: ``` flow-cli test-locally --mode=opt-split-dwarf f278102738 --name "PyPer OOM repro f277966799 f63b1f9c5c0c" --run-as-secure-group oncall_pytorch_jit --entitlement default ``` Reviewed By: ZolotukhinM Differential Revision: D29019751 fbshipit-source-id: b27f2bb5000e31a7b49ea86a6928faa0ae2ead24 --- torch/csrc/jit/passes/tensorexpr_fuser.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.cpp b/torch/csrc/jit/passes/tensorexpr_fuser.cpp index 3c1a154e7f2df..bba6971375cd9 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.cpp +++ b/torch/csrc/jit/passes/tensorexpr_fuser.cpp @@ -1116,6 +1116,17 @@ class TensorExprFuser { TORCH_INTERNAL_ASSERT( consumer->kind() == prim::TensorExprGroup || canHandle(consumer)); + // nvrtc has a limit on the number of arguments allowed in a CUDA kernel. + // The specific limit is a function of constant memory size, amount + // available to pass arguments, and some implementation dependence. Select a + // safe limit here. + constexpr size_t subgraphArgLimit = 128; + if ((consumer->inputs().size() + consumer->outputs().size() + + producer->inputs().size() + producer->outputs().size()) > + subgraphArgLimit) { + return false; + } + // Device checks if (consumer->kind() != aten::cat && producer->kind() != aten::cat) { // aten::cat needs a special handling because it takes a Tensor[] as its From b4c35d7ae78e7742dc11d5b767196d4450527d2c Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 11 Jun 2021 05:04:04 -0700 Subject: [PATCH 040/305] Remove USE_CUDA from ProcessGroupGloo (#59561) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59561 Needed to merge c10d into libtorch(_cuda). ghstack-source-id: 131169544 Test Plan: CI Reviewed By: agolynski Differential Revision: D28931379 fbshipit-source-id: 9bd68477ae7bb870b6737a555edd5696149ff5d6 --- torch/lib/c10d/ProcessGroupGloo.cpp | 310 +++++++----------- torch/lib/c10d/ProcessGroupGloo.hpp | 5 - .../c10d/test/ProcessGroupGlooAsyncTest.cpp | 1 + 3 files changed, 115 insertions(+), 201 deletions(-) diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp index ac15a06923b0d..1b413c84bb84d 100644 --- a/torch/lib/c10d/ProcessGroupGloo.cpp +++ b/torch/lib/c10d/ProcessGroupGloo.cpp @@ -33,15 +33,6 @@ #include -#ifdef USE_CUDA -#include -#include -#include -#include -#include -#include -#endif - #include #include #include @@ -309,10 +300,8 @@ void setOutput(O& opts, at::Tensor& tensor, std::vector& counts) { opts.setOutput(getDataPointer(tensor), counts); } -#ifdef USE_CUDA - at::Tensor pinnedLike(at::Tensor& tensor) { - auto* allocator = at::cuda::getPinnedMemoryAllocator(); + auto* allocator = at::detail::getCUDAHooks().getPinnedMemoryAllocator(); auto storage = c10::Storage( c10::Storage::use_byte_size_t(), at::detail::computeStorageNbytes( @@ -330,20 +319,20 @@ at::Tensor pinnedLike(at::Tensor& tensor) { // on the tensors. void initializeStreamsEvents( const std::vector& tensors, - std::vector& streams, - std::vector& events) { - at::cuda::OptionalCUDAGuard guard; + std::vector& streams, + std::vector& events) { streams.reserve(tensors.size()); - events.resize(tensors.size()); + events.reserve(tensors.size()); for(const auto i : c10::irange(tensors.size())) { - guard.set_index(tensors[i].device().index()); + c10::Device device = tensors[i].device(); + c10::impl::VirtualGuardImpl impl(device.type()); // Record event on current stream - events[i].record(at::cuda::getCurrentCUDAStream()); + events.emplace_back(device.type()); + events[i].record(impl.getStream(device)); // Get a non-default stream to execute asynchronous CUDA operations // on for this device. This ensures that the default stream used // by the caller is not occupied by c10d related operations. - streams.push_back(at::cuda::getStreamFromPool( - /* isHighPriority */ true, tensors[i].device().index())); + streams.push_back(impl.getStreamFromGlobalPool(device, /*isHighPriority=*/true)); // Ensure the new stream is synchronized with the current stream. events[i].block(streams[i]); @@ -351,18 +340,15 @@ void initializeStreamsEvents( // new streams in this Work to prevent being freed before the Work finishes. if (tensors[i].is_sparse()) { if (tensors[i].is_coalesced()) { - c10::cuda::CUDACachingAllocator::recordStream( - tensors[i].indices().storage().data_ptr(), streams[i]); - c10::cuda::CUDACachingAllocator::recordStream( - tensors[i].values().storage().data_ptr(), streams[i]); + impl.recordDataPtrOnStream(tensors[i].indices().storage().data_ptr(), streams[i]); + impl.recordDataPtrOnStream(tensors[i].values().storage().data_ptr(), streams[i]); } else { // We will need to coalesce first, which means new tensors will // be allocated on the streams we just allocated, and there // is no need to record them separately. } } else { - c10::cuda::CUDACachingAllocator::recordStream( - tensors[i].storage().data_ptr(), streams[i]); + impl.recordDataPtrOnStream(tensors[i].storage().data_ptr(), streams[i]); } } } @@ -373,8 +359,8 @@ void initializeStreamsEvents( // on the same device. void initializeStreamsEvents( std::vector>& tensors, - std::vector& streams, - std::vector& events) { + std::vector& streams, + std::vector& events) { // Ensure that the tensors in the nested tensor vectors are on the same // device. for (const auto& tensorgroup : tensors) { @@ -388,18 +374,18 @@ void initializeStreamsEvents( } } - at::cuda::OptionalCUDAGuard guard; streams.reserve(tensors.size()); - events.resize(tensors.size()); + events.reserve(tensors.size()); for(const auto i : c10::irange(tensors.size())) { - guard.set_index(tensors[i][0].device().index()); + c10::Device device = tensors[i][0].device(); + c10::impl::VirtualGuardImpl impl(device.type()); // Record event on current stream - events[i].record(at::cuda::getCurrentCUDAStream()); + events.emplace_back(device.type()); + events[i].record(impl.getStream(device)); // Get a non-default stream to execute asynchronous CUDA operations // on for this output. This ensures that the default stream used // by the caller is not occupied by c10d related operations. - streams.push_back(at::cuda::getStreamFromPool( - /* isHighPriority */ true, tensors[i][0].device().index())); + streams.push_back(impl.getStreamFromGlobalPool(device, /*isHighPriority=*/true)); // Ensure the new stream is synchronized with the current stream. events[i].block(streams[i]); @@ -407,14 +393,11 @@ void initializeStreamsEvents( // `tensors` are created on a different stream. Hence, they must record // new streams in this Work to prevent being freed before the Work // finishes. - c10::cuda::CUDACachingAllocator::recordStream( - tensor.storage().data_ptr(), streams[i]); + impl.recordDataPtrOnStream(tensor.storage().data_ptr(), streams[i]); } } } -#endif - const auto kLoopbackAddress = "127.0.0.1"; } // namespace @@ -844,8 +827,6 @@ class AsyncBroadcastWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - class AsyncBroadcastCUDAWork : public AsyncBroadcastWork { public: AsyncBroadcastCUDAWork( @@ -859,7 +840,7 @@ class AsyncBroadcastCUDAWork : public AsyncBroadcastWork { // Create pinned host side tensors. tmp = pinnedLike(inputs[rootTensor]); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; if (context->rank == rootRank) { guard.reset_stream(streams[rootTensor]); tmp.copy_(inputs[rootTensor], /* non_blocking */ true); @@ -867,18 +848,16 @@ class AsyncBroadcastCUDAWork : public AsyncBroadcastWork { } void run() override { - at::cuda::OptionalCUDAStreamGuard guard; - // Synchronize with copy operation if applicable. if (context->rank == rootRank) { - guard.reset_stream(streams[rootTensor]); - AT_CUDA_CHECK(cudaStreamSynchronize(streams[rootTensor])); + streams[rootTensor].synchronize(); } // Run broadcast on host side tensors. broadcast(tmp); // Kick off copy back to the CUDA tensors. + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(streams[i]); inputs[i].copy_(tmp, /* non_blocking */ true); @@ -887,22 +866,18 @@ class AsyncBroadcastCUDAWork : public AsyncBroadcastWork { } void synchronize() override { - at::cuda::OptionalCUDAGuard guard; - // Synchronize with the copy back to CUDA tensors. for(const auto i : c10::irange(inputs.size())) { - guard.set_index(inputs[i].device().index()); - events[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = inputs[i].device(); + events[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } at::Tensor tmp; - std::vector streams; - std::vector events; + std::vector streams; + std::vector events; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::broadcast( @@ -920,9 +895,10 @@ c10::intrusive_ptr ProcessGroupGloo::broadcast( const auto& device = inputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -934,11 +910,9 @@ c10::intrusive_ptr ProcessGroupGloo::broadcast( if (device.type() == at::kCPU) { work = c10::make_intrusive( std::move(context), inputs, opts.rootRank, opts.rootTensor, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), inputs, opts.rootRank, opts.rootTensor, tag); -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -1294,8 +1268,6 @@ class AsyncSparseAllreduceWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - class AsyncAllreduceCUDAWork : public AsyncAllreduceWork { public: AsyncAllreduceCUDAWork( @@ -1308,7 +1280,7 @@ class AsyncAllreduceCUDAWork : public AsyncAllreduceWork { // Kick off copy from CUDA tensors to pinned CPU tensors. tmp.reserve(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(streams[i]); tmp.push_back(pinnedLike(inputs[i]).copy_(inputs[i], true)); @@ -1317,18 +1289,16 @@ class AsyncAllreduceCUDAWork : public AsyncAllreduceWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i].device().index()); - AT_CUDA_CHECK(cudaStreamSynchronize(streams[i])); + streams[i].synchronize(); } // Run allreduce on host side tensors. allreduce(tmp); - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { - stream_guard.reset_stream(streams[i]); + guard.reset_stream(streams[i]); inputs[i].copy_(tmp[i], /* non_blocking */ true); events[i].record(streams[i]); } @@ -1336,16 +1306,15 @@ class AsyncAllreduceCUDAWork : public AsyncAllreduceWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(inputs.size())) { - guard.set_index(inputs[i].device().index()); - events[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = inputs[i].device(); + events[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmp; - std::vector streams; - std::vector events; + std::vector streams; + std::vector events; }; class AsyncSparseAllreduceCUDAWork : public AsyncSparseAllreduceWork { @@ -1361,7 +1330,7 @@ class AsyncSparseAllreduceCUDAWork : public AsyncSparseAllreduceWork { // Note that both coalescing the sparse tensor and copying it to CPU // memory must be performed asynchronously, or we block the caller. tmp.reserve(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(streams[i]); tmp.push_back( @@ -1371,19 +1340,17 @@ class AsyncSparseAllreduceCUDAWork : public AsyncSparseAllreduceWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i].device().index()); - AT_CUDA_CHECK(cudaStreamSynchronize(streams[i])); + streams[i].synchronize(); } // Run allreduce on host side tensors. auto output = allreduce(tmp); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { - stream_guard.reset_stream(streams[i]); + guard.reset_stream(streams[i]); inputs[i].copy_(output, /*non_blocking=*/true); events[i].record(streams[i]); } @@ -1391,20 +1358,17 @@ class AsyncSparseAllreduceCUDAWork : public AsyncSparseAllreduceWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(inputs.size())) { - guard.set_index(inputs[i].device().index()); - events[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = inputs[i].device(); + events[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmp; - std::vector streams; - std::vector events; + std::vector streams; + std::vector events; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::allreduce( @@ -1421,9 +1385,10 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( const auto& device = inputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -1449,7 +1414,6 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( } else { invalidArgument("unsupported layout"); } -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { if (layout == c10::kStrided) { work = c10::make_intrusive( @@ -1460,7 +1424,6 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( } else { invalidArgument("unsupported layout"); } -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -1584,8 +1547,6 @@ class AsyncReduceWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - class AsyncReduceCUDAWork : public AsyncReduceWork { public: AsyncReduceCUDAWork( @@ -1600,7 +1561,7 @@ class AsyncReduceCUDAWork : public AsyncReduceWork { // Kick off copy from CUDA tensors to pinned CPU tensors. tmp.reserve(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(streams[i]); tmp.push_back(pinnedLike(inputs[i]).copy_(inputs[i], true)); @@ -1609,19 +1570,17 @@ class AsyncReduceCUDAWork : public AsyncReduceWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i].device().index()); - AT_CUDA_CHECK(cudaStreamSynchronize(streams[i])); + streams[i].synchronize(); } // Run reduce on host side tensors. reduce(tmp); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { - stream_guard.reset_stream(streams[i]); + guard.reset_stream(streams[i]); inputs[i].copy_(tmp[i], /* non_blocking */ true); events[i].record(streams[i]); } @@ -1629,20 +1588,17 @@ class AsyncReduceCUDAWork : public AsyncReduceWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(inputs.size())) { - guard.set_index(inputs[i].device().index()); - events[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = inputs[i].device(); + events[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmp; - std::vector streams; - std::vector events; + std::vector streams; + std::vector events; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::reduce( @@ -1660,9 +1616,10 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( const auto& device = inputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -1679,7 +1636,6 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( opts.rootTensor, opts.reduceOp, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), @@ -1688,7 +1644,6 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( opts.rootTensor, opts.reduceOp, tag); -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -1747,8 +1702,6 @@ class AsyncAllgatherWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - // Note: current CUDA implementation holds the assumption that the // tensors in the nested output tensor vectors are on the same device. class AsyncAllgatherCUDAWork : public AsyncAllgatherWork { @@ -1764,7 +1717,7 @@ class AsyncAllgatherCUDAWork : public AsyncAllgatherWork { // Kick off copy from CUDA tensors to pinned CPU tensors. tmpInputs.reserve(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(inputStreams[i]); tmpInputs.push_back(pinnedLike(inputs[i]).copy_(inputs[i], true)); @@ -1781,24 +1734,21 @@ class AsyncAllgatherCUDAWork : public AsyncAllgatherWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i].device().index()); - AT_CUDA_CHECK(cudaStreamSynchronize(inputStreams[i])); + inputStreams[i].synchronize(); } for(const auto i : c10::irange(outputs.size())) { - device_guard.set_index(outputs[i][0].device().index()); - AT_CUDA_CHECK(cudaStreamSynchronize(outputStreams[i])); + outputStreams[i].synchronize(); } // Run allgather on host side tensors. allgather(tmpOutputs, tmpInputs); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(outputs.size())) { - stream_guard.reset_stream(outputStreams[i]); + guard.reset_stream(outputStreams[i]); for(const auto j : c10::irange(outputs[i].size())) { outputs[i][j].copy_(tmpOutputs[i][j], /* non_blocking */ true); } @@ -1808,24 +1758,21 @@ class AsyncAllgatherCUDAWork : public AsyncAllgatherWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(outputs.size())) { - guard.set_index(outputs[i][0].device().index()); - outputEvents[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = outputs[i][0].device(); + outputEvents[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmpInputs; - std::vector inputStreams; - std::vector inputEvents; + std::vector inputStreams; + std::vector inputEvents; std::vector> tmpOutputs; - std::vector outputStreams; - std::vector outputEvents; + std::vector outputStreams; + std::vector outputEvents; }; -#endif - } // namespace // Note: current CUDA implementation holds the assumption that the @@ -1871,9 +1818,10 @@ c10::intrusive_ptr ProcessGroupGloo::allgather( const auto& device = inputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -1885,11 +1833,9 @@ c10::intrusive_ptr ProcessGroupGloo::allgather( if (device.type() == at::kCPU) { work = c10::make_intrusive( std::move(context), outputs, inputs, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), outputs, inputs, tag); -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -2085,8 +2031,6 @@ class AsyncGatherWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - // Note: current CUDA implementation holds the assumptions: // - inputs.size() is 1 // - outputs.size() is 1 @@ -2106,7 +2050,7 @@ class AsyncGatherCUDAWork : public AsyncGatherWork { // Kick off copy from CUDA tensors to pinned CPU tensors. tmpInputs.reserve(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(inputStreams[i]); tmpInputs.push_back(pinnedLike(inputs[i]).copy_(inputs[i], true)); @@ -2123,24 +2067,21 @@ class AsyncGatherCUDAWork : public AsyncGatherWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i].get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(inputStreams[i])); + inputStreams[i].synchronize(); } for(const auto i : c10::irange(outputs.size())) { - device_guard.set_index(outputs[i][0].get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(outputStreams[i])); + outputStreams[i].synchronize(); } // Run gather on host side tensors. gather(tmpOutputs, tmpInputs); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(outputs.size())) { - stream_guard.reset_stream(outputStreams[i]); + guard.reset_stream(outputStreams[i]); for(const auto j : c10::irange(outputs[i].size())) { outputs[i][j].copy_(tmpOutputs[i][j], /* non_blocking */ true); } @@ -2150,24 +2091,21 @@ class AsyncGatherCUDAWork : public AsyncGatherWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(outputs.size())) { - guard.set_index(static_cast(outputs[i][0].get_device())); - outputEvents[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = outputs[i][0].device(); + outputEvents[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmpInputs; - std::vector inputStreams; - std::vector inputEvents; + std::vector inputStreams; + std::vector inputEvents; std::vector> tmpOutputs; - std::vector outputStreams; - std::vector outputEvents; + std::vector outputStreams; + std::vector outputEvents; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::gather( @@ -2208,9 +2146,10 @@ c10::intrusive_ptr ProcessGroupGloo::gather( const auto& device = inputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -2222,11 +2161,9 @@ c10::intrusive_ptr ProcessGroupGloo::gather( if (device.type() == at::kCPU) { work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -2285,8 +2222,6 @@ class AsyncScatterWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - class AsyncScatterCUDAWork : public AsyncScatterWork { public: AsyncScatterCUDAWork( @@ -2301,7 +2236,7 @@ class AsyncScatterCUDAWork : public AsyncScatterWork { // Kick off copy from CUDA tensors to pinned CPU tensors. tmpInputs.resize(inputs.size()); - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(inputs.size())) { guard.reset_stream(inputStreams[i]); tmpInputs[i].reserve(inputs[i].size()); @@ -2319,23 +2254,20 @@ class AsyncScatterCUDAWork : public AsyncScatterWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; for(const auto i : c10::irange(inputs.size())) { - device_guard.set_index(inputs[i][0].get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(inputStreams[i])); + inputStreams[i].synchronize(); } for(const auto i : c10::irange(outputs.size())) { - device_guard.set_index(outputs[i].get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(outputStreams[i])); + outputStreams[i].synchronize(); } // Run scatter on host side tensors. scatter(tmpOutputs, tmpInputs); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; + c10::OptionalStreamGuard guard; for(const auto i : c10::irange(outputs.size())) { - stream_guard.reset_stream(outputStreams[i]); + guard.reset_stream(outputStreams[i]); outputs[i].copy_(tmpOutputs[i], /* non_blocking */ true); outputEvents[i].record(outputStreams[i]); } @@ -2343,24 +2275,21 @@ class AsyncScatterCUDAWork : public AsyncScatterWork { void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; for(const auto i : c10::irange(outputs.size())) { - guard.set_index(static_cast(outputs[i].get_device())); - outputEvents[i].block(at::cuda::getCurrentCUDAStream()); + c10::Device device = outputs[i].device(); + outputEvents[i].block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } } std::vector tmpOutputs; - std::vector outputStreams; - std::vector outputEvents; + std::vector outputStreams; + std::vector outputEvents; std::vector> tmpInputs; - std::vector inputStreams; - std::vector inputEvents; + std::vector inputStreams; + std::vector inputEvents; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::scatter( @@ -2400,9 +2329,10 @@ c10::intrusive_ptr ProcessGroupGloo::scatter( const auto& device = outputs[0].device(); switch (device.type()) { case at::kCPU: -#ifdef USE_CUDA + break; case at::kCUDA: -#endif + // If the user gave us a CUDA tensor then CUDA must be loaded. + TORCH_INTERNAL_ASSERT(at::hasCUDA()); break; default: invalidArgument(c10::str("unsupported device type ", device.type())); @@ -2414,11 +2344,9 @@ c10::intrusive_ptr ProcessGroupGloo::scatter( if (device.type() == at::kCPU) { work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); -#endif } else { throw std::runtime_error("Invalid backend"); } @@ -2496,8 +2424,6 @@ class AsyncAlltoallWork : public ProcessGroupGloo::AsyncWork { } }; -#ifdef USE_CUDA - class AsyncAlltoallCUDAWork : public AsyncAlltoallWork { public: AsyncAlltoallCUDAWork( @@ -2518,7 +2444,7 @@ class AsyncAlltoallCUDAWork : public AsyncAlltoallWork { initializeStreamsEvents({outputTensor}, outputStreams, outputEvents); // Kick off copy from CUDA tensors to pinned CPU tensors. - at::cuda::OptionalCUDAStreamGuard guard; + c10::OptionalStreamGuard guard; guard.reset_stream(inputStreams.front()); cpuInput = pinnedLike(inputTensor).copy_(inputTensor, true); @@ -2528,40 +2454,34 @@ class AsyncAlltoallCUDAWork : public AsyncAlltoallWork { void run() override { // Synchronize with copy operations. - at::cuda::OptionalCUDAGuard device_guard; - device_guard.set_index(inputTensor.get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(inputStreams.front())); - device_guard.set_index(outputTensor.get_device()); - AT_CUDA_CHECK(cudaStreamSynchronize(outputStreams.front())); + inputStreams.front().synchronize(); + outputStreams.front().synchronize(); // Run alltoall on host side tensors. alltoall(cpuOutput, cpuInput); // Kick off copy back to the CUDA tensors. - at::cuda::OptionalCUDAStreamGuard stream_guard; - stream_guard.reset_stream(outputStreams.front()); + c10::OptionalStreamGuard guard; + guard.reset_stream(outputStreams.front()); outputTensor.copy_(cpuOutput, /* non_blocking */ true); outputEvents.front().record(outputStreams.front()); } void synchronize() override { // Synchronize with the copy back to CUDA tensors. - at::cuda::OptionalCUDAGuard guard; - guard.set_index(static_cast(outputTensor.get_device())); - outputEvents.front().block(at::cuda::getCurrentCUDAStream()); + c10::Device device = outputTensor.device(); + outputEvents.front().block(c10::impl::VirtualGuardImpl(device.type()).getStream(device)); } at::Tensor cpuOutput; - std::vector outputStreams; - std::vector outputEvents; + std::vector outputStreams; + std::vector outputEvents; at::Tensor cpuInput; - std::vector inputStreams; - std::vector inputEvents; + std::vector inputStreams; + std::vector inputEvents; }; -#endif - } // namespace c10::intrusive_ptr ProcessGroupGloo::alltoall_base( @@ -2593,7 +2513,6 @@ c10::intrusive_ptr ProcessGroupGloo::alltoall_base( outputCounts, inputCounts, tag); -#ifdef USE_CUDA } else if (device.type() == at::kCUDA) { work = c10::make_intrusive( std::move(context), @@ -2602,7 +2521,6 @@ c10::intrusive_ptr ProcessGroupGloo::alltoall_base( outputCounts, inputCounts, tag); -#endif } else { invalidArgument(c10::str("unsupported device type ", device.type())); } diff --git a/torch/lib/c10d/ProcessGroupGloo.hpp b/torch/lib/c10d/ProcessGroupGloo.hpp index fad1a3cb27ed1..eeaf86869c1dd 100644 --- a/torch/lib/c10d/ProcessGroupGloo.hpp +++ b/torch/lib/c10d/ProcessGroupGloo.hpp @@ -16,11 +16,6 @@ #include -#ifdef USE_CUDA -#include -#include -#endif - #include #include #include diff --git a/torch/lib/c10d/test/ProcessGroupGlooAsyncTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooAsyncTest.cpp index 8b26758aa5396..edfdbe8e88427 100644 --- a/torch/lib/c10d/test/ProcessGroupGlooAsyncTest.cpp +++ b/torch/lib/c10d/test/ProcessGroupGlooAsyncTest.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include #include From cbcae46fa5df47c5b587b3904c7d33ec75655fa6 Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 11 Jun 2021 05:04:04 -0700 Subject: [PATCH 041/305] Remove USE_CUDA from c10d reducer/logger (#59562) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59562 Needed to merge c10d into libtorch(_cuda). ghstack-source-id: 131169542 Test Plan: CI Reviewed By: agolynski Differential Revision: D28931378 fbshipit-source-id: 71376b862ff6ef7dbfa7331ec8d269bd3fcc7e0d --- tools/build_variables.bzl | 1 + torch/lib/c10d/logger.cpp | 157 ++++++++------------------------ torch/lib/c10d/logger.hpp | 14 +-- torch/lib/c10d/reducer.cpp | 131 +++++++++++++++----------- torch/lib/c10d/reducer.hpp | 59 +++++------- torch/lib/c10d/reducer_cuda.cpp | 89 ++++++++++++++++++ 6 files changed, 237 insertions(+), 214 deletions(-) create mode 100644 torch/lib/c10d/reducer_cuda.cpp diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 718809fa4da21..0a9aa427ef801 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -669,6 +669,7 @@ libtorch_python_distributed_core_sources = [ "torch/lib/c10d/comm.cpp", "torch/lib/c10d/default_comm_hooks.cpp", "torch/lib/c10d/reducer.cpp", + "torch/lib/c10d/reducer_cuda.cpp", "torch/lib/c10d/logger.cpp", "torch/csrc/distributed/c10d/python_comm_hook.cpp", "torch/csrc/distributed/c10d/init.cpp", diff --git a/torch/lib/c10d/logger.cpp b/torch/lib/c10d/logger.cpp index 6cc4614947fcf..14a1295574630 100644 --- a/torch/lib/c10d/logger.cpp +++ b/torch/lib/c10d/logger.cpp @@ -9,12 +9,6 @@ namespace c10d { // stats. const int LoggingIterations[] = {10, 20, 100, 1000}; -namespace { - -const int kMilliSecondToNanosSecond = 1000000; - -} // anonymous namespace - std::ostream& operator<<(std::ostream& output, const Logger& logger) { auto& ddp_logging_data = (*logger.ddp_logging_data_); @@ -176,48 +170,21 @@ void Logger::set_construction_data_and_log( at::LogPyTorchDDPUsage(*ddp_logging_data_); } -void Logger::calculate_avg_cpu_time( - int64_t& avg_time, - int64_t& time_duration, - int64_t cpu_start_time, - int64_t cpu_end_time) { - // If cpu_end_time is not recorded in this iteration, - // avg_time will return invalid value. - // For some cases like DDP runs on non-sync mode, backward compute - // end time can not be recorded in this iteration and thus can not - // calculate the valid avg_time. - // In this case, skip calculating the avg_time and return. - TORCH_CHECK(num_iterations_stats_recorded_ > 0); - if (cpu_end_time < cpu_start_time) { - return; - } - time_duration = cpu_end_time - cpu_start_time; - avg_time = (time_duration + avg_time * (num_iterations_stats_recorded_ - 1)) / - num_iterations_stats_recorded_; -} - -#ifdef USE_CUDA -void Logger::calculate_avg_gpu_time( +void Logger::calculate_avg_time( int64_t& avg_time, int64_t& time_duration, - at::cuda::CUDAEvent& gpu_start, - at::cuda::CUDAEvent& gpu_end) { + Timer& timer, + Timer::Event start_event, + Timer::Event end_event) { TORCH_CHECK(num_iterations_stats_recorded_ > 0); - float milliseconds = gpu_start.elapsed_time(gpu_end); - // If gpu_end is not recorded in this iteration, - // milliseconds will have invalid value. - // For some cases like DDP runs on non-sync mode, - // gpu_end can not be recorded in this iteration and thus can not - // calculate the valid avg_time. - // In this case, skip calculating the avg_time and return. - if (milliseconds < 0) { + c10::optional maybe_time_duration = timer.measureDifference(start_event, end_event); + if (!maybe_time_duration.has_value()) { return; } - time_duration = int64_t(milliseconds * kMilliSecondToNanosSecond); + time_duration = maybe_time_duration.value(); avg_time = (time_duration + avg_time * (num_iterations_stats_recorded_ - 1)) / num_iterations_stats_recorded_; } -#endif void Logger::reset_performance_stats() { ddp_logging_data_->ints_map["forward_compute_time"] = 0; @@ -260,85 +227,39 @@ void Logger::set_runtime_stats_and_log() { reset_performance_stats(); - if (reducer_->replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Cuda time stats are only collected for single device modules. - if (reducer_->is_multi_device_module_) { - TORCH_WARN_ONCE( - "Cuda time stats are not collected for multi-device modules." - ); - return; - } - // Check events on the replicas_[0][0].device(). - at::DeviceGuard g(reducer_->replicas_[0][0].device()); - // It is possible users did not call backward or run codes in - // no-sync mode, in this case, some cudaEvents like "backward_compute_end" - // or "backward_comm_start" or "backward_comm_end" will not be recorded. - // cudaEvent is created when it is first time to be recorded. - // If it is never recorded/created, skip synchronize and calculation. - // Otherwise it will throw cuda errors. - if (!reducer_->gpu_timer_.forward_start.isCreated() || - !reducer_->gpu_timer_.backward_compute_start.isCreated() || - !reducer_->gpu_timer_.backward_compute_end.isCreated() || - !reducer_->gpu_timer_.backward_comm_start.isCreated() || - !reducer_->gpu_timer_.backward_comm_end.isCreated()) { - return; - } - - // set_runtime_stats_and_log is called at the beginning of forward call, - // when it is cheap to synchronize the cuda events of previous iteration, - // as mostly all cuda operations are finished in previous iteration. - reducer_->gpu_timer_.forward_start.synchronize(); - reducer_->gpu_timer_.backward_compute_start.synchronize(); - reducer_->gpu_timer_.backward_compute_end.synchronize(); - reducer_->gpu_timer_.backward_comm_start.synchronize(); - reducer_->gpu_timer_.backward_comm_end.synchronize(); - calculate_avg_gpu_time( - ddp_logging_data_->ints_map["avg_forward_compute_time"], - ddp_logging_data_->ints_map["forward_compute_time"], - reducer_->gpu_timer_.forward_start, - reducer_->gpu_timer_.backward_compute_start); - calculate_avg_gpu_time( - ddp_logging_data_->ints_map["avg_backward_compute_time"], - ddp_logging_data_->ints_map["backward_compute_time"], - reducer_->gpu_timer_.backward_compute_start, - reducer_->gpu_timer_.backward_compute_end); - calculate_avg_gpu_time( - ddp_logging_data_->ints_map["avg_backward_comm_time"], - ddp_logging_data_->ints_map["backward_comm_time"], - reducer_->gpu_timer_.backward_comm_start, - reducer_->gpu_timer_.backward_comm_end); - calculate_avg_gpu_time( - ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"], - ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"], - reducer_->gpu_timer_.backward_comm_start, - reducer_->gpu_timer_.backward_compute_end); -#endif - } else { - calculate_avg_cpu_time( - ddp_logging_data_->ints_map["avg_forward_compute_time"], - ddp_logging_data_->ints_map["forward_compute_time"], - reducer_->cpu_timer_.forward_start_time, - reducer_->cpu_timer_.backward_compute_start_time); - - calculate_avg_cpu_time( - ddp_logging_data_->ints_map["avg_backward_compute_time"], - ddp_logging_data_->ints_map["backward_compute_time"], - reducer_->cpu_timer_.backward_compute_start_time, - reducer_->cpu_timer_.backward_compute_end_time); - - calculate_avg_cpu_time( - ddp_logging_data_->ints_map["avg_backward_comm_time"], - ddp_logging_data_->ints_map["backward_comm_time"], - reducer_->cpu_timer_.backward_comm_start_time, - reducer_->cpu_timer_.backward_comm_end_time); - - calculate_avg_cpu_time( - ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"], - ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"], - reducer_->cpu_timer_.backward_comm_start_time, - reducer_->cpu_timer_.backward_compute_end_time); + // Cuda time stats are only collected for single device modules. + if (reducer_->replicas_[0][0].is_cuda() && reducer_->is_multi_device_module_) { + TORCH_WARN_ONCE( + "Cuda time stats are not collected for multi-device modules." + ); + return; } + TORCH_INTERNAL_ASSERT(reducer_->timer_); + calculate_avg_time( + ddp_logging_data_->ints_map["avg_forward_compute_time"], + ddp_logging_data_->ints_map["forward_compute_time"], + *reducer_->timer_, + Timer::Event::kForwardStart, + Timer::Event::kBackwardComputeStart); + calculate_avg_time( + ddp_logging_data_->ints_map["avg_backward_compute_time"], + ddp_logging_data_->ints_map["backward_compute_time"], + *reducer_->timer_, + Timer::Event::kBackwardComputeStart, + Timer::Event::kBackwardComputeEnd); + calculate_avg_time( + ddp_logging_data_->ints_map["avg_backward_comm_time"], + ddp_logging_data_->ints_map["backward_comm_time"], + *reducer_->timer_, + Timer::Event::kBackwardCommStart, + Timer::Event::kBackwardCommEnd); + calculate_avg_time( + ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"], + ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"], + *reducer_->timer_, + Timer::Event::kBackwardCommStart, + Timer::Event::kBackwardComputeEnd); + // Log runtime stats to stderr if TORCH_DISTRIBUTED_DEBUG=DETAIL is enabled. if (parseDistDebugLevel() == DistributedDebugLevel::DETAIL) { LOG(INFO) << *this; diff --git a/torch/lib/c10d/logger.hpp b/torch/lib/c10d/logger.hpp index fe400267ab0b3..1895e0aabdfcb 100644 --- a/torch/lib/c10d/logger.hpp +++ b/torch/lib/c10d/logger.hpp @@ -41,18 +41,12 @@ class Logger { // Calculate avg stats using cpu timer and gpu timer // that has been recorded in reducer. - void calculate_avg_cpu_time( + void calculate_avg_time( int64_t& avg_time, int64_t& time_duration, - int64_t cpu_start_time, - int64_t cpu_end_time); -#ifdef USE_CUDA - void calculate_avg_gpu_time( - int64_t& avg_time, - int64_t& time_duration, - at::cuda::CUDAEvent& gpu_start, - at::cuda::CUDAEvent& gpu_end); -#endif + Timer& timer, + Timer::Event start_event, + Timer::Event end_event); // Set stats that can be collected only during // training loop. It is called at the beginning of forward call // to record the run time stats of sampled iterations that previouly ran. diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp index 383c91b12efde..28312b6306101 100644 --- a/torch/lib/c10d/reducer.cpp +++ b/torch/lib/c10d/reducer.cpp @@ -38,6 +38,67 @@ constexpr int kUnsetDivFactor = -1; } // namespace +C10_DEFINE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device); + +namespace { + +class CpuTimer : public Timer { + private: + // The timestamp of forward call start time in each iteration. + int64_t forward_start_time = -1; + // The timestamp of backward computation start and end time in each + // iteration. + int64_t backward_compute_start_time = -1; + int64_t backward_compute_end_time = -1; + // The timestamp of first communication call start time in each iteration. + int64_t backward_comm_start_time = -1; + // The timestamp of last communication call end time in each iteration. + int64_t backward_comm_end_time = -1; + + int64_t& getTime(Event event) { + switch (event) { + case Event::kForwardStart: + return forward_start_time; + case Event::kBackwardComputeStart: + return backward_compute_start_time; + case Event::kBackwardComputeEnd: + return backward_compute_end_time; + case Event::kBackwardCommStart: + return backward_comm_start_time; + case Event::kBackwardCommEnd: + return backward_comm_end_time; + default: + TORCH_INTERNAL_ASSERT(false); + } + } + + public: + explicit CpuTimer(c10::Device /* unused */) {} + + void record(Event event) override { + getTime(event) = current_time_in_nanos(); + } + + c10::optional measureDifference(Event start, Event end) override { + int64_t start_time = getTime(start); + int64_t end_time = getTime(end); + // If cpu_end_time is not recorded in this iteration, + // avg_time will return invalid value. + // For some cases like DDP runs on non-sync mode, backward compute + // end time can not be recorded in this iteration and thus can not + // calculate the valid avg_time. + // In this case, skip calculating the avg_time and return. + if (end_time < start_time) { + return c10::nullopt; + } + return end_time - start_time; + } +}; + +C10_REGISTER_TYPED_CLASS(TimerRegistry, c10::kCPU, CpuTimer); + +} // namespace + Reducer::Reducer( std::vector> replicas, std::vector> bucket_indices, @@ -88,6 +149,12 @@ Reducer::Reducer( } } + // For CUDA, record events only for single device module. + c10::Device device = replicas_[0][0].device(); + if (!(device.is_cuda() && is_multi_device_module_)) { + timer_ = TimerRegistry()->Create(device.type(), device); + } + // If `expect_sparse_gradients` is not specified, initialize it such that // we do not expect sparse gradients for any parameter. if (expect_sparse_gradients_.empty()) { @@ -724,7 +791,7 @@ void Reducer::mark_variable_ready(size_t variable_index) { checkAndRaiseMarkedTwiceError(variable_index); perIterationReadyParams_.insert(variable_index); backward_stats_[0][variable_index] = - current_time_in_nanos() - cpu_timer_.backward_compute_start_time; + current_time_in_nanos() - backward_compute_start_time_; // Any time we mark a variable ready (be it in line due to unused parameters, // or via an autograd hook), we require a call to the finalize function. If @@ -1180,7 +1247,7 @@ void Reducer::prepare_for_backward( const std::vector& outputs) { std::lock_guard lock(mutex_); - cpu_timer_.backward_compute_start_time = current_time_in_nanos(); + backward_compute_start_time_ = current_time_in_nanos(); if (should_collect_runtime_stats()) { record_backward_compute_start_time(); } @@ -1723,72 +1790,32 @@ bool Reducer::should_collect_runtime_stats() { } void Reducer::record_forward_compute_start_time() { - if (replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Record event only for single device module. - if (!is_multi_device_module_) { - // Create and record event on the replicas_[0][0].device(). - at::DeviceGuard g(replicas_[0][0].device()); - gpu_timer_.forward_start.record(); - } -#endif - } else { - cpu_timer_.forward_start_time = current_time_in_nanos(); + if (timer_) { + timer_->record(Timer::Event::kForwardStart); } } void Reducer::record_backward_compute_start_time() { - if (replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Record event only for single device module. - if (!is_multi_device_module_) { - // Create and record event on the replicas_[0][0].device(). - at::DeviceGuard g(replicas_[0][0].device()); - gpu_timer_.backward_compute_start.record(); - } -#endif + if (timer_) { + timer_->record(Timer::Event::kBackwardComputeStart); } } void Reducer::record_backward_compute_end_time() { - if (replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Record event only for single device module. - if (!is_multi_device_module_) { - at::DeviceGuard g(replicas_[0][0].device()); - gpu_timer_.backward_compute_end.record(); - } -#endif - } else { - cpu_timer_.backward_compute_end_time = current_time_in_nanos(); + if (timer_) { + timer_->record(Timer::Event::kBackwardComputeEnd); } } void Reducer::record_backward_comm_start_time() { - if (replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Record event only for single device module - if (!is_multi_device_module_) { - at::DeviceGuard g(replicas_[0][0].device()); - gpu_timer_.backward_comm_start.record(); - } -#endif - } else { - cpu_timer_.backward_comm_start_time = current_time_in_nanos(); + if (timer_) { + timer_->record(Timer::Event::kBackwardCommStart); } } void Reducer::record_backward_comm_end_time() { - if (replicas_[0][0].is_cuda()) { -#ifdef USE_CUDA - // Record event only for single device module. - if (!is_multi_device_module_) { - at::DeviceGuard g(replicas_[0][0].device()); - gpu_timer_.backward_comm_end.record(); - } -#endif - } else { - cpu_timer_.backward_comm_end_time = current_time_in_nanos(); + if (timer_) { + timer_->record(Timer::Event::kBackwardCommEnd); } } diff --git a/torch/lib/c10d/reducer.hpp b/torch/lib/c10d/reducer.hpp index 72d925c1e2adf..9d0380800cb1e 100644 --- a/torch/lib/c10d/reducer.hpp +++ b/torch/lib/c10d/reducer.hpp @@ -1,9 +1,6 @@ #pragma once #include -#ifdef USE_CUDA -#include -#endif #include #include #include @@ -29,6 +26,28 @@ constexpr int kDDPRuntimeLoggingSampleRate = 100; // Forward declaration class Logger; +class Timer { + public: + enum class Event { + kForwardStart, + kBackwardComputeStart, + kBackwardComputeEnd, + kBackwardCommStart, + kBackwardCommEnd, + }; + + // Record the current event, i.e., mark it as having occurred now. + virtual void record(Event event) = 0; + + // Return the difference between when two events occurred, in nanoseconds. + // Or nullopt if one of them hasn't been recorded. + virtual c10::optional measureDifference(Event start, Event end) = 0; + + virtual ~Timer() = default; +}; + +C10_DECLARE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device); + class Reducer { public: // The constructor takes a list of variables for every model replica. @@ -335,37 +354,9 @@ class Reducer { // communication calls like allReduce or communication hooks. int num_buckets_ready_; - // CPU timestamp to record event start and end time. - struct CPUTimer { - // The timestamp of forward call start time in each iteration. - int64_t forward_start_time; - // The timestamp of backward computation start and end time in each - // iteration. - int64_t backward_compute_start_time; - int64_t backward_compute_end_time; - // The timestamp of first communication call start time in each iteration. - int64_t backward_comm_start_time; - // The timestamp of last communication call end time in each iteration. - int64_t backward_comm_end_time; - }; - - CPUTimer cpu_timer_{}; - -#ifdef USE_CUDA - // GPU events to record event start and end time. - struct GPUTimer { - at::cuda::CUDAEvent forward_start = at::cuda::CUDAEvent(cudaEventDefault); - at::cuda::CUDAEvent backward_compute_start = - at::cuda::CUDAEvent(cudaEventDefault); - at::cuda::CUDAEvent backward_compute_end = - at::cuda::CUDAEvent(cudaEventDefault); - at::cuda::CUDAEvent backward_comm_start = - at::cuda::CUDAEvent(cudaEventDefault); - at::cuda::CUDAEvent backward_comm_end = - at::cuda::CUDAEvent(cudaEventDefault); - }; - GPUTimer gpu_timer_; -#endif + // Timing information. + int64_t backward_compute_start_time_ = -1; + std::unique_ptr timer_; // We collect the relative timestamp of every gradient being ready // when executing autograd. This can be used to derive a timeline of diff --git a/torch/lib/c10d/reducer_cuda.cpp b/torch/lib/c10d/reducer_cuda.cpp new file mode 100644 index 0000000000000..0f55b5a131181 --- /dev/null +++ b/torch/lib/c10d/reducer_cuda.cpp @@ -0,0 +1,89 @@ +#include + +#ifdef USE_CUDA + +#include +#include + +namespace c10d { +namespace { + +const int kMilliSecondToNanosSecond = 1000000; + +class CudaTimer : public Timer { + private: + c10::Device device; + + at::cuda::CUDAEvent forward_start = at::cuda::CUDAEvent(cudaEventDefault); + at::cuda::CUDAEvent backward_compute_start = + at::cuda::CUDAEvent(cudaEventDefault); + at::cuda::CUDAEvent backward_compute_end = + at::cuda::CUDAEvent(cudaEventDefault); + at::cuda::CUDAEvent backward_comm_start = + at::cuda::CUDAEvent(cudaEventDefault); + at::cuda::CUDAEvent backward_comm_end = + at::cuda::CUDAEvent(cudaEventDefault); + + at::cuda::CUDAEvent& getEvent(Event event) { + switch (event) { + case Event::kForwardStart: + return forward_start; + case Event::kBackwardComputeStart: + return backward_compute_start; + case Event::kBackwardComputeEnd: + return backward_compute_end; + case Event::kBackwardCommStart: + return backward_comm_start; + case Event::kBackwardCommEnd: + return backward_comm_end; + default: + TORCH_INTERNAL_ASSERT(false); + } + } + + public: + explicit CudaTimer(c10::Device dev) : device(dev) {} + + void record(Event event) override { + c10::DeviceGuard g(device); + getEvent(event).record(); + } + + c10::optional measureDifference(Event start, Event end) override { + c10::DeviceGuard g(device); + at::cuda::CUDAEvent& start_event = getEvent(start); + at::cuda::CUDAEvent& end_event = getEvent(end); + // It is possible users did not call backward or run codes in + // no-sync mode, in this case, some cudaEvents like "backward_compute_end" + // or "backward_comm_start" or "backward_comm_end" will not be recorded. + // cudaEvent is created when it is first time to be recorded. + // If it is never recorded/created, skip synchronize and calculation. + // Otherwise it will throw cuda errors. + if (!start_event.isCreated() || !end_event.isCreated()) { + return c10::nullopt; + } + // set_runtime_stats_and_log is called at the beginning of forward call, + // when it is cheap to synchronize the cuda events of previous iteration, + // as mostly all cuda operations are finished in previous iteration. + start_event.synchronize(); + end_event.synchronize(); + float milliseconds = start_event.elapsed_time(end_event); + // If gpu_end is not recorded in this iteration, + // milliseconds will have invalid value. + // For some cases like DDP runs on non-sync mode, + // gpu_end can not be recorded in this iteration and thus can not + // calculate the valid avg_time. + // In this case, skip calculating the avg_time and return. + if (milliseconds < 0) { + return c10::nullopt; + } + return int64_t(milliseconds * kMilliSecondToNanosSecond); + } +}; + +C10_REGISTER_TYPED_CLASS(TimerRegistry, c10::kCUDA, CudaTimer); + +} // namespace +} // namespace c10d + +#endif From 773b56e7191f0e8412da5c8fa17252341eff536d Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 11 Jun 2021 05:04:04 -0700 Subject: [PATCH 042/305] Fix Windows guards in c10d (#59696) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59696 Some files in c10d refer to dist autograd. However, on Windows, dist autograd isn't built. Hence we need to "mask out" those references under Windows. This was already partly done, but when moving c10d to libtorch some issues came up, possibly due to the different way in which linking happens. Hence I masked out the remaining references. ghstack-source-id: 131169541 Test Plan: CI Reviewed By: agolynski Differential Revision: D28987579 fbshipit-source-id: c29c5330f8429d699554972d30f99a89b2e3971d --- torch/lib/c10d/reducer.cpp | 8 ++++++-- torch/lib/c10d/reducer.hpp | 23 ++++++++++++++++++++--- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp index 28312b6306101..c7f2990f59b77 100644 --- a/torch/lib/c10d/reducer.cpp +++ b/torch/lib/c10d/reducer.cpp @@ -1493,17 +1493,20 @@ void Reducer::finalize_backward() { void Reducer::runGradCallbackForVariable( at::Tensor& variable, GradCallback&& cb) { +#ifdef _WIN32 + cb(variable.mutable_grad()); +#else auto context_ptr = rpc_context_.context_ptr.load(); if (context_ptr == nullptr) { cb(variable.mutable_grad()); } else { // Under distributed autograd -#ifndef _WIN32 context_ptr->runGradCallbackForVariable(variable, std::move(cb)); -#endif } +#endif } +#ifndef _WIN32 void Reducer::RpcContext::set(ContextPtr&& new_context_ptr) { // We should set 'new_context_ptr' even if it's nullptr. That means the // reducer is under a local backward run. @@ -1515,6 +1518,7 @@ void Reducer::RpcContext::set(ContextPtr&& new_context_ptr) { context_ptr_holder = std::move(new_context_ptr); } } +#endif void Reducer::sync_bucket_indices( std::vector>& bucket_indices) { diff --git a/torch/lib/c10d/reducer.hpp b/torch/lib/c10d/reducer.hpp index 9d0380800cb1e..e9275d0d0555b 100644 --- a/torch/lib/c10d/reducer.hpp +++ b/torch/lib/c10d/reducer.hpp @@ -7,6 +7,8 @@ #include #include +#include +#include #include #include #include @@ -14,7 +16,9 @@ #include #include #include +#ifndef _WIN32 #include +#endif namespace c10d { @@ -224,8 +228,19 @@ class Reducer { // the buckets void sync_bucket_indices(std::vector>& bucket_indices); - using GradCallback = - torch::distributed::autograd::DistAutogradContext::GradCallback; + // We'd like to use DistAutogradContext::GradCallback here but dist autograd + // doesn't exist under Windows. So we just directly use the concrete type but + // to preserve and enforce our original intent we do a static assert when dist + // autograd is available. + using GradCallback = std::function; +#ifndef _WIN32 + static_assert( + std::is_same< + GradCallback, + torch::distributed::autograd::DistAutogradContext::GradCallback>:: + value, + ""); +#endif void runGradCallbackForVariable(at::Tensor& variable, GradCallback&& cb); // A bucket replica represents [1..N] gradients to be reduced, @@ -318,7 +333,7 @@ class Reducer { // Keep future work handle around DDP comm hook. // If no hook is registered, a temporary vanilla allreduce hook will be // used. - c10::intrusive_ptr future_work; + c10::intrusive_ptr future_work; // If this bucket should expect a single sparse gradient. // Implies: replicas[i].variables.size() == 1. @@ -381,6 +396,7 @@ class Reducer { std::vector rebuilt_param_indices_; const int64_t bucket_bytes_cap_; +#ifndef _WIN32 struct RpcContext { using ContextPtr = torch::distributed::autograd::ContextPtr; // The shared_ptr is to hold the context instance. @@ -390,6 +406,7 @@ class Reducer { void set(ContextPtr&& new_context_ptr); }; RpcContext rpc_context_; +#endif // A struct containing work handle and tensor for allreduce scheduled in // forward pass, if applicable. From c9e4d1372fdbbbd908661a61de9a68e0044a0a1d Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 11 Jun 2021 05:04:04 -0700 Subject: [PATCH 043/305] Add guards for USE_C10D_FOO in relevant c10d files (#59697) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59697 The c10d build process selectively adds files based on the `USE_C10D_FOO` flags (where `FOO` is one of `GLOO`, `NCCL` or `MPI`). Replicating this logic inside libtorch will be harder, since libtorch uses a simpler approach (i.e., it lists the files in `build_variables.bzl`). So instead we could always include all files, and "disable" each file as needed using `#ifdef`s. Note that this is not a new approach: we already do the same for all the files of the TensorPipe agent based on the flag `USE_TENSORPIPE`. ghstack-source-id: 131169540 Test Plan: CI Reviewed By: agolynski Differential Revision: D28987577 fbshipit-source-id: 4c6195de4e9a58101dad9379537e8d055dfd38af --- BUILD.bazel | 2 -- torch/lib/c10d/GlooDeviceFactory.cpp | 4 ++++ torch/lib/c10d/GlooDeviceFactory.hpp | 4 ++++ torch/lib/c10d/NCCLUtils.cpp | 5 +++++ torch/lib/c10d/NCCLUtils.hpp | 4 ++++ torch/lib/c10d/ProcessGroupGloo.cpp | 5 ++++- torch/lib/c10d/ProcessGroupGloo.hpp | 4 ++++ torch/lib/c10d/ProcessGroupMPI.cpp | 4 ++++ torch/lib/c10d/ProcessGroupMPI.hpp | 4 ++++ torch/lib/c10d/ProcessGroupNCCL.cpp | 8 ++++++-- torch/lib/c10d/ProcessGroupNCCL.hpp | 4 ++++ torch/lib/c10d/ProcessGroupWrapper.cpp | 7 ++++++- torch/lib/c10d/ProcessGroupWrapper.hpp | 4 ++++ 13 files changed, 53 insertions(+), 6 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 2cdf157555aa9..8a116d6c5cc8e 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1875,8 +1875,6 @@ cc_library( "torch/lib/c10d/*.hpp", ], exclude = [ - "torch/lib/c10d/ProcessGroupMPI.hpp", - "torch/lib/c10d/ProcessGroupNCCL.hpp", "torch/csrc/autograd/generated/VariableType.h", "torch/csrc/autograd/generated/RegistrationDeclarations.h", "torch/csrc/autograd/generated/variable_factories.h", diff --git a/torch/lib/c10d/GlooDeviceFactory.cpp b/torch/lib/c10d/GlooDeviceFactory.cpp index 01a6445a8dd3f..416676483e182 100644 --- a/torch/lib/c10d/GlooDeviceFactory.cpp +++ b/torch/lib/c10d/GlooDeviceFactory.cpp @@ -1,5 +1,7 @@ #include +#ifdef USE_C10D_GLOO + #include #include @@ -162,3 +164,5 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: } } // namespace c10d + +#endif // USE_C10D_GLOO diff --git a/torch/lib/c10d/GlooDeviceFactory.hpp b/torch/lib/c10d/GlooDeviceFactory.hpp index 58821a9340a6b..7d038180bfdb0 100644 --- a/torch/lib/c10d/GlooDeviceFactory.hpp +++ b/torch/lib/c10d/GlooDeviceFactory.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_GLOO + #include #include @@ -26,3 +28,5 @@ C10_DECLARE_SHARED_REGISTRY( const std::string& /* hostname */); } // namespace c10d + +#endif // USE_C10D_GLOO diff --git a/torch/lib/c10d/NCCLUtils.cpp b/torch/lib/c10d/NCCLUtils.cpp index d0245fe545221..9e0566a8ae528 100644 --- a/torch/lib/c10d/NCCLUtils.cpp +++ b/torch/lib/c10d/NCCLUtils.cpp @@ -1,4 +1,7 @@ #include + +#ifdef USE_C10D_NCCL + #include namespace c10d { @@ -32,3 +35,5 @@ std::string ncclGetErrorWithVersion(ncclResult_t error) { } } // namespace c10d + +#endif // USE_C10D_NCCL diff --git a/torch/lib/c10d/NCCLUtils.hpp b/torch/lib/c10d/NCCLUtils.hpp index 846293086b5ee..0dec4573112a1 100644 --- a/torch/lib/c10d/NCCLUtils.hpp +++ b/torch/lib/c10d/NCCLUtils.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_NCCL + #include #include @@ -200,3 +202,5 @@ class NCCLComm { }; } // namespace c10d + +#endif // USE_C10D_NCCL diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp index 1b413c84bb84d..d423271192db8 100644 --- a/torch/lib/c10d/ProcessGroupGloo.cpp +++ b/torch/lib/c10d/ProcessGroupGloo.cpp @@ -1,6 +1,7 @@ -#include #include +#ifdef USE_C10D_GLOO + #include #include #include @@ -2816,3 +2817,5 @@ uint64_t ProcessGroupGloo::getSequenceNumberForGroup() { } } // namespace c10d + +#endif // USE_C10D_GLOO diff --git a/torch/lib/c10d/ProcessGroupGloo.hpp b/torch/lib/c10d/ProcessGroupGloo.hpp index eeaf86869c1dd..32e3799e94f3d 100644 --- a/torch/lib/c10d/ProcessGroupGloo.hpp +++ b/torch/lib/c10d/ProcessGroupGloo.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_GLOO + #include #include #include @@ -357,3 +359,5 @@ class ProcessGroupGloo : public ProcessGroup { }; } // namespace c10d + +#endif // USE_C10D_GLOO diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp index 1c56e686822fe..0c471216dffa7 100644 --- a/torch/lib/c10d/ProcessGroupMPI.cpp +++ b/torch/lib/c10d/ProcessGroupMPI.cpp @@ -1,5 +1,7 @@ #include +#ifdef USE_C10D_MPI + #include #include @@ -920,3 +922,5 @@ c10::intrusive_ptr ProcessGroupMPI::_allgather_base( } } // namespace c10d + +#endif // USE_C10D_MPI diff --git a/torch/lib/c10d/ProcessGroupMPI.hpp b/torch/lib/c10d/ProcessGroupMPI.hpp index b63448e9970ba..95363313b20fd 100644 --- a/torch/lib/c10d/ProcessGroupMPI.hpp +++ b/torch/lib/c10d/ProcessGroupMPI.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_MPI + #include #include #include @@ -263,3 +265,5 @@ class ProcessGroupMPI : public ProcessGroup { }; } // namespace c10d + +#endif // USE_C10D_MPI diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp index 7e0320d028aac..3f62cab44602b 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.cpp +++ b/torch/lib/c10d/ProcessGroupNCCL.cpp @@ -1,7 +1,7 @@ -#include -#include #include +#ifdef USE_C10D_NCCL + #include #include #include @@ -12,7 +12,9 @@ #include #include +#include #include +#include #include #include @@ -1925,3 +1927,5 @@ c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( } } // namespace c10d + +#endif // USE_C10D_NCCL diff --git a/torch/lib/c10d/ProcessGroupNCCL.hpp b/torch/lib/c10d/ProcessGroupNCCL.hpp index 5a467ff9df3ec..bafb76d7c6dd0 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.hpp +++ b/torch/lib/c10d/ProcessGroupNCCL.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_NCCL + #include #include #include @@ -560,3 +562,5 @@ class ProcessGroupNCCL : public ProcessGroup { }; } // namespace c10d + +#endif // USE_C10D_NCCL diff --git a/torch/lib/c10d/ProcessGroupWrapper.cpp b/torch/lib/c10d/ProcessGroupWrapper.cpp index 7d0da4d1c9922..c31e6f6263777 100644 --- a/torch/lib/c10d/ProcessGroupWrapper.cpp +++ b/torch/lib/c10d/ProcessGroupWrapper.cpp @@ -1,3 +1,7 @@ +#include + +#ifdef USE_C10D_GLOO + #include #include #include @@ -6,7 +10,6 @@ #include #include #include -#include #include namespace c10d { @@ -321,3 +324,5 @@ void ProcessGroupWrapper::runCollectiveChecks( } } // namespace c10d + +#endif // USE_C10D_GLOO diff --git a/torch/lib/c10d/ProcessGroupWrapper.hpp b/torch/lib/c10d/ProcessGroupWrapper.hpp index ea80ea04a82f0..9ee435593681d 100644 --- a/torch/lib/c10d/ProcessGroupWrapper.hpp +++ b/torch/lib/c10d/ProcessGroupWrapper.hpp @@ -1,5 +1,7 @@ #pragma once +#ifdef USE_C10D_GLOO + #include #include #include @@ -124,3 +126,5 @@ class ProcessGroupWrapper : public ProcessGroup { const std::vector& tensors) const; }; } // namespace c10d + +#endif // USE_C10D_GLOO From e6110d4d5ddbac8ddbfa613161be9315e60d621d Mon Sep 17 00:00:00 2001 From: albanD Date: Fri, 11 Jun 2021 07:27:53 -0700 Subject: [PATCH 044/305] Fix input_buffer check if inplace update is valid (#59817) Summary: Fixes an issue introduced in https://github.com/pytorch/pytorch/issues/17182 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59817 Reviewed By: bdhirsh Differential Revision: D29040738 Pulled By: albanD fbshipit-source-id: 67fd4e9fa0dadf507ddd954d20e119d8781c4de0 --- test/test_autograd.py | 19 +++++++++++++++++++ torch/csrc/autograd/input_buffer.cpp | 7 ++++--- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/test/test_autograd.py b/test/test_autograd.py index 10abd8bf596ae..31e8b7d6faab2 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -5421,6 +5421,25 @@ def get_refs(with_backward): gc.collect() self.assertTrue(ref.expired()) + def test_input_buffer_accum(self): + leaf = torch.rand(2, 2, requires_grad=True) + + # An op that returns sparse gradients + ind = torch.tensor([[0, 0]], dtype=torch.long) + out2 = leaf.gather(0, ind, sparse_grad=True) + + # An op that returns the gradients as-is + out1 = leaf.clone() + + grad_out1_original = torch.rand_like(out1) + grad_out1 = grad_out1_original.clone() + grad_out2 = torch.rand_like(out2) + + torch.autograd.backward((out1, out2), (grad_out1, grad_out2)) + + # Given gradients should not be modified inplace + self.assertEqual(grad_out1, grad_out1_original) + def index_perm_variable(shape, max_indices): if not isinstance(shape, tuple): diff --git a/torch/csrc/autograd/input_buffer.cpp b/torch/csrc/autograd/input_buffer.cpp index 3f8fab6c69f2f..f010b7261409c 100644 --- a/torch/csrc/autograd/input_buffer.cpp +++ b/torch/csrc/autograd/input_buffer.cpp @@ -19,14 +19,15 @@ namespace torch { namespace autograd { // ATen doesn't route sparse additions correctly... // do dense + sparse in-place if possible if (old_var.is_sparse()) { - //storage use_count is a big hammer, but for anything lighter there's an adversarial example with unexpected inplace modification - if (!var.is_sparse() && var.is_contiguous() && var.storage().use_count() == 1) { + // It is safe to change the Tensor inplace if the Tensor is only used in this buffer (this could be the gradient passed by the + // user) and that no other Tensor is using the same storage. + if (!var.is_sparse() && var.is_contiguous() && var.use_count() == 1 && var.storage().use_count() == 1) { buffer[pos] = var.add_(old_var); } else { buffer[pos] = var + old_var; } } else { - if (var.is_sparse() && !old_var.is_sparse() && old_var.is_contiguous() && old_var.storage().use_count() == 1) { + if (var.is_sparse() && !old_var.is_sparse() && old_var.is_contiguous() && old_var.use_count() == 1 && old_var.storage().use_count() == 1) { buffer[pos] = old_var.add_(var); } else { buffer[pos] = old_var + var; From d75e99b709d138dbbf294eeb6f7ee927d93533d5 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Fri, 11 Jun 2021 08:52:21 -0700 Subject: [PATCH 045/305] fx quant: enable qconfig_dict to target function invocations by order (#59605) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59605 Enables targeting of individual function invocations by execution order. For example, given a module such as ``` class M1(torch.nn.Module): def forward(self, x): x = torch.add(x, x) x = torch.add(x, x) return x class M2(torch.nn.Module): def __init__(self): self.m1 = M1() def forward(self, x): x = self.m1(x) return x ``` We can now target the first add of `m1` with ``` qconfig_dict = { "module_name_function_order": ("m1", torch.add, 0, custom_qconfig), } ``` Test Plan: ``` python test/test_quantization.py TestQuantizeFx.test_qconfig_module_name_function_order ``` Imported from OSS Reviewed By: hx89 Differential Revision: D28951077 fbshipit-source-id: 311d423724a31193d4fa4bbf3a712b46464b5a29 --- test/quantization/fx/test_quantize_fx.py | 132 +++++++++++++++++++++++ torch/quantization/fx/qconfig_utils.py | 72 +++++++++++-- torch/quantization/quantize_fx.py | 14 ++- 3 files changed, 210 insertions(+), 8 deletions(-) diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py index de82e4b19ec4e..173adb8bb48f8 100644 --- a/test/quantization/fx/test_quantize_fx.py +++ b/test/quantization/fx/test_quantize_fx.py @@ -1033,6 +1033,138 @@ def forward(self, x): self.assertEqual(m.module_conv1.qconfig, module_name_regex_qconfig) self.assertEqual(m.module_conv2.qconfig, module_name_qconfig) + def test_qconfig_module_name_object_type_order(self): + class M1(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1, 1) + self.fc2 = nn.Linear(1, 1) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + x = torch.add(x, x) + x = torch.add(x, x) + return x + + class M2(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1, 1) + self.fc2 = nn.Linear(1, 1) + self.m1 = M1() + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + x = torch.add(x, x) + x = torch.add(x, x) + x = self.m1(x) + return x + + class M3(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1, 1) + self.fc2 = nn.Linear(1, 1) + self.m2 = M2() + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + x = torch.add(x, x) + x = torch.add(x, x) + x = self.m2(x) + return x + + m = M3().eval() + qconfig_dict = { + "module_name_object_type_order": [ + # test various FQNs: global, single child, multiple children + ("", nn.Linear, 0, torch.quantization.default_qconfig), + ("", torch.add, 0, torch.quantization.default_qconfig), + ("m2", nn.Linear, 1, torch.quantization.default_qconfig), + ("m2", torch.add, 1, torch.quantization.default_qconfig), + ("m2.m1", nn.Linear, 0, torch.quantization.default_qconfig), + ("m2.m1", torch.add, 0, torch.quantization.default_qconfig), + ], + } + m = prepare_fx(m, qconfig_dict) + data = torch.randn(1, 1, 1, 1) + m(data) + m = convert_fx(m) + m(data) + + node_list = [ + # m3 + ns.call_function(torch.quantize_per_tensor), + ns.call_module(nnq.Linear), + ns.call_method("dequantize"), + ns.call_module(nn.Linear), + ns.call_function(torch.quantize_per_tensor), + ns.call_function(torch.ops.quantized.add), + ns.call_method("dequantize"), + ns.call_function(torch.add), + # m2 + ns.call_module(nn.Linear), + ns.call_function(torch.quantize_per_tensor), + ns.call_module(nnq.Linear), + ns.call_method("dequantize"), + ns.call_function(torch.add), + ns.call_function(torch.quantize_per_tensor), + ns.call_function(torch.ops.quantized.add), + # m1 + ns.call_module(nnq.Linear), + ns.call_method("dequantize"), + ns.call_module(nn.Linear), + ns.call_function(torch.quantize_per_tensor), + ns.call_function(torch.ops.quantized.add), + ns.call_method("dequantize"), + ns.call_function(torch.add), + ] + self.checkGraphModuleNodes(m, expected_node_list=node_list) + + # test that function order overrides global qconfig + class M4(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(1, 1) + self.fc2 = nn.Linear(1, 1) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + x = torch.add(x, x) + x = torch.add(x, x) + return x + + m = M4().eval() + qconfig_dict = { + "": torch.quantization.default_qconfig, + "module_name_object_type_order": [ + ("", nn.Linear, 1, None), + ("", torch.add, 1, None), + ], + } + m = prepare_fx(m, qconfig_dict) + data = torch.randn(1, 1, 1, 1) + m(data) + m = convert_fx(m) + m(data) + + node_list = [ + ns.call_function(torch.quantize_per_tensor), + ns.call_module(nnq.Linear), + ns.call_method("dequantize"), + ns.call_module(nn.Linear), + ns.call_function(torch.quantize_per_tensor), + ns.call_function(torch.ops.quantized.add), + ns.call_method("dequantize"), + ns.call_function(torch.add), + ] + self.checkGraphModuleNodes(m, expected_node_list=node_list) + + def test_qconfig_dict_validity(self): r""" Verifies that if a user passes an invalid key or makes a typo when diff --git a/torch/quantization/fx/qconfig_utils.py b/torch/quantization/fx/qconfig_utils.py index 1b528d7677d3c..4aee9ad4acfc5 100644 --- a/torch/quantization/fx/qconfig_utils.py +++ b/torch/quantization/fx/qconfig_utils.py @@ -1,5 +1,5 @@ import torch -from collections import OrderedDict +from collections import OrderedDict, defaultdict from typing import Union, Callable, Any, Dict, Tuple, Set import re @@ -86,6 +86,27 @@ def get_module_name_regex_qconfig(qconfig_dict, module_name, fallback_qconfig): return fallback_qconfig +def maybe_adjust_qconfig_for_module_name_object_type_order( + qconfig_dict: Any, + cur_module_path: str, + cur_object_type: Callable, + cur_object_type_idx: int, + fallback_qconfig: QConfigAny, +) -> QConfigAny: + qconfig_module_name_object_type_order = \ + qconfig_dict.get('module_name_object_type_order', {}) + for module_path, object_type, object_type_idx, qconfig in \ + qconfig_module_name_object_type_order: + if ( + (module_path == cur_module_path) and + (object_type == cur_object_type) and + (object_type_idx == cur_object_type_idx) + ): + return qconfig + + return fallback_qconfig + + def get_module_name_qconfig(qconfig_dict, module_name, fallback_qconfig): if module_name == '': # module name qconfig not found @@ -101,7 +122,7 @@ def get_module_name_qconfig(qconfig_dict, module_name, fallback_qconfig): # global_qconfig if necessary -def get_qconfig(qconfig_dict, module_type, module_name, global_qconfig): +def maybe_adjust_qconfig_for_module_type_or_name(qconfig_dict, module_type, module_name, global_qconfig): module_type_qconfig = get_object_type_qconfig( qconfig_dict, module_type, global_qconfig) module_name_regex_qconfig = get_module_name_regex_qconfig( @@ -119,11 +140,21 @@ def generate_qconfig_map( node_name_to_scope: Dict[str, Tuple[str, type]]) -> Dict[str, QConfigAny]: global_qconfig = qconfig_dict.get("", None) qconfig_map = dict() + + # example: + # + # {'foo.bar': {F.linear: 0, F.conv2d: 1, ...}, ...} + # + # meaning in submodule 'foo.bar', we have seen 0 F.linear and + # 1 F.conv2d invocations so far. + submodule_to_object_type_to_cur_idx: Dict[str, Dict[Callable, int]] = \ + defaultdict(lambda: defaultdict(int)) + for node in input_graph.nodes: qconfig = None if node.op == "get_attr": module_name, _ = _parent_name(node.target) - qconfig = get_qconfig( + qconfig = maybe_adjust_qconfig_for_module_type_or_name( qconfig_dict, type(modules[module_name]), module_name, global_qconfig) elif node.op == "call_function": # precedence: module_name_qconfig @@ -132,20 +163,45 @@ def generate_qconfig_map( function_qconfig = get_object_type_qconfig( qconfig_dict, node.target, global_qconfig) module_path, module_type = node_name_to_scope[node.name] - qconfig = get_qconfig( + qconfig = maybe_adjust_qconfig_for_module_type_or_name( qconfig_dict, module_type, module_path, function_qconfig) + + cur_object_type_idx = \ + submodule_to_object_type_to_cur_idx[module_path][node.target] + submodule_to_object_type_to_cur_idx[module_path][node.target] += 1 + qconfig = maybe_adjust_qconfig_for_module_name_object_type_order( + qconfig_dict, module_path, node.target, cur_object_type_idx, + qconfig) + elif node.op == "call_method": module_path, module_type = node_name_to_scope[node.name] # use the qconfig of the module that the node belongs to - qconfig = get_qconfig( + qconfig = maybe_adjust_qconfig_for_module_type_or_name( qconfig_dict, module_type, module_path, global_qconfig) + # Currently call_method does not support modifying qconfig + # by order, we can add this later if it is needed. + elif node.op == 'call_module': - qconfig = get_qconfig( + qconfig = maybe_adjust_qconfig_for_module_type_or_name( qconfig_dict, type(modules[node.target]), node.target, global_qconfig) + + module_path, module_type = node_name_to_scope[node.name] + # Note: for call_module, the module_path is the current module's name. + # to meaningfully count invocations, we need to count them in the parent + # module. + parent_name, _ = _parent_name(module_path) + cur_object_type_idx = \ + submodule_to_object_type_to_cur_idx[parent_name][module_type] + submodule_to_object_type_to_cur_idx[parent_name][module_type] += 1 + qconfig = maybe_adjust_qconfig_for_module_name_object_type_order( + qconfig_dict, parent_name, module_type, cur_object_type_idx, + qconfig) + # regex is not supported eager mode propagate_qconfig_, we'll # need to set the qconfig explicitly here in case regex # is used modules[node.target].qconfig = qconfig + qconfig_map[node.name] = qconfig return qconfig_map @@ -172,7 +228,9 @@ def check_is_valid_qconfig_dict(qconfig_dict: Any) -> None: `qconfig_dict`: dictionary whose keys we want to check """ - qconfig_dict_allowed_keys = {"", "object_type", "module_name_regex", "module_name"} + qconfig_dict_allowed_keys = { + "", "object_type", "module_name_regex", "module_name", + "module_name_object_type_order"} check_is_valid_config_dict(qconfig_dict, qconfig_dict_allowed_keys, "qconfig_dict") diff --git a/torch/quantization/quantize_fx.py b/torch/quantization/quantize_fx.py index 0d6bb075dd75f..692e7375ad70c 100644 --- a/torch/quantization/quantize_fx.py +++ b/torch/quantization/quantize_fx.py @@ -289,7 +289,19 @@ def prepare_fx( ("foo.*bar.*conv[0-9]+", qconfig?) ..., ], - # priority (in increasing order): global, object_type, module_name_regex, module_name + + # optional, used for matching object type invocations in a submodule by + # order + # TODO(future PR): potentially support multiple indices ('0,1') and/or + # ranges ('0:3'). + "module_name_object_type_order": [ + # fully_qualified_name, object_type, index, qconfig + ("foo.bar", torch.nn.functional.linear, 0, qconfig?), + ], + + # priority (in increasing order): + # global, object_type, module_name_regex, module_name, + # module_name_object_type_order # qconfig == None means fusion and quantization should be skipped for anything # matching the rule } From ba09355b12bc4ef0993b9d67e983c9bbbd07bbb5 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Fri, 11 Jun 2021 09:08:16 -0700 Subject: [PATCH 046/305] Upgrade Windows CI Python to 3.8 (#59729) Summary: Python 3.6 EOL is end of this year--we should use newer Python in CI. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59729 Reviewed By: bdhirsh Differential Revision: D29006807 Pulled By: janeyx99 fbshipit-source-id: c79214b02a72656058ba5d199141f8838212b3b6 --- .../cimodel/data/windows_build_definitions.py | 8 +- .circleci/config.yml | 104 +++++++++--------- .../build-parameters/pytorch-build-params.yml | 2 +- .../job-specs/pytorch-job-specs.yml | 4 +- .../workflows/workflows-scheduled-ci.yml | 22 ++-- .github/templates/windows_ci_workflow.yml.j2 | 2 +- .../workflows/pytorch-win-vs2019-cpu-py3.yml | 2 +- .../pytorch-win-vs2019-cuda10-cudnn7-py3.yml | 2 +- .../win-test-helpers/setup_pytorch_env.bat | 4 +- 9 files changed, 75 insertions(+), 75 deletions(-) diff --git a/.circleci/cimodel/data/windows_build_definitions.py b/.circleci/cimodel/data/windows_build_definitions.py index ab9aaec98dc7e..c6b680a1e0472 100644 --- a/.circleci/cimodel/data/windows_build_definitions.py +++ b/.circleci/cimodel/data/windows_build_definitions.py @@ -40,11 +40,13 @@ def gen_tree(self): target_arch = self.cuda_version.render_dots() if self.cuda_version else "cpu" + python_version = "3.8" + base_name_parts = [ "pytorch", "windows", self.vscode_spec.render(), - "py36", + "py" + python_version.replace(".", ""), target_arch, ] @@ -65,7 +67,7 @@ def gen_tree(self): ["pytorch", "win"] + self.vscode_spec.get_elements() + arch_env_elements - + ["py3"] + + ["py" + python_version.split(".")[0]] ) is_running_on_cuda = bool(self.cuda_version) and not self.force_on_cpu @@ -75,7 +77,7 @@ def gen_tree(self): else: props_dict = { "build_environment": build_environment_string, - "python_version": miniutils.quote("3.6"), + "python_version": miniutils.quote(python_version), "vc_version": miniutils.quote(self.vscode_spec.dotted_version()), "vc_year": miniutils.quote(str(self.vscode_spec.year)), "vc_product": self.vscode_spec.get_product(), diff --git a/.circleci/config.yml b/.circleci/config.yml index 72bcc82787474..41aa42b1228ef 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -352,7 +352,7 @@ pytorch_windows_params: &pytorch_windows_params default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" @@ -715,7 +715,7 @@ jobs: default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" @@ -782,7 +782,7 @@ jobs: default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" @@ -7631,8 +7631,8 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda10.1_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "" @@ -7647,10 +7647,10 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda10.1_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -7666,10 +7666,10 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda10.1_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools @@ -7684,10 +7684,10 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda10.1_on_cpu_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_on_cpu_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test1 use_cuda: "0" vc_product: BuildTools @@ -7696,8 +7696,8 @@ workflows: - pytorch_windows_build: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.1" - name: pytorch_windows_vs2019_py36_cuda11.1_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "" @@ -7712,10 +7712,10 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda11.1_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.1_build + - pytorch_windows_vs2019_py38_cuda11.1_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -7731,10 +7731,10 @@ workflows: - master - /ci-all\/.*/ - /release\/.*/ - name: pytorch_windows_vs2019_py36_cuda11.1_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.1_build + - pytorch_windows_vs2019_py38_cuda11.1_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools @@ -9272,8 +9272,8 @@ workflows: - pytorch_windows_build: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" - name: pytorch_windows_vs2019_py36_cuda10.1_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "" @@ -9282,10 +9282,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda10.1_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -9295,10 +9295,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda10.1_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools @@ -9307,10 +9307,10 @@ workflows: - pytorch_windows_test: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" - name: pytorch_windows_vs2019_py36_cuda10.1_on_cpu_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda10.1_on_cpu_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda10.1_build + - pytorch_windows_vs2019_py38_cuda10.1_build test_name: pytorch-windows-test1 use_cuda: "0" vc_product: BuildTools @@ -9319,8 +9319,8 @@ workflows: - pytorch_windows_build: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.1" - name: pytorch_windows_vs2019_py36_cuda11.1_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "" @@ -9329,10 +9329,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.1" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.1_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.1_build + - pytorch_windows_vs2019_py38_cuda11.1_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -9342,10 +9342,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.1" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.1_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.1_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.1_build + - pytorch_windows_vs2019_py38_cuda11.1_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools @@ -9414,7 +9414,7 @@ workflows: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" name: periodic_pytorch_windows_cuda11.3_build - python_version: "3.6" + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "14.28.29333" @@ -9424,7 +9424,7 @@ workflows: cuda_version: "11.3" executor: windows-with-nvidia-gpu name: periodic_pytorch_windows_cuda11.3_test1 - python_version: "3.6" + python_version: "3.8" requires: - periodic_pytorch_windows_cuda11.3_build test_name: pytorch-windows-test1 @@ -9437,7 +9437,7 @@ workflows: cuda_version: "11.3" executor: windows-with-nvidia-gpu name: periodic_pytorch_windows_cuda11.3_test2 - python_version: "3.6" + python_version: "3.8" requires: - periodic_pytorch_windows_cuda11.3_build test_name: pytorch-windows-test2 @@ -9495,8 +9495,8 @@ workflows: - pytorch_windows_build: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" - name: pytorch_windows_vs2019_py36_cuda11.3_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "14.28.29333" @@ -9510,10 +9510,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.3_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.3_build + - pytorch_windows_vs2019_py38_cuda11.3_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -9528,10 +9528,10 @@ workflows: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.3_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.3_build + - pytorch_windows_vs2019_py38_cuda11.3_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools diff --git a/.circleci/verbatim-sources/build-parameters/pytorch-build-params.yml b/.circleci/verbatim-sources/build-parameters/pytorch-build-params.yml index dc82098261d4f..22e8cc3e31861 100644 --- a/.circleci/verbatim-sources/build-parameters/pytorch-build-params.yml +++ b/.circleci/verbatim-sources/build-parameters/pytorch-build-params.yml @@ -84,7 +84,7 @@ pytorch_windows_params: &pytorch_windows_params default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" diff --git a/.circleci/verbatim-sources/job-specs/pytorch-job-specs.yml b/.circleci/verbatim-sources/job-specs/pytorch-job-specs.yml index 52448d6dcdf03..1d77d107e66b0 100644 --- a/.circleci/verbatim-sources/job-specs/pytorch-job-specs.yml +++ b/.circleci/verbatim-sources/job-specs/pytorch-job-specs.yml @@ -253,7 +253,7 @@ jobs: default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" @@ -320,7 +320,7 @@ jobs: default: "10.1" python_version: type: string - default: "3.6" + default: "3.8" vc_version: type: string default: "14.16" diff --git a/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml b/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml index 8a2115a796807..e36550129e162 100644 --- a/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml +++ b/.circleci/verbatim-sources/workflows/workflows-scheduled-ci.yml @@ -35,7 +35,7 @@ build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" name: periodic_pytorch_windows_cuda11.3_build - python_version: "3.6" + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "14.28.29333" @@ -45,7 +45,7 @@ cuda_version: "11.3" executor: windows-with-nvidia-gpu name: periodic_pytorch_windows_cuda11.3_test1 - python_version: "3.6" + python_version: "3.8" requires: - periodic_pytorch_windows_cuda11.3_build test_name: pytorch-windows-test1 @@ -58,7 +58,7 @@ cuda_version: "11.3" executor: windows-with-nvidia-gpu name: periodic_pytorch_windows_cuda11.3_test2 - python_version: "3.6" + python_version: "3.8" requires: - periodic_pytorch_windows_cuda11.3_build test_name: pytorch-windows-test2 @@ -116,8 +116,8 @@ - pytorch_windows_build: build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" - name: pytorch_windows_vs2019_py36_cuda11.3_build - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_build + python_version: "3.8" use_cuda: "1" vc_product: BuildTools vc_version: "14.28.29333" @@ -131,10 +131,10 @@ build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.3_test1 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_test1 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.3_build + - pytorch_windows_vs2019_py38_cuda11.3_build test_name: pytorch-windows-test1 use_cuda: "1" vc_product: BuildTools @@ -149,10 +149,10 @@ build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 cuda_version: "11.3" executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py36_cuda11.3_test2 - python_version: "3.6" + name: pytorch_windows_vs2019_py38_cuda11.3_test2 + python_version: "3.8" requires: - - pytorch_windows_vs2019_py36_cuda11.3_build + - pytorch_windows_vs2019_py38_cuda11.3_build test_name: pytorch-windows-test2 use_cuda: "1" vc_product: BuildTools diff --git a/.github/templates/windows_ci_workflow.yml.j2 b/.github/templates/windows_ci_workflow.yml.j2 index 2d792b1de5fde..65d6949f9cd84 100644 --- a/.github/templates/windows_ci_workflow.yml.j2 +++ b/.github/templates/windows_ci_workflow.yml.j2 @@ -19,7 +19,7 @@ env: IN_CI: 1 INSTALL_WINDOWS_SDK: 1 JOB_BASE_NAME: test - PYTHON_VERSION: "3.6" + PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" VC_VERSION: "" diff --git a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml index 8cfeea07e597c..6ff79a6795b94 100644 --- a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml @@ -18,7 +18,7 @@ env: IN_CI: 1 INSTALL_WINDOWS_SDK: 1 JOB_BASE_NAME: test - PYTHON_VERSION: "3.6" + PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" VC_VERSION: "" diff --git a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml index 9d3b9ed3aff53..9cd66ebd4b723 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml @@ -17,7 +17,7 @@ env: IN_CI: 1 INSTALL_WINDOWS_SDK: 1 JOB_BASE_NAME: test - PYTHON_VERSION: "3.6" + PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" VC_VERSION: "" diff --git a/.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat b/.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat index 9302c1fa961e2..371883a9f07ec 100644 --- a/.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat +++ b/.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat @@ -20,9 +20,7 @@ if NOT "%BUILD_ENVIRONMENT%"=="" ( ) call %CONDA_PARENT_DIR%\Miniconda3\Scripts\activate.bat %CONDA_PARENT_DIR%\Miniconda3 if NOT "%BUILD_ENVIRONMENT%"=="" ( - :: We have to pin Python version to 3.6.7, until mkl supports Python 3.7 - :: Numba is pinned to 0.44.0 to avoid https://github.com/numba/numba/issues/4352 - call conda install -y -q python=3.6.7 numpy mkl cffi pyyaml boto3 protobuf numba==0.44.0 scipy==1.5.0 typing_extensions dataclasses libuv + call conda install -y -q python=3.8 numpy mkl cffi pyyaml boto3 protobuf numba scipy typing_extensions dataclasses libuv if %errorlevel% neq 0 ( exit /b %errorlevel% ) call conda install -y -q -c conda-forge cmake if %errorlevel% neq 0 ( exit /b %errorlevel% ) From 23c232554bfc138d0d082ea8ea8e8ecec51e30ae Mon Sep 17 00:00:00 2001 From: Heitor Schueroff Date: Fri, 11 Jun 2021 09:39:21 -0700 Subject: [PATCH 047/305] Implemented torch.cov (#58311) Summary: Based from https://github.com/pytorch/pytorch/pull/50466 Adds the initial implementation of `torch.cov` similar to `numpy.cov`. For simplicity, we removed support for many parameters in `numpy.cov` that are either redundant such as `bias`, or have simple workarounds such as `y` and `rowvar`. cc PandaBoi TODO - [x] Improve documentation Pull Request resolved: https://github.com/pytorch/pytorch/pull/58311 Reviewed By: mruberry Differential Revision: D28994140 Pulled By: heitorschueroff fbshipit-source-id: 1890166c0a9c01e0a536acd91571cd704d632f44 --- aten/src/ATen/core/aten_interned_strings.h | 1 + aten/src/ATen/native/Correlation.cpp | 109 ++++++++++++++++++ aten/src/ATen/native/native_functions.yaml | 3 + docs/source/tensors.rst | 1 + docs/source/torch.rst | 1 + test/test_torch.py | 47 ++++++++ tools/build_variables.bzl | 1 + torch/_tensor_docs.py | 6 + torch/_torch_docs.py | 69 +++++++++++ torch/overrides.py | 1 + .../_internal/common_methods_invocations.py | 23 ++++ 11 files changed, 262 insertions(+) create mode 100644 aten/src/ATen/native/Correlation.cpp diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 1a7486a019a06..3bb9f66d0c958 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -260,6 +260,7 @@ _(aten, cosine_embedding_loss) \ _(aten, cosine_similarity) \ _(aten, count_nonzero) \ _(aten, cross) \ +_(aten, cov) \ _(aten, std_mean) \ _(aten, var_mean) \ _(aten, ctc_loss) \ diff --git a/aten/src/ATen/native/Correlation.cpp b/aten/src/ATen/native/Correlation.cpp new file mode 100644 index 0000000000000..9e1f67b915dd2 --- /dev/null +++ b/aten/src/ATen/native/Correlation.cpp @@ -0,0 +1,109 @@ +#include +#include + +namespace at { +namespace native { + +Tensor cov( + const Tensor& self, + int64_t correction, + const c10::optional& fweights, + const c10::optional& aweights) { + constexpr int64_t OBSERVATIONS_DIM = 1; + + TORCH_CHECK( + self.ndimension() <= 2, + "cov(): expected input to have two or fewer dimensions but got an input with ", + self.ndimension(), + " dimensions"); + + TORCH_CHECK( + self.scalar_type() != kBool, "cov(): bool dtype is not supported for input"); + + // View input tensor as 2D (variables, observations) + auto in = self.ndimension() < 2 ? self.view({1, -1}) : self; + const auto num_observations = in.size(OBSERVATIONS_DIM); + + // The product of frequencies (fweights) and weights (aweights). + Tensor w; + + if (fweights.has_value()) { + w = fweights.value(); + TORCH_CHECK( + w.ndimension() <= 1, + "cov(): expected fweights to have one or fewer dimensions but got fweights with ", + w.ndimension(), + " dimensions"); + TORCH_CHECK( + at::isIntegralType(w.scalar_type(), false), + "cov(): expected fweights to have integral dtype but got fweights with ", + w.scalar_type(), + " dtype"); + TORCH_CHECK( + w.numel() == num_observations, + "cov(): expected fweights to have the same numel as there are observations in the input but got ", + w.numel(), + " != ", + num_observations); + TORCH_CHECK( + num_observations == 0 || w.min().ge(0).item(), + "cov(): fweights cannot be negative"); + } + + if (aweights.has_value()) { + const auto& aw = aweights.value(); + TORCH_CHECK( + aw.ndimension() <= 1, + "cov(): expected aweights to have one or fewer dimensions but got aweights with ", + aw.ndimension(), + " dimensions"); + TORCH_CHECK( + at::isFloatingType(aw.scalar_type()), + "cov(): expected aweights to have floating point dtype but got aweights with ", + aw.scalar_type(), + " dtype"); + TORCH_CHECK( + aw.numel() == num_observations, + "cov(): expected aweights to have the same numel as there are observations in the input but got ", + aw.numel(), + " != ", + num_observations); + TORCH_CHECK( + num_observations == 0 || aw.min().ge(0).item(), + "cov(): aweights cannot be negative"); + w = w.defined() ? w * aw : aw; + } + + // Compute a weighted average of the observations + const auto w_sum = w.defined() + ? w.sum() + : at::scalar_tensor(num_observations, in.options().dtype(kLong)); + + TORCH_CHECK( + !w.defined() || w_sum.ne(0).item(), + "cov(): weights sum to zero, can't be normalized"); + + const auto avg = (w.defined() ? in * w : in).sum(OBSERVATIONS_DIM) / w_sum; + + // Compute the normalization factor + Tensor norm_factor; + + if (w.defined() && aweights.has_value() && correction != 0) { + norm_factor = w_sum - correction * (w * aweights.value()).sum() / w_sum; + } else { + norm_factor = w_sum - correction; + } + + if (norm_factor.le(0).item()) { + TORCH_WARN("cov(): degrees of freedom is <= 0"); + norm_factor.zero_(); + } + + // Compute covariance matrix + in = in - avg.unsqueeze(1); + const auto c = at::mm(in, (w.defined() ? in * w : in).t().conj()); + return at::true_divide(c, norm_factor).squeeze(); +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 5e0dd9917dd9f..f2e4150b67429 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -1277,6 +1277,9 @@ dispatch: CompositeExplicitAutograd: count_nonzero +- func: cov(Tensor self, *, int correction=1, Tensor? fweights=None, Tensor? aweights=None) -> Tensor + variants: function, method + - func: cudnn_affine_grid_generator(Tensor theta, int N, int C, int H, int W) -> Tensor grid dispatch: CUDA: cudnn_affine_grid_generator_forward diff --git a/docs/source/tensors.rst b/docs/source/tensors.rst index dc210c78ce6d0..36686b98e54c0 100644 --- a/docs/source/tensors.rst +++ b/docs/source/tensors.rst @@ -286,6 +286,7 @@ Tensor class reference Tensor.cosh Tensor.cosh_ Tensor.count_nonzero + Tensor.cov Tensor.acosh Tensor.acosh_ Tensor.arccosh diff --git a/docs/source/torch.rst b/docs/source/torch.rst index 94b288920ca6b..ec9ff083514fa 100644 --- a/docs/source/torch.rst +++ b/docs/source/torch.rst @@ -477,6 +477,7 @@ Other Operations cdist clone combinations + cov cross cummax cummin diff --git a/test/test_torch.py b/test/test_torch.py index b30697e099443..30ad853b30c42 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -4324,6 +4324,53 @@ def test_exponential_no_zero(self, device, dtype): x = torch.empty(50000000, device=device, dtype=dtype).exponential_() self.assertTrue(x.min() > 0) + @dtypes(torch.float, torch.cfloat) + def test_cov(self, device, dtype): + def check(t, correction=1, fweights=None, aweights=None): + actual = torch.cov(t, correction=correction, fweights=fweights, aweights=aweights) + t = t.cpu().numpy() + fweights = fweights.cpu().numpy() if fweights is not None else None + aweights = aweights.cpu().numpy() if aweights is not None else None + expected = np.cov(t, ddof=correction, fweights=fweights, aweights=aweights) + expected = torch.from_numpy(np.array(expected)).to(dtype=actual.dtype) + self.assertEqual(actual, expected, atol=1e-05, rtol=1e-05) + + def generate_input_tensors(): + yield make_tensor((0, 0), device, dtype) + yield make_tensor((1, 0), device, dtype) + yield make_tensor((0, 1), device, dtype) + yield make_tensor((2), device, dtype) + yield make_tensor((2, 1), device, dtype) + yield make_tensor((2, 2), device, dtype) + yield make_tensor((2, 3), device, dtype) + yield make_tensor((5, 10), device, dtype) + yield make_tensor((5, 10), device, dtype, noncontiguous=True) + yield torch.tensor([0, -2, nan, 10.2, inf], dtype=dtype, device=device) + + for t in generate_input_tensors(): + check(t) + num_observations = t.numel() if t.ndim < 2 else t.size(1) + if num_observations > 0: + fweights = torch.randint(1, 10, (num_observations,), device=device) + aweights = make_tensor((num_observations,), device, torch.float, low=1) + for correction, fw, aw in product([0, 1, 2], [None, fweights], [None, aweights]): + check(t, correction, fweights, aweights) + + def test_cov_error(self, device): + def check(msg, *args, **kwargs): + with self.assertRaisesRegex(RuntimeError, r'cov\(\):.*' + msg + r'.*'): + torch.cov(*args, **kwargs) + + a = torch.rand(2) + check(r'expected input to have two or fewer dimensions', torch.rand(2, 2, 2)) + check(r'expected fweights to have one or fewer dimensions', a, fweights=torch.rand(2, 2)) + check(r'expected aweights to have one or fewer dimensions', a, aweights=torch.rand(2, 2)) + check(r'expected fweights to have integral dtype', a, fweights=torch.rand(2)) + check(r'expected aweights to have floating point dtype', a, aweights=torch.tensor([1, 1])) + check(r'expected fweights to have the same numel', a, fweights=torch.tensor([1])) + check(r'expected aweights to have the same numel', a, aweights=torch.rand(1)) + check(r'fweights cannot be negative', a, fweights=torch.tensor([-1, -2])) + check(r'aweights cannot be negative', a, aweights=torch.tensor([-1., -2.])) @skipIfNoSciPy @dtypes(*torch.testing.get_all_fp_dtypes()) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 0a9aa427ef801..8901e02fa95a4 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -932,6 +932,7 @@ aten_native_source_non_codegen_list = [ "aten/src/ATen/native/ConvolutionMM3d.cpp", "aten/src/ATen/native/ConvolutionTBC.cpp", "aten/src/ATen/native/Copy.cpp", + "aten/src/ATen/native/Correlation.cpp", "aten/src/ATen/native/Cross.cpp", "aten/src/ATen/native/DilatedMaxPool2d.cpp", "aten/src/ATen/native/DilatedMaxPool3d.cpp", diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py index a6b0c0ef6c454..4d3a8cafe9d3f 100644 --- a/torch/_tensor_docs.py +++ b/torch/_tensor_docs.py @@ -976,6 +976,12 @@ def add_docstr_all(method, docstr): See :func:`torch.count_nonzero` """) +add_docstr_all('cov', r""" +cov(*, correction=1, fweights=None, aweights=None) -> Tensor + +See :func:`torch.cov` +""") + add_docstr_all('cross', r""" cross(other, dim=-1) -> Tensor diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index 4e38a487f0e78..579041087a3ee 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -1638,6 +1638,75 @@ def merge_dicts(*dicts): False """) +add_docstr(torch.cov, r""" +cov(input, *, correction=1, fweights=None, aweights=None) -> Tensor + +Estimates the covariance matrix of the variables given by the :attr:`input` matrix, where rows are +the variables and columns are the observations. + +A covariance matrix is a square matrix giving the covariance of each pair of variables. The diagonal contains +the variance of each variable (covariance of a variable with itself). By definition, if :attr:`input` represents +a single variable (Scalar or 1D) then its variance is returned. + +The unbiased sample covariance of the variables :math:`x` and :math:`y` is given by: + +.. math:: + \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}(x_{i} - \bar{x})(y_{i} - \bar{y})}{N~-~1} + +where :math:`\bar{x}` and :math:`\bar{y}` are the simple means of the :math:`x` and :math:`y` respectively. + +If :attr:`fweights` and/or :attr:`aweights` are provided, the unbiased weighted covariance +is calculated, which is given by: + +.. math:: + \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}w_i(x_{i} - \mu_x^*)(y_{i} - \mu_y^*)}{\sum^{N}_{i = 1}w_i~-~1} + +where :math:`w` denotes :attr:`fweights` or :attr:`aweights` based on whichever is provided, or +:math:`w = fweights \times aweights` if both are provided, and +:math:`\mu_x^* = \frac{\sum^{N}_{i = 1}w_ix_{i} }{\sum^{N}_{i = 1}w_i}` is the weighted mean of the variable. + +Args: + input (Tensor): A 2D matrix containing multiple variables and observations, or a + Scalar or 1D vector representing a single variable. + +Keyword Args: + correction (int, optional): difference between the sample size and sample degrees of freedom. + Defaults to Bessel's correction, ``correction = 1`` which returns the unbiased estimate, + even if both :attr:`fweights` and :attr:`aweights` are specified. ``correction = 0`` + will return the simple average. Defaults to ``1``. + fweights (tensor, optional): A Scalar or 1D tensor of observation vector frequencies representing the number of + times each observation should be repeated. Its numel must equal the number of columns of :attr:`input`. + Must have integral dtype. Ignored if ``None``. `Defaults to ``None``. + aweights (tensor, optional): A Scalar or 1D array of observation vector weights. + These relative weights are typically large for observations considered “important” and smaller for + observations considered less “important”. Its numel must equal the number of columns of :attr:`input`. + Must have floating point dtype. Ignored if ``None``. `Defaults to ``None``. + +Returns: + (Tensor) The covariance matrix of the variables. + +Example:: + >>> x = torch.tensor([[0, 2], [1, 1], [2, 0]]).T + >>> x + tensor([[0, 1, 2], + [2, 1, 0]]) + >>> torch.cov(x) + tensor([[ 1., -1.], + [-1., 1.]]) + >>> torch.cov(x, correction=0) + tensor([[ 0.6667, -0.6667], + [-0.6667, 0.6667]]) + >>> fw = torch.randint(1, 10, (3,)) + >>> fw + tensor([1, 6, 9]) + >>> aw = torch.rand(3) + >>> aw + tensor([0.4282, 0.0255, 0.4144]) + >>> torch.cov(x, fweights=fw, aweights=aw) + tensor([[ 0.4169, -0.4169], + [-0.4169, 0.4169]]) +""") + add_docstr(torch.cat, r""" cat(tensors, dim=0, *, out=None) -> Tensor diff --git a/torch/overrides.py b/torch/overrides.py index 75bde5decb787..18c9cca37ceea 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -344,6 +344,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.clamp_min: lambda input, min, out=None: -1, torch.clamp_max: lambda input, max, out=None: -1, torch.column_stack: lambda tensors, out=None: -1, + torch.cov: lambda input, correction=1, fweights=None, aweights=None: -1, torch.clone: lambda input: -1, torch.combinations: lambda input, r=2, with_replacement=False: -1, torch.complex: lambda real, imag: -1, diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 1cda7f822db50..2881b5e7b63d7 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -2677,6 +2677,22 @@ def sample_inputs_std_var(op_info, device, dtype, requires_grad, **kwargs): ] +def sample_inputs_cov(op_info, device, dtype, requires_grad, **kwargs): + shapes = [(2,), (1, 2), (3, 2), (2, 3)] + + inputs = [] + for shape in shapes: + t = make_tensor(shape, device, dtype, requires_grad=requires_grad) + inputs.append(SampleInput(t)) + num_observations = t.numel() if t.ndimension() < 2 else t.size(1) + fweights = make_tensor((num_observations,), device, torch.int, low=0, high=10, requires_grad=requires_grad) + aweights = make_tensor((num_observations,), device, torch.float, low=0, high=1, requires_grad=requires_grad) + for correction, fw, aw in product(range(num_observations), [None, fweights], [None, aweights]): + inputs.append(SampleInput(t, kwargs={'correction': correction, 'fweights': fw, 'aweights': aw})) + + return inputs + + def _sample_inputs_svd(op_info, device, dtype, requires_grad=False, is_linalg_svd=False): """ This function generates input for torch.svd with distinct singular values so that autograd is always stable. @@ -4907,6 +4923,13 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), )), + OpInfo('cov', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.half, *[torch.bfloat16] if CUDA11OrLater else []), + sample_inputs_func=sample_inputs_cov, + supports_out=False, + # JIT test not working for tensor kwargs (https://github.com/pytorch/pytorch/issues/58507) + skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit'),)), OpInfo('cross', dtypes=all_types_and_complex(), dtypesIfCUDA=all_types_and(torch.half), From 2ce21b2e615668ab3718de99f8fa3b282a431dbb Mon Sep 17 00:00:00 2001 From: Kimish Patel Date: Fri, 11 Jun 2021 10:14:56 -0700 Subject: [PATCH 048/305] [Pytorch backend delegation] Preprocess to accept (#58873) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58873 BackenDebugInforRecorder Prior to this PR: In order to generate debug handles corresponding to the graph being lowered, backend's preprocess will call generate_debug_handles and will get map of Node*-to-debug_handles. In order to facilitate this, to_backend will own BackendDebugInfoRecorder and initialize thread local pointer to it. generate_debug_handle function will query thread local pointer to see if there is a valid BackendDebugInforRecorder for the context. If there is it will generate debug handles. After this PR: Signature of preprocess is changed such that backends have to register preprocess that accepts instance of BackendDebugInfoRecorder by reference. generate_debug_handles is no more a free function but becomes part of the API of BackendDebugInfoRecorder. Now backend's preprocess function will call generate_debug_handles on BackendDebugInfoRecorder instead of free function. Reason for this change: With RAII that initializes thread local pointer, results in a lose contract with backends, which may result in backends not storing debug information. Making it part of API results in backends having to be aware of BackendDebugInfoRecorder and explicitly chosing not to generate/store debug information if they chose to do so. Test Plan: backend tests Imported from OSS Reviewed By: jbschlosser, raziel Differential Revision: D28648613 fbshipit-source-id: c9b7e7bf0f78e87023ea7bc08612cf893b08cb98 --- .../jit/test_backend_compiler_preprocess.cpp | 4 +- test/cpp/jit/test_backend_lib.cpp | 4 +- test/custom_backend/custom_backend.h | 4 +- tools/build_variables.bzl | 1 - .../jit/backends/backend_debug_handler.cpp | 24 +------- .../csrc/jit/backends/backend_debug_handler.h | 29 +++------- torch/csrc/jit/backends/backend_detail.cpp | 55 ++++++++++++++++++- torch/csrc/jit/backends/backend_detail.h | 14 ++++- torch/csrc/jit/backends/backend_init.cpp | 6 -- .../jit/backends/generate_debug_handles.cpp | 39 ------------- .../jit/backends/generate_debug_handles.h | 39 ------------- 11 files changed, 82 insertions(+), 137 deletions(-) delete mode 100644 torch/csrc/jit/backends/generate_debug_handles.cpp delete mode 100644 torch/csrc/jit/backends/generate_debug_handles.h diff --git a/test/cpp/jit/test_backend_compiler_preprocess.cpp b/test/cpp/jit/test_backend_compiler_preprocess.cpp index 3de692bd48561..18217c0471bc8 100644 --- a/test/cpp/jit/test_backend_compiler_preprocess.cpp +++ b/test/cpp/jit/test_backend_compiler_preprocess.cpp @@ -1,6 +1,5 @@ #include #include -#include #include #include @@ -13,7 +12,8 @@ namespace { // can be passed when there's no usage of compilation in runtime backend lib. c10::IValue preprocess( const Module& mod, - const c10::Dict& method_compile_spec) { + const c10::Dict& method_compile_spec, + const BackendDebugHandleGenerator& generate_debug_handles) { // The output of this process would produce a dictionary // Key: method name. // Val: compiled blob (represented by a string). diff --git a/test/cpp/jit/test_backend_lib.cpp b/test/cpp/jit/test_backend_lib.cpp index f9563e741f750..9e6252dc9a232 100644 --- a/test/cpp/jit/test_backend_lib.cpp +++ b/test/cpp/jit/test_backend_lib.cpp @@ -1,4 +1,5 @@ #include +#include #include namespace torch { @@ -72,7 +73,8 @@ class TestBackend : public PyTorchBackendInterface { namespace { c10::IValue preprocess( const Module& mod, - const c10::Dict& method_compile_spec) { + const c10::Dict& method_compile_spec, + const BackendDebugHandleGenerator& generate_debug_handles) { return mod._ivalue(); } diff --git a/test/custom_backend/custom_backend.h b/test/custom_backend/custom_backend.h index 4829d2169d521..e3fc32eba634e 100644 --- a/test/custom_backend/custom_backend.h +++ b/test/custom_backend/custom_backend.h @@ -1,4 +1,5 @@ #include +#include #include namespace torch { @@ -68,7 +69,8 @@ class CustomBackend : public torch::jit::PyTorchBackendInterface { c10::IValue preprocess( const torch::jit::Module& mod, - const c10::Dict& method_compile_spec) { + const c10::Dict& method_compile_spec, + const torch::jit::BackendDebugHandleGenerator& generate_debug_handles) { return mod._ivalue(); } diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 8901e02fa95a4..ce043f2ebcf86 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -146,7 +146,6 @@ core_sources_full_mobile = [ "torch/csrc/jit/backends/backend_detail.cpp", "torch/csrc/jit/backends/backend_interface.cpp", "torch/csrc/jit/backends/backend_resolver.cpp", - "torch/csrc/jit/backends/generate_debug_handles.cpp", "torch/csrc/jit/codegen/fuser/codegen.cpp", "torch/csrc/jit/codegen/fuser/compiler.cpp", "torch/csrc/jit/codegen/fuser/executor.cpp", diff --git a/torch/csrc/jit/backends/backend_debug_handler.cpp b/torch/csrc/jit/backends/backend_debug_handler.cpp index d21e4efd5681d..f30702825d697 100644 --- a/torch/csrc/jit/backends/backend_debug_handler.cpp +++ b/torch/csrc/jit/backends/backend_debug_handler.cpp @@ -1,12 +1,10 @@ #include +#include + namespace torch { namespace jit { -namespace { -thread_local BackendDebugInfoRecorder* debug_info_recorder_ptr{nullptr}; -} // namespace - // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) std::atomic BackendDebugInfoRecorder::unique_debug_handle_{0}; @@ -36,23 +34,5 @@ BackendDebugInfoMapType BackendDebugInfoRecorder::stopRecording() { return handles_to_inlined_callstack_ptrs_; } -WithBackendDebugInfoRecorder::WithBackendDebugInfoRecorder( - BackendDebugInfoRecorder* recorder) throw() { - TORCH_CHECK( - debug_info_recorder_ptr == nullptr, - "Module debug recording already in progress."); - debug_info_recorder_ptr = recorder; -} - -WithBackendDebugInfoRecorder::~WithBackendDebugInfoRecorder() { - // If due to some exception within preprocess, such as compilation failure - // we throw, then we want to make sure the exit is clean - debug_info_recorder_ptr = nullptr; -} - -BackendDebugInfoRecorder* getBackendDebugInfoRecorder() { - return debug_info_recorder_ptr; -} - } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/backends/backend_debug_handler.h b/torch/csrc/jit/backends/backend_debug_handler.h index 60727bfcc242a..d25ce2f8cb041 100644 --- a/torch/csrc/jit/backends/backend_debug_handler.h +++ b/torch/csrc/jit/backends/backend_debug_handler.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -109,47 +110,31 @@ namespace jit { * [M.forward, source range] -> [aten::mul's source range] We need to track * mul's source range and inlined CS both. */ -using DebugHandleType = int64_t; using BackendDebugInfoMapType = - std::unordered_map; + std::unordered_map; /* * This class is used to generate debug info map. - * It instantiates debug_handle_manager and initialize thread local pointer to - * it. backend's preprocess will call generate_debug_handles, which uses - * debug_handle_manager to generate debug handles. When lowering process - * finishes, calling stopRecording will return debug info map from - * debug_handle_manager + * backend's preprocess will call generate_debug_handles (see + * backend_detail.cpp), which uses debug_handle_manager to generate debug + * handles. When lowering process finishes, calling stopRecording will + * return debug info map from debug_handle_manager */ class TORCH_API BackendDebugInfoRecorder { public: BackendDebugInfoRecorder() = default; - int64_t getNextDebugHandle(const Node* node); // Reason this is not done as RAII is that work done in stopRecording // can throw, and throwing with dtor will call terminate and thus voids any // exception catching at a higher level. BackendDebugInfoMapType stopRecording(); + NodeToDebugHandle generate_debug_handles(const std::shared_ptr& graph); private: static std::atomic unique_debug_handle_; BackendDebugInfoMapType handles_to_inlined_callstack_ptrs_; }; -// This is a RAII class that on ctor captures pointer to -// BackendDebugInfoRecorder and initializes thread_local pointer -// debug_info_recorder to it. Upon dtor it sets debug_info_recorder -// pointer back to null. Note that this context manager always requires -// that debug_info_recorder be nullptr when initializing the context. -// This is because nested scopes with debug_info_recorder are not yet allowed. -class WithBackendDebugInfoRecorder { - public: - WithBackendDebugInfoRecorder(BackendDebugInfoRecorder* recorder) throw(); - ~WithBackendDebugInfoRecorder(); -}; - -BackendDebugInfoRecorder* getBackendDebugInfoRecorder(); - } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/backends/backend_detail.cpp b/torch/csrc/jit/backends/backend_detail.cpp index 60cf0358ff8e6..5682d8e7163e5 100644 --- a/torch/csrc/jit/backends/backend_detail.cpp +++ b/torch/csrc/jit/backends/backend_detail.cpp @@ -7,12 +7,60 @@ #include #include +#include +#include #include namespace torch { namespace jit { namespace detail { namespace { + +/* + * This is the API via which backend's preprocess function will obtain debug + * handles corresponding to the nodes of the graph for the lowered methods of + * the module. + * Implementation: Given graph + * For each node of the graph, request debug handle via debug_info_recorder. + * debug_info_recorder returns the next debug handle and record node with + * corresponding debug info, such as source range and inlined callstack. + * + * Backend code for lowering module, preprocess, calls + * generate_debug_handles(graph)) which will return debug handles corresponding + * to the Node* of the said graph. + * + * In to_backend, after lowering, stopRecording is called on + * BackendModuleDebugInfoRecorder: It will extract debug map. This map gets + * stored as part of the lowered module. + * During serialization, specifically for bytecode serialization, check is made + * to see if the model being serialized has any lowered modules. If so + * corresponding debug map is extracted and serialized. + */ + +NodeToDebugHandle generate_debug_handles( + BackendDebugInfoRecorder& debug_info_recorder, + const std::shared_ptr& graph) { + NodeToDebugHandle node_to_debug_handles; + + std::stack blocks_to_visit; + // TODO: Look into using DepthFirstGraphNodeIterator + // At the moment it takes non-const graph but maybe we can make it + // general such that it can work with both. + blocks_to_visit.push(graph->block()); + while (!blocks_to_visit.empty()) { + Block* b = blocks_to_visit.top(); + blocks_to_visit.pop(); + for (Node* n : b->nodes()) { + DebugHandleType debug_handle = debug_info_recorder.getNextDebugHandle(n); + node_to_debug_handles.emplace(n, debug_handle); + for (Block* subblock : n->blocks()) { + blocks_to_visit.push(subblock); + } + } + } + return node_to_debug_handles; +} + std::unordered_map& backendPreprocessFunctions() { static std::unordered_map @@ -69,7 +117,6 @@ Module codegen_backend_module( // 2. Later call debug_info_recorder.stopRecording() to gather // recorded debug info and save it in __backend_debug_info. BackendDebugInfoRecorder debug_info_recorder; - WithBackendDebugInfoRecorder recorder_context(&debug_info_recorder); // Generate attributes. // This is the preprocessed module. @@ -77,11 +124,15 @@ Module codegen_backend_module( // the backend interface rather than as a separate function, we just pass // the cloned original Module. + BackendDebugHandleGenerator debug_handle_generator = + [&](const std::shared_ptr& g) { + return generate_debug_handles(debug_info_recorder, g); + }; loweredModule.register_attribute( "__processed_module", AnyType::get(), detail::getBackendPreprocessFunction(backend_name)( - cloned_module, method_compile_spec), + cloned_module, method_compile_spec, debug_handle_generator), /*is_param=*/false); // This is for the method_compile_spec passed in to to_ or diff --git a/torch/csrc/jit/backends/backend_detail.h b/torch/csrc/jit/backends/backend_detail.h index ffc9ce07b34ef..7299ce259bc8f 100644 --- a/torch/csrc/jit/backends/backend_detail.h +++ b/torch/csrc/jit/backends/backend_detail.h @@ -8,10 +8,20 @@ namespace torch { namespace jit { + +using DebugHandleType = int64_t; + +using NodeToDebugHandle = std::unordered_map; + +using BackendDebugHandleGenerator = + std::function&)>; + namespace detail { -using BackendPreprocessFunction = - std::function&)>; +using BackendPreprocessFunction = std::function&, + const BackendDebugHandleGenerator& generate_debug_handles)>; TORCH_API void registerBackendPreprocessFunction( const std::string& name, diff --git a/torch/csrc/jit/backends/backend_init.cpp b/torch/csrc/jit/backends/backend_init.cpp index deb0b950782d2..d6f4171f2261f 100644 --- a/torch/csrc/jit/backends/backend_init.cpp +++ b/torch/csrc/jit/backends/backend_init.cpp @@ -3,7 +3,6 @@ #include #include #include -#include #include #include #include @@ -189,11 +188,6 @@ void initJitBackendBindings(PyObject* module) { throw py::cast_error(c10::str( "Object ", py::str(orig_module), " is not a ScriptModule")); }); - - m.def( - "_jit_backend_generate_debug_handles", [](std::shared_ptr& graph) { - return generate_debug_handles(graph); - }); } } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/backends/generate_debug_handles.cpp b/torch/csrc/jit/backends/generate_debug_handles.cpp deleted file mode 100644 index 70b8384f04f39..0000000000000 --- a/torch/csrc/jit/backends/generate_debug_handles.cpp +++ /dev/null @@ -1,39 +0,0 @@ -#include - -#include - -namespace torch { -namespace jit { - -NodeToDebugHandle generate_debug_handles(const std::shared_ptr& graph) { - NodeToDebugHandle node_to_debug_handles; - auto* debug_info_recorder_ptr = getBackendDebugInfoRecorder(); - - // Note now we make having a valid debug_handle_manager a must. - // This avoids silently failing when say some code change results in - // to_backend not creating appropriate debug_handle_manager to - // be used with backend's preprocess function. - TORCH_CHECK( - debug_info_recorder_ptr, "Valid debug info recorder must be available."); - std::stack blocks_to_visit; - // TODO: Look into using DepthFirstGraphNodeIterator - // At the moment it takes non-const graph but maybe we can make it - // general such that it can work with both. - blocks_to_visit.push(graph->block()); - while (!blocks_to_visit.empty()) { - Block* b = blocks_to_visit.top(); - blocks_to_visit.pop(); - for (Node* n : b->nodes()) { - DebugHandleType debug_handle = - debug_info_recorder_ptr->getNextDebugHandle(n); - node_to_debug_handles.emplace(n, debug_handle); - for (Block* subblock : n->blocks()) { - blocks_to_visit.push(subblock); - } - } - } - return node_to_debug_handles; -} - -} // namespace jit -} // namespace torch diff --git a/torch/csrc/jit/backends/generate_debug_handles.h b/torch/csrc/jit/backends/generate_debug_handles.h deleted file mode 100644 index a2205093c9f0e..0000000000000 --- a/torch/csrc/jit/backends/generate_debug_handles.h +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once -#include - -#include -#include -#include - -namespace torch { -namespace jit { - -using NodeToDebugHandle = std::unordered_map; - -/* - * This is the API via which backend's preprocess function will obtain debug - * handles corresponding to the nodes of the graph for the lowered methods of - * the module. It is expected that the graphs of the methods are inlined. If - * graph is not inlined, this method will throw exception. Implementation: Given - * moudle with inlined methods: - * 1. Query if a valid debug handle manager has been initialized - * 2. If so use debug handle manager to generate debug handles, else all handles - * are -1. -1 is not quite the great constant for invalid handle, so we will - * probably fix this later. This will be used to generate debug handles and - * debug info map: - * 1. Inside to_backend, use BackendModuleDebugInfoRecorder to initialize thread - * local debug handler context. for the lowered module ptr. - * 2. Backend code for lowering module, preprocess, calls - * generate_debug_handles(graph)) which will return debug handles - * corresponding to the Node* of the said graph. - * 3. In to_backend, after lowering, call stopRecording on - * BackendModuleDebugInfoRecorder: It will extract debug map. This map gets - * stored in static instance of ModuleDebugInfoMap. Now there is a global map in - * which module's callstack ptr maps are stored and can be queried during - * serialization. - */ -NodeToDebugHandle TORCH_API -generate_debug_handles(const std::shared_ptr& graph); - -} // namespace jit -} // namespace torch From 9cdbddb3f72444cd335510982eb503e8d1b60e90 Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Fri, 11 Jun 2021 10:52:52 -0700 Subject: [PATCH 049/305] Fix `Vectorize::trunc` on ARM platform (#59858) Summary: Use `vrndq_f32`, which corresponds to `VRINTZ` instruction, which rounds floating point value towards zero, which matches `std::trunc` behaviour. This makes trunc implementation correct even for values that fit into float32, but can not be converted to int32, for example `-1.0e+20`, see the following [gist](https://gist.github.com/malfet/c612c9f4b3b5681ca1b2a69930825871): ``` inp= 3.1 2.7 -2.9 -1e+20 old_trunc= 3 2 -2 -2.14748e+09 new_trunc= 3 2 -2 -1e+20 ``` Fixes `test_reference_numerics_hard_trunc_cpu_float32` on M1 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59858 Reviewed By: kimishpatel Differential Revision: D29052008 Pulled By: malfet fbshipit-source-id: 6b567f39151538be1aa3890e3b4e1e978e598657 --- aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h index 08638bba66aa4..9ffd413346c0d 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h @@ -469,8 +469,8 @@ template <> class Vectorized { return map(std::tanh); } Vectorized trunc() const { - float32x4_t r0 = vcvtq_f32_s32(vcvtq_s32_f32(values.val[0])); - float32x4_t r1 = vcvtq_f32_s32(vcvtq_s32_f32(values.val[1])); + float32x4_t r0 = vrndq_f32(values.val[0]); + float32x4_t r1 = vrndq_f32(values.val[1]); return Vectorized(r0, r1); } Vectorized lgamma() const { From d433a55c948b1f7d188ff1fd55f18fac7677f056 Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Fri, 11 Jun 2021 11:14:26 -0700 Subject: [PATCH 050/305] Replace throw std::runtime_error with torch_check in torch/csrc/distributed (#59683) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59683 Replaces usages of throw std::runtime_error("foo") with the better torch_check(false, "foo") which allows C++ stacktraces to show up when TORCH_SHOW_CPP_STACKTRACES=1. This will hopefully provide much better debugging information when debugging crashes/flaky tests. ghstack-source-id: 131167210 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D28981327 fbshipit-source-id: 677f569e28600263cab18759eb1b282e0391aa7b --- test/cpp/rpc/test_wire_serialization.cpp | 6 ++++-- torch/csrc/distributed/autograd/autograd.cpp | 2 +- torch/csrc/distributed/c10d/init.cpp | 10 +++++----- torch/csrc/distributed/c10d/python_comm_hook.cpp | 2 +- torch/csrc/distributed/rpc/process_group_agent.cpp | 2 +- torch/csrc/distributed/rpc/py_rref.cpp | 2 +- torch/csrc/distributed/rpc/rref_context.cpp | 2 +- torch/csrc/distributed/rpc/tensorpipe_agent.cpp | 2 +- torch/csrc/distributed/rpc/utils.cpp | 6 +++--- 9 files changed, 18 insertions(+), 16 deletions(-) diff --git a/test/cpp/rpc/test_wire_serialization.cpp b/test/cpp/rpc/test_wire_serialization.cpp index c9d64694a2d61..9c949dd94c03a 100644 --- a/test/cpp/rpc/test_wire_serialization.cpp +++ b/test/cpp/rpc/test_wire_serialization.cpp @@ -7,6 +7,8 @@ #include #include +using ::testing::IsSubstring; + TEST(WireSerialize, Base) { auto run = [](const std::string& payload, const std::vector& tensors) { @@ -71,8 +73,8 @@ TEST(WireSerialize, Errors) { try { f(); FAIL(); - } catch (const std::runtime_error& e) { - EXPECT_STREQ(e.what(), msg); + } catch (const std::exception& e) { + EXPECT_PRED_FORMAT2(IsSubstring, msg, e.what()); } catch (...) { FAIL(); } diff --git a/torch/csrc/distributed/autograd/autograd.cpp b/torch/csrc/distributed/autograd/autograd.cpp index c361b40af86c3..115c8c4fd6047 100644 --- a/torch/csrc/distributed/autograd/autograd.cpp +++ b/torch/csrc/distributed/autograd/autograd.cpp @@ -18,7 +18,7 @@ void backward( DistEngine::getInstance().execute(context_id, roots, retain_graph); } catch (std::exception& e) { // FIXME: crashes if exception type is not RuntimeError - throw std::runtime_error(e.what()); + TORCH_CHECK(false, e.what()); } } diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp index e6dab63b0a147..1995d46ad7793 100644 --- a/torch/csrc/distributed/c10d/init.cpp +++ b/torch/csrc/distributed/c10d/init.cpp @@ -1592,19 +1592,19 @@ Example:: add("key3", 3); add("key3", 2); if (get("key") != "6") { - throw std::runtime_error("assertion failed"); + TORCH_CHECK(false, "assertion failed"); } if (get("key0") != "value0") { - throw std::runtime_error("assertion failed"); + TORCH_CHECK(false, "assertion failed"); } if (get("key1") != "value1") { - throw std::runtime_error("assertion failed"); + TORCH_CHECK(false, "assertion failed"); } if (get("key2") != "value2") { - throw std::runtime_error("assertion failed"); + TORCH_CHECK(false, "assertion failed"); } if (get("key3") != "15") { - throw std::runtime_error("assertion failed"); + TORCH_CHECK(false, "assertion failed"); } }, py::call_guard()); diff --git a/torch/csrc/distributed/c10d/python_comm_hook.cpp b/torch/csrc/distributed/c10d/python_comm_hook.cpp index 594fc99bbba96..9526501f1a052 100644 --- a/torch/csrc/distributed/c10d/python_comm_hook.cpp +++ b/torch/csrc/distributed/c10d/python_comm_hook.cpp @@ -35,7 +35,7 @@ c10::intrusive_ptr PythonCommHook::runHook( type.attr("__module__").cast(), ".", type.attr("__qualname__").cast()); - throw std::runtime_error(errMsg); + TORCH_CHECK(false, errMsg); } } diff --git a/torch/csrc/distributed/rpc/process_group_agent.cpp b/torch/csrc/distributed/rpc/process_group_agent.cpp index 86c665a327056..9985d7f016e8a 100644 --- a/torch/csrc/distributed/rpc/process_group_agent.cpp +++ b/torch/csrc/distributed/rpc/process_group_agent.cpp @@ -290,7 +290,7 @@ c10::intrusive_ptr ProcessGroupAgent::send( "tried to send() a message of type ", message->type(), " but RPC is no longer running on this node."); - throw std::runtime_error(err); + TORCH_CHECK(false, err); } TORCH_CHECK( to.id_ < (worker_id_t)pg_->getSize(), diff --git a/torch/csrc/distributed/rpc/py_rref.cpp b/torch/csrc/distributed/rpc/py_rref.cpp index 720a6304827cf..0fbbe9a75d265 100644 --- a/torch/csrc/distributed/rpc/py_rref.cpp +++ b/torch/csrc/distributed/rpc/py_rref.cpp @@ -318,7 +318,7 @@ void PyRRef::backwardOwnerRRef( try { value = torch::jit::toIValue(obj, c10::TensorType::get()); } catch (py::cast_error& e) { - throw std::runtime_error("RRef should contain a tensor for .backward()"); + TORCH_CHECK(false, "RRef should contain a tensor for .backward()"); } } diff --git a/torch/csrc/distributed/rpc/rref_context.cpp b/torch/csrc/distributed/rpc/rref_context.cpp index 659147e816e83..7e68d90965a00 100644 --- a/torch/csrc/distributed/rpc/rref_context.cpp +++ b/torch/csrc/distributed/rpc/rref_context.cpp @@ -110,7 +110,7 @@ void RRefContext::handleException(const JitFuture& jitFuture) { if (jitFuture.hasError()) { auto errMsg = jitFuture.tryRetrieveErrorMessage(); VLOG(1) << "Got exception: " << errMsg; - throw std::runtime_error(errMsg); + TORCH_CHECK(false, errMsg); } } diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp index 59dfa8f2bb3f7..ea88629f0b55d 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp @@ -875,7 +875,7 @@ c10::intrusive_ptr TensorPipeAgent::send( "tried to send() a message of type ", requestMessage->type(), " but RPC is no longer running on this node."); - throw std::runtime_error(err); + TORCH_CHECK(false, err); } const auto& url = findWorkerURL(toWorkerInfo); diff --git a/torch/csrc/distributed/rpc/utils.cpp b/torch/csrc/distributed/rpc/utils.cpp index 087a3459d6c91..615abbf300666 100644 --- a/torch/csrc/distributed/rpc/utils.cpp +++ b/torch/csrc/distributed/rpc/utils.cpp @@ -302,7 +302,7 @@ parseWireSections(const void* data, size_t data_size) { ++ptr; // past the '\n' } if (!ok) { - throw std::runtime_error("failed parse"); + TORCH_CHECK(false, "failed parse"); } std::unordered_map> out; @@ -311,7 +311,7 @@ parseWireSections(const void* data, size_t data_size) { ptr += headerEnt.second; } if (ptr != endp) { - throw std::runtime_error("failed bounds"); + TORCH_CHECK(false, "failed bounds"); } return out; } @@ -445,7 +445,7 @@ std::pair, std::vector> wireDeserialize( auto sectionReadFunc = [&](const std::string& ename) -> at::DataPtr { auto it = sections.find(ename); if (it == sections.end()) { - throw std::runtime_error("Couldn't find entity " + ename); + TORCH_CHECK(false, "Couldn't find entity " + ename); } const auto& idat = it->second; auto dptr = at::getCPUAllocator()->allocate(idat.second); From 6ea607500284d74d9f7f6f91b489b355a1fa7fad Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Fri, 11 Jun 2021 11:14:26 -0700 Subject: [PATCH 051/305] torch/lib/c10d: Use torch_check instead of throwing runtime_error (#59684) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59684 Same reasoning as in the below diff. ghstack-source-id: 131167212 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D28981326 fbshipit-source-id: 264a7f787ea8be76f743a2eaca67ae1d3bd8073a --- torch/lib/c10d/FileStore.cpp | 6 +- torch/lib/c10d/GlooDeviceFactory.cpp | 4 +- torch/lib/c10d/NCCLUtils.hpp | 5 +- torch/lib/c10d/ProcessGroup.cpp | 8 +-- torch/lib/c10d/ProcessGroup.hpp | 12 ++-- torch/lib/c10d/ProcessGroupGloo.cpp | 42 ++++++------- torch/lib/c10d/ProcessGroupMPI.cpp | 46 +++++++------- torch/lib/c10d/ProcessGroupNCCL.cpp | 66 ++++++++++---------- torch/lib/c10d/ProcessGroupRoundRobin.cpp | 10 +-- torch/lib/c10d/TCPStore.cpp | 10 +-- torch/lib/c10d/UnixSockUtils.hpp | 2 +- torch/lib/c10d/Utils.cpp | 8 +-- torch/lib/c10d/Utils.hpp | 16 ++--- torch/lib/c10d/WinSockUtils.hpp | 2 +- torch/lib/c10d/frontend.cpp | 2 +- torch/lib/c10d/test/ProcessGroupGlooTest.cpp | 2 +- torch/lib/c10d/test/ProcessGroupMPITest.cpp | 22 +++---- torch/lib/c10d/test/TCPStoreTest.cpp | 9 +-- 18 files changed, 137 insertions(+), 135 deletions(-) diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp index ea98963ee9df8..73342272c54c0 100644 --- a/torch/lib/c10d/FileStore.cpp +++ b/torch/lib/c10d/FileStore.cpp @@ -273,7 +273,7 @@ FileStore::FileStore(const std::string& path, int numWorkers) cleanupKey_("cleanup/"), regularPrefix_("/") { if (numWorkers_ < 1) { - throw std::runtime_error( + TORCH_CHECK(false, "Number of workers for FileStore should be greater than zero"); } } @@ -341,7 +341,7 @@ std::vector FileStore::get(const std::string& key) { const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout_ != kNoTimeout && elapsed > timeout_) { - throw std::runtime_error("Timeout waiting for key: " + key); + TORCH_CHECK(false, "Timeout waiting for key: " + key); } std::this_thread::sleep_for(std::chrono::milliseconds(10)); continue; @@ -424,7 +424,7 @@ void FileStore::wait( const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout != kNoTimeout && elapsed > timeout) { - throw std::runtime_error("Wait timeout"); + TORCH_CHECK(false, "Wait timeout"); } /* sleep override */ diff --git a/torch/lib/c10d/GlooDeviceFactory.cpp b/torch/lib/c10d/GlooDeviceFactory.cpp index 416676483e182..cb83a99838520 100644 --- a/torch/lib/c10d/GlooDeviceFactory.cpp +++ b/torch/lib/c10d/GlooDeviceFactory.cpp @@ -149,7 +149,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForInterface(const std::string& interfaceName) { auto device = makeGlooDevice(interfaceName, ""); if (!device) { - throw std::runtime_error("makeDeviceForInterface(): unsupported gloo device"); + TORCH_CHECK(false, "makeDeviceForInterface(): unsupported gloo device"); } return device; } @@ -158,7 +158,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForHostname(const std::string& hostname) { auto device = makeGlooDevice("", hostname); if (!device) { - throw std::runtime_error("makeDeviceForHostname(): unsupported gloo device"); + TORCH_CHECK(false, "makeDeviceForHostname(): unsupported gloo device"); } return device; } diff --git a/torch/lib/c10d/NCCLUtils.hpp b/torch/lib/c10d/NCCLUtils.hpp index 0dec4573112a1..e3ee14da0f542 100644 --- a/torch/lib/c10d/NCCLUtils.hpp +++ b/torch/lib/c10d/NCCLUtils.hpp @@ -9,6 +9,7 @@ #include #include +#include namespace { // Provides additional detail into NCCL error codes based on when these are @@ -57,7 +58,7 @@ const inline char* getNcclErrorDetailStr(ncclResult_t error) { std::string err = "NCCL error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \ "\n" + getNcclErrorDetailStr(result); \ - throw std::runtime_error(err); \ + TORCH_CHECK(false, err); \ } \ } while (0) @@ -142,7 +143,7 @@ class NCCLComm { ncclComm_t getNcclComm() { std::unique_lock lock(mutex_); if (aborted_) { - throw std::runtime_error( + TORCH_CHECK(false, "NCCL communicator was aborted on rank " + std::to_string(rank_) + "."); } diff --git a/torch/lib/c10d/ProcessGroup.cpp b/torch/lib/c10d/ProcessGroup.cpp index 39ae2bf71c598..4e03824eb12da 100644 --- a/torch/lib/c10d/ProcessGroup.cpp +++ b/torch/lib/c10d/ProcessGroup.cpp @@ -107,13 +107,13 @@ std::exception_ptr ProcessGroup::Work::exception() const { } int ProcessGroup::Work::sourceRank() const { - throw std::runtime_error( + TORCH_CHECK(false, "sourceRank() may only be called on work objects " "that correspond to a recv or recv-from-any call."); } std::vector ProcessGroup::Work::result() { - throw std::runtime_error("result() not implemented."); + TORCH_CHECK(false, "result() not implemented."); } void ProcessGroup::Work::synchronize() {} @@ -129,7 +129,7 @@ bool ProcessGroup::Work::wait(std::chrono::milliseconds timeout) { if (!completed_) { // Throw exception if the wait operation timed out and the work was not // completed. - throw std::runtime_error("Operation timed out!"); + TORCH_CHECK(false, "Operation timed out!"); } } if (exception_) { @@ -186,7 +186,7 @@ c10::intrusive_ptr ProcessGroup::allgather_coalesced( std::vector>& /* usused */, std::vector& /* usused */, const AllgatherOptions& /* usused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for allgather_coalesced in this process group"); } diff --git a/torch/lib/c10d/ProcessGroup.hpp b/torch/lib/c10d/ProcessGroup.hpp index ee2990fd33975..3a3ffa6b95d67 100644 --- a/torch/lib/c10d/ProcessGroup.hpp +++ b/torch/lib/c10d/ProcessGroup.hpp @@ -264,7 +264,7 @@ class ProcessGroup : public torch::CustomClassHolder { at::Tensor&, at::Tensor&, const ReduceScatterOptions& opts = ReduceScatterOptions()) { - throw std::runtime_error("ProcessGroup does not support reduce_scatter_base"); + TORCH_CHECK(false, "ProcessGroup does not support reduce_scatter_base"); } @@ -274,20 +274,20 @@ class ProcessGroup : public torch::CustomClassHolder { std::vector& outputSplitSizes, std::vector& inputSplitSizes, const AllToAllOptions& opts = AllToAllOptions()) { - throw std::runtime_error("ProcessGroup does not support alltoall"); + TORCH_CHECK(false, "ProcessGroup does not support alltoall"); } virtual c10::intrusive_ptr alltoall( std::vector& outputTensors, std::vector& inputTensors, const AllToAllOptions& opts = AllToAllOptions()) { - throw std::runtime_error("ProcessGroup does not support alltoall"); + TORCH_CHECK(false, "ProcessGroup does not support alltoall"); } virtual void monitoredBarrier( const BarrierOptions& /* unused */, bool /* unused */ = false ) { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not support monitoredBarrier, only GLOO supports monitored barrier.") @@ -299,7 +299,7 @@ class ProcessGroup : public torch::CustomClassHolder { // for GLOO and NCCL backends currently. virtual void setSequenceNumberForGroup() { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") @@ -311,7 +311,7 @@ class ProcessGroup : public torch::CustomClassHolder { // may indicate that there is some sort of collective desynchronization. virtual uint64_t getSequenceNumberForGroup() { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp index d423271192db8..98164237feb9c 100644 --- a/torch/lib/c10d/ProcessGroupGloo.cpp +++ b/torch/lib/c10d/ProcessGroupGloo.cpp @@ -66,7 +66,7 @@ func(__VA_ARGS__); \ break; \ default: \ - throw std::runtime_error("Invalid scalar type"); \ + TORCH_CHECK(false, "Invalid scalar type"); \ } #define HOST_NAME_MAX 256 @@ -95,7 +95,7 @@ func(args); \ break; \ default: \ - throw std::runtime_error("Invalid scalar type"); \ + TORCH_CHECK(false, "Invalid scalar type"); \ } #endif @@ -178,22 +178,22 @@ ReduceFunc toFunction(const ReduceOp& r) { case ReduceOp::MAX: return ReduceFunc(&::gloo::max); case ReduceOp::BAND: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BAND with non-integral dtype"); break; case ReduceOp::BOR: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BOR with non-integral dtype"); break; case ReduceOp::BXOR: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with non-integral dtype"); break; case ReduceOp::UNUSED: break; } - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); } // Bitwise AND with SFINAE guard for integral types. @@ -258,7 +258,7 @@ ReduceFunc toFunction(const ReduceOp& r) { break; } - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); } template @@ -368,7 +368,7 @@ void initializeStreamsEvents( const auto device_id = tensorgroup[0].device().index(); for (const auto& tensor : tensorgroup) { if (tensor.device().index() != device_id) { - throw std::runtime_error( + TORCH_CHECK(false, "tensors in the nested tensor vectors need to " "be on the same device"); } @@ -683,7 +683,7 @@ ProcessGroupGloo::ProcessGroupGloo( collectiveCounter_(0) { auto& devices = options->devices; if (devices.empty()) { - throw std::runtime_error("No device(s) specified"); + TORCH_CHECK(false, "No device(s) specified"); } // Create and connect a context for every device. @@ -915,7 +915,7 @@ c10::intrusive_ptr ProcessGroupGloo::broadcast( work = c10::make_intrusive( std::move(context), inputs, opts.rootRank, opts.rootTensor, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); @@ -1426,7 +1426,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( invalidArgument("unsupported layout"); } } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); @@ -1487,7 +1487,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce_coalesced( invalidArgument("unsupported layout"); } } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1646,7 +1646,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( opts.reduceOp, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1838,7 +1838,7 @@ c10::intrusive_ptr ProcessGroupGloo::allgather( work = c10::make_intrusive( std::move(context), outputs, inputs, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1972,7 +1972,7 @@ c10::intrusive_ptr ProcessGroupGloo::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in Gloo process group"); } @@ -2166,7 +2166,7 @@ c10::intrusive_ptr ProcessGroupGloo::gather( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -2349,7 +2349,7 @@ c10::intrusive_ptr ProcessGroupGloo::scatter( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -2359,7 +2359,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce_scatter( std::vector& outputs, std::vector>& inputs, const ReduceScatterOptions& opts) { - throw std::runtime_error("ProcessGroupGloo does not support reduce_scatter"); + TORCH_CHECK(false, "ProcessGroupGloo does not support reduce_scatter"); } namespace { @@ -2531,14 +2531,14 @@ c10::intrusive_ptr ProcessGroupGloo::alltoall_base( at::Tensor& checkSingleTensor(std::vector& tensors) { if (tensors.size() != 1) { - throw std::runtime_error("ProcessGroupGloo::send takes a single tensor"); + TORCH_CHECK(false, "ProcessGroupGloo::send takes a single tensor"); } auto& tensor = tensors[0]; if (!tensor.is_contiguous()) { - throw std::runtime_error("input tensor has to be contiguous"); + TORCH_CHECK(false, "input tensor has to be contiguous"); } if (tensor.is_sparse()) { - throw std::runtime_error("input tensor has to be dense"); + TORCH_CHECK(false, "input tensor has to be dense"); } return tensor; } diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp index 0c471216dffa7..aa6d81bbe4a13 100644 --- a/torch/lib/c10d/ProcessGroupMPI.cpp +++ b/torch/lib/c10d/ProcessGroupMPI.cpp @@ -21,7 +21,7 @@ namespace c10d { std::string err = "MPI error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + \ ", with error code: " + std::to_string(mpiStatus); \ - throw std::runtime_error(err); \ + TORCH_CHECK(false, err); \ } \ } while (0) @@ -63,13 +63,13 @@ bool cudaAwareMpiCheck() { // Checking the input tensor's validity void checkSingleTensorHelper(const at::Tensor& tensor) { if (!tensor.is_contiguous()) { - throw std::runtime_error("input tensor has to be contiguous"); + TORCH_CHECK(false, "input tensor has to be contiguous"); } if (tensor.is_sparse()) { - throw std::runtime_error("input tensor has to be dense"); + TORCH_CHECK(false, "input tensor has to be dense"); } if (tensor.is_cuda() && !cudaAwareMpiCheck()) { - throw std::runtime_error( + TORCH_CHECK(false, "CUDA tensor detected and the MPI used doesn't " "have CUDA-aware MPI support"); } @@ -77,7 +77,7 @@ void checkSingleTensorHelper(const at::Tensor& tensor) { void checkSingleTensor(const std::vector& tensors) { if (tensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "MPI process group does not support multi-GPU collectives"); } checkSingleTensorHelper(tensors[0]); @@ -89,7 +89,7 @@ void checkSameSizeAndType( for (const auto& tensor : tensors) { if ((tensor.numel() != t_in.numel()) || (tensor.scalar_type() != t_in.scalar_type())) { - throw std::runtime_error("Tensors are not equal in size or data type"); + TORCH_CHECK(false, "Tensors are not equal in size or data type"); } checkSingleTensorHelper(tensor); } @@ -158,7 +158,7 @@ bool ProcessGroupMPI::AsyncWork::isCompleted() { bool ProcessGroupMPI::AsyncWork::isSuccess() const { if (request_ != MPI_REQUEST_NULL) { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid call to AsyncWork::isSuccess before work has completed"); } @@ -232,14 +232,14 @@ void ProcessGroupMPI::initMPIOnce() { MPI_CHECK(MPI_Init_thread( nullptr, nullptr, MPI_THREAD_SERIALIZED, &mpiThreadSupport_)); if (mpiThreadSupport_ < MPI_THREAD_SERIALIZED) { - throw std::runtime_error( + TORCH_CHECK(false, "Used MPI implementation doesn't have the " "minimum level of threading support: " "MPI_THREAD_SERIALIZED. This is required by " "c10d package"); } if (std::atexit(ProcessGroupMPI::mpiExit)) { - throw std::runtime_error("Fail to register the MPI exit handler"); + TORCH_CHECK(false, "Fail to register the MPI exit handler"); } }); } @@ -285,7 +285,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( MPI_CHECK(MPI_Comm_size(groupComm, &size)); if (rank < 0 || size < 0) { - throw std::runtime_error("Failed to get the world_size / rank"); + TORCH_CHECK(false, "Failed to get the world_size / rank"); } } } @@ -303,7 +303,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( ProcessGroupMPI::ProcessGroupMPI(int rank, int size, MPI_Comm pgComm) : ProcessGroup(rank, size), stop_(false), pgComm_(pgComm) { if (pgComm_ == MPI_COMM_NULL) { - throw std::runtime_error("pgComm_ must not be MPI_COMM_NULL"); + TORCH_CHECK(false, "pgComm_ must not be MPI_COMM_NULL"); } // Start the worker thread accepting MPI calls @@ -427,7 +427,7 @@ c10::intrusive_ptr ProcessGroupMPI::allreduce( c10::intrusive_ptr ProcessGroupMPI::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - throw std::runtime_error( + TORCH_CHECK(false, "allreduce_coalesced is currently not supported with MPI"); } @@ -467,12 +467,12 @@ c10::intrusive_ptr ProcessGroupMPI::allgather( const AllgatherOptions& opts) { checkSingleTensor(inputTensors); if (outputTensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "MPI process group only supports a single " "tensor op"); } if (static_cast(size_) != outputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "All gather: number of output tensors should equal " "to the world size"); } @@ -512,7 +512,7 @@ c10::intrusive_ptr ProcessGroupMPI::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupMPI does not support allgather_coalesced"); } @@ -524,16 +524,16 @@ c10::intrusive_ptr ProcessGroupMPI::gather( if (rank_ != opts.rootRank) { if (outputTensors.size() > 0) { - throw std::runtime_error( + TORCH_CHECK(false, "Gather: number of output tensors should be 0 " "for non-root"); } } else { if (outputTensors.size() != 1) { - throw std::runtime_error("Gather: multi-GPU collective is not supported"); + TORCH_CHECK(false, "Gather: multi-GPU collective is not supported"); } if (static_cast(size_) != outputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Gather: number of output tensors should equal " "to the world size"); } @@ -598,17 +598,17 @@ c10::intrusive_ptr ProcessGroupMPI::scatter( if (rank_ != opts.rootRank) { if (inputTensors.size() > 0) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: number of input tensors should be 0 " "for non-root"); } } else { if (inputTensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: multi-GPU collective is not supported"); } if (static_cast(size_) != inputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: number of input tensors should equal " "to the world size"); } @@ -670,7 +670,7 @@ c10::intrusive_ptr ProcessGroupMPI::reduce_scatter( std::vector& outputTensors, std::vector>& inputTensors, const ReduceScatterOptions& opts) { - throw std::runtime_error("ProcessGroupMPI does not support reduce_scatter"); + TORCH_CHECK(false, "ProcessGroupMPI does not support reduce_scatter"); } c10::intrusive_ptr ProcessGroupMPI::alltoall_base( @@ -917,7 +917,7 @@ c10::intrusive_ptr ProcessGroupMPI::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in MPI process group"); } diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp index 3f62cab44602b..f538e2f4ea560 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.cpp +++ b/torch/lib/c10d/ProcessGroupNCCL.cpp @@ -90,16 +90,16 @@ ncclRedOp_t getNcclReduceOp(const ReduceOp reduceOp, at::Tensor& input) { } catch (const std::out_of_range& e) { switch (reduceOp) { case ReduceOp::BAND: - throw std::runtime_error("Cannot use ReduceOp.BAND with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BAND with NCCL"); break; case ReduceOp::BOR: - throw std::runtime_error("Cannot use ReduceOp.BOR with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BOR with NCCL"); break; case ReduceOp::BXOR: - throw std::runtime_error("Cannot use ReduceOp.BXOR with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with NCCL"); break; default: - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); break; } } @@ -396,7 +396,7 @@ void ProcessGroupNCCL::WorkNCCL::synchronizeInternal( " ran for ", timeElapsed.count(), " milliseconds before timing out."); - throw std::runtime_error(exceptionMsg); + TORCH_CHECK(false, exceptionMsg); } // Check for errors and throw appropriate exception. checkAndThrowException(); @@ -819,7 +819,7 @@ std::vector>& ProcessGroupNCCL::getNCCLComm( bool isSendRecvSelf) { // Sanity check if (devicesKey.empty()) { - throw std::runtime_error( + TORCH_CHECK(false, "Not able to create/get the NCCL Communicator since " "the GPU devices are not known"); } @@ -945,10 +945,10 @@ namespace { // Check validity of tensor void check_gpu_single_tensor(const at::Tensor& tensor) { if (!tensor.is_cuda() || tensor.is_sparse()) { - throw std::runtime_error("Tensors must be CUDA and dense"); + TORCH_CHECK(false, "Tensors must be CUDA and dense"); } if (!tensor.is_contiguous()) { - throw std::runtime_error("Tensors must be contiguous"); + TORCH_CHECK(false, "Tensors must be contiguous"); } } @@ -956,10 +956,10 @@ void check_gpu_single_tensor(const at::Tensor& tensor) { // across distinct GPUs. void check_gpu_tensors(const std::vector& tensors) { if (tensors.size() == 0) { - throw std::runtime_error("Tensor list must be nonempty"); + TORCH_CHECK(false, "Tensor list must be nonempty"); } if (tensors.size() > static_cast(at::cuda::getNumGPUs())) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list mustn't be larger than the number of available GPUs"); } @@ -971,23 +971,23 @@ void check_gpu_tensors(const std::vector& tensors) { for (const auto& t : tensors) { if (!t.is_cuda() || t.is_sparse()) { - throw std::runtime_error("Tensors must be CUDA and dense"); + TORCH_CHECK(false, "Tensors must be CUDA and dense"); } if (t.scalar_type() != first.scalar_type()) { - throw std::runtime_error("Tensors must have identical type"); + TORCH_CHECK(false, "Tensors must have identical type"); } if (t.sizes() != first.sizes()) { - throw std::runtime_error("Tensors must have identical size"); + TORCH_CHECK(false, "Tensors must have identical size"); } if (t.strides() != first.strides()) { - throw std::runtime_error("Tensors must have identical strides"); + TORCH_CHECK(false, "Tensors must have identical strides"); } if (!t.is_non_overlapping_and_dense()) { - throw std::runtime_error("Tensors must be non-overlapping and dense"); + TORCH_CHECK(false, "Tensors must be non-overlapping and dense"); } const auto inserted = usedDevices.insert(t.get_device()).second; if (!inserted) { - throw std::runtime_error("Tensors must be on distinct GPU devices"); + TORCH_CHECK(false, "Tensors must be on distinct GPU devices"); } } } @@ -999,7 +999,7 @@ std::vector flatten_for_scatter_gather( std::vector& other, size_t world_size) { if (tensor_lists.size() != other.size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list operands to scatter/gather must have the same length"); } const auto num_devices = tensor_lists.size(); @@ -1009,7 +1009,7 @@ std::vector flatten_for_scatter_gather( for (auto i = size_t{}; i < num_devices; ++i) { if (tensor_lists[i].size() != world_size * num_devices) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list input to scatter/gather must match number of collective" " participants"); } @@ -1017,14 +1017,14 @@ std::vector flatten_for_scatter_gather( // Only check device match for the first tensor in the list; the call to // newLikeFlat() below will check the rest. if (tensor_lists[i].front().get_device() != other[i].get_device()) { - throw std::runtime_error( + TORCH_CHECK(false, "Corresponding input/output tensors to scatter/gather must all reside" " on the same device"); } for (const auto& t : tensor_lists[i]) { if (t.numel() != other[i].numel()) { - throw std::runtime_error( + TORCH_CHECK(false, "All tensor operands to scatter/gather must have the same number of elements"); } } @@ -1343,7 +1343,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allreduce( c10::intrusive_ptr ProcessGroupNCCL::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - throw std::runtime_error( + TORCH_CHECK(false, "allreduce_coalesced is currently not supported with NCCL"); } @@ -1481,7 +1481,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL does not support allgather_coalesced"); } @@ -1549,11 +1549,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_reduce_scatter_base( const ReduceScatterOptions& opts) { if (inputTensor.dtype() != outputTensor.dtype()) { - throw std::runtime_error("input tensor must be the same type as the outut tensor."); + TORCH_CHECK(false, "input tensor must be the same type as the outut tensor."); } if (inputTensor.numel() != outputTensor.numel() * size_) { - throw std::runtime_error("input tensor must be the same size as output size times world size"); + TORCH_CHECK(false, "input tensor must be the same size as output size times world size"); } // @lint-ignore CLANGTIDY @@ -1821,7 +1821,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall_base( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1829,7 +1829,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1837,7 +1837,7 @@ c10::intrusive_ptr ProcessGroupNCCL::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports send for NCCL lib version >= 2.7.0"); } @@ -1845,7 +1845,7 @@ c10::intrusive_ptr ProcessGroupNCCL::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports recv for NCCL lib version >= 2.7.0"); } #endif @@ -1868,20 +1868,20 @@ c10::intrusive_ptr ProcessGroupNCCL::gather( std::vector>& /* unused */, std::vector& /* unused */, const GatherOptions& /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support gather"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support gather"); } c10::intrusive_ptr ProcessGroupNCCL::scatter( std::vector& /* unused */, std::vector>& /* unused */, const ScatterOptions& /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support scatter"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support scatter"); } c10::intrusive_ptr ProcessGroupNCCL::recvAnysource( std::vector& /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support recvAnysource"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support recvAnysource"); } c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( @@ -1892,11 +1892,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( check_gpu_single_tensor(output_tensor); if (input_tensor.dtype() != output_tensor.dtype()) { - throw std::runtime_error("output tensor must have the same type as input tensor"); + TORCH_CHECK(false, "output tensor must have the same type as input tensor"); } if (input_tensor.numel() * size_ != output_tensor.numel()) { - throw std::runtime_error("output tensor size must be equal to world_size times input tensor size"); + TORCH_CHECK(false, "output tensor size must be equal to world_size times input tensor size"); } // just a wrapper to fit the collective interface diff --git a/torch/lib/c10d/ProcessGroupRoundRobin.cpp b/torch/lib/c10d/ProcessGroupRoundRobin.cpp index a55eea968b1e1..c439cf771a147 100644 --- a/torch/lib/c10d/ProcessGroupRoundRobin.cpp +++ b/torch/lib/c10d/ProcessGroupRoundRobin.cpp @@ -90,25 +90,25 @@ c10::intrusive_ptr ProcessGroupRoundRobin::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support send"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support send"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recvAnysource( std::vector& /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::barrier( const BarrierOptions& /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support barrier"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support barrier"); }; const c10::intrusive_ptr& ProcessGroupRoundRobin::next() { @@ -124,7 +124,7 @@ c10::intrusive_ptr ProcessGroupRoundRobin::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in RoundRobin process group"); } diff --git a/torch/lib/c10d/TCPStore.cpp b/torch/lib/c10d/TCPStore.cpp index 4958c47b79a71..6498f8bcbe633 100644 --- a/torch/lib/c10d/TCPStore.cpp +++ b/torch/lib/c10d/TCPStore.cpp @@ -133,7 +133,7 @@ void BackgroundThread::join() { void BackgroundThread::initStopSignal() { ghStopEvent_ = CreateEvent(NULL, TRUE, FALSE, NULL); if (ghStopEvent_ == NULL) { - throw std::runtime_error( + TORCH_CHECK(false, "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -149,7 +149,7 @@ void BackgroundThread::stop() { #else void BackgroundThread::initStopSignal() { if (pipe(controlPipeFd_.data()) == -1) { - throw std::runtime_error( + TORCH_CHECK(false, "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -336,7 +336,7 @@ void TCPStoreMasterDaemon::query(int socket) { watchHandler(socket); } else { - throw std::runtime_error("Unexpected query type"); + TORCH_CHECK(false, "Unexpected query type"); } } @@ -1126,7 +1126,7 @@ bool TCPStore::check(const std::vector& keys) { if (response == detail::CheckResponseType::NOT_READY) { return false; } - throw std::runtime_error("ready or not_ready response expected"); + TORCH_CHECK(false, "ready or not_ready response expected"); } void TCPStore::wait(const std::vector& keys) { @@ -1156,7 +1156,7 @@ void TCPStore::doWait( auto response = client_->receiveValue(); if (response != detail::WaitResponseType::STOP_WAITING) { - throw std::runtime_error("Stop_waiting response is expected"); + TORCH_CHECK(false, "Stop_waiting response is expected"); } } diff --git a/torch/lib/c10d/UnixSockUtils.hpp b/torch/lib/c10d/UnixSockUtils.hpp index fa74be27f889e..b75bddb763787 100644 --- a/torch/lib/c10d/UnixSockUtils.hpp +++ b/torch/lib/c10d/UnixSockUtils.hpp @@ -56,7 +56,7 @@ inline void waitSocketConnected( throw std::system_error(errno, std::system_category()); } else if (numReady == 0) { errno = 0; - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } socklen_t errLen = sizeof(errno); diff --git a/torch/lib/c10d/Utils.cpp b/torch/lib/c10d/Utils.cpp index f8adc58746c66..5d9aa744dbacd 100644 --- a/torch/lib/c10d/Utils.cpp +++ b/torch/lib/c10d/Utils.cpp @@ -118,7 +118,7 @@ PortType getSocketPort(int fd) { listenPort = ntohs(addr->sin6_port); } else { - throw std::runtime_error("unsupported protocol"); + TORCH_CHECK(false, "unsupported protocol"); } return listenPort; } @@ -140,7 +140,7 @@ std::string sockaddrToString(struct ::sockaddr* addr) { __output != nullptr) address[INET6_ADDRSTRLEN] = '\0'; } else { - throw std::runtime_error("unsupported protocol"); + TORCH_CHECK(false, "unsupported protocol"); } return address; } @@ -229,7 +229,7 @@ void handleConnectException( if (timeout != kNoTimeout) { const auto elapsed = std::chrono::high_resolution_clock::now() - start; if (elapsed > timeout) { - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::seconds(1)); @@ -346,7 +346,7 @@ std::tuple accept( while (true) { int res = tcputil::poll(events.get(), 1, timeout.count()); if (res == 0) { - throw std::runtime_error( + TORCH_CHECK(false, "waiting for processes to " "connect has timed out"); } else if (res == -1) { diff --git a/torch/lib/c10d/Utils.hpp b/torch/lib/c10d/Utils.hpp index 55edff85606cf..5beb5f1c6708b 100644 --- a/torch/lib/c10d/Utils.hpp +++ b/torch/lib/c10d/Utils.hpp @@ -92,7 +92,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { try { val = std::stoi(stringValue); } catch (std::exception& e) { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid value for environment variable: " + std::string(envVarName)); } if (val == 1) { @@ -100,7 +100,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { } else if (val == 0) { return false; } else { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid value for environment variable: " + std::string(envVarName)); } } @@ -340,16 +340,16 @@ inline at::Tensor newLikeFlat( std::vector>& tensors, size_t deviceIdx) { if (tensors.size() == 0 || tensors[0].size() == 0) { - throw std::runtime_error("Received an empty list"); + TORCH_CHECK(false, "Received an empty list"); } if (deviceIdx >= tensors.size()) { - throw std::runtime_error("Invalid device index"); + TORCH_CHECK(false, "Invalid device index"); } auto& t = tensors[deviceIdx][0]; auto device = t.device(); for (size_t i = 1; i < tensors[deviceIdx].size(); ++i) { if (tensors[deviceIdx][i].device() != device) { - throw std::runtime_error("Expecting all tensors on the same device"); + TORCH_CHECK(false, "Expecting all tensors on the same device"); } } at::DeviceGuard gpuGuard(device); @@ -363,7 +363,7 @@ inline at::Tensor newLikeFlat( inline at::Tensor newLikeFlat(std::vector& tensors) { if (tensors.size() == 0) { - throw std::runtime_error("Received an empty list"); + TORCH_CHECK(false, "Received an empty list"); } auto& t = tensors[0]; at::DeviceGuard gpuGuard(t.device()); @@ -504,7 +504,7 @@ using SizeType = uint64_t; continue; \ } else if ( \ errno_local == WSAETIMEDOUT || errno_local == WSAEWOULDBLOCK) { \ - throw std::runtime_error("Socket Timeout"); \ + TORCH_CHECK(false, "Socket Timeout"); \ } else { \ throw std::system_error(errno_local, std::system_category()); \ } \ @@ -521,7 +521,7 @@ using SizeType = uint64_t; if (errno == EINTR) { \ continue; \ } else if (errno == EAGAIN || errno == EWOULDBLOCK) { \ - throw std::runtime_error("Socket Timeout"); \ + TORCH_CHECK(false, "Socket Timeout"); \ } else { \ throw std::system_error(errno, std::system_category()); \ } \ diff --git a/torch/lib/c10d/WinSockUtils.hpp b/torch/lib/c10d/WinSockUtils.hpp index cd37695845ab1..793a0dc7640f2 100644 --- a/torch/lib/c10d/WinSockUtils.hpp +++ b/torch/lib/c10d/WinSockUtils.hpp @@ -46,7 +46,7 @@ inline void waitSocketConnected( std::chrono::high_resolution_clock::now() - startTime; if (elapsed > timeout) { errno = 0; - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::milliseconds(10)); diff --git a/torch/lib/c10d/frontend.cpp b/torch/lib/c10d/frontend.cpp index 86a78b6fcebb5..b65cba79884af 100644 --- a/torch/lib/c10d/frontend.cpp +++ b/torch/lib/c10d/frontend.cpp @@ -146,7 +146,7 @@ c10::intrusive_ptr DistributedC10d::newProcessGroupHelper( pg_name) { return pg_name.second == *group_name; }); if (it != pg_names_.end()) { - throw std::runtime_error( + TORCH_CHECK(false, "The specified group name has already been " "created, please use a different group name"); } diff --git a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp index f3a44cbcad4ae..a158d2c9685df 100644 --- a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp +++ b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp @@ -221,7 +221,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - throw std::runtime_error("future result should be tensor list or none"); + TORCH_CHECK(false, "future result should be tensor list or none"); } } return copyTensors(outputTensors); diff --git a/torch/lib/c10d/test/ProcessGroupMPITest.cpp b/torch/lib/c10d/test/ProcessGroupMPITest.cpp index bfefbbba2945e..b8538a016d5b7 100644 --- a/torch/lib/c10d/test/ProcessGroupMPITest.cpp +++ b/torch/lib/c10d/test/ProcessGroupMPITest.cpp @@ -48,7 +48,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - throw std::runtime_error("future result should be tensor list or none"); + TORCH_CHECK(false, "future result should be tensor list or none"); } } return outputTensors; @@ -80,7 +80,7 @@ void testAllreduce(int iter = 1000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -113,7 +113,7 @@ void testBroadcast(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -143,7 +143,7 @@ void testReduce(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -183,7 +183,7 @@ void testAllgather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -227,7 +227,7 @@ void testGather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -235,7 +235,7 @@ void testGather(int iter = 10000) { } else { for (const auto i : c10::irange(iter)) { if (outputTensors[i].size() != 0) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -277,7 +277,7 @@ void testScatter(int iter = 1) { auto data = outputTensors[i][0].data_ptr(); for (auto k = 0; k < outputTensors[i][0].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -333,13 +333,13 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { // Verify outputs for (const auto i : c10::irange(iter)) { if (recvAnysource && srcRanks[i] != 0) { - throw std::runtime_error("src rank is wrong for recvAnysource"); + TORCH_CHECK(false, "src rank is wrong for recvAnysource"); } const auto expected = i; auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -348,7 +348,7 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { void testBackendName() { auto pg = c10d::ProcessGroupMPI::createProcessGroupMPI(); if (pg->getBackendName() != std::string(c10d::MPI_BACKEND_NAME)) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } diff --git a/torch/lib/c10d/test/TCPStoreTest.cpp b/torch/lib/c10d/test/TCPStoreTest.cpp index e5b7eaf35cc5b..65fb425022b24 100644 --- a/torch/lib/c10d/test/TCPStoreTest.cpp +++ b/torch/lib/c10d/test/TCPStoreTest.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -73,7 +74,7 @@ void testHelper(const std::string& prefix = "") { EXPECT_EQ(numKeys, 4); auto timeout = std::chrono::milliseconds(kShortStoreTimeoutMillis); serverStore->setTimeout(timeout); - EXPECT_THROW(serverStore->get("key0"), std::runtime_error); + EXPECT_THROW(serverStore->get("key0"), c10::Error); }); // Hammer on TCPStore @@ -238,7 +239,7 @@ void testWatchKeyCallback(const std::string& prefix = "") { numCallbacksExecutedPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (numCallbacksExecutedFuture.wait_for(span) == std::future_status::timeout) - throw std::runtime_error("Callback execution timed out."); + TORCH_CHECK(false, "Callback execution timed out."); // Check number of callbacks executed equal to number of key change operations // Wait for all callbacks to be triggered @@ -302,7 +303,7 @@ void testKeyChangeHelper( std::future callbackFuture = callbackPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (callbackFuture.wait_for(span) == std::future_status::timeout) - throw std::runtime_error("Callback execution timed out."); + TORCH_CHECK(false, "Callback execution timed out."); // Any exceptions raised from asserts should be rethrown if (eptr) @@ -373,7 +374,7 @@ TEST(TCPStoreTest, testCleanShutdown) { clientTCPStore->get("key"); auto clientThread = std::thread([&clientTCPStore] { - EXPECT_THROW(clientTCPStore->get("invalid_key"), std::runtime_error); + EXPECT_THROW(clientTCPStore->get("invalid_key"), std::system_error); }); // start server shutdown during a client request From ee3025f734903916f261a75925a19b06cb4cdc39 Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Fri, 11 Jun 2021 11:14:39 -0700 Subject: [PATCH 052/305] Give clearer lint error messages (#59876) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59876 Test Plan: Imported from OSS Reviewed By: janeyx99 Differential Revision: D29067747 Pulled By: samestep fbshipit-source-id: cce7195467b5f9286d55a9d0c1655b4f92d4fbaf --- .github/workflows/lint.yml | 55 +++++++++++++++++++++++++++++++++----- 1 file changed, 48 insertions(+), 7 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 352c389260383..550f8abb98da8 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -99,7 +99,10 @@ jobs: uses: actions/checkout@v2 - name: Attempt to run setup.py run: | - python2 setup.py | grep "Python 2 has reached end-of-life and is no longer supported by PyTorch." + if ! python2 setup.py | grep -q "Python 2 has reached end-of-life and is no longer supported by PyTorch."; then + echo 'Running setup.py with Python 2 did not give the expected error message.' + false + fi shellcheck: runs-on: ubuntu-18.04 @@ -124,7 +127,19 @@ jobs: id: generate_workflows run: .github/scripts/generate_ci_workflows.py - name: Assert that regenerating the workflows didn't change them - run: .github/scripts/report_git_status.sh .github/workflows + run: | + if ! .github/scripts/report_git_status.sh .github/workflows; then + echo + echo 'As shown by the above diff, the committed .github/workflows' + echo 'are not up to date according to .github/templates.' + echo 'Please run this command, commit, and push again to your PR:' + echo + echo ' .github/scripts/generate_ci_workflows.py' + echo + echo 'If running that command does nothing, you may need to rebase' + echo 'onto a more recent commit from the PyTorch master branch.' + false + fi - name: Install ShellCheck id: install_shellcheck if: always() @@ -145,7 +160,16 @@ jobs: - name: Run ShellCheck if: always() && steps.install_shellcheck.outcome == 'success' run: | - tools/run_shellcheck.sh .extracted_scripts .jenkins/pytorch + if ! tools/run_shellcheck.sh .extracted_scripts .jenkins/pytorch; then + echo + echo 'ShellCheck gave a nonzero exit code. Please fix the warnings' + echo 'listed above. Note that if a path in one of the above warning' + echo 'messages starts with .extracted_scripts/ then that means it' + echo 'is referring to a shell script embedded within another file,' + echo 'whose path is given by the path components immediately' + echo 'following the .extracted_scripts/ prefix.' + false + fi - name: Check that jobs will be cancelled if: always() && steps.generate_workflows.outcome == 'success' run: | @@ -165,13 +189,27 @@ jobs: run: npm install -g markdown-toc - name: Regenerate ToCs and check that they didn't change run: | - set -eux + set -eu export PATH=~/.npm-global/bin:"$PATH" for FILE in $(git grep -Il '' -- '**.md'); do markdown-toc --bullets='-' -i "$FILE" done - .github/scripts/report_git_status.sh + if ! .github/scripts/report_git_status.sh .; then + echo + echo 'As shown by the above diff, the table of contents in one or' + echo 'more Markdown files is not up to date with the file contents.' + echo 'You can either apply that Git diff directly to correct the' + echo 'table of contents, or if you have npm installed, you can' + echo 'install the npm package markdown-toc and run the following' + # shellcheck disable=SC2016 + echo 'command (replacing $FILE with the filename for which you want' + echo 'to regenerate the table of contents):' + echo + # shellcheck disable=SC2016 + echo " markdown-toc --bullets='-' -i \"\$FILE\"" + false + fi flake8-py3: runs-on: ubuntu-18.04 @@ -219,10 +257,13 @@ jobs: path: flake8-output/ - name: Fail if there were any warnings run: | - set -eux + set -eu # Re-output flake8 status so GitHub logs show it on the step that actually failed cat "${GITHUB_WORKSPACE}"/flake8-output.txt - [ ! -s "${GITHUB_WORKSPACE}"/flake8-output.txt ] + if [ -s "${GITHUB_WORKSPACE}"/flake8-output.txt ]; then + echo 'Please fix the above Flake8 warnings.' + false + fi clang-tidy: if: github.event_name == 'pull_request' From f4fdc4995733c9346ea57d61a4944b17f075be4d Mon Sep 17 00:00:00 2001 From: Hui Guo Date: Fri, 11 Jun 2021 11:27:33 -0700 Subject: [PATCH 053/305] [NNC] Add python bindings for loopnest.compress_buffer (#59681) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59681 Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D28981573 Pulled By: huiguoo fbshipit-source-id: 003d66df576903c71bf46c95851fe6ccbba76f29 --- torch/csrc/jit/tensorexpr/expr.h | 3 +++ torch/csrc/jit/tensorexpr/tensorexpr_init.cpp | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/torch/csrc/jit/tensorexpr/expr.h b/torch/csrc/jit/tensorexpr/expr.h index 4b955e49708da..6e36be55fe713 100644 --- a/torch/csrc/jit/tensorexpr/expr.h +++ b/torch/csrc/jit/tensorexpr/expr.h @@ -251,6 +251,9 @@ class TORCH_API BufHandle : public ExprHandle { const Buf* node() const { return static_cast(ExprHandle::node()); } + Buf* node() { + return static_cast(ExprHandle::node()); + } template inline ExprHandle load(const Ts&... ts) const; diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp index b38a26eaae116..d1aad75b5a571 100644 --- a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp +++ b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp @@ -503,6 +503,12 @@ void initTensorExprBindings(PyObject* module) { "vectorize", [](const LoopNest& self, For* f) { self.vectorize(f); }, py::return_value_policy::reference) + .def_static( + "compress_buffer", + [](BufHandle& buf, Stmt* stmt) { + return LoopNest::compressBuffer(buf.node(), stmt); + }, + py::return_value_policy::reference) .def( "cache_accesses", [](LoopNest& self, From a9e136a61ee6ea8234da8085f5c59c740ba4d548 Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Fri, 11 Jun 2021 11:54:12 -0700 Subject: [PATCH 054/305] Remove ci/no-build (#59889) Summary: This reverts https://github.com/pytorch/pytorch/issues/58778, since triggering our primary CircleCI workflow only via pytorch-probot has been causing more problems than it's worth. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59889 Reviewed By: walterddr, seemethere Differential Revision: D29070418 Pulled By: samestep fbshipit-source-id: 0b47121b190c2e9efa27f38000ca362e634876dc --- .circleci/config.yml | 2 +- .circleci/verbatim-sources/header-section.yml | 2 +- .github/pytorch-circleci-labels.yml | 11 ----------- 3 files changed, 2 insertions(+), 13 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 41aa42b1228ef..a3df5be23b111 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -13,7 +13,7 @@ parameters: default: false run_build: type: boolean - default: false + default: true run_master_build: type: boolean default: false diff --git a/.circleci/verbatim-sources/header-section.yml b/.circleci/verbatim-sources/header-section.yml index 527340d542907..ff2081b7031c9 100644 --- a/.circleci/verbatim-sources/header-section.yml +++ b/.circleci/verbatim-sources/header-section.yml @@ -13,7 +13,7 @@ parameters: default: false run_build: type: boolean - default: false + default: true run_master_build: type: boolean default: false diff --git a/.github/pytorch-circleci-labels.yml b/.github/pytorch-circleci-labels.yml index 0c030a6154057..987d52b0c37b2 100644 --- a/.github/pytorch-circleci-labels.yml +++ b/.github/pytorch-circleci-labels.yml @@ -1,7 +1,5 @@ # For documentation concerning this configuration please refer to, # https://github.com/pytorch/pytorch-probot#trigger-circleci-workflows -default_params: - run_build: true labels_to_circle_params: ci/binaries: parameter: run_binary_tests @@ -11,15 +9,6 @@ labels_to_circle_params: - release/.* tags: - v[0-9]+(\.[0-9]+)*-rc[0-9]+ - ci/no-build: - default_true_on: - branches: - - nightly - - release/.* - tags: - - v[0-9]+(\.[0-9]+)*-rc[0-9]+ - set_to_false: - - run_build ci/master: parameter: run_master_build ci/slow-gradcheck: From 30e24b2d2ba4f9626825f11cf8c30f9993867b6d Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 11 Jun 2021 12:01:03 -0700 Subject: [PATCH 055/305] [nnc] Modified vectorize API to return bool (#59422) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59422 Test Plan: Imported from OSS Reviewed By: huiguoo Differential Revision: D28886980 Pulled By: navahgar fbshipit-source-id: 58cc3ecd86564a312a132f8260d836b096505095 --- test/cpp/tensorexpr/test_approx.cpp | 2 +- test/cpp/tensorexpr/test_llvm.cpp | 14 ++++++++------ test/cpp/tensorexpr/test_loopnest.cpp | 7 ++++--- test/cpp/tensorexpr/test_reductions.cpp | 10 ++++------ torch/csrc/jit/tensorexpr/loopnest.cpp | 23 +++++++++++++++-------- torch/csrc/jit/tensorexpr/loopnest.h | 2 +- 6 files changed, 33 insertions(+), 25 deletions(-) diff --git a/test/cpp/tensorexpr/test_approx.cpp b/test/cpp/tensorexpr/test_approx.cpp index 6bd31e2ef047e..2005f1e0fe7f7 100644 --- a/test/cpp/tensorexpr/test_approx.cpp +++ b/test/cpp/tensorexpr/test_approx.cpp @@ -15,7 +15,7 @@ static void vectorize(te::LoopNest* ln, te::Tensor* target, int width) { auto loops = ln->getLoopStmtsFor(target); te::For *inner, *tail; ln->splitWithTail(loops[0], width, &inner, &tail); - ln->vectorize(inner); + ASSERT_TRUE(te::LoopNest::vectorize(inner)); } std::string diffs(const at::Tensor& a, const at::Tensor& b) { diff --git a/test/cpp/tensorexpr/test_llvm.cpp b/test/cpp/tensorexpr/test_llvm.cpp index b1aebf2276815..52b5ccc6d63ee 100644 --- a/test/cpp/tensorexpr/test_llvm.cpp +++ b/test/cpp/tensorexpr/test_llvm.cpp @@ -599,7 +599,8 @@ TEST(LLVM, VectorizerLoadStoreTest) { Placeholder c_buf(BufHandle(c->buf())); LoopNest l({c}); Stmt* s = l.root_stmt(); - l.vectorize(dynamic_cast(dynamic_cast(s)->front())); + ASSERT_TRUE(LoopNest::vectorize( + dynamic_cast(dynamic_cast(s)->front()))); ASSERT_TRUE(dynamic_cast(dynamic_cast(s)->front()) == nullptr); @@ -623,7 +624,8 @@ TEST(LLVM, VectorizeBitCast) { Placeholder c_buf(BufHandle(c->buf())); LoopNest l({c}); Stmt* s = l.root_stmt(); - l.vectorize(dynamic_cast(dynamic_cast(s)->front())); + ASSERT_TRUE(LoopNest::vectorize( + dynamic_cast(dynamic_cast(s)->front()))); ASSERT_TRUE(dynamic_cast(dynamic_cast(s)->front()) == nullptr); LLVMCodeGen cg(s, {a, c_buf}); @@ -1550,9 +1552,9 @@ TEST(LLVM, RFactorVectorizedReduction) { auto distributed_loops = loopnest.distributeLoop(all_loops[1][1]); // Vectorize initializer of rfac_buf - loopnest.vectorize(distributed_loops[0]); + ASSERT_TRUE(LoopNest::vectorize(distributed_loops[0])); // Vectorize producer of rfac_buf - loopnest.vectorize(distributed_loops[1]); + ASSERT_TRUE(LoopNest::vectorize(distributed_loops[1])); loopnest.simplify(); loopnest.prepareForCodegen(); @@ -1755,8 +1757,8 @@ TEST(LLVM, VectorizedGEMM) { } { auto loops = NodeFinder::find(loop.root_stmt()); - loop.vectorize(loops[3]); - loop.vectorize(loops.back()); + ASSERT_TRUE(LoopNest::vectorize(loops[3])); + ASSERT_TRUE(LoopNest::vectorize(loops.back())); } loop.prepareForCodegen(); diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index c0860bc0d4787..c89e560e462b6 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -2662,7 +2662,8 @@ TEST(LoopNest, OuterLoopVectorization) { }); LoopNest l({tensor}); - l.vectorize(l.getAllLoopNestsWritingToBuf(tensor->buf())[0][0]); + ASSERT_TRUE( + LoopNest::vectorize(l.getAllLoopNestsWritingToBuf(tensor->buf())[0][0])); Stmt* root_stmt = l.root_stmt(); Block* outer_block = dynamic_cast(root_stmt); @@ -4727,9 +4728,9 @@ TEST(LoopNest, VectorizeUse) { "c", {{N, "n"}}, [&](const VarHandle& n) { return b->load(n) + 2.0f; }); LoopNest nest({c}, {b, c}); auto loops = nest.getAllLoopNestsWritingToBuf(b->buf())[0]; - nest.vectorize(loops[0]); + ASSERT_TRUE(LoopNest::vectorize(loops[0])); loops = nest.getAllLoopNestsWritingToBuf(c->buf())[0]; - nest.vectorize(loops[0]); + ASSERT_TRUE(LoopNest::vectorize(loops[0])); nest.prepareForCodegen(); // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores) Stmt* s = nest.root_stmt(); diff --git a/test/cpp/tensorexpr/test_reductions.cpp b/test/cpp/tensorexpr/test_reductions.cpp index 9bf56284b5b48..73e565157609b 100644 --- a/test/cpp/tensorexpr/test_reductions.cpp +++ b/test/cpp/tensorexpr/test_reductions.cpp @@ -1927,7 +1927,7 @@ TEST(Reductions, ReductionVectorize) { SimpleIREvaluator cg_before(l_before.root_stmt(), {in, tensor}); cg_before.call({in_, out_before}); - l.vectorize(l.getLoopStmtsFor(tensor)[0]); + ASSERT_TRUE(LoopNest::vectorize(l.getLoopStmtsFor(tensor)[0])); Stmt* s = l.root_stmt(); s = IRSimplifier::simplify(s); @@ -1962,8 +1962,7 @@ TEST(Reductions, ReductionVectorizeInner) { Tensor* tensor = Reduce("sum", {{8, "m"}}, Sum(), in, {{8, "n"}}); LoopNest l({tensor}); - ASSERT_THROWS_WITH( - l.vectorize(l.getLoopStmtsFor(tensor)[1]), "reduction axis"); + ASSERT_FALSE(LoopNest::vectorize(l.getLoopStmtsFor(tensor)[1])); } // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) @@ -1989,8 +1988,7 @@ TEST(Reductions, ReductionVectorizeRfactor) { SimpleIREvaluator cg_before(l_before.root_stmt(), {in, tensor}); cg_before.call({in_, out_before}); - ASSERT_THROWS_WITH( - l.vectorize(l.getLoopStmtsFor(tensor)[1]), "reduction axis"); + ASSERT_FALSE(LoopNest::vectorize(l.getLoopStmtsFor(tensor)[1])); // But if we rfactor this so it's not a reduce axis we can vectorize that // loop. @@ -2005,7 +2003,7 @@ TEST(Reductions, ReductionVectorizeRfactor) { l.distributeLoop(loops.at(0)); auto rfac_loops = l.getAllLoopNestsWritingToBuf(rfac_buf); - l.vectorize(rfac_loops[1][0]); + ASSERT_TRUE(LoopNest::vectorize(rfac_loops[1][0])); l.simplify(); Stmt* s = l.root_stmt(); diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index bfa11d374ddd3..f6ab0fcc5c068 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -422,10 +422,10 @@ class Vectorizer : public IRMutator { const Expr* start_ = nullptr; }; -void LoopNest::vectorize(For* f) { +bool LoopNest::vectorize(For* f) { Block* b = dynamic_cast(f->get_parent()); if (!b) { - return; + return false; } // Can't vectorize reduction axes. @@ -433,22 +433,29 @@ void LoopNest::vectorize(For* f) { for (auto* r : reductions) { if (std::find(r->reduce_args().begin(), r->reduce_args().end(), f->var()) != r->reduce_args().end()) { - throw std::logic_error("Cannot vectorize reduction axis - rfactor first"); + return false; } } Vectorizer v; - Stmt* old_f = Stmt::clone(f); Stmt* new_f = nullptr; try { - new_f = FlattenIndexes(f); + new_f = FlattenIndexes(Stmt::clone(f)); new_f = v.vectorize(dynamic_cast(new_f)); } catch (std::runtime_error& e) { - // Partial vectorization may have corrupted f - new_f = old_f; + // We clone f before vectorizing. So, any partial vectorization will + // have modified the clone. In case of an exception, we can continue + // using f. + new_f = f; } - b->replace_stmt(f, IRSimplifier::simplify(new_f)); + if (new_f != f) { + b->replace_stmt(f, IRSimplifier::simplify(new_f)); + return true; + } + + // Vectorization was not successful. + return false; } void LoopNest::initialize( diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h index bb5db759162e9..0799823c12f70 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.h +++ b/torch/csrc/jit/tensorexpr/loopnest.h @@ -86,7 +86,7 @@ class TORCH_API LoopNest { // getAllLoopNestsWritingToBuf(a) => {{i1,j1}, {i2,j2,k2}, {i2,j3}} std::vector> getAllLoopNestsWritingToBuf(const Buf*) const; - static void vectorize(For*); + static bool vectorize(For*); Stmt* simplify(); bool computeInline(Stmt* s); From b83ac0cc4eb59fbd70fbbdbb4960eb3c3eb26091 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 11 Jun 2021 12:01:03 -0700 Subject: [PATCH 056/305] [nnc] Added a check to vectorize only those loops that are normalized. (#59423) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59423 Test Plan: Imported from OSS Reviewed By: huiguoo Differential Revision: D28886979 Pulled By: navahgar fbshipit-source-id: edfc61feaf5efe22d4f367ac718b83b3d0f47cb3 --- test/cpp/tensorexpr/test_loopnest.cpp | 24 ++++++++++++++++++++++++ torch/csrc/jit/tensorexpr/loopnest.cpp | 8 +++++--- torch/csrc/jit/tensorexpr/loopnest.h | 14 +++++++++----- 3 files changed, 38 insertions(+), 8 deletions(-) diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index c89e560e462b6..ebb7f5cef4a65 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -2682,6 +2682,30 @@ TEST(LoopNest, OuterLoopVectorization) { ASSERT_EQ(dynamic_cast(for_body->front()), nullptr); } +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST(LoopNest, VectorizeLoopNotNormalized) { + KernelScope kernel_scope; + + // Input IR: + // for (int i = 0; i < 10; i++) { + // for (int j = 1; j < 5; j++) { + // A[i,j] = i * j; + // } + // } + BufHandle a_buf("A", {10, 5}, kInt); + VarHandle i("i", kInt); + VarHandle j("j", kInt); + auto for_body = Block::make({Store::make(a_buf, {i, j}, i * j)}); + auto inner_for = For::make(j, 1, 5, for_body); + auto outer_for = For::make(i, 0, 10, inner_for); + auto block = Block::make({outer_for}); + LoopNest l(block, {a_buf.node()}); + + ASSERT_TRUE(LoopNest::vectorize(inner_for)); + ASSERT_EQ(outer_for->body()->nstmts(), 1); + ASSERT_EQ(dynamic_cast(outer_for->body()->front()), nullptr); +} + namespace { std::string constantUpperBoundLoopIR(int upper_bound_val) { diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index f6ab0fcc5c068..2dc1bbaa4df5c 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -440,7 +440,9 @@ bool LoopNest::vectorize(For* f) { Vectorizer v; Stmt* new_f = nullptr; try { - new_f = FlattenIndexes(Stmt::clone(f)); + new_f = Stmt::clone(f); + normalize(dynamic_cast(new_f)); + new_f = FlattenIndexes(new_f); new_f = v.vectorize(dynamic_cast(new_f)); } catch (std::runtime_error& e) { // We clone f before vectorizing. So, any partial vectorization will @@ -1997,8 +1999,8 @@ bool LoopNest::normalize(For* f) { auto for_body_normalized = Substitute( f->body(), {{f->var(), (VarHandle(f->var()) + ExprHandle(f->start())).node()}}); - f->setBody(for_body_normalized); - f->setStop(new Sub(f->stop(), f->start())); + f->setBody(IRSimplifier::simplify(for_body_normalized)); + f->setStop(IRSimplifier::simplify(new Sub(f->stop(), f->start()))); f->setStart(new IntImm(0)); return true; } diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h index 0799823c12f70..717f9e8a39465 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.h +++ b/torch/csrc/jit/tensorexpr/loopnest.h @@ -86,7 +86,6 @@ class TORCH_API LoopNest { // getAllLoopNestsWritingToBuf(a) => {{i1,j1}, {i2,j2,k2}, {i2,j3}} std::vector> getAllLoopNestsWritingToBuf(const Buf*) const; - static bool vectorize(For*); Stmt* simplify(); bool computeInline(Stmt* s); @@ -403,6 +402,15 @@ class TORCH_API LoopNest { bool rfactor(Stmt* s, For* outer_reduction_for); bool rfactor(Stmt* s, For* outer_reduction_for, Buf** rfac_buf_ptr); + // Vectorize the given loop. This method requires that the given loop + // does not perform a reduction. + // It returns true if vectorization is successful and false otherwise. + static bool vectorize(For*); + + // Find the inner-most loops and vectorize them. Currently, this only works + // for the LLVM backend, when no reductions are involved. + void vectorizeInnerLoops(); + void setBufferMap( For* f, const std::unordered_map& map); @@ -410,10 +418,6 @@ class TORCH_API LoopNest { void eliminateDeadStores(); void prepareForCodegen(); - // Find the inner-most loops and vectorize them. Currently, this only works - // for the LLVM backend, when no reductions are involved. - void vectorizeInnerLoops(); - const std::unordered_set getInputBufs() const; const std::unordered_set getOutputBufs() const { return output_bufs_; From cf0c4ac25811cf93e51b4be6eb58bbdb95963b3b Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Fri, 11 Jun 2021 13:14:00 -0700 Subject: [PATCH 057/305] Fix some issues in CUDACachingAllocator (#59819) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59819 Test Plan: Sandcastle Reviewed By: ngimel Differential Revision: D29034650 fbshipit-source-id: 7e9689fc1ae121432e9421fa4a9ae00f7f78caca --- c10/cuda/CUDACachingAllocator.cpp | 41 ++++++++++++++----------------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp index 6681c2bdf4e3a..b877f6e8e8ca9 100644 --- a/c10/cuda/CUDACachingAllocator.cpp +++ b/c10/cuda/CUDACachingAllocator.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -130,7 +131,7 @@ void update_stat_array( StatArray& stat_array, int64_t amount, const StatTypes& stat_types) { - for (size_t stat_type = 0; stat_type < stat_types.size(); ++stat_type) { + for (const auto stat_type : c10::irange(stat_types.size())) { if (stat_types[stat_type]) { update_stat(stat_array[stat_type], amount); } @@ -614,9 +615,8 @@ class DeviceCachingAllocator { void resetAccumulatedStats() { std::lock_guard lock(mutex); - for (size_t statType = 0; - statType < static_cast(StatType::NUM_TYPES); - ++statType) { + for (const auto statType : + c10::irange(static_cast(StatType::NUM_TYPES))) { reset_accumulated_stat(stats.allocation[statType]); reset_accumulated_stat(stats.segment[statType]); reset_accumulated_stat(stats.active[statType]); @@ -635,9 +635,8 @@ class DeviceCachingAllocator { void resetPeakStats() { std::lock_guard lock(mutex); - for (size_t statType = 0; - statType < static_cast(StatType::NUM_TYPES); - ++statType) { + for (const auto statType : + c10::irange(static_cast(StatType::NUM_TYPES))) { reset_peak_stat(stats.allocation[statType]); reset_peak_stat(stats.segment[statType]); reset_peak_stat(stats.active[statType]); @@ -1090,11 +1089,11 @@ class DeviceCachingAllocator { stream_set streams(std::move(block->stream_uses)); AT_ASSERT(block->stream_uses.empty()); - for (auto it = streams.begin(); it != streams.end(); ++it) { - C10_CUDA_CHECK(cudaSetDevice(it->device_index())); + for (auto& stream : streams) { + C10_CUDA_CHECK(cudaSetDevice(stream.device_index())); cudaEvent_t event = create_event_internal(); - C10_CUDA_CHECK(cudaEventRecord(event, it->stream())); + C10_CUDA_CHECK(cudaEventRecord(event, stream.stream())); block->event_count++; cuda_events.emplace_back(event, block); @@ -1148,7 +1147,7 @@ class DeviceCachingAllocator { // Accumulates sizes of all memory blocks for given device in given pool void cache_info_aux(const BlockPool& pool, size_t* total, size_t* largest) { for (const auto& block : pool.blocks) { - size_t blocksize = block->size; + const auto blocksize = block->size; *total += blocksize; if (blocksize > *largest) { *largest = blocksize; @@ -1193,10 +1192,10 @@ class THCCachingAllocator { } void init(int device_count) { - int size = device_allocator.size(); + const auto size = device_allocator.size(); if (size < device_count) { device_allocator.resize(device_count); - for (int i = size; i < device_count; i++) { + for (const auto i : c10::irange(device_count)) { device_allocator[i] = std::unique_ptr( new DeviceCachingAllocator()); } @@ -1206,7 +1205,7 @@ class THCCachingAllocator { /** allocates a block which is safe to use from the provided stream */ void malloc(void** devPtr, int device, size_t size, cudaStream_t stream) { TORCH_INTERNAL_ASSERT( - 0 <= device && device < device_allocator.size(), + 0 <= device && static_cast(device) < device_allocator.size(), "Allocator not initialized for device ", device, ": did you call init?"); @@ -1228,7 +1227,7 @@ class THCCachingAllocator { void setMemoryFraction(double fraction, int device) { TORCH_INTERNAL_ASSERT( - 0 <= device && device < device_allocator.size(), + 0 <= device && static_cast(device) < device_allocator.size(), "Allocator not initialized for device ", device, ": did you call init?"); @@ -1246,9 +1245,8 @@ class THCCachingAllocator { } void emptyCache() { - int count = device_allocator.size(); - for (int i = 0; i < count; i++) - device_allocator[i]->emptyCache(); + for (auto& da : device_allocator) + da->emptyCache(); } void* getBaseAllocation(void* ptr, size_t* outSize) { @@ -1282,9 +1280,8 @@ class THCCachingAllocator { std::vector snapshot() { std::vector result; - int count = device_allocator.size(); - for (int i = 0; i < count; i++) { - auto snap = device_allocator[i]->snapshot(); + for (auto& da : device_allocator) { + auto snap = da->snapshot(); result.insert(result.end(), snap.begin(), snap.end()); } @@ -1377,7 +1374,7 @@ std::mutex* getFreeMutex() { } static inline void assertValidDevice(int device) { - int device_num = caching_allocator.device_allocator.size(); + const auto device_num = caching_allocator.device_allocator.size(); TORCH_CHECK(0 <= device && device < device_num, "Invalid device argument."); } From e41bc31eb2b66b15c586d2c1c26eaba9c35480b0 Mon Sep 17 00:00:00 2001 From: Rong Rong Date: Fri, 11 Jun 2021 13:56:06 -0700 Subject: [PATCH 058/305] make --run-specified-test-case use --include (#59704) Summary: instead of having specific logic to handle run-specific-test-case, we provide the flag to override include or bring-to-front with the SPECIFIED_TEST_CASES_FILE. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59704 Reviewed By: janeyx99 Differential Revision: D29038425 Pulled By: walterddr fbshipit-source-id: 803d3555813437c7f287a22f7704106b0c609919 --- test/run_test.py | 43 ++++++++++++++++++++++++++++++++----------- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/test/run_test.py b/test/run_test.py index ff3318dcd1b53..38563b2700504 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -397,7 +397,7 @@ # "test_nn": ["test_doubletensor_avg_pool3d", "test_share_memory", "test_hook_requires_grad"], # ... # } -# For test_nn.py, we would ONLY run test_doubletensor_avg_pool3d, test_share_memory, and test_hook_requires_grad. +# then for test_nn.py, we would ONLY run test_doubletensor_avg_pool3d, test_share_memory, and test_hook_requires_grad. SPECIFIED_TEST_CASES_DICT: Dict[str, List[str]] = {} # The file from which the SPECIFIED_TEST_CASES_DICT will be filled, a CSV of test cases that would be run when @@ -532,10 +532,10 @@ def get_slow_tests_based_on_S3() -> List[str]: def get_test_case_args(test_module, using_pytest) -> List[str]: - if test_module not in SPECIFIED_TEST_CASES_DICT: - sys.exit(f'Warning! Test module {test_module} is not found in the specified tests dict. This should never' - 'happen as we make a check for that before entering this function.') args = [] + # if test_module not specified or specified with '__all__' then run all tests + if test_module not in SPECIFIED_TEST_CASES_DICT or '__all__' in SPECIFIED_TEST_CASES_DICT[test_module]: + return args if using_pytest: args.append('-k') @@ -583,6 +583,7 @@ def run_test(test_module, test_directory, options, launcher_cmd=None, extra_unit executable = get_executable_command(options, allow_pytest=not extra_unittest_args, disable_coverage=disable_coverage) + # TODO: move this logic into common_utils.py instead of passing in "-k" individually # The following logic for running specified tests will only run for non-distributed tests, as those are dispatched # to test_distributed and not run_test (this function) if options.run_specified_test_cases: @@ -739,7 +740,8 @@ def __contains__(self, item): def parse_args(): parser = argparse.ArgumentParser( description='Run the PyTorch unit test suite', - epilog='where TESTS is any of: {}'.format(', '.join(TESTS))) + epilog='where TESTS is any of: {}'.format(', '.join(TESTS)), + formatter_class=argparse.RawTextHelpFormatter) parser.add_argument( '-v', '--verbose', @@ -838,8 +840,25 @@ def parse_args(): nargs='?', type=str, const=SPECIFIED_TEST_CASES_FILE, - help='runs specified test cases from previous OSS CI stats from a file, format CSV', - + help='load specified test cases file dumped from previous OSS CI stats, format CSV. ' + ' If all test cases should run for a please add a single row: \n' + ' test_filename,test_case_name\n' + ' ...\n' + ' ,__all__\n' + ' ...\n' + 'how we use the stats will be based on option "--use-specified-test-cases-by".' + ) + parser.add_argument( + '--use-specified-test-cases-by', + type=str, + choices=['include', 'bring-to-front'], + default='include', + help='used together with option "--run-specified-test-cases". When specified test case ' + 'file is set, this option allows the user to control whether to only run the specified test ' + 'modules or to simply bring the specified modules to front and also run the remaining ' + 'modules. Note: regardless of this option, we will only run the specified test cases ' + ' within a specified test module. For unspecified test modules with the bring-to-front ' + 'option, all test cases will be run, as one may expect.', ) return parser.parse_args() @@ -893,6 +912,12 @@ def exclude_tests(exclude_list, selected_tests, exclude_message=None): def get_selected_tests(options): + if options.run_specified_test_cases: + if options.use_specified_test_cases_for == 'include': + options.include = list(SPECIFIED_TEST_CASES_DICT.keys()) + elif options.use_specified_test_cases_for == 'bring-to-front': + options.bring_to_front = list(SPECIFIED_TEST_CASES_DICT.keys()) + selected_tests = options.include if options.bring_to_front: @@ -911,10 +936,6 @@ def get_selected_tests(options): if options.exclude_jit_executor: options.exclude.extend(JIT_EXECUTOR_TESTS) - if options.run_specified_test_cases: - # Filter out any unspecified test modules. - selected_tests = [t for t in selected_tests if t in SPECIFIED_TEST_CASES_DICT] - selected_tests = exclude_tests(options.exclude, selected_tests) if sys.platform == 'win32' and not options.ignore_win_blocklist: From 60eb22e45e73946e4b97f5a2d4370af57d2ac492 Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Fri, 11 Jun 2021 16:10:38 -0700 Subject: [PATCH 059/305] Build an -Wextra around c10 (#59853) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59853 Test Plan: Sandcastle Reviewed By: ngimel Differential Revision: D29016682 fbshipit-source-id: f6c5f32464d57dbd60b59b5f9e2234ef2c39f1c1 --- c10/cuda/CUDACachingAllocator.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp index b877f6e8e8ca9..07b0c3746ba70 100644 --- a/c10/cuda/CUDACachingAllocator.cpp +++ b/c10/cuda/CUDACachingAllocator.cpp @@ -1192,7 +1192,7 @@ class THCCachingAllocator { } void init(int device_count) { - const auto size = device_allocator.size(); + const auto size = static_cast(device_allocator.size()); if (size < device_count) { device_allocator.resize(device_count); for (const auto i : c10::irange(device_count)) { @@ -1375,7 +1375,9 @@ std::mutex* getFreeMutex() { static inline void assertValidDevice(int device) { const auto device_num = caching_allocator.device_allocator.size(); - TORCH_CHECK(0 <= device && device < device_num, "Invalid device argument."); + TORCH_CHECK( + 0 <= device && device < static_cast(device_num), + "Invalid device argument."); } DeviceStats getDeviceStats(int device) { From 864d129bae3c9f1fa41ed1ffce476c8958fd570e Mon Sep 17 00:00:00 2001 From: Supriya Rao Date: Fri, 11 Jun 2021 16:21:12 -0700 Subject: [PATCH 060/305] [quant][fx] Remove extra q-dq for weight bias in normalization ops (#59882) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59882 Currently for normalization ops, the weight and bias arguments are treated as activationn inputs which require observers. This results in adding extra quant-dequant ops for the weight and bias inputs. This PR adds support to skip observing weight/bias inputs of norm operators, thus removing the redundant q-dq ops Quantized graph with F.layer_norm Before this PR ``` def forward(self, x): _input_scale_0 = self._input_scale_0 _input_zero_point_0 = self._input_zero_point_0 quantize_per_tensor = torch.quantize_per_tensor(x, _input_scale_0, _input_zero_point_0, torch.quint8); x = _input_scale_0 = _input_zero_point_0 = None scale = self.scale _input_scale_1 = self._input_scale_1 _input_zero_point_1 = self._input_zero_point_1 quantize_per_tensor_1 = torch.quantize_per_tensor(scale, _input_scale_1, _input_zero_point_1, torch.quint8); scale = _input_scale_1 = _input_zero_point_1 = None bias = self.bias _input_scale_2 = self._input_scale_2 _input_zero_point_2 = self._input_zero_point_2 quantize_per_tensor_2 = torch.quantize_per_tensor(bias, _input_scale_2, _input_zero_point_2, torch.quint8); bias = _input_scale_2 = _input_zero_point_2 = None _scale_0 = self._scale_0 _zero_point_0 = self._zero_point_0 dequantize = quantize_per_tensor_1.dequantize(); quantize_per_tensor_1 = None dequantize_1 = quantize_per_tensor_2.dequantize(); quantize_per_tensor_2 = None layer_norm = torch.ops.quantized.layer_norm(quantize_per_tensor, [2, 5, 5], weight = dequantize, bias = dequantize_1, eps = 1e-05, output_scale = _scale_0, output_zero_point = _zero_point_0); quantize_per_tensor = dequantize = dequantize_1 = _scale_0 = _zero_point_0 = None dequantize_2 = layer_norm.dequantize(); layer_norm = None return dequantize_2 ``` After ``` def forward(self, x): _input_scale_0 = self._input_scale_0 _input_zero_point_0 = self._input_zero_point_0 quantize_per_tensor = torch.quantize_per_tensor(x, _input_scale_0, _input_zero_point_0, torch.quint8); x = _input_scale_0 = _input_zero_point_0 = None scale = self.scale bias = self.bias _scale_0 = self._scale_0 _zero_point_0 = self._zero_point_0 layer_norm = torch.ops.quantized.layer_norm(quantize_per_tensor, [2, 5, 5], weight = scale, bias = bias, eps = 1e-05, output_scale = _scale_0, output_zero_point = _zero_point_0); quantize_per_tensor = scale = bias = _scale_0 = _zero_point_0 = None dequantize = layer_norm.dequantize(); layer_norm = None return dequantize ``` Test Plan: python test/test_quantization.py TestQuantizeFxOps.test_norm_weight_bias Imported from OSS Reviewed By: HDCharles, ailzhang Differential Revision: D29068203 fbshipit-source-id: 24b5c38bbea5fd355d34522bfa654c9db18607da --- test/quantization/fx/test_quantize_fx.py | 37 ++++++++++++++++++++++++ torch/quantization/fx/prepare.py | 9 ++++-- torch/quantization/fx/utils.py | 12 ++++++++ 3 files changed, 56 insertions(+), 2 deletions(-) diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py index 173adb8bb48f8..b8c64ed8f8ba0 100644 --- a/test/quantization/fx/test_quantize_fx.py +++ b/test/quantization/fx/test_quantize_fx.py @@ -3404,6 +3404,43 @@ def test_instance_norm(self): quantized_module, torch.ops.quantized.instance_norm, skip_op_arg_for_functional=True) + def test_norm_weight_bias(self): + class Linear(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.ones(5, 5) + self.b = torch.zeros(5) + + def forward(self, x): + return torch.nn.functional.linear(x, self.w, self.b) + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.mods1 = Linear() + self.scale = torch.randn(5, 5) + self.bias = torch.randn(5, 5) + + def forward(self, x): + x1 = self.mods1(x) + y = F.layer_norm(x1, [5, 5], weight=self.scale, bias=self.bias) + return y + + model = M() + expected_occurrence = { + ns.call_function(torch.quantize_per_tensor): 1, + ns.call_function(torch.ops.quantized.linear): 1, + ns.call_function(torch.ops.quantized.layer_norm): 1, + ns.call_method("dequantize"): 1, + } + + self.checkGraphModeFxOp( + model, + (torch.rand(5, 5),), + QuantType.STATIC, + expected_node_occurrence=expected_occurrence + ) + def _test_default_node_quant_handler_ops( self, module, functional, qconfig, is_reference=True, node_list=None, additional_quant_pattern_dict=None ): diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index 8cc3c268f5928..2433012a4f7ca 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -52,7 +52,9 @@ assert_and_get_unique_device, node_bool_tensor_arg_indexes, get_new_attr_name_with_prefix, + NON_QUANTIZABLE_WEIGHT_OPS, WEIGHT_INDEX_DICT, + FUNCTIONAL_OPS_WITH_BIAS, ) from ..quantization_mappings import ( @@ -88,6 +90,9 @@ def node_arg_is_weight(node: Node, arg: Any) -> bool: if arg is node_arg and i in \ WEIGHT_INDEX_DICT[node.target]: # type: ignore[index] return True + for kwarg_name, kwarg_value in node.kwargs.items(): + if kwarg_name == 'weight' and arg is kwarg_value: + return True return False CONV_OPS_WITH_BIAS = { @@ -103,7 +108,7 @@ def node_arg_is_bias(node: Node, arg: Any) -> bool: for i, node_arg in enumerate(node.args): if arg is node_arg and i == CONV_BIAS_ARG_INDEX: return True - elif node.target is torch.nn.functional.linear: + elif node.target in FUNCTIONAL_OPS_WITH_BIAS: for kwarg_name, kwarg_value in node.kwargs.items(): if kwarg_name == 'bias' and arg is kwarg_value: return True @@ -277,7 +282,7 @@ def maybe_insert_input_observer_for_arg_or_kwarg( qconfig.activation is_bias = node_arg_is_bias(node, arg) is_activation = not (is_weight or is_bias) - weight_needs_obs = is_weight and weight_is_quantized(qconfig) + weight_needs_obs = is_weight and weight_is_quantized(qconfig) and node.target not in NON_QUANTIZABLE_WEIGHT_OPS bias_needs_obs = \ (is_bias and activation_dtype(qconfig) == torch.float16) and \ weight_dtype(qconfig) == torch.float16 diff --git a/torch/quantization/fx/utils.py b/torch/quantization/fx/utils.py index c96891a0405eb..76b481f8b38e0 100644 --- a/torch/quantization/fx/utils.py +++ b/torch/quantization/fx/utils.py @@ -19,6 +19,18 @@ torch.nn.functional.conv2d : [1], torch.nn.functional.conv3d : [1], torch.nn.functional.linear : [1], + torch.nn.functional.layer_norm : [2], + torch.nn.functional.group_norm : [2], + torch.nn.functional.instance_norm : [3], +} + +NON_QUANTIZABLE_WEIGHT_OPS = {torch.nn.functional.layer_norm, torch.nn.functional.group_norm, torch.nn.functional.instance_norm} + +FUNCTIONAL_OPS_WITH_BIAS = { + torch.nn.functional.linear, + torch.nn.functional.layer_norm, + torch.nn.functional.group_norm, + torch.nn.functional.instance_norm } # turn foo.bar -> ['foo', 'bar'] From f3218568ada0d13fb5f849fa23460a28ad3e571b Mon Sep 17 00:00:00 2001 From: mingfeima Date: Fri, 11 Jun 2021 16:21:38 -0700 Subject: [PATCH 061/305] optimize channels last for BatchNorm2d on CPU (#59286) Summary: replacement of https://github.com/pytorch/pytorch/issues/48919 optimize channels last performance for BatchNorm2 on CPU. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59286 Reviewed By: bdhirsh Differential Revision: D29008198 Pulled By: VitalyFedyunin fbshipit-source-id: 8a7d020bd6a42ab5c21ffe788b79a22f4ec82ac0 --- aten/src/ATen/cpu/vec/vec256/functional.h | 29 + aten/src/ATen/native/Normalization.cpp | 178 ++--- aten/src/ATen/native/batch_norm.h | 24 +- .../src/ATen/native/cpu/batch_norm_kernel.cpp | 651 +++++++++++++++++- aten/src/ATen/test/vec_test_all_types.cpp | 43 ++ aten/src/ATen/test/vec_test_all_types.h | 1 + test/test_nn.py | 33 + 7 files changed, 805 insertions(+), 154 deletions(-) diff --git a/aten/src/ATen/cpu/vec/vec256/functional.h b/aten/src/ATen/cpu/vec/vec256/functional.h index 05971b14b049d..e6ac5c5cee548 100644 --- a/aten/src/ATen/cpu/vec/vec256/functional.h +++ b/aten/src/ATen/cpu/vec/vec256/functional.h @@ -231,4 +231,33 @@ inline void map3( } } +template +inline void map4( + const Op& vec_fun, + scalar_t* output_data, + const scalar_t* input_data1, + const scalar_t* input_data2, + const scalar_t* input_data3, + const scalar_t* input_data4, + int64_t size) { + using Vec = vec::Vectorized; + int64_t d = 0; + for (; d < size - (size % Vec::size()); d += Vec::size()) { + Vec data_vec1 = Vec::loadu(input_data1 + d); + Vec data_vec2 = Vec::loadu(input_data2 + d); + Vec data_vec3 = Vec::loadu(input_data3 + d); + Vec data_vec4 = Vec::loadu(input_data4 + d); + Vec output_vec = vec_fun(data_vec1, data_vec2, data_vec3, data_vec4); + output_vec.store(output_data + d); + } + if (size - d > 0) { + Vec data_vec1 = Vec::loadu(input_data1 + d, size - d); + Vec data_vec2 = Vec::loadu(input_data2 + d, size - d); + Vec data_vec3 = Vec::loadu(input_data3 + d, size - d); + Vec data_vec4 = Vec::loadu(input_data4 + d, size - d); + Vec output_vec = vec_fun(data_vec1, data_vec2, data_vec3, data_vec4); + output_vec.store(output_data + d, size - d); + } +} + }} // namespace at::vec diff --git a/aten/src/ATen/native/Normalization.cpp b/aten/src/ATen/native/Normalization.cpp index 4b5b345616b5b..d48ce03ee45cc 100644 --- a/aten/src/ATen/native/Normalization.cpp +++ b/aten/src/ATen/native/Normalization.cpp @@ -35,7 +35,11 @@ TORCH_META_FUNC(renorm)(const Tensor& self, const Scalar& p, int64_t dim, const namespace native { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -DEFINE_DISPATCH(batch_norm_cpu_inference_contiguous_stub); +DEFINE_DISPATCH(batch_norm_cpu_stub); +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +DEFINE_DISPATCH(batch_norm_cpu_collect_stats_stub); +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +DEFINE_DISPATCH(batch_norm_cpu_backward_stub); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(renorm_scale_factor_stub); @@ -53,15 +57,6 @@ namespace { } } -// TensorAccessor when it is defined to work around undefined... -template -static TensorAccessor conditional_accessor_1d(const Tensor& t) { - if (! t.defined()) { - return TensorAccessor(nullptr, nullptr, nullptr); - } - return t.accessor(); -} - template struct InvStd { T operator()(T var, double epsilon) const { @@ -80,87 +75,8 @@ struct Var { } }; -template -void batch_norm_cpu_inference_collect_linear_and_constant_terms( - scalar_t* alpha, scalar_t* beta, int64_t n_channel, - const Tensor& weight /* optional */, const Tensor& bias /* optional */, - const Tensor& mean, const Tensor& variance, double eps) { - - const scalar_t* weight_data = weight.defined() ? weight.data_ptr() : nullptr; - const scalar_t* bias_data = bias.defined() ? bias.data_ptr() : nullptr; - const scalar_t* mean_data = mean.data_ptr(); - const scalar_t* var_data = variance.data_ptr(); - - /// Collect the linear and constant terms regarding the input. - /// output(n, c, h, w) - /// = (input(n, c, h, w) - mean(c)) / sqrt(var(c) + eps) * weight(c) - /// + bias(c) - /// = input(n, c, h, w) * inv_var(c) * weight(c) - /// - mean(c) * inv_var(c) * weight(c) + bias(c), - /// where inv_var(c) = 1 / sqrt(var(c) + eps). - /// So the linear term, alpha(c) = inv_var(c) * weight(c), - /// the constant term beta(c) = bias(c) - mean(c) * inv_var(c) * weight(c) - /// Note that this is only a good idea if (input_size >> c), in degenerate - /// cases where image_size == 1 && batch_size == 1, it is slow. - for (int64_t c = 0; c < n_channel; c++) { - scalar_t inv_var = 1 / std::sqrt(var_data[c] + static_cast(eps)); - scalar_t weight_v = weight_data ? weight_data[c] : 1; - scalar_t bias_v = bias_data ? bias_data[c] : 0; - alpha[c] = inv_var * weight_v; - beta[c] = bias_v - mean_data[c] * inv_var * weight_v; - } -} - -/// A fast path for CPU inference when all tensors are channels last contiguous. -/// This code achieves machine bandwidth peak without AVX support. -/// If this changes for future architectures, we can move it to the cpu/ -/// directory. -template -void batch_norm_cpu_inference_channels_last(Tensor& output, const Tensor& input, - const Tensor& weight /* optional */, const Tensor& bias /* optional */, - const Tensor& mean, const Tensor& variance, double eps) { - - int64_t n_batch = input.size(0); - int64_t n_channel = input.size(1); - int64_t image_size = input.numel() / n_batch / n_channel; - - scalar_t* output_data = output.data_ptr(); - const scalar_t* input_data = input.data_ptr(); - - Tensor alpha = at::empty_like(mean, LEGACY_CONTIGUOUS_MEMORY_FORMAT); - Tensor beta = at::empty_like(mean, LEGACY_CONTIGUOUS_MEMORY_FORMAT); - scalar_t* alpha_data = alpha.data_ptr(); - scalar_t* beta_data = beta.data_ptr(); - - batch_norm_cpu_inference_collect_linear_and_constant_terms( - alpha_data, beta_data, n_channel, weight, bias, mean, variance, eps); - - // Apply the linear terms to the input, - // output(n, c, h, w) = input(n, c, h, w) * alpha(c) + beta(c) - // No need to use parallel_for as this function is supposed to be - // memory-limited. - // Keep the loop structure simple to make sure compiler vectorization kicks in. - if (n_channel != 1) { - for (int64_t n = 0; n < n_batch; ++n) { - for (int64_t i = 0; i < image_size; ++i) { - for (int64_t c = 0; c < n_channel; ++c) { - // Keep all the offset calculation within the inner loop for - // simplicity. Compilers are very good at hoisting the common part - // outside. - int64_t offset = n * image_size * n_channel + i * n_channel + c; - output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c]; - } - } - } - } else { - // n_channel == 1 - for (int64_t n = 0; n < n_batch; ++n) { - for (int64_t i = 0; i < image_size; ++i) { - int64_t offset = n * image_size + i; - output_data[offset] = input_data[offset] * alpha_data[0] + beta_data[0]; - } - } - } +static inline bool is_contiguous(const Tensor& t) { + return t.is_contiguous() || t.is_contiguous(at::MemoryFormat::ChannelsLast); } template @@ -170,29 +86,18 @@ std::tuple batch_norm_cpu_transform_input_template( const Tensor& running_mean /* optional */, const Tensor& running_var /* optional */, bool train, double eps) { - // Check if we should use the fast path for contiguous memory format - if (!train && input.is_contiguous() + bool all_contiguous = is_contiguous(input) && (!weight.defined() || weight.is_contiguous()) && (!bias.defined() || bias.is_contiguous()) && running_mean.is_contiguous() - && running_var.is_contiguous()) { - - Tensor output = at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT); - batch_norm_cpu_inference_contiguous_stub(kCPU, output, input, weight, - bias, running_mean, running_var, eps); - return std::make_tuple(output, save_mean, save_invstd); - } + && running_var.is_contiguous(); - // Check if we should use the fast path for channel last memory format - if (!train && input.is_contiguous(at::MemoryFormat::ChannelsLast) - && (!weight.defined() || weight.is_contiguous()) - && (!bias.defined() || bias.is_contiguous()) - && running_mean.is_contiguous() - && running_var.is_contiguous()) { + Tensor output = at::empty_like(input, input.suggest_memory_format()); - Tensor output = at::empty_like(input, at::MemoryFormat::ChannelsLast); - batch_norm_cpu_inference_channels_last( - output, input, weight, bias, running_mean, running_var, eps); + // inference contiguous path + if (all_contiguous) { + batch_norm_cpu_stub(kCPU, output, input, weight, bias, + save_mean, save_invstd, running_mean, running_var, train, eps); return std::make_tuple(output, save_mean, save_invstd); } @@ -220,7 +125,6 @@ std::tuple batch_norm_cpu_transform_input_template( auto b = bias.defined() ? as_nd(bias) : at::detail::scalar_tensor_static(0, input.scalar_type(), kCPU); - Tensor output = at::empty(input.sizes(), input.options()); auto iter = TensorIteratorConfig() .add_output(output) .add_input(input) @@ -262,6 +166,34 @@ std::tuple batch_norm_cpu_update_stats_template( auto running_mean_a = conditional_accessor_1d(running_mean); auto running_var_a = conditional_accessor_1d(running_var); + bool all_contiguous = is_contiguous(input); + if (all_contiguous) { + auto _mean = at::empty({n_input}, input.options()); + auto _var_sum = at::empty({n_input}, input.options()); + auto _mean_a = _mean.accessor(); + auto _var_sum_a = _var_sum.accessor(); + + batch_norm_cpu_collect_stats_stub(kCPU, _mean, _var_sum, input); + + parallel_for(0, n_input, 1, [&](int64_t b_begin, int64_t b_end) { + for (int64_t f = b_begin; f < b_end; ++f) { + save_mean_a[f] = _mean_a[f]; + save_var_transform_a[f] = VarTransform{}(_var_sum_a[f] / n, eps); + + if (running_mean.defined()) { + running_mean_a[f] = momentum * _mean_a[f] + (1 - momentum) * running_mean_a[f]; + } + if (running_var.defined()) { + accscalar_t unbiased_var = _var_sum_a[f] / (n - 1); + running_var_a[f] = momentum * unbiased_var + (1 - momentum) * running_var_a[f]; + } + } + }); + + return std::make_tuple(save_mean, save_var_transform); + } + + // non-contiguous path parallel_for(0, n_input, 1, [&](int64_t b_begin, int64_t b_end) { for (int64_t f = b_begin; f < b_end; ++f) { Tensor in = input.select(1, f); @@ -290,11 +222,11 @@ std::tuple batch_norm_cpu_update_stats_template( return std::make_tuple(save_mean, save_var_transform); } - template -std::tuple batch_norm_backward_cpu_template(const Tensor& grad_out_, const Tensor& input, const Tensor& weight, - const Tensor& running_mean, const Tensor& running_var, const Tensor& save_mean, const Tensor& save_invstd, - bool train, double eps, std::array grad_input_mask) { +std::tuple batch_norm_backward_cpu_template( + const Tensor& grad_out_, const Tensor& input, const Tensor& weight, + const Tensor& running_mean, const Tensor& running_var, const Tensor& save_mean, const Tensor& save_invstd, + bool train, double eps, std::array grad_input_mask) { using accscalar_t = at::acc_type; @@ -302,13 +234,25 @@ std::tuple batch_norm_backward_cpu_template(const Tensor Tensor grad_weight; Tensor grad_bias; if (grad_input_mask[0]) { - grad_input = at::empty_like(input, LEGACY_CONTIGUOUS_MEMORY_FORMAT); + grad_input = at::empty_like(input, input.suggest_memory_format()); } if (grad_input_mask[1]) { - grad_weight = at::empty_like(weight, LEGACY_CONTIGUOUS_MEMORY_FORMAT); + grad_weight = at::empty_like(weight, at::MemoryFormat::Contiguous); } if (grad_input_mask[2]) { - grad_bias = at::empty_like(weight, LEGACY_CONTIGUOUS_MEMORY_FORMAT); + grad_bias = at::empty_like(weight, at::MemoryFormat::Contiguous); + } + + // since we are directly manipulating pointers in contiguous path, + // need to make sure input and grad_out have the same memory format. + bool all_contiguous = is_contiguous(input) + && is_contiguous(grad_out_) + && input.suggest_memory_format() == grad_out_.suggest_memory_format(); + + if (all_contiguous) { + batch_norm_cpu_backward_stub(kCPU, grad_input, grad_weight, grad_bias, + grad_out_, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps); + return std::make_tuple(grad_input, grad_weight, grad_bias); } auto weight_a = conditional_accessor_1d(weight); diff --git a/aten/src/ATen/native/batch_norm.h b/aten/src/ATen/native/batch_norm.h index bd78a5b7ebe24..4c25b08aa684d 100644 --- a/aten/src/ATen/native/batch_norm.h +++ b/aten/src/ATen/native/batch_norm.h @@ -8,9 +8,29 @@ namespace at { namespace native { using batch_norm_fn = void (*)(Tensor&, const Tensor&, const Tensor&, - const Tensor&, const Tensor&, const Tensor&, double); + const Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, bool, double); +using batch_norm_collect_stats_fn = void (*)(Tensor&, Tensor&, const Tensor&); +using batch_norm_backward_fn = void(*)(Tensor&, Tensor&, Tensor&, const Tensor&, + const Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&, bool, double); -DECLARE_DISPATCH(batch_norm_fn, batch_norm_cpu_inference_contiguous_stub); +DECLARE_DISPATCH(batch_norm_fn, batch_norm_cpu_stub); +DECLARE_DISPATCH(batch_norm_collect_stats_fn, batch_norm_cpu_collect_stats_stub); +DECLARE_DISPATCH(batch_norm_backward_fn, batch_norm_cpu_backward_stub); + +// TensorAccessor when it is defined to work around undefined... +template +static TensorAccessor conditional_accessor_1d(const Tensor& t) { + if (! t.defined()) { + return TensorAccessor(nullptr, nullptr, nullptr); + } + return t.accessor(); +} + +template +static scalar_t* conditional_data_ptr(const Tensor& t) { + return t.defined() ? t.contiguous().data_ptr() + : nullptr; +} } // namespace native diff --git a/aten/src/ATen/native/cpu/batch_norm_kernel.cpp b/aten/src/ATen/native/cpu/batch_norm_kernel.cpp index 940329c13ea45..ff4c6f2dee075 100644 --- a/aten/src/ATen/native/cpu/batch_norm_kernel.cpp +++ b/aten/src/ATen/native/cpu/batch_norm_kernel.cpp @@ -1,10 +1,15 @@ #include #include +#include #include #include +#include #include #include +#include +#include +#include namespace at { namespace native { namespace { @@ -12,15 +17,19 @@ namespace { using namespace vec; template -void batch_norm_cpu_inference_collect_linear_and_constant_terms( - TensorAccessor alpha, TensorAccessor beta, int64_t n_channel, +void batch_norm_cpu_collect_linear_and_constant_terms( + scalar_t* alpha, scalar_t* beta, int64_t n_channel, const Tensor& weight /* optional */, const Tensor& bias /* optional */, - const Tensor& mean, const Tensor& variance, double eps) { + const Tensor& save_mean, const Tensor& save_invstd, + const Tensor& running_mean, const Tensor& running_var, bool train, double eps) { const scalar_t* weight_data = weight.defined() ? weight.data_ptr() : nullptr; const scalar_t* bias_data = bias.defined() ? bias.data_ptr() : nullptr; - auto mean_data = mean.accessor(); - auto var_data = variance.accessor(); + + auto save_mean_a = conditional_accessor_1d(save_mean); + auto save_invstd_a = conditional_accessor_1d(save_invstd); + auto running_mean_a = conditional_accessor_1d(running_mean); + auto running_var_a = conditional_accessor_1d(running_var); /// Collect the linear and constant terms regarding the input. /// output(n, c, h, w) @@ -34,48 +43,57 @@ void batch_norm_cpu_inference_collect_linear_and_constant_terms( /// Note that this is only a good idea if (input_size >> c), in degenerate /// cases where image_size == 1 && batch_size == 1, it is slow. for (int64_t c = 0; c < n_channel; c++) { - scalar_t inv_var = 1 / std::sqrt(var_data[c] + static_cast(eps)); + scalar_t mean, invstd; + if (train) { + mean = save_mean_a[c]; + invstd = save_invstd_a[c]; + } else { + mean = running_mean_a[c]; + invstd = 1 / std::sqrt(running_var_a[c] + static_cast(eps)); + } scalar_t weight_v = weight_data ? weight_data[c] : 1; scalar_t bias_v = bias_data ? bias_data[c] : 0; - alpha[c] = inv_var * weight_v; - beta[c] = bias_v - mean_data[c] * alpha[c]; + alpha[c] = invstd * weight_v; + beta[c] = bias_v - mean * alpha[c]; } } -/// A fast path for CPU inference when all tensors are contiguous. +/// A fast path for CPU inference and training forward when all tensors are contiguous. template -void batch_norm_cpu_inference_contiguous_impl(Tensor& output, - const Tensor& input, const Tensor& weight, const Tensor& bias, - const Tensor& mean, const Tensor& variance, double eps) { +void batch_norm_cpu_contiguous_impl(Tensor& output, const Tensor& input, + const Tensor& weight, const Tensor& bias, const Tensor& save_mean, const Tensor& save_invstd, + const Tensor& running_mean, const Tensor& running_var, bool train, double eps) { using Vec = Vectorized; int64_t n_batch = input.size(0); int64_t n_channel = input.size(1); int64_t image_size = input.numel() / n_batch / n_channel; - Tensor alpha = at::empty_like(mean, LEGACY_CONTIGUOUS_MEMORY_FORMAT); - Tensor beta = at::empty_like(mean, LEGACY_CONTIGUOUS_MEMORY_FORMAT); - auto alpha_data = alpha.accessor(); - auto beta_data = beta.accessor(); + Tensor alpha = at::empty({n_channel}, input.options()); + Tensor beta = at::empty({n_channel}, input.options()); + scalar_t* alpha_data = alpha.data_ptr(); + scalar_t* beta_data = beta.data_ptr(); - batch_norm_cpu_inference_collect_linear_and_constant_terms( - alpha_data, beta_data, n_channel, weight, bias, mean, variance, eps); + batch_norm_cpu_collect_linear_and_constant_terms( + alpha_data, beta_data, n_channel, weight, bias, + save_mean, save_invstd, running_mean, running_var, train, eps); scalar_t* output_data = output.data_ptr(); const scalar_t* input_data = input.data_ptr(); // Apply the linear terms to the input, // output(n, c, h, w) = input(n, c, h, w) * alpha(c) + beta(c) - // No need to use parallel_for as this function is supposed to be - // memory-limited. if (image_size != 1) { - const int64_t n_offset = n_channel * image_size; const int64_t loop_size = image_size - (image_size % Vec::size()); - for (int64_t n = 0; n < n_batch; n++) { - for (int64_t c = 0; c < n_channel; c++) { + at::parallel_for(0, n_batch * n_channel, 1, [&](int64_t begin, int64_t end) { + int64_t n = 0; + int64_t c = 0; + data_index_init(begin, n, n_batch, c, n_channel); + + for (int64_t i = begin; i < end; i++) { const Vec alpha_vec(alpha_data[c]); const Vec beta_vec(beta_data[c]); - int64_t offset = n * n_offset + c * image_size; + int64_t offset = i * image_size; int64_t d = 0; for (; d < loop_size; d += Vec::size()) { Vec data_vec = Vec::loadu(input_data + offset + d); @@ -87,29 +105,592 @@ void batch_norm_cpu_inference_contiguous_impl(Tensor& output, Vec output_vec = data_vec * alpha_vec + beta_vec; output_vec.store(output_data + offset + d, image_size - d); } + // move on to next index + data_index_step(n, n_batch, c, n_channel); } - } + }); } else { // image_size == 1 - for (int64_t n = 0; n < n_batch; ++n) { - for (int64_t c = 0; c < n_channel; ++c) { - int64_t offset = n * n_channel + c; - output_data[offset] = input_data[offset] * alpha_data[c] + beta_data[c]; + const int64_t loop_size = n_channel - (n_channel % Vec::size()); + at::parallel_for(0, n_batch, 1, [&](int64_t begin, int64_t end) { + for (int64_t n = begin; n < end; n++) { + int64_t offset = n * n_channel; + int64_t d = 0; + for (; d < loop_size; d += Vec::size()) { + Vec alpha_vec = Vec::loadu(alpha_data + d); + Vec beta_vec = Vec::loadu(beta_data + d); + Vec data_vec = Vec::loadu(input_data + offset + d); + Vec output_vec = data_vec * alpha_vec + beta_vec; + output_vec.store(output_data + offset + d); + } + if (n_channel - d > 0) { + Vec alpha_vec = Vec::loadu(alpha_data + d, n_channel - d); + Vec beta_vec = Vec::loadu(beta_data + d, n_channel - d); + Vec data_vec = Vec::loadu(input_data + offset + d, n_channel - d); + Vec output_vec = data_vec * alpha_vec + beta_vec; + output_vec.store(output_data + offset + d, n_channel - d); + } } - } + }); } } -void batch_norm_cpu_inference_contiguous_kernel(Tensor& output, const Tensor& input, - const Tensor& weight, const Tensor& bias, const Tensor& mean, const Tensor& variance, double eps) { - AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_inference_contiguous", [&] { - batch_norm_cpu_inference_contiguous_impl(output, input, weight, bias, mean, variance, eps); +template +void batch_norm_cpu_channels_last_impl(Tensor& output, const Tensor& input, + const Tensor& weight, const Tensor& bias, const Tensor& save_mean, const Tensor& save_invstd, + const Tensor& running_mean, const Tensor& runnning_var, bool train, double eps) { + + using Vec = Vectorized; + int64_t n_batch = input.size(0); + int64_t n_channel = input.size(1); + int64_t image_size = input.numel() / n_batch / n_channel; + + Tensor alpha = at::empty({n_channel}, input.options()); + Tensor beta = at::empty({n_channel}, input.options()); + scalar_t* alpha_data = alpha.data_ptr(); + scalar_t* beta_data = beta.data_ptr(); + + batch_norm_cpu_collect_linear_and_constant_terms( + alpha_data, beta_data, n_channel, weight, bias, + save_mean, save_invstd, running_mean, runnning_var, train, eps); + + scalar_t* output_data = output.data_ptr(); + const scalar_t* input_data = input.data_ptr(); + + // Apply the linear terms to the input, + // output(n, c, h, w) = input(n, c, h, w) * alpha(c) + beta(c) + const int64_t loop_size = n_channel - (n_channel % Vec::size()); + at::parallel_for(0, n_batch * image_size, 1, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; i++) { + int64_t offset = i * n_channel; + int64_t d = 0; + // vectorize on channel dimension, for normal batch_norm input size, + // alpha/beta should fit in L1 cache, otherwise consider blocking. + for (; d < loop_size; d += Vec::size()) { + Vec alpha_vec = Vec::loadu(alpha_data + d); + Vec beta_vec = Vec::loadu(beta_data + d); + Vec data_vec = Vec::loadu(input_data + offset + d); + Vec output_vec = data_vec * alpha_vec + beta_vec; + output_vec.store(output_data + offset + d); + } + if (n_channel - d > 0) { + Vec alpha_vec = Vec::loadu(alpha_data + d, n_channel - d); + Vec beta_vec = Vec::loadu(beta_data + d, n_channel - d); + Vec data_vec = Vec::loadu(input_data + offset + d, n_channel - d); + Vec output_vec = data_vec * alpha_vec + beta_vec; + output_vec.store(output_data + offset + d, n_channel - d); + } + } }); } +template +void batch_norm_cpu_collect_stats_contiguous_impl( + Tensor& mean, Tensor& var_sum, const Tensor& input) { + + using accscalar_t = at::acc_type; + int64_t n_batch = input.size(0); + int64_t n_channel = input.size(1); + int64_t image_size = input.numel() / n_batch / n_channel; + int64_t N = input.numel() / n_channel; + + const scalar_t* input_data = input.data_ptr(); + scalar_t* mean_data = mean.data_ptr(); + scalar_t* var_sum_data = var_sum.data_ptr(); + + // parallel dim reduce on 'channel' + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (int64_t c = begin; c < end; c++) { + // compute mean per input + accscalar_t sum = 0; + for (int64_t n = 0; n < n_batch; n++) { + for (int64_t i = 0; i < image_size; i++) { + auto offset = n * n_channel * image_size + c * image_size + i; + sum += input_data[offset]; + } + } + scalar_t mean = sum / N; + mean_data[c] = mean; + + // compute variance per input + accscalar_t _var_sum = 0; + for (int64_t n = 0; n < n_batch; n++) { + for (int64_t i = 0; i < image_size; i++) { + auto offset = n * n_channel * image_size + c * image_size + i; + auto x = input_data[offset]; + _var_sum += (x - mean) * (x - mean); + } + } + var_sum_data[c] = _var_sum; + } + }); +} + +template +void batch_norm_cpu_collect_stats_channels_last_impl( + Tensor& mean, Tensor& var_sum, const Tensor& input) { + + using Vec = Vectorized; + using accscalar_t = at::acc_type; + int64_t n_channel = input.size(1); + int64_t N = input.numel() / n_channel; + + const scalar_t* input_data = input.data_ptr(); + scalar_t* mean_data = mean.data_ptr(); + scalar_t* var_sum_data = var_sum.data_ptr(); + + // Typical vertical reduce from shape of {NHW, C} to {C}. + // Apply two path parallel reduction: + // First path: allocate an immediate buffer of size {max_threads, C}, parallel along dim0, + // {NHW, C} => {max_threads, C} + // + // Second path: parallel along dim1 of the immediate buffer, + // {max_threads, C} => {C} + // + // Normal size of C should fit in L1, otherwise consider blocking on C. + // + int num_threads = at::get_num_threads(); + Tensor buffer = at::empty({num_threads, n_channel}, input.options()).zero_(); + scalar_t* buffer_data = buffer.data_ptr(); + + // compute mean per input + at::parallel_for(0, N, 1, [&](int64_t begin, int64_t end) { + int tid = at::get_thread_num(); + TORCH_CHECK(tid < num_threads, + "expect thread id smaller than ", num_threads, ", got thread id ", tid); + scalar_t* buffer_ptr = buffer_data + tid * n_channel; + for (int64_t i = begin; i < end; i++) { + const scalar_t* x_ptr = input_data + i * n_channel; + vec::map2( + [](Vec x, Vec y) { return x + y; }, + buffer_ptr, + x_ptr, + buffer_ptr, + n_channel); + } + }); + + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (int64_t c = begin; c < end; c++) { + accscalar_t sum = 0; + for (int64_t t = 0; t < num_threads; t++) { + sum += buffer_data[t * n_channel + c]; + } + scalar_t mean = sum / N; + mean_data[c] = mean; + } + }); + + // compute variance per input, reuse the immediate buffer + buffer.zero_(); + at::parallel_for(0, N, 1, [&](int64_t begin, int64_t end) { + int tid = at::get_thread_num(); + TORCH_CHECK(tid < num_threads, "expect thread id smaller than ", num_threads, ", got thread id ", tid); + scalar_t* buffer_ptr = buffer_data + tid * n_channel; + for (int64_t i = begin; i < end; i++) { + const scalar_t* x_ptr = input_data + i * n_channel; + vec::map3( + [](Vec x, Vec y, Vec mean) { return y + (x - mean) * (x - mean); }, + buffer_ptr, + x_ptr, + buffer_ptr, + mean_data, + n_channel); + } + }); + + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (int64_t c = begin; c < end; c++) { + accscalar_t _var_sum = 0; + for (int64_t t = 0; t < num_threads; t++) { + _var_sum += buffer_data[t * n_channel + c]; + } + var_sum_data[c] = _var_sum; + } + }); +} + +template +void batch_norm_cpu_backward_contiguous_impl(Tensor& grad_input, Tensor& grad_weight, Tensor& grad_bias, + const Tensor& grad_output, const Tensor& input, const Tensor& weight, + const Tensor& running_mean, const Tensor& running_var, const Tensor& save_mean, const Tensor& save_invstd, + bool train, double eps) { + + using Vec = Vectorized; + using accscalar_t = at::acc_type; + int64_t n_batch = input.size(0); + int64_t n_channel = input.size(1); + int64_t image_size = input.numel() / n_batch / n_channel; + int64_t N = input.numel() / n_channel; + + const scalar_t* grad_output_data = grad_output.data_ptr(); + const scalar_t* input_data = input.data_ptr(); + + scalar_t* grad_input_data = grad_input.defined() ? grad_input.data_ptr() : nullptr; + scalar_t* grad_weight_data = grad_weight.defined() ? grad_weight.data_ptr() : nullptr; + scalar_t* grad_bias_data = grad_bias.defined() ? grad_bias.data_ptr() : nullptr; + const bool grad_input_null = grad_input_data == nullptr; + const bool grad_weight_null = grad_weight_data == nullptr; + const bool grad_bias_null = grad_bias_data == nullptr; + + auto weight_a = conditional_accessor_1d(weight); + auto save_mean_a = conditional_accessor_1d(save_mean); + auto save_invstd_a = conditional_accessor_1d(save_invstd); + auto running_mean_a = conditional_accessor_1d(running_mean); + auto running_var_a = conditional_accessor_1d(running_var); + + // parallel dim reduce on 'channel' + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (int64_t c = begin; c < end; c++) { + scalar_t w = weight.defined() ? weight_a[c] : 1; + + scalar_t mean, invstd; + if (train) { + mean = save_mean_a[c]; + invstd = save_invstd_a[c]; + } else { + mean = running_mean_a[c]; + invstd = 1 / std::sqrt(running_var_a[c] + eps); + } + + // reduce over grad_output in feature plane + // compute 1) sum; 2) dot product of Q(X) and dY. + // fuse into a single loop to reuse dY + // + accscalar_t sum = 0; + accscalar_t dotp = 0; + for (int64_t n = 0; n < n_batch; n++) { + const scalar_t* x_ptr = input_data + n * n_channel * image_size + c * image_size; + const scalar_t* dy_ptr = grad_output_data + n * n_channel * image_size + c * image_size; + + sum += vec::reduce_all( + [](Vec& x, Vec& y) { return x + y; }, + dy_ptr, + image_size); + + dotp += vec::map2_reduce_all( + [mean](Vec x, Vec dy) { return (x - Vec(mean)) * dy; }, + [](Vec x, Vec y) { return x + y; }, + x_ptr, + dy_ptr, + image_size); + } + + if (!grad_input_null) { + if (train) { + scalar_t k = (scalar_t) dotp * invstd * invstd / N; + scalar_t grad_mean = sum / N; + + for (int64_t n = 0; n < n_batch; n++) { + const scalar_t* x_ptr = input_data + n * n_channel * image_size + c * image_size; + scalar_t* dx_ptr = grad_input_data + n * n_channel * image_size + c * image_size; + const scalar_t* dy_ptr = grad_output_data + n * n_channel * image_size + c * image_size; + + // Scalar math: + // for (int64_t j = 0; j < image_size; ++j) { + // scalar_t dx = (x_ptr[j] - mean) * k; + // dx_ptr[j] = (dy_ptr[j] - grad_mean - dx) * invstd * w; + // } + vec::map2( + [=](Vec x, Vec dy) { + Vec dx = (x - Vec(mean)) * Vec(k); + return (dy - Vec(grad_mean) - dx) * Vec(invstd) * Vec(w); + }, + dx_ptr, + x_ptr, + dy_ptr, + image_size); + } + } else { // evaluation mode + for (int64_t n = 0; n < n_batch; n++) { + scalar_t* dx_ptr = grad_input_data + n * n_channel * image_size + c * image_size; + const scalar_t* dy_ptr = grad_output_data + n * n_channel * image_size + c * image_size; + + // Scalar math: + // for (int64_t j = 0; j < image_size; ++j) { + // dx_ptr[j] = dy_ptr[j] * invstd * w; + // } + vec::map( + [=](Vec dy) { return dy * Vec(invstd) * Vec(w); }, + dx_ptr, + dy_ptr, + image_size); + } + } + } + + if (!grad_weight_null) { + grad_weight_data[c] = dotp * invstd; + } + + if (!grad_bias_null) { + grad_bias_data[c] = sum; + } + } + }); +} + +template +void batch_norm_cpu_backward_channels_last_impl(Tensor& grad_input, Tensor& grad_weight, Tensor& grad_bias, + const Tensor& grad_output, const Tensor& input, const Tensor& weight, + const Tensor& running_mean, const Tensor& running_var, const Tensor& save_mean, const Tensor& save_invstd, + bool train, double eps) { + + using Vec = Vectorized; + using accscalar_t = at::acc_type; + int64_t n_channel = input.size(1); + int64_t N = input.numel() / n_channel; + + const scalar_t* grad_output_data = grad_output.data_ptr(); + const scalar_t* input_data = input.data_ptr(); + + scalar_t* grad_input_data = grad_input.defined() ? grad_input.data_ptr() : nullptr; + scalar_t* grad_weight_data = grad_weight.defined() ? grad_weight.data_ptr() : nullptr; + scalar_t* grad_bias_data = grad_bias.defined() ? grad_bias.data_ptr() : nullptr; + + scalar_t* save_mean_data = conditional_data_ptr(save_mean); + scalar_t* save_invstd_data = conditional_data_ptr(save_invstd); + scalar_t* running_mean_data = conditional_data_ptr(running_mean); + scalar_t* running_var_data = conditional_data_ptr(running_var); + + Tensor weight_ = weight.defined() ? weight : at::ones({n_channel}, input.options()); + const scalar_t* weight_data = weight_.data_ptr(); + + scalar_t* mean_ptr = nullptr; + scalar_t* invstd_ptr = nullptr; + Tensor invstd = at::empty({0}, input.options()); + if (train) { + mean_ptr = save_mean_data; + invstd_ptr = save_invstd_data; + } else { + mean_ptr = running_mean_data; + + invstd.resize_({n_channel}); + invstd_ptr = invstd.data_ptr(); + for (int64_t c = 0; c < n_channel; c++) { + invstd_ptr[c] = 1 / std::sqrt(running_var_data[c] + eps); + } + } + + // Typical vertical reduce from shape of {NHW, C} to {C}. + // Apply two path parallel reduction: + // First path: allocate an immediate buffer of size {2, max_threads, C}, parallel along dim0, + // sum = buffer[0], dotp = buffer[2] + // + // Second path: parallel along dim1 of the immediate buffer. + // + int num_threads = at::get_num_threads(); + Tensor buffer = at::empty({2, num_threads, n_channel}, input.options()).zero_(); + scalar_t* sum_data = buffer.data_ptr(); + scalar_t* dotp_data = sum_data + num_threads * n_channel; + + // compute sum and dotp per feature plain, + // fuse into a single loop to reuse grad_output in L1. + at::parallel_for(0, N, 1, [&](int64_t begin, int64_t end) { + int tid = at::get_thread_num(); + TORCH_CHECK(tid < num_threads, "expect thread id smaller than ", num_threads, ", got thread id ", tid); + scalar_t* sum_ptr = sum_data + tid * n_channel; + scalar_t* dotp_ptr = dotp_data + tid * n_channel; + for (int64_t i = begin; i < end; i++) { + const scalar_t* x_ptr = input_data + i * n_channel; + const scalar_t* dy_ptr = grad_output_data + i * n_channel; + + vec::map2( + [](Vec sum, Vec dy) { return sum + dy; }, + sum_ptr, + sum_ptr, + dy_ptr, + n_channel); + + vec::map4( + [](Vec dotp, Vec x, Vec mean, Vec dy) { return dotp + (x - mean) * dy; }, + dotp_ptr, + dotp_ptr, + x_ptr, + mean_ptr, + dy_ptr, + n_channel); + } + }); + + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (int64_t c = begin; c < end; c++) { + // store the final result of sum and dotp in the 1st lane of immediate buffer, + // so that we won't need to allocate anther buffer to store the temp values. + accscalar_t _sum = 0; + for (int64_t t = 0; t < num_threads; t++) { + _sum += sum_data[t * n_channel + c]; + } + sum_data[/* 0 * n_channel + */c] = _sum; + + accscalar_t _dotp = 0; + for (int64_t t = 0; t < num_threads; t++) { + _dotp += dotp_data[t * n_channel + c]; + } + dotp_data[/* 0 * n_channel + */c] = _dotp; + } + }); + + // compute grad_input + const int64_t loop_size = n_channel - (n_channel % Vec::size()); + if (grad_input.defined()) { + at::parallel_for(0, N, 1, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; i++) { + scalar_t* dx_ptr = grad_input_data + i * n_channel; + const scalar_t* x_ptr = input_data + i * n_channel; + const scalar_t* dy_ptr = grad_output_data + i * n_channel; + if (train) { + int64_t d = 0; + for (; d < loop_size; d += Vec::size()) { + Vec x = Vec::loadu(x_ptr + d); + Vec mean = Vec::loadu(mean_ptr + d); + Vec dotp = Vec::loadu(dotp_data + d); + Vec invstd = Vec::loadu(invstd_ptr + d); + Vec k = dotp * invstd * invstd / Vec(N); + Vec dx = (x - mean) * k; + Vec dy = Vec::loadu(dy_ptr + d); + Vec grad_mean = Vec::loadu(sum_data + d) / Vec(N); + Vec w = Vec::loadu(weight_data + d); + dx = (dy - grad_mean - dx) * invstd * w; + dx.store(dx_ptr + d); + } + if (n_channel - d > 0) { + Vec x = Vec::loadu(x_ptr + d, n_channel - d); + Vec mean = Vec::loadu(mean_ptr + d, n_channel - d); + Vec dotp = Vec::loadu(dotp_data + d, n_channel - d); + Vec invstd = Vec::loadu(invstd_ptr + d, n_channel - d); + Vec k = dotp * invstd * invstd / Vec(N); + Vec dx = (x - mean) * k; + Vec dy = Vec::loadu(dy_ptr + d, n_channel - d); + Vec grad_mean = Vec::loadu(sum_data + d, n_channel - d) / Vec(N); + Vec w = Vec::loadu(weight_data + d, n_channel - d); + dx = (dy - grad_mean - dx) * invstd * w; + dx.store(dx_ptr + d, n_channel - d); + } + } else { // evaluation mode + int64_t d = 0; + for (; d < loop_size; d += Vec::size()) { + Vec dy = Vec::loadu(dy_ptr + d); + Vec invstd = Vec::loadu(invstd_ptr + d); + Vec w = Vec::loadu(weight_data + d); + Vec dx = dy * invstd * w; + dx.store(dx_ptr + d); + } + if (n_channel - d > 0) { + Vec dy = Vec::loadu(dy_ptr + d, n_channel - d); + Vec invstd = Vec::loadu(invstd_ptr + d, n_channel - d); + Vec w = Vec::loadu(weight_data + d, n_channel - d); + Vec dx = dy * invstd * w; + dx.store(dx_ptr + d, n_channel - d); + } + } + } + }); + } + + if (grad_weight.defined()) { + // grad_weight = dotp * invstd + vec::map2( + [](Vec dotp, Vec invstd) { return dotp * invstd; }, + grad_weight_data, + dotp_data, + invstd_ptr, + n_channel); + } + + // grad_bias = sum + if (grad_bias.defined()) { + vec::map( + [](Vec sum) { return sum; }, + grad_bias_data, + sum_data, + n_channel); + } +} + +void batch_norm_cpu_kernel(Tensor& output, const Tensor& input, + const Tensor& weight, const Tensor& bias, const Tensor& save_mean, const Tensor& save_invstd, + const Tensor& running_mean, const Tensor& running_var, bool train, double eps) { + switch (input.suggest_memory_format()) { + case at::MemoryFormat::Contiguous: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_contiguous", [&] { + batch_norm_cpu_contiguous_impl(output, input, weight, bias, + save_mean, save_invstd, running_mean, running_var, train, eps); + }); + break; + } + case at::MemoryFormat::ChannelsLast: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_channels_last", [&] { + batch_norm_cpu_channels_last_impl(output, input, weight, bias, + save_mean, save_invstd, running_mean, running_var, train, eps); + }); + break; + } + default: + TORCH_CHECK(false, "Unsupported memory format. Supports only ChannelsLast, Contiguous"); + } +} + +void batch_norm_cpu_collect_stats_kernel( + Tensor& mean, Tensor& var_sum, const Tensor& input) { + int64_t image_size = input.numel() / input.size(0) / input.size(1); + switch (input.suggest_memory_format()) { + case at::MemoryFormat::Contiguous: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_collect_stats_contiguous", [&] { + if (image_size == 1) { // NC11 is also channels last + batch_norm_cpu_collect_stats_channels_last_impl(mean, var_sum, input); + } else { + batch_norm_cpu_collect_stats_contiguous_impl(mean, var_sum, input); + } + }); + break; + } + case at::MemoryFormat::ChannelsLast: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_collect_stats_channels_last", [&] { + batch_norm_cpu_collect_stats_channels_last_impl(mean, var_sum, input); + }); + break; + } + default: + TORCH_CHECK(false, "Unsupported memory format. Supports only ChannelsLast, Contiguous"); + } +} + +void batch_norm_cpu_backward_kernel(Tensor& grad_input, Tensor& grad_weight, Tensor& grad_bias, + const Tensor& grad_output, const Tensor& input, const Tensor& weight, + const Tensor& running_mean, const Tensor& running_var, const Tensor& save_mean, const Tensor& save_invstd, + bool train, double eps) { + int64_t image_size = input.numel() / input.size(0) / input.size(1); + switch (input.suggest_memory_format()) { + case at::MemoryFormat::Contiguous: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_backward_contiguous", [&] { + if (image_size == 1) { // NC11 is also channels last + batch_norm_cpu_backward_channels_last_impl(grad_input, grad_weight, grad_bias, + grad_output, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps); + } else { + batch_norm_cpu_backward_contiguous_impl(grad_input, grad_weight, grad_bias, + grad_output, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps); + } + }); + break; + } + case at::MemoryFormat::ChannelsLast: { + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "batch_norm_cpu_backward_channels_last", [&] { + batch_norm_cpu_backward_channels_last_impl(grad_input, grad_weight, grad_bias, + grad_output, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps); + }); + break; + } + default: + TORCH_CHECK(false, "Unsupported memory format. Supports only ChannelsLast, Contiguous"); + } +} + }// anonymous namespace // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -REGISTER_DISPATCH(batch_norm_cpu_inference_contiguous_stub, &batch_norm_cpu_inference_contiguous_kernel); +REGISTER_DISPATCH(batch_norm_cpu_stub, &batch_norm_cpu_kernel); +REGISTER_DISPATCH(batch_norm_cpu_collect_stats_stub, &batch_norm_cpu_collect_stats_kernel); +REGISTER_DISPATCH(batch_norm_cpu_backward_stub, &batch_norm_cpu_backward_kernel); }} // namespace at::native diff --git a/aten/src/ATen/test/vec_test_all_types.cpp b/aten/src/ATen/test/vec_test_all_types.cpp index 844e8d84bd378..1f1698ae33635 100644 --- a/aten/src/ATen/test/vec_test_all_types.cpp +++ b/aten/src/ATen/test/vec_test_all_types.cpp @@ -57,6 +57,8 @@ namespace { class ComplexTests : public ::testing::Test {}; template class QuantizationTests : public ::testing::Test {}; + template + class FunctionalTests : public ::testing::Test {}; using RealFloatTestedTypes = ::testing::Types; using FloatTestedTypes = ::testing::Types; using ALLTestedTypes = ::testing::Types; @@ -91,6 +93,7 @@ namespace { TYPED_TEST_CASE(BitwiseFloatsAdditional, RealFloatTestedTypes); TYPED_TEST_CASE(BitwiseFloatsAdditional2, FloatTestedTypes); TYPED_TEST_CASE(QuantizationTests, QuantTestedTypes); + TYPED_TEST_CASE(FunctionalTests, RealFloatIntTestedTypes); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) TYPED_TEST(Memory, UnAlignedLoadStore) { using vec = TypeParam; @@ -1297,6 +1300,46 @@ namespace { }, test_case); } + TYPED_TEST(FunctionalTests, Map) { + using vec = TypeParam; + using VT = ValueType; + constexpr auto R = 2LL; // residual + constexpr auto N = vec::size() + R; + CACHE_ALIGN VT x1[N]; + CACHE_ALIGN VT x2[N]; + CACHE_ALIGN VT x3[N]; + CACHE_ALIGN VT x4[N]; + CACHE_ALIGN VT y[N]; + CACHE_ALIGN VT ref_y[N]; + auto seed = TestSeed(); + ValueGen generator(VT(-100), VT(100), seed); + for (int64_t i = 0; i < N; i++) { + x1[i] = generator.get(); + x2[i] = generator.get(); + x3[i] = generator.get(); + x4[i] = generator.get(); + } + auto cmp = [&](VT* y, VT* ref_y) { + AssertVectorized(NAME_INFO(Map), vec::loadu(y), vec::loadu(ref_y)).check(true); + AssertVectorized(NAME_INFO(Map), vec::loadu(y + vec::size(), R), vec::loadu(ref_y + vec::size(), R)).check(true); + }; + // test map: y = x1 + at::vec::map([](vec x) { return x; }, y, x1, N); + for (int64_t i = 0; i < N; i++) { ref_y[i] = x1[i]; } + cmp(y, ref_y); + // test map2: y = x1 + x2 + at::vec::map2([](vec x1, vec x2) { return x1 + x2; }, y, x1, x2, N); + for (int64_t i = 0; i < N; i++) { ref_y[i] = x1[i] + x2[i]; } + cmp(y, ref_y); + // test map3: y = x1 + x2 + x3 + at::vec::map3([](vec x1, vec x2, vec x3) { return x1 + x2 + x3; }, y, x1, x2, x3, N); + for (int64_t i = 0; i < N; i++) { ref_y[i] = x1[i] + x2[i] + x3[i]; } + cmp(y, ref_y); + // test map3: y = x1 + x2 + x3 + x4 + at::vec::map4([](vec x1, vec x2, vec x3, vec x4) { return x1 + x2 + x3 + x4; }, y, x1, x2, x3, x4, N); + for (int64_t i = 0; i < N; i++) { ref_y[i] = x1[i] + x2[i] + x3[i] + x4[i]; } + cmp(y, ref_y); + } #else #error GTEST does not have TYPED_TEST diff --git a/aten/src/ATen/test/vec_test_all_types.h b/aten/src/ATen/test/vec_test_all_types.h index 4150e55bb54bb..ce96b514cc0a2 100644 --- a/aten/src/ATen/test/vec_test_all_types.h +++ b/aten/src/ATen/test/vec_test_all_types.h @@ -1,5 +1,6 @@ #pragma once #include +#include #include #include #include diff --git a/test/test_nn.py b/test/test_nn.py index 650463748ff6d..29357bd13ec51 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -8508,6 +8508,39 @@ def test_hardtanh_backward(self): x_grad_ref = torch.where(mask, grad, z) self.assertEqual(x.grad, x_grad_ref) + def test_batchnorm_nhwc_cpu(self): + def helper(self, size): + channels = size[1] + input = torch.randn(size, dtype=torch.float32, device='cpu', requires_grad=True) + input = input.contiguous(memory_format=torch.channels_last) + input.retain_grad() + grad = torch.randn(size, dtype=torch.float32, device='cpu') + grad = grad.contiguous(memory_format=torch.channels_last) + bn = nn.BatchNorm2d(channels).cpu().float() + bn.weight.data.uniform_() + bn.bias.data.uniform_() + + ref_input = input.detach().clone().contiguous().requires_grad_(True) + ref_grad = grad.detach().clone().contiguous() + ref_bn = nn.BatchNorm2d(channels).cpu().float() + ref_bn.load_state_dict(bn.state_dict()) + + out = bn(input) + out.backward(grad) + ref_out = ref_bn(ref_input) + ref_out.backward(ref_grad) + + self.assertTrue(out.is_contiguous(memory_format=torch.channels_last)) + self.assertTrue(ref_out.is_contiguous()) + self.assertEqual(out, ref_out) + self.assertEqual(bn.weight.grad, ref_bn.weight.grad) + self.assertEqual(bn.bias.grad, ref_bn.bias.grad) + self.assertEqual(input.grad, ref_input.grad) + + helper(self, (4, 8, 10, 10)) + helper(self, (4, 1, 9, 9)) + helper(self, (4, 9, 1, 1)) + @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") @unittest.skipIf(not TEST_CUDNN, "needs cudnn") @skipIfRocm From 3529a48ebb3605b9a0578f1392c0a3bdbf3b207b Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Fri, 11 Jun 2021 17:15:55 -0700 Subject: [PATCH 062/305] Revert D28981326: torch/lib/c10d: Use torch_check instead of throwing runtime_error Test Plan: revert-hammer Differential Revision: D28981326 (https://github.com/pytorch/pytorch/commit/6ea607500284d74d9f7f6f91b489b355a1fa7fad) Original commit changeset: 264a7f787ea8 fbshipit-source-id: 75625b76dfbd0cbaf59705d621ef9e2d1677c482 --- torch/lib/c10d/FileStore.cpp | 6 +- torch/lib/c10d/GlooDeviceFactory.cpp | 4 +- torch/lib/c10d/NCCLUtils.hpp | 5 +- torch/lib/c10d/ProcessGroup.cpp | 8 +-- torch/lib/c10d/ProcessGroup.hpp | 12 ++-- torch/lib/c10d/ProcessGroupGloo.cpp | 42 ++++++------- torch/lib/c10d/ProcessGroupMPI.cpp | 46 +++++++------- torch/lib/c10d/ProcessGroupNCCL.cpp | 66 ++++++++++---------- torch/lib/c10d/ProcessGroupRoundRobin.cpp | 10 +-- torch/lib/c10d/TCPStore.cpp | 10 +-- torch/lib/c10d/UnixSockUtils.hpp | 2 +- torch/lib/c10d/Utils.cpp | 8 +-- torch/lib/c10d/Utils.hpp | 16 ++--- torch/lib/c10d/WinSockUtils.hpp | 2 +- torch/lib/c10d/frontend.cpp | 2 +- torch/lib/c10d/test/ProcessGroupGlooTest.cpp | 2 +- torch/lib/c10d/test/ProcessGroupMPITest.cpp | 22 +++---- torch/lib/c10d/test/TCPStoreTest.cpp | 9 ++- 18 files changed, 135 insertions(+), 137 deletions(-) diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp index 73342272c54c0..ea98963ee9df8 100644 --- a/torch/lib/c10d/FileStore.cpp +++ b/torch/lib/c10d/FileStore.cpp @@ -273,7 +273,7 @@ FileStore::FileStore(const std::string& path, int numWorkers) cleanupKey_("cleanup/"), regularPrefix_("/") { if (numWorkers_ < 1) { - TORCH_CHECK(false, + throw std::runtime_error( "Number of workers for FileStore should be greater than zero"); } } @@ -341,7 +341,7 @@ std::vector FileStore::get(const std::string& key) { const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout_ != kNoTimeout && elapsed > timeout_) { - TORCH_CHECK(false, "Timeout waiting for key: " + key); + throw std::runtime_error("Timeout waiting for key: " + key); } std::this_thread::sleep_for(std::chrono::milliseconds(10)); continue; @@ -424,7 +424,7 @@ void FileStore::wait( const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout != kNoTimeout && elapsed > timeout) { - TORCH_CHECK(false, "Wait timeout"); + throw std::runtime_error("Wait timeout"); } /* sleep override */ diff --git a/torch/lib/c10d/GlooDeviceFactory.cpp b/torch/lib/c10d/GlooDeviceFactory.cpp index cb83a99838520..416676483e182 100644 --- a/torch/lib/c10d/GlooDeviceFactory.cpp +++ b/torch/lib/c10d/GlooDeviceFactory.cpp @@ -149,7 +149,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForInterface(const std::string& interfaceName) { auto device = makeGlooDevice(interfaceName, ""); if (!device) { - TORCH_CHECK(false, "makeDeviceForInterface(): unsupported gloo device"); + throw std::runtime_error("makeDeviceForInterface(): unsupported gloo device"); } return device; } @@ -158,7 +158,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForHostname(const std::string& hostname) { auto device = makeGlooDevice("", hostname); if (!device) { - TORCH_CHECK(false, "makeDeviceForHostname(): unsupported gloo device"); + throw std::runtime_error("makeDeviceForHostname(): unsupported gloo device"); } return device; } diff --git a/torch/lib/c10d/NCCLUtils.hpp b/torch/lib/c10d/NCCLUtils.hpp index e3ee14da0f542..0dec4573112a1 100644 --- a/torch/lib/c10d/NCCLUtils.hpp +++ b/torch/lib/c10d/NCCLUtils.hpp @@ -9,7 +9,6 @@ #include #include -#include namespace { // Provides additional detail into NCCL error codes based on when these are @@ -58,7 +57,7 @@ const inline char* getNcclErrorDetailStr(ncclResult_t error) { std::string err = "NCCL error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \ "\n" + getNcclErrorDetailStr(result); \ - TORCH_CHECK(false, err); \ + throw std::runtime_error(err); \ } \ } while (0) @@ -143,7 +142,7 @@ class NCCLComm { ncclComm_t getNcclComm() { std::unique_lock lock(mutex_); if (aborted_) { - TORCH_CHECK(false, + throw std::runtime_error( "NCCL communicator was aborted on rank " + std::to_string(rank_) + "."); } diff --git a/torch/lib/c10d/ProcessGroup.cpp b/torch/lib/c10d/ProcessGroup.cpp index 4e03824eb12da..39ae2bf71c598 100644 --- a/torch/lib/c10d/ProcessGroup.cpp +++ b/torch/lib/c10d/ProcessGroup.cpp @@ -107,13 +107,13 @@ std::exception_ptr ProcessGroup::Work::exception() const { } int ProcessGroup::Work::sourceRank() const { - TORCH_CHECK(false, + throw std::runtime_error( "sourceRank() may only be called on work objects " "that correspond to a recv or recv-from-any call."); } std::vector ProcessGroup::Work::result() { - TORCH_CHECK(false, "result() not implemented."); + throw std::runtime_error("result() not implemented."); } void ProcessGroup::Work::synchronize() {} @@ -129,7 +129,7 @@ bool ProcessGroup::Work::wait(std::chrono::milliseconds timeout) { if (!completed_) { // Throw exception if the wait operation timed out and the work was not // completed. - TORCH_CHECK(false, "Operation timed out!"); + throw std::runtime_error("Operation timed out!"); } } if (exception_) { @@ -186,7 +186,7 @@ c10::intrusive_ptr ProcessGroup::allgather_coalesced( std::vector>& /* usused */, std::vector& /* usused */, const AllgatherOptions& /* usused */) { - TORCH_CHECK(false, + throw std::runtime_error( "no support for allgather_coalesced in this process group"); } diff --git a/torch/lib/c10d/ProcessGroup.hpp b/torch/lib/c10d/ProcessGroup.hpp index 3a3ffa6b95d67..ee2990fd33975 100644 --- a/torch/lib/c10d/ProcessGroup.hpp +++ b/torch/lib/c10d/ProcessGroup.hpp @@ -264,7 +264,7 @@ class ProcessGroup : public torch::CustomClassHolder { at::Tensor&, at::Tensor&, const ReduceScatterOptions& opts = ReduceScatterOptions()) { - TORCH_CHECK(false, "ProcessGroup does not support reduce_scatter_base"); + throw std::runtime_error("ProcessGroup does not support reduce_scatter_base"); } @@ -274,20 +274,20 @@ class ProcessGroup : public torch::CustomClassHolder { std::vector& outputSplitSizes, std::vector& inputSplitSizes, const AllToAllOptions& opts = AllToAllOptions()) { - TORCH_CHECK(false, "ProcessGroup does not support alltoall"); + throw std::runtime_error("ProcessGroup does not support alltoall"); } virtual c10::intrusive_ptr alltoall( std::vector& outputTensors, std::vector& inputTensors, const AllToAllOptions& opts = AllToAllOptions()) { - TORCH_CHECK(false, "ProcessGroup does not support alltoall"); + throw std::runtime_error("ProcessGroup does not support alltoall"); } virtual void monitoredBarrier( const BarrierOptions& /* unused */, bool /* unused */ = false ) { auto backendName = getBackendName(); - TORCH_CHECK(false, + throw std::runtime_error( c10::str("ProcessGroup ", backendName, " does not support monitoredBarrier, only GLOO supports monitored barrier.") @@ -299,7 +299,7 @@ class ProcessGroup : public torch::CustomClassHolder { // for GLOO and NCCL backends currently. virtual void setSequenceNumberForGroup() { auto backendName = getBackendName(); - TORCH_CHECK(false, + throw std::runtime_error( c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") @@ -311,7 +311,7 @@ class ProcessGroup : public torch::CustomClassHolder { // may indicate that there is some sort of collective desynchronization. virtual uint64_t getSequenceNumberForGroup() { auto backendName = getBackendName(); - TORCH_CHECK(false, + throw std::runtime_error( c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp index 98164237feb9c..d423271192db8 100644 --- a/torch/lib/c10d/ProcessGroupGloo.cpp +++ b/torch/lib/c10d/ProcessGroupGloo.cpp @@ -66,7 +66,7 @@ func(__VA_ARGS__); \ break; \ default: \ - TORCH_CHECK(false, "Invalid scalar type"); \ + throw std::runtime_error("Invalid scalar type"); \ } #define HOST_NAME_MAX 256 @@ -95,7 +95,7 @@ func(args); \ break; \ default: \ - TORCH_CHECK(false, "Invalid scalar type"); \ + throw std::runtime_error("Invalid scalar type"); \ } #endif @@ -178,22 +178,22 @@ ReduceFunc toFunction(const ReduceOp& r) { case ReduceOp::MAX: return ReduceFunc(&::gloo::max); case ReduceOp::BAND: - TORCH_CHECK(false, + throw std::runtime_error( "Cannot use ReduceOp.BAND with non-integral dtype"); break; case ReduceOp::BOR: - TORCH_CHECK(false, + throw std::runtime_error( "Cannot use ReduceOp.BOR with non-integral dtype"); break; case ReduceOp::BXOR: - TORCH_CHECK(false, + throw std::runtime_error( "Cannot use ReduceOp.BXOR with non-integral dtype"); break; case ReduceOp::UNUSED: break; } - TORCH_CHECK(false, "Unhandled ReduceOp"); + throw std::runtime_error("Unhandled ReduceOp"); } // Bitwise AND with SFINAE guard for integral types. @@ -258,7 +258,7 @@ ReduceFunc toFunction(const ReduceOp& r) { break; } - TORCH_CHECK(false, "Unhandled ReduceOp"); + throw std::runtime_error("Unhandled ReduceOp"); } template @@ -368,7 +368,7 @@ void initializeStreamsEvents( const auto device_id = tensorgroup[0].device().index(); for (const auto& tensor : tensorgroup) { if (tensor.device().index() != device_id) { - TORCH_CHECK(false, + throw std::runtime_error( "tensors in the nested tensor vectors need to " "be on the same device"); } @@ -683,7 +683,7 @@ ProcessGroupGloo::ProcessGroupGloo( collectiveCounter_(0) { auto& devices = options->devices; if (devices.empty()) { - TORCH_CHECK(false, "No device(s) specified"); + throw std::runtime_error("No device(s) specified"); } // Create and connect a context for every device. @@ -915,7 +915,7 @@ c10::intrusive_ptr ProcessGroupGloo::broadcast( work = c10::make_intrusive( std::move(context), inputs, opts.rootRank, opts.rootTensor, tag); } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); @@ -1426,7 +1426,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( invalidArgument("unsupported layout"); } } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); @@ -1487,7 +1487,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce_coalesced( invalidArgument("unsupported layout"); } } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); return work; @@ -1646,7 +1646,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( opts.reduceOp, tag); } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); return work; @@ -1838,7 +1838,7 @@ c10::intrusive_ptr ProcessGroupGloo::allgather( work = c10::make_intrusive( std::move(context), outputs, inputs, tag); } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); return work; @@ -1972,7 +1972,7 @@ c10::intrusive_ptr ProcessGroupGloo::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "no support for _allgather_base in Gloo process group"); } @@ -2166,7 +2166,7 @@ c10::intrusive_ptr ProcessGroupGloo::gather( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); return work; @@ -2349,7 +2349,7 @@ c10::intrusive_ptr ProcessGroupGloo::scatter( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - TORCH_CHECK(false, "Invalid backend"); + throw std::runtime_error("Invalid backend"); } enqueue(work); return work; @@ -2359,7 +2359,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce_scatter( std::vector& outputs, std::vector>& inputs, const ReduceScatterOptions& opts) { - TORCH_CHECK(false, "ProcessGroupGloo does not support reduce_scatter"); + throw std::runtime_error("ProcessGroupGloo does not support reduce_scatter"); } namespace { @@ -2531,14 +2531,14 @@ c10::intrusive_ptr ProcessGroupGloo::alltoall_base( at::Tensor& checkSingleTensor(std::vector& tensors) { if (tensors.size() != 1) { - TORCH_CHECK(false, "ProcessGroupGloo::send takes a single tensor"); + throw std::runtime_error("ProcessGroupGloo::send takes a single tensor"); } auto& tensor = tensors[0]; if (!tensor.is_contiguous()) { - TORCH_CHECK(false, "input tensor has to be contiguous"); + throw std::runtime_error("input tensor has to be contiguous"); } if (tensor.is_sparse()) { - TORCH_CHECK(false, "input tensor has to be dense"); + throw std::runtime_error("input tensor has to be dense"); } return tensor; } diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp index aa6d81bbe4a13..0c471216dffa7 100644 --- a/torch/lib/c10d/ProcessGroupMPI.cpp +++ b/torch/lib/c10d/ProcessGroupMPI.cpp @@ -21,7 +21,7 @@ namespace c10d { std::string err = "MPI error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + \ ", with error code: " + std::to_string(mpiStatus); \ - TORCH_CHECK(false, err); \ + throw std::runtime_error(err); \ } \ } while (0) @@ -63,13 +63,13 @@ bool cudaAwareMpiCheck() { // Checking the input tensor's validity void checkSingleTensorHelper(const at::Tensor& tensor) { if (!tensor.is_contiguous()) { - TORCH_CHECK(false, "input tensor has to be contiguous"); + throw std::runtime_error("input tensor has to be contiguous"); } if (tensor.is_sparse()) { - TORCH_CHECK(false, "input tensor has to be dense"); + throw std::runtime_error("input tensor has to be dense"); } if (tensor.is_cuda() && !cudaAwareMpiCheck()) { - TORCH_CHECK(false, + throw std::runtime_error( "CUDA tensor detected and the MPI used doesn't " "have CUDA-aware MPI support"); } @@ -77,7 +77,7 @@ void checkSingleTensorHelper(const at::Tensor& tensor) { void checkSingleTensor(const std::vector& tensors) { if (tensors.size() != 1) { - TORCH_CHECK(false, + throw std::runtime_error( "MPI process group does not support multi-GPU collectives"); } checkSingleTensorHelper(tensors[0]); @@ -89,7 +89,7 @@ void checkSameSizeAndType( for (const auto& tensor : tensors) { if ((tensor.numel() != t_in.numel()) || (tensor.scalar_type() != t_in.scalar_type())) { - TORCH_CHECK(false, "Tensors are not equal in size or data type"); + throw std::runtime_error("Tensors are not equal in size or data type"); } checkSingleTensorHelper(tensor); } @@ -158,7 +158,7 @@ bool ProcessGroupMPI::AsyncWork::isCompleted() { bool ProcessGroupMPI::AsyncWork::isSuccess() const { if (request_ != MPI_REQUEST_NULL) { - TORCH_CHECK(false, + throw std::runtime_error( "Invalid call to AsyncWork::isSuccess before work has completed"); } @@ -232,14 +232,14 @@ void ProcessGroupMPI::initMPIOnce() { MPI_CHECK(MPI_Init_thread( nullptr, nullptr, MPI_THREAD_SERIALIZED, &mpiThreadSupport_)); if (mpiThreadSupport_ < MPI_THREAD_SERIALIZED) { - TORCH_CHECK(false, + throw std::runtime_error( "Used MPI implementation doesn't have the " "minimum level of threading support: " "MPI_THREAD_SERIALIZED. This is required by " "c10d package"); } if (std::atexit(ProcessGroupMPI::mpiExit)) { - TORCH_CHECK(false, "Fail to register the MPI exit handler"); + throw std::runtime_error("Fail to register the MPI exit handler"); } }); } @@ -285,7 +285,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( MPI_CHECK(MPI_Comm_size(groupComm, &size)); if (rank < 0 || size < 0) { - TORCH_CHECK(false, "Failed to get the world_size / rank"); + throw std::runtime_error("Failed to get the world_size / rank"); } } } @@ -303,7 +303,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( ProcessGroupMPI::ProcessGroupMPI(int rank, int size, MPI_Comm pgComm) : ProcessGroup(rank, size), stop_(false), pgComm_(pgComm) { if (pgComm_ == MPI_COMM_NULL) { - TORCH_CHECK(false, "pgComm_ must not be MPI_COMM_NULL"); + throw std::runtime_error("pgComm_ must not be MPI_COMM_NULL"); } // Start the worker thread accepting MPI calls @@ -427,7 +427,7 @@ c10::intrusive_ptr ProcessGroupMPI::allreduce( c10::intrusive_ptr ProcessGroupMPI::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - TORCH_CHECK(false, + throw std::runtime_error( "allreduce_coalesced is currently not supported with MPI"); } @@ -467,12 +467,12 @@ c10::intrusive_ptr ProcessGroupMPI::allgather( const AllgatherOptions& opts) { checkSingleTensor(inputTensors); if (outputTensors.size() != 1) { - TORCH_CHECK(false, + throw std::runtime_error( "MPI process group only supports a single " "tensor op"); } if (static_cast(size_) != outputTensors[0].size()) { - TORCH_CHECK(false, + throw std::runtime_error( "All gather: number of output tensors should equal " "to the world size"); } @@ -512,7 +512,7 @@ c10::intrusive_ptr ProcessGroupMPI::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupMPI does not support allgather_coalesced"); } @@ -524,16 +524,16 @@ c10::intrusive_ptr ProcessGroupMPI::gather( if (rank_ != opts.rootRank) { if (outputTensors.size() > 0) { - TORCH_CHECK(false, + throw std::runtime_error( "Gather: number of output tensors should be 0 " "for non-root"); } } else { if (outputTensors.size() != 1) { - TORCH_CHECK(false, "Gather: multi-GPU collective is not supported"); + throw std::runtime_error("Gather: multi-GPU collective is not supported"); } if (static_cast(size_) != outputTensors[0].size()) { - TORCH_CHECK(false, + throw std::runtime_error( "Gather: number of output tensors should equal " "to the world size"); } @@ -598,17 +598,17 @@ c10::intrusive_ptr ProcessGroupMPI::scatter( if (rank_ != opts.rootRank) { if (inputTensors.size() > 0) { - TORCH_CHECK(false, + throw std::runtime_error( "Scatter: number of input tensors should be 0 " "for non-root"); } } else { if (inputTensors.size() != 1) { - TORCH_CHECK(false, + throw std::runtime_error( "Scatter: multi-GPU collective is not supported"); } if (static_cast(size_) != inputTensors[0].size()) { - TORCH_CHECK(false, + throw std::runtime_error( "Scatter: number of input tensors should equal " "to the world size"); } @@ -670,7 +670,7 @@ c10::intrusive_ptr ProcessGroupMPI::reduce_scatter( std::vector& outputTensors, std::vector>& inputTensors, const ReduceScatterOptions& opts) { - TORCH_CHECK(false, "ProcessGroupMPI does not support reduce_scatter"); + throw std::runtime_error("ProcessGroupMPI does not support reduce_scatter"); } c10::intrusive_ptr ProcessGroupMPI::alltoall_base( @@ -917,7 +917,7 @@ c10::intrusive_ptr ProcessGroupMPI::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "no support for _allgather_base in MPI process group"); } diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp index f538e2f4ea560..3f62cab44602b 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.cpp +++ b/torch/lib/c10d/ProcessGroupNCCL.cpp @@ -90,16 +90,16 @@ ncclRedOp_t getNcclReduceOp(const ReduceOp reduceOp, at::Tensor& input) { } catch (const std::out_of_range& e) { switch (reduceOp) { case ReduceOp::BAND: - TORCH_CHECK(false, "Cannot use ReduceOp.BAND with NCCL"); + throw std::runtime_error("Cannot use ReduceOp.BAND with NCCL"); break; case ReduceOp::BOR: - TORCH_CHECK(false, "Cannot use ReduceOp.BOR with NCCL"); + throw std::runtime_error("Cannot use ReduceOp.BOR with NCCL"); break; case ReduceOp::BXOR: - TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with NCCL"); + throw std::runtime_error("Cannot use ReduceOp.BXOR with NCCL"); break; default: - TORCH_CHECK(false, "Unhandled ReduceOp"); + throw std::runtime_error("Unhandled ReduceOp"); break; } } @@ -396,7 +396,7 @@ void ProcessGroupNCCL::WorkNCCL::synchronizeInternal( " ran for ", timeElapsed.count(), " milliseconds before timing out."); - TORCH_CHECK(false, exceptionMsg); + throw std::runtime_error(exceptionMsg); } // Check for errors and throw appropriate exception. checkAndThrowException(); @@ -819,7 +819,7 @@ std::vector>& ProcessGroupNCCL::getNCCLComm( bool isSendRecvSelf) { // Sanity check if (devicesKey.empty()) { - TORCH_CHECK(false, + throw std::runtime_error( "Not able to create/get the NCCL Communicator since " "the GPU devices are not known"); } @@ -945,10 +945,10 @@ namespace { // Check validity of tensor void check_gpu_single_tensor(const at::Tensor& tensor) { if (!tensor.is_cuda() || tensor.is_sparse()) { - TORCH_CHECK(false, "Tensors must be CUDA and dense"); + throw std::runtime_error("Tensors must be CUDA and dense"); } if (!tensor.is_contiguous()) { - TORCH_CHECK(false, "Tensors must be contiguous"); + throw std::runtime_error("Tensors must be contiguous"); } } @@ -956,10 +956,10 @@ void check_gpu_single_tensor(const at::Tensor& tensor) { // across distinct GPUs. void check_gpu_tensors(const std::vector& tensors) { if (tensors.size() == 0) { - TORCH_CHECK(false, "Tensor list must be nonempty"); + throw std::runtime_error("Tensor list must be nonempty"); } if (tensors.size() > static_cast(at::cuda::getNumGPUs())) { - TORCH_CHECK(false, + throw std::runtime_error( "Tensor list mustn't be larger than the number of available GPUs"); } @@ -971,23 +971,23 @@ void check_gpu_tensors(const std::vector& tensors) { for (const auto& t : tensors) { if (!t.is_cuda() || t.is_sparse()) { - TORCH_CHECK(false, "Tensors must be CUDA and dense"); + throw std::runtime_error("Tensors must be CUDA and dense"); } if (t.scalar_type() != first.scalar_type()) { - TORCH_CHECK(false, "Tensors must have identical type"); + throw std::runtime_error("Tensors must have identical type"); } if (t.sizes() != first.sizes()) { - TORCH_CHECK(false, "Tensors must have identical size"); + throw std::runtime_error("Tensors must have identical size"); } if (t.strides() != first.strides()) { - TORCH_CHECK(false, "Tensors must have identical strides"); + throw std::runtime_error("Tensors must have identical strides"); } if (!t.is_non_overlapping_and_dense()) { - TORCH_CHECK(false, "Tensors must be non-overlapping and dense"); + throw std::runtime_error("Tensors must be non-overlapping and dense"); } const auto inserted = usedDevices.insert(t.get_device()).second; if (!inserted) { - TORCH_CHECK(false, "Tensors must be on distinct GPU devices"); + throw std::runtime_error("Tensors must be on distinct GPU devices"); } } } @@ -999,7 +999,7 @@ std::vector flatten_for_scatter_gather( std::vector& other, size_t world_size) { if (tensor_lists.size() != other.size()) { - TORCH_CHECK(false, + throw std::runtime_error( "Tensor list operands to scatter/gather must have the same length"); } const auto num_devices = tensor_lists.size(); @@ -1009,7 +1009,7 @@ std::vector flatten_for_scatter_gather( for (auto i = size_t{}; i < num_devices; ++i) { if (tensor_lists[i].size() != world_size * num_devices) { - TORCH_CHECK(false, + throw std::runtime_error( "Tensor list input to scatter/gather must match number of collective" " participants"); } @@ -1017,14 +1017,14 @@ std::vector flatten_for_scatter_gather( // Only check device match for the first tensor in the list; the call to // newLikeFlat() below will check the rest. if (tensor_lists[i].front().get_device() != other[i].get_device()) { - TORCH_CHECK(false, + throw std::runtime_error( "Corresponding input/output tensors to scatter/gather must all reside" " on the same device"); } for (const auto& t : tensor_lists[i]) { if (t.numel() != other[i].numel()) { - TORCH_CHECK(false, + throw std::runtime_error( "All tensor operands to scatter/gather must have the same number of elements"); } } @@ -1343,7 +1343,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allreduce( c10::intrusive_ptr ProcessGroupNCCL::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - TORCH_CHECK(false, + throw std::runtime_error( "allreduce_coalesced is currently not supported with NCCL"); } @@ -1481,7 +1481,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupNCCL does not support allgather_coalesced"); } @@ -1549,11 +1549,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_reduce_scatter_base( const ReduceScatterOptions& opts) { if (inputTensor.dtype() != outputTensor.dtype()) { - TORCH_CHECK(false, "input tensor must be the same type as the outut tensor."); + throw std::runtime_error("input tensor must be the same type as the outut tensor."); } if (inputTensor.numel() != outputTensor.numel() * size_) { - TORCH_CHECK(false, "input tensor must be the same size as output size times world size"); + throw std::runtime_error("input tensor must be the same size as output size times world size"); } // @lint-ignore CLANGTIDY @@ -1821,7 +1821,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall_base( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1829,7 +1829,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1837,7 +1837,7 @@ c10::intrusive_ptr ProcessGroupNCCL::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupNCCL only supports send for NCCL lib version >= 2.7.0"); } @@ -1845,7 +1845,7 @@ c10::intrusive_ptr ProcessGroupNCCL::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "ProcessGroupNCCL only supports recv for NCCL lib version >= 2.7.0"); } #endif @@ -1868,20 +1868,20 @@ c10::intrusive_ptr ProcessGroupNCCL::gather( std::vector>& /* unused */, std::vector& /* unused */, const GatherOptions& /* unused */) { - TORCH_CHECK(false, "ProcessGroupNCCL does not support gather"); + throw std::runtime_error("ProcessGroupNCCL does not support gather"); } c10::intrusive_ptr ProcessGroupNCCL::scatter( std::vector& /* unused */, std::vector>& /* unused */, const ScatterOptions& /* unused */) { - TORCH_CHECK(false, "ProcessGroupNCCL does not support scatter"); + throw std::runtime_error("ProcessGroupNCCL does not support scatter"); } c10::intrusive_ptr ProcessGroupNCCL::recvAnysource( std::vector& /* unused */, int /* unused */) { - TORCH_CHECK(false, "ProcessGroupNCCL does not support recvAnysource"); + throw std::runtime_error("ProcessGroupNCCL does not support recvAnysource"); } c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( @@ -1892,11 +1892,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( check_gpu_single_tensor(output_tensor); if (input_tensor.dtype() != output_tensor.dtype()) { - TORCH_CHECK(false, "output tensor must have the same type as input tensor"); + throw std::runtime_error("output tensor must have the same type as input tensor"); } if (input_tensor.numel() * size_ != output_tensor.numel()) { - TORCH_CHECK(false, "output tensor size must be equal to world_size times input tensor size"); + throw std::runtime_error("output tensor size must be equal to world_size times input tensor size"); } // just a wrapper to fit the collective interface diff --git a/torch/lib/c10d/ProcessGroupRoundRobin.cpp b/torch/lib/c10d/ProcessGroupRoundRobin.cpp index c439cf771a147..a55eea968b1e1 100644 --- a/torch/lib/c10d/ProcessGroupRoundRobin.cpp +++ b/torch/lib/c10d/ProcessGroupRoundRobin.cpp @@ -90,25 +90,25 @@ c10::intrusive_ptr ProcessGroupRoundRobin::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - TORCH_CHECK(false, "ProcessGroupRoundRobin does not support send"); + throw std::runtime_error("ProcessGroupRoundRobin does not support send"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); + throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recvAnysource( std::vector& /* unused */, int /* unused */) { - TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); + throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::barrier( const BarrierOptions& /* unused */) { - TORCH_CHECK(false, "ProcessGroupRoundRobin does not support barrier"); + throw std::runtime_error("ProcessGroupRoundRobin does not support barrier"); }; const c10::intrusive_ptr& ProcessGroupRoundRobin::next() { @@ -124,7 +124,7 @@ c10::intrusive_ptr ProcessGroupRoundRobin::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - TORCH_CHECK(false, + throw std::runtime_error( "no support for _allgather_base in RoundRobin process group"); } diff --git a/torch/lib/c10d/TCPStore.cpp b/torch/lib/c10d/TCPStore.cpp index 6498f8bcbe633..4958c47b79a71 100644 --- a/torch/lib/c10d/TCPStore.cpp +++ b/torch/lib/c10d/TCPStore.cpp @@ -133,7 +133,7 @@ void BackgroundThread::join() { void BackgroundThread::initStopSignal() { ghStopEvent_ = CreateEvent(NULL, TRUE, FALSE, NULL); if (ghStopEvent_ == NULL) { - TORCH_CHECK(false, + throw std::runtime_error( "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -149,7 +149,7 @@ void BackgroundThread::stop() { #else void BackgroundThread::initStopSignal() { if (pipe(controlPipeFd_.data()) == -1) { - TORCH_CHECK(false, + throw std::runtime_error( "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -336,7 +336,7 @@ void TCPStoreMasterDaemon::query(int socket) { watchHandler(socket); } else { - TORCH_CHECK(false, "Unexpected query type"); + throw std::runtime_error("Unexpected query type"); } } @@ -1126,7 +1126,7 @@ bool TCPStore::check(const std::vector& keys) { if (response == detail::CheckResponseType::NOT_READY) { return false; } - TORCH_CHECK(false, "ready or not_ready response expected"); + throw std::runtime_error("ready or not_ready response expected"); } void TCPStore::wait(const std::vector& keys) { @@ -1156,7 +1156,7 @@ void TCPStore::doWait( auto response = client_->receiveValue(); if (response != detail::WaitResponseType::STOP_WAITING) { - TORCH_CHECK(false, "Stop_waiting response is expected"); + throw std::runtime_error("Stop_waiting response is expected"); } } diff --git a/torch/lib/c10d/UnixSockUtils.hpp b/torch/lib/c10d/UnixSockUtils.hpp index b75bddb763787..fa74be27f889e 100644 --- a/torch/lib/c10d/UnixSockUtils.hpp +++ b/torch/lib/c10d/UnixSockUtils.hpp @@ -56,7 +56,7 @@ inline void waitSocketConnected( throw std::system_error(errno, std::system_category()); } else if (numReady == 0) { errno = 0; - TORCH_CHECK(false, kConnectTimeoutMsg); + throw std::runtime_error(kConnectTimeoutMsg); } socklen_t errLen = sizeof(errno); diff --git a/torch/lib/c10d/Utils.cpp b/torch/lib/c10d/Utils.cpp index 5d9aa744dbacd..f8adc58746c66 100644 --- a/torch/lib/c10d/Utils.cpp +++ b/torch/lib/c10d/Utils.cpp @@ -118,7 +118,7 @@ PortType getSocketPort(int fd) { listenPort = ntohs(addr->sin6_port); } else { - TORCH_CHECK(false, "unsupported protocol"); + throw std::runtime_error("unsupported protocol"); } return listenPort; } @@ -140,7 +140,7 @@ std::string sockaddrToString(struct ::sockaddr* addr) { __output != nullptr) address[INET6_ADDRSTRLEN] = '\0'; } else { - TORCH_CHECK(false, "unsupported protocol"); + throw std::runtime_error("unsupported protocol"); } return address; } @@ -229,7 +229,7 @@ void handleConnectException( if (timeout != kNoTimeout) { const auto elapsed = std::chrono::high_resolution_clock::now() - start; if (elapsed > timeout) { - TORCH_CHECK(false, kConnectTimeoutMsg); + throw std::runtime_error(kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::seconds(1)); @@ -346,7 +346,7 @@ std::tuple accept( while (true) { int res = tcputil::poll(events.get(), 1, timeout.count()); if (res == 0) { - TORCH_CHECK(false, + throw std::runtime_error( "waiting for processes to " "connect has timed out"); } else if (res == -1) { diff --git a/torch/lib/c10d/Utils.hpp b/torch/lib/c10d/Utils.hpp index 5beb5f1c6708b..55edff85606cf 100644 --- a/torch/lib/c10d/Utils.hpp +++ b/torch/lib/c10d/Utils.hpp @@ -92,7 +92,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { try { val = std::stoi(stringValue); } catch (std::exception& e) { - TORCH_CHECK(false, + throw std::runtime_error( "Invalid value for environment variable: " + std::string(envVarName)); } if (val == 1) { @@ -100,7 +100,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { } else if (val == 0) { return false; } else { - TORCH_CHECK(false, + throw std::runtime_error( "Invalid value for environment variable: " + std::string(envVarName)); } } @@ -340,16 +340,16 @@ inline at::Tensor newLikeFlat( std::vector>& tensors, size_t deviceIdx) { if (tensors.size() == 0 || tensors[0].size() == 0) { - TORCH_CHECK(false, "Received an empty list"); + throw std::runtime_error("Received an empty list"); } if (deviceIdx >= tensors.size()) { - TORCH_CHECK(false, "Invalid device index"); + throw std::runtime_error("Invalid device index"); } auto& t = tensors[deviceIdx][0]; auto device = t.device(); for (size_t i = 1; i < tensors[deviceIdx].size(); ++i) { if (tensors[deviceIdx][i].device() != device) { - TORCH_CHECK(false, "Expecting all tensors on the same device"); + throw std::runtime_error("Expecting all tensors on the same device"); } } at::DeviceGuard gpuGuard(device); @@ -363,7 +363,7 @@ inline at::Tensor newLikeFlat( inline at::Tensor newLikeFlat(std::vector& tensors) { if (tensors.size() == 0) { - TORCH_CHECK(false, "Received an empty list"); + throw std::runtime_error("Received an empty list"); } auto& t = tensors[0]; at::DeviceGuard gpuGuard(t.device()); @@ -504,7 +504,7 @@ using SizeType = uint64_t; continue; \ } else if ( \ errno_local == WSAETIMEDOUT || errno_local == WSAEWOULDBLOCK) { \ - TORCH_CHECK(false, "Socket Timeout"); \ + throw std::runtime_error("Socket Timeout"); \ } else { \ throw std::system_error(errno_local, std::system_category()); \ } \ @@ -521,7 +521,7 @@ using SizeType = uint64_t; if (errno == EINTR) { \ continue; \ } else if (errno == EAGAIN || errno == EWOULDBLOCK) { \ - TORCH_CHECK(false, "Socket Timeout"); \ + throw std::runtime_error("Socket Timeout"); \ } else { \ throw std::system_error(errno, std::system_category()); \ } \ diff --git a/torch/lib/c10d/WinSockUtils.hpp b/torch/lib/c10d/WinSockUtils.hpp index 793a0dc7640f2..cd37695845ab1 100644 --- a/torch/lib/c10d/WinSockUtils.hpp +++ b/torch/lib/c10d/WinSockUtils.hpp @@ -46,7 +46,7 @@ inline void waitSocketConnected( std::chrono::high_resolution_clock::now() - startTime; if (elapsed > timeout) { errno = 0; - TORCH_CHECK(false, kConnectTimeoutMsg); + throw std::runtime_error(kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::milliseconds(10)); diff --git a/torch/lib/c10d/frontend.cpp b/torch/lib/c10d/frontend.cpp index b65cba79884af..86a78b6fcebb5 100644 --- a/torch/lib/c10d/frontend.cpp +++ b/torch/lib/c10d/frontend.cpp @@ -146,7 +146,7 @@ c10::intrusive_ptr DistributedC10d::newProcessGroupHelper( pg_name) { return pg_name.second == *group_name; }); if (it != pg_names_.end()) { - TORCH_CHECK(false, + throw std::runtime_error( "The specified group name has already been " "created, please use a different group name"); } diff --git a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp index a158d2c9685df..f3a44cbcad4ae 100644 --- a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp +++ b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp @@ -221,7 +221,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - TORCH_CHECK(false, "future result should be tensor list or none"); + throw std::runtime_error("future result should be tensor list or none"); } } return copyTensors(outputTensors); diff --git a/torch/lib/c10d/test/ProcessGroupMPITest.cpp b/torch/lib/c10d/test/ProcessGroupMPITest.cpp index b8538a016d5b7..bfefbbba2945e 100644 --- a/torch/lib/c10d/test/ProcessGroupMPITest.cpp +++ b/torch/lib/c10d/test/ProcessGroupMPITest.cpp @@ -48,7 +48,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - TORCH_CHECK(false, "future result should be tensor list or none"); + throw std::runtime_error("future result should be tensor list or none"); } } return outputTensors; @@ -80,7 +80,7 @@ void testAllreduce(int iter = 1000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -113,7 +113,7 @@ void testBroadcast(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -143,7 +143,7 @@ void testReduce(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -183,7 +183,7 @@ void testAllgather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -227,7 +227,7 @@ void testGather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -235,7 +235,7 @@ void testGather(int iter = 10000) { } else { for (const auto i : c10::irange(iter)) { if (outputTensors[i].size() != 0) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -277,7 +277,7 @@ void testScatter(int iter = 1) { auto data = outputTensors[i][0].data_ptr(); for (auto k = 0; k < outputTensors[i][0].numel(); ++k) { if (data[k] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -333,13 +333,13 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { // Verify outputs for (const auto i : c10::irange(iter)) { if (recvAnysource && srcRanks[i] != 0) { - TORCH_CHECK(false, "src rank is wrong for recvAnysource"); + throw std::runtime_error("src rank is wrong for recvAnysource"); } const auto expected = i; auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } } @@ -348,7 +348,7 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { void testBackendName() { auto pg = c10d::ProcessGroupMPI::createProcessGroupMPI(); if (pg->getBackendName() != std::string(c10d::MPI_BACKEND_NAME)) { - TORCH_CHECK(false, "BOOM!"); + throw std::runtime_error("BOOM!"); } } diff --git a/torch/lib/c10d/test/TCPStoreTest.cpp b/torch/lib/c10d/test/TCPStoreTest.cpp index 65fb425022b24..e5b7eaf35cc5b 100644 --- a/torch/lib/c10d/test/TCPStoreTest.cpp +++ b/torch/lib/c10d/test/TCPStoreTest.cpp @@ -4,7 +4,6 @@ #include #include #include -#include #include #include @@ -74,7 +73,7 @@ void testHelper(const std::string& prefix = "") { EXPECT_EQ(numKeys, 4); auto timeout = std::chrono::milliseconds(kShortStoreTimeoutMillis); serverStore->setTimeout(timeout); - EXPECT_THROW(serverStore->get("key0"), c10::Error); + EXPECT_THROW(serverStore->get("key0"), std::runtime_error); }); // Hammer on TCPStore @@ -239,7 +238,7 @@ void testWatchKeyCallback(const std::string& prefix = "") { numCallbacksExecutedPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (numCallbacksExecutedFuture.wait_for(span) == std::future_status::timeout) - TORCH_CHECK(false, "Callback execution timed out."); + throw std::runtime_error("Callback execution timed out."); // Check number of callbacks executed equal to number of key change operations // Wait for all callbacks to be triggered @@ -303,7 +302,7 @@ void testKeyChangeHelper( std::future callbackFuture = callbackPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (callbackFuture.wait_for(span) == std::future_status::timeout) - TORCH_CHECK(false, "Callback execution timed out."); + throw std::runtime_error("Callback execution timed out."); // Any exceptions raised from asserts should be rethrown if (eptr) @@ -374,7 +373,7 @@ TEST(TCPStoreTest, testCleanShutdown) { clientTCPStore->get("key"); auto clientThread = std::thread([&clientTCPStore] { - EXPECT_THROW(clientTCPStore->get("invalid_key"), std::system_error); + EXPECT_THROW(clientTCPStore->get("invalid_key"), std::runtime_error); }); // start server shutdown during a client request From e49f0f4ffdd7d349113a0465606cef296ba4a8aa Mon Sep 17 00:00:00 2001 From: Facebook Community Bot Date: Fri, 11 Jun 2021 17:48:06 -0700 Subject: [PATCH 063/305] Automated submodule update: FBGEMM (#59874) Summary: This is an automated pull request to update the first-party submodule for [pytorch/FBGEMM](https://github.com/pytorch/FBGEMM). New submodule commit: https://github.com/pytorch/FBGEMM/commit/ae8ad8fd04eacdcfc5fd979170f0ca08a9e9f0fb Pull Request resolved: https://github.com/pytorch/pytorch/pull/59874 Test Plan: Ensure that CI jobs succeed on GitHub before landing. Reviewed By: jspark1105 Differential Revision: D29064980 fbshipit-source-id: 593f08361817fb771afcf2732f0f647d7c2c72c3 --- third_party/fbgemm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/fbgemm b/third_party/fbgemm index 0520ad5f95db7..ae8ad8fd04eac 160000 --- a/third_party/fbgemm +++ b/third_party/fbgemm @@ -1 +1 @@ -Subproject commit 0520ad5f95db754fbc0ccfb7b563986b6d77bb20 +Subproject commit ae8ad8fd04eacdcfc5fd979170f0ca08a9e9f0fb From 10a3a3d363e05f479ce109347c026e8258a6fb2e Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Fri, 11 Jun 2021 18:18:31 -0700 Subject: [PATCH 064/305] Fix bad change in a CUDACachingAllocator loop (#59903) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59903 D29034650 (https://github.com/pytorch/pytorch/commit/cf0c4ac25811cf93e51b4be6eb58bbdb95963b3b) probably breaks something because it changes a `for` loop on ~Line 1200 from `[size,max)` to `[0,max)`. This fixes that Test Plan: Sandcastle Reviewed By: ngimel Differential Revision: D29081688 fbshipit-source-id: 21f08e3f244fc02cf97d137b3cc80d4378d17185 --- c10/cuda/CUDACachingAllocator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp index 07b0c3746ba70..f7d84d96722ae 100644 --- a/c10/cuda/CUDACachingAllocator.cpp +++ b/c10/cuda/CUDACachingAllocator.cpp @@ -1195,7 +1195,7 @@ class THCCachingAllocator { const auto size = static_cast(device_allocator.size()); if (size < device_count) { device_allocator.resize(device_count); - for (const auto i : c10::irange(device_count)) { + for (const auto i : c10::irange(size, device_count)) { device_allocator[i] = std::unique_ptr( new DeviceCachingAllocator()); } From 1f6e39336f4333c71e58db652bc80d3ec509bfd4 Mon Sep 17 00:00:00 2001 From: lezcano Date: Fri, 11 Jun 2021 19:51:41 -0700 Subject: [PATCH 065/305] Simplify parametrizations.SpectralNorm and improve its initialization (#59564) Summary: Implements a number of changes discussed with soulitzer offline. In particular: - Initialise `u`, `v` in `__init__` rather than in `_update_vectors` - Initialise `u`, `v` to some reasonable vectors by doing 15 power iterations at the start - Simplify the code of `_reshape_weight_to_matrix` (and make it faster) by using `flatten` Pull Request resolved: https://github.com/pytorch/pytorch/pull/59564 Reviewed By: ailzhang Differential Revision: D29066238 Pulled By: soulitzer fbshipit-source-id: 6a58e39ddc7f2bf989ff44fb387ab408d4a1ce3d --- aten/src/ATen/native/LinearAlgebra.cpp | 4 +- test/test_nn.py | 72 +++++++------- torch/nn/utils/parametrizations.py | 125 +++++++++++++++---------- 3 files changed, 117 insertions(+), 84 deletions(-) diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp index 6bce8120e9e39..2ae7b1325d1cb 100644 --- a/aten/src/ATen/native/LinearAlgebra.cpp +++ b/aten/src/ATen/native/LinearAlgebra.cpp @@ -562,7 +562,7 @@ Tensor multi_dot_impl(TensorList _tensors, c10::optional _out) { TORCH_CHECK( false, "multi_dot(): the last tensor must be 1D or 2D but got ", - _tensors[0].dim(), + _tensors[n - 1].dim(), "D"); } @@ -573,7 +573,7 @@ Tensor multi_dot_impl(TensorList _tensors, c10::optional _out) { "multi_dot(): tensor ", i, " must be 2D but got ", - _tensors[0].dim(), + _tensors[i].dim(), "D"); tensors[i] = _tensors[i]; } diff --git a/test/test_nn.py b/test/test_nn.py index 29357bd13ec51..7ce0da8a32f66 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -3730,16 +3730,16 @@ def test_new_spectral_norm(self): m = torch.nn.utils.parametrizations.spectral_norm(m) spectral_norm_m = m.parametrizations.weight[0] - self.assertEqual(spectral_norm_m.u.size(), torch.Size([m.weight.size(0)])) + self.assertEqual(spectral_norm_m._u.size(), torch.Size([m.weight.size(0)])) # .parametrizations.weight.original should be trainable self.assertTrue(hasattr(m.parametrizations.weight, 'original')) self.assertTrue('original' in m.parametrizations.weight._parameters) # u should be just a reused buffer - self.assertTrue(hasattr(spectral_norm_m, 'u')) - self.assertTrue('u' in spectral_norm_m._buffers) - self.assertTrue('v' in spectral_norm_m._buffers) + self.assertTrue(hasattr(spectral_norm_m, '_u')) + self.assertTrue('_u' in spectral_norm_m._buffers) + self.assertTrue('_v' in spectral_norm_m._buffers) # weight should be a plain attribute, not counted as a buffer or a param self.assertIsNotNone(m.weight) @@ -3774,6 +3774,7 @@ def test_new_spectral_norm(self): # Neither weight and bias are parametrized self.assertFalse(hasattr(m, 'parametrizations')) self.assertTrue('weight' in m._parameters) + self.assertFalse(torch.nn.utils.parametrize.is_parametrized(m)) # test correctness in training/eval modes and cpu/multi-gpu settings for apply_dp in (True, False): @@ -3803,16 +3804,24 @@ def get_modules(): m, wrapped_m, spectral_norm_m = get_modules() - self.assertTrue(hasattr(spectral_norm_m, 'u')) - u0 = spectral_norm_m.u.clone() - v0 = spectral_norm_m.v.clone() + self.assertTrue(hasattr(spectral_norm_m, '_u')) + u0 = spectral_norm_m._u.clone() + v0 = spectral_norm_m._v.clone() # TEST TRAINING BEHAVIOR - # run forward again and assert that u and v are updated + # We perform GD first to modify the initial matrix + opt = torch.optim.SGD(wrapped_m.parameters(), lr=0.1) + + opt.zero_grad() + wrapped_m(input).sum().backward() + opt.step() + out = wrapped_m(input) - self.assertNotEqual(u0, spectral_norm_m.u) - self.assertNotEqual(v0, spectral_norm_m.v) + if requires_grad: + # run forward again and assert that u and v are updated + self.assertNotEqual(u0, spectral_norm_m._u) + self.assertNotEqual(v0, spectral_norm_m._v) # assert that backprop reaches original weight # can't use gradcheck because the function changes as we @@ -3823,12 +3832,12 @@ def get_modules(): # test backward works with multiple forwards # it uses training mode so we need to reset `u` and `v` vectors # to same value at beginning for finite difference test to pass - saved_u = spectral_norm_m.u.clone() - saved_v = spectral_norm_m.v.clone() + saved_u = spectral_norm_m._u.clone() + saved_v = spectral_norm_m._v.clone() def fn(input): - spectral_norm_m.u.data.copy_(saved_u) - spectral_norm_m.v.data.copy_(saved_v) + spectral_norm_m._u.data.copy_(saved_u) + spectral_norm_m._v.data.copy_(saved_v) out0 = wrapped_m(input) out1 = wrapped_m(input) return out0 + out1 @@ -3844,13 +3853,6 @@ def fn(input): m = torch.nn.utils.parametrize.remove_parametrizations(m, 'weight') self.assertEqual(wrapped_m(input), pre_remove_out) - torch.nn.utils.parametrizations.spectral_norm(m) - pre_remove_out = wrapped_m(input) - m.train() - self.assertTrue(m.parametrizations.weight[0].training) - m = torch.nn.utils.parametrize.remove_parametrizations(m, 'weight') - self.assertNotEqual(wrapped_m(input), pre_remove_out) - torch.nn.utils.parametrizations.spectral_norm(m) for _ in range(3): pre_remove_out = wrapped_m(input) @@ -3862,8 +3864,8 @@ def fn(input): m, wrapped_m, spectral_norm_m = get_modules() wrapped_m(input) last_train_out = wrapped_m(input) - last_train_u = spectral_norm_m.u.clone() - last_train_v = spectral_norm_m.v.clone() + last_train_u = spectral_norm_m._u.clone() + last_train_v = spectral_norm_m._v.clone() wrapped_m.zero_grad() wrapped_m.eval() @@ -3872,8 +3874,8 @@ def fn(input): self.assertEqual(eval_out0, last_train_out) # assert doing more iteartion in eval don't change things self.assertEqual(eval_out0, wrapped_m(input)) - self.assertEqual(last_train_u, spectral_norm_m.u) - self.assertEqual(last_train_v, spectral_norm_m.v) + self.assertEqual(last_train_u, spectral_norm_m._u) + self.assertEqual(last_train_v, spectral_norm_m._v) # FIXME: the code below is flaky when executed with DataParallel # see https://github.com/pytorch/pytorch/issues/13818 @@ -3884,12 +3886,12 @@ def fn(input): # and eval modes # it uses training mode so we need to reset `u` and `v` vectors # to same value at beginning for finite difference test to pass - saved_u = spectral_norm_m.u.clone() - saved_v = spectral_norm_m.v.clone() + saved_u = spectral_norm_m._u.clone() + saved_v = spectral_norm_m._v.clone() def fn(input): - spectral_norm_m.u.data.copy_(saved_u) - spectral_norm_m.v.data.copy_(saved_v) + spectral_norm_m._u.data.copy_(saved_u) + spectral_norm_m._v.data.copy_(saved_v) wrapped_m.train() out0 = wrapped_m(input) wrapped_m.eval() @@ -3923,8 +3925,8 @@ def test_new_spectral_norm_load_state_dict(self): self.assertEqual({ 'parametrizations.weight.original', 'bias', - 'parametrizations.weight.0.v', - 'parametrizations.weight.0.u' + 'parametrizations.weight.0._v', + 'parametrizations.weight.0._u' }, set(state_dict.keys())) # test that non-strict loading works @@ -3935,9 +3937,9 @@ def test_new_spectral_norm_load_state_dict(self): snm.load_state_dict(non_strict_state_dict, strict=False) del non_strict_state_dict['parametrizations.weight.original'] snm.load_state_dict(non_strict_state_dict, strict=False) - del non_strict_state_dict['parametrizations.weight.0.u'] + del non_strict_state_dict['parametrizations.weight.0._u'] snm.load_state_dict(non_strict_state_dict, strict=False) - del non_strict_state_dict['parametrizations.weight.0.v'] + del non_strict_state_dict['parametrizations.weight.0._v'] snm.load_state_dict(non_strict_state_dict, strict=False) non_strict_state_dict['weight'] = snm.weight.detach().clone() # set W as a buffer snm.load_state_dict(non_strict_state_dict, strict=False) @@ -4094,7 +4096,7 @@ def test_new_spectral_norm_dim(self): # this should not run into incompatible shapes x = m(inp) # check that u refers to the same dimension - self.assertEqual(snm.u.shape, m.parametrizations.weight.original[0, :, 0, 0].shape) + self.assertEqual(snm._u.shape, m.parametrizations.weight.original[0, :, 0, 0].shape) def test_spectral_norm_forward(self): input = torch.randn(3, 5) @@ -4119,7 +4121,7 @@ def test_new_spectral_norm_forward(self): snm = m.parametrizations.weight[0] # naive forward _weight = m.parametrizations.weight.original - _bias, _v = m.bias, snm.v + _bias, _v = m.bias, snm._v _weight_mat = _weight.view(_weight.size(0), -1) _u = torch.mv(_weight_mat, _v) _u = F.normalize(_u, dim=0, eps=1e-12) diff --git a/torch/nn/utils/parametrizations.py b/torch/nn/utils/parametrizations.py index a44b3c79899f5..baf634563ca3b 100644 --- a/torch/nn/utils/parametrizations.py +++ b/torch/nn/utils/parametrizations.py @@ -14,29 +14,43 @@ def __init__( eps: float = 1e-12 ) -> None: super().__init__() - self.dim = dim + ndim = weight.ndim + if dim >= ndim or dim < -ndim: + raise IndexError("Dimension out of range (expected to be in range of " + f"[-{ndim}, {ndim - 1}] but got {dim})") + if n_power_iterations <= 0: raise ValueError('Expected n_power_iterations to be positive, but ' 'got n_power_iterations={}'.format(n_power_iterations)) - self.n_power_iterations = n_power_iterations + self.dim = dim if dim >= 0 else dim + ndim self.eps = eps - self.register_buffer('u', None) - self.register_buffer('v', None) + if ndim > 1: + # For ndim == 1 we do not need to approximate anything (see _SpectralNorm.forward) + self.n_power_iterations = n_power_iterations + weight_mat = self._reshape_weight_to_matrix(weight) + h, w = weight_mat.size() + + u = weight_mat.new_empty(h).normal_(0, 1) + v = weight_mat.new_empty(w).normal_(0, 1) + self.register_buffer('_u', F.normalize(u, dim=0, eps=self.eps)) + self.register_buffer('_v', F.normalize(v, dim=0, eps=self.eps)) - weight_mat = self._reshape_weight_to_matrix(weight) - self._update_vectors(weight_mat) + # Start with u, v initialized to some reasonable values by performing a number + # of iterations of the power method + self._power_method(weight_mat, 15) def _reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor: - weight_mat = weight + # Precondition + assert weight.ndim > 1 + if self.dim != 0: # permute dim to front - weight_mat = weight_mat.permute(self.dim, - *[d for d in range(weight_mat.dim()) if d != self.dim]) - height = weight_mat.size(0) - return weight_mat.reshape(height, -1) + weight = weight.permute(self.dim, *(d for d in range(weight.dim()) if d != self.dim)) + + return weight.flatten(1) @torch.autograd.no_grad() - def _update_vectors(self, weight_mat: torch.Tensor) -> None: + def _power_method(self, weight_mat: torch.Tensor, n_power_iterations: int) -> None: # See original note at torch/nn/utils/spectral_norm.py # NB: If `do_power_iteration` is set, the `u` and `v` vectors are # updated in power iteration **in-place**. This is very important @@ -67,30 +81,34 @@ def _update_vectors(self, weight_mat: torch.Tensor) -> None: # GAN training: loss = D(real) - D(fake). Otherwise, engine will # complain that variables needed to do backward for the first forward # (i.e., the `u` and `v` vectors) are changed in the second forward. - if self.u is None or self.v is None: # type: ignore[has-type] - # randomly initialize `u` and `v` - h, w = weight_mat.size() - self.u = F.normalize(weight_mat.new_empty(h).normal_(0, 1), dim=0, eps=self.eps) - self.v = F.normalize(weight_mat.new_empty(w).normal_(0, 1), dim=0, eps=self.eps) - for _ in range(self.n_power_iterations): + # Precondition + assert weight_mat.ndim > 1 + for _ in range(n_power_iterations): # Spectral norm of weight equals to `u^T W v`, where `u` and `v` # are the first left and right singular vectors. # This power iteration produces approximations of `u` and `v`. - self.u = F.normalize(torch.mv(weight_mat, self.v), - dim=0, eps=self.eps, out=self.u) # type: ignore[has-type] - self.v = F.normalize(torch.mv(weight_mat.t(), self.u), # type: ignore[has-type] - dim=0, eps=self.eps, out=self.v) # type: ignore[has-type] + self._u = F.normalize(torch.mv(weight_mat, self._v), # type: ignore[has-type] + dim=0, eps=self.eps, out=self._u) # type: ignore[has-type] + self._v = F.normalize(torch.mv(weight_mat.t(), self._u), # type: ignore[has-type] + dim=0, eps=self.eps, out=self._v) # type: ignore[has-type] # See above on why we need to clone - self.u = self.u.clone(memory_format=torch.contiguous_format) - self.v = self.v.clone(memory_format=torch.contiguous_format) + self._u = self._u.clone(memory_format=torch.contiguous_format) + self._v = self._v.clone(memory_format=torch.contiguous_format) def forward(self, weight: torch.Tensor) -> torch.Tensor: - weight_mat = self._reshape_weight_to_matrix(weight) - if self.training: - self._update_vectors(weight_mat) - sigma = torch.dot(self.u, torch.mv(weight_mat, self.v)) - return weight / sigma + if weight.ndim == 1: + # Faster and more exact path, no need to approximate anything + return F.normalize(weight, dim=0, eps=self.eps) + else: + weight_mat = self._reshape_weight_to_matrix(weight) + if self.training: + self._power_method(weight_mat, self.n_power_iterations) + # The proper way of computing this should be through F.bilinear, but + # it seems to have some efficiency issues: + # https://github.com/pytorch/pytorch/issues/58093 + sigma = torch.dot(self._u, torch.mv(weight_mat, self._v)) + return weight / sigma def right_inverse(self, value: torch.Tensor) -> torch.Tensor: # we may want to assert here that the passed value already @@ -109,17 +127,41 @@ def spectral_norm(module: Module, \mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})}, \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2} + When applied on a vector, it simplifies to + + .. math:: + \mathbf{x}_{SN} = \dfrac{\mathbf{x}}{\|\mathbf{x}\|_2} + Spectral normalization stabilizes the training of discriminators (critics) - in Generative Adversarial Networks (GANs) by rescaling the weight tensor - with spectral norm :math:`\sigma` of the weight matrix calculated using - power iteration method. If the dimension of the weight tensor is greater - than 2, it is reshaped to 2D in power iteration method to get spectral - norm. + in Generative Adversarial Networks (GANs) by reducing the Lipschitz constant + of the model. :math:`\sigma` is approximated performing one iteration of the + `power method`_ every time the weight is accessed. If the dimension of the + weight tensor is greater than 2, it is reshaped to 2D in power iteration + method to get spectral norm. + See `Spectral Normalization for Generative Adversarial Networks`_ . + .. _`power method`: https://en.wikipedia.org/wiki/Power_iteration .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957 + .. note:: + This function is implemented using the new parametrization functionality + in :func:`torch.nn.utils.parametrize.register_parametrization`. It is a + reimplementation of :func:`torch.nn.utils.spectral_norm`. + + .. note:: + When this constraint is registered, the singular vectors associated to the largest + singular value are estimated rather than sampled at random. These are then updated + performing :attr:`n_power_iterations` of the `power method`_ whenever the tensor + is accessed with the module on `training` mode. + + .. note:: + If the `_SpectralNorm` module, i.e., `module.parametrization.weight[idx]`, + is in training mode on removal, it will perform another power iteration. + If you'd like to avoid this iteration, set the module to eval mode + before its removal. + Args: module (nn.Module): containing module name (str, optional): name of weight parameter @@ -135,17 +177,6 @@ def spectral_norm(module: Module, The original module with a new parametrization registered to the specified weight - .. note:: - This function is implemented using the new parametrization functionality - in :func:`torch.nn.utils.parametrize.register_parametrization`. It is a - reimplementation of :func:`torch.nn.utils.spectral_norm`. - - .. note:: - If the `_SpectralNorm` module, i.e., `module.parametrization.weight[idx]`, - is in training mode on removal, it will perform another power iteration. - If you'd like to avoid this iteration, set the module to eval mode - before its removal. - Example:: >>> snm = spectral_norm(nn.Linear(20, 40)) @@ -158,8 +189,8 @@ def spectral_norm(module: Module, ) ) ) - >>> snm.parametrizations.weight[0].u.size() - torch.Size([40]) + >>> torch.linalg.matrix_norm(snm.weight, 2) + tensor(1.0000, grad_fn=) """ if not hasattr(module, name): raise ValueError( From 7143a6a1895127e9a87d6dbf6637857aaafe92ed Mon Sep 17 00:00:00 2001 From: albanD Date: Sat, 12 Jun 2021 06:55:44 -0700 Subject: [PATCH 066/305] Avoid unnecessary re-computation autograd codegen 21s -> 15s (#59847) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59847 This whole stack does not change anything to the codegened code Test Plan: Imported from OSS Reviewed By: ailzhang Differential Revision: D29063817 Pulled By: albanD fbshipit-source-id: 284c3e057029b7a67f43a1b034bb30863bd68c71 --- tools/autograd/gen_python_functions.py | 10 +-- tools/autograd/load_derivatives.py | 89 +++++++++++++++----------- tools/codegen/gen.py | 48 ++++++++------ tools/pyi/gen_pyi.py | 9 ++- 4 files changed, 88 insertions(+), 68 deletions(-) diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index a59f67bc244d8..9a0eaecc06998 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -157,12 +157,14 @@ def is_py_special_function(f: NativeFunction) -> bool: def gen(out: str, native_yaml_path: str, deprecated_yaml_path: str, template_path: str) -> None: fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + native_functions = parse_native_yaml(native_yaml_path).native_functions + native_functions = list(filter(should_generate_py_binding, native_functions)) - methods = load_signatures(native_yaml_path, deprecated_yaml_path, method=True) + methods = load_signatures(native_functions, deprecated_yaml_path, method=True) create_python_bindings( fm, methods, is_py_variable_method, None, 'python_variable_methods.cpp', method=True) - functions = load_signatures(native_yaml_path, deprecated_yaml_path, method=False) + functions = load_signatures(native_functions, deprecated_yaml_path, method=False) create_python_bindings( fm, functions, is_py_torch_function, 'torch', 'python_torch_functions.cpp', method=False) @@ -211,15 +213,13 @@ def create_python_bindings( }) def load_signatures( - native_yaml_path: str, + native_functions: List[NativeFunction], deprecated_yaml_path: str, *, method: bool, skip_deprecated: bool = False, pyi: bool = False, ) -> Sequence[PythonSignatureNativeFunctionPair]: - native_functions = parse_native_yaml(native_yaml_path).native_functions - native_functions = list(filter(should_generate_py_binding, native_functions)) @with_native_function def gen_signature_pairs(f: NativeFunction) -> PythonSignatureNativeFunctionPair: diff --git a/tools/autograd/load_derivatives.py b/tools/autograd/load_derivatives.py index 59ba25a196e75..77faf020435d9 100644 --- a/tools/autograd/load_derivatives.py +++ b/tools/autograd/load_derivatives.py @@ -25,46 +25,57 @@ except ImportError: from yaml import SafeLoader as Loader # type: ignore[misc] +_GLOBAL_LOAD_DERIVATIVE_CACHE = {} + def load_derivatives(derivatives_yaml_path: str, native_yaml_path: str) -> Sequence[DifferentiabilityInfo]: - with open(derivatives_yaml_path, 'r') as f: - definitions = yaml.load(f, Loader=Loader) - - functions = parse_native_yaml(native_yaml_path).native_functions - - # What's the difference between function schema v.s. signature? - # function schema is the complete declaration including mutability annotation / default value and etc. - # signature is the canonical schema for a group of functions (in-place/out/functional variants) - # that are semantically related. - functions_by_signature: Dict[FunctionSchema, List[NativeFunction]] = defaultdict(list) - functions_by_schema: Dict[str, NativeFunction] = dict() - for function in functions: - functions_by_signature[function.func.signature()].append(function) - assert str(function.func) not in functions_by_schema - functions_by_schema[str(function.func)] = function - - infos = [ - create_differentiability_info(defn, functions_by_signature, functions_by_schema) - for defn in definitions] - - # To keep it byte-for-byte compatible with the old codegen, we assign op names as a separate - # step. We only assign op names to those with differentiable args, and only append suffix to - # duplicated op names. This can be simplified if the first of the duplicates can be named - # 'XyzBackward' instead of 'XyzBackward0' or unconditionally append '0' to singletons. - op_names = create_op_names(infos) - return [ - DifferentiabilityInfo( - name=info.name, - func=info.func, - op=op_name, - derivatives=info.derivatives, - forward_derivatives=info.forward_derivatives, - all_saved_inputs=info.all_saved_inputs, - all_saved_outputs=info.all_saved_outputs, - args_with_derivatives=info.args_with_derivatives, - non_differentiable_arg_names=info.non_differentiable_arg_names, - output_differentiability=info.output_differentiability, - ) - for info, op_name in zip(infos, op_names)] + # Do some caching as this is a deterministic function + global _GLOBAL_LOAD_DERIVATIVE_CACHE + key = (derivatives_yaml_path, native_yaml_path) + if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE: + + with open(derivatives_yaml_path, 'r') as f: + definitions = yaml.load(f, Loader=Loader) + + functions = parse_native_yaml(native_yaml_path).native_functions + + # What's the difference between function schema v.s. signature? + # function schema is the complete declaration including mutability annotation / default value and etc. + # signature is the canonical schema for a group of functions (in-place/out/functional variants) + # that are semantically related. + functions_by_signature: Dict[FunctionSchema, List[NativeFunction]] = defaultdict(list) + functions_by_schema: Dict[str, NativeFunction] = dict() + for function in functions: + functions_by_signature[function.func.signature()].append(function) + assert str(function.func) not in functions_by_schema + functions_by_schema[str(function.func)] = function + + infos = [ + create_differentiability_info(defn, functions_by_signature, functions_by_schema) + for defn in definitions] + + # To keep it byte-for-byte compatible with the old codegen, we assign op names as a separate + # step. We only assign op names to those with differentiable args, and only append suffix to + # duplicated op names. This can be simplified if the first of the duplicates can be named + # 'XyzBackward' instead of 'XyzBackward0' or unconditionally append '0' to singletons. + op_names = create_op_names(infos) + res = [ + DifferentiabilityInfo( + name=info.name, + func=info.func, + op=op_name, + derivatives=info.derivatives, + forward_derivatives=info.forward_derivatives, + all_saved_inputs=info.all_saved_inputs, + all_saved_outputs=info.all_saved_outputs, + args_with_derivatives=info.args_with_derivatives, + non_differentiable_arg_names=info.non_differentiable_arg_names, + output_differentiability=info.output_differentiability, + ) + for info, op_name in zip(infos, op_names)] + + _GLOBAL_LOAD_DERIVATIVE_CACHE[key] = res + + return _GLOBAL_LOAD_DERIVATIVE_CACHE[key] @with_native_function def cpp_arguments(f: NativeFunction) -> Sequence[Binding]: diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index 24f6f9503774b..a414f78b31708 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -77,30 +77,36 @@ def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def] mapping['__line__'] = node.start_mark.line + 1 return mapping +_GLOBAL_PARSE_NATIVE_YAML_CACHE = {} + # Parse native_functions.yaml into a sequence of NativeFunctions and Backend Indices. ParsedYaml = namedtuple('ParsedYaml', ['native_functions', 'backend_indices']) def parse_native_yaml(path: str) -> ParsedYaml: - with open(path, 'r') as f: - es = yaml.load(f, Loader=LineLoader) - assert isinstance(es, list) - rs: List[NativeFunction] = [] - bs: Dict[DispatchKey, Dict[OperatorName, BackendMetadata]] = defaultdict(dict) - for e in es: - assert isinstance(e.get('__line__'), int), e - loc = Location(path, e['__line__']) - funcs = e.get('func') - with context(f'in {loc}:\n {funcs}'): - func, m = NativeFunction.from_yaml(e, loc) - rs.append(func) - BackendIndex.grow_index(bs, m) - error_check_native_functions(rs) - # Default dict is to prevent the codegen from barfing when we have a dispatch key that has no kernels yet. - indices: Dict[DispatchKey, BackendIndex] = defaultdict(lambda: BackendIndex( - dispatch_key=DispatchKey.Undefined, use_out_as_primary=True, external=False, index={})) - for k, v in bs.items(): - # All structured in-tree operators are implemented in terms of their out operator. - indices[k] = BackendIndex(dispatch_key=k, use_out_as_primary=True, external=False, index=v) - return ParsedYaml(rs, indices) + global _GLOBAL_PARSE_NATIVE_YAML_CACHE + if path not in _GLOBAL_PARSE_NATIVE_YAML_CACHE: + with open(path, 'r') as f: + es = yaml.load(f, Loader=LineLoader) + assert isinstance(es, list) + rs: List[NativeFunction] = [] + bs: Dict[DispatchKey, Dict[OperatorName, BackendMetadata]] = defaultdict(dict) + for e in es: + assert isinstance(e.get('__line__'), int), e + loc = Location(path, e['__line__']) + funcs = e.get('func') + with context(f'in {loc}:\n {funcs}'): + func, m = NativeFunction.from_yaml(e, loc) + rs.append(func) + BackendIndex.grow_index(bs, m) + error_check_native_functions(rs) + # Default dict is to prevent the codegen from barfing when we have a dispatch key that has no kernels yet. + indices: Dict[DispatchKey, BackendIndex] = defaultdict(lambda: BackendIndex( + dispatch_key=DispatchKey.Undefined, use_out_as_primary=True, external=False, index={})) + for k, v in bs.items(): + # All structured in-tree operators are implemented in terms of their out operator. + indices[k] = BackendIndex(dispatch_key=k, use_out_as_primary=True, external=False, index=v) + _GLOBAL_PARSE_NATIVE_YAML_CACHE[path] = ParsedYaml(rs, indices) + + return _GLOBAL_PARSE_NATIVE_YAML_CACHE[path] # Some assertions are already performed during parsing, but those are only within a single NativeFunction. # Assertions here are meant to be performed across NativeFunctions. diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py index 0b68b4c5fdcbd..37ba0e33afdfd 100644 --- a/tools/pyi/gen_pyi.py +++ b/tools/pyi/gen_pyi.py @@ -6,7 +6,7 @@ from tools.codegen.model import Variant from tools.codegen.api.python import (PythonSignatureGroup, PythonSignatureNativeFunctionPair) -from tools.codegen.gen import FileManager +from tools.codegen.gen import FileManager, parse_native_yaml from typing import Sequence, List, Dict from ..autograd.gen_python_functions import should_generate_py_binding, load_signatures, group_overloads @@ -376,7 +376,10 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) - ' other: Union[Tensor, Number],' ' *, alpha: Optional[Number]=1, out: Optional[Tensor]=None) -> Tensor: ...'.format(binop)) - function_signatures = load_signatures(native_yaml_path, deprecated_yaml_path, method=False, pyi=True) + native_functions = parse_native_yaml(native_yaml_path).native_functions + native_functions = list(filter(should_generate_py_binding, native_functions)) + + function_signatures = load_signatures(native_functions, deprecated_yaml_path, method=False, pyi=True) sig_groups = get_py_torch_functions(function_signatures) for group in sorted(sig_groups, key=lambda g: g.signature.name): name = group.signature.name @@ -501,7 +504,7 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) - # pyi tensor methods don't currently include deprecated signatures for some reason # TODO: we should probably add them in - tensor_method_signatures = load_signatures(native_yaml_path, deprecated_yaml_path, method=True, skip_deprecated=True, pyi=True) + tensor_method_signatures = load_signatures(native_functions, deprecated_yaml_path, method=True, skip_deprecated=True, pyi=True) tensor_method_sig_groups = get_py_torch_functions(tensor_method_signatures, method=True) for group in sorted(tensor_method_sig_groups, key=lambda g: g.signature.name): From 504ec30109df623ccafd336e31b032825d3235f4 Mon Sep 17 00:00:00 2001 From: albanD Date: Sat, 12 Jun 2021 06:55:44 -0700 Subject: [PATCH 067/305] avoid error string formatting aten codegen 28s -> 23s (#59848) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59848 This whole stack does not change anything to the codegened code Test Plan: Imported from OSS Reviewed By: ailzhang Differential Revision: D29063818 Pulled By: albanD fbshipit-source-id: c68734672eeacd212d7bd9bebe3d53aaa20c3c24 --- tools/codegen/context.py | 2 +- tools/codegen/gen.py | 2 +- tools/codegen/gen_backend_stubs.py | 4 ++-- tools/codegen/utils.py | 3 ++- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tools/codegen/context.py b/tools/codegen/context.py index dfb7c54f4f347..ba21c86c79345 100644 --- a/tools/codegen/context.py +++ b/tools/codegen/context.py @@ -25,7 +25,7 @@ def native_function_manager(g: Union[NativeFunctionsGroup, NativeFunction]) -> I f = g.out else: f = g - with context(f'in native_functions.yaml line {f.loc}:\n {f.func}'): + with context(lambda: f'in native_functions.yaml line {f.loc}:\n {f.func}'): with local.parametrize(use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors): yield diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index a414f78b31708..81a096f9df82d 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -93,7 +93,7 @@ def parse_native_yaml(path: str) -> ParsedYaml: assert isinstance(e.get('__line__'), int), e loc = Location(path, e['__line__']) funcs = e.get('func') - with context(f'in {loc}:\n {funcs}'): + with context(lambda: f'in {loc}:\n {funcs}'): func, m = NativeFunction.from_yaml(e, loc) rs.append(func) BackendIndex.grow_index(bs, m) diff --git a/tools/codegen/gen_backend_stubs.py b/tools/codegen/gen_backend_stubs.py index 2602570b0cd4d..5aaf66f668c2f 100644 --- a/tools/codegen/gen_backend_stubs.py +++ b/tools/codegen/gen_backend_stubs.py @@ -78,7 +78,7 @@ def create_backend_index(backend_ops: List[str], dispatch_key: DispatchKey) -> B backend_key: Optional[DispatchKey] = None if len(supported) > 0: - with context(f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'): + with context(lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'): backend_key = DispatchKey.parse(backend) backend_idx = create_backend_index(supported, backend_key) @@ -87,7 +87,7 @@ def create_backend_index(backend_ops: List[str], dispatch_key: DispatchKey) -> B autograd_key: Optional[DispatchKey] = None if len(supported_autograd) > 0: - with context(f'The "autograd" key was specified, which indicates that you would like to override \ + with context(lambda: f'The "autograd" key was specified, which indicates that you would like to override \ the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.'): autograd_key = DispatchKey.parse(f'Autograd{backend}') diff --git a/tools/codegen/utils.py b/tools/codegen/utils.py index 5a8b88f5a2081..19f058450cfac 100644 --- a/tools/codegen/utils.py +++ b/tools/codegen/utils.py @@ -60,11 +60,12 @@ def concatMap(func: Callable[[T], Sequence[S]], xs: Iterable[T]) -> Iterator[S]: # easily say that an error occurred while processing a specific # context. @contextlib.contextmanager -def context(msg: str) -> Iterator[None]: +def context(msg_fn: Callable[[], str]) -> Iterator[None]: try: yield except Exception as e: # TODO: this does the wrong thing with KeyError + msg = msg_fn() msg = textwrap.indent(msg, ' ') msg = f'{e.args[0]}\n{msg}' if e.args else msg e.args = (msg,) + e.args[1:] From c60d1ac9cf4e2127d1b7f89be355570928379f45 Mon Sep 17 00:00:00 2001 From: albanD Date: Sat, 12 Jun 2021 06:55:44 -0700 Subject: [PATCH 068/305] Use C dumper if possible aten codegen 23s -> 13s (#59849) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59849 This whole stack does not change anything to the codegened code Test Plan: Imported from OSS Reviewed By: ailzhang Differential Revision: D29063815 Pulled By: albanD fbshipit-source-id: c4baa72594bd2fe50ac67f513916f2b2ccb7488c --- tools/codegen/gen.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index 81a096f9df82d..2adce0440c0b4 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -40,6 +40,12 @@ except ImportError: from yaml import SafeLoader as Loader # type: ignore[misc] +try: + # use faster C Dumper if available + from yaml import CSafeDumper as Dumper +except ImportError: + from yaml import SafeDumper as Dumper # type: ignore[misc] + # Welcome to the ATen code generator v2! The ATen code generator is # responsible for parsing native_functions.yaml and then generating # various generated files (e.g., TypeDefault.cpp) based on the operators @@ -531,14 +537,14 @@ def dict_representer(dumper: Any, data: Any) -> Any: return dumper.represent_dict(data.items()) def format_yaml(data: object) -> str: - noalias_dumper = yaml.dumper.SafeDumper + noalias_dumper = Dumper noalias_dumper.ignore_aliases = lambda self, data: True # type: ignore[assignment] # Support serializing OrderedDict noalias_dumper.add_representer(OrderedDict, dict_representer) # type: ignore[no-untyped-call] # Some yaml parsers (e.g. Haskell's) don't understand line breaks. - # width=float('Inf') turns off optional line breaks and improves + # width=1e9 turns off optional line breaks and improves # the portability of the outputted yaml. - return yaml.dump(data, default_flow_style=False, Dumper=noalias_dumper, width=float('Inf')) # type: ignore[no-any-return] + return yaml.dump(data, default_flow_style=False, Dumper=noalias_dumper, width=1e9) # type: ignore[no-any-return] # For some reason, some defaults we write to YAML are written as native # YAML objects, rather than doing them uniformly as strings. This From 30a18fe3189d7c619641cf84a4d55e89834a08ff Mon Sep 17 00:00:00 2001 From: albanD Date: Sat, 12 Jun 2021 06:55:44 -0700 Subject: [PATCH 069/305] refactor yaml loader import, no runtime change (#59850) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59850 This whole stack does not change anything to the codegened code Test Plan: Imported from OSS Reviewed By: ailzhang Differential Revision: D29063816 Pulled By: albanD fbshipit-source-id: ca3067443d8e6282c1077d3dafa3b4f330d43b28 --- tools/autograd/gen_python_functions.py | 10 ++------- tools/autograd/load_derivatives.py | 10 ++------- tools/codegen/gen.py | 30 ++++++++------------------ tools/codegen/gen_backend_stubs.py | 10 ++------- tools/codegen/utils.py | 13 +++++++++++ 5 files changed, 28 insertions(+), 45 deletions(-) diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index 9a0eaecc06998..bf622b4e42196 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -56,16 +56,10 @@ from tools.codegen.context import with_native_function from tools.codegen.model import (Argument, BaseOperatorName, NativeFunction, Type, Variant) -from tools.codegen.utils import split_name_params +from tools.codegen.utils import split_name_params, YamlLoader from typing import Dict, Optional, List, Tuple, Set, Sequence, Callable -try: - # use faster C loader if available - from yaml import CSafeLoader as Loader -except ImportError: - from yaml import SafeLoader as Loader # type: ignore[misc] - # # declarations blocklist # We skip codegen for these functions, for various reasons. @@ -282,7 +276,7 @@ def signature_deprecated(opname: str, params: List[str], call_args: List[str]) - results: List[PythonSignatureNativeFunctionPair] = [] with open(deprecated_yaml_path, 'r') as f: - deprecated_defs = yaml.load(f, Loader=Loader) + deprecated_defs = yaml.load(f, Loader=YamlLoader) for deprecated in deprecated_defs: _, params = split_name_params(deprecated['name']) diff --git a/tools/autograd/load_derivatives.py b/tools/autograd/load_derivatives.py index 77faf020435d9..03eaa1974a29a 100644 --- a/tools/autograd/load_derivatives.py +++ b/tools/autograd/load_derivatives.py @@ -17,13 +17,7 @@ from tools.codegen.gen import parse_native_yaml from tools.codegen.context import with_native_function from tools.codegen.model import FunctionSchema, NativeFunction, Variant, Type, SchemaKind -from tools.codegen.utils import IDENT_REGEX, split_name_params - -try: - # use faster C loader if available - from yaml import CSafeLoader as Loader -except ImportError: - from yaml import SafeLoader as Loader # type: ignore[misc] +from tools.codegen.utils import IDENT_REGEX, split_name_params, YamlLoader _GLOBAL_LOAD_DERIVATIVE_CACHE = {} @@ -34,7 +28,7 @@ def load_derivatives(derivatives_yaml_path: str, native_yaml_path: str) -> Seque if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE: with open(derivatives_yaml_path, 'r') as f: - definitions = yaml.load(f, Loader=Loader) + definitions = yaml.load(f, Loader=YamlLoader) functions = parse_native_yaml(native_yaml_path).native_functions diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index 2adce0440c0b4..f3e4d726d9628 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -27,25 +27,13 @@ import tools.codegen.api.structured as structured from tools.codegen.api.translate import translate from tools.codegen.selective_build.selector import SelectiveBuilder -from tools.codegen.utils import Target, concatMap, context, mapMaybe +from tools.codegen.utils import Target, concatMap, context, mapMaybe, YamlDumper, YamlLoader from tools.codegen.context import (method_with_native_function, native_function_manager, with_native_function_and_indices, with_native_function) import tools.codegen.dest as dest -try: - # use faster C loader if available - from yaml import CSafeLoader as Loader -except ImportError: - from yaml import SafeLoader as Loader # type: ignore[misc] - -try: - # use faster C Dumper if available - from yaml import CSafeDumper as Dumper -except ImportError: - from yaml import SafeDumper as Dumper # type: ignore[misc] - # Welcome to the ATen code generator v2! The ATen code generator is # responsible for parsing native_functions.yaml and then generating # various generated files (e.g., TypeDefault.cpp) based on the operators @@ -76,7 +64,7 @@ # A custom loader for YAML to let us also keep track of line numbers # of each entry in the YAML file -class LineLoader(Loader): +class LineLoader(YamlLoader): def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def] mapping = super().construct_mapping(node, deep=deep) # type: ignore[no-untyped-call] # Add 1 so line numbering starts at 1 @@ -533,18 +521,18 @@ def __call__(self, f: NativeFunction) -> Optional[str]: # # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # -def dict_representer(dumper: Any, data: Any) -> Any: - return dumper.represent_dict(data.items()) - def format_yaml(data: object) -> str: - noalias_dumper = Dumper - noalias_dumper.ignore_aliases = lambda self, data: True # type: ignore[assignment] + # Ignore alias in Dumper + YamlDumper.ignore_aliases = lambda self, data: True # type: ignore[assignment] + # Support serializing OrderedDict - noalias_dumper.add_representer(OrderedDict, dict_representer) # type: ignore[no-untyped-call] + def dict_representer(dumper: Any, data: Any) -> Any: + return dumper.represent_dict(data.items()) + YamlDumper.add_representer(OrderedDict, dict_representer) # type: ignore[no-untyped-call] # Some yaml parsers (e.g. Haskell's) don't understand line breaks. # width=1e9 turns off optional line breaks and improves # the portability of the outputted yaml. - return yaml.dump(data, default_flow_style=False, Dumper=noalias_dumper, width=1e9) # type: ignore[no-any-return] + return yaml.dump(data, default_flow_style=False, Dumper=YamlDumper, width=1e9) # type: ignore[no-any-return] # For some reason, some defaults we write to YAML are written as native # YAML objects, rather than doing them uniformly as strings. This diff --git a/tools/codegen/gen_backend_stubs.py b/tools/codegen/gen_backend_stubs.py index 5aaf66f668c2f..6d3724a61efed 100644 --- a/tools/codegen/gen_backend_stubs.py +++ b/tools/codegen/gen_backend_stubs.py @@ -8,16 +8,10 @@ from tools.codegen.model import (BackendIndex, BackendMetadata, DispatchKey, NativeFunction, NativeFunctionsGroup, OperatorName) from tools.codegen.selective_build.selector import SelectiveBuilder -from tools.codegen.utils import Target, concatMap, context +from tools.codegen.utils import Target, concatMap, context, YamlLoader import tools.codegen.dest as dest import tools.codegen.api.dispatcher as dispatcher -try: - # use faster C loader if available - from yaml import CSafeLoader as Loader -except ImportError: - from yaml import SafeLoader as Loader # type: ignore[misc] - # Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key. # Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping) @@ -35,7 +29,7 @@ def parse_backend_yaml( } with open(backend_yaml_path, 'r') as f: - yaml_values = yaml.load(f, Loader=Loader) + yaml_values = yaml.load(f, Loader=YamlLoader) assert isinstance(yaml_values, dict) valid_keys = ['backend', 'cpp_namespace', 'extra_headers', 'supported', 'autograd'] diff --git a/tools/codegen/utils.py b/tools/codegen/utils.py index 19f058450cfac..dff1dc1f747b2 100644 --- a/tools/codegen/utils.py +++ b/tools/codegen/utils.py @@ -4,6 +4,19 @@ import contextlib import textwrap +# Safely load fast C Yaml loader/dumper if they are available +try: + from yaml import CSafeLoader as Loader +except ImportError: + from yaml import SafeLoader as Loader # type: ignore[misc] +YamlLoader = Loader + +try: + from yaml import CSafeDumper as Dumper +except ImportError: + from yaml import SafeDumper as Dumper # type: ignore[misc] +YamlDumper = Dumper + # Many of these functions share logic for defining both the definition # and declaration (for example, the function signature is the same), so # we organize them into one function that takes a Target to say which From d03ff1a17dcc0fa4dc992ef84a58cbd76f5864d9 Mon Sep 17 00:00:00 2001 From: albanD Date: Sat, 12 Jun 2021 06:55:44 -0700 Subject: [PATCH 070/305] pre compute regex and match simple signature autograd codegen 15s -> 12s (#59852) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59852 This whole stack does not change anything to the codegened code Test Plan: Imported from OSS Reviewed By: ezyang Differential Revision: D29063814 Pulled By: albanD fbshipit-source-id: a751047526f8d58f4760ee6f9ae906675bed5d75 --- tools/autograd/gen_python_functions.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py index bf622b4e42196..874e54d74a69d 100644 --- a/tools/autograd/gen_python_functions.py +++ b/tools/autograd/gen_python_functions.py @@ -69,7 +69,7 @@ # # These functions require manual Python bindings or are not exposed to Python -SKIP_PYTHON_BINDINGS = [ +_SKIP_PYTHON_BINDINGS = [ 'alias', 'contiguous', 'is_cuda', 'is_sparse', 'is_sparse_csr', 'size', 'stride', '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward', '.*_forward_out', '_unsafe_view', 'tensor', '_?sparse_coo_tensor.*', @@ -94,27 +94,31 @@ 'fake_quantize_per_channel_affine_cachemask', ] +SKIP_PYTHON_BINDINGS = list(map(lambda pattern: re.compile(rf'^{pattern}$'), _SKIP_PYTHON_BINDINGS)) + # These function signatures are not exposed to Python. Note that this signature # list does not support regex. SKIP_PYTHON_BINDINGS_SIGNATURES = [ - 'add(Tensor, Scalar, Scalar)', 'add_(Tensor, Scalar, Scalar)', - 'sub(Tensor, Scalar, Scalar)', 'sub_(Tensor, Scalar, Scalar)', - 'mul(Tensor, Scalar)', 'mul_(Tensor, Scalar)', - 'div(Tensor, Scalar)', 'div_(Tensor, Scalar)', + 'add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor', + 'add_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)', + 'sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor', + 'sub_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)', + 'mul.Scalar(Tensor self, Scalar other) -> Tensor', + 'mul_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)', + 'div.Scalar(Tensor self, Scalar other) -> Tensor', + 'div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)', ] @with_native_function def should_generate_py_binding(f: NativeFunction) -> bool: name = cpp.name(f.func) - for pattern in SKIP_PYTHON_BINDINGS: - if re.match('^' + pattern + '$', name): + for skip_regex in SKIP_PYTHON_BINDINGS: + if skip_regex.match(name): return False - args = ', '.join(argument_type_str(arg.type) - for arg in signature(f).arguments()) - sig = f'{name}({args})' + signature = str(f.func) for pattern in SKIP_PYTHON_BINDINGS_SIGNATURES: - if pattern == sig: + if pattern == signature: return False return True From 0ceea7faf4f864ff834617150f1d58d9c1c56b2b Mon Sep 17 00:00:00 2001 From: Victor Quach Date: Sat, 12 Jun 2021 23:20:19 -0700 Subject: [PATCH 071/305] Refactor SavedVariable (#59836) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59836 Preparing for #58500 Test Plan: Imported from OSS Reviewed By: ailzhang Differential Revision: D29069159 fbshipit-source-id: dd4d870c8ae10a4bd7f12be127e093f60fa072fa --- torch/csrc/autograd/saved_variable.cpp | 45 ++++++++++++++------------ 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/torch/csrc/autograd/saved_variable.cpp b/torch/csrc/autograd/saved_variable.cpp index a667c6040debe..bbc6afe6750c3 100644 --- a/torch/csrc/autograd/saved_variable.cpp +++ b/torch/csrc/autograd/saved_variable.cpp @@ -37,28 +37,34 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i "you can make a clone to get a normal tensor and use it in autograd.") was_default_constructed_ = false; + is_inplace_view_ = is_inplace_view; + version_counter_ = impl::version_counter(variable); + saved_version_ = version_counter_.current_version(); + output_nr_ = variable.output_nr(); requires_grad_ = variable.requires_grad(); has_grad_fn_ = !variable.is_leaf(); - is_inplace_view_ = is_inplace_view; + // These copies are all shared_ptr copies, so slightly more expensive. // Do them here instead of in the init list in case data is undefined. data_ = variable.tensor_data(); - // TODO(albanD) This needs to be updated when moving to multiple levels - const auto& fw_grad = variable._fw_grad(/* level */ 0); - if (fw_grad.defined()) { - fw_grad_ = std::make_shared(); - fw_grad_->set_value(fw_grad, /* level */ 0); - } + if (variable.is_leaf()) { grad_accumulator_ = impl::grad_accumulator(variable); } else if (!is_output) { grad_fn_ = variable.grad_fn(); - } else if (is_inplace_view) { + } + + if(is_output && is_inplace_view) { weak_grad_fn_ = variable.grad_fn(); } - version_counter_ = impl::version_counter(variable); - saved_version_ = version_counter_.current_version(); + + // TODO(albanD) This needs to be updated when moving to multiple levels + const auto& fw_grad = variable._fw_grad(/* level */ 0); + if (fw_grad.defined()) { + fw_grad_ = std::make_shared(); + fw_grad_->set_value(fw_grad, /* level */ 0); + } } } @@ -67,19 +73,15 @@ SavedVariable::SavedVariable(const c10::optional& variable, bool is_ou Variable SavedVariable::unpack(std::shared_ptr saved_for) const { if (!data_.defined()) { - if (!was_default_constructed_) { - throw std::runtime_error(ERR_BACKWARD_TWICE); - } + TORCH_CHECK(was_default_constructed_, ERR_BACKWARD_TWICE); return Variable(); } + // We want grad_fn here to provide the most hlepful debug message to the user + // if versions don't match auto grad_fn = is_inplace_view_ ? weak_grad_fn_.lock() : grad_fn_; if (has_grad_fn_ && !grad_fn) { - if (!saved_for) { - // If saving the grad_fn would create a circular reference, then it must - // be passed in to the unpack function. - throw std::runtime_error("No grad_fn for non-leaf saved variable"); - } + TORCH_CHECK(saved_for,"No grad_fn for non-leaf saved variable"); grad_fn = std::move(saved_for); } @@ -104,7 +106,7 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { "that failed to compute its gradient. The variable in question " "was changed in there or anywhere later. Good luck!"; } - throw std::runtime_error(message.str()); + TORCH_CHECK(false, message.str()); } // NB: saved views are unpacked as normal Variables (not views) even though @@ -122,8 +124,9 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { // should have saved the grad accumulator. Even if the Variable no longer // alive, the accumulator should be kept alive by the references in the // graph). - if (requires_grad_ && !var.grad_fn() && grad_accumulator_.expired()) - throw std::logic_error("No grad accumulator for a saved leaf!"); + if (requires_grad_ && !var.grad_fn() && grad_accumulator_.expired()) { + TORCH_CHECK(false, "No grad accumulator for a saved leaf!"); + } impl::set_grad_accumulator(var, grad_accumulator_); // NB: var here is never a view so there is no need to make anything special From 92513038e8bf59c512e4e5b7d95b3fe14e5c33a6 Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Sun, 13 Jun 2021 02:32:23 -0700 Subject: [PATCH 072/305] Revert D28994140: [pytorch][PR] Implemented torch.cov Test Plan: revert-hammer Differential Revision: D28994140 (https://github.com/pytorch/pytorch/commit/23c232554bfc138d0d082ea8ea8e8ecec51e30ae) Original commit changeset: 1890166c0a9c fbshipit-source-id: 73dfe1b00464e38f004f99960cdeeb604ed4b20a --- aten/src/ATen/core/aten_interned_strings.h | 1 - aten/src/ATen/native/Correlation.cpp | 109 ------------------ aten/src/ATen/native/native_functions.yaml | 3 - docs/source/tensors.rst | 1 - docs/source/torch.rst | 1 - test/test_torch.py | 47 -------- tools/build_variables.bzl | 1 - torch/_tensor_docs.py | 6 - torch/_torch_docs.py | 69 ----------- torch/overrides.py | 1 - .../_internal/common_methods_invocations.py | 23 ---- 11 files changed, 262 deletions(-) delete mode 100644 aten/src/ATen/native/Correlation.cpp diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 3bb9f66d0c958..1a7486a019a06 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -260,7 +260,6 @@ _(aten, cosine_embedding_loss) \ _(aten, cosine_similarity) \ _(aten, count_nonzero) \ _(aten, cross) \ -_(aten, cov) \ _(aten, std_mean) \ _(aten, var_mean) \ _(aten, ctc_loss) \ diff --git a/aten/src/ATen/native/Correlation.cpp b/aten/src/ATen/native/Correlation.cpp deleted file mode 100644 index 9e1f67b915dd2..0000000000000 --- a/aten/src/ATen/native/Correlation.cpp +++ /dev/null @@ -1,109 +0,0 @@ -#include -#include - -namespace at { -namespace native { - -Tensor cov( - const Tensor& self, - int64_t correction, - const c10::optional& fweights, - const c10::optional& aweights) { - constexpr int64_t OBSERVATIONS_DIM = 1; - - TORCH_CHECK( - self.ndimension() <= 2, - "cov(): expected input to have two or fewer dimensions but got an input with ", - self.ndimension(), - " dimensions"); - - TORCH_CHECK( - self.scalar_type() != kBool, "cov(): bool dtype is not supported for input"); - - // View input tensor as 2D (variables, observations) - auto in = self.ndimension() < 2 ? self.view({1, -1}) : self; - const auto num_observations = in.size(OBSERVATIONS_DIM); - - // The product of frequencies (fweights) and weights (aweights). - Tensor w; - - if (fweights.has_value()) { - w = fweights.value(); - TORCH_CHECK( - w.ndimension() <= 1, - "cov(): expected fweights to have one or fewer dimensions but got fweights with ", - w.ndimension(), - " dimensions"); - TORCH_CHECK( - at::isIntegralType(w.scalar_type(), false), - "cov(): expected fweights to have integral dtype but got fweights with ", - w.scalar_type(), - " dtype"); - TORCH_CHECK( - w.numel() == num_observations, - "cov(): expected fweights to have the same numel as there are observations in the input but got ", - w.numel(), - " != ", - num_observations); - TORCH_CHECK( - num_observations == 0 || w.min().ge(0).item(), - "cov(): fweights cannot be negative"); - } - - if (aweights.has_value()) { - const auto& aw = aweights.value(); - TORCH_CHECK( - aw.ndimension() <= 1, - "cov(): expected aweights to have one or fewer dimensions but got aweights with ", - aw.ndimension(), - " dimensions"); - TORCH_CHECK( - at::isFloatingType(aw.scalar_type()), - "cov(): expected aweights to have floating point dtype but got aweights with ", - aw.scalar_type(), - " dtype"); - TORCH_CHECK( - aw.numel() == num_observations, - "cov(): expected aweights to have the same numel as there are observations in the input but got ", - aw.numel(), - " != ", - num_observations); - TORCH_CHECK( - num_observations == 0 || aw.min().ge(0).item(), - "cov(): aweights cannot be negative"); - w = w.defined() ? w * aw : aw; - } - - // Compute a weighted average of the observations - const auto w_sum = w.defined() - ? w.sum() - : at::scalar_tensor(num_observations, in.options().dtype(kLong)); - - TORCH_CHECK( - !w.defined() || w_sum.ne(0).item(), - "cov(): weights sum to zero, can't be normalized"); - - const auto avg = (w.defined() ? in * w : in).sum(OBSERVATIONS_DIM) / w_sum; - - // Compute the normalization factor - Tensor norm_factor; - - if (w.defined() && aweights.has_value() && correction != 0) { - norm_factor = w_sum - correction * (w * aweights.value()).sum() / w_sum; - } else { - norm_factor = w_sum - correction; - } - - if (norm_factor.le(0).item()) { - TORCH_WARN("cov(): degrees of freedom is <= 0"); - norm_factor.zero_(); - } - - // Compute covariance matrix - in = in - avg.unsqueeze(1); - const auto c = at::mm(in, (w.defined() ? in * w : in).t().conj()); - return at::true_divide(c, norm_factor).squeeze(); -} - -} // namespace native -} // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index f2e4150b67429..5e0dd9917dd9f 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -1277,9 +1277,6 @@ dispatch: CompositeExplicitAutograd: count_nonzero -- func: cov(Tensor self, *, int correction=1, Tensor? fweights=None, Tensor? aweights=None) -> Tensor - variants: function, method - - func: cudnn_affine_grid_generator(Tensor theta, int N, int C, int H, int W) -> Tensor grid dispatch: CUDA: cudnn_affine_grid_generator_forward diff --git a/docs/source/tensors.rst b/docs/source/tensors.rst index 36686b98e54c0..dc210c78ce6d0 100644 --- a/docs/source/tensors.rst +++ b/docs/source/tensors.rst @@ -286,7 +286,6 @@ Tensor class reference Tensor.cosh Tensor.cosh_ Tensor.count_nonzero - Tensor.cov Tensor.acosh Tensor.acosh_ Tensor.arccosh diff --git a/docs/source/torch.rst b/docs/source/torch.rst index ec9ff083514fa..94b288920ca6b 100644 --- a/docs/source/torch.rst +++ b/docs/source/torch.rst @@ -477,7 +477,6 @@ Other Operations cdist clone combinations - cov cross cummax cummin diff --git a/test/test_torch.py b/test/test_torch.py index 30ad853b30c42..b30697e099443 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -4324,53 +4324,6 @@ def test_exponential_no_zero(self, device, dtype): x = torch.empty(50000000, device=device, dtype=dtype).exponential_() self.assertTrue(x.min() > 0) - @dtypes(torch.float, torch.cfloat) - def test_cov(self, device, dtype): - def check(t, correction=1, fweights=None, aweights=None): - actual = torch.cov(t, correction=correction, fweights=fweights, aweights=aweights) - t = t.cpu().numpy() - fweights = fweights.cpu().numpy() if fweights is not None else None - aweights = aweights.cpu().numpy() if aweights is not None else None - expected = np.cov(t, ddof=correction, fweights=fweights, aweights=aweights) - expected = torch.from_numpy(np.array(expected)).to(dtype=actual.dtype) - self.assertEqual(actual, expected, atol=1e-05, rtol=1e-05) - - def generate_input_tensors(): - yield make_tensor((0, 0), device, dtype) - yield make_tensor((1, 0), device, dtype) - yield make_tensor((0, 1), device, dtype) - yield make_tensor((2), device, dtype) - yield make_tensor((2, 1), device, dtype) - yield make_tensor((2, 2), device, dtype) - yield make_tensor((2, 3), device, dtype) - yield make_tensor((5, 10), device, dtype) - yield make_tensor((5, 10), device, dtype, noncontiguous=True) - yield torch.tensor([0, -2, nan, 10.2, inf], dtype=dtype, device=device) - - for t in generate_input_tensors(): - check(t) - num_observations = t.numel() if t.ndim < 2 else t.size(1) - if num_observations > 0: - fweights = torch.randint(1, 10, (num_observations,), device=device) - aweights = make_tensor((num_observations,), device, torch.float, low=1) - for correction, fw, aw in product([0, 1, 2], [None, fweights], [None, aweights]): - check(t, correction, fweights, aweights) - - def test_cov_error(self, device): - def check(msg, *args, **kwargs): - with self.assertRaisesRegex(RuntimeError, r'cov\(\):.*' + msg + r'.*'): - torch.cov(*args, **kwargs) - - a = torch.rand(2) - check(r'expected input to have two or fewer dimensions', torch.rand(2, 2, 2)) - check(r'expected fweights to have one or fewer dimensions', a, fweights=torch.rand(2, 2)) - check(r'expected aweights to have one or fewer dimensions', a, aweights=torch.rand(2, 2)) - check(r'expected fweights to have integral dtype', a, fweights=torch.rand(2)) - check(r'expected aweights to have floating point dtype', a, aweights=torch.tensor([1, 1])) - check(r'expected fweights to have the same numel', a, fweights=torch.tensor([1])) - check(r'expected aweights to have the same numel', a, aweights=torch.rand(1)) - check(r'fweights cannot be negative', a, fweights=torch.tensor([-1, -2])) - check(r'aweights cannot be negative', a, aweights=torch.tensor([-1., -2.])) @skipIfNoSciPy @dtypes(*torch.testing.get_all_fp_dtypes()) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index ce043f2ebcf86..8ae8c3805381d 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -931,7 +931,6 @@ aten_native_source_non_codegen_list = [ "aten/src/ATen/native/ConvolutionMM3d.cpp", "aten/src/ATen/native/ConvolutionTBC.cpp", "aten/src/ATen/native/Copy.cpp", - "aten/src/ATen/native/Correlation.cpp", "aten/src/ATen/native/Cross.cpp", "aten/src/ATen/native/DilatedMaxPool2d.cpp", "aten/src/ATen/native/DilatedMaxPool3d.cpp", diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py index 4d3a8cafe9d3f..a6b0c0ef6c454 100644 --- a/torch/_tensor_docs.py +++ b/torch/_tensor_docs.py @@ -976,12 +976,6 @@ def add_docstr_all(method, docstr): See :func:`torch.count_nonzero` """) -add_docstr_all('cov', r""" -cov(*, correction=1, fweights=None, aweights=None) -> Tensor - -See :func:`torch.cov` -""") - add_docstr_all('cross', r""" cross(other, dim=-1) -> Tensor diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index 579041087a3ee..4e38a487f0e78 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -1638,75 +1638,6 @@ def merge_dicts(*dicts): False """) -add_docstr(torch.cov, r""" -cov(input, *, correction=1, fweights=None, aweights=None) -> Tensor - -Estimates the covariance matrix of the variables given by the :attr:`input` matrix, where rows are -the variables and columns are the observations. - -A covariance matrix is a square matrix giving the covariance of each pair of variables. The diagonal contains -the variance of each variable (covariance of a variable with itself). By definition, if :attr:`input` represents -a single variable (Scalar or 1D) then its variance is returned. - -The unbiased sample covariance of the variables :math:`x` and :math:`y` is given by: - -.. math:: - \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}(x_{i} - \bar{x})(y_{i} - \bar{y})}{N~-~1} - -where :math:`\bar{x}` and :math:`\bar{y}` are the simple means of the :math:`x` and :math:`y` respectively. - -If :attr:`fweights` and/or :attr:`aweights` are provided, the unbiased weighted covariance -is calculated, which is given by: - -.. math:: - \text{cov}_w(x,y) = \frac{\sum^{N}_{i = 1}w_i(x_{i} - \mu_x^*)(y_{i} - \mu_y^*)}{\sum^{N}_{i = 1}w_i~-~1} - -where :math:`w` denotes :attr:`fweights` or :attr:`aweights` based on whichever is provided, or -:math:`w = fweights \times aweights` if both are provided, and -:math:`\mu_x^* = \frac{\sum^{N}_{i = 1}w_ix_{i} }{\sum^{N}_{i = 1}w_i}` is the weighted mean of the variable. - -Args: - input (Tensor): A 2D matrix containing multiple variables and observations, or a - Scalar or 1D vector representing a single variable. - -Keyword Args: - correction (int, optional): difference between the sample size and sample degrees of freedom. - Defaults to Bessel's correction, ``correction = 1`` which returns the unbiased estimate, - even if both :attr:`fweights` and :attr:`aweights` are specified. ``correction = 0`` - will return the simple average. Defaults to ``1``. - fweights (tensor, optional): A Scalar or 1D tensor of observation vector frequencies representing the number of - times each observation should be repeated. Its numel must equal the number of columns of :attr:`input`. - Must have integral dtype. Ignored if ``None``. `Defaults to ``None``. - aweights (tensor, optional): A Scalar or 1D array of observation vector weights. - These relative weights are typically large for observations considered “important” and smaller for - observations considered less “important”. Its numel must equal the number of columns of :attr:`input`. - Must have floating point dtype. Ignored if ``None``. `Defaults to ``None``. - -Returns: - (Tensor) The covariance matrix of the variables. - -Example:: - >>> x = torch.tensor([[0, 2], [1, 1], [2, 0]]).T - >>> x - tensor([[0, 1, 2], - [2, 1, 0]]) - >>> torch.cov(x) - tensor([[ 1., -1.], - [-1., 1.]]) - >>> torch.cov(x, correction=0) - tensor([[ 0.6667, -0.6667], - [-0.6667, 0.6667]]) - >>> fw = torch.randint(1, 10, (3,)) - >>> fw - tensor([1, 6, 9]) - >>> aw = torch.rand(3) - >>> aw - tensor([0.4282, 0.0255, 0.4144]) - >>> torch.cov(x, fweights=fw, aweights=aw) - tensor([[ 0.4169, -0.4169], - [-0.4169, 0.4169]]) -""") - add_docstr(torch.cat, r""" cat(tensors, dim=0, *, out=None) -> Tensor diff --git a/torch/overrides.py b/torch/overrides.py index 18c9cca37ceea..75bde5decb787 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -344,7 +344,6 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.clamp_min: lambda input, min, out=None: -1, torch.clamp_max: lambda input, max, out=None: -1, torch.column_stack: lambda tensors, out=None: -1, - torch.cov: lambda input, correction=1, fweights=None, aweights=None: -1, torch.clone: lambda input: -1, torch.combinations: lambda input, r=2, with_replacement=False: -1, torch.complex: lambda real, imag: -1, diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 2881b5e7b63d7..1cda7f822db50 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -2677,22 +2677,6 @@ def sample_inputs_std_var(op_info, device, dtype, requires_grad, **kwargs): ] -def sample_inputs_cov(op_info, device, dtype, requires_grad, **kwargs): - shapes = [(2,), (1, 2), (3, 2), (2, 3)] - - inputs = [] - for shape in shapes: - t = make_tensor(shape, device, dtype, requires_grad=requires_grad) - inputs.append(SampleInput(t)) - num_observations = t.numel() if t.ndimension() < 2 else t.size(1) - fweights = make_tensor((num_observations,), device, torch.int, low=0, high=10, requires_grad=requires_grad) - aweights = make_tensor((num_observations,), device, torch.float, low=0, high=1, requires_grad=requires_grad) - for correction, fw, aw in product(range(num_observations), [None, fweights], [None, aweights]): - inputs.append(SampleInput(t, kwargs={'correction': correction, 'fweights': fw, 'aweights': aw})) - - return inputs - - def _sample_inputs_svd(op_info, device, dtype, requires_grad=False, is_linalg_svd=False): """ This function generates input for torch.svd with distinct singular values so that autograd is always stable. @@ -4923,13 +4907,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), )), - OpInfo('cov', - dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), - dtypesIfCUDA=all_types_and_complex_and(torch.half, *[torch.bfloat16] if CUDA11OrLater else []), - sample_inputs_func=sample_inputs_cov, - supports_out=False, - # JIT test not working for tensor kwargs (https://github.com/pytorch/pytorch/issues/58507) - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit'),)), OpInfo('cross', dtypes=all_types_and_complex(), dtypesIfCUDA=all_types_and(torch.half), From be038d89898d0d2111b8acedefd08ceed62664cb Mon Sep 17 00:00:00 2001 From: Michael Carilli Date: Sun, 13 Jun 2021 12:08:43 -0700 Subject: [PATCH 073/305] [CUDA graphs] Make stream semantics of backward calls consistent with other cuda ops (ci-all edition) (#57833) Summary: ci-all resubmit of https://github.com/pytorch/pytorch/pull/54227. Tests look good except for a few distributed autograd failures (pytorch_linux_xenial_cuda10_2_cudnn7_py3_multigpu_test) and rocm failures (pr/pytorch-linux-bionic-rocm4.1-py3.6). The common denominator in rocm failures appears to be multi-gpu activity: some [multiprocess DDP failures](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-linux-bionic-rocm4.1-py3.6-test1/8115/console), some [single-process failures](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-linux-bionic-rocm4.1-py3.6-test2/8115/console) where the single process has autograd ops that span devices. jeffdaily jithunnair-amd sunway513, could one of you take a look? The streaming backward change is also beneficial to rocm, I expect. For debugging rocm failures, I think we should ignore the multiprocess/DDP tests and focus on the single process cases. The root cause is probably the same and the single process cases are simpler. ---------------------------------- Update: Rocm failures are due to https://github.com/pytorch/pytorch/issues/59750. https://github.com/pytorch/pytorch/pull/57833/commits/2718a54032d0791ce90a9a95d15150c53727713e is a workaround, to be updated once https://github.com/pytorch/pytorch/issues/59750 is fixed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/57833 Reviewed By: mruberry Differential Revision: D28942391 Pulled By: ngimel fbshipit-source-id: d6047e971c5f1c6386334bf3641402a92f12e2f8 --- docs/source/notes/cuda.rst | 54 +++++-- test/test_autograd.py | 27 ++++ test/test_cuda.py | 146 ++++++++++++++++-- torch/csrc/autograd/engine.cpp | 144 +++++++++++++---- torch/csrc/autograd/engine.h | 8 + .../autograd/engine/dist_engine.cpp | 13 ++ torch/lib/c10d/reducer.cpp | 19 --- 7 files changed, 338 insertions(+), 73 deletions(-) diff --git a/docs/source/notes/cuda.rst b/docs/source/notes/cuda.rst index 34ee143a77d57..41b6b1c9257a7 100644 --- a/docs/source/notes/cuda.rst +++ b/docs/source/notes/cuda.rst @@ -201,36 +201,66 @@ ensure proper synchronization. Stream semantics of backward passes ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Internally, each backward CUDA op runs on the same stream that was used for its corresponding forward op. - -When manually supplying CUDA tensor(s) as a backward pass's initial gradient(s) (e.g., +A. Each backward CUDA op runs on the same stream that was used for its corresponding forward op. + +B. The stream semantics of a backward call with respect to surrounding ops are the same +as for any other call. More concretely, when calling +:func:`autograd.backward`, +:func:`autograd.grad`, or +:meth:`tensor.backward`, +and optionally supplying CUDA tensor(s) as the initial gradient(s) (e.g., :func:`autograd.backward(..., grad_tensors=initial_grads)`, :func:`autograd.grad(..., grad_outputs=initial_grads)`, or :meth:`tensor.backward(..., gradient=initial_grad)`), the acts of -1. populating the initial gradient(s) and -2. invoking the backward pass +1. optionally populating initial gradient(s), +2. invoking the backward pass, and +3. using the gradients + +have the same stream-semantics relationship as any group of ops:: + + s = torch.cuda.Stream() + + # Safe, grads are used in the same stream context as backward() + with torch.cuda.stream(s): + loss.backward() + use grads + + # Unsafe + with torch.cuda.stream(s): + loss.backward() + use grads -have the same stream-semantics relationship as any pair of ops:: + # Safe, with synchronization + with torch.cuda.stream(s): + loss.backward() + torch.cuda.current_stream().wait_stream(s) + use grads - # Safe, populating initial_grad and invoking backward are in the same stream context - with torch.cuda.stream(strm): + # Safe, populating initial grad and invoking backward are in the same stream context + with torch.cuda.stream(s): loss.backward(gradient=torch.ones_like(loss)) # Unsafe, populating initial_grad and invoking backward are in different stream contexts, # without synchronization initial_grad = torch.ones_like(loss) - with torch.cuda.stream(strm): + with torch.cuda.stream(s): loss.backward(gradient=initial_grad) # Safe, with synchronization initial_grad = torch.ones_like(loss) - strm.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(strm): - initial_grad.record_stream(strm) + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + initial_grad.record_stream(s) loss.backward(gradient=initial_grad) +If your forward pass runs some independent ops in parallel on different streams, +A. helps the backward pass exploit that same parallelism. + +The backward call inserts internal syncs as needed to ensure B. holds true even if A. +makes some backward ops run on assorted side streams. + .. _CUDA stream: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#streams .. _cuda-memory-management: diff --git a/test/test_autograd.py b/test/test_autograd.py index 31e8b7d6faab2..8fcf42cac1796 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -4512,6 +4512,33 @@ def test_checkpoint_valid_reset_on_error(self): c = checkpoint(torch.exp, a).sum() c.backward() + def test_callback_adds_callback(self): + called = [0] + + def callback_final(): + called[0] += 1 + + def callback_adds_callback(): + called[0] += 1 + Variable._execution_engine.queue_callback(callback_final) + + class MyFunc(Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, grad): + Variable._execution_engine.queue_callback(callback_adds_callback) + return grad + + a = torch.rand((3, 3), requires_grad=True) + b = MyFunc.apply(a) + b.sum().backward() + + self.assertEqual(called[0], 2) + def _test_reentrant_with_callbacks(self, install_callbacks_in_depths): counter = {} counter["inner"] = 0 diff --git a/test/test_cuda.py b/test/test_cuda.py index 8eae7796898fa..328aa1ad5a132 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -1672,34 +1672,67 @@ def test_cuda_kernel_loop_overflow_large(self): torch.cuda.synchronize() self.assertEqual(y[0, 0, 0, 2**31 - 2], expected) - @skipCUDANonDefaultStreamIf(True) - def test_streaming_backwards_sync(self): - default_stream = torch.cuda.current_stream() - stream = torch.cuda.Stream() - + # this might create a reference cycle on self... + def _make_multiply_in_stream(self): class MultiplyInStream(torch.autograd.Function): @staticmethod - def forward(ctx, x): - return x * 2 + def forward(ctx, x, val): + ctx.val = val + ctx.stream = torch.cuda.current_stream() + return x * val @staticmethod def backward(ctx, grad): - self.assertEqual(torch.cuda.current_stream(), stream) + self.assertEqual(torch.cuda.current_stream(), ctx.stream) # delays the operation in the the background stream - torch.cuda._sleep(1000 * 1000) - return grad * 2 + torch.cuda._sleep(1000 * 5000) + return grad * ctx.val, None + + return MultiplyInStream + @skipCUDANonDefaultStreamIf(True) + def test_streaming_backwards_sync(self): + default_stream = torch.cuda.current_stream() + stream = torch.cuda.Stream() + + MultiplyInStream = self._make_multiply_in_stream() + + # Tests using grads outside the backward() stream context + # See "Stream semantics of backward passes" on https://pytorch.org/docs/stable/notes/cuda.html x = torch.randn(5, 5, device='cuda', requires_grad=True) with torch.cuda.stream(stream): stream.wait_stream(default_stream) - output = MultiplyInStream.apply(x) + output = MultiplyInStream.apply(x, 2) output.sum().backward() - + # sync needed + default_stream.wait_stream(stream) self.assertEqual(x.grad, torch.ones_like(x) * 2) self.assertEqual(torch.cuda.current_stream(), default_stream) - def test_streaming_backwards_multiple_streams(self): + # Tests that using grads in the same stream context as backward() + # is safe regardless what streams bwd ops ran on + bwd_ambient_stream = torch.cuda.Stream() + x = torch.randn(5, 5, device='cuda', requires_grad=True) + with torch.cuda.stream(stream): + stream.wait_stream(default_stream) + output = MultiplyInStream.apply(x, 3) + with torch.cuda.stream(bwd_ambient_stream): + bwd_ambient_stream.wait_stream(stream) + output.sum().backward() + # x was first used on "stream" so its AccumulateGrad leaf should run on "stream". + # The end of backward() should have synced "bwd_ambient_stream" with "stream" + # so it should be safe to use x.grad here without any syncs. + self.assertEqual(x.grad, torch.ones_like(x) * 3) + self.assertEqual(torch.cuda.current_stream(), bwd_ambient_stream) + # Skip the test for ROCm as per https://github.com/pytorch/pytorch/issues/53190 + @skipIfRocm + def test_streaming_backwards_multiple_streams_legacy(self): + # Tests calling backward() under a side stream then using a grad + # on the default stream without syncing. Right now, this pattern is safe, + # but only for BC. In a future PR, this pattern will become unsafe, + # a sync will be required, and this test will be deleted in favor of + # test_streaming_backward_multiple_streams below. class StreamModel(torch.nn.Module): def __init__(self): super(StreamModel, self).__init__() @@ -1732,6 +1765,49 @@ def accum_hook(grad): self.assertEqual(x.grad, torch.ones_like(x) * 5) + # Skip the test for ROCm as per https://github.com/pytorch/pytorch/issues/53190 + @skipIfRocm + def test_streaming_backwards_multiple_streams(self): + MultiplyInStream = self._make_multiply_in_stream() + + class StreamModel(torch.nn.Module): + def __init__(self): + super(StreamModel, self).__init__() + self.event = torch.cuda.Event() + self.stream0 = torch.cuda.Stream() + self.stream1 = torch.cuda.Stream() + + def forward(self, x, x_first_use_on_ambient): + if x_first_use_on_ambient: + x0 = x.clone() + self.stream0.wait_stream(torch.cuda.current_stream()) + self.stream1.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.stream0): + if not x_first_use_on_ambient: + x0 = x.clone() + y0 = MultiplyInStream.apply(x0, 2) + self.event.record(stream=torch.cuda.current_stream()) + + with torch.cuda.stream(self.stream1): + y1 = MultiplyInStream.apply(x, 3) + self.stream1.wait_event(self.event) + return y0 + y1 + + stream = torch.cuda.Stream() + + for x_first_use_on_ambient in (True, False): + with torch.cuda.stream(stream): + x = torch.randn(5, 5, device='cuda', requires_grad=True) + model = StreamModel().cuda() + x.register_hook(lambda grad: self.assertEqual(torch.cuda.current_stream(), + stream if x_first_use_on_ambient else model.stream0)) + for i in range(5): + model(x, x_first_use_on_ambient).sum().backward() + # See "Stream semantics of backward passes" on https://pytorch.org/docs/stable/notes/cuda.html + torch.cuda.current_stream().wait_stream(stream) + + self.assertEqual(x.grad, torch.ones_like(x) * 5 * 5) + @unittest.skipIf(not TEST_MULTIGPU, "only one GPU detected") def test_streaming_backwards_device_transfer(self): # This function must run with non-default current streams on all devices, otherwise it's meaningless. @@ -1774,7 +1850,7 @@ def test_streaming_backwards_device_transfer(self): self.assertTrue(a.grad.sum().item() == 4 * size) self.assertTrue(b.grad.sum().item() == 4 * size) - def test_streaming_backward_sync_graph_root(self): + def test_streaming_backwards_sync_graph_root(self): # This function tests if bwd ops running on a side stream properly sync with the GraphRoot. # The potential bug it targets is a race condition. The test uses multiple trials and # torch.cuda._sleep such that if the race condition exists, the test will almost certainly fail, @@ -1819,6 +1895,48 @@ def test_streaming_backward_sync_graph_root(self): self.assertEqual(a.grad, grad * b) self.assertEqual(b.grad, grad * a) + def test_streaming_backwards_callback(self): + # Tests if autograd callbacks sync properly with respect to leaf streams and + # the user-facing stream surrounding backward(). If it fails, first suspect is + # sync logic where "final_callbacks_" are called in torch/csrc/autograd/engine.cpp + MultiplyInStream = self._make_multiply_in_stream() + + size = int(1e3) + a = torch.full((size,), 1, device="cuda", dtype=torch.float, requires_grad=True) + b = torch.full((size,), 1, device="cuda", dtype=torch.float, requires_grad=True) + + s0 = torch.cuda.Stream() + s1 = torch.cuda.Stream() + s2 = torch.cuda.Stream() + + stash = [] + + # sets up a nontrivial structure of leaf streams + s0.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s0): + c = MultiplyInStream.apply(a, 2) + + s1.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s1): + d = MultiplyInStream.apply(b, 3) + s1.wait_stream(s0) + e = c * d + + def clone_leaf_grads(): + stash.append(a.grad.clone()) + stash.append(b.grad.clone()) + + # Use a hook on e to install the callback + e.register_hook(lambda grad: torch.autograd.Variable._execution_engine.queue_callback(clone_leaf_grads)) + + s2.wait_stream(s1) + with torch.cuda.stream(s2): + e.sum().backward() + # The autograd engine should sync s2 with all leaf streams then run the callback clone_leaf_grads on s2. + # If those things happened properly, checking the values of the cloned grads on s2 should be safe: + self.assertEqual(stash[0], torch.full_like(a, 6)) + self.assertEqual(stash[1], torch.full_like(a, 6)) + @unittest.skipIf(not TEST_MULTIGPU, "only one GPU detected") @unittest.skipIf(IS_SANDCASTLE or IS_REMOTE_GPU, "Does not work on Sandcastle") def test_cuda_init_race(self): diff --git a/torch/csrc/autograd/engine.cpp b/torch/csrc/autograd/engine.cpp index 364e17fb9e304..d8550ea2eb823 100644 --- a/torch/csrc/autograd/engine.cpp +++ b/torch/csrc/autograd/engine.cpp @@ -11,6 +11,8 @@ #include #include #include +#include + #include #include #include @@ -155,18 +157,25 @@ C10_DEFINE_TLS_static(std::shared_ptr, tls_local_ready_queue); // stream used to run the function OR the inputs are on different devices // and the function is responsible for properly acquiring them. // -// Historically, the autograd engine ran all CUDA operations on their -// device's DEFAULT stream. This meant that syncing (implicitly or -// explicitly) with the default streams was required before and after -// calling backward(). It also meant, however, that syncing with -// the default streams after backward() was sufficient to ensure -// that backward() had finished running. To preserve this historic -// behavior the engine records "leaf streams," the streams of the -// leaf variables, and syncs them with their device's default stream -// at the end of backward. All other streams are already synchronized -// to happen before at least one leaf stream (per the above), so syncing -// the leaf streams with the default streams is sufficient to implement -// the historic behavior. +// User-facing stream semantics of a backward() (or torch.autograd.grad()) +// call with respect to surrounding ops are the same as for any other call. +// See "Stream semantics of backward passes" on +// https://pytorch.org/docs/stable/notes/cuda.html +// +// Internally, backward() runs ops (including leaf nodes) on side threads. +// And streams are thread local. So GraphTask achieves the above semantics by +// 1. remembering the current and default streams on all active CUDA devices +// in the user-facing thread (aka, the thread that called execute() to +// launch the GraphTask) +// 2. remembering the "leaf streams" (streams each backward leaf node ran on) +// 3. during exec_post_processing, for each leaf stream, sync the remembered +// current and default streams (on the leaf stream's device) with that +// leaf stream. +// +// Syncing default streams (as well as current streams) with leaf streams is +// done for temporary BC, and is more conservative than the usage guidance +// (https://pytorch.org/docs/stable/notes/cuda.html) requires. +// TODO: change 1, 2, 3 to sync only current streams with leaf streams. int NodeTask::getReentrantDepth() const { std::shared_ptr graph_task = base_.lock(); @@ -533,24 +542,65 @@ void GraphTask::exec_post_processing() { // more callbacks (or they can be registered from other threads // while it's waiting. std::unique_lock cb_lock(final_callbacks_lock_); - // WARNING: Don't use a range-for loop here because more callbacks may be - // added in between callback calls, so iterators may become invalidated. - // NOLINTNEXTLINE(modernize-loop-convert) - for (const auto i : c10::irange(final_callbacks_.size())) { - cb_lock.unlock(); - final_callbacks_[i](); - cb_lock.lock(); + + // caller_current_streams_ with nullopt entries removed + std::vector caller_current_streams_filtered; + + // See Note [Streaming backwards]. + // Syncs caller_current_stream with leaf streams, so final_callbacks may use + // any grad on its device's current stream. + if (leaf_streams.size() > 0) { + for (const auto& leaf_stream : leaf_streams) { + // stash_current_streams() stashed streams for all device IDs that already had a + // CUDA context before the GraphTask executed. For inactive devices, it stashed + // a c10::nullopt. I don't expect GraphTask's backward pass ran leaf nodes on + // any new devices, so the stashed streams should be enough. + // If leaf_stream.device_index() happens to be for a new device, + // operator* on the c10::nullopt should throw an error. + const auto caller_current_stream = *caller_current_streams_[leaf_stream.device_index()]; + + if (caller_current_stream != leaf_stream) { + auto event = c10::Event{c10::DeviceType::CUDA}; + event.record(leaf_stream); + caller_current_stream.wait(event); + } + } + + caller_current_streams_filtered.reserve(caller_current_streams_.size()); + for (const auto& opt_stream : caller_current_streams_) { + if (opt_stream.has_value()) { + caller_current_streams_filtered.push_back(*opt_stream); + } + } + } + + { + // final_callbacks run on the per-device caller_current_streams (the ambient streams + // surrounding the user's call to backward()). This has two benefits: + // 1. caller_current_streams have been synced with leaf_streams, so callbacks may + // safely access any grad. + // 2. The callback's results can safely be used on (user-facing) caller_current_streams + // after backward(). + c10::MultiStreamGuard g(caller_current_streams_filtered); + // WARNING: Don't use a range-for loop here because more callbacks may be + // added in between callback calls, so iterators may become invalidated. + // NOLINTNEXTLINE(modernize-loop-convert) + for (size_t i = 0; i < final_callbacks_.size(); ++i) { + cb_lock.unlock(); + final_callbacks_[i](); + cb_lock.lock(); + } } - // Syncs leaf streams with default streams (if necessary) - // See note "Streaming backwards" - for (const auto& leaf_stream : leaf_streams) { - const auto guard = c10::impl::VirtualGuardImpl{c10::DeviceType::CUDA}; - const auto default_stream = guard.getDefaultStream(leaf_stream.device()); - if (leaf_stream != default_stream) { + // For temporary BC, syncs default streams with caller_current_streams so callback results are also + // usable on user-facing default streams after backward() + for (const auto& caller_current_stream : caller_current_streams_filtered) { + const auto caller_default_stream = *caller_default_streams_[caller_current_stream.device_index()]; + + if (caller_current_stream != caller_default_stream) { auto event = c10::Event{c10::DeviceType::CUDA}; - event.record(leaf_stream); - default_stream.wait(event); + event.record(caller_current_stream); + caller_default_stream.wait(event); } } } @@ -772,7 +822,7 @@ void Engine::evaluate_function( int num_outputs = outputs.size(); if (num_outputs == 0) { // Note: doesn't acquire the mutex // Records leaf stream (if applicable) - // See note "Streaming backwards" + // See Note [Streaming backwards] if (opt_parent_stream) { std::lock_guard lock(graph_task->mutex_); graph_task->leaf_streams.emplace(*opt_parent_stream); @@ -878,6 +928,8 @@ auto Engine::compute_dependencies(Node* root, GraphTask& task, uint64_t min_topo // Computes the number of dependencies for each function which requires grad std::unordered_set seen; std::vector queue { root }; + bool might_use_cuda = at::globalContext().hasCUDA(); + bool will_use_cuda = false; // Queue contains all nodes that will start propagating gradients. // We no longer have to expand functions that don't require grad. @@ -887,6 +939,9 @@ auto Engine::compute_dependencies(Node* root, GraphTask& task, uint64_t min_topo if (fn->topological_nr() < min_topo_nr) { continue; } + if (might_use_cuda && !will_use_cuda) { + will_use_cuda = fn->stream(c10::DeviceType::CUDA).has_value(); + } for (const auto& edge : fn->next_edges()) { if (auto next_ptr = edge.function.get()) { dependencies[next_ptr] += 1; @@ -895,6 +950,12 @@ auto Engine::compute_dependencies(Node* root, GraphTask& task, uint64_t min_topo } } } + + if (will_use_cuda) { + // Collects current and default streams for devices where this process has a context, + // so GraphTask::exec_post_processing can sync them with leaf_streams. + task.stash_current_streams(); + } } auto Engine::execute(const edge_list& roots, @@ -1179,6 +1240,33 @@ void Engine::add_thread_pool_task(const std::weak_ptr& graph_task) { thread_pool_shared_->work_.notify_one(); } +// Remembers current and default streams on all devices where a context has been created. +// Only called if Engine::execute detects at least one node runs on a cuda stream. +void GraphTask::stash_current_streams() { + const auto guard = c10::impl::VirtualGuardImpl{c10::DeviceType::CUDA}; + auto num_gpus = guard.deviceCount(); + caller_current_streams_.resize(num_gpus); + caller_default_streams_.resize(num_gpus); + if (num_gpus > 0) { + for (c10::DeviceIndex idx = 0; idx < num_gpus; idx++) { +#ifdef __HIP_PLATFORM_HCC__ + // If the build targets ROCM, stash streams for all visible devices unconditionally, to work around + // https://github.com/pytorch/pytorch/issues/59750. + // TODO: Remove ROCM-specific behavior when https://github.com/pytorch/pytorch/issues/59750 is fixed. + if (true) { +#else + if (at::detail::getCUDAHooks().hasPrimaryContext(idx)) { +#endif + caller_current_streams_[idx] = guard.getStream({c10::DeviceType::CUDA, idx}); + caller_default_streams_[idx] = guard.getDefaultStream({c10::DeviceType::CUDA, idx}); + } else { + caller_current_streams_[idx] = c10::nullopt; + caller_default_streams_[idx] = c10::nullopt; + } + } + } +} + void GraphTask::init_to_execute(Node& graph_root, const edge_list& outputs, bool accumulate_grad, uint64_t min_topo_nr) { // Populates exec_info so nodes that should be executed have `exec_info[node].needed_ = true` // Only nodes that have a path to any edge in `outputs` should be executed. diff --git a/torch/csrc/autograd/engine.h b/torch/csrc/autograd/engine.h index 637192035c0ce..7c810cf9891d8 100644 --- a/torch/csrc/autograd/engine.h +++ b/torch/csrc/autograd/engine.h @@ -112,6 +112,14 @@ struct GraphTask: std::enable_shared_from_this { std::unordered_set leaf_streams; + // Per-device current and default streams of the execute() that called this GraphTask. + // These will be synced with leaf_streams in exec_post_processing. + std::vector> caller_current_streams_; + std::vector> caller_default_streams_; + + // Collects caller_current_streams_ and caller_default_streams_ + void stash_current_streams(); + void init_to_execute(Node& graph_root, const edge_list& outputs, bool accumulate_grad, uint64_t min_topo_nr); // The value of worker_device in the thread that created this task. diff --git a/torch/csrc/distributed/autograd/engine/dist_engine.cpp b/torch/csrc/distributed/autograd/engine/dist_engine.cpp index 2f1f7ef7323b8..6c9feda380747 100644 --- a/torch/csrc/distributed/autograd/engine/dist_engine.cpp +++ b/torch/csrc/distributed/autograd/engine/dist_engine.cpp @@ -212,6 +212,9 @@ void DistEngine::computeDependencies( queue.push(mapEntry.second.get()); } + bool might_use_cuda = at::globalContext().hasCUDA(); + bool will_use_cuda = false; + edge_list recvBackwardEdges; // Traverse the graph. auto& dependencies = graphTask->dependencies_; @@ -219,6 +222,10 @@ void DistEngine::computeDependencies( auto fn = queue.front(); queue.pop(); + if (might_use_cuda && !will_use_cuda) { + will_use_cuda = fn->stream(c10::DeviceType::CUDA).has_value(); + } + for (const auto& edge : fn->next_edges()) { if (auto nextFn = edge.function.get()) { dependencies[nextFn] += 1; @@ -255,6 +262,12 @@ void DistEngine::computeDependencies( } } + if (will_use_cuda) { + // Collects current and default streams for devices where this process has a context, + // so graphTask::exec_post_processing can sync them with leaf_streams. + graphTask->stash_current_streams(); + } + // Now lets compute which functions need to be executed. The algorithm is as // follows: // 1. Create a dummy GraphRoot which points to all 'send' functions for this diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp index c7f2990f59b77..1f556cc126976 100644 --- a/torch/lib/c10d/reducer.cpp +++ b/torch/lib/c10d/reducer.cpp @@ -504,25 +504,11 @@ void Reducer::set_divide_factor() { } } -const c10::Stream Reducer::get_current_stream() { - const auto& device = replicas_[0][0].device(); - c10::DeviceType deviceType = device.type(); - const c10::impl::VirtualGuardImpl guard = - c10::impl::VirtualGuardImpl{deviceType}; - return guard.getStream(device); -} - // Right now delay_all_reduce is only called when static_graph_=true and // num_iterations_==1. void Reducer::delay_all_reduce() { std::lock_guard lock(this->mutex_); - // The autograd engine uses the default stream when running callbacks, so we - // pass in the current CUDA stream in case it is not the default. - const c10::Stream currentStream = get_current_stream(); - // Run callback with the current stream - c10::OptionalStreamGuard currentStreamGuard{currentStream}; - if (should_collect_runtime_stats()) { record_backward_compute_end_time(); record_backward_comm_start_time(); @@ -832,13 +818,8 @@ void Reducer::mark_variable_ready(size_t variable_index) { all_reduce_local_used_map(); } - // The autograd engine uses the default stream when running callbacks, so we - // pass in the current CUDA stream in case it is not the default. - const c10::Stream currentStream = get_current_stream(); torch::autograd::Engine::get_default_engine().queue_callback([=] { std::lock_guard lock(this->mutex_); - // Run callback with the current stream - c10::OptionalStreamGuard currentStreamGuard{currentStream}; if (should_collect_runtime_stats()) { record_backward_compute_end_time(); } From 095cd6a0da255a5ed24d146898c53790f0e9c3f0 Mon Sep 17 00:00:00 2001 From: Peter Bell Date: Sun, 13 Jun 2021 12:29:06 -0700 Subject: [PATCH 074/305] MemoryOverlap: Avoid has_storage calls (#59013) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59013 Test Plan: Imported from OSS Reviewed By: bdhirsh Differential Revision: D29040929 Pulled By: ngimel fbshipit-source-id: 69745e7abbaf523795a90f68cf01d3d94508210e --- aten/src/ATen/MemoryOverlap.cpp | 6 ++---- c10/core/TensorImpl.h | 9 +++++++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/aten/src/ATen/MemoryOverlap.cpp b/aten/src/ATen/MemoryOverlap.cpp index 4b90f59f5adab..c2388ebf8d968 100644 --- a/aten/src/ATen/MemoryOverlap.cpp +++ b/aten/src/ATen/MemoryOverlap.cpp @@ -48,15 +48,13 @@ MemOverlapStatus get_overlap_status(TensorImpl* a, TensorImpl* b) { if (!a->is_non_overlapping_and_dense() || !b->is_non_overlapping_and_dense()) { return MemOverlapStatus::TOO_HARD; } - if (!a->has_storage() || !b->has_storage()) { - return MemOverlapStatus::NO; - } // Test for storage equality, rather than pointer equality. // This reduces precision, but if people are aliasing the // same pointer across multiple storages there are many // similar situations (e.g., storage().data() == storage().data()+1) // which we will miss. - if (a->storage().is_alias_of(b->storage())) { + auto a_storage = a->unsafe_storage(); + if (a_storage && a_storage.is_alias_of(b->unsafe_storage())) { const auto a_begin = static_cast(a->data()); const auto a_end = a_begin + a->numel() * a->itemsize(); const auto b_begin = static_cast(b->data()); diff --git a/c10/core/TensorImpl.h b/c10/core/TensorImpl.h index 0ff3411bb9136..458090d446c77 100644 --- a/c10/core/TensorImpl.h +++ b/c10/core/TensorImpl.h @@ -672,6 +672,15 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target { return storage_; } + /** + * Return the underlying storage, unsafely assuming this is a basic strided + * tensor. In cases where `storage` access would throw, this returns a + * default-constructed Storage. + */ + inline const Storage& unsafe_storage() const { + return storage_; + } + /** * The number of elements in a tensor. * From ff15d93b88c0e1e32d2553f272bc654d3b00606f Mon Sep 17 00:00:00 2001 From: Xiaomeng Yang Date: Sun, 13 Jun 2021 16:12:15 -0700 Subject: [PATCH 075/305] Improve numerical stability of GroupNorm (#54921) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/54921 Improve numerical stability of GroupNorm Test Plan: buck test mode/dev-nosan //caffe2/test:nn -- "GroupNorm" Reviewed By: ngimel Differential Revision: D27414438 fbshipit-source-id: 815517240ca5ea3e2beb77ced3bd862e9c83d445 --- aten/src/ATen/native/SharedReduceOps.h | 13 +- aten/src/ATen/native/cpu/SumKernel.cpp | 21 +-- .../src/ATen/native/cpu/group_norm_kernel.cpp | 59 +++---- aten/src/ATen/native/cpu/moments_utils.h | 147 ++++++++++++++++++ aten/src/ATen/native/cpu/utils.h | 11 +- aten/src/ATen/native/cuda/block_reduce.cuh | 31 ++++ .../src/ATen/native/cuda/group_norm_kernel.cu | 45 ++++-- aten/src/ATen/native/group_norm.cpp | 42 +++-- test/test_nn.py | 31 +++- torch/testing/_internal/common_nn.py | 11 ++ 10 files changed, 319 insertions(+), 92 deletions(-) create mode 100644 aten/src/ATen/native/cpu/moments_utils.h diff --git a/aten/src/ATen/native/SharedReduceOps.h b/aten/src/ATen/native/SharedReduceOps.h index 571ce4c269ed6..5f5cdf9bf1c47 100644 --- a/aten/src/ATen/native/SharedReduceOps.h +++ b/aten/src/ATen/native/SharedReduceOps.h @@ -80,8 +80,15 @@ struct WelfordData { scalar_t m2; index_t n; combine_t nf; - C10_HOST_DEVICE WelfordData() : mean(0), m2(0), n(0), nf(0) {} - C10_DEVICE WelfordData(scalar_t mean, scalar_t m2, index_t n, combine_t nf) : mean(mean), m2(m2), n(n), nf(nf) {} + + C10_HOST_DEVICE WelfordData() : mean(0), m2(0), n(0), nf(0) {} + + C10_HOST_DEVICE WelfordData( + scalar_t mean, + scalar_t m2, + index_t n, + combine_t nf) + : mean(mean), m2(m2), n(n), nf(nf) {} }; @@ -145,7 +152,7 @@ struct WelfordOps { }; } #endif - WelfordOps(index_t correction, bool take_sqrt) + C10_HOST_DEVICE WelfordOps(index_t correction, bool take_sqrt) : correction(correction), take_sqrt(take_sqrt) {} }; diff --git a/aten/src/ATen/native/cpu/SumKernel.cpp b/aten/src/ATen/native/cpu/SumKernel.cpp index cd50cd39f8fd6..e0e5e21069f17 100644 --- a/aten/src/ATen/native/cpu/SumKernel.cpp +++ b/aten/src/ATen/native/cpu/SumKernel.cpp @@ -1,11 +1,11 @@ -#include -#include #include -#include -#include #include +#include +#include +#include +#include namespace at { namespace native { @@ -48,17 +48,6 @@ void accumulate_result(char * C10_RESTRICT data, int64_t stride, int64_t index, } } -int64_t ceil_log2(int64_t x) { - if (x <= 2) { - return 1; - } - - auto ux = static_cast(x); - // Last set bit is floor(log2(x)), floor + 1 is ceil - // except when x is an exact powers of 2, so subtract 1 first - return static_cast(llvm::findLastSet(ux - 1)) + 1; -} - /** Simultaneously sum over n rows at once This algorithm calculates the sum without loss of precision over large axes. It @@ -101,7 +90,7 @@ std::array multi_row_sum( constexpr int64_t num_levels = 4; const int64_t level_power = - std::max(int64_t(4), ceil_log2(size) / num_levels); + std::max(int64_t(4), utils::CeilLog2(size) / num_levels); const int64_t level_step = (1 << level_power); const int64_t level_mask = level_step - 1; diff --git a/aten/src/ATen/native/cpu/group_norm_kernel.cpp b/aten/src/ATen/native/cpu/group_norm_kernel.cpp index c59b162f00f67..97a59cafd5a9d 100644 --- a/aten/src/ATen/native/cpu/group_norm_kernel.cpp +++ b/aten/src/ATen/native/cpu/group_norm_kernel.cpp @@ -8,6 +8,7 @@ #include #include #include +#include namespace at { namespace native { @@ -38,47 +39,33 @@ void GroupNormKernelImplInternal( T* Y_data = Y.data_ptr(); T* mean_data = mean.data_ptr(); T* rstd_data = rstd.data_ptr(); - const T s = T(1) / static_cast(D * HxW); const bool gamma_null = (gamma_data == nullptr); const bool beta_null = beta_data == nullptr; + const int64_t inner_size = D * HxW; at::parallel_for(0, N * G, 1, [&](int64_t start, int64_t end) { - constexpr int64_t K = vec::Vectorized::size(); - const int64_t inner_size = D * HxW / K * K; - // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) - std::array mean_arr; - // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) - std::array rstd_arr; for (int64_t i = start; i < end; ++i) { - const T* X_ptr = X_data + i * D * HxW; - vec::Vectorized mean_vec(0); - vec::Vectorized rstd_vec(0); - for (int64_t j = 0; j < inner_size; j += K) { - const vec::Vectorized x_vec = vec::Vectorized::loadu(X_ptr + j); - mean_vec = mean_vec + x_vec; - rstd_vec = rstd_vec + x_vec * x_vec; - } - mean_vec.store(mean_arr.data()); - rstd_vec.store(rstd_arr.data()); - T mean_val = std::accumulate(mean_arr.cbegin(), mean_arr.cend(), T(0)); - T rstd_val = std::accumulate(rstd_arr.cbegin(), rstd_arr.cend(), T(0)); - for (int64_t j = inner_size; j < D * HxW; ++j) { - mean_val += X_ptr[j]; - rstd_val += X_ptr[j] * X_ptr[j]; - } - mean_val *= s; - rstd_val = std::max(rstd_val * s - mean_val * mean_val, T(0)); - rstd_val = T(1) / std::sqrt(rstd_val + eps); - - const int64_t g = i % G; - for (int64_t j = 0; j < D; ++j) { - const int64_t c = g * D + j; - const T scale = rstd_val * (gamma_null ? T(1) : gamma_data[c]); - const T bias = -scale * mean_val + (beta_null ? T(0) : beta_data[c]); - X_ptr = X_data + (i * D + j) * HxW; - T* Y_ptr = Y_data + (i * D + j) * HxW; - for (int64_t k = 0; k < HxW; ++k) { - Y_ptr[k] = scale * X_ptr[k] + bias; + const T* X_ptr = X_data + i * inner_size; + T mean_val; + T rstd_val; + std::tie(mean_val, rstd_val) = utils::RowwiseMoments(X_ptr, inner_size); + rstd_val = T(1) / std::sqrt(std::max(rstd_val, T(0)) + eps); + if (gamma_null && beta_null) { + T* Y_ptr = Y_data + i * inner_size; + for (int j = 0; j < inner_size; ++j) { + Y_ptr[j] = (X_ptr[j] - mean_val) * rstd_val; + } + } else { + const int64_t g = i % G; + for (int64_t j = 0; j < D; ++j) { + const int64_t c = g * D + j; + const T scale = rstd_val * (gamma_null ? T(1) : gamma_data[c]); + const T bias = -scale * mean_val + (beta_null ? T(0) : beta_data[c]); + X_ptr = X_data + (i * D + j) * HxW; + T* Y_ptr = Y_data + (i * D + j) * HxW; + for (int64_t k = 0; k < HxW; ++k) { + Y_ptr[k] = scale * X_ptr[k] + bias; + } } } mean_data[i] = mean_val; diff --git a/aten/src/ATen/native/cpu/moments_utils.h b/aten/src/ATen/native/cpu/moments_utils.h new file mode 100644 index 0000000000000..20d79ff479292 --- /dev/null +++ b/aten/src/ATen/native/cpu/moments_utils.h @@ -0,0 +1,147 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace at { +namespace native { +namespace utils { + +constexpr int64_t kChunkSize = 16; + +template +void AddMoments( + int64_t m0_add, + const T& m1_add, + const T& m2_add, + int64_t& m0, + T& m1, + T& m2) { + const int64_t n = m0 + m0_add; + const T c = n == 0 ? 0 : static_cast(m0_add) / static_cast(n); + const T delta = m1_add - m1; + m1 += c * delta; + m2 += m2_add + delta * delta * c * static_cast(m0); + m0 = n; +} + +template +void AddMomentsVec( + int64_t m0_add, + const vec::Vectorized& m1_add, + const vec::Vectorized& m2_add, + int64_t& m0, + vec::Vectorized& m1, + vec::Vectorized& m2) { + using Vec = vec::Vectorized; + const int64_t n = m0 + m0_add; + const T c = n == 0 ? 0 : static_cast(m0_add) / static_cast(n); + const Vec c_vec(c); + const Vec delta = m1_add - m1; + m1 += c_vec * delta; + m2 += m2_add + delta * delta * c_vec * Vec(static_cast(m0)); + m0 = n; +} + +// Compute rowwise moments by Welford algorithm and cascade sum to improve +// numerical stability. +// https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance +// https://en.wikipedia.org/wiki/Pairwise_summation +template +std::pair RowwiseMomentsImpl(const T* X, int64_t N) { + using Vec = vec::Vectorized; + + constexpr int64_t kVecSize = Vec::size(); + const int64_t n = N / kVecSize; + const int64_t m = divup(n, kChunkSize); + const int64_t depth = CeilLog2(m); + + const Vec kZeroVec(T(0)); + c10::SmallVector m0_stk(depth, 0); + c10::SmallVector m1_stk(depth, kZeroVec); + c10::SmallVector m2_stk(depth, kZeroVec); + + for (int64_t i = 0; i < m; ++i) { + const T* X_ptr = X + i * kChunkSize * kVecSize; + const int64_t m0 = std::min(kChunkSize, n - i * kChunkSize); + Vec m1_vec(0); + Vec m2_vec(0); + for (int64_t j = 0; j < m0; ++j) { + const Vec x_vec = Vec::loadu(X_ptr + j * kVecSize); + const Vec delta_vec = x_vec - m1_vec; + const Vec c_vec = Vec(T(1) / static_cast(j + 1)); + m1_vec += delta_vec * c_vec; + m2_vec += delta_vec * (x_vec - m1_vec); + } + AddMomentsVec(m0, m1_vec, m2_vec, m0_stk[0], m1_stk[0], m2_stk[0]); + int64_t mask = i + 1; + for (int64_t j = 1; j < depth && (mask & 1) == 0; ++j) { + AddMomentsVec( + m0_stk[j - 1], + m1_stk[j - 1], + m2_stk[j - 1], + m0_stk[j], + m1_stk[j], + m2_stk[j]); + m0_stk[j - 1] = 0; + m1_stk[j - 1] = kZeroVec; + m2_stk[j - 1] = kZeroVec; + mask >>= 1; + } + } + for (int64_t i = 1; i < depth; ++i) { + AddMomentsVec( + m0_stk[i], m1_stk[i], m2_stk[i], m0_stk[0], m1_stk[0], m2_stk[0]); + } + + std::array m1_arr{}; + std::array m2_arr{}; + m1_stk[0].store(m1_arr.data()); + m2_stk[0].store(m2_arr.data()); + + int64_t m0 = 0; + T m1 = 0; + T m2 = 0; + for (int64_t i = n * kVecSize; i < N; ++i) { + const T delta = X[i] - m1; + ++m0; + m1 += delta / static_cast(m0); + m2 += delta * (X[i] - m1); + } + for (int64_t i = 0; i < kVecSize; ++i) { + AddMoments(n, m1_arr[i], m2_arr[i], m0, m1, m2); + } + + return std::make_pair(m1, m2 / static_cast(N)); +} + +template +std::pair RowwiseMoments(const T* X, int64_t N) { + using Vec = vec::Vectorized; + constexpr int64_t kVecSize = Vec::size(); + const int64_t n = N / kVecSize; + const int64_t m = divup(n, kChunkSize); + const int64_t depth = CeilLog2(m); + if (depth <= 4) { + return RowwiseMomentsImpl(X, N); + } else if (depth <= 8) { + return RowwiseMomentsImpl(X, N); + } else if (depth <= 16) { + return RowwiseMomentsImpl(X, N); + } else if (depth <= 32) { + return RowwiseMomentsImpl(X, N); + } else { + return RowwiseMomentsImpl(X, N); + } +} + +} // namespace utils +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/cpu/utils.h b/aten/src/ATen/native/cpu/utils.h index 566a322047b97..5a76f751adac2 100644 --- a/aten/src/ATen/native/cpu/utils.h +++ b/aten/src/ATen/native/cpu/utils.h @@ -3,7 +3,10 @@ #include #include -namespace at { namespace native { namespace { +namespace at { +namespace native { + +namespace { template inline T data_index_init(T offset) { @@ -11,7 +14,7 @@ inline T data_index_init(T offset) { } template -inline T data_index_init(T offset, T &x, const T &X, Args &&... args) { +inline T data_index_init(T offset, T& x, const T& X, Args&&... args) { offset = data_index_init(offset, std::forward(args)...); x = offset % X; return offset / X; @@ -22,7 +25,7 @@ inline bool data_index_step() { } template -inline bool data_index_step(T &x, const T &X, Args &&... args) { +inline bool data_index_step(T& x, const T& X, Args&&... args) { if (data_index_step(std::forward(args)...)) { x = ((x + 1) == X) ? 0 : (x + 1); return x == 0; @@ -47,4 +50,4 @@ T CeilLog2(const T& x) { } // namespace utils } // namespace native -} // namespace at// namespace at::native:: +} // namespace at diff --git a/aten/src/ATen/native/cuda/block_reduce.cuh b/aten/src/ATen/native/cuda/block_reduce.cuh index 325bf616d89aa..a3f600e48848a 100644 --- a/aten/src/ATen/native/cuda/block_reduce.cuh +++ b/aten/src/ATen/native/cuda/block_reduce.cuh @@ -1,5 +1,8 @@ #pragma once +#include + +#include #include namespace at { @@ -45,6 +48,34 @@ __inline__ __device__ T BlockReduceSum(T val, T* shared) { return val; } +template +__inline__ __device__ T WarpReduce(T val, const ReduceOp& op) { +#pragma unroll + for (int offset = (C10_WARP_SIZE >> 1); offset > 0; offset >>= 1) { + val = op.combine(val, op.warp_shfl_down(val, offset)); + } + return val; +} + +template +__inline__ __device__ T +BlockReduce(T val, const ReduceOp& op, const T& identity_element, T* shared) { + const int lid = threadIdx.x % C10_WARP_SIZE; + const int wid = threadIdx.x / C10_WARP_SIZE; + val = WarpReduce(val, op); + __syncthreads(); + if (lid == 0) { + shared[wid] = val; + } + __syncthreads(); + val = (threadIdx.x < blockDim.x / C10_WARP_SIZE) ? shared[lid] + : identity_element; + if (wid == 0) { + val = WarpReduce(val, op); + } + return val; +} + } // namespace cuda_utils } // namespace native } // namespace at diff --git a/aten/src/ATen/native/cuda/group_norm_kernel.cu b/aten/src/ATen/native/cuda/group_norm_kernel.cu index 46601eee24125..81d87bfb51d56 100644 --- a/aten/src/ATen/native/cuda/group_norm_kernel.cu +++ b/aten/src/ATen/native/cuda/group_norm_kernel.cu @@ -8,14 +8,14 @@ #include #include #include +#include #include +#include #include #include #include #include -#include - namespace at { namespace native { @@ -32,29 +32,38 @@ __global__ void RowwiseMomentsCUDAKernel( T* mean, T* rstd) { using T_ACC = acc_type; + using WelfordType = WelfordData; + using WelfordOp = + WelfordOps>; + const int64_t i = blockIdx.x; - T_ACC sum1 = 0; - T_ACC sum2 = 0; + WelfordOp welford_op = {/*correction=*/0, /*take_sqrt=*/false}; + WelfordType val(0, 0, 0, 0); for (int64_t j = threadIdx.x; j < N; j += blockDim.x) { const int64_t index = i * N + j; - sum1 += static_cast(X[index]); - sum2 += static_cast(X[index]) * static_cast(X[index]); + val = welford_op.reduce(val, static_cast(X[index]), index); } if (blockDim.x <= C10_WARP_SIZE) { - sum1 = cuda_utils::WarpReduceSum(sum1); - sum2 = cuda_utils::WarpReduceSum(sum2); + val = cuda_utils::WarpReduce(val, welford_op); } else { - __shared__ T_ACC m_shared[C10_WARP_SIZE]; - __shared__ T_ACC v_shared[C10_WARP_SIZE]; - sum1 = cuda_utils::BlockReduceSum(sum1, m_shared); - sum2 = cuda_utils::BlockReduceSum(sum2, v_shared); + // There will be a warning if we declare a __shared__ WelfordType array. + // https://github.com/pytorch/pytorch/pull/13967 + __shared__ typename std::aligned_storage< + sizeof(WelfordType), + alignof(WelfordType)>::type val_shared[C10_WARP_SIZE]; + WelfordType* val_shared_ptr = reinterpret_cast(val_shared); + val = cuda_utils::BlockReduce( + val, + welford_op, + /*identity_element=*/WelfordType(0, 0, 0, 0), + val_shared_ptr); } if (threadIdx.x == 0) { - const T_ACC scale = T_ACC(1) / static_cast(N); - sum1 *= scale; - sum2 = c10::cuda::compat::max(sum2 * scale - sum1 * sum1, T_ACC(0)); - mean[i] = sum1; - rstd[i] = c10::cuda::compat::rsqrt(sum2 + static_cast(eps)); + T_ACC m1; + T_ACC m2; + thrust::tie(m2, m1) = welford_op.project(val); + mean[i] = m1; + rstd[i] = c10::cuda::compat::rsqrt(m2 + static_cast(eps)); } } @@ -605,6 +614,7 @@ void GroupNormKernelImplInternal( ComputeFusedParamsCUDAKernel<<>>( N, C, G, mean_data, rstd_data, gamma_data, beta_data, a_data, b_data); C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto iter = TensorIteratorConfig() .check_all_same_dtype(std::is_same::value) .resize_outputs(false) @@ -617,6 +627,7 @@ void GroupNormKernelImplInternal( return a * static_cast(x) + b; }); } + AT_CUDA_CHECK(cudaGetLastError()); } void GroupNormKernelImpl( diff --git a/aten/src/ATen/native/group_norm.cpp b/aten/src/ATen/native/group_norm.cpp index 83faa8b957032..5c205d8ce6f5b 100644 --- a/aten/src/ATen/native/group_norm.cpp +++ b/aten/src/ATen/native/group_norm.cpp @@ -1,10 +1,10 @@ -#include #include -#include +#include #include -#include +#include #include #include +#include #include #include @@ -13,21 +13,25 @@ #include #include - namespace at { namespace native { std::tuple native_group_norm( - const Tensor& X, const c10::optional& gamma_opt /* optional */, const c10::optional& beta_opt /* optional */, + const Tensor& X, + const c10::optional& gamma_opt /* optional */, + const c10::optional& beta_opt /* optional */, int64_t N, int64_t C, int64_t HxW, int64_t group, double eps) { // See [Note: hacky wrapper removal for optional tensor] - c10::MaybeOwned gamma_maybe_owned = at::borrow_from_optional_tensor(gamma_opt); + c10::MaybeOwned gamma_maybe_owned = + at::borrow_from_optional_tensor(gamma_opt); const Tensor& gamma = *gamma_maybe_owned; - const Tensor& beta = c10::value_or_else(beta_opt, [] {return Tensor();}); + const Tensor& beta = c10::value_or_else(beta_opt, [] { return Tensor(); }); + + TORCH_CHECK(X.is_contiguous()); Tensor Y = at::native::empty_like( X, @@ -47,14 +51,16 @@ std::tuple native_group_norm_backward( const Tensor& dY, const Tensor& X, const Tensor& mean, - const Tensor& rstd, const c10::optional& gamma_opt, + const Tensor& rstd, + const c10::optional& gamma_opt, int64_t N, int64_t C, int64_t HxW, int64_t group, std::array grad_input_mask) { // See [Note: hacky wrapper removal for optional tensor] - c10::MaybeOwned gamma_maybe_owned = at::borrow_from_optional_tensor(gamma_opt); + c10::MaybeOwned gamma_maybe_owned = + at::borrow_from_optional_tensor(gamma_opt); const Tensor& gamma = *gamma_maybe_owned; Tensor dX; @@ -106,13 +112,16 @@ std::tuple native_group_norm_backward( Tensor group_norm( const Tensor& input, - int64_t num_groups, const c10::optional& weight_opt /* optional */, const c10::optional& bias_opt /* optional */, + int64_t num_groups, + const c10::optional& weight_opt /* optional */, + const c10::optional& bias_opt /* optional */, double eps, bool /* cudnn_enabled, deprecated */) { // See [Note: hacky wrapper removal for optional tensor] - c10::MaybeOwned weight_maybe_owned = at::borrow_from_optional_tensor(weight_opt); + c10::MaybeOwned weight_maybe_owned = + at::borrow_from_optional_tensor(weight_opt); const Tensor& weight = *weight_maybe_owned; - const Tensor& bias = c10::value_or_else(bias_opt, [] {return Tensor();}); + const Tensor& bias = c10::value_or_else(bias_opt, [] { return Tensor(); }); const int64_t N = input.size(0); const int64_t C = input.size(1); @@ -160,16 +169,19 @@ DEFINE_DISPATCH(GroupNormBackwardKernel); // Ported from pytorch/xla repo std::tuple math_group_norm( - const Tensor& input, const c10::optional& weight_opt, const c10::optional& bias_opt, + const Tensor& input, + const c10::optional& weight_opt, + const c10::optional& bias_opt, int64_t N, int64_t C, int64_t HxW, int64_t group, double eps) { // See [Note: hacky wrapper removal for optional tensor] - c10::MaybeOwned weight_maybe_owned = at::borrow_from_optional_tensor(weight_opt); + c10::MaybeOwned weight_maybe_owned = + at::borrow_from_optional_tensor(weight_opt); const Tensor& weight = *weight_maybe_owned; - const Tensor& bias = c10::value_or_else(bias_opt, [] {return Tensor();}); + const Tensor& bias = c10::value_or_else(bias_opt, [] { return Tensor(); }); auto input_shape = input.sizes(); at::Tensor input_reshaped = input.view({1, N * group, N ? -1 : 1}); diff --git a/test/test_nn.py b/test/test_nn.py index 7ce0da8a32f66..3dde5053ca282 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -5762,7 +5762,6 @@ def test_RNN_cell_forward_hidden_size(self): nn.GRUCell(*cell_shared_param)): self.assertRaises(Exception, lambda: cell(input, hx)) - def _test_loss_equal_input_target_shape(self, cast): # Tests losses whose inputs should have the same size. losses = { @@ -12549,6 +12548,36 @@ def test_GroupNorm_empty(self, device): with torch.backends.cudnn.flags(enabled=False): self._test_module_empty_input(mod, inp) + @onlyOnCPUAndCUDA + def test_GroupNorm_numeric(self, device): + def group_norm_ref(X, gamma, beta, groups, channels, eps): + batch_size = X.size()[0] + X_view = X.view(batch_size, groups, -1) + mean = X_view.mean(dim=-1, keepdim=True) + var = X_view.var(dim=-1, unbiased=False, keepdim=True) + Y = ((X_view - mean) / torch.sqrt(var + eps)).view( + batch_size, channels, -1) + Y = Y * gamma.view(channels, 1) + beta.view(channels, 1) + return Y.view(*X.size()) + + batch_size = 1 + groups = 4 + channels = 32 + group_norm = nn.GroupNorm(groups, channels).float().to(device) + X = torch.rand(batch_size, channels, 256, 256, 72, + dtype=torch.float32, device=device) + + Y = group_norm(X) + Y_ref = group_norm_ref( + X, group_norm.weight.data, group_norm.bias.data, groups, + channels, group_norm.eps) + self.assertEqual(Y, Y_ref, rtol=0, atol=1e-5) + + if self.device_type == 'cuda': + group_norm.cpu() + Y_cpu = group_norm(X.cpu()) + self.assertEqual(Y_cpu, Y, rtol=0, atol=1e-5) + @onlyOnCPUAndCUDA @dtypes(torch.float64, torch.complex128) def test_pad(self, device, dtype): diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py index ad4a0390d4585..a6a8045b1148f 100644 --- a/torch/testing/_internal/common_nn.py +++ b/torch/testing/_internal/common_nn.py @@ -1647,6 +1647,17 @@ def fractional_max_pool3d_test(test_case): desc='2d_affine_large_feature', test_cpu=False, ), + dict( + module_name='GroupNorm', + constructor_args=(3, 51, 1e-5, False), + cpp_constructor_args='torch::nn::GroupNormOptions(3, 51).eps(1e-5).affine(false)', + input_size=(2, 51, 28, 28), + cudnn=True, + check_eval=True, + check_bfloat16=True, + desc='2d_no_affine_large_feature', + test_cpu=False, + ), dict( module_name='GroupNorm', constructor_args=(3, 3, 1e-3, False), From cf38b20c61e2b08496c51b0d879892f388d6e03b Mon Sep 17 00:00:00 2001 From: Kushashwa Ravi Shrimali Date: Mon, 14 Jun 2021 03:03:46 -0700 Subject: [PATCH 076/305] Alias for `digamma` as `psi` to `special` namespace (#59143) Summary: See https://github.com/pytorch/pytorch/issues/50345 cc: mruberry kshitij12345 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59143 Reviewed By: jbschlosser Differential Revision: D28986909 Pulled By: mruberry fbshipit-source-id: bc8ff0375de968f3662b224689fa0a6b117f9c4e --- aten/src/ATen/core/aten_interned_strings.h | 1 - aten/src/ATen/core/interned_strings.h | 3 ++ aten/src/ATen/native/UnaryOps.cpp | 7 ++++ aten/src/ATen/native/native_functions.yaml | 16 +++++++++ docs/source/special.rst | 2 ++ torch/_torch_docs.py | 25 ++----------- torch/csrc/api/include/torch/special.h | 32 +++++++++++++++++ torch/csrc/jit/passes/normalize_ops.cpp | 2 ++ torch/overrides.py | 2 ++ torch/special/__init__.py | 35 +++++++++++++++++++ .../_internal/common_methods_invocations.py | 1 + 11 files changed, 102 insertions(+), 24 deletions(-) diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 1a7486a019a06..25cda648c89a6 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -296,7 +296,6 @@ _(aten, diagonal) \ _(aten, fill_diagonal_) \ _(aten, diff) \ _(aten, frexp) \ -_(aten, digamma) \ _(aten, dim) \ _(aten, dist) \ _(aten, dot) \ diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h index 96a49f4426bec..d06618271bc77 100644 --- a/aten/src/ATen/core/interned_strings.h +++ b/aten/src/ATen/core/interned_strings.h @@ -322,6 +322,9 @@ namespace c10 { _(aten, moveaxis) \ _(aten, lgamma) \ _(aten, special_gammaln) \ + _(aten, digamma) \ + _(aten, special_psi) \ + _(aten, special_digamma) \ _(aten, erf) \ _(aten, special_erf) \ _(aten, erfc) \ diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index 624881a01dbf6..4ba59ef3d05dd 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -441,6 +441,13 @@ Tensor special_erfc(const Tensor& self) { return self.erfc(); } Tensor& special_erfinv_out(const Tensor& self, Tensor& result) { return at::erfinv_out(result, self); } Tensor special_erfinv(const Tensor& self) { return self.erfinv(); } +// special_psi, alias for digamma +Tensor& special_psi_out(const Tensor& self, Tensor& result) { return at::digamma_out(result, self); } +Tensor special_psi(const Tensor& self) { return self.digamma(); } +// special_digamma, alias for digamma +Tensor& special_digamma_out(const Tensor& self, Tensor& result) { return at::digamma_out(result, self); } +Tensor special_digamma(const Tensor& self) { return self.digamma(); } + // special_i0, alias for i0 Tensor& special_i0_out(const Tensor& self, Tensor& result) { return at::i0_out(result, self); } Tensor special_i0(const Tensor& self) { return self.i0(); } diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 5e0dd9917dd9f..f0762dfc535c3 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -9426,6 +9426,22 @@ python_module: special variants: function +- func: special_psi(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_psi.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_digamma(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_digamma.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + - func: special_gammaln(Tensor self) -> Tensor python_module: special variants: function diff --git a/docs/source/special.rst b/docs/source/special.rst index 39aa0640c953e..cc173dbc65bad 100644 --- a/docs/source/special.rst +++ b/docs/source/special.rst @@ -26,6 +26,8 @@ Functions .. autofunction:: expm1 .. autofunction:: exp2 .. autofunction:: gammaln +.. autofunction:: digamma +.. autofunction:: psi .. autofunction:: i0 .. autofunction:: i0e .. autofunction:: i1 diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index 4e38a487f0e78..81fe8c007b956 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -2874,29 +2874,8 @@ def merge_dicts(*dicts): add_docstr(torch.digamma, r""" digamma(input, *, out=None) -> Tensor -Computes the logarithmic derivative of the gamma function on `input`. - -.. math:: - \psi(x) = \frac{d}{dx} \ln\left(\Gamma\left(x\right)\right) = \frac{\Gamma'(x)}{\Gamma(x)} -""" + r""" -Args: - input (Tensor): the tensor to compute the digamma function on - -Keyword args: - {out} - -.. note:: This function is similar to SciPy's `scipy.special.digamma`. - -.. note:: From PyTorch 1.8 onwards, the digamma function returns `-Inf` for `0`. - Previously it returned `NaN` for `0`. - -Example:: - - >>> a = torch.tensor([1, 0.5]) - >>> torch.digamma(a) - tensor([-0.5772, -1.9635]) -""".format(**common_args)) - +Alias for :func:`torch.special.digamma`. +""") add_docstr(torch.dist, r""" diff --git a/torch/csrc/api/include/torch/special.h b/torch/csrc/api/include/torch/special.h index d80a43981ae6e..cf667f9412a79 100644 --- a/torch/csrc/api/include/torch/special.h +++ b/torch/csrc/api/include/torch/special.h @@ -21,6 +21,38 @@ inline Tensor& gammaln_out(Tensor& result, const Tensor& self) { return torch::special_gammaln_out(result, self); } +/// Computes the logarithmic derivative of the gamma function on input +/// See https://pytorch.org/docs/master/special.html#torch.special.psi +/// +/// Example: +/// ``` +/// auto t = torch::randn(128, dtype=kDouble); +/// torch::special::psi(t); +/// ``` +inline Tensor psi(const Tensor& self) { + return torch::special_psi(self); +} + +inline Tensor& psi_out(Tensor& result, const Tensor& self) { + return torch::special_psi_out(result, self); +} + +/// Computes the logarithmic derivative of the gamma function on input +/// See https://pytorch.org/docs/master/special.html#torch.special.digamma +/// +/// Example: +/// ``` +/// auto t = torch::randn(128, dtype=kDouble); +/// torch::special::digamma(t); +/// ``` +inline Tensor digamma(const Tensor& self) { + return torch::special_digamma(self); +} + +inline Tensor& digamma_out(Tensor& result, const Tensor& self) { + return torch::special_digamma_out(result, self); +} + /// Computes entropy of input, elementwise /// See https://pytorch.org/docs/master/special.html#torch.special.entr. /// diff --git a/torch/csrc/jit/passes/normalize_ops.cpp b/torch/csrc/jit/passes/normalize_ops.cpp index eda56fe22777e..4e59467d7c138 100644 --- a/torch/csrc/jit/passes/normalize_ops.cpp +++ b/torch/csrc/jit/passes/normalize_ops.cpp @@ -117,6 +117,8 @@ const std::unordered_map& getOperatorAliasMap() { {aten::special_exp2, aten::exp2}, {aten::special_expm1, aten::expm1}, {aten::special_logit, aten::logit}, + {aten::special_digamma, aten::digamma}, + {aten::special_psi, aten::digamma}, {aten::special_i0, aten::i0}, {aten::orgqr, aten::linalg_householder_product}, {aten::special_gammaln, aten::lgamma}}; diff --git a/torch/overrides.py b/torch/overrides.py index 75bde5decb787..aa6876d43ea71 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -873,6 +873,8 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.special.exp2: lambda input: -1, torch.special.expm1: lambda input: -1, torch.special.expit: lambda input: -1, + torch.special.digamma: lambda input: -1, + torch.special.psi: lambda input: -1, torch.special.gammaln: lambda input: -1, torch.special.i0: lambda input: -1, torch.special.i0e: lambda input: -1, diff --git a/torch/special/__init__.py b/torch/special/__init__.py index d0aae87a0b239..68133bbe66f75 100644 --- a/torch/special/__init__.py +++ b/torch/special/__init__.py @@ -35,6 +35,41 @@ tensor([ -inf, 0.0000, 0.3466]) """) +psi = _add_docstr(_special.special_psi, + r""" +psi(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.digamma`. +""") + +digamma = _add_docstr(_special.special_digamma, + r""" +digamma(input, *, out=None) -> Tensor + +Computes the logarithmic derivative of the gamma function on `input`. + +.. math:: + \digamma(x) = \frac{d}{dx} \ln\left(\Gamma\left(x\right)\right) = \frac{\Gamma'(x)}{\Gamma(x)} +""" + r""" +Args: + input (Tensor): the tensor to compute the digamma function on + +Keyword args: + {out} + +.. note:: This function is similar to SciPy's `scipy.special.digamma`. + +.. note:: From PyTorch 1.8 onwards, the digamma function returns `-Inf` for `0`. + Previously it returned `NaN` for `0`. + +Example:: + + >>> a = torch.tensor([1, 0.5]) + >>> torch.special.digamma(a) + tensor([-0.5772, -1.9635]) + +""".format(**common_args)) + gammaln = _add_docstr(_special.special_gammaln, r""" gammaln(input, *, out=None) -> Tensor diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 1cda7f822db50..08e35af018b0a 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -7071,6 +7071,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True), UnaryUfuncInfo('digamma', ref=scipy.special.digamma if TEST_SCIPY else _NOTHING, + aliases=('special.psi', 'special.digamma',), decorators=(precisionOverride({torch.float16: 5e-1}),), dtypes=all_types_and(torch.bool), dtypesIfCUDA=all_types_and(torch.bool, torch.half), From 700add0737351f4db8f5e56d9ae6762007ab0f0d Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 14 Jun 2021 07:21:57 -0700 Subject: [PATCH 077/305] Fix expecttest accept on Python 3.8 and later (#59709) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59709 Fixes #59705. Python 3.8 fixed tracebacks to report the beginning of the line that raised an error, rather than the end. This makes for a simpler implementation (no more string reversing) but need to actually implement. This wasn't caught by tests because we hard coded line numbers to do substitutions, so I also added a little smoketest to detect future changes to traceback line number behavior. Signed-off-by: Edward Z. Yang Test Plan: Imported from OSS Reviewed By: bdhirsh Differential Revision: D28994919 Pulled By: ezyang fbshipit-source-id: 1fb0a782e17c55c13d668fabd04766d2b3811962 --- test/test_expecttest.py | 79 +++++++++++++++++++++++-- torch/testing/_internal/expecttest.py | 85 ++++++++++++++++++++------- 2 files changed, 137 insertions(+), 27 deletions(-) diff --git a/test/test_expecttest.py b/test/test_expecttest.py index 39b6f44136761..dbf24325d9835 100644 --- a/test/test_expecttest.py +++ b/test/test_expecttest.py @@ -5,6 +5,7 @@ import textwrap import doctest from typing import Dict, Any +import traceback import hypothesis from hypothesis.strategies import text, integers, composite, sampled_from, booleans @@ -29,8 +30,8 @@ def nth_line_ref(src, lineno): return len("\n".join(xs)) self.assertEqual(expecttest.nth_line(t, lineno), nth_line_ref(t, lineno)) - @hypothesis.given(text(string.printable), booleans(), sampled_from(['"', "'"])) - def test_replace_string_literal_roundtrip(self, t, raw, quote): + @hypothesis.given(text(string.printable), booleans(), sampled_from(['"', "'"]), booleans()) + def test_replace_string_literal_roundtrip(self, t, raw, quote, lineno_at_start): if raw: hypothesis.assume(expecttest.ok_for_raw_triple_quoted_string(t, quote=quote)) prog = """\ @@ -38,7 +39,8 @@ def test_replace_string_literal_roundtrip(self, t, raw, quote): r2 = {r}{quote}placeholder2{quote} r3 = {r}{quote}placeholder3{quote} """.format(r='r' if raw else '', quote=quote * 3) - new_prog = expecttest.replace_string_literal(textwrap.dedent(prog), 2, t)[0] + new_prog = expecttest.replace_string_literal( + textwrap.dedent(prog), 2, t, lineno_at_start=lineno_at_start)[0] ns : Dict[str, Any] = {} exec(new_prog, ns) msg = "program was:\n{}".format(new_prog) @@ -46,7 +48,7 @@ def test_replace_string_literal_roundtrip(self, t, raw, quote): self.assertEqual(ns['r2'], expecttest.normalize_nl(t), msg=msg) # noqa: F821 self.assertEqual(ns['r3'], 'placeholder3', msg=msg) # noqa: F821 - def test_sample(self): + def test_sample_lineno_at_end(self): prog = r""" single_single('''0''') single_multi('''1''') @@ -76,7 +78,8 @@ def test_sample(self): fn = 'not_a_real_file.py' for lineno, actual in edits: lineno = history.adjust_lineno(fn, lineno) - prog, delta = expecttest.replace_string_literal(prog, lineno, actual) + prog, delta = expecttest.replace_string_literal( + prog, lineno, actual, lineno_at_start=False) history.record_edit(fn, lineno, delta) self.assertExpectedInline(prog, r""" single_single('''a''') @@ -96,6 +99,72 @@ def test_sample(self): ''') """) + def test_sample_lineno_at_start(self): + prog = r""" +single_single('''0''') +single_multi('''1''') +multi_single('''\ +2 +''') +multi_multi_less('''\ +3 +4 +''') +multi_multi_same('''\ +5 +''') +multi_multi_more('''\ +6 +''') +""" + # NB: These are the beginning of the statements + edits = [(2, "a"), + (3, "b\n"), + (4, "c"), + (7, "d\n"), + (11, "e\n"), + (14, "f\ng\n")] + history = expecttest.EditHistory() + fn = 'not_a_real_file.py' + for lineno, actual in edits: + lineno = history.adjust_lineno(fn, lineno) + prog, delta = expecttest.replace_string_literal( + prog, lineno, actual, lineno_at_start=True) + history.record_edit(fn, lineno, delta) + self.assertExpectedInline(prog, r""" +single_single('''a''') +single_multi('''\ +b +''') +multi_single('''c''') +multi_multi_less('''\ +d +''') +multi_multi_same('''\ +e +''') +multi_multi_more('''\ +f +g +''') +""") + + def test_lineno_assumptions(self): + def get_tb(s): + return traceback.extract_stack(limit=2) + + tb1 = get_tb("") + tb2 = get_tb("""a +b +c""") + + if expecttest.LINENO_AT_START: + # tb2's stack starts on the next line + self.assertEqual(tb1[0].lineno + 1, tb2[0].lineno) + else: + # starts at the end here + self.assertEqual(tb1[0].lineno + 1 + 2, tb2[0].lineno) + def load_tests(loader, tests, ignore): tests.addTests(doctest.DocTestSuite(expecttest)) diff --git a/torch/testing/_internal/expecttest.py b/torch/testing/_internal/expecttest.py index 6e1f2f01b5c80..7fff670fe9269 100644 --- a/torch/testing/_internal/expecttest.py +++ b/torch/testing/_internal/expecttest.py @@ -3,6 +3,7 @@ import traceback import os import string +import sys from typing import Tuple @@ -51,6 +52,8 @@ ACCEPT = os.getenv('EXPECTTEST_ACCEPT') +LINENO_AT_START = sys.version_info >= (3, 8) + def nth_line(src, lineno): """ @@ -132,16 +135,29 @@ def ok_for_raw_triple_quoted_string(s, quote): return quote * 3 not in s and (not s or s[-1] not in [quote, '\\']) +RE_EXPECT = re.compile( + ( + r"^(?P[^\n]*?)" + r"(?Pr?)" + r"(?P'''|" r'""")' + r"(?P.*?)" + r"(?P=quote)" + ), + re.DOTALL +) + + # This operates on the REVERSED string (that's why suffix is first) -RE_EXPECT = re.compile(r"^(?P[^\n]*?)" - r"(?P'''|" r'""")' - r"(?P.*?)" - r"(?P=quote)" - r"(?Pr?)", re.DOTALL) +RE_REVERSED_EXPECT = \ + re.compile(r"^(?P[^\n]*?)" + r"(?P'''|" r'""")' + r"(?P.*?)" + r"(?P=quote)" + r"(?Pr?)", re.DOTALL) def replace_string_literal(src : str, lineno : int, - new_string : str) -> Tuple[str, int]: + new_string : str, *, lineno_at_start: bool) -> Tuple[str, int]: r""" Replace a triple quoted string literal with new contents. Only handles printable ASCII correctly at the moment. This @@ -152,9 +168,9 @@ def replace_string_literal(src : str, lineno : int, Returns a tuple of the replaced string, as well as a delta of number of lines added/removed. - >>> replace_string_literal("'''arf'''", 1, "barf") + >>> replace_string_literal("'''arf'''", 1, "barf", lineno_at_start=False) ("'''barf'''", 0) - >>> r = replace_string_literal(" moo = '''arf'''", 1, "'a'\n\\b\n") + >>> r = replace_string_literal(" moo = '''arf'''", 1, "'a'\n\\b\n", lineno_at_start=False) >>> print(r[0]) moo = '''\ 'a' @@ -162,21 +178,21 @@ def replace_string_literal(src : str, lineno : int, ''' >>> r[1] 3 - >>> replace_string_literal(" moo = '''\\\narf'''", 2, "'a'\n\\b\n")[1] + >>> replace_string_literal(" moo = '''\\\narf'''", 2, "'a'\n\\b\n", lineno_at_start=False)[1] 2 - >>> print(replace_string_literal(" f('''\"\"\"''')", 1, "a ''' b")[0]) + >>> print(replace_string_literal(" f('''\"\"\"''')", 1, "a ''' b", lineno_at_start=False)[0]) f('''a \'\'\' b''') """ # Haven't implemented correct escaping for non-printable characters assert all(c in string.printable for c in new_string) - i = nth_eol(src, lineno) + new_string = normalize_nl(new_string) delta = [new_string.count("\n")] if delta[0] > 0: delta[0] += 1 # handle the extra \\\n - def replace(m): + def compute_raw_new_body_and_adjust_delta(m): s = new_string raw = m.group('raw') == 'r' if not raw or not ok_for_raw_triple_quoted_string(s, quote=m.group('quote')[0]): @@ -189,17 +205,39 @@ def replace(m): new_body = "\\\n" + s if "\n" in s and not raw else s delta[0] -= m.group('body').count("\n") + return raw, new_body + + if lineno_at_start: + i = nth_line(src, lineno) + + # i points to the start of the string + def replace(m): + raw, new_body = compute_raw_new_body_and_adjust_delta(m) + return ''.join([m.group('prefix'), + 'r' if raw else '', + m.group('quote'), + new_body, + m.group('quote'), + ]) + + return (src[:i] + RE_EXPECT.sub(replace, src[i:], count=1), delta[0]) + else: + i = nth_eol(src, lineno) - return ''.join([m.group('suffix'), - m.group('quote'), - new_body[::-1], - m.group('quote'), - 'r' if raw else '', - ]) + # i points to the END of the string. Do some funny + # business with reversing the string to do the replace + def replace(m): + raw, new_body = compute_raw_new_body_and_adjust_delta(m) + return ''.join([m.group('suffix'), + m.group('quote'), + new_body[::-1], + m.group('quote'), + 'r' if raw else '', + ]) - # Having to do this in reverse is very irritating, but it's the - # only way to make the non-greedy matches work correctly. - return (RE_EXPECT.sub(replace, src[:i][::-1], count=1)[::-1] + src[i:], delta[0]) + # Having to do this in reverse is very irritating, but it's the + # only way to make the non-greedy matches work correctly. + return (RE_REVERSED_EXPECT.sub(replace, src[:i][::-1], count=1)[::-1] + src[i:], delta[0]) class TestCase(unittest.TestCase): @@ -228,7 +266,10 @@ def assertExpectedInline(self, actual, expect, skip=0): # compute the change in lineno lineno = EDIT_HISTORY.adjust_lineno(fn, lineno) - new, delta = replace_string_literal(old, lineno, actual) + new, delta = replace_string_literal( + old, lineno, actual, + lineno_at_start=LINENO_AT_START + ) assert old != new, f"Failed to substitute string at {fn}:{lineno}; did you use triple quotes?" From d60d81b5a79885514b32b7727c96df68e642ff82 Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Mon, 14 Jun 2021 07:21:57 -0700 Subject: [PATCH 078/305] Make PyObject_FastGetAttrString accept const char* (#59758) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59758 The underlying call to tp_getattr is const safe but CPython has not fixed it due to BC problems. No reason not to advertise the better type here though! Signed-off-by: Edward Z. Yang Test Plan: Imported from OSS Reviewed By: albanD Differential Revision: D29017911 Pulled By: ezyang fbshipit-source-id: 8d55983fe6416c03eb69c6367bcc431c30000133 --- torch/csrc/utils/python_strings.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/torch/csrc/utils/python_strings.h b/torch/csrc/utils/python_strings.h index d373c7625f1c7..9162d02e6ebe2 100644 --- a/torch/csrc/utils/python_strings.h +++ b/torch/csrc/utils/python_strings.h @@ -102,14 +102,15 @@ inline void THPUtils_internStringInPlace(PyObject** obj) { */ // NOLINTNEXTLINE(clang-diagnostic-unused-function) -static py::object PyObject_FastGetAttrString(PyObject *obj, char *name) +static py::object PyObject_FastGetAttrString(PyObject *obj, const char *name) { PyTypeObject *tp = Py_TYPE(obj); PyObject *res = (PyObject *)nullptr; /* Attribute referenced by (char *)name */ if (tp->tp_getattr != nullptr) { - res = (*tp->tp_getattr)(obj, name); + // This is OK per https://bugs.python.org/issue39620 + res = (*tp->tp_getattr)(obj, const_cast(name)); if (res == nullptr) { PyErr_Clear(); } From 68d690ffbd64d0fb697dc3da1635216366649787 Mon Sep 17 00:00:00 2001 From: leslie-fang-intel Date: Mon, 14 Jun 2021 07:52:52 -0700 Subject: [PATCH 079/305] Vectorize the softmax calculation when not along the last dim (#59195) Summary: Currently, if we do softmax which are not along the last dim, the calculation will fall to a [scalar version](https://github.com/pytorch/pytorch/blob/d417a094f398f1c4efd7f818b14b8471a597fbcc/aten/src/ATen/native/SoftMax.cpp#L14-L64). And we find actually we have the chance to vectorize the calculation along the inner_size dim. Changes we made: - Use vectorized softmax_kernel instead of host_softmax when not along the last dim. Performance data on 28 cores' Intel 8280 CPU when the Input size is [32, 81, 15130] and do softmax along the second dim(81). - FP32 Baseline: 24.67 ms - FP32 optimized: 9.2 ms Pull Request resolved: https://github.com/pytorch/pytorch/pull/59195 Reviewed By: ailzhang Differential Revision: D28854796 Pulled By: cpuhrsch fbshipit-source-id: 18477acc3963754c59009b1794f080496ae16c3d --- aten/src/ATen/native/SoftMax.cpp | 8 +- aten/src/ATen/native/cpu/SoftMaxKernel.cpp | 116 +++++++++++++++++++++ aten/src/ATen/native/cpu/SoftmaxKernel.h | 3 + 3 files changed, 123 insertions(+), 4 deletions(-) diff --git a/aten/src/ATen/native/SoftMax.cpp b/aten/src/ATen/native/SoftMax.cpp index ac2a91d9cd54c..d1258a82a2326 100644 --- a/aten/src/ATen/native/SoftMax.cpp +++ b/aten/src/ATen/native/SoftMax.cpp @@ -132,7 +132,7 @@ Tensor softmax_cpu(const Tensor& input_, const int64_t dim_, const bool half_to_ if (input.numel() == 0) { return output; } - if (input.dim() == 0) + if (input.dim() == 0) input = input.view(1); TORCH_CHECK( dim >= 0 && dim < input.dim(), @@ -140,9 +140,7 @@ Tensor softmax_cpu(const Tensor& input_, const int64_t dim_, const bool half_to_ if (input.ndimension() > 0 && dim == input.ndimension() - 1) { softmax_lastdim_kernel(kCPU, output, input); } else { - AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "softmax", [&] { - host_softmax(output, input, dim); - }); + softmax_kernel(kCPU, output, input, dim); } return output; } @@ -310,6 +308,8 @@ DEFINE_DISPATCH(softmax_backward_lastdim_kernel); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(log_softmax_backward_lastdim_kernel); +DEFINE_DISPATCH(softmax_kernel); + Tensor softmax(const Tensor& self, Dimname dim, optional dtype) { return at::softmax(self, dimname_to_position(self, dim), dtype); } diff --git a/aten/src/ATen/native/cpu/SoftMaxKernel.cpp b/aten/src/ATen/native/cpu/SoftMaxKernel.cpp index 61442c5637e59..eaa895d33bc07 100644 --- a/aten/src/ATen/native/cpu/SoftMaxKernel.cpp +++ b/aten/src/ATen/native/cpu/SoftMaxKernel.cpp @@ -10,6 +10,7 @@ #include #include +#include // [Note AVX-SSE transitions] In general we avoid calls into cmath for code // compiled with AVX/AVX2 This is because of SSE-AVX transitions and a bug in // Glibc2.23 See https://bugs.launchpad.net/ubuntu/+source/glibc/+bug/1663280 @@ -206,6 +207,113 @@ struct vec_host_softmax_lastdim { } }; +template +inline void _vec_softmax( + scalar_t* input_data_base, + scalar_t* output_data_base, + int64_t outer_size, + int64_t inner_size, + int64_t dim_size) { + using Vec = vec::Vectorized; + int64_t dim_stride = inner_size; + int64_t outer_stride = dim_size * dim_stride; + int64_t grain_size = std::min(internal::GRAIN_SIZE / dim_size, (int64_t)1); + int vectorized_step = Vec().size(); // Currently, we only support scalar_t with double or float32 + TORCH_CHECK( + (vectorized_step == 8) || (vectorized_step == 4), + "vectorized_step must be 8 with dtype float or 4 with dtype double"); + parallel_for( + 0, outer_size * inner_size, grain_size, [&](int64_t begin, int64_t end) { + int64_t idx = begin; + while (idx < end) { + int64_t outer_idx = idx / inner_size; + int64_t inner_idx = idx % inner_size; + if (((inner_idx + vectorized_step) <= inner_size) && ((idx + vectorized_step) <= end)) { + // Vectorization + scalar_t* input_data = + input_data_base + outer_idx * outer_stride + inner_idx; + scalar_t* output_data = + output_data_base + outer_idx * outer_stride + inner_idx; + // Step 1: Get max Score + Vec max_m256 = Vec::loadu(input_data); + for (int64_t d = 1; d < dim_size; d += 1) { + Vec input_m256 = Vec::loadu(input_data + d * dim_stride); + max_m256 = vec::maximum(max_m256, input_m256); + } + // Step2: Calculate sum + Vec sum_m256 = Vec(0.0); + for (int64_t d = 0; d < dim_size; d += 1) { + Vec output_m256 = + (Vec::loadu(input_data + d * dim_stride) - max_m256).exp(); + output_m256.store(output_data + d * dim_stride); + sum_m256 = sum_m256 + output_m256; + } + // Step3: Unify + for (int64_t d = 0; d < dim_size; d += 1) { + Vec output_m256 = + Vec::loadu(output_data + d * dim_stride) / sum_m256; + output_m256.store(output_data + d * dim_stride); + } + idx += vectorized_step; + } else { + // Tail case(Scalar): it is exactly same logic as host_softmax + // inside aten/src/ATen/native/SoftMax.cpp. There are 2 kind of + // cases which will fall through this part: + // Case 1: For the idx at the end of total chunk for each thread, there are not enough numbers for parallization. + // Case 2: For the idx at the end of each inner_size inside thread, there are not enough numbers for parallization. + int64_t tail_number = ((idx+vectorized_step) > end) ? /*Case1*/ (end - idx) : /*Case2*/ (inner_size - inner_idx); + for (int64_t i=0; i < tail_number; i++) { + outer_idx = (idx + i) / inner_size; + inner_idx = (idx + i) % inner_size; + scalar_t* input_data = + input_data_base + outer_idx * outer_stride + inner_idx; + scalar_t* output_data = + output_data_base + outer_idx * outer_stride + inner_idx; + // Step1: Get max score + scalar_t max_input = input_data[0]; + for (int64_t d = 1; d < dim_size; d += 1) { + max_input = std::max(max_input, input_data[d * dim_stride]); + } + // Step2: Calculate the Sum + scalar_t sum_data = 0; + for (int64_t d = 0; d < dim_size; d += 1) { + output_data[d * dim_stride] = + std::exp(input_data[d * dim_stride] - max_input); + sum_data += output_data[d * dim_stride]; + } + // Step3: Unify + for (int64_t d = 0; d < dim_size; d += 1) { + output_data[d * dim_stride] = + output_data[d * dim_stride]/sum_data; + } + } + idx += tail_number; + } + } + }); +} + +template +struct vec_softmax { + static void apply(Tensor& output, const Tensor& input, int64_t dim) { + int64_t outer_size = 1; + int64_t dim_size = input.size(dim); + int64_t inner_size = 1; + for (int64_t i = 0; i < dim; ++i) + outer_size *= input.size(i); + for (int64_t i = dim + 1; i < input.dim(); ++i) + inner_size *= input.size(i); + scalar_t* input_data_base = input.data_ptr(); + scalar_t* output_data_base = output.data_ptr(); + if (LogSoftMax) { + AT_ERROR("vec_softmax not implemented for LogSoftMax"); + } else { + _vec_softmax( + input_data_base, output_data_base, outer_size, inner_size, dim_size); + } + } +}; + template struct vec_host_softmax_backward_lastdim { static void @@ -232,6 +340,12 @@ static void softmax_lastdim_kernel_impl(Tensor& result, const Tensor& self) { }); } +static void softmax_kernel_impl(Tensor& result, const Tensor& self, int64_t dim) { + AT_DISPATCH_FLOATING_TYPES(self.scalar_type(), "softmax_kernel_impl", [&] { + vec_softmax::apply(result, self, dim); + }); +} + static void log_softmax_lastdim_kernel_impl( Tensor& result, const Tensor& self) { @@ -279,4 +393,6 @@ REGISTER_DISPATCH( log_softmax_backward_lastdim_kernel, &log_softmax_backward_lastdim_kernel_impl); +REGISTER_DISPATCH(softmax_kernel, &softmax_kernel_impl); + }} // namespace at::native diff --git a/aten/src/ATen/native/cpu/SoftmaxKernel.h b/aten/src/ATen/native/cpu/SoftmaxKernel.h index 0fb2a8e18a5ff..9490d6f5f90d3 100644 --- a/aten/src/ATen/native/cpu/SoftmaxKernel.h +++ b/aten/src/ATen/native/cpu/SoftmaxKernel.h @@ -14,5 +14,8 @@ DECLARE_DISPATCH(forward_fn, log_softmax_lastdim_kernel); DECLARE_DISPATCH(backward_fn, softmax_backward_lastdim_kernel); DECLARE_DISPATCH(backward_fn, log_softmax_backward_lastdim_kernel); +using forward_fn_with_dim = void(*)(Tensor &, const Tensor &, const int64_t); +DECLARE_DISPATCH(forward_fn_with_dim, softmax_kernel); + } } From 3d90c82a5c7ed1e769e29a92d81816dc91e76461 Mon Sep 17 00:00:00 2001 From: Jason Ansel Date: Mon, 14 Jun 2021 09:29:34 -0700 Subject: [PATCH 080/305] [TensorExpr] Python binding improvements (#59920) Summary: Some minor quality of life improvements for the NNC python bindings: - expose `call_raw()` - support passing integers to `call()` (for dynamic shapes) - implicit conversions to cleanup `[BufferArg(x) for x in [A, B, C]]` into just `[A, B, C]` - don't silently default to "ir_eval" for unknown mode (e.g. "LLVM") Pull Request resolved: https://github.com/pytorch/pytorch/pull/59920 Reviewed By: ZolotukhinM Differential Revision: D29090904 Pulled By: jansel fbshipit-source-id: 154ace82725ae2046cfe2e6eb324fd37f5d209a7 --- test/test_tensorexpr_pybind.py | 104 +++++++++++++----- torch/csrc/jit/tensorexpr/tensorexpr_init.cpp | 33 +++++- 2 files changed, 105 insertions(+), 32 deletions(-) diff --git a/test/test_tensorexpr_pybind.py b/test/test_tensorexpr_pybind.py index 3515671781ea9..fceec6ecd9d15 100644 --- a/test/test_tensorexpr_pybind.py +++ b/test/test_tensorexpr_pybind.py @@ -1,5 +1,6 @@ import torch import numpy as np +import torch._C._te as te from torch.testing._internal.common_utils import run_tests from torch.testing._internal.jit_utils import JitTestCase @@ -7,6 +8,7 @@ LLVM_ENABLED = torch._C._llvm_enabled() + class kernel_arena_scope(object): def __enter__(self): self.scope = torch._C._te.KernelScope() @@ -14,48 +16,62 @@ def __enter__(self): def __exit__(self, typ, val, traceback): self.scope = None -class TestTensorExprPyBind(JitTestCase): - def test_simple_sum(self): - with kernel_arena_scope(): - dtype = torch._C._te.Dtype.Float - N = 32 - dN = torch._C._te.ExprHandle.int(N) - A = torch._C._te.Placeholder('A', dtype, [dN]) - B = torch._C._te.Placeholder('B', dtype, [dN]) +def construct_adder(n: int, dtype=te.Dtype.Float): + dN = te.ExprHandle.int(n) + A = te.Placeholder('A', dtype, [dN]) + B = te.Placeholder('B', dtype, [dN]) - def compute(i): - return A.load([i]) + B.load([i]) - C = torch._C._te.Compute('C', [torch._C._te.DimArg(dN, 'i')], compute) + def compute(i): + return A.load([i]) + B.load([i]) - loopnest = torch._C._te.LoopNest([C]) - loopnest.prepare_for_codegen() - stmt = torch._C._te.simplify(loopnest.root_stmt()) + C = te.Compute('C', [te.DimArg(dN, 'i')], compute) + + loopnest = te.LoopNest([C]) + loopnest.prepare_for_codegen() + stmt = te.simplify(loopnest.root_stmt()) - cg = torch._C._te.construct_codegen('ir_eval', stmt, [torch._C._te.BufferArg(x) for x in [A, B, C]]) + return te.construct_codegen('ir_eval', stmt, [A, B, C]) - tA = torch.rand(N) * 5 - tB = torch.rand(N) * 6 - tC = torch.empty(N) + +class TestTensorExprPyBind(JitTestCase): + def test_simple_sum(self): + with kernel_arena_scope(): + n = 32 + cg = construct_adder(n) + + tA = torch.randn(n) + tB = torch.randn(n) + tC = torch.empty(n) cg.call([tA, tB, tC]) torch.testing.assert_allclose(tA + tB, tC) + def test_call_raw(self): + with kernel_arena_scope(): + n = 16 + cg = construct_adder(n, dtype=te.Dtype.Double) + + tA = torch.randn(n, dtype=torch.float64) + tB = torch.randn(n, dtype=torch.float64) + tC = torch.empty(n, dtype=torch.float64) + cg.call_raw([tA.data_ptr(), tB.data_ptr(), tC.data_ptr()]) + torch.testing.assert_allclose(tA + tB, tC) + def test_external_calls(self): with kernel_arena_scope(): - dtype = torch._C._te.Dtype.Float + dtype = te.Dtype.Float - ZERO = torch._C._te.ExprHandle.int(0) - ONE = torch._C._te.ExprHandle.int(1) - FOUR = torch._C._te.ExprHandle.int(4) - A = torch._C._te.BufHandle('A', [ONE, FOUR], dtype) - B = torch._C._te.BufHandle('B', [FOUR, ONE], dtype) - C = torch._C._te.BufHandle('C', [ONE, ONE], dtype) + ONE = te.ExprHandle.int(1) + FOUR = te.ExprHandle.int(4) + A = te.BufHandle('A', [ONE, FOUR], dtype) + B = te.BufHandle('B', [FOUR, ONE], dtype) + C = te.BufHandle('C', [ONE, ONE], dtype) - s = torch._C._te.ExternalCall(C, "nnc_aten_matmul", [A, B], []) + s = te.ExternalCall(C, "nnc_aten_matmul", [A, B], []) - loopnest = torch._C._te.LoopNest(s, [C]) + loopnest = te.LoopNest(s, [C]) loopnest.prepare_for_codegen() - codegen = torch._C._te.construct_codegen('ir_eval', s, [torch._C._te.BufferArg(x) for x in [A, B, C]]) + codegen = te.construct_codegen('ir_eval', s, [te.BufferArg(x) for x in [A, B, C]]) tA = torch.ones(1, 4) tB = torch.ones(4, 1) @@ -63,10 +79,41 @@ def test_external_calls(self): codegen.call([tA, tB, tC]) torch.testing.assert_allclose(torch.matmul(tA, tB), tC) + def test_dynamic_shape(self): + with kernel_arena_scope(): + dN = te.VarHandle("n", te.Dtype.Int) + A = te.Placeholder('A', te.Dtype.Double, [dN]) + B = te.Placeholder('B', te.Dtype.Double, [dN]) + + def compute(i): + return A.load([i]) - B.load([i]) + + C = te.Compute('C', [te.DimArg(dN, 'i')], compute) + + loopnest = te.LoopNest([C]) + loopnest.prepare_for_codegen() + stmt = te.simplify(loopnest.root_stmt()) + + cg = te.construct_codegen( + 'ir_eval', + stmt, + [A, B, C, dN]) + + def test_with_shape(n): + tA = torch.randn(n, dtype=torch.double) + tB = torch.randn(n, dtype=torch.double) + tC = torch.empty(n, dtype=torch.double) + cg.call([tA, tB, tC, n]) + torch.testing.assert_allclose(tA - tB, tC) + + test_with_shape(8) + test_with_shape(31) + @unittest.skipIf(not LLVM_ENABLED, "LLVM backend not enabled") def test_kernel_with_tensor_inputs(self): def f(a, b, c): return a + b + c + device, size = 'cpu', (4, 4) x = torch.rand(size, device=device) y = torch.rand(size, device=device) @@ -94,6 +141,7 @@ def f(a, b, c): def test_kernel_with_scalar_inputs(self): def f(a, b, c): return a + b + c + x = torch.tensor(0.1, dtype=torch.float, device='cpu') y = torch.tensor(0.6, dtype=torch.float, device='cpu') z = torch.tensor(0.7, dtype=torch.float, device='cpu') diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp index d1aad75b5a571..f73fe02a3b3b8 100644 --- a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp +++ b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp @@ -47,6 +47,7 @@ ArgValue convertPyToArgValue(py::handle inp) { throw std::runtime_error("nyi"); } } + void initTensorExprBindings(PyObject* module) { auto m = py::handle(module).cast(); @@ -653,14 +654,30 @@ void initTensorExprBindings(PyObject* module) { py::class_(te, "CodeGen") .def( "call", - [](CodeGen& self, const std::vector& values) { + [](CodeGen& self, const py::sequence& values) { std::vector value_ptrs; - value_ptrs.reserve(values.size()); + value_ptrs.reserve(py::len(values)); for (const auto& value : values) { - value_ptrs.emplace_back(CodeGen::CallArg(value.data_ptr())); + if (py::isinstance(value)) { + value_ptrs.emplace_back(value.cast()); + } else { + value_ptrs.emplace_back(value.cast().data_ptr()); + } } self.call(value_ptrs); }) + .def( + "call_raw", + [](CodeGen& self, const py::sequence& values) { + std::vector value_ptrs; + value_ptrs.reserve(py::len(values)); + for (const auto& value : values) { + // Tensor.data_ptr() returns an int in python + value_ptrs.emplace_back( + reinterpret_cast(value.cast())); + } + self.call_raw(value_ptrs); + }) .def( "get_code_text", [](CodeGen& self, const std::string& attr = "") { @@ -678,6 +695,11 @@ void initTensorExprBindings(PyObject* module) { .def(py::init()) .def(py::init()); + py::implicitly_convertible(); + py::implicitly_convertible(); + py::implicitly_convertible(); + py::implicitly_convertible(); + te.def( "construct_codegen", [](const std::string& name, @@ -696,8 +718,11 @@ void initTensorExprBindings(PyObject* module) { #else throw std::runtime_error("PyTorch not compiled with CUDA support!"); #endif - } else { + } else if (name == "ir_eval") { cg = new SimpleIREvaluator(stmt, args); + } else { + throw std::runtime_error( + "construct_codegen() expects 'llvm', 'cuda', or 'ir_eval'"); } return cg; }); From 580a20f33b530bdbfdc6af1dc7908403144eb92f Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Mon, 14 Jun 2021 09:51:39 -0700 Subject: [PATCH 081/305] [reland] torch/lib/c10d: Use torch_check instead of throwing runtime_error (#59918) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59918 Reland of https://github.com/pytorch/pytorch/pull/59684 ghstack-source-id: 131303057 Test Plan: ci Reviewed By: cbalioglu Differential Revision: D29081452 fbshipit-source-id: 419df79341f702e796f7adf5f1071a6cd1dcd8d1 --- torch/lib/c10d/FileStore.cpp | 6 +- torch/lib/c10d/GlooDeviceFactory.cpp | 4 +- torch/lib/c10d/NCCLUtils.hpp | 5 +- torch/lib/c10d/ProcessGroup.cpp | 8 +-- torch/lib/c10d/ProcessGroup.hpp | 12 ++-- torch/lib/c10d/ProcessGroupGloo.cpp | 42 ++++++------ torch/lib/c10d/ProcessGroupMPI.cpp | 46 ++++++------- torch/lib/c10d/ProcessGroupNCCL.cpp | 66 +++++++++---------- torch/lib/c10d/ProcessGroupRoundRobin.cpp | 10 +-- torch/lib/c10d/TCPStore.cpp | 10 +-- torch/lib/c10d/UnixSockUtils.hpp | 2 +- torch/lib/c10d/Utils.cpp | 8 +-- torch/lib/c10d/Utils.hpp | 16 ++--- torch/lib/c10d/WinSockUtils.hpp | 2 +- torch/lib/c10d/frontend.cpp | 2 +- torch/lib/c10d/test/ProcessGroupGlooTest.cpp | 2 +- torch/lib/c10d/test/ProcessGroupMPITest.cpp | 22 +++---- .../c10d/test/ProcessGroupNCCLErrorsTest.cpp | 2 +- torch/lib/c10d/test/TCPStoreTest.cpp | 9 +-- 19 files changed, 138 insertions(+), 136 deletions(-) diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp index ea98963ee9df8..73342272c54c0 100644 --- a/torch/lib/c10d/FileStore.cpp +++ b/torch/lib/c10d/FileStore.cpp @@ -273,7 +273,7 @@ FileStore::FileStore(const std::string& path, int numWorkers) cleanupKey_("cleanup/"), regularPrefix_("/") { if (numWorkers_ < 1) { - throw std::runtime_error( + TORCH_CHECK(false, "Number of workers for FileStore should be greater than zero"); } } @@ -341,7 +341,7 @@ std::vector FileStore::get(const std::string& key) { const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout_ != kNoTimeout && elapsed > timeout_) { - throw std::runtime_error("Timeout waiting for key: " + key); + TORCH_CHECK(false, "Timeout waiting for key: " + key); } std::this_thread::sleep_for(std::chrono::milliseconds(10)); continue; @@ -424,7 +424,7 @@ void FileStore::wait( const auto elapsed = std::chrono::duration_cast( std::chrono::steady_clock::now() - start); if (timeout != kNoTimeout && elapsed > timeout) { - throw std::runtime_error("Wait timeout"); + TORCH_CHECK(false, "Wait timeout"); } /* sleep override */ diff --git a/torch/lib/c10d/GlooDeviceFactory.cpp b/torch/lib/c10d/GlooDeviceFactory.cpp index 416676483e182..cb83a99838520 100644 --- a/torch/lib/c10d/GlooDeviceFactory.cpp +++ b/torch/lib/c10d/GlooDeviceFactory.cpp @@ -149,7 +149,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForInterface(const std::string& interfaceName) { auto device = makeGlooDevice(interfaceName, ""); if (!device) { - throw std::runtime_error("makeDeviceForInterface(): unsupported gloo device"); + TORCH_CHECK(false, "makeDeviceForInterface(): unsupported gloo device"); } return device; } @@ -158,7 +158,7 @@ std::shared_ptr<::gloo::transport::Device> GlooDeviceFactory:: makeDeviceForHostname(const std::string& hostname) { auto device = makeGlooDevice("", hostname); if (!device) { - throw std::runtime_error("makeDeviceForHostname(): unsupported gloo device"); + TORCH_CHECK(false, "makeDeviceForHostname(): unsupported gloo device"); } return device; } diff --git a/torch/lib/c10d/NCCLUtils.hpp b/torch/lib/c10d/NCCLUtils.hpp index 0dec4573112a1..e3ee14da0f542 100644 --- a/torch/lib/c10d/NCCLUtils.hpp +++ b/torch/lib/c10d/NCCLUtils.hpp @@ -9,6 +9,7 @@ #include #include +#include namespace { // Provides additional detail into NCCL error codes based on when these are @@ -57,7 +58,7 @@ const inline char* getNcclErrorDetailStr(ncclResult_t error) { std::string err = "NCCL error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + ", " + ncclGetErrorWithVersion(result) + \ "\n" + getNcclErrorDetailStr(result); \ - throw std::runtime_error(err); \ + TORCH_CHECK(false, err); \ } \ } while (0) @@ -142,7 +143,7 @@ class NCCLComm { ncclComm_t getNcclComm() { std::unique_lock lock(mutex_); if (aborted_) { - throw std::runtime_error( + TORCH_CHECK(false, "NCCL communicator was aborted on rank " + std::to_string(rank_) + "."); } diff --git a/torch/lib/c10d/ProcessGroup.cpp b/torch/lib/c10d/ProcessGroup.cpp index 39ae2bf71c598..4e03824eb12da 100644 --- a/torch/lib/c10d/ProcessGroup.cpp +++ b/torch/lib/c10d/ProcessGroup.cpp @@ -107,13 +107,13 @@ std::exception_ptr ProcessGroup::Work::exception() const { } int ProcessGroup::Work::sourceRank() const { - throw std::runtime_error( + TORCH_CHECK(false, "sourceRank() may only be called on work objects " "that correspond to a recv or recv-from-any call."); } std::vector ProcessGroup::Work::result() { - throw std::runtime_error("result() not implemented."); + TORCH_CHECK(false, "result() not implemented."); } void ProcessGroup::Work::synchronize() {} @@ -129,7 +129,7 @@ bool ProcessGroup::Work::wait(std::chrono::milliseconds timeout) { if (!completed_) { // Throw exception if the wait operation timed out and the work was not // completed. - throw std::runtime_error("Operation timed out!"); + TORCH_CHECK(false, "Operation timed out!"); } } if (exception_) { @@ -186,7 +186,7 @@ c10::intrusive_ptr ProcessGroup::allgather_coalesced( std::vector>& /* usused */, std::vector& /* usused */, const AllgatherOptions& /* usused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for allgather_coalesced in this process group"); } diff --git a/torch/lib/c10d/ProcessGroup.hpp b/torch/lib/c10d/ProcessGroup.hpp index ee2990fd33975..3a3ffa6b95d67 100644 --- a/torch/lib/c10d/ProcessGroup.hpp +++ b/torch/lib/c10d/ProcessGroup.hpp @@ -264,7 +264,7 @@ class ProcessGroup : public torch::CustomClassHolder { at::Tensor&, at::Tensor&, const ReduceScatterOptions& opts = ReduceScatterOptions()) { - throw std::runtime_error("ProcessGroup does not support reduce_scatter_base"); + TORCH_CHECK(false, "ProcessGroup does not support reduce_scatter_base"); } @@ -274,20 +274,20 @@ class ProcessGroup : public torch::CustomClassHolder { std::vector& outputSplitSizes, std::vector& inputSplitSizes, const AllToAllOptions& opts = AllToAllOptions()) { - throw std::runtime_error("ProcessGroup does not support alltoall"); + TORCH_CHECK(false, "ProcessGroup does not support alltoall"); } virtual c10::intrusive_ptr alltoall( std::vector& outputTensors, std::vector& inputTensors, const AllToAllOptions& opts = AllToAllOptions()) { - throw std::runtime_error("ProcessGroup does not support alltoall"); + TORCH_CHECK(false, "ProcessGroup does not support alltoall"); } virtual void monitoredBarrier( const BarrierOptions& /* unused */, bool /* unused */ = false ) { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not support monitoredBarrier, only GLOO supports monitored barrier.") @@ -299,7 +299,7 @@ class ProcessGroup : public torch::CustomClassHolder { // for GLOO and NCCL backends currently. virtual void setSequenceNumberForGroup() { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") @@ -311,7 +311,7 @@ class ProcessGroup : public torch::CustomClassHolder { // may indicate that there is some sort of collective desynchronization. virtual uint64_t getSequenceNumberForGroup() { auto backendName = getBackendName(); - throw std::runtime_error( + TORCH_CHECK(false, c10::str("ProcessGroup ", backendName, " does not yet support sequence numbers.") diff --git a/torch/lib/c10d/ProcessGroupGloo.cpp b/torch/lib/c10d/ProcessGroupGloo.cpp index d423271192db8..98164237feb9c 100644 --- a/torch/lib/c10d/ProcessGroupGloo.cpp +++ b/torch/lib/c10d/ProcessGroupGloo.cpp @@ -66,7 +66,7 @@ func(__VA_ARGS__); \ break; \ default: \ - throw std::runtime_error("Invalid scalar type"); \ + TORCH_CHECK(false, "Invalid scalar type"); \ } #define HOST_NAME_MAX 256 @@ -95,7 +95,7 @@ func(args); \ break; \ default: \ - throw std::runtime_error("Invalid scalar type"); \ + TORCH_CHECK(false, "Invalid scalar type"); \ } #endif @@ -178,22 +178,22 @@ ReduceFunc toFunction(const ReduceOp& r) { case ReduceOp::MAX: return ReduceFunc(&::gloo::max); case ReduceOp::BAND: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BAND with non-integral dtype"); break; case ReduceOp::BOR: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BOR with non-integral dtype"); break; case ReduceOp::BXOR: - throw std::runtime_error( + TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with non-integral dtype"); break; case ReduceOp::UNUSED: break; } - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); } // Bitwise AND with SFINAE guard for integral types. @@ -258,7 +258,7 @@ ReduceFunc toFunction(const ReduceOp& r) { break; } - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); } template @@ -368,7 +368,7 @@ void initializeStreamsEvents( const auto device_id = tensorgroup[0].device().index(); for (const auto& tensor : tensorgroup) { if (tensor.device().index() != device_id) { - throw std::runtime_error( + TORCH_CHECK(false, "tensors in the nested tensor vectors need to " "be on the same device"); } @@ -683,7 +683,7 @@ ProcessGroupGloo::ProcessGroupGloo( collectiveCounter_(0) { auto& devices = options->devices; if (devices.empty()) { - throw std::runtime_error("No device(s) specified"); + TORCH_CHECK(false, "No device(s) specified"); } // Create and connect a context for every device. @@ -915,7 +915,7 @@ c10::intrusive_ptr ProcessGroupGloo::broadcast( work = c10::make_intrusive( std::move(context), inputs, opts.rootRank, opts.rootTensor, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); @@ -1426,7 +1426,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce( invalidArgument("unsupported layout"); } } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); @@ -1487,7 +1487,7 @@ c10::intrusive_ptr ProcessGroupGloo::allreduce_coalesced( invalidArgument("unsupported layout"); } } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1646,7 +1646,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce( opts.reduceOp, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1838,7 +1838,7 @@ c10::intrusive_ptr ProcessGroupGloo::allgather( work = c10::make_intrusive( std::move(context), outputs, inputs, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -1972,7 +1972,7 @@ c10::intrusive_ptr ProcessGroupGloo::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in Gloo process group"); } @@ -2166,7 +2166,7 @@ c10::intrusive_ptr ProcessGroupGloo::gather( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -2349,7 +2349,7 @@ c10::intrusive_ptr ProcessGroupGloo::scatter( work = c10::make_intrusive( std::move(context), outputs, inputs, opts.rootRank, tag); } else { - throw std::runtime_error("Invalid backend"); + TORCH_CHECK(false, "Invalid backend"); } enqueue(work); return work; @@ -2359,7 +2359,7 @@ c10::intrusive_ptr ProcessGroupGloo::reduce_scatter( std::vector& outputs, std::vector>& inputs, const ReduceScatterOptions& opts) { - throw std::runtime_error("ProcessGroupGloo does not support reduce_scatter"); + TORCH_CHECK(false, "ProcessGroupGloo does not support reduce_scatter"); } namespace { @@ -2531,14 +2531,14 @@ c10::intrusive_ptr ProcessGroupGloo::alltoall_base( at::Tensor& checkSingleTensor(std::vector& tensors) { if (tensors.size() != 1) { - throw std::runtime_error("ProcessGroupGloo::send takes a single tensor"); + TORCH_CHECK(false, "ProcessGroupGloo::send takes a single tensor"); } auto& tensor = tensors[0]; if (!tensor.is_contiguous()) { - throw std::runtime_error("input tensor has to be contiguous"); + TORCH_CHECK(false, "input tensor has to be contiguous"); } if (tensor.is_sparse()) { - throw std::runtime_error("input tensor has to be dense"); + TORCH_CHECK(false, "input tensor has to be dense"); } return tensor; } diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp index 0c471216dffa7..aa6d81bbe4a13 100644 --- a/torch/lib/c10d/ProcessGroupMPI.cpp +++ b/torch/lib/c10d/ProcessGroupMPI.cpp @@ -21,7 +21,7 @@ namespace c10d { std::string err = "MPI error in: " + std::string(__FILE__) + ":" + \ std::to_string(__LINE__) + \ ", with error code: " + std::to_string(mpiStatus); \ - throw std::runtime_error(err); \ + TORCH_CHECK(false, err); \ } \ } while (0) @@ -63,13 +63,13 @@ bool cudaAwareMpiCheck() { // Checking the input tensor's validity void checkSingleTensorHelper(const at::Tensor& tensor) { if (!tensor.is_contiguous()) { - throw std::runtime_error("input tensor has to be contiguous"); + TORCH_CHECK(false, "input tensor has to be contiguous"); } if (tensor.is_sparse()) { - throw std::runtime_error("input tensor has to be dense"); + TORCH_CHECK(false, "input tensor has to be dense"); } if (tensor.is_cuda() && !cudaAwareMpiCheck()) { - throw std::runtime_error( + TORCH_CHECK(false, "CUDA tensor detected and the MPI used doesn't " "have CUDA-aware MPI support"); } @@ -77,7 +77,7 @@ void checkSingleTensorHelper(const at::Tensor& tensor) { void checkSingleTensor(const std::vector& tensors) { if (tensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "MPI process group does not support multi-GPU collectives"); } checkSingleTensorHelper(tensors[0]); @@ -89,7 +89,7 @@ void checkSameSizeAndType( for (const auto& tensor : tensors) { if ((tensor.numel() != t_in.numel()) || (tensor.scalar_type() != t_in.scalar_type())) { - throw std::runtime_error("Tensors are not equal in size or data type"); + TORCH_CHECK(false, "Tensors are not equal in size or data type"); } checkSingleTensorHelper(tensor); } @@ -158,7 +158,7 @@ bool ProcessGroupMPI::AsyncWork::isCompleted() { bool ProcessGroupMPI::AsyncWork::isSuccess() const { if (request_ != MPI_REQUEST_NULL) { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid call to AsyncWork::isSuccess before work has completed"); } @@ -232,14 +232,14 @@ void ProcessGroupMPI::initMPIOnce() { MPI_CHECK(MPI_Init_thread( nullptr, nullptr, MPI_THREAD_SERIALIZED, &mpiThreadSupport_)); if (mpiThreadSupport_ < MPI_THREAD_SERIALIZED) { - throw std::runtime_error( + TORCH_CHECK(false, "Used MPI implementation doesn't have the " "minimum level of threading support: " "MPI_THREAD_SERIALIZED. This is required by " "c10d package"); } if (std::atexit(ProcessGroupMPI::mpiExit)) { - throw std::runtime_error("Fail to register the MPI exit handler"); + TORCH_CHECK(false, "Fail to register the MPI exit handler"); } }); } @@ -285,7 +285,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( MPI_CHECK(MPI_Comm_size(groupComm, &size)); if (rank < 0 || size < 0) { - throw std::runtime_error("Failed to get the world_size / rank"); + TORCH_CHECK(false, "Failed to get the world_size / rank"); } } } @@ -303,7 +303,7 @@ c10::intrusive_ptr ProcessGroupMPI::createProcessGroupMPI( ProcessGroupMPI::ProcessGroupMPI(int rank, int size, MPI_Comm pgComm) : ProcessGroup(rank, size), stop_(false), pgComm_(pgComm) { if (pgComm_ == MPI_COMM_NULL) { - throw std::runtime_error("pgComm_ must not be MPI_COMM_NULL"); + TORCH_CHECK(false, "pgComm_ must not be MPI_COMM_NULL"); } // Start the worker thread accepting MPI calls @@ -427,7 +427,7 @@ c10::intrusive_ptr ProcessGroupMPI::allreduce( c10::intrusive_ptr ProcessGroupMPI::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - throw std::runtime_error( + TORCH_CHECK(false, "allreduce_coalesced is currently not supported with MPI"); } @@ -467,12 +467,12 @@ c10::intrusive_ptr ProcessGroupMPI::allgather( const AllgatherOptions& opts) { checkSingleTensor(inputTensors); if (outputTensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "MPI process group only supports a single " "tensor op"); } if (static_cast(size_) != outputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "All gather: number of output tensors should equal " "to the world size"); } @@ -512,7 +512,7 @@ c10::intrusive_ptr ProcessGroupMPI::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupMPI does not support allgather_coalesced"); } @@ -524,16 +524,16 @@ c10::intrusive_ptr ProcessGroupMPI::gather( if (rank_ != opts.rootRank) { if (outputTensors.size() > 0) { - throw std::runtime_error( + TORCH_CHECK(false, "Gather: number of output tensors should be 0 " "for non-root"); } } else { if (outputTensors.size() != 1) { - throw std::runtime_error("Gather: multi-GPU collective is not supported"); + TORCH_CHECK(false, "Gather: multi-GPU collective is not supported"); } if (static_cast(size_) != outputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Gather: number of output tensors should equal " "to the world size"); } @@ -598,17 +598,17 @@ c10::intrusive_ptr ProcessGroupMPI::scatter( if (rank_ != opts.rootRank) { if (inputTensors.size() > 0) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: number of input tensors should be 0 " "for non-root"); } } else { if (inputTensors.size() != 1) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: multi-GPU collective is not supported"); } if (static_cast(size_) != inputTensors[0].size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Scatter: number of input tensors should equal " "to the world size"); } @@ -670,7 +670,7 @@ c10::intrusive_ptr ProcessGroupMPI::reduce_scatter( std::vector& outputTensors, std::vector>& inputTensors, const ReduceScatterOptions& opts) { - throw std::runtime_error("ProcessGroupMPI does not support reduce_scatter"); + TORCH_CHECK(false, "ProcessGroupMPI does not support reduce_scatter"); } c10::intrusive_ptr ProcessGroupMPI::alltoall_base( @@ -917,7 +917,7 @@ c10::intrusive_ptr ProcessGroupMPI::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in MPI process group"); } diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp index 3f62cab44602b..f538e2f4ea560 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.cpp +++ b/torch/lib/c10d/ProcessGroupNCCL.cpp @@ -90,16 +90,16 @@ ncclRedOp_t getNcclReduceOp(const ReduceOp reduceOp, at::Tensor& input) { } catch (const std::out_of_range& e) { switch (reduceOp) { case ReduceOp::BAND: - throw std::runtime_error("Cannot use ReduceOp.BAND with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BAND with NCCL"); break; case ReduceOp::BOR: - throw std::runtime_error("Cannot use ReduceOp.BOR with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BOR with NCCL"); break; case ReduceOp::BXOR: - throw std::runtime_error("Cannot use ReduceOp.BXOR with NCCL"); + TORCH_CHECK(false, "Cannot use ReduceOp.BXOR with NCCL"); break; default: - throw std::runtime_error("Unhandled ReduceOp"); + TORCH_CHECK(false, "Unhandled ReduceOp"); break; } } @@ -396,7 +396,7 @@ void ProcessGroupNCCL::WorkNCCL::synchronizeInternal( " ran for ", timeElapsed.count(), " milliseconds before timing out."); - throw std::runtime_error(exceptionMsg); + TORCH_CHECK(false, exceptionMsg); } // Check for errors and throw appropriate exception. checkAndThrowException(); @@ -819,7 +819,7 @@ std::vector>& ProcessGroupNCCL::getNCCLComm( bool isSendRecvSelf) { // Sanity check if (devicesKey.empty()) { - throw std::runtime_error( + TORCH_CHECK(false, "Not able to create/get the NCCL Communicator since " "the GPU devices are not known"); } @@ -945,10 +945,10 @@ namespace { // Check validity of tensor void check_gpu_single_tensor(const at::Tensor& tensor) { if (!tensor.is_cuda() || tensor.is_sparse()) { - throw std::runtime_error("Tensors must be CUDA and dense"); + TORCH_CHECK(false, "Tensors must be CUDA and dense"); } if (!tensor.is_contiguous()) { - throw std::runtime_error("Tensors must be contiguous"); + TORCH_CHECK(false, "Tensors must be contiguous"); } } @@ -956,10 +956,10 @@ void check_gpu_single_tensor(const at::Tensor& tensor) { // across distinct GPUs. void check_gpu_tensors(const std::vector& tensors) { if (tensors.size() == 0) { - throw std::runtime_error("Tensor list must be nonempty"); + TORCH_CHECK(false, "Tensor list must be nonempty"); } if (tensors.size() > static_cast(at::cuda::getNumGPUs())) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list mustn't be larger than the number of available GPUs"); } @@ -971,23 +971,23 @@ void check_gpu_tensors(const std::vector& tensors) { for (const auto& t : tensors) { if (!t.is_cuda() || t.is_sparse()) { - throw std::runtime_error("Tensors must be CUDA and dense"); + TORCH_CHECK(false, "Tensors must be CUDA and dense"); } if (t.scalar_type() != first.scalar_type()) { - throw std::runtime_error("Tensors must have identical type"); + TORCH_CHECK(false, "Tensors must have identical type"); } if (t.sizes() != first.sizes()) { - throw std::runtime_error("Tensors must have identical size"); + TORCH_CHECK(false, "Tensors must have identical size"); } if (t.strides() != first.strides()) { - throw std::runtime_error("Tensors must have identical strides"); + TORCH_CHECK(false, "Tensors must have identical strides"); } if (!t.is_non_overlapping_and_dense()) { - throw std::runtime_error("Tensors must be non-overlapping and dense"); + TORCH_CHECK(false, "Tensors must be non-overlapping and dense"); } const auto inserted = usedDevices.insert(t.get_device()).second; if (!inserted) { - throw std::runtime_error("Tensors must be on distinct GPU devices"); + TORCH_CHECK(false, "Tensors must be on distinct GPU devices"); } } } @@ -999,7 +999,7 @@ std::vector flatten_for_scatter_gather( std::vector& other, size_t world_size) { if (tensor_lists.size() != other.size()) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list operands to scatter/gather must have the same length"); } const auto num_devices = tensor_lists.size(); @@ -1009,7 +1009,7 @@ std::vector flatten_for_scatter_gather( for (auto i = size_t{}; i < num_devices; ++i) { if (tensor_lists[i].size() != world_size * num_devices) { - throw std::runtime_error( + TORCH_CHECK(false, "Tensor list input to scatter/gather must match number of collective" " participants"); } @@ -1017,14 +1017,14 @@ std::vector flatten_for_scatter_gather( // Only check device match for the first tensor in the list; the call to // newLikeFlat() below will check the rest. if (tensor_lists[i].front().get_device() != other[i].get_device()) { - throw std::runtime_error( + TORCH_CHECK(false, "Corresponding input/output tensors to scatter/gather must all reside" " on the same device"); } for (const auto& t : tensor_lists[i]) { if (t.numel() != other[i].numel()) { - throw std::runtime_error( + TORCH_CHECK(false, "All tensor operands to scatter/gather must have the same number of elements"); } } @@ -1343,7 +1343,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allreduce( c10::intrusive_ptr ProcessGroupNCCL::allreduce_coalesced( std::vector& tensors, const AllreduceCoalescedOptions& opts) { - throw std::runtime_error( + TORCH_CHECK(false, "allreduce_coalesced is currently not supported with NCCL"); } @@ -1481,7 +1481,7 @@ c10::intrusive_ptr ProcessGroupNCCL::allgather_coalesced( std::vector>& /* unused */, std::vector& /* unused */, const AllgatherOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL does not support allgather_coalesced"); } @@ -1549,11 +1549,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_reduce_scatter_base( const ReduceScatterOptions& opts) { if (inputTensor.dtype() != outputTensor.dtype()) { - throw std::runtime_error("input tensor must be the same type as the outut tensor."); + TORCH_CHECK(false, "input tensor must be the same type as the outut tensor."); } if (inputTensor.numel() != outputTensor.numel() * size_) { - throw std::runtime_error("input tensor must be the same size as output size times world size"); + TORCH_CHECK(false, "input tensor must be the same size as output size times world size"); } // @lint-ignore CLANGTIDY @@ -1821,7 +1821,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall_base( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1829,7 +1829,7 @@ c10::intrusive_ptr ProcessGroupNCCL::alltoall( std::vector& /* unused */, std::vector& /* unused */, const AllToAllOptions& /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports alltoall* for NCCL lib version >= 2.7.0"); } @@ -1837,7 +1837,7 @@ c10::intrusive_ptr ProcessGroupNCCL::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports send for NCCL lib version >= 2.7.0"); } @@ -1845,7 +1845,7 @@ c10::intrusive_ptr ProcessGroupNCCL::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "ProcessGroupNCCL only supports recv for NCCL lib version >= 2.7.0"); } #endif @@ -1868,20 +1868,20 @@ c10::intrusive_ptr ProcessGroupNCCL::gather( std::vector>& /* unused */, std::vector& /* unused */, const GatherOptions& /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support gather"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support gather"); } c10::intrusive_ptr ProcessGroupNCCL::scatter( std::vector& /* unused */, std::vector>& /* unused */, const ScatterOptions& /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support scatter"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support scatter"); } c10::intrusive_ptr ProcessGroupNCCL::recvAnysource( std::vector& /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupNCCL does not support recvAnysource"); + TORCH_CHECK(false, "ProcessGroupNCCL does not support recvAnysource"); } c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( @@ -1892,11 +1892,11 @@ c10::intrusive_ptr ProcessGroupNCCL::_allgather_base( check_gpu_single_tensor(output_tensor); if (input_tensor.dtype() != output_tensor.dtype()) { - throw std::runtime_error("output tensor must have the same type as input tensor"); + TORCH_CHECK(false, "output tensor must have the same type as input tensor"); } if (input_tensor.numel() * size_ != output_tensor.numel()) { - throw std::runtime_error("output tensor size must be equal to world_size times input tensor size"); + TORCH_CHECK(false, "output tensor size must be equal to world_size times input tensor size"); } // just a wrapper to fit the collective interface diff --git a/torch/lib/c10d/ProcessGroupRoundRobin.cpp b/torch/lib/c10d/ProcessGroupRoundRobin.cpp index a55eea968b1e1..c439cf771a147 100644 --- a/torch/lib/c10d/ProcessGroupRoundRobin.cpp +++ b/torch/lib/c10d/ProcessGroupRoundRobin.cpp @@ -90,25 +90,25 @@ c10::intrusive_ptr ProcessGroupRoundRobin::send( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support send"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support send"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recv( std::vector& /* unused */, int /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::recvAnysource( std::vector& /* unused */, int /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support recv"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support recv"); }; c10::intrusive_ptr ProcessGroupRoundRobin::barrier( const BarrierOptions& /* unused */) { - throw std::runtime_error("ProcessGroupRoundRobin does not support barrier"); + TORCH_CHECK(false, "ProcessGroupRoundRobin does not support barrier"); }; const c10::intrusive_ptr& ProcessGroupRoundRobin::next() { @@ -124,7 +124,7 @@ c10::intrusive_ptr ProcessGroupRoundRobin::_allgather_base( at::Tensor& /*unused */, at::Tensor& /*unused */, const AllgatherOptions& /*unused */) { - throw std::runtime_error( + TORCH_CHECK(false, "no support for _allgather_base in RoundRobin process group"); } diff --git a/torch/lib/c10d/TCPStore.cpp b/torch/lib/c10d/TCPStore.cpp index 4958c47b79a71..6498f8bcbe633 100644 --- a/torch/lib/c10d/TCPStore.cpp +++ b/torch/lib/c10d/TCPStore.cpp @@ -133,7 +133,7 @@ void BackgroundThread::join() { void BackgroundThread::initStopSignal() { ghStopEvent_ = CreateEvent(NULL, TRUE, FALSE, NULL); if (ghStopEvent_ == NULL) { - throw std::runtime_error( + TORCH_CHECK(false, "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -149,7 +149,7 @@ void BackgroundThread::stop() { #else void BackgroundThread::initStopSignal() { if (pipe(controlPipeFd_.data()) == -1) { - throw std::runtime_error( + TORCH_CHECK(false, "Failed to create the control pipe to start the " "BackgroundThread run"); } @@ -336,7 +336,7 @@ void TCPStoreMasterDaemon::query(int socket) { watchHandler(socket); } else { - throw std::runtime_error("Unexpected query type"); + TORCH_CHECK(false, "Unexpected query type"); } } @@ -1126,7 +1126,7 @@ bool TCPStore::check(const std::vector& keys) { if (response == detail::CheckResponseType::NOT_READY) { return false; } - throw std::runtime_error("ready or not_ready response expected"); + TORCH_CHECK(false, "ready or not_ready response expected"); } void TCPStore::wait(const std::vector& keys) { @@ -1156,7 +1156,7 @@ void TCPStore::doWait( auto response = client_->receiveValue(); if (response != detail::WaitResponseType::STOP_WAITING) { - throw std::runtime_error("Stop_waiting response is expected"); + TORCH_CHECK(false, "Stop_waiting response is expected"); } } diff --git a/torch/lib/c10d/UnixSockUtils.hpp b/torch/lib/c10d/UnixSockUtils.hpp index fa74be27f889e..b75bddb763787 100644 --- a/torch/lib/c10d/UnixSockUtils.hpp +++ b/torch/lib/c10d/UnixSockUtils.hpp @@ -56,7 +56,7 @@ inline void waitSocketConnected( throw std::system_error(errno, std::system_category()); } else if (numReady == 0) { errno = 0; - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } socklen_t errLen = sizeof(errno); diff --git a/torch/lib/c10d/Utils.cpp b/torch/lib/c10d/Utils.cpp index f8adc58746c66..5d9aa744dbacd 100644 --- a/torch/lib/c10d/Utils.cpp +++ b/torch/lib/c10d/Utils.cpp @@ -118,7 +118,7 @@ PortType getSocketPort(int fd) { listenPort = ntohs(addr->sin6_port); } else { - throw std::runtime_error("unsupported protocol"); + TORCH_CHECK(false, "unsupported protocol"); } return listenPort; } @@ -140,7 +140,7 @@ std::string sockaddrToString(struct ::sockaddr* addr) { __output != nullptr) address[INET6_ADDRSTRLEN] = '\0'; } else { - throw std::runtime_error("unsupported protocol"); + TORCH_CHECK(false, "unsupported protocol"); } return address; } @@ -229,7 +229,7 @@ void handleConnectException( if (timeout != kNoTimeout) { const auto elapsed = std::chrono::high_resolution_clock::now() - start; if (elapsed > timeout) { - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::seconds(1)); @@ -346,7 +346,7 @@ std::tuple accept( while (true) { int res = tcputil::poll(events.get(), 1, timeout.count()); if (res == 0) { - throw std::runtime_error( + TORCH_CHECK(false, "waiting for processes to " "connect has timed out"); } else if (res == -1) { diff --git a/torch/lib/c10d/Utils.hpp b/torch/lib/c10d/Utils.hpp index 55edff85606cf..5beb5f1c6708b 100644 --- a/torch/lib/c10d/Utils.hpp +++ b/torch/lib/c10d/Utils.hpp @@ -92,7 +92,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { try { val = std::stoi(stringValue); } catch (std::exception& e) { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid value for environment variable: " + std::string(envVarName)); } if (val == 1) { @@ -100,7 +100,7 @@ inline bool parseEnvVarFlag(const char* envVarName) { } else if (val == 0) { return false; } else { - throw std::runtime_error( + TORCH_CHECK(false, "Invalid value for environment variable: " + std::string(envVarName)); } } @@ -340,16 +340,16 @@ inline at::Tensor newLikeFlat( std::vector>& tensors, size_t deviceIdx) { if (tensors.size() == 0 || tensors[0].size() == 0) { - throw std::runtime_error("Received an empty list"); + TORCH_CHECK(false, "Received an empty list"); } if (deviceIdx >= tensors.size()) { - throw std::runtime_error("Invalid device index"); + TORCH_CHECK(false, "Invalid device index"); } auto& t = tensors[deviceIdx][0]; auto device = t.device(); for (size_t i = 1; i < tensors[deviceIdx].size(); ++i) { if (tensors[deviceIdx][i].device() != device) { - throw std::runtime_error("Expecting all tensors on the same device"); + TORCH_CHECK(false, "Expecting all tensors on the same device"); } } at::DeviceGuard gpuGuard(device); @@ -363,7 +363,7 @@ inline at::Tensor newLikeFlat( inline at::Tensor newLikeFlat(std::vector& tensors) { if (tensors.size() == 0) { - throw std::runtime_error("Received an empty list"); + TORCH_CHECK(false, "Received an empty list"); } auto& t = tensors[0]; at::DeviceGuard gpuGuard(t.device()); @@ -504,7 +504,7 @@ using SizeType = uint64_t; continue; \ } else if ( \ errno_local == WSAETIMEDOUT || errno_local == WSAEWOULDBLOCK) { \ - throw std::runtime_error("Socket Timeout"); \ + TORCH_CHECK(false, "Socket Timeout"); \ } else { \ throw std::system_error(errno_local, std::system_category()); \ } \ @@ -521,7 +521,7 @@ using SizeType = uint64_t; if (errno == EINTR) { \ continue; \ } else if (errno == EAGAIN || errno == EWOULDBLOCK) { \ - throw std::runtime_error("Socket Timeout"); \ + TORCH_CHECK(false, "Socket Timeout"); \ } else { \ throw std::system_error(errno, std::system_category()); \ } \ diff --git a/torch/lib/c10d/WinSockUtils.hpp b/torch/lib/c10d/WinSockUtils.hpp index cd37695845ab1..793a0dc7640f2 100644 --- a/torch/lib/c10d/WinSockUtils.hpp +++ b/torch/lib/c10d/WinSockUtils.hpp @@ -46,7 +46,7 @@ inline void waitSocketConnected( std::chrono::high_resolution_clock::now() - startTime; if (elapsed > timeout) { errno = 0; - throw std::runtime_error(kConnectTimeoutMsg); + TORCH_CHECK(false, kConnectTimeoutMsg); } } std::this_thread::sleep_for(std::chrono::milliseconds(10)); diff --git a/torch/lib/c10d/frontend.cpp b/torch/lib/c10d/frontend.cpp index 86a78b6fcebb5..b65cba79884af 100644 --- a/torch/lib/c10d/frontend.cpp +++ b/torch/lib/c10d/frontend.cpp @@ -146,7 +146,7 @@ c10::intrusive_ptr DistributedC10d::newProcessGroupHelper( pg_name) { return pg_name.second == *group_name; }); if (it != pg_names_.end()) { - throw std::runtime_error( + TORCH_CHECK(false, "The specified group name has already been " "created, please use a different group name"); } diff --git a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp index f3a44cbcad4ae..a158d2c9685df 100644 --- a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp +++ b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp @@ -221,7 +221,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - throw std::runtime_error("future result should be tensor list or none"); + TORCH_CHECK(false, "future result should be tensor list or none"); } } return copyTensors(outputTensors); diff --git a/torch/lib/c10d/test/ProcessGroupMPITest.cpp b/torch/lib/c10d/test/ProcessGroupMPITest.cpp index bfefbbba2945e..b8538a016d5b7 100644 --- a/torch/lib/c10d/test/ProcessGroupMPITest.cpp +++ b/torch/lib/c10d/test/ProcessGroupMPITest.cpp @@ -48,7 +48,7 @@ std::vector> waitFuture( } else if (result.isTensorList()) { outputTensors.emplace_back(result.toTensorVector()); } else { - throw std::runtime_error("future result should be tensor list or none"); + TORCH_CHECK(false, "future result should be tensor list or none"); } } return outputTensors; @@ -80,7 +80,7 @@ void testAllreduce(int iter = 1000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -113,7 +113,7 @@ void testBroadcast(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -143,7 +143,7 @@ void testReduce(int iter = 10000) { auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -183,7 +183,7 @@ void testAllgather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -227,7 +227,7 @@ void testGather(int iter = 10000) { auto data = outputTensors[i][j].data_ptr(); for (auto k = 0; k < outputTensors[i][j].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -235,7 +235,7 @@ void testGather(int iter = 10000) { } else { for (const auto i : c10::irange(iter)) { if (outputTensors[i].size() != 0) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -277,7 +277,7 @@ void testScatter(int iter = 1) { auto data = outputTensors[i][0].data_ptr(); for (auto k = 0; k < outputTensors[i][0].numel(); ++k) { if (data[k] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -333,13 +333,13 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { // Verify outputs for (const auto i : c10::irange(iter)) { if (recvAnysource && srcRanks[i] != 0) { - throw std::runtime_error("src rank is wrong for recvAnysource"); + TORCH_CHECK(false, "src rank is wrong for recvAnysource"); } const auto expected = i; auto data = outputTensors[i][0].data_ptr(); for (auto j = 0; j < outputTensors[i][0].numel(); ++j) { if (data[j] != expected) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } } @@ -348,7 +348,7 @@ void testSendRecv(bool recvAnysource, int iter = 10000) { void testBackendName() { auto pg = c10d::ProcessGroupMPI::createProcessGroupMPI(); if (pg->getBackendName() != std::string(c10d::MPI_BACKEND_NAME)) { - throw std::runtime_error("BOOM!"); + TORCH_CHECK(false, "BOOM!"); } } diff --git a/torch/lib/c10d/test/ProcessGroupNCCLErrorsTest.cpp b/torch/lib/c10d/test/ProcessGroupNCCLErrorsTest.cpp index b93f45dadcdd9..fb7baa41c5b87 100644 --- a/torch/lib/c10d/test/ProcessGroupNCCLErrorsTest.cpp +++ b/torch/lib/c10d/test/ProcessGroupNCCLErrorsTest.cpp @@ -220,7 +220,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLTimedoutErrorsBlocking) { // Now run all reduce with errors. pg.set_timedout_error(); work = pg.allreduce(tensors_); - EXPECT_THROW(work->wait(), std::runtime_error); + EXPECT_THROW(work->wait(), c10::Error); // Communicators might be aborted here, further operations would fail. } diff --git a/torch/lib/c10d/test/TCPStoreTest.cpp b/torch/lib/c10d/test/TCPStoreTest.cpp index e5b7eaf35cc5b..65fb425022b24 100644 --- a/torch/lib/c10d/test/TCPStoreTest.cpp +++ b/torch/lib/c10d/test/TCPStoreTest.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -73,7 +74,7 @@ void testHelper(const std::string& prefix = "") { EXPECT_EQ(numKeys, 4); auto timeout = std::chrono::milliseconds(kShortStoreTimeoutMillis); serverStore->setTimeout(timeout); - EXPECT_THROW(serverStore->get("key0"), std::runtime_error); + EXPECT_THROW(serverStore->get("key0"), c10::Error); }); // Hammer on TCPStore @@ -238,7 +239,7 @@ void testWatchKeyCallback(const std::string& prefix = "") { numCallbacksExecutedPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (numCallbacksExecutedFuture.wait_for(span) == std::future_status::timeout) - throw std::runtime_error("Callback execution timed out."); + TORCH_CHECK(false, "Callback execution timed out."); // Check number of callbacks executed equal to number of key change operations // Wait for all callbacks to be triggered @@ -302,7 +303,7 @@ void testKeyChangeHelper( std::future callbackFuture = callbackPromise.get_future(); std::chrono::milliseconds span(kStoreCallbackTimeoutMillis); if (callbackFuture.wait_for(span) == std::future_status::timeout) - throw std::runtime_error("Callback execution timed out."); + TORCH_CHECK(false, "Callback execution timed out."); // Any exceptions raised from asserts should be rethrown if (eptr) @@ -373,7 +374,7 @@ TEST(TCPStoreTest, testCleanShutdown) { clientTCPStore->get("key"); auto clientThread = std::thread([&clientTCPStore] { - EXPECT_THROW(clientTCPStore->get("invalid_key"), std::runtime_error); + EXPECT_THROW(clientTCPStore->get("invalid_key"), std::system_error); }); // start server shutdown during a client request From c50c77b444812d6f28bb7a9071ca0eb643ab6f78 Mon Sep 17 00:00:00 2001 From: cyy Date: Mon, 14 Jun 2021 10:32:41 -0700 Subject: [PATCH 082/305] remove unused variables (#59912) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59912 Reviewed By: soulitzer Differential Revision: D29100518 Pulled By: albanD fbshipit-source-id: b86a4aa9050e4fa70a0872c1d8799e5953cd2bc8 --- aten/src/ATen/native/Convolution.cpp | 3 --- aten/src/ATen/native/TensorShape.cpp | 1 - torch/csrc/autograd/python_function.cpp | 1 - 3 files changed, 5 deletions(-) diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp index 7b87405ba8ee8..e84258b3a561d 100644 --- a/aten/src/ATen/native/Convolution.cpp +++ b/aten/src/ATen/native/Convolution.cpp @@ -476,8 +476,6 @@ static void check_shape_forward(const at::Tensor& input, int64_t weight_dim = weight_sizes.size(); int64_t groups = params.groups; auto padding = params.padding; - auto output_padding = params.output_padding; - auto stride = params.stride; auto dilation = params.dilation; bool transposed = params.transposed; @@ -527,7 +525,6 @@ static void check_shape_forward(const at::Tensor& input, // If kernel size is incorrect std::ostringstream input_ss; std::ostringstream kernel_ss; - std::ostringstream output_ss; std::string separator = ""; for (int i = 0, len = input_shape.size(); i < len; ++i) { diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp index a85acce882cd7..d9bef09a4d07a 100644 --- a/aten/src/ATen/native/TensorShape.cpp +++ b/aten/src/ATen/native/TensorShape.cpp @@ -623,7 +623,6 @@ std::vector unsafe_chunk(const Tensor& self, int64_t chunks, int64_t dim TORCH_CHECK(chunks > 0, "chunk expects `chunks` to be greater than 0, got: ", chunks); - std::vector result; const auto dim_size = self.size(dim); int64_t split_size = (dim_size + chunks - 1) / chunks; diff --git a/torch/csrc/autograd/python_function.cpp b/torch/csrc/autograd/python_function.cpp index 04fa9787ccfab..ac379b5a9dc9d 100644 --- a/torch/csrc/autograd/python_function.cpp +++ b/torch/csrc/autograd/python_function.cpp @@ -497,7 +497,6 @@ static void _trace_post_record( // Isolate C variable ptrs in a vector int num_outputs = PyTuple_GET_SIZE(output_objects); - variable_list output_vars(num_outputs); auto graph = node->owningGraph(); node->addOutput(); if (!unpack_output) { From 1f7251df90afa7a32d37fdcef1aa8021f8a6e17d Mon Sep 17 00:00:00 2001 From: jiej Date: Mon, 14 Jun 2021 10:34:27 -0700 Subject: [PATCH 083/305] fixing DifferentiableGraphOp updating requires_grad on input tensor list; python test added to verify the test (#57574) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/57574 Test Plan: Imported from OSS Reviewed By: bdhirsh Differential Revision: D29038774 Pulled By: Krovatkin fbshipit-source-id: cb342c1b04fa3713a8166b39213437bc9f2d8606 --- test/jit/test_autodiff_subgraph_slicing.py | 22 +++++++++++++++++++++- torch/csrc/jit/runtime/graph_executor.cpp | 6 +++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/test/jit/test_autodiff_subgraph_slicing.py b/test/jit/test_autodiff_subgraph_slicing.py index 71d07ae4154cc..af9b351306998 100644 --- a/test/jit/test_autodiff_subgraph_slicing.py +++ b/test/jit/test_autodiff_subgraph_slicing.py @@ -12,7 +12,7 @@ from torch.testing._internal.jit_utils import JitTestCase, disable_autodiff_subgraph_inlining from torch.testing import FileCheck -from typing import Optional +from typing import List, Tuple, Optional if __name__ == '__main__': raise RuntimeError("This test file is not meant to be run directly, use:\n\n" @@ -138,6 +138,26 @@ def method1(x, weight, bias: Optional[torch.Tensor]): # check_types requires last_graph on scripted to be set, so we just skip it check_against_reference(self, scripted, method1, lambda x: x, (x, weight, bias), check_types=False) + def test_requires_grad_for_tensor_list(self): + + with enable_profiling_mode_for_profiling_tests(): + + # output & var_list[0] should have requires_grad set to True + def func(input0: torch.Tensor, input1: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]: + var_list = [input0, input1] + var = torch.cat(var_list) + output = var + 1.0 + return output, var_list + jit_f = torch.jit.script(func) + input0 = torch.randn((2,), requires_grad=True) + input1 = torch.randn((2,)) + output_ref = func(input0, input1) + for i in range(2): + output = jit_f(input0, input1) + assert(output_ref[0].requires_grad == output[0].requires_grad) + assert(output_ref[1][0].requires_grad == output[1][0].requires_grad) + assert(output_ref[1][1].requires_grad == output[1][1].requires_grad) + def test_simple_merge(self): # o --> o def fn(x, y, z): diff --git a/torch/csrc/jit/runtime/graph_executor.cpp b/torch/csrc/jit/runtime/graph_executor.cpp index 5c77fd5df9fb2..df6a22d73b9d8 100644 --- a/torch/csrc/jit/runtime/graph_executor.cpp +++ b/torch/csrc/jit/runtime/graph_executor.cpp @@ -440,9 +440,9 @@ struct DifferentiableGraphOp { if (v.isTensor()) { v = IValue(detach(std::move(v).toTensor())); } else if (v.isTensorList()) { - c10::List lst = std::move(v).toTensorList(); - for (size_t i = 0; i < lst.size(); ++i) { - lst.set(i, detach(lst.extract(i))); + std::vector lst = v.toTensorVector(); + for (auto& tensor : lst) { + tensor = detach(tensor); } v = std::move(lst); } From 9ad0de3c6f4459ded04004e78e335f633cf05e92 Mon Sep 17 00:00:00 2001 From: jiej Date: Mon, 14 Jun 2021 10:34:27 -0700 Subject: [PATCH 084/305] Rework requires_grad on DifferentiableGraphOp (#57575) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/57575 This PR does two things: 1. reverts "Manual revert of D27369251 (https://github.com/pytorch/pytorch/commit/f88a3fff65b35cb6d4968fc54a9a0a1314a9a3b7) (#56080)" in commit 92a09fb87a567100122b872613344d3a422abc9f. 2. fixing DifferentiableGraph output with wrong requires_grad flag Fixing requires_grad on outputs from DifferentiableGraph, the proper flag is retrieved from profiling information. We previously only retrieves the profiling information on the first profile node in all its uses. However, in case where control flows are present, we need to iteratively search for profile node with profiling information available, in case the first use is in an inactive code path. e.g. ``` graph(%0 : Tensor, %1 : Bool): ..., %2 : Tensor = prim::DifferentiableGraph_0(%0) %3 : Tensor = prim::If(%1) block0(): %4 : Tensor = prim::DifferentiableGraph_1(%2) -> (%4) block1(): %5 : Tensor = prim::DifferentiableGraph_2(%2) -> (%5) -> (%3) with prim::DifferentiableGraph_0 = graph(%0 : Tensor): ... %out : Tensor = aten::operation(...) ... return (..., %out) with prim::DifferentiableGraph_1 = graph(%0 : Tensor): %temp : Tensor = prim::profile[profiled_type=Tensor](%0) ... with prim::DifferentiableGraph_2 = graph(%0 : Tensor): %temp : Tensor = prim::profile[profiled_type=Float(...)](%0) ... ``` Test Plan: Imported from OSS Reviewed By: bdhirsh Differential Revision: D29038773 Pulled By: Krovatkin fbshipit-source-id: 6c0a851119f6b8f2f1afae5c74532407aae238fe --- test/jit/test_autodiff_subgraph_slicing.py | 66 +++++++++++++++++ .../runtime/profiling_graph_executor_impl.cpp | 74 +++++++++++++++++++ 2 files changed, 140 insertions(+) diff --git a/test/jit/test_autodiff_subgraph_slicing.py b/test/jit/test_autodiff_subgraph_slicing.py index af9b351306998..42640ac2fa5bc 100644 --- a/test/jit/test_autodiff_subgraph_slicing.py +++ b/test/jit/test_autodiff_subgraph_slicing.py @@ -158,6 +158,72 @@ def func(input0: torch.Tensor, input1: torch.Tensor) -> Tuple[torch.Tensor, List assert(output_ref[1][0].requires_grad == output[1][0].requires_grad) assert(output_ref[1][1].requires_grad == output[1][1].requires_grad) + @unittest.skipIf(GRAPH_EXECUTOR != ProfilingMode.PROFILING, + "Requires fusion optimization pass to be effective") + def test_differentiable_graph_ops_requires_grad(self): + x = torch.randn(8, 2, dtype=torch.float).requires_grad_() + y = torch.randn(8, 2, dtype=torch.float) + + def t(x : torch.Tensor, y : torch.Tensor, flag : bool): + o = x + 1.0 + o1 = torch.relu(o) + o = y + 1.5 + o2 = torch.relu(o) + o3 = o1 + o2 + + if flag: + o = o1 + 1.0 + oo1 = torch.relu(o) + o = o2 + 2.5 + oo2 = torch.relu(o) + oo3 = oo1 + oo2 + else: + o = o1 * 1.0 + oo1 = torch.relu(o) + o = o2 * 2.0 + oo2 = torch.relu(o) + oo3 = oo1 + oo2 + + return o1, o2, o3, oo1, oo2, oo3 + + with enable_profiling_mode_for_profiling_tests(): + + t_jit = torch.jit.script(t) + jit_o = t_jit(x, y, False) + jit_o = t_jit(x, y, False) + o = t(x, y, False) + + FileCheck().check("prim::DifferentiableGraph").run(t_jit.graph_for(x, y, False)) + # validate the differentiableGraphOps are marking proper requires_grad + for oo, jit_oo in zip(o, jit_o): + self.assertEqual(oo.requires_grad, jit_oo.requires_grad) + self.assertEqual(oo, jit_oo) + # one more runs to trigger fusion + jit_o = t_jit(x, y, False) + for oo, jit_oo in zip(o, jit_o): + self.assertEqual(oo.dtype, jit_oo.dtype) + self.assertEqual(oo.requires_grad, jit_oo.requires_grad) + self.assertEqual(oo, jit_oo) + + @unittest.skipIf(GRAPH_EXECUTOR == ProfilingMode.PROFILING, "Simple Executor doesn't support gradients") + def test_prune_grad(self): + @torch.jit.script + def t(input, bias): + return torch.nn.functional.relu(input + bias) + input = torch.randn(2, 8, requires_grad=True) + bias = torch.randn(8, requires_grad=False) # bias does NOT require grad + NUM_PROFILED_RUNS = 1 + with num_profiled_runs(NUM_PROFILED_RUNS): + WARMUP = 3 # 2 runs to reach backward + 1 to optimize it + for x in range(WARMUP): + o = t(input, bias) + o.sum().backward() + + fwd_plan = list(t.get_debug_state().execution_plans.values())[0] + bwd_graph = list(fwd_plan.code.grad_executor_states()[0].execution_plans.values())[0].graph + tup = next(bwd_graph.outputs()) + self.assertEqual(len(list(tup.node().inputs())), 1) + def test_simple_merge(self): # o --> o def fn(x, y, z): diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp index 0716b9ff46288..2342dbb56e814 100644 --- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp @@ -130,6 +130,79 @@ static bool needsGradientInProfilingMode(Block* b) { return false; } +// `prim::RequiresGradCheck` guarantees that requires_grad properties +// of input tensors will match the profiled, otherwise a fallback path +// will be triggered. This allow us to prune off gradients in backward +// graph for inputs that don't need gradients. We transfer requires_grad +// properties from inputs to the `prim::DifferentiableGraph` onto inputs to the +// differentiable graph. Autodiff will inspect these properties and prune +// off gradients that aren't required +// `requires_grad` properties from `dnode->outputs()` will also be transferred +static void setRequiresGradOnDiffGraph(Node* dnode) { + auto gi = dnode->g(attr::Subgraph)->inputs(); + for (size_t i = 0; i < dnode->inputs().size(); i++) { + if (auto ty = dnode->input(i)->type()->cast()) { + auto gi_ty = gi[i]->type()->expect(); + gi[i]->setType(gi_ty->withRequiresGrad(ty->requires_grad())); + GRAPH_DEBUG( + "Setting ", + *gi_ty->withRequiresGrad(ty->requires_grad()), + " on ", + gi[i], + " ", + gi[i]->debugName()); + } + } + + // We also need to put requires_grad on outputs within subgraph, so autodiff + // can set df_input_vjps and DifferentiableGraphOp can set `requires_grad=` + // properly + auto go = dnode->g(attr::Subgraph)->outputs(); + auto set_requires_grad = [](const TensorTypePtr& t, Value* val) -> bool { + if (t && t->requiresGrad().has_value()) { + GRAPH_DEBUG("setting type ", *t); + val->setType(t); + return true; + } + return false; + }; + + for (size_t i = 0; i < go.size(); i++) { + auto ty = go[i]->type()->cast(); + if (ty) { + auto n = go[i]->node(); + auto dno = dnode->outputs().at(i); + for (auto dno_use : dno->uses()) { + GRAPH_DEBUG("found user of ", i, " as ", *dno_use.user); + if (n->kind() == prim::profile) { + if (set_requires_grad( + n->ty(attr::profiled_type)->expect(), go[i])) { + break; + } + } else if (dno_use.user->kind() == prim::profile) { + if (set_requires_grad( + dno_use.user->ty(attr::profiled_type)->expect(), + go[i])) { + break; + } + } else if (dno_use.user->kind() == prim::DifferentiableGraph) { + Value* o = + dno_use.user->g(attr::Subgraph)->inputs().at(dno_use.offset); + // Is it safe to not check other uses, because we are inside a + // DifferentiableGraph? + auto nn = o->uses().at(0).user; + if (nn->kind() == prim::profile) { + if (set_requires_grad( + nn->ty(attr::profiled_type)->expect(), go[i])) { + break; + } + } + } + } + } + } +} + bool guardDifferentiableGraph(Node* dnode) { auto gi = dnode->g(attr::Subgraph)->inputs(); bool all_inputs_seen = true; @@ -163,6 +236,7 @@ bool guardDifferentiableGraph(Node* dnode) { } } if (all_inputs_seen) { + setRequiresGradOnDiffGraph(dnode); // we may have seen both true and false for requires_grad. In this case // we guard with true here and the other case is in the fallback. This // will give us trouble when we get "alternating patterns" of gradients From ab70e1e9848cf470fa97ee50a5de5e1f27d2c22d Mon Sep 17 00:00:00 2001 From: Jason Ansel Date: Mon, 14 Jun 2021 10:35:22 -0700 Subject: [PATCH 085/305] [TensorExpr] Add error checking in mem_arena (#59922) Summary: Gives an error message (rather than a segfault) if you forget `KernelScope()`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59922 Reviewed By: bertmaher Differential Revision: D29091303 Pulled By: jansel fbshipit-source-id: a24ee2385cae1f210b0cbc3f8860948fc052b655 --- test/test_tensorexpr_pybind.py | 3 +++ torch/csrc/jit/tensorexpr/mem_arena.cpp | 32 ++++++++++++++++++------- torch/csrc/jit/tensorexpr/mem_arena.h | 8 +++---- 3 files changed, 31 insertions(+), 12 deletions(-) diff --git a/test/test_tensorexpr_pybind.py b/test/test_tensorexpr_pybind.py index fceec6ecd9d15..a3efb8416a37d 100644 --- a/test/test_tensorexpr_pybind.py +++ b/test/test_tensorexpr_pybind.py @@ -352,6 +352,9 @@ def f(a): np.testing.assert_allclose(res1.numpy(), correct.numpy(), atol=2e-3) np.testing.assert_allclose(res2.numpy(), correct.numpy(), atol=2e-3) + def test_forgot_kernel_arena(self): + self.assertRaises(RuntimeError, lambda: torch._C._te.VarHandle("n", torch._C._te.Dtype.Int)) + if __name__ == '__main__': run_tests() diff --git a/torch/csrc/jit/tensorexpr/mem_arena.cpp b/torch/csrc/jit/tensorexpr/mem_arena.cpp index 65ccf1ce483ab..625dd6473734f 100644 --- a/torch/csrc/jit/tensorexpr/mem_arena.cpp +++ b/torch/csrc/jit/tensorexpr/mem_arena.cpp @@ -1,4 +1,6 @@ +#include #include +#include namespace torch { namespace jit { @@ -19,6 +21,10 @@ KernelArena::~KernelArena() { KernelScopedObject::KernelScopedObject() { KernelArena* kernel = KernelArena::GetCurrentKernelArena(); + if (kernel == nullptr) { + throw std::runtime_error( + "KernelScope() must be constructed before calling this"); + } kernel->kernel_objects_.push_back(this); } @@ -30,21 +36,31 @@ KernelArena* KernelArena::GetCurrentKernelArena() { return current_arena; } -KernelScope::KernelScope() : owning_(true) { - old_kernel_arena_ = KernelArena::GetCurrentKernelArena(); - KernelArena::SetCurrentKernelArena(new KernelArena); +KernelScope::KernelScope() + : kernel_arena_(new KernelArena()), + old_kernel_arena_(KernelArena::GetCurrentKernelArena()), + owning_(true) { + KernelArena::SetCurrentKernelArena(kernel_arena_); } -KernelScope::KernelScope(KernelArena* arena_) : owning_(false) { - old_kernel_arena_ = KernelArena::GetCurrentKernelArena(); - KernelArena::SetCurrentKernelArena(arena_); +KernelScope::KernelScope(KernelArena* arena_) + : kernel_arena_(arena_), + old_kernel_arena_(KernelArena::GetCurrentKernelArena()), + owning_(false) { + KernelArena::SetCurrentKernelArena(kernel_arena_); } KernelScope::~KernelScope() { - if (owning_) { - delete KernelArena::GetCurrentKernelArena(); + if (KernelArena::GetCurrentKernelArena() != kernel_arena_) { + // This should be an error, but it gets triggered in + // caffe2/benchmarks/static_runtime:static_runtime_cpptest + TORCH_WARN("KernelScope() destructed out of order, leaking memory"); + return; } KernelArena::SetCurrentKernelArena(old_kernel_arena_); + if (owning_) { + delete kernel_arena_; + } } } // namespace tensorexpr diff --git a/torch/csrc/jit/tensorexpr/mem_arena.h b/torch/csrc/jit/tensorexpr/mem_arena.h index 0b1e96c59662c..3b8c329c3b420 100644 --- a/torch/csrc/jit/tensorexpr/mem_arena.h +++ b/torch/csrc/jit/tensorexpr/mem_arena.h @@ -36,10 +36,10 @@ class KernelScope { KernelScope& operator=(const KernelScope&) = delete; private: - KernelArena* old_kernel_arena_ = - nullptr; // previous arena, will be restored in destructor - bool owning_ = false; // determines whether the arena will be freed along with - // the scope object + KernelArena* kernel_arena_; // maybe owned + KernelArena* old_kernel_arena_; // previous arena, restored in destructor + bool owning_; // determines whether the arena will be freed along with + // the scope object }; // The base object managed by the Kernel. From 061e71b1994bcd8b73971f3c365d952b0bf563a3 Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 14 Jun 2021 11:10:34 -0700 Subject: [PATCH 086/305] Parametrizations depending on several inputs (#58488) Summary: Makes possible that the first register parametrization depends on a number of parameters rather than just one. Examples of these types of parametrizations are `torch.nn.utils.weight_norm` and low rank parametrizations via the multiplication of a `n x k` tensor by a `k x m` tensor with `k <= m, n`. Follows the plan outlined in https://github.com/pytorch/pytorch/pull/33344#issuecomment-768574924. A short summary of the idea is: we call `right_inverse` when registering a parametrization to generate the tensors that we are going to save. If `right_inverse` returns a sequence of tensors, then we save them as `original0`, `original1`... If it returns a `Tensor` or a sequence of length 1, we save it as `original`. We only allow to have many-to-one parametrizations in the first parametrization registered. The next parametrizations would need to be one-to-one. There were a number of choices in the implementation: If the `right_inverse` returns a sequence of parameters, then we unpack it in the forward. This is to allow to write code as: ```python class Sum(nn.Module): def forward(self, X, Y): return X + Y def right_inverse(Z): return Z, torch.zeros_like(Z) ``` rather than having to unpack manually a list or a tuple within the `forward` function. At the moment the errors are a bit all over the place. This is to avoid having to check some properties of `forward` and `right_inverse` when they are registered. I left this like this for now, but I believe it'd be better to call these functions when they are registered to make sure the invariants hold and throw errors as soon as possible. The invariants are the following: 1. The following code should be well-formed ```python X = module.weight Y = param.right_inverse(X) assert isinstance(Y, Tensor) or isinstance(Y, collections.Sequence) Z = param(Y) if isisntance(Y, Tensor) else param(*Y) ``` in other words, if `Y` is a `Sequence` of `Tensor`s (we check also that the elements of the sequence are Tensors), then it is of the same length as the number parameters `param.forward` accepts. 2. Always: `X.dtype == Z.dtype and X.shape == Z.shape`. This is to protect the user from shooting themselves in the foot, as it's too odd for a parametrization to change the metadata of a tensor. 3. If it's one-to-one: `X.dtype == Y.dtype`. This is to be able to do `X.set_(Y)` so that if a user first instantiates the optimiser and then puts the parametrisation, then we reuse `X` and the user does not need to add a new parameter to the optimiser. Alas, this is not possible when the parametrisation is many-to-one. The current implementation of `spectral_norm` and `weight_norm` does not seem to care about this, so this would not be a regression. I left a warning in the documentation though, as this case is a bit tricky. I'm still missing to go over the formatting of the documentation, I'll do that tomorrow. Pull Request resolved: https://github.com/pytorch/pytorch/pull/58488 Reviewed By: soulitzer Differential Revision: D29100708 Pulled By: albanD fbshipit-source-id: b9e91f439cf6b5b54d5fa210ec97c889efb9da38 --- test/test_nn.py | 411 ++++++++++++++++++++++++++++------ torch/nn/utils/parametrize.py | 391 +++++++++++++++++++++++++------- 2 files changed, 654 insertions(+), 148 deletions(-) diff --git a/test/test_nn.py b/test/test_nn.py index 3dde5053ca282..0c6fa695d5f27 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -2155,7 +2155,11 @@ def forward(self, X): # Cayley map # If X is skew-symmetric it returns an orthogonal matrix Id = torch.eye(X.size(0), device=X.device) - return torch.linalg.solve(Id + X, Id - X) + # We call contiguous because solve returns a tensor with strides that are Fortran-contiguous + # and autograd raises a performance warning. + # This happens when we remove the parametrization with leave_parametrized=True, + # which does a set_ with a non-contiguous tensor while the gradient is contiguous + return torch.linalg.solve(Id + X, Id - X).contiguous() # Define a couple vector parametrizations class FirstZero(nn.Module): @@ -2223,10 +2227,16 @@ def forward(self, x): self.assertEqual(model.bias[-1].item(), 0.) self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happpened # Should not throw + + sgd = torch.optim.SGD(model.parameters(), lr=0.01) + + weight_copy = model.weight.clone() + bias_copy = model.bias.clone() + sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - with torch.no_grad(): - for p in model.parameters(): - p.add_(- p.grad, alpha=0.01) + sgd.step() + self.assertNotEqual(model.weight, weight_copy) + self.assertNotEqual(model.bias, bias_copy) # Remove first parametrization. # Check that the model is still parametrized and so is the second parameter @@ -2240,10 +2250,13 @@ def forward(self, x): self.assertEqual(id(model.weight), initial_weight_id) # Keeps the same id self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happened # Should not throw + weight_copy = model.weight.clone() + bias_copy = model.bias.clone() + sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - with torch.no_grad(): - for p in model.parameters(): - p.add_(- p.grad, alpha=0.01) + sgd.step() + self.assertNotEqual(model.weight, weight_copy) + self.assertNotEqual(model.bias, bias_copy) # Remove the second parametrization. # Check that the module is not parametrized @@ -2256,22 +2269,33 @@ def forward(self, x): self.assertFalse(hasattr(model, "parametrizations")) # Not parametrized the module self.assertEqual(model.__class__, nn.Linear) # Resores the previous class self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happeed - # Should not throw + + # Should not throw things are updated + weight_copy = model.weight.clone() + bias_copy = model.bias.clone() + sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - with torch.no_grad(): - for p in model.parameters(): - p.add_(- p.grad, alpha=0.01) + sgd.step() + self.assertNotEqual(model.weight, weight_copy) + self.assertNotEqual(model.bias, bias_copy) # Test leave_parametrized=True for _ in range(2): parametrize.register_parametrization(model, "weight", Skew()) parametrize.register_parametrization(model, "weight", Orthogonal()) parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) - # Should not throw + # We didn't change the dtype nor had multiple inputs, so the id should be the same + self.assertEqual(id(model.weight), initial_weight_id) + self.assertEqual(id(model.bias), initial_bias_id) + + # Should not throw. Things are updated + weight_copy = model.weight.clone() + bias_copy = model.bias.clone() + sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - with torch.no_grad(): - for p in model.parameters(): - p.add_(- p.grad, alpha=0.01) + sgd.step() + self.assertNotEqual(model.weight, weight_copy) + self.assertNotEqual(model.bias, bias_copy) def test_register_and_remove_buffer_parametrization(self): r"""Test that it is possible to add and remove parametrizations on buffers""" @@ -2392,8 +2416,12 @@ def right_inverse(self, X): N = 5 model = nn.Linear(N, N) - # Register the skew-symmetric onstraint. The result is now skew-symmetric - parametrize.register_parametrization(model, "weight", Skew()) + # Register the skew-symmetric constraint. The result is now skew-symmetric + skew = Skew() + # Make the weight skew-symmetric before registering the parametrization + with torch.no_grad(): + model.weight.set_(skew(model.weight)) + parametrize.register_parametrization(model, "weight", skew) X = torch.rand(N, N) # X is not skew-symmetric, so it throws an error with self.assertRaises(ValueError): @@ -2416,46 +2444,321 @@ def right_inverse(self, X): self.assertEqual(model.weight, X) self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X)) - def test_errors_parametrization(self): - # A parametrization shall not change the size of the parameter - class ChangeSize(nn.Module): - def forward(self, x): - return x[:-1] + def test_errors_unparametrized_tensor_parametrization(self): + # Test errors when registering a parametrization on an unparametrized tensor + module = nn.Linear(3, 4) + weight_init = module.weight.clone() - # A simple parametrization that does not implement a right_inverse - class Double(nn.Module): + class Identity(nn.Module): def forward(self, x): - return 2 * x + return x - module = nn.Linear(3, 4) - # This should not throw when registering - parametrize.register_parametrization(module, "weight", ChangeSize()) - # It throws in the forward - with self.assertRaisesRegex(RuntimeError, "may not change the size"): - module(torch.rand(2)) - # Undo - parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) + # Register a parametrization on a non-existing parameter throws + with self.assertRaisesRegex(ValueError, "does not have a parameter"): + parametrize.register_parametrization(module, "foo", Identity()) self.assertFalse(parametrize.is_parametrized(module)) - # Removing a parametrization from an unparametrized tensor throws + # Removing parametrizations from an unparametrized tensor throws with self.assertRaisesRegex(ValueError, "does not have a parametrization"): parametrize.remove_parametrizations(module, "bias") - # Nothing odd happens self.assertFalse(parametrize.is_parametrized(module)) - # Register a parametrization on a non-existing parameter breaks - with self.assertRaisesRegex(ValueError, "does not have a parameter"): - parametrize.register_parametrization(module, "foo", ChangeSize()) + # A correct parametrization with several outputs + class Sum(nn.Module): + def forward(self, x, y): + return x + y + + def right_inverse(self, z): + return z, torch.zeros_like(z) + + parametrize.register_parametrization(module, "weight", Sum()) + # Cannot remove a parametrization with several outputs with `leave_parametrized=False` + with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): + parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) + parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) + + # A parametrization with an incorrect number of outputs + class WrongNumberParams(nn.Module): + def forward(self, x, y, z): + return x + y + z + + def right_inverse(self, w): + return w, torch.zeros_like(w) + + # Makes param(*param.right_inverse(X)) fail + with self.assertRaisesRegex(TypeError, "positional argument"): + parametrize.register_parametrization(module, "weight", WrongNumberParams()) + self.assertFalse(parametrize.is_parametrized(module)) + + # A parametrization with a right_inverse that does not return a Tensor or Sequence[Tensor] + class WrongRightInverse(Identity): + def right_inverse(self, z): + return None + + # right_inverse should return a Tensor or a Sequence[Tensor] + with self.assertRaisesRegex(ValueError, "Tensor or a Sequence of"): + parametrize.register_parametrization(module, "weight", WrongRightInverse()) self.assertFalse(parametrize.is_parametrized(module)) - # Try to assign to a parametrization that does not implement `right_inverse` - parametrize.register_parametrization(module, "weight", Double()) - with self.assertRaisesRegex(RuntimeError, "right_inverse"): - module.weight = torch.rand(4, 3) - # Undo - parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) + # If it's a sequence, it must to be a sequence of tensors + class WrongRightInverseSequence(nn.Module): + def forward(self, x, y): + return x + + def right_inverse(self, z): + return None, z + + with self.assertRaisesRegex(ValueError, "of the sequence with type"): + parametrize.register_parametrization(module, "weight", WrongRightInverseSequence()) self.assertFalse(parametrize.is_parametrized(module)) + # A parametrization from one tensor to one tensor that changes the dtype + class ChangeDtypeInverse(nn.Module): + def forward(self, x): + return x.float() + + def right_inverse(self, w): + return w.bool() + + # For parametrizations that return one tensor, right_inverse may not change the dtype + with self.assertRaisesRegex(ValueError, "outputs one tensor, it may not change the dtype"): + parametrize.register_parametrization(module, "weight", ChangeDtypeInverse()) + self.assertFalse(parametrize.is_parametrized(module)) + + # Doesn't return a tensor + class NotTensor(nn.Module): + def forward(self, x): + return 2 + + # Forward must return a tensor + with self.assertRaisesRegex(ValueError, "must return a tensor"): + parametrize.register_parametrization(module, "weight", NotTensor()) + self.assertFalse(parametrize.is_parametrized(module)) + + # A parametrization from one tensor to one tensor that changes the dtype + class ChangeDtype(nn.Module): + def forward(self, x): + return x.bool() + + # forward should not change the initial dtype + with self.assertRaisesRegex(ValueError, "may not change the dtype"): + parametrize.register_parametrization(module, "weight", ChangeDtype()) + self.assertFalse(parametrize.is_parametrized(module)) + + # Change shape + class ChangeShape(nn.Module): + def forward(self, x): + return x[:-1] + + # forward should not change the original shape + with self.assertRaisesRegex(ValueError, "may not change the shape"): + parametrize.register_parametrization(module, "weight", ChangeShape()) + self.assertFalse(parametrize.is_parametrized(module)) + + # Many to one that changes dtype + class ChangeDtypeMulti(nn.Module): + def forward(self, x, y): + return (x + y).bool() + + def right_inverse(self, w): + return w, w + 1 + + # forward should not change the original shape even for parametrizations with many inputs + with self.assertRaisesRegex(ValueError, "may not change the dtype"): + parametrize.register_parametrization(module, "weight", ChangeDtypeMulti()) + self.assertFalse(parametrize.is_parametrized(module)) + + # Returning a sequence of size one, although weird, it's correct + class SequenceLen1(nn.Module): + def forward(self, x): + return x + + def right_inverse(self, w): + return (w,) + + parametrize.register_parametrization(module, "weight", SequenceLen1()) + self.assertTrue(hasattr(module.parametrizations.weight, "original0")) + self.assertFalse(hasattr(module.parametrizations.weight, "original1")) + _ = module.weight # Does not throw + self.assertTrue(parametrize.is_parametrized(module)) + parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) + + # None of the operations above should have altered the weight + self.assertFalse(parametrize.is_parametrized(module)) + self.assertEqual(module.weight, weight_init) + + def test_errors_parametrized_tensor_parametrization(self): + # Test errors when registering a parametrization on a parametrized tensor + + class Identity(nn.Module): + def forward(self, x): + return x + + module = nn.Linear(3, 4) + parametrize.register_parametrization(module, "weight", Identity()) + + # Has to return a tensor + class WrongReturn(nn.Module): + def forward(self, x): + return x, x + + with self.assertRaisesRegex(ValueError, "must return a tensor"): + parametrize.register_parametrization(module, "weight", WrongReturn()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + # Cannot change dtype + class ChangeDtype(nn.Module): + def forward(self, x): + return x.bool() + + with self.assertRaisesRegex(ValueError, "may not change the dtype"): + parametrize.register_parametrization(module, "weight", ChangeDtype()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + # Cannot change shape + class ChangeShape(nn.Module): + def forward(self, x): + return x[:-1] + + with self.assertRaisesRegex(ValueError, "may not change the shape"): + parametrize.register_parametrization(module, "weight", ChangeShape()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + # The following checks are mostly due to bugs in the code of the parametrization + + # right_inverse has to return a tensor + class WrongReturnInverse(Identity): + def right_inverse(self, x): + return x, x + + with self.assertRaisesRegex(ValueError, "right_inverse must return a tensor"): + parametrize.register_parametrization(module, "weight", WrongReturnInverse()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + # Cannot change dtype + class ChangeDtypeInverse(Identity): + def right_inverse(self, x): + return x.bool() + + with self.assertRaisesRegex(ValueError, "must have the same dtype"): + parametrize.register_parametrization(module, "weight", ChangeDtypeInverse()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + # Cannot change shape + class ChangeShapeInverse(Identity): + def right_inverse(self, x): + return x[:-1] + + with self.assertRaisesRegex(ValueError, "must have the same shape"): + parametrize.register_parametrization(module, "weight", ChangeShapeInverse()) + self.assertTrue(parametrize.is_parametrized(module)) + self.assertEqual(len(module.parametrizations.weight), 1) + self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) + + def test_multiple_inputs_parametrization(self): + # A parametrization with several outputs + class RankOne(nn.Module): + def forward(self, x, y): + # Form a rank-1 matrix from a pair of vectors + return x.unsqueeze(-1) @ y.unsqueeze(-2) + + def right_inverse(self, Y): + # We project the given matrix onto the rank 1 matrices + U, S, Vh = torch.linalg.svd(Y, full_matrices=False) + # S is ordered in a decreasing way. + s0_sqrt = S[0].sqrt().unsqueeze(-1) + return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt + + # Simple parametrisation + class Double(nn.Module): + def forward(self, x): + return 2.0 * x + + def right_inverse(self, w): + return 0.5 * w + + model = nn.Linear(3, 3) + # Test one parametrization + parametrize.register_parametrization(model, "weight", RankOne()) + self.assertTrue(hasattr(model, "parametrizations")) + self.assertTrue(parametrize.is_parametrized(model)) + self.assertTrue(parametrize.is_parametrized(model, "weight")) + self.assertTrue(hasattr(model.parametrizations.weight, "original0")) + self.assertIn("original0", model.parametrizations.weight._parameters) + self.assertTrue(hasattr(model.parametrizations.weight, "original1")) + self.assertIn("original1", model.parametrizations.weight._parameters) + self.assertFalse(parametrize.is_parametrized(model, "bias")) + self.assertNotIn("weight", model._parameters) + # Result should be rank 1 + self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) + + with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): + # Cannot remove a parametrization with multiple inputs and not leave it parametrized + parametrize.remove_parametrizations(model, "weight", leave_parametrized=False) + # Remove parametrization and check consistency + parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) + self.assertFalse(hasattr(model, "parametrizations")) + self.assertEqual(model.__class__, nn.Linear) + self.assertFalse(parametrize.is_parametrized(model)) + self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) + self.assertIn("weight", model._parameters) + + # Registering parametrizations with one input on top of one with multiple inputs should work + init_weight = model.weight.clone() + parametrize.register_parametrization(model, "weight", RankOne()) + # Projecting a rank 1 matrix onto the matrices of rank one does not change the matrix + self.assertTrue(torch.allclose(init_weight, model.weight)) + parametrize.register_parametrization(model, "weight", Double()) + # The matrix now is twice the initial matrix + self.assertTrue(torch.allclose(2.0 * init_weight, model.weight)) + # Multiplying by a scalar does not change the rank + self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) + + # The model has now three parameters + self.assertEqual(len(list(model.parameters())), 3) + + sgd = torch.optim.SGD(model.parameters(), lr=0.1) + + # Test backward. Should not throw + for _ in range(2): + sgd.zero_grad() + loss = (model.weight.T @ model.bias).sum() + loss.backward() + sgd.step() + + # Same drill as before, removing should work as expected + with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): + # Cannot remove a parametrization with multiple inputs and not leave it parametrized + parametrize.remove_parametrizations(model, "weight", leave_parametrized=False) + # Remove parametrization and check consistency + parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) + self.assertFalse(hasattr(model, "parametrizations")) + self.assertEqual(model.__class__, nn.Linear) + self.assertFalse(parametrize.is_parametrized(model)) + self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) + self.assertIn("weight", model._parameters) + + # The model has now two parameters + self.assertEqual(len(list(model.parameters())), 2) + + # Test backward. Should not throw + sgd = torch.optim.SGD(model.parameters(), lr=0.1) + for _ in range(2): + sgd.zero_grad() + loss = (model.weight.T @ model.bias).sum() + loss.backward() + sgd.step() + def test_caching_parametrization(self): r"""Test the caching system of a parametrization""" # Define a couple matrix parametrizations @@ -2479,24 +2782,6 @@ def forward(self, X): Y = model.weight self.assertEqual(id(X), id(Y)) - def test_dtype_parametrization(self): - r"""Test a case that is not allowed when removing a parametrization""" - class ChangeType(nn.Module): - def forward(self, X): - return X.double() - - module = nn.Linear(4, 4).float() - input_ = torch.rand(4).double() - # It is allowed to register a parametrization that changes the dtype - parametrize.register_parametrization(module, "weight", ChangeType()) - module(input_) - # We can remove it leaving the original tensor - parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) - # But leaving it parametrized breaks - parametrize.register_parametrization(module, "weight", ChangeType()) - with self.assertRaisesRegex(ValueError, "changes the dtype"): - parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) - def test_parametrization_same_training_mode(self): r"""Test training mode updated on parametrization registration""" class Identity(nn.Module): @@ -7211,7 +7496,7 @@ def test_cudnn_weight_format(self): weight = all_vars[4] weight_data = weight.data.clone() with torch.no_grad(): - weight.set_(weight_data) + weight.copy_(weight_data) for _ in range(2): with warnings.catch_warnings(record=True) as w: diff --git a/torch/nn/utils/parametrize.py b/torch/nn/utils/parametrize.py index a68a944824773..6765dd938fef9 100644 --- a/torch/nn/utils/parametrize.py +++ b/torch/nn/utils/parametrize.py @@ -2,8 +2,10 @@ from torch.nn.modules.container import ModuleList, ModuleDict, Module from torch.nn.parameter import Parameter from torch import Tensor -from typing import Union, Optional, Iterable, Dict, Tuple + +import collections from contextlib import contextmanager +from typing import Union, Optional, Dict, Tuple, Sequence _cache_enabled = 0 @@ -53,72 +55,208 @@ def cached(): _cache = {} +def _register_parameter_or_buffer(module, name, X): + if isinstance(X, Parameter): + module.register_parameter(name, X) + else: + module.register_buffer(name, X) + + class ParametrizationList(ModuleList): - r"""A sequential container that holds and manages the ``original`` parameter or buffer of - a parametrized :class:`torch.nn.Module`. It is the type of - ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` has been parametrized - with :func:`register_parametrization`. + r"""A sequential container that holds and manages the ``original`` or ``original0``, ``original1``, ... + parameters or buffers of a parametrized :class:`torch.nn.Module`. - .. note :: + It is the type of ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` + has been parametrized with :func:`register_parametrization`. + + If the first registered parmetrization has a ``right_inverse`` that returns one tensor or + does not have a ``right_inverse`` (in which case we assume that ``right_inverse`` is the identity), + it will hold the tensor under the name ``original``. + If it has a ``right_inverse`` that returns more than one tensor, these will be registered as + ``original0``, ``original1``, ... + + .. warning:: This class is used internally by :func:`register_parametrization`. It is documented - here for completeness. It should not be instantiated by the user. + here for completeness. It shall not be instantiated by the user. Args: - modules (iterable): an iterable of modules representing the parametrizations + modules (sequence): sequence of modules representing the parametrizations original (Parameter or Tensor): parameter or buffer that is parametrized """ original: Tensor def __init__( - self, modules: Iterable[Module], original: Union[Tensor, Parameter] + self, modules: Sequence[Module], original: Union[Tensor, Parameter] ) -> None: + # We require this because we need to treat differently the first parametrization + # This should never throw, unless this class is used from the outside + if len(modules) == 0: + raise ValueError("ParametrizationList requires one or more modules.") + super().__init__(modules) - if isinstance(original, Parameter): - self.register_parameter("original", original) - else: - self.register_buffer("original", original) - def set_original_(self, value: Tensor) -> None: - r"""This method is called when assigning to a parametrized tensor. + # In plain words: + # module.weight must keep its dtype and shape. + # Furthermore, if there is no right_inverse or the right_inverse returns a tensor, + # this should be of the same dtype as the original tensor + # + # We check that the following invariants hold: + # X = module.weight + # Y = param.right_inverse(X) + # assert isinstance(Y, Tensor) or + # (isinstance(Y, collections.abc.Sequence) and all(isinstance(t, Tensor) for t in Y)) + # Z = param(Y) if isisntance(Y, Tensor) else param(*Y) + # # Consistency checks + # assert X.dtype == Z.dtype and X.shape == Z.shape + # # If it has one input, this allows to be able to use set_ to be able to + # # move data to/from the original tensor without changing its id (which is what the + # # optimiser uses to track parameters) + # if isinstance(Y, Tensor) + # assert X.dtype == Y.dtype + # Below we use original = X, new = Y + + original_shape = original.shape + original_dtype = original.dtype + + # Compute new + with torch.no_grad(): + new = original + for module in reversed(self): # type: ignore[call-overload] + if hasattr(module, "right_inverse"): + new = module.right_inverse(new) + # else, we assume that right_inverse is the identity + + if not isinstance(new, Tensor) and not isinstance(new, collections.abc.Sequence): + raise ValueError("'right_inverse' must return a Tensor or a Sequence of tensors (list, tuple...). " + f"Got {type(new).__name__}") + + # Set the number of original tensors + self.is_tensor = isinstance(new, Tensor) + self.ntensors = 1 if self.is_tensor else len(new) + + # Register the tensor(s) + if self.is_tensor: + if original.dtype != new.dtype: + raise ValueError( + "When `right_inverse` outputs one tensor, it may not change the dtype.\n" + f"original.dtype: {original.dtype}\n" + f"right_inverse(original).dtype: {new.dtype}" + ) + # Set the original to original so that the user does not need to re-register the parameter + # manually in the optimiser + with torch.no_grad(): + original.set_(new) # type: ignore[call-overload] + _register_parameter_or_buffer(self, "original", original) + else: + for i, originali in enumerate(new): + if not isinstance(originali, Tensor): + raise ValueError("'right_inverse' must return a Tensor or a Sequence of tensors " + "(list, tuple...). " + f"Got element {i} of the sequence with type {type(originali).__name__}.") + + # If the original tensor was a Parameter that required grad, we expect the user to + # add the new parameters to the optimizer after registering the parametrization + # (this is documented) + if isinstance(original, Parameter): + originali = Parameter(originali) + originali.requires_grad_(original.requires_grad) + _register_parameter_or_buffer(self, f"original{i}", originali) + + # Consistency checks: + # Since f : A -> B, right_inverse : B -> A, Z and original should live in B + # Z = forward(right_inverse(original)) + Z = self() + if not isinstance(Z, Tensor): + raise ValueError( + f"A parametrization must return a tensor. Got {type(Z).__name__}." + ) + if Z.dtype != original_dtype: + raise ValueError( + "Registering a parametrization may not change the dtype of the tensor.\n" + f"unparametrized dtype: {original_dtype}\n" + f"parametrized dtype: {Z.dtype}" + ) + if Z.shape != original_shape: + raise ValueError( + "Registering a parametrization may not change the shape of the tensor.\n" + f"unarametrized shape: {original_shape}\n" + f"parametrized shape: {Z.shape}" + ) - It calls the methods ``right_inverse`` (see :func:`register_parametrization`) - of the parametrizations in the inverse order that they have been registered. - Then, it assigns the result to ``self.original``. + def right_inverse(self, value: Tensor) -> None: + r"""Calls the methods ``right_inverse`` (see :func:`register_parametrization`) + of the parametrizations in the inverse order they were registered in. + Then, it stores the result in ``self.original`` if ``right_inverse`` outputs one tensor + or in ``self.original0``, ``self.original1``, ... if it outputs several. Args: value (Tensor): Value to which initialize the module - - Raises: - RuntimeError: if any of the parametrizations do not implement a ``right_inverse`` method """ + # All the exceptions in this function should almost never throw. + # They could throw if, for example, right_inverse function returns a different + # dtype when given a different input, which should most likely be caused by a + # bug in the user's code + with torch.no_grad(): # See https://github.com/pytorch/pytorch/issues/53103 for module in reversed(self): # type: ignore[call-overload] if hasattr(module, "right_inverse"): value = module.right_inverse(value) - else: - raise RuntimeError( - "The parametrization '{}' does not implement a 'right_inverse' method. " - "Assigning to a parametrized tensor is only possible when all the parametrizations " - "implement a 'right_inverse' method.".format(module.__class__.__name__) + # else we assume that right_inverse is the identity + if self.is_tensor: + # These exceptions should only throw when a right_inverse function does not + # return the same dtype for every input, which should most likely be caused by a bug + if not isinstance(value, Tensor): + raise ValueError( + f"`right_inverse` should return a tensor. Got {type(value).__name__}" + ) + if value.dtype != self.original.dtype: + raise ValueError( + f"The tensor returned by `right_inverse` has dtype {value.dtype} " + f"while `original` has dtype {self.original.dtype}" + ) + # We know that the result is going to have the same dtype + self.original.set_(value) # type: ignore[call-overload] + else: + if not isinstance(value, collections.abc.Sequence): + raise ValueError( + "'right_inverse' must return a sequence of tensors. " + f"Got {type(value).__name__}." ) - self.original.copy_(value) + if len(value) != self.ntensors: + raise ValueError( + "'right_inverse' must return a sequence of tensors of length " + f"{self.ntensors}. Got a sequence of lenght {len(value)}." + ) + for i, tensor in enumerate(value): + original_i = getattr(self, f"original{i}") + if not isinstance(tensor, Tensor): + raise ValueError( + f"`right_inverse` must return a sequence of tensors. " + f"Got element {i} of type {type(tensor).__name__}" + ) + if original_i.dtype != tensor.dtype: + raise ValueError( + f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} " + f"while `original{i}` has dtype {original_i.dtype}" + ) + original_i.set_(tensor) def forward(self) -> Tensor: - x = self.original - for module in self: + # Unpack the originals for the first parametrization + if self.is_tensor: + x = self[0](self.original) + else: + originals = (getattr(self, f"original{i}") for i in range(self.ntensors)) + x = self[0](*originals) + # It's not possible to call self[1:] here, so we have to be a bit more cryptic + for module in list(self._modules.values())[1:]: x = module(x) - if x.size() != self.original.size(): - raise RuntimeError( - "The parametrization may not change the size of the parametrized tensor. " - "Size of original tensor: {} " - "Size of parametrized tensor: {}".format(self.original.size(), x.size()) - ) return x def _inject_new_class(module: Module) -> None: - r"""Sets up the parametrization mechanism used by parametrizations. + r"""Sets up a module to be parametrized. This works by substituting the class of the module by a class that extends it to be able to inject a property @@ -137,7 +275,7 @@ def getstate(self): ) param_cls = type( - "Parametrized{}".format(cls.__name__), + f"Parametrized{cls.__name__}", (cls,), { "__getstate__": getstate, @@ -178,11 +316,10 @@ def get_parametrized(self) -> Tensor: return parametrization() def set_original(self, value: Tensor) -> None: - self.parametrizations[tensor_name].set_original_(value) + self.parametrizations[tensor_name].right_inverse(value) setattr(module.__class__, tensor_name, property(get_parametrized, set_original)) - def register_parametrization( module: Module, tensor_name: str, parametrization: Module ) -> Module: @@ -191,12 +328,12 @@ def register_parametrization( Assume that ``tensor_name="weight"`` for simplicity. When accessing ``module.weight``, the module will return the parametrized version ``parametrization(module.weight)``. If the original tensor requires a gradient, the backward pass will differentiate - through the :attr:`parametrization`, and the optimizer will update the tensor accordingly. + through :attr:`parametrization`, and the optimizer will update the tensor accordingly. The first time that a module registers a parametrization, this function will add an attribute ``parametrizations`` to the module of type :class:`~ParametrizationList`. - The list of parametrizations on a tensor will be accessible under + The list of parametrizations on the tensor ``weight`` will be accessible under ``module.parametrizations.weight``. The original tensor will be accessible under @@ -205,8 +342,8 @@ def register_parametrization( Parametrizations may be concatenated by registering several parametrizations on the same attribute. - The training mode of the registered parametrizations are updated on registration - if necessary to match the training mode of the host module + The training mode of a registered parametrization is updated on registration + to match the training mode of the host module Parametrized parameters and buffers have an inbuilt caching system that can be activated using the context manager :func:`cached`. @@ -215,16 +352,37 @@ def register_parametrization( .. code-block:: python - def right_inverse(self, X: Tensor) -> Tensor + def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] - If :attr:`parametrization` implements this method, it will be possible to assign - to the parametrized tensor. This may be used to initialize the tensor, as shown in the example. + If this method is not implemented, it defaults to the identity. + This method is called on the unparametrized tensor when the first parametrization + is registered. In most situations, ``right_inverse`` will be a function such that ``forward(right_inverse(X)) == X`` (see `right inverse `_). Sometimes, when the parametrization is not surjective, it may be reasonable - to relax this, as shown in the example below. + to relax this. + This may be used to initialize the tensor, as shown in the example below. + + It is possible for the first parametrization to depend on several inputs. + This may be implemented returning a tuple of tensors from ``right_inverse`` + (see the example implementation of a ``RankOne`` parametrization below). + + In this case, the unconstrained tensors are also located under ``module.parametrizations.weight`` + with names ``original0``, ``original1``,... + + .. note:: + + Whenever a parametrization is registered, both its forward and backward method will be called + once to perform a number of consistency checks. + + .. warning:: + + If a parametrization depends on several inputs, :func:`~register_parametrization` + will register a number of new parameters. If such parametrization is registered + after the optimizer is created, these new parameters will need to be added manually + to the optimizer. See :meth:`torch.Optimizer.add_param_group`. Args: module (nn.Module): module on which to register the parametrization @@ -232,24 +390,22 @@ def right_inverse(self, X: Tensor) -> Tensor the parametrization parametrization (nn.Module): the parametrization to register - Returns: - Module: module - Raises: ValueError: if the module does not have a parameter or a buffer named :attr:`tensor_name` Examples: >>> import torch + >>> import torch.nn as nn >>> import torch.nn.utils.parametrize as P >>> - >>> class Symmetric(torch.nn.Module): + >>> class Symmetric(nn.Module): >>> def forward(self, X): >>> return X.triu() + X.triu(1).T # Return a symmetric matrix >>> >>> def right_inverse(self, A): >>> return A.triu() >>> - >>> m = torch.nn.Linear(5, 5) + >>> m = nn.Linear(5, 5) >>> P.register_parametrization(m, "weight", Symmetric()) >>> print(torch.allclose(m.weight, m.weight.T)) # m.weight is now symmetric True @@ -258,15 +414,80 @@ def right_inverse(self, X: Tensor) -> Tensor >>> m.weight = A # Initialize the weight to be the symmetric matrix A >>> print(torch.allclose(m.weight, A)) True + + >>> class RankOne(nn.Module): + >>> def forward(self, x, y): + >>> # Form a rank 1 matrix multiplying two vectors + >>> return x.unsqueeze(-1) @ y.unsqueeze(-2) + >>> + >>> def right_inverse(self, Z): + >>> # Project Z onto the rank 1 matrices + >>> U, S, Vh = torch.linalg.svd(Z, full_matrices=False) + >>> # Return rescaled singular vectors + >>> s0_sqrt = S[0].sqrt().unsqueeze(-1) + >>> return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt + >>> + >>> linear_rank_one = P.register_parametrization(nn.Linear(4, 4), "weight", RankOne()) + >>> print(torch.linalg.matrix_rank(linear_rank_one.weight).item()) + 1 + """ parametrization.train(module.training) if is_parametrized(module, tensor_name): - # Just add the new parametrization to the parametrization list - module.parametrizations[tensor_name].append(parametrization) # type: ignore[index, union-attr] + # Correctness checks. + # If A is the space of tensors with shape and dtype equal to module.weight + # we check that parametrization.forward and parametrization.right_inverse are + # functions from A to A + + Y = getattr(module, tensor_name) + X = parametrization(Y) + if not isinstance(X, Tensor): + raise ValueError( + f"A parametrization must return a tensor. Got {type(X).__name__}." + ) + if X.dtype != Y.dtype: + raise ValueError( + "Registering a parametrization may not change the dtype of the tensor.\n" + f"module.{tensor_name}.dtype: {Y.dtype}\n" + f"parametrization(module.{tensor_name}).dtype: {X.dtype}" + ) + if X.shape != Y.shape: + raise ValueError( + "Registering a parametrization may not change the shape of the tensor.\n" + f"module.{tensor_name}.shape: {Y.shape}\n" + f"parametrization(module.{tensor_name}).shape: {X.shape}" + ) + if hasattr(parametrization, "right_inverse"): + Z = parametrization.right_inverse(X) # type: ignore[operator] + if not isinstance(Z, Tensor): + raise ValueError( + f"parametrization.right_inverse must return a tensor. Got: {type(Z).__name__}" + ) + if Z.dtype != Y.dtype: + raise ValueError( + "The tensor returned by parametrization.right_inverse must have the same dtype " + f"as module.{tensor_name}.\n" + f"module.{tensor_name}.dtype: {Y.dtype}\n" + f"returned dtype: {Z.dtype}" + ) + if Z.shape != Y.shape: + raise ValueError( + "The tensor returned by parametrization.right_inverse must have the same shape " + f"as module.{tensor_name}.\n" + f"module.{tensor_name}.shape: {Y.shape}\n" + f"returned shape: {Z.shape}" + ) + # else right_inverse is assumed to be the identity + + # add the new parametrization to the parametrization list + assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy + module.parametrizations[tensor_name].append(parametrization) elif tensor_name in module._buffers or tensor_name in module._parameters: # Set the parametrization mechanism # Fetch the original buffer or parameter original = getattr(module, tensor_name) + # We create this early to check for possible errors + parametrizations = ParametrizationList([parametrization], original) # Delete the previous parameter or buffer delattr(module, tensor_name) # If this is the first parametrization registered on the module, @@ -274,18 +495,17 @@ def right_inverse(self, X: Tensor) -> Tensor if not is_parametrized(module): # Change the class _inject_new_class(module) - # Inject the a ``ModuleDict`` into the instance under module.parametrizations + # Inject a ``ModuleDict`` into the instance under module.parametrizations module.parametrizations = ModuleDict() # Add a property into the class _inject_property(module, tensor_name) # Add a ParametrizationList - module.parametrizations[tensor_name] = ParametrizationList( # type: ignore[assignment, index, operator] - [parametrization], original - ) + assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy + module.parametrizations[tensor_name] = parametrizations else: raise ValueError( - "Module '{}' does not have a parameter, a buffer, or a " - "parametrized element with name '{}'".format(module, tensor_name) + f"Module '{module}' does not have a parameter, a buffer, or a " + f"parametrized element with name '{tensor_name}'" ) return module @@ -321,6 +541,7 @@ def remove_parametrizations( of the tensor. - If ``leave_parametrized=False``, ``module[tensor_name]`` will be set to the unparametrised tensor in ``module.parametrizations[tensor_name].original``. + This is only possible when the parametrization depends on just one tensor. Args: module (nn.Module): module from which remove the parametrization @@ -333,44 +554,44 @@ def remove_parametrizations( Raises: ValueError: if ``module[tensor_name]`` is not parametrized - ValueError: if ``leave_parametrized=True`` and the parametrization changes the size or dtype - of the tensor + ValueError: if ``leave_parametrized=False`` and the parametrization depends on several tensors """ if not is_parametrized(module, tensor_name): - raise ValueError( - "Module {} does not have a parametrization on {}".format( - module, tensor_name - ) - ) + raise ValueError(f"Module {module} does not have a parametrization on {tensor_name}") # Fetch the original tensor - original = module.parametrizations[tensor_name].original # type: ignore[index, union-attr] - if leave_parametrized: - with torch.no_grad(): - t = getattr(module, tensor_name) - # If they have the same dtype, we reuse the original tensor. - # We do this so that the parameter does not to change the id() - # This way the user does not need to update the optimizer - if t.dtype == original.dtype: + assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy + parametrizations = module.parametrizations[tensor_name] + if parametrizations.is_tensor: + original = parametrizations.original + if leave_parametrized: + with torch.no_grad(): + t = getattr(module, tensor_name) + # We know they have the same dtype because we have checked this when registering the + # parametrizations. As such, we can use set_ + # We do this so that the parameter does not to change the id() + # This way the user does not need to update the optimizer with torch.no_grad(): original.set_(t) + else: + if leave_parametrized: + # We cannot use no_grad because we need to know whether one or more + # original tensors required grad + t = getattr(module, tensor_name) + # We'll have to trust the user to add it to the optimizer + original = Parameter(t) if t.requires_grad else t else: - raise ValueError( - "The parametrization changes the dtype of the tensor from {} to {}. " - "It is not supported to leave the tensor parametrized (`leave_parametrized=True`) " - "in this case.".format(original.dtype, t.dtype) - ) + raise ValueError("Cannot leave unparametrized (`leave_parametrized=False`) a tensor " + "that is parametrized in terms of a sequence of tensors.") + # Delete the property that manages the parametrization delattr(module.__class__, tensor_name) # Delete the ParametrizationList - del module.parametrizations[tensor_name] # type: ignore[operator, union-attr] + del module.parametrizations[tensor_name] # Restore the parameter / buffer into the main class - if isinstance(original, Parameter): - module.register_parameter(tensor_name, original) - else: - module.register_buffer(tensor_name, original) + _register_parameter_or_buffer(module, tensor_name, original) # Roll back the parametrized class if no other buffer or parameter # is currently parametrized in this class From ef13341a8d1128a5fbd54716c89358cbd20c0a1d Mon Sep 17 00:00:00 2001 From: yanbing-j Date: Mon, 14 Jun 2021 11:54:22 -0700 Subject: [PATCH 087/305] upgrade onednn to v2.2.3 (#57928) Summary: This PR is to upgrade onednn to v2.2.3 (including v2.2 and v2.2.3 changes) which has the following main changes about CPU: v2.2 changes: Improved performance of compute functionality for future Intel Core processor with Intel AVX2 and Intel DL Boost instructions support (code name Alder Lake). Improved fp32 inner product forward propagation performance for processors with Intel AVX-512 support. Improved dnnl_gemm performance for cases with n=1 on all supported processors. v2.2.3 changes: Fixed a bug in int8 depthwise convolution ptimitive with groups and 1d spatial size for processors with Intel AVX-512 and Intel AVX2 support Fixed correctness issue for PReLU primitive on Intel Processor Graphics Fixed corretness issue in reorder for blocked layouts with zero padding Improved performance of weights reorders used by BRGEMM-based convolution primitive for processors with Intel AVX-512 support More changes can be found in https://github.com/oneapi-src/oneDNN/releases. Ideep used version is pytorch-rls-v2.2.3. OneDNN used version is v2.2.3. Pull Request resolved: https://github.com/pytorch/pytorch/pull/57928 Reviewed By: bdhirsh Differential Revision: D29037857 Pulled By: VitalyFedyunin fbshipit-source-id: db74534858bdcf5d6c7dcf58e224fc756188bc31 --- third_party/ideep | 2 +- third_party/mkl-dnn.BUILD | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/third_party/ideep b/third_party/ideep index f9468ff1a3d60..9ca27bbfd88fa 160000 --- a/third_party/ideep +++ b/third_party/ideep @@ -1 +1 @@ -Subproject commit f9468ff1a3d601b509ebe2c17d2ed0a58dffacee +Subproject commit 9ca27bbfd88fa1469cbf0467bd6f14cd1738fa40 diff --git a/third_party/mkl-dnn.BUILD b/third_party/mkl-dnn.BUILD index 54f5fffeb59e2..4be75e2080ec7 100644 --- a/third_party/mkl-dnn.BUILD +++ b/third_party/mkl-dnn.BUILD @@ -16,9 +16,9 @@ template_rule( out = "include/oneapi/dnnl/dnnl_version.h", substitutions = { "@DNNL_VERSION_MAJOR@": "2", - "@DNNL_VERSION_MINOR@": "1", - "@DNNL_VERSION_PATCH@": "2", - "@DNNL_VERSION_HASH@": "98be7e8afa711dc9b66c8ff3504129cb82013cdb", + "@DNNL_VERSION_MINOR@": "2", + "@DNNL_VERSION_PATCH@": "3", + "@DNNL_VERSION_HASH@": "7336ca9f055cf1bfa13efb658fe15dc9b41f0740", }, ) @@ -45,6 +45,7 @@ cc_library( "src/cpu/**/*.hpp", "src/cpu/**/*.h", "src/common/*.hpp", + "src/common/ittnotify/jitprofiling.h", ], exclude=[ "src/cpu/aarch64/**/*.hpp", "src/cpu/aarch64/**/*.h", From 8e92a3a8b026178d65fb995e2a4b750fea93a093 Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Mon, 14 Jun 2021 12:57:30 -0700 Subject: [PATCH 088/305] [docs] Add pickle security warning to package docs (#59959) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59959 **Summary** This commit replaces the warning on the `torch.package` documentation page about the module not being publicly released (which will no longer be true as of 1.9) with one that warns about security issues caused by the use of the `pickle` module. **Test Plan** 1) Built the docs locally. 2) Continuous integration. Captura de Pantalla 2021-06-14 a la(s) 11 22 05 a m Test Plan: Imported from OSS Reviewed By: suo Differential Revision: D29108429 Pulled By: SplitInfinity fbshipit-source-id: 3a0aeac0dc804a31203bc5071efb1c5bd6ef9725 --- docs/source/package.rst | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/source/package.rst b/docs/source/package.rst index ed4d6457846bf..d21f7b016efce 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -2,11 +2,6 @@ torch.package ============= - -.. warning:: - - This module is experimental and has not yet been publicly released. - ``torch.package`` adds support for creating hermetic packages containing arbitrary PyTorch code. These packages can be saved, shared, used to load and execute models at a later date or on a different machine, and can even be deployed to production using @@ -16,6 +11,16 @@ This document contains tutorials, how-to guides, explanations, and an API refere will help you learn more about ``torch.package`` and how to use it. +.. warning:: + + This module depends on the ``pickle`` module which is is not secure. Only unpackage data you trust. + + It is possible to construct malicious pickle data which will **execute arbitrary code during unpickling**. + Never unpackage data that could have come from an untrusted source, or that could have been tampered with. + + For more information, review the `documentation `_ for the ``pickle`` module. + + .. contents:: :local: :depth: 2 From f9ec86a6c6863862194b1d5f76f5950ff3c4cdd7 Mon Sep 17 00:00:00 2001 From: Emilio Castillo Date: Mon, 14 Jun 2021 13:44:28 -0700 Subject: [PATCH 089/305] External stream (#59527) Summary: Previous is https://github.com/pytorch/pytorch/issues/57781 We add now two CUDA bindings to avoid using ctypes to fix a windows issue. However, we use ctypes to allocate the stream and create its pointer (we can do this with a 0-dim tensor too if it feels better). CC. ezyang rgommers ngimel mruberry Pull Request resolved: https://github.com/pytorch/pytorch/pull/59527 Reviewed By: albanD Differential Revision: D29053062 Pulled By: ezyang fbshipit-source-id: 661e7e58de98b1bdb7a0871808cd41d91fe8f13f --- .../hip/impl/HIPStreamMasqueradingAsCUDA.h | 5 ++ c10/core/Stream.h | 35 +++++++--- c10/cuda/CUDAStream.cpp | 69 ++++++++++++++----- c10/cuda/CUDAStream.h | 10 +++ caffe2/contrib/opencl/context.h | 2 +- caffe2/core/context.h | 2 +- caffe2/core/context_base.h | 2 +- caffe2/ideep/utils/ideep_context.h | 2 +- test/test_cuda.py | 36 ++++++++++ torch/_C/__init__.pyi.in | 2 +- torch/csrc/cuda/Stream.cpp | 15 ++-- torch/csrc/cuda/shared/cudart.cpp | 6 ++ torch/cuda/streams.py | 23 +++++++ torch/utils/hipify/cuda_to_hip_mappings.py | 9 +++ 14 files changed, 184 insertions(+), 34 deletions(-) diff --git a/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h b/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h index 13a88e6aea5f1..417943da3777d 100644 --- a/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h +++ b/aten/src/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h @@ -90,6 +90,11 @@ inline getStreamFromPoolMasqueradingAsCUDA(const bool isHighPriority = false, De return HIPStreamMasqueradingAsCUDA(getStreamFromPool(isHighPriority, device)); } +HIPStreamMasqueradingAsCUDA +inline getStreamFromExternalMasqueradingAsCUDA(hipStream_t ext_stream, DeviceIndex device) { + return HIPStreamMasqueradingAsCUDA(getStreamFromExternal(ext_stream, device)); +} + inline HIPStreamMasqueradingAsCUDA getDefaultHIPStreamMasqueradingAsCUDA(DeviceIndex device_index = -1) { return HIPStreamMasqueradingAsCUDA(getDefaultHIPStream(device_index)); } diff --git a/c10/core/Stream.h b/c10/core/Stream.h index 446d590b588c0..d0abcb1e21269 100644 --- a/c10/core/Stream.h +++ b/c10/core/Stream.h @@ -12,7 +12,7 @@ namespace c10 { /// numbering system which is not visible to the user. HOWEVER, we /// guarantee that StreamId 0 is always a valid stream, and corresponds /// to some sort of "default" stream. -using StreamId = int32_t; +using StreamId = int64_t; // NB: I decided not to call the above StreamIndex to avoid confusion with // DeviceIndex. This way, you access device index with index(), and stream id @@ -127,21 +127,38 @@ class C10_API Stream final { // that the bitmasking code below is updated accordingly! static_assert(sizeof(DeviceType) == 1, "DeviceType is not 8-bit"); static_assert(sizeof(DeviceIndex) == 1, "DeviceIndex is not 8-bit"); - static_assert(sizeof(StreamId) == 4, "DeviceIndex is not 32-bit"); + static_assert(sizeof(StreamId) == 8, "StreamId is not 64-bit"); // Concat these together into a 64-bit integer // See Note [Hazard when concatenating signed integers] uint64_t bits = static_cast(static_cast(device_type())) - << 48 | - static_cast(static_cast(device_index())) << 32 | - static_cast(static_cast(id())); + << 56 | + static_cast(static_cast(device_index())) << 48 | + // Remove the sign extension part of the 64-bit address because + // the id might be used to hold a pointer. + (static_cast(id()) & ((1ull << 48) - 1)); + TORCH_INTERNAL_ASSERT( + static_cast((bits >> 48) & 0xFFull) == device_index(), + "DeviceIndex is not correctly packed"); + TORCH_INTERNAL_ASSERT( + static_cast((bits >> 56)) == device_type(), + "DeviceType is not correctly packed"); + // Re-extend the sign of stream_id for checking + uint64_t mask = (1ull << 47); + TORCH_INTERNAL_ASSERT( + static_cast(((bits & 0xFFFFFFFFFFFFull) ^ mask) - mask) == + id(), + "DeviceType is not correctly packed"); return bits; } static Stream unpack(uint64_t bits) { - const auto stream_id = static_cast(bits & 0xFFFFFFFFull); - bits >>= 32; - const auto device_index = static_cast(bits & 0xFFFFull); - bits >>= 16; + // Re-extend the sign of stream_id + uint64_t mask = (1ull << 47); + const auto stream_id = + (static_cast(bits & 0xFFFFFFFFFFFFull) ^ mask) - mask; + bits >>= 48; + const auto device_index = static_cast(bits & 0xFFull); + bits >>= 8; const auto device_type = static_cast(bits); TORCH_CHECK(isValidDeviceType(device_type)); // Unfortunately, we can't check if the StreamId is valid here; it diff --git a/c10/cuda/CUDAStream.cpp b/c10/cuda/CUDAStream.cpp index c32580bf79908..4ef4e9eb4e878 100644 --- a/c10/cuda/CUDAStream.cpp +++ b/c10/cuda/CUDAStream.cpp @@ -10,6 +10,7 @@ #include #include +#include namespace c10 { namespace cuda { @@ -41,6 +42,7 @@ static DeviceIndex num_gpus = -1; static constexpr int kStreamsPerPoolBits = 5; static constexpr int kStreamsPerPool = 1 << kStreamsPerPoolBits; static constexpr unsigned int kDefaultFlags = cudaStreamNonBlocking; +static constexpr int kStreamTypeBits = 3; // Note: lower numbers are higher priorities, zero is default priority static int kHighPriority = -1; @@ -73,13 +75,13 @@ static std::array // ~~~~~~~~~~~~~~~~~~~~~~~~~~ // How do we assign stream IDs? // -// -- 25 bits -- -- 2 bits -- -- 5 bits ----- -// zeros StreamIdType stream id index +// -- 57 bits -- -- 5 bits ----- -- 3 bits -- +// zeros stream id index StreamIdType // // Where StreamIdType: -// 00 = default stream -// 01 = low priority stream -// 10 = high priority stream +// 000 = default stream or externally allocated if id[63:3] != 0 +// 001 = low priority stream +// 010 = high priority stream // // This is not really for efficiency; it's just easier to write the code // to extract the index if we do this with bitmasks :) @@ -95,11 +97,16 @@ static std::array // could work around this with something like // https://stackoverflow.com/questions/13150449/efficient-unsigned-to-signed-cast-avoiding-implementation-defined-behavior // but it seems a bit overkill for this. - +// +// Also, external managed stream pointers (cudaStream_t) can be directly stored +// in the Id field so in this case, we need to check the stream alignment. +// The IdType uses an additional bit to match with the 64-bit address alignment +// making easy to identify an external stream when its value (X & 7) > 0 enum class StreamIdType : uint8_t { DEFAULT = 0x0, LOW = 0x1, HIGH = 0x2, + EXT = 0x3, }; std::ostream& operator<<(std::ostream& stream, StreamIdType s) { @@ -113,6 +120,9 @@ std::ostream& operator<<(std::ostream& stream, StreamIdType s) { case StreamIdType::HIGH: stream << "HIGH"; break; + case StreamIdType::EXT: + stream << "EXT"; + break; default: stream << static_cast(s); break; @@ -120,21 +130,29 @@ std::ostream& operator<<(std::ostream& stream, StreamIdType s) { return stream; } -// StreamId is 32-bit, so we can just rely on regular promotion rules. +// StreamId is 64-bit, so we can just rely on regular promotion rules. // We rely on streamIdIndex and streamIdType being non-negative; // see Note [Hazard when concatenating signed integers] static inline StreamIdType streamIdType(StreamId s) { - return static_cast(s >> kStreamsPerPoolBits); + int mask_for_type = (1 << kStreamTypeBits) - 1; + if (s && ((s & mask_for_type) == 0)) { + // Externally allocated streams have their id being the cudaStream_ptr + // so the bits corresponding to the type will be 0 and will collide with + // the default stream. + return StreamIdType::EXT; + } + return static_cast(s & mask_for_type); } static inline size_t streamIdIndex(StreamId s) { - return static_cast(s & ((1 << kStreamsPerPoolBits) - 1)); + return static_cast( + (s >> kStreamTypeBits) & ((1 << kStreamsPerPoolBits) - 1)); } StreamId makeStreamId(StreamIdType st, size_t si) { - return (static_cast(st) << kStreamsPerPoolBits) | - static_cast(si); + return (static_cast(si) << kStreamTypeBits) | + static_cast(st); } template @@ -251,7 +269,7 @@ static void initCUDAStreamsOnce() { // Helper to verify the GPU index is valid static inline void check_gpu(DeviceIndex device_index) { - AT_ASSERT(device_index >= 0 && device_index < num_gpus); + TORCH_INTERNAL_ASSERT(device_index >= 0 && device_index < num_gpus); } // Helper to determine the index of the stream to return @@ -305,9 +323,16 @@ CUDAStream CUDAStream_fromInternals(const LeakyStreamInternals* ptr) { } // anonymous namespace cudaStream_t CUDAStream::stream() const { - auto ptr = CUDAStream_internals(*this); - AT_ASSERT(ptr); - return ptr->stream; + int64_t stream_id = unwrap().id(); + if (streamIdType(stream_id) == StreamIdType::EXT) { + // In this case this is a externally allocated stream + // we don't need to manage its life cycle + return reinterpret_cast(stream_id); + } else { + auto ptr = CUDAStream_internals(*this); + TORCH_INTERNAL_ASSERT(ptr); + return ptr->stream; + } } // Returns a stream from the requested pool @@ -334,6 +359,18 @@ CUDAStream getStreamFromPool( return CUDAStream_fromInternals(&low_priority_streams[device_index][idx]); } +CUDAStream getStreamFromExternal( + cudaStream_t ext_stream, + DeviceIndex device_index) { + return CUDAStream( + CUDAStream::UNCHECKED, + // The stream pointer will be the actual id + Stream( + Stream::UNSAFE, + c10::Device(DeviceType::CUDA, device_index), + reinterpret_cast(ext_stream))); +} + CUDAStream getDefaultCUDAStream(DeviceIndex device_index) { initCUDAStreamsOnce(); if (device_index == -1) { @@ -354,7 +391,7 @@ CUDAStream getCurrentCUDAStream(DeviceIndex device_index) { void setCurrentCUDAStream(CUDAStream stream) { initCUDAStreamsOnce(); auto ptr = CUDAStream_internals(stream); - AT_ASSERT(ptr); + TORCH_INTERNAL_ASSERT(ptr); current_streams[ptr->device_index] = ptr; } diff --git a/c10/cuda/CUDAStream.h b/c10/cuda/CUDAStream.h index 2e00ecc4a02c9..1b6075bee10f5 100644 --- a/c10/cuda/CUDAStream.h +++ b/c10/cuda/CUDAStream.h @@ -195,6 +195,16 @@ class C10_CUDA_API CUDAStream { TORCH_API CUDAStream getStreamFromPool(const bool isHighPriority = false, DeviceIndex device = -1); +/** + * Get a CUDAStream from a externally allocated one. + * + * This is mainly for interoperability with different libraries where we + * want to operate on a non-torch allocated stream for data exchange or similar + * purposes + */ +TORCH_API CUDAStream +getStreamFromExternal(cudaStream_t ext_stream, DeviceIndex device_index); + /** * Get the default CUDA stream, for the passed CUDA device, or for the * current device if no device index is passed. The default stream is diff --git a/caffe2/contrib/opencl/context.h b/caffe2/contrib/opencl/context.h index 15bfda2203f06..b1e61c2124adc 100644 --- a/caffe2/contrib/opencl/context.h +++ b/caffe2/contrib/opencl/context.h @@ -64,7 +64,7 @@ class OpenCLContext final { CopyBytes(n * meta.itemsize(), src, dst); } - void SwitchToDevice(int a, ...) { + void SwitchToDevice(int64_t a, ...) { auto& ctx = GetSingleton(); CAFFE_ENFORCE(a < ctx.devices.size()); ctx.device = ctx.devices[a]; diff --git a/caffe2/core/context.h b/caffe2/core/context.h index d5fe10820152c..27e1fb79fa9fe 100644 --- a/caffe2/core/context.h +++ b/caffe2/core/context.h @@ -63,7 +63,7 @@ class TORCH_API CPUContext final : public BaseContext { ~CPUContext() noexcept override {} - inline void SwitchToDevice(int /*stream_id*/) override {} + inline void SwitchToDevice(int64_t /*stream_id*/) override {} using BaseContext::SwitchToDevice; diff --git a/caffe2/core/context_base.h b/caffe2/core/context_base.h index dfc1504e2092d..cc8cc4c5bb608 100644 --- a/caffe2/core/context_base.h +++ b/caffe2/core/context_base.h @@ -42,7 +42,7 @@ class TORCH_API BaseContext { /* Sorry for the naming, will get rid of this in future diff */ virtual DeviceType device_type() const = 0; - virtual void SwitchToDevice(int /*stream_id*/) = 0; + virtual void SwitchToDevice(int64_t /*stream_id*/) = 0; inline void SwitchToDevice() { SwitchToDevice(0); diff --git a/caffe2/ideep/utils/ideep_context.h b/caffe2/ideep/utils/ideep_context.h index d0f1207a08f69..7df6763b1baf7 100644 --- a/caffe2/ideep/utils/ideep_context.h +++ b/caffe2/ideep/utils/ideep_context.h @@ -23,7 +23,7 @@ class IDEEPContext final : public BaseContext { ~IDEEPContext() noexcept override {} - inline void SwitchToDevice(int /*stream_id*/) {} + inline void SwitchToDevice(int64_t /*stream_id*/) {} using BaseContext::SwitchToDevice; inline void WaitEvent(const Event& ev) { diff --git a/test/test_cuda.py b/test/test_cuda.py index 328aa1ad5a132..6f35cbf2ffe63 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -1,6 +1,8 @@ from itertools import repeat, chain, product from typing import NamedTuple import collections +import contextlib +import ctypes import gc import io import pickle @@ -1312,6 +1314,40 @@ def test_record_stream_on_shifted_view(self): self.assertNotEqual(try_realloc.data_ptr(), data_ptr) + @contextlib.contextmanager + def _get_external_stream(self, device): + cudart = torch.cuda.cudart() + stream = ctypes.c_ulonglong(0) + stream_p = ctypes.POINTER(ctypes.c_void_p)(stream) + stream_p_int = ctypes.cast(stream_p, ctypes.c_void_p).value + with device: + try: + out = cudart.cudaStreamCreate(stream_p_int) + self.assertEqual(out, 0) + self.assertNotEqual(stream.value, 0) + yield stream.value + finally: + out = cudart.cudaStreamDestroy(stream.value) + self.assertEqual(out, 0) + + @skipIfRocm + def test_external_streams(self): + device = torch.cuda.device(0) + with self._get_external_stream(device) as stream_v: + ext_stream = torch.cuda.streams.ExternalStream(stream_v) + self.assertEqual(stream_v, ext_stream.cuda_stream) + self.assertEqual(ext_stream.device.index, device.idx) + + @skipIfRocm + @unittest.skipIf(not TEST_MULTIGPU, "detected only one GPU") + def test_external_streams_multi_device(self): + device = torch.cuda.device(1) + with self._get_external_stream(device) as stream_v: + ext_stream = torch.cuda.streams.ExternalStream( + stream_v, device=device) + self.assertEqual(stream_v, ext_stream.cuda_stream) + self.assertEqual(ext_stream.device.index, device.idx) + def test_noncontiguous_pinned_memory(self): # See issue #3266 x = torch.arange(0, 10).view((2, 5)) diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in index ec390cec9be47..1fa986b0fd72e 100644 --- a/torch/_C/__init__.pyi.in +++ b/torch/_C/__init__.pyi.in @@ -840,7 +840,7 @@ class _CudaStreamBase: cuda_stream: _int priority: _int - def __new__(self, priority: _int = 0, _cdata: _int = 0) -> _CudaStreamBase: ... + def __new__(self, priority: _int = 0, _cdata: _int = 0, stream_ptr: _int = 0) -> _CudaStreamBase: ... def query(self) -> _bool: ... def synchronize(self) -> None: ... def priority_range(self) -> Tuple[_int, _int]: ... diff --git a/torch/csrc/cuda/Stream.cpp b/torch/csrc/cuda/Stream.cpp index 254ddc6508bbf..f358e59728347 100644 --- a/torch/csrc/cuda/Stream.cpp +++ b/torch/csrc/cuda/Stream.cpp @@ -22,11 +22,12 @@ static PyObject * THCPStream_pynew( int priority = 0; uint64_t cdata = 0; + uint64_t stream_ptr = 0; // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays) - static char *kwlist[] = {"priority", "_cdata", nullptr}; + static char *kwlist[] = {"priority", "_cdata", "stream_ptr", nullptr}; if (!PyArg_ParseTupleAndKeywords( - args, kwargs, "|iK", kwlist, &priority, &cdata)) { + args, kwargs, "|iKK", kwlist, &priority, &cdata, &stream_ptr)) { return nullptr; } @@ -35,11 +36,17 @@ static PyObject * THCPStream_pynew( return nullptr; } + if (stream_ptr) { + TORCH_CHECK(priority == 0, "Priority was explicitly set for a external stream") + } + at::cuda::CUDAStream stream = cdata ? at::cuda::CUDAStream::unpack(cdata) : - at::cuda::getStreamFromPool( - /* isHighPriority */ priority < 0 ? true : false); + stream_ptr ? + at::cuda::getStreamFromExternal(reinterpret_cast(stream_ptr), current_device) : + at::cuda::getStreamFromPool( + /* isHighPriority */ priority < 0 ? true : false); THCPStream* self = (THCPStream *)ptr.get(); self->cdata = stream.pack(); diff --git a/torch/csrc/cuda/shared/cudart.cpp b/torch/csrc/cuda/shared/cudart.cpp index 30a43bed05346..9dfa57357f046 100644 --- a/torch/csrc/cuda/shared/cudart.cpp +++ b/torch/csrc/cuda/shared/cudart.cpp @@ -36,6 +36,12 @@ void initCudartBindings(PyObject* module) { cudart.def("cuda" "HostUnregister", [](uintptr_t ptr) -> cudaError_t { return cudaHostUnregister((void*)ptr); }); + cudart.def("cuda" "StreamCreate", [](uintptr_t ptr) -> cudaError_t { + return cudaStreamCreate((cudaStream_t*)ptr); + }); + cudart.def("cuda" "StreamDestroy", [](uintptr_t ptr) -> cudaError_t { + return cudaStreamDestroy((cudaStream_t)ptr); + }); #ifndef __HIP_PLATFORM_HCC__ cudart.def("cuda" "ProfilerInitialize", cudaProfilerInitialize); #endif diff --git a/torch/cuda/streams.py b/torch/cuda/streams.py index 55b5b10e35537..0f983728f630a 100644 --- a/torch/cuda/streams.py +++ b/torch/cuda/streams.py @@ -112,6 +112,29 @@ def __repr__(self): .format(self.device, self.cuda_stream)) +class ExternalStream(Stream): + r"""Wrapper around an externally allocated CUDA stream. + + This class is used to wrap streams allocated in other libraries in order + to facilitate data exchange and multi-library interactions. + + .. note:: This class doesn't manage the stream life-cycle, it is the user + responsibility to keep the referenced stream alive while this class is + being used. + + Args: + stream_ptr(int): Integer representation of the `cudaStream_t` value. + allocated externally. + device(torch.device or int, optional): the device where the stream + was originally allocated. if device is specified incorrectly, + subsequent launches using this stream may fail. + """ + + def __new__(cls, stream_ptr, device=None, **kwargs): + with torch.cuda.device(device): + return super(Stream, cls).__new__(cls, stream_ptr=stream_ptr, **kwargs) + + class Event(torch._C._CudaEventBase): r"""Wrapper around a CUDA event. diff --git a/torch/utils/hipify/cuda_to_hip_mappings.py b/torch/utils/hipify/cuda_to_hip_mappings.py index 3f901b5115138..8c15491bf8573 100644 --- a/torch/utils/hipify/cuda_to_hip_mappings.py +++ b/torch/utils/hipify/cuda_to_hip_mappings.py @@ -7961,6 +7961,15 @@ "cuda::getDefaultCUDAStream", ("hip::getDefaultHIPStreamMasqueradingAsCUDA", API_PYTORCH), ), + ( + "cuda::getStreamFromExternal", + ("hip::getStreamFromExternalMasqueradingAsCUDA", API_PYTORCH), + ), + ("getStreamFromExternal", ("getStreamFromExternalMasqueradingAsCUDA", API_PYTORCH)), + ( + "cuda::getDefaultCUDAStream", + ("hip::getDefaultHIPStreamMasqueradingAsCUDA", API_PYTORCH), + ), ( "getDefaultCUDAStream", ("getDefaultHIPStreamMasqueradingAsCUDA", API_PYTORCH), From c645d39a7708f99616e7fa3c79eb573db5b0e1e3 Mon Sep 17 00:00:00 2001 From: Joel Schlosser Date: Mon, 14 Jun 2021 13:49:38 -0700 Subject: [PATCH 090/305] Implementation of torch.isin() (#53125) Summary: Fixes https://github.com/pytorch/pytorch/issues/3025 ## Background This PR implements a function similar to numpy's [`isin()`](https://numpy.org/doc/stable/reference/generated/numpy.isin.html#numpy.isin). The op supports integral and floating point types on CPU and CUDA (+ half & bfloat16 for CUDA). Inputs can be one of: * (Tensor, Tensor) * (Tensor, Scalar) * (Scalar, Tensor) Internally, one of two algorithms is selected based on the number of elements vs. test elements. The heuristic for deciding which algorithm to use is taken from [numpy's implementation](https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/arraysetops.py#L575): if `len(test_elements) < 10 * len(elements) ** 0.145`, then a naive brute-force checking algorithm is used. Otherwise, a stablesort-based algorithm is used. I've done some preliminary benchmarking to verify this heuristic on a devgpu, and determined for a limited set of tests that a power value of `0.407` instead of `0.145` is a better inflection point. For now, the heuristic has been left to match numpy's, but input is welcome for the best way to select it or whether it should be left the same as numpy's. Tests are adapted from numpy's [isin and in1d tests](https://github.com/numpy/numpy/blob/7dcd29aaafe1ab8be4be04d3c793e5bcaf17459f/numpy/lib/tests/test_arraysetops.py). Note: my locally generated docs look terrible for some reason, so I'm not including the screenshot for them until I figure out why. Pull Request resolved: https://github.com/pytorch/pytorch/pull/53125 Test Plan: ``` python test/test_ops.py # Ex: python test/test_ops.py TestOpInfoCPU.test_supported_dtypes_isin_cpu_int32 python test/test_sort_and_select.py # Ex: python test/test_sort_and_select.py TestSortAndSelectCPU.test_isin_cpu_int32 ``` Reviewed By: soulitzer Differential Revision: D29101165 Pulled By: jbschlosser fbshipit-source-id: 2dcc38d497b1e843f73f332d837081e819454b4e --- aten/src/ATen/native/TensorCompare.cpp | 133 +++++++++++++++++- aten/src/ATen/native/TensorCompare.h | 2 + .../ATen/native/cpu/TensorCompareKernel.cpp | 32 +++++ aten/src/ATen/native/cuda/TensorCompare.cu | 10 ++ aten/src/ATen/native/native_functions.yaml | 30 ++++ docs/source/torch.rst | 1 + test/test_sort_and_select.py | 122 ++++++++++++++++ torch/_torch_docs.py | 29 ++++ torch/overrides.py | 1 + .../_internal/common_methods_invocations.py | 12 ++ 10 files changed, 370 insertions(+), 2 deletions(-) diff --git a/aten/src/ATen/native/TensorCompare.cpp b/aten/src/ATen/native/TensorCompare.cpp index f12e16ae88773..33bcaebe8101f 100644 --- a/aten/src/ATen/native/TensorCompare.cpp +++ b/aten/src/ATen/native/TensorCompare.cpp @@ -5,10 +5,50 @@ #include #include #include +#include #include #include +#include -namespace at { namespace native { +namespace at { +namespace meta { + +static inline void check_for_unsupported_isin_dtype(const ScalarType type) { + // Bail out for dtypes unsupported by the sorting algorithm to keep the interface consistent. + TORCH_CHECK(type != ScalarType::Bool && + type != ScalarType::BFloat16 && + type != ScalarType::ComplexFloat && + type != ScalarType::ComplexDouble, + "Unsupported input type encountered for isin(): ", type); +} + +TORCH_META_FUNC2(isin, Tensor_Tensor) ( + const Tensor& elements, const Tensor& test_elements, bool assume_unique, bool invert +) { + check_for_unsupported_isin_dtype(elements.scalar_type()); + check_for_unsupported_isin_dtype(test_elements.scalar_type()); + set_output(elements.sizes(), TensorOptions(elements.device()).dtype(ScalarType::Bool)); +} + +TORCH_META_FUNC2(isin, Tensor_Scalar) ( + const Tensor& elements, const c10::Scalar& test_elements, bool assume_unique, bool invert +) { + check_for_unsupported_isin_dtype(elements.scalar_type()); + check_for_unsupported_isin_dtype(test_elements.type()); + set_output(elements.sizes(), TensorOptions(elements.device()).dtype(ScalarType::Bool)); +} + +TORCH_META_FUNC2(isin, Scalar_Tensor) ( + const c10::Scalar& elements, const Tensor& test_elements, bool assume_unique, bool invert +) { + check_for_unsupported_isin_dtype(elements.type()); + check_for_unsupported_isin_dtype(test_elements.scalar_type()); + set_output({0}, TensorOptions(test_elements.device()).dtype(ScalarType::Bool)); +} + +} // namespace meta + +namespace native { DEFINE_DISPATCH(where_kernel); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(max_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) @@ -23,6 +63,7 @@ DEFINE_DISPATCH(clamp_max_stub); // NOLINT(cppcoreguidelines-avoid-non-const-glo DEFINE_DISPATCH(clamp_scalar_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(clamp_min_scalar_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(clamp_max_scalar_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) +DEFINE_DISPATCH(isin_default_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) bool allclose(const Tensor& self, const Tensor& other, double rtol, double atol, bool equal_nan) { return at::isclose(self, other, rtol, atol, equal_nan).all().item(); @@ -245,6 +286,56 @@ Tensor wrapped_scalar_tensor_default_dtype( } // anonymous namespace +// Sorting-based algorithm for isin(); used when the number of test elements is large. +static void isin_sorting( + const Tensor& elements, + const Tensor& test_elements, + bool assume_unique, + bool invert, + const Tensor& out) { + // 1. Concatenate unique elements with unique test elements in 1D form. If + // assume_unique is true, skip calls to unique(). + Tensor elements_flat, test_elements_flat, unique_order; + if (assume_unique) { + elements_flat = elements.ravel(); + test_elements_flat = test_elements.ravel(); + } else { + std::tie (elements_flat, unique_order) = at::_unique( + elements, /*sorted=*/ false, /*return_inverse=*/ true); + std::tie (test_elements_flat, std::ignore) = at::_unique(test_elements, /*sorted=*/ false); + } + + // 2. Stable sort all elements, maintaining order indices to reverse the + // operation. Stable sort is necessary to keep elements before test + // elements within the sorted list. + Tensor all_elements = at::_cat({elements_flat, test_elements_flat}); + Tensor sorted_elements, sorted_order; + std::tie (sorted_elements, sorted_order) = all_elements.sort( + /*stable=*/ true, /*dim=*/ 0, /*descending=*/ false); + + // 3. Create a mask for locations of adjacent duplicate values within the + // sorted list. Duplicate values are in both elements and test elements. + Tensor duplicate_mask = at::empty_like(sorted_elements, TensorOptions(ScalarType::Bool)); + Tensor sorted_except_first = sorted_elements.slice(0, 1, at::indexing::None); + Tensor sorted_except_last = sorted_elements.slice(0, 0, -1); + duplicate_mask.slice(0, 0, -1).copy_( + invert ? sorted_except_first.ne(sorted_except_last) : sorted_except_first.eq(sorted_except_last)); + duplicate_mask.index_put_({-1}, invert); + + // 4. Reorder the mask to match the pre-sorted element order. + Tensor mask = at::empty_like(duplicate_mask); + mask.index_copy_(0, sorted_order, duplicate_mask); + + // 5. Index the mask to match the pre-unique element order. If + // assume_unique is true, just take the first N items of the mask, + // where N is the original number of elements. + if (assume_unique) { + out.copy_(mask.slice(0, 0, elements.numel()).view_as(out)); + } else { + out.copy_(at::index(mask, {c10::optional(unique_order)})); + } +} + Tensor where(const Tensor& condition, const Tensor& self, const Tensor& other) { TORCH_CHECK(condition.device() == self.device() && self.device() == other.device(), "Expected condition, x and y to be on the same device, but condition is on ", @@ -659,4 +750,42 @@ std::tuple mode_out(const Tensor& self, Dimname dim, bool kee return at::mode_out(values, indices, self, dimname_to_position(self, dim), keepdim); } -}} // namespace at::native +TORCH_IMPL_FUNC(isin_Tensor_Tensor_out) ( + const Tensor& elements, const Tensor& test_elements, bool assume_unique, bool invert, const Tensor& out +) { + if (elements.numel() == 0) { + return; + } + + // Heuristic taken from numpy's implementation. + // See https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/arraysetops.py#L575 + if (test_elements.numel() < static_cast( + 10.0f * std::pow(static_cast(elements.numel()), 0.145))) { + out.fill_(invert); + isin_default_stub(elements.device().type(), elements, test_elements, invert, out); + } else { + isin_sorting(elements, test_elements, assume_unique, invert, out); + } +} + +TORCH_IMPL_FUNC(isin_Tensor_Scalar_out) ( + const Tensor& elements, const c10::Scalar& test_elements, bool assume_unique, bool invert, const Tensor& out +) { + // redispatch to eq / ne + if (invert) { + at::ne_out(const_cast(out), elements, test_elements); + } else { + at::eq_out(const_cast(out), elements, test_elements); + } +} + +TORCH_IMPL_FUNC(isin_Scalar_Tensor_out) ( + const c10::Scalar& elements, const Tensor& test_elements, bool assume_unique, bool invert, const Tensor& out +) { + // redispatch + at::isin_out(const_cast(out), wrapped_scalar_tensor(elements, test_elements.device()), + test_elements, assume_unique, invert); +} + +} // namespace native +} // namespace at diff --git a/aten/src/ATen/native/TensorCompare.h b/aten/src/ATen/native/TensorCompare.h index 9ffbfe7105cb5..6bcd38648f661 100644 --- a/aten/src/ATen/native/TensorCompare.h +++ b/aten/src/ATen/native/TensorCompare.h @@ -33,4 +33,6 @@ DECLARE_DISPATCH(void (*)(TensorIterator &, Scalar, Scalar), clamp_scalar_stub); DECLARE_DISPATCH(void (*)(TensorIterator &, Scalar), clamp_min_scalar_stub); DECLARE_DISPATCH(void (*)(TensorIterator &, Scalar), clamp_max_scalar_stub); +using isin_default_fn = void (*)(const Tensor&, const Tensor&, bool, const Tensor&); +DECLARE_DISPATCH(isin_default_fn, isin_default_stub); }} // namespace at::native diff --git a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp index 48eacf6533184..d774ccd3c1892 100644 --- a/aten/src/ATen/native/cpu/TensorCompareKernel.cpp +++ b/aten/src/ATen/native/cpu/TensorCompareKernel.cpp @@ -301,6 +301,37 @@ static void mode_kernel_impl( }); } +// Default brute force implementation of isin(). Used when the number of test elements is small. +// Iterates through each element and checks it against each test element. +static void isin_default_kernel_cpu( + const Tensor& elements, + const Tensor& test_elements, + bool invert, + const Tensor& out) { + // Since test elements is not an input of the TensorIterator, type promotion + // must be done manually. + ScalarType common_type = at::result_type(elements, test_elements); + Tensor test_elements_flat = test_elements.to(common_type).ravel(); + Tensor promoted_elements = elements.to(common_type); + auto iter = TensorIteratorConfig() + .add_output(out) + .add_input(promoted_elements) + .check_all_same_dtype(false) + .build(); + // Dispatch based on promoted type. + AT_DISPATCH_ALL_TYPES(iter.dtype(1), "isin_default_cpu", [&]() { + cpu_kernel(iter, [&](scalar_t element_val) -> bool { + const auto* test_element_data = reinterpret_cast(test_elements_flat.data_ptr()); + for (auto j = 0; j < test_elements_flat.numel(); ++j) { + if (element_val == test_element_data[j]) { + return !invert; + } + } + return invert; + }); + }); +} + static void clamp_kernel_impl(TensorIterator& iter) { AT_DISPATCH_ALL_TYPES_AND(kBFloat16, iter.common_dtype(), "clamp_cpu", [&]() { cpu_kernel_vec(iter, @@ -403,5 +434,6 @@ REGISTER_DISPATCH(clamp_max_stub, &clamp_max_kernel_impl); REGISTER_DISPATCH(clamp_scalar_stub, &clamp_scalar_kernel_impl); REGISTER_DISPATCH(clamp_min_scalar_stub, &clamp_min_scalar_kernel_impl); REGISTER_DISPATCH(clamp_max_scalar_stub, &clamp_max_scalar_kernel_impl); +REGISTER_DISPATCH(isin_default_stub, &isin_default_kernel_cpu); }} // namespace at::native diff --git a/aten/src/ATen/native/cuda/TensorCompare.cu b/aten/src/ATen/native/cuda/TensorCompare.cu index 0b4be605fc297..66f6c21a81faa 100644 --- a/aten/src/ATen/native/cuda/TensorCompare.cu +++ b/aten/src/ATen/native/cuda/TensorCompare.cu @@ -129,6 +129,15 @@ void clamp_max_scalar_kernel_impl(TensorIterator& iter, Scalar max) { }); } +// Composite op implementation for simplicity. This materializes the cross product of elements and test elements, +// so it is not very memory efficient, but it is fast on CUDA. +void isin_default_kernel_gpu(const Tensor& elements, const Tensor& test_elements, bool invert, const Tensor& out) { + std::vector bc_shape(elements.dim(), 1); + bc_shape.push_back(-1); + out.copy_(invert ? elements.unsqueeze(-1).ne(test_elements.view(bc_shape)).all(-1) + : elements.unsqueeze(-1).eq(test_elements.view(bc_shape)).any(-1)); +} + } // anonymous namespace @@ -141,6 +150,7 @@ REGISTER_DISPATCH(clamp_max_stub, &clamp_max_kernel_impl); REGISTER_DISPATCH(clamp_scalar_stub, &clamp_scalar_kernel_impl); REGISTER_DISPATCH(clamp_min_scalar_stub, &clamp_min_scalar_kernel_impl); REGISTER_DISPATCH(clamp_max_scalar_stub, &clamp_max_scalar_kernel_impl); +REGISTER_DISPATCH(isin_default_stub, &isin_default_kernel_gpu); template __global__ void _assert_async_cuda_kernel(scalar_t* input) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index f0762dfc535c3..d5711da70dc27 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -2248,6 +2248,36 @@ - func: isclose(Tensor self, Tensor other, float rtol=1e-05, float atol=1e-08, bool equal_nan=False) -> Tensor variants: function, method +- func: isin.Tensor_Tensor_out(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA: isin_Tensor_Tensor_out + +- func: isin.Tensor_Tensor(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Tensor_Tensor_out + +- func: isin.Tensor_Scalar_out(Tensor elements, Scalar test_element, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA: isin_Tensor_Scalar_out + +- func: isin.Tensor_Scalar(Tensor elements, Scalar test_element, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Tensor_Scalar_out + +- func: isin.Scalar_Tensor_out(Scalar element, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA: isin_Scalar_Tensor_out + +- func: isin.Scalar_Tensor(Scalar element, Tensor test_elements, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Scalar_Tensor_out + - func: isnan(Tensor self) -> Tensor variants: function, method device_check: NoCheck diff --git a/docs/source/torch.rst b/docs/source/torch.rst index 94b288920ca6b..620b2ac07e6ea 100644 --- a/docs/source/torch.rst +++ b/docs/source/torch.rst @@ -421,6 +421,7 @@ Comparison Ops greater isclose isfinite + isin isinf isposinf isneginf diff --git a/test/test_sort_and_select.py b/test/test_sort_and_select.py index a5cf5ec3306b5..526aaa6ba6601 100644 --- a/test/test_sort_and_select.py +++ b/test/test_sort_and_select.py @@ -5,6 +5,7 @@ from torch._six import nan from itertools import permutations, product +from torch.testing import all_types, all_types_and from torch.testing._internal.common_utils import \ (TEST_WITH_ROCM, TestCase, run_tests, make_tensor, slowTest) from torch.testing._internal.common_device_type import \ @@ -843,6 +844,127 @@ def test_kthvalue_scalar(self, device, dtype): self.assertEqual(res[0], ref[0].squeeze()) self.assertEqual(res[1], ref[1].squeeze()) + @dtypes(*all_types()) + @dtypesIfCUDA(*all_types_and(torch.half)) + def test_isin(self, device, dtype): + def assert_isin_equal(a, b): + # Compare to the numpy reference implementation. + x = torch.isin(a, b) + a = a.cpu().numpy() if torch.is_tensor(a) else np.array(a) + b = b.cpu().numpy() if torch.is_tensor(b) else np.array(b) + y = np.isin(a, b) + self.assertEqual(x, y) + + # multi-dim tensor, multi-dim tensor + a = torch.arange(24, device=device, dtype=dtype).reshape([2, 3, 4]) + b = torch.tensor([[10, 20, 30], [0, 1, 3], [11, 22, 33]], device=device, dtype=dtype) + assert_isin_equal(a, b) + + # zero-dim tensor + zero_d = torch.tensor(3, device=device, dtype=dtype) + assert_isin_equal(zero_d, b) + assert_isin_equal(a, zero_d) + assert_isin_equal(zero_d, zero_d) + + # empty tensor + empty = torch.tensor([], device=device, dtype=dtype) + assert_isin_equal(empty, b) + assert_isin_equal(a, empty) + assert_isin_equal(empty, empty) + + # scalar + assert_isin_equal(a, 6) + assert_isin_equal(5, b) + + def define_expected(lst, invert=False): + expected = torch.tensor(lst, device=device) + if invert: + expected = expected.logical_not() + return expected + + # Adapted from numpy's in1d tests + for mult in [1, 10]: + for invert in [False, True]: + a = torch.tensor([5, 7, 1, 2], device=device, dtype=dtype) + b = torch.tensor([2, 4, 3, 1, 5] * mult, device=device, dtype=dtype) + ec = define_expected([True, False, True, True], invert=invert) + c = torch.isin(a, b, assume_unique=True, invert=invert) + self.assertEqual(c, ec) + + a[0] = 8 + ec = define_expected([False, False, True, True], invert=invert) + c = torch.isin(a, b, assume_unique=True, invert=invert) + self.assertEqual(c, ec) + + a[0], a[3] = 4, 8 + ec = define_expected([True, False, True, False], invert=invert) + c = torch.isin(a, b, assume_unique=True, invert=invert) + self.assertEqual(c, ec) + + a = torch.tensor([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5], device=device, dtype=dtype) + b = torch.tensor([2, 3, 4] * mult, device=device, dtype=dtype) + ec = define_expected([False, True, False, True, True, True, True, True, True, + False, True, False, False, False], invert=invert) + c = torch.isin(a, b, invert=invert) + self.assertEqual(c, ec) + + b = torch.tensor([2, 3, 4] * mult + [5, 5, 4] * mult, device=device, dtype=dtype) + ec = define_expected([True, True, True, True, True, True, True, True, True, True, + True, False, True, True], invert=invert) + c = torch.isin(a, b, invert=invert) + self.assertEqual(c, ec) + + a = torch.tensor([5, 7, 1, 2], device=device, dtype=dtype) + b = torch.tensor([2, 4, 3, 1, 5] * mult, device=device, dtype=dtype) + ec = define_expected([True, False, True, True], invert=invert) + c = torch.isin(a, b, invert=invert) + self.assertEqual(c, ec) + + a = torch.tensor([5, 7, 1, 1, 2], device=device, dtype=dtype) + b = torch.tensor([2, 4, 3, 3, 1, 5] * mult, device=device, dtype=dtype) + ec = define_expected([True, False, True, True, True], invert=invert) + c = torch.isin(a, b, invert=invert) + self.assertEqual(c, ec) + + a = torch.tensor([5, 5], device=device, dtype=dtype) + b = torch.tensor([2, 2] * mult, device=device, dtype=dtype) + ec = define_expected([False, False], invert=invert) + c = torch.isin(a, b, invert=invert) + self.assertEqual(c, ec) + + # multi-dimensional input case using sort-based algo + for assume_unique in [False, True]: + a = torch.arange(6, device=device, dtype=dtype).reshape([2, 3]) + b = torch.arange(3, 30, device=device, dtype=dtype) + ec = define_expected([[False, False, False], [True, True, True]], invert=invert) + c = torch.isin(a, b, invert=invert, assume_unique=assume_unique) + self.assertEqual(c, ec) + + def test_isin_different_dtypes(self, device): + supported_types = all_types() if device == 'cpu' else all_types_and(torch.half) + for mult in [1, 10]: + for assume_unique in [False, True]: + for dtype1, dtype2 in product(supported_types, supported_types): + a = torch.tensor([1, 2, 3], device=device, dtype=dtype1) + b = torch.tensor([3, 4, 5] * mult, device=device, dtype=dtype2) + ec = torch.tensor([False, False, True], device=device) + c = torch.isin(a, b, assume_unique=assume_unique) + self.assertEqual(c, ec) + + @onlyCUDA + @dtypes(*all_types()) + def test_isin_different_devices(self, device, dtype): + a = torch.arange(6, device=device, dtype=dtype).reshape([2, 3]) + b = torch.arange(3, 30, device='cpu', dtype=dtype) + with self.assertRaises(RuntimeError): + torch.isin(a, b) + + c = torch.arange(6, device='cpu', dtype=dtype).reshape([2, 3]) + d = torch.arange(3, 30, device=device, dtype=dtype) + with self.assertRaises(RuntimeError): + torch.isin(c, d) + + instantiate_device_type_tests(TestSortAndSelect, globals()) if __name__ == '__main__': diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index 81fe8c007b956..772aef7220f81 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -4042,6 +4042,35 @@ def merge_dicts(*dicts): Alias for :func:`torch.linalg.inv` """.format(**common_args)) +add_docstr(torch.isin, r""" +isin(elements, test_elements, *, assume_unique=False, invert=False) -> Tensor + +Tests if each element of :attr:`elements` is in :attr:`test_elements`. Returns +a boolean tensor of the same shape as :attr:`elements` that is True for elements +in :attr:`test_elements` and False otherwise. + +.. note:: + One of :attr:`elements` or :attr:`test_elements` can be a scalar, but not both. + +Args: + elements (Tensor or Scalar): Input elements + test_elements (Tensor or Scalar): Values against which to test for each input element + assume_unique (bool, optional): If True, assumes both :attr:`elements` and + :attr:`test_elements` contain unique elements, which can speed up the + calculation. Default: False + invert (bool, optional): If True, inverts the boolean return tensor, resulting in True + values for elements *not* in :attr:`test_elements`. Default: False + +Returns: + A boolean tensor of the same shape as :attr:`elements` that is True for elements in + :attr:`test_elements` and False otherwise + +Example: + >>> torch.isin(torch.tensor([[1, 2], [3, 4]]), torch.tensor([2, 3])) + tensor([[False, True], + [ True, False]]) +""") + add_docstr(torch.isinf, r""" isinf(input) -> Tensor diff --git a/torch/overrides.py b/torch/overrides.py index aa6876d43ea71..a0353fda65ab8 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -493,6 +493,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.index_select: lambda input, dim, index, out=None: -1, torch.index_fill: lambda input, dim, index, value: -1, torch.isfinite: lambda tensor: -1, + torch.isin: lambda e, te, assume_unique=False, invert=False: -1, torch.isinf: lambda tensor: -1, torch.isreal: lambda tensor: -1, torch.isposinf: lambda input, out=None: -1, diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 08e35af018b0a..eb60f1f041f49 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -3283,6 +3283,13 @@ def sample_inputs_floor_divide(op_info, device, dtype, requires_grad, **kwargs): SampleInput(lhs, args=(3.14,)), ] +def sample_inputs_isin(op_info, device, dtype, requires_grad): + element = make_tensor((L,), device, dtype, low=None, high=None, requires_grad=requires_grad) + indices = torch.randint(0, L, size=[S]) + test_elements = element[indices].clone() + return [ + SampleInput(element, args=(test_elements,)) + ] def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs): make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) @@ -5336,6 +5343,11 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, sample_inputs_func=sample_inputs_linalg_invertible, decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCUDAIfRocm, skipCPUIfNoLapack]), + OpInfo('isin', + dtypesIfCPU=all_types(), + dtypesIfCUDA=all_types_and(torch.half), + supports_autograd=False, + sample_inputs_func=sample_inputs_isin), OpInfo('kthvalue', dtypes=all_types(), dtypesIfCUDA=all_types_and(torch.float16), From 5e993e6c81202ad6c734196f158bb13daab0d5df Mon Sep 17 00:00:00 2001 From: Shiyan Deng Date: Mon, 14 Jun 2021 14:02:24 -0700 Subject: [PATCH 091/305] [fx2trt] Make TRTInterpreter don't need concrete tensor as arg (#59948) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59948 1. We have two Interpreters. One for vanilla op and one for acc op. Some of the logic between them are similar and in this diff we extract out the similar logic to a Base Interpreter. This makes any future general feature change could benefit both Interpreters. 2. Make TRT Interpreter not depending on concrete tensor arg. We will use `InputTensorSpec` to create necessary inputs for acc tracer. 3. Add unittests for acc op converter. Test Plan: ``` buck test mode/opt caffe2/torch/fb/fx2trt:test_linear buck test mode/opt caffe2/torch/fb/fx2trt:test_batchnorm buck test mode/opt caffe2/torch/fb/fx2trt:test_convolution buck test mode/opt caffe2/torch/fb/fx2trt:test_reshape buck test mode/opt caffe2/torch/fb/fx2trt:test_relu buck test mode/opt caffe2/torch/fb/fx2trt:test_add buck test mode/opt caffe2/torch/fb/fx2trt:test_maxpool ``` Reviewed By: jackm321 Differential Revision: D28749682 fbshipit-source-id: 830d845aede7203f6e56eb1c4e6776af197a0fc3 --- torch/fx/experimental/fx2trt/fx2trt.py | 42 ++++++++++++++++---------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/torch/fx/experimental/fx2trt/fx2trt.py b/torch/fx/experimental/fx2trt/fx2trt.py index fcb79f4e3db68..3a8bebec8b218 100644 --- a/torch/fx/experimental/fx2trt/fx2trt.py +++ b/torch/fx/experimental/fx2trt/fx2trt.py @@ -1,10 +1,10 @@ +import copy import warnings from typing import List, NamedTuple, Iterable, Any, Optional import torch import torch.fx import tensorrt as trt -import copy from torch.fx.experimental.normalize import NormalizeArgs @@ -141,6 +141,7 @@ def register_converter(converter): class InputTensorSpec(NamedTuple): shape : torch.Size dtype : torch.dtype + has_batch_dim : bool = True @classmethod def from_tensor(cls, tensor: torch.Tensor): @@ -151,12 +152,8 @@ def from_tensors(cls, tensors: Iterable[torch.Tensor]): return [cls.from_tensor(t) for t in tensors] -class TRTInterpreter(torch.fx.Interpreter): - def __init__(self, module : torch.fx.GraphModule, input_shapes : List[InputTensorSpec], logger_level=trt.Logger.WARNING): - # Preprocess the model - module = copy.deepcopy(module) - module = module.cpu().float() - module = NormalizeArgs(module).transform() +class BaseTRTInterpreter(torch.fx.Interpreter): + def __init__(self, module : torch.fx.GraphModule, input_specs : List[InputTensorSpec], logger_level=trt.Logger.WARNING): super().__init__(module) self.logger = trt.Logger(logger_level) @@ -168,16 +165,13 @@ def __init__(self, module : torch.fx.GraphModule, input_shapes : List[InputTenso self.network = self.builder.create_network() - self.input_shape_itr = iter(input_shapes) - + self.input_specs_iter = iter(input_specs) self._cur_node_name: Optional[str] = None - self._input_names: List[str] = [] self._output_names: List[str] = [] def run( self, - *args, max_batch_size=64, max_workspace_size=1 << 25, fp16_mode=True, @@ -193,7 +187,7 @@ def run( if fp16_mode and not self.builder.platform_has_fast_fp16: warnings.warn("Current platform doesn't support fast native fp16!") - super().run(*args) + super().run() self.builder.max_batch_size = max_batch_size builder_config = self.builder.create_builder_config() @@ -216,9 +210,11 @@ def run_node(self, n): return super().run_node(n) def placeholder(self, target, args, kwargs): - shape, dtype = next(self.input_shape_itr) self._input_names.append(target) - return self.network.add_input(name=target, shape=tuple(shape[1:]), dtype=torch_dtype_to_trt(dtype)) + shape, dtype, has_batch_dim = next(self.input_specs_iter) + if has_batch_dim: + shape = shape[1:] + return self.network.add_input(name=target, shape=tuple(shape), dtype=torch_dtype_to_trt(dtype)) def call_module(self, target, args, kwargs): assert isinstance(target, str) @@ -255,12 +251,26 @@ def output(self, target, args, kwargs): raise RuntimeError('TensorRT requires all outputs to be Tensor!') for i, output in enumerate(outputs): - # TODO: set location and dtype? name = f'output{i}' output.name = name + self.network.mark_output(output) if self.fp16_mode: output.dtype = trt.float16 else: output.dtype = trt.float32 - self.network.mark_output(output) self._output_names.append(name) + + +class TRTInterpreter(BaseTRTInterpreter): + """ + Use this for general case where there're PyTorch vanilla ops in the FX mdoule. + """ + def __init__(self, module : torch.nn.Module, input_specs : List[InputTensorSpec], logger_level=trt.Logger.WARNING): + # Preprocess the model + if not isinstance(module, torch.fx.GraphModule): + module = torch.fx.symbolic_trace(module) + else: + module = copy.deepcopy(module) + module = module.cpu().float() + module = NormalizeArgs(module).transform() + super().__init__(module, input_specs, logger_level) From 5c1d17e6978f9b2f09fd19614bbc915e245199b6 Mon Sep 17 00:00:00 2001 From: Alban Desmaison Date: Mon, 14 Jun 2021 14:07:41 -0700 Subject: [PATCH 092/305] Revert D29100708: [pytorch][PR] Parametrizations depending on several inputs Test Plan: revert-hammer Differential Revision: D29100708 (https://github.com/pytorch/pytorch/commit/061e71b1994bcd8b73971f3c365d952b0bf563a3) Original commit changeset: b9e91f439cf6 fbshipit-source-id: bff6d8a3d7b24f4beb976383912033c250d91a53 --- test/test_nn.py | 411 ++++++---------------------------- torch/nn/utils/parametrize.py | 391 +++++++------------------------- 2 files changed, 148 insertions(+), 654 deletions(-) diff --git a/test/test_nn.py b/test/test_nn.py index 0c6fa695d5f27..3dde5053ca282 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -2155,11 +2155,7 @@ def forward(self, X): # Cayley map # If X is skew-symmetric it returns an orthogonal matrix Id = torch.eye(X.size(0), device=X.device) - # We call contiguous because solve returns a tensor with strides that are Fortran-contiguous - # and autograd raises a performance warning. - # This happens when we remove the parametrization with leave_parametrized=True, - # which does a set_ with a non-contiguous tensor while the gradient is contiguous - return torch.linalg.solve(Id + X, Id - X).contiguous() + return torch.linalg.solve(Id + X, Id - X) # Define a couple vector parametrizations class FirstZero(nn.Module): @@ -2227,16 +2223,10 @@ def forward(self, x): self.assertEqual(model.bias[-1].item(), 0.) self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happpened # Should not throw - - sgd = torch.optim.SGD(model.parameters(), lr=0.01) - - weight_copy = model.weight.clone() - bias_copy = model.bias.clone() - sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - sgd.step() - self.assertNotEqual(model.weight, weight_copy) - self.assertNotEqual(model.bias, bias_copy) + with torch.no_grad(): + for p in model.parameters(): + p.add_(- p.grad, alpha=0.01) # Remove first parametrization. # Check that the model is still parametrized and so is the second parameter @@ -2250,13 +2240,10 @@ def forward(self, x): self.assertEqual(id(model.weight), initial_weight_id) # Keeps the same id self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happened # Should not throw - weight_copy = model.weight.clone() - bias_copy = model.bias.clone() - sgd.zero_grad() (model.weight.T @ model.bias).sum().backward() - sgd.step() - self.assertNotEqual(model.weight, weight_copy) - self.assertNotEqual(model.bias, bias_copy) + with torch.no_grad(): + for p in model.parameters(): + p.add_(- p.grad, alpha=0.01) # Remove the second parametrization. # Check that the module is not parametrized @@ -2269,33 +2256,22 @@ def forward(self, x): self.assertFalse(hasattr(model, "parametrizations")) # Not parametrized the module self.assertEqual(model.__class__, nn.Linear) # Resores the previous class self.assertEqual(len(list(model.parameters())), 2) # Nothing weird has happeed - - # Should not throw things are updated - weight_copy = model.weight.clone() - bias_copy = model.bias.clone() - sgd.zero_grad() + # Should not throw (model.weight.T @ model.bias).sum().backward() - sgd.step() - self.assertNotEqual(model.weight, weight_copy) - self.assertNotEqual(model.bias, bias_copy) + with torch.no_grad(): + for p in model.parameters(): + p.add_(- p.grad, alpha=0.01) # Test leave_parametrized=True for _ in range(2): parametrize.register_parametrization(model, "weight", Skew()) parametrize.register_parametrization(model, "weight", Orthogonal()) parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) - # We didn't change the dtype nor had multiple inputs, so the id should be the same - self.assertEqual(id(model.weight), initial_weight_id) - self.assertEqual(id(model.bias), initial_bias_id) - - # Should not throw. Things are updated - weight_copy = model.weight.clone() - bias_copy = model.bias.clone() - sgd.zero_grad() + # Should not throw (model.weight.T @ model.bias).sum().backward() - sgd.step() - self.assertNotEqual(model.weight, weight_copy) - self.assertNotEqual(model.bias, bias_copy) + with torch.no_grad(): + for p in model.parameters(): + p.add_(- p.grad, alpha=0.01) def test_register_and_remove_buffer_parametrization(self): r"""Test that it is possible to add and remove parametrizations on buffers""" @@ -2416,12 +2392,8 @@ def right_inverse(self, X): N = 5 model = nn.Linear(N, N) - # Register the skew-symmetric constraint. The result is now skew-symmetric - skew = Skew() - # Make the weight skew-symmetric before registering the parametrization - with torch.no_grad(): - model.weight.set_(skew(model.weight)) - parametrize.register_parametrization(model, "weight", skew) + # Register the skew-symmetric onstraint. The result is now skew-symmetric + parametrize.register_parametrization(model, "weight", Skew()) X = torch.rand(N, N) # X is not skew-symmetric, so it throws an error with self.assertRaises(ValueError): @@ -2444,320 +2416,45 @@ def right_inverse(self, X): self.assertEqual(model.weight, X) self.assertEqual(model.parametrizations.weight.original, torch.zeros_like(X)) - def test_errors_unparametrized_tensor_parametrization(self): - # Test errors when registering a parametrization on an unparametrized tensor - module = nn.Linear(3, 4) - weight_init = module.weight.clone() + def test_errors_parametrization(self): + # A parametrization shall not change the size of the parameter + class ChangeSize(nn.Module): + def forward(self, x): + return x[:-1] - class Identity(nn.Module): + # A simple parametrization that does not implement a right_inverse + class Double(nn.Module): def forward(self, x): - return x + return 2 * x - # Register a parametrization on a non-existing parameter throws - with self.assertRaisesRegex(ValueError, "does not have a parameter"): - parametrize.register_parametrization(module, "foo", Identity()) + module = nn.Linear(3, 4) + # This should not throw when registering + parametrize.register_parametrization(module, "weight", ChangeSize()) + # It throws in the forward + with self.assertRaisesRegex(RuntimeError, "may not change the size"): + module(torch.rand(2)) + # Undo + parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) self.assertFalse(parametrize.is_parametrized(module)) - # Removing parametrizations from an unparametrized tensor throws + # Removing a parametrization from an unparametrized tensor throws with self.assertRaisesRegex(ValueError, "does not have a parametrization"): parametrize.remove_parametrizations(module, "bias") + # Nothing odd happens self.assertFalse(parametrize.is_parametrized(module)) - # A correct parametrization with several outputs - class Sum(nn.Module): - def forward(self, x, y): - return x + y - - def right_inverse(self, z): - return z, torch.zeros_like(z) - - parametrize.register_parametrization(module, "weight", Sum()) - # Cannot remove a parametrization with several outputs with `leave_parametrized=False` - with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): - parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) - parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) - - # A parametrization with an incorrect number of outputs - class WrongNumberParams(nn.Module): - def forward(self, x, y, z): - return x + y + z - - def right_inverse(self, w): - return w, torch.zeros_like(w) - - # Makes param(*param.right_inverse(X)) fail - with self.assertRaisesRegex(TypeError, "positional argument"): - parametrize.register_parametrization(module, "weight", WrongNumberParams()) - self.assertFalse(parametrize.is_parametrized(module)) - - # A parametrization with a right_inverse that does not return a Tensor or Sequence[Tensor] - class WrongRightInverse(Identity): - def right_inverse(self, z): - return None - - # right_inverse should return a Tensor or a Sequence[Tensor] - with self.assertRaisesRegex(ValueError, "Tensor or a Sequence of"): - parametrize.register_parametrization(module, "weight", WrongRightInverse()) - self.assertFalse(parametrize.is_parametrized(module)) - - # If it's a sequence, it must to be a sequence of tensors - class WrongRightInverseSequence(nn.Module): - def forward(self, x, y): - return x - - def right_inverse(self, z): - return None, z - - with self.assertRaisesRegex(ValueError, "of the sequence with type"): - parametrize.register_parametrization(module, "weight", WrongRightInverseSequence()) - self.assertFalse(parametrize.is_parametrized(module)) - - # A parametrization from one tensor to one tensor that changes the dtype - class ChangeDtypeInverse(nn.Module): - def forward(self, x): - return x.float() - - def right_inverse(self, w): - return w.bool() - - # For parametrizations that return one tensor, right_inverse may not change the dtype - with self.assertRaisesRegex(ValueError, "outputs one tensor, it may not change the dtype"): - parametrize.register_parametrization(module, "weight", ChangeDtypeInverse()) - self.assertFalse(parametrize.is_parametrized(module)) - - # Doesn't return a tensor - class NotTensor(nn.Module): - def forward(self, x): - return 2 - - # Forward must return a tensor - with self.assertRaisesRegex(ValueError, "must return a tensor"): - parametrize.register_parametrization(module, "weight", NotTensor()) - self.assertFalse(parametrize.is_parametrized(module)) - - # A parametrization from one tensor to one tensor that changes the dtype - class ChangeDtype(nn.Module): - def forward(self, x): - return x.bool() - - # forward should not change the initial dtype - with self.assertRaisesRegex(ValueError, "may not change the dtype"): - parametrize.register_parametrization(module, "weight", ChangeDtype()) - self.assertFalse(parametrize.is_parametrized(module)) - - # Change shape - class ChangeShape(nn.Module): - def forward(self, x): - return x[:-1] - - # forward should not change the original shape - with self.assertRaisesRegex(ValueError, "may not change the shape"): - parametrize.register_parametrization(module, "weight", ChangeShape()) - self.assertFalse(parametrize.is_parametrized(module)) - - # Many to one that changes dtype - class ChangeDtypeMulti(nn.Module): - def forward(self, x, y): - return (x + y).bool() - - def right_inverse(self, w): - return w, w + 1 - - # forward should not change the original shape even for parametrizations with many inputs - with self.assertRaisesRegex(ValueError, "may not change the dtype"): - parametrize.register_parametrization(module, "weight", ChangeDtypeMulti()) + # Register a parametrization on a non-existing parameter breaks + with self.assertRaisesRegex(ValueError, "does not have a parameter"): + parametrize.register_parametrization(module, "foo", ChangeSize()) self.assertFalse(parametrize.is_parametrized(module)) - # Returning a sequence of size one, although weird, it's correct - class SequenceLen1(nn.Module): - def forward(self, x): - return x - - def right_inverse(self, w): - return (w,) - - parametrize.register_parametrization(module, "weight", SequenceLen1()) - self.assertTrue(hasattr(module.parametrizations.weight, "original0")) - self.assertFalse(hasattr(module.parametrizations.weight, "original1")) - _ = module.weight # Does not throw - self.assertTrue(parametrize.is_parametrized(module)) - parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) - - # None of the operations above should have altered the weight + # Try to assign to a parametrization that does not implement `right_inverse` + parametrize.register_parametrization(module, "weight", Double()) + with self.assertRaisesRegex(RuntimeError, "right_inverse"): + module.weight = torch.rand(4, 3) + # Undo + parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) self.assertFalse(parametrize.is_parametrized(module)) - self.assertEqual(module.weight, weight_init) - - def test_errors_parametrized_tensor_parametrization(self): - # Test errors when registering a parametrization on a parametrized tensor - - class Identity(nn.Module): - def forward(self, x): - return x - - module = nn.Linear(3, 4) - parametrize.register_parametrization(module, "weight", Identity()) - - # Has to return a tensor - class WrongReturn(nn.Module): - def forward(self, x): - return x, x - - with self.assertRaisesRegex(ValueError, "must return a tensor"): - parametrize.register_parametrization(module, "weight", WrongReturn()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - # Cannot change dtype - class ChangeDtype(nn.Module): - def forward(self, x): - return x.bool() - - with self.assertRaisesRegex(ValueError, "may not change the dtype"): - parametrize.register_parametrization(module, "weight", ChangeDtype()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - # Cannot change shape - class ChangeShape(nn.Module): - def forward(self, x): - return x[:-1] - - with self.assertRaisesRegex(ValueError, "may not change the shape"): - parametrize.register_parametrization(module, "weight", ChangeShape()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - # The following checks are mostly due to bugs in the code of the parametrization - - # right_inverse has to return a tensor - class WrongReturnInverse(Identity): - def right_inverse(self, x): - return x, x - - with self.assertRaisesRegex(ValueError, "right_inverse must return a tensor"): - parametrize.register_parametrization(module, "weight", WrongReturnInverse()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - # Cannot change dtype - class ChangeDtypeInverse(Identity): - def right_inverse(self, x): - return x.bool() - - with self.assertRaisesRegex(ValueError, "must have the same dtype"): - parametrize.register_parametrization(module, "weight", ChangeDtypeInverse()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - # Cannot change shape - class ChangeShapeInverse(Identity): - def right_inverse(self, x): - return x[:-1] - - with self.assertRaisesRegex(ValueError, "must have the same shape"): - parametrize.register_parametrization(module, "weight", ChangeShapeInverse()) - self.assertTrue(parametrize.is_parametrized(module)) - self.assertEqual(len(module.parametrizations.weight), 1) - self.assertTrue(isinstance(module.parametrizations.weight[0], Identity)) - - def test_multiple_inputs_parametrization(self): - # A parametrization with several outputs - class RankOne(nn.Module): - def forward(self, x, y): - # Form a rank-1 matrix from a pair of vectors - return x.unsqueeze(-1) @ y.unsqueeze(-2) - - def right_inverse(self, Y): - # We project the given matrix onto the rank 1 matrices - U, S, Vh = torch.linalg.svd(Y, full_matrices=False) - # S is ordered in a decreasing way. - s0_sqrt = S[0].sqrt().unsqueeze(-1) - return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt - - # Simple parametrisation - class Double(nn.Module): - def forward(self, x): - return 2.0 * x - - def right_inverse(self, w): - return 0.5 * w - - model = nn.Linear(3, 3) - # Test one parametrization - parametrize.register_parametrization(model, "weight", RankOne()) - self.assertTrue(hasattr(model, "parametrizations")) - self.assertTrue(parametrize.is_parametrized(model)) - self.assertTrue(parametrize.is_parametrized(model, "weight")) - self.assertTrue(hasattr(model.parametrizations.weight, "original0")) - self.assertIn("original0", model.parametrizations.weight._parameters) - self.assertTrue(hasattr(model.parametrizations.weight, "original1")) - self.assertIn("original1", model.parametrizations.weight._parameters) - self.assertFalse(parametrize.is_parametrized(model, "bias")) - self.assertNotIn("weight", model._parameters) - # Result should be rank 1 - self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) - - with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): - # Cannot remove a parametrization with multiple inputs and not leave it parametrized - parametrize.remove_parametrizations(model, "weight", leave_parametrized=False) - # Remove parametrization and check consistency - parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) - self.assertFalse(hasattr(model, "parametrizations")) - self.assertEqual(model.__class__, nn.Linear) - self.assertFalse(parametrize.is_parametrized(model)) - self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) - self.assertIn("weight", model._parameters) - - # Registering parametrizations with one input on top of one with multiple inputs should work - init_weight = model.weight.clone() - parametrize.register_parametrization(model, "weight", RankOne()) - # Projecting a rank 1 matrix onto the matrices of rank one does not change the matrix - self.assertTrue(torch.allclose(init_weight, model.weight)) - parametrize.register_parametrization(model, "weight", Double()) - # The matrix now is twice the initial matrix - self.assertTrue(torch.allclose(2.0 * init_weight, model.weight)) - # Multiplying by a scalar does not change the rank - self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) - - # The model has now three parameters - self.assertEqual(len(list(model.parameters())), 3) - - sgd = torch.optim.SGD(model.parameters(), lr=0.1) - - # Test backward. Should not throw - for _ in range(2): - sgd.zero_grad() - loss = (model.weight.T @ model.bias).sum() - loss.backward() - sgd.step() - - # Same drill as before, removing should work as expected - with self.assertRaisesRegex(ValueError, "leave_parametrized=False"): - # Cannot remove a parametrization with multiple inputs and not leave it parametrized - parametrize.remove_parametrizations(model, "weight", leave_parametrized=False) - # Remove parametrization and check consistency - parametrize.remove_parametrizations(model, "weight", leave_parametrized=True) - self.assertFalse(hasattr(model, "parametrizations")) - self.assertEqual(model.__class__, nn.Linear) - self.assertFalse(parametrize.is_parametrized(model)) - self.assertEqual(torch.linalg.matrix_rank(model.weight).item(), 1) - self.assertIn("weight", model._parameters) - - # The model has now two parameters - self.assertEqual(len(list(model.parameters())), 2) - - # Test backward. Should not throw - sgd = torch.optim.SGD(model.parameters(), lr=0.1) - for _ in range(2): - sgd.zero_grad() - loss = (model.weight.T @ model.bias).sum() - loss.backward() - sgd.step() def test_caching_parametrization(self): r"""Test the caching system of a parametrization""" @@ -2782,6 +2479,24 @@ def forward(self, X): Y = model.weight self.assertEqual(id(X), id(Y)) + def test_dtype_parametrization(self): + r"""Test a case that is not allowed when removing a parametrization""" + class ChangeType(nn.Module): + def forward(self, X): + return X.double() + + module = nn.Linear(4, 4).float() + input_ = torch.rand(4).double() + # It is allowed to register a parametrization that changes the dtype + parametrize.register_parametrization(module, "weight", ChangeType()) + module(input_) + # We can remove it leaving the original tensor + parametrize.remove_parametrizations(module, "weight", leave_parametrized=False) + # But leaving it parametrized breaks + parametrize.register_parametrization(module, "weight", ChangeType()) + with self.assertRaisesRegex(ValueError, "changes the dtype"): + parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) + def test_parametrization_same_training_mode(self): r"""Test training mode updated on parametrization registration""" class Identity(nn.Module): @@ -7496,7 +7211,7 @@ def test_cudnn_weight_format(self): weight = all_vars[4] weight_data = weight.data.clone() with torch.no_grad(): - weight.copy_(weight_data) + weight.set_(weight_data) for _ in range(2): with warnings.catch_warnings(record=True) as w: diff --git a/torch/nn/utils/parametrize.py b/torch/nn/utils/parametrize.py index 6765dd938fef9..a68a944824773 100644 --- a/torch/nn/utils/parametrize.py +++ b/torch/nn/utils/parametrize.py @@ -2,10 +2,8 @@ from torch.nn.modules.container import ModuleList, ModuleDict, Module from torch.nn.parameter import Parameter from torch import Tensor - -import collections +from typing import Union, Optional, Iterable, Dict, Tuple from contextlib import contextmanager -from typing import Union, Optional, Dict, Tuple, Sequence _cache_enabled = 0 @@ -55,208 +53,72 @@ def cached(): _cache = {} -def _register_parameter_or_buffer(module, name, X): - if isinstance(X, Parameter): - module.register_parameter(name, X) - else: - module.register_buffer(name, X) - - class ParametrizationList(ModuleList): - r"""A sequential container that holds and manages the ``original`` or ``original0``, ``original1``, ... - parameters or buffers of a parametrized :class:`torch.nn.Module`. + r"""A sequential container that holds and manages the ``original`` parameter or buffer of + a parametrized :class:`torch.nn.Module`. It is the type of + ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` has been parametrized + with :func:`register_parametrization`. - It is the type of ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` - has been parametrized with :func:`register_parametrization`. - - If the first registered parmetrization has a ``right_inverse`` that returns one tensor or - does not have a ``right_inverse`` (in which case we assume that ``right_inverse`` is the identity), - it will hold the tensor under the name ``original``. - If it has a ``right_inverse`` that returns more than one tensor, these will be registered as - ``original0``, ``original1``, ... - - .. warning:: + .. note :: This class is used internally by :func:`register_parametrization`. It is documented - here for completeness. It shall not be instantiated by the user. + here for completeness. It should not be instantiated by the user. Args: - modules (sequence): sequence of modules representing the parametrizations + modules (iterable): an iterable of modules representing the parametrizations original (Parameter or Tensor): parameter or buffer that is parametrized """ original: Tensor def __init__( - self, modules: Sequence[Module], original: Union[Tensor, Parameter] + self, modules: Iterable[Module], original: Union[Tensor, Parameter] ) -> None: - # We require this because we need to treat differently the first parametrization - # This should never throw, unless this class is used from the outside - if len(modules) == 0: - raise ValueError("ParametrizationList requires one or more modules.") - super().__init__(modules) - - # In plain words: - # module.weight must keep its dtype and shape. - # Furthermore, if there is no right_inverse or the right_inverse returns a tensor, - # this should be of the same dtype as the original tensor - # - # We check that the following invariants hold: - # X = module.weight - # Y = param.right_inverse(X) - # assert isinstance(Y, Tensor) or - # (isinstance(Y, collections.abc.Sequence) and all(isinstance(t, Tensor) for t in Y)) - # Z = param(Y) if isisntance(Y, Tensor) else param(*Y) - # # Consistency checks - # assert X.dtype == Z.dtype and X.shape == Z.shape - # # If it has one input, this allows to be able to use set_ to be able to - # # move data to/from the original tensor without changing its id (which is what the - # # optimiser uses to track parameters) - # if isinstance(Y, Tensor) - # assert X.dtype == Y.dtype - # Below we use original = X, new = Y - - original_shape = original.shape - original_dtype = original.dtype - - # Compute new - with torch.no_grad(): - new = original - for module in reversed(self): # type: ignore[call-overload] - if hasattr(module, "right_inverse"): - new = module.right_inverse(new) - # else, we assume that right_inverse is the identity - - if not isinstance(new, Tensor) and not isinstance(new, collections.abc.Sequence): - raise ValueError("'right_inverse' must return a Tensor or a Sequence of tensors (list, tuple...). " - f"Got {type(new).__name__}") - - # Set the number of original tensors - self.is_tensor = isinstance(new, Tensor) - self.ntensors = 1 if self.is_tensor else len(new) - - # Register the tensor(s) - if self.is_tensor: - if original.dtype != new.dtype: - raise ValueError( - "When `right_inverse` outputs one tensor, it may not change the dtype.\n" - f"original.dtype: {original.dtype}\n" - f"right_inverse(original).dtype: {new.dtype}" - ) - # Set the original to original so that the user does not need to re-register the parameter - # manually in the optimiser - with torch.no_grad(): - original.set_(new) # type: ignore[call-overload] - _register_parameter_or_buffer(self, "original", original) + if isinstance(original, Parameter): + self.register_parameter("original", original) else: - for i, originali in enumerate(new): - if not isinstance(originali, Tensor): - raise ValueError("'right_inverse' must return a Tensor or a Sequence of tensors " - "(list, tuple...). " - f"Got element {i} of the sequence with type {type(originali).__name__}.") - - # If the original tensor was a Parameter that required grad, we expect the user to - # add the new parameters to the optimizer after registering the parametrization - # (this is documented) - if isinstance(original, Parameter): - originali = Parameter(originali) - originali.requires_grad_(original.requires_grad) - _register_parameter_or_buffer(self, f"original{i}", originali) - - # Consistency checks: - # Since f : A -> B, right_inverse : B -> A, Z and original should live in B - # Z = forward(right_inverse(original)) - Z = self() - if not isinstance(Z, Tensor): - raise ValueError( - f"A parametrization must return a tensor. Got {type(Z).__name__}." - ) - if Z.dtype != original_dtype: - raise ValueError( - "Registering a parametrization may not change the dtype of the tensor.\n" - f"unparametrized dtype: {original_dtype}\n" - f"parametrized dtype: {Z.dtype}" - ) - if Z.shape != original_shape: - raise ValueError( - "Registering a parametrization may not change the shape of the tensor.\n" - f"unarametrized shape: {original_shape}\n" - f"parametrized shape: {Z.shape}" - ) + self.register_buffer("original", original) + + def set_original_(self, value: Tensor) -> None: + r"""This method is called when assigning to a parametrized tensor. - def right_inverse(self, value: Tensor) -> None: - r"""Calls the methods ``right_inverse`` (see :func:`register_parametrization`) - of the parametrizations in the inverse order they were registered in. - Then, it stores the result in ``self.original`` if ``right_inverse`` outputs one tensor - or in ``self.original0``, ``self.original1``, ... if it outputs several. + It calls the methods ``right_inverse`` (see :func:`register_parametrization`) + of the parametrizations in the inverse order that they have been registered. + Then, it assigns the result to ``self.original``. Args: value (Tensor): Value to which initialize the module - """ - # All the exceptions in this function should almost never throw. - # They could throw if, for example, right_inverse function returns a different - # dtype when given a different input, which should most likely be caused by a - # bug in the user's code + Raises: + RuntimeError: if any of the parametrizations do not implement a ``right_inverse`` method + """ with torch.no_grad(): # See https://github.com/pytorch/pytorch/issues/53103 for module in reversed(self): # type: ignore[call-overload] if hasattr(module, "right_inverse"): value = module.right_inverse(value) - # else we assume that right_inverse is the identity - if self.is_tensor: - # These exceptions should only throw when a right_inverse function does not - # return the same dtype for every input, which should most likely be caused by a bug - if not isinstance(value, Tensor): - raise ValueError( - f"`right_inverse` should return a tensor. Got {type(value).__name__}" - ) - if value.dtype != self.original.dtype: - raise ValueError( - f"The tensor returned by `right_inverse` has dtype {value.dtype} " - f"while `original` has dtype {self.original.dtype}" - ) - # We know that the result is going to have the same dtype - self.original.set_(value) # type: ignore[call-overload] - else: - if not isinstance(value, collections.abc.Sequence): - raise ValueError( - "'right_inverse' must return a sequence of tensors. " - f"Got {type(value).__name__}." + else: + raise RuntimeError( + "The parametrization '{}' does not implement a 'right_inverse' method. " + "Assigning to a parametrized tensor is only possible when all the parametrizations " + "implement a 'right_inverse' method.".format(module.__class__.__name__) ) - if len(value) != self.ntensors: - raise ValueError( - "'right_inverse' must return a sequence of tensors of length " - f"{self.ntensors}. Got a sequence of lenght {len(value)}." - ) - for i, tensor in enumerate(value): - original_i = getattr(self, f"original{i}") - if not isinstance(tensor, Tensor): - raise ValueError( - f"`right_inverse` must return a sequence of tensors. " - f"Got element {i} of type {type(tensor).__name__}" - ) - if original_i.dtype != tensor.dtype: - raise ValueError( - f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} " - f"while `original{i}` has dtype {original_i.dtype}" - ) - original_i.set_(tensor) + self.original.copy_(value) def forward(self) -> Tensor: - # Unpack the originals for the first parametrization - if self.is_tensor: - x = self[0](self.original) - else: - originals = (getattr(self, f"original{i}") for i in range(self.ntensors)) - x = self[0](*originals) - # It's not possible to call self[1:] here, so we have to be a bit more cryptic - for module in list(self._modules.values())[1:]: + x = self.original + for module in self: x = module(x) + if x.size() != self.original.size(): + raise RuntimeError( + "The parametrization may not change the size of the parametrized tensor. " + "Size of original tensor: {} " + "Size of parametrized tensor: {}".format(self.original.size(), x.size()) + ) return x def _inject_new_class(module: Module) -> None: - r"""Sets up a module to be parametrized. + r"""Sets up the parametrization mechanism used by parametrizations. This works by substituting the class of the module by a class that extends it to be able to inject a property @@ -275,7 +137,7 @@ def getstate(self): ) param_cls = type( - f"Parametrized{cls.__name__}", + "Parametrized{}".format(cls.__name__), (cls,), { "__getstate__": getstate, @@ -316,10 +178,11 @@ def get_parametrized(self) -> Tensor: return parametrization() def set_original(self, value: Tensor) -> None: - self.parametrizations[tensor_name].right_inverse(value) + self.parametrizations[tensor_name].set_original_(value) setattr(module.__class__, tensor_name, property(get_parametrized, set_original)) + def register_parametrization( module: Module, tensor_name: str, parametrization: Module ) -> Module: @@ -328,12 +191,12 @@ def register_parametrization( Assume that ``tensor_name="weight"`` for simplicity. When accessing ``module.weight``, the module will return the parametrized version ``parametrization(module.weight)``. If the original tensor requires a gradient, the backward pass will differentiate - through :attr:`parametrization`, and the optimizer will update the tensor accordingly. + through the :attr:`parametrization`, and the optimizer will update the tensor accordingly. The first time that a module registers a parametrization, this function will add an attribute ``parametrizations`` to the module of type :class:`~ParametrizationList`. - The list of parametrizations on the tensor ``weight`` will be accessible under + The list of parametrizations on a tensor will be accessible under ``module.parametrizations.weight``. The original tensor will be accessible under @@ -342,8 +205,8 @@ def register_parametrization( Parametrizations may be concatenated by registering several parametrizations on the same attribute. - The training mode of a registered parametrization is updated on registration - to match the training mode of the host module + The training mode of the registered parametrizations are updated on registration + if necessary to match the training mode of the host module Parametrized parameters and buffers have an inbuilt caching system that can be activated using the context manager :func:`cached`. @@ -352,37 +215,16 @@ def register_parametrization( .. code-block:: python - def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] + def right_inverse(self, X: Tensor) -> Tensor - If this method is not implemented, it defaults to the identity. - This method is called on the unparametrized tensor when the first parametrization - is registered. + If :attr:`parametrization` implements this method, it will be possible to assign + to the parametrized tensor. This may be used to initialize the tensor, as shown in the example. In most situations, ``right_inverse`` will be a function such that ``forward(right_inverse(X)) == X`` (see `right inverse `_). Sometimes, when the parametrization is not surjective, it may be reasonable - to relax this. - This may be used to initialize the tensor, as shown in the example below. - - It is possible for the first parametrization to depend on several inputs. - This may be implemented returning a tuple of tensors from ``right_inverse`` - (see the example implementation of a ``RankOne`` parametrization below). - - In this case, the unconstrained tensors are also located under ``module.parametrizations.weight`` - with names ``original0``, ``original1``,... - - .. note:: - - Whenever a parametrization is registered, both its forward and backward method will be called - once to perform a number of consistency checks. - - .. warning:: - - If a parametrization depends on several inputs, :func:`~register_parametrization` - will register a number of new parameters. If such parametrization is registered - after the optimizer is created, these new parameters will need to be added manually - to the optimizer. See :meth:`torch.Optimizer.add_param_group`. + to relax this, as shown in the example below. Args: module (nn.Module): module on which to register the parametrization @@ -390,22 +232,24 @@ def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] the parametrization parametrization (nn.Module): the parametrization to register + Returns: + Module: module + Raises: ValueError: if the module does not have a parameter or a buffer named :attr:`tensor_name` Examples: >>> import torch - >>> import torch.nn as nn >>> import torch.nn.utils.parametrize as P >>> - >>> class Symmetric(nn.Module): + >>> class Symmetric(torch.nn.Module): >>> def forward(self, X): >>> return X.triu() + X.triu(1).T # Return a symmetric matrix >>> >>> def right_inverse(self, A): >>> return A.triu() >>> - >>> m = nn.Linear(5, 5) + >>> m = torch.nn.Linear(5, 5) >>> P.register_parametrization(m, "weight", Symmetric()) >>> print(torch.allclose(m.weight, m.weight.T)) # m.weight is now symmetric True @@ -414,80 +258,15 @@ def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] >>> m.weight = A # Initialize the weight to be the symmetric matrix A >>> print(torch.allclose(m.weight, A)) True - - >>> class RankOne(nn.Module): - >>> def forward(self, x, y): - >>> # Form a rank 1 matrix multiplying two vectors - >>> return x.unsqueeze(-1) @ y.unsqueeze(-2) - >>> - >>> def right_inverse(self, Z): - >>> # Project Z onto the rank 1 matrices - >>> U, S, Vh = torch.linalg.svd(Z, full_matrices=False) - >>> # Return rescaled singular vectors - >>> s0_sqrt = S[0].sqrt().unsqueeze(-1) - >>> return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt - >>> - >>> linear_rank_one = P.register_parametrization(nn.Linear(4, 4), "weight", RankOne()) - >>> print(torch.linalg.matrix_rank(linear_rank_one.weight).item()) - 1 - """ parametrization.train(module.training) if is_parametrized(module, tensor_name): - # Correctness checks. - # If A is the space of tensors with shape and dtype equal to module.weight - # we check that parametrization.forward and parametrization.right_inverse are - # functions from A to A - - Y = getattr(module, tensor_name) - X = parametrization(Y) - if not isinstance(X, Tensor): - raise ValueError( - f"A parametrization must return a tensor. Got {type(X).__name__}." - ) - if X.dtype != Y.dtype: - raise ValueError( - "Registering a parametrization may not change the dtype of the tensor.\n" - f"module.{tensor_name}.dtype: {Y.dtype}\n" - f"parametrization(module.{tensor_name}).dtype: {X.dtype}" - ) - if X.shape != Y.shape: - raise ValueError( - "Registering a parametrization may not change the shape of the tensor.\n" - f"module.{tensor_name}.shape: {Y.shape}\n" - f"parametrization(module.{tensor_name}).shape: {X.shape}" - ) - if hasattr(parametrization, "right_inverse"): - Z = parametrization.right_inverse(X) # type: ignore[operator] - if not isinstance(Z, Tensor): - raise ValueError( - f"parametrization.right_inverse must return a tensor. Got: {type(Z).__name__}" - ) - if Z.dtype != Y.dtype: - raise ValueError( - "The tensor returned by parametrization.right_inverse must have the same dtype " - f"as module.{tensor_name}.\n" - f"module.{tensor_name}.dtype: {Y.dtype}\n" - f"returned dtype: {Z.dtype}" - ) - if Z.shape != Y.shape: - raise ValueError( - "The tensor returned by parametrization.right_inverse must have the same shape " - f"as module.{tensor_name}.\n" - f"module.{tensor_name}.shape: {Y.shape}\n" - f"returned shape: {Z.shape}" - ) - # else right_inverse is assumed to be the identity - - # add the new parametrization to the parametrization list - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - module.parametrizations[tensor_name].append(parametrization) + # Just add the new parametrization to the parametrization list + module.parametrizations[tensor_name].append(parametrization) # type: ignore[index, union-attr] elif tensor_name in module._buffers or tensor_name in module._parameters: # Set the parametrization mechanism # Fetch the original buffer or parameter original = getattr(module, tensor_name) - # We create this early to check for possible errors - parametrizations = ParametrizationList([parametrization], original) # Delete the previous parameter or buffer delattr(module, tensor_name) # If this is the first parametrization registered on the module, @@ -495,17 +274,18 @@ def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] if not is_parametrized(module): # Change the class _inject_new_class(module) - # Inject a ``ModuleDict`` into the instance under module.parametrizations + # Inject the a ``ModuleDict`` into the instance under module.parametrizations module.parametrizations = ModuleDict() # Add a property into the class _inject_property(module, tensor_name) # Add a ParametrizationList - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - module.parametrizations[tensor_name] = parametrizations + module.parametrizations[tensor_name] = ParametrizationList( # type: ignore[assignment, index, operator] + [parametrization], original + ) else: raise ValueError( - f"Module '{module}' does not have a parameter, a buffer, or a " - f"parametrized element with name '{tensor_name}'" + "Module '{}' does not have a parameter, a buffer, or a " + "parametrized element with name '{}'".format(module, tensor_name) ) return module @@ -541,7 +321,6 @@ def remove_parametrizations( of the tensor. - If ``leave_parametrized=False``, ``module[tensor_name]`` will be set to the unparametrised tensor in ``module.parametrizations[tensor_name].original``. - This is only possible when the parametrization depends on just one tensor. Args: module (nn.Module): module from which remove the parametrization @@ -554,44 +333,44 @@ def remove_parametrizations( Raises: ValueError: if ``module[tensor_name]`` is not parametrized - ValueError: if ``leave_parametrized=False`` and the parametrization depends on several tensors + ValueError: if ``leave_parametrized=True`` and the parametrization changes the size or dtype + of the tensor """ if not is_parametrized(module, tensor_name): - raise ValueError(f"Module {module} does not have a parametrization on {tensor_name}") + raise ValueError( + "Module {} does not have a parametrization on {}".format( + module, tensor_name + ) + ) # Fetch the original tensor - assert isinstance(module.parametrizations, ModuleDict) # Make mypy happy - parametrizations = module.parametrizations[tensor_name] - if parametrizations.is_tensor: - original = parametrizations.original - if leave_parametrized: - with torch.no_grad(): - t = getattr(module, tensor_name) - # We know they have the same dtype because we have checked this when registering the - # parametrizations. As such, we can use set_ - # We do this so that the parameter does not to change the id() - # This way the user does not need to update the optimizer + original = module.parametrizations[tensor_name].original # type: ignore[index, union-attr] + if leave_parametrized: + with torch.no_grad(): + t = getattr(module, tensor_name) + # If they have the same dtype, we reuse the original tensor. + # We do this so that the parameter does not to change the id() + # This way the user does not need to update the optimizer + if t.dtype == original.dtype: with torch.no_grad(): original.set_(t) - else: - if leave_parametrized: - # We cannot use no_grad because we need to know whether one or more - # original tensors required grad - t = getattr(module, tensor_name) - # We'll have to trust the user to add it to the optimizer - original = Parameter(t) if t.requires_grad else t else: - raise ValueError("Cannot leave unparametrized (`leave_parametrized=False`) a tensor " - "that is parametrized in terms of a sequence of tensors.") - + raise ValueError( + "The parametrization changes the dtype of the tensor from {} to {}. " + "It is not supported to leave the tensor parametrized (`leave_parametrized=True`) " + "in this case.".format(original.dtype, t.dtype) + ) # Delete the property that manages the parametrization delattr(module.__class__, tensor_name) # Delete the ParametrizationList - del module.parametrizations[tensor_name] + del module.parametrizations[tensor_name] # type: ignore[operator, union-attr] # Restore the parameter / buffer into the main class - _register_parameter_or_buffer(module, tensor_name, original) + if isinstance(original, Parameter): + module.register_parameter(tensor_name, original) + else: + module.register_buffer(tensor_name, original) # Roll back the parametrized class if no other buffer or parameter # is currently parametrized in this class From c2098487e8005837c1c4c5ed76df7df4cac796f2 Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Mon, 14 Jun 2021 15:04:36 -0700 Subject: [PATCH 093/305] [c10d] Move pg wrapper tests to their own file. (#59840) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59840 moving these tests to their own standalone file. No meaningful code changes. ghstack-source-id: 131359162 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D29012664 fbshipit-source-id: 348870016509a6ed7e69240fa82bccef4a12d674 --- .jenkins/pytorch/multigpu-test.sh | 1 + .../win-test-helpers/test_distributed.bat | 3 + test/distributed/test_c10d_common.py | 119 ------- test/distributed/test_c10d_gloo.py | 88 ----- test/distributed/test_c10d_nccl.py | 85 +---- test/distributed/test_pg_wrapper.py | 324 ++++++++++++++++++ test/run_test.py | 2 + torch/distributed/CONTRIBUTING.md | 3 + 8 files changed, 334 insertions(+), 291 deletions(-) create mode 100644 test/distributed/test_pg_wrapper.py diff --git a/.jenkins/pytorch/multigpu-test.sh b/.jenkins/pytorch/multigpu-test.sh index daa0fb8556e93..7ed43a9f8a852 100755 --- a/.jenkins/pytorch/multigpu-test.sh +++ b/.jenkins/pytorch/multigpu-test.sh @@ -26,6 +26,7 @@ time python test/run_test.py --verbose -i distributed/test_c10d_nccl time python test/run_test.py --verbose -i distributed/test_c10d_spawn_gloo time python test/run_test.py --verbose -i distributed/test_c10d_spawn_nccl time python test/run_test.py --verbose -i distributed/test_store +time python test/run_test.py --verbose -i distributed/test_pg_wrapper time python test/run_test.py --verbose -i distributed/rpc/cuda/test_process_group_agent time python test/run_test.py --verbose -i distributed/rpc/cuda/test_tensorpipe_agent assert_git_not_dirty diff --git a/.jenkins/pytorch/win-test-helpers/test_distributed.bat b/.jenkins/pytorch/win-test-helpers/test_distributed.bat index a50c153822561..53ebee897c35d 100644 --- a/.jenkins/pytorch/win-test-helpers/test_distributed.bat +++ b/.jenkins/pytorch/win-test-helpers/test_distributed.bat @@ -22,3 +22,6 @@ if %errorlevel% neq 0 ( exit /b %errorlevel% ) %1\python.exe test/run_test.py --verbose -i distributed/test_store if %errorlevel% neq 0 ( exit /b %errorlevel% ) + +%1\python.exe test/run_test.py --verbose -i distributed/test_pg_wrapper +if %errorlevel% neq 0 ( exit /b %errorlevel% ) diff --git a/test/distributed/test_c10d_common.py b/test/distributed/test_c10d_common.py index 4d2dc0e1767a9..a25ef9a046bf5 100644 --- a/test/distributed/test_c10d_common.py +++ b/test/distributed/test_c10d_common.py @@ -233,125 +233,6 @@ def forward(self, x): return F.softmax(self.embedding(x), dim=1) -class AbstractProcessGroupWrapperTest(MultiProcessTestCase): - def setUp(self): - super(AbstractProcessGroupWrapperTest, self).setUp() - # For Windows platform, Python does not support fork, change it to spawn here. - if sys.platform == "win32": - self._spawn_processes() - else: - self._fork_processes() - - def _test_collective_hang(self, wrapper_pg, use_cuda=False): - # All ranks besides 1 call allreduce and wrapper_pg should detect a hang - # and report an issue with rank 1. - faulty_rank = 1 - if self.rank != faulty_rank: - tensor = torch.randn(20, 10) - if use_cuda: - tensor = tensor.to(self.rank) - - if self.rank == 0: - # Rank 0 reports faulty ranks - err = f"Ranks {faulty_rank} failed to pass monitoredBarrier" - else: - err = "Please check rank 0 logs for faulty rank" - with self.assertRaisesRegex(RuntimeError, err): - wrapper_pg.allreduce([tensor]) - - def _test_collectives_op_mismatch(self, wrapper_pg, use_cuda=False): - tensor = torch.randn(20, 10) - if use_cuda: - tensor = tensor.to(self.rank) - works = [] - # Run a few successful collectives - for _ in range(10): - work = wrapper_pg.allreduce([tensor]) - works.append(work) - - for w in works: - w.wait() - - # Simulate mismatch: allreduce vs reduce. - with self.assertRaisesRegex( - RuntimeError, "Mismatch between collective operation types" - ): - if self.rank == 0: - wrapper_pg.allreduce([tensor]) - else: - wrapper_pg.reduce([tensor]) - - # Check additional mismatches - - with self.assertRaisesRegex( - RuntimeError, "Mismatch between collective operation types" - ): - if self.rank == 0: - wrapper_pg.reduce([tensor]) - else: - wrapper_pg.barrier() - - with self.assertRaisesRegex( - RuntimeError, "Mismatch between collective operation types" - ): - scatter_result = [torch.ones(4) * i for i in range(self.world_size)] - scattered_tensor = torch.empty(4) - if self.rank == 0: - wrapper_pg.scatter(scattered_tensor, scatter_result, 0) - else: - wrapper_pg.reduce_scatter(scattered_tensor, scatter_result) - - with self.assertRaisesRegex( - RuntimeError, "Mismatch between collective operation types" - ): - if self.rank == 0: - wrapper_pg.broadcast(tensor, 0) - else: - output_tensors = [ - torch.zeros_like(tensor) for _ in range(self.world_size) - ] - wrapper_pg.allgather([output_tensors], [tensor]) - - def _test_collective_shape_mismatch(self, wrapper_pg, use_cuda=False): - wrapper_pg.barrier() - dim = 2 if self.rank == 0 else 10 - tensor = torch.randn(20, dim) - if use_cuda: - tensor = tensor.to(self.rank) - with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): - wrapper_pg.allreduce([tensor]) - # Check errors are raised when dimensionality of shapes is different - tensor = torch.randn(20, 10, 2) if self.rank == 0 else torch.randn(20, 10) - if use_cuda: - tensor = tensor.to(self.rank) - with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): - wrapper_pg.allreduce([tensor]) - - # Check shape errors with scatter - input = [ - torch.tensor( - [self.rank] if self.rank == 0 else [self.rank, self.rank], - device=self.rank if use_cuda else "cpu", - ) - for _ in range(self.world_size) - ] - outputs = [ - torch.tensor( - [-1] if self.rank == 0 else [-1, -1], - device=self.rank if use_cuda else "cpu", - ) - for _ in range(self.world_size) - ] - root_rank = 0 - opts = c10d.ScatterOptions() - opts.rootRank = root_rank - with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): - if self.rank == root_rank: - wrapper_pg.scatter([outputs[self.rank]], [input], opts).wait() - else: - wrapper_pg.scatter([outputs[self.rank]], [], opts).wait() - - class AbstractDistributedDataParallelTest(object): def tearDown(self): # DistributedDataParallel test doesn't seem to call FileStore destructor diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py index 8ce6443d48bb1..a4c2b855f2fc7 100644 --- a/test/distributed/test_c10d_gloo.py +++ b/test/distributed/test_c10d_gloo.py @@ -29,7 +29,6 @@ simple_sparse_reduce_tests, skip_if_win32, create_device, - with_dist_debug_levels, verify_ddp_error_logged, ) from torch.testing._internal.common_utils import ( @@ -45,7 +44,6 @@ Task, ModuleForDdpCommHook, SparseGradientModule, - AbstractProcessGroupWrapperTest, ) @@ -203,92 +201,6 @@ class TimeoutTest(test_c10d_common.AbstractTimeoutTest, TestCase): def test_default_store_timeout_gloo(self): self._test_default_store_timeout("gloo") -@requires_gloo() -@unittest.skipIf( - TEST_WITH_TSAN, - "TSAN is not fork-safe since we're forking in a multi-threaded environment", -) -class ProcessGroupGlooWrapperTest(AbstractProcessGroupWrapperTest): - def setUp(self): - super(ProcessGroupGlooWrapperTest, self).setUp() - - def opts(self, threads=2, timeout=10.0): - opts = c10d.ProcessGroupGloo._Options() - opts._timeout = timeout - opts._devices = [create_device(interface=LOOPBACK)] - opts._threads = threads - return opts - - def _create_wrapper_pg(self, with_new_group=False, timeout=10.0): - store = c10d.FileStore(self.file_name, self.world_size) - c10d.init_process_group( - backend="gloo", rank=self.rank, world_size=self.world_size, store=store - ) - if with_new_group: - pg = c10d.new_group(backend="gloo") - else: - _pg = c10d.ProcessGroupGloo(store, self.rank, self.world_size, self.opts(timeout=timeout)) - pg = c10d._create_process_group_wrapper( - _pg, - "unused", - store, - self.rank, - self.world_size, - timeout=timeout, - ) - return pg - - def test_collective_hang(self): - pg = self._create_wrapper_pg(timeout=2.0) - self._test_collective_hang(pg) - - # NOTE: these tests are separated by debug level instead of combined into - # one due to https://github.com/pytorch/pytorch/issues/55967, they can be - # combined after that is resolved. - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collectives_op_mismatch_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collectives_op_mismatch(pg) - - @with_dist_debug_levels(levels=["OFF"]) - def test_collectives_op_mismatch(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collectives_op_mismatch(pg) - - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collective_shape_mismatch_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collective_shape_mismatch(pg) - - @with_dist_debug_levels(levels=["OFF"]) - def test_collective_shape_mismatch(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collective_shape_mismatch(pg) - - @skip_if_lt_x_gpu(4) - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collectives_op_mismatch_cuda_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collectives_op_mismatch(pg, use_cuda=True) - - @skip_if_lt_x_gpu(4) - @with_dist_debug_levels(levels=["OFF"]) - def test_collectives_op_mismatch_cuda(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collectives_op_mismatch(pg, use_cuda=True) - - @skip_if_lt_x_gpu(4) - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collective_shape_mismatch_cuda_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collective_shape_mismatch(pg, use_cuda=True) - - @skip_if_lt_x_gpu(4) - @with_dist_debug_levels(levels=["OFF"]) - def test_collective_shape_mismatch_cuda(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collective_shape_mismatch(pg, use_cuda=True) - @requires_gloo() @unittest.skipIf( TEST_WITH_TSAN, diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py index 244b50649beec..5583cbb8a32b1 100644 --- a/test/distributed/test_c10d_nccl.py +++ b/test/distributed/test_c10d_nccl.py @@ -30,7 +30,6 @@ from torch.utils.checkpoint import checkpoint from torch.testing._internal.common_distributed import ( MultiProcessTestCase, - requires_gloo, requires_nccl, requires_nccl_version, skip_if_lt_x_gpu, @@ -46,7 +45,7 @@ TEST_WITH_TSAN, ) import test_c10d_common -from test_c10d_common import gpus_for_rank, DoubleGpuNet, ConvNet, ModuleForDdpCommHook, AbstractProcessGroupWrapperTest +from test_c10d_common import gpus_for_rank, DoubleGpuNet, ConvNet, ModuleForDdpCommHook class RendezvousEnvTest(TestCase): @@ -159,88 +158,6 @@ def test_default_store_timeout_nccl(self): raise unittest.SkipTest("No GPUs available, skipping test") self._test_default_store_timeout("nccl") -@requires_gloo() -@requires_nccl() -@unittest.skipIf( - TEST_WITH_TSAN, - "TSAN is not fork-safe since we're forking in a multi-threaded environment", -) -class ProcessGroupNCCLWrapperTest(AbstractProcessGroupWrapperTest): - def setUp(self): - self.num_gpus = torch.cuda.device_count() - if self.num_gpus < 2: - raise unittest.SkipTest("NCCL test requires 2+ GPUs") - super(AbstractProcessGroupWrapperTest, self).setUp() - self._spawn_processes() - # NCCL_BLOCKING_WAIT overrides NCCL_ASYNC_ERROR_HANDLING hence tests - # that use NCCL_BLOCKING_WAIT will test it as expected. - os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1" - - @property - def world_size(self) -> int: - return 2 - - def _create_wrapper_pg(self, with_new_group=False, timeout=10.0): - store = c10d.FileStore(self.file_name, self.world_size) - c10d.init_process_group( - backend="nccl", - rank=self.rank, - world_size=self.world_size, - store=store, - timeout=timedelta(seconds=timeout) - ) - if with_new_group: - pg = c10d.new_group(backend="nccl", timeout=timedelta(seconds=timeout)) - else: - _pg = c10d.ProcessGroupNCCL(store, self.rank, self.world_size, timeout=timedelta(seconds=timeout)) - pg = c10d._create_process_group_wrapper( - _pg, - "unused", - store, - self.rank, - self.world_size, - timeout=timeout, - ) - return pg - - @requires_nccl() - @skip_if_lt_x_gpu(2) - def test_collective_hang(self): - pg = self._create_wrapper_pg(timeout=2.0) - self._test_collective_hang(pg) - - # NOTE: these tests are separated by debug level instead of combined into - # one due to https://github.com/pytorch/pytorch/issues/55967, they can be - # combined after that is resolved. - @requires_nccl() - @skip_if_lt_x_gpu(2) - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collectives_op_mismatch_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collectives_op_mismatch(pg, use_cuda=True) - - @requires_nccl() - @skip_if_lt_x_gpu(2) - @with_dist_debug_levels(levels=["OFF"]) - def test_collectives_op_mismatch(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collectives_op_mismatch(pg, use_cuda=True) - - @requires_nccl() - @skip_if_lt_x_gpu(2) - @with_dist_debug_levels(levels=["DETAIL"]) - def test_collective_shape_mismatch_debug_mode(self): - pg = self._create_wrapper_pg(with_new_group=True) - self._test_collective_shape_mismatch(pg, use_cuda=True) - - @requires_nccl() - @skip_if_lt_x_gpu(2) - @with_dist_debug_levels(levels=["OFF"]) - def test_collective_shape_mismatch(self): - pg = self._create_wrapper_pg(with_new_group=False) - self._test_collective_shape_mismatch(pg, use_cuda=True) - - class ProcessGroupNCCLNoGPUTest(TestCase): MAIN_PROCESS_RANK = 0 diff --git a/test/distributed/test_pg_wrapper.py b/test/distributed/test_pg_wrapper.py new file mode 100644 index 0000000000000..aa32a5b9bcef1 --- /dev/null +++ b/test/distributed/test_pg_wrapper.py @@ -0,0 +1,324 @@ +import os +import sys +import unittest +from datetime import timedelta + +import torch +import torch.distributed as c10d + +if not c10d.is_available(): + print("c10d not available, skipping tests", file=sys.stderr) + sys.exit(0) + +from torch.testing._internal.common_distributed import ( + MultiProcessTestCase, + requires_nccl, + requires_gloo, + skip_if_lt_x_gpu, + with_dist_debug_levels, + create_device, +) +from torch.testing._internal.common_utils import ( + run_tests, + TEST_WITH_TSAN, +) +from test_c10d_common import LOOPBACK + + +class AbstractProcessGroupWrapperTest(MultiProcessTestCase): + def setUp(self): + super(AbstractProcessGroupWrapperTest, self).setUp() + # For Windows platform, Python does not support fork, change it to spawn here. + if sys.platform == "win32": + self._spawn_processes() + else: + self._fork_processes() + + def _test_collective_hang(self, wrapper_pg, use_cuda=False): + # All ranks besides 1 call allreduce and wrapper_pg should detect a hang + # and report an issue with rank 1. + faulty_rank = 1 + if self.rank != faulty_rank: + tensor = torch.randn(20, 10) + if use_cuda: + tensor = tensor.to(self.rank) + + if self.rank == 0: + # Rank 0 reports faulty ranks + err = f"Ranks {faulty_rank} failed to pass monitoredBarrier" + else: + err = "Please check rank 0 logs for faulty rank" + with self.assertRaisesRegex(RuntimeError, err): + wrapper_pg.allreduce([tensor]) + + def _test_collectives_op_mismatch(self, wrapper_pg, use_cuda=False): + tensor = torch.randn(20, 10) + if use_cuda: + tensor = tensor.to(self.rank) + works = [] + # Run a few successful collectives + for _ in range(10): + work = wrapper_pg.allreduce([tensor]) + works.append(work) + + for w in works: + w.wait() + + # Simulate mismatch: allreduce vs reduce. + with self.assertRaisesRegex( + RuntimeError, "Mismatch between collective operation types" + ): + if self.rank == 0: + wrapper_pg.allreduce([tensor]) + else: + wrapper_pg.reduce([tensor]) + + # Check additional mismatches + + with self.assertRaisesRegex( + RuntimeError, "Mismatch between collective operation types" + ): + if self.rank == 0: + wrapper_pg.reduce([tensor]) + else: + wrapper_pg.barrier() + + with self.assertRaisesRegex( + RuntimeError, "Mismatch between collective operation types" + ): + scatter_result = [torch.ones(4) * i for i in range(self.world_size)] + scattered_tensor = torch.empty(4) + if self.rank == 0: + wrapper_pg.scatter(scattered_tensor, scatter_result, 0) + else: + wrapper_pg.reduce_scatter(scattered_tensor, scatter_result) + + with self.assertRaisesRegex( + RuntimeError, "Mismatch between collective operation types" + ): + if self.rank == 0: + wrapper_pg.broadcast(tensor, 0) + else: + output_tensors = [ + torch.zeros_like(tensor) for _ in range(self.world_size) + ] + wrapper_pg.allgather([output_tensors], [tensor]) + + def _test_collective_shape_mismatch(self, wrapper_pg, use_cuda=False): + wrapper_pg.barrier() + dim = 2 if self.rank == 0 else 10 + tensor = torch.randn(20, dim) + if use_cuda: + tensor = tensor.to(self.rank) + with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): + wrapper_pg.allreduce([tensor]) + # Check errors are raised when dimensionality of shapes is different + tensor = torch.randn(20, 10, 2) if self.rank == 0 else torch.randn(20, 10) + if use_cuda: + tensor = tensor.to(self.rank) + with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): + wrapper_pg.allreduce([tensor]) + + # Check shape errors with scatter + input = [ + torch.tensor( + [self.rank] if self.rank == 0 else [self.rank, self.rank], + device=self.rank if use_cuda else "cpu", + ) + for _ in range(self.world_size) + ] + outputs = [ + torch.tensor( + [-1] if self.rank == 0 else [-1, -1], + device=self.rank if use_cuda else "cpu", + ) + for _ in range(self.world_size) + ] + root_rank = 0 + opts = c10d.ScatterOptions() + opts.rootRank = root_rank + with self.assertRaisesRegex(RuntimeError, "Error when verifying shape tensors"): + if self.rank == root_rank: + wrapper_pg.scatter([outputs[self.rank]], [input], opts).wait() + else: + wrapper_pg.scatter([outputs[self.rank]], [], opts).wait() + + +@requires_gloo() +@requires_nccl() +@unittest.skipIf( + TEST_WITH_TSAN, + "TSAN is not fork-safe since we're forking in a multi-threaded environment", +) +class ProcessGroupNCCLWrapperTest(AbstractProcessGroupWrapperTest): + def setUp(self): + self.num_gpus = torch.cuda.device_count() + if self.num_gpus < 2: + raise unittest.SkipTest("NCCL test requires 2+ GPUs") + super(AbstractProcessGroupWrapperTest, self).setUp() + self._spawn_processes() + # NCCL_BLOCKING_WAIT overrides NCCL_ASYNC_ERROR_HANDLING hence tests + # that use NCCL_BLOCKING_WAIT will test it as expected. + os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1" + + @property + def world_size(self) -> int: + return 2 + + def _create_wrapper_pg(self, with_new_group=False, timeout=10.0): + store = c10d.FileStore(self.file_name, self.world_size) + c10d.init_process_group( + backend="nccl", + rank=self.rank, + world_size=self.world_size, + store=store, + timeout=timedelta(seconds=timeout), + ) + if with_new_group: + pg = c10d.new_group(backend="nccl", timeout=timedelta(seconds=timeout)) + else: + _pg = c10d.ProcessGroupNCCL( + store, self.rank, self.world_size, timeout=timedelta(seconds=timeout) + ) + pg = c10d._create_process_group_wrapper( + _pg, + "unused", + store, + self.rank, + self.world_size, + timeout=timeout, + ) + return pg + + @requires_nccl() + @skip_if_lt_x_gpu(2) + def test_collective_hang(self): + pg = self._create_wrapper_pg(timeout=2.0) + self._test_collective_hang(pg) + + # NOTE: these tests are separated by debug level instead of combined into + # one due to https://github.com/pytorch/pytorch/issues/55967, they can be + # combined after that is resolved. + @requires_nccl() + @skip_if_lt_x_gpu(2) + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collectives_op_mismatch_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collectives_op_mismatch(pg, use_cuda=True) + + @requires_nccl() + @skip_if_lt_x_gpu(2) + @with_dist_debug_levels(levels=["OFF"]) + def test_collectives_op_mismatch(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collectives_op_mismatch(pg, use_cuda=True) + + @requires_nccl() + @skip_if_lt_x_gpu(2) + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collective_shape_mismatch_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collective_shape_mismatch(pg, use_cuda=True) + + @requires_nccl() + @skip_if_lt_x_gpu(2) + @with_dist_debug_levels(levels=["OFF"]) + def test_collective_shape_mismatch(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collective_shape_mismatch(pg, use_cuda=True) + + +@requires_gloo() +@unittest.skipIf( + TEST_WITH_TSAN, + "TSAN is not fork-safe since we're forking in a multi-threaded environment", +) +class ProcessGroupGlooWrapperTest(AbstractProcessGroupWrapperTest): + def setUp(self): + super(ProcessGroupGlooWrapperTest, self).setUp() + + def opts(self, threads=2, timeout=10.0): + opts = c10d.ProcessGroupGloo._Options() + opts._timeout = timeout + opts._devices = [create_device(interface=LOOPBACK)] + opts._threads = threads + return opts + + def _create_wrapper_pg(self, with_new_group=False, timeout=10.0): + store = c10d.FileStore(self.file_name, self.world_size) + c10d.init_process_group( + backend="gloo", rank=self.rank, world_size=self.world_size, store=store + ) + if with_new_group: + pg = c10d.new_group(backend="gloo") + else: + _pg = c10d.ProcessGroupGloo( + store, self.rank, self.world_size, self.opts(timeout=timeout) + ) + pg = c10d._create_process_group_wrapper( + _pg, + "unused", + store, + self.rank, + self.world_size, + timeout=timeout, + ) + return pg + + def test_collective_hang(self): + pg = self._create_wrapper_pg(timeout=2.0) + self._test_collective_hang(pg) + + # NOTE: these tests are separated by debug level instead of combined into + # one due to https://github.com/pytorch/pytorch/issues/55967, they can be + # combined after that is resolved. + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collectives_op_mismatch_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collectives_op_mismatch(pg) + + @with_dist_debug_levels(levels=["OFF"]) + def test_collectives_op_mismatch(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collectives_op_mismatch(pg) + + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collective_shape_mismatch_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collective_shape_mismatch(pg) + + @with_dist_debug_levels(levels=["OFF"]) + def test_collective_shape_mismatch(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collective_shape_mismatch(pg) + + @skip_if_lt_x_gpu(4) + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collectives_op_mismatch_cuda_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collectives_op_mismatch(pg, use_cuda=True) + + @skip_if_lt_x_gpu(4) + @with_dist_debug_levels(levels=["OFF"]) + def test_collectives_op_mismatch_cuda(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collectives_op_mismatch(pg, use_cuda=True) + + @skip_if_lt_x_gpu(4) + @with_dist_debug_levels(levels=["DETAIL"]) + def test_collective_shape_mismatch_cuda_debug_mode(self): + pg = self._create_wrapper_pg(with_new_group=True) + self._test_collective_shape_mismatch(pg, use_cuda=True) + + @skip_if_lt_x_gpu(4) + @with_dist_debug_levels(levels=["OFF"]) + def test_collective_shape_mismatch_cuda(self): + pg = self._create_wrapper_pg(with_new_group=False) + self._test_collective_shape_mismatch(pg, use_cuda=True) + +if __name__ == "__main__": + assert ( + not torch.cuda._initialized + ), "test_pg_wrapper must not have initialized CUDA context on main process" + + run_tests() diff --git a/test/run_test.py b/test/run_test.py index 38563b2700504..5670da354ce01 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -49,6 +49,7 @@ 'distributed/test_c10d_spawn_gloo', 'distributed/test_c10d_spawn_nccl', 'distributed/test_store', + 'distributed/test_pg_wrapper', 'test_cuda', 'test_jit_cuda_fuser', 'test_cuda_primary_ctx', @@ -311,6 +312,7 @@ 'distributed/test_c10d_spawn_gloo', 'distributed/test_c10d_spawn_nccl', 'distributed/test_store', + 'distributed/test_pg_wrapper', 'test_quantization', 'test_pruning_op', 'test_determination', diff --git a/torch/distributed/CONTRIBUTING.md b/torch/distributed/CONTRIBUTING.md index f6213332670d1..0f4428a9594f6 100644 --- a/torch/distributed/CONTRIBUTING.md +++ b/torch/distributed/CONTRIBUTING.md @@ -81,6 +81,9 @@ python test/distributed/test_c10d_nccl.py # Run the Store tests. python test/distributed/test_store.py +# Run Process Group Wrapper tests. +python test/distributed/test_pg_wrapper.py + # Run distributed tests, including tests for Distributed Data Parallel. python test/run_test.py --verbose -i distributed/test_distributed_fork python test/run_test.py --verbose -i distributed/test_distributed_spawn From 117b7ae38a7895909fa384e5ba0caf2983af8bf9 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Mon, 14 Jun 2021 15:24:22 -0700 Subject: [PATCH 094/305] Remove update-disabled-tests workflow as it is migrated to test-infra (#59986) Summary: Will be replaced by https://github.com/pytorch/test-infra/pull/37 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59986 Reviewed By: seemethere, soulitzer Differential Revision: D29115397 Pulled By: janeyx99 fbshipit-source-id: 2c1a88d6a3fec8cef57818a360884644ec2c7b79 --- .github/workflows/update_disabled_tests.yml | 30 --------------------- 1 file changed, 30 deletions(-) delete mode 100644 .github/workflows/update_disabled_tests.yml diff --git a/.github/workflows/update_disabled_tests.yml b/.github/workflows/update_disabled_tests.yml deleted file mode 100644 index 71c96e5dc2ffb..0000000000000 --- a/.github/workflows/update_disabled_tests.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Update disabled tests - -on: - issues: - types: [opened, edited, labeled, unlabeled, closed, reopened] - # Have the ability to trigger this job manually through the API - workflow_dispatch: - -jobs: - update-disabled-tests: - if: ${{ github.repository_owner == 'pytorch' }} - runs-on: ubuntu-18.04 - steps: - - name: Generate new disabled test list - run: | - # score changes every request, so we strip it out to avoid creating a commit every time we query. - curl 'https://api.github.com/search/issues?q=is%3Aissue+is%3Aopen+label%3A%22module%3A+flaky-tests%22+repo:pytorch/pytorch+in%3Atitle+DISABLED' \ - | sed 's/"score": [0-9\.]*/"score": 0.0/g' > disabled-tests.json - - name: Push file to test-infra repository - uses: dmnemec/copy_file_to_another_repo_action@5f40763ccee2954067adba7fb8326e4df33bcb92 - env: - API_TOKEN_GITHUB: ${{ secrets.TEST_INFRA_TOKEN }} - with: - source_file: 'disabled-tests.json' - destination_repo: 'pytorch/test-infra' - destination_folder: 'stats' - destination_branch: master - user_email: 'test-infra@pytorch.org' - user_name: 'Pytorch Test Infra' - commit_message: 'Updating disabled tests stats' From 5b9fced70aefc7f8648bb1e31b037011811dc4cd Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Mon, 14 Jun 2021 16:18:57 -0700 Subject: [PATCH 095/305] add output_process_fn_grad before sum().backward() (#59971) Summary: This should fix `to_sparse` test issue. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59971 Test Plan: CI Also: directly examine the RuntimeError thrown from test_unsupported_backward - Before: ``` NotImplementedError: Could not run 'aten::sum' with arguments from the 'SparseCPU' backend. ``` - After: ``` to_dense() not supported for float16 on CPU ``` Reviewed By: soulitzer Differential Revision: D29112558 Pulled By: walterddr fbshipit-source-id: c2acd22cd18d5b34d25209b8415feb3ba28fa104 --- test/test_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/test_ops.py b/test/test_ops.py index e355be89c58b6..da6f8462e1e17 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -76,6 +76,8 @@ def test_unsupported_backward(self, device, dtype, op): # TODO: handle non-tensor outputs if not isinstance(result, torch.Tensor): self.skipTest("Skipped! Test does not handle non-tensor outputs") + if sample.output_process_fn_grad is not None: + result = sample.output_process_fn_grad(result) result.sum().backward() # Verifies that backward for each supported floating or complex dtype From 087ac75b26db696169d1b48c20074c4a204633d8 Mon Sep 17 00:00:00 2001 From: Denis Kokarev Date: Mon, 14 Jun 2021 17:29:00 -0700 Subject: [PATCH 096/305] Fix quantized mean operator in QNNPACK backend (#59761) Summary: cc: kimishpatel Fixes https://github.com/pytorch/pytorch/issues/58668 Test it with `pytest -k test_quantized_mean test/test_quantization.py` or `buck test //caffe2/test:quantization -- test_quantized_mean` Pull Request resolved: https://github.com/pytorch/pytorch/pull/59761 Reviewed By: bdhirsh Differential Revision: D29013271 Pulled By: kimishpatel fbshipit-source-id: 020956fb63bd5078856ca17b137be016d3fc29b8 --- .../ATen/native/quantized/cpu/qreduction.cpp | 22 ++++++++++--------- test/quantization/core/test_quantized_op.py | 19 +++++++++++++++- .../testing/_internal/common_quantization.py | 16 ++++++++++++++ 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/aten/src/ATen/native/quantized/cpu/qreduction.cpp b/aten/src/ATen/native/quantized/cpu/qreduction.cpp index 66733b4e75292..d8090d04ac67d 100644 --- a/aten/src/ATen/native/quantized/cpu/qreduction.cpp +++ b/aten/src/ATen/native/quantized/cpu/qreduction.cpp @@ -8,7 +8,7 @@ namespace at { namespace native { #ifdef USE_PYTORCH_QNNPACK -Tensor qnnpack_mean(const Tensor& input, IntArrayRef dim) { +Tensor qnnpack_mean(const Tensor& input, IntArrayRef dim, bool keepdim) { Tensor output; TORCH_CHECK( input.ndimension() == 4, @@ -31,10 +31,15 @@ Tensor qnnpack_mean(const Tensor& input, IntArrayRef dim) { initQNNPACK(); const auto scale = input_contig.q_scale(); const auto zero_point = input_contig.q_zero_point(); - const auto outC = inC; + output = at::_empty_affine_quantized( - {batch_size, outC}, at::device(kCPU).dtype(kQUInt8), scale, zero_point); + keepdim ? IntArrayRef{batch_size, outC, 1, 1} + : IntArrayRef{batch_size, outC}, + at::device(kCPU).dtype(kQUInt8), + scale, + zero_point); + pytorch_qnnp_operator_t qnnpack_operator{nullptr}; const pytorch_qnnp_status createStatus = pytorch_qnnp_create_global_average_pooling_nwc_q8( @@ -75,7 +80,8 @@ Tensor qnnpack_mean(const Tensor& input, IntArrayRef dim) { return output; } #endif -Tensor& mean_out_quantized_cpu(const Tensor& self, +Tensor& mean_out_quantized_cpu( + const Tensor& self, IntArrayRef dim, bool keepdim, c10::optional opt_dtype, @@ -85,12 +91,8 @@ Tensor& mean_out_quantized_cpu(const Tensor& self, self.scalar_type() == kQUInt8 && // QNNPACK currently is only supported for NCHW + dim=(2, 3) // Remove these checks after generic version is implemented. - self.ndimension() == 4 && - dim.size() == 2 && - dim[0] == 2 && - dim[1] == 3 - ){ - result = qnnpack_mean(self, dim); + self.ndimension() == 4 && dim.size() == 2 && dim[0] == 2 && dim[1] == 3) { + result = qnnpack_mean(self, dim, keepdim); return result; } #endif diff --git a/test/quantization/core/test_quantized_op.py b/test/quantization/core/test_quantized_op.py index e1081bbc6a41b..1aff284defbb5 100644 --- a/test/quantization/core/test_quantized_op.py +++ b/test/quantization/core/test_quantized_op.py @@ -21,7 +21,7 @@ from torch.testing._internal.common_utils import TestCase from torch.testing._internal.common_utils import IS_PPC, TEST_WITH_UBSAN, IS_MACOS -from torch.testing._internal.common_quantization import skipIfNoFBGEMM +from torch.testing._internal.common_quantization import skipIfNoFBGEMM, skipIfNoQNNPACK from torch.testing._internal.common_quantized import _quantize, _dequantize, _calculate_dynamic_qparams, \ override_quantized_engine, supported_qengines, override_qengines, _snr from torch.testing._internal.common_quantized import qengine_is_qnnpack @@ -1838,6 +1838,23 @@ def test_mean(self, X, dim): self.assertEqual(Y, qY.dequantize()) + @skipIfNoQNNPACK + @given(keep=st.booleans()) + def test_quantized_mean_qnnpack(self, keep): + with override_quantized_engine("qnnpack"): + # using multiple of 4 sizes to satisfy pytorch_q8gavgpool_ukernel_up8xm__sse2() 4-byte alignment demand under ASAN + in_dim = (4, 4, 4, 4) + if keep: + out_dim = (4, 4, 1, 1) + else: + out_dim = (4, 4) + X = torch.ones(in_dim) + Y = torch.ones(out_dim) + XQ = torch.quantize_per_tensor(X, scale=0.2, zero_point=0, dtype=torch.quint8) + YQ = torch.quantize_per_tensor(Y, scale=0.2, zero_point=0, dtype=torch.quint8) + MQ = XQ.mean((2, 3), keepdim=keep) + self.assertTrue(torch.equal(MQ, YQ)) + """Tests the correctness of the quantized equal op.""" @given(X=hu.tensor(shapes=hu.array_shapes(1, 5, 1, 5), qparams=hu.qparams()), diff --git a/torch/testing/_internal/common_quantization.py b/torch/testing/_internal/common_quantization.py index 03d265b5f19cb..0df65a094fd3b 100644 --- a/torch/testing/_internal/common_quantization.py +++ b/torch/testing/_internal/common_quantization.py @@ -266,6 +266,22 @@ def wrapper(*args, **kwargs): fn(*args, **kwargs) return wrapper +def skipIfNoQNNPACK(fn): + reason = 'Quantized operations require QNNPACK.' + if isinstance(fn, type): + if 'qnnpack' not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if 'qnnpack' not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + return wrapper + try: import torchvision # noqa: F401 HAS_TORCHVISION = True From cbd1e8c3350bb48d590b296b0f7d9eea4ad945c6 Mon Sep 17 00:00:00 2001 From: Hao Lu Date: Mon, 14 Jun 2021 22:53:30 -0700 Subject: [PATCH 097/305] [Static Runtime] Fix bug in aten::to (#59995) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59995 Reviewed By: ajyu Differential Revision: D29083106 fbshipit-source-id: 687ffb121af2716d606c145474942650a2d9ac7e --- torch/csrc/jit/runtime/static/ops.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index 7e22b8a6b397a..3c9590fa9934c 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -1256,6 +1256,11 @@ std::function getNativeOperation(Node* n) { const auto in1_i = p_node->Input(1).toScalarType(); p_node->Output(0) = at::native::to(in0_t, in1_i, in2_i, in3_i, in4_o); } + // in case that Output(0) is an alias of in0_t, copy the tensor. + if (p_node->Output(0).toTensor().unsafeGetTensorImpl() == + in0_t.unsafeGetTensorImpl()) { + p_node->Output(0) = in0_t.clone(); + } }; } return nullptr; From 8d50a4e326e10fe29e322753bb90be15546e5435 Mon Sep 17 00:00:00 2001 From: Garret Catron Date: Mon, 14 Jun 2021 23:30:27 -0700 Subject: [PATCH 098/305] Add support for embeddingBagBytewise in FXGlow Summary: This adds support for embeddingBagBytewise with fp32 scale/bias to FXGlow. Test Plan: buck run //glow/fb/fx/fx_glow:test_fx_glow Reviewed By: jfix71 Differential Revision: D29075288 fbshipit-source-id: 4145486505a903129678216b133bbb8ad71f4fef --- torch/fx/experimental/graph_manipulation.py | 63 ++++++++++++++++----- 1 file changed, 50 insertions(+), 13 deletions(-) diff --git a/torch/fx/experimental/graph_manipulation.py b/torch/fx/experimental/graph_manipulation.py index 6f54cba5d2fd6..b005624a47b87 100644 --- a/torch/fx/experimental/graph_manipulation.py +++ b/torch/fx/experimental/graph_manipulation.py @@ -90,7 +90,9 @@ def get_size_of_node(fx_module: GraphModule, node: Node) -> size_bytes: total_num_of_elems += output_elem # Assume for now if it's quantized then it's qint8 or quint8 if tensor_meta.is_quantized: - size_per_elem_bytes = torch._empty_affine_quantized([], dtype=tensor_meta.dtype).element_size() + size_per_elem_bytes = torch._empty_affine_quantized( + [], dtype=tensor_meta.dtype + ).element_size() else: size_per_elem_bytes = torch.tensor([], dtype=tensor_meta.dtype).element_size() total_size = size_per_elem_bytes * total_num_of_elems @@ -106,7 +108,9 @@ def serialize_stride(stride: Tuple[int]) -> str: return str(list(stride)) -def serialize_tensor_quantization(tensor: torch.Tensor, weights: Dict, pcq_prefix: str) -> Tuple[Dict, Dict]: +def serialize_tensor_quantization( + tensor: torch.Tensor, weights: Dict, pcq_prefix: str +) -> Tuple[Dict, Dict]: """ Args: tensor: The tensor from which we try to extract quantization information. @@ -174,19 +178,29 @@ def serialize_tensor_quantization(tensor: torch.Tensor, weights: Dict, pcq_prefi torch.per_channel_symmetric, }: # per_channel_scales is float64. Here we save it as float32. - weights[f"{pcq_prefix}_per_channel_scales"] = tensor.q_per_channel_scales().float() + weights[ + f"{pcq_prefix}_per_channel_scales" + ] = tensor.q_per_channel_scales().float() scheme["q_per_channel_scales"] = f"{pcq_prefix}_per_channel_scales" per_channel_dict.update( - serialize_weight(weights[f"{pcq_prefix}_per_channel_scales"], weights, f"{pcq_prefix}_per_channel_scales") + serialize_weight( + weights[f"{pcq_prefix}_per_channel_scales"], + weights, + f"{pcq_prefix}_per_channel_scales", + ) ) # per_channel_zero_point is int64. Here we save it as int32. - weights[f"{pcq_prefix}_per_channel_zero_points"] = tensor.q_per_channel_zero_points().int() - scheme[ - "q_per_channel_zero_points" - ] = f"{pcq_prefix}_per_channel_zero_points" + weights[ + f"{pcq_prefix}_per_channel_zero_points" + ] = tensor.q_per_channel_zero_points().int() + scheme["q_per_channel_zero_points"] = f"{pcq_prefix}_per_channel_zero_points" per_channel_dict.update( - serialize_weight(weights[f"{pcq_prefix}_per_channel_zero_points"], weights, f"{pcq_prefix}_per_channel_zero_points") + serialize_weight( + weights[f"{pcq_prefix}_per_channel_zero_points"], + weights, + f"{pcq_prefix}_per_channel_zero_points", + ) ) scheme["q_per_channel_axis"] = tensor.q_per_channel_axis() @@ -201,7 +215,9 @@ def serialize_weight(tensor: torch.Tensor, weights: Dict, name: str) -> Dict: weight_dict[name]["stride"] = serialize_stride(tensor.stride()) if tensor.is_quantized: - quantization_info, per_channel_dict = serialize_tensor_quantization(tensor, weights, name) + quantization_info, per_channel_dict = serialize_tensor_quantization( + tensor, weights, name + ) weight_dict[name].update(quantization_info) weight_dict.update(per_channel_dict) @@ -215,7 +231,9 @@ def serialize_leaf_module( for p_name, p_value in node.attrs_for_lowering.items(): # type: ignore[attr-defined] if isinstance(p_value, torch.Tensor): - weights_metadata.update(serialize_weight(p_value, weights, f"{name_prefix}.{p_name}")) + weights_metadata.update( + serialize_weight(p_value, weights, f"{name_prefix}.{p_name}") + ) weights[f"{name_prefix}.{p_name}"] = p_value else: parameters[p_name] = str(p_value) @@ -341,10 +359,29 @@ def get_node_info(node): if node.op == "get_attr": # If we are targeting a parent constant we update the target. if node.target.startswith("parent."): - stripped_name = node.target[len("parent."):] + stripped_name = node.target[len("parent.") :] node.name = stripped_name node_rep["target"] = stripped_name - weight = serialize_weight(weights[stripped_name], weights, node.target[len("parent."):]) + weight = serialize_weight( + weights[stripped_name], weights, node.target[len("parent.") :] + ) + # For quantized embedding tables we need to update the shape/type, + # so we check if the users of this get_attr is a quantized EB and this is the weight for the EB. + user_targets = { + _get_qualified_name(n.target).replace("glow.fb.fx.", ""): n + for n in node.users.keys() + } + if ( + "acc_ops.embedding_bag_byte_rowwise_offsets" in user_targets + and str( + user_targets[ + "acc_ops.embedding_bag_byte_rowwise_offsets" + ].kwargs["weight"] + ) + == stripped_name + ): + weight[stripped_name]["dtype"] = "acc.uint8fused" + serialized_dict["weights"].update(weight) else: # Find the actual target parameter/buffer from the fx_module. From a1780432fa1b21d2f724b97baea7fbd7427525b8 Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 02:00:08 -0700 Subject: [PATCH 099/305] Move c10d to libtorch(_cuda) (#59563) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59563 ghstack-source-id: 131331264 Test Plan: CI Reviewed By: malfet Differential Revision: D28932239 fbshipit-source-id: 5df6cdfa5253b15cbbc97039fe672d6d97321e34 --- CMakeLists.txt | 6 + caffe2/CMakeLists.txt | 63 +++++++-- test/cpp/rpc/CMakeLists.txt | 2 +- tools/build_variables.bzl | 48 ++++++- torch/CMakeLists.txt | 12 +- torch/csrc/cuda/nccl.h | 1 - torch/lib/c10d/CMakeLists.txt | 164 ---------------------- torch/lib/c10d/FileStore.hpp | 2 +- torch/lib/c10d/GlooDeviceFactory.hpp | 2 +- torch/lib/c10d/HashStore.hpp | 2 +- torch/lib/c10d/ParamCommsUtils.hpp | 5 +- torch/lib/c10d/PrefixStore.hpp | 2 +- torch/lib/c10d/ProcessGroup.hpp | 11 +- torch/lib/c10d/ProcessGroupGloo.hpp | 12 +- torch/lib/c10d/ProcessGroupMPI.hpp | 2 +- torch/lib/c10d/ProcessGroupNCCL.hpp | 2 +- torch/lib/c10d/ProcessGroupRoundRobin.hpp | 2 +- torch/lib/c10d/ProcessGroupWrapper.hpp | 2 +- torch/lib/c10d/Store.hpp | 3 +- torch/lib/c10d/TCPStore.hpp | 2 +- torch/lib/c10d/Utils.hpp | 6 +- torch/lib/c10d/comm.cpp | 1 - torch/lib/c10d/comm.hpp | 4 +- torch/lib/c10d/example/CMakeLists.txt | 5 +- torch/lib/c10d/logger.hpp | 2 +- torch/lib/c10d/reducer.hpp | 9 +- torch/lib/c10d/reducer_cuda.cpp | 4 - torch/lib/c10d/sequence_num.hpp | 3 +- torch/lib/c10d/test/CMakeLists.txt | 43 +++--- 29 files changed, 183 insertions(+), 239 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9fca32640ccc7..af3f21be8fc3b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -289,6 +289,12 @@ cmake_dependent_option( cmake_dependent_option( USE_GLOO_WITH_OPENSSL "Use Gloo with OpenSSL. Only available if USE_GLOO is on." OFF "USE_GLOO AND LINUX AND NOT INTERN_BUILD_MOBILE" OFF) +cmake_dependent_option( + USE_C10D_GLOO "USE C10D GLOO" ON "USE_DISTRIBUTED;USE_GLOO" OFF) +cmake_dependent_option( + USE_C10D_NCCL "USE C10D NCCL" ON "USE_DISTRIBUTED;USE_NCCL" OFF) +cmake_dependent_option( + USE_C10D_MPI "USE C10D MPI" ON "USE_DISTRIBUTED;USE_MPI" OFF) cmake_dependent_option( USE_TENSORPIPE "Use TensorPipe. Only available if USE_DISTRIBUTED is on." ON "USE_DISTRIBUTED" OFF) diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 7b34d8e19b108..20ef2a6376c0c 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -357,8 +357,8 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.cpp" "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.h" ) - target_link_libraries(process_group_agent PRIVATE torch c10d fmt::fmt-header-only) - add_dependencies(process_group_agent torch c10d) + target_link_libraries(process_group_agent PRIVATE torch fmt::fmt-header-only) + add_dependencies(process_group_agent torch) if(USE_TENSORPIPE) add_library(tensorpipe_agent @@ -370,8 +370,8 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp" "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h" ) - target_link_libraries(tensorpipe_agent PRIVATE torch c10d tensorpipe fmt::fmt-header-only) - add_dependencies(tensorpipe_agent torch c10d) + target_link_libraries(tensorpipe_agent PRIVATE torch tensorpipe fmt::fmt-header-only) + add_dependencies(tensorpipe_agent torch) if(USE_CUDA) target_compile_definitions(tensorpipe_agent PUBLIC USE_CUDA) endif() @@ -621,8 +621,11 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) PROPERTIES COMPILE_FLAGS "-DC10_DISABLE_LEGACY_IMPORT" ) endif() - if(USE_DISTRIBUTED AND NOT WIN32) - append_filelist("libtorch_distributed_sources" TORCH_SRCS) + if(USE_DISTRIBUTED) + append_filelist("libtorch_distributed_base_sources" TORCH_SRCS) + if(NOT WIN32) + append_filelist("libtorch_distributed_extra_sources" TORCH_SRCS) + endif() endif() endif() @@ -653,6 +656,17 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) list(APPEND Caffe2_GPU_SRCS ${TORCH_SRC_DIR}/csrc/cuda/nccl.cpp) endif() + if(USE_DISTRIBUTED) + if(BUILD_SPLIT_CUDA) + set(_target "Caffe2_GPU_SRCS_CPP") + else() + set(_target "Caffe2_GPU_SRCS") + endif() + append_filelist("libtorch_cuda_distributed_base_sources" ${_target}) + if(NOT WIN32) + append_filelist("libtorch_cuda_distributed_extra_sources" ${_target}) + endif() + endif() set_source_files_properties( ${TORCH_ROOT}/aten/src/ATen/cuda/detail/LazyNVRTC.cpp PROPERTIES COMPILE_DEFINITIONS "NVRTC_SHORTHASH=${CUDA_NVRTC_SHORTHASH}" @@ -670,6 +684,12 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) list(APPEND Caffe2_HIP_SRCS ${TORCH_SRC_DIR}/csrc/cuda/nccl.cpp) endif() + if(USE_DISTRIBUTED) + append_filelist("libtorch_cuda_distributed_base_sources" Caffe2_HIP_SRCS) + if(NOT WIN32) + append_filelist("libtorch_cuda_distributed_extra_sources" Caffe2_HIP_SRCS) + endif() + endif() # caffe2_nvrtc's stubs to driver APIs are useful for HIP. # See NOTE [ ATen NVRTC Stub and HIP ] add_library(caffe2_nvrtc SHARED ${ATen_NVRTC_STUB_SRCS}) @@ -1047,6 +1067,9 @@ endif() install(DIRECTORY "${TORCH_SRC_DIR}/csrc" DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch FILES_MATCHING PATTERN "*.h") + install(DIRECTORY "${TORCH_SRC_DIR}/lib/c10d" + DESTINATION ${TORCH_INSTALL_INCLUDE_DIR} + FILES_MATCHING PATTERN "*.hpp") install(FILES "${TORCH_SRC_DIR}/script.h" "${TORCH_SRC_DIR}/extension.h" @@ -1210,9 +1233,31 @@ endif() # Pass USE_DISTRIBUTED to torch_cpu, as some codes in jit/pickler.cpp and # jit/unpickler.cpp need to be compiled only when USE_DISTRIBUTED is set if(USE_DISTRIBUTED) - target_compile_definitions(torch_cpu PRIVATE - USE_DISTRIBUTED - ) + # Needed to support the inclusion of c10d/Foo.hpp headers. + target_include_directories(torch_cpu PUBLIC ${TORCH_SRC_DIR}/lib) + target_compile_definitions(torch_cpu PRIVATE USE_DISTRIBUTED) + if(USE_GLOO AND USE_C10D_GLOO) + target_compile_definitions(torch_cpu PUBLIC USE_C10D_GLOO) + endif() + if(USE_NCCL AND USE_C10D_NCCL) + if(USE_ROCM) + target_compile_definitions(torch_hip PUBLIC USE_C10D_NCCL) + else() + if(BUILD_SPLIT_CUDA) + target_compile_definitions(torch_cuda_cpp PUBLIC USE_C10D_NCCL) + else() + target_compile_definitions(torch_cuda PUBLIC USE_C10D_NCCL) + endif() + endif() + endif() + if(USE_MPI AND USE_C10D_MPI) + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + set_source_files_properties( + "${TORCH_SRC_DIR}/lib/c10d/ProcessGroupMPI.cpp" + PROPERTIES COMPILE_FLAGS -Wno-deprecated-declarations) + endif() + target_compile_definitions(torch_cpu PUBLIC USE_C10D_MPI) + endif() # Pass USE_RPC in order to reduce use of # #if defined(USE_DISTRIBUTED) && !defined(_WIN32) # need to be removed when RPC is supported diff --git a/test/cpp/rpc/CMakeLists.txt b/test/cpp/rpc/CMakeLists.txt index caa642aeb6812..0eff382d2b1b8 100644 --- a/test/cpp/rpc/CMakeLists.txt +++ b/test/cpp/rpc/CMakeLists.txt @@ -5,7 +5,7 @@ set(TORCH_RPC_TEST_SOURCES ${TORCH_RPC_TEST_DIR}/test_wire_serialization.cpp ) set(TORCH_RPC_TEST_DEPENDENCY_LIBS - torch c10d gtest process_group_agent + torch gtest process_group_agent ) if(USE_GLOO) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 8ae8c3805381d..d7f2d8bf95ab3 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -313,7 +313,28 @@ core_sources_full = core_sources_full_mobile + [ libtorch_core_sources = sorted(core_sources_common + core_sources_full + core_trainer_sources) -libtorch_distributed_sources = [ +# These files are the only ones that are supported on Windows. +libtorch_distributed_base_sources = [ + "torch/lib/c10d/comm.cpp", + "torch/lib/c10d/default_comm_hooks.cpp", + "torch/lib/c10d/FileStore.cpp", + "torch/lib/c10d/GlooDeviceFactory.cpp", + "torch/lib/c10d/logger.cpp", + "torch/lib/c10d/ParamCommsUtils.cpp", + "torch/lib/c10d/PrefixStore.cpp", + "torch/lib/c10d/ProcessGroup.cpp", + "torch/lib/c10d/ProcessGroupGloo.cpp", + "torch/lib/c10d/ProcessGroupMPI.cpp", + "torch/lib/c10d/ProcessGroupWrapper.cpp", + "torch/lib/c10d/reducer.cpp", + "torch/lib/c10d/sequence_num.cpp", + "torch/lib/c10d/Store.cpp", + "torch/lib/c10d/TCPStore.cpp", + "torch/lib/c10d/Utils.cpp", +] + +# These files are only supported on Linux (and others) but not on Windows. +libtorch_distributed_extra_sources = [ "torch/csrc/distributed/autograd/autograd.cpp", "torch/csrc/distributed/autograd/utils.cpp", "torch/csrc/distributed/autograd/context/container.cpp", @@ -350,8 +371,12 @@ libtorch_distributed_sources = [ "torch/csrc/distributed/rpc/types.cpp", "torch/csrc/distributed/rpc/utils.cpp", "torch/csrc/distributed/rpc/metrics/registry.cpp", + "torch/lib/c10d/HashStore.cpp", + "torch/lib/c10d/ProcessGroupRoundRobin.cpp", ] +libtorch_distributed_sources = libtorch_distributed_base_sources + libtorch_distributed_extra_sources + jit_sources_full = [ "torch/csrc/jit/codegen/cuda/interface.cpp", "torch/csrc/jit/passes/lower_graph.cpp", @@ -490,7 +515,20 @@ libtorch_cuda_core_sources = [ "torch/csrc/jit/runtime/register_cuda_ops.cpp", ] -libtorch_cuda_sources = libtorch_cuda_core_sources + [ +# These files are the only ones that are supported on Windows. +libtorch_cuda_distributed_base_sources = [ + "torch/lib/c10d/reducer_cuda.cpp", +] + +# These files are only supported on Linux (and others) but not on Windows. +libtorch_cuda_distributed_extra_sources = [ + "torch/lib/c10d/NCCLUtils.cpp", + "torch/lib/c10d/ProcessGroupNCCL.cpp", +] + +libtorch_cuda_distributed_sources = libtorch_cuda_distributed_base_sources + libtorch_cuda_distributed_extra_sources + +libtorch_cuda_sources = libtorch_cuda_core_sources + libtorch_cuda_distributed_sources + [ "torch/csrc/cuda/nccl.cpp", ] @@ -665,13 +703,9 @@ libtorch_python_core_sources = [ ] libtorch_python_distributed_core_sources = [ - "torch/lib/c10d/comm.cpp", - "torch/lib/c10d/default_comm_hooks.cpp", - "torch/lib/c10d/reducer.cpp", - "torch/lib/c10d/reducer_cuda.cpp", - "torch/lib/c10d/logger.cpp", "torch/csrc/distributed/c10d/python_comm_hook.cpp", "torch/csrc/distributed/c10d/init.cpp", + "torch/lib/c10d/frontend.cpp", ] libtorch_python_distributed_sources = libtorch_python_distributed_core_sources + [ diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index f2e61e27ac26d..197926f309838 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -278,7 +278,17 @@ if(USE_DISTRIBUTED) list(APPEND TORCH_PYTHON_LINK_LIBRARIES tensorpipe) list(APPEND TORCH_PYTHON_PUBLIC_COMPILE_DEFINITIONS USE_TENSORPIPE) endif() - list(APPEND TORCH_PYTHON_LINK_LIBRARIES c10d) + # NCCL is a private dependency of libtorch, but libtorch_python includes + # some private headers of libtorch, which in turn include NCCL. As a hacky + # alternative to making NCCL a public dependency of libtorch, we make it + # a private dependency of libtorch_python as well. + if(USE_NCCL) + list(APPEND TORCH_PYTHON_LINK_LIBRARIES __caffe2_nccl) + endif() + # Same for MPI. + if(USE_MPI) + list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${MPI_CXX_LIBRARIES}) + endif() list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_C10D) endif() diff --git a/torch/csrc/cuda/nccl.h b/torch/csrc/cuda/nccl.h index f331a5ce8542c..2ff7273345d35 100644 --- a/torch/csrc/cuda/nccl.h +++ b/torch/csrc/cuda/nccl.h @@ -2,7 +2,6 @@ #include #include -#include #include #include diff --git a/torch/lib/c10d/CMakeLists.txt b/torch/lib/c10d/CMakeLists.txt index 3e93108e2a414..20b8b2502d16f 100644 --- a/torch/lib/c10d/CMakeLists.txt +++ b/torch/lib/c10d/CMakeLists.txt @@ -1,166 +1,5 @@ cmake_minimum_required(VERSION 3.2 FATAL_ERROR) -# Find modules. -list(APPEND CMAKE_MODULE_PATH - ${CMAKE_CURRENT_SOURCE_DIR}/../../../cmake/public - ${CMAKE_CURRENT_SOURCE_DIR}/../../../cmake/Modules - ${CMAKE_CURRENT_SOURCE_DIR}/../../../cmake/Modules_CUDA_fix) - -if(USE_CUDA) - add_definitions(-DUSE_CUDA=1) -elseif(USE_ROCM) - add_definitions(-DUSE_ROCM=1) - add_definitions(-D__HIP_PLATFORM_HCC__=1) -else() - message(STATUS "Building c10d without CUDA/ROCm support") -endif() - -if(USE_TBB) - include_directories(${TBB_ROOT_DIR}/include) -endif() - -if(USE_GLOO) - option(USE_C10D_GLOO "USE C10D GLOO" ON) -endif() - -if(USE_NCCL) - option(USE_C10D_NCCL "USE C10D NCCL" ON) -endif() - -if(USE_MPI) - find_package(MPI) - if(MPI_FOUND) - message(STATUS "MPI_INCLUDE_PATH: ${MPI_INCLUDE_PATH}") - message(STATUS "MPI_LIBRARIES: ${MPI_LIBRARIES}") - message(STATUS "MPIEXEC: ${MPIEXEC}") - option(USE_C10D_MPI "USE C10D MPI" ON) - else() - message(STATUS "Not able to find MPI, will compile c10d without MPI support") - endif() -endif() - -function(copy_header file) - configure_file(${file} ${CMAKE_BINARY_DIR}/include/c10d/${file} COPYONLY) -endfunction() - -set(C10D_SRCS - frontend.cpp - FileStore.cpp - ParamCommsUtils.cpp - PrefixStore.cpp - ProcessGroup.cpp - sequence_num.cpp - Store.cpp - TCPStore.cpp - Utils.cpp - ) - -if(NOT WIN32) - list(APPEND C10D_SRCS HashStore.cpp ProcessGroupRoundRobin.cpp) -endif() - -set(C10D_LIBS torch) - -if(USE_C10D_NCCL) - list(APPEND C10D_SRCS ProcessGroupNCCL.cpp NCCLUtils.cpp) - list(APPEND C10D_LIBS __caffe2_nccl) -endif() - -if(USE_C10D_MPI) - list(APPEND C10D_SRCS ProcessGroupMPI.cpp) - if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - set_source_files_properties(ProcessGroupMPI.cpp PROPERTIES COMPILE_FLAGS -Wno-deprecated-declarations) - endif() - list(APPEND C10D_LIBS ${MPI_LIBRARIES}) -endif() - -if(USE_C10D_GLOO) - list(APPEND C10D_SRCS ProcessGroupGloo.cpp GlooDeviceFactory.cpp ProcessGroupWrapper.cpp) - list(APPEND C10D_LIBS gloo) - if(USE_CUDA) - list(APPEND C10D_LIBS gloo_cuda) - endif() -endif() - -add_library(c10d STATIC ${C10D_SRCS}) -set_property(TARGET c10d PROPERTY POSITION_INDEPENDENT_CODE ON) -set_property(TARGET c10d PROPERTY CXX_STANDARD 14) - -if(NOT MSVC) - target_compile_options(c10d PUBLIC - -Wall - -Wextra - -Wno-unused-parameter - -Wno-missing-field-initializers - -Wno-write-strings - -Wno-unknown-pragmas - ) -endif() - -add_dependencies(c10d torch) - -if(USE_C10D_GLOO) - add_dependencies(c10d gloo) - if(USE_CUDA) - add_dependencies(c10d gloo_cuda) - endif() -endif() - -target_include_directories(c10d PUBLIC - ${CMAKE_BINARY_DIR}/aten/src # provides "ATen/TypeExtendedInterface.h" to ATen.h - ${CMAKE_BINARY_DIR}/caffe2/aten/src # provides to THC.h - ) - -# For -target_include_directories(c10d PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/..) - -if(USE_C10D_NCCL) - target_compile_definitions(c10d PUBLIC USE_C10D_NCCL) -endif() - -if(USE_C10D_MPI) - target_compile_definitions(c10d PUBLIC USE_C10D_MPI) -endif() - -if(USE_C10D_GLOO) - target_compile_definitions(c10d PUBLIC USE_C10D_GLOO) -endif() - -copy_header(FileStore.hpp) -copy_header(ParamCommsUtils.hpp) -copy_header(PrefixStore.hpp) -copy_header(ProcessGroup.hpp) -copy_header(Store.hpp) -copy_header(TCPStore.hpp) -copy_header(Types.hpp) -copy_header(Utils.hpp) -copy_header(sequence_num.hpp) -if(USE_GLOO) - copy_header(ProcessGroupGloo.hpp) - copy_header(GlooDeviceFactory.hpp) - copy_header(ProcessGroupWrapper.hpp) -endif() -if(NOT WIN32) - copy_header(HashStore.hpp) - copy_header(UnixSockUtils.hpp) -else() - copy_header(WinSockUtils.hpp) -endif() - -if(USE_C10D_NCCL) - copy_header(ProcessGroupNCCL.hpp) - copy_header(NCCLUtils.hpp) -endif() - -if(USE_C10D_MPI) - target_include_directories(c10d PUBLIC ${MPI_INCLUDE_PATH}) - copy_header(ProcessGroupMPI.hpp) -endif() - -target_link_libraries(c10d PUBLIC ${C10D_LIBS}) - -install(TARGETS c10d DESTINATION lib) - option(BUILD_EXAMPLES "Build examples" OFF) if(BUILD_EXAMPLES) add_subdirectory(example) @@ -171,6 +10,3 @@ if(BUILD_TEST) enable_testing() add_subdirectory(test) endif() - -# Install all header files that were prepared in the build directory -install(DIRECTORY ${CMAKE_BINARY_DIR}/include/ DESTINATION include) diff --git a/torch/lib/c10d/FileStore.hpp b/torch/lib/c10d/FileStore.hpp index 636f36aef9faf..814dc4823e947 100644 --- a/torch/lib/c10d/FileStore.hpp +++ b/torch/lib/c10d/FileStore.hpp @@ -9,7 +9,7 @@ namespace c10d { -class FileStore : public Store { +class TORCH_API FileStore : public Store { public: explicit FileStore(const std::string& path, int numWorkers); diff --git a/torch/lib/c10d/GlooDeviceFactory.hpp b/torch/lib/c10d/GlooDeviceFactory.hpp index 7d038180bfdb0..dd37b261062f0 100644 --- a/torch/lib/c10d/GlooDeviceFactory.hpp +++ b/torch/lib/c10d/GlooDeviceFactory.hpp @@ -10,7 +10,7 @@ namespace c10d { -class GlooDeviceFactory { +class TORCH_API GlooDeviceFactory { public: // Create new device instance for specific interface. static std::shared_ptr<::gloo::transport::Device> makeDeviceForInterface( diff --git a/torch/lib/c10d/HashStore.hpp b/torch/lib/c10d/HashStore.hpp index 623836f2d758c..94b01cd171b8a 100644 --- a/torch/lib/c10d/HashStore.hpp +++ b/torch/lib/c10d/HashStore.hpp @@ -10,7 +10,7 @@ namespace c10d { -class HashStore : public Store { +class TORCH_API HashStore : public Store { public: ~HashStore() override {} diff --git a/torch/lib/c10d/ParamCommsUtils.hpp b/torch/lib/c10d/ParamCommsUtils.hpp index 31936cb323200..ac06d8fd2ab7f 100644 --- a/torch/lib/c10d/ParamCommsUtils.hpp +++ b/torch/lib/c10d/ParamCommsUtils.hpp @@ -2,14 +2,15 @@ #include #include +#include #include #include namespace torch { -extern const std::string kParamCommsCallName; +extern TORCH_API const std::string kParamCommsCallName; -class ParamCommsDebugInfo +class TORCH_API ParamCommsDebugInfo : public c10::DebugInfoBase { public: diff --git a/torch/lib/c10d/PrefixStore.hpp b/torch/lib/c10d/PrefixStore.hpp index ea3a4af05b56f..c9e57312fac6a 100644 --- a/torch/lib/c10d/PrefixStore.hpp +++ b/torch/lib/c10d/PrefixStore.hpp @@ -5,7 +5,7 @@ namespace c10d { -class PrefixStore : public Store { +class TORCH_API PrefixStore : public Store { public: explicit PrefixStore( const std::string& prefix, diff --git a/torch/lib/c10d/ProcessGroup.hpp b/torch/lib/c10d/ProcessGroup.hpp index 3a3ffa6b95d67..e7739bdce7f00 100644 --- a/torch/lib/c10d/ProcessGroup.hpp +++ b/torch/lib/c10d/ProcessGroup.hpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -50,10 +51,10 @@ enum class OpType : std::uint8_t { }; // Converts OpType to human readable string. -std::string opTypeToString(OpType opType); +TORCH_API std::string opTypeToString(OpType opType); // Whether or not an OP is an p2p op (SEND, RECV, RECVANYSOURCE) -bool isP2POp(OpType opType); +TORCH_API bool isP2POp(OpType opType); // ProcessGroup is a base class that captures collective and point to // point communication in a fixed set of processes. @@ -75,13 +76,13 @@ bool isP2POp(OpType opType); // process group to find each other (referred to as rendezvous from // hereon) // -class ProcessGroup : public torch::CustomClassHolder { +class TORCH_API ProcessGroup : public torch::CustomClassHolder { public: // Please do not use ProcessGroup::Work API, it is going away, to be // replaced by ivalue::Future. // Python binding for this class might change, please do not assume // this will be bound using pybind. - class Work : public torch::CustomClassHolder { + class TORCH_API Work : public torch::CustomClassHolder { public: Work( int rank = -1, @@ -176,7 +177,7 @@ class ProcessGroup : public torch::CustomClassHolder { // when constructing a ProcessGroup. Each ProcessGroup subclass should // extend this struct and define its options if it wants to provide more // config options (beyond basic ones defined here) to end user. - struct Options : torch::CustomClassHolder { + struct TORCH_API Options : torch::CustomClassHolder { explicit Options( std::string backend, std::chrono::milliseconds timeout = kProcessGroupDefaultTimeout) diff --git a/torch/lib/c10d/ProcessGroupGloo.hpp b/torch/lib/c10d/ProcessGroupGloo.hpp index 32e3799e94f3d..19ebed4856bfd 100644 --- a/torch/lib/c10d/ProcessGroupGloo.hpp +++ b/torch/lib/c10d/ProcessGroupGloo.hpp @@ -50,7 +50,7 @@ constexpr const char* GLOO_BACKEND_NAME = "gloo"; // number can be automatically tuned, but only if we let a single // process take charge, and have it broadcast the limits. // -class ProcessGroupGloo : public ProcessGroup { +class TORCH_API ProcessGroupGloo : public ProcessGroup { public: // AsyncWork is the Gloo specific superclass for asynchronous work items. // We can split asynchronous work into 3 phases: @@ -68,7 +68,7 @@ class ProcessGroupGloo : public ProcessGroup { // // FIXME: This probably should be called WorkGloo since the work is executed in sync mode // by a background thread. - class AsyncWork : public ProcessGroup::Work { + class TORCH_API AsyncWork : public ProcessGroup::Work { public: explicit AsyncWork( std::vector> outputTensors, @@ -97,7 +97,7 @@ class ProcessGroupGloo : public ProcessGroup { }; // Wrap c10d store as Gloo store - class GlooStore : public ::gloo::rendezvous::Store { + class TORCH_API GlooStore : public ::gloo::rendezvous::Store { public: GlooStore(const c10::intrusive_ptr<::c10d::Store>& store) : store_(store) {} @@ -140,7 +140,7 @@ class ProcessGroupGloo : public ProcessGroup { // recv operation. It keeps a reference to the tensor it is // operating on to prevent it from being deallocated while the // operation is still in flight. - class SendWork : public ProcessGroup::Work { + class TORCH_API SendWork : public ProcessGroup::Work { public: explicit SendWork( at::Tensor& tensor, @@ -155,7 +155,7 @@ class ProcessGroupGloo : public ProcessGroup { std::unique_ptr<::gloo::transport::UnboundBuffer> buffer_; }; - class RecvWork : public ProcessGroup::Work { + class TORCH_API RecvWork : public ProcessGroup::Work { public: explicit RecvWork( at::Tensor& tensor, @@ -174,7 +174,7 @@ class ProcessGroupGloo : public ProcessGroup { int srcRank_; }; - struct Options : public ProcessGroup::Options { + struct TORCH_API Options : public ProcessGroup::Options { explicit Options( std::chrono::milliseconds timeout = kProcessGroupDefaultTimeout); diff --git a/torch/lib/c10d/ProcessGroupMPI.hpp b/torch/lib/c10d/ProcessGroupMPI.hpp index 95363313b20fd..8067b13b7bf6d 100644 --- a/torch/lib/c10d/ProcessGroupMPI.hpp +++ b/torch/lib/c10d/ProcessGroupMPI.hpp @@ -78,7 +78,7 @@ struct WorkEntry { // // CUDA tensor can be supported if the MPI used is CUDA-aware MPI, and // ProcessGroupMPI will automatically detect this support. -class ProcessGroupMPI : public ProcessGroup { +class TORCH_API ProcessGroupMPI : public ProcessGroup { public: class WorkMPI : public ProcessGroup::Work { public: diff --git a/torch/lib/c10d/ProcessGroupNCCL.hpp b/torch/lib/c10d/ProcessGroupNCCL.hpp index bafb76d7c6dd0..a35b4681acf93 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.hpp +++ b/torch/lib/c10d/ProcessGroupNCCL.hpp @@ -70,7 +70,7 @@ constexpr const char* NCCL_BACKEND_NAME = "nccl"; // work->wait() // // // Now continue on other work in the current stream. -class ProcessGroupNCCL : public ProcessGroup { +class TORCH_API ProcessGroupNCCL : public ProcessGroup { public: class WorkNCCL : public ProcessGroup::Work, public std::enable_shared_from_this { diff --git a/torch/lib/c10d/ProcessGroupRoundRobin.hpp b/torch/lib/c10d/ProcessGroupRoundRobin.hpp index 393bb9cda3ac3..d5450badaac52 100644 --- a/torch/lib/c10d/ProcessGroupRoundRobin.hpp +++ b/torch/lib/c10d/ProcessGroupRoundRobin.hpp @@ -18,7 +18,7 @@ constexpr const char* ROUND_ROBIN_BACKEND_NAME = "round_robin"; // across all processes in the process group. This is the only way that we // can guarantee to match up the same calls among all processes. // -class ProcessGroupRoundRobin final : public ProcessGroup { +class TORCH_API ProcessGroupRoundRobin final : public ProcessGroup { public: explicit ProcessGroupRoundRobin( int rank, diff --git a/torch/lib/c10d/ProcessGroupWrapper.hpp b/torch/lib/c10d/ProcessGroupWrapper.hpp index 9ee435593681d..a8f9b454b0037 100644 --- a/torch/lib/c10d/ProcessGroupWrapper.hpp +++ b/torch/lib/c10d/ProcessGroupWrapper.hpp @@ -9,7 +9,7 @@ namespace c10d { -class ProcessGroupWrapper : public ProcessGroup { +class TORCH_API ProcessGroupWrapper : public ProcessGroup { public: explicit ProcessGroupWrapper( c10::intrusive_ptr pg, diff --git a/torch/lib/c10d/Store.hpp b/torch/lib/c10d/Store.hpp index 0bac043e638ac..ee15f767dee46 100644 --- a/torch/lib/c10d/Store.hpp +++ b/torch/lib/c10d/Store.hpp @@ -6,6 +6,7 @@ #include #include +#include #include namespace c10d { @@ -15,7 +16,7 @@ namespace c10d { using WatchKeyCallback = std::function, c10::optional)>; -class Store : public torch::CustomClassHolder { +class TORCH_API Store : public torch::CustomClassHolder { public: static constexpr std::chrono::milliseconds kDefaultTimeout = std::chrono::seconds(300); diff --git a/torch/lib/c10d/TCPStore.hpp b/torch/lib/c10d/TCPStore.hpp index 6a4155989748f..4457d830ea0c2 100644 --- a/torch/lib/c10d/TCPStore.hpp +++ b/torch/lib/c10d/TCPStore.hpp @@ -36,7 +36,7 @@ struct TCPStoreOptions { bool multiTenant = false; }; -class TCPStore : public Store { +class TORCH_API TCPStore : public Store { public: explicit TCPStore(std::string host, const TCPStoreOptions& opts = {}); diff --git a/torch/lib/c10d/Utils.hpp b/torch/lib/c10d/Utils.hpp index 5beb5f1c6708b..aa27c09086812 100644 --- a/torch/lib/c10d/Utils.hpp +++ b/torch/lib/c10d/Utils.hpp @@ -45,12 +45,12 @@ extern const char* kDistDebugDetailLogLevel; extern const char* kDistDebugInfoLogLevel; extern const char* kDistDebugOffLogLevel; -std::string parse_env(const char* env_var_name); +TORCH_API std::string parse_env(const char* env_var_name); -DistributedDebugLevel parseDistDebugLevel(); +TORCH_API DistributedDebugLevel parseDistDebugLevel(); // Retrieve tensor shapes from a given tensor. -std::vector getTensorShapes(const std::vector& tensors); +TORCH_API std::vector getTensorShapes(const std::vector& tensors); // Turns at::IntArrayRef into "(1, 2, 3, 4)". inline std::string toString(at::IntArrayRef l) { diff --git a/torch/lib/c10d/comm.cpp b/torch/lib/c10d/comm.cpp index afc0b1d8c42bd..c2d96c39d89df 100644 --- a/torch/lib/c10d/comm.cpp +++ b/torch/lib/c10d/comm.cpp @@ -5,7 +5,6 @@ #include #include #include -#include #include namespace c10d { diff --git a/torch/lib/c10d/comm.hpp b/torch/lib/c10d/comm.hpp index 67f403971fb3e..dfe741a8bb05a 100644 --- a/torch/lib/c10d/comm.hpp +++ b/torch/lib/c10d/comm.hpp @@ -7,14 +7,14 @@ namespace c10d { // Broadcast many tensors to all processes in the process group. -void broadcast_coalesced( +TORCH_API void broadcast_coalesced( c10::intrusive_ptr process_group, at::TensorList tensors, size_t buffer_size, int rank = 0); // This class passes bucket contents tensor to DDP communication hook. -class GradBucket { +class TORCH_API GradBucket { public: explicit GradBucket( size_t index, diff --git a/torch/lib/c10d/example/CMakeLists.txt b/torch/lib/c10d/example/CMakeLists.txt index ba5f058148099..2a56efacd1cd9 100644 --- a/torch/lib/c10d/example/CMakeLists.txt +++ b/torch/lib/c10d/example/CMakeLists.txt @@ -1,3 +1,6 @@ add_executable(allreduce allreduce.cpp) target_include_directories(allreduce PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) -target_link_libraries(allreduce pthread c10d) +target_link_libraries(allreduce pthread torch_cpu) +if(USE_CUDA) + target_link_libraries(allreduce torch_cuda) +endif() diff --git a/torch/lib/c10d/logger.hpp b/torch/lib/c10d/logger.hpp index 1895e0aabdfcb..e56f1573dc808 100644 --- a/torch/lib/c10d/logger.hpp +++ b/torch/lib/c10d/logger.hpp @@ -3,7 +3,7 @@ namespace c10d { -class Logger { +class TORCH_API Logger { public: explicit Logger(std::shared_ptr reducer); // Set logging data that can be got during DistributedDataParallel diff --git a/torch/lib/c10d/reducer.hpp b/torch/lib/c10d/reducer.hpp index e9275d0d0555b..41900f19cbb45 100644 --- a/torch/lib/c10d/reducer.hpp +++ b/torch/lib/c10d/reducer.hpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -30,7 +31,7 @@ constexpr int kDDPRuntimeLoggingSampleRate = 100; // Forward declaration class Logger; -class Timer { +class TORCH_API Timer { public: enum class Event { kForwardStart, @@ -52,7 +53,7 @@ class Timer { C10_DECLARE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device); -class Reducer { +class TORCH_API Reducer { public: // The constructor takes a list of variables for every model replica. // The bucket assignment for this reducer is specified as a list of @@ -492,7 +493,7 @@ class Reducer { // The index of tensors[i] assigned to bucket is tensor_indices[i], // when tensor_indices is empty, the index of tensors[i] assigned to // bucket is i. -std::vector> compute_bucket_assignment_by_size( +TORCH_API std::vector> compute_bucket_assignment_by_size( const std::vector& tensors, const std::vector& bucket_size, const std::vector& expect_sparse_gradient = {}, @@ -500,7 +501,7 @@ std::vector> compute_bucket_assignment_by_size( // Verify models across all processes are the same as model on rank 0 with // respect to no. of params and matching dtype/size/layout. -void verify_replica0_across_processes( +TORCH_API void verify_replica0_across_processes( c10::intrusive_ptr process_group, std::vector> model_replicas); } // namespace c10d diff --git a/torch/lib/c10d/reducer_cuda.cpp b/torch/lib/c10d/reducer_cuda.cpp index 0f55b5a131181..c21546ef0dc12 100644 --- a/torch/lib/c10d/reducer_cuda.cpp +++ b/torch/lib/c10d/reducer_cuda.cpp @@ -1,7 +1,5 @@ #include -#ifdef USE_CUDA - #include #include @@ -85,5 +83,3 @@ C10_REGISTER_TYPED_CLASS(TimerRegistry, c10::kCUDA, CudaTimer); } // namespace } // namespace c10d - -#endif diff --git a/torch/lib/c10d/sequence_num.hpp b/torch/lib/c10d/sequence_num.hpp index 71732bcbe4a00..c4e0b98bd6773 100644 --- a/torch/lib/c10d/sequence_num.hpp +++ b/torch/lib/c10d/sequence_num.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include @@ -36,7 +37,7 @@ inline uint64_t fromVec(const std::vector& values) { return num; } -class SequenceNum { +class TORCH_API SequenceNum { public: SequenceNum(); explicit SequenceNum(const uint64_t num); diff --git a/torch/lib/c10d/test/CMakeLists.txt b/torch/lib/c10d/test/CMakeLists.txt index b74d4b65f70f7..10df84d86d681 100644 --- a/torch/lib/c10d/test/CMakeLists.txt +++ b/torch/lib/c10d/test/CMakeLists.txt @@ -1,7 +1,7 @@ if(USE_CUDA) cuda_add_library(c10d_cuda_test CUDATest.cu) - target_link_libraries(c10d_cuda_test c10d) - add_dependencies(c10d_cuda_test c10d) + target_link_libraries(c10d_cuda_test torch_cuda) + add_dependencies(c10d_cuda_test torch_cuda) endif() function(c10d_add_test test_src) @@ -16,29 +16,40 @@ function(c10d_add_test test_src) add_test(NAME ${test_name} COMMAND $) endfunction() -c10d_add_test(FileStoreTest.cpp c10d gtest_main) -c10d_add_test(TCPStoreTest.cpp c10d gtest_main) +c10d_add_test(FileStoreTest.cpp torch_cpu gtest_main) +c10d_add_test(TCPStoreTest.cpp torch_cpu gtest_main) if(NOT WIN32) - c10d_add_test(HashStoreTest.cpp c10d gtest_main) + c10d_add_test(HashStoreTest.cpp torch_cpu gtest_main) endif() if(USE_CUDA) - if(USE_C10D_GLOO) - c10d_add_test(ProcessGroupGlooTest.cpp c10d c10d_cuda_test gtest_main) - c10d_add_test(ProcessGroupGlooAsyncTest.cpp c10d c10d_cuda_test gtest_main) + if(USE_GLOO AND USE_C10D_GLOO) + c10d_add_test(ProcessGroupGlooTest.cpp torch_cpu c10d_cuda_test gtest_main) + c10d_add_test(ProcessGroupGlooAsyncTest.cpp torch_cpu c10d_cuda_test gtest_main) endif() - if(USE_C10D_NCCL) - c10d_add_test(ProcessGroupNCCLTest.cpp c10d c10d_cuda_test gtest_main) - c10d_add_test(ProcessGroupNCCLErrorsTest.cpp c10d c10d_cuda_test - gtest_main) + if(USE_NCCL AND USE_C10D_NCCL) + # NCCL is a private dependency of libtorch, but the tests include some + # private headers of libtorch, which in turn include NCCL. As a hacky + # alternative to making NCCL a public dependency of libtorch, we make it + # a private dependency of the tests as well. + c10d_add_test( + ProcessGroupNCCLTest.cpp + torch_cpu c10d_cuda_test gtest_main __caffe2_nccl) + c10d_add_test( + ProcessGroupNCCLErrorsTest.cpp + torch_cpu c10d_cuda_test gtest_main __caffe2_nccl) endif() else() - if(USE_C10D_GLOO) - c10d_add_test(ProcessGroupGlooTest.cpp c10d gtest_main) + if(USE_GLOO AND USE_C10D_GLOO) + c10d_add_test(ProcessGroupGlooTest.cpp torch_cpu gtest_main) endif() endif() -if(USE_C10D_MPI) +if(USE_MPI AND USE_C10D_MPI) add_definitions(-DMPIEXEC=${MPIEXEC}) - c10d_add_test(ProcessGroupMPITest.cpp c10d) + # MPI is a private dependency of libtorch, but the tests include some + # private headers of libtorch, which in turn include MPI. As a hacky + # alternative to making MPI a public dependency of libtorch, we make it + # a private dependency of the tests as well. + c10d_add_test(ProcessGroupMPITest.cpp torch_cpu ${MPI_CXX_LIBRARIES}) endif() From cf63893211e98474ac3bd52de5b3fcd4272accc1 Mon Sep 17 00:00:00 2001 From: Martin Yuan Date: Tue, 15 Jun 2021 02:06:48 -0700 Subject: [PATCH 100/305] Enable implicit operator versioning via number of arguments (#58852) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58852 Enable implicit operator versioning via number of arguments from Mobile. 1. By default, TS doesn't emit instructions for tailing default args and the provided number of specified args is serialized to bytecode. From interpreter the default values are fetched from operator schema. The implementation has been landed in #56845. Please refer to #56845 for details. 2. Since there is bytecode schema change, the bytecode version is bumped from 5 to 6. 3. The corresponding backport function is provided, for forward compatibility use. Note that because there is instruction change, a global flag is used as the switch to control the two versions. Test Plan: Imported from OSS Reviewed By: raziel Differential Revision: D28789746 Pulled By: iseeyuan fbshipit-source-id: 6e5f16460c79b2bd3312de02d0f57b79f50bf66b --- caffe2/serialize/versions.h | 2 +- torch/csrc/jit/mobile/backport_manager.cpp | 157 +++++++++++++++++- torch/csrc/jit/serialization/export.h | 20 +++ .../csrc/jit/serialization/export_module.cpp | 38 +++-- 4 files changed, 193 insertions(+), 24 deletions(-) diff --git a/caffe2/serialize/versions.h b/caffe2/serialize/versions.h index eee98e4fa9cd3..61c8c46666e67 100644 --- a/caffe2/serialize/versions.h +++ b/caffe2/serialize/versions.h @@ -74,7 +74,7 @@ constexpr uint64_t kProducedFileFormatVersion = 0x3L; // 0x6L: Implicit opereator versioning using number of specified argument. // Refer to the summary of https://github.com/pytorch/pytorch/pull/56845 // for details. -constexpr uint64_t kProducedBytecodeVersion = 0x5L; +constexpr uint64_t kProducedBytecodeVersion = 0x6L; static_assert(kProducedBytecodeVersion >= kProducedFileFormatVersion, "kProducedBytecodeVersion must be higher or equal to kProducedFileFormatVersion."); diff --git a/torch/csrc/jit/mobile/backport_manager.cpp b/torch/csrc/jit/mobile/backport_manager.cpp index 37eb4b781c38f..87233ff1e5625 100644 --- a/torch/csrc/jit/mobile/backport_manager.cpp +++ b/torch/csrc/jit/mobile/backport_manager.cpp @@ -25,6 +25,7 @@ using caffe2::serialize::ReadAdapterInterface; namespace { constexpr int64_t kBytecodeVersionV4 = 0x4L; constexpr int64_t kBytecodeVersionV5 = 0x5L; +constexpr int64_t kBytecodeVersionV6 = 0x6L; } // namespace // Utility function that can be reused by backport_vn_to_vn-1(). If any utility @@ -61,15 +62,7 @@ void selective_copy( // constants.pkl // bytecode.pkl // version - bool skip = false; - - // Skip files (exaxt path) - for (const auto& excluded_file : excluded_files) { - if (record == excluded_file) { - skip = true; - break; - } - } + bool skip = excluded_files.count(record) > 0; // Skip dirs, find the last '/' and compare it with record for (const auto& excluded_dir : excluded_dirs) { @@ -211,6 +204,151 @@ std::stringstream backport_v5_to_v4(std::stringstream& input_model_stream) { return ouput_model_stream; } +void writeArchiveV5( + PyTorchStreamWriter& writer, + const IValue& value, + const std::string& archive_name, + const std::string& archive_dir, + const std::string& tensor_dir, + bool tensor_cdata_naming_scheme, + StorageContext& storage_context) { + std::vector data; + // Vector to capture the run-time class types during pickling the IValues + std::vector memoizedClassTypes; + std::vector tensor_names; + Pickler data_pickle( + [&](const char* buf, size_t size) { + data.insert(data.end(), buf, buf + size); + }, + nullptr, + nullptr, + &memoizedClassTypes, + [&](const at::Tensor& tensor) { + // returns a string to use in picker.cpp as storage obj key + if (tensor_cdata_naming_scheme) { + std::string string_id = + std::to_string(reinterpret_cast( + tensor.storage().unsafeGetStorageImpl())); + tensor_names.push_back(string_id + ".storage"); + storage_context.addStorage(string_id, tensor.storage()); + } else { + tensor_names.push_back(std::to_string(tensor_names.size())); + } + return tensor_names.back(); + }); + data_pickle.protocol(); + data_pickle.pushIValue(value); + data_pickle.stop(); + // write out tensor data + size_t i = 0; + std::string prefix = archive_name + "/"; + + TORCH_INTERNAL_ASSERT(tensor_names.size() == data_pickle.tensorData().size()); + const std::vector& pre_serialized_files = + writer.getAllWrittenRecords(); + + for (const auto& td : data_pickle.tensorData()) { + WriteableTensorData writable_td = getWriteableTensorData(td); + std::string fname = tensor_dir + tensor_names[i++]; + if (tensor_cdata_naming_scheme && + std::find( + pre_serialized_files.begin(), pre_serialized_files.end(), fname) != + pre_serialized_files.end()) { + // storage has been serialzed already, skip + continue; + } + writer.writeRecord(fname, writable_td.data(), writable_td.sizeInBytes()); + } + + std::string fname = archive_dir + archive_name + ".pkl"; + writer.writeRecord(fname, data.data(), data.size()); +} + +std::stringstream backport_v6_to_v5(std::stringstream& input_model_stream) { + std::shared_ptr rai = + std::make_shared(&input_model_stream); + auto reader = std::make_shared(rai); + std::vector constants_values = + readArchive(kArchiveNameConstants, *reader.get()).toTuple()->elements(); + + // If there are debug info files in the original model file, it should also + // show up in the backported model + bool hasBytecodeDebug = reader->hasRecord("mobile_debug_handles.pkl"); + + // extra_files are kept + auto records = reader->getAllRecords(); + ExtraFilesMap extra_files; + for (const auto& record : records) { + std::size_t found = record.find_last_of("/\\"); + auto path = record.substr(0, found); + if ("extra" == path) { + extra_files.emplace(record.substr(found + 1), ""); + } + } + // Loading the TS module is required for this backport, because bytecode needs + // to be re-emitted (refer to the comments below) + Module torch_script = torch::jit::load(rai, c10::nullopt, extra_files); + + // The RAII guard to change the flag, emitBytecodeDefaultInputs, to true, so + // that TS stores the default argument values in the constant table, and emits + // the instructions (LOADC, for example), to push the values to the stack. It + // restores the behavior of V5 and before. For V6, the default arg values are + // resolved at runtime init stage for better operator compatibility. + std::stringstream intermediate_model_stream; + { + BytecodeEmitDefaultInputsGuard argNumGuard(true); + torch_script._save_for_mobile( + intermediate_model_stream, extra_files, hasBytecodeDebug); + } + + // Update the bytecode version (from 6 to 5) + + PyTorchStreamReader reader_bytecode(&intermediate_model_stream); + std::vector bytecode_values = get_bytecode_ivalues(reader_bytecode); + std::unordered_set excluded_files{ + "constants.pkl", + "bytecode.pkl", + "version", + }; + + std::unordered_set excluded_dirs{ + "constants", + "bytecode", + }; + + std::stringstream ouput_model_stream; + auto writer_func = [&](const void* buf, size_t nbytes) -> size_t { + ouput_model_stream.write(static_cast(buf), nbytes); + return !ouput_model_stream ? 0 : nbytes; + }; + + PyTorchStreamWriter writer_bytecode(writer_func); + + selective_copy( + reader_bytecode, writer_bytecode, excluded_files, excluded_dirs); + + update_bytecode_version(bytecode_values, kBytecodeVersionV5); + auto bytecode_tuple = c10::ivalue::Tuple::create(std::move(bytecode_values)); + StorageContext storage_context; + writeArchiveV5( + writer_bytecode, + c10::ivalue::Tuple::create(constants_values), + /*archive_name=*/"constants", + /*archive_dir=*/"", + /*tensor_dir=*/"constants/", + /*tensor_cdata_naming_scheme=*/true, + storage_context); + writeArchiveV5( + writer_bytecode, + bytecode_tuple, + /*archive_name=*/"bytecode", + /*archive_dir=*/"", + /*tensor_dir=*/"constants/", + /*tensor_cdata_naming_scheme=*/true, + storage_context); + + return ouput_model_stream; +} } // namespace // A generic contract for backport logic to the previous bytecode version. @@ -223,6 +361,7 @@ using BytecodeBackportFunction = BackportManager::BackportManager() { registerBytecodeBackportFunction(kBytecodeVersionV5, backport_v5_to_v4); + registerBytecodeBackportFunction(kBytecodeVersionV6, backport_v6_to_v5); } std::unordered_map< diff --git a/torch/csrc/jit/serialization/export.h b/torch/csrc/jit/serialization/export.h index 5fb898ca55cf1..f4f4da3b4dd87 100644 --- a/torch/csrc/jit/serialization/export.h +++ b/torch/csrc/jit/serialization/export.h @@ -193,5 +193,25 @@ TORCH_API void SetExportModuleMobileInfoConverter( */ TORCH_API std::vector export_opnames(const Module& m); +struct TORCH_API BytecodeEmitDefaultValueForUnspecifiedArgMode { + static bool is_enabled(); + static void set_enabled(bool enabled); +}; + +// RAII guard to switch the way JIT emits the bytecode for inputs. +// true: instruction of default argument values (like LOADC) is emitted. +// false: instruction of default argument values are not emitted. Instead +// they are fetched from operator schema. +struct TORCH_API BytecodeEmitDefaultInputsGuard { + BytecodeEmitDefaultInputsGuard(bool enable) + : prev_mode(BytecodeEmitDefaultValueForUnspecifiedArgMode::is_enabled()) { + BytecodeEmitDefaultValueForUnspecifiedArgMode::set_enabled(enable); + } + ~BytecodeEmitDefaultInputsGuard() { + BytecodeEmitDefaultValueForUnspecifiedArgMode::set_enabled(prev_mode); + } + bool prev_mode; +}; + } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/serialization/export_module.cpp b/torch/csrc/jit/serialization/export_module.cpp index 68b7f20af4047..52155cac38581 100644 --- a/torch/csrc/jit/serialization/export_module.cpp +++ b/torch/csrc/jit/serialization/export_module.cpp @@ -71,13 +71,11 @@ std::pair getFunctionTuple( Inline(*graph); std::shared_ptr code; - if (caffe2::serialize::kProducedBytecodeVersion == 6) { - code = std::make_shared( - graph, func.name(), false /* emit_default_input_instructions */); - } else { - code = std::make_shared( - graph, func.name(), true /* emit_default_input_instructions */); - } + code = std::make_shared( + graph, + func.name(), + BytecodeEmitDefaultValueForUnspecifiedArgMode:: + is_enabled() /* emit_default_input_instructions */); auto instructions_copy = code->instructions(); // operator names @@ -173,11 +171,11 @@ std::pair getFunctionTuple( if (it != op_to_specified_args.end()) { num_args = it->second; } - if (caffe2::serialize::kProducedBytecodeVersion == 6) { + if (BytecodeEmitDefaultValueForUnspecifiedArgMode::is_enabled()) { + operators.emplace_back(Tup({opname.name, opname.overload_name})); + } else { operators.emplace_back( Tup({opname.name, opname.overload_name, num_args})); - } else { - operators.emplace_back(Tup({opname.name, opname.overload_name})); } } @@ -630,12 +628,12 @@ void ScriptModuleSerializer::writeByteCode( const bool save_mobile_debug_info) { std::vector elements; BackendDebugInfoRecorder debug_info_recorder; - elements.emplace_back( - static_cast(caffe2::serialize::kProducedBytecodeVersion)); + int64_t version_to_write = caffe2::serialize::kProducedBytecodeVersion; + + elements.emplace_back(static_cast(version_to_write)); std::vector debug_info_elements; // Always save debug handles - debug_info_elements.emplace_back( - static_cast(caffe2::serialize::kProducedBytecodeVersion)); + debug_info_elements.emplace_back(static_cast(version_to_write)); moduleMethodsTuple( module, elements, debug_info_elements, debug_info_recorder); @@ -850,5 +848,17 @@ std::vector export_opnames(const script::Module& m) { return std::vector(names.begin(), names.end()); } +// Thread local flag (only happens in export, i.e. on server side) +// to control if instructions for bytecode default inputs are emitted +// or not. It's the major difference between bytecode v5 and v6. +thread_local bool emitBytecodeDefaultInputs = + caffe2::serialize::kProducedBytecodeVersion <= 5 ? true : false; +bool BytecodeEmitDefaultValueForUnspecifiedArgMode::is_enabled() { + return emitBytecodeDefaultInputs; +} +void BytecodeEmitDefaultValueForUnspecifiedArgMode::set_enabled(bool enabled) { + emitBytecodeDefaultInputs = enabled; +} + } // namespace jit } // namespace torch From 83ba71aa0e8ad1a9b47fd3a340a67597aa8540df Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 03:25:29 -0700 Subject: [PATCH 101/305] Make CUDA serde support for TP agent pluggable (#59376) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59376 This is an experiment. The end goal is to separate the CUDA-specific aspects of the TensorPipe agent so that they can be plugged "on top" of the CPU-only parts. This will then allow to move the TP agent to libtorch (because libtorch is split into a CPU and a CUDA part; now it's in libtorch_python), although unfortunately other conditions need to also be met for this to happen. The only instance where we had CPU and CUDA logic within the same code, guarded by `#ifdef USE_CUDA`, is the serialization/deserialization code. I'm thus introducing a sort-of registry in order to "decentralize it". It's not a c10::Registry, because that's overkill (it uses an unordered_map, with strings as keys): here we can just use an array with integers as "keys". ghstack-source-id: 131326167 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28796428 fbshipit-source-id: b52df832e0c0abf489a9e418353103496382ea41 --- .../csrc/distributed/rpc/tensorpipe_utils.cpp | 258 ++++++++++++------ torch/csrc/distributed/rpc/tensorpipe_utils.h | 48 ++++ 2 files changed, 222 insertions(+), 84 deletions(-) diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp index 109b16ee0dd4d..085d55e592e62 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp @@ -39,6 +39,131 @@ inline c10::Device indexToDevice(c10::DeviceIndex index) { } } +class TensorpipeCpuConverter : public TensorpipeDeviceTypeConverter { + public: + c10::optional> prepareTensorForSending( + const c10::Storage& storage, + const std::vector& /* streams */, + tensorpipe::Message& message) const override { + // Enforce memory copy if tensor is created from torch::from_blob, means + // that the tensor doesn't own the memory. + bool storageHasDeleter = storage.data_ptr().get_context() != nullptr; + if (!storageHasDeleter) { + std::vector storageData( + storage.data(), storage.data() + storage.nbytes()); + + tensorpipe::CpuBuffer buffer; + buffer.ptr = storageData.data(); + + tensorpipe::Message::Tensor tensor; + tensor.buffer = buffer; + tensor.length = storageData.size(); + + message.tensors.push_back(std::move(tensor)); + + return c10::make_optional(std::move(storageData)); + } else { + tensorpipe::CpuBuffer buffer; + buffer.ptr = storage.data(); + + tensorpipe::Message::Tensor tensor; + tensor.buffer = buffer; + tensor.length = storage.nbytes(); + + message.tensors.push_back(std::move(tensor)); + + return c10::nullopt; + } + } + + at::DataPtr allocateTensorForReceiving( + int /* deviceIndex */, + size_t length, + const std::vector& /* streams */, + tensorpipe::Allocation& allocation) const override { + at::DataPtr dataPtr = at::getCPUAllocator()->allocate(length); + + tensorpipe::CpuBuffer buffer; + buffer.ptr = dataPtr.get(); + + tensorpipe::Allocation::Tensor tensor; + tensor.buffer = buffer; + + allocation.tensors.push_back(std::move(tensor)); + + return dataPtr; + } +}; + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CPU, TensorpipeCpuConverter); + +#ifdef USE_CUDA_NOT_ROCM +class TensorpipeCudaConverter : public TensorpipeDeviceTypeConverter { + public: + c10::optional> prepareTensorForSending( + const c10::Storage& storage, + const std::vector& streams, + tensorpipe::Message& message) const override { + auto stream = + at::cuda::CUDAStream(getStreamForDevice(streams, storage.device())); + // record tensor data ptrs on TensorPipe streams, so that the tensors + // won't be destructed before TensorPipe finishing sending them. + c10::cuda::CUDACachingAllocator::recordStream(storage.data_ptr(), stream); + + tensorpipe::CudaBuffer buffer; + buffer.ptr = storage.data(); + buffer.stream = stream.stream(); + + tensorpipe::Message::Tensor tensor; + tensor.buffer = buffer; + tensor.length = storage.nbytes(); + + message.tensors.push_back(std::move(tensor)); + + return c10::nullopt; + } + + at::DataPtr allocateTensorForReceiving( + int deviceIndex, + size_t length, + const std::vector& streams, + tensorpipe::Allocation& allocation) const override { + c10::Device device(c10::kCUDA, deviceIndex); + at::cuda::CUDAStream stream(getStreamForDevice(streams, device)); + // CUDACachingAllocator will call recordStream accordingly on the current + // stream. + at::cuda::CUDAStreamGuard guard(stream); + at::DataPtr dataPtr = + c10::cuda::CUDACachingAllocator::get()->allocate(length); + + tensorpipe::CudaBuffer buffer; + buffer.ptr = dataPtr.get(); + buffer.stream = stream.stream(); + + tensorpipe::Allocation::Tensor tensor; + tensor.buffer = buffer; + + allocation.tensors.push_back(tensor); + + return dataPtr; + } +}; + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CUDA, TensorpipeCudaConverter); +#endif + +c10::DeviceType convertDeviceType(const std::string& tpDeviceType) { + if (tpDeviceType == tensorpipe::kCpuDeviceType) { + return c10::kCPU; + } else if (tpDeviceType == tensorpipe::kCudaDeviceType) { + return c10::kCUDA; + } else { + TORCH_INTERNAL_ASSERT(false, "Unrecognized TensorPipe buffer type."); + } +} + } // namespace // As the vector of streams will typically be very small (1-8 items) we expect @@ -54,6 +179,18 @@ const c10::Stream& getStreamForDevice( TORCH_INTERNAL_ASSERT(false, "No stream found for device ", device); } +std::array< + std::atomic, + static_cast(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> + // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) + device_type_converter_registry; + +TensorpipeDeviceTypeConverterRegistrar::TensorpipeDeviceTypeConverterRegistrar( + DeviceType type, + const TensorpipeDeviceTypeConverter* impl) { + device_type_converter_registry[static_cast(type)].store(impl); +} + std::tuple tensorpipeSerialize( c10::intrusive_ptr rpcMessage, std::vector devices, @@ -103,71 +240,31 @@ std::tuple tensorpipeSerialize( // kTpMessagePickleIdx = 3 tpMessage.payloads.push_back(tensorpipe::Message::Payload{ buffers.pickle.data(), buffers.pickle.size()}); - const auto& tensorDataVec = pickler.tensorData(); + const std::vector& tensorDataVec = pickler.tensorData(); + tpMessage.tensors.reserve(tensorDataVec.size()); for (const auto i : c10::irange(tensorDataVec.size())) { - // This is different from jit::getWriteableTensorData as it avoids copying - // tensor to CPU. - const auto& tensorData = - jit::getWriteableTensorData(tensorDataVec[i], /* toCpu */ false); + const torch::Tensor& tensor = tensorDataVec[i]; + + const TensorpipeDeviceTypeConverter* converter = + getDeviceTypeConverter(tensor.device().type()); + TORCH_CHECK( + converter != nullptr, + "Attempting to send a Tensor with unexpected device type ", + tensor.device()); + + TORCH_INTERNAL_ASSERT(tpMessage.tensors.size() == i); + c10::optional> maybeCopiedTensor = + converter->prepareTensorForSending( + tensor.storage(), streams, tpMessage); + TORCH_INTERNAL_ASSERT(tpMessage.tensors.size() == i + 1); + tensorpipe::Device targetDevice = devices.empty() || devices[i].is_cpu() ? tensorpipe::Device{tensorpipe::kCpuDeviceType, 0} : tensorpipe::Device{tensorpipe::kCudaDeviceType, devices[i].index()}; + tpMessage.tensors.back().targetDevice = std::move(targetDevice); - // Enforce memory copy if tensor is created from torch::from_blob, means - // that the tensor doesn't own the memory. - if (!tensorData.storageHasDeleter()) { - std::vector storageData( - tensorData.data(), tensorData.data() + tensorData.sizeInBytes()); - tensorpipe::CpuBuffer buffer; - buffer.ptr = storageData.data(); - - tensorpipe::Message::Tensor tensor; - tensor.buffer = buffer; - tensor.length = storageData.size(); - tensor.targetDevice = std::move(targetDevice); - - tpMessage.tensors.push_back(std::move(tensor)); - buffers.copiedTensors.push_back(std::move(storageData)); - } else { - // TensorPipe uses the same Message class for both reading and writing, so - // it uses non-const ptrs even though it doesn't modify them when writing. - // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) - char* tensorPtr = const_cast(tensorData.data()); - if (tensorDataVec[i].device().is_cpu()) { - tensorpipe::CpuBuffer buffer; - buffer.ptr = tensorPtr; - - tensorpipe::Message::Tensor tensor; - tensor.buffer = buffer; - tensor.length = tensorData.sizeInBytes(); - tensor.targetDevice = std::move(targetDevice); - - tpMessage.tensors.push_back(std::move(tensor)); -#ifdef USE_CUDA_NOT_ROCM - } else if (tensorDataVec[i].device().is_cuda()) { - auto stream = at::cuda::CUDAStream( - getStreamForDevice(streams, tensorDataVec[i].device())); - tensorpipe::CudaBuffer buffer; - buffer.ptr = tensorPtr; - buffer.stream = stream.stream(); - - tensorpipe::Message::Tensor tensor; - tensor.buffer = buffer; - tensor.length = tensorData.sizeInBytes(); - tensor.targetDevice = std::move(targetDevice); - - tpMessage.tensors.push_back(std::move(tensor)); - // record tensor data ptrs on TensorPipe streams, so that the tensors - // won't be destructed before TensorPipe finishing sending them. - c10::cuda::CUDACachingAllocator::recordStream( - tensorDataVec[i].storage().data_ptr(), stream); -#endif - } else { - TORCH_CHECK( - false, - "Attempting to send a Tensor with unexpected device type ", - tensorDataVec[i].device()); - } + if (maybeCopiedTensor.has_value()) { + buffers.copiedTensors.push_back(std::move(maybeCopiedTensor).value()); } } @@ -215,34 +312,27 @@ std::pair tensorpipeAllocate( tpAllocation.payloads[kTpMessagePickleIdx].data = buffers.pickle.data(); size_t numTensors = tpDescriptor.tensors.size(); - tpAllocation.tensors.resize(numTensors); + tpAllocation.tensors.reserve(numTensors); for (const auto tensorIdx : c10::irange(numTensors)) { const tensorpipe::Descriptor::Tensor& tensor = tpDescriptor.tensors[tensorIdx]; TORCH_INTERNAL_ASSERT(tensor.targetDevice.has_value()); - if (tensor.targetDevice->type == tensorpipe::kCpuDeviceType) { - buffers.tensors.emplace_back( - at::getCPUAllocator()->allocate(tensor.length)); - tensorpipe::CpuBuffer buffer; - buffer.ptr = buffers.tensors.back().get(); - tpAllocation.tensors[tensorIdx].buffer = buffer; -#ifdef USE_CUDA_NOT_ROCM - } else if (tensor.targetDevice->type == tensorpipe::kCudaDeviceType) { - c10::Device device(c10::kCUDA, tensor.targetDevice->index); - auto stream = at::cuda::CUDAStream(getStreamForDevice(streams, device)); - // CUDACachingAllocator will call recordStream accordingly on the current - // stream. - at::cuda::CUDAStreamGuard guard(stream); - buffers.tensors.emplace_back( - c10::cuda::CUDACachingAllocator::get()->allocate(tensor.length)); - tensorpipe::CudaBuffer buffer; - buffer.ptr = buffers.tensors.back().get(); - buffer.stream = stream.stream(); - tpAllocation.tensors[tensorIdx].buffer = buffer; -#endif - } else { - TORCH_INTERNAL_ASSERT(false, "Unrecognized TensorPipe buffer type."); - } + c10::DeviceType targetDeviceType = + convertDeviceType(tensor.targetDevice->type); + + const TensorpipeDeviceTypeConverter* converter = + getDeviceTypeConverter(targetDeviceType); + TORCH_INTERNAL_ASSERT( + converter != nullptr, + "Attempting to receive a Tensor with unexpected device type ", + targetDeviceType); + + TORCH_INTERNAL_ASSERT(tpAllocation.tensors.size() == tensorIdx); + at::DataPtr dataPtr = converter->allocateTensorForReceiving( + tensor.targetDevice->index, tensor.length, streams, tpAllocation); + TORCH_INTERNAL_ASSERT(tpAllocation.tensors.size() == tensorIdx + 1); + + buffers.tensors.push_back(std::move(dataPtr)); } return {std::move(tpAllocation), std::move(buffers)}; diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index 5f3d356330751..7c6549115fb45 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -23,6 +23,54 @@ const c10::Stream& getStreamForDevice( const std::vector& streams, const c10::Device& device); +// Inspired by c10/core/impl/DeviceGuardImplInterface.h. + +class TensorpipeDeviceTypeConverter { + public: + // Ideally we'd want this to also return a tensorpipe::Message::Tensor object + // but we cannot forward-declare that class (because it's nested), and we + // cannot include the TensorPipe headers because it's a private dependency. + // Thus we bend over backwards and entrust this method with appending that + // object to the `tensors` field of the tensorpipe::Message object we pass. + virtual c10::optional> prepareTensorForSending( + const c10::Storage& storage, + const std::vector& streams, + tensorpipe::Message& message) const = 0; + + // Same as above: this method cannot return a tensorpipe::Allocation::Tensor, + // thus it appends it to the `tensors` field of the tensorpipe::Allocation. + virtual at::DataPtr allocateTensorForReceiving( + int deviceIndex, + size_t length, + const std::vector& streams, + tensorpipe::Allocation& allocation) const = 0; + + virtual ~TensorpipeDeviceTypeConverter() = default; +}; + +extern C10_API std::array< + std::atomic, + static_cast(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> + device_type_converter_registry; + +class C10_API TensorpipeDeviceTypeConverterRegistrar { + public: + TensorpipeDeviceTypeConverterRegistrar( + DeviceType, + const TensorpipeDeviceTypeConverter*); +}; + +#define C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER( \ + DevType, TensorpipeDeviceTypeConverter) \ + static ::torch::distributed::rpc::TensorpipeDeviceTypeConverterRegistrar \ + C10_ANONYMOUS_VARIABLE(g_##DeviceType)( \ + ::c10::DeviceType::DevType, new TensorpipeDeviceTypeConverter()); + +inline const TensorpipeDeviceTypeConverter* getDeviceTypeConverter( + DeviceType type) { + return device_type_converter_registry[static_cast(type)].load(); +} + // A struct that holds pointers that keep alive all the memory that will be // accessed by TensorPipe during a write operation. struct TensorpipeWriteBuffers { From 5e5ca0682bb9a4798820a7e177c7ffe091c7615c Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 03:25:29 -0700 Subject: [PATCH 102/305] Move CUDA-related stuff of TP agent to separate file (#59377) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59377 This PR demonstrates that now the CUDA parts of the TensorPipe agent just "plug on top" of the CPU-only parts. Thus ideally the CPU-only parts could go in libtorch while the CUDA-only parts could go in libtorch_cuda. Unfortunately we can't do that just yet, because the TensorPipe agent depends on c10d (for its Store and its ProcessGroup), which lives in libtorch_python. ghstack-source-id: 131326168 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D28796429 fbshipit-source-id: 41b2eb8400c0da282f3750a4eea21ad83ee4a175 --- caffe2/CMakeLists.txt | 1 + tools/build_variables.bzl | 1 + .../csrc/distributed/rpc/tensorpipe_agent.cpp | 99 +------------ torch/csrc/distributed/rpc/tensorpipe_agent.h | 28 +++- .../csrc/distributed/rpc/tensorpipe_cuda.cpp | 136 ++++++++++++++++++ .../csrc/distributed/rpc/tensorpipe_utils.cpp | 61 -------- torch/csrc/distributed/rpc/tensorpipe_utils.h | 4 - 7 files changed, 163 insertions(+), 167 deletions(-) create mode 100644 torch/csrc/distributed/rpc/tensorpipe_cuda.cpp diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 20ef2a6376c0c..174018456efd8 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -367,6 +367,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) "${TORCH_SRC_DIR}/csrc/distributed/rpc/macros.h" "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.cpp" "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.h" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_cuda.cpp" "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp" "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h" ) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index d7f2d8bf95ab3..b62aa38db2190 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -718,6 +718,7 @@ libtorch_python_distributed_sources = libtorch_python_distributed_core_sources + "torch/csrc/distributed/rpc/python_rpc_handler.cpp", "torch/csrc/distributed/rpc/request_callback_impl.cpp", "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", + "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/testing/init.cpp", diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp index ea88629f0b55d..0f6645cdcd5d5 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp @@ -136,9 +136,9 @@ std::vector getDevicesOfTensors( } std::vector devices; devices.reserve(deviceCount); - for (c10::DeviceIndex idx = 0; idx < indexBitset.size(); idx++) { + for (const auto idx : c10::irange(indexBitset.size())) { if (indexBitset[idx]) { - devices.emplace_back(impl->type(), idx); + devices.emplace_back(impl->type(), static_cast(idx)); } } return devices; @@ -195,36 +195,6 @@ const std::string& TensorPipeAgent::guessAddress() { namespace { -// These priorities instruct TensorPipe on which transport/channel to pick -// during handshake. Higher priorities will take precedence over lower ones. -// The transport with lowest priority will be the one used to bootstrap pipes. - -constexpr int64_t kShmTransportPriority = 200; -constexpr int64_t kIbvTransportPriority = 100; -// The UV transport just uses TCP and should work everywhere, thus keep it last. -constexpr int64_t kUvTransportPriority = 0; - -constexpr int64_t kCmaChannelPriority = 1200; -constexpr int64_t kMultiplexedUvChannelPriority = 1100; -// The basic channel reuses a transport as a channel, and is thus our fallback. -constexpr int64_t kBasicChannelPriority = 1000; - -// CPU channel have higher priority than CUDA channels, since the latter might -// handle CPU-to-CPU transfers, but will always be less efficient than their -// CPU-only counterparts. -#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL && defined(USE_CUDA_NOT_ROCM) -constexpr int64_t kCudaIpcChannelPriority = 300; -#endif - -#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM) -constexpr int64_t kCudaGdrChannelPriority = 200; -#endif - -#ifdef USE_CUDA_NOT_ROCM -constexpr int64_t kCudaXthChannelPriority = 400; -constexpr int64_t kCudaBasicChannelPriority = 0; -#endif - std::unique_ptr makeUvTransport() { auto context = tensorpipe::transport::uv::create(); std::string address = TensorPipeAgent::guessAddress(); @@ -319,7 +289,7 @@ constexpr static int kNumUvThreads = 16; std::unique_ptr makeMultiplexedUvChannel() { std::vector> contexts; std::vector> listeners; - for (const auto laneIdx : c10::irange(kNumUvThreads)) { + for (const auto laneIdx C10_UNUSED : c10::irange(kNumUvThreads)) { auto context = tensorpipe::transport::uv::create(); std::string address = TensorPipeAgent::guessAddress(); contexts.push_back(std::move(context)); @@ -343,69 +313,6 @@ C10_REGISTER_CREATOR( mpt_uv, makeMultiplexedUvChannel); -#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL && defined(USE_CUDA_NOT_ROCM) - -std::unique_ptr makeCudaIpcChannel() { - auto context = tensorpipe::channel::cuda_ipc::create(); - return std::make_unique( - ChannelRegistration{std::move(context), kCudaIpcChannelPriority}); -} - -// The cuda_ipc channels use cudaMemcpy to transmit CUDA tensor across processes -// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_ipc, makeCudaIpcChannel); - -#endif - -#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL && defined(USE_CUDA_NOT_ROCM) - -std::unique_ptr makeCudaGdrChannel() { - auto context = tensorpipe::channel::cuda_gdr::create(); - return std::make_unique( - ChannelRegistration{std::move(context), kCudaGdrChannelPriority}); -} - -// The cuda_gdr channel sends CUDA memory over InfiniBand using GPUDirect RDMA. -// It directly registers the user-provided tensor with libibverbs, an operation -// which is expensive the first time, but it then caches the registration in -// order to amortize the cost and get low latency for subsequent transfers. A -// ready-to-send/ready-to-receive handshake is still needed before the transfer -// in order to ensure readiness and to agree on the device indices and thus the -// queue pair to use. It automatically pairs each GPU to the "closest" NIC if -// there are multiple of them (closest = longest prefix match in PCI tree). -// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_gdr, makeCudaGdrChannel); - -#endif - -#ifdef USE_CUDA_NOT_ROCM - -std::unique_ptr makeCudaXthChannel() { - auto context = tensorpipe::channel::cuda_xth::create(); - return std::make_unique( - ChannelRegistration{std::move(context), kCudaXthChannelPriority}); -} - -// The cuda_xth channel supports same-process GPU-to-GPU comm -// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_xth, makeCudaXthChannel); - -std::unique_ptr makeCudaBasicChannel() { - auto context = tensorpipe::channel::cuda_basic::create( - tensorpipe::channel::basic::create()); - return std::make_unique( - ChannelRegistration{std::move(context), kCudaBasicChannelPriority}); -} - -// The cuda_basic is the fallback channel for GPU-to-GPU comm -// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -C10_REGISTER_CREATOR( - TensorPipeChannelRegistry, - cuda_basic, - makeCudaBasicChannel); - -#endif - } // namespace ////////////////////////// MetricsTracker ///////////////////////////////// diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index 2d2a238e8a9e7..df3328793fa11 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -17,12 +17,6 @@ namespace tensorpipe { -struct CpuBuffer; - -#ifdef USE_CUDA_NOT_ROCM -struct CudaBuffer; -#endif - class Context; class Error; class Listener; @@ -43,6 +37,28 @@ namespace torch { namespace distributed { namespace rpc { +// These priorities instruct TensorPipe on which transport/channel to pick +// during handshake. Higher priorities will take precedence over lower ones. +// The transport with lowest priority will be the one used to bootstrap pipes. + +constexpr int64_t kShmTransportPriority = 200; +constexpr int64_t kIbvTransportPriority = 100; +// The UV transport just uses TCP and should work everywhere, thus keep it last. +constexpr int64_t kUvTransportPriority = 0; + +constexpr int64_t kCmaChannelPriority = 1200; +constexpr int64_t kMultiplexedUvChannelPriority = 1100; +// The basic channel reuses a transport as a channel, and is thus our fallback. +constexpr int64_t kBasicChannelPriority = 1000; + +// CPU channel have higher priority than CUDA channels, since the latter might +// handle CPU-to-CPU transfers, but will always be less efficient than their +// CPU-only counterparts. +constexpr int64_t kCudaIpcChannelPriority = 300; +constexpr int64_t kCudaGdrChannelPriority = 200; +constexpr int64_t kCudaXthChannelPriority = 400; +constexpr int64_t kCudaBasicChannelPriority = 0; + using steady_clock_time_point = std::chrono::time_point; diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp new file mode 100644 index 0000000000000..ebae1942de8ed --- /dev/null +++ b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp @@ -0,0 +1,136 @@ +#include +#include +#include + +#if defined(USE_TENSORPIPE) && defined(USE_CUDA_NOT_ROCM) + +#include +#include +#include + +#include + +namespace torch { +namespace distributed { +namespace rpc { +namespace { + +#if TENSORPIPE_HAS_CUDA_IPC_CHANNEL + +std::unique_ptr makeCudaIpcChannel() { + auto context = tensorpipe::channel::cuda_ipc::create(); + return std::make_unique( + ChannelRegistration{std::move(context), kCudaIpcChannelPriority}); +} + +// The cuda_ipc channels use cudaMemcpy to transmit CUDA tensor across processes +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_ipc, makeCudaIpcChannel); + +#endif + +#if TENSORPIPE_HAS_CUDA_GDR_CHANNEL + +std::unique_ptr makeCudaGdrChannel() { + auto context = tensorpipe::channel::cuda_gdr::create(); + return std::make_unique( + ChannelRegistration{std::move(context), kCudaGdrChannelPriority}); +} + +// The cuda_gdr channel sends CUDA memory over InfiniBand using GPUDirect RDMA. +// It directly registers the user-provided tensor with libibverbs, an operation +// which is expensive the first time, but it then caches the registration in +// order to amortize the cost and get low latency for subsequent transfers. A +// ready-to-send/ready-to-receive handshake is still needed before the transfer +// in order to ensure readiness and to agree on the device indices and thus the +// queue pair to use. It automatically pairs each GPU to the "closest" NIC if +// there are multiple of them (closest = longest prefix match in PCI tree). +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_gdr, makeCudaGdrChannel); + +#endif + +std::unique_ptr makeCudaXthChannel() { + auto context = tensorpipe::channel::cuda_xth::create(); + return std::make_unique( + ChannelRegistration{std::move(context), kCudaXthChannelPriority}); +} + +// The cuda_xth channel supports same-process GPU-to-GPU comm +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_CREATOR(TensorPipeChannelRegistry, cuda_xth, makeCudaXthChannel); + +std::unique_ptr makeCudaBasicChannel() { + auto context = tensorpipe::channel::cuda_basic::create( + tensorpipe::channel::basic::create()); + return std::make_unique( + ChannelRegistration{std::move(context), kCudaBasicChannelPriority}); +} + +// The cuda_basic is the fallback channel for GPU-to-GPU comm +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_CREATOR( + TensorPipeChannelRegistry, + cuda_basic, + makeCudaBasicChannel); + +class TensorpipeCudaConverter : public TensorpipeDeviceTypeConverter { + public: + c10::optional> prepareTensorForSending( + const c10::Storage& storage, + const std::vector& streams, + tensorpipe::Message& message) const override { + auto stream = + at::cuda::CUDAStream(getStreamForDevice(streams, storage.device())); + // record tensor data ptrs on TensorPipe streams, so that the tensors + // won't be destructed before TensorPipe finishing sending them. + c10::cuda::CUDACachingAllocator::recordStream(storage.data_ptr(), stream); + + tensorpipe::CudaBuffer buffer; + buffer.ptr = storage.data(); + buffer.stream = stream.stream(); + + tensorpipe::Message::Tensor tensor; + tensor.buffer = buffer; + tensor.length = storage.nbytes(); + + message.tensors.push_back(std::move(tensor)); + + return c10::nullopt; + } + + at::DataPtr allocateTensorForReceiving( + int deviceIndex, + size_t length, + const std::vector& streams, + tensorpipe::Allocation& allocation) const override { + c10::Device device(c10::kCUDA, deviceIndex); + at::cuda::CUDAStream stream(getStreamForDevice(streams, device)); + // CUDACachingAllocator will call recordStream accordingly on the current + // stream. + at::cuda::CUDAStreamGuard guard(stream); + at::DataPtr dataPtr = + c10::cuda::CUDACachingAllocator::get()->allocate(length); + + tensorpipe::CudaBuffer buffer; + buffer.ptr = dataPtr.get(); + buffer.stream = stream.stream(); + + tensorpipe::Allocation::Tensor tensor; + tensor.buffer = buffer; + + allocation.tensors.push_back(tensor); + + return dataPtr; + } +}; + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CUDA, TensorpipeCudaConverter); + +} // namespace +} // namespace rpc +} // namespace distributed +} // namespace torch + +#endif diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp index 085d55e592e62..55b8554f66d28 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp @@ -3,12 +3,7 @@ #ifdef USE_TENSORPIPE -#ifdef USE_CUDA_NOT_ROCM -#include -#include -#include #include -#endif #include @@ -98,62 +93,6 @@ class TensorpipeCpuConverter : public TensorpipeDeviceTypeConverter { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CPU, TensorpipeCpuConverter); -#ifdef USE_CUDA_NOT_ROCM -class TensorpipeCudaConverter : public TensorpipeDeviceTypeConverter { - public: - c10::optional> prepareTensorForSending( - const c10::Storage& storage, - const std::vector& streams, - tensorpipe::Message& message) const override { - auto stream = - at::cuda::CUDAStream(getStreamForDevice(streams, storage.device())); - // record tensor data ptrs on TensorPipe streams, so that the tensors - // won't be destructed before TensorPipe finishing sending them. - c10::cuda::CUDACachingAllocator::recordStream(storage.data_ptr(), stream); - - tensorpipe::CudaBuffer buffer; - buffer.ptr = storage.data(); - buffer.stream = stream.stream(); - - tensorpipe::Message::Tensor tensor; - tensor.buffer = buffer; - tensor.length = storage.nbytes(); - - message.tensors.push_back(std::move(tensor)); - - return c10::nullopt; - } - - at::DataPtr allocateTensorForReceiving( - int deviceIndex, - size_t length, - const std::vector& streams, - tensorpipe::Allocation& allocation) const override { - c10::Device device(c10::kCUDA, deviceIndex); - at::cuda::CUDAStream stream(getStreamForDevice(streams, device)); - // CUDACachingAllocator will call recordStream accordingly on the current - // stream. - at::cuda::CUDAStreamGuard guard(stream); - at::DataPtr dataPtr = - c10::cuda::CUDACachingAllocator::get()->allocate(length); - - tensorpipe::CudaBuffer buffer; - buffer.ptr = dataPtr.get(); - buffer.stream = stream.stream(); - - tensorpipe::Allocation::Tensor tensor; - tensor.buffer = buffer; - - allocation.tensors.push_back(tensor); - - return dataPtr; - } -}; - -// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -C10_REGISTER_TENSORPIPE_DEVICE_TYPE_CONVERTER(CUDA, TensorpipeCudaConverter); -#endif - c10::DeviceType convertDeviceType(const std::string& tpDeviceType) { if (tpDeviceType == tensorpipe::kCpuDeviceType) { return c10::kCPU; diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index 7c6549115fb45..3f41b351c9898 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -5,10 +5,6 @@ #include #include -#ifdef USE_CUDA_NOT_ROCM -#include -#endif - namespace tensorpipe { class Message; class Allocation; From f4f795081215944c72fce26f83cc9757da18837b Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 03:25:29 -0700 Subject: [PATCH 103/305] Prepare for TensorPipe separating its CUDA-specific headers (#59788) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59788 This one line is all we need to "migrate" PyTorch to the "new API" of TensorPipe that splits the CUDA-specific stuff in a separate top-level header. (The motivation behind that is that it will allow us to "stack" the CUDA code on top of the CPU one). ghstack-source-id: 131326166 Test Plan: None yet Reviewed By: beauby Differential Revision: D28875277 fbshipit-source-id: ecfd0b7fc0218ab7899bfe64ffe73c1417b897db --- torch/csrc/distributed/rpc/tensorpipe_cuda.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp index ebae1942de8ed..9489fcd222bbd 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp @@ -9,6 +9,7 @@ #include #include +#include namespace torch { namespace distributed { From f9445c8a6b606dec9e59a8d8fd4e90a7d4e128d1 Mon Sep 17 00:00:00 2001 From: Serhat Yilmaz Date: Tue, 15 Jun 2021 06:57:15 -0700 Subject: [PATCH 104/305] [torch][segment_reduce] Add cuda support for mean reduction (#59543) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59543 Building on top of previous PR: https://github.com/pytorch/pytorch/pull/59521 This diff is adding support for mean reduction for Cuda (fwd only currently). Will add cuda backward implementation in subsequent PR. Next Steps: cuda backward support for mean 2d data input support more testing benchmarking Test Plan: update unit test to cover this part as well. Reviewed By: ngimel Differential Revision: D28922838 fbshipit-source-id: 72b7e5e79db967116b96ad010f290c9f057232d4 --- aten/src/ATen/native/cuda/SegmentReduce.cu | 91 ++++++++++++++++++---- test/test_segment_reductions.py | 3 - 2 files changed, 77 insertions(+), 17 deletions(-) diff --git a/aten/src/ATen/native/cuda/SegmentReduce.cu b/aten/src/ATen/native/cuda/SegmentReduce.cu index 1728961960573..5b4a8c40634e2 100644 --- a/aten/src/ATen/native/cuda/SegmentReduce.cu +++ b/aten/src/ATen/native/cuda/SegmentReduce.cu @@ -23,6 +23,14 @@ struct CustomMax { } }; +struct CustomSum { + template + __host__ __device__ __forceinline__ OutputT + operator()(const OutputT& a, const OutputT& b) const { + return a + b; + } +}; + Tensor _get_complete_sum(const Tensor& lengths) { int64_t segment_count = lengths.numel(); TORCH_CHECK(segment_count < INT_MAX); @@ -41,6 +49,27 @@ Tensor _get_complete_sum(const Tensor& lengths) { return offsets; } +template +__global__ static void post_sum_div_kernel( + scalar_t* output_data, + const int64_t* lengths_data, + const int64_t segment_count, + bool is_initial_set, + scalar_t initial) { + CUDA_KERNEL_LOOP(index, segment_count) { + CUDA_KERNEL_ASSERT(lengths_data[index] >= 0); + if (lengths_data[index] == 0) { + if (is_initial_set) { + output_data[index] = initial; + } else { + output_data[index] = NAN; + } + } else if (!at::_isnan(output_data[index])) { + output_data[index] = output_data[index] / lengths_data[index]; + } + } +} + Tensor _segment_reduce_cuda_kernel( SegmentReductionType reduction, const Tensor& data, @@ -53,6 +82,12 @@ Tensor _segment_reduce_cuda_kernel( auto offsets = _get_complete_sum(lengths); auto* offsets_data_ptr = offsets.data_ptr(); + constexpr int threads_per_block = 256; + int64_t num_blocks = + (segment_count + threads_per_block - 1) / threads_per_block; + num_blocks = std::max(num_blocks, (int64_t)1); + auto* lengths_data_ptr = lengths.data_ptr(); + AT_DISPATCH_ALL_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, @@ -62,20 +97,48 @@ Tensor _segment_reduce_cuda_kernel( auto* data_data_ptr = data.data_ptr(); auto* output_data_ptr = output.data_ptr(); - CustomMax max_op{}; - scalar_t initial_value = initial.has_value() - ? initial.value().to() - : std::numeric_limits::lowest(); - CUB_WRAPPER( - cub::DeviceSegmentedReduce::Reduce, - data_data_ptr, - output_data_ptr, - segment_count, - offsets_data_ptr, - offsets_data_ptr + 1, - max_op, - initial_value, - at::cuda::getCurrentCUDAStream()); + if (reduction == SegmentReductionType::MAX) { + CustomMax max_op{}; + scalar_t initial_value = initial.has_value() + ? initial.value().to() + : std::numeric_limits::lowest(); + CUB_WRAPPER( + cub::DeviceSegmentedReduce::Reduce, + data_data_ptr, + output_data_ptr, + segment_count, + offsets_data_ptr, + offsets_data_ptr + 1, + max_op, + initial_value, + at::cuda::getCurrentCUDAStream()); + } else if (reduction == SegmentReductionType::MEAN) { + CustomSum sum_op{}; + scalar_t initial_value = initial.has_value() + ? initial.value().to() + : (scalar_t)0; + CUB_WRAPPER( + cub::DeviceSegmentedReduce::Reduce, + data_data_ptr, + output_data_ptr, + segment_count, + offsets_data_ptr, + offsets_data_ptr + 1, + sum_op, + initial_value, + at::cuda::getCurrentCUDAStream()); + + post_sum_div_kernel + <<>>( + output_data_ptr, + lengths_data_ptr, + segment_count, + initial.has_value(), + initial_value); + } }); return output; diff --git a/test/test_segment_reductions.py b/test/test_segment_reductions.py index 7338d462fcd84..d6a23fff63544 100644 --- a/test/test_segment_reductions.py +++ b/test/test_segment_reductions.py @@ -85,9 +85,6 @@ def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) def test_simple_1d(self, device, dtype): for reduction in ("max", "mean"): - # TODO: Remove if once mean reduction for cuda is implemented - if reduction == "mean" and device != "cpu": - continue self._test_simple_1d(reduction, device, dtype, False, 0) self._test_simple_1d(reduction, device, dtype, False, -1) self._test_simple_1d(reduction, device, dtype, True, 0) From 90cf76dde55e304e512e85fc83a767f7b6ad8983 Mon Sep 17 00:00:00 2001 From: nikithamalgi Date: Tue, 15 Jun 2021 07:21:07 -0700 Subject: [PATCH 105/305] Support torch.nn.parameter type for PDT (#59249) Summary: ========= Support torch.nn.parameter type for PDT Pull Request resolved: https://github.com/pytorch/pytorch/pull/59249 Test Plan: ==== with-proxy python test/test_jit.py -k TestPDT Reviewed By: ZolotukhinM Differential Revision: D29124413 Pulled By: nikithamalgifb fbshipit-source-id: b486b82c897dbc2b55fbacd5d610bdb700ddc9fa --- test/jit/test_pdt.py | 17 +++++++++++++++++ torch/jit/_monkeytype_config.py | 6 ++++++ 2 files changed, 23 insertions(+) diff --git a/test/jit/test_pdt.py b/test/jit/test_pdt.py index 4833b98d340c7..d75365d62e748 100644 --- a/test/jit/test_pdt.py +++ b/test/jit/test_pdt.py @@ -405,3 +405,20 @@ def test_model_with_args(a, m): user_class = ClassWithArgs(False) scripted_fn = torch.jit._script_pdt(test_model_with_args, example_inputs=[(10, user_class, ), (10.9, user_class, ), ]) self.assertEqual(scripted_fn(100, ClassWithArgs(True), ), test_model_with_args(100, ClassWithArgs(True))) + + def test_nn_parameter_as_arg(self): + class TestNNParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.inp = torch.nn.Parameter(torch.ones(2, 3)) + + def add_nn_parameter_with_int(self, x, y): + return torch.add(x, y) + + def forward(self, y): + return self.add_nn_parameter_with_int(self.inp, y) + + make_global(TestNNParameter) + pdt_model = TestNNParameter() + scripted_fn = torch.jit._script_pdt(pdt_model, example_inputs={pdt_model: [(10, ), ], }) + self.assertEqual(scripted_fn(20), pdt_model(20)) diff --git a/torch/jit/_monkeytype_config.py b/torch/jit/_monkeytype_config.py index 29ead382cfe32..c124602d0cfd5 100644 --- a/torch/jit/_monkeytype_config.py +++ b/torch/jit/_monkeytype_config.py @@ -1,5 +1,6 @@ import inspect import typing +import torch from typing import Optional, Iterable, List, Dict from collections import defaultdict @@ -72,6 +73,11 @@ def consolidate_types(self, qualified_name: str) -> Dict: if inspect.getmodule(_type) == typing: _type_to_string = str(_type) _all_type += _type_to_string.replace('typing.', '') + ',' + elif _type is torch.nn.parameter.Parameter: + # Check if the type is torch.nn.parameter.Parameter, + # use the entire quaalified name `torch.nn.parameter.Parameter` + # for type + _all_type += 'torch.nn.parameter.Parameter' + ',' else: _all_type += _type.__name__ + ',' _all_type = _all_type.lstrip(" ") # Remove any trailing spaces From 1c502d1f8ec861c31a08d580ae7b73b7fbebebed Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Tue, 15 Jun 2021 07:38:37 -0700 Subject: [PATCH 106/305] Don't run_build when run_binary_tests (#59982) Summary: https://github.com/pytorch/pytorch/issues/59889 wasn't a proper revert of https://github.com/pytorch/pytorch/issues/58778. This PR fixes that. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59982 Reviewed By: seemethere Differential Revision: D29114129 Pulled By: samestep fbshipit-source-id: b40563db6ff1153a5f759639978279f5fcbccaa9 --- .github/pytorch-circleci-labels.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/pytorch-circleci-labels.yml b/.github/pytorch-circleci-labels.yml index 987d52b0c37b2..de990146d027f 100644 --- a/.github/pytorch-circleci-labels.yml +++ b/.github/pytorch-circleci-labels.yml @@ -9,6 +9,8 @@ labels_to_circle_params: - release/.* tags: - v[0-9]+(\.[0-9]+)*-rc[0-9]+ + set_to_false: + - run_build ci/master: parameter: run_master_build ci/slow-gradcheck: From 4afd0b7952c5e5cb5908441704fc63846d053ed8 Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Tue, 15 Jun 2021 10:21:11 -0700 Subject: [PATCH 107/305] .github: Add Windows CUDA 11.1 workflow (#59960) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59960 Adds the CUDA 11.1 workflow to GHA Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: janeyx99 Differential Revision: D29116814 Pulled By: seemethere fbshipit-source-id: 90601610e481e1f70a60eaa1b640373ecb89bdb9 --- .circleci/scripts/windows_cuda_install.sh | 2 +- .github/scripts/generate_ci_workflows.py | 5 + .../pytorch-win-vs2019-cuda11-cudnn8-py3.yml | 188 ++++++++++++++++++ 3 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml diff --git a/.circleci/scripts/windows_cuda_install.sh b/.circleci/scripts/windows_cuda_install.sh index 8f28a4a0a18a1..f2d5eb6dcf7b2 100644 --- a/.circleci/scripts/windows_cuda_install.sh +++ b/.circleci/scripts/windows_cuda_install.sh @@ -25,7 +25,7 @@ else exit 1 fi -if [[ "$cuda_major_version" == "11" && "${JOB_EXECUTOR}" == "windows-with-nvidia-gpu" ]]; then +if [[ "$cuda_major_version" == "11" && "${JOB_EXECUTOR:-}" == "windows-with-nvidia-gpu" ]]; then cuda_install_packages="${cuda_install_packages} Display.Driver" fi diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index d27f593d67137..aae7ef2685fec 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -80,6 +80,11 @@ def generate_workflow_file( build_environment="pytorch-win-vs2019-cuda10-cudnn7-py3", cuda_version="10.1", test_runner_type=WINDOWS_CUDA_TEST_RUNNER, + ), + PyTorchWindowsWorkflow( + build_environment="pytorch-win-vs2019-cuda11-cudnn8-py3", + cuda_version="11.1", + test_runner_type=WINDOWS_CUDA_TEST_RUNNER, ) ] diff --git a/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml new file mode 100644 index 0000000000000..0de71b92f0add --- /dev/null +++ b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml @@ -0,0 +1,188 @@ +# @generated DO NOT EDIT MANUALLY +# Template is at: .github/templates/windows_ci_workflow.yml.j2 +# Generation script: .github/scripts/generate_ci_workflows.py +name: Windows CI (pytorch-win-vs2019-cuda11-cudnn8-py3) + +on: + push: + branches: + - master + - release/* + workflow_dispatch: + +env: + BUILD_ENVIRONMENT: pytorch-win-vs2019-cuda11-cudnn8-py3 + BUILD_WHEEL: 1 + CUDA_VERSION: "11.1" + IN_CI: 1 + INSTALL_WINDOWS_SDK: 1 + JOB_BASE_NAME: test + PYTHON_VERSION: "3.8" + SCCACHE_BUCKET: "ossci-compiler-cache" + VC_PRODUCT: "BuildTools" + VC_VERSION: "" + VC_YEAR: "2019" + TORCH_CUDA_ARCH_LIST: "7.0" + USE_CUDA: 1 + +concurrency: + group: pytorch-win-vs2019-cuda11-cudnn8-py3-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: true + +jobs: + build: + runs-on: "windows.4xlarge" + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf + - name: Install Visual Studio 2019 toolchain + shell: powershell + run: | + .\.circleci\scripts\vs_install.ps1 + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh + - name: Build + shell: bash + run: | + .jenkins/pytorch/win-build.sh + # Upload to github so that people can click and download artifacts + - name: Upload artifacts to Github + if: always() + uses: actions/upload-artifact@v2 + # Don't fail on upload to GH since it's only for user convenience + continue-on-error: true + with: + retention-days: 14 + if-no-files-found: error + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\w\build-results + - name: Upload artifacts to s3 + if: always() + uses: seemethere/upload-artifact-s3@9d7ceb0ab39c2c88d93ef7792b27425b27d59162 + with: + retention-days: 14 + if-no-files-found: error + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\w\build-results + + test: + runs-on: windows.8xlarge.nvidia.gpu + env: + JOB_BASE_NAME: pytorch-win-vs2019-cuda11-cudnn8-py3-test + needs: + - build + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + submodules: recursive + - name: Clean workspace (including things in .gitignore) + shell: bash + run: | + git clean -xdf + - name: Install Visual Studio 2019 toolchain + shell: powershell + run: | + .\.circleci\scripts\vs_install.ps1 + - name: Install Cuda + shell: bash + run: | + .circleci/scripts/windows_cuda_install.sh + - name: Install Cudnn + shell: bash + run: | + .circleci/scripts/windows_cudnn_install.sh + - uses: seemethere/download-artifact-s3@0504774707cbc8603d7dca922e8026eb8bf3b47b + name: Download PyTorch Build Artifacts + with: + name: ${{ env.BUILD_ENVIRONMENT }} + path: C:\${{ github.run_id }}\build-results + - name: Check build-results folder + shell: powershell + run: | + tree /F C:\$Env:GITHUB_RUN_ID\build-results + # Needed for coverage in win-test.sh + - uses: actions/setup-python@v2 + name: Setup Python3 + with: + python-version: '3.x' + - name: Run test scripts + shell: bash + env: + PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ + run: | + .jenkins/pytorch/win-test.sh + - uses: actions/upload-artifact@v2 + name: Store PyTorch Test Reports + if: always() + with: + name: test-reports + retention-days: 14 + if-no-files-found: error + path: + test/**/*.xml + + # this is a separate step from test because the log files from test are too + # long: basically, GitHub tries to render all of the log files when you click + # through an action causing extreme slowdown on actions that contain too many + # logs (like test); we can always move it back to the other one, but it + # doesn't create the best experience + render_test_results: + if: always() + needs: + - test + runs-on: ubuntu-18.04 + # TODO: Make this into a composite step + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + # deep clone, to allow tools/print_test_stats.py to use Git commands + fetch-depth: 0 + - uses: actions/download-artifact@v2 + name: Download PyTorch Test Reports + with: + name: test-reports + path: test/test-reports + - uses: actions/setup-python@v2 + with: + python-version: 3.9 + - name: Install dependencies + # boto3 version copied from .circleci/docker/common/install_conda.sh + run: | + pip install -r requirements.txt + pip install boto3==1.16.34 junitparser rich + - name: Output Test Results (Click Me) + run: | + python tools/render_junit.py test + - name: Parse ref + id: parse-ref + run: .github/scripts/parse_ref.py + - name: Display and upload test statistics (Click Me) + # temporary hack: set CIRCLE_* vars, until we update + # tools/print_test_stats.py to natively support GitHub Actions + env: + SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_OSSCI_METRICS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_OSSCI_METRICS_SECRET_ACCESS_KEY }} + CIRCLE_BRANCH: ${{ steps.parse-ref.outputs.branch }} + CIRCLE_JOB: pytorch-win-vs2019-cuda11-cudnn8-py3 + CIRCLE_PR_NUMBER: ${{ github.event.pull_request.number }} + CIRCLE_SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + CIRCLE_TAG: ${{ steps.parse-ref.outputs.tag }} + CIRCLE_WORKFLOW_ID: ${{ github.run_id }} # dunno if this corresponds + run: | + export PYTHONPATH=$PWD + python tools/print_test_stats.py --upload-to-s3 --compare-with-s3 test From aa163aeff5a9b4919142efaf1aa212a4dc423733 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Tue, 15 Jun 2021 10:34:45 -0700 Subject: [PATCH 108/305] [nnc] Made several LoopNest APIs static (#59494) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59494 Test Plan: Imported from OSS Reviewed By: ZolotukhinM Differential Revision: D28915959 Pulled By: navahgar fbshipit-source-id: bf52e30d893f4d86812219b538a14307f347f10b --- test/cpp/tensorexpr/test_boundsinference.cpp | 6 +- test/cpp/tensorexpr/test_loopnest.cpp | 134 +++++++++--------- test/cpp/tensorexpr/test_reductions.cpp | 58 ++++---- torch/csrc/jit/tensorexpr/kernel.cpp | 8 +- torch/csrc/jit/tensorexpr/loopnest.h | 18 +-- torch/csrc/jit/tensorexpr/tensorexpr_init.cpp | 35 +++-- 6 files changed, 128 insertions(+), 131 deletions(-) diff --git a/test/cpp/tensorexpr/test_boundsinference.cpp b/test/cpp/tensorexpr/test_boundsinference.cpp index 87fb244e0cb9a..7af4801850130 100644 --- a/test/cpp/tensorexpr/test_boundsinference.cpp +++ b/test/cpp/tensorexpr/test_boundsinference.cpp @@ -222,7 +222,7 @@ TEST(BoundsInference, _5) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getLoopStmtsFor(b); - l.splitWithTail(loops[0], 16, &inner, &tail); + LoopNest::splitWithTail(loops[0], 16, &inner, &tail); For* outer = loops[0]; { @@ -533,7 +533,7 @@ TEST(BoundsInference, CacheReads) { auto bounds_info_before = inferBounds(l.root_stmt()); Stmt* j_loop = l.getLoopStmtsFor(B)[1]; - l.cacheAccesses(A->buf(), "A_local", j_loop); + LoopNest::cacheAccesses(A->buf(), "A_local", j_loop); auto bounds_info_after = inferBounds(l.root_stmt()); @@ -734,7 +734,7 @@ TEST(BoundsInference, GetPotentialHazardsLoopSplit) { // A. For* outer = l.getLoopStmtsFor(A)[0]; // `outer` loop get transformed to the outer loop after splitting. - l.splitWithTail(outer, 5, &inner, &tail); + LoopNest::splitWithTail(outer, 5, &inner, &tail); using namespace analysis; diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index ebb7f5cef4a65..60da47d776284 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -38,8 +38,8 @@ TEST(LoopNest, ExprSimple01) { LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.splitWithTail(loops[0], 2); - l.splitWithTail(loops[0], 2); + LoopNest::splitWithTail(loops[0], 2); + LoopNest::splitWithTail(loops[0], 2); } // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) @@ -67,7 +67,7 @@ TEST(LoopNest, ExprSimple02) { LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.splitWithTail(loops[0], 4); + LoopNest::splitWithTail(loops[0], 4); Stmt* stmt = l.root_stmt(); std::ostringstream oss; @@ -166,7 +166,7 @@ TEST(LoopNest, ExprSliceHeadWithLoopOptions) { For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); - l.sliceHead(loops[0], 2, &head, &tail); + LoopNest::sliceHead(loops[0], 2, &head, &tail); Block* body = getSimplifiedBody(l); assertForRanges(body, {{0, 2}, {0, 8}}); @@ -190,14 +190,14 @@ TEST(LoopNest, ExprSliceTailWithLoopOptions) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceTail(loops[0], 4, &head, &tail); + LoopNest::sliceTail(loops[0], 4, &head, &tail); // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail_head; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail_tail; l.setGPUBlockIndex(tail, LoopOptions::IDX_Y); - l.sliceTail(tail, 2, &tail_head, &tail_tail); + LoopNest::sliceTail(tail, 2, &tail_head, &tail_tail); Block* body = getSimplifiedBody(l); assertForRanges(body, {{0, 6}, {0, 2}, {8, 10}}); @@ -224,7 +224,7 @@ TEST(LoopNest, ExprSliceHeadWhenFactorEqualsSize) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceHead(loops[0], 10, &head, &tail); + LoopNest::sliceHead(loops[0], 10, &head, &tail); ASSERT_EQ(head, loops[0]); ASSERT_EQ(tail, nullptr); @@ -246,7 +246,7 @@ TEST(LoopNest, ExprSliceHeadWhenFactorLargerThanSize) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceHead(loops[0], 100, &head, &tail); + LoopNest::sliceHead(loops[0], 100, &head, &tail); ASSERT_EQ(head, loops[0]); ASSERT_EQ(tail, nullptr); @@ -268,7 +268,7 @@ TEST(LoopNest, ExprSliceHead) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceHead(loops[0], 4, &head, &tail); + LoopNest::sliceHead(loops[0], 4, &head, &tail); ASSERT_NE(head, nullptr); ASSERT_NE(head, loops[0]); @@ -293,11 +293,11 @@ TEST(LoopNest, ExprSliceHeadWithNonZeroStart) { For* head; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; - l.sliceTail(loops[0], 4, &head, &tail); + LoopNest::sliceTail(loops[0], 4, &head, &tail); // head: [0, 6) // tail: [6, 10) - l.sliceHead(tail, 2); + LoopNest::sliceHead(tail, 2); // tail_head: [6, 8) // tail_tail: [8, 10) @@ -320,7 +320,7 @@ TEST(LoopNest, ExprSliceTailWhenFactorEqualsSize) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceTail(loops[0], 10, &head, &tail); + LoopNest::sliceTail(loops[0], 10, &head, &tail); ASSERT_EQ(head, nullptr); ASSERT_EQ(tail, loops[0]); @@ -344,7 +344,7 @@ TEST(LoopNest, ExprSliceTailWhenFactorLargerThanSize) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceTail(loops[0], 100, &head, &tail); + LoopNest::sliceTail(loops[0], 100, &head, &tail); ASSERT_EQ(head, nullptr); ASSERT_EQ(tail, loops[0]); @@ -366,7 +366,7 @@ TEST(LoopNest, ExprSliceTail) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.sliceTail(loops[0], 4, &head, &tail); + LoopNest::sliceTail(loops[0], 4, &head, &tail); ASSERT_NE(head, nullptr); ASSERT_NE(head, loops[0]); @@ -397,9 +397,9 @@ TEST(LoopNest, ExprSplitAndSlice) { // outer: [0, 4) // inner: [0, 21) // tail: [84, 100) - l.splitWithTail(loops[0], 21, &inner, &tail); - l.sliceTail(inner, 2); - l.sliceHead(loops[0], 2); + LoopNest::splitWithTail(loops[0], 21, &inner, &tail); + LoopNest::sliceTail(inner, 2); + LoopNest::sliceHead(loops[0], 2); // for (int x_outer = 0; x_outer < 2; x_outer++) { // for (int x_inner = 0; x_inner < 19; x_inner++) { @@ -448,7 +448,7 @@ TEST(LoopNest, ExprSliceAndNormalize) { For* head; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; - l.sliceHead(loops[0], 2, &head, &tail); + LoopNest::sliceHead(loops[0], 2, &head, &tail); // head: [0, 2) // tail: [2, 10) @@ -482,9 +482,9 @@ TEST(LoopNest, ExprSliceWithVariableDimension) { For* head; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; - l.sliceHead(loops[0], 2, &head, &tail); + LoopNest::sliceHead(loops[0], 2, &head, &tail); - l.sliceTail(tail, 2); + LoopNest::sliceTail(tail, 2); Block* body = getSimplifiedBody(l); ASSERT_EQ(expected_for_ranges.size(), 3); @@ -516,9 +516,9 @@ TEST(LoopNest, ExprSplitWithTail) { LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) - l.splitWithTail(loops[0], 17); + LoopNest::splitWithTail(loops[0], 17); // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) - l.splitWithTail(loops[0], 7); + LoopNest::splitWithTail(loops[0], 7); Stmt* stmt = l.root_stmt(); Stmt* simplified = IRSimplifier::simplify(stmt); @@ -546,7 +546,7 @@ TEST(LoopNest, ExprSplitWithTailNone) { Tensor* tensor = Compute("f", {{24, "x"}, {5, "y"}}, func); LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.splitWithTail(loops[0], 4); + LoopNest::splitWithTail(loops[0], 4); Stmt* stmt = l.root_stmt(); std::ostringstream oss; @@ -610,7 +610,7 @@ TEST(LoopNest, ExprSplitWithMask01) { LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.splitWithMask(loops[1], 4); + LoopNest::splitWithMask(loops[1], 4); Stmt* stmt = l.root_stmt(); @@ -645,8 +645,8 @@ TEST(LoopNest, ExprSplitWithMaskRepeatedNoMask) { LoopNest l({tensor}); std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.splitWithMask(loops[0], 4); - l.splitWithMask(loops[0], 4); + LoopNest::splitWithMask(loops[0], 4); + LoopNest::splitWithMask(loops[0], 4); Stmt* stmt1 = IRSimplifier::simplify(l.root_stmt()); @@ -677,7 +677,7 @@ TEST(LoopNest, SplitWithTailWithLoopOptions) { auto loops = NodeFinder::find(l.root_stmt()); ASSERT_GT(loops.size(), 0); l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); - l.splitWithTail(loops[0], 4, &inner, &tail); + LoopNest::splitWithTail(loops[0], 4, &inner, &tail); ASSERT_NE(inner, nullptr); ASSERT_NE(tail, nullptr); For* outer = loops[0]; @@ -708,7 +708,7 @@ TEST(LoopNest, SplitWithMaskWithLoopOptions) { LoopNest l({tensor}); auto loops = NodeFinder::find(l.root_stmt()); l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); - l.splitWithMask(loops[0], 4, &inner); + LoopNest::splitWithMask(loops[0], 4, &inner); For* outer = loops[0]; // Outer loop carries loop axis bindings. @@ -1246,7 +1246,7 @@ TEST(LoopNest, ScheduleSplitAThenInline) { LoopNest l({b}, {a, b}); std::vector loops = l.getAllLoopNestsWritingToBuf(a->buf()).at(0); - l.splitWithMask(loops[0], 4); + LoopNest::splitWithMask(loops[0], 4); ASSERT_THROWS_WITH(l.computeInline(a->buf()), "compound indices"); } @@ -1262,7 +1262,7 @@ TEST(LoopNest, ScheduleSplitBThenInline) { LoopNest l({b}, {a, b}); std::vector loops = l.getAllLoopNestsWritingToBuf(b->buf()).at(0); - l.splitWithMask(loops[0], 3); + LoopNest::splitWithMask(loops[0], 3); l.computeInline(a->buf()); l.prepareForCodegen(); Stmt* s = IRSimplifier::simplify(l.root_stmt()); @@ -1290,8 +1290,8 @@ TEST(LoopNest, ScheduleSplitTwiceThenInline) { LoopNest l({b}, {a, b}); std::vector loops = l.getAllLoopNestsWritingToBuf(a->buf()).at(0); - l.splitWithMask(loops[0], 4, &i_inner); - l.splitWithMask(i_inner, 2); + LoopNest::splitWithMask(loops[0], 4, &i_inner); + LoopNest::splitWithMask(i_inner, 2); ASSERT_THROWS_WITH(l.computeInline(a->buf()), "compound indices"); } @@ -1309,7 +1309,7 @@ TEST(LoopNest, ScheduleInlineThenSplit) { l.computeInline(a->buf()); std::vector loops = NodeFinder::find(l.root_stmt()); - l.splitWithMask(loops.back(), 3); + LoopNest::splitWithMask(loops.back(), 3); l.prepareForCodegen(); Stmt* s = IRSimplifier::simplify(l.root_stmt()); std::vector output(6, 0); @@ -1333,11 +1333,11 @@ TEST(LoopNest, ScheduleSplitInlineThenSplit) { LoopNest l({b}, {a, b}); auto loops = NodeFinder::find(l.root_stmt()); - l.splitWithMask(loops.back(), 2); + LoopNest::splitWithMask(loops.back(), 2); l.computeInline(a->buf()); loops = NodeFinder::find(l.root_stmt()); - l.splitWithMask(loops.front(), 2); + LoopNest::splitWithMask(loops.front(), 2); l.prepareForCodegen(); Stmt* s = IRSimplifier::simplify(l.root_stmt()); std::vector output(16, 0); @@ -1362,7 +1362,7 @@ TEST(LoopNest, ScheduleSplitInlineSimplify) { LoopNest l({b}, {a, b}); std::vector loops = l.getAllLoopNestsWritingToBuf(a->buf()).at(0); - l.splitWithMask(loops[0], 4); + LoopNest::splitWithMask(loops[0], 4); ASSERT_THROWS_WITH(l.computeInline(a->buf()), "compound indices"); } @@ -1476,11 +1476,11 @@ TEST(LoopNest, ScheduleInlineThreeMixedSplit) { LoopNest l({c}, {a, b, c}); std::vector loops = l.getAllLoopNestsWritingToBuf(a->buf()).at(0); - l.splitWithMask(loops[0], 4); + LoopNest::splitWithMask(loops[0], 4); loops = l.getAllLoopNestsWritingToBuf(b->buf()).at(0); - l.splitWithMask(loops[0], 3); + LoopNest::splitWithMask(loops[0], 3); loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.splitWithMask(loops[0], 2); + LoopNest::splitWithMask(loops[0], 2); ASSERT_THROWS_WITH(l.computeInline(a->buf()), "compound indices"); } @@ -1649,7 +1649,7 @@ TEST(LoopNest, LoopNestComputeAt_1) { "B", {{N, "i_b"}}, [&](const VarHandle& i_b) { return A->load(i_b); }); LoopNest l({B}, {A, B}); std::vector loops = l.getAllLoopNestsWritingToBuf(B->buf()).at(0); - l.computeAt(l.getLoopBodyFor(A), loops[0]); + LoopNest::computeAt(l.getLoopBodyFor(A), loops[0]); l.prepareForCodegen(); Stmt* s = l.root_stmt(); @@ -1717,7 +1717,7 @@ TEST(LoopNest, LoopNestComputeAt_2) { // First let's try to compute P at axis cy (the outer loop) LoopNest l(orig_loopnest); std::vector loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.computeAt(l.getLoopBodyFor(p), loops[0]); + LoopNest::computeAt(l.getLoopBodyFor(p), loops[0]); l.prepareForCodegen(); Stmt* s = l.root_stmt(); @@ -1743,7 +1743,7 @@ TEST(LoopNest, LoopNestComputeAt_2) { // Now let's try to compute P at axis cx (the inner loop) LoopNest l(orig_loopnest); std::vector loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.computeAt(l.getLoopBodyFor(p), loops[1]); + LoopNest::computeAt(l.getLoopBodyFor(p), loops[1]); l.prepareForCodegen(); Stmt* s = l.root_stmt(); @@ -1817,7 +1817,7 @@ TEST(LoopNest, LoopNestComputeAt_3) { // First let's try to compute A at axis dy (the outer loop) LoopNest l(orig_loopnest); std::vector loops = l.getAllLoopNestsWritingToBuf(D->buf()).at(0); - l.computeAt(l.getLoopBodyFor(A), loops[0]); + LoopNest::computeAt(l.getLoopBodyFor(A), loops[0]); l.prepareForCodegen(); Stmt* s = l.root_stmt(); @@ -1848,7 +1848,7 @@ TEST(LoopNest, LoopNestComputeAt_3) { // Now let's try to compute A at axis dx (the inner loop) LoopNest l(orig_loopnest); std::vector loops = l.getAllLoopNestsWritingToBuf(D->buf()).at(0); - l.computeAt(l.getLoopBodyFor(A), loops[1]); + LoopNest::computeAt(l.getLoopBodyFor(A), loops[1]); l.prepareForCodegen(); Stmt* s = l.root_stmt(); @@ -1927,7 +1927,7 @@ TEST(LoopNest, Reduce2dComputeAt) { // First let's try to compute P at axis cy (the outer loop) LoopNest l(orig_loopnest); auto loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.computeAt(l.getLoopBodyFor(p), loops[0]); + LoopNest::computeAt(l.getLoopBodyFor(p), loops[0]); // FIXME: Calling simplify here breaks the IR: // MALFORMED INPUT: could not find base node in Load - temp[...] // l.simplify(); @@ -1964,7 +1964,7 @@ TEST(LoopNest, Reduce2dComputeAt) { // Now let's try to compute P at axis cx (the inner loop) LoopNest l(orig_loopnest); std::vector loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.computeAt(l.getLoopBodyFor(p), loops[1]); + LoopNest::computeAt(l.getLoopBodyFor(p), loops[1]); l.simplify(); l.eliminateDeadStores(); l.prepareForCodegen(); @@ -2039,7 +2039,7 @@ TEST(LoopNest, DISABLED_Conv1d_NH) { # CHECK: } )IR"); std::vector loops = l.getAllLoopNestsWritingToBuf(B->buf()).at(0); - l.computeAt(l.getLoopBodyFor(A), loops[0]); + LoopNest::computeAt(l.getLoopBodyFor(A), loops[0]); // FIXME: The current IR is totally broken. The body of the inlined loop is: // temp[idx0, idx1] = IfThenElse(idx0 + n>=257 ? 1 : (idx0 + n<1 ? 1 : 0), @@ -2110,7 +2110,7 @@ TEST(LoopNest, LoopNestReorderAxis1) { cg.call({stmt1_output}); auto loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[0], loops[1]); + LoopNest::reorderAxis(loops[0], loops[1]); Stmt* stmt2 = Stmt::clone(l.root_stmt()); ASSERT_NE(stmt1, stmt2); @@ -2131,7 +2131,7 @@ TEST(LoopNest, LoopNestReorderAxis1) { // Reorder them back. loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[0], loops[1]); + LoopNest::reorderAxis(loops[0], loops[1]); Stmt* stmt3 = l.root_stmt(); std::string order3 = loopOrderHelper.getOrder(stmt3); @@ -2166,7 +2166,7 @@ TEST(LoopNest, LoopNestReorderPartialAxes) { cg.call({stmt1_output}); auto loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[0], loops[1]); + LoopNest::reorderAxis(loops[0], loops[1]); ASSERT_EQ(loopOrderHelper.getOrder(l.root_stmt()), "y,x,z,"); Stmt* stmt2 = Stmt::clone(l.root_stmt()); @@ -2180,7 +2180,7 @@ TEST(LoopNest, LoopNestReorderPartialAxes) { } loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[1], loops[2]); + LoopNest::reorderAxis(loops[1], loops[2]); ASSERT_EQ(loopOrderHelper.getOrder(l.root_stmt()), "y,z,x,"); Stmt* stmt3 = Stmt::clone(l.root_stmt()); @@ -2218,7 +2218,7 @@ TEST(LoopNest, LoopNestReorderInternalAxis) { cg.call({stmt1_output}); auto loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[2], loops[1]); + LoopNest::reorderAxis(loops[2], loops[1]); ASSERT_EQ(loopOrderHelper.getOrder(l.root_stmt()), "w,y,x,z,"); Stmt* stmt2 = l.root_stmt(); @@ -2255,7 +2255,7 @@ TEST(LoopNest, LoopNestReorderEnclosingAxis) { cg.call({stmt1_output}); auto loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[0], loops[3]); + LoopNest::reorderAxis(loops[0], loops[3]); ASSERT_EQ(loopOrderHelper.getOrder(l.root_stmt()), "z,x,y,w,"); Stmt* stmt2 = l.root_stmt(); @@ -2280,7 +2280,7 @@ TEST(LoopNest, LoopNestReorderSameAxis) { Stmt* stmt1 = Stmt::clone(l.root_stmt()); auto loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[1], loops[1]); + LoopNest::reorderAxis(loops[1], loops[1]); Stmt* stmt2 = Stmt::clone(l.root_stmt()); std::ostringstream oss, oss2; @@ -2350,7 +2350,7 @@ TEST(LoopNest, LoopNestReorderExtraStatements) { * */ - l.reorderAxis(loops[1], loops[2]); + LoopNest::reorderAxis(loops[1], loops[2]); Stmt* stmt2 = Stmt::clone(l.root_stmt()); // Check the IR we produced @@ -2398,7 +2398,7 @@ TEST(LoopNest, LoopNestReorderExtraStatements) { * */ loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.reorderAxis(loops[0], loops[2]); + LoopNest::reorderAxis(loops[0], loops[2]); Stmt* stmt3 = Stmt::clone(l.root_stmt()); // Check the IR we produced @@ -2484,7 +2484,7 @@ void LoopNestReorderTestHelper( } loops = l.getAllLoopNestsWritingToBuf(c->buf()).at(0); - l.reorderAxis(loops[index1], loops[index2]); + LoopNest::reorderAxis(loops[index1], loops[index2]); Stmt* stmt2 = Stmt::clone(l.root_stmt()); std::ostringstream oss, oss2; @@ -2592,7 +2592,7 @@ TEST(LoopNest, LoopNestReorderInternalLoopNest) { b = f; } } - l.reorderAxis(a, b); + LoopNest::reorderAxis(a, b); l.prepareForCodegen(); Stmt* stmt = IRSimplifier::simplify(l.root_stmt()); @@ -3173,7 +3173,7 @@ TEST(LoopNest, NormalizeAndSplitWithTail) { For* x_inner; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* x_tail; - l.splitWithTail(for_stmt, 10, &x_inner, &x_tail); + LoopNest::splitWithTail(for_stmt, 10, &x_inner, &x_tail); auto x_outer_result = IRSimplifier::simplify(for_stmt); std::ostringstream oss_outer; @@ -3541,7 +3541,7 @@ TEST(LoopNest, CacheReadsSimple) { LoopNest l({B, C}, {A, B, C}); Stmt* j_loop = l.getAllLoopNestsWritingToBuf(B->buf())[0][1]; - l.cacheAccesses(A->buf(), "A_local", j_loop); + LoopNest::cacheAccesses(A->buf(), "A_local", j_loop); l.prepareForCodegen(); Stmt* result = IRSimplifier::simplify(l.root_stmt()); @@ -3610,7 +3610,7 @@ TEST(LoopNest, CacheReadsOuter) { LoopNest l({B, C}, {A, B, C}); Stmt* i_loop = l.getAllLoopNestsWritingToBuf(B->buf())[0][0]; - l.cacheAccesses(A->buf(), "A_local", i_loop); + LoopNest::cacheAccesses(A->buf(), "A_local", i_loop); l.prepareForCodegen(); Stmt* result = IRSimplifier::simplify(l.root_stmt()); @@ -3659,7 +3659,7 @@ TEST(LoopNest, CacheReadsInternal) { LoopNest l({B, C}, {A, B, C}); Stmt* j_loop = l.getAllLoopNestsWritingToBuf(B->buf())[0][1]; - l.cacheAccesses(A->buf(), "A_local", j_loop); + LoopNest::cacheAccesses(A->buf(), "A_local", j_loop); l.prepareForCodegen(); Stmt* result = IRSimplifier::simplify(l.root_stmt()); @@ -3708,7 +3708,7 @@ TEST(LoopNest, CacheReadsInner) { LoopNest l({B, C}, {A, B, C}); Stmt* body = l.getLoopBodyFor(B); - l.cacheAccesses(A->buf(), "A_local", body); + LoopNest::cacheAccesses(A->buf(), "A_local", body); l.prepareForCodegen(); Stmt* result = IRSimplifier::simplify(l.root_stmt()); @@ -3756,7 +3756,7 @@ TEST(LoopNest, CacheWritesSimple) { LoopNest l({B, C}, {A, B, C}); Stmt* a_loop = l.getAllLoopNestsWritingToBuf(A->buf())[0][1]; - l.cacheAccesses(A->buf(), "A_local", a_loop); + LoopNest::cacheAccesses(A->buf(), "A_local", a_loop); l.prepareForCodegen(); Stmt* result = IRSimplifier::simplify(l.root_stmt()); @@ -4593,7 +4593,7 @@ static Stmt* splitTailReorder(Tensor* b) { // Loopnest #2: {n_outer, n_inner, m}; // We will have to reorder n_inner and m. auto loopnests = nest.getAllLoopNestsWritingToBuf(b->buf()); - nest.reorderAxis(loopnests[1][1], loopnests[1][2]); + LoopNest::reorderAxis(loopnests[1][1], loopnests[1][2]); nest.prepareForCodegen(); return nest.root_stmt(); } @@ -4604,7 +4604,7 @@ static Stmt* splitMaskReorder(Tensor* b) { auto loops = nest.getAllLoopNestsWritingToBuf(b->buf())[1]; nest.splitWithMask(loops[0], kVectorWidth); loops = nest.getAllLoopNestsWritingToBuf(b->buf())[1]; - nest.reorderAxis(loops[1], loops[2]); + LoopNest::reorderAxis(loops[1], loops[2]); nest.prepareForCodegen(); return nest.root_stmt(); } @@ -4723,7 +4723,7 @@ TEST(LoopNest, ReorderAxisWithMultipleConds) { auto forI = For::make(i, 0, 20, outer_cond); Stmt* par = Block::make({forI}); LoopNest l(par, {a_buf.node()}); - l.reorderAxis(forI, forJ); + LoopNest::reorderAxis(forI, forJ); ASSERT_EQ(par, l.root_stmt()); par = IRSimplifier::simplify(par); diff --git a/test/cpp/tensorexpr/test_reductions.cpp b/test/cpp/tensorexpr/test_reductions.cpp index 73e565157609b..ccbd9c63e3590 100644 --- a/test/cpp/tensorexpr/test_reductions.cpp +++ b/test/cpp/tensorexpr/test_reductions.cpp @@ -642,7 +642,7 @@ TEST(Reductions, SplitReduceAxis) { Tensor* tensor = Reduce("sum", {{16, "m"}}, Sum(), in, {{8, "n"}}); LoopNest l({tensor}); std::vector loops = l.getLoopStmtsFor(tensor); - l.splitWithTail(loops[1], 2); + LoopNest::splitWithTail(loops[1], 2); l.prepareForCodegen(); @@ -673,8 +673,8 @@ TEST(Reductions, SplitNonReduceAxis) { Tensor* tensor = Reduce("sum", {{16, "m"}}, Sum(), in, {{8, "n"}}); LoopNest l({tensor}); std::vector loops = l.getLoopStmtsFor(tensor); - l.splitWithTail(loops[0], 2); - l.splitWithTail(loops[0], 2); + LoopNest::splitWithTail(loops[0], 2); + LoopNest::splitWithTail(loops[0], 2); l.prepareForCodegen(); @@ -716,7 +716,7 @@ TEST(Reductions, ReorderedReductionInitializer) { l.setGPUBlockIndex(loops[0], 0); l.setGPUThreadIndex(loops[1], 0); - l.reorderAxis(loops[1], loops[2]); + LoopNest::reorderAxis(loops[1], loops[2]); Stmt* s = l.root_stmt(); // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores) @@ -918,7 +918,7 @@ TEST(Reductions, ReduceSplitTail) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithTail(loops[i], 8); + LoopNest::splitWithTail(loops[i], 8); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -951,7 +951,7 @@ TEST(Reductions, ReduceSplitNoTail) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithTail(loops[i], 5); + LoopNest::splitWithTail(loops[i], 5); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -986,7 +986,7 @@ TEST(Reductions, ReduceOverSplitTail) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithTail(loops[i], 16); + LoopNest::splitWithTail(loops[i], 16); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -1020,7 +1020,7 @@ TEST(Reductions, ReduceSplitMask) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithMask(loops[i], 8); + LoopNest::splitWithMask(loops[i], 8); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -1053,7 +1053,7 @@ TEST(Reductions, ReduceSplitNoMask) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithMask(loops[i], 5); + LoopNest::splitWithMask(loops[i], 5); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -1087,7 +1087,7 @@ TEST(Reductions, ReduceOverSplitMask) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithMask(loops[i], 16); + LoopNest::splitWithMask(loops[i], 16); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -1124,13 +1124,13 @@ TEST(Reductions, ReduceSplitRfactor) { Tensor* c = Reduce("sum", {{M, "m"}}, Sum(), b, {{N, "n"}, {K, "k"}}); LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.splitWithTail(loops[2], SPLIT_FACTOR); + LoopNest::splitWithTail(loops[2], SPLIT_FACTOR); // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto c_body = const_cast(loop.getAllWritesToBuf(c->buf())[2]); auto all_loops = loop.getAllLoopNestsWritingToBuf(c->buf()); ASSERT_TRUE(all_loops.size() == 3 && all_loops.at(2).size() == 3); - loop.reorderAxis(all_loops[2][1], all_loops[2][2]); + LoopNest::reorderAxis(all_loops[2][1], all_loops[2][2]); all_loops = loop.getAllLoopNestsWritingToBuf(c->buf()); ASSERT_TRUE(all_loops.size() == 3 && all_loops.at(2).size() == 3); ASSERT_TRUE(loop.rfactor(c_body, all_loops[2][1])); @@ -1169,15 +1169,15 @@ TEST(Reductions, ReduceOverSplitRfactor) { std::vector loops = loop.getLoopStmtsFor(c); // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For *i, *t; - loop.splitWithTail(loops[1], SPLIT_FACTOR, &i, &t); - loop.reorderAxis(loops[0], i); + LoopNest::splitWithTail(loops[1], SPLIT_FACTOR, &i, &t); + LoopNest::reorderAxis(loops[0], i); auto all_loops = loop.getAllLoopNestsWritingToBuf(c->buf()); ASSERT_TRUE(all_loops.size() == 3 && all_loops.at(1).size() == 3); // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto c_body = const_cast(loop.getAllWritesToBuf(c->buf())[1]); ASSERT_TRUE(loop.rfactor(c_body, all_loops[1][0])); - loop.reorderAxis(all_loops[1][0], all_loops[1][2]); + LoopNest::reorderAxis(all_loops[1][0], all_loops[1][2]); loop.prepareForCodegen(); loop.simplify(); @@ -1648,7 +1648,7 @@ TEST(Reductions, ReductionCacheConsumerAccess) { LoopNest l({e}, {c, d, e}); - l.splitWithMask(l.getLoopStmtsFor(e)[0], 4); + LoopNest::splitWithMask(l.getLoopStmtsFor(e)[0], 4); Stmt* e_loop = l.getLoopStmtsFor(e)[1]; l.cacheAccesses(d->buf(), "sum_local", e_loop); @@ -1694,10 +1694,10 @@ TEST(Reductions, ReductionSplitCacheConsumerAccess) { For* inner; // Split outer reduction axis. - l.splitWithMask(l.getLoopStmtsFor(d)[0], 4, &inner); + LoopNest::splitWithMask(l.getLoopStmtsFor(d)[0], 4, &inner); // Split reduction consumer. - l.splitWithMask(l.getLoopStmtsFor(e)[0], 4, &inner); + LoopNest::splitWithMask(l.getLoopStmtsFor(e)[0], 4, &inner); l.cacheAccesses(d->buf(), "sum_local", inner); l.prepareForCodegen(); @@ -1744,10 +1744,10 @@ TEST(Reductions, ReductionReorderCacheConsumerAccess) { // reorder outer reduction axes. auto loops = l.getLoopStmtsFor(d); - l.reorderAxis(loops[0], loops[1]); + LoopNest::reorderAxis(loops[0], loops[1]); // Split reduction consumer. - l.splitWithMask(l.getLoopStmtsFor(e)[0], 4, &inner); + LoopNest::splitWithMask(l.getLoopStmtsFor(e)[0], 4, &inner); l.cacheAccesses(d->buf(), "sum_local", inner); l.prepareForCodegen(); @@ -1791,7 +1791,7 @@ TEST(Reductions, ReductionRfactorCacheTempOuter) { LoopNest loop({c}); std::vector loops = loop.getLoopStmtsFor(c); - loop.reorderAxis(loops.at(0), loops.at(1)); + LoopNest::reorderAxis(loops.at(0), loops.at(1)); loops = loop.getLoopStmtsFor(c); // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto c_body = const_cast(loop.getAllWritesToBuf(c->buf())[1]); @@ -1802,10 +1802,10 @@ TEST(Reductions, ReductionRfactorCacheTempOuter) { auto all_loops = loop.getAllLoopNestsWritingToBuf(rfac_buf); ASSERT_TRUE(all_loops.size() == 2 && all_loops.at(1).size() == 3); - loop.reorderAxis(all_loops[1][0], all_loops[1][1]); + LoopNest::reorderAxis(all_loops[1][0], all_loops[1][1]); all_loops = loop.getAllLoopNestsWritingToBuf(rfac_buf); - loop.cacheAccesses(rfac_buf, "tmp", all_loops[1][1]); + LoopNest::cacheAccesses(rfac_buf, "tmp", all_loops[1][1]); loop.simplify(); loop.prepareForCodegen(); Stmt* s = loop.root_stmt(); @@ -1864,7 +1864,7 @@ TEST(Reductions, ReductionRfactorCacheTempInner) { // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto c_body = const_cast(loop.getAllWritesToBuf(c->buf())[1]); - loop.reorderAxis(loops.at(0), loops.at(1)); + LoopNest::reorderAxis(loops.at(0), loops.at(1)); loops = loop.getLoopStmtsFor(c); // NOLINTNEXTLINE(cppcoreguidelines-init-variables) Buf* rfac_buf; @@ -1872,11 +1872,11 @@ TEST(Reductions, ReductionRfactorCacheTempInner) { loop.distributeLoop(loops.at(0)); auto all_loops = loop.getAllLoopNestsWritingToBuf(rfac_buf); ASSERT_TRUE(all_loops.size() == 2 && all_loops.at(1).size() == 3); - loop.reorderAxis(all_loops[1][0], all_loops[1][1]); + LoopNest::reorderAxis(all_loops[1][0], all_loops[1][1]); all_loops = loop.getAllLoopNestsWritingToBuf(rfac_buf); ASSERT_TRUE(all_loops.size() == 2 && all_loops.at(1).size() == 3); - loop.cacheAccesses(rfac_buf, "tmp", all_loops[1][2]); + LoopNest::cacheAccesses(rfac_buf, "tmp", all_loops[1][2]); loop.prepareForCodegen(); loop.simplify(); Stmt* s = loop.root_stmt(); @@ -1993,14 +1993,14 @@ TEST(Reductions, ReductionVectorizeRfactor) { // But if we rfactor this so it's not a reduce axis we can vectorize that // loop. std::vector loops = l.getLoopStmtsFor(tensor); - l.reorderAxis(loops[0], loops[1]); + LoopNest::reorderAxis(loops[0], loops[1]); loops = l.getLoopStmtsFor(tensor); // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) auto tensor_body = const_cast(l.getAllWritesToBuf(tensor->buf())[1]); Buf* rfac_buf = nullptr; - ASSERT_TRUE(l.rfactor(tensor_body, loops.at(0), &rfac_buf)); + ASSERT_TRUE(LoopNest::rfactor(tensor_body, loops.at(0), &rfac_buf)); - l.distributeLoop(loops.at(0)); + LoopNest::distributeLoop(loops.at(0)); auto rfac_loops = l.getAllLoopNestsWritingToBuf(rfac_buf); ASSERT_TRUE(LoopNest::vectorize(rfac_loops[1][0])); diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 315f93062bbe6..2bca7421a330f 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -2984,7 +2984,7 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { if (blockSize < 0) { blockSize = kDefaultBlockSize; } - l.splitWithMask(flattened, blockSize, &inner); + LoopNest::splitWithMask(flattened, blockSize, &inner); l.setGPUBlockIndex(flattened, 0); l.setGPUThreadIndex(inner, 0); } else if (loopLevels == 3) { @@ -2997,8 +2997,8 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { const int kDefaultBlockSize = 256; blockCount = (blockCount > 0) ? blockCount : kDefaultBlockCount; blockSize = (blockSize > 0) ? blockSize : kDefaultBlockSize; - l.splitWithMask(flattened, blockCount * blockSize, &inner); - l.splitWithMask(inner, blockSize, &inner1); + LoopNest::splitWithMask(flattened, blockCount * blockSize, &inner); + LoopNest::splitWithMask(inner, blockSize, &inner1); l.setGPUBlockIndex(inner, 0); l.setGPUThreadIndex(inner1, 0); } else { @@ -3024,7 +3024,7 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { assert(flattened); For* inner = nullptr; - l.splitWithMask(flattened, blockSize, &inner); + LoopNest::splitWithMask(flattened, blockSize, &inner); l.setGPUBlockIndex(flattened, 0); l.setGPUThreadIndex(inner, 0); l.setBufferMap(flattened, block_analysis->getBufferMap()); diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h index 717f9e8a39465..4157716bc2e94 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.h +++ b/torch/csrc/jit/tensorexpr/loopnest.h @@ -264,7 +264,7 @@ class TORCH_API LoopNest { // * Fusing the loops does not violate or add any dependencies. static bool fuseLoops(const std::vector& loops, For** fused); - void reorderAxis(For* a, For* b); + static void reorderAxis(For* a, For* b); // Reorder the given list of loops according to the permutation specified. // Here permutation[i] represents the location of the loop i in the result. @@ -333,11 +333,11 @@ class TORCH_API LoopNest { static std::vector getLoopStmtsInLoopNest(For* f, size_t num); // LoopOptions are propagated to tail. - void sliceHead(For* f, int factor, For** head, For** tail); - void sliceHead(For* f, int factor); + static void sliceHead(For* f, int factor, For** head, For** tail); + static void sliceHead(For* f, int factor); // LoopOptions are propagated to head. - void sliceTail(For* f, int factor, For** head, For** tail); - void sliceTail(For* f, int factor); + static void sliceTail(For* f, int factor, For** head, For** tail); + static void sliceTail(For* f, int factor); void setGPUBlockIndex(For* f, int idx); void setGPUThreadIndex(For* f, int idx); @@ -346,7 +346,7 @@ class TORCH_API LoopNest { // Insert a cache for the consumer's usages of the buffer produced in // consumer, and redirect reads and writes in the consumer to that cache. // Returns a pair of the new cache buffer, and the new rewritten consumer. - AccessResult cacheAccesses( + static AccessResult cacheAccesses( const Buf* producer, const std::string& name, Stmt* consumer); @@ -355,7 +355,7 @@ class TORCH_API LoopNest { // S is assumed to be a Store or a Block containing a Store. Along with the // computation itself, this transformation inserts Alloc/Free statements for // the temporary buffer used in the computation. - void computeAt(Stmt* s, For* at); + static void computeAt(Stmt* s, For* at); // Rfactor a reduction axis into a normal axis. // @@ -399,8 +399,8 @@ class TORCH_API LoopNest { // S4: for k # reduction axis // X_rfac[i,j] = ReduceOp(X_rfac[i,j] + Y[i,j,k], reduce_axis={k}) // X[i] = ReduceOp(X[i] + X_rfac[i,j], reduce_axis={j}) - bool rfactor(Stmt* s, For* outer_reduction_for); - bool rfactor(Stmt* s, For* outer_reduction_for, Buf** rfac_buf_ptr); + static bool rfactor(Stmt* s, For* outer_reduction_for); + static bool rfactor(Stmt* s, For* outer_reduction_for, Buf** rfac_buf_ptr); // Vectorize the given loop. This method requires that the given loop // does not perform a reduction. diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp index f73fe02a3b3b8..05adfd70fa9eb 100644 --- a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp +++ b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp @@ -426,33 +426,33 @@ void initTensorExprBindings(PyObject* module) { py::return_value_policy::reference) .def( "split_with_tail", - [](const LoopNest& self, For* f, int factor) { + [](For* f, int factor) { For *inner = nullptr, *tail = nullptr; - self.splitWithTail(f, factor, &inner, &tail); + LoopNest::splitWithTail(f, factor, &inner, &tail); return std::make_tuple(inner, tail); }, py::return_value_policy::reference) .def( "split_with_mask", - [](const LoopNest& self, For* f, int factor) { + [](For* f, int factor) { For* inner = nullptr; - self.splitWithMask(f, factor, &inner); + LoopNest::splitWithMask(f, factor, &inner); return inner; }, py::return_value_policy::reference) .def( "slice_head", - [](LoopNest& self, For* f, int factor) { + [](For* f, int factor) { For *head = nullptr, *tail = nullptr; - self.sliceHead(f, factor, &head, &tail); + LoopNest::sliceHead(f, factor, &head, &tail); return std::make_tuple(head, tail); }, py::return_value_policy::reference) .def( "slice_tail", - [](LoopNest& self, For* f, int factor) { + [](For* f, int factor) { For *head = nullptr, *tail = nullptr; - self.sliceTail(f, factor, &head, &tail); + LoopNest::sliceTail(f, factor, &head, &tail); return std::make_tuple(head, tail); }, py::return_value_policy::reference) @@ -502,7 +502,7 @@ void initTensorExprBindings(PyObject* module) { py::return_value_policy::reference) .def( "vectorize", - [](const LoopNest& self, For* f) { self.vectorize(f); }, + [](For* f) { LoopNest::vectorize(f); }, py::return_value_policy::reference) .def_static( "compress_buffer", @@ -512,18 +512,15 @@ void initTensorExprBindings(PyObject* module) { py::return_value_policy::reference) .def( "cache_accesses", - [](LoopNest& self, - const BufHandle& producer, + [](const BufHandle& producer, const std::string& name, Stmt* consumer) { std::pair ret = - self.cacheAccesses(producer.node(), name, consumer); + LoopNest::cacheAccesses(producer.node(), name, consumer); return std::make_pair(BufHandle(ret.first), ret.second); }, py::return_value_policy::reference) - .def( - "compute_at", - [](LoopNest& self, Stmt* s, For* at) { self.computeAt(s, at); }) + .def("compute_at", [](Stmt* s, For* at) { LoopNest::computeAt(s, at); }) .def( "compute_inline", [](LoopNest& self, Stmt* s) { self.computeInline(s); }, @@ -536,17 +533,17 @@ void initTensorExprBindings(PyObject* module) { py::return_value_policy::reference) .def( "rfactor", - [](LoopNest& self, Stmt* s, For* target_for) { + [](Stmt* s, For* target_for) { Buf* rfac_buf = nullptr; - self.rfactor(s, target_for, &rfac_buf); + LoopNest::rfactor(s, target_for, &rfac_buf); return BufHandle(rfac_buf); }, py::return_value_policy::reference) .def( "flatten", - [](const LoopNest& self, const std::vector& loops) { + [](const std::vector& loops) { For* flattened = nullptr; - self.flatten(loops, &flattened); + LoopNest::flatten(loops, &flattened); return flattened; }, py::return_value_policy::reference) From b822928e33605ee88141c6a852882741dab06325 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Tue, 15 Jun 2021 10:34:45 -0700 Subject: [PATCH 109/305] [nnc] Removed setGPUBlockIndex and setGPUThreadIndex methods from LoopNest (#59495) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59495 Test Plan: Imported from OSS Reviewed By: ZolotukhinM Differential Revision: D28915960 Pulled By: navahgar fbshipit-source-id: 20a4032b031aba6e43d85433ade5f0680c65fbc0 --- test/cpp/tensorexpr/test_cuda.cpp | 68 +++++++++---------- test/cpp/tensorexpr/test_loopnest.cpp | 8 +-- test/cpp/tensorexpr/test_reductions.cpp | 4 +- torch/csrc/jit/tensorexpr/kernel.cpp | 12 ++-- torch/csrc/jit/tensorexpr/loopnest.cpp | 8 --- torch/csrc/jit/tensorexpr/loopnest.h | 3 - torch/csrc/jit/tensorexpr/tensorexpr_init.cpp | 18 ++--- 7 files changed, 56 insertions(+), 65 deletions(-) diff --git a/test/cpp/tensorexpr/test_cuda.cpp b/test/cpp/tensorexpr/test_cuda.cpp index a12592939a0f6..f75c60e1b38ec 100644 --- a/test/cpp/tensorexpr/test_cuda.cpp +++ b/test/cpp/tensorexpr/test_cuda.cpp @@ -46,8 +46,8 @@ static void testCudaTestVectorAdd01_impl() { }); LoopNest l({c}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[1], 0); - l.setGPUThreadIndex(loops[2], 0); + loops[1]->set_gpu_block_index(0); + loops[2]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); CudaCodeGen cuda_cg(stmt, c, a_buf, b_buf); @@ -111,8 +111,8 @@ TEST(Cuda, Sigmoid_CUDA) { }); LoopNest l({c}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[1], 0); - l.setGPUThreadIndex(loops[2], 0); + loops[1]->set_gpu_block_index(0); + loops[2]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); CudaCodeGen cuda_cg(stmt, c, a_buf); @@ -175,8 +175,8 @@ static void testCudaTestVectorAdd02_impl(int N, int block_size) { For* n_inner; std::vector loops = l.getLoopStmtsFor(c); l.splitWithMask(loops[0], block_size, &n_inner); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(n_inner, 0); + loops[0]->set_gpu_block_index(0); + n_inner->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); CudaCodeGen cuda_cg(stmt, c, a_buf, b_buf); @@ -340,8 +340,8 @@ TEST(Cuda, TestRand01_CUDA) { }); LoopNest l({c}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[1], 0); - l.setGPUThreadIndex(loops[2], 0); + loops[1]->set_gpu_block_index(0); + loops[2]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); CudaCodeGen cuda_cg(stmt, c); @@ -393,8 +393,8 @@ TEST(Cuda, DynamicShapeSplit_CUDA) { For* inner; std::vector loops = l.getLoopStmtsFor(b); l.splitWithMask(loops[0], 1024, &inner); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(inner, 0); + loops[0]->set_gpu_block_index(0); + inner->set_gpu_thread_index(0); Stmt* s = l.root_stmt(); CudaCodeGen cg(s, {a, b, n}); @@ -1176,9 +1176,9 @@ TEST(Cuda, MaskBlockDim_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); + loops[0]->set_gpu_block_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 0); + loops[0]->set_gpu_block_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -1269,9 +1269,9 @@ TEST(Cuda, MaskThreadDim_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUThreadIndex(loops[0], 0); + loops[0]->set_gpu_thread_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUThreadIndex(loops[0], 0); + loops[0]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -1364,9 +1364,9 @@ TEST(Cuda, MaskMultiBlockDim_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); + loops[0]->set_gpu_block_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 1); + loops[0]->set_gpu_block_index(1); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -1458,9 +1458,9 @@ TEST(Cuda, MaskBlockAndThreadDim_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); + loops[0]->set_gpu_block_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUThreadIndex(loops[0], 0); + loops[0]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -1557,11 +1557,11 @@ TEST(Cuda, MaskMultiDim_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -1687,11 +1687,11 @@ TEST(Cuda, MaskMultiDimSymbolic_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -2101,11 +2101,11 @@ TEST(Cuda, MaskMultiDimMultiAxis_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 1); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(1); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); @@ -2232,11 +2232,11 @@ TEST(Cuda, MaskMultiDimMultiLevel_CUDA) { LoopNest l({c, d}); std::vector loops = l.getLoopStmtsFor(c); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); loops = l.getLoopStmtsFor(d); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); l.prepareForCodegen(); Stmt* stmt = l.root_stmt(); diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index 60da47d776284..12e87801a9d80 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -165,7 +165,7 @@ TEST(LoopNest, ExprSliceHeadWithLoopOptions) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail; std::vector loops = l.getAllLoopNestsWritingToBuf(tensor->buf()).at(0); - l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); + loops[0]->set_gpu_block_index(LoopOptions::IDX_Y); LoopNest::sliceHead(loops[0], 2, &head, &tail); Block* body = getSimplifiedBody(l); @@ -196,7 +196,7 @@ TEST(LoopNest, ExprSliceTailWithLoopOptions) { For* tail_head; // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* tail_tail; - l.setGPUBlockIndex(tail, LoopOptions::IDX_Y); + tail->set_gpu_block_index(LoopOptions::IDX_Y); LoopNest::sliceTail(tail, 2, &tail_head, &tail_tail); Block* body = getSimplifiedBody(l); @@ -676,7 +676,7 @@ TEST(LoopNest, SplitWithTailWithLoopOptions) { LoopNest l({tensor}); auto loops = NodeFinder::find(l.root_stmt()); ASSERT_GT(loops.size(), 0); - l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); + loops[0]->set_gpu_block_index(LoopOptions::IDX_Y); LoopNest::splitWithTail(loops[0], 4, &inner, &tail); ASSERT_NE(inner, nullptr); ASSERT_NE(tail, nullptr); @@ -707,7 +707,7 @@ TEST(LoopNest, SplitWithMaskWithLoopOptions) { LoopNest l({tensor}); auto loops = NodeFinder::find(l.root_stmt()); - l.setGPUBlockIndex(loops[0], LoopOptions::IDX_Y); + loops[0]->set_gpu_block_index(LoopOptions::IDX_Y); LoopNest::splitWithMask(loops[0], 4, &inner); For* outer = loops[0]; diff --git a/test/cpp/tensorexpr/test_reductions.cpp b/test/cpp/tensorexpr/test_reductions.cpp index ccbd9c63e3590..6068a21be9eb1 100644 --- a/test/cpp/tensorexpr/test_reductions.cpp +++ b/test/cpp/tensorexpr/test_reductions.cpp @@ -713,8 +713,8 @@ TEST(Reductions, ReorderedReductionInitializer) { LoopNest l({tensor}); auto loops = l.getLoopStmtsFor(tensor); - l.setGPUBlockIndex(loops[0], 0); - l.setGPUThreadIndex(loops[1], 0); + loops[0]->set_gpu_block_index(0); + loops[1]->set_gpu_thread_index(0); LoopNest::reorderAxis(loops[1], loops[2]); diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 2bca7421a330f..7b832080dadc4 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -2985,8 +2985,8 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { blockSize = kDefaultBlockSize; } LoopNest::splitWithMask(flattened, blockSize, &inner); - l.setGPUBlockIndex(flattened, 0); - l.setGPUThreadIndex(inner, 0); + flattened->set_gpu_block_index(0); + inner->set_gpu_thread_index(0); } else if (loopLevels == 3) { // NOLINTNEXTLINE(cppcoreguidelines-init-variables) For* inner; @@ -2999,8 +2999,8 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { blockSize = (blockSize > 0) ? blockSize : kDefaultBlockSize; LoopNest::splitWithMask(flattened, blockCount * blockSize, &inner); LoopNest::splitWithMask(inner, blockSize, &inner1); - l.setGPUBlockIndex(inner, 0); - l.setGPUThreadIndex(inner1, 0); + inner->set_gpu_block_index(0); + inner1->set_gpu_thread_index(0); } else { throw std::runtime_error( "Invalid loop-level: " + c10::to_string(loopLevels)); @@ -3025,8 +3025,8 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { For* inner = nullptr; LoopNest::splitWithMask(flattened, blockSize, &inner); - l.setGPUBlockIndex(flattened, 0); - l.setGPUThreadIndex(inner, 0); + flattened->set_gpu_block_index(0); + inner->set_gpu_thread_index(0); l.setBufferMap(flattened, block_analysis->getBufferMap()); } } diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index 2dc1bbaa4df5c..cee97c8ac90f8 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -2241,14 +2241,6 @@ std::vector LoopNest::getLoopStmtsFor(Stmt* s) const { return result; } -void LoopNest::setGPUBlockIndex(For* f, int block_index) { - f->set_gpu_block_index(block_index); -} - -void LoopNest::setGPUThreadIndex(For* f, int thread_index) { - f->set_gpu_thread_index(thread_index); -} - void LoopNest::setBufferMap( For* f, const std::unordered_map& map) { diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h index 4157716bc2e94..74edb026a03f1 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.h +++ b/torch/csrc/jit/tensorexpr/loopnest.h @@ -339,9 +339,6 @@ class TORCH_API LoopNest { static void sliceTail(For* f, int factor, For** head, For** tail); static void sliceTail(For* f, int factor); - void setGPUBlockIndex(For* f, int idx); - void setGPUThreadIndex(For* f, int idx); - using AccessResult = std::pair; // Insert a cache for the consumer's usages of the buffer produced in // consumer, and redirect reads and writes in the consumer to that cache. diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp index 05adfd70fa9eb..c247ecb9b7202 100644 --- a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp +++ b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp @@ -326,6 +326,16 @@ void initTensorExprBindings(PyObject* module) { py::return_value_policy::reference) .def("body", &For::body, py::return_value_policy::reference) .def("set_parallel", &For::set_parallel) + .def( + "set_gpu_block_index", + [](For& self, int block_index) { + self.set_gpu_block_index(block_index); + }) + .def( + "set_gpu_thread_index", + [](For& self, int thread_index) { + self.set_gpu_thread_index(thread_index); + }) .def_static( "make", [](const VarHandle& var, @@ -552,14 +562,6 @@ void initTensorExprBindings(PyObject* module) { &LoopNest::reorderAxis, py::return_value_policy::reference) .def("simplify", &LoopNest::simplify, py::return_value_policy::reference) - .def( - "set_GPU_block_index", - &LoopNest::setGPUBlockIndex, - py::return_value_policy::reference) - .def( - "set_GPU_thread_index", - &LoopNest::setGPUThreadIndex, - py::return_value_policy::reference) .def( "inline_intermediate_bufs", [](LoopNest& self, bool allow_duplicated_work) { From 20460b0c0591759c73901d449faaa5e679015075 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Tue, 15 Jun 2021 10:34:45 -0700 Subject: [PATCH 110/305] [nnc] Removed setBufferMap method from LoopNest (#59496) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59496 Test Plan: Imported from OSS Reviewed By: ZolotukhinM Differential Revision: D28915958 Pulled By: navahgar fbshipit-source-id: 71e649c93fc67b36c37373f043c729aa835968a0 --- torch/csrc/jit/tensorexpr/kernel.cpp | 2 +- torch/csrc/jit/tensorexpr/loopnest.cpp | 6 ------ torch/csrc/jit/tensorexpr/loopnest.h | 4 ---- 3 files changed, 1 insertion(+), 11 deletions(-) diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 7b832080dadc4..0022cec9bb598 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -3027,7 +3027,7 @@ Stmt* TensorExprKernel::transformLoops(BackendType backendType, Stmt* st) { LoopNest::splitWithMask(flattened, blockSize, &inner); flattened->set_gpu_block_index(0); inner->set_gpu_thread_index(0); - l.setBufferMap(flattened, block_analysis->getBufferMap()); + flattened->set_buffer_map(block_analysis->getBufferMap()); } } diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index cee97c8ac90f8..985dee1aa9148 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -2241,12 +2241,6 @@ std::vector LoopNest::getLoopStmtsFor(Stmt* s) const { return result; } -void LoopNest::setBufferMap( - For* f, - const std::unordered_map& map) { - f->set_buffer_map(map); -} - Stmt* LoopNest::getLoopBodyFor(Tensor* t) const { return getLoopBodyFor(t->buf()); } diff --git a/torch/csrc/jit/tensorexpr/loopnest.h b/torch/csrc/jit/tensorexpr/loopnest.h index 74edb026a03f1..a49aa5a027ef1 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.h +++ b/torch/csrc/jit/tensorexpr/loopnest.h @@ -408,10 +408,6 @@ class TORCH_API LoopNest { // for the LLVM backend, when no reductions are involved. void vectorizeInnerLoops(); - void setBufferMap( - For* f, - const std::unordered_map& map); - void eliminateDeadStores(); void prepareForCodegen(); From d7eb5836bb18e4bce9d39800e56df88d96a80e5e Mon Sep 17 00:00:00 2001 From: Pritam Damania Date: Tue, 15 Jun 2021 10:48:08 -0700 Subject: [PATCH 111/305] Add RRef support to ShardedTensor. (#59776) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59776 Overall design: https://github.com/pytorch/pytorch/issues/55207. In this PR, I've added support to ShardedTensor such that it also creates RRefs pointing to the remote shards if the RPC framework is initialized. As a result, this provides more flexiblity for ShardedTensor such that users can use collectives with local shards or use the RPC framework to interact with remote shards. ghstack-source-id: 131381914 Test Plan: 1) unit tests 2) waitforbuildbot Reviewed By: SciPioneer Differential Revision: D29020844 fbshipit-source-id: acb308d0029a5e486c464d93189b5de1ba680c85 --- .../_sharded_tensor/test_sharded_tensor.py | 268 +++++++++++++++++- torch/distributed/_sharded_tensor/api.py | 143 +++++++++- 2 files changed, 392 insertions(+), 19 deletions(-) diff --git a/test/distributed/_sharded_tensor/test_sharded_tensor.py b/test/distributed/_sharded_tensor/test_sharded_tensor.py index acd751f37d356..98a4b87ec1f9e 100644 --- a/test/distributed/_sharded_tensor/test_sharded_tensor.py +++ b/test/distributed/_sharded_tensor/test_sharded_tensor.py @@ -1,5 +1,7 @@ +from functools import wraps import torch import torch.distributed as dist +from torch.distributed import rpc from torch.distributed import _sharded_tensor from torch.distributed._sharding_spec import ( ChunkShardingSpec, @@ -30,19 +32,52 @@ def init_pg(self): init_method=f"file://{self.file_name}", ) + def init_rpc(self): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions() + rpc_backend_options.init_method = f"file://{self.file_name}" + for rank in range(self.world_size): + rpc_backend_options.set_device_map(f'worker{rank}', {rank : self.rank, self.rank : rank}) + + rpc.init_rpc( + name="worker%d" % self.rank, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + def init_comms(self): + self.init_rpc() + self.init_pg() + + def destroy_comms(self): + # Wait for all ranks to reach here before starting shutdown. + dist.barrier() + + rpc.shutdown() + dist.destroy_process_group() + def setUp(self) -> None: super().setUp() self._spawn_processes() +def with_comms(func): + @wraps(func) + def wrapper(self): + self.init_comms() + func(self) + self.destroy_comms() + return wrapper + + @unittest.skipIf( TEST_WITH_ASAN, "Skip ASAN as torch + multiprocessing spawn have known issues" ) class TestShardedTensorChunked(ShardedTensorTestBase, MultiProcessTestCase): + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_complete_world_size(self): - self.init_pg() for dim in [0, -2]: spec = ChunkShardingSpec( @@ -78,10 +113,25 @@ def test_complete_world_size(self): self.assertEqual([3, 20], shard_metadata.shard_lengths) self.assertEqual(f'rank:{rank}/cuda:{rank}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + self.assertEqual(3, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual(f'rank:{rpc_rank}/cuda:{rpc_rank}', shard.metadata.placement) + if rpc_rank == 3: + self.assertEqual((1, 20), shard.tensor.size()) + else: + self.assertEqual((3, 20), shard.tensor.size()) + + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_partial_world_size(self): - self.init_pg() spec = ChunkShardingSpec( dim=0, @@ -111,10 +161,25 @@ def test_partial_world_size(self): self.assertEqual([5, 20], shard_metadata.shard_lengths) self.assertEqual(f'rank:{shard_rank + 2}/cuda:{shard_rank + 2}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + if self.rank >= 2: + self.assertEqual(1, len(remote_shards)) + else: + self.assertEqual(2, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual(f'rank:{rpc_rank}/cuda:{rpc_rank}', shard.metadata.placement) + self.assertEqual((5, 20), shard.tensor.size()) + + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_new_group(self): - self.init_pg() spec = ChunkShardingSpec( dim=0, @@ -147,11 +212,25 @@ def test_new_group(self): self.assertEqual([5, 20], shard_metadata.shard_lengths) self.assertEqual(f'rank:{shard_rank + 1}/cuda:{shard_rank + 2}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + if self.rank >= 2: + self.assertEqual(1, len(remote_shards)) + else: + self.assertEqual(2, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + for remote_shard in shards: + shard = remote_shard.to_here() + self.assertEqual(rpc_rank, remote_shard.owner().id) + self.assertEqual(f'rank:{rpc_rank - 1}/cuda:{rpc_rank}', shard.metadata.placement) + self.assertEqual((5, 20), shard.tensor.size()) + + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_multiple_local_shards(self): - self.init_pg() - spec = ChunkShardingSpec( dim=0, placements=[ @@ -183,6 +262,18 @@ def test_multiple_local_shards(self): self.assertEqual([2, 20], shard_metadata.shard_lengths) self.assertEqual(f'rank:{shard_idx % 4}/cuda:{shard_idx % 4}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + self.assertEqual(3, len(remote_shards)) + owners = {} + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(2, len(shards)) + for remote_shard in shards: + shard = remote_shard.to_here() + self.assertEqual((2, 20), shard.tensor.size()) + self.assertEqual(rpc_rank, remote_shard.owner().id) + + @skip_if_lt_x_gpu(4) @requires_nccl() def test_sharding_columns(self): @@ -255,6 +346,45 @@ def test_invalid_sharding(self): with self.assertRaisesRegex(ValueError, 'Only torch.contiguous_format memory_format is currently supported'): _sharded_tensor.empty(spec, 10, 20, memory_format=torch.channels_last) + spec = ChunkShardingSpec(dim=0, placements=["worker0/cuda:1"]) + with self.assertRaisesRegex(RuntimeError, 'RPC framework needs to be initialized'): + _sharded_tensor.empty(spec, 10, 20) + + spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"]) + with self.assertRaisesRegex(RuntimeError, 'RPC was not initialized'): + st = _sharded_tensor.empty(spec, 10, 20) + st.remote_shards + + self.init_rpc() + + # ShardedTensor was initialized before RPC. + with self.assertRaisesRegex(RuntimeError, 'RPC was not initialized'): + st.remote_shards + + spec = ChunkShardingSpec(dim=0, placements=["workerfoo/cuda:1"]) + with self.assertRaisesRegex(ValueError, 'Invalid worker name'): + _sharded_tensor.empty(spec, 10, 20) + + @skip_if_lt_x_gpu(4) + @requires_nccl() + def test_invalid_pg_rpc_ranks(self): + self.init_pg() + + # Init RPC with different ranks. + rpc_backend_options = rpc.TensorPipeRpcBackendOptions() + rpc_backend_options.init_method = f"file://{self.file_name}" + rank = (self.rank + 1) % self.world_size + rpc.init_rpc( + name=f'worker{rank}', + rank=rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + spec = ChunkShardingSpec(dim=0, placements=["rank:1/cuda:1"]) + with self.assertRaisesRegex(ValueError, 'Default ProcessGroup and RPC ranks must be the same'): + _sharded_tensor.empty(spec, 10, 20) + @skip_if_lt_x_gpu(4) @requires_nccl() def test_insufficient_sharding_dims(self): @@ -296,10 +426,10 @@ def test_insufficient_sharding_dims(self): ) class TestShardedTensorEnumerable(ShardedTensorTestBase, MultiProcessTestCase): + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_grid_sharding(self): - self.init_pg() spec = EnumerableShardingSpec([ ShardMetadata( @@ -346,6 +476,17 @@ def test_grid_sharding(self): self.assertEqual((5, 5), shard_metadata.shard_lengths) self.assertEqual(f'rank:{rank}/cuda:{rank}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + self.assertEqual(3, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual((5, 5), shard.tensor.size()) + @skip_if_lt_x_gpu(4) @requires_nccl() def test_uneven_shards(self): @@ -416,11 +557,10 @@ def verify_offsets(rank, offsets): verify_size(rank, shard_metadata.shard_lengths) self.assertEqual(f'rank:{rank}/cuda:{rank}', shard_metadata.placement) + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_partial_world_size(self): - self.init_pg() - spec = EnumerableShardingSpec([ ShardMetadata( shard_offsets=[0, 0], @@ -460,11 +600,25 @@ def test_partial_world_size(self): self.assertEqual((5, 5), shard_metadata.shard_lengths) self.assertEqual(f'rank:{rank}/cuda:{rank}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + if self.rank <= 1: + self.assertEqual(1, len(remote_shards)) + else: + self.assertEqual(2, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual((5, 5), shard.tensor.size()) + + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_new_group(self): - self.init_pg() - spec = EnumerableShardingSpec([ ShardMetadata( shard_offsets=[0, 0], @@ -502,11 +656,26 @@ def test_new_group(self): self.assertEqual((5, 5), shard_metadata.shard_lengths) self.assertEqual(f'rank:{rank * 2}/cuda:{rank * 2 + 1}', shard_metadata.placement) + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + if self.rank == 1 or self.rank == 3: + self.assertEqual(1, len(remote_shards)) + else: + self.assertEqual(2, len(remote_shards)) + + owners = {} + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual((5, 5), shard.tensor.size()) + + @with_comms @skip_if_lt_x_gpu(4) @requires_nccl() def test_multiple_local_shards(self): - self.init_pg() - spec = EnumerableShardingSpec([ ShardMetadata( shard_offsets=[0, 0], @@ -555,3 +724,78 @@ def test_multiple_local_shards(self): self.assertEqual((shard_rank // 2 * 5, (shard_rank % 2) * 5), shard_metadata.shard_offsets) self.assertEqual((5, 5), shard_metadata.shard_lengths) self.assertEqual(f'rank:{shard_rank % 2}/cuda:{shard_rank % 2}', shard_metadata.placement) + + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + if self.rank <= 1: + self.assertEqual(1, len(remote_shards)) + else: + self.assertEqual(2, len(remote_shards)) + + owners = {} + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(2, len(shards)) + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual((5, 5), shard.tensor.size()) + + @with_comms + @skip_if_lt_x_gpu(4) + @requires_nccl() + def test_with_rpc_names(self): + spec = EnumerableShardingSpec([ + ShardMetadata( + shard_offsets=[0, 0], + shard_lengths=[5, 5], + placement="worker0/cuda:0", + ), + ShardMetadata( + shard_offsets=[0, 5], + shard_lengths=[5, 5], + placement="worker1/cuda:1", + ), + ShardMetadata( + shard_offsets=[5, 0], + shard_lengths=[5, 5], + placement="worker2/cuda:2", + ), + ShardMetadata( + shard_offsets=[5, 5], + shard_lengths=[5, 5], + placement="worker3/cuda:3", + ) + ]) + + sharded_tensor = _sharded_tensor.empty(spec, 10, 10) + self.assertEqual((10, 10), sharded_tensor.size()) + self.assertEqual(1, len(sharded_tensor.local_shards())) + + # Verify local shard. + local_shard = sharded_tensor.local_shards()[0] + self.assertEqual(torch.device(f'cuda:{self.rank}'), local_shard.tensor.device) + self.assertEqual((5, 5), local_shard.tensor.size()) + + # Verify local shard metadata. + self.assertEqual((self.rank // 2 * 5, (self.rank % 2) * 5), local_shard.metadata.shard_offsets) + self.assertEqual((5, 5), local_shard.metadata.shard_lengths) + self.assertEqual(f'worker{self.rank}/cuda:{self.rank}', local_shard.metadata.placement) + + # Verify global metadata. + sharding_metadata = sharded_tensor.sharding_metadata() + self.assertEqual(4, len(sharding_metadata)) + for rank, shard_metadata in enumerate(sharding_metadata): + self.assertEqual((rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets) + self.assertEqual((5, 5), shard_metadata.shard_lengths) + self.assertEqual(f'worker{rank}/cuda:{rank}', shard_metadata.placement) + + # Validate remote shards. + remote_shards = sharded_tensor.remote_shards + self.assertEqual(3, len(remote_shards)) + + for rpc_rank, shards in remote_shards.items(): + self.assertEqual(1, len(shards)) + for remote_shard in shards: + self.assertEqual(rpc_rank, remote_shard.owner().id) + shard = remote_shard.to_here() + self.assertEqual((5, 5), shard.tensor.size()) diff --git a/torch/distributed/_sharded_tensor/api.py b/torch/distributed/_sharded_tensor/api.py index b492ae25abee2..6716e6cc9bfd7 100644 --- a/torch/distributed/_sharded_tensor/api.py +++ b/torch/distributed/_sharded_tensor/api.py @@ -1,8 +1,14 @@ from dataclasses import dataclass -from typing import List +from typing import ( + Dict, + List +) +import threading import torch import torch.distributed as dist +from torch.distributed import rpc +from torch.distributed import distributed_c10d from torch.distributed._sharding_spec import ( ChunkShardingSpec, EnumerableShardingSpec, @@ -12,6 +18,11 @@ from torch.distributed._sharding_spec._internals import is_valid_device from torch.distributed.utils import _parse_remote_device +# Tracking for sharded tensor objects. +_sharded_tensor_lock = threading.Lock() +_sharded_tensor_current_id = 0 +_sharded_tensor_map: Dict[int, 'ShardedTensor'] = {} + @dataclass class Shard(object): @@ -25,6 +36,15 @@ class Shard(object): metadata: ShardMetadata +def _register_remote_shards(sharded_tensor_id: int, rrefs: List[rpc.RRef[Shard]], rpc_rank: int): + with _sharded_tensor_lock: + if sharded_tensor_id not in _sharded_tensor_map: + raise RuntimeError( + f'Could not find sharded_tensor_id: {sharded_tensor_id} in map: {_sharded_tensor_map.keys()}') + + _sharded_tensor_map[sharded_tensor_id]._register_remote_shards(rrefs, rpc_rank) + + class ShardedTensor(object): """ ShardedTensor is an abstraction to represent Tensors that are sharded @@ -75,6 +95,19 @@ def __init__( memory_format=torch.contiguous_format, process_group=None, ): + self._rpc_initialized = False + self._sharded_tensor_id = None + if rpc._is_current_rpc_agent_set(): + # Validate PG and RPC ranks match. + pg_rank = dist.get_rank() + rpc_rank = rpc.get_worker_info().id + if pg_rank != rpc_rank: + raise ValueError( + f'Default ProcessGroup and RPC ranks must be ' + f'the same for ShardedTensor, found process group rank: ' + f'{pg_rank} and RPC rank: {rpc_rank}' + ) + if layout != torch.strided: raise ValueError('Only torch.strided layout is currently supported') @@ -86,13 +119,14 @@ def __init__( self._process_group = ( process_group if process_group is not None - else torch.distributed.distributed_c10d._get_default_group() + else distributed_c10d._get_default_group() ) - if torch.distributed.distributed_c10d._rank_not_in_group(self._process_group): + if distributed_c10d._rank_not_in_group(self._process_group): raise ValueError(f'Global rank: {dist.get_rank()} not part of process group') self._local_shards: List[Shard] = [] + self._remote_shards: Dict[int, List[rpc.RRef[Shard]]] = {} self._sharding_metadata: List[ShardMetadata] = [] if isinstance(self._sharding_spec, ChunkShardingSpec): self._init_chunked( @@ -113,6 +147,73 @@ def __init__( else: raise ValueError(f'Unsupported sharding_spec: {self._sharding_spec}') + with _sharded_tensor_lock: + global _sharded_tensor_current_id, _sharded_tensor_map + self._sharded_tensor_id = _sharded_tensor_current_id + _sharded_tensor_map[self._sharded_tensor_id] = self + _sharded_tensor_current_id += 1 + + # Initialize RPC if available. + if rpc._is_current_rpc_agent_set(): + self._init_rpc() + + def __del__(self): + # Clean up the global map. + with _sharded_tensor_lock: + global _sharded_tensor_current_id, _sharded_tensor_map + if self._sharded_tensor_id in _sharded_tensor_map: + _sharded_tensor_map.pop(self._sharded_tensor_id) + + def _init_rpc(self): + self._rpc_initialized = True + self._remote_shards = {} + + # Gather all the sharded tensor ids. + world_size = dist.get_world_size(self._process_group) + worker_infos = rpc._get_current_rpc_agent().get_worker_infos() + rank_to_name = {} + name_to_rank = {} + + for worker_info in worker_infos: + rank_to_name[worker_info.id] = worker_info.name + name_to_rank[worker_info.name] = worker_info.id + + rpc_workers = set() + for rank in range(world_size): + if self._process_group == distributed_c10d._get_default_group(): + global_rank = rank + else: + global_rank = distributed_c10d._get_global_rank(self._process_group, rank) + rpc_workers.add(rank_to_name[global_rank]) + + all_tensor_ids = rpc.api._all_gather(self._sharded_tensor_id, rpc_workers) + + # Share the local shards to the entire world. + futs = [] + rpc_rank = rpc.get_worker_info().id + for rank in range(world_size): + # Skip self. + if rank == dist.get_rank(self._process_group): + continue + + if self._process_group == distributed_c10d._get_default_group(): + global_rank = rank + else: + global_rank = distributed_c10d._get_global_rank(self._process_group, rank) + + if len(self.local_shards()) != 0: + rrefs: List[rpc.RRef[Shard]] = [rpc.RRef(shard) for shard in self.local_shards()] + fut = rpc.rpc_async( + global_rank, + _register_remote_shards, + args=(all_tensor_ids[rank_to_name[global_rank]], rrefs, rpc_rank)) + futs.append(fut) + + torch.futures.wait_all(futs) + + # Barrier for all RPCs to finish on all ranks. + rpc.api._barrier(rpc_workers) + def _init_chunked( self, dtype, @@ -208,13 +309,24 @@ def _init_enumerable( def _parse_and_validate_remote_device(self, device): - rank, local_device = _parse_remote_device(device) # type: ignore[arg-type] + on, local_device = _parse_remote_device(device) # type: ignore[arg-type] # Validate rank. - if not isinstance(rank, int) or (rank < 0 or rank >= dist.get_world_size(self._process_group)): - raise ValueError(f'Invalid rank: {rank}') + if isinstance(on, int) and (on < 0 or on >= dist.get_world_size(self._process_group)): + raise ValueError(f'Invalid rank: {on}') - return rank, local_device + if isinstance(on, str): + if not rpc._is_current_rpc_agent_set(): + raise RuntimeError(f'RPC framework needs to be initialized for using worker names: {on}') + + workers = rpc._get_current_rpc_agent().get_worker_infos() + for worker in workers: + if worker.name == on: + return worker.id, local_device + + raise ValueError(f'Invalid worker name: {on}') + + return on, local_device def sharding_spec(self) -> ShardingSpec: """ @@ -245,3 +357,20 @@ def size(self) -> torch.Size: Returns the size of the self tensor. The returned value is a subclass of tuple. """ return torch.Size(self._dims) + + def _register_remote_shards(self, remote_shards: List[rpc.RRef[Shard]], rpc_rank: int): + self._remote_shards[rpc_rank] = remote_shards + + @property + def remote_shards(self) -> Dict[int, List[rpc.RRef[Shard]]]: + """ + Returns a Dict[int, RRef] with keys being the RPC rank and values + being RRefs to shards on that rank. Need to initialize the + RPC framework for this functionality. + """ + if not self._rpc_initialized: + raise RuntimeError( + "RPC was not initialized before creating the ShardedTensor. Please initialize it using " + "torch.distributed.rpc.init_rpc before creating the ShardedTensor for remote_shards support" + ) + return self._remote_shards From 79d7c15dc54bd89293b678f860d4e96442975656 Mon Sep 17 00:00:00 2001 From: Scott Wolchok Date: Tue, 15 Jun 2021 11:24:20 -0700 Subject: [PATCH 112/305] [PyTorch] Add ExclusivelyOwned (#59419) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59419 This introduces ExclusivelyOwned, which allows isolated pieces of code that can make ownership guarantees to opt out of reference counting operations on `intrusive_ptr` and `Tensor` entirely. To elaborate, if you know you are the exclusive owner of an `intrusive_ptr` or `Tensor`, moving it into an `ExclusivelyOwned` will avoid performing atomic reference counting operations at destruction time. The documentation comment should provide sufficient explanation; please request changes if not. ghstack-source-id: 131376658 Test Plan: Added `ExclusivelyOwned_test.cpp`. It passes. When I ran it under valgrind, valgrind reported no leaks. Inspected assembly from `inspect` functions in `ExclusivelyOwned_test.cpp` in an optimized (opt-clang) build. As expected, `ExclusivelyOwned` calls `release_resources()` and the `TensorImpl` virtual destructor without including any atomic reference counting operations. Reviewed By: ezyang Differential Revision: D28885314 fbshipit-source-id: 20bf6c82b0966aaa635ab0233974781ed15f93c1 --- aten/src/ATen/templates/TensorBody.h | 54 +++++++ aten/src/ATen/test/ExclusivelyOwned_test.cpp | 133 +++++++++++++++++ c10/util/ExclusivelyOwned.h | 146 +++++++++++++++++++ c10/util/intrusive_ptr.h | 83 ++++++++++- 4 files changed, 414 insertions(+), 2 deletions(-) create mode 100644 aten/src/ATen/test/ExclusivelyOwned_test.cpp create mode 100644 c10/util/ExclusivelyOwned.h diff --git a/aten/src/ATen/templates/TensorBody.h b/aten/src/ATen/templates/TensorBody.h index 0fa56a6624efe..fa879d656ab51 100644 --- a/aten/src/ATen/templates/TensorBody.h +++ b/aten/src/ATen/templates/TensorBody.h @@ -1021,6 +1021,60 @@ struct MaybeOwnedTraits { return true; } }; + +template <> +struct ExclusivelyOwnedTraits { + using repr_type = at::Tensor; + using pointer_type = at::Tensor*; + using const_pointer_type = const at::Tensor*; + + static repr_type nullRepr() { + return at::Tensor(); + } + + template + static repr_type createInPlace(Args&&... args) { + return at::Tensor(std::forward(args)...); + } + + static repr_type moveToRepr(at::Tensor&& x) { + return std::move(x); + } + + static void destroyOwned(at::Tensor& x) { + TensorImpl*const toDestroy = x.unsafeReleaseTensorImpl(); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(toDestroy != nullptr, "Tensor somehow got null TensorImpl?"); + // May be 0 because UndefinedTensorImpl doesn't get its refcount + // incremented. + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + toDestroy->refcount_ == 1 || (toDestroy->refcount_ == 0 && toDestroy == UndefinedTensorImpl::singleton()), + "ExclusivelyOwned destroyed with refcount ", toDestroy->refcount_, ", expected 1!"); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + toDestroy->weakcount_ == 1 || (toDestroy->weakcount_ == 0 && toDestroy == UndefinedTensorImpl::singleton()), + "ExclusivelyOwned destroyed with weakcount ", toDestroy->weakcount_, ", expected 1!"); + if (toDestroy != UndefinedTensorImpl::singleton()) { +#ifndef NDEBUG + // Needed to pass the debug assertions in ~intrusive_ptr_target. + toDestroy->refcount_ = 0; + toDestroy->weakcount_ = 0; +#endif + toDestroy->release_resources(); + delete toDestroy; + } + } + + static at::Tensor take(at::Tensor& x) { + return std::move(x); + } + + static pointer_type getImpl(repr_type& x) { + return &x; + } + + static const_pointer_type getImpl(const repr_type& x) { + return &x; + } +}; } // namespace c10 namespace at { diff --git a/aten/src/ATen/test/ExclusivelyOwned_test.cpp b/aten/src/ATen/test/ExclusivelyOwned_test.cpp new file mode 100644 index 0000000000000..973359fe37e9e --- /dev/null +++ b/aten/src/ATen/test/ExclusivelyOwned_test.cpp @@ -0,0 +1,133 @@ +#include + +#include +#include +#include +#include + +#include + +namespace { + +using at::Tensor; + +struct MyString : public c10::intrusive_ptr_target, public std::string { + using std::string::string; +}; + +template +class ExclusivelyOwnedTest : public ::testing::Test { + public: + c10::ExclusivelyOwned defaultConstructed; + c10::ExclusivelyOwned sample; + protected: + void SetUp() override; // defined below helpers + void TearDown() override { + defaultConstructed = c10::ExclusivelyOwned(); + sample = c10::ExclusivelyOwned(); + } +}; + +template +T getSampleValue(); + +template <> +c10::intrusive_ptr getSampleValue() { + return c10::make_intrusive("hello"); +} + +template <> +Tensor getSampleValue() { + return at::native::zeros({2, 2}).to(at::kCPU); +} + +template +void assertIsSampleObject(const T& eo); + +template <> +void assertIsSampleObject(const MyString& s) { + EXPECT_STREQ(s.c_str(), "hello"); +} + +template <> +void assertIsSampleObject>(const c10::intrusive_ptr& s) { + assertIsSampleObject(*s); +} + +template <> +void assertIsSampleObject(const Tensor& t) { + EXPECT_EQ(t.sizes(), (c10::IntArrayRef{2, 2})); + EXPECT_EQ(t.strides(), (c10::IntArrayRef{2, 1})); + ASSERT_EQ(t.scalar_type(), at::ScalarType::Float); + static const float zeros[4] = {0}; + EXPECT_EQ(memcmp(zeros, t.data_ptr(), 4 * sizeof(float)), 0); +} + + +template +void ExclusivelyOwnedTest::SetUp() { + defaultConstructed = c10::ExclusivelyOwned(); + sample = c10::ExclusivelyOwned(getSampleValue()); +} + +using ExclusivelyOwnedTypes = ::testing::Types< + c10::intrusive_ptr, + Tensor + >; + +TYPED_TEST_CASE(ExclusivelyOwnedTest, ExclusivelyOwnedTypes); + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TYPED_TEST(ExclusivelyOwnedTest, DefaultConstructor) { + c10::ExclusivelyOwned defaultConstructed; +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TYPED_TEST(ExclusivelyOwnedTest, MoveConstructor) { + auto movedDefault = std::move(this->defaultConstructed); + auto movedSample = std::move(this->sample); + + assertIsSampleObject(*movedSample); +} + +TYPED_TEST(ExclusivelyOwnedTest, MoveAssignment) { + // Move assignment from a default-constructed ExclusivelyOwned is handled in + // TearDown at the end of every test! + c10::ExclusivelyOwned anotherSample = c10::ExclusivelyOwned(getSampleValue()); + anotherSample = std::move(this->sample); + assertIsSampleObject(*anotherSample); +} + +TYPED_TEST(ExclusivelyOwnedTest, MoveAssignmentFromContainedType) { + c10::ExclusivelyOwned anotherSample = c10::ExclusivelyOwned(getSampleValue()); + anotherSample = getSampleValue(); + assertIsSampleObject(*anotherSample); +} + +TYPED_TEST(ExclusivelyOwnedTest, Take) { + auto x = std::move(this->sample).take(); + assertIsSampleObject(x); +} + +} // namespace + +extern "C" void inspectTensor() { + auto t = getSampleValue(); +} + +extern "C" void inspectExclusivelyOwnedTensor() { + c10::ExclusivelyOwned t(getSampleValue()); +} + + +extern "C" void inspectIntrusivePtr() { + auto p = getSampleValue>(); +} + +extern "C" void inspectExclusivelyOwnedIntrusivePtr() { + c10::ExclusivelyOwned> p(getSampleValue>()); +} + +extern "C" void inspectUniquePtr() { + std::unique_ptr p(getSampleValue>().release()); +} diff --git a/c10/util/ExclusivelyOwned.h b/c10/util/ExclusivelyOwned.h new file mode 100644 index 0000000000000..46f4845285a2f --- /dev/null +++ b/c10/util/ExclusivelyOwned.h @@ -0,0 +1,146 @@ +#pragma once + +#include + +namespace c10 { + +// See example implementations in TensorBody.h and intrusive_ptr.h. +// Synopsis: +// +// repr_type -- type to use to store an owned T in ExclusivelyOwned. +// +// pointer_type -- pointer-esque type to return from +// ExclusivelyOwned's get() and operator*() methods. +// +// const_pointer_type -- similar to pointer_type, used for the const methods. +// +// static repr_type nullRepr() -- return a null instance of repr_type. +// +// template +// static repr_type createInPlace(Args&&... args) -- used by the in-place +// ExclusivelyOwned constructor. +// +// static repr_type moveToRepr(T&& x) -- move the given x into an +// instance of repr_type. used by the ExclusivelyOwned(T&&) +// constructor. +// +// static void destroyOwned(repr_type x) -- free memory for a +// known-exclusively-owned instance of x. Replaces calling repr_type's +// destructor. Being able to implement this more efficiently than +// repr_type's destructor is the main reason to use ExclusivelyOwned +// for a type. +// +// static T take(repr_type&) -- move out of the given repr_type into an owned T. +// +// static pointer_type getImpl(const repr_type&) -- return a pointer +// to the given repr_type. May take repr_type by value if that is more +// efficient. +template +struct ExclusivelyOwnedTraits; + +/// ExclusivelyOwned is a smart-pointer-like wrapper around an +/// exclusively-owned instance of some type T that normally has +/// mandatory reference counting (currently Tensor or +/// c10::intrusive_ptr). If you have an isolated piece of code that +/// knows that it has sole ownership of an object of one of these +/// types (i.e., because you created it directly or using a factory +/// function) and that object will not escape from that isolated piece +/// of code, then moving the object into an ExclusivelyOwned will +/// avoid an atomic reference count decrement at destruction time. +/// +/// If you directly create the Tensor/intrusive_ptr in the first +/// place, you can use the in_place constructor of ExclusivelyOwned to +/// additionally avoid doing any stores to initialize the refcount & +/// weakcount. (Do note, however, that in this case you should +/// probably just use std::unique_ptr instead of intrusive_ptr if applicable.) +template +class ExclusivelyOwned { + using EOT = ExclusivelyOwnedTraits; + union { + char dummy_; + typename ExclusivelyOwnedTraits::repr_type repr_; + }; + + public: + ExclusivelyOwned() : repr_(EOT::nullRepr()) {} + + explicit ExclusivelyOwned(T&& t) : repr_(EOT::moveToRepr(std::move(t))) {} + + template + explicit ExclusivelyOwned(in_place_t, Args&&... args) + : repr_(EOT::createInPlace(std::forward(args)...)) {} + + ExclusivelyOwned(const ExclusivelyOwned&) = delete; + + ExclusivelyOwned(ExclusivelyOwned&& rhs) noexcept + : repr_(std::move(rhs.repr_)) { + rhs.repr_ = EOT::nullRepr(); + } + + ExclusivelyOwned& operator=(const ExclusivelyOwned&) = delete; + + ExclusivelyOwned& operator=(ExclusivelyOwned&& rhs) noexcept { + EOT::destroyOwned(repr_); + repr_ = std::move(rhs.repr_); + rhs.repr_ = EOT::nullRepr(); + return *this; + } + + ExclusivelyOwned& operator=(T&& rhs) noexcept { + EOT::destroyOwned(repr_); + repr_ = EOT::moveToRepr(std::move(rhs)); + return *this; + } + + ~ExclusivelyOwned() { + EOT::destroyOwned(repr_); + // Don't bother to call the destructor of repr_, since we already + // did specialized destruction for the exclusively-owned case in + // destroyOwned! + } + + // We don't provide this because it would require us to be able to + // differentiate an owned-but-empty T from a lack of T. This is + // particularly problematic for Tensor, which wants to use an + // undefined Tensor as its null state. + explicit operator bool() const noexcept = delete; + + operator T() && { + return take(); + } + + // NOTE: the equivalent operation on MaybeOwned is a moving + // operator*. For ExclusivelyOwned, take() and operator*() may well + // have different return types (e.g., for intrusive_ptr, take() + // returns c10::intrusive_ptr whereas operator* returns T&), so + // they are different functions. + T take() && { + return EOT::take(repr_); + } + + typename EOT::const_pointer_type operator->() const { + return get(); + } + + typename EOT::const_pointer_type get() const { + return EOT::getImpl(repr_); + } + + typename EOT::pointer_type operator->() { + return get(); + } + + typename EOT::pointer_type get() { + return EOT::getImpl(repr_); + } + + std::remove_pointer_t& operator*() const { + return *get(); + } + + std::remove_pointer_t& operator*() { + return *get(); + } +}; + +} // namespace c10 diff --git a/c10/util/intrusive_ptr.h b/c10/util/intrusive_ptr.h index 1be2207598385..0e2e2c2e3de57 100644 --- a/c10/util/intrusive_ptr.h +++ b/c10/util/intrusive_ptr.h @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -21,6 +22,9 @@ namespace intrusive_ptr { inline void incref(intrusive_ptr_target* self); } +template +struct ExclusivelyOwnedTraits; + // constructor tag used by intrusive_ptr constructors struct DontIncreaseRefcount {}; } // namespace raw @@ -85,6 +89,9 @@ class C10_API intrusive_ptr_target { friend inline void raw::weak_intrusive_ptr::incref( intrusive_ptr_target* self); + template + friend struct ExclusivelyOwnedTraits; + protected: // protected destructor. We never want to destruct intrusive_ptr_target* // directly. @@ -226,6 +233,8 @@ class intrusive_ptr final { TTarget* target_; + template + friend struct ExclusivelyOwnedTraits; template friend class intrusive_ptr; friend class weak_intrusive_ptr; @@ -295,8 +304,8 @@ class intrusive_ptr final { TORCH_INTERNAL_ASSERT_DEBUG_ONLY( target_->refcount_ == 0 && target_->weakcount_ == 0, "intrusive_ptr: Newly-created target had non-zero refcounts. Does its " - "constructor do something strange like incref or create an intrusive_ptr" - "from `this`?"); + "constructor do something strange like incref or create an " + "intrusive_ptr from `this`?"); target_->refcount_.store(1, std::memory_order_relaxed); target_->weakcount_.store(1, std::memory_order_relaxed); } @@ -555,6 +564,76 @@ struct MaybeOwnedTraits> { } }; +template +struct ExclusivelyOwnedTraits> { + using repr_type = T*; + using pointer_type = T*; + // You can still have non-const access to the T in the const methods + // because it's not stored by value. + using const_pointer_type = T*; + + static constexpr repr_type nullRepr() { + return nullptr; + } + + template + static repr_type createInPlace(Args&&... args) { + return new T(std::forward(args)...); + } + + static repr_type moveToRepr(c10::intrusive_ptr&& x) { + return x.release(); + } + + static void destroyOwned(repr_type x) { + if (!x) { + return; + } + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + x->refcount_ == 1, + "ExclusivelyOwned> destroyed with refcount other than 1!"); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + x->weakcount_ == 1, + "ExclusivelyOwned> destroyed with weakcount other than 1!"); + const_cast*>(x)->release_resources(); +#ifndef NDEBUG + // Needed to pass the debug assertions in ~intrusive_ptr_target. + x->refcount_ = 0; + x->weakcount_ = 0; +#endif + x->release_resources(); + delete x; + } + + static c10::intrusive_ptr take(repr_type& x) { + // May need to do reference count initialization, so use the regular + // intrusive_ptr ctor. + + // Refcount would be zero if the ExclusivelyOwned was created + // in-place (so that the underlying T was never owned by an + // intrusive_ptr), and it would be 1 if it was created as an + // intrusive_ptr and then moved into the ExclusivelyOwned. + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + x->refcount_ == 1 || x->refcount_ == 0, + "take() from ExclusivelyOwned> with refcount other than 0 or 1!"); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY( + x->weakcount_ == 1 || x->weakcount_ == 0, + "take() from ExclusivelyOwned> with weakcount other than 0 or 1!"); +#ifndef NDEBUG + // Needed to pass the debug assertions in ~intrusive_ptr_target. + x->refcount_ = 0; + x->weakcount_ = 0; +#endif + auto result = c10::intrusive_ptr(x); + x = nullptr; + return result; + } + + static pointer_type getImpl(repr_type x) { + return x; + } +}; + template < typename TTarget, class NullType = detail::intrusive_target_default_null_type> From a120a12ab4d5b355e212fcc4c1f13f309792297e Mon Sep 17 00:00:00 2001 From: Jiong Gu Date: Tue, 15 Jun 2021 11:42:23 -0700 Subject: [PATCH 113/305] [Bootcamp][pytorch]Add WebIterDataPipe and ToBytesIterDataPipe to the datapipes. (#59816) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59816 Add two new DataPipes, one for getting web file urls to yield streams and one for getting streams to yield bytes. Test Plan: Add test_web_iterable_datapipe in test/test_datapipes.py. The test initiates a local http server for serving test files. Test below locally ok. 1. create and load 16M localhost file urls (each of size 10 Bytes) 2. create and load a 64GB localhost file in the unit test, for sake of testing time, disabling both stress test and large file test Imported from OSS Reviewed By: VitalyFedyunin Differential Revision: D29051186 fbshipit-source-id: f8e44491e670560bf445af96f94d98230436f396 --- test/test_datapipe.py | 143 ++++++++++++++++++ torch/utils/data/datapipes/iter/__init__.py | 31 +++- torch/utils/data/datapipes/iter/httpreader.py | 39 +++++ .../data/datapipes/iter/readlinesfromfile.py | 19 +++ torch/utils/data/datapipes/iter/tobytes.py | 24 +++ 5 files changed, 248 insertions(+), 8 deletions(-) create mode 100644 torch/utils/data/datapipes/iter/httpreader.py create mode 100644 torch/utils/data/datapipes/iter/readlinesfromfile.py create mode 100644 torch/utils/data/datapipes/iter/tobytes.py diff --git a/test/test_datapipe.py b/test/test_datapipe.py index afdfc41a4f07f..251b9cd26e0d0 100644 --- a/test/test_datapipe.py +++ b/test/test_datapipe.py @@ -1,6 +1,7 @@ import itertools import numpy as np import os +import os.path import pickle import random import sys @@ -9,10 +10,16 @@ import warnings import zipfile +import unittest from unittest import skipIf from typing import ( Any, Awaitable, Dict, Generic, Iterator, List, NamedTuple, Optional, Tuple, Type, TypeVar, Set, Union) +import http.server +import socketserver +import threading +import time +from functools import partial import torch import torch.nn as nn @@ -68,6 +75,7 @@ def create_temp_dir_and_files(): return [(temp_dir, temp_file1_name, temp_file2_name, temp_file3_name), (temp_sub_dir, temp_sub_file1_name, temp_sub_file2_name)] + class TestIterableDataPipeBasic(TestCase): def setUp(self): @@ -248,6 +256,139 @@ def test_groupbykey_iterable_datapipe(self): self.assertEqual(count, 8) +class FileLoggerSimpleHTTPRequestHandler(http.server.SimpleHTTPRequestHandler): + def __init__(self, *args, logfile=None, **kwargs): + self.__loggerHandle = None + if logfile is not None: + self.__loggerHandle = open(logfile, 'a+') + super().__init__(*args, **kwargs) + + def log_message(self, format, *args): + if self.__loggerHandle is not None: + self.__loggerHandle.write("%s - - [%s] %s\n" % + (self.address_string(), + self.log_date_time_string(), + format % args)) + return + + def finish(self): + if self.__loggerHandle is not None: + self.__loggerHandle.close() + super().finish() + + +def setUpLocalServerInThread(): + try: + Handler = partial(FileLoggerSimpleHTTPRequestHandler, logfile=None) + socketserver.TCPServer.allow_reuse_address = True + + server = socketserver.TCPServer(("", 0), Handler) + server_addr = "{host}:{port}".format(host=server.server_address[0], port=server.server_address[1]) + server_thread = threading.Thread(target=server.serve_forever) + server_thread.start() + + # Wait a bit for the server to come up + time.sleep(3) + + return (server_thread, server_addr, server) + except Exception: + raise + + +def create_temp_files_for_serving(tmp_dir, file_count, file_size, + file_url_template): + furl_local_file = os.path.join(tmp_dir, "urls_list") + with open(furl_local_file, 'w') as fsum: + for i in range(0, file_count): + f = os.path.join(tmp_dir, "webfile_test_{num}.data".format(num=i)) + + write_chunk = 1024 * 1024 * 16 + rmn_size = file_size + while rmn_size > 0: + with open(f, 'ab+') as fout: + fout.write(os.urandom(min(rmn_size, write_chunk))) + rmn_size = rmn_size - min(rmn_size, write_chunk) + + fsum.write(file_url_template.format(num=i)) + + +class TestIterableDataPipeHttp(TestCase): + __server_thread: threading.Thread + __server_addr: str + __server: socketserver.TCPServer + + @classmethod + def setUpClass(cls): + try: + (cls.__server_thread, cls.__server_addr, + cls.__server) = setUpLocalServerInThread() + except Exception as e: + warnings.warn("TestIterableDataPipeHttp could\ + not set up due to {0}".format(str(e))) + + @classmethod + def tearDownClass(cls): + try: + cls.__server.shutdown() + cls.__server_thread.join(timeout=15) + except Exception as e: + warnings.warn("TestIterableDataPipeHttp could\ + not tear down (clean up temp directory or terminate\ + local server) due to {0}".format(str(e))) + + def _http_test_base(self, test_file_size, test_file_count, timeout=None, + chunk=None): + + def _get_data_from_tuple_fn(data, *args, **kwargs): + return data[args[0]] + + with tempfile.TemporaryDirectory(dir=os.getcwd()) as tmpdir: + # create tmp dir and files for test + base_tmp_dir = os.path.basename(os.path.normpath(tmpdir)) + file_url_template = ("http://{server_addr}/{tmp_dir}/" + "/webfile_test_{num}.data\n")\ + .format(server_addr=self.__server_addr, tmp_dir=base_tmp_dir, + num='{num}') + create_temp_files_for_serving(tmpdir, test_file_count, + test_file_size, file_url_template) + + datapipe_dir_f = dp.iter.ListDirFiles(tmpdir, '*_list') + datapipe_f_lines = dp.iter.ReadLinesFromFile(datapipe_dir_f) + datapipe_line_url: IterDataPipe[str] = \ + dp.iter.Map(datapipe_f_lines, _get_data_from_tuple_fn, (1,)) + datapipe_http = dp.iter.HttpReader(datapipe_line_url, + timeout=timeout) + datapipe_tob = dp.iter.ToBytes(datapipe_http, chunk=chunk) + + for (url, data) in datapipe_tob: + self.assertGreater(len(url), 0) + self.assertRegex(url, r'^http://.+\d+.data$') + if chunk is not None: + self.assertEqual(len(data), chunk) + else: + self.assertEqual(len(data), test_file_size) + + @unittest.skip("Stress test on large amount of files skipped\ + due to the CI timing constraint.") + def test_stress_http_reader_iterable_datapipes(self): + test_file_size = 10 + # STATS: It takes about 5 hours to stress test 16 * 1024 * 1024 + # files locally + test_file_count = 1024 + self._http_test_base(test_file_size, test_file_count) + + @unittest.skip("Test on the very large file skipped\ + due to the CI timing constraint.") + def test_large_files_http_reader_iterable_datapipes(self): + # STATS: It takes about 11 mins to test a large file of 64GB locally + test_file_size = 1024 * 1024 * 128 + test_file_count = 1 + timeout = 30 + chunk = 1024 * 1024 * 8 + self._http_test_base(test_file_size, test_file_count, timeout=timeout, + chunk=chunk) + + class IDP_NoLen(IterDataPipe): def __init__(self, input_dp): super().__init__() @@ -288,9 +429,11 @@ def __len__(self) -> int: def _fake_fn(data, *args, **kwargs): return data + def _fake_filter_fn(data, *args, **kwargs): return data >= 5 + def _worker_init_fn(worker_id): random.seed(123) diff --git a/torch/utils/data/datapipes/iter/__init__.py b/torch/utils/data/datapipes/iter/__init__.py index 7b79243f25267..0fb668b805f7e 100644 --- a/torch/utils/data/datapipes/iter/__init__.py +++ b/torch/utils/data/datapipes/iter/__init__.py @@ -1,12 +1,24 @@ -from torch.utils.data.datapipes.iter.listdirfiles import ListDirFilesIterDataPipe as ListDirFiles -from torch.utils.data.datapipes.iter.loadfilesfromdisk import LoadFilesFromDiskIterDataPipe as LoadFilesFromDisk -from torch.utils.data.datapipes.iter.readfilesfromtar import ReadFilesFromTarIterDataPipe as ReadFilesFromTar -from torch.utils.data.datapipes.iter.readfilesfromzip import ReadFilesFromZipIterDataPipe as ReadFilesFromZip -from torch.utils.data.datapipes.iter.routeddecoder import RoutedDecoderIterDataPipe as RoutedDecoder +from torch.utils.data.datapipes.iter.listdirfiles import \ + ListDirFilesIterDataPipe as ListDirFiles +from torch.utils.data.datapipes.iter.loadfilesfromdisk import \ + LoadFilesFromDiskIterDataPipe as LoadFilesFromDisk +from torch.utils.data.datapipes.iter.readfilesfromtar import \ + ReadFilesFromTarIterDataPipe as ReadFilesFromTar +from torch.utils.data.datapipes.iter.readfilesfromzip import \ + ReadFilesFromZipIterDataPipe as ReadFilesFromZip +from torch.utils.data.datapipes.iter.routeddecoder import \ + RoutedDecoderIterDataPipe as RoutedDecoder +from torch.utils.data.datapipes.iter.httpreader import \ + HTTPReaderIterDataPipe as HttpReader +from torch.utils.data.datapipes.iter.tobytes import \ + ToBytesIterDataPipe as ToBytes +from torch.utils.data.datapipes.iter.readlinesfromfile import \ + ReadLinesFromFileIterDataPipe as ReadLinesFromFile # Functional DataPipe from torch.utils.data.datapipes.iter.callable import \ - (MapIterDataPipe as Map, CollateIterDataPipe as Collate, TransformsIterDataPipe as Transforms) + (MapIterDataPipe as Map, CollateIterDataPipe as Collate, + TransformsIterDataPipe as Transforms) from torch.utils.data.datapipes.iter.combining import \ (ConcatIterDataPipe as Concat, ZipIterDataPipe as Zip) from torch.utils.data.datapipes.iter.combinatorics import \ @@ -18,5 +30,8 @@ (FilterIterDataPipe as Filter) -__all__ = ['ListDirFiles', 'LoadFilesFromDisk', 'ReadFilesFromTar', 'ReadFilesFromZip', 'RoutedDecoder', 'GroupByKey', - 'Batch', 'BucketBatch', 'Collate', 'Concat', 'Filter', 'Map', 'Sampler', 'Shuffle', 'Transforms', 'Zip'] +__all__ = ['ListDirFiles', 'LoadFilesFromDisk', 'ReadFilesFromTar', + 'ReadFilesFromZip', 'RoutedDecoder', 'GroupByKey', + 'Batch', 'BucketBatch', 'Collate', 'Concat', 'Filter', 'Map', + 'Sampler', 'Shuffle', 'Transforms', 'Zip', + 'HttpReader', 'ToBytes', 'ReadLinesFromFile'] diff --git a/torch/utils/data/datapipes/iter/httpreader.py b/torch/utils/data/datapipes/iter/httpreader.py new file mode 100644 index 0000000000000..c663a18cdaab8 --- /dev/null +++ b/torch/utils/data/datapipes/iter/httpreader.py @@ -0,0 +1,39 @@ +from io import IOBase +from typing import Tuple +from urllib.error import HTTPError, URLError +import urllib.request as urllib +from torch.utils.data import IterDataPipe + + +class HTTPReaderIterDataPipe(IterDataPipe[Tuple[str, IOBase]]): + r""" :class:`HTTPReaderIterDataPipe` + + Iterable DataPipe to load file url(s) (http url(s) pointing to file(s)), + yield file url and IO stream in a tuple + args: + timeout : timeout for http request + """ + + def __init__(self, source_datapipe, timeout=None): + self.source_datapipe = source_datapipe + self.timeout = timeout + + def __iter__(self): + for furl in self.source_datapipe: + try: + if self.timeout is None: + r = urllib.urlopen(furl) + else: + r = urllib.urlopen(furl, timeout=self.timeout) + + yield(furl, r) + except HTTPError as e: + raise Exception("Could not get the file.\ + [HTTP Error] {code}: {reason}." + .format(code=e.code, reason=e.reason)) + except URLError as e: + raise Exception("Could not get the file at {url}.\ + [URL Error] {reason}." + .format(reason=e.reason, url=furl)) + except Exception: + raise diff --git a/torch/utils/data/datapipes/iter/readlinesfromfile.py b/torch/utils/data/datapipes/iter/readlinesfromfile.py new file mode 100644 index 0000000000000..c8366af3b475f --- /dev/null +++ b/torch/utils/data/datapipes/iter/readlinesfromfile.py @@ -0,0 +1,19 @@ +from typing import Tuple +from torch.utils.data import IterDataPipe + + +class ReadLinesFromFileIterDataPipe(IterDataPipe[Tuple[str, str]]): + r""" :class:`ReadLinesFromFileDataPipe` + + Iterable DataPipe to load file names as source iter data pipe + and yield filename and line(s). + """ + + def __init__(self, source_datapipe): + self.source_datapipe = source_datapipe + + def __iter__(self): + for file_name in self.source_datapipe: + with open(file_name) as file: + for line in file: + yield (file_name, line) diff --git a/torch/utils/data/datapipes/iter/tobytes.py b/torch/utils/data/datapipes/iter/tobytes.py new file mode 100644 index 0000000000000..21fd82d381bcb --- /dev/null +++ b/torch/utils/data/datapipes/iter/tobytes.py @@ -0,0 +1,24 @@ +from typing import Tuple +from torch.utils.data import IterDataPipe + + +class ToBytesIterDataPipe(IterDataPipe[Tuple[str, bytes]]): + r""" :class:`ToBytesIterDataPipe` + + Iterable DataPipe to load IO stream with label name, + and to yield bytes with label name in a tuple + args: + chunk : bytes to read from stream on each iteration. + If None, stream reads to the EOF. + """ + def __init__(self, source_datapipe, chunk=None): + self.source_datapipe = source_datapipe + self.chunk = chunk + + def __iter__(self): + for (furl, stream) in self.source_datapipe: + while True: + d = stream.read(self.chunk) + if not d: + break + yield (furl, d) From 55530e227610a1b0827250a9e091045106935ec7 Mon Sep 17 00:00:00 2001 From: BowenBao Date: Tue, 15 Jun 2021 12:21:31 -0700 Subject: [PATCH 114/305] Update Autograd Export Docs (#56594) (#59534) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59534 Update autograd export docs Test Plan: Imported from OSS Reviewed By: nikithamalgifb, ansley Differential Revision: D29046606 Pulled By: SplitInfinity fbshipit-source-id: 36057f6bdfd3e5c071dbca05d327de7952904120 Co-authored-by: neginraoof --- docs/source/onnx.rst | 95 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/docs/source/onnx.rst b/docs/source/onnx.rst index ad309cb6ba67b..7ad71687cd56d 100644 --- a/docs/source/onnx.rst +++ b/docs/source/onnx.rst @@ -872,6 +872,101 @@ The interface for specifying operator definitions is experimental; adventurous users should note that the APIs will probably change in a future interface. +Autograd Function +~~~~~~~~~~~~~~~~~ + +Autograd functions can be used to compute operation results and gradients and save the history. +More information on extending the torch.autograd engine can be found on this `page `_. +There are two ways to export a subclass of the torch.autograd.function. + +Symbolic Static Method +^^^^^^^^^^^^^^^^^^^^^^ + +You can add a symbolic static method to your function class. The symbolic method should contain a set +of ONNX operators that represent operator's behavior in ONNX. Here's an example for adding symbolic to +the your autograd function: :: + + + class MyRelu(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + ctx.save_for_backward(input) + return input.clamp(min=0) + + @staticmethod + def backward(ctx, grad_output): + input, = ctx.saved_tensors + grad_input = grad_output.clone() + grad_input[input < 0] = 0 + return grad_input + + @staticmethod + def symbolic(g, self): + return g.op("Clip", self, g.op('Constant', value_t=torch.tensor(0, dtype=torch.float))) + + class MyModel(torch.nn.Module): + def forward(self, x): + my_relu = MyRelu.apply + return my_relu(x) + + input_tensor = torch.randn(2, 3, 224, 224, requires_grad=True) + torch.onnx.export(MyModel(), input_tensor, 'test.onnx', opset_version=12) + +Export as Custom Operator +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Alternatively, you can register a custom symbolic function to export the autograd function as custom operator. +This is designed for more advanced usage, because it gives you access to more info on the export symbolic side. +Autograd functions are emitted in the IR graph as ``prim::PythonOp`` nodes. The attribute ``name`` identifies the +original module name. The ``prim::PythonOp`` node object can be accessed by argument ``n``. The example below shows +how you can access ``requires_grad`` info of the inputs and outputs of the autograd function. :: + + class MyClip(torch.autograd.Function): + @staticmethod + def forward(ctx, input, scalar): + ctx.save_for_backward(input) + return input.clamp(min=scalar) + + class MyRelu(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + ctx.save_for_backward(input) + return input.clamp(min=0) + + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.clip = MyClip.apply + self.relu = MyRelu.apply + + def forward(self, x): + h = self.clip(x, 2) + h = self.relu(h) + return h + + def symbolic_pythonop(g, n, *args, **kwargs): + # print information + print('original node: ', n) + for i, out in enumerate(n.outputs()): + print('original output {}: {}, requires grad: {}'.format(i, out, out.requiresGrad())) + import torch.onnx.symbolic_helper as sym_helper + for i, arg in enumerate(args): + print('arg {}: {}, requires grad: {}'.format(i, arg, arg.requiresGrad() if sym_helper._is_value(arg) else False)) + + name = kwargs['name'] + if name == "MyClip": + return g.op("Clip", args[0], min_f=args[1]) + elif name == "MyRelu": + return g.op("Relu", args[0]) + else: + return _unimplemented("prim::PythonOp", "unknown node kind: " + name) + + from torch.onnx import register_custom_op_symbolic + register_custom_op_symbolic('::prim_PythonOp', symbolic_pythonop, 1) + + input_tensor = torch.randn(2, 3, 224, 224, requires_grad=True) + torch.onnx.export(MyModule(), input_tensor, 'test.onnx', opset_version=12) + Custom operators ~~~~~~~~~~~~~~~~ From cd5f142af46e62d289f5c3339e03910e7e5d7663 Mon Sep 17 00:00:00 2001 From: BowenBao Date: Tue, 15 Jun 2021 12:21:31 -0700 Subject: [PATCH 115/305] fix error message for type_as (#57948) (#59535) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59535 Improve error message for type_as and add unit test. Test Plan: Imported from OSS Reviewed By: nikithamalgifb, ansley Differential Revision: D29046605 Pulled By: SplitInfinity fbshipit-source-id: 978bceeb62e4d3c68815cd5fdf160909a99d00f2 Co-authored-by: hwangdeyu --- test/onnx/test_pytorch_onnx_onnxruntime.py | 15 +++++++++++++++ torch/onnx/symbolic_opset9.py | 5 +++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index 6914444b26790..24f3f295c900c 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -5876,6 +5876,21 @@ def forward(self, input, other): model = MyModule() self.run_test(model, (x, y)) + @skipIfUnsupportedMinOpsetVersion(9) + def test_type_as(self): + class MyModule(torch.nn.Module): + def forward(self, x): + y = torch.tensor([1.0]) + return x.type_as(y) + + a = torch.tensor([True, False], dtype=torch.bool) + b = torch.randn(3, 4, dtype=torch.double) + c = torch.ones((2, 2), dtype=torch.int64) + model = MyModule() + self.run_test(model, a) + self.run_test(model, b) + self.run_test(model, c) + @skipIfUnsupportedMinOpsetVersion(9) def test_ones_bool(self): class MyModule(torch.nn.Module): diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py index c3aa444623bdd..d4514dce4fba5 100644 --- a/torch/onnx/symbolic_opset9.py +++ b/torch/onnx/symbolic_opset9.py @@ -1526,8 +1526,9 @@ def type_as(g, self, other): # We don't know the type of other, bail by emitting ATen return g.op("ATen", self, other, operator_s="type_as") else: - raise RuntimeError("Unsupported: ONNX export of type_as for tensor " - "of unknown dtype.") + raise RuntimeError('Unsupported: ONNX export of type_as for tensor ' + 'of unknown dtype. Please check if the dtype of the ' + 'parameter passed to the type_as function is correct.') @parse_args("v", "v", "i", "f") From 83450aa11d244e426c9a9d2d6d780efe056d326c Mon Sep 17 00:00:00 2001 From: BowenBao Date: Tue, 15 Jun 2021 12:21:31 -0700 Subject: [PATCH 116/305] [ONNX] Add support for torch.bernoulli() export (#57003) (#59536) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59536 Support export HuggingFace - Training DeBERTa model. Test Plan: Imported from OSS Reviewed By: nikithamalgifb, ansley Differential Revision: D29046609 Pulled By: SplitInfinity fbshipit-source-id: df87e0c6ed0f13463297bdeba73967fcf2aa37ca Co-authored-by: hwangdeyu --- test/onnx/test_pytorch_onnx_onnxruntime.py | 11 +++++++++++ torch/onnx/symbolic_opset9.py | 16 +++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index 24f3f295c900c..bf7ee661e1857 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -2400,6 +2400,17 @@ def forward(self, x): x = torch.randn(2, 3, 4) self.run_test(RandLike(), x) + def test_bernoulli(self): + class Bernoulli(torch.nn.Module): + def forward(self, x): + return torch.mul(x, torch.bernoulli(x).size(0)) + + x = torch.empty(3, 3).uniform_(0, 1) + self.run_test(Bernoulli(), x) + + x = torch.empty(2, 3, 3, dtype=torch.double).uniform_(0, 1) + self.run_test(Bernoulli(), x) + def test_reshape_different_rank(self): class ReshapeModel(torch.nn.Module): def forward(self, x): diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py index d4514dce4fba5..f23842761b1ad 100644 --- a/torch/onnx/symbolic_opset9.py +++ b/torch/onnx/symbolic_opset9.py @@ -2412,7 +2412,21 @@ def rrelu(g, input, lower, upper, training, generator): return g.op("PRelu", input, p) -@parse_args("v") +def bernoulli(g, input, generator=None, out=None): + if out is not None: + _unimplemented("Bernoulli", "out parameter is not supported for bernoulli") + if generator is not None and not sym_help._is_none(generator): + _unimplemented("Bernoulli", "generator is not supported for bernoulli") + + dtype = sym_help._try_get_scalar_type(input) + if dtype is None: + return _unimplemented("Bernoulli", "input dtype not accessible") + p = g.op('RandomUniformLike', input, high_f=1.0, low_f=0.0, dtype_i=sym_help.cast_pytorch_to_onnx[dtype]) + output = g.op('Less', p, input) + return g.op("Cast", output, to_i=sym_help.cast_pytorch_to_onnx[dtype]) + + +@parse_args('v') def log_sigmoid(g, input): p = g.op("Sigmoid", input) return g.op("Log", p) From 5d00c374ddf03ec8784aec934b7ed28abd3af76f Mon Sep 17 00:00:00 2001 From: BowenBao Date: Tue, 15 Jun 2021 12:21:31 -0700 Subject: [PATCH 117/305] [ONNX] Sum empty tensor could not be exported to ONNX successfully. (#58141) (#59537) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59537 PyTorch sum over empty tensor gives 0, while ONNX produces an error. torch.sum will be translated into onnx::ReduceSum op. Per the definition of ReduceSum, update the keepdims attribute for this scenario. Test Plan: Imported from OSS Reviewed By: nikithamalgifb, ansley Differential Revision: D29046604 Pulled By: SplitInfinity fbshipit-source-id: 6f5f3a66cb8eda8b5114b8474dda6fcdbae73469 Co-authored-by: fatcat-z --- test/onnx/test_pytorch_onnx_onnxruntime.py | 11 +++++++++++ torch/onnx/symbolic_helper.py | 7 +++++++ torch/onnx/symbolic_opset13.py | 4 ++-- torch/onnx/symbolic_opset9.py | 2 +- 4 files changed, 21 insertions(+), 3 deletions(-) diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index bf7ee661e1857..28d36e26c6337 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -9036,6 +9036,17 @@ def forward(self, x): self.run_test(M(2, 1), (x,)) self.run_test(M([-1, 3], [-2, -1]), (x,)) + def test_sum_empty_tensor(self): + class M(torch.nn.Module): + def forward(self, x): + return x[0:0].sum() + + x = torch.ones(12) + self.run_test(M(), (x,)) + + x = torch.ones(2, 0, 3) + self.run_test(M(), (x,)) + def make_test(name, base, layer, bidirectional, initial_state, variable_length, dropout, script_test_min_opset_version, **extra_kwargs): diff --git a/torch/onnx/symbolic_helper.py b/torch/onnx/symbolic_helper.py index 2af83c7604d1f..0bab2b2d220e2 100644 --- a/torch/onnx/symbolic_helper.py +++ b/torch/onnx/symbolic_helper.py @@ -738,6 +738,13 @@ def _optional_input_placeholder_tensor(g): n.setType(OptionalType.ofTensor()) return n +def _handle_reduce_dim_none(g, self, op_name): + dim_size = _get_tensor_dim_size(self, 0) + if dim_size is None or dim_size == 0: + # If input tensor is empty, according to ONNX ReduceSum definition, + # set keepdims=1 so that the resulted tensor has the same rank as the input. + return g.op(op_name, self, keepdims_i=1) + return g.op(op_name, self, keepdims_i=0) # --------------------------------------------------------------------- # ONNX operator version diff --git a/torch/onnx/symbolic_opset13.py b/torch/onnx/symbolic_opset13.py index a5e47dacb9974..7f20833571a53 100644 --- a/torch/onnx/symbolic_opset13.py +++ b/torch/onnx/symbolic_opset13.py @@ -147,9 +147,9 @@ def symbolic(g, self, dim=None, keepdim=None): self = _maybe_cast_reduce_op_input(g, self) if dim is None: # all-reduce path - return g.op(onnx_op_name, self, keepdims_i=0) + return sym_help._handle_reduce_dim_none(g, self, onnx_op_name) else: - keepdim = sym_help._get_const(keepdim, "i", "keepdim") + keepdim = sym_help._get_const(keepdim, 'i', 'keepdim') return g.op(onnx_op_name, self, dim, keepdims_i=keepdim) return symbolic diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py index f23842761b1ad..d6177d16226ae 100644 --- a/torch/onnx/symbolic_opset9.py +++ b/torch/onnx/symbolic_opset9.py @@ -368,7 +368,7 @@ def symbolic(g, self, dim=None, keepdim=None): self = _maybe_cast_reduce_op_input(g, self) if dim is None: # all-reduce path - return g.op(onnx_op_name, self, keepdims_i=0) + return sym_help._handle_reduce_dim_none(g, self, onnx_op_name) else: # dim-reduce path desc = "is" if allow_multi_dim_support else "i" From 044b519a80459f6787f6723c1c091a18b153d184 Mon Sep 17 00:00:00 2001 From: BowenBao Date: Tue, 15 Jun 2021 12:21:31 -0700 Subject: [PATCH 118/305] Symbolic for ReLu6 (#58560) (#59538) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59538 Four mealv2 models can export in torch 1.8.1, but fails when torch master introduces relu6 a few months back. Test Plan: Imported from OSS Reviewed By: nikithamalgifb, ansley Differential Revision: D29046607 Pulled By: SplitInfinity fbshipit-source-id: d9cf7050e4ac0dad892441305ffebc19ba84e2be Co-authored-by: David --- test/onnx/test_pytorch_onnx_onnxruntime.py | 15 +++++++++++++++ torch/onnx/symbolic_opset11.py | 12 ++++++++++++ torch/onnx/symbolic_opset9.py | 3 +++ 3 files changed, 30 insertions(+) diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index 28d36e26c6337..ecccccc0b704e 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -5553,6 +5553,21 @@ def forward(self, x): dynamic_axes={"x": [1, 2]}, test_with_inputs=[y]) + def test_relu6(self): + class Relu6Model(torch.nn.Module): + def __init__(self): + super(Relu6Model, self).__init__() + self.relu6 = torch.nn.ReLU6() + + def forward(self, x): + return self.relu6(x) + + x = torch.randn(2, 3, 4) * 100.0 + y = torch.randn(2, 4, 5) * 100.0 + self.run_test(Relu6Model(), x, input_names=['x'], + dynamic_axes={'x': [1, 2]}, + test_with_inputs=[y]) + def test_silu(self): class SiLUModel(torch.nn.Module): def __init__(self): diff --git a/torch/onnx/symbolic_opset11.py b/torch/onnx/symbolic_opset11.py index c0aba4fb8b359..bc1e1166887d6 100644 --- a/torch/onnx/symbolic_opset11.py +++ b/torch/onnx/symbolic_opset11.py @@ -75,6 +75,18 @@ def clamp_max(g, self, max): return g.op("Min", self, max) +def relu6(g, input): + relu = g.op("Relu", input) + dtype = input.type().scalarType() + if dtype is None: + dtype = 6 # float + else: + dtype = sym_help.scalar_type_to_onnx.index(sym_help.cast_pytorch_to_onnx[dtype]) + min_val = g.op("Constant", value_t=torch.tensor(0, dtype=sym_help.scalar_type_to_pytorch_type[dtype])) + max_val = g.op("Constant", value_t=torch.tensor(6, dtype=sym_help.scalar_type_to_pytorch_type[dtype])) + return clamp(g, relu, min_val, max_val) + + # Opset 11 gather accepts negative indices @parse_args("v", "i", "v") def select(g, self, dim, index): diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py index d6177d16226ae..22e3eaa4b57b8 100644 --- a/torch/onnx/symbolic_opset9.py +++ b/torch/onnx/symbolic_opset9.py @@ -729,6 +729,9 @@ def mish(g, input): def relu(g, input): return g.op("Relu", input) +def relu6(g, input): + relu = g.op("Relu", input) + return clamp_max(g, relu, 6) def ceil(g, input): return g.op("Ceil", input) From c23624351a0d8268910e3e3b05de33c8225799c6 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 15 Jun 2021 12:33:52 -0700 Subject: [PATCH 119/305] disable test_sparse_allreduce_basics (#60029) Summary: This test will be disabled due to intermittent failures in https://circleci.com/gh/pytorch/pytorch/14155828?utm_campaign=vcs-integration-link&utm_medium=referral&utm_source=github-build-link as per https://hud.pytorch.org/build2/pytorch-master Pull Request resolved: https://github.com/pytorch/pytorch/pull/60029 Reviewed By: seemethere Differential Revision: D29139042 Pulled By: Krovatkin fbshipit-source-id: 105000e8636f17846be31f517abdf56ea0a994e9 --- test/distributed/test_c10d_gloo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py index a4c2b855f2fc7..32f049f084c64 100644 --- a/test/distributed/test_c10d_gloo.py +++ b/test/distributed/test_c10d_gloo.py @@ -640,6 +640,7 @@ def _test_sparse_allreduce_basics(self, fn): self.assertEqual(tensors, outputs) self.assertEqual(result, outputs) + @unittest.skip("intermittent failures on Windows, in CI") def test_sparse_allreduce_basics(self): self._test_sparse_allreduce_basics(lambda t: t) From a0e62c4da4ba82efb5b354737f9ac0f12b1d2da5 Mon Sep 17 00:00:00 2001 From: Ailing Zhang Date: Tue, 15 Jun 2021 12:58:35 -0700 Subject: [PATCH 120/305] Reuse run_torch_xla_tests from pytorch/xla (#59888) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59888 Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29114274 Pulled By: ailzhang fbshipit-source-id: d2845c7fc95d038cd68c10e22b68be8ad3cae736 --- .jenkins/pytorch/test.sh | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 3bce691f8cf85..7a1ca1cba8b95 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -37,8 +37,6 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then # mainly used so that we're not spending extra cycles testing cpu # devices on expensive gpu machines export PYTORCH_TESTING_DEVICE_ONLY_FOR="cuda" -elif [[ "$BUILD_ENVIRONMENT" == *xla* ]]; then - export PYTORCH_TESTING_DEVICE_ONLY_FOR="xla" fi if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then @@ -339,23 +337,9 @@ test_torch_function_benchmark() { } test_xla() { - export XLA_USE_XRT=1 XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0" - # Issue #30717: randomize the port of XLA/gRPC workers is listening on to reduce flaky tests. - XLA_PORT=$(shuf -i 40701-40999 -n 1) - export XRT_WORKERS="localservice:0;grpc://localhost:$XLA_PORT" - pushd xla - echo "Running Python Tests" - ./test/run_tests.sh - - # Disabled due to MNIST download issue. - # See https://github.com/pytorch/pytorch/issues/53267 - # echo "Running MNIST Test" - # python test/test_train_mnist.py --tidy - - echo "Running C++ Tests" - pushd test/cpp - CC=clang-9 CXX=clang++-9 ./run_tests.sh - popd + # shellcheck disable=SC1091 + source "./xla/.circleci/common.sh" + run_torch_xla_tests "$(pwd)" "$(pwd)/xla" assert_git_not_dirty } From b162d95e461a5ea22f6840bf492a5dbb2ebbd151 Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Tue, 15 Jun 2021 13:13:43 -0700 Subject: [PATCH 121/305] Fix a number of lint perf and safety issues in torch (#59897) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59897 Test Plan: Sandcastle Reviewed By: ngimel Differential Revision: D29037012 fbshipit-source-id: 7c16286d5fc2b67964fb65f8374dfff4d1a7aefb --- torch/csrc/CudaIPCTypes.cpp | 3 +- torch/csrc/CudaIPCTypes.h | 5 +-- torch/csrc/DynamicTypes.cpp | 11 +++-- torch/csrc/Exceptions.h | 3 +- .../include/torch/nn/functional/activation.h | 43 ++++++++++--------- torch/csrc/api/include/torch/nn/utils/rnn.h | 3 +- torch/csrc/api/src/nn/modules/_functions.cpp | 6 --- torch/csrc/api/src/optim/adamw.cpp | 3 +- torch/csrc/autograd/FunctionsManual.cpp | 23 ++++------ torch/csrc/autograd/FunctionsManual.h | 2 +- torch/csrc/autograd/VariableTypeUtils.h | 3 +- torch/csrc/autograd/custom_function.h | 1 - torch/csrc/autograd/python_function.cpp | 2 - torch/csrc/autograd/variable.cpp | 3 +- torch/csrc/copy_utils.h | 5 +-- torch/csrc/cuda/comm.cpp | 8 +--- torch/csrc/deploy/example/benchmark.cpp | 3 +- .../autograd/engine/dist_engine.cpp | 2 - torch/csrc/jit/codegen/cuda/ir_cloner.h | 2 +- torch/csrc/jit/frontend/edit_distance.cpp | 3 +- torch/csrc/jit/frontend/lexer.h | 3 +- torch/csrc/jit/frontend/sugared_value.cpp | 3 +- .../csrc/jit/passes/constant_propagation.cpp | 3 +- .../csrc/jit/passes/frozen_ops_to_mkldnn.cpp | 3 +- .../passes/onnx/fixup_onnx_controlflow.cpp | 4 -- .../jit/passes/onnx/shape_type_inference.cpp | 38 +++++++--------- .../quantization/quantization_patterns.h | 16 ++++--- .../passes/utils/check_alias_annotation.cpp | 3 +- torch/csrc/jit/python/pybind_utils.h | 2 +- torch/csrc/jit/serialization/python_print.cpp | 7 ++- torch/csrc/jit/tensorexpr/cuda_codegen.cpp | 6 --- torch/csrc/jit/tensorexpr/eval.cpp | 3 +- torch/csrc/jit/tensorexpr/expr.cpp | 2 - torch/csrc/jit/tensorexpr/hash_provider.h | 3 +- torch/csrc/jit/tensorexpr/ir_simplifier.cpp | 8 ++-- torch/csrc/jit/tensorexpr/kernel.cpp | 3 +- torch/csrc/python_dimname.cpp | 2 +- torch/csrc/serialization.cpp | 6 +-- torch/csrc/utils.cpp | 4 +- torch/csrc/utils/invalid_arguments.cpp | 2 +- torch/csrc/utils/pybind.h | 2 +- torch/csrc/utils/python_arg_parser.cpp | 6 +-- torch/csrc/utils/python_arg_parser.h | 16 +++---- torch/csrc/utils/throughput_benchmark-inl.h | 2 +- 44 files changed, 112 insertions(+), 169 deletions(-) diff --git a/torch/csrc/CudaIPCTypes.cpp b/torch/csrc/CudaIPCTypes.cpp index dbda3580622e0..ebdcdbbd89990 100644 --- a/torch/csrc/CudaIPCTypes.cpp +++ b/torch/csrc/CudaIPCTypes.cpp @@ -132,8 +132,7 @@ void ReturnRefCounter(const std::string& handle, uint64_t offset /* unused */) { // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) CudaIPCSentData::CudaIPCSentData( - // NOLINTNEXTLINE(modernize-pass-by-value) - std::string handle, + const std::string& handle, int64_t offset, int64_t* counter_ptr, at::Device device) diff --git a/torch/csrc/CudaIPCTypes.h b/torch/csrc/CudaIPCTypes.h index 66dd0bf0680cb..d0048b40ce89d 100644 --- a/torch/csrc/CudaIPCTypes.h +++ b/torch/csrc/CudaIPCTypes.h @@ -29,7 +29,7 @@ struct CudaIPCSentData final { at::Device device_; CudaIPCSentData( - std::string handle, + const std::string& handle, int64_t offset, int64_t* counter_ptr, at::Device device); @@ -78,8 +78,7 @@ struct CudaIPCSentDataLimbo final { struct CudaIPCRefCountersFile final { CudaIPCRefCountersFile( - // NOLINTNEXTLINE(modernize-pass-by-value) - std::string handle, + const std::string& handle, uint64_t size, at::DataPtr data_ptr) : next_offset_(0), diff --git a/torch/csrc/DynamicTypes.cpp b/torch/csrc/DynamicTypes.cpp index 5bb3224a2f64f..8535b8b6f3c70 100644 --- a/torch/csrc/DynamicTypes.cpp +++ b/torch/csrc/DynamicTypes.cpp @@ -12,6 +12,7 @@ #include +#include #include #include #include @@ -30,13 +31,11 @@ std::unordered_map attype_to_py_st // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) std::unordered_map py_storage_type_to_attype; -// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables,cppcoreguidelines-avoid-c-arrays) -THPDtype* dtype_registry - [static_cast(at::ScalarType::NumOptions)] = {}; +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +std::array(at::ScalarType::NumOptions)> dtype_registry = {}; -// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables,cppcoreguidelines-avoid-c-arrays) -THPLayout* layout_registry - [static_cast(at::Layout::NumOptions)] = {}; +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +std::array(at::Layout::NumOptions)> layout_registry = {}; at::Backend get_backend(bool is_cuda, bool is_sparse) { if (is_cuda) { diff --git a/torch/csrc/Exceptions.h b/torch/csrc/Exceptions.h index f6072c878b233..4bdb5014b2054 100644 --- a/torch/csrc/Exceptions.h +++ b/torch/csrc/Exceptions.h @@ -294,8 +294,7 @@ struct ValueError : public PyTorchError { // Translates to Python NotImplementedError struct NotImplementedError : public PyTorchError { - // NOLINTNEXTLINE(modernize-use-equals-default) - NotImplementedError() {} + NotImplementedError() = default; PyObject* python_type() override { return PyExc_NotImplementedError; } diff --git a/torch/csrc/api/include/torch/nn/functional/activation.h b/torch/csrc/api/include/torch/nn/functional/activation.h index a0487c61835ef..42ade6ddcb879 100644 --- a/torch/csrc/api/include/torch/nn/functional/activation.h +++ b/torch/csrc/api/include/torch/nn/functional/activation.h @@ -690,26 +690,29 @@ inline std::tuple multi_head_attention_forward( v = F::linear(value, _w, _b); } } else { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto q_proj_weight_non_opt = q_proj_weight; - auto sizes = q_proj_weight_non_opt.sizes(); - auto len1 = sizes[0]; - auto len2 = sizes[1]; - TORCH_CHECK(len1 == embed_dim && len2 == query.size(-1)); - - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto k_proj_weight_non_opt = k_proj_weight; - sizes = k_proj_weight_non_opt.sizes(); - len1 = sizes[0]; - len2 = sizes[1]; - TORCH_CHECK(len1 == embed_dim && len2 == key.size(-1)); - - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto v_proj_weight_non_opt = v_proj_weight; - sizes = v_proj_weight_non_opt.sizes(); - len1 = sizes[0]; - len2 = sizes[1]; - TORCH_CHECK(len1 == embed_dim && len2 == value.size(-1)); + const auto& q_proj_weight_non_opt = q_proj_weight; + { + const auto sizes = q_proj_weight_non_opt.sizes(); + const auto len1 = sizes[0]; + const auto len2 = sizes[1]; + TORCH_CHECK(len1 == embed_dim && len2 == query.size(-1)); + } + + const auto& k_proj_weight_non_opt = k_proj_weight; + { + const auto sizes = k_proj_weight_non_opt.sizes(); + const auto len1 = sizes[0]; + const auto len2 = sizes[1]; + TORCH_CHECK(len1 == embed_dim && len2 == key.size(-1)); + } + + const auto& v_proj_weight_non_opt = v_proj_weight; + { + const auto sizes = v_proj_weight_non_opt.sizes(); + const auto len1 = sizes[0]; + const auto len2 = sizes[1]; + TORCH_CHECK(len1 == embed_dim && len2 == value.size(-1)); + } if (in_proj_bias.defined()) { q = F::linear(query, q_proj_weight_non_opt, in_proj_bias.slice(/*dim=*/0, 0, embed_dim)); diff --git a/torch/csrc/api/include/torch/nn/utils/rnn.h b/torch/csrc/api/include/torch/nn/utils/rnn.h index 6bcff2f2e2d1c..e6bcf51821c5c 100644 --- a/torch/csrc/api/include/torch/nn/utils/rnn.h +++ b/torch/csrc/api/include/torch/nn/utils/rnn.h @@ -242,8 +242,7 @@ inline std::tuple pad_packed_sequence( Tensor padded_output, lengths; std::tie(padded_output, lengths) = torch::_pad_packed_sequence( sequence.data(), sequence.batch_sizes(), batch_first, padding_value, max_seq_length); - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - Tensor unsorted_indices = sequence.unsorted_indices(); + const Tensor& unsorted_indices = sequence.unsorted_indices(); if (unsorted_indices.defined()) { int64_t batch_dim = batch_first ? 0 : 1; return std::make_tuple(padded_output.index_select(batch_dim, unsorted_indices), lengths.index({unsorted_indices})); diff --git a/torch/csrc/api/src/nn/modules/_functions.cpp b/torch/csrc/api/src/nn/modules/_functions.cpp index bb357181fc77e..1a81b44afc7d7 100644 --- a/torch/csrc/api/src/nn/modules/_functions.cpp +++ b/torch/csrc/api/src/nn/modules/_functions.cpp @@ -23,13 +23,7 @@ Variable CrossMapLRN2d::forward( torch::Tensor output = torch::empty({0}, input.options()); - // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores,clang-diagnostic-unused-variable) - int64_t batch_size = input.size(0); int64_t channels = input.size(1); - // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores,clang-diagnostic-unused-variable) - int64_t input_height = input.size(2); - // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores,clang-diagnostic-unused-variable) - int64_t input_width = input.size(3); output.resize_as_(input); ctx->saved_data["scale"].toTensor().resize_as_(input); diff --git a/torch/csrc/api/src/optim/adamw.cpp b/torch/csrc/api/src/optim/adamw.cpp index 52a678f66cbbe..61dd7f8a25e85 100644 --- a/torch/csrc/api/src/optim/adamw.cpp +++ b/torch/csrc/api/src/optim/adamw.cpp @@ -82,8 +82,7 @@ Tensor AdamW::step(LossClosure closure) { if (!p.grad().defined()) { continue; } - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto grad = p.grad(); + const auto& grad = p.grad(); TORCH_CHECK(!grad.is_sparse(), "AdamW does not support sparse gradients"/*, please consider SparseAdamW instead*/); auto param_state = state_.find(c10::guts::to_string(p.unsafeGetTensorImpl())); auto& options = static_cast(group.options()); diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp index f4b57ccf2c337..e738d21bccc4a 100644 --- a/torch/csrc/autograd/FunctionsManual.cpp +++ b/torch/csrc/autograd/FunctionsManual.cpp @@ -536,8 +536,7 @@ Tensor logcumsumexp_backward(Tensor grad, const Tensor & self, Tensor result, in Tensor unbind_backward(const variable_list& grads, int64_t dim) { IntArrayRef sizes; at::TensorOptions o; - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto v : grads) { + for (const auto& v : grads) { if (v.defined()) { sizes = v.sizes(); o = static_cast(v).options(); @@ -1216,10 +1215,8 @@ Tensor log_sigmoid_double_backward(const Tensor & grad, const Tensor & input) { } Tensor softmax_double_backward(const Tensor & grad, const Tensor & grad_output, int dim, const Tensor & output) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto gO = grad_output; - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto ggI = grad; + const auto& gO = grad_output; + const auto& ggI = grad; auto ggI_output = ggI * output; auto ggI_out_sum = ggI_output.sum(dim, true); @@ -2201,8 +2198,7 @@ Tensor eig_backward(const std::vector &grads, const T // variable names correspond to the ones in the reference document auto D = eigenvalues; - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto U = eigenvectors; + const auto& U = eigenvectors; auto D_grad = grads[0]; auto U_grad = grads[1]; @@ -3316,9 +3312,8 @@ Tensor _cudnn_ctc_loss_backward(const Tensor& grad_out, const Tensor& loss, cons } } -bool any_variable_defined(variable_list& variables) { - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto variable : variables) { +bool any_variable_defined(const variable_list& variables) { + for (const auto& variable : variables) { if (variable.defined()) { return true; } @@ -3383,9 +3378,9 @@ std::tuple householder_product_backward(const Tensor& grad, cons auto start_j = tau.size(-1) - 1; for (int64_t j = start_j; j >= 0; j--) { - auto v = input_.index({"...", Slice(), j}); - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto v1 = v, v2 = v; + const auto v = input_.index({"...", Slice(), j}); + const auto& v1 = v; + const auto& v2 = v; // we need to recompute input[j] * at::outer(v, v) auto tau_unsqueezed = tau.index({"...", j}).unsqueeze(-1); // tau[..., j][:, None] diff --git a/torch/csrc/autograd/FunctionsManual.h b/torch/csrc/autograd/FunctionsManual.h index 94d86bbd55cfb..dfc7836a2d7db 100644 --- a/torch/csrc/autograd/FunctionsManual.h +++ b/torch/csrc/autograd/FunctionsManual.h @@ -37,7 +37,7 @@ Tensor toNonOptFwGrad(const c10::optional& t); Tensor toNonOptPrimal(const c10::optional& t); Tensor toNonOptTensor(const c10::optional& t); -bool any_variable_defined(variable_list& variables); +bool any_variable_defined(const variable_list& variables); void copy_range(variable_list& out, IndexRange range, const at::Tensor & t); void copy_range(variable_list& out, IndexRange range, at::ArrayRef t); at::Tensor copysign_tensor_self_backward(const Tensor & grad, const Tensor & self, const Tensor & result); diff --git a/torch/csrc/autograd/VariableTypeUtils.h b/torch/csrc/autograd/VariableTypeUtils.h index d4db6eb12c692..dd4200360d0c8 100644 --- a/torch/csrc/autograd/VariableTypeUtils.h +++ b/torch/csrc/autograd/VariableTypeUtils.h @@ -85,8 +85,7 @@ inline void throw_error_for_complex_autograd(const Tensor& tensor, const char* n } inline void throw_error_for_complex_autograd(const TensorList& tensorlist, const char* name) { - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto tensor: tensorlist) { + for (const auto& tensor: tensorlist) { throw_error_for_complex_autograd(tensor, name); } } diff --git a/torch/csrc/autograd/custom_function.h b/torch/csrc/autograd/custom_function.h index cec33d05029c7..be39da30f9723 100644 --- a/torch/csrc/autograd/custom_function.h +++ b/torch/csrc/autograd/custom_function.h @@ -331,7 +331,6 @@ variable_list CppNode::apply(variable_list&& inputs) { variable_list results; results.reserve(num_outputs); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) for (const auto i : c10::irange(num_outputs)) { if (!is_variable_input_[i]) { if (outputs[i].defined()) { diff --git a/torch/csrc/autograd/python_function.cpp b/torch/csrc/autograd/python_function.cpp index ac379b5a9dc9d..50f0453ad0711 100644 --- a/torch/csrc/autograd/python_function.cpp +++ b/torch/csrc/autograd/python_function.cpp @@ -114,8 +114,6 @@ auto PyNode::apply(variable_list&& inputs) -> variable_list { // Massage the Python results tuple back into a C++ variable_list variable_list results; results.reserve(num_outputs); - // NOLINTNEXTLINE(clang-diagnostic-unused-variable) - auto& input_info = py_fn->input_info; for (int i = 0; i != num_outputs; ++i) { PyObject* output = PyTuple_GET_ITEM(r.get(), i); bool was_variable = is_variable_input[i]; diff --git a/torch/csrc/autograd/variable.cpp b/torch/csrc/autograd/variable.cpp index abbdd7880c0f1..3a95803e79242 100644 --- a/torch/csrc/autograd/variable.cpp +++ b/torch/csrc/autograd/variable.cpp @@ -175,8 +175,7 @@ namespace impl { std::unique_ptr hook_ptr(new CppFunctionPreHook(list, self.output_nr())); clear_hooks(self); add_hook(self, std::make_shared(list, 0)); - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto fn = self.grad_fn(); + const auto& fn = self.grad_fn(); if (fn) { fn->add_pre_hook(std::move(hook_ptr)); } diff --git a/torch/csrc/copy_utils.h b/torch/csrc/copy_utils.h index 09d6bf60f1b4d..f04e63368ef28 100644 --- a/torch/csrc/copy_utils.h +++ b/torch/csrc/copy_utils.h @@ -42,10 +42,9 @@ inline PyObject * THPStorageCopyMethod(const THPCopyList& v, PyObject *self, PyO // NOLINTNEXTLINE(cppcoreguidelines-init-variables) PyObject *src; int non_blocking = 0; - // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays,clang-diagnostic-writable-strings) - static char *kwlist[] = {"source", "non_blocking", nullptr}; + static std::array kwlist = {"source", "non_blocking", nullptr}; // use int as parse type because bool not available in python2. - if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|i:copy_", kwlist, &src, &non_blocking)) { + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|i:copy_", kwlist.data(), &src, &non_blocking)) { return nullptr; } diff --git a/torch/csrc/cuda/comm.cpp b/torch/csrc/cuda/comm.cpp index 46c85ee1dc1fd..f450257cc4152 100644 --- a/torch/csrc/cuda/comm.cpp +++ b/torch/csrc/cuda/comm.cpp @@ -190,11 +190,9 @@ tensor_list2d broadcast_coalesced( auto& device_outputs = outputs[i]; auto& inds = broadcast_indices[i]; auto& vals = broadcast_values[i]; - for (auto& t : + for (const auto& var : utils::unflatten_sparse_tensors(inds, vals, chunk.tensors)) { // See NOTE [ Version Counter in comm.*_coalesced ] - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - Variable var = t; device_outputs.push_back(make_variable(var.tensor_data(), false)); } } @@ -204,11 +202,9 @@ tensor_list2d broadcast_coalesced( for (size_t i = 1, num_devices = devices.size(); i < num_devices; ++i) { device_guard.set_index(devices[i]); auto& device_outputs = outputs[i]; - for (auto& t : + for (auto& var : utils::unflatten_dense_tensors(results[i], chunk.tensors)) { // See NOTE [ Version Counter in comm.*_coalesced ] - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - Variable var = t; device_outputs.push_back(make_variable(var.tensor_data(), false)); } } diff --git a/torch/csrc/deploy/example/benchmark.cpp b/torch/csrc/deploy/example/benchmark.cpp index 919e6a70b831f..af3be7d6f5729 100644 --- a/torch/csrc/deploy/example/benchmark.cpp +++ b/torch/csrc/deploy/example/benchmark.cpp @@ -146,8 +146,7 @@ struct RunJIT { } void operator()(int i) { if (cuda) { - // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions) - int device_id = i % models_.size(); + const auto device_id = i % models_.size(); auto d = torch::Device(torch::DeviceType::CUDA, device_id); to_device( models_[device_id].forward(to_device_vec(eg_, d)), diff --git a/torch/csrc/distributed/autograd/engine/dist_engine.cpp b/torch/csrc/distributed/autograd/engine/dist_engine.cpp index 6c9feda380747..50c24458017b6 100644 --- a/torch/csrc/distributed/autograd/engine/dist_engine.cpp +++ b/torch/csrc/distributed/autograd/engine/dist_engine.cpp @@ -619,8 +619,6 @@ size_t DistEngine::numBackwardPasses() const { std::unordered_map DistEngine::getDebugInfo() const { std::unordered_map debugInfo; - // NOLINTNEXTLINE(clang-diagnostic-unused-variable) - auto& DistAutogradContainer = DistAutogradContainer::getInstance(); debugInfo[kNumBackwardPasses] = numBackwardPasses(); debugInfo[kNumAutogradContexts] = DistAutogradContainer::getInstance().numAutogradContexts(); diff --git a/torch/csrc/jit/codegen/cuda/ir_cloner.h b/torch/csrc/jit/codegen/cuda/ir_cloner.h index 61baef3290cd4..41c038e4772c7 100644 --- a/torch/csrc/jit/codegen/cuda/ir_cloner.h +++ b/torch/csrc/jit/codegen/cuda/ir_cloner.h @@ -31,8 +31,8 @@ class TORCH_CUDA_CU_API IrCloner : private OptInConstDispatch { template std::vector clone(const std::vector& container) { std::vector copy; + copy.reserve(container.size()); for (auto p : container) { - // NOLINTNEXTLINE(performance-inefficient-vector-operation) copy.push_back(clone(p)); } return copy; diff --git a/torch/csrc/jit/frontend/edit_distance.cpp b/torch/csrc/jit/frontend/edit_distance.cpp index b7a6b76bf95f6..1931c503ce81e 100644 --- a/torch/csrc/jit/frontend/edit_distance.cpp +++ b/torch/csrc/jit/frontend/edit_distance.cpp @@ -36,8 +36,7 @@ size_t ComputeEditDistance( unsigned previous = y - 1; for (size_t x = 1; x <= n; ++x) { - // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions) - int old_row = row[x]; + const auto old_row = row[x]; row[x] = std::min( previous + (word1[y - 1] == word2[x - 1] ? 0u : 1u), std::min(row[x - 1], row[x]) + 1); diff --git a/torch/csrc/jit/frontend/lexer.h b/torch/csrc/jit/frontend/lexer.h index eb89c6340c61f..c23023732e616 100644 --- a/torch/csrc/jit/frontend/lexer.h +++ b/torch/csrc/jit/frontend/lexer.h @@ -474,8 +474,7 @@ struct Lexer { break; case TK_WHITESPACE: case TK_WHITESPACE_EOF: { - int depth = - // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions) + const auto depth = r.kind == TK_WHITESPACE_EOF ? indent_stack.front() : r.range.size(); // note: TK_WHITESPACE_EOF is whitespace right before the EOF token // just like we allow the code to be indented to a particular initial diff --git a/torch/csrc/jit/frontend/sugared_value.cpp b/torch/csrc/jit/frontend/sugared_value.cpp index 7aa932c15dc82..bfbb13b5bcac6 100644 --- a/torch/csrc/jit/frontend/sugared_value.cpp +++ b/torch/csrc/jit/frontend/sugared_value.cpp @@ -536,8 +536,7 @@ SugaredValuePtr RangeValue::getitem( std::vector IterableTree::get_base_iterables() { std::vector base_iters{}; - // NOLINTNEXTLINE(performance-for-range-copy) - for (SugaredValuePtr sv : children_) { + for (SugaredValuePtr& sv : children_) { if (auto iv = std::dynamic_pointer_cast(sv)) { std::vector child_iters = iv->get_base_iterables(); // merge child iters with the base_iters diff --git a/torch/csrc/jit/passes/constant_propagation.cpp b/torch/csrc/jit/passes/constant_propagation.cpp index ec8c79fe93d5b..323e747fbd87e 100644 --- a/torch/csrc/jit/passes/constant_propagation.cpp +++ b/torch/csrc/jit/passes/constant_propagation.cpp @@ -86,8 +86,7 @@ c10::optional> runNodeIfInputsAreConstant( for (const IValue& v : stack) { if (v.isTensor()) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - at::Tensor t = v.toTensor(); + const at::Tensor& t = v.toTensor(); if (t.defined() && t.requires_grad()) { // requires grad tensors cannot be constants return c10::nullopt; diff --git a/torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp b/torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp index e59bc29b83e7e..b4acab100fdcf 100644 --- a/torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp +++ b/torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp @@ -85,8 +85,7 @@ void assertNonTensorTypeDoesNotContainTensors(TypePtr type) { if (type->cast()) { return; } - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto t : type->containedTypes()) { + for (const auto& t : type->containedTypes()) { TORCH_INTERNAL_ASSERT(!t->cast()); } } diff --git a/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp b/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp index b7085c54cd0f2..d87cccba978e1 100644 --- a/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp +++ b/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp @@ -119,8 +119,6 @@ std::vector ConvertSequenceDependencies(Node* node, int opset_version) { } auto* loop_node = node; - // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores,clang-diagnostic-unused-variable) - auto* graph = loop_node->owningGraph(); TORCH_INTERNAL_ASSERT(loop_node->blocks().size() == 1); auto* sub_block = loop_node->blocks()[0]; @@ -379,8 +377,6 @@ std::vector FixupONNXIfNode(Node* node, int opset_version) { } GRAPH_DUMP("Graph before fixing controlflow: ", node->owningGraph()); auto* if_node = node; - // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores,clang-diagnostic-unused-variable) - auto* graph = if_node->owningGraph(); FixupONNXSubblockOutputs(node); ONNXFixupUninitializedOutput(if_node); // Copy type of block output to node output. diff --git a/torch/csrc/jit/passes/onnx/shape_type_inference.cpp b/torch/csrc/jit/passes/onnx/shape_type_inference.cpp index 0e4ee107ce91a..c81946ba67fa9 100644 --- a/torch/csrc/jit/passes/onnx/shape_type_inference.cpp +++ b/torch/csrc/jit/passes/onnx/shape_type_inference.cpp @@ -94,8 +94,7 @@ TensorTypePtr TorchTensorTypeFromONNX( {}); if (onnx_tensor_type.has_shape()) { std::vector sizes; - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto onnx_shape = onnx_tensor_type.shape(); + const auto& onnx_shape = onnx_tensor_type.shape(); for (int i = 0; i < onnx_shape.dim_size(); ++i) { auto& dim = onnx_shape.dim(i); @@ -108,8 +107,7 @@ TensorTypePtr TorchTensorTypeFromONNX( // Search if this is already known, // and assign the same Symbol. GRAPH_UPDATE("Got dim_param:", dim.dim_param()); - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto pair : symbol_map) { + for (const auto& pair : symbol_map) { if (pair.second == dim.dim_param()) { sym = pair.first; break; @@ -146,12 +144,10 @@ ListTypePtr TorchListTypeFromONNX( SymbolDimMap& symbol_map) { c10::optional scalar_type; if (onnx_sequence_type.has_elem_type()) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto onnx_seq_elem_type = onnx_sequence_type.elem_type(); + const auto& onnx_seq_elem_type = onnx_sequence_type.elem_type(); if (onnx_seq_elem_type.has_tensor_type()) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto onnx_tensor_type = onnx_seq_elem_type.tensor_type(); - auto v_tensor_type = + const auto& onnx_tensor_type = onnx_seq_elem_type.tensor_type(); + const auto v_tensor_type = TorchTensorTypeFromONNX(onnx_tensor_type, symbol_map); auto v_type = ListType::create(v_tensor_type); return v_type; @@ -168,16 +164,15 @@ void UpdateTorchValueByOnnxValueInfo( return; } - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto p_type = p_info.type(); + const auto& p_type = p_info.type(); if (p_type.has_tensor_type()) { - auto torch_tensor_type = + const auto torch_tensor_type = TorchTensorTypeFromONNX(p_type.tensor_type(), symbol_map); if (torch_tensor_type) { v->setType(MergeInferredType(v->type(), torch_tensor_type)); } } else if (p_type.has_sequence_type()) { - auto torch_list_type = + const auto torch_list_type = TorchListTypeFromONNX(p_type.sequence_type(), symbol_map); if (torch_list_type) { v->setType(MergeInferredType(v->type(), torch_list_type)); @@ -1316,11 +1311,10 @@ void UpdateOutputTypeByONNXProto( Node* clone_node, const onnx::ModelProto& model_proto, SymbolDimMap& symbol_map) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto graph_proto = model_proto.graph(); + const auto& graph_proto = model_proto.graph(); // get data from value_info and updated original graph. - auto updateNodeOutputsByONNXValueInfo = + const auto updateNodeOutputsByONNXValueInfo = [&](const onnx::ValueInfoProto& v_info) { for (size_t i = 0; i < n->outputs().size(); ++i) { if (clone_node->output(i)->debugName() == v_info.name()) { @@ -1486,10 +1480,9 @@ void ONNXSetDynamicInputShape( shape_ref.has_value(), "Input tensor shape should have value."); auto shape = shape_ref.value(); - // NOLINTNEXTLINE(performance-for-range-copy) - for (auto pair : axes_names) { - auto axis = pair.first; - auto name = pair.second; + for (const auto& pair : axes_names) { + const auto axis = pair.first; + const auto name = pair.second; if (name_to_sym.find(name) == name_to_sym.end()) { name_to_sym[name] = ::c10::ShapeSymbol::newSymbol(); } @@ -1547,8 +1540,7 @@ size_t ONNXAssignOutputShape( index_check(); if (THPVariable_Check(output_obj)) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - at::Tensor var = THPVariable_Unpack(output_obj); + const at::Tensor& var = THPVariable_Unpack(output_obj); ONNXUpdateTypeFromTensor( graph->outputs().at(outputs_index), var, onnx_shape_inference); outputs_index++; @@ -1562,7 +1554,7 @@ size_t ONNXAssignOutputShape( onnx_shape_inference); } } else if (PyList_Check(output_obj)) { - size_t list_len = PyList_GET_SIZE(output_obj); + const auto list_len = PyList_GET_SIZE(output_obj); if (HasSequenceTypeOutput(graph->outputs().at(outputs_index)->node())) { auto output_type = graph->outputs().at(outputs_index)->type(); TORCH_CHECK( diff --git a/torch/csrc/jit/passes/quantization/quantization_patterns.h b/torch/csrc/jit/passes/quantization/quantization_patterns.h index 48b0a33afcb04..606be7a609126 100644 --- a/torch/csrc/jit/passes/quantization/quantization_patterns.h +++ b/torch/csrc/jit/passes/quantization/quantization_patterns.h @@ -37,15 +37,17 @@ std::string getAtenOpPattern( std::string aten_op_pattern = graph_header; if (scalar_args) { for (const auto& extra_arg : _extra_op_args) { - aten_op_pattern += R"( - )" + - // NOLINTNEXTLINE(performance-inefficient-string-concatenation) - extra_arg + "_scalar = aten::item(" + extra_arg + ")"; + aten_op_pattern + .append(R"( + )") + .append(extra_arg) + .append("_scalar = aten::item(") + .append(extra_arg) + .append(")"); } for (auto& _extra_op_arg : _extra_op_args) { - // NOLINTNEXTLINE(performance-inefficient-string-concatenation) - _extra_op_arg = _extra_op_arg + "_scalar"; + _extra_op_arg.append("_scalar"); } } const auto& extra_op_arg_list = getExtraArgList(_extra_op_args); @@ -171,8 +173,8 @@ QuantFusionInfo getClampOpFusionInfo( op_pattern += R"( %r = )"; std::vector scalar_extra_args; + scalar_extra_args.reserve(extra_op_args.size()); for (const auto& arg : extra_op_args) { - // NOLINTNEXTLINE(performance-inefficient-vector-operation) scalar_extra_args.push_back(arg + "_scalar"); } op_pattern += diff --git a/torch/csrc/jit/passes/utils/check_alias_annotation.cpp b/torch/csrc/jit/passes/utils/check_alias_annotation.cpp index 9f88f644feddc..cd894b46ff69b 100644 --- a/torch/csrc/jit/passes/utils/check_alias_annotation.cpp +++ b/torch/csrc/jit/passes/utils/check_alias_annotation.cpp @@ -240,8 +240,7 @@ void checkAliasAnnotation( // it was created by the op. checkInputPreconditions(stack); - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - const auto schema = node->schema(); + const auto& schema = node->schema(); std::vector inputsToCheck; for (const auto i : c10::irange(schema.arguments().size())) { diff --git a/torch/csrc/jit/python/pybind_utils.h b/torch/csrc/jit/python/pybind_utils.h index accfbb01adc8d..68a816b7834de 100644 --- a/torch/csrc/jit/python/pybind_utils.h +++ b/torch/csrc/jit/python/pybind_utils.h @@ -953,8 +953,8 @@ inline py::object runAndInsertCall( auto return_type = callee.getSchema().returns().at(0).type(); auto graph = tracing_state->graph; std::vector named_values; + named_values.reserve(input_values.size()); for (Value* v : input_values) { - // NOLINTNEXTLINE(performance-inefficient-vector-operation) named_values.emplace_back(v); } diff --git a/torch/csrc/jit/serialization/python_print.cpp b/torch/csrc/jit/serialization/python_print.cpp index d8f5b5d80a44a..f43a33ee8d1c0 100644 --- a/torch/csrc/jit/serialization/python_print.cpp +++ b/torch/csrc/jit/serialization/python_print.cpp @@ -1475,11 +1475,10 @@ struct PythonPrintImpl { method.arguments().at(0).name() == "self"); for (const Argument& arg : at::ArrayRef(method.arguments()).slice(1)) { - // NOLINTNEXTLINE(performance-unnecessary-copy-initialization) - auto type = arg.type(); - registerClassDependencies(type); + const auto& arg_type = arg.type(); + registerClassDependencies(arg_type); body_ << ", " << arg.name() << ": " - << type->annotation_str(type_printer_); + << arg_type->annotation_str(type_printer_); } auto return_type = method.returns().at(0).type(); registerClassDependencies(return_type); diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp index 38201c1c9493d..0e8a60c144119 100644 --- a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp +++ b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp @@ -363,9 +363,6 @@ class AtomicAddFuser : public IRMutator { const std::unordered_set& thread_local_bufs, const GPUMetaVarRewriter& metavars) : thread_local_bufs_(thread_local_bufs) { - // NOLINTNEXTLINE(clang-diagnostic-unused-variable) - size_t DIMS = 3; - const std::vector& block_extents = metavars.gpu_block_extents(); const std::vector& block_vars = metavars.gpu_block_vars(); @@ -612,9 +609,6 @@ class PrioritizeLoad : public IRMutator { } Stmt* mutate(const Block* v) override { - // NOLINTNEXTLINE(clang-diagnostic-unused-variable) - bool any_change = false; - Block* v1 = const_cast(v); // NOLINT assert(v1); std::list stmts = v1->stmts(); diff --git a/torch/csrc/jit/tensorexpr/eval.cpp b/torch/csrc/jit/tensorexpr/eval.cpp index eeb1400866f8a..f2deb19beb3cc 100644 --- a/torch/csrc/jit/tensorexpr/eval.cpp +++ b/torch/csrc/jit/tensorexpr/eval.cpp @@ -794,8 +794,7 @@ class SimpleIREvaluatorImpl : public IRVisitor { dim->accept(this); total_byte_size *= value_.as(); } - // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - int int_count = (total_byte_size + sizeof(int) - 1) / sizeof(int); + const auto int_count = (total_byte_size + sizeof(int) - 1) / sizeof(int); std::unique_ptr> buffer(new std::vector(int_count)); auto iter = buffer_mapping_.find(b); if (iter != buffer_mapping_.end() && iter->second != nullptr) { diff --git a/torch/csrc/jit/tensorexpr/expr.cpp b/torch/csrc/jit/tensorexpr/expr.cpp index 4c0e81fb80e91..a812b4985102d 100644 --- a/torch/csrc/jit/tensorexpr/expr.cpp +++ b/torch/csrc/jit/tensorexpr/expr.cpp @@ -147,8 +147,6 @@ ExprHandle abs(const ExprHandle& v) { // The default tanh is quite slow, use the Eigen version from here: // https://bitbucket.org/eigen/eigen/src/94875feeeeb9abe5509b314197da1991ba2070f5/Eigen/src/Core/MathFunctionsImpl.h#lines-26 ExprHandle fast_tanh(const ExprHandle& v) { - // NOLINTNEXTLINE(clang-diagnostic-unused-variable) - Dtype dtype = v.dtype(); // TODO: use a dedicated bind-var to make sure v is not evalualted multiple // times. Clamp the input expression to [-9, 9] ExprHandle plus_9 = FloatImm::make(9.0f); diff --git a/torch/csrc/jit/tensorexpr/hash_provider.h b/torch/csrc/jit/tensorexpr/hash_provider.h index 18064a146ce03..92943b0ce5d1f 100644 --- a/torch/csrc/jit/tensorexpr/hash_provider.h +++ b/torch/csrc/jit/tensorexpr/hash_provider.h @@ -233,8 +233,7 @@ class TORCH_API HashProvider : public IRVisitor { size_t te_hash(std::string val) { size_t hash{0}; int64_t intval{0}; - // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions) - int s = val.size() - 1; + int64_t s = val.size() - 1; while (s >= 0) { for (unsigned int i = 0; i < 8; ++i) { if (s < 0) diff --git a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp index 560175548f348..3474ee6331d6b 100644 --- a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp +++ b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp @@ -1921,7 +1921,7 @@ const Expr* simplifyRoundModPattern(const Polynomial* poly) { while (!mods.empty() && repeat) { repeat = false; // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - for (int i = mods.size() - 1; i >= 0; i--) { + for (int64_t i = mods.size() - 1; i >= 0; i--) { const Term* m = mods[i]; const Mod* mod = dynamic_cast(m->variables()[0]); CHECK(mod); @@ -1929,7 +1929,7 @@ const Expr* simplifyRoundModPattern(const Polynomial* poly) { const Expr* mod_rhs = IRSimplifier::simplify(mod->rhs()); bool merged = false; // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - for (int j = mod_rounds.size() - 1; j >= 0; j--) { + for (int64_t j = mod_rounds.size() - 1; j >= 0; j--) { const Term* mr = mod_rounds[j]; auto a = isModRound(mr); CHECK(a); @@ -1968,7 +1968,7 @@ const Expr* simplifyRoundModPattern(const Polynomial* poly) { } // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - for (int k = rounds.size() - 1; k >= 0; k--) { + for (int64_t k = rounds.size() - 1; k >= 0; k--) { const Term* r = rounds[k]; const RoundOff* roundoff = dynamic_cast(r->variables()[0]); @@ -2045,9 +2045,9 @@ const Term* IRSimplifierBase::factorizePolynomial(const Polynomial* poly) { // Create new struture. std::vector newPolyTerms; + newPolyTerms.reserve(variables.size()); for (auto* t : variables) { // New term with the scalar divided by the GCD. - // NOLINTNEXTLINE(performance-inefficient-vector-operation) newPolyTerms.push_back(new Term( poly->hasher(), evaluateOp(new Div(t->scalar(), GCD)), t->variables())); } diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 0022cec9bb598..3884ce1c762fb 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -3115,8 +3115,7 @@ void TensorExprKernel::genInputDebugNames() { std::string sanitized_name = sanitizeName(input->debugName()); // we could get fancier here, but name conflict is extremely unlikely while (name_set.count(sanitized_name)) { - // NOLINTNEXTLINE(performance-inefficient-string-concatenation) - sanitized_name = sanitized_name + "_"; + sanitized_name.append("_"); } value_to_name[input] = sanitized_name; name_set.insert(sanitized_name); diff --git a/torch/csrc/python_dimname.cpp b/torch/csrc/python_dimname.cpp index cf51a3965d7f6..351fc33934110 100644 --- a/torch/csrc/python_dimname.cpp +++ b/torch/csrc/python_dimname.cpp @@ -63,7 +63,7 @@ bool THPUtils_checkDimnameList(PyObject* obj) { return false; } // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); + const auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); if (size == 0) { return true; } diff --git a/torch/csrc/serialization.cpp b/torch/csrc/serialization.cpp index 715e39d5af4bf..40905d773c255 100644 --- a/torch/csrc/serialization.cpp +++ b/torch/csrc/serialization.cpp @@ -82,13 +82,11 @@ static inline ssize_t doPartialPythonIO(PyObject* fildes, void* buf, size_t nbyt reinterpret_cast(buf), nbytes, rw_flag)); if (!memview) throw python_error(); - // NOLINTNEXTLINE(clang-diagnostic-writable-strings) - char* method = "write"; + std::string method = "write"; if (is_read) { - // NOLINTNEXTLINE(clang-diagnostic-writable-strings) method = "readinto"; } - THPObjectPtr r(PyObject_CallMethod(fildes, method, "O", memview.get())); + THPObjectPtr r(PyObject_CallMethod(fildes, method.c_str(), "O", memview.get())); if (r) { return PyLong_AsSsize_t(r.get()); } diff --git a/torch/csrc/utils.cpp b/torch/csrc/utils.cpp index e7217e77bdb74..414c4069b3a63 100644 --- a/torch/csrc/utils.cpp +++ b/torch/csrc/utils.cpp @@ -57,7 +57,7 @@ bool THPUtils_tryUnpackLongs(PyObject *arg, THLongStoragePtr& result) { bool list = PyList_Check(arg); if (tuple || list) { // NOLINTNEXTLINE(bugprone-branch-clone) - int nDim = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); + const auto nDim = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); THLongStoragePtr storage(THLongStorage_newWithSize(nDim)); for (int i = 0; i != nDim; ++i) { PyObject* item = tuple ? PyTuple_GET_ITEM(arg, i) : PyList_GET_ITEM(arg, i); @@ -77,7 +77,7 @@ std::vector THPUtils_unpackLongs(PyObject *arg) { bool list = PyList_Check(arg); if (tuple || list) { // NOLINTNEXTLINE(bugprone-branch-clone) - int nDim = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); + const auto nDim = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); std::vector sizes(nDim); for (int i = 0; i != nDim; ++i) { PyObject* item = tuple ? PyTuple_GET_ITEM(arg, i) : PyList_GET_ITEM(arg, i); diff --git a/torch/csrc/utils/invalid_arguments.cpp b/torch/csrc/utils/invalid_arguments.cpp index 1daf400a8d861..2b955125dff37 100644 --- a/torch/csrc/utils/invalid_arguments.cpp +++ b/torch/csrc/utils/invalid_arguments.cpp @@ -298,7 +298,7 @@ std::vector _tryMatchKwargs(const Option& option, const std::unordered_map& kwargs) { std::vector unmatched; // NOLINTNEXTLINE(cppcoreguidelines-narrowing-conversions,bugprone-narrowing-conversions) - int start_idx = option.arguments.size() - kwargs.size(); + int64_t start_idx = option.arguments.size() - kwargs.size(); if (option.has_out && kwargs.count("out") == 0) start_idx--; if (start_idx < 0) diff --git a/torch/csrc/utils/pybind.h b/torch/csrc/utils/pybind.h index b9ffc70a769a6..14ddb0eeb22b9 100644 --- a/torch/csrc/utils/pybind.h +++ b/torch/csrc/utils/pybind.h @@ -79,7 +79,7 @@ template<> struct type_caster { auto tuple = PyTuple_Check(source); if (tuple || PyList_Check(source)) { // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(source) : PyList_GET_SIZE(source); + const auto size = tuple ? PyTuple_GET_SIZE(source) : PyList_GET_SIZE(source); v_value.resize(size); for(const auto idx : c10::irange(size)) { PyObject* obj = tuple ? PyTuple_GET_ITEM(source, idx) : PyList_GET_ITEM(source, idx); diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp index 138c8a7ee7911..4e12d35889529 100644 --- a/torch/csrc/utils/python_arg_parser.cpp +++ b/torch/csrc/utils/python_arg_parser.cpp @@ -383,7 +383,7 @@ bool is_scalar_list(PyObject* obj) { return false; } // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); + const auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); for (const auto idx : c10::irange(size)) { PyObject* iobj = tuple ? PyTuple_GET_ITEM(obj, idx) : PyList_GET_ITEM(obj, idx); if (!THPUtils_checkScalar(iobj)) { @@ -399,7 +399,7 @@ bool is_tensor_list_and_append_overloaded(PyObject* obj, std::vector return false; } // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); +const auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); for (long idx = 0; idx < size; idx++) { PyObject* iobj = tuple ? PyTuple_GET_ITEM(obj, idx) : PyList_GET_ITEM(obj, idx); if (!is_tensor_and_append_overloaded(iobj, overloaded_args)) { @@ -420,7 +420,7 @@ bool is_float_or_complex_list(PyObject* obj) { } // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); + const auto size = tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj); if (size > 0) { PyObject* iobj = tuple ? PyTuple_GET_ITEM(obj, 0) : PyList_GET_ITEM(obj, 0); if (!THPUtils_checkDouble(iobj) && !PyComplex_Check(iobj)) { diff --git a/torch/csrc/utils/python_arg_parser.h b/torch/csrc/utils/python_arg_parser.h index 2abacdcb03e52..1f3a21b25c315 100644 --- a/torch/csrc/utils/python_arg_parser.h +++ b/torch/csrc/utils/python_arg_parser.h @@ -384,15 +384,15 @@ inline std::vector PythonArgs::intlist(int i) { inline std::vector PythonArgs::intlistWithDefault(int i, std::vector default_intlist) { if (!args[i]) return default_intlist; PyObject* arg = args[i]; - auto size = signature.params[i].size; - if (size > 0 && THPUtils_checkLong(arg)) { - return std::vector(size, THPUtils_unpackIndex(arg)); + const auto size1 = signature.params[i].size; + if (size1 > 0 && THPUtils_checkLong(arg)) { + return std::vector(size1, THPUtils_unpackIndex(arg)); } auto tuple = PyTuple_Check(arg); // NOLINTNEXTLINE(bugprone-branch-clone) - size = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); - std::vector res(size); - for(const auto idx : c10::irange(size)) { + const auto size2 = tuple ? PyTuple_GET_SIZE(arg) : PyList_GET_SIZE(arg); + std::vector res(size2); + for(const auto idx : c10::irange(size2)) { PyObject* obj = tuple ? PyTuple_GET_ITEM(arg, idx) : PyList_GET_ITEM(arg, idx); try { // Elements of torch.Size are tensors during tracing, and we need to record extra @@ -400,14 +400,14 @@ inline std::vector PythonArgs::intlistWithDefault(int i, std::vector(); continue; } else { res[idx] = THPUtils_unpackIndex(obj); } } catch (const std::exception &e) { - throw TypeError("%s(): argument '%s' must be %s, but found element of type %s at pos %d", + throw TypeError("%s(): argument '%s' must be %s, but found element of type %s at pos %ld", signature.name.c_str(), signature.params[i].name.c_str(), signature.params[i].type_name().c_str(), Py_TYPE(obj)->tp_name, idx + 1); } diff --git a/torch/csrc/utils/throughput_benchmark-inl.h b/torch/csrc/utils/throughput_benchmark-inl.h index 5b2c7b81d9766..908f092fe8601 100644 --- a/torch/csrc/utils/throughput_benchmark-inl.h +++ b/torch/csrc/utils/throughput_benchmark-inl.h @@ -58,8 +58,8 @@ BenchmarkExecutionStats BenchmarkHelper::benchmark( std::atomic num_attempted_iters{0}; std::vector callers; + callers.reserve(config.num_calling_threads); for (const auto thread_id : c10::irange(config.num_calling_threads)) { - // NOLINTNEXTLINE(performance-inefficient-vector-operation) callers.emplace_back([&, thread_id]() { // We use conditional variable as a barrier to make sure each thread // performs required warmeup iterations before we start measuring From 8dd0570b34c7c378ae9729c21267546cba07fdc9 Mon Sep 17 00:00:00 2001 From: Ailing Zhang Date: Tue, 15 Jun 2021 13:18:39 -0700 Subject: [PATCH 122/305] Reuse build_torch_xla from pytorch/xla repo. (#59989) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59989 Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29138211 Pulled By: ailzhang fbshipit-source-id: 349d307c510e7fad266822e320f0d6904fa00239 --- .jenkins/pytorch/build.sh | 40 +++++++------------------------- .jenkins/pytorch/common_utils.sh | 4 ++++ 2 files changed, 12 insertions(+), 32 deletions(-) diff --git a/.jenkins/pytorch/build.sh b/.jenkins/pytorch/build.sh index c2be6c96b3e72..01f5125a2b4bf 100755 --- a/.jenkins/pytorch/build.sh +++ b/.jenkins/pytorch/build.sh @@ -200,8 +200,10 @@ fi # Patch required to build xla if [[ "${BUILD_ENVIRONMENT}" == *xla* ]]; then - git clone --recursive https://github.com/pytorch/xla.git - ./xla/scripts/apply_patches.sh + clone_pytorch_xla + # shellcheck disable=SC1091 + source "xla/.circleci/common.sh" + apply_patches fi if [[ "${BUILD_ENVIRONMENT}" == pytorch-linux-xenial-py3.6-gcc7-build || "${BUILD_ENVIRONMENT}" == pytorch-linux-xenial-py3.6-gcc5.4-build ]]; then @@ -311,36 +313,10 @@ fi # Test XLA build if [[ "${BUILD_ENVIRONMENT}" == *xla* ]]; then - # TODO: Move this to Dockerfile. - - pip_install lark-parser - pip_install cloud-tpu-client - - sudo apt-get -qq update - sudo apt-get -qq install npm nodejs - - # XLA build requires Bazel - # We use bazelisk to avoid updating Bazel version manually. - sudo npm install -g @bazel/bazelisk - sudo ln -s "$(command -v bazelisk)" /usr/bin/bazel - - # Install bazels3cache for cloud cache - sudo npm install -g bazels3cache - BAZELS3CACHE="$(which bazels3cache)" - if [ -z "${BAZELS3CACHE}" ]; then - echo "Unable to find bazels3cache..." - exit 1 - fi - - bazels3cache --bucket="${XLA_CLANG_CACHE_S3_BUCKET_NAME}" --maxEntrySizeBytes=0 - pushd xla - export CC=clang-9 CXX=clang++-9 - # Use cloud cache to build when available. - # shellcheck disable=SC1003 - sed -i '/bazel build/ a --remote_http_cache=http://localhost:7777 \\' build_torch_xla_libs.sh - - python setup.py install - popd + XLA_DIR=xla + # These functions are defined in .circleci/common.sh in pytorch/xla repo + install_deps_pytorch_xla $XLA_DIR + build_torch_xla $XLA_DIR assert_git_not_dirty fi diff --git a/.jenkins/pytorch/common_utils.sh b/.jenkins/pytorch/common_utils.sh index ea8d0c6282714..fd94ce14a1c5f 100644 --- a/.jenkins/pytorch/common_utils.sh +++ b/.jenkins/pytorch/common_utils.sh @@ -86,3 +86,7 @@ function checkout_install_torchvision() { time python setup.py install popd } + +function clone_pytorch_xla() { + git clone --recursive https://github.com/pytorch/xla.git +} From dc1f60a9a2616946e2c3e82b915871e8a8993f31 Mon Sep 17 00:00:00 2001 From: Zafar Takhirov Date: Tue, 15 Jun 2021 13:35:53 -0700 Subject: [PATCH 123/305] [sparsity][refactor] Restructure the tests folders (#60032) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60032 There will be more sparse tests coming. This PR creates a separate folder for the sparse tests Test Plan: `python test/test_ao.py` Reviewed By: raghuramank100 Differential Revision: D29139265 fbshipit-source-id: d0db915f00e6bc8d89a5651f08f72e362a912a6b --- .../{test_ao_sparse.py => ao/sparsity/test_kernels.py} | 0 test/test_ao_sparsity.py | 10 ++++++++++ 2 files changed, 10 insertions(+) rename test/{test_ao_sparse.py => ao/sparsity/test_kernels.py} (100%) create mode 100644 test/test_ao_sparsity.py diff --git a/test/test_ao_sparse.py b/test/ao/sparsity/test_kernels.py similarity index 100% rename from test/test_ao_sparse.py rename to test/ao/sparsity/test_kernels.py diff --git a/test/test_ao_sparsity.py b/test/test_ao_sparsity.py new file mode 100644 index 0000000000000..9aad623d38bb8 --- /dev/null +++ b/test/test_ao_sparsity.py @@ -0,0 +1,10 @@ +# -*- coding: utf-8 -*- + +from torch.testing._internal.common_utils import run_tests + +# Kernels +from ao.sparsity.test_kernels import TestQuantizedSparseKernels # noqa: F401 +from ao.sparsity.test_kernels import TestQuantizedSparseLayers # noqa: F401 + +if __name__ == '__main__': + run_tests() From 1d5a577f04fbd67824e7b74aaa1ce3ec72f5557d Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Tue, 15 Jun 2021 13:37:16 -0700 Subject: [PATCH 124/305] Fix some items identified as problematic by Wextra and other clean-up (#59909) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59909 Test Plan: Sandcastle Reviewed By: vkuzo Differential Revision: D29073150 fbshipit-source-id: 500a92ccb57b0e40277863a3b235099fd66ab8ad --- aten/src/ATen/BatchedTensorImpl.cpp | 6 +- .../quantized/cpu/qlinear_prepack.cpp | 11 +-- aten/src/ATen/native/cpu/SumKernel.cpp | 7 +- .../cpu/kernels/QuantizedOpKernels.cpp | 85 ++++++++----------- .../ATen/native/quantized/cpu/qnnpack_utils.h | 29 ++++--- 5 files changed, 63 insertions(+), 75 deletions(-) diff --git a/aten/src/ATen/BatchedTensorImpl.cpp b/aten/src/ATen/BatchedTensorImpl.cpp index 7deb84a5b7ec7..db06930247426 100644 --- a/aten/src/ATen/BatchedTensorImpl.cpp +++ b/aten/src/ATen/BatchedTensorImpl.cpp @@ -2,6 +2,7 @@ #include #include +#include namespace at { @@ -23,8 +24,7 @@ BatchedTensorImpl::BatchedTensorImpl(Tensor value, BatchDims bdims) const auto value_sizes = value_.sizes(); const auto value_strides = value_.strides(); sizes_and_strides_.resize(public_dims); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (int64_t dim = 0; dim < public_dims; dim++) { + for (const auto dim : c10::irange(public_dims)) { auto actual_dim = actualDim(dim, /*wrap_dim=*/false); sizes_and_strides_.size_at_unchecked(dim) = value_sizes.at(actual_dim); sizes_and_strides_.stride_at_unchecked(dim) = value_strides.at(actual_dim); @@ -51,7 +51,7 @@ int64_t BatchedTensorImpl::actualDim(int64_t dim, bool wrap_dim) const { // but it might require newer (>= ~2015) CPUs. We should clean this up // if/when we have dropped support for older CPUs. int64_t non_bdim_count = 0; - for (int64_t actual_dim = 0; actual_dim < kVmapMaxTensorDims; actual_dim++) { + for (const auto actual_dim : c10::irange(kVmapMaxTensorDims)) { if (is_bdim[actual_dim]) { continue; } diff --git a/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_prepack.cpp b/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_prepack.cpp index 6f91164acbd2d..7c780806c5d07 100644 --- a/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_prepack.cpp +++ b/aten/src/ATen/native/ao_sparse/quantized/cpu/qlinear_prepack.cpp @@ -1,4 +1,5 @@ #include +#include #include #include @@ -28,10 +29,10 @@ void calc_col_offsets_transpose( int32_t* col_offsets, c10::QScheme qtype) { // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (size_t i = 0; i < N; ++i) { + for (const auto i : c10::irange(N)) { int32_t sum = 0; // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (size_t j = 0; j < K; ++j) { + for (const auto j : c10::irange(K)) { sum += Bint8[i * K + j]; } if (qtype == c10::kPerTensorAffine) { @@ -64,7 +65,7 @@ c10::intrusive_ptr PackedLinearWeight:: weight_zero_points_int32[0] = weight.q_zero_point(); } else if (qtype == c10::kPerChannelAffine) { weight_zero_points_int32.resize(N, 0); - for (int i = 0; i < N; ++i) { + for (const auto i : c10::irange(N)) { weight_zero_points_int32[i] = weight.q_per_channel_zero_points()[i].item(); } @@ -80,7 +81,7 @@ c10::intrusive_ptr PackedLinearWeight:: weight_scales_float[0] = weight.q_scale(); } else if (qtype == c10::kPerChannelAffine) { weight_scales_float.resize(N, 0.0); - for (int i = 0; i < N; ++i) { + for (const auto i : c10::irange(N)) { weight_scales_float[i] = weight.q_per_channel_scales()[i].item(); } } @@ -185,7 +186,7 @@ PackedLinearWeightQnnp::PackedLinearWeightQnnp( auto wt_numel = weight_contig.numel(); int8_t* w_data = reinterpret_cast(weight_contig.data_ptr()); - for (int i = 0; i < wt_numel; ++i) { + for (const auto i : c10::irange(wt_numel)) { qnnp_w_data[i] = static_cast(w_data[i] + 128); } bcsr_matrix_ = qnnpack::generateBlockCSRMatrix( diff --git a/aten/src/ATen/native/cpu/SumKernel.cpp b/aten/src/ATen/native/cpu/SumKernel.cpp index e0e5e21069f17..45399da3a3f84 100644 --- a/aten/src/ATen/native/cpu/SumKernel.cpp +++ b/aten/src/ATen/native/cpu/SumKernel.cpp @@ -34,16 +34,15 @@ T load(const char * C10_RESTRICT data, int64_t stride, int64_t index) { template void accumulate_result(char * C10_RESTRICT data, int64_t stride, int64_t index, scalar_t value) { - auto * ptr = reinterpret_cast(data + index * stride); + auto *const ptr = reinterpret_cast(data + index * stride); *ptr += value; } template void accumulate_result(char * C10_RESTRICT data, int64_t stride, int64_t index, const std::array &values) { - auto *base_ptr = data + stride * index; - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (int64_t k = 0; k < numel; ++k) { + auto *const base_ptr = data + stride * index; + for (const auto k : c10::irange(numel)) { accumulate_result(base_ptr, stride, k, values[k]); } } diff --git a/aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp b/aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp index 7ee96da8db12c..9594883ab680d 100644 --- a/aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp +++ b/aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #ifdef USE_FBGEMM @@ -106,8 +107,7 @@ Tensor qcat_nhwc_kernel( for (int64_t row = 0; row < H; ++row) { for (int64_t col = 0; col < W; ++col) { // loop over input tensors - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (int64_t tidx = 0; tidx < Cs_in.size(); ++tidx) { + for (const auto tidx : c10::irange(Cs_in.size())) { scalar_t::underlying* optr = reinterpret_cast(output.data_ptr()) + batch * H * W * C_out + row * W * C_out + col * C_out + @@ -493,12 +493,10 @@ static void leaky_qrelu_out_kernel(Tensor& out, const Tensor& qx, */ auto dx_vec_vec = qx_vec.dequantize(i_scale_vec, i_zp_vec, i_scale_zp_neg_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < dx_vec_vec.size(); ++idx) { - const auto dx_vec = dx_vec_vec[idx]; + for (auto& dx_vec: dx_vec_vec) { const auto multiplicand = Vec::blendv(negval_vec, one_vec, dx_vec > zero_vec); - dx_vec_vec[idx] = dx_vec * multiplicand; + dx_vec = dx_vec * multiplicand; } return qVec::quantize(dx_vec_vec, o_scale, o_zp, o_inv_scale); }); @@ -539,12 +537,11 @@ void qsigmoid_kernel( [&](Vec value_qx) -> Vec { auto value_dx = value_qx.dequantize( scale_vec, zero_point_vec, scale_neg_zp_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < value_dx.size(); ++idx) { - value_dx[idx] = value_dx[idx].neg(); - value_dx[idx] = value_dx[idx].exp(); - value_dx[idx] = Vectorized(1.0f) + value_dx[idx]; - value_dx[idx] = value_dx[idx].reciprocal(); + for (auto& value: value_dx) { + value = value.neg(); + value = value.exp(); + value = Vectorized(1.0f) + value; + value = value.reciprocal(); } return Vec::quantize( value_dx, output_scale, output_zero_point, inv_output_scale); @@ -604,11 +601,10 @@ void qhardsigmoid_kernel(const Tensor& qx, Tensor& qy) { [&](qVec value_qx) -> qVec { auto value_dx = value_qx.dequantize( scale_vec, zero_point_vec, scale_neg_zp_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < value_dx.size(); ++idx) { - value_dx[idx] = + for (auto& value : value_dx) { + value = vec::minimum( - vec::maximum(value_dx[idx] + kThreeVec, kZeroVec), + vec::maximum(value + kThreeVec, kZeroVec), kSixVec) / kSixVec; } @@ -765,15 +761,14 @@ void qthreshold_kernel( // dequantize auto dx_vec = value_qx.dequantize( input_scale_vec, input_zero_point_vec, input_scale_neg_zp_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < dx_vec.size(); ++idx) { + for (auto& value : dx_vec) { // check if any elements are below threshold - auto cmp_to_threshold = dx_vec[idx] > threshold_vec; + const auto cmp_to_threshold = value > threshold_vec; if (cmp_to_threshold.zero_mask()) { // blend - dx_vec[idx] = Vec::blendv(value_vec, dx_vec[idx], cmp_to_threshold); - } + value = Vec::blendv(value_vec, value, cmp_to_threshold); } + } // quantize return qVec::quantize(dx_vec, output_scale, output_zero_point, inv_output_scale); }); @@ -812,10 +807,9 @@ void qhardswish_kernel(const Tensor& qx, Tensor& qy) { [&](qVec value) -> qVec { auto value_dx = value.dequantize(i_scale_vec, i_zero_point_vec, i_scale_neg_zp_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < value_dx.size(); idx++) { - value_dx[idx] = value_dx[idx] * vec::minimum( - vec::maximum(value_dx[idx] + three_vec, zero_vec), + for (auto& value: value_dx) { + value = value * vec::minimum( + vec::maximum(value + three_vec, zero_vec), six_vec ) / six_vec; } @@ -869,7 +863,7 @@ void qtanh_kernel(const Tensor& qx, Tensor& qy) { const auto value_dx = value_qx.dequantize( scale_vec, zero_point_vec, scale_neg_zp_premul_vec); Vec::float_vec_return_type retvals; - for (int idx = 0; idx < Vec::float_num_vecs(); ++idx) { + for (const auto idx : c10::irange(Vec::float_num_vecs())) { retvals[idx] = value_dx[idx].tanh(); } return Vec::quantize( @@ -940,26 +934,23 @@ void qelu_kernel( // dequantize auto dx_vec_vec = value_qx.dequantize(i_scale_vec, i_zero_point_vec, i_scale_neg_zp_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare,modernize-loop-convert) - for (int idx = 0; idx < dx_vec_vec.size(); idx++) { - + for (auto& value : dx_vec_vec) { // quickly check if any elements are below zero - auto cmp_to_zero = dx_vec_vec[idx] > zero_vec; + const auto cmp_to_zero = value > zero_vec; if (cmp_to_zero.zero_mask()) { - - Vec dx_vec_copy_neg_elu = dx_vec_vec[idx] * one_vec; + Vec dx_vec_copy_neg_elu = value * one_vec; // calculate the negative part of ELU on the copy dx_vec_copy_neg_elu = dx_vec_copy_neg_elu * input_scale_coef_vec; dx_vec_copy_neg_elu = dx_vec_copy_neg_elu.exp(); dx_vec_copy_neg_elu = dx_vec_copy_neg_elu - one_vec; dx_vec_copy_neg_elu = dx_vec_copy_neg_elu * alpha_vec; // blend - dx_vec_vec[idx] = Vec::blendv(dx_vec_copy_neg_elu, dx_vec_vec[idx], - dx_vec_vec[idx] > zero_vec); + value = Vec::blendv(dx_vec_copy_neg_elu, value, + value > zero_vec); } - dx_vec_vec[idx] = dx_vec_vec[idx] * scale_coef_vec; + value = value * scale_coef_vec; } // quantize return qVec::quantize(dx_vec_vec, o_scale, o_zp, inv_o_scale); @@ -1007,7 +998,7 @@ void qadd_scalar_kernel(Tensor& out, const Tensor& self, const Scalar& other) { Vec::int_vec_return_type a_sub_z = a.widening_subtract(Vec(static_cast(self_zero_point))); Vec::int_vec_return_type c; - for (int i = 0; i < Vec::int_num_vecs(); ++i) { + for (const auto i : c10::irange(Vec::int_num_vecs())) { c[i] = a_sub_z[i] + other_vec; } Vec rv = Vec::requantize_from_int(c, multiplier, zero_point); @@ -1068,7 +1059,7 @@ void qadd_kernel(Tensor& out, const Tensor& self, const Tensor& other) { const auto db = b.dequantize( other_scale_vec, other_zero_point_vec, other_scale_zp_premul_vec); Vec::float_vec_return_type retvals; - for (int i = 0; i < Vec::float_num_vecs(); ++i) { + for (const auto i : c10::irange(Vec::float_num_vecs())) { auto c = da[i] + db[i]; if (ReLUFused) { c = vec::maximum(c, Vectorized(0.0f)); @@ -1130,7 +1121,7 @@ void qmul_kernel(Tensor& out, const Tensor& self, const Tensor& other) { Vec::int_vec_return_type b_sub_zp = b.widening_subtract(Vec(static_cast(other_zero_point))); Vec::int_vec_return_type c; - for (int i = 0; i < Vec::int_num_vecs(); ++i) { + for (const auto i : c10::irange(Vec::int_num_vecs())) { c[i] = a_sub_zp[i] * b_sub_zp[i]; } Vec rv = Vec::requantize_from_int(c, multiplier, zero_point); @@ -1962,8 +1953,7 @@ inline void do_bn_compute( auto vals_q = Vec::loadu(X_ptr); // Fake scale of 1.0 here, should not affect performance (FMA in place of sub) auto vals_dq = vals_q.dequantize(fake_scale, in_zp_vec, scale_neg_zp_premul); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (size_t idx = 0; idx < vec_num; ++idx) { + for (const auto idx : c10::irange(vec_num)) { auto alpha_v = Vectorized::loadu(alpha + idx * kVLen); auto beta_v = Vectorized::loadu(beta + idx * kVLen); vals_dq[idx] = vec::fmadd(alpha_v, vals_dq[idx], beta_v); @@ -2006,13 +1996,12 @@ void q_batch_norm_kernel( auto fake_scale = Vectorized(1.0f); auto scale_neg_zp_premul = fake_scale * in_zp_vec.neg(); auto out_zero_point_v = Vec(scalar_t(out_zero_point)); - size_t lanes = Vec::float_num_vecs() * kVLen; - for (int64_t i = 0; i < outer_size; ++i) { + const auto lanes = static_cast(Vec::float_num_vecs() * kVLen); + for (const auto i : c10::irange(outer_size)) { auto* X_ptr = reinterpret_cast(X + i * C); auto* Y_ptr = reinterpret_cast(Y + i * C); int64_t ch = 0; - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) for(; ch + lanes <= C; ch += lanes ) { do_bn_compute( X_ptr + ch, @@ -2330,10 +2319,9 @@ void quantized_normalize_kernel( auto qXVec = qVec::loadu(X_ptr + vecStartIdx); auto dqXVec = qXVec.dequantize(x_fake_scale_vec, x_zp_vec, x_fake_scale_zp_neg_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (int dqXVecIdx = 0; dqXVecIdx < dqXVec.size(); dqXVecIdx++) { - dqXVec[dqXVecIdx] = - (dqXVec[dqXVecIdx] - layer_mean_div_scale_xVec) * + for (auto &dq : dqXVec) { + dq = + (dq - layer_mean_div_scale_xVec) * gamma_p_vec + beta_vec; qVec::quantize(dqXVec, y_scale, y_zp, y_inv_scale) .store(Y_ptr + vecStartIdx); @@ -2357,8 +2345,7 @@ void quantized_normalize_kernel( auto qXVec = qVec::loadu(X_ptr + vecStartIdx); auto dqXVec = qXVec.dequantize(x_fake_scale_vec, x_zp_vec, x_fake_scale_zp_neg_premul_vec); - // NOLINTNEXTLINE(clang-diagnostic-sign-compare) - for (int dqXVecIdx = 0; dqXVecIdx < dqXVec.size(); dqXVecIdx++) { + for (const auto dqXVecIdx : c10::irange(dqXVec.size())) { int64_t vecVecStartIdx = vecStartIdx + dqXVecIdx * kFloatVLen; auto gammaVec = gamma_null ? one_vec diff --git a/aten/src/ATen/native/quantized/cpu/qnnpack_utils.h b/aten/src/ATen/native/quantized/cpu/qnnpack_utils.h index 43ce06b1323ab..14eea471e5ac9 100644 --- a/aten/src/ATen/native/quantized/cpu/qnnpack_utils.h +++ b/aten/src/ATen/native/quantized/cpu/qnnpack_utils.h @@ -2,6 +2,7 @@ #ifdef USE_PYTORCH_QNNPACK #include +#include #include #include @@ -367,13 +368,13 @@ std::vector generate_requantization_scales( std::vector& requant_scales) { // Since weight scale is allocated with padding // weight_scales.numel() gives us padded num elements. - auto num_output_channels_padded = weight_scales.numel(); - float* weight_scales_data = weight_scales.data_ptr(); - if (requant_scales.size() < num_output_channels_padded) { + const auto num_output_channels_padded = weight_scales.numel(); + float *const weight_scales_data = weight_scales.data_ptr(); + if (static_cast(requant_scales.size()) < num_output_channels_padded) { requant_scales.resize(num_output_channels_padded); } - for (int i = 0; i < num_output_channels_padded; ++i) { - auto inverse_output_scale = 1.f /output_scale; + for (const auto i : c10::irange(num_output_channels_padded)) { + const auto inverse_output_scale = 1.f /output_scale; requant_scales[i] = (weight_scales_data[i] * input_scale) * inverse_output_scale; TORCH_CHECK( (requant_scales[i] > 0.0f && std::isnormal(requant_scales[i])), @@ -390,14 +391,14 @@ std::pair, at::Tensor> make_zero_points_and_scales_tensor( uint32_t groups = 1 ) { const int out_ch_idx = transpose ? 1 : 0; - auto num_output_channels = weight_contig.size(out_ch_idx) * (transpose ? groups : 1); + const auto num_output_channels = weight_contig.size(out_ch_idx) * (transpose ? groups : 1); // Add 8 to account for bufferring needed by QNNPACK. - auto num_output_channels_padded = num_output_channels + 8; + const auto num_output_channels_padded = num_output_channels + 8; const auto qtype = weight_contig.qscheme(); std::vector weight_zp(num_output_channels_padded, 0); // Adjust weight zero point, similar to weight data. if (qtype == at::kPerTensorAffine) { - for (int i = 0; i < num_output_channels; ++i) { + for (const auto i : c10::irange(num_output_channels)) { weight_zp[i] = (uint8_t)(weight_contig.q_zero_point() + 128); } } else if (qtype == at::kPerChannelAffine) { @@ -406,7 +407,7 @@ std::pair, at::Tensor> make_zero_points_and_scales_tensor( "Per channel zero points dtype must be long int."); const int64_t* per_channel_zero_points = weight_contig.q_per_channel_zero_points().data_ptr(); - for (int i = 0; i < num_output_channels; ++i) { + for (const auto i : c10::irange(num_output_channels)) { weight_zp[i] = (uint8_t)(per_channel_zero_points[i] + 128); } } else { @@ -416,24 +417,24 @@ std::pair, at::Tensor> make_zero_points_and_scales_tensor( at::empty( {num_output_channels_padded}, at::device(at::kCPU).dtype(at::kFloat)); - float* weight_scales_data = weight_scales.data_ptr(); + float *const weight_scales_data = weight_scales.data_ptr(); if (qtype == at::kPerTensorAffine) { - for (int i = 0; i < num_output_channels; ++i) { + for (const auto i : c10::irange(num_output_channels)) { weight_scales_data[i] = weight_contig.q_scale(); } } else if (qtype == at::kPerChannelAffine) { TORCH_CHECK( weight_contig.q_per_channel_scales().scalar_type() == at::kDouble, "Per channel scales dtype must be double."); - const double* per_channel_scales = + const double *const per_channel_scales = weight_contig.q_per_channel_scales().data_ptr(); - for (int i = 0; i < num_output_channels; ++i) { + for (const auto i : c10::irange(num_output_channels)) { weight_scales_data[i] = static_cast(per_channel_scales[i]); } } else { TORCH_INTERNAL_ASSERT("Unsupported quantization scheme."); } - for (int i = num_output_channels; i < num_output_channels_padded; ++i) { + for (const auto i : c10::irange(num_output_channels, num_output_channels_padded)) { weight_scales_data[i] = 1.f; } return {weight_zp, weight_scales}; From 50229b5250950505108347b39da02a04747e4c06 Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Tue, 15 Jun 2021 14:09:49 -0700 Subject: [PATCH 125/305] Fix some typing issues (#59952) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59952 Test Plan: Sandcastle Reviewed By: swolchok Differential Revision: D29083423 fbshipit-source-id: 7a13d6ba60808bcf88d809db194d0f873605172c --- caffe2/proto/caffe2_pb.h | 2 +- torch/csrc/jit/runtime/argument_spec.h | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/caffe2/proto/caffe2_pb.h b/caffe2/proto/caffe2_pb.h index fc8acab2d62ab..f99f3778effee 100644 --- a/caffe2/proto/caffe2_pb.h +++ b/caffe2/proto/caffe2_pb.h @@ -111,7 +111,7 @@ inline TORCH_API caffe2::DeviceOption DeviceToOption( inline TORCH_API at::Device OptionToDevice(const caffe2::DeviceOption option) { auto type = option.device_type(); - int32_t id = -1; + c10::DeviceIndex id = -1; switch (type) { case caffe2::PROTO_CPU: if (option.has_numa_node_id()) { diff --git a/torch/csrc/jit/runtime/argument_spec.h b/torch/csrc/jit/runtime/argument_spec.h index 34c63e7bb285e..05a1ef0cb995c 100644 --- a/torch/csrc/jit/runtime/argument_spec.h +++ b/torch/csrc/jit/runtime/argument_spec.h @@ -344,7 +344,9 @@ struct CompleteArgumentInfo { return pod(i).requires_grad; } at::Device device() const { - return at::Device(DeviceType(pod(i).dev_type), pod(i).device); + return at::Device( + DeviceType(pod(i).dev_type), + static_cast(pod(i).device)); } int ndimension() const { // See [valid range], it is always valid to ask for offset for (i + 1) From f232b052a61ac2aeefc56bdc32ceb5b87b8db451 Mon Sep 17 00:00:00 2001 From: Hangchen Yu Date: Tue, 15 Jun 2021 16:12:53 -0700 Subject: [PATCH 126/305] [fx-acc][easy] Format FX experimental partitioner code (#60030) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60030 As titled. Non-functional re-format. Test Plan: NA Reviewed By: gcatron Differential Revision: D29038449 fbshipit-source-id: a7c94eaab86850ef57b51ec66bfe8ea0e68d2dc8 --- test/test_fx_experimental.py | 388 ++++++++----- .../experimental/accelerator_partitioner.py | 536 ++++++++++-------- torch/fx/experimental/normalize.py | 101 ++-- torch/fx/experimental/partitioner_utils.py | 158 ++++-- 4 files changed, 723 insertions(+), 460 deletions(-) diff --git a/test/test_fx_experimental.py b/test/test_fx_experimental.py index 353bcafe90856..7329be167914d 100644 --- a/test/test_fx_experimental.py +++ b/test/test_fx_experimental.py @@ -1,52 +1,61 @@ -import torch -import operator -import unittest -import sys import math import numbers +import operator +import sys +import unittest from typing import Callable, Dict, Union, List, Optional + +import torch +import torch.fx.experimental.optimization as optimization from torch.fx._symbolic_trace import symbolic_trace -from torch.fx.graph_module import GraphModule -from torch.fx.node import Node from torch.fx.experimental import graph_manipulation +from torch.fx.experimental import merge_matmul from torch.fx.experimental.accelerator_partitioner import Partitioner -from torch.fx.experimental.rewriter import RewritingTracer +from torch.fx.experimental.normalize import NormalizeOperators, NormalizeArgs from torch.fx.experimental.param_fetch import lift_lowering_attrs_to_nodes -from torch.testing._internal.common_utils import run_tests -from torch.testing._internal.jit_utils import JitTestCase -from torch.testing._internal.common_methods_invocations import op_db -from torch.testing._internal.common_device_type import ops, onlyCPU, instantiate_device_type_tests -from torch.fx.passes.split_module import split_module from torch.fx.experimental.partitioner_utils import ( NodeLatency, get_partition_to_latency_mapping, get_latency_of_partitioned_graph, Device, PartitionerConfig, - PartitionMode + PartitionMode, ) -import torch.fx.experimental.optimization as optimization -from torch.fx.experimental import merge_matmul -from torch.fx.experimental.normalize import NormalizeOperators, NormalizeArgs +from torch.fx.experimental.rewriter import RewritingTracer from torch.fx.experimental.schema_type_annotation import AnnotateTypesWithSchema -from torch.testing._internal.common_nn import module_tests, new_module_tests +from torch.fx.graph_module import GraphModule +from torch.fx.node import Node from torch.fx.operator_schemas import ( _torchscript_type_to_python_type, normalize_function, normalize_module, type_matches, - create_type_hint + create_type_hint, ) from torch.fx.passes.shape_prop import extract_tensor_metadata, ShapeProp +from torch.fx.passes.split_module import split_module +from torch.testing._internal.common_device_type import ( + ops, + onlyCPU, + instantiate_device_type_tests, +) +from torch.testing._internal.common_methods_invocations import op_db +from torch.testing._internal.common_nn import module_tests, new_module_tests +from torch.testing._internal.common_utils import run_tests +from torch.testing._internal.jit_utils import JitTestCase try: - from torchvision.models import resnet18 import torchvision.models + from torchvision.models import resnet18 + HAS_TORCHVISION = True except ImportError: HAS_TORCHVISION = False skipIfNoTorchVision = unittest.skipIf(not HAS_TORCHVISION, "no torchvision") -skipIfNoMkldnn = unittest.skipIf(not (torch.backends.mkldnn.enabled and torch.backends.mkldnn.is_available()), "no MKLDNN") +skipIfNoMkldnn = unittest.skipIf( + not (torch.backends.mkldnn.enabled and torch.backends.mkldnn.is_available()), + "no MKLDNN", +) def symbolic_trace_with_rewrite(root: Union[torch.nn.Module, Callable]) -> GraphModule: @@ -55,6 +64,7 @@ def symbolic_trace_with_rewrite(root: Union[torch.nn.Module, Callable]) -> Graph RewritingTracer().trace(root), ) + class TestFXExperimental(JitTestCase): def test_serialize_graph(self): class TestModule(torch.nn.Module): @@ -86,19 +96,20 @@ def forward(self, a, b, c): # Fix for now to add type/shape to output for node in traced.graph.nodes: if node.op == "output": - node.meta['tensor_meta'] = extract_tensor_metadata(a) + node.meta["tensor_meta"] = extract_tensor_metadata(a) for mod in module_with_submodules.modules(): if isinstance(mod, GraphModule): for node in mod.graph.nodes: - node.meta['tensor_meta'] = extract_tensor_metadata(a) + node.meta["tensor_meta"] = extract_tensor_metadata(a) for node in module_with_submodules.graph.nodes: - node.meta['tensor_meta'] = extract_tensor_metadata(a) - + node.meta["tensor_meta"] = extract_tensor_metadata(a) weights1 = {} weights2 = {} serialized_graph1 = graph_manipulation.serialize_module(traced, weights1) - serialized_graph2 = graph_manipulation.serialize_module(module_with_submodules, weights2) + serialized_graph2 = graph_manipulation.serialize_module( + module_with_submodules, weights2 + ) assert len(weights1) == 4 assert len(weights2) == 4 assert len(serialized_graph1["nodes"]) == 10 @@ -108,13 +119,8 @@ def forward(self, a, b, c): assert len(serialized_graph2["weights"]) == 4 assert len(serialized_graph2["modules"]) == 1 assert serialized_graph1["weights"]["linear.weight"]["shape"] == "[4, 4]" - assert ( - serialized_graph1["weights"]["linear.weight"]["dtype"] - == "torch.float32" - ) - assert ( - serialized_graph1["weights"]["linear.weight"]["is_quantized"] is False - ) + assert serialized_graph1["weights"]["linear.weight"]["dtype"] == "torch.float32" + assert serialized_graph1["weights"]["linear.weight"]["is_quantized"] is False assert serialized_graph1["nodes"][0]["shape"] == "[4]" assert serialized_graph1["nodes"][0]["dtype"] == "torch.float32" assert serialized_graph1["nodes"][0]["target"] == "a" @@ -135,8 +141,12 @@ def forward(self, a, b, c): q_tensor_channel = torch.quantize_per_channel( x, torch.tensor([0.1, 0.01]), torch.tensor([10, 0]), 0, torch.quint8 ) - result, _ = graph_manipulation.serialize_tensor_quantization(q_tensor, weights={}, pcq_prefix="foo") - result2, per_channel_dict = graph_manipulation.serialize_tensor_quantization(q_tensor_channel, weights={}, pcq_prefix="bar") + result, _ = graph_manipulation.serialize_tensor_quantization( + q_tensor, weights={}, pcq_prefix="foo" + ) + result2, per_channel_dict = graph_manipulation.serialize_tensor_quantization( + q_tensor_channel, weights={}, pcq_prefix="bar" + ) assert result["qscheme"] == "torch.per_tensor_affine" assert result["q_scale"] == 1.0 assert result2["qscheme"] == "torch.per_channel_affine" @@ -157,7 +167,7 @@ def forward(self, a, b): devices = [ Device("dev_0", 125, 0), Device("dev_1", 125, 1), - Device("dev_2", 125, 2) + Device("dev_2", 125, 2), ] partitioner_config = PartitionerConfig(devices) ret = partitioner.partition_graph(traced, m, partitioner_config) @@ -207,7 +217,7 @@ def forward(self, a): Device("dev_1", 40, 0), Device("dev_2", 40, 0), Device("dev_3", 40, 0), - Device("dev_4", 40, 0) + Device("dev_4", 40, 0), ] partitioner_config = PartitionerConfig(devices, PartitionMode.size_based) catch_runtime_error = False @@ -230,7 +240,7 @@ def forward(self, a, b): a, b = torch.rand(4), torch.rand(4) graph_manipulation.get_size_of_all_nodes(traced, [a, b]) partitioner = Partitioner() - devices = [Device('dev_0', 1000, 0)] + devices = [Device("dev_0", 1000, 0)] partitioner_config = PartitionerConfig(devices) ret = partitioner.partition_graph(traced, m, partitioner_config) partition = partitioner.partitions[0] @@ -238,10 +248,10 @@ def forward(self, a, b): # Select add_2 node to remove selected_node = None for node in partition.nodes: - if node.name == 'add_2': + if node.name == "add_2": selected_node = node partition.remove_node(selected_node) - assert(partition.used_mem_bytes == 80) + assert partition.used_mem_bytes == 80 def test_size_based_partition(self): class TestModule(torch.nn.Module): @@ -265,7 +275,7 @@ def forward(self, a, b): devices = [ Device("dev_0", 125, 0), Device("dev_1", 125, 1), - Device("dev_2", 125, 2) + Device("dev_2", 125, 2), ] partitioner_config = PartitionerConfig(devices, PartitionMode.size_based) ret = partitioner.partition_graph(traced, m, partitioner_config) @@ -358,7 +368,7 @@ def forward(self, a, b, offset): devices = [ Device("dev_0", 33000000, 0), Device("dev_1", 33000000, 1), - Device("dev_2", 33000000, 2) + Device("dev_2", 33000000, 2), ] partitioner_config = PartitionerConfig(devices, PartitionMode.sparse_nn) partitioner = Partitioner() @@ -443,11 +453,15 @@ def forward(self, a): def get_node_to_latency_mapping(fx_module: GraphModule): node_to_latency_mapping: Dict[Node, Nodelatency] = {} for node in fx_module.graph.nodes: - if node.op not in {'output', 'placeholder', 'get_attr'}: + if node.op not in {"output", "placeholder", "get_attr"}: if node.size_bytes.total_size == node.size_bytes.output_size: - node_to_latency_mapping[node] = NodeLatency(node.size_bytes.total_size, 1) + node_to_latency_mapping[node] = NodeLatency( + node.size_bytes.total_size, 1 + ) else: - node_to_latency_mapping[node] = NodeLatency(node.size_bytes.total_size, node.size_bytes.output_size) + node_to_latency_mapping[node] = NodeLatency( + node.size_bytes.total_size, node.size_bytes.output_size + ) return node_to_latency_mapping m = MyModule() @@ -455,17 +469,17 @@ def get_node_to_latency_mapping(fx_module: GraphModule): a = torch.rand(4) graph_manipulation.get_size_of_all_nodes(traced, [a]) devices = [ - Device('dev_0', 125, 0), - Device('dev_1', 125, 1), - Device('dev_2', 125, 2), - Device('dev_3', 125, 3) + Device("dev_0", 125, 0), + Device("dev_1", 125, 1), + Device("dev_2", 125, 2), + Device("dev_3", 125, 3), ] node_to_latency_mapping = get_node_to_latency_mapping(traced) partitioner_config = PartitionerConfig( devices, mode=PartitionMode.cost_aware, transfer_rate_bytes_per_sec=2, - node_to_latency_mapping=node_to_latency_mapping + node_to_latency_mapping=node_to_latency_mapping, ) partitioner = Partitioner() ret = partitioner.partition_graph(traced, m, partitioner_config) @@ -473,13 +487,15 @@ def get_node_to_latency_mapping(fx_module: GraphModule): dag = ret.dag self.assertEqual(traced(a), module_with_submodules(a)) partitions = partitioner.partitions - partition_to_latency_mapping = get_partition_to_latency_mapping(partitions, node_to_latency_mapping) + partition_to_latency_mapping = get_partition_to_latency_mapping( + partitions, node_to_latency_mapping + ) critical_path_latency_sec = get_latency_of_partitioned_graph( partitions, partition_to_latency_mapping, - partitioner_config.transfer_rate_bytes_per_sec + partitioner_config.transfer_rate_bytes_per_sec, ) - assert critical_path_latency_sec == 160. + assert critical_path_latency_sec == 160.0 def test_kl_based_partition(self): class TestModule(torch.nn.Module): @@ -498,6 +514,7 @@ def forward(self, a): add_4 = add_2 + self.d add_5 = add_3 + add_4 return add_4 + m = TestModule() traced = symbolic_trace(m) a = torch.rand(4) @@ -505,17 +522,17 @@ def forward(self, a): node_to_latency_mapping = get_node_to_latency_mapping(traced) transfer_rate_bytes_per_sec = 2 devices = [ - Device('dev_0', 200, 0), - Device('dev_1', 200, 1), - Device('dev_2', 200, 2), - Device('dev_3', 200, 3) + Device("dev_0", 200, 0), + Device("dev_1", 200, 1), + Device("dev_2", 200, 2), + Device("dev_3", 200, 3), ] partitioner = Partitioner() partitioner_config = PartitionerConfig( devices, mode=PartitionMode.kl_based, transfer_rate_bytes_per_sec=transfer_rate_bytes_per_sec, - node_to_latency_mapping=node_to_latency_mapping + node_to_latency_mapping=node_to_latency_mapping, ) ret = partitioner.partition_graph(traced, m, partitioner_config) module_with_submodules = ret.module_with_submodules @@ -524,15 +541,14 @@ def forward(self, a): assert dag.nodes[0] == 176 assert dag.nodes[1] == 112 partition_to_latency_mapping = get_partition_to_latency_mapping( - partitioner.partitions, - node_to_latency_mapping + partitioner.partitions, node_to_latency_mapping ) cost = get_latency_of_partitioned_graph( partitioner.partitions, partition_to_latency_mapping, - transfer_rate_bytes_per_sec + transfer_rate_bytes_per_sec, ) - assert cost == 208. + assert cost == 208.0 def test_aot_based_partition(self): class TestModule(torch.nn.Module): @@ -545,6 +561,7 @@ def forward(self, a): add_1 = a + self.b add_2 = self.c + add_1 return add_2 + m = TestModule() traced = symbolic_trace(m) a = torch.rand(4) @@ -553,16 +570,16 @@ def forward(self, a): count = 0 GraphManipulation.get_size_of_all_nodes(traced, [a]) for node in traced.graph.nodes: - if node.op not in {'placeholder', 'get_attr', 'output'}: + if node.op not in {"placeholder", "get_attr", "output"}: node_to_partition_id[node] = count partition_to_logical_devices[count] = [0] count += 1 - devices = [Device('dev_0', 200, 0)] + devices = [Device("dev_0", 200, 0)] partitioner_config = PartitionerConfig( devices=devices, mode=PartitionMode.aot_based, node_to_partition_mapping=node_to_partition_id, - partition_to_logical_device_mapping=partition_to_logical_devices + partition_to_logical_device_mapping=partition_to_logical_devices, ) partitioner = Partitioner() ret = partitioner.partition_graph(traced, m, partitioner_config) @@ -577,6 +594,7 @@ def test_replace_target_nodes_with(self): class testModule(torch.nn.Module): def forward(self, a, b): return a + b + m = testModule() traced = symbolic_trace(m) input1 = torch.randn(1) @@ -597,7 +615,9 @@ def test_conv_bn_fusion(self): traced = symbolic_trace(rn18) fused = optimization.fuse(traced) - self.assertTrue(all(not isinstance(m, torch.nn.BatchNorm2d) for m in fused.modules())) + self.assertTrue( + all(not isinstance(m, torch.nn.BatchNorm2d) for m in fused.modules()) + ) N, C, H, W = 20, 3, 224, 224 inp = torch.randn(N, C, H, W) @@ -756,7 +776,9 @@ def mod_partition(node: Node): return partition # split module in module with submodules - module_with_submodules = split_module(my_module_traced, my_module, mod_partition) + module_with_submodules = split_module( + my_module_traced, my_module, mod_partition + ) # Check that test_meta_info was still on all nodes. submodules = dict(module_with_submodules.named_modules()) @@ -805,6 +827,7 @@ def test_normalize_binary_operators(self): # Test Tensor/Tensor callsite for op in ops_to_test: + class WrapperMod(torch.nn.Module): def forward(self, x, y): return op(x, y) @@ -813,11 +836,13 @@ def forward(self, x, y): normalized = NormalizeOperators(traced).transform() x, y = torch.randn(3, 4), torch.randn(3, 4) torch.testing.assert_allclose(traced(x, y), normalized(x, y)) - self.assertFalse(any(n.target in ops_to_test for n in normalized.graph.nodes)) - + self.assertFalse( + any(n.target in ops_to_test for n in normalized.graph.nodes) + ) # Test Tensor/scalar callsite for op in ops_to_test: + class WrapperMod(torch.nn.Module): def forward(self, x): return op(x, 42) @@ -826,14 +851,18 @@ def forward(self, x): normalized = NormalizeOperators(traced).transform() x = torch.randn(3, 4) torch.testing.assert_allclose(traced(x), normalized(x)) - self.assertFalse(any(n.target in ops_to_test for n in normalized.graph.nodes)) + self.assertFalse( + any(n.target in ops_to_test for n in normalized.graph.nodes) + ) @skipIfNoTorchVision def test_normalize_args(self): m = resnet18() class FunctionalTracer(torch.fx.Tracer): - def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> bool: + def is_leaf_module( + self, m: torch.nn.Module, module_qualified_name: str + ) -> bool: # `leaves` contains the set of standard `nn.Modules` that are not # currently symbolically traceable. Ideally this set would be empty leaves = set([torch.nn.BatchNorm2d]) @@ -847,13 +876,12 @@ def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> boo ShapeProp(traced).propagate(input) traced = NormalizeArgs(traced).transform() - modules = dict(traced.named_modules()) for node in traced.graph.nodes: - if node.op == 'call_function' and node.target != operator.add: + if node.op == "call_function" and node.target != operator.add: self.assertEqual(len(node.args), 0) - elif node.op == 'call_module': + elif node.op == "call_module": submod_class = modules[node.target].__class__ nn_class = getattr(torch.nn, submod_class.__name__) if submod_class == nn_class: @@ -867,15 +895,15 @@ def test_normalize_modules_exhaustive(self): torch.nn Module classes """ for test_params in module_tests + new_module_tests: - if 'constructor' not in test_params: - constructor = getattr(torch.nn, test_params['module_name']) + if "constructor" not in test_params: + constructor = getattr(torch.nn, test_params["module_name"]) else: - constructor = test_params['constructor'] + constructor = test_params["constructor"] - if 'constructor_args' not in test_params: + if "constructor_args" not in test_params: args = () else: - args = test_params['constructor_args'] + args = test_params["constructor_args"] mod = constructor(*args) # Skip modules that are not standard `torch.nn` @@ -884,18 +912,18 @@ def test_normalize_modules_exhaustive(self): if mod.__class__.__name__ not in dir(torch.nn): continue - if 'input_fn' not in test_params: - inputs = torch.randn(test_params['input_size']) + if "input_fn" not in test_params: + inputs = torch.randn(test_params["input_size"]) else: - inputs = test_params['input_fn']() + inputs = test_params["input_fn"]() if not isinstance(inputs, (tuple, list)): inputs = (inputs,) - params = ', '.join(f'v{i}' for i in range(len(inputs))) + params = ", ".join(f"v{i}" for i in range(len(inputs))) # Generate a class to wrap this standard `nn.Module` instance - test_classname = f'Test{mod.__class__.__name__}' + test_classname = f"Test{mod.__class__.__name__}" test_mod_code = f""" class {test_classname}(torch.nn.Module): def __init__(self, mod): @@ -906,7 +934,7 @@ def forward(self, {params}): return self.mod({params}) """ - gbls = {'torch' : torch} + gbls = {"torch": torch} exec(test_mod_code, gbls) test_instance = gbls[test_classname](mod) @@ -917,13 +945,15 @@ def forward(self, {params}): # in those arguments as kwargs modules = dict(traced.named_modules()) for node in traced.graph.nodes: - if node.op == 'call_module': + if node.op == "call_module": submod_class = modules[node.target].__class__ nn_class = getattr(torch.nn, submod_class.__name__) if submod_class == nn_class: normalized_args = node.normalized_arguments(traced) - normalized_args2 = normalize_module(traced, node.target, node.args, node.kwargs) - assert(normalized_args == normalized_args2) + normalized_args2 = normalize_module( + traced, node.target, node.args, node.kwargs + ) + assert normalized_args == normalized_args2 assert normalized_args node.args = normalized_args.args node.kwargs = normalized_args.kwargs @@ -933,8 +963,7 @@ def forward(self, {params}): # These Modules have an RNG in their forward, so testing # correctness by comparing outputs is not correct. Skip that # check for these - stochastic_modules = {'FractionalMaxPool2d', 'FractionalMaxPool3d', - 'RReLU'} + stochastic_modules = {"FractionalMaxPool2d", "FractionalMaxPool3d", "RReLU"} if mod.__class__.__name__ not in stochastic_modules: self.assertEqual(traced(*inputs), mod(*inputs)) @@ -942,14 +971,12 @@ def forward(self, {params}): traced = NormalizeArgs(symbolic_trace(test_instance)).transform() modules = dict(traced.named_modules()) for node in traced.graph.nodes: - if node.op == 'call_module': + if node.op == "call_module": submod_class = modules[node.target].__class__ nn_class = getattr(torch.nn, submod_class.__name__) if submod_class == nn_class: self.assertEqual(len(node.args), 0) - - @skipIfNoTorchVision def test_annotate_returns_with_schema(self): m = resnet18() @@ -959,14 +986,23 @@ def test_annotate_returns_with_schema(self): for node in traced_modules_annotated.graph.nodes: if node.type is None: check = (node.op, node.target) - self.assertTrue(check in {('placeholder', 'x'), ('call_function', operator.add), - ('call_function', torch.flatten), ('output', 'output')}) + self.assertTrue( + check + in { + ("placeholder", "x"), + ("call_function", operator.add), + ("call_function", torch.flatten), + ("output", "output"), + } + ) # Smoke test torchscript compilation since now we're emitting type annotations torch.jit.script(traced_modules_annotated) class FunctionalTracer(torch.fx.Tracer): - def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> bool: + def is_leaf_module( + self, m: torch.nn.Module, module_qualified_name: str + ) -> bool: # `leaves` contains the set of standard `nn.Modules` that are not # currently symbolically traceable. Ideally this set would be empty leaves = set([torch.nn.BatchNorm2d]) @@ -974,18 +1010,20 @@ def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> boo traced_functionals = torch.fx.GraphModule(m, FunctionalTracer().trace(m)) - traced_functionals_annotated = AnnotateTypesWithSchema(traced_functionals).transform() + traced_functionals_annotated = AnnotateTypesWithSchema( + traced_functionals + ).transform() for node in traced_functionals_annotated.graph.nodes: if node.type is None: check = (node.op, node.target) excluded_nodes = { - ('placeholder', 'x'), - ('call_function', torch.conv2d), + ("placeholder", "x"), + ("call_function", torch.conv2d), # Return type differs based on boolean dispatch :( - ('call_function', torch.nn.functional.max_pool2d), - ('call_function', operator.add), - ('call_function', torch.flatten), - ('output', 'output'), + ("call_function", torch.nn.functional.max_pool2d), + ("call_function", operator.add), + ("call_function", torch.flatten), + ("output", "output"), } self.assertTrue(check in excluded_nodes) @@ -1011,11 +1049,12 @@ def forward(self, a, b, c, d): mm = MyModule() traced = symbolic_trace(mm) - def split_cb(node : torch.fx.Node): - if node.name == 'a' or node.name == 'b' or node.name == 'add': + def split_cb(node: torch.fx.Node): + if node.name == "a" or node.name == "b" or node.name == "add": return 0 else: return 1 + module_with_submodule = split_module(traced, mm, split_cb) self.assertEqual(module_with_submodule(a, b, c, d), traced(a, b, c, d)) @@ -1033,22 +1072,26 @@ def __init__(self): self.seq = torch.nn.Sequential(torch.nn.BatchNorm1d(2, 2)) self.linear = torch.nn.Linear(2, 2) self.attr = torch.randn(2) - self.register_buffer('attr2', torch.randn(2)) + self.register_buffer("attr2", torch.randn(2)) def forward(self, x): return self.linear(self.seq(self.W + self.attr + self.attr2 + x)) mod = symbolic_trace(Test()) - module_name = 'Foo' + module_name = "Foo" import tempfile from pathlib import Path + with tempfile.TemporaryDirectory() as tmp_dir: tmp_dir = Path(tmp_dir) mod.to_folder(tmp_dir, module_name) # Recipe taken from here: # https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly import importlib.util - spec = importlib.util.spec_from_file_location(module_name, tmp_dir / '__init__.py') + + spec = importlib.util.spec_from_file_location( + module_name, tmp_dir / "__init__.py" + ) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) @@ -1058,10 +1101,21 @@ def forward(self, x): def test_fetch(self): attrs_for_lowering: Dict[str, List[str]] = { "torch.nn.modules.conv.Conv2d": [ - "weight", "bias", "kernel_size", "stride", "padding", "dilation", "groups", "padding_mode" + "weight", + "bias", + "kernel_size", + "stride", + "padding", + "dilation", + "groups", + "padding_mode", ], "torch.nn.modules.batchnorm.BatchNorm2d": [ - "weight", "bias", "running_mean", "running_var", "eps" + "weight", + "bias", + "running_mean", + "running_var", + "eps", ], } @@ -1225,12 +1279,18 @@ def test_type_matches(self): (List[int], create_type_hint([int, int])), (List[int], create_type_hint((int, int))), (List[torch.Tensor], create_type_hint([torch.Tensor, torch.Tensor])), - (List[torch.Tensor], create_type_hint([torch.nn.Parameter, torch.nn.Parameter])), + ( + List[torch.Tensor], + create_type_hint([torch.nn.Parameter, torch.nn.Parameter]), + ), (torch.Tensor, torch.nn.Parameter), (List[torch.Tensor], create_type_hint([torch.nn.Parameter, torch.Tensor])), (List[torch.Tensor], create_type_hint([torch.Tensor, torch.nn.Parameter])), (List[torch.Tensor], create_type_hint((torch.Tensor, torch.Tensor))), - (List[torch.Tensor], create_type_hint((torch.nn.Parameter, torch.nn.Parameter))), + ( + List[torch.Tensor], + create_type_hint((torch.nn.Parameter, torch.nn.Parameter)), + ), (torch.Tensor, torch.nn.Parameter), (List[torch.Tensor], create_type_hint((torch.nn.Parameter, torch.Tensor))), (List[torch.Tensor], create_type_hint((torch.Tensor, torch.nn.Parameter))), @@ -1243,7 +1303,7 @@ def test_type_matches(self): should_fail = [ (int, float), (Union[int, float], str), - (List[torch.Tensor], List[int]) + (List[torch.Tensor], List[int]), ] for sig_type, arg_type in should_fail: @@ -1272,19 +1332,23 @@ def __init__(self): def forward(self, x): return self.model(x) + self.model2(x) - - N, C, H, W, = 1, 3, 224, 224 + N, C, H, W, = ( + 1, + 3, + 224, + 224, + ) inp = torch.randn(N, C, H, W) with torch.no_grad(): model = Foo().eval() optimized_model = optimization.optimize_for_inference(model) torch.testing.assert_allclose(model(inp), optimized_model(inp)) - optimized_model2 = \ - optimization.optimize_for_inference(model, pass_config={"remove_dropout": False}) + optimized_model2 = optimization.optimize_for_inference( + model, pass_config={"remove_dropout": False} + ) torch.testing.assert_allclose(model(inp), optimized_model2(inp)) - @skipIfNoTorchVision @skipIfNoMkldnn def test_optimize_for_inference_cpu_torchvision(self): @@ -1296,12 +1360,16 @@ def test_optimize_for_inference_cpu_torchvision(self): torchvision.models.vgg16, torchvision.models.mobilenet_v2, torchvision.models.mnasnet1_0, - torchvision.models.resnext50_32x4d + torchvision.models.resnext50_32x4d, ] with torch.no_grad(): for model_type in models: model = model_type() - C, H, W, = 3, 224, 224 + C, H, W, = ( + 3, + 224, + 224, + ) inp = torch.randn(3, C, H, W) model(inp) model.eval() @@ -1319,40 +1387,41 @@ class TestNormalizeOperators(JitTestCase): @ops(op_db, allowed_dtypes=(torch.float,)) def test_normalize_operator_exhaustive(self, device, dtype, op): # Sorted and one entry on each line to minimize merge conflicts. - op_skip = {'contiguous', - 'einsum', - 'expand', - 'expand_as', - 'fill_', - 'gradient', - 'index_put', - 'polygamma', - 'repeat', - 'reshape_as', - 'resize_', - 'resize_as_', - 'to_sparse', - 'view', - 'view_as', - 'unfold', - 'where', - 'zero_', - '__getitem__', - '__radd__', - '__rsub__', - '__rmul__', - '__rdiv__', - '__rmod__', - '__rpow__', - '__rmatmul__'} + op_skip = { + "contiguous", + "einsum", + "expand", + "expand_as", + "fill_", + "gradient", + "index_put", + "polygamma", + "repeat", + "reshape_as", + "resize_", + "resize_as_", + "to_sparse", + "view", + "view_as", + "unfold", + "where", + "zero_", + "__getitem__", + "__radd__", + "__rsub__", + "__rmul__", + "__rdiv__", + "__rmod__", + "__rpow__", + "__rmatmul__", + } # Unsupported input types if op.name in op_skip: return # These ops currently don't trace in FX for various reasons (i.e. they take a list of tensors) - fx_fail = {'stack', 'hstack', 'vstack', 'dstack', - 'linalg.multi_dot'} + fx_fail = {"stack", "hstack", "vstack", "dstack", "linalg.multi_dot"} sample_inputs_itr = op.sample_inputs(device, dtype, requires_grad=False) for sample_input in sample_inputs_itr: unsupported_arg_type = False @@ -1389,14 +1458,17 @@ def jit_infer_type(v): continue # Test normalize_function by itself ref_out = op.op(*arg_values, **kwarg_values) - norm_args_and_kwargs = normalize_function(op.op, arg_values, kwarg_values, arg_types, kwarg_types) + norm_args_and_kwargs = normalize_function( + op.op, arg_values, kwarg_values, arg_types, kwarg_types + ) if norm_args_and_kwargs is None: raise RuntimeError( """ FX failed to normalize op - add the op to the op_skip list. A common reason is if your OpInfo was implemented with a lambda - otherwise, file an issue - """) + """ + ) test_out = op.op(*norm_args_and_kwargs.args, **norm_args_and_kwargs.kwargs) self.assertEqual(test_out, ref_out) @@ -1412,15 +1484,15 @@ def jit_infer_type(v): param_values.append(v) fx_args.append(param_names[-1]) else: - fx_args.append(f'{repr(v)}') + fx_args.append(f"{repr(v)}") for k, v in kwarg_values.items(): if isinstance(v, torch.Tensor): param_names.append(k) param_values.append(v) - fx_args.append(f'{k} = {k}') + fx_args.append(f"{k} = {k}") else: - fx_args.append(f'{k} = {repr(v)}') + fx_args.append(f"{k} = {repr(v)}") code = f""" class TestModule(torch.nn.Module): @@ -1428,18 +1500,19 @@ def forward(self, {', '.join(param_names)}): return torch.{op.name}({', '.join(fx_args)}) """ - g = {'torch': torch, 'inf' : math.inf} + g = {"torch": torch, "inf": math.inf} exec(code, g) - TestModule = g['TestModule'] - + TestModule = g["TestModule"] m = TestModule() traced = torch.fx.symbolic_trace(m) ref_out = traced(*param_values) for node in traced.graph.nodes: - if node.op == 'call_function': - normalized_args = node.normalized_arguments(traced, arg_types, kwarg_types) + if node.op == "call_function": + normalized_args = node.normalized_arguments( + traced, arg_types, kwarg_types + ) assert normalized_args node.args = normalized_args.args node.kwargs = normalized_args.kwargs @@ -1448,6 +1521,7 @@ def forward(self, {', '.join(param_names)}): test_out = traced(*param_values) self.assertEqual(test_out, ref_out) + instantiate_device_type_tests(TestNormalizeOperators, globals()) if __name__ == "__main__": diff --git a/torch/fx/experimental/accelerator_partitioner.py b/torch/fx/experimental/accelerator_partitioner.py index ce7165ea9c532..c16ba8c097957 100644 --- a/torch/fx/experimental/accelerator_partitioner.py +++ b/torch/fx/experimental/accelerator_partitioner.py @@ -1,26 +1,35 @@ -from torch.fx.graph_module import GraphModule -from torch.fx.node import Node, map_arg +import operator from typing import Dict, List, Set, NamedTuple, Tuple + import torch -from torch.fx.passes.split_module import split_module -import operator from torch.fx.experimental.graph_manipulation import get_size_of_all_nodes -from torch.fx.experimental.partitioner_utils import Partition, \ - Device, PartitionerConfig, get_partition_to_latency_mapping,\ - get_latency_of_partitioned_graph, NodeLatency, get_extra_size_of, \ - PartitionMode +from torch.fx.experimental.partitioner_utils import ( + Partition, + Device, + PartitionerConfig, + get_partition_to_latency_mapping, + get_latency_of_partitioned_graph, + NodeLatency, + get_extra_size_of, + PartitionMode, +) +from torch.fx.graph_module import GraphModule +from torch.fx.node import Node, map_arg +from torch.fx.passes.split_module import split_module + -class DAGNode(): +class DAGNode: """DAGNode class maintains useful information for a partition (submodule), - and its input submodules and output submodules. + and its input submodules and output submodules. """ + def __init__( self, submodule_node: Node, input_nodes: List[Node], output_nodes: List[Node], logical_device_ids: List[int], - size_bytes: int + size_bytes: int, ) -> None: self.submodule_node: Node = submodule_node self.input_nodes: List[Node] = input_nodes @@ -31,8 +40,10 @@ def __init__( def __str__(self) -> str: return str(self.submodule_node) + class DAG: """DAG class contains all the DAG nodes""" + def __init__(self) -> None: self.nodes: List[DAGNode] = [] @@ -42,30 +53,35 @@ def create_node( input_nodes: List[Node], output_nodes: List[Node], logical_devices: List[int], - size_bytes: int + size_bytes: int, ) -> None: - node = DAGNode(submodule_node, input_nodes, output_nodes, logical_devices, size_bytes) + node = DAGNode( + submodule_node, input_nodes, output_nodes, logical_devices, size_bytes + ) self.nodes.append(node) + class PartitionResult(NamedTuple): - """NameTuple used for returning DAG and a new fx module - """ + """NameTuple used for returning DAG and a new fx module""" + dag: DAG module_with_submodules: GraphModule + """Followings are some helper functions for partition manipulation""" + + def reset_partition_device(partitions): for partition in partitions: partition.logical_device_ids = [] + def combine_two_partitions( - partition_0: Partition, - partition_1: Partition, - partitions: List[Partition] + partition_0: Partition, partition_1: Partition, partitions: List[Partition] ) -> None: """Given a list of partitions and its two partitions, - combine these two partitions into a new one appending to the partitions - and remove the previous two partitions from the list of partitions + combine these two partitions into a new one appending to the partitions + and remove the previous two partitions from the list of partitions """ partition = Partition(len(partitions)) partition.nodes = partition_0.nodes.union(partition_1.nodes) @@ -76,9 +92,9 @@ def combine_two_partitions( reorganize_partitions(partitions) return + def set_parents_and_children(partitions: List[Partition]) -> None: - """Given a list of partitions, mark parents and children for each partition - """ + """Given a list of partitions, mark parents and children for each partition""" # Go through all nodes in a partition. # If a node's user is in other partition, # then the other partition is this partition's children. @@ -100,9 +116,10 @@ def set_parents_and_children(partitions: List[Partition]) -> None: p.parents.add(partition) return + def reorganize_partitions(partitions: List[Partition]) -> None: """Given a list of partitions, reorganzie partiton id, - its parents and its children for each partition + its parents and its children for each partition """ # Rearrange partition ids for i, partition in enumerate(partitions): @@ -110,9 +127,10 @@ def reorganize_partitions(partitions: List[Partition]) -> None: set_parents_and_children(partitions) return + def get_bfs_level_partition(partitions: List[Partition]) -> None: """Given a list of partitions, - mark the bfs level for each partition + mark the bfs level for each partition """ current_level: Set[Partition] = set() visited: Set[Partition] = set() @@ -137,20 +155,26 @@ def get_bfs_level_partition(partitions: List[Partition]) -> None: level += 1 return + def get_node_to_partition_mapping(partitions: List[Partition]) -> Dict[Node, int]: - """Given a list of partitions,return node to partition mapping - """ + """Given a list of partitions,return node to partition mapping""" node_to_partition: Dict[Node, int] = {} for partition in partitions: for node in partition.nodes: node_to_partition[node] = partition.partition_id return node_to_partition -def get_device_to_partitions_mapping(partitions: List[Partition], devices: List[Device]): + +def get_device_to_partitions_mapping( + partitions: List[Partition], devices: List[Device] +): """Given a list of partitions and a list of devices, - map each partition into a device. + map each partition into a device. """ - def calculate_extra_mem_bytes_needed_for(partition: Partition, partitions: List[Partition]): + + def calculate_extra_mem_bytes_needed_for( + partition: Partition, partitions: List[Partition] + ): all_nodes: Set[Node] = set() for p in partitions: all_nodes = all_nodes.union(p.nodes) @@ -164,19 +188,22 @@ def calculate_extra_mem_bytes_needed_for(partition: Partition, partitions: List[ def find_device_for(partition: Partition): """Given a partition, find a logical device for the partition - The algorithm is to put the partition on the device - that has just enough mem left for that partition. - device_to_left_mem_bytes is a dictionary between device and its left mem size - sorted by its left mem size + The algorithm is to put the partition on the device + that has just enough mem left for that partition. + device_to_left_mem_bytes is a dictionary between device and its left mem size + sorted by its left mem size """ for d in device_to_left_mem_bytes: - extra_size_needed = calculate_extra_mem_bytes_needed_for(partition, device_to_partitions[d]) + extra_size_needed = calculate_extra_mem_bytes_needed_for( + partition, device_to_partitions[d] + ) if extra_size_needed < device_to_left_mem_bytes[d]: device_to_partitions[d].append(partition) partition.logical_device_ids.append(d.logical_id) device_to_left_mem_bytes[d] -= extra_size_needed return True return False + # logical id to device logical_id_to_device: Dict[int, Device] = {} # Track partitions on device @@ -195,24 +222,29 @@ def find_device_for(partition: Partition): logical_id = partition.logical_device_ids[0] device = logical_id_to_device[logical_id] device_to_partitions[device] = [partition] - device_to_left_mem_bytes[device] = d.available_mem_bytes - partition.used_mem_bytes + device_to_left_mem_bytes[device] = ( + d.available_mem_bytes - partition.used_mem_bytes + ) else: no_device_partitions.append(partition) # Find devices for all the partitions without a device found_device = True for partition in no_device_partitions: device_to_left_mem_bytes = { - d: left_mem_bytes for d, left_mem_bytes - in sorted(device_to_left_mem_bytes.items(), key=lambda item: item[1]) + d: left_mem_bytes + for d, left_mem_bytes in sorted( + device_to_left_mem_bytes.items(), key=lambda item: item[1] + ) } found_device = find_device_for(partition) if not found_device: break return found_device + def check_dependency(partition): """Given a partition,check if there is a circular dependency on - this partition using bfs + this partition using bfs """ visited: Set[Partition] = set([partition]) queue: List[Partition] = [partition] @@ -227,15 +259,17 @@ def check_dependency(partition): queue.append(child) return False + class Partitioner: """A fx module may not fit into one device. - Partitioner class helps partition one fx module into submodules (partitions), - so that the submodules can be executed crossing different accelerators. - The main function of this class is self.partition_graph. - It partitions the fx module based on the scheme specified in partition_config - A DAG structure is returned - along with a new fx module with submodule nodes. + Partitioner class helps partition one fx module into submodules (partitions), + so that the submodules can be executed crossing different accelerators. + The main function of this class is self.partition_graph. + It partitions the fx module based on the scheme specified in partition_config + A DAG structure is returned + along with a new fx module with submodule nodes. """ + def __init__(self) -> None: self.partitions: List[Partition] = [] self.node_to_partition: Dict[Node, int] = {} @@ -245,27 +279,27 @@ def partition_graph( self, fx_module: GraphModule, torch_module: torch.nn.Module, - partitioner_config: PartitionerConfig + partitioner_config: PartitionerConfig, ) -> PartitionResult: """Given the fx module, torch module and partitioner_config, - find the partitions, do the partitions, - and then return a DAG and a new fx module with submodule nodes (partitions) + find the partitions, do the partitions, + and then return a DAG and a new fx module with submodule nodes (partitions) """ self.graph_module = fx_module self.torch_module = torch_module self.devices = partitioner_config.devices if len(self.devices) == 0: - raise RuntimeError('No devices') + raise RuntimeError("No devices") # Tag the size in bytes to all nodes in the graph_module. get_size_of_all_nodes(self.graph_module) # Check if there are op nodes in the fx module nodes = self.graph_module.graph.nodes - if all(node.op in {'placeholder', 'get_attr', 'output'} for node in nodes): - raise RuntimeError('No Partition since no operations in the module') + if all(node.op in {"placeholder", "get_attr", "output"} for node in nodes): + raise RuntimeError("No Partition since no operations in the module") # Calculate total size of the fx module total_size_of_graph = 0 for node in nodes: - if node.op == 'output': + if node.op == "output": break total_size_of_graph += node.size_bytes.total_size # Find the device with the max mem size @@ -274,19 +308,22 @@ def partition_graph( if partitioner_config.mode == PartitionMode.aot_based: self.aot_based_partition( partitioner_config.node_to_partition_mapping, - partitioner_config.partition_to_logical_device_mapping + partitioner_config.partition_to_logical_device_mapping, ) # Single partition if the whole module can be fit into one device elif total_size_of_graph <= device_with_max_mem.available_mem_bytes: self.find_single_partition(total_size_of_graph) elif total_size_of_graph > sum([d.available_mem_bytes for d in self.devices]): - raise RuntimeError('Devices have no enough memory for the module') + raise RuntimeError("Devices have no enough memory for the module") else: # Sparse nn based partition if partitioner_config.mode == PartitionMode.sparse_nn: available_mem_bytes = self.devices[0].available_mem_bytes - if not all(device.available_mem_bytes == available_mem_bytes for device in self.devices): - raise RuntimeError('All devices must have same memory size!') + if not all( + device.available_mem_bytes == available_mem_bytes + for device in self.devices + ): + raise RuntimeError("All devices must have same memory size!") # sparse_nn_partition only support same memory size # TODO: add different size support for sparse_nn_partition self.sparse_nn_partition(available_mem_bytes) @@ -294,13 +331,13 @@ def partition_graph( elif partitioner_config.mode == PartitionMode.cost_aware: self.cost_aware_partition( partitioner_config.transfer_rate_bytes_per_sec, - partitioner_config.node_to_latency_mapping + partitioner_config.node_to_latency_mapping, ) # KL based partition elif partitioner_config.mode == PartitionMode.kl_based: self.kl_based_partition( partitioner_config.transfer_rate_bytes_per_sec, - partitioner_config.node_to_latency_mapping + partitioner_config.node_to_latency_mapping, ) else: self.size_based_partition() @@ -312,11 +349,10 @@ def partition_graph( return ret def find_single_partition(self, total_size_of_graph) -> None: - """Fit the whole fx module into one device - """ + """Fit the whole fx module into one device""" partition_0 = self.create_partition() for node in self.graph_module.graph.nodes: - if node.op == 'output': + if node.op == "output": break partition_0.nodes.add(node) partition_0.used_mem_bytes = total_size_of_graph @@ -327,30 +363,34 @@ def find_single_partition(self, total_size_of_graph) -> None: def size_based_partition(self) -> None: """This method is to partition the fx module based on memory size. - It uses greedy approach. The result may not be the best. - The basic idea is: - Step 1: - Find a device which has enough memory to fit the current node, create a empty partition - with the size of that device. - Then keep adding the following nodes into the partition until the partition is full. - Step 2: - Repeat Step 1 until no device left - Step 3: - If some nodes are left, create a partition for each left node (single node partition). - and then try to map those partitions into logical devices with enough mem left. + It uses greedy approach. The result may not be the best. + The basic idea is: + Step 1: + Find a device which has enough memory to fit the current node, create a empty partition + with the size of that device. + Then keep adding the following nodes into the partition until the partition is full. + Step 2: + Repeat Step 1 until no device left + Step 3: + If some nodes are left, create a partition for each left node (single node partition). + and then try to map those partitions into logical devices with enough mem left. """ + def find_device_based_on_size(node) -> Device: """Given a node, this function is to find a logical device - that could fit the node. + that could fit the node. """ mem_size_needed = get_extra_size_of(node, set()) - device = Device('', -1, -1) + device = Device("", -1, -1) for d in self.devices: - if d not in occupied_devices and d.available_mem_bytes >= mem_size_needed: + if ( + d not in occupied_devices + and d.available_mem_bytes >= mem_size_needed + ): device = d break if device.available_mem_bytes < 0: - raise RuntimeError(str(node) + 'is too large to fit any device') + raise RuntimeError(str(node) + "is too large to fit any device") occupied_devices.append(device) return device @@ -360,7 +400,7 @@ def find_device_based_on_size(node) -> Device: occupied_devices: List[Device] = [] partition = self.create_partition() for node in self.graph_module.graph.nodes: - if node.op in {'call_module', 'call_method', 'call_function'}: + if node.op in {"call_module", "call_method", "call_function"}: # Check if there are devices left if len(self.partitions) <= len(self.devices): total_size_of_input_nodes = get_extra_size_of(node, partition.nodes) @@ -370,13 +410,18 @@ def find_device_based_on_size(node) -> Device: device = find_device_based_on_size(node) occupied_devices.append(device) # Update partition and its left mem size - partition_to_left_mem_bytes[partition] = device.available_mem_bytes + partition_to_left_mem_bytes[ + partition + ] = device.available_mem_bytes # Update available mem for the current partitio partition.logical_device_ids.append(device.logical_id) else: # The current partition is not the first partition # Check if the current node can fit into current partition - if partition_to_left_mem_bytes[partition] < total_size_of_input_nodes: + if ( + partition_to_left_mem_bytes[partition] + < total_size_of_input_nodes + ): # Check if no device is left if len(self.partitions) == len(self.devices): # No device is left @@ -389,8 +434,12 @@ def find_device_based_on_size(node) -> Device: # Create a new partition with a mem size that is enough for the current node device = find_device_based_on_size(node) partition = self.create_partition() - total_size_of_input_nodes = get_extra_size_of(node, partition.nodes) - partition_to_left_mem_bytes[partition] = device.available_mem_bytes + total_size_of_input_nodes = get_extra_size_of( + node, partition.nodes + ) + partition_to_left_mem_bytes[ + partition + ] = device.available_mem_bytes partition.logical_device_ids.append(device.logical_id) partition.add_node(node) partition_to_left_mem_bytes[partition] -= total_size_of_input_nodes @@ -401,7 +450,9 @@ def find_device_based_on_size(node) -> Device: # Get the node to partition mapping self.node_to_partition = get_node_to_partition_mapping(self.partitions) # Mapping all partitions into device - found_partition_to_device_mapping = get_device_to_partitions_mapping(self.partitions, self.devices) + found_partition_to_device_mapping = get_device_to_partitions_mapping( + self.partitions, self.devices + ) if not found_partition_to_device_mapping: raise RuntimeError("Cannot Get a Valid Partition to Logical Device Mapping") return @@ -411,7 +462,7 @@ def do_partition(self) -> GraphModule: module_with_submodules = split_module( self.graph_module, self.torch_module, - lambda node: self.node_to_partition[node] + lambda node: self.node_to_partition[node], ) return module_with_submodules @@ -419,13 +470,13 @@ def dump_dag(self, module_with_submodules: GraphModule) -> DAG: """Return the dag structure and the new fx module with submodules""" dag = DAG() for node in module_with_submodules.graph.nodes: - if node.op == 'output': + if node.op == "output": break - if node.op in {'placeholder', 'get_attr'}: + if node.op in {"placeholder", "get_attr"}: continue if node.target == operator.__getitem__: continue - input_nodes : Dict[Node, None] = {} + input_nodes: Dict[Node, None] = {} map_arg(node.args, lambda n: input_nodes.setdefault(n)) map_arg(node.kwargs, lambda n: input_nodes.setdefault(n)) # When a node has two or more output nodes, @@ -436,10 +487,12 @@ def dump_dag(self, module_with_submodules: GraphModule) -> DAG: output_nodes = list(node.users) else: output_nodes = [node] - partition_id = int(node.name.rsplit('_', 1)[-1]) + partition_id = int(node.name.rsplit("_", 1)[-1]) device_ids = self.partitions[partition_id].logical_device_ids size_bytes = self.partitions[partition_id].used_mem_bytes - dag.create_node(node, list(input_nodes), output_nodes, device_ids, size_bytes) + dag.create_node( + node, list(input_nodes), output_nodes, device_ids, size_bytes + ) return dag def create_partition(self) -> Partition: @@ -457,35 +510,38 @@ def create_single_node_partition(self, node): def sparse_nn_partition(self, available_mem_bytes: int) -> None: """This method partition a sparse nn module. - It is size based partition but different from size_based_partition, - it only works when all the devices have same memory size (available_mem_bytes). - In the future, devices with different mem sizes will be supported like size_based_partition. - It first traverse all the nodes and do the partitions based on the same memory size. - If the current partition has no enough memory left for a new op node - (call_module, call_method, call_function), a new partition is created. - When crossing the boundary between non-embedding nodes and embedding nodes, - a new partition is created regardlessly. - For example, if the current node is a non-embedding node but the next node is an - embedding node, a new partition is created for the next node. - After the partition, the partitions are combined as much as possible. - The rule is that a non-embedding partition only - combines with another non-embedding one. - So as the embedding partitions. + It is size based partition but different from size_based_partition, + it only works when all the devices have same memory size (available_mem_bytes). + In the future, devices with different mem sizes will be supported like size_based_partition. + It first traverse all the nodes and do the partitions based on the same memory size. + If the current partition has no enough memory left for a new op node + (call_module, call_method, call_function), a new partition is created. + When crossing the boundary between non-embedding nodes and embedding nodes, + a new partition is created regardlessly. + For example, if the current node is a non-embedding node but the next node is an + embedding node, a new partition is created for the next node. + After the partition, the partitions are combined as much as possible. + The rule is that a non-embedding partition only + combines with another non-embedding one. + So as the embedding partitions. """ - def combine_partitions_based_on_size(partitions: List[Partition], available_mem_bytes: int) -> None: + + def combine_partitions_based_on_size( + partitions: List[Partition], available_mem_bytes: int + ) -> None: """Combining small partitions together to keep as less partitions as possible. - Here is an example of the algorithm to do this: - Assume some partitions, we first sort them based on partiiton used memory size. - [(partition_4, 1), (partition_3, 1), (partition_2, 2), (partition_1, 7), (partition_0, 9)] - The available memory is 10. - step 1: self.find_partition_to_combine_based_on_size() - First, mark bfs level for each partition - Second, look the smallest partition, partition_4: 10 - 1 = 9 - It means any partition has a used memory equal or less than 9 could combine this partition - We go from the largest and selection partition_0. - Check the bfs level for two partitions, if the level difference is less than 2, - it can be combined. - step 2: repeat step 1 until no partitions can be combined + Here is an example of the algorithm to do this: + Assume some partitions, we first sort them based on partiiton used memory size. + [(partition_4, 1), (partition_3, 1), (partition_2, 2), (partition_1, 7), (partition_0, 9)] + The available memory is 10. + step 1: self.find_partition_to_combine_based_on_size() + First, mark bfs level for each partition + Second, look the smallest partition, partition_4: 10 - 1 = 9 + It means any partition has a used memory equal or less than 9 could combine this partition + We go from the largest and selection partition_0. + Check the bfs level for two partitions, if the level difference is less than 2, + it can be combined. + step 2: repeat step 1 until no partitions can be combined """ find_combination = True while find_combination: @@ -493,17 +549,14 @@ def combine_partitions_based_on_size(partitions: List[Partition], available_mem_ sorted_partitions = sorted(partitions, key=lambda p: p.used_mem_bytes) # Mark bfs level get_bfs_level_partition(self.partitions) - find_combination, partitions = \ - find_partition_to_combine_based_on_size( - sorted_partitions, - available_mem_bytes, - partitions - ) + find_combination, partitions = find_partition_to_combine_based_on_size( + sorted_partitions, available_mem_bytes, partitions + ) return def calculate_mem_bytes_needed(p1, p2): """Given two partitions, calculate how many mem bytes - are needed if two partitions are combined + are needed if two partitions are combined """ nodes = p1.nodes.union(p2.nodes) mem_bytes_needed = 0 @@ -514,7 +567,7 @@ def calculate_mem_bytes_needed(p1, p2): def find_partition_to_combine_based_on_size( sorted_partitions: List[Partition], available_mem_bytes: int, - partitions: List[Partition] + partitions: List[Partition], ) -> Tuple[bool, List[Partition]]: """step 1 in combine_partition_based_on_size()""" find_combination = False @@ -534,7 +587,7 @@ def find_partition_to_combine_based_on_size( def reset_partition_in_sparse_nn(partition, new_partition=True): """If crossing the boudary between non-embedding nodes and - embedding nodes, create a new partition + embedding nodes, create a new partition """ if in_embedding_region: embedding_partitions.append(partition) @@ -548,13 +601,15 @@ def reset_partition_in_sparse_nn(partition, new_partition=True): def is_embedding_node(node: Node) -> bool: """Check if a node is an embedding node""" - if node.op == 'call_module': + if node.op == "call_module": submodule = self.graph_module - for atom in str(node.target).split('.'): + for atom in str(node.target).split("."): if not hasattr(submodule, atom): - raise RuntimeError(f'Module {submodule} has no attribute {atom}') + raise RuntimeError( + f"Module {submodule} has no attribute {atom}" + ) submodule = getattr(submodule, atom) - if 'Embedding' in str(submodule): + if "Embedding" in str(submodule): return True return False @@ -565,7 +620,7 @@ def is_embedding_node(node: Node) -> bool: in_embedding_region: bool = False partition = self.create_partition() for node in self.graph_module.graph.nodes: - if node.op in {'call_module', 'call_method', 'call_function'}: + if node.op in {"call_module", "call_method", "call_function"}: # Check if crossing the boundary between embedding nodes and non embedding nodes if is_embedding_node(node) != in_embedding_region: # Crossing the boundary @@ -575,11 +630,16 @@ def is_embedding_node(node: Node) -> bool: partition = reset_partition_in_sparse_nn(partition) in_embedding_region = not in_embedding_region total_size_of_input_nodes = get_extra_size_of(node, partition.nodes) - if total_size_of_input_nodes + partition.used_mem_bytes > available_mem_bytes: + if ( + total_size_of_input_nodes + partition.used_mem_bytes + > available_mem_bytes + ): partition = reset_partition_in_sparse_nn(partition) total_size_of_input_nodes = get_extra_size_of(node, partition.nodes) if total_size_of_input_nodes > available_mem_bytes: - raise RuntimeError(node.target + 'is too large to fit into a device') + raise RuntimeError( + node.target + "is too large to fit into a device" + ) partition.add_node(node) reset_partition_in_sparse_nn(partition, new_partition=False) # Set parents and children for partitions @@ -593,17 +653,25 @@ def is_embedding_node(node: Node) -> bool: total_size_of_non_embedding_partitions += partition.used_mem_bytes # Check if devices are enough for all partitions if len(embedding_partitions) > len(self.devices): - msg = 'Need ' + str(len(embedding_partitions)) + ' devices, but only ' \ - + str(len(self.devices)) + ' provided' + msg = ( + "Need " + + str(len(embedding_partitions)) + + " devices, but only " + + str(len(self.devices)) + + " provided" + ) raise RuntimeError(msg) occupied_devices = [] for i, partition in enumerate(embedding_partitions): # Check if all non-embedding partitions can fit into embedding partition devices - if total_size_of_non_embedding_partitions + partition.used_mem_bytes > available_mem_bytes: + if ( + total_size_of_non_embedding_partitions + partition.used_mem_bytes + > available_mem_bytes + ): raise RuntimeError( - 'partition_' + - str(partition.partition_id) + - '(embedding partition) and non embedding partitions can not fit into one device' + "partition_" + + str(partition.partition_id) + + "(embedding partition) and non embedding partitions can not fit into one device" ) else: # Add logical device to the partition @@ -619,68 +687,82 @@ def is_embedding_node(node: Node) -> bool: def cost_aware_partition( self, transfer_rate_bytes_per_sec: float, - node_to_latency_mapping: Dict[Node, NodeLatency] + node_to_latency_mapping: Dict[Node, NodeLatency], ) -> None: """This method is to partition the fx module based on the cost. - The cost is the total latency of running the whole fx module. - In partitioner_utils.py, the cost model is built. - The cost aware partition algorithm is: - #1. At every begining, each node is a partition. - Then we map all the partitions to the devices - and calculate the cost - #2. Then try to pre-combine any two of the partitions if the two - partitions can be combined. - (the bfs level is less than 2 or two partitions are connected and - can find partition to device mapping) - See if any partition pair could reduce the current cost. - Choose the pair that shows the minimum cost and then combine them - #3. Repeat #2 until the cost cannot be reduced. + The cost is the total latency of running the whole fx module. + In partitioner_utils.py, the cost model is built. + The cost aware partition algorithm is: + #1. At every begining, each node is a partition. + Then we map all the partitions to the devices + and calculate the cost + #2. Then try to pre-combine any two of the partitions if the two + partitions can be combined. + (the bfs level is less than 2 or two partitions are connected and + can find partition to device mapping) + See if any partition pair could reduce the current cost. + Choose the pair that shows the minimum cost and then combine them + #3. Repeat #2 until the cost cannot be reduced. """ - def try_combining_partitions( - p0_index, - p1_index, - partitions - ) -> float: + + def try_combining_partitions(p0_index, p1_index, partitions) -> float: """Given two partitions and a list of partitions, combine these two partitions - and see what is the cost of the modified partition list + and see what is the cost of the modified partition list """ p0 = partitions[p0_index] p1 = partitions[p1_index] """If two partitions' bfs level are less than 2 or two partitions are connected to each other, then they can be combined """ - if (abs(p0.bfs_level - p1.bfs_level) <= 1) or (p0 in p1.parents) or p0 in (p1.children): + if ( + (abs(p0.bfs_level - p1.bfs_level) <= 1) + or (p0 in p1.parents) + or p0 in (p1.children) + ): combine_two_partitions(p0, p1, partitions) # Check if a circular dependency exists after combining if check_dependency(partitions[-1]): - return float('inf') + return float("inf") # Check if the modified partition list can be mapped to devices after combination reset_partition_device(partitions) - found_deivce = get_device_to_partitions_mapping(partitions, self.devices) + found_deivce = get_device_to_partitions_mapping( + partitions, self.devices + ) if not found_deivce: - return float('inf') + return float("inf") # Calculate the new cost - partition_to_latency_mapping = get_partition_to_latency_mapping(partitions, node_to_latency_mapping) - cost = get_latency_of_partitioned_graph(partitions, partition_to_latency_mapping, transfer_rate_bytes_per_sec) + partition_to_latency_mapping = get_partition_to_latency_mapping( + partitions, node_to_latency_mapping + ) + cost = get_latency_of_partitioned_graph( + partitions, + partition_to_latency_mapping, + transfer_rate_bytes_per_sec, + ) return cost # If two partition can not be combined, the cost is inf - return float('inf') + return float("inf") def search_combination( - transfer_rate_bytes_per_sec, - node_to_latency_mapping + transfer_rate_bytes_per_sec, node_to_latency_mapping ) -> bool: """Given transfer rate between partitions and each node's latency, - find two partitions to combine so the cost of the partitions can - be reduced. - The algorithm is : - 1. Go through all the partition pairs and see - if any pair of partitions can be combined. - 2. Calculate the cost after the combination. - 3. Select the minimum cost and combine its cooresponding partition pair. + find two partitions to combine so the cost of the partitions can + be reduced. + The algorithm is : + 1. Go through all the partition pairs and see + if any pair of partitions can be combined. + 2. Calculate the cost after the combination. + 3. Select the minimum cost and combine its cooresponding partition pair. """ - partition_to_latency_mapping = get_partition_to_latency_mapping(self.partitions, node_to_latency_mapping) - cost = get_latency_of_partitioned_graph(self.partitions, partition_to_latency_mapping, transfer_rate_bytes_per_sec) + partition_to_latency_mapping = get_partition_to_latency_mapping( + self.partitions, node_to_latency_mapping + ) + cost = get_latency_of_partitioned_graph( + self.partitions, + partition_to_latency_mapping, + transfer_rate_bytes_per_sec, + ) if len(self.partitions) == 1: return False partition_pair: List[int] = [] @@ -688,11 +770,7 @@ def search_combination( for j in range(i + 1, len(self.partitions)): # Try to combine the partition pair # and see the new cost after combination - new_cost = try_combining_partitions( - i, - j, - self.partitions[:] - ) + new_cost = try_combining_partitions(i, j, self.partitions[:]) if new_cost <= cost: partition_pair = [i, j] cost = new_cost @@ -708,7 +786,7 @@ def search_combination( return len(partition_pair) != 0 for node in self.graph_module.graph.nodes: - if node.op not in {'placeholder', 'get_attr', 'output'}: + if node.op not in {"placeholder", "get_attr", "output"}: self.create_single_node_partition(node) # Set up parent partitions and children partitions for each partition set_parents_and_children(self.partitions) @@ -719,8 +797,7 @@ def search_combination( # Search for a pair partition to generate the minimum new cost, # then combine them find_combination = search_combination( - transfer_rate_bytes_per_sec, - node_to_latency_mapping + transfer_rate_bytes_per_sec, node_to_latency_mapping ) # Make sure all partitions are set up correctly reorganize_partitions(self.partitions) @@ -731,26 +808,27 @@ def search_combination( def kl_based_partition( self, transfer_rate_bytes_per_sec: float, - node_to_latency_mapping: Dict[Node, NodeLatency] + node_to_latency_mapping: Dict[Node, NodeLatency], ) -> None: """This function is a cost aware partition based - on Kernighan-Lin algorithm. - First, the graph is partitioned using size_based_partition. - Then, each node is swapped with any other node in a different - partition, and at the same time, the cost is estimated after - the swapping. - For example, we have nodes n0, n1, n2, n3 and n4. - Using size_based_partition, n0 and n1 are in Partition p0. - n2, n3 and n4 in Partition p1. The current cost is esimated. - We first tried using n0 to swap with n2 from the other partiton. - Then we see that swapping n0 and n2 shows a lower cost - than the current cost and it is the minimum among other pairs like - (n0, None)(This means moving n0 to Partition without swapping other nodes), - (n0, n3) and (n0, n4). We swap n0 and n2 and set the new cost - as the current cost. - Then We repeat this process for all the other nodes until all swapping pairs - are tried. + on Kernighan-Lin algorithm. + First, the graph is partitioned using size_based_partition. + Then, each node is swapped with any other node in a different + partition, and at the same time, the cost is estimated after + the swapping. + For example, we have nodes n0, n1, n2, n3 and n4. + Using size_based_partition, n0 and n1 are in Partition p0. + n2, n3 and n4 in Partition p1. The current cost is esimated. + We first tried using n0 to swap with n2 from the other partiton. + Then we see that swapping n0 and n2 shows a lower cost + than the current cost and it is the minimum among other pairs like + (n0, None)(This means moving n0 to Partition without swapping other nodes), + (n0, n3) and (n0, n4). We swap n0 and n2 and set the new cost + as the current cost. + Then We repeat this process for all the other nodes until all swapping pairs + are tried. """ + def swap_nodes(n0, n1, p0, p1): # Either n0 or n1 could be None # That means we simply move the node @@ -762,8 +840,10 @@ def swap_nodes(n0, n1, p0, p1): p0.add_node(n1) p1.remove_node(n1) - def try_swap_nodes(n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec): - cost = float('inf') + def try_swap_nodes( + n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec + ): + cost = float("inf") swap_nodes(n0, n1, p0, p1) # Reorganize partitions after swapping reorganize_partitions(self.partitions) @@ -771,18 +851,19 @@ def try_swap_nodes(n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_se if (not check_dependency(p0)) and (not check_dependency(p1)): reset_partition_device(self.partitions) partition_to_latency_mapping = get_partition_to_latency_mapping( - self.partitions, - node_to_latency_mapping + self.partitions, node_to_latency_mapping ) # Check if all partitions can be mapped to logical devices after swapping - found_device = get_device_to_partitions_mapping(self.partitions, self.devices) + found_device = get_device_to_partitions_mapping( + self.partitions, self.devices + ) if not found_device: - cost = float('inf') + cost = float("inf") else: cost = get_latency_of_partitioned_graph( self.partitions, partition_to_latency_mapping, - transfer_rate_bytes_per_sec + transfer_rate_bytes_per_sec, ) # Swap back and reset all partitions back to original swap_nodes(n1, n0, p0, p1) @@ -791,19 +872,23 @@ def try_swap_nodes(n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_se get_device_to_partitions_mapping(self.partitions, self.devices) return cost - def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_per_sec): + def swap_node_to_partition( + node, p0, p1, node_to_latency_mapping, transfer_rate_per_sec + ): """This function helps to swap one node from partition p0 - with all the nodes in another partition p1 + with all the nodes in another partition p1 """ p1_nodes = list(p1.nodes) + [None] - min_cost = float('inf') + min_cost = float("inf") node_pair: List[Node] = [] for n1 in p1_nodes: # Ignore the node if it is not a op node - if n1 is not None and n1.op in {'placeholder', 'get_attr'}: + if n1 is not None and n1.op in {"placeholder", "get_attr"}: continue # Try swapping node in p0 with n1 in p1 - cost = try_swap_nodes(node, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec) + cost = try_swap_nodes( + node, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec + ) if cost < min_cost: node_pair = [node, n1] min_cost = cost @@ -812,14 +897,11 @@ def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_ # First use size_base_partition self.size_based_partition() partition_to_latency_mapping = get_partition_to_latency_mapping( - self.partitions, - node_to_latency_mapping + self.partitions, node_to_latency_mapping ) # Calculate the cost of the partitions cost = get_latency_of_partitioned_graph( - self.partitions, - partition_to_latency_mapping, - transfer_rate_bytes_per_sec + self.partitions, partition_to_latency_mapping, transfer_rate_bytes_per_sec ) # Keep tracking the node pair that shows the better cost node_pair: List[Node] = [] @@ -828,7 +910,7 @@ def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_ # Collect all the op nodes from the graph op_nodes = [] for n in self.graph_module.graph.nodes: - if n.op not in {'placeholder', 'get_attr', 'output'}: + if n.op not in {"placeholder", "get_attr", "output"}: op_nodes.append(n) for node in op_nodes: # Find which partition the current node belongs @@ -844,7 +926,7 @@ def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_ p0, p1, node_to_latency_mapping, - transfer_rate_bytes_per_sec + transfer_rate_bytes_per_sec, ) # Update the cost # Track the swapped node pair and their partitions @@ -854,7 +936,9 @@ def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_ partition_pair = [p0, p1] # Do the swapping after trying all the nodes from a partition if len(node_pair) != 0: - swap_nodes(node_pair[0], node_pair[1], partition_pair[0], partition_pair[1]) + swap_nodes( + node_pair[0], node_pair[1], partition_pair[0], partition_pair[1] + ) reorganize_partitions(self.partitions) get_device_to_partitions_mapping(self.partitions, self.devices) reorganize_partitions(self.partitions) @@ -862,9 +946,11 @@ def swap_node_to_partition(node, p0, p1, node_to_latency_mapping, transfer_rate_ get_device_to_partitions_mapping(self.partitions, self.devices) return - def aot_based_partition(self, node_to_partition_mapping, partition_to_logical_device_mapping): + def aot_based_partition( + self, node_to_partition_mapping, partition_to_logical_device_mapping + ): """This function helps to rebuild the partitions given the nodes and its - corresponding partition id + corresponding partition id """ partition_id_to_partition_mapping: Dict[int, Partition] = {} self.node_to_partition = node_to_partition_mapping @@ -875,8 +961,12 @@ def aot_based_partition(self, node_to_partition_mapping, partition_to_logical_de partition = Partition(partition_id) self.partitions.append(partition) partition_id_to_partition_mapping[partition_id] = partition - partition.logical_device_ids = partition_to_logical_device_mapping[partition_id] + partition.logical_device_ids = partition_to_logical_device_mapping[ + partition_id + ] else: - partition = partition_id_to_partition_mapping[self.node_to_partition[node]] + partition = partition_id_to_partition_mapping[ + self.node_to_partition[node] + ] # Add the current node into the partition partition.add_node(node) diff --git a/torch/fx/experimental/normalize.py b/torch/fx/experimental/normalize.py index 143d4fdb76895..4b757fd0d9993 100644 --- a/torch/fx/experimental/normalize.py +++ b/torch/fx/experimental/normalize.py @@ -1,14 +1,20 @@ +import operator +from typing import Any, Callable, Dict, Tuple, Optional + import torch import torch.fx import torch.fx as fx -import operator -from typing import Any, Callable, Dict, Tuple, Optional +from torch.fx import Transformer, Proxy from torch.fx.node import Argument, Target, Node, map_aggregate -from torch.fx.operator_schemas import normalize_module, normalize_function, create_type_hint +from torch.fx.operator_schemas import ( + normalize_module, + normalize_function, + create_type_hint, +) -from torch.fx import Transformer, Proxy from .schema_type_annotation import AnnotateTypesWithSchema + class NormalizeArgs(Transformer): """ Normalize arguments to Python targets. This means that @@ -26,8 +32,10 @@ class NormalizeArgs(Transformer): traced = torch.fx.symbolic_trace(m) traced = NormalizeArgs(traced).transform() """ - def __init__(self, module : torch.nn.Module, - normalize_to_only_use_kwargs : bool = True): + + def __init__( + self, module: torch.nn.Module, normalize_to_only_use_kwargs: bool = True + ): super().__init__(module) self.node_map: Dict[Proxy, Node] = {} self.normalize_to_only_use_kwargs = normalize_to_only_use_kwargs @@ -37,14 +45,14 @@ def run_node(self, n: Node) -> Any: def get_type(arg): if isinstance(arg, fx.Node): - return n.meta['type'] if 'type' in n.meta else None + return n.meta["type"] if "type" in n.meta else None return type(arg) arg_types = map_aggregate(n.args, get_type) - assert(isinstance(arg_types, tuple)) + assert isinstance(arg_types, tuple) arg_types = tuple([create_type_hint(i) for i in arg_types]) kwarg_types = {k: get_type(v) for k, v in kwargs.items()} - if n.op == 'call_function': + if n.op == "call_function": out = self.call_function(n.target, args, kwargs, arg_types, kwarg_types) else: out = super().run_node(n) @@ -53,27 +61,48 @@ def get_type(arg): return out def call_function( - self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any], - arg_types: Optional[Tuple[Any, ...]] = None, kwarg_types : Optional[Dict[str, Any]] = None): + self, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Any], + arg_types: Optional[Tuple[Any, ...]] = None, + kwarg_types: Optional[Dict[str, Any]] = None, + ): assert callable(target) - new_args_and_kwargs = normalize_function(target, args, kwargs, arg_types, kwarg_types, # type: ignore[arg-type] - self.normalize_to_only_use_kwargs) + new_args_and_kwargs = normalize_function( + target, + args, # type: ignore[arg-type] + kwargs, + arg_types, # type: ignore[arg-type] + kwarg_types, # type: ignore[arg-type] + self.normalize_to_only_use_kwargs, + ) if new_args_and_kwargs: new_args, new_kwargs = new_args_and_kwargs - return self.tracer.create_proxy('call_function', target, new_args, new_kwargs) + return self.tracer.create_proxy( + "call_function", target, new_args, new_kwargs + ) else: return super().call_function(target, args, kwargs) - def call_module(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]): + def call_module( + self, target: Target, args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ): assert isinstance(target, str) - new_args_and_kwargs = normalize_module(self.module, target, args, kwargs, # type: ignore[arg-type] - self.normalize_to_only_use_kwargs) + new_args_and_kwargs = normalize_module( + self.module, + target, + args, # type: ignore[arg-type] + kwargs, # type: ignore[arg-type] + self.normalize_to_only_use_kwargs, + ) if new_args_and_kwargs: new_args, new_kwargs = new_args_and_kwargs return super().call_module(target, new_args, new_kwargs) else: return super().call_module(target, args, kwargs) + class NormalizeOperators(AnnotateTypesWithSchema): """ Normalize callsites that are different ways of "spelling" the same @@ -91,22 +120,27 @@ class NormalizeOperators(AnnotateTypesWithSchema): traced = NormalizeOperators(traced).transform() """ - binary_magic_method_remap : Dict[Callable[[Any, Any], Any], Callable[[Any, Any], Any]] = { - torch.add : operator.add, - torch.mul : operator.mul, - torch.sub : operator.sub, - torch.div : operator.truediv, - torch.floor_divide : operator.floordiv, - torch.remainder : operator.mod, - torch.eq : operator.eq, - torch.ne : operator.ne, - torch.lt : operator.lt, - torch.le : operator.le, - torch.gt : operator.gt, - torch.ge : operator.ge, + + binary_magic_method_remap: Dict[ + Callable[[Any, Any], Any], Callable[[Any, Any], Any] + ] = { + torch.add: operator.add, + torch.mul: operator.mul, + torch.sub: operator.sub, + torch.div: operator.truediv, + torch.floor_divide: operator.floordiv, + torch.remainder: operator.mod, + torch.eq: operator.eq, + torch.ne: operator.ne, + torch.lt: operator.lt, + torch.le: operator.le, + torch.gt: operator.gt, + torch.ge: operator.ge, } - def call_function(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]): + def call_function( + self, target: Target, args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ): # Normalize operators according to the magic methods implemented on tensors here: # https://github.com/pytorch/pytorch/blob/28c5d90b679c6b38bf4183ec99f16d933c2f1bcd/tools/autograd/templates/python_variable_methods.cpp#L1137 # noqa: B950 @@ -118,6 +152,9 @@ def call_function(self, target : Target, args : Tuple[Argument, ...], kwargs : D lhs, rhs = args return super().call_function( - target=self.binary_magic_method_remap[target], args=(lhs, rhs), kwargs={}) + target=self.binary_magic_method_remap[target], + args=(lhs, rhs), + kwargs={}, + ) return super().call_function(target, args, kwargs) diff --git a/torch/fx/experimental/partitioner_utils.py b/torch/fx/experimental/partitioner_utils.py index 7d54924bda463..cffba192b6b79 100644 --- a/torch/fx/experimental/partitioner_utils.py +++ b/torch/fx/experimental/partitioner_utils.py @@ -1,15 +1,19 @@ +from enum import Enum from typing import NamedTuple, Dict, List, Set + from torch.fx.node import Node, map_arg -from enum import Enum + + class Partition: """Partition class contains all the information about an individual partition. It also provides necessary methods for manipulation the partition. """ + def __init__(self, partition_id: int) -> None: self.nodes: Set[Node] = set() self.partition_id = partition_id - self.parents: Set['Partition'] = set() - self.children: Set['Partition'] = set() + self.parents: Set["Partition"] = set() + self.children: Set["Partition"] = set() self.bfs_level: int = -1 self.used_mem_bytes: int = 0 self.logical_device_ids: List[int] = [] @@ -28,7 +32,7 @@ def add_node(self, node): map_arg(node.kwargs, lambda n: input_nodes.setdefault(n)) # Add current node's input nodes if they are placeholder or constants for n in input_nodes: - if n.op in {'placeholder', 'get_attr'}: + if n.op in {"placeholder", "get_attr"}: self.nodes.add(n) self.nodes.add(node) self.recalculate_mem_size() @@ -45,21 +49,26 @@ def remove_node(self, node): # and this input node is not used by some other nodes in this partition, # the remove this input node for input_node in input_nodes: - if all([n not in self.nodes for n in input_node.users]) and input_node.op in {'placeholder', 'get_attr'}: + if all( + [n not in self.nodes for n in input_node.users] + ) and input_node.op in {"placeholder", "get_attr"}: self.nodes.remove(input_node) self.recalculate_mem_size() + class Device(NamedTuple): name: str available_mem_bytes: int logical_id: int + class NodeLatency(NamedTuple): # Latency due to the memory bandwidth mem_latency_sec: float # Latency due to the computation computer_latency_sec: float + class PartitionLatency(NamedTuple): # Sum of all nodes' memory latency on the critical path mem_latency_sec: float @@ -68,6 +77,7 @@ class PartitionLatency(NamedTuple): # Latency of the critical path overall_latency_sec: float + class PartitionMode(Enum): size_based = 0 sparse_nn = 1 @@ -75,18 +85,20 @@ class PartitionMode(Enum): kl_based = 3 aot_based = 4 + class PartitionerConfig(NamedTuple): devices: List[Device] mode: PartitionMode = PartitionMode.size_based - transfer_rate_bytes_per_sec: float = 0. + transfer_rate_bytes_per_sec: float = 0.0 node_to_latency_mapping: Dict[Node, NodeLatency] = {} node_to_partition_mapping: Dict[Node, int] = {} partition_to_logical_device_mapping: Dict[int, List[int]] = {} + def get_extra_size_of(node: Node, nodes: Set[Node]) -> int: """Given a node and a set of nodes, - this function return the extra size that needed - if this node is included in this set. + this function return the extra size that needed + if this node is included in this set. """ # Find all its input nodes input_nodes: Dict[Node, None] = {} @@ -97,22 +109,22 @@ def get_extra_size_of(node: Node, nodes: Set[Node]) -> int: for n in input_nodes: # Make sure this node hasn't been in this set yet if n not in nodes: - size_bytes = getattr(n, 'size_bytes', None) + size_bytes = getattr(n, "size_bytes", None) if size_bytes: total_size_of_input_nodes += size_bytes.output_size else: - raise RuntimeError('node has no size_bytes attr') + raise RuntimeError("node has no size_bytes attr") # Don't forget the op node itself - size_bytes = getattr(node, 'size_bytes', None) + size_bytes = getattr(node, "size_bytes", None) if size_bytes: total_size_of_input_nodes += size_bytes.total_size else: - raise RuntimeError('node has no size_bytes attr') + raise RuntimeError("node has no size_bytes attr") return total_size_of_input_nodes + def get_latency_of_one_partition( - partition: Partition, - node_to_latency_mapping: Dict[Node, NodeLatency] + partition: Partition, node_to_latency_mapping: Dict[Node, NodeLatency] ) -> PartitionLatency: """Given a partiton and its nodes' latency, return a PartitionLatency for this partition""" @@ -121,7 +133,7 @@ def get_top_nodes(partition: Partition) -> List[Node]: top_nodes: List[Node] = [] for node in partition.nodes: # Skip placeholder and get_attr nodes - if node.op in {'placeholder', 'get_attr'}: + if node.op in {"placeholder", "get_attr"}: continue input_nodes: Dict[Node, None] = {} map_arg(node.args, lambda n: input_nodes.setdefault(n)) @@ -129,67 +141,110 @@ def get_top_nodes(partition: Partition) -> List[Node]: # If a node has no input nodes in this partition, # or its input nodes in this partition are placeholders and get_attrs # this node is on the top bfs level in this partition - if not any([n in partition.nodes and n.op not in {'placeholder', 'get_attr'} for n in input_nodes]): + if not any( + [ + n in partition.nodes and n.op not in {"placeholder", "get_attr"} + for n in input_nodes + ] + ): top_nodes.append(node) return top_nodes def dfs_helper(node: Node, partition_latency) -> PartitionLatency: """Given a top node of a partition, this function returns - the latency of the critical path in the partition + the latency of the critical path in the partition """ node_latency = node_to_latency_mapping[node] # Calculate the current overall latency of the partition - overall_latency_sec = partition_latency.overall_latency_sec + \ - max(node_latency.computer_latency_sec, node_latency.mem_latency_sec) + overall_latency_sec = partition_latency.overall_latency_sec + max( + node_latency.computer_latency_sec, node_latency.mem_latency_sec + ) # Update the mem latency of this path - mem_latency_sec = partition_latency.mem_latency_sec + node_latency.mem_latency_sec + mem_latency_sec = ( + partition_latency.mem_latency_sec + node_latency.mem_latency_sec + ) # Update the compute latency of this path - computer_latency_sec = partition_latency.computer_latency_sec + node_latency.computer_latency_sec + computer_latency_sec = ( + partition_latency.computer_latency_sec + node_latency.computer_latency_sec + ) # Get all users of this node that are in this partition users = set(node.users).intersection(partition.nodes) if users: - max_latency = PartitionLatency(mem_latency_sec=0., computer_latency_sec=0., overall_latency_sec=0.) + max_latency = PartitionLatency( + mem_latency_sec=0.0, computer_latency_sec=0.0, overall_latency_sec=0.0 + ) for n in users: # Get new partition latency recursively - new_partition_latency = dfs_helper(n, PartitionLatency(mem_latency_sec, computer_latency_sec, overall_latency_sec)) - if new_partition_latency.overall_latency_sec > max_latency.overall_latency_sec: + new_partition_latency = dfs_helper( + n, + PartitionLatency( + mem_latency_sec, computer_latency_sec, overall_latency_sec + ), + ) + if ( + new_partition_latency.overall_latency_sec + > max_latency.overall_latency_sec + ): max_latency = new_partition_latency return max_latency # If there is no user, the node is at bottom of the partition - return PartitionLatency(mem_latency_sec, computer_latency_sec, overall_latency_sec) + return PartitionLatency( + mem_latency_sec, computer_latency_sec, overall_latency_sec + ) + # Main part starts # Get all top level nodes of this partition top_nodes = get_top_nodes(partition) - critical_path_latency = PartitionLatency(mem_latency_sec=0., computer_latency_sec=0., overall_latency_sec=0.) + critical_path_latency = PartitionLatency( + mem_latency_sec=0.0, computer_latency_sec=0.0, overall_latency_sec=0.0 + ) # Go through all top nodes and find the largest latency (critical pass latency) for node in top_nodes: - partition_latency = dfs_helper(node, PartitionLatency(mem_latency_sec=0., computer_latency_sec=0., overall_latency_sec=0.)) - if partition_latency.overall_latency_sec > critical_path_latency.overall_latency_sec: + partition_latency = dfs_helper( + node, + PartitionLatency( + mem_latency_sec=0.0, computer_latency_sec=0.0, overall_latency_sec=0.0 + ), + ) + if ( + partition_latency.overall_latency_sec + > critical_path_latency.overall_latency_sec + ): critical_path_latency = partition_latency return critical_path_latency + def get_partition_to_latency_mapping( - partitions: List[Partition], - node_to_latency_mapping: Dict[Node, NodeLatency] + partitions: List[Partition], node_to_latency_mapping: Dict[Node, NodeLatency] ) -> Dict[Partition, PartitionLatency]: """Given all the partitions and node_to_latency_mapping dictionary, - return a mapping dictionary of each partition to its overall latency + return a mapping dictionary of each partition to its overall latency """ partition_to_latency_mapping: Dict[Partition, PartitionLatency] = {} # Go through each partition and get its latency for partition in partitions: - partition_latency = get_latency_of_one_partition(partition, node_to_latency_mapping) + partition_latency = get_latency_of_one_partition( + partition, node_to_latency_mapping + ) partition_to_latency_mapping[partition] = partition_latency return partition_to_latency_mapping -def get_comm_latency_between(parent_partition: Partition, child_partition: Partition, transfer_rate_bytes_per_sec: float): + +def get_comm_latency_between( + parent_partition: Partition, + child_partition: Partition, + transfer_rate_bytes_per_sec: float, +): """Given two partitions (parent and child), - calculate the communication latency between the two. + calculate the communication latency between the two. """ # If two partitions are on the same device, the comm latency is 0. - if parent_partition.logical_device_ids != [] and child_partition.logical_device_ids != [] \ - and parent_partition.logical_device_ids == child_partition.logical_device_ids: - return 0. + if ( + parent_partition.logical_device_ids != [] + and child_partition.logical_device_ids != [] + and parent_partition.logical_device_ids == child_partition.logical_device_ids + ): + return 0.0 # Keep tracking the communication size between parent and child comm_size = 0 # Keep tracking all the counted node @@ -210,26 +265,33 @@ def get_comm_latency_between(parent_partition: Partition, child_partition: Parti visited_nodes.add(n) return comm_size / transfer_rate_bytes_per_sec + def get_latency_of_partitioned_graph( partitions: List[Partition], partition_to_latency_mapping: Dict[Partition, PartitionLatency], - transfer_rate_bytes_per_sec: float + transfer_rate_bytes_per_sec: float, ): """Given all paritions in a graph, find the critical path among all partitions - and return its latency as the latency of the whole graph + and return its latency as the latency of the whole graph """ + def dfs_helper(partition: Partition, latency_so_far_sec: float) -> float: - """This function helps to recursively get the latency of a path of partitions - """ + """This function helps to recursively get the latency of a path of partitions""" # Update latency by adding current partition's latency - latency_so_far_sec += partition_to_latency_mapping[partition].overall_latency_sec + latency_so_far_sec += partition_to_latency_mapping[ + partition + ].overall_latency_sec children = partition.children if partition.children: - max_latency_sec = 0. + max_latency_sec = 0.0 for child in partition.children: # Calculate latency between - comm_latency_sec = get_comm_latency_between(partition, child, transfer_rate_bytes_per_sec) - new_latency_sec = dfs_helper(child, latency_so_far_sec + comm_latency_sec) + comm_latency_sec = get_comm_latency_between( + partition, child, transfer_rate_bytes_per_sec + ) + new_latency_sec = dfs_helper( + child, latency_so_far_sec + comm_latency_sec + ) if new_latency_sec > max_latency_sec: max_latency_sec = new_latency_sec return max_latency_sec @@ -237,7 +299,7 @@ def dfs_helper(partition: Partition, latency_so_far_sec: float) -> float: def get_top_partitions(partitions: List[Partition]) -> List[Partition]: """This function is to return all the partitions without parents - as the starting points of all the paths + as the starting points of all the paths """ top_partitions = [] for partition in partitions: @@ -247,9 +309,9 @@ def get_top_partitions(partitions: List[Partition]) -> List[Partition]: return top_partitions top_partitions = get_top_partitions(partitions) - critical_path_latency_sec = 0. + critical_path_latency_sec = 0.0 for partition in top_partitions: - latency_sec = dfs_helper(partition, 0.) + latency_sec = dfs_helper(partition, 0.0) if latency_sec > critical_path_latency_sec: critical_path_latency_sec = latency_sec return critical_path_latency_sec From cbbb7e145e770213146bab1fc36c6abe1865d857 Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 16:16:54 -0700 Subject: [PATCH 127/305] Pass RequestCallback to FaultyPG RPC agent Summary: This is needed to avoid FaultyPG from including and depending on RequestCallbackImpl, which is Python-only. The other RPC agents accept an explicit (upcast) pointer as an argument, and we can do the same for FaultyPG. Test Plan: Later in the stack. Reviewed By: mrshenli Differential Revision: D29132955 fbshipit-source-id: bb7554b84bcbf39750af637e6480515ac8b92b86 --- .../distributed/rpc/testing/faulty_process_group_agent.cpp | 4 ++-- .../csrc/distributed/rpc/testing/faulty_process_group_agent.h | 1 + torch/csrc/distributed/rpc/testing/init.cpp | 2 ++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp index f51de9d870971..bb980ee8cef08 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp @@ -1,4 +1,3 @@ -#include #include #include @@ -16,6 +15,7 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( c10::intrusive_ptr<::c10d::ProcessGroup> pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, + std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends) @@ -25,7 +25,7 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( std::move(pg), numSendRecvThreads, rpcTimeout, - std::make_unique()), + std::move(cb)), failNumSends_(failNumSends), messageTypesToFail_(parseMessagesToFailInput(messagesToFail)), messageTypesToDelay_(parseMessagesToDelay(messageTypesToDelay)) {} diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index b80bd78c3e1de..ee589072f2ddd 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -39,6 +39,7 @@ class FaultyProcessGroupAgent : public ProcessGroupAgent { c10::intrusive_ptr pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, + std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends = 0); diff --git a/torch/csrc/distributed/rpc/testing/init.cpp b/torch/csrc/distributed/rpc/testing/init.cpp index 28344bb5b1978..bccaa1f2b4232 100644 --- a/torch/csrc/distributed/rpc/testing/init.cpp +++ b/torch/csrc/distributed/rpc/testing/init.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -82,6 +83,7 @@ PyObject* faulty_agent_init(PyObject* _unused, PyObject* noargs) { process_group, num_send_recv_threads, rpc_timeout, + std::make_unique(), messages_to_fail, messages_to_delay, failNumSends), From 04ec122868d436e32d81345343725e271d94cdfc Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 16:16:54 -0700 Subject: [PATCH 128/305] Add some TORCH_API annotations to RPC Summary: They will be needed when RPC gets merged into libtorch Test Plan: CI later in the stack Reviewed By: mrshenli Differential Revision: D29132956 fbshipit-source-id: 8637640d56a1744a5dca5eb7d4b8ad0860c6b67c --- torch/csrc/distributed/rpc/process_group_agent.h | 6 +++--- torch/csrc/distributed/rpc/tensorpipe_agent.h | 12 ++++++------ torch/csrc/distributed/rpc/tensorpipe_utils.h | 6 +++--- .../rpc/testing/faulty_process_group_agent.h | 4 ++-- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/torch/csrc/distributed/rpc/process_group_agent.h b/torch/csrc/distributed/rpc/process_group_agent.h index 5706870988140..a6d1115f4074c 100644 --- a/torch/csrc/distributed/rpc/process_group_agent.h +++ b/torch/csrc/distributed/rpc/process_group_agent.h @@ -15,7 +15,7 @@ namespace rpc { constexpr auto kDefaultNumSendRecvThreads = 4; -struct ProcessGroupRpcBackendOptions : public RpcBackendOptions { +struct TORCH_API ProcessGroupRpcBackendOptions : public RpcBackendOptions { ProcessGroupRpcBackendOptions( int num_send_recv_threads, float rpc_timeout, @@ -34,7 +34,7 @@ struct ProcessGroupRpcBackendOptions : public RpcBackendOptions { // SendWork and RecvWork will be put into a task queue, and later picked up by // worker threads from the same ThreadPool. -struct SendWork { +struct TORCH_API SendWork { SendWork(const WorkerInfo& to, c10::intrusive_ptr message) : to_(to), message_(std::move(message)) {} @@ -44,7 +44,7 @@ struct SendWork { // SendWork wraps a Message and RecvWork wraps a Tensor. The difference here is // to allow us to run serialization/deserialization in the worker threads. -struct RecvWork { +struct TORCH_API RecvWork { RecvWork( const WorkerInfo& from, MessageType type, diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index df3328793fa11..9462c396b0f3b 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -62,7 +62,7 @@ constexpr int64_t kCudaBasicChannelPriority = 0; using steady_clock_time_point = std::chrono::time_point; -struct TransportRegistration { +struct TORCH_API TransportRegistration { std::shared_ptr transport; int64_t priority; std::string address; @@ -71,7 +71,7 @@ struct TransportRegistration { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) C10_DECLARE_REGISTRY(TensorPipeTransportRegistry, TransportRegistration); -struct ChannelRegistration { +struct TORCH_API ChannelRegistration { std::shared_ptr channel; int64_t priority; }; @@ -81,7 +81,7 @@ C10_DECLARE_REGISTRY(TensorPipeChannelRegistry, ChannelRegistration); constexpr auto kDefaultNumWorkerThreads = 16; -struct TensorPipeRpcBackendOptions : public RpcBackendOptions { +struct TORCH_API TensorPipeRpcBackendOptions : public RpcBackendOptions { TensorPipeRpcBackendOptions( int numWorkerThreads, optional> transports, @@ -146,13 +146,13 @@ struct TensorPipeRpcBackendOptions : public RpcBackendOptions { }; // Struct to track the network source metrics -struct NetworkSourceInfo { +struct TORCH_API NetworkSourceInfo { worker_id_t srcRank; std::vector srcMachineAddr; }; // Struct to track aggregated network metrics -struct AggregatedNetworkData { +struct TORCH_API AggregatedNetworkData { uint64_t numCalls{0}; uint64_t totalSentBytes{0}; uint64_t totalRecvBytes{0}; @@ -163,7 +163,7 @@ struct AggregatedNetworkData { // to transparently move tensors and payloads through the fastest available // transport or channel. It acts like a hybrid RPC transport, providing shared // memory (linux) and TCP (linux & mac) support. CUDA support is in progress. -class TensorPipeAgent : public RpcAgent { +class TORCH_API TensorPipeAgent : public RpcAgent { public: TensorPipeAgent( const c10::intrusive_ptr<::c10d::Store>& store, diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index 3f41b351c9898..ab328b9dca1a1 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -15,7 +15,7 @@ namespace torch { namespace distributed { namespace rpc { -const c10::Stream& getStreamForDevice( +TORCH_API const c10::Stream& getStreamForDevice( const std::vector& streams, const c10::Device& device); @@ -44,12 +44,12 @@ class TensorpipeDeviceTypeConverter { virtual ~TensorpipeDeviceTypeConverter() = default; }; -extern C10_API std::array< +extern TORCH_API std::array< std::atomic, static_cast(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> device_type_converter_registry; -class C10_API TensorpipeDeviceTypeConverterRegistrar { +class TORCH_API TensorpipeDeviceTypeConverterRegistrar { public: TensorpipeDeviceTypeConverterRegistrar( DeviceType, diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index ee589072f2ddd..d0bbb33fe3df2 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -7,7 +7,7 @@ namespace torch { namespace distributed { namespace rpc { -struct FaultyProcessGroupRpcBackendOptions +struct TORCH_API FaultyProcessGroupRpcBackendOptions : public ProcessGroupRpcBackendOptions { FaultyProcessGroupRpcBackendOptions( int num_send_recv_threads, @@ -31,7 +31,7 @@ struct FaultyProcessGroupRpcBackendOptions int numFailSends; }; -class FaultyProcessGroupAgent : public ProcessGroupAgent { +class TORCH_API FaultyProcessGroupAgent : public ProcessGroupAgent { public: FaultyProcessGroupAgent( const c10::intrusive_ptr<::c10d::Store>& store, From fc50f91929b67f1ccf34890416b46a17b399e21b Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Tue, 15 Jun 2021 16:16:54 -0700 Subject: [PATCH 129/305] Move RPC agents to libtorch (#59939) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59939 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28875276 fbshipit-source-id: f2f6970fd74de5f112636e78edaa4410c61d8c45 --- BUILD.bazel | 2 +- caffe2/CMakeLists.txt | 57 +-------- cmake/Dependencies.cmake | 7 ++ test/cpp/rpc/CMakeLists.txt | 4 +- third_party/tensorpipe | 2 +- third_party/tensorpipe.BUILD | 112 ++++++++++-------- tools/build_variables.bzl | 12 +- torch/CMakeLists.txt | 6 - torch/csrc/distributed/rpc/macros.h | 5 - .../csrc/distributed/rpc/tensorpipe_agent.cpp | 1 - torch/csrc/distributed/rpc/tensorpipe_agent.h | 1 - .../csrc/distributed/rpc/tensorpipe_cuda.cpp | 3 +- .../csrc/distributed/rpc/tensorpipe_utils.cpp | 1 - torch/csrc/distributed/rpc/tensorpipe_utils.h | 1 - 14 files changed, 86 insertions(+), 128 deletions(-) delete mode 100644 torch/csrc/distributed/rpc/macros.h diff --git a/BUILD.bazel b/BUILD.bazel index 8a116d6c5cc8e..d7289ce854602 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1726,7 +1726,7 @@ cc_library( ], [ ":aten", - "@tensorpipe", + "@tensorpipe//:tensorpipe_cpu", ], ), alwayslink = True, diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 174018456efd8..88cffd1a75d1c 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -344,53 +344,6 @@ endif() if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) - if(USE_DISTRIBUTED) - - # Define this target even if we're building without TensorPipe, to make life - # easier to other targets that depend on this. However, in that case, by not - # setting the USE_TENSORPIPE compile definition, this target will just end - # up being empty. Downstream targets should also add a #ifdef guard. - if(NOT WIN32) - add_library(process_group_agent - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.h" - ) - target_link_libraries(process_group_agent PRIVATE torch fmt::fmt-header-only) - add_dependencies(process_group_agent torch) - - if(USE_TENSORPIPE) - add_library(tensorpipe_agent - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/macros.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_cuda.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h" - ) - target_link_libraries(tensorpipe_agent PRIVATE torch tensorpipe fmt::fmt-header-only) - add_dependencies(tensorpipe_agent torch) - if(USE_CUDA) - target_compile_definitions(tensorpipe_agent PUBLIC USE_CUDA) - endif() - - if(USE_ROCM) - target_compile_definitions(tensorpipe_agent PRIVATE - USE_ROCM - __HIP_PLATFORM_HCC__ - ) - endif() - - target_compile_definitions(tensorpipe_agent PUBLIC USE_TENSORPIPE) - target_link_libraries(tensorpipe_agent PRIVATE tensorpipe) - add_dependencies(tensorpipe_agent tensorpipe) - endif() - endif() - endif() - set(CMAKE_POSITION_INDEPENDENT_CODE TRUE) # Generate files @@ -1236,7 +1189,7 @@ endif() if(USE_DISTRIBUTED) # Needed to support the inclusion of c10d/Foo.hpp headers. target_include_directories(torch_cpu PUBLIC ${TORCH_SRC_DIR}/lib) - target_compile_definitions(torch_cpu PRIVATE USE_DISTRIBUTED) + target_compile_definitions(torch_cpu PUBLIC USE_DISTRIBUTED) if(USE_GLOO AND USE_C10D_GLOO) target_compile_definitions(torch_cpu PUBLIC USE_C10D_GLOO) endif() @@ -1263,16 +1216,12 @@ if(USE_DISTRIBUTED) # #if defined(USE_DISTRIBUTED) && !defined(_WIN32) # need to be removed when RPC is supported if(NOT WIN32) - target_compile_definitions(torch_cpu PRIVATE - USE_RPC - ) + target_compile_definitions(torch_cpu PUBLIC USE_RPC) endif() # Pass USE_TENSORPIPE to torch_cpu as some parts of rpc/utils.cpp # can only be compiled with USE_TENSORPIPE is set. if(USE_TENSORPIPE) - target_compile_definitions(torch_cpu PRIVATE - USE_TENSORPIPE - ) + target_compile_definitions(torch_cpu PUBLIC USE_TENSORPIPE) endif() endif() diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index ab4cd32c40bce..70b6d71face6b 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1377,6 +1377,13 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE) add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/tensorpipe) list(APPEND Caffe2_DEPENDENCY_LIBS tensorpipe) + if(USE_CUDA) + list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS tensorpipe_cuda) + elseif(USE_ROCM) + message(WARNING "TensorPipe doesn't yet support ROCm") + # Not yet... + # list(APPEND Caffe2_HIP_DEPENDENCY_LIBS tensorpipe_hip) + endif() endif() endif() diff --git a/test/cpp/rpc/CMakeLists.txt b/test/cpp/rpc/CMakeLists.txt index 0eff382d2b1b8..c9fb1b0e7f17a 100644 --- a/test/cpp/rpc/CMakeLists.txt +++ b/test/cpp/rpc/CMakeLists.txt @@ -5,7 +5,7 @@ set(TORCH_RPC_TEST_SOURCES ${TORCH_RPC_TEST_DIR}/test_wire_serialization.cpp ) set(TORCH_RPC_TEST_DEPENDENCY_LIBS - torch gtest process_group_agent + torch gtest ) if(USE_GLOO) @@ -20,7 +20,7 @@ if(USE_TENSORPIPE) ${TORCH_RPC_TEST_DIR}/test_tensorpipe_serialization.cpp ) list(APPEND TORCH_RPC_TEST_DEPENDENCY_LIBS - tensorpipe_agent tensorpipe + tensorpipe ) endif() diff --git a/third_party/tensorpipe b/third_party/tensorpipe index 42a67277c1882..c0e7623adb05f 160000 --- a/third_party/tensorpipe +++ b/third_party/tensorpipe @@ -1 +1 @@ -Subproject commit 42a67277c1882c90cec0da6e57afb20247424994 +Subproject commit c0e7623adb05f36311c7cde6dac8fc4c290419d9 diff --git a/third_party/tensorpipe.BUILD b/third_party/tensorpipe.BUILD index d9e4bdb395741..ae210f473933d 100644 --- a/third_party/tensorpipe.BUILD +++ b/third_party/tensorpipe.BUILD @@ -71,63 +71,82 @@ cc_library( ) header_template_rule( - name = "tensorpipe_config_header", + name = "tensorpipe_cpu_config_header", src = "tensorpipe/config.h.in", out = "tensorpipe/config.h", substitutions = { - "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "", - "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "", - "#cmakedefine01 TENSORPIPE_SUPPORTS_CUDA": "", + "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "#define TENSORPIPE_HAS_SHM_TRANSPORT 1", + "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "#define TENSORPIPE_HAS_IBV_TRANSPORT 1", + "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "#define TENSORPIPE_HAS_CMA_CHANNEL 1", }, ) -TENSORPIPE_HEADERS = glob([ - "tensorpipe/*.h", - "tensorpipe/channel/*.h", - "tensorpipe/channel/*/*.h", - "tensorpipe/common/*.h", - "tensorpipe/core/*.h", - "tensorpipe/transport/*.h", - "tensorpipe/transport/*/*.h", - "tensorpipe/util/*/*.h", -]) +header_template_rule( + name = "tensorpipe_cuda_config_header", + src = "tensorpipe/config_cuda.h.in", + out = "tensorpipe/config_cuda.h", + substitutions = { + "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "#define TENSORPIPE_HAS_CUDA_IPC_CHANNEL 1", + "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "#define TENSORPIPE_HAS_CUDA_GDR_CHANNEL 1", + }, +) -TENSORPIPE_BASE_SRCS = glob([ - "tensorpipe/*.cc", - "tensorpipe/channel/*.cc", - "tensorpipe/common/address.cc", - "tensorpipe/common/epoll_loop.cc", - "tensorpipe/common/error.cc", - "tensorpipe/common/fd.cc", - "tensorpipe/common/ibv.cc", - "tensorpipe/common/socket.cc", - "tensorpipe/common/system.cc", - "tensorpipe/core/*.cc", - "tensorpipe/transport/*.cc", - "tensorpipe/util/*/*.cc", -]) +# We explicitly list the CUDA headers & sources, and we consider everything else +# as CPU (using a catch-all glob). This is both because there's fewer CUDA files +# (thus making it easier to list them exhaustively) and because it will make it +# more likely to catch a misclassified file: if we forget to mark a file as CUDA +# we'll try to build it on CPU and that's likely to fail. -TENSORPIPE_SRCS = TENSORPIPE_BASE_SRCS + glob([ - "tensorpipe/channel/basic/*.cc", - "tensorpipe/channel/mpt/*.cc", - "tensorpipe/channel/xth/*.cc", - "tensorpipe/transport/uv/*.cc", -]) +TENSORPIPE_CUDA_HEADERS = [ + "tensorpipe/tensorpipe_cuda.h", + "tensorpipe/channel/cuda_basic/*.h", + "tensorpipe/channel/cuda_gdr/*.h", + "tensorpipe/channel/cuda_ipc/*.h", + "tensorpipe/channel/cuda_xth/*.h", + "tensorpipe/common/cuda.h", + "tensorpipe/common/cuda_buffer.h", + "tensorpipe/common/cuda_lib.h", + "tensorpipe/common/cuda_loop.h", + "tensorpipe/common/nvml_lib.h", +] -TENSORPIPE_SRCS_CUDA = TENSORPIPE_SRCS + glob([ - "tensorpipe/common/cuda_loop.cc", +TENSORPIPE_CUDA_SOURCES = [ "tensorpipe/channel/cuda_basic/*.cc", + "tensorpipe/channel/cuda_gdr/*.cc", "tensorpipe/channel/cuda_ipc/*.cc", "tensorpipe/channel/cuda_xth/*.cc", -]) + "tensorpipe/common/cuda_buffer.cc", + "tensorpipe/common/cuda_loop.cc", +] + +TENSORPIPE_CPU_HEADERS = glob( + [ + "tensorpipe/*.h", + "tensorpipe/channel/*.h", + "tensorpipe/channel/*/*.h", + "tensorpipe/common/*.h", + "tensorpipe/core/*.h", + "tensorpipe/transport/*.h", + "tensorpipe/transport/*/*.h", + ], + exclude=TENSORPIPE_CUDA_HEADERS) + +TENSORPIPE_CPU_SOURCES = glob( + [ + "tensorpipe/*.cc", + "tensorpipe/channel/*.cc", + "tensorpipe/channel/*/*.cc", + "tensorpipe/common/*.cc", + "tensorpipe/core/*.cc", + "tensorpipe/transport/*.cc", + "tensorpipe/transport/*/*.cc", + ], + exclude=TENSORPIPE_CUDA_SOURCES) cc_library( - name = "tensorpipe", - srcs = TENSORPIPE_SRCS + [":tensorpipe_config_header"], - hdrs = TENSORPIPE_HEADERS, + name = "tensorpipe_cpu", + srcs = TENSORPIPE_CPU_SOURCES, + hdrs = TENSORPIPE_CPU_HEADERS + [":tensorpipe_cpu_config_header"], includes = [ ".", ], @@ -143,8 +162,8 @@ cc_library( cc_library( name = "tensorpipe_cuda", - srcs = TENSORPIPE_SRCS_CUDA + [":tensorpipe_config_header"], - hdrs = TENSORPIPE_HEADERS, + srcs = TENSORPIPE_CUDA_SOURCES, + hdrs = TENSORPIPE_CUDA_HEADERS + [":tensorpipe_cuda_config_header"], includes = [ ".", ], @@ -153,8 +172,7 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ - ":libnop", - ":libuv", + ":tensorpipe_cpu", "@cuda", ], ) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index b62aa38db2190..8258e0e6429eb 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -352,12 +352,14 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.cpp", + "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/message.cpp", "torch/csrc/distributed/rpc/profiler/remote_profiler_manager.cpp", "torch/csrc/distributed/rpc/profiler/server_process_global_profiler.cpp", "torch/csrc/distributed/rpc/python_call.cpp", "torch/csrc/distributed/rpc/python_remote_call.cpp", "torch/csrc/distributed/rpc/python_resp.cpp", + "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/request_callback.cpp", "torch/csrc/distributed/rpc/request_callback_no_python.cpp", "torch/csrc/distributed/rpc/rpc_agent.cpp", @@ -367,6 +369,9 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/rpc/script_call.cpp", "torch/csrc/distributed/rpc/script_remote_call.cpp", "torch/csrc/distributed/rpc/script_resp.cpp", + "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", + "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", + "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/torchscript_functions.cpp", "torch/csrc/distributed/rpc/types.cpp", "torch/csrc/distributed/rpc/utils.cpp", @@ -522,6 +527,7 @@ libtorch_cuda_distributed_base_sources = [ # These files are only supported on Linux (and others) but not on Windows. libtorch_cuda_distributed_extra_sources = [ + "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", "torch/lib/c10d/NCCLUtils.cpp", "torch/lib/c10d/ProcessGroupNCCL.cpp", ] @@ -710,17 +716,11 @@ libtorch_python_distributed_core_sources = [ libtorch_python_distributed_sources = libtorch_python_distributed_core_sources + [ "torch/csrc/distributed/autograd/init.cpp", - "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/init.cpp", - "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/py_rref.cpp", "torch/csrc/distributed/rpc/python_functions.cpp", "torch/csrc/distributed/rpc/python_rpc_handler.cpp", "torch/csrc/distributed/rpc/request_callback_impl.cpp", - "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", - "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", - "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", - "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/testing/init.cpp", "torch/csrc/distributed/rpc/unpickled_python_call.cpp", "torch/csrc/distributed/rpc/unpickled_python_remote_call.cpp", diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index 197926f309838..ce0f16bf5abeb 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -261,11 +261,9 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") endif() if(USE_DISTRIBUTED) - list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_DISTRIBUTED) if(WIN32) append_filelist("libtorch_python_distributed_core_sources" TORCH_PYTHON_SRCS) else() - list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_RPC) append_filelist("libtorch_python_distributed_sources" TORCH_PYTHON_SRCS) endif() # Disable certain warnings for GCC-9.X @@ -274,10 +272,6 @@ if(USE_DISTRIBUTED) set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/rpc/testing/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/c10d/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") endif() - if(USE_TENSORPIPE) - list(APPEND TORCH_PYTHON_LINK_LIBRARIES tensorpipe) - list(APPEND TORCH_PYTHON_PUBLIC_COMPILE_DEFINITIONS USE_TENSORPIPE) - endif() # NCCL is a private dependency of libtorch, but libtorch_python includes # some private headers of libtorch, which in turn include NCCL. As a hacky # alternative to making NCCL a public dependency of libtorch, we make it diff --git a/torch/csrc/distributed/rpc/macros.h b/torch/csrc/distributed/rpc/macros.h deleted file mode 100644 index 2763dd0207bef..0000000000000 --- a/torch/csrc/distributed/rpc/macros.h +++ /dev/null @@ -1,5 +0,0 @@ -#pragma once - -#if defined(USE_CUDA) && !defined(__HIP_PLATFORM_HCC__) -#define USE_CUDA_NOT_ROCM -#endif diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp index 0f6645cdcd5d5..74c279425658b 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp @@ -10,7 +10,6 @@ #include #include -#include #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index 9462c396b0f3b..4450792a0f06d 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -9,7 +9,6 @@ #include #include #include -#include #include // Forward-declare the TensorPipe classes we need, to avoid including its diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp index 9489fcd222bbd..03ec63d8ddc88 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp @@ -1,8 +1,7 @@ -#include #include #include -#if defined(USE_TENSORPIPE) && defined(USE_CUDA_NOT_ROCM) +#if defined(USE_TENSORPIPE) && !defined(__HIP_PLATFORM_HCC__) #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp index 55b8554f66d28..32f3a132f8f50 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp @@ -1,4 +1,3 @@ -#include #include #ifdef USE_TENSORPIPE diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index ab328b9dca1a1..bf5d87cacc4b5 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -2,7 +2,6 @@ #ifdef USE_TENSORPIPE -#include #include namespace tensorpipe { From 5fd6ead0970a539f0bd1b24efbe6fc73bd70603f Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 15 Jun 2021 16:21:23 -0700 Subject: [PATCH 130/305] refine disabled test (#60040) Summary: This is to refine: https://github.com/pytorch/pytorch/pull/60029 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60040 Reviewed By: ezyang Differential Revision: D29147009 Pulled By: Krovatkin fbshipit-source-id: 37e01ac6e8d6f7e6b5c517f7804704f9136a56f5 --- test/distributed/test_c10d_gloo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py index 32f049f084c64..22de4440cd755 100644 --- a/test/distributed/test_c10d_gloo.py +++ b/test/distributed/test_c10d_gloo.py @@ -640,7 +640,7 @@ def _test_sparse_allreduce_basics(self, fn): self.assertEqual(tensors, outputs) self.assertEqual(result, outputs) - @unittest.skip("intermittent failures on Windows, in CI") + @skip_if_win32() def test_sparse_allreduce_basics(self): self._test_sparse_allreduce_basics(lambda t: t) From e341bab8aee3e2b6f3704a1f000b1c799aa905e2 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Tue, 15 Jun 2021 16:51:52 -0700 Subject: [PATCH 131/305] bugfix: ensure that at::{dispatch_key}:: API gets external linkage (#58569) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58569 This should allow external C++ files that aren't compiled into `libtorch.so`/`libtorch_cpu.so` (including all of fbcode) to use fast path functions like `at::cpu::add()`, which skip the dispatcher. So, after spending way too much time trying to figure out why I was getting linker errors when calling `at::meta::{op}` and `at::cpu::{op}` from C++ test files, I realized that we're not including the header files for C++ for the namespaced operator definitions. I.e. `RegisterCPU.cpp`, which provides definitions for the `at::cpu::{op}` fast path functions, wasn't including the `CPUFunctions.h` header. Why that breaks stuff: the `CPUFunctions.h` header file is what marks each function with the `TORCH_API` macro, so without including it, when we build `libtorch.so` and `libtorch_cpu.so`, the compiler will look at the definition in `RegisterCPU.cpp`, not see a `TORCH_API`, and decide that the function should get internal linkage. An alternative would be to directly mark the function definitions in `RegisterCPU.cpp` with `TORCH_API`, but this seemed cleaner. Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D28711300 Pulled By: bdhirsh fbshipit-source-id: 535f245c20e977ff566d6da0757b3cefa137040b --- BUILD.bazel | 4 +++- aten/src/ATen/templates/RegisterDispatchKey.cpp | 1 + tools/codegen/gen.py | 1 + tools/codegen/gen_backend_stubs.py | 1 + 4 files changed, 6 insertions(+), 1 deletion(-) diff --git a/BUILD.bazel b/BUILD.bazel index d7289ce854602..51ea3e914b892 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -132,11 +132,13 @@ genrule( "aten/src/ATen/RegisterSparseCPU.cpp", "aten/src/ATen/RegisterSparseCsrCPU.cpp", "aten/src/ATen/RegisterCompositeImplicitAutograd.cpp", - "aten/src/ATen/RegisterMeta.cpp", "aten/src/ATen/RegisterCompositeExplicitAutograd.cpp", + "aten/src/ATen/RegisterMeta.cpp", "aten/src/ATen/RegisterSchema.cpp", "aten/src/ATen/CPUFunctions.h", "aten/src/ATen/CUDAFunctions.h", + "aten/src/ATen/CompositeExplicitAutogradFunctions.h", + "aten/src/ATen/CompositeImplicitAutogradFunctions.h", "aten/src/ATen/Functions.h", "aten/src/ATen/Functions.cpp", "aten/src/ATen/RedispatchFunctions.h", diff --git a/aten/src/ATen/templates/RegisterDispatchKey.cpp b/aten/src/ATen/templates/RegisterDispatchKey.cpp index 9b0c3587daa62..6adb80586f37a 100644 --- a/aten/src/ATen/templates/RegisterDispatchKey.cpp +++ b/aten/src/ATen/templates/RegisterDispatchKey.cpp @@ -35,6 +35,7 @@ $extra_cuda_headers $legacy_th_headers $external_backend_headers +$namespaced_headers namespace at { diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index f3e4d726d9628..79c209e12d358 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -1016,6 +1016,7 @@ def make_file_manager(install_dir: str) -> FileManager: '#include ' if dispatch_key == DispatchKey.CUDA else '', 'external_backend_headers': '', + 'namespaced_headers': f'#include ' if dispatch_key in functions_keys else '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_definitions': list(concatMap( diff --git a/tools/codegen/gen_backend_stubs.py b/tools/codegen/gen_backend_stubs.py index 6d3724a61efed..a3b9ac254d1ce 100644 --- a/tools/codegen/gen_backend_stubs.py +++ b/tools/codegen/gen_backend_stubs.py @@ -175,6 +175,7 @@ def make_file_manager(install_dir: str) -> FileManager: 'legacy_th_headers': '', 'external_backend_headers': f'''#include "{output_dir}/{backend_key}NativeFunctions.h" #include ''', + 'namespaced_headers': '', 'DispatchKey': dispatch_key, 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_definitions': list(concatMap( From 27a3204982c3985f2d9639dbfed174da267bcd47 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Tue, 15 Jun 2021 16:51:52 -0700 Subject: [PATCH 132/305] generate C++ API for meta functions using at::meta:: (#58570) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58570 **What the PR does** Generate a fast-path `at::meta::{op}` API for calling meta functions without having to go through the dispatcher. This will be important for perf for external backends that want to use meta functions for shape checking (which seems likely to be what we end up doing for LazyTensorCore). **Details** In order to avoid naming collisions I had to make two small changes: - rename `MetaFunctions.h` template -> `NativeMetaFunctions.h` (this is the file that declares the impl() function for every structured operator). - rename the meta class: `at::meta::{op}::meta()` -> `at::meta::structured_{op}::meta()` I also deleted a few unnecessary includes, since any file that includes NativeFunctions.h will automatically include NativeMetaFunctions.h. **Why I made the change** This change isn't actually immediately used anywhere; I already started writing it because I thought it would be useful for structured composite ops, but that isn't actually true (see [comment](https://github.com/pytorch/pytorch/pull/58266#issuecomment-843213147)). The change feels useful and unambiguous though so I think it's safe to add. I added explicit tests for C++ meta function calls just to ensure that I wrote it correctly - which is actually how I hit the internal linkage issue in the PR below this in the stack. Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D28711299 Pulled By: bdhirsh fbshipit-source-id: d410d17358c2b406f0191398093f17308b3c6b9e --- BUILD.bazel | 1 + aten/src/ATen/TensorMeta.h | 4 +-- aten/src/ATen/native/Normalization.cpp | 1 - aten/src/ATen/native/UpSampleNearest1d.cpp | 1 - aten/src/ATen/templates/NativeFunctions.h | 2 +- ...{MetaFunctions.h => NativeMetaFunctions.h} | 0 .../ATen/templates/RegisterDispatchKey.cpp | 1 - test/cpp/api/CMakeLists.txt | 1 + test/cpp/api/meta_tensor.cpp | 35 +++++++++++++++++++ tools/codegen/dest/native_functions.py | 2 +- tools/codegen/dest/register_dispatch_key.py | 4 +-- tools/codegen/gen.py | 5 +-- 12 files changed, 46 insertions(+), 11 deletions(-) rename aten/src/ATen/templates/{MetaFunctions.h => NativeMetaFunctions.h} (100%) create mode 100644 test/cpp/api/meta_tensor.cpp diff --git a/BUILD.bazel b/BUILD.bazel index 51ea3e914b892..1f840807ec46d 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -147,6 +147,7 @@ genrule( "aten/src/ATen/Operators.cpp", "aten/src/ATen/NativeFunctions.h", "aten/src/ATen/MetaFunctions.h", + "aten/src/ATen/NativeMetaFunctions.h", "aten/src/ATen/core/TensorBody.h", "aten/src/ATen/core/TensorMethods.cpp", "aten/src/ATen/core/ATenOpList.cpp", diff --git a/aten/src/ATen/TensorMeta.h b/aten/src/ATen/TensorMeta.h index f719a7df64466..ac295ec9bde79 100644 --- a/aten/src/ATen/TensorMeta.h +++ b/aten/src/ATen/TensorMeta.h @@ -23,8 +23,8 @@ namespace impl { // set_output(sizes, options); // } // -#define TORCH_META_FUNC(name) void name::meta -#define TORCH_META_FUNC2(name, overload) void name##_##overload::meta +#define TORCH_META_FUNC(name) void structured_##name::meta +#define TORCH_META_FUNC2(name, overload) void structured_##name##_##overload::meta // Use this to define the prototype for an implementation. This takes only // one argument, which is the name of the dispatch key entry you're diff --git a/aten/src/ATen/native/Normalization.cpp b/aten/src/ATen/native/Normalization.cpp index d48ce03ee45cc..a1d751f43fcbd 100644 --- a/aten/src/ATen/native/Normalization.cpp +++ b/aten/src/ATen/native/Normalization.cpp @@ -1,5 +1,4 @@ #include -#include #include #include #include diff --git a/aten/src/ATen/native/UpSampleNearest1d.cpp b/aten/src/ATen/native/UpSampleNearest1d.cpp index b20f5dee9e4e6..46c0d68ca3dd4 100644 --- a/aten/src/ATen/native/UpSampleNearest1d.cpp +++ b/aten/src/ATen/native/UpSampleNearest1d.cpp @@ -1,7 +1,6 @@ #include #include #include -#include namespace at { namespace meta { diff --git a/aten/src/ATen/templates/NativeFunctions.h b/aten/src/ATen/templates/NativeFunctions.h index 2e35fde1b95e9..0acc828312cec 100644 --- a/aten/src/ATen/templates/NativeFunctions.h +++ b/aten/src/ATen/templates/NativeFunctions.h @@ -3,7 +3,7 @@ // ${generated_comment} #include -#include +#include #include #include #include diff --git a/aten/src/ATen/templates/MetaFunctions.h b/aten/src/ATen/templates/NativeMetaFunctions.h similarity index 100% rename from aten/src/ATen/templates/MetaFunctions.h rename to aten/src/ATen/templates/NativeMetaFunctions.h diff --git a/aten/src/ATen/templates/RegisterDispatchKey.cpp b/aten/src/ATen/templates/RegisterDispatchKey.cpp index 6adb80586f37a..5eac5c51965f9 100644 --- a/aten/src/ATen/templates/RegisterDispatchKey.cpp +++ b/aten/src/ATen/templates/RegisterDispatchKey.cpp @@ -11,7 +11,6 @@ #include #include #include -#include #include #include #include diff --git a/test/cpp/api/CMakeLists.txt b/test/cpp/api/CMakeLists.txt index ebc3dd5192392..9bd9d6780fe7d 100644 --- a/test/cpp/api/CMakeLists.txt +++ b/test/cpp/api/CMakeLists.txt @@ -12,6 +12,7 @@ set(TORCH_API_TEST_SOURCES ${TORCH_API_TEST_DIR}/init.cpp ${TORCH_API_TEST_DIR}/jit.cpp ${TORCH_API_TEST_DIR}/memory.cpp + ${TORCH_API_TEST_DIR}/meta_tensor.cpp ${TORCH_API_TEST_DIR}/misc.cpp ${TORCH_API_TEST_DIR}/module.cpp ${TORCH_API_TEST_DIR}/moduledict.cpp diff --git a/test/cpp/api/meta_tensor.cpp b/test/cpp/api/meta_tensor.cpp new file mode 100644 index 0000000000000..286eaf2c5d5fb --- /dev/null +++ b/test/cpp/api/meta_tensor.cpp @@ -0,0 +1,35 @@ +#include + +#include +#include + +#include + +TEST(MetaTensorTest, MetaDeviceApi) { + auto a = at::ones({4}, at::kFloat); + auto b = at::ones({3, 4}, at::kFloat); + // at::add() will return a meta tensor if its inputs are also meta tensors. + auto out_meta = at::add(a.to(c10::kMeta), b.to(c10::kMeta)); + + ASSERT_EQ(a.device(), c10::kCPU); + ASSERT_EQ(b.device(), c10::kCPU); + ASSERT_EQ(out_meta.device(), c10::kMeta); + c10::IntArrayRef sizes_actual = out_meta.sizes(); + std::vector sizes_expected = std::vector{3, 4}; + ASSERT_EQ(sizes_actual, sizes_expected); +} + +TEST(MetaTensorTest, MetaNamespaceApi) { + auto a = at::ones({4}, at::kFloat); + auto b = at::ones({3, 4}, at::kFloat); + // The at::meta:: namespace take in tensors from any backend + // and return a meta tensor. + auto out_meta = at::meta::add(a, b); + + ASSERT_EQ(a.device(), c10::kCPU); + ASSERT_EQ(b.device(), c10::kCPU); + ASSERT_EQ(out_meta.device(), c10::kMeta); + c10::IntArrayRef sizes_actual = out_meta.sizes(); + std::vector sizes_expected = std::vector{3, 4}; + ASSERT_EQ(sizes_actual, sizes_expected); +} diff --git a/tools/codegen/dest/native_functions.py b/tools/codegen/dest/native_functions.py index c643ea0eebe01..fbb894f815243 100644 --- a/tools/codegen/dest/native_functions.py +++ b/tools/codegen/dest/native_functions.py @@ -28,7 +28,7 @@ def gen_structured(g: NativeFunctionsGroup, backend_index: BackendIndex) -> List return [] prefix = '' if backend_index.external else 'TORCH_API ' return [f"""\ -struct {prefix}structured_{metadata.kernel} : public at::meta::{meta_name} {{ +struct {prefix}structured_{metadata.kernel} : public at::meta::structured_{meta_name} {{ void impl({', '.join(a.decl() for a in out_args)}); }}; """] diff --git a/tools/codegen/dest/register_dispatch_key.py b/tools/codegen/dest/register_dispatch_key.py index 69bd17df79ece..adaba8c4bb67c 100644 --- a/tools/codegen/dest/register_dispatch_key.py +++ b/tools/codegen/dest/register_dispatch_key.py @@ -502,11 +502,11 @@ def generate_defn(cpp_sig: CppSignature) -> str: # operator; feeding it the output argument(s) if it is known if self.backend_index.dispatch_key is DispatchKey.Meta: class_name = f"structured_{meta.name(self.g)}_meta_{k.name}" - parent_class = f"at::meta::{meta.name(self.g)}" + parent_class = f"at::meta::structured_{meta.name(self.g)}" elif self.backend_index.dispatch_key is DispatchKey.CompositeExplicitAutograd: # TODO: dedup this branch class_name = f"structured_{meta.name(self.g)}_default_backend_{k.name}" - parent_class = f"at::meta::{meta.name(self.g)}" + parent_class = f"at::meta::structured_{meta.name(self.g)}" else: metadata = self.backend_index.get_kernel(self.g) assert metadata is not None diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index 79c209e12d358..193a4d35d74bb 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -441,7 +441,7 @@ def compute_meta_function_declaration(g: NativeFunctionsGroup) -> Optional[str]: if parent_class is None: parent_class = "at::impl::MetaBase" return f"""\ -struct TORCH_API {name} : public {parent_class} {{ +struct TORCH_API structured_{name} : public {parent_class} {{ void meta({args_str}); }}; """ @@ -998,6 +998,7 @@ def make_file_manager(install_dir: str) -> FileManager: DispatchKey.CUDA, DispatchKey.CompositeImplicitAutograd, DispatchKey.CompositeExplicitAutograd, + DispatchKey.Meta, } if options.backend_whitelist: dispatch_keys = [k for k in dispatch_keys if is_generic_dispatch_key(k) or str(k) in options.backend_whitelist] @@ -1072,7 +1073,7 @@ def make_file_manager(install_dir: str) -> FileManager: list(mapMaybe(ComputeBackendSelect(Target.REGISTRATION, selector), native_functions)), }) - cpu_fm.write('MetaFunctions.h', lambda: { + cpu_fm.write('NativeMetaFunctions.h', lambda: { 'declarations': list(mapMaybe(compute_meta_function_declaration, structured_native_functions)), }) From 1207745e98d740c6c275fbf80193cb537f6edf6f Mon Sep 17 00:00:00 2001 From: jjsjann123 Date: Tue, 15 Jun 2021 16:52:39 -0700 Subject: [PATCH 133/305] fixing illegal memory access on NHWC BN kernel (#59981) Summary: adding an early exit in the kernel to avoid reading out of bound. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59981 Reviewed By: ezyang Differential Revision: D29147349 Pulled By: ngimel fbshipit-source-id: b36a6a9e2526c609ff98fb5a44468f3257e0af67 --- aten/src/ATen/native/cuda/Normalization.cuh | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/aten/src/ATen/native/cuda/Normalization.cuh b/aten/src/ATen/native/cuda/Normalization.cuh index edd7a809d4a05..24faf9fa7913a 100644 --- a/aten/src/ATen/native/cuda/Normalization.cuh +++ b/aten/src/ATen/native/cuda/Normalization.cuh @@ -1067,6 +1067,10 @@ __global__ void batch_norm_transform_input_channels_last_kernel( int m_offset = blockIdx.y * blockDim.y + threadIdx.y; int c_offset = blockIdx.x * blockDim.x + threadIdx.x; + if (c_offset >= stride || m_offset >= reduction_size) { + return; + } + auto m_c = mean[c_offset]; auto inv_std_c = static_cast(inv_std[c_offset]); auto w_c = weight == nullptr ? accscalar_t(1.0) : static_cast(weight[c_offset]); @@ -1155,6 +1159,10 @@ __global__ void batch_norm_backward_reduce_channels_last_kernel( int m_offset = blockIdx.y * blockDim.y + threadIdx.y; int c_offset = blockIdx.x * blockDim.x + threadIdx.x; + if (c_offset >= stride || m_offset >= reduction_size) { + return; + } + int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS); int address_base = m_offset * stride + c_offset; int address_increment = inner_loop_stride * stride; @@ -1296,6 +1304,10 @@ __device__ __forceinline__ void batch_norm_backward_elemt_channels_last_kernel_i int m_offset = blockIdx.y * blockDim.y + threadIdx.y; int c_offset = blockIdx.x * blockDim.x + threadIdx.x; + if (c_offset >= stride || m_offset >= reduction_size) { + return; + } + auto m_c = mean[c_offset]; auto m_dy_c = sum_dy[c_offset] * norm_fct; auto factor_1_c = inv_std[c_offset]; From 469f0e42d6e2b3cd8c78b224b97d45be2dc7d0ee Mon Sep 17 00:00:00 2001 From: Bert Maher Date: Tue, 15 Jun 2021 18:18:33 -0700 Subject: [PATCH 134/305] [nnc] Handle more cases of excessive # of cat args (#60043) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60043 And add a unit test Test Plan: new unit test Reviewed By: navahgar Differential Revision: D29146547 fbshipit-source-id: 31532926032dbef70d163930f3d8be160f5eacc3 --- test/test_jit_fuser_te.py | 9 +++++++++ torch/csrc/jit/passes/tensorexpr_fuser.cpp | 11 ++++++----- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py index 66aae87b9e47d..2521ff95350e1 100644 --- a/test/test_jit_fuser_te.py +++ b/test/test_jit_fuser_te.py @@ -1871,6 +1871,15 @@ def eager(b: float): script = self.checkScript(eager, (1.0,)) + def test_cat_2k_args(self): + with inline_fusion_groups(): + def eager(x): + return torch.relu(torch.cat([x for _ in range(2000)])) + x = torch.randn(1) + trace = self.checkTrace(eager, (x,)) + fusion_groups = self.findFusionGroups(trace.graph_for(x)) + self.assertEqual(len(fusion_groups), 0) + works_list = [ '__radd__', diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.cpp b/torch/csrc/jit/passes/tensorexpr_fuser.cpp index bba6971375cd9..32983f996cdbf 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.cpp +++ b/torch/csrc/jit/passes/tensorexpr_fuser.cpp @@ -1121,11 +1121,10 @@ class TensorExprFuser { // available to pass arguments, and some implementation dependence. Select a // safe limit here. constexpr size_t subgraphArgLimit = 128; - if ((consumer->inputs().size() + consumer->outputs().size() + - producer->inputs().size() + producer->outputs().size()) > - subgraphArgLimit) { - return false; - } + auto const nInputs = consumer->inputs().size() + + consumer->outputs().size() + producer->inputs().size() + + producer->outputs().size(); + REQ(nInputs <= subgraphArgLimit); // Device checks if (consumer->kind() != aten::cat && producer->kind() != aten::cat) { @@ -1179,6 +1178,7 @@ class TensorExprFuser { for (auto const& input : listConstruct->inputs()) { REQ(isFusableOnDevice(input->node())); } + REQ((nInputs + listConstruct->inputs().size()) <= subgraphArgLimit); } else if (consumer->kind() == aten::cat) { REQ(consumer->input(0)->node()->kind() == prim::ListConstruct); REQ(consumer->input(0)->uses().size() == 1); @@ -1191,6 +1191,7 @@ class TensorExprFuser { auto listconstruct_device = tensorexpr::pickDeviceType(listConstruct->inputs()); REQ(listconstruct_device); + REQ((nInputs + listConstruct->inputs().size()) <= subgraphArgLimit); } else { REQ(isFusableOnDevice(producer)); } From 95257e8a62fdb0016859b691ef2b01241f905fb4 Mon Sep 17 00:00:00 2001 From: Hangchen Yu Date: Tue, 15 Jun 2021 19:35:34 -0700 Subject: [PATCH 135/305] [fx-acc] Fix wrong device assignment in find_single_partition (#60056) 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/60056 Previously we put the whole graph as a single partition onto a device with maximum memory if possible, but the code assumed that the first logical device always has the maximum memory. This diff fixes this issue and updates the unittest to reflect such a corner case. Test Plan: ``` buck test mode/opt //caffe2/test:test_fx_experimental -- --exact 'caffe2/test:test_fx_experimental - test_find_single_partition (test_fx_experimental.TestFXExperimental)' Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/6473924507772744 ✓ ListingSuccess: caffe2/test:test_fx_experimental - main (1.357) ✓ Pass: caffe2/test:test_fx_experimental - test_find_single_partition (test_fx_experimental.TestFXExperimental) (1.206) Summary Pass: 1 ListingSuccess: 1 ``` Reviewed By: gcatron Differential Revision: D29118715 fbshipit-source-id: cac6a1f0d2f47717446dcc80093bbcf362663859 --- test/test_fx_experimental.py | 4 ++-- torch/fx/experimental/accelerator_partitioner.py | 16 +++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/test/test_fx_experimental.py b/test/test_fx_experimental.py index 7329be167914d..49a2e3f9d080d 100644 --- a/test/test_fx_experimental.py +++ b/test/test_fx_experimental.py @@ -166,7 +166,7 @@ def forward(self, a, b): partitioner = Partitioner() devices = [ Device("dev_0", 125, 0), - Device("dev_1", 125, 1), + Device("dev_1", 150, 1), Device("dev_2", 125, 2), ] partitioner_config = PartitionerConfig(devices) @@ -174,7 +174,7 @@ def forward(self, a, b): module_with_submodules = ret.module_with_submodules dag = ret.dag self.assertEqual(traced(a, b), module_with_submodules(a, b)) - assert dag.nodes[0].logical_device_ids == [0] + assert dag.nodes[0].logical_device_ids == [1] def test_lack_of_devices(self): class TestModule(torch.nn.Module): diff --git a/torch/fx/experimental/accelerator_partitioner.py b/torch/fx/experimental/accelerator_partitioner.py index c16ba8c097957..c6904b042875e 100644 --- a/torch/fx/experimental/accelerator_partitioner.py +++ b/torch/fx/experimental/accelerator_partitioner.py @@ -222,9 +222,7 @@ def find_device_for(partition: Partition): logical_id = partition.logical_device_ids[0] device = logical_id_to_device[logical_id] device_to_partitions[device] = [partition] - device_to_left_mem_bytes[device] = ( - d.available_mem_bytes - partition.used_mem_bytes - ) + device_to_left_mem_bytes[device] -= partition.used_mem_bytes else: no_device_partitions.append(partition) # Find devices for all the partitions without a device @@ -312,7 +310,9 @@ def partition_graph( ) # Single partition if the whole module can be fit into one device elif total_size_of_graph <= device_with_max_mem.available_mem_bytes: - self.find_single_partition(total_size_of_graph) + self.find_single_partition( + total_size_of_graph, logical_device_id=device_with_max_mem.logical_id + ) elif total_size_of_graph > sum([d.available_mem_bytes for d in self.devices]): raise RuntimeError("Devices have no enough memory for the module") else: @@ -348,7 +348,9 @@ def partition_graph( ret = PartitionResult(dag, module_with_submodules) return ret - def find_single_partition(self, total_size_of_graph) -> None: + def find_single_partition( + self, total_size_of_graph, logical_device_id: int = 0 + ) -> None: """Fit the whole fx module into one device""" partition_0 = self.create_partition() for node in self.graph_module.graph.nodes: @@ -356,7 +358,7 @@ def find_single_partition(self, total_size_of_graph) -> None: break partition_0.nodes.add(node) partition_0.used_mem_bytes = total_size_of_graph - partition_0.logical_device_ids = [0] + partition_0.logical_device_ids = [logical_device_id] # Get the node to partition mapping self.node_to_partition = get_node_to_partition_mapping(self.partitions) return @@ -413,7 +415,7 @@ def find_device_based_on_size(node) -> Device: partition_to_left_mem_bytes[ partition ] = device.available_mem_bytes - # Update available mem for the current partitio + # Update available mem for the current partition partition.logical_device_ids.append(device.logical_id) else: # The current partition is not the first partition From eda2ddb5b06dce13bafd2a745e4634802e4640ef Mon Sep 17 00:00:00 2001 From: Hao Lu Date: Tue, 15 Jun 2021 20:03:26 -0700 Subject: [PATCH 136/305] [ATen] Fix aten::to schema (#60001) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60001 Fix the aten::to schema to reflect that the output may alias input. Test Plan: Added new unit tests. Reviewed By: ezyang Differential Revision: D29121620 fbshipit-source-id: c29b6aa22d367ffedf06e47116bc46b3e188c39c --- aten/src/ATen/native/native_functions.yaml | 8 ++--- .../check_backward_compatibility.py | 1 + test/cpp/jit/test_alias_analysis.cpp | 35 +++++++++++++++++++ torch/csrc/jit/runtime/static/ops.cpp | 12 +++---- 4 files changed, 46 insertions(+), 10 deletions(-) diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index d5711da70dc27..df977933ecd53 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -5251,22 +5251,22 @@ # to(Device) must not exist because all constructors of Device also works for # TensorOptions. Otherwise, an ambiguity error is thrown. # See NOTE [ TensorOptions Constructors ]. -- func: to.dtype_layout(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor +- func: to.dtype_layout(Tensor(a) self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) variants: method device_check: NoCheck device_guard: False -- func: to.device(Tensor self, Device device, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor +- func: to.device(Tensor(a) self, Device device, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) variants: method device_check: NoCheck device_guard: False -- func: to.dtype(Tensor self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor +- func: to.dtype(Tensor(a) self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) variants: method device_check: NoCheck device_guard: False -- func: to.other(Tensor self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor +- func: to.other(Tensor(a) self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) variants: method device_check: NoCheck device_guard: False diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index caa79a9c41823..d15d446cfbd98 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -97,6 +97,7 @@ ("aten::segment_reduce_backward", datetime.date(2021, 6, 15)), ("aten::segment_reduce", datetime.date(2021, 8, 26)), ("aten::_segment_reduce_backward", datetime.date(2021, 8, 26)), + ("aten::to", datetime.date(2021, 6, 22)), ] def allow_listed(schema, allow_list): diff --git a/test/cpp/jit/test_alias_analysis.cpp b/test/cpp/jit/test_alias_analysis.cpp index 5477a0b5515b0..9ef090e1a75c7 100644 --- a/test/cpp/jit/test_alias_analysis.cpp +++ b/test/cpp/jit/test_alias_analysis.cpp @@ -625,6 +625,41 @@ TEST(ContainerAliasingTest, MayContainAlias) { EXPECT_FALSE(aliasDb.mayContainAlias(str_output, graph->outputs())); } +TEST(ContainerAliasingTest, MayContainAlias_cast) { + auto graph = std::make_shared(); + std::unordered_map vmap; + parseIR( + R"IR( + graph(%input.1 : Tensor): + %2 : NoneType = prim::Constant() + %3 : bool = prim::Constant[value=0]() + %4 : int = prim::Constant[value=6]() + %5 : int = prim::Constant[value=1]() + %a.1 : Tensor = aten::add(%input.1, %input.1, %5) + %b.1 : Tensor = aten::to(%a.1, %4, %3, %3, %2) + %c.1 : Tensor = aten::mul(%b.1, %b.1) + return (%c.1) + )IR", + &*graph, + vmap); + + auto a = vmap["a.1"]; + auto b = vmap["b.1"]; + auto c = vmap["c.1"]; + AliasDb aliasDb(graph); + + EXPECT_TRUE(graph->outputs().size() == 1); + for (auto out : graph->outputs()) { + EXPECT_TRUE(aliasDb.mayContainAlias(c, out)); + } + + EXPECT_TRUE(aliasDb.mayContainAlias(a, b)); + EXPECT_FALSE(aliasDb.mayContainAlias(b, graph->inputs())); + + EXPECT_TRUE(aliasDb.mayContainAlias({c}, graph->outputs())); + EXPECT_FALSE(aliasDb.mayContainAlias(b, graph->outputs())); +} + // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) TEST(ContainerAliasingTest, PrimitveValuesDontAliasContainers) { auto graph = std::make_shared(); diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index 3c9590fa9934c..97734aa441765 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -1234,9 +1234,9 @@ std::function getNativeOperation(Node* n) { }; } else if (n->kind() == c10::Symbol::fromQualString("aten::to")) { if (!n->matches(torch::schema( - "aten::to.other(Tensor self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor")) && + "aten::to.other(Tensor(a) self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a)")) && !n->matches(torch::schema( - "aten::to.dtype(Tensor self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor"))) { + "aten::to.dtype(Tensor(a) self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a)"))) { LogAndDumpSchema(n); return nullptr; } @@ -1246,13 +1246,13 @@ std::function getNativeOperation(Node* n) { const auto in3_i = p_node->Input(3).toBool(); const auto in4_o = p_node->Input(4).toOptional(); if (p_node->Input(1).isTensor()) { - // to.other(Tensor self, Tensor other, bool non_blocking=False, bool - // copy=False, MemoryFormat? memory_format=None) -> Tensor + // to.other(Tensor(a) self, Tensor other, bool non_blocking=False, bool + // copy=False, MemoryFormat? memory_format=None) -> Tensor(a) const auto in1_t = p_node->Input(1).toTensor(); p_node->Output(0) = at::native::to(in0_t, in1_t, in2_i, in3_i, in4_o); } else { - // to.dtype(Tensor self, ScalarType dtype, bool non_blocking=False, bool - // copy=False, MemoryFormat? memory_format=None) -> Tensor + // to.dtype(Tensor(a) self, ScalarType dtype, bool non_blocking=False, + // bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) const auto in1_i = p_node->Input(1).toScalarType(); p_node->Output(0) = at::native::to(in0_t, in1_i, in2_i, in3_i, in4_o); } From e50f264b512ace9e170d733f2ffd81f9e341c6e2 Mon Sep 17 00:00:00 2001 From: Stephen Macke Date: Tue, 15 Jun 2021 20:24:39 -0700 Subject: [PATCH 137/305] [caffe2] make MulGradient implementation in-place compatible (#60035) 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/60035 In Caffe2, the operator schema for the MulGradient op indicates that MulGradient may be performed in-place, overwriting one of its inputs as the output. The implementation is not safe to perform in-place however, due to an accidentally-introduced write-read dependency on the overwriten input in the in-place case. We fix it here. Test Plan: ``` buck test //caffe2/caffe2/python/operator_test:elementwise_ops_test ``` Note that the newly added test fails without this change, but passes with this change: ``` ✓ ListingSuccess: caffe2/caffe2/python/operator_test:elementwise_ops_test - main (24.992) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_exp (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_log1p (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_abs (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_bitwise_and (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_reciprocal (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_sqr (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_rsqrt (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_mul (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_sqrt (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_add (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_swish_gradient_inplace (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_sigmoid (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_bitwise_or (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_cbrt_grad (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_not (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_sub (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_div (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_eq (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_softsign (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_eq_bcast (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_powt (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ************************************************************************************************************************************************************************************* *********************************************************************************************************************************************************************** ************************************************************************************************************************************************************************************* ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_mul_gradient_inplace (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ************************************************************************************************************************************************************************************* ********************************************************************************************************************************************************************** ************************************************************************************************************************************************************************************* ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_hard_sigmoid (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_bitwise_xor (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_log (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_cube (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_swish (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_cbrt (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - test_div_legacy_grad (caffe2.caffe2.python.operator_test.elementwise_ops_test.TestElementwiseOps) (125.898) ✓ Pass: caffe2/caffe2/python/operator_test:elementwise_ops_test - main (125.898) Summary Pass: 30 ListingSuccess: 1 ``` Reviewed By: clrfb Differential Revision: D29034265 fbshipit-source-id: 98550e1d5976398e45d37ff2120591af1439c42a --- .../operators/elementwise_mul_gradient_op.cc | 23 ++++++++++++ caffe2/operators/elementwise_mul_op.cu | 9 +++++ .../operator_test/elementwise_ops_test.py | 37 +++++++++++++++++++ 3 files changed, 69 insertions(+) diff --git a/caffe2/operators/elementwise_mul_gradient_op.cc b/caffe2/operators/elementwise_mul_gradient_op.cc index dc2f599631de8..b7256159587a8 100644 --- a/caffe2/operators/elementwise_mul_gradient_op.cc +++ b/caffe2/operators/elementwise_mul_gradient_op.cc @@ -74,6 +74,15 @@ void ComputeMulGradient( const float* B, float* dA, float* dB) { + if (dA != nullptr) { + CAFFE_ENFORCE_NE(dA, dB, "Outputs dA and dB should point to distinct blobs"); + } + if (dC == dA) { + // Ensure operation can be performed in-place. + // See below comment in `MulFunctor::Backward`. + std::swap(A, B); + std::swap(dA, dB); + } for (int i = 0; i < size; ++i) { dA[i] = dC[i] * B[i]; dB[i] = dC[i] * A[i]; @@ -94,8 +103,22 @@ bool MulFunctor::Backward( TGrad* dA, TGrad* dB, CPUContext* context) const { + if (dA != nullptr) { + CAFFE_ENFORCE_NE(dA, dB, "Outputs dA and dB should point to distinct blobs"); + } if (A_dims == B_dims) { const auto size = c10::multiply_integers(A_dims); + if (dC == dA) { + // A, B, and dC are inputs (dC is the output of the previous gradient op + // in the dag), and dA and dB are outputs. If the op is performed + // in-place, either dA or dB could alias dC. In the dC == dA case, we need + // to make sure we don't overwrite dC when we write to dA, so swap the + // inputs to avoid clobbering dC. Semantically this is equivalent with + // writing to dB first. The other case (dC == dB) is already safe because + // we are writing to dA first. + std::swap(A, B); + std::swap(dA, dB); + } math::Mul(size, dC, B, dA, context); math::Mul(size, dC, A, dB, context); return true; diff --git a/caffe2/operators/elementwise_mul_op.cu b/caffe2/operators/elementwise_mul_op.cu index 85576772b55a1..bdbf760cf95bd 100644 --- a/caffe2/operators/elementwise_mul_op.cu +++ b/caffe2/operators/elementwise_mul_op.cu @@ -217,7 +217,16 @@ bool MulFunctor::Backward( TGrad* dA, TGrad* dB, CUDAContext* context) const { + if (dA != nullptr) { + CAFFE_ENFORCE_NE(dA, dB, "Outputs dA and dB should point to distinct blobs"); + } if (A_dims == B_dims) { + if (dC == dA) { + // Ensure operation can be performed in-place. + // We want to avoid clobbering dC if it aliases dA. + std::swap(A, B); + std::swap(dA, dB); + } const int size = std::accumulate( A_dims.cbegin(), A_dims.cend(), 1, std::multiplies()); math::Mul(size, dC, B, dA, context); diff --git a/caffe2/python/operator_test/elementwise_ops_test.py b/caffe2/python/operator_test/elementwise_ops_test.py index 130ebade010ba..3c38e08490703 100644 --- a/caffe2/python/operator_test/elementwise_ops_test.py +++ b/caffe2/python/operator_test/elementwise_ops_test.py @@ -352,6 +352,43 @@ def swish_gradient(X, Y, dY): reference=swish_gradient, ) + @given(n=st.integers(1, 6), m=st.integers(4, 6), + seed=st.integers(0, 1000), **hu.gcs) + @settings(deadline=10000) + def test_mul_gradient_inplace(self, n, m, gc, dc, seed): + np.random.seed(seed) + + def mul_gradient(dC, A, B): + return [B * dC, A * dC] + + A = np.random.rand(n, m).astype(np.float32) + B = np.random.rand(n, m).astype(np.float32) + dC = np.random.rand(n, m).astype(np.float32) + op_dA_inplace = core.CreateOperator( + "MulGradient", + ["dC", "A", "B"], + ["dC", "dB"], + ) + op_dB_inplace = core.CreateOperator( + "MulGradient", + ["dC", "A", "B"], + ["dA", "dC"], + ) + + self.assertReferenceChecks( + device_option=gc, + op=op_dA_inplace, + inputs=[dC, A, B], + reference=mul_gradient, + ) + + self.assertReferenceChecks( + device_option=gc, + op=op_dB_inplace, + inputs=[dC, A, B], + reference=mul_gradient, + ) + @given(X=hu.tensor(dtype=np.float32), inplace=st.booleans(), engine=st.sampled_from(["", "CUDNN"]), **hu.gcs) @settings(deadline=10000) From 78011bc0ce9e11997b5bde51fb1c881b5eb7df47 Mon Sep 17 00:00:00 2001 From: clint Date: Tue, 15 Jun 2021 21:11:27 -0700 Subject: [PATCH 138/305] typofix (torch.zero to torch.zeros) in docstring (#59703) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59703 Reviewed By: ezyang Differential Revision: D29145998 Pulled By: H-Huang fbshipit-source-id: f2670502170aa100fb02408046b7f6850f9379cf --- torch/distributed/distributed_c10d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index 61f9cea2256a5..a43c57cbd5c46 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -1918,7 +1918,7 @@ def all_gather(tensor_list, tensor, group=None, async_op=False): Examples: >>> # All tensors below are of torch.int64 dtype. >>> # We have 2 process groups, 2 ranks. - >>> tensor_list = [torch.zero(2, dtype=torch.int64) for _ in range(2)] + >>> tensor_list = [torch.zeros(2, dtype=torch.int64) for _ in range(2)] >>> tensor_list [tensor([0, 0]), tensor([0, 0])] # Rank 0 and 1 >>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank @@ -1932,7 +1932,7 @@ def all_gather(tensor_list, tensor, group=None, async_op=False): >>> # All tensors below are of torch.cfloat dtype. >>> # We have 2 process groups, 2 ranks. - >>> tensor_list = [torch.zero(2, dtype=torch.cfloat) for _ in range(2)] + >>> tensor_list = [torch.zeros(2, dtype=torch.cfloat) for _ in range(2)] >>> tensor_list [tensor([0.+0.j, 0.+0.j]), tensor([0.+0.j, 0.+0.j])] # Rank 0 and 1 >>> tensor = torch.tensor([1+1j, 2+2j], dtype=torch.cfloat) + 2 * rank * (1+1j) @@ -1986,7 +1986,7 @@ def _all_gather_base(output_tensor, input_tensor, group=None, async_op=False): Examples: >>> # All tensors below are of torch.int64 dtype. >>> # We have 2 process groups, 2 ranks. - >>> output_tensor = torch.zero(2, dtype=torch.int64) + >>> output_tensor = torch.zeros(2, dtype=torch.int64) >>> output_tensor [tensor([0, 0])] # Rank 0 and 1 >>> tensor = torch.arange(1, dtype=torch.int64) + 1 + rank From a344b09db24dfab104f59a8a372ab6a51d5fcae6 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Tue, 15 Jun 2021 21:53:54 -0700 Subject: [PATCH 139/305] [quant][fx][graphmode] Remove Quantizer class (#59606) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59606 Test Plan: python test/test_quantization.py TestQuantizeFx Imported from OSS Reviewed By: vkuzo Differential Revision: D28951432 fbshipit-source-id: 3301f7200a4c7166673c27f9ac7ff559f1e6935d --- torch/quantization/fx/__init__.py | 3 ++- torch/quantization/fx/convert.py | 8 ++++---- torch/quantization/fx/prepare.py | 6 +++--- torch/quantization/fx/quantize.py | 31 ------------------------------- torch/quantization/quantize_fx.py | 11 +++++------ 5 files changed, 14 insertions(+), 45 deletions(-) delete mode 100644 torch/quantization/fx/quantize.py diff --git a/torch/quantization/fx/__init__.py b/torch/quantization/fx/__init__.py index 35e47dc07d39e..8ba1b6d79dd71 100644 --- a/torch/quantization/fx/__init__.py +++ b/torch/quantization/fx/__init__.py @@ -1,2 +1,3 @@ -from .quantize import Quantizer +from .prepare import prepare +from .convert import convert from .fuse import Fuser diff --git a/torch/quantization/fx/convert.py b/torch/quantization/fx/convert.py index b6a94e40b9e49..1906c37ee1e9d 100644 --- a/torch/quantization/fx/convert.py +++ b/torch/quantization/fx/convert.py @@ -138,10 +138,10 @@ def restore_state( patterns: Dict[Pattern, QuantizeHandler] = observed._patterns # type: ignore[assignment] return patterns, node_name_to_scope, prepare_custom_config_dict -def _convert(model: GraphModule, is_reference: bool = False, - convert_custom_config_dict: Dict[str, Any] = None, - is_standalone_module: bool = False, - _remove_qconfig_flag: bool = True) -> QuantizedGraphModule: +def convert(model: GraphModule, is_reference: bool = False, + convert_custom_config_dict: Dict[str, Any] = None, + is_standalone_module: bool = False, + _remove_qconfig_flag: bool = True) -> QuantizedGraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index 2433012a4f7ca..a2a7fe842c943 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -918,12 +918,12 @@ def save_state( prepare_custom_config_dict # type: ignore[assignment] observed._node_name_to_scope = node_name_to_scope # type: ignore[assignment] -def _prepare( +def prepare( model: GraphModule, qconfig_dict: Any, node_name_to_scope: Dict[str, Tuple[str, type]], - prepare_custom_config_dict: Optional[Dict[str, Any]], - is_standalone_module: bool) -> ObservedGraphModule: + prepare_custom_config_dict: Optional[Dict[str, Any]] = None, + is_standalone_module: bool = False) -> ObservedGraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. diff --git a/torch/quantization/fx/quantize.py b/torch/quantization/fx/quantize.py deleted file mode 100644 index dd228c50a4988..0000000000000 --- a/torch/quantization/fx/quantize.py +++ /dev/null @@ -1,31 +0,0 @@ -from torch.fx import ( - GraphModule, -) -from .prepare import _prepare -from .convert import _convert -from .graph_module import ( - ObservedGraphModule, - QuantizedGraphModule, -) - -from typing import Any, Dict, Tuple - -class Quantizer: - def prepare( - self, - model: GraphModule, - qconfig_dict: Any, - node_name_to_scope: Dict[str, Tuple[str, type]], - prepare_custom_config_dict: Dict[str, Any] = None, - is_standalone_module: bool = False) -> ObservedGraphModule: - return _prepare( - model, qconfig_dict, node_name_to_scope, prepare_custom_config_dict, - is_standalone_module) - - def convert(self, model: GraphModule, is_reference: bool = False, - convert_custom_config_dict: Dict[str, Any] = None, - is_standalone_module: bool = False, - _remove_qconfig: bool = True) -> QuantizedGraphModule: - quantized = _convert( - model, is_reference, convert_custom_config_dict, is_standalone_module, _remove_qconfig_flag=_remove_qconfig) - return quantized diff --git a/torch/quantization/quantize_fx.py b/torch/quantization/quantize_fx.py index 692e7375ad70c..82fd0b49c099f 100644 --- a/torch/quantization/quantize_fx.py +++ b/torch/quantization/quantize_fx.py @@ -3,7 +3,7 @@ from torch.fx._symbolic_trace import Tracer from torch.fx.node import Target, Node, Argument from .fx import Fuser # noqa: F401 -from .fx import Quantizer # noqa: F401 +from .fx import prepare, convert # noqa: F401 from .fx.utils import graph_pretty_str # noqa: F401 from .fx.utils import get_custom_module_class_keys # noqa: F401 from .fx.graph_module import ObservedGraphModule, QuantizedGraphModule @@ -183,8 +183,7 @@ def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, for attr_name in preserved_attributes: setattr(graph_module, attr_name, getattr(model, attr_name)) graph_module = _fuse_fx(graph_module, prepare_custom_config_dict) - quantizer = Quantizer() - prepared = quantizer.prepare( + prepared = prepare( graph_module, qconfig_dict, tracer.node_name_to_scope, @@ -462,9 +461,9 @@ def _convert_fx( _check_is_graph_module(graph_module) check_is_valid_convert_custom_config_dict(convert_custom_config_dict) - quantizer = Quantizer() - quantized = quantizer.convert(graph_module, is_reference, convert_custom_config_dict, - is_standalone_module, _remove_qconfig=_remove_qconfig) + quantized = convert( + graph_module, is_reference, convert_custom_config_dict, + is_standalone_module, _remove_qconfig_flag=_remove_qconfig) preserved_attributes = convert_custom_config_dict.get("preserved_attributes", []) for attr_name in preserved_attributes: From 9fbbab88da8c1affa0fe6f71ea0951a549f7751c Mon Sep 17 00:00:00 2001 From: Hangchen Yu Date: Tue, 15 Jun 2021 23:03:33 -0700 Subject: [PATCH 140/305] [fx-acc] Saturate host by replicating partitions onto idle devices (#60064) 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/60064 This implements a host saturation optimization to maximize the utilization of the available devices. It uses a greedy heuristic to replicate all partitions on the used devices to another set of idle devices with enough memory. The added unittest shows an example as follows: ``` partition_0: 192 bytes; partition_1: 48 bytes dev_0: 200 bytes, [partition_0] dev_1: 200 bytes, [partition_1] dev_2: 100 bytes, dev_3: 100 bytes, dev_4: 200 bytes, dev_5: 100 bytes ``` Before host saturation, `partition_0` is assigned to dev_0 and `partition_1` is assigned to dev_1. After host saturation, `partition_0` is replicated to dev_4 simply because it's the only device that can hold all partitions on dev_0. `partition_1` is replicated to dev_2 because it has minimal but large enough memory to hold all partitions on dev_1. Test Plan: ``` buck test mode/opt //caffe2/test:test_fx_experimental -- --exact 'caffe2/test:test_fx_experimental - test_saturate_host (test_fx_experimental.TestFXExperimental)' Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/8444249343103429 ✓ ListingSuccess: caffe2/test:test_fx_experimental - main (1.322) ✓ Pass: caffe2/test:test_fx_experimental - test_saturate_host (test_fx_experimental.TestFXExperimental) (1.322) Summary Pass: 1 ListingSuccess: 1 ``` An e2e test will be added to `test_fx_glow.py` in a followup diff. Reviewed By: gcatron Differential Revision: D29039998 fbshipit-source-id: 57518aadf668f7f05abd6ff73224c16b5d2a12ac --- test/test_fx_experimental.py | 41 +++++ .../experimental/accelerator_partitioner.py | 150 +++++++++++++++--- torch/fx/experimental/partitioner_utils.py | 2 + 3 files changed, 171 insertions(+), 22 deletions(-) diff --git a/test/test_fx_experimental.py b/test/test_fx_experimental.py index 49a2e3f9d080d..4d36095b48e23 100644 --- a/test/test_fx_experimental.py +++ b/test/test_fx_experimental.py @@ -609,6 +609,47 @@ def forward(self, a, b): ) assert (input1 * input2) == traced(input1, input2) + def test_saturate_host(self): + class TestModule(torch.nn.Module): + def __init__(self): + super(TestModule, self).__init__() + self.linear = torch.nn.Linear(4, 4) + + def forward(self, a): + add_1 = a + torch.rand(4) + add_2 = add_1 + torch.rand(4) + linear_1 = self.linear(add_1) + add_3 = add_2 + linear_1 + add_4 = add_2 + add_3 + return add_4 + + m = TestModule() + traced = symbolic_trace(m) + a = torch.rand(4) + graph_manipulation.get_size_of_all_nodes(traced, [a]) + devices = [ + Device("dev_0", 200, 0), + Device("dev_1", 200, 1), + Device("dev_2", 100, 2), + Device("dev_3", 100, 3), + Device("dev_4", 200, 4), + Device("dev_5", 100, 5), + ] + partitioner = Partitioner() + # Without host saturation, the model will be split into two partitions. + # dev_0 holds partition 0 of 192 bytes and dev_1 holds partition 1 of 48 bytes. + partitioner_config = PartitionerConfig(devices, saturate_host=True) + ret = partitioner.partition_graph(traced, m, partitioner_config) + module_with_submodules = ret.module_with_submodules + self.assertEqual(traced(a), module_with_submodules(a)) + + partitions = partitioner.partitions + self.assertEqual(len(partitions), 2) + # With host saturation, partition 1 will be replicated to dev_4, and partition 2 + # will be replicated to dev_2. + self.assertEqual(partitions[0].logical_device_ids, [0, 4]) + self.assertEqual(partitions[1].logical_device_ids, [1, 2]) + @skipIfNoTorchVision def test_conv_bn_fusion(self): rn18 = resnet18().eval() diff --git a/torch/fx/experimental/accelerator_partitioner.py b/torch/fx/experimental/accelerator_partitioner.py index c6904b042875e..e7a16a5ac8eb1 100644 --- a/torch/fx/experimental/accelerator_partitioner.py +++ b/torch/fx/experimental/accelerator_partitioner.py @@ -165,6 +165,51 @@ def get_node_to_partition_mapping(partitions: List[Partition]) -> Dict[Node, int return node_to_partition +def get_logical_id_to_device(devices: List[Device]) -> Dict[int, Device]: + """Get a mapping from device logical ID to Device object.""" + logical_id_to_device: Dict[int, Device] = {} + for d in devices: + logical_id_to_device[d.logical_id] = d + return logical_id_to_device + + +def get_device_partition_stats( + partitions: List[Partition], devices: List[Device] +) -> Tuple[Dict[Device, List[Partition]], Dict[Device, int], List[Partition]]: + """Given a list of partitions and a list of devices, returns: + 1. A mapping from device to partitions on it; + 2. A mapping from device to its remaining memory size; + 3. A list of partitions that do not have a device. + """ + # logical id to device + logical_id_to_device = get_logical_id_to_device(devices) + # Track partitions on device + device_to_partitions: Dict[Device, List[Partition]] = {} + # Track device's left mem size + device_to_left_mem_bytes: Dict[Device, int] = {} + for d in devices: + device_to_partitions[d] = [] + device_to_left_mem_bytes[d] = d.available_mem_bytes + + # Deal with the partitions that already have a device + # and also collect all partitions without a device (no_device_partitions) + no_device_partitions = [] + for partition in partitions: + if partition.logical_device_ids != []: + for logical_id in partition.logical_device_ids: + device = logical_id_to_device[logical_id] + device_to_partitions[device].append(partition) + device_to_left_mem_bytes[device] -= partition.used_mem_bytes + else: + no_device_partitions.append(partition) + + return ( + device_to_partitions, + device_to_left_mem_bytes, + no_device_partitions, + ) + + def get_device_to_partitions_mapping( partitions: List[Partition], devices: List[Device] ): @@ -204,27 +249,12 @@ def find_device_for(partition: Partition): return True return False - # logical id to device - logical_id_to_device: Dict[int, Device] = {} - # Track partitions on device - device_to_partitions: Dict[Device, List[Partition]] = {} - # Track device's left mem size - device_to_left_mem_bytes: Dict[Device, int] = {} - for d in devices: - logical_id_to_device[d.logical_id] = d - device_to_partitions[d] = [] - device_to_left_mem_bytes[d] = d.available_mem_bytes - # Deal with the partitions that already have a device - # and also collect all partitions without a device (no_device_partitions) - no_device_partitions = [] - for partition in partitions: - if partition.logical_device_ids != []: - logical_id = partition.logical_device_ids[0] - device = logical_id_to_device[logical_id] - device_to_partitions[device] = [partition] - device_to_left_mem_bytes[device] -= partition.used_mem_bytes - else: - no_device_partitions.append(partition) + ( + device_to_partitions, + device_to_left_mem_bytes, + no_device_partitions, + ) = get_device_partition_stats(partitions, devices) + # Find devices for all the partitions without a device found_device = True for partition in no_device_partitions: @@ -341,7 +371,14 @@ def partition_graph( ) else: self.size_based_partition() + + # Saturate host if possible. + if partitioner_config.saturate_host: + self.saturate_host() + + # Partition the graph module based on the partition assignment. module_with_submodules = self.do_partition() + # The DAG contains DAGNodes with info of each partition's input nodes, output nodes # and how partitions are connected. dag = self.dump_dag(module_with_submodules) @@ -459,6 +496,75 @@ def find_device_based_on_size(node) -> Device: raise RuntimeError("Cannot Get a Valid Partition to Logical Device Mapping") return + def saturate_host(self) -> None: + """Saturate host by assigning replicates to unused devices with enough memory. + It uses a greedy approach to find a next available set of devices to place all split + partitions: For each used device, it searches for an idle device with minimal memory + size that can hold all the partition located on that device; If the search is successful + for all used devices, it then assigns the new devices' logical ID to the corresponding + partition. + """ + ( + device_to_partitions, + device_to_left_mem_bytes, + no_device_partitions, + ) = get_device_partition_stats(self.partitions, self.devices) + + assert ( + len(no_device_partitions) == 0 + ), f"Expect no_device_partitions has 0 device, but get {len(no_device_partitions)}" + + # Devices that hold partitions + used_devices = [d for d in self.devices if len(device_to_partitions[d]) > 0] + # Track replicates of the assigned devices + replicated_device_to_used_device: Dict[Device, Device] = {} + + while len(used_devices) * 2 + len(replicated_device_to_used_device) <= len( + self.devices + ): + # Success flag for this round + success = True + # Devices that have not been assigned + idle_devices = [ + d + for d in self.devices + if d not in used_devices and d not in replicated_device_to_used_device + ] + # Temporary mapping from replicated device to original device + temp_replicate_mapping = {} + + # Find a new device to replicate all partitions on an used device + for used_device in used_devices: + # Idle devices that have enough memory + available_devices = [ + d + for d in idle_devices + if d.available_mem_bytes + >= used_device.available_mem_bytes + - device_to_left_mem_bytes[used_device] + ] + if len(available_devices) == 0: + success = False + break + new_device = min(available_devices, key=lambda d: d.available_mem_bytes) + idle_devices.remove(new_device) + temp_replicate_mapping[new_device] = used_device + + if not success: + break + replicated_device_to_used_device.update(temp_replicate_mapping) + + # Update logical device IDs assigned to the partitions + for ( + replicate_device, + original_device, + ) in replicated_device_to_used_device.items(): + logical_id = replicate_device.logical_id + for partition in device_to_partitions[original_device]: + partition.logical_device_ids.append(logical_id) + for p in self.partitions: + print(p.logical_device_ids) + def do_partition(self) -> GraphModule: """Return a new fx module with submodule nodes (partitions).""" module_with_submodules = split_module( @@ -469,7 +575,7 @@ def do_partition(self) -> GraphModule: return module_with_submodules def dump_dag(self, module_with_submodules: GraphModule) -> DAG: - """Return the dag structure and the new fx module with submodules""" + """Return the dag structure and the new fx module with submodules.""" dag = DAG() for node in module_with_submodules.graph.nodes: if node.op == "output": diff --git a/torch/fx/experimental/partitioner_utils.py b/torch/fx/experimental/partitioner_utils.py index cffba192b6b79..b99eb906e4dcf 100644 --- a/torch/fx/experimental/partitioner_utils.py +++ b/torch/fx/experimental/partitioner_utils.py @@ -93,6 +93,8 @@ class PartitionerConfig(NamedTuple): node_to_latency_mapping: Dict[Node, NodeLatency] = {} node_to_partition_mapping: Dict[Node, int] = {} partition_to_logical_device_mapping: Dict[int, List[int]] = {} + # Saturate host by replicating partitions to the remaining idle devices. + saturate_host: bool = False def get_extra_size_of(node: Node, nodes: Set[Node]) -> int: From 36a5647e30e46172c985f6bb422b2bd5d746b96e Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Wed, 16 Jun 2021 00:40:00 -0700 Subject: [PATCH 141/305] Handle exceptions from THPModule_setQEngine (#60073) Summary: Prevents Python runtime crashes when `torch._C._set_qengine(2**65)` or `torch.backends.quantized.engine="fbgemm"` if PyTorch was compiled without fbgemm Pull Request resolved: https://github.com/pytorch/pytorch/pull/60073 Reviewed By: supriyar Differential Revision: D29156430 Pulled By: malfet fbshipit-source-id: 95b97352a52a262f1634b72da64a0c950eaf2373 --- torch/csrc/Module.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp index 921dc8cac51d0..ed69a6542c057 100644 --- a/torch/csrc/Module.cpp +++ b/torch/csrc/Module.cpp @@ -574,9 +574,11 @@ PyObject *THPModule_setQEngine(PyObject */* unused */, PyObject *arg) { THPUtils_assert(THPUtils_checkLong(arg), "set_qengine expects an int, " "but got %s", THPUtils_typename(arg)); + HANDLE_TH_ERRORS auto qengine = static_cast(THPUtils_unpackLong(arg)); at::globalContext().setQEngine(static_cast(qengine)); Py_RETURN_NONE; + END_HANDLE_TH_ERRORS } PyObject *THPModule_qEngine(PyObject *_unused, PyObject *noargs) From f43ff754ca0021afb6bb5b83ac0f8e27301ba49c Mon Sep 17 00:00:00 2001 From: lezcano Date: Wed, 16 Jun 2021 01:20:08 -0700 Subject: [PATCH 142/305] [docs] Correct errata in linalg.eigh and add a bit more information (#59784) Summary: Add extra information about the returned elements of the spectral decompositions Resolves https://github.com/pytorch/pytorch/issues/59718 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59784 Reviewed By: soulitzer Differential Revision: D29088998 Pulled By: mruberry fbshipit-source-id: 58a191c41ff5e4c9d9675e5b3d7cbbcf16be4da1 --- torch/linalg/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/torch/linalg/__init__.py b/torch/linalg/__init__.py index 2fed1a5a8b212..8a0f9a8ca047e 100644 --- a/torch/linalg/__init__.py +++ b/torch/linalg/__init__.py @@ -476,7 +476,8 @@ Returns: A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`V` above. - `eigenvalues` and `eigenvectors` will always be complex-valued, even when :attr:`A` is real. + `eigenvalues` and `eigenvectors` will always be complex-valued, even when :attr:`A` is real. The eigenvectors + will be given by the columns of `eigenvectors`. Examples:: @@ -634,12 +635,12 @@ out (tuple, optional): output tuple of two tensors. Ignored if `None`. Default: `None`. Returns: - A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`V` above. + A named tuple `(eigenvalues, eigenvectors)` which corresponds to :math:`\Lambda` and :math:`Q` above. `eigenvalues` will always be real-valued, even when :attr:`A` is complex. It will also be ordered in ascending order. - `eigenvectors` will have the same dtype as :attr:`A`. + `eigenvectors` will have the same dtype as :attr:`A` and will contain the eigenvectors as its columns. Examples:: @@ -1522,7 +1523,8 @@ `S` will always be real-valued, even when :attr:`A` is complex. It will also be ordered in descending order. - `U` and `Vh` will have the same dtype as :attr:`A`. + `U` and `Vh` will have the same dtype as :attr:`A`. The left / right singular vectors will be given by + the columns of `U` and the rows of `Vh` respectively. Examples:: From bda40639c53a96cecc9eb1c21d47257d927ebe74 Mon Sep 17 00:00:00 2001 From: Bert Maher Date: Wed, 16 Jun 2021 05:07:43 -0700 Subject: [PATCH 143/305] [nnc] Move operator implementations into a subdirectory (#59988) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59988 As we broaden operator support, putting all the implementations into kernel.cpp is getting unwieldy. Let's factor them out into the "operators" subdirectory. This diff is big but it's entirely code movement; I didn't change anything, other than to expose a few utilities in kernel.h. ghstack-source-id: 131405139 Test Plan: CI Reviewed By: ZolotukhinM Differential Revision: D29115916 fbshipit-source-id: ba0df1d8dd4a108b584da3baf168407e966b2c78 --- tools/build_variables.bzl | 3 + torch/csrc/jit/tensorexpr/kernel.cpp | 381 +++--------------- torch/csrc/jit/tensorexpr/kernel.h | 19 + .../csrc/jit/tensorexpr/operators/matmul.cpp | 55 +++ torch/csrc/jit/tensorexpr/operators/matmul.h | 16 + .../csrc/jit/tensorexpr/operators/operators.h | 6 + .../jit/tensorexpr/operators/reduction.cpp | 105 +++++ .../csrc/jit/tensorexpr/operators/reduction.h | 15 + .../csrc/jit/tensorexpr/operators/softmax.cpp | 160 ++++++++ torch/csrc/jit/tensorexpr/operators/softmax.h | 16 + 10 files changed, 442 insertions(+), 334 deletions(-) create mode 100644 torch/csrc/jit/tensorexpr/operators/matmul.cpp create mode 100644 torch/csrc/jit/tensorexpr/operators/matmul.h create mode 100644 torch/csrc/jit/tensorexpr/operators/operators.h create mode 100644 torch/csrc/jit/tensorexpr/operators/reduction.cpp create mode 100644 torch/csrc/jit/tensorexpr/operators/reduction.h create mode 100644 torch/csrc/jit/tensorexpr/operators/softmax.cpp create mode 100644 torch/csrc/jit/tensorexpr/operators/softmax.h diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 8258e0e6429eb..e7421d8a292e8 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -291,6 +291,9 @@ core_sources_full_mobile = [ "torch/csrc/jit/tensorexpr/mem_arena.cpp", "torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp", "torch/csrc/jit/tensorexpr/operators/conv2d.cpp", + "torch/csrc/jit/tensorexpr/operators/matmul.cpp", + "torch/csrc/jit/tensorexpr/operators/reduction.cpp", + "torch/csrc/jit/tensorexpr/operators/softmax.cpp", "torch/csrc/jit/tensorexpr/reduction.cpp", "torch/csrc/jit/tensorexpr/registerizer.cpp", "torch/csrc/jit/tensorexpr/tensor.cpp", diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 3884ce1c762fb..482c2560f7ae8 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -11,8 +11,7 @@ #include #include #include -#include -#include +#include using namespace torch::jit; using namespace torch::jit::tensorexpr; @@ -338,9 +337,21 @@ std::shared_ptr removeUnusedSelfArgument( return res; } -} // namespace tensorexpr -} // namespace jit -} // namespace torch +std::vector valueShape(const ArgValue& v) { + if (auto b = c10::get_if(&v)) { + return b->dims(); + } + return {}; +} + +ExprHandle tensorOrConstant( + const ArgValue& v, + const std::vector& axes) { + if (auto b = c10::get_if(&v)) { + return broadcast(*b, axes); + } + return constant(v); +} size_t normalizeAndCheckIndex(int64_t idx, int64_t list_size) { if (idx < 0) { @@ -354,11 +365,30 @@ size_t normalizeAndCheckIndex(int64_t idx, int64_t list_size) { return static_cast(idx); } -static at::ScalarType tensorType(const Buf* b) { - return static_cast(b->dtype().scalar_type()); +ExprHandle broadcast(BufHandle b, const std::vector& axes) { + return b.load(computeIndicesToBroadcast(axes, b.dims())); } -static std::vector computeIndicesToBroadcast( +ExprHandle constant(const ArgValue& v) { + if (auto s = c10::get_if(&v)) { + return *s; + } else if (auto d = c10::get_if(&v)) { + return DoubleImm::make(*d); + } else if (auto i = c10::get_if(&v)) { + return LongImm::make(*i); + } else if (auto b = c10::get_if(&v)) { + return BoolImm::make(*b); + } else if (c10::get_if(&v)) { + // This is just a placeholder so we don't throw. None-handling + // is operator-specific and should be handled properly in + // the operator-specific lowering code. + return IntImm::make(0); + } else { + throw unsupported_dtype("Trying to convert unsupported dtype to constant"); + } +} + +std::vector computeIndicesToBroadcast( const std::vector& outputAxes, const std::vector& inputSizes) { if (outputAxes.size() < inputSizes.size()) { @@ -381,6 +411,14 @@ static std::vector computeIndicesToBroadcast( return bcast; } +} // namespace tensorexpr +} // namespace jit +} // namespace torch + +static at::ScalarType tensorType(const Buf* b) { + return static_cast(b->dtype().scalar_type()); +} + std::vector bufferSizes(const Buf* b) { std::vector sizes; for (size_t i = 0; i < b->ndim(); i++) { @@ -430,37 +468,6 @@ ExprHandle promoteToDtype(ExprHandle e, ScalarType dt) { return e; } -ExprHandle broadcast(BufHandle b, const std::vector& axes) { - return b.load(computeIndicesToBroadcast(axes, b.dims())); -} - -ExprHandle constant(const ArgValue& v) { - if (auto s = c10::get_if(&v)) { - return *s; - } else if (auto d = c10::get_if(&v)) { - return DoubleImm::make(*d); - } else if (auto i = c10::get_if(&v)) { - return LongImm::make(*i); - } else if (auto b = c10::get_if(&v)) { - return BoolImm::make(*b); - } else if (c10::get_if(&v)) { - // This is just a placeholder so we don't throw. None-handling - // is operator-specific and should be handled properly in - // the operator-specific lowering code. - return IntImm::make(0); - } else { - throw unsupported_dtype("Trying to convert unsupported dtype to constant"); - } -} - -ExprHandle tensorOrConstant( - const ArgValue& v, - const std::vector& axes) { - if (auto b = c10::get_if(&v)) { - return broadcast(*b, axes); - } - return constant(v); -} ExprHandle TensorExprKernel::constant(const torch::jit::Value* v) { if (v->node()->kind() == prim::Constant) { const auto val = toIValue(v).value(); @@ -950,13 +957,6 @@ std::vector TensorExprKernel::broadcastShapesMut( return res.first; } -std::vector valueShape(const ArgValue& v) { - if (auto b = c10::get_if(&v)) { - return b->dims(); - } - return {}; -} - Tensor* computeOneOperand( const std::string& name, const std::vector& inputValues, @@ -1271,294 +1271,6 @@ Tensor* computeCat( }); } -// Remove all indices from axes positions. -std::vector squeezeIndices( - const ParameterList& indices, - const std::vector& axes) { - std::vector indices_squeezed; - for (size_t dim = 0; dim < indices.size(); ++dim) { - if (!std::count(axes.begin(), axes.end(), dim)) { - indices_squeezed.push_back(indices[dim]); - } - } - return indices_squeezed; -} - -Tensor* computeSoftmax( - const std::vector& inputs, - const std::vector& outputShape, - bool log_softmax) { - // Softmax is computed as follows: - // softmax(vi) = exp(vi) / sum(exp(vi)) - // - // In order to avoid overflow issues due to exp of a large number, we - // subtract the max of that dim before computing exp. - // softmax(vi) = exp(vi - max(vi)) / sum(exp(vi - max(vi))) - // - // This is implemented as 4 loopnests: - // - First loop computes the max over the softmax dim. - // - Second loop computes exp for every element in v after subtracting - // the max of the softmax dim it belongs to. - // - Third loop computes the sum over the softmax dim. - // - Final loop computes softmax for every element in v. - - // LogSoftmax is computed as follows: - // log_softmax(vi) = log(softmax(vi)) - // = vi - log(sum(exp(vi))) - // - // Using the same max trick as above: - // log_softmax(vi) = vi - max(vi) - log(sum(exp(vi - max(vi)))) - // - // This is implemented as 5 loopnests: - // - First loop computes the max over the softmax dim. - // - Second loop computes exp for every element in v after subtracting - // the max of the softmax dim it belongs to. - // - Third loop computes the sum over the softmax dim. - // - Fourth loop computes log for every element in the sum. - // - Final loop computes the log_softmax for every element in v. - - TORCH_INTERNAL_ASSERT(inputs.size() == 3); - auto output_dims = c10::fmap(outputShape); - - // We do not handle None for dims (input 1) because that is supposed to - // be deprecated. - TORCH_INTERNAL_ASSERT(c10::get_if(&inputs[1])); - int64_t rank = valueShape(inputs[0]).size(); - size_t softmax_dim = - normalizeAndCheckIndex(c10::get(inputs[1]), rank); - std::vector non_softmax_dims; - for (size_t i = 0; i < output_dims.size(); ++i) { - if (i != softmax_dim) { - non_softmax_dims.push_back(output_dims[i]); - } - } - - // Softmax implementation includes two reductions, one to find the max and - // the other to calculate the sum along the softmax dim. These reductions - // will have the softmax dimension as the inner most loop. So, the innermost - // index in the indices will refer to the softmax dimension. - - // Update the indices by moving the softmax dimension index to the - // appropriate position. - auto move_softmax_dim_index_to_pos = [&](const ParameterList& indices) { - std::vector new_indices; - for (auto ind : indices) { - new_indices.push_back(ind); - } - for (size_t i = softmax_dim; i < indices.size() - 1; ++i) { - new_indices[i + 1] = indices[i]; - } - new_indices[softmax_dim] = indices[indices.size() - 1]; - return new_indices; - }; - - // Remove the index corresponding to the softmax dimension. - auto remove_softmax_dim_index = [&](const ParameterList& indices) { - std::vector new_indices; - for (size_t i = 0; i < indices.size(); ++i) { - if (i != softmax_dim) { - new_indices.push_back(indices[i]); - } - } - return new_indices; - }; - - auto convert_indices_to_expr_handle = [&](const ParameterList& indices) { - std::vector new_indices(indices.size()); - for (size_t i = 0; i < indices.size(); ++i) { - new_indices[i] = indices[i]; - } - return new_indices; - }; - - c10::optional dtype = ToDtype(ScalarType::Undefined); - if (auto d = c10::get_if(&inputs[2])) { - dtype = ToDtype(static_cast(*d)); - } - - auto max = Reduce( - "aten_softmax_max", - non_softmax_dims, - Maximum(dtype.value()), - [&](ParameterList& indices) { - return tensorOrConstant( - inputs[0], move_softmax_dim_index_to_pos(indices)); - }, - {output_dims[softmax_dim]}); - auto e = - Compute("aten_softmax_exp", output_dims, [&](ParameterList& indices) { - auto inp = tensorOrConstant( - inputs[0], convert_indices_to_expr_handle(indices)); - return exp(inp - max->load(remove_softmax_dim_index(indices))); - }); - auto sum = Reduce( - "aten_softmax_sum", - non_softmax_dims, - Sum(), - [&](ParameterList& indices) { - return e->load(move_softmax_dim_index_to_pos(indices)); - }, - {output_dims[softmax_dim]}); - if (!log_softmax) { - auto result = - Compute("aten_softmax", output_dims, [&](ParameterList& indices) { - return e->load(indices) / - sum->load(remove_softmax_dim_index(indices)); - }); - return new Tensor( - result->buf(), - new tensorexpr::Block( - {max->stmt(), e->stmt(), sum->stmt(), result->stmt()})); - } - - auto log_sum = Compute( - "aten_softmax_log_sum", non_softmax_dims, [&](ParameterList& indices) { - return log(sum->load(indices)); - }); - auto result = - Compute("aten_log_softmax", output_dims, [&](ParameterList& indices) { - auto inp = tensorOrConstant( - inputs[0], convert_indices_to_expr_handle(indices)); - auto non_softmax_indices = remove_softmax_dim_index(indices); - return inp - max->load(non_softmax_indices) - - log_sum->load(non_softmax_indices); - }); - return new Tensor( - result->buf(), - new tensorexpr::Block( - {max->stmt(), - e->stmt(), - sum->stmt(), - log_sum->stmt(), - result->stmt()})); -} - -Tensor* computeSum( - const std::vector& inputs, - const c10::optional& outputType) { - std::vector axes; - bool keepdim = false; - // aten::sum takes the input tensor named self. - auto sizes = valueShape(inputs[0]); - - size_t rank = sizes.size(); - if (inputs.size() > 2) { - if (auto emptyAxes = c10::get_if(&inputs[1])) { - // If dim-array is an empty list, it will appear as BufList instead of - // IntList, and hence we need a special handling for it. - // In that case, we need to sum over all axes. - TORCH_INTERNAL_ASSERT(emptyAxes->empty()); - axes.resize(rank); - std::iota(axes.begin(), axes.end(), 0); - } else if (rank > 0) { - auto nodeAxes = c10::get(inputs[1]); - // Canonicalize axes: wrap around, sort and make unique. - for (auto axis : nodeAxes) { - axes.push_back(at::maybe_wrap_dim(axis, rank)); - } - std::sort(axes.begin(), axes.end()); - axes.erase(std::unique(axes.begin(), axes.end()), axes.end()); - } - keepdim = c10::get(inputs[2]); - } else { - axes.resize(rank); - std::iota(axes.begin(), axes.end(), 0); - } - // Axes go into reduction dimensions. - std::vector reductionDims; - reductionDims.reserve(rank); - for (size_t axis : axes) { - reductionDims.emplace_back(sizes[axis]); - } - std::vector outputDims; - // Output dimensions are the complement of axes. When keepdim is set, a - // one-sized dimension is inserted for each axis. - for (size_t dim = 0; dim < rank; ++dim) { - if (!std::count(axes.begin(), axes.end(), dim)) { - outputDims.emplace_back(sizes[dim]); - } else if (keepdim) { - outputDims.emplace_back(1); - } - } - - return Reduce( - "sum", - outputDims, - Sum(), - [&](ParameterList& indices) { - // "Squeeze" out indices inserted when keepdim is set. - auto indices_squeezed = - keepdim ? squeezeIndices(indices, axes) : indices; - TORCH_INTERNAL_ASSERT(axes.size() <= indices_squeezed.size()); - // Move innermost indices into axes positions: - // 1. Fill the outermost indices first. - // 2. Insert the innermost indices into the correct axis position, - // displacing the outermost indices as needed. - std::vector indices_exprs; - size_t i = 0; - for (; i < indices_squeezed.size() - axes.size(); ++i) { - indices_exprs.push_back(indices_squeezed[i]); - } - for (auto axis : axes) { - indices_exprs.insert( - indices_exprs.begin() + axis, indices_squeezed[i]); - ++i; - } - auto indexed = tensorOrConstant(inputs[0], indices_exprs); - if (outputType) { - return Cast::make(ToDtype(*outputType), indexed); - } else { - return indexed; - } - }, - reductionDims); -} - -Tensor* computeMatmul( - const std::vector& inputs, - const std::vector& outputShape, - const c10::optional& outputType) { - Dtype dtype = kFloat; - if (outputType) { - dtype = Dtype(*outputType); - } - BufHandle ResultBuf("matmul", outputShape, dtype); - const BufHandle a = c10::get(inputs[0]); - const BufHandle b = c10::get(inputs[1]); - - auto size_a = a.dims(); - auto size_b = b.dims(); - // We currently only support rank 2 matmuls - TORCH_INTERNAL_ASSERT(size_a.size() == 2 && size_b.size() == 2); - auto total_size = dynamic_cast( - IRSimplifier::simplify( - cast(size_a[0]) * cast(size_a[1]) * - cast(size_b[1])) - .node()); - - // For small sizes, where N*M*K < 1000, lower matmul to a naive 3-level - // loopnest. The number is not tuned very carefully, and in future we should - // fine-tune it as well as we should add more advanced native TE lowerings for - // matmuls. For bigger sizes we generate a TE ExternalCall, which would call - // an aten::matmul. - // Native, even naive, lowering is beneficial when the sizes are small because - // it allows to eliminate dispatch overhead. - if (total_size && total_size->value() < 1000) { - return Reduce( - "nnc_matmul", - {{size_a[0], "M"}, {size_b[1], "N"}}, - Sum(), - [&](const ExprHandle& m, const ExprHandle& n, const ExprHandle& k) { - return Load::make(a, {m, k}) * Load::make(b, {k, n}); - }, - {{size_a[1], "K"}}); - } else { - return new Tensor( - ResultBuf.node(), - ExternalCall::make(ResultBuf, "nnc_aten_matmul", {a, b}, {})); - } -} - Tensor* computeConv2d( const std::vector& inputs, const std::vector& outputShape, @@ -1957,6 +1669,7 @@ Tensor* tensorexpr::computeOperandValue( return a * point_five * (one + erf(a * m_sqrt1_2)); }); } break; + case aten::batch_norm: { bool hasWeight = true; bool hasBias = true; diff --git a/torch/csrc/jit/tensorexpr/kernel.h b/torch/csrc/jit/tensorexpr/kernel.h index f464773bf5e6b..88333e0c2c756 100644 --- a/torch/csrc/jit/tensorexpr/kernel.h +++ b/torch/csrc/jit/tensorexpr/kernel.h @@ -36,6 +36,25 @@ using ArgValue = c10::variant< IntList, ArgNone>; +// Get the dimensions of a value. +std::vector valueShape(const ArgValue& v); + +// If v is a tensor, broadcast it to match the shape of axes, or return +// directly if v is a constant. +ExprHandle tensorOrConstant( + const ArgValue& v, + const std::vector& axes); + +size_t normalizeAndCheckIndex(int64_t idx, int64_t list_size); + +ExprHandle broadcast(BufHandle b, const std::vector& axes); + +ExprHandle constant(const ArgValue& v); + +std::vector computeIndicesToBroadcast( + const std::vector& outputAxes, + const std::vector& inputSizes); + inline std::string getArgValueName(const ArgValue& a) { if (c10::get_if(&a)) { return "BufHandle"; diff --git a/torch/csrc/jit/tensorexpr/operators/matmul.cpp b/torch/csrc/jit/tensorexpr/operators/matmul.cpp new file mode 100644 index 0000000000000..30c34978a066f --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/matmul.cpp @@ -0,0 +1,55 @@ +#include +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeMatmul( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType) { + Dtype dtype = kFloat; + if (outputType) { + dtype = Dtype(*outputType); + } + BufHandle ResultBuf("matmul", outputShape, dtype); + const BufHandle a = c10::get(inputs[0]); + const BufHandle b = c10::get(inputs[1]); + + auto size_a = a.dims(); + auto size_b = b.dims(); + // We currently only support rank 2 matmuls + TORCH_INTERNAL_ASSERT(size_a.size() == 2 && size_b.size() == 2); + auto total_size = dynamic_cast( + IRSimplifier::simplify( + cast(size_a[0]) * cast(size_a[1]) * + cast(size_b[1])) + .node()); + + // For small sizes, where N*M*K < 1000, lower matmul to a naive 3-level + // loopnest. The number is not tuned very carefully, and in future we should + // fine-tune it as well as we should add more advanced native TE lowerings for + // matmuls. For bigger sizes we generate a TE ExternalCall, which would call + // an aten::matmul. + // Native, even naive, lowering is beneficial when the sizes are small because + // it allows to eliminate dispatch overhead. + if (total_size && total_size->value() < 1000) { + return Reduce( + "nnc_matmul", + {{size_a[0], "M"}, {size_b[1], "N"}}, + Sum(), + [&](const ExprHandle& m, const ExprHandle& n, const ExprHandle& k) { + return Load::make(a, {m, k}) * Load::make(b, {k, n}); + }, + {{size_a[1], "K"}}); + } else { + return new Tensor( + ResultBuf.node(), + ExternalCall::make(ResultBuf, "nnc_aten_matmul", {a, b}, {})); + } +} + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/matmul.h b/torch/csrc/jit/tensorexpr/operators/matmul.h new file mode 100644 index 0000000000000..893b3a9c820de --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/matmul.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeMatmul( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType); + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/operators.h b/torch/csrc/jit/tensorexpr/operators/operators.h new file mode 100644 index 0000000000000..05705e2d4e89c --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/operators.h @@ -0,0 +1,6 @@ +#pragma once + +#include +#include +#include +#include diff --git a/torch/csrc/jit/tensorexpr/operators/reduction.cpp b/torch/csrc/jit/tensorexpr/operators/reduction.cpp new file mode 100644 index 0000000000000..2f78298a060e9 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/reduction.cpp @@ -0,0 +1,105 @@ +#include + +using namespace torch::jit::tensorexpr; + +// Remove all indices from axes positions. +static std::vector squeezeIndices( + const ParameterList& indices, + const std::vector& axes) { + std::vector indices_squeezed; + for (size_t dim = 0; dim < indices.size(); ++dim) { + if (!std::count(axes.begin(), axes.end(), dim)) { + indices_squeezed.push_back(indices[dim]); + } + } + return indices_squeezed; +} + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeSum( + const std::vector& inputs, + const c10::optional& outputType) { + std::vector axes; + bool keepdim = false; + // aten::sum takes the input tensor named self. + auto sizes = valueShape(inputs[0]); + + size_t rank = sizes.size(); + if (inputs.size() > 2) { + if (auto emptyAxes = c10::get_if(&inputs[1])) { + // If dim-array is an empty list, it will appear as BufList instead of + // IntList, and hence we need a special handling for it. + // In that case, we need to sum over all axes. + TORCH_INTERNAL_ASSERT(emptyAxes->empty()); + axes.resize(rank); + std::iota(axes.begin(), axes.end(), 0); + } else if (rank > 0) { + auto nodeAxes = c10::get(inputs[1]); + // Canonicalize axes: wrap around, sort and make unique. + for (auto axis : nodeAxes) { + axes.push_back(at::maybe_wrap_dim(axis, rank)); + } + std::sort(axes.begin(), axes.end()); + axes.erase(std::unique(axes.begin(), axes.end()), axes.end()); + } + keepdim = c10::get(inputs[2]); + } else { + axes.resize(rank); + std::iota(axes.begin(), axes.end(), 0); + } + // Axes go into reduction dimensions. + std::vector reductionDims; + reductionDims.reserve(rank); + for (size_t axis : axes) { + reductionDims.emplace_back(sizes[axis]); + } + std::vector outputDims; + // Output dimensions are the complement of axes. When keepdim is set, a + // one-sized dimension is inserted for each axis. + for (size_t dim = 0; dim < rank; ++dim) { + if (!std::count(axes.begin(), axes.end(), dim)) { + outputDims.emplace_back(sizes[dim]); + } else if (keepdim) { + outputDims.emplace_back(1); + } + } + + return Reduce( + "sum", + outputDims, + Sum(), + [&](ParameterList& indices) { + // "Squeeze" out indices inserted when keepdim is set. + auto indices_squeezed = + keepdim ? squeezeIndices(indices, axes) : indices; + TORCH_INTERNAL_ASSERT(axes.size() <= indices_squeezed.size()); + // Move innermost indices into axes positions: + // 1. Fill the outermost indices first. + // 2. Insert the innermost indices into the correct axis position, + // displacing the outermost indices as needed. + std::vector indices_exprs; + size_t i = 0; + for (; i < indices_squeezed.size() - axes.size(); ++i) { + indices_exprs.push_back(indices_squeezed[i]); + } + for (auto axis : axes) { + indices_exprs.insert( + indices_exprs.begin() + axis, indices_squeezed[i]); + ++i; + } + auto indexed = tensorOrConstant(inputs[0], indices_exprs); + if (outputType) { + return Cast::make(ToDtype(*outputType), indexed); + } else { + return indexed; + } + }, + reductionDims); +} + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/reduction.h b/torch/csrc/jit/tensorexpr/operators/reduction.h new file mode 100644 index 0000000000000..aeb4cd35b8765 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/reduction.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeSum( + const std::vector& inputs, + const c10::optional& outputType); + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/softmax.cpp b/torch/csrc/jit/tensorexpr/operators/softmax.cpp new file mode 100644 index 0000000000000..edb911e0f5cb7 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/softmax.cpp @@ -0,0 +1,160 @@ +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +using namespace torch::jit::tensorexpr; + +Tensor* computeSoftmax( + const std::vector& inputs, + const std::vector& outputShape, + bool log_softmax) { + // Softmax is computed as follows: + // softmax(vi) = exp(vi) / sum(exp(vi)) + // + // In order to avoid overflow issues due to exp of a large number, we + // subtract the max of that dim before computing exp. + // softmax(vi) = exp(vi - max(vi)) / sum(exp(vi - max(vi))) + // + // This is implemented as 4 loopnests: + // - First loop computes the max over the softmax dim. + // - Second loop computes exp for every element in v after subtracting + // the max of the softmax dim it belongs to. + // - Third loop computes the sum over the softmax dim. + // - Final loop computes softmax for every element in v. + + // LogSoftmax is computed as follows: + // log_softmax(vi) = log(softmax(vi)) + // = vi - log(sum(exp(vi))) + // + // Using the same max trick as above: + // log_softmax(vi) = vi - max(vi) - log(sum(exp(vi - max(vi)))) + // + // This is implemented as 5 loopnests: + // - First loop computes the max over the softmax dim. + // - Second loop computes exp for every element in v after subtracting + // the max of the softmax dim it belongs to. + // - Third loop computes the sum over the softmax dim. + // - Fourth loop computes log for every element in the sum. + // - Final loop computes the log_softmax for every element in v. + + TORCH_INTERNAL_ASSERT(inputs.size() == 3); + auto output_dims = c10::fmap(outputShape); + + // We do not handle None for dims (input 1) because that is supposed to + // be deprecated. + TORCH_INTERNAL_ASSERT(c10::get_if(&inputs[1])); + int64_t rank = valueShape(inputs[0]).size(); + size_t softmax_dim = + normalizeAndCheckIndex(c10::get(inputs[1]), rank); + std::vector non_softmax_dims; + for (size_t i = 0; i < output_dims.size(); ++i) { + if (i != softmax_dim) { + non_softmax_dims.push_back(output_dims[i]); + } + } + + // Softmax implementation includes two reductions, one to find the max and + // the other to calculate the sum along the softmax dim. These reductions + // will have the softmax dimension as the inner most loop. So, the innermost + // index in the indices will refer to the softmax dimension. + + // Update the indices by moving the softmax dimension index to the + // appropriate position. + auto move_softmax_dim_index_to_pos = [&](const ParameterList& indices) { + std::vector new_indices; + for (auto ind : indices) { + new_indices.push_back(ind); + } + for (size_t i = softmax_dim; i < indices.size() - 1; ++i) { + new_indices[i + 1] = indices[i]; + } + new_indices[softmax_dim] = indices[indices.size() - 1]; + return new_indices; + }; + + // Remove the index corresponding to the softmax dimension. + auto remove_softmax_dim_index = [&](const ParameterList& indices) { + std::vector new_indices; + for (size_t i = 0; i < indices.size(); ++i) { + if (i != softmax_dim) { + new_indices.push_back(indices[i]); + } + } + return new_indices; + }; + + auto convert_indices_to_expr_handle = [&](const ParameterList& indices) { + std::vector new_indices(indices.size()); + for (size_t i = 0; i < indices.size(); ++i) { + new_indices[i] = indices[i]; + } + return new_indices; + }; + + c10::optional dtype = ToDtype(ScalarType::Undefined); + if (auto d = c10::get_if(&inputs[2])) { + dtype = ToDtype(static_cast(*d)); + } + + auto max = Reduce( + "aten_softmax_max", + non_softmax_dims, + Maximum(dtype.value()), + [&](ParameterList& indices) { + return tensorOrConstant( + inputs[0], move_softmax_dim_index_to_pos(indices)); + }, + {output_dims[softmax_dim]}); + auto e = + Compute("aten_softmax_exp", output_dims, [&](ParameterList& indices) { + auto inp = tensorOrConstant( + inputs[0], convert_indices_to_expr_handle(indices)); + return exp(inp - max->load(remove_softmax_dim_index(indices))); + }); + auto sum = Reduce( + "aten_softmax_sum", + non_softmax_dims, + Sum(), + [&](ParameterList& indices) { + return e->load(move_softmax_dim_index_to_pos(indices)); + }, + {output_dims[softmax_dim]}); + if (!log_softmax) { + auto result = + Compute("aten_softmax", output_dims, [&](ParameterList& indices) { + return e->load(indices) / + sum->load(remove_softmax_dim_index(indices)); + }); + return new Tensor( + result->buf(), + new tensorexpr::Block( + {max->stmt(), e->stmt(), sum->stmt(), result->stmt()})); + } + + auto log_sum = Compute( + "aten_softmax_log_sum", non_softmax_dims, [&](ParameterList& indices) { + return log(sum->load(indices)); + }); + auto result = + Compute("aten_log_softmax", output_dims, [&](ParameterList& indices) { + auto inp = tensorOrConstant( + inputs[0], convert_indices_to_expr_handle(indices)); + auto non_softmax_indices = remove_softmax_dim_index(indices); + return inp - max->load(non_softmax_indices) - + log_sum->load(non_softmax_indices); + }); + return new Tensor( + result->buf(), + new tensorexpr::Block( + {max->stmt(), + e->stmt(), + sum->stmt(), + log_sum->stmt(), + result->stmt()})); +} + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/softmax.h b/torch/csrc/jit/tensorexpr/operators/softmax.h new file mode 100644 index 0000000000000..07ddd0f95b355 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/softmax.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeSoftmax( + const std::vector& inputs, + const std::vector& outputShape, + bool log_softmax); + +} // namespace tensorexpr +} // namespace jit +} // namespace torch From 842a831f53bc9305c0553f3b2096afc6c1c6efe4 Mon Sep 17 00:00:00 2001 From: Bert Maher Date: Wed, 16 Jun 2021 05:07:43 -0700 Subject: [PATCH 144/305] [nnc] Move batchnorm to operators library (#59992) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59992 Wrapped batch norm in function `computeBatchNorm`. ghstack-source-id: 131407851 Test Plan: CI Reviewed By: ZolotukhinM Differential Revision: D29116661 fbshipit-source-id: 2873a9a3e70f31db1988787160fc96c388ea3d4a --- tools/build_variables.bzl | 1 + torch/csrc/jit/tensorexpr/kernel.cpp | 234 +++++++----------- torch/csrc/jit/tensorexpr/kernel.h | 29 ++- torch/csrc/jit/tensorexpr/operators/norm.cpp | 75 ++++++ torch/csrc/jit/tensorexpr/operators/norm.h | 16 ++ .../csrc/jit/tensorexpr/operators/operators.h | 1 + 6 files changed, 201 insertions(+), 155 deletions(-) create mode 100644 torch/csrc/jit/tensorexpr/operators/norm.cpp create mode 100644 torch/csrc/jit/tensorexpr/operators/norm.h diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index e7421d8a292e8..06acafd645eab 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -292,6 +292,7 @@ core_sources_full_mobile = [ "torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp", "torch/csrc/jit/tensorexpr/operators/conv2d.cpp", "torch/csrc/jit/tensorexpr/operators/matmul.cpp", + "torch/csrc/jit/tensorexpr/operators/norm.cpp", "torch/csrc/jit/tensorexpr/operators/reduction.cpp", "torch/csrc/jit/tensorexpr/operators/softmax.cpp", "torch/csrc/jit/tensorexpr/reduction.cpp", diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 482c2560f7ae8..9eb2cc7e27706 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -16,6 +16,47 @@ using namespace torch::jit; using namespace torch::jit::tensorexpr; +namespace { + +static bool checkTypes(const ScalarType highType, const int typeConstraints) { + if (typeConstraints == kAllTypes) { + return true; + } + + if (c10::isIntegralType(highType, false)) { + return (typeConstraints & kIntegralTypes) != 0; + } else if (c10::isFloatingType(highType)) { + return (typeConstraints & kFloatingPointTypes) != 0; + } else if (highType == ScalarType::Bool) { + return (typeConstraints & kBoolType) != 0; + } + + // assume JIT not supporting complex and qint yet + TORCH_INTERNAL_ASSERT((typeConstraints & (kQintTypes | kComplexTypes)) == 0); + return false; +} + +static ExprHandle promoteToDtype(ExprHandle e, ScalarType dt) { + if (e.dtype().scalar_type() == dt) { + return e; + } + + switch (dt) { +// NOLINTNEXTLINE +#define TYPE_CASE(Type, Name) \ + case ScalarType::Name: \ + e = cast(e); \ + break; + AT_FORALL_SCALAR_TYPES_AND2(Half, Bool, TYPE_CASE); +#undef TYPE_CASE + default: + throw unsupported_dtype(); + } + return e; +} + +} // namespace + namespace torch { namespace jit { namespace tensorexpr { @@ -411,6 +452,52 @@ std::vector computeIndicesToBroadcast( return bcast; } +void promoteInputs(std::vector& inputs, const int typeConstraints) { + if (inputs.empty()) { + return; + } + + // Find the highest type among the inputs. + ScalarType highType = inputs[0].dtype().scalar_type(); + for (const auto input : inputs) { + highType = promoteTypes(highType, input.dtype().scalar_type()); + } + + if (!checkTypes(highType, typeConstraints)) { + throw unsupported_dtype(); + } + + for (ExprHandle& e : inputs) { + e = promoteToDtype(e, highType); + } +} + +ExprHandle demoteOutput( + const ExprHandle& e, + const c10::optional type) { + if (!type.has_value()) { + return e; + } + if (*type == e.dtype().scalar_type()) { + return e; + } + + switch (*type) { +// NOLINTNEXTLINE +#define TYPE_CASE(Type, Name) \ + case ScalarType::Name: \ + return cast(e); + AT_FORALL_SCALAR_TYPES_AND(Half, TYPE_CASE); +#undef TYPE_CASE + case ScalarType::Bool: + return cast(e); + default: + throw unsupported_dtype(); + } + + return e; +} + } // namespace tensorexpr } // namespace jit } // namespace torch @@ -449,25 +536,6 @@ ExprHandle TensorExprKernel::chunk( return BufHandle(b).load(indices); } -ExprHandle promoteToDtype(ExprHandle e, ScalarType dt) { - if (e.dtype().scalar_type() == dt) { - return e; - } - - switch (dt) { -// NOLINTNEXTLINE -#define TYPE_CASE(Type, Name) \ - case ScalarType::Name: \ - e = cast(e); \ - break; - AT_FORALL_SCALAR_TYPES_AND2(Half, Bool, TYPE_CASE); -#undef TYPE_CASE - default: - throw unsupported_dtype(); - } - return e; -} - ExprHandle TensorExprKernel::constant(const torch::jit::Value* v) { if (v->node()->kind() == prim::Constant) { const auto val = toIValue(v).value(); @@ -804,72 +872,6 @@ ExprHandle clamp( return CompareSelect::make(mm, cmax, cmax, mm, kGT); } -bool checkTypes(const ScalarType highType, const int typeConstraints) { - if (typeConstraints == kAllTypes) { - return true; - } - - if (c10::isIntegralType(highType, false)) { - return (typeConstraints & kIntegralTypes) != 0; - } else if (c10::isFloatingType(highType)) { - return (typeConstraints & kFloatingPointTypes) != 0; - } else if (highType == ScalarType::Bool) { - return (typeConstraints & kBoolType) != 0; - } - - // assume JIT not supporting complex and qint yet - TORCH_INTERNAL_ASSERT((typeConstraints & (kQintTypes | kComplexTypes)) == 0); - return false; -} - -void promoteInputs( - std::vector& inputs, - const int typeConstraints = kAllTypes) { - if (inputs.empty()) { - return; - } - - // Find the highest type among the inputs. - ScalarType highType = inputs[0].dtype().scalar_type(); - for (const auto input : inputs) { - highType = promoteTypes(highType, input.dtype().scalar_type()); - } - - if (!checkTypes(highType, typeConstraints)) { - throw unsupported_dtype(); - } - - for (ExprHandle& e : inputs) { - e = promoteToDtype(e, highType); - } -} - -ExprHandle demoteOutput( - const ExprHandle& e, - const c10::optional type) { - if (!type.has_value()) { - return e; - } - if (*type == e.dtype().scalar_type()) { - return e; - } - - switch (*type) { -// NOLINTNEXTLINE -#define TYPE_CASE(Type, Name) \ - case ScalarType::Name: \ - return cast(e); - AT_FORALL_SCALAR_TYPES_AND(Half, TYPE_CASE); -#undef TYPE_CASE - case ScalarType::Bool: - return cast(e); - default: - throw unsupported_dtype(); - } - - return e; -} - static bool isOne(ExprHandle e) { auto const& n = e.AsNode(); if (!n) { @@ -1671,67 +1673,9 @@ Tensor* tensorexpr::computeOperandValue( } break; case aten::batch_norm: { - bool hasWeight = true; - bool hasBias = true; - - if (c10::get_if(&inputs[1])) { - hasWeight = false; - } - - if (c10::get_if(&inputs[2])) { - hasBias = false; - } - - return Compute( - "aten_batch_norm", - c10::fmap(outputShape), - [&](const std::vector& axes) { - TORCH_INTERNAL_ASSERT(axes.size() >= 2); - // axes: N, C, H, W - std::vector indices(axes.begin(), axes.end()); - ExprHandle c = indices[1]; - - // Parameter list: - // input, weight, bias, mean, var, training, momentum, eps, - // cudnn_enabled - std::vector exprInputs = { - tensorOrConstant(inputs[0], indices), // input - tensorOrConstant(inputs[3], {c}), // mean - tensorOrConstant(inputs[4], {c}), // var - constant(inputs[7]) // eps - }; - - if (hasWeight) { - exprInputs.push_back(tensorOrConstant(inputs[1], {c})); - } - if (hasBias) { - exprInputs.push_back(tensorOrConstant(inputs[2], {c})); - } - promoteInputs(exprInputs); - - ExprHandle input = exprInputs[0]; - ExprHandle mean = exprInputs[1]; - ExprHandle var = exprInputs[2]; - ExprHandle eps = exprInputs[3]; - ExprHandle weight = FloatImm::make(1); - ExprHandle bias = FloatImm::make(0); - - if (hasWeight) { - weight = exprInputs[4]; - } - // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks) - if (hasBias) { - bias = exprInputs[5]; - } + return computeBatchNorm(inputs, outputShape, outputType); + } - // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks) - auto inv_var = rsqrt(var + eps); - auto alpha = inv_var * weight; - auto beta = bias - mean * alpha; - auto output = input * alpha + beta; - return demoteOutput(output, outputType); - }); - } break; case aten::log: { return computeOneOperand( "aten_log", inputs, outputShape, outputType, [](const ExprHandle& a) { diff --git a/torch/csrc/jit/tensorexpr/kernel.h b/torch/csrc/jit/tensorexpr/kernel.h index 88333e0c2c756..b85b536d88f41 100644 --- a/torch/csrc/jit/tensorexpr/kernel.h +++ b/torch/csrc/jit/tensorexpr/kernel.h @@ -23,6 +23,17 @@ inline std::vector bufferSizes(const T& t) { } return sizes; } + +enum ElementType { + kAllTypes = 0, + kIntegralTypes = 1 << 0, + kFloatingPointTypes = 1 << 1, + kBoolType = 1 << 2, + kComplexTypes = 1 << 3, + kQintTypes = 1 << 4, + kNonComplexOrQintTypes = kIntegralTypes | kBoolType | kFloatingPointTypes, +}; + using ArgNone = c10::monostate; using BufList = std::vector; using IntList = std::vector; @@ -55,6 +66,14 @@ std::vector computeIndicesToBroadcast( const std::vector& outputAxes, const std::vector& inputSizes); +void promoteInputs( + std::vector& inputs, + const int typeConstraints = kAllTypes); + +ExprHandle demoteOutput( + const ExprHandle& e, + const c10::optional type); + inline std::string getArgValueName(const ArgValue& a) { if (c10::get_if(&a)) { return "BufHandle"; @@ -98,16 +117,6 @@ struct TensorInfo { c10::ScalarType dtype; }; -enum ElementType { - kAllTypes = 0, - kIntegralTypes = 1 << 0, - kFloatingPointTypes = 1 << 1, - kBoolType = 1 << 2, - kComplexTypes = 1 << 3, - kQintTypes = 1 << 4, - kNonComplexOrQintTypes = kIntegralTypes | kBoolType | kFloatingPointTypes, -}; - TORCH_API Tensor* computeOperandValue( c10::Symbol op, const std::vector& inputs, diff --git a/torch/csrc/jit/tensorexpr/operators/norm.cpp b/torch/csrc/jit/tensorexpr/operators/norm.cpp new file mode 100644 index 0000000000000..d96ebcd9447db --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/norm.cpp @@ -0,0 +1,75 @@ +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeBatchNorm( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType) { + bool hasWeight = true; + bool hasBias = true; + + if (c10::get_if(&inputs[1])) { + hasWeight = false; + } + + if (c10::get_if(&inputs[2])) { + hasBias = false; + } + + return Compute( + "aten_batch_norm", + c10::fmap(outputShape), + [&](const std::vector& axes) { + TORCH_INTERNAL_ASSERT(axes.size() >= 2); + // axes: N, C, H, W + std::vector indices(axes.begin(), axes.end()); + ExprHandle c = indices[1]; + + // Parameter list: + // input, weight, bias, mean, var, training, momentum, eps, + // cudnn_enabled + std::vector exprInputs = { + tensorOrConstant(inputs[0], indices), // input + tensorOrConstant(inputs[3], {c}), // mean + tensorOrConstant(inputs[4], {c}), // var + constant(inputs[7]) // eps + }; + + if (hasWeight) { + exprInputs.push_back(tensorOrConstant(inputs[1], {c})); + } + if (hasBias) { + exprInputs.push_back(tensorOrConstant(inputs[2], {c})); + } + promoteInputs(exprInputs); + + ExprHandle input = exprInputs[0]; + ExprHandle mean = exprInputs[1]; + ExprHandle var = exprInputs[2]; + ExprHandle eps = exprInputs[3]; + ExprHandle weight = FloatImm::make(1); + ExprHandle bias = FloatImm::make(0); + + if (hasWeight) { + weight = exprInputs[4]; + } + // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks) + if (hasBias) { + bias = exprInputs[5]; + } + + // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks) + auto inv_var = rsqrt(var + eps); + auto alpha = inv_var * weight; + auto beta = bias - mean * alpha; + auto output = input * alpha + beta; + return demoteOutput(output, outputType); + }); +} + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/norm.h b/torch/csrc/jit/tensorexpr/operators/norm.h new file mode 100644 index 0000000000000..98d53b4c306e3 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/operators/norm.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +Tensor* computeBatchNorm( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType); + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/operators.h b/torch/csrc/jit/tensorexpr/operators/operators.h index 05705e2d4e89c..d94c9589e5e0e 100644 --- a/torch/csrc/jit/tensorexpr/operators/operators.h +++ b/torch/csrc/jit/tensorexpr/operators/operators.h @@ -2,5 +2,6 @@ #include #include +#include #include #include From a6ecfb32961b920fc2da0030a1dc5ed2aa9e2b39 Mon Sep 17 00:00:00 2001 From: Elton Leander Pinto Date: Wed, 16 Jun 2021 08:44:25 -0700 Subject: [PATCH 145/305] Update lint.yml to use custom clang-tidy build (#59967) Summary: Related: https://github.com/pytorch/pytorch/issues/59815 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59967 Reviewed By: samestep Differential Revision: D29164686 Pulled By: 1ntEgr8 fbshipit-source-id: b6f9fb6fa4280f757a54a37b30b027b7504bef63 --- .github/workflows/lint.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 550f8abb98da8..917a8c70e0b0e 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -270,7 +270,7 @@ jobs: runs-on: ubuntu-18.04 container: # ubuntu18.04-cuda10.2-py3.6-tidy11 - image: ghcr.io/pytorch/cilint-clang-tidy:e2cfc57ce4fa3a257a4b78fdfdc2b065c167b9c5 + image: ghcr.io/pytorch/cilint-clang-tidy:52a8ad78d49fc9f40241fee7988db48c920499df steps: - name: Checkout PyTorch uses: actions/checkout@v2 From 241aac3ef81358577df40a38348c6a5744bed317 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Wed, 16 Jun 2021 08:56:53 -0700 Subject: [PATCH 146/305] Add GITHUB_HEAD_REF in check for IN_PULL_REQUEST (#60047) Summary: I believe IN_PULL_REQUEST is unset for some GHA test runs because we don't also check GITHUB_HEAD_REF. This PR is a small fix for that. Example: https://github.com/pytorch/pytorch/pull/60023/checks?check_run_id=2831813860 doesn't set it properly Pull Request resolved: https://github.com/pytorch/pytorch/pull/60047 Reviewed By: walterddr Differential Revision: D29148233 Pulled By: janeyx99 fbshipit-source-id: 7c8c1866f39ce8af8d13c34ddc0c5786a829321e --- .jenkins/pytorch/test.sh | 2 +- .jenkins/pytorch/win-test.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 7a1ca1cba8b95..7cfd9d34d3689 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -125,7 +125,7 @@ fi # Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second # CIRCLE_PULL_REQUEST comes from CircleCI # GITHUB_HEAD_REF comes from Github Actions -IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} +IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-${GITHUB_HEAD_REF:-}} if [ -n "$IN_PULL_REQUEST" ] && [[ "$BUILD_ENVIRONMENT" != *coverage* ]]; then DETERMINE_FROM=$(mktemp) file_diff_from_base "$DETERMINE_FROM" diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index 96fe2e6225be4..c49d60b2443f2 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -43,7 +43,7 @@ fi export SCRIPT_HELPERS_DIR=$SCRIPT_PARENT_DIR/win-test-helpers # Try to pull value from CIRCLE_PULL_REQUEST -IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} +IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-${GITHUB_HEAD_REF:-}} if [ -n "$IN_PULL_REQUEST" ]; then DETERMINE_FROM="${TMP_DIR}/determine_from" file_diff_from_base "$DETERMINE_FROM" From d88fbf0fbc69ec55eb407822dbd174a2d8a9ccd2 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Wed, 16 Jun 2021 09:29:48 -0700 Subject: [PATCH 147/305] fix minor typo in run_test.py (#60055) Summary: Fixes typo in run_test.py for option use_specified_test_cases_by Pull Request resolved: https://github.com/pytorch/pytorch/pull/60055 Reviewed By: walterddr Differential Revision: D29150156 Pulled By: janeyx99 fbshipit-source-id: 375e594d09c83188bfa80762c8b833a0b7c5cca4 --- test/run_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/run_test.py b/test/run_test.py index 5670da354ce01..a0e73b9c86ed1 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -915,9 +915,9 @@ def exclude_tests(exclude_list, selected_tests, exclude_message=None): def get_selected_tests(options): if options.run_specified_test_cases: - if options.use_specified_test_cases_for == 'include': + if options.use_specified_test_cases_by == 'include': options.include = list(SPECIFIED_TEST_CASES_DICT.keys()) - elif options.use_specified_test_cases_for == 'bring-to-front': + elif options.use_specified_test_cases_by == 'bring-to-front': options.bring_to_front = list(SPECIFIED_TEST_CASES_DICT.keys()) selected_tests = options.include From bac6bcd6d8658a01ecf356453d6a29010867f31d Mon Sep 17 00:00:00 2001 From: Cao Gao Date: Wed, 16 Jun 2021 10:14:18 -0700 Subject: [PATCH 148/305] Update call site for FBGemm quantization util functions. (#624) Summary: Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/624 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59637 Replace FloatToFusedNBitRowwiseQuantizedSBHalf, FusedNBitRowwiseQuantizedSBHalfToFloat, FloatToFused8BitRowwiseQuantizedSBFloat, and Fused8BitRowwiseQuantizedSBFloatToFloat with newer version. Test Plan: CI tests. Reviewed By: dskhudia Differential Revision: D28918581 fbshipit-source-id: a21274add71439c5e51287a0e2ec918a8d8e5392 --- .../ATen/native/quantized/cpu/qembeddingbag_prepack.cpp | 4 ++-- .../ATen/native/quantized/cpu/qembeddingbag_unpack.cpp | 4 ++-- caffe2/perfkernels/fused_nbit_rowwise_conversion.cc | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/aten/src/ATen/native/quantized/cpu/qembeddingbag_prepack.cpp b/aten/src/ATen/native/quantized/cpu/qembeddingbag_prepack.cpp index 106908c94b186..ed41b851ab1a9 100644 --- a/aten/src/ATen/native/quantized/cpu/qembeddingbag_prepack.cpp +++ b/aten/src/ATen/native/quantized/cpu/qembeddingbag_prepack.cpp @@ -222,7 +222,7 @@ Tensor qembeddingbag_byte_prepack(const Tensor& weight) { at::parallel_for( 0, embedding_rows, 1, [&](int32_t start_idx, int32_t end_idx) { for (int64_t row = start_idx; row < end_idx; ++row) { - fbgemm::FloatToFused8BitRowwiseQuantizedSBFloat( + fbgemm::FloatOrHalfToFused8BitRowwiseQuantizedSBFloat( weight_data + row * embedding_cols, 1, embedding_cols, output_data + row * output_columns); } @@ -302,7 +302,7 @@ Tensor _qembeddingbag_nbit_prepack_helper( at::parallel_for( 0, embedding_rows, 1, [&](int32_t start_idx, int32_t end_idx) { for (int64_t row = start_idx; row < end_idx; ++row) { - fbgemm::FloatToFusedNBitRowwiseQuantizedSBHalf( + fbgemm::FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf( bit_width, weight_data + row * embedding_cols, 1, embedding_cols, output_data + row * output_shape[1]); } diff --git a/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp b/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp index f74d3b7ac8431..0db8905a1844e 100644 --- a/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp +++ b/aten/src/ATen/native/quantized/cpu/qembeddingbag_unpack.cpp @@ -128,7 +128,7 @@ Tensor qembeddingbag_byte_unpack(const Tensor& packed_weight) { at::parallel_for( 0, input_rows, 1, [&](int32_t start_idx, int32_t end_idx) { for (int64_t row = start_idx; row < end_idx; ++row) { - fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloat( + fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf( input_data + row * input_columns, 1, input_columns, @@ -175,7 +175,7 @@ Tensor _qembeddingbag_nbit_unpack_helper( at::parallel_for( 0, input_rows, 1, [&](int32_t start_idx, int32_t end_idx) { for (int64_t row = start_idx; row < end_idx; ++row) { - fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloat(BIT_RATE, + fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloatOrHalf(BIT_RATE, input_data + row * input_columns, 1, input_columns, diff --git a/caffe2/perfkernels/fused_nbit_rowwise_conversion.cc b/caffe2/perfkernels/fused_nbit_rowwise_conversion.cc index 461827a3cae14..446f02873f9b1 100644 --- a/caffe2/perfkernels/fused_nbit_rowwise_conversion.cc +++ b/caffe2/perfkernels/fused_nbit_rowwise_conversion.cc @@ -75,7 +75,7 @@ void FloatToFused8BitRowwiseQuantized( int input_columns, std::uint8_t* output) { #ifdef USE_FBGEMM - fbgemm::FloatToFused8BitRowwiseQuantizedSBFloat( + fbgemm::FloatOrHalfToFused8BitRowwiseQuantizedSBFloat( input, input_rows, input_columns, output); #else FloatToFused8BitRowwiseQuantized__base( @@ -89,7 +89,7 @@ void Fused8BitRowwiseQuantizedToFloat( int input_columns, float* output) { #ifdef USE_FBGEMM - fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloat( + fbgemm::Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf( input, input_rows, input_columns, output); #else Fused8BitRowwiseQuantizedToFloat__base( @@ -196,7 +196,7 @@ void FloatToFusedNBitRowwiseQuantizedSBHalf( int input_columns, std::uint8_t* output) { #ifdef USE_FBGEMM - fbgemm::FloatToFusedNBitRowwiseQuantizedSBHalf( + fbgemm::FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf( bit_rate, input, input_rows, input_columns, output); #else FloatToFusedNBitRowwiseQuantizedSBHalf__base( @@ -211,7 +211,7 @@ void FusedNBitRowwiseQuantizedSBHalfToFloat( int input_columns, float* output) { #ifdef USE_FBGEMM - fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloat( + fbgemm::FusedNBitRowwiseQuantizedSBHalfToFloatOrHalf( bit_rate, input, input_rows, input_columns, output); #else FusedNBitRowwiseQuantizedSBHalfToFloat__base( From 74ea1f23b46f5defd71cdb24c3be1c6da9fd7594 Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Wed, 16 Jun 2021 10:40:28 -0700 Subject: [PATCH 149/305] Revert D29148233: [pytorch][PR] Add GITHUB_HEAD_REF in check for IN_PULL_REQUEST Test Plan: revert-hammer Differential Revision: D29148233 (https://github.com/pytorch/pytorch/commit/241aac3ef81358577df40a38348c6a5744bed317) Original commit changeset: 7c8c1866f39c fbshipit-source-id: f32c6c6decd737ef290d3e83c9d021475aabaab0 --- .jenkins/pytorch/test.sh | 2 +- .jenkins/pytorch/win-test.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 7cfd9d34d3689..7a1ca1cba8b95 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -125,7 +125,7 @@ fi # Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second # CIRCLE_PULL_REQUEST comes from CircleCI # GITHUB_HEAD_REF comes from Github Actions -IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-${GITHUB_HEAD_REF:-}} +IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} if [ -n "$IN_PULL_REQUEST" ] && [[ "$BUILD_ENVIRONMENT" != *coverage* ]]; then DETERMINE_FROM=$(mktemp) file_diff_from_base "$DETERMINE_FROM" diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index c49d60b2443f2..96fe2e6225be4 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -43,7 +43,7 @@ fi export SCRIPT_HELPERS_DIR=$SCRIPT_PARENT_DIR/win-test-helpers # Try to pull value from CIRCLE_PULL_REQUEST -IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-${GITHUB_HEAD_REF:-}} +IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} if [ -n "$IN_PULL_REQUEST" ]; then DETERMINE_FROM="${TMP_DIR}/determine_from" file_diff_from_base "$DETERMINE_FROM" From 8c4e78129ec8d71d587ac5d143ad17e4b95b3576 Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Wed, 16 Jun 2021 10:43:47 -0700 Subject: [PATCH 150/305] .circleci: Disable Windows GPU jobs (#60024) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60024 Disables windows GPU jobs on CircleCI since they have been migrated to GHA Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: malfet Differential Revision: D29137287 Pulled By: seemethere fbshipit-source-id: 204e0c9232201a36a557cd0843e31d34269cc722 --- .../cimodel/data/windows_build_definitions.py | 6 - .circleci/config.yml | 146 ------------------ 2 files changed, 152 deletions(-) diff --git a/.circleci/cimodel/data/windows_build_definitions.py b/.circleci/cimodel/data/windows_build_definitions.py index c6b680a1e0472..173e669cc68d5 100644 --- a/.circleci/cimodel/data/windows_build_definitions.py +++ b/.circleci/cimodel/data/windows_build_definitions.py @@ -147,14 +147,8 @@ def render(self): WORKFLOW_DATA = [ # VS2019 CUDA-10.1 WindowsJob(None, _VC2019, CudaVersion(10, 1), master_only=True), - WindowsJob(1, _VC2019, CudaVersion(10, 1), master_only=True), - WindowsJob(2, _VC2019, CudaVersion(10, 1), master_only=True), # VS2019 CUDA-10.1 force on cpu WindowsJob(1, _VC2019, CudaVersion(10, 1), force_on_cpu=True, master_only=True), - # VS2019 CUDA-11.1 - WindowsJob(None, _VC2019, CudaVersion(11, 1)), - WindowsJob(1, _VC2019, CudaVersion(11, 1), master_only=True), - WindowsJob(2, _VC2019, CudaVersion(11, 1), master_only=True), # TODO: This test is disabled due to https://github.com/pytorch/pytorch/issues/59724 # WindowsJob('_azure_multi_gpu', _VC2019, CudaVersion(11, 1), multi_gpu=True, master_and_nightly=True), diff --git a/.circleci/config.yml b/.circleci/config.yml index a3df5be23b111..8c3060b43646e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -7637,44 +7637,6 @@ workflows: vc_product: BuildTools vc_version: "" vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 - cuda_version: "10.1" - executor: windows-with-nvidia-gpu - filters: - branches: - only: - - master - - /ci-all\/.*/ - - /release\/.*/ - name: pytorch_windows_vs2019_py38_cuda10.1_test1 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda10.1_build - test_name: pytorch-windows-test1 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 - cuda_version: "10.1" - executor: windows-with-nvidia-gpu - filters: - branches: - only: - - master - - /ci-all\/.*/ - - /release\/.*/ - name: pytorch_windows_vs2019_py38_cuda10.1_test2 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda10.1_build - test_name: pytorch-windows-test2 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - pytorch_windows_test: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" @@ -7693,53 +7655,6 @@ workflows: vc_product: BuildTools vc_version: "" vc_year: "2019" - - pytorch_windows_build: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - name: pytorch_windows_vs2019_py38_cuda11.1_build - python_version: "3.8" - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - executor: windows-with-nvidia-gpu - filters: - branches: - only: - - master - - /ci-all\/.*/ - - /release\/.*/ - name: pytorch_windows_vs2019_py38_cuda11.1_test1 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda11.1_build - test_name: pytorch-windows-test1 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - executor: windows-with-nvidia-gpu - filters: - branches: - only: - - master - - /ci-all\/.*/ - - /release\/.*/ - name: pytorch_windows_vs2019_py38_cuda11.1_test2 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda11.1_build - test_name: pytorch-windows-test2 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - update_s3_htmls: context: org-member filters: @@ -9278,32 +9193,6 @@ workflows: vc_product: BuildTools vc_version: "" vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 - cuda_version: "10.1" - executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py38_cuda10.1_test1 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda10.1_build - test_name: pytorch-windows-test1 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 - cuda_version: "10.1" - executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py38_cuda10.1_test2 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda10.1_build - test_name: pytorch-windows-test2 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - pytorch_windows_test: build_environment: pytorch-win-vs2019-cuda10-cudnn7-py3 cuda_version: "10.1" @@ -9316,41 +9205,6 @@ workflows: vc_product: BuildTools vc_version: "" vc_year: "2019" - - pytorch_windows_build: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - name: pytorch_windows_vs2019_py38_cuda11.1_build - python_version: "3.8" - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py38_cuda11.1_test1 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda11.1_build - test_name: pytorch-windows-test1 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" - - pytorch_windows_test: - build_environment: pytorch-win-vs2019-cuda11-cudnn8-py3 - cuda_version: "11.1" - executor: windows-with-nvidia-gpu - name: pytorch_windows_vs2019_py38_cuda11.1_test2 - python_version: "3.8" - requires: - - pytorch_windows_vs2019_py38_cuda11.1_build - test_name: pytorch-windows-test2 - use_cuda: "1" - vc_product: BuildTools - vc_version: "" - vc_year: "2019" when: << pipeline.parameters.run_master_build >> slow_gradcheck_build: jobs: From 9f68f93aca0f4cc9726ef697db13715f699fb234 Mon Sep 17 00:00:00 2001 From: Zhuangzhuang Zhang Date: Wed, 16 Jun 2021 11:21:16 -0700 Subject: [PATCH 151/305] Training resnext with msuru_suru_union and ig_msuru_suru_union datasets Summary: We updated the training scripts and re-trained the Resnext model with msuru_suru_union and ig_msuru_suru_union datasets Test Plan: Main command line to run: *./deeplearning/projects/classy_vision/fb/projects/msuru_suru/scripts/train_cluster.sh* Config we used is *msuru_suru_config.json*, which is "Normal ResNeXt101 with finetunable head". Experiments: - msuru_suru_union f279939874 - Train/test split - msuru_suru_union_dataset_train_w_shard: 143,632,674 rows - msuru_suru_union_dataset_test_w_shard: 1,831,236 rows - Results {F625232741} {F625232819} - ig_msuru_suru_union f279964200 - Train/test split - ig_msuru_suru_union_dataset_train_w_shard: 241,884,760 rows - ig_msuru_suru_union_dataset_test_w_shard: 3,477,181 rows - Results {F625234126} {F625234457} Differential Revision: D29154971 fbshipit-source-id: d534d830020f4f8e596bb6b941966eb84a1e8adb --- torch/nn/parallel/distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py index b46cae2e742d3..2285a07f329cd 100644 --- a/torch/nn/parallel/distributed.py +++ b/torch/nn/parallel/distributed.py @@ -406,7 +406,7 @@ def __init__( broadcast_buffers=True, process_group=None, bucket_cap_mb=25, - find_unused_parameters=False, + find_unused_parameters=True, check_reduction=False, gradient_as_bucket_view=False, ): From 15f236f3e35d8ca4099336aa34097b210b24bb5b Mon Sep 17 00:00:00 2001 From: Michael Suo Date: Wed, 16 Jun 2021 11:22:35 -0700 Subject: [PATCH 152/305] [package] fix tutorial link (#60113) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60113 The tutorial link in the docs was to an fb-only colab. Test Plan: Imported from OSS Reviewed By: SplitInfinity Differential Revision: D29169818 Pulled By: suo fbshipit-source-id: 374807c234a185bd515b8ffe1300e6cf8d821636 --- docs/source/package.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/package.rst b/docs/source/package.rst index d21f7b016efce..ea303037cab52 100644 --- a/docs/source/package.rst +++ b/docs/source/package.rst @@ -30,7 +30,7 @@ Tutorials Packaging your first model ^^^^^^^^^^^^^^^^^^^^^^^^^^ A tutorial that guides you through packaging and unpackaging a simple model is available -`on Colab `_. +`on Colab `_. After completing this exercise, you will be familiar with the basic API for creating and using Torch packages. From 84688b0c40abf23cf4d6f2f03f285e8ad1bc4da5 Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Wed, 16 Jun 2021 11:29:57 -0700 Subject: [PATCH 153/305] ci: Add note about file_diff_from_base for GHA (#60110) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60110 file_diff_from_base is currently bugged for ghstack PRs since it fails to find a merge base Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: driazati Differential Revision: D29168767 Pulled By: seemethere fbshipit-source-id: 580a909aa392541769cbbfdc6acce1e6c5d1c341 --- .jenkins/pytorch/macos-test.sh | 3 ++- .jenkins/pytorch/test.sh | 3 ++- .jenkins/pytorch/win-test.sh | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.jenkins/pytorch/macos-test.sh b/.jenkins/pytorch/macos-test.sh index d1ae02525847b..0ee446c4d25fc 100755 --- a/.jenkins/pytorch/macos-test.sh +++ b/.jenkins/pytorch/macos-test.sh @@ -53,7 +53,8 @@ test_python_all() { # Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second # CIRCLE_PULL_REQUEST comes from CircleCI - # GITHUB_HEAD_REF comes from Github Actions + # NOTE: file_diff_from_base is currently bugged for GHA due to an issue finding a merge base for ghstack PRs + # see https://github.com/pytorch/pytorch/issues/60111 IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-${GITHUB_HEAD_REF:-}} if [ -n "$IN_PULL_REQUEST" ]; then DETERMINE_FROM=$(mktemp) diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 7a1ca1cba8b95..6610fe03aafe9 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -124,7 +124,8 @@ fi # Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second # CIRCLE_PULL_REQUEST comes from CircleCI -# GITHUB_HEAD_REF comes from Github Actions +# NOTE: file_diff_from_base is currently bugged for GHA due to an issue finding a merge base for ghstack PRs +# see https://github.com/pytorch/pytorch/issues/60111 IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} if [ -n "$IN_PULL_REQUEST" ] && [[ "$BUILD_ENVIRONMENT" != *coverage* ]]; then DETERMINE_FROM=$(mktemp) diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index 96fe2e6225be4..736886b105c37 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -43,6 +43,8 @@ fi export SCRIPT_HELPERS_DIR=$SCRIPT_PARENT_DIR/win-test-helpers # Try to pull value from CIRCLE_PULL_REQUEST +# NOTE: file_diff_from_base is currently bugged for GHA due to an issue finding a merge base for ghstack PRs +# see https://github.com/pytorch/pytorch/issues/60111 IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} if [ -n "$IN_PULL_REQUEST" ]; then DETERMINE_FROM="${TMP_DIR}/determine_from" From 2c5db9a40a881de7f683469d12dd10a95817645c Mon Sep 17 00:00:00 2001 From: Neel Pragnesh Gandhi Date: Wed, 16 Jun 2021 12:12:15 -0700 Subject: [PATCH 154/305] Add c10d filestore functionality to the current c10d_rendezvous_backend (#59719) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59719 Added filestore functionality to the c10d backend. FileStore will create a temporary file in the /tmp directory to use if it is selected as the store type. Appropriate tests were added as well. FileStore was modified to expose the path field for testing. It was also modified so that the numWorkers field in the constructor is optional (defaulting to -1). A negative value indicates there is not a fixed number of workers. In this case, the file is not attempted to be cleaned at the end. Test Plan: Unit tests for creating a c10d backend with filestore and simple error handling. Reviewed By: cbalioglu, H-Huang Differential Revision: D28997436 fbshipit-source-id: 24c9b2c9b13ea6c947e8b1207beda892bdca2217 --- .../c10d_rendezvous_backend_test.py | 135 +++++++++++++++--- torch/_C/_distributed_c10d.pyi | 2 +- torch/csrc/distributed/c10d/init.cpp | 12 +- .../rendezvous/c10d_rendezvous_backend.py | 55 +++++-- torch/lib/c10d/FileStore.cpp | 10 +- torch/lib/c10d/FileStore.hpp | 5 + 6 files changed, 180 insertions(+), 39 deletions(-) diff --git a/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py b/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py index 5d1455ffd5a92..1190c9ee9eeb0 100644 --- a/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py +++ b/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py @@ -4,14 +4,20 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import os +import tempfile + from base64 import b64encode from datetime import timedelta -from typing import ClassVar, cast -from unittest import TestCase +from typing import ClassVar, cast, Callable +from unittest import TestCase, mock -from torch.distributed import TCPStore +from torch.distributed import TCPStore, FileStore -from torch.distributed.elastic.rendezvous import RendezvousConnectionError, RendezvousParameters +from torch.distributed.elastic.rendezvous import ( + RendezvousConnectionError, + RendezvousParameters, + RendezvousError) from torch.distributed.elastic.rendezvous.c10d_rendezvous_backend import ( C10dRendezvousBackend, create_backend, @@ -20,7 +26,7 @@ from rendezvous_backend_test import RendezvousBackendTestMixin -class C10dRendezvousBackendTest(TestCase, RendezvousBackendTestMixin): +class TCPStoreBackendTest(TestCase, RendezvousBackendTestMixin): _store: ClassVar[TCPStore] @classmethod @@ -36,9 +42,31 @@ def setUp(self) -> None: def _corrupt_state(self) -> None: self._store.set("torch.rendezvous.dummy_run_id", "non_base64") +class FileStoreBackendTest(TestCase, RendezvousBackendTestMixin): + _store: ClassVar[FileStore] + + def setUp(self) -> None: + _, path = tempfile.mkstemp() + self._path = path + + # Currently, filestore doesn't implement a delete_key method, so a new + # filestore has to be initialized for every test in order to have a + # clean slate. + self._store = FileStore(path) + self._backend = C10dRendezvousBackend(self._store, "dummy_run_id") + + def tearDown(self) -> None: + os.remove(self._path) + + def _corrupt_state(self) -> None: + self._store.set("torch.rendezvous.dummy_run_id", "non_base64") + class CreateBackendTest(TestCase): def setUp(self) -> None: + # For testing, the default parameters used are for tcp. If a test + # uses parameters for file store, we set the self._params to + # self._params_filestore. self._params = RendezvousParameters( backend="dummy_backend", endpoint="localhost:29300", @@ -50,23 +78,58 @@ def setUp(self) -> None: read_timeout="10", ) + _, tmp_path = tempfile.mkstemp() + + # Parameters for filestore testing. + self._params_filestore = RendezvousParameters( + backend="dummy_backend", + endpoint=tmp_path, + run_id="dummy_run_id", + min_nodes=1, + max_nodes=1, + store_type="fIlE", + ) + self._expected_endpoint_file = tmp_path + self._expected_temp_dir = tempfile.gettempdir() + self._expected_endpoint_host = "localhost" self._expected_endpoint_port = 29300 self._expected_store_type = TCPStore self._expected_read_timeout = timedelta(seconds=10) - def test_create_backend_returns_backend(self) -> None: + def tearDown(self) -> None: + os.remove(self._expected_endpoint_file) + + + def _run_test_with_store(self, store_type: str, test_to_run: Callable): + """ + Use this function to specify the store type to use in a test. If + not used, the test will default to TCPStore. + """ + if store_type == "file": + self._params = self._params_filestore + self._expected_store_type = FileStore + self._expected_read_timeout = timedelta(seconds=300) + + test_to_run() + + def _assert_create_backend_returns_backend(self) -> None: backend, store = create_backend(self._params) self.assertEqual(backend.name, "c10d") self.assertIsInstance(store, self._expected_store_type) - tcp_store = cast(TCPStore, store) - - self.assertEqual(tcp_store.host, self._expected_endpoint_host) # type: ignore[attr-defined] - self.assertEqual(tcp_store.port, self._expected_endpoint_port) # type: ignore[attr-defined] - self.assertEqual(tcp_store.timeout, self._expected_read_timeout) # type: ignore[attr-defined] + typecast_store = cast(self._expected_store_type, store) + self.assertEqual(typecast_store.timeout, self._expected_read_timeout) # type: ignore[attr-defined] + if (self._expected_store_type == TCPStore): + self.assertEqual(typecast_store.host, self._expected_endpoint_host) # type: ignore[attr-defined] + self.assertEqual(typecast_store.port, self._expected_endpoint_port) # type: ignore[attr-defined] + if (self._expected_store_type == FileStore): + if self._params.endpoint: + self.assertEqual(typecast_store.path, self._expected_endpoint_file) # type: ignore[attr-defined] + else: + self.assertTrue(typecast_store.path.startswith(self._expected_temp_dir)) # type: ignore[attr-defined] backend.set_state(b"dummy_state") @@ -74,6 +137,11 @@ def test_create_backend_returns_backend(self) -> None: self.assertEqual(state, b64encode(b"dummy_state")) + def test_create_backend_returns_backend(self) -> None: + for store_type in ["tcp", "file"]: + with self.subTest(store_type=store_type): + self._run_test_with_store(store_type, self._assert_create_backend_returns_backend) + def test_create_backend_returns_backend_if_is_host_is_false(self) -> None: store = TCPStore( # type: ignore[call-arg] # noqa: F841 self._expected_endpoint_host, self._expected_endpoint_port, is_master=True @@ -81,12 +149,12 @@ def test_create_backend_returns_backend_if_is_host_is_false(self) -> None: self._params.config["is_host"] = "false" - self.test_create_backend_returns_backend() + self._assert_create_backend_returns_backend() def test_create_backend_returns_backend_if_is_host_is_not_specified(self) -> None: del self._params.config["is_host"] - self.test_create_backend_returns_backend() + self._assert_create_backend_returns_backend() def test_create_backend_returns_backend_if_is_host_is_not_specified_and_store_already_exists( self, @@ -97,26 +165,35 @@ def test_create_backend_returns_backend_if_is_host_is_not_specified_and_store_al del self._params.config["is_host"] - self.test_create_backend_returns_backend() + self._assert_create_backend_returns_backend() def test_create_backend_returns_backend_if_endpoint_port_is_not_specified(self) -> None: self._params.endpoint = self._expected_endpoint_host self._expected_endpoint_port = 29400 - self.test_create_backend_returns_backend() + self._assert_create_backend_returns_backend() + + def test_create_backend_returns_backend_if_endpoint_file_is_not_specified(self) -> None: + self._params_filestore.endpoint = "" + + self._run_test_with_store("file", self._assert_create_backend_returns_backend) def test_create_backend_returns_backend_if_store_type_is_not_specified(self) -> None: del self._params.config["store_type"] - self.test_create_backend_returns_backend() + self._expected_store_type = TCPStore + if (not self._params.get("read_timeout")): + self._expected_read_timeout = timedelta(seconds=60) + + self._assert_create_backend_returns_backend() def test_create_backend_returns_backend_if_read_timeout_is_not_specified(self) -> None: del self._params.config["read_timeout"] self._expected_read_timeout = timedelta(seconds=60) - self.test_create_backend_returns_backend() + self._assert_create_backend_returns_backend() def test_create_backend_raises_error_if_store_is_unreachable(self) -> None: self._params.config["is_host"] = "false" @@ -146,7 +223,7 @@ def test_create_backend_raises_error_if_store_type_is_invalid(self) -> None: self._params.config["store_type"] = "dummy_store_type" with self.assertRaisesRegex( - ValueError, r"^The store type must be 'tcp'. Other store types are not supported yet.$" + ValueError, r"^Invalid store type given. Currently only supports file and tcp.$" ): create_backend(self._params) @@ -159,3 +236,25 @@ def test_create_backend_raises_error_if_read_timeout_is_invalid(self) -> None: ValueError, r"^The read timeout must be a positive integer.$" ): create_backend(self._params) + + @mock.patch("tempfile.mkstemp") + def test_create_backend_raises_error_if_tempfile_creation_fails(self, tempfile_mock) -> None: + tempfile_mock.side_effect = OSError("test error") + # Set the endpoint to empty so it defaults to creating a temp file + self._params_filestore.endpoint = "" + with self.assertRaisesRegex( + RendezvousError, + r"The file creation for C10d store has failed. See inner exception for details." + ): + create_backend(self._params_filestore) + + @mock.patch("torch.distributed.elastic.rendezvous.c10d_rendezvous_backend.FileStore") + def test_create_backend_raises_error_if_file_path_is_invalid(self, filestore_mock) -> None: + filestore_mock.side_effect = RuntimeError("test error") + self._params_filestore.endpoint = "bad file path" + with self.assertRaisesRegex( + RendezvousConnectionError, + r"^The connection to the C10d store has failed. See inner exception for " + r"details.$", + ): + create_backend(self._params_filestore) diff --git a/torch/_C/_distributed_c10d.pyi b/torch/_C/_distributed_c10d.pyi index a47a1e7bc0baa..40368bf6478a0 100644 --- a/torch/_C/_distributed_c10d.pyi +++ b/torch/_C/_distributed_c10d.pyi @@ -128,7 +128,7 @@ class Store: def wait(self, keys: List[str], timeout: timedelta): ... class FileStore(Store): - def __init__(self, path: str, numWorkers: int): ... + def __init__(self, path: str, numWorkers: int = ...): ... class HashStore(Store): def __init__(self): ... diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp index 1995d46ad7793..493d5a180af0f 100644 --- a/torch/csrc/distributed/c10d/init.cpp +++ b/torch/csrc/distributed/c10d/init.cpp @@ -807,7 +807,7 @@ A store implementation that uses a file to store the underlying key-value pairs. Arguments: file_name (str): path of the file in which to store the key-value pairs - world_size (int): The total number of processes using the store + world_size (int, optional): The total number of processes using the store. Default is -1 (a negative value indicates a non-fixed number of store users). Example:: >>> import torch.distributed as dist @@ -818,7 +818,13 @@ Example:: >>> store2.get("first_key") )") - .def(py::init()); + .def(py::init(), + py::arg("file_name"), + py::arg("world_size") = -1) + .def_property_readonly( + "path", + &::c10d::FileStore::getPath, + R"(Gets the path of the file used by FileStore to store key-value pairs.)"); #ifndef _WIN32 intrusive_ptr_class_<::c10d::HashStore>( @@ -854,7 +860,7 @@ the server to establish a connection. Arguments: host_name (str): The hostname or IP Address the server store should run on. port (int): The port on which the server store should listen for incoming requests. - world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates an non-fixed number of store users). + world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates a non-fixed number of store users). is_master (bool, optional): True when initializing the server store and False for client stores. Default is False. timeout (timedelta, optional): Timeout used by the store during initialization and for methods such as :meth:`~torch.distributed.store.get` and :meth:`~torch.distributed.store.wait`. Default is timedelta(seconds=300) wait_for_worker (bool, optional): Whether to wait for all the workers to connect with the server store. This is only applicable when world_size is a fixed value. Default is True. diff --git a/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py b/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py index ee3453edbda00..4db9b669c314b 100644 --- a/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py +++ b/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py @@ -7,13 +7,14 @@ import binascii import logging import os +import tempfile from base64 import b64decode, b64encode from datetime import timedelta from typing import Any, Optional, Tuple, cast -from torch.distributed import Store, TCPStore +from torch.distributed import Store, TCPStore, FileStore -from .api import RendezvousConnectionError, RendezvousParameters, RendezvousStateError +from .api import RendezvousConnectionError, RendezvousParameters, RendezvousStateError, RendezvousError from .dynamic_rendezvous import RendezvousBackend, Token from .utils import _matches_machine_hostname, parse_rendezvous_endpoint @@ -166,6 +167,28 @@ def _create_tcp_store(params: RendezvousParameters) -> TCPStore: return store +def _create_file_store(params: RendezvousParameters) -> FileStore: + # If a user specifies an endpoint, we treat it as a path to a file. + if params.endpoint: + path = params.endpoint + else: + try: + # The temporary file is readable and writable only by the user of + # this process. + _, path = tempfile.mkstemp() + except OSError as exc: + raise RendezvousError( + "The file creation for C10d store has failed. See inner exception for details." + ) from exc + + try: + store = FileStore(path) + except (ValueError, RuntimeError) as exc: + raise RendezvousConnectionError( + "The connection to the C10d store has failed. See inner exception for details." + ) from exc + + return store def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend, Store]: """Creates a new :py:class:`C10dRendezvousBackend` from the specified @@ -174,12 +197,19 @@ def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend, +--------------+-----------------------------------------------------------+ | Parameter | Description | +==============+===========================================================+ - | store_type | The type of the C10d store. As of today the only | - | | supported type is "tcp" which corresponds to | - | | :py:class:`torch.distributed.TCPStore`. Defaults to "tcp".| + | store_type | The type of the C10d store. The currently supported types | + | | are "tcp" and "file" which correspond to | + | | :py:class:`torch.distributed.TCPStore` and | + | | :py:class:`torch.distributed.FileStore`, respectively. | + | | Defaults to "tcp". | +--------------+-----------------------------------------------------------+ | read_timeout | The read timeout, in seconds, for store operations. | | | Defaults to 60 seconds. | + | | | + | | Note this only applies to | + | | :py:class:`torch.distributed.TCPStore`. It is not relevant| + | | to :py:class:`torch.distributed.FileStore` which does not | + | | take in timeout as a parameter. | +--------------+-----------------------------------------------------------+ | is_host | A boolean value indicating whether this backend instance | | | will host the C10d store. If not specified it will be | @@ -195,12 +225,15 @@ def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend, | | the hostname or does not match the FQDN of the machine). | +--------------+-----------------------------------------------------------+ """ - # As of today we only support TCPStore. Other store types do not have the - # required functionality (e.g. compare_set) yet. + # As of today we only support TCPStore and FileStore. Other store types do + # not have the required functionality (e.g. compare_set) yet. store_type = params.get("store_type", "tcp").strip().lower() - if store_type != "tcp": - raise ValueError("The store type must be 'tcp'. Other store types are not supported yet.") - - store = _create_tcp_store(params) + store: Store + if store_type == "file": + store = _create_file_store(params) + elif store_type == "tcp": + store = _create_tcp_store(params) + else: + raise ValueError("Invalid store type given. Currently only supports file and tcp.") return C10dRendezvousBackend(store, params.run_id), store diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp index 73342272c54c0..77fab9565d6c9 100644 --- a/torch/lib/c10d/FileStore.cpp +++ b/torch/lib/c10d/FileStore.cpp @@ -272,18 +272,16 @@ FileStore::FileStore(const std::string& path, int numWorkers) numWorkers_(numWorkers), cleanupKey_("cleanup/"), regularPrefix_("/") { - if (numWorkers_ < 1) { - TORCH_CHECK(false, - "Number of workers for FileStore should be greater than zero"); - } } FileStore::~FileStore() { // cleanup key will be different from all rest keys since all rest keys will // have a regular prefix. auto numFinishedWorker = addHelper(cleanupKey_, 1); - // The last worker cleans up the file - if (numFinishedWorker == numWorkers_) { + // The last worker cleans up the file. If numWorkers was not initialized to + // a specific postive value (i.e. meaning that there was not a fixed number + // of workers), we don't attempt to clean. + if (numWorkers_ >= 0 && numFinishedWorker == numWorkers_) { // Best effort removal without checking the return std::remove(path_.c_str()); } diff --git a/torch/lib/c10d/FileStore.hpp b/torch/lib/c10d/FileStore.hpp index 814dc4823e947..98c4f23430f56 100644 --- a/torch/lib/c10d/FileStore.hpp +++ b/torch/lib/c10d/FileStore.hpp @@ -38,6 +38,11 @@ class TORCH_API FileStore : public Store { const std::vector& keys, const std::chrono::milliseconds& timeout) override; + // Returns the path used by the FileStore. + const std::string& getPath() const noexcept { + return path_; + } + protected: int64_t addHelper(const std::string& key, int64_t i); From 15dbc566c57eedbd0245e786912e94586eba0fd2 Mon Sep 17 00:00:00 2001 From: Serhat Yilmaz Date: Wed, 16 Jun 2021 12:15:12 -0700 Subject: [PATCH 155/305] [torch][segment_reduce] Add missing cuda kernel launch check (#60114) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60114 Same as title. Test Plan: Unit test (test_kernel_launch_checks.py) is passing. Reviewed By: ngimel Differential Revision: D29169538 fbshipit-source-id: ba4518dcb1a4713144d92faec2bb5bdf656ff7c5 --- aten/src/ATen/native/cuda/SegmentReduce.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/aten/src/ATen/native/cuda/SegmentReduce.cu b/aten/src/ATen/native/cuda/SegmentReduce.cu index 5b4a8c40634e2..be3fc4003129d 100644 --- a/aten/src/ATen/native/cuda/SegmentReduce.cu +++ b/aten/src/ATen/native/cuda/SegmentReduce.cu @@ -138,6 +138,7 @@ Tensor _segment_reduce_cuda_kernel( segment_count, initial.has_value(), initial_value); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } }); From 691183bb740a890d6972803e996bcca63673a7c4 Mon Sep 17 00:00:00 2001 From: Erjia Guan Date: Wed, 16 Jun 2021 12:18:33 -0700 Subject: [PATCH 156/305] Fix compile failure on CUDA92 (#60017) Summary: Fixes https://github.com/pytorch/pytorch/issues/60016 For CUDA 92 - OptionalBase was not check if `is_arrayref` - constexpr seems not expect to raise Exception for cuda 92 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60017 Reviewed By: malfet Differential Revision: D29139515 Pulled By: ejguan fbshipit-source-id: 4f4f6d9fe6a5f2eadf913de0a9781cc9f2e6ac6f --- c10/macros/Macros.h | 6 ++++++ c10/util/ArrayRef.h | 6 ++++-- c10/util/Optional.h | 28 +++++++++++++++++----------- 3 files changed, 27 insertions(+), 13 deletions(-) diff --git a/c10/macros/Macros.h b/c10/macros/Macros.h index 7dff18da80f30..a4aa2faf6a5c7 100644 --- a/c10/macros/Macros.h +++ b/c10/macros/Macros.h @@ -409,6 +409,12 @@ __host__ __device__ #define C10_HOST_CONSTEXPR_VAR constexpr #endif +#if defined(__CUDA_ARCH__) && defined(CUDA_VERSION) && CUDA_VERSION <= 9200 +#define C10_CONSTEXPR_EXCEPT_CUDA92 +#else +#define C10_CONSTEXPR_EXCEPT_CUDA92 constexpr +#endif + #if !defined(__clang__) && !defined(_MSC_VER) && defined(__GNUC__) && \ __GNUC__ < 6 #define CONSTEXPR_EXCEPT_GCC5 diff --git a/c10/util/ArrayRef.h b/c10/util/ArrayRef.h index a3a85aa3ad1d7..edfca99622f80 100644 --- a/c10/util/ArrayRef.h +++ b/c10/util/ArrayRef.h @@ -72,13 +72,15 @@ class ArrayRef final { constexpr ArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {} /// Construct an ArrayRef from a pointer and length. - constexpr ArrayRef(const T* data, size_t length) + /// CUDA 9.2 fails to compile constexpr that throws exception + C10_CONSTEXPR_EXCEPT_CUDA92 ArrayRef(const T* data, size_t length) : Data(data), Length(length) { debugCheckNullptrInvariant(); } /// Construct an ArrayRef from a range. - constexpr ArrayRef(const T* begin, const T* end) + /// CUDA 9.2 fails to compile constexpr that throws exception + C10_CONSTEXPR_EXCEPT_CUDA92 ArrayRef(const T* begin, const T* end) : Data(begin), Length(end - begin) { debugCheckNullptrInvariant(); } diff --git a/c10/util/Optional.h b/c10/util/Optional.h index 15f59a61c3846..5e0684bb7d2f5 100644 --- a/c10/util/Optional.h +++ b/c10/util/Optional.h @@ -526,11 +526,14 @@ using OptionalBase = std::conditional_t< optional_base>>>>; #else template -using OptionalBase = typename std::conditional< - std::is_trivially_destructible::value, // if possible - constexpr_optional_base::type>, // use base with trivial destructor - optional_base::type>>::type; +using OptionalBase = std::conditional_t< + detail_::is_arrayref::value, + arrayref_optional_base, + std::conditional_t< + std::is_trivially_destructible::value, // if possible + constexpr_optional_base>, // use base with + // trivial destructor + optional_base>>>; #endif template @@ -564,11 +567,14 @@ class optional : private OptionalBase { optional_base>>>>; #else template - using OptionalBase = typename std::conditional< - std::is_trivially_destructible::value, // if possible - constexpr_optional_base::type>, // use base with trivial destructor - optional_base::type>>::type; + using OptionalBase = std::conditional_t< + detail_::is_arrayref::value, + arrayref_optional_base, + std::conditional_t< + std::is_trivially_destructible::value, // if possible + constexpr_optional_base>, // use base with + // trivial destructor + optional_base>>>; #endif static_assert( @@ -638,7 +644,7 @@ class optional : private OptionalBase { std::is_nothrow_move_constructible::value) { if (rhs.initialized()) { ::new (static_cast(dataptr())) T(std::move(*rhs)); - setInitialized(true); + OptionalBase::setInitialized(true); } } #endif From b2fc6de2c442569533ccc7f7db121f274bc2556b Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Wed, 16 Jun 2021 13:31:14 -0700 Subject: [PATCH 157/305] support parsing of PR stats in run_test.py (#60026) Summary: Currently S3 test stats doesn't support PR stats parisng. Changes to s3_stats_parser: 1. they are uploaded to `test_times/{sha1}/{job}` and `pr_test_times/{pr}/{sha1}/{job}` separately. Thus we need parsing logics for both 2. need to attach time for PR stats parsing for ordering since PR commits can be force-pushed Changes to run_test.py 1. Reordering based on previous PR stats if available 2. Falling back to file change option if not enabled. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60026 Test Plan: - CI. - local repro: plz run: ``` CIRCLE_JOB="pytorch_linux_bionic_py3_6_clang9_noarch_test" CIRCLE_PR_NUMBER=60057 IN_CI=1 ENABLE_PR_HISTORY_REORDERING=1 python test/run_test.py ``` Reviewed By: samestep Differential Revision: D29164754 Pulled By: walterddr fbshipit-source-id: 206688e0fb0b78d1c9042c07243da1fbf88a924b --- test/run_test.py | 72 +++++++++++++++++++++-------- tools/stats_utils/s3_stat_parser.py | 38 +++++++++++++-- 2 files changed, 88 insertions(+), 22 deletions(-) diff --git a/test/run_test.py b/test/run_test.py index a0e73b9c86ed1..997d43c0cad27 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -23,7 +23,11 @@ try: sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) - from tools.stats_utils.s3_stat_parser import (get_previous_reports_for_branch, Report, HAVE_BOTO3) + from tools.stats_utils.s3_stat_parser import ( + get_previous_reports_for_branch, + get_previous_reports_for_pr, + Report, + HAVE_BOTO3) except ImportError: print("Unable to import s3_stat_parser from tools. Running without S3 stats...") HAVE_BOTO3 = False @@ -385,6 +389,8 @@ PYTORCH_COLLECT_COVERAGE = bool(os.environ.get("PYTORCH_COLLECT_COVERAGE")) +ENABLE_PR_HISTORY_REORDERING = bool(os.environ.get("ENABLE_PR_HISTORY_REORDERING", "0") == "1") + JIT_EXECUTOR_TESTS = [ 'test_jit_cuda_fuser', 'test_jit_profiling', @@ -1163,36 +1169,66 @@ def query_changed_test_files() -> List[str]: lines = [line.strip() for line in lines] return lines +def query_failure_test_module(reports: List[Tuple["Report", str]]) -> List[str]: + test_modules = [] + if len(reports) == 0 or len(reports[0]) == 0: + return test_modules + report = reports[0][0] + assert report.get('format_version') == 2, "S3 format currently handled is version 2 only" + files: Dict[str, Any] = report['files'] + for fname, file in files.items(): + contains_failure = any( + any(case['status'] == 'errored' or case['status'] == 'failed' + for _, case in suite['cases'].items()) + for _, suite in file['suites'].items()) + if contains_failure: + test_modules.append(fname) + return test_modules + def reorder_tests(tests: List[str]) -> List[str]: - try: - changed_files = query_changed_test_files() - except Exception: - # If unable to get changed files from git, quit without doing any sorting - return tests + prioritized_tests = [] + # Try using historic stats from PR. + if ENABLE_PR_HISTORY_REORDERING and HAVE_BOTO3: + pr_number = os.environ.get("CIRCLE_PR_NUMBER", "") + if len(pr_number): + ci_job_prefix = get_stripped_CI_job() + s3_reports: List[Tuple["Report", str]] = get_previous_reports_for_pr( + pr_number, ci_job_prefix) + prioritized_tests = query_failure_test_module(s3_reports) + print("Prioritized test from previous CI info.") + + # Using file changes priority if no stats found from previous PR. + if len(prioritized_tests) == 0: + try: + changed_files = query_changed_test_files() + except Exception: + # If unable to get changed files from git, quit without doing any sorting + return tests - prefix = f"test{os.path.sep}" - changed_tests = [f for f in changed_files if f.startswith(prefix) and f.endswith(".py")] - changed_tests = [f[len(prefix):] for f in changed_tests] - changed_tests = [f[:-len(".py")] for f in changed_tests] + prefix = f"test{os.path.sep}" + prioritized_tests = [f for f in changed_files if f.startswith(prefix) and f.endswith(".py")] + prioritized_tests = [f[len(prefix):] for f in prioritized_tests] + prioritized_tests = [f[:-len(".py")] for f in prioritized_tests] + print("Prioritized test from test file changes.") bring_to_front = [] the_rest = [] for test in tests: - if test in changed_tests: + if test in prioritized_tests: bring_to_front.append(test) else: the_rest.append(test) - - sorted_tests = bring_to_front + the_rest - - if len(sorted_tests) != len(tests): - # Something went wrong, bail out without doing any sorting + if len(tests) == len(bring_to_front) + len(the_rest): + print(f"reordering tests for PR:\n" + f"prioritized: {bring_to_front}\nthe rest: {the_rest}\n") + return bring_to_front + the_rest + else: + print(f"Something went wrong in CI reordering, expecting total of {len(tests)}:\n" + f"but found prioritized: {len(bring_to_front)}\nthe rest: {len(the_rest)}\n") return tests - return sorted_tests - def main(): options = parse_args() diff --git a/tools/stats_utils/s3_stat_parser.py b/tools/stats_utils/s3_stat_parser.py index 89f9eed730ee9..3dccac4efda23 100644 --- a/tools/stats_utils/s3_stat_parser.py +++ b/tools/stats_utils/s3_stat_parser.py @@ -4,7 +4,7 @@ import subprocess from collections import defaultdict from datetime import datetime, timedelta -from typing import Dict, List, Optional, Union, Any, cast +from typing import Dict, List, Optional, Tuple, Union, Any, cast from typing_extensions import Literal, TypedDict try: @@ -143,9 +143,10 @@ def get_cases( return cases -def _parse_s3_summaries(summaries: Any, jobs: List[str]) -> Dict[str, List[Report]]: +def _parse_master_summaries(summaries: Any, jobs: List[str]) -> Dict[str, List[Report]]: summary_dict = defaultdict(list) for summary in summaries: + # master summary format: "test_time/{sha}/{job}/file" summary_job = summary.key.split('/')[2] if summary_job in jobs or len(jobs) == 0: binary = summary.get()["Body"].read() @@ -153,19 +154,37 @@ def _parse_s3_summaries(summaries: Any, jobs: List[str]) -> Dict[str, List[Repor summary_dict[summary_job].append(json.loads(string)) return summary_dict +def _parse_pr_summaries(summaries: Any, job_prefix: str) -> Dict[str, List[Tuple[Report, str]]]: + summary_dict = defaultdict(list) + for summary in summaries: + # PR summary format: "pr_test_time/{pr}/{sha}/{job}/file" + summary_job = summary.key.split('/')[3] + summary_timestamp = summary.key.split('/')[4][:len("YYYY-MM-ddTHH:mm:ss")] + if not job_prefix or len(job_prefix) == 0 or summary_job.startswith(job_prefix): + binary = summary.get()["Body"].read() + string = bz2.decompress(binary).decode("utf-8") + summary_dict[summary_job].append((json.loads(string), summary_timestamp)) + return summary_dict + + # Collect and decompress S3 test stats summaries into JSON. # data stored on S3 buckets are pathed by {sha}/{job} so we also allow # optional jobs filter def get_test_stats_summaries(*, sha: str, jobs: Optional[List[str]] = None) -> Dict[str, List[Report]]: bucket = get_S3_bucket_readonly(OSSCI_METRICS_BUCKET) summaries = bucket.objects.filter(Prefix=f"test_time/{sha}") - return _parse_s3_summaries(summaries, jobs=list(jobs or [])) + return _parse_master_summaries(summaries, jobs=list(jobs or [])) def get_test_stats_summaries_for_job(*, sha: str, job_prefix: str) -> Dict[str, List[Report]]: bucket = get_S3_bucket_readonly(OSSCI_METRICS_BUCKET) summaries = bucket.objects.filter(Prefix=f"test_time/{sha}/{job_prefix}") - return _parse_s3_summaries(summaries, jobs=list()) + return _parse_master_summaries(summaries, jobs=list()) + +def get_test_stats_summaries_for_pr(*, pr: str, job_prefix: str) -> Dict[str, List[Tuple[Report, str]]]: + bucket = get_S3_bucket_readonly(OSSCI_METRICS_BUCKET) + summaries = bucket.objects.filter(Prefix=f"pr_test_time/{pr}/") + return _parse_pr_summaries(summaries, job_prefix=job_prefix) # This function returns a list of S3 test time reports. This function can run into errors if HAVE_BOTO3 = False @@ -195,3 +214,14 @@ def get_previous_reports_for_branch(branch: str, ci_job_prefix: str = "") -> Lis logger.warning(f'WARNING: Multiple summary objects found for {commit}/{job_name}') commit_index += 1 return reports + + +def get_previous_reports_for_pr(pr: str, ci_job_prefix: str = "") -> List[Tuple[Report, str]]: + reports: List[Tuple[Report, str]] = [] + logger.info(f'Grabbing reports from PR: {[pr]}') + summaries = get_test_stats_summaries_for_pr(pr=pr, job_prefix=ci_job_prefix) + for _, summary in summaries.items(): + reports.extend(summary) + # sort by summary_timestamp + reports.sort(reverse=True, key=lambda s: s[1]) + return reports From 9d7764642b4a353e25c1545c7f110bb91f643210 Mon Sep 17 00:00:00 2001 From: David Riazati Date: Wed, 16 Jun 2021 13:38:41 -0700 Subject: [PATCH 158/305] Use GitHub's diff directly in clang-tidy (#60048) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60048 This changes clang-tidy in lint.yml to pull the raw diff from GitHub and parse that rather than use the PRs base revision. The base revision can cause the spurious inclusion of files not changed in the PR as in https://github.com/pytorch/pytorch/pull/59967/checks?check_run_id=2832565901. We could be smarter about how we query git, but this approach ends up being simpler since we just need to search for the diff headers in the .diff file. See https://github.com/pytorch/pytorch/pull/60049/checks?check_run_id=2834140350 for an example CI run with this on Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29148886 Pulled By: driazati fbshipit-source-id: ca23446d5cc8938d1345f272afe77b9ee8898b74 --- .github/workflows/lint.yml | 47 +++++++++++++++------------- tools/clang_tidy.py | 64 ++++++++++++++++++-------------------- 2 files changed, 55 insertions(+), 56 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 917a8c70e0b0e..52fcbc1a0f56b 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -312,12 +312,14 @@ jobs: fi - name: Run clang-tidy env: - BASE_SHA: ${{ github.event.pull_request.base.sha }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} + PR_NUMBER: ${{ github.event.pull_request.number }} run: | cd "${GITHUB_WORKSPACE}" set -eux + wget -O pr.diff "https://patch-diff.githubusercontent.com/raw/pytorch/pytorch/pull/$PR_NUMBER.diff" + # Run Clang-Tidy # 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. @@ -326,27 +328,28 @@ jobs: # /torch/csrc/generic/*.cpp is excluded because those files aren't actually built. # deploy/interpreter files are excluded due to using macros and other techniquies # that are not easily converted to accepted c++ - python3 tools/clang_tidy.py \ - --verbose \ - --paths torch/csrc/ \ - --diff "$BASE_SHA" \ - -g"-torch/csrc/jit/passes/onnx/helper.cpp" \ - -g"-torch/csrc/jit/passes/onnx/shape_type_inference.cpp"\ - -g"-torch/csrc/jit/serialization/onnx.cpp" \ - -g"-torch/csrc/jit/serialization/export.cpp" \ - -g"-torch/csrc/jit/serialization/import.cpp" \ - -g"-torch/csrc/jit/serialization/import_legacy.cpp" \ - -g"-torch/csrc/onnx/init.cpp" \ - -g"-torch/csrc/cuda/nccl.*" \ - -g"-torch/csrc/cuda/python_nccl.cpp" \ - -g"-torch/csrc/autograd/FunctionsManual.cpp" \ - -g"-torch/csrc/generic/*.cpp" \ - -g"-torch/csrc/jit/codegen/cuda/runtime/*" \ - -g"-torch/csrc/deploy/interpreter/interpreter.cpp" \ - -g"-torch/csrc/deploy/interpreter/interpreter.h" \ - -g"-torch/csrc/deploy/interpreter/interpreter_impl.h" \ - -g"-torch/csrc/deploy/interpreter/test_main.cpp" \ - "$@" > "${GITHUB_WORKSPACE}"/clang-tidy-output.txt + python3 tools/clang_tidy.py \ + --verbose \ + --paths torch/csrc/ \ + --diff-file pr.diff \ + -g"-torch/csrc/jit/passes/onnx/helper.cpp" \ + -g"-torch/csrc/jit/passes/onnx/shape_type_inference.cpp" \ + -g"-torch/csrc/jit/serialization/onnx.cpp" \ + -g"-torch/csrc/jit/serialization/export.cpp" \ + -g"-torch/csrc/jit/serialization/import.cpp" \ + -g"-torch/csrc/jit/serialization/import_legacy.cpp" \ + -g"-torch/csrc/onnx/init.cpp" \ + -g"-torch/csrc/cuda/nccl.*" \ + -g"-torch/csrc/cuda/python_nccl.cpp" \ + -g"-torch/csrc/autograd/FunctionsManual.cpp" \ + -g"-torch/csrc/generic/*.cpp" \ + -g"-torch/csrc/jit/codegen/cuda/runtime/*" \ + -g"-torch/csrc/deploy/interpreter/interpreter.cpp" \ + -g"-torch/csrc/deploy/interpreter/interpreter.h" \ + -g"-torch/csrc/deploy/interpreter/interpreter_impl.h" \ + -g"-torch/csrc/deploy/interpreter/test_main.cpp" \ + "$@" >"${GITHUB_WORKSPACE}"/clang-tidy-output.txt + cat "${GITHUB_WORKSPACE}"/clang-tidy-output.txt diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index f5c71f41cd3d2..7574c4f3b538e 100755 --- a/tools/clang_tidy.py +++ b/tools/clang_tidy.py @@ -21,7 +21,6 @@ import os import os.path import re -import shlex import shutil import subprocess import sys @@ -32,7 +31,7 @@ except ImportError: from pipes import quote -from typing import Any, Dict, Iterable, List, Set, Union +from typing import Any, Dict, Iterable, List, Set, Tuple Patterns = collections.namedtuple("Patterns", "positive, negative") @@ -42,8 +41,13 @@ # (c/cc/cpp) file. DEFAULT_FILE_PATTERN = re.compile(r".*\.c(c|pp)?") -# @@ -start,count +start,count @@ -CHUNK_PATTERN = r"^@@\s+-\d+(?:,\d+)?\s+\+(\d+)(?:,(\d+))?\s+@@" +# Search for: +# diff --git ... +# index ... +# --- ... +# +++ ... +CHUNK_HEADER_RE = r"diff --git .*?\nindex.*?\n---.*?\n\+\+\+ b/(.*?)\n@@ -(\d+,\d+) \+(\d+,\d+) @@" + CLANG_WARNING_PATTERN = re.compile(r"([^:]+):(\d+):\d+:\s+warning:.*\[([^\]]+)\]") @@ -125,35 +129,25 @@ def filter_files(files: Iterable[str], file_patterns: Patterns) -> Iterable[str] print("{} omitted due to file filters".format(file)) -def get_changed_files(revision: str, paths: List[str]) -> List[str]: - """Runs git diff to get the paths of all changed files.""" - # --diff-filter AMU gets us files that are (A)dded, (M)odified or (U)nmerged (in the working copy). - # --name-only makes git diff return only the file paths, without any of the source changes. - command = "git diff-index --diff-filter=AMU --ignore-all-space --name-only" - output = run_shell_command(shlex.split(command) + [revision] + paths) - return output.split("\n") - - def get_all_files(paths: List[str]) -> List[str]: """Returns all files that are tracked by git in the given paths.""" output = run_shell_command(["git", "ls-files"] + paths) return output.split("\n") -def get_changed_lines(revision: str, filename: str) -> Dict[str, Union[str, List[List[int]]]]: - """Runs git diff to get the line ranges of all file changes.""" - command = shlex.split("git diff-index --unified=0") + [revision, filename] - output = run_shell_command(command) - changed_lines = [] - for chunk in re.finditer(CHUNK_PATTERN, output, re.MULTILINE): - start = int(chunk.group(1)) - count = int(chunk.group(2) or 1) - # If count == 0, a chunk was removed and can be ignored. - if count == 0: - continue - changed_lines.append([start, start + count]) +def find_changed_lines(diff: str) -> Dict[str, List[Tuple[int, int]]]: + files = collections.defaultdict(list) + + matches = re.findall(CHUNK_HEADER_RE, diff, re.MULTILINE) + for file, start, end in matches: + start_line, _ = start.split(",") + end_line, _ = end.split(",") + print(file, start_line, end_line) + + files[file].append((start_line, end_line)) + + return dict(files) - return {"name": filename, "lines": changed_lines} ninja_template = """ rule do_cmd @@ -180,7 +174,7 @@ def run_shell_commands_in_parallel(commands: Iterable[List[str]]) -> str: return run_shell_command(['ninja', '-f', f.name]) -def run_clang_tidy(options: Any, line_filters: Any, files: Iterable[str]) -> str: +def run_clang_tidy(options: Any, line_filters: List[Dict[str, Any]], files: Iterable[str]) -> str: """Executes the actual clang-tidy command in the shell.""" command = [options.clang_tidy_exe, "-p", options.compile_commands_dir] if not options.config_file and os.path.exists(".clang-tidy"): @@ -283,7 +277,7 @@ def parse_options() -> Any: help="Path to the folder containing compile_commands.json", ) parser.add_argument( - "-d", "--diff", help="Git revision to diff against to get changes" + "--diff-file", help="File containing diff to use for determining files to lint and line filters" ) parser.add_argument( "-p", @@ -333,9 +327,15 @@ def main() -> None: # Normalize the paths first. paths = [path.rstrip("/") for path in options.paths] - if options.diff: - files = get_changed_files(options.diff, paths) + if options.diff_file: + with open(options.diff_file, "r") as f: + changed_files = find_changed_lines(f.read()) + line_filters = [ + {"name": name, "lines": lines} for name, lines, in changed_files.items() + ] + files = list(changed_files.keys()) else: + line_filters = [] files = get_all_files(paths) file_patterns = get_file_patterns(options.glob, options.regex) files = list(filter_files(files, file_patterns)) @@ -345,10 +345,6 @@ def main() -> None: print("No files detected.") sys.exit() - line_filters = [] - if options.diff: - line_filters = [get_changed_lines(options.diff, f) for f in files] - clang_tidy_output = run_clang_tidy(options, line_filters, files) if options.suppress_diagnostics: warnings = extract_warnings(clang_tidy_output, base_dir=options.compile_commands_dir) From fd1e9253fffea2e48bd20f6f691a3aeef5e797c8 Mon Sep 17 00:00:00 2001 From: Gisle Dankel Date: Wed, 16 Jun 2021 14:24:07 -0700 Subject: [PATCH 159/305] [Profiler] Fix timestamp discrepancy in profiler_kineto.cpp (#60070) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60070 PyTorch pull request https://github.com/pytorch/pytorch/pull/57333 changed high_resolution_clock to system_clock but missed one location in profiler_kineto.cpp. On some platforms (e.g. Windows), high_resolution_clock and system_clock do not map to the same underlying clock and therefore we get mixed timestamps on some platforms. Reviewed By: wesolwsk Differential Revision: D29155809 fbshipit-source-id: a6de6b4d550613f26f5577487c3c53716896e219 --- torch/csrc/autograd/profiler_kineto.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torch/csrc/autograd/profiler_kineto.cpp b/torch/csrc/autograd/profiler_kineto.cpp index ed1cb660f17d9..e3d50227c86e3 100644 --- a/torch/csrc/autograd/profiler_kineto.cpp +++ b/torch/csrc/autograd/profiler_kineto.cpp @@ -9,6 +9,7 @@ #ifdef USE_KINETO #include +#include #ifndef _MSC_VER // TODO: TO be removed, once this properly works from libkineto @@ -32,8 +33,7 @@ uint64_t next_correlation_id() { } inline int64_t getTimeUs() { - using namespace std::chrono; - return duration_cast(high_resolution_clock::now().time_since_epoch()).count(); + return libkineto::timeSinceEpoch(std::chrono::system_clock::now()); } std::string shapesToStr(const std::vector>& shapes); From 8ce6d0c42f0da27d4d5a01bea7d4b7016b5aa38e Mon Sep 17 00:00:00 2001 From: Zachary DeVito Date: Wed, 16 Jun 2021 14:40:10 -0700 Subject: [PATCH 160/305] [torch deploy] add register_module_source (#58290) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/58290 this is a helper function to get some python source code loaded on each interpreter without having to use the standard import system or packages. Useful for debugging or for writing wrapper classes for handling loaded modules. Test Plan: Imported from OSS Reviewed By: wconstab Differential Revision: D28435306 Pulled By: zdevito fbshipit-source-id: b85c16346b9001cd7350d65879cb990098060813 --- torch/csrc/deploy/deploy.h | 20 +++++++++ .../deploy/interpreter/interpreter_impl.cpp | 41 +++++++++++++++++++ .../deploy/interpreter/interpreter_impl.h | 3 ++ torch/csrc/deploy/test_deploy.cpp | 9 ++++ 4 files changed, 73 insertions(+) diff --git a/torch/csrc/deploy/deploy.h b/torch/csrc/deploy/deploy.h index 92cdc3e67cc88..640f3da66323b 100644 --- a/torch/csrc/deploy/deploy.h +++ b/torch/csrc/deploy/deploy.h @@ -116,6 +116,15 @@ struct TORCH_API InterpreterManager { // can be used for balancing work across GPUs I.global("torch", "version").attr("__setattr__")({"interp", int(i)}); // std::cerr << "Interpreter " << i << " initialized\n"; + instances_.back().pImpl_->set_find_module( + [this](const std::string& name) -> at::optional { + auto it = registered_module_sources_.find(name); + if (it != registered_module_sources_.end()) { + return it->second; + } else { + return at::nullopt; + } + }); } TORCH_DEPLOY_SAFE_CATCH_RETHROW } @@ -146,6 +155,16 @@ struct TORCH_API InterpreterManager { Package load_package(const std::string& uri); Package load_package( std::shared_ptr reader); + + // convience function for loading some python source code as a module across + // all interpreters. this can be used for writing tests of deploy that need to + // execute python code, or for small amounts of application logic that are + // best written in Python. For larger amounts of code, prefer creating and + // loading them as packages. + void register_module_source(std::string name, std::string src) { + registered_module_sources_[std::move(name)] = std::move(src); + } + InterpreterManager(const InterpreterManager&) = delete; InterpreterManager& operator=(const InterpreterManager&) = delete; InterpreterManager& operator=(InterpreterManager&&) = delete; @@ -156,6 +175,7 @@ struct TORCH_API InterpreterManager { size_t next_object_id_ = 0; std::vector instances_; LoadBalancer resources_; + std::unordered_map registered_module_sources_; }; struct TORCH_API ReplicatedObjImpl { diff --git a/torch/csrc/deploy/interpreter/interpreter_impl.cpp b/torch/csrc/deploy/interpreter/interpreter_impl.cpp index 01e6ca8be9e2f..f247cb70028ca 100644 --- a/torch/csrc/deploy/interpreter/interpreter_impl.cpp +++ b/torch/csrc/deploy/interpreter/interpreter_impl.cpp @@ -8,6 +8,7 @@ // NOLINTNEXTLINE(modernize-deprecated-headers) #include #include +#include // NOLINTNEXTLINE(modernize-deprecated-headers) #include #include @@ -112,6 +113,8 @@ extern "C" struct _frozen _PyImport_FrozenModules_torch[]; const char* startup = R"RAW( import sys +import importlib.abc +import linecache # We need to register a custom meta path finder because we are registering # `torch._C` as a builtin module. @@ -129,6 +132,28 @@ class F: return None sys.meta_path.insert(0, F()) +class RegisterModuleImporter(importlib.abc.InspectLoader): + def __init__(self, find_module_source): + self.find_module_source = find_module_source + + def create_module(self, spec): + return None + + def get_source(self, name): + return self.find_module_source(name) + + def exec_module(self, module): + filename = f"_deploy_internal.{module.__name__}" + linecache.lazycache(filename, module.__dict__) + code = compile(self.get_source(module.__name__), filename, "exec", dont_inherit=True) + exec(code, module.__dict__) + + def find_spec(self, fullname, path, target=None): + r = self.find_module_source(fullname) + if r is not None: + return importlib.util.spec_from_loader(fullname, self) + return None + # print("exec_prefix:", sys.base_exec_prefix) # print("_base_executable:", sys._base_executable) # print("base_prefix:", sys.base_prefix) @@ -328,6 +353,22 @@ struct ConcreteInterpreterImpl : public torch::deploy::InterpreterImpl { } } + void set_find_module( + std::function(const std::string&)> find_module) + override { + std::function wrapped_find_module = + [=](const std::string& name) -> py::object { + auto r = find_module(name); + return r ? py::cast(*r) : py::none(); + }; + py::object register_module_importer = + py::module::import("__main__") + .attr("RegisterModuleImporter")(wrapped_find_module); + py::module::import("sys") + .attr("meta_path") + .attr("append")(register_module_importer); + } + torch::deploy::InterpreterSessionImpl* acquire_session() override; py::object save_storage; py::object load_storage; diff --git a/torch/csrc/deploy/interpreter/interpreter_impl.h b/torch/csrc/deploy/interpreter/interpreter_impl.h index 0148b73e47e7e..3a0b3854ac44e 100644 --- a/torch/csrc/deploy/interpreter/interpreter_impl.h +++ b/torch/csrc/deploy/interpreter/interpreter_impl.h @@ -125,6 +125,9 @@ struct InterpreterSessionImpl { struct InterpreterImpl { virtual InterpreterSessionImpl* acquire_session() = 0; + virtual void set_find_module( + std::function(const std::string&)> + find_module) = 0; virtual ~InterpreterImpl() = default; // this will uninitialize python }; diff --git a/torch/csrc/deploy/test_deploy.cpp b/torch/csrc/deploy/test_deploy.cpp index 3631534f0ec42..bdc91a2571d5c 100644 --- a/torch/csrc/deploy/test_deploy.cpp +++ b/torch/csrc/deploy/test_deploy.cpp @@ -229,3 +229,12 @@ TEST(TorchpyTest, DisarmHook) { auto I = m.acquire_one(); ASSERT_THROW(I.from_ivalue(t), c10::Error); // NOT a segfault } + +TEST(TorchpyTest, RegisterModule) { + torch::deploy::InterpreterManager m(2); + m.register_module_source("foomodule", "def add1(x): return x + 1\n"); + for (const auto& interp : m.all_instances()) { + auto I = interp.acquire_session(); + AT_ASSERT(3 == I.global("foomodule", "add1")({2}).toIValue().toInt()); + } +} From 4c8c61f200042a19d90580b00b507ed9ed9e524c Mon Sep 17 00:00:00 2001 From: Richard Barnes Date: Wed, 16 Jun 2021 15:15:59 -0700 Subject: [PATCH 161/305] Some fixes to vec256_bfloat16.h (#59957) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59957 Test Plan: Sandcastle Reviewed By: VitalyFedyunin Differential Revision: D29073913 fbshipit-source-id: dc01a2015e4ff42daa1d69443460182744c06e90 --- aten/src/ATen/cpu/vec/vec256/vec256_base.h | 4 ++-- .../src/ATen/cpu/vec/vec256/vec256_bfloat16.h | 20 +++++++++---------- .../cpu/vec/vec256/vec256_complex_double.h | 2 +- .../cpu/vec/vec256/vec256_complex_float.h | 2 +- aten/src/ATen/cpu/vec/vec256/vec256_double.h | 2 +- aten/src/ATen/cpu/vec/vec256/vec256_float.h | 2 +- .../ATen/cpu/vec/vec256/vec256_float_neon.h | 2 +- .../vec256/vsx/vec256_complex_double_vsx.h | 4 ++-- .../vec/vec256/vsx/vec256_complex_float_vsx.h | 4 ++-- .../cpu/vec/vec256/vsx/vec256_double_vsx.h | 4 ++-- .../cpu/vec/vec256/vsx/vec256_float_vsx.h | 4 ++-- 11 files changed, 25 insertions(+), 25 deletions(-) diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_base.h b/aten/src/ATen/cpu/vec/vec256/vec256_base.h index 39dc995f20bf5..596dac67c2cd3 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_base.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_base.h @@ -213,14 +213,14 @@ struct Vectorized { } return vec; } - Vectorized map(T (*f)(T)) const { + Vectorized map(T (*const f)(T)) const { Vectorized ret; for (int64_t i = 0; i != size(); i++) { ret[i] = f(values[i]); } return ret; } - Vectorized map(T (*f)(const T &)) const { + Vectorized map(T (*const f)(const T &)) const { Vectorized ret; for (int64_t i = 0; i != size(); i++) { ret[i] = f(values[i]); diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_bfloat16.h b/aten/src/ATen/cpu/vec/vec256/vec256_bfloat16.h index f12879b5bc7c2..d19afb0d34701 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_bfloat16.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_bfloat16.h @@ -197,19 +197,19 @@ template <> class Vectorized { } return b; } - Vectorized map(const __m256 (*vop)(__m256)) const { + Vectorized map(const __m256 (*const vop)(__m256)) const { __m256 lo, hi; cvtbf16_fp32(values, lo, hi); - auto o1 = vop(lo); - auto o2 = vop(hi); + const auto o1 = vop(lo); + const auto o2 = vop(hi); return cvtfp32_bf16(o1, o2); } Vectorized abs() const { __m256 lo, hi; cvtbf16_fp32(values, lo, hi); - auto mask = _mm256_set1_ps(-0.f); - auto o1 = _mm256_andnot_ps(mask, lo); - auto o2 = _mm256_andnot_ps(mask, hi); + const auto mask = _mm256_set1_ps(-0.f); + const auto o1 = _mm256_andnot_ps(mask, lo); + const auto o2 = _mm256_andnot_ps(mask, hi); return cvtfp32_bf16(o1, o2); } Vectorized angle() const { @@ -328,17 +328,17 @@ template <> class Vectorized { Vectorized i0e() const { __m256 lo, hi; cvtbf16_fp32(values, lo, hi); - auto sz = size(); + constexpr auto sz = size(); __at_align32__ float tmp1[sz / 2], tmp2[sz / 2]; _mm256_storeu_ps(reinterpret_cast(tmp1), lo); _mm256_storeu_ps(reinterpret_cast(tmp2), hi); - for (decltype(sz) i = 0; i < sz / 2; i++) { + for (auto i = decltype(sz){0}; i < sz / 2; i++) { tmp1[i] = calc_i0e(tmp1[i]); tmp2[i] = calc_i0e(tmp2[i]); } - auto o1 = _mm256_loadu_ps(tmp1); - auto o2 = _mm256_loadu_ps(tmp2); + const auto o1 = _mm256_loadu_ps(tmp1); + const auto o2 = _mm256_loadu_ps(tmp2); return cvtfp32_bf16(o1, o2); } Vectorized igamma(const Vectorized &x) const { diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_complex_double.h b/aten/src/ATen/cpu/vec/vec256/vec256_complex_double.h index 2400afd57d3a8..f96aea6e09ebd 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_complex_double.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_complex_double.h @@ -105,7 +105,7 @@ template <> class Vectorized> { } const c10::complex& operator[](int idx) const = delete; c10::complex& operator[](int idx) = delete; - Vectorized> map(c10::complex (*f)(const c10::complex &)) const { + Vectorized> map(c10::complex (*const f)(const c10::complex &)) const { __at_align32__ c10::complex tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_complex_float.h b/aten/src/ATen/cpu/vec/vec256/vec256_complex_float.h index 2c79d520cf228..5494828b56501 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_complex_float.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_complex_float.h @@ -141,7 +141,7 @@ template <> class Vectorized> { } const c10::complex& operator[](int idx) const = delete; c10::complex& operator[](int idx) = delete; - Vectorized> map(c10::complex (*f)(const c10::complex &)) const { + Vectorized> map(c10::complex (*const f)(const c10::complex &)) const { __at_align32__ c10::complex tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_double.h b/aten/src/ATen/cpu/vec/vec256/vec256_double.h index c5eacdbeb2dc9..1c575b9a28c7a 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_double.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_double.h @@ -99,7 +99,7 @@ template <> class Vectorized { Vectorized isnan() const { return _mm256_cmp_pd(values, _mm256_set1_pd(0.0), _CMP_UNORD_Q); } - Vectorized map(double (*f)(double)) const { + Vectorized map(double (*const f)(double)) const { __at_align32__ double tmp[size()]; store(tmp); for (int64_t i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_float.h b/aten/src/ATen/cpu/vec/vec256/vec256_float.h index 5ef85266d2afc..1f4c3f63477c1 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_float.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_float.h @@ -106,7 +106,7 @@ template <> class Vectorized { Vectorized isnan() const { return _mm256_cmp_ps(values, _mm256_set1_ps(0.0f), _CMP_UNORD_Q); } - Vectorized map(float (*f)(float)) const { + Vectorized map(float (*const f)(float)) const { __at_align32__ float tmp[size()]; store(tmp); for (int64_t i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h index 9ffd413346c0d..b39d808a13a3c 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h @@ -296,7 +296,7 @@ template <> class Vectorized { } return loadu(res); }; - Vectorized map(float (*f)(float)) const { + Vectorized map(float (*const f)(float)) const { __at_align32__ float tmp[size()]; store(tmp); for (int64_t i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h index 7b1d467292b2f..ce59bae3f4ffc 100644 --- a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h +++ b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h @@ -164,7 +164,7 @@ class Vectorized { const ComplexDbl& operator[](int idx) const = delete; ComplexDbl& operator[](int idx) = delete; - Vectorized map(ComplexDbl (*f)(ComplexDbl)) const { + Vectorized map(ComplexDbl (*const f)(ComplexDbl)) const { __at_align32__ ComplexDbl tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { @@ -173,7 +173,7 @@ class Vectorized { return loadu(tmp); } - Vectorized map(ComplexDbl (*f)(const ComplexDbl&)) const { + Vectorized map(ComplexDbl (*const f)(const ComplexDbl&)) const { __at_align32__ ComplexDbl tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h index 046fbc1e1edae..f96488964bb9f 100644 --- a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h +++ b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h @@ -220,7 +220,7 @@ class Vectorized { const ComplexFlt& operator[](int idx) const = delete; ComplexFlt& operator[](int idx) = delete; - Vectorized map(ComplexFlt (*f)(ComplexFlt)) const { + Vectorized map(ComplexFlt (*const f)(ComplexFlt)) const { __at_align32__ ComplexFlt tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { @@ -229,7 +229,7 @@ class Vectorized { return loadu(tmp); } - Vectorized map(ComplexFlt (*f)(const ComplexFlt&)) const { + Vectorized map(ComplexFlt (*const f)(const ComplexFlt&)) const { __at_align32__ ComplexFlt tmp[size()]; store(tmp); for (int i = 0; i < size(); i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h index d42c291309bd8..ac0a131878a02 100644 --- a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h +++ b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h @@ -191,7 +191,7 @@ class Vectorized { void dump() const { std::cout << _vec0[0] << "," << _vec0[1] << "," << _vec1[0] << "," << _vec1[1] << std::endl; } - Vectorized map(double (*f)(double)) const { + Vectorized map(double (*const f)(double)) const { Vectorized ret; for (int i = 0; i < size()/2; i++) { ret._vec0[i] = f(_vec0[i]); @@ -202,7 +202,7 @@ class Vectorized { return ret; } - Vectorized mapbi(double (*f)(double, double), const Vectorized& other) + Vectorized mapbi(double (*const f)(double, double), const Vectorized& other) const { Vectorized ret; for (int i = 0; i < size()/2; i++) { diff --git a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h index 56d2593e864aa..5fd1fb9afc80b 100644 --- a/aten/src/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h +++ b/aten/src/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h @@ -201,7 +201,7 @@ class Vectorized { const float& operator[](int idx) const = delete; float& operator[](int idx) = delete; - Vectorized map(float (*f)(float)) const { + Vectorized map(float (*const f)(float)) const { Vectorized ret; for (int i = 0; i < size() / 2; i++) { ret._vec0[i] = f(_vec0[i]); @@ -212,7 +212,7 @@ class Vectorized { return ret; } - Vectorized mapbi(float (*f)(float, float), const Vectorized& other) + Vectorized mapbi(float (*const f)(float, float), const Vectorized& other) const { Vectorized ret; for (int i = 0; i < size() / 2; i++) { From 5686fe5817538d7854fbe0a8b5f26c94fac14917 Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Wed, 16 Jun 2021 15:32:04 -0700 Subject: [PATCH 162/305] Revert D29154971: Training resnext with msuru_suru_union and ig_msuru_suru_union datasets Test Plan: revert-hammer Differential Revision: D29154971 (https://github.com/pytorch/pytorch/commit/9f68f93aca0f4cc9726ef697db13715f699fb234) Original commit changeset: d534d830020f fbshipit-source-id: a3d16acc8e6b66a6010b501c28dbe295f573bc86 --- torch/nn/parallel/distributed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py index 2285a07f329cd..b46cae2e742d3 100644 --- a/torch/nn/parallel/distributed.py +++ b/torch/nn/parallel/distributed.py @@ -406,7 +406,7 @@ def __init__( broadcast_buffers=True, process_group=None, bucket_cap_mb=25, - find_unused_parameters=True, + find_unused_parameters=False, check_reduction=False, gradient_as_bucket_view=False, ): From a95207dad4c6b970d74c69e69d43a1a7514626b4 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Wed, 16 Jun 2021 16:06:05 -0700 Subject: [PATCH 163/305] [quant] Add a quantize_per_tensor overload that takes Tensor quantization parameters (#59773) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59773 Current quantize_per_tensor takes float scale and int zero_point, which does not work with Proxy, this PR adds a quantize_per_tensor overload that takes Tensor scale and zero_point instead. Test Plan: Tested locally that following runs without errors: ```python import torch from torch.quantization.quantize_fx import prepare_fx, convert_fx from torch.fx.experimental import normalize class TestModule(torch.nn.Module): def forward(self, x): return x + x mod = TestModule() mod.eval() config = {"": torch.quantization.get_default_qconfig("fbgemm")} mod = prepare_fx(mod, config) mod = convert_fx(mod) mod = torch.fx.Transformer(mod).transform() ``` Imported from OSS Reviewed By: vkuzo Differential Revision: D29019862 fbshipit-source-id: c0176040f3b73f0a30516ed17d261b44cc658407 --- aten/src/ATen/native/native_functions.yaml | 5 +++++ aten/src/ATen/native/quantized/QTensor.cpp | 9 +++++++++ test/quantization/fx/test_quantize_fx.py | 15 +++++++++++++++ 3 files changed, 29 insertions(+) diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index df977933ecd53..45f4ed2261f99 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -5132,6 +5132,11 @@ dispatch: CPU, CUDA: quantize_per_tensor +- func: quantize_per_tensor.tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, ScalarType dtype) -> Tensor + variants: function + dispatch: + CPU, CUDA: quantize_per_tensor_tensor_qparams + - func: quantize_per_tensor.tensors(Tensor[] tensors, Tensor scales, Tensor zero_points, ScalarType dtype) -> Tensor[] variants: function dispatch: diff --git a/aten/src/ATen/native/quantized/QTensor.cpp b/aten/src/ATen/native/quantized/QTensor.cpp index cc82179d3ee38..17736b28962f5 100644 --- a/aten/src/ATen/native/quantized/QTensor.cpp +++ b/aten/src/ATen/native/quantized/QTensor.cpp @@ -20,6 +20,15 @@ Tensor quantize_per_tensor( return quantizer->quantize(self); } +Tensor quantize_per_tensor_tensor_qparams( + const Tensor& self, + const Tensor& scale, + const Tensor& zero_point, + ScalarType dtype) { + auto quantizer = make_per_tensor_affine_quantizer(scale.item().toDouble(), zero_point.item().toLong(), dtype); + return quantizer->quantize(self); +} + std::vector quantize_per_tensor_list_cpu( TensorList tensors, const Tensor& scales, diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py index b8c64ed8f8ba0..e6f6cad3408d6 100644 --- a/test/quantization/fx/test_quantize_fx.py +++ b/test/quantization/fx/test_quantize_fx.py @@ -2511,6 +2511,21 @@ def forward(self, x): } self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence) + def test_trace_quantize_per_tensor(self): + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(1, 1, 1) + + def forward(self, x): + x = self.conv(x) + return x + + m = M().eval() + m = prepare_fx(m, {"": default_qconfig}) + m = convert_fx(m) + # Make sure this runs without error + m = torch.fx.Transformer(m).transform() @skipIfNoFBGEMM class TestQuantizeFxOps(QuantizationTestCase): From 5948e6f653a4021d95544c2f62358f9c0e16e8dc Mon Sep 17 00:00:00 2001 From: Patrick Date: Wed, 16 Jun 2021 16:28:52 -0700 Subject: [PATCH 164/305] removed gelu from autocast fp32 list (#59639) Summary: Fixes #{issue number} Pull Request resolved: https://github.com/pytorch/pytorch/pull/59639 Reviewed By: H-Huang Differential Revision: D29155914 Pulled By: ezyang fbshipit-source-id: feb117181894c2355768d5b1189b3d5f1649fc0b --- aten/src/ATen/autocast_mode.cpp | 1 - docs/source/amp.rst | 1 - torch/testing/_internal/autocast_test_lists.py | 1 - 3 files changed, 3 deletions(-) diff --git a/aten/src/ATen/autocast_mode.cpp b/aten/src/ATen/autocast_mode.cpp index 0e81c969064b3..922ea976c708e 100644 --- a/aten/src/ATen/autocast_mode.cpp +++ b/aten/src/ATen/autocast_mode.cpp @@ -368,7 +368,6 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) { KERNEL(ADD_NS(pow), "pow.Tensor_Tensor", Tensor (const Tensor &, const Tensor &), fp32) KERNEL(ADD_NS(pow), "pow.Scalar", Tensor (const Scalar&, const Tensor &), fp32) KERNEL(ADD_NS(softplus), "softplus", Tensor (const Tensor &, const Scalar&, const Scalar&), fp32) - KERNEL(ADD_NS(gelu), "gelu", Tensor (const Tensor &), fp32) KERNEL(ADD_NS(layer_norm), "layer_norm", Tensor (const Tensor &, IntArrayRef, const c10::optional&, const c10::optional&, double, bool), fp32) // The macro doesn't like this one (I think it chokes on commas inside <>) so write it manually m.impl(TORCH_SELECTIVE_NAME("aten::native_layer_norm"), diff --git a/docs/source/amp.rst b/docs/source/amp.rst index 4ebd123610069..539fd07be97c5 100644 --- a/docs/source/amp.rst +++ b/docs/source/amp.rst @@ -141,7 +141,6 @@ Ops that can autocast to ``float32`` ``erfinv``, ``exp``, ``expm1``, -``gelu``, ``grid_sample``, ``group_norm``, ``hinge_embedding_loss``, diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py index 5283ccecd886d..e7ac13d5ce7fb 100644 --- a/torch/testing/_internal/autocast_test_lists.py +++ b/torch/testing/_internal/autocast_test_lists.py @@ -212,7 +212,6 @@ def __init__(self, dev): ] self.nn_fp32 = [ ("softplus", pointwise0_fp16), - ("gelu", pointwise0_fp16), ("nll_loss", (torch.rand((n, n), device=dev, dtype=torch.float), torch.zeros((n,), device=dev, dtype=torch.long))), ("nll_loss2d", (torch.rand((n, n, n, n), device=dev, dtype=torch.half), From 1efa863837aca33a611c36f030de9f30b7087ca4 Mon Sep 17 00:00:00 2001 From: Victor Quach Date: Wed, 16 Jun 2021 16:39:47 -0700 Subject: [PATCH 165/305] Avoid un-necessary unwrapping of Tensor in SavedVariable (#59837) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59837 Fixes #58500 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29069215 fbshipit-source-id: 603db3c8a64b729e86385ed774825f01c6ce0f20 --- test/test_autograd.py | 27 ++++++++++++++++ torch/csrc/autograd/saved_variable.cpp | 45 +++++++++++++++++++------- torch/csrc/autograd/saved_variable.h | 21 ++++++++++-- 3 files changed, 79 insertions(+), 14 deletions(-) diff --git a/test/test_autograd.py b/test/test_autograd.py index 8fcf42cac1796..dfba768f33a42 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -5467,6 +5467,33 @@ def test_input_buffer_accum(self): # Given gradients should not be modified inplace self.assertEqual(grad_out1, grad_out1_original) + def test_no_unnecessary_unwrapping(self): + a = torch.randn(5, requires_grad=True) + a_orig = a.detach().clone() + b = a * a + c = a * b + d = torch.exp(a) + + # a is leaf + self.assertIs(b.grad_fn._saved_self, a) + self.assertIs(b.grad_fn._saved_other, a) + self.assertIs(c.grad_fn._saved_self, a) + + # b is not an output + self.assertIs(c.grad_fn._saved_other, b) + + # d is an output + self.assertEqual(d.grad_fn._saved_result, d) + self.assertIsNot(d.grad_fn._saved_result, d) + + c.sum().backward() + + with self.assertRaisesRegex(RuntimeError, "after they have already been freed"): + c.grad_fn._saved_self + + # a is left untouched + self.assertEqual(a, a_orig) + def index_perm_variable(shape, max_indices): if not isinstance(shape, tuple): diff --git a/torch/csrc/autograd/saved_variable.cpp b/torch/csrc/autograd/saved_variable.cpp index bbc6afe6750c3..0f51c102ea24d 100644 --- a/torch/csrc/autograd/saved_variable.cpp +++ b/torch/csrc/autograd/saved_variable.cpp @@ -14,7 +14,7 @@ namespace torch { namespace autograd { -SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_inplace_view) { +SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_inplace_on_view) { if (variable.defined()) { // Note [Inference tensor cannot be saved for backward] // Invariant: @@ -37,13 +37,27 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i "you can make a clone to get a normal tensor and use it in autograd.") was_default_constructed_ = false; - is_inplace_view_ = is_inplace_view; - version_counter_ = impl::version_counter(variable); - saved_version_ = version_counter_.current_version(); + is_inplace_on_view_ = is_inplace_on_view; + const auto& version_counter = impl::version_counter(variable); + saved_version_ = version_counter.current_version(); + + // If the variable is a leaf or is not an output, we can safely save the + // original variable without running the risk of reference cycles. + // 1. If the variable is not an output, its grad_fn has already been fully + // created and in particular will be a different Node than the one + // we are currently constructing (the one that owns this SavedVariable). + // 2. If the variable is a leaf, it only has weak reference to the grad_accumulator + // which cannot create a cycle. + if (!is_output || variable.is_leaf()) { + saved_original_ = true; + data_ = variable; + return; + } output_nr_ = variable.output_nr(); requires_grad_ = variable.requires_grad(); has_grad_fn_ = !variable.is_leaf(); + version_counter_ = version_counter; // These copies are all shared_ptr copies, so slightly more expensive. // Do them here instead of in the init list in case data is undefined. @@ -55,7 +69,7 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i grad_fn_ = variable.grad_fn(); } - if(is_output && is_inplace_view) { + if(is_output && is_inplace_on_view) { weak_grad_fn_ = variable.grad_fn(); } @@ -68,8 +82,8 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i } } -SavedVariable::SavedVariable(const c10::optional& variable, bool is_output, bool is_inplace_view) - : SavedVariable(variable.has_value() ? *variable : Variable(), is_output, is_inplace_view) {} +SavedVariable::SavedVariable(const c10::optional& variable, bool is_output, bool is_inplace_on_view) + : SavedVariable(variable.has_value() ? *variable : Variable(), is_output, is_inplace_on_view) {} Variable SavedVariable::unpack(std::shared_ptr saved_for) const { if (!data_.defined()) { @@ -77,15 +91,20 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { return Variable(); } - // We want grad_fn here to provide the most hlepful debug message to the user + // We want grad_fn here to provide the most helpful debug message to the user // if versions don't match - auto grad_fn = is_inplace_view_ ? weak_grad_fn_.lock() : grad_fn_; + auto grad_fn = saved_original_ ? data_.grad_fn() + : is_inplace_on_view_ ? weak_grad_fn_.lock() + : grad_fn_; if (has_grad_fn_ && !grad_fn) { TORCH_CHECK(saved_for,"No grad_fn for non-leaf saved variable"); grad_fn = std::move(saved_for); } - if (saved_version_ != version_counter_.current_version()) { + auto current_version = saved_original_ ? impl::version_counter(data_).current_version() + : version_counter_.current_version(); + + if (saved_version_ != current_version) { std::stringstream message; message << "one of the variables needed for gradient computation has been " "modified by an inplace operation: [" << data_.toString() << " " @@ -94,7 +113,7 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { message << ", which is output " << output_nr_ << " of " << grad_fn->name() << ","; } - message << " is at version " << version_counter_.current_version() + message << " is at version " << current_version << "; expected version " << saved_version_ << " instead."; if (!AnomalyMode::is_enabled()) { message << " Hint: enable anomaly detection to find the operation " @@ -109,6 +128,10 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { TORCH_CHECK(false, message.str()); } + if (saved_original_) { + return data_; + } + // NB: saved views are unpacked as normal Variables (not views) even though // they still share the same storage. This works only because we never call // in-place functions on unpacked variables. diff --git a/torch/csrc/autograd/saved_variable.h b/torch/csrc/autograd/saved_variable.h index 9f1f38ee7c357..b90dc43136190 100644 --- a/torch/csrc/autograd/saved_variable.h +++ b/torch/csrc/autograd/saved_variable.h @@ -21,8 +21,8 @@ TORCH_API extern const char* ERR_BACKWARD_TWICE; class TORCH_API SavedVariable { public: SavedVariable() = default; - SavedVariable(const Variable& variable, bool is_output, bool is_inplace_view=false); - SavedVariable(const c10::optional& variable, bool is_output, bool is_inplace_view=false); + SavedVariable(const Variable& variable, bool is_output, bool is_inplace_on_view=false); + SavedVariable(const c10::optional& variable, bool is_output, bool is_inplace_on_view=false); SavedVariable(SavedVariable&&) = default; SavedVariable& operator=(SavedVariable&&) = default; ~SavedVariable() { @@ -46,6 +46,16 @@ class TORCH_API SavedVariable { } private: + // This field contains either: + // 1. the variable to save + // 2. or its tensor_data. + // If storing the variable itself would create a circular reference, + // we fall into the second case and its metadata is also saved separately. + // In that case, the grad_fn must be passed in to the unpack function when + // reconstructing the Variable (except when we are doing an inplace operation on + // a view, see below). + // The field saved_orignal_ below reflects the two cases: its value is true + // in the first case and false in the second case. at::Tensor data_; // This field is used to store the forward AD gradients associated with @@ -60,6 +70,10 @@ class TORCH_API SavedVariable { std::shared_ptr grad_fn_; // Weak version of grad_fn_ that prevents leaks in rebase_history() for // inplace views. + // This variable is used when the user chooses to create a SavedVariable with + // is_inplace_on_view = true. + // In that case, the grad_fn passed in to the unpack function at unwrapping + // time is unused. std::weak_ptr weak_grad_fn_; std::weak_ptr grad_accumulator_; c10::VariableVersion version_counter_; @@ -69,6 +83,7 @@ class TORCH_API SavedVariable { bool was_default_constructed_ = true; bool requires_grad_ = false; bool has_grad_fn_ = false; - bool is_inplace_view_ = false; + bool is_inplace_on_view_ = false; + bool saved_original_ = false; }; }} // namespace torch::autograd From dab1e596525e7a1f767e4b207c13836fd009309e Mon Sep 17 00:00:00 2001 From: Victor Quach Date: Wed, 16 Jun 2021 16:39:47 -0700 Subject: [PATCH 166/305] Remove dead code in SavedVariable (#59838) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59838 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29069214 fbshipit-source-id: 5debf93a6c3d1c3d585efbe54438e8df92646d62 --- tools/autograd/gen_autograd_functions.py | 1 - torch/csrc/autograd/saved_variable.cpp | 44 +++++++++--------------- torch/csrc/autograd/saved_variable.h | 12 ------- 3 files changed, 16 insertions(+), 41 deletions(-) diff --git a/tools/autograd/gen_autograd_functions.py b/tools/autograd/gen_autograd_functions.py index 4f1cb242ac8bd..dee03398b2d5d 100644 --- a/tools/autograd/gen_autograd_functions.py +++ b/tools/autograd/gen_autograd_functions.py @@ -324,7 +324,6 @@ def save_var(var: SavedAttribute, is_output: bool) -> None: (type == BaseCType(scalarT) and is_output): saved_variables.append(f'SavedVariable {name}_;') release_variables.append(f'{name}_.reset_data();') - release_variables.append(f'{name}_.reset_grad_function();') ptr = 'shared_from_this()' if is_output else '' unpack.append(f'auto {name} = {name}_.unpack({ptr});') getter_definitions.append(GETTER_DEFINITION_SAVEDVAR.substitute( diff --git a/torch/csrc/autograd/saved_variable.cpp b/torch/csrc/autograd/saved_variable.cpp index 0f51c102ea24d..56bd9744d2ad2 100644 --- a/torch/csrc/autograd/saved_variable.cpp +++ b/torch/csrc/autograd/saved_variable.cpp @@ -37,7 +37,6 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i "you can make a clone to get a normal tensor and use it in autograd.") was_default_constructed_ = false; - is_inplace_on_view_ = is_inplace_on_view; const auto& version_counter = impl::version_counter(variable); saved_version_ = version_counter.current_version(); @@ -48,28 +47,24 @@ SavedVariable::SavedVariable(const Variable& variable, bool is_output, bool is_i // we are currently constructing (the one that owns this SavedVariable). // 2. If the variable is a leaf, it only has weak reference to the grad_accumulator // which cannot create a cycle. + // In those cases, we save the original variable and don't need further processing. if (!is_output || variable.is_leaf()) { saved_original_ = true; data_ = variable; return; } + // From now on, we can assume the variable is not a leaf and is an output. + + is_inplace_on_view_ = is_inplace_on_view; output_nr_ = variable.output_nr(); - requires_grad_ = variable.requires_grad(); - has_grad_fn_ = !variable.is_leaf(); version_counter_ = version_counter; // These copies are all shared_ptr copies, so slightly more expensive. // Do them here instead of in the init list in case data is undefined. data_ = variable.tensor_data(); - if (variable.is_leaf()) { - grad_accumulator_ = impl::grad_accumulator(variable); - } else if (!is_output) { - grad_fn_ = variable.grad_fn(); - } - - if(is_output && is_inplace_on_view) { + if(is_inplace_on_view) { weak_grad_fn_ = variable.grad_fn(); } @@ -94,9 +89,10 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { // We want grad_fn here to provide the most helpful debug message to the user // if versions don't match auto grad_fn = saved_original_ ? data_.grad_fn() - : is_inplace_on_view_ ? weak_grad_fn_.lock() - : grad_fn_; - if (has_grad_fn_ && !grad_fn) { + : is_inplace_on_view_ ? weak_grad_fn_.lock() + : nullptr; + + if (!saved_original_ && !grad_fn) { TORCH_CHECK(saved_for,"No grad_fn for non-leaf saved variable"); grad_fn = std::move(saved_for); } @@ -128,30 +124,22 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { TORCH_CHECK(false, message.str()); } + // The version counter is correct. If we have the original variable, we simply return it + if (saved_original_) { return data_; } + // From now on, we can assume the variable is not a leaf and is an output. + // Additionnally, because the variable is not a leaf, we have its grad_fn + // (computed above) and need to attach it to the returned tensor. + // NB: saved views are unpacked as normal Variables (not views) even though // they still share the same storage. This works only because we never call // in-place functions on unpacked variables. - Variable var; - if (grad_fn) { - var = make_variable(data_, Edge(std::move(grad_fn), output_nr_)); - } else { - var = make_variable(data_, requires_grad_); - } + Variable var = make_variable(data_, Edge(std::move(grad_fn), output_nr_)); impl::set_version_counter(var, saved_version_); - // If a Variable is a leaf (no grad_fn saved), and it requires_grad, then we - // should have saved the grad accumulator. Even if the Variable no longer - // alive, the accumulator should be kept alive by the references in the - // graph). - if (requires_grad_ && !var.grad_fn() && grad_accumulator_.expired()) { - TORCH_CHECK(false, "No grad accumulator for a saved leaf!"); - } - impl::set_grad_accumulator(var, grad_accumulator_); - // NB: var here is never a view so there is no need to make anything special // for the case where the saved Tensor was a view. This whole argument relies // on the fact that the Tensor returned by this function is never diff --git a/torch/csrc/autograd/saved_variable.h b/torch/csrc/autograd/saved_variable.h index b90dc43136190..82184c6f89be9 100644 --- a/torch/csrc/autograd/saved_variable.h +++ b/torch/csrc/autograd/saved_variable.h @@ -41,10 +41,6 @@ class TORCH_API SavedVariable { return data_.reset(); } - void reset_grad_function() { - grad_fn_.reset(); - } - private: // This field contains either: // 1. the variable to save @@ -63,11 +59,6 @@ class TORCH_API SavedVariable { // either the saved Tensor or the unpacked Tensor. See note [ Using ForwardGrad ] std::shared_ptr fw_grad_; - // The gradient function associated with this node. If has_grad_fn - // is false, then this is a leaf node. Note that the grad_fn is not saved if - // it would create a circular reference. In that case, the grad_fn must be - // passed in to the unpack function when reconstructing the Variable. - std::shared_ptr grad_fn_; // Weak version of grad_fn_ that prevents leaks in rebase_history() for // inplace views. // This variable is used when the user chooses to create a SavedVariable with @@ -75,14 +66,11 @@ class TORCH_API SavedVariable { // In that case, the grad_fn passed in to the unpack function at unwrapping // time is unused. std::weak_ptr weak_grad_fn_; - std::weak_ptr grad_accumulator_; c10::VariableVersion version_counter_; uint32_t saved_version_ = 0; uint32_t output_nr_ = 0; bool was_default_constructed_ = true; - bool requires_grad_ = false; - bool has_grad_fn_ = false; bool is_inplace_on_view_ = false; bool saved_original_ = false; }; From 0bf1260795445622adc3ab58e3c35e18baf24fee Mon Sep 17 00:00:00 2001 From: Edward Yang Date: Wed, 16 Jun 2021 17:08:59 -0700 Subject: [PATCH 167/305] Fix Python 3.8 expecttest machinery again, this time for good. (#60044) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60044 In #59709 I attempted to fix the expecttest machinery to work in Python 3.8. However, I noticed that it would fail to do substitutions in this case: ``` self.assertExpectedInline( foo(), """bar""" ) ``` This is because the triple quoted string is not on the same line as the backtrace line number (at the very beginning), and for safety reasons the preexisting regex refused to search beyond the first line. This wasn't a big deal prior to Python 3.8 because the flipped version of the regex simply required the triple quoted string to be flush with the end of the statement (which it typically was!) But it is a big deal now that we only have the start of the statement. I couldn't think of a way to fix this in the current model, so I decided to call in the big guns. Instead of trying to do the regex with only the start xor end line number, I now require you provide BOTH line numbers, and we will only regex within this range. The way we compute these line numbers is by parsing the Python test file with ast, and then searching through statements until we find one that is consistent with the line number reported by the backtrace. If we don't find anything, we conservatively assume that the string lies exactly in the backtrace (and you'll probably fail the substitution in that case.) The resulting code is quite a lot simpler (no more reversed regex) and hopefully more robust, although I suppose we are going to have to do some field testing. Signed-off-by: Edward Z. Yang Test Plan: Imported from OSS Reviewed By: walterddr Differential Revision: D29146943 Pulled By: ezyang fbshipit-source-id: 2c24abc3acd4275c5b3a8f222d2a60cbad5e8c78 --- test/test_expecttest.py | 91 +++++++---------------- torch/testing/_internal/expecttest.py | 100 ++++++++++++-------------- 2 files changed, 70 insertions(+), 121 deletions(-) diff --git a/test/test_expecttest.py b/test/test_expecttest.py index dbf24325d9835..d193dbb2dc907 100644 --- a/test/test_expecttest.py +++ b/test/test_expecttest.py @@ -30,8 +30,8 @@ def nth_line_ref(src, lineno): return len("\n".join(xs)) self.assertEqual(expecttest.nth_line(t, lineno), nth_line_ref(t, lineno)) - @hypothesis.given(text(string.printable), booleans(), sampled_from(['"', "'"]), booleans()) - def test_replace_string_literal_roundtrip(self, t, raw, quote, lineno_at_start): + @hypothesis.given(text(string.printable), booleans(), sampled_from(['"', "'"])) + def test_replace_string_literal_roundtrip(self, t, raw, quote): if raw: hypothesis.assume(expecttest.ok_for_raw_triple_quoted_string(t, quote=quote)) prog = """\ @@ -40,7 +40,7 @@ def test_replace_string_literal_roundtrip(self, t, raw, quote, lineno_at_start): r3 = {r}{quote}placeholder3{quote} """.format(r='r' if raw else '', quote=quote * 3) new_prog = expecttest.replace_string_literal( - textwrap.dedent(prog), 2, t, lineno_at_start=lineno_at_start)[0] + textwrap.dedent(prog), 2, 2, t)[0] ns : Dict[str, Any] = {} exec(new_prog, ns) msg = "program was:\n{}".format(new_prog) @@ -48,7 +48,7 @@ def test_replace_string_literal_roundtrip(self, t, raw, quote, lineno_at_start): self.assertEqual(ns['r2'], expecttest.normalize_nl(t), msg=msg) # noqa: F821 self.assertEqual(ns['r3'], 'placeholder3', msg=msg) # noqa: F821 - def test_sample_lineno_at_end(self): + def test_sample_lineno(self): prog = r""" single_single('''0''') single_multi('''1''') @@ -65,72 +65,27 @@ def test_sample_lineno_at_end(self): multi_multi_more('''\ 6 ''') +different_indent( + RuntimeError, + '''7''' +) """ - # NB: These are the end of the statements, not beginning - # TODO: Test other permutations of these edits - edits = [(2, "a"), - (3, "b\n"), - (6, "c"), - (10, "d\n"), - (13, "e\n"), - (16, "f\ng\n")] + edits = [(2, 2, "a"), + (3, 3, "b\n"), + (4, 6, "c"), + (7, 10, "d\n"), + (11, 13, "e\n"), + (14, 16, "f\ng\n"), + (17, 20, "h")] history = expecttest.EditHistory() fn = 'not_a_real_file.py' - for lineno, actual in edits: - lineno = history.adjust_lineno(fn, lineno) + for start_lineno, end_lineno, actual in edits: + start_lineno = history.adjust_lineno(fn, start_lineno) + end_lineno = history.adjust_lineno(fn, end_lineno) prog, delta = expecttest.replace_string_literal( - prog, lineno, actual, lineno_at_start=False) - history.record_edit(fn, lineno, delta) - self.assertExpectedInline(prog, r""" -single_single('''a''') -single_multi('''\ -b -''') -multi_single('''c''') -multi_multi_less('''\ -d -''') -multi_multi_same('''\ -e -''') -multi_multi_more('''\ -f -g -''') -""") - - def test_sample_lineno_at_start(self): - prog = r""" -single_single('''0''') -single_multi('''1''') -multi_single('''\ -2 -''') -multi_multi_less('''\ -3 -4 -''') -multi_multi_same('''\ -5 -''') -multi_multi_more('''\ -6 -''') -""" - # NB: These are the beginning of the statements - edits = [(2, "a"), - (3, "b\n"), - (4, "c"), - (7, "d\n"), - (11, "e\n"), - (14, "f\ng\n")] - history = expecttest.EditHistory() - fn = 'not_a_real_file.py' - for lineno, actual in edits: - lineno = history.adjust_lineno(fn, lineno) - prog, delta = expecttest.replace_string_literal( - prog, lineno, actual, lineno_at_start=True) - history.record_edit(fn, lineno, delta) + prog, start_lineno, end_lineno, actual) + # NB: it doesn't really matter start/end you record edit at + history.record_edit(fn, start_lineno, delta) self.assertExpectedInline(prog, r""" single_single('''a''') single_multi('''\ @@ -147,6 +102,10 @@ def test_sample_lineno_at_start(self): f g ''') +different_indent( + RuntimeError, + '''h''' +) """) def test_lineno_assumptions(self): diff --git a/torch/testing/_internal/expecttest.py b/torch/testing/_internal/expecttest.py index 7fff670fe9269..900274e063c3b 100644 --- a/torch/testing/_internal/expecttest.py +++ b/torch/testing/_internal/expecttest.py @@ -4,6 +4,7 @@ import os import string import sys +import ast from typing import Tuple @@ -137,7 +138,6 @@ def ok_for_raw_triple_quoted_string(s, quote): RE_EXPECT = re.compile( ( - r"^(?P[^\n]*?)" r"(?Pr?)" r"(?P'''|" r'""")' r"(?P.*?)" @@ -147,17 +147,8 @@ def ok_for_raw_triple_quoted_string(s, quote): ) -# This operates on the REVERSED string (that's why suffix is first) -RE_REVERSED_EXPECT = \ - re.compile(r"^(?P[^\n]*?)" - r"(?P'''|" r'""")' - r"(?P.*?)" - r"(?P=quote)" - r"(?Pr?)", re.DOTALL) - - -def replace_string_literal(src : str, lineno : int, - new_string : str, *, lineno_at_start: bool) -> Tuple[str, int]: +def replace_string_literal(src : str, start_lineno : int, end_lineno : int, + new_string : str) -> Tuple[str, int]: r""" Replace a triple quoted string literal with new contents. Only handles printable ASCII correctly at the moment. This @@ -168,9 +159,9 @@ def replace_string_literal(src : str, lineno : int, Returns a tuple of the replaced string, as well as a delta of number of lines added/removed. - >>> replace_string_literal("'''arf'''", 1, "barf", lineno_at_start=False) + >>> replace_string_literal("'''arf'''", 1, 1, "barf") ("'''barf'''", 0) - >>> r = replace_string_literal(" moo = '''arf'''", 1, "'a'\n\\b\n", lineno_at_start=False) + >>> r = replace_string_literal(" moo = '''arf'''", 1, 1, "'a'\n\\b\n") >>> print(r[0]) moo = '''\ 'a' @@ -178,9 +169,9 @@ def replace_string_literal(src : str, lineno : int, ''' >>> r[1] 3 - >>> replace_string_literal(" moo = '''\\\narf'''", 2, "'a'\n\\b\n", lineno_at_start=False)[1] + >>> replace_string_literal(" moo = '''\\\narf'''", 1, 2, "'a'\n\\b\n")[1] 2 - >>> print(replace_string_literal(" f('''\"\"\"''')", 1, "a ''' b", lineno_at_start=False)[0]) + >>> print(replace_string_literal(" f('''\"\"\"''')", 1, 1, "a ''' b")[0]) f('''a \'\'\' b''') """ # Haven't implemented correct escaping for non-printable characters @@ -192,7 +183,12 @@ def replace_string_literal(src : str, lineno : int, if delta[0] > 0: delta[0] += 1 # handle the extra \\\n - def compute_raw_new_body_and_adjust_delta(m): + assert start_lineno <= end_lineno + start = nth_line(src, start_lineno) + end = nth_eol(src, end_lineno) + assert start <= end + + def replace(m): s = new_string raw = m.group('raw') == 'r' if not raw or not ok_for_raw_triple_quoted_string(s, quote=m.group('quote')[0]): @@ -205,39 +201,13 @@ def compute_raw_new_body_and_adjust_delta(m): new_body = "\\\n" + s if "\n" in s and not raw else s delta[0] -= m.group('body').count("\n") - return raw, new_body - - if lineno_at_start: - i = nth_line(src, lineno) - - # i points to the start of the string - def replace(m): - raw, new_body = compute_raw_new_body_and_adjust_delta(m) - return ''.join([m.group('prefix'), - 'r' if raw else '', - m.group('quote'), - new_body, - m.group('quote'), - ]) - - return (src[:i] + RE_EXPECT.sub(replace, src[i:], count=1), delta[0]) - else: - i = nth_eol(src, lineno) - - # i points to the END of the string. Do some funny - # business with reversing the string to do the replace - def replace(m): - raw, new_body = compute_raw_new_body_and_adjust_delta(m) - return ''.join([m.group('suffix'), - m.group('quote'), - new_body[::-1], - m.group('quote'), - 'r' if raw else '', - ]) + return ''.join(['r' if raw else '', + m.group('quote'), + new_body, + m.group('quote'), + ]) - # Having to do this in reverse is very irritating, but it's the - # only way to make the non-greedy matches work correctly. - return (RE_REVERSED_EXPECT.sub(replace, src[:i][::-1], count=1)[::-1] + src[i:], delta[0]) + return (src[:start] + RE_EXPECT.sub(replace, src[start:end], count=1) + src[end:], delta[0]) class TestCase(unittest.TestCase): @@ -263,15 +233,35 @@ def assertExpectedInline(self, actual, expect, skip=0): print("Accepting new output for {} at {}:{}".format(self.id(), fn, lineno)) with open(fn, 'r+') as f: old = f.read() + old_ast = ast.parse(old) - # compute the change in lineno + # NB: it's only the traceback line numbers that are wrong; + # we reread the file every time we write to it, so AST's + # line numbers are correct lineno = EDIT_HISTORY.adjust_lineno(fn, lineno) - new, delta = replace_string_literal( - old, lineno, actual, - lineno_at_start=LINENO_AT_START - ) - assert old != new, f"Failed to substitute string at {fn}:{lineno}; did you use triple quotes?" + # Conservative assumption to start + start_lineno = lineno + end_lineno = lineno + # Try to give a more accurate bounds based on AST + # NB: this walk is in no specified order (in practice it's + # breadth first) + for n in ast.walk(old_ast): + if isinstance(n, ast.Expr): + if hasattr(n, 'end_lineno'): + assert LINENO_AT_START + if n.lineno == start_lineno: + end_lineno = n.end_lineno # type: ignore[attr-defined] + else: + if n.lineno == end_lineno: + start_lineno = n.lineno + + new, delta = replace_string_literal(old, start_lineno, end_lineno, actual) + + assert old != new, f"Failed to substitute string at {fn}:{lineno}; did you use triple quotes? " \ + "If this is unexpected, please file a bug report at " \ + "https://github.com/pytorch/pytorch/issues/new?labels=module:%20expecttest " \ + f"with the contents of the source file near {fn}:{lineno}" # Only write the backup file the first time we hit the # file From 64aec8d2cac444319275b0f3a59053d19b86ec73 Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Wed, 16 Jun 2021 17:15:59 -0700 Subject: [PATCH 168/305] [testing] OpInfoHelper tool (#58698) Summary: Fixes: https://github.com/pytorch/pytorch/issues/57577 Usage: Add OpInfo entry to `common_methods_invocations` with `dtypes=_DYNAMIC_DYTPES` Eg. ``` OpInfo('atan2', dtypes=_DYNAMIC_DTYPES, sample_inputs_func=sample_inputs_atan2,) ``` Run the helper with `python -m torch.testing._internal.opinfo_helper` Output ``` OpInfo(atan2, # hint: all_types + (torch.bool,), dtypes=[torch.float32, torch.float64, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64, torch.bool], # hint: all_types + (torch.bool, torch.bfloat16, torch.float16), dtypesIfCUDA=[torch.float32, torch.float64, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64, torch.bool, torch.bfloat16, torch.float16], sample_inputs_func=sample_inputs_atan2) ``` Output without CUDA (run with `$ CUDA_VISIBLE_DEVICES=-1 python -m torch.testing._internal.opinfo_helper`) ``` UserWarning: WARNING: CUDA is not available, information pertaining to CUDA could be wrong warnings.warn("WARNING: CUDA is not available, information pertaining to CUDA could be wrong") OpInfo(atan2, # hint: all_types + (torch.bool,), dtypes=[torch.float32, torch.float64, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64, torch.bool], sample_inputs_func=sample_inputs_atan2) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/58698 Reviewed By: H-Huang Differential Revision: D29160668 Pulled By: mruberry fbshipit-source-id: 707370a83b451b02ad2fe539775c8c50ecf90be8 --- test/test_ops.py | 20 ++- test/test_testing.py | 19 +++ torch/testing/_core.py | 2 + .../_internal/common_methods_invocations.py | 17 ++- torch/testing/_internal/opinfo_helper.py | 138 ++++++++++++++++++ 5 files changed, 193 insertions(+), 3 deletions(-) create mode 100644 torch/testing/_internal/opinfo_helper.py diff --git a/test/test_ops.py b/test/test_ops.py index da6f8462e1e17..41efa551c46cb 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from functools import partial, wraps import warnings @@ -7,7 +8,7 @@ (FileCheck, floating_and_complex_types_and) from torch.testing._internal.common_utils import \ (TestCase, is_iterable_of_tensors, run_tests, IS_SANDCASTLE, clone_input_helper, make_tensor, - gradcheck, gradgradcheck) + gradcheck, gradgradcheck, IS_PYTORCH_CI) from torch.testing._internal.common_methods_invocations import \ (op_db,) from torch.testing._internal.common_device_type import \ @@ -17,13 +18,28 @@ from torch.testing._internal.jit_metaprogramming_utils import create_script_fn, create_traced_fn, \ check_alias_annotation from torch.testing._internal.jit_utils import disable_autodiff_subgraph_inlining -from collections.abc import Sequence +import torch.testing._internal.opinfo_helper as opinfo_helper # Tests that apply to all operators class TestOpInfo(TestCase): exact_dtype = True + @classmethod + def tearDownClass(cls): + super().tearDownClass() + + if IS_PYTORCH_CI: + err_msg = ("The operator(s) below is(are) using dynamic_dtypes in the OpInfo entries." + "This is OK for testing, but be sure to set the dtypes manually before landing your PR!") + # Assure no opinfo entry has dynamic_dtypes + filtered_ops = list(filter(opinfo_helper.is_dynamic_dtype_set, op_db)) + for op in filtered_ops: + fmt_str = opinfo_helper.str_format_dynamic_dtype(op) + err_msg += "\n" + fmt_str + + assert len(filtered_ops) == 0, err_msg + # Verifies that ops have their unsupported dtypes # registered correctly by testing that each claimed unsupported dtype # throws a runtime error diff --git a/test/test_testing.py b/test/test_testing.py index 9fdcd88a884e0..aefe037841ef9 100644 --- a/test/test_testing.py +++ b/test/test_testing.py @@ -17,6 +17,8 @@ (PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY, PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, dtypes, get_device_type_test_bases, instantiate_device_type_tests, onlyCUDA, onlyOnCPUAndCUDA, deviceCountAtLeast) +from torch.testing._internal.common_methods_invocations import op_db +import torch.testing._internal.opinfo_helper as opinfo_helper from torch.testing._asserts import UsageError # For testing TestCase methods and torch.testing functions @@ -595,6 +597,23 @@ def test_trivial_passing_test_case_on_cpu_cuda(self, device): # we are currently disabling CUDA early termination for distributed tests. self.assertIn('Ran 2 test', stderr) + @onlyOnCPUAndCUDA + def test_get_supported_dtypes(self, device): + # Test the `get_supported_dtypes` helper function. + # We acquire the dtypes for few Ops dynamically and verify them against + # the correct statically described values. + ops_to_test = list(filter(lambda op: op.name in ['atan2', 'topk', 'xlogy'], op_db)) + + for op in ops_to_test: + dynamic_dtypes = opinfo_helper.get_supported_dtypes(op.op, op.sample_inputs_func, self.device_type) + dynamic_dispatch = opinfo_helper.dtypes_dispatch_hint(dynamic_dtypes) + if self.device_type == 'cpu': + dtypes = op.dtypesIfCPU + else: # device_type ='cuda' + dtypes = op.dtypesIfCUDA + + self.assertTrue(set(dtypes) == set(dynamic_dtypes)) + self.assertTrue(set(dtypes) == set(dynamic_dispatch.dispatch_fn())) instantiate_device_type_tests(TestTesting, globals()) diff --git a/torch/testing/_core.py b/torch/testing/_core.py index 1522c16aa6ced..d58dbe71041b6 100644 --- a/torch/testing/_core.py +++ b/torch/testing/_core.py @@ -17,12 +17,14 @@ "all_types_and", "all_types_and_complex", "all_types_and_complex_and", + "all_types_and_half", "assert_allclose", "complex_types", "floating_and_complex_types", "floating_and_complex_types_and", "floating_types", "floating_types_and", + "floating_types_and_half", "get_all_complex_dtypes", "get_all_dtypes", "get_all_device_types", diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index eb60f1f041f49..60ca5bd4b8bad 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -33,6 +33,7 @@ TEST_WITH_ROCM, IS_WINDOWS, IS_MACOS, make_tensor, TEST_SCIPY, torch_to_numpy_dtype_dict, slowTest, TEST_WITH_ASAN, GRADCHECK_NONDET_TOL,) +import torch.testing._internal.opinfo_helper as opinfo_helper from setuptools import distutils @@ -232,14 +233,28 @@ def __init__(self, test_conjugated_samples=True, ): + dtypes_args = (dtypes, dtypesIfCPU, dtypesIfCUDA, dtypesIfROCM) # Validates the dtypes are generated from the dispatch-related functions - for dtype_list in (dtypes, dtypesIfCPU, dtypesIfCUDA, dtypesIfROCM): + for dtype_list in dtypes_args: assert isinstance(dtype_list, (_dispatch_dtypes, type(None))) self.name = name self.aten_name = aten_name if aten_name is not None else name self.variant_test_name = variant_test_name + # Attribute to verify dynamic_dtypes are used. + self.dynamic_dtypes = any(map(lambda dtypes: isinstance( + dtypes, opinfo_helper._dynamic_dispatch_dtypes), dtypes_args)) + + if self.dynamic_dtypes: + # Make sure `dtyesIfCUDA` is dynamic, if dynamic dispatch is used for CPU + # This is because, below we set dtypesIfCUDA to dtypes if they are None. + assert isinstance(dtypesIfCUDA, opinfo_helper._dynamic_dispatch_dtypes), \ + (f"To use dynamic dypes for operator {name}, " + "acquire the dtypes dynamically for argument `dtypesIfCUDA`." + "This is to ensure that CUDA dtypes are acquired correctly as they" + "differ from CPU dtypes occasionally") + self.dtypes = set(dtypes) self.dtypesIfCPU = set(dtypesIfCPU) if dtypesIfCPU is not None else self.dtypes self.dtypesIfCUDA = set(dtypesIfCUDA) if dtypesIfCUDA is not None else self.dtypes diff --git a/torch/testing/_internal/opinfo_helper.py b/torch/testing/_internal/opinfo_helper.py new file mode 100644 index 0000000000000..5129af4f99e34 --- /dev/null +++ b/torch/testing/_internal/opinfo_helper.py @@ -0,0 +1,138 @@ +import collections +import warnings +from functools import partial + +import torch +from torch.testing._internal.common_cuda import (TEST_CUDA) +from torch.testing._core import _dispatch_dtypes +from torch.testing import (all_types_and_complex_and, + all_types_and_complex, + all_types_and_half, + all_types, + complex_types, + floating_and_complex_types, + floating_types_and_half, + floating_types, + integral_types, + floating_types_and, + floating_and_complex_types_and, + integral_types_and, + all_types_and, + ) + +COMPLETE_DTYPES_DISPATCH = ( + all_types, + all_types_and_complex, + all_types_and_half, + floating_types, + floating_and_complex_types, + floating_types_and_half, + integral_types, + complex_types, +) + +EXTENSIBLE_DTYPE_DISPATCH = ( + all_types_and_complex_and, + floating_types_and, + floating_and_complex_types_and, + integral_types_and, + all_types_and, +) + +# Better way to acquire devices? +DEVICES = ['cpu'] + (['cuda'] if TEST_CUDA else []) + +class _dynamic_dispatch_dtypes(_dispatch_dtypes): + # Class to tag the dynamically generated types. + pass + + +def get_supported_dtypes(op, sample_inputs_fn, device_type): + # Returns the supported dtypes for the given operator and device_type pair. + assert device_type in ['cpu', 'cuda'] + if not TEST_CUDA and device_type == 'cuda': + warnings.warn("WARNING: CUDA is not available, empty_dtypes dispatch will be returned!") + return _dynamic_dispatch_dtypes(()) + + supported_dtypes = set() + for dtype in all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half): + try: + samples = sample_inputs_fn(op, device_type, dtype, False) + except RuntimeError: + # If `sample_inputs_fn` doesn't support sampling for a given + # `dtype`, we assume that the `dtype` is not supported. + # We raise a warning, so that user knows that this was the case + # and can investigate if there was an issue with the `sample_inputs_fn`. + warnings.warn(f"WARNING: Unable to generate sample for device:{device_type} and dtype:{dtype}") + continue + + # We assume the dtype is supported + # only if all samples pass for the given dtype. + supported = True + for sample in samples: + try: + op(sample.input, *sample.args, **sample.kwargs) + except RuntimeError as re: + # dtype is not supported + supported = False + break + + if supported: + supported_dtypes.add(dtype) + + return _dynamic_dispatch_dtypes(supported_dtypes) + + +def dtypes_dispatch_hint(dtypes): + # Function returns the appropriate dispatch function (from COMPLETE_DTYPES_DISPATCH and EXTENSIBLE_DTYPE_DISPATCH) + # and its string representation for the passed `dtypes`. + return_type = collections.namedtuple('return_type', 'dispatch_fn dispatch_fn_str') + + # CUDA is not available, dtypes will be empty. + if len(dtypes) == 0: + return return_type((), str(tuple())) + + set_dtypes = set(dtypes) + for dispatch in COMPLETE_DTYPES_DISPATCH: + # Short circuit if we get an exact match. + if set(dispatch()) == set_dtypes: + return return_type(dispatch, dispatch.__name__ + "()") + + chosen_dispatch = None + chosen_dispatch_score = 0. + for dispatch in EXTENSIBLE_DTYPE_DISPATCH: + dispatch_dtypes = set(dispatch()) + if not dispatch_dtypes.issubset(set_dtypes): + continue + + score = len(dispatch_dtypes) + if score > chosen_dispatch_score: + chosen_dispatch_score = score + chosen_dispatch = dispatch + + # If user passed dtypes which are lower than the lowest + # dispatch type available (not likely but possible in code path). + if chosen_dispatch is None: + return return_type((), str(dtypes)) + + return return_type(partial(dispatch, *tuple(set(dtypes) - set(dispatch()))), + dispatch.__name__ + str(tuple(set(dtypes) - set(dispatch())))) + + +def is_dynamic_dtype_set(op): + # Detect if the OpInfo entry acquired dtypes dynamically + # using `get_supported_dtypes`. + return op.dynamic_dtypes + + +def str_format_dynamic_dtype(op): + fmt_str = """ + OpInfo({name}, + dtypes={dtypesIfCPU}, + dtypesIfCUDA={dtypesIfCUDA}, + ) + """.format(name=op.name, + dtypesIfCPU=dtypes_dispatch_hint(op.dtypesIfCPU).dispatch_fn_str, + dtypesIfCUDA=dtypes_dispatch_hint(op.dtypesIfCUDA).dispatch_fn_str) + + return fmt_str From 5f017e91b890c897107e2a479d274fcda5a3024b Mon Sep 17 00:00:00 2001 From: cyy Date: Wed, 16 Jun 2021 17:21:16 -0700 Subject: [PATCH 169/305] don't use moved field in the second lambda (#59914) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59914 Reviewed By: H-Huang Differential Revision: D29147018 Pulled By: ezyang fbshipit-source-id: 04fe52fb8cf3cc8f3a538a2dddb13c52cf558549 --- torch/custom_class.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/torch/custom_class.h b/torch/custom_class.h index cbbbae1a38697..f39695a89881e 100644 --- a/torch/custom_class.h +++ b/torch/custom_class.h @@ -223,12 +223,12 @@ class class_ : public ::torch::detail::class_base { /// Property registration API for properties with read-write access. template class_& def_readwrite(const std::string& name, T CurClass::*field) { - auto getter_func = - [field = std::move(field)](const c10::intrusive_ptr& self) { - return self.get()->*field; - }; + auto getter_func = [field = + field](const c10::intrusive_ptr& self) { + return self.get()->*field; + }; - auto setter_func = [field = std::move(field)]( + auto setter_func = [field = field]( const c10::intrusive_ptr& self, T value) { self.get()->*field = value; }; From f65793507d7831c001010d4a9fbd7f90f89603cc Mon Sep 17 00:00:00 2001 From: Jordan Fix Date: Wed, 16 Jun 2021 17:23:23 -0700 Subject: [PATCH 170/305] [fx][Transformer] Add override for call_function (#60057) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60057 This ensures that if a function was `wrap`'d before symbolic tracing + being passed into the transformer then it will still be wrapped. Test Plan: Added test to `test_fx.py` Reviewed By: jamesr66a Differential Revision: D29151191 fbshipit-source-id: 93560be59505bdcfe8d4f013e21d4719788afd59 --- test/test_fx.py | 18 ++++++++++++++++++ torch/fx/interpreter.py | 4 ++++ 2 files changed, 22 insertions(+) diff --git a/test/test_fx.py b/test/test_fx.py index 636297ba4ca11..853d01a28d415 100644 --- a/test/test_fx.py +++ b/test/test_fx.py @@ -309,6 +309,24 @@ def to_trace(y): self.assertIs(wrapped_via_decorator, real_wrapped_via_decorator) self.assertFalse(hasattr(wrapped_via_decorator, "__fx_already_patched")) + def test_wrapped_via_decorator_and_transformed(self): + self.assertEqual(wrapped_via_decorator(0), 1) + + def to_trace(y): + return wrapped_via_decorator(y) + + m = symbolic_trace(to_trace) + self.assertIn('wrapped_via_decorator', m.code) + self.assertEqual(m(0), 1) + self.assertIs(wrapped_via_decorator, real_wrapped_via_decorator) + self.assertFalse(hasattr(wrapped_via_decorator, "__fx_already_patched")) + + transformed = torch.fx.Transformer(m).transform() + self.assertIn('wrapped_via_decorator', transformed.code) + self.assertEqual(transformed(0), 1) + self.assertIs(wrapped_via_decorator, real_wrapped_via_decorator) + self.assertFalse(hasattr(wrapped_via_decorator, "__fx_already_patched")) + def test_wrap_with_submodule(self): class M(torch.nn.Module): diff --git a/torch/fx/interpreter.py b/torch/fx/interpreter.py index b9a7cfceb985a..20dcf62e0c3cb 100644 --- a/torch/fx/interpreter.py +++ b/torch/fx/interpreter.py @@ -409,6 +409,10 @@ def call_module(self, target : 'Target', args : Tuple[Argument, ...], kwargs : D submod = self.fetch_attr(target) return self.tracer.call_module(submod, submod.forward, args, kwargs) + def call_function(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any: + # Override so that functions that were wrapped are still wrapped. + return self.tracer.create_proxy('call_function', target, args, kwargs) + def transform(self) -> GraphModule: """ Transform ``self.module`` and return the transformed From 3288c9d304574de82d01e3f8200197c3080220f4 Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Wed, 16 Jun 2021 17:43:11 -0700 Subject: [PATCH 171/305] [numpy] mvlgamma: int -> float promotion (#59934) Summary: Reference: https://github.com/pytorch/pytorch/issues/42515 Last int->float promotion as per the tracker! Pull Request resolved: https://github.com/pytorch/pytorch/pull/59934 Reviewed By: H-Huang Differential Revision: D29160008 Pulled By: mruberry fbshipit-source-id: 389a5a7683e0c00d474da913012768bf2a212ef0 --- aten/src/ATen/native/UnaryOps.cpp | 9 ++++++--- .../_internal/common_methods_invocations.py | 18 +++++++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index 4ba59ef3d05dd..5a01df1f9b87d 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -645,8 +645,6 @@ constexpr double QUARTER = 0.25; } static inline void mvlgamma_check(const Tensor& self, int64_t p) { - TORCH_CHECK(at::isFloatingType(self.scalar_type()), - "mvlgamma is not implemented for ", self.scalar_type()); TORCH_CHECK((self > HALF * (p - 1)).all().item(), "All elements must be greater than (p-1)/2"); TORCH_CHECK(p >= 1, "p has to be greater than or equal to 1"); @@ -654,11 +652,16 @@ static inline void mvlgamma_check(const Tensor& self, int64_t p) { Tensor mvlgamma(const Tensor& self, int64_t p) { mvlgamma_check(self, p); + auto dtype = c10::scalarTypeToTypeMeta(self.scalar_type()); + if (at::isIntegralType(self.scalar_type(), /*include_bool=*/true)) { + // int -> float promotion + dtype = c10::get_default_dtype(); + } Tensor args = native::arange( -p * HALF + HALF, HALF, HALF, - optTypeMetaToScalarType(self.options().dtype_opt()), + optTypeMetaToScalarType(dtype), self.options().layout_opt(), self.options().device_opt(), self.options().pinned_memory_opt()); diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 60ca5bd4b8bad..179d6c3e51d56 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -3451,6 +3451,9 @@ def compute_min_val(p): def generator(): for shape, n in product(tensor_shapes, ns): min_val = compute_min_val(n) + if not dtype.is_floating_point: + # Round-up minimum value for integral dtypes + min_val += 1 yield SampleInput(make_arg(shape, low=min_val), args=(n,)) return list(generator()) @@ -3488,10 +3491,11 @@ def __init__(self, variant_test_name, domain, skips, sample_kwargs): variant_test_name=variant_test_name, domain=domain, decorators=(precisionOverride({torch.float16: 5e-2}),), - dtypes=floating_types(), - dtypesIfCUDA=floating_types_and(torch.half), + dtypes=all_types_and(torch.bool), + dtypesIfCUDA=all_types_and(torch.bool, torch.half), sample_inputs_func=sample_inputs_mvlgamma, supports_out=False, + safe_casts_outputs=True, skips=skips, sample_kwargs=sample_kwargs) @@ -5934,19 +5938,23 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_forward_ad=True, sample_inputs_func=sample_inputs_mode,), MvlGammaInfo(variant_test_name='mvlgamma_p_1', - domain=(1e-4, float('inf')), + domain=(1, float('inf')), skips=skips_mvlgamma(), sample_kwargs=lambda device, dtype, input: ({'p': 1}, {'d': 1})), MvlGammaInfo(variant_test_name='mvlgamma_p_3', - domain=(1.1, float('inf')), + domain=(2, float('inf')), skips=skips_mvlgamma(skip_redundant=True) + ( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', dtypes=(torch.float16,)), + # bool can't represent the low value from the domain + SkipInfo('TestOpInfo', 'test_supported_dtypes', dtypes=(torch.bool,)), ), sample_kwargs=lambda device, dtype, input: ({'p': 3}, {'d': 3})), MvlGammaInfo(variant_test_name='mvlgamma_p_5', - domain=(2.1, float('inf')), + domain=(3, float('inf')), skips=skips_mvlgamma(skip_redundant=True) + ( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', dtypes=(torch.float16,)), + # bool can't represent the low value from the domain + SkipInfo('TestOpInfo', 'test_supported_dtypes', dtypes=(torch.bool,)), ), sample_kwargs=lambda device, dtype, input: ({'p': 5}, {'d': 5})), OpInfo('ne', From c458bb985e620376e60493ab83e426f098926634 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Wed, 16 Jun 2021 17:47:55 -0700 Subject: [PATCH 172/305] make it easier to grep for unary/binary op kernels (#60128) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60128 Test Plan: Imported from OSS Reviewed By: wenleix Differential Revision: D29175499 Pulled By: bdhirsh fbshipit-source-id: 1838900276e0b956edf25cdddcff438ff685a50e --- aten/src/ATen/native/BinaryOps.cpp | 36 ++++++------ aten/src/ATen/native/UnaryOps.cpp | 94 +++++++++++++++--------------- 2 files changed, 65 insertions(+), 65 deletions(-) diff --git a/aten/src/ATen/native/BinaryOps.cpp b/aten/src/ATen/native/BinaryOps.cpp index ec42935f26ea4..59dd175b663cf 100644 --- a/aten/src/ATen/native/BinaryOps.cpp +++ b/aten/src/ATen/native/BinaryOps.cpp @@ -250,24 +250,24 @@ TORCH_IMPL_FUNC(special_xlog1py_out) (const Tensor& self, const Tensor& other, c xlog1py_stub(device_type(), *this); } -#define CREATE_BINARY_TORCH_IMPL_FUNC(func) \ -TORCH_IMPL_FUNC(func##_out) (const Tensor& self, const Tensor& other, const Tensor& result) { \ - func##_stub(device_type(), *this); \ -} - -CREATE_BINARY_TORCH_IMPL_FUNC(maximum); -CREATE_BINARY_TORCH_IMPL_FUNC(minimum); -CREATE_BINARY_TORCH_IMPL_FUNC(fmax); -CREATE_BINARY_TORCH_IMPL_FUNC(fmin); -CREATE_BINARY_TORCH_IMPL_FUNC(logaddexp); -CREATE_BINARY_TORCH_IMPL_FUNC(logaddexp2); -CREATE_BINARY_TORCH_IMPL_FUNC(gcd); -CREATE_BINARY_TORCH_IMPL_FUNC(lcm); -CREATE_BINARY_TORCH_IMPL_FUNC(hypot); -CREATE_BINARY_TORCH_IMPL_FUNC(igamma); -CREATE_BINARY_TORCH_IMPL_FUNC(igammac); -CREATE_BINARY_TORCH_IMPL_FUNC(nextafter); -CREATE_BINARY_TORCH_IMPL_FUNC(remainder); +#define CREATE_BINARY_TORCH_IMPL_FUNC(func_out, func_stub) \ +TORCH_IMPL_FUNC(func_out) (const Tensor& self, const Tensor& other, const Tensor& result) { \ + func_stub(device_type(), *this); \ +} + +CREATE_BINARY_TORCH_IMPL_FUNC(maximum_out, maximum_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(minimum_out, minimum_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(fmax_out, fmax_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(fmin_out, fmin_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(logaddexp_out, logaddexp_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(logaddexp2_out, logaddexp2_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(gcd_out, gcd_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(lcm_out, lcm_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(hypot_out, hypot_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(igamma_out, igamma_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(igammac_out, igammac_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(nextafter_out, nextafter_stub); +CREATE_BINARY_TORCH_IMPL_FUNC(remainder_out, remainder_stub); Tensor special_xlog1py(const Scalar& x, const Tensor& y) { return at::special_xlog1py(wrapped_scalar_tensor(x), y); diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index 5a01df1f9b87d..aeb83025784ba 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -127,53 +127,53 @@ namespace native { // them work for your case, but just write something new instead. Here we use helper functions instead of a flat fat // macro that implements everything, because the former allows some simple preprocessing that are unique to some // operators (more is foreseeable) and is more flexible and elegant than the latter. -#define CREATE_UNARY_TORCH_IMPL_FUNC(func) \ -TORCH_IMPL_FUNC(func##_out) (const Tensor& self, const Tensor& result) { \ - func##_stub(device_type(), *this); \ -} - -CREATE_UNARY_TORCH_IMPL_FUNC(acos) -CREATE_UNARY_TORCH_IMPL_FUNC(acosh) -CREATE_UNARY_TORCH_IMPL_FUNC(asin) -CREATE_UNARY_TORCH_IMPL_FUNC(asinh) -CREATE_UNARY_TORCH_IMPL_FUNC(atan) -CREATE_UNARY_TORCH_IMPL_FUNC(atanh) -CREATE_UNARY_TORCH_IMPL_FUNC(bitwise_not) -CREATE_UNARY_TORCH_IMPL_FUNC(ceil) -CREATE_UNARY_TORCH_IMPL_FUNC(cos) -CREATE_UNARY_TORCH_IMPL_FUNC(cosh) -CREATE_UNARY_TORCH_IMPL_FUNC(digamma) -CREATE_UNARY_TORCH_IMPL_FUNC(erf) -CREATE_UNARY_TORCH_IMPL_FUNC(erfc) -CREATE_UNARY_TORCH_IMPL_FUNC(erfinv) -CREATE_UNARY_TORCH_IMPL_FUNC(exp) -CREATE_UNARY_TORCH_IMPL_FUNC(exp2) -CREATE_UNARY_TORCH_IMPL_FUNC(expm1) -CREATE_UNARY_TORCH_IMPL_FUNC(floor) -CREATE_UNARY_TORCH_IMPL_FUNC(frac) -CREATE_UNARY_TORCH_IMPL_FUNC(i0) -CREATE_UNARY_TORCH_IMPL_FUNC(lgamma) -CREATE_UNARY_TORCH_IMPL_FUNC(log) -CREATE_UNARY_TORCH_IMPL_FUNC(log10) -CREATE_UNARY_TORCH_IMPL_FUNC(log1p) -CREATE_UNARY_TORCH_IMPL_FUNC(log2) -CREATE_UNARY_TORCH_IMPL_FUNC(neg) -CREATE_UNARY_TORCH_IMPL_FUNC(reciprocal) -CREATE_UNARY_TORCH_IMPL_FUNC(round) -CREATE_UNARY_TORCH_IMPL_FUNC(rsqrt) -CREATE_UNARY_TORCH_IMPL_FUNC(sigmoid) -CREATE_UNARY_TORCH_IMPL_FUNC(sign) -CREATE_UNARY_TORCH_IMPL_FUNC(sin) -CREATE_UNARY_TORCH_IMPL_FUNC(sinc) -CREATE_UNARY_TORCH_IMPL_FUNC(sinh) -CREATE_UNARY_TORCH_IMPL_FUNC(special_entr) -CREATE_UNARY_TORCH_IMPL_FUNC(special_i0e) -CREATE_UNARY_TORCH_IMPL_FUNC(special_i1e) -CREATE_UNARY_TORCH_IMPL_FUNC(special_i1) -CREATE_UNARY_TORCH_IMPL_FUNC(sqrt) -CREATE_UNARY_TORCH_IMPL_FUNC(tan) -CREATE_UNARY_TORCH_IMPL_FUNC(tanh) -CREATE_UNARY_TORCH_IMPL_FUNC(trunc) +#define CREATE_UNARY_TORCH_IMPL_FUNC(func_out, func_stub) \ +TORCH_IMPL_FUNC(func_out) (const Tensor& self, const Tensor& result) { \ + func_stub(device_type(), *this); \ +} + +CREATE_UNARY_TORCH_IMPL_FUNC(acos_out, acos_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(acosh_out, acosh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(asin_out, asin_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(asinh_out, asinh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(atan_out, atan_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(atanh_out, atanh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(bitwise_not_out, bitwise_not_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(ceil_out, ceil_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(cos_out, cos_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(cosh_out, cosh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(digamma_out, digamma_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(erf_out, erf_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(erfc_out, erfc_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(erfinv_out, erfinv_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(exp_out, exp_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(exp2_out, exp2_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(expm1_out, expm1_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(floor_out, floor_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(frac_out, frac_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(i0_out, i0_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(lgamma_out, lgamma_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(log_out, log_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(log10_out, log10_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(log1p_out, log1p_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(log2_out, log2_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(neg_out, neg_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(reciprocal_out, reciprocal_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(round_out, round_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(rsqrt_out, rsqrt_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sigmoid_out, sigmoid_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sign_out, sign_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sin_out, sin_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sinc_out, sinc_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sinh_out, sinh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(special_entr_out, special_entr_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(special_i0e_out, special_i0e_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(special_i1e_out, special_i1e_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(special_i1_out, special_i1_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(sqrt_out, sqrt_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(tan_out, tan_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(tanh_out, tanh_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(trunc_out, trunc_stub) TORCH_IMPL_FUNC(polygamma_out) (int64_t n, const Tensor& self, const Tensor& result) { From d99a8a31b1536a19f3ed5a20d189ebdeab11e903 Mon Sep 17 00:00:00 2001 From: Ivan Yashchuk Date: Wed, 16 Jun 2021 18:03:09 -0700 Subject: [PATCH 173/305] Fix version comparison for defining CUDA11OrLater (#60010) Summary: Before this PR `CUDA11OrLater` was incorrectly set to `False` when `torch.version.cuda == "11.0"`. `torch.version.cuda` returns major and minor CUDA versions, it doesn't return patch info. LooseVersion comparison was calling `[11, 0] >= [11, 0, 0]` which evaluates to `False`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60010 Reviewed By: mruberry Differential Revision: D29147107 Pulled By: ezyang fbshipit-source-id: bd9ed076337b4d32bf1c3376b8f7ae15dbc4d08d --- torch/testing/_internal/common_cuda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/testing/_internal/common_cuda.py b/torch/testing/_internal/common_cuda.py index 30610ec5db1a3..60195a1eaa22c 100644 --- a/torch/testing/_internal/common_cuda.py +++ b/torch/testing/_internal/common_cuda.py @@ -16,7 +16,7 @@ TEST_CUDNN = TEST_CUDA and torch.backends.cudnn.is_acceptable(torch.tensor(1., device=CUDA_DEVICE)) TEST_CUDNN_VERSION = torch.backends.cudnn.version() if TEST_CUDNN else 0 -CUDA11OrLater = torch.version.cuda and distutils.version.LooseVersion(torch.version.cuda) >= "11.0.0" +CUDA11OrLater = torch.version.cuda and distutils.version.LooseVersion(torch.version.cuda) >= "11.0" CUDA9 = torch.version.cuda and torch.version.cuda.startswith('9.') SM53OrLater = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (5, 3) From c01939a9b1424c39c8e6a8f15b4118c36a3a49da Mon Sep 17 00:00:00 2001 From: Meghan Lele Date: Wed, 16 Jun 2021 19:59:36 -0700 Subject: [PATCH 174/305] [JIT] Handle modules that already have __constants__ (#60003) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60003 **Summary** `infer_concrete_type_builder` in `_recursive.py` assumes `__constants__` is a `set` if it exists as an attribute on the module being scripted. Instead, it should create a set out of whatever `__constants__` is. **Test Plan** Ran code from the issue. **Fixes** This commit fixes #59947. Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D29174243 Pulled By: SplitInfinity fbshipit-source-id: aeb8bded80038da35478714b6a697a766ac447f5 --- test/jit/test_modules.py | 31 +++++++++++++++++++++++++++++++ test/test_jit.py | 1 + torch/jit/_recursive.py | 2 +- 3 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 test/jit/test_modules.py diff --git a/test/jit/test_modules.py b/test/jit/test_modules.py new file mode 100644 index 0000000000000..320a1a4357f33 --- /dev/null +++ b/test/jit/test_modules.py @@ -0,0 +1,31 @@ +import torch +import os +import sys +from torch.testing._internal.jit_utils import JitTestCase + +# Make the helper files in test/ importable +pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) +sys.path.append(pytorch_test_dir) + +if __name__ == '__main__': + raise RuntimeError("This test file is not meant to be run directly, use:\n\n" + "\tpython test/test_jit.py TESTNAME\n\n" + "instead.") + +class TestModules(JitTestCase): + def test_script_module_with_constants_list(self): + """ + Test that a module that has __constants__ set to something + that is not a set can be scripted. + """ + + # torch.nn.Linear has a __constants__ attribute defined + # and intialized to a list. + class Net(torch.nn.Linear): + x: torch.jit.Final[int] + + def __init__(self): + super().__init__(5, 10) + self.x = 0 + + self.checkModule(Net(), (torch.randn(5),)) diff --git a/test/test_jit.py b/test/test_jit.py index 256183a6cb73b..27dd5a47ffeb7 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -13,6 +13,7 @@ from jit.test_async import TestAsync # noqa: F401 from jit.test_data_parallel import TestDataParallel # noqa: F401 from jit.test_models import TestModels # noqa: F401 +from jit.test_modules import TestModules # noqa: F401 from jit.test_autodiff_subgraph_slicing import TestAutodiffSubgraphSlicing # noqa: F401 from jit.test_custom_operators import TestCustomOperators # noqa: F401 from jit.test_export_modes import TestExportModes # noqa: F401 diff --git a/torch/jit/_recursive.py b/torch/jit/_recursive.py index a466ff8aee447..7cd0b2280133b 100644 --- a/torch/jit/_recursive.py +++ b/torch/jit/_recursive.py @@ -198,7 +198,7 @@ def infer_type(name, item): added_names.add(name) # populate constants_set - constants_set = getattr(nn_module, "__constants__", set()) + constants_set = set(getattr(nn_module, "__constants__", ())) # Constants annotated via `Final[T]` rather than being added to `__constants__` for name, ann in class_annotations.items(): From 85517a2b700a5abc0b38f53ce8c99404cd67db79 Mon Sep 17 00:00:00 2001 From: Jason Ansel Date: Wed, 16 Jun 2021 19:59:42 -0700 Subject: [PATCH 175/305] [TensorExpr] More python binding cleanups (#60058) Summary: A few more quality of life improvements for NNC's python bindings: - Use standard `torch.dtype`s (rather than `te.Dtype`) - Make names optional (they don't seem to matter) - Make shapes optional - A few implicit conversions to make code cleaner Followup to https://github.com/pytorch/pytorch/issues/59920 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60058 Reviewed By: bertmaher Differential Revision: D29151953 Pulled By: jansel fbshipit-source-id: c8286e329eb4ee3921ca0786e17248cf6a898bd8 --- test/test_tensorexpr_pybind.py | 25 +++++++++------ torch/csrc/jit/tensorexpr/expr.h | 5 +++ torch/csrc/jit/tensorexpr/tensor.h | 6 ++++ torch/csrc/jit/tensorexpr/tensorexpr_init.cpp | 31 ++++++++++++++++--- 4 files changed, 53 insertions(+), 14 deletions(-) diff --git a/test/test_tensorexpr_pybind.py b/test/test_tensorexpr_pybind.py index a3efb8416a37d..cc4551515bb48 100644 --- a/test/test_tensorexpr_pybind.py +++ b/test/test_tensorexpr_pybind.py @@ -49,7 +49,7 @@ def test_simple_sum(self): def test_call_raw(self): with kernel_arena_scope(): n = 16 - cg = construct_adder(n, dtype=te.Dtype.Double) + cg = construct_adder(n, dtype=torch.float64) tA = torch.randn(n, dtype=torch.float64) tB = torch.randn(n, dtype=torch.float64) @@ -59,7 +59,7 @@ def test_call_raw(self): def test_external_calls(self): with kernel_arena_scope(): - dtype = te.Dtype.Float + dtype = torch.float32 ONE = te.ExprHandle.int(1) FOUR = te.ExprHandle.int(4) @@ -81,22 +81,21 @@ def test_external_calls(self): def test_dynamic_shape(self): with kernel_arena_scope(): - dN = te.VarHandle("n", te.Dtype.Int) - A = te.Placeholder('A', te.Dtype.Double, [dN]) - B = te.Placeholder('B', te.Dtype.Double, [dN]) + dN = te.VarHandle(torch.int32) + A = te.BufHandle(torch.float64) + B = te.BufHandle(torch.float64) def compute(i): - return A.load([i]) - B.load([i]) + return A.load(i) - B.load(i) - C = te.Compute('C', [te.DimArg(dN, 'i')], compute) + C = te.Compute('C', [dN], compute) loopnest = te.LoopNest([C]) loopnest.prepare_for_codegen() - stmt = te.simplify(loopnest.root_stmt()) cg = te.construct_codegen( 'ir_eval', - stmt, + loopnest.simplify(), [A, B, C, dN]) def test_with_shape(n): @@ -109,6 +108,14 @@ def test_with_shape(n): test_with_shape(8) test_with_shape(31) + def test_dtype_error(self): + with kernel_arena_scope(): + one = te.ExprHandle.int(1) + te.Placeholder([one], torch.float32) # ok + te.Placeholder([one]) # ok + self.assertRaises(TypeError, + lambda: te.Placeholder([one], "float55")) + @unittest.skipIf(not LLVM_ENABLED, "LLVM backend not enabled") def test_kernel_with_tensor_inputs(self): def f(a, b, c): diff --git a/torch/csrc/jit/tensorexpr/expr.h b/torch/csrc/jit/tensorexpr/expr.h index 6e36be55fe713..2f96d8c4de0a4 100644 --- a/torch/csrc/jit/tensorexpr/expr.h +++ b/torch/csrc/jit/tensorexpr/expr.h @@ -247,6 +247,11 @@ class TORCH_API BufHandle : public ExprHandle { Dtype dtype) : ExprHandle(Buf::make(name_hint, dims, dtype)) {} + BufHandle(const std::vector& dims, Dtype dtype) + : ExprHandle(Buf::make("_", dims, dtype)) {} + + explicit BufHandle(Dtype dtype) : ExprHandle(Buf::make("_", {}, dtype)) {} + explicit BufHandle(const Buf* node) : ExprHandle(node) {} const Buf* node() const { return static_cast(ExprHandle::node()); diff --git a/torch/csrc/jit/tensorexpr/tensor.h b/torch/csrc/jit/tensorexpr/tensor.h index 436f3052db9e1..95c98af0bdce5 100644 --- a/torch/csrc/jit/tensorexpr/tensor.h +++ b/torch/csrc/jit/tensorexpr/tensor.h @@ -82,6 +82,12 @@ class Placeholder { const std::vector& dims) : Placeholder(BufHandle(name, dims, dtype)) {} + Placeholder(const std::vector& dims, const Dtype& dtype) + : Placeholder(BufHandle("_", dims, dtype)) {} + + explicit Placeholder(const std::vector& dims) + : Placeholder(BufHandle("_", dims, kFloat)) {} + const Buf* data() const { return data_; } diff --git a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp index c247ecb9b7202..c1e5fc6aa4f0c 100644 --- a/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp +++ b/torch/csrc/jit/tensorexpr/tensorexpr_init.cpp @@ -44,7 +44,15 @@ ArgValue convertPyToArgValue(py::handle inp) { throw std::runtime_error("vector conversion failed"); } } else { - throw std::runtime_error("nyi"); + throw std::runtime_error("conversion not yet implemented"); + } +} + +Dtype parsePythonDtype(py::handle obj) { + if (py::isinstance(obj, py::module_::import("torch").attr("dtype"))) { + return Dtype(reinterpret_cast(obj.ptr())->scalar_type); + } else { + throw std::runtime_error("expected a torch.dtype instance"); } } @@ -55,7 +63,9 @@ void initTensorExprBindings(PyObject* module) { auto te = m.def_submodule("_te"); py::class_(te, "KernelScope").def(py::init<>()); - auto dtype_class = py::class_(te, "Dtype"); + auto dtype_class = + py::class_(te, "Dtype").def(py::init(&parsePythonDtype)); + py::implicitly_convertible(); #define DTYPE_SINGLETON_ACCESSOR(ctype, name) \ dtype_class.def_property_readonly_static( \ @@ -139,21 +149,31 @@ void initTensorExprBindings(PyObject* module) { #undef EXPRHANDLE_CTOR py::class_(te, "VarHandle") + .def(py::init()) .def(py::init()); py::class_( // NOLINT te, "BufHandle") .def( py::init&, Dtype>()) - .def("load", [](BufHandle& self, const std::vector& v) { - return Load::make(self, v); + .def(py::init&, Dtype>()) + .def(py::init()) + .def( + "load", + [](BufHandle& self, const std::vector& v) { + return Load::make(self, v); + }) + .def("load", [](BufHandle& self, const ExprHandle& v) { + return Load::make(self, {v}); }); py::class_(te, "Placeholder") .def(py::init< const std::string&, const Dtype&, - std::vector&>()) + const std::vector&>()) + .def(py::init&, const Dtype&>()) + .def(py::init&>()) .def( "load", [](Placeholder& self, const std::vector& v) { @@ -183,6 +203,7 @@ void initTensorExprBindings(PyObject* module) { py::class_(te, "DimArg") .def(py::init()) .def(py::init()); + py::implicitly_convertible(); te.def( "Compute", From c6cdb4f1130a2404c863445028b7715bf31c7051 Mon Sep 17 00:00:00 2001 From: Andrew Gu Date: Wed, 16 Jun 2021 20:51:11 -0700 Subject: [PATCH 176/305] Refactor ZeroRedundancyOptimizer Assuming SPSD (#59834) Summary: **Overview:** This refactors the `ZeroRedundancyOptimizer` implementation to assume single-process single-device (SPSD) instead of accommodating single-process multiple-device (SPMD). `DistributedDataParallel` [retired SPMD recently](https://github.com/pytorch/pytorch/issues/47012), so this change follows the same spirit. **Changes:** The parent-class `Optimizer` constructor permits the input argument `params` to be both an `iterable` of `torch.Tensor` and an `iterable` of `dict`. The latter usage is for initializing the optimizer with multiple `param_group`s to start. However, currently, `ZeroRedundancyOptimizer` only supports the former usage, requiring explicit calls to `add_param_group()` for multiple `param_group`s. Given the existing implementation, the type error would be silent and not manifest until much later (e.g. since `super().__init__()` would have no issue). Hence, I added a series of checks to begin the `__init__()` function (encapsulated in `_verify_and_init_params()`). A postcondition of this validation is that `self._all_params` is a non-empty list of all model parameters. Additionally, I added a check for SPSD usage assuming that all model parameters exist on the same device. This logic is included in `_verify_same_param_device()` and is called immediately after the `params` type-checking. Support for SPSD with model parameters sharded across devices may be added in the future. Related to that aforementioned post-condition on `self._all_params`, previously there was undefined behavior resulting from different typing of the passed in `params` input argument. If `params` was a `List`, then the usage of `self._reference_is_trainable_mask` was as expected. However, if `params` was a generator (e.g. as in the canonical usage of passing `model.parameters()`), then the ensuing behavior was divergent. This is because after a generator is iterated over, it is empty. As a result, when we set `self._all_params = params` [in the old code](https://github.com/pytorch/pytorch/blob/68d690ffbd64d0fb697dc3da1635216366649787/torch/distributed/optim/zero_redundancy_optimizer.py#L165), `self._all_params` is empty, reducing `training_mask` to always be the empty list. This causes missed calls to `_update_trainable()` in `step()`. (A consequence of this is that `test_pytorch_parity()`, which is renamed to `test_local_optimizer_parity()`, now outputs warnings about the trainable parameters changing.) The existing implementation assumes that all parameters share the same dense type when allocating the bucket buffers. This change preserves this assumption, which may be removed in the future. I added a check for this in `_verify_same_dense_param_type()` to avoid erroring silently later on. Note that it is insufficient to simply check for the same `dtype` since dense and sparse tensors may share the same `dtype` but require differing storage sizes. One solution is to use `torch.typename()` as the means for comparison. --- The primary change in this refactor is with respect to `self._per_device_params` and `self.buckets`. `self._per_device_params` mapped `torch.device` to `List[List[Parameter]]`. The keys were the devices that the model parameters exist on, and the values designated which ranks are assigned to updating those parameters. `self.buckets` mapped `torch.device` to `List[torch.Tensor]`. The keys were the same as `self._per_device_params`, and the values were the buckets for that device. The usage of these two data structures were confined to each other only. Hence, because the notions of device and rank are now in 1:1 correspondence, we can eliminate the former completely and only use rank. As such, I removed `self._per_device_params` and made `self.buckets` directly a list of buckets (i.e. `torch.Tensor`s). Iteration over the parameters of a rank for a given device could be simplified to just iteration over the parameters of a rank. Hence, I relied on `self.partition_parameters()` now for that iteration. Refer to `_setup_flat_buffers()` and `step()` for these changes. One convenient side effect of removing `self._per_device_params` is that there is no longer the re-computation of the parameter partitions mentioned at the end of this [PR](https://github.com/pytorch/pytorch/pull/59410). --- I changed the data structure `self._index_to_param_cache` from a `dict` to a `List` because the domain is `0`, `1`, ..., `k-1` where `k` is the number of parameters. This should yield marginal improvements in memory usage and access speed. `_sync_param_groups()` is a static method, meaning it can be called either via `self._sync_param_groups()` or `ZeroRedundancyOptimizer._sync_param_groups()` when inside the class. I made the usage consistently `self._sync_param_groups()` rather than have instances of both. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59834 Test Plan: I ran through the existing test suite on an AI AWS cluster: ``` srun -p $DEV_QUEUE --cpus-per-task=16 -t 5:00:00 --gpus-per-node=4 python test/distributed/optim/test_zero_redundancy_optimizer.py ``` Note: The only test where `parameters_as_bucket_view` is `True` is `test_step_with_closure()`, meaning that that is the test that exercises the core changes of removing `self._per_device_params` and changing `self.buckets`. Also, I added tests for the `ZeroRedundancyOptimizer` constructor changes and the assumption checks. Reviewed By: mrshenli Differential Revision: D29177065 Pulled By: andwgu fbshipit-source-id: 0ff004ae3959d6d3b521024028c7156bfddc93d8 --- .../optim/test_zero_redundancy_optimizer.py | 62 +++- .../optim/zero_redundancy_optimizer.py | 269 +++++++++++------- 2 files changed, 227 insertions(+), 104 deletions(-) diff --git a/test/distributed/optim/test_zero_redundancy_optimizer.py b/test/distributed/optim/test_zero_redundancy_optimizer.py index 730d893367810..9cc9a4e410c98 100644 --- a/test/distributed/optim/test_zero_redundancy_optimizer.py +++ b/test/distributed/optim/test_zero_redundancy_optimizer.py @@ -67,7 +67,6 @@ def dist_init(self, rank, world_size=-1): return dist.init_process_group(backend=BACKEND, store=store, rank=rank, world_size=world_size) - class TestZeroRedundancyOptimizerSingleRank(TestZeroRedundancyOptimizer): def test_state_dict(self): """Check that the ZeroRedundancyOptimizer exposes the expected state dict interface, @@ -204,6 +203,65 @@ def test_zero_grad(self): self.assertFalse(m.weight.grad) self.assertFalse(m.bias.grad) + def test_constructor(self): + """Check the robustness of the ZeroRedundancyOptimizer constructor by + passing different values for `params`""" + self.dist_init(self.rank) + + m = torch.nn.Linear(1, 1) + # (input, expected error) + inputs = [ + ([], ValueError), # empty parameter list + (torch.randn(1), TypeError), # non-iterable: `torch.Tensor` + (1.2, TypeError), # non-iterable: `float` + ([{"params": m.parameters()}], TypeError), # iterable of dict + (list(m.parameters()) + [42], TypeError), # iterable containing non-`torch.Tensor` + (m.parameters(), None), # `params` as a generator + (list(m.parameters()), None) # `params` as a list + ] + + for input, error in inputs: + if (error): + with self.assertRaises(error): + ZeroRedundancyOptimizer(input, optimizer_class=SGD, lr=0.1) + else: + ZeroRedundancyOptimizer(input, optimizer_class=SGD, lr=0.1) + + def test_same_dense_param_type(self): + """Check that ZeroRedundancyOptimizer raises an exception if the input + parameters include sparse tensors or different dense types. + + NOTE: This test should be removed once support for sparse parameters + and varying parameter types is added. + """ + self.dist_init(self.rank) + + inputs = [ + [torch.sparse_coo_tensor(size=(2, 3))], + [torch.FloatTensor(1), torch.DoubleTensor(1)], + [torch.FloatTensor(1), torch.FloatTensor(1), + torch.sparse_coo_tensor(size=(2, 3))] + ] + for input in inputs: + with self.assertRaises(ValueError): + ZeroRedundancyOptimizer(input, optimizer_class=SGD, lr=0.1) + + def test_same_param_device(self): + """Check that ZeroRedundancyOptimizer raises an exception if the input + parameters are sharded on multiple devices. + + NOTE: This test should be removed once support for sharding a rank's + model parameters across multiple devices is added. + """ + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + return + self.dist_init(self.rank) + + # Move the parameters to cuda:0 and cuda:1 respectively + params = [torch.Tensor(1).to(0), torch.Tensor(1).to(1)] + with self.assertRaises(ValueError): + ZeroRedundancyOptimizer(params, optimizer_class=SGD, lr=0.1) + class TestZeroRedundancyOptimizerDistributed(TestZeroRedundancyOptimizer): @property @@ -499,7 +557,7 @@ def closure(): check(optimizer) @common_distributed.skip_if_no_gpu - def test_pytorch_parity(self): + def test_local_optimizer_parity(self): """When combined with DDP, check that ZeroRedundancyOptimizer(optimizer) and the same monolithic optimizer give the exact same results """ diff --git a/torch/distributed/optim/zero_redundancy_optimizer.py b/torch/distributed/optim/zero_redundancy_optimizer.py index eb7eb0f0ddb8b..d299c4f986640 100644 --- a/torch/distributed/optim/zero_redundancy_optimizer.py +++ b/torch/distributed/optim/zero_redundancy_optimizer.py @@ -6,13 +6,11 @@ import collections import copy import io -from collections import OrderedDict from itertools import chain from typing import Any, Callable, Dict, List, Optional, Type import torch import torch.distributed as dist -from torch.nn import Parameter from torch.optim import Optimizer import logging @@ -96,7 +94,7 @@ class ZeroRedundancyOptimizer(Optimizer): in conjunction with :class:`torch.nn.parallel.DistributedDataparallel` to reduce per-rank peak memory consumption. - ``ZeroRedundancyOptimizer`` use a greedy algorithm to pack a number of + ``ZeroRedundancyOptimizer`` uses a sorted-greedy algorithm to pack a number of parameters at each rank. Each parameter belongs to a single rank and is not divided among ranks. The partition is arbitrary and might not match the the parameter registration or usage order. @@ -110,7 +108,7 @@ class ZeroRedundancyOptimizer(Optimizer): group (``ProcessGroup``, optional): ``torch.distributed`` ``ProcessGroup`` (default: ``group.WORLD`` initialized by :meth:`torch.distributed.init_process_group`). - parameters_as_bucket_views (bool): when enabled, parameters will + parameters_as_bucket_view (bool): when enabled, parameters will be packed into larger buckets to speed up communication and ``param.data`` fields will point to bucket views at different offsets. When disabled, each individual parameter will be @@ -133,6 +131,10 @@ class ZeroRedundancyOptimizer(Optimizer): >>> ddp(inputs).sum().backward() >>> opt.step() + .. note: Currently, ``ZeroRedundancyOptimizer`` requires that all of the + passed-in parameters are on the same device and that they are the same + dense type. + .. warning: ZeroRedundancyOptimizer is experimental and subject to change. .. _ZeRO: https://arxiv.org/abs/1910.02054 @@ -147,23 +149,24 @@ def __init__( parameters_as_bucket_view: bool = False, **default: Any, ): + # Perform type and assumption checks on the input parameters + self._verify_and_init_params(params) + self._verify_same_param_device() + self._verify_same_dense_param_type() + self._device = self._all_params[0].device + # Hold all the model params in the root .param_groups # NOTE: the default constructor uses `add_param_group` which is partially overloaded here # we introduce the `initialized` flag for be able to dissociate the behaviour of # `add_param_group` in between super() and ZeroRedundancyOptimizer self.initialized = False - super().__init__(params, default) + super().__init__(self._all_params, default) - # Partition information. lazy evaluation, computed if requested - self._per_device_params_cache: "OrderedDict[torch.device, List[List[Parameter]]]" = ( - OrderedDict() - ) # device, rank, params - self._param_rank_cache: Dict[torch.Tensor, int] = {} + # Partition information (evaluated lazily) + self._param_to_rank_cache: Dict[torch.Tensor, int] = {} self._param_to_index_cache: Dict[int, int] = {} self._partition_parameters_cache: List[List[Dict]] = [] - self._index_to_param_cache: Dict[int, torch.Tensor] = {} - self._all_params = params - self._reference_is_trainable_mask = list(map(_is_trainable, self._all_params)) + self._index_to_param_cache: List[torch.Tensor] = [] # Build the wrapped optimizer, responsible for a shard of the params self.group = group if group is not None else dist.group.WORLD @@ -175,20 +178,21 @@ def __init__( self._optim_defaults = default self._optim_constructor = optimizer_class - # Optional consolidated optimizer state + # Optional consolidated optimizer state self._all_states: List[Dict[str, Any]] = [] - # Current default device is set by the parameters allocated to this rank - self._device = list(self._per_device_params.keys())[0] - self.buckets: Dict[torch.device, List[torch.Tensor]] = {} + self._reference_is_trainable_mask = list(map(_is_trainable, self._all_params)) + self.buckets: List[torch.Tensor] = [] self._update_trainable() self.initialized = True def _clear_cache(self) -> None: + r""" + Clears the cached data structures giving partition information. + """ self._partition_parameters_cache.clear() - self._per_device_params_cache.clear() - self._param_rank_cache.clear() + self._param_to_rank_cache.clear() self._index_to_param_cache.clear() self._param_to_index_cache.clear() @@ -329,62 +333,38 @@ def partition_parameters(self) -> List[List[Dict]]: return self._partition_parameters_cache @property - def _per_device_params(self) -> Dict[torch.device, List[List[Parameter]]]: + def _param_to_rank(self) -> Dict[torch.Tensor, int]: r""" - Sorted list of all the params, first per device then per rank. - - Within a list params are sorted per number of elements to allow for an easy bucketing. + Hash table mapping parameters to their assigned data parallel rank in + the partition. """ - if len(self._per_device_params_cache) == 0: - # Go through all params, log them per device - # The ordering is important here, needs to be the same on all ranks - # So that ulterior broadcast calls are matching - for param_group in self.param_groups: - for param in param_group["params"]: - device = param.device - if self._per_device_params_cache.get(device) is None: - self._per_device_params_cache[device] = [[] for _ in range(self.world_size)] - self._per_device_params_cache[device][self._param_to_rank[param]] += [param] - - # Sort param_lists by size - for k in self._per_device_params_cache.keys(): - for r in self._per_device_params_cache[k]: - r.sort(key=lambda x: x.numel()) - - return self._per_device_params_cache - - @property - def _param_to_rank(self) -> Dict[torch.Tensor, int]: - r"""Look up table to match a given param with a data parallel rank""" - if len(self._param_rank_cache) == 0: + if len(self._param_to_rank_cache) == 0: for rank, param_groups in enumerate(self.partition_parameters()): for param_group in param_groups: for param in param_group["params"]: - self._param_rank_cache[param] = rank - return self._param_rank_cache + self._param_to_rank_cache[param] = rank + return self._param_to_rank_cache @property def _param_to_index(self) -> Dict[int, int]: r""" - Hash table in between parameter indices in the global optimizer scheme, - and the actual params. + Hash table mapping parameters to their indices in the global optimizer + scheme. """ if len(self._param_to_index_cache) == 0: self._param_to_index_cache = { id(p): i for i, p in enumerate(chain(*(g["params"] for g in self.param_groups))) } - return self._param_to_index_cache @property def _index_to_param(self) -> Dict[int, torch.Tensor]: r""" - Hash table in between parameter indices in the global optimizer scheme, - and the actual params. + List mapping parameter indices in the global optimizer scheme to the + actual params. """ if len(self._index_to_param_cache) == 0: - self._index_to_param_cache = {i: p for i, p in enumerate(chain(*(g["params"] for g in self.param_groups)))} - + self._index_to_param_cache = list(chain(*(g["params"] for g in self.param_groups))) return self._index_to_param_cache def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]: @@ -404,7 +384,8 @@ def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> trainable_mask = list(map(_is_trainable, self._all_params)) if trainable_mask != self._reference_is_trainable_mask: logging.warning( - "ZeroRedundancyOptimizer detected that the trainable params changed, updating the partitioning" + "ZeroRedundancyOptimizer detected that the trainable params " + "changed, updating the partitioning" ) self._update_trainable() self._reference_is_trainable_mask = trainable_mask @@ -421,19 +402,21 @@ def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> # Sync all the updated shards in between the ranks handles = [] if self.parameters_as_bucket_view: - for device in self.buckets.keys(): - for src_rank, bucket in enumerate(self.buckets[device]): - global_src_rank = _get_global_rank(self.group, src_rank) - handles.append(dist.broadcast(tensor=bucket, src=global_src_rank, group=self.group, async_op=True)) + for rank, bucket in enumerate(self.buckets): + global_rank = _get_global_rank(self.group, rank) + handles.append( + dist.broadcast(tensor=bucket, src=global_rank, + group=self.group, async_op=True) + ) else: - for device, per_rank_params in self._per_device_params.items(): - for dst_rank, params in enumerate(per_rank_params): - global_dst_rank = _get_global_rank(self.group, dst_rank) - for param in params: + for rank, param_groups in enumerate(self.partition_parameters()): + global_rank = _get_global_rank(self.group, rank) + for param_group in param_groups: + for param in param_group["params"]: handles.append( - dist.broadcast(tensor=param.data, src=global_dst_rank, group=self.group, async_op=True) + dist.broadcast(tensor=param.data, src=global_rank, + group=self.group, async_op=True) ) - _ = list(map(lambda x: x.wait(), handles)) # Sync hypothethical new results from the wrapped optimizer to the exposed param_groups @@ -462,8 +445,8 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None: super().load_state_dict(state_dict) # Sync with the optimizer param groups - ZeroRedundancyOptimizer._sync_param_groups(state_dict["param_groups"], self.param_groups) - ZeroRedundancyOptimizer._sync_param_groups(self.param_groups, self.optim.param_groups) + self._sync_param_groups(state_dict["param_groups"], self.param_groups) + self._sync_param_groups(self.param_groups, self.optim.param_groups) def local_state_dict(self) -> Dict: r""" @@ -548,44 +531,45 @@ def _sync_param_groups(source: List[Dict[Any, Any]], destination: List[Dict[Any, def _setup_flat_buffers(self) -> None: r""" Make all params which are on the same device and tied to the same rank - views of a single buffer. This is used at construction time, and anytime - parameter trainability is changed (frozen or unfrozen) and + views of a single buffer. This is used at construction time, and + anytime parameter trainability is changed (frozen or unfrozen) and ``_update_trainable`` is called. """ - - for device, per_rank_params in self._per_device_params.items(): - # Only wipe the existing buckets if there are none - # (could be that this is called twice, when trainability changes) - if device not in self.buckets.keys(): - self.buckets[device] = [] - - # Make parameters a view of the bucket - for dst_rank, params in enumerate(per_rank_params): - if len(params) > 0: - - # Clone the non-trainable params, if in a bucket it will get destroyed - for param in filter(lambda x: not x.requires_grad, params): + for rank, param_groups in enumerate(self.partition_parameters()): + # Clone the non-trainable params, find the buffer size and dtype + # for the trainable params' bucket, and compile a list of the + # trainable params + buffer_size = 0 + dtype = None + trainable_params = [] + for param_group in param_groups: + for param in param_group["params"]: + if not _is_trainable(param): param.data = param.data.detach().clone() - - # Merge all the trainable params in a single bucket - trainable_params = list(filter(_is_trainable, params)) - buffer_size = sum(map(lambda x: x.numel(), trainable_params)) - bucket = torch.empty(buffer_size, dtype=params[0].dtype, device=device) - offset = 0 - - for param in trainable_params: - offset_next = offset + param.numel() - bucket[offset:offset_next].copy_(param.data.flatten()) - param.data = bucket[offset:offset_next].view_as(param.data) - offset = offset_next - - # Either replace the existing bucket, or create it - if len(self.buckets[device]) == dst_rank: - self.buckets[device].append(bucket) else: - self.buckets[device][dst_rank] = bucket - else: - self.buckets[device].append(torch.zeros(1, device=device)) + buffer_size += param.numel() + trainable_params.append(param) + dtype = param.dtype # assumes all dense and same dtype + + # Create a dummy bucket if there are no params + if buffer_size == 0: + self.buckets.append(torch.zeros(1, device=self._device)) + continue + + # Otherwise, construct the bucket + bucket = torch.empty(buffer_size, dtype=dtype, device=self._device) + offset = 0 + for param in trainable_params: + offset_next = offset + param.numel() + bucket[offset:offset_next].copy_(param.data.flatten()) + param.data = bucket[offset:offset_next].view_as(param.data) + offset = offset_next + + # Either replace the existing bucket or create it + if len(self.buckets) != rank: + self.buckets[rank] = bucket + else: + self.buckets.append(bucket) def _update_trainable(self) -> None: r""" @@ -596,9 +580,90 @@ def _update_trainable(self) -> None: # Create the optim which will work on the param shard if not hasattr(self, "optim"): self._clear_cache() - self._default_device = list(self._per_device_params.keys())[0] self.optim = self._optim_constructor(self.partition_parameters()[self.rank], **self._optim_defaults) self._sync_param_groups(self.optim.param_groups, self.param_groups) if self.parameters_as_bucket_view: self._setup_flat_buffers() + + def _verify_and_init_params(self, params: Any) -> None: + r""" + Verifies the type of ``params`` and initializes ``self._all_params`` + if ``params`` is valid. + + While :class:`optim.Optimizer ` allows + ``params`` to be an iterable of :class:`dict` s, currently + ``ZeroRedundancyOptimizer`` strictly requires ``params`` to be an + iterable of :class:`torch.Tensor` s. + + Raises: + TypeError: ``params`` has an invalid type. + ValueError: ``params`` is empty. + """ + if isinstance(params, torch.Tensor): + raise TypeError("params argument should be an iterable of " + f"Tensors, but got {torch.typename(params)}") + try: + self._all_params = list(params) + except TypeError: + raise TypeError("params argument should be an iterable of " + f"Tensors, but got {torch.typename(params)}") + if len(self._all_params) == 0: + raise ValueError("ZeroRedundancyOptimizer got an empty parameter " + "list") + for param in self._all_params: + if not isinstance(param, torch.Tensor): + raise TypeError("params argument should be an iterable of " + "Tensors, but got an iterable containing " + f"{torch.typename(param)}") + + def _verify_same_param_device(self) -> None: + r""" + Verifies that ZeRO is being used under the single-process single- + device regime where a process operates exclusively on a full model + replica on a single device. + + The function assumes that ``self._all_params`` has been initialized + and is non-empty. + + Raises: + ValueError: ``params`` contains parameters across multiple + devices. + + NOTE: This function can be removed once support for sharding a rank's + model parameters across multiple devices is added. + """ + device = self._all_params[0].device + for param in self._all_params[1:]: + if param.device != device: + raise ValueError("ZeroRedundancyOptimizer assumes that each " + "rank's model parameters are on the same " + f"device but got both {device} and " + f"{param.device}") + + def _verify_same_dense_param_type(self) -> None: + r""" + Verifies that all parameters are of the same dense type. + + The function assumes that ``self._all_params`` has been initialized + and is non-empty. + + Raises: + ValueError: ``params`` contains sparse parameters or parameters + of varying dense types. + + NOTE: This function can be removed once support for sparse parameters + and varying parameter types is added. + """ + typename = torch.typename(self._all_params[0]) + if self._all_params[0].is_sparse: + raise ValueError("ZeroRedundancyOptimizer only supports using " + "the same dense type for all parameters but got " + f"{typename}") + for param in self._all_params[1:]: + other_typename = torch.typename(param) + if other_typename != typename: + raise ValueError("ZeroRedundancyOptimizer only supports " + "using the same dense type for all " + f"parameters but got both {typename} and " + f"{other_typename}") From 7ce74f3339d66cbf4d78b9cfdddf71847194fadb Mon Sep 17 00:00:00 2001 From: Angela Yi Date: Wed, 16 Jun 2021 22:30:02 -0700 Subject: [PATCH 177/305] [quant] EqualizationQConfig to distinguish input/output activations (#59739) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59739 Created an EqualizationQConfig specifically for equalization. This inherits from QConfig and is used to distinguish between inserting an input observer with an output observer. Since the output observer field is included in the EqualizationQConfig, we no longer need an output observer field in the _InputEqualizationObserver Test Plan: compiles Imported from OSS Reviewed By: ezyang Differential Revision: D29135298 fbshipit-source-id: 3dde9c029c291467ff0a0845f0fc9c44573fc6f6 --- torch/quantization/fx/_equalize.py | 58 ++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 14 deletions(-) diff --git a/torch/quantization/fx/_equalize.py b/torch/quantization/fx/_equalize.py index fdbf6c977bec1..0869546759423 100644 --- a/torch/quantization/fx/_equalize.py +++ b/torch/quantization/fx/_equalize.py @@ -1,7 +1,11 @@ import torch import torch.nn as nn -from torch.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver +from ..observer import ( + PerChannelMinMaxObserver, _with_args +) + +from collections import namedtuple import warnings @@ -16,8 +20,6 @@ class _InputEqualizationObserver(nn.Module): follow the 8-bit setup. quant_max: Maximum quantization value. If unspecified, it will follow the 8-bit setup. - output_obs: For the user to specify what kind of output observer they - would like to use The running minimum/maximum :math:`x_\text{min/max}` are computed in the same way as :class:`~torch.quantization.observer.PerChannelMinMaxObserver`, @@ -33,8 +35,7 @@ class _InputEqualizationObserver(nn.Module): """ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine, - quant_min=None, quant_max=None, output_obs=None, - factory_kwargs=None) -> None: + quant_min=None, quant_max=None, factory_kwargs=None) -> None: super(_InputEqualizationObserver, self).__init__() if qscheme not in {torch.per_tensor_affine, torch.per_tensor_symmetric}: @@ -46,15 +47,6 @@ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine, quant_max=quant_max, factory_kwargs=factory_kwargs) - if output_obs is None: - self.output_obs = MinMaxObserver(dtype=dtype, - qscheme=qscheme, - quant_min=quant_min, - quant_max=quant_max, - factory_kwargs=factory_kwargs) - else: - self.output_obs = output_obs - self.equalization_scale = torch.empty(0) def forward(self, x_orig): @@ -92,6 +84,8 @@ def calculate_qparams(self): return scale_input, zero_point_input + with_args = classmethod(_with_args) + class _WeightEqualizationObserver(nn.Module): r"""Observer for tracking the running min/max values of weight columns and @@ -206,6 +200,8 @@ def calculate_qparams(self): return scale_weight, zero_point_weight + with_args = classmethod(_with_args) + def calculate_equalization_scale(input_obs: _InputEqualizationObserver, weight_obs: _WeightEqualizationObserver) -> torch.Tensor: @@ -229,3 +225,37 @@ def calculate_equalization_scale(input_obs: _InputEqualizationObserver, equalization_scale = torch.sqrt((max_weights - min_weights) / (max_inputs - min_inputs)) return equalization_scale + + +class EqualizationQConfig(namedtuple('EqualizationQConfig', ['input_activation', 'weight'])): + """ + Describes how to quantize a layer or a part of the network specifically for + input-weight equalization by providing settings (observer classes) for + inputs, outputs, and weights. + + Note that EqualizationQConfig needs to contain observer **classes** (like + MinMaxObserver) or a callable that returns instances on invocation, not the + concrete observer instances themselves. + Quantization function will instantiate observers multiple times for each of + the layers. + + Observer classes have usually reasonable default arguments, but they can be + overwritten with `with_args` method (that behaves like functools.partial): + + my_qconfig = EqualizationQConfig(input_activation=_InputEqualizationObserver.with_args(dtype=torch.qint8), + weight=_WeightEqualizationObserver.with_args(dtype=torch.qint8)) + """ + def __new__(cls, input_activation=torch.nn.Identity, weight=torch.nn.Identity): + if isinstance(input_activation, nn.Module) or isinstance(weight, nn.Module): + raise ValueError("EqualizationQConfig received observer instance, please pass observer class instead. " + + "Use MyObserver.with_args(x=1) to override arguments to constructor if needed") + self = super(EqualizationQConfig, cls).__new__(cls, input_activation, weight) + return self + + +input_equalization_observer = _InputEqualizationObserver.with_args( + dtype=torch.quint8, qscheme=torch.per_tensor_symmetric) +weight_equalization_observer = _WeightEqualizationObserver.with_args( + dtype=torch.qint8, qscheme=torch.per_channel_symmetric) +default_equalization_qconfig = EqualizationQConfig(input_activation=input_equalization_observer, + weight=weight_equalization_observer) From 45c31cabb57c9ccd8f29686f0b944be7ff0be010 Mon Sep 17 00:00:00 2001 From: Angela Yi Date: Wed, 16 Jun 2021 22:30:02 -0700 Subject: [PATCH 178/305] [quant] Input Weight Equalization - prepare modifications (#59747) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59747 Modifies prepare_fx for input-weight equalization. If a current node is being equalized (there exists a EqualizationQConfig), then the EqualizationObserver will be inserted before its quantization observer. For a singular linear layer, the general flow looks like: Original graph: `x0 -> linear -> x1`, `w -> linear` After prepare: `x0 -> InpEqObs -> MinMaxObs -> linear1 -> MinMaxObs -> x1` `w -> WeightEqObs -> MinMaxObs -> linear1` For two connected linear layers, the general flow looks like: Original graph: `x0 -> linear1 -> linear2 -> x1`, `w1 -> linear1`, `w2 -> linear2` After prepare: `x0 -> InpEqObs -> MinMaxObs -> linear1 -> MinMaxObs -> InpEqObs -> linear2 -> MinMaxObs -> x1` `w1 -> WeightEqObs -> MinMaxObs -> linear1`, `w2 -> WeightEqObs -> MinMaxObs -> linear2 Test Plan: `python test/test_quantization.py TestEqualizeFx.test_input_equalization_prepare` Original model with one `nn.Linear` layer ``` LinearModule( (linear): Linear(in_features=1, out_features=1, bias=True) ) ``` Graph after `prepare_fx`: ``` graph(): %x : [#users=1] = placeholder[target=x] %x_equalization_process_0 : [#users=1] = call_module[target=x_equalization_process_0](args = (%x,), kwargs = {}) %x_activation_post_process_0 : [#users=1] = call_module[target=x_activation_post_process_00](args = (%x_equalization_process_0,), kwargs = {}) %linear : [#users=1] = call_module[target=linear](args = (%x_activation_post_process_0,), kwargs = {}) %linear_activation_post_process_0 : [#users=1] = call_module[target=linear_activation_post_process_0](args = (%linear,), kwargs = {}) return linear_activation_post_process_0 ``` -------------------------------------- Original model with two connected functional linear layers ``` FunctionalLinearModule( (linear1): Linear() (linear2): Linear() ) ``` Graph after `prepare_fx`: ``` graph(): %x : [#users=1] = placeholder[target=x] %x_equalization_process_0 : [#users=1] = call_module[target=x_equalization_process_0](args = (%x,), kwargs = {}) %x_activation_post_process_0 : [#users=1] = call_module[target=x_activation_post_process_00](args = (%x_equalization_process_0,), kwargs = {}) %linear1_w : [#users=1] = get_attr[target=linear1.w] %linear1_w_equalization_process_0 : [#users=1] = call_module[target=linear1_w_equalization_process_0](args = (%linear1_w,), kwargs = {}) %linear1_w_activation_post_process_0 : [#users=1] = call_module[target=linear1_w_activation_post_process_00](args = (%linear1_w_equalization_process_0,), kwargs = {}) %linear1_b : [#users=1] = get_attr[target=linear1.b] %linear : [#users=1] = call_function[target=torch.nn.functional.linear](args = (%x_activation_post_process_0, %linear1_w_activation_post_process_0), kwargs = {bias: %linear1_b}) %linear_activation_post_process_0 : [#users=1] = call_module[target=linear_activation_post_process_0](args = (%linear,), kwargs = {}) %linear_activation_post_process_0_equalization_process_0 : [#users=1] = call_module[target=linear_activation_post_process_0_equalization_process_0](args = (%linear_activation_post_process_0,), kwargs = {}) %linear2_w : [#users=1] = get_attr[target=linear2.w] %linear2_w_equalization_process_0 : [#users=1] = call_module[target=linear2_w_equalization_process_0](args = (%linear2_w,), kwargs = {}) %linear2_w_activation_post_process_0 : [#users=1] = call_module[target=linear2_w_activation_post_process_00](args = (%linear2_w_equalization_process_0,), kwargs = {}) %linear2_b : [#users=1] = get_attr[target=linear2.b] %linear_1 : [#users=1] = call_function[target=torch.nn.functional.linear](args = (%linear_activation_post_process_0_equalization_process_0, %linear2_w_activation_post_process_0), kwargs = {bias: %linear2_b}) %linear_1_activation_post_process_0 : [#users=1] = call_module[target=linear_1_activation_post_process_0](args = (%linear_1,), kwargs = {}) return linear_1_activation_post_process_0 ``` Imported from OSS Reviewed By: jerryzh168 Differential Revision: D29135316 fbshipit-source-id: 91697e805ede254dbb2a42ee4c23eb1c1c64590e --- test/quantization/fx/test_equalize_fx.py | 124 ++++++++++++++++++++++- torch/quantization/fx/_equalize.py | 20 ++++ torch/quantization/fx/graph_module.py | 1 + torch/quantization/fx/prepare.py | 71 ++++++++++++- torch/quantization/quantize_fx.py | 17 +++- 5 files changed, 225 insertions(+), 8 deletions(-) diff --git a/test/quantization/fx/test_equalize_fx.py b/test/quantization/fx/test_equalize_fx.py index e49aaef5113a2..3980d1bf42df5 100644 --- a/test/quantization/fx/test_equalize_fx.py +++ b/test/quantization/fx/test_equalize_fx.py @@ -1,9 +1,18 @@ import torch -from torch.testing._internal.common_quantization import QuantizationTestCase +import torch.nn as nn +from torch.quantization import default_qconfig +from torch.quantization.observer import MinMaxObserver +from torch.quantization.quantize_fx import prepare_fx from torch.quantization.fx._equalize import ( - _InputEqualizationObserver, _WeightEqualizationObserver, calculate_equalization_scale + _InputEqualizationObserver, + _WeightEqualizationObserver, + calculate_equalization_scale, + default_equalization_qconfig, ) +from torch.testing._internal.common_quantization import NodeSpec as ns +from torch.testing._internal.common_quantization import QuantizationTestCase + # Standard Libraries import numpy as np @@ -19,6 +28,9 @@ class TestEqualizeFx(QuantizationTestCase): weight_qscheme=st.sampled_from((torch.per_channel_affine, torch.per_channel_symmetric, torch.per_channel_affine_float_qparams))) def test_input_weight_observer(self, input_qdtype, input_qscheme, weight_qdtype, weight_qscheme): + """ Tests that the Input- and Weight- EqualizationObservers perform as expected + """ + input_obs = _InputEqualizationObserver(dtype=input_qdtype, qscheme=input_qscheme) weight_obs = _WeightEqualizationObserver(dtype=weight_qdtype, qscheme=weight_qscheme) @@ -110,3 +122,111 @@ def test_input_weight_observer(self, input_qdtype, input_qscheme, weight_qdtype, ref_scales, dtype=weight_qparams[0].dtype), atol=0.0001)) self.assertTrue(torch.allclose(weight_qparams[1], torch.tensor( ref_zero_points, dtype=weight_qparams[1].dtype), atol=1)) + + def test_input_weight_equalization_prepare(self): + """ Tests that graphs created after prepare_fx is as expected + """ + qconfig_dict = {"": None, + "object_type": [(nn.Linear, default_qconfig), (nn.functional.linear, default_qconfig)]} + + default_equalization_qconfig_dict = { + "": default_qconfig, + "object_type": [(nn.Linear, default_equalization_qconfig), + (nn.functional.linear, default_equalization_qconfig)] + } + + # Basic test with one linear layer + class LinearModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(1, 1) + + def forward(self, x): + return self.linear(x) + + linear_node_occurrence = { + ns.call_module(_InputEqualizationObserver): 1, + ns.call_module(MinMaxObserver): 2, + } + + # Test with two linear layers + class Linear2Module(nn.Module): + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(1, 1) + self.linear2 = nn.Linear(1, 1) + + def forward(self, x): + x = self.linear1(x) + x = self.linear2(x) + return x + + linear2_node_occurrence = { + ns.call_module(_InputEqualizationObserver): 2, + ns.call_module(MinMaxObserver): 3, + } + + class Linear(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.ones(5, 5) + self.b = torch.zeros(5) + + def forward(self, x): + return nn.functional.linear(x, self.w, self.b) + + # Test where we have two functional linear layers + class FunctionalLinearModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = Linear() + self.linear2 = Linear() + + def forward(self, x): + x = self.linear1(x) + x = self.linear2(x) + return x + + functionalLinear_node_occurrence = { + ns.call_module(_InputEqualizationObserver): 2, + ns.call_module(_WeightEqualizationObserver): 2, + ns.call_module(MinMaxObserver): 5, + } + + # Test where we have a Linear layer followed by a fp32 operation + # (conv layer) without a qconfig + class FunctionalLinearFP32Module(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = Linear() + self.conv = nn.Conv2d(3, 3, 1, 1) + self.linear2 = Linear() + + def forward(self, x): + x = self.linear1(x) + x = torch.add(x, 5) + x = self.linear2(x) + return x + + fp32_equalization_qconfig_dict = { + "": None, + "object_type": [(nn.Linear, default_equalization_qconfig), + (nn.functional.linear, default_equalization_qconfig)] + } + + functionalLinearFP32_node_occurrence = { + ns.call_module(_InputEqualizationObserver): 2, + ns.call_module(_WeightEqualizationObserver): 2, + ns.call_module(MinMaxObserver): 6, + } + + tests = [(LinearModule, default_equalization_qconfig_dict, linear_node_occurrence), + (Linear2Module, fp32_equalization_qconfig_dict, linear2_node_occurrence), + (FunctionalLinearModule, default_equalization_qconfig_dict, functionalLinear_node_occurrence), + (FunctionalLinearFP32Module, fp32_equalization_qconfig_dict, functionalLinearFP32_node_occurrence)] + + + for (M, equalization_qconfig_dict, node_occurrence) in tests: + m = M().eval() + prepared = prepare_fx(m, qconfig_dict, equalization_qconfig_dict=equalization_qconfig_dict) + self.checkGraphModuleNodes(prepared, expected_node_occurrence=node_occurrence) diff --git a/torch/quantization/fx/_equalize.py b/torch/quantization/fx/_equalize.py index 0869546759423..9259422fcc2b2 100644 --- a/torch/quantization/fx/_equalize.py +++ b/torch/quantization/fx/_equalize.py @@ -1,6 +1,8 @@ import torch import torch.nn as nn +from torch.fx.graph import Node + from ..observer import ( PerChannelMinMaxObserver, _with_args ) @@ -41,6 +43,8 @@ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine, if qscheme not in {torch.per_tensor_affine, torch.per_tensor_symmetric}: raise TypeError("Input qscheme must be per-tensor") + self.dtype = dtype + self.input_obs = PerChannelMinMaxObserver(ch_axis=1, dtype=dtype, qscheme=qscheme, quant_min=quant_min, @@ -120,6 +124,8 @@ def __init__(self, dtype=torch.qint8, qscheme=torch.per_tensor_affine, quant_min quant_max=None, factory_kwargs=None) -> None: super(_WeightEqualizationObserver, self).__init__() + self.dtype = dtype + self.weight_col_obs = PerChannelMinMaxObserver(ch_axis=1, dtype=dtype, qscheme=qscheme, quant_min=quant_min, @@ -259,3 +265,17 @@ def __new__(cls, input_activation=torch.nn.Identity, weight=torch.nn.Identity): dtype=torch.qint8, qscheme=torch.per_channel_symmetric) default_equalization_qconfig = EqualizationQConfig(input_activation=input_equalization_observer, weight=weight_equalization_observer) + +def node_supports_equalization(node: Node, modules) -> bool: + """ Checks if the current node supports equalization + Currently we only support nn.Linear and F.Linear layers + """ + if node.op == 'call_module': + return isinstance(modules[node.target], nn.Linear) + elif node.op == 'call_function': + return node.target == nn.functional.linear + return False + +def is_equalization_observer(observer: nn.Module) -> bool: + return (isinstance(observer, _InputEqualizationObserver) or + isinstance(observer, _WeightEqualizationObserver)) diff --git a/torch/quantization/fx/graph_module.py b/torch/quantization/fx/graph_module.py index 0b33221d25898..05a73bf041ee1 100644 --- a/torch/quantization/fx/graph_module.py +++ b/torch/quantization/fx/graph_module.py @@ -29,6 +29,7 @@ def __init__(self, root: Union[torch.nn.Module, Dict[str, Any]], graph: Graph, p '_patterns', '_qconfig_map', '_prepare_custom_config_dict', + '_equalization_qconfig_map', '_node_name_to_scope']).union(preserved_attr_names) preserved_attrs = {attr: getattr(root, attr) for attr in self.preserved_attr_names if hasattr(root, attr)} super().__init__(root, graph) diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index a2a7fe842c943..122d36aca68d0 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -30,6 +30,11 @@ from .quantization_types import Pattern +from ._equalize import ( + is_equalization_observer, + node_supports_equalization, +) + from .graph_module import ( ObservedGraphModule, ObservedStandaloneGraphModule, @@ -160,7 +165,10 @@ def insert_observer( if model_device: observer.to(model_device) # add observer module as attribute - prefix = node.name + '_activation_post_process_' + if is_equalization_observer(observer): + prefix = node.name + '_equalization_process_' + else: + prefix = node.name + '_activation_post_process_' get_new_observer_name = get_new_attr_name_with_prefix(prefix) observer_name = get_new_observer_name(model) setattr(model, observer_name, observer) @@ -278,8 +286,10 @@ def maybe_insert_input_observer_for_arg_or_kwarg( # regular flow for most nodes, except standalone modules is_weight = node_arg_is_weight(node, arg) assert qconfig is not None + act_post_process_ctr = qconfig.weight if is_weight else \ qconfig.activation + is_bias = node_arg_is_bias(node, arg) is_activation = not (is_weight or is_bias) weight_needs_obs = is_weight and weight_is_quantized(qconfig) and node.target not in NON_QUANTIZABLE_WEIGHT_OPS @@ -417,6 +427,45 @@ def maybe_insert_input_observers_for_node( node.args = tuple(new_args) node.kwargs = new_kwargs # type: ignore[assignment] +def maybe_insert_input_equalization_observers_for_node( + node: Node, + equalization_qconfig: Any, + model: torch.nn.Module, + modules: Dict[str, torch.nn.Module], + graph: Graph, + node_name_to_target_dtype: Dict[str, Any], +) -> None: + """ + If `node` needs to be equalized, find the input/weight observers it needs in + `equalization_qconfig`, creates them, and inserts it into `graph`. + + If `node` does not need an equalization observer, returns None. + """ + if equalization_qconfig is None or not node_supports_equalization(node, modules): + return + + new_args = [] + for arg in node.args: + if not isinstance(arg, Node) or node_arg_is_bias(node, arg): + new_args.append(arg) + continue + + is_weight = node_arg_is_weight(node, arg) + + act_eq_process_ctr = equalization_qconfig.weight if is_weight else \ + equalization_qconfig.input_activation + + new_eq_obs_mod = act_eq_process_ctr() + new_eq_obs_node = insert_observer( + arg, new_eq_obs_mod, model, modules, graph) + + # set the type, so the next node can read it + node_name_to_target_dtype[new_eq_obs_node.name] = node_name_to_target_dtype[arg.name] + + new_args.append(new_eq_obs_node) + + # assign the new args and kwargs to the node, inplace + node.args = tuple(new_args) def maybe_insert_output_observer_for_node( node: Node, @@ -701,6 +750,7 @@ def insert_observers_for_model( qconfig_map: Dict[str, QConfigAny], graph: Graph, prepare_custom_config_dict: Dict[str, Any], + equalization_config_map: Dict[str, Any], input_quantized_idxs: List[int], output_quantized_idxs: List[int], ) -> Optional[Node]: @@ -775,6 +825,7 @@ def insert_observers_for_model( # check for matches root_node, matched_nodes, pattern, qhandler, qconfig = matches.get( node.name, (None, None, None, None, None)) + equalization_qconfig = equalization_config_map.get(node.name, None) if node.op == 'placeholder': # if a graph input is in fp32, it does not need observation @@ -804,6 +855,11 @@ def insert_observers_for_model( node_name_to_target_dtype, qhandler, prepare_custom_config_dict) + # Insert equalization input observers if needed + maybe_insert_input_equalization_observers_for_node( + node, equalization_qconfig, model, modules, graph, + node_name_to_target_dtype) + is_last_node_of_pattern = root_node is node is_like_copy_node = \ (qhandler is not None and ( @@ -910,19 +966,22 @@ def save_state( qconfig_map: Dict[str, QConfigAny], node_name_to_scope: Dict[str, Tuple[str, type]], patterns: Dict[Pattern, QuantizeHandler], - prepare_custom_config_dict: Dict[str, Any] + prepare_custom_config_dict: Dict[str, Any], + equalization_qconfig_map: Dict[str, Any], ) -> None: observed._patterns = patterns # type: ignore[assignment] observed._qconfig_map = qconfig_map # type: ignore[assignment] observed._prepare_custom_config_dict = \ prepare_custom_config_dict # type: ignore[assignment] observed._node_name_to_scope = node_name_to_scope # type: ignore[assignment] + observed._equalization_qconfig_map = equalization_qconfig_map # type: ignore[assignment] def prepare( model: GraphModule, qconfig_dict: Any, node_name_to_scope: Dict[str, Tuple[str, type]], prepare_custom_config_dict: Optional[Dict[str, Any]] = None, + equalization_qconfig_dict: Optional[Dict[str, Any]] = None, is_standalone_module: bool = False) -> ObservedGraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. @@ -946,6 +1005,8 @@ def prepare( """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} + if equalization_qconfig_dict is None: + equalization_qconfig_dict = {} additional_quant_patterns = \ prepare_custom_config_dict.get("additional_quant_pattern", {}) @@ -963,6 +1024,7 @@ def prepare( get_default_quant_patterns(), additional_quant_patterns) convert_dict_to_ordered_dict(qconfig_dict) + convert_dict_to_ordered_dict(equalization_qconfig_dict) flattened_qconfig_dict = get_flattened_qconfig_dict(qconfig_dict) # TODO: support regex as well propagate_qconfig_(model, flattened_qconfig_dict) @@ -981,6 +1043,7 @@ def prepare( modules = dict(model.named_modules()) # fill qconfig_map, a map from node name to qconfig, used in find_matches + equalization_qconfig_map = generate_qconfig_map(model, modules, model.graph, equalization_qconfig_dict, node_name_to_scope) qconfig_map = generate_qconfig_map(model, modules, model.graph, qconfig_dict, node_name_to_scope) # match the patterns that will get quantized @@ -1008,9 +1071,11 @@ def prepare( result_node = insert_observers_for_model( model, modules, matches, qconfig_map, model.graph, prepare_custom_config_dict, + equalization_qconfig_map, input_quantized_idxs, output_quantized_idxs) - save_state(model, qconfig_map, node_name_to_scope, patterns, prepare_custom_config_dict) + save_state(model, qconfig_map, node_name_to_scope, patterns, + prepare_custom_config_dict, equalization_qconfig_map) preserved_attributes = set(prepare_custom_config_dict.get("preserved_attributes", [])) model = ObservedGraphModule(model, model.graph, preserved_attributes) if is_standalone_module: diff --git a/torch/quantization/quantize_fx.py b/torch/quantization/quantize_fx.py index 82fd0b49c099f..aa8edbba64e49 100644 --- a/torch/quantization/quantize_fx.py +++ b/torch/quantization/quantize_fx.py @@ -141,10 +141,12 @@ def create_node(self, kind : str, target : Target, def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, prepare_custom_config_dict: Dict[str, Any] = None, + equalization_qconfig_dict: Dict[str, Any] = None, is_standalone_module: bool = False) -> ObservedGraphModule: r""" Internal helper function for prepare_fx Args: - `model`, `qconfig_dict`, `prepare_custom_config_dict`: see docs for :func:`~torch.quantization.prepare_fx` + `model`, `qconfig_dict`, `prepare_custom_config_dict`, `equalization_qonfig_dict`: + see docs for :func:`~torch.quantization.prepare_fx` `is_standalone_module`: a boolean flag indicates whether we are quantizing a standalone module or not, a standalone module is a submodule of the parent module that is not inlined in the @@ -154,9 +156,12 @@ def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, """ if prepare_custom_config_dict is None: prepare_custom_config_dict = {} + if equalization_qconfig_dict is None: + equalization_qconfig_dict = {} check_is_valid_qconfig_dict(qconfig_dict) check_is_valid_prepare_custom_config_dict(prepare_custom_config_dict) + check_is_valid_qconfig_dict(equalization_qconfig_dict) skipped_module_names = prepare_custom_config_dict.get("non_traceable_module_name", []) skipped_module_classes = prepare_custom_config_dict.get("non_traceable_module_class", []) @@ -188,6 +193,7 @@ def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, qconfig_dict, tracer.node_name_to_scope, prepare_custom_config_dict=prepare_custom_config_dict, + equalization_qconfig_dict=equalization_qconfig_dict, is_standalone_module=is_standalone_module) for attr_name in preserved_attributes: @@ -259,7 +265,8 @@ def fuse_fx(model: torch.nn.Module, def prepare_fx( model: torch.nn.Module, qconfig_dict: Any, - prepare_custom_config_dict: Dict[str, Any] = None) -> ObservedGraphModule: + prepare_custom_config_dict: Dict[str, Any] = None, + equalization_qconfig_dict: Dict[str, Any] = None) -> ObservedGraphModule: r""" Prepare a model for post training static quantization Args: @@ -381,6 +388,10 @@ def prepare_fx( # not used in the code, these attributes will also persist through deepcopy "preserved_attributes": ["preserved_attr"], } + `equalization_qconfig_dict`: equalization_qconfig_dict is a dictionary + with a similar structure as qconfig_dict except it will contain + configurations specific to equalization techniques such as input-weight + equalization. Return: @@ -409,7 +420,7 @@ def calibrate(model, data_loader): torch._C._log_api_usage_once("quantization_api.quantize_fx.prepare_fx") assert not model.training, 'prepare_fx only works for models in ' + \ 'eval mode' - return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict) + return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, equalization_qconfig_dict) def prepare_qat_fx( model: torch.nn.Module, qconfig_dict: Any, From c0b7c59e559c499db1b097191e404381ca1c35b6 Mon Sep 17 00:00:00 2001 From: Angela Yi Date: Wed, 16 Jun 2021 22:30:02 -0700 Subject: [PATCH 179/305] [quant] Equalization Observer modifications (#59953) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59953 The following modifications were made to the equalization observers due to design changes: - [InputEqualizationObserver] Replaced `calculate_qparams()` with `calculate_scaled_minmax()` since we will need to return the scaled min/max values to update the following input quantization observer - [WeightEqualizationObserver] We no longer need a row observer since this will be taken care of by the following weight quantization observer - [WeightEqualizationObserver] Following the previous comment, we no longer need to calculate the scaled qparam values. Instead, we will use the equalization scale to later scale the weights and the qparams will be taken care of by the weight quantization observer. Test Plan: `python test/test_quantization.py TestEqualizeFx.test_input_weight_eq_observer` Imported from OSS Reviewed By: supriyar Differential Revision: D29135332 fbshipit-source-id: be7e468273c8b62fc183b1e1ec50f6bd6d8cf831 --- test/quantization/fx/test_equalize_fx.py | 94 +++++++++++++----------- torch/quantization/fx/_equalize.py | 85 ++++----------------- torch/quantization/observer.py | 23 +----- torch/quantization/utils.py | 32 ++++++++ 4 files changed, 100 insertions(+), 134 deletions(-) diff --git a/test/quantization/fx/test_equalize_fx.py b/test/quantization/fx/test_equalize_fx.py index 3980d1bf42df5..99b74379c9d27 100644 --- a/test/quantization/fx/test_equalize_fx.py +++ b/test/quantization/fx/test_equalize_fx.py @@ -1,7 +1,7 @@ import torch import torch.nn as nn from torch.quantization import default_qconfig -from torch.quantization.observer import MinMaxObserver +from torch.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver from torch.quantization.quantize_fx import prepare_fx from torch.quantization.fx._equalize import ( _InputEqualizationObserver, @@ -27,12 +27,12 @@ class TestEqualizeFx(QuantizationTestCase): weight_qdtype=st.sampled_from((torch.qint8, torch.quint8)), weight_qscheme=st.sampled_from((torch.per_channel_affine, torch.per_channel_symmetric, torch.per_channel_affine_float_qparams))) - def test_input_weight_observer(self, input_qdtype, input_qscheme, weight_qdtype, weight_qscheme): + def test_input_weight_eq_observer(self, input_qdtype, input_qscheme, weight_qdtype, weight_qscheme): """ Tests that the Input- and Weight- EqualizationObservers perform as expected """ - input_obs = _InputEqualizationObserver(dtype=input_qdtype, qscheme=input_qscheme) - weight_obs = _WeightEqualizationObserver(dtype=weight_qdtype, qscheme=weight_qscheme) + input_eq_obs = _InputEqualizationObserver(dtype=input_qdtype, qscheme=input_qscheme) + weight_eq_obs = _WeightEqualizationObserver(dtype=weight_qdtype, qscheme=weight_qscheme) width = np.random.randint(1, 10) x_height = np.random.randint(2, 10) @@ -41,82 +41,90 @@ def test_input_weight_observer(self, input_qdtype, input_qscheme, weight_qdtype, x = (np.random.random(size=(x_height, width)) * 10).round(decimals=2).astype(np.float32) w = (np.random.random(size=(w_height, width)) * 10).round(decimals=2).astype(np.float32) - ret_x = input_obs(torch.tensor(x)) - ret_w = weight_obs(torch.tensor(w)) + ret_x = input_eq_obs(torch.tensor(x)) + ret_w = weight_eq_obs(torch.tensor(w)) self.assertEqual((ret_x, ret_w), (x, w)) # Check the min/max input columns are correct ref_min_inputs = x.min(axis=0) ref_max_inputs = x.max(axis=0) - self.assertEqual(input_obs.get_input_minmax(), (ref_min_inputs, ref_max_inputs)) + self.assertEqual(input_eq_obs.get_input_minmax(), (ref_min_inputs, ref_max_inputs)) # Check the min/max weight columns are correct ref_min_weights_col = w.min(axis=0) ref_max_weights_col = w.max(axis=0) - self.assertEqual(weight_obs.get_weight_col_minmax(), (ref_min_weights_col, ref_max_weights_col)) - - # Check the min/max weight rows are correct - ref_min_weights_row = w.min(axis=1) - ref_max_weights_row = w.max(axis=1) - self.assertEqual(weight_obs.get_weight_row_minmax(), (ref_min_weights_row, ref_max_weights_row)) - - # Check the column indices of the min/max weight rows are correct - ref_min_weights_ind = w.argmin(axis=1) - ref_max_weights_ind = w.argmax(axis=1) - self.assertEqual((weight_obs.min_weights_ind, weight_obs.max_weights_ind), - (ref_min_weights_ind, ref_max_weights_ind)) + self.assertEqual(weight_eq_obs.get_weight_col_minmax(), (ref_min_weights_col, ref_max_weights_col)) # Check the equalization scale is correct - equalization_scale = calculate_equalization_scale(input_obs, weight_obs) + equalization_scale = calculate_equalization_scale(input_eq_obs, weight_eq_obs) ref_equalization_scale = np.sqrt((ref_max_weights_col - ref_min_weights_col) / (ref_max_inputs - ref_min_inputs)) self.assertEqual(equalization_scale, ref_equalization_scale) - input_obs.set_equalization_scale(equalization_scale) - weight_obs.set_equalization_scale(equalization_scale) + input_eq_obs.set_equalization_scale(equalization_scale) + weight_eq_obs.set_equalization_scale(equalization_scale) - # check the input scale/zero-point values - input_qparams = input_obs.calculate_qparams() + # Check the input scale/zero-point values + min_input_scaled, max_input_scaled = input_eq_obs.calculate_scaled_minmax() + input_quant_obs = MinMaxObserver(dtype=input_qdtype, qscheme=input_qscheme) + input_quant_obs.min_val = min_input_scaled + input_quant_obs.max_val = max_input_scaled + input_qparams = input_quant_obs.calculate_qparams() - min_input_scaled = np.min(ref_min_inputs * ref_equalization_scale) - min_input_scaled = min(0, min_input_scaled) - max_input_scaled = np.max(ref_max_inputs * ref_equalization_scale) - max_input_scaled = max(0, max_input_scaled) + ref_min_input_scaled = np.min(ref_min_inputs * ref_equalization_scale) + ref_min_input_scaled = min(0, ref_min_input_scaled) + ref_max_input_scaled = np.max(ref_max_inputs * ref_equalization_scale) + ref_max_input_scaled = max(0, ref_max_input_scaled) if input_qscheme == torch.per_tensor_symmetric: - ref_scale = 2 * max(abs(min_input_scaled), max_input_scaled) / 255 + ref_scale = 2 * max(abs(ref_min_input_scaled), ref_max_input_scaled) / 255 ref_zero_point = 0 if input_qdtype is torch.qint8 else 128 else: - ref_scale = (max_input_scaled - min_input_scaled) / 255 + ref_scale = (ref_max_input_scaled - ref_min_input_scaled) / 255 ref_zero_point = -128 if input_qdtype is torch.qint8 else 0 self.assertEqual(input_qparams[0].item(), ref_scale, atol=1e-5, rtol=0) self.assertEqual(input_qparams[1].item(), ref_zero_point) - # check the weight scale/zero-point values - weight_qparams = weight_obs.calculate_qparams() + # During input-weight equalization, we will scale the weights so that + # the following weight quantized observer will have the correct scaled qparams + # Check the weight scale/zero-point values of the quantized observer + weight_quant_obs = PerChannelMinMaxObserver(dtype=weight_qdtype, qscheme=weight_qscheme) + + # Scale the weights for input-weight equalization + ref_w_scaled = w * np.reciprocal(ref_equalization_scale) + w_scaled = torch.mul(torch.tensor(w), torch.reciprocal(equalization_scale)) + self.assertEqual(ref_w_scaled, w_scaled) + + # Call forward on the weight quantization observer + weight_quant_obs(w_scaled) + + # Check the min/max weight rows are correct + ref_min_weights_scaled = ref_w_scaled.min(axis=1) + ref_max_weights_scaled = ref_w_scaled.max(axis=1) + self.assertEqual(weight_quant_obs.min_vals, ref_min_weights_scaled) + self.assertEqual(weight_quant_obs.max_vals, ref_max_weights_scaled) - min_weights_scaled = ref_min_weights_row * (1 / ref_equalization_scale[ref_min_weights_ind]) - max_weights_scaled = ref_max_weights_row * (1 / ref_equalization_scale[ref_max_weights_ind]) + weight_qparams = weight_quant_obs.calculate_qparams() if weight_qscheme == torch.per_channel_symmetric: - min_weights_scaled = np.minimum(np.zeros(min_weights_scaled.shape), min_weights_scaled) - max_weights_scaled = np.maximum(np.zeros(max_weights_scaled.shape), max_weights_scaled) + ref_min_weights_scaled = np.minimum(np.zeros(ref_min_weights_scaled.shape), ref_min_weights_scaled) + ref_max_weights_scaled = np.maximum(np.zeros(ref_max_weights_scaled.shape), ref_max_weights_scaled) - ref_scales = 2 * np.maximum(np.abs(min_weights_scaled), max_weights_scaled) / 255 + ref_scales = 2 * np.maximum(np.abs(ref_min_weights_scaled), ref_max_weights_scaled) / 255 ref_zero_points = np.zeros_like( ref_scales) if weight_qdtype is torch.qint8 else np.ones_like(ref_scales) * 128 elif weight_qscheme == torch.per_channel_affine_float_qparams: - ref_scales = (max_weights_scaled - min_weights_scaled) / 255 + ref_scales = (ref_max_weights_scaled - ref_min_weights_scaled) / 255 ref_scales = np.where(ref_scales > 1e-7, ref_scales, np.ones_like(ref_scales)) - ref_zero_points = -1 * min_weights_scaled / ref_scales + ref_zero_points = -1 * ref_min_weights_scaled / ref_scales else: - min_weights_scaled = np.minimum(np.zeros_like(min_weights_scaled), min_weights_scaled) - max_weights_scaled = np.maximum(np.zeros_like(max_weights_scaled), max_weights_scaled) + ref_min_weights_scaled = np.minimum(np.zeros_like(ref_min_weights_scaled), ref_min_weights_scaled) + ref_max_weights_scaled = np.maximum(np.zeros_like(ref_max_weights_scaled), ref_max_weights_scaled) - ref_scales = (max_weights_scaled - min_weights_scaled) / 255 + ref_scales = (ref_max_weights_scaled - ref_min_weights_scaled) / 255 ref_zero_points = -128 if weight_qdtype is torch.qint8 else 0 - ref_zero_points = ref_zero_points - np.round(min_weights_scaled / ref_scales) + ref_zero_points = ref_zero_points - np.round(ref_min_weights_scaled / ref_scales) self.assertTrue(torch.allclose(weight_qparams[0], torch.tensor( ref_scales, dtype=weight_qparams[0].dtype), atol=0.0001)) diff --git a/torch/quantization/fx/_equalize.py b/torch/quantization/fx/_equalize.py index 9259422fcc2b2..9fbe1fc9bcbb8 100644 --- a/torch/quantization/fx/_equalize.py +++ b/torch/quantization/fx/_equalize.py @@ -4,8 +4,10 @@ from torch.fx.graph import Node from ..observer import ( - PerChannelMinMaxObserver, _with_args + PerChannelMinMaxObserver, + _with_args, ) +from ..utils import check_min_max_valid from collections import namedtuple import warnings @@ -44,6 +46,7 @@ def __init__(self, dtype=torch.quint8, qscheme=torch.per_tensor_affine, raise TypeError("Input qscheme must be per-tensor") self.dtype = dtype + self.qscheme = qscheme self.input_obs = PerChannelMinMaxObserver(ch_axis=1, dtype=dtype, qscheme=qscheme, @@ -66,17 +69,16 @@ def get_input_minmax(self): def set_equalization_scale(self, equalization_scale): self.equalization_scale = equalization_scale - def calculate_qparams(self): + def calculate_scaled_minmax(self): r""" - Returns the scale/zero_point for the input and weight rows + Returns the scaled min/max inputs """ - if self.equalization_scale.nelement() == 0: warnings.warn( "Must call calculate_scale before calling calculate_qparams.\ - Returning default scale and zero point. " + Returning default min and max input." ) - return torch.tensor([1.0]), torch.tensor([0]), torch.tensor([1.0]), torch.tensor([0]) + return torch.tensor([0]), torch.tensor([0]) # Calculate qparams for the scaled min/max inputs # Scale the input by the equalization scale located at the same column @@ -84,9 +86,8 @@ def calculate_qparams(self): (min_inputs, max_inputs) = self.get_input_minmax() min_input_scaled = torch.min(torch.mul(min_inputs, self.equalization_scale)) max_input_scaled = torch.max(torch.mul(max_inputs, self.equalization_scale)) - (scale_input, zero_point_input) = self.input_obs._calculate_qparams(min_input_scaled, max_input_scaled) - return scale_input, zero_point_input + return min_input_scaled, max_input_scaled with_args = classmethod(_with_args) @@ -125,6 +126,8 @@ def __init__(self, dtype=torch.qint8, qscheme=torch.per_tensor_affine, quant_min super(_WeightEqualizationObserver, self).__init__() self.dtype = dtype + self.qscheme = qscheme + self.ch_axis = 0 self.weight_col_obs = PerChannelMinMaxObserver(ch_axis=1, dtype=dtype, qscheme=qscheme, @@ -132,80 +135,20 @@ def __init__(self, dtype=torch.qint8, qscheme=torch.per_tensor_affine, quant_min quant_max=quant_max, factory_kwargs=factory_kwargs) - self.weight_row_obs = PerChannelMinMaxObserver(ch_axis=0, dtype=dtype, - qscheme=qscheme, - quant_min=quant_min, - quant_max=quant_max, - factory_kwargs=factory_kwargs) - self.equalization_scale = torch.empty(0) def forward(self, w_orig): # TODO: Allow for convoluational layers if not (w_orig.ndim == 2): raise ValueError("WeightEqualizationObserver only supports Linear layers") - - return self._forward(w_orig) - - def _forward(self, w_orig): - r""" - Calculates the min/max values of each weight column and weight row. - """ - - w_orig = self.weight_col_obs(w_orig) - w_orig = self.weight_row_obs(w_orig) - - # Calculate the column indices of the min/max weight in each row - num_row, _ = w_orig.shape - min_weights_ind = [] - max_weights_ind = [] - for i in range(num_row): - min_weights_ind.append(torch.nonzero(w_orig[i] == self.weight_row_obs.min_vals[i])[0][0]) - max_weights_ind.append(torch.nonzero(w_orig[i] == self.weight_row_obs.max_vals[i])[0][0]) - self.min_weights_ind = torch.tensor(min_weights_ind) - self.max_weights_ind = torch.tensor(max_weights_ind) - - return w_orig + return self.weight_col_obs(w_orig) def get_weight_col_minmax(self): return (self.weight_col_obs.min_vals, self.weight_col_obs.max_vals) - def get_weight_row_minmax(self): - return (self.weight_row_obs.min_vals, self.weight_row_obs.max_vals) - def set_equalization_scale(self, equalization_scale): self.equalization_scale = equalization_scale - def calculate_qparams(self): - r""" - Returns the scale/zero_point for the input and weight rows - """ - - if self.equalization_scale.nelement() == 0: - warnings.warn( - "Must call calculate_scale before calling calculate_qparams.\ - Returning default scale and zero point. " - ) - return torch.tensor([1.0]), torch.tensor([0]), torch.tensor([1.0]), torch.tensor([0]) - - if self.min_weights_ind is None or self.max_weights_ind is None: - warnings.warn( - "Must find the column indicies of the minimum of each row in the \ - weights in order to calculate the qparams calculate the \ - qparams. Returning default scale and zero point. " - ) - return torch.tensor([1.0]), torch.tensor([0]), torch.tensor([1.0]), torch.tensor([0]) - - # Calculate the qparams for weights by using the rows - # Scale the weight rows by the reciprocal of the equalization scale - # located at the same column index - (min_weights, max_weights) = self.get_weight_row_minmax() - min_weights_scaled = torch.mul(min_weights, torch.reciprocal(self.equalization_scale[self.min_weights_ind])) - max_weights_scaled = torch.mul(max_weights, torch.reciprocal(self.equalization_scale[self.max_weights_ind])) - (scale_weight, zero_point_weight) = self.weight_row_obs._calculate_qparams(min_weights_scaled, max_weights_scaled) - - return scale_weight, zero_point_weight - with_args = classmethod(_with_args) @@ -222,6 +165,9 @@ def calculate_equalization_scale(input_obs: _InputEqualizationObserver, (min_inputs, max_inputs) = input_obs.get_input_minmax() (min_weights, max_weights) = weight_obs.get_weight_col_minmax() + if not (check_min_max_valid(min_inputs, max_inputs) and check_min_max_valid(min_weights, max_weights)): + return torch.tensor(1) + if not (min_inputs.shape == min_weights.shape): raise ValueError( "Input and Weight must have the same column dimension. " + @@ -229,7 +175,6 @@ def calculate_equalization_scale(input_obs: _InputEqualizationObserver, ) equalization_scale = torch.sqrt((max_weights - min_weights) / (max_inputs - min_inputs)) - return equalization_scale diff --git a/torch/quantization/observer.py b/torch/quantization/observer.py index 26779590ca6d7..d57e90c4c49cd 100644 --- a/torch/quantization/observer.py +++ b/torch/quantization/observer.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from .utils import check_min_max_valid class _PartialWrapper(object): @@ -271,29 +272,9 @@ def _calculate_qparams( scales: Scales tensor of shape (#channels,) zero_points: Zero points tensor of shape (#channels,) """ - if min_val.numel() == 0 or max_val.numel() == 0: - warnings.warn( - "must run observer before calling calculate_qparams.\ - Returning default scale and zero point " - ) + if not check_min_max_valid(min_val, max_val): return torch.tensor([1.0]), torch.tensor([0]) - if min_val.dim() == 0 or max_val.dim() == 0: - if min_val == float("inf") and max_val == float("-inf"): - warnings.warn( - "must run observer before calling calculate_qparams.\ - Returning default scale and zero point " - ) - return torch.tensor([1.0]), torch.tensor([0]) - - assert min_val <= max_val, "min {} should be less than max {}".format( - min_val, max_val - ) - else: - assert torch.all( - min_val <= max_val - ), "min {} should be less than max {}".format(min_val, max_val) - quant_min, quant_max = self._calculate_qmin_qmax() min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) diff --git a/torch/quantization/utils.py b/torch/quantization/utils.py index 131339dcfa480..5544c9718c742 100644 --- a/torch/quantization/utils.py +++ b/torch/quantization/utils.py @@ -1,6 +1,8 @@ """ Utils shared by different modes of quantization (eager/graph) """ +import warnings + import torch from .quant_type import QuantType, quant_type_to_str @@ -101,3 +103,33 @@ def get_quant_type(qconfig): raise Exception("Unrecognized dtype combination in get_quant_type: activation({})," "weight({})".format(activation.dtype, weight.dtype)) + +def check_min_max_valid(min_val: torch.Tensor, max_val: torch.Tensor) -> bool: + """ Checks if the given minimum and maximum values are valid, meaning that + they exist and the min value is less than the max value. + """ + if min_val.numel() == 0 or max_val.numel() == 0: + warnings.warn( + "must run observer before calling calculate_qparams. " + + "Returning default values." + ) + return False + + if min_val.dim() == 0 or max_val.dim() == 0: + if min_val == float("inf") and max_val == float("-inf"): + warnings.warn( + "must run observer before calling calculate_qparams. " + + "Returning default values." + ) + + return False + + assert min_val <= max_val, "min {} should be less than max {}".format( + min_val, max_val + ) + else: + assert torch.all( + min_val <= max_val + ), "min {} should be less than max {}".format(min_val, max_val) + + return True From 0e7b5ea6c003b763603d8ca0fe12f476fdf9bf32 Mon Sep 17 00:00:00 2001 From: Peter Bell Date: Wed, 16 Jun 2021 22:49:27 -0700 Subject: [PATCH 180/305] nonzero: Default to transposed output strides (#59370) Summary: Fixes https://github.com/pytorch/pytorch/issues/46224 cc ailzhang Pull Request resolved: https://github.com/pytorch/pytorch/pull/59370 Reviewed By: ezyang Differential Revision: D29143842 Pulled By: ngimel fbshipit-source-id: 5aa7a247b4a70cd816d0eed368ab4c445568c986 --- aten/src/ATen/native/TensorAdvancedIndexing.cpp | 5 ++++- aten/src/ATen/native/TensorConversions.cpp | 2 +- c10/core/Device.h | 5 +++++ c10/core/TensorImpl.h | 2 +- 4 files changed, 11 insertions(+), 3 deletions(-) diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp index 3484681e56607..3f85dae9677c9 100644 --- a/aten/src/ATen/native/TensorAdvancedIndexing.cpp +++ b/aten/src/ATen/native/TensorAdvancedIndexing.cpp @@ -1568,7 +1568,10 @@ Tensor& nonzero_out_cpu(const Tensor& self, Tensor& result) { const auto self_sizes = self.sizes(); const auto total_nonzero = thread_count_nonzero.back(); const int64_t ndim = self_sizes.size(); - resize_output(result, {total_nonzero, ndim}); + if (resize_output(result, {total_nonzero, ndim})) { + // Default to fortran-contiguous output (see gh-46224) + result.as_strided_({total_nonzero, ndim}, {1, total_nonzero}); + } if (result.numel() == 0) { return result; diff --git a/aten/src/ATen/native/TensorConversions.cpp b/aten/src/ATen/native/TensorConversions.cpp index 3e4a86a8155cf..876f9f867dccb 100644 --- a/aten/src/ATen/native/TensorConversions.cpp +++ b/aten/src/ATen/native/TensorConversions.cpp @@ -33,7 +33,7 @@ static inline Tensor to_impl(const Tensor& self, const TensorOptions& options, b (options.layout() == c10::kStrided)); if (memory_format == MemoryFormat::Preserve) { - if (self.is_non_overlapping_and_dense()) { + if (self.is_non_overlapping_and_dense() && options.device().supports_as_strided()) { // Copy all strides auto r = at::empty_strided(self.sizes(), self.strides(), diff --git a/c10/core/Device.h b/c10/core/Device.h index 3f904b389e464..3cacbb480cb4a 100644 --- a/c10/core/Device.h +++ b/c10/core/Device.h @@ -96,6 +96,11 @@ struct C10_API Device final { return type_ == DeviceType::CPU; } + /// Return true if the device supports arbirtary strides. + bool supports_as_strided() const noexcept { + return type_ != DeviceType::XLA; + } + /// Same string as returned from operator<<. std::string str() const; diff --git a/c10/core/TensorImpl.h b/c10/core/TensorImpl.h index 458090d446c77..6008826600bb8 100644 --- a/c10/core/TensorImpl.h +++ b/c10/core/TensorImpl.h @@ -914,7 +914,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target { * It can be expanded as needed in the future, e.g sparse Tensor. */ inline bool support_as_strided() const { - return device().type() != at::kXLA; + return device().supports_as_strided(); } // ~~~~~ Autograd API ~~~~~ From 9e79a8a54f3e06ed613a0194e640c4622412dcfd Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Thu, 17 Jun 2021 00:23:45 -0700 Subject: [PATCH 181/305] [iOS GPU][MaskRCNN] Force the temporaryImage to become static when doing synchronization (#60155) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60155 For intermediate tensors, we need to convert them to static images when doing GPU -> CPU synchronization. ghstack-source-id: 131540760 Test Plan: - CI - buck test pp-macos Reviewed By: SS-JIA Differential Revision: D29126278 fbshipit-source-id: cd50b5f104e0161ec7fcfcc2c51785f241e48704 --- .../native/metal/mpscnn/MPSImageWrapper.mm | 4 +++ .../src/ATen/native/metal/ops/MetalReshape.mm | 30 +++++++++++-------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm index ba1e4d01ce015..fd02432c75a27 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm @@ -118,6 +118,10 @@ - (void)endSynchronization:(NSError*)error { TORCH_CHECK(_buffer, "Allocate GPU memory failed!"); } copyToMetalBuffer(_commandBuffer, _buffer, _image); + if (_image.isTemporaryImage && _image.readCount != 0) { + _image = + createStaticImage((MPSTemporaryImage*)_image, _commandBuffer, false); + } } void MPSImageWrapper::synchronize() { diff --git a/aten/src/ATen/native/metal/ops/MetalReshape.mm b/aten/src/ATen/native/metal/ops/MetalReshape.mm index e5525dfb1edc5..817b81e85a538 100644 --- a/aten/src/ATen/native/metal/ops/MetalReshape.mm +++ b/aten/src/ATen/native/metal/ops/MetalReshape.mm @@ -32,18 +32,18 @@ Tensor view(const Tensor& input, IntArrayRef size) { MetalTensorImplStorage mt{inferred_size, stride_value}; mt.texture()->allocateTemporaryStorage(inferred_size, commandBuffer); MPSImage* Y = mt.texture()->image(); - id state = [[MPSCNNContext sharedInstance] - specializedPipelineState:"reshape" - Constants:@[ - @(Y.height), - @(Y.width), - @(Y.featureChannels), - @(Y.numberOfImages), - @(X.height), - @(X.width), - @(X.featureChannels), - @(X.numberOfImages), - ]]; + id state = + [[MPSCNNContext sharedInstance] specializedPipelineState:"reshape" + Constants:@[ + @(Y.height), + @(Y.width), + @(Y.featureChannels), + @(Y.numberOfImages), + @(X.height), + @(X.width), + @(X.featureChannels), + @(X.numberOfImages), + ]]; id encoder = [commandBuffer.buffer computeCommandEncoder]; [encoder setComputePipelineState:state]; @@ -95,7 +95,13 @@ Tensor flatten_using_ints( return input.reshape(shape); } +Tensor detach(const Tensor& input) { + TORCH_CHECK(input.is_metal()); + return input; +} + TORCH_LIBRARY_IMPL(aten, Metal, m) { + m.impl("detach", TORCH_FN(detach)); m.impl("view", TORCH_FN(view)); m.impl("reshape", TORCH_FN(reshape)); m.impl("flatten.using_ints", TORCH_FN(flatten_using_ints)); From a0ad4c24d1eafad579e949261db980c5c79bf2c7 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Thu, 17 Jun 2021 00:34:03 -0700 Subject: [PATCH 182/305] MAINT Migrates rrelu_with_noise from THC to ATen on Cuda (#57864) Summary: Fixes https://github.com/pytorch/pytorch/issues/24618 Related to https://github.com/pytorch/pytorch/issues/24507
Benchmark script: ```py import torch import torch.nn as nn import time torch.manual_seed(0) def _time(): torch.cuda.synchronize() return time.time() device = "cuda" m = nn.RReLU().cuda() for n in [100, 10_000, 100_000]: fwd_t = 0 bwd_t = 0 input = torch.randn(128, n, device=device) grad_output = torch.ones(128, n, device=device) for i in range(10000): t1 = _time() output = m(input) t2 = _time() fwd_t = fwd_t + (t2 -t1) fwd_avg = fwd_t / 10000 * 1000 print(f"input size(128, {n}) forward time is {fwd_avg:.2f} (ms)") ```
### Results from benchmark: #### This PR ``` input size(128, 100) forward time is 0.01 (ms) input size(128, 10000) forward time is 0.06 (ms) input size(128, 100000) forward time is 0.54 (ms) ``` #### On master ``` input size(128, 100) forward time is 0.01 (ms) input size(128, 10000) forward time is 0.08 (ms) input size(128, 100000) forward time is 0.66 (ms) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/57864 Reviewed By: H-Huang Differential Revision: D29177169 Pulled By: ngimel fbshipit-source-id: 4572133db06f143d27e70a91ade977ea962c8f77 --- BUILD.bazel | 1 - aten/src/ATen/LegacyTHFunctionsCUDA.h | 3 - aten/src/ATen/cuda/LegacyTHFunctionsCUDA.cpp | 106 ----------- aten/src/ATen/native/cuda/Activation.cu | 177 +++++++++++++++++++ aten/src/ATen/native/native_functions.yaml | 6 +- aten/src/THCUNN/CMakeLists.txt | 1 - aten/src/THCUNN/RReLU.cu | 107 ----------- aten/src/THCUNN/generic/RReLU.cu | 72 -------- aten/src/THCUNN/generic/THCUNN.h | 11 -- 9 files changed, 180 insertions(+), 304 deletions(-) delete mode 100644 aten/src/THCUNN/RReLU.cu delete mode 100644 aten/src/THCUNN/generic/RReLU.cu diff --git a/BUILD.bazel b/BUILD.bazel index 1f840807ec46d..45d71a8d4626e 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -421,7 +421,6 @@ filegroup( "aten/src/THCUNN/LogSigmoid.cu.cc", "aten/src/THCUNN/MultiLabelMarginCriterion.cu.cc", "aten/src/THCUNN/MultiMarginCriterion.cu.cc", - "aten/src/THCUNN/RReLU.cu.cc", "aten/src/THCUNN/SoftMarginCriterion.cu.cc", "aten/src/THCUNN/SoftPlus.cu.cc", "aten/src/THCUNN/SoftShrink.cu.cc", diff --git a/aten/src/ATen/LegacyTHFunctionsCUDA.h b/aten/src/ATen/LegacyTHFunctionsCUDA.h index 785333c30ab8c..5dea6aff4d68c 100644 --- a/aten/src/ATen/LegacyTHFunctionsCUDA.h +++ b/aten/src/ATen/LegacyTHFunctionsCUDA.h @@ -51,10 +51,7 @@ std::tuple _thnn_log_sigmoid_forward_out(const Tensor & self, std::tuple _thnn_log_sigmoid_forward(const Tensor & self); Tensor & _thnn_log_sigmoid_backward_out(const Tensor & grad_output, const Tensor & self, const Tensor & buffer, Tensor & grad_input); Tensor _thnn_log_sigmoid_backward(const Tensor & grad_output, const Tensor & self, const Tensor & buffer); -Tensor & _thnn_rrelu_with_noise_forward_out(const Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator, Tensor & output); -Tensor _thnn_rrelu_with_noise_forward(const Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator); Tensor _thnn_rrelu_with_noise_backward(const Tensor & grad_output, const Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training); -Tensor & _thnn_rrelu_with_noise_forward_(Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator); std::tuple _thnn_conv2d_forward_out(const Tensor & self, const Tensor & weight, IntArrayRef kernel_size, const c10::optional& bias_opt, IntArrayRef stride, IntArrayRef padding, Tensor & output, Tensor & columns, Tensor & ones); std::tuple _thnn_conv2d_forward(const Tensor & self, const Tensor & weight, IntArrayRef kernel_size, const optional & bias, IntArrayRef stride, IntArrayRef padding); std::tuple _thnn_conv2d_backward_out(Tensor & grad_input, Tensor & grad_weight, Tensor & grad_bias, const Tensor & grad_output, const Tensor & self, const Tensor & weight, IntArrayRef kernel_size, IntArrayRef stride, IntArrayRef padding, const Tensor & columns, const Tensor & ones); diff --git a/aten/src/ATen/cuda/LegacyTHFunctionsCUDA.cpp b/aten/src/ATen/cuda/LegacyTHFunctionsCUDA.cpp index 5b0208a9ee050..a93e000dcb999 100644 --- a/aten/src/ATen/cuda/LegacyTHFunctionsCUDA.cpp +++ b/aten/src/ATen/cuda/LegacyTHFunctionsCUDA.cpp @@ -1285,112 +1285,6 @@ Tensor _thnn_log_sigmoid_backward(const Tensor & grad_output, const Tensor & sel } return grad_input; } -Tensor & _thnn_rrelu_with_noise_forward_out(const Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator, Tensor & output) { - const OptionalDeviceGuard device_guard(device_of(self)); - auto dispatch_scalar_type = infer_scalar_type(self); - - switch (dispatch_scalar_type) { - case ScalarType::Double: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - auto output_ = checked_dense_tensor_unwrap(output, "output", 6, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - THNN_CudaDoubleRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - case ScalarType::Float: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - auto output_ = checked_dense_tensor_unwrap(output, "output", 6, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - THNN_CudaRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - case ScalarType::Half: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - auto output_ = checked_dense_tensor_unwrap(output, "output", 6, "_thnn_rrelu_with_noise_forward_out", false, DeviceType::CUDA, dispatch_scalar_type); - THNN_CudaHalfRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - default: - AT_ERROR("_thnn_rrelu_with_noise_forward_out not supported on CUDAType for ", dispatch_scalar_type); - } - return output; -} -Tensor _thnn_rrelu_with_noise_forward(const Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator) { - const OptionalDeviceGuard device_guard(device_of(self)); - auto dispatch_scalar_type = infer_scalar_type(self); - auto output_ = c10::make_intrusive(c10::Storage(c10::Storage::use_byte_size_t(), 0, allocator(), true),DispatchKey::CUDA, scalarTypeToTypeMeta(dispatch_scalar_type)).release(); - auto output = Tensor(c10::intrusive_ptr::reclaim(output_)); - switch (dispatch_scalar_type) { - case ScalarType::Double: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaDoubleRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - case ScalarType::Float: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - case ScalarType::Half: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaHalfRReLU_updateOutput(globalContext().getTHCState(), self_, output_, noise_, lower_, upper_, training, false, generator); - break; - } - default: - AT_ERROR("_thnn_rrelu_with_noise_forward not supported on CUDAType for ", dispatch_scalar_type); - } - return output; -} -Tensor & _thnn_rrelu_with_noise_forward_(Tensor & self, const Tensor & noise, const Scalar& lower, const Scalar& upper, bool training, c10::optional generator) { - const OptionalDeviceGuard device_guard(device_of(self)); - auto dispatch_scalar_type = infer_scalar_type(self); - - switch (dispatch_scalar_type) { - case ScalarType::Double: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaDoubleRReLU_updateOutput(globalContext().getTHCState(), self_, self_, noise_, lower_, upper_, training, true, generator); - break; - } - case ScalarType::Float: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaRReLU_updateOutput(globalContext().getTHCState(), self_, self_, noise_, lower_, upper_, training, true, generator); - break; - } - case ScalarType::Half: { - auto self_ = checked_dense_tensor_unwrap(self, "self", 1, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto noise_ = checked_dense_tensor_unwrap(noise, "noise", 2, "_thnn_rrelu_with_noise_forward_", false, DeviceType::CUDA, dispatch_scalar_type); - auto lower_ = lower.toDouble(); - auto upper_ = upper.toDouble(); - THNN_CudaHalfRReLU_updateOutput(globalContext().getTHCState(), self_, self_, noise_, lower_, upper_, training, true, generator); - break; - } - default: - AT_ERROR("_thnn_rrelu_with_noise_forward_ not supported on CUDAType for ", dispatch_scalar_type); - } - return self; -} std::tuple _thnn_conv2d_forward_out(const Tensor & self, const Tensor & weight, IntArrayRef kernel_size, const c10::optional& bias_opt, IntArrayRef stride, IntArrayRef padding, Tensor & output, Tensor & columns, Tensor & ones) { // See [Note: hacky wrapper removal for optional tensor] c10::MaybeOwned bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt); diff --git a/aten/src/ATen/native/cuda/Activation.cu b/aten/src/ATen/native/cuda/Activation.cu index d400c336c0329..072cafc5a5128 100644 --- a/aten/src/ATen/native/cuda/Activation.cu +++ b/aten/src/ATen/native/cuda/Activation.cu @@ -8,13 +8,17 @@ #include #include +#include #include #include +#include #include #include #include #include +#include #include +#include #include namespace at { @@ -243,6 +247,179 @@ std::tuple prelu_backward_cuda(const Tensor& grad_out_, const Te return std::tuple{input_grad, weight_grad}; } +// ----------------------------------- +// rrelu +// ----------------------------------- +template +#if __CUDA_ARCH__ >= 350 || defined __HIP_PLATFORM_HCC__ +C10_LAUNCH_BOUNDS_2(256, 4) +#endif +__global__ void rrelu_with_noise_cuda_kernel( + int numel, + PhiloxCudaState philox_args, + scalar_t* output, + scalar_t* input, + scalar_t* noise, + double lower, + double upper, + const F& random_func) { + auto seeds = at::cuda::philox::unpack(philox_args); + int idx = blockIdx.x * blockDim.x + threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(std::get<0>(seeds), + idx, + std::get<1>(seeds), + &state); + + int grid_stride = blockDim.x * gridDim.x * unroll_factor; + int rounded_size = ((numel - 1) / grid_stride + 1) * grid_stride; + double range = upper - lower; + + for (int linear_index = idx; linear_index < rounded_size; linear_index += grid_stride) { + auto rand = random_func(&state); + + // ensure that (&rand.x)[ii] is safe + static_assert(sizeof(rand)/sizeof(rand.x) == unroll_factor, ""); + + #pragma unroll + for (int ii = 0; ii < unroll_factor; ii++) { + int li = linear_index + blockDim.x * gridDim.x * ii; + if (li >= numel) { + continue; + } + scalar_t r = static_cast((&rand.x)[ii]); + r = r * range + lower; + if (input[li] <= 0) { + output[li] = input[li] * r; + noise[li] = r; + } else { + output[li] = input[li]; + noise[li] = static_cast(0); + } + } + __syncthreads(); + } +} + +template +inline void _rrelu_with_noise_cuda_train( + Tensor& output, + const Tensor& input_, + const Tensor& noise_, + const Scalar& lower_, + const Scalar& upper_, + c10::optional generator) { + auto input = input_.contiguous(); + auto noise = noise_.contiguous(); + Tensor tmp_output = output.contiguous(); + + int64_t numel = input.numel(); + auto execution_policy = calc_execution_policy(numel); + + auto counter_offset = std::get<0>(execution_policy); + auto grid = std::get<1>(execution_policy); + auto block = std::get<2>(execution_policy); + + auto gen = get_generator_or_default( + generator, cuda::detail::getDefaultCUDAGenerator()); + PhiloxCudaState rng_engine_inputs; + { + // See Note [Acquire lock when using random generators] + std::lock_guard lock(gen->mutex_); + rng_engine_inputs = gen->philox_cuda_state(counter_offset); + } + + scalar_t* input_data = input.data_ptr(); + scalar_t* noise_data = noise.data_ptr(); + scalar_t* output_data = tmp_output.data_ptr(); + + double lower = lower_.to(); + double upper = upper_.to(); + + auto stream = at::cuda::getCurrentCUDAStream(); + + if (std::is_same::value) { + rrelu_with_noise_cuda_kernel<<>>( + numel, + rng_engine_inputs, + output_data, + input_data, + noise_data, + lower, + upper, + [] __device__ (curandStatePhilox4_32_10_t* state) { + return curand_uniform2_double(state); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + // half and float + rrelu_with_noise_cuda_kernel<<>>( + numel, + rng_engine_inputs, + output_data, + input_data, + noise_data, + lower, upper, + [] __device__ (curandStatePhilox4_32_10_t* state) { + return curand_uniform4(state); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + if (!output.is_contiguous()) { + output.copy_(tmp_output); + } +} + +Tensor& rrelu_with_noise_out_cuda(const Tensor& self, + const Tensor& noise, + const Scalar& lower, + const Scalar& upper, + bool training, + c10::optional generator, + Tensor& output) { + TensorArg self_arg{self, "self", 1}, noise_arg{noise, "noise", 2}, + output_arg{output, "output", 3}; + checkAllSameGPU("rrelu_with_noise_out_cuda", {self_arg, noise_arg, output_arg}); + + if (training) { + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + self.scalar_type(), "rrelu_with_noise_out_cuda", [&] { + _rrelu_with_noise_cuda_train( + output, self, noise, lower, upper, generator); + }); + } + else { + auto lower_tensor = lower.to(); + auto upper_tensor = upper.to(); + Scalar negative_slope = (lower_tensor + upper_tensor) / 2; + at::leaky_relu_out(output, self, negative_slope); + } + return output; +} + +Tensor rrelu_with_noise_cuda( + const Tensor& self, + const Tensor& noise, + const Scalar& lower, + const Scalar& upper, + bool training, + c10::optional generator) { + Tensor output = at::empty_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT); + return at::native::rrelu_with_noise_out_cuda(self, noise, lower, upper, training, generator, output); +} + +Tensor& rrelu_with_noise_cuda_( + Tensor& self, + const Tensor& noise, + const Scalar& lower, + const Scalar& upper, + bool training, + c10::optional generator) { + return at::native::rrelu_with_noise_out_cuda( + self, noise, lower, upper, training, generator, self); +} + // ----------------------------------- // hardshrink // ----------------------------------- diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 45f4ed2261f99..953fdbf9b1433 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -8387,13 +8387,13 @@ python_module: nn dispatch: CPU: rrelu_with_noise_out_cpu - CUDA: legacy::cuda::_thnn_rrelu_with_noise_forward_out + CUDA: rrelu_with_noise_out_cuda - func: rrelu_with_noise(Tensor self, Tensor noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor python_module: nn dispatch: CPU: rrelu_with_noise_cpu - CUDA: legacy::cuda::_thnn_rrelu_with_noise_forward + CUDA: rrelu_with_noise_cuda - func: rrelu_with_noise_backward(Tensor grad_output, Tensor self, Tensor noise, Scalar lower, Scalar upper, bool training, bool self_is_result) -> Tensor python_module: nn @@ -8404,7 +8404,7 @@ python_module: nn dispatch: CPU: rrelu_with_noise_cpu_ - CUDA: legacy::cuda::_thnn_rrelu_with_noise_forward_ + CUDA: rrelu_with_noise_cuda_ - func: softplus.out(Tensor self, Scalar beta=1, Scalar threshold=20, *, Tensor(a!) out) -> Tensor(a!) structured: True diff --git a/aten/src/THCUNN/CMakeLists.txt b/aten/src/THCUNN/CMakeLists.txt index 565ebc430b83a..b530096753cf3 100644 --- a/aten/src/THCUNN/CMakeLists.txt +++ b/aten/src/THCUNN/CMakeLists.txt @@ -4,7 +4,6 @@ ${CMAKE_CURRENT_SOURCE_DIR}/GatedLinearUnit.cu ${CMAKE_CURRENT_SOURCE_DIR}/LogSigmoid.cu ${CMAKE_CURRENT_SOURCE_DIR}/MultiLabelMarginCriterion.cu ${CMAKE_CURRENT_SOURCE_DIR}/MultiMarginCriterion.cu -${CMAKE_CURRENT_SOURCE_DIR}/RReLU.cu ${CMAKE_CURRENT_SOURCE_DIR}/SpatialClassNLLCriterion.cu ${CMAKE_CURRENT_SOURCE_DIR}/SpatialConvolutionMM.cu ${CMAKE_CURRENT_SOURCE_DIR}/SpatialDepthwiseConvolution.cu diff --git a/aten/src/THCUNN/RReLU.cu b/aten/src/THCUNN/RReLU.cu deleted file mode 100644 index 1de3e289db699..0000000000000 --- a/aten/src/THCUNN/RReLU.cu +++ /dev/null @@ -1,107 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// copied from cutorch/lib/THC/THCTensorRandom.cu -#define MAX_NUM_BLOCKS 64 -#define BLOCK_SIZE 256 -#define NUM_BLOCKS(n) \ - (std::min((int)THCCeilDiv(n, (ptrdiff_t)BLOCK_SIZE), MAX_NUM_BLOCKS)) - -template -inline T __device__ curand_uniform_type(curandStatePhilox4_32_10_t *state); - -template <> -inline THHalf __device__ curand_uniform_type(curandStatePhilox4_32_10_t *state) { - auto rand = curand_uniform4(state); - return ScalarConvert::to(rand.x); -} - -template <> -inline float __device__ curand_uniform_type(curandStatePhilox4_32_10_t *state) { - auto rand = curand_uniform4(state); - return rand.x; -} - -template <> -inline double __device__ curand_uniform_type(curandStatePhilox4_32_10_t *state) { - auto rand = curand_uniform2_double(state); - return rand.x; -} - -template -__global__ void rreluUpdateOutputTrain(int n, at::PhiloxCudaState philox_args, - T *input, T* noise, T *output, double a, double b) -{ - auto seeds = at::cuda::philox::unpack(philox_args); - int idx = blockIdx.x * blockDim.x + threadIdx.x; - curandStatePhilox4_32_10_t state; - curand_init(std::get<0>(seeds), - idx, - std::get<1>(seeds), - &state); - - CUDA_KERNEL_LOOP(i, n) - { - if (input[i] <= 0) - { - T r = curand_uniform_type(&state); - r = ScalarConvert::to(r * (b-a) + a); - output[i] = input[i] * r; - noise[i] = r; - } - else - { - output[i] = input[i]; - noise[i] = ScalarConvert::to(1); - } - } -} - -template -struct RReLUUpdateOutputEval_functor -{ - const T negSlope_; - - RReLUUpdateOutputEval_functor(T negSlope) - : negSlope_(negSlope) - {} - - __device__ __forceinline__ void operator()(T *out, T *in) - { - const T x = *in; - const T r = x <= 0 ? negSlope_ : ScalarConvert::to(1); - *out = x * r; - } -}; - -template -struct RReLUUpdateOutputEvalIP_functor -{ - const T negSlope_; - - RReLUUpdateOutputEvalIP_functor(T negSlope) - : negSlope_(negSlope) - {} - - __device__ __forceinline__ void operator()(T *x) - { - if (*x <= 0) - { - *x = *x * negSlope_; - } - } -}; - -#include -#include diff --git a/aten/src/THCUNN/generic/RReLU.cu b/aten/src/THCUNN/generic/RReLU.cu deleted file mode 100644 index 7cd1876ba2c1d..0000000000000 --- a/aten/src/THCUNN/generic/RReLU.cu +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef THC_GENERIC_FILE -#define THC_GENERIC_FILE "THCUNN/generic/RReLU.cu" -#else - -#include -#include - -void THNN_(RReLU_updateOutput)( - THCState *state, - THCTensor *input, - THCTensor *output, - THCTensor *noise, - double lower, - double upper, - bool train, - bool inplace, - c10::optional generator) -{ - THCUNN_assertSameGPU(state, 3, input, output, noise); - auto gen = at::get_generator_or_default(generator, at::cuda::detail::getDefaultCUDAGenerator()); - if (train) - { - auto inputTensor = THTensor_wrap(input).contiguous(); - input = inputTensor.unsafeGetTensorImpl(); - THCTensor_(resizeAs)(state, noise, input); - scalar_t *input_data = THCTensor_(data)(state, input); - scalar_t *noise_data = THCTensor_(data)(state, noise); - ptrdiff_t n = THCTensor_(nElement)(state, input); - - // philox offset calculation for grid-stride loop utilizing curand4 - const uint32_t curand4_engine_calls = 4; - dim3 grid = NUM_BLOCKS(n); - uint64_t counter_offset = ((n - 1) / (BLOCK_SIZE * grid.x) + 1) * curand4_engine_calls; - at::PhiloxCudaState rng_engine_inputs; - { - // See Note [Acquire lock when using random generators] - std::lock_guard lock(gen->mutex_); - rng_engine_inputs = gen->philox_cuda_state(counter_offset); - } - if (inplace) - { - rreluUpdateOutputTrain<<>>( - n, rng_engine_inputs, input_data, noise_data, input_data, lower, upper); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - THCTensor_(set)(state, output, input); - } - else - { - THCTensor_(resizeAs)(state, output, input); - scalar_t *output_data = THCTensor_(data)(state, output); - rreluUpdateOutputTrain<<>>( - n, rng_engine_inputs, input_data, noise_data, output_data, lower, upper); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - } - } - else - { - const scalar_t negSlope = ScalarConvert::to((lower + upper) / 2); - if (inplace) - { - THC_pointwiseApply1(state, input, RReLUUpdateOutputEvalIP_functor(negSlope)); - THCTensor_(set)(state, output, input); - } - else - { - THCTensor_(resizeAs)(state, output, input); - THC_pointwiseApply2(state, output, input, RReLUUpdateOutputEval_functor(negSlope)); - } - } -} -#endif diff --git a/aten/src/THCUNN/generic/THCUNN.h b/aten/src/THCUNN/generic/THCUNN.h index ff8ebf618cdd2..3ef791ae41e65 100644 --- a/aten/src/THCUNN/generic/THCUNN.h +++ b/aten/src/THCUNN/generic/THCUNN.h @@ -200,15 +200,4 @@ TORCH_CUDA_CU_API void THNN_(SpatialDepthwiseConvolution_accGradParameters)( int padH, int dilationW, int dilationH); - -TORCH_CUDA_CU_API void THNN_(RReLU_updateOutput)( - THCState* state, - THCTensor* input, - THCTensor* output, - THCTensor* noise, - double lower, - double upper, - bool train, - bool inplace, - c10::optional generator); #endif From e5c99d990880c259fd83a740e7e576f74ca99602 Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Thu, 17 Jun 2021 00:44:22 -0700 Subject: [PATCH 183/305] Revert D29147009: [pytorch][PR] refine disabled test Test Plan: revert-hammer Differential Revision: D29147009 (https://github.com/pytorch/pytorch/commit/5fd6ead0970a539f0bd1b24efbe6fc73bd70603f) Original commit changeset: 37e01ac6e8d6 fbshipit-source-id: e9cd819fd819e3d653deda3b7a981c39ec0452f4 --- test/distributed/test_c10d_gloo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/distributed/test_c10d_gloo.py b/test/distributed/test_c10d_gloo.py index 22de4440cd755..32f049f084c64 100644 --- a/test/distributed/test_c10d_gloo.py +++ b/test/distributed/test_c10d_gloo.py @@ -640,7 +640,7 @@ def _test_sparse_allreduce_basics(self, fn): self.assertEqual(tensors, outputs) self.assertEqual(result, outputs) - @skip_if_win32() + @unittest.skip("intermittent failures on Windows, in CI") def test_sparse_allreduce_basics(self): self._test_sparse_allreduce_basics(lambda t: t) From f233274f3072ea52d506531d70a3a065b1c5069c Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Thu, 17 Jun 2021 00:47:42 -0700 Subject: [PATCH 184/305] Revert D28875276: Move RPC agents to libtorch Test Plan: revert-hammer Differential Revision: D28875276 (https://github.com/pytorch/pytorch/commit/fc50f91929b67f1ccf34890416b46a17b399e21b) Original commit changeset: f2f6970fd74d fbshipit-source-id: 3c52af652579733ebea8ddfb06576a0ce262bf78 --- BUILD.bazel | 2 +- caffe2/CMakeLists.txt | 57 ++++++++- cmake/Dependencies.cmake | 7 -- test/cpp/rpc/CMakeLists.txt | 4 +- third_party/tensorpipe | 2 +- third_party/tensorpipe.BUILD | 112 ++++++++---------- tools/build_variables.bzl | 12 +- torch/CMakeLists.txt | 6 + torch/csrc/distributed/rpc/macros.h | 5 + .../csrc/distributed/rpc/tensorpipe_agent.cpp | 1 + torch/csrc/distributed/rpc/tensorpipe_agent.h | 1 + .../csrc/distributed/rpc/tensorpipe_cuda.cpp | 3 +- .../csrc/distributed/rpc/tensorpipe_utils.cpp | 1 + torch/csrc/distributed/rpc/tensorpipe_utils.h | 1 + 14 files changed, 128 insertions(+), 86 deletions(-) create mode 100644 torch/csrc/distributed/rpc/macros.h diff --git a/BUILD.bazel b/BUILD.bazel index 45d71a8d4626e..b7e16ac1c915c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1728,7 +1728,7 @@ cc_library( ], [ ":aten", - "@tensorpipe//:tensorpipe_cpu", + "@tensorpipe", ], ), alwayslink = True, diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 88cffd1a75d1c..174018456efd8 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -344,6 +344,53 @@ endif() if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) + if(USE_DISTRIBUTED) + + # Define this target even if we're building without TensorPipe, to make life + # easier to other targets that depend on this. However, in that case, by not + # setting the USE_TENSORPIPE compile definition, this target will just end + # up being empty. Downstream targets should also add a #ifdef guard. + if(NOT WIN32) + add_library(process_group_agent + "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.h" + ) + target_link_libraries(process_group_agent PRIVATE torch fmt::fmt-header-only) + add_dependencies(process_group_agent torch) + + if(USE_TENSORPIPE) + add_library(tensorpipe_agent + "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/macros.h" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.h" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_cuda.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h" + ) + target_link_libraries(tensorpipe_agent PRIVATE torch tensorpipe fmt::fmt-header-only) + add_dependencies(tensorpipe_agent torch) + if(USE_CUDA) + target_compile_definitions(tensorpipe_agent PUBLIC USE_CUDA) + endif() + + if(USE_ROCM) + target_compile_definitions(tensorpipe_agent PRIVATE + USE_ROCM + __HIP_PLATFORM_HCC__ + ) + endif() + + target_compile_definitions(tensorpipe_agent PUBLIC USE_TENSORPIPE) + target_link_libraries(tensorpipe_agent PRIVATE tensorpipe) + add_dependencies(tensorpipe_agent tensorpipe) + endif() + endif() + endif() + set(CMAKE_POSITION_INDEPENDENT_CODE TRUE) # Generate files @@ -1189,7 +1236,7 @@ endif() if(USE_DISTRIBUTED) # Needed to support the inclusion of c10d/Foo.hpp headers. target_include_directories(torch_cpu PUBLIC ${TORCH_SRC_DIR}/lib) - target_compile_definitions(torch_cpu PUBLIC USE_DISTRIBUTED) + target_compile_definitions(torch_cpu PRIVATE USE_DISTRIBUTED) if(USE_GLOO AND USE_C10D_GLOO) target_compile_definitions(torch_cpu PUBLIC USE_C10D_GLOO) endif() @@ -1216,12 +1263,16 @@ if(USE_DISTRIBUTED) # #if defined(USE_DISTRIBUTED) && !defined(_WIN32) # need to be removed when RPC is supported if(NOT WIN32) - target_compile_definitions(torch_cpu PUBLIC USE_RPC) + target_compile_definitions(torch_cpu PRIVATE + USE_RPC + ) endif() # Pass USE_TENSORPIPE to torch_cpu as some parts of rpc/utils.cpp # can only be compiled with USE_TENSORPIPE is set. if(USE_TENSORPIPE) - target_compile_definitions(torch_cpu PUBLIC USE_TENSORPIPE) + target_compile_definitions(torch_cpu PRIVATE + USE_TENSORPIPE + ) endif() endif() diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 70b6d71face6b..ab4cd32c40bce 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1377,13 +1377,6 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE) add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/tensorpipe) list(APPEND Caffe2_DEPENDENCY_LIBS tensorpipe) - if(USE_CUDA) - list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS tensorpipe_cuda) - elseif(USE_ROCM) - message(WARNING "TensorPipe doesn't yet support ROCm") - # Not yet... - # list(APPEND Caffe2_HIP_DEPENDENCY_LIBS tensorpipe_hip) - endif() endif() endif() diff --git a/test/cpp/rpc/CMakeLists.txt b/test/cpp/rpc/CMakeLists.txt index c9fb1b0e7f17a..0eff382d2b1b8 100644 --- a/test/cpp/rpc/CMakeLists.txt +++ b/test/cpp/rpc/CMakeLists.txt @@ -5,7 +5,7 @@ set(TORCH_RPC_TEST_SOURCES ${TORCH_RPC_TEST_DIR}/test_wire_serialization.cpp ) set(TORCH_RPC_TEST_DEPENDENCY_LIBS - torch gtest + torch gtest process_group_agent ) if(USE_GLOO) @@ -20,7 +20,7 @@ if(USE_TENSORPIPE) ${TORCH_RPC_TEST_DIR}/test_tensorpipe_serialization.cpp ) list(APPEND TORCH_RPC_TEST_DEPENDENCY_LIBS - tensorpipe + tensorpipe_agent tensorpipe ) endif() diff --git a/third_party/tensorpipe b/third_party/tensorpipe index c0e7623adb05f..42a67277c1882 160000 --- a/third_party/tensorpipe +++ b/third_party/tensorpipe @@ -1 +1 @@ -Subproject commit c0e7623adb05f36311c7cde6dac8fc4c290419d9 +Subproject commit 42a67277c1882c90cec0da6e57afb20247424994 diff --git a/third_party/tensorpipe.BUILD b/third_party/tensorpipe.BUILD index ae210f473933d..d9e4bdb395741 100644 --- a/third_party/tensorpipe.BUILD +++ b/third_party/tensorpipe.BUILD @@ -71,82 +71,63 @@ cc_library( ) header_template_rule( - name = "tensorpipe_cpu_config_header", + name = "tensorpipe_config_header", src = "tensorpipe/config.h.in", out = "tensorpipe/config.h", substitutions = { - "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "#define TENSORPIPE_HAS_SHM_TRANSPORT 1", - "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "#define TENSORPIPE_HAS_IBV_TRANSPORT 1", - "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "#define TENSORPIPE_HAS_CMA_CHANNEL 1", + "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "", + "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "", + "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "", + "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "", + "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "", + "#cmakedefine01 TENSORPIPE_SUPPORTS_CUDA": "", }, ) -header_template_rule( - name = "tensorpipe_cuda_config_header", - src = "tensorpipe/config_cuda.h.in", - out = "tensorpipe/config_cuda.h", - substitutions = { - "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "#define TENSORPIPE_HAS_CUDA_IPC_CHANNEL 1", - "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "#define TENSORPIPE_HAS_CUDA_GDR_CHANNEL 1", - }, -) +TENSORPIPE_HEADERS = glob([ + "tensorpipe/*.h", + "tensorpipe/channel/*.h", + "tensorpipe/channel/*/*.h", + "tensorpipe/common/*.h", + "tensorpipe/core/*.h", + "tensorpipe/transport/*.h", + "tensorpipe/transport/*/*.h", + "tensorpipe/util/*/*.h", +]) -# We explicitly list the CUDA headers & sources, and we consider everything else -# as CPU (using a catch-all glob). This is both because there's fewer CUDA files -# (thus making it easier to list them exhaustively) and because it will make it -# more likely to catch a misclassified file: if we forget to mark a file as CUDA -# we'll try to build it on CPU and that's likely to fail. +TENSORPIPE_BASE_SRCS = glob([ + "tensorpipe/*.cc", + "tensorpipe/channel/*.cc", + "tensorpipe/common/address.cc", + "tensorpipe/common/epoll_loop.cc", + "tensorpipe/common/error.cc", + "tensorpipe/common/fd.cc", + "tensorpipe/common/ibv.cc", + "tensorpipe/common/socket.cc", + "tensorpipe/common/system.cc", + "tensorpipe/core/*.cc", + "tensorpipe/transport/*.cc", + "tensorpipe/util/*/*.cc", +]) -TENSORPIPE_CUDA_HEADERS = [ - "tensorpipe/tensorpipe_cuda.h", - "tensorpipe/channel/cuda_basic/*.h", - "tensorpipe/channel/cuda_gdr/*.h", - "tensorpipe/channel/cuda_ipc/*.h", - "tensorpipe/channel/cuda_xth/*.h", - "tensorpipe/common/cuda.h", - "tensorpipe/common/cuda_buffer.h", - "tensorpipe/common/cuda_lib.h", - "tensorpipe/common/cuda_loop.h", - "tensorpipe/common/nvml_lib.h", -] +TENSORPIPE_SRCS = TENSORPIPE_BASE_SRCS + glob([ + "tensorpipe/channel/basic/*.cc", + "tensorpipe/channel/mpt/*.cc", + "tensorpipe/channel/xth/*.cc", + "tensorpipe/transport/uv/*.cc", +]) -TENSORPIPE_CUDA_SOURCES = [ +TENSORPIPE_SRCS_CUDA = TENSORPIPE_SRCS + glob([ + "tensorpipe/common/cuda_loop.cc", "tensorpipe/channel/cuda_basic/*.cc", - "tensorpipe/channel/cuda_gdr/*.cc", "tensorpipe/channel/cuda_ipc/*.cc", "tensorpipe/channel/cuda_xth/*.cc", - "tensorpipe/common/cuda_buffer.cc", - "tensorpipe/common/cuda_loop.cc", -] - -TENSORPIPE_CPU_HEADERS = glob( - [ - "tensorpipe/*.h", - "tensorpipe/channel/*.h", - "tensorpipe/channel/*/*.h", - "tensorpipe/common/*.h", - "tensorpipe/core/*.h", - "tensorpipe/transport/*.h", - "tensorpipe/transport/*/*.h", - ], - exclude=TENSORPIPE_CUDA_HEADERS) - -TENSORPIPE_CPU_SOURCES = glob( - [ - "tensorpipe/*.cc", - "tensorpipe/channel/*.cc", - "tensorpipe/channel/*/*.cc", - "tensorpipe/common/*.cc", - "tensorpipe/core/*.cc", - "tensorpipe/transport/*.cc", - "tensorpipe/transport/*/*.cc", - ], - exclude=TENSORPIPE_CUDA_SOURCES) +]) cc_library( - name = "tensorpipe_cpu", - srcs = TENSORPIPE_CPU_SOURCES, - hdrs = TENSORPIPE_CPU_HEADERS + [":tensorpipe_cpu_config_header"], + name = "tensorpipe", + srcs = TENSORPIPE_SRCS + [":tensorpipe_config_header"], + hdrs = TENSORPIPE_HEADERS, includes = [ ".", ], @@ -162,8 +143,8 @@ cc_library( cc_library( name = "tensorpipe_cuda", - srcs = TENSORPIPE_CUDA_SOURCES, - hdrs = TENSORPIPE_CUDA_HEADERS + [":tensorpipe_cuda_config_header"], + srcs = TENSORPIPE_SRCS_CUDA + [":tensorpipe_config_header"], + hdrs = TENSORPIPE_HEADERS, includes = [ ".", ], @@ -172,7 +153,8 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ - ":tensorpipe_cpu", + ":libnop", + ":libuv", "@cuda", ], ) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 06acafd645eab..c0b206fa48e05 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -356,14 +356,12 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.cpp", - "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/message.cpp", "torch/csrc/distributed/rpc/profiler/remote_profiler_manager.cpp", "torch/csrc/distributed/rpc/profiler/server_process_global_profiler.cpp", "torch/csrc/distributed/rpc/python_call.cpp", "torch/csrc/distributed/rpc/python_remote_call.cpp", "torch/csrc/distributed/rpc/python_resp.cpp", - "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/request_callback.cpp", "torch/csrc/distributed/rpc/request_callback_no_python.cpp", "torch/csrc/distributed/rpc/rpc_agent.cpp", @@ -373,9 +371,6 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/rpc/script_call.cpp", "torch/csrc/distributed/rpc/script_remote_call.cpp", "torch/csrc/distributed/rpc/script_resp.cpp", - "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", - "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", - "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/torchscript_functions.cpp", "torch/csrc/distributed/rpc/types.cpp", "torch/csrc/distributed/rpc/utils.cpp", @@ -531,7 +526,6 @@ libtorch_cuda_distributed_base_sources = [ # These files are only supported on Linux (and others) but not on Windows. libtorch_cuda_distributed_extra_sources = [ - "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", "torch/lib/c10d/NCCLUtils.cpp", "torch/lib/c10d/ProcessGroupNCCL.cpp", ] @@ -720,11 +714,17 @@ libtorch_python_distributed_core_sources = [ libtorch_python_distributed_sources = libtorch_python_distributed_core_sources + [ "torch/csrc/distributed/autograd/init.cpp", + "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/init.cpp", + "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/py_rref.cpp", "torch/csrc/distributed/rpc/python_functions.cpp", "torch/csrc/distributed/rpc/python_rpc_handler.cpp", "torch/csrc/distributed/rpc/request_callback_impl.cpp", + "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", + "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", + "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", + "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/testing/init.cpp", "torch/csrc/distributed/rpc/unpickled_python_call.cpp", "torch/csrc/distributed/rpc/unpickled_python_remote_call.cpp", diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index ce0f16bf5abeb..197926f309838 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -261,9 +261,11 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") endif() if(USE_DISTRIBUTED) + list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_DISTRIBUTED) if(WIN32) append_filelist("libtorch_python_distributed_core_sources" TORCH_PYTHON_SRCS) else() + list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_RPC) append_filelist("libtorch_python_distributed_sources" TORCH_PYTHON_SRCS) endif() # Disable certain warnings for GCC-9.X @@ -272,6 +274,10 @@ if(USE_DISTRIBUTED) set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/rpc/testing/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/c10d/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") endif() + if(USE_TENSORPIPE) + list(APPEND TORCH_PYTHON_LINK_LIBRARIES tensorpipe) + list(APPEND TORCH_PYTHON_PUBLIC_COMPILE_DEFINITIONS USE_TENSORPIPE) + endif() # NCCL is a private dependency of libtorch, but libtorch_python includes # some private headers of libtorch, which in turn include NCCL. As a hacky # alternative to making NCCL a public dependency of libtorch, we make it diff --git a/torch/csrc/distributed/rpc/macros.h b/torch/csrc/distributed/rpc/macros.h new file mode 100644 index 0000000000000..2763dd0207bef --- /dev/null +++ b/torch/csrc/distributed/rpc/macros.h @@ -0,0 +1,5 @@ +#pragma once + +#if defined(USE_CUDA) && !defined(__HIP_PLATFORM_HCC__) +#define USE_CUDA_NOT_ROCM +#endif diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp index 74c279425658b..0f6645cdcd5d5 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index 4450792a0f06d..9462c396b0f3b 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -9,6 +9,7 @@ #include #include #include +#include #include // Forward-declare the TensorPipe classes we need, to avoid including its diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp index 03ec63d8ddc88..9489fcd222bbd 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp @@ -1,7 +1,8 @@ +#include #include #include -#if defined(USE_TENSORPIPE) && !defined(__HIP_PLATFORM_HCC__) +#if defined(USE_TENSORPIPE) && defined(USE_CUDA_NOT_ROCM) #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp index 32f3a132f8f50..55b8554f66d28 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp @@ -1,3 +1,4 @@ +#include #include #ifdef USE_TENSORPIPE diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index bf5d87cacc4b5..ab328b9dca1a1 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -2,6 +2,7 @@ #ifdef USE_TENSORPIPE +#include #include namespace tensorpipe { From 3c3bb9110365a953a277d553b12eb64f9a363a5f Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Thu, 17 Jun 2021 00:47:42 -0700 Subject: [PATCH 185/305] Revert D29132956: Add some TORCH_API annotations to RPC Test Plan: revert-hammer Differential Revision: D29132956 (https://github.com/pytorch/pytorch/commit/04ec122868d436e32d81345343725e271d94cdfc) Original commit changeset: 8637640d56a1 fbshipit-source-id: f497adcbfd5a6b5a46b8689b1943ae2687ea737b --- torch/csrc/distributed/rpc/process_group_agent.h | 6 +++--- torch/csrc/distributed/rpc/tensorpipe_agent.h | 12 ++++++------ torch/csrc/distributed/rpc/tensorpipe_utils.h | 6 +++--- .../rpc/testing/faulty_process_group_agent.h | 4 ++-- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/torch/csrc/distributed/rpc/process_group_agent.h b/torch/csrc/distributed/rpc/process_group_agent.h index a6d1115f4074c..5706870988140 100644 --- a/torch/csrc/distributed/rpc/process_group_agent.h +++ b/torch/csrc/distributed/rpc/process_group_agent.h @@ -15,7 +15,7 @@ namespace rpc { constexpr auto kDefaultNumSendRecvThreads = 4; -struct TORCH_API ProcessGroupRpcBackendOptions : public RpcBackendOptions { +struct ProcessGroupRpcBackendOptions : public RpcBackendOptions { ProcessGroupRpcBackendOptions( int num_send_recv_threads, float rpc_timeout, @@ -34,7 +34,7 @@ struct TORCH_API ProcessGroupRpcBackendOptions : public RpcBackendOptions { // SendWork and RecvWork will be put into a task queue, and later picked up by // worker threads from the same ThreadPool. -struct TORCH_API SendWork { +struct SendWork { SendWork(const WorkerInfo& to, c10::intrusive_ptr message) : to_(to), message_(std::move(message)) {} @@ -44,7 +44,7 @@ struct TORCH_API SendWork { // SendWork wraps a Message and RecvWork wraps a Tensor. The difference here is // to allow us to run serialization/deserialization in the worker threads. -struct TORCH_API RecvWork { +struct RecvWork { RecvWork( const WorkerInfo& from, MessageType type, diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index 9462c396b0f3b..df3328793fa11 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -62,7 +62,7 @@ constexpr int64_t kCudaBasicChannelPriority = 0; using steady_clock_time_point = std::chrono::time_point; -struct TORCH_API TransportRegistration { +struct TransportRegistration { std::shared_ptr transport; int64_t priority; std::string address; @@ -71,7 +71,7 @@ struct TORCH_API TransportRegistration { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) C10_DECLARE_REGISTRY(TensorPipeTransportRegistry, TransportRegistration); -struct TORCH_API ChannelRegistration { +struct ChannelRegistration { std::shared_ptr channel; int64_t priority; }; @@ -81,7 +81,7 @@ C10_DECLARE_REGISTRY(TensorPipeChannelRegistry, ChannelRegistration); constexpr auto kDefaultNumWorkerThreads = 16; -struct TORCH_API TensorPipeRpcBackendOptions : public RpcBackendOptions { +struct TensorPipeRpcBackendOptions : public RpcBackendOptions { TensorPipeRpcBackendOptions( int numWorkerThreads, optional> transports, @@ -146,13 +146,13 @@ struct TORCH_API TensorPipeRpcBackendOptions : public RpcBackendOptions { }; // Struct to track the network source metrics -struct TORCH_API NetworkSourceInfo { +struct NetworkSourceInfo { worker_id_t srcRank; std::vector srcMachineAddr; }; // Struct to track aggregated network metrics -struct TORCH_API AggregatedNetworkData { +struct AggregatedNetworkData { uint64_t numCalls{0}; uint64_t totalSentBytes{0}; uint64_t totalRecvBytes{0}; @@ -163,7 +163,7 @@ struct TORCH_API AggregatedNetworkData { // to transparently move tensors and payloads through the fastest available // transport or channel. It acts like a hybrid RPC transport, providing shared // memory (linux) and TCP (linux & mac) support. CUDA support is in progress. -class TORCH_API TensorPipeAgent : public RpcAgent { +class TensorPipeAgent : public RpcAgent { public: TensorPipeAgent( const c10::intrusive_ptr<::c10d::Store>& store, diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index ab328b9dca1a1..3f41b351c9898 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -15,7 +15,7 @@ namespace torch { namespace distributed { namespace rpc { -TORCH_API const c10::Stream& getStreamForDevice( +const c10::Stream& getStreamForDevice( const std::vector& streams, const c10::Device& device); @@ -44,12 +44,12 @@ class TensorpipeDeviceTypeConverter { virtual ~TensorpipeDeviceTypeConverter() = default; }; -extern TORCH_API std::array< +extern C10_API std::array< std::atomic, static_cast(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> device_type_converter_registry; -class TORCH_API TensorpipeDeviceTypeConverterRegistrar { +class C10_API TensorpipeDeviceTypeConverterRegistrar { public: TensorpipeDeviceTypeConverterRegistrar( DeviceType, diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index d0bbb33fe3df2..ee589072f2ddd 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -7,7 +7,7 @@ namespace torch { namespace distributed { namespace rpc { -struct TORCH_API FaultyProcessGroupRpcBackendOptions +struct FaultyProcessGroupRpcBackendOptions : public ProcessGroupRpcBackendOptions { FaultyProcessGroupRpcBackendOptions( int num_send_recv_threads, @@ -31,7 +31,7 @@ struct TORCH_API FaultyProcessGroupRpcBackendOptions int numFailSends; }; -class TORCH_API FaultyProcessGroupAgent : public ProcessGroupAgent { +class FaultyProcessGroupAgent : public ProcessGroupAgent { public: FaultyProcessGroupAgent( const c10::intrusive_ptr<::c10d::Store>& store, From 6b1712019ad6f60f42f0c8cfb85d1de3a9125f2c Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Thu, 17 Jun 2021 00:47:42 -0700 Subject: [PATCH 186/305] Revert D29132955: Pass RequestCallback to FaultyPG RPC agent Test Plan: revert-hammer Differential Revision: D29132955 (https://github.com/pytorch/pytorch/commit/cbbb7e145e770213146bab1fc36c6abe1865d857) Original commit changeset: bb7554b84bcb fbshipit-source-id: 4dfa2fbe7b8f58c951991c79aa9e2aa819793013 --- .../distributed/rpc/testing/faulty_process_group_agent.cpp | 4 ++-- .../csrc/distributed/rpc/testing/faulty_process_group_agent.h | 1 - torch/csrc/distributed/rpc/testing/init.cpp | 2 -- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp index bb980ee8cef08..f51de9d870971 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp @@ -1,3 +1,4 @@ +#include #include #include @@ -15,7 +16,6 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( c10::intrusive_ptr<::c10d::ProcessGroup> pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, - std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends) @@ -25,7 +25,7 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( std::move(pg), numSendRecvThreads, rpcTimeout, - std::move(cb)), + std::make_unique()), failNumSends_(failNumSends), messageTypesToFail_(parseMessagesToFailInput(messagesToFail)), messageTypesToDelay_(parseMessagesToDelay(messageTypesToDelay)) {} diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index ee589072f2ddd..b80bd78c3e1de 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -39,7 +39,6 @@ class FaultyProcessGroupAgent : public ProcessGroupAgent { c10::intrusive_ptr pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, - std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends = 0); diff --git a/torch/csrc/distributed/rpc/testing/init.cpp b/torch/csrc/distributed/rpc/testing/init.cpp index bccaa1f2b4232..28344bb5b1978 100644 --- a/torch/csrc/distributed/rpc/testing/init.cpp +++ b/torch/csrc/distributed/rpc/testing/init.cpp @@ -1,7 +1,6 @@ #include #include -#include #include #include #include @@ -83,7 +82,6 @@ PyObject* faulty_agent_init(PyObject* _unused, PyObject* noargs) { process_group, num_send_recv_threads, rpc_timeout, - std::make_unique(), messages_to_fail, messages_to_delay, failNumSends), From eb36f67dcc7000caa58fa4dfa9c089ce17c6d523 Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Thu, 17 Jun 2021 01:22:17 -0700 Subject: [PATCH 187/305] [TensorExpr] Minor cleanup in TensorExprKernel::computeValue (#60041) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60041 Differential Revision: D29146709 D29146709 Test Plan: Imported from OSS Reviewed By: navahgar Pulled By: ZolotukhinM fbshipit-source-id: 49ac919c18f669d7fda1a26c5a74e62ea752df4f --- test/cpp/tensorexpr/test_kernel.cpp | 4 ++-- torch/csrc/jit/tensorexpr/kernel.cpp | 30 +++++++++------------------- 2 files changed, 11 insertions(+), 23 deletions(-) diff --git a/test/cpp/tensorexpr/test_kernel.cpp b/test/cpp/tensorexpr/test_kernel.cpp index d1ca8b86bd826..e3c5ce0a19975 100644 --- a/test/cpp/tensorexpr/test_kernel.cpp +++ b/test/cpp/tensorexpr/test_kernel.cpp @@ -808,13 +808,13 @@ TEST_F(Kernel, SumOneAxis) { TEST_F(Kernel, SumMultipleAxes) { // Test lowering of sum on multiple axes. const auto graph_template = R"IR( - graph(%0 : Float(2, 3, 2, 3, strides=[18, 6, 3, 1], device=cpu)): + graph(%0 : Float(2, 3, 2, 3, strides=[18, 6, 3, 1], requires_grad=0, device=cpu)): %1 : int = prim::Constant[value=${dim1}]() %2 : int = prim::Constant[value=${dim2}]() %3 : int[] = prim::ListConstruct(%1, %2) %4 : bool = prim::Constant[value=${keepdim}]() %5 : ${dtype} - %6 : Float(${size}, strides=[${strides}]) = aten::sum(%0, %3, %4, %5) + %6 : Float(${size}, strides=[${strides}], requires_grad=0, device=cpu) = aten::sum(%0, %3, %4, %5) return (%6))IR"; auto a = iotaTensor({2, 3, 2, 3}, TensorOptions(kCPU).dtype(at::kFloat)); diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 9eb2cc7e27706..1462e82540d29 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -2466,30 +2466,18 @@ Tensor* TensorExprKernel::computeValue(const torch::jit::Value* v) { case aten::sum: case aten::softmax: case aten::log_softmax: - case aten::conv2d: { - std::vector argInputs; - for (auto inp : inputs) { - argInputs.push_back(toArg(inp)); - } - auto outputType = findDtypeForValue(v->node()->output()); - std::vector outputShape = {}; - // shape inference not implemented for sum - if (v->node()->kind() != aten::sum) { - outputShape = sizesForValue(v); - } - return computeOperandValue( - v->node()->kind(), argInputs, outputShape, outputType, device_); - } break; - + case aten::conv2d: case aten::to: { std::vector argInputs; - argInputs.push_back(toArg(inputs[0])); - auto outputType = findDtypeForValue(v->node()->output()); - std::vector outputShape = {}; - // shape inference not implemented for sum - if (v->node()->kind() != aten::sum) { - outputShape = sizesForValue(v); + if (v->node()->kind() != aten::to) { + for (auto inp : inputs) { + argInputs.push_back(toArg(inp)); + } + } else { + argInputs.push_back(toArg(inputs[0])); } + auto outputType = findDtypeForValue(v->node()->output()); + std::vector outputShape = sizesForValue(v); return computeOperandValue( v->node()->kind(), argInputs, outputShape, outputType, device_); } break; From 4caca7a15b6edd9f99bba1fd0a5b8fb64ed991c3 Mon Sep 17 00:00:00 2001 From: Heitor Schueroff Date: Thu, 17 Jun 2021 04:47:33 -0700 Subject: [PATCH 188/305] Improved torch.einsum testing and fixed bug (#59731) Summary: Improved torch.einsum testing and fixed a bug where lower case letters appeared before upper case letters in the sorted order which is inconsistent with NumPy. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59731 Reviewed By: SplitInfinity, ansley Differential Revision: D29183078 Pulled By: heitorschueroff fbshipit-source-id: a33980d273707da2d60a387a2af2fa41527ddb68 --- aten/src/ATen/native/Linear.cpp | 16 +- test/test_linalg.py | 286 +++++++++++++++++++------------- torch/functional.py | 4 +- 3 files changed, 182 insertions(+), 124 deletions(-) diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp index 7f9a110b7a5ce..fb8ece5d63606 100644 --- a/aten/src/ATen/native/Linear.cpp +++ b/aten/src/ATen/native/Linear.cpp @@ -154,7 +154,12 @@ bool einsum_check_label(unsigned char label) { uint8_t einsum_label_to_index(unsigned char label) { constexpr uint8_t NUM_OF_LETTERS = 'z' - 'a' + 1; - return std::islower(label) ? label - 'a' : NUM_OF_LETTERS + (label - 'A'); + return std::isupper(label) ? label - 'A' : NUM_OF_LETTERS + (label - 'a'); +} + +unsigned char einsum_index_to_label(uint8_t index) { + constexpr uint8_t NUM_OF_LETTERS = 'z' - 'a' + 1; + return index < NUM_OF_LETTERS ? index + 'A' : index - NUM_OF_LETTERS + 'a'; } } // namespace @@ -166,7 +171,6 @@ uint8_t einsum_label_to_index(unsigned char label) { // dimensions We do the last part by reducing to bmm. Tensor einsum(c10::string_view equation, TensorList operands) { TORCH_CHECK(!operands.empty(), "einsum(): must provide at least one operand"); - checkDeviceType("einsum():", operands, operands[0].device().type()); // Code used to identify ELLIPSIS ("...") constexpr uint8_t ELLIPSIS = 52; @@ -245,7 +249,7 @@ Tensor einsum(c10::string_view equation, TensorList operands) { const auto operand = operands[i]; const auto labels = op_labels[i]; const auto ndims = operand.dim(); - int64_t nlabels = labels.size(); + int64_t nlabels = static_cast(labels.size()); bool has_ellipsis = false; for (const auto& label : labels) { @@ -364,11 +368,11 @@ Tensor einsum(c10::string_view equation, TensorList operands) { const auto labels = op_labels[i]; const auto original_sizes = operand.sizes(); - std::size_t j = 0; + int64_t j = 0; for (const auto& label : labels) { if (label == ELLIPSIS) { // Add missing dimensions covered by the ellipsis - const int64_t num_missing_dim = + const auto num_missing_dim = ell_num_dim - (original_sizes.size() - labels.size() + 1); for (const auto k : c10::irange(num_missing_dim)) { (void)k; //Suppress unused warning @@ -383,7 +387,7 @@ Tensor einsum(c10::string_view equation, TensorList operands) { TORCH_CHECK( operand.size(j) == operand.size(dim), "einsum(): subscript ", - char(label + 'a'), + einsum_index_to_label(label), " is repeated for operand ", i, " but the sizes don't match, ", diff --git a/test/test_linalg.py b/test/test_linalg.py index 3bb75bc9d67e0..359b8033778a5 100644 --- a/test/test_linalg.py +++ b/test/test_linalg.py @@ -23,6 +23,7 @@ onlyCUDA, skipCUDAVersionIn, skipMeta, skipCUDAIfNoCusolver) from torch.testing import floating_and_complex_types, floating_types, all_types from torch.testing._internal.common_cuda import SM53OrLater, tf32_on_and_off, CUDA11OrLater, CUDA9 +from torch.distributions.binomial import Binomial # Protects against includes accidentally setting the default dtype # NOTE: jit_metaprogramming_utils sets the default dtype to double! @@ -4477,14 +4478,15 @@ def test_qr_error_cases(self, device, dtype): with self.assertRaisesRegex(RuntimeError, "qr received unrecognized mode 'hello'"): torch.linalg.qr(t2, mode='hello') - @dtypes(torch.double, torch.cdouble) - def test_einsum(self, device, dtype): - def check(*args): + def _check_einsum(self, *args, np_args=None): + if np_args is None: np_args = [arg.cpu().numpy() if isinstance(arg, torch.Tensor) else arg for arg in args] - ref = np.einsum(*np_args) - res = torch.einsum(*args) - self.assertEqual(torch.from_numpy(np.array(ref)), res) + res = torch.einsum(*args) + ref = np.einsum(*np_args) + self.assertEqual(torch.from_numpy(np.array(ref)), res) + @dtypes(torch.double, torch.cdouble) + def test_einsum(self, device, dtype): # Test cases from https://gist.github.com/rockt/15ee013889d65342088e9260a377dc8f x = make_tensor((5,), device, dtype) y = make_tensor((7,), device, dtype) @@ -4499,102 +4501,186 @@ def check(*args): I = make_tensor((2, 3, 2), device, dtype) # Vector operations - check('i->', x) # sum - check('i,i->', x, x) # dot - check('i,i->i', x, x) # vector element-wisem mul - check('i,j->ij', x, y) # outer + self._check_einsum('i->', x) # sum + self._check_einsum('i,i->', x, x) # dot + self._check_einsum('i,i->i', x, x) # vector element-wisem mul + self._check_einsum('i,j->ij', x, y) # outer # Matrix operations - check("ij->ji", A) # transpose - check("ij->j", A) # row sum - check("ij->i", A) # col sum - check("ij,ij->ij", A, A) # matrix element-wise mul - check("ij,j->i", A, x) # matrix vector multiplication - check("ij,kj->ik", A, B) # matmul - check("ij,ab->ijab", A, E) # matrix outer product + self._check_einsum("ij->ji", A) # transpose + self._check_einsum("ij->j", A) # row sum + self._check_einsum("ij->i", A) # col sum + self._check_einsum("ij,ij->ij", A, A) # matrix element-wise mul + self._check_einsum("ij,j->i", A, x) # matrix vector multiplication + self._check_einsum("ij,kj->ik", A, B) # matmul + self._check_einsum("ij,ab->ijab", A, E) # matrix outer product # Tensor operations - check("Aij,Ajk->Aik", C, D) # batch matmul - check("ijk,jk->i", C, A) # tensor matrix contraction - check("aij,jk->aik", D, E) # tensor matrix contraction - check("abCd,dFg->abCFg", F, G) # tensor tensor contraction - check("ijk,jk->ik", C, A) # tensor matrix contraction with double indices - check("ijk,jk->ij", C, A) # tensor matrix contraction with double indices - check("ijk,ik->j", C, B) # non contiguous - check("ijk,ik->jk", C, B) # non contiguous with double indices + self._check_einsum("Aij,Ajk->Aik", C, D) # batch matmul + self._check_einsum("ijk,jk->i", C, A) # tensor matrix contraction + self._check_einsum("aij,jk->aik", D, E) # tensor matrix contraction + self._check_einsum("abCd,dFg->abCFg", F, G) # tensor tensor contraction + self._check_einsum("ijk,jk->ik", C, A) # tensor matrix contraction with double indices + self._check_einsum("ijk,jk->ij", C, A) # tensor matrix contraction with double indices + self._check_einsum("ijk,ik->j", C, B) # non contiguous + self._check_einsum("ijk,ik->jk", C, B) # non contiguous with double indices # Test diagonals - check("ii", H) # trace - check("ii->i", H) # diagonal - check('iji->j', I) # non-contiguous trace - check('ngrg...->nrg...', make_tensor((2, 1, 3, 1, 4), device, dtype)) + self._check_einsum("ii", H) # trace + self._check_einsum("ii->i", H) # diagonal + self._check_einsum('iji->j', I) # non-contiguous trace + self._check_einsum('ngrg...->nrg...', make_tensor((2, 1, 3, 1, 4), device, dtype)) # Test ellipsis - check("i...->...", H) - check("ki,...k->i...", A.t(), B) - check("k...,jk->...", A.t(), B) - check('...ik, ...j -> ...ij', C, x) - check('Bik,k...j->i...j', C, make_tensor((5, 3), device, dtype)) - check('i...j, ij... -> ...ij', C, make_tensor((2, 5, 2, 3), device, dtype)) + self._check_einsum("i...->...", H) + self._check_einsum("ki,...k->i...", A.t(), B) + self._check_einsum("k...,jk->...", A.t(), B) + self._check_einsum('...ik, ...j -> ...ij', C, x) + self._check_einsum('Bik,k...j->i...j', C, make_tensor((5, 3), device, dtype)) + self._check_einsum('i...j, ij... -> ...ij', C, make_tensor((2, 5, 2, 3), device, dtype)) # torch.bilinear with noncontiguous tensors l = make_tensor((5, 10), device, dtype, noncontiguous=True) r = make_tensor((5, 20), device, dtype, noncontiguous=True) w = make_tensor((15, 10, 20), device, dtype) - check("bn,anm,bm->ba", l, w, r) + self._check_einsum("bn,anm,bm->ba", l, w, r) # with strided tensors - check("bn,Anm,bm->bA", l[:, ::2], w[:, ::2, ::2], r[:, ::2]) + self._check_einsum("bn,Anm,bm->bA", l[:, ::2], w[:, ::2, ::2], r[:, ::2]) + + @dtypes(torch.double, torch.cdouble) + def test_einsum_sublist_format(self, device, dtype): + x = make_tensor((5,), device, dtype) + y = make_tensor((7,), device, dtype) + A = make_tensor((3, 5), device, dtype) + B = make_tensor((2, 5), device, dtype) + C = make_tensor((2, 1, 3, 1, 4), device, dtype) + + self._check_einsum(x, [0]) + self._check_einsum(x, [0], []) + self._check_einsum(x, [0], y, [1], [0, 1]) + self._check_einsum(A, [0, 1], [1, 0]) + self._check_einsum(A, [0, 1], x, [1], [0]) + self._check_einsum(A, [0, 1], B, [2, 1]) + self._check_einsum(A, [0, 1], B, [2, 1], [0, 2]) + self._check_einsum(C, [0, 1, 2, 1, Ellipsis], [0, 2, 1, Ellipsis]) + self._check_einsum(A.t(), [0, 1], B, [Ellipsis, 0]) + self._check_einsum(A.t(), [0, 1], B, [Ellipsis, 0], [1, Ellipsis]) + self._check_einsum(A.t(), [0, Ellipsis], B, [1, 0], [Ellipsis]) + + # torch.bilinear with noncontiguous tensors + l = make_tensor((5, 10), device, dtype, noncontiguous=True) + r = make_tensor((5, 20), device, dtype, noncontiguous=True) + w = make_tensor((15, 10, 20), device, dtype) + self._check_einsum(l, [40, 41], w, [2, 41, 50], r, [40, 50], [40, 2]) @dtypes(torch.double, torch.cdouble) def test_einsum_random(self, device, dtype): - def check(equation, *operands): - ref = np.einsum(equation, *[op.cpu().numpy() for op in operands]) - res = torch.einsum(equation, operands) - self.assertEqual(res.cpu(), torch.from_numpy(np.array(ref))) - - for _ in range(20): - # Create a random number of input operands, each with a random - # number of dimensions randomly labeled. - op_labels = [] - valid_labels = set() - for _ in range(random.randint(1, 3)): - labels = np.random.randint(0, 10, random.randint(1, 5)) - op_labels.append(labels) - valid_labels.update(labels) - label_size = np.random.randint(1, 5, 10) - ell_sizes = np.random.randint(1, 5, 3) - - # Build equation and tensors from input operand labels. - ops = [] - equation = '' - for labels in op_labels: - sizes = [label_size[label] for label in labels] - labels = [chr(ord('a') + label) for label in labels] - - # Add ellipsis dimensions at random - ell_num_dim = random.randint(0, 3) - if ell_num_dim > 0: - ell_index = random.randint(0, len(labels)) - sizes[ell_index:ell_index] = ell_sizes[-ell_num_dim:] - labels.insert(ell_index, "...") - - equation += ''.join(labels) + ',' - ops.append(make_tensor(sizes, device, dtype)) - equation = equation[:-1] - - # Test with implicit output - check(equation, *ops) - - # Randomly choose some labels to be part of the output - out_labels = np.unique(np.random.choice(list(valid_labels), random.randint(1, len(valid_labels)))) - out_labels = [chr(ord('a') + label) for label in out_labels] - ell_index = random.randint(0, len(out_labels)) - out_labels.insert(ell_index, '...') - equation += '->' + ''.join(out_labels) - - # Randomly test the output - check(equation, *ops) + def convert_label(label): + if label == ...: + return '...' + elif label < 26: + return chr(ord('A') + label) + else: + return chr(ord('a') + label - 26) + + def convert_sublist(sublist): + return ''.join(convert_label(label) for label in sublist) + + def test(n=10, # how many tests to generate + n_labels=5, # how many labels available + min_ops=1, max_ops=3, # min and max number of operands per test + min_dims=1, max_dims=3, # min and max number of dimensions per operand + min_size=1, max_size=8, # min and max size of each dimension + max_out_dim=3, # max number of dimensions for the output + enable_diagonals=True, # controls if labels can be repeated for diagonals + ellipsis_prob=0.5, # probability of including ellipsis in operand + broadcasting_prob=0.1): # probability of turning some dim sizes 1 for broadcasting + + all_labels = torch.arange(52) + + assert 0 <= n + assert 0 <= n_labels < len(all_labels) + assert 0 < min_ops <= max_ops + assert 0 <= min_dims <= max_dims + assert 0 <= min_size <= max_size + assert 0 <= max_out_dim + assert enable_diagonals or max_dims <= n_labels + + for _ in range(n): + + # Select a subset of labels for this test and give them random sizes + possible_labels = all_labels[torch.randperm(len(all_labels))[:n_labels]] + labels_size = torch.randint_like(all_labels, min_size, max_size + 1) + ellipsis_shape = torch.randint(min_size, max_size + 1, (max_dims - min_dims,)) + + operands = [] + sublists = [] + + ell_size = 0 + valid_labels = set() + + # create random input operands + for _ in range(random.randint(min_ops, max_ops)): + n_dim = random.randint(min_dims, max_dims) + labels_idx = torch.ones(len(possible_labels)).multinomial(n_dim, enable_diagonals) + labels = possible_labels[labels_idx] + valid_labels.update(labels.tolist()) + shape = labels_size[labels] + + # turn some dimensions to size 1 for testing broadcasting + mask = Binomial(probs=broadcasting_prob).sample((n_dim,)) + broadcast_labels = torch.unique(labels[mask == 1]) + shape[(labels[..., None] == broadcast_labels).any(-1)] = 1 + + labels = labels.tolist() + shape = shape.tolist() + + # include ellipsis if not all dimensions were assigned a label already + if n_dim < max_dims and torch.rand(1) < ellipsis_prob: + ell_num_dim = random.randint(1, max_dims - n_dim) + ell_size = max(ell_size, ell_num_dim) + ell_shape = ellipsis_shape[-ell_num_dim:] + # again, turn some dimensions to size 1 for broadcasting + mask = Binomial(probs=broadcasting_prob).sample((ell_num_dim,)) + ell_shape[mask == 1] = 1 + ell_index = random.randint(0, n_dim) + shape[ell_index:ell_index] = ell_shape + labels.insert(ell_index, ...) + + operands.append(make_tensor(shape, device, dtype)) + sublists.append(labels) + + # NumPy has a bug with the sublist format so for now we compare PyTorch sublist + # implementation against the equation format implementation of NumPy + # see https://github.com/numpy/numpy/issues/10926 + np_operands = [op.cpu().numpy() for op in operands] + + # test equation format + equation = ','.join(convert_sublist(l) for l in sublists) + self._check_einsum(equation, *operands, np_args=(equation, *np_operands)) + + # test sublist format + args = [*itertools.chain(*zip(operands, sublists))] + self._check_einsum(*args, np_args=(equation, *np_operands)) + + # generate an explicit output + out_sublist = [] + num_out_labels = max(0, random.randint(0, min(max_out_dim, len(valid_labels))) - ell_size) + if num_out_labels > 0: + out_labels_idx = torch.ones(len(valid_labels)).multinomial(num_out_labels) + out_sublist = torch.tensor(list(valid_labels))[out_labels_idx].tolist() + out_sublist.insert(random.randint(0, num_out_labels), ...) + + # test equation format with explicit output + equation += '->' + convert_sublist(out_sublist) + self._check_einsum(equation, *operands, np_args=(equation, *np_operands)) + + # test sublist format with explicit output + args.append(out_sublist) + self._check_einsum(*args, np_args=(equation, *np_operands)) + + test(100) def test_einsum_corner_cases(self, device): def check(equation, *operands, expected_output): @@ -4668,38 +4754,6 @@ def check(*args, regex, exception=RuntimeError): check(x, [-1], regex=r'not within the valid range \[0, 52\)', exception=ValueError) check(x, [52], regex=r'not within the valid range \[0, 52\)', exception=ValueError) - @dtypes(torch.double, torch.cdouble) - def test_einsum_sublist_format(self, device, dtype): - def check(*args): - np_args = [arg.cpu().numpy() if isinstance(arg, torch.Tensor) else arg for arg in args] - ref = np.einsum(*np_args) - res = torch.einsum(*args) - self.assertEqual(torch.from_numpy(np.array(ref)), res) - - x = make_tensor((5,), device, dtype) - y = make_tensor((7,), device, dtype) - A = make_tensor((3, 5), device, dtype) - B = make_tensor((2, 5), device, dtype) - C = make_tensor((2, 1, 3, 1, 4), device, dtype) - - check(x, [0]) - check(x, [0], []) - check(x, [0], y, [1], [0, 1]) - check(A, [0, 1], [1, 0]) - check(A, [0, 1], x, [1], [0]) - check(A, [0, 1], B, [2, 1]) - check(A, [0, 1], B, [2, 1], [0, 2]) - check(C, [0, 1, 2, 1, Ellipsis], [0, 2, 1, Ellipsis]) - check(A.t(), [0, 1], B, [Ellipsis, 0]) - check(A.t(), [0, 1], B, [Ellipsis, 0], [1, Ellipsis]) - check(A.t(), [0, Ellipsis], B, [1, 0], [Ellipsis]) - - # torch.bilinear with noncontiguous tensors - l = make_tensor((5, 10), device, dtype, noncontiguous=True) - r = make_tensor((5, 20), device, dtype, noncontiguous=True) - w = make_tensor((15, 10, 20), device, dtype) - check(l, [40, 41], w, [2, 41, 50], r, [40, 50], [40, 2]) - def triangular_solve_test_helper(self, A_dims, b_dims, upper, unitriangular, device, dtype): triangle_function = torch.triu if upper else torch.tril diff --git a/torch/functional.py b/torch/functional.py index 8f50ae80ccad0..840357d238835 100644 --- a/torch/functional.py +++ b/torch/functional.py @@ -294,9 +294,9 @@ def parse_subscript(n: int) -> str: if n == Ellipsis: return '...' if n >= 0 and n < 26: - return chr(n + ord('a')) + return chr(ord('A') + n) if n >= 26 and n < 52: - return chr(n - 26 + ord('A')) + return chr(ord('a') + n - 26) raise ValueError('einsum(): subscript in subscript list is not within the valid range [0, 52)') # Parse subscripts for input operands From 59b10036d557cf99df90a25d3a3a04d6312ac4f6 Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Thu, 17 Jun 2021 06:33:48 -0700 Subject: [PATCH 189/305] Unifies OpInfo dtype tests (#60157) Summary: Simplifies the OpInfo dtype tests and produces nicer error messages, like: ``` AssertionError: Items in the first set but not the second: torch.bfloat16 Items in the second set but not the first: torch.int64 : Attempted to compare [set] types: Expected: {torch.float64, torch.float32, torch.float16, torch.bfloat16}; Actual: {torch.float64, torch.float32, torch.float16, torch.int64}. The supported dtypes for logcumsumexp on cuda according to its OpInfo are {torch.float64, torch.float32, torch.float16, torch.int64}, but the detected supported dtypes are {torch.float64, torch.float32, torch.float16, torch.bfloat16}. The following dtypes should be added to the OpInfo: {torch.bfloat16}. The following dtypes should be removed from the OpInfo: {torch.int64}. ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/60157 Reviewed By: ngimel Differential Revision: D29188665 Pulled By: mruberry fbshipit-source-id: e84c9892c6040ea47adb027cfef3a6c0fd2f9f3c --- test/test_ops.py | 1162 +++++++++-------- torch/testing/_internal/common_device_type.py | 53 +- .../_internal/common_methods_invocations.py | 391 +++--- 3 files changed, 833 insertions(+), 773 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 41efa551c46cb..9b5f5dab6a9e4 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -5,7 +5,7 @@ import torch from torch.testing import \ - (FileCheck, floating_and_complex_types_and) + (FileCheck, floating_and_complex_types_and, get_all_dtypes) from torch.testing._internal.common_utils import \ (TestCase, is_iterable_of_tensors, run_tests, IS_SANDCASTLE, clone_input_helper, make_tensor, gradcheck, gradgradcheck, IS_PYTORCH_CI) @@ -14,17 +14,23 @@ from torch.testing._internal.common_device_type import \ (instantiate_device_type_tests, ops, onlyOnCPUAndCUDA, skipCUDAIfRocm, OpDTypes) from torch.testing._internal.common_jit import JitCommonTestCase, check_against_reference - from torch.testing._internal.jit_metaprogramming_utils import create_script_fn, create_traced_fn, \ check_alias_annotation from torch.testing._internal.jit_utils import disable_autodiff_subgraph_inlining import torch.testing._internal.opinfo_helper as opinfo_helper +# variant testing is only done with torch.float and torch.cfloat to avoid +# excessive test times and maximize signal to noise ratio +_variant_ops = partial(ops, dtypes=OpDTypes.supported, + allowed_dtypes=(torch.float, torch.cfloat)) + -# Tests that apply to all operators -class TestOpInfo(TestCase): +# Tests that apply to all operators and aren't related to any particular +# system +class TestCommon(TestCase): exact_dtype = True + # Verifies, on teardown, that no OpInfo is still using dynamic dtypes in CI @classmethod def tearDownClass(cls): super().tearDownClass() @@ -40,254 +46,394 @@ def tearDownClass(cls): assert len(filtered_ops) == 0, err_msg - # Verifies that ops have their unsupported dtypes - # registered correctly by testing that each claimed unsupported dtype - # throws a runtime error + # Validates that each OpInfo specifies its forward and backward dtypes + # correctly for CPU and CUDA devices @skipCUDAIfRocm @onlyOnCPUAndCUDA - @ops(op_db, dtypes=OpDTypes.unsupported) - def test_unsupported_dtypes(self, device, dtype, op): - # sample_inputs can have a function for generating the input that doesn't work for specified dtype - # https://github.com/pytorch/pytorch/issues/49024 - with self.assertRaises(RuntimeError): - samples = op.sample_inputs(device, dtype) - for sample in samples: - op(sample.input, *sample.args, **sample.kwargs) - - # Verifies that ops have their supported dtypes - # registered correctly by testing that each claimed supported dtype - # does NOT throw a runtime error - # In addition verifies that the generated sample_inputs have the requested device and dtype - @onlyOnCPUAndCUDA - @ops(op_db, dtypes=OpDTypes.supported) - def test_supported_dtypes(self, device, dtype, op): - for sample in op.sample_inputs(device, dtype): - op(sample.input, *sample.args, **sample.kwargs) - # NOTE: only check the first tensor in the iterable of tensors - sample_input = sample.input[0] if is_iterable_of_tensors(sample.input) else sample.input - self.assertTrue(sample_input.dtype == dtype) - self.assertTrue(sample_input.device.type == self.device_type) - - # Verifies that backward for each unsupported floating or complex dtype - # throw a runtime error. - @onlyOnCPUAndCUDA - @ops(op_db, dtypes=OpDTypes.unsupported_backward, - allowed_dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16)) - def test_unsupported_backward(self, device, dtype, op): - if not op.supports_autograd: - self.skipTest("Skipped! Autograd not supported.") - - try: - samples = op.sample_inputs(device, dtype, requires_grad=True) - except RuntimeError as e: - self.skipTest(f"Skipped! unable to generate sample. {e}") + @ops(op_db, dtypes=OpDTypes.none) + def test_dtypes(self, device, dtype, op): + # dtypes to try to backward in + allowed_backward_dtypes = floating_and_complex_types_and(torch.bfloat16, torch.float16) + + # lists for (un)supported dtypes + supported_dtypes = [] + unsupported_dtypes = [] + supported_backward_dtypes = [] + unsupported_backward_dtypes = [] + + def unsupported(dtype): + unsupported_dtypes.append(dtype) + if dtype in allowed_backward_dtypes: + unsupported_backward_dtypes.append(dtype) + + for dtype in get_all_dtypes(): + # tries to acquire samples - failure indicates lack of support + requires_grad = (dtype in allowed_backward_dtypes and op.supports_autograd) + try: + samples = op.sample_inputs(device, dtype, requires_grad=requires_grad) + except Exception as e: + unsupported(dtype) + continue - if len(samples) == 0: - self.skipTest("Skipped! No sample inputs!") + # Counts number of successful backward attempts + # NOTE: This exists as a kludge because this only understands how to + # request a gradient if the output is a tensor or a sequence with + # a tensor as its first element. + num_backward_successes = 0 + for sample in samples: + # tries to call operator with the sample - failure indicates + # lack of support + try: + result = op(sample.input, *sample.args, **sample.kwargs) + except Exception as e: + # NOTE: some ops will fail in forward if their inputs + # require grad but they don't support computing the gradient + # in that type! This is a bug in the op! + unsupported(dtype) + + # Short-circuits testing this dtype -- it doesn't work + if dtype in unsupported_dtypes: + break + + # Short-circuits if the dtype isn't a backward dtype or + # it's already identified as not supported + if dtype not in allowed_backward_dtypes or dtype in unsupported_backward_dtypes: + continue - # NOTE: assert exception raised on ANY sample input - with self.assertRaises(RuntimeError): - for sample in op.sample_inputs(device, dtype, requires_grad=True): - result = op(sample.input, *sample.args, **sample.kwargs) - # TODO: handle non-tensor outputs - if not isinstance(result, torch.Tensor): - self.skipTest("Skipped! Test does not handle non-tensor outputs") - if sample.output_process_fn_grad is not None: + # Checks for backward support in the same dtype + try: result = sample.output_process_fn_grad(result) - result.sum().backward() + if isinstance(result, torch.Tensor): + backward_tensor = result + elif isinstance(result, Sequence) and isinstance(result[0], torch.Tensor): + backward_tensor = result[0] + else: + continue - # Verifies that backward for each supported floating or complex dtype - # does NOT throw a runtime error. - # TODO: support multi-tensor outputs - @onlyOnCPUAndCUDA - @ops(op_db, dtypes=OpDTypes.supported_backward, - allowed_dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16)) - def test_supported_backward(self, device, dtype, op): - if not op.supports_autograd: - self.skipTest("Skipped! Autograd not supported.") + # Note: this grad may not have the same dtype as dtype + # For functions like complex (float -> complex) or abs + # (complex -> float) the grad tensor will have a + # different dtype than the input. + # For simplicity, this is still modeled as these ops + # supporting grad in the input dtype. + grad = torch.randn_like(backward_tensor) + backward_tensor.backward(grad) + num_backward_successes += 1 + except Exception as e: + unsupported_backward_dtypes.append(dtype) + + if dtype not in unsupported_dtypes: + supported_dtypes.append(dtype) + if num_backward_successes > 0 and dtype not in unsupported_backward_dtypes: + supported_backward_dtypes.append(dtype) + + # Checks that dtypes are listed correctly and generates an informative + # error message + device_type = torch.device(device).type + claimed_supported = set(op.supported_dtypes(device_type)) + supported_dtypes = set(supported_dtypes) + + supported_but_unclaimed = supported_dtypes - claimed_supported + claimed_but_unsupported = claimed_supported - supported_dtypes + msg = """The supported dtypes for {0} on {1} according to its OpInfo are + {2}, but the detected supported dtypes are {3}. + """.format(op.name, device_type, claimed_supported, supported_dtypes) + + if len(supported_but_unclaimed) > 0: + msg += "The following dtypes should be added to the OpInfo: {0}. ".format(supported_but_unclaimed) + if len(claimed_but_unsupported) > 0: + msg += "The following dtypes should be removed from the OpInfo: {0}.".format(claimed_but_unsupported) + + self.assertEqual(supported_dtypes, claimed_supported, msg=msg) + + # Checks that backward dtypes are listed correctly and generates an + # informative error message + # NOTE: this code is nearly identical to the check + msg generation + claimed_backward_supported = set(op.supported_backward_dtypes(device_type)) + supported_backward_dtypes = set(supported_backward_dtypes) + + supported_but_unclaimed = supported_backward_dtypes - claimed_backward_supported + claimed_but_unsupported = claimed_backward_supported - supported_backward_dtypes + msg = """The supported backward dtypes for {0} on {1} according to its OpInfo are + {2}, but the detected supported backward dtypes are {3}. + """.format(op.name, device_type, claimed_backward_supported, supported_backward_dtypes) + + if len(supported_but_unclaimed) > 0: + msg += "The following backward dtypes should be added to the OpInfo: {0}. ".format(supported_but_unclaimed) + if len(claimed_but_unsupported) > 0: + msg += "The following backward dtypes should be removed from the OpInfo: {0}.".format(claimed_but_unsupported) + + self.assertEqual(supported_backward_dtypes, claimed_backward_supported, msg=msg) - for sample in op.sample_inputs(device, dtype, requires_grad=True): - result = op(sample.input, *sample.args, **sample.kwargs) - if not isinstance(result, torch.Tensor): - continue - if sample.output_process_fn_grad is not None: - result = sample.output_process_fn_grad(result) - result.sum().backward() + # Validates ops implement the correct out= behavior + # See https://github.com/pytorch/pytorch/wiki/Developer-FAQ#how-does-out-work-in-pytorch + # for a description of the correct behavior + # TODO: operations that support out= but don't support float + # are not covered by this test. + @ops(op_db, allowed_dtypes=(torch.float,)) + def test_out(self, device, dtype, op): + # TODO: verify the op doesn't support the out= kwarg + if not op.supports_out: + self.skipTest("Skipped! Op doesn't support out= kwarg.") + # NOTE: only tests on first sample + samples = op.sample_inputs(device, dtype) + sample = samples[0] -# gradcheck requires double precision -_gradcheck_ops = partial(ops, dtypes=OpDTypes.supported, - allowed_dtypes=[torch.double, torch.cdouble]) + # calls it normally to get the expected result + expected = op(sample.input, *sample.args, **sample.kwargs) + op_out = partial(op, sample.input, *sample.args, **sample.kwargs) + # Short-circuits if output is not a single tensor or an + # iterable of tensors -class TestGradients(TestCase): - exact_dtype = True + if not isinstance(expected, torch.Tensor) and not is_iterable_of_tensors(expected, include_empty=True): + self.skipTest("Skipped! Only supports single tensor or iterable of tensor outputs.") - # Copies inputs to inplace operations to avoid inplace modifications - # to leaves requiring gradient - def _get_safe_inplace(self, inplace_variant): - @wraps(inplace_variant) - def _fn(t, *args, **kwargs): - return inplace_variant(t.clone(), *args, **kwargs) + # A wrapper around map that works with single tensors and always + # instantiates the map. Used below to apply transforms to + # single tensor and iterable tensor outputs. + def _apply_out_transform(fn, out): + if isinstance(out, torch.Tensor): + return fn(out) - return _fn + # assumes (see above) that out is an iterable of tensors + return tuple(map(fn, out)) - def _check_helper(self, device, dtype, op, variant, check, *, check_forward_ad=False): - if variant is None: - self.skipTest("Skipped! Variant not implemented.") - if not op.supports_dtype(dtype, torch.device(device).type): - self.skipTest(f"Skipped! {op.name} does not support dtype {str(dtype)}") + # Case 0: out= with the correct shape, dtype, and device + # but NaN values for floating point and complex tensors, and + # maximum values for integer tensors. + # Expected behavior: out= values have no effect on the computation. + def _case_zero_transform(t): + try: + info = torch.iinfo(t.dtype) + return torch.full_like(t, info.max) + except TypeError as te: + # for non-integer types fills with NaN + return torch.full_like(t, float('nan')) - def is_inplace(variant): - if hasattr(variant, "__wrapped__"): - return variant.__wrapped__ is op.get_inplace() - return variant is op.get_inplace() + out = _apply_out_transform(_case_zero_transform, expected) + result = op_out(out=out) + self.assertEqual(expected, out) - include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex - samples = op.sample_inputs(device, dtype, requires_grad=True, include_conjugated_inputs=include_conjugated_inputs) + # Checks that the returned value shares storage with out + # NOTE: only checks on the CPU and CUDA device types since some + # device types don't have storage + if self.device_type == 'cpu' or self.device_type == 'cuda': + if isinstance(out, torch.Tensor): + self.assertEqual(out.storage().data_ptr(), result.storage().data_ptr()) + else: + for out_t, result_t in zip(out, result): + self.assertEqual(out_t.storage().data_ptr(), result_t.storage().data_ptr()) - for sample in samples: - if sample.broadcasts_input and is_inplace(variant): - continue + # Case 1: out= with the correct shape, dtype, and device, + # but noncontiguous. + # Expected behavior: strides are respected and `out` storage is not changed. + def _case_one_transform(t): + return make_tensor(t.shape, + dtype=t.dtype, + device=t.device, + noncontiguous=True) - # Note on TensorList inputs - # - # gradcheck does not support TensorList inputs so here we pass TensorList - # inputs of size n as n single Tensor inputs to gradcheck and wrap the op - # in a function that puts the n Tensor inputs back into a TensorList - def fn(*inputs): - # Put tensors back into TensorList since we splat them when passing to gradcheck - if is_iterable_of_tensors(sample.input): - n = len(sample.input) - inputs = (inputs[:n], *inputs[n:]) - output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs) - if sample.output_process_fn_grad is not None: - return sample.output_process_fn_grad(output) - return output + # Extracts strides from a tensor or iterable of tensors into a tuple + def _extract_strides(out): + if isinstance(out, torch.Tensor): + return (out.stride(),) - # Splat TensorList inputs into single Tensor inputs - gradcheck_args = (sample.input,) if isinstance(sample.input, torch.Tensor) else tuple(sample.input) - gradcheck_args += sample.args + # assumes (see above) that out is an iterable of tensors + return tuple(map(lambda t: t.stride(), out)) - if check == 'gradcheck': - self.assertTrue(gradcheck(fn, gradcheck_args, - check_batched_grad=op.check_batched_grad, - check_grad_dtypes=True, - nondet_tol=op.gradcheck_nondet_tol, - fast_mode=op.gradcheck_fast_mode, - check_forward_ad=check_forward_ad)) - elif check == 'gradgradcheck': - self.assertFalse(check_forward_ad, msg="Cannot run forward AD check for gradgradcheck") - self.assertTrue(gradgradcheck(fn, gradcheck_args, - gen_non_contig_grad_outputs=False, - check_batched_grad=op.check_batched_gradgrad, - check_grad_dtypes=True, - nondet_tol=op.gradcheck_nondet_tol, - fast_mode=op.gradcheck_fast_mode)) - self.assertTrue(gradgradcheck(fn, gradcheck_args, - gen_non_contig_grad_outputs=True, - check_batched_grad=op.check_batched_gradgrad, - check_grad_dtypes=True, - nondet_tol=op.gradcheck_nondet_tol, - fast_mode=op.gradcheck_fast_mode)) - else: - self.assertTrue(False, msg="Unknown check requested!") + def _extract_data_ptrs(out): + if isinstance(out, torch.Tensor): + return (out.data_ptr(),) - def _grad_test_helper(self, device, dtype, op, variant, *, check_forward_ad=False): - return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad) + # assumes (see above) that out is an iterable of tensors + return tuple(map(lambda t: t.data_ptr(), out)) - def _gradgrad_test_helper(self, device, dtype, op, variant): - return self._check_helper(device, dtype, op, variant, 'gradgradcheck') - def _skip_helper(self, op, device, dtype): - if not op.supports_autograd: - self.skipTest("Skipped! autograd not supported.") - if not op.supports_complex_autograd(torch.device(device).type) and dtype.is_complex: - self.skipTest("Skipped! Complex autograd not supported.") + out = _apply_out_transform(_case_one_transform, expected) + original_strides = _extract_strides(out) + original_ptrs = _extract_data_ptrs(out) - # Tests that gradients are computed correctly - @_gradcheck_ops(op_db) - def test_fn_grad(self, device, dtype, op): - self._skip_helper(op, device, dtype) - self._grad_test_helper(device, dtype, op, op.get_op()) + op_out(out=out) + final_strides = _extract_strides(out) + final_ptrs = _extract_data_ptrs(out) - # Method grad (and gradgrad, see below) tests are disabled since they're - # costly and redundant with function grad (and gradgad) tests - # @_gradcheck_ops(op_db) - # def test_method_grad(self, device, dtype, op): - # self._skip_helper(op, device, dtype) - # self._grad_test_helper(device, dtype, op, op.get_method()) + self.assertEqual(expected, out) + self.assertEqual(original_strides, final_strides) + self.assertEqual(original_ptrs, final_ptrs) - @_gradcheck_ops(op_db) - def test_inplace_grad(self, device, dtype, op): - self._skip_helper(op, device, dtype) - if not op.inplace_variant or not op.supports_inplace_autograd: - self.skipTest("Skipped! Operation does not support inplace autograd.") - self._grad_test_helper(device, dtype, op, self._get_safe_inplace(op.get_inplace())) + # Case 2: out= with the correct dtype and device, but the wrong shape + # Expected behavior: resize with a warning. + def _case_two_transform(t): + wrong_shape = list(t.shape) - # Test that gradients of gradients are computed correctly - @_gradcheck_ops(op_db) - def test_fn_gradgrad(self, device, dtype, op): - self._skip_helper(op, device, dtype) - if not op.supports_gradgrad: - self.skipTest("Skipped! Operation does not support gradgrad") - self._gradgrad_test_helper(device, dtype, op, op.get_op()) + if len(wrong_shape) == 0: + # Handles scalar tensor case (empty list) + wrong_shape = [2] + else: + wrong_shape[-1] = wrong_shape[-1] + 1 + return make_tensor(wrong_shape, dtype=t.dtype, device=t.device) - # Test that gradients of gradients are properly raising - @_gradcheck_ops(op_db) - def test_fn_fail_gradgrad(self, device, dtype, op): - self._skip_helper(op, device, dtype) - if op.supports_gradgrad: - self.skipTest("Skipped! Operation does support gradgrad") + out = _apply_out_transform(_case_two_transform, expected) + msg_fail = "Resized a non-empty tensor but did not warn about it." + with self.assertWarnsRegex(UserWarning, "An output with one or more elements", msg=msg_fail): + op_out(out=out) + self.assertEqual(expected, out) - err_msg = r"derivative for .* is not implemented" - with self.assertRaisesRegex(RuntimeError, err_msg): - self._gradgrad_test_helper(device, dtype, op, op.get_op()) + # Case 3: out= with the correct dtype and device, but an empty + # tensor. + # Expected behavior: resize without warning. + def _case_three_transform(t): + return make_tensor((0,), + dtype=t.dtype, + device=t.device) - # Method gradgrad (and grad, see above) tests are disabled since they're - # costly and redundant with function gradgrad (and grad) tests - # @_gradcheck_ops(op_db) - # def test_method_gradgrad(self, device, dtype, op): - # self._skip_helper(op, device, dtype) - # self._gradgrad_test_helper(device, dtype, op, op.get_method()) + out = _apply_out_transform(_case_three_transform, expected) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + op_out(out=out) - @_gradcheck_ops(op_db) - def test_inplace_gradgrad(self, device, dtype, op): - self._skip_helper(op, device, dtype) - if not op.inplace_variant or not op.supports_inplace_autograd: - self.skipTest("Skipped! Operation does not support inplace autograd.") - self._gradgrad_test_helper(device, dtype, op, self._get_safe_inplace(op.get_inplace())) + # Verifies no warning is a resize warning + for w in caught: + if "An output with one or more elements" in str(w.message): + self.fail("Resizing an out= argument with no elements threw a resize warning!") - @_gradcheck_ops(op_db) - def test_forward_mode_AD(self, device, dtype, op): - self._skip_helper(op, device, dtype) + self.assertEqual(expected, out) - if op.supports_forward_ad: - self._grad_test_helper(device, dtype, op, op.get_op(), check_forward_ad=True) - else: - err_msg = r"Trying to use forward AD with .* that does not support it\." - hint_msg = ("Running forward AD for an OP that has does not support it did not " - "raise any error. If your op supports forward AD, you should set supports_forward_ad=True") - with self.assertRaisesRegex(NotImplementedError, err_msg, msg=hint_msg): - self._grad_test_helper(device, dtype, op, op.get_op(), check_forward_ad=True) + # Case 4: out= with correct shape and dtype, but wrong device. + wrong_device = None + if torch.device(device).type != 'cpu': + wrong_device = 'cpu' + elif torch.cuda.is_available(): + wrong_device = 'cuda' + if wrong_device is not None: + def _case_four_transform(t): + return make_tensor(t.shape, dtype=t.dtype, device=wrong_device) -# Tests operators for consistency between JIT and eager, also checks -# correctness of JIT specific alias schemas and intended -# autodifferentiation behavior. -# Inherits from JitCommonTestCase instead of TestCase directly to share -# functionality with original test_jit.py method operator tests -class TestCommon(JitCommonTestCase): - exact_dtype = True + out = _apply_out_transform(_case_four_transform, expected) + msg_fail = f"Expected RuntimeError when calling with input.device={device} and out.device={wrong_device}" + with self.assertRaises(RuntimeError, msg=msg_fail): + op_out(out=out) + + # Case 5: out= with correct shape and device, but a dtype + # that output cannot be "safely" cast to (long). + # Expected behavior: error. + # NOTE: this case is filtered by dtype since some ops produce + # bool tensors, for example, which can be safely cast to any + # dtype. It is applied when single tensors are floating point or complex + # dtypes, or if an op returns multiple tensors when at least one such + # tensor is a floating point or complex dtype. + _dtypes = floating_and_complex_types_and(torch.float16, torch.bfloat16) + if (isinstance(expected, torch.Tensor) and expected.dtype in _dtypes or + (not isinstance(expected, torch.Tensor) and any(t.dtype in _dtypes for t in expected))): + def _case_five_transform(t): + return make_tensor(t.shape, dtype=torch.long, device=t.device) - # variant testing is only done with torch.float and torch.cfloat to avoid - # excessive test times and maximize signal to noise ratio - _variant_ops = partial(ops, dtypes=OpDTypes.supported, - allowed_dtypes=(torch.float, torch.cfloat)) + out = _apply_out_transform(_case_five_transform, expected) + msg_fail = "" if not isinstance(expected, torch.Tensor) else \ + ("Expected RuntimeError when doing an unsafe cast from a result of dtype " + f"{expected.dtype} into an out= with dtype torch.long") + with self.assertRaises(RuntimeError, msg=msg_fail): + op_out(out=out) - # alias testing is only done with troch.float for the same reason - _alias_ops = partial(ops, dtypes=OpDTypes.supported, - allowed_dtypes=(torch.float,)) + # Tests that + # 1. The operator's output for physically conjugated tensors and conjugate view tensors + # produces the same value + # 2. The gradients are same in both cases mentioned in (1) + # 3. If the operator's inplace variant is supported, tests that the inplace operation + # produces the correct value when called on a conjugate view tensor and that the output + # has its conj bit set to true + # This test only runs for C -> R and C -> C functions + # TODO: add tests for `R->C` functions + # Note: This test runs for functions that take both tensors and tensorlists as input. + @ops(op_db, allowed_dtypes=(torch.cfloat,)) + def test_conj_view(self, device, dtype, op): + if not op.test_conjugated_samples: + self.skipTest("Operation doesn't support conjugated inputs.") + _requires_grad = (op.supports_autograd and op.supports_complex_autograd(torch.device(device).type)) + samples = op.sample_inputs(device, dtype, requires_grad=_requires_grad) + inplace_variant = op.inplace_variant + + # helper function to physically conjugate the tensor + def conjugate_physical(input): + if isinstance(input, torch.Tensor): + tensor_requires_grad = input.requires_grad + with torch.no_grad(): + input = input.conj_physical() + return input.requires_grad_(tensor_requires_grad) + + if isinstance(input, Sequence): + out = list(map(clone_input_helper, input)) + out[0] = conjugate_physical(out[0]) + return tuple(out) + + # helper function to clone and conjugate the input if its a tensor + # else clone the sequence and conjugate the first element in the sequence + # If a requires_grad argument is provided the tensor being conjugated will + # have its requires_grad set to that value. + def clone_conj_input_helper(input, **kwargs): + if isinstance(input, torch.Tensor): + requires_grad = kwargs.get('requires_grad', input.requires_grad) + with torch.no_grad(): + input = input.clone() + # Note: .conj() is not called under no_grad mode since it's not allowed to modify a + # view created in no_grad mode. Here it's ok to do so, so as a workaround we call conj + # before resetting the requires_grad field for input + input = input.conj() + assert input.is_leaf + return input.requires_grad_(requires_grad) + + if isinstance(input, Sequence): + out = list(map(clone_input_helper, input)) + out[0] = clone_conj_input_helper(out[0]) + return tuple(out) + + for sample in samples: + tensor = sample.input if isinstance(sample.input, torch.Tensor) else sample.input[0] + cloned1 = clone_conj_input_helper(sample.input) + sample.input = conjugate_physical(sample.input) + + # Computes function forward value with a physically conjugated tensor and + # a conj view tensor and verifies that the output in both case are equal. + expected_forward = op(sample.input, *sample.args, **sample.kwargs) + forward_with_conjview = op(cloned1, *sample.args, **sample.kwargs) + self.assertEqual(expected_forward, forward_with_conjview) + + # If the op has an inplace variant, and the input doesn't require broadcasting + # and has the same dtype as output, verify that the inplace operation on a conjugated + # input produces correct output, and the output tensor has the conj bit set to True + if inplace_variant is not None and not sample.broadcasts_input: + cloned2 = clone_conj_input_helper(tensor, requires_grad=False) + if (isinstance(expected_forward, torch.Tensor) and + expected_forward.dtype is tensor.dtype): + inplace_forward = inplace_variant(cloned2, *sample.args, **sample.kwargs) + self.assertTrue(inplace_forward.is_conj()) + self.assertEqual(inplace_forward, expected_forward) + + # TODO: backward consistency only supported for single tensor outputs + # TODO: backward consistency only checked on sample.input, not all + # tensor inputs + # TODO: update to handle checking grads of all tensor inputs as + # derived from each tensor output + if isinstance(expected_forward, torch.Tensor) and expected_forward.requires_grad: + tensor = sample.input if isinstance(sample.input, torch.Tensor) else sample.input[0] + expected_forward.sum().backward(retain_graph=True) + forward_with_conjview.sum().backward(retain_graph=True) + if tensor.grad is not None: + cloned1_tensor = cloned1 if isinstance(cloned1, torch.Tensor) else cloned1[0] + self.assertEqual(tensor.grad, cloned1_tensor.grad) + + tensor.grad, cloned1_tensor.grad = None, None + + # a repeat of the above test if output is not complex valued + if (expected_forward.is_complex()): + grad = torch.randn_like(expected_forward) + expected_forward.backward(grad.conj_physical()) + forward_with_conjview.backward(grad.conj()) + + self.assertEqual(tensor.grad, cloned1_tensor.grad) # Tests that the forward and backward passes of operations produce the # same values for the cross-product of op variants (method, inplace) @@ -411,89 +557,259 @@ def _test_inplace_preserve_storage(samples, variants): inplace_samples = list(filter(lambda sample: not sample.broadcasts_input, samples)) _test_inplace_preserve_storage(inplace_samples, inplace_variants) - # Tests that the forward and backward passes of operations produce the - # same values for the cross-product of op variants (function, method, inplace) - # and runtimes (eager, traced, scripted). - # TODO WARNING: inplace x {traced, scripted} not currently tested - @_variant_ops(op_db) - def test_variant_consistency_jit(self, device, dtype, op): - _requires_grad = op.supports_autograd and (dtype.is_floating_point or - op.supports_complex_autograd(torch.device(device).type)) - include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex - samples = op.sample_inputs(device, dtype, requires_grad=_requires_grad, include_conjugated_inputs=include_conjugated_inputs) +# gradcheck requires double precision +_gradcheck_ops = partial(ops, dtypes=OpDTypes.supported, + allowed_dtypes=[torch.double, torch.cdouble]) - for sample in samples: - # Acquires variants to test - func = op.get_op() - method = op.get_method() - variants = { - # TODO: inplace tests currently fail, fix and add inplace variant - 'function': func, 'method': method, - } - # Test traced and scripted consistency - for func_type, variant in variants.items(): - if variant is None: - continue +class TestGradients(TestCase): + exact_dtype = True - # Create accessor for script function variant - name = op.name + '_' if func_type == 'inplace' else op.name + # Copies inputs to inplace operations to avoid inplace modifications + # to leaves requiring gradient + def _get_safe_inplace(self, inplace_variant): + @wraps(inplace_variant) + def _fn(t, *args, **kwargs): + return inplace_variant(t.clone(), *args, **kwargs) - # run with disable_autodiff_subgraph_inlining(True) to test - # autodiff support. Context manager forces the graph to contain - # DifferentiableGraph nodes if they are present - with disable_autodiff_subgraph_inlining(): - # Check scripted forward, grad, and grad grad - script_fn = create_script_fn(self, name, func_type) + return _fn - def out_fn(output): - # Processes the output for autograd - if sample.output_process_fn_grad is not None: - return sample.output_process_fn_grad(output) - return output + def _check_helper(self, device, dtype, op, variant, check, *, check_forward_ad=False): + if variant is None: + self.skipTest("Skipped! Variant not implemented.") + if not op.supports_dtype(dtype, torch.device(device).type): + self.skipTest(f"Skipped! {op.name} does not support dtype {str(dtype)}") - check_against_reference(self, - script_fn, - func, - out_fn, - (sample.input,) + sample.args, - sample.kwargs, - no_grad=not _requires_grad, no_gradgrad=not op.supports_gradgrad) + def is_inplace(variant): + if hasattr(variant, "__wrapped__"): + return variant.__wrapped__ is op.get_inplace() + return variant is op.get_inplace() - # Check traced forward, grad, and grad grad - traced_fn = create_traced_fn(self, variant) - check_against_reference(self, - traced_fn, - func, - out_fn, - (sample.input,) + sample.args, - sample.kwargs, - no_grad=not _requires_grad, no_gradgrad=not op.supports_gradgrad) + include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex + samples = op.sample_inputs(device, dtype, requires_grad=True, include_conjugated_inputs=include_conjugated_inputs) - # Check alias annotation schema for correctness (make - # sure inputs that aren't supposed to be modified aren't) - # Note: only runs in float32 and int64 because schema isn't affected by dtype, - # so running it on all dtypes is would be excessive - if dtype in [torch.float32, torch.int32]: - check_alias_annotation(name, (sample.input,) + sample.args, sample.kwargs, - func_type=func_type, aten_name=op.aten_name) + for sample in samples: + if sample.broadcasts_input and is_inplace(variant): + continue - # Check autodifferentiation of nodes for traced and scripted graphs, only need to check once per sample - if dtype is torch.float32: - # Sandcastle doesn't fuse nodes - if IS_SANDCASTLE: - # fusible nodes are expected to be found in FusionGroups in the DifferentiableGraphs - nonfusible_nodes = op.autodiff_nonfusible_nodes + op.autodiff_fusible_nodes - fusible_nodes = [] - else: - nonfusible_nodes = op.autodiff_nonfusible_nodes - fusible_nodes = op.autodiff_fusible_nodes + # Note on TensorList inputs + # + # gradcheck does not support TensorList inputs so here we pass TensorList + # inputs of size n as n single Tensor inputs to gradcheck and wrap the op + # in a function that puts the n Tensor inputs back into a TensorList + def fn(*inputs): + # Put tensors back into TensorList since we splat them when passing to gradcheck + if is_iterable_of_tensors(sample.input): + n = len(sample.input) + inputs = (inputs[:n], *inputs[n:]) + output = op.gradcheck_wrapper(variant, *inputs, **sample.kwargs) + if sample.output_process_fn_grad is not None: + return sample.output_process_fn_grad(output) + return output - self.assertAutodiffNode(traced_fn.last_graph, op.assert_autodiffed, nonfusible_nodes, fusible_nodes) - self.assertAutodiffNode(script_fn.last_graph, op.assert_autodiffed, nonfusible_nodes, fusible_nodes) + # Splat TensorList inputs into single Tensor inputs + gradcheck_args = (sample.input,) if isinstance(sample.input, torch.Tensor) else tuple(sample.input) + gradcheck_args += sample.args - @_alias_ops((op for op in op_db if op.aliases)) + if check == 'gradcheck': + self.assertTrue(gradcheck(fn, gradcheck_args, + check_batched_grad=op.check_batched_grad, + check_grad_dtypes=True, + nondet_tol=op.gradcheck_nondet_tol, + fast_mode=op.gradcheck_fast_mode, + check_forward_ad=check_forward_ad)) + elif check == 'gradgradcheck': + self.assertFalse(check_forward_ad, msg="Cannot run forward AD check for gradgradcheck") + self.assertTrue(gradgradcheck(fn, gradcheck_args, + gen_non_contig_grad_outputs=False, + check_batched_grad=op.check_batched_gradgrad, + check_grad_dtypes=True, + nondet_tol=op.gradcheck_nondet_tol, + fast_mode=op.gradcheck_fast_mode)) + self.assertTrue(gradgradcheck(fn, gradcheck_args, + gen_non_contig_grad_outputs=True, + check_batched_grad=op.check_batched_gradgrad, + check_grad_dtypes=True, + nondet_tol=op.gradcheck_nondet_tol, + fast_mode=op.gradcheck_fast_mode)) + else: + self.assertTrue(False, msg="Unknown check requested!") + + def _grad_test_helper(self, device, dtype, op, variant, *, check_forward_ad=False): + return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad) + + def _gradgrad_test_helper(self, device, dtype, op, variant): + return self._check_helper(device, dtype, op, variant, 'gradgradcheck') + + def _skip_helper(self, op, device, dtype): + if not op.supports_autograd: + self.skipTest("Skipped! autograd not supported.") + if not op.supports_complex_autograd(torch.device(device).type) and dtype.is_complex: + self.skipTest("Skipped! Complex autograd not supported.") + + # Tests that gradients are computed correctly + @_gradcheck_ops(op_db) + def test_fn_grad(self, device, dtype, op): + self._skip_helper(op, device, dtype) + self._grad_test_helper(device, dtype, op, op.get_op()) + + # Method grad (and gradgrad, see below) tests are disabled since they're + # costly and redundant with function grad (and gradgad) tests + # @_gradcheck_ops(op_db) + # def test_method_grad(self, device, dtype, op): + # self._skip_helper(op, device, dtype) + # self._grad_test_helper(device, dtype, op, op.get_method()) + + @_gradcheck_ops(op_db) + def test_inplace_grad(self, device, dtype, op): + self._skip_helper(op, device, dtype) + if not op.inplace_variant or not op.supports_inplace_autograd: + self.skipTest("Skipped! Operation does not support inplace autograd.") + self._grad_test_helper(device, dtype, op, self._get_safe_inplace(op.get_inplace())) + + # Test that gradients of gradients are computed correctly + @_gradcheck_ops(op_db) + def test_fn_gradgrad(self, device, dtype, op): + self._skip_helper(op, device, dtype) + if not op.supports_gradgrad: + self.skipTest("Skipped! Operation does not support gradgrad") + self._gradgrad_test_helper(device, dtype, op, op.get_op()) + + # Test that gradients of gradients are properly raising + @_gradcheck_ops(op_db) + def test_fn_fail_gradgrad(self, device, dtype, op): + self._skip_helper(op, device, dtype) + if op.supports_gradgrad: + self.skipTest("Skipped! Operation does support gradgrad") + + err_msg = r"derivative for .* is not implemented" + with self.assertRaisesRegex(RuntimeError, err_msg): + self._gradgrad_test_helper(device, dtype, op, op.get_op()) + + # Method gradgrad (and grad, see above) tests are disabled since they're + # costly and redundant with function gradgrad (and grad) tests + # @_gradcheck_ops(op_db) + # def test_method_gradgrad(self, device, dtype, op): + # self._skip_helper(op, device, dtype) + # self._gradgrad_test_helper(device, dtype, op, op.get_method()) + + @_gradcheck_ops(op_db) + def test_inplace_gradgrad(self, device, dtype, op): + self._skip_helper(op, device, dtype) + if not op.inplace_variant or not op.supports_inplace_autograd: + self.skipTest("Skipped! Operation does not support inplace autograd.") + self._gradgrad_test_helper(device, dtype, op, self._get_safe_inplace(op.get_inplace())) + + @_gradcheck_ops(op_db) + def test_forward_mode_AD(self, device, dtype, op): + self._skip_helper(op, device, dtype) + + if op.supports_forward_ad: + self._grad_test_helper(device, dtype, op, op.get_op(), check_forward_ad=True) + else: + err_msg = r"Trying to use forward AD with .* that does not support it\." + hint_msg = ("Running forward AD for an OP that has does not support it did not " + "raise any error. If your op supports forward AD, you should set supports_forward_ad=True") + with self.assertRaisesRegex(NotImplementedError, err_msg, msg=hint_msg): + self._grad_test_helper(device, dtype, op, op.get_op(), check_forward_ad=True) + + +# Tests operators for consistency between JIT and eager, also checks +# correctness of JIT specific alias schemas and intended +# autodifferentiation behavior. +# Inherits from JitCommonTestCase instead of TestCase directly to share +# functionality with original test_jit.py method operator tests +class TestJit(JitCommonTestCase): + exact_dtype = True + + # Tests that the forward and backward passes of operations produce the + # same values for the cross-product of op variants (function, method, inplace) + # and runtimes (eager, traced, scripted). + # TODO WARNING: inplace x {traced, scripted} not currently tested + @_variant_ops(op_db) + def test_variant_consistency_jit(self, device, dtype, op): + _requires_grad = op.supports_autograd and (dtype.is_floating_point or + op.supports_complex_autograd(torch.device(device).type)) + + include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex + samples = op.sample_inputs(device, dtype, requires_grad=_requires_grad, include_conjugated_inputs=include_conjugated_inputs) + + for sample in samples: + # Acquires variants to test + func = op.get_op() + method = op.get_method() + variants = { + # TODO: inplace tests currently fail, fix and add inplace variant + 'function': func, 'method': method, + } + + # Test traced and scripted consistency + for func_type, variant in variants.items(): + if variant is None: + continue + + # Create accessor for script function variant + name = op.name + '_' if func_type == 'inplace' else op.name + + # run with disable_autodiff_subgraph_inlining(True) to test + # autodiff support. Context manager forces the graph to contain + # DifferentiableGraph nodes if they are present + with disable_autodiff_subgraph_inlining(): + # Check scripted forward, grad, and grad grad + script_fn = create_script_fn(self, name, func_type) + + def out_fn(output): + # Processes the output for autograd + if sample.output_process_fn_grad is not None: + return sample.output_process_fn_grad(output) + return output + + check_against_reference(self, + script_fn, + func, + out_fn, + (sample.input,) + sample.args, + sample.kwargs, + no_grad=not _requires_grad, no_gradgrad=not op.supports_gradgrad) + + # Check traced forward, grad, and grad grad + traced_fn = create_traced_fn(self, variant) + check_against_reference(self, + traced_fn, + func, + out_fn, + (sample.input,) + sample.args, + sample.kwargs, + no_grad=not _requires_grad, no_gradgrad=not op.supports_gradgrad) + + # Check alias annotation schema for correctness (make + # sure inputs that aren't supposed to be modified aren't) + # Note: only runs in float32 and int64 because schema isn't affected by dtype, + # so running it on all dtypes is would be excessive + if dtype in [torch.float32, torch.int32]: + check_alias_annotation(name, (sample.input,) + sample.args, sample.kwargs, + func_type=func_type, aten_name=op.aten_name) + + # Check autodifferentiation of nodes for traced and scripted graphs, only need to check once per sample + if dtype is torch.float32: + # Sandcastle doesn't fuse nodes + if IS_SANDCASTLE: + # fusible nodes are expected to be found in FusionGroups in the DifferentiableGraphs + nonfusible_nodes = op.autodiff_nonfusible_nodes + op.autodiff_fusible_nodes + fusible_nodes = [] + else: + nonfusible_nodes = op.autodiff_nonfusible_nodes + fusible_nodes = op.autodiff_fusible_nodes + + self.assertAutodiffNode(traced_fn.last_graph, op.assert_autodiffed, nonfusible_nodes, fusible_nodes) + self.assertAutodiffNode(script_fn.last_graph, op.assert_autodiffed, nonfusible_nodes, fusible_nodes) + + # alias testing is only done with torch.float for the same reason + _alias_ops = partial(ops, dtypes=OpDTypes.supported, + allowed_dtypes=(torch.float,)) + + @_alias_ops((op for op in op_db if op.aliases)) def test_jit_alias_remapping(self, device, dtype, op): samples = op.sample_inputs(device, dtype, requires_grad=True) if len(samples) == 0: @@ -601,280 +917,10 @@ def _fn(*sample_args, **sample_kwargs): graph = traced.graph_for(*inp) FileCheck().check(op_name).check_not(variant_name).run(graph) - # Validates ops implement the correct out= behavior - # See https://github.com/pytorch/pytorch/wiki/Developer-FAQ#how-does-out-work-in-pytorch - # for a description of the correct behavior - # TODO: operations that support out= but don't support float - # are not covered by this test. - @ops(op_db, allowed_dtypes=(torch.float,)) - def test_out(self, device, dtype, op): - # TODO: verify the op doesn't support the out= kwarg - if not op.supports_out: - self.skipTest("Skipped! Op doesn't support out= kwarg.") - - # NOTE: only tests on first sample - samples = op.sample_inputs(device, dtype) - sample = samples[0] - - # calls it normally to get the expected result - expected = op(sample.input, *sample.args, **sample.kwargs) - op_out = partial(op, sample.input, *sample.args, **sample.kwargs) - - # Short-circuits if output is not a single tensor or an - # iterable of tensors - - if not isinstance(expected, torch.Tensor) and not is_iterable_of_tensors(expected, include_empty=True): - self.skipTest("Skipped! Only supports single tensor or iterable of tensor outputs.") - - # A wrapper around map that works with single tensors and always - # instantiates the map. Used below to apply transforms to - # single tensor and iterable tensor outputs. - def _apply_out_transform(fn, out): - if isinstance(out, torch.Tensor): - return fn(out) - - # assumes (see above) that out is an iterable of tensors - return tuple(map(fn, out)) - - # Case 0: out= with the correct shape, dtype, and device - # but NaN values for floating point and complex tensors, and - # maximum values for integer tensors. - # Expected behavior: out= values have no effect on the computation. - def _case_zero_transform(t): - try: - info = torch.iinfo(t.dtype) - return torch.full_like(t, info.max) - except TypeError as te: - # for non-integer types fills with NaN - return torch.full_like(t, float('nan')) - - out = _apply_out_transform(_case_zero_transform, expected) - result = op_out(out=out) - self.assertEqual(expected, out) - - # Checks that the returned value shares storage with out - # NOTE: only checks on the CPU and CUDA device types since some - # device types don't have storage - if self.device_type == 'cpu' or self.device_type == 'cuda': - if isinstance(out, torch.Tensor): - self.assertEqual(out.storage().data_ptr(), result.storage().data_ptr()) - else: - for out_t, result_t in zip(out, result): - self.assertEqual(out_t.storage().data_ptr(), result_t.storage().data_ptr()) - - # Case 1: out= with the correct shape, dtype, and device, - # but noncontiguous. - # Expected behavior: strides are respected and `out` storage is not changed. - def _case_one_transform(t): - return make_tensor(t.shape, - dtype=t.dtype, - device=t.device, - noncontiguous=True) - - # Extracts strides from a tensor or iterable of tensors into a tuple - def _extract_strides(out): - if isinstance(out, torch.Tensor): - return (out.stride(),) - - # assumes (see above) that out is an iterable of tensors - return tuple(map(lambda t: t.stride(), out)) - - def _extract_data_ptrs(out): - if isinstance(out, torch.Tensor): - return (out.data_ptr(),) - - # assumes (see above) that out is an iterable of tensors - return tuple(map(lambda t: t.data_ptr(), out)) - - - out = _apply_out_transform(_case_one_transform, expected) - original_strides = _extract_strides(out) - original_ptrs = _extract_data_ptrs(out) - - op_out(out=out) - final_strides = _extract_strides(out) - final_ptrs = _extract_data_ptrs(out) - - self.assertEqual(expected, out) - self.assertEqual(original_strides, final_strides) - self.assertEqual(original_ptrs, final_ptrs) - - # Case 2: out= with the correct dtype and device, but the wrong shape - # Expected behavior: resize with a warning. - def _case_two_transform(t): - wrong_shape = list(t.shape) - - if len(wrong_shape) == 0: - # Handles scalar tensor case (empty list) - wrong_shape = [2] - else: - wrong_shape[-1] = wrong_shape[-1] + 1 - return make_tensor(wrong_shape, dtype=t.dtype, device=t.device) - - out = _apply_out_transform(_case_two_transform, expected) - msg_fail = "Resized a non-empty tensor but did not warn about it." - with self.assertWarnsRegex(UserWarning, "An output with one or more elements", msg=msg_fail): - op_out(out=out) - self.assertEqual(expected, out) - - # Case 3: out= with the correct dtype and device, but an empty - # tensor. - # Expected behavior: resize without warning. - def _case_three_transform(t): - return make_tensor((0,), - dtype=t.dtype, - device=t.device) - - out = _apply_out_transform(_case_three_transform, expected) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - op_out(out=out) - - # Verifies no warning is a resize warning - for w in caught: - if "An output with one or more elements" in str(w.message): - self.fail("Resizing an out= argument with no elements threw a resize warning!") - - self.assertEqual(expected, out) - - # Case 4: out= with correct shape and dtype, but wrong device. - wrong_device = None - if torch.device(device).type != 'cpu': - wrong_device = 'cpu' - elif torch.cuda.is_available(): - wrong_device = 'cuda' - - if wrong_device is not None: - def _case_four_transform(t): - return make_tensor(t.shape, dtype=t.dtype, device=wrong_device) - - out = _apply_out_transform(_case_four_transform, expected) - msg_fail = f"Expected RuntimeError when calling with input.device={device} and out.device={wrong_device}" - with self.assertRaises(RuntimeError, msg=msg_fail): - op_out(out=out) - - # Case 5: out= with correct shape and device, but a dtype - # that output cannot be "safely" cast to (long). - # Expected behavior: error. - # NOTE: this case is filtered by dtype since some ops produce - # bool tensors, for example, which can be safely cast to any - # dtype. It is applied when single tensors are floating point or complex - # dtypes, or if an op returns multiple tensors when at least one such - # tensor is a floating point or complex dtype. - _dtypes = floating_and_complex_types_and(torch.float16, torch.bfloat16) - if (isinstance(expected, torch.Tensor) and expected.dtype in _dtypes or - (not isinstance(expected, torch.Tensor) and any(t.dtype in _dtypes for t in expected))): - def _case_five_transform(t): - return make_tensor(t.shape, dtype=torch.long, device=t.device) - - out = _apply_out_transform(_case_five_transform, expected) - msg_fail = "" if not isinstance(expected, torch.Tensor) else \ - ("Expected RuntimeError when doing an unsafe cast from a result of dtype " - f"{expected.dtype} into an out= with dtype torch.long") - with self.assertRaises(RuntimeError, msg=msg_fail): - op_out(out=out) - # Tests that - # 1. The operator's output for physically conjugated tensors and conjugate view tensors - # produces the same value - # 2. The gradients are same in both cases mentioned in (1) - # 3. If the operator's inplace variant is supported, tests that the inplace operation - # produces the correct value when called on a conjugate view tensor and that the output - # has its conj bit set to true - # This test only runs for C -> R and C -> C functions - # TODO: add tests for `R->C` functions - # Note: This test runs for functions that take both tensors and tensorlists as input. - @ops(op_db, allowed_dtypes=(torch.cfloat,)) - def test_conj_view(self, device, dtype, op): - if not op.test_conjugated_samples: - self.skipTest("Operation doesn't support conjugated inputs.") - _requires_grad = (op.supports_autograd and op.supports_complex_autograd(torch.device(device).type)) - samples = op.sample_inputs(device, dtype, requires_grad=_requires_grad) - inplace_variant = op.inplace_variant - - # helper function to physically conjugate the tensor - def conjugate_physical(input): - if isinstance(input, torch.Tensor): - tensor_requires_grad = input.requires_grad - with torch.no_grad(): - input = input.conj_physical() - return input.requires_grad_(tensor_requires_grad) - - if isinstance(input, Sequence): - out = list(map(clone_input_helper, input)) - out[0] = conjugate_physical(out[0]) - return tuple(out) - - # helper function to clone and conjugate the input if its a tensor - # else clone the sequence and conjugate the first element in the sequence - # If a requires_grad argument is provided the tensor being conjugated will - # have its requires_grad set to that value. - def clone_conj_input_helper(input, **kwargs): - if isinstance(input, torch.Tensor): - requires_grad = kwargs.get('requires_grad', input.requires_grad) - with torch.no_grad(): - input = input.clone() - # Note: .conj() is not called under no_grad mode since it's not allowed to modify a - # view created in no_grad mode. Here it's ok to do so, so as a workaround we call conj - # before resetting the requires_grad field for input - input = input.conj() - assert input.is_leaf - return input.requires_grad_(requires_grad) - - if isinstance(input, Sequence): - out = list(map(clone_input_helper, input)) - out[0] = clone_conj_input_helper(out[0]) - return tuple(out) - - for sample in samples: - tensor = sample.input if isinstance(sample.input, torch.Tensor) else sample.input[0] - cloned1 = clone_conj_input_helper(sample.input) - sample.input = conjugate_physical(sample.input) - - # Computes function forward value with a physically conjugated tensor and - # a conj view tensor and verifies that the output in both case are equal. - expected_forward = op(sample.input, *sample.args, **sample.kwargs) - forward_with_conjview = op(cloned1, *sample.args, **sample.kwargs) - self.assertEqual(expected_forward, forward_with_conjview) - - # If the op has an inplace variant, and the input doesn't require broadcasting - # and has the same dtype as output, verify that the inplace operation on a conjugated - # input produces correct output, and the output tensor has the conj bit set to True - if inplace_variant is not None and not sample.broadcasts_input: - cloned2 = clone_conj_input_helper(tensor, requires_grad=False) - if (isinstance(expected_forward, torch.Tensor) and - expected_forward.dtype is tensor.dtype): - inplace_forward = inplace_variant(cloned2, *sample.args, **sample.kwargs) - self.assertTrue(inplace_forward.is_conj()) - self.assertEqual(inplace_forward, expected_forward) - - # TODO: backward consistency only supported for single tensor outputs - # TODO: backward consistency only checked on sample.input, not all - # tensor inputs - # TODO: update to handle checking grads of all tensor inputs as - # derived from each tensor output - if isinstance(expected_forward, torch.Tensor) and expected_forward.requires_grad: - tensor = sample.input if isinstance(sample.input, torch.Tensor) else sample.input[0] - expected_forward.sum().backward(retain_graph=True) - forward_with_conjview.sum().backward(retain_graph=True) - if tensor.grad is not None: - cloned1_tensor = cloned1 if isinstance(cloned1, torch.Tensor) else cloned1[0] - self.assertEqual(tensor.grad, cloned1_tensor.grad) - - tensor.grad, cloned1_tensor.grad = None, None - - # a repeat of the above test if output is not complex valued - if (expected_forward.is_complex()): - grad = torch.randn_like(expected_forward) - expected_forward.backward(grad.conj_physical()) - forward_with_conjview.backward(grad.conj()) - - self.assertEqual(tensor.grad, cloned1_tensor.grad) - - -instantiate_device_type_tests(TestOpInfo, globals()) -instantiate_device_type_tests(TestGradients, globals()) instantiate_device_type_tests(TestCommon, globals()) +instantiate_device_type_tests(TestGradients, globals()) +instantiate_device_type_tests(TestJit, globals()) if __name__ == '__main__': run_tests() diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py index 73b6b59a4e518..7d814d25c645a 100644 --- a/torch/testing/_internal/common_device_type.py +++ b/torch/testing/_internal/common_device_type.py @@ -177,7 +177,7 @@ def _construct_test_name(test_name, op, device_type, dtype): test_name += "_" + device_type - if dtype is not None: + if dtype is not None and dtype is not _NO_DTYPES: if isinstance(dtype, (list, tuple)): for d in dtype: test_name += "_" + str(d).split('.')[1] @@ -186,6 +186,12 @@ def _construct_test_name(test_name, op, device_type, dtype): return test_name + +# Marker class to signify an absense of dtypes +class _NO_DTYPES(object): + pass + + class DeviceTypeTestBase(TestCase): device_type: str = 'generic_device_type' @@ -319,6 +325,8 @@ def instantiated_test(self, name=name, test=test_fn, dtype=dtype, op=op): dtypes = op.supported_dtypes(cls.device_type) elif test.opinfo_dtypes == OpDTypes.basic: dtypes = op.default_test_dtypes(cls.device_type) + elif test.opinfo_dtypes == OpDTypes.none: + dtypes = _NO_DTYPES else: raise RuntimeError(f"Unknown OpDType: {test.opinfo_dtypes}") @@ -328,12 +336,19 @@ def instantiated_test(self, name=name, test=test_fn, dtype=dtype, op=op): assert test.allowed_dtypes is None, "ops(allowed_dtypes=[...]) and the dtypes decorator are incompatible" assert test.opinfo_dtypes == OpDTypes.basic, "ops(dtypes=...) and the dtypes decorator are incompatible" - for dtype in dtypes: + if dtypes is _NO_DTYPES: instantiate_test_helper(cls, name, test=test, - dtype=dtype, + dtype=_NO_DTYPES, op=op) + else: + for dtype in dtypes: + instantiate_test_helper(cls, + name, + test=test, + dtype=dtype, + op=op) else: # Handles tests that don't use the ops decorator dtypes = cls._get_dtypes(test) @@ -569,6 +584,7 @@ class OpDTypes(Enum): unsupported = 2 # Test only unsupported dtypes supported_backward = 3 # Test all supported backward dtypes unsupported_backward = 4 # Test only unsupported backward dtypes + none = 5 # Instantiate no dtype variants (the dtype kwarg will be None) # Decorator that defines the ops a test should be run with @@ -578,6 +594,37 @@ class OpDTypes(Enum): # @ops(unary_ufuncs) # def test_numerics(self, device, dtype, op): # +# +# This will instantiate variants of test_numerics for each given operator, +# on each device that operator supports, and for every dtype supported by +# that operator. There are a few caveats to the dtype rule, explained below. +# +# First, if the OpInfo defines "default_test_dtypes" then then the test +# is instantiated for the intersection of default_test_dtypes and the +# dtypes the operator supports. Second, the @ops decorator can accept two +# additional arguments, "dtypes" and "allowed_dtypes". If "dtypes" is specified +# then the test variants are instantiated for those dtypes, regardless of +# what the operator supports. If given "allowed_dtypes" then test variants +# are instantiated only for the intersection of allowed_dtypes and the dtypes +# they would otherwise be instantiated with. That is, allowed_dtypes composes +# with the options listed above and below. +# +# The "dtypes" argument can also accept additional values (see OpDTypes above): +# OpDTypes.supported - the test is instantiated for all dtypes the operator +# supports +# OpDTypes.unsupported - the test is instantiated for all dtypes the operator +# doesn't support +# OpDTypes.supported_backward - the test is instantiated for all dtypes the +# operator's gradient formula supports +# OpDTypes.unsupported_backward - the test is instantiated for all dtypes the +# operator's gradient formula doesn't support +# OpDTypes.none - the test is instantied without any dtype. The dtype +# arg will be set to _NO_DTYPES. +# +# These options allow tests to have considerable control over the dtypes +# they're instantiated for. Finally, the @dtypes decorator composes with the +# @ops decorator, and works the same as the "dtypes" argument to @ops. + class ops(object): def __init__(self, op_list, *, dtypes: OpDTypes = OpDTypes.basic, allowed_dtypes: Optional[Sequence[torch.dtype]] = None): diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 179d6c3e51d56..686f4830966cf 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -84,7 +84,7 @@ class SampleInput(object): __slots__ = ['input', 'args', 'kwargs', 'output_process_fn_grad', 'broadcasts_input', 'name'] - def __init__(self, input, *, args=tuple(), kwargs=None, output_process_fn_grad=None, broadcasts_input=False, name=""): + def __init__(self, input, *, args=tuple(), kwargs=None, output_process_fn_grad=lambda x: x, broadcasts_input=False, name=""): # input is the first input to the op and must be either a Tensor or TensorList (Sequence[Tensor]). # This follows the typical pattern where for Tensor inputs op(t, ...) = t.op(...). # op with TensorList inputs do not support method or inplace variants. @@ -186,8 +186,32 @@ class OpInfo(object): def __init__(self, name, # the string name of the function *, + # the following metadata describes the operator, its variants, + # and its aliases, if any + aliases=None, # iterable of aliases, e.g. ("absolute",) for torch.abs + variant_test_name='', # additional string to include in the test name + # this is useful when an op needs multiple OpInfos, + # like divide does, often because it's really several + # different ops behind the scenes op=None, # the function variant of the operation, populated as torch. if None + method_variant=_NOTHING, # explicitly specifies the method variant of the operator + # if _NOTHING (default), the method variant will be autopopulated + # if None, then the OpInfo specifies no method variant + inplace_variant=_NOTHING, # explicitly specifies the inplace variant of the operator + # if _NOTHING (default), the method variant will be autopopulated + # if None, then the OpInfo specifies no method variant + + # the following metadata are test directives for skipping or + # modifying tests and a pointer to the op's sample inputs function + # this function lets the OpInfo generate valid inputs + skips=tuple(), # information about which tests to skip + decorators=None, # decorators to apply to generated tests + sample_inputs_func=None, # function to generate sample inputs + + # the following metadata relates to dtype support and is tested for correctness in test_ops.py dtypes=floating_types(), # dtypes this function is expected to work with + # the following dtypesIf... options override the dtypes value + # on their respective device types dtypesIfCPU=None, # dtypes this function is expected to work with on CPU dtypesIfCUDA=None, # dtypes this function is expected to work with on CUDA dtypesIfROCM=None, # dtypes this function is expected to work with on ROCM @@ -195,41 +219,51 @@ def __init__(self, backward_dtypesIfCPU=None, # backward dtypes this function is expected to work with on CPU backward_dtypesIfCUDA=None, # backward dtypes this function is expected to work with on CUDA backward_dtypesIfROCM=None, # backward dtypes this function is expected to work with on ROCM - default_test_dtypes=None, # dtypes to test with by default. Gets intersected - # with the dtypes support on the tested device - assert_autodiffed=False, # if a op's aten::node is expected to be symbolically autodiffed - autodiff_nonfusible_nodes=None, # a list of strings with node names that are expected to be in a - # DifferentiableGraph when autodiffed. Ex: ['aten::add', 'aten::mm'], - # default is populated to be ['aten::(name of Python operator)'] - autodiff_fusible_nodes=None, # a list of strings with node names that are expected to be in FusionGroups - # inside of DifferentiableGraphs when this operation is autodiffed. - # Ex: ['aten::add', 'aten::mm'], defaults to an empty list - # Note: currently no ops use fusible nodes + default_test_dtypes=None, # dtypes to test with by default. Tests are instantiated with + # these dtypes for the op unless otherwise specified. + # This is helpful in reducing the test matrix. + # the following metadata describes the operators out= support supports_out=True, # whether the op supports the out kwarg - skips=tuple(), # information about which tests to skip - decorators=None, # decorators to apply to generated tests + # defaults to True, if the op does not allow the out kwarg or + # supports it incorrectly then test_out in test_ops.py should fail safe_casts_outputs=False, # whether op allows safe casting when writing to out arguments - sample_inputs_func=None, # function to generate sample inputs - aten_name=None, # name of the corresponding aten:: operator - aliases=None, # iterable of aliases, e.g. ("absolute",) for torch.abs - variant_test_name='', # additional string to include in the test name - supports_autograd=True, # support for autograd - supports_gradgrad=True, # support second order gradients (this value is ignored if supports_autograd=False) + + # the following metadata relates to autograd support + supports_autograd=True, # whether the operation supports gradient computations + # if true, gradient correctness is tested in test_ops.py + # using the op's sample inputs + supports_gradgrad=True, # whether the op supports second order gradients + # if true, gradgrad correctness is tested in test_ops.py + # (this value is ignored if supports_autograd=False) supports_inplace_autograd=None, # whether the operation supports inplace autograd + # if true, tested in test_ops.py # defaults to supports_autograd's value supports_forward_ad=False, # Whether the operation support forward mode AD # If the value is True, we check that the gradients are correct # If the value is False, we test that forward grad is not implemented - supports_sparse=False, # whether the op supports sparse inputs gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs), # wrapper function for gradcheck - check_batched_grad=True, # check batched grad when doing gradcheck - check_batched_gradgrad=True, # check batched grad grad when doing gradgradcheck + check_batched_grad=True, # whether to check batched grad when doing gradcheck + check_batched_gradgrad=True, # whether to check batched grad grad when doing gradgradcheck gradcheck_nondet_tol=0.0, # tolerance for nondeterminism while performing gradcheck gradcheck_fast_mode=None, # Whether to use the fast implmentation for gradcheck/gradgradcheck. # When set to None, defers to the default value provided by the wrapper # function around gradcheck (testing._internal.common_utils.gradcheck) - inplace_variant=_NOTHING, # explicitly pass the inplace variant of the operator if required - method_variant=_NOTHING, # explicitly pass the method variant of the operator if required + + # the following metadata relates to JIT support and is tested for correctness in test_ops.py + aten_name=None, # name of the corresponding aten:: operator + assert_autodiffed=False, # if a op's aten::node is expected to be symbolically autodiffed + autodiff_nonfusible_nodes=None, # a list of strings with node names that are expected to be in a + # DifferentiableGraph when autodiffed. Ex: ['aten::add', 'aten::mm'], + # default is populated to be ['aten::(name of Python operator)'] + autodiff_fusible_nodes=None, # a list of strings with node names that are expected to be in FusionGroups + # inside of DifferentiableGraphs when this operation is autodiffed. + # Ex: ['aten::add', 'aten::mm'], defaults to an empty list + # Note: currently no ops use fusible nodes + + # the following metadata relates to sparse support and is used in test_sparse.py + supports_sparse=False, # whether the op supports sparse inputs + + # the following metadata relates to complex support and is checked in test_ops.py test_conjugated_samples=True, ): @@ -256,17 +290,29 @@ def __init__(self, "differ from CPU dtypes occasionally") self.dtypes = set(dtypes) - self.dtypesIfCPU = set(dtypesIfCPU) if dtypesIfCPU is not None else self.dtypes - self.dtypesIfCUDA = set(dtypesIfCUDA) if dtypesIfCUDA is not None else self.dtypes - self.dtypesIfROCM = set(dtypesIfROCM) if dtypesIfROCM is not None else self.dtypesIfCUDA + # NOTE: backward dtypes must be acquired before forward dtypes + # since they fallback to explicit (not implicit!) specifications of + # forward dtypes self.backward_dtypes = set(backward_dtypes) if backward_dtypes is not None else self.dtypes self.backward_dtypesIfCPU = set(backward_dtypesIfCPU) if backward_dtypesIfCPU is not None else ( - self.dtypesIfCPU if dtypesIfCPU is not None else self.backward_dtypes) + backward_dtypes if backward_dtypes is not None + else dtypesIfCPU if dtypesIfCPU is not None + else dtypes) self.backward_dtypesIfCUDA = set(backward_dtypesIfCUDA) if backward_dtypesIfCUDA is not None else ( - self.dtypesIfCUDA if dtypesIfCUDA is not None else self.backward_dtypes) + backward_dtypes if backward_dtypes is not None + else dtypesIfCUDA if dtypesIfCUDA is not None + else dtypes) self.backward_dtypesIfROCM = set(backward_dtypesIfROCM) if backward_dtypesIfROCM is not None else ( - self.dtypesIfROCM if dtypesIfROCM is not None else self.backward_dtypesIfCUDA) + backward_dtypesIfCUDA if backward_dtypesIfCUDA is not None + else backward_dtypes if backward_dtypes is not None + else dtypesIfROCM if dtypesIfROCM is not None + else dtypesIfCUDA if dtypesIfCUDA is not None + else dtypes) + + self.dtypesIfCPU = set(dtypesIfCPU) if dtypesIfCPU is not None else self.dtypes + self.dtypesIfCUDA = set(dtypesIfCUDA) if dtypesIfCUDA is not None else self.dtypes + self.dtypesIfROCM = set(dtypesIfROCM) if dtypesIfROCM is not None else self.dtypesIfCUDA self._default_test_dtypes = set(default_test_dtypes) if default_test_dtypes is not None else None @@ -409,12 +455,19 @@ def supported_dtypes(self, device_type): return self.dtypes def supported_backward_dtypes(self, device_type): + if not self.supports_autograd: + return set() + + backward_dtypes = None if device_type == 'cpu': - return self.backward_dtypesIfCPU - if device_type == 'cuda': - return self.backward_dtypesIfROCM if TEST_WITH_ROCM else self.backward_dtypesIfCUDA + backward_dtypes = self.backward_dtypesIfCPU + elif device_type == 'cuda': + backward_dtypes = self.backward_dtypesIfROCM if TEST_WITH_ROCM else self.backward_dtypesIfCUDA else: - return self.backward_dtypes + backward_dtypes = self.backward_dtypes + + allowed_backward_dtypes = floating_and_complex_types_and(torch.bfloat16, torch.float16) + return set(allowed_backward_dtypes).intersection(backward_dtypes) def supports_complex_autograd(self, device_type): if device_type == 'cpu': @@ -3471,7 +3524,7 @@ def skips_mvlgamma(skip_redundant=False): # Redundant tests skips = skips + ( # type: ignore[assignment] SkipInfo('TestGradients'), - SkipInfo('TestOpInfo'), + SkipInfo('TestJit'), SkipInfo('TestCommon'), ) return skips @@ -3491,8 +3544,8 @@ def __init__(self, variant_test_name, domain, skips, sample_kwargs): variant_test_name=variant_test_name, domain=domain, decorators=(precisionOverride({torch.float16: 5e-2}),), - dtypes=all_types_and(torch.bool), - dtypesIfCUDA=all_types_and(torch.bool, torch.half), + dtypes=all_types(), + dtypesIfCUDA=all_types_and(torch.half), sample_inputs_func=sample_inputs_mvlgamma, supports_out=False, safe_casts_outputs=True, @@ -4509,27 +4562,21 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestCommon', 'test_out', dtypes=(torch.float32,)), # Reference: https://github.com/pytorch/pytorch/issues/55589 SkipInfo('TestCommon', 'test_variant_consistency_eager'), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16, torch.complex64, torch.complex128), active_if=TEST_WITH_ROCM), ), sample_inputs_func=sample_inputs_addmv), OpInfo('addbmm', dtypes=floating_types(), dtypesIfCPU=all_types_and_complex_and(torch.float16, torch.bfloat16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if SM53OrLater else []), dtypesIfROCM=floating_types_and(torch.half), + backward_dtypesIfROCM=floating_types_and(torch.half), supports_forward_ad=True, skips=( # addbmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), # https://github.com/pytorch/pytorch/issues/55907 SkipInfo('TestCommon', 'test_variant_consistency_eager'), - SkipInfo('TestOpInfo', 'test_supported_backward', dtypes=(torch.bfloat16, ), - device_type='cuda', active_if=not SM53OrLater), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16, torch.complex64, torch.complex128), active_if=TEST_WITH_ROCM), ), sample_inputs_func=sample_inputs_addbmm), OpInfo('baddbmm', @@ -4537,15 +4584,13 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCPU=all_types_and_complex_and(torch.float16, torch.bfloat16), dtypesIfCUDA=floating_types_and(torch.float16, torch.complex64, torch.complex128, *[torch.bfloat16] if CUDA11OrLater else []), + backward_dtypesIfCUDA=floating_types_and(torch.float16, + *[torch.bfloat16] if SM53OrLater else [], + torch.complex64, torch.complex128), supports_forward_ad=True, skips=( # baddbmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), - SkipInfo('TestOpInfo', 'test_supported_backward', dtypes=(torch.bfloat16, ), - device_type='cuda', active_if=not SM53OrLater), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), ), sample_inputs_func=sample_inputs_baddbmm), OpInfo('dot', @@ -4554,49 +4599,30 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True, sample_inputs_func=sample_inputs_dot_vdot, supports_forward_ad=True, - skips=( - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), - )), + ), OpInfo('vdot', dtypes=all_types_and_complex_and(torch.float16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), sample_inputs_func=sample_inputs_dot_vdot, supports_forward_ad=True, - skips=( - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), - )), + ), OpInfo('bmm', dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if SM53OrLater else []), assert_autodiffed=True, supports_forward_ad=True, skips=( # bmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), - SkipInfo('TestOpInfo', 'test_supported_backward', dtypes=(torch.bfloat16, ), - device_type='cuda', active_if=not SM53OrLater), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), ), sample_inputs_func=sample_inputs_bmm), OpInfo('mv', - dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypes=all_types_and_complex_and(torch.bfloat16), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), skips=( # bmm does not correctly warn when resizing out= inputs - SkipInfo('TestCommon', 'test_out'), - SkipInfo('TestOpInfo', 'test_supported_backward', dtypes=(torch.float16,)), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), - # mv calls into addmv which doesn't fully support float16 - # RuntimeError: "addmv_impl_cpu" not implemented for 'Half' - SkipInfo('TestOpInfo', 'test_supported_dtypes', dtypes=(torch.float16,)),), + SkipInfo('TestCommon', 'test_out'),), assert_autodiffed=True, sample_inputs_func=sample_inputs_mv), OpInfo('addr', @@ -4610,8 +4636,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): # Reference: https://github.com/pytorch/pytorch/issues/50747 SkipInfo('TestCommon', 'test_variant_consistency_eager', dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16)), - SkipInfo('TestOpInfo', 'test_unsupported_backward', - device_type='cuda', dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), ), sample_inputs_func=sample_inputs_addr, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), @@ -4794,6 +4818,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): gradcheck_wrapper=gradcheck_wrapper_triangular_input, decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], skips=( + # TODO: FIXME: cholesky_inverse throws an error in forward when requires_grad=True + # for complex tensors + SkipInfo('TestCommon', 'test_dtypes'), # cholesky_inverse does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'),)), OpInfo('chunk', @@ -4812,7 +4839,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_forward_ad=True, skips=( # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), supports_out=False), OpInfo('symeig', @@ -4864,7 +4891,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): torch.bfloat16, torch.half), supports_forward_ad=True, skips=( - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.float32, )), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32, )), )), OpInfo('resolve_conj', dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), @@ -4882,10 +4909,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=floating_types_and(torch.half), supports_out=False, supports_forward_ad=True, - skips=( - # "sum_cpu/sum_cuda" not implemented for 'ComplexHalf' - SkipInfo('TestOpInfo', 'test_supported_backward', dtypes=(torch.half,)), - ), sample_inputs_func=sample_inputs_view_as_complex), OpInfo('complex', dtypes=floating_types(), @@ -4945,25 +4968,19 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): # CUDA illegal memory access on Windows SkipInfo(device_type='cuda', active_if=IS_WINDOWS))), OpInfo('cumsum', - dtypesIfCPU=all_types_and_complex_and(torch.bool), - dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCPU=all_types_and_complex(), + dtypesIfCUDA=all_types_and_complex_and(torch.half, torch.bfloat16), supports_forward_ad=True, skips=( - # "cumsum_out_{cpu,cuda}" not implemented for 'Bool' - SkipInfo('TestOpInfo', 'test_supported_dtypes', - dtypes=(torch.bool,)), # cumsum does not handle correctly out= dtypes SkipInfo('TestCommon', 'test_out'), ), sample_inputs_func=sample_inputs_cumulative_ops), OpInfo('cumprod', - dtypes=all_types_and_complex_and(torch.bool), - dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypes=all_types_and_complex(), + dtypesIfCUDA=all_types_and_complex_and(torch.float16, torch.bfloat16), supports_forward_ad=True, skips=( - # "cumprod_out_{cpu, cuda}" not implemented for 'Bool' - SkipInfo('TestOpInfo', 'test_supported_dtypes', - dtypes=(torch.bool,)), # cumprod does not handle correctly out= dtypes SkipInfo('TestCommon', 'test_out', dtypes=[torch.float32]), @@ -5015,7 +5032,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/issues/59174 - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), assert_autodiffed=True), OpInfo('div', @@ -5027,7 +5044,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_forward_ad=True, skips=( # Reference: https://github.com/pytorch/pytorch/issues/59174 - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), assert_autodiffed=True), OpInfo('true_divide', @@ -5048,9 +5065,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), ), assert_autodiffed=True, supports_forward_ad=True, @@ -5061,7 +5075,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): sample_inputs_func=sample_inputs_expand, skips=( # Because expand does not have a function variant. - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), supports_forward_ad=True, supports_out=False), OpInfo('expand_as', @@ -5071,7 +5085,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): sample_inputs_func=sample_inputs_expand_as, skips=( # Because expand_as does not have a function variant. - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), supports_out=False), OpInfo('diag', dtypes=all_types_and_complex_and(torch.bool), @@ -5283,10 +5297,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('floor_divide', dtypes=all_types_and(torch.half, torch.bfloat16), sample_inputs_func=sample_inputs_floor_divide, - skips=( - # `test_duplicate_method_tests` doesn't raise any warning, as it doesn't actually - # call the operator. - SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),), supports_autograd=False, ), UnaryUfuncInfo('frexp', @@ -5351,7 +5361,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # following tests give a runtime error with undefined value tensor # see discussion : https://github.com/pytorch/pytorch/issues/56660 - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.float32, torch.complex64)), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32, torch.complex64)), ), supports_inplace_autograd=False, sample_inputs_func=sample_inputs_gradient), @@ -5388,6 +5398,8 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], supports_inplace_autograd=False, skips=( + # linalg.det throwns an error when given complex inputs that require grad + SkipInfo('TestCommon', 'test_dtypes'), # The following tests fail only on ROCm. This is probably # related to the fact that the current linalg.det backward is # unstable if the matrix has repeated singular values, see @@ -5396,7 +5408,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=(torch.float64,), active_if=TEST_WITH_ROCM), SkipInfo('TestGradients', 'test_fn_gradgrad', device_type='cuda', dtypes=(torch.float64,), active_if=TEST_WITH_ROCM), - SkipInfo('TestCommon', 'test_variant_consistency_jit', device_type='cuda', + SkipInfo('TestJit', 'test_variant_consistency_jit', device_type='cuda', dtypes=(torch.float64, torch.float32), active_if=TEST_WITH_ROCM), )), OpInfo('linalg.cholesky', @@ -5475,7 +5487,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_autograd=False, decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], skips=( - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), )), OpInfo('linalg.matrix_power', aliases=('matrix_power',), @@ -5497,11 +5509,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): check_batched_gradgrad=False, sample_inputs_func=sample_inputs_linalg_multi_dot, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, - skips=( - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), - )), + ), OpInfo('linalg.norm', op=torch.linalg.norm, dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), @@ -5652,7 +5660,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # we skip jit tests because lu_backward is impelemented as autograd.Function, # which does not support autograd with scripting - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), # Skip operator schema test because this is a functional and not an operator SkipInfo('TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), )), @@ -5690,36 +5698,19 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), sample_inputs_func=sample_inputs_matrix_exp, supports_out=False, - skips=( - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), - )), + ), OpInfo('matmul', dtypes=floating_types(), dtypesIfCPU=all_types_and_complex(), dtypesIfCUDA=floating_and_complex_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else []), dtypesIfROCM=floating_types_and(torch.half, torch.bfloat16), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16), assert_autodiffed=True, sample_inputs_func=sample_inputs_matmul, skips=( # matmul does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), - # https://github.com/pytorch/pytorch/issues/55755 - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', - device_type='cpu', dtypes=(torch.float16,)), - # Backward for BFloat16 isn't supported because of the error - # "RuntimeError: CUDA error: CUBLAS_STATUS_NOT_SUPPORTED when - # calling cublasGemmStridedBatchedExFix." - SkipInfo('TestOpInfo', 'test_supported_backward', - device_type='cuda', dtypes=(torch.bfloat16,)), SkipInfo('TestCommon', 'test_conj_view', device_type='cpu'), - # "addmv_impl_cpu" not implemented for 'Half' - SkipInfo('TestOpInfo', 'test_unsupported_backward', - device_type='cpu', dtypes=(torch.float16,)), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.complex64, torch.complex128), active_if=TEST_WITH_ROCM), )), OpInfo('max', op=torch.max, @@ -5766,8 +5757,10 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, supports_forward_ad=True, skips=( + # TODO: FIXME: complex inputs requiring grad error in forward + SkipInfo('TestCommon', 'test_dtypes'), # TODO: review with var_mean tests in test_autograd.py - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), SkipInfo('TestGradients', 'test_fn_grad'), SkipInfo('TestGradients', 'test_fn_gradgrad'), SkipInfo('TestGradients', 'test_forward_mode_AD'))), @@ -5781,8 +5774,10 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, supports_forward_ad=True, skips=( + # TODO: FIXME: complex inputs requiring grad error in forward + SkipInfo('TestCommon', 'test_dtypes'), # TODO: fix along with var_mean autograd tests - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), SkipInfo('TestGradients', 'test_fn_grad'), SkipInfo('TestGradients', 'test_fn_gradgrad'), SkipInfo('TestGradients', 'test_forward_mode_AD'))), @@ -5928,9 +5923,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # mm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), - # some test samples works for ROCM backward but not all - SkipInfo('TestOpInfo', 'test_unsupported_backward', device_type='cuda', - dtypes=(torch.bfloat16,), active_if=TEST_WITH_ROCM), )), OpInfo('mode', op=torch.mode, @@ -5945,16 +5937,12 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): domain=(2, float('inf')), skips=skips_mvlgamma(skip_redundant=True) + ( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', dtypes=(torch.float16,)), - # bool can't represent the low value from the domain - SkipInfo('TestOpInfo', 'test_supported_dtypes', dtypes=(torch.bool,)), ), sample_kwargs=lambda device, dtype, input: ({'p': 3}, {'d': 3})), MvlGammaInfo(variant_test_name='mvlgamma_p_5', domain=(3, float('inf')), skips=skips_mvlgamma(skip_redundant=True) + ( SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_hard', dtypes=(torch.float16,)), - # bool can't represent the low value from the domain - SkipInfo('TestOpInfo', 'test_supported_dtypes', dtypes=(torch.bool,)), ), sample_kwargs=lambda device, dtype, input: ({'p': 5}, {'d': 5})), OpInfo('ne', @@ -6153,7 +6141,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), sample_inputs_func=sample_inputs_rbinops, supports_out=False, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit',),), assert_autodiffed=True, supports_forward_ad=True, autodiff_nonfusible_nodes=['aten::add'],), @@ -6162,7 +6150,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), sample_inputs_func=sample_inputs_rbinops, supports_out=False, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit',),), assert_autodiffed=True, autodiff_nonfusible_nodes=['aten::mul', 'aten::reciprocal'],), OpInfo('__rmul__', @@ -6170,7 +6158,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), sample_inputs_func=sample_inputs_rbinops, supports_out=False, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit',),), assert_autodiffed=True, supports_forward_ad=True, autodiff_nonfusible_nodes=['aten::mul'],), @@ -6178,21 +6166,14 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): op=torch.Tensor.__rmatmul__, dtypes=floating_types(), dtypesIfCPU=all_types_and_complex(), - dtypesIfCUDA=floating_types_and(torch.float16, torch.complex64, torch.complex128), + dtypesIfCUDA=floating_types_and(torch.float16, *[torch.bfloat16] if CUDA11OrLater else [], + torch.complex64, torch.complex128), + backward_dtypesIfCUDA=floating_types_and(torch.float16, torch.complex64, torch.complex128), assert_autodiffed=True, sample_inputs_func=sample_inputs_matmul, supports_out=False, skips=( - SkipInfo('TestCommon', 'test_variant_consistency_jit',), - # https://github.com/pytorch/pytorch/issues/55755 - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', - device_type='cpu', dtypes=(torch.float16,)), - # https://github.com/pytorch/pytorch/pull/57934#issuecomment-840091579 - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', - device_type='cuda', dtypes=(torch.bfloat16,)), - # addmv_impl_cpu" not implemented for 'Half' - SkipInfo('TestOpInfo', 'test_unsupported_backward', - dtypes=(torch.float16, torch.bfloat16)), + SkipInfo('TestJit', 'test_variant_consistency_jit',), )), OpInfo('__rmod__', op=torch.Tensor.__rmod__, @@ -6201,7 +6182,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypesIfCUDA=all_types_and(torch.bfloat16, torch.half, torch.bool), sample_inputs_func=sample_inputs_rbinops, supports_out=False, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit',),), # Support autograd after torch.remainder(Tensor, Tensor) supports # autograd of the second argument. # https://github.com/pytorch/pytorch/pull/58476/files#r637167630 @@ -6211,15 +6192,13 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('__rpow__', op=torch.Tensor.__rpow__, dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + # Reference: https://github.com/pytorch/pytorch/issues/54774 + # "log2" "_vml_cpu" not implemented for Half + backward_dtypesIfCPU=all_types_and_complex_and(torch.bfloat16, torch.bool), sample_inputs_func=sample_inputs_rbinops, supports_out=False, skips=( - # Reference: https://github.com/pytorch/pytorch/issues/54774 - # "log2" "_vml_cpu" not implemented for Half - SkipInfo('TestOpInfo', 'test_supported_backward', device_type='cpu', - dtypes=(torch.float16,)), - - SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + SkipInfo('TestJit', 'test_variant_consistency_jit',),), assert_autodiffed=True, autodiff_nonfusible_nodes=['aten::pow'],), OpInfo('__rsub__', @@ -6227,7 +6206,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), sample_inputs_func=sample_inputs_rbinops, supports_out=False, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit',),), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit',),), assert_autodiffed=True, autodiff_nonfusible_nodes=['aten::rsub'],), OpInfo('rsub', @@ -6238,7 +6217,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Reference: https://github.com/pytorch/pytorch/issues/53797 # JIT doesn't understand complex literals - SkipInfo('TestCommon', 'test_variant_consistency_jit', + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=[torch.cfloat, torch.cdouble]), ), sample_inputs_func=partial(sample_inputs_rsub, variant='tensor'),), @@ -6251,7 +6230,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Reference: https://github.com/pytorch/pytorch/issues/53797 # JIT doesn't understand complex literals - SkipInfo('TestCommon', 'test_variant_consistency_jit', + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=all_types_and_complex_and(torch.bfloat16, torch.half)),), assert_autodiffed=True,), OpInfo('select', @@ -6543,21 +6522,14 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): # TODO(@heitorschueroff) update SampleInput to handle such cases op=lambda tensors, equation: torch.einsum(equation, tensors), dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), - dtypesIfCUDA=floating_and_complex_types_and(torch.half), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, *[torch.bfloat16] if CUDA11OrLater else []), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.half), supports_out=False, sample_inputs_func=sample_inputs_einsum, skips=( # test does not work with passing lambda for op # there's a test `test_einsum` in `test_jit.py` to handle this case - SkipInfo('TestCommon', 'test_variant_consistency_jit'), - # The following dtypes are only supported for some inputs, ideally we should have - # checked this in the einsum code but to keep BC we'll just skip the tests for now. - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', - dtypes=[torch.bool]), - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', - device_type='cuda', dtypes=integral_types_and(torch.bfloat16)), - SkipInfo('TestOpInfo', 'test_unsupported_backward', - device_type='cuda', dtypes=(torch.bfloat16,)), + SkipInfo('TestJit', 'test_variant_consistency_jit'), )), OpInfo('svd', op=torch.svd, @@ -6618,7 +6590,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): # def the_method(i0): # return torch.polygamma(i0, 1) # ~~~~~~~~~~~~~~~ <--- HERE - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), sample_kwargs=lambda device, dtype, input: ({'n': 0}, {'n': 0})), UnaryUfuncInfo('polygamma', op=lambda x, n, **kwargs: torch.polygamma(n, x, **kwargs), @@ -6632,7 +6604,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Redundant tests SkipInfo('TestGradients'), - SkipInfo('TestOpInfo'), + SkipInfo('TestJit'), SkipInfo('TestCommon'), # Mismatch: https://github.com/pytorch/pytorch/issues/55357 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal'), @@ -6652,7 +6624,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Redundant tests SkipInfo('TestGradients'), - SkipInfo('TestOpInfo'), + SkipInfo('TestJit'), SkipInfo('TestCommon'), # Mismatch: https://github.com/pytorch/pytorch/issues/55357 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal'), @@ -6673,7 +6645,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Redundant tests SkipInfo('TestGradients'), - SkipInfo('TestOpInfo'), + SkipInfo('TestJit'), SkipInfo('TestCommon'), # Mismatch: https://github.com/pytorch/pytorch/issues/55357 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal'), @@ -6695,7 +6667,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Redundant tests SkipInfo('TestGradients'), - SkipInfo('TestOpInfo'), + SkipInfo('TestJit'), SkipInfo('TestCommon'), # Mismatch: https://github.com/pytorch/pytorch/issues/55357 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal'), @@ -6720,7 +6692,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): sample_inputs_func=sample_inputs_view_as_reshape_as, skips=( # Because reshape_as does not have a function variant. - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), supports_out=False, ), OpInfo('view', @@ -6729,7 +6701,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # Because view does not have a function variant. - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), sample_inputs_func=sample_inputs_view_reshape, ), OpInfo('view_as', @@ -6738,7 +6710,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # Because view_as does not have a function variant. - SkipInfo('TestCommon', 'test_variant_consistency_jit'),), + SkipInfo('TestJit', 'test_variant_consistency_jit'),), sample_inputs_func=sample_inputs_view_as_reshape_as, ), OpInfo('pinverse', @@ -6760,7 +6732,6 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): OpInfo('index_fill', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_inplace_autograd=False, - skips=(SkipInfo('TestOpInfo', 'test_duplicate_method_tests'),), supports_out=False, supports_forward_ad=True, sample_inputs_func=sample_inputs_index_fill), @@ -6788,7 +6759,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_inplace_autograd=False, op=torch.Tensor.__getitem__, sample_inputs_func=sample_inputs_getitem, - skips=(SkipInfo('TestCommon', 'test_variant_consistency_jit'),)), + skips=(SkipInfo('TestJit', 'test_variant_consistency_jit'),)), OpInfo('index_put', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), supports_out=False, @@ -6796,7 +6767,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_forward_ad=True, sample_inputs_func=sample_inputs_index_put, skips=( - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), )), OpInfo('sort', dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), @@ -6854,7 +6825,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestCommon', 'test_out'), # RuntimeError: _fn() Expected a value of type # 'Tensor (inferred)' for argument 't0' but instead found type 'tuple'. - SkipInfo('TestCommon', 'test_jit_alias_remapping'))), + SkipInfo('TestJit', 'test_jit_alias_remapping'))), OpInfo('dstack', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), sample_inputs_func=sample_inputs_hstack_dstack_vstack, @@ -6868,22 +6839,21 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): check_batched_gradgrad=False, skips=( # torch.unfold does not exist so we get a RuntimeError. - SkipInfo('TestCommon', 'test_variant_consistency_jit', + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16)), # Skip operator schema test because this is a functional and not an operator SkipInfo('TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), ), sample_inputs_func=sample_inputs_unfold), OpInfo('msort', - dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), dtypesIfROCM=all_types_and(torch.float16), check_batched_gradgrad=False, skips=( # msort does not correctly warn when resizing out= inputs. SkipInfo('TestCommon', 'test_out', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16)), - # msort does not raise expected Runtime Error. - SkipInfo('TestOpInfo', 'test_unsupported_dtypes', dtypes=[torch.bool]), ), sample_inputs_func=sample_inputs_msort), OpInfo('movedim', @@ -6894,7 +6864,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Expected a value of type 'int' for argument 'source' # but instead found type 'list'. - SkipInfo('TestCommon', 'test_jit_alias_remapping'), + SkipInfo('TestJit', 'test_jit_alias_remapping'), )), OpInfo('renorm', dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), @@ -6906,7 +6876,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # torch.repeat does not exist so we get a RuntimeError. - SkipInfo('TestCommon', 'test_variant_consistency_jit', + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16)), ), sample_inputs_func=sample_repeat_tile), @@ -6923,7 +6893,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), sample_inputs_func=sample_inputs_fill_), OpInfo('resize_', @@ -6935,7 +6905,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_autograd=False, skips=( # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), sample_inputs_func=sample_inputs_resize_ops), OpInfo('resize_as_', @@ -6947,7 +6917,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_autograd=False, skips=( # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), sample_inputs_func=sample_inputs_resize_ops), OpInfo('take_along_dim', @@ -6992,18 +6962,15 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), sample_inputs_func=sample_inputs_zero_), OpInfo('special.xlog1py', aten_name='special_xlog1py', dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + backward_dtypesIfCPU=all_types_and(torch.bool, torch.bfloat16), safe_casts_outputs=True, supports_forward_ad=True, - skips=( - SkipInfo('TestOpInfo', 'test_supported_backward', - device_type='cpu', dtypes=[torch.float16]), - ), sample_inputs_func=sample_inputs_xlog1py), OpInfo('logsumexp', dtypes=floating_types_and(torch.bfloat16), @@ -7050,7 +7017,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): skips=( # Currently failing due to an INTERNAL_ASSERT_FAILED error. # Reference: https://github.com/pytorch/pytorch/issues/56314 - SkipInfo("TestCommon", "test_variant_consistency_jit", dtypes=[torch.float32]), + SkipInfo("TestJit", "test_variant_consistency_jit", dtypes=[torch.float32]), # Skip operator schema test because this is a functional and not an operator. # Reference: https://github.com/pytorch/pytorch/issues/54574 SkipInfo('TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), @@ -7066,8 +7033,10 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): check_batched_grad=False, check_batched_gradgrad=False, skips=( + # TODO: FIXME: complex inputs requiring grad error in forward + SkipInfo('TestCommon', 'test_dtypes'), # JIT has issue when op is passed as lambda - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ) ), OpInfo('logcumsumexp', @@ -7086,6 +7055,8 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): torch.complex64: 1e-1, torch.bfloat16: 1e-2}),), skips=( + # TODO: FIXME: sigmoid fails on complex inputs that require grad + SkipInfo('TestCommon', 'test_dtypes'), # Reference: https://github.com/pytorch/pytorch/issues/56012 SkipInfo('TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda', dtypes=[torch.complex64]), @@ -7229,7 +7200,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, skips=( # test does not work with passing lambda for op - SkipInfo('TestCommon', 'test_variant_consistency_jit'), + SkipInfo('TestJit', 'test_variant_consistency_jit'), ), dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16)), # `torch.norm` has multiple code paths depending on the value of `p`. @@ -7257,11 +7228,11 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestCommon', 'test_out'), # RuntimeError: # Arguments for call are not valid. - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.complex64,)), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.complex64,)), # RuntimeError: aliasOp != torch::jit::getOperatorAliasMap().end() # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":157, # please report a bug to PyTorch. - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), ) ), OpInfo('norm', @@ -7275,15 +7246,11 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): SkipInfo('TestCommon', 'test_out'), # RuntimeError: # Arguments for call are not valid. - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.complex64,)), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.complex64,)), # RuntimeError: aliasOp != torch::jit::getOperatorAliasMap().end() # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":157, # please report a bug to PyTorch. - SkipInfo('TestCommon', 'test_variant_consistency_jit', dtypes=(torch.float32,)), - # t = torch.randn((2, 2), dtype=torch.float16) - # torch.norm(t) # Works - # torch.norm(t, 'fro', [0, 1]) # Errors - SkipInfo('TestOpInfo', 'test_unsupported_dtypes'), + SkipInfo('TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), ) ), OpInfo('norm', @@ -7293,7 +7260,7 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): backward_dtypesIfCPU=floating_and_complex_types_and(torch.float16, torch.bfloat16), skips=( # following 3 tests failed intermittenly - SkipInfo('TestCommon', 'test_variant_consistency_jit', + SkipInfo('TestJit', 'test_variant_consistency_jit', device_type='cpu', dtypes=(torch.complex64,)), SkipInfo('TestGradients', 'test_fn_grad', device_type='cpu', dtypes=(torch.complex128,)), From b8ab98626b36ed2724a44cb9e9c958e745f39152 Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Thu, 17 Jun 2021 07:55:18 -0700 Subject: [PATCH 190/305] only runs mem leak check on master (#60023) Summary: setting environment variable to only do cuda mem leak check on master CI jobs. See discussion in https://github.com/pytorch/pytorch/pull/59402#issuecomment-860773034 See stats before/after disabling mem leak check: https://github.com/pytorch/pytorch/pull/59942#issuecomment-860947095 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60023 Test Plan: https://github.com/pytorch/pytorch/issues/60108 https://github.com/pytorch/pytorch/issues/60116 Reviewed By: janeyx99 Differential Revision: D29164182 Pulled By: walterddr fbshipit-source-id: dfe88c2c1275b6eb35f18b58aacdc220f34ccb59 --- .jenkins/pytorch/test.sh | 17 ++++++++++++----- torch/testing/_internal/common_utils.py | 15 +++++++++++---- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 6610fe03aafe9..da1c766f63235 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -17,6 +17,12 @@ echo "Testing pytorch" export LANG=C.UTF-8 +# Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second +# CIRCLE_PULL_REQUEST comes from CircleCI +# NOTE: file_diff_from_base is currently bugged for GHA due to an issue finding a merge base for ghstack PRs +# see https://github.com/pytorch/pytorch/issues/60111 +IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} + if [[ "$BUILD_ENVIRONMENT" == *-slow-* ]]; then export PYTORCH_TEST_WITH_SLOW=1 export PYTORCH_TEST_SKIP_FAST=1 @@ -49,6 +55,12 @@ else export PYTORCH_TEST_SKIP_NOARCH=1 fi +if [[ -n "$IN_PULL_REQUEST" ]]; then + export PYTORCH_TEST_SKIP_CUDA_MEM_LEAK_CHECK=1 +else + export PYTORCH_TEST_SKIP_CUDA_MEM_LEAK_CHECK=0 +fi + if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then # Print GPU info rocminfo | grep -E 'Name:.*\sgfx|Marketing' @@ -122,11 +134,6 @@ elif [[ "${BUILD_ENVIRONMENT}" == *-NO_AVX2-* ]]; then export ATEN_CPU_CAPABILITY=avx fi -# Try to pull value from CIRCLE_PULL_REQUEST first then GITHUB_HEAD_REF second -# CIRCLE_PULL_REQUEST comes from CircleCI -# NOTE: file_diff_from_base is currently bugged for GHA due to an issue finding a merge base for ghstack PRs -# see https://github.com/pytorch/pytorch/issues/60111 -IN_PULL_REQUEST=${CIRCLE_PULL_REQUEST:-} if [ -n "$IN_PULL_REQUEST" ] && [[ "$BUILD_ENVIRONMENT" != *coverage* ]]; then DETERMINE_FROM=$(mktemp) file_diff_from_base "$DETERMINE_FROM" diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 506f6dfe4e8ef..6d7ca8ac5ef82 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -391,6 +391,12 @@ def _check_module_exists(name): # (unlike slow tests!) TEST_SKIP_NOARCH = os.getenv('PYTORCH_TEST_SKIP_NOARCH', '0') == '1' +# Determine whether to enable cuda memory leak check. +# CUDA mem leak check is expensive and thus we don't want to execute it on every +# test case / configuration. +# See: https://github.com/pytorch/pytorch/pull/59402#issuecomment-858811135 +TEST_SKIP_CUDA_MEM_LEAK_CHECK = os.getenv('PYTORCH_TEST_SKIP_CUDA_MEM_LEAK_CHECK', '0') == '1' + # Disables tests for when on Github Actions ON_GHA = os.getenv('GITHUB_ACTIONS', '0') == '1' @@ -995,10 +1001,11 @@ def __init__(self, method_name='runTest'): test_method = getattr(self, method_name, None) if test_method is not None: # Wraps the tested method if we should do CUDA memory check. - self._do_cuda_memory_leak_check &= getattr(test_method, '_do_cuda_memory_leak_check', True) - # FIXME: figure out the flaky -1024 anti-leaks on windows. See #8044 - if self._do_cuda_memory_leak_check and not IS_WINDOWS: - self.wrap_with_cuda_policy(method_name, self.assertLeaksNoCudaTensors) + if not TEST_SKIP_CUDA_MEM_LEAK_CHECK: + self._do_cuda_memory_leak_check &= getattr(test_method, '_do_cuda_memory_leak_check', True) + # FIXME: figure out the flaky -1024 anti-leaks on windows. See #8044 + if self._do_cuda_memory_leak_check and not IS_WINDOWS: + self.wrap_with_cuda_policy(method_name, self.assertLeaksNoCudaTensors) # Wraps the tested method if we should enforce non default CUDA stream. self._do_cuda_non_default_stream &= getattr(test_method, '_do_cuda_non_default_stream', True) From 7809494c68dd885392871e7dbc82c27ae0de3727 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Thu, 17 Jun 2021 08:15:05 -0700 Subject: [PATCH 191/305] Port `all` kernel to structured kernels. (#59371) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59371 Tracking issue: #55070 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29104399 Pulled By: ezyang fbshipit-source-id: 18bb747b7a19d873427d52c1145ef7cede333a0e --- aten/src/ATen/native/ReduceOps.cpp | 94 +++++++++++++--------- aten/src/ATen/native/ReduceOpsUtils.h | 49 ++++++++++- aten/src/ATen/native/native_functions.yaml | 4 +- 3 files changed, 105 insertions(+), 42 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index b755314b11764..a75511ebfc63c 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -29,6 +29,42 @@ #include namespace at { +namespace meta { + +TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { + dim = at::maybe_wrap_dim(dim, self.dim()); + // Refer [all, any : uint8 compatibility] + TORCH_CHECK( + self.layout() == Layout::Strided, + "all only supports strided layout, got: ", + self.layout()); + + const auto& result = maybe_get_output(); + ScalarType out_dtype; + + if (result.defined()) { + // Refer [all, any : uint8 compatibility] + TORCH_CHECK( + result.scalar_type() == ScalarType::Bool || + result.scalar_type() == ScalarType::Byte, + "all only supports bool tensor for result, got: ", + result.scalar_type()); + out_dtype = result.scalar_type(); + } else { + if (self.scalar_type() == ScalarType::Byte) { + out_dtype = self.scalar_type(); + } else { + out_dtype = ScalarType::Bool; + } + } + + auto shape = get_reduction_shape(self, dim, keepdim); + set_output(shape, self.options().dtype(out_dtype)); + namedinference::propagate_names_for_reduction(result, self, dim, keepdim); +} + +} // namespace meta + namespace native { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) @@ -1101,13 +1137,25 @@ Tensor norm(const Tensor& self, const Scalar& p) { return at::native::_norm(self, p); } +inline TensorIterator get_reduction_iter( + const Tensor& self, + const Tensor& result, + int64_t dim, + bool keepdim) { + if (self.is_cuda()) { + return meta::make_reduction(self, result, dim, keepdim, self.scalar_type()); + } + return meta::make_reduction_from_out_ty( + self, result, dim, keepdim, result.scalar_type()); +} + // Note [all, any : uint8 compatibility]: // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ // For NumPy comptability, `all` and `any` return // Tensor of dtype `bool`. However for compatibility reason, // for `uint8`, they return Tensor of same dtype `uint8`. // Reference: https://github.com/pytorch/pytorch/pull/47878#issuecomment-747108561 -inline Tensor & _all(Tensor & result, TensorIterator & iter) { +inline const Tensor & _all(const Tensor & result, TensorIterator & iter) { if (iter.numel() == 0) { result.fill_(1); } else { @@ -1148,44 +1196,12 @@ Tensor all(const Tensor& self) { return _all(result, iter); } -Tensor all(const Tensor& self, int64_t dim, bool keepdim) { - // Refer [all, any : uint8 compatibility] - Tensor result; - if (self.scalar_type() == ScalarType::Byte){ - result = at::empty({0}, self.options()); - } else { - result = at::empty({0}, self.options().dtype(kBool)); - } - - return at::native::all_out(self, dim, keepdim, result); -} - -Tensor &all_out(const Tensor &self, int64_t dim, bool keepdim, Tensor &result) { - TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), - "all only supports CPU AND CUDA device type, got: ", self.device().type()); - TORCH_CHECK(self.layout() == Layout::Strided, - "all only supports strided layout, got: ", self.layout()); - // Refer [all, any : uint8 compatibility] - TORCH_CHECK(result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, - "all only supports bool tensor for result, got: ", result.scalar_type()); - - auto out_dtype = result.scalar_type(); - dim = maybe_wrap_dim(dim, self.dim()); - if (_dimreduce_return_trivial(result, self, 1, dim, keepdim)) { - return result; - } else { - if (self.is_cuda()) { - // As CUDA supports dynamic type casting, we use this overload of - // `make_reduction`, which doesn't cast input to the result type i.e. kBool., - // otherwise we use the overload below which casts the input to kBool (which is - // an extra operation). - auto iter = make_reduction( - "all", result, self, dim, keepdim, self.scalar_type(), out_dtype); - return _all(result, iter); - } - auto iter = - make_reduction("all", result, self, dim, keepdim, /*out_dtype=*/out_dtype); - return _all(result, iter); +TORCH_IMPL_FUNC(all_out) +(const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { + auto iter = get_reduction_iter(self, result, dim, keepdim); + auto mut_result = const_cast(result); + if (!_dimreduce_return_trivial(mut_result, self, 1, dim, keepdim)) { + _all(mut_result, iter); } } diff --git a/aten/src/ATen/native/ReduceOpsUtils.h b/aten/src/ATen/native/ReduceOpsUtils.h index b443ed7ed9aa4..45c2553f34fa3 100644 --- a/aten/src/ATen/native/ReduceOpsUtils.h +++ b/aten/src/ATen/native/ReduceOpsUtils.h @@ -301,5 +301,52 @@ static void zero_numel_tensor_resize(Tensor& result, Tensor& result_indices, at::native::resize_output(result, sizes); at::native::resize_output(result_indices, sizes); } +} // native -}} // at::native +namespace meta { + +static DimVector get_reduction_shape( + const Tensor& self, + IntArrayRef dims, + bool keepdim) { + auto mask = native::make_dim_mask(dims, self.dim()); + return native::shape_from_dim_mask(self, mask, keepdim); +} + +static TensorIterator make_reduction( + const Tensor& self, + const Tensor& result, + c10::optional dim_opt, + bool keepdim, + ScalarType in_dtype) { + IntArrayRef dim = dim_opt.value_or(IntArrayRef{}); + int64_t ndim = self.dim(); + auto mask = at::native::make_dim_mask(dim, ndim); + auto viewed_result = + at::native::review_reduce_result(result, ndim, mask, keepdim); + if (self.scalar_type() == in_dtype) { + return TensorIterator::reduce_op(viewed_result, self); + } + return TensorIterator::reduce_op(viewed_result, self.to(in_dtype)); +} + +static TensorIterator make_reduction_from_out_ty( + const Tensor& self, + const Tensor& result, + c10::optional dim, + bool keepdim, + ScalarType out_dtype) { + // special case for type promotion in mixed precision, improves computational + // efficiency. + // not generalize this to common mismatched input/output types to avoid cross + // product of templated kernel launches. + const bool gpu_lowp_to_f32 = + (self.is_cuda() && + (self.scalar_type() == kHalf || self.scalar_type() == kBFloat16) && + out_dtype == kFloat); + auto in_dtype = gpu_lowp_to_f32 ? self.scalar_type() : out_dtype; + return make_reduction(self, result, dim, keepdim, in_dtype); +} + +} // namespace meta +} // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 953fdbf9b1433..786beea784843 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -468,12 +468,12 @@ - func: all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor device_check: NoCheck # TensorIterator + structured_delegate: all.out variants: function, method - dispatch: - CPU, CUDA: all - func: all.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) device_check: NoCheck # TensorIterator + structured: True dispatch: CPU, CUDA: all_out From 519698362dd23808a093480986b0a4ba0b1044a8 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Thu, 17 Jun 2021 08:15:05 -0700 Subject: [PATCH 192/305] Port `any` kernel to structured kernels. (#59372) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59372 Tracking issue: #55070 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29104395 Pulled By: ezyang fbshipit-source-id: 0cfde57c22ba88607945c98f28b18df7709becd0 --- aten/src/ATen/native/ReduceOps.cpp | 71 ++++++++-------------- aten/src/ATen/native/native_functions.yaml | 4 +- 2 files changed, 28 insertions(+), 47 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index a75511ebfc63c..8dcf7d26c2968 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -31,15 +31,20 @@ namespace at { namespace meta { -TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { - dim = at::maybe_wrap_dim(dim, self.dim()); +void check_all_any( + impl::MetaBase& meta, + const char* name, + const Tensor& self, + int64_t raw_dim, + bool keepdim) { + auto dim = at::maybe_wrap_dim(raw_dim, self.dim()); // Refer [all, any : uint8 compatibility] TORCH_CHECK( self.layout() == Layout::Strided, - "all only supports strided layout, got: ", + name, " only supports strided layout, got: ", self.layout()); - const auto& result = maybe_get_output(); + const auto& result = meta.maybe_get_output(); ScalarType out_dtype; if (result.defined()) { @@ -47,7 +52,7 @@ TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { TORCH_CHECK( result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, - "all only supports bool tensor for result, got: ", + name, " only supports bool tensor for result, got: ", result.scalar_type()); out_dtype = result.scalar_type(); } else { @@ -59,10 +64,18 @@ TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { } auto shape = get_reduction_shape(self, dim, keepdim); - set_output(shape, self.options().dtype(out_dtype)); + meta.set_output(shape, self.options().dtype(out_dtype)); namedinference::propagate_names_for_reduction(result, self, dim, keepdim); } +TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { + check_all_any(*this, "all", self, dim, keepdim); +} + +TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { + check_all_any(*this, "any", self, dim, keepdim); +} + } // namespace meta namespace native { @@ -1205,7 +1218,7 @@ TORCH_IMPL_FUNC(all_out) } } -inline Tensor & _any(Tensor & result, TensorIterator & iter) { +inline const Tensor & _any(const Tensor & result, TensorIterator & iter) { if (iter.numel() == 0) { result.fill_(0); } else { @@ -1246,44 +1259,12 @@ Tensor any(const Tensor& self) { return _any(result, iter); } -Tensor any(const Tensor& self, int64_t dim, bool keepdim) { - // Refer [all, any : uint8 compatibility] - Tensor result; - if (self.scalar_type() == ScalarType::Byte){ - result = at::empty({0}, self.options()); - } else { - result = at::empty({0}, self.options().dtype(kBool)); - } - - return at::native::any_out(self, dim, keepdim, result); -} - -Tensor &any_out(const Tensor &self, int64_t dim, bool keepdim, Tensor &result) { - TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), - "any only supports CPU AND CUDA device type, got: ", self.device().type()); - TORCH_CHECK(self.layout() == Layout::Strided, - "any only supports strided layout, got: ", self.layout()); - // Refer [all, any : uint8 compatibility] - TORCH_CHECK(result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, - "any only supports bool tensor for result, got: ", result.scalar_type()); - - auto out_dtype = result.scalar_type(); - dim = maybe_wrap_dim(dim, self.dim()); - if (_dimreduce_return_trivial(result, self, 0, dim, keepdim)) { - return result; - } else { - if (self.is_cuda()) { - // As CUDA supports dynamic type casting, we use this overload of - // `make_reduction`, which doesn't cast input to the result type i.e. kBool., - // otherwise we use the overload below which casts the input to kBool (which is - // an extra operation). - auto iter = make_reduction( - "any", result, self, dim, keepdim, self.scalar_type(), out_dtype); - return _any(result, iter); - } - auto iter = - make_reduction("any", result, self, dim, keepdim, /*out_dtype=*/out_dtype); - return _any(result, iter); +TORCH_IMPL_FUNC(any_out) +(const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { + auto iter = get_reduction_iter(self, result, dim, keepdim); + auto mut_result = const_cast(result); + if (!_dimreduce_return_trivial(mut_result, self, 0, dim, keepdim)) { + _any(mut_result, iter); } } diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 786beea784843..87cd40c5f6003 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -489,12 +489,12 @@ - func: any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor device_check: NoCheck # TensorIterator + structured_delegate: any.out variants: function, method - dispatch: - CPU, CUDA: any - func: any.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) device_check: NoCheck # TensorIterator + structured: True dispatch: CPU, CUDA: any_out From c078cefa7d90357bfb871096efd2685163181723 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Thu, 17 Jun 2021 08:15:05 -0700 Subject: [PATCH 193/305] Using meta checks for unary `torch.all` and `torch.any`. (#59373) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59373 This PR makes use of the newly implemented unified `at::meta::check_reduction` for validating the inputs and configuring its `TensorIterator`. Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29104398 Pulled By: ezyang fbshipit-source-id: 6771b80130c91c2f1360853127de0acebcfff183 --- aten/src/ATen/native/ReduceOps.cpp | 114 ++++++++++++----------------- 1 file changed, 46 insertions(+), 68 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index 8dcf7d26c2968..3de8a461c9ac4 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -31,20 +31,18 @@ namespace at { namespace meta { -void check_all_any( - impl::MetaBase& meta, +ScalarType check_allany_and_get_output_dtype( const char* name, const Tensor& self, - int64_t raw_dim, + const Tensor& result, + IntArrayRef dims, bool keepdim) { - auto dim = at::maybe_wrap_dim(raw_dim, self.dim()); // Refer [all, any : uint8 compatibility] TORCH_CHECK( self.layout() == Layout::Strided, name, " only supports strided layout, got: ", self.layout()); - const auto& result = meta.maybe_get_output(); ScalarType out_dtype; if (result.defined()) { @@ -63,17 +61,29 @@ void check_all_any( } } + return out_dtype; +} + +void check_allany_for_meta( + impl::MetaBase& meta, + const char* name, + const Tensor& self, + int64_t dim, + bool keepdim) { + dim = maybe_wrap_dim(dim, self.dim()); + const auto& result = meta.maybe_get_output(); + auto out_dtype = check_allany_and_get_output_dtype(name, self, result, dim, keepdim); auto shape = get_reduction_shape(self, dim, keepdim); meta.set_output(shape, self.options().dtype(out_dtype)); namedinference::propagate_names_for_reduction(result, self, dim, keepdim); } TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_all_any(*this, "all", self, dim, keepdim); + check_allany_for_meta(*this, "all", self, dim, keepdim); } TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_all_any(*this, "any", self, dim, keepdim); + check_allany_for_meta(*this, "any", self, dim, keepdim); } } // namespace meta @@ -1150,18 +1160,6 @@ Tensor norm(const Tensor& self, const Scalar& p) { return at::native::_norm(self, p); } -inline TensorIterator get_reduction_iter( - const Tensor& self, - const Tensor& result, - int64_t dim, - bool keepdim) { - if (self.is_cuda()) { - return meta::make_reduction(self, result, dim, keepdim, self.scalar_type()); - } - return meta::make_reduction_from_out_ty( - self, result, dim, keepdim, result.scalar_type()); -} - // Note [all, any : uint8 compatibility]: // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ // For NumPy comptability, `all` and `any` return @@ -1178,40 +1176,38 @@ inline const Tensor & _all(const Tensor & result, TensorIterator & iter) { return result; } -Tensor all(const Tensor& self) { - TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), - "all only supports CPU AND CUDA device type, got: ", self.device().type()); - TORCH_CHECK(self.layout() == Layout::Strided, - "all only supports strided layout, got: ", self.layout()); - - // Refer [all, any : uint8 compatibility] - Tensor result; - ScalarType out_dtype; - if (self.scalar_type() == ScalarType::Byte){ - result = at::empty({0}, self.options()); - out_dtype = self.scalar_type(); - } else { - result = at::empty({0}, self.options().dtype(kBool)); - out_dtype = ScalarType::Bool; - } - +inline TensorIterator get_allany_iter( + const Tensor& self, + const Tensor& result, + IntArrayRef dims, + bool keepdim) { if (self.is_cuda()) { // As CUDA supports dynamic type casting, we use this overload of // `make_reduction`, which doesn't cast input to the result type i.e. kBool., // otherwise we use the overload below which casts the input to kBool (which is // an extra operation). - auto iter = make_reduction( - "all", result, self, {}, false, self.scalar_type(), out_dtype); - return _all(result, iter); + return meta::make_reduction(self, result, dims, keepdim, self.scalar_type()); } - auto iter = - make_reduction("all", result, self, {}, false, /*out_dtype=*/out_dtype); + return meta::make_reduction_from_out_ty( + self, result, dims, keepdim, result.scalar_type()); +} + +Tensor all(const Tensor& self) { + Tensor result; + + auto out_dtype = + meta::check_allany_and_get_output_dtype("all", self, result, {}, false); + auto shape = meta::get_reduction_shape(self, {}, false); + + result = at::empty(shape, self.options().dtype(out_dtype)); + auto iter = get_allany_iter(self, result, {}, false); + return _all(result, iter); } TORCH_IMPL_FUNC(all_out) (const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_reduction_iter(self, result, dim, keepdim); + auto iter = get_allany_iter(self, result, dim, keepdim); auto mut_result = const_cast(result); if (!_dimreduce_return_trivial(mut_result, self, 1, dim, keepdim)) { _all(mut_result, iter); @@ -1229,39 +1225,21 @@ inline const Tensor & _any(const Tensor & result, TensorIterator & iter) { } Tensor any(const Tensor& self) { - TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), - "any only supports CPU AND CUDA device type, got: ", self.device().type()); - TORCH_CHECK(self.layout() == Layout::Strided || self.layout() == Layout::Sparse, - "any only supports strided AND sparse layout, got: ", self.layout()); - - // Refer [all, any : uint8 compatibility] Tensor result; - ScalarType out_dtype; - if (self.scalar_type() == ScalarType::Byte){ - result = at::empty({0}, self.options()); - out_dtype = self.scalar_type(); - } else { - result = at::empty({0}, self.options().dtype(kBool)); - out_dtype = ScalarType::Bool; - } - if (self.is_cuda()) { - // As CUDA supports dynamic type casting, we use this overload of - // `make_reduction`, which doesn't cast input to the result type i.e. kBool., - // otherwise we use the overload below which casts the input to kBool (which is - // an extra operation). - auto iter = make_reduction( - "any", result, self, {}, false, self.scalar_type(), out_dtype); - return _any(result, iter); - } - auto iter = - make_reduction("any", result, self, {}, false, /*out_dtype=*/out_dtype); + auto out_dtype = + meta::check_allany_and_get_output_dtype("any", self, result, {}, false); + auto shape = meta::get_reduction_shape(self, {}, false); + + result = at::empty(shape, self.options().dtype(out_dtype)); + auto iter = get_allany_iter(self, result, {}, false); + return _any(result, iter); } TORCH_IMPL_FUNC(any_out) (const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_reduction_iter(self, result, dim, keepdim); + auto iter = get_allany_iter(self, result, dim, keepdim); auto mut_result = const_cast(result); if (!_dimreduce_return_trivial(mut_result, self, 0, dim, keepdim)) { _any(mut_result, iter); From 6f3da4f4bf0ddecdb13b006a1bb4b7ee9cf473a4 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Thu, 17 Jun 2021 08:15:05 -0700 Subject: [PATCH 194/305] Port `argmax` to structured kernels. (#59937) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59937 Tracking issue: #55070 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29104397 Pulled By: ezyang fbshipit-source-id: 580355cf3b4e9e5c934b4e51a16196087bcb3459 --- aten/src/ATen/native/ReduceOps.cpp | 61 +++++++++++++--------- aten/src/ATen/native/native_functions.yaml | 4 +- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index 3de8a461c9ac4..a0c092297c5dc 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -86,6 +86,22 @@ TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { check_allany_for_meta(*this, "any", self, dim, keepdim); } +TORCH_META_FUNC(argmax) +(const Tensor& self, c10::optional dim, bool keepdim) { + DimVector shape; + + if (dim.has_value()) { + native::zero_numel_check_dims(self, dim.value(), "argmax()"); + shape = get_reduction_shape(self, dim.value(), keepdim); + } else { + TORCH_CHECK_INDEX( + self.numel() != 0, + "argmax(): Expected reduction dim to be specified for input.numel() == 0."); + } + + set_output(shape, self.options().dtype(kLong)); +} + } // namespace meta namespace native { @@ -1284,40 +1300,37 @@ Tensor amax(const Tensor& self, IntArrayRef dim, bool keepdim) { return at::amax_out(result, self, dim, keepdim); } -Tensor& argmax_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { +TORCH_IMPL_FUNC(argmax_out) +(const Tensor& self, + c10::optional dim, + bool keepdim, + const Tensor& result) { c10::MaybeOwned in; - if (dim) { + DimVector dims; + int64_t wrapped_dim = 0; + + if (dim.has_value()) { + wrapped_dim = maybe_wrap_dim(dim.value(), self.dim()); auto sizes = self.sizes(); - zero_numel_check_dims(self, dim.value(), "argmax()"); - auto wrap_dim = maybe_wrap_dim(dim.value(), self.dim()); - if (sizes[wrap_dim] == 1) { - if (keepdim) { - result = at::zeros(sizes, self.options().dtype(at::kLong)); - } else { - auto sizes_vec = sizes.vec(); - sizes_vec.erase(sizes_vec.begin() + wrap_dim); - result = at::zeros(sizes_vec, self.options().dtype(at::kLong)); - } - return result; + if (sizes[wrapped_dim] == 1) { + result.fill_(0); + return; } + + dims = IntArrayRef(wrapped_dim); in = c10::MaybeOwned::borrowed(self); } else { - TORCH_CHECK_INDEX(self.numel() != 0, "argmax_out(): Expected reduction dim to be specified for input.numel() == 0."); in = c10::MaybeOwned::owned(self.reshape({-1})); keepdim = false; } - auto itr = make_reduction("argmax", result, *in, dim.value_or(0), keepdim, - self.scalar_type(), at::kLong); - if (itr.numel() != 0) { - argmax_stub(itr.device_type(), itr); - } - return result; -} -Tensor argmax(const Tensor& self, c10::optional dim, bool keepdims) { - Tensor result = at::empty({0}, self.options().dtype(at::kLong)); - return at::native::argmax_out(self, dim, keepdims, result); + auto iter = + meta::make_reduction(*in, result, dims, keepdim, self.scalar_type()); + + if (iter.numel() != 0) { + argmax_stub(iter.device_type(), iter); + } } Tensor& argmin_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 87cd40c5f6003..dbfe057b426a0 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -526,12 +526,12 @@ - func: _dim_arange(Tensor like, int dim) -> Tensor - func: argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor + structured_delegate: argmax.out device_check: NoCheck # TensorIterator variants: function, method - dispatch: - CPU, CUDA: argmax - func: argmax.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True dispatch: CPU, CUDA: argmax_out From 226d745a0bf6ba174a08b92659613f4174aa393a Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Thu, 17 Jun 2021 08:15:05 -0700 Subject: [PATCH 195/305] Port `argmin` kernel to structured kernels. (#59938) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59938 Tracking issue: #55070 Test Plan: Imported from OSS Reviewed By: soulitzer Differential Revision: D29104396 Pulled By: ezyang fbshipit-source-id: 39c59bcc044649c1ec9c9685366c4dda87f76aa7 --- aten/src/ATen/native/ReduceOps.cpp | 82 ++++++++++------------ aten/src/ATen/native/native_functions.yaml | 4 +- torch/csrc/jit/runtime/static/ops.cpp | 2 +- 3 files changed, 42 insertions(+), 46 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index a0c092297c5dc..f67aa6465872d 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -86,20 +86,34 @@ TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { check_allany_for_meta(*this, "any", self, dim, keepdim); } -TORCH_META_FUNC(argmax) -(const Tensor& self, c10::optional dim, bool keepdim) { +void check_argmax_argmin( + impl::MetaBase& meta, + const char* name, + const Tensor& self, + c10::optional dim, + bool keepdim) { DimVector shape; if (dim.has_value()) { - native::zero_numel_check_dims(self, dim.value(), "argmax()"); + native::zero_numel_check_dims(self, dim.value(), name); shape = get_reduction_shape(self, dim.value(), keepdim); } else { TORCH_CHECK_INDEX( self.numel() != 0, - "argmax(): Expected reduction dim to be specified for input.numel() == 0."); + name, ": Expected reduction dim to be specified for input.numel() == 0."); } - set_output(shape, self.options().dtype(kLong)); + meta.set_output(shape, self.options().dtype(kLong)); +} + +TORCH_META_FUNC(argmax) +(const Tensor& self, c10::optional dim, bool keepdim) { + check_argmax_argmin(*this, "argmax", self, dim, keepdim); +} + +TORCH_META_FUNC(argmin) +(const Tensor& self, c10::optional dim, bool keepdim) { + check_argmax_argmin(*this, "argmin", self, dim, keepdim); } } // namespace meta @@ -1300,11 +1314,13 @@ Tensor amax(const Tensor& self, IntArrayRef dim, bool keepdim) { return at::amax_out(result, self, dim, keepdim); } -TORCH_IMPL_FUNC(argmax_out) -(const Tensor& self, - c10::optional dim, - bool keepdim, - const Tensor& result) { +template +void argmax_argmin_impl( + const Tensor& self, + c10::optional dim, + bool keepdim, + const Tensor& result, + Stub& stub) { c10::MaybeOwned in; DimVector dims; int64_t wrapped_dim = 0; @@ -1329,44 +1345,24 @@ TORCH_IMPL_FUNC(argmax_out) meta::make_reduction(*in, result, dims, keepdim, self.scalar_type()); if (iter.numel() != 0) { - argmax_stub(iter.device_type(), iter); + stub(iter.device_type(), iter); } } -Tensor& argmin_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { - c10::MaybeOwned in; - if (dim) { - auto sizes = self.sizes(); - zero_numel_check_dims(self, dim.value(), "argmin()"); - - auto wrap_dim = maybe_wrap_dim(dim.value(), self.dim()); - if (sizes[wrap_dim] == 1) { - if (keepdim) { - result = at::zeros(sizes, self.options().dtype(at::kLong)); - } else { - auto sizes_vec = sizes.vec(); - sizes_vec.erase(sizes_vec.begin() + wrap_dim); - result = at::zeros(sizes_vec, self.options().dtype(at::kLong)); - } - return result; - } - in = c10::MaybeOwned::borrowed(self); - } else { - TORCH_CHECK_INDEX(self.numel() != 0, "argmin_out(): Expected reduction dim to be specified for input.numel() == 0."); - in = c10::MaybeOwned::owned(self.reshape({-1})); - keepdim = false; - } - auto itr = make_reduction("argmin", result, *in, dim.value_or(0), keepdim, - self.scalar_type(), at::kLong); - if (itr.numel() != 0) { - argmin_stub(itr.device_type(), itr); - } - return result; +TORCH_IMPL_FUNC(argmax_out) +(const Tensor& self, + c10::optional dim, + bool keepdim, + const Tensor& result) { + argmax_argmin_impl(self, dim, keepdim, result, argmax_stub); } -Tensor argmin(const Tensor& self, c10::optional dim, bool keepdims) { - Tensor result = at::empty({0}, self.options().dtype(at::kLong)); - return at::native::argmin_out(self, dim, keepdims, result); +TORCH_IMPL_FUNC(argmin_out) +(const Tensor& self, + c10::optional dim, + bool keepdim, + const Tensor& result) { + argmax_argmin_impl(self, dim, keepdim, result, argmin_stub); } static double std_var_all_cpu(const Tensor& self, int64_t correction, bool take_sqrt) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index dbfe057b426a0..6b1d1cf7c6725 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -536,12 +536,12 @@ CPU, CUDA: argmax_out - func: argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor + structured_delegate: argmin.out device_check: NoCheck # TensorIterator variants: function, method - dispatch: - CPU, CUDA: argmin - func: argmin.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True dispatch: CPU, CUDA: argmin_out diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index 97734aa441765..908a188300f28 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -1484,7 +1484,7 @@ REGISTER_OPERATOR_FUNCTOR(aten::argmin, aten_argmin, [](Node* n) -> SROperator { } auto& out_t = p_node->Output(0).toTensor(); fastResizeToZero(out_t); - at::native::argmin_out(in0_t, dim, keepdim, out_t); + at::cpu::argmin_out(out_t, in0_t, dim, keepdim); }; }); From 010f4b6f2d37f46e48b6422e353dbfe6bfea3a1e Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Thu, 17 Jun 2021 09:01:57 -0700 Subject: [PATCH 196/305] Add .isort.cfg (#60119) Summary: This adds the `.isort.cfg` file from https://github.com/pytorch/pytorch/issues/55928, but doesn't try to enforce it in CI because as that PR showed, that is currently difficult to do. We could use this to gradually sort the codebase according to this configuration (enforcing bits and pieces in CI) but I don't do that here. The advantage of including this file (even if we don't enforce it) is that it affects how certain tools work, thus encouraging a specific import style for people who happen to use those tools. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60119 Test Plan: Open `test/run_test.py` in VS Code and run the **Python Refactor: Sort Imports** command. Compare with and without this PR. Reviewed By: 1ntEgr8 Differential Revision: D29199504 Pulled By: samestep fbshipit-source-id: 83e937b0f517c60e3e7dedb6c0306173908fbbb0 --- .isort.cfg | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .isort.cfg diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 0000000000000..d14d9bf207e6f --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,6 @@ +[settings] +include_trailing_comma=True +multi_line_output=3 +skip=third_party +skip_gitignore=True +use_parentheses=True From ed1da5be210c31cc07b033ac0f19f3dd6366feac Mon Sep 17 00:00:00 2001 From: Alexander Golynski Date: Thu, 17 Jun 2021 09:03:10 -0700 Subject: [PATCH 197/305] PG NCCL cleanup: remove usage of completed_ in WorkNCCL copies (#59899) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59899 Test Plan: Imported from OSS Reviewed By: cbalioglu, osalpekar Differential Revision: D29080299 Pulled By: agolynski fbshipit-source-id: 9ae368f91e81f19471e0a20fc913d8e9df1b9dec --- torch/lib/c10d/ProcessGroupNCCL.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp index f538e2f4ea560..2b96874ffdc6f 100644 --- a/torch/lib/c10d/ProcessGroupNCCL.cpp +++ b/torch/lib/c10d/ProcessGroupNCCL.cpp @@ -249,7 +249,6 @@ ProcessGroupNCCL::WorkNCCL::WorkNCCL(const WorkNCCL& w) blockingWait_(w.blockingWait_), opTimeout_(w.opTimeout_), workStartTime_(w.workStartTime_) { - completed_ = w.completed_; exception_ = w.exception_; } @@ -319,7 +318,6 @@ void ProcessGroupNCCL::WorkNCCL::checkAndThrowException() { void ProcessGroupNCCL::WorkNCCL::handleNCCLGuard() { std::lock_guard lock(mutex_); - completed_ = true; if (exception_) { auto exceptionMsg = c10::str( "Some NCCL operations have failed or timed out. Due to the ", From 96b3537e71ed1c5a2aa5af183c83dc6497ce6174 Mon Sep 17 00:00:00 2001 From: Bin Bao Date: Thu, 17 Jun 2021 09:33:46 -0700 Subject: [PATCH 198/305] [NNC] Add a dtypeToCppString virtual method in IRPrinter (#59449) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59449 Make dtypeToCppString as a virtual method so that a child class can easily override the dtype string generation rule. This is needed as a preparation to make loop and tensor index as int64_t. Test Plan: ``` build/bin/test_tensorexpr ``` Reviewed By: H-Huang Differential Revision: D29173969 Pulled By: desertfire fbshipit-source-id: a447badba76788354da1c79f80c834c99f105776 --- torch/csrc/jit/tensorexpr/cuda_codegen.cpp | 19 ++++++++++--------- torch/csrc/jit/tensorexpr/cuda_codegen.h | 4 ++++ torch/csrc/jit/tensorexpr/ir_printer.cpp | 12 ++++++++---- torch/csrc/jit/tensorexpr/ir_printer.h | 4 ++++ 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp index 0e8a60c144119..639e76f86be79 100644 --- a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp +++ b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp @@ -107,7 +107,7 @@ static void codegenOutputQuery( compile_to_sass = (major == prop->major) && (minor == prop->minor); } -std::string cudaDtypeCppString(const Dtype& dtype) { +std::string CudaPrinter::dtypeToCppString(const Dtype& dtype) { switch (dtype.scalar_type()) { case ScalarType::Bool: return "bool"; @@ -219,7 +219,7 @@ void CudaAnalysis::visit(const For* v) { } } -static void print_flat_alloc(std::ostream& os, const Allocate* alloc) { +void CudaPrinter::print_flat_alloc(const Allocate* alloc) { std::vector dims = alloc->dims(); // TODO: this should be merged with the storage flattener. int64_t flat_size = 1; @@ -231,8 +231,8 @@ static void print_flat_alloc(std::ostream& os, const Allocate* alloc) { throw std::runtime_error("Only IntImm dimensions are supported for now"); } } - os << cudaDtypeCppString(alloc->dtype()) << " " << (*alloc->buffer_var()) - << "[" << flat_size << "];" << std::endl; + os() << dtypeToCppString(alloc->dtype()) << " " << (*alloc->buffer_var()) + << "[" << flat_size << "];" << std::endl; } void CudaPrinter::visit(const Allocate* v) { @@ -240,13 +240,13 @@ void CudaPrinter::visit(const Allocate* v) { if (cuda_analysis_->cross_block_bufs().count(v->buffer_var()) != 0) { emitIndent(); os() << "__shared__ "; - print_flat_alloc(os(), v); + print_flat_alloc(v); return; } if (cuda_analysis_->thread_local_bufs().count(v->buffer_var()) != 0) { emitIndent(); - print_flat_alloc(os(), v); + print_flat_alloc(v); return; } @@ -274,7 +274,7 @@ void CudaPrinter::visit(const Cast* v) { return; } - os() << "(" << cudaDtypeCppString(v->dtype()) << ")"; + os() << "(" << dtypeToCppString(v->dtype()) << ")"; os() << "("; v->src_value()->accept(this); os() << ")"; @@ -522,7 +522,7 @@ void CudaPrinter::visit(const Block* v) { void CudaPrinter::visit(const Let* v) { emitIndent(); - os() << cudaDtypeCppString(v->dtype()); + os() << dtypeToCppString(v->dtype()); os() << " " << *v->var() << " = "; v->value()->accept(this); os() << ";" << std::endl; @@ -971,7 +971,8 @@ void CudaCodeGen::Initialize() { const Var* var = buffer_arg.var(); Dtype dtype = buffer_arg.dtype(); - os() << cudaDtypeCppString(dtype) << (buffer_arg.isVar() ? " " : "* ") + os() << printer_->dtypeToCppString(dtype) + << (buffer_arg.isVar() ? " " : "* ") << name_manager()->get_unique_name(var); } // NOLINTNEXTLINE(cppcoreguidelines-init-variables) diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.h b/torch/csrc/jit/tensorexpr/cuda_codegen.h index 3539c214e43ec..c2c38cd489e6a 100644 --- a/torch/csrc/jit/tensorexpr/cuda_codegen.h +++ b/torch/csrc/jit/tensorexpr/cuda_codegen.h @@ -177,12 +177,16 @@ class CudaPrinter : public IRPrinter { return rand_func_; } + std::string dtypeToCppString(const Dtype& dtype) override; + using IRPrinter::name_manager; using IRPrinter::visit; private: const Var* rand_func_; const CudaAnalysis* cuda_analysis_; + + void print_flat_alloc(const Allocate* alloc); }; // Construct Cuda C from the buffer and tensor input, and invoke the kernel diff --git a/torch/csrc/jit/tensorexpr/ir_printer.cpp b/torch/csrc/jit/tensorexpr/ir_printer.cpp index 645a2583a4125..f99ec2ef26655 100644 --- a/torch/csrc/jit/tensorexpr/ir_printer.cpp +++ b/torch/csrc/jit/tensorexpr/ir_printer.cpp @@ -10,6 +10,10 @@ namespace torch { namespace jit { namespace tensorexpr { +std::string IRPrinter::dtypeToCppString(const Dtype& dtype) { + return dtype.ToCppString(); +} + void IRPrinter::print(ExprHandle expr) { expr.node()->accept(this); } @@ -217,7 +221,7 @@ AT_FORALL_SCALAR_TYPES_AND2(Bool, Half, IMM_PRINT_VISIT); void IRPrinter::visit(const Cast* v) { auto dtype = v->dtype(); - os() << dtype.ToCppString() << "("; + os() << dtypeToCppString(dtype) << "("; v->src_value()->accept(this); os() << ")"; } @@ -383,7 +387,7 @@ void IRPrinter::visit(const Store* v) { void IRPrinter::visit(const For* v) { const Var* var = v->var(); VarHandle vv(var); - os() << "for (" << var->dtype().ToCppString() << " " << vv << " = " + os() << "for (" << dtypeToCppString(var->dtype()) << " " << vv << " = " << ExprHandle(v->start()) << "; " << vv << " < " << ExprHandle(v->stop()) << "; " << vv << "++) "; std::string loop_options_str = v->loop_options().ToString(); @@ -412,7 +416,7 @@ void IRPrinter::visit(const Block* v) { void IRPrinter::visit(const Allocate* v) { os() << "Allocate(" << *v->buffer_var() - << "); // dtype=" << v->dtype().ToCppString(); + << "); // dtype=" << dtypeToCppString(v->dtype()); os() << ", dims=["; const std::vector& dims = v->dims(); for (const auto i : c10::irange(dims.size())) { @@ -429,7 +433,7 @@ void IRPrinter::visit(const Free* v) { } void IRPrinter::visit(const Let* v) { - os() << v->dtype().ToCppString() << " " << *v->var(); + os() << dtypeToCppString(v->dtype()) << " " << *v->var(); os() << " = " << *v->value(); os() << ";"; } diff --git a/torch/csrc/jit/tensorexpr/ir_printer.h b/torch/csrc/jit/tensorexpr/ir_printer.h index 3f36777939878..a8622aeef1454 100644 --- a/torch/csrc/jit/tensorexpr/ir_printer.h +++ b/torch/csrc/jit/tensorexpr/ir_printer.h @@ -60,6 +60,10 @@ class TORCH_API IRPrinter : public IRVisitor { void visit(const Free* v) override; void visit(const Let* v) override; + // A child class may have a difference rule for generating dtype + // string, e.g. CUDA needs int64_t to be generated as long long. + virtual std::string dtypeToCppString(const Dtype& dtype); + std::ostream& os() { return printer_os_; } From 3dc8112187c5a4162581b9725695455ca959e752 Mon Sep 17 00:00:00 2001 From: Bin Bao Date: Thu, 17 Jun 2021 09:33:46 -0700 Subject: [PATCH 199/305] [NNC] Handle int64 indices and loop bounds (#59769) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59769 Allow loop bound and tensor indice to be either int32 or int64, and avoid unnecessary cast op. Test Plan: ``` build/bin/test_tensorexpr ``` Reviewed By: H-Huang Differential Revision: D29173970 Pulled By: desertfire fbshipit-source-id: 859a876ddb1b41535b2266089aa1222884295c78 --- test/cpp/tensorexpr/test_loopnest.cpp | 16 +++++++--------- torch/csrc/jit/tensorexpr/ir.cpp | 11 +++++++++-- torch/csrc/jit/tensorexpr/ir_verifier.cpp | 10 ++++++---- torch/csrc/jit/tensorexpr/tensor.h | 7 +++++-- 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index 12e87801a9d80..2d7184f849259 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -4769,15 +4769,13 @@ TEST(LoopNest, VectorizeUse) { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) const char* int64Loop = R"IR( -{ - for (int64_t n = 0; n < 12; n++) { - b[n] = (a[n]) + 1; - } -} +# CHECK: for (int64_t n = 0; n < 12; n++) { +# CHECK: b[n] = (a[n]) + 1; +# CHECK: } )IR"; // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -TEST(LoopNest, DISABLED_Int64Direct) { +TEST(LoopNest, Int64Direct) { KernelScope kernel_scope; constexpr int64_t N = 12; @@ -4788,11 +4786,11 @@ TEST(LoopNest, DISABLED_Int64Direct) { s = IRSimplifier::simplify(s); std::ostringstream oss; oss << *s; - ASSERT_EQ(oss.str(), int64Loop); + torch::jit::testing::FileCheck().run(int64Loop, oss.str()); } // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) -TEST(LoopNest, DISABLED_Int64Compute) { +TEST(LoopNest, Int64Compute) { KernelScope kernel_scope; constexpr int64_t N = 12; @@ -4805,7 +4803,7 @@ TEST(LoopNest, DISABLED_Int64Compute) { nest.simplify(); std::ostringstream oss; oss << *nest.root_stmt(); - ASSERT_EQ(oss.str(), int64Loop); + torch::jit::testing::FileCheck().run(int64Loop, oss.str()); } // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) diff --git a/torch/csrc/jit/tensorexpr/ir.cpp b/torch/csrc/jit/tensorexpr/ir.cpp index f818a43e2f363..5f40b709429b6 100644 --- a/torch/csrc/jit/tensorexpr/ir.cpp +++ b/torch/csrc/jit/tensorexpr/ir.cpp @@ -21,9 +21,16 @@ static Dtype dtypeOfIndices(const std::vector& indices) { } void castIndicesToInts(std::vector& indices) { - // Cast all indices to Int - // TODO: Should we use int64 here? + // Cast all indices to either Int or Long auto index_dtype = ScalarType::Int; + for (auto& index : indices) { + if (index->dtype().scalar_type() == ScalarType::Long) { + // If any of the indexes is Long, cast all of them to Long + index_dtype = ScalarType::Long; + break; + } + } + for (auto& index : indices) { const Dtype& dt = index->dtype(); if (c10::isIntegralType(dt.scalar_type(), true) && diff --git a/torch/csrc/jit/tensorexpr/ir_verifier.cpp b/torch/csrc/jit/tensorexpr/ir_verifier.cpp index 6e7dd4bebaa43..b3036869bdd5e 100644 --- a/torch/csrc/jit/tensorexpr/ir_verifier.cpp +++ b/torch/csrc/jit/tensorexpr/ir_verifier.cpp @@ -86,8 +86,9 @@ void IRVerifier::visit(const Load* v) { if (indices.size() > 1 && index_dtype.lanes() > 1) { throw malformed_ir("Multilane is only allowed in a flattened index"); } - if (index_dtype.scalar_type() != ScalarType::Int) { - throw malformed_ir("Index scalar dtype is not Int!"); + if (index_dtype.scalar_type() != ScalarType::Int && + index_dtype.scalar_type() != ScalarType::Long) { + throw malformed_ir("Index scalar dtype is not Int or Long!"); } IRVisitor::visit(v); @@ -129,8 +130,9 @@ void IRVerifier::visit(const Store* v) { if (indices.size() > 1 && index_dtype.lanes() > 1) { throw malformed_ir("Multilane is only allowed in a flattened index"); } - if (index_dtype.scalar_type() != ScalarType::Int) { - throw malformed_ir("Index scalar dtype is not Int!"); + if (index_dtype.scalar_type() != ScalarType::Int && + index_dtype.scalar_type() != ScalarType::Long) { + throw malformed_ir("Index scalar dtype is not Int or Long!"); } if (v->buf()->dtype() != v->value()->dtype()) { throw malformed_ir("buf and value dtype mismatch in Store"); diff --git a/torch/csrc/jit/tensorexpr/tensor.h b/torch/csrc/jit/tensorexpr/tensor.h index 95c98af0bdce5..cb91ff0c98774 100644 --- a/torch/csrc/jit/tensorexpr/tensor.h +++ b/torch/csrc/jit/tensorexpr/tensor.h @@ -161,8 +161,11 @@ inline void unpack_dim_args( dims->clear(); vars->clear(); for (const DimArg& dim_arg : dim_args) { - dims->push_back(dim_arg.dim().node()); - vars->push_back(new Var(dim_arg.name_hint(), kInt)); + const Expr* expr = dim_arg.dim().node(); + dims->push_back(expr); + vars->push_back(new Var( + dim_arg.name_hint(), + expr->dtype().scalar_type() == ScalarType::Long ? kLong : kInt)); } } From 6b5e77904f8d2477cbbff4a9c59a3479f3a0b770 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 10:24:26 -0700 Subject: [PATCH 200/305] Revert D29104396: Port `argmin` kernel to structured kernels. Test Plan: revert-hammer Differential Revision: D29104396 (https://github.com/pytorch/pytorch/commit/226d745a0bf6ba174a08b92659613f4174aa393a) Original commit changeset: 39c59bcc0446 fbshipit-source-id: 82de26f925a885f65572a785fa45a9980d3a974b --- aten/src/ATen/native/ReduceOps.cpp | 82 ++++++++++++---------- aten/src/ATen/native/native_functions.yaml | 4 +- torch/csrc/jit/runtime/static/ops.cpp | 2 +- 3 files changed, 46 insertions(+), 42 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index f67aa6465872d..a0c092297c5dc 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -86,34 +86,20 @@ TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { check_allany_for_meta(*this, "any", self, dim, keepdim); } -void check_argmax_argmin( - impl::MetaBase& meta, - const char* name, - const Tensor& self, - c10::optional dim, - bool keepdim) { +TORCH_META_FUNC(argmax) +(const Tensor& self, c10::optional dim, bool keepdim) { DimVector shape; if (dim.has_value()) { - native::zero_numel_check_dims(self, dim.value(), name); + native::zero_numel_check_dims(self, dim.value(), "argmax()"); shape = get_reduction_shape(self, dim.value(), keepdim); } else { TORCH_CHECK_INDEX( self.numel() != 0, - name, ": Expected reduction dim to be specified for input.numel() == 0."); + "argmax(): Expected reduction dim to be specified for input.numel() == 0."); } - meta.set_output(shape, self.options().dtype(kLong)); -} - -TORCH_META_FUNC(argmax) -(const Tensor& self, c10::optional dim, bool keepdim) { - check_argmax_argmin(*this, "argmax", self, dim, keepdim); -} - -TORCH_META_FUNC(argmin) -(const Tensor& self, c10::optional dim, bool keepdim) { - check_argmax_argmin(*this, "argmin", self, dim, keepdim); + set_output(shape, self.options().dtype(kLong)); } } // namespace meta @@ -1314,13 +1300,11 @@ Tensor amax(const Tensor& self, IntArrayRef dim, bool keepdim) { return at::amax_out(result, self, dim, keepdim); } -template -void argmax_argmin_impl( - const Tensor& self, - c10::optional dim, - bool keepdim, - const Tensor& result, - Stub& stub) { +TORCH_IMPL_FUNC(argmax_out) +(const Tensor& self, + c10::optional dim, + bool keepdim, + const Tensor& result) { c10::MaybeOwned in; DimVector dims; int64_t wrapped_dim = 0; @@ -1345,24 +1329,44 @@ void argmax_argmin_impl( meta::make_reduction(*in, result, dims, keepdim, self.scalar_type()); if (iter.numel() != 0) { - stub(iter.device_type(), iter); + argmax_stub(iter.device_type(), iter); } } -TORCH_IMPL_FUNC(argmax_out) -(const Tensor& self, - c10::optional dim, - bool keepdim, - const Tensor& result) { - argmax_argmin_impl(self, dim, keepdim, result, argmax_stub); +Tensor& argmin_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { + c10::MaybeOwned in; + if (dim) { + auto sizes = self.sizes(); + zero_numel_check_dims(self, dim.value(), "argmin()"); + + auto wrap_dim = maybe_wrap_dim(dim.value(), self.dim()); + if (sizes[wrap_dim] == 1) { + if (keepdim) { + result = at::zeros(sizes, self.options().dtype(at::kLong)); + } else { + auto sizes_vec = sizes.vec(); + sizes_vec.erase(sizes_vec.begin() + wrap_dim); + result = at::zeros(sizes_vec, self.options().dtype(at::kLong)); + } + return result; + } + in = c10::MaybeOwned::borrowed(self); + } else { + TORCH_CHECK_INDEX(self.numel() != 0, "argmin_out(): Expected reduction dim to be specified for input.numel() == 0."); + in = c10::MaybeOwned::owned(self.reshape({-1})); + keepdim = false; + } + auto itr = make_reduction("argmin", result, *in, dim.value_or(0), keepdim, + self.scalar_type(), at::kLong); + if (itr.numel() != 0) { + argmin_stub(itr.device_type(), itr); + } + return result; } -TORCH_IMPL_FUNC(argmin_out) -(const Tensor& self, - c10::optional dim, - bool keepdim, - const Tensor& result) { - argmax_argmin_impl(self, dim, keepdim, result, argmin_stub); +Tensor argmin(const Tensor& self, c10::optional dim, bool keepdims) { + Tensor result = at::empty({0}, self.options().dtype(at::kLong)); + return at::native::argmin_out(self, dim, keepdims, result); } static double std_var_all_cpu(const Tensor& self, int64_t correction, bool take_sqrt) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 6b1d1cf7c6725..dbfe057b426a0 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -536,12 +536,12 @@ CPU, CUDA: argmax_out - func: argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor - structured_delegate: argmin.out device_check: NoCheck # TensorIterator variants: function, method + dispatch: + CPU, CUDA: argmin - func: argmin.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) - structured: True dispatch: CPU, CUDA: argmin_out diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index 908a188300f28..97734aa441765 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -1484,7 +1484,7 @@ REGISTER_OPERATOR_FUNCTOR(aten::argmin, aten_argmin, [](Node* n) -> SROperator { } auto& out_t = p_node->Output(0).toTensor(); fastResizeToZero(out_t); - at::cpu::argmin_out(out_t, in0_t, dim, keepdim); + at::native::argmin_out(in0_t, dim, keepdim, out_t); }; }); From 873dac4b5a11ec82904a5dfc6fba6f169280e93f Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 10:24:26 -0700 Subject: [PATCH 201/305] Revert D29104397: Port `argmax` to structured kernels. Test Plan: revert-hammer Differential Revision: D29104397 (https://github.com/pytorch/pytorch/commit/6f3da4f4bf0ddecdb13b006a1bb4b7ee9cf473a4) Original commit changeset: 580355cf3b4e fbshipit-source-id: e51fb79329066bc1a6364cfa44a8732908a684ed --- aten/src/ATen/native/ReduceOps.cpp | 61 +++++++++------------- aten/src/ATen/native/native_functions.yaml | 4 +- 2 files changed, 26 insertions(+), 39 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index a0c092297c5dc..3de8a461c9ac4 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -86,22 +86,6 @@ TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { check_allany_for_meta(*this, "any", self, dim, keepdim); } -TORCH_META_FUNC(argmax) -(const Tensor& self, c10::optional dim, bool keepdim) { - DimVector shape; - - if (dim.has_value()) { - native::zero_numel_check_dims(self, dim.value(), "argmax()"); - shape = get_reduction_shape(self, dim.value(), keepdim); - } else { - TORCH_CHECK_INDEX( - self.numel() != 0, - "argmax(): Expected reduction dim to be specified for input.numel() == 0."); - } - - set_output(shape, self.options().dtype(kLong)); -} - } // namespace meta namespace native { @@ -1300,37 +1284,40 @@ Tensor amax(const Tensor& self, IntArrayRef dim, bool keepdim) { return at::amax_out(result, self, dim, keepdim); } -TORCH_IMPL_FUNC(argmax_out) -(const Tensor& self, - c10::optional dim, - bool keepdim, - const Tensor& result) { +Tensor& argmax_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { c10::MaybeOwned in; - DimVector dims; - int64_t wrapped_dim = 0; - - if (dim.has_value()) { - wrapped_dim = maybe_wrap_dim(dim.value(), self.dim()); + if (dim) { auto sizes = self.sizes(); + zero_numel_check_dims(self, dim.value(), "argmax()"); - if (sizes[wrapped_dim] == 1) { - result.fill_(0); - return; + auto wrap_dim = maybe_wrap_dim(dim.value(), self.dim()); + if (sizes[wrap_dim] == 1) { + if (keepdim) { + result = at::zeros(sizes, self.options().dtype(at::kLong)); + } else { + auto sizes_vec = sizes.vec(); + sizes_vec.erase(sizes_vec.begin() + wrap_dim); + result = at::zeros(sizes_vec, self.options().dtype(at::kLong)); + } + return result; } - - dims = IntArrayRef(wrapped_dim); in = c10::MaybeOwned::borrowed(self); } else { + TORCH_CHECK_INDEX(self.numel() != 0, "argmax_out(): Expected reduction dim to be specified for input.numel() == 0."); in = c10::MaybeOwned::owned(self.reshape({-1})); keepdim = false; } - - auto iter = - meta::make_reduction(*in, result, dims, keepdim, self.scalar_type()); - - if (iter.numel() != 0) { - argmax_stub(iter.device_type(), iter); + auto itr = make_reduction("argmax", result, *in, dim.value_or(0), keepdim, + self.scalar_type(), at::kLong); + if (itr.numel() != 0) { + argmax_stub(itr.device_type(), itr); } + return result; +} + +Tensor argmax(const Tensor& self, c10::optional dim, bool keepdims) { + Tensor result = at::empty({0}, self.options().dtype(at::kLong)); + return at::native::argmax_out(self, dim, keepdims, result); } Tensor& argmin_out(const Tensor& self, c10::optional dim, bool keepdim, Tensor& result) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index dbfe057b426a0..87cd40c5f6003 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -526,12 +526,12 @@ - func: _dim_arange(Tensor like, int dim) -> Tensor - func: argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor - structured_delegate: argmax.out device_check: NoCheck # TensorIterator variants: function, method + dispatch: + CPU, CUDA: argmax - func: argmax.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) - structured: True dispatch: CPU, CUDA: argmax_out From 81baa7fb0d346d0f87c3f1935019193a1025ac71 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 10:24:26 -0700 Subject: [PATCH 202/305] Revert D29104398: Using meta checks for unary `torch.all` and `torch.any`. Test Plan: revert-hammer Differential Revision: D29104398 (https://github.com/pytorch/pytorch/commit/c078cefa7d90357bfb871096efd2685163181723) Original commit changeset: 6771b80130c9 fbshipit-source-id: 10e5a34370113fcd2f87aea2c2e76108fa9328d8 --- aten/src/ATen/native/ReduceOps.cpp | 114 +++++++++++++++++------------ 1 file changed, 68 insertions(+), 46 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index 3de8a461c9ac4..8dcf7d26c2968 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -31,18 +31,20 @@ namespace at { namespace meta { -ScalarType check_allany_and_get_output_dtype( +void check_all_any( + impl::MetaBase& meta, const char* name, const Tensor& self, - const Tensor& result, - IntArrayRef dims, + int64_t raw_dim, bool keepdim) { + auto dim = at::maybe_wrap_dim(raw_dim, self.dim()); // Refer [all, any : uint8 compatibility] TORCH_CHECK( self.layout() == Layout::Strided, name, " only supports strided layout, got: ", self.layout()); + const auto& result = meta.maybe_get_output(); ScalarType out_dtype; if (result.defined()) { @@ -61,29 +63,17 @@ ScalarType check_allany_and_get_output_dtype( } } - return out_dtype; -} - -void check_allany_for_meta( - impl::MetaBase& meta, - const char* name, - const Tensor& self, - int64_t dim, - bool keepdim) { - dim = maybe_wrap_dim(dim, self.dim()); - const auto& result = meta.maybe_get_output(); - auto out_dtype = check_allany_and_get_output_dtype(name, self, result, dim, keepdim); auto shape = get_reduction_shape(self, dim, keepdim); meta.set_output(shape, self.options().dtype(out_dtype)); namedinference::propagate_names_for_reduction(result, self, dim, keepdim); } TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_allany_for_meta(*this, "all", self, dim, keepdim); + check_all_any(*this, "all", self, dim, keepdim); } TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_allany_for_meta(*this, "any", self, dim, keepdim); + check_all_any(*this, "any", self, dim, keepdim); } } // namespace meta @@ -1160,6 +1150,18 @@ Tensor norm(const Tensor& self, const Scalar& p) { return at::native::_norm(self, p); } +inline TensorIterator get_reduction_iter( + const Tensor& self, + const Tensor& result, + int64_t dim, + bool keepdim) { + if (self.is_cuda()) { + return meta::make_reduction(self, result, dim, keepdim, self.scalar_type()); + } + return meta::make_reduction_from_out_ty( + self, result, dim, keepdim, result.scalar_type()); +} + // Note [all, any : uint8 compatibility]: // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ // For NumPy comptability, `all` and `any` return @@ -1176,38 +1178,40 @@ inline const Tensor & _all(const Tensor & result, TensorIterator & iter) { return result; } -inline TensorIterator get_allany_iter( - const Tensor& self, - const Tensor& result, - IntArrayRef dims, - bool keepdim) { +Tensor all(const Tensor& self) { + TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), + "all only supports CPU AND CUDA device type, got: ", self.device().type()); + TORCH_CHECK(self.layout() == Layout::Strided, + "all only supports strided layout, got: ", self.layout()); + + // Refer [all, any : uint8 compatibility] + Tensor result; + ScalarType out_dtype; + if (self.scalar_type() == ScalarType::Byte){ + result = at::empty({0}, self.options()); + out_dtype = self.scalar_type(); + } else { + result = at::empty({0}, self.options().dtype(kBool)); + out_dtype = ScalarType::Bool; + } + if (self.is_cuda()) { // As CUDA supports dynamic type casting, we use this overload of // `make_reduction`, which doesn't cast input to the result type i.e. kBool., // otherwise we use the overload below which casts the input to kBool (which is // an extra operation). - return meta::make_reduction(self, result, dims, keepdim, self.scalar_type()); + auto iter = make_reduction( + "all", result, self, {}, false, self.scalar_type(), out_dtype); + return _all(result, iter); } - return meta::make_reduction_from_out_ty( - self, result, dims, keepdim, result.scalar_type()); -} - -Tensor all(const Tensor& self) { - Tensor result; - - auto out_dtype = - meta::check_allany_and_get_output_dtype("all", self, result, {}, false); - auto shape = meta::get_reduction_shape(self, {}, false); - - result = at::empty(shape, self.options().dtype(out_dtype)); - auto iter = get_allany_iter(self, result, {}, false); - + auto iter = + make_reduction("all", result, self, {}, false, /*out_dtype=*/out_dtype); return _all(result, iter); } TORCH_IMPL_FUNC(all_out) (const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_allany_iter(self, result, dim, keepdim); + auto iter = get_reduction_iter(self, result, dim, keepdim); auto mut_result = const_cast(result); if (!_dimreduce_return_trivial(mut_result, self, 1, dim, keepdim)) { _all(mut_result, iter); @@ -1225,21 +1229,39 @@ inline const Tensor & _any(const Tensor & result, TensorIterator & iter) { } Tensor any(const Tensor& self) { - Tensor result; - - auto out_dtype = - meta::check_allany_and_get_output_dtype("any", self, result, {}, false); - auto shape = meta::get_reduction_shape(self, {}, false); + TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), + "any only supports CPU AND CUDA device type, got: ", self.device().type()); + TORCH_CHECK(self.layout() == Layout::Strided || self.layout() == Layout::Sparse, + "any only supports strided AND sparse layout, got: ", self.layout()); - result = at::empty(shape, self.options().dtype(out_dtype)); - auto iter = get_allany_iter(self, result, {}, false); + // Refer [all, any : uint8 compatibility] + Tensor result; + ScalarType out_dtype; + if (self.scalar_type() == ScalarType::Byte){ + result = at::empty({0}, self.options()); + out_dtype = self.scalar_type(); + } else { + result = at::empty({0}, self.options().dtype(kBool)); + out_dtype = ScalarType::Bool; + } + if (self.is_cuda()) { + // As CUDA supports dynamic type casting, we use this overload of + // `make_reduction`, which doesn't cast input to the result type i.e. kBool., + // otherwise we use the overload below which casts the input to kBool (which is + // an extra operation). + auto iter = make_reduction( + "any", result, self, {}, false, self.scalar_type(), out_dtype); + return _any(result, iter); + } + auto iter = + make_reduction("any", result, self, {}, false, /*out_dtype=*/out_dtype); return _any(result, iter); } TORCH_IMPL_FUNC(any_out) (const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_allany_iter(self, result, dim, keepdim); + auto iter = get_reduction_iter(self, result, dim, keepdim); auto mut_result = const_cast(result); if (!_dimreduce_return_trivial(mut_result, self, 0, dim, keepdim)) { _any(mut_result, iter); From 3ff5507fb037e489487adcc6026520c3be29f3b1 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 10:24:26 -0700 Subject: [PATCH 203/305] Revert D29104395: Port `any` kernel to structured kernels. Test Plan: revert-hammer Differential Revision: D29104395 (https://github.com/pytorch/pytorch/commit/519698362dd23808a093480986b0a4ba0b1044a8) Original commit changeset: 0cfde57c22ba fbshipit-source-id: ac5ebdc4b9d3aeb4c5eeab55c92ac931599d39d1 --- aten/src/ATen/native/ReduceOps.cpp | 71 ++++++++++++++-------- aten/src/ATen/native/native_functions.yaml | 4 +- 2 files changed, 47 insertions(+), 28 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index 8dcf7d26c2968..a75511ebfc63c 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -31,20 +31,15 @@ namespace at { namespace meta { -void check_all_any( - impl::MetaBase& meta, - const char* name, - const Tensor& self, - int64_t raw_dim, - bool keepdim) { - auto dim = at::maybe_wrap_dim(raw_dim, self.dim()); +TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { + dim = at::maybe_wrap_dim(dim, self.dim()); // Refer [all, any : uint8 compatibility] TORCH_CHECK( self.layout() == Layout::Strided, - name, " only supports strided layout, got: ", + "all only supports strided layout, got: ", self.layout()); - const auto& result = meta.maybe_get_output(); + const auto& result = maybe_get_output(); ScalarType out_dtype; if (result.defined()) { @@ -52,7 +47,7 @@ void check_all_any( TORCH_CHECK( result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, - name, " only supports bool tensor for result, got: ", + "all only supports bool tensor for result, got: ", result.scalar_type()); out_dtype = result.scalar_type(); } else { @@ -64,18 +59,10 @@ void check_all_any( } auto shape = get_reduction_shape(self, dim, keepdim); - meta.set_output(shape, self.options().dtype(out_dtype)); + set_output(shape, self.options().dtype(out_dtype)); namedinference::propagate_names_for_reduction(result, self, dim, keepdim); } -TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_all_any(*this, "all", self, dim, keepdim); -} - -TORCH_META_FUNC2(any, dim)(const Tensor& self, int64_t dim, bool keepdim) { - check_all_any(*this, "any", self, dim, keepdim); -} - } // namespace meta namespace native { @@ -1218,7 +1205,7 @@ TORCH_IMPL_FUNC(all_out) } } -inline const Tensor & _any(const Tensor & result, TensorIterator & iter) { +inline Tensor & _any(Tensor & result, TensorIterator & iter) { if (iter.numel() == 0) { result.fill_(0); } else { @@ -1259,12 +1246,44 @@ Tensor any(const Tensor& self) { return _any(result, iter); } -TORCH_IMPL_FUNC(any_out) -(const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_reduction_iter(self, result, dim, keepdim); - auto mut_result = const_cast(result); - if (!_dimreduce_return_trivial(mut_result, self, 0, dim, keepdim)) { - _any(mut_result, iter); +Tensor any(const Tensor& self, int64_t dim, bool keepdim) { + // Refer [all, any : uint8 compatibility] + Tensor result; + if (self.scalar_type() == ScalarType::Byte){ + result = at::empty({0}, self.options()); + } else { + result = at::empty({0}, self.options().dtype(kBool)); + } + + return at::native::any_out(self, dim, keepdim, result); +} + +Tensor &any_out(const Tensor &self, int64_t dim, bool keepdim, Tensor &result) { + TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), + "any only supports CPU AND CUDA device type, got: ", self.device().type()); + TORCH_CHECK(self.layout() == Layout::Strided, + "any only supports strided layout, got: ", self.layout()); + // Refer [all, any : uint8 compatibility] + TORCH_CHECK(result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, + "any only supports bool tensor for result, got: ", result.scalar_type()); + + auto out_dtype = result.scalar_type(); + dim = maybe_wrap_dim(dim, self.dim()); + if (_dimreduce_return_trivial(result, self, 0, dim, keepdim)) { + return result; + } else { + if (self.is_cuda()) { + // As CUDA supports dynamic type casting, we use this overload of + // `make_reduction`, which doesn't cast input to the result type i.e. kBool., + // otherwise we use the overload below which casts the input to kBool (which is + // an extra operation). + auto iter = make_reduction( + "any", result, self, dim, keepdim, self.scalar_type(), out_dtype); + return _any(result, iter); + } + auto iter = + make_reduction("any", result, self, dim, keepdim, /*out_dtype=*/out_dtype); + return _any(result, iter); } } diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 87cd40c5f6003..786beea784843 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -489,12 +489,12 @@ - func: any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor device_check: NoCheck # TensorIterator - structured_delegate: any.out variants: function, method + dispatch: + CPU, CUDA: any - func: any.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) device_check: NoCheck # TensorIterator - structured: True dispatch: CPU, CUDA: any_out From ef09428804d9b2b580f988c723b3e4cc479d03ec Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 10:24:26 -0700 Subject: [PATCH 204/305] Revert D29104399: Port `all` kernel to structured kernels. Test Plan: revert-hammer Differential Revision: D29104399 (https://github.com/pytorch/pytorch/commit/7809494c68dd885392871e7dbc82c27ae0de3727) Original commit changeset: 18bb747b7a19 fbshipit-source-id: f57043df5646f1e675e8a555cb4fa0e436953751 --- aten/src/ATen/native/ReduceOps.cpp | 94 +++++++++------------- aten/src/ATen/native/ReduceOpsUtils.h | 49 +---------- aten/src/ATen/native/native_functions.yaml | 4 +- 3 files changed, 42 insertions(+), 105 deletions(-) diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index a75511ebfc63c..b755314b11764 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -29,42 +29,6 @@ #include namespace at { -namespace meta { - -TORCH_META_FUNC2(all, dim)(const Tensor& self, int64_t dim, bool keepdim) { - dim = at::maybe_wrap_dim(dim, self.dim()); - // Refer [all, any : uint8 compatibility] - TORCH_CHECK( - self.layout() == Layout::Strided, - "all only supports strided layout, got: ", - self.layout()); - - const auto& result = maybe_get_output(); - ScalarType out_dtype; - - if (result.defined()) { - // Refer [all, any : uint8 compatibility] - TORCH_CHECK( - result.scalar_type() == ScalarType::Bool || - result.scalar_type() == ScalarType::Byte, - "all only supports bool tensor for result, got: ", - result.scalar_type()); - out_dtype = result.scalar_type(); - } else { - if (self.scalar_type() == ScalarType::Byte) { - out_dtype = self.scalar_type(); - } else { - out_dtype = ScalarType::Bool; - } - } - - auto shape = get_reduction_shape(self, dim, keepdim); - set_output(shape, self.options().dtype(out_dtype)); - namedinference::propagate_names_for_reduction(result, self, dim, keepdim); -} - -} // namespace meta - namespace native { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) @@ -1137,25 +1101,13 @@ Tensor norm(const Tensor& self, const Scalar& p) { return at::native::_norm(self, p); } -inline TensorIterator get_reduction_iter( - const Tensor& self, - const Tensor& result, - int64_t dim, - bool keepdim) { - if (self.is_cuda()) { - return meta::make_reduction(self, result, dim, keepdim, self.scalar_type()); - } - return meta::make_reduction_from_out_ty( - self, result, dim, keepdim, result.scalar_type()); -} - // Note [all, any : uint8 compatibility]: // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ // For NumPy comptability, `all` and `any` return // Tensor of dtype `bool`. However for compatibility reason, // for `uint8`, they return Tensor of same dtype `uint8`. // Reference: https://github.com/pytorch/pytorch/pull/47878#issuecomment-747108561 -inline const Tensor & _all(const Tensor & result, TensorIterator & iter) { +inline Tensor & _all(Tensor & result, TensorIterator & iter) { if (iter.numel() == 0) { result.fill_(1); } else { @@ -1196,12 +1148,44 @@ Tensor all(const Tensor& self) { return _all(result, iter); } -TORCH_IMPL_FUNC(all_out) -(const Tensor& self, int64_t dim, bool keepdim, const Tensor& result) { - auto iter = get_reduction_iter(self, result, dim, keepdim); - auto mut_result = const_cast(result); - if (!_dimreduce_return_trivial(mut_result, self, 1, dim, keepdim)) { - _all(mut_result, iter); +Tensor all(const Tensor& self, int64_t dim, bool keepdim) { + // Refer [all, any : uint8 compatibility] + Tensor result; + if (self.scalar_type() == ScalarType::Byte){ + result = at::empty({0}, self.options()); + } else { + result = at::empty({0}, self.options().dtype(kBool)); + } + + return at::native::all_out(self, dim, keepdim, result); +} + +Tensor &all_out(const Tensor &self, int64_t dim, bool keepdim, Tensor &result) { + TORCH_CHECK(self.device().is_cpu() || self.is_cuda(), + "all only supports CPU AND CUDA device type, got: ", self.device().type()); + TORCH_CHECK(self.layout() == Layout::Strided, + "all only supports strided layout, got: ", self.layout()); + // Refer [all, any : uint8 compatibility] + TORCH_CHECK(result.scalar_type() == ScalarType::Bool || result.scalar_type() == ScalarType::Byte, + "all only supports bool tensor for result, got: ", result.scalar_type()); + + auto out_dtype = result.scalar_type(); + dim = maybe_wrap_dim(dim, self.dim()); + if (_dimreduce_return_trivial(result, self, 1, dim, keepdim)) { + return result; + } else { + if (self.is_cuda()) { + // As CUDA supports dynamic type casting, we use this overload of + // `make_reduction`, which doesn't cast input to the result type i.e. kBool., + // otherwise we use the overload below which casts the input to kBool (which is + // an extra operation). + auto iter = make_reduction( + "all", result, self, dim, keepdim, self.scalar_type(), out_dtype); + return _all(result, iter); + } + auto iter = + make_reduction("all", result, self, dim, keepdim, /*out_dtype=*/out_dtype); + return _all(result, iter); } } diff --git a/aten/src/ATen/native/ReduceOpsUtils.h b/aten/src/ATen/native/ReduceOpsUtils.h index 45c2553f34fa3..b443ed7ed9aa4 100644 --- a/aten/src/ATen/native/ReduceOpsUtils.h +++ b/aten/src/ATen/native/ReduceOpsUtils.h @@ -301,52 +301,5 @@ static void zero_numel_tensor_resize(Tensor& result, Tensor& result_indices, at::native::resize_output(result, sizes); at::native::resize_output(result_indices, sizes); } -} // native -namespace meta { - -static DimVector get_reduction_shape( - const Tensor& self, - IntArrayRef dims, - bool keepdim) { - auto mask = native::make_dim_mask(dims, self.dim()); - return native::shape_from_dim_mask(self, mask, keepdim); -} - -static TensorIterator make_reduction( - const Tensor& self, - const Tensor& result, - c10::optional dim_opt, - bool keepdim, - ScalarType in_dtype) { - IntArrayRef dim = dim_opt.value_or(IntArrayRef{}); - int64_t ndim = self.dim(); - auto mask = at::native::make_dim_mask(dim, ndim); - auto viewed_result = - at::native::review_reduce_result(result, ndim, mask, keepdim); - if (self.scalar_type() == in_dtype) { - return TensorIterator::reduce_op(viewed_result, self); - } - return TensorIterator::reduce_op(viewed_result, self.to(in_dtype)); -} - -static TensorIterator make_reduction_from_out_ty( - const Tensor& self, - const Tensor& result, - c10::optional dim, - bool keepdim, - ScalarType out_dtype) { - // special case for type promotion in mixed precision, improves computational - // efficiency. - // not generalize this to common mismatched input/output types to avoid cross - // product of templated kernel launches. - const bool gpu_lowp_to_f32 = - (self.is_cuda() && - (self.scalar_type() == kHalf || self.scalar_type() == kBFloat16) && - out_dtype == kFloat); - auto in_dtype = gpu_lowp_to_f32 ? self.scalar_type() : out_dtype; - return make_reduction(self, result, dim, keepdim, in_dtype); -} - -} // namespace meta -} // namespace at +}} // at::native diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 786beea784843..953fdbf9b1433 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -468,12 +468,12 @@ - func: all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor device_check: NoCheck # TensorIterator - structured_delegate: all.out variants: function, method + dispatch: + CPU, CUDA: all - func: all.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) device_check: NoCheck # TensorIterator - structured: True dispatch: CPU, CUDA: all_out From ebafd2aadfcf04c0918197598a063e80aa7580f7 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Thu, 17 Jun 2021 10:33:08 -0700 Subject: [PATCH 205/305] Stop warning on .names() access in max_pool2d and max_pool2d_backward (#60059) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60059 Fixes #60053. The problem is that `.names()` always triggers the named tensor warning. To not trigger it, one has to guard it with has_names: `x.has_names() ? x.names() : DimnameList{}` This is not the first time this has happened; we should probably make it so that .names() doesn't raise a warning unless it is actually populated with names. That's a little tricky to implement so I'm leaving it for the future. Test Plan: - New test, also run `python test/test_nn.py -v -k "max_pool"` and confirm there are no warnings. Reviewed By: gchanan Differential Revision: D29152737 Pulled By: zou3519 fbshipit-source-id: 89a2fdbe6a6064a7044b5b75f7d0c58e51e57509 --- aten/src/ATen/native/DilatedMaxPool2d.cpp | 12 +++++++----- test/test_namedtensor.py | 9 +++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/aten/src/ATen/native/DilatedMaxPool2d.cpp b/aten/src/ATen/native/DilatedMaxPool2d.cpp index 3b9faba377b4b..1ede82dc605b1 100644 --- a/aten/src/ATen/native/DilatedMaxPool2d.cpp +++ b/aten/src/ATen/native/DilatedMaxPool2d.cpp @@ -66,14 +66,15 @@ bool ceil_mode) { outputHeight, outputWidth, memory_format); /* resize output and indices */ + DimnameList maybe_names = input.has_names() ? input.names() : DimnameList{}; if (input.ndimension() == 3) { - set_output(0, {nInputPlane, outputHeight, outputWidth}, {}, input.options().memory_format(memory_format), input.names()); + set_output(0, {nInputPlane, outputHeight, outputWidth}, {}, input.options().memory_format(memory_format), maybe_names); /* indices will contain the locations for each output point */ - set_output(1, {nInputPlane, outputHeight, outputWidth}, {}, input.options().dtype(kLong), input.names()); + set_output(1, {nInputPlane, outputHeight, outputWidth}, {}, input.options().dtype(kLong), maybe_names); } else { - set_output(0, {nbatch, nInputPlane, outputHeight, outputWidth}, {}, input.options().memory_format(memory_format), input.names()); + set_output(0, {nbatch, nInputPlane, outputHeight, outputWidth}, {}, input.options().memory_format(memory_format), maybe_names); /* indices will contain the locations for each output point */ - set_output(1, {nbatch, nInputPlane, outputHeight, outputWidth}, {}, input.options().dtype(kLong), input.names()); + set_output(1, {nbatch, nInputPlane, outputHeight, outputWidth}, {}, input.options().dtype(kLong), maybe_names); } } @@ -148,7 +149,8 @@ const Tensor& indices) { outputHeight_for_shape_check, outputWidth_for_shape_check, memory_format); - set_output(0, input.sizes(), {}, input.options().memory_format(memory_format), input.names()); + set_output(0, input.sizes(), {}, input.options().memory_format(memory_format), + input.has_names() ? input.names() : DimnameList{}); } } // namespace meta diff --git a/test/test_namedtensor.py b/test/test_namedtensor.py index daabc1060afd4..b5e7aac402abb 100644 --- a/test/test_namedtensor.py +++ b/test/test_namedtensor.py @@ -296,6 +296,15 @@ def check_tuple_return(op, inputs, expected_names): check_tuple_return(F.max_pool2d_with_indices, [named_tensor_2d, [2, 2]], named_tensor_2d.names) check_tuple_return(F.max_pool3d_with_indices, [named_tensor_3d, [2, 2, 2]], named_tensor_3d.names) + def test_max_pooling_without_names_does_not_warn(self): + for device in torch.testing.get_all_device_types(): + tensor_2d = torch.zeros(2, 3, 5, 7, device=device, requires_grad=True) + with warnings.catch_warnings(record=True) as warns: + warnings.simplefilter("always") + result = F.max_pool2d(tensor_2d, [2, 2]) + result.sum().backward() + self.assertEqual(len(warns), 0) + def test_no_save_support(self): named_tensor = torch.zeros(2, 3, names=('N', 'C')) buf = io.BytesIO() From bbedfd913d53d677f9128caf3b8b6ea6311fe3b3 Mon Sep 17 00:00:00 2001 From: Shen Li Date: Thu, 17 Jun 2021 11:46:48 -0700 Subject: [PATCH 206/305] Run an dummy rpc._all_gather in init_rpc to avoid shutdown timeout (#59801) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59801 Fixes https://github.com/pytorch/pytorch/issues/59795. The RPC calls in shutdown no longer able to finish within 5s if there is no other RPCs before `rpc.shutdown()` in that process, because agent initialization can take longer than 5s. We don't have this problem previously, because TensorPipe's backend registry used to use RPC to communicate CUDA devices in `init_rpc`. However, after #58753, `init_rpc` uses ProcessGroup to communicate devices, and hence the channels/transport could be uninitialized after `init_rpc`. Differential Revision: D29039238 D29039238 Test Plan: Imported from OSS Reviewed By: rohan-varma Pulled By: mrshenli fbshipit-source-id: 46f89b01a058a51d271ddef9084a67b220a067b7 --- torch/distributed/rpc/backend_registry.py | 9 +++++++++ torch/testing/_internal/distributed/rpc/rpc_test.py | 5 +++++ 2 files changed, 14 insertions(+) diff --git a/torch/distributed/rpc/backend_registry.py b/torch/distributed/rpc/backend_registry.py index 9dc4f8fa0ef9f..6b278f1b4b6c6 100644 --- a/torch/distributed/rpc/backend_registry.py +++ b/torch/distributed/rpc/backend_registry.py @@ -353,6 +353,15 @@ def _tensorpipe_init_backend_handler(store, name, rank, world_size, rpc_backend_ api._init_rpc_states(agent) + # Run one dummy round of RPC to initialize channels/transports. Without + # this, it's easy to hit timeout in rpc.shutdown() if there is no other RPC + # on that process before rpc.shutdown(), as the agent initialization can + # take longer than 5s. + api._all_gather(None, timeout=rpc_constants.DEFAULT_RPC_TIMEOUT_SEC) + # Need a barrier here to make sure no peers leave before the rank0 finishes + # _all_gather + group.barrier().wait() + return agent diff --git a/torch/testing/_internal/distributed/rpc/rpc_test.py b/torch/testing/_internal/distributed/rpc/rpc_test.py index e91b366eea970..7b54149abd8bc 100644 --- a/torch/testing/_internal/distributed/rpc/rpc_test.py +++ b/torch/testing/_internal/distributed/rpc/rpc_test.py @@ -4921,6 +4921,11 @@ def forward(self, x, is_rref=False): torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) return self.net(x) + def __getstate__(self): + # return an empty dict to avoid inspecting the model contents on the + # owner + return {} + class TensorPipeAgentCudaRpcTest(RpcAgentTestFixture): From 462448f07ab9f2f2909e062185832e33843431fa Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Thu, 17 Jun 2021 12:59:01 -0700 Subject: [PATCH 207/305] Enable GHA sharding on linux (#60124) Summary: This is branch off of https://github.com/pytorch/pytorch/issues/59970 to only shard on linux so far (we're running in issues with windows gflags). This would enable sharding of tests on a few Linux jobs on GHA, allowing tts to be essentially halved. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60124 Reviewed By: zou3519 Differential Revision: D29204211 Pulled By: janeyx99 fbshipit-source-id: 1cc31d1eccd564d96e2aef14c0acae96a3f0fcd0 --- .github/scripts/generate_ci_workflows.py | 9 ++++- .../scripts/generate_pytorch_test_matrix.py | 31 +++++++++++++++ .github/templates/linux_ci_workflow.yml.j2 | 38 ++++++++++++++++++- .github/templates/windows_ci_workflow.yml.j2 | 34 ++++++++++++++++- ...inux-xenial-cuda10.2-cudnn7-py3.6-gcc7.yml | 38 ++++++++++++++++++- .../pytorch-linux-xenial-py3.6-gcc5.4.yml | 38 ++++++++++++++++++- .../workflows/pytorch-win-vs2019-cpu-py3.yml | 34 ++++++++++++++++- .../pytorch-win-vs2019-cuda10-cudnn7-py3.yml | 34 ++++++++++++++++- .../pytorch-win-vs2019-cuda11-cudnn8-py3.yml | 34 ++++++++++++++++- .jenkins/pytorch/test.sh | 4 +- .jenkins/pytorch/win-test.sh | 6 +-- test/run_test.py | 6 +-- 12 files changed, 289 insertions(+), 17 deletions(-) create mode 100755 .github/scripts/generate_pytorch_test_matrix.py diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index aae7ef2685fec..d0e1293acb793 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -24,13 +24,15 @@ def PyTorchWindowsWorkflow( build_environment: str, test_runner_type: str, cuda_version: str, - on_pull_request: bool = False + on_pull_request: bool = False, + num_test_shards: int = 1, ) -> PyTorchWorkflow: return { "build_environment": build_environment, "test_runner_type": test_runner_type, "cuda_version": cuda_version, "on_pull_request": on_pull_request, + "num_test_shards": num_test_shards, } @@ -45,6 +47,7 @@ def PyTorchLinuxWorkflow( test_runner_type: str, on_pull_request: bool = False, enable_doc_jobs: bool = False, + num_test_shards: int = 1, ) -> PyTorchWorkflow: return { "build_environment": build_environment, @@ -52,6 +55,7 @@ def PyTorchLinuxWorkflow( "test_runner_type": test_runner_type, "on_pull_request": on_pull_request, "enable_doc_jobs": enable_doc_jobs, + "num_test_shards": num_test_shards, } @@ -74,7 +78,7 @@ def generate_workflow_file( build_environment="pytorch-win-vs2019-cpu-py3", cuda_version="cpu", test_runner_type=WINDOWS_CPU_TEST_RUNNER, - on_pull_request=True + on_pull_request=True, ), PyTorchWindowsWorkflow( build_environment="pytorch-win-vs2019-cuda10-cudnn7-py3", @@ -95,6 +99,7 @@ def generate_workflow_file( test_runner_type=LINUX_CPU_TEST_RUNNER, on_pull_request=True, enable_doc_jobs=True, + num_test_shards=2, ), # PyTorchLinuxWorkflow( # build_environment="pytorch-paralleltbb-linux-xenial-py3.6-gcc5.4", diff --git a/.github/scripts/generate_pytorch_test_matrix.py b/.github/scripts/generate_pytorch_test_matrix.py new file mode 100755 index 0000000000000..7ab2099b8e727 --- /dev/null +++ b/.github/scripts/generate_pytorch_test_matrix.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 + +"""Generates a matrix to be utilized through github actions + +Will output a matrix to represent our testing configurations, which is currently +dictated by just sharding. + +""" + +import json +import os +from typing import List + + +NUM_TEST_SHARDS = int(os.getenv('NUM_TEST_SHARDS', '1')) + +def generate_sharding_list() -> List[int]: + return list(range(1, NUM_TEST_SHARDS + 1)) + + +def main() -> None: + print(json.dumps( + { + 'test_config': generate_sharding_list() + }, + sort_keys=True, + )) + + +if __name__ == "__main__": + main() diff --git a/.github/templates/linux_ci_workflow.yml.j2 b/.github/templates/linux_ci_workflow.yml.j2 index c28a2c63072a8..d9262feba32b1 100644 --- a/.github/templates/linux_ci_workflow.yml.j2 +++ b/.github/templates/linux_ci_workflow.yml.j2 @@ -102,6 +102,7 @@ jobs: needs: calculate-docker-image env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: !{{ build_environment }}-build steps: - name: Log in to ECR run: | @@ -127,6 +128,7 @@ jobs: run: | docker run \ -e BUILD_ENVIRONMENT \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ -e SCCACHE_BUCKET \ -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ @@ -163,7 +165,7 @@ jobs: docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . - name: Archive artifacts into zip run: | - zip -r artifacts.zip dist/ build/ + zip -r artifacts.zip dist/ build/ .pytorch-test-times.json # Upload to github so that people can click and download artifacts - uses: actions/upload-artifact@v2 # Don't fail on upload to GH since it's only for user convenience @@ -189,13 +191,40 @@ jobs: # Prune all of the docker images docker system prune -af + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: !{{ num_test_shards }} + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: !{{ test_runner_type }} needs: - calculate-docker-image - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: !{{ build_environment }}-test + NUM_TEST_SHARDS: !{{ num_test_shards }} + TEST_CONFIG: ${{ matrix.test_config }} steps: - name: Log in to ECR run: | @@ -245,6 +274,11 @@ jobs: env | grep '^GITHUB' > "/tmp/github_env_${GITHUB_RUN_ID}" - name: Test PyTorch run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi # TODO: Stop building test binaries as part of the build phase # Used for GPU_FLAG since that doesn't play nice # shellcheck disable=SC2086 @@ -254,6 +288,8 @@ jobs: -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ -e GITHUB_ACTIONS \ -e IN_CI \ + -e SHARD_NUMBER \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ --security-opt seccomp=unconfined \ diff --git a/.github/templates/windows_ci_workflow.yml.j2 b/.github/templates/windows_ci_workflow.yml.j2 index 65d6949f9cd84..2d09f9bb511fe 100644 --- a/.github/templates/windows_ci_workflow.yml.j2 +++ b/.github/templates/windows_ci_workflow.yml.j2 @@ -18,7 +18,6 @@ env: CUDA_VERSION: "!{{ cuda_version }}" IN_CI: 1 INSTALL_WINDOWS_SDK: 1 - JOB_BASE_NAME: test PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" @@ -36,6 +35,8 @@ concurrency: jobs: build: runs-on: "windows.4xlarge" + env: + JOB_BASE_NAME: !{{ build_environment }}-build steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -83,12 +84,38 @@ jobs: name: ${{ env.BUILD_ENVIRONMENT }} path: C:\w\build-results + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: !{{ num_test_shards }} + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: !{{ test_runner_type }} env: JOB_BASE_NAME: !{{ build_environment }}-test + NUM_TEST_SHARDS: !{{ num_test_shards }} + TEST_CONFIG: ${{ matrix.test_config }} needs: - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -131,6 +158,11 @@ jobs: env: PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi .jenkins/pytorch/win-test.sh - uses: actions/upload-artifact@v2 name: Store PyTorch Test Reports diff --git a/.github/workflows/pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7.yml b/.github/workflows/pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7.yml index 39677325aa546..ff8f6484ce8f8 100644 --- a/.github/workflows/pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7.yml +++ b/.github/workflows/pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7.yml @@ -100,6 +100,7 @@ jobs: needs: calculate-docker-image env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7-build steps: - name: Log in to ECR run: | @@ -125,6 +126,7 @@ jobs: run: | docker run \ -e BUILD_ENVIRONMENT \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ -e SCCACHE_BUCKET \ -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ @@ -161,7 +163,7 @@ jobs: docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . - name: Archive artifacts into zip run: | - zip -r artifacts.zip dist/ build/ + zip -r artifacts.zip dist/ build/ .pytorch-test-times.json # Upload to github so that people can click and download artifacts - uses: actions/upload-artifact@v2 # Don't fail on upload to GH since it's only for user convenience @@ -187,13 +189,40 @@ jobs: # Prune all of the docker images docker system prune -af + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 1 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: linux.8xlarge.nvidia.gpu needs: - calculate-docker-image - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-cuda10.2-cudnn7-py3.6-gcc7-test + NUM_TEST_SHARDS: 1 + TEST_CONFIG: ${{ matrix.test_config }} steps: - name: Log in to ECR run: | @@ -243,6 +272,11 @@ jobs: env | grep '^GITHUB' > "/tmp/github_env_${GITHUB_RUN_ID}" - name: Test PyTorch run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi # TODO: Stop building test binaries as part of the build phase # Used for GPU_FLAG since that doesn't play nice # shellcheck disable=SC2086 @@ -252,6 +286,8 @@ jobs: -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ -e GITHUB_ACTIONS \ -e IN_CI \ + -e SHARD_NUMBER \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ --security-opt seccomp=unconfined \ diff --git a/.github/workflows/pytorch-linux-xenial-py3.6-gcc5.4.yml b/.github/workflows/pytorch-linux-xenial-py3.6-gcc5.4.yml index 3e55acb94bda1..733316517f6a6 100644 --- a/.github/workflows/pytorch-linux-xenial-py3.6-gcc5.4.yml +++ b/.github/workflows/pytorch-linux-xenial-py3.6-gcc5.4.yml @@ -101,6 +101,7 @@ jobs: needs: calculate-docker-image env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-py3.6-gcc5.4-build steps: - name: Log in to ECR run: | @@ -126,6 +127,7 @@ jobs: run: | docker run \ -e BUILD_ENVIRONMENT \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ -e SCCACHE_BUCKET \ -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ @@ -162,7 +164,7 @@ jobs: docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . - name: Archive artifacts into zip run: | - zip -r artifacts.zip dist/ build/ + zip -r artifacts.zip dist/ build/ .pytorch-test-times.json # Upload to github so that people can click and download artifacts - uses: actions/upload-artifact@v2 # Don't fail on upload to GH since it's only for user convenience @@ -188,13 +190,40 @@ jobs: # Prune all of the docker images docker system prune -af + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 2 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: linux.2xlarge needs: - calculate-docker-image - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false env: DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-py3.6-gcc5.4-test + NUM_TEST_SHARDS: 2 + TEST_CONFIG: ${{ matrix.test_config }} steps: - name: Log in to ECR run: | @@ -244,6 +273,11 @@ jobs: env | grep '^GITHUB' > "/tmp/github_env_${GITHUB_RUN_ID}" - name: Test PyTorch run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi # TODO: Stop building test binaries as part of the build phase # Used for GPU_FLAG since that doesn't play nice # shellcheck disable=SC2086 @@ -253,6 +287,8 @@ jobs: -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ -e GITHUB_ACTIONS \ -e IN_CI \ + -e SHARD_NUMBER \ + -e JOB_BASE_NAME \ -e MAX_JOBS="$(nproc --ignore=2)" \ --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ --security-opt seccomp=unconfined \ diff --git a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml index 6ff79a6795b94..2d9cd8535b009 100644 --- a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml @@ -17,7 +17,6 @@ env: CUDA_VERSION: "cpu" IN_CI: 1 INSTALL_WINDOWS_SDK: 1 - JOB_BASE_NAME: test PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" @@ -31,6 +30,8 @@ concurrency: jobs: build: runs-on: "windows.4xlarge" + env: + JOB_BASE_NAME: pytorch-win-vs2019-cpu-py3-build steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -68,12 +69,38 @@ jobs: name: ${{ env.BUILD_ENVIRONMENT }} path: C:\w\build-results + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 1 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: windows.4xlarge env: JOB_BASE_NAME: pytorch-win-vs2019-cpu-py3-test + NUM_TEST_SHARDS: 1 + TEST_CONFIG: ${{ matrix.test_config }} needs: - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -106,6 +133,11 @@ jobs: env: PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi .jenkins/pytorch/win-test.sh - uses: actions/upload-artifact@v2 name: Store PyTorch Test Reports diff --git a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml index 9cd66ebd4b723..c1c476e1cfa85 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml @@ -16,7 +16,6 @@ env: CUDA_VERSION: "10.1" IN_CI: 1 INSTALL_WINDOWS_SDK: 1 - JOB_BASE_NAME: test PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" @@ -32,6 +31,8 @@ concurrency: jobs: build: runs-on: "windows.4xlarge" + env: + JOB_BASE_NAME: pytorch-win-vs2019-cuda10-cudnn7-py3-build steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -77,12 +78,38 @@ jobs: name: ${{ env.BUILD_ENVIRONMENT }} path: C:\w\build-results + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 1 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: windows.8xlarge.nvidia.gpu env: JOB_BASE_NAME: pytorch-win-vs2019-cuda10-cudnn7-py3-test + NUM_TEST_SHARDS: 1 + TEST_CONFIG: ${{ matrix.test_config }} needs: - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -123,6 +150,11 @@ jobs: env: PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi .jenkins/pytorch/win-test.sh - uses: actions/upload-artifact@v2 name: Store PyTorch Test Reports diff --git a/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml index 0de71b92f0add..2a78b77c1bc26 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml @@ -16,7 +16,6 @@ env: CUDA_VERSION: "11.1" IN_CI: 1 INSTALL_WINDOWS_SDK: 1 - JOB_BASE_NAME: test PYTHON_VERSION: "3.8" SCCACHE_BUCKET: "ossci-compiler-cache" VC_PRODUCT: "BuildTools" @@ -32,6 +31,8 @@ concurrency: jobs: build: runs-on: "windows.4xlarge" + env: + JOB_BASE_NAME: pytorch-win-vs2019-cuda11-cudnn8-py3-build steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -77,12 +78,38 @@ jobs: name: ${{ env.BUILD_ENVIRONMENT }} path: C:\w\build-results + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 1 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + test: runs-on: windows.8xlarge.nvidia.gpu env: JOB_BASE_NAME: pytorch-win-vs2019-cuda11-cudnn8-py3-test + NUM_TEST_SHARDS: 1 + TEST_CONFIG: ${{ matrix.test_config }} needs: - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false steps: - name: Checkout PyTorch uses: actions/checkout@v2 @@ -123,6 +150,11 @@ jobs: env: PYTORCH_FINAL_PACKAGE_DIR: /c/${{ github.run_id }}/build-results/ run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi .jenkins/pytorch/win-test.sh - uses: actions/upload-artifact@v2 name: Store PyTorch Test Reports diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index da1c766f63235..eeef1b9272485 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -443,7 +443,7 @@ elif [[ "${BUILD_ENVIRONMENT}" == *jit_legacy-test || "${JOB_BASE_NAME}" == *jit elif [[ "${BUILD_ENVIRONMENT}" == *libtorch* ]]; then # TODO: run some C++ tests echo "no-op at the moment" -elif [[ "${BUILD_ENVIRONMENT}" == *-test1 || "${JOB_BASE_NAME}" == *-test1 ]]; then +elif [[ "${BUILD_ENVIRONMENT}" == *-test1 || "${SHARD_NUMBER}" == 1 ]]; then if [[ "${BUILD_ENVIRONMENT}" == pytorch-linux-xenial-cuda11.1-cudnn8-py3-gcc7-test1 ]]; then test_torch_deploy fi @@ -451,7 +451,7 @@ elif [[ "${BUILD_ENVIRONMENT}" == *-test1 || "${JOB_BASE_NAME}" == *-test1 ]]; t install_torchvision test_python_shard1 test_aten -elif [[ "${BUILD_ENVIRONMENT}" == *-test2 || "${JOB_BASE_NAME}" == *-test2 ]]; then +elif [[ "${BUILD_ENVIRONMENT}" == *-test2 || "${SHARD_NUMBER}" == 2 ]]; then install_torchvision test_python_shard2 test_libtorch diff --git a/.jenkins/pytorch/win-test.sh b/.jenkins/pytorch/win-test.sh index 736886b105c37..a27867e727d21 100755 --- a/.jenkins/pytorch/win-test.sh +++ b/.jenkins/pytorch/win-test.sh @@ -64,7 +64,7 @@ run_tests() { fi done - if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then + if [[ ( -z "${JOB_BASE_NAME}" || "${JOB_BASE_NAME}" == *-test ) && $NUM_TEST_SHARDS -eq 1 ]]; then "$SCRIPT_HELPERS_DIR"/test_python.bat "$DETERMINE_FROM" "$SCRIPT_HELPERS_DIR"/test_custom_script_ops.bat "$SCRIPT_HELPERS_DIR"/test_custom_backend.bat @@ -74,13 +74,13 @@ run_tests() { export PYTORCH_COLLECT_COVERAGE=1 export COVERAGE_RCFILE=$PWD/.coveragerc # coverage config file needed for plug-ins and settings to work fi - if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then + if [[ "${JOB_BASE_NAME}" == *-test1 || "${SHARD_NUMBER}" == 1 ]]; then "$SCRIPT_HELPERS_DIR"/test_python_first_shard.bat "$DETERMINE_FROM" "$SCRIPT_HELPERS_DIR"/test_libtorch.bat if [[ "${USE_CUDA}" == "1" ]]; then "$SCRIPT_HELPERS_DIR"/test_python_jit_legacy.bat "$DETERMINE_FROM" fi - elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then + elif [[ "${JOB_BASE_NAME}" == *-test2 || "${SHARD_NUMBER}" == 2 ]]; then "$SCRIPT_HELPERS_DIR"/test_python_second_shard.bat "$DETERMINE_FROM" "$SCRIPT_HELPERS_DIR"/test_custom_backend.bat "$SCRIPT_HELPERS_DIR"/test_custom_script_ops.bat diff --git a/test/run_test.py b/test/run_test.py index 997d43c0cad27..d385224373f7b 100755 --- a/test/run_test.py +++ b/test/run_test.py @@ -419,12 +419,12 @@ def print_to_stderr(message): # Convert something like pytorch_windows_vs2019_py36_cuda10.1_build to pytorch_windows_vs2019_py36_cuda10.1 def get_stripped_CI_job() -> str: - job = os.environ.get("CIRCLE_JOB", "").rstrip('0123456789') + job = os.environ.get("JOB_BASE_NAME", os.environ.get("CIRCLE_JOB", "")).rstrip('0123456789') if job.endswith('_slow_test'): job = job[:len(job) - len('_slow_test')] - elif job.endswith('_test'): + elif job.endswith('_test') or job.endswith('-test'): job = job[:len(job) - len('_test')] - elif job.endswith('_build'): + elif job.endswith('_build') or job.endswith('-build'): job = job[:len(job) - len('_build')] return job From e2129d1c067326efba4eac53255b94af05a45b1b Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 13:08:37 -0700 Subject: [PATCH 208/305] beef up at::_ops API (#59115) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59115 This PR beefs up the `at::_ops::` API as a source of truth for compile-time information about each operator. ### Changes For every op defined in native_functions.yaml, e.g. `at::_ops::add_Tensor` previously defined an unambiguous function; effectively an unambiguously named version of the C++ API that you could decltype() successfully because it had no overloads with a user-facing macro: `decltype(ATEN_FN2(add, Tensor)) // expands to decltype(at::_ops::add_Tensor)`. Now, `at::_ops::add_Tensor` is a struct containing a few static fields and methods (declared in `Operators.h`, defined in `Operators.cpp`): ``` struct TORCH_API add_Tensor { using schema = at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &); using ptr_schema = at::Tensor (*)(const at::Tensor &, const at::Tensor &, const at::Scalar &); static constexpr const char* name = "aten::add"; static constexpr const char* overload_name = "Tensor"; static constexpr const char* schema_str = "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"; static at::Tensor call(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha); static at::Tensor redispatch(c10::DispatchKeySet dispatchKeySet, const at::Tensor & self, const at::Tensor & ot }; ``` What used to be the function `at::_ops::add_Tensor` can now be accessed as `at::_ops::add_Tensor::call`, and I've added a new macro to access the entire struct (naming suggestions welcome) - `ATEN_OP2(add, Tensor)`. ### Motivation There were two motivations for this change: **Codegen refactor** The `at::_ops::` API as it exists now is (yet another) C++ entry point into the dispatcher, in addition to the Function, Method, and Redispatch APIs. Instead, after this PR, the existing three API's are all inline-able wrapper API's that call into the `at::_ops` API to do the real work. The function and method API's call into `at::_ops::{op}::call`, while the redispatch API calls into `at::_ops::{op}::redispatch`. This will hopefully make it easier to pile in any future C++ API's that we want to code-generate. It also means that stuff like the string name, overload name, and schema of each operator is consolidated in a single place, rather than having the codegen hardcode various strings in multiple codegen output files. **Extra compile-time metadata** In the [boxed CPU fallback PR](https://github.com/pytorch/pytorch/pull/58065/files#diff-c9b55f0d692a9bea8019c6f19bc46877f1efa0f9d4fc2086cf299b52768343b4R31) above this in the stack, I added a new API that external backends can use to call directly into their boxed fallback from an unboxed context. Adding extra metadata to `at::_ops` means that XLA's usage of that API doesn't require passing in the string name and overload of each name as arguments; we can just infer them. The updated API looks like this (see [the XLA-side PR ](https://github.com/pytorch/xla/pull/2945/files#diff-5e65c3c1d847191cb691d1874732e971f09fa1aad7a980a555c3b0504a5b6470R250) for more examples) ``` return at::native::call_fallback_fn<&xla_cpu_fallback, ATEN_OP2(add, Tensor)>::call(a, b, 1.0); ``` **Characteristics of the `at::_ops` API** (I also commented this in the codegen) (1) It follows the Dispatcher API. This means, e.g., that it takes in the expanded arguments rather than `TensorOptions`. This is kind of necessary for perf, if we want to `at::_ops` to serve as the main implementation of the existing C++ API's. For example: if it followed the C++ API, then all of the faithful C++ factory functions would need to wrap their arguments into TensorOptions only to unwrap them again. (2) Overload names are disambiguated. This is the same as before; it's helpful for pytorch extenders who would like to decltype() an aten operator, that has overloads, e.g. decltype(at::_ops::mul_Tensor::call) (3) No argument defaulting is allowed. This is more of an implementation detail to avoid #include cycles, since TensorBody.h (which defines the Tensor class) needs to include this file. The #include situation is precarious though! (4) manual_cpp_bindings and faithful names are not included in the API. I think that this is one we have a choice with. This applies to stuff like __dispatch__is_complex(), and add_outf(). These aren't "real native_functions.yaml ops", they're just additional functions provided by the C++ API. They're implemented as wrappers in Functions.h that call into the actual operators defined here, i.e. at::_ops::is_complex::call() and at::_ops::add_out::call(). This means that ATEN_OP(is_complex) will not fastpath, and will go through the dispatcher. It also means that `ATEN_OP2(add, out)` is automatically faithful and takes its out argument at the end (this is just because it follows the dispatcher API). **Details** Instead of codegen'ing the existing 3 API's in `Functions.cpp`, `TensorMethods.cpp` and `RedispatchFunctions.cpp`, I codegen them directly into the headers: `Functions.h`, `TensorBody.h`, and `RedispatchFunctions.h`. I mostly did this for perf, since we want to avoid introducing an extra function call in the hot path of every operator. These functions are also now all one-liners that call into `at::_ops`, so the compiler should just inline them all anyway. The main downside in doing that though was that I had to bend over backwards in a few cases to avoid cyclical #include statements. The issue is that `TensorBody.h` now includes `Operators.h` (because the codegen'd method API is implemented by calling into `at::_ops`), but `TensorBody.h` also includes the definition of the Tensor class. That means that `Operators.h` can't be aware of the Tensor class; it needs to forward declare everything and avoid using the Tensor class directly. To fix cyclic includes, I had to: - Not allow defaulting in the `at::_ops` API - Move some code that was called when translating from C++ to Dispatcher API's directly into the codegen template (`check_tensor_options_and_extract_memory_format`) It's not great, but I don't think this specific include cycle will break down in the near future; the only code that we need to call before getting to `Operators.cpp` is the translations from various API's to the dispatcher API; there aren't many of them, and there's no major reason for them to live an external utils file somewhere. Moving the code into the headers also meant that the codegen no longer needs to deal with `Functions.cpp`/`TensorMethods.cpp`/`RedispatchFunctions.cpp`. All of the functions that used to be defined in `TensorMethods.cpp` seemed small enough for me to lump into `TensorBody.h`, but some of the functions in `Functions.cpp` looked pretty big to put in a header, so I moved the file to `aten/src/ATen/native/Functions.cpp`. It might be worth keeping `TensorMethods.cpp` there and leaving it too, in-case we have any beefy hand-written tensor methods that we don't want to put in a header. **Perf** I ran a few benchmarks in callgrind, and didn't see a noticeable instruction count change when calling `at::add()`. I also saw in the output that `at::add()` was successfully getting inlined. There's also probably a light risk of binary size increase; I think that there's a binary size regression test that I can run in phabricator (going to try it). I can also try inspecting `libtorch.so` directly and seeing if it's any bigger, but my hope is that the inline-ing means that we aren't generated separate symbols for `at::add` and `at::_ops::add_Tensor::call`. Test Plan: Imported from OSS Reviewed By: ezyang Differential Revision: D28833086 Pulled By: bdhirsh fbshipit-source-id: 55f322a8378cb9a3cb6642f72aa291be381dd95b --- BUILD.bazel | 8 +- aten/src/ATen/core/op_registration/adaption.h | 20 -- .../src/ATen/templates/DispatchKeyFunctions.h | 37 ++- .../ATen/templates/DispatchKeyFunctions_inl.h | 16 + aten/src/ATen/templates/Functions.cpp | 160 --------- aten/src/ATen/templates/Functions.h | 121 +++++-- aten/src/ATen/templates/Operators.cpp | 2 + aten/src/ATen/templates/Operators.h | 85 ++++- aten/src/ATen/templates/RedispatchFunctions.h | 2 +- .../ATen/templates/RegisterBackendSelect.cpp | 1 + aten/src/ATen/templates/RegisterSchema.cpp | 1 + aten/src/ATen/templates/TensorBody.h | 75 ++++- aten/src/ATen/templates/TensorMethods.cpp | 85 ----- caffe2/contrib/aten/gen_op.py | 8 +- tools/codegen/api/python.py | 22 +- tools/codegen/api/types.py | 38 ++- tools/codegen/gen.py | 312 +++++++++--------- tools/codegen/model.py | 12 + 18 files changed, 495 insertions(+), 510 deletions(-) create mode 100644 aten/src/ATen/templates/DispatchKeyFunctions_inl.h delete mode 100644 aten/src/ATen/templates/Functions.cpp delete mode 100644 aten/src/ATen/templates/TensorMethods.cpp diff --git a/BUILD.bazel b/BUILD.bazel index b7e16ac1c915c..217b20fb54b9d 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -136,20 +136,22 @@ genrule( "aten/src/ATen/RegisterMeta.cpp", "aten/src/ATen/RegisterSchema.cpp", "aten/src/ATen/CPUFunctions.h", + "aten/src/ATen/CPUFunctions_inl.h", "aten/src/ATen/CUDAFunctions.h", + "aten/src/ATen/CUDAFunctions_inl.h", "aten/src/ATen/CompositeExplicitAutogradFunctions.h", + "aten/src/ATen/CompositeExplicitAutogradFunctions_inl.h", "aten/src/ATen/CompositeImplicitAutogradFunctions.h", + "aten/src/ATen/CompositeImplicitAutogradFunctions_inl.h", "aten/src/ATen/Functions.h", - "aten/src/ATen/Functions.cpp", "aten/src/ATen/RedispatchFunctions.h", - "aten/src/ATen/RedispatchFunctions.cpp", "aten/src/ATen/Operators.h", "aten/src/ATen/Operators.cpp", "aten/src/ATen/NativeFunctions.h", "aten/src/ATen/MetaFunctions.h", + "aten/src/ATen/MetaFunctions_inl.h", "aten/src/ATen/NativeMetaFunctions.h", "aten/src/ATen/core/TensorBody.h", - "aten/src/ATen/core/TensorMethods.cpp", "aten/src/ATen/core/ATenOpList.cpp", ], cmd = "$(location :gen) --source-path aten/src/ATen --install_dir `dirname $(location aten/src/ATen/Declarations.yaml)`", diff --git a/aten/src/ATen/core/op_registration/adaption.h b/aten/src/ATen/core/op_registration/adaption.h index 327c1d10e5a28..5bf1b691ebad3 100644 --- a/aten/src/ATen/core/op_registration/adaption.h +++ b/aten/src/ATen/core/op_registration/adaption.h @@ -43,26 +43,6 @@ namespace c10 { namespace impl { -inline c10::optional -check_tensor_options_and_extract_memory_format( - const TensorOptions& options, - c10::optional memory_format) { - TORCH_CHECK( - options.requires_grad_opt() == c10::nullopt || - options.requires_grad_opt().value() == false, - "Operators taking TensorOptions cannot take a TensorOptions with " - "options.requires_grad set as true. This isn't implemented yet."); - TORCH_CHECK( - !(options.has_memory_format() && memory_format.has_value()), - "Cannot set memory_format both in TensorOptions and explicit argument; please delete " - "the redundant setter."); - if (memory_format.has_value()) { - return memory_format; - } else { - return options.memory_format_opt(); - } -} - TORCH_API void common_device_check_failure(optional& common_device, const at::Tensor& tensor, at::CheckedFrom methodName, at::CheckedFrom argName); inline void check_and_update_common_device(optional& common_device, const at::Tensor& tensor, at::CheckedFrom methodName, at::CheckedFrom argName) { diff --git a/aten/src/ATen/templates/DispatchKeyFunctions.h b/aten/src/ATen/templates/DispatchKeyFunctions.h index e72b39b5ae5f0..1718b4be8274c 100644 --- a/aten/src/ATen/templates/DispatchKeyFunctions.h +++ b/aten/src/ATen/templates/DispatchKeyFunctions.h @@ -1,14 +1,23 @@ -// ${generated_comment} - -// NB: The implementing C++ file is RegisterDispatchKey.cpp - -// TODO: tighten this include -#include - -namespace at { -namespace ${dispatch_namespace} { - -${dispatch_namespaced_declarations} - -} // namespace ${dispatch_namespace} -} // namespace at +#include +// Note [Avoiding Include Cycles In Static Dispatch] +// In order to avoid #include cycles in the static dispatch build, we've carefully split out +// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h. +// +// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h. +// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods +// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all +// directly inlined into TensorBody.h. +// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API, +// which include functions that have defaultable optional arguments. +// That requires knowing the full Tensor class definition. +// +// We break the cycle by doing the following: +// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h +// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl., +// - CPUFunctions_inl.h includes everything else +// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class, +// and then it includes CPUFunctions_inl.h. +// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly. +// - This also means that static dispatch build, CPUFunctions.h only needs to +// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h. +${inline_headers_for_nonstatic_build} diff --git a/aten/src/ATen/templates/DispatchKeyFunctions_inl.h b/aten/src/ATen/templates/DispatchKeyFunctions_inl.h new file mode 100644 index 0000000000000..365ce8b98d614 --- /dev/null +++ b/aten/src/ATen/templates/DispatchKeyFunctions_inl.h @@ -0,0 +1,16 @@ +// ${generated_comment} + +// NB: The implementing C++ file is RegisterDispatchKey.cpp + +// The only #includes we need are for custom classes that have defaults in the C++ API +#include +#include +#include + +namespace at { +namespace ${dispatch_namespace} { + +${dispatch_namespaced_declarations} + +} // namespace ${dispatch_namespace} +} // namespace at diff --git a/aten/src/ATen/templates/Functions.cpp b/aten/src/ATen/templates/Functions.cpp deleted file mode 100644 index 3d119b6314abb..0000000000000 --- a/aten/src/ATen/templates/Functions.cpp +++ /dev/null @@ -1,160 +0,0 @@ -// ${generated_comment} - -#include - -#include -#include - -#include -#include - -${static_dispatch_extra_headers} - -namespace at { - -Tensor var(const Tensor& self, int dim) { - return at::var(self, IntArrayRef{dim}); -} - -std::tuple var_mean(const Tensor& self, int dim) { - return at::var_mean(self, IntArrayRef{dim}); -} - -Tensor std(const Tensor& self, int dim) { - return at::std(self, IntArrayRef{dim}); -} - -std::tuple std_mean(const Tensor& self, int dim) { - return at::std_mean(self, IntArrayRef{dim}); -} - -at::Tensor conv1d( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - std::initializer_list padding_, - IntArrayRef dilation, - int64_t groups) { - auto padding = IntArrayRef(padding_); - return at::conv1d(input, weight, bias, stride, padding, dilation, groups); -} - -at::Tensor conv2d( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - std::initializer_list padding_, - IntArrayRef dilation, - int64_t groups) { - auto padding = IntArrayRef(padding_); - return at::conv2d(input, weight, bias, stride, padding, dilation, groups); -} - -at::Tensor conv3d( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - std::initializer_list padding_, - IntArrayRef dilation, - int64_t groups) { - auto padding = IntArrayRef(padding_); - return at::conv3d(input, weight, bias, stride, padding, dilation, groups); -} - -namespace detail { - -void noopDelete(void*) {} - -} // namespace detail - -Tensor TensorMaker::make_tensor() { - AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove. - tracer::impl::NoTracerDispatchMode tracer_guard{}; - - check_size_nonnegative(sizes_); - - TORCH_CHECK_VALUE( - !deleter_ || !ctx_, - "The deleter and context arguments are mutually exclusive."); - - if (device_ == nullopt) { - device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type()); - } - - if (opts_.device().has_index()) { - // clang-format off - TORCH_CHECK_VALUE( - opts_.device() == *device_, - "Specified device ", opts_.device(), " does not match device of data ", *device_); - // clang-format on - } - - std::size_t size_bytes = computeStorageSize(); - - DataPtr data_ptr{}; - if (deleter_) { - data_ptr = makeDataPtrFromDeleter(); - } else { - data_ptr = makeDataPtrFromContext(); - } - - Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr)}; - - Tensor tensor = detail::make_tensor( - std::move(storage), opts_.computeDispatchKey(), opts_.dtype()); - - if (sizes_.size() != 1 || sizes_[0] != 0) { - TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl(); - - if (strides_) { - tensor_impl->set_sizes_and_strides(sizes_, *strides_); - } else { - tensor_impl->set_sizes_contiguous(sizes_); - } - } - - return tensor; -} - -std::size_t TensorMaker::computeStorageSize() const noexcept { - std::size_t itemsize = opts_.dtype().itemsize(); - - if (strides_) { - return detail::computeStorageNbytes(sizes_, *strides_, itemsize); - } - - std::size_t size = 1; - for (std::int64_t s : sizes_) { - size *= static_cast(s); - } - return size * itemsize; -} - -inline DataPtr TensorMaker::makeDataPtrFromDeleter() const { - return InefficientStdFunctionContext::makeDataPtr(data_, deleter_, *device_); -} - -inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept { - return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_}; -} - -IntArrayRef TensorMaker::makeTempSizes() const noexcept { - static std::int64_t zeros[5] = {0, 0, 0, 0, 0}; - if (opts_.has_memory_format()) { - MemoryFormat format = *opts_.memory_format_opt(); - if (format == MemoryFormat::ChannelsLast) { - return IntArrayRef(zeros, 4); - } - if (format == MemoryFormat::ChannelsLast3d) { - return IntArrayRef(zeros, 5); - } - } - return IntArrayRef(zeros, 1); -} - -${function_definitions} - -} // namespace at diff --git a/aten/src/ATen/templates/Functions.h b/aten/src/ATen/templates/Functions.h index 18e7c00a0bb0a..caeb328e0c5bd 100644 --- a/aten/src/ATen/templates/Functions.h +++ b/aten/src/ATen/templates/Functions.h @@ -14,6 +14,9 @@ #include #include #include +#include + +${static_dispatch_extra_headers} namespace at { @@ -40,32 +43,27 @@ AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR) AT_FORALL_COMPLEX_TYPES(TENSOR) #undef TENSOR -${function_declarations} +${function_definitions} // Special C++ only overloads for std()-like functions (See gh-40287) // These are needed because int -> bool conversion takes precedence over int -> IntArrayRef // So, for example std(0) would select the std(unbiased=False) overload -TORCH_API Tensor var(const Tensor& self, int dim); -TORCH_API std::tuple var_mean(const Tensor& self, int dim); -TORCH_API Tensor std(const Tensor& self, int dim); -TORCH_API std::tuple std_mean(const Tensor& self, int dim); - - -// Special C++ only overloads for convnd functions (See gh-45667) -// These are needed because {1, 2} is ambiguous between string and IntArrayRef overloads -TORCH_API at::Tensor conv1d( - const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, - std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); -TORCH_API at::Tensor conv2d( - const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, - std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); -TORCH_API at::Tensor conv3d( - const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, - std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); +TORCH_API inline Tensor var(const Tensor& self, int dim) { + return at::var(self, IntArrayRef{dim}); +} +TORCH_API inline std::tuple var_mean(const Tensor& self, int dim) { + return at::var_mean(self, IntArrayRef{dim}); +} +TORCH_API inline Tensor std(const Tensor& self, int dim) { + return at::std(self, IntArrayRef{dim}); +} +TORCH_API inline std::tuple std_mean(const Tensor& self, int dim) { + return at::std_mean(self, IntArrayRef{dim}); +} namespace detail { -TORCH_API void noopDelete(void*); +TORCH_API inline void noopDelete(void*) {} } // namespace detail @@ -117,19 +115,94 @@ class TORCH_API TensorMaker { return *this; } - Tensor make_tensor(); + Tensor make_tensor() { + AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove. + tracer::impl::NoTracerDispatchMode tracer_guard{}; + + check_size_nonnegative(sizes_); + + TORCH_CHECK_VALUE( + !deleter_ || !ctx_, + "The deleter and context arguments are mutually exclusive."); + + if (device_ == nullopt) { + device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type()); + } + + if (opts_.device().has_index()) { + // clang-format off + TORCH_CHECK_VALUE( + opts_.device() == *device_, + "Specified device ", opts_.device(), " does not match device of data ", *device_); + // clang-format on + } + + std::size_t size_bytes = computeStorageSize(); + + DataPtr data_ptr{}; + if (deleter_) { + data_ptr = makeDataPtrFromDeleter(); + } else { + data_ptr = makeDataPtrFromContext(); + } + + Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr)}; + + Tensor tensor = detail::make_tensor( + std::move(storage), opts_.computeDispatchKey(), opts_.dtype()); + + if (sizes_.size() != 1 || sizes_[0] != 0) { + TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl(); + + if (strides_) { + tensor_impl->set_sizes_and_strides(sizes_, *strides_); + } else { + tensor_impl->set_sizes_contiguous(sizes_); + } + } + + return tensor; + } private: explicit TensorMaker(void* data, IntArrayRef sizes) noexcept : data_{data}, sizes_{sizes} {} - std::size_t computeStorageSize() const noexcept; + std::size_t computeStorageSize() const noexcept { + std::size_t itemsize = opts_.dtype().itemsize(); - DataPtr makeDataPtrFromDeleter() const; + if (strides_) { + return detail::computeStorageNbytes(sizes_, *strides_, itemsize); + } - DataPtr makeDataPtrFromContext() noexcept; + std::size_t size = 1; + for (std::int64_t s : sizes_) { + size *= static_cast(s); + } + return size * itemsize; + } - IntArrayRef makeTempSizes() const noexcept; + inline DataPtr makeDataPtrFromDeleter() const { + return InefficientStdFunctionContext::makeDataPtr(data_, deleter_, *device_); + } + + inline DataPtr makeDataPtrFromContext() noexcept { + return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_}; + } + + IntArrayRef makeTempSizes() const noexcept { + static std::int64_t zeros[5] = {0, 0, 0, 0, 0}; + if (opts_.has_memory_format()) { + MemoryFormat format = *opts_.memory_format_opt(); + if (format == MemoryFormat::ChannelsLast) { + return IntArrayRef(zeros, 4); + } + if (format == MemoryFormat::ChannelsLast3d) { + return IntArrayRef(zeros, 5); + } + } + return IntArrayRef(zeros, 1); + } void* data_; IntArrayRef sizes_; diff --git a/aten/src/ATen/templates/Operators.cpp b/aten/src/ATen/templates/Operators.cpp index 4d50c5a2e0ba2..c0d46f58a848a 100644 --- a/aten/src/ATen/templates/Operators.cpp +++ b/aten/src/ATen/templates/Operators.cpp @@ -1,4 +1,6 @@ #include +#include +#include namespace at { namespace _ops { diff --git a/aten/src/ATen/templates/Operators.h b/aten/src/ATen/templates/Operators.h index 39eaa2bd15a49..a92b7503ad7a2 100644 --- a/aten/src/ATen/templates/Operators.h +++ b/aten/src/ATen/templates/Operators.h @@ -2,8 +2,11 @@ // ${generated_comment} -#include -#include +#include +#include +#include +#include +#include // Extension writers: do you write wrapper functions? Are you frustrated with // resolving overloads of operators? Are you frustrated with dealing with @@ -24,23 +27,77 @@ // ATEN_FN2(sin, out) gives a function that is *faithful* to the schema; // that is, the order of arguments is exactly what it looks like in the schema. -#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload -#define ATEN_FN(op_name) at::_ops::op_name +#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload::call +#define ATEN_FN(op_name) at::_ops::op_name::call + +// Separately, ATEN_OP(op) and ATEN_OP2(op, overload) define a class containing compile-time +// metadata about a given aten operator. +// Notable data on the class includes: +// - ATEN_OP2(add, Tensor)::name // returns the string name: "add" +// - ATEN_OP2(add, Tensor)::overload_name // returns the string overload name: "Tensor" +// - ATEN_OP2(add, Tensor)::schema // returns the C++ schema type: at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &) +// - ATEN_OP2(add, Tensor)::schema_str // returns the string jit type: "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor" + +#define ATEN_OP2(op_name, overload) at::_ops::op_name##_##overload +#define ATEN_OP(op_name) at::_ops::op_name // WARNING: Please do not call any of the ops in the _ops namespace directly. // Use the ATEN_FN macros. We do not guarantee stability of the naming // scheme for the functions in at::_ops -namespace at { namespace _ops { -// NB: We are forced to special case requires_grad_. This is because all -// of the auto-generated inplace method signatures in TensorMethods.h are -// codegen'ed to return Tensor&, but requires_grad_ has a `manual_cpp_binding` -// with a different signature that returns `const Tensor&`. -// -// Eventually, the plan is to kill Tensor& from all C++ signatures and use -// const Tensor&. When that happens, we can remove this special case and just -// let the codegen handle it. -TORCH_API Tensor & requires_grad_(Tensor & self, bool requires_grad); +// See Note [The ATen Operators API] for details of the at::_ops namespace + +namespace c10 { namespace impl { + +inline c10::optional +check_tensor_options_and_extract_memory_format( + const TensorOptions& options, + c10::optional memory_format) { + TORCH_CHECK( + options.requires_grad_opt() == c10::nullopt || + options.requires_grad_opt().value() == false, + "Operators taking TensorOptions cannot take a TensorOptions with " + "options.requires_grad set as true. This isn't implemented yet."); + TORCH_CHECK( + !(options.has_memory_format() && memory_format.has_value()), + "Cannot set memory_format both in TensorOptions and explicit argument; please delete " + "the redundant setter."); + if (memory_format.has_value()) { + return memory_format; + } else { + return options.memory_format_opt(); + } +} + +}} // namespace impl namespace c10 + + +// Forward declarations of any types needed in the operator signatures. +// We can't directly include these classes because it will cause circular include dependencies. +// This file is included by TensorBody.h, which defines the Tensor class. +namespace c10 { + +template +class optional; +template +class List; +class Stream; +struct Storage; + +} + +namespace at { + +class Tensor; +struct Dimname; +struct Generator; +using TensorList = c10::ArrayRef; +using DimnameList = c10::ArrayRef; +using Stream = c10::Stream; +using Storage = c10::Storage; +using QScheme = c10::QScheme; + +namespace _ops { ${declarations} diff --git a/aten/src/ATen/templates/RedispatchFunctions.h b/aten/src/ATen/templates/RedispatchFunctions.h index ecaf8f05162e1..b5219425e8c8c 100644 --- a/aten/src/ATen/templates/RedispatchFunctions.h +++ b/aten/src/ATen/templates/RedispatchFunctions.h @@ -18,7 +18,7 @@ namespace at { namespace redispatch { - ${function_redispatch_declarations} + ${function_redispatch_definitions} } // namespace redispatch } diff --git a/aten/src/ATen/templates/RegisterBackendSelect.cpp b/aten/src/ATen/templates/RegisterBackendSelect.cpp index bf5cc8683ecbc..27b8e2bcc5125 100644 --- a/aten/src/ATen/templates/RegisterBackendSelect.cpp +++ b/aten/src/ATen/templates/RegisterBackendSelect.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include diff --git a/aten/src/ATen/templates/RegisterSchema.cpp b/aten/src/ATen/templates/RegisterSchema.cpp index c9dbf5880a7e1..6861b5cbbcf3e 100644 --- a/aten/src/ATen/templates/RegisterSchema.cpp +++ b/aten/src/ATen/templates/RegisterSchema.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include diff --git a/aten/src/ATen/templates/TensorBody.h b/aten/src/ATen/templates/TensorBody.h index fa879d656ab51..775d2e6803aa6 100644 --- a/aten/src/ATen/templates/TensorBody.h +++ b/aten/src/ATen/templates/TensorBody.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -366,8 +367,15 @@ class TORCH_API Tensor { bool is_alias_of(const at::Tensor& other) const{ return impl_->storage().is_alias_of(other.storage()); } - Tensor toType(ScalarType t) const; - Tensor toBackend(Backend b) const; + + Tensor toType(ScalarType t) const { + return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false); + } + + // TODO: Deprecate me + Tensor toBackend(Backend b) const { + return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false); + } C10_DEPRECATED_MESSAGE("Tensor.is_variable() is deprecated; everything is a variable now. (If you want to assert that variable has been appropriately handled already, use at::impl::variable_excluded_from_dispatch())") bool is_variable() const noexcept { @@ -515,7 +523,11 @@ class TORCH_API Tensor { /// Returns the `TensorOptions` corresponding to this `Tensor`. Defined in /// TensorOptions.h. - TensorOptions options() const; + TensorOptions options() const { + return TensorOptions().dtype(dtype()) + .device(device()) + .layout(layout()); + } void* data_ptr() const { return this->unsafeGetTensorImpl()->data(); @@ -609,11 +621,26 @@ class TORCH_API Tensor { Tensor & index_put_(std::initializer_list indices, Tensor const & rhs); Tensor & index_put_(std::initializer_list indices, const Scalar& v); - Tensor cpu() const; - Tensor cuda() const; - Tensor hip() const; - Tensor vulkan() const; - Tensor metal() const; + Tensor cpu() const { + return to(options().device(DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false); + } + + // TODO: The Python version also accepts arguments + Tensor cuda() const { + return to(options().device(DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor hip() const { + return to(options().device(DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor vulkan() const { + return to(options().device(DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor metal() const { + return to(options().device(DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false); + } // ~~~~~ Autograd API ~~~~~ @@ -944,6 +971,31 @@ inline int64_t get_device(const Tensor& self) { return self.get_device(); } +#define DEFINE_CAST(T, name) \ + template <> \ + TORCH_API inline T* Tensor::data_ptr() const { \ + TORCH_CHECK( \ + scalar_type() == ScalarType::name, \ + "expected scalar type " \ + #name \ + " but found ", \ + scalar_type()); \ + return this->unsafeGetTensorImpl()->data_ptr_impl(); \ + } + +AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_CAST) +AT_FORALL_QINT_TYPES(DEFINE_CAST) +#undef DEFINE_CAST + +#define DEFINE_ITEM(T, name) \ + template <> \ + TORCH_API inline T Tensor::item() const { \ + return item().to##name(); \ + } + +AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_ITEM) +#undef DEFINE_ITEM + template auto Tensor::register_hook(T&& hook) const -> Tensor::hook_return_void_t { // Return the grad argument in case of a hook with void return type to have an @@ -978,6 +1030,13 @@ static inline DispatchKey legacyExtractDispatchKey(const Tensor& t) { } // namespace at +// See Note [Avoiding Include Cycles In Static Dispatch] +${static_dispatch_extra_headers} +namespace at { +${tensor_method_definitions} +} // namespace at + + namespace c10 { template <> struct MaybeOwnedTraits { diff --git a/aten/src/ATen/templates/TensorMethods.cpp b/aten/src/ATen/templates/TensorMethods.cpp deleted file mode 100644 index 0eba7dc65d737..0000000000000 --- a/aten/src/ATen/templates/TensorMethods.cpp +++ /dev/null @@ -1,85 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -${static_dispatch_extra_headers} - -namespace at { - -using Stream = c10::Stream; - -Tensor Tensor::cpu() const { - return to(options().device(DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false); -} - -// TODO: The Python version also accepts arguments -Tensor Tensor::cuda() const { - return to(options().device(DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false); -} - -Tensor Tensor::hip() const { - return to(options().device(DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false); -} - -Tensor Tensor::vulkan() const { - return to(options().device(DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false); -} - -Tensor Tensor::metal() const { - return to(options().device(DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false); -} - -Tensor Tensor::toType(ScalarType t) const { - return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false); -} - -// TODO: Deprecate me -Tensor Tensor::toBackend(Backend b) const { - return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false); -} - -TensorOptions Tensor::options() const { - return TensorOptions().dtype(dtype()) - .device(device()) - .layout(layout()); -} - -${tensor_method_definitions} - -#define DEFINE_CAST(T, name) \ - template <> \ - TORCH_API T* Tensor::data_ptr() const { \ - TORCH_CHECK( \ - scalar_type() == ScalarType::name, \ - "expected scalar type " \ - #name \ - " but found ", \ - scalar_type()); \ - return this->unsafeGetTensorImpl()->data_ptr_impl(); \ - } - -AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_CAST) -AT_FORALL_QINT_TYPES(DEFINE_CAST) -#undef DEFINE_CAST - -#define DEFINE_ITEM(T, name) \ - template <> \ - TORCH_API T Tensor::item() const { \ - return item().to##name(); \ - } - -AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_ITEM) -#undef DEFINE_ITEM - -} //namespace at diff --git a/caffe2/contrib/aten/gen_op.py b/caffe2/contrib/aten/gen_op.py index 94a1f1fedc3b2..13908902e2084 100755 --- a/caffe2/contrib/aten/gen_op.py +++ b/caffe2/contrib/aten/gen_op.py @@ -77,7 +77,7 @@ def value_is_tensor_type(v): 'at::Scalar': 'assignTo(Output(${offset}),${output}.type(), ${output});', 'bool': 'assignToValue(Output(${offset}),${output});', 'int64_t': 'assignToValue(Output(${offset}),${output});', - 'std::vector': 'assignListStartingAt(${offset}, ${output});', + '::std::vector': 'assignListStartingAt(${offset}, ${output});', } # for each non-Tensor aten argument, how to we read it from caffe2's @@ -90,8 +90,8 @@ def value_is_tensor_type(v): 'double': 'double ${arg} = readAttribute("${arg}");', 'int64_t': 'int64_t ${arg} = readAttribute("${arg}");', 'at::IntArrayRef': 'auto ${arg} = readIntArrayRef("${arg}");', - 'std::array': 'auto ${arg} = readBoolMask<2>("${arg}");', - 'std::array': 'auto ${arg} = readBoolMask<3>("${arg}");', + '::std::array': 'auto ${arg} = readBoolMask<2>("${arg}");', + '::std::array': 'auto ${arg} = readBoolMask<3>("${arg}");', } # for BC reasons we want to route some of the functions to different @@ -189,7 +189,7 @@ def get_output(o, i): if len(o['returns']) == 1: return 'the_result' else: - return 'std::get<{}>(the_result)'.format(i) + return '::std::get<{}>(the_result)'.format(i) def attribute_names(o): diff --git a/tools/codegen/api/python.py b/tools/codegen/api/python.py index 8eb2f3696f517..5064f069a0643 100644 --- a/tools/codegen/api/python.py +++ b/tools/codegen/api/python.py @@ -613,7 +613,7 @@ def argument_type_str(t: Type, *, simple_type: bool = False) -> str: size = t.size if not simple_type else None if str(t.elem) == 'bool': assert t.size is not None - return f'std::array' + return f'::std::array' elif str(t.elem) == 'int': return f'IntArrayRef[{size}]' if size is not None else 'IntArrayRef' elif str(t.elem) == 'Tensor': @@ -910,16 +910,16 @@ def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument: # to add an appropriate wrap() overload in torch/csrc/autograd/utils/wrap_outputs.h. SUPPORTED_RETURN_TYPES = { 'at::Tensor', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::tuple', - 'std::vector', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::tuple', + '::std::vector', 'at::Scalar', 'bool', 'int64_t', 'void*', 'void', 'at::QScheme', 'double', 'at::IntArrayRef', diff --git a/tools/codegen/api/types.py b/tools/codegen/api/types.py index 850b6a50df285..675db01cd3243 100644 --- a/tools/codegen/api/types.py +++ b/tools/codegen/api/types.py @@ -166,10 +166,10 @@ class VectorCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'std::vector<{self.elem.cpp_type()}>' + return f'::std::vector<{self.elem.cpp_type()}>' def cpp_type_registration_declarations(self) -> str: - return f'std::vector<{self.elem.cpp_type_registration_declarations()}>' + return f'::std::vector<{self.elem.cpp_type_registration_declarations()}>' def remove_const_ref(self) -> 'CType': return VectorCType(self.elem.remove_const_ref()) @@ -181,10 +181,10 @@ class ArrayCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'std::array<{self.elem.cpp_type()},{self.size}>' + return f'::std::array<{self.elem.cpp_type()},{self.size}>' def cpp_type_registration_declarations(self) -> str: - return f'std::array<{self.elem.cpp_type_registration_declarations()},{self.size}>' + return f'::std::array<{self.elem.cpp_type_registration_declarations()},{self.size}>' def remove_const_ref(self) -> 'CType': return ArrayCType(self.elem.remove_const_ref(), self.size) @@ -195,10 +195,10 @@ class TupleCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'std::tuple<{",".join([e.cpp_type() for e in self.elems])}>' + return f'::std::tuple<{",".join([e.cpp_type() for e in self.elems])}>' def cpp_type_registration_declarations(self) -> str: - return f'std::tuple<{",".join([e.cpp_type_registration_declarations() for e in self.elems])}>' + return f'::std::tuple<{",".join([e.cpp_type_registration_declarations() for e in self.elems])}>' def remove_const_ref(self) -> 'CType': return TupleCType([e.remove_const_ref() for e in self.elems]) @@ -338,26 +338,37 @@ def name(self) -> str: return n # Render the C++ declaration for this signature - def decl(self, *, prefix: str = "", is_redispatching_fn: bool = False) -> str: + def decl(self, *, name: Optional[str] = None, prefix: str = "", is_redispatching_fn: bool = False) -> str: returns_type = cpp.returns_type(self.func.returns).cpp_type() cpp_args = [a.decl() for a in self.arguments()] if is_redispatching_fn: cpp_args = ['c10::DispatchKeySet dispatchKeySet'] + cpp_args cpp_args_str = ', '.join(cpp_args) - name = prefix + self.name() + if name is None: + name = prefix + self.name() return f"{returns_type} {name}({cpp_args_str})" # Render the C++ definition for this signature, not including # the body (with curly braces) - def defn(self, *, prefix: str = "", is_redispatching_fn: bool = False) -> str: + def defn(self, *, name: Optional[str] = None, prefix: str = "", is_redispatching_fn: bool = False) -> str: returns_type = cpp.returns_type(self.func.returns).cpp_type() cpp_args = [a.defn() for a in self.arguments()] if is_redispatching_fn: cpp_args = ['c10::DispatchKeySet dispatchKeySet'] + cpp_args cpp_args_str = ', '.join(cpp_args) - name = prefix + self.name() + if name is None: + name = prefix + self.name() return f"{returns_type} {name}({cpp_args_str})" + def ptr_type(self) -> str: + args_types_str = ', '.join(a.type for a in self.arguments()) + return f'{cpp.returns_type(self.func.returns).cpp_type()} (*)({args_types_str})' + + # Return the C++ function type, e.g., something like int(bool) + def type(self) -> str: + args_types_str = ', '.join(a.type for a in self.arguments()) + return f'{cpp.returns_type(self.func.returns).cpp_type()} ({args_types_str})' + # Represents group of all CppSignatures associated with a # FunctionSchema. Right now, that's the regular, user-visible @@ -424,8 +435,11 @@ def decl(self, name: Optional[str] = None) -> str: name = self.name() return f"{self.returns_type().cpp_type()} {name}({args_str})" - def defn(self, name: Optional[str] = None) -> str: - args_str = ', '.join(a.defn() for a in self.arguments()) + def defn(self, name: Optional[str] = None, *, is_redispatching_fn: bool = False) -> str: + args = [a.defn() for a in self.arguments()] + if is_redispatching_fn: + args = ['c10::DispatchKeySet dispatchKeySet'] + args + args_str = ', '.join(args) if name is None: name = self.name() return f"{self.returns_type().cpp_type()} {name}({args_str})" diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index 193a4d35d74bb..e69dc62c456fd 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -140,14 +140,24 @@ def cpp_string(s: str) -> str: # to be generated. This pattern makes it convenient to use map, concatMap # and similar functional combinators. -def static_dispatch_extra_headers(backend: Optional[BackendIndex]) -> str: +def static_dispatch_keys(backend: Optional[BackendIndex]) -> List[DispatchKey]: if backend is None: - return '' - return f""" -#include -#include -#include -""" + return [] + else: + return [ + backend.dispatch_key, + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeExplicitAutograd + ] + +def static_dispatch_extra_headers(backend: Optional[BackendIndex], skip_tensor_include: bool = False) -> str: + if skip_tensor_include: + # See Note [Avoiding Include Cycles In Static Dispatch] + maybe_inl = '_inl' + else: + maybe_inl = '' + return '\n'.join([ + f'#include ' for dispatch_key in static_dispatch_keys(backend)]) def static_dispatch( f: NativeFunction, cpp_sig: CppSignature, @@ -187,23 +197,7 @@ class RegisterSchema: def __call__(self, f: NativeFunction) -> Optional[str]: if not self.selector.is_native_function_selected(f): return None - schema_str = cpp_string(str(f.func)) - schema_str = '"' + "aten::" + schema_str[1:] - return f'm.def({schema_str});\n' - - -def _num_leading_spaces(line: str) -> int: - return len(line) - len(line.lstrip()) - - -# Unindents all lines in code. Each line gets unindented the same amount; -# that amount is equal to the smallest number of leading spaces across all lines -def deindent(code: str) -> str: - lines = code.split('\n') - min_leading_spaces = min(map(_num_leading_spaces, lines)) - lines = [line[min_leading_spaces:] for line in lines] - return '\n'.join(lines) - + return f'm.def({cpp_string(str(f.func))});\n' # Generates Operators.h and Operators.cpp. # These provide macros that, given an operator and overload name, allow users @@ -219,141 +213,119 @@ class ComputeOperators: @method_with_native_function def __call__(self, f: NativeFunction) -> Optional[str]: - # NB: requires_grad is the only exception to the rule because - # its const correctness is questionable. - if str(f.func.name) in set(['requires_grad_']): - return None + sig = DispatcherSignature.from_schema(f.func) + name = f.func.name.unambiguous_name() + call_method_name = 'call' + redispatch_method_name = 'redispatch' if self.target is Target.DECLARATION: - return self.gen_declaration(f) - if self.target is Target.DEFINITION: - return self.gen_definition(f) + # Note [The ATen Operators API] + # The ATen Operators API lives in the at::_ops namespace, and contains compile-time + # metadata about each operator + entry points into the Dispatcher. + # The C++ function, method, and redispatch API's are all implemented as wrappers + # into various bits of the structs defined here. + # + # Important characteristics about the Operators API: + # (1) It follows the Dispatcher API. + # This is kind of necessary to avoid overhead. + # For example: if it followed the C++ API, then all of the faithful C++ factory functions + # would need to wrap their arguments into TensorOptions only to unwrap them again. + # (2) Overload names are disambiguated. + # This is helpful for pytorch extenders who would like to decltype() an aten operator, + # that has overloads, e.g. decltype(at::_ops::mul_Tensor::call) + # (3) No argument defaulting is allowed. + # This is more of an implementation detail to avoid #include cycles, + # since TensorBody.h (which defines the Tensor class) needs to include this file. + # (4) manual_cpp_bindings and faithful names are not included in the API. + # This applies to stuff like __dispatch__is_complex(), and add_outf(). + # These aren't "real aten ops", they're just additional functions provided by the C++ API. + # They're implemented as wrappers in Functions.h that call into the actual operators + # defined here, i.e. at::_ops::is_complex::call() and at::_ops::add_out::call(). + # This means that ATEN_OP(is_complex) will not fastpath, and will go through the dispatcher. + return f""" +struct TORCH_API {name} {{ + using schema = {sig.type()}; + using ptr_schema = schema*; + static constexpr const char* name = "aten::{str(f.func.name.name)}"; + static constexpr const char* overload_name = "{f.func.name.overload_name}"; + static constexpr const char* schema_str = {cpp_string(str(f.func))}; + static {sig.defn(name=call_method_name, is_redispatching_fn=False)}; + static {sig.defn(name=redispatch_method_name, is_redispatching_fn=True)}; +}};""" + elif self.target is Target.DEFINITION: + defns = '' + for is_redispatching_fn in [False, True]: + if is_redispatching_fn: + dispatcher_exprs_str = ', '.join(['dispatchKeySet'] + [a.name for a in sig.arguments()]) + dispatcher_call = 'redispatch' + method_name = f'{name}::{redispatch_method_name}' + else: + dispatcher_exprs_str = ', '.join([a.name for a in sig.arguments()]) + dispatcher_call = 'call' + method_name = f'{name}::{call_method_name}' + + defns += f""" +// aten::{f.func} +{sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{ + static auto op = c10::Dispatcher::singleton() + .findSchemaOrThrow(name, overload_name) + .typed(); + return op.{dispatcher_call}({dispatcher_exprs_str}); +}} +""" + return defns else: assert_never(self.target) - # NB: This must be synchronized with the naming scheme in - # aten/src/ATen/templates/Operators.h - # Given a function schema "aten::op.overload(...)", - # If there is no overload name, this returns f"{op}" - # If there is an overload name, this returns f"{op}_{overload}" - def unambiguous_function_name(self, f: NativeFunction) -> str: - base_name = str(f.func.name.name) - overload_name = f.func.name.overload_name - if overload_name: - return f'{base_name}_{overload_name}' - return base_name - - def gen_declaration(self, f: NativeFunction) -> str: - unambiguous_name = self.unambiguous_function_name(f) - sig = DispatcherSignature.from_schema(f.func) - return f"TORCH_API {sig.decl(unambiguous_name)};" - - def most_faithful_name(self, f: NativeFunction) -> str: - sig_group = CppSignatureGroup.from_native_function(f, method=False) - sig = sig_group.most_faithful_signature() - return sig.name() - def invocation(self, f: NativeFunction) -> str: - faithful_op_name = self.most_faithful_name(f) - args = tuple(arg.name for arg in dispatcher.arguments(f.func)) - # Method only - if Variant.function not in f.variants: - return f"{args[0]}.{faithful_op_name}({', '.join(args[1:])})" - return f"at::{faithful_op_name}({', '.join(args)})" - - def gen_definition(self, f: NativeFunction) -> str: - unambiguous_name = self.unambiguous_function_name(f) - args = dispatcher.arguments(f.func) - sig = DispatcherSignature.from_schema(f.func) - - return deindent(f"""\ - {sig.defn(unambiguous_name)} {{ - return {self.invocation(f)}; - }}\ - """) - - -# Generates Function.cpp and Function.h. These files provide the -# functional public C++ API, and the scaffolding to call into -# the dispatcher from these functions. See also compute_tensor_method. +# Generates Function.h, which provides the functional public C++ API, +# and the scaffolding to call into the dispatcher from these functions. @dataclass(frozen=True) class ComputeFunction: - target: Union[ - Literal[Target.DECLARATION], - Literal[Target.DEFINITION] - ] static_dispatch_backend_index: Optional[BackendIndex] - is_redispatching_fn: bool @method_with_native_function def __call__(self, f: NativeFunction) -> Optional[str]: - # We unconditionally generate function variants of the redispatch API. - # This is mainly because we can namespace functions separately, but not methods, - if Variant.function not in f.variants and not self.is_redispatching_fn: + if Variant.function not in f.variants: return None - with native_function_manager(f): - return self.callImpl(f) - - def callImpl(self, f: NativeFunction) -> str: - name = cpp.name(f.func) - sig_group = CppSignatureGroup.from_native_function(f, method=False, fallback_binding=f.manual_cpp_binding) - if self.target is Target.DECLARATION: - sig_str = sig_group.signature.decl(is_redispatching_fn=self.is_redispatching_fn) - result = f"TORCH_API {sig_str};\n" - if sig_group.faithful_signature is not None: - sig_str = sig_group.faithful_signature.decl(is_redispatching_fn=self.is_redispatching_fn) - result += f"TORCH_API {sig_str};\n" - return result - - if self.target is not Target.DEFINITION: - assert_never(self.target) - def generate_defn(faithful: bool) -> str: - dispatcher_sig = DispatcherSignature.from_schema(f.func) - - if faithful and sig_group.faithful_signature is not None: + if faithful: sig = sig_group.faithful_signature + assert sig is not None else: sig = sig_group.signature - dispatcher_exprs = translate(sig.arguments(), dispatcher_sig.arguments()) - if self.is_redispatching_fn: - dispatcher_exprs_str = ', '.join(['dispatchKeySet'] + [a.expr for a in dispatcher_exprs]) - dispatcher_call = 'redispatch' - else: - dispatcher_exprs_str = ', '.join(a.expr for a in dispatcher_exprs) - dispatcher_call = 'call' + # See Note [The ATen Operators API] + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments()) + exprs_str = ', '.join([e.expr for e in exprs]) static_dispatch_block = static_dispatch(f, sig, method=False, backend_index=self.static_dispatch_backend_index) if static_dispatch_block is None: return f""" // aten::{f.func} -{sig.defn(is_redispatching_fn=self.is_redispatching_fn)} {{ - static auto op = c10::Dispatcher::singleton() - .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") - .typed<{dispatcher_sig.type()}>(); - return op.{dispatcher_call}({dispatcher_exprs_str}); +TORCH_API inline {sig.decl()} {{ + return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); }} """ else: return f""" // aten::{f.func} -{sig.defn(is_redispatching_fn=self.is_redispatching_fn)} {{ +TORCH_API inline {sig.decl()} {{ {static_dispatch_block} }} """ - result = generate_defn(sig_group.faithful_signature is None) + result = generate_defn(False) if sig_group.faithful_signature is not None: result += generate_defn(True) return result -# Generates TensorBody.h (sic) and TensorMethods.cpp. These files provide the -# object-oriented (method-based) public C++ API, and the scaffolding to call into -# the dispatcher from these functions. See also compute_function. +# Generates TensorBody.h. This file provides the object-oriented (method-based) +# public C++ API, and the scaffolding to call into the dispatcher from these functions. @dataclass(frozen=True) class ComputeTensorMethod: target: Union[ @@ -370,8 +342,6 @@ def __call__(self, f: NativeFunction) -> Optional[str]: assert not f.func.is_out_fn() assert f.func.arguments.self_arg is not None - name = cpp.name(f.func) - sig_group = CppSignatureGroup.from_native_function(f, method=True, fallback_binding=f.manual_cpp_binding) if self.target is Target.DECLARATION: @@ -384,32 +354,28 @@ def __call__(self, f: NativeFunction) -> Optional[str]: assert_never(self.target) def generate_defn(faithful: bool) -> str: - dispatcher_sig = DispatcherSignature.from_schema(f.func) - if faithful: sig = sig_group.faithful_signature assert sig is not None else: sig = sig_group.signature - dispatcher_exprs = translate(sig.arguments(), dispatcher_sig.arguments(), method=True) - dispatcher_exprs_str = ', '.join(a.expr for a in dispatcher_exprs) + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments(), method=True) + exprs_str = ', '.join([e.expr for e in exprs]) static_dispatch_block = static_dispatch(f, sig, method=True, backend_index=self.static_dispatch_backend_index) if static_dispatch_block is None: return f""" // aten::{f.func} -{sig.defn(prefix="Tensor::")} const {{ - static auto op = c10::Dispatcher::singleton() - .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") - .typed<{dispatcher_sig.type()}>(); - return op.call({dispatcher_exprs_str}); +inline {sig.defn(prefix="Tensor::")} const {{ + return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); }} """ else: return f""" // aten::{f.func} -{sig.defn(prefix="Tensor::")} const {{ +inline {sig.defn(prefix="Tensor::")} const {{ {static_dispatch_block} }} """ @@ -420,6 +386,42 @@ def generate_defn(faithful: bool) -> str: return result +# Generates RedispatchFunctions.h. +# This is similar to the C++ API defined in Functions.h, but provides access +# to the dispatcher's redispatch API. +@dataclass(frozen=True) +class ComputeRedispatchFunction: + + @method_with_native_function + def __call__(self, f: NativeFunction) -> Optional[str]: + # We unconditionally generate function variants of the redispatch API. + # This is mainly because we can namespace functions separately, but not methods, + sig_group = CppSignatureGroup.from_native_function(f, method=False, fallback_binding=f.manual_cpp_binding) + + def generate_defn(faithful: bool) -> str: + if faithful: + sig = sig_group.faithful_signature + assert sig is not None + else: + sig = sig_group.signature + + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments()) + exprs_str = ', '.join(['dispatchKeySet'] + [a.expr for a in exprs]) + + return f""" +// aten::{f.func} +TORCH_API inline {sig.decl(is_redispatching_fn=True)} {{ + return at::_ops::{f.func.name.unambiguous_name()}::redispatch({exprs_str}); +}} +""" + result = generate_defn(False) + if sig_group.faithful_signature is not None: + result += generate_defn(True) + + return result + + # Generates ATenOpList.cpp, a runtime accessible list of all aten # operators. # TODO: This was historically used to help some JIT interop code @@ -504,8 +506,8 @@ def __call__(self, f: NativeFunction) -> Optional[str]: C10_ALWAYS_INLINE {sig.defn(name)} {{ static auto op = c10::Dispatcher::singleton() - .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") - .typed<{dispatcher_sig.type()}>(); + .findSchemaOrThrow(at::_ops::{f.func.name.unambiguous_name()}::name, at::_ops::{f.func.name.unambiguous_name()}::overload_name) + .typed(); {compute_dk} return op.redispatch(_dk, {', '.join(a.expr for a in dispatcher_exprs)}); }} @@ -1050,7 +1052,17 @@ def make_file_manager(install_dir: str) -> FileManager: }) if dispatch_key in functions_keys: + if dispatch_key in static_dispatch_keys(static_dispatch_idx): + # See Note [Avoiding Include Cycles In Static Dispatch] + inl_headers = '' + else: + inl_headers = f'#include ' + fm.write_with_template(f'{dispatch_key}Functions.h', 'DispatchKeyFunctions.h', lambda: { + 'dispatch_key': str(dispatch_key), + 'inline_headers_for_nonstatic_build': inl_headers, + }) + fm.write_with_template(f'{dispatch_key}Functions_inl.h', 'DispatchKeyFunctions_inl.h', lambda: { 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_declarations': list(concatMap( dest.RegisterDispatchKey( @@ -1094,35 +1106,27 @@ def make_file_manager(install_dir: str) -> FileManager: }) cpu_fm.write('Functions.h', lambda: { - 'function_declarations': list(mapMaybe(ComputeFunction( - Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=False), native_functions)), - }) - cpu_fm.write('Functions.cpp', lambda: { 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), 'function_definitions': list(mapMaybe(ComputeFunction( - Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=False), native_functions)), - }) - cpu_fm.write('RedispatchFunctions.h', lambda: { - 'function_redispatch_declarations': list(mapMaybe(ComputeFunction( - Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=True), native_functions)), - }) - cpu_fm.write('RedispatchFunctions.cpp', lambda: { - 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), - 'function_redispatch_definitions': list(mapMaybe(ComputeFunction( - Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=True), native_functions)), + static_dispatch_backend_index=static_dispatch_idx), native_functions)), }) + core_fm.write('TensorBody.h', lambda: { - 'tensor_method_declarations': list(mapMaybe( - ComputeTensorMethod(Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), + 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx, skip_tensor_include=True), + 'tensor_method_declarations': list(mapMaybe(ComputeTensorMethod( + target=Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), + 'tensor_method_definitions': list(mapMaybe(ComputeTensorMethod( + target=Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), }) - core_fm.write('TensorMethods.cpp', lambda: { - 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), - 'tensor_method_definitions': list(mapMaybe( - ComputeTensorMethod(Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), + + cpu_fm.write('RedispatchFunctions.h', lambda: { + 'function_redispatch_definitions': list(mapMaybe(ComputeRedispatchFunction(), native_functions)), }) + core_fm.write('ATenOpList.cpp', lambda: { 'aten_ops': list(mapMaybe(compute_aten_op, native_functions)), }) + cpu_fm.write('NativeFunctions.h', lambda: { 'native_function_declarations': list(concatMap( # Convert to a set first to remove duplicate kernel names. diff --git a/tools/codegen/model.py b/tools/codegen/model.py index fba3756fab181..acde0827b0e6c 100644 --- a/tools/codegen/model.py +++ b/tools/codegen/model.py @@ -1454,6 +1454,18 @@ def __str__(self) -> str: else: return f"{self.name}" + # NB: This must be synchronized with the naming scheme in + # aten/src/ATen/templates/Operators.h + # Given a function schema "aten::op.overload(...)", + # If there is no overload name, this returns f"{op}" + # If there is an overload name, this returns f"{op}_{overload}" + def unambiguous_name(self) -> str: + if self.overload_name: + return f"{self.name}_{self.overload_name}" + else: + return f"{self.name}" + + def gets_generated_out_inplace_wrapper(f: NativeFunction, g: NativeFunctionsGroup, b: BackendIndex) -> bool: return f.func.kind() is not SchemaKind.functional and \ not b.has_kernel(f) and \ From 2062cafaa5ede56d63ecfc8b9edc2b69494f2247 Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Thu, 17 Jun 2021 13:28:24 -0700 Subject: [PATCH 209/305] [iOS GPU][MaskRCNN] Implement RoIAlign in Metal shaders using Sampler (#56075) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56075 Inspired by the CUDA implementation - https://fburl.com/diffusion/e90tabkj. The main difference is the way we implement bilinear interpolation. CUDA does this manually by iterating every point in each bin box. Whereas, Metal does this by calling sampler's sample function, which is a bit easier and faster. The result is almost identical to the result from CPU - P365102522. We'll do another round of refactor once we have figured out how to support custom ops on GPU. ghstack-source-id: 131720620 Test Plan: 1. Circle CI 2. Sandcastle Reviewed By: ajtulloch Differential Revision: D27485068 fbshipit-source-id: 31e831aead9d3799a3fde96e99dd677d96bd3da1 --- aten/src/ATen/native/metal/MetalAten.mm | 4 +- aten/src/ATen/native/metal/MetalShaders.h | 74 +++++++++++++++++++ aten/src/ATen/native/metal/MetalTensorImpl.h | 7 -- .../native/metal/MetalTensorImplStorage.h | 4 +- aten/src/ATen/native/metal/MetalUtils.h | 7 ++ .../native/metal/mpscnn/MPSImageWrapper.h | 6 +- .../native/metal/mpscnn/MPSImageWrapper.mm | 6 +- aten/src/ATen/native/metal/ops/MetalAddmm.mm | 12 +++ .../metal/ops/MetalBinaryElementwise.mm | 12 +++ .../ATen/native/metal/ops/MetalConvolution.mm | 11 ++- .../src/ATen/native/metal/ops/MetalNeurons.mm | 6 ++ .../src/ATen/native/metal/ops/MetalPooling.mm | 22 ++++-- .../src/ATen/native/metal/ops/MetalReshape.mm | 3 + .../src/ATen/native/metal/ops/MetalSoftmax.mm | 3 + .../metal/ops/MetalUpsamplingNearest.mm | 3 + 15 files changed, 152 insertions(+), 28 deletions(-) diff --git a/aten/src/ATen/native/metal/MetalAten.mm b/aten/src/ATen/native/metal/MetalAten.mm index 0300cc5a3718d..c9cee4092dd23 100644 --- a/aten/src/ATen/native/metal/MetalAten.mm +++ b/aten/src/ATen/native/metal/MetalAten.mm @@ -26,7 +26,9 @@ TORCH_INTERNAL_ASSERT( dst.is_contiguous(), "copy_from_metal is implemented only for contiguous output tensor"); - + if(dst.numel() == 0){ + return dst; + } MetalTensorImplStorage& tensorImplStorage = getTensorImplStorage(src); tensorImplStorage.copy_data_to_host(dst.data_ptr()); return dst; diff --git a/aten/src/ATen/native/metal/MetalShaders.h b/aten/src/ATen/native/metal/MetalShaders.h index bd56b22169b3f..25952b252e339 100644 --- a/aten/src/ATen/native/metal/MetalShaders.h +++ b/aten/src/ATen/native/metal/MetalShaders.h @@ -873,6 +873,80 @@ kernel void split_channels(texture2d_array in_arr[[texture(0 } } } + +constant bool ra_has_in_arr = (ushort_arg_3 > 1 || ushort_arg_2 > 4); +constant bool ra_has_out_arr = (ushort_arg_4 > 1 || ushort_arg_2 > 4); +constant bool ra_has_in_tex = (!ra_has_in_arr); +constant bool ra_has_out_tex = (!ra_has_out_arr); +kernel void roi_align(texture2d_array ina[[texture(0), function_constant(ra_has_in_arr)]], + texture2d in[[texture(0), function_constant(ra_has_in_tex)]], + texture2d_array outa[[texture(1), function_constant(ra_has_out_arr)]], + texture2d out[[texture(1), function_constant(ra_has_out_tex)]], + constant half4* rois[[buffer(0)]], + ushort3 gid[[thread_position_in_grid]]) { + + ushort out_width, out_height; + if (ra_has_out_arr) { + out_width = outa.get_width(); + out_height = outa.get_height(); + } else { + out_width = out.get_width(); + out_height = out.get_height(); + } + if (gid.x >= out_width || gid.y >= out_height) { + return; + } + const half spatial_scale = half(ushort_arg_0) / 10000; + const ushort sampling_ratio = ushort_arg_1; + const ushort C = ushort_arg_2; + const ushort pw = gid.x; + const ushort ph = gid.y; + const ushort n = gid.z / divRoundUp(C, 4); + const ushort c = gid.z % divRoundUp(C, 4); + + const half4 roi_scaled = rois[n] * spatial_scale; + const half roi_start_w = roi_scaled[0]; + const half roi_start_h = roi_scaled[1]; + const half roi_end_w = roi_scaled[2]; + const half roi_end_h = roi_scaled[3]; + + // Force malformed ROIs to be 1x1 + const half roi_width = max(roi_end_w - roi_start_w, (half)1.); + const half roi_height = max(roi_end_h - roi_start_h, (half)1.); + + const half bin_size_h = static_cast(roi_height) / static_cast(out_height); + const half bin_size_w = static_cast(roi_width) / static_cast(out_width); + + const ushort roi_bin_grid_h = sampling_ratio > 0 ? sampling_ratio : ceil(roi_height / static_cast(out_height)); + const ushort roi_bin_grid_w = sampling_ratio > 0 ? sampling_ratio : ceil(roi_width / static_cast(out_width)); + + const half count = roi_bin_grid_h * roi_bin_grid_w; + half4 output_val = 0.0; + + constexpr sampler s2(coord::pixel, address::clamp_to_edge, filter::linear); + + for (int iy = 0; iy < roi_bin_grid_h; iy++) { + for (int ix = 0; ix < roi_bin_grid_w; ix++) { + // Shift the pixel by 0.5. This is critical to achieve high accuracy. + const half y = + roi_start_h + ph * bin_size_h + (iy+0.5) * bin_size_h / static_cast(roi_bin_grid_h); + const half x = + roi_start_w + pw * bin_size_w + (ix+0.5) * bin_size_w / static_cast(roi_bin_grid_w); + if (ra_has_in_arr) { + output_val += ina.sample(s2, float2(x, y), c); + } else { + output_val += in.sample(s2, float2(x, y)); + } + } + } + output_val /= count; + if (ra_has_out_arr) { + outa.write(static_cast(output_val), gid.xy, gid.z); + } else { + out.write(static_cast(output_val), gid.xy); + } +} + )PT_METAL_SHADERS"; #endif /* MPSCNNShaders_h */ diff --git a/aten/src/ATen/native/metal/MetalTensorImpl.h b/aten/src/ATen/native/metal/MetalTensorImpl.h index 7e76390e841e9..865e466a8de75 100644 --- a/aten/src/ATen/native/metal/MetalTensorImpl.h +++ b/aten/src/ATen/native/metal/MetalTensorImpl.h @@ -40,13 +40,6 @@ struct TORCH_API MetalTensorImpl : public OpaqueTensorImpl { return strides_[d]; } - void release_resources() override { - using MetalTensorImplStorage = at::native::metal::MetalTensorImplStorage; - auto&& handle = (MetalTensorImplStorage)this->opaque_handle(); - handle.texture()->release(); - OpaqueTensorImpl::release_resources(); - } - private: const char* tensorimpl_type_name() const override { return "MetalTensorImpl"; diff --git a/aten/src/ATen/native/metal/MetalTensorImplStorage.h b/aten/src/ATen/native/metal/MetalTensorImplStorage.h index 5fbe429aeae5e..1ac7d126de95f 100644 --- a/aten/src/ATen/native/metal/MetalTensorImplStorage.h +++ b/aten/src/ATen/native/metal/MetalTensorImplStorage.h @@ -11,8 +11,8 @@ class MetalTensorImplStorage final { public: MetalTensorImplStorage(){}; - explicit MetalTensorImplStorage(const std::vector& sizes); - explicit MetalTensorImplStorage( + MetalTensorImplStorage(const std::vector& sizes); + MetalTensorImplStorage( const std::vector& sizes, const std::vector& strides); ~MetalTensorImplStorage() = default; diff --git a/aten/src/ATen/native/metal/MetalUtils.h b/aten/src/ATen/native/metal/MetalUtils.h index 662459113e4bd..e110da1bfcf4a 100644 --- a/aten/src/ATen/native/metal/MetalUtils.h +++ b/aten/src/ATen/native/metal/MetalUtils.h @@ -85,6 +85,13 @@ idmakeMTLBuffer(const std::vector& src) { return buffer; } +static inline idmakeMTLBuffer(int64_t bytes) { + id buffer = [[MPSCNNContext sharedInstance].device + newBufferWithLength:bytes + options:MTLResourceOptionCPUCacheModeWriteCombined]; + return buffer; +} + } // namespace metal } // namespace native } // namespace at diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.h b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.h index f4c310ea46335..bba2a525429a6 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.h +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.h @@ -30,10 +30,10 @@ class API_AVAILABLE(ios(10.0), macos(10.13)) MPSImageWrapper { private: std::vector _imageSizes; - MPSImage* _image = nullptr; + MPSImage* _image = nil; id _buffer = nil; - MetalCommandBuffer* _commandBuffer; - id _delegate; + MetalCommandBuffer* _commandBuffer = nil; + id _delegate = nil; }; } // namespace metal diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm index fd02432c75a27..c5931b2870fd7 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm @@ -131,10 +131,8 @@ - (void)endSynchronization:(NSError*)error { } void MPSImageWrapper::release() { - if ([_image isTemporaryImage]) { - [_image recycle]; - [_commandBuffer remove:(MPSTemporaryImage*)_image]; - } + [_image recycle]; + [_commandBuffer remove:(MPSTemporaryImage*)_image]; [_commandBuffer removeSubscriber:_delegate]; _delegate = nil; _commandBuffer = nil; diff --git a/aten/src/ATen/native/metal/ops/MetalAddmm.mm b/aten/src/ATen/native/metal/ops/MetalAddmm.mm index 85c9b587c20a4..768a721451270 100644 --- a/aten/src/ATen/native/metal/ops/MetalAddmm.mm +++ b/aten/src/ATen/native/metal/ops/MetalAddmm.mm @@ -28,6 +28,9 @@ Tensor addmm( TORCH_CHECK(bias.device() == kCPU); TORCH_CHECK(beta.toFloat() == 1.0f); TORCH_CHECK(alpha.toFloat() == 1.0f); + if(input.numel() == 0 || weight.numel() == 0){ + return makeTensor({{input.size(0), weight.size(0)}}, input.options()); + } // Here we treat the matrix multiplication as convolution auto weight_ = weight.t().view({weight.size(1), weight.size(0), 1, 1}).contiguous(); @@ -64,6 +67,9 @@ Tensor linear(const Tensor& input, LinearOpContext& context) { TORCH_CHECK(input.is_metal()); TORCH_CHECK(context.get_weight().device() == kCPU); TORCH_CHECK(context.get_weight().dim() == 4); + if(input.numel() == 0 || context.get_weight().numel() == 0){ + return makeTensor({{input.size(0), context.get_weight().size(0)}}, input.options()); + } // Reshape the input tensor to {N, C, 1, 1} auto input_ = input.view({input.size(0), input.size(1), 1, 1}); MPSImage* X = imageFromTensor(input_); @@ -98,6 +104,12 @@ Tensor linear(const Tensor& input, LinearOpContext& context) { MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input_); mt.texture()->allocateTemporaryStorage(textureSize, commandBuffer); MPSImage* Y1 = mt.texture()->image(); + // HACK alert: + // Here we force X to become static before encoding. + // We've seen weird crashes in the MaskRCNN model complaining about + // a "sub-image" was released before its readCount was zero. + // TODO[T93395421]: Figure out the root cause and remove this line. + X = createStaticImage((MPSTemporaryImage* )X, commandBuffer, NO); [op encode:commandBuffer.buffer sourceImage:X destinationImage:Y1]; if (nt == NeuronType::Clamp) { MPSImage* Y2 = createTemporaryImage(commandBuffer, [Y1 sizes]); diff --git a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm index c28f22310e3bb..ad611679b9881 100644 --- a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm +++ b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm @@ -72,6 +72,9 @@ Tensor binaryElementwiseShaderKernel( if (broadCastFirstInput(X1, X2)) { outputSize = input2.sizes(); } + if(c10::multiply_integers(outputSize) == 0){ + return makeTensor({outputSize.vec()}, input1.options()); + } MetalTensorImplStorage mt{outputSize.vec()}; MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); @@ -111,6 +114,9 @@ Tensor binaryElementwiseShaderKernel( if (broadCastFirstInput(X1, X2)) { outputSize = input2.sizes(); } + if(c10::multiply_integers(outputSize) == 0){ + return input1; + } MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); TORCH_CHECK( @@ -149,6 +155,9 @@ Tensor binaryElementwiseMPSCNNKernel( if (broadCastFirstInput(X1, X2)) { outputSize = input2.sizes(); } + if(c10::multiply_integers(outputSize) == 0){ + return makeTensor({outputSize.vec()}, input1.options()); + } MetalTensorImplStorage mt{outputSize.vec()}; MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); @@ -180,6 +189,9 @@ Tensor binaryElementwiseMPSCNNKernel( if (broadCastFirstInput(X1, X2)) { outputSize = input2.sizes(); } + if(c10::multiply_integers(outputSize) == 0){ + return input1; + } MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); TORCH_CHECK( diff --git a/aten/src/ATen/native/metal/ops/MetalConvolution.mm b/aten/src/ATen/native/metal/ops/MetalConvolution.mm index b7b02d9c19f76..98fc87e84be73 100644 --- a/aten/src/ATen/native/metal/ops/MetalConvolution.mm +++ b/aten/src/ATen/native/metal/ops/MetalConvolution.mm @@ -28,6 +28,10 @@ Tensor conv2d( TORCH_INTERNAL_ASSERT(input.dim() == 4, "Expected 4-dimensional input"); TORCH_INTERNAL_ASSERT(weight.dim() == 4, "Expected 4-dimensional weight"); TORCH_CHECK(weight.device().type() == kCPU); + auto outputSize = params.output_sizes(); + if(c10::multiply_integers(outputSize) == 0){ + return makeTensor({outputSize}, input.options()); + } MPSImage* X = imageFromTensor(input); auto packedWeights = weight.contiguous(c10::MemoryFormat::ChannelsLast); // MPSCNN Convolution @@ -37,7 +41,6 @@ Tensor conv2d( weights:w bias:b neuronFilter:NeuronType::None]; - auto outputSize = params.output_sizes(); MetalTensorImplStorage mt{outputSize}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); @@ -57,6 +60,10 @@ Tensor conv2d(const Tensor& input, Conv2dOpContext& context) { context.stride, context.dilation, context.groups}; + auto outputSize = params.output_sizes(); + if(c10::multiply_integers(outputSize) == 0){ + return makeTensor({outputSize}, input.options()); + } MPSCNNConvOp* op = (__bridge MPSCNNConvOp*)(context.conv2dOp); NeuronType nt = neuronType(context.output_min, context.output_max); if (!op) { @@ -71,8 +78,6 @@ Tensor conv2d(const Tensor& input, Conv2dOpContext& context) { } }; } - - auto outputSize = params.output_sizes(); MetalTensorImplStorage mt{outputSize}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); diff --git a/aten/src/ATen/native/metal/ops/MetalNeurons.mm b/aten/src/ATen/native/metal/ops/MetalNeurons.mm index 486823471bf13..fa72db1afcdb0 100644 --- a/aten/src/ATen/native/metal/ops/MetalNeurons.mm +++ b/aten/src/ATen/native/metal/ops/MetalNeurons.mm @@ -18,6 +18,9 @@ Tensor neuronKernel(const Tensor& input, MPSCNNNeuron* neuron) { MPSImage* X = imageFromTensor(input); IntArrayRef outputSize = input.sizes(); + if(input.numel() == 0){ + return makeTensor({outputSize.vec()}, input.options()); + } IntArrayRef textureSize = outputSize; MetalTensorImplStorage mt{outputSize.vec()}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); @@ -33,6 +36,9 @@ Tensor neuronKernel(const Tensor& input, MPSCNNNeuron* neuron) { Tensor& neuronKernel_(Tensor& input, MPSCNNNeuron* neuron) { MPSImage* X = imageFromTensor(input); IntArrayRef outputSize = input.sizes(); + if(input.numel() == 0){ + return input; + } IntArrayRef textureSize = outputSize; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); MPSImage* Y = createTemporaryImage(commandBuffer, textureSize); diff --git a/aten/src/ATen/native/metal/ops/MetalPooling.mm b/aten/src/ATen/native/metal/ops/MetalPooling.mm index 0e237a91293b1..35c6c34ffc6dd 100644 --- a/aten/src/ATen/native/metal/ops/MetalPooling.mm +++ b/aten/src/ATen/native/metal/ops/MetalPooling.mm @@ -39,6 +39,14 @@ Tensor max_pool2d( const int64_t pW = padding[1]; const int64_t dH = dilation[0]; const int64_t dW = dilation[1]; + int64_t oN = iN; + int64_t oC = iC; + int64_t oH = pooling_output_shape(iH, kH, pH, sH, dH, ceil_mode); + int64_t oW = pooling_output_shape(iW, kW, pW, sW, dW, ceil_mode); + SmallVectoroutputSize{oN, oC, oH, oW}; + if(input.numel() == 0){ + return makeTensor({IntArrayRef(outputSize).vec()}, input.options()); + } MPSImage* X = imageFromTensor(input); MPSCNNPoolingMax* pool = [[MPSCNNPoolingMax alloc] initWithDevice:[MPSCNNContext sharedInstance].device @@ -51,12 +59,6 @@ Tensor max_pool2d( setOffset:{.x = mpscnn::computeMPSAlignOffset(kernel_size[0], padding[0]), .y = mpscnn::computeMPSAlignOffset(kernel_size[1], padding[1]), .z = 0}]; - int64_t oN = iN; - int64_t oC = iC; - int64_t oH = pooling_output_shape(iH, kH, pH, sH, dH, ceil_mode); - int64_t oW = pooling_output_shape(iW, kW, pW, sW, dW, ceil_mode); - - SmallVectoroutputSize{oN, oC, oH, oW}; MetalTensorImplStorage mt{IntArrayRef(outputSize).vec()}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); @@ -73,6 +75,11 @@ Tensor adaptive_avg_pool2d(const Tensor& input, IntArrayRef output_size) { // averages across the width and height, and outputs a 1x1xC image. TORCH_CHECK(output_size[0] == 1 && output_size[1] == 1); TORCH_CHECK(input.is_metal()); + SmallVector outputSize{ + input.sizes()[0], input.sizes()[1], output_size[0], output_size[1]}; + if(input.numel() == 0){ + return makeTensor({IntArrayRef(outputSize).vec()}, input.options()); + } MPSImage* X = imageFromTensor(input); MPSCNNPoolingAverage* pool = [[MPSCNNPoolingAverage alloc] initWithDevice:[MPSCNNContext sharedInstance].device @@ -84,8 +91,7 @@ Tensor adaptive_avg_pool2d(const Tensor& input, IntArrayRef output_size) { [pool setOffset:{.x = static_cast(X.width / 2), .y = static_cast(X.height / 2), .z = 0}]; - SmallVector outputSize{ - input.sizes()[0], input.sizes()[1], output_size[0], output_size[1]}; + MetalTensorImplStorage mt{IntArrayRef(outputSize).vec()}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); diff --git a/aten/src/ATen/native/metal/ops/MetalReshape.mm b/aten/src/ATen/native/metal/ops/MetalReshape.mm index 817b81e85a538..15ae3c50d60e4 100644 --- a/aten/src/ATen/native/metal/ops/MetalReshape.mm +++ b/aten/src/ATen/native/metal/ops/MetalReshape.mm @@ -27,6 +27,9 @@ Tensor view(const Tensor& input, IntArrayRef size) { "not compatible with input tensor's size and stride (at least one dimension" " spans across two contiguous subspaces). Use .reshape(...) instead."); auto stride_value = *stride; + if(input.numel() == 0) { + return makeTensor({inferred_size, stride_value}, input.options()); + } MPSImage* X = imageFromTensor(input); MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); MetalTensorImplStorage mt{inferred_size, stride_value}; diff --git a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm index d27e5980297c0..f3e1ef5e362fe 100644 --- a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm +++ b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm @@ -22,6 +22,9 @@ Tensor mpscnn_softmax( TORCH_CHECK(input.is_metal()); // TODO: [T87180544] Implment softmax/log_softmax in metal shaders TORCH_CHECK(input.dim() == 2); + if(input.numel() == 0){ + return makeTensor({input.sizes().vec()}, input.options()); + } std::vector newSize(4, 1); if (dim == 0) { newSize[1] = input.size(0); diff --git a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm index 359cfb63e0828..25138eddddae8 100644 --- a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm +++ b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm @@ -41,6 +41,9 @@ Tensor upsample_nearest2d_vec( output_width); std::vector outputSizes{ nbatch, channels, output_height, output_width}; + if(input.numel() == 0){ + return makeTensor({outputSizes}, input.options()); + } MPSImage* X = imageFromTensor(input); MetalTensorImplStorage mt{outputSizes}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); From acd914f03909a70631ecadde121f8a771876cd9f Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Thu, 17 Jun 2021 13:38:03 -0700 Subject: [PATCH 210/305] Fix Pipe + DDP for unused parameters, static graph (#60118) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60118 Pipe + DDP has a few issues: 1) with static graph, does not synchronize gradients on first backward pass (i.e. delay allreduce is not run). does not work since https://github.com/pytorch/pytorch/pull/55248 2) when find_unused_parameters=True, also does not results in gradient synchronization. does not work since https://github.com/pytorch/pytorch/pull/57081 The reason for both cases is that calling `DDPSink.apply(output_tensor)` does not call the custom `backward` of `DDPSink` when the `output_tensor` is actually an `OwnerRRef`, which is the case when running DDP in `Pipe`. This is because we do `backward` on the `rref.local_value()` which does not have this autograd recording. To fix, we unwrap the RRef and reconstruct it as needed, similar to the fix in https://github.com/pytorch/pytorch/pull/49908. to test: All tests in pipe_with_ddp_test pass. The reason these tests did not catch the errors earlier is because all ranks received the same model inputs. So if gradient synchronization did not occur, then grads would still be the same because the model is the same on all ranks (guaranteed by ddp). Fixed the tests to use different inputs across ranks. ghstack-source-id: 131688187 Test Plan: CI Reviewed By: pritamdamania87 Differential Revision: D29167283 fbshipit-source-id: fe62310db2dc6de8519eb361b1df8ae4dfce3ab8 --- torch/nn/parallel/distributed.py | 25 +++++++++++++++++-- .../distributed/pipe_with_ddp_test.py | 18 +++++++++++-- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py index b46cae2e742d3..e69dcc9006ac1 100644 --- a/torch/nn/parallel/distributed.py +++ b/torch/nn/parallel/distributed.py @@ -25,6 +25,22 @@ from ._functions import _get_stream from .scatter_gather import scatter_kwargs, gather, is_namedtuple +def _tree_flatten_with_rref(output): + output_is_rref = RPC_AVAILABLE and isinstance(output, RRef) + if output_is_rref: + output_tensor_list, treespec = tree_flatten(output.local_value()) + else: + output_tensor_list, treespec = tree_flatten(output) + # Need to return flattened tensors, spec to re-pack them, as well + # as if the return type was actually an RRef to reconstruct. + return output_tensor_list, treespec, output_is_rref + +def _tree_unflatten_with_rref(output, treespec, output_is_rref): + output = tree_unflatten(output, treespec) + if output_is_rref: + output = RRef(output) + return output + def _find_tensors(obj): r""" @@ -863,14 +879,19 @@ def forward(self, *inputs, **kwargs): 'find_unused': find_unused, 'num_iterations': self.num_iterations, } - output_tensor_list, treespec = tree_flatten(output) + + output_tensor_list, treespec, output_is_rref = _tree_flatten_with_rref( + output + ) passthrough_tensor_list = _DDPSink.apply( self.reducer, state_dict, *output_tensor_list, ) # Reconstruct output data structure. - output = tree_unflatten(passthrough_tensor_list, treespec) + output = _tree_unflatten_with_rref( + passthrough_tensor_list, treespec, output_is_rref + ) return output def scatter(self, inputs, kwargs, device_ids): diff --git a/torch/testing/_internal/distributed/pipe_with_ddp_test.py b/torch/testing/_internal/distributed/pipe_with_ddp_test.py index a17e880cef0fe..81bc262e497d8 100644 --- a/torch/testing/_internal/distributed/pipe_with_ddp_test.py +++ b/torch/testing/_internal/distributed/pipe_with_ddp_test.py @@ -109,12 +109,26 @@ def forward(self, inp): model = DistributedDataParallel(model, find_unused_parameters=find_unused_parameters) if static_graph: model._set_static_graph() - out = model(torch.rand(16, 16).cuda(2 * self.rank)).local_value() + + # Ensure inputs are different across ranks to verify that gradient + # sync indeed occurs. + model_input = torch.rand(16, 16).cuda(2 * self.rank) * (self.rank + 1) + out = model(model_input).local_value() out.sum().backward() # Run forward again for find_unused_parameters to trigger any potential errors. if find_unused_parameters: - model(torch.rand(16, 16).cuda(2 * self.rank)) + # Ensure inputs are different across ranks to verify that gradient + # sync indeed occurs. + unused_param_input = torch.rand(16, 16).cuda(2 * self.rank) * (self.rank + 1) + model(unused_param_input).local_value().sum().backward() + + # Run a few more iterations of fwd + bwd to ensure gradient synchronization + # occurs properly across iterations via delay_all_reduce/bucketized allreduce. + for _ in range(3): + model_input = torch.rand(16, 16).cuda(2 * self.rank) * (self.rank + 1) + out = model(model_input).local_value() + out.sum().backward() # Check grads output = [torch.empty_like(fc1.weight.grad), torch.empty_like(fc1.weight.grad)] From 0cbb5e15d75c803a629dd35f1a1f62ba9dd81e44 Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Thu, 17 Jun 2021 13:38:03 -0700 Subject: [PATCH 211/305] Correct backend in pipe_with_ddp_test (#60123) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60123 All of the tests would run with gloo, but some tests specify a different backend param which we should respect. ghstack-source-id: 131688188 Test Plan: CI Reviewed By: cbalioglu Differential Revision: D29171549 fbshipit-source-id: 3e306060df189c0e38d5ca6dd34f4b4fbca052b9 --- torch/testing/_internal/distributed/pipe_with_ddp_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/testing/_internal/distributed/pipe_with_ddp_test.py b/torch/testing/_internal/distributed/pipe_with_ddp_test.py index 81bc262e497d8..a89db18188f5c 100644 --- a/torch/testing/_internal/distributed/pipe_with_ddp_test.py +++ b/torch/testing/_internal/distributed/pipe_with_ddp_test.py @@ -78,7 +78,7 @@ def test_basic_gloo_ckpt_except_last(self): def _run_basic_test(self, backend, checkpoint, find_unused_parameters=False, static_graph=False): dist.init_process_group( - backend="nccl", + backend=backend, init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), world_size=self.world_size, rank=self.rank, From 6d0fb85a623f5ef3f3f1a2afc3660cb71fa70511 Mon Sep 17 00:00:00 2001 From: Brian Hirsh Date: Thu, 17 Jun 2021 14:26:51 -0700 Subject: [PATCH 212/305] Revert D28833086: beef up at::_ops API Test Plan: revert-hammer Differential Revision: D28833086 (https://github.com/pytorch/pytorch/commit/e2129d1c067326efba4eac53255b94af05a45b1b) Original commit changeset: 55f322a8378c fbshipit-source-id: e55bf812ec411bb6bee87654f1d65ff10c046106 --- BUILD.bazel | 8 +- aten/src/ATen/core/op_registration/adaption.h | 20 ++ .../src/ATen/templates/DispatchKeyFunctions.h | 37 +-- .../ATen/templates/DispatchKeyFunctions_inl.h | 16 - aten/src/ATen/templates/Functions.cpp | 160 +++++++++ aten/src/ATen/templates/Functions.h | 121 ++----- aten/src/ATen/templates/Operators.cpp | 2 - aten/src/ATen/templates/Operators.h | 85 +---- aten/src/ATen/templates/RedispatchFunctions.h | 2 +- .../ATen/templates/RegisterBackendSelect.cpp | 1 - aten/src/ATen/templates/RegisterSchema.cpp | 1 - aten/src/ATen/templates/TensorBody.h | 75 +---- aten/src/ATen/templates/TensorMethods.cpp | 85 +++++ caffe2/contrib/aten/gen_op.py | 8 +- tools/codegen/api/python.py | 22 +- tools/codegen/api/types.py | 38 +-- tools/codegen/gen.py | 312 +++++++++--------- tools/codegen/model.py | 12 - 18 files changed, 510 insertions(+), 495 deletions(-) delete mode 100644 aten/src/ATen/templates/DispatchKeyFunctions_inl.h create mode 100644 aten/src/ATen/templates/Functions.cpp create mode 100644 aten/src/ATen/templates/TensorMethods.cpp diff --git a/BUILD.bazel b/BUILD.bazel index 217b20fb54b9d..b7e16ac1c915c 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -136,22 +136,20 @@ genrule( "aten/src/ATen/RegisterMeta.cpp", "aten/src/ATen/RegisterSchema.cpp", "aten/src/ATen/CPUFunctions.h", - "aten/src/ATen/CPUFunctions_inl.h", "aten/src/ATen/CUDAFunctions.h", - "aten/src/ATen/CUDAFunctions_inl.h", "aten/src/ATen/CompositeExplicitAutogradFunctions.h", - "aten/src/ATen/CompositeExplicitAutogradFunctions_inl.h", "aten/src/ATen/CompositeImplicitAutogradFunctions.h", - "aten/src/ATen/CompositeImplicitAutogradFunctions_inl.h", "aten/src/ATen/Functions.h", + "aten/src/ATen/Functions.cpp", "aten/src/ATen/RedispatchFunctions.h", + "aten/src/ATen/RedispatchFunctions.cpp", "aten/src/ATen/Operators.h", "aten/src/ATen/Operators.cpp", "aten/src/ATen/NativeFunctions.h", "aten/src/ATen/MetaFunctions.h", - "aten/src/ATen/MetaFunctions_inl.h", "aten/src/ATen/NativeMetaFunctions.h", "aten/src/ATen/core/TensorBody.h", + "aten/src/ATen/core/TensorMethods.cpp", "aten/src/ATen/core/ATenOpList.cpp", ], cmd = "$(location :gen) --source-path aten/src/ATen --install_dir `dirname $(location aten/src/ATen/Declarations.yaml)`", diff --git a/aten/src/ATen/core/op_registration/adaption.h b/aten/src/ATen/core/op_registration/adaption.h index 5bf1b691ebad3..327c1d10e5a28 100644 --- a/aten/src/ATen/core/op_registration/adaption.h +++ b/aten/src/ATen/core/op_registration/adaption.h @@ -43,6 +43,26 @@ namespace c10 { namespace impl { +inline c10::optional +check_tensor_options_and_extract_memory_format( + const TensorOptions& options, + c10::optional memory_format) { + TORCH_CHECK( + options.requires_grad_opt() == c10::nullopt || + options.requires_grad_opt().value() == false, + "Operators taking TensorOptions cannot take a TensorOptions with " + "options.requires_grad set as true. This isn't implemented yet."); + TORCH_CHECK( + !(options.has_memory_format() && memory_format.has_value()), + "Cannot set memory_format both in TensorOptions and explicit argument; please delete " + "the redundant setter."); + if (memory_format.has_value()) { + return memory_format; + } else { + return options.memory_format_opt(); + } +} + TORCH_API void common_device_check_failure(optional& common_device, const at::Tensor& tensor, at::CheckedFrom methodName, at::CheckedFrom argName); inline void check_and_update_common_device(optional& common_device, const at::Tensor& tensor, at::CheckedFrom methodName, at::CheckedFrom argName) { diff --git a/aten/src/ATen/templates/DispatchKeyFunctions.h b/aten/src/ATen/templates/DispatchKeyFunctions.h index 1718b4be8274c..e72b39b5ae5f0 100644 --- a/aten/src/ATen/templates/DispatchKeyFunctions.h +++ b/aten/src/ATen/templates/DispatchKeyFunctions.h @@ -1,23 +1,14 @@ -#include -// Note [Avoiding Include Cycles In Static Dispatch] -// In order to avoid #include cycles in the static dispatch build, we've carefully split out -// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h. -// -// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h. -// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods -// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all -// directly inlined into TensorBody.h. -// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API, -// which include functions that have defaultable optional arguments. -// That requires knowing the full Tensor class definition. -// -// We break the cycle by doing the following: -// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h -// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl., -// - CPUFunctions_inl.h includes everything else -// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class, -// and then it includes CPUFunctions_inl.h. -// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly. -// - This also means that static dispatch build, CPUFunctions.h only needs to -// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h. -${inline_headers_for_nonstatic_build} +// ${generated_comment} + +// NB: The implementing C++ file is RegisterDispatchKey.cpp + +// TODO: tighten this include +#include + +namespace at { +namespace ${dispatch_namespace} { + +${dispatch_namespaced_declarations} + +} // namespace ${dispatch_namespace} +} // namespace at diff --git a/aten/src/ATen/templates/DispatchKeyFunctions_inl.h b/aten/src/ATen/templates/DispatchKeyFunctions_inl.h deleted file mode 100644 index 365ce8b98d614..0000000000000 --- a/aten/src/ATen/templates/DispatchKeyFunctions_inl.h +++ /dev/null @@ -1,16 +0,0 @@ -// ${generated_comment} - -// NB: The implementing C++ file is RegisterDispatchKey.cpp - -// The only #includes we need are for custom classes that have defaults in the C++ API -#include -#include -#include - -namespace at { -namespace ${dispatch_namespace} { - -${dispatch_namespaced_declarations} - -} // namespace ${dispatch_namespace} -} // namespace at diff --git a/aten/src/ATen/templates/Functions.cpp b/aten/src/ATen/templates/Functions.cpp new file mode 100644 index 0000000000000..3d119b6314abb --- /dev/null +++ b/aten/src/ATen/templates/Functions.cpp @@ -0,0 +1,160 @@ +// ${generated_comment} + +#include + +#include +#include + +#include +#include + +${static_dispatch_extra_headers} + +namespace at { + +Tensor var(const Tensor& self, int dim) { + return at::var(self, IntArrayRef{dim}); +} + +std::tuple var_mean(const Tensor& self, int dim) { + return at::var_mean(self, IntArrayRef{dim}); +} + +Tensor std(const Tensor& self, int dim) { + return at::std(self, IntArrayRef{dim}); +} + +std::tuple std_mean(const Tensor& self, int dim) { + return at::std_mean(self, IntArrayRef{dim}); +} + +at::Tensor conv1d( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + std::initializer_list padding_, + IntArrayRef dilation, + int64_t groups) { + auto padding = IntArrayRef(padding_); + return at::conv1d(input, weight, bias, stride, padding, dilation, groups); +} + +at::Tensor conv2d( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + std::initializer_list padding_, + IntArrayRef dilation, + int64_t groups) { + auto padding = IntArrayRef(padding_); + return at::conv2d(input, weight, bias, stride, padding, dilation, groups); +} + +at::Tensor conv3d( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + std::initializer_list padding_, + IntArrayRef dilation, + int64_t groups) { + auto padding = IntArrayRef(padding_); + return at::conv3d(input, weight, bias, stride, padding, dilation, groups); +} + +namespace detail { + +void noopDelete(void*) {} + +} // namespace detail + +Tensor TensorMaker::make_tensor() { + AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove. + tracer::impl::NoTracerDispatchMode tracer_guard{}; + + check_size_nonnegative(sizes_); + + TORCH_CHECK_VALUE( + !deleter_ || !ctx_, + "The deleter and context arguments are mutually exclusive."); + + if (device_ == nullopt) { + device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type()); + } + + if (opts_.device().has_index()) { + // clang-format off + TORCH_CHECK_VALUE( + opts_.device() == *device_, + "Specified device ", opts_.device(), " does not match device of data ", *device_); + // clang-format on + } + + std::size_t size_bytes = computeStorageSize(); + + DataPtr data_ptr{}; + if (deleter_) { + data_ptr = makeDataPtrFromDeleter(); + } else { + data_ptr = makeDataPtrFromContext(); + } + + Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr)}; + + Tensor tensor = detail::make_tensor( + std::move(storage), opts_.computeDispatchKey(), opts_.dtype()); + + if (sizes_.size() != 1 || sizes_[0] != 0) { + TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl(); + + if (strides_) { + tensor_impl->set_sizes_and_strides(sizes_, *strides_); + } else { + tensor_impl->set_sizes_contiguous(sizes_); + } + } + + return tensor; +} + +std::size_t TensorMaker::computeStorageSize() const noexcept { + std::size_t itemsize = opts_.dtype().itemsize(); + + if (strides_) { + return detail::computeStorageNbytes(sizes_, *strides_, itemsize); + } + + std::size_t size = 1; + for (std::int64_t s : sizes_) { + size *= static_cast(s); + } + return size * itemsize; +} + +inline DataPtr TensorMaker::makeDataPtrFromDeleter() const { + return InefficientStdFunctionContext::makeDataPtr(data_, deleter_, *device_); +} + +inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept { + return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_}; +} + +IntArrayRef TensorMaker::makeTempSizes() const noexcept { + static std::int64_t zeros[5] = {0, 0, 0, 0, 0}; + if (opts_.has_memory_format()) { + MemoryFormat format = *opts_.memory_format_opt(); + if (format == MemoryFormat::ChannelsLast) { + return IntArrayRef(zeros, 4); + } + if (format == MemoryFormat::ChannelsLast3d) { + return IntArrayRef(zeros, 5); + } + } + return IntArrayRef(zeros, 1); +} + +${function_definitions} + +} // namespace at diff --git a/aten/src/ATen/templates/Functions.h b/aten/src/ATen/templates/Functions.h index caeb328e0c5bd..18e7c00a0bb0a 100644 --- a/aten/src/ATen/templates/Functions.h +++ b/aten/src/ATen/templates/Functions.h @@ -14,9 +14,6 @@ #include #include #include -#include - -${static_dispatch_extra_headers} namespace at { @@ -43,27 +40,32 @@ AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR) AT_FORALL_COMPLEX_TYPES(TENSOR) #undef TENSOR -${function_definitions} +${function_declarations} // Special C++ only overloads for std()-like functions (See gh-40287) // These are needed because int -> bool conversion takes precedence over int -> IntArrayRef // So, for example std(0) would select the std(unbiased=False) overload -TORCH_API inline Tensor var(const Tensor& self, int dim) { - return at::var(self, IntArrayRef{dim}); -} -TORCH_API inline std::tuple var_mean(const Tensor& self, int dim) { - return at::var_mean(self, IntArrayRef{dim}); -} -TORCH_API inline Tensor std(const Tensor& self, int dim) { - return at::std(self, IntArrayRef{dim}); -} -TORCH_API inline std::tuple std_mean(const Tensor& self, int dim) { - return at::std_mean(self, IntArrayRef{dim}); -} +TORCH_API Tensor var(const Tensor& self, int dim); +TORCH_API std::tuple var_mean(const Tensor& self, int dim); +TORCH_API Tensor std(const Tensor& self, int dim); +TORCH_API std::tuple std_mean(const Tensor& self, int dim); + + +// Special C++ only overloads for convnd functions (See gh-45667) +// These are needed because {1, 2} is ambiguous between string and IntArrayRef overloads +TORCH_API at::Tensor conv1d( + const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, + std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); +TORCH_API at::Tensor conv2d( + const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, + std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); +TORCH_API at::Tensor conv3d( + const Tensor& input, const Tensor& weight, const Tensor& bias, IntArrayRef stride, + std::initializer_list padding, IntArrayRef dilation = 1, int64_t groups = 1); namespace detail { -TORCH_API inline void noopDelete(void*) {} +TORCH_API void noopDelete(void*); } // namespace detail @@ -115,94 +117,19 @@ class TORCH_API TensorMaker { return *this; } - Tensor make_tensor() { - AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove. - tracer::impl::NoTracerDispatchMode tracer_guard{}; - - check_size_nonnegative(sizes_); - - TORCH_CHECK_VALUE( - !deleter_ || !ctx_, - "The deleter and context arguments are mutually exclusive."); - - if (device_ == nullopt) { - device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type()); - } - - if (opts_.device().has_index()) { - // clang-format off - TORCH_CHECK_VALUE( - opts_.device() == *device_, - "Specified device ", opts_.device(), " does not match device of data ", *device_); - // clang-format on - } - - std::size_t size_bytes = computeStorageSize(); - - DataPtr data_ptr{}; - if (deleter_) { - data_ptr = makeDataPtrFromDeleter(); - } else { - data_ptr = makeDataPtrFromContext(); - } - - Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr)}; - - Tensor tensor = detail::make_tensor( - std::move(storage), opts_.computeDispatchKey(), opts_.dtype()); - - if (sizes_.size() != 1 || sizes_[0] != 0) { - TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl(); - - if (strides_) { - tensor_impl->set_sizes_and_strides(sizes_, *strides_); - } else { - tensor_impl->set_sizes_contiguous(sizes_); - } - } - - return tensor; - } + Tensor make_tensor(); private: explicit TensorMaker(void* data, IntArrayRef sizes) noexcept : data_{data}, sizes_{sizes} {} - std::size_t computeStorageSize() const noexcept { - std::size_t itemsize = opts_.dtype().itemsize(); + std::size_t computeStorageSize() const noexcept; - if (strides_) { - return detail::computeStorageNbytes(sizes_, *strides_, itemsize); - } + DataPtr makeDataPtrFromDeleter() const; - std::size_t size = 1; - for (std::int64_t s : sizes_) { - size *= static_cast(s); - } - return size * itemsize; - } + DataPtr makeDataPtrFromContext() noexcept; - inline DataPtr makeDataPtrFromDeleter() const { - return InefficientStdFunctionContext::makeDataPtr(data_, deleter_, *device_); - } - - inline DataPtr makeDataPtrFromContext() noexcept { - return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_}; - } - - IntArrayRef makeTempSizes() const noexcept { - static std::int64_t zeros[5] = {0, 0, 0, 0, 0}; - if (opts_.has_memory_format()) { - MemoryFormat format = *opts_.memory_format_opt(); - if (format == MemoryFormat::ChannelsLast) { - return IntArrayRef(zeros, 4); - } - if (format == MemoryFormat::ChannelsLast3d) { - return IntArrayRef(zeros, 5); - } - } - return IntArrayRef(zeros, 1); - } + IntArrayRef makeTempSizes() const noexcept; void* data_; IntArrayRef sizes_; diff --git a/aten/src/ATen/templates/Operators.cpp b/aten/src/ATen/templates/Operators.cpp index c0d46f58a848a..4d50c5a2e0ba2 100644 --- a/aten/src/ATen/templates/Operators.cpp +++ b/aten/src/ATen/templates/Operators.cpp @@ -1,6 +1,4 @@ #include -#include -#include namespace at { namespace _ops { diff --git a/aten/src/ATen/templates/Operators.h b/aten/src/ATen/templates/Operators.h index a92b7503ad7a2..39eaa2bd15a49 100644 --- a/aten/src/ATen/templates/Operators.h +++ b/aten/src/ATen/templates/Operators.h @@ -2,11 +2,8 @@ // ${generated_comment} -#include -#include -#include -#include -#include +#include +#include // Extension writers: do you write wrapper functions? Are you frustrated with // resolving overloads of operators? Are you frustrated with dealing with @@ -27,77 +24,23 @@ // ATEN_FN2(sin, out) gives a function that is *faithful* to the schema; // that is, the order of arguments is exactly what it looks like in the schema. -#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload::call -#define ATEN_FN(op_name) at::_ops::op_name::call - -// Separately, ATEN_OP(op) and ATEN_OP2(op, overload) define a class containing compile-time -// metadata about a given aten operator. -// Notable data on the class includes: -// - ATEN_OP2(add, Tensor)::name // returns the string name: "add" -// - ATEN_OP2(add, Tensor)::overload_name // returns the string overload name: "Tensor" -// - ATEN_OP2(add, Tensor)::schema // returns the C++ schema type: at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &) -// - ATEN_OP2(add, Tensor)::schema_str // returns the string jit type: "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor" - -#define ATEN_OP2(op_name, overload) at::_ops::op_name##_##overload -#define ATEN_OP(op_name) at::_ops::op_name +#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload +#define ATEN_FN(op_name) at::_ops::op_name // WARNING: Please do not call any of the ops in the _ops namespace directly. // Use the ATEN_FN macros. We do not guarantee stability of the naming // scheme for the functions in at::_ops +namespace at { namespace _ops { -// See Note [The ATen Operators API] for details of the at::_ops namespace - -namespace c10 { namespace impl { - -inline c10::optional -check_tensor_options_and_extract_memory_format( - const TensorOptions& options, - c10::optional memory_format) { - TORCH_CHECK( - options.requires_grad_opt() == c10::nullopt || - options.requires_grad_opt().value() == false, - "Operators taking TensorOptions cannot take a TensorOptions with " - "options.requires_grad set as true. This isn't implemented yet."); - TORCH_CHECK( - !(options.has_memory_format() && memory_format.has_value()), - "Cannot set memory_format both in TensorOptions and explicit argument; please delete " - "the redundant setter."); - if (memory_format.has_value()) { - return memory_format; - } else { - return options.memory_format_opt(); - } -} - -}} // namespace impl namespace c10 - - -// Forward declarations of any types needed in the operator signatures. -// We can't directly include these classes because it will cause circular include dependencies. -// This file is included by TensorBody.h, which defines the Tensor class. -namespace c10 { - -template -class optional; -template -class List; -class Stream; -struct Storage; - -} - -namespace at { - -class Tensor; -struct Dimname; -struct Generator; -using TensorList = c10::ArrayRef; -using DimnameList = c10::ArrayRef; -using Stream = c10::Stream; -using Storage = c10::Storage; -using QScheme = c10::QScheme; - -namespace _ops { +// NB: We are forced to special case requires_grad_. This is because all +// of the auto-generated inplace method signatures in TensorMethods.h are +// codegen'ed to return Tensor&, but requires_grad_ has a `manual_cpp_binding` +// with a different signature that returns `const Tensor&`. +// +// Eventually, the plan is to kill Tensor& from all C++ signatures and use +// const Tensor&. When that happens, we can remove this special case and just +// let the codegen handle it. +TORCH_API Tensor & requires_grad_(Tensor & self, bool requires_grad); ${declarations} diff --git a/aten/src/ATen/templates/RedispatchFunctions.h b/aten/src/ATen/templates/RedispatchFunctions.h index b5219425e8c8c..ecaf8f05162e1 100644 --- a/aten/src/ATen/templates/RedispatchFunctions.h +++ b/aten/src/ATen/templates/RedispatchFunctions.h @@ -18,7 +18,7 @@ namespace at { namespace redispatch { - ${function_redispatch_definitions} + ${function_redispatch_declarations} } // namespace redispatch } diff --git a/aten/src/ATen/templates/RegisterBackendSelect.cpp b/aten/src/ATen/templates/RegisterBackendSelect.cpp index 27b8e2bcc5125..bf5cc8683ecbc 100644 --- a/aten/src/ATen/templates/RegisterBackendSelect.cpp +++ b/aten/src/ATen/templates/RegisterBackendSelect.cpp @@ -5,7 +5,6 @@ #include #include -#include #include #include #include diff --git a/aten/src/ATen/templates/RegisterSchema.cpp b/aten/src/ATen/templates/RegisterSchema.cpp index 6861b5cbbcf3e..c9dbf5880a7e1 100644 --- a/aten/src/ATen/templates/RegisterSchema.cpp +++ b/aten/src/ATen/templates/RegisterSchema.cpp @@ -9,7 +9,6 @@ #include #include #include -#include #include #include #include diff --git a/aten/src/ATen/templates/TensorBody.h b/aten/src/ATen/templates/TensorBody.h index 775d2e6803aa6..fa879d656ab51 100644 --- a/aten/src/ATen/templates/TensorBody.h +++ b/aten/src/ATen/templates/TensorBody.h @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include @@ -367,15 +366,8 @@ class TORCH_API Tensor { bool is_alias_of(const at::Tensor& other) const{ return impl_->storage().is_alias_of(other.storage()); } - - Tensor toType(ScalarType t) const { - return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false); - } - - // TODO: Deprecate me - Tensor toBackend(Backend b) const { - return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false); - } + Tensor toType(ScalarType t) const; + Tensor toBackend(Backend b) const; C10_DEPRECATED_MESSAGE("Tensor.is_variable() is deprecated; everything is a variable now. (If you want to assert that variable has been appropriately handled already, use at::impl::variable_excluded_from_dispatch())") bool is_variable() const noexcept { @@ -523,11 +515,7 @@ class TORCH_API Tensor { /// Returns the `TensorOptions` corresponding to this `Tensor`. Defined in /// TensorOptions.h. - TensorOptions options() const { - return TensorOptions().dtype(dtype()) - .device(device()) - .layout(layout()); - } + TensorOptions options() const; void* data_ptr() const { return this->unsafeGetTensorImpl()->data(); @@ -621,26 +609,11 @@ class TORCH_API Tensor { Tensor & index_put_(std::initializer_list indices, Tensor const & rhs); Tensor & index_put_(std::initializer_list indices, const Scalar& v); - Tensor cpu() const { - return to(options().device(DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false); - } - - // TODO: The Python version also accepts arguments - Tensor cuda() const { - return to(options().device(DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false); - } - - Tensor hip() const { - return to(options().device(DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false); - } - - Tensor vulkan() const { - return to(options().device(DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false); - } - - Tensor metal() const { - return to(options().device(DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false); - } + Tensor cpu() const; + Tensor cuda() const; + Tensor hip() const; + Tensor vulkan() const; + Tensor metal() const; // ~~~~~ Autograd API ~~~~~ @@ -971,31 +944,6 @@ inline int64_t get_device(const Tensor& self) { return self.get_device(); } -#define DEFINE_CAST(T, name) \ - template <> \ - TORCH_API inline T* Tensor::data_ptr() const { \ - TORCH_CHECK( \ - scalar_type() == ScalarType::name, \ - "expected scalar type " \ - #name \ - " but found ", \ - scalar_type()); \ - return this->unsafeGetTensorImpl()->data_ptr_impl(); \ - } - -AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_CAST) -AT_FORALL_QINT_TYPES(DEFINE_CAST) -#undef DEFINE_CAST - -#define DEFINE_ITEM(T, name) \ - template <> \ - TORCH_API inline T Tensor::item() const { \ - return item().to##name(); \ - } - -AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_ITEM) -#undef DEFINE_ITEM - template auto Tensor::register_hook(T&& hook) const -> Tensor::hook_return_void_t { // Return the grad argument in case of a hook with void return type to have an @@ -1030,13 +978,6 @@ static inline DispatchKey legacyExtractDispatchKey(const Tensor& t) { } // namespace at -// See Note [Avoiding Include Cycles In Static Dispatch] -${static_dispatch_extra_headers} -namespace at { -${tensor_method_definitions} -} // namespace at - - namespace c10 { template <> struct MaybeOwnedTraits { diff --git a/aten/src/ATen/templates/TensorMethods.cpp b/aten/src/ATen/templates/TensorMethods.cpp new file mode 100644 index 0000000000000..0eba7dc65d737 --- /dev/null +++ b/aten/src/ATen/templates/TensorMethods.cpp @@ -0,0 +1,85 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +${static_dispatch_extra_headers} + +namespace at { + +using Stream = c10::Stream; + +Tensor Tensor::cpu() const { + return to(options().device(DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false); +} + +// TODO: The Python version also accepts arguments +Tensor Tensor::cuda() const { + return to(options().device(DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false); +} + +Tensor Tensor::hip() const { + return to(options().device(DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false); +} + +Tensor Tensor::vulkan() const { + return to(options().device(DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false); +} + +Tensor Tensor::metal() const { + return to(options().device(DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false); +} + +Tensor Tensor::toType(ScalarType t) const { + return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false); +} + +// TODO: Deprecate me +Tensor Tensor::toBackend(Backend b) const { + return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false); +} + +TensorOptions Tensor::options() const { + return TensorOptions().dtype(dtype()) + .device(device()) + .layout(layout()); +} + +${tensor_method_definitions} + +#define DEFINE_CAST(T, name) \ + template <> \ + TORCH_API T* Tensor::data_ptr() const { \ + TORCH_CHECK( \ + scalar_type() == ScalarType::name, \ + "expected scalar type " \ + #name \ + " but found ", \ + scalar_type()); \ + return this->unsafeGetTensorImpl()->data_ptr_impl(); \ + } + +AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_CAST) +AT_FORALL_QINT_TYPES(DEFINE_CAST) +#undef DEFINE_CAST + +#define DEFINE_ITEM(T, name) \ + template <> \ + TORCH_API T Tensor::item() const { \ + return item().to##name(); \ + } + +AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF(DEFINE_ITEM) +#undef DEFINE_ITEM + +} //namespace at diff --git a/caffe2/contrib/aten/gen_op.py b/caffe2/contrib/aten/gen_op.py index 13908902e2084..94a1f1fedc3b2 100755 --- a/caffe2/contrib/aten/gen_op.py +++ b/caffe2/contrib/aten/gen_op.py @@ -77,7 +77,7 @@ def value_is_tensor_type(v): 'at::Scalar': 'assignTo(Output(${offset}),${output}.type(), ${output});', 'bool': 'assignToValue(Output(${offset}),${output});', 'int64_t': 'assignToValue(Output(${offset}),${output});', - '::std::vector': 'assignListStartingAt(${offset}, ${output});', + 'std::vector': 'assignListStartingAt(${offset}, ${output});', } # for each non-Tensor aten argument, how to we read it from caffe2's @@ -90,8 +90,8 @@ def value_is_tensor_type(v): 'double': 'double ${arg} = readAttribute("${arg}");', 'int64_t': 'int64_t ${arg} = readAttribute("${arg}");', 'at::IntArrayRef': 'auto ${arg} = readIntArrayRef("${arg}");', - '::std::array': 'auto ${arg} = readBoolMask<2>("${arg}");', - '::std::array': 'auto ${arg} = readBoolMask<3>("${arg}");', + 'std::array': 'auto ${arg} = readBoolMask<2>("${arg}");', + 'std::array': 'auto ${arg} = readBoolMask<3>("${arg}");', } # for BC reasons we want to route some of the functions to different @@ -189,7 +189,7 @@ def get_output(o, i): if len(o['returns']) == 1: return 'the_result' else: - return '::std::get<{}>(the_result)'.format(i) + return 'std::get<{}>(the_result)'.format(i) def attribute_names(o): diff --git a/tools/codegen/api/python.py b/tools/codegen/api/python.py index 5064f069a0643..8eb2f3696f517 100644 --- a/tools/codegen/api/python.py +++ b/tools/codegen/api/python.py @@ -613,7 +613,7 @@ def argument_type_str(t: Type, *, simple_type: bool = False) -> str: size = t.size if not simple_type else None if str(t.elem) == 'bool': assert t.size is not None - return f'::std::array' + return f'std::array' elif str(t.elem) == 'int': return f'IntArrayRef[{size}]' if size is not None else 'IntArrayRef' elif str(t.elem) == 'Tensor': @@ -910,16 +910,16 @@ def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument: # to add an appropriate wrap() overload in torch/csrc/autograd/utils/wrap_outputs.h. SUPPORTED_RETURN_TYPES = { 'at::Tensor', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::tuple', - '::std::vector', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::tuple', + 'std::vector', 'at::Scalar', 'bool', 'int64_t', 'void*', 'void', 'at::QScheme', 'double', 'at::IntArrayRef', diff --git a/tools/codegen/api/types.py b/tools/codegen/api/types.py index 675db01cd3243..850b6a50df285 100644 --- a/tools/codegen/api/types.py +++ b/tools/codegen/api/types.py @@ -166,10 +166,10 @@ class VectorCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'::std::vector<{self.elem.cpp_type()}>' + return f'std::vector<{self.elem.cpp_type()}>' def cpp_type_registration_declarations(self) -> str: - return f'::std::vector<{self.elem.cpp_type_registration_declarations()}>' + return f'std::vector<{self.elem.cpp_type_registration_declarations()}>' def remove_const_ref(self) -> 'CType': return VectorCType(self.elem.remove_const_ref()) @@ -181,10 +181,10 @@ class ArrayCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'::std::array<{self.elem.cpp_type()},{self.size}>' + return f'std::array<{self.elem.cpp_type()},{self.size}>' def cpp_type_registration_declarations(self) -> str: - return f'::std::array<{self.elem.cpp_type_registration_declarations()},{self.size}>' + return f'std::array<{self.elem.cpp_type_registration_declarations()},{self.size}>' def remove_const_ref(self) -> 'CType': return ArrayCType(self.elem.remove_const_ref(), self.size) @@ -195,10 +195,10 @@ class TupleCType: def cpp_type(self, *, strip_ref: bool = False) -> str: # Do not pass `strip_ref` recursively. - return f'::std::tuple<{",".join([e.cpp_type() for e in self.elems])}>' + return f'std::tuple<{",".join([e.cpp_type() for e in self.elems])}>' def cpp_type_registration_declarations(self) -> str: - return f'::std::tuple<{",".join([e.cpp_type_registration_declarations() for e in self.elems])}>' + return f'std::tuple<{",".join([e.cpp_type_registration_declarations() for e in self.elems])}>' def remove_const_ref(self) -> 'CType': return TupleCType([e.remove_const_ref() for e in self.elems]) @@ -338,37 +338,26 @@ def name(self) -> str: return n # Render the C++ declaration for this signature - def decl(self, *, name: Optional[str] = None, prefix: str = "", is_redispatching_fn: bool = False) -> str: + def decl(self, *, prefix: str = "", is_redispatching_fn: bool = False) -> str: returns_type = cpp.returns_type(self.func.returns).cpp_type() cpp_args = [a.decl() for a in self.arguments()] if is_redispatching_fn: cpp_args = ['c10::DispatchKeySet dispatchKeySet'] + cpp_args cpp_args_str = ', '.join(cpp_args) - if name is None: - name = prefix + self.name() + name = prefix + self.name() return f"{returns_type} {name}({cpp_args_str})" # Render the C++ definition for this signature, not including # the body (with curly braces) - def defn(self, *, name: Optional[str] = None, prefix: str = "", is_redispatching_fn: bool = False) -> str: + def defn(self, *, prefix: str = "", is_redispatching_fn: bool = False) -> str: returns_type = cpp.returns_type(self.func.returns).cpp_type() cpp_args = [a.defn() for a in self.arguments()] if is_redispatching_fn: cpp_args = ['c10::DispatchKeySet dispatchKeySet'] + cpp_args cpp_args_str = ', '.join(cpp_args) - if name is None: - name = prefix + self.name() + name = prefix + self.name() return f"{returns_type} {name}({cpp_args_str})" - def ptr_type(self) -> str: - args_types_str = ', '.join(a.type for a in self.arguments()) - return f'{cpp.returns_type(self.func.returns).cpp_type()} (*)({args_types_str})' - - # Return the C++ function type, e.g., something like int(bool) - def type(self) -> str: - args_types_str = ', '.join(a.type for a in self.arguments()) - return f'{cpp.returns_type(self.func.returns).cpp_type()} ({args_types_str})' - # Represents group of all CppSignatures associated with a # FunctionSchema. Right now, that's the regular, user-visible @@ -435,11 +424,8 @@ def decl(self, name: Optional[str] = None) -> str: name = self.name() return f"{self.returns_type().cpp_type()} {name}({args_str})" - def defn(self, name: Optional[str] = None, *, is_redispatching_fn: bool = False) -> str: - args = [a.defn() for a in self.arguments()] - if is_redispatching_fn: - args = ['c10::DispatchKeySet dispatchKeySet'] + args - args_str = ', '.join(args) + def defn(self, name: Optional[str] = None) -> str: + args_str = ', '.join(a.defn() for a in self.arguments()) if name is None: name = self.name() return f"{self.returns_type().cpp_type()} {name}({args_str})" diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py index e69dc62c456fd..193a4d35d74bb 100644 --- a/tools/codegen/gen.py +++ b/tools/codegen/gen.py @@ -140,24 +140,14 @@ def cpp_string(s: str) -> str: # to be generated. This pattern makes it convenient to use map, concatMap # and similar functional combinators. -def static_dispatch_keys(backend: Optional[BackendIndex]) -> List[DispatchKey]: +def static_dispatch_extra_headers(backend: Optional[BackendIndex]) -> str: if backend is None: - return [] - else: - return [ - backend.dispatch_key, - DispatchKey.CompositeImplicitAutograd, - DispatchKey.CompositeExplicitAutograd - ] - -def static_dispatch_extra_headers(backend: Optional[BackendIndex], skip_tensor_include: bool = False) -> str: - if skip_tensor_include: - # See Note [Avoiding Include Cycles In Static Dispatch] - maybe_inl = '_inl' - else: - maybe_inl = '' - return '\n'.join([ - f'#include ' for dispatch_key in static_dispatch_keys(backend)]) + return '' + return f""" +#include +#include +#include +""" def static_dispatch( f: NativeFunction, cpp_sig: CppSignature, @@ -197,7 +187,23 @@ class RegisterSchema: def __call__(self, f: NativeFunction) -> Optional[str]: if not self.selector.is_native_function_selected(f): return None - return f'm.def({cpp_string(str(f.func))});\n' + schema_str = cpp_string(str(f.func)) + schema_str = '"' + "aten::" + schema_str[1:] + return f'm.def({schema_str});\n' + + +def _num_leading_spaces(line: str) -> int: + return len(line) - len(line.lstrip()) + + +# Unindents all lines in code. Each line gets unindented the same amount; +# that amount is equal to the smallest number of leading spaces across all lines +def deindent(code: str) -> str: + lines = code.split('\n') + min_leading_spaces = min(map(_num_leading_spaces, lines)) + lines = [line[min_leading_spaces:] for line in lines] + return '\n'.join(lines) + # Generates Operators.h and Operators.cpp. # These provide macros that, given an operator and overload name, allow users @@ -213,119 +219,141 @@ class ComputeOperators: @method_with_native_function def __call__(self, f: NativeFunction) -> Optional[str]: - sig = DispatcherSignature.from_schema(f.func) - name = f.func.name.unambiguous_name() - call_method_name = 'call' - redispatch_method_name = 'redispatch' + # NB: requires_grad is the only exception to the rule because + # its const correctness is questionable. + if str(f.func.name) in set(['requires_grad_']): + return None if self.target is Target.DECLARATION: - # Note [The ATen Operators API] - # The ATen Operators API lives in the at::_ops namespace, and contains compile-time - # metadata about each operator + entry points into the Dispatcher. - # The C++ function, method, and redispatch API's are all implemented as wrappers - # into various bits of the structs defined here. - # - # Important characteristics about the Operators API: - # (1) It follows the Dispatcher API. - # This is kind of necessary to avoid overhead. - # For example: if it followed the C++ API, then all of the faithful C++ factory functions - # would need to wrap their arguments into TensorOptions only to unwrap them again. - # (2) Overload names are disambiguated. - # This is helpful for pytorch extenders who would like to decltype() an aten operator, - # that has overloads, e.g. decltype(at::_ops::mul_Tensor::call) - # (3) No argument defaulting is allowed. - # This is more of an implementation detail to avoid #include cycles, - # since TensorBody.h (which defines the Tensor class) needs to include this file. - # (4) manual_cpp_bindings and faithful names are not included in the API. - # This applies to stuff like __dispatch__is_complex(), and add_outf(). - # These aren't "real aten ops", they're just additional functions provided by the C++ API. - # They're implemented as wrappers in Functions.h that call into the actual operators - # defined here, i.e. at::_ops::is_complex::call() and at::_ops::add_out::call(). - # This means that ATEN_OP(is_complex) will not fastpath, and will go through the dispatcher. - return f""" -struct TORCH_API {name} {{ - using schema = {sig.type()}; - using ptr_schema = schema*; - static constexpr const char* name = "aten::{str(f.func.name.name)}"; - static constexpr const char* overload_name = "{f.func.name.overload_name}"; - static constexpr const char* schema_str = {cpp_string(str(f.func))}; - static {sig.defn(name=call_method_name, is_redispatching_fn=False)}; - static {sig.defn(name=redispatch_method_name, is_redispatching_fn=True)}; -}};""" - elif self.target is Target.DEFINITION: - defns = '' - for is_redispatching_fn in [False, True]: - if is_redispatching_fn: - dispatcher_exprs_str = ', '.join(['dispatchKeySet'] + [a.name for a in sig.arguments()]) - dispatcher_call = 'redispatch' - method_name = f'{name}::{redispatch_method_name}' - else: - dispatcher_exprs_str = ', '.join([a.name for a in sig.arguments()]) - dispatcher_call = 'call' - method_name = f'{name}::{call_method_name}' - - defns += f""" -// aten::{f.func} -{sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{ - static auto op = c10::Dispatcher::singleton() - .findSchemaOrThrow(name, overload_name) - .typed(); - return op.{dispatcher_call}({dispatcher_exprs_str}); -}} -""" - return defns + return self.gen_declaration(f) + if self.target is Target.DEFINITION: + return self.gen_definition(f) else: assert_never(self.target) + # NB: This must be synchronized with the naming scheme in + # aten/src/ATen/templates/Operators.h + # Given a function schema "aten::op.overload(...)", + # If there is no overload name, this returns f"{op}" + # If there is an overload name, this returns f"{op}_{overload}" + def unambiguous_function_name(self, f: NativeFunction) -> str: + base_name = str(f.func.name.name) + overload_name = f.func.name.overload_name + if overload_name: + return f'{base_name}_{overload_name}' + return base_name + + def gen_declaration(self, f: NativeFunction) -> str: + unambiguous_name = self.unambiguous_function_name(f) + sig = DispatcherSignature.from_schema(f.func) + return f"TORCH_API {sig.decl(unambiguous_name)};" + + def most_faithful_name(self, f: NativeFunction) -> str: + sig_group = CppSignatureGroup.from_native_function(f, method=False) + sig = sig_group.most_faithful_signature() + return sig.name() -# Generates Function.h, which provides the functional public C++ API, -# and the scaffolding to call into the dispatcher from these functions. + def invocation(self, f: NativeFunction) -> str: + faithful_op_name = self.most_faithful_name(f) + args = tuple(arg.name for arg in dispatcher.arguments(f.func)) + # Method only + if Variant.function not in f.variants: + return f"{args[0]}.{faithful_op_name}({', '.join(args[1:])})" + return f"at::{faithful_op_name}({', '.join(args)})" + + def gen_definition(self, f: NativeFunction) -> str: + unambiguous_name = self.unambiguous_function_name(f) + args = dispatcher.arguments(f.func) + sig = DispatcherSignature.from_schema(f.func) + + return deindent(f"""\ + {sig.defn(unambiguous_name)} {{ + return {self.invocation(f)}; + }}\ + """) + + +# Generates Function.cpp and Function.h. These files provide the +# functional public C++ API, and the scaffolding to call into +# the dispatcher from these functions. See also compute_tensor_method. @dataclass(frozen=True) class ComputeFunction: + target: Union[ + Literal[Target.DECLARATION], + Literal[Target.DEFINITION] + ] static_dispatch_backend_index: Optional[BackendIndex] + is_redispatching_fn: bool @method_with_native_function def __call__(self, f: NativeFunction) -> Optional[str]: - if Variant.function not in f.variants: + # We unconditionally generate function variants of the redispatch API. + # This is mainly because we can namespace functions separately, but not methods, + if Variant.function not in f.variants and not self.is_redispatching_fn: return None + with native_function_manager(f): + return self.callImpl(f) + + def callImpl(self, f: NativeFunction) -> str: + name = cpp.name(f.func) + sig_group = CppSignatureGroup.from_native_function(f, method=False, fallback_binding=f.manual_cpp_binding) + if self.target is Target.DECLARATION: + sig_str = sig_group.signature.decl(is_redispatching_fn=self.is_redispatching_fn) + result = f"TORCH_API {sig_str};\n" + if sig_group.faithful_signature is not None: + sig_str = sig_group.faithful_signature.decl(is_redispatching_fn=self.is_redispatching_fn) + result += f"TORCH_API {sig_str};\n" + return result + + if self.target is not Target.DEFINITION: + assert_never(self.target) + def generate_defn(faithful: bool) -> str: - if faithful: + dispatcher_sig = DispatcherSignature.from_schema(f.func) + + if faithful and sig_group.faithful_signature is not None: sig = sig_group.faithful_signature - assert sig is not None else: sig = sig_group.signature - # See Note [The ATen Operators API] - target_sig = DispatcherSignature.from_schema(f.func) - exprs = translate(sig.arguments(), target_sig.arguments()) - exprs_str = ', '.join([e.expr for e in exprs]) + dispatcher_exprs = translate(sig.arguments(), dispatcher_sig.arguments()) + if self.is_redispatching_fn: + dispatcher_exprs_str = ', '.join(['dispatchKeySet'] + [a.expr for a in dispatcher_exprs]) + dispatcher_call = 'redispatch' + else: + dispatcher_exprs_str = ', '.join(a.expr for a in dispatcher_exprs) + dispatcher_call = 'call' static_dispatch_block = static_dispatch(f, sig, method=False, backend_index=self.static_dispatch_backend_index) if static_dispatch_block is None: return f""" // aten::{f.func} -TORCH_API inline {sig.decl()} {{ - return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); +{sig.defn(is_redispatching_fn=self.is_redispatching_fn)} {{ + static auto op = c10::Dispatcher::singleton() + .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") + .typed<{dispatcher_sig.type()}>(); + return op.{dispatcher_call}({dispatcher_exprs_str}); }} """ else: return f""" // aten::{f.func} -TORCH_API inline {sig.decl()} {{ +{sig.defn(is_redispatching_fn=self.is_redispatching_fn)} {{ {static_dispatch_block} }} """ - result = generate_defn(False) + result = generate_defn(sig_group.faithful_signature is None) if sig_group.faithful_signature is not None: result += generate_defn(True) return result -# Generates TensorBody.h. This file provides the object-oriented (method-based) -# public C++ API, and the scaffolding to call into the dispatcher from these functions. +# Generates TensorBody.h (sic) and TensorMethods.cpp. These files provide the +# object-oriented (method-based) public C++ API, and the scaffolding to call into +# the dispatcher from these functions. See also compute_function. @dataclass(frozen=True) class ComputeTensorMethod: target: Union[ @@ -342,6 +370,8 @@ def __call__(self, f: NativeFunction) -> Optional[str]: assert not f.func.is_out_fn() assert f.func.arguments.self_arg is not None + name = cpp.name(f.func) + sig_group = CppSignatureGroup.from_native_function(f, method=True, fallback_binding=f.manual_cpp_binding) if self.target is Target.DECLARATION: @@ -354,28 +384,32 @@ def __call__(self, f: NativeFunction) -> Optional[str]: assert_never(self.target) def generate_defn(faithful: bool) -> str: + dispatcher_sig = DispatcherSignature.from_schema(f.func) + if faithful: sig = sig_group.faithful_signature assert sig is not None else: sig = sig_group.signature - target_sig = DispatcherSignature.from_schema(f.func) - exprs = translate(sig.arguments(), target_sig.arguments(), method=True) - exprs_str = ', '.join([e.expr for e in exprs]) + dispatcher_exprs = translate(sig.arguments(), dispatcher_sig.arguments(), method=True) + dispatcher_exprs_str = ', '.join(a.expr for a in dispatcher_exprs) static_dispatch_block = static_dispatch(f, sig, method=True, backend_index=self.static_dispatch_backend_index) if static_dispatch_block is None: return f""" // aten::{f.func} -inline {sig.defn(prefix="Tensor::")} const {{ - return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); +{sig.defn(prefix="Tensor::")} const {{ + static auto op = c10::Dispatcher::singleton() + .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") + .typed<{dispatcher_sig.type()}>(); + return op.call({dispatcher_exprs_str}); }} """ else: return f""" // aten::{f.func} -inline {sig.defn(prefix="Tensor::")} const {{ +{sig.defn(prefix="Tensor::")} const {{ {static_dispatch_block} }} """ @@ -386,42 +420,6 @@ def generate_defn(faithful: bool) -> str: return result -# Generates RedispatchFunctions.h. -# This is similar to the C++ API defined in Functions.h, but provides access -# to the dispatcher's redispatch API. -@dataclass(frozen=True) -class ComputeRedispatchFunction: - - @method_with_native_function - def __call__(self, f: NativeFunction) -> Optional[str]: - # We unconditionally generate function variants of the redispatch API. - # This is mainly because we can namespace functions separately, but not methods, - sig_group = CppSignatureGroup.from_native_function(f, method=False, fallback_binding=f.manual_cpp_binding) - - def generate_defn(faithful: bool) -> str: - if faithful: - sig = sig_group.faithful_signature - assert sig is not None - else: - sig = sig_group.signature - - target_sig = DispatcherSignature.from_schema(f.func) - exprs = translate(sig.arguments(), target_sig.arguments()) - exprs_str = ', '.join(['dispatchKeySet'] + [a.expr for a in exprs]) - - return f""" -// aten::{f.func} -TORCH_API inline {sig.decl(is_redispatching_fn=True)} {{ - return at::_ops::{f.func.name.unambiguous_name()}::redispatch({exprs_str}); -}} -""" - result = generate_defn(False) - if sig_group.faithful_signature is not None: - result += generate_defn(True) - - return result - - # Generates ATenOpList.cpp, a runtime accessible list of all aten # operators. # TODO: This was historically used to help some JIT interop code @@ -506,8 +504,8 @@ def __call__(self, f: NativeFunction) -> Optional[str]: C10_ALWAYS_INLINE {sig.defn(name)} {{ static auto op = c10::Dispatcher::singleton() - .findSchemaOrThrow(at::_ops::{f.func.name.unambiguous_name()}::name, at::_ops::{f.func.name.unambiguous_name()}::overload_name) - .typed(); + .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}") + .typed<{dispatcher_sig.type()}>(); {compute_dk} return op.redispatch(_dk, {', '.join(a.expr for a in dispatcher_exprs)}); }} @@ -1052,17 +1050,7 @@ def make_file_manager(install_dir: str) -> FileManager: }) if dispatch_key in functions_keys: - if dispatch_key in static_dispatch_keys(static_dispatch_idx): - # See Note [Avoiding Include Cycles In Static Dispatch] - inl_headers = '' - else: - inl_headers = f'#include ' - fm.write_with_template(f'{dispatch_key}Functions.h', 'DispatchKeyFunctions.h', lambda: { - 'dispatch_key': str(dispatch_key), - 'inline_headers_for_nonstatic_build': inl_headers, - }) - fm.write_with_template(f'{dispatch_key}Functions_inl.h', 'DispatchKeyFunctions_inl.h', lambda: { 'dispatch_namespace': dispatch_key.lower(), 'dispatch_namespaced_declarations': list(concatMap( dest.RegisterDispatchKey( @@ -1106,27 +1094,35 @@ def make_file_manager(install_dir: str) -> FileManager: }) cpu_fm.write('Functions.h', lambda: { + 'function_declarations': list(mapMaybe(ComputeFunction( + Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=False), native_functions)), + }) + cpu_fm.write('Functions.cpp', lambda: { 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), 'function_definitions': list(mapMaybe(ComputeFunction( - static_dispatch_backend_index=static_dispatch_idx), native_functions)), + Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=False), native_functions)), + }) + cpu_fm.write('RedispatchFunctions.h', lambda: { + 'function_redispatch_declarations': list(mapMaybe(ComputeFunction( + Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=True), native_functions)), + }) + cpu_fm.write('RedispatchFunctions.cpp', lambda: { + 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), + 'function_redispatch_definitions': list(mapMaybe(ComputeFunction( + Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx, is_redispatching_fn=True), native_functions)), }) - core_fm.write('TensorBody.h', lambda: { - 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx, skip_tensor_include=True), - 'tensor_method_declarations': list(mapMaybe(ComputeTensorMethod( - target=Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), - 'tensor_method_definitions': list(mapMaybe(ComputeTensorMethod( - target=Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), + 'tensor_method_declarations': list(mapMaybe( + ComputeTensorMethod(Target.DECLARATION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), }) - - cpu_fm.write('RedispatchFunctions.h', lambda: { - 'function_redispatch_definitions': list(mapMaybe(ComputeRedispatchFunction(), native_functions)), + core_fm.write('TensorMethods.cpp', lambda: { + 'static_dispatch_extra_headers': static_dispatch_extra_headers(static_dispatch_idx), + 'tensor_method_definitions': list(mapMaybe( + ComputeTensorMethod(Target.DEFINITION, static_dispatch_backend_index=static_dispatch_idx), native_functions)), }) - core_fm.write('ATenOpList.cpp', lambda: { 'aten_ops': list(mapMaybe(compute_aten_op, native_functions)), }) - cpu_fm.write('NativeFunctions.h', lambda: { 'native_function_declarations': list(concatMap( # Convert to a set first to remove duplicate kernel names. diff --git a/tools/codegen/model.py b/tools/codegen/model.py index acde0827b0e6c..fba3756fab181 100644 --- a/tools/codegen/model.py +++ b/tools/codegen/model.py @@ -1454,18 +1454,6 @@ def __str__(self) -> str: else: return f"{self.name}" - # NB: This must be synchronized with the naming scheme in - # aten/src/ATen/templates/Operators.h - # Given a function schema "aten::op.overload(...)", - # If there is no overload name, this returns f"{op}" - # If there is an overload name, this returns f"{op}_{overload}" - def unambiguous_name(self) -> str: - if self.overload_name: - return f"{self.name}_{self.overload_name}" - else: - return f"{self.name}" - - def gets_generated_out_inplace_wrapper(f: NativeFunction, g: NativeFunctionsGroup, b: BackendIndex) -> bool: return f.func.kind() is not SchemaKind.functional and \ not b.has_kernel(f) and \ From ef1c107be523c23f883a7f76963b70a47638b5f5 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 17 Jun 2021 15:18:34 -0700 Subject: [PATCH 213/305] [vulkan] Do not use memcmp to compare structs (#60199) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60199 It isn't safe to use `memcmp` to determine the equality of structs due to potential random padding between fields of the struct. This can cause overloaded equality operators to return false when comparing structs with equivalent fields. This bug appears to be responsible for the Vulkan backend crashing on WorkVC release builds. Test Plan: Run Vulkan unit tests: ``` cd ~/fbsource buck build -c ndk.custom_libcxx=false -c pt.enable_qpl=0 //xplat/caffe2:pt_vulkan_api_test_binAndroid\#android-arm64 --show-output adb push buck-out/gen/xplat/caffe2/pt_vulkan_api_test_binAndroid\#android-arm64 /data/local/tmp/vulkan_api_test adb shell "/data/local/tmp/vulkan_api_test" cd - ``` Test on workvc rdk build, first ensure you are receiving the Vulkan models. ``` buck install fbsource//fbandroid/mode/opt fbsource//fbandroid/mode/aloha_build_rdk fbsource//fbandroid/mode/no_obfuscation fbandroid/buck-configs/buckconfig.caffe2_pkg_snpe_libs_android aloha_workvc_rdk --deep --show-full-output ``` Reviewed By: IvanKobzarev Differential Revision: D29203177 fbshipit-source-id: e0ee79d4e635174e165b250f2cee842a09092df9 --- aten/src/ATen/native/vulkan/api/Pipeline.h | 12 ++++------ aten/src/ATen/native/vulkan/api/Resource.h | 8 +++---- aten/src/ATen/native/vulkan/api/Shader.h | 28 ++++++++++++---------- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/aten/src/ATen/native/vulkan/api/Pipeline.h b/aten/src/ATen/native/vulkan/api/Pipeline.h index 5ac4e1f34e684..bbff4fa914a37 100644 --- a/aten/src/ATen/native/vulkan/api/Pipeline.h +++ b/aten/src/ATen/native/vulkan/api/Pipeline.h @@ -196,11 +196,8 @@ inline Pipeline::Barrier::operator bool() const { inline bool operator==( const Pipeline::Layout::Descriptor& _1, const Pipeline::Layout::Descriptor& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(Pipeline::Layout::Descriptor))); + return (_1.descriptor_set_layout == _2.descriptor_set_layout); } inline size_t Pipeline::Layout::Factory::Hasher::operator()( @@ -211,11 +208,10 @@ inline size_t Pipeline::Layout::Factory::Hasher::operator()( inline bool operator==( const Pipeline::Descriptor& _1, const Pipeline::Descriptor& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(Pipeline::Descriptor))); + return (_1.pipeline_layout == _2.pipeline_layout && \ + _1.shader_module == _2.shader_module && \ + _1.local_work_group == _2.local_work_group); } inline size_t Pipeline::Factory::Hasher::operator()( diff --git a/aten/src/ATen/native/vulkan/api/Resource.h b/aten/src/ATen/native/vulkan/api/Resource.h index 0726c3215d1fc..192ddd7d5f03b 100644 --- a/aten/src/ATen/native/vulkan/api/Resource.h +++ b/aten/src/ATen/native/vulkan/api/Resource.h @@ -396,11 +396,11 @@ inline Resource::Buffer::operator bool() const { inline bool operator==( const Resource::Image::Sampler::Descriptor& _1, const Resource::Image::Sampler::Descriptor& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(Resource::Image::Sampler::Descriptor))); + return (_1.filter == _2.filter && \ + _1.mipmap_mode == _2.mipmap_mode && \ + _1.address_mode == _2.address_mode && \ + _1.border == _2.border); } inline size_t Resource::Image::Sampler::Factory::Hasher::operator()( diff --git a/aten/src/ATen/native/vulkan/api/Shader.h b/aten/src/ATen/native/vulkan/api/Shader.h index edec81f4feab1..e68061a320b70 100644 --- a/aten/src/ATen/native/vulkan/api/Shader.h +++ b/aten/src/ATen/native/vulkan/api/Shader.h @@ -221,11 +221,8 @@ inline Shader::Layout::Object Shader::Layout::Cache::retrieve( inline bool operator==( const Shader::WorkGroup& _1, const Shader::WorkGroup& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(Shader::WorkGroup))); + return (_1.data[0u] == _2.data[0u] && _1.data[1u] == _2.data[1u] && _1.data[2u] == _2.data[2u]); } inline Shader::Descriptor::Descriptor(const char* const glsl) @@ -259,11 +256,17 @@ inline Shader::Descriptor::Descriptor( inline bool operator==( const Shader::Descriptor& _1, const Shader::Descriptor& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(Shader::Descriptor))); + if (_1.type != _2.type) + return false; + + if (_1.type == Shader::Descriptor::Type::Binary) { + return (_1.shader.binary.spirv == _2.shader.binary.spirv && \ + _1.shader.binary.size == _2.shader.binary.size); + } + else { + return (_1.shader.source.glsl == _2.shader.source.glsl); + } } inline size_t Shader::Factory::Hasher::operator()( @@ -286,11 +289,12 @@ inline size_t Shader::Factory::Hasher::operator()( inline bool operator==( const VkDescriptorSetLayoutBinding& _1, const VkDescriptorSetLayoutBinding& _2) { - static_assert( - std::is_trivially_copyable::value, - "This implementation is no longer valid!"); - return (0 == memcmp(&_1, &_2, sizeof(VkDescriptorSetLayoutBinding))); + return (_1.binding == _2.binding && \ + _1.descriptorType == _2.descriptorType && \ + _1.descriptorCount == _2.descriptorCount && \ + _1.stageFlags == _2.stageFlags && \ + _1.pImmutableSamplers == _2.pImmutableSamplers); } #endif /* USE_VULKAN_API */ From 3995fb1840bfaeb4901cf0965b9098c86443d0bf Mon Sep 17 00:00:00 2001 From: BowenBao Date: Thu, 17 Jun 2021 15:47:49 -0700 Subject: [PATCH 214/305] Add new_ones symbolic (#59255) (#59539) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59539 Add new_ones symbolic in PT-ONNX exporter Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D29046603 Pulled By: SplitInfinity fbshipit-source-id: e7420c7b543c33e3640e62461d08ff4d5843eda7 Co-authored-by: Shubham Bhokare --- test/onnx/test_pytorch_onnx_onnxruntime.py | 10 ++++++++++ torch/onnx/symbolic_opset9.py | 6 ++++++ 2 files changed, 16 insertions(+) diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index ecccccc0b704e..75da49c3e14f3 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -5018,6 +5018,16 @@ def forward(self, x): self.run_test(Zero_(), x, input_names=["x"], dynamic_axes={"x": [0, 1, 2]}) self.run_test(Zero_(), x, remained_onnx_input_idx=[]) + @skipIfUnsupportedMinOpsetVersion(9) + def test_new_ones(self): + class OnesModel(torch.nn.Module): + def forward(self, x): + return x.new_ones(x.shape[1:2]), x.new_ones(x.shape[2:], dtype=torch.long) + + x = torch.randn(2, 3, 4) + self.run_test(OnesModel(), x, input_names=["x"], dynamic_axes={"x": [0, 1, 2]}) + self.run_test(OnesModel(), x, remained_onnx_input_idx=[]) + @skipIfONNXShapeInference(True) @skipIfUnsupportedMinOpsetVersion(9) def test_tolist(self): diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py index 22e3eaa4b57b8..5d655e67c9db1 100644 --- a/torch/onnx/symbolic_opset9.py +++ b/torch/onnx/symbolic_opset9.py @@ -1816,6 +1816,12 @@ def ones_like(g, input, dtype=None, layout=None, device=None, pin_memory=False, return g.op("ConstantOfShape", shape, value_t=torch.tensor([1], dtype=sym_help.scalar_type_to_pytorch_type[dtype])) +def new_ones(g, self, sizes, dtype, layout, device, pin_memory=False): + self_dtype = sym_help._try_get_scalar_type(self) + if dtype is None and self_dtype is not None: + dtype = self_dtype + dtype = sym_help.scalar_type_to_onnx.index(sym_help.cast_pytorch_to_onnx[dtype]) + return ones(g, sizes, dtype, layout, device, pin_memory) def full(g, sizes, value, dtype, layout, device, pin_memory=False): const_value = sym_help._maybe_get_const(value, "t") From 9c03de1dde1f72a9b3961250666d6c949e15bc03 Mon Sep 17 00:00:00 2001 From: zhouzhuojie Date: Thu, 17 Jun 2021 16:17:47 -0700 Subject: [PATCH 215/305] Use mirrors for ubuntu apt source (#60216) Summary: Fixes https://github.com/pytorch/pytorch/issues/60135 Experimented on circleci https://app.circleci.com/pipelines/github/zhouzhuojie/gha-ci-playground/7/workflows/965c95b8-2186-434a-92ca-9cd9c8aaafdc/jobs/7 Sample logs ``` Need to get 1,389 kB of archives. After this operation, 5,495 kB of additional disk space will be used. Get:1 http://mirrors.ubuntu.com/mirrors.txt Mirrorlist [3,270 B] Get:2 http://mirror.lstn.net/ubuntu focal/main amd64 libtcl8.6 amd64 8.6.10+dfsg-1 [902 kB] Get:7 http://ubuntu.securedservers.com focal/main amd64 libipc-run-perl all 20180523.0-2 [89.7 kB] Get:5 http://mirrors.edge.kernel.org/ubuntu focal/universe amd64 expect amd64 5.45.4-2build1 [137 kB] Get:4 http://mirror.pnl.gov/ubuntu focal/universe amd64 tcl-expect amd64 5.45.4-2build1 [105 kB] Get:6 http://mirror.lstn.net/ubuntu focal/main amd64 libio-pty-perl amd64 1:1.12-1 [32.4 kB] Get:9 https://mirrors.bloomu.edu/ubuntu focal/main amd64 libtimedate-perl all 2.3200-1 [34.0 kB] Get:8 http://la-mirrors.evowise.com/ubuntu focal/universe amd64 libtime-duration-perl all 1.21-1 [13.1 kB] Get:3 http://mirrors.ocf.berkeley.edu/ubuntu focal/main amd64 tcl8.6 amd64 8.6.10+dfsg-1 [14.8 kB] Get:10 http://mirrors.ocf.berkeley.edu/ubuntu focal/universe amd64 moreutils amd64 0.63-1 [60.5 kB] Fetched 1,392 kB in 3s (464 kB/s) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/60216 Reviewed By: seemethere Differential Revision: D29214661 Pulled By: zhouzhuojie fbshipit-source-id: ed2d85f8c0c23af4bcf33558c57472fcf9d913e8 --- .circleci/scripts/setup_ci_environment.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.circleci/scripts/setup_ci_environment.sh b/.circleci/scripts/setup_ci_environment.sh index 59bbe84e9ebe6..d7ab8e2776780 100755 --- a/.circleci/scripts/setup_ci_environment.sh +++ b/.circleci/scripts/setup_ci_environment.sh @@ -7,6 +7,9 @@ sudo rm -f /etc/apt/heroku.list sudo rm -f /etc/apt/openjdk-r-ubuntu-ppa-xenial.list sudo rm -f /etc/apt/partner.list +# To increase the network reliability, let apt decide which mirror is best to use +sudo sed -i -e 's/http:\/\/.*archive/mirror:\/\/mirrors/' -e 's/\/ubuntu\//\/mirrors.txt/' /etc/apt/sources.list + retry () { $* || $* || $* || $* || $* } From acf04cdedf29836833d799595d4caaa72dc1d22c Mon Sep 17 00:00:00 2001 From: Zhuojie Zhou Date: Thu, 17 Jun 2021 16:21:41 -0700 Subject: [PATCH 216/305] Fix default DEFAULT_FILE_PATTERN in clang-tidy (#60212) Summary: Without the change, clang-tidy also checks folders like `.circleci/...` Example of the clang-tidy that looked into `.circleci` changes https://github.com/pytorch/pytorch/runs/2844682644?check_suite_focus=true [skip ci] Pull Request resolved: https://github.com/pytorch/pytorch/pull/60212 Reviewed By: seemethere Differential Revision: D29214728 Pulled By: zhouzhuojie fbshipit-source-id: fd53f7b2f7d88936264db1effdc06cc4fc271ca4 --- tools/clang_tidy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index 7574c4f3b538e..dd8b85e8465be 100755 --- a/tools/clang_tidy.py +++ b/tools/clang_tidy.py @@ -39,7 +39,7 @@ # NOTE: Clang-tidy cannot lint headers directly, because headers are not # compiled -- translation units are, of which there is one per implementation # (c/cc/cpp) file. -DEFAULT_FILE_PATTERN = re.compile(r".*\.c(c|pp)?") +DEFAULT_FILE_PATTERN = re.compile(r"^.*\.c(c|pp)?$") # Search for: # diff --git ... From 8e6798199577f217a2023fc20c6dfcf5678bccae Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Thu, 17 Jun 2021 16:22:11 -0700 Subject: [PATCH 217/305] .github: Disable clang-tidy for now (#60219) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60219 Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: driazati Differential Revision: D29214928 Pulled By: seemethere fbshipit-source-id: 20cf38ebfe77ed646e25293c577937c56bd930d3 --- .github/workflows/lint.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 52fcbc1a0f56b..68b9853868ff1 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -266,7 +266,9 @@ jobs: fi clang-tidy: - if: github.event_name == 'pull_request' + # if: github.event_name == 'pull_request' + # TODO: Fix clang-tidy, see https://github.com/pytorch/pytorch/issues/60192 + if: ${{ false }} runs-on: ubuntu-18.04 container: # ubuntu18.04-cuda10.2-py3.6-tidy11 From a727f655c8bb0e732a49597c9efde02631f3ce9a Mon Sep 17 00:00:00 2001 From: Serhat Yilmaz Date: Thu, 17 Jun 2021 16:23:15 -0700 Subject: [PATCH 218/305] [torch][segment_reduce] Support for multi dimension (cpu only) (#59951) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59951 Add support for multi-d input for cpu forward/backward implementation. Next step: Adding cuda support for multi-d input. Test Plan: Added unit tests. Reviewed By: ngimel Differential Revision: D29105457 fbshipit-source-id: a389ba4cc10f02434a336b8e7d36259f32552e11 --- aten/src/ATen/native/SegmentReduce.cpp | 178 ++++++++++++--------- aten/src/ATen/native/SegmentReduce.h | 3 +- aten/src/ATen/native/native_functions.yaml | 2 +- test/test_segment_reductions.py | 174 +++++++++++++++++--- 4 files changed, 254 insertions(+), 103 deletions(-) diff --git a/aten/src/ATen/native/SegmentReduce.cpp b/aten/src/ATen/native/SegmentReduce.cpp index c17cb67b5b60e..449942680a776 100644 --- a/aten/src/ATen/native/SegmentReduce.cpp +++ b/aten/src/ATen/native/SegmentReduce.cpp @@ -29,52 +29,63 @@ Tensor _segment_reduce_cpu_kernel( const Tensor& lengths, int64_t axis, const c10::optional& initial) { - int64_t batch_size = lengths.numel(); - auto output = at::empty({batch_size}, data.options()); + int64_t segment_count = lengths.numel(); + auto output_shape = data.sizes().vec(); + output_shape[axis] = segment_count; + auto output = at::empty(output_shape, data.options()); + int64_t stride_count = data.numel() / data.size(axis); const auto* lengths_data = lengths.data_ptr(); AT_DISPATCH_ALL_TYPES_AND2( kBFloat16, kHalf, data.scalar_type(), "_segment_reduce_cpu", ([&]() { auto* output_data = output.data_ptr(); const auto* values_data = data.data_ptr(); - int64_t k = 0; - for (int64_t i = 0; i < batch_size; ++i) { - // ===== step1: initialize starting value - scalar_t initial_value; - if (initial.has_value()) { - initial_value = initial.value().to(); - } else if (reduction == SegmentReductionType::MAX) { - initial_value = std::numeric_limits::lowest(); - } else if (reduction == SegmentReductionType::MEAN) { - initial_value = 0; - } - - // ===== step2: apply reduction - for (int64_t j = 0; j < lengths_data[i]; ++j) { - const auto data = values_data[k]; - // TODO: There is no need to branch with every element - if (reduction == SegmentReductionType::MAX) { - initial_value = at::_isnan(data) - ? data - : std::max(initial_value, data); + int64_t lengths_cum_sum = 0; + for (int64_t i = 0; i < segment_count; ++i) { + for (int64_t l = 0; l < stride_count; ++l) { + // ===== step1: initialize starting value + scalar_t initial_value; + if (initial.has_value()) { + initial_value = initial.value().to(); + } else if (reduction == SegmentReductionType::MAX) { + initial_value = std::numeric_limits::lowest(); } else if (reduction == SegmentReductionType::MEAN) { - initial_value = at::_isnan(data) ? data : (initial_value + data); + initial_value = 0; } - k++; - } - // ===== step3: finalize reduction - TORCH_CHECK(lengths_data[i] >= 0); - if (lengths_data[i] == 0 && !initial.has_value()) { - output_data[i] = static_cast(NAN); - continue; - } - output_data[i] = initial_value; - if (reduction == SegmentReductionType::MEAN && lengths_data[i] > 0 && - !at::_isnan(output_data[i])) { - output_data[i] = output_data[i] / lengths_data[i]; + // ===== step2: apply reduction + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = + ((lengths_cum_sum + j) * stride_count) + l; + const auto data = values_data[starting_index]; + // TODO: There is no need to branch with every element + if (reduction == SegmentReductionType::MAX) { + initial_value = at::_isnan(data) + ? data + : std::max(initial_value, data); + } else if (reduction == SegmentReductionType::MEAN) { + initial_value = + at::_isnan(data) ? data : (initial_value + data); + } + } + + // ===== step3: finalize reduction + TORCH_CHECK(lengths_data[i] >= 0); + int64_t output_index = (i * stride_count) + l; + if (lengths_data[i] == 0 && !initial.has_value()) { + output_data[output_index] = static_cast(NAN); + } else { + output_data[output_index] = initial_value; + if (reduction == SegmentReductionType::MEAN && + lengths_data[i] > 0 && + !at::_isnan(output_data[output_index])) { + output_data[output_index] = + output_data[output_index] / lengths_data[i]; + } + } } + lengths_cum_sum += lengths_data[i]; } })); @@ -86,56 +97,70 @@ Tensor _segment_reduce_cpu_backward_kernel( const Tensor& output_contig, const Tensor& data_contig, SegmentReductionType reduction, - const Tensor& lengths_contig) { + const Tensor& lengths_contig, + int64_t axis) { + int64_t segment_count = lengths_contig.numel(); + auto output_shape = data_contig.sizes().vec(); + output_shape[axis] = segment_count; auto grad_input = at::zeros({data_contig.sizes()}, grad_contig.options()); - int64_t batch_size = lengths_contig.numel(); + int64_t stride_count = data_contig.numel() / data_contig.size(axis); const auto* lengths_data = lengths_contig.data_ptr(); + // TODO: Swtich to TensorIterator for better maintainablility and readability AT_DISPATCH_ALL_TYPES_AND2( kBFloat16, kHalf, data_contig.scalar_type(), "_segment_reduce_cpu", ([&]() { - auto* output_data = output_contig.data_ptr(); - auto* grad_data = grad_contig.data_ptr(); - auto* grad_input_data = grad_input.data_ptr(); - const auto* values_data = data_contig.data_ptr(); - int64_t k = 0; - for (int64_t i = 0; i < batch_size; ++i) { - if (lengths_data[i] == 0) { + auto* output_data = output_contig.data_ptr(); + auto* grad_data = grad_contig.data_ptr(); + auto* grad_input_data = grad_input.data_ptr(); + const auto* values_data = data_contig.data_ptr(); + + int64_t lengths_cum_sum = 0; + for (int64_t i = 0; i < segment_count; ++i) { + if (lengths_data[i] == 0) { + continue; + } + + for (int64_t l = 0; l < stride_count; ++l) { + int64_t output_index = (i * stride_count) + l; + + if (reduction == SegmentReductionType::MAX) { + int64_t counter = 0; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; + if (at::_isnan(values_data[starting_index]) || + values_data[starting_index] == output_data[output_index]) { + grad_input_data[starting_index] = grad_data[output_index]; + counter++; + } + } + // Average gradient based on number of maximum elements in the + // segment + if (counter < 2) { continue; } - if (reduction == SegmentReductionType::MAX) { - int64_t counter = 0; - for (int64_t j = 0; j < lengths_data[i]; ++j) { - if (at::_isnan(values_data[k]) || - values_data[k] == output_data[i]) { - grad_input_data[k] = grad_data[i]; - counter++; - } - k++; - } - // Average gradient based on number of maximum elements in the - // segment - if (counter < 2) { - continue; - } - for (int64_t j = 0; j < lengths_data[i]; ++j) { - int64_t index = k - j - 1; - if (grad_input_data[index] > 0) { - grad_input_data[index] = grad_input_data[index] / counter; - } - } - } else if (reduction == SegmentReductionType::MEAN) { - auto grad_val = grad_data[i] / lengths_data[i]; - for (int64_t j = 0; j < lengths_data[i]; ++j) { - grad_input_data[k] = grad_val; - k++; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; + if (grad_input_data[starting_index] > 0) { + grad_input_data[starting_index] = + grad_input_data[starting_index] / counter; } } + } else if (reduction == SegmentReductionType::MEAN) { + auto grad_val = grad_data[output_index] / lengths_data[i]; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; + grad_input_data[starting_index] = grad_val; + } } + } + + lengths_cum_sum += lengths_data[i]; + } })); return grad_input; @@ -152,8 +177,7 @@ Tensor segment_reduce_kernel( bool unsafe, const c10::optional& initial) { axis = maybe_wrap_dim(axis, data.ndimension()); - TORCH_CHECK(axis == 0, "Currently only dim=0 is supported!"); - TORCH_CHECK(data.dim() == 1); + TORCH_CHECK(axis == 0, "Currently only dim=0 is supported! ", axis); TORCH_CHECK(data.numel() > 0); // length related checks @@ -169,7 +193,7 @@ Tensor segment_reduce_kernel( auto min_length = lengths_value.min().item(); TORCH_CHECK((min_length >= 0), "lengths contains negative value!"); TORCH_CHECK(min_length != 0 || initial.has_value()); - TORCH_CHECK(lengths_value.sum().item() == data.numel()); + TORCH_CHECK(lengths_value.sum().item() == data.size(axis)); } auto reduction = get_reduction_enum(reduce); @@ -203,7 +227,10 @@ Tensor _segment_reduce_backward_kernel( const Tensor& output, const Tensor& data, c10::string_view reduce, - const c10::optional& lengths) { + const c10::optional& lengths, + int64_t axis) { + axis = maybe_wrap_dim(axis, data.ndimension()); + TORCH_CHECK(axis == 0, "Currently only dim=0 is supported! ", axis); TORCH_CHECK( lengths.has_value(), "Currently only lengths based reduction is supported!") @@ -221,7 +248,8 @@ Tensor _segment_reduce_backward_kernel( output_contig, data_contig, reduction, - lengths_contig); + lengths_contig, + axis); } REGISTER_ARCH_DISPATCH( diff --git a/aten/src/ATen/native/SegmentReduce.h b/aten/src/ATen/native/SegmentReduce.h index 8bb7ece4e7d07..5eb87d798eeb4 100644 --- a/aten/src/ATen/native/SegmentReduce.h +++ b/aten/src/ATen/native/SegmentReduce.h @@ -22,7 +22,8 @@ using segment_reduce_backward_fn = Tensor (*)( const Tensor&, const Tensor&, SegmentReductionType, - const Tensor&); + const Tensor&, + int64_t); DECLARE_DISPATCH(segment_reduce_backward_fn, _segment_reduce_backward_stub); } // namespace native diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 953fdbf9b1433..cefe0726e4745 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -10125,7 +10125,7 @@ dispatch: CPU, CUDA: segment_reduce_kernel -- func: _segment_reduce_backward(Tensor grad, Tensor output, Tensor data, str reduce, *, Tensor? lengths=None) -> Tensor +- func: _segment_reduce_backward(Tensor grad, Tensor output, Tensor data, str reduce, *, Tensor? lengths=None, int axis=0) -> Tensor variants: function dispatch: CPU, CUDA: _segment_reduce_backward_kernel diff --git a/test/test_segment_reductions.py b/test/test_segment_reductions.py index d6a23fff63544..35e5deaa1bef7 100644 --- a/test/test_segment_reductions.py +++ b/test/test_segment_reductions.py @@ -1,8 +1,10 @@ +import numpy as np import torch from torch.testing._internal.common_device_type import ( instantiate_device_type_tests, dtypes, dtypesIfCUDA, + onlyCPU, ) from torch.testing._internal.common_utils import ( TestCase, @@ -12,29 +14,29 @@ class TestSegmentReductions(TestCase): - def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): - lengths = torch.tensor([1, 2, 3, 0], device=device) + def _test_common( + self, + reduction, + device, + dtype, + unsafe, + axis, + initial_value, + data_arr, + lengths_arr, + expected_arr, + expected_grad_arr, + check_backward, + ): + lengths = torch.tensor(lengths_arr, device=device) data = torch.tensor( - [1, float("nan"), 3, 4, 5, 5], + data_arr, device=device, dtype=dtype, requires_grad=True, ) - initial_value = 0 - if reduction == "max": - expected_result = torch.tensor( - [1, float("nan"), 5, initial_value], device=device, dtype=dtype - ) - expected_grad = torch.tensor( - [1, 1, 0, 0, 0.5, 0.5], device=device, dtype=dtype - ) - elif reduction == "mean": - expected_result = torch.tensor( - [1, float("nan"), 4.666, initial_value], device=device, dtype=dtype - ) - expected_grad = torch.tensor( - [1.0, 0.5, 0.5, 0.333, 0.333, 0.333], device=device, dtype=dtype - ) + expected_result = torch.tensor(expected_arr, device=device, dtype=dtype) + expected_grad = torch.tensor(expected_grad_arr, device=device, dtype=dtype) actual_result = torch.segment_reduce( data=data, reduce=reduction, @@ -47,8 +49,7 @@ def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): expected_result, actual_result, rtol=1e-02, atol=1e-05, equal_nan=True ) - # TODO: Remove this check once cuda backward support is implemented - if data.is_cuda: + if not check_backward: return # Test backward @@ -60,9 +61,11 @@ def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): # gradcheck does not work well with bfloat16 or fp16 cpu types # also there is small numerical difference with fp32 if dtype not in [torch.half, torch.bfloat16, torch.float]: - # gradcheck does not like "nan" input + # gradcheck does not like "nan" input, setting to random 10 + d_non_nan = np.nan_to_num(data_arr, nan=10) data = torch.tensor( - [1, 10, 3, 4, 5, 5], + # [10 if v == float("nan") else v for v in data], + d_non_nan, device=device, dtype=dtype, requires_grad=True, @@ -84,11 +87,130 @@ def _test_simple_1d(self, reduction, device, dtype, unsafe, axis): @dtypesIfCUDA(torch.half, torch.bfloat16, torch.float, torch.double) @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) def test_simple_1d(self, device, dtype): + lengths = [1, 2, 3, 0] + data = [1, float("nan"), 3, 4, 5, 5] + initial_value = 0 + + # TODO: Set this to true once cuda backward support is implemented + check_backward = device == "cpu" + for reduction in ("max", "mean"): - self._test_simple_1d(reduction, device, dtype, False, 0) - self._test_simple_1d(reduction, device, dtype, False, -1) - self._test_simple_1d(reduction, device, dtype, True, 0) - self._test_simple_1d(reduction, device, dtype, True, -1) + if reduction == "max": + expected_result = [1, float("nan"), 5, initial_value] + expected_grad = [1, 1, 0, 0, 0.5, 0.5] + elif reduction == "mean": + expected_result = [1, float("nan"), 4.666, initial_value] + expected_grad = [1.0, 0.5, 0.5, 0.333, 0.333, 0.333] + for axis in [0, -1]: + for unsafe in [True, False]: + self._test_common( + reduction, + device, + dtype, + unsafe, + axis, + initial_value, + data, + lengths, + expected_result, + expected_grad, + check_backward, + ) + + @onlyCPU + @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) + def test_multi_d_simple(self, device, dtype): + initial_value = 0 + axis = 0 + lengths = [1, 2, 3, 0] + data = [[1, 1], [float("nan"), 1], [3, float("nan")], [4, 1], [3, 2], [2, 3]] + + # TODO: Set this to true once cuda backward support is implemented + check_backward = device == "cpu" + + for reduction in ["max", "mean"]: + if reduction == "max": + expected_result = [ + [1, 1], + [float("nan"), float("nan")], + [4, 3], + [initial_value, initial_value], + ] + expected_grad = [ + [1, 1], + [1, 0], + [0, 1], + [1, 0], + [0, 0], + [0, 1], + ] + elif reduction == "mean": + expected_result = [ + [1, 1], + [float("nan"), float("nan")], + [3, 2], + [initial_value, initial_value], + ] + expected_grad = [ + [1.0, 1.0], + [0.5, 0.5], + [0.5, 0.5], + [0.333, 0.333], + [0.333, 0.333], + [0.333, 0.333], + ] + for unsafe in [True, False]: + self._test_common( + reduction, + device, + dtype, + unsafe, + axis, + initial_value, + data, + lengths, + expected_result, + expected_grad, + check_backward, + ) + + @onlyCPU + @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) + def test_multi_d(self, device, dtype): + initial_value = 0 + axis = 0 + lengths = [0, 2] + data = np.arange(20).reshape(2, 2, 5).tolist() + expected_grad = [] + + # TODO: calculate grad and check correctness + check_backward = False + + for reduction in ["max", "mean"]: + if reduction == "max": + expected_result = [ + np.full((2, 5), initial_value).tolist(), + np.max(data, axis=0).tolist(), + ] + elif reduction == "mean": + expected_result = [ + np.full((2, 5), initial_value).tolist(), + np.mean(data, axis=0).tolist(), + ] + for unsafe in [True, False]: + self._test_common( + reduction, + device, + dtype, + unsafe, + axis, + initial_value, + data, + lengths, + expected_result, + expected_grad, + check_backward, + ) instantiate_device_type_tests(TestSegmentReductions, globals()) From 6af5d00e4b6c57dcb6d49180e67f6dacdeaf7fb4 Mon Sep 17 00:00:00 2001 From: Serhat Yilmaz Date: Thu, 17 Jun 2021 16:23:15 -0700 Subject: [PATCH 219/305] [torch][segment_reduce] Add support for multi-dimensional input (cuda) (#60018) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60018 Same as title. This diff finishes cuda support for currently implemented reductions and input parameters. Next Steps: - Add support for sum/min - More testing and benchmarking - Cleanup - Update default values when length is 0 - Use TensorIterator - Update documentation Test Plan: Unit test to cover cuda forward path. Reviewed By: ngimel Differential Revision: D29135373 fbshipit-source-id: d070727eeb660f56782e7ac8a5b0798be688480a --- aten/src/ATen/native/SegmentReduce.cpp | 108 ++++----- aten/src/ATen/native/cuda/SegmentReduce.cu | 262 ++++++++++++++++++--- test/test_segment_reductions.py | 15 +- 3 files changed, 288 insertions(+), 97 deletions(-) diff --git a/aten/src/ATen/native/SegmentReduce.cpp b/aten/src/ATen/native/SegmentReduce.cpp index 449942680a776..fb6ebea07d044 100644 --- a/aten/src/ATen/native/SegmentReduce.cpp +++ b/aten/src/ATen/native/SegmentReduce.cpp @@ -65,25 +65,22 @@ Tensor _segment_reduce_cpu_kernel( ? data : std::max(initial_value, data); } else if (reduction == SegmentReductionType::MEAN) { - initial_value = - at::_isnan(data) ? data : (initial_value + data); + initial_value = initial_value + data; } } // ===== step3: finalize reduction TORCH_CHECK(lengths_data[i] >= 0); - int64_t output_index = (i * stride_count) + l; + if (lengths_data[i] == 0 && !initial.has_value()) { - output_data[output_index] = static_cast(NAN); - } else { - output_data[output_index] = initial_value; - if (reduction == SegmentReductionType::MEAN && - lengths_data[i] > 0 && - !at::_isnan(output_data[output_index])) { - output_data[output_index] = - output_data[output_index] / lengths_data[i]; - } + initial_value = static_cast(NAN); + } else if ( + reduction == SegmentReductionType::MEAN && + lengths_data[i] > 0 && !at::_isnan(initial_value)) { + initial_value = initial_value / lengths_data[i]; } + int64_t output_index = (i * stride_count) + l; + output_data[output_index] = initial_value; } lengths_cum_sum += lengths_data[i]; } @@ -114,53 +111,56 @@ Tensor _segment_reduce_cpu_backward_kernel( data_contig.scalar_type(), "_segment_reduce_cpu", ([&]() { - auto* output_data = output_contig.data_ptr(); - auto* grad_data = grad_contig.data_ptr(); - auto* grad_input_data = grad_input.data_ptr(); - const auto* values_data = data_contig.data_ptr(); - - int64_t lengths_cum_sum = 0; - for (int64_t i = 0; i < segment_count; ++i) { - if (lengths_data[i] == 0) { - continue; - } - - for (int64_t l = 0; l < stride_count; ++l) { - int64_t output_index = (i * stride_count) + l; - - if (reduction == SegmentReductionType::MAX) { - int64_t counter = 0; - for (int64_t j = 0; j < lengths_data[i]; ++j) { - int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; - if (at::_isnan(values_data[starting_index]) || - values_data[starting_index] == output_data[output_index]) { - grad_input_data[starting_index] = grad_data[output_index]; - counter++; - } - } - // Average gradient based on number of maximum elements in the - // segment - if (counter < 2) { + auto* output_data = output_contig.data_ptr(); + auto* grad_data = grad_contig.data_ptr(); + auto* grad_input_data = grad_input.data_ptr(); + const auto* values_data = data_contig.data_ptr(); + + int64_t lengths_cum_sum = 0; + for (int64_t i = 0; i < segment_count; ++i) { + if (lengths_data[i] == 0) { continue; } - for (int64_t j = 0; j < lengths_data[i]; ++j) { - int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; - if (grad_input_data[starting_index] > 0) { - grad_input_data[starting_index] = - grad_input_data[starting_index] / counter; + + for (int64_t l = 0; l < stride_count; ++l) { + int64_t output_index = (i * stride_count) + l; + + if (reduction == SegmentReductionType::MAX) { + int64_t counter = 0; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = + ((lengths_cum_sum + j) * stride_count) + l; + if (at::_isnan(values_data[starting_index]) || + values_data[starting_index] == output_data[output_index]) { + grad_input_data[starting_index] = grad_data[output_index]; + counter++; + } + } + // Average gradient based on number of maximum elements in the + // segment + if (counter < 2) { + continue; + } + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = + ((lengths_cum_sum + j) * stride_count) + l; + if (grad_input_data[starting_index] > 0) { + grad_input_data[starting_index] = + grad_input_data[starting_index] / counter; + } + } + } else if (reduction == SegmentReductionType::MEAN) { + auto grad_val = grad_data[output_index] / lengths_data[i]; + for (int64_t j = 0; j < lengths_data[i]; ++j) { + int64_t starting_index = + ((lengths_cum_sum + j) * stride_count) + l; + grad_input_data[starting_index] = grad_val; + } } } - } else if (reduction == SegmentReductionType::MEAN) { - auto grad_val = grad_data[output_index] / lengths_data[i]; - for (int64_t j = 0; j < lengths_data[i]; ++j) { - int64_t starting_index = ((lengths_cum_sum + j) * stride_count) + l; - grad_input_data[starting_index] = grad_val; - } - } - } - lengths_cum_sum += lengths_data[i]; - } + lengths_cum_sum += lengths_data[i]; + } })); return grad_input; diff --git a/aten/src/ATen/native/cuda/SegmentReduce.cu b/aten/src/ATen/native/cuda/SegmentReduce.cu index be3fc4003129d..36ae51a4eae48 100644 --- a/aten/src/ATen/native/cuda/SegmentReduce.cu +++ b/aten/src/ATen/native/cuda/SegmentReduce.cu @@ -10,6 +10,7 @@ namespace at { namespace native { +namespace { struct CustomMax { template __host__ __device__ __forceinline__ OutputT @@ -70,6 +71,169 @@ __global__ static void post_sum_div_kernel( } } +template +__global__ static void segment_reduce_forward_kernel( + SegmentReductionType reduction, + scalar_t* output_data, + scalar_t* values_data, + const int64_t* lengths_data, + const int64_t* lengths_cumsum_data, + const int64_t segment_count, + const int64_t stride_count, + bool is_initial_set, + scalar_t initial_value) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t row_id = idx / stride_count; + int64_t lane_id = idx % stride_count; + if (idx >= (segment_count * stride_count)) { + return; + } + int64_t offset_start = lengths_cumsum_data[row_id]; + int64_t offset_end = lengths_cumsum_data[row_id + 1]; + + // ===== step2: apply reduction + for (int64_t j = offset_start; j < offset_end; ++j) { + int64_t starting_index = (j * stride_count) + lane_id; + const auto data = values_data[starting_index]; + // TODO: There is no need to branch with every element + if (reduction == SegmentReductionType::MAX) { + initial_value = + at::_isnan(data) ? data : std::max(initial_value, data); + } else if (reduction == SegmentReductionType::MEAN) { + initial_value = initial_value + data; + } + } + + // ===== step3: finalize reduction + CUDA_KERNEL_ASSERT(lengths_data[row_id] >= 0); + if (lengths_data[row_id] == 0 && !is_initial_set) { + initial_value = static_cast(NAN); + } else if ( + reduction == SegmentReductionType::MEAN && lengths_data[row_id] > 0 && + !at::_isnan(initial_value)) { + initial_value = initial_value / lengths_data[row_id]; + } + int64_t output_index = (row_id * stride_count) + lane_id; + output_data[output_index] = initial_value; +} + +template +__global__ static void segment_reduce_backward_kernel( + SegmentReductionType reduction, + scalar_t* grad_input_data, + scalar_t* grad_data, + scalar_t* output_data, + const scalar_t* values_data, + const int64_t* lengths_data, + const int64_t* lengths_cumsum_data, + const int64_t segment_count, + const int64_t stride_count) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t row_id = idx / stride_count; + int64_t lane_id = idx % stride_count; + + if (idx >= (segment_count * stride_count)) { + return; + } + if (lengths_data[row_id] == 0) { + return; + } + + int64_t offset_start = lengths_cumsum_data[row_id]; + int64_t offset_end = lengths_cumsum_data[row_id + 1]; + + int64_t output_index = (row_id * stride_count) + lane_id; + + if (reduction == SegmentReductionType::MAX) { + int64_t counter = 0; + for (int64_t j = offset_start; j < offset_end; ++j) { + int64_t starting_index = (j * stride_count) + lane_id; + if (at::_isnan(values_data[starting_index]) || + values_data[starting_index] == output_data[output_index]) { + grad_input_data[starting_index] = grad_data[output_index]; + counter++; + } + } + // Average gradient based on number of maximum elements in the + // segment + if (counter < 2) { + return; + } + for (int64_t j = offset_start; j < offset_end; ++j) { + int64_t starting_index = (j * stride_count) + lane_id; + if (grad_input_data[starting_index] > 0) { + grad_input_data[starting_index] = + grad_input_data[starting_index] / counter; + } + } + } else if (reduction == SegmentReductionType::MEAN) { + auto grad_val = grad_data[output_index] / lengths_data[row_id]; + for (int64_t j = offset_start; j < offset_end; ++j) { + int64_t starting_index = (j * stride_count) + lane_id; + grad_input_data[starting_index] = grad_val; + } + } +} + +} // namespace + +Tensor _segment_reduce_cuda_backward_kernel( + const Tensor& grad_contig, + const Tensor& output_contig, + const Tensor& data_contig, + SegmentReductionType reduction, + const Tensor& lengths_contig, + int64_t axis) { + int64_t segment_count = lengths_contig.numel(); + auto output_shape = data_contig.sizes().vec(); + output_shape[axis] = segment_count; + auto grad_input = at::zeros({data_contig.sizes()}, grad_contig.options()); + + int64_t stride_count = data_contig.numel() / data_contig.size(axis); + const auto* lengths_data = lengths_contig.data_ptr(); + + auto offsets = _get_complete_sum(lengths_contig); + auto* offsets_data = offsets.data_ptr(); + + constexpr int threads_per_block = 256; + int64_t num_blocks = + ((segment_count * stride_count) + threads_per_block - 1) / + threads_per_block; + + num_blocks = std::max(num_blocks, (int64_t)1); + + // TODO: Swtich to TensorIterator for better maintainablility and readability + AT_DISPATCH_ALL_TYPES_AND2( + kBFloat16, + kHalf, + data_contig.scalar_type(), + "_segment_reduce_cpu", + ([&]() { + auto* output_data = output_contig.data_ptr(); + auto* grad_data = grad_contig.data_ptr(); + auto* grad_input_data = grad_input.data_ptr(); + const auto* values_data = data_contig.data_ptr(); + + segment_reduce_backward_kernel + <<>>( + reduction, + grad_input_data, + grad_data, + output_data, + values_data, + lengths_data, + offsets_data, + segment_count, + stride_count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + })); + + return grad_input; +} + Tensor _segment_reduce_cuda_kernel( SegmentReductionType reduction, const Tensor& data, @@ -77,14 +241,21 @@ Tensor _segment_reduce_cuda_kernel( int64_t axis, const c10::optional& initial) { int64_t segment_count = lengths.numel(); - auto output = at::empty({segment_count}, data.options()); + auto output_shape = data.sizes().vec(); + output_shape[axis] = segment_count; + auto output = at::empty(output_shape, data.options()); + + int64_t stride_count = data.numel() / data.size(axis); + const auto* lengths_data = lengths.data_ptr(); auto offsets = _get_complete_sum(lengths); auto* offsets_data_ptr = offsets.data_ptr(); constexpr int threads_per_block = 256; int64_t num_blocks = - (segment_count + threads_per_block - 1) / threads_per_block; + ((segment_count * stride_count) + threads_per_block - 1) / + threads_per_block; + num_blocks = std::max(num_blocks, (int64_t)1); auto* lengths_data_ptr = lengths.data_ptr(); @@ -97,48 +268,70 @@ Tensor _segment_reduce_cuda_kernel( auto* data_data_ptr = data.data_ptr(); auto* output_data_ptr = output.data_ptr(); - if (reduction == SegmentReductionType::MAX) { - CustomMax max_op{}; - scalar_t initial_value = initial.has_value() - ? initial.value().to() - : std::numeric_limits::lowest(); - CUB_WRAPPER( - cub::DeviceSegmentedReduce::Reduce, - data_data_ptr, - output_data_ptr, - segment_count, - offsets_data_ptr, - offsets_data_ptr + 1, - max_op, - initial_value, - at::cuda::getCurrentCUDAStream()); + // initialize starting value + scalar_t initial_value; + if (initial.has_value()) { + initial_value = initial.value().to(); + } else if (reduction == SegmentReductionType::MAX) { + initial_value = std::numeric_limits::lowest(); } else if (reduction == SegmentReductionType::MEAN) { - CustomSum sum_op{}; - scalar_t initial_value = initial.has_value() - ? initial.value().to() - : (scalar_t)0; - CUB_WRAPPER( - cub::DeviceSegmentedReduce::Reduce, - data_data_ptr, - output_data_ptr, - segment_count, - offsets_data_ptr, - offsets_data_ptr + 1, - sum_op, - initial_value, - at::cuda::getCurrentCUDAStream()); - - post_sum_div_kernel + initial_value = 0; + } + + if (output_shape.size() > 1) { + segment_reduce_forward_kernel <<>>( + reduction, output_data_ptr, + data_data_ptr, lengths_data_ptr, + offsets_data_ptr, segment_count, + stride_count, initial.has_value(), initial_value); C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + if (reduction == SegmentReductionType::MAX) { + CustomMax max_op{}; + CUB_WRAPPER( + cub::DeviceSegmentedReduce::Reduce, + data_data_ptr, + output_data_ptr, + segment_count, + offsets_data_ptr, + offsets_data_ptr + 1, + max_op, + initial_value, + at::cuda::getCurrentCUDAStream()); + } else if (reduction == SegmentReductionType::MEAN) { + CustomSum sum_op{}; + CUB_WRAPPER( + cub::DeviceSegmentedReduce::Reduce, + data_data_ptr, + output_data_ptr, + segment_count, + offsets_data_ptr, + offsets_data_ptr + 1, + sum_op, + initial_value, + at::cuda::getCurrentCUDAStream()); + + post_sum_div_kernel + <<>>( + output_data_ptr, + lengths_data_ptr, + segment_count, + initial.has_value(), + initial_value); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } } }); @@ -146,6 +339,9 @@ Tensor _segment_reduce_cuda_kernel( } REGISTER_DISPATCH(_segment_reduce_stub, &_segment_reduce_cuda_kernel); +REGISTER_DISPATCH( + _segment_reduce_backward_stub, + &_segment_reduce_cuda_backward_kernel); } // namespace native } // namespace at diff --git a/test/test_segment_reductions.py b/test/test_segment_reductions.py index 35e5deaa1bef7..8b349b4c0b3fd 100644 --- a/test/test_segment_reductions.py +++ b/test/test_segment_reductions.py @@ -4,7 +4,6 @@ instantiate_device_type_tests, dtypes, dtypesIfCUDA, - onlyCPU, ) from torch.testing._internal.common_utils import ( TestCase, @@ -90,9 +89,7 @@ def test_simple_1d(self, device, dtype): lengths = [1, 2, 3, 0] data = [1, float("nan"), 3, 4, 5, 5] initial_value = 0 - - # TODO: Set this to true once cuda backward support is implemented - check_backward = device == "cpu" + check_backward = True for reduction in ("max", "mean"): if reduction == "max": @@ -117,18 +114,16 @@ def test_simple_1d(self, device, dtype): check_backward, ) - @onlyCPU + @dtypesIfCUDA(torch.half, torch.bfloat16, torch.float, torch.double) @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) def test_multi_d_simple(self, device, dtype): initial_value = 0 + check_backward = True axis = 0 lengths = [1, 2, 3, 0] data = [[1, 1], [float("nan"), 1], [3, float("nan")], [4, 1], [3, 2], [2, 3]] - # TODO: Set this to true once cuda backward support is implemented - check_backward = device == "cpu" - - for reduction in ["max", "mean"]: + for reduction in ("max", "mean"): if reduction == "max": expected_result = [ [1, 1], @@ -174,7 +169,7 @@ def test_multi_d_simple(self, device, dtype): check_backward, ) - @onlyCPU + @dtypesIfCUDA(torch.half, torch.bfloat16, torch.float, torch.double) @dtypes(torch.half, torch.bfloat16, torch.float, torch.double) def test_multi_d(self, device, dtype): initial_value = 0 From 047925dac1c07a0ad2c86c281fac5610b084d1bd Mon Sep 17 00:00:00 2001 From: Eli Uriegas Date: Thu, 17 Jun 2021 16:25:58 -0700 Subject: [PATCH 220/305] .github: Run Windows CUDA build on pull requests (#60215) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60215 Signed-off-by: Eli Uriegas Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29214519 Pulled By: seemethere fbshipit-source-id: 58df5ee49cc5cd46f48938f023f87a6da958f3b6 --- .github/scripts/generate_ci_workflows.py | 4 ++++ .github/templates/windows_ci_workflow.yml.j2 | 11 +++++++++++ .../pytorch-win-vs2019-cuda10-cudnn7-py3.yml | 6 ++++-- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index d0e1293acb793..84f6b58be574e 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -25,6 +25,7 @@ def PyTorchWindowsWorkflow( test_runner_type: str, cuda_version: str, on_pull_request: bool = False, + only_build_on_pull_request: bool = False, num_test_shards: int = 1, ) -> PyTorchWorkflow: return { @@ -32,6 +33,7 @@ def PyTorchWindowsWorkflow( "test_runner_type": test_runner_type, "cuda_version": cuda_version, "on_pull_request": on_pull_request, + "only_build_on_pull_request": only_build_on_pull_request and on_pull_request, "num_test_shards": num_test_shards, } @@ -84,6 +86,8 @@ def generate_workflow_file( build_environment="pytorch-win-vs2019-cuda10-cudnn7-py3", cuda_version="10.1", test_runner_type=WINDOWS_CUDA_TEST_RUNNER, + on_pull_request=True, + only_build_on_pull_request=True ), PyTorchWindowsWorkflow( build_environment="pytorch-win-vs2019-cuda11-cudnn8-py3", diff --git a/.github/templates/windows_ci_workflow.yml.j2 b/.github/templates/windows_ci_workflow.yml.j2 index 2d09f9bb511fe..ec57ceaf6a27c 100644 --- a/.github/templates/windows_ci_workflow.yml.j2 +++ b/.github/templates/windows_ci_workflow.yml.j2 @@ -85,7 +85,11 @@ jobs: path: C:\w\build-results generate-test-matrix: +{%- if only_build_on_pull_request %} + if: ${{ github.repository_owner == 'pytorch' && github.event_name == 'push' }} +{%- else %} if: ${{ github.repository_owner == 'pytorch' }} +{%- endif %} runs-on: ubuntu-18.04 env: NUM_TEST_SHARDS: !{{ num_test_shards }} @@ -105,6 +109,9 @@ jobs: echo "::set-output name=matrix::${MATRIX}" test: +{%- if only_build_on_pull_request %} + if: ${{ github.event_name == 'push' }} +{%- endif %} runs-on: !{{ test_runner_type }} env: JOB_BASE_NAME: !{{ build_environment }}-test @@ -180,7 +187,11 @@ jobs: # logs (like test); we can always move it back to the other one, but it # doesn't create the best experience render_test_results: +{%- if only_build_on_pull_request %} + if: ${{ github.event_name == 'push' && always() }} +{%- else %} if: always() +{%- endif %} needs: - test runs-on: ubuntu-18.04 diff --git a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml index c1c476e1cfa85..8329a1904d804 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml @@ -4,6 +4,7 @@ name: Windows CI (pytorch-win-vs2019-cuda10-cudnn7-py3) on: + pull_request: push: branches: - master @@ -79,7 +80,7 @@ jobs: path: C:\w\build-results generate-test-matrix: - if: ${{ github.repository_owner == 'pytorch' }} + if: ${{ github.repository_owner == 'pytorch' && github.event_name == 'push' }} runs-on: ubuntu-18.04 env: NUM_TEST_SHARDS: 1 @@ -99,6 +100,7 @@ jobs: echo "::set-output name=matrix::${MATRIX}" test: + if: ${{ github.event_name == 'push' }} runs-on: windows.8xlarge.nvidia.gpu env: JOB_BASE_NAME: pytorch-win-vs2019-cuda10-cudnn7-py3-test @@ -172,7 +174,7 @@ jobs: # logs (like test); we can always move it back to the other one, but it # doesn't create the best experience render_test_results: - if: always() + if: ${{ github.event_name == 'push' && always() }} needs: - test runs-on: ubuntu-18.04 From 7e032f18cf1405804c4f787b05ea2de5e08a091e Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Thu, 17 Jun 2021 19:16:33 -0700 Subject: [PATCH 221/305] DOC Describes behavior for None in module.register_* (#60125) Summary: Fixes https://github.com/pytorch/pytorch/issues/45834 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60125 Reviewed By: zou3519 Differential Revision: D29196138 Pulled By: jbschlosser fbshipit-source-id: af736c0d66005ec33412860f00b233a5d2922137 --- torch/nn/modules/module.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py index 2c25a5926ddad..94a3eb10ffd6b 100644 --- a/torch/nn/modules/module.py +++ b/torch/nn/modules/module.py @@ -285,7 +285,8 @@ def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool name (string): name of the buffer. The buffer can be accessed from this module using the given name tensor (Tensor or None): buffer to be registered. If ``None``, then operations - that run on buffers, such as :attr:`cuda`, are ignored. + that run on buffers, such as :attr:`cuda`, are ignored. If ``None``, + the buffer is **not** included in the module's :attr:`state_dict`. persistent (bool): whether the buffer is part of this module's :attr:`state_dict`. @@ -330,7 +331,8 @@ def register_parameter(self, name: str, param: Optional[Parameter]) -> None: from this module using the given name param (Parameter or None): parameter to be added to the module. If ``None``, then operations that run on parameters, such as :attr:`cuda`, - are ignored. + are ignored. If ``None``, the parameter is **not** included in the + module's :attr:`state_dict`. """ if '_parameters' not in self.__dict__: raise AttributeError( From bcf8752fb2f467a80049414d18f28be289481b27 Mon Sep 17 00:00:00 2001 From: Patrick Date: Thu, 17 Jun 2021 21:01:04 -0700 Subject: [PATCH 222/305] updated launch bounds for trilinear 3d (#59999) Summary: Updates launch bounds for upsample_trilinear_3d forward and backward kernel to remove register spilling into local memory. Improves runtime for forward pass by 3-4x factor, backward pass has same runtime (probably different bottleneck). Timing data: (Using Nvidia Titan-V GPU) ![TrilinearTimingData](https://user-images.githubusercontent.com/22803332/121979658-72f19200-cd3f-11eb-9363-c00e2c4eea6d.PNG) Pull Request resolved: https://github.com/pytorch/pytorch/pull/59999 Reviewed By: zou3519 Differential Revision: D29185976 Pulled By: ngimel fbshipit-source-id: 0b2313e70e45c53938cd7262464d3aa4fab8da4a --- aten/src/ATen/native/cuda/UpSampleTrilinear3d.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/aten/src/ATen/native/cuda/UpSampleTrilinear3d.cu b/aten/src/ATen/native/cuda/UpSampleTrilinear3d.cu index 60bf6c250556c..856d21944126c 100644 --- a/aten/src/ATen/native/cuda/UpSampleTrilinear3d.cu +++ b/aten/src/ATen/native/cuda/UpSampleTrilinear3d.cu @@ -27,7 +27,7 @@ idx_3d(const size_t nc, } template -C10_LAUNCH_BOUNDS_1(1024) +C10_LAUNCH_BOUNDS_1(512) __global__ void upsample_trilinear3d_out_frame( const int n, const accscalar_t rdepth, @@ -111,7 +111,7 @@ __global__ void upsample_trilinear3d_out_frame( // Backward (adjoint) operation 1 <- 2 (accumulates) template -C10_LAUNCH_BOUNDS_1(1024) +C10_LAUNCH_BOUNDS_1(256) __global__ void upsample_trilinear3d_backward_out_frame( const int num_kernels, const accscalar_t rdepth, @@ -254,7 +254,7 @@ static void upsample_trilinear3d_out_cuda_template( const int num_kernels = output_depth * output_height * output_width; const int num_threads = std::min( - at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 1024); + at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 512); cudaStream_t stream = at::cuda::getCurrentCUDAStream(); AT_DISPATCH_FLOATING_TYPES_AND_HALF( @@ -322,7 +322,7 @@ static void upsample_trilinear3d_backward_out_cuda_template( const int num_kernels = output_depth * output_height * output_width; const int num_threads = std::min( - at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 1024); + at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 256); cudaStream_t stream = at::cuda::getCurrentCUDAStream(); AT_DISPATCH_FLOATING_TYPES_AND_HALF( From faf459f13ecb6fe8f6947959cbad7c6cf681bd19 Mon Sep 17 00:00:00 2001 From: Gisle Dankel Date: Thu, 17 Jun 2021 21:04:19 -0700 Subject: [PATCH 223/305] [Profiler] Fix memory profiler merge issue (#60037) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60037 The memory profiler was broken due to a mis-merge during rebase. Add lost line back. Reviewed By: ezyang Differential Revision: D29143469 fbshipit-source-id: c3bf0088ca12e7535eeddbede24e28201eccd5f4 --- torch/csrc/autograd/profiler_kineto.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/torch/csrc/autograd/profiler_kineto.cpp b/torch/csrc/autograd/profiler_kineto.cpp index e3d50227c86e3..7910a4ca79284 100644 --- a/torch/csrc/autograd/profiler_kineto.cpp +++ b/torch/csrc/autograd/profiler_kineto.cpp @@ -385,6 +385,7 @@ void prepareProfiler( std::set cpuTypes = { libkineto::ActivityType::CPU_OP, + libkineto::ActivityType::CPU_INSTANT_EVENT, #ifdef USE_KINETO_UPDATED libkineto::ActivityType::USER_ANNOTATION, #endif From 8b55e9feafb3eef8f0fc86ec16b52b36374fe158 Mon Sep 17 00:00:00 2001 From: Patrick Wang Date: Thu, 17 Jun 2021 21:12:09 -0700 Subject: [PATCH 224/305] removed cat, equal, and stack from autocast promote list (#59497) Summary: Fixes #{issue number} Pull Request resolved: https://github.com/pytorch/pytorch/pull/59497 Reviewed By: zou3519 Differential Revision: D29185909 Pulled By: ngimel fbshipit-source-id: db96239106d9e46a2704b8f457fd0463dacc1f5c --- aten/src/ATen/autocast_mode.cpp | 5 ----- docs/source/amp.rst | 3 --- torch/testing/_internal/autocast_test_lists.py | 6 +++--- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/aten/src/ATen/autocast_mode.cpp b/aten/src/ATen/autocast_mode.cpp index 922ea976c708e..b2d3617da51bc 100644 --- a/aten/src/ATen/autocast_mode.cpp +++ b/aten/src/ATen/autocast_mode.cpp @@ -433,14 +433,9 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) { KERNEL(ADD_NS(addcmul), "addcmul", Tensor (const Tensor &, const Tensor &, const Tensor &, const Scalar&), promote) KERNEL(ADD_NS(atan2), "atan2", Tensor (const Tensor &, const Tensor &), promote) KERNEL(ADD_NS(bilinear), "bilinear", Tensor (const Tensor &, const Tensor &, const Tensor &, const c10::optional&), promote) - KERNEL(ADD_NS(cat), "cat", Tensor (TensorList, int64_t), promote) - KERNEL(ADD_NS(cat), "cat.names", Tensor (TensorList, Dimname), promote) - KERNEL(ADD_NS(_cat), "_cat", Tensor (TensorList, int64_t), promote) KERNEL(ADD_NS(cross), "cross", Tensor (const Tensor &, const Tensor &, c10::optional), promote) KERNEL(ADD_NS(dot), "dot", Tensor (const Tensor &, const Tensor &), promote) - KERNEL(ADD_NS(equal), "equal", bool (const Tensor &, const Tensor &), promote) KERNEL(ADD_NS(index_put), "index_put", Tensor (const Tensor &, const torch::List>&, const Tensor &, bool), promote) - KERNEL(ADD_NS(stack), "stack", Tensor (TensorList, int64_t), promote) KERNEL(ADD_NS(tensordot), "tensordot", Tensor (const Tensor &, const Tensor &, IntArrayRef, IntArrayRef), promote) KERNEL(ADD_NS(scatter_add), "scatter_add", Tensor (const Tensor&, int64_t, const Tensor&, const Tensor&), promote) diff --git a/docs/source/amp.rst b/docs/source/amp.rst index 539fd07be97c5..8e376c6c2a1bd 100644 --- a/docs/source/amp.rst +++ b/docs/source/amp.rst @@ -187,13 +187,10 @@ autocast casts all inputs to ``float32`` and runs the op in ``float32``. ``addcmul``, ``atan2``, ``bilinear``, -``cat``, ``cross``, ``dot``, -``equal``, ``index_put``, ``scatter_add``, -``stack``, ``tensordot`` Some ops not listed here (e.g., binary ops like ``add``) natively promote diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py index e7ac13d5ce7fb..cfb1f33b7cf1f 100644 --- a/torch/testing/_internal/autocast_test_lists.py +++ b/torch/testing/_internal/autocast_test_lists.py @@ -61,6 +61,9 @@ def __init__(self, dev): ("add", pointwise0_fp32 + pointwise1_fp16, torch.float32), ("div", pointwise0_fp32 + pointwise1_fp16, torch.float32), ("mul", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("cat", (pointwise0_fp16 + pointwise1_fp32,), torch.float32), + ("equal", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("stack", (pointwise0_fp16 + pointwise1_fp32,), torch.float32), ] self.methods_expect_builtin_promote = [ ("__eq__", pointwise0_fp32 + pointwise1_fp16, torch.bool), @@ -188,14 +191,11 @@ def __init__(self, dev): torch.randn((1,), dtype=torch.float32, device=dev))), ("cross", (torch.randn(3, dtype=torch.float32, device=dev), torch.randn(3, dtype=torch.float16, device=dev))), - ("cat", (pointwise0_fp16 + pointwise1_fp32,)), ("dot", pointwise0_fp16 + pointwise1_fp32), - ("equal", pointwise0_fp32 + pointwise1_fp16), ("index_put", pointwise0_fp32 + ((torch.tensor([1], device=dev, dtype=torch.long),), torch.randn(1, device=dev, dtype=torch.float16))), ("index_put", pointwise0_fp16 + ((torch.tensor([1], device=dev, dtype=torch.long),), torch.randn(1, device=dev, dtype=torch.float32))), - ("stack", (pointwise0_fp16 + pointwise1_fp32,)), ("tensordot", (torch.randn((2, 2, 2), dtype=torch.float32, device=dev), torch.randn((2, 2, 2), dtype=torch.float16, device=dev))), ("scatter_add", (torch.zeros(2, 2, 2, dtype=torch.float32, device=dev), From 38c31168133e3a19531581d2624496509edcef03 Mon Sep 17 00:00:00 2001 From: Ruilin Chen Date: Thu, 17 Jun 2021 22:22:57 -0700 Subject: [PATCH 225/305] [hierarchical sharding 5/n] enable table-wise -> col-wise sharding in embedding table lookup Summary: This diff add table-wise -> col-wise sharding support in GroupedShardedEmbeddingBag. Changes includes: 1. Add necessary member variables set up. 2. Create new fast kernel and add fast kernel lookup support 3. Add intra-host all2all and cross-host all2all logic. Test Plan: UT ``` buck test mode/dev-nosan //caffe2/torch/fb/training_toolkit/backend/tests:test_model_materializer_full_sync_spawn ``` ``` buck test caffe2/torch/fb/hpc/tests:model_sharder_test ``` QPS check: ``` buck run mode/dev-nosan -c python.package_style=inplace caffe2/torch/fb/training_toolkit/examples:sync_sgd_local_driver -- prod-preset --num-trainers 32 --use-shrunk-model false --model-version=inline_cvr_dec_2020 --fast-kernel table_batched --max-batches 10000 --num-dpp-worker-threads 16 --num-readers 100 --hpc-identity ads_model_platform --table-partition hierarchical_based --hierarchical-options "["table_based", "column_based"]" --flow-entitlement ads_global_qps ``` with diff: dec inline_cvr: table-wise -> table-wise (82K): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_d0a0cba5?version=0&tab=status&env=PRODUCTION table-wise -> column-wise (80k): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_b1ac5873 column-wise: dec inline_cvr: gpu trace: https://www.internalfb.com/intern/perfdoctor/trace_view?filepath=tree%2Ftraces%2Fdynocli%2F0%2F1623827677%2F127.0.0.1%2Flibkineto_activities_4550.json.gz&bucket=gpu_traces https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_a79e1522 (81k) https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_2dacc13e (88k) row-wise(62k): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_4e349cab table-wise(90k): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_5d51b608 10x ctr_mbl_feed: ``` buck run mode/dev-nosan -c python.package_style=inplace caffe2/torch/fb/training_toolkit/examples:sync_sgd_local_driver -- prod-preset --num-trainers 128 --use-shrunk-model false --model-version=ctr_mbl_oct_2020_10x_3tb --num-dpp-worker-threads 16 --num-readers 200 --fast-kernel table_batched --max-batches 5000000 --hpc-identity ads_model_platform --table-partition column_based --flow-entitlement ads_global_tc_mimo ``` column-wise: https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_f05fb306?version=0&tab=status&env=PRODUCTION (290k) w/o diff: dec inline_cvr: column-wise (87K): gpu trace: https://www.internalfb.com/intern/perfdoctor/trace_view?filepath=tree%2Ftraces%2Fdynocli%2F0%2F1623864444%2F127.0.0.1%2Flibkineto_activities_4451.json.gz&bucket=gpu_traces https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_e1315f14 row-wise (60k): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_8fcc0adf table-wise (91k): https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_cb94ff41 10x ctr_mbl_feed: https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_203ef35b?version=0&tab=status&env=PRODUCTION (281k) NE check(use deterministic reading D28711400) ``` buck run mode/dev-nosan -c python.package_style=inplace caffe2/torch/fb/training_toolkit/examples:sync_sgd_local_driver -- prod-preset --num-trainers 32 --use-shrunk-model false --model-version=inline_cvr_dec_2020 --fast-kernel table_batched --max-batches 100000 --num-dpp-worker-threads 16 --num-readers 64 --hpc-identity ads_model_platform --table-partition hierarchical_based --hierarchical-options "[table_based, column_based]" --flow-entitlement ads_global_qps --use-deterministic-model --use-deterministic-reading --model-entity-id 995557193 ``` w/o this diff: ``` I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: ne-ne|lifetime_ne 0.8660048340401448 I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: ne-ne|window_ne 0.8660048340401447 I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: qps-qps|total_examples 1867776.0 I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: qps-qps|window_qps 491.5199890136719 ``` https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_15bc6243?version=0&tab=status&env=PRODUCTION w this diff: ``` I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: ne-ne|lifetime_ne 0.8660048340401448 I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: ne-ne|window_ne 0.8660048340401447 I0611 12:19:18.766000 647 print_publisher.py:33 master ] Publishing batch metrics: qps-qps|total_examples 1867776.0 ``` https://www.internalfb.com/mast/job/tsm_ruilinchen-SparseNNApplication_15bc6243?version=0&tab=status&env=PRODUCTION Reviewed By: JadeNie Differential Revision: D28689126 fbshipit-source-id: 1c7879d4e3ee2b90aaf2a89e87f7b827d54173b3 --- torch/distributed/distributed_c10d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index a43c57cbd5c46..c77826690d362 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -244,7 +244,7 @@ def _store_based_barrier(rank, store, timeout): ) ) - logger.info(f"Rank {rank}: Completed store-based barrier for {world_size} nodes.") + logger.info(f"Rank {rank}: Completed store-based barrier for key:{store_key} with {world_size} nodes.") def _rank_not_in_group(group: ProcessGroup): From ecc37184a54a5a2c9cf5814af6c4a3b5dc546fe4 Mon Sep 17 00:00:00 2001 From: driazati Date: Thu, 17 Jun 2021 23:02:18 -0700 Subject: [PATCH 226/305] Fix clang-tidy path filtering (#60225) Summary: PR https://github.com/pytorch/pytorch/issues/60048 neglected to include the `--paths` option for file filtering, so it ended up passing every changed file in the diff to clang-tidy (cpp files outside `torch/csrc/`, yaml/sh files, etc.). This adds that back in to make the filtering work properly again. Tested it manually by printing out the files to lint and running ```bash curl -L https://github.com/pytorch/pytorch/pull/60018.diff > diff python tools/clang_tidy.py --diff-file diff --paths torch/csrc/ curl -L https://github.com/pytorch/pytorch/pull/60222.diff > diff python tools/clang_tidy.py --diff-file diff --paths torch/csrc/ ``` Should fix https://github.com/pytorch/pytorch/issues/60192 and fix https://github.com/pytorch/pytorch/issues/60193, the files tripping errors there shouldn't have been passed to clang-tidy in the first place (supporting aten/ for clang-tidy is a separate task) Pull Request resolved: https://github.com/pytorch/pytorch/pull/60225 Reviewed By: zhouzhuojie Differential Revision: D29216251 Pulled By: driazati fbshipit-source-id: b5d7fb7161d33eb7958a6f1ccc25809942045209 --- tools/clang_tidy.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index dd8b85e8465be..0087b68e70339 100755 --- a/tools/clang_tidy.py +++ b/tools/clang_tidy.py @@ -142,7 +142,6 @@ def find_changed_lines(diff: str) -> Dict[str, List[Tuple[int, int]]]: for file, start, end in matches: start_line, _ = start.split(",") end_line, _ = end.split(",") - print(file, start_line, end_line) files[file].append((start_line, end_line)) @@ -330,6 +329,11 @@ def main() -> None: if options.diff_file: with open(options.diff_file, "r") as f: changed_files = find_changed_lines(f.read()) + changed_files = { + filename: v + for filename, v in changed_files.items() + if any(filename.startswith(path) for path in options.paths) + } line_filters = [ {"name": name, "lines": lines} for name, lines, in changed_files.items() ] From 5609c2e59cdea2c6eefc7461d4ad3e4fe0792832 Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Fri, 18 Jun 2021 03:39:22 -0700 Subject: [PATCH 227/305] Adds an OpInfo note (#57428) Summary: Like the title says. The OpInfo pattern can be confusing when first encountered, so this note links the Developer Wiki and tracking issue, plus elaborates on the goals and structure of the OpInfo pattern. cc imaginary-person, who I can't add as a reviewer, unfortunately Pull Request resolved: https://github.com/pytorch/pytorch/pull/57428 Reviewed By: SplitInfinity Differential Revision: D29221874 Pulled By: mruberry fbshipit-source-id: aa73228748c9c96eadf2b2397a8b2ec31383971e --- torch/testing/_internal/common_device_type.py | 246 ++++++++++++------ .../_internal/common_methods_invocations.py | 244 +++++++++++++++++ 2 files changed, 408 insertions(+), 82 deletions(-) diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py index 7d814d25c645a..a06aaf9a280ce 100644 --- a/torch/testing/_internal/common_device_type.py +++ b/torch/testing/_internal/common_device_type.py @@ -22,103 +22,188 @@ except ImportError: HAS_PSUTIL = False -# Note: Generic Device-Type Testing +# Note [Writing Test Templates] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # -# [WRITING TESTS] +# This note was written shortly after the PyTorch 1.9 release. +# If you notice it's out-of-date or think it could be improved then please +# file an issue. # -# Write your test class as usual except: -# (1) Each test method should have one of following five signatures: +# PyTorch has its own framework for instantiating test templates. That is, for +# taking test classes that look similar to unittest or pytest +# compatible test classes and optionally doing the following: # -# (1a) testX(self, device) +# - instantiating a version of the test class for each available device type +# (often the CPU, CUDA, and META device types) +# - further instantiating a version of each test that's always specialized +# on the test class's device type, and optionally specialized further +# on datatypes or operators # -# (1b) @deviceCountAtLeast() -# testX(self, devices) +# This functionality is similar to pytest's parametrize functionality +# (see https://docs.pytest.org/en/6.2.x/parametrize.html), but with considerable +# additional logic that specializes the instantiated test classes for their +# device types (see CPUTestBase and CUDATestBase below), supports a variety +# of composable decorators that allow for test filtering and setting +# tolerances, and allows tests parametrized by operators to instantiate +# only the subset of device type x dtype that operator supports. # -# (1c) @dtypes( or ) -# testX(self, device, dtype) +# This framework was built to make it easier to write tests that run on +# multiple device types, multiple datatypes (dtypes), and for multiple +# operators. It's also useful for controlling which tests are fun. For example, +# only tests that use a CUDA device can be run on platforms with CUDA. +# Let's dive in with an example to get an idea for how it works: # -# (1d) @deviceCountAtLeast() -# @dtypes( or ) -# testX(self, devices, dtype) +# -------------------------------------------------------- +# A template class (looks like a regular unittest TestCase) +# class TestClassFoo(TestCase): # -# (1e) @ops() -# testX(self, device, dtype, op) +# # A template test that can be specialized with a device +# # NOTE: this test case is not runnably by unittest or pytest because it +# # accepts an extra positional argument, "device", they do not understand +# def test_bar(self, device): +# pass # -# Note that the decorators are required for signatures 1b--1e. +# # Function that instantiates a template class and its tests +# instantiate_device_type_tests(TestCommon, globals()) +# -------------------------------------------------------- # -# When a test like (1a) is called it will be given a device string, -# like 'cpu' or 'cuda:0.' +# In the above code example we see a template class and a single test template +# that can be instantiated with a device. The function +# instantiate_device_type_tests(), called at file scope, instantiates +# new test classes, one per available device type, and new tests in those +# classes from these templates. It actually does this by removing +# the class TestClassFoo and replacing it with classes like TestClassFooCPU +# and TestClassFooCUDA, instantiated test classes that inherit from CPUTestBase +# and CUDATestBase respectively. Additional device types, like XLA, +# (see https://github.com/pytorch/xla) can further extend the set of +# instantiated test classes to create classes like TestClassFooXLA. # -# Tests like (1b) are called with a list of device strings, like -# ['cuda:0', 'cuda:1']. The first device string will be the -# primary device. These tests will be skipped if the device type -# has fewer available devices than the argument to @deviceCountAtLeast. +# The test template, test_bar(), is also instantiated. In this case the template +# is only specialized on a device, so (depending on the available device +# types) it might become test_bar_cpu() in TestClassFooCPU and test_bar_cuda() +# in TestClassFooCUDA. We can think of the instantiated test classes as +# looking like this: # -# Tests like (1c) are called with a device string and a torch.dtype (or -# a tuple of torch.dtypes) from the list of dtypes (or list of tuples -# of torch.dtypes) specified in the @dtypes decorator. Device-specific -# dtype overrides can be specified using @dtypesIfCPU and @dtypesIfCUDA. +# -------------------------------------------------------- +# # An instantiated test class for the CPU device type +# class TestClassFooCPU(CPUTestBase): # -# Tests like (1d) take a devices argument like (1b) and a dtype -# argument from (1c). +# # An instantiated test that calls the template with the string representation +# # of a device from the test class's device type +# def test_bar_cpu(self): +# test_bar(self, 'cpu') # -# Tests like (1e) are instantiated for each provided OpInfo instance, -# with dtypes specified by the OpInfo instance (unless overridden with -# an additional @dtypes decorator). +# # An instantiated test class for the CUDA device type +# class TestClassFooCUDA(CUDATestBase): # -# (2) Prefer using test decorators defined in this file to others. -# For example, using the @skipIfNoLapack decorator instead of the -# @skipCPUIfNoLapack will cause the test to not run on CUDA if -# LAPACK is not available, which is wrong. If you need to use a decorator -# you may want to ask about porting it to this framework. +# # An instantiated test that calls the template with the string representation +# # of a device from the test class's device type +# def test_bar_cuda(self): +# test_bar(self, 'cuda:0') +# -------------------------------------------------------- # -# See the TestTorchDeviceType class in test_torch.py for an example. +# These instantiated test classes are discoverable and runnable by both +# unittest and pytest. One thing that may be confusing, however, is that +# attempting to run "test_bar" will not work, despite it appearing in the +# original template code. This is because "test_bar" is no longer discoverable +# after instantiate_device_type_tests() runs, as the above snippet shows. +# Instead "test_bar_cpu" and "test_bar_cuda" may be run directly, or both +# can be run with the option "-k test_bar". # -# [RUNNING TESTS] +# Removing the template class and adding the instantiated classes requires +# passing "globals()" to instantiate_device_type_tests(), because it +# edits the file's Python objects. # -# After defining your test class call instantiate_device_type_tests on it -# and pass in globals() for the second argument. This will instantiate -# discoverable device-specific test classes from your generic class. It will -# also hide the tests in your generic class so they're not run. +# As mentioned, tests can be additionally parametrized on dtypes or +# operators. Datatype parametrization uses the @dtypes decorator and +# require a test template like this: # -# If you device-generic test class is TestClass then new classes with names -# TestClass will be created for each available device type. -# TestClassCPU and TestClassCUDA, for example. Tests in these classes also -# have the device type and dtype, if provided, appended to their original -# name. testX, for instance, becomes testX_ or -# testX__. +# -------------------------------------------------------- +# # A template test that can be specialized with a device and a datatype (dtype) +# @dtypes(torch.float32, torch.int64) +# def test_car(self, device, dtype) +# pass +# -------------------------------------------------------- # -# More concretely, TestTorchDeviceType becomes TestTorchDeviceTypeCPU, -# TestTorchDeviceTypeCUDA, ... test_diagonal in TestTorchDeviceType becomes -# test_diagonal_cpu, test_diagonal_cuda, ... test_erfinv, which accepts a dtype, -# becomes test_erfinv_cpu_float, test_erfinv_cpu_double, test_erfinv_cuda_half, -# ... +# If the CPU and CUDA device types are available this test would be +# instantiated as 4 tests that cover the cross-product of the two dtypes +# and two device types: # -# In short, if you write a test signature like -# def textX(self, device) -# You are effectively writing -# def testX_cpu(self, device='cpu') -# def textX_cuda(self, device='cuda') -# def testX_xla(self, device='xla') -# ... +# - test_car_cpu_float32 +# - test_car_cpu_int64 +# - test_car_cuda_float32 +# - test_car_cuda_int64 +# +# The dtype is passed as a torch.dtype object. +# +# Tests parametrized on operators (actually on OpInfos, more on that in a +# moment...) use the @ops decorator and require a test template like this: +# -------------------------------------------------------- +# # A template test that can be specialized with a device, dtype, and OpInfo +# @ops(op_db) +# def test_car(self, device, dtype, op) +# pass +# -------------------------------------------------------- # -# These tests can be run directly like normal tests: -# "python test_torch.py TestTorchDeviceTypeCPU.test_diagonal_cpu" +# See the documentation for the @ops decorator below for additional details +# on how to use it and see the note [OpInfos] in +# common_methods_invocations.py for more details on OpInfos. # -# All the tests for a particular device type can be run using the class, and -# other collections of tests can be run using pytest filtering, like +# A test parametrized over the entire "op_db", which contains hundreds of +# OpInfos, will likely have hundreds or thousands of instantiations. The +# test will be instantiated on the cross-product of device types, operators, +# and the dtypes the operator supports on that device type. The instantiated +# tests will have names like: # -# "pytest test_torch.py -k 'test_diag'" +# - test_car_add_cpu_float32 +# - test_car_sub_cuda_int64 # -# which will run test_diag on every available device. +# The first instantiated test calls the original test_car() with the OpInfo +# for torch.add as its "op" argument, the string 'cpu' for its "device" argument, +# and the dtype torch.float32 for is "dtype" argument. The second instantiated +# test calls the test_car() with the OpInfo for torch.sub, a CUDA device string +# like 'cuda:0' or 'cuda:1' for its "device" argument, and the dtype +# torch.int64 for its "dtype argument." # -# To specify particular device types the 'and' keyword can be used: +# Clever test filtering can be very useful when working with parametrized +# tests. "-k test_car" would run every instantiated variant of the test_car() +# test template, and "-k test_car_add" runs every variant instantiated with +# torch.add. # -# "pytest test_torch.py -k 'test_erfinv and cpu'" +# It is important to use the passed device and dtype as appropriate. Use +# helper functions like make_tensor() that require explicitly specifying +# the device and dtype so they're not forgotten. # -# will run test_erfinv on all cpu dtypes. +# Test templates can use a variety of composable decorators to specify +# additional options and requirements, some are listed here: # -# [ADDING A DEVICE TYPE] +# - @deviceCountAtLeast() +# Passes a list of strings representing all available devices of +# the test class's device type as the test template's "device" argument. +# If there are a fewer devices than the value passed to the decorator +# the test is skipped. +# - @dtypes() +# In addition to accepting multiple dtypes, the @dtypes decorator +# can accept a sequence of tuple pairs of dtypes. The test template +# will be called with each tuple for its "dtype" argument. +# - @onlyOnCPUAndCUDA +# Skips the test if the device is not a CPU or CUDA device +# - @onlyCPU +# Skips the test if the device is not a CPU device +# - @onlyCUDA +# Skips the test if the device is not a CUDA device +# - @skipCPUIfNoLapack +# Skips the test if the device is a CPU device and LAPACK is not installed +# - @skipCPUIfNoMkl +# Skips the test if the device is a CPU device and MKL is not installed +# - @skipCUDAIfNoMagma +# Skips the test if the device is a CUDA device and MAGMA is not installed +# - @skipCUDAIfRocm +# Skips the test if the device is a CUDA device and ROCm is being used + + +# Note [Adding a Device Type] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~ # # To add a device type: # @@ -135,9 +220,11 @@ # # setUpClass is called AFTER tests have been created and BEFORE and ONLY IF # they are run. This makes it useful for initializing devices and dependencies. -# + + # Note [Overriding methods in generic tests] # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# # Device generic tests look a lot like normal test classes, but they differ # from ordinary classes in some important ways. In particular, overriding # methods in generic tests doesn't work quite the way you expect. @@ -162,11 +249,6 @@ # is to either (1) add your functionality to TestCase and make it toggled # by a class attribute, or (2) create your own subclass of TestCase, and # then inherit from it for your generic test. -# - -# List of device type test bases that can be used to instantiate tests. -# See below for how this list is populated. If you're adding a device type -# you should check if it's available and (if it is) add it to this list. def _construct_test_name(test_name, op, device_type, dtype): @@ -587,19 +669,19 @@ class OpDTypes(Enum): none = 5 # Instantiate no dtype variants (the dtype kwarg will be None) -# Decorator that defines the ops a test should be run with -# The test signature must be: -# (self, device, dtype, op) -# For example: +# Decorator that defines the OpInfos a test template should be instantiated for. +# +# Example usage: +# # @ops(unary_ufuncs) # def test_numerics(self, device, dtype, op): # # -# This will instantiate variants of test_numerics for each given operator, -# on each device that operator supports, and for every dtype supported by +# This will instantiate variants of test_numerics for each given OpInfo, +# on each device the OpInfo's operator supports, and for every dtype supported by # that operator. There are a few caveats to the dtype rule, explained below. # -# First, if the OpInfo defines "default_test_dtypes" then then the test +# First, if the OpInfo defines "default_test_dtypes" then the test # is instantiated for the intersection of default_test_dtypes and the # dtypes the operator supports. Second, the @ops decorator can accept two # additional arguments, "dtypes" and "allowed_dtypes". If "dtypes" is specified diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 686f4830966cf..385950c4ffabb 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -179,6 +179,232 @@ def _getattr_qual(obj, name, default=_NOTHING): else: raise +# Note [OpInfos] +# ~~~~~~~~~~~~~~ +# +# This note was written shortly after the PyTorch 1.9 release. +# If you notice it's out-of-date or think it could be improved then please +# file an issue. +# +# See also: the OpInfo tracker (https://github.com/pytorch/pytorch/issues/54261) +# See also: "Writing Test Templates" in common_device_type.py to learn how to +# parametrize a test template using OpInfos. +# +# An OpInfo is a collection of metadata related to a PyTorch operator. This +# metadata is used to generate tests that validate properties of the operator, +# like if it implements the correct gradient formula. +# +# WHY OPINFOS? +# ~~~~~~~~~~~~ +# +# OpInfos are principally intended to do two things: +# +# 1) to simplify testing an operator +# 2) to allow systems (like autograd, torchscript, fx, nnc...) to test +# against every PyTorch operator +# +# Both these goals are still a work in progress. Not every operator has an +# OpInfo, and some operator tests still have to be written manually. +# +# The utility of OpInfos can also be motivated from a different perspective. +# PyTorch is a complicated framework with many interrelated systems, too +# many for any one person to keep track of. An OpInfo can be thought of as the +# interface between an operator implementer and those other systems. Instead of +# requiring the implementer of torch.foo understand how to test its forward +# mode AD or NNC support that's typically handled automatically just by +# defining an OpInfo. This is a helpful perspective to have, because it's often +# surprising to OpInfo writers that just implementing an OpInfo typically can't +# verify an operator is actually implemented correctly. "If an OpInfo doesn't +# validate my op works as expected, what's the point of it?" But the point of +# it is that it lets engineers focus on testing their operator logic instead +# of having to write tests for how the operator interacts with each of +# PyTorch's many systems. And, OK, sometimes it validates your op works +# the way you want and all you have to do is write an OpInfo and you're done +# testing... more on that below. +# +# WHAT'S AN OPINFO? +# ~~~~~~~~~~~~~~~~~ +# +# So what is an OpInfo? It's a Python class that describes an operator's properties, +# like which dtypes it supports on the CPU and whether it has any aliases. +# These properties can be divided into three categories: +# +# 1) Metadata describing the operator, like the operator's name and if it +# "supports" the out kwarg. +# 2) Test directives, like "skips" that tell the test suite to skip some +# tests. +# 3) A "sample inputs" function that generates valid inputs for the operator. +# +# OpInfo attributes are described in more detail below. +# +# THE SAMPLE INPUTS FUNCTION +# ~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The "sample inputs" function merits special elaboration. This function is +# crucial to testing with OpInfos. A typical OpInfo test has to treat the operator +# as a black box. There's no structure for the test to understand or exploit. +# Without "sample inputs" it wouldn't even know how to call the OpInfo's +# operator. The sample input function saves the day by providing different +# "SampleInputs" that can be used to call the operator. A sample input +# function should have the following signature: +# +# def sample_inputs_foo(op_info, device, dtype, requires_grad, **kwargs): +# +# And should return a list of SampleInputs (see the class description above). +# Each SampleInput defines an "input", "args", "kwargs", +# an "output_process_fn_grad" function, the "broadcasts_input" bool and +# a "name". +# +# The "input" is the first argument to the operator, or the tensor that +# the method or inplace variants of the operator should be called on, and +# should be on the requested device, of the requested dtype, and its +# requires_grad attribute should be set to the requires_grad argument. +# +# "args" should contain positional arguments, and "kwargs" keyword arguments. +# +# "output_process_fn_grad" has an interesting name. It's a function that maps +# the operator's output (when given the input, args, and kwargs) to the +# portion of the output to gradcheck. For example, consider an operator +# like torch.linalg.slogdet +# (https://pytorch.org/docs/master/generated/torch.linalg.slogdet.html). +# This operator returns a tuple of two tensors, but the first tensor +# cannot be backwarded through. Its "output_process_fn_grad" filters +# this output tuple to just the second argument, which we can call backward +# on. Functions that produce a single tensor can ignore this argument. +# +# "broadcasts_input" is a bool indicated if the SampleInput causes the operator +# to broadcast the "input" argument. This is important for tests to understand +# because inplace variants of operations throw a runtime error if they +# would broadcast their input arguments, so tests that work with inplace +# variants filter SampleInputs that broadcast their input. +# +# "name" is a string that's just used for debugging. It appears when printing +# the SampleInput. +# +# OPINFO FILE ORGANIZATION +# ~~~~~~~~~~~~~~~~~~~~~~~~ +# +# All OpInfos are currently defined in this file. Most OpInfo tests are defined +# in test_ops.py, but some system-specific tests are defined in those +# systems' test files, and subclass-specific tests are defined in the test +# file that corresponds to that subclass (see the below). +# Expect a reorganization in the future. +# +# WHAT'S TESTED? +# ~~~~~~~~~~~~~~ +# +# Every OpInfo in the op_db sequence has the following properties validated in +# test_ops.py: +# +# - that its supported dtypes are specified correctly +# - that it supports the out= argument properly (if it allows out=), +# see https://github.com/pytorch/pytorch/wiki/Developer-FAQ#how-does-out-work-in-pytorch +# - that it works with the conjugate view bit properly +# - that its function, method, and inplace variants perform the same operation +# (that is, that torch.add, torch.Tensor.add, and torch.Tensor.add_ all +# do the same thing). +# - that its inplace variant preserves the input's storage +# - that its gradient formula is implemented correctly, and that it supports +# gradgrad and complex grad and gradgrad and forward mode AD properly for +# the op's function and inplace variants (method variants are skipped +# to reduce test time). +# - that the operation performs the same operation when traced or scripted +# using the jit +# - that the operation is autodifferentiated by the jit as expected +# - that the operator's aliases, if any, perform the same operation and that +# the jit understands the alias +# +# Additional OpInfo tests are in test_jit_fuser_te.py, test_fx_experimental.py, +# and test_fx.py. These tests validate that operators work with NNC and FX +# as expected. +# +# For performance, some of the above tests may only run on the first +# SampleInput returned by an OpInfo's sample input function. +# +# In addition to these tests, some subclasses (discussed in the next section) +# define additional tests. +# +# Critically, as mentioned above, what's not tested is that the operator +# works as expected. When implementing an OpInfo an engineer must still +# typically write one or more tests validating the operator's behavior. +# +# OPINFO (SUB)CLASSES +# ~~~~~~~~~~~~~~~~~~~ +# +# In addition to the OpInfo base class there are several specialized OpInfo +# subclasses. For example, the UnaryUfuncInfo subclass is used for +# unary elementwise operations. These operations have a common structure +# that test_unary_ufuncs.py exploits with additional automated testing. +# The automated testing in test_unary_ufuncs.py is so thorough, comparing +# the operator to a NumPy reference function on a plethora of values, that +# just implementing an OpInfo for a unary elementwise operation is often +# sufficient testing. +# +# The ForeachFuncInfo is another OpInfo subclass that is hyper-specialized to a +# very unique class of operations. These OpInfos aren't included in the +# op_db sequence and have their own tests. +# +# Other OpInfo subclasses, like SpectralFuncInfo, are just for convenience +# when writing OpInfos. +# +# TESTING A NEW OPERATOR +# ~~~~~~~~~~~~~~~~~~~~~~ +# +# If you're adding a new operator to the torch, torch.fft, torch.linalg, +# or torch.special namespaces then you should add an OpInfo for it. As +# mentioned a couple times above, implementing an OpInfo is not usually +# sufficient testing (unless the operator is a unary elementwise operator). +# The OpInfo will only test the properties described in the "WHAT'S TESTED" +# section. It DOES NOT verify that the operator is implemented correctly. +# +# We are currently reviewing if operators in the torch.nn.functional namespace +# will be added as OpInfos, but you are encouraged to add an OpInfo for +# such operators, too. +# +# TIPS FOR WRITING AN OPINFO AND OPINFO TESTS +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# Writing an OpInfo can be a little daunting. Since the point of an OpInfo is to +# be consumed by a variety of systems it can be hard to understand how to +# deal with test failures or how to set the OpInfo metadata properly. +# +# Before adding an OpInfo it helps to look at other OpInfos. A sample inputs +# function must be defined, and the operator's dtypes must be specified. +# Once that's done you should run the operator's tests in test_ops.py +# (these can be filtered using the "-k" argument in pytest). Tests that +# fail should provide an error message that describes what to change about +# your OpInfo. You don't need to worry about changing an OpInfo's default +# values unless a test yells at you. +# +# Similarly, if you're writing a test that consumes OpInfos then it's critical +# your test provides a clear error message describing what to do when it +# fails. You should not assume the OpInfo implementer is familiar with your +# system. +# +# If you see a confusing error message while developing an OpInfo then please +# file an issue describing what happened. +# +# This trial-and-error approach can be frustrating to writing an OpInfo can +# be frustrating, but it's probably necessary as long as OpInfos don't require +# learning about all the systems that consume them. One thing that can help +# is the get_supported_dtypes() function defined in opinfo_helper.py. This +# function can be used to programmatically specify the dtypes an operator +# supports, and is especially useful if writing an OpInfo on a machine +# without a CUDA device. See its documentation for more details. +# +# THE FUTURE OF OPINFOS AND OPINFO TESTING +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# In the future we expect OpInfo coverage to improve, particularly for the +# torch, torch.fft, torch.linalg, and torch.special namespaces, and possibly +# for the torch.nn.functional namespace, too. In addition an analogous class, +# ModuleInfo, will be developed to improve module testing. +# +# We also expect at least two new OpInfo subclasses: BinaryUfuncInfo and +# ReductionInfo. Both will have new automated tests for correctness, too, +# which might make testing binary elementwise operations and reductions as +# simple as testing unary elementwise operations today. + # Classes and methods for the operator database class OpInfo(object): """Operator information and helper functions for acquiring it.""" @@ -4573,6 +4799,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): backward_dtypesIfROCM=floating_types_and(torch.half), supports_forward_ad=True, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), # addbmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), # https://github.com/pytorch/pytorch/issues/55907 @@ -4589,6 +4818,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): torch.complex64, torch.complex128), supports_forward_ad=True, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), # baddbmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), ), @@ -4613,6 +4845,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True, supports_forward_ad=True, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), # bmm does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), ), @@ -5708,6 +5943,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): assert_autodiffed=True, sample_inputs_func=sample_inputs_matmul, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), # matmul does not correctly warn when resizing out= inputs SkipInfo('TestCommon', 'test_out'), SkipInfo('TestCommon', 'test_conj_view', device_type='cpu'), @@ -6173,6 +6411,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): sample_inputs_func=sample_inputs_matmul, supports_out=False, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), SkipInfo('TestJit', 'test_variant_consistency_jit',), )), OpInfo('__rmod__', @@ -6527,6 +6768,9 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_out=False, sample_inputs_func=sample_inputs_einsum, skips=( + # FIXME: bfloat16 backward support likely depends on CUDA11+ + # and SM53+ + SkipInfo('TestCommon', 'test_dtypes', active_if=IS_WINDOWS), # test does not work with passing lambda for op # there's a test `test_einsum` in `test_jit.py` to handle this case SkipInfo('TestJit', 'test_variant_consistency_jit'), From 8a839c54788e6551ead9a018993c4995e02f3219 Mon Sep 17 00:00:00 2001 From: albanD Date: Fri, 18 Jun 2021 04:35:34 -0700 Subject: [PATCH 228/305] Fix saved variable unpacking version counter (#60195) Summary: We only set the value and not the actual VC. This means that in the context of double backward, if that saved tensor is saved again and the original Tensor is modified inplace, we would not detect it. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60195 Reviewed By: Varal7 Differential Revision: D29208766 Pulled By: albanD fbshipit-source-id: 81175f8e3f111f89524f8e46f47577b2ea4fc945 --- test/test_autograd.py | 15 +++++++++++++++ torch/csrc/autograd/saved_variable.cpp | 2 +- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/test/test_autograd.py b/test/test_autograd.py index dfba768f33a42..72239d0585fe8 100644 --- a/test/test_autograd.py +++ b/test/test_autograd.py @@ -5494,6 +5494,21 @@ def test_no_unnecessary_unwrapping(self): # a is left untouched self.assertEqual(a, a_orig) + def test_saved_variable_version_counter(self): + a = torch.rand(2, requires_grad=True) + + b = torch.exp(a) + + b_unpacked = b.grad_fn._saved_result + self.assertEqual(b, b_unpacked) + self.assertEqual(b._version, b_unpacked._version) + + with torch.no_grad(): + b += 1 + + self.assertEqual(b, b_unpacked) + self.assertEqual(b._version, b_unpacked._version) + def index_perm_variable(shape, max_indices): if not isinstance(shape, tuple): diff --git a/torch/csrc/autograd/saved_variable.cpp b/torch/csrc/autograd/saved_variable.cpp index 56bd9744d2ad2..b81ccf2e68324 100644 --- a/torch/csrc/autograd/saved_variable.cpp +++ b/torch/csrc/autograd/saved_variable.cpp @@ -138,7 +138,7 @@ Variable SavedVariable::unpack(std::shared_ptr saved_for) const { // they still share the same storage. This works only because we never call // in-place functions on unpacked variables. Variable var = make_variable(data_, Edge(std::move(grad_fn), output_nr_)); - impl::set_version_counter(var, saved_version_); + impl::set_version_counter(var, version_counter_); // NB: var here is never a view so there is no need to make anything special // for the case where the saved Tensor was a view. This whole argument relies From 83fde5d981ba8afd53faca1eaf86338936f0dddc Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 18 Jun 2021 05:12:27 -0700 Subject: [PATCH 229/305] [reland] Pass RequestCallback to FaultyPG RPC agent (#60168) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60168 Reland of #59939. ghstack-source-id: 131706860 Test Plan: CI Reviewed By: mrshenli Differential Revision: D29193235 fbshipit-source-id: 170108956a041f6a91b2b21c76ab1a0e0cdd34a2 --- .../distributed/rpc/testing/faulty_process_group_agent.cpp | 4 ++-- .../csrc/distributed/rpc/testing/faulty_process_group_agent.h | 1 + torch/csrc/distributed/rpc/testing/init.cpp | 2 ++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp index f51de9d870971..bb980ee8cef08 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp @@ -1,4 +1,3 @@ -#include #include #include @@ -16,6 +15,7 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( c10::intrusive_ptr<::c10d::ProcessGroup> pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, + std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends) @@ -25,7 +25,7 @@ FaultyProcessGroupAgent::FaultyProcessGroupAgent( std::move(pg), numSendRecvThreads, rpcTimeout, - std::make_unique()), + std::move(cb)), failNumSends_(failNumSends), messageTypesToFail_(parseMessagesToFailInput(messagesToFail)), messageTypesToDelay_(parseMessagesToDelay(messageTypesToDelay)) {} diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index b80bd78c3e1de..ee589072f2ddd 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -39,6 +39,7 @@ class FaultyProcessGroupAgent : public ProcessGroupAgent { c10::intrusive_ptr pg, int numSendRecvThreads, std::chrono::milliseconds rpcTimeout, + std::unique_ptr cb, const std::vector& messagesToFail, const std::unordered_map& messageTypesToDelay, int failNumSends = 0); diff --git a/torch/csrc/distributed/rpc/testing/init.cpp b/torch/csrc/distributed/rpc/testing/init.cpp index 28344bb5b1978..bccaa1f2b4232 100644 --- a/torch/csrc/distributed/rpc/testing/init.cpp +++ b/torch/csrc/distributed/rpc/testing/init.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -82,6 +83,7 @@ PyObject* faulty_agent_init(PyObject* _unused, PyObject* noargs) { process_group, num_send_recv_threads, rpc_timeout, + std::make_unique(), messages_to_fail, messages_to_delay, failNumSends), From 958b881d70a813242b61f55063b03fc58454d8ff Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 18 Jun 2021 05:12:27 -0700 Subject: [PATCH 230/305] [reland] Add some TORCH_API annotations to RPC (#60169) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60169 Reland of #59939. ghstack-source-id: 131706861 Test Plan: CI Reviewed By: mrshenli Differential Revision: D29193233 fbshipit-source-id: 91d3ef9003b9da7b99e1b9310b7f5a6c505d3b99 --- torch/csrc/distributed/rpc/process_group_agent.h | 6 +++--- torch/csrc/distributed/rpc/tensorpipe_agent.h | 12 ++++++------ torch/csrc/distributed/rpc/tensorpipe_utils.h | 6 +++--- .../rpc/testing/faulty_process_group_agent.h | 4 ++-- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/torch/csrc/distributed/rpc/process_group_agent.h b/torch/csrc/distributed/rpc/process_group_agent.h index 5706870988140..a6d1115f4074c 100644 --- a/torch/csrc/distributed/rpc/process_group_agent.h +++ b/torch/csrc/distributed/rpc/process_group_agent.h @@ -15,7 +15,7 @@ namespace rpc { constexpr auto kDefaultNumSendRecvThreads = 4; -struct ProcessGroupRpcBackendOptions : public RpcBackendOptions { +struct TORCH_API ProcessGroupRpcBackendOptions : public RpcBackendOptions { ProcessGroupRpcBackendOptions( int num_send_recv_threads, float rpc_timeout, @@ -34,7 +34,7 @@ struct ProcessGroupRpcBackendOptions : public RpcBackendOptions { // SendWork and RecvWork will be put into a task queue, and later picked up by // worker threads from the same ThreadPool. -struct SendWork { +struct TORCH_API SendWork { SendWork(const WorkerInfo& to, c10::intrusive_ptr message) : to_(to), message_(std::move(message)) {} @@ -44,7 +44,7 @@ struct SendWork { // SendWork wraps a Message and RecvWork wraps a Tensor. The difference here is // to allow us to run serialization/deserialization in the worker threads. -struct RecvWork { +struct TORCH_API RecvWork { RecvWork( const WorkerInfo& from, MessageType type, diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index df3328793fa11..9462c396b0f3b 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -62,7 +62,7 @@ constexpr int64_t kCudaBasicChannelPriority = 0; using steady_clock_time_point = std::chrono::time_point; -struct TransportRegistration { +struct TORCH_API TransportRegistration { std::shared_ptr transport; int64_t priority; std::string address; @@ -71,7 +71,7 @@ struct TransportRegistration { // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) C10_DECLARE_REGISTRY(TensorPipeTransportRegistry, TransportRegistration); -struct ChannelRegistration { +struct TORCH_API ChannelRegistration { std::shared_ptr channel; int64_t priority; }; @@ -81,7 +81,7 @@ C10_DECLARE_REGISTRY(TensorPipeChannelRegistry, ChannelRegistration); constexpr auto kDefaultNumWorkerThreads = 16; -struct TensorPipeRpcBackendOptions : public RpcBackendOptions { +struct TORCH_API TensorPipeRpcBackendOptions : public RpcBackendOptions { TensorPipeRpcBackendOptions( int numWorkerThreads, optional> transports, @@ -146,13 +146,13 @@ struct TensorPipeRpcBackendOptions : public RpcBackendOptions { }; // Struct to track the network source metrics -struct NetworkSourceInfo { +struct TORCH_API NetworkSourceInfo { worker_id_t srcRank; std::vector srcMachineAddr; }; // Struct to track aggregated network metrics -struct AggregatedNetworkData { +struct TORCH_API AggregatedNetworkData { uint64_t numCalls{0}; uint64_t totalSentBytes{0}; uint64_t totalRecvBytes{0}; @@ -163,7 +163,7 @@ struct AggregatedNetworkData { // to transparently move tensors and payloads through the fastest available // transport or channel. It acts like a hybrid RPC transport, providing shared // memory (linux) and TCP (linux & mac) support. CUDA support is in progress. -class TensorPipeAgent : public RpcAgent { +class TORCH_API TensorPipeAgent : public RpcAgent { public: TensorPipeAgent( const c10::intrusive_ptr<::c10d::Store>& store, diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index 3f41b351c9898..ab328b9dca1a1 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -15,7 +15,7 @@ namespace torch { namespace distributed { namespace rpc { -const c10::Stream& getStreamForDevice( +TORCH_API const c10::Stream& getStreamForDevice( const std::vector& streams, const c10::Device& device); @@ -44,12 +44,12 @@ class TensorpipeDeviceTypeConverter { virtual ~TensorpipeDeviceTypeConverter() = default; }; -extern C10_API std::array< +extern TORCH_API std::array< std::atomic, static_cast(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)> device_type_converter_registry; -class C10_API TensorpipeDeviceTypeConverterRegistrar { +class TORCH_API TensorpipeDeviceTypeConverterRegistrar { public: TensorpipeDeviceTypeConverterRegistrar( DeviceType, diff --git a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h index ee589072f2ddd..d0bbb33fe3df2 100644 --- a/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h +++ b/torch/csrc/distributed/rpc/testing/faulty_process_group_agent.h @@ -7,7 +7,7 @@ namespace torch { namespace distributed { namespace rpc { -struct FaultyProcessGroupRpcBackendOptions +struct TORCH_API FaultyProcessGroupRpcBackendOptions : public ProcessGroupRpcBackendOptions { FaultyProcessGroupRpcBackendOptions( int num_send_recv_threads, @@ -31,7 +31,7 @@ struct FaultyProcessGroupRpcBackendOptions int numFailSends; }; -class FaultyProcessGroupAgent : public ProcessGroupAgent { +class TORCH_API FaultyProcessGroupAgent : public ProcessGroupAgent { public: FaultyProcessGroupAgent( const c10::intrusive_ptr<::c10d::Store>& store, From 08ce5eedf526181bccc64a3b9a3bb869bc06c38d Mon Sep 17 00:00:00 2001 From: Luca Wehrstedt Date: Fri, 18 Jun 2021 05:12:27 -0700 Subject: [PATCH 231/305] [reland] Move RPC agents to libtorch (#60170) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60170 Reland of #59939. Test Plan: CI Reviewed By: mrshenli Differential Revision: D29193234 fbshipit-source-id: ee2a90d6be961c10f91361512bdd4cadca43dd60 --- BUILD.bazel | 2 +- caffe2/CMakeLists.txt | 57 +-------- cmake/Dependencies.cmake | 7 ++ test/cpp/rpc/CMakeLists.txt | 4 +- third_party/tensorpipe | 2 +- third_party/tensorpipe.BUILD | 112 ++++++++++-------- tools/build_variables.bzl | 12 +- torch/CMakeLists.txt | 6 - torch/csrc/distributed/rpc/macros.h | 5 - .../csrc/distributed/rpc/tensorpipe_agent.cpp | 1 - torch/csrc/distributed/rpc/tensorpipe_agent.h | 1 - .../csrc/distributed/rpc/tensorpipe_cuda.cpp | 3 +- .../csrc/distributed/rpc/tensorpipe_utils.cpp | 1 - torch/csrc/distributed/rpc/tensorpipe_utils.h | 1 - 14 files changed, 86 insertions(+), 128 deletions(-) delete mode 100644 torch/csrc/distributed/rpc/macros.h diff --git a/BUILD.bazel b/BUILD.bazel index b7e16ac1c915c..45d71a8d4626e 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1728,7 +1728,7 @@ cc_library( ], [ ":aten", - "@tensorpipe", + "@tensorpipe//:tensorpipe_cpu", ], ), alwayslink = True, diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 174018456efd8..88cffd1a75d1c 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -344,53 +344,6 @@ endif() if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) - if(USE_DISTRIBUTED) - - # Define this target even if we're building without TensorPipe, to make life - # easier to other targets that depend on this. However, in that case, by not - # setting the USE_TENSORPIPE compile definition, this target will just end - # up being empty. Downstream targets should also add a #ifdef guard. - if(NOT WIN32) - add_library(process_group_agent - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/process_group_agent.h" - ) - target_link_libraries(process_group_agent PRIVATE torch fmt::fmt-header-only) - add_dependencies(process_group_agent torch) - - if(USE_TENSORPIPE) - add_library(tensorpipe_agent - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/agent_utils.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/macros.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_agent.h" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_cuda.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.cpp" - "${TORCH_SRC_DIR}/csrc/distributed/rpc/tensorpipe_utils.h" - ) - target_link_libraries(tensorpipe_agent PRIVATE torch tensorpipe fmt::fmt-header-only) - add_dependencies(tensorpipe_agent torch) - if(USE_CUDA) - target_compile_definitions(tensorpipe_agent PUBLIC USE_CUDA) - endif() - - if(USE_ROCM) - target_compile_definitions(tensorpipe_agent PRIVATE - USE_ROCM - __HIP_PLATFORM_HCC__ - ) - endif() - - target_compile_definitions(tensorpipe_agent PUBLIC USE_TENSORPIPE) - target_link_libraries(tensorpipe_agent PRIVATE tensorpipe) - add_dependencies(tensorpipe_agent tensorpipe) - endif() - endif() - endif() - set(CMAKE_POSITION_INDEPENDENT_CODE TRUE) # Generate files @@ -1236,7 +1189,7 @@ endif() if(USE_DISTRIBUTED) # Needed to support the inclusion of c10d/Foo.hpp headers. target_include_directories(torch_cpu PUBLIC ${TORCH_SRC_DIR}/lib) - target_compile_definitions(torch_cpu PRIVATE USE_DISTRIBUTED) + target_compile_definitions(torch_cpu PUBLIC USE_DISTRIBUTED) if(USE_GLOO AND USE_C10D_GLOO) target_compile_definitions(torch_cpu PUBLIC USE_C10D_GLOO) endif() @@ -1263,16 +1216,12 @@ if(USE_DISTRIBUTED) # #if defined(USE_DISTRIBUTED) && !defined(_WIN32) # need to be removed when RPC is supported if(NOT WIN32) - target_compile_definitions(torch_cpu PRIVATE - USE_RPC - ) + target_compile_definitions(torch_cpu PUBLIC USE_RPC) endif() # Pass USE_TENSORPIPE to torch_cpu as some parts of rpc/utils.cpp # can only be compiled with USE_TENSORPIPE is set. if(USE_TENSORPIPE) - target_compile_definitions(torch_cpu PRIVATE - USE_TENSORPIPE - ) + target_compile_definitions(torch_cpu PUBLIC USE_TENSORPIPE) endif() endif() diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index ab4cd32c40bce..70b6d71face6b 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1377,6 +1377,13 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE) add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/tensorpipe) list(APPEND Caffe2_DEPENDENCY_LIBS tensorpipe) + if(USE_CUDA) + list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS tensorpipe_cuda) + elseif(USE_ROCM) + message(WARNING "TensorPipe doesn't yet support ROCm") + # Not yet... + # list(APPEND Caffe2_HIP_DEPENDENCY_LIBS tensorpipe_hip) + endif() endif() endif() diff --git a/test/cpp/rpc/CMakeLists.txt b/test/cpp/rpc/CMakeLists.txt index 0eff382d2b1b8..c9fb1b0e7f17a 100644 --- a/test/cpp/rpc/CMakeLists.txt +++ b/test/cpp/rpc/CMakeLists.txt @@ -5,7 +5,7 @@ set(TORCH_RPC_TEST_SOURCES ${TORCH_RPC_TEST_DIR}/test_wire_serialization.cpp ) set(TORCH_RPC_TEST_DEPENDENCY_LIBS - torch gtest process_group_agent + torch gtest ) if(USE_GLOO) @@ -20,7 +20,7 @@ if(USE_TENSORPIPE) ${TORCH_RPC_TEST_DIR}/test_tensorpipe_serialization.cpp ) list(APPEND TORCH_RPC_TEST_DEPENDENCY_LIBS - tensorpipe_agent tensorpipe + tensorpipe ) endif() diff --git a/third_party/tensorpipe b/third_party/tensorpipe index 42a67277c1882..c0e7623adb05f 160000 --- a/third_party/tensorpipe +++ b/third_party/tensorpipe @@ -1 +1 @@ -Subproject commit 42a67277c1882c90cec0da6e57afb20247424994 +Subproject commit c0e7623adb05f36311c7cde6dac8fc4c290419d9 diff --git a/third_party/tensorpipe.BUILD b/third_party/tensorpipe.BUILD index d9e4bdb395741..ae210f473933d 100644 --- a/third_party/tensorpipe.BUILD +++ b/third_party/tensorpipe.BUILD @@ -71,63 +71,82 @@ cc_library( ) header_template_rule( - name = "tensorpipe_config_header", + name = "tensorpipe_cpu_config_header", src = "tensorpipe/config.h.in", out = "tensorpipe/config.h", substitutions = { - "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "", - "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "", - "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "", - "#cmakedefine01 TENSORPIPE_SUPPORTS_CUDA": "", + "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "#define TENSORPIPE_HAS_SHM_TRANSPORT 1", + "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "#define TENSORPIPE_HAS_IBV_TRANSPORT 1", + "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "#define TENSORPIPE_HAS_CMA_CHANNEL 1", }, ) -TENSORPIPE_HEADERS = glob([ - "tensorpipe/*.h", - "tensorpipe/channel/*.h", - "tensorpipe/channel/*/*.h", - "tensorpipe/common/*.h", - "tensorpipe/core/*.h", - "tensorpipe/transport/*.h", - "tensorpipe/transport/*/*.h", - "tensorpipe/util/*/*.h", -]) +header_template_rule( + name = "tensorpipe_cuda_config_header", + src = "tensorpipe/config_cuda.h.in", + out = "tensorpipe/config_cuda.h", + substitutions = { + "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "#define TENSORPIPE_HAS_CUDA_IPC_CHANNEL 1", + "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "#define TENSORPIPE_HAS_CUDA_GDR_CHANNEL 1", + }, +) -TENSORPIPE_BASE_SRCS = glob([ - "tensorpipe/*.cc", - "tensorpipe/channel/*.cc", - "tensorpipe/common/address.cc", - "tensorpipe/common/epoll_loop.cc", - "tensorpipe/common/error.cc", - "tensorpipe/common/fd.cc", - "tensorpipe/common/ibv.cc", - "tensorpipe/common/socket.cc", - "tensorpipe/common/system.cc", - "tensorpipe/core/*.cc", - "tensorpipe/transport/*.cc", - "tensorpipe/util/*/*.cc", -]) +# We explicitly list the CUDA headers & sources, and we consider everything else +# as CPU (using a catch-all glob). This is both because there's fewer CUDA files +# (thus making it easier to list them exhaustively) and because it will make it +# more likely to catch a misclassified file: if we forget to mark a file as CUDA +# we'll try to build it on CPU and that's likely to fail. -TENSORPIPE_SRCS = TENSORPIPE_BASE_SRCS + glob([ - "tensorpipe/channel/basic/*.cc", - "tensorpipe/channel/mpt/*.cc", - "tensorpipe/channel/xth/*.cc", - "tensorpipe/transport/uv/*.cc", -]) +TENSORPIPE_CUDA_HEADERS = [ + "tensorpipe/tensorpipe_cuda.h", + "tensorpipe/channel/cuda_basic/*.h", + "tensorpipe/channel/cuda_gdr/*.h", + "tensorpipe/channel/cuda_ipc/*.h", + "tensorpipe/channel/cuda_xth/*.h", + "tensorpipe/common/cuda.h", + "tensorpipe/common/cuda_buffer.h", + "tensorpipe/common/cuda_lib.h", + "tensorpipe/common/cuda_loop.h", + "tensorpipe/common/nvml_lib.h", +] -TENSORPIPE_SRCS_CUDA = TENSORPIPE_SRCS + glob([ - "tensorpipe/common/cuda_loop.cc", +TENSORPIPE_CUDA_SOURCES = [ "tensorpipe/channel/cuda_basic/*.cc", + "tensorpipe/channel/cuda_gdr/*.cc", "tensorpipe/channel/cuda_ipc/*.cc", "tensorpipe/channel/cuda_xth/*.cc", -]) + "tensorpipe/common/cuda_buffer.cc", + "tensorpipe/common/cuda_loop.cc", +] + +TENSORPIPE_CPU_HEADERS = glob( + [ + "tensorpipe/*.h", + "tensorpipe/channel/*.h", + "tensorpipe/channel/*/*.h", + "tensorpipe/common/*.h", + "tensorpipe/core/*.h", + "tensorpipe/transport/*.h", + "tensorpipe/transport/*/*.h", + ], + exclude=TENSORPIPE_CUDA_HEADERS) + +TENSORPIPE_CPU_SOURCES = glob( + [ + "tensorpipe/*.cc", + "tensorpipe/channel/*.cc", + "tensorpipe/channel/*/*.cc", + "tensorpipe/common/*.cc", + "tensorpipe/core/*.cc", + "tensorpipe/transport/*.cc", + "tensorpipe/transport/*/*.cc", + ], + exclude=TENSORPIPE_CUDA_SOURCES) cc_library( - name = "tensorpipe", - srcs = TENSORPIPE_SRCS + [":tensorpipe_config_header"], - hdrs = TENSORPIPE_HEADERS, + name = "tensorpipe_cpu", + srcs = TENSORPIPE_CPU_SOURCES, + hdrs = TENSORPIPE_CPU_HEADERS + [":tensorpipe_cpu_config_header"], includes = [ ".", ], @@ -143,8 +162,8 @@ cc_library( cc_library( name = "tensorpipe_cuda", - srcs = TENSORPIPE_SRCS_CUDA + [":tensorpipe_config_header"], - hdrs = TENSORPIPE_HEADERS, + srcs = TENSORPIPE_CUDA_SOURCES, + hdrs = TENSORPIPE_CUDA_HEADERS + [":tensorpipe_cuda_config_header"], includes = [ ".", ], @@ -153,8 +172,7 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ - ":libnop", - ":libuv", + ":tensorpipe_cpu", "@cuda", ], ) diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index c0b206fa48e05..06acafd645eab 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -356,12 +356,14 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.cpp", "torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.cpp", + "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/message.cpp", "torch/csrc/distributed/rpc/profiler/remote_profiler_manager.cpp", "torch/csrc/distributed/rpc/profiler/server_process_global_profiler.cpp", "torch/csrc/distributed/rpc/python_call.cpp", "torch/csrc/distributed/rpc/python_remote_call.cpp", "torch/csrc/distributed/rpc/python_resp.cpp", + "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/request_callback.cpp", "torch/csrc/distributed/rpc/request_callback_no_python.cpp", "torch/csrc/distributed/rpc/rpc_agent.cpp", @@ -371,6 +373,9 @@ libtorch_distributed_extra_sources = [ "torch/csrc/distributed/rpc/script_call.cpp", "torch/csrc/distributed/rpc/script_remote_call.cpp", "torch/csrc/distributed/rpc/script_resp.cpp", + "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", + "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", + "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/torchscript_functions.cpp", "torch/csrc/distributed/rpc/types.cpp", "torch/csrc/distributed/rpc/utils.cpp", @@ -526,6 +531,7 @@ libtorch_cuda_distributed_base_sources = [ # These files are only supported on Linux (and others) but not on Windows. libtorch_cuda_distributed_extra_sources = [ + "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", "torch/lib/c10d/NCCLUtils.cpp", "torch/lib/c10d/ProcessGroupNCCL.cpp", ] @@ -714,17 +720,11 @@ libtorch_python_distributed_core_sources = [ libtorch_python_distributed_sources = libtorch_python_distributed_core_sources + [ "torch/csrc/distributed/autograd/init.cpp", - "torch/csrc/distributed/rpc/agent_utils.cpp", "torch/csrc/distributed/rpc/init.cpp", - "torch/csrc/distributed/rpc/process_group_agent.cpp", "torch/csrc/distributed/rpc/py_rref.cpp", "torch/csrc/distributed/rpc/python_functions.cpp", "torch/csrc/distributed/rpc/python_rpc_handler.cpp", "torch/csrc/distributed/rpc/request_callback_impl.cpp", - "torch/csrc/distributed/rpc/tensorpipe_agent.cpp", - "torch/csrc/distributed/rpc/tensorpipe_cuda.cpp", - "torch/csrc/distributed/rpc/tensorpipe_utils.cpp", - "torch/csrc/distributed/rpc/testing/faulty_process_group_agent.cpp", "torch/csrc/distributed/rpc/testing/init.cpp", "torch/csrc/distributed/rpc/unpickled_python_call.cpp", "torch/csrc/distributed/rpc/unpickled_python_remote_call.cpp", diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index 197926f309838..ce0f16bf5abeb 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -261,11 +261,9 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") endif() if(USE_DISTRIBUTED) - list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_DISTRIBUTED) if(WIN32) append_filelist("libtorch_python_distributed_core_sources" TORCH_PYTHON_SRCS) else() - list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_RPC) append_filelist("libtorch_python_distributed_sources" TORCH_PYTHON_SRCS) endif() # Disable certain warnings for GCC-9.X @@ -274,10 +272,6 @@ if(USE_DISTRIBUTED) set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/rpc/testing/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") set_source_files_properties(${TORCH_SRC_DIR}/csrc/distributed/c10d/init.cpp PROPERTIES COMPILE_FLAGS "-Wno-cast-function-type") endif() - if(USE_TENSORPIPE) - list(APPEND TORCH_PYTHON_LINK_LIBRARIES tensorpipe) - list(APPEND TORCH_PYTHON_PUBLIC_COMPILE_DEFINITIONS USE_TENSORPIPE) - endif() # NCCL is a private dependency of libtorch, but libtorch_python includes # some private headers of libtorch, which in turn include NCCL. As a hacky # alternative to making NCCL a public dependency of libtorch, we make it diff --git a/torch/csrc/distributed/rpc/macros.h b/torch/csrc/distributed/rpc/macros.h deleted file mode 100644 index 2763dd0207bef..0000000000000 --- a/torch/csrc/distributed/rpc/macros.h +++ /dev/null @@ -1,5 +0,0 @@ -#pragma once - -#if defined(USE_CUDA) && !defined(__HIP_PLATFORM_HCC__) -#define USE_CUDA_NOT_ROCM -#endif diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp index 0f6645cdcd5d5..74c279425658b 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.cpp @@ -10,7 +10,6 @@ #include #include -#include #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_agent.h b/torch/csrc/distributed/rpc/tensorpipe_agent.h index 9462c396b0f3b..4450792a0f06d 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_agent.h +++ b/torch/csrc/distributed/rpc/tensorpipe_agent.h @@ -9,7 +9,6 @@ #include #include #include -#include #include // Forward-declare the TensorPipe classes we need, to avoid including its diff --git a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp index 9489fcd222bbd..03ec63d8ddc88 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_cuda.cpp @@ -1,8 +1,7 @@ -#include #include #include -#if defined(USE_TENSORPIPE) && defined(USE_CUDA_NOT_ROCM) +#if defined(USE_TENSORPIPE) && !defined(__HIP_PLATFORM_HCC__) #include #include diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp index 55b8554f66d28..32f3a132f8f50 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.cpp +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.cpp @@ -1,4 +1,3 @@ -#include #include #ifdef USE_TENSORPIPE diff --git a/torch/csrc/distributed/rpc/tensorpipe_utils.h b/torch/csrc/distributed/rpc/tensorpipe_utils.h index ab328b9dca1a1..bf5d87cacc4b5 100644 --- a/torch/csrc/distributed/rpc/tensorpipe_utils.h +++ b/torch/csrc/distributed/rpc/tensorpipe_utils.h @@ -2,7 +2,6 @@ #ifdef USE_TENSORPIPE -#include #include namespace tensorpipe { From 7c29ca7f2b49a91dbfc6ee35bf6da55616bd2960 Mon Sep 17 00:00:00 2001 From: TJ-coding <44460325+TJ-coding@users.noreply.github.com> Date: Fri, 18 Jun 2021 07:06:17 -0700 Subject: [PATCH 232/305] Fix Subset of a Subset not sliceable issue (#59513) Summary: Dataset can be indexed by a list, but a list can not be indexed by a list. This gives error when slicing a Subset initialised with a Subset, instead of a dataset. Fixed the issue by changing the indices to a Tensor which can be indexed by a list. Fixes https://github.com/pytorch/pytorch/issues/59512 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59513 Reviewed By: zou3519 Differential Revision: D29196891 Pulled By: ejguan fbshipit-source-id: ccde6e474fbcbddd2e9c7c107bc8b5de1307cdb9 --- test/test_dataloader.py | 31 ++++++++++++++++++++++++++++++- torch/utils/data/dataset.py | 2 ++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/test/test_dataloader.py b/test/test_dataloader.py index 7b3ae5b6de409..3c80106ab008b 100644 --- a/test/test_dataloader.py +++ b/test/test_dataloader.py @@ -13,7 +13,7 @@ import warnings import tempfile from torch import multiprocessing as mp -from torch.utils.data import _utils, Dataset, IterableDataset, TensorDataset, DataLoader, ConcatDataset, ChainDataset +from torch.utils.data import _utils, Dataset, IterableDataset, TensorDataset, DataLoader, ConcatDataset, ChainDataset, Subset from torch.utils.data._utils import MP_STATUS_CHECK_INTERVAL from torch.utils.data.dataset import random_split from torch._utils import ExceptionWrapper @@ -151,6 +151,35 @@ def test_splits_generator(self): b = torch.rand(10) self.assertEqual(a, b) + def test_slicing_of_subset_of_dataset(self): + # Testing slicing a subset initialized with a dataset + dataset = TensorDataset(torch.tensor([1, 2, 3, 4, 5])) + subset_of_dataset = Subset(dataset, [0, 1, 2, 3, 4]) + self.assertEqual(subset_of_dataset[:], dataset[:]) + self.assertEqual(subset_of_dataset[1:2], dataset[1:2]) + self.assertEqual(subset_of_dataset[0:-1:2], dataset[0:-1:2]) + # Testing slicing of subset from random split + subset1, subset2 = random_split(dataset, [3, 2]) + self.assertEqual(subset1[:], dataset[subset1.indices[:]]) + self.assertEqual(subset1[0:2], dataset[subset1.indices[0:2]]) + self.assertEqual(subset1[0:-1:2], dataset[subset1.indices[0:-1:2]]) + + def test_slicing_of_subset_of_subset(self): + # Testing slicing a subset initialized with a subset + dataset = TensorDataset(torch.tensor([1, 2, 3, 4, 5])) + subset_of_dataset = Subset(dataset, [0, 1, 2, 3, 4]) + subset_of_subset = Subset(subset_of_dataset, [0, 1, 2, 3, 4]) + self.assertEqual(subset_of_subset[:], dataset[:]) + self.assertEqual(subset_of_subset[0:2], dataset[0:2]) + self.assertEqual(subset_of_subset[0:-1:2], dataset[0:-1:2]) + # Testing slicing of subset of subset from random split + subset1, subset2 = random_split(dataset, [4, 1]) + subset_of_subset1, subset_of_subset2 = random_split(subset1, [3, 1]) + idx = [subset1.indices[i] for i in subset_of_subset1.indices] + self.assertEqual(subset_of_subset1[:], dataset[idx[:]]) + self.assertEqual(subset_of_subset1[0:2], dataset[idx[0:2]]) + self.assertEqual(subset_of_subset1[0:-1:2], dataset[idx[0:-1:2]]) + class CUDACountingDataset(Dataset): def __init__(self, n): diff --git a/torch/utils/data/dataset.py b/torch/utils/data/dataset.py index 73f71f809218f..a1549a0e97b0e 100644 --- a/torch/utils/data/dataset.py +++ b/torch/utils/data/dataset.py @@ -316,6 +316,8 @@ def __init__(self, dataset: Dataset[T_co], indices: Sequence[int]) -> None: self.indices = indices def __getitem__(self, idx): + if isinstance(idx, list): + return self.dataset[[self.indices[i] for i in idx]] return self.dataset[self.indices[idx]] def __len__(self): From d5988c5eca0221e9ef58918e4f0b504940cb926a Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Fri, 18 Jun 2021 07:22:22 -0700 Subject: [PATCH 233/305] remove unused `type: ignore` directives (#60006) Summary: During development it is common practice to put `type: ignore` comments on lines that are correct, but `mypy` doesn't recognize this. This often stems from the fact, that the used `mypy` version wasn't able to handle the used pattern. With every new release `mypy` gets better at handling complex code. In addition to fix all the previously accepted but now failing patterns, we should also revisit all `type: ignore` comments to see if they are still needed or not. Fortunately, we don't need to do it manually: by adding `warn_unused_ignores = True` to the configuration, `mypy` will error out in case it encounters an `type: ignore` that is no longer needed. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60006 Reviewed By: jbschlosser, malfet Differential Revision: D29133237 Pulled By: albanD fbshipit-source-id: 41e82edc5cd5affa7ccedad044b59b94dad4425a --- .github/scripts/lint_native_functions.py | 2 +- caffe2/contrib/aten/gen_op.py | 2 +- .../distributed/file_store_handler_op_test.py | 2 +- .../redis_store_handler_op_test.py | 2 +- mypy-strict.ini | 1 + mypy.ini | 14 ++++++++ test/test_futures.py | 2 +- test/test_utils.py | 2 +- tools/pyi/gen_pyi.py | 12 +++++-- tools/render_junit.py | 2 +- torch/distributed/_sharded_tensor/api.py | 2 +- torch/distributed/distributed_c10d.py | 6 ++-- .../rendezvous/c10d_rendezvous_backend.py | 2 +- .../utils/data/elastic_distributed_sampler.py | 4 +-- torch/distributed/launcher/api.py | 4 +-- torch/distributed/nn/api/remote_module.py | 14 ++++---- torch/distributed/rendezvous.py | 3 +- torch/distributions/utils.py | 2 +- torch/fx/experimental/normalize.py | 4 +-- torch/fx/passes/net_min_base.py | 2 +- torch/fx/passes/split_module.py | 34 +++++++++---------- torch/jit/_monkeytype_config.py | 4 +-- torch/jit/mobile/__init__.py | 20 +++++------ torch/nn/quantized/modules/conv.py | 2 +- torch/nn/utils/parametrizations.py | 2 +- torch/package/package_importer.py | 2 +- torch/quantization/fx/convert.py | 2 +- torch/quantization/fx/prepare.py | 6 ++-- torch/testing/_internal/common_distributed.py | 2 +- .../_internal/common_methods_invocations.py | 10 +++--- torch/testing/_internal/common_utils.py | 2 +- torch/testing/_internal/dist_utils.py | 2 +- torch/utils/data/dataloader.py | 6 ++-- torch/utils/data/dataset.py | 3 +- torch/utils/data/distributed.py | 12 +++---- torch/utils/model_dump/__init__.py | 6 ++-- torch/utils/tensorboard/writer.py | 2 +- 37 files changed, 108 insertions(+), 93 deletions(-) diff --git a/.github/scripts/lint_native_functions.py b/.github/scripts/lint_native_functions.py index 2e6d4e3e76757..86a31615de638 100755 --- a/.github/scripts/lint_native_functions.py +++ b/.github/scripts/lint_native_functions.py @@ -26,7 +26,7 @@ def fn(base: str) -> str: with open(Path(__file__).parent.parent.parent / fn('.'), "r") as f: contents = f.read() -yaml = ruamel.yaml.YAML() # type: ignore[attr-defined] +yaml = ruamel.yaml.YAML() yaml.preserve_quotes = True yaml.width = 1000 yaml.boolean_representation = ['False', 'True'] diff --git a/caffe2/contrib/aten/gen_op.py b/caffe2/contrib/aten/gen_op.py index 94a1f1fedc3b2..003d570b6eee2 100755 --- a/caffe2/contrib/aten/gen_op.py +++ b/caffe2/contrib/aten/gen_op.py @@ -39,7 +39,7 @@ sys.path.insert(0, os.path.join(args.aten_root, '..')) from tools.codegen.code_template import CodeTemplate as CT else: - from tools.codegen.code_template import CodeTemplate as CT # type: ignore[import,no-redef] + from tools.codegen.code_template import CodeTemplate as CT OP_TEMPLATE = CT.from_file( os.path.join(args.template_dir, 'aten_op_template.h')) diff --git a/caffe2/distributed/file_store_handler_op_test.py b/caffe2/distributed/file_store_handler_op_test.py index 427b68420d398..72f8e456292de 100644 --- a/caffe2/distributed/file_store_handler_op_test.py +++ b/caffe2/distributed/file_store_handler_op_test.py @@ -8,7 +8,7 @@ import tempfile import shutil -from caffe2.distributed.python import StoreHandlerTimeoutError # type: ignore[import] +from caffe2.distributed.python import StoreHandlerTimeoutError from caffe2.distributed.store_ops_test_util import StoreOpsTests from caffe2.python import core, workspace, dyndep from caffe2.python.test_util import TestCase diff --git a/caffe2/distributed/redis_store_handler_op_test.py b/caffe2/distributed/redis_store_handler_op_test.py index 8f5d58e851853..2eb6c9adb7050 100644 --- a/caffe2/distributed/redis_store_handler_op_test.py +++ b/caffe2/distributed/redis_store_handler_op_test.py @@ -6,7 +6,7 @@ import os import uuid -from caffe2.distributed.python import StoreHandlerTimeoutError # type: ignore[import] +from caffe2.distributed.python import StoreHandlerTimeoutError from caffe2.distributed.store_ops_test_util import StoreOpsTests from caffe2.python import core, workspace, dyndep from caffe2.python.test_util import TestCase diff --git a/mypy-strict.ini b/mypy-strict.ini index cb8ef8f59c30e..5d165092287c7 100644 --- a/mypy-strict.ini +++ b/mypy-strict.ini @@ -19,6 +19,7 @@ disallow_any_unimported = True # Across versions of mypy, the flags toggled by --strict vary. To ensure # we have reproducible type check, we instead manually specify the flags warn_unused_configs = True +warn_unused_ignores = True disallow_any_generics = True disallow_subclassing_any = True disallow_untyped_calls = True diff --git a/mypy.ini b/mypy.ini index 1002b7da06856..efacfae88be25 100644 --- a/mypy.ini +++ b/mypy.ini @@ -6,6 +6,7 @@ plugins = mypy_plugins/check_mypy_version.py cache_dir = .mypy_cache/normal warn_unused_configs = True +warn_unused_ignores = True warn_redundant_casts = True show_error_codes = True show_column_numbers = True @@ -95,6 +96,19 @@ ignore_errors = True [mypy-torch.overrides] ignore_errors = True +# +# Files with 'type: ignore' comments that are needed if checked with mypy-strict.ini +# + +[mypy-tools.render_junit] +warn_unused_ignores = False + +[mypy-tools.generate_torch_version] +warn_unused_ignores = False + +[mypy-tools.stats_utils.s3_stat_parser] +warn_unused_ignores = False + # # Adding type annotations to caffe2 is probably not worth the effort # only work on this if you have a specific reason for it, otherwise diff --git a/test/test_futures.py b/test/test_futures.py index e39bae5f9f57e..ee17f38460278 100644 --- a/test/test_futures.py +++ b/test/test_futures.py @@ -30,7 +30,7 @@ def test_set_exception(self) -> None: f = Future() f.set_exception(value_error) with self.assertRaisesRegex(ValueError, "Intentional"): - f.value() # type: ignore[attr-defined] + f.value() def cb(fut): fut.value() diff --git a/test/test_utils.py b/test/test_utils.py index 8d4814f60c5b3..89318cdf88112 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -743,7 +743,7 @@ def forward(self, x): # data can be passed without errors x = torch.randn(4, 4).fill_(1.0) ms(x) - with self.assertRaisesRegex(torch.jit.Error, "foo"): # type: ignore[type-var] + with self.assertRaisesRegex(torch.jit.Error, "foo"): ms(torch.tensor([False], dtype=torch.bool)) diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py index 37ba0e33afdfd..8be769b7a5e1c 100644 --- a/tools/pyi/gen_pyi.py +++ b/tools/pyi/gen_pyi.py @@ -126,7 +126,10 @@ def should_bind_method(python_func: PythonSignatureNativeFunctionPair) -> bool: 'iadd', 'iand', 'idiv', 'ilshift', 'imul', 'ior', 'irshift', 'isub', 'ixor', 'ifloordiv', 'imod', # inplace ops ) -comparison_ops = ('eq', 'ne', 'ge', 'gt', 'lt', 'le') +symmetric_comparison_ops = ('eq', 'ne') +asymmetric_comparison_ops = ('ge', 'gt', 'lt', 'le') +comparison_ops = symmetric_comparison_ops + asymmetric_comparison_ops + unary_ops = ('neg', 'abs', 'invert') to_py_type_ops = ('bool', 'float', 'complex', 'long', 'index', 'int', 'nonzero') all_ops = binary_ops + comparison_ops + unary_ops + to_py_type_ops @@ -145,8 +148,11 @@ def sig_for_ops(opname: str) -> List[str]: if name in binary_ops: return ['def {}(self, other: Any) -> Tensor: ...'.format(opname)] elif name in comparison_ops: - # unsafe override https://github.com/python/mypy/issues/5704 - return ['def {}(self, other: Any) -> Tensor: ... # type: ignore[override]'.format(opname)] + sig = 'def {}(self, other: Any) -> Tensor: ...'.format(opname) + if name in symmetric_comparison_ops: + # unsafe override https://github.com/python/mypy/issues/5704 + sig += ' # type: ignore[override]' + return [sig] elif name in unary_ops: return ['def {}(self) -> Tensor: ...'.format(opname)] elif name in to_py_type_ops: diff --git a/tools/render_junit.py b/tools/render_junit.py index eac873d321cab..28e617af0e8f4 100644 --- a/tools/render_junit.py +++ b/tools/render_junit.py @@ -12,7 +12,7 @@ ) try: - import rich # type: ignore[import] + import rich except ImportError: print("rich not found, for color output use 'pip install rich'") diff --git a/torch/distributed/_sharded_tensor/api.py b/torch/distributed/_sharded_tensor/api.py index 6716e6cc9bfd7..f8cc49ea0d401 100644 --- a/torch/distributed/_sharded_tensor/api.py +++ b/torch/distributed/_sharded_tensor/api.py @@ -309,7 +309,7 @@ def _init_enumerable( def _parse_and_validate_remote_device(self, device): - on, local_device = _parse_remote_device(device) # type: ignore[arg-type] + on, local_device = _parse_remote_device(device) # Validate rank. if isinstance(on, int) and (on < 0 or on >= dist.get_world_size(self._process_group)): diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index c77826690d362..f48fc54a2c091 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -1591,7 +1591,7 @@ def all_gather_object(object_list, obj, group=None): all_gather(output_tensors, input_tensor, group=group) # Deserialize outputs back to object. for i, tensor in enumerate(output_tensors): - tensor = tensor.type(torch.uint8) # type:ignore[call-overload] + tensor = tensor.type(torch.uint8) if tensor.device != torch.device("cpu"): tensor = tensor.cpu() tensor_size = object_size_list[i] @@ -1695,7 +1695,7 @@ def gather_object(obj, object_gather_list=None, dst=0, group=None): if my_rank != dst: return for i, tensor in enumerate(output_tensors): - tensor = tensor.type(torch.uint8) # type: ignore[call-overload] + tensor = tensor.type(torch.uint8) tensor_size = object_size_list[i] object_gather_list[i] = _tensor_to_object(tensor, tensor_size) @@ -1790,7 +1790,7 @@ def broadcast_object_list(object_list, src=0, group=None): if my_rank != src: for i, obj_size in enumerate(object_sizes_tensor): obj_view = object_tensor[offset : offset + obj_size] - obj_view = obj_view.type(torch.uint8) # type: ignore[call-overload] + obj_view = obj_view.type(torch.uint8) if obj_view.device != torch.device("cpu"): obj_view = obj_view.cpu() offset += obj_size diff --git a/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py b/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py index 4db9b669c314b..df61186421e78 100644 --- a/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py +++ b/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py @@ -143,7 +143,7 @@ def _create_tcp_store(params: RendezvousParameters) -> TCPStore: # see the explanation in the except clause below. for is_server in [is_host, False]: try: - store = TCPStore( # type: ignore[call-arg] + store = TCPStore( host, port, is_master=is_server, timeout=timedelta(seconds=read_timeout) ) diff --git a/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py b/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py index dcf20dc7ad3ad..c4ff898823e94 100644 --- a/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py +++ b/torch/distributed/elastic/utils/data/elastic_distributed_sampler.py @@ -44,7 +44,7 @@ def __init__(self, dataset, num_replicas=None, rank=None, start_index=0): self.start_index = start_index self.num_samples = int( - math.ceil(float(len(self.dataset) - self.start_index) / self.num_replicas) # type: ignore[arg-type] + math.ceil(float(len(self.dataset) - self.start_index) / self.num_replicas) ) self.total_size = self.num_samples * self.num_replicas @@ -53,7 +53,7 @@ def __iter__(self): g = torch.Generator() g.manual_seed(self.epoch) indices = ( - torch.randperm(len(self.dataset) - self.start_index, generator=g) # type: ignore[arg-type] + torch.randperm(len(self.dataset) - self.start_index, generator=g) .add(self.start_index) .tolist() ) diff --git a/torch/distributed/launcher/api.py b/torch/distributed/launcher/api.py index 00ae8dad03f44..336474b1b99c5 100644 --- a/torch/distributed/launcher/api.py +++ b/torch/distributed/launcher/api.py @@ -12,8 +12,8 @@ import torch.distributed.elastic.rendezvous.registry as rdzv_registry from torch.distributed.elastic import events, metrics -from torch.distributed.elastic.agent.server.api import WorkerSpec, WorkerState # type: ignore[import] -from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent # type: ignore[import] +from torch.distributed.elastic.agent.server.api import WorkerSpec, WorkerState +from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent from torch.distributed.elastic.multiprocessing import Std from torch.distributed.elastic.multiprocessing.errors import ChildFailedError, record from torch.distributed.elastic.rendezvous import RendezvousParameters diff --git a/torch/distributed/nn/api/remote_module.py b/torch/distributed/nn/api/remote_module.py index 2d27d25442408..177a846fbc1fb 100644 --- a/torch/distributed/nn/api/remote_module.py +++ b/torch/distributed/nn/api/remote_module.py @@ -395,10 +395,10 @@ def named_modules( ): _raise_not_supported(self.named_modules.__name__) - def train(self: T, mode: bool = True) -> T: # type: ignore[return] + def train(self: T, mode: bool = True) -> T: return self.module_rref.rpc_sync().train() # type: ignore[operator, union-attr] - def eval(self: T) -> T: # type: ignore[return] + def eval(self: T) -> T: return self.module_rref.rpc_sync().eval() # type: ignore[operator, union-attr] def requires_grad_(self: T, requires_grad: bool = True) -> T: # type: ignore[return] @@ -413,7 +413,7 @@ def share_memory(self: T) -> T: # type: ignore[return] def extra_repr(self) -> str: # type: ignore[return] _raise_not_supported(self.extra_repr.__name__) - def _prepare_init(self, remote_device: str) -> bool: # type: ignore[return] + def _prepare_init(self, remote_device: str) -> bool: """ Prepares the initializaiton and returns whether to enable automatically moving CPU tensors to CUDA devices. """ @@ -639,7 +639,7 @@ def __init__( args: Tuple = None, kwargs: Dict[str, Any] = None, ): - super().__init__(remote_device, module_cls, args, kwargs) # type: ignore[arg-type] + super().__init__(remote_device, module_cls, args, kwargs) def _remote_module_receiver( @@ -651,7 +651,7 @@ def _remote_module_receiver( serialized_remote_module = _SerializedRemoteModule._make( remote_module_pickled_attrs ) - m = object.__new__(RemoteModule) # type: ignore[attr-defined] + m = object.__new__(RemoteModule) m.__dict__.update(serialized_remote_module._asdict()) # Unpickling the attribute `module_rref` must invoke RRef's `_deserialize()` method. @@ -675,10 +675,10 @@ def _remote_module_reducer(remote_module): # Pickling the attribute `module_rref` must invoke RRef's `_serialize()` method. if k == "module_rref": pickled_attrs[k] = v._serialize() - elif k in _REMOTE_MODULE_PICKLED_ATTRIBUTES: # type: ignore[attr-defined] + elif k in _REMOTE_MODULE_PICKLED_ATTRIBUTES: pickled_attrs[k] = v # Check if unpickled attributes are all in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING. - elif k not in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING: # type: ignore[attr-defined] + elif k not in _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING: print( "The new attribute ``{}`` of RemoteModule is ignored during RPC pickling. " "To pickle this attribute, please add it to ``_REMOTE_MODULE_PICKLED_ATTRIBUTES``. " diff --git a/torch/distributed/rendezvous.py b/torch/distributed/rendezvous.py index b9380a65dec0f..99fa0e9836b16 100644 --- a/torch/distributed/rendezvous.py +++ b/torch/distributed/rendezvous.py @@ -60,8 +60,7 @@ def rendezvous(url: str, rank: int = -1, world_size: int = -1, **kwargs): result = urlparse(url) if rank != -1 or world_size != -1: query_dict: Dict[str, Union[int, str]] = dict( - # mypy doesn't allow dict() to accept List of values (#257) - pair.split("=") for pair in filter(None, result.query.split("&")) # type: ignore[arg-type, misc] + pair.split("=") for pair in filter(None, result.query.split("&")) ) assert ( "rank" not in query_dict and "world_size" not in query_dict diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py index 84f45f1d33cf4..f5843adec1a61 100644 --- a/torch/distributions/utils.py +++ b/torch/distributions/utils.py @@ -101,7 +101,7 @@ class lazy_property(object): """ def __init__(self, wrapped): self.wrapped = wrapped - update_wrapper(self, wrapped) # type: ignore[arg-type] + update_wrapper(self, wrapped) def __get__(self, instance, obj_type=None): if instance is None: diff --git a/torch/fx/experimental/normalize.py b/torch/fx/experimental/normalize.py index 4b757fd0d9993..dc4454b2db8e0 100644 --- a/torch/fx/experimental/normalize.py +++ b/torch/fx/experimental/normalize.py @@ -74,7 +74,7 @@ def call_function( args, # type: ignore[arg-type] kwargs, arg_types, # type: ignore[arg-type] - kwarg_types, # type: ignore[arg-type] + kwarg_types, self.normalize_to_only_use_kwargs, ) if new_args_and_kwargs: @@ -93,7 +93,7 @@ def call_module( self.module, target, args, # type: ignore[arg-type] - kwargs, # type: ignore[arg-type] + kwargs, self.normalize_to_only_use_kwargs, ) if new_args_and_kwargs: diff --git a/torch/fx/passes/net_min_base.py b/torch/fx/passes/net_min_base.py index 81ab0e475d1ae..2a093bea49a4c 100644 --- a/torch/fx/passes/net_min_base.py +++ b/torch/fx/passes/net_min_base.py @@ -256,7 +256,7 @@ def _tag_nodes(self, selected_nodes: NodeSet): if node in selected_nodes: node.tag = "minimize" elif any( - n.tag in {"minimize", "main_1"} # type: ignore[attr-defined] + n.tag in {"minimize", "main_1"} for n in node.all_input_nodes if n.op in CALLABLE_NODE_OPS ): diff --git a/torch/fx/passes/split_module.py b/torch/fx/passes/split_module.py index 238f2a947c0b2..989ec92777cc3 100644 --- a/torch/fx/passes/split_module.py +++ b/torch/fx/passes/split_module.py @@ -10,8 +10,8 @@ def __init__(self, name: str): self.outputs: Dict[str, None] = {} self.partitions_dependent_on: Dict[str, None] = {} self.partition_dependents: Dict[str, None] = {} - self.graph : torch.fx.graph.Graph = torch.fx.graph.Graph() # type: ignore[attr-defined, name-defined] - self.environment : Dict[torch.fx.node.Node, torch.fx.node.Node] = {} # type: ignore[name-defined] + self.graph : torch.fx.graph.Graph = torch.fx.graph.Graph() + self.environment : Dict[torch.fx.node.Node, torch.fx.node.Node] = {} self.targets : Dict[str, Any] = {} def __repr__(self) -> str: @@ -26,12 +26,12 @@ def __repr__(self) -> str: def split_module( m: GraphModule, root_m: torch.nn.Module, - split_callback: Callable[[torch.fx.node.Node], int], # type: ignore[name-defined] + split_callback: Callable[[torch.fx.node.Node], int], ): partitions: Dict[str, Partition] = {} - orig_nodes: Dict[str, torch.fx.node.Node] = {} # type: ignore[name-defined] + orig_nodes: Dict[str, torch.fx.node.Node] = {} - def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optional[torch.fx.node.Node]): # type: ignore[name-defined] # noqa: B950 + def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optional[torch.fx.node.Node]): # noqa: B950 def_partition_name = getattr(def_node, '_fx_partition', None) use_partition_name = getattr(use_node, '_fx_partition', None) if def_partition_name != use_partition_name: @@ -56,7 +56,7 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona if node.op in ["placeholder", "get_attr"]: continue if node.op == 'output': - torch.fx.graph.map_arg(node.args[0], lambda n: record_cross_partition_use(n, None)) # type: ignore[attr-defined] + torch.fx.graph.map_arg(node.args[0], lambda n: record_cross_partition_use(n, None)) continue partition_name = str(split_callback(node)) @@ -68,8 +68,8 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona partition.node_names.append(node.name) node._fx_partition = partition_name - torch.fx.graph.map_arg(node.args, lambda def_node: record_cross_partition_use(def_node, node)) # type: ignore[attr-defined] - torch.fx.graph.map_arg(node.kwargs, lambda def_node: record_cross_partition_use(def_node, node)) # type: ignore[attr-defined] # noqa: B950 + torch.fx.graph.map_arg(node.args, lambda def_node: record_cross_partition_use(def_node, node)) + torch.fx.graph.map_arg(node.kwargs, lambda def_node: record_cross_partition_use(def_node, node)) # noqa: B950 # find partitions with no dependencies root_partitions : List[str] = [] @@ -104,8 +104,8 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona # swap out old graph nodes in kw/args with references to new nodes in this submodule environment = partition.environment - gathered_args = torch.fx.graph.map_arg(node.args, lambda n : environment[n]) # type: ignore[attr-defined] - gathered_kwargs = torch.fx.graph.map_arg(node.kwargs, lambda n : environment[n]) # type: ignore[attr-defined] + gathered_args = torch.fx.graph.map_arg(node.args, lambda n : environment[n]) + gathered_kwargs = torch.fx.graph.map_arg(node.kwargs, lambda n : environment[n]) if node.op not in ['call_module', 'get_attr']: target = node.target @@ -128,9 +128,9 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona partition.environment[node] = new_node # Set up values to construct base module - base_mod_env : Dict[str, torch.fx.node.Node] = {} # type: ignore[name-defined] - base_mod_graph : torch.fx.graph.Graph = torch.fx.graph.Graph() # type: ignore[attr-defined, name-defined] - base_mod_attrs : Dict[str, torch.fx.graph_module.GraphModule] = {} # type: ignore[name-defined] + base_mod_env : Dict[str, torch.fx.node.Node] = {} + base_mod_graph : torch.fx.graph.Graph = torch.fx.graph.Graph() + base_mod_attrs : Dict[str, torch.fx.graph_module.GraphModule] = {} for node in m.graph.nodes: if node.op == 'placeholder': base_mod_env[node.name] = base_mod_graph.placeholder(node.name) @@ -159,14 +159,14 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona # Construct GraphModule for this partition submod_name = f'submod_{partition_name}' - base_mod_attrs[submod_name] = torch.fx.graph_module.GraphModule(partition.targets, partition.graph) # type: ignore[attr-defined] # noqa: B950 + base_mod_attrs[submod_name] = torch.fx.graph_module.GraphModule(partition.targets, partition.graph) # noqa: B950 # Emit call in base graph to this submodule output_val = base_mod_graph.call_module(submod_name, tuple(base_mod_env[name] for name in partition.inputs)) if len(partition.outputs) > 1: # Unpack multiple return values from submodule - output_val_proxy = torch.fx.proxy.Proxy(output_val) # type: ignore[attr-defined] + output_val_proxy = torch.fx.proxy.Proxy(output_val) for i, output_name in enumerate(partition.outputs): base_mod_env[output_name] = output_val_proxy[i].node # type: ignore[index] else: @@ -174,6 +174,6 @@ def record_cross_partition_use(def_node : torch.fx.node.Node, use_node : Optiona for node in m.graph.nodes: if node.op == 'output': - base_mod_graph.output(torch.fx.graph.map_arg(node.args[0], lambda n : base_mod_env[n.name])) # type: ignore[attr-defined] # noqa: B950 + base_mod_graph.output(torch.fx.graph.map_arg(node.args[0], lambda n : base_mod_env[n.name])) # noqa: B950 - return torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph) # type: ignore[attr-defined] + return torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph) diff --git a/torch/jit/_monkeytype_config.py b/torch/jit/_monkeytype_config.py index c124602d0cfd5..f71192fabb5e7 100644 --- a/torch/jit/_monkeytype_config.py +++ b/torch/jit/_monkeytype_config.py @@ -53,7 +53,7 @@ def analyze(self, qualified_name: str) -> Dict: # and create a dictionary of all the types # for arguments. records = self.trace_records[qualified_name] - all_args = defaultdict(set) # type: ignore[var-annotated] + all_args = defaultdict(set) for record in records: for arg, arg_type in record.arg_types.items(): all_args[arg].add(arg_type) @@ -123,4 +123,4 @@ class JitTypeTraceConfig: # type: ignore[no-redef] def __init__(self): pass - monkeytype_trace = None # type: ignore[assignment] # noqa: F811 + monkeytype_trace = None # noqa: F811 diff --git a/torch/jit/mobile/__init__.py b/torch/jit/mobile/__init__.py index ef77e78c8db2b..0172d52fa63ae 100644 --- a/torch/jit/mobile/__init__.py +++ b/torch/jit/mobile/__init__.py @@ -45,9 +45,9 @@ def _load_for_lite_interpreter(f, map_location=None): map_location = validate_map_location(map_location) if isinstance(f, str) or isinstance(f, pathlib.Path): - cpp_module = torch._C._load_for_lite_interpreter(f, map_location) # type: ignore[attr-defined] + cpp_module = torch._C._load_for_lite_interpreter(f, map_location) else: - cpp_module = torch._C._load_for_lite_interpreter_from_buffer(f.read(), map_location) # type: ignore[attr-defined] + cpp_module = torch._C._load_for_lite_interpreter_from_buffer(f.read(), map_location) return LiteScriptModule(cpp_module) @@ -102,9 +102,9 @@ def _get_model_bytecode_version(f_input) -> int: raise ValueError(f"The provided filename {f_input} is a directory") if (isinstance(f_input, str) or isinstance(f_input, pathlib.Path)): - return torch._C._get_model_bytecode_version(str(f_input)) # type: ignore[attr-defined] + return torch._C._get_model_bytecode_version(str(f_input)) else: - return torch._C._get_model_bytecode_version_from_buffer(f_input.read()) # type: ignore[attr-defined] + return torch._C._get_model_bytecode_version_from_buffer(f_input.read()) def _backport_for_mobile(f_input, f_output, to_version): r""" @@ -124,9 +124,9 @@ def _backport_for_mobile(f_input, f_output, to_version): if ((isinstance(f_input, str) or isinstance(f_input, pathlib.Path)) and ( isinstance(f_output, str) or isinstance(f_output, pathlib.Path))): - return torch._C._backport_for_mobile(str(f_input), str(f_output), to_version) # type: ignore[attr-defined] + return torch._C._backport_for_mobile(str(f_input), str(f_output), to_version) else: - return torch._C._backport_for_mobile_from_buffer(f_input.read(), str(f_output), to_version) # type: ignore[attr-defined] + return torch._C._backport_for_mobile_from_buffer(f_input.read(), str(f_output), to_version) def _backport_for_mobile_to_buffer(f_input, to_version): r""" @@ -142,9 +142,9 @@ def _backport_for_mobile_to_buffer(f_input, to_version): raise ValueError(f"The provided filename {f_input} is a directory") if (isinstance(f_input, str) or isinstance(f_input, pathlib.Path)): - return torch._C._backport_for_mobile_to_buffer(str(f_input), to_version) # type: ignore[attr-defined] + return torch._C._backport_for_mobile_to_buffer(str(f_input), to_version) else: - return torch._C._backport_for_mobile_from_buffer_to_buffer(f_input.read(), to_version) # type: ignore[attr-defined] + return torch._C._backport_for_mobile_from_buffer_to_buffer(f_input.read(), to_version) def _get_model_ops_and_info(f_input): r""" @@ -182,6 +182,6 @@ def _get_model_ops_and_info(f_input): raise ValueError(f"The provided filename {f_input} is a directory") if (isinstance(f_input, str) or isinstance(f_input, pathlib.Path)): - return torch._C._get_model_ops_and_info(str(f_input)) # type: ignore[attr-defined] + return torch._C._get_model_ops_and_info(str(f_input)) else: - return torch._C._get_model_ops_and_info(f_input.read()) # type: ignore[attr-defined] + return torch._C._get_model_ops_and_info(f_input.read()) diff --git a/torch/nn/quantized/modules/conv.py b/torch/nn/quantized/modules/conv.py index a458ac5142714..9d5ccaea034ef 100644 --- a/torch/nn/quantized/modules/conv.py +++ b/torch/nn/quantized/modules/conv.py @@ -201,7 +201,7 @@ def get_qconv(cls, mod, activation_post_process, weight_post_process=None): 'Weight observer must have a dtype of qint8' qweight = _quantize_weight(mod.weight.float(), weight_post_process) # the __init__ call used is the one from derived classes and not the one from _ConvNd - qconv = cls(mod.in_channels, mod.out_channels, mod.kernel_size, # type: ignore[call-arg] + qconv = cls(mod.in_channels, mod.out_channels, mod.kernel_size, mod.stride, mod.padding, mod.dilation, mod.groups, mod.bias is not None, mod.padding_mode) qconv.set_weight_bias(qweight, mod.bias) diff --git a/torch/nn/utils/parametrizations.py b/torch/nn/utils/parametrizations.py index baf634563ca3b..7941f41f19cac 100644 --- a/torch/nn/utils/parametrizations.py +++ b/torch/nn/utils/parametrizations.py @@ -90,7 +90,7 @@ def _power_method(self, weight_mat: torch.Tensor, n_power_iterations: int) -> No # This power iteration produces approximations of `u` and `v`. self._u = F.normalize(torch.mv(weight_mat, self._v), # type: ignore[has-type] dim=0, eps=self.eps, out=self._u) # type: ignore[has-type] - self._v = F.normalize(torch.mv(weight_mat.t(), self._u), # type: ignore[has-type] + self._v = F.normalize(torch.mv(weight_mat.t(), self._u), dim=0, eps=self.eps, out=self._v) # type: ignore[has-type] # See above on why we need to clone self._u = self._u.clone(memory_format=torch.contiguous_format) diff --git a/torch/package/package_importer.py b/torch/package/package_importer.py index f901d46a3e96c..e793df3ff342b 100644 --- a/torch/package/package_importer.py +++ b/torch/package/package_importer.py @@ -366,7 +366,7 @@ def _install_on_parent(self, parent: str, name: str, module: types.ModuleType): return # Set the module as an attribute on its parent. parent_module = self.modules[parent] - if parent_module.__loader__ is self: # type: ignore[union-attr] + if parent_module.__loader__ is self: setattr(parent_module, name.rpartition(".")[2], module) # note: copied from cpython's import code, with call to create module replaced with _make_module diff --git a/torch/quantization/fx/convert.py b/torch/quantization/fx/convert.py index 1906c37ee1e9d..75d95b1b9e57c 100644 --- a/torch/quantization/fx/convert.py +++ b/torch/quantization/fx/convert.py @@ -394,7 +394,7 @@ def insert_quantize_node(node: Node, modules: Dict[str, torch.nn.Module]) -> Non # for non-standalone module, since _standalone_module_output_quantized_idxs # is only available in observed standalone module if is_observed_standalone_module_node: - out_quant_idxs = modules[node.target]._standalone_module_output_quantized_idxs.tolist() # type: ignore[operator] # noqa: B950 + out_quant_idxs = modules[node.target]._standalone_module_output_quantized_idxs.tolist() # noqa: B950 assert len(out_quant_idxs) <= 1, "Currently standalone only support one output" quantized = 0 in out_quant_idxs diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index 122d36aca68d0..39a0f69033420 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -425,7 +425,7 @@ def maybe_insert_input_observers_for_node( # assign the new args and kwargs to the node, inplace node.args = tuple(new_args) - node.kwargs = new_kwargs # type: ignore[assignment] + node.kwargs = new_kwargs def maybe_insert_input_equalization_observers_for_node( node: Node, @@ -946,7 +946,7 @@ def run_prepare_fx_on_standalone_modules( get_standalone_module_configs( root_node, modules, prepare_custom_config_dict, qconfig) - standalone_module = modules[root_node.target] # type: ignore[index] + standalone_module = modules[root_node.target] prepare = \ torch.quantization.quantize_fx._prepare_standalone_module_fx # type: ignore[attr-defined] observed_standalone_module = \ @@ -959,7 +959,7 @@ def run_prepare_fx_on_standalone_modules( parent_name, name = _parent_name(root_node.target) setattr(modules[parent_name], name, observed_standalone_module) - modules[root_node.target] = observed_standalone_module # type: ignore[index] + modules[root_node.target] = observed_standalone_module def save_state( observed: GraphModule, diff --git a/torch/testing/_internal/common_distributed.py b/torch/testing/_internal/common_distributed.py index 216b60137c031..169ba99a2bd7a 100644 --- a/torch/testing/_internal/common_distributed.py +++ b/torch/testing/_internal/common_distributed.py @@ -508,7 +508,7 @@ def run_test(self, test_name: str, pipe) -> None: if sys.platform != 'win32' and sys.platform != 'darwin': # Register signal handler to dump stack traces on FATALs. # Windows and MacOS do not support the signal handlers. - torch._C._set_print_stack_traces_on_fatal_signal(True) # type: ignore[attr-defined] + torch._C._set_print_stack_traces_on_fatal_signal(True) # self.id() == e.g. '__main__.TestDistributed.test_get_rank' # We're retrieving a corresponding test and executing it. diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 385950c4ffabb..8faa1f1f535a0 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -3292,24 +3292,24 @@ def sample_inputs_fmod_remainder(op_info, device, dtype, requires_grad, *, autod make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) if autodiffed: - samples = ( # type: ignore[assignment] + samples = ( ((S, S, S), 1.5, False), ((), 1.5, False), ) else: - cases = ( # type: ignore[assignment] + cases = ( ((S, S, S), (), False), ((S, S, S), (S, S, S), False), ((S, S, S), (S,), False), ) # Sample inputs with scalars as torch tensors - cases_with_tensor_scalar = ( # type: ignore[assignment] + cases_with_tensor_scalar = ( ((), torch.tensor(1, dtype=dtype, device=device, requires_grad=False), False), ) # Sample inputs with broadcasting - cases_with_broadcasting = ( # type: ignore[assignment] + cases_with_broadcasting = ( ((S,), (S, S, S), True), ((S, 1, S), (S, S, S), True), ((), (S, S, S), True), @@ -3978,7 +3978,7 @@ def sample_inputs_split(op_info, device, dtype, requires_grad, *, list_args=Fals make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) if list_args: - cases = ( # type: ignore[assignment] + cases = ( ((S, S, S), ([int(S / 3), S - int(S / 3) * 2, int(S / 3)],)), ((S, S, S), ([int(S / 2), S - int(S / 2) * 2, int(S / 2)], 2),), ((S, S, S), ([int(S / 2), S - int(S / 2) * 2, int(S / 2)], -2),) diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 6d7ca8ac5ef82..47c36568f50d3 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -1440,7 +1440,7 @@ def assertEqual(self, x, y, msg: Optional[str] = None, *, super().assertEqual(x, y, msg=msg) def assertNotEqual(self, x, y, msg: Optional[str] = None, *, # type: ignore[override] - atol: Optional[float] = None, rtol: Optional[float] = None, **kwargs) -> None: # type: ignore[override] + atol: Optional[float] = None, rtol: Optional[float] = None, **kwargs) -> None: with self.assertRaises(AssertionError, msg=msg): self.assertEqual(x, y, msg, atol=atol, rtol=rtol, **kwargs) diff --git a/torch/testing/_internal/dist_utils.py b/torch/testing/_internal/dist_utils.py index 40fc21c229b47..1c5bf2b2f8c0c 100644 --- a/torch/testing/_internal/dist_utils.py +++ b/torch/testing/_internal/dist_utils.py @@ -6,7 +6,7 @@ import torch.distributed as dist import torch.distributed.rpc as rpc -from torch.distributed.rpc import _rref_context_get_debug_info # type: ignore[attr-defined] +from torch.distributed.rpc import _rref_context_get_debug_info from torch.testing._internal.common_utils import FILE_SCHEMA diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py index b80acbde7cabd..50535d3144bd7 100644 --- a/torch/utils/data/dataloader.py +++ b/torch/utils/data/dataloader.py @@ -265,11 +265,9 @@ def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1, sampler = _InfiniteConstantSampler() else: # map-style if shuffle: - # Cannot statically verify that dataset is Sized - # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] - sampler = RandomSampler(dataset, generator=generator) # type: ignore[arg-type] + sampler = RandomSampler(dataset, generator=generator) else: - sampler = SequentialSampler(dataset) # type: ignore[arg-type] + sampler = SequentialSampler(dataset) if batch_size is not None and batch_sampler is None: # auto_collation without custom batch_sampler diff --git a/torch/utils/data/dataset.py b/torch/utils/data/dataset.py index a1549a0e97b0e..b875da6ce52fd 100644 --- a/torch/utils/data/dataset.py +++ b/torch/utils/data/dataset.py @@ -295,7 +295,6 @@ def __len__(self): total = 0 for d in self.datasets: assert isinstance(d, IterableDataset), "ChainDataset only supports IterableDataset" - # Cannot verify that all self.datasets are Sized total += len(d) return total @@ -338,7 +337,7 @@ def random_split(dataset: Dataset[T], lengths: Sequence[int], generator (Generator): Generator used for the random permutation. """ # Cannot verify that dataset is Sized - if sum(lengths) != len(dataset): # type: ignore[arg-type] + if sum(lengths) != len(dataset): raise ValueError("Sum of input lengths does not equal the length of the input dataset!") indices = randperm(sum(lengths), generator=generator).tolist() diff --git a/torch/utils/data/distributed.py b/torch/utils/data/distributed.py index 7ef638ada118e..929cf9401b2a7 100644 --- a/torch/utils/data/distributed.py +++ b/torch/utils/data/distributed.py @@ -78,17 +78,15 @@ def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None, self.drop_last = drop_last # If the dataset length is evenly divisible by # of replicas, then there # is no need to drop any data, since the dataset will be split equally. - if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type] + if self.drop_last and len(self.dataset) % self.num_replicas != 0: # Split to nearest available length that is evenly divisible. # This is to ensure each rank receives the same amount of data when # using this Sampler. self.num_samples = math.ceil( - # `type:ignore` is required because Dataset cannot provide a default __len__ - # see NOTE in pytorch/torch/utils/data/sampler.py - (len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type] + (len(self.dataset) - self.num_replicas) / self.num_replicas ) else: - self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type] + self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) self.total_size = self.num_samples * self.num_replicas self.shuffle = shuffle self.seed = seed @@ -98,9 +96,9 @@ def __iter__(self) -> Iterator[T_co]: # deterministically shuffle based on epoch and seed g = torch.Generator() g.manual_seed(self.seed + self.epoch) - indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] + indices = torch.randperm(len(self.dataset), generator=g).tolist() else: - indices = list(range(len(self.dataset))) # type: ignore[arg-type] + indices = list(range(len(self.dataset))) if not self.drop_last: # add extra samples to make it evenly divisible diff --git a/torch/utils/model_dump/__init__.py b/torch/utils/model_dump/__init__.py index 8c23e6e395967..2922567a42f26 100644 --- a/torch/utils/model_dump/__init__.py +++ b/torch/utils/model_dump/__init__.py @@ -334,10 +334,10 @@ def get_inline_skeleton(): import importlib.resources - skeleton = importlib.resources.read_text(__package__, "skeleton.html") # type: ignore[attr-defined] - js_code = importlib.resources.read_text(__package__, "code.js") # type: ignore[attr-defined] + skeleton = importlib.resources.read_text(__package__, "skeleton.html") + js_code = importlib.resources.read_text(__package__, "code.js") for js_module in ["preact", "htm"]: - js_lib = importlib.resources.read_binary(__package__, f"{js_module}.mjs") # type: ignore[attr-defined] + js_lib = importlib.resources.read_binary(__package__, f"{js_module}.mjs") js_url = "data:application/javascript," + urllib.parse.quote(js_lib) js_code = js_code.replace(f"https://unpkg.com/{js_module}?module", js_url) skeleton = skeleton.replace(' src="./code.js">', ">\n" + js_code) diff --git a/torch/utils/tensorboard/writer.py b/torch/utils/tensorboard/writer.py index d2362cc4eecef..a37af381978e1 100644 --- a/torch/utils/tensorboard/writer.py +++ b/torch/utils/tensorboard/writer.py @@ -837,7 +837,7 @@ def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, ta metadata, label_img, fs, subdir, global_step, tag) self._projector_config.embeddings.extend([embedding_info]) - from google.protobuf import text_format # type: ignore[attr-defined] + from google.protobuf import text_format config_pbtxt = text_format.MessageToString(self._projector_config) write_pbtxt(self._get_file_writer().get_logdir(), config_pbtxt) From 3815a013ed76e2a0a0b15fbcd0d419d9be322254 Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Fri, 18 Jun 2021 09:17:47 -0700 Subject: [PATCH 234/305] Enable xenial-cuda11.1-cudnn8-py3.6-gcc7 in GHA (#60196) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60196 Test Plan: https://github.com/pytorch/pytorch/issues/60198: https://github.com/pytorch/pytorch/actions/runs/947796763 I should have used `ghstack` but I forgot; will do that in the future. Reviewed By: walterddr Differential Revision: D29231161 Pulled By: samestep fbshipit-source-id: 8299a248ca9c1d36c3845d1c8a10ca9bf7101124 --- .github/scripts/generate_ci_workflows.py | 10 +- ...inux-xenial-cuda11.1-cudnn8-py3.6-gcc7.yml | 374 ++++++++++++++++++ 2 files changed, 379 insertions(+), 5 deletions(-) create mode 100644 .github/workflows/pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7.yml diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index 84f6b58be574e..ca15fbbb1fdc6 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -140,11 +140,11 @@ def generate_workflow_file( docker_image_base=f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-xenial-cuda10.2-cudnn7-py3-gcc7", test_runner_type=LINUX_CUDA_TEST_RUNNER, ), - # PyTorchLinuxWorkflow( - # build_environment="pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7", - # docker_image_base=f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-xenial-cuda11.1-cudnn8-py3-gcc7", - # test_runner_type=LINUX_CUDA_TEST_RUNNER, - # ), + PyTorchLinuxWorkflow( + build_environment="pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7", + docker_image_base=f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-xenial-cuda11.1-cudnn8-py3-gcc7", + test_runner_type=LINUX_CUDA_TEST_RUNNER, + ), # PyTorchLinuxWorkflow( # build_environment="pytorch-libtorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7", # docker_image_base=f"{DOCKER_REGISTRY}/pytorch/pytorch-linux-xenial-cuda11.1-cudnn8-py3-gcc7", diff --git a/.github/workflows/pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7.yml b/.github/workflows/pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7.yml new file mode 100644 index 0000000000000..f5e6f2fd28021 --- /dev/null +++ b/.github/workflows/pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7.yml @@ -0,0 +1,374 @@ +# @generated DO NOT EDIT MANUALLY +# Template is at: .github/templates/linux_ci_workflow.yml.j2 +# Generation script: .github/scripts/generate_ci_workflows.py +name: Linux CI (pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7) + +on: + # TODO: Enable pull_request builds when we can verify capacity can be met by auto-scalers + push: + branches: + - master + - release/* + workflow_dispatch: + +env: + BUILD_ENVIRONMENT: pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7 + DOCKER_IMAGE_BASE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/pytorch-linux-xenial-cuda11.1-cudnn8-py3-gcc7 + SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 + TORCH_CUDA_ARCH_LIST: 5.2 + IN_CI: 1 + # Used for custom_opertor, jit_hooks, custom_backend, see .jenkins/pytorch/build.sh + CUSTOM_TEST_ARTIFACT_BUILD_DIR: build/custom_test_artifacts + ALPINE_IMAGE: "308535385114.dkr.ecr.us-east-1.amazonaws.com/tool/alpine" + +concurrency: + group: pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: true + +jobs: + calculate-docker-image: + runs-on: linux.2xlarge + env: + DOCKER_BUILDKIT: 1 + timeout-minutes: 90 + outputs: + docker_image: ${{ steps.calculate-tag.outputs.docker_image }} + steps: + - name: Log in to ECR + run: | + aws ecr get-login --no-include-email --region us-east-1 > /tmp/ecr-login.sh + bash /tmp/ecr-login.sh + rm /tmp/ecr-login.sh + - name: Chown workspace + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + # deep clone, to allow use of git merge-base + fetch-depth: 0 + - name: Calculate docker image tag + id: calculate-tag + run: | + DOCKER_TAG=$(git rev-parse HEAD:.circleci/docker) + echo "::set-output name=docker_tag::${DOCKER_TAG}" + echo "::set-output name=docker_image::${DOCKER_IMAGE_BASE}:${DOCKER_TAG}" + - name: Check if image should be built + id: check + env: + DOCKER_TAG: ${{ steps.calculate-tag.outputs.docker_tag }} + BASE_REVISION: ${{ github.event.pull_request.base.sha || github.sha }} + run: | + set -x + # Check if image already exists, if it does then skip building it + if docker manifest inspect "${DOCKER_IMAGE_BASE}:${DOCKER_TAG}"; then + exit 0 + fi + if [[ "$BASE_REVISION" = "$(git rev-parse HEAD)" ]]; then + # if we're on the base branch then use the parent commit + MERGE_BASE=$(git rev-parse HEAD~) + else + # otherwise we're on a PR, so use the most recent base commit + MERGE_BASE=$(git merge-base HEAD "$BASE_REVISION") + fi + # Covers the case where a previous tag doesn't exist for the tree + # this is only really applicable on trees that don't have `.circleci/docker` at its merge base, i.e. nightly + if ! git rev-parse "$MERGE_BASE:.circleci/docker"; then + echo "Directory '.circleci/docker' not found in commit $MERGE_BASE, you should probably rebase onto a more recent commit" + exit 1 + fi + PREVIOUS_DOCKER_TAG=$(git rev-parse "$MERGE_BASE:.circleci/docker") + # If no image exists but the hash is the same as the previous hash then we should error out here + if [[ "${PREVIOUS_DOCKER_TAG}" = "${DOCKER_TAG}" ]]; then + echo "ERROR: Something has gone wrong and the previous image isn't available for the merge-base of your branch" + echo " contact the PyTorch team to restore the original images" + exit 1 + fi + echo ::set-output name=rebuild::yes + - name: Build and push docker image + if: steps.check.outputs.rebuild + env: + DOCKER_TAG: ${{ steps.calculate-tag.outputs.docker_tag }} + DOCKER_SKIP_S3_UPLOAD: 1 + run: | + export IMAGE_NAME=${DOCKER_IMAGE_BASE#308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/} + cd .circleci/docker && ./build_docker.sh + + build: + runs-on: linux.2xlarge + needs: calculate-docker-image + env: + DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7-build + steps: + - name: Log in to ECR + run: | + aws ecr get-login --no-include-email --region us-east-1 > /tmp/ecr-login.sh + bash /tmp/ecr-login.sh + rm /tmp/ecr-login.sh + - name: Chown workspace + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + fetch-depth: 0 # deep clone, to allow sharding to use git rev-list + submodules: recursive + - name: Pull docker image + run: | + docker pull "${DOCKER_IMAGE}" + - name: Preserve github env variables for use in docker + run: | + env | grep '^GITHUB' > "/tmp/github_env_${GITHUB_RUN_ID}" + - name: Build PyTorch + run: | + docker run \ + -e BUILD_ENVIRONMENT \ + -e JOB_BASE_NAME \ + -e MAX_JOBS="$(nproc --ignore=2)" \ + -e SCCACHE_BUCKET \ + -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ + -e SKIP_SCCACHE_INITIALIZATION=1 \ + -e TORCH_CUDA_ARCH_LIST \ + --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ + --security-opt seccomp=unconfined \ + --cap-add=SYS_PTRACE \ + --tty \ + --user jenkins \ + -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \ + -w /var/lib/jenkins/workspace \ + "${DOCKER_IMAGE}" \ + sh -c 'sudo chown -R jenkins . && .jenkins/pytorch/build.sh' + - name: Display and upload binary build size statistics (Click Me) + # temporary hack: set CIRCLE_* vars, until we update + # tools/print_test_stats.py to natively support GitHub Actions + env: + SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} + CIRCLE_BRANCH: ${{ steps.parse-ref.outputs.branch }} + CIRCLE_PR_NUMBER: ${{ github.event.pull_request.number }} + CIRCLE_SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + CIRCLE_TAG: ${{ steps.parse-ref.outputs.tag }} + CIRCLE_WORKFLOW_ID: ${{ github.run_id }} # dunno if this corresponds + run: | + export PYTHONPATH=$PWD + COMMIT_TIME=$(git log --max-count=1 --format=%ct || echo 0) + export COMMIT_TIME + pip3 install requests + python3 .circleci/scripts/upload_binary_size_to_scuba.py || exit 0 + - name: Chown workspace + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + - name: Archive artifacts into zip + run: | + zip -r artifacts.zip dist/ build/ .pytorch-test-times.json + # Upload to github so that people can click and download artifacts + - uses: actions/upload-artifact@v2 + # Don't fail on upload to GH since it's only for user convenience + continue-on-error: true + name: Store PyTorch Build Artifacts on Github + with: + name: ${{ env.BUILD_ENVIRONMENT }} + retention-days: 14 + if-no-files-found: error + path: + artifacts.zip + - uses: seemethere/upload-artifact-s3@9d7ceb0ab39c2c88d93ef7792b27425b27d59162 + name: Store PyTorch Build Artifacts on S3 + with: + name: ${{ env.BUILD_ENVIRONMENT }} + retention-days: 14 + if-no-files-found: error + path: + artifacts.zip + - name: Clean up docker images + if: always() + run: | + # Prune all of the docker images + docker system prune -af + + generate-test-matrix: + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-18.04 + env: + NUM_TEST_SHARDS: 1 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + container: + image: python:3.9 + steps: + - name: Clone pytorch/pytorch + uses: actions/checkout@v2 + - name: Generating test matrix + id: set-matrix + run: | + # outputting for debugging purposes + MATRIX=$(python .github/scripts/generate_pytorch_test_matrix.py) + echo "${MATRIX}" + echo "::set-output name=matrix::${MATRIX}" + + test: + runs-on: linux.8xlarge.nvidia.gpu + needs: + - calculate-docker-image + - build + - generate-test-matrix + strategy: + matrix: ${{ fromJson(needs.generate-test-matrix.outputs.matrix) }} + fail-fast: false + env: + DOCKER_IMAGE: ${{ needs.calculate-docker-image.outputs.docker_image }} + JOB_BASE_NAME: pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7-test + NUM_TEST_SHARDS: 1 + TEST_CONFIG: ${{ matrix.test_config }} + steps: + - name: Log in to ECR + run: | + aws ecr get-login --no-include-email --region us-east-1 > /tmp/ecr-login.sh + bash /tmp/ecr-login.sh + rm /tmp/ecr-login.sh + - name: Chown workspace + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)/../":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + submodules: recursive + - name: Pull docker image + run: | + docker pull "${DOCKER_IMAGE}" + - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG + if: ${{ contains(env.BUILD_ENVIRONMENT, 'cuda') }} + run: | + bash .github/scripts/install_nvidia_utils_linux.sh + echo "GPU_FLAG=--gpus all" >> "${GITHUB_ENV}" + - name: Determine shm-size + run: | + shm_size="1g" + case "${BUILD_ENVIRONMENT}" in + *cuda*) + shm_size="2g" + ;; + *rocm*) + shm_size="8g" + ;; + esac + echo "SHM_SIZE=${shm_size}" >> "${GITHUB_ENV}" + - uses: seemethere/download-artifact-s3@0504774707cbc8603d7dca922e8026eb8bf3b47b + name: Download PyTorch Build Artifacts + with: + name: ${{ env.BUILD_ENVIRONMENT }} + - name: Unzip artifacts + run: | + unzip -o artifacts.zip + - name: Output disk space left + run: | + sudo df -H + - name: Preserve github env variables for use in docker + run: | + env | grep '^GITHUB' > "/tmp/github_env_${GITHUB_RUN_ID}" + - name: Test PyTorch + run: | + if [[ $NUM_TEST_SHARDS -eq 2 ]]; then + export SHARD_NUMBER=$TEST_CONFIG + else + export SHARD_NUMBER=0 + fi + # TODO: Stop building test binaries as part of the build phase + # Used for GPU_FLAG since that doesn't play nice + # shellcheck disable=SC2086 + docker run \ + ${GPU_FLAG:-} \ + -e BUILD_ENVIRONMENT \ + -e CUSTOM_TEST_ARTIFACT_BUILD_DIR \ + -e GITHUB_ACTIONS \ + -e IN_CI \ + -e SHARD_NUMBER \ + -e JOB_BASE_NAME \ + -e MAX_JOBS="$(nproc --ignore=2)" \ + --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ + --security-opt seccomp=unconfined \ + --cap-add=SYS_PTRACE \ + --shm-size="${SHM_SIZE}" \ + --tty \ + --user jenkins \ + -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \ + -w /var/lib/jenkins/workspace \ + "${DOCKER_IMAGE}" \ + sh -c 'sudo chown -R jenkins . && pip install dist/*.whl && .jenkins/pytorch/test.sh' + - name: Chown workspace + if: always() + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + - uses: actions/upload-artifact@v2 + name: Store PyTorch Test Reports + if: always() + with: + name: test-reports + retention-days: 14 + if-no-files-found: error + path: + test/**/*.xml + - name: Clean up docker images + if: always() + run: | + # Ensure the working directory gets chowned back to the current user + docker run --rm -v "$(pwd)":/v -w /v "${ALPINE_IMAGE}" chown -R "$(id -u):$(id -g)" . + # Prune all of the docker images + docker system prune -af + + # this is a separate step from test because the log files from test are too + # long: basically, GitHub tries to render all of the log files when you click + # through an action causing extreme slowdown on actions that contain too many + # logs (like test); we can always move it back to the other one, but it + # doesn't create the best experience + render_test_results: + if: always() + needs: + - test + runs-on: ubuntu-18.04 + steps: + - name: Checkout PyTorch + uses: actions/checkout@v2 + with: + # deep clone, to allow tools/print_test_stats.py to use Git commands + fetch-depth: 0 + - uses: actions/download-artifact@v2 + name: Download PyTorch Test Reports + with: + name: test-reports + path: test/test-reports + - uses: actions/setup-python@v2 + with: + python-version: 3.9 + - name: Install dependencies + # boto3 version copied from .circleci/docker/common/install_conda.sh + run: | + pip install -r requirements.txt + pip install boto3==1.16.34 junitparser rich + - name: Output Test Results (Click Me) + run: | + python tools/render_junit.py test + - name: Parse ref + id: parse-ref + run: .github/scripts/parse_ref.py + - name: Display and upload test statistics (Click Me) + # temporary hack: set CIRCLE_* vars, until we update + # tools/print_test_stats.py to natively support GitHub Actions + env: + SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_OSSCI_METRICS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_OSSCI_METRICS_SECRET_ACCESS_KEY }} + CIRCLE_BRANCH: ${{ steps.parse-ref.outputs.branch }} + CIRCLE_JOB: pytorch-linux-xenial-cuda11.1-cudnn8-py3.6-gcc7 + CIRCLE_PR_NUMBER: ${{ github.event.pull_request.number }} + CIRCLE_SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + CIRCLE_TAG: ${{ steps.parse-ref.outputs.tag }} + CIRCLE_WORKFLOW_ID: ${{ github.run_id }} # dunno if this corresponds + run: | + export PYTHONPATH=$PWD + python tools/print_test_stats.py --upload-to-s3 --compare-with-s3 test From d5df274ea51ccf020480d69a830d038949e1a008 Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Fri, 18 Jun 2021 09:18:28 -0700 Subject: [PATCH 235/305] [DDP] Support for multiple backwards (#59359) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59359 Move `prepare_for_backward` into `_DDPSink` backward instead of calling it in DDP forward pass so that we can run multiple backwards in DDP with `retain_graph=True`. ghstack-source-id: 131774159 Test Plan: CI Reviewed By: zhaojuanmao Differential Revision: D28855226 fbshipit-source-id: 6b7b25d75b7696f5b5629078233433f97663d61c --- test/distributed/test_c10d_nccl.py | 20 +-- torch/csrc/distributed/c10d/init.cpp | 8 +- torch/lib/c10d/reducer.cpp | 83 ++++++---- torch/lib/c10d/reducer.hpp | 27 +++- torch/nn/parallel/distributed.py | 85 ++++++----- .../_internal/distributed/distributed_test.py | 144 +++++++++++++++++- 6 files changed, 284 insertions(+), 83 deletions(-) diff --git a/test/distributed/test_c10d_nccl.py b/test/distributed/test_c10d_nccl.py index 5583cbb8a32b1..a22cfa0352ec2 100644 --- a/test/distributed/test_c10d_nccl.py +++ b/test/distributed/test_c10d_nccl.py @@ -1024,7 +1024,7 @@ def forward(self, x): # Compute loss and gradients for both outputs output1, output2 = model(input) loss1 = criterion(output1, target) - loss1.backward() + loss1.backward(retain_graph=True) loss2 = criterion(output2, target) loss2.backward() @@ -1087,9 +1087,8 @@ def check_no_grads(): check_no_grads() def _test_accumulate_gradients_module(self, gradient_as_bucket_view=False): - # This is NOT the recommended way to implement accumulating grads, but - # we would like to make sure DDP does not mess up with the underlying - # module. + # Test gradient accumulation via model.no_sync context manager, which is + # the supported way of implementing gradient accumulation. int_devices = gpus_for_rank(self.world_size)[self.rank][:1] devices = [torch.device("cuda:" + str(i)) for i in int_devices] store = c10d.FileStore(self.file_name, self.world_size) @@ -1118,12 +1117,13 @@ def step_model(model, input, target): step_model(model, input, target) if iteration % 2 == 0: - # Skip gradients sync without calling prepare_for_backward - step_model( - ddp_model.module, - input[self.rank: (self.rank + 1)], - target[self.rank: (self.rank + 1)], - ) + # Skip gradients sync using no_sync context manager. + with ddp_model.no_sync(): + step_model( + ddp_model, + input[self.rank: (self.rank + 1)], + target[self.rank: (self.rank + 1)], + ) for i, j in zip(model.parameters(), ddp_model.parameters()): self.assertNotEqual(i.grad, j.grad) else: diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp index 493d5a180af0f..a9f78e9e2b794 100644 --- a/torch/csrc/distributed/c10d/init.cpp +++ b/torch/csrc/distributed/c10d/init.cpp @@ -360,6 +360,7 @@ An enum-like class for built-in communication hooks: ``ALLREDUCE`` and ``FP16_CO .def( "prepare_for_forward", &::c10d::Reducer::prepare_for_forward, + py::arg("will_run_grad_reduction") = true, py::call_guard()) .def( "prepare_for_backward", @@ -414,7 +415,12 @@ An enum-like class for built-in communication hooks: ``ALLREDUCE`` and ``FP16_CO const std::shared_ptr<::c10d::Logger> logger) { std::weak_ptr<::c10d::Logger> logger_weakref = logger; reducer.set_logger(logger_weakref); - }); + }) + .def( + "_static_graph_first_bwd", + &::c10d::Reducer::static_graph_first_bwd, + py::call_guard() + ); shared_ptr_class_<::c10d::Logger>(module, "Logger") .def( diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp index 1f556cc126976..70317ce8e3708 100644 --- a/torch/lib/c10d/reducer.cpp +++ b/torch/lib/c10d/reducer.cpp @@ -119,6 +119,7 @@ Reducer::Reducer( gradient_as_bucket_view_(gradient_as_bucket_view), local_used_maps_reduced_(false), num_iterations_(0), + num_backward_calls_(0), num_buckets_ready_(0), has_rebuilt_bucket_(false), bucket_bytes_cap_(bucket_bytes_cap), @@ -299,12 +300,12 @@ bool Reducer::dynamic_graph_find_unused() { return !static_graph_ && find_unused_parameters_; } -bool Reducer::static_graph_first_iteration() { - return static_graph_ && num_iterations_ == 1; +bool Reducer::static_graph_first_bwd() { + return static_graph_ && num_backward_calls_ == 1; } -bool Reducer::static_graph_after_first_iteration() { - return static_graph_ && num_iterations_ > 1; +bool Reducer::static_graph_after_first_bwd() { + return static_graph_ && num_backward_calls_ > 1; } void Reducer::initialize_local_used_map() { @@ -402,8 +403,7 @@ void Reducer::mark_variable_ready_dense(size_t variable_index) { // Gradient is undefined. When find_unused_parameters=True, ensure it is // not marked as locally used, otherwise we will be allreducing zero's // instead of not touching .grad field of parameter. - if (this->dynamic_graph_find_unused() || - this->static_graph_first_iteration()) { + if (this->dynamic_graph_find_unused() || this->static_graph_first_bwd()) { REDUCER_CHECK( local_used_maps_[0][variable_index].item() == 0, logger_, @@ -483,6 +483,15 @@ void Reducer::push_rebuilt_params_for_all_indices() { } void Reducer::push_rebuilt_params(const size_t& index) { + // NOTE: We don't check this in should_rebuild_bucket because that controls + // whether we should push rebuilt params and whether to actually kick off + // process to rebuild buckets, if we check this in should_rebuild_buckets then + // the latter would break. + if (all_rebuilt_params_pushed_) { + // We only enter here in the case we are calling multiple backwards with + // retain_graph=True in the iteration before rebuilding buckets. + return; + } rebuilt_params_.push_back(replicas_[0][index]); rebuilt_param_indices_.push_back(index); } @@ -569,7 +578,7 @@ void Reducer::autograd_hook(size_t index) { } // See Note [Skip allreducing local_used_maps_dev] - if (dynamic_graph_find_unused() || static_graph_first_iteration()) { + if (dynamic_graph_find_unused() || static_graph_first_bwd()) { // Since it gets here, this param has been used for this iteration. We want // to mark it in local_used_maps_. During no_sync session, the same var can // be set multiple times, which is OK as does not affect correctness. As @@ -587,7 +596,7 @@ void Reducer::autograd_hook(size_t index) { }); } - if (static_graph_first_iteration()) { + if (static_graph_first_bwd()) { numGradHooksTriggeredMap_[index] += 1; return; } @@ -614,7 +623,7 @@ void Reducer::autograd_hook(size_t index) { // will be broadcasted and initialized. // If it is static graph, after 1st iteration, check if a variable // is ready for communication based on numGradHooksTriggeredMap_. - if (static_graph_after_first_iteration()) { + if (static_graph_after_first_bwd()) { REDUCER_CHECK( numGradHooksTriggeredMapPerIteration_[index] > 0, logger_, @@ -825,7 +834,7 @@ void Reducer::mark_variable_ready(size_t variable_index) { } // Check that all buckets were completed and had their work kicked off. TORCH_INTERNAL_ASSERT(next_bucket_ == buckets_.size()); - if (static_graph_after_first_iteration() && should_rebuild_buckets()) { + if (static_graph_after_first_bwd() && should_rebuild_buckets()) { for (const auto& unused_index : unused_parameters_) { push_rebuilt_params(unused_index); } @@ -906,9 +915,11 @@ void Reducer::initialize_buckets( this->rpc_context_.set(ThreadLocalDistAutogradContext::getContextPtr()); #endif - // This shouldn't be called if we're expecting autograd hooks to fire. + // Note that we check !require_finalize instead of !expect_autograd_hooks + // since the latter is set in forward pass, and the former indicates + // at least one gradient hook has fired and we are in autograd execution. REDUCER_CHECK( - !expect_autograd_hooks_, + !require_finalize_, logger_, "`initialize_buckets` must NOT be called during autograd execution."); @@ -1061,6 +1072,10 @@ void Reducer::initialize_buckets( buckets_.push_back(std::move(bucket)); } + // Need to reset bucket.pending and variable.pending as buckets have been + // re-initialized and they must be appropriately set before the next backward + // pass. + reset_bucket_counting(); } // (see Note: "Gradient Layout Contract" in initialize_buckets). @@ -1138,14 +1153,29 @@ void Reducer::populate_bucket_views_out( } } -void Reducer::prepare_for_forward() { +void Reducer::prepare_for_forward(bool will_run_grad_reduction) { std::lock_guard lock(mutex_); + expect_autograd_hooks_ = will_run_grad_reduction; + // To maintain compatibility with current version, where prepare_for_forward + // is not called if will_run_grad_reduction is False. + if (!expect_autograd_hooks_) { + return; + } num_iterations_++; if (should_collect_runtime_stats()) { record_forward_compute_start_time(); } } +void Reducer::reset_variable_counting() { + // Reset unused parameter accounting. + has_marked_unused_parameters_ = false; + // Reset per iteration marked ready parameters. + perIterationReadyParams_.clear(); + // Reset bucket counting. + reset_bucket_counting(); +} + void Reducer::reset_bucket_counting() { next_bucket_ = 0; // Reset num_buckets_ready_ at the beginning of backward computation @@ -1227,22 +1257,12 @@ void Reducer::search_unused_parameters( void Reducer::prepare_for_backward( const std::vector& outputs) { std::lock_guard lock(mutex_); - + ++num_backward_calls_; backward_compute_start_time_ = current_time_in_nanos(); if (should_collect_runtime_stats()) { record_backward_compute_start_time(); } - // Reset accounting. - expect_autograd_hooks_ = true; - - reset_bucket_counting(); - - // Reset unused parameter accounting. - has_marked_unused_parameters_ = false; - // Reset per iteration marked ready parameters. - perIterationReadyParams_.clear(); - // If static graph is not set, search graph to detect unused parameters. // When static graph is set, unused_parameters_ will be detected and will // not change after 1st iteration. @@ -1402,9 +1422,9 @@ void Reducer::save_thread_local_state() { } void Reducer::finalize_backward() { - // No longer expect autograd hooks to fire after this function returns. + // Note that we don't reset expect_autograd_hooks_ so that we can re-run + // backwards with retain_graph=True. TORCH_INTERNAL_ASSERT(expect_autograd_hooks_); - expect_autograd_hooks_ = false; // No longer require call to finalize after this function returns. TORCH_INTERNAL_ASSERT(require_finalize_); @@ -1445,7 +1465,7 @@ void Reducer::finalize_backward() { } // See Note [Skip allreducing local_used_maps_dev] - if (dynamic_graph_find_unused() || static_graph_first_iteration()) { + if (dynamic_graph_find_unused() || static_graph_first_bwd()) { // Due to the lazy wait, it is possible that reduction of the current // iteration is still going when the one for next iteration gets kicked off. // For such case, we want to wait explicitly to make sure the reduction does @@ -1466,6 +1486,15 @@ void Reducer::finalize_backward() { local_used_maps_reduced_ = false; } + // Reset various accounting variables including bucket counting to ensure we + // can appropriately launch allreduce for each bucket in the next backwards. + reset_variable_counting(); + // If we populated rebuilt params list in this backward call, avoid + // repopulating in subsequent backward calls. In particular this is needed to + // avoid re-pushing parameters when calling multiple backwards with + // retain_graph=True. + all_rebuilt_params_pushed_ = all_rebuilt_params_pushed_ || !rebuilt_params_.empty(); + if (should_collect_runtime_stats()) { record_backward_comm_end_time(); } diff --git a/torch/lib/c10d/reducer.hpp b/torch/lib/c10d/reducer.hpp index 41900f19cbb45..c9a02aef8550b 100644 --- a/torch/lib/c10d/reducer.hpp +++ b/torch/lib/c10d/reducer.hpp @@ -83,9 +83,9 @@ class TORCH_API Reducer { // a call to this function can simply be omitted. void prepare_for_backward(const std::vector& outputs); - // Called at the begginning of forward() inside DistributedDataParallel, + // Called at the beginning of forward() inside DistributedDataParallel, // right now it caputures the starting time of forward in each iteration. - void prepare_for_forward(); + void prepare_for_forward(bool will_run_grad_reduction = true); // Returns the relative time in nanoseconds when gradients were ready, // with respect to the time `prepare_for_backward` was called. The outer @@ -157,6 +157,12 @@ class TORCH_API Reducer { // Delay all reduce to be after all gradients' calculation is complete. void delay_all_reduce(); + bool static_graph_first_bwd(); + + // Resets various counters Reducer uses to manager internal state such as + // buckets that need to be reduced across workers. + void reset_variable_counting(); + // Weak reference to associated DDP logger. The reference is weak to avoid // refcycle between reducer and logger. void set_logger(std::weak_ptr logger); @@ -178,6 +184,8 @@ class TORCH_API Reducer { std::vector>> hooks_; + // Whether we need to run autograd hooks (only false if user runs with + // no_grad or no_sync context manager) bool expect_autograd_hooks_; bool require_finalize_; size_t next_bucket_; @@ -365,7 +373,13 @@ class TORCH_API Reducer { std::vector variable_locators_; // track the number of iterations to synchronize grads in training so far. + // This is the number of calls to the forward pass, not necessarily equal to + // number of calls to backward pass. long num_iterations_; + // Number of times backward() has been called. This is mainly used for static + // graph training to know when to populate the map of how many times grad + // hooks have been triggered. + long num_backward_calls_; // track the number of buckets that have been ready for // communication calls like allReduce or communication hooks. int num_buckets_ready_; @@ -392,7 +406,13 @@ class TORCH_API Reducer { bool is_multi_device_module_ = false; // Following variables are to help build dynamic bucket order + // Whether the process of rebuilding buckets has occured. bool has_rebuilt_bucket_; + // Flag indicating all rebuilt param indices have been pushed. This is needed + // because there can be multiple calls to backward with retain_graph=True + // without a forward that actually rebuilds the buckets. In this case, we use + // this flag to avoid pushing parameters multiple times. + bool all_rebuilt_params_pushed_{false}; std::vector rebuilt_params_; std::vector rebuilt_param_indices_; const int64_t bucket_bytes_cap_; @@ -457,8 +477,7 @@ class TORCH_API Reducer { // get current cuda stream const c10::Stream get_current_stream(); bool dynamic_graph_find_unused(); - bool static_graph_first_iteration(); - bool static_graph_after_first_iteration(); + bool static_graph_after_first_bwd(); // comm_hook_ is used to access the DDP communication hook if registered. std::unique_ptr comm_hook_; diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py index e69dcc9006ac1..0f99de8d99c7f 100644 --- a/torch/nn/parallel/distributed.py +++ b/torch/nn/parallel/distributed.py @@ -142,23 +142,35 @@ def forward(ctx, reducer, state_dict, *inputs): @staticmethod def backward(ctx, *grad_outputs): state_dict = ctx.state_dict + + grad_enabled = state_dict['grad_enabled'] + require_backward_grad_sync = state_dict['require_backward_grad_sync'] static_graph_training = ctx.state_dict['static_graph'] - if static_graph_training and ctx.state_dict['num_iterations'] == 1: + if grad_enabled and require_backward_grad_sync: + if static_graph_training or not state_dict['find_unused']: + ctx.reducer.prepare_for_backward([]) + else: + # First type of unused params: parameters that did not participate + # in computing model outputs. These are found by the below call to + # prepare_for_backward. + # Second type of unused params: params that won't get gradient + # because outputs they produced do not get used in computing loss + # for this call to backward. Due to this passthrough autograd + # function, autograd hooks for these parameters are now triggered + # with undefined gradient to maintain parity with local training. + # DDP takes care of undefined grads in this case to ensure the .grad + # field of the param is not touched. + ctx.reducer.prepare_for_backward(list(_find_tensors(ctx.inputs))) + + # Note that we enqueue delay allreduce after prepare_for_backward in + # static graph training as prepare_for_backward sets the + # num_backwards_call counter in the reducer. + static_graph_first_bwd = ( + static_graph_training and ctx.reducer._static_graph_first_bwd() + ) + if static_graph_first_bwd: Variable._execution_engine.queue_callback(ctx.reducer._delay_all_reduce) - elif state_dict['find_unused'] and not static_graph_training: - # First type of unused params: parameters that did not participate - # in computing model outputs. These are found by the below call to - # prepare_for_backward. - # Second type of unused params: params that won't get gradient - # because outputs they produced do not get used in computing loss - # for this call to backward. Due to this passthrough autograd - # function, autograd hooks for these parameters are now triggered - # with undefined gradient to maintain parity with local training. - # DDP takes care of undefined grads in this case to ensure the .grad - # field of the param is not touched. - ctx.reducer.prepare_for_backward(list(_find_tensors(ctx.inputs))) - return (None, None, *grad_outputs) class DistributedDataParallel(Module): @@ -573,7 +585,6 @@ def _ddp_init_helper(self, parameters, expect_sparse_gradient, param_to_name_map (4) Logging constructin-time DDP logging data (5) passing a handle of DDP to SyncBatchNorm Layer """ - self.num_iterations = 0 # The bucket size limit is specified in the constructor. # Additionally, we allow for a single small bucket for parameters # that are defined first, such that their gradients don't spill into @@ -807,10 +818,10 @@ def forward(self, *inputs, **kwargs): with torch.autograd.profiler.record_function("DistributedDataParallel.forward"): self.reducer.save_thread_local_state() grad_enabled = torch.is_grad_enabled() - if grad_enabled and self.require_backward_grad_sync: + will_run_grad_reduction = grad_enabled and self.require_backward_grad_sync + if will_run_grad_reduction: self.logger.set_runtime_stats_and_log() - self.num_iterations += 1 - self.reducer.prepare_for_forward() + self.reducer.prepare_for_forward(will_run_grad_reduction) if self.ddp_uneven_inputs_config.ddp_join_enabled: ones = torch.ones(1, device=self.device) work = dist.all_reduce(ones, group=self.process_group, async_op=True) @@ -856,33 +867,29 @@ def forward(self, *inputs, **kwargs): else: output = self.module(*inputs, **kwargs) - if grad_enabled and self.require_backward_grad_sync: - self.require_forward_param_sync = True - if self.static_graph or not self.find_unused_parameters: - self.reducer.prepare_for_backward([]) - else: - self.require_forward_param_sync = False + self.require_forward_param_sync = ( + grad_enabled and self.require_backward_grad_sync + ) + + if not grad_enabled: + # Don't need to run through DDPSink as there will be no backward + # pass. + return output - # TODO: DDPSink is currently enabled for unused parameter detection and - # static graph training for first iteration, in the future we plan to - # enable this passthrough for all training use cases. - if (self.find_unused_parameters and not self.static_graph) or ( - self.static_graph and self.num_iterations == 1 - ): - find_unused = all([ - grad_enabled, - self.require_backward_grad_sync, - self.find_unused_parameters, - ]) state_dict = { 'static_graph': self.static_graph, - 'find_unused': find_unused, - 'num_iterations': self.num_iterations, + 'find_unused': self.find_unused_parameters, + 'grad_enabled': grad_enabled, + 'require_backward_grad_sync': self.require_backward_grad_sync, } - output_tensor_list, treespec, output_is_rref = _tree_flatten_with_rref( output ) + # Note: DDPSink helps to ensure that prepare_for_backward is called + # immediately before the backwards pass, to support a variety of + # features such as: enqueue delay allreduce for static graph, support + # multiple calls to backwards with retain_graph=True, and support + # finding all parameters that will not receive gradient. passthrough_tensor_list = _DDPSink.apply( self.reducer, state_dict, @@ -892,7 +899,7 @@ def forward(self, *inputs, **kwargs): output = _tree_unflatten_with_rref( passthrough_tensor_list, treespec, output_is_rref ) - return output + return output def scatter(self, inputs, kwargs, device_ids): return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim) diff --git a/torch/testing/_internal/distributed/distributed_test.py b/torch/testing/_internal/distributed/distributed_test.py index 79ac138095c41..c41580f701454 100644 --- a/torch/testing/_internal/distributed/distributed_test.py +++ b/torch/testing/_internal/distributed/distributed_test.py @@ -6982,8 +6982,8 @@ def test_ddp_sync_bn_training_vs_eval(self): for i in range(6): inp = torch.randn(10, 2, 4, 4).cuda(rank) out = model_inference(inp) - loss = out.sum() - loss.backward() + # Do not need to run backward as we are testing only + # inference mode here. # Ensure sync does not occur in eval() mode. if BACKEND == "nccl": @@ -7091,6 +7091,9 @@ def get_loss(model_output): loss = get_loss(out) loss.backward() self._model_step(model) + # Test non 1:1 calls to fwd/backward to ensure + # https://github.com/pytorch/pytorch/issues/58111 is fixed. + model_static_graph(inp, output_type=output_type) out_static = model_static_graph(inp, output_type=output_type) self.assertTrue(isinstance(out_static, type_mapping[output_type])) loss_static = get_loss(out_static) @@ -7100,3 +7103,140 @@ def get_loss(model_output): model.parameters(), model_static_graph.parameters() ): self.assertEqual(p, p_static) + + def _verify_ddp_model(self, ddp_model, local_model): + # Verify weights are appropriately synchronized. + all_params = [None for _ in range(dist.get_world_size())] + dist.all_gather_object(all_params, list(ddp_model.parameters())) + rank_0_params = all_params[0] + for param_list in all_params[1:]: + for i, p in enumerate(param_list): + rank_0_param = rank_0_params[i] + self.assertTrue(torch.equal(rank_0_param.data.cpu(), p.data.cpu())) + if self.rank == 0: + local_params = list(local_model.parameters()) + for dist_param, local_param in zip(rank_0_params, local_params): + self.assertTrue(torch.equal(dist_param.data.cpu(), local_param.data.cpu())) + + def _test_ddp_bwd_with_retain_graph(self, static_graph, find_unused_parameters): + # Ensures that calling backward multiple times with retain_graph=True + # is supported in DDP and verifies parity with local training. + class ToyModel(nn.Module): + def __init__(self): + super(ToyModel, self).__init__() + self.net1 = nn.Linear(10, 10, bias=False) + + def forward(self, x): + return self.net1(x) + + rank = self.rank + torch.cuda.set_device(rank) + model = ToyModel().cuda(torch.cuda.current_device()) + local_model = copy.deepcopy(model) + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[torch.cuda.current_device()], + find_unused_parameters=find_unused_parameters, + ) + if static_graph: + ddp_model._set_static_graph() + + # Run multiple backwards for DDP and local model. + inp = torch.randn(20, 10, device=rank) + for _ in range(3): + loss = ddp_model(inp).sum() + loss.backward(retain_graph=True) + loss.backward(retain_graph=True) + loss.backward() + + for _ in range(3): + local_loss = local_model(inp).sum() + local_loss.backward(retain_graph=True) + local_loss.backward(retain_graph=True) + local_loss.backward() + + # Run additional forward/backward steps to ensure that things like + # rebuild_buckets work appropriately. + for _ in range(3): + loss = ddp_model(inp).sum() + local_loss = local_model(inp).sum() + loss.backward() + local_loss.backward() + # Compare models. + self._verify_ddp_model(ddp_model, local_model) + + @skip_if_lt_x_gpu(2) + @unittest.skipIf( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_ddp_bwd_with_retain_graph(self): + self._test_ddp_bwd_with_retain_graph( + static_graph=False, find_unused_parameters=False + ) + + @skip_if_lt_x_gpu(2) + @unittest.skipIf( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_ddp_bwd_with_retain_graph_find_unused_params(self): + self._test_ddp_bwd_with_retain_graph( + static_graph=False, find_unused_parameters=True + ) + + @skip_if_lt_x_gpu(2) + @unittest.skipIf( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_ddp_bwd_with_retain_graph_static(self): + self._test_ddp_bwd_with_retain_graph( + static_graph=True, find_unused_parameters=False + ) + + @skip_if_lt_x_gpu(2) + @unittest.skipIf( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_ddp_bwd_with_retain_graph_static_find_unused_params(self): + self._test_ddp_bwd_with_retain_graph( + static_graph=True, find_unused_parameters=True + ) + + @skip_if_lt_x_gpu(2) + @unittest.skipIf( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_ddp_bwd_with_retain_graph_static_single_fwd(self): + # Ensures that if we do 1 forward pass immediately followed by 2 + # backward passes, there is no issue. In particular, verifies that + # delay allreduce is enqueued only once. + rank = self.rank + model = TwoLinLayerNet().cuda(rank) + model_ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank] + ) + model_ddp._set_static_graph() + model_local = copy.deepcopy(model) + inp = torch.randn(2, 10, device=rank) + # Run single forward pass followed by 2 calls to backward with + # retain_graph=True. + out_ddp = model_ddp(inp) + out_ddp = torch.add(out_ddp[0], out_ddp[1]).sum() + out_local = model_local(inp) + out_local = torch.add(out_local[0], out_local[1]).sum() + out_ddp.backward(retain_graph=True) + out_ddp.backward() + out_local.backward(retain_graph=True) + out_local.backward() + dist_grad_tensor = torch.cat( + [param.grad for param in model_ddp.module.parameters()] + ) + local_grad_tensor = torch.cat( + [param.grad for param in model_local.parameters()] + ) + self.assertEqual(dist_grad_tensor, local_grad_tensor) From 0baad214b07ad35be1f10100168ed761cc7c51c0 Mon Sep 17 00:00:00 2001 From: Ansha Yu Date: Fri, 18 Jun 2021 09:55:12 -0700 Subject: [PATCH 236/305] [static runtime][fix] resize to the input tensor size for full_like (#60229) 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/60229 Fix bug where we did not resize to the input tensor size, causing the output to be incorrect Test Plan: Test on replayer, rebased on D29217781, with model 278203319_26. Verify with jit outputs (D28583950) `./buck-out/gen/admarket/lib/ranking/prediction_replayer/replayer --model_inference_type_target=DISAGG_ACCELERATOR --prediction_replayer_force_model_type=inline_cvr_post_imp_model --prediction_replayer_force_model=278203319_26 --prediction_replayer_target_tier=sigrid.predictor.perf.dianshi_staticruntime_debug_0604.test --prediction_replayer_input_stream_filename=/data/users/ansha/tmp/adfinder/filtered_requests_inline_cvr_100 --ignore_model_id_mismatch --check_performance --fully_remote_sr_connection_options="overall_timeout:10000000,processing_timeout:10000000" --use_new_encoding_for_ads_services --use_new_encoding_from_model_id_to_shard_id --sigrid_force_model_dir=/data/users/ansha/tmp/adfinder/278203319_26/ --sigrid_predictor_model_suffix=.predictor.disagg.local —use_new_encoding_from_model_id_to_shard_id=true --prediction_replayer_force_model_kind=19 --pytorch_predictor_static_runtime_enable=true --prediction_replayer_target_qps=1` Reviewed By: hlu1, movefast1990 Differential Revision: D29218918 fbshipit-source-id: dab4bbbabeaa8367174ed90edca43d6204c65409 --- torch/csrc/jit/runtime/static/ops.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torch/csrc/jit/runtime/static/ops.cpp b/torch/csrc/jit/runtime/static/ops.cpp index 97734aa441765..fb230bcf3158c 100644 --- a/torch/csrc/jit/runtime/static/ops.cpp +++ b/torch/csrc/jit/runtime/static/ops.cpp @@ -1662,8 +1662,8 @@ REGISTER_OPERATOR_FUNCTOR(aten::full_like, aten_full_like, [](Node* n) -> SROper } return [](ProcessedNode* p_node) { const auto in1_s = p_node->Input(1).toScalar(); + const auto& in0_t = p_node->Input(0).toTensor(); if (p_node->Output(0).isNone()) { - const auto& in0_t = p_node->Input(0).toTensor(); const auto dtype = p_node->Input(2).toOptional(); const auto layout = p_node->Input(3).toOptional(); const auto device = p_node->Input(4).toOptional(); @@ -1675,6 +1675,7 @@ REGISTER_OPERATOR_FUNCTOR(aten::full_like, aten_full_like, [](Node* n) -> SROper in0_t, dtype, layout, device, pin_memory, memory_format); } auto& out_t = p_node->Output(0).toTensor(); + at::native::resize_(out_t, in0_t.sizes(), c10::nullopt); at::native::fill_out(out_t, in1_s); }; }); From 5a45103139b52917703630b4f055df056d6f6260 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Fri, 18 Jun 2021 10:22:55 -0700 Subject: [PATCH 237/305] ns for fx: add API usage logging (#60103) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60103 Adds internal logging for NS for FX API usage. Test Plan: CI Reviewed By: jerryzh168 Differential Revision: D29166710 fbshipit-source-id: 2a1bf2f6038b0c6c5945b57b2db2de25c585a04a --- torch/quantization/_numeric_suite_fx.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/torch/quantization/_numeric_suite_fx.py b/torch/quantization/_numeric_suite_fx.py index c4cc2d4ab9b1e..7da94bf2bfc2d 100644 --- a/torch/quantization/_numeric_suite_fx.py +++ b/torch/quantization/_numeric_suite_fx.py @@ -121,6 +121,7 @@ def _extract_weights_one_model( nodes_and_names_to_instrument: List[Tuple[Node, str]], results: NSResultsType, ) -> None: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._extract_weights_one_model") base_name_to_sets_of_related_ops = get_base_name_to_sets_of_related_ops() type_a_related_to_b = \ get_type_a_related_to_b(base_name_to_sets_of_related_ops) @@ -143,6 +144,7 @@ def _extract_weights_impl( base_name_to_sets_of_related_ops: Optional[Dict[str, Set[NSNodeTargetType]]] = None, unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> NSResultsType: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._extract_weights_impl") matched_subgraph_pairs = get_matching_subgraph_pairs( gm_a, gm_b, base_name_to_sets_of_related_ops, unmatchable_types_map) @@ -173,6 +175,7 @@ def extract_weights( base_name_to_sets_of_related_ops: Optional[Dict[str, Set[NSNodeTargetType]]] = None, unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> NSResultsType: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx.extract_weights") base_name_to_sets_of_related_ops = get_base_name_to_sets_of_related_ops() type_a_related_to_b = \ get_type_a_related_to_b(base_name_to_sets_of_related_ops) @@ -196,6 +199,7 @@ def _add_loggers_one_model( nodes_and_names_to_instrument_outputs: List[Tuple[Node, str]], logger_cls: Callable, ) -> nn.Module: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._add_loggers_one_model") # TODO(future PR): do not observe nodes we do not care # about (both fp32, denylist, etc) @@ -222,6 +226,7 @@ def _add_loggers_impl( base_name_to_sets_of_related_ops: Optional[Dict[str, Set[NSNodeTargetType]]] = None, unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> Tuple[nn.Module, nn.Module]: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._add_loggers_impl") matched_subgraph_pairs = get_matching_subgraph_pairs( gm_a, gm_b, base_name_to_sets_of_related_ops, unmatchable_types_map) @@ -259,6 +264,7 @@ def add_loggers( base_name_to_sets_of_related_ops: Optional[Dict[str, Set[NSNodeTargetType]]] = None, unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> Tuple[nn.Module, nn.Module]: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx.add_loggers") # TODO(future PR): expose these skipped_module_names: List[str] = [] skipped_module_classes: List[Callable] = [] @@ -278,6 +284,7 @@ def _extract_logger_info_one_model( results: NSResultsType, logger_cls: Callable, ) -> None: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._extract_logger_info_one_model") for gm_name, mod in model.named_modules(): # TODO(future PR): better check when scripted is_logger = ( @@ -331,6 +338,7 @@ def extract_logger_info( Output format: NSResultsType """ + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx.extract_logger_info") results: NSResultsType = {} for model in (model_a, model_b): _extract_logger_info_one_model(model, results, logger_cls) @@ -348,6 +356,7 @@ def _add_shadow_loggers_impl( node_type_to_io_type_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, unmatchable_types_map: Optional[Dict[str, Set[NSNodeTargetType]]] = None, ) -> nn.Module: + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx._add_shadow_loggers_impl") matched_subgraph_pairs = get_matching_subgraph_pairs( gm_a, gm_b, base_name_to_sets_of_related_ops, unmatchable_types_map) @@ -373,6 +382,7 @@ def add_shadow_loggers( Same thing as add_loggers, but for an `a_shadows_b` model. TODO(future PR): real docblock """ + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx.add_shadow_loggers") # TODO(future PR): expose these skipped_module_names: List[str] = [] skipped_module_classes: List[Callable] = [] @@ -396,6 +406,7 @@ def extract_shadow_logger_info( Same thing as extract_logger_info, but for an `a_shadows_b` model. TODO(future PR): real docblock """ + torch._C._log_api_usage_once("quantization_api._numeric_suite_fx.extract_shadow_logger_info") results: NSResultsType = collections.defaultdict(dict) _extract_logger_info_one_model(model_a_shadows_b, results, logger_cls) return dict(results) From 5f010c066fcb6964285ee4b881b1c211376ef811 Mon Sep 17 00:00:00 2001 From: Zhengxu Chen Date: Fri, 18 Jun 2021 11:10:32 -0700 Subject: [PATCH 238/305] [package] Bring back save_source_file (#59962) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59962 This reverts commit 44b021d21b5681c105529881bdbaefb6d3e335f6. Test Plan: Imported from OSS Reviewed By: H-Huang Differential Revision: D29113224 Pulled By: zhxchen17 fbshipit-source-id: 55d42acc421c5f4abbbad9d9ed4d32b615939463 --- test/package/test_save_load.py | 15 +++++++++ torch/package/package_exporter.py | 54 +++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/test/package/test_save_load.py b/test/package/test_save_load.py index 729f35e8f8d12..e50b42fa5e48d 100644 --- a/test/package/test_save_load.py +++ b/test/package/test_save_load.py @@ -20,6 +20,21 @@ class TestSaveLoad(PackageTestCase): """Core save_* and loading API tests.""" + @skipIf( + IS_FBCODE or IS_SANDCASTLE, + "Tests that use temporary files are disabled in fbcode", + ) + def test_saving_source(self): + filename = self.temp() + with PackageExporter(filename, verbose=False) as he: + he.save_source_file("foo", str(packaging_directory / "module_a.py")) + he.save_source_file("foodir", str(packaging_directory / "package_a")) + hi = PackageImporter(filename) + foo = hi.import_module("foo") + s = hi.import_module("foodir.subpackage") + self.assertEqual(foo.result, "module_a") + self.assertEqual(s.result, "package_a.subpackage") + @skipIf( IS_FBCODE or IS_SANDCASTLE, "Tests that use temporary files are disabled in fbcode", diff --git a/torch/package/package_exporter.py b/torch/package/package_exporter.py index 9444fae168fd6..d3640ba9f43d7 100644 --- a/torch/package/package_exporter.py +++ b/torch/package/package_exporter.py @@ -223,6 +223,60 @@ def __init__( self.patterns: Dict[GlobGroup, _PatternInfo] = {} self._unique_id = 0 + def save_source_file( + self, module_name: str, file_or_directory: str, dependencies=True + ): + """Adds the local file system ``file_or_directory`` to the source package to provide the code + for ``module_name``. + + Args: + module_name (str): e.g. ``"my_package.my_subpackage"``, code will be saved to provide code for this package. + file_or_directory (str): the path to a file or directory of code. When a directory, all python files in the directory + are recursively copied using :meth:`save_source_file`. If a file is named ``"/__init__.py"`` the code is treated + as a package. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. + """ + path = Path(file_or_directory) + if path.is_dir(): + to_save = [] # list of tuples with arguments to save_source_string + module_path = module_name.replace(".", "/") + for filename in path.glob("**/*.py"): + relative_path = filename.relative_to(path).as_posix() + archivename = module_path + "/" + relative_path + submodule_name = None + if filename.name == "__init__.py": + submodule_name = archivename[: -len("/__init__.py")].replace( + "/", "." + ) + is_package = True + else: + submodule_name = archivename[: -len(".py")].replace("/", ".") + is_package = False + + # we delay the call to save_source_string so that we record all the source files + # being provided by this directory structure _before_ attempting to resolve the dependencies + # on the source. This makes sure we don't try to copy over modules that will just get + # overwritten by this directory blob + to_save.append( + ( + submodule_name, + _read_file(str(filename)), + is_package, + dependencies, + ) + ) + + for item in to_save: + self.save_source_string(*item) + else: + is_package = path.name == "__init__.py" + self.save_source_string( + module_name, + _read_file(file_or_directory), + is_package, + dependencies, + ) + def get_unique_id(self) -> str: """Get an id. This id is guaranteed to only be handed out once for this package.""" ret = str(self._unique_id) From 3870e68644ce545a42bf939627f816cd7534dfab Mon Sep 17 00:00:00 2001 From: Eddie Yan Date: Fri, 18 Jun 2021 11:40:06 -0700 Subject: [PATCH 239/305] TF32 threshold twiddling for tests (#60209) Summary: Following https://github.com/pytorch/pytorch/issues/59624 I observed some straggling failing tests on Ampere due to TF32 thresholds. This PR just twiddles some more thresholds to fix the (6) failing tests I saw on A100. CC Flamefire ptrblck ngimel Pull Request resolved: https://github.com/pytorch/pytorch/pull/60209 Reviewed By: gchanan Differential Revision: D29220508 Pulled By: ngimel fbshipit-source-id: 7c83187a246e1b3a24b181334117c0ccf2baf311 --- test/test_nn.py | 2 +- torch/testing/_internal/common_nn.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/test/test_nn.py b/test/test_nn.py index 3dde5053ca282..dfeeb8fe3ab52 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -5652,7 +5652,7 @@ def test_Conv2d_depthwise_naive_groups_cuda(self, dtype=torch.float): @unittest.skipIf(not TEST_CUDA, 'CUDA not available') @repeat_test_for_types(ALL_TENSORTYPES) - @tf32_on_and_off(0.001) + @tf32_on_and_off(0.005) def test_Conv3d_depthwise_naive_groups_cuda(self, dtype=torch.float): for depth_multiplier in [1, 2]: m = nn.Conv3d(2, 2 * depth_multiplier, kernel_size=3, groups=2).to("cuda", dtype) diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py index a6a8045b1148f..8e56c4537d360 100644 --- a/torch/testing/_internal/common_nn.py +++ b/torch/testing/_internal/common_nn.py @@ -1862,6 +1862,7 @@ def fractional_max_pool3d_test(test_case): cudnn=True, check_with_long_tensor=True, with_tf32=True, + tf32_precision=0.005, ), dict( module_name='Conv2d', @@ -1925,6 +1926,7 @@ def fractional_max_pool3d_test(test_case): cudnn=True, check_with_long_tensor=True, with_tf32=True, + tf32_precision=0.005, ), dict( fullname='Conv2d_groups_thnn', @@ -1933,6 +1935,7 @@ def fractional_max_pool3d_test(test_case): input_size=(2, 4, 6, 5), check_with_long_tensor=True, with_tf32=True, + tf32_precision=0.005, ), dict( fullname='Conv2d_pad_valid', @@ -1950,7 +1953,7 @@ def fractional_max_pool3d_test(test_case): input_size=(2, 2, 6, 5), cudnn=True, with_tf32=True, - tf32_precision=0.005, + tf32_precision=0.01, ), dict( fullname='Conv2d_pad_same_dilated', From f042455a8d50ca83eb657e29d854df69034011cb Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Fri, 18 Jun 2021 11:53:39 -0700 Subject: [PATCH 240/305] [JIT] ShapeProp: add missing ops from mobilenet v3. (#59163) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59163 Test Plan: Imported from OSS Reviewed By: ejguan Differential Revision: D28853833 Pulled By: ZolotukhinM fbshipit-source-id: 451fb9ee848968049d26fb5623a904d8fa7bd6fc --- .../jit/runtime/symbolic_shape_registry.cpp | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/torch/csrc/jit/runtime/symbolic_shape_registry.cpp b/torch/csrc/jit/runtime/symbolic_shape_registry.cpp index a4ca941a90242..fa65872e4abe2 100644 --- a/torch/csrc/jit/runtime/symbolic_shape_registry.cpp +++ b/torch/csrc/jit/runtime/symbolic_shape_registry.cpp @@ -49,6 +49,43 @@ const std::string shape_compute_functions = out.append(elem) return out + def unary_one_unused_input(self: List[int], inp0: Any): + out: List[int] = [] + for elem in self: + out.append(elem) + return out + + def unary(self: List[int]): + out: List[int] = [] + for elem in self: + out.append(elem) + return out + + def view(self: List[int], sizes: List[int]): + # TODO: add assertions to check whether requested dims are valid + out: List[int] = [] + for elem in sizes: + if elem == -1: + # TODO: support -1 in view dimensions + raise AssertionError("Shape function doesn't support -1 view dims yet") + out.append(elem) + return out + + def mean_dim(self: List[int], dims: List[int], keep_dim: bool, dt : Any): + out: List[int] = [] + idx : int = 0 + for elem in self: + is_mean_dim : bool = False + for reduce_dim in dims: + if idx == reduce_dim: + is_mean_dim = True + if is_mean_dim: + if keep_dim: + out.append(1) + else: + out.append(elem) + return out + def broadcast_one_unused_input(self: List[int], other: List[int], unused: Any): return broadcast(self, other) @@ -150,6 +187,12 @@ const std::string shape_compute_functions = assert broadcast(bias, out) == out return out + def addmm(self: List[int], mat1: List[int], mat2: List[int], beta: Any, alpha: Any): + out = matmul(mat1, t(mat2)) + if self is not None: + assert broadcast(self, out) == out + return out + def check_non_negative(array: List[int]) -> bool: for val in array: if val < 0: @@ -264,9 +307,12 @@ static const OperatorMap& get_schema_to_function_graph() { // clang-format off static const OperatorMap schema_to_function_graph{ {"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor", "broadcast"}, + {"aten::mul.Scalar(Tensor self, Scalar other) -> Tensor", "unary_one_unused_input"}, {"aten::div.Tensor(Tensor self, Tensor other) -> Tensor", "broadcast"}, + {"aten::div.Scalar(Tensor self, Scalar other) -> Tensor", "unary_one_unused_input"}, {"aten::gt.Tensor(Tensor self, Tensor other) -> Tensor", "broadcast"}, {"aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor", "broadcast_one_unused_input"}, + {"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor", "unary_two_unused_inputs"}, {"aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor", "unary_two_unused_inputs"}, {"aten::adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor", "adaptive_avg_pool2d"}, {"aten::mm(Tensor self, Tensor mat2) -> Tensor", "mm"}, @@ -279,6 +325,11 @@ static const OperatorMap& get_schema_to_function_graph() { {"aten::conv2d(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] dilation=1, int groups=1) -> Tensor", "conv2d"}, {"aten::conv3d(Tensor input, Tensor weight, Tensor? bias=None, int[3] stride=1, int[3] padding=0, int[3] dilation=1, int groups=1) -> Tensor", "conv3d"}, {"aten::flatten.using_ints(Tensor(a) self, int start_dim=0, int end_dim=-1) -> Tensor(a)", "flatten"}, + {"aten::relu(Tensor self) -> Tensor", "unary"}, + {"aten::view(Tensor(a) self, int[] size) -> Tensor(a)", "view"}, + {"aten::expand_as(Tensor(a) self, Tensor other) -> Tensor(a)", "view"}, + {"aten::mean.dim(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor", "mean_dim"}, + {"aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor", "addmm"}, }; // clang-format on return schema_to_function_graph; From c6bb9409b8511cddf18c79753b0d68eb8875595e Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Fri, 18 Jun 2021 11:53:39 -0700 Subject: [PATCH 241/305] [TensorExpr] Handle not-specified dtypes and strides. (#59346) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59346 Currently JIT has a pass to propagate shapes, but doesn't have a capability to fill in strides and dtypes. This PR works around that by assuming default dtype to be Float and strides corresponding to contiguous layout, unless otherwise specified. Ideally, we won't need this, and this is done simply as a workaround unless the corresponding features are implemented on JIT side. This is required for AOT compilation of mobilenet v3 with NNC. Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D28853831 Pulled By: ZolotukhinM fbshipit-source-id: 81adb59409684f39b444909ab8ec58ee4a39d496 --- torch/csrc/jit/tensorexpr/kernel.cpp | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 1462e82540d29..761c226037e18 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -182,20 +182,26 @@ c10::optional pickDeviceType(const std::shared_ptr& graph) { // nullopt. c10::optional getTensorInfoJit(torch::jit::Value* v) { auto const& it = v->type()->cast(); + + c10::ScalarType dtype = c10::ScalarType::Float; + if (!it) { return c10::nullopt; } if (!it->isComplete()) { return c10::nullopt; } - if (!it->scalarType()) { - return c10::nullopt; + if (it->scalarType()) { + // TODO: ideally we should be strict here and return nullopt if the dtype is + // absent in the JIT IR. We're assuming a default Float dtype for now, until + // dtype propagation is implemented. + dtype = *it->scalarType(); } auto concrete_sizes = it->sizes().concrete_sizes(); if (!concrete_sizes) { return c10::nullopt; } - return TensorInfo{*concrete_sizes, *it->scalarType()}; + return TensorInfo{*concrete_sizes, dtype}; } c10::optional getTensorInfo(BufHandle b) { std::vector dims; @@ -641,7 +647,7 @@ std::vector TensorExprKernel::sizesForValue( // need to infer it. if (v->type()->kind() == TypeKind::TensorType) { auto tt = v->type()->cast(); - if (tt->isComplete()) { + if (tt->sizes().concrete_sizes()) { return sizesFromVaryingShape(tt->sizes()); } } @@ -2883,6 +2889,9 @@ Tensor* TensorExprKernel::convertOutputToCorrectStrides(torch::jit::Value* v) { TORCH_INTERNAL_ASSERT(tt->sizes().concrete_sizes()); const auto sizes = *tt->sizes().concrete_sizes(); std::vector default_strides = TensorType::contiguousStridesOf(sizes); + if (!tt->strides().concrete_sizes()) { + return new Tensor(buf, nullptr); + } TORCH_INTERNAL_ASSERT(tt->strides().concrete_sizes()); const std::vector strides = *tt->strides().concrete_sizes(); // All Tensors in NNC are layed out in default, contiguous layout. @@ -3031,12 +3040,12 @@ void TensorExprKernel::compile() { const auto& tt = output->type()->expect(); auto sizes = *tt->sizes().concrete_sizes(); tensorOutputSizes_.push_back(sizes); - auto strides = *tt->strides().concrete_sizes(); + auto strides = tt->strides().concrete_sizes(); // If the tensor is not dense or overlaps, we have // no way of matching the profiled striding - if (denseAndNonOverlapping(sizes, strides)) { - tensorOutputStrides_.push_back(*tt->strides().concrete_sizes()); + if (strides && denseAndNonOverlapping(sizes, *strides)) { + tensorOutputStrides_.push_back(*strides); } else { tensorOutputStrides_.push_back(TensorType::contiguousStridesOf(sizes)); } From d9e7df707bdc737a10b3af15f08143cc300ace39 Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Fri, 18 Jun 2021 11:53:39 -0700 Subject: [PATCH 242/305] [TensorExpr] Add NNC lowerings for `aten::mean`, `aten::addmm`, and `aten::adaptive_avg_pool2d`. (#59347) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59347 We had external call wrappers for them, but they were not used in NNC. This PR adds lowerings using these ext calls and fixes some bugs in them. Test Plan: Imported from OSS Reviewed By: jbschlosser Differential Revision: D28853832 Pulled By: ZolotukhinM fbshipit-source-id: 1718400368e1a9cf3f19180ee2290a4ed9c99d41 --- test/test_jit_fuser_te.py | 23 +++++++++- .../jit/tensorexpr/external_functions.cpp | 17 +++++-- torch/csrc/jit/tensorexpr/kernel.cpp | 12 +++++ .../csrc/jit/tensorexpr/operators/matmul.cpp | 22 ++++++++++ torch/csrc/jit/tensorexpr/operators/matmul.h | 4 ++ .../jit/tensorexpr/operators/reduction.cpp | 44 +++++++++++++++++++ .../csrc/jit/tensorexpr/operators/reduction.h | 8 ++++ 7 files changed, 124 insertions(+), 6 deletions(-) diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py index 2521ff95350e1..21e8d306d158a 100644 --- a/test/test_jit_fuser_te.py +++ b/test/test_jit_fuser_te.py @@ -1880,6 +1880,23 @@ def eager(x): fusion_groups = self.findFusionGroups(trace.graph_for(x)) self.assertEqual(len(fusion_groups), 0) + def test_adaptive_avg_pool2d(self): + # TODO: once the adaptive_avg_pool2d is available in OpInfo DB, this + # test should be moved there + with inline_fusion_groups(): + def foo1(x): + return torch.nn.functional.adaptive_avg_pool2d(x, (2, 2)) + + def foo2(x): + return torch.nn.functional.adaptive_avg_pool2d(x, (2)) + + x = torch.randn(4, 4, 4) + for foo in [foo1, foo2]: + f = torch.jit.trace(foo, (x,)) + kernel = torch._C._te.TensorExprKernel(f.graph) + correct_val = f(x) + self.assertEqual(kernel.run((x,)), correct_val) + works_list = [ '__radd__', @@ -1890,6 +1907,7 @@ def eager(x): 'acos', 'add', 'addcmul', + 'addmm.decomposed', 'asin', 'atan', 'atan2', @@ -1922,6 +1940,7 @@ def eager(x): 'lt', 'masked_fill', 'max.binary', + 'mean', 'min.binary', 'mm', 'mul', @@ -1959,9 +1978,9 @@ def eager(x): ] known_failures = [ - 'matmul', - 'frac', '__rmatmul__' + 'frac', + 'matmul', ] # If your OpInfo test causes this test to fail, add it here diff --git a/torch/csrc/jit/tensorexpr/external_functions.cpp b/torch/csrc/jit/tensorexpr/external_functions.cpp index e0fe18a979cfe..833916fe5738d 100644 --- a/torch/csrc/jit/tensorexpr/external_functions.cpp +++ b/torch/csrc/jit/tensorexpr/external_functions.cpp @@ -110,7 +110,10 @@ void nnc_aten_adaptive_avg_pool2d( at::Tensor& r = tensors[0]; const at::Tensor& x = tensors[1]; int64_t H = extra_args[0]; - int64_t W = extra_args[1]; + int64_t W = H; + if (args_num > 1) { + W = extra_args[1]; + } try { at::adaptive_avg_pool2d_out(r, x, {H, W}); } catch (...) { @@ -130,9 +133,12 @@ void nnc_aten_mean( at::Tensor& r = tensors[0]; const at::Tensor& x = tensors[1]; - int64_t dim = extra_args[0]; + std::vector mean_dims(args_num); + if (args_num > 0) { + memcpy(mean_dims.data(), extra_args, sizeof(int64_t) * args_num); + } try { - at::mean_out(r, x, {dim}); + at::mean_out(r, x, mean_dims); } catch (...) { } } @@ -152,8 +158,11 @@ void nnc_aten_addmm( const at::Tensor& x = tensors[1]; const at::Tensor& y = tensors[2]; const at::Tensor& z = tensors[3]; + // TODO: handle other alpha and beta dtypes, e.g. alpha=0.6, beta=0.2 + int64_t alpha = extra_args[0], beta = extra_args[1]; + try { - at::addmm_out(r, x, y, z, extra_args[0], extra_args[1]); + at::addmm_out(r, x, y, z, alpha, beta); } catch (...) { } } diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index 761c226037e18..d716c8320f369 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -2364,6 +2364,15 @@ Tensor* tensorexpr::computeOperandValue( case aten::conv2d: { return computeConv2d(inputs, outputShape, outputType); } break; + case aten::addmm: { + return computeAddMM(inputs, outputShape, outputType); + } break; + case aten::mean: { + return computeMean(inputs, outputShape, outputType); + } break; + case aten::adaptive_avg_pool2d: { + return computeAdaptiveAvgPool2d(inputs, outputShape, outputType); + } break; default: { std::string msg = std::string("Unhandled node kind (in computeOperandValue): ") + @@ -2473,6 +2482,9 @@ Tensor* TensorExprKernel::computeValue(const torch::jit::Value* v) { case aten::softmax: case aten::log_softmax: case aten::conv2d: + case aten::addmm: + case aten::mean: + case aten::adaptive_avg_pool2d: case aten::to: { std::vector argInputs; if (v->node()->kind() != aten::to) { diff --git a/torch/csrc/jit/tensorexpr/operators/matmul.cpp b/torch/csrc/jit/tensorexpr/operators/matmul.cpp index 30c34978a066f..3fdd086826222 100644 --- a/torch/csrc/jit/tensorexpr/operators/matmul.cpp +++ b/torch/csrc/jit/tensorexpr/operators/matmul.cpp @@ -50,6 +50,28 @@ Tensor* computeMatmul( } } +Tensor* computeAddMM( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType) { + Dtype dtype = kFloat; + if (outputType) { + dtype = Dtype(*outputType); + } + BufHandle ResultBuf("addmm", outputShape, dtype); + return new Tensor( + ResultBuf.node(), + ExternalCall::make( + ResultBuf, + "nnc_aten_addmm", + {c10::get(inputs[0]), + c10::get(inputs[1]), + c10::get(inputs[2])}, + {c10::get(inputs[3]), + c10::get( + inputs[4])})); // TODO: handle other dtypes of alpha and beta +} + } // namespace tensorexpr } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/matmul.h b/torch/csrc/jit/tensorexpr/operators/matmul.h index 893b3a9c820de..35b30f4168914 100644 --- a/torch/csrc/jit/tensorexpr/operators/matmul.h +++ b/torch/csrc/jit/tensorexpr/operators/matmul.h @@ -10,6 +10,10 @@ Tensor* computeMatmul( const std::vector& inputs, const std::vector& outputShape, const c10::optional& outputType); +Tensor* computeAddMM( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType); } // namespace tensorexpr } // namespace jit diff --git a/torch/csrc/jit/tensorexpr/operators/reduction.cpp b/torch/csrc/jit/tensorexpr/operators/reduction.cpp index 2f78298a060e9..1b28698aba41c 100644 --- a/torch/csrc/jit/tensorexpr/operators/reduction.cpp +++ b/torch/csrc/jit/tensorexpr/operators/reduction.cpp @@ -100,6 +100,50 @@ Tensor* computeSum( reductionDims); } +Tensor* computeMean( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType) { + Dtype dtype = kFloat; + if (outputType) { + dtype = Dtype(*outputType); + } + BufHandle ResultBuf("mean", outputShape, dtype); + BufHandle InputBuf = c10::get(inputs[0]); + std::vector mean_dims_expr; + if (auto mean_dims = c10::get_if(&inputs[1])) { + mean_dims_expr = c10::fmap(*mean_dims); + } else { + // When dims argument is not specified, reduce over all dimensions + for (int64_t idx = 0; idx < InputBuf.ndim(); idx++) { + mean_dims_expr.push_back(idx); + } + } + return new Tensor( + ResultBuf.node(), + ExternalCall::make( + ResultBuf, "nnc_aten_mean", {InputBuf}, mean_dims_expr)); +} + +Tensor* computeAdaptiveAvgPool2d( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType) { + Dtype dtype = kFloat; + if (outputType) { + dtype = Dtype(*outputType); + } + BufHandle ResultBuf("adaptive_avgpool2d", outputShape, dtype); + auto out_size_param = c10::get(inputs[1]); + return new Tensor( + ResultBuf.node(), + ExternalCall::make( + ResultBuf, + "nnc_aten_adaptive_avg_pool2d", + {c10::get(inputs[0])}, + c10::fmap(out_size_param))); +} + } // namespace tensorexpr } // namespace jit } // namespace torch diff --git a/torch/csrc/jit/tensorexpr/operators/reduction.h b/torch/csrc/jit/tensorexpr/operators/reduction.h index aeb4cd35b8765..29f051f323b28 100644 --- a/torch/csrc/jit/tensorexpr/operators/reduction.h +++ b/torch/csrc/jit/tensorexpr/operators/reduction.h @@ -9,6 +9,14 @@ namespace tensorexpr { Tensor* computeSum( const std::vector& inputs, const c10::optional& outputType); +Tensor* computeMean( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType); +Tensor* computeAdaptiveAvgPool2d( + const std::vector& inputs, + const std::vector& outputShape, + const c10::optional& outputType); } // namespace tensorexpr } // namespace jit From c0f8cad0f0ed273a0db6a11933c9b657db981251 Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Fri, 18 Jun 2021 12:48:04 -0700 Subject: [PATCH 243/305] Be fix shard inbalance (#60206) Summary: First step to address https://github.com/pytorch/pytorch/issues/60136 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60206 Reviewed By: janeyx99 Differential Revision: D29215237 Pulled By: walterddr fbshipit-source-id: ec25beb57366ef2eaf37878cdea391b245de9bef --- tools/print_test_stats.py | 55 ++++++++++++++++++++------------------- 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/tools/print_test_stats.py b/tools/print_test_stats.py index 20c36e0d816ed..57567c2dc5429 100755 --- a/tools/print_test_stats.py +++ b/tools/print_test_stats.py @@ -599,16 +599,13 @@ def append(self, test_case: TestCase) -> None: self.skipped_count += 1 if test_case.skipped else 0 self.errored_count += 1 if test_case.errored else 0 - def replace(self, test_case: TestCase) -> float: + def update(self, test_case: TestCase) -> None: name = test_case.name assert name in self.test_cases, f'Error: attempting to replace nonexistent test case {name}' - old_time = self.test_cases[name].time - # We don't replace anything if the old test case was not shorter. - if old_time >= test_case.time: - return 0.0 - self.total_time = self.total_time + test_case.time - old_time - self.test_cases[name] = test_case - return test_case.time - old_time + self.test_cases[name].time += test_case.time + self.test_cases[name].failed |= test_case.failed + self.test_cases[name].errored |= test_case.errored + self.test_cases[name].skipped |= test_case.skipped def print_report(self, num_longest: int = 3) -> None: sorted_tests = sorted(self.test_cases.values(), key=lambda x: x.time) @@ -632,23 +629,22 @@ def __init__(self, name: str) -> None: self.total_time = 0.0 self.test_suites: Dict[str, TestSuite] = dict() - def append(self, test_case: TestCase) -> None: - suite_name = test_case.class_name + def append(self, test_case: TestCase, test_type: str) -> None: + is_multi_test = self.name == 'test_cpp_extensions_aot' or \ + self.name == 'distributed/test_distributed_fork' or \ + self.name == 'distributed/test_distributed_spawn' or \ + self.name == 'distributed/test_c10d_gloo' or \ + self.name == 'cpp' # The caffe2 cpp tests spawn duplicate test cases as well. + if is_multi_test: + suite_name = test_case.class_name + '__' + test_type + else: + suite_name = test_case.class_name if suite_name not in self.test_suites: self.test_suites[suite_name] = TestSuite(suite_name) if test_case.name in self.test_suites[suite_name].test_cases: - # We expect duplicate tests for test_cpp_extensions_aot, distributed/test_distributed_fork, - # and distributed/test_distributed_spawn and test_c10d_gloo. - # In these cases, we store the test case that took the longest, - # as in these jobs, the duplicate tests are run in parallel. - # For other unexpected cases, we should raise a warning. - if self.name == 'test_cpp_extensions_aot' or \ - self.name == 'distributed/test_distributed_fork' or \ - self.name == 'distributed/test_distributed_spawn' or \ - self.name == 'distributed/test_c10d_gloo' or \ - self.name == 'cpp': # The caffe2 cpp tests spawn duplicate test cases as well. - time_difference = self.test_suites[suite_name].replace(test_case) - self.total_time += time_difference + if is_multi_test: + self.test_suites[suite_name].update(test_case) + self.total_time += test_case.time else: raise RuntimeWarning(f'Duplicate test case {test_case.name} in suite {suite_name} called from {self.name}') else: @@ -670,11 +666,16 @@ def parse_reports(folder: str) -> Dict[str, TestFile]: reports = glob(os.path.join(folder, '**', '*.xml'), recursive=True) tests_by_file = dict() for report in reports: - test_filename = re.sub(r'\.', '/', os.path.basename(os.path.dirname(report))) + report_path = Path(report) + # basename of the directory of test-report is the test filename + test_filename = re.sub(r'\.', '/', report_path.parent.name) + # test type is the parent directory (only applies to dist-*) + # See: CUSTOM_HANDLERS in test/run_test.py + test_type = report_path.parent.parent.name if test_filename not in tests_by_file: tests_by_file[test_filename] = TestFile(test_filename) for test_case in parse_report(report): - tests_by_file[test_filename].append(test_case) + tests_by_file[test_filename].append(test_case, test_type) return tests_by_file def build_info() -> ReportMetaMeta: @@ -761,9 +762,9 @@ def assemble_s3_object( 'cases': { name: { 'seconds': case.time, - 'status': 'skipped' if case.skipped else - 'errored' if case.errored else - 'failed' if case.failed else None + 'status': 'errored' if case.errored else + 'failed' if case.failed else + 'skipped' if case.skipped else None } for name, case in suite.test_cases.items() }, From a029422cae70b019222c00558da5437020550173 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Fri, 18 Jun 2021 13:28:41 -0700 Subject: [PATCH 244/305] [quant][graphmode][fx][refactor] Change the env map to add dtype as a key (#60054) 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/60054 Previously env in convert is Dict[str, Tuple[Node, torch.dtype]], that is, at a given time each node can only have one dtype, this causes a problem for the following case: ``` class M(torch.nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(1, 1, 1) def forward(self, x): x = self.conv(x) x1 = x.expand_as(x) x2 = torch.add(x, x1) return x2 def forward(self, x): x = self.activation_post_process_0(x) x = self.conv(x) x = self.activation_post_process_1(x) x1 = x.expand_as(x) x1 = self.activation_post_process_2(x1) x2 = torch.add(x, x1) x2 = self.activation_post_process_3(x2) return x2 def forward(self, x): x = torch.quantize_per_tensor(x, ...) x = self.conv(x). # quantized conv x = torch.dequantize(x) x1 = x.expand_as(x) x1 = torch.quantize_per_tensor(x1, ...) # Error: x is dequantized x2 = torch.ops.quantized.add(x, x1) return x2 Currently we have a env that is a map from node name of the observed graph to the Node in the quantized graph, here the problem is that following a quantized operator conv, we have two operators, one is expecting float input (expand_as), the other is expecting quantized input (quantized add), and in the quantized graph, ideally, expand_as should consume the dequantized output, and quantized add should consume the quantized output: quantized_conv - dequantize - expand_as \ ------- quantized_add But currently in env, each node needs to either be quantized or not quantized. Therefore we will need to change env to include dtype as well: env: Dict[str, Dict[dtype, Node]], e.g. {‘x’: {torch.float: dequantized_node, torch.quint8: quantized_node}} And when we load from the env, we will need to provide the dtype of the Node that we want to load as well. We can have a separate pass to figure out this information for each node. ``` Test Plan: python test/test_quantization.py TestQuantizeFx python test/test_quantization.py TestQuantizeFxOps Imported from OSS Reviewed By: vkuzo Differential Revision: D29149408 fbshipit-source-id: c9e4b7d65444ab6a6f573929bae1db5037629892 --- torch/quantization/fx/convert.py | 147 +++++++++++------- .../quantization/fx/quantization_patterns.py | 139 +++++++++-------- 2 files changed, 162 insertions(+), 124 deletions(-) diff --git a/torch/quantization/fx/convert.py b/torch/quantization/fx/convert.py index 75d95b1b9e57c..641e738e80746 100644 --- a/torch/quantization/fx/convert.py +++ b/torch/quantization/fx/convert.py @@ -1,4 +1,5 @@ from typing import Any, Dict, Tuple, List, Callable, Optional, Union +from collections import defaultdict import torch from torch.fx import ( GraphModule, @@ -180,7 +181,7 @@ def convert(model: GraphModule, is_reference: bool = False, custom_module_classes=custom_module_classes) quantized_graph = Graph() - env: Dict[str, Tuple[Node, Optional[torch.dtype]]] = {} + env: Dict[str, Dict[Optional[torch.dtype], Node]] = defaultdict(lambda: defaultdict(Node)) # type: ignore[arg-type] graph_inputs: List[str] = [] for node in model.graph.nodes: @@ -193,71 +194,106 @@ def load_non_quantized(n: Node) -> Node: 'node:' + n.name + \ ' in env: ' + \ str(env) - quantized_node, dtype = env[n.name] - if dtype and dtype != torch.float: - env[n.name] = Proxy(quantized_node).dequantize().node, torch.float - return env[n.name][0] - - def load_quantized(n: Node) -> Node: - assert n.name in env, \ - 'trying to load quantized node but did not find node:' + \ - n.name + ' in environment:' + str(env) - quantized_node, dtype = env[n.name] - assert dtype in [torch.quint8, torch.qint8, torch.float16], \ - f'Expecting node {quantized_node} to be quantized but got dtype: {dtype}' - return quantized_node + dtype_to_node = env[n.name] + if torch.float in dtype_to_node: + return dtype_to_node[torch.float] + elif None in dtype_to_node: + return dtype_to_node[None] + else: + quantized_node = None + for dtype in [torch.quint8, torch.qint8, torch.float16]: + if dtype in dtype_to_node: + quantized_node = dtype_to_node[dtype] + break + assert quantized_node is not None, "Did not find a supported quantized dtype:{}".format(dtype_to_node) + env[n.name][torch.float] = Proxy(quantized_node).dequantize().node + return env[n.name][torch.float] + + def load_quantized(dtype: torch.dtype): + def load_quantized_impl(n: Node): + assert n.name in env, \ + 'trying to load quantized node but did not find node:' + \ + n.name + ' in environment:' + str(env) + dtype_to_node = env[n.name] + local_dtype : Optional[torch.dtype] = dtype + if local_dtype == torch.float and local_dtype not in dtype_to_node: + local_dtype = None + if local_dtype in [torch.float, None]: + return load_non_quantized(n) + assert local_dtype in dtype_to_node, f'Expecting {dtype} in {dtype_to_node}' + return dtype_to_node[local_dtype] + + return load_quantized_impl def load_x(n: Node) -> Node: assert n.name in env, \ 'node ' + n.name + ' does not exist in environment' - return env[n.name][0] - - def load_arg(quantized: Optional[Union[List[int], bool, Tuple[int, ...]]] - ) -> Callable[[Node], Argument]: + dtype_to_node = env[n.name] + dtypes = [torch.quint8, torch.qint8, torch.float16, torch.float32, None] + for dtype in dtypes: + if dtype in dtype_to_node: + return dtype_to_node[dtype] + raise Exception(f'dtype {dtype} not found in environment: {dtype_to_node} for node {n.name}') + + def load_arg( + quantized: Optional[Union[List[int], Dict[int, torch.dtype], torch.dtype, Tuple[int, ...]]] + ) -> Callable[[Node], Argument]: """ - Input: quantized, which can be None, list, boolean or tuple + Input: quantized, which can be None, torch.dtype, list or tuple - if quantized is None, then we'll load the node as long as it exists - - if quantized is a boolean, then all args will be - quantized/not quantized - - if quantized is an empty list or tuple, then it is the same as load_arg(quantized=False) + - if quantized is a dtype, then all args will be + quantized to the specific dtype + - if quantized is an empty list or tuple, then it is the same as load_arg(quantized=torch.float) - if quantized is a list or tuple, then arg should be a list and - the args with corresponding indexes will be quantized + the args with corresponding indexes will be quantized to torch.quint8 Output: fn which takes arg_or_args, and loads them from the corresponding environment depending on the value of quantized. """ assert quantized is None or \ - isinstance(quantized, (tuple, list, bool)), type(quantized) - if isinstance(quantized, (tuple, list)) and len(quantized) == 0: + isinstance(quantized, (tuple, list, dict, torch.dtype)), type(quantized) + if isinstance(quantized, (tuple, list, dict)) and len(quantized) == 0: # empty tuple or list means nothing is quantized - quantized = False + quantized = torch.float def load_arg_impl(arg_or_args): # we'll update the format of `quantized` # to better match arg_or_args - updated_quantized: Optional[Union[List[int], bool, Tuple[int, ...]]] = quantized + updated_quantized: Optional[Union[List[int], torch.dtype, Dict[int, torch.dtype], Tuple[int, ...]]] = quantized if isinstance(quantized, (tuple, list)) and \ len(quantized) == 1 and isinstance(arg_or_args, Node): # when argument is one Node instead of tuple, we just need to check # 0 is in the quantized list - updated_quantized = 0 in quantized + if 0 in quantized: + updated_quantized = torch.quint8 if updated_quantized is None: return map_arg(arg_or_args, load_x) - if isinstance(updated_quantized, bool): + if isinstance(updated_quantized, torch.dtype): return map_arg( arg_or_args, - load_quantized if updated_quantized else load_non_quantized) + load_quantized(updated_quantized)) elif isinstance(updated_quantized, (tuple, list)): assert isinstance(arg_or_args, (tuple, list)), arg_or_args loaded_args = [] # for now, we only support quantizing positional arguments for i, a in enumerate(arg_or_args): if i in updated_quantized: - loaded_args.append(map_arg(a, load_quantized)) + # Currently it's hardcoded to torch.quint8, we can extend this + # in the future to support all quantized + # dtypes + loaded_args.append(map_arg(a, load_quantized(torch.quint8))) + else: + loaded_args.append(map_arg(a, load_non_quantized)) + return type(arg_or_args)(loaded_args) + elif isinstance(updated_quantized, dict): + loaded_args = [] + for i, a in enumerate(arg_or_args): + if i in updated_quantized: + loaded_args.append(map_arg(a, load_quantized(updated_quantized[i]))) else: loaded_args.append(map_arg(a, load_non_quantized)) return type(arg_or_args)(loaded_args) @@ -268,8 +304,8 @@ def node_arg_is_quantized(node_arg: Any) -> bool: assert node_arg.name in env, \ 'Expecting node_arg to be in the environment' if node_arg.name in env: - _, dtype = env[node_arg.name] - return dtype != torch.float + dtype_to_node = env[node_arg.name] + return any([x in dtype_to_node for x in [torch.quint8, torch.qint8, torch.float16]]) else: return False elif isinstance(node_arg, list): @@ -322,37 +358,35 @@ def insert_quantize_node(node: Node, modules: Dict[str, torch.nn.Module]) -> Non prev_node = node.args[0] if observer_module.dtype == torch.float32: # copy the observer for fp32 dtype - env[node.name] = quantized_graph.node_copy( - node, load_non_quantized), torch.float + env[node.name][torch.float] = quantized_graph.node_copy( + node, load_non_quantized) elif isinstance(prev_node, Node) and prev_node.name in env: # if previous node is already quantized, we'll just remove the # activation_post_process - _, prev_dtype = env[prev_node.name] - current_dtype = observer_module.dtype - if prev_dtype == current_dtype: - env[node.name] = env[prev_node.name] + prev_dtype_to_node: Dict[Optional[torch.dtype], Node] = env[prev_node.name] + current_dtype: Optional[torch.dtype] = observer_module.dtype # type: ignore[assignment] + if current_dtype in prev_dtype_to_node: + env[node.name][current_dtype] = prev_dtype_to_node[current_dtype] else: root_module = modules[""] assert isinstance(prev_node, Node) observer_dtype: torch.dtype = observer_module.dtype # type: ignore[assignment] - env[node.name] = ( + env[node.name][observer_dtype] = \ quantize_node( load_non_quantized(prev_node), observer_module, node, modules, quantized_graph, - node_name_to_scope, is_input=True), - observer_dtype) + node_name_to_scope, is_input=True) else: # replace activation post process with quantization ops root_module = modules[""] assert isinstance(node.args[0], Node) dtype: torch.dtype = observer_module.dtype # type: ignore[assignment] - env[node.name] = ( + env[node.name][dtype] = \ quantize_node( load_non_quantized(node.args[0]), observer_module, node, modules, quantized_graph, - node_name_to_scope, is_input=True), - dtype) + node_name_to_scope, is_input=True) # additional state to override inputs to be quantized, if specified # by the user @@ -406,9 +440,9 @@ def insert_quantize_node(node: Node, modules: Dict[str, torch.nn.Module]) -> Non quantized = is_output_quantized(node, obj, qconfig, modules) if quantized: - env[node.name] = result, activation_dtype(qconfig) + env[node.name][activation_dtype(qconfig)] = result else: - env[node.name] = result, torch.float + env[node.name][torch.float] = result continue elif root_node is not None: if qconfig is None: @@ -422,7 +456,7 @@ def insert_quantize_node(node: Node, modules: Dict[str, torch.nn.Module]) -> Non # function will not be called. result = quantized_graph.node_copy( node, load_non_quantized) - env[node.name] = result, torch.float + env[node.name][torch.float] = result continue # handle activation post process calls @@ -433,23 +467,22 @@ def insert_quantize_node(node: Node, modules: Dict[str, torch.nn.Module]) -> Non cur_placeholder_node_idx = placeholder_node_seen_cnt placeholder_node_seen_cnt += 1 if cur_placeholder_node_idx in input_quantized_idxs: - env[node.name] = \ - quantized_graph.node_copy( - node, load_non_quantized), torch.quint8 + env[node.name][torch.quint8] = quantized_graph.node_copy( + node, load_non_quantized) else: - env[node.name] = \ - quantized_graph.node_copy(node, load_non_quantized), torch.float + env[node.name][torch.float] = \ + quantized_graph.node_copy(node, load_non_quantized) else: # copy quantized or non-quantized node # get_tensor_info_node like shape works for both # quantized and non-quantized input and output a non-Tensor # (we use None for dtype currently for non-Tensors) if is_get_tensor_info_node(node): - env[node.name] = \ - quantized_graph.node_copy(node, load_x), None + env[node.name][None] = \ + quantized_graph.node_copy(node, load_x) else: - env[node.name] = \ - quantized_graph.node_copy(node, load_non_quantized), torch.float + env[node.name][torch.float] = \ + quantized_graph.node_copy(node, load_non_quantized) # remove activation post process act_post_process_removed_graph = Graph() diff --git a/torch/quantization/fx/quantization_patterns.py b/torch/quantization/fx/quantization_patterns.py index cc893c49155ad..7023b6cfdb8d7 100644 --- a/torch/quantization/fx/quantization_patterns.py +++ b/torch/quantization/fx/quantization_patterns.py @@ -290,10 +290,10 @@ def convert(self, if is_reference and self.binary_op in binary_reference_op_supported_dtypes and \ dtypes in binary_reference_op_supported_dtypes[self.binary_op]: if dtypes in binary_op_int8_dtypes: - args = load_arg(quantized=[0, 1])(node.args) - args = load_arg(quantized=False)(node.args) - kwargs = load_arg(quantized=False)(node.kwargs) - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + args = load_arg(quantized=[torch.quint8, torch.qint8])(node.args) + args = load_arg(quantized=torch.float)(node.args) + kwargs = load_arg(quantized=torch.float)(node.kwargs) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) activation_post_process = \ self._maybe_get_last_node_only_observer(modules) assert activation_post_process is not None @@ -305,7 +305,7 @@ def convert(self, "No implementation found for dtype combination: {}" "for op {} with is_reference={} despite it being listed as supported" "this should not happen".format(dtypes, self.binary_op, is_reference)) - return quantized_graph.node_copy(node, load_arg(quantized=False)) + return quantized_graph.node_copy(node, load_arg(quantized=torch.float)) elif not is_reference and self.binary_op in binary_op_supported_dtypes and \ dtypes in binary_op_supported_dtypes[self.binary_op]: if dtypes in [(torch.quint8, torch.qint8, None)]: @@ -336,7 +336,7 @@ def convert(self, node.name, scale, zero_point, modules, quantized_graph, node_name_to_scope) kwargs = {**self.binary_op_node.kwargs} - add_args = (*load_arg(quantized=True)(self.binary_op_node.args), scale_arg, zero_point_arg) + add_args = (*load_arg(quantized=activation_dtype(qconfig))(self.binary_op_node.args), scale_arg, zero_point_arg) op = quantized_graph.create_node( 'call_function', self.quantized_binary_op, add_args, kwargs) return op @@ -344,14 +344,14 @@ def convert(self, assert dtypes == (torch.float16, torch.float16, None) # TODO (refactor) this is duplicated, maybe have a helper function if self.relu_node: - op_out = quantized_graph.node_copy(self.binary_op_node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(self.binary_op_node, load_arg(quantized=torch.float)) relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) op_out = quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) else: - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) return quantized_graph.create_node( "call_method", "to", (op_out, torch.float16,), {} ) @@ -373,14 +373,14 @@ def convert(self, ) ) if self.relu_node: - op_out = quantized_graph.node_copy(self.binary_op_node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(self.binary_op_node, load_arg(quantized=torch.float)) relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) return quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) else: - return quantized_graph.node_copy(node, load_arg(quantized=False)) + return quantized_graph.node_copy(node, load_arg(quantized=torch.float)) @register_quant_pattern(torch.cat) @@ -396,7 +396,7 @@ def convert(self, convert_custom_config_dict: Dict[str, Any] = None) -> Node: if not self.all_node_args_are_tensors: return NotImplemented - return quantized_graph.node_copy(node, load_arg(quantized=True)) + return quantized_graph.node_copy(node, load_arg(quantized=torch.quint8)) # handle conv, maybe followed by relu # NB: matching order is reversed, that is we match from the bottom of this list to the beginning @@ -472,14 +472,14 @@ def convert(self, "supported by Conv " "supported dtype combinations are: {}".format(dtypes, supported_dtypes)) if self.relu_node: - conv_out = quantized_graph.node_copy(self.conv_node, load_arg(quantized=False)) + conv_out = quantized_graph.node_copy(self.conv_node, load_arg(quantized=torch.float)) relu_args = [conv_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) return quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) else: - return quantized_graph.node_copy(node, load_arg(quantized=False)) + return quantized_graph.node_copy(node, load_arg(quantized=torch.float)) activation_int8_quantized = activation_is_int8_quantized(qconfig) @@ -504,20 +504,20 @@ def convert(self, return quantized_graph.create_node( 'call_module', self.conv_node.target, - (load_arg(quantized=True)(self.conv_node.args[0]),), + (load_arg(quantized=torch.quint8)(self.conv_node.args[0]),), {}) else: # call_function assert self.conv_node.op == "call_function" if is_reference: - args = load_arg(quantized=[0, 1])(self.conv_node.args) - args = load_arg(quantized=False)(self.conv_node.args) - kwargs = load_arg(quantized=False)(self.conv_node.kwargs) + args = load_arg(quantized=[torch.quint8, torch.qint8])(self.conv_node.args) + args = load_arg(quantized=torch.float)(self.conv_node.args) + kwargs = load_arg(quantized=torch.float)(self.conv_node.kwargs) op_out = quantized_graph.create_node( "call_function", self.conv, args, kwargs) if self.relu_node: relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) op_out = quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) @@ -542,10 +542,10 @@ def convert(self, else: assert len(self.conv_node.args) >= 7, \ "only conv2d calls with all arguments specified is supported right now in is_reference=False option" - args = load_arg(quantized=[0, 1])(self.conv_node.args) + args = load_arg(quantized={0: torch.quint8, 1: torch.qint8})(self.conv_node.args) # pack weight - weight = load_arg(quantized=True)(self.conv_node.args[1]) - other_args = load_arg(quantized=False)(self.conv_node.args[2:]) + weight = load_arg(quantized=torch.qint8)(self.conv_node.args[1]) + other_args = load_arg(quantized=torch.float)(self.conv_node.args[2:]) bias, stride, padding, dilation, groups = other_args if self.conv == torch.nn.functional.conv1d: # F.conv1d can take `int` as well as `list[int]` for stride, @@ -563,7 +563,7 @@ def convert(self, # construct conv input if activation_int8_quantized: qconv_op = get_qconv_op(self.conv, self.relu_node is not None) - conv_input = load_arg(quantized=True)(self.conv_node.args[0]) + conv_input = load_arg(quantized=torch.quint8)(self.conv_node.args[0]) activation_post_process = \ self._maybe_get_last_node_only_observer(modules) @@ -575,7 +575,7 @@ def convert(self, self.conv_node.name, scale, zero_point, modules, quantized_graph, node_name_to_scope) qconv_args = (conv_input, packed_weight, scale_node, zero_point_node) - kwargs = load_arg(quantized=False)(self.conv_node.kwargs) + kwargs = load_arg(quantized=torch.float)(self.conv_node.kwargs) op = quantized_graph.create_node( 'call_function', qconv_op, qconv_args, kwargs) # Store the name of the fused op to get the path of node after fusion as well. @@ -644,10 +644,10 @@ def convert(self, "supported by Linear " "supported dtype combinations are: {}".format(dtypes, supported_dtypes)) if self.relu_node: - op_out = quantized_graph.node_copy(self.linear_node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(self.linear_node, load_arg(quantized=torch.float)) relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) return quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) else: @@ -683,27 +683,30 @@ def convert(self, parent_name, name = _parent_name(self.linear_node.target) setattr(modules[parent_name], name, quantized) # activation needs to be quantized for static quantization + dtype = torch.float + if activation_int8_quantized: + dtype = activation_dtype(qconfig) return quantized_graph.create_node( 'call_module', self.linear_node.target, - (load_arg(quantized=activation_int8_quantized)(self.linear_node.args[0]),), {}) + (load_arg(quantized=dtype)(self.linear_node.args[0]),), {}) else: # call_function assert self.linear_node.op == 'call_function' if is_reference: - quantized_input_idxs = [] + quantized_input_dtypes = [torch.float, torch.float] if activation_int8_quantized: - quantized_input_idxs.append(0) + quantized_input_dtypes[0] = torch.quint8 if weight_is_statically_quantized(qconfig): - quantized_input_idxs.append(1) - args = load_arg(quantized=quantized_input_idxs)(self.linear_node.args) - args = load_arg(quantized=False)(self.linear_node.args) - kwargs = load_arg(quantized=False)(self.linear_node.kwargs) + quantized_input_dtypes[1] = torch.qint8 + args = load_arg(quantized=quantized_input_dtypes)(self.linear_node.args) + args = load_arg(quantized=torch.float)(self.linear_node.args) + kwargs = load_arg(quantized=torch.float)(self.linear_node.kwargs) op_out = quantized_graph.create_node( "call_function", torch.nn.functional.linear, args, kwargs) if self.relu_node: relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) op_out = quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) @@ -731,17 +734,19 @@ def convert(self, if dtypes != (torch.float16, torch.float16, None): # linear args # (x, weight, bias, ...) + # TODO: the name should be weight is int8 quantized weight_quantized = weight_is_statically_quantized(qconfig) - linear_weight = load_arg(quantized=weight_quantized)(self.linear_node.args[1]) + dtype = weight_dtype if weight_quantized else torch.float + linear_weight = load_arg(quantized=dtype)(self.linear_node.args[1]) # get other arguments - kwargs = {**load_arg(quantized=False)(self.linear_node.kwargs)} + kwargs = {**load_arg(quantized=torch.float)(self.linear_node.kwargs)} # pack weight bias = None # all args after bias, including bias - other_args = load_arg(quantized=False)(self.linear_node.args[2:]) + other_args = load_arg(quantized=torch.float)(self.linear_node.args[2:]) if len(self.linear_node.args) > 2: - bias = load_arg(quantized=False)(self.linear_node.args[2]) + bias = load_arg(quantized=torch.float)(self.linear_node.args[2]) other_args = other_args[1:] # remove the bias argument else: assert 'bias' in kwargs, \ @@ -755,7 +760,7 @@ def convert(self, # construct linear input if activation_int8_quantized: qlinear_op = torch.ops.quantized.linear_relu if self.relu_node else torch.ops.quantized.linear - linear_input = load_arg(quantized=True)(self.linear_node.args[0]) + linear_input = load_arg(quantized=torch.quint8)(self.linear_node.args[0]) activation_post_process = \ self._maybe_get_last_node_only_observer(modules) assert activation_post_process is not None @@ -779,7 +784,7 @@ def convert(self, qlinear_op = torch.ops.quantized.linear_dynamic \ if weight_dtype == torch.qint8 \ else torch.ops.quantized.linear_dynamic_fp16 - linear_input = load_arg(quantized=False)(self.linear_node.args[0]) + linear_input = load_arg(quantized=torch.float)(self.linear_node.args[0]) qlinear_args = (linear_input, packed_weight) # type: ignore[assignment] op_out = quantized_graph.create_node( "call_function", qlinear_op, qlinear_args, kwargs) @@ -794,14 +799,14 @@ def convert(self, assert dtypes == (torch.float16, torch.float16, None) # TODO (refactor) this is duplicated, maybe have a helper function if self.relu_node: - op_out = quantized_graph.node_copy(self.linear_node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(self.linear_node, load_arg(quantized=torch.float)) relu_args = [op_out] - relu_args.extend(load_arg(quantized=False)(self.relu_node.args[1:])) - relu_kwargs = load_arg(quantized=False)(self.relu_node.kwargs) + relu_args.extend(load_arg(quantized=torch.float)(self.relu_node.args[1:])) + relu_kwargs = load_arg(quantized=torch.float)(self.relu_node.kwargs) op_out = quantized_graph.create_node( "call_function", torch.nn.functional.relu, tuple(relu_args), relu_kwargs) else: - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) return quantized_graph.create_node( "call_method", "to", (op_out, torch.float16), {}) @@ -844,7 +849,7 @@ def convert(self, 'call_module', self.bn_node.target, load_arg(quantized=[0])(self.bn_node.args), - load_arg(quantized=False)(self.bn_node.kwargs)) + load_arg(quantized=torch.float)(self.bn_node.kwargs)) @register_quant_pattern(torch.nn.Embedding) @register_quant_pattern(torch.nn.EmbeddingBag) @@ -895,8 +900,8 @@ def convert(self, return quantized_graph.create_node( 'call_module', emb_node.target, - load_arg(quantized=False)(emb_node.args), - load_arg(quantized=False)(emb_node.kwargs)) + load_arg(quantized=torch.float)(emb_node.args), + load_arg(quantized=torch.float)(emb_node.kwargs)) # TODO (maybe): merge with embedding quantize handler @register_quant_pattern(torch.nn.GRUCell) @@ -949,8 +954,8 @@ def convert(self, return quantized_graph.create_node( 'call_module', node.target, - load_arg(quantized=False)(node.args), - load_arg(quantized=False)(node.kwargs)) + load_arg(quantized=torch.float)(node.args), + load_arg(quantized=torch.float)(node.kwargs)) ARGS_TO_SKIP = { torch._ops.ops.quantized.hardswish: ['inplace'], @@ -1052,7 +1057,7 @@ def convert(self, "dtype combination: {} is not " "supported by {} " "supported dtype combinations are: {}".format(dtypes, self.op, supported_dtypes[self.op])) - return quantized_graph.node_copy(node, load_arg(quantized=False)) + return quantized_graph.node_copy(node, load_arg(quantized=torch.float)) # TODO: make helper functions for (torch.quint8, torch.qint8, None) if not is_reference: if dtypes in [(torch.quint8, torch.qint8, None)]: @@ -1071,7 +1076,7 @@ def convert(self, 'call_module', node.target, load_arg(quantized=[0])(node.args), - load_arg(quantized=False)(node.kwargs)) + load_arg(quantized=torch.float)(node.kwargs)) else: assert node.op == "call_function" # call_function @@ -1087,7 +1092,7 @@ def convert(self, "call_function to be a function instead of a string" quantized_op = get_quantized_operator(node.target) args = load_arg(quantized=[0])(node.args) - kwargs = {**load_arg(quantized=False)(node.kwargs), "output_scale": scale_arg, + kwargs = {**load_arg(quantized=torch.float)(node.kwargs), "output_scale": scale_arg, "output_zero_point": zero_point_arg} if quantized_op in ARGS_TO_SKIP: args_to_skip = ARGS_TO_SKIP[quantized_op] @@ -1102,16 +1107,16 @@ def convert(self, warnings.warn( "Only reference patterns are currently supported for {dtype} dtype with {op} op" "".format(dtype=dtypes, op=self.op)) - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) return quantized_graph.create_node( "call_method", "to", (op_out, torch.float16), {}) else: assert is_reference if dtypes in [(torch.quint8, torch.qint8, None)]: load_arg(quantized=[0])(node.args) - args = load_arg(quantized=False)(node.args) - kwargs = load_arg(quantized=False)(node.kwargs) - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + args = load_arg(quantized=torch.float)(node.args) + kwargs = load_arg(quantized=torch.float)(node.kwargs) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) activation_post_process = \ self._maybe_get_last_node_only_observer(modules) assert activation_post_process is not None @@ -1120,7 +1125,7 @@ def convert(self, node, modules, quantized_graph, node_name_to_scope, is_input=False) else: assert dtypes in [(torch.float16, torch.float16, None)] - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) return quantized_graph.create_node( "call_method", "to", (op_out, torch.float16), {}) @@ -1148,7 +1153,7 @@ def convert(self, quantized_op = get_quantized_operator(node.target) args = load_arg(quantized=[0])(node.args) - kwargs = {**load_arg(quantized=False)(node.kwargs), 'output_scale': scale_arg, 'output_zero_point': zero_point_arg} + kwargs = {**load_arg(quantized=torch.float)(node.kwargs), 'output_scale': scale_arg, 'output_zero_point': zero_point_arg} kwargs.pop('inplace') return quantized_graph.create_node( 'call_function', quantized_op, args, kwargs) # type: ignore[arg-type] @@ -1210,7 +1215,7 @@ def convert(self, convert_custom_config_dict: Dict[str, Any] = None) -> Node: dtypes = get_qconfig_dtypes(qconfig) if dtypes == (torch.float16, torch.float16, None): - op_out = quantized_graph.node_copy(node, load_arg(quantized=False)) + op_out = quantized_graph.node_copy(node, load_arg(quantized=torch.float)) return quantized_graph.create_node( "call_method", "to", (op_out, torch.float16,), {} ) From 55755edc60aee98a769f9e8123f975316247600c Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 18 Jun 2021 14:29:14 -0700 Subject: [PATCH 245/305] [jit] Made a list for element-wise ops. (#59579) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59579 Test Plan: Imported from OSS Reviewed By: eellison Differential Revision: D28955319 Pulled By: navahgar fbshipit-source-id: 605531aedf9250a226b0401d55fda3427bdc6f33 --- torch/csrc/jit/passes/tensorexpr_fuser.cpp | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.cpp b/torch/csrc/jit/passes/tensorexpr_fuser.cpp index 32983f996cdbf..2909428c3d39c 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.cpp +++ b/torch/csrc/jit/passes/tensorexpr_fuser.cpp @@ -164,21 +164,28 @@ static const OperatorSet& supported_eltwise_set() { "aten::where.ScalarSelf(Tensor condition, Scalar self, Tensor other) -> Tensor", "aten::where.ScalarOther(Tensor condition, Tensor self, Scalar other) -> Tensor", "aten::where.Scalar(Tensor condition, Scalar self, Scalar other) -> Tensor", - "aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor", // TODO: enable other min/max variants, operators that can be both // elementwise or reductions: "aten::min.other(Tensor self, Tensor other) -> Tensor", "aten::max.other(Tensor self, Tensor other) -> Tensor", // TODO: enable slice, shape inference is not implemented for this op yet - - "aten::conv2d(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] dilation=1, int groups=1) -> Tensor", - "aten::matmul(Tensor self, Tensor other) -> Tensor", }; // clang-format on return supported_eltwise_set; } +static const OperatorSet& supported_non_eltwise_set() { + // clang-format off + static const OperatorSet supported_non_eltwise_set{ + "aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor", + "aten::conv2d(Tensor input, Tensor weight, Tensor? bias=None, int[2] stride=1, int[2] padding=0, int[2] dilation=1, int groups=1) -> Tensor", + "aten::matmul(Tensor self, Tensor other) -> Tensor", + }; + // clang-format on + return supported_non_eltwise_set; +}; + bool isSupported(Node* node) { // For Block codegen we allow limited ops. if (tensorexpr::getTEGenerateBlockCode()) { @@ -198,6 +205,7 @@ bool isSupported(Node* node) { // clang-format on if (node->isMemberOf(supported_eltwise_set()) || + node->isMemberOf(supported_non_eltwise_set()) || node->isMemberOf(supported_misc_set) || (texpr_reductions_enabled && node->isMemberOf(supported_reduction_set))) { // We only insert guards on Tensor types, so we rely on the output @@ -527,9 +535,11 @@ class TensorExprFuser { continue; } - // we only support shape calculations for elementwise and + // we only support shape calculations for elementwise, some + // non-elementwise like batch_norm, conv, matmul, and // a few exceptions (e.g. prim::ConstantChunk, etc) listed above - if (!n->isMemberOf(tensorexpr::supported_eltwise_set())) { + if (!n->isMemberOf(tensorexpr::supported_eltwise_set()) && + !n->isMemberOf(tensorexpr::supported_non_eltwise_set())) { continue; } From d4c626a346ce6c7160ea07ae2d798072eab0f633 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 18 Jun 2021 14:29:14 -0700 Subject: [PATCH 246/305] [jit] Exported a method to get the supported list of elementwise ops (#60162) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60162 Test Plan: Imported from OSS Reviewed By: bertmaher Differential Revision: D29190841 Pulled By: navahgar fbshipit-source-id: bb786a653441c5b586509e25cc80d357d2223af3 --- torch/csrc/jit/passes/tensorexpr_fuser.cpp | 2 +- torch/csrc/jit/passes/tensorexpr_fuser.h | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.cpp b/torch/csrc/jit/passes/tensorexpr_fuser.cpp index 2909428c3d39c..5188275d002da 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.cpp +++ b/torch/csrc/jit/passes/tensorexpr_fuser.cpp @@ -60,7 +60,7 @@ Value* broadcastSizes(at::ArrayRef sizes, AliasDb* db) { namespace tensorexpr { -static const OperatorSet& supported_eltwise_set() { +const OperatorSet& supported_eltwise_set() { // clang-format off // breaks up the schema strings so they are no longer discoverable with ctrl-F static const OperatorSet supported_eltwise_set{ diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.h b/torch/csrc/jit/passes/tensorexpr_fuser.h index cc8b427030de9..55239b298c828 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.h +++ b/torch/csrc/jit/passes/tensorexpr_fuser.h @@ -60,6 +60,7 @@ TORCH_API bool usedOnlyInSize(Value* v); TORCH_API Value* broadcastSizes(at::ArrayRef sizes, AliasDb* db); namespace tensorexpr { +TORCH_API const OperatorSet& supported_eltwise_set(); TORCH_API bool isSupported(Node* node); } } // namespace jit From d0c4ace00f189782c8f848c37cef70de671a0aa4 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 18 Jun 2021 14:29:14 -0700 Subject: [PATCH 247/305] [jit] Added a tranformation to move consumers of aten::cat to its inputs, in the fused subgraphs (#59580) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59580 Test Plan: Imported from OSS Reviewed By: bertmaher Differential Revision: D28955318 Pulled By: navahgar fbshipit-source-id: 7504d5aea441920f4eb9234cdfa17077161ab13c --- test/cpp/tensorexpr/CMakeLists.txt | 1 + test/cpp/tensorexpr/test_graph_opt.cpp | 317 ++++++++++++++++++++++++ tools/build_variables.bzl | 1 + torch/csrc/jit/tensorexpr/graph_opt.cpp | 164 ++++++++++++ torch/csrc/jit/tensorexpr/graph_opt.h | 63 +++++ torch/csrc/jit/tensorexpr/kernel.cpp | 2 + 6 files changed, 548 insertions(+) create mode 100644 test/cpp/tensorexpr/test_graph_opt.cpp create mode 100644 torch/csrc/jit/tensorexpr/graph_opt.cpp create mode 100644 torch/csrc/jit/tensorexpr/graph_opt.h diff --git a/test/cpp/tensorexpr/CMakeLists.txt b/test/cpp/tensorexpr/CMakeLists.txt index ce87e9ad802b9..516eb3a85b05e 100644 --- a/test/cpp/tensorexpr/CMakeLists.txt +++ b/test/cpp/tensorexpr/CMakeLists.txt @@ -7,6 +7,7 @@ set(TENSOREXPR_TEST_SRCS ${TENSOREXPR_TEST_ROOT}/test_cpp_codegen.cpp ${TENSOREXPR_TEST_ROOT}/test_expr.cpp ${TENSOREXPR_TEST_ROOT}/test_external_calls.cpp + ${TENSOREXPR_TEST_ROOT}/test_graph_opt.cpp ${TENSOREXPR_TEST_ROOT}/test_ir_printer.cpp ${TENSOREXPR_TEST_ROOT}/test_ir_verifier.cpp ${TENSOREXPR_TEST_ROOT}/test_kernel.cpp diff --git a/test/cpp/tensorexpr/test_graph_opt.cpp b/test/cpp/tensorexpr/test_graph_opt.cpp new file mode 100644 index 0000000000000..685251b0e9b29 --- /dev/null +++ b/test/cpp/tensorexpr/test_graph_opt.cpp @@ -0,0 +1,317 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include + +namespace torch { +namespace jit { + +using namespace torch::jit::tensorexpr; + +class GraphOpt : public ::testing::Test { + public: + // NOLINTNEXTLINE(modernize-use-override,cppcoreguidelines-explicit-virtual-functions) + void SetUp() { + old_cat_wo_conditionals_ = getCatWoConditionals(); + getCatWoConditionals() = true; + } + + void TearDown() { + getCatWoConditionals() = old_cat_wo_conditionals_; + } + + private: + bool old_cat_wo_conditionals_; +}; + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCat) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Float(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::log(%cat) + return (%5))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // The `aten::log` op must be moved to the inputs of `aten::cat`. + testing::FileCheck() + .check("aten::log") + ->check("aten::log") + ->check("aten::log") + ->check("aten::cat") + ->check_not("aten::log") + ->run(*kernel.graph()); + + auto x = at::rand({10}, at::kFloat); + auto y = at::rand({20}, at::kFloat); + auto z = at::rand({30}, at::kFloat); + auto ref = at::log(at::cat({x, y, z}, 0)); + + std::vector inputs = {x, y, z}; + std::vector stack = fmap(inputs); + kernel.run(stack); + auto out = stack[0].toTensor(); + ASSERT_EQ(out.sizes(), ref.sizes()); + ASSERT_EQ(out.dtype(), ref.dtype()); + ASSERT_TRUE(at::allclose(out, ref)); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCat2) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Float(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::log(%cat) + %6 : Float(60, strides=[1], device=cpu) = aten::tanh(%5) + return (%6))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // The `aten::log` and `aten::tanh` ops must be moved to the inputs of + // `aten::cat`. + testing::FileCheck() + .check("aten::log") + ->check("aten::log") + ->check("aten::log") + ->check("aten::tanh") + ->check("aten::tanh") + ->check("aten::tanh") + ->check("aten::cat") + ->check_not("aten::log") + ->check_not("aten::tanh") + ->run(*kernel.graph()); + + auto x = at::rand({10}, at::kFloat); + auto y = at::rand({20}, at::kFloat); + auto z = at::rand({30}, at::kFloat); + auto ref = at::tanh(at::log(at::cat({x, y, z}, 0))); + + std::vector inputs = {x, y, z}; + std::vector stack = fmap(inputs); + kernel.run(stack); + auto out = stack[0].toTensor(); + ASSERT_EQ(out.sizes(), ref.sizes()); + ASSERT_EQ(out.dtype(), ref.dtype()); + ASSERT_TRUE(at::allclose(out, ref)); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCat3) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%a : Float(60, strides=[1], device=cpu), + %x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Float(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::tanh(%cat) + %6 : Float(60, strides=[1], device=cpu) = aten::mul(%a, %5) + return (%6))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // The `aten::tanh` op must be moved to the inputs of `aten::cat`. + // But the `aten::mul` op must not be moved since it is not a single-tensor + // op (it has 2 tensor inputs). + testing::FileCheck() + .check("aten::tanh") + ->check("aten::tanh") + ->check("aten::tanh") + ->check("aten::cat") + ->check("aten::mul") + ->check_not("aten::tanh") + ->run(*kernel.graph()); + + auto a = at::rand({60}, at::kFloat); + auto x = at::rand({10}, at::kFloat); + auto y = at::rand({20}, at::kFloat); + auto z = at::rand({30}, at::kFloat); + auto ref = at::tanh(at::cat({x, y, z}, 0)) * a; + + std::vector inputs = {a, x, y, z}; + std::vector stack = fmap(inputs); + kernel.run(stack); + auto out = stack[0].toTensor(); + ASSERT_EQ(out.sizes(), ref.sizes()); + ASSERT_EQ(out.dtype(), ref.dtype()); + ASSERT_TRUE(at::allclose(out, ref)); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCatWithTypePromotionInUser) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%x : Int(10, strides=[1], device=cpu), + %y : Int(20, strides=[1], device=cpu), + %z : Int(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Int(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::tanh(%cat) + return (%5))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // The `aten::tanh` op must be moved to the inputs of `aten::cat`. + // The scalar type of the inputs to `cat` should now be `Float` since they + // are the result of `tanh` which does the type promotion. + testing::FileCheck() + .check("aten::tanh") + ->check("aten::tanh") + ->check("aten::tanh") + ->check("aten::cat") + ->check_not("aten::tanh") + ->run(*kernel.graph()); + + auto x = at::randint(std::numeric_limits::max(), {10}, at::kInt); + auto y = at::randint(std::numeric_limits::max(), {20}, at::kInt); + auto z = at::randint(std::numeric_limits::max(), {30}, at::kInt); + auto ref = at::tanh(at::cat({x, y, z}, 0)); + + std::vector inputs = {x, y, z}; + std::vector stack = fmap(inputs); + kernel.run(stack); + auto out = stack[0].toTensor(); + ASSERT_EQ(out.sizes(), ref.sizes()); + ASSERT_EQ(out.dtype(), ref.dtype()); + ASSERT_TRUE(at::allclose(out, ref)); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCatWithTypePromotionInCat) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Double(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Double(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Double(60, strides=[1], device=cpu) = aten::log(%cat) + return (%5))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // No transformation should have happened because the `aten::cat` op performs + // type promotion. This case is currently not handled. + testing::FileCheck() + .check("aten::cat") + ->check("aten::log") + ->check_not("aten::cat") + ->check_not("aten::log") + ->run(*kernel.graph()); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCatNoSingleTensorElementwiseOp) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%0 : Float(60, strides=[1], device=cpu), + %x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Float(30, strides=[1], device=cpu)): + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::mul(%0, %cat) + return (%5))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // No transformation is expected since the consumers of cat are not + // single-tensor element-wise ops. + testing::FileCheck() + .check("aten::cat") + ->check("aten::mul") + ->check_not("aten::cat") + ->check_not("aten::mul") + ->run(*kernel.graph()); +#endif +} + +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(GraphOpt, OptimizeCatNoSingleTensorElementwiseOp2) { +#ifdef TORCH_ENABLE_LLVM + const auto graph_string = R"IR( + graph(%0 : Float(60, strides=[1], device=cpu), + %1 : Float(60, strides=[1], device=cpu), + %x : Float(10, strides=[1], device=cpu), + %y : Float(20, strides=[1], device=cpu), + %z : Float(30, strides=[1], device=cpu)): + %one : int = prim::Constant[value=1]() + %dim : int = prim::Constant[value=0]() + %xyz_list : Tensor[] = prim::ListConstruct(%x, %y, %z) + %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xyz_list, %dim) + %5 : Float(60, strides=[1], device=cpu) = aten::mul(%0, %cat) + %6 : Float(60, strides=[1], device=cpu) = aten::add(%5, %1, %one) + return (%6))IR"; + auto g = std::make_shared(); + torch::jit::parseIR(graph_string, g.get()); + g->lint(); + + KernelScope kernel_scope; + TensorExprKernel kernel(g); + + // No transformation is expected since the consumers of cat are not + // single-tensor element-wise ops. + testing::FileCheck() + .check("aten::cat") + ->check("aten::mul") + ->check("aten::add") + ->check_not("aten::cat") + ->check_not("aten::mul") + ->check_not("aten::add") + ->run(*kernel.graph()); +#endif +} + +} // namespace jit +} // namespace torch diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index 06acafd645eab..cf7e67fa77797 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -276,6 +276,7 @@ core_sources_full_mobile = [ "torch/csrc/jit/tensorexpr/eval.cpp", "torch/csrc/jit/tensorexpr/expr.cpp", "torch/csrc/jit/tensorexpr/external_functions_registry.cpp", + "torch/csrc/jit/tensorexpr/graph_opt.cpp", "torch/csrc/jit/tensorexpr/hash_provider.cpp", "torch/csrc/jit/tensorexpr/intrinsic_symbols.cpp", "torch/csrc/jit/tensorexpr/ir.cpp", diff --git a/torch/csrc/jit/tensorexpr/graph_opt.cpp b/torch/csrc/jit/tensorexpr/graph_opt.cpp new file mode 100644 index 0000000000000..67f9a671bfa20 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/graph_opt.cpp @@ -0,0 +1,164 @@ +#include + +#include +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +// Move the given user of `aten::cat` op to its inputs. +Node* moveCatAfterUse(Node* cat, Node* user, std::shared_ptr subgraph) { + // Example IR: + // %1 = ... + // %2 = ... + // %3 = prim::ListConstruct(%1, %2) + // %4 = aten::cat(%3, ...) + // %5 = aten::relu(%4) + // return (%5) + // + // To be transformed to: + // %1 = ... + // %2 = ... + // %5.1 = aten::relu(%1) + // %5.2 = aten::relu(%2) + // %3 = prim::ListConstruct(%5.1, %5.2) + // %4 = aten::cat(%3, ...) + // return (%4) + + TORCH_INTERNAL_ASSERT(cat->output()->hasUses()); + TORCH_INTERNAL_ASSERT(cat->output()->uses().size() == 1); + TORCH_INTERNAL_ASSERT(cat->input(0)->node()->kind() == prim::ListConstruct); + auto cat_list = cat->input(0)->node(); + auto cat_inputs = cat_list->inputs(); + + auto user_tensor_type = user->output()->type()->cast(); + TORCH_INTERNAL_ASSERT(user_tensor_type); + std::unordered_map new_cat_inputs; + for (auto inp : cat_inputs) { + auto new_cat_input = subgraph->createClone( + user, [&](Value* k) { return (k == cat->output()) ? inp : k; }); + // Since we are cloning user, its result should be the same scalar type + // as the user. But the dims should correspond to that of the input. + auto input_tensor_type = inp->type()->cast(); + TORCH_INTERNAL_ASSERT(input_tensor_type); + auto new_input_type = + input_tensor_type->withScalarType(user_tensor_type->scalarType()); + new_cat_input->output()->setType(new_input_type); + new_cat_input->insertBefore(cat_list); + new_cat_inputs[inp] = new_cat_input->output(); + } + auto new_cat_list = subgraph->createClone( + cat_list, [&](Value* k) { return new_cat_inputs[k]; }); + new_cat_list->insertBefore(cat); + auto new_cat = subgraph->createClone(cat, [&](Value* k) { + return (k == cat_list->output()) ? new_cat_list->output() : k; + }); + new_cat->output()->setType(user_tensor_type); + new_cat->insertBefore(cat); + + user->output()->replaceAllUsesWith(new_cat->output()); + user->destroy(); + + TORCH_INTERNAL_ASSERT(!cat->output()->hasUses()); + cat->destroy(); + + if (!cat_list->output()->hasUses()) { + cat_list->destroy(); + } + + return new_cat; +} + +int numTensorInputs(Node* node) { + int count = 0; + for (auto v : node->inputs()) { + if (v->type()->cast()) { + ++count; + } + } + return count; +} + +// Returns true if the given `cat` node promotes types. +// If the inputs to `cat` are of different types, then the implementation +// of `cat` is expected to promote type. +bool doesCatPromoteTypes(Node* node) { + TORCH_INTERNAL_ASSERT(node->kind() == aten::cat); + TORCH_INTERNAL_ASSERT(node->input(0)->node()->kind() == prim::ListConstruct); + auto inputs = node->input(0)->node()->inputs(); + TORCH_INTERNAL_ASSERT(!inputs.empty()); + auto scalar_type = + inputs.front()->type()->cast()->scalarType(); + for (size_t i = 1; i < inputs.size(); ++i) { + auto inp_scalar_type = + inputs[i]->type()->cast()->scalarType(); + if (scalar_type != inp_scalar_type) { + return true; + } + } + return false; +} + +// Move the users of the given `aten::cat` op to its inputs. +// The following constraints need to be satisfied on the cat op and its user. +// * the cat op should have only one use. +// * the user should be an element-wise op. +// * the user should have only one tensor input. +// - If the user has > 1 tensor inputs, that user op cannot be applied on +// the inputs of cat because the other tensor inputs will not be split, +// and hence the shape of those tensors would not match that of the +// inputs of cat. +// For example: +// %1 = ... +// %2 = ... +// %3 = prim::ListConstruct([%1, %2]) +// %4 = aten::cat(%3, ...) +// %5 = aten::add(%4, %0) +// In this example, we cannot move `aten::add` to the inputs of +// `aten::cat`, %1 and %2, because the shape of %0 will be different. +// * the cat op does not promote types. +// - When the cat op promote types, the type of inputs to cat after moving +// it user needs to reflect the original type. This is currently not +// handled. TODO +void moveCatOpToEnd(Node* cat, std::shared_ptr subgraph) { + TORCH_INTERNAL_ASSERT(cat->kind() == aten::cat); + if (cat->output()->uses().size() == 1) { + auto use = cat->output()->uses().front(); + if (use.user->isMemberOf(supported_eltwise_set()) && + numTensorInputs(use.user) == 1) { + if (!doesCatPromoteTypes(cat)) { + TORCH_INTERNAL_ASSERT( + use.user->output()->owningGraph() == subgraph.get()); + auto new_cat = moveCatAfterUse(cat, use.user, subgraph); + moveCatOpToEnd(new_cat, subgraph); + } + } + } +} + +// Moves the users of `aten::cat` ops to its inputs whenever possible +// in the given subgraph. +void moveCatOpsToEnd(std::shared_ptr subgraph) { + std::vector cat_nodes; + for (Node* n : subgraph->nodes()) { + if (n->kind() == aten::cat) { + cat_nodes.push_back(n); + } + } + for (auto cat : cat_nodes) { + moveCatOpToEnd(cat, subgraph); + } +} + +bool OptimizeCat(const std::shared_ptr& graph) { + if (getCatWoConditionals()) { + moveCatOpsToEnd(graph); + return true; + } + return false; +} + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/graph_opt.h b/torch/csrc/jit/tensorexpr/graph_opt.h new file mode 100644 index 0000000000000..5a81553b425f3 --- /dev/null +++ b/torch/csrc/jit/tensorexpr/graph_opt.h @@ -0,0 +1,63 @@ +#pragma once + +#include + +namespace torch { +namespace jit { +namespace tensorexpr { + +// Optimize aten::cat ops in the given subgraph. +// +// Moving users of cat to its inputs. +// Cat ops get lowered into multiple loops, one per input. When the result +// of cat is used by some other op, it results in a situation where inlining +// of cat does not happen. This in turn results in intermediate buffers +// being created for the result of cat, since it is not inlined. +// +// For example, consider the following graph: +// graph(%x : Float(10, strides=[1], device=cpu), +// %y : Float(20, strides=[1], device=cpu)): +// %dim : int = prim::Constant[value=0]() +// %xy_list : Tensor[] = prim::ListConstruct(%x, %y) +// %cat : Float(60, strides=[1], device=cpu) = aten::cat(%xy_list, %dim) +// %5 : Float(60, strides=[1], device=cpu) = aten::log(%cat) +// return (%5))IR"; +// +// This will get lowered into: +// Allocate(aten_cat); +// for (...) +// aten_cat[...] = x[...] +// for (...) +// aten_cat[...] = y[...] +// for (...) +// aten_log[...] = log(aten_cat[...]) +// Free(aten_cat); +// Note that aten_cat is not inlined into aten_log and it results in +// an intermediate buffer allocation as well. +// +// Optimization: +// We move the ops that use the result of `cat` into its inputs whenever +// possible. +// +// The graph above will be transformed to: +// graph(%x : Float(10, strides=[1], device=cpu), +// %y : Float(20, strides=[1], device=cpu)): +// %3 : int = prim::Constant[value=0]() +// %7 : Float(10, strides=[1], device=cpu) = aten::log(%x) +// %8 : Float(20, strides=[1], device=cpu) = aten::log(%y) +// %9 : Tensor[] = prim::ListConstruct(%7, %8) +// %10 : Float(60, strides=[1], device=cpu) = aten::cat(%9, %3) +// return (%10) +// +// This will get lowered into: +// for (...) +// aten_cat[...] = log(x[...]) +// for (...) +// aten_cat[...] = log(y[...]) +// aten_cat is the output buffer here. + +bool OptimizeCat(const std::shared_ptr& graph); + +} // namespace tensorexpr +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/tensorexpr/kernel.cpp b/torch/csrc/jit/tensorexpr/kernel.cpp index d716c8320f369..30385c1428928 100644 --- a/torch/csrc/jit/tensorexpr/kernel.cpp +++ b/torch/csrc/jit/tensorexpr/kernel.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -2999,6 +3000,7 @@ void TensorExprKernel::compile() { GRAPH_DUMP("TensorExprKernel graph:", graph_); device_ = *pickDeviceType(graph_); + OptimizeCat(graph_); // Block to collect the Stmts corresponding to all tensors. auto block = new Block({}); From 47bbc01e0b5efb69c93c6e6d0b46251b55049e5c Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Fri, 18 Jun 2021 14:29:14 -0700 Subject: [PATCH 248/305] [nnc] Added micro-benchmark to show perf improvement with cat subgraph optimization (#59581) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59581 Test Plan: Imported from OSS Reviewed By: bertmaher Differential Revision: D28955317 Pulled By: navahgar fbshipit-source-id: 53bb3dbfafbd3b146063f305523c2e6ec96cf6b8 --- benchmarks/tensorexpr/concat.py | 63 +++++++++++++++++++++++++++++- benchmarks/tensorexpr/pt_engine.py | 9 +++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/benchmarks/tensorexpr/concat.py b/benchmarks/tensorexpr/concat.py index fe63454337aae..a6d9b857f1e80 100644 --- a/benchmarks/tensorexpr/concat.py +++ b/benchmarks/tensorexpr/concat.py @@ -1,5 +1,6 @@ from . import benchmark import numpy as np +import torch class Concat2D2InputBench(benchmark.Benchmark): def __init__(self, mode, device, dtype, I1_D1, I1_D2, I2_D1, I2_D2, concat_dim): @@ -50,7 +51,67 @@ def default_configs(): [1, 580, 1, 174, 1], [20, 160, 20, 14, 1], [20, 580, 20, 174, 1], - [8, 512, 8, 512, 1] + [8, 512, 8, 512, 1], + [1 << 13, 1060, 1 << 13, 1040, 1], + [1 << 13, 2000, 1 << 13, 1074, 1], + [1 << 15, 1060, 1 << 15, 2670, 1], + [1 << 15, 5120, 1 << 15, 2512, 1] ] benchmark.register_benchmark_class(Concat2D2InputBench) + +class ConcatGraphOptBench(benchmark.Benchmark): + def __init__(self, mode, device, dtype, I1_D1, I1_D2, I2_D1, I2_D2, concat_dim): + super().__init__(mode, device, dtype) + self.I1_D1 = I1_D1 + self.I1_D2 = I1_D2 + self.I2_D1 = I2_D1 + self.I2_D2 = I2_D2 + self.concat_dim = concat_dim + self.input1 = self.randn([I1_D1, I1_D2], device=device, dtype=dtype, requires_grad=self.requires_grad) + self.input2 = self.randn([I2_D1, I2_D2], device=device, dtype=dtype, requires_grad=self.requires_grad) + self.inputs = [self.input1, self.input2] + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_cat_wo_conditionals(True) + + def forward(self, input1, input2): + x1 = self.add(input1, 0.00001) + x2 = self.add(input2, 0.00001) + y = self.cat((x1, x2), dim=self.concat_dim) + z = self.relu(y) + return z + + def reference(self): + return np.concatenate((self.numpy(self.input1), self.numpy(self.input2)), axis=concat_dim) + + def config(self): + return [self.I1_D1, self.I1_D2, self.I2_D1, self.I2_D2, self.concat_dim] + + @staticmethod + def module(): + return "concatGraphOpt" + + def memory_workload(self): + if self.mode == "fwd": + sol_count = 1 + 1 + algorithmic_count = 3 + 1 + else: + sol_count = (1 + 1) + (1 + 1) + algorithmic_count = (3 + 1) + (3 + 1) + + buffer_size = self.I1_D1 * self.I1_D2 + self.I2_D1 * self.I2_D2 + return { + "sol": buffer_size * sol_count, + "algorithmic": buffer_size * algorithmic_count, + } + + @staticmethod + def default_configs(): + return [ + [1 << 13, 1060, 1 << 13, 1040, 1], + [1 << 13, 2000, 1 << 13, 1074, 1], + [1 << 15, 1060, 1 << 15, 2670, 1], + [1 << 15, 5120, 1 << 15, 2512, 1] + ] + +benchmark.register_benchmark_class(ConcatGraphOptBench) diff --git a/benchmarks/tensorexpr/pt_engine.py b/benchmarks/tensorexpr/pt_engine.py index cfe283216b964..c25b568a2271f 100644 --- a/benchmarks/tensorexpr/pt_engine.py +++ b/benchmarks/tensorexpr/pt_engine.py @@ -50,6 +50,15 @@ def softmax(self, data, dim=None, dtype=None): def cat(self, inputs, dim=0): return torch.cat(inputs, dim=dim) + def clamp(self, data, min, max): + return torch.clamp(data, min=min, max=max) + + def relu(self, data): + return torch.nn.functional.relu(data) + + def tanh(self, data): + return torch.tanh(data) + def max_pool2d(self, data, kernel_size, stride=1): return torch.nn.functional.max_pool2d(data, kernel_size, stride=stride) From e8e3394ea860c68ce8fafff7e52acb93b913ae51 Mon Sep 17 00:00:00 2001 From: Peter Bell Date: Fri, 18 Jun 2021 16:28:00 -0700 Subject: [PATCH 249/305] Recognize transposed dense tensors as a form of partial overlap (#59014) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59014 Fixes #48401 `assert_no_overlap` currently has a false-negative where it recognizes the transpose of a contiguous tensor as fully overlapping. This happens because the memory regions do fully overlap, but of course the strides are different so the actual elements don't all overlap. This goes slightly in the other direction, by requiring strides to exactly match we get false-positives for some unusual situations, e.g. ``` torch.add(a, a, out=a.view([1, *a.shape])) ``` Or replacing strides of length-1 dimensions, etc. However, I think these are sufficiently obscure that it's okay to error and the common cases like inplace operations still work as before. Test Plan: Imported from OSS Reviewed By: gchanan Differential Revision: D29040928 Pulled By: ngimel fbshipit-source-id: 5a636c67536a3809c83f0d3117d2fdf49c0a45e6 --- aten/src/ATen/MemoryOverlap.cpp | 3 ++- test/test_torch.py | 24 +++++++++++++++++++----- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/aten/src/ATen/MemoryOverlap.cpp b/aten/src/ATen/MemoryOverlap.cpp index c2388ebf8d968..b7d98932f7074 100644 --- a/aten/src/ATen/MemoryOverlap.cpp +++ b/aten/src/ATen/MemoryOverlap.cpp @@ -61,7 +61,8 @@ MemOverlapStatus get_overlap_status(TensorImpl* a, TensorImpl* b) { const auto b_end = b_begin + b->numel() * b->itemsize(); if (a_begin == b_begin && a_end == b_end) { - return MemOverlapStatus::FULL; + return (a->strides() == b->strides()) ? + MemOverlapStatus::FULL : MemOverlapStatus::PARTIAL; } if (a_begin < b_end && b_begin < a_end) { return MemOverlapStatus::PARTIAL; diff --git a/test/test_torch.py b/test/test_torch.py index b30697e099443..b5f2f85034c15 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -3047,24 +3047,38 @@ def _test(op, output, input): with self.assertRaises(AssertionError): with self.assertRaisesRegex(RuntimeError, 'unsupported operation'): _test(op, data[0:sz], data[1:sz + 1]) + # output is transpose of input: + length = int(math.sqrt(sz)) + input = data[:length**2].view([length, length]) + out = input.t() + if not expected_failure: + with self.assertRaisesRegex(RuntimeError, 'unsupported operation'): + _test(op, out, input) + else: + with self.assertRaises(AssertionError): + with self.assertRaisesRegex(RuntimeError, 'unsupported operation'): + _test(op, out, input) def ternary_check_input_output_mem_overlap(self, op, device, expected_failure=False): - sz = 3 + sz = 9 data = torch.randn(2 * sz, device=device) other1 = torch.randn(sz, device=device) other2 = torch.randn(sz, device=device) self.unary_check_input_output_mem_overlap( - data, sz, lambda input, out: op(input, other1, other2, out=out), + data, sz, lambda input, out: + op(input, other1.view(input.shape), other2.view(input.shape), out=out), expected_failure=expected_failure) self.unary_check_input_output_mem_overlap( - data, sz, lambda input, out: op(other1, input, other2, out=out), + data, sz, lambda input, out: + op(other1.view(input.shape), input, other2.view(input.shape), out=out), expected_failure=expected_failure) self.unary_check_input_output_mem_overlap( - data, sz, lambda input, out: op(other1, other2, input, out=out), + data, sz, lambda input, out: + op(other1.view(input.shape), other2.view(input.shape), input, out=out), expected_failure=expected_failure) @@ -6242,7 +6256,7 @@ def test_ternary_op_mem_overlap(self, device, dtype): def test_copy_mem_overlap(self, device, dtype): self.check_internal_mem_overlap( torch.Tensor.copy_, num_inputs=2, dtype=dtype, device=device) - sz = 3 + sz = 9 doubles = torch.randn(2 * sz, dtype=dtype, device=device) self.unary_check_input_output_mem_overlap( doubles, sz, lambda input, out: out.copy_(input)) From aae2a3c95ee6d62e834a5e6890a12f7ecf0dd17f Mon Sep 17 00:00:00 2001 From: Joel Schlosser Date: Fri, 18 Jun 2021 17:16:55 -0700 Subject: [PATCH 250/305] Clarify ConvTransposeNd + reference links (#60291) Summary: Fixes https://github.com/pytorch/pytorch/issues/56873 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60291 Reviewed By: gchanan Differential Revision: D29239199 Pulled By: jbschlosser fbshipit-source-id: 9b2de1a8b1a7444797f82c73195c5efc929562eb --- torch/nn/modules/conv.py | 42 +++++++++++++++++++++++----------------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/torch/nn/modules/conv.py b/torch/nn/modules/conv.py index ba1e6e0d09bdc..9ccfa48bec7e7 100644 --- a/torch/nn/modules/conv.py +++ b/torch/nn/modules/conv.py @@ -653,7 +653,9 @@ class ConvTranspose1d(_ConvTransposeNd): This module can be seen as the gradient of Conv1d with respect to its input. It is also known as a fractionally-strided convolution or - a deconvolution (although it is not an actual deconvolution operation). + a deconvolution (although it is not an actual deconvolution operation as it does + not compute a true inverse of convolution). For more information, see the visualizations + `here`_ and the `Deconvolutional Networks`_ paper. This module supports :ref:`TensorFloat32`. @@ -667,7 +669,7 @@ class ConvTranspose1d(_ConvTransposeNd): of the output shape. See note below for details. * :attr:`dilation` controls the spacing between the kernel points; also known as the à trous algorithm. - It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does. + It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does. {groups_note} @@ -726,11 +728,11 @@ class ConvTranspose1d(_ConvTransposeNd): sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where :math:`k = \frac{groups}{C_\text{out} * \text{kernel\_size}}` - .. _cross-correlation: - https://en.wikipedia.org/wiki/Cross-correlation - - .. _link: + .. _`here`: https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md + + .. _`Deconvolutional Networks`: + https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf """ def __init__( @@ -778,7 +780,9 @@ class ConvTranspose2d(_ConvTransposeNd): This module can be seen as the gradient of Conv2d with respect to its input. It is also known as a fractionally-strided convolution or - a deconvolution (although it is not an actual deconvolution operation). + a deconvolution (although it is not an actual deconvolution operation as it does + not compute a true inverse of convolution). For more information, see the visualizations + `here`_ and the `Deconvolutional Networks`_ paper. This module supports :ref:`TensorFloat32`. @@ -792,7 +796,7 @@ class ConvTranspose2d(_ConvTransposeNd): of the output shape. See note below for details. * :attr:`dilation` controls the spacing between the kernel points; also known as the à trous algorithm. - It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does. + It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does. {groups_note} @@ -874,11 +878,11 @@ class ConvTranspose2d(_ConvTransposeNd): >>> output.size() torch.Size([1, 16, 12, 12]) - .. _cross-correlation: - https://en.wikipedia.org/wiki/Cross-correlation - - .. _link: + .. _`here`: https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md + + .. _`Deconvolutional Networks`: + https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf """ def __init__( @@ -929,7 +933,9 @@ class ConvTranspose3d(_ConvTransposeNd): This module can be seen as the gradient of Conv3d with respect to its input. It is also known as a fractionally-strided convolution or - a deconvolution (although it is not an actual deconvolution operation). + a deconvolution (although it is not an actual deconvolution operation as it does + not compute a true inverse of convolution). For more information, see the visualizations + `here`_ and the `Deconvolutional Networks`_ paper. This module supports :ref:`TensorFloat32`. @@ -943,7 +949,7 @@ class ConvTranspose3d(_ConvTransposeNd): of the output shape. See note below for details. * :attr:`dilation` controls the spacing between the kernel points; also known as the à trous algorithm. - It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does. + It is harder to describe, but the link `here`_ has a nice visualization of what :attr:`dilation` does. {groups_note} @@ -1019,11 +1025,11 @@ class ConvTranspose3d(_ConvTransposeNd): >>> input = torch.randn(20, 16, 10, 50, 100) >>> output = m(input) - .. _cross-correlation: - https://en.wikipedia.org/wiki/Cross-correlation - - .. _link: + .. _`here`: https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md + + .. _`Deconvolutional Networks`: + https://www.matthewzeiler.com/mattzeiler/deconvolutionalnetworks.pdf """ def __init__( From 1c97c3e3a4c5d7c99d1bdfe5368cc666c7ec3132 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Fri, 18 Jun 2021 17:27:41 -0700 Subject: [PATCH 251/305] DOC Adds LSTM docs for defined variables when bidirectional=True (#60120) Summary: Fixes https://github.com/pytorch/pytorch/issues/59332 Pull Request resolved: https://github.com/pytorch/pytorch/pull/60120 Reviewed By: gchanan Differential Revision: D29240245 Pulled By: jbschlosser fbshipit-source-id: acad9c24f41f7253a7d42cd940e54bb66e083ecf --- torch/nn/modules/rnn.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py index 7a6fde26ba9d9..0d92428dc555f 100644 --- a/torch/nn/modules/rnn.py +++ b/torch/nn/modules/rnn.py @@ -574,6 +574,16 @@ class LSTM(RNNBase): weight_hr_l[k] : the learnable projection weights of the :math:`\text{k}^{th}` layer of shape `(proj_size, hidden_size)`. Only present when ``proj_size > 0`` was specified. + weight_ih_l[k]_reverse: Analogous to `weight_ih_l[k]` for the reverse direction. + Only present when ``bidirectional=True``. + weight_hh_l[k]_reverse: Analogous to `weight_hh_l[k]` for the reverse direction. + Only present when ``bidirectional=True``. + bias_ih_l[k]_reverse: Analogous to `bias_ih_l[k]` for the reverse direction. + Only present when ``bidirectional=True``. + bias_hh_l[k]_reverse: Analogous to `bias_hh_l[k]` for the reverse direction. + Only present when ``bidirectional=True``. + weight_hr_l[k]_reverse: Analogous to `weight_hr_l[k]` for the reverse direction. + Only present when ``bidirectional=True`` and ``proj_size > 0`` was specified. .. note:: All the weights and biases are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` From a8430f1076c2a2e61d83879acf2fe6be48bec889 Mon Sep 17 00:00:00 2001 From: Pritam Damania Date: Fri, 18 Jun 2021 17:36:11 -0700 Subject: [PATCH 252/305] Remove PlacementSpec from ShardingSpecs. (#59990) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59990 ShardingSpecs accepted a Device/PlacementSpec and was initially written this way for flexibility. Although, it is slightly confusing given there is no general use case for this. As a result, to keep things simple I've ensured that both specs only accept devices for now. We can always extend this to include a general PlacementSpec later on. ghstack-source-id: 131842525 Test Plan: waitforbuildbot Reviewed By: SciPioneer, rohan-varma Differential Revision: D29116463 fbshipit-source-id: a6f2b3f1346ac6afab91c9595d4cae4f4da04fda --- torch/distributed/_sharding_spec/api.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/torch/distributed/_sharding_spec/api.py b/torch/distributed/_sharding_spec/api.py index bed86ef46f638..38dd58ba6a6d2 100644 --- a/torch/distributed/_sharding_spec/api.py +++ b/torch/distributed/_sharding_spec/api.py @@ -62,7 +62,7 @@ class ChunkShardingSpec(ShardingSpec): The dimension to shard on, could be an integer representing the dimension or a string in case of named tensors where dimensions are named. - placement(List[Device] or List[PlacementSpec]): + placement(List[Device]): Specifies the placement of each shard of the Tensor. The size of the list represents the number of shards to be created. This parameter can be a list of devices @@ -77,10 +77,9 @@ class ChunkShardingSpec(ShardingSpec): """ ShardingDim = Union[int, str] - ShardPlacements = List[Union[Device, PlacementSpec]] dim: ShardingDim - placements: ShardPlacements + placements: List[Device] def __post_init__(self): self._verify_dim(self.dim) @@ -91,7 +90,7 @@ def _verify_devices(placements): if placements is None or len(placements) == 0: raise ValueError(f'None/Empty placement provided: {placements}') for dev in placements: - if not isinstance(dev, PlacementSpec) and not is_valid_device(dev): + if not is_valid_device(dev): raise ValueError(f'{dev} is not a valid device') @staticmethod @@ -113,7 +112,7 @@ class ShardMetadata(object): shard_lengths(List[int]): Lengths indicating the length of each dimension for this shard. Should have the same rank as the original tensor. - placement(List[Device or PlacementSpec]): + placement(List[Device]): Specifies the placement of each shard of the Tensor. The size of the list represents the number of shards to be created. This parameter can be a list of devices @@ -127,16 +126,14 @@ class ShardMetadata(object): 2. "/" (ex: "trainer0/cuda:0"). """ - ShardPlacement = Union[Device, PlacementSpec] - __slots__ = ['shard_offsets', 'shard_lengths', 'placement'] shard_offsets: List[int] shard_lengths: List[int] - placement: ShardPlacement + placement: Device def __post_init__(self): - if not isinstance(self.placement, PlacementSpec) and not is_valid_device(self.placement): + if not is_valid_device(self.placement): raise ValueError(f'{self.placement} is not a valid device') if len(self.shard_offsets) != len(self.shard_lengths): From 319890b1b2eb3c5063ffded7bf3ba94200e21b45 Mon Sep 17 00:00:00 2001 From: Pritam Damania Date: Fri, 18 Jun 2021 17:52:26 -0700 Subject: [PATCH 253/305] Support *args in Pipe.forward API. (#55441) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55441 This is the first step towards supporting the proposal outlined in https://github.com/pytorch/pytorch/issues/53952. In this PR I've ensured Pipe.forward() accepts a *inputs argument instead of just a single input as previously. This lays the groundwork for supporting non-Tensors and generic arguments to the Pipe API. In this PR we still only support Tensors and non-Tensor support will come in future PRs. For backward compatibility I've ensured a single Tuple[Tensor] input still works as expected previously. ghstack-source-id: 130767499 Test Plan: waitforbuildbot Reviewed By: SciPioneer Differential Revision: D27613887 fbshipit-source-id: 05e19e537e6d7fe4999745fc4ba9941ac54906de --- .../distributed/pipeline/sync/test_balance.py | 6 +- test/distributed/pipeline/sync/test_bugs.py | 8 +-- .../pipeline/sync/test_microbatch.py | 11 +-- test/distributed/pipeline/sync/test_pipe.py | 68 ++++++++++--------- torch/distributed/pipeline/sync/checkpoint.py | 13 +++- torch/distributed/pipeline/sync/microbatch.py | 36 +++++----- torch/distributed/pipeline/sync/pipe.py | 36 ++++++++-- torch/distributed/pipeline/sync/pipeline.py | 10 +-- .../distributed/pipeline/sync/skip/layout.py | 10 ++- 9 files changed, 115 insertions(+), 83 deletions(-) diff --git a/test/distributed/pipeline/sync/test_balance.py b/test/distributed/pipeline/sync/test_balance.py index 5aa9ec6d454e9..140e5b75497da 100644 --- a/test/distributed/pipeline/sync/test_balance.py +++ b/test/distributed/pipeline/sync/test_balance.py @@ -191,8 +191,7 @@ def forward(self, x): return x, x.detach() class Add(nn.Module): - def forward(self, a_b): - a, b = a_b + def forward(self, a, b): return a + b model = nn.Sequential(Twin(), Add()) @@ -207,8 +206,7 @@ def forward(self, x): return x, x.detach() class Add(nn.Module): - def forward(self, a_b): - a, b = a_b + def forward(self, a, b): return a + b model = nn.Sequential(Twin(), Add()) diff --git a/test/distributed/pipeline/sync/test_bugs.py b/test/distributed/pipeline/sync/test_bugs.py index 580e58bf58bcb..875317dc12a26 100644 --- a/test/distributed/pipeline/sync/test_bugs.py +++ b/test/distributed/pipeline/sync/test_bugs.py @@ -90,8 +90,7 @@ def __init__(self): super().__init__() self.ones = nn.Parameter(torch.ones(32, 3, 32, 32, requires_grad=True)) - def forward(self, pair): - a, b = pair + def forward(self, a, b): a = a * self.ones return a * 1, b * 2, b * 3 @@ -100,8 +99,7 @@ def __init__(self): super().__init__() self.ones = nn.Parameter(torch.ones(32, 3, 32, 32, requires_grad=True)) - def forward(self, triple): - a, b, c = triple + def forward(self, a, b, c): a = a * self.ones b = Sleep.apply(b) return a + b + c @@ -112,7 +110,7 @@ def forward(self, triple): a = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True) b = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True) - y = model((a, b)) + y = model(a, b) y.local_value().norm().backward() torch.cuda.synchronize(0) diff --git a/test/distributed/pipeline/sync/test_microbatch.py b/test/distributed/pipeline/sync/test_microbatch.py index 914e9e8e8ae2a..4a5c67cceaaeb 100644 --- a/test/distributed/pipeline/sync/test_microbatch.py +++ b/test/distributed/pipeline/sync/test_microbatch.py @@ -49,8 +49,11 @@ def test_batch_call(): def f(x): return x + def g(x, y): + return x, y + assert a.call(f).atomic - assert not b.call(f).atomic + assert not b.call(g).atomic def test_batch_setitem_by_index(): @@ -86,7 +89,7 @@ def test_batch_setitem_by_slice(): def test_check(): check(torch.tensor(42)) - check((torch.tensor(4), torch.tensor(2))) + check(torch.tensor(4), torch.tensor(2)) with pytest.raises(TypeError): check(42) @@ -127,10 +130,10 @@ def test_scatter_tensor(): assert b.tensor.size() == (1, 1) -def test_scatter_tuple(): +def test_scatter_multiple_tensors(): ab = (torch.zeros(2, 1), torch.zeros(4, 2)) - a, b = scatter(ab, chunks=2) + a, b = scatter(*ab, chunks=2) assert a.tensors[0].size() == (1, 1) assert b.tensors[0].size() == (1, 1) diff --git a/test/distributed/pipeline/sync/test_pipe.py b/test/distributed/pipeline/sync/test_pipe.py index 12516768a98c2..c1777a38891ef 100644 --- a/test/distributed/pipeline/sync/test_pipe.py +++ b/test/distributed/pipeline/sync/test_pipe.py @@ -13,6 +13,7 @@ from torch import nn from torch.distributed.pipeline.sync import Pipe +from torch.distributed.pipeline.sync.pipe import PipeSequential skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required") @@ -177,8 +178,8 @@ def forward(self, input): return (input * 2, torch.tensor([False])) class JoinNonFloat(nn.Module): - def forward(self, input): - return input[0] * 2 + def forward(self, input, non_float): + return input * 2 model = nn.Sequential(ForkNonFloat(), JoinNonFloat()) model = Pipe(model, chunks=1, checkpoint="always") @@ -295,8 +296,7 @@ def __init__(self): self.fc_a = nn.Linear(1, 1) self.fc_b = nn.Linear(1, 1) - def forward(self, a_and_b): - a, b = a_and_b + def forward(self, a, b): return (self.fc_a(a), self.fc_b(b)) model = nn.Sequential(Two()) @@ -305,24 +305,24 @@ def forward(self, a_and_b): a = torch.rand(10, 1, requires_grad=True) b = torch.rand(10, 1, requires_grad=True) - a_out, b_out = model((a, b)).local_value() - loss = (a_out + b_out).mean() - loss.backward() - - assert a.grad is not None - assert b.grad is not None - - # Test with list. - a.grad = None - b.grad = None - a_out, b_out = model([a, b]).local_value() + a_out, b_out = model(a, b).local_value() loss = (a_out + b_out).mean() loss.backward() assert a.grad is not None assert b.grad is not None +def test_multi_sequence_input(setup_rpc): + class MultiSeq(nn.Module): + def forward(self, tup1, tup2): + return tup1, tup2 + model = Pipe(nn.Sequential(MultiSeq())) + with pytest.raises(TypeError): + model( + [torch.rand(10), torch.rand(10)], + [torch.rand(10), torch.rand(10)] + ) def test_input_singleton(setup_rpc): class One(nn.Module): @@ -330,8 +330,7 @@ def __init__(self): super().__init__() self.fc = nn.Linear(1, 1) - def forward(self, only_a): - (a,) = only_a + def forward(self, a): return (self.fc(a),) model = nn.Sequential(One()) @@ -339,19 +338,7 @@ def forward(self, only_a): a = torch.rand(10, 1, requires_grad=True) - (a_out,) = model((a,)).local_value() - loss = a_out.mean() - loss.backward() - - assert all(p.grad is not None for p in model.parameters()) - assert a.grad is not None - - # Test with list - a.grad = None - for p in model.parameters(): - p.grad = None - - (a_out,) = model([a]).local_value() + (a_out,) = model(a).local_value() loss = a_out.mean() loss.backward() @@ -485,8 +472,8 @@ def test_merged_partitions(setup_rpc): model = Pipe(model) assert isinstance(model.partitions, nn.ModuleList) - assert isinstance(model.partitions[0], nn.Sequential) - assert isinstance(model.partitions[1], nn.Sequential) + assert isinstance(model.partitions[0], PipeSequential) + assert isinstance(model.partitions[1], PipeSequential) assert list(model.partitions[0]) == [a, b[0], b[1]] assert list(model.partitions[1]) == [c] assert list(model.partitions[2]) == [d] @@ -536,7 +523,6 @@ def test_empty_module(setup_rpc): model = Pipe(model) assert model(torch.tensor(42)).local_value() == torch.tensor(42) - assert model((torch.tensor(42),)).local_value() == (torch.tensor(42),) # But only tensor or tensors is legal in Pipe. with pytest.raises(TypeError): @@ -650,3 +636,19 @@ def forward(self, x): # Partition #1: 000! 111! 222! # assert timeline == [(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (2, 1)] + +@pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"]) +@skip_if_no_cuda +def test_multiple_inputs(checkpoint, setup_rpc): + class Module1(nn.Module): + def forward(self, a, b, c): + return a + b + c, a * b * c + + class Module2(nn.Module): + def forward(self, a, b): + return a + b + + model = Pipe(nn.Sequential(Module1().cuda(0), Module2().cuda(0)), chunks=2, checkpoint=checkpoint) + t = torch.rand(10) + res = model(t, t, t).local_value() + assert torch.equal(res, (t + t + t) + (t * t * t)) diff --git a/torch/distributed/pipeline/sync/checkpoint.py b/torch/distributed/pipeline/sync/checkpoint.py index 7d2faa8db041c..8dd7a488ff0a7 100644 --- a/torch/distributed/pipeline/sync/checkpoint.py +++ b/torch/distributed/pipeline/sync/checkpoint.py @@ -268,8 +268,11 @@ def forward( ctx.save_for_backward(*input) with torch.no_grad(), enable_checkpointing(): - output = function(input[0] if input_atomic else input) - + if input_atomic: + assert len(input) == 1 + output = function(input[0]) + else: + output = function(*input) return output @staticmethod @@ -317,7 +320,11 @@ def backward(ctx: Context, *grad_output: Tensor) -> Tuple[None, ...]: # pragma: with restore_rng_states(input[0].device, ctx.rng_states): with torch.enable_grad(), enable_recomputing(): - output = ctx.function(input_leaf[0] if ctx.input_atomic else input_leaf) + if ctx.input_atomic: + assert len(input_leaf) == 1 + output = ctx.function(input_leaf[0]) + else: + output = ctx.function(*input_leaf) ctx.recomputed.append((output, input_leaf)) diff --git a/torch/distributed/pipeline/sync/microbatch.py b/torch/distributed/pipeline/sync/microbatch.py index f422bfc3e0a36..db5733c937429 100644 --- a/torch/distributed/pipeline/sync/microbatch.py +++ b/torch/distributed/pipeline/sync/microbatch.py @@ -6,7 +6,7 @@ # LICENSE file in the root directory of this source tree. """Manipulation of micro-batches.""" import typing -from typing import Callable, Iterable, Iterator, List, Union, cast, Sequence +from typing import Callable, Iterator, List, Union, cast, Sequence import torch from torch import Tensor @@ -67,7 +67,10 @@ def call(self, function: Function) -> "Batch": """Calls a function by the underlying tensor or tensors. It also wraps the output with :class:`Batch`. """ - return Batch(function(self.value)) + if self.atomic: + return Batch(function(self.value)) + else: + return Batch(function(*self.value)) def __repr__(self) -> str: return f"Batch[atomic={self.atomic!r}]({self.value!r})" @@ -132,39 +135,32 @@ def _setitem_by_slice(self, index: slice, value: Tensors) -> None: self.value = value[0] -def check(input: TensorOrTensors) -> None: +def check(*inputs) -> None: """Checks whether the input is a tensor or tensors. Raises: TypeError: input is not a tensor or tensors. """ - if isinstance(input, Sequence): - for x in input: - if not isinstance(x, Tensor): - raise TypeError(f"expected Tensor, but got {input.__class__.__name__}") - return + for input in inputs: + if not isinstance(input, Tensor): + raise TypeError(f"expected Tensor, but got {input.__class__.__name__}") - if not isinstance(input, Tensor): - raise TypeError(f"expected Tensor, but got {input.__class__.__name__}") - -def scatter(input: TensorOrTensors, chunks: int) -> List[Batch]: +def scatter(*inputs, chunks: int) -> List[Batch]: """Splits an input mini-batch into multiple micro-batches.""" - inputs: Iterable[TensorOrTensors] - - if isinstance(input, Tensor): - inputs = input.chunk(chunks) + if len(inputs) == 1 and isinstance(inputs[0], Tensor): + unwrapped_inputs = inputs[0].chunk(chunks) else: rotated: List[Tensors] = [] - for tensor in input: + for tensor in inputs: tensors = tensor.chunk(chunks) - rotated.append(cast(Tensors, tensors)) + rotated.append(tensors) - inputs = zip(*rotated) + unwrapped_inputs = zip(*rotated) # type: ignore[assignment] - return [Batch(x) for x in inputs] + return [Batch(x) for x in unwrapped_inputs] def gather(outputs: List[Batch]) -> TensorOrTensors: diff --git a/torch/distributed/pipeline/sync/pipe.py b/torch/distributed/pipeline/sync/pipe.py index 22b980d0f06f7..55675e57ec6ef 100644 --- a/torch/distributed/pipeline/sync/pipe.py +++ b/torch/distributed/pipeline/sync/pipe.py @@ -6,7 +6,7 @@ # LICENSE file in the root directory of this source tree. """The Pipe interface.""" from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Union, cast, Sequence +from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Union, Sequence, Tuple, cast import torch from torch import Tensor, nn @@ -117,6 +117,22 @@ def _retrieve_device(module: nn.Module) -> torch.device: return device if device is not None else torch.device("cpu") + +class PipeSequential(nn.Sequential): + """ + Pipe variant of ``nn.Sequential`` which supports multiple inputs. + """ + + def forward(self, *inputs): + for module in self: + if isinstance(inputs, Tuple): # type: ignore[arg-type] + inputs = module(*inputs) + else: + # Don't expand single variables (ex: lists/Tensor) + inputs = module(inputs) + return inputs + + def _assemble_partition(modules: List[nn.Module]): modules_list: List[nn.Module] = [] for module in modules: @@ -124,7 +140,7 @@ def _assemble_partition(modules: List[nn.Module]): modules_list.extend(module.children()) else: modules_list.append(module) - return nn.Sequential(*modules_list) + return PipeSequential(*modules_list) def _split_module(modules: nn.Sequential) -> Tuple[List[nn.Sequential], List[torch.device]]: partitions = [] @@ -353,7 +369,7 @@ def _ensure_copy_streams(self) -> List[List[AbstractStream]]: return self._copy_streams - def forward(self, input) -> RRef: + def forward(self, *inputs) -> RRef: """ Processes a single input mini-batch through the pipe and returns an :class:`~torch.distributed.rpc.RRef` pointing to the output. @@ -363,6 +379,12 @@ def forward(self, input) -> RRef: :class:`~torch.Tensor` or a sequence of tensors. This restriction is applied at partition boundaries too. + The sequence of inputs are fed into the first stage of the pipeline as + ``*inputs``. As a result the positional args for this function should + match the positional args for the first stage of the pipeline. The same + condition applies for output of one stage of the pipeline which is the + input for the next stage. + The input tensor is split into multiple micro-batches based on the ``chunks`` parameter used to initialize :class:`Pipe`. The batch size is assumed to be the first dimension of the tensor and if the batch @@ -370,7 +392,7 @@ def forward(self, input) -> RRef: the batch size. Args: - input (torch.Tensor or sequence of :class:`~torch.Tensor`): input mini-batch + inputs (torch.Tensor or sequence of :class:`~torch.Tensor`): input mini-batch Returns: :class:`~torch.distributed.rpc.RRef` to the output of the mini-batch @@ -379,14 +401,14 @@ def forward(self, input) -> RRef: TypeError: input is not a tensor or sequence of tensors. """ - microbatch.check(input) + microbatch.check(*inputs) if not self.devices: # Empty sequential module is not illegal. - return RRef(input) + return RRef(*inputs) # Divide a mini-batch into micro-batches. - batches = microbatch.scatter(input, self.chunks) + batches = microbatch.scatter(*inputs, chunks=self.chunks) # Run pipeline parallelism. self.pipeline.run(batches) diff --git a/torch/distributed/pipeline/sync/pipeline.py b/torch/distributed/pipeline/sync/pipeline.py index 4c5fd033057a5..579e1286e7b4a 100644 --- a/torch/distributed/pipeline/sync/pipeline.py +++ b/torch/distributed/pipeline/sync/pipeline.py @@ -194,16 +194,16 @@ def compute( if checkpoint: def function( - input: TensorOrTensors, - partition: nn.Sequential = partition, + *inputs: TensorOrTensors, + partition: nn.Module = partition, skip_tracker: SkipTrackerThroughPotals = skip_trackers[i], chunk_id: int = i, part_id: int = j, ) -> TensorOrTensors: with use_skip_tracker(skip_tracker), record_function("chunk%d-part%d" % (chunk_id, part_id)): - return partition(input) + return partition(*inputs) - chk = Checkpointing(function, batch) + chk = Checkpointing(function, batch) # type: ignore[arg-type] task = Task(streams[j], compute=chk.checkpoint, finalize=chk.recompute) del function, chk @@ -211,7 +211,7 @@ def function( def compute( batch: Batch = batch, - partition: nn.Sequential = partition, + partition: nn.Module = partition, skip_tracker: SkipTrackerThroughPotals = skip_trackers[i], chunk_id: int = i, part_id: int = j, diff --git a/torch/distributed/pipeline/sync/skip/layout.py b/torch/distributed/pipeline/sync/skip/layout.py index bff417bfbd65b..04d76d34ea166 100644 --- a/torch/distributed/pipeline/sync/skip/layout.py +++ b/torch/distributed/pipeline/sync/skip/layout.py @@ -72,9 +72,9 @@ def inspect_skip_layout(partitions: List[nn.Sequential]) -> SkipLayout: stashed_at: Dict[Tuple[Namespace, str], int] = {} for j, partition in enumerate(partitions): - for layer in partition: + def inspect_layer(layer): if not isinstance(layer, Skippable): - continue + return for ns, name in layer.stashable(): stashed_at[(ns, name)] = j @@ -83,4 +83,10 @@ def inspect_skip_layout(partitions: List[nn.Sequential]) -> SkipLayout: prev_j = stashed_at.pop((ns, name)) skip_routes[(ns, name)] = (prev_j, j) + if isinstance(partition, nn.Sequential): + for layer in partition: + inspect_layer(layer) + else: + inspect_layer(partition) + return SkipLayout(len(partitions), skip_routes) From 80e6e3f1da625fc4e446a36c9ec1ff86543c1ba4 Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Fri, 18 Jun 2021 18:43:11 -0700 Subject: [PATCH 254/305] [iOS GPU][BE][1/n] Rename MPSCNNContext to MetalContext (#60280) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60280 No significant changes besides renaming the class. In the future, we'll convert this objc class to c++. ghstack-source-id: 131827490 Test Plan: - CircleCI - buck test pp-mac Reviewed By: SS-JIA Differential Revision: D29231824 fbshipit-source-id: a0d1327a55a0414011c78a7144d3b05f1579cf42 --- aten/src/ATen/native/metal/MetalAten.mm | 4 +-- .../ATen/native/metal/MetalCommandBuffer.mm | 4 +-- .../MPSCNNContext.h => MetalContext.h} | 2 +- .../MPSCNNContext.mm => MetalContext.mm} | 8 +++--- aten/src/ATen/native/metal/MetalUtils.h | 6 ++-- aten/src/ATen/native/metal/MetalUtils.mm | 2 +- .../ATen/native/metal/mpscnn/MPSCNNClampOp.mm | 6 ++-- .../ATen/native/metal/mpscnn/MPSCNNConvOp.mm | 6 ++-- .../metal/mpscnn/MPSCNNFullyConnectedOp.mm | 6 ++-- .../native/metal/mpscnn/MPSCNNNeuronOp.mm | 10 +++---- .../ATen/native/metal/mpscnn/MPSImageUtils.mm | 28 +++++++++---------- .../native/metal/mpscnn/MPSImageWrapper.mm | 4 +-- aten/src/ATen/native/metal/ops/MetalAddmm.mm | 2 +- .../metal/ops/MetalBinaryElementwise.mm | 10 +++---- aten/src/ATen/native/metal/ops/MetalChunk.mm | 4 +-- aten/src/ATen/native/metal/ops/MetalConcat.mm | 10 +++---- aten/src/ATen/native/metal/ops/MetalCopy.mm | 4 +-- .../ATen/native/metal/ops/MetalHardswish.mm | 4 +-- .../src/ATen/native/metal/ops/MetalNeurons.mm | 2 +- .../src/ATen/native/metal/ops/MetalPadding.mm | 4 +-- .../src/ATen/native/metal/ops/MetalPooling.mm | 6 ++-- aten/src/ATen/native/metal/ops/MetalReduce.mm | 4 +-- .../src/ATen/native/metal/ops/MetalReshape.mm | 4 +-- .../src/ATen/native/metal/ops/MetalSoftmax.mm | 4 +-- .../ATen/native/metal/ops/MetalTranspose.mm | 6 ++-- .../metal/ops/MetalUpsamplingNearest.mm | 6 ++-- 26 files changed, 78 insertions(+), 78 deletions(-) rename aten/src/ATen/native/metal/{mpscnn/MPSCNNContext.h => MetalContext.h} (95%) rename aten/src/ATen/native/metal/{mpscnn/MPSCNNContext.mm => MetalContext.mm} (97%) diff --git a/aten/src/ATen/native/metal/MetalAten.mm b/aten/src/ATen/native/metal/MetalAten.mm index c9cee4092dd23..2677fe5a96eee 100644 --- a/aten/src/ATen/native/metal/MetalAten.mm +++ b/aten/src/ATen/native/metal/MetalAten.mm @@ -1,6 +1,6 @@ #import #import -#import +#import #import #include #include @@ -114,7 +114,7 @@ Tensor empty( struct MetalImpl : public at::metal::MetalInterface { bool is_metal_available() const override { #if defined(USE_PYTORCH_METAL) - return [[MPSCNNContext sharedInstance] available]; + return [[MetalContext sharedInstance] available]; #else return false; #endif diff --git a/aten/src/ATen/native/metal/MetalCommandBuffer.mm b/aten/src/ATen/native/metal/MetalCommandBuffer.mm index 900efd08730f2..3bbbe18a48b22 100644 --- a/aten/src/ATen/native/metal/MetalCommandBuffer.mm +++ b/aten/src/ATen/native/metal/MetalCommandBuffer.mm @@ -1,5 +1,5 @@ #import -#import +#import #import NSString* thread_local_storage_key = @"PTMetalCommandBuffer"; @@ -10,7 +10,7 @@ @implementation MetalCommandBuffer { + (MetalCommandBuffer*)newBuffer { MetalCommandBuffer* cb = [MetalCommandBuffer new]; - cb->_buffer = [[MPSCNNContext sharedInstance].commandQueue commandBuffer]; + cb->_buffer = [[MetalContext sharedInstance].commandQueue commandBuffer]; cb->_images = [NSMutableArray new]; cb->_delegates = [NSMutableSet new]; return cb; diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.h b/aten/src/ATen/native/metal/MetalContext.h similarity index 95% rename from aten/src/ATen/native/metal/mpscnn/MPSCNNContext.h rename to aten/src/ATen/native/metal/MetalContext.h index 876448ea18e3e..c459537799855 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.h +++ b/aten/src/ATen/native/metal/MetalContext.h @@ -5,7 +5,7 @@ API_AVAILABLE(ios(10.0), macos(10.13)) // TODO[T79947194]: Convert this class to C++ -@interface MPSCNNContext : NSObject +@interface MetalContext : NSObject @property(nonatomic, strong, readonly) id device; @property(nonatomic, strong, readonly) id commandQueue; @property(nonatomic, strong, readonly) id library; diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.mm b/aten/src/ATen/native/metal/MetalContext.mm similarity index 97% rename from aten/src/ATen/native/metal/mpscnn/MPSCNNContext.mm rename to aten/src/ATen/native/metal/MetalContext.mm index 3e0075def9bbe..064e59efe2c19 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.mm +++ b/aten/src/ATen/native/metal/MetalContext.mm @@ -1,6 +1,6 @@ #import #import -#import +#import #include @@ -14,7 +14,7 @@ #endif using namespace at::native::metal; -@implementation MPSCNNContext { +@implementation MetalContext { std::mutex _pipelineCacheMutex; MetalDeviceInfo _deviceInfo; std::unordered_map> _pipelineCache; @@ -22,9 +22,9 @@ @implementation MPSCNNContext { + (instancetype)sharedInstance { static dispatch_once_t onceToken; - static MPSCNNContext* instance = nil; + static MetalContext* instance = nil; dispatch_once(&onceToken, ^{ - instance = [[MPSCNNContext alloc] init]; + instance = [[MetalContext alloc] init]; id device = MTLCreateSystemDefaultDevice(); instance->_device = device; instance->_deviceInfo = createDeviceInfo(device); diff --git a/aten/src/ATen/native/metal/MetalUtils.h b/aten/src/ATen/native/metal/MetalUtils.h index e110da1bfcf4a..a94fce6f76e33 100644 --- a/aten/src/ATen/native/metal/MetalUtils.h +++ b/aten/src/ATen/native/metal/MetalUtils.h @@ -1,5 +1,5 @@ #include -#include +#include #include #include #include @@ -78,7 +78,7 @@ static inline MetalCommandBuffer* getCommandBufferFromTensor( template idmakeMTLBuffer(const std::vector& src) { - id buffer = [[MPSCNNContext sharedInstance].device + id buffer = [[MetalContext sharedInstance].device newBufferWithLength:src.size() * sizeof(T) options:MTLResourceOptionCPUCacheModeWriteCombined]; memcpy(buffer.contents, src.data(), src.size() * sizeof(T)); @@ -86,7 +86,7 @@ idmakeMTLBuffer(const std::vector& src) { } static inline idmakeMTLBuffer(int64_t bytes) { - id buffer = [[MPSCNNContext sharedInstance].device + id buffer = [[MetalContext sharedInstance].device newBufferWithLength:bytes options:MTLResourceOptionCPUCacheModeWriteCombined]; return buffer; diff --git a/aten/src/ATen/native/metal/MetalUtils.mm b/aten/src/ATen/native/metal/MetalUtils.mm index 2751ca3be3ee6..a082c153a235f 100644 --- a/aten/src/ATen/native/metal/MetalUtils.mm +++ b/aten/src/ATen/native/metal/MetalUtils.mm @@ -1,5 +1,5 @@ #import -#import +#import #import namespace at { diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm index acfa5e4c601be..2d380dd3cda1e 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import @implementation MPSCNNClampOp { @@ -28,14 +28,14 @@ - (void)encode:(id)cb { have to use `clamp(vector, half4, half4)` instead. */ id encoder = [cb computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:at::native::metal::mpscnn::kernelFor( _X, "clamp_half4", "clamp_half4_nonarray")]; [encoder setComputePipelineState:state]; [encoder setTexture:[_X texture] atIndex:0]; [encoder setTexture:[_Y texture] atIndex:1]; - id clampBuffer = [[MPSCNNContext sharedInstance].device + id clampBuffer = [[MetalContext sharedInstance].device newBufferWithLength:2 * sizeof(fp16_t) options:MTLResourceOptionCPUCacheModeWriteCombined]; fp16_t* clampBufferPtr = (fp16_t*)[clampBuffer contents]; diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNConvOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNConvOp.mm index 32bb8d293a0b4..83fd0d3c6c6de 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNConvOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNConvOp.mm @@ -1,5 +1,5 @@ #import -#import +#import #import #import @@ -120,14 +120,14 @@ + (MPSCNNConvOp*)conv2d:(const Conv2DParams&)params MPSCNNConvolution* conv = nil; if (@available(iOS 11.0, *)) { conv = [[MPSCNNConvolution alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device weights:dataSource]; } else { #if TARGET_OS_IPHONE // Fallback on earlier versions conv = [[MPSCNNConvolution alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device convolutionDescriptor:desc kernelWeights:w biasTerms:b diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNFullyConnectedOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNFullyConnectedOp.mm index e9c32595774a4..7e4d5974bbb82 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNFullyConnectedOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNFullyConnectedOp.mm @@ -1,4 +1,4 @@ -#import +#import #import #import @@ -28,12 +28,12 @@ + (MPSCNNFullyConnectedOp*)linear:(const Conv2DParams&)params Bias:(float*)b Desc:desc]; fc = [[MPSCNNFullyConnected alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device weights:ds]; } else { #if TARGET_OS_IPHONE fc = [[MPSCNNFullyConnected alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device convolutionDescriptor:desc kernelWeights:w biasTerms:b diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNNeuronOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNNeuronOp.mm index 69cd62df229ae..5e208731c88ce 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNNeuronOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNNeuronOp.mm @@ -1,4 +1,4 @@ -#import +#import #import @implementation MPSCNNNeuronOp @@ -8,7 +8,7 @@ + (MPSCNNNeuronHardSigmoid*)hardSigmoid API_AVAILABLE(ios(11.0), macos(10.13)) { static MPSCNNNeuronHardSigmoid* neuron = nil; dispatch_once(&onceToken, ^{ neuron = [[MPSCNNNeuronHardSigmoid alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device a:1.0 / 6.0 b:0.5]; }); @@ -20,7 +20,7 @@ + (MPSCNNNeuronReLU*)relu { static dispatch_once_t onceToken; dispatch_once(&onceToken, ^{ relu = [[MPSCNNNeuronReLU alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device a:0]; }); return relu; @@ -31,7 +31,7 @@ + (MPSCNNNeuronSigmoid*)sigmoid { static MPSCNNNeuronSigmoid* sigmoid = nil; dispatch_once(&onceToken, ^{ sigmoid = [[MPSCNNNeuronSigmoid alloc] - initWithDevice:[MPSCNNContext sharedInstance].device]; + initWithDevice:[MetalContext sharedInstance].device]; }); return sigmoid; } @@ -41,7 +41,7 @@ + (MPSCNNNeuronTanH*)tanh { static MPSCNNNeuronTanH* tanh = nil; dispatch_once(&onceToken, ^{ tanh = [[MPSCNNNeuronTanH alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device a:1 b:1]; }); diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm index 30083d0cf2e9c..8d2f17163780c 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm @@ -1,5 +1,5 @@ #import -#import +#import #import #import #import @@ -20,7 +20,7 @@ numberOfImages:sizes[0] usage:MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite]; - return [[MPSImage alloc] initWithDevice:[MPSCNNContext sharedInstance].device + return [[MPSImage alloc] initWithDevice:[MetalContext sharedInstance].device imageDescriptor:desc]; } @@ -38,7 +38,7 @@ usage:MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite]; MPSImage* image = - [[MPSImage alloc] initWithDevice:[MPSCNNContext sharedInstance].device + [[MPSImage alloc] initWithDevice:[MetalContext sharedInstance].device imageDescriptor:desc]; int64_t slices = (C + 3) / 4 * N; @@ -59,12 +59,12 @@ MPSImage* createStaticImage(const float* src, IntArrayRef sizes) { int64_t size_bytes = c10::multiply_integers(sizes) * sizeof(float); - id buff = [[MPSCNNContext sharedInstance].device + id buff = [[MetalContext sharedInstance].device newBufferWithLength:size_bytes options:MTLResourceOptionCPUCacheModeWriteCombined]; memcpy(buff.contents, src, size_bytes); MPSImage* output = createStaticImage(sizes); - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:metal::mpscnn::kernelFor( output, "copy_nchw_to_metal", @@ -103,7 +103,7 @@ MPSImage* Y = createStaticImage([image sizes]); MetalCommandBuffer* cb = [MetalCommandBuffer newBuffer]; id encoder = [cb.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:mpscnn::kernelFor(image, "copy", "copy_nonarray")]; [encoder setComputePipelineState:state]; [encoder setTexture:[image texture] atIndex:0]; @@ -125,7 +125,7 @@ TORCH_CHECK(buffer); MPSImage* Y = createStaticImage([image sizes]); id encoder = [buffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:mpscnn::kernelFor(image, "copy", "copy_nonarray")]; [encoder setComputePipelineState:state]; @@ -165,12 +165,12 @@ MPSTemporaryImage* createTemporaryImage(MetalCommandBuffer* buffer, IntArrayRef sizes, const float* src) { TORCH_CHECK(buffer); int64_t size_bytes = c10::multiply_integers(sizes) * sizeof(float); - id buff = [[MPSCNNContext sharedInstance].device + id buff = [[MetalContext sharedInstance].device newBufferWithBytes:src length:size_bytes options:MTLResourceStorageModeShared]; MPSTemporaryImage* output = createTemporaryImage(buffer, sizes); - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:metal::mpscnn::kernelFor( output, "copy_nchw_to_metal", @@ -199,7 +199,7 @@ TORCH_CHECK(buffer); MPSTemporaryImage* Y = createTemporaryImage(buffer, [image sizes]); id encoder = [buffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:metal::mpscnn::kernelFor(image, "copy", "copy_nonarray")]; [encoder setComputePipelineState:state]; [encoder setTexture:[image texture] atIndex:0]; @@ -215,14 +215,14 @@ void copyToHost(float* dst, MPSImage* image) { int64_t size_bytes = c10::multiply_integers([image sizes]) * sizeof(float); - id buffer = [[MPSCNNContext sharedInstance].device + id buffer = [[MetalContext sharedInstance].device newBufferWithLength:size_bytes options:MTLResourceOptionCPUCacheModeDefault]; id cb = - [MPSCNNContext sharedInstance].commandQueue.commandBuffer; + [MetalContext sharedInstance].commandQueue.commandBuffer; id encoder = [cb computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:metal::mpscnn::kernelFor( image, "copy_metal_to_nchw", @@ -254,7 +254,7 @@ void copyToMetalBuffer( TORCH_CHECK(cmdBuffer.buffer); id encoder = [cmdBuffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:metal::mpscnn::kernelFor( image, "copy_metal_to_nchw", diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm index c5931b2870fd7..10f396713e2a6 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm @@ -1,6 +1,6 @@ #import #import -#import +#import #import #import #import @@ -112,7 +112,7 @@ - (void)endSynchronization:(NSError*)error { void MPSImageWrapper::prepare() { if (!_buffer) { int64_t size_bytes = c10::multiply_integers([_image sizes]) * sizeof(float); - _buffer = [[MPSCNNContext sharedInstance].device + _buffer = [[MetalContext sharedInstance].device newBufferWithLength:size_bytes options:MTLResourceCPUCacheModeWriteCombined]; TORCH_CHECK(_buffer, "Allocate GPU memory failed!"); diff --git a/aten/src/ATen/native/metal/ops/MetalAddmm.mm b/aten/src/ATen/native/metal/ops/MetalAddmm.mm index 768a721451270..c023d91f7c1a1 100644 --- a/aten/src/ATen/native/metal/ops/MetalAddmm.mm +++ b/aten/src/ATen/native/metal/ops/MetalAddmm.mm @@ -5,7 +5,7 @@ #import #import #import -#import +#import #import #import #import diff --git a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm index ad611679b9881..294913d3244f3 100644 --- a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm +++ b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -82,7 +82,7 @@ Tensor binaryElementwiseShaderKernel( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); mt.texture()->allocateTemporaryStorage(outputSize, cb1); MPSImage* Y = mt.texture()->image(); - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:mpscnn::kernelFor(X1, arrayKernel, nonarrayKernel)]; id encoder = [cb1.buffer computeCommandEncoder]; [encoder setComputePipelineState:state]; @@ -122,7 +122,7 @@ Tensor binaryElementwiseShaderKernel( TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); MPSImage* Y = createTemporaryImage(cb1, outputSize.vec()); - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:mpscnn::kernelFor(X1, arrayKernel, nonarrayKernel)]; id encoder = [cb1.buffer computeCommandEncoder]; [encoder setComputePipelineState:state]; @@ -165,7 +165,7 @@ Tensor binaryElementwiseMPSCNNKernel( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); mt.texture()->allocateTemporaryStorage(outputSize, cb1); MPSImage* Y = mt.texture()->image(); - T* kernel = [[T alloc] initWithDevice:[MPSCNNContext sharedInstance].device]; + T* kernel = [[T alloc] initWithDevice:[MetalContext sharedInstance].device]; kernel.primaryStrideInPixelsY = X1.height == 1 ? 0 : 1; kernel.primaryStrideInPixelsX = X1.width == 1 ? 0 : 1; kernel.secondaryStrideInPixelsY = X2.height == 1 ? 0 : 1; @@ -197,7 +197,7 @@ Tensor binaryElementwiseMPSCNNKernel( TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); MPSImage* Y = createTemporaryImage(cb1, outputSize.vec()); - T* kernel = [[T alloc] initWithDevice:[MPSCNNContext sharedInstance].device]; + T* kernel = [[T alloc] initWithDevice:[MetalContext sharedInstance].device]; kernel.primaryStrideInPixelsY = X1.height == 1 ? 0 : 1; kernel.primaryStrideInPixelsX = X1.width == 1 ? 0 : 1; kernel.secondaryStrideInPixelsY = X2.height == 1 ? 0 : 1; diff --git a/aten/src/ATen/native/metal/ops/MetalChunk.mm b/aten/src/ATen/native/metal/ops/MetalChunk.mm index 04a73f78b86b5..89d8d3647e93c 100644 --- a/aten/src/ATen/native/metal/ops/MetalChunk.mm +++ b/aten/src/ATen/native/metal/ops/MetalChunk.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import #import @@ -37,7 +37,7 @@ mt2.texture()->allocateTemporaryStorage(outputSize2, commandBuffer); MPSImage* Y1 = mt1.texture()->image(); MPSImage* Y2 = mt2.texture()->image(); - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:"split_channels" Constants:@[ @(X.featureChannels), diff --git a/aten/src/ATen/native/metal/ops/MetalConcat.mm b/aten/src/ATen/native/metal/ops/MetalConcat.mm index f9fc1f834b065..7e143dac60910 100644 --- a/aten/src/ATen/native/metal/ops/MetalConcat.mm +++ b/aten/src/ATen/native/metal/ops/MetalConcat.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import #import @@ -30,10 +30,10 @@ Tensor cat_batch(const TensorList tensors, MetalTensorImplStorage& mt) { @"inputs have different Metal command buffers"); id encoder = [commandBuffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] pipelineState:mpscnn::kernelFor( X, "copy_offset", "copy_offset_nonarray")]; - id offsetBuffer = [[MPSCNNContext sharedInstance].device + id offsetBuffer = [[MetalContext sharedInstance].device newBufferWithLength:1 * sizeof(ushort) options:MTLResourceOptionCPUCacheModeWriteCombined]; ushort* offsetBufferPtr = (ushort*)[offsetBuffer contents]; @@ -85,8 +85,8 @@ Tensor cat_feature(const TensorList tensors, MetalTensorImplStorage& mt) { } id state = - [[MPSCNNContext sharedInstance] pipelineState:kernelString]; - id offsetBuffer = [[MPSCNNContext sharedInstance].device + [[MetalContext sharedInstance] pipelineState:kernelString]; + id offsetBuffer = [[MetalContext sharedInstance].device newBufferWithLength:5 * sizeof(ushort) options:MTLResourceOptionCPUCacheModeWriteCombined]; ushort* offsetBufferPtr = (ushort*)[offsetBuffer contents]; diff --git a/aten/src/ATen/native/metal/ops/MetalCopy.mm b/aten/src/ATen/native/metal/ops/MetalCopy.mm index ac5d174f07739..a51eab9693af3 100644 --- a/aten/src/ATen/native/metal/ops/MetalCopy.mm +++ b/aten/src/ATen/native/metal/ops/MetalCopy.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -28,7 +28,7 @@ Tensor copy_to_host(const Tensor& input) { id encoder = [commandBuffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:metal::mpscnn::kernelFor( X, "copy", "copy_nonarray") Constants:@[ diff --git a/aten/src/ATen/native/metal/ops/MetalHardswish.mm b/aten/src/ATen/native/metal/ops/MetalHardswish.mm index 234e80256e2cb..66bf36230eead 100644 --- a/aten/src/ATen/native/metal/ops/MetalHardswish.mm +++ b/aten/src/ATen/native/metal/ops/MetalHardswish.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import #import @@ -23,7 +23,7 @@ MPSImage* Y = createTemporaryImage(commandBuffer, imageSize); id encoder = [commandBuffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:mpscnn::kernelFor( X, "hardswish", "hardswish_nonarray") Constants:@[ diff --git a/aten/src/ATen/native/metal/ops/MetalNeurons.mm b/aten/src/ATen/native/metal/ops/MetalNeurons.mm index fa72db1afcdb0..b095d33288bc2 100644 --- a/aten/src/ATen/native/metal/ops/MetalNeurons.mm +++ b/aten/src/ATen/native/metal/ops/MetalNeurons.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import #import diff --git a/aten/src/ATen/native/metal/ops/MetalPadding.mm b/aten/src/ATen/native/metal/ops/MetalPadding.mm index 4ab0b99bdbc93..2610790e169d1 100644 --- a/aten/src/ATen/native/metal/ops/MetalPadding.mm +++ b/aten/src/ATen/native/metal/ops/MetalPadding.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -55,7 +55,7 @@ Tensor reflection_pad2d(const Tensor& input, IntArrayRef padding) { id encoder = [commandBuffer.buffer computeCommandEncoder]; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:"reflection_pad2d" Constants:@[ @(Y.height), diff --git a/aten/src/ATen/native/metal/ops/MetalPooling.mm b/aten/src/ATen/native/metal/ops/MetalPooling.mm index 35c6c34ffc6dd..945fc844d43ce 100644 --- a/aten/src/ATen/native/metal/ops/MetalPooling.mm +++ b/aten/src/ATen/native/metal/ops/MetalPooling.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -49,7 +49,7 @@ Tensor max_pool2d( } MPSImage* X = imageFromTensor(input); MPSCNNPoolingMax* pool = [[MPSCNNPoolingMax alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device kernelWidth:kernel_size[0] kernelHeight:kernel_size[1] strideInPixelsX:stride[0] @@ -82,7 +82,7 @@ Tensor adaptive_avg_pool2d(const Tensor& input, IntArrayRef output_size) { } MPSImage* X = imageFromTensor(input); MPSCNNPoolingAverage* pool = [[MPSCNNPoolingAverage alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device kernelWidth:X.width kernelHeight:X.height strideInPixelsX:X.width diff --git a/aten/src/ATen/native/metal/ops/MetalReduce.mm b/aten/src/ATen/native/metal/ops/MetalReduce.mm index 5dfe974efb0b3..29a4bd9dc579e 100644 --- a/aten/src/ATen/native/metal/ops/MetalReduce.mm +++ b/aten/src/ATen/native/metal/ops/MetalReduce.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import @@ -17,7 +17,7 @@ API_AVAILABLE(ios(11.3), macos(10.13)) static inline MPSNNReduceUnary* kernelForReducedDim(int dim) { - id device = [MPSCNNContext sharedInstance].device; + id device = [MetalContext sharedInstance].device; if (dim == 3) { return [[MPSNNReduceRowMean alloc] initWithDevice:device]; } else if (dim == 2) { diff --git a/aten/src/ATen/native/metal/ops/MetalReshape.mm b/aten/src/ATen/native/metal/ops/MetalReshape.mm index 15ae3c50d60e4..64b3c8d8e7659 100644 --- a/aten/src/ATen/native/metal/ops/MetalReshape.mm +++ b/aten/src/ATen/native/metal/ops/MetalReshape.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -36,7 +36,7 @@ Tensor view(const Tensor& input, IntArrayRef size) { mt.texture()->allocateTemporaryStorage(inferred_size, commandBuffer); MPSImage* Y = mt.texture()->image(); id state = - [[MPSCNNContext sharedInstance] specializedPipelineState:"reshape" + [[MetalContext sharedInstance] specializedPipelineState:"reshape" Constants:@[ @(Y.height), @(Y.width), diff --git a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm index f3e1ef5e362fe..181d1fff9ae34 100644 --- a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm +++ b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import @@ -37,7 +37,7 @@ Tensor mpscnn_softmax( MPSImage* X = imageFromTensor(input_); // MPSCNNSoftmax kernels operate on feature channels // https://developer.apple.com/documentation/metalperformanceshaders/mpscnnsoftmax?changes=_1&language=objc - T* softmax = [[T alloc] initWithDevice:[MPSCNNContext sharedInstance].device]; + T* softmax = [[T alloc] initWithDevice:[MetalContext sharedInstance].device]; MetalTensorImplStorage mt{newSize}; MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input_); mt.texture()->allocateTemporaryStorage(newSize, commandBuffer); diff --git a/aten/src/ATen/native/metal/ops/MetalTranspose.mm b/aten/src/ATen/native/metal/ops/MetalTranspose.mm index d10a74eed3011..3adb0e0c1bc89 100644 --- a/aten/src/ATen/native/metal/ops/MetalTranspose.mm +++ b/aten/src/ATen/native/metal/ops/MetalTranspose.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -33,7 +33,7 @@ Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { mt.texture()->allocateTemporaryStorage(outputSizes, commandBuffer); MPSImage* Y = mt.texture()->image(); MPSImageTranspose* transpose = [[MPSImageTranspose alloc] - initWithDevice:[MPSCNNContext sharedInstance].device]; + initWithDevice:[MetalContext sharedInstance].device]; [transpose encodeToCommandBuffer:commandBuffer.buffer sourceImage:X destinationImage:Y]; @@ -50,7 +50,7 @@ Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { id encoder = [commandBuffer.buffer computeCommandEncoder]; id state = - [[MPSCNNContext sharedInstance] specializedPipelineState:"transpose" + [[MetalContext sharedInstance] specializedPipelineState:"transpose" Constants:@[ @(dim0), @(dim1), diff --git a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm index 25138eddddae8..049aefec168fa 100644 --- a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm +++ b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -51,7 +51,7 @@ Tensor upsample_nearest2d_vec( MPSImage* Y = mt.texture()->image(); if (@available(iOS 11.0, *)) { MPSCNNUpsamplingNearest* kernel = [[MPSCNNUpsamplingNearest alloc] - initWithDevice:[MPSCNNContext sharedInstance].device + initWithDevice:[MetalContext sharedInstance].device integerScaleFactorX:(NSUInteger)scale_w.value() integerScaleFactorY:(NSUInteger)scale_h.value()]; [kernel encodeToCommandBuffer:commandBuffer.buffer @@ -60,7 +60,7 @@ Tensor upsample_nearest2d_vec( } else { NSUInteger sh = scale_h.value() * 10000; NSUInteger sw = scale_w.value() * 10000; - id state = [[MPSCNNContext sharedInstance] + id state = [[MetalContext sharedInstance] specializedPipelineState:mpscnn::kernelFor( Y, "resize_nearest", From b9cd97c94bf558463bac1ca2e63d94d77e9d0b02 Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Fri, 18 Jun 2021 18:43:11 -0700 Subject: [PATCH 255/305] [iOS GPU][BE][2/n] Remove unused APIs (#60281) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60281 1. REmove unused APIs from MPSImageUtils. 2. Move tensor related APIs from MetalUtils to MetalTensorUtils. Delete MetalUtils.h/mm 3. Move metal buffer related APIs to MetalContext ghstack-source-id: 131839559 Test Plan: 1. CircleCI 2. buck test pp-mac Reviewed By: SS-JIA Differential Revision: D29232973 fbshipit-source-id: a4c0c848883b8ef615eeb2936c1f3d18cddcb318 --- aten/src/ATen/native/metal/MetalAten.mm | 2 +- aten/src/ATen/native/metal/MetalContext.h | 1 + aten/src/ATen/native/metal/MetalContext.mm | 27 +++-- .../native/metal/MetalTensorImplStorage.mm | 4 +- aten/src/ATen/native/metal/MetalTensorUtils.h | 61 +++++++++++ aten/src/ATen/native/metal/MetalUtils.h | 97 ----------------- aten/src/ATen/native/metal/MetalUtils.mm | 100 ----------------- .../ATen/native/metal/mpscnn/MPSCNNClampOp.mm | 2 +- .../ATen/native/metal/mpscnn/MPSImageUtils.h | 12 +- .../ATen/native/metal/mpscnn/MPSImageUtils.mm | 103 +----------------- .../native/metal/mpscnn/MPSImageWrapper.mm | 2 +- .../native/metal/mpscnn/tests/MPSCNNTests.h | 1 - .../native/metal/mpscnn/tests/MPSCNNTests.mm | 33 +----- .../metal/mpscnn/tests/MetalOpTestRunner.mm | 15 +-- aten/src/ATen/native/metal/ops/MetalAddmm.mm | 6 +- .../metal/ops/MetalBinaryElementwise.mm | 16 +-- aten/src/ATen/native/metal/ops/MetalChunk.mm | 4 +- aten/src/ATen/native/metal/ops/MetalClamp.mm | 4 +- aten/src/ATen/native/metal/ops/MetalConcat.mm | 12 +- .../ATen/native/metal/ops/MetalConvolution.mm | 6 +- aten/src/ATen/native/metal/ops/MetalCopy.mm | 4 +- .../ATen/native/metal/ops/MetalHardswish.mm | 4 +- .../src/ATen/native/metal/ops/MetalNeurons.mm | 6 +- .../src/ATen/native/metal/ops/MetalPadding.mm | 4 +- .../src/ATen/native/metal/ops/MetalPooling.mm | 6 +- aten/src/ATen/native/metal/ops/MetalReduce.mm | 4 +- .../src/ATen/native/metal/ops/MetalReshape.mm | 4 +- .../src/ATen/native/metal/ops/MetalSoftmax.mm | 4 +- .../ATen/native/metal/ops/MetalTranspose.mm | 18 ++- .../metal/ops/MetalUpsamplingNearest.mm | 4 +- 30 files changed, 156 insertions(+), 410 deletions(-) delete mode 100644 aten/src/ATen/native/metal/MetalUtils.h delete mode 100644 aten/src/ATen/native/metal/MetalUtils.mm diff --git a/aten/src/ATen/native/metal/MetalAten.mm b/aten/src/ATen/native/metal/MetalAten.mm index 2677fe5a96eee..f1cbdaf5d3a7f 100644 --- a/aten/src/ATen/native/metal/MetalAten.mm +++ b/aten/src/ATen/native/metal/MetalAten.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #include #include diff --git a/aten/src/ATen/native/metal/MetalContext.h b/aten/src/ATen/native/metal/MetalContext.h index c459537799855..ca58eb9a433a0 100644 --- a/aten/src/ATen/native/metal/MetalContext.h +++ b/aten/src/ATen/native/metal/MetalContext.h @@ -16,5 +16,6 @@ API_AVAILABLE(ios(10.0), macos(10.13)) - (id)specializedPipelineState:(const std::string&)kernel Constants:(NSArray*) constants; +- (id)emptyMTLBuffer:(int64_t) size; @end diff --git a/aten/src/ATen/native/metal/MetalContext.mm b/aten/src/ATen/native/metal/MetalContext.mm index 064e59efe2c19..80ee55efa591e 100644 --- a/aten/src/ATen/native/metal/MetalContext.mm +++ b/aten/src/ATen/native/metal/MetalContext.mm @@ -65,7 +65,7 @@ - (BOOL)available { #else return false; #endif - NSError* error = [self compileProgram]; + NSError* error = [self _compileProgram]; if (error) { std::string compilationError = error.localizedDescription.UTF8String; std::string deviceInfo = self.description.UTF8String; @@ -139,7 +139,22 @@ - (BOOL)available { return state; } -- (NSError*)compileProgram { +- (id)emptyMTLBuffer:(int64_t) size { + TORCH_CHECK(_device); + id buffer = [_device newBufferWithLength:size + options:MTLResourceOptionCPUCacheModeWriteCombined]; + return buffer; +} + +- (NSString*)description { + NSString* desc = + [NSString stringWithFormat:@"DeviceName: %s, LanguageVersion: %lu", + _deviceInfo.name.c_str(), + (unsigned long)_deviceInfo.languageVersion]; + return desc; +} + +- (NSError*)_compileProgram { __block NSError* compilationError = nil; static dispatch_once_t onceToken; dispatch_once(&onceToken, ^{ @@ -156,12 +171,6 @@ - (NSError*)compileProgram { return compilationError; } -- (NSString*)description { - NSString* desc = - [NSString stringWithFormat:@"DeviceName: %s, LanguageVersion: %lu", - _deviceInfo.name.c_str(), - (unsigned long)_deviceInfo.languageVersion]; - return desc; -} + @end diff --git a/aten/src/ATen/native/metal/MetalTensorImplStorage.mm b/aten/src/ATen/native/metal/MetalTensorImplStorage.mm index 91b336cede615..cd73ba4eddb3c 100644 --- a/aten/src/ATen/native/metal/MetalTensorImplStorage.mm +++ b/aten/src/ATen/native/metal/MetalTensorImplStorage.mm @@ -1,6 +1,6 @@ #import #import -#import +#import #import #include @@ -49,7 +49,7 @@ void copy_data_to_host(float* host) { MetalTensorImplStorage::MetalTensorImplStorage( const std::vector& sizes) - : MetalTensorImplStorage(sizes, compute_strides(sizes)) {} + : MetalTensorImplStorage(sizes, computeStrides(sizes)) {} MetalTensorImplStorage::MetalTensorImplStorage( const std::vector& sizes, diff --git a/aten/src/ATen/native/metal/MetalTensorUtils.h b/aten/src/ATen/native/metal/MetalTensorUtils.h index aaa143432bd5a..318da09d86b28 100644 --- a/aten/src/ATen/native/metal/MetalTensorUtils.h +++ b/aten/src/ATen/native/metal/MetalTensorUtils.h @@ -1,4 +1,14 @@ #include +#include +#include +#include +#include + +#if (defined(__ARM_NEON__) || defined(__ARM_NEON)) +typedef float16_t fp16_t; +#else +typedef uint16_t fp16_t; +#endif namespace at { namespace native { @@ -9,6 +19,57 @@ uint32_t channelsSize(const Tensor& tensor); uint32_t heightSize(const Tensor& tensor); uint32_t widthSize(const Tensor& tensor); +// When copying the result back to a CPU tensor, the memory format becomes NCHW. +// Thus,we compute the strides based on contiguous memory format. +static inline std::vector computeStrides( + const std::vector& sizes) { + const auto dim = sizes.size(); + std::vector strides(dim, 0); + if (dim > 0) { + const auto last_idx = dim - 1; + strides[last_idx] = 1; + for (int64_t i = last_idx - 1; i >= 0; --i) { + strides[i] = strides[i + 1] * std::max(sizes[i + 1], 1); + } + } + return strides; +} + +static inline MetalTensorImplStorage& getTensorImplStorage( + const at::Tensor& tensor) { + using MetalTensorImpl = at::MetalTensorImpl; + TORCH_CHECK(tensor.is_metal()); + MetalTensorImpl* impl = + static_cast(tensor.unsafeGetTensorImpl()); + return impl->unsafe_opaque_handle(); +} + +static inline at::Tensor makeTensor( + MetalTensorImplStorage&& mt, + const TensorOptions& options) { + using MetalTensorImpl = at::MetalTensorImpl; + auto sizes = mt.sizes(); // sizes is stored in TensorImpl + auto strides = mt.strides(); // strides is stored in MetalTensorImpl + return detail::make_tensor( + DispatchKeySet(DispatchKey::Metal), + options.dtype(), + at::Device(at::kMetal), + std::move(mt), + std::vector(sizes.begin(), sizes.end()), + std::vector(strides.begin(), strides.end())); +} + +static inline MetalCommandBuffer* getCommandBuffer( + const Tensor& tensor) { + TORCH_CHECK(tensor.is_metal()); + auto implStorage = getTensorImplStorage(tensor); + MetalCommandBuffer* cmdBuffer = implStorage.texture()->commandBuffer(); + if (!cmdBuffer || !cmdBuffer.valid) { + cmdBuffer = [MetalCommandBuffer currentBuffer]; + } + return cmdBuffer; +} + } // namespace metal } // namespace native } // namespace at diff --git a/aten/src/ATen/native/metal/MetalUtils.h b/aten/src/ATen/native/metal/MetalUtils.h deleted file mode 100644 index a94fce6f76e33..0000000000000 --- a/aten/src/ATen/native/metal/MetalUtils.h +++ /dev/null @@ -1,97 +0,0 @@ -#include -#include -#include -#include -#include -#include - -#if (defined(__ARM_NEON__) || defined(__ARM_NEON)) -typedef float16_t fp16_t; -#else -typedef uint16_t fp16_t; -#endif - -namespace at { -namespace native { -namespace metal { - -std::vector Fp32ToFp16(const std::vector& src); -std::vector Fp16ToFp32(const std::vector& src); - -std::vector NCHWToNC4( - const float* src, - const std::vector& sizes); -std::vector NC4ToNCHW( - const float* src, - const std::vector& sizes); - -// When copying the result back to a CPU tensor, the memory format becomes NCHW. -// Thus,we compute the strides based on contiguous memory format. -static inline std::vector compute_strides( - const std::vector& sizes) { - const auto dim = sizes.size(); - std::vector strides(dim, 0); - if (dim > 0) { - const auto last_idx = dim - 1; - strides[last_idx] = 1; - for (int i = last_idx - 1; i >= 0; --i) { - strides[i] = strides[i + 1] * std::max(sizes[i + 1], 1); - } - } - return strides; -} - -static inline MetalTensorImplStorage& getTensorImplStorage( - const at::Tensor& tensor) { - using MetalTensorImpl = at::MetalTensorImpl; - TORCH_CHECK(tensor.is_metal()); - MetalTensorImpl* impl = - static_cast(tensor.unsafeGetTensorImpl()); - return impl->unsafe_opaque_handle(); -} - -static inline at::Tensor makeTensor( - MetalTensorImplStorage&& mt, - const TensorOptions& options) { - using MetalTensorImpl = at::MetalTensorImpl; - auto sizes = mt.sizes(); // sizes is stored in TensorImpl - auto strides = mt.strides(); // strides is stored in MetalTensorImpl - return detail::make_tensor( - DispatchKeySet(DispatchKey::Metal), - options.dtype(), - at::Device(at::kMetal), - std::move(mt), - std::vector(sizes.begin(), sizes.end()), - std::vector(strides.begin(), strides.end())); -} - -static inline MetalCommandBuffer* getCommandBufferFromTensor( - const Tensor& tensor) { - TORCH_CHECK(tensor.is_metal()); - auto implStorage = getTensorImplStorage(tensor); - MetalCommandBuffer* cmdBuffer = implStorage.texture()->commandBuffer(); - if (!cmdBuffer || !cmdBuffer.valid) { - cmdBuffer = [MetalCommandBuffer currentBuffer]; - } - return cmdBuffer; -} - -template -idmakeMTLBuffer(const std::vector& src) { - id buffer = [[MetalContext sharedInstance].device - newBufferWithLength:src.size() * sizeof(T) - options:MTLResourceOptionCPUCacheModeWriteCombined]; - memcpy(buffer.contents, src.data(), src.size() * sizeof(T)); - return buffer; -} - -static inline idmakeMTLBuffer(int64_t bytes) { - id buffer = [[MetalContext sharedInstance].device - newBufferWithLength:bytes - options:MTLResourceOptionCPUCacheModeWriteCombined]; - return buffer; -} - -} // namespace metal -} // namespace native -} // namespace at diff --git a/aten/src/ATen/native/metal/MetalUtils.mm b/aten/src/ATen/native/metal/MetalUtils.mm deleted file mode 100644 index a082c153a235f..0000000000000 --- a/aten/src/ATen/native/metal/MetalUtils.mm +++ /dev/null @@ -1,100 +0,0 @@ -#import -#import -#import - -namespace at { -namespace native { -namespace metal { - -std::vector Fp32ToFp16(const std::vector& src) { - unsigned long count = src.size(); - std::vector output(count, 0); - vImage_Buffer float32{(void*)src.data(), 1, count, count * sizeof(float)}; - vImage_Buffer float16{(void*)output.data(), 1, count, count * sizeof(fp16_t)}; - if (vImageConvert_PlanarFtoPlanar16F(&float32, &float16, 0) != - kvImageNoError) { - TORCH_CHECK(false); - } - return output; -} - -std::vector Fp16ToFp32(const std::vector& src) { - unsigned long count = src.size(); - std::vector output(count, 0); - vImage_Buffer float16{(void*)src.data(), 1, count, count * sizeof(fp16_t)}; - vImage_Buffer float32{(void*)output.data(), 1, count, count * sizeof(float)}; - if (vImageConvert_Planar16FtoPlanarF(&float16, &float32, 0) != - kvImageNoError) { - TORCH_CHECK(false); - } - return output; -} - -std::vector NCHWToNC4( - const float* src, - const std::vector& sizes) { - int64_t N = sizes[0]; - int64_t C = sizes[1]; - int64_t H = sizes[2]; - int64_t W = sizes[3]; - int64_t src_image_count = C * H * W; - int64_t src_count = N * src_image_count; - int64_t slices = (C + 3) / 4; - int64_t numComponents = C < 3 ? C : 4; - int64_t dst_image_count = slices * numComponents * W * H; - int64_t dst_count = N * dst_image_count; - std::vector output(dst_count, 0.0f); - for (int n = 0; n < N; ++n) { - int64_t src_image = n * src_image_count; - int64_t dst_image = n * dst_image_count; - for (int i = 0; i < slices; ++i) { - int64_t slice = i * W * H * numComponents; - for (int j = 0; j < W * H; ++j) { - for (int k = 0; k < numComponents; ++k) { - int ii = src_image + slice + k * W * H + j; - int oi = dst_image + slice + j * numComponents + k; - if (k < C && ii < src_count) { - output[oi] = src[ii]; - } - } - } - } - } - return output; -} - -std::vector NC4ToNCHW( - const float* src, - const std::vector& sizes) { - int64_t N = sizes[0]; - int64_t C = sizes[1]; - int64_t H = sizes[2]; - int64_t W = sizes[3]; - int64_t slices = (C + 3) / 4; - int64_t numComponents = C < 3 ? C : 4; - int64_t src_image_count = slices * numComponents * W * H; - int64_t dst_image_count = C * H * W; - int64_t dst_count = N * dst_image_count; - std::vector output(dst_count, 0.0f); - for (int n = 0; n < N; ++n) { - int64_t src_image = n * src_image_count; - int64_t dst_image = n * dst_image_count; - for (int i = 0; i < slices; ++i) { - int64_t slice = i * W * H * numComponents; - for (int j = 0; j < numComponents; ++j) { - for (int k = 0; k < W * H; ++k) { - int ii = src_image + slice + k * numComponents + j; - int oi = dst_image + slice + j * W * H + k; - if (j < C && oi < dst_count) { - output[oi] = src[ii]; - } - } - } - } - } - return output; -} - -} -} -} diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm index 2d380dd3cda1e..93218bd5e0bcf 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm @@ -1,4 +1,4 @@ -#import +#import #import #import #import diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h index 4ef2247c938ea..53065e4be3fa8 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import @@ -10,14 +10,9 @@ namespace native { namespace metal { MPSImage* createStaticImage(IntArrayRef sizes); -MPSImage* createStaticImage( - const fp16_t* src, - const IntArrayRef sizes); MPSImage* createStaticImage( const float* src, const IntArrayRef sizes); -MPSImage* createStaticImage(const at::Tensor& tensor); -MPSImage* createStaticImage(MPSImage* image); MPSImage* createStaticImage( MPSTemporaryImage* image, MetalCommandBuffer* buffer, @@ -37,9 +32,6 @@ MPSTemporaryImage* createTemporaryImage( void copyToHost(float* dst, MPSImage* image); void copyToMetalBuffer(MetalCommandBuffer* buffer, id dst, MPSImage* image); -std::vector staticImageToFp16Array(MPSImage* image); -at::Tensor staticImageToTensor(MPSImage* image); - static inline MPSImage* imageFromTensor(const Tensor& tensor) { TORCH_CHECK(tensor.is_metal()); using MetalTensorImplStorage = at::native::metal::MetalTensorImplStorage; @@ -63,7 +55,7 @@ static inline std::vector computeImageSize(IntArrayRef sizes) { std::vector imageSize(4, 1); int64_t index = 3; int64_t batch = 1; - for (int i = sizes.size() - 1; i >= 0; i--) { + for (int64_t i = sizes.size() - 1; i >= 0; i--) { if (index != 0) { imageSize[index] = sizes[i]; index--; diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm index 8d2f17163780c..817672a0f5524 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm @@ -1,4 +1,4 @@ -#import +#import #import #import #import @@ -24,39 +24,6 @@ imageDescriptor:desc]; } -MPSImage* createStaticImage(const fp16_t* src, IntArrayRef sizes) { - int64_t N = sizes[0]; - int64_t C = sizes[1]; - int64_t H = sizes[2]; - int64_t W = sizes[3]; - MPSImageDescriptor* desc = [MPSImageDescriptor - imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 - width:W - height:H - featureChannels:C - numberOfImages:N - usage:MTLTextureUsageShaderRead | - MTLTextureUsageShaderWrite]; - MPSImage* image = - [[MPSImage alloc] initWithDevice:[MetalContext sharedInstance].device - imageDescriptor:desc]; - - int64_t slices = (C + 3) / 4 * N; - int64_t numComponents = image.featureChannels < 3 ? image.featureChannels : 4; - int64_t bytesPerRow = W * numComponents * sizeof(fp16_t); - uint8_t* ptr = (uint8_t*)src; - for (int i = 0; i < slices; ++i) { - [image.texture replaceRegion:MTLRegionMake2D(0, 0, W, H) - mipmapLevel:0 - slice:i - withBytes:ptr - bytesPerRow:bytesPerRow - bytesPerImage:0]; - ptr += H * bytesPerRow; - } - return image; -} - MPSImage* createStaticImage(const float* src, IntArrayRef sizes) { int64_t size_bytes = c10::multiply_integers(sizes) * sizeof(float); id buff = [[MetalContext sharedInstance].device @@ -88,36 +55,6 @@ return output; } -MPSImage* createStaticImage(const at::Tensor& tensor) { - TORCH_CHECK(tensor.device().is_cpu()); - TORCH_CHECK(tensor.dim() == 4); - auto contiguousTensor = tensor.contiguous(); - float* src = tensor.data_ptr(); - std::vector sizes = tensor.sizes().vec(); - auto c4 = NCHWToNC4(src, sizes); - auto c4fp16 = Fp32ToFp16(c4); - return createStaticImage(c4fp16.data(), sizes); -} - -MPSImage* createStaticImage(MPSImage* image) { - MPSImage* Y = createStaticImage([image sizes]); - MetalCommandBuffer* cb = [MetalCommandBuffer newBuffer]; - id encoder = [cb.buffer computeCommandEncoder]; - id state = [[MetalContext sharedInstance] - pipelineState:mpscnn::kernelFor(image, "copy", "copy_nonarray")]; - [encoder setComputePipelineState:state]; - [encoder setTexture:[image texture] atIndex:0]; - [encoder setTexture:[Y texture] atIndex:1]; - - const auto& launchParams = - mpscnn::spatialPointwiseKernelLaunchParams(state, image); - [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid - threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; - [encoder endEncoding]; - [cb commit]; - return Y; -} - MPSImage* createStaticImage( MPSTemporaryImage* image, MetalCommandBuffer* buffer, @@ -276,44 +213,6 @@ void copyToMetalBuffer( [encoder endEncoding]; } -std::vector staticImageToFp16Array(MPSImage* image) { - if (image.pixelFormat == MTLPixelFormatR16Float || - image.pixelFormat == MTLPixelFormatRG16Float || - image.pixelFormat == MTLPixelFormatRGBA16Float) { - int64_t slices = (image.featureChannels + 3) / 4; - int64_t C = image.featureChannels < 3 ? image.featureChannels : slices * 4; - int64_t numComponents = - image.featureChannels < 3 ? image.featureChannels : 4; - int64_t count = image.width * image.height * image.numberOfImages * C; - std::vector output(count, 0); - int64_t bytesPerRow = image.width * numComponents * sizeof(fp16_t); - uint8_t* buffer = (uint8_t*)output.data(); - for (int i = 0; i < slices * image.numberOfImages; ++i) { - [image.texture getBytes:buffer - bytesPerRow:bytesPerRow - bytesPerImage:0 - fromRegion:MTLRegionMake2D(0, 0, image.width, image.height) - mipmapLevel:0 - slice:i]; - buffer += image.height * bytesPerRow; - } - return output; - } - TORCH_CHECK( - false, "Copy to float buffer failed: The pixel format didn't match"); -} - -at::Tensor staticImageToTensor(MPSImage* image) { - auto outputSize = [image sizes]; - std::vector fp16Array = staticImageToFp16Array(image); - auto fp32 = metal::Fp16ToFp32(fp16Array); - std::vector fp32_nchw = metal::NC4ToNCHW(fp32.data(), outputSize); - auto tensor = at::empty(outputSize); - int64_t size_bytes = c10::multiply_integers(outputSize) * sizeof(float); - memcpy(tensor.data_ptr(), fp32_nchw.data(), size_bytes); - return tensor; -} - } } } diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm index 10f396713e2a6..1d94bdd554a25 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm @@ -1,5 +1,5 @@ #import -#import +#import #import #import #import diff --git a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.h b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.h index 95e8f16448373..57c1c8a5f2c6b 100644 --- a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.h +++ b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.h @@ -2,7 +2,6 @@ #define MPSCNNTests_h bool test_synchronization(); -bool test_nchw_to_nc4_cpu(); bool test_copy_nchw_to_metal(); bool test_conv2d(); bool test_depthwiseConv(); diff --git a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm index f21fb1390387e..5df31eb320d47 100644 --- a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm +++ b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm @@ -1,5 +1,5 @@ #import -#import +#import #import #import #import @@ -86,13 +86,7 @@ void PRINT_TENSOR(std::string name, const at::Tensor& tensor) { } std::cout << str << std::endl; }; - if (tensor.is_metal()) { - MPSImage* image = at::native::metal::imageFromTensor(tensor); - auto t = at::native::metal::staticImageToTensor(image); - print(t); - } else { - print(tensor); - } + print(tensor); } } @@ -111,29 +105,6 @@ bool test_synchronization() { }); } -bool test_nchw_to_nc4_cpu() { - bool result = true; - for (int i = 0; i < ITER_COUNT; ++i) { - int64_t N = rand(1, 24); - int64_t C = rand(1, 48); - int64_t H = rand(1, 320); - int64_t W = rand(1, 320); - __block std::vector size{N, C, H, W}; - bool b = TEST(size, __PRETTY_FUNCTION__, ^bool { - auto t = at::rand(size, at::TensorOptions(at::kCPU).dtype(at::kFloat)); - const auto len = c10::multiply_integers(std::begin(size), std::end(size)); - auto buf = - std::vector{t.data_ptr(), t.data_ptr() + len}; - auto c4 = NCHWToNC4((float*)t.data_ptr(), t.sizes().vec()); - auto n4 = NC4ToNCHW((float*)c4.data(), t.sizes().vec()); - return n4 == buf; - }); - if (!b) { - result = false; - } - } - return result; -} bool test_copy_nchw_to_metal() { __block std::vector size{1, 3, 224, 224}; diff --git a/aten/src/ATen/native/metal/mpscnn/tests/MetalOpTestRunner.mm b/aten/src/ATen/native/metal/mpscnn/tests/MetalOpTestRunner.mm index eb36e4de30708..5c9ecfb3fce0b 100644 --- a/aten/src/ATen/native/metal/mpscnn/tests/MetalOpTestRunner.mm +++ b/aten/src/ATen/native/metal/mpscnn/tests/MetalOpTestRunner.mm @@ -6,13 +6,12 @@ #import @implementation MetalOpTestRunner { - NSMutableDictionary *_tests; + NSMutableDictionary* _tests; } -+ (instancetype)sharedInstance -{ ++ (instancetype)sharedInstance { static dispatch_once_t onceToken; - static MetalOpTestRunner *instance = nil; + static MetalOpTestRunner* instance = nil; dispatch_once(&onceToken, ^{ instance = [MetalOpTestRunner new]; }); @@ -29,9 +28,11 @@ - (instancetype)init { - (void)registerTests { _tests = [NSMutableDictionary dictionary]; -#define REG_TEST(arg1, arg2) _tests[@arg1] = ^BOOL(void){return arg2();} +#define REG_TEST(arg1, arg2) \ + _tests[@arg1] = ^BOOL(void) { \ + return arg2(); \ + } REG_TEST("test_synchronization", test_synchronization); - REG_TEST("test_nchw_to_nc4_cpu", test_nchw_to_nc4_cpu); REG_TEST("test_copy_nchw_to_metal", test_copy_nchw_to_metal); REG_TEST("test_conv2d", test_conv2d); REG_TEST("test_depthwiseConv", test_depthwiseConv); @@ -81,7 +82,7 @@ - (void)registerTests { REG_TEST("test_reflection_pad2d", test_reflection_pad2d); } -- (NSDictionary *) tests { +- (NSDictionary*)tests { return _tests; } diff --git a/aten/src/ATen/native/metal/ops/MetalAddmm.mm b/aten/src/ATen/native/metal/ops/MetalAddmm.mm index c023d91f7c1a1..8086b17eba476 100644 --- a/aten/src/ATen/native/metal/ops/MetalAddmm.mm +++ b/aten/src/ATen/native/metal/ops/MetalAddmm.mm @@ -3,7 +3,7 @@ #import #import #import -#import +#import #import #import #import @@ -45,7 +45,7 @@ Tensor addmm( auto packedWeights = weight_.contiguous(c10::MemoryFormat::ChannelsLast); MetalTensorImplStorage mt{{params.N, params.OC}}; SmallVector textureSize = {params.N, params.OC, 1, 1}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input_); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input_); mt.texture()->allocateTemporaryStorage(textureSize, commandBuffer); MPSImage* Y = mt.texture()->image(); float* w = packedWeights.data_ptr(); @@ -101,7 +101,7 @@ Tensor linear(const Tensor& input, LinearOpContext& context) { } MetalTensorImplStorage mt{{params.N, params.OC}}; SmallVector textureSize = {params.N, params.OC, 1, 1}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input_); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input_); mt.texture()->allocateTemporaryStorage(textureSize, commandBuffer); MPSImage* Y1 = mt.texture()->image(); // HACK alert: diff --git a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm index 294913d3244f3..97294a11b9be5 100644 --- a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm +++ b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm @@ -76,8 +76,8 @@ Tensor binaryElementwiseShaderKernel( return makeTensor({outputSize.vec()}, input1.options()); } MetalTensorImplStorage mt{outputSize.vec()}; - MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); - MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); + MetalCommandBuffer* cb1 = getCommandBuffer(input1); + MetalCommandBuffer* cb2 = getCommandBuffer(input2); TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); mt.texture()->allocateTemporaryStorage(outputSize, cb1); @@ -117,8 +117,8 @@ Tensor binaryElementwiseShaderKernel( if(c10::multiply_integers(outputSize) == 0){ return input1; } - MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); - MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); + MetalCommandBuffer* cb1 = getCommandBuffer(input1); + MetalCommandBuffer* cb2 = getCommandBuffer(input2); TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); MPSImage* Y = createTemporaryImage(cb1, outputSize.vec()); @@ -159,8 +159,8 @@ Tensor binaryElementwiseMPSCNNKernel( return makeTensor({outputSize.vec()}, input1.options()); } MetalTensorImplStorage mt{outputSize.vec()}; - MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); - MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); + MetalCommandBuffer* cb1 = getCommandBuffer(input1); + MetalCommandBuffer* cb2 = getCommandBuffer(input2); TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); mt.texture()->allocateTemporaryStorage(outputSize, cb1); @@ -192,8 +192,8 @@ Tensor binaryElementwiseMPSCNNKernel( if(c10::multiply_integers(outputSize) == 0){ return input1; } - MetalCommandBuffer* cb1 = getCommandBufferFromTensor(input1); - MetalCommandBuffer* cb2 = getCommandBufferFromTensor(input2); + MetalCommandBuffer* cb1 = getCommandBuffer(input1); + MetalCommandBuffer* cb2 = getCommandBuffer(input2); TORCH_CHECK( [cb1 isEqual:cb2], @"inputs have different Metal command buffers"); MPSImage* Y = createTemporaryImage(cb1, outputSize.vec()); diff --git a/aten/src/ATen/native/metal/ops/MetalChunk.mm b/aten/src/ATen/native/metal/ops/MetalChunk.mm index 89d8d3647e93c..3da3e683ee727 100644 --- a/aten/src/ATen/native/metal/ops/MetalChunk.mm +++ b/aten/src/ATen/native/metal/ops/MetalChunk.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -28,7 +28,7 @@ std::vector splits(num_splits); int64_t last_split_size = split_size - (split_size * num_splits - dim_size); MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); auto outputSize1 = {input.size(0), split_size, input.size(2), input.size(3)}; auto outputSize2 = {input.size(0), last_split_size, input.size(2), input.size(3)}; MetalTensorImplStorage mt1(outputSize1); diff --git a/aten/src/ATen/native/metal/ops/MetalClamp.mm b/aten/src/ATen/native/metal/ops/MetalClamp.mm index bf96813cd74e9..23ed28d2401d2 100644 --- a/aten/src/ATen/native/metal/ops/MetalClamp.mm +++ b/aten/src/ATen/native/metal/ops/MetalClamp.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -15,7 +15,7 @@ Tensor& hardtanh_(Tensor& input, const Scalar& min_val, const Scalar& max_val) { TORCH_CHECK(input.is_metal()); MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); MPSImage* Y = createTemporaryImage(commandBuffer, input.sizes().vec()); float min = min_val.toFloat(); float max = max_val.toFloat(); diff --git a/aten/src/ATen/native/metal/ops/MetalConcat.mm b/aten/src/ATen/native/metal/ops/MetalConcat.mm index 7e143dac60910..2b34bc4efef53 100644 --- a/aten/src/ATen/native/metal/ops/MetalConcat.mm +++ b/aten/src/ATen/native/metal/ops/MetalConcat.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -18,13 +18,13 @@ Tensor cat_batch(const TensorList tensors, MetalTensorImplStorage& mt) { at::Tensor tensor = tensors[0]; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(tensor); + MetalCommandBuffer* commandBuffer = getCommandBuffer(tensor); MPSImage* Y = mt.texture()->image(); ushort cat_dim4_pointer = 0; for (int i = 0; i < tensors.size(); ++i) { const auto& t = tensors[i]; MPSImage* X = imageFromTensor(t); - MetalCommandBuffer* Xcb = getCommandBufferFromTensor(t); + MetalCommandBuffer* Xcb = getCommandBuffer(t); TORCH_CHECK( [commandBuffer isEqual:Xcb], @"inputs have different Metal command buffers"); @@ -58,13 +58,13 @@ Tensor cat_batch(const TensorList tensors, MetalTensorImplStorage& mt) { Tensor cat_feature(const TensorList tensors, MetalTensorImplStorage& mt) { at::Tensor tensor = tensors[0]; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(tensor); + MetalCommandBuffer* commandBuffer = getCommandBuffer(tensor); MPSImage* Y = mt.texture()->image(); ushort channel_offset = 0; for (int i = 0; i < tensors.size(); ++i) { const auto& t = tensors[i]; MPSImage* X = imageFromTensor(t); - MetalCommandBuffer* Xcb = getCommandBufferFromTensor(t); + MetalCommandBuffer* Xcb = getCommandBuffer(t); TORCH_CHECK( [commandBuffer isEqual:Xcb], @"inputs have different Metal command buffers"); @@ -124,7 +124,7 @@ Tensor cat(const TensorList tensors, int64_t dim) { "Metal cat is implemented only for batch dimension"); int64_t cat_dim_size = 0; at::Tensor tensor = tensors[0]; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(tensor); + MetalCommandBuffer* commandBuffer = getCommandBuffer(tensor); for (int i = 0; i < tensors.size(); ++i) { const auto& t = tensors[i]; TORCH_CHECK(t.dim() == 4, "Metal cat expects 4 dimensional inputs"); diff --git a/aten/src/ATen/native/metal/ops/MetalConvolution.mm b/aten/src/ATen/native/metal/ops/MetalConvolution.mm index 98fc87e84be73..c726382dde45f 100644 --- a/aten/src/ATen/native/metal/ops/MetalConvolution.mm +++ b/aten/src/ATen/native/metal/ops/MetalConvolution.mm @@ -1,6 +1,6 @@ #import #import -#import +#import #import #import #import @@ -42,7 +42,7 @@ Tensor conv2d( bias:b neuronFilter:NeuronType::None]; MetalTensorImplStorage mt{outputSize}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); MPSImage* Y = mt.texture()->image(); [op encode:commandBuffer.buffer sourceImage:X destinationImage:Y]; @@ -79,7 +79,7 @@ Tensor conv2d(const Tensor& input, Conv2dOpContext& context) { }; } MetalTensorImplStorage mt{outputSize}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); MPSImage* Y1 = mt.texture()->image(); [op encode:commandBuffer.buffer sourceImage:X destinationImage:Y1]; diff --git a/aten/src/ATen/native/metal/ops/MetalCopy.mm b/aten/src/ATen/native/metal/ops/MetalCopy.mm index a51eab9693af3..b6c783b1e1579 100644 --- a/aten/src/ATen/native/metal/ops/MetalCopy.mm +++ b/aten/src/ATen/native/metal/ops/MetalCopy.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -19,7 +19,7 @@ Tensor copy_to_host(const Tensor& input) { if (X && !X.isTemporaryImage) { return input; } - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); auto&& sizes = [X sizes]; MetalTensorImplStorage mt{sizes}; mt.texture()->setCommandBuffer(commandBuffer); diff --git a/aten/src/ATen/native/metal/ops/MetalHardswish.mm b/aten/src/ATen/native/metal/ops/MetalHardswish.mm index 66bf36230eead..f446b4c49bcc4 100644 --- a/aten/src/ATen/native/metal/ops/MetalHardswish.mm +++ b/aten/src/ATen/native/metal/ops/MetalHardswish.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -17,7 +17,7 @@ Tensor& hardswish_(Tensor& input) { MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); IntArrayRef outputSize = input.sizes(); std::vector imageSize = computeImageSize(outputSize); MPSImage* Y = createTemporaryImage(commandBuffer, imageSize); diff --git a/aten/src/ATen/native/metal/ops/MetalNeurons.mm b/aten/src/ATen/native/metal/ops/MetalNeurons.mm index b095d33288bc2..5ecbf2b9af333 100644 --- a/aten/src/ATen/native/metal/ops/MetalNeurons.mm +++ b/aten/src/ATen/native/metal/ops/MetalNeurons.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -23,7 +23,7 @@ Tensor neuronKernel(const Tensor& input, MPSCNNNeuron* neuron) { } IntArrayRef textureSize = outputSize; MetalTensorImplStorage mt{outputSize.vec()}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(textureSize, commandBuffer); MPSImage* Y = mt.texture()->image(); [neuron encodeToCommandBuffer:commandBuffer.buffer @@ -40,7 +40,7 @@ Tensor neuronKernel(const Tensor& input, MPSCNNNeuron* neuron) { return input; } IntArrayRef textureSize = outputSize; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); MPSImage* Y = createTemporaryImage(commandBuffer, textureSize); [neuron encodeToCommandBuffer:commandBuffer.buffer sourceImage:X diff --git a/aten/src/ATen/native/metal/ops/MetalPadding.mm b/aten/src/ATen/native/metal/ops/MetalPadding.mm index 2610790e169d1..9a37f7e0abfba 100644 --- a/aten/src/ATen/native/metal/ops/MetalPadding.mm +++ b/aten/src/ATen/native/metal/ops/MetalPadding.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -48,7 +48,7 @@ Tensor reflection_pad2d(const Tensor& input, IntArrayRef padding) { } MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); MetalTensorImplStorage mt{output_size}; mt.texture()->allocateTemporaryStorage(output_size, commandBuffer); MPSImage* Y = mt.texture()->image(); diff --git a/aten/src/ATen/native/metal/ops/MetalPooling.mm b/aten/src/ATen/native/metal/ops/MetalPooling.mm index 945fc844d43ce..db8f8fd77f87e 100644 --- a/aten/src/ATen/native/metal/ops/MetalPooling.mm +++ b/aten/src/ATen/native/metal/ops/MetalPooling.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -60,7 +60,7 @@ Tensor max_pool2d( .y = mpscnn::computeMPSAlignOffset(kernel_size[1], padding[1]), .z = 0}]; MetalTensorImplStorage mt{IntArrayRef(outputSize).vec()}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); MPSImage* Y = mt.texture()->image(); [pool encodeToCommandBuffer:commandBuffer.buffer @@ -93,7 +93,7 @@ Tensor adaptive_avg_pool2d(const Tensor& input, IntArrayRef output_size) { .z = 0}]; MetalTensorImplStorage mt{IntArrayRef(outputSize).vec()}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(outputSize, commandBuffer); MPSImage* Y = mt.texture()->image(); [pool encodeToCommandBuffer:commandBuffer.buffer diff --git a/aten/src/ATen/native/metal/ops/MetalReduce.mm b/aten/src/ATen/native/metal/ops/MetalReduce.mm index 29a4bd9dc579e..5c3129b0f374a 100644 --- a/aten/src/ATen/native/metal/ops/MetalReduce.mm +++ b/aten/src/ATen/native/metal/ops/MetalReduce.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -40,7 +40,7 @@ Tensor wrapper_mean_dim( // TODO: [T87340633] Support reducing the batch dimension TORCH_CHECK(imageSize[0] == 1); auto mask = make_dim_mask(dims, input.dim()); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); MPSImage* Y = nil; for (int dim : dims) { imageSize[dim] = 1; diff --git a/aten/src/ATen/native/metal/ops/MetalReshape.mm b/aten/src/ATen/native/metal/ops/MetalReshape.mm index 64b3c8d8e7659..28dbae22d4d2a 100644 --- a/aten/src/ATen/native/metal/ops/MetalReshape.mm +++ b/aten/src/ATen/native/metal/ops/MetalReshape.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -31,7 +31,7 @@ Tensor view(const Tensor& input, IntArrayRef size) { return makeTensor({inferred_size, stride_value}, input.options()); } MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); MetalTensorImplStorage mt{inferred_size, stride_value}; mt.texture()->allocateTemporaryStorage(inferred_size, commandBuffer); MPSImage* Y = mt.texture()->image(); diff --git a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm index 181d1fff9ae34..bd22a0abe5e1e 100644 --- a/aten/src/ATen/native/metal/ops/MetalSoftmax.mm +++ b/aten/src/ATen/native/metal/ops/MetalSoftmax.mm @@ -2,7 +2,7 @@ #import #import #import -#import +#import #import #import #import @@ -39,7 +39,7 @@ Tensor mpscnn_softmax( // https://developer.apple.com/documentation/metalperformanceshaders/mpscnnsoftmax?changes=_1&language=objc T* softmax = [[T alloc] initWithDevice:[MetalContext sharedInstance].device]; MetalTensorImplStorage mt{newSize}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input_); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input_); mt.texture()->allocateTemporaryStorage(newSize, commandBuffer); MPSImage* Y = mt.texture()->image(); [softmax encodeToCommandBuffer:commandBuffer.buffer diff --git a/aten/src/ATen/native/metal/ops/MetalTranspose.mm b/aten/src/ATen/native/metal/ops/MetalTranspose.mm index 3adb0e0c1bc89..a7017fb24cd0b 100644 --- a/aten/src/ATen/native/metal/ops/MetalTranspose.mm +++ b/aten/src/ATen/native/metal/ops/MetalTranspose.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -14,6 +14,16 @@ namespace native { namespace metal { +// TODO: Move this function to MetalContext +template +id _makeMTLBuffer(const std::vector& src) { + id buffer = [[MetalContext sharedInstance].device + newBufferWithLength:src.size() * sizeof(T) + options:MTLResourceOptionCPUCacheModeWriteCombined]; + memcpy(buffer.contents, src.data(), src.size() * sizeof(T)); + return buffer; +} + Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { TORCH_CHECK(input.is_metal()); auto ndims = input.dim(); @@ -27,7 +37,7 @@ Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { auto outputSizes = input.sizes().vec(); std::swap(outputSizes[dim0], outputSizes[dim1]); MPSImage* X = imageFromTensor(input); - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); if (input.dim() == 2) { MetalTensorImplStorage mt{outputSizes}; mt.texture()->allocateTemporaryStorage(outputSizes, commandBuffer); @@ -40,9 +50,9 @@ Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { auto output = makeTensor(std::move(mt), input.options()); return output; } else { - id sizeBuf1 = makeMTLBuffer( + id sizeBuf1 = _makeMTLBuffer( std::vector{input.sizes().begin(), input.sizes().end()}); - id sizeBuf2 = makeMTLBuffer( + id sizeBuf2 = _makeMTLBuffer( std::vector{outputSizes.begin(), outputSizes.end()}); MetalTensorImplStorage mt{outputSizes}; mt.texture()->allocateTemporaryStorage(outputSizes, commandBuffer); diff --git a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm index 049aefec168fa..c5c008bf05d0d 100644 --- a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm +++ b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm @@ -1,7 +1,7 @@ #import #import #import -#import +#import #import #import #import @@ -46,7 +46,7 @@ Tensor upsample_nearest2d_vec( } MPSImage* X = imageFromTensor(input); MetalTensorImplStorage mt{outputSizes}; - MetalCommandBuffer* commandBuffer = getCommandBufferFromTensor(input); + MetalCommandBuffer* commandBuffer = getCommandBuffer(input); mt.texture()->allocateTemporaryStorage(outputSizes, commandBuffer); MPSImage* Y = mt.texture()->image(); if (@available(iOS 11.0, *)) { From cac9ae1506feabfc87d37a208b3d39ed46c59483 Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Fri, 18 Jun 2021 18:43:11 -0700 Subject: [PATCH 256/305] [iOS GPU][BE][3/n] Give MPSImage objects a label for better debugging experience (#60282) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60282 1. Adds a label to the MPSImage objects. The label describes the size of the image. 2. Remove `[image markRead]`. 3. Rename two APIs for better naming convention. ghstack-source-id: 131839557 Test Plan: 1. CircleCI 2. buck test pp-mac Reviewed By: SS-JIA Differential Revision: D29232975 fbshipit-source-id: 075175c4b5a1c5b79e795f4860e1694d7c06d4f2 --- .../ATen/native/metal/mpscnn/MPSCNNClampOp.mm | 2 - .../ATen/native/metal/mpscnn/MPSImageUtils.h | 14 ++--- .../ATen/native/metal/mpscnn/MPSImageUtils.mm | 51 ++++++++++++------- .../native/metal/mpscnn/MPSImageWrapper.mm | 2 +- .../native/metal/mpscnn/tests/MPSCNNTests.mm | 2 +- .../metal/ops/MetalBinaryElementwise.mm | 4 -- aten/src/ATen/native/metal/ops/MetalChunk.mm | 3 -- aten/src/ATen/native/metal/ops/MetalConcat.mm | 2 - aten/src/ATen/native/metal/ops/MetalCopy.mm | 1 - .../ATen/native/metal/ops/MetalHardswish.mm | 1 - .../src/ATen/native/metal/ops/MetalPadding.mm | 1 - .../src/ATen/native/metal/ops/MetalReshape.mm | 2 - .../ATen/native/metal/ops/MetalTranspose.mm | 3 -- .../metal/ops/MetalUpsamplingNearest.mm | 2 - 14 files changed, 44 insertions(+), 46 deletions(-) diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm index 93218bd5e0bcf..2f303a3b135f9 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSCNNClampOp.mm @@ -47,8 +47,6 @@ - (void)encode:(id)cb { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [_X markRead]; - [_Y markRead]; } @end diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h index 53065e4be3fa8..e81d600974243 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.h @@ -10,9 +10,7 @@ namespace native { namespace metal { MPSImage* createStaticImage(IntArrayRef sizes); -MPSImage* createStaticImage( - const float* src, - const IntArrayRef sizes); +MPSImage* createStaticImage(const float* src, const IntArrayRef sizes); MPSImage* createStaticImage( MPSTemporaryImage* image, MetalCommandBuffer* buffer, @@ -29,8 +27,12 @@ MPSTemporaryImage* createTemporaryImage( MetalCommandBuffer* buffer, MPSImage* image); -void copyToHost(float* dst, MPSImage* image); -void copyToMetalBuffer(MetalCommandBuffer* buffer, id dst, MPSImage* image); +void copyImageToFloatBuffer(float* dst, MPSImage* image); + +void copyImageToMetalBuffer( + MetalCommandBuffer* buffer, + id dst, + MPSImage* image); static inline MPSImage* imageFromTensor(const Tensor& tensor) { TORCH_CHECK(tensor.is_metal()); @@ -57,7 +59,7 @@ static inline std::vector computeImageSize(IntArrayRef sizes) { int64_t batch = 1; for (int64_t i = sizes.size() - 1; i >= 0; i--) { if (index != 0) { - imageSize[index] = sizes[i]; + imageSize[index] = sizes[i]; index--; continue; } diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm index 817672a0f5524..b72da18e4e920 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageUtils.mm @@ -1,5 +1,5 @@ -#import #import +#import #import #import #import @@ -12,16 +12,24 @@ namespace metal { MPSImage* createStaticImage(IntArrayRef sizes) { + int64_t N = sizes[0]; + int64_t C = sizes[1]; + int64_t H = sizes[2]; + int64_t W = sizes[3]; MPSImageDescriptor* desc = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 - width:sizes[3] - height:sizes[2] - featureChannels:sizes[1] - numberOfImages:sizes[0] + width:W + height:H + featureChannels:C + numberOfImages:N usage:MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite]; - return [[MPSImage alloc] initWithDevice:[MetalContext sharedInstance].device - imageDescriptor:desc]; + MPSImage* image = + [[MPSImage alloc] initWithDevice:[MetalContext sharedInstance].device + imageDescriptor:desc]; + image.label = [NSString + stringWithFormat:@"[%d, %d, %d, %d]", (int)N, (int)C, (int)H, (int)W]; + return image; } MPSImage* createStaticImage(const float* src, IntArrayRef sizes) { @@ -74,32 +82,42 @@ [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [image markRead]; if (waitUntilCompleted) { [buffer commit]; } return Y; } -MPSTemporaryImage* createTemporaryImage(MetalCommandBuffer* buffer, IntArrayRef sizes) { +MPSTemporaryImage* createTemporaryImage( + MetalCommandBuffer* buffer, + IntArrayRef sizes) { TORCH_CHECK(buffer); + int64_t N = sizes[0]; + int64_t C = sizes[1]; + int64_t H = sizes[2]; + int64_t W = sizes[3]; MPSImageDescriptor* desc = [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat16 - width:sizes[3] - height:sizes[2] - featureChannels:sizes[1] - numberOfImages:sizes[0] + width:W + height:H + featureChannels:C + numberOfImages:N usage:MTLTextureUsageShaderRead | MTLTextureUsageShaderWrite]; MPSTemporaryImage* image = [MPSTemporaryImage temporaryImageWithCommandBuffer:buffer.buffer imageDescriptor:desc]; image.readCount = INT_MAX; + image.label = [NSString + stringWithFormat:@"[%d, %d, %d, %d]", (int)N, (int)C, (int)H, (int)W]; [buffer add:image]; return image; } -MPSTemporaryImage* createTemporaryImage(MetalCommandBuffer* buffer, IntArrayRef sizes, const float* src) { +MPSTemporaryImage* createTemporaryImage( + MetalCommandBuffer* buffer, + IntArrayRef sizes, + const float* src) { TORCH_CHECK(buffer); int64_t size_bytes = c10::multiply_integers(sizes) * sizeof(float); id buff = [[MetalContext sharedInstance].device @@ -126,7 +144,6 @@ [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [output markRead]; return output; } @@ -150,7 +167,7 @@ return Y; } -void copyToHost(float* dst, MPSImage* image) { +void copyImageToFloatBuffer(float* dst, MPSImage* image) { int64_t size_bytes = c10::multiply_integers([image sizes]) * sizeof(float); id buffer = [[MetalContext sharedInstance].device newBufferWithLength:size_bytes @@ -184,7 +201,7 @@ void copyToHost(float* dst, MPSImage* image) { memcpy(dst, buffer.contents, buffer.length); } -void copyToMetalBuffer( +void copyImageToMetalBuffer( MetalCommandBuffer* cmdBuffer, id dst, MPSImage* image) { diff --git a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm index 1d94bdd554a25..287f94dde778f 100644 --- a/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm +++ b/aten/src/ATen/native/metal/mpscnn/MPSImageWrapper.mm @@ -117,7 +117,7 @@ - (void)endSynchronization:(NSError*)error { options:MTLResourceCPUCacheModeWriteCombined]; TORCH_CHECK(_buffer, "Allocate GPU memory failed!"); } - copyToMetalBuffer(_commandBuffer, _buffer, _image); + copyImageToMetalBuffer(_commandBuffer, _buffer, _image); if (_image.isTemporaryImage && _image.readCount != 0) { _image = createStaticImage((MPSTemporaryImage*)_image, _commandBuffer, false); diff --git a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm index 5df31eb320d47..7f6b738ac64e4 100644 --- a/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm +++ b/aten/src/ATen/native/metal/mpscnn/tests/MPSCNNTests.mm @@ -115,7 +115,7 @@ bool test_copy_nchw_to_metal() { createTemporaryImage(cb, t1.sizes().vec(), t1.data_ptr()); MPSImage* img2 = createStaticImage(img1, cb, true); auto t2 = at::zeros(size); - copyToHost(t2.data_ptr(), img2); + copyImageToFloatBuffer(t2.data_ptr(), img2); return almostEqual(t1, t2); }); } diff --git a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm index 97294a11b9be5..6aa4b7ae33bf0 100644 --- a/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm +++ b/aten/src/ATen/native/metal/ops/MetalBinaryElementwise.mm @@ -94,8 +94,6 @@ Tensor binaryElementwiseShaderKernel( [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X1 markRead]; - [X2 markRead]; auto output = makeTensor(std::move(mt), input1.options()); return output; } @@ -134,8 +132,6 @@ Tensor binaryElementwiseShaderKernel( [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X1 markRead]; - [X2 markRead]; MetalTensorImpl* impl = (MetalTensorImpl*)input1.unsafeGetTensorImpl(); MetalTensorImplStorage& implStorage = impl->unsafe_opaque_handle(); implStorage.texture()->setImage(Y); diff --git a/aten/src/ATen/native/metal/ops/MetalChunk.mm b/aten/src/ATen/native/metal/ops/MetalChunk.mm index 3da3e683ee727..14e54642817a9 100644 --- a/aten/src/ATen/native/metal/ops/MetalChunk.mm +++ b/aten/src/ATen/native/metal/ops/MetalChunk.mm @@ -54,9 +54,6 @@ [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; - [Y1 markRead]; - [Y2 markRead]; auto output1 = makeTensor(std::move(mt1), input.options()); auto output2 = makeTensor(std::move(mt2), input.options()); return {output1, output2}; diff --git a/aten/src/ATen/native/metal/ops/MetalConcat.mm b/aten/src/ATen/native/metal/ops/MetalConcat.mm index 2b34bc4efef53..e3f7592b8996a 100644 --- a/aten/src/ATen/native/metal/ops/MetalConcat.mm +++ b/aten/src/ATen/native/metal/ops/MetalConcat.mm @@ -49,7 +49,6 @@ Tensor cat_batch(const TensorList tensors, MetalTensorImplStorage& mt) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; cat_dim4_pointer += t.size(0) * ((t.size(1) + 3) / 4); } auto output = makeTensor(std::move(mt), tensor.options()); @@ -111,7 +110,6 @@ Tensor cat_feature(const TensorList tensors, MetalTensorImplStorage& mt) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; channel_offset += X.featureChannels; } auto output = makeTensor(std::move(mt), tensor.options()); diff --git a/aten/src/ATen/native/metal/ops/MetalCopy.mm b/aten/src/ATen/native/metal/ops/MetalCopy.mm index b6c783b1e1579..82bb164019f9c 100644 --- a/aten/src/ATen/native/metal/ops/MetalCopy.mm +++ b/aten/src/ATen/native/metal/ops/MetalCopy.mm @@ -46,7 +46,6 @@ Tensor copy_to_host(const Tensor& input) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; auto output = makeTensor(std::move(mt), input.options()); return output; } diff --git a/aten/src/ATen/native/metal/ops/MetalHardswish.mm b/aten/src/ATen/native/metal/ops/MetalHardswish.mm index f446b4c49bcc4..8d3526a4c6b2a 100644 --- a/aten/src/ATen/native/metal/ops/MetalHardswish.mm +++ b/aten/src/ATen/native/metal/ops/MetalHardswish.mm @@ -41,7 +41,6 @@ [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; MetalTensorImpl* impl = (MetalTensorImpl*)input.unsafeGetTensorImpl(); MetalTensorImplStorage& implStorage = impl->unsafe_opaque_handle(); implStorage.texture()->setImage(Y); diff --git a/aten/src/ATen/native/metal/ops/MetalPadding.mm b/aten/src/ATen/native/metal/ops/MetalPadding.mm index 9a37f7e0abfba..8d38a9d3fde02 100644 --- a/aten/src/ATen/native/metal/ops/MetalPadding.mm +++ b/aten/src/ATen/native/metal/ops/MetalPadding.mm @@ -81,7 +81,6 @@ Tensor reflection_pad2d(const Tensor& input, IntArrayRef padding) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; auto output = makeTensor(std::move(mt), input.options()); return output; } diff --git a/aten/src/ATen/native/metal/ops/MetalReshape.mm b/aten/src/ATen/native/metal/ops/MetalReshape.mm index 28dbae22d4d2a..185747ac1ee9a 100644 --- a/aten/src/ATen/native/metal/ops/MetalReshape.mm +++ b/aten/src/ATen/native/metal/ops/MetalReshape.mm @@ -57,8 +57,6 @@ Tensor view(const Tensor& input, IntArrayRef size) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; - [Y markRead]; auto output = makeTensor(std::move(mt), input.options()); return output; } diff --git a/aten/src/ATen/native/metal/ops/MetalTranspose.mm b/aten/src/ATen/native/metal/ops/MetalTranspose.mm index a7017fb24cd0b..f7f7890155f72 100644 --- a/aten/src/ATen/native/metal/ops/MetalTranspose.mm +++ b/aten/src/ATen/native/metal/ops/MetalTranspose.mm @@ -82,9 +82,6 @@ Tensor transpose(const Tensor& input, int64_t dim0, int64_t dim1) { [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; - [Y markRead]; - auto output = makeTensor(std::move(mt), input.options()); return output; } diff --git a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm index c5c008bf05d0d..c7979f5c7a267 100644 --- a/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm +++ b/aten/src/ATen/native/metal/ops/MetalUpsamplingNearest.mm @@ -81,8 +81,6 @@ Tensor upsample_nearest2d_vec( [encoder dispatchThreadgroups:launchParams.threadgroupsPerGrid threadsPerThreadgroup:launchParams.threadsPerThreadgroup]; [encoder endEncoding]; - [X markRead]; - [Y markRead]; } auto output = makeTensor(std::move(mt), input.options()); return output; From 5824a866b72c251ad47a9c16dc652e49cfd7e234 Mon Sep 17 00:00:00 2001 From: Jiakai Liu Date: Sat, 19 Jun 2021 06:09:31 -0700 Subject: [PATCH 257/305] [pytorch][nnc] support custom class parameters (#59466) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59466 Change saved parameter type from at::Tensor to at::IValue to support custom class parameters, e.g. `__torch__.torch.classes.xnnpack.Conv2dOpContext`. The NNC produced kernels won't deal with custom class parameters directly. They simply pass through to the external operators that take these custom class parameters, e.g. `prepacked::conv2d_clamp_run`. It will reuse the `__getstate__` and `__setstate__` methods on the custom class to persist and restore the state of the parameters. When calling into the kernel, it will pass in the untyped raw pointer of the custom class objects to the kernel as `void*`. It's similar to the regular tensor parameters, for which it will pass in the raw data pointer of the tensor storage. The generated kernel needs to hardcode the expected type for each parameter and cast before calling the external ops. ghstack-source-id: 131897904 Test Plan: - unit tests Reviewed By: kimishpatel Differential Revision: D28902496 fbshipit-source-id: 4b2c0895dd28f0b7d344aa08183d42ad6a355dae --- .../macos-lite-interpreter-build-test.sh | 1 - .jenkins/pytorch/macos-test.sh | 1 - .jenkins/pytorch/test.sh | 1 + caffe2/CMakeLists.txt | 11 +- test/mobile/nnc/CMakeLists.txt | 4 +- test/mobile/nnc/test_context.cpp | 19 +- test/mobile/nnc/test_nnc_backend.cpp | 236 ++++++++++++++++++ tools/build_variables.bzl | 5 +- torch/csrc/jit/mobile/nnc/backend.cpp | 1 - torch/csrc/jit/mobile/nnc/context.cpp | 15 +- torch/csrc/jit/mobile/nnc/context.h | 8 +- 11 files changed, 275 insertions(+), 27 deletions(-) create mode 100644 test/mobile/nnc/test_nnc_backend.cpp diff --git a/.jenkins/pytorch/macos-lite-interpreter-build-test.sh b/.jenkins/pytorch/macos-lite-interpreter-build-test.sh index 0e23f13bba1db..901f4517ddbd5 100644 --- a/.jenkins/pytorch/macos-lite-interpreter-build-test.sh +++ b/.jenkins/pytorch/macos-lite-interpreter-build-test.sh @@ -28,7 +28,6 @@ if [ "${BUILD_LITE_INTERPRETER}" == 1 ]; then popd || exit "${CPP_BUILD}/caffe2/build/bin/test_lite_interpreter_runtime" - "${CPP_BUILD}/caffe2/build/bin/test_mobile_nnc" # Change the permission manually from 755 to 644 to keep git clean chmod 644 "${HOME}/project/.jenkins/pytorch/macos-lite-interpreter-build-test.sh" diff --git a/.jenkins/pytorch/macos-test.sh b/.jenkins/pytorch/macos-test.sh index 0ee446c4d25fc..14334a75d2c23 100755 --- a/.jenkins/pytorch/macos-test.sh +++ b/.jenkins/pytorch/macos-test.sh @@ -159,7 +159,6 @@ test_jit_hooks() { assert_git_not_dirty } - if [ -z "${BUILD_ENVIRONMENT}" ] || [[ "${BUILD_ENVIRONMENT}" == *-test ]]; then test_python_all test_libtorch diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index eeef1b9272485..8338340af4f13 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -236,6 +236,7 @@ test_libtorch() { wait OMP_NUM_THREADS=2 TORCH_CPP_TEST_MNIST_PATH="test/cpp/api/mnist" build/bin/test_api --gtest_output=xml:$TEST_REPORTS_DIR/test_api.xml build/bin/test_tensorexpr --gtest_output=xml:$TEST_REPORTS_DIR/test_tensorexpr.xml + build/bin/test_mobile_nnc --gtest_output=xml:$TEST_REPORTS_DIR/test_mobile_nnc.xml assert_git_not_dirty fi } diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 88cffd1a75d1c..d54f0a95e2bb9 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -527,8 +527,6 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE) ${TORCH_SRC_DIR}/csrc/jit/mobile/module.cpp ${TORCH_SRC_DIR}/csrc/jit/mobile/observer.cpp ${TORCH_SRC_DIR}/csrc/jit/mobile/interpreter.cpp - ${TORCH_SRC_DIR}/csrc/jit/mobile/nnc/context.cpp - ${TORCH_SRC_DIR}/csrc/jit/mobile/nnc/registry.cpp ${TORCH_SRC_DIR}/csrc/jit/mobile/train/export_data.cpp ${TORCH_SRC_DIR}/csrc/jit/mobile/train/optim/sgd.cpp ${TORCH_SRC_DIR}/csrc/jit/mobile/train/random.cpp @@ -1039,10 +1037,6 @@ endif() ${TORCH_ROOT}/test/cpp/lite_interpreter_runtime ${CMAKE_BINARY_DIR}/test_lite_interpreter_runtime ) - add_subdirectory( - ${TORCH_ROOT}/test/mobile/nnc - ${CMAKE_BINARY_DIR}/test_mobile_nnc - ) else() add_subdirectory(${TORCH_ROOT}/test/cpp/jit ${CMAKE_BINARY_DIR}/test_jit) add_subdirectory( @@ -1052,6 +1046,11 @@ endif() if(USE_DISTRIBUTED AND NOT WIN32) add_subdirectory(${TORCH_ROOT}/test/cpp/rpc ${CMAKE_BINARY_DIR}/test_cpp_rpc) endif() + + add_subdirectory( + ${TORCH_ROOT}/test/mobile/nnc + ${CMAKE_BINARY_DIR}/test_mobile_nnc + ) endif() endif() diff --git a/test/mobile/nnc/CMakeLists.txt b/test/mobile/nnc/CMakeLists.txt index 001c7f32d5b8b..5d37f8bb99d9a 100644 --- a/test/mobile/nnc/CMakeLists.txt +++ b/test/mobile/nnc/CMakeLists.txt @@ -2,12 +2,14 @@ set(MOBILE_NNC_TEST_ROOT ${TORCH_ROOT}/test/mobile/nnc) set(MOBILE_NNC_TEST_SRCS ${MOBILE_NNC_TEST_ROOT}/test_context.cpp + ${MOBILE_NNC_TEST_ROOT}/test_nnc_backend.cpp ${MOBILE_NNC_TEST_ROOT}/test_registry.cpp ) add_executable(test_mobile_nnc ${TORCH_ROOT}/test/cpp/lite_interpreter_runtime/main.cpp - ${MOBILE_NNC_TEST_SRCS}) + ${MOBILE_NNC_TEST_SRCS} +) target_link_libraries(test_mobile_nnc PRIVATE torch gtest) target_include_directories(test_mobile_nnc PRIVATE ${ATen_CPU_INCLUDE}) diff --git a/test/mobile/nnc/test_context.cpp b/test/mobile/nnc/test_context.cpp index 3022695f51399..a9f68ee58324b 100644 --- a/test/mobile/nnc/test_context.cpp +++ b/test/mobile/nnc/test_context.cpp @@ -63,8 +63,10 @@ TEST(Function, ExecuteSlowMul) { f.set_nnc_kernel_id("slow_mul"); f.set_input_specs({create_test_input_spec({size})}); - f.set_output_spec({create_test_output_spec({size})}); - f.set_parameters({at::ones({1}, at::kInt).mul(n)}); + f.set_output_specs({create_test_output_spec({size})}); + f.set_parameters(c10::impl::toList(c10::List({ + at::ones({1}, at::kInt).mul(n) + }))); f.set_memory_plan(create_test_memory_plan({sizeof(float) * size})); c10::List input({ @@ -81,12 +83,13 @@ TEST(Function, Serialization) { f.set_name("test_function"); f.set_nnc_kernel_id("test_kernel"); f.set_input_specs({create_test_input_spec({1, 3, 224, 224})}); - f.set_output_spec({create_test_output_spec({1000})}); - f.set_parameters({ + f.set_output_specs({create_test_output_spec({1000})}); + + f.set_parameters(c10::impl::toList(c10::List({ at::ones({1, 16, 3, 3}, at::kFloat), at::ones({16, 32, 1, 1}, at::kFloat), at::ones({32, 1, 3, 3}, at::kFloat) - }); + }))); f.set_memory_plan(create_test_memory_plan({ sizeof(float) * 1024, sizeof(float) * 2048, @@ -105,9 +108,9 @@ TEST(Function, Serialization) { EXPECT_EQ(f2.output_specs()[0].dtype_, at::kFloat); EXPECT_EQ(f2.parameters().size(), 3); - EXPECT_EQ(f2.parameters()[0].sizes(), at::IntArrayRef({1, 16, 3, 3})); - EXPECT_EQ(f2.parameters()[1].sizes(), at::IntArrayRef({16, 32, 1, 1})); - EXPECT_EQ(f2.parameters()[2].sizes(), at::IntArrayRef({32, 1, 3, 3})); + EXPECT_EQ(f2.parameters()[0].toTensor().sizes(), at::IntArrayRef({1, 16, 3, 3})); + EXPECT_EQ(f2.parameters()[1].toTensor().sizes(), at::IntArrayRef({16, 32, 1, 1})); + EXPECT_EQ(f2.parameters()[2].toTensor().sizes(), at::IntArrayRef({32, 1, 3, 3})); EXPECT_EQ(f2.memory_plan().buffer_sizes_.size(), 2); EXPECT_EQ(f2.memory_plan().buffer_sizes_[0], sizeof(float) * 1024); diff --git a/test/mobile/nnc/test_nnc_backend.cpp b/test/mobile/nnc/test_nnc_backend.cpp new file mode 100644 index 0000000000000..0e59aaa5547dd --- /dev/null +++ b/test/mobile/nnc/test_nnc_backend.cpp @@ -0,0 +1,236 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace torch { +namespace jit { +namespace mobile { +namespace nnc { + +namespace { + +c10::Dict create_compile_spec( + const std::string& method_name, + const std::string& nnc_kernel_id, + const std::vector>& input_shapes, + const std::vector>& output_shapes, + const c10::impl::GenericList& parameters, + const std::vector& buffer_sizes) { + c10::Dict method_spec( + c10::StringType::get(), c10::AnyType::get()); + method_spec.insert("nnc_kernel_id", nnc_kernel_id); + method_spec.insert("input_sizes", input_shapes); + method_spec.insert("output_sizes", output_shapes); + + // For testing purpose we don't call the real NNC so pass in these directly. + method_spec.insert("parameters", parameters); + method_spec.insert("buffer_sizes", buffer_sizes); + + c10::Dict compile_spec( + c10::StringType::get(), c10::AnyType::get()); + compile_spec.insert(method_name, method_spec); + return compile_spec; +} + +std::vector get_input_specs( + const c10::Dict& method_compile_spec) { + auto input_shapes = method_compile_spec.at("input_sizes").toList(); + + std::vector specs; + for (const auto& input_shape : input_shapes) { + mobile::nnc::InputSpec spec; + spec.sizes_ = ((c10::IValue) input_shape).toIntVector(); + spec.dtype_ = c10::ScalarType::Float; + specs.emplace_back(std::move(spec)); + } + return specs; +} + +std::vector get_output_specs( + const c10::Dict& method_compile_spec) { + auto output_shapes = method_compile_spec.at("output_sizes").toList(); + + std::vector specs; + for (const auto& output_shape : output_shapes) { + mobile::nnc::OutputSpec spec; + spec.sizes_ = ((c10::IValue) output_shape).toIntVector(); + spec.dtype_ = c10::ScalarType::Float; + specs.emplace_back(std::move(spec)); + } + return specs; +} + +// A fake NNC preprocess method, which only produces the compiled model but +// does not produce the assembly with the NNC compiler. +c10::IValue preprocess( + const torch::jit::Module& /* mod */, + const c10::Dict& method_compile_spec, + const torch::jit::BackendDebugHandleGenerator&) { + torch::jit::mobile::nnc::CompilationUnit cu; + for (const auto& entry : method_compile_spec) { + const std::string& method_name = entry.key().toStringRef(); + auto compile_spec = entry.value().toGenericDict(); + + auto func = std::make_unique(); + func->set_name(method_name); + func->set_nnc_kernel_id(compile_spec.at("nnc_kernel_id").toStringRef()); + func->set_input_specs(get_input_specs(compile_spec)); + func->set_output_specs(get_output_specs(compile_spec)); + + func->set_parameters(compile_spec.at("parameters").toList()); + + mobile::nnc::MemoryPlan plan; + plan.buffer_sizes_ = compile_spec.at("buffer_sizes").toIntVector(); + func->set_memory_plan(plan); + + cu.register_function(std::move(func)); + } + return cu.serialize(); +} + +static auto reg = torch::jit::backend_preprocess_register("nnc", preprocess); + +struct FakeTensor : torch::CustomClassHolder { + explicit FakeTensor(std::vector data) : data_(std::move(data)) {} + int64_t get() { + return data_[0]; + } + std::vector data_; +}; + +TORCH_LIBRARY(_TorchScriptTesting, m) { + m.class_("_MobileNNCFakeTensor") + .def(torch::init>()) + .def("get", &FakeTensor::get) + .def_pickle( + [](c10::intrusive_ptr self) { // __getstate__ + return self->data_; + }, + [](std::vector state) { // __setstate__ + return c10::make_intrusive(std::move(state)); + }); +} + +} // namespace + +extern "C" { + +// The test kernels are supposed to be generated by the NNC compiler ahead-of- +// time. For integration test purpose we manually wrote instead. +int add_kernel(void** args) { + // out = input + param + at::Tensor input = at::from_blob(args[0], {4, 4}, at::kFloat); + at::Tensor out = at::from_blob(args[1], {4, 4}, at::kFloat); + at::Tensor param = at::from_blob(args[2], {1}, at::kFloat); + out.copy_(at::add(input, param)); + return 0; +} + +int fake_tensor_add_kernel(void** args) { + // out = input + param.get() + at::Tensor input = at::from_blob(args[0], {4, 4}, at::kFloat); + at::Tensor out = at::from_blob(args[1], {4, 4}, at::kFloat); + FakeTensor* param = reinterpret_cast(args[2]); + out.copy_(at::add(input, param->get())); + return 0; +} + +} // extern "C" + +REGISTER_NNC_KERNEL("_add_kernel", add_kernel) +REGISTER_NNC_KERNEL("_fake_tensor_add_kernel", fake_tensor_add_kernel) + +TEST(NNCBackendTest, AOTCompileThenExecute) { + torch::jit::Module m("m"); + auto param = torch::ones({}); + m.register_parameter("param", param, false); + m.define(R"( + def forward(self, input): + return input + self.param + )"); + + // Run the TorchScript module to get reference result. + std::vector inputs; + inputs.emplace_back(2.0 * torch::ones({4, 4})); + auto reference = m.forward(inputs); + + // Compile the model with NNC. + auto compile_spec = create_compile_spec( + "forward", + "_add_kernel", + {{4, 4}}, + {{4, 4}}, + c10::impl::toList(c10::List({param})), + {}); + auto any_dict_ty = + c10::DictType::create(c10::StringType::get(), c10::AnyType::get()); + auto compiled_module = torch::jit::detail::codegen_backend_module( + "nnc", m, compile_spec, any_dict_ty); + + // Save the compiled model. + std::stringstream ss; + compiled_module._save_for_mobile(ss); + + // Load and run the saved model. + auto loaded_module = _load_for_mobile(ss); + auto result = loaded_module.forward(inputs); + EXPECT_TRUE(result.toTensor().equal(3.0 * torch::ones({4, 4}))); + EXPECT_TRUE(result.toTensor().equal(reference.toTensor())); +} + +TEST(NNCBackendTest, FakeTensor) { + script::Module m("m"); + auto param_cls = getCustomClass( + "__torch__.torch.classes._TorchScriptTesting._MobileNNCFakeTensor"); + auto param_value = c10::make_intrusive(std::vector({3})); + m.register_attribute("param", param_cls, param_value, false); + m.define( + R"( + def forward(self, input): + return input + self.param.get() + )"); + + // Run the TorchScript module to get reference result. + std::vector inputs; + inputs.emplace_back(2.0 * torch::ones({4, 4})); + auto reference = m.forward(inputs); + + // Compile the model with NNC. + auto params = c10::impl::GenericList(c10::AnyType::get()); + params.emplace_back(param_value); + auto compile_spec = create_compile_spec( + "forward", + "_fake_tensor_add_kernel", + {{4, 4}}, + {{4, 4}}, + params, + {}); + auto any_dict_ty = + c10::DictType::create(c10::StringType::get(), c10::AnyType::get()); + auto compiled_module = torch::jit::detail::codegen_backend_module( + "nnc", m, compile_spec, any_dict_ty); + + // Save the compiled model. + std::stringstream ss; + compiled_module._save_for_mobile(ss); + + // Load and run the saved model. + auto loaded_module = _load_for_mobile(ss); + auto result = loaded_module.forward(inputs); + EXPECT_TRUE(result.toTensor().equal(5.0 * torch::ones({4, 4}))); + EXPECT_TRUE(result.toTensor().equal(reference.toTensor())); +} + +} // namespace nnc +} // namespace mobile +} // namespace jit +} // namespace torch diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl index cf7e67fa77797..afe86e06ea281 100644 --- a/tools/build_variables.bzl +++ b/tools/build_variables.bzl @@ -174,6 +174,9 @@ core_sources_full_mobile = [ "torch/csrc/jit/ir/subgraph_matcher.cpp", "torch/csrc/jit/jit_log.cpp", "torch/csrc/jit/jit_opt_limit.cpp", + "torch/csrc/jit/mobile/nnc/backend.cpp", + "torch/csrc/jit/mobile/nnc/context.cpp", + "torch/csrc/jit/mobile/nnc/registry.cpp", "torch/csrc/jit/passes/annotate_warns.cpp", "torch/csrc/jit/passes/bailout_graph.cpp", "torch/csrc/jit/passes/batch_mm.cpp", @@ -451,8 +454,6 @@ libtorch_extra_sources = libtorch_core_jit_sources + [ "torch/csrc/jit/mobile/interpreter.cpp", "torch/csrc/jit/mobile/model_compatibility.cpp", "torch/csrc/jit/mobile/module.cpp", - "torch/csrc/jit/mobile/nnc/context.cpp", - "torch/csrc/jit/mobile/nnc/registry.cpp", "torch/csrc/jit/mobile/observer.cpp", "torch/csrc/jit/mobile/train/export_data.cpp", "torch/csrc/jit/mobile/train/optim/sgd.cpp", diff --git a/torch/csrc/jit/mobile/nnc/backend.cpp b/torch/csrc/jit/mobile/nnc/backend.cpp index 1909ab4f6c3ad..77742faa5cc38 100644 --- a/torch/csrc/jit/mobile/nnc/backend.cpp +++ b/torch/csrc/jit/mobile/nnc/backend.cpp @@ -2,7 +2,6 @@ #include #include -#include namespace torch { namespace jit { diff --git a/torch/csrc/jit/mobile/nnc/context.cpp b/torch/csrc/jit/mobile/nnc/context.cpp index 8b20f0ca919c7..15e7bc9d48228 100644 --- a/torch/csrc/jit/mobile/nnc/context.cpp +++ b/torch/csrc/jit/mobile/nnc/context.cpp @@ -89,7 +89,7 @@ Function::Function(const c10::IValue& value) { auto dict = value.toGenericDict(); name_ = c10::QualifiedName(dict.at("name").toStringRef()); nnc_kernel_id_ = dict.at("nnc_kernel_id").toStringRef(); - parameters_ = dict.at("parameters").toTensorVector(); + parameters_ = dict.at("parameters").toList(); // input_specs_ for (const auto& input_value : dict.at("input_specs").toTuple()->elements()) { @@ -157,9 +157,18 @@ void Function::init_execution_state() const { // Keep empty slots to fill in inputs/outputs pointers at execution time. arguments.resize(input_args + output_args); - // Fill in parameter pointers. + // Fill in parameters as untyped raw pointers. + // The underlying storage of the parameters should be owned by `parameters_`, + // which should be alive when `execution_state_` is being used. for (const auto& param : parameters_) { - arguments.emplace_back(param.data_ptr()); + const c10::IValue& ivalue = (c10::IValue)param; + if (ivalue.isTensor()) { + arguments.emplace_back(ivalue.toTensor().data_ptr()); + } else if (torch::isCustomClass(ivalue)) { + arguments.emplace_back(ivalue.toObjectRef().getSlot(0).toCapsule().get()); + } else { + TORCH_CHECK(false, "Invalid parameter: ", ivalue); + } } // Fill in preallocated buffer pointers. diff --git a/torch/csrc/jit/mobile/nnc/context.h b/torch/csrc/jit/mobile/nnc/context.h index a2ec1760eb517..e301846ee8aea 100644 --- a/torch/csrc/jit/mobile/nnc/context.h +++ b/torch/csrc/jit/mobile/nnc/context.h @@ -124,11 +124,11 @@ class TORCH_API Function { // The parameters (e.g. weights / bias tensors) to be passed to the generated // NNC kernel. - const std::vector& parameters() const { + const c10::impl::GenericList& parameters() const { return parameters_; } - void set_parameters(const std::vector& parameters) { + void set_parameters(const c10::impl::GenericList& parameters) { parameters_ = parameters; } @@ -144,7 +144,7 @@ class TORCH_API Function { return output_specs_; } - void set_output_spec(const std::vector& output_specs) { + void set_output_specs(const std::vector& output_specs) { output_specs_ = output_specs; } @@ -161,7 +161,7 @@ class TORCH_API Function { c10::QualifiedName name_; std::string nnc_kernel_id_; - std::vector parameters_; + c10::impl::GenericList parameters_{at::AnyType::get()}; std::vector input_specs_; std::vector output_specs_; MemoryPlan memory_plan_; From 5ec4ad7f54cc35ee40b5991bd65c974056048f84 Mon Sep 17 00:00:00 2001 From: kshitij12345 Date: Sat, 19 Jun 2021 18:35:11 -0700 Subject: [PATCH 258/305] [special] Add special.ndtri (#58650) Summary: Reference: https://github.com/pytorch/pytorch/issues/50345 TODO * [x] Add docs https://13865352-65600975-gh.circle-artifacts.com/0/docs/special.html#torch.special.ndtri * [x] Add comments on implementation * [x] Clean-up Pull Request resolved: https://github.com/pytorch/pytorch/pull/58650 Reviewed By: H-Huang Differential Revision: D29160170 Pulled By: mruberry fbshipit-source-id: 50e4ea663920e97b8437d03d5b52bcd9dedc1a8d --- aten/src/ATen/core/aten_interned_strings.h | 1 + aten/src/ATen/native/Distributions.h | 10 +- aten/src/ATen/native/Math.h | 182 ++++++++++++++++-- aten/src/ATen/native/UnaryOps.cpp | 3 + aten/src/ATen/native/UnaryOps.h | 1 + aten/src/ATen/native/cpu/UnaryOpsKernel.cpp | 9 + .../ATen/native/cuda/UnarySpecialOpsKernel.cu | 8 + aten/src/ATen/native/native_functions.yaml | 13 ++ docs/source/special.rst | 1 + tools/autograd/derivatives.yaml | 3 + torch/csrc/api/include/torch/special.h | 8 + torch/overrides.py | 1 + torch/special/__init__.py | 30 +++ .../_internal/common_methods_invocations.py | 6 + 14 files changed, 246 insertions(+), 30 deletions(-) diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 25cda648c89a6..562a982ab86d2 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -533,6 +533,7 @@ _(aten, native_tensor) \ _(aten, native_zero) \ _(aten, special_ndtr) \ _(aten, nextafter) \ +_(aten, special_ndtri) \ _(aten, bitwise_and) \ _(aten, bitwise_not) \ _(aten, bitwise_or) \ diff --git a/aten/src/ATen/native/Distributions.h b/aten/src/ATen/native/Distributions.h index b7a1f52bbfd0c..ebfaf46313697 100644 --- a/aten/src/ATen/native/Distributions.h +++ b/aten/src/ATen/native/Distributions.h @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -118,15 +119,6 @@ C10_DEVICE scalar_t sample_gamma(scalar_t alpha, BaseSampler -C10_DEVICE static inline scalar_t polevl(const scalar_t x, const scalar_t A[], size_t len) { - scalar_t result = 0; - for (size_t i = 0; i <= len; i++) { - result = result * x + A[i]; - } - return result; -} - /* the functions stirling_approx_tail, binomial_inversion, and btrs are adapted * from TensorFlow's random_binomial_op.cc implementation. That code is under * copyright: 2019 The TensorFlow Authors. diff --git a/aten/src/ATen/native/Math.h b/aten/src/ATen/native/Math.h index a492ac02dd4e5..8a23289690929 100644 --- a/aten/src/ATen/native/Math.h +++ b/aten/src/ATen/native/Math.h @@ -220,16 +220,24 @@ static inline double zeta(double x, double q) { return s; } -static inline double polevl(double x, double *A, size_t len) { - double result = 0; - for (size_t i = 0; i <= len; i++) { - result = result * x + A[i]; - } - return result; -} - -static inline float polevlf(float x, float *A, size_t len) { - float result = 0; +/* + * This function is derived from the implementation of the digamma function in the Cephes Math Library. + * See note [3-Clause BSD License for the Cephes Math Library]. + * + * Evaluates polynomial of degree N: + * + * 2 N + * y = C + C x + C x +...+ C x + * 0 1 2 N + * + * Coefficients are stored in reverse order: + * + * coef[0] = C , ..., coef[N] = C . + * N 0 + */ +template +C10_HOST_DEVICE static inline T polevl(const T x, const T A[], size_t len) { + T result = 0; for (size_t i = 0; i <= len; i++) { result = result * x + A[i]; } @@ -312,7 +320,7 @@ static inline double calc_digamma(double x) { } // Compute asymptotic digamma - static double A[] = { + static const double A[] = { 8.33333333333333333333E-2, -2.10927960927960927961E-2, 7.57575757575757575758E-3, @@ -371,7 +379,7 @@ static inline float calc_digamma(float x) { } // Compute asymptotic digamma - static float A[] = { + static const float A[] = { 8.33333333333333333333E-2f, -2.10927960927960927961E-2f, 7.57575757575757575758E-3f, @@ -384,7 +392,7 @@ static inline float calc_digamma(float x) { float y = 0; if (x < 1.0e17f) { float z = 1 / (x * x); - y = z * polevlf(z, A, 6); + y = z * polevl(z, A, 6); } return result + logf(x) - (0.5f / x) - y; } @@ -1196,7 +1204,7 @@ chbevl(const T x, const T array[], size_t len) { * of all inputs to convert them into the domain of the approximation. */ template -inline std::tuple chebyshev_coefficients_i0e_A() { +static inline std::tuple chebyshev_coefficients_i0e_A() { /* Chebyshev coefficients for exp(-x) I0(x) * in the interval [0,8]. * @@ -1222,7 +1230,7 @@ inline std::tuple chebyshev_coefficients_i0e_A() { }; template -inline std::tuple chebyshev_coefficients_i0e_B() { +static inline std::tuple chebyshev_coefficients_i0e_B() { /* Chebyshev coefficients for exp(-x) sqrt(x) I0(x) * in the inverted interval [8,infinity]. * @@ -1247,7 +1255,7 @@ inline std::tuple chebyshev_coefficients_i0e_B() { }; template -inline typename std::enable_if::value, std::tuple>::type +static inline typename std::enable_if::value, std::tuple>::type chebyshev_coefficients_i1e_A() { /* Chebyshev coefficients for exp(-x) I1(x) * in the interval [0,8]. @@ -1274,7 +1282,7 @@ chebyshev_coefficients_i1e_A() { }; template -inline typename std::enable_if::value, std::tuple>::type +static inline typename std::enable_if::value, std::tuple>::type chebyshev_coefficients_i1e_A() { /* Chebyshev coefficients for exp(-x) I1(x) * in the interval [0,8]. @@ -1303,7 +1311,7 @@ chebyshev_coefficients_i1e_A() { }; template -inline typename std::enable_if::value, std::tuple>::type +static inline typename std::enable_if::value, std::tuple>::type chebyshev_coefficients_i1e_B() { /* Chebyshev coefficients for exp(-x) sqrt(x) I1(x) * in the inverted interval [8,infinity]. @@ -1329,7 +1337,7 @@ chebyshev_coefficients_i1e_B() { }; template -inline typename std::enable_if::value, std::tuple>::type +static inline typename std::enable_if::value, std::tuple>::type chebyshev_coefficients_i1e_B() { /* Chebyshev coefficients for exp(-x) sqrt(x) I1(x) * in the inverted interval [8,infinity]. @@ -1368,7 +1376,7 @@ calc_i0(T _x) { } // Upcast bfloat16 input to float for numerical accuracy purposes -inline c10::BFloat16 calc_i0(c10::BFloat16 a) { return calc_i0(static_cast(a)); } +static inline c10::BFloat16 calc_i0(c10::BFloat16 a) { return calc_i0(static_cast(a)); } /* * This function is derived from the implementation of the i0e function in the Cephes Math Library. @@ -1400,7 +1408,7 @@ calc_i0e(T _x) { } // Upcast bfloat16 input to float for numerical accuracy purposes -inline c10::BFloat16 calc_i0e(c10::BFloat16 a) { return calc_i0e(static_cast(a)); } +static inline c10::BFloat16 calc_i0e(c10::BFloat16 a) { return calc_i0e(static_cast(a)); } /* * This function is derived from the implementation of the i1 function in the Cephes Math Library. @@ -1461,3 +1469,135 @@ calc_i1e(T _x) { static_cast(chbevl(static_cast(32.0 / x - 2.0), B, len) / std::sqrt(x)); return (_x < 0.0) ? -out : out; } + +/* + * This function is derived from the implementation of the i1e function in the Cephes Math Library. + * See note [3-Clause BSD License for the Cephes Math Library]. + * + * Computes the argument, x, for which the area under the Gaussian probability density function + * (integrated from minus infinity to x) is equal to y. + */ +template +static inline C10_HOST_DEVICE T calc_ndtri(T y0) { + + /* sqrt(2pi) */ + constexpr T s2pi = 2.50662827463100050242E0; + constexpr T one = 1; + constexpr T zero = 0; + + /* approximation for 0 <= |y - 0.5| <= 3/8 */ + static const T P0[5] = { + -5.99633501014107895267E1, + 9.80010754185999661536E1, + -5.66762857469070293439E1, + 1.39312609387279679503E1, + -1.23916583867381258016E0, + }; + + static const T Q0[9] = { + 1.00000000000000000000E0, + 1.95448858338141759834E0, + 4.67627912898881538453E0, + 8.63602421390890590575E1, + -2.25462687854119370527E2, + 2.00260212380060660359E2, + -8.20372256168333339912E1, + 1.59056225126211695515E1, + -1.18331621121330003142E0, + }; + + /* Approximation for interval z = sqrt(-2 log y ) between 2 and 8 + * i.e., y between exp(-2) = .135 and exp(-32) = 1.27e-14. + */ + static const T P1[9] = { + 4.05544892305962419923E0, + 3.15251094599893866154E1, + 5.71628192246421288162E1, + 4.40805073893200834700E1, + 1.46849561928858024014E1, + 2.18663306850790267539E0, + -1.40256079171354495875E-1, + -3.50424626827848203418E-2, + -8.57456785154685413611E-4, + }; + + static const T Q1[9] = { + 1.00000000000000000000E0, + 1.57799883256466749731E1, + 4.53907635128879210584E1, + 4.13172038254672030440E1, + 1.50425385692907503408E1, + 2.50464946208309415979E0, + -1.42182922854787788574E-1, + -3.80806407691578277194E-2, + -9.33259480895457427372E-4, + }; + + /* Approximation for interval z = sqrt(-2 log y ) between 8 and 64 + * i.e., y between exp(-32) = 1.27e-14 and exp(-2048) = 3.67e-890. + */ + + static const T P2[9] = { + 3.23774891776946035970E0, + 6.91522889068984211695E0, + 3.93881025292474443415E0, + 1.33303460815807542389E0, + 2.01485389549179081538E-1, + 1.23716634817820021358E-2, + 3.01581553508235416007E-4, + 2.65806974686737550832E-6, + 6.23974539184983293730E-9, + }; + + static const T Q2[9] = { + 1.00000000000000000000E0, + 6.02427039364742014255E0, + 3.67983563856160859403E0, + 1.37702099489081330271E0, + 2.16236993594496635890E-1, + 1.34204006088543189037E-2, + 3.28014464682127739104E-4, + 2.89247864745380683936E-6, + 6.79019408009981274425E-9, + }; + + if (y0 == zero) { + return -std::numeric_limits::infinity(); + } + if (y0 == one) { + return std::numeric_limits::infinity(); + } + if (y0 < zero || y0 > one) { + return std::numeric_limits::quiet_NaN(); + } + bool code = true; + T y = y0; + if (y > one - T{0.13533528323661269189}) { /* 0.135... = exp(-2) */ + y = one - y; + code = false; + } + + if (y > T{0.13533528323661269189}) { + y = y - T{0.5}; + const T y2 = y * y; + T x = y + y * (y2 * polevl(y2, P0, 4) / polevl(y2, Q0, 8)); + return (x * s2pi); + } + + T x = ::sqrt(T{-2.0} * ::log(y)); + const T x0 = x - ::log(x) / x; + + const T z = one / x; + T x1; + if (x < T{8.0}) /* y > exp(-32) = 1.2664165549e-14 */ + { + x1 = z * polevl(z, P1, 8) / polevl(z, Q1, 8); + } else { + x1 = z * polevl(z, P2, 8) / polevl(z, Q2, 8); + } + x = x0 - x1; + if (code) { + x = -x; + } + return x; +} diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index aeb83025784ba..c7c2b545db013 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -66,6 +66,7 @@ CREATE_UNARY_FLOAT_META_FUNC(special_entr) CREATE_UNARY_FLOAT_META_FUNC(special_i0e) CREATE_UNARY_FLOAT_META_FUNC(special_i1) CREATE_UNARY_FLOAT_META_FUNC(special_i1e) +CREATE_UNARY_FLOAT_META_FUNC(special_ndtri) CREATE_UNARY_FLOAT_META_FUNC(sqrt) CREATE_UNARY_FLOAT_META_FUNC(tan) CREATE_UNARY_FLOAT_META_FUNC(tanh) @@ -170,6 +171,7 @@ CREATE_UNARY_TORCH_IMPL_FUNC(special_entr_out, special_entr_stub) CREATE_UNARY_TORCH_IMPL_FUNC(special_i0e_out, special_i0e_stub) CREATE_UNARY_TORCH_IMPL_FUNC(special_i1e_out, special_i1e_stub) CREATE_UNARY_TORCH_IMPL_FUNC(special_i1_out, special_i1_stub) +CREATE_UNARY_TORCH_IMPL_FUNC(special_ndtri_out, special_ndtri_stub) CREATE_UNARY_TORCH_IMPL_FUNC(sqrt_out, sqrt_stub) CREATE_UNARY_TORCH_IMPL_FUNC(tan_out, tan_stub) CREATE_UNARY_TORCH_IMPL_FUNC(tanh_out, tanh_stub) @@ -759,6 +761,7 @@ DEFINE_DISPATCH(log10_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global- DEFINE_DISPATCH(log1p_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(log2_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(logical_not_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) +DEFINE_DISPATCH(special_ndtri_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(neg_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(nan_to_num_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) DEFINE_DISPATCH(polygamma_stub); // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) diff --git a/aten/src/ATen/native/UnaryOps.h b/aten/src/ATen/native/UnaryOps.h index 27654bc9f816f..5239847356fea 100644 --- a/aten/src/ATen/native/UnaryOps.h +++ b/aten/src/ATen/native/UnaryOps.h @@ -45,6 +45,7 @@ DECLARE_DISPATCH(unary_fn, log_stub); DECLARE_DISPATCH(unary_fn, log10_stub); DECLARE_DISPATCH(unary_fn, log1p_stub); DECLARE_DISPATCH(unary_fn, log2_stub); +DECLARE_DISPATCH(unary_fn, special_ndtri_stub); DECLARE_DISPATCH(unary_fn, neg_stub); DECLARE_DISPATCH(unary_fn, reciprocal_stub); diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp index 190160219049b..63b8836c05c80 100644 --- a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp @@ -600,6 +600,13 @@ static void frexp_kernel(TensorIteratorBase& iter) { }); } +static void ndtri_kernel(TensorIteratorBase& iter) { + TORCH_INTERNAL_ASSERT(iter.ntensors() == 2); + AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "ndtri_cpu", [&]() { + cpu_kernel(iter, [](scalar_t x) { return calc_ndtri(x); }); + }); +} + static void i0e_kernel(TensorIteratorBase& iter) { TORCH_INTERNAL_ASSERT(iter.ntensors() == 2); AT_DISPATCH_FLOATING_TYPES_AND( @@ -765,6 +772,8 @@ REGISTER_DISPATCH(frexp_stub, &CPU_CAPABILITY::frexp_kernel); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_DISPATCH(special_i0e_stub, &CPU_CAPABILITY::i0e_kernel); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +REGISTER_DISPATCH(special_ndtri_stub, &CPU_CAPABILITY::ndtri_kernel); +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_DISPATCH(special_i1_stub, &CPU_CAPABILITY::i1_kernel); // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) REGISTER_DISPATCH(special_i1e_stub, &CPU_CAPABILITY::i1e_kernel); diff --git a/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu b/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu index cf1a11cab688e..0cff8649f42ba 100644 --- a/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu +++ b/aten/src/ATen/native/cuda/UnarySpecialOpsKernel.cu @@ -114,6 +114,13 @@ void logit_kernel_cuda(TensorIteratorBase& iter, const Scalar& eps_scalar) { }); } +void ndtri_kernel_cuda(TensorIteratorBase& iter) { + AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "ndtri_cuda", [&]() { + gpu_kernel( + iter, [] GPU_LAMBDA(scalar_t a) -> scalar_t { return calc_ndtri(a); }); + }); +} + void erf_kernel_cuda(TensorIteratorBase& iter) { AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.common_dtype(), "erf_cuda", [&]() { gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t { @@ -188,6 +195,7 @@ REGISTER_DISPATCH(erfc_stub, &erfc_kernel_cuda); REGISTER_DISPATCH(erfinv_stub, &erfinv_kernel_cuda); REGISTER_DISPATCH(kaiser_window_stub, &kaiser_window_kernel_cuda); REGISTER_DISPATCH(special_entr_stub, &entr_kernel_cuda); +REGISTER_DISPATCH(special_ndtri_stub, &ndtri_kernel_cuda); } // namespace native } // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index cefe0726e4745..965ecbe17e316 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -9445,6 +9445,19 @@ dispatch: CPU, CUDA: special_entr_out +- func: special_ndtri(Tensor self) -> Tensor + structured_delegate: special_ndtri.out + python_module: special + variants: function + +- func: special_ndtri.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA: special_ndtri_out + - func: special_expm1(Tensor self) -> Tensor python_module: special variants: function diff --git a/docs/source/special.rst b/docs/source/special.rst index cc173dbc65bad..4a1f0fa5e5136 100644 --- a/docs/source/special.rst +++ b/docs/source/special.rst @@ -34,4 +34,5 @@ Functions .. autofunction:: i1e .. autofunction:: logit .. autofunction:: ndtr +.. autofunction:: ndtri .. autofunction:: xlog1py diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index c67c4e6812622..634d6d01c6109 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -1108,6 +1108,9 @@ - name: special_entr(Tensor self) -> Tensor self: grad * (-(1 + self.log())) +- name: special_ndtri(Tensor self) -> Tensor + self: grad * std::sqrt(2 * M_PI) * (result.square() / 2).exp() + # DO NOT define a backward for reshape! # reshape is special in that it sometimes returns a view, and sometimes not. # Defining a backward will make codegen spit out the forward call as diff --git a/torch/csrc/api/include/torch/special.h b/torch/csrc/api/include/torch/special.h index cf667f9412a79..c224c19b3b3dd 100644 --- a/torch/csrc/api/include/torch/special.h +++ b/torch/csrc/api/include/torch/special.h @@ -117,6 +117,14 @@ inline Tensor& erfinv_out(Tensor& result, const Tensor& self) { return torch::special_erfinv_out(result, self); } +inline Tensor ndtri(const Tensor& self) { + return torch::special_ndtri(self); +} + +inline Tensor& ndtri_out(Tensor& result, const Tensor& self) { + return torch::special_ndtri_out(result, self); +} + /// Computes the logit of input, elementwise. /// See https://pytorch.org/docs/master/special.html#torch.special.logit. /// diff --git a/torch/overrides.py b/torch/overrides.py index a0353fda65ab8..69acbb84c5dbd 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -882,6 +882,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.special.i1: lambda input: -1, torch.special.i1e: lambda input: -1, torch.special.logit: lambda input: -1, + torch.special.ndtri: lambda input: -1, torch.special.ndtr: lambda input: -1, torch.special.xlog1py: lambda input, other, out=None: -1, torch.t: lambda input: -1, diff --git a/torch/special/__init__.py b/torch/special/__init__.py index 68133bbe66f75..939b09ef70a6f 100644 --- a/torch/special/__init__.py +++ b/torch/special/__init__.py @@ -342,8 +342,10 @@ """ + r""" Args: {input} + Keyword args: {out} + Example:: >>> torch.special.i0e(torch.arange(5, dtype=torch.float32)) tensor([1.0000, 0.4658, 0.3085, 0.2430, 0.2070]) @@ -361,8 +363,10 @@ """ + r""" Args: {input} + Keyword args: {out} + Example:: >>> torch.special.i1(torch.arange(5, dtype=torch.float32)) tensor([0.0000, 0.5652, 1.5906, 3.9534, 9.7595]) @@ -381,8 +385,10 @@ """ + r""" Args: {input} + Keyword args: {out} + Example:: >>> torch.special.i1e(torch.arange(5, dtype=torch.float32)) tensor([0.0000, 0.2079, 0.2153, 0.1968, 0.1788]) @@ -408,3 +414,27 @@ >>> torch.special.ndtr(torch.tensor([-3., -2, -1, 0, 1, 2, 3])) tensor([0.0013, 0.0228, 0.1587, 0.5000, 0.8413, 0.9772, 0.9987]) """.format(**common_args)) + +ndtri = _add_docstr(_special.special_ndtri, + r""" +ndtri(input, *, out=None) -> Tensor +Computes the argument, x, for which the area under the Gaussian probability density function +(integrated from minus infinity to x) is equal to :attr:`input`, elementwise. + +.. math:: + \text{ndtri}(p) = \sqrt{2}\text{erf}^{-1}(2p - 1) + +.. note:: + Also known as quantile function for Normal Distribution. + +""" + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + >>> torch.special.ndtri(torch.tensor([0, 0.25, 0.5, 0.75, 1])) + tensor([ -inf, -0.6745, 0.0000, 0.6745, inf]) +""".format(**common_args)) diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 8faa1f1f535a0..5184a69301031 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -7341,6 +7341,12 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs): supports_inplace_autograd=False, safe_casts_outputs=True, sample_inputs_func=sample_inputs_entr), + UnaryUfuncInfo('special.ndtri', + ref=scipy.special.ndtri if TEST_SCIPY else _NOTHING, + domain=(0, 1), + aten_name='special_ndtri', + dtypes=all_types_and(torch.bool), + safe_casts_outputs=True), UnaryUfuncInfo('erf', ref=scipy.special.erf if TEST_SCIPY else _NOTHING, aliases=('special.erf', ), From c19acf816f984124bb00fa7d32f4376d2d76c331 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Sat, 19 Jun 2021 19:55:26 -0700 Subject: [PATCH 259/305] Replace TensorRT's deprecated API in `caffe2/python/trt/test_pt_onnx_trt.py` (#60236) Summary: TensorRT v8 is going to remove some functions/methods that used in test. ref: - getMaxWorkspaceSize deprecation: https://github.com/NVIDIA/TensorRT/blob/b2d60b6e1003f973983903de154a274b569006b8/include/NvInfer.h#L6984-L6993 - buildCudaEngine deprecation: https://github.com/NVIDIA/TensorRT/blob/b2d60b6e1003f973983903de154a274b569006b8/include/NvInfer.h#L7079-L7087 cc ptrblck Pull Request resolved: https://github.com/pytorch/pytorch/pull/60236 Reviewed By: gchanan Differential Revision: D29232376 Pulled By: ngimel fbshipit-source-id: 2b8a48787bf61c68a81568b6026d6afd5a83e751 --- caffe2/python/trt/test_pt_onnx_trt.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/caffe2/python/trt/test_pt_onnx_trt.py b/caffe2/python/trt/test_pt_onnx_trt.py index e066e8363a179..3ba67ea21244c 100644 --- a/caffe2/python/trt/test_pt_onnx_trt.py +++ b/caffe2/python/trt/test_pt_onnx_trt.py @@ -61,16 +61,18 @@ def setUp(self): self.image_files[index] = os.path.abspath(os.path.join(data_path, f)) if not os.path.exists(self.image_files[index]): raise FileNotFoundError(self.image_files[index] + " does not exist.") - self.labels = open(os.path.abspath(os.path.join(data_path, "class_labels.txt")), 'r').read().split('\n') + with open(os.path.abspath(os.path.join(data_path, "class_labels.txt")), 'r') as f: + self.labels = f.read().split('\n') def build_engine_onnx(self, model_file): with trt.Builder(TRT_LOGGER) as builder, builder.create_network(flags = 1) as network, trt.OnnxParser(network, TRT_LOGGER) as parser: - builder.max_workspace_size = 1 << 33 + builder_config = builder.create_builder_config() + builder_config.max_workspace_size = 1 << 33 with open(model_file, 'rb') as model: if not parser.parse(model.read()): for error in range(parser.num_errors): self.fail("ERROR: {}".format(parser.get_error(error))) - return builder.build_cuda_engine(network) + return builder.build_engine(network, builder_config) def _test_model(self, model_name, input_shape = (3, 224, 224), normalization_hint = 0): From 69b2bf70f9c0e591ce5e566afa59e19618031ead Mon Sep 17 00:00:00 2001 From: Jiakai Liu Date: Sun, 20 Jun 2021 00:38:07 -0700 Subject: [PATCH 260/305] [pytorch] fix tools/code_analyzer for llvm 11 (#60322) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60322 Test Plan: Imported from OSS Reviewed By: iseeyuan Differential Revision: D29250420 Pulled By: ljk53 fbshipit-source-id: ff7f9cbacd1d9518ed81c06fc843a90d6948f760 --- tools/code_analyzer/op_deps_pass.cpp | 53 ++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/tools/code_analyzer/op_deps_pass.cpp b/tools/code_analyzer/op_deps_pass.cpp index 565deb67b47cc..f41e032969d78 100644 --- a/tools/code_analyzer/op_deps_pass.cpp +++ b/tools/code_analyzer/op_deps_pass.cpp @@ -66,10 +66,13 @@ #include "llvm/Demangle/Demangle.h" #include "llvm/Analysis/LazyCallGraph.h" +#if LLVM_VERSION_MAJOR < 8 #include "llvm/IR/CallSite.h" +#endif #include "llvm/IR/Constant.h" #include "llvm/IR/Function.h" #include "llvm/IR/Instruction.h" +#include "llvm/IR/Instructions.h" #include "llvm/IR/Intrinsics.h" #include "llvm/IR/Module.h" #include "llvm/Pass.h" @@ -194,6 +197,10 @@ using VALUE_SET = std::unordered_set; using PATH = std::unordered_map>; +inline std::string _name(const Value* V) { + return V->getName().str(); +} + // Referenced the logic in llvm-cxxfilt.cpp. // Starting from LLVM 9 it provides a `demangle()` API. Here we keep our ad-hoc // version for backward compatibility. @@ -213,6 +220,22 @@ std::string _demangle(const std::string& mangled) { return result; } +inline bool _isCallSite(Value* V) { +#if LLVM_VERSION_MAJOR >= 8 + return isa(V); +#else + return !!CallSite(V); +#endif +} + +inline Function* _getCalledFunction(Value* V) { +#if LLVM_VERSION_MAJOR >= 8 + return dyn_cast(V)->getCalledFunction(); +#else + return CallSite(V).getCalledFunction(); +#endif +} + // LLVM_DEBUG needs opt to be built with debug support. template< typename T, @@ -276,7 +299,7 @@ class OpDependency : public ModulePass { } SET roots; for (const auto& F : visibleFuncs) { - std::string name = F->getName(); + std::string name = _name(F); auto demangled = _demangle(name); if (RootSymbolPatternLoc.pattern->match(demangled)) { roots.insert(name); @@ -300,12 +323,12 @@ class OpDependency : public ModulePass { if (F.hasDefaultVisibility()) { visibleFuncs->insert(&F); } - std::string caller = F.getName(); + std::string caller = _name(&F); std::string callerDemangled = _demangle(caller); for (BasicBlock& BB : F) { for (Instruction& I : BB) { scanReferredFunctions(I, [&](Function* func) -> void { - std::string callee = func->getName(); + std::string callee = _name(func); std::string calleeDemangled = _demangle(callee); (*deps)[caller].insert(callee); if (Verbose > 1) { @@ -363,8 +386,8 @@ class OpDependency : public ModulePass { SmallVector worklist; SmallPtrSet visited; - if (auto CS = CallSite(&I)) { - Function* callee = CS.getCalledFunction(); + if (_isCallSite(&I)) { + Function* callee = _getCalledFunction(&I); if (callee && !callee->isIntrinsic() && visited.insert(callee).second) { CB(callee); } @@ -553,7 +576,7 @@ class OpDependency : public ModulePass { if (!visitedOps->empty()) { if (Verbose) { std::cerr << "[INFO] ignore extra op schema str: " << *schemaStr - << " in: " << _demangle(src->getFunction()->getName()) + << " in: " << _demangle(_name(src->getFunction())) << ", because already found valid op schema str: " << *visitedOps->begin() << std::endl; } @@ -570,10 +593,10 @@ class OpDependency : public ModulePass { return; } if (visitedFunctions) { - (*visitedFunctions).insert(F->getName()); + (*visitedFunctions).insert(_name(F)); } if (Verbose > 1) { - std::cerr << "[DEBUG][FUNC] " << _demangle(F->getName()) << std::endl; + std::cerr << "[DEBUG][FUNC] " << _demangle(_name(F)) << std::endl; printDebugPath(debugPath.get(), src, V); } } @@ -624,7 +647,7 @@ class OpDependency : public ModulePass { for (auto V : instructions) { auto I = dyn_cast(V); // We only need to process call/invoke instructions. - if (!I || !CallSite(I)) { + if (!I || !_isCallSite(I)) { continue; } auto contextualNamespace = inferContextualNamespace(I); @@ -648,7 +671,7 @@ class OpDependency : public ModulePass { std::cerr << op << " "; } std::cerr << ") in a registration call in function: " - << _demangle(I->getFunction()->getName()) + << _demangle(_name(I->getFunction())) << " contextualNamespace: " << contextualNamespace << std::endl; } @@ -657,7 +680,7 @@ class OpDependency : public ModulePass { if (visitedFunctions.empty()) { std::cerr << "[WARNING] could not find registered function for op: " << op << " in function: " - << _demangle(I->getFunction()->getName()) + << _demangle(_name(I->getFunction())) << " contextualNamespace: " << contextualNamespace << std::endl; } @@ -673,7 +696,7 @@ class OpDependency : public ModulePass { } static std::string inferContextualNamespace(Instruction* I) { - auto functionName = _demangle(I->getFunction()->getName()); + auto functionName = _demangle(_name(I->getFunction())); for (auto& pattern : TorchLibraryInitPattern) { if (!pattern.pattern->match(functionName)) { continue; @@ -717,13 +740,13 @@ class OpDependency : public ModulePass { for (auto V : instructions) { auto I = dyn_cast(V); // We only need to process call/invoke instructions. - if (!I || !CallSite(I)) { + if (!I || !_isCallSite(I)) { continue; } if (Verbose > 2) { std::cerr << "[DEBUG][CALL][INST] " << *I << std::endl; } - std::string caller = I->getFunction()->getName(); + std::string caller = _name(I->getFunction()); SET visitedOps; scanOpSchemaStrAndFunction(I, {}, {}, &visitedOps, nullptr); if (visitedOps.size() != 1) { @@ -820,7 +843,7 @@ class OpDependency : public ModulePass { static void printDebugValue(Value* V) { if (auto F = dyn_cast(V)) { - std::cerr << "[FUNC] " << _demangle(F->getName()); + std::cerr << "[FUNC] " << _demangle(_name(F)); } else if (isa(V)) { std::cerr << "[CONST] " << *V; } else if (isa(V)) { From 0131a5972d586a5236aef01a8b92427f41b45a8e Mon Sep 17 00:00:00 2001 From: Rohan Varma Date: Sun, 20 Jun 2021 12:01:41 -0700 Subject: [PATCH 261/305] [DDP] Test inference works with eval() and no_grad() (#59666) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59666 Tests that inference with DDP model won't hang when user sets eval() or no_grad(). Note that if the model has a syncBN layer, they need both eval() and no_grad() as eval() makes SyncBN work like a regular BN layer. ghstack-source-id: 131906625 Test Plan: CI Reviewed By: SciPioneer Differential Revision: D28974146 fbshipit-source-id: 137f8245b1c303beb2416518476e70fe67c73376 --- .../_internal/distributed/distributed_test.py | 56 +++++++++++++++++-- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/torch/testing/_internal/distributed/distributed_test.py b/torch/testing/_internal/distributed/distributed_test.py index c41580f701454..401be4ae2876c 100644 --- a/torch/testing/_internal/distributed/distributed_test.py +++ b/torch/testing/_internal/distributed/distributed_test.py @@ -6943,10 +6943,58 @@ def test_ddp_multiple_nested_unused_params_err_ignore_params(self): # certain parameters. self._test_ddp_multiple_nested_unused_params_error(ignore_sparse=True) - @unittest.skipIf( - BACKEND != "nccl" and BACKEND != "gloo", - "Only Nccl & Gloo backend support DistributedDataParallel", - ) + @unittest.skipIf(BACKEND != 'nccl' and BACKEND != 'gloo', + "Only Nccl & Gloo backend support DistributedDataParallel") + @skip_if_lt_x_gpu(2) + def test_ddp_inference(self): + # tests that DDP module can be run on a single node with no_grad + # or eval setting and there is no hang. + rank = self.rank + torch.cuda.set_device(rank) + model = Net().cuda() + local_model = copy.deepcopy(model) + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank], + ) + syncbn_model = nn.SyncBatchNorm( + 2, momentum=0.99, track_running_stats=False + ).cuda() + local_syncbn_model = copy.deepcopy(model) + syncbn_model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank] + ) + inp = torch.randn(10, 2, device=rank) + inp_syncbn = torch.randn(10, 2, 4, 4, device=rank) + tests = [ + (model, local_model, inp), + (syncbn_model, local_syncbn_model, inp_syncbn), + ] + for test in tests: + test_model, test_local_model, test_inp = test + if self.rank == 0: + with torch.no_grad(): + for _ in range(6): + self.assertEqual( + test_model(test_inp), + test_local_model(test_inp) + ) + + model.eval() + for _ in range(6): + self.assertEqual( + test_model(test_inp), + test_local_model(test_inp) + ) + + # Barrier since only rank 0 runs inference. Test should be + # much faster than 30s, but this is to avoid flakiness. + self._barrier(timeout=30) + + + @unittest.skipIf(BACKEND != 'nccl' and BACKEND != 'gloo', + "Only Nccl & Gloo backend support DistributedDataParallel") @skip_if_lt_x_gpu(2) def test_ddp_sync_bn_training_vs_eval(self): rank = self.rank From b298013cd5b8b081f7b512e9b0ea7da05dcbe63b Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Sun, 20 Jun 2021 19:04:04 -0700 Subject: [PATCH 262/305] [add/sub] Cast `alpha` to `acc_type` (#60227) Summary: This PR lets `torch.add` & `torch.sub` CUDA kernels cast `alpha` to `acc_type`, not `scalar_t`. I do not remove `cast`s from `test/test_foreach.py` because I'll do this in https://github.com/pytorch/pytorch/issues/59907 or follow-up for it. Current upstream `torch._foreach_add` & `torch._foreach_sub` upcast `alpha` parameter to `acc_type` while `torch.add` & `torch.sub` not. This is kind of problematic because outputs of `torch.add` and `torch.sub` are different from `torch._foreach_add` and `torch._foreach_sub`, respectively if the dtype of input tensors is either `torch.half` or `torch.bfloat16`. The discrepancy is proportional-ish to `abs(alpha)` except when `alpha` is representable with 16 bits. ref: - `torch._foreach_add` & `torch._foreach_sub` cast `alpha`: https://github.com/pytorch/pytorch/blob/6d0fb85a623f5ef3f3f1a2afc3660cb71fa70511/aten/src/ATen/native/cuda/ForeachBinaryOpList.cu#L21-L28, `BinaryOpListAlphaFunctor` is defined here: https://github.com/pytorch/pytorch/blob/6d0fb85a623f5ef3f3f1a2afc3660cb71fa70511/aten/src/ATen/native/cuda/ForeachFunctors.cuh#L202 related: https://github.com/pytorch/pytorch/issues/58833, https://github.com/pytorch/pytorch/pull/59907 cc ngimel ptrblck mcarilli Pull Request resolved: https://github.com/pytorch/pytorch/pull/60227 Reviewed By: mruberry Differential Revision: D29252759 Pulled By: ngimel fbshipit-source-id: 847f3b9493ae30a900f7445af00aef1abcc1ab21 --- aten/src/ATen/native/cuda/BinaryAddSubKernel.cu | 10 ++++++---- test/test_binary_ufuncs.py | 6 ++++++ 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/aten/src/ATen/native/cuda/BinaryAddSubKernel.cu b/aten/src/ATen/native/cuda/BinaryAddSubKernel.cu index b8152da0819bf..da4735b5c4554 100644 --- a/aten/src/ATen/native/cuda/BinaryAddSubKernel.cu +++ b/aten/src/ATen/native/cuda/BinaryAddSubKernel.cu @@ -1,3 +1,4 @@ +#include #include #include #include @@ -8,19 +9,20 @@ namespace at { namespace native { -template +template struct AddFunctor { - AddFunctor(scalar_t a): alpha(a) {} + AddFunctor(accscalar_t a): alpha(a) {} __device__ __forceinline__ scalar_t operator() (const scalar_t a, const scalar_t b) const { return a + alpha * b; } private: - scalar_t alpha; + accscalar_t alpha; }; void add_kernel_cuda(TensorIteratorBase& iter, const Scalar& alpha_scalar) { AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(kHalf, kBool, kBFloat16, iter.common_dtype(), "add_cuda/sub_cuda", [&]() { - AddFunctor f(alpha_scalar.to()); + using accscalar_t = at::acc_type; + AddFunctor f(alpha_scalar.to()); gpu_kernel_with_scalars(iter, f); }); } diff --git a/test/test_binary_ufuncs.py b/test/test_binary_ufuncs.py index 741014d3f3503..e675649aa4c94 100644 --- a/test/test_binary_ufuncs.py +++ b/test/test_binary_ufuncs.py @@ -2355,6 +2355,12 @@ def test_add(self, device): self.assertRaisesRegex(RuntimeError, r"result type ComplexFloat can't be cast to the desired output type Double", lambda: torch.add(m1, m1, out=m2)) + @onlyCUDA + def test_add_half_tensor_with_alpha(self, device): + x = torch.tensor([60000.0], dtype=torch.half, device=device) + y = torch.tensor([-60000.0], dtype=torch.half, device=device) + actual = torch.add(x, y, alpha=2) + self.assertTrue(not (actual.isnan() or actual.isinf())) def test_sub_typing(self, device): m1 = torch.tensor([True, False, False, True, False, False], dtype=torch.bool, device=device) From 47d727fe1bc8d320c66c305a23b09d0f218fac96 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Sun, 20 Jun 2021 19:24:42 -0700 Subject: [PATCH 263/305] [quant][graphmode][fx] Produce conv reference static quant modules (#60138) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60138 Test Plan: python test/test_quantization.py TestQuantizeFx Imported from OSS Reviewed By: vkuzo Differential Revision: D29184791 fbshipit-source-id: 971a40012dbba0cf687c62a3a4af9358513c253b --- test/quantization/fx/test_quantize_fx.py | 85 +++++++++++++++---- .../quantized/_reference/modules/conv_relu.py | 10 +-- .../quantization/fx/quantization_patterns.py | 2 +- torch/quantization/quantization_mappings.py | 34 +++++++- 4 files changed, 108 insertions(+), 23 deletions(-) diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py index e6f6cad3408d6..8b49d1f4f503e 100644 --- a/test/quantization/fx/test_quantize_fx.py +++ b/test/quantization/fx/test_quantize_fx.py @@ -2,6 +2,7 @@ import torch.nn.functional as F import torch.nn as nn import torch.nn.quantized as nnq +import torch.nn.quantized._reference as nnqr import torch.nn.quantized.dynamic as nnqd import torch.nn.intrinsic as nni import torch.nn.intrinsic.quantized as nniq @@ -312,12 +313,12 @@ def forward(self, x, y): if n.op == 'call_module' and type(modules[n.target]) == nn.ReLU: self.assertTrue(is_match(modules, n, pattern)) - def _get_conv_linear_test_cases(self): + def _get_conv_linear_test_cases(self, is_reference): """ Returns a list of test cases, with format: is_dynamic, ModuleClass, module_constructor_inputs, inputs, quantized_node, weight_prepack_op """ - class Conv1d(torch.nn.Module): + class FunctionalConv1d(torch.nn.Module): def __init__(self, weight): super().__init__() self.weight = torch.nn.Parameter(weight) @@ -329,10 +330,20 @@ def __init__(self, weight): def forward(self, x): return F.conv1d(x, self.weight, None, self.stride, self.padding, self.dilation, self.groups) + + class Conv1d(torch.nn.Module): + def __init__(self, *args): + super().__init__() + self.conv = torch.nn.Conv1d(*args) + + def forward(self, x): + return self.conv(x) + conv1d_input = torch.rand(1, 3, 224) conv1d_weight = torch.rand(3, 3, 3) + conv1d_module_args = (3, 3, 3) - class Conv2d(torch.nn.Module): + class FunctionalConv2d(torch.nn.Module): def __init__(self, weight): super().__init__() self.weight = torch.nn.Parameter(weight) @@ -344,10 +355,19 @@ def __init__(self, weight): def forward(self, x): return F.conv2d(x, self.weight, None, self.stride, self.padding, self.dilation, self.groups) + class Conv2d(torch.nn.Module): + def __init__(self, *args): + super().__init__() + self.conv = torch.nn.Conv2d(*args) + + def forward(self, x): + return self.conv(x) + conv2d_input = torch.rand(1, 3, 224, 224) conv2d_weight = torch.rand(3, 3, 3, 3) + conv2d_module_args = (3, 3, 3) - class Conv3d(torch.nn.Module): + class FunctionalConv3d(torch.nn.Module): def __init__(self, weight): super().__init__() self.weight = torch.nn.Parameter(weight) @@ -367,8 +387,17 @@ def forward(self, x): self.groups, ) + class Conv3d(torch.nn.Module): + def __init__(self, *args): + super().__init__() + self.conv = torch.nn.Conv3d(*args) + + def forward(self, x): + return self.conv(x) + conv3d_input = torch.rand(1, 3, 32, 224, 224) conv3d_weight = torch.rand(3, 3, 3, 3, 3) + conv3d_module_args = (3, 3, 3) class Linear(torch.nn.Module): def __init__(self, weight): @@ -391,37 +420,63 @@ def forward(self, x): linear_module_input = torch.rand(8, 5) + # is_dynamic, ModuleClass, module_constructor_inputs, + # inputs, quantized_node, weight_prepack_node tests = [ ( False, - Conv1d, + FunctionalConv1d, (conv1d_weight,), (conv1d_input,), - ns.call_function(torch.ops.quantized.conv1d), + ns.call_function(torch.nn.functional.conv1d if is_reference else torch.ops.quantized.conv1d) , ns.call_function(torch.ops.quantized.conv1d_prepack), ), ( False, - Conv2d, + FunctionalConv2d, (conv2d_weight,), (conv2d_input,), - ns.call_function(torch.ops.quantized.conv2d), + ns.call_function(torch.nn.functional.conv2d if is_reference else torch.ops.quantized.conv2d), ns.call_function(torch.ops.quantized.conv2d_prepack), ), ( False, - Conv3d, + FunctionalConv3d, (conv3d_weight,), (conv3d_input,), - ns.call_function(torch.ops.quantized.conv3d), + ns.call_function(torch.nn.functional.conv3d if is_reference else torch.ops.quantized.conv3d), ns.call_function(torch.ops.quantized.conv3d_prepack), ), + ( + False, + Conv1d, + conv1d_module_args, + (conv1d_input,), + ns.call_module(nnqr.Conv1d if is_reference else nnq.Conv1d), + None + ), + ( + False, + Conv2d, + conv2d_module_args, + (conv2d_input,), + ns.call_module(nnqr.Conv2d if is_reference else nnq.Conv2d), + None + ), + ( + False, + Conv3d, + conv3d_module_args, + (conv3d_input,), + ns.call_module(nnqr.Conv3d if is_reference else nnq.Conv3d), + None + ), ( True, Linear, (linear_weight,), (linear_input,), - ns.call_function(torch.ops.quantized.linear_dynamic), + None if is_reference else ns.call_function(torch.ops.quantized.linear_dynamic), ns.call_function(torch.ops.quantized.linear_prepack), ), ( @@ -429,7 +484,7 @@ def forward(self, x): Linear, (linear_weight,), (linear_input,), - ns.call_function(torch.ops.quantized.linear), + ns.call_function(torch.nn.functional.linear if is_reference else torch.ops.quantized.linear), ns.call_function(torch.ops.quantized.linear_prepack), ), ( @@ -458,7 +513,7 @@ def forward(self, x): def test_functional_not_reference(self): """ Test quantizing functional conv and linear """ - tests = self._get_conv_linear_test_cases() + tests = self._get_conv_linear_test_cases(is_reference=False) for (is_dynamic, ModuleClass, module_constructor_inputs, inputs, quantized_node, weight_prepack_node) in tests: quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC @@ -476,17 +531,17 @@ def test_functional_not_reference(self): def test_functional_reference(self): """ Test quantizing functional conv and linear with reference option """ - tests = self._get_conv_linear_test_cases() + tests = self._get_conv_linear_test_cases(is_reference=True) for (is_dynamic, ModuleClass, module_constructor_inputs, inputs, quantized_node, weight_prepack_node) in tests: quant_type = QuantType.DYNAMIC if is_dynamic else QuantType.STATIC node_occurrence = dict() if weight_prepack_node: node_occurrence[weight_prepack_node] = 0 - node_occurrence[quantized_node] = 0 self.checkGraphModeFxOp( ModuleClass(*module_constructor_inputs), inputs, quant_type, + expected_node=quantized_node, expected_node_occurrence=node_occurrence, is_reference=True) diff --git a/torch/nn/intrinsic/quantized/_reference/modules/conv_relu.py b/torch/nn/intrinsic/quantized/_reference/modules/conv_relu.py index 712f048ddca2e..b0305f6207d95 100644 --- a/torch/nn/intrinsic/quantized/_reference/modules/conv_relu.py +++ b/torch/nn/intrinsic/quantized/_reference/modules/conv_relu.py @@ -3,14 +3,14 @@ import torch.nn.functional as F class ConvReLU1d(nnqr.Conv1d): - _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU1d # type: ignore[assignment] + _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU1d def forward(self, x: torch.Tensor) -> torch.Tensor: x_dequant = x.dequantize() weight_dequant = self._qweight.dequantize() float_result = F.conv1d( - x_dequant, weight_dequant, self._bias, self._conv1d_stride, - self._conv1d_padding, self._conv1d_dilation, self.groups) + x_dequant, weight_dequant, self._bias, self._conv1d_stride, # type: ignore[has-type] + self._conv1d_padding, self._conv1d_dilation, self.groups) # type: ignore[has-type] float_result = F.relu(float_result, inplace=True) # NEEDFIX: we don't have dtype in the Linear module APIs right now! result = torch.quantize_per_tensor( @@ -22,7 +22,7 @@ def _get_name(self): class ConvReLU2d(nnqr.Conv2d): - _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU2d # type: ignore[assignment] + _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU2d def forward(self, x: torch.Tensor) -> torch.Tensor: x_dequant = x.dequantize() @@ -40,7 +40,7 @@ def _get_name(self): return "QuantizedConvReLU2d(Reference)" class ConvReLU3d(nnqr.Conv3d): - _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU3d # type: ignore[assignment] + _FLOAT_MODULE = torch.nn.intrinsic.ConvReLU3d def forward(self, x: torch.Tensor) -> torch.Tensor: x_dequant = x.dequantize() diff --git a/torch/quantization/fx/quantization_patterns.py b/torch/quantization/fx/quantization_patterns.py index 7023b6cfdb8d7..ed1474e7acf3f 100644 --- a/torch/quantization/fx/quantization_patterns.py +++ b/torch/quantization/fx/quantization_patterns.py @@ -497,7 +497,7 @@ def convert(self, self.conv.activation_post_process = output_activation_post_process # 2. select quantized class qconv_cls = get_static_quant_module_class( - type(self.conv), additional_static_quant_mapping) + type(self.conv), additional_static_quant_mapping, is_reference=is_reference) quantized = qconv_cls.from_float(self.conv) parent_name, name = _parent_name(self.conv_node.target) setattr(modules[parent_name], name, quantized) diff --git a/torch/quantization/quantization_mappings.py b/torch/quantization/quantization_mappings.py index 581299a45d89f..a80e8ec2a4464 100644 --- a/torch/quantization/quantization_mappings.py +++ b/torch/quantization/quantization_mappings.py @@ -6,8 +6,10 @@ import torch.nn.functional as F import torch.nn.intrinsic as nni import torch.nn.intrinsic.quantized as nniq +import torch.nn.intrinsic.quantized._reference as nniqr import torch.nn.intrinsic.qat as nniqat import torch.nn.quantized as nnq +import torch.nn.quantized._reference as nnqr import torch.nn.quantized.dynamic as nnqd import torch.nn.qat as nnqat @@ -20,6 +22,31 @@ ) from .utils import get_combined_dict +# Default map for swapping float module to reference quantized modules +DEFAULT_REFERENCE_STATIC_QUANT_MODULE_MAPPINGS : Dict[Callable, Any] = { + nn.Conv1d: nnqr.Conv1d, + nn.Conv2d: nnqr.Conv2d, + nn.Conv3d: nnqr.Conv3d, + # nn.Linear, nnqr.Linear, + nni.ConvReLU1d: nniqr.ConvReLU1d, + nni.ConvReLU2d: nniqr.ConvReLU2d, + nni.ConvReLU3d: nniqr.ConvReLU3d, + nni.LinearReLU: nniqr.LinearReLU, + # QAT Modules + # nnqat.Linear: nnqr.Linear, + nnqat.Conv2d: nnqr.Conv2d, + nnqat.Conv3d: nnqr.Conv3d, + nniqat.ConvBn1d: nnqr.Conv1d, + nniqat.ConvBn2d: nnqr.Conv2d, + nniqat.ConvBn3d: nnqr.Conv3d, + nniqat.ConvBnReLU1d: nniqr.ConvReLU1d, + nniqat.ConvBnReLU2d: nniqr.ConvReLU2d, + nniqat.ConvBnReLU3d: nniqr.ConvReLU3d, + nniqat.ConvReLU2d: nniqr.ConvReLU2d, + nniqat.ConvReLU3d: nniqr.ConvReLU3d, + # nniqat.LinearReLU: nniqr.LinearReLU, +} + # Default map for swapping float module to quantized ones DEFAULT_STATIC_QUANT_MODULE_MAPPINGS : Dict[Callable, Any] = { QuantStub: nnq.Quantize, @@ -134,13 +161,16 @@ def get_default_static_quant_module_mappings() -> Dict[Callable, Any]: def get_static_quant_module_class( float_module_class: Callable, - additional_static_quant_mapping: Optional[Dict[Callable, Any]] = None) -> Any: + additional_static_quant_mapping: Optional[Dict[Callable, Any]] = None, + is_reference: bool = False) -> Any: r"""n Get the statically quantized module class corresponding to the floating point module class """ if additional_static_quant_mapping is None: additional_static_quant_mapping = {} - all_mappings = get_combined_dict(DEFAULT_STATIC_QUANT_MODULE_MAPPINGS, additional_static_quant_mapping) + all_mappings = get_combined_dict( + DEFAULT_REFERENCE_STATIC_QUANT_MODULE_MAPPINGS if is_reference + else DEFAULT_STATIC_QUANT_MODULE_MAPPINGS, additional_static_quant_mapping) static_quant_module_class = all_mappings.get(float_module_class, None) assert static_quant_module_class is not None, \ "Floating point module class {}".format(str(float_module_class)) + \ From a516424a70f03cef10c7274642eb12ca3398592f Mon Sep 17 00:00:00 2001 From: Ivan Yashchuk Date: Sun, 20 Jun 2021 19:36:04 -0700 Subject: [PATCH 264/305] Update internal code for torch.linalg.solve (#56613) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56613 Replace linalg_solve_helper with `lu_stub` + `lu_solve_stub`. Once `lu_stub` and `lu_solve_stub` have cuSOLVER-based codepath, `torch.linalg.solve` will have it as well. Test Plan: Imported from OSS Reviewed By: zou3519 Differential Revision: D28627408 Pulled By: mruberry fbshipit-source-id: b95bbdf35f845a56a1489c04b53742a01b36e789 --- aten/src/ATen/native/BatchLinearAlgebra.cpp | 24 +++++++------------ .../ATen/native/cuda/BatchLinearAlgebra.cu | 20 +++++++--------- aten/src/ATen/native/native_functions.yaml | 10 ++------ .../check_backward_compatibility.py | 1 + 4 files changed, 20 insertions(+), 35 deletions(-) diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp index f3e1a1a28a3e0..fad0976060ea3 100644 --- a/aten/src/ATen/native/BatchLinearAlgebra.cpp +++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp @@ -792,19 +792,6 @@ std::tuple solve_out(const Tensor& self, const Tensor& A, Tenso return std::tuple(solution, lu); } - -// This is a type dispatching helper function for 'apply_solve' -Tensor& _linalg_solve_out_helper_cpu(Tensor& result, Tensor& input, Tensor& infos) { - // 'result' and 'input' should be in column major order (it should be checked before calling this function) - // the content of 'result', 'input' and 'infos' is overwritten by 'apply_solve' - // 'result' should contain data of 'other' tensor (right-hand-side of the linear system of equations) - // 'input' should contain data of original 'input' tensor (left-hand-side of the linear system of equations) - AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(result.scalar_type(), "linalg_solve_out_cpu", [&]{ - apply_solve(result, input, infos); - }); - return result; -} - // Solves a system of linear equations matmul(input, x) = other in-place // LAPACK/MAGMA error codes are saved in 'infos' tensor, they are not checked here static Tensor& linalg_solve_out_info(Tensor& result, Tensor& infos, const Tensor& input, const Tensor& other) { @@ -896,7 +883,7 @@ static Tensor& linalg_solve_out_info(Tensor& result, Tensor& infos, const Tensor result = result.unsqueeze_(-1); } - // _linalg_solve_out_helper_ (apply_solve) performs calculations in-place and result must be a copy of other_broadcasted + // lu_stub+lu_solve_stub perform calculations in-place and 'result' must be a copy of 'other_broadcasted' result.copy_(other_broadcasted); auto input_working_copy = cloneBatchedColumnMajor(input_broadcasted); @@ -909,7 +896,14 @@ static Tensor& linalg_solve_out_info(Tensor& result, Tensor& infos, const Tensor infos.fill_(0); } - result = at::_linalg_solve_out_helper_(result, input_working_copy, infos); + // compute the LU factorization of 'input_working_copy' + auto pivots_shape = IntArrayRef(input_broadcasted.sizes().data(), input_broadcasted.dim() - 2).vec(); // input_broadcasted.shape[:-2] + pivots_shape.push_back(std::min(input.size(-2), input.size(-1))); + Tensor pivots = at::empty(pivots_shape, input.options().dtype(kInt)); + lu_stub(input.device().type(), input_working_copy, pivots, infos, /*compute_pivots=*/true); + + // solve the linear system using the LU factorization + lu_solve_stub(input.device().type(), result, input_working_copy, pivots); // for 1-dimensional 'other', we need to squeeze the result after "apply_solve" if (vector_case) { diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu index 0daace358c76c..9b22fbe1ee94b 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu @@ -1320,18 +1320,6 @@ std::tuple _solve_helper_cuda(const Tensor& self, const Tensor& return std::tuple(self_working_copy, A_working_copy); } -// This is a type dispatching helper function for 'apply_solve' -Tensor& _linalg_solve_out_helper_cuda(Tensor& result, Tensor& input, Tensor& infos) { - // 'result' and 'input' should be in column major order (it should be checked before calling this function) - // the content of 'result', 'input' and 'infos' is overwritten by 'apply_solve' - // 'result' should contain data of 'other' tensor (right-hand-side of the linear system of equations) - // 'input' should contain data of origianl 'input' tensor (left-hand-side of the linear system) - AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(result.scalar_type(), "linalg_solve_out_cpu", [&]{ - apply_solve(result, input, infos); - }); - return result; -} - // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ inverse ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ /* @@ -1900,6 +1888,10 @@ static void apply_lu_batched_magma(const Tensor& input, const Tensor& pivots, co if (compute_pivots) { auto pivots_data = pivots.data_ptr(); auto pivots_stride = pivots.size(-1); + // fill pivots with ones to avoid memory access violations inside magma kernels + // magmaLuBatched might not set the values for it + // see https://github.com/pytorch/pytorch/pull/53064 + pivots.fill_(1); magma_int_t** pivots_array; ALLOCATE_ARRAY(pivots_array, magma_int_t*, batch_size); for (int64_t i = 0; i < batch_size; i++) { @@ -1914,6 +1906,10 @@ static void apply_lu_batched_magma(const Tensor& input, const Tensor& pivots, co Tensor pivots_tmp = at::arange(1, k + 1, input.options().dtype(at::kInt)).expand_as(pivots); pivots.copy_(pivots_tmp); } + + // block CPU until all operations on the queue are finished + // this explicit sync prevents garbage results from the subsequent magmaLuSolveBatched call from a different queue + magma_queue_sync(magma_queue.get_queue()); #endif } diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 965ecbe17e316..3bd8018c8d0c0 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -10019,22 +10019,16 @@ python_module: linalg variants: function -- func: _linalg_solve_out_helper_(Tensor(a!) self, Tensor(b!) other, Tensor(c!) infos) -> Tensor(a!) - variants: function - dispatch: - CPU: _linalg_solve_out_helper_cpu - CUDA: _linalg_solve_out_helper_cuda - - func: linalg_solve(Tensor input, Tensor other) -> Tensor python_module: linalg variants: function dispatch: - CompositeExplicitAutograd: linalg_solve + CPU, CUDA: linalg_solve - func: linalg_solve.out(Tensor input, Tensor other, *, Tensor(a!) out) -> Tensor(a!) python_module: linalg dispatch: - CompositeExplicitAutograd: linalg_solve_out + CPU, CUDA: linalg_solve_out - func: linalg_tensorinv(Tensor self, int ind=2) -> Tensor python_module: linalg diff --git a/test/backward_compatibility/check_backward_compatibility.py b/test/backward_compatibility/check_backward_compatibility.py index d15d446cfbd98..3ea5764f5bb1a 100644 --- a/test/backward_compatibility/check_backward_compatibility.py +++ b/test/backward_compatibility/check_backward_compatibility.py @@ -47,6 +47,7 @@ ("aten::_svd_helper", datetime.date(2021, 1, 31)), ("aten::_syevd_helper", datetime.date(9999, 1, 1)), ("aten::_lu_solve_helper", datetime.date(9999, 1, 1)), + ("aten::_linalg_solve_out_helper_", datetime.date(9999, 1, 1)), ("aten::_cudnn_rnn_flatten_weight", datetime.date(2020, 12, 31)), ("aten::_cudnn_rnn", datetime.date(2020, 12, 31)), ("aten::_cudnn_rnn_backward", datetime.date(2020, 12, 31)), From 2293ab4e53229e8729e0a89478b97f146e224362 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Sun, 20 Jun 2021 19:40:12 -0700 Subject: [PATCH 265/305] [quant][graphmode][fx] Refactor convert for linear to use get_static_module_mapping and get_dynamic_module_mapping (#60151) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60151 Test Plan: ``` python test/test_quantization.py TestQuantizeFx python test/test_quantization.py TestQuantizeFxOps ``` Imported from OSS Reviewed By: vkuzo Differential Revision: D29188264 fbshipit-source-id: d2b77ffcf4b7446fc6c43248e43218092d2a6aea --- .../quantization/fx/quantization_patterns.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/torch/quantization/fx/quantization_patterns.py b/torch/quantization/fx/quantization_patterns.py index ed1474e7acf3f..699b2407124b1 100644 --- a/torch/quantization/fx/quantization_patterns.py +++ b/torch/quantization/fx/quantization_patterns.py @@ -4,8 +4,6 @@ Node, Graph, ) -import torch.nn.quantized as nnq -import torch.nn.quantized.dynamic as nnqd from torch.quantization import ( default_affine_fixed_qparams_fake_quant, default_symmetric_fixed_qparams_fake_quant, @@ -623,6 +621,8 @@ def convert(self, load_arg: Callable, is_reference: bool = False, convert_custom_config_dict: Dict[str, Any] = None) -> Node: + if convert_custom_config_dict is None: + convert_custom_config_dict = {} # Supported combinations are: # quant_type | activation (compute_type) | weight # static quint8 qint8 @@ -671,14 +671,17 @@ def convert(self, self.linear.activation_post_process = output_activation_post_process # 2. select corresponding quantized linear class for the float linear class - if type(self.linear) in [torch.nn.Linear, torch.nn.qat.Linear]: - qlinear = nnq.Linear if activation_int8_quantized else nnqd.Linear - elif type(self.linear) in [torch.nn.intrinsic.LinearReLU, torch.nn.intrinsic.qat.LinearReLU]: - assert activation_int8_quantized, \ - 'Only int8 static quantization is supported for LinearReLU' - qlinear = torch.nn.intrinsic.quantized.LinearReLU + if activation_int8_quantized: + additional_static_quant_mapping = convert_custom_config_dict.get("static", {}) + qlinear = get_static_quant_module_class(type(self.linear), additional_static_quant_mapping) else: - raise Exception("unhandled linear type:", type(self.linear)) + assert dtypes in [ + (torch.float32, torch.qint8, torch.quint8), + (torch.float32, torch.float16, None), + ], f"dtype {dtypes} not supported yet" + additional_dynamic_quant_mapping = convert_custom_config_dict.get("dynamic", {}) + qlinear = get_dynamic_quant_module_class(type(self.linear), additional_dynamic_quant_mapping) + quantized = qlinear.from_float(self.linear) parent_name, name = _parent_name(self.linear_node.target) setattr(modules[parent_name], name, quantized) From 510334f34ba1f6ff5c1418f92bfa31c0d7399b6f Mon Sep 17 00:00:00 2001 From: "Rong Rong (AI Infra)" Date: Sun, 20 Jun 2021 19:41:57 -0700 Subject: [PATCH 266/305] [BE] clean up IS_PYTORCH_CI and IN_CI (#60279) Summary: `IS_PYTORCH_CI` and `IN_CI` are used randomly, however in some cases IN_CI is not currently set because it only exist in .circleci/scripts/setup_ci_environment.sh. This cleans up the 2 flags and only use IN_CI Pull Request resolved: https://github.com/pytorch/pytorch/pull/60279 Test Plan: CI Reviewed By: seemethere Differential Revision: D29239545 Pulled By: walterddr fbshipit-source-id: a069424a2bb8790a3adfdaf0dc460301026bf8c7 --- .jenkins/pytorch/common.sh | 5 ++++- test/distributed/elastic/multiprocessing/api_test.py | 4 ++-- test/distributed/rpc/test_faulty_agent.py | 4 ++-- test/distributed/rpc/test_process_group_agent.py | 4 ++-- test/distributed/rpc/test_tensorpipe_agent.py | 4 ++-- test/test_dataloader.py | 4 ++-- test/test_ops.py | 4 ++-- torch/testing/_internal/common_utils.py | 7 ++----- 8 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.jenkins/pytorch/common.sh b/.jenkins/pytorch/common.sh index ed80984041bac..0d31dcebbc90f 100644 --- a/.jenkins/pytorch/common.sh +++ b/.jenkins/pytorch/common.sh @@ -29,7 +29,10 @@ fi # system; to find out more, grep for this string in ossci-job-dsl. echo "ENTERED_USER_LAND" -export IS_PYTORCH_CI=1 +# Previously IN_CI is only set in .circleci/scripts/setup_ci_environment.sh, +# this means other CI system doesn't actually have this flag properly set. +# Now we explicitly export IN_CI environment variable here. +export IN_CI=1 # compositional trap taken from https://stackoverflow.com/a/7287873/23845 diff --git a/test/distributed/elastic/multiprocessing/api_test.py b/test/distributed/elastic/multiprocessing/api_test.py index 3ae3a4b278398..13726ac4d7a52 100644 --- a/test/distributed/elastic/multiprocessing/api_test.py +++ b/test/distributed/elastic/multiprocessing/api_test.py @@ -35,7 +35,7 @@ NO_MULTIPROCESSING_SPAWN, TEST_WITH_ASAN, TEST_WITH_TSAN, - IS_PYTORCH_CI, + IS_IN_CI, IS_WINDOWS, IS_MACOS, ) @@ -593,7 +593,7 @@ def test_binary_redirect_and_tee(self): @unittest.skipIf( - TEST_WITH_ASAN or TEST_WITH_TSAN or IS_WINDOWS or IS_MACOS or IS_PYTORCH_CI, + TEST_WITH_ASAN or TEST_WITH_TSAN or IS_WINDOWS or IS_MACOS or IS_IN_CI, "tests incompatible with tsan or asan, the redirect functionality does not work on macos or windows", ) class StartProcessesNotCITest(StartProcessesTest): diff --git a/test/distributed/rpc/test_faulty_agent.py b/test/distributed/rpc/test_faulty_agent.py index 5c5c1b0654535..7c26643ab6b60 100644 --- a/test/distributed/rpc/test_faulty_agent.py +++ b/test/distributed/rpc/test_faulty_agent.py @@ -9,7 +9,7 @@ print("Distributed not available, skipping tests", file=sys.stderr) sys.exit(0) -from torch.testing._internal.common_utils import IS_PYTORCH_CI, run_tests +from torch.testing._internal.common_utils import IS_IN_CI, run_tests from torch.testing._internal.distributed.rpc.faulty_rpc_agent_test_fixture import ( FaultyRpcAgentTestFixture, ) @@ -22,7 +22,7 @@ # On CircleCI these tests are already run on CPU jobs, thus to save resources do # not run them on GPU jobs, since thet wouldn't provide additional test signal. -if not (IS_PYTORCH_CI and torch.cuda.is_available()): +if not (IS_IN_CI and torch.cuda.is_available()): globals().update( generate_tests( "Faulty", diff --git a/test/distributed/rpc/test_process_group_agent.py b/test/distributed/rpc/test_process_group_agent.py index 4e9c1c8341a32..98fa6fc7ab085 100644 --- a/test/distributed/rpc/test_process_group_agent.py +++ b/test/distributed/rpc/test_process_group_agent.py @@ -9,7 +9,7 @@ print("Distributed not available, skipping tests", file=sys.stderr) sys.exit(0) -from torch.testing._internal.common_utils import IS_PYTORCH_CI, run_tests +from torch.testing._internal.common_utils import IS_IN_CI, run_tests from torch.testing._internal.distributed.rpc.process_group_agent_test_fixture import ( ProcessGroupRpcAgentTestFixture, ) @@ -23,7 +23,7 @@ # On CircleCI these tests are already run on CPU jobs, thus to save resources do # not run them on GPU jobs, since thet wouldn't provide additional test signal. -if not (IS_PYTORCH_CI and torch.cuda.is_available()): +if not (IS_IN_CI and torch.cuda.is_available()): globals().update( generate_tests( "ProcessGroup", diff --git a/test/distributed/rpc/test_tensorpipe_agent.py b/test/distributed/rpc/test_tensorpipe_agent.py index 7d425933198c3..32b0e1c69357a 100644 --- a/test/distributed/rpc/test_tensorpipe_agent.py +++ b/test/distributed/rpc/test_tensorpipe_agent.py @@ -9,7 +9,7 @@ print("Distributed not available, skipping tests", file=sys.stderr) sys.exit(0) -from torch.testing._internal.common_utils import IS_PYTORCH_CI, run_tests +from torch.testing._internal.common_utils import IS_IN_CI, run_tests from torch.testing._internal.distributed.rpc.tensorpipe_rpc_agent_test_fixture import ( TensorPipeRpcAgentTestFixture, ) @@ -23,7 +23,7 @@ # On CircleCI these tests are already run on CPU jobs, thus to save resources do # not run them on GPU jobs, since thet wouldn't provide additional test signal. -if not (IS_PYTORCH_CI and torch.cuda.is_available()): +if not (IS_IN_CI and torch.cuda.is_available()): globals().update( generate_tests( "TensorPipe", diff --git a/test/test_dataloader.py b/test/test_dataloader.py index 3c80106ab008b..68998bb1bbf86 100644 --- a/test/test_dataloader.py +++ b/test/test_dataloader.py @@ -18,7 +18,7 @@ from torch.utils.data.dataset import random_split from torch._utils import ExceptionWrapper from torch.testing._internal.common_utils import (TestCase, run_tests, TEST_NUMPY, IS_WINDOWS, - IS_PYTORCH_CI, NO_MULTIPROCESSING_SPAWN, skipIfRocm, slowTest, + IS_IN_CI, NO_MULTIPROCESSING_SPAWN, skipIfRocm, slowTest, load_tests, TEST_WITH_TSAN, IS_SANDCASTLE) try: @@ -28,7 +28,7 @@ HAS_PSUTIL = False err_msg = ("psutil not found. Some critical data loader tests relying on it " "(e.g., TestDataLoader.test_proper_exit) will not run.") - if IS_PYTORCH_CI: + if IS_IN_CI: raise ImportError(err_msg) from None else: warnings.warn(err_msg) diff --git a/test/test_ops.py b/test/test_ops.py index 9b5f5dab6a9e4..8d5c4ec0d22b9 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -8,7 +8,7 @@ (FileCheck, floating_and_complex_types_and, get_all_dtypes) from torch.testing._internal.common_utils import \ (TestCase, is_iterable_of_tensors, run_tests, IS_SANDCASTLE, clone_input_helper, make_tensor, - gradcheck, gradgradcheck, IS_PYTORCH_CI) + gradcheck, gradgradcheck, IS_IN_CI) from torch.testing._internal.common_methods_invocations import \ (op_db,) from torch.testing._internal.common_device_type import \ @@ -35,7 +35,7 @@ class TestCommon(TestCase): def tearDownClass(cls): super().tearDownClass() - if IS_PYTORCH_CI: + if IS_IN_CI: err_msg = ("The operator(s) below is(are) using dynamic_dtypes in the OpInfo entries." "This is OK for testing, but be sure to set the dtypes manually before landing your PR!") # Assure no opinfo entry has dynamic_dtypes diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 47c36568f50d3..f00bccd168ca9 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -63,6 +63,7 @@ if sys.platform == 'win32': FILE_SCHEMA = "file:///" +# Environment variable `IN_CI` is set in `.jenkins/common.sh`. IS_IN_CI = os.getenv('IN_CI') == '1' IS_SANDCASTLE = os.getenv('SANDCASTLE') == '1' or os.getenv('TW_JOB_USER') == 'sandcastle' IS_FBCODE = os.getenv('PYTORCH_TEST_FBCODE') == '1' @@ -235,9 +236,6 @@ def call_helper(self, *args): return call_helper return repeat_helper -# Environment variable `IS_PYTORCH_CI` is set in `.jenkins/common.sh`. -IS_PYTORCH_CI = bool(os.environ.get('IS_PYTORCH_CI')) - def discover_test_cases_recursively(suite_or_case): if isinstance(suite_or_case, unittest.TestCase): @@ -832,8 +830,7 @@ def settings(*args, **kwargs): verbosity=hypothesis.Verbosity.verbose)) hypothesis.settings.load_profile( - "pytorch_ci" if IS_PYTORCH_CI else os.getenv('PYTORCH_HYPOTHESIS_PROFILE', - 'dev') + "pytorch_ci" if IS_IN_CI else os.getenv('PYTORCH_HYPOTHESIS_PROFILE', 'dev') ) except ImportError: print('Fail to import hypothesis in common_utils, tests are not derandomized') From 4a3eea9a6a6c5e71b45401915618873a7b35365f Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Sun, 20 Jun 2021 20:07:11 -0700 Subject: [PATCH 267/305] [quant][graphmode][fx] Produce reference linear module in convert (#60152) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60152 Test Plan: python test/test_quantization.py TestQuantizeFx Imported from OSS Reviewed By: vkuzo Differential Revision: D29188263 fbshipit-source-id: f7bbbef5d4d747eadf7a627a4e77a5ec9bb0bc94 --- test/quantization/fx/test_quantize_fx.py | 2 +- .../intrinsic/quantized/_reference/modules/linear_relu.py | 3 +++ torch/nn/quantized/modules/linear.py | 2 +- torch/quantization/fx/quantization_patterns.py | 4 +++- torch/quantization/quantization_mappings.py | 6 +++--- 5 files changed, 11 insertions(+), 6 deletions(-) diff --git a/test/quantization/fx/test_quantize_fx.py b/test/quantization/fx/test_quantize_fx.py index 8b49d1f4f503e..9555a86c4a608 100644 --- a/test/quantization/fx/test_quantize_fx.py +++ b/test/quantization/fx/test_quantize_fx.py @@ -500,7 +500,7 @@ def forward(self, x): LinearModule, (), (linear_module_input,), - ns.call_module(nnq.Linear), + ns.call_module(nnqr.Linear if is_reference else nnq.Linear), None, ), ] diff --git a/torch/nn/intrinsic/quantized/_reference/modules/linear_relu.py b/torch/nn/intrinsic/quantized/_reference/modules/linear_relu.py index f8dab5900d246..39c595376fded 100644 --- a/torch/nn/intrinsic/quantized/_reference/modules/linear_relu.py +++ b/torch/nn/intrinsic/quantized/_reference/modules/linear_relu.py @@ -1,8 +1,11 @@ import torch +import torch.nn.intrinsic as nni import torch.nn.quantized._reference as nnqr import torch.nn.functional as F class LinearReLU(nnqr.Linear): + _FLOAT_MODULE = nni.LinearReLU + def __init__( self, in_features, diff --git a/torch/nn/quantized/modules/linear.py b/torch/nn/quantized/modules/linear.py index e7f5c97d51ed6..4abd2115e4125 100644 --- a/torch/nn/quantized/modules/linear.py +++ b/torch/nn/quantized/modules/linear.py @@ -259,7 +259,7 @@ def from_float(cls, mod): if not isinstance(cls._FLOAT_MODULE, Iterable): cls._FLOAT_MODULE = [cls._FLOAT_MODULE] # type: ignore[assignment] supported_modules = ', '.join([float_mod.__name__ for float_mod in cls._FLOAT_MODULE]) # type: ignore[attr-defined] - error_msg = 'nnq.{}.from_float only works for {}'.format(cls.__name__, supported_modules) + error_msg = 'nnq.{}.from_float only works for {}, but got: {}'.format(cls.__name__, supported_modules, type(mod)) assert type(mod) in cls._FLOAT_MODULE, error_msg.format() # type: ignore[attr-defined] assert hasattr(mod, 'qconfig'), 'Input float module must have qconfig defined' activation_post_process = mod.activation_post_process diff --git a/torch/quantization/fx/quantization_patterns.py b/torch/quantization/fx/quantization_patterns.py index 699b2407124b1..86e282e829aa0 100644 --- a/torch/quantization/fx/quantization_patterns.py +++ b/torch/quantization/fx/quantization_patterns.py @@ -673,7 +673,9 @@ def convert(self, # 2. select corresponding quantized linear class for the float linear class if activation_int8_quantized: additional_static_quant_mapping = convert_custom_config_dict.get("static", {}) - qlinear = get_static_quant_module_class(type(self.linear), additional_static_quant_mapping) + qlinear = get_static_quant_module_class( + type(self.linear), additional_static_quant_mapping, + is_reference=is_reference) else: assert dtypes in [ (torch.float32, torch.qint8, torch.quint8), diff --git a/torch/quantization/quantization_mappings.py b/torch/quantization/quantization_mappings.py index a80e8ec2a4464..6179398b7398a 100644 --- a/torch/quantization/quantization_mappings.py +++ b/torch/quantization/quantization_mappings.py @@ -27,13 +27,13 @@ nn.Conv1d: nnqr.Conv1d, nn.Conv2d: nnqr.Conv2d, nn.Conv3d: nnqr.Conv3d, - # nn.Linear, nnqr.Linear, + nn.Linear: nnqr.Linear, nni.ConvReLU1d: nniqr.ConvReLU1d, nni.ConvReLU2d: nniqr.ConvReLU2d, nni.ConvReLU3d: nniqr.ConvReLU3d, nni.LinearReLU: nniqr.LinearReLU, # QAT Modules - # nnqat.Linear: nnqr.Linear, + nnqat.Linear: nnqr.Linear, nnqat.Conv2d: nnqr.Conv2d, nnqat.Conv3d: nnqr.Conv3d, nniqat.ConvBn1d: nnqr.Conv1d, @@ -44,7 +44,7 @@ nniqat.ConvBnReLU3d: nniqr.ConvReLU3d, nniqat.ConvReLU2d: nniqr.ConvReLU2d, nniqat.ConvReLU3d: nniqr.ConvReLU3d, - # nniqat.LinearReLU: nniqr.LinearReLU, + nniqat.LinearReLU: nniqr.LinearReLU, } # Default map for swapping float module to quantized ones From 1dee99c973fda55e1e9cac3d50b4d4982b6c6c26 Mon Sep 17 00:00:00 2001 From: Sameer Deshmukh Date: Sun, 20 Jun 2021 21:26:35 -0700 Subject: [PATCH 268/305] LU Solve using cublas and cusolver (#59148) Summary: This PR introduces cuSOLVER and cuBLAS for the `lu_solve` routine. Solves a part of https://github.com/pytorch/pytorch/issues/47953. Since usage of cuSOLVER with MAGMA introduces performance regressions in MAGMA (https://github.com/pytorch/pytorch/issues/56590), we use heuristics for determining when to call cuSOLVER, cuBLAS or MAGMA depending on the batch and matrix sizes. The 64-bit cuSOLVER API is not introduced in this PR since there are several problems with the LU factorization using cusolver (https://github.com/pytorch/pytorch/pull/59148). The following are performance benchmarks using various configurations:
``` [--------------------------------------------------------- LU solve CUDA torch.float64 ----------------------------------------------------------] | lu_solve CUSOLVER | lu_solve MAGMA | lu_solve CUBLAS | lu_solve cuSOLVER/MAGMA | lu_solve TEST ALL 1 threads: --------------------------------------------------------------------------------------------------------------------------------------- torch.Size([1, 1, 1]) | 703.4 | 489.8 | 511.8 | 710.1 | 487.1 torch.Size([2, 1, 1]) | 738.9 | 504.1 | 513.0 | 958.2 | 494.4 torch.Size([4, 1, 1]) | 790.7 | 514.7 | 506.8 | 983.9 | 540.2 torch.Size([8, 1, 1]) | 865.3 | 496.4 | 514.7 | 975.2 | 520.0 torch.Size([16, 1, 1]) | 955.5 | 483.9 | 508.3 | 937.6 | 526.5 torch.Size([32, 1, 1]) | 1167.7 | 495.2 | 511.2 | 934.0 | 528.7 torch.Size([64, 1, 1]) | 1730.0 | 492.1 | 537.8 | 936.4 | 533.2 torch.Size([128, 1, 1]) | 2748.4 | 499.7 | 526.5 | 982.9 | 540.8 torch.Size([1, 2, 2]) | 724.6 | 498.2 | 541.7 | 715.0 | 504.7 torch.Size([2, 2, 2]) | 737.0 | 514.3 | 527.6 | 934.5 | 524.5 torch.Size([4, 2, 2]) | 750.5 | 524.1 | 537.4 | 935.5 | 543.0 torch.Size([8, 2, 2]) | 844.8 | 513.7 | 538.9 | 953.3 | 534.4 torch.Size([16, 2, 2]) | 1013.1 | 521.9 | 530.0 | 932.2 | 537.9 torch.Size([32, 2, 2]) | 1335.8 | 515.1 | 544.4 | 939.9 | 559.5 torch.Size([64, 2, 2]) | 1819.6 | 511.8 | 534.1 | 973.9 | 540.0 torch.Size([128, 2, 2]) | 3018.7 | 526.3 | 546.1 | 979.3 | 543.5 torch.Size([1, 8, 8]) | 732.5 | 524.9 | 532.9 | 762.4 | 516.8 torch.Size([2, 8, 8]) | 771.2 | 514.9 | 538.7 | 1007.5 | 531.1 torch.Size([4, 8, 8]) | 811.3 | 507.7 | 534.6 | 1002.2 | 548.5 torch.Size([8, 8, 8]) | 866.6 | 530.0 | 532.0 | 1016.1 | 562.9 torch.Size([16, 8, 8]) | 991.8 | 533.6 | 548.0 | 1022.6 | 548.5 torch.Size([32, 8, 8]) | 1271.7 | 541.2 | 534.7 | 1013.8 | 545.6 torch.Size([64, 8, 8]) | 1817.2 | 530.2 | 520.6 | 1008.7 | 566.3 torch.Size([128, 8, 8]) | 2678.7 | 531.6 | 552.2 | 1006.2 | 555.0 torch.Size([1, 16, 16]) | 738.2 | 546.1 | 536.6 | 775.6 | 540.1 torch.Size([2, 16, 16]) | 782.6 | 543.5 | 539.6 | 1010.9 | 541.1 torch.Size([4, 16, 16]) | 815.2 | 546.1 | 560.9 | 1012.5 | 553.1 torch.Size([8, 16, 16]) | 877.7 | 543.0 | 547.9 | 1012.8 | 551.5 torch.Size([16, 16, 16]) | 1008.7 | 549.2 | 562.7 | 1016.6 | 546.8 torch.Size([32, 16, 16]) | 1291.9 | 540.8 | 560.3 | 1055.8 | 539.3 torch.Size([64, 16, 16]) | 1846.3 | 553.5 | 556.0 | 1010.8 | 551.9 torch.Size([128, 16, 16]) | 2953.8 | 562.7 | 547.5 | 1026.2 | 555.8 torch.Size([1, 32, 32]) | 789.1 | 590.6 | 590.9 | 790.5 | 579.0 torch.Size([2, 32, 32]) | 806.9 | 596.6 | 600.2 | 1085.6 | 573.8 torch.Size([4, 32, 32]) | 852.0 | 597.9 | 588.2 | 1098.9 | 574.7 torch.Size([8, 32, 32]) | 914.2 | 597.8 | 591.4 | 1090.3 | 585.7 torch.Size([16, 32, 32]) | 1063.0 | 604.6 | 597.3 | 1094.0 | 580.5 torch.Size([32, 32, 32]) | 1302.0 | 602.0 | 598.9 | 1090.3 | 583.6 torch.Size([64, 32, 32]) | 1861.7 | 601.1 | 599.8 | 1113.4 | 588.6 torch.Size([128, 32, 32]) | 3251.0 | 619.6 | 595.3 | 1106.8 | 608.9 torch.Size([1, 64, 64]) | 978.6 | 842.7 | 778.6 | 1071.4 | 825.8 torch.Size([2, 64, 64]) | 1072.3 | 845.7 | 785.4 | 1400.6 | 829.0 torch.Size([4, 64, 64]) | 1051.9 | 842.9 | 796.1 | 1352.2 | 788.2 torch.Size([8, 64, 64]) | 1090.3 | 834.1 | 805.2 | 1382.6 | 804.7 torch.Size([16, 64, 64]) | 1206.9 | 835.7 | 802.2 | 1365.6 | 801.2 torch.Size([32, 64, 64]) | 1671.2 | 846.5 | 794.5 | 1345.1 | 814.2 torch.Size([64, 64, 64]) | 2759.3 | 848.5 | 795.4 | 1409.7 | 832.9 torch.Size([128, 64, 64]) | 4928.6 | 877.4 | 848.3 | 1439.0 | 883.9 torch.Size([1, 128, 128]) | 1315.6 | 1158.4 | 1130.0 | 1301.3 | 1177.1 torch.Size([2, 128, 128]) | 1334.7 | 1198.2 | 1186.6 | 1703.9 | 1209.5 torch.Size([4, 128, 128]) | 1374.6 | 1200.7 | 1266.2 | 1640.6 | 1272.3 torch.Size([8, 128, 128]) | 1453.6 | 1215.9 | 1287.3 | 1669.1 | 1288.7 torch.Size([16, 128, 128]) | 1882.1 | 1244.9 | 1337.6 | 1698.8 | 1347.1 torch.Size([32, 128, 128]) | 2789.0 | 1284.5 | 1398.6 | 1747.6 | 1396.3 torch.Size([64, 128, 128]) | 4763.0 | 1425.2 | 1581.7 | 1921.0 | 1584.1 torch.Size([128, 128, 128]) | 8835.9 | 1808.9 | 1968.7 | 2197.6 | 1961.8 torch.Size([1, 512, 512]) | 4369.9 | 4577.6 | 4804.0 | 4331.4 | 4599.0 torch.Size([2, 512, 512]) | 4635.9 | 4850.1 | 5159.1 | 5315.4 | 4845.5 torch.Size([4, 512, 512]) | 5367.5 | 5261.6 | 6134.7 | 5807.8 | 5345.2 torch.Size([8, 512, 512]) | 7025.2 | 6184.5 | 7065.6 | 6711.6 | 6303.9 torch.Size([16, 512, 512]) | 10221.3 | 7849.7 | 8820.1 | 8323.6 | 7992.1 torch.Size([32, 512, 512]) | 16574.8 | 11208.4 | 12284.3 | 11704.7 | 11394.4 torch.Size([64, 512, 512]) | 29500.1 | 18043.1 | 19249.3 | 18744.0 | 18242.1 torch.Size([128, 512, 512]) | 56783.3 | 33903.9 | 34713.5 | 33893.8 | 34041.8 torch.Size([1, 1024, 1024]) | 14864.5 | 15714.6 | 16128.1 | 14726.7 | 14992.6 torch.Size([2, 1024, 1024]) | 17891.0 | 18553.3 | 19111.6 | 19271.5 | 19283.0 torch.Size([4, 1024, 1024]) | 22143.4 | 21909.2 | 23667.1 | 22698.9 | 22713.8 torch.Size([8, 1024, 1024]) | 30621.1 | 28669.9 | 30822.9 | 29725.0 | 29760.8 torch.Size([16, 1024, 1024]) | 47045.9 | 41900.0 | 44353.8 | 43215.6 | 43237.5 torch.Size([32, 1024, 1024]) | 79245.5 | 68316.9 | 70959.0 | 69506.4 | 69876.7 torch.Size([64, 1024, 1024]) | 147973.9 | 121120.6 | 124601.1 | 122084.4 | 122578.7 torch.Size([128, 1024, 1024]) | 295586.2 | 232871.8 | 237421.8 | 233765.3 | 234704.6 Times are in microseconds (us). ```
Here's the details of how the tests were performed: * CUSOLVER - Only call `cusolver` for all problem sizes. * MAGMA - Only call `magma` for all problem sizes (this is the current master branch). * CUBLAS - Only call `cublas` for all problem sizes. * cuSOLVER / MAGMA - Use cusolver for `batch_size == 1` and magma for all others. * TEST ALL - Employ heuristics to switch between cublas/cusolver/magma. This yields the best overall results (this PR). Script for reproducing the results:
``` python import torch import pickle import itertools from torch.utils.benchmark import Timer import sys shapes = [1, 2, 8, 16, 32, 64, 128, 512, 1024] batches = [(1,), (2,), (4,), (8,), (16,), (32,), (64,), (128,)] results = [] num_threads = 1 dtype = torch.float64 repeats = 2 from torch.testing._internal.common_utils import random_hermitian_pd_matrix def lu_factorize_solve(mat, b): lu_data = torch.lu(mat) x = torch.lu_solve(b, *lu_data) for shape, batch in itertools.product(shapes, batches): mat = torch.randn(*batch, shape, shape, dtype=dtype, device='cuda') b = torch.randn(*batch, shape, 1, dtype=dtype, device='cuda') tasks = [("lu_factorize_solve(mat, b)", "lu_solve CUSOLVER")] print("shape: ", shape, " batch: ", batch) timers = [Timer(stmt=stmt, num_threads=num_threads, label=f"LU solve CUDA {dtype}", sub_label=f"{mat.shape}", description=label, globals=globals()) for stmt, label in tasks] for i, timer in enumerate(timers * repeats): results.append( pickle.dumps(timer.blocked_autorange()) ) print(f"\r{i + 1} / {len(timers) * repeats}", end="") sys.stdout.flush() f = open("cusolver_lu_solve.pickle", "wb") pickle.dump(results, f) f.close() ```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59148 Reviewed By: H-Huang Differential Revision: D29160609 Pulled By: mruberry fbshipit-source-id: 7280f25db1e66aa650ea15608a6dc5d688fb4db2 --- aten/src/ATen/cuda/CUDABlas.cpp | 66 ++++++++++++++++++- aten/src/ATen/cuda/CUDABlas.h | 22 ++++++- .../ATen/native/cuda/BatchLinearAlgebra.cu | 43 ++++++++---- .../ATen/native/cuda/BatchLinearAlgebraLib.cu | 65 ++++++++++++++++++ .../ATen/native/cuda/BatchLinearAlgebraLib.h | 2 + 5 files changed, 184 insertions(+), 14 deletions(-) diff --git a/aten/src/ATen/cuda/CUDABlas.cpp b/aten/src/ATen/cuda/CUDABlas.cpp index 67595def245a3..a876db0003fff 100644 --- a/aten/src/ATen/cuda/CUDABlas.cpp +++ b/aten/src/ATen/cuda/CUDABlas.cpp @@ -908,9 +908,73 @@ void vdot>(CUDABLAS_DOT_ARGTYPES(c10::complex)) { reinterpret_cast(result))); } -// This guards blocks use of geqrfBatched, getrfBatched, getriBatched on platforms other than cuda +// This guards blocks use of getrsBatched, geqrfBatched, getrfBatched, getriBatched on platforms other than cuda #ifdef CUDART_VERSION +template <> +void getrsBatched(CUDABLAS_GETRS_ARGTYPES(float)) { + TORCH_CUDABLAS_CHECK(cublasSgetrsBatched( + handle, + trans, + n, + nrhs, + dA_array, + lda, + ipiv_array, + dB_array, + ldb, + info_array, + batchsize)); +} + +template <> +void getrsBatched(CUDABLAS_GETRS_ARGTYPES(double)) { + TORCH_CUDABLAS_CHECK(cublasDgetrsBatched( + handle, + trans, + n, + nrhs, + dA_array, + lda, + ipiv_array, + dB_array, + ldb, + info_array, + batchsize)); +} + +template <> +void getrsBatched>(CUDABLAS_GETRS_ARGTYPES(c10::complex)) { + TORCH_CUDABLAS_CHECK(cublasCgetrsBatched( + handle, + trans, + n, + nrhs, + reinterpret_cast(dA_array), + lda, + ipiv_array, + reinterpret_cast(dB_array), + ldb, + info_array, + batchsize)); +} + +template <> +void getrsBatched>(CUDABLAS_GETRS_ARGTYPES(c10::complex)) { + TORCH_CUDABLAS_CHECK(cublasZgetrsBatched( + handle, + trans, + n, + nrhs, + reinterpret_cast(dA_array), + lda, + ipiv_array, + reinterpret_cast(dB_array), + ldb, + info_array, + batchsize)); +} + template <> void geqrfBatched(CUDABLAS_GEQRF_BATCHED_ARGTYPES(float)) { TORCH_CUDABLAS_CHECK(cublasSgeqrfBatched( diff --git a/aten/src/ATen/cuda/CUDABlas.h b/aten/src/ATen/cuda/CUDABlas.h index f579c4f7e8bbe..b594362c5e41c 100644 --- a/aten/src/ATen/cuda/CUDABlas.h +++ b/aten/src/ATen/cuda/CUDABlas.h @@ -199,9 +199,28 @@ void vdot>(CUDABLAS_DOT_ARGTYPES(c10::complex)); template <> void vdot>(CUDABLAS_DOT_ARGTYPES(c10::complex)); -// This guards blocks use of geqrfBatched, getrfBatched, getriBatched on platforms other than cuda +// This guards blocks use of getrsBatched, geqrfBatched, getrfBatched, getriBatched on platforms other than cuda #ifdef CUDART_VERSION +#define CUDABLAS_GETRS_ARGTYPES(Dtype) \ + cublasHandle_t handle, cublasOperation_t trans, \ + int n, int nrhs, Dtype** dA_array, int lda, int* ipiv_array, \ + Dtype** dB_array, int ldb, int* info_array, int batchsize + +template +void getrsBatched(CUDABLAS_GETRS_ARGTYPES(Dtype)) { + TORCH_INTERNAL_ASSERT(false, "at::cuda::blas::getrsBatched: not implemented for ", + typeid(Dtype).name()); +} +template<> +void getrsBatched(CUDABLAS_GETRS_ARGTYPES(float)); +template<> +void getrsBatched(CUDABLAS_GETRS_ARGTYPES(double)); +template<> +void getrsBatched>(CUDABLAS_GETRS_ARGTYPES(c10::complex)); +template<> +void getrsBatched>(CUDABLAS_GETRS_ARGTYPES(c10::complex)); + #define CUDABLAS_GEQRF_BATCHED_ARGTYPES(Dtype) \ cublasHandle_t handle, int m, int n, Dtype **A_array, int lda, \ Dtype **tau_array, int *info, int batchsize @@ -240,7 +259,6 @@ void getrfBatched>(CUDABLAS_GETRF_ARGTYPES(c10::complex void getrfBatched>(CUDABLAS_GETRF_ARGTYPES(c10::complex)); - #define CUDABLAS_GETRI_ARGTYPES(Dtype) \ int n, Dtype** dA_array, int ldda, int* ipiv_array, Dtype** dC_array, int lddc, int* info_array, int batchsize diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu index 9b22fbe1ee94b..2cb743380d5aa 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebra.cu @@ -2840,20 +2840,41 @@ static void apply_lu_solve_batched_magma(const Tensor& b, const Tensor& lu, cons #endif } -static void lu_solve_magma(const Tensor& b, const Tensor& lu, const Tensor& pivots) { - // TODO: compare performance and use the best performing option based on lu's sizes - if (b.dim() == 2) { - AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(b.scalar_type(), "lu_solve_magma", [&]{ - apply_lu_solve_looped_magma(b, lu, pivots); - }); - } else { - AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(b.scalar_type(), "lu_solve_magma", [&]{ - apply_lu_solve_batched_magma(b, lu, pivots); - }); +static void lu_solve_batched_magma(const Tensor& b, const Tensor& lu, const Tensor& pivots) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(b.scalar_type(), "lu_solve_batched_magma", [&]{ + apply_lu_solve_batched_magma(b, lu, pivots); + }); +} + +static void lu_solve_looped_magma(const Tensor& b, const Tensor& lu, const Tensor& pivots) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(b.scalar_type(), "lu_solve_looped_magma", [&]{ + apply_lu_solve_looped_magma(b, lu, pivots); + }); +} + +static void lu_solve_dispatch(const Tensor& b, const Tensor& lu, const Tensor& pivots) { + auto batch_size = batchCount(lu); + auto m = lu.size(-2); +#ifdef USE_CUSOLVER + if (batch_size == 1 && m > 512) { + lu_solve_looped_cusolver(b, lu, pivots); + } +#else + if (batch_size == 1) { + lu_solve_looped_magma(b, lu, pivots); + } +#endif // ifdef USE_CUSOLVER +#ifdef CUDART_VERSION + else if (batch_size > 2 && m <= 128) { + lu_solve_batched_cublas(b, lu, pivots); + } +#endif // ifdef CUDART_VERSION + else { + lu_solve_batched_magma(b, lu, pivots); } } -REGISTER_DISPATCH(lu_solve_stub, &lu_solve_magma); +REGISTER_DISPATCH(lu_solve_stub, &lu_solve_dispatch); // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ lstsq ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu index ae1d0281bf595..0b10b348332ea 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cu @@ -69,6 +69,38 @@ void geqrf_batched_cublas(const Tensor& input, const Tensor& tau) { }); } +template +static void apply_lu_solve_batched_cublas(const Tensor& b, const Tensor& lu, const Tensor& pivots) { +#ifndef CUDART_VERSION + TORCH_CHECK(false, "lu_solve: cuBLAS backend for lu_solve is not available.") +#else + cublasOperation_t trans = CUBLAS_OP_N; + + auto pivots_data = pivots.data_ptr(); + auto batch_size = cuda_int_cast(batchCount(lu), "batch_size");; + auto m = cuda_int_cast(lu.size(-2), "m"); + auto nrhs = cuda_int_cast(b.size(-1), "nrhs"); + auto lda = cuda_int_cast(std::max(1, m), "lda"); + int info = 0; + + Tensor lu_ptr_array = get_device_pointers(lu); + Tensor b_ptr_array = get_device_pointers(b); + auto lu_ptr_array_data = reinterpret_cast(lu_ptr_array.data_ptr()); + auto b_ptr_array_data = reinterpret_cast(b_ptr_array.data_ptr()); + + auto handle = at::cuda::getCurrentCUDABlasHandle(); + at::cuda::blas::getrsBatched(handle, trans, m, nrhs, lu_ptr_array_data, + lda, pivots_data, b_ptr_array_data, lda, &info, batch_size); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(info == 0); +#endif +} + +void lu_solve_batched_cublas(const Tensor& b, const Tensor& lu, const Tensor& pivots) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(lu.scalar_type(), "lu_solve_cublas", [&]{ + apply_lu_solve_batched_cublas(b, lu, pivots); + }); +} + template static void apply_triangular_solve(Tensor& A, Tensor& B, bool upper, bool transpose, bool conjugate_transpose, bool unitriangular) { cublasFillMode_t uplo = upper ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER; @@ -1223,6 +1255,39 @@ void linalg_eigh_cusolver(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& inf } } +void lu_solve_looped_cusolver(const Tensor& b, const Tensor& lu, const Tensor& pivots) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(b.scalar_type(), "lu_solve_cusolver", [&] { + int n = cuda_int_cast(lu.size(-2), "n"); + int nrhs = cuda_int_cast(b.size(-1), "nrhs"); + auto batch_size = batchCount(lu); + auto info = at::zeros({1}, lu.options().dtype(kInt)); + auto info_data = info.data_ptr(); + auto b_data = b.data_ptr(); + auto lu_data = lu.data_ptr(); + auto pivots_data = pivots.data_ptr(); + auto pivots_stride = pivots.size(-1); + auto lu_stride = matrixStride(lu); + auto b_stride = matrixStride(b); + int leading_dimension = cuda_int_cast(std::max(1, n), "leading_dimension"); + + auto handle = at::cuda::getCurrentCUDASolverDnHandle(); + for (auto batch = decltype(batch_size){0}; batch < batch_size; ++batch) { + at::cuda::solver::getrs( + handle, + n, + nrhs, + lu_data + batch * lu_stride, + leading_dimension, + pivots_data + batch * pivots_stride, + b_data + batch * b_stride, + leading_dimension, + info_data); + + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(info.item().toInt() == 0); + } + }); +} + #endif // USE_CUSOLVER }} // namespace at::native diff --git a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h index a2095641e637c..87377cae46a43 100644 --- a/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h +++ b/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.h @@ -29,6 +29,7 @@ void geqrf_batched_cublas(const Tensor& input, const Tensor& tau); void triangular_solve_cublas(Tensor& A, Tensor& B, Tensor& infos, bool upper, bool transpose, bool conjugate_transpose, bool unitriangular); void triangular_solve_batched_cublas(Tensor& A, Tensor& B, Tensor& infos, bool upper, bool transpose, bool conjugate_transpose, bool unitriangular); void gels_batched_cublas(const Tensor& a, Tensor& b, Tensor& infos); +void lu_solve_batched_cublas(const Tensor& b, const Tensor& lu, const Tensor& pivots); #ifdef USE_CUSOLVER @@ -49,6 +50,7 @@ void ormqr_cusolver(const Tensor& input, const Tensor& tau, const Tensor& other, Tensor& orgqr_helper_cusolver(Tensor& result, const Tensor& tau); void linalg_eigh_cusolver(Tensor& eigenvalues, Tensor& eigenvectors, Tensor& infos, bool upper, bool compute_eigenvectors); +void lu_solve_looped_cusolver(const Tensor& b, const Tensor& lu, const Tensor& pivots); #endif // USE_CUSOLVER From 5ff407df675b3f178080f56abc86442da428370f Mon Sep 17 00:00:00 2001 From: Mike Ruberry Date: Mon, 21 Jun 2021 01:34:37 -0700 Subject: [PATCH 269/305] Skips failing MacOS tests (#60348) Summary: Mitigates, but does not fix https://github.com/pytorch/pytorch/issues/60347. Pull Request resolved: https://github.com/pytorch/pytorch/pull/60348 Reviewed By: ngimel Differential Revision: D29257917 Pulled By: mruberry fbshipit-source-id: de9be93ddeda1ca27ea2ff4650162f886d10f1e2 --- test/distributions/test_distributions.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/distributions/test_distributions.py b/test/distributions/test_distributions.py index 350738439dc9e..87ba61348f4c3 100644 --- a/test/distributions/test_distributions.py +++ b/test/distributions/test_distributions.py @@ -36,8 +36,9 @@ torch.set_default_dtype(torch.double) from torch._six import inf -from torch.testing._internal.common_utils import TestCase, run_tests, set_rng_seed, TEST_WITH_UBSAN, load_tests, \ - gradcheck +from torch.testing._internal.common_utils import \ + (TestCase, run_tests, set_rng_seed, TEST_WITH_UBSAN, load_tests, + gradcheck, IS_MACOS) from torch.testing._internal.common_cuda import TEST_CUDA from torch.autograd import grad from torch.autograd.functional import jacobian @@ -1027,6 +1028,7 @@ def ref_log_prob(idx, val, log_prob): self.assertEqual(Geometric(p).entropy(), scipy.stats.geom(p.detach().numpy(), loc=-1).entropy(), atol=1e-3, rtol=0) self.assertEqual(float(Geometric(s).entropy()), scipy.stats.geom(s, loc=-1).entropy().item(), atol=1e-3, rtol=0) + @unittest.skipIf(IS_MACOS, "See https://github.com/pytorch/pytorch/issues/60347") @unittest.skipIf(not TEST_NUMPY, "NumPy not found") def test_geometric_sample(self): set_rng_seed(0) # see Note [Randomized statistical tests] @@ -1043,6 +1045,7 @@ def test_binomial(self): self.assertRaises(NotImplementedError, Binomial(10, p).rsample) self.assertRaises(NotImplementedError, Binomial(10, p).entropy) + @unittest.skipIf(IS_MACOS, "See https://github.com/pytorch/pytorch/issues/60347") @unittest.skipIf(not TEST_NUMPY, "NumPy not found") def test_binomial_sample(self): set_rng_seed(0) # see Note [Randomized statistical tests] @@ -1323,6 +1326,7 @@ def ref_log_prob(idx, x, log_prob): self._gradcheck_log_prob(Poisson, (rate,)) self._gradcheck_log_prob(Poisson, (rate_1d,)) + @unittest.skipIf(IS_MACOS, "See https://github.com/pytorch/pytorch/issues/60347") @unittest.skipIf(not TEST_NUMPY, "Numpy not found") def test_poisson_sample(self): set_rng_seed(1) # see Note [Randomized statistical tests] From f2bb0932da3bc61a75867e4fec241ae8ca65b63e Mon Sep 17 00:00:00 2001 From: CodemodService FBSourceClangFormatLinterBot <> Date: Mon, 21 Jun 2021 03:55:42 -0700 Subject: [PATCH 270/305] [AutoAccept][Codemod][FBSourceClangFormatLinter] Daily `arc lint --take CLANGFORMAT` Reviewed By: zertosh Differential Revision: D29259226 fbshipit-source-id: 15fd79f6fed38d6ed2d84018852806683d5a09fa --- torch/csrc/jit/passes/tensorexpr_fuser.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.h b/torch/csrc/jit/passes/tensorexpr_fuser.h index 55239b298c828..3f6538b7e587a 100644 --- a/torch/csrc/jit/passes/tensorexpr_fuser.h +++ b/torch/csrc/jit/passes/tensorexpr_fuser.h @@ -62,6 +62,6 @@ TORCH_API Value* broadcastSizes(at::ArrayRef sizes, AliasDb* db); namespace tensorexpr { TORCH_API const OperatorSet& supported_eltwise_set(); TORCH_API bool isSupported(Node* node); -} +} // namespace tensorexpr } // namespace jit } // namespace torch From 61e0bc19554179f83d77a803429ba5f12cd58605 Mon Sep 17 00:00:00 2001 From: Raghavan Raman Date: Mon, 21 Jun 2021 09:54:20 -0700 Subject: [PATCH 271/305] [nnc] Remove check on initializer in compressBuffer (#60194) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60194 Test Plan: Imported from OSS Reviewed By: bertmaher, huiguoo Differential Revision: D29206255 Pulled By: navahgar fbshipit-source-id: 0a68ec4067c37f06ca1ea9ddeeb5ad5e0dcb0639 --- torch/csrc/jit/tensorexpr/loopnest.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index 985dee1aa9148..9cc52ccf6ce56 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -2093,10 +2093,6 @@ bool LoopNest::flatten(const std::vector& loops) { } void LoopNest::compressBuffer(Buf* buf, Stmt* stmt) { - if (buf->initializer()) { - throw malformed_input("Can't compress buffer whose initializer is set"); - } - // Loop iterations in NNC IR do not follow sequential semantics by default. // In other words, the iterations of the loops could be executed in any // random order without affecting correctness. This constraint in turn From f89ae9cb8dbd442ee85f39a0bb9572f9be823c21 Mon Sep 17 00:00:00 2001 From: Michael Carilli Date: Mon, 21 Jun 2021 10:19:30 -0700 Subject: [PATCH 272/305] Moves grid_sampler to autocast promote list (#58618) Summary: Should close https://github.com/pytorch/pytorch/issues/42218 Numerically, `grid_sampler` is fine in fp16 or fp32, but takes several inputs and expects their dtypes to match, so it belongs on the autocast promote list. `grid_sampler` currently uses `gpuAtomicAdd`, notoriously slow in fp16 because it calls cuda's atomicAdd __half overload which uses a software compare-and-swap loop internally. To allow good performance if both inputs happen to be FP16, the PR also modifies `grid_sampler_[2,3]d_backward_kernel`s to use `fastAtomicAdd` instead. Pull Request resolved: https://github.com/pytorch/pytorch/pull/58618 Reviewed By: mruberry Differential Revision: D29257199 Pulled By: ngimel fbshipit-source-id: 3cc7505945b480427f2fc1beb36bee80bf3853b3 --- aten/src/ATen/autocast_mode.cpp | 2 +- aten/src/ATen/native/cuda/GridSampler.cu | 101 +++++++++++------- aten/src/ATen/native/cuda/GridSampler.cuh | 34 ++++-- aten/src/ATen/native/cuda/IndexKernel.cu | 2 +- aten/src/ATen/native/cuda/KernelUtils.cuh | 16 +-- docs/source/amp.rst | 2 +- .../testing/_internal/autocast_test_lists.py | 6 +- 7 files changed, 102 insertions(+), 61 deletions(-) diff --git a/aten/src/ATen/autocast_mode.cpp b/aten/src/ATen/autocast_mode.cpp index b2d3617da51bc..be8ff9dfc8b3d 100644 --- a/aten/src/ATen/autocast_mode.cpp +++ b/aten/src/ATen/autocast_mode.cpp @@ -401,7 +401,6 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) { KERNEL(ADD_NS(pdist), "pdist", Tensor (const Tensor &, double), fp32) KERNEL(ADD_NS(cdist), "cdist", Tensor (const Tensor &, const Tensor &, double, c10::optional), fp32) KERNEL(ADD_NS(renorm), "renorm", Tensor (const Tensor &, const Scalar&, int64_t, const Scalar&), fp32) - KERNEL(ADD_NS(grid_sampler), "grid_sampler", Tensor (const Tensor &, const Tensor &, int64_t, int64_t, bool), fp32) // fp32_set_opt_dtype KERNEL(ADD_NS(prod), "prod", Tensor (const Tensor &, c10::optional), fp32_set_opt_dtype) KERNEL(ADD_NS(prod), "prod.dim_int", Tensor (const Tensor &, int64_t, bool, c10::optional), fp32_set_opt_dtype) @@ -435,6 +434,7 @@ TORCH_LIBRARY_IMPL(aten, Autocast, m) { KERNEL(ADD_NS(bilinear), "bilinear", Tensor (const Tensor &, const Tensor &, const Tensor &, const c10::optional&), promote) KERNEL(ADD_NS(cross), "cross", Tensor (const Tensor &, const Tensor &, c10::optional), promote) KERNEL(ADD_NS(dot), "dot", Tensor (const Tensor &, const Tensor &), promote) + KERNEL(ADD_NS(grid_sampler), "grid_sampler", Tensor (const Tensor &, const Tensor &, int64_t, int64_t, bool), promote) KERNEL(ADD_NS(index_put), "index_put", Tensor (const Tensor &, const torch::List>&, const Tensor &, bool), promote) KERNEL(ADD_NS(tensordot), "tensordot", Tensor (const Tensor &, const Tensor &, IntArrayRef, IntArrayRef), promote) KERNEL(ADD_NS(scatter_add), "scatter_add", Tensor (const Tensor&, int64_t, const Tensor&, const Tensor&), promote) diff --git a/aten/src/ATen/native/cuda/GridSampler.cu b/aten/src/ATen/native/cuda/GridSampler.cu index a08c13037e348..be8c15d0d8ada 100644 --- a/aten/src/ATen/native/cuda/GridSampler.cu +++ b/aten/src/ATen/native/cuda/GridSampler.cu @@ -290,6 +290,12 @@ namespace { } } +// Note [Passing pointer and offset to fastAtomicAdd] +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +// For its internal bounds checking, fastAtomicAdd needs to know where the destination address +// lies relative to the entire tensor, so we pass the base grad_input.data and full offset information, +// including batch * channel offset (NC_offset). + template C10_LAUNCH_BOUNDS_1(1024) __global__ void grid_sampler_2d_backward_kernel( @@ -301,7 +307,8 @@ namespace { TensorInfo grad_grid, // initialized to empty const GridSamplerInterpolation interpolation_mode, const GridSamplerPadding padding_mode, - bool align_corners) { + bool align_corners, + const index_t grad_input_memory_span) { index_t C = input.sizes[1]; index_t inp_H = input.sizes[2]; @@ -360,16 +367,16 @@ namespace { scalar_t gix = static_cast(0), giy = static_cast(0); scalar_t *gOut_ptr_NCHW = grad_output.data + n * gOut_sN + h * gOut_sH + w * gOut_sW; - scalar_t *gInp_ptr_NC = grad_input.data + n * gInp_sN; + index_t NC_offset = n * gInp_sN; scalar_t *inp_ptr_NC = input.data + n * inp_sN; - for (index_t c = 0; c < C; ++c, inp_ptr_NC += inp_sC, gInp_ptr_NC += gInp_sC, gOut_ptr_NCHW += gOut_sC) { + for (index_t c = 0; c < C; ++c, inp_ptr_NC += inp_sC, NC_offset += gInp_sC, gOut_ptr_NCHW += gOut_sC) { scalar_t gOut = *gOut_ptr_NCHW; - // calculate and set grad_input - safe_add_2d(gInp_ptr_NC, iy_nw, ix_nw, gInp_sH, gInp_sW, inp_H, inp_W, nw * gOut); - safe_add_2d(gInp_ptr_NC, iy_ne, ix_ne, gInp_sH, gInp_sW, inp_H, inp_W, ne * gOut); - safe_add_2d(gInp_ptr_NC, iy_sw, ix_sw, gInp_sH, gInp_sW, inp_H, inp_W, sw * gOut); - safe_add_2d(gInp_ptr_NC, iy_se, ix_se, gInp_sH, gInp_sW, inp_H, inp_W, se * gOut); + // calculate and set grad_input. See Note [Passing pointer and offset to fastAtomicAdd]. + safe_add_2d(grad_input.data, iy_nw, ix_nw, gInp_sH, gInp_sW, inp_H, inp_W, nw * gOut, NC_offset, grad_input_memory_span); + safe_add_2d(grad_input.data, iy_ne, ix_ne, gInp_sH, gInp_sW, inp_H, inp_W, ne * gOut, NC_offset, grad_input_memory_span); + safe_add_2d(grad_input.data, iy_sw, ix_sw, gInp_sH, gInp_sW, inp_H, inp_W, sw * gOut, NC_offset, grad_input_memory_span); + safe_add_2d(grad_input.data, iy_se, ix_se, gInp_sH, gInp_sW, inp_H, inp_W, se * gOut, NC_offset, grad_input_memory_span); // calculate grad_grid if (within_bounds_2d(iy_nw, ix_nw, inp_H, inp_W)) { @@ -407,10 +414,10 @@ namespace { // assign nearest neighor pixel value to output pixel scalar_t *gOut_ptr_NCHW = grad_output.data + n * gOut_sN + h * gOut_sH + w * gOut_sW; - scalar_t *gInp_ptr_NC = grad_input.data + n * gInp_sN; - for (index_t c = 0; c < C; ++c, gInp_ptr_NC += gInp_sC, gOut_ptr_NCHW += gOut_sC) { - // calculate and set grad_input - safe_add_2d(gInp_ptr_NC, iy_nearest, ix_nearest, gInp_sH, gInp_sW, inp_H, inp_W, *gOut_ptr_NCHW); + index_t NC_offset = n * gInp_sN; + for (index_t c = 0; c < C; ++c, NC_offset += gInp_sC, gOut_ptr_NCHW += gOut_sC) { + // calculate and set grad_input. See Note [Passing pointer and offset to fastAtomicAdd]. + safe_add_2d(grad_input.data, iy_nearest, ix_nearest, gInp_sH, gInp_sW, inp_H, inp_W, *gOut_ptr_NCHW, NC_offset, grad_input_memory_span); } // assuming grad_grid is contiguous @@ -445,18 +452,22 @@ namespace { scalar_t giy = static_cast(0); scalar_t *gOut_ptr_NCHW = grad_output.data + n * gOut_sN + h * gOut_sH + w * gOut_sW; - scalar_t *gInp_ptr_NC = grad_input.data + n * gInp_sN; + index_t NC_offset = n * gInp_sN; scalar_t *inp_ptr_NC = input.data + n * inp_sN; - for (index_t c = 0; c < C; ++c, gOut_ptr_NCHW += gOut_sC, gInp_ptr_NC += gInp_sC, inp_ptr_NC+= inp_sC) { + for (index_t c = 0; c < C; ++c, gOut_ptr_NCHW += gOut_sC, NC_offset += gInp_sC, inp_ptr_NC+= inp_sC) { scalar_t gOut = *gOut_ptr_NCHW; for (index_t i = 0; i < 4; ++i) { for (index_t j = 0; j < 4; ++j) { - // set input gradient - add_value_bounded(gInp_ptr_NC, ix_nw - 1 + i, iy_nw - 1 + j, inp_W, inp_H, - gInp_sW, gInp_sH, gOut * x_coeffs[i] * y_coeffs[j], padding_mode, align_corners); + // set input gradient. See Note [Passing pointer and offset to fastAtomicAdd]. + add_value_bounded(grad_input.data, ix_nw - 1 + i, iy_nw - 1 + j, inp_W, inp_H, gInp_sW, gInp_sH, + gOut * x_coeffs[i] * y_coeffs[j], + padding_mode, + align_corners, + NC_offset, + grad_input_memory_span); // set grid gradient scalar_t val = get_value_bounded(inp_ptr_NC, ix_nw - 1 + i, iy_nw - 1 + j, @@ -486,7 +497,8 @@ namespace { TensorInfo grad_grid, // initialized to empty const GridSamplerInterpolation interpolation_mode, const GridSamplerPadding padding_mode, - bool align_corners) { + bool align_corners, + const index_t grad_input_memory_span) { index_t C = input.sizes[1]; index_t inp_D = input.sizes[2]; @@ -583,21 +595,29 @@ namespace { scalar_t gix = static_cast(0), giy = static_cast(0), giz = static_cast(0); scalar_t *gOut_ptr_NCDHW = grad_output.data + n * gOut_sN + d * gOut_sD + h * gOut_sH + w * gOut_sW; - scalar_t *gInp_ptr_NC = grad_input.data + n * gInp_sN; + index_t NC_offset = n * gInp_sN; scalar_t *inp_ptr_NC = input.data + n * inp_sN; // calculate bilinear weighted pixel value and set output pixel - for (index_t c = 0; c < C; ++c, gOut_ptr_NCDHW += gOut_sC, gInp_ptr_NC += gInp_sC, inp_ptr_NC += inp_sC) { + for (index_t c = 0; c < C; ++c, gOut_ptr_NCDHW += gOut_sC, NC_offset += gInp_sC, inp_ptr_NC += inp_sC) { scalar_t gOut = *gOut_ptr_NCDHW; - // calculate and set grad_input - safe_add_3d(gInp_ptr_NC, iz_tnw, iy_tnw, ix_tnw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tnw * gOut); - safe_add_3d(gInp_ptr_NC, iz_tne, iy_tne, ix_tne, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tne * gOut); - safe_add_3d(gInp_ptr_NC, iz_tsw, iy_tsw, ix_tsw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tsw * gOut); - safe_add_3d(gInp_ptr_NC, iz_tse, iy_tse, ix_tse, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tse * gOut); - safe_add_3d(gInp_ptr_NC, iz_bnw, iy_bnw, ix_bnw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bnw * gOut); - safe_add_3d(gInp_ptr_NC, iz_bne, iy_bne, ix_bne, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bne * gOut); - safe_add_3d(gInp_ptr_NC, iz_bsw, iy_bsw, ix_bsw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bsw * gOut); - safe_add_3d(gInp_ptr_NC, iz_bse, iy_bse, ix_bse, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bse * gOut); + // calculate and set grad_input. See Note [Passing pointer and offset to fastAtomicAdd]. + safe_add_3d(grad_input.data, iz_tnw, iy_tnw, ix_tnw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tnw * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_tne, iy_tne, ix_tne, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tne * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_tsw, iy_tsw, ix_tsw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tsw * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_tse, iy_tse, ix_tse, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, tse * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_bnw, iy_bnw, ix_bnw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bnw * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_bne, iy_bne, ix_bne, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bne * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_bsw, iy_bsw, ix_bsw, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bsw * gOut, + NC_offset, grad_input_memory_span); + safe_add_3d(grad_input.data, iz_bse, iy_bse, ix_bse, gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, bse * gOut, + NC_offset, grad_input_memory_span); // calculate grad_grid if (within_bounds_3d(iz_tnw, iy_tnw, ix_tnw, inp_D, inp_H, inp_W)) { @@ -665,11 +685,12 @@ namespace { // assign nearest neighor pixel value to output pixel scalar_t *gOut_ptr_NCDHW = grad_output.data + n * gOut_sN + d * gOut_sD + h * gOut_sH + w * gOut_sW; - scalar_t *gInp_ptr_NC = grad_input.data + n * gInp_sN; - for (index_t c = 0; c < C; ++c, gOut_ptr_NCDHW += gOut_sC, gInp_ptr_NC += gInp_sC) { - // calculate and set grad_input - safe_add_3d(gInp_ptr_NC, iz_nearest, iy_nearest, ix_nearest, - gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, *gOut_ptr_NCDHW); + index_t NC_offset = n * gInp_sN; + for (index_t c = 0; c < C; ++c, gOut_ptr_NCDHW += gOut_sC, NC_offset += gInp_sC) { + // calculate and set grad_input. See Note [Passing pointer and offset to fastAtomicAdd]. + safe_add_3d(grad_input.data, iz_nearest, iy_nearest, ix_nearest, + gInp_sD, gInp_sH, gInp_sW, inp_D, inp_H, inp_W, *gOut_ptr_NCDHW, + NC_offset, grad_input_memory_span); } // assuming grad_grid is contiguous @@ -795,7 +816,8 @@ grid_sampler_2d_backward_cuda(const Tensor& grad_output, const Tensor& input, getTensorInfo(grad_grid), static_cast(interpolation_mode), static_cast(padding_mode), - align_corners); + align_corners, + /*grad_input_memory_span =*/static_cast(grad_input.numel())); C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { grid_sampler_2d_backward_kernel @@ -808,7 +830,8 @@ grid_sampler_2d_backward_cuda(const Tensor& grad_output, const Tensor& input, getTensorInfo(grad_grid), static_cast(interpolation_mode), static_cast(padding_mode), - align_corners); + align_corners, + /*grad_input_memory_span =*/grad_input.numel()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } }); @@ -845,7 +868,8 @@ grid_sampler_3d_backward_cuda(const Tensor& grad_output, const Tensor& input, getTensorInfo(grad_grid), static_cast(interpolation_mode), static_cast(padding_mode), - align_corners); + align_corners, + /*grad_input_memory_span =*/static_cast(grad_input.numel())); C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { grid_sampler_3d_backward_kernel @@ -858,7 +882,8 @@ grid_sampler_3d_backward_cuda(const Tensor& grad_output, const Tensor& input, getTensorInfo(grad_grid), static_cast(interpolation_mode), static_cast(padding_mode), - align_corners); + align_corners, + /*grad_input_memory_span =*/grad_input.numel()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } }); diff --git a/aten/src/ATen/native/cuda/GridSampler.cuh b/aten/src/ATen/native/cuda/GridSampler.cuh index 3897f3b10a27a..8f9f700335a75 100644 --- a/aten/src/ATen/native/cuda/GridSampler.cuh +++ b/aten/src/ATen/native/cuda/GridSampler.cuh @@ -1,7 +1,7 @@ #include #include #include -#include +#include namespace at { namespace native { @@ -252,33 +252,47 @@ scalar_t get_value_bounded( return static_cast(0); } -template +template static __forceinline__ __device__ void safe_add_2d(scalar_t *data, int h, int w, int sH, int sW, int H, int W, - scalar_t delta) { + scalar_t delta, + const index_t NC_offset, + const index_t memory_span) { if (within_bounds_2d(h, w, H, W)) { - gpuAtomicAdd(data + h * sH + w * sW, delta); + fastAtomicAdd(data, + NC_offset + h * sH + w * sW, + memory_span, + delta, + true); } } -template +template static __forceinline__ __device__ void safe_add_3d(scalar_t *data, int d, int h, int w, int sD, int sH, int sW, int D, int H, int W, - scalar_t delta) { + scalar_t delta, + const index_t NC_offset, + const index_t memory_span) { if (within_bounds_3d(d, h, w, D, H, W)) { - gpuAtomicAdd(data + d * sD + h * sH + w * sW, delta); + fastAtomicAdd(data, + NC_offset + d * sD + h * sH + w * sW, + memory_span, + delta, + true); } } -template +template static __forceinline__ __device__ void add_value_bounded( scalar_t* data, scalar_t x, scalar_t y, int W, int H, int sW, int sH, scalar_t delta, GridSamplerPadding padding_mode, - bool align_corners) { + bool align_corners, + const index_t NC_offset, + const index_t memory_span) { x = compute_coordinates(x, W, padding_mode, align_corners); y = compute_coordinates(y, H, padding_mode, align_corners); @@ -286,7 +300,7 @@ void add_value_bounded( int ix = static_cast(x); int iy = static_cast(y); - safe_add_2d(data, iy, ix, sH, sW, H, W, delta); + safe_add_2d(data, iy, ix, sH, sW, H, W, delta, NC_offset, memory_span); } // Calculate the differential of the cubic convolution, i.e. `d coeff / d x` diff --git a/aten/src/ATen/native/cuda/IndexKernel.cu b/aten/src/ATen/native/cuda/IndexKernel.cu index 68eb9b98d4a77..f42c1c4988ee9 100644 --- a/aten/src/ATen/native/cuda/IndexKernel.cu +++ b/aten/src/ATen/native/cuda/IndexKernel.cu @@ -332,7 +332,7 @@ void put_kernel(TensorIterator& iter, const Tensor& output, const bool accumulat "put_cuda_index", [&] { auto* __restrict__ indexed_ptr = output.template data(); if (accumulate) { - const auto numel = output.numel(); + index_t numel = output.numel(); cuda_take_put_kernel(iter, output, [numel, indexed_ptr] __device__(scalar_t& iterated, const index_t offset) { fastSpecializedAtomicAdd(indexed_ptr, offset, numel, iterated); diff --git a/aten/src/ATen/native/cuda/KernelUtils.cuh b/aten/src/ATen/native/cuda/KernelUtils.cuh index cdf2f39f11f74..89f35052c7a2b 100644 --- a/aten/src/ATen/native/cuda/KernelUtils.cuh +++ b/aten/src/ATen/native/cuda/KernelUtils.cuh @@ -7,12 +7,13 @@ namespace native { template < typename scalar_t, + typename index_t, typename std::enable_if::value>::type* = nullptr> __device__ __forceinline__ void fastSpecializedAtomicAdd( scalar_t* tensor, - size_t index, - const size_t numel, + index_t index, + const index_t numel, scalar_t value) { #if ( \ (CUDA_VERSION < 10000) || \ @@ -46,21 +47,22 @@ __device__ __forceinline__ void fastSpecializedAtomicAdd( template < typename scalar_t, + typename index_t, typename std::enable_if::value>::type* = nullptr> __device__ __forceinline__ void fastSpecializedAtomicAdd( scalar_t* tensor, - size_t index, - const size_t numel, + index_t index, + const index_t numel, scalar_t value) { gpuAtomicAdd(tensor + index, value); } -template +template __device__ __forceinline__ void fastAtomicAdd( scalar_t* tensor, - size_t index, - const size_t numel, + index_t index, + const index_t numel, scalar_t value, bool fast_atomics) { if (fast_atomics) { diff --git a/docs/source/amp.rst b/docs/source/amp.rst index 8e376c6c2a1bd..d256171d7fa27 100644 --- a/docs/source/amp.rst +++ b/docs/source/amp.rst @@ -141,7 +141,6 @@ Ops that can autocast to ``float32`` ``erfinv``, ``exp``, ``expm1``, -``grid_sample``, ``group_norm``, ``hinge_embedding_loss``, ``kl_div``, @@ -189,6 +188,7 @@ autocast casts all inputs to ``float32`` and runs the op in ``float32``. ``bilinear``, ``cross``, ``dot``, +``grid_sample``, ``index_put``, ``scatter_add``, ``tensordot`` diff --git a/torch/testing/_internal/autocast_test_lists.py b/torch/testing/_internal/autocast_test_lists.py index cfb1f33b7cf1f..754ccca11ed9d 100644 --- a/torch/testing/_internal/autocast_test_lists.py +++ b/torch/testing/_internal/autocast_test_lists.py @@ -177,9 +177,6 @@ def __init__(self, dev): ("renorm", mat0_fp16 + (2, 0, 1.0)), ("sum", pointwise0_fp16), ("sum", mat0_fp16 + (1,)), - ("grid_sampler", (torch.randn((2, 3, 33, 22), dtype=torch.float16, device=dev), - torch.randn((2, 22, 11, 2), dtype=torch.float16, device=dev), - 0, 0, False)), ] self.torch_need_autocast_promote = [ ("addcdiv", pointwise0_fp32 + pointwise1_fp16 + (pointwise2_fp16[0].clamp(0.1, 100),)), @@ -192,6 +189,9 @@ def __init__(self, dev): ("cross", (torch.randn(3, dtype=torch.float32, device=dev), torch.randn(3, dtype=torch.float16, device=dev))), ("dot", pointwise0_fp16 + pointwise1_fp32), + ("grid_sampler", (torch.randn((2, 3, 33, 22), dtype=torch.float16, device=dev), + torch.randn((2, 22, 11, 2), dtype=torch.float32, device=dev), + 0, 0, False)), ("index_put", pointwise0_fp32 + ((torch.tensor([1], device=dev, dtype=torch.long),), torch.randn(1, device=dev, dtype=torch.float16))), ("index_put", pointwise0_fp16 + ((torch.tensor([1], device=dev, dtype=torch.long),), From c16f87949fb9759a71e6e0bb46a62193e4efc90b Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 21 Jun 2021 10:51:49 -0700 Subject: [PATCH 273/305] ENH Adds nn.ReflectionPad3d (#59791) Summary: Fixes https://github.com/pytorch/pytorch/issues/27655 This PR adds a C++ and Python version of ReflectionPad3d with structured kernels. The implementation uses lambdas extensively to better share code from the backward and forward pass. Pull Request resolved: https://github.com/pytorch/pytorch/pull/59791 Reviewed By: gchanan Differential Revision: D29242015 Pulled By: jbschlosser fbshipit-source-id: 18e692d3b49b74082be09f373fc95fb7891e1b56 --- aten/src/ATen/core/aten_interned_strings.h | 3 + aten/src/ATen/native/ReflectionPad.cpp | 432 ++++++++++++++++++ aten/src/ATen/native/cuda/ReflectionPad.cu | 215 +++++++++ aten/src/ATen/native/native_functions.yaml | 22 + docs/source/nn.rst | 1 + test/cpp/api/functional.cpp | 30 ++ test/cpp/api/modules.cpp | 53 +++ test/cpp_api_parity/parity-tracker.md | 1 + test/test_jit.py | 1 + test/test_module_init.py | 1 + test/test_nn.py | 28 +- test/test_torch.py | 12 + tools/autograd/derivatives.yaml | 7 + tools/autograd/gen_variable_type.py | 4 +- torch/__init__.py | 1 + .../api/include/torch/nn/functional/padding.h | 2 +- .../api/include/torch/nn/modules/padding.h | 27 ++ .../api/include/torch/nn/options/padding.h | 10 + torch/csrc/api/src/nn/modules/padding.cpp | 1 + torch/csrc/jit/passes/shape_analysis.cpp | 1 + torch/nn/functional.py | 9 +- torch/nn/modules/__init__.py | 4 +- torch/nn/modules/padding.py | 51 +++ torch/testing/_internal/common_nn.py | 14 + 24 files changed, 919 insertions(+), 11 deletions(-) diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 562a982ab86d2..244cb6e1bb53c 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -593,6 +593,9 @@ _(aten, reflection_pad1d_forward) \ _(aten, reflection_pad2d) \ _(aten, reflection_pad2d_backward) \ _(aten, reflection_pad2d_forward) \ +_(aten, reflection_pad3d) \ +_(aten, reflection_pad3d_backward) \ +_(aten, reflection_pad3d_forward) \ _(aten, relu) \ _(aten, remainder) \ _(aten, renorm) \ diff --git a/aten/src/ATen/native/ReflectionPad.cpp b/aten/src/ATen/native/ReflectionPad.cpp index bbbb840116c91..d8b4ba5d901da 100644 --- a/aten/src/ATen/native/ReflectionPad.cpp +++ b/aten/src/ATen/native/ReflectionPad.cpp @@ -96,6 +96,117 @@ TORCH_META_FUNC(reflection_pad1d_backward)(const Tensor& grad_output, set_output(input.sizes(), input.options()); } +TORCH_META_FUNC(reflection_pad3d)(const Tensor& input, IntArrayRef padding) { + TORCH_CHECK(padding.size() == 6, "padding size is expected to be 6"); + int64_t pad_left = padding[0]; + int64_t pad_right = padding[1]; + int64_t pad_top = padding[2]; + int64_t pad_bottom = padding[3]; + int64_t pad_front = padding[4]; + int64_t pad_back = padding[5]; + int64_t dim_w = 3; + int64_t dim_h = 2; + int64_t dim_d = 1; + int64_t dim_plane = 0; + + // allow batch size of 0-dim. + bool valid_dims = + input.size(1) != 0 && input.size(2) != 0 && input.size(3) != 0; + bool valid_single = input.dim() == 4 && input.size(0) != 0 && valid_dims; + bool valid_batch = input.dim() == 5 && valid_dims && input.size(4) != 0; + + TORCH_CHECK( + valid_single || valid_batch, + "Expected 4D or 5D (batch mode) tensor with possibly 0 batch size and other non-zero dimensions for input, but got: ", + input.sizes()); + + bool batch_mode = (input.dim() == 5); + if (batch_mode) { + dim_w++; + dim_h++; + dim_d++; + dim_plane++; + } + + int64_t nplane = input.size(dim_plane); + int64_t input_d = input.size(dim_d); + int64_t input_h = input.size(dim_h); + int64_t input_w = input.size(dim_w); + int64_t output_d = input_d + pad_front + pad_back; + int64_t output_h = input_h + pad_top + pad_bottom; + int64_t output_w = input_w + pad_left + pad_right; + + TORCH_CHECK( + pad_left < input_w && pad_right < input_w, + "Argument #4: Padding size " + "should be less than the corresponding input dimension, but got: padding (", + pad_left, ", ", pad_right, ") at dimension ", dim_w, " of input ", input.sizes()); + TORCH_CHECK( + pad_top < input_h && pad_bottom < input_h, + "Argument #6: Padding size " + "should be less than the corresponding input dimension, but got: padding (", + pad_top, ", ", pad_bottom, ") at dimension ", dim_h, " of input ", input.sizes()); + TORCH_CHECK( + pad_front < input_d && pad_back < input_d, + "Argument #8: Padding size " + "should be less than the corresponding input dimension, but got: padding (", + pad_front, ", ", pad_back, ") at dimension ", dim_d, " of input ", input.sizes()); + + TORCH_CHECK(output_w >= 1 || output_h >=1 || output_d >= 1, + "input (D: ", input_d, " H: ", input_h, ", W: ", input_w, + ") is too small." + " Calculated output D: ", output_d, " H: ", output_h, " W: ", output_w); + + if (batch_mode) { + set_output({input.size(0), nplane, output_d, output_h, output_w}, input.options()); + } else { + set_output({nplane, output_d, output_h, output_w}, input.options()); + } +} + +TORCH_META_FUNC(reflection_pad3d_backward)( + const Tensor& grad_output, + const Tensor& input, + IntArrayRef padding +) { + TORCH_CHECK(padding.size() == 6, "padding size is expected to be 6"); + TORCH_CHECK(input.dim() > 3); + TORCH_CHECK(grad_output.dim() == input.dim()); + + int64_t pad_left = padding[0]; + int64_t pad_right = padding[1]; + int64_t pad_top = padding[2]; + int64_t pad_bottom = padding[3]; + int64_t pad_front = padding[4]; + int64_t pad_back = padding[5]; + int64_t dim_w = 3; + int64_t dim_h = 2; + int64_t dim_d = 1; + + if (input.dim() == 5) + { + // batch mode + dim_w++; + dim_h++; + dim_d++; + } + + int64_t input_d = input.size(dim_d); + int64_t input_h = input.size(dim_h); + int64_t input_w = input.size(dim_w); + int64_t output_d = input_d + pad_front + pad_back; + int64_t output_h = input_h + pad_top + pad_bottom; + int64_t output_w = input_w + pad_left + pad_right; + + TORCH_CHECK(output_w == grad_output.size(dim_w), "grad_output width unexpected." + " Expected: ", output_w, ", Got: ", grad_output.size(dim_w)); + TORCH_CHECK(output_h == grad_output.size(dim_h), "grad_output height unexpected." + " Expected: ", output_h, ", Got: ", grad_output.size(dim_h)); + TORCH_CHECK(output_d == grad_output.size(dim_d), "grad_output depth unexpected." + " Expected: ", output_h, ", Got: ", grad_output.size(dim_d)); + + set_output(input.sizes(), input.options()); +} } // namespace meta namespace native { @@ -560,6 +671,184 @@ void reflection_pad2d_backward_out_template( ); } } +template +inline void parallel_reflection_pad3d( + int64_t nplane, + int64_t input_w, int64_t input_h, int64_t input_d, + int64_t output_w, int64_t output_h, int64_t output_d, + int64_t pad_left, int64_t pad_top, int64_t pad_front, + const F& f) { + + auto i_start_x = std::max(int64_t(0), -pad_left); + auto i_start_y = std::max(int64_t(0), -pad_top); + auto i_start_z = std::max(int64_t(0), -pad_front); + auto o_start_x = std::max(int64_t(0), pad_left); + auto o_start_y = std::max(int64_t(0), pad_top); + auto o_start_z = std::max(int64_t(0), pad_front); + + at::parallel_for(0, nplane, 0, [&](int64_t start, int64_t end) { + // NOLINTNEXTLINE(cppcoreguidelines-init-variables) + int64_t ip_x, ip_y, ip_z; + for (int64_t k = start; k < end; k++) { + for (int64_t op_z = 0; op_z < output_d; op_z++) { + for (int64_t op_y = 0; op_y < output_h; op_y++) { + for (int64_t op_x = 0; op_x < output_w; op_x++) { + if (op_x < pad_left) { + ip_x = pad_left * 2 - op_x; + } else if (op_x >= pad_left && op_x < input_w + pad_left) { + ip_x = op_x; + } else { + ip_x = (input_w + pad_left - 1) * 2 - op_x; + } + ip_x = ip_x - o_start_x + i_start_x; + + if (op_y < pad_top) { + ip_y = pad_top * 2 - op_y; + } else if (op_y >= pad_top && op_y < input_h + pad_top) { + ip_y = op_y; + } else { + ip_y = (input_h + pad_top - 1) * 2 - op_y; + } + ip_y = ip_y - o_start_y + i_start_y; + + if (op_z < pad_front) { + ip_z = pad_front * 2 - op_z; + } else if (op_z >= pad_front && op_z < input_d + pad_front) { + ip_z = op_z; + } else { + ip_z = (input_d + pad_front - 1) * 2 - op_z; + } + ip_z = ip_z - o_start_z + i_start_z; + + f(k, op_z, op_y, op_x, ip_z, ip_y, ip_x); + } + } + } + } + }); +} + +template +static void reflection_pad3d_out_frame( + scalar_t *input_p, scalar_t *output_p, + int64_t nplane, + int64_t input_w, int64_t input_h, int64_t input_d, + int64_t output_w, int64_t output_h, int64_t output_d, + int64_t pad_left, int64_t pad_top, int64_t pad_front) +{ + parallel_reflection_pad3d( + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front, + [&](int64_t k, + int64_t op_z, + int64_t op_y, + int64_t op_x, + int64_t ip_z, + int64_t ip_y, + int64_t ip_x) { + scalar_t* dest_p = output_p + k * output_w * output_h * output_d + + op_z * output_w * output_h + op_y * output_w + op_x; + scalar_t* src_p = input_p + k * input_w * input_h * input_d + + ip_z * input_w * input_h + ip_y * input_w + ip_x; + *dest_p = *src_p; + }); +} + +template +static void reflection_pad3d_out_loop( + scalar_t *input_p, scalar_t *output_p, + int64_t nbatch, int64_t nplane, + int64_t input_w, int64_t input_h, int64_t input_d, + int64_t output_w, int64_t output_h, int64_t output_d, + int64_t pad_left, int64_t pad_top, int64_t pad_front) +{ + at::parallel_for(0, nbatch, 0, [&](int64_t start, int64_t end) { + for (int64_t p = start; p < end; p++) { + reflection_pad3d_out_frame( + input_p + p * nplane * input_w * input_h * input_d, + output_p + p * nplane * output_w * output_h * output_d, + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + } + }); +} + +template +static void reflection_pad3d_backward_out_frame( + scalar_t *grad_input, scalar_t *grad_output, + int64_t nplane, + int64_t input_w, int64_t input_h, int64_t input_d, + int64_t output_w, int64_t output_h, int64_t output_d, + int64_t pad_left, int64_t pad_top, int64_t pad_front +) { + parallel_reflection_pad3d( + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front, + [&](int64_t k, + int64_t op_z, + int64_t op_y, + int64_t op_x, + int64_t ip_z, + int64_t ip_y, + int64_t ip_x) { + scalar_t* src_p = grad_output + k * output_w * output_h * output_d + + op_z * output_w * output_h + op_y * output_w + op_x; + scalar_t* dest_p = grad_input + k * input_w * input_h * input_d + + ip_z * input_w * input_h + ip_y * input_w + ip_x; + *dest_p += *src_p; + }); +} + +template +static void reflection_pad3d_backward_out_loop( + scalar_t *grad_input, scalar_t *grad_output, + int64_t nbatch, int64_t nplane, + int64_t input_w, int64_t input_h, int64_t input_d, + int64_t output_w, int64_t output_h, int64_t output_d, + int64_t pad_left, int64_t pad_top, int64_t pad_front +) { + at::parallel_for(0, nbatch, 0, [&](int64_t start, int64_t end) { + for (int64_t p = start; p < end; p++) { + reflection_pad3d_backward_out_frame( + grad_input + p * nplane * input_w * input_h * input_d, + grad_output + p * nplane * output_w * output_h * output_d, + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + } + }); +} } // namespace @@ -686,5 +975,148 @@ Tensor reflection_pad2d_backward_cpu( return grad_input; } +TORCH_IMPL_FUNC(reflection_pad3d_out_cpu) +(const Tensor& input_, IntArrayRef padding, const Tensor& output) { + int64_t pad_left = padding[0]; + int64_t pad_top = padding[2]; + int64_t pad_front = padding[4]; + int64_t dim_w = 3; + int64_t dim_h = 2; + int64_t dim_d = 1; + int64_t dim_plane = 0; + bool batch_mode = (input_.dim() == 5); + + if (batch_mode) { + dim_w++; + dim_h++; + dim_d++; + dim_plane++; + } + + int64_t nplane = input_.size(dim_plane); + int64_t input_w = input_.size(dim_w); + int64_t input_h = input_.size(dim_h); + int64_t input_d = input_.size(dim_d); + int64_t output_w = output.size(dim_w); + int64_t output_h = output.size(dim_h); + int64_t output_d = output.size(dim_d); + + auto input = input_.contiguous(); + + if (batch_mode) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1( + kHalf, input.scalar_type(), "replication_pad3d_cpu", [&] { + auto input_data = input.data_ptr(); + auto output_data = output.data_ptr(); + auto nbatch = input.size(0); + reflection_pad3d_out_loop( + input_data, + output_data, + nbatch, + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + }); + } else { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1( + kHalf, input.scalar_type(), "replication_pad3d_cpu", [&] { + auto input_data = input.data_ptr(); + auto output_data = output.data_ptr(); + reflection_pad3d_out_frame( + input_data, + output_data, + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + }); + } +} + +TORCH_IMPL_FUNC(reflection_pad3d_backward_out_cpu)(const Tensor& grad_output, + const Tensor& input, + IntArrayRef padding, + const Tensor& grad_input) { + int64_t pad_left = padding[0]; + int64_t pad_top = padding[2]; + int64_t pad_front = padding[4]; + int64_t dim_w = 3; + int64_t dim_h = 2; + int64_t dim_d = 1; + int64_t dim_plane = 0; + bool batch_mode = (input.dim() == 5); + + if (batch_mode) { + dim_w++; + dim_h++; + dim_d++; + dim_plane++; + } + + int64_t nplane = input.size(dim_plane); + int64_t input_d = input.size(dim_d); + int64_t input_h = input.size(dim_h); + int64_t input_w = input.size(dim_w); + int64_t output_d = grad_output.size(dim_d); + int64_t output_h = grad_output.size(dim_h); + int64_t output_w = grad_output.size(dim_w); + + auto grad_output_ = grad_output.contiguous(); + if (grad_output_.numel() == 0) { + return; + } + + grad_input.zero_(); + + if (batch_mode) { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1( + kHalf, input.scalar_type(), "replication_pad3d_backward_cpu", [&] { + reflection_pad3d_backward_out_loop( + grad_input.data_ptr(), + grad_output_.data_ptr(), + input.size(0), + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + }); + } else { + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1( + kHalf, input.scalar_type(), "replication_pad3d_backward_cpu", [&] { + reflection_pad3d_backward_out_frame( + grad_input.data_ptr(), + grad_output_.data_ptr(), + nplane, + input_w, + input_h, + input_d, + output_w, + output_h, + output_d, + pad_left, + pad_top, + pad_front); + }); + } +} } // namespace native } // namespace at diff --git a/aten/src/ATen/native/cuda/ReflectionPad.cu b/aten/src/ATen/native/cuda/ReflectionPad.cu index 8c9d6fe0fda17..e1bd873ecbb9e 100644 --- a/aten/src/ATen/native/cuda/ReflectionPad.cu +++ b/aten/src/ATen/native/cuda/ReflectionPad.cu @@ -145,6 +145,113 @@ __global__ void reflection_pad2d_backward_out_kernel( gpuAtomicAdd(&grad_input[index_pair.first], grad_output[index_pair.second]); } } +template +__device__ inline void parallel_reflection_pad3d( + PackedTensorAccessor64 input, + PackedTensorAccessor64 output, + int64_t pad_left, + int64_t pad_top, + int64_t pad_front, + int64_t y_shift, + int64_t z_shift, + const F& f) { + int64_t output_id = threadIdx.x + blockIdx.x * blockDim.x; + + if (output_id >= (output.size(2) * output.size(3) * output.size(4))) { + return; + } + + int64_t output_x = output_id % output.size(4); + int64_t output_y = (output_id / output.size(4)) % output.size(3); + int64_t output_z = output_id / (output.size(3) * output.size(4)); + + int64_t i_start_x = ::max(int64_t(0), -pad_left); + int64_t o_start_x = ::max(int64_t(0), pad_left); + int64_t i_start_y = ::max(int64_t(0), -pad_top); + int64_t o_start_y = ::max(int64_t(0), pad_top); + int64_t i_start_z = ::max(int64_t(0), -pad_front); + int64_t o_start_z = ::max(int64_t(0), pad_front); + + int64_t input_x = ::abs(output_x - pad_left) + - ::abs(output_x - (input.size(4) + pad_left - 1)) + - output_x + + 2 * pad_left + input.size(4) - 1 + - o_start_x + i_start_x; + int64_t input_y = ::abs(output_y - pad_top) + - ::abs(output_y - (input.size(3) + pad_top - 1)) + - output_y + + 2 * pad_top + input.size(3) - 1 + - o_start_y + i_start_y; + + int64_t input_z = ::abs(output_z - pad_front) + - ::abs(output_z - (input.size(2) + pad_front - 1)) + - output_z + + 2 * pad_front + input.size(2) - 1 + - o_start_z + i_start_z; + + int64_t plane = blockIdx.y + y_shift; + int64_t batch = blockIdx.z + z_shift; + f(plane, batch, output_z, output_y, output_x, input_z, input_y, input_x); +} + +template +__global__ void reflection_pad3d_out_kernel( + PackedTensorAccessor64 input, + PackedTensorAccessor64 output, + int64_t pad_left, int64_t pad_top, int64_t pad_front, + int64_t y_shift, int64_t z_shift +){ + parallel_reflection_pad3d( + input, + output, + pad_left, + pad_top, + pad_front, + y_shift, + z_shift, + [&] __device__( + int64_t plane, + int64_t batch, + int64_t output_z, + int64_t output_y, + int64_t output_x, + int64_t input_z, + int64_t input_y, + int64_t input_x) { + auto value_to_copy = input[batch][plane][input_z][input_y][input_x]; + output[batch][plane][output_z][output_y][output_x] = value_to_copy; + }); +} + +template +__global__ void reflection_pad3d_backward_out_kernel( + PackedTensorAccessor64 grad_input, + PackedTensorAccessor64 grad_output, + int64_t pad_left, int64_t pad_top, int64_t pad_front, + int64_t y_shift, int64_t z_shift +) { + parallel_reflection_pad3d( + grad_input, + grad_output, + pad_left, + pad_top, + pad_front, + y_shift, + z_shift, + [&] __device__( + int64_t plane, + int64_t batch, + int64_t output_z, + int64_t output_y, + int64_t output_x, + int64_t input_z, + int64_t input_y, + int64_t input_x) { + auto value_to_add = grad_output[batch][plane][output_z][output_y][output_x]; + auto target = &grad_input[batch][plane][input_z][input_y][input_x]; + gpuAtomicAdd(target, value_to_add); + }); +} void reflection_pad2d_out_template( Tensor &output, const Tensor &input_, IntArrayRef padding) { @@ -448,5 +555,113 @@ Tensor reflection_pad2d_backward_cuda( return grad_input; } + +TORCH_IMPL_FUNC(reflection_pad3d_out_cuda) ( + const Tensor& input_, IntArrayRef padding, const Tensor& output + ) { + TORCH_CHECK( + canUse32BitIndexMath(input_), + "input tensor must fit into 32-bit index math"); + + if (output.numel() == 0) { + return; + } + + int64_t pad_left = padding[0]; + int64_t pad_top = padding[2]; + int64_t pad_front = padding[4]; + + auto input = input_.contiguous(); + bool batch_mode = (input.dim() == 5); + + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1(kHalf, + input.scalar_type(), "reflection_pad3d_out_cuda", [&] { + auto input_inner = input; + auto output_inner = output; + if (!batch_mode) { + // non-batch mode + auto input_inner = input.unsqueeze(0); + auto output_inner = output.unsqueeze(0); + } + + auto input_packed = input_inner.packed_accessor64(); + auto output_packed = output_inner.packed_accessor64(); + + int64_t output_plane_size = output_packed.size(2) * output_packed.size(3) * output_packed.size(4); + int64_t size_y = input_packed.size(1); + int64_t size_z = input_packed.size(0); + dim3 block_size(output_plane_size > 256 ? 256 : output_plane_size); + + for (int64_t block_y = 0; block_y < size_y; block_y += 65535) { + int64_t block_y_size = std::min(size_y - block_y, static_cast(65535)); + for (int64_t block_z = 0; block_z < size_z; block_z += 65535) { + int64_t block_z_size = std::min(size_z - block_z, static_cast(65535)); + + dim3 grid_size(at::cuda::ATenCeilDiv(output_plane_size, static_cast(256)), \ + block_y_size, block_z_size); + + reflection_pad3d_out_kernel<<< + grid_size, block_size,0, at::cuda::getCurrentCUDAStream()>>>( + input_packed, output_packed, pad_left, pad_top, pad_front, + block_y, block_z); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + } + }); +} + +TORCH_IMPL_FUNC(reflection_pad3d_backward_out_cuda) ( + const Tensor& grad_output, const Tensor& input, IntArrayRef padding, + const Tensor& grad_input) { + globalContext().alertNotDeterministic("reflection_pad3d_backward_out_cuda"); + TORCH_CHECK(canUse32BitIndexMath(input), "input tensor must fit into 32-bit index math"); + TORCH_CHECK(canUse32BitIndexMath(grad_output), "input tensor must fit into 32-bit index math"); + + if (grad_input.numel() == 0) { + return; + } + grad_input.zero_(); + + int64_t pad_left = padding[0]; + int64_t pad_top = padding[2]; + int64_t pad_front = padding[4]; + + AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1(kHalf, + input.scalar_type(), "reflection_pad3d_backward_out_cuda", [&] { + auto grad_input_ = grad_input; + auto grad_output_ = grad_output; + if (input.dim() == 4) { + // non-batch mode + grad_input_ = grad_input.unsqueeze(0); + grad_output_ = grad_output.unsqueeze(0); + } + + auto grad_input_packed = grad_input_.packed_accessor64(); + auto grad_output_packed = grad_output_.packed_accessor64(); + + int64_t output_plane_size = grad_output_packed.size(2) * + grad_output_packed.size(3) * grad_output_packed.size(4); + int64_t size_y = grad_input_packed.size(1); + int64_t size_z = grad_input_packed.size(0); + dim3 block_size(output_plane_size > 256 ? 256 : output_plane_size); + + for (int64_t block_y = 0; block_y < size_y; block_y += 65535) { + int64_t block_y_size = std::min(size_y - block_y, static_cast(65535)); + for (int64_t block_z = 0; block_z < size_z; block_z += 65535) { + int64_t block_z_size = std::min(size_z - block_z, static_cast(65535)); + + dim3 grid_size(at::cuda::ATenCeilDiv(output_plane_size, static_cast(256)), \ + block_y_size, block_z_size); + + reflection_pad3d_backward_out_kernel<<< + grid_size, block_size,0, at::cuda::getCurrentCUDAStream()>>>( + grad_input_packed, grad_output_packed, pad_left, pad_top, pad_front, + block_y, block_z); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + } + }); +} + } // namespace native } // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 3bd8018c8d0c0..bd6f49b418047 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -8814,6 +8814,28 @@ CPU: reflection_pad2d_backward_cpu CUDA: reflection_pad2d_backward_cuda +- func: reflection_pad3d.out(Tensor self, int[6] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad3d_out_cpu + CUDA: reflection_pad3d_out_cuda + +- func: reflection_pad3d(Tensor self, int[6] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad3d.out + +- func: reflection_pad3d_backward.grad_input(Tensor grad_output, Tensor self, int[6] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad3d_backward_out_cpu + CUDA: reflection_pad3d_backward_out_cuda + +- func: reflection_pad3d_backward(Tensor grad_output, Tensor self, int[6] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad3d_backward.grad_input + - func: replication_pad1d.out(Tensor self, int[2] padding, *, Tensor(a!) out) -> Tensor(a!) python_module: nn structured: True diff --git a/docs/source/nn.rst b/docs/source/nn.rst index 15b0f17f62036..ced9c8968d1ab 100644 --- a/docs/source/nn.rst +++ b/docs/source/nn.rst @@ -114,6 +114,7 @@ Padding Layers nn.ReflectionPad1d nn.ReflectionPad2d + nn.ReflectionPad3d nn.ReplicationPad1d nn.ReplicationPad2d nn.ReplicationPad3d diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp index 74e32bf343f23..eeef66dbd384c 100644 --- a/test/cpp/api/functional.cpp +++ b/test/cpp/api/functional.cpp @@ -2332,6 +2332,36 @@ TEST_F(FunctionalTest, Pad) { ASSERT_EQ(output.sizes(), std::vector({1, 1, 5, 5, 6})); ASSERT_TRUE(output.allclose(expected, 1e-04)); } + { + auto input = torch::arange(18, torch::kDouble).reshape({1, 1, 3, 2, 3}); + auto output = F::pad(input, F::PadFuncOptions({0, 2, 1, 0, 1, 2}).mode(torch::kReflect)); + auto expected = torch::tensor( + {{{{{ 9., 10., 11., 10., 9.}, + { 6., 7., 8., 7., 6.}, + { 9., 10., 11., 10., 9.}}, + + {{ 3., 4., 5., 4., 3.}, + { 0., 1., 2., 1., 0.}, + { 3., 4., 5., 4., 3.}}, + + {{ 9., 10., 11., 10., 9.}, + { 6., 7., 8., 7., 6.}, + { 9., 10., 11., 10., 9.}}, + + {{ 15., 16., 17., 16., 15.}, + { 12., 13., 14., 13., 12.}, + { 15., 16., 17., 16., 15.}}, + + {{ 9., 10., 11., 10., 9.}, + { 6., 7., 8., 7., 6.}, + { 9., 10., 11., 10., 9.}}, + + {{ 3., 4., 5., 4., 3.}, + { 0., 1., 2., 1., 0.}, + { 3., 4., 5., 4., 3.}}}}}, torch::kDouble); + ASSERT_EQ(output.sizes(), std::vector({1, 1, 6, 3, 5})); + ASSERT_TRUE(output.allclose(expected, 1e-04)); + } { auto input = torch::ones({1, 1, 1, 1}, torch::kDouble); auto output = F::pad(input, F::PadFuncOptions({1, 1}).mode(torch::kConstant).value(0)); diff --git a/test/cpp/api/modules.cpp b/test/cpp/api/modules.cpp index 4b22a38343762..bae4fe264fd1b 100644 --- a/test/cpp/api/modules.cpp +++ b/test/cpp/api/modules.cpp @@ -3987,6 +3987,59 @@ TEST_F(ModulesTest, ReflectionPad2d) { } } +// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) +TEST_F(ModulesTest, ReflectionPad3d) { + { + ReflectionPad3d m(ReflectionPad3dOptions(1)); + auto input = torch::arange(8, torch::kFloat).reshape({1, 1, 2, 2, 2}); + auto output = m(input); + auto expected = torch::tensor({{{{{7., 6., 7., 6.}, + {5., 4., 5., 4.}, + {7., 6., 7., 6.}, + {5., 4., 5., 4.}}, + {{3., 2., 3., 2.}, + {1., 0., 1., 0.}, + {3., 2., 3., 2.}, + {1., 0., 1., 0.}}, + {{7., 6., 7., 6.}, + {5., 4., 5., 4.}, + {7., 6., 7., 6.}, + {5., 4., 5., 4.}}, + {{3., 2., 3., 2.}, + {1., 0., 1., 0.}, + {3., 2., 3., 2.}, + {1., 0., 1., 0.}}}}}, torch::kFloat); + ASSERT_TRUE(output.allclose(expected)); + } + { + ReflectionPad3d m(ReflectionPad3dOptions({0, 1, 1, 0, 1, 2})); + auto input = torch::arange(16, torch::kFloat).reshape({1, 1, 4, 2, 2}); + auto output = m(input); + auto expected = torch::tensor({{{{{6., 7., 6.}, + {4., 5., 4.}, + {6., 7., 6.}}, + {{2., 3., 2.}, + {0., 1., 0.}, + {2., 3., 2.}}, + {{6., 7., 6.}, + {4., 5., 4.}, + {6., 7., 6.}}, + {{10., 11., 10.}, + {8., 9., 8.}, + {10., 11., 10.}}, + {{14., 15., 14.}, + {12., 13., 12.}, + {14., 15., 14.}}, + {{10., 11., 10.}, + {8., 9., 8.}, + {10., 11., 10.}}, + {{6., 7., 6.}, + {4., 5., 4.}, + {6., 7., 6.}}}}}, torch::kFloat); + ASSERT_EQ(output.sizes(), std::vector({1, 1, 7, 3, 3})); + ASSERT_TRUE(output.allclose(expected)); + } +} // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) TEST_F(ModulesTest, ReplicationPad1d) { { diff --git a/test/cpp_api_parity/parity-tracker.md b/test/cpp_api_parity/parity-tracker.md index 9252c7fa3adc6..869ef300f6c85 100644 --- a/test/cpp_api_parity/parity-tracker.md +++ b/test/cpp_api_parity/parity-tracker.md @@ -37,6 +37,7 @@ torch::nn::AdaptiveAvgPool2d|Yes|No torch::nn::AdaptiveAvgPool3d|Yes|No torch::nn::ReflectionPad1d|Yes|No torch::nn::ReflectionPad2d|Yes|No +torch::nn::ReflectionPad3d|Yes|No torch::nn::ReplicationPad1d|Yes|No torch::nn::ReplicationPad2d|Yes|No torch::nn::ReplicationPad3d|Yes|No diff --git a/test/test_jit.py b/test/test_jit.py index 27dd5a47ffeb7..fd9a77cc12211 100644 --- a/test/test_jit.py +++ b/test/test_jit.py @@ -897,6 +897,7 @@ def forward(self, input): (Mod(nn.ConstantPad3d(3, 3.5)), torch.randn(16, 3, 10, 20, 30)), (Mod(nn.ReflectionPad1d(2)), torch.arange(8, dtype=torch.float).reshape(1, 2, 4)), (Mod(nn.ReflectionPad2d(2)), torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)), + (Mod(nn.ReflectionPad3d(3)), torch.randn(16, 3, 8, 32, 48)), (Mod(nn.ReplicationPad1d(2)), torch.arange(8, dtype=torch.float).reshape(1, 2, 4)), (Mod(nn.ReplicationPad2d(2)), torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)), (Mod(nn.ReplicationPad3d(3)), torch.randn(16, 3, 8, 32, 48)), diff --git a/test/test_module_init.py b/test/test_module_init.py index 80bb7c7c6662f..984ea50303f4e 100644 --- a/test/test_module_init.py +++ b/test/test_module_init.py @@ -129,6 +129,7 @@ def build_constructor_arg_db(): torch.nn.ReLU: ((), {}), torch.nn.ReflectionPad1d: ((2,), {}), torch.nn.ReflectionPad2d: ((2,), {}), + torch.nn.ReflectionPad3d: ((2,), {}), torch.nn.ReplicationPad1d: ((2,), {}), torch.nn.ReplicationPad2d: ((2,), {}), torch.nn.ReplicationPad3d: ((2,), {}), diff --git a/test/test_nn.py b/test/test_nn.py index dfeeb8fe3ab52..db3e05b7bc2a6 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -12749,7 +12749,8 @@ def test_ReplicationPad3d_large(self, device): def test_ReflectionPad_empty(self, device, dtype): for mod, inp in [ (torch.nn.ReflectionPad1d(2), torch.randn(0, 3, 10, device=device, dtype=dtype)), - (torch.nn.ReflectionPad2d(2), torch.randn(0, 3, 10, 10, device=device, dtype=dtype))]: + (torch.nn.ReflectionPad2d(2), torch.randn(0, 3, 10, 10, device=device, dtype=dtype)), + (torch.nn.ReflectionPad3d(3), torch.randn(0, 3, 10, 10, 10, device=device, dtype=dtype))]: self._test_module_empty_input(mod, inp, check_size=False) with self.assertRaisesRegex(RuntimeError, '2D or 3D'): @@ -12762,6 +12763,11 @@ def test_ReflectionPad_empty(self, device, dtype): inp = torch.randn(3, 0, 10, 10, device=device, dtype=dtype) mod(inp) + with self.assertRaisesRegex(RuntimeError, '4D or 5D'): + mod = torch.nn.ReflectionPad3d(3) + inp = torch.randn(3, 0, 10, 10, 10, device=device, dtype=dtype) + mod(inp) + @onlyCUDA # Test if CPU and GPU results match def test_ReflectionPad2d_large(self, device): shapes = ([2, 65736, 6, 6], [65736, 2, 6, 6]) @@ -12783,6 +12789,26 @@ def test_ReflectionPad2d_large(self, device): self.assertEqual(x.grad, ref_x.grad) + @onlyCUDA # Test if CPU and GPU results match + def test_ReflectionPad3d_large(self, device): + shapes = ([2, 1000, 7, 7, 7], [1000, 2, 7, 7, 7]) + pad = (1, 2, 3, 4, 5, 6) + for shape in shapes: + x = torch.randn(shape, device=device, requires_grad=True) + ref_x = x.detach().cpu().requires_grad_() + + out = F.pad(x, pad, mode='reflect') + ref_out = F.pad(ref_x, pad, mode='reflect') + + self.assertEqual(out, ref_out) + + g = torch.randn_like(out) + ref_g = g.cpu() + + out.backward(g) + ref_out.backward(ref_g) + + self.assertEqual(x.grad, ref_x.grad) @onlyOnCPUAndCUDA @dtypes(torch.float, torch.double) diff --git a/test/test_torch.py b/test/test_torch.py index b5f2f85034c15..d60ccb029f132 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -3891,6 +3891,18 @@ def backward_func(slf, device): backward_func(self, device) + def test_nondeterministic_alert_ReflectionPad3d(self, device): + module = torch.nn.ReflectionPad3d((1, 2, 3, 4, 5, 6)) + input = torch.randn(2, 3, 8, 8, 8, device=device, requires_grad=True) + res = module(input) + grad = torch.ones_like(res) + + @expectedAlertNondeterministic('reflection_pad3d_backward_out_cuda', 'cuda') + def backward_func(slf, device): + res.backward(grad) + + backward_func(self, device) + def test_nondeterministic_alert_ReplicationPad1d(self, device): module = torch.nn.ReplicationPad1d((1, 2)) input = torch.randn(2, 3, 4, device=device, requires_grad=True) diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml index 634d6d01c6109..2707240b75088 100644 --- a/tools/autograd/derivatives.yaml +++ b/tools/autograd/derivatives.yaml @@ -1606,6 +1606,9 @@ - name: reflection_pad2d(Tensor self, int[4] padding) -> Tensor self: reflection_pad2d_backward(grad, self, padding) +- name: reflection_pad3d(Tensor self, int[6] padding) -> Tensor + self: reflection_pad3d_backward(grad, self, padding) + - name: replication_pad1d(Tensor self, int[2] padding) -> Tensor self: replication_pad1d_backward(grad, self, padding) @@ -1875,6 +1878,10 @@ grad_output: reflection_pad2d(grad, padding) self: zeros_like(self) +- name: reflection_pad3d_backward(Tensor grad_output, Tensor self, int[6] padding) -> Tensor + grad_output: reflection_pad3d(grad, padding) + self: zeros_like(self) + - name: replication_pad1d_backward(Tensor grad_output, Tensor self, int[2] padding) -> Tensor grad_output: replication_pad1d(grad, padding) self: zeros_like(self) diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py index 5b667a6cf9692..b62577ba0b0be 100644 --- a/tools/autograd/gen_variable_type.py +++ b/tools/autograd/gen_variable_type.py @@ -94,8 +94,8 @@ 'matrix_exp', 'linalg_eigh', 'cholesky_solve', 'linalg_qr', '_svd_helper', '_fft_c2c', '_fft_r2c', 'linalg_solve', 'sqrt', 'stack', 'gather', 'index_select', 'index_add_', 'linalg_inv', 'linalg_inv_ex', 'l1_loss_backward', 'baddbmm', 'addbmm', 'addmm', 'addmv', 'addr', 'linalg_householder_product', - 'constant_pad_nd', 'reflection_pad1d', 'reflection_pad2d', 'linalg_cholesky_ex', 'linalg_eig', - 'reflection_pad1d_backward', 'reflection_pad2d_backward', 'symeig', '_sparse_sparse_matmul', + 'constant_pad_nd', 'reflection_pad1d', 'reflection_pad2d', 'reflection_pad3d', 'linalg_cholesky_ex', 'linalg_eig', + 'reflection_pad1d_backward', 'reflection_pad2d_backward', 'reflection_pad3d_backward', 'symeig', '_sparse_sparse_matmul', 'replication_pad1d', 'replication_pad2d', 'replication_pad3d', 'take', 'put_', 'replication_pad1d_backward', 'replication_pad2d_backward', 'replication_pad3d_backward', 'diag', 'masked_scatter', 'masked_select', 'index_fill', 'trace', 'polar', 'cumsum', 'rsub', diff --git a/torch/__init__.py b/torch/__init__.py index 605c152122ff7..39ecbcdef4b5b 100644 --- a/torch/__init__.py +++ b/torch/__init__.py @@ -405,6 +405,7 @@ def use_deterministic_algorithms(mode): * :class:`torch.nn.ReflectionPad1d` when attempting to differentiate a CUDA tensor * :class:`torch.nn.ReflectionPad2d` when attempting to differentiate a CUDA tensor + * :class:`torch.nn.ReflectionPad3d` when attempting to differentiate a CUDA tensor * :class:`torch.nn.ReplicationPad1d` when attempting to differentiate a CUDA tensor * :class:`torch.nn.ReplicationPad2d` when attempting to differentiate a CUDA tensor * :class:`torch.nn.ReplicationPad3d` when attempting to differentiate a CUDA tensor diff --git a/torch/csrc/api/include/torch/nn/functional/padding.h b/torch/csrc/api/include/torch/nn/functional/padding.h index 431ffd852eff8..7a9554b03774c 100644 --- a/torch/csrc/api/include/torch/nn/functional/padding.h +++ b/torch/csrc/api/include/torch/nn/functional/padding.h @@ -69,7 +69,7 @@ inline Tensor pad(const Tensor& input, } else if (input.dim() == 5) { TORCH_CHECK(pad.size() == 6, "5D tensors expect 6 values for padding"); if (c10::get_if(&mode)) { - TORCH_CHECK(false, "NotImplementedError"); + return torch::reflection_pad3d(input, pad); } else if (c10::get_if(&mode)) { return torch::replication_pad3d(input, pad); } else if (c10::get_if(&mode)) { diff --git a/torch/csrc/api/include/torch/nn/modules/padding.h b/torch/csrc/api/include/torch/nn/modules/padding.h index 19328a093d7c5..c873e2badfc63 100644 --- a/torch/csrc/api/include/torch/nn/modules/padding.h +++ b/torch/csrc/api/include/torch/nn/modules/padding.h @@ -81,6 +81,33 @@ class TORCH_API ReflectionPad2dImpl : public ReflectionPadImpl<2, ReflectionPad2 /// module storage semantics. TORCH_MODULE(ReflectionPad2d); +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ReflectionPad3d ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +/// Applies ReflectionPad over a 3-D input. +/// See https://pytorch.org/docs/master/nn.html#torch.nn.ReflectionPad3d to learn +/// about the exact behavior of this module. +/// +/// See the documentation for `torch::nn::ReflectionPad3dOptions` class to learn what +/// constructor arguments are supported for this module. +/// +/// Example: +/// ``` +/// ReflectionPad3d model(ReflectionPad3dOptions(1)); +/// ReflectionPad3d model(ReflectionPad3dOptions({1, 1, 2, 0, 1, 2})); +/// ``` +// NOLINTNEXTLINE(bugprone-exception-escape) +class TORCH_API ReflectionPad3dImpl : public ReflectionPadImpl<3, ReflectionPad3dImpl> { + public: + using ReflectionPadImpl<3, ReflectionPad3dImpl>::ReflectionPadImpl; +}; + +/// A `ModuleHolder` subclass for `ReflectionPad3dImpl`. +/// See the documentation for `ReflectionPad3dImpl` class to learn what methods it +/// provides, and examples of how to use `ReflectionPad3d` with `torch::nn::ReflectionPad3dOptions`. +/// See the documentation for `ModuleHolder` to learn about PyTorch's +/// module storage semantics. +TORCH_MODULE(ReflectionPad3d); + // ============================================================================ /// Base class for all (dimension-specialized) ReplicationPad modules. diff --git a/torch/csrc/api/include/torch/nn/options/padding.h b/torch/csrc/api/include/torch/nn/options/padding.h index 517e16e742afa..3265cd6438cef 100644 --- a/torch/csrc/api/include/torch/nn/options/padding.h +++ b/torch/csrc/api/include/torch/nn/options/padding.h @@ -19,6 +19,8 @@ struct TORCH_API ReflectionPadOptions { /// If it is `int`, uses the same padding in all boundaries. /// If it is a 2-`tuple` (for ReflectionPad1d), uses (padding_left, padding_right). /// If it is a 4-`tuple` (for ReflectionPad2d), uses (padding_left, padding_right, padding_top, padding_bottom). + /// If it is a 6-`tuple` (for ReflectionPad3d), uses (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back). + TORCH_ARG(ExpandingArray, padding); }; @@ -38,6 +40,14 @@ using ReflectionPad1dOptions = ReflectionPadOptions<1>; /// ``` using ReflectionPad2dOptions = ReflectionPadOptions<2>; +/// `ReflectionPadOptions` specialized for the `ReflectionPad3d` module. +/// +/// Example: +/// ``` +/// ReflectionPad3d model(ReflectionPad3dOptions({1, 1, 2, 0, 1, 1})); +/// ``` +using ReflectionPad3dOptions = ReflectionPadOptions<3>; + // ============================================================================ /// Options for a `D`-dimensional ReplicationPad module. diff --git a/torch/csrc/api/src/nn/modules/padding.cpp b/torch/csrc/api/src/nn/modules/padding.cpp index 71b860318217e..c2b3238ab6db6 100644 --- a/torch/csrc/api/src/nn/modules/padding.cpp +++ b/torch/csrc/api/src/nn/modules/padding.cpp @@ -27,6 +27,7 @@ void ReflectionPadImpl::pretty_print(std::ostream& stream) const { template class ReflectionPadImpl<1, ReflectionPad1dImpl>; template class ReflectionPadImpl<2, ReflectionPad2dImpl>; +template class ReflectionPadImpl<3, ReflectionPad3dImpl>; // ============================================================================ diff --git a/torch/csrc/jit/passes/shape_analysis.cpp b/torch/csrc/jit/passes/shape_analysis.cpp index d3e60baf8d144..2c42db578a4de 100644 --- a/torch/csrc/jit/passes/shape_analysis.cpp +++ b/torch/csrc/jit/passes/shape_analysis.cpp @@ -1171,6 +1171,7 @@ class ShapePropagator { "aten::max_unpool3d(Tensor self, Tensor indices, int[] output_size, int[] stride, int[] padding) -> Tensor", "aten::reflection_pad1d(Tensor self, int[] padding) -> Tensor", "aten::reflection_pad2d(Tensor self, int[] padding) -> Tensor", + "aten::reflection_pad3d(Tensor self, int[] padding) -> Tensor", "aten::replication_pad1d(Tensor self, int[] padding) -> Tensor", "aten::replication_pad2d(Tensor self, int[] padding) -> Tensor", "aten::replication_pad3d(Tensor self, int[] padding) -> Tensor", diff --git a/torch/nn/functional.py b/torch/nn/functional.py index ac612dcd5914f..2162be77b00ce 100644 --- a/torch/nn/functional.py +++ b/torch/nn/functional.py @@ -4109,10 +4109,9 @@ def _pad(input: Tensor, pad: List[int], mode: str = "constant", value: float = 0 See :class:`torch.nn.ConstantPad2d`, :class:`torch.nn.ReflectionPad2d`, and :class:`torch.nn.ReplicationPad2d` for concrete examples on how each of the padding modes works. Constant padding is implemented for arbitrary dimensions. - Replicate padding is implemented for padding the last 3 dimensions of 5D input - tensor, or the last 2 dimensions of 4D input tensor, or the last dimension of - 3D input tensor. Reflect padding is only implemented for padding the last 2 - dimensions of 4D input tensor, or the last dimension of 3D input tensor. + Replicate and reflection padding is implemented for padding the last 3 + dimensions of 5D input tensor, or the last 2 dimensions of 4D input + tensor, or the last dimension of 3D input tensor. Note: When using the CUDA backend, this operation may induce nondeterministic @@ -4178,7 +4177,7 @@ def _pad(input: Tensor, pad: List[int], mode: str = "constant", value: float = 0 elif input.dim() == 5: assert len(pad) == 6, "5D tensors expect 6 values for padding" if mode == "reflect": - raise NotImplementedError + return torch._C._nn.reflection_pad3d(input, pad) elif mode == "replicate": return torch._C._nn.replication_pad3d(input, pad) elif mode == "circular": diff --git a/torch/nn/modules/__init__.py b/torch/nn/modules/__init__.py index c566d2b87aa9e..f0e68233dd0eb 100644 --- a/torch/nn/modules/__init__.py +++ b/torch/nn/modules/__init__.py @@ -20,7 +20,7 @@ from .instancenorm import InstanceNorm1d, InstanceNorm2d, InstanceNorm3d from .normalization import LocalResponseNorm, CrossMapLRN2d, LayerNorm, GroupNorm from .dropout import Dropout, Dropout2d, Dropout3d, AlphaDropout, FeatureAlphaDropout -from .padding import ReflectionPad1d, ReflectionPad2d, ReplicationPad1d, ReplicationPad2d, \ +from .padding import ReflectionPad1d, ReflectionPad2d, ReflectionPad3d, ReplicationPad1d, ReplicationPad2d, \ ReplicationPad3d, ZeroPad2d, ConstantPad1d, ConstantPad2d, ConstantPad3d from .sparse import Embedding, EmbeddingBag from .rnn import RNNBase, RNN, LSTM, GRU, \ @@ -49,7 +49,7 @@ 'LPPool1d', 'LPPool2d', 'LocalResponseNorm', 'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d', 'LayerNorm', 'GroupNorm', 'SyncBatchNorm', 'Dropout', 'Dropout2d', 'Dropout3d', 'AlphaDropout', 'FeatureAlphaDropout', - 'ReflectionPad1d', 'ReflectionPad2d', 'ReplicationPad2d', 'ReplicationPad1d', 'ReplicationPad3d', + 'ReflectionPad1d', 'ReflectionPad2d', 'ReflectionPad3d', 'ReplicationPad2d', 'ReplicationPad1d', 'ReplicationPad3d', 'CrossMapLRN2d', 'Embedding', 'EmbeddingBag', 'RNNBase', 'RNN', 'LSTM', 'GRU', 'RNNCellBase', 'RNNCell', 'LSTMCell', 'GRUCell', 'PixelShuffle', 'PixelUnshuffle', 'Upsample', 'UpsamplingNearest2d', 'UpsamplingBilinear2d', 'PairwiseDistance', 'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d', 'AdaptiveAvgPool1d', diff --git a/torch/nn/modules/padding.py b/torch/nn/modules/padding.py index 5c6539154d951..3156bbbc72915 100644 --- a/torch/nn/modules/padding.py +++ b/torch/nn/modules/padding.py @@ -267,6 +267,57 @@ def __init__(self, padding: _size_4_t) -> None: self.padding = _quadruple(padding) +class ReflectionPad3d(_ReflectionPadNd): + r"""Pads the input tensor using the reflection of the input boundary. + + For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`. + + Args: + padding (int, tuple): the size of the padding. If is `int`, uses the same + padding in all boundaries. If a 6-`tuple`, uses + (:math:`\text{padding\_left}`, :math:`\text{padding\_right}`, + :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`, + :math:`\text{padding\_front}`, :math:`\text{padding\_back}`) + + Shape: + - Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` + - Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` where + + :math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}` + + :math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}` + + :math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}` + + Examples:: + + >>> m = nn.ReflectionPad3d(1) + >>> input = torch.arange(8, dtype=torch.float).reshape(1, 1, 2, 2, 2) + >>> m(input) + tensor([[[[[7., 6., 7., 6.], + [5., 4., 5., 4.], + [7., 6., 7., 6.], + [5., 4., 5., 4.]], + [[3., 2., 3., 2.], + [1., 0., 1., 0.], + [3., 2., 3., 2.], + [1., 0., 1., 0.]], + [[7., 6., 7., 6.], + [5., 4., 5., 4.], + [7., 6., 7., 6.], + [5., 4., 5., 4.]], + [[3., 2., 3., 2.], + [1., 0., 1., 0.], + [3., 2., 3., 2.], + [1., 0., 1., 0.]]]]]) + """ + padding: Tuple[int, int, int, int, int, int] + + def __init__(self, padding: _size_6_t) -> None: + super(ReflectionPad3d, self).__init__() + self.padding = _ntuple(6)(padding) + + class _ReplicationPadNd(Module): __constants__ = ['padding'] padding: Sequence[int] diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py index 8e56c4537d360..aeef8c4f2a0c5 100644 --- a/torch/testing/_internal/common_nn.py +++ b/torch/testing/_internal/common_nn.py @@ -2205,6 +2205,20 @@ def fractional_max_pool3d_test(test_case): skip_half=True, desc='complex' ), + dict( + module_name='ReflectionPad3d', + constructor_args=((1, 2, 0, 2, 1, 2),), + cpp_constructor_args='torch::nn::ReflectionPad3dOptions({1, 2, 0, 2, 1, 2})', + input_size=(2, 3, 8, 8, 8), + ), + dict( + module_name='ReflectionPad3d', + constructor_args=((1, 2, 0, 2, 1, 2),), + cpp_constructor_args='torch::nn::ReflectionPad3dOptions({1, 2, 0, 2, 1, 2})', + input_fn=lambda: torch.rand(2, 3, 8, 8, 8, dtype=torch.complex128, requires_grad=True), + skip_half=True, + desc='complex' + ), dict( module_name='ReplicationPad1d', constructor_args=((1, 2),), From 27e34f731ac55e2f492af7a0db812e7028bb3a5e Mon Sep 17 00:00:00 2001 From: David Riazati Date: Mon, 21 Jun 2021 11:02:10 -0700 Subject: [PATCH 274/305] Re-enable clang-tidy on PRs (#60297) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60297 This switches clang-tidy to the fresh tag from https://github.com/pytorch/test-infra/runs/2860763986 which has a fix for the missing OMP headers we were seeing. Along with #60225 this should restore clang-tidy to normal functionality and we shouldn't see any spurious warnings. Test Plan: Imported from OSS Reviewed By: seemethere, 1ntEgr8 Differential Revision: D29239783 Pulled By: driazati fbshipit-source-id: b1893256fdb27436af03d6c5279e81f64b47fe6b --- .github/workflows/lint.yml | 6 ++---- tools/clang_tidy.py | 27 +++++++++++++++------------ 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 68b9853868ff1..e4aef1bd1eb67 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -266,13 +266,11 @@ jobs: fi clang-tidy: - # if: github.event_name == 'pull_request' - # TODO: Fix clang-tidy, see https://github.com/pytorch/pytorch/issues/60192 - if: ${{ false }} + if: github.event_name == 'pull_request' runs-on: ubuntu-18.04 container: # ubuntu18.04-cuda10.2-py3.6-tidy11 - image: ghcr.io/pytorch/cilint-clang-tidy:52a8ad78d49fc9f40241fee7988db48c920499df + image: ghcr.io/pytorch/cilint-clang-tidy:7f0b4616100071a4813318bfdbd5b06ae36c5272 steps: - name: Checkout PyTorch uses: actions/checkout@v2 diff --git a/tools/clang_tidy.py b/tools/clang_tidy.py index 0087b68e70339..7fd40acc480e0 100755 --- a/tools/clang_tidy.py +++ b/tools/clang_tidy.py @@ -41,13 +41,6 @@ # (c/cc/cpp) file. DEFAULT_FILE_PATTERN = re.compile(r"^.*\.c(c|pp)?$") -# Search for: -# diff --git ... -# index ... -# --- ... -# +++ ... -CHUNK_HEADER_RE = r"diff --git .*?\nindex.*?\n---.*?\n\+\+\+ b/(.*?)\n@@ -(\d+,\d+) \+(\d+,\d+) @@" - CLANG_WARNING_PATTERN = re.compile(r"([^:]+):(\d+):\d+:\s+warning:.*\[([^\]]+)\]") @@ -136,14 +129,24 @@ def get_all_files(paths: List[str]) -> List[str]: def find_changed_lines(diff: str) -> Dict[str, List[Tuple[int, int]]]: + # Delay import since this isn't required unless using the --diff-file + # argument, which for local runs people don't care about + try: + import unidiff # type: ignore[import] + except ImportError as e: + e.msg += ", run 'pip install unidiff'" # type: ignore[attr-defined] + raise e + files = collections.defaultdict(list) - matches = re.findall(CHUNK_HEADER_RE, diff, re.MULTILINE) - for file, start, end in matches: - start_line, _ = start.split(",") - end_line, _ = end.split(",") + for file in unidiff.PatchSet(diff): + for hunk in file: + start = hunk[0].target_line_no + if start is None: + start = 1 + end = hunk[-1].target_line_no - files[file].append((start_line, end_line)) + files[file.path].append((start, end)) return dict(files) From 023907a6fe4e60dec1c5bb0e8ddcd5880e1e00fd Mon Sep 17 00:00:00 2001 From: Sam Estep Date: Mon, 21 Jun 2021 11:26:32 -0700 Subject: [PATCH 275/305] Allow Docker build on macOS (#60375) Summary: This PR allows developers using macOS to build Docker images locally. The `basename $(mktemp -u)` part was suggested by seemethere; I modified it slightly to appease ShellCheck and because [Docker doesn't allow uppercase characters in tags](https://stackoverflow.com/a/54291205). Pull Request resolved: https://github.com/pytorch/pytorch/pull/60375 Test Plan: On a Mac: ``` cd .circleci/docker ./build.sh pytorch-linux-xenial-py3.6-gcc5.4 ``` Reviewed By: driazati Differential Revision: D29267025 Pulled By: samestep fbshipit-source-id: ba27d2fb108f573a50db069cf9ddea0414ed6074 --- .circleci/docker/build.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/docker/build.sh b/.circleci/docker/build.sh index b1072256b4a43..df41b93e4e277 100755 --- a/.circleci/docker/build.sh +++ b/.circleci/docker/build.sh @@ -320,7 +320,7 @@ if [ -n "${JENKINS:-}" ]; then JENKINS_GID=$(id -g jenkins) fi -tmp_tag="tmp-$(cat /dev/urandom | tr -dc 'a-z' | head -c 32)" +tmp_tag=$(basename "$(mktemp -u)" | tr '[:upper:]' '[:lower:]') # Build image # TODO: build-arg THRIFT is not turned on for any image, remove it once we confirm From eaa36ee679aac320edddca31d55201973626ab55 Mon Sep 17 00:00:00 2001 From: Jane Xu Date: Mon, 21 Jun 2021 11:41:02 -0700 Subject: [PATCH 276/305] Enable sharding for Windows GHA CI (#59970) Summary: Enables sharding for Windows on CI. To make that possible, we currently remove the smoke tests tested in shard 1 which don't seem all that important as they are 1. tested on nightlies 2. seems to be tested anyway by running the test suite Pull Request resolved: https://github.com/pytorch/pytorch/pull/59970 Reviewed By: seemethere Differential Revision: D29268484 Pulled By: janeyx99 fbshipit-source-id: 7f90d73037cfeb2c267b28714550316eb471b4dd --- .github/scripts/generate_ci_workflows.py | 5 ++++- .github/workflows/pytorch-win-vs2019-cpu-py3.yml | 4 ++-- .../workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml | 4 ++-- .../workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml | 4 ++-- .../win-test-helpers/test_python_first_shard.bat | 10 ---------- 5 files changed, 10 insertions(+), 17 deletions(-) diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index ca15fbbb1fdc6..80df24533f0bf 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -81,18 +81,21 @@ def generate_workflow_file( cuda_version="cpu", test_runner_type=WINDOWS_CPU_TEST_RUNNER, on_pull_request=True, + num_test_shards=2, ), PyTorchWindowsWorkflow( build_environment="pytorch-win-vs2019-cuda10-cudnn7-py3", cuda_version="10.1", test_runner_type=WINDOWS_CUDA_TEST_RUNNER, on_pull_request=True, - only_build_on_pull_request=True + only_build_on_pull_request=True, + num_test_shards=2, ), PyTorchWindowsWorkflow( build_environment="pytorch-win-vs2019-cuda11-cudnn8-py3", cuda_version="11.1", test_runner_type=WINDOWS_CUDA_TEST_RUNNER, + num_test_shards=2, ) ] diff --git a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml index 2d9cd8535b009..c9eaca4e3d3de 100644 --- a/.github/workflows/pytorch-win-vs2019-cpu-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cpu-py3.yml @@ -73,7 +73,7 @@ jobs: if: ${{ github.repository_owner == 'pytorch' }} runs-on: ubuntu-18.04 env: - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} container: @@ -93,7 +93,7 @@ jobs: runs-on: windows.4xlarge env: JOB_BASE_NAME: pytorch-win-vs2019-cpu-py3-test - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 TEST_CONFIG: ${{ matrix.test_config }} needs: - build diff --git a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml index 8329a1904d804..46653bc786a5b 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda10-cudnn7-py3.yml @@ -83,7 +83,7 @@ jobs: if: ${{ github.repository_owner == 'pytorch' && github.event_name == 'push' }} runs-on: ubuntu-18.04 env: - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} container: @@ -104,7 +104,7 @@ jobs: runs-on: windows.8xlarge.nvidia.gpu env: JOB_BASE_NAME: pytorch-win-vs2019-cuda10-cudnn7-py3-test - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 TEST_CONFIG: ${{ matrix.test_config }} needs: - build diff --git a/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml index 2a78b77c1bc26..54e55a52769b6 100644 --- a/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml +++ b/.github/workflows/pytorch-win-vs2019-cuda11-cudnn8-py3.yml @@ -82,7 +82,7 @@ jobs: if: ${{ github.repository_owner == 'pytorch' }} runs-on: ubuntu-18.04 env: - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 outputs: matrix: ${{ steps.set-matrix.outputs.matrix }} container: @@ -102,7 +102,7 @@ jobs: runs-on: windows.8xlarge.nvidia.gpu env: JOB_BASE_NAME: pytorch-win-vs2019-cuda11-cudnn8-py3-test - NUM_TEST_SHARDS: 1 + NUM_TEST_SHARDS: 2 TEST_CONFIG: ${{ matrix.test_config }} needs: - build diff --git a/.jenkins/pytorch/win-test-helpers/test_python_first_shard.bat b/.jenkins/pytorch/win-test-helpers/test_python_first_shard.bat index 05203a60ccf39..f580dd8d76a75 100644 --- a/.jenkins/pytorch/win-test-helpers/test_python_first_shard.bat +++ b/.jenkins/pytorch/win-test-helpers/test_python_first_shard.bat @@ -4,16 +4,6 @@ echo Copying over test times file copy /Y "%PYTORCH_FINAL_PACKAGE_DIR_WIN%\.pytorch-test-times.json" "%TEST_DIR_WIN%" pushd test - -echo Some smoke tests -"C:\Program Files (x86)\Windows Kits\10\Debuggers\x64\gflags.exe" /i python.exe +sls -python %SCRIPT_HELPERS_DIR%\run_python_nn_smoketests.py -if ERRORLEVEL 1 exit /b 1 - -"C:\Program Files (x86)\Windows Kits\10\Debuggers\x64\gflags.exe" /i python.exe -sls -if ERRORLEVEL 1 exit /b 1 - -echo Run nn tests python run_test.py --exclude-jit-executor --shard 1 2 --verbose --determine-from="%1" if ERRORLEVEL 1 exit /b 1 From 2f3be2735f85525b4b51ebdb97d4004a06d42533 Mon Sep 17 00:00:00 2001 From: Michael Wootton Date: Mon, 21 Jun 2021 11:43:11 -0700 Subject: [PATCH 277/305] Don't split oversize cached blocks (#44742) Summary: Fixes https://github.com/pytorch/pytorch/issues/35901 This change is designed to prevent fragmentation in the Caching Allocator. Permissive block splitting in the allocator allows very large blocks to be split into many pieces. Once split too finely it is unlikely all pieces will be 'free' at that same time so the original allocation can never be returned. Anecdotally, we've seen a model run out of memory failing to alloc a 50 MB block on a 32 GB card while the caching allocator is holding 13 GB of 'split free blocks' Approach: - Large blocks above a certain size are designated "oversize". This limit is currently set 1 decade above large, 200 MB - Oversize blocks can not be split - Oversize blocks must closely match the requested size (e.g. a 200 MB request will match an existing 205 MB block, but not a 300 MB block) - In lieu of splitting oversize blocks there is a mechanism to quickly free a single oversize block (to the system allocator) to allow an appropriate size block to be allocated. This will be activated under memory pressure and will prevent _release_cached_blocks()_ from triggering Initial performance tests show this is similar or quicker than the original strategy. Additional tests are ongoing. Pull Request resolved: https://github.com/pytorch/pytorch/pull/44742 Reviewed By: zou3519 Differential Revision: D29186394 Pulled By: ezyang fbshipit-source-id: c88918836db3f51df59de6d1b3e03602ebe306a9 --- c10/cuda/CUDACachingAllocator.cpp | 211 +++++++++++++++++++++++++----- c10/cuda/CUDACachingAllocator.h | 9 ++ docs/source/notes/cuda.rst | 15 +++ torch/csrc/cuda/Module.cpp | 3 + torch/cuda/memory.py | 34 +++++ torch/utils/collect_env.py | 7 + 6 files changed, 245 insertions(+), 34 deletions(-) diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp index f7d84d96722ae..b48f75560f58b 100644 --- a/c10/cuda/CUDACachingAllocator.cpp +++ b/c10/cuda/CUDACachingAllocator.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -35,13 +36,18 @@ namespace CUDACachingAllocator { // - The allocator attempts to find the smallest cached block that will fit the // requested size. If the block is larger than the requested size, it may be // split. If no block is found, the allocator will delegate to cudaMalloc. -// - If the cudaMalloc fails, the allocator will free all cached blocks that -// are not split and retry the allocation. +// - If the cudaMalloc fails, the allocator will attempt to free one cached +// block of sufficient size that is not split and retry the allocation. +// If this also fails, the allocator will attempt to free all cached blocks +// that are not split and retry the allocation. // - Large (>1MB) and small allocations are stored in separate pools. // Small requests are packed into 2MB buffers. Large requests will use the // smallest available free block or allocate a new block using cudaMalloc. -// To reduce fragmentation, requests between 1MB and 10MB will allocate and +// - To reduce fragmentation, requests between 1MB and 10MB will allocate and // split a 20MB block, if no free block of sufficient size is available. +// - To further reduce fragmentation, blocks >= 200MB are not allowed to be +// split. These oversize cached blocks will still satisfy requests within +// 20MB of the oversize cached block size. // // With this allocator, allocations and frees should logically be considered // "usages" of the memory segment associated with streams, just like kernel @@ -243,13 +249,13 @@ struct AllocParams { block(nullptr), err(cudaSuccess) {} - int device() { + int device() const { return search_key.device; } - cudaStream_t stream() { + cudaStream_t stream() const { return search_key.stream; } - size_t size() { + size_t size() const { return search_key.size; } @@ -310,6 +316,67 @@ cudaError_t cudaMallocMaybeCapturing(void** p, size_t size) { } // namespace +class CachingAllocatorConfig { + public: + static size_t max_split_size() { + return instance().m_max_split_size; + } + + private: + static std::once_flag s_flag; + static CachingAllocatorConfig* s_instance; + static CachingAllocatorConfig& instance() { + std::call_once(s_flag, &CachingAllocatorConfig::init); + return *s_instance; + } + static void init() { + s_instance = new CachingAllocatorConfig(); + s_instance->parseArgs(); + } + + CachingAllocatorConfig() + : m_max_split_size(std::numeric_limits::max()) {} + size_t m_max_split_size; + + void parseArgs() { + const char* val = getenv("PYTORCH_CUDA_ALLOC_CONF"); + if (val != NULL) { + const std::string config(val); + + std::regex exp("[\\s,]+"); + std::sregex_token_iterator it(config.begin(), config.end(), exp, -1); + std::sregex_token_iterator end; + std::vector options(it, end); + + for (auto option : options) { + std::regex exp2("[:]+"); + std::sregex_token_iterator it2(option.begin(), option.end(), exp2, -1); + std::sregex_token_iterator end2; + std::vector kv(it2, end2); + if (kv.size() >= 2) { + /* Maximum split size in MB. Limited to large size blocks */ + if (kv[0].compare("max_split_size_mb") == 0) { + size_t val2 = stoi(kv[1]); + TORCH_CHECK( + val2 > kLargeBuffer / (1024 * 1024), + "CachingAllocator option max_split_size_mb too small, must be >= ", + kLargeBuffer / (1024 * 1024), + ""); + val2 = std::max(val2, kLargeBuffer / (1024 * 1024)); + val2 = std::min( + val2, (std::numeric_limits::max() / (1024 * 1024))); + m_max_split_size = val2 * 1024 * 1024; + } else { + TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", kv[0]); + } + } + } + } + } +}; +CachingAllocatorConfig* CachingAllocatorConfig::s_instance; +std::once_flag CachingAllocatorConfig::s_flag; + class DeviceCachingAllocator { private: // lock around all operations @@ -363,7 +430,9 @@ class DeviceCachingAllocator { public: DeviceCachingAllocator() : large_blocks(BlockComparator, /*is_small=*/false), - small_blocks(BlockComparator, /*is_small=*/true) {} + small_blocks(BlockComparator, /*is_small=*/true) { + stats.max_split_size = CachingAllocatorConfig::max_split_size(); + } // All public methods (except the above) acquire the allocator mutex. // Thus, do not call a public method from another public method. @@ -399,8 +468,11 @@ class DeviceCachingAllocator { || (trigger_free_memory_callbacks(params) && get_free_block(params)) // Attempt allocate || alloc_block(params, false) + // Free enough available cached blocks to satisfy alloc and retry alloc. + || + (release_available_cached_blocks(params) && alloc_block(params, false)) // Free all non-split cached blocks and retry alloc. - || (free_cached_blocks() && alloc_block(params, true)); + || (release_cached_blocks() && alloc_block(params, true)); if (!block_found) { // For any error code other than cudaErrorMemoryAllocation, @@ -456,7 +528,10 @@ class DeviceCachingAllocator { format_size( stats.reserved_bytes[static_cast(StatType::AGGREGATE)] .current), - " reserved in total by PyTorch)"); + " reserved in total by PyTorch)", + " If reserved memory is >> allocated memory try setting max_split_size_mb to avoid" + " fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF", + ""); } TORCH_INTERNAL_ASSERT( @@ -511,6 +586,8 @@ class DeviceCachingAllocator { update_stat_array(stats.allocated_bytes, block->size, params.stat_types); update_stat_array(stats.active, 1, params.stat_types); update_stat_array(stats.active_bytes, block->size, params.stat_types); + if (block->size >= CachingAllocatorConfig::max_split_size()) + update_stat(stats.oversize_allocations, 1); return block; } @@ -529,6 +606,8 @@ class DeviceCachingAllocator { true; update_stat_array(stats.allocation, -1, {stat_types}); update_stat_array(stats.allocated_bytes, -block->size, {stat_types}); + if (block->size >= CachingAllocatorConfig::max_split_size()) + update_stat(stats.oversize_allocations, -1); if (!block->stream_uses.empty()) { if (C10_UNLIKELY(captures_underway)) { @@ -584,7 +663,7 @@ class DeviceCachingAllocator { /** returns cached blocks to the system allocator **/ void emptyCache() { std::lock_guard lock(mutex); - free_cached_blocks(); + release_cached_blocks(); } /** Retrieves info (total size + largest block) of the memory cache **/ @@ -629,6 +708,8 @@ class DeviceCachingAllocator { stats.num_alloc_retries = 0; stats.num_ooms = 0; + reset_accumulated_stat(stats.oversize_allocations); + reset_accumulated_stat(stats.oversize_segments); } /** Resets the historical peak stats for the device **/ @@ -646,6 +727,8 @@ class DeviceCachingAllocator { reset_peak_stat(stats.active_bytes[statType]); reset_peak_stat(stats.inactive_split_bytes[statType]); } + reset_peak_stat(stats.oversize_allocations); + reset_peak_stat(stats.oversize_segments); } /** Dump a complete snapshot of the memory held by the allocator. Potentially @@ -905,8 +988,12 @@ class DeviceCachingAllocator { bool should_split(const Block* block, size_t size) { size_t remaining = block->size - size; - return (block->pool->is_small) ? (remaining >= kMinBlockSize) - : (remaining > kSmallSize); + if (block->pool->is_small) { + return remaining >= kMinBlockSize; + } else { + return (size < CachingAllocatorConfig::max_split_size()) && + (remaining > kSmallSize); + } } static size_t get_allocation_size(size_t size) { @@ -924,6 +1011,14 @@ class DeviceCachingAllocator { auto it = pool.blocks.lower_bound(&p.search_key); if (it == pool.blocks.end() || (*it)->stream != p.stream()) return false; + // Do not return an oversized block for a large request + if ((p.size() < CachingAllocatorConfig::max_split_size()) && + ((*it)->size >= CachingAllocatorConfig::max_split_size())) + return false; + // Allow oversized block size to be rounded up but within a limit + if ((p.size() >= CachingAllocatorConfig::max_split_size()) && + ((*it)->size >= p.size() + kLargeBuffer)) + return false; p.block = *it; pool.blocks.erase(it); return true; @@ -985,27 +1080,71 @@ class DeviceCachingAllocator { p.block = new Block(p.device(), p.stream(), size, p.pool, (char*)ptr); update_stat_array(stats.segment, 1, p.stat_types); update_stat_array(stats.reserved_bytes, size, p.stat_types); + if (size >= CachingAllocatorConfig::max_split_size()) + update_stat(stats.oversize_segments, 1); // p.block came from new, not cudaMalloc. It should not be nullptr here. TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr); return true; } - bool free_cached_blocks() { + /** Free one or more oversize blocks to the system allocator. But only enough + * **/ + /** to satisfy the target size **/ + bool release_available_cached_blocks(const AllocParams& p) { + if (CachingAllocatorConfig::max_split_size() == + std::numeric_limits::max()) + return false; + BlockPool& pool = *p.pool; + Block key = p.search_key; + key.size = (key.size < CachingAllocatorConfig::max_split_size()) + ? CachingAllocatorConfig::max_split_size() + : key.size; + auto it = pool.blocks.lower_bound(&key); + if (it == pool.blocks.end() || (*it)->stream != p.stream()) { + // No single block is large enough; free multiple oversize blocks, + // starting with the largest + if (it == pool.blocks.begin()) + return false; + size_t totalReleased = 0; + --it; // Back up one item. Now on the largest block for the correct + // stream + while ((totalReleased < key.size) && + ((*it)->size >= CachingAllocatorConfig::max_split_size()) && + ((*it)->stream == p.stream())) { + auto cur = it; + totalReleased += (*it)->size; + if (it != pool.blocks.begin()) { + --it; + release_block(*cur); + } else { + release_block(*cur); + break; + } + } + if (totalReleased < key.size) + return false; + } else { + release_block(*it); + } + return true; + } + + bool release_cached_blocks() { // First ensure that all blocks that can't currently be allocated due to // outstanding events are returned to the pool. synchronize_and_free_events(); - // Free all non-split cached blocks - free_blocks(large_blocks); - free_blocks(small_blocks); + // Free all non-split cached blocks to system allocator + release_blocks(large_blocks); + release_blocks(small_blocks); for (auto it = graph_pools_freeable.begin(); it != graph_pools_freeable.end();) { // See notifyCaptureDestroy for the strategy here. TORCH_INTERNAL_ASSERT(it->second->use_count == 0); - free_blocks(it->second->small_blocks); - free_blocks(it->second->large_blocks); + release_blocks(it->second->small_blocks); + release_blocks(it->second->large_blocks); if (it->second->cudaMalloc_count == 0) { auto erase_count = graph_pools.erase(it->first); TORCH_INTERNAL_ASSERT(erase_count == 1); @@ -1018,33 +1157,37 @@ class DeviceCachingAllocator { return true; } - void free_blocks(BlockPool& pool) { + void release_block(Block* block) { + C10_CUDA_CHECK(cudaFree((void*)block->ptr)); + total_allocated_memory -= block->size; + + StatTypes stat_types; + stat_types[static_cast(StatType::AGGREGATE)] = true; + stat_types[static_cast(get_stat_type_for_pool(*(block->pool)))] = + true; + update_stat_array(stats.segment, -1, stat_types); + update_stat_array(stats.reserved_bytes, -block->size, stat_types); + if (block->size >= CachingAllocatorConfig::max_split_size()) + update_stat(stats.oversize_segments, -1); + + block->pool->blocks.erase(block); + delete block; + } + + void release_blocks(BlockPool& pool) { // Frees all non-split blocks auto it = pool.blocks.begin(); while (it != pool.blocks.end()) { Block* block = *it; + ++it; if (!block->prev && !block->next) { - C10_CUDA_CHECK(cudaFree((void*)block->ptr)); - total_allocated_memory -= block->size; + release_block(block); if (pool.owner_PrivatePool) { // The cudaFreed block belonged to a CUDA graph's PrivatePool. TORCH_INTERNAL_ASSERT(pool.owner_PrivatePool->cudaMalloc_count > 0); pool.owner_PrivatePool->cudaMalloc_count--; } - - StatTypes stat_types; - stat_types[static_cast(StatType::AGGREGATE)] = true; - stat_types[static_cast(get_stat_type_for_pool(pool))] = true; - update_stat_array(stats.segment, -1, stat_types); - update_stat_array(stats.reserved_bytes, -block->size, stat_types); - - auto cur = it; - ++it; - pool.blocks.erase(cur); - delete block; - } else { - ++it; } } } diff --git a/c10/cuda/CUDACachingAllocator.h b/c10/cuda/CUDACachingAllocator.h index 51f2c87e62871..05d3702e0b380 100644 --- a/c10/cuda/CUDACachingAllocator.h +++ b/c10/cuda/CUDACachingAllocator.h @@ -87,6 +87,15 @@ struct DeviceStats { // COUNT: total number of OOMs (i.e. failed calls to CUDA after cache flush) int64_t num_ooms = 0; + + // COUNT: total number of oversize blocks allocated from pool + Stat oversize_allocations; + + // COUNT: total number of oversize blocks requiring malloc + Stat oversize_segments; + + // SIZE: maximum block size that is allowed to be split. + int64_t max_split_size = 0; }; // Struct containing info of an allocation block (i.e. a fractional part of a diff --git a/docs/source/notes/cuda.rst b/docs/source/notes/cuda.rst index 41b6b1c9257a7..d19b4601150ac 100644 --- a/docs/source/notes/cuda.rst +++ b/docs/source/notes/cuda.rst @@ -290,6 +290,21 @@ Use of a caching allocator can interfere with memory checking tools such as ``cuda-memcheck``. To debug memory errors using ``cuda-memcheck``, set ``PYTORCH_NO_CUDA_MEMORY_CACHING=1`` in your environment to disable caching. +The behavior of caching allocator can be controlled via environment variable +``PYTORCH_CUDA_ALLOC_CONF``. +The format is ``PYTORCH_CUDA_ALLOC_CONF=