1 | #include <ATen/cuda/PinnedMemoryAllocator.h> |
2 | #include <ATen/Context.h> |
3 | #include <ATen/Config.h> |
4 | #include <ATen/TensorUtils.h> |
5 | #include <c10/core/Storage.h> |
6 | #include <ATen/ATen.h> |
7 | #include <ATen/CPUFunctions.h> |
8 | |
9 | namespace at { |
10 | |
11 | namespace native { |
12 | |
13 | bool is_pinned_cuda(const Tensor& self, c10::optional<Device> device) { |
14 | TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!device.has_value() || device->is_cuda()); |
15 | // TODO: unhook this |
16 | return detail::getCUDAHooks().isPinnedPtr(self.storage().data()); |
17 | } |
18 | |
19 | Tensor _pin_memory_cuda(const Tensor& self, c10::optional<Device> device) { |
20 | TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!device.has_value() || device->is_cuda()); |
21 | auto* allocator = at::cuda::getPinnedMemoryAllocator(); |
22 | auto storage = Storage( |
23 | Storage::use_byte_size_t(), |
24 | detail::computeStorageNbytes( |
25 | self.sizes(), self.strides(), self.dtype().itemsize()), |
26 | allocator, |
27 | /*resizable=*/false); |
28 | auto tensor = at::cpu::empty({0}, self.options()).set_(storage, 0, self.sizes(), self.strides()); |
29 | tensor.copy_(self); |
30 | return tensor; |
31 | } |
32 | |
33 | |
34 | } // namespace native |
35 | } // namespace at |
36 | |