GCC Code Coverage Report

ernie/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
1 / 12
8.3%
Branches:
0 / 24
0.0%
Functions:
-
-%
device/cuda.h
Line Branch Exec Source
1 // SPDX-License-Identifier: BSD-3-Clause
2 /// @file cuda.h
3 /// @brief Zero-copy interop with CUDA, and through it PyTorch, Warp and any DLPack consumer:
4 /// Ernie's shared memory mapped into CUDA in place, and CUDA streams ordered with Ernie's queues
5 /// through shared timelines, so neither the host nor a copy sits between them. The CUDA driver
6 /// is loaded at run time: Ernie builds without the CUDA toolkit and runs without a CUDA driver.
7 #pragma once
8
9 #include "ernie/device/column.h"
10 #include "ernie/util/library.h"
11 #include "ernie/util/tensor.h"
12
13 #include <cstdint>
14 #include <memory>
15 #include <mutex>
16 #include <optional>
17 #include <unordered_map>
18
19 namespace ernie {
20
21 /// @brief The CUDA driver API's handles, declared here so no CUDA header is needed.
22 using CUresult = int;
23 using CUdevice = int;
24 using CUdeviceptr = std::uint64_t;
25 using CUcontext = struct CUctx_st*;
26 using CUstream = struct CUstream_st*;
27 using CUexternalMemory = struct CUextMemory_st*;
28 using CUexternalSemaphore = struct CUextSemaphore_st*;
29 using CUmemGenericAllocationHandle = unsigned long long;
30
31 /// @brief The CUDA driver API functions Ernie uses, by their ABI names.
32 6 struct CudaDriver {
33 SharedLibrary library;
34 CUresult (*init)(unsigned int) = nullptr;
35 CUresult (*deviceGetCount)(int*) = nullptr;
36 CUresult (*deviceGet)(CUdevice*, int) = nullptr;
37 CUresult (*deviceGetUuid)(unsigned char*, CUdevice) = nullptr;
38 CUresult (*primaryCtxRetain)(CUcontext*, CUdevice) = nullptr;
39 CUresult (*primaryCtxRelease)(CUdevice) = nullptr;
40 CUresult (*ctxPushCurrent)(CUcontext) = nullptr;
41 CUresult (*ctxPopCurrent)(CUcontext*) = nullptr;
42 CUresult (*ctxSynchronize)() = nullptr;
43 CUresult (*importExternalMemory)(CUexternalMemory*, const void*) = nullptr;
44 CUresult (*externalMemoryGetMappedBuffer)(CUdeviceptr*, CUexternalMemory,
45 const void*) = nullptr;
46 CUresult (*destroyExternalMemory)(CUexternalMemory) = nullptr;
47 CUresult (*importExternalSemaphore)(CUexternalSemaphore*, const void*) = nullptr;
48 CUresult (*signalExternalSemaphoresAsync)(const CUexternalSemaphore*, const void*, unsigned int,
49 CUstream) = nullptr;
50 CUresult (*waitExternalSemaphoresAsync)(const CUexternalSemaphore*, const void*, unsigned int,
51 CUstream) = nullptr;
52 CUresult (*destroyExternalSemaphore)(CUexternalSemaphore) = nullptr;
53 CUresult (*memFree)(CUdeviceptr) = nullptr;
54 CUresult (*memAlloc)(CUdeviceptr*, std::size_t) = nullptr;
55 CUresult (*memcpyDtoDAsync)(CUdeviceptr, CUdeviceptr, std::size_t, CUstream) = nullptr;
56 CUresult (*memcpyDtoH)(void*, CUdeviceptr, std::size_t) = nullptr;
57 CUresult (*memsetD32Async)(CUdeviceptr, unsigned int, std::size_t, CUstream) = nullptr;
58 CUresult (*streamSynchronize)(CUstream) = nullptr;
59 CUresult (*streamCreate)(CUstream*, unsigned int) = nullptr;
60 CUresult (*streamDestroy)(CUstream) = nullptr;
61 // Virtual memory management, whose allocations can be exported.
62 CUresult (*memCreate)(CUmemGenericAllocationHandle*, std::size_t, const void*,
63 unsigned long long) = nullptr;
64 CUresult (*memRelease)(CUmemGenericAllocationHandle) = nullptr;
65 CUresult (*memAddressReserve)(CUdeviceptr*, std::size_t, std::size_t, CUdeviceptr,
66 unsigned long long) = nullptr;
67 CUresult (*memAddressFree)(CUdeviceptr, std::size_t) = nullptr;
68 CUresult (*memMap)(CUdeviceptr, std::size_t, std::size_t, CUmemGenericAllocationHandle,
69 unsigned long long) = nullptr;
70 CUresult (*memUnmap)(CUdeviceptr, std::size_t) = nullptr;
71 CUresult (*memSetAccess)(CUdeviceptr, std::size_t, const void*, std::size_t) = nullptr;
72 CUresult (*memGetAllocationGranularity)(std::size_t*, const void*, int) = nullptr;
73 CUresult (*memRetainAllocationHandle)(CUmemGenericAllocationHandle*, void*) = nullptr;
74 CUresult (*memGetAllocationPropertiesFromHandle)(void*, CUmemGenericAllocationHandle) = nullptr;
75 CUresult (*memExportToShareableHandle)(void*, CUmemGenericAllocationHandle, int,
76 unsigned long long) = nullptr;
77 CUresult (*memGetAddressRange)(CUdeviceptr*, std::size_t*, CUdeviceptr) = nullptr;
78 CUresult (*getErrorString)(CUresult, const char**) = nullptr;
79
80 /// @brief Exportable device memory, as `AllocateShareable` makes it.
81 struct Shareable {
82 CUdeviceptr pointer = 0;
83 CUmemGenericAllocationHandle handle = 0;
84 std::size_t size = 0;
85 };
86
87 /// @return The driver, loaded and initialized once per process; null when there is none.
88 static const CudaDriver* Get();
89 /// @brief Allocates at least `bytes` on device `ordinal` through virtual memory management,
90 /// with handles another API imports in place. Needs a current context.
91 /// @return The allocation, or none on failure.
92 std::optional<Shareable> AllocateShareable(int ordinal, std::size_t bytes) const;
93 void FreeShareable(const Shareable& allocation) const;
94 /// @throws DeviceError when `result` is an error.
95 void Check(CUresult result, std::string_view what) const;
96 };
97
98 /// @brief Exportable allocations by address, as `CudaDriver::AllocateShareable` makes them, for
99 /// whoever hands out addresses and takes them back: Ernie's interop, and PyTorch's allocator hook.
100 ✗ class ShareableAllocations {
101 public:
102
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ explicit ShareableAllocations(const CudaDriver& driver) : driver_(driver) {}
103 ShareableAllocations(const ShareableAllocations&) = delete;
104 ShareableAllocations& operator=(const ShareableAllocations&) = delete;
105
106 /// @return The allocation's address, or 0 on failure. Needs a current context.
107 CUdeviceptr Allocate(int ordinal, std::size_t bytes);
108 /// @brief Frees what `Allocate` returned; other addresses are left alone.
109 void Free(CUdeviceptr pointer);
110 /// @brief Frees every allocation left.
111 void FreeAll();
112
113 private:
114 const CudaDriver& driver_;
115 std::mutex mutex_;
116 std::unordered_map<CUdeviceptr, CudaDriver::Shareable> live_;
117 };
118
119 /// @brief The CUDA device paired with a context's device by UUID. Columns in
120 /// `MemoryUse::Shared` arenas are viewed as CUDA tensors in place; each arena is mapped once.
121 /// Streams are CUDA streams (`torch.cuda.current_stream().cuda_stream`, `wp.get_stream()`), 0
122 /// for the default one.
123 class CudaInterop {
124 public:
125 /// @return The interop, or null when there is no CUDA driver, no CUDA device with the
126 /// context's UUID, or the device shares no memory or timelines.
127 static std::unique_ptr<CudaInterop> Create(DeviceContext& context);
128 ~CudaInterop();
129 CudaInterop(const CudaInterop&) = delete;
130 CudaInterop& operator=(const CudaInterop&) = delete;
131
132 /// @return The CUDA device ordinal, as `cuda:N` names it.
133 ✗ int Ordinal() const {
134
0/12
✗ Branch 1 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 ✗ return ordinal_;
135 }
136 ✗ const CudaDriver& Driver() const {
137
0/10
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 ✗ return driver_;
138 }
139
140 /// @brief A CUDA tensor over the column's rows, in place, shaped (rows, element...). It is
141 /// valid while the column keeps its range; order access with `Wait` and `Signal`.
142 /// @throws DeviceError unless the column's arenas are `MemoryUse::Shared`.
143 Tensor View(const DeviceColumn& column);
144 /// @brief Makes work enqueued later on `stream` wait for `ticket`, without the host.
145 void Wait(CUstream stream, const Ticket& ticket);
146 /// @return A ticket reached once the work enqueued on `stream` so far is done, for Ernie's
147 /// queues to wait on.
148 Ticket Signal(CUstream stream);
149 /// @brief Makes `tensor`'s values the column's rows, after the work on `stream` so far: in
150 /// place, with no copy, when `tensor` is the column's own view; otherwise by one copy on the
151 /// device, counted in `DeviceContext::Transfers`.
152 /// @return The ticket Ernie's queues wait on before reading the rows.
153 /// @throws std::invalid_argument unless `tensor` is a contiguous CUDA tensor on this device
154 /// with the column's byte size.
155 Ticket Assign(const DeviceColumn& column, const Tensor& tensor, CUstream stream);
156 /// @brief A device column over a CUDA tensor's own memory, in place, when that memory is
157 /// exportable: allocated through CUDA's virtual memory management with shareable handles, as
158 /// `Allocate` and PyTorch's expandable segments allocate. Ordinary `cudaMalloc` memory is
159 /// not; assign such tensors to an Ernie column instead.
160 /// @return The column, shaped (rows, element...) as the tensor is and holding it alive; null
161 /// when its memory cannot be shared.
162 std::unique_ptr<DeviceColumn> Wrap(const Tensor& tensor, std::string_view name);
163 /// @brief Allocates `bytes` of exportable device memory, which `Wrap` shares in place.
164 /// @return Its address, or 0 on failure.
165 CUdeviceptr Allocate(std::size_t bytes);
166 /// @brief Frees what `Allocate` returned.
167 void Free(CUdeviceptr pointer);
168
169 /// @brief Makes the device's primary context current on this thread for a scope, for calls
170 /// to the driver directly.
171 class Current {
172 public:
173 explicit Current(const CudaInterop& interop);
174 ~Current();
175 Current(const Current&) = delete;
176 Current& operator=(const Current&) = delete;
177
178 private:
179 const CudaDriver& driver_;
180 };
181
182 private:
183 struct Mapping {
184 CUexternalMemory memory = nullptr;
185 CUdeviceptr pointer = 0;
186 };
187
188 /// A stream's own timeline, which only that stream signals, so its values only increase.
189 ✗ struct StreamTimeline {
190 std::unique_ptr<Timeline> timeline;
191 CUexternalSemaphore semaphore = nullptr;
192 std::uint64_t signalled = 0;
193 };
194
195 CudaInterop(DeviceContext& context, const CudaDriver& driver, CUdevice device, int ordinal);
196 /// @throws std::invalid_argument unless `tensor` is a contiguous tensor on this device.
197 void CheckOnDevice(const Tensor& tensor, std::string_view what) const;
198 CUdeviceptr Map(const Buffer& buffer);
199 /// Imports one of the context's timelines into CUDA, once. Takes `mutex_`'s lock as held.
200 CUexternalSemaphore Import(const Timeline& timeline);
201
202 DeviceContext& context_;
203 const CudaDriver& driver_;
204 CUdevice device_;
205 int ordinal_;
206 CUcontext cuda_ = nullptr;
207 std::mutex mutex_;
208 /// Mapped buffers by `Buffer::Id`, which no later buffer reuses.
209 std::unordered_map<std::uint64_t, Mapping> mappings_;
210 std::unordered_map<const Timeline*, CUexternalSemaphore> semaphores_;
211 std::unordered_map<CUstream, StreamTimeline> streams_;
212 ShareableAllocations allocations_;
213 };
214
215 } // namespace ernie
216