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
9namespace at {
10
11namespace native {
12
13bool 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
19Tensor _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