1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/csrc/utils/pybind.h>
#if !defined(USE_ROCM)
#include <cuda_profiler_api.h>
#else
#include <hip/hip_runtime_api.h>
#endif
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
namespace torch::cuda::shared {
#ifdef USE_ROCM
namespace {
hipError_t hipReturnSuccess() {
return hipSuccess;
}
} // namespace
#endif
void initCudartBindings(PyObject* module) {
auto m = py::handle(module).cast<py::module>();
auto cudart = m.def_submodule("_cudart", "libcudart.so bindings");
// By splitting the names of these objects into two literals we prevent the
// HIP rewrite rules from changing these names when building with HIP.
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION < 12000
// cudaOutputMode_t is used in cudaProfilerInitialize only. The latter is gone
// in CUDA 12.
py::enum_<cudaOutputMode_t>(
cudart,
"cuda"
"OutputMode")
.value("KeyValuePair", cudaKeyValuePair)
.value("CSV", cudaCSV);
#endif
py::enum_<cudaError_t>(
cudart,
"cuda"
"Error")
.value("success", cudaSuccess);
cudart.def(
"cuda"
"GetErrorString",
cudaGetErrorString);
cudart.def(
"cuda"
"ProfilerStart",
#ifdef USE_ROCM
hipReturnSuccess
#else
cudaProfilerStart
#endif
);
cudart.def(
"cuda"
"ProfilerStop",
#ifdef USE_ROCM
hipReturnSuccess
#else
cudaProfilerStop
#endif
);
cudart.def(
"cuda"
"HostRegister",
[](uintptr_t ptr, size_t size, unsigned int flags) -> cudaError_t {
py::gil_scoped_release no_gil;
return C10_CUDA_ERROR_HANDLED(
// NOLINTNEXTLINE(performance-no-int-to-ptr)
cudaHostRegister((void*)ptr, size, flags));
});
cudart.def(
"cuda"
"HostUnregister",
[](uintptr_t ptr) -> cudaError_t {
py::gil_scoped_release no_gil;
// NOLINTNEXTLINE(performance-no-int-to-ptr)
return C10_CUDA_ERROR_HANDLED(cudaHostUnregister((void*)ptr));
});
cudart.def(
"cuda"
"StreamCreate",
[](uintptr_t ptr) -> cudaError_t {
py::gil_scoped_release no_gil;
// NOLINTNEXTLINE(performance-no-int-to-ptr)
return C10_CUDA_ERROR_HANDLED(cudaStreamCreate((cudaStream_t*)ptr));
});
cudart.def(
"cuda"
"StreamDestroy",
[](uintptr_t ptr) -> cudaError_t {
py::gil_scoped_release no_gil;
// NOLINTNEXTLINE(performance-no-int-to-ptr)
return C10_CUDA_ERROR_HANDLED(cudaStreamDestroy((cudaStream_t)ptr));
});
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION < 12000
// cudaProfilerInitialize is no longer needed after CUDA 12:
// https://forums.developer.nvidia.com/t/cudaprofilerinitialize-is-deprecated-alternative/200776/3
cudart.def(
"cuda"
"ProfilerInitialize",
cudaProfilerInitialize,
py::call_guard<py::gil_scoped_release>());
#endif
cudart.def(
"cuda"
"MemGetInfo",
[](c10::DeviceIndex device) -> std::pair<size_t, size_t> {
c10::cuda::CUDAGuard guard(device);
size_t device_free = 0;
size_t device_total = 0;
py::gil_scoped_release no_gil;
C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total));
return {device_free, device_total};
});
}
} // namespace torch::cuda::shared
|
