1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
2 #include <ATen/core/Tensor.h>
3 #include <ATen/MemoryOverlap.h>
4 #include <ATen/Context.h>
5 #include <c10/core/Storage.h>
6 #include <ATen/EmptyTensor.h>
7
8 #ifndef AT_PER_OPERATOR_HEADERS
9 #include <ATen/Functions.h>
10 #include <ATen/NativeFunctions.h>
11 #include <ATen/CPUFunctions.h>
12 #else
13 #include <ATen/ops/_debug_has_internal_overlap_native.h>
14 #include <ATen/ops/_pin_memory.h>
15 #include <ATen/ops/is_pinned_native.h>
16 #include <ATen/ops/pin_memory_native.h>
17 #include <ATen/ops/_pin_memory_native.h>
18 #include <ATen/ops/empty_cpu_dispatch.h>
19 #endif
20
21 namespace at::native {
22
23 // Exposes at::has_internal_overlap as an operator for testing purposes
_debug_has_internal_overlap(const Tensor & self)24 int64_t _debug_has_internal_overlap(const Tensor& self) {
25 return static_cast<int64_t>(at::has_internal_overlap(self));
26 }
27
is_pinned(const Tensor & self,std::optional<c10::Device> device)28 bool is_pinned(const Tensor& self, std::optional<c10::Device> device) {
29 std::optional<c10::DeviceType> opt_device_type;
30 if (device.has_value()) {
31 TORCH_WARN_DEPRECATION(
32 "The argument 'device' of Tensor.is_pinned() ",
33 "is deprecated. Please do not pass this argument.")
34 opt_device_type = device.value().type();
35 }
36 // Only CPU tensors can be pinned
37 if (!self.is_cpu()) {
38 return false;
39 }
40 // Use getAcceleratorHooksInterface to make is_pinned device-agnostic
41 return at::globalContext().isPinnedPtr(self.storage().data(), opt_device_type);
42 }
43
pin_memory(const Tensor & self,std::optional<c10::Device> device)44 Tensor pin_memory(const Tensor& self, std::optional<c10::Device> device) {
45 if (device.has_value()) {
46 TORCH_WARN_DEPRECATION(
47 "The argument 'device' of Tensor.pin_memory() ",
48 "is deprecated. Please do not pass this argument.")
49 }
50 // Kind of mad that I have to do two dynamic dispatches here, pretty
51 // annoying
52 if (self.is_pinned(device)) {
53 return self;
54 }
55 return at::_pin_memory(self, device);
56 }
57
_pin_memory(const Tensor & self,std::optional<c10::Device> device)58 Tensor _pin_memory(const Tensor& self, std::optional<c10::Device> device) {
59 TORCH_CHECK(self.device().is_cpu(), "cannot pin '", self.toString(), "' only dense CPU tensors can be pinned");
60 // Use getAcceleratorHooksInterface to make pin_memory device-agnostic
61 auto* allocator = device.has_value()?
62 at::globalContext().getPinnedMemoryAllocator(device.value().type()):
63 at::globalContext().getPinnedMemoryAllocator();
64 auto storage = Storage(
65 Storage::use_byte_size_t(),
66 detail::computeStorageNbytes(
67 self.sizes(), self.strides(), self.dtype().itemsize()),
68 allocator,
69 /*resizable=*/false);
70 auto tensor = at::cpu::empty({0}, self.options()).set_(storage, 0, self.sizes(), self.strides());
71 tensor.copy_(self);
72 return tensor;
73 }
74
75 } // namespace at::native
76