| 1 | #include <torch/csrc/stable/tensor.h> |
| 2 | #include <torch/csrc/stable/ops.h> |
| 3 | #include <torch/csrc/stable/accelerator.h> |
| 4 | #include <torch/headeronly/core/ScalarType.h> |
| 5 | #include <torch/csrc/stable/device.h> |
| 6 | #include <torch/csrc/stable/c/shim.h> |
| 7 | #include <torch/headeronly/version.h> |
| 8 | #include <cuda_runtime.h> |
| 9 | |
| 10 | #include <array> |
| 11 | #include <optional> |
| 12 | |
| 13 | // This function assumes that `cpu_tensor` is a CPU tensor, |
| 14 | // and that UVA (Unified Virtual Addressing) is enabled. |
| 15 | torch::stable::Tensor get_cuda_view_from_cpu_tensor( |
| 16 | torch::stable::Tensor& cpu_tensor) { |
| 17 | STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU"); |
| 18 | |
| 19 | const auto dtype = cpu_tensor.scalar_type(); |
| 20 | const auto layout = cpu_tensor.layout(); |
| 21 | const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA); |
| 22 | |
| 23 | // handle empty tensor |
| 24 | if (cpu_tensor.numel() == 0) { |
| 25 | return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev); |
| 26 | } |
| 27 | |
| 28 | std::array<StableIValue, 2> is_pinned_stack{ |
| 29 | torch::stable::detail::from(cpu_tensor), |
| 30 | torch::stable::detail::from(std::nullopt)}; |
| 31 | TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( |
| 32 | "aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION)); |
| 33 | if (torch::stable::detail::to<bool>(is_pinned_stack[0])) { |
| 34 | // If CPU tensor is pinned, directly get the device pointer. |
| 35 | void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr()); |
| 36 | void* device_ptr = nullptr; |
| 37 | cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); |
| 38 | STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ", |
| 39 | cudaGetErrorString(err)); |
| 40 | |
| 41 | return torch::stable::from_blob( |
| 42 | device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype, |
| 43 | [base = cpu_tensor](void*) {}); // keep cpu tensor alive |
| 44 | } |
| 45 | |
| 46 | // If CPU tensor is not pinned, allocate a new pinned memory buffer. |
| 47 | torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor); |
| 48 | size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size(); |
| 49 | |
| 50 | void* host_ptr = nullptr; |
| 51 | cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped); |
| 52 | if (err != cudaSuccess) { |
| 53 | STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err)); |
| 54 | } |
| 55 | |
| 56 | err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes, |
| 57 | cudaMemcpyDefault); |
| 58 | if (err != cudaSuccess) { |
| 59 | cudaFreeHost(host_ptr); |
| 60 | STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err)); |
| 61 | } |
| 62 | |
| 63 | void* device_ptr = nullptr; |
| 64 | err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0); |
| 65 | if (err != cudaSuccess) { |
| 66 | cudaFreeHost(host_ptr); |
| 67 | STD_TORCH_CHECK( |
| 68 | false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err)); |
| 69 | } |
| 70 | |
| 71 | auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); }; |
| 72 | |
| 73 | return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(), |
| 74 | contiguous_cpu.strides(), cuda_dev, |
| 75 | contiguous_cpu.scalar_type(), deleter); |
| 76 | } |