From 4883d39c6fd38431bdc60e1db6402e251429b1e1 Mon Sep 17 00:00:00 2001 From: Martin Yuan Date: Tue, 17 Nov 2020 17:30:37 -0800 Subject: [PATCH] Avoid direct reference to at::native::tensor from TensorDataContainer (#47567) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/47567 Test Plan: Imported from OSS Reviewed By: ezyang Differential Revision: D24822517 Pulled By: iseeyuan fbshipit-source-id: f69bfc029aae5199dbc63193fc7a5e5e6feb5790 --- aten/src/ATen/Utils.cpp | 80 +++++++++++++++++++++-- aten/src/ATen/Utils.h | 17 +++++ aten/src/ATen/native/TensorFactories.cpp | 42 ++---------- aten/src/ATen/templates/Functions.cpp | 2 - aten/src/ATen/templates/Functions.h | 2 - aten/src/ATen/templates/NativeFunctions.h | 18 ++--- 6 files changed, 103 insertions(+), 58 deletions(-) diff --git a/aten/src/ATen/Utils.cpp b/aten/src/ATen/Utils.cpp index 9e1c33a4dbb9..a2e5a82c5d06 100644 --- a/aten/src/ATen/Utils.cpp +++ b/aten/src/ATen/Utils.cpp @@ -1,10 +1,12 @@ #include +#include +#include +#include +#include #include +#include #include #include -#include -#include -#include namespace at { @@ -15,9 +17,16 @@ int _crash_if_asan(int arg) { } namespace detail { -// empty_cpu is used in ScalarOps.h, which can be referenced by other ATen files. Since we want to decouple direct referencing native symbols and only access native symbols through dispatching, we move its implementation here. -Tensor empty_cpu(IntArrayRef size, c10::optional dtype_opt, c10::optional layout_opt, - c10::optional device_opt, c10::optional pin_memory_opt, c10::optional memory_format_opt) { +// empty_cpu is used in ScalarOps.h, which can be referenced by other ATen +// files. Since we want to decouple direct referencing native symbols and only +// access native symbols through dispatching, we move its implementation here. +Tensor empty_cpu( + IntArrayRef size, + c10::optional dtype_opt, + c10::optional layout_opt, + c10::optional device_opt, + c10::optional pin_memory_opt, + c10::optional memory_format_opt) { Device device = device_or_default(device_opt); TORCH_CHECK(device.type() == DeviceType::CPU); @@ -53,6 +62,63 @@ Tensor empty_cpu(IntArrayRef size, c10::optional dtype_opt, c10::opt return tensor; } + +template +Tensor tensor_cpu(ArrayRef values, const TensorOptions& options) { + auto result = at::empty(values.size(), options); + AT_ASSERT(result.is_contiguous()); + AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "tensor_cpu", [&] { + std::copy( + values.begin(), values.end(), result.template data_ptr()); + }); + return result; +} + +template +Tensor tensor_backend(ArrayRef values, const TensorOptions& options) { + auto cpu_tensor = tensor_cpu(values, options.device(DeviceType::CPU)); + return cpu_tensor.to(options.device()); +} + +template +Tensor tensor_complex_cpu(ArrayRef values, const TensorOptions& options) { + auto result = at::empty(values.size(), options); + AT_ASSERT(result.is_contiguous()); + AT_DISPATCH_COMPLEX_TYPES(result.scalar_type(), "tensor_cpu", [&] { + std::copy( + values.begin(), values.end(), result.template data_ptr()); + }); + return result; +} + +template +Tensor tensor_complex_backend( + ArrayRef values, + const TensorOptions& options) { + auto cpu_tensor = tensor_complex_cpu(values, options.device(DeviceType::CPU)); + return cpu_tensor.to(options.device()); +} } // namespace detail -} // at +#define TENSOR(T, _1) \ + Tensor tensor(ArrayRef values, const TensorOptions& options) { \ + if (options.device().type() != c10::DeviceType::CPU) { \ + return at::detail::tensor_backend(values, options); \ + } else { \ + return at::detail::tensor_cpu(values, options); \ + } \ + } +AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR) +#undef TENSOR + +#define TENSOR(T, _1) \ + Tensor tensor(ArrayRef values, const TensorOptions& options) { \ + if (options.device().type() != c10::DeviceType::CPU) { \ + return at::detail::tensor_complex_backend(values, options); \ + } else { \ + return at::detail::tensor_complex_cpu(values, options); \ + } \ + } +AT_FORALL_COMPLEX_TYPES(TENSOR) +#undef TENSOR +} // namespace at diff --git a/aten/src/ATen/Utils.h b/aten/src/ATen/Utils.h index 14e8fa49c1b1..e814d52efe95 100644 --- a/aten/src/ATen/Utils.h +++ b/aten/src/ATen/Utils.h @@ -138,6 +138,23 @@ namespace detail { CAFFE2_API Tensor empty_cpu(IntArrayRef size, c10::optional dtype_opt, c10::optional layout_opt, c10::optional device_opt, c10::optional pin_memory_opt, c10::optional memory_format_opt); + +template +CAFFE2_API +Tensor tensor_cpu(ArrayRef values, const TensorOptions& options); + +template +CAFFE2_API +Tensor tensor_backend(ArrayRef values, const TensorOptions& options); + +template +CAFFE2_API +Tensor tensor_complex_cpu(ArrayRef values, const TensorOptions& options); + +template +CAFFE2_API +Tensor tensor_complex_backend(ArrayRef values, const TensorOptions& options); } // namespace detail + } // at diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp index 42d98336e5cd..10c2d6ee999c 100644 --- a/aten/src/ATen/native/TensorFactories.cpp +++ b/aten/src/ATen/native/TensorFactories.cpp @@ -1046,58 +1046,24 @@ Tensor vander(const Tensor& x, c10::optional N, bool increasing) { template Tensor tensor_cpu(ArrayRef values, const TensorOptions& options) { - auto result = at::empty(values.size(), options); - AT_ASSERT(result.is_contiguous()); - AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "tensor_cpu", [&] { - std::copy(values.begin(), values.end(), result.template data_ptr()); - }); - return result; + return at::detail::tensor_cpu(values, options); } template Tensor tensor_backend(ArrayRef values, const TensorOptions& options) { - auto cpu_tensor = tensor_cpu(values, options.device(DeviceType::CPU)); - return cpu_tensor.to(options.device()); + return at::detail::tensor_backend(values, options); } template Tensor tensor_complex_cpu(ArrayRef values, const TensorOptions& options) { - auto result = at::empty(values.size(), options); - AT_ASSERT(result.is_contiguous()); - AT_DISPATCH_COMPLEX_TYPES(result.scalar_type(), "tensor_cpu", [&] { - std::copy(values.begin(), values.end(), result.template data_ptr()); - }); - return result; + return at::detail::tensor_complex_cpu(values, options); } template Tensor tensor_complex_backend(ArrayRef values, const TensorOptions& options) { - auto cpu_tensor = tensor_complex_cpu(values, options.device(DeviceType::CPU)); - return cpu_tensor.to(options.device()); + return at::detail::tensor_complex_backend(values, options); } -#define TENSOR(T, _1) \ - Tensor tensor(ArrayRef values, const TensorOptions& options) { \ - if (options.device().type() != c10::DeviceType::CPU) { \ - return tensor_backend(values, options); \ - } else { \ - return tensor_cpu(values, options); \ - } \ - } -AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR) -#undef TENSOR - -#define TENSOR(T, _1) \ - Tensor tensor(ArrayRef values, const TensorOptions& options) { \ - if (options.device().type() != c10::DeviceType::CPU) { \ - return tensor_complex_backend(values, options); \ - } else { \ - return tensor_complex_cpu(values, options); \ - } \ - } -AT_FORALL_COMPLEX_TYPES(TENSOR) -#undef TENSOR - Tensor from_file(std::string filename, c10::optional shared, c10::optional size, const TensorOptions& options) { TORCH_CHECK(!options.pinned_memory(), "tensors constructed from a file cannot be pinned"); int64_t my_size = size.value_or(0); diff --git a/aten/src/ATen/templates/Functions.cpp b/aten/src/ATen/templates/Functions.cpp index 589121af07ef..81e2a9f6d406 100644 --- a/aten/src/ATen/templates/Functions.cpp +++ b/aten/src/ATen/templates/Functions.cpp @@ -7,8 +7,6 @@ namespace at { -using native::tensor; - ${function_definitions} } diff --git a/aten/src/ATen/templates/Functions.h b/aten/src/ATen/templates/Functions.h index 81e46642ad58..5ec90e7e3d61 100644 --- a/aten/src/ATen/templates/Functions.h +++ b/aten/src/ATen/templates/Functions.h @@ -19,8 +19,6 @@ namespace at { -using native::tensor; - ${function_declarations} // Special C++ only overloads for std()-like functions (See gh-40287) diff --git a/aten/src/ATen/templates/NativeFunctions.h b/aten/src/ATen/templates/NativeFunctions.h index 7d0dab41a640..0244efcb3a6b 100644 --- a/aten/src/ATen/templates/NativeFunctions.h +++ b/aten/src/ATen/templates/NativeFunctions.h @@ -3,9 +3,9 @@ // ${generated_comment} #include +#include #include #include -#include #include #include @@ -23,31 +23,31 @@ struct Type; } // namespace at namespace at { -namespace native { - -// These functions are defined in native/TensorFactories.cpp. +// These functions are defined in ATen/Utils.cpp. #define TENSOR(T, S) \ CAFFE2_API Tensor tensor(ArrayRef values, const TensorOptions& options); \ inline Tensor tensor( \ std::initializer_list values, const TensorOptions& options) { \ - return native::tensor(ArrayRef(values), options); \ + return at::tensor(ArrayRef(values), options); \ } \ inline Tensor tensor(T value, const TensorOptions& options) { \ - return native::tensor(ArrayRef(value), options); \ + return at::tensor(ArrayRef(value), options); \ } \ inline Tensor tensor(ArrayRef values) { \ - return native::tensor(std::move(values), at::dtype(k##S)); \ + return at::tensor(std::move(values), at::dtype(k##S)); \ } \ inline Tensor tensor(std::initializer_list values) { \ - return native::tensor(ArrayRef(values)); \ + return at::tensor(ArrayRef(values)); \ } \ inline Tensor tensor(T value) { \ - return native::tensor(ArrayRef(value)); \ + return at::tensor(ArrayRef(value)); \ } AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR) AT_FORALL_COMPLEX_TYPES(TENSOR) #undef TENSOR +namespace native { + ${native_function_declarations} } // namespace native