xref: /aosp_15_r20/external/pytorch/aten/src/ATen/Utils.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <ATen/Context.h>
2 #include <ATen/Dispatch.h>
3 #include <ATen/Functions.h>
4 #include <ATen/Utils.h>
5 #include <c10/util/accumulate.h>
6 
7 
8 #include <cstdlib>
9 
10 namespace at {
11 
_crash_if_asan(int arg)12 int _crash_if_asan(int arg) {
13   // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
14   volatile char x[3];
15   x[arg] = 0;
16   return x[0];
17 }
18 
19 namespace detail {
20 
21 template <typename T>
tensor_cpu(ArrayRef<T> values,const TensorOptions & options)22 Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options) {
23   auto result = at::empty(values.size(), options);
24   AT_ASSERT(result.is_contiguous());
25   AT_DISPATCH_ALL_TYPES_AND_COMPLEX(result.scalar_type(), "tensor_cpu", [&] {
26     std::copy(
27         values.begin(), values.end(), result.template data_ptr<scalar_t>());
28   });
29   return result;
30 }
31 
32 template <typename T>
tensor_backend(ArrayRef<T> values,const TensorOptions & options)33 Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options) {
34   auto cpu_tensor = tensor_cpu(values, options.device(DeviceType::CPU));
35   return cpu_tensor.to(options.device());
36 }
37 
38 template <typename T>
tensor_complex_cpu(ArrayRef<T> values,const TensorOptions & options)39 Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options) {
40   auto result = at::empty(values.size(), options);
41   AT_ASSERT(result.is_contiguous());
42   AT_DISPATCH_COMPLEX_TYPES(result.scalar_type(), "tensor_cpu", [&] {
43     std::copy(
44         values.begin(), values.end(), result.template data_ptr<scalar_t>());
45   });
46   return result;
47 }
48 
49 template <typename T>
tensor_complex_backend(ArrayRef<T> values,const TensorOptions & options)50 Tensor tensor_complex_backend(
51     ArrayRef<T> values,
52     const TensorOptions& options) {
53   auto cpu_tensor = tensor_complex_cpu(values, options.device(DeviceType::CPU));
54   return cpu_tensor.to(options.device());
55 }
56 } // namespace detail
57 
58 #define TENSOR(T, _1)                                               \
59   Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
60     if (options.device().type() != c10::DeviceType::CPU) {          \
61       return at::detail::tensor_backend(values, options);           \
62     } else {                                                        \
63       return at::detail::tensor_cpu(values, options);               \
64     }                                                               \
65   }
66 AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
67 #undef TENSOR
68 
69 #define TENSOR(T, _1)                                               \
70   Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
71     if (options.device().type() != c10::DeviceType::CPU) {          \
72       return at::detail::tensor_complex_backend(values, options);   \
73     } else {                                                        \
74       return at::detail::tensor_complex_cpu(values, options);       \
75     }                                                               \
76   }
77 AT_FORALL_COMPLEX_TYPES(TENSOR)
78 #undef TENSOR
79 } // namespace at
80