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