device/cuda.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file cuda.cpp | ||
| 3 | /// @brief Loading the CUDA driver, pairing devices by UUID, and importing Ernie's memory and | ||
| 4 | /// timelines into CUDA. | ||
| 5 | #include "ernie/device/cuda.h" | ||
| 6 | |||
| 7 | #include "ernie/util/on_failure.h" | ||
| 8 | #include "ernie/util/text.h" | ||
| 9 | |||
| 10 | #include <cstring> | ||
| 11 | #include <stdexcept> | ||
| 12 | |||
| 13 | namespace ernie { | ||
| 14 | namespace { | ||
| 15 | |||
| 16 | // The driver API's descriptor layouts, as cuda.h's _v1 structs declare them. | ||
| 17 | struct MemoryHandleDesc { | ||
| 18 | int type; | ||
| 19 | union { | ||
| 20 | int fd; | ||
| 21 | struct { | ||
| 22 | void* handle; | ||
| 23 | const void* name; | ||
| 24 | } win32; | ||
| 25 | const void* nvSciBufObject; | ||
| 26 | } handle; | ||
| 27 | unsigned long long size; | ||
| 28 | unsigned int flags; | ||
| 29 | unsigned int reserved[16]; | ||
| 30 | }; | ||
| 31 | struct BufferDesc { | ||
| 32 | unsigned long long offset; | ||
| 33 | unsigned long long size; | ||
| 34 | unsigned int flags; | ||
| 35 | unsigned int reserved[16]; | ||
| 36 | }; | ||
| 37 | struct SemaphoreHandleDesc { | ||
| 38 | int type; | ||
| 39 | union { | ||
| 40 | int fd; | ||
| 41 | struct { | ||
| 42 | void* handle; | ||
| 43 | const void* name; | ||
| 44 | } win32; | ||
| 45 | const void* nvSciSyncObj; | ||
| 46 | } handle; | ||
| 47 | unsigned int flags; | ||
| 48 | unsigned int reserved[16]; | ||
| 49 | }; | ||
| 50 | struct SemaphoreSignalParams { | ||
| 51 | struct { | ||
| 52 | struct { | ||
| 53 | unsigned long long value; | ||
| 54 | } fence; | ||
| 55 | union { | ||
| 56 | void* fence; | ||
| 57 | unsigned long long reserved; | ||
| 58 | } nvSciSync; | ||
| 59 | struct { | ||
| 60 | unsigned long long key; | ||
| 61 | } keyedMutex; | ||
| 62 | unsigned int reserved[12]; | ||
| 63 | } params; | ||
| 64 | unsigned int flags; | ||
| 65 | unsigned int reserved[16]; | ||
| 66 | }; | ||
| 67 | struct SemaphoreWaitParams { | ||
| 68 | struct { | ||
| 69 | struct { | ||
| 70 | unsigned long long value; | ||
| 71 | } fence; | ||
| 72 | union { | ||
| 73 | void* fence; | ||
| 74 | unsigned long long reserved; | ||
| 75 | } nvSciSync; | ||
| 76 | struct { | ||
| 77 | unsigned long long key; | ||
| 78 | unsigned int timeoutMs; | ||
| 79 | } keyedMutex; | ||
| 80 | unsigned int reserved[10]; | ||
| 81 | } params; | ||
| 82 | unsigned int flags; | ||
| 83 | unsigned int reserved[16]; | ||
| 84 | }; | ||
| 85 | |||
| 86 | // Virtual memory management's layouts, as cuda.h's _v1 structs declare them. | ||
| 87 | struct MemoryLocation { | ||
| 88 | int type; | ||
| 89 | int id; | ||
| 90 | }; | ||
| 91 | struct AllocationProperties { | ||
| 92 | int type; | ||
| 93 | int requestedHandleTypes; | ||
| 94 | MemoryLocation location; | ||
| 95 | void* win32HandleMetaData; | ||
| 96 | struct { | ||
| 97 | unsigned char compressionType; | ||
| 98 | unsigned char gpuDirectRDMACapable; | ||
| 99 | unsigned short usage; | ||
| 100 | unsigned char reserved[4]; | ||
| 101 | } allocFlags; | ||
| 102 | }; | ||
| 103 | struct AccessDescription { | ||
| 104 | MemoryLocation location; | ||
| 105 | int flags; | ||
| 106 | }; | ||
| 107 | constexpr int kPinned = 1; // CU_MEM_ALLOCATION_TYPE_PINNED | ||
| 108 | constexpr int kOnDevice = 1; // CU_MEM_LOCATION_TYPE_DEVICE | ||
| 109 | constexpr int kReadWrite = 3; // CU_MEM_ACCESS_FLAGS_PROT_READWRITE | ||
| 110 | constexpr int kMinimumGranule = 0; // CU_MEM_ALLOC_GRANULARITY_MINIMUM | ||
| 111 | |||
| 112 | #if defined(_WIN32) | ||
| 113 | constexpr int kShareableHandle = 2; // CU_MEM_HANDLE_TYPE_WIN32 | ||
| 114 | constexpr int kMemoryHandleType = 2; // CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32 | ||
| 115 | constexpr int kSemaphoreHandleType = | ||
| 116 | 10; // CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_TIMELINE_SEMAPHORE_WIN32 | ||
| 117 | #else | ||
| 118 | constexpr int kShareableHandle = 1; // CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR | ||
| 119 | constexpr int kMemoryHandleType = 1; // CU_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD | ||
| 120 | constexpr int kSemaphoreHandleType = 9; // CU_EXTERNAL_SEMAPHORE_HANDLE_TYPE_TIMELINE_SEMAPHORE_FD | ||
| 121 | #endif | ||
| 122 | constexpr unsigned int kDedicated = 1; // CUDA_EXTERNAL_MEMORY_DEDICATED | ||
| 123 | |||
| 124 | /// Fills a descriptor's handle from an exported one, which CUDA takes where the platform does. | ||
| 125 | ✗ | template <typename Desc> void SetHandle(Desc& desc, const ExternalHandle& handle) { | |
| 126 | #if defined(_WIN32) | ||
| 127 | desc.handle.win32.handle = reinterpret_cast<void*>(handle.Get()); | ||
| 128 | #else | ||
| 129 | ✗ | desc.handle.fd = static_cast<int>(handle.Get()); | |
| 130 | #endif | ||
| 131 | ✗ | } | |
| 132 | |||
| 133 | 6 | bool LoadDriver(CudaDriver& d) { | |
| 134 | 6 | if (!d.library.Open({"libcuda.so.1", "libcuda.so", "nvcuda.dll"})) { | |
| 135 | return false; | ||
| 136 | } | ||
| 137 | ✗ | const SharedLibrary& l = d.library; | |
| 138 | ✗ | return l.Load(d.init, "cuInit") && l.Load(d.deviceGetCount, "cuDeviceGetCount") && | |
| 139 | ✗ | l.Load(d.deviceGet, "cuDeviceGet") && l.Load(d.deviceGetUuid, "cuDeviceGetUuid_v2") && | |
| 140 | ✗ | l.Load(d.primaryCtxRetain, "cuDevicePrimaryCtxRetain") && | |
| 141 | ✗ | l.Load(d.primaryCtxRelease, "cuDevicePrimaryCtxRelease_v2") && | |
| 142 | ✗ | l.Load(d.ctxPushCurrent, "cuCtxPushCurrent_v2") && | |
| 143 | ✗ | l.Load(d.ctxPopCurrent, "cuCtxPopCurrent_v2") && | |
| 144 | ✗ | l.Load(d.ctxSynchronize, "cuCtxSynchronize") && | |
| 145 | ✗ | l.Load(d.importExternalMemory, "cuImportExternalMemory") && | |
| 146 | ✗ | l.Load(d.externalMemoryGetMappedBuffer, "cuExternalMemoryGetMappedBuffer") && | |
| 147 | ✗ | l.Load(d.destroyExternalMemory, "cuDestroyExternalMemory") && | |
| 148 | ✗ | l.Load(d.importExternalSemaphore, "cuImportExternalSemaphore") && | |
| 149 | ✗ | l.Load(d.signalExternalSemaphoresAsync, "cuSignalExternalSemaphoresAsync") && | |
| 150 | ✗ | l.Load(d.waitExternalSemaphoresAsync, "cuWaitExternalSemaphoresAsync") && | |
| 151 | ✗ | l.Load(d.destroyExternalSemaphore, "cuDestroyExternalSemaphore") && | |
| 152 | ✗ | l.Load(d.memFree, "cuMemFree_v2") && l.Load(d.memAlloc, "cuMemAlloc_v2") && | |
| 153 | ✗ | l.Load(d.memcpyDtoDAsync, "cuMemcpyDtoDAsync_v2") && | |
| 154 | ✗ | l.Load(d.memcpyDtoH, "cuMemcpyDtoH_v2") && | |
| 155 | ✗ | l.Load(d.memsetD32Async, "cuMemsetD32Async") && | |
| 156 | ✗ | l.Load(d.streamSynchronize, "cuStreamSynchronize") && | |
| 157 | ✗ | l.Load(d.streamCreate, "cuStreamCreate") && | |
| 158 | ✗ | l.Load(d.streamDestroy, "cuStreamDestroy_v2") && l.Load(d.memCreate, "cuMemCreate") && | |
| 159 | ✗ | l.Load(d.memRelease, "cuMemRelease") && | |
| 160 | ✗ | l.Load(d.memAddressReserve, "cuMemAddressReserve") && | |
| 161 | ✗ | l.Load(d.memAddressFree, "cuMemAddressFree") && l.Load(d.memMap, "cuMemMap") && | |
| 162 | ✗ | l.Load(d.memUnmap, "cuMemUnmap") && l.Load(d.memSetAccess, "cuMemSetAccess") && | |
| 163 | ✗ | l.Load(d.memGetAllocationGranularity, "cuMemGetAllocationGranularity") && | |
| 164 | ✗ | l.Load(d.memRetainAllocationHandle, "cuMemRetainAllocationHandle") && | |
| 165 | ✗ | l.Load(d.memGetAllocationPropertiesFromHandle, | |
| 166 | ✗ | "cuMemGetAllocationPropertiesFromHandle") && | |
| 167 | ✗ | l.Load(d.memExportToShareableHandle, "cuMemExportToShareableHandle") && | |
| 168 | ✗ | l.Load(d.memGetAddressRange, "cuMemGetAddressRange_v2") && | |
| 169 | ✗ | l.Load(d.getErrorString, "cuGetErrorString") && d.init(0) == 0; | |
| 170 | } | ||
| 171 | |||
| 172 | } // namespace | ||
| 173 | |||
| 174 | 6 | const CudaDriver* CudaDriver::Get() { | |
| 175 | 6 | static const std::unique_ptr<CudaDriver> driver = [] { | |
| 176 | 6 | auto loaded = std::make_unique<CudaDriver>(); | |
| 177 | 12 | return LoadDriver(*loaded) ? std::move(loaded) : nullptr; | |
| 178 | 12 | }(); | |
| 179 | 6 | return driver.get(); | |
| 180 | } | ||
| 181 | |||
| 182 | ✗ | void CudaDriver::Check(CUresult result, std::string_view what) const { | |
| 183 | ✗ | if (result != 0) { | |
| 184 | ✗ | const char* message = nullptr; | |
| 185 | ✗ | getErrorString(result, &message); | |
| 186 | ✗ | throw DeviceError(DeviceErrorKind::Failed, | |
| 187 | ✗ | std::string(what) + ": " + (message ? message : "CUDA error")); | |
| 188 | } | ||
| 189 | ✗ | } | |
| 190 | |||
| 191 | ✗ | CudaInterop::Current::Current(const CudaInterop& interop) : driver_(interop.driver_) { | |
| 192 | ✗ | driver_.Check(driver_.ctxPushCurrent(interop.cuda_), "Making the CUDA context current"); | |
| 193 | ✗ | } | |
| 194 | |||
| 195 | ✗ | CudaInterop::Current::~Current() { | |
| 196 | ✗ | CUcontext popped = nullptr; | |
| 197 | ✗ | driver_.ctxPopCurrent(&popped); | |
| 198 | ✗ | } | |
| 199 | |||
| 200 | 6 | std::unique_ptr<CudaInterop> CudaInterop::Create(DeviceContext& context) { | |
| 201 | 6 | const CudaDriver* driver = CudaDriver::Get(); | |
| 202 | 6 | const DeviceCapabilities& capabilities = context.Capabilities(); | |
| 203 | 6 | if (driver == nullptr || !capabilities.externalMemory || !capabilities.externalTimeline) { | |
| 204 | 6 | return nullptr; | |
| 205 | } | ||
| 206 | ✗ | int count = 0; | |
| 207 | ✗ | driver->Check(driver->deviceGetCount(&count), "Counting CUDA devices"); | |
| 208 | ✗ | for (int ordinal = 0; ordinal < count; ++ordinal) { | |
| 209 | ✗ | CUdevice device = 0; | |
| 210 | ✗ | unsigned char uuid[16] = {}; | |
| 211 | ✗ | if (driver->deviceGet(&device, ordinal) == 0 && driver->deviceGetUuid(uuid, device) == 0 && | |
| 212 | ✗ | Hex(uuid, sizeof(uuid)) == context.Info().uuid) { | |
| 213 | ✗ | return std::unique_ptr<CudaInterop>(new CudaInterop(context, *driver, device, ordinal)); | |
| 214 | } | ||
| 215 | } | ||
| 216 | ✗ | return nullptr; | |
| 217 | } | ||
| 218 | |||
| 219 | ✗ | CudaInterop::CudaInterop(DeviceContext& context, const CudaDriver& driver, CUdevice device, | |
| 220 | ✗ | int ordinal) | |
| 221 | ✗ | : context_(context), driver_(driver), device_(device), ordinal_(ordinal), allocations_(driver) { | |
| 222 | // The primary context, which PyTorch, Warp and the runtime API share. | ||
| 223 | ✗ | driver_.Check(driver_.primaryCtxRetain(&cuda_, device_), "Retaining the CUDA context"); | |
| 224 | ✗ | } | |
| 225 | |||
| 226 | ✗ | CudaInterop::~CudaInterop() { | |
| 227 | ✗ | { | |
| 228 | // Neither API may still use what goes: CUDA's streams finish, then Ernie's queues, which | ||
| 229 | // may wait on the streams' timelines. | ||
| 230 | ✗ | const Current current(*this); | |
| 231 | ✗ | driver_.ctxSynchronize(); | |
| 232 | ✗ | Quietly([this] { context_.WaitIdle(); }); | |
| 233 | ✗ | for (const auto& [id, mapping] : mappings_) { | |
| 234 | ✗ | driver_.memFree(mapping.pointer); | |
| 235 | ✗ | driver_.destroyExternalMemory(mapping.memory); | |
| 236 | } | ||
| 237 | // The context's queue timelines and the streams' own. | ||
| 238 | ✗ | for (const auto& [timeline, semaphore] : semaphores_) { | |
| 239 | ✗ | driver_.destroyExternalSemaphore(semaphore); | |
| 240 | } | ||
| 241 | ✗ | allocations_.FreeAll(); | |
| 242 | ✗ | } | |
| 243 | ✗ | driver_.primaryCtxRelease(device_); | |
| 244 | ✗ | } | |
| 245 | |||
| 246 | ✗ | CUdeviceptr CudaInterop::Map(const Buffer& buffer) { | |
| 247 | ✗ | const std::lock_guard lock(mutex_); | |
| 248 | ✗ | const auto found = mappings_.find(buffer.Id()); | |
| 249 | ✗ | if (found != mappings_.end()) { | |
| 250 | ✗ | return found->second.pointer; | |
| 251 | } | ||
| 252 | ✗ | ERNIE_ZONE("Map memory into CUDA"); | |
| 253 | ✗ | ExternalHandle handle = buffer.Export(); | |
| 254 | ✗ | MemoryHandleDesc memory{.type = kMemoryHandleType, | |
| 255 | .handle = {}, | ||
| 256 | ✗ | .size = buffer.MemorySize(), | |
| 257 | .flags = kDedicated, | ||
| 258 | ✗ | .reserved = {}}; | |
| 259 | ✗ | SetHandle(memory, handle); | |
| 260 | ✗ | const Current current(*this); | |
| 261 | ✗ | Mapping mapping; | |
| 262 | ✗ | driver_.Check(driver_.importExternalMemory(&mapping.memory, &memory), | |
| 263 | "Importing memory into CUDA"); | ||
| 264 | ✗ | handle.Imported(); | |
| 265 | ✗ | const BufferDesc whole{.offset = 0, .size = buffer.Size(), .flags = 0, .reserved = {}}; | |
| 266 | ✗ | const CUresult mapped = | |
| 267 | ✗ | driver_.externalMemoryGetMappedBuffer(&mapping.pointer, mapping.memory, &whole); | |
| 268 | ✗ | if (mapped != 0) { | |
| 269 | ✗ | driver_.destroyExternalMemory(mapping.memory); | |
| 270 | ✗ | driver_.Check(mapped, "Mapping memory in CUDA"); | |
| 271 | } | ||
| 272 | ✗ | mappings_.emplace(buffer.Id(), mapping); | |
| 273 | ✗ | return mapping.pointer; | |
| 274 | ✗ | } | |
| 275 | |||
| 276 | ✗ | CUexternalSemaphore CudaInterop::Import(const Timeline& timeline) { | |
| 277 | ✗ | const auto found = semaphores_.find(&timeline); | |
| 278 | ✗ | if (found != semaphores_.end()) { | |
| 279 | ✗ | return found->second; | |
| 280 | } | ||
| 281 | ✗ | ExternalHandle handle = timeline.Export(); | |
| 282 | ✗ | SemaphoreHandleDesc desc{ | |
| 283 | .type = kSemaphoreHandleType, .handle = {}, .flags = 0, .reserved = {}}; | ||
| 284 | ✗ | SetHandle(desc, handle); | |
| 285 | ✗ | const Current current(*this); | |
| 286 | ✗ | CUexternalSemaphore semaphore = nullptr; | |
| 287 | ✗ | driver_.Check(driver_.importExternalSemaphore(&semaphore, &desc), | |
| 288 | ✗ | "Importing " + timeline.Name() + " into CUDA"); | |
| 289 | ✗ | handle.Imported(); | |
| 290 | ✗ | semaphores_.emplace(&timeline, semaphore); | |
| 291 | ✗ | return semaphore; | |
| 292 | ✗ | } | |
| 293 | |||
| 294 | ✗ | Tensor CudaInterop::View(const DeviceColumn& column) { | |
| 295 | ✗ | const CUdeviceptr base = Map(column.Storage()); | |
| 296 | ✗ | return Tensor(reinterpret_cast<void*>(base + column.Range().offset), column.Type(), | |
| 297 | ✗ | column.Shape(), {}, Device{DeviceKind::Cuda, ordinal_}); | |
| 298 | } | ||
| 299 | |||
| 300 | ✗ | void CudaInterop::CheckOnDevice(const Tensor& tensor, std::string_view what) const { | |
| 301 | ✗ | if (tensor.Location() != Device{DeviceKind::Cuda, ordinal_} || !tensor.IsContiguous() || | |
| 302 | ✗ | tensor.Rank() == 0) { | |
| 303 | ✗ | throw std::invalid_argument( | |
| 304 | ✗ | std::string(what) + " takes a contiguous tensor on cuda:" + std::to_string(ordinal_)); | |
| 305 | } | ||
| 306 | ✗ | } | |
| 307 | |||
| 308 | ✗ | void CudaInterop::Wait(CUstream stream, const Ticket& ticket) { | |
| 309 | ✗ | if (ticket.timeline != nullptr && ticket.timeline->Device() != context_.Device()) { | |
| 310 | ✗ | throw DeviceError(DeviceErrorKind::Failed, | |
| 311 | ✗ | "CUDA cannot wait on another device's " + ticket.timeline->Name()); | |
| 312 | } | ||
| 313 | ✗ | if (ticket.Done()) { | |
| 314 | ✗ | return; | |
| 315 | } | ||
| 316 | ✗ | SemaphoreWaitParams wait{}; | |
| 317 | ✗ | wait.params.fence.value = ticket.value; | |
| 318 | ✗ | const std::lock_guard lock(mutex_); | |
| 319 | ✗ | const CUexternalSemaphore semaphore = Import(*ticket.timeline); | |
| 320 | ✗ | const Current current(*this); | |
| 321 | ✗ | driver_.Check(driver_.waitExternalSemaphoresAsync(&semaphore, &wait, 1, stream), | |
| 322 | ✗ | "Waiting on " + ticket.timeline->Name() + " in CUDA"); | |
| 323 | ✗ | } | |
| 324 | |||
| 325 | ✗ | Ticket CudaInterop::Signal(CUstream stream) { | |
| 326 | ✗ | const std::lock_guard lock(mutex_); | |
| 327 | ✗ | StreamTimeline& timeline = streams_[stream]; | |
| 328 | ✗ | if (!timeline.timeline) { | |
| 329 | ✗ | auto made = std::make_unique<Timeline>(context_.Device(), "a CUDA stream", true); | |
| 330 | ✗ | timeline.semaphore = Import(*made); | |
| 331 | ✗ | timeline.timeline = std::move(made); | |
| 332 | ✗ | } | |
| 333 | ✗ | SemaphoreSignalParams signal{}; | |
| 334 | ✗ | signal.params.fence.value = timeline.signalled + 1; | |
| 335 | ✗ | const Current current(*this); | |
| 336 | ✗ | driver_.Check(driver_.signalExternalSemaphoresAsync(&timeline.semaphore, &signal, 1, stream), | |
| 337 | "Signalling from CUDA"); | ||
| 338 | ✗ | return timeline.timeline->At(++timeline.signalled); | |
| 339 | ✗ | } | |
| 340 | |||
| 341 | ✗ | Ticket CudaInterop::Assign(const DeviceColumn& column, const Tensor& tensor, CUstream stream) { | |
| 342 | ✗ | const std::uint64_t bytes = column.Bytes(); | |
| 343 | ✗ | CheckOnDevice(tensor, "Assigning " + column.Name()); | |
| 344 | ✗ | if (tensor.Bytes() != bytes) { | |
| 345 | ✗ | throw std::invalid_argument("Assigning " + column.Name() + " takes a tensor of its size"); | |
| 346 | } | ||
| 347 | ✗ | const CUdeviceptr rows = Map(column.Storage()) + column.Range().offset; | |
| 348 | ✗ | const auto source = reinterpret_cast<CUdeviceptr>(tensor.Data()); | |
| 349 | ✗ | if (source != rows) { | |
| 350 | // After Ernie's work that may still read the rows. | ||
| 351 | ✗ | for (const Queue* queue : context_.Queues()) { | |
| 352 | ✗ | Wait(stream, queue->Last()); | |
| 353 | ✗ | } | |
| 354 | ✗ | const Current current(*this); | |
| 355 | ✗ | driver_.Check(driver_.memcpyDtoDAsync(rows, source, bytes, stream), | |
| 356 | ✗ | "Copying a tensor into " + column.Name()); | |
| 357 | ✗ | context_.CountCopy(bytes); | |
| 358 | ✗ | } | |
| 359 | ✗ | return Signal(stream); | |
| 360 | } | ||
| 361 | |||
| 362 | ✗ | std::unique_ptr<DeviceColumn> CudaInterop::Wrap(const Tensor& tensor, std::string_view name) { | |
| 363 | ✗ | ERNIE_ZONE("Wrap a CUDA tensor"); | |
| 364 | ✗ | CheckOnDevice(tensor, "Wrapping " + std::string(name)); | |
| 365 | ✗ | const auto pointer = reinterpret_cast<CUdeviceptr>(tensor.Data()); | |
| 366 | ✗ | const std::uint64_t bytes = tensor.Bytes(); | |
| 367 | ✗ | const Current current(*this); | |
| 368 | // Only virtual memory allocations have handles, and only those asked to be shareable export. | ||
| 369 | ✗ | CUmemGenericAllocationHandle allocation = 0; | |
| 370 | ✗ | if (driver_.memRetainAllocationHandle(&allocation, reinterpret_cast<void*>(pointer)) != 0) { | |
| 371 | ✗ | return nullptr; | |
| 372 | } | ||
| 373 | ✗ | AllocationProperties properties{}; | |
| 374 | ✗ | CUdeviceptr base = 0; | |
| 375 | ✗ | std::size_t size = 0; | |
| 376 | ✗ | ExternalHandle exported; | |
| 377 | ✗ | if (driver_.memGetAllocationPropertiesFromHandle(&properties, allocation) == 0 && | |
| 378 | ✗ | (properties.requestedHandleTypes & kShareableHandle) != 0 && | |
| 379 | ✗ | driver_.memGetAddressRange(&base, &size, pointer) == 0 && pointer + bytes <= base + size) { | |
| 380 | #if defined(_WIN32) | ||
| 381 | void* handle = nullptr; | ||
| 382 | if (driver_.memExportToShareableHandle(&handle, allocation, kShareableHandle, 0) == 0) { | ||
| 383 | exported = ExternalHandle(reinterpret_cast<std::intptr_t>(handle)); | ||
| 384 | } | ||
| 385 | #else | ||
| 386 | ✗ | int handle = -1; | |
| 387 | ✗ | if (driver_.memExportToShareableHandle(&handle, allocation, kShareableHandle, 0) == 0) { | |
| 388 | ✗ | exported = ExternalHandle(handle); | |
| 389 | } | ||
| 390 | #endif | ||
| 391 | } | ||
| 392 | ✗ | driver_.memRelease(allocation); | |
| 393 | ✗ | if (!exported) { | |
| 394 | ✗ | return nullptr; | |
| 395 | } | ||
| 396 | ✗ | auto storage = std::make_shared<const Buffer>(context_, std::move(exported), size, name, | |
| 397 | ✗ | SharedFrom::OtherApi); | |
| 398 | ✗ | std::vector<std::int64_t> element(tensor.Shape().begin() + 1, tensor.Shape().end()); | |
| 399 | ✗ | return std::make_unique<DeviceColumn>( | |
| 400 | ✗ | context_, std::move(storage), pointer - base, name, tensor.Type(), std::move(element), | |
| 401 | ✗ | static_cast<std::uint64_t>(tensor.Shape().front()), std::make_shared<const Tensor>(tensor)); | |
| 402 | ✗ | } | |
| 403 | |||
| 404 | ✗ | std::optional<CudaDriver::Shareable> CudaDriver::AllocateShareable(int ordinal, | |
| 405 | std::size_t bytes) const { | ||
| 406 | ✗ | const AllocationProperties properties{.type = kPinned, | |
| 407 | .requestedHandleTypes = kShareableHandle, | ||
| 408 | .location = {kOnDevice, ordinal}, | ||
| 409 | .win32HandleMetaData = nullptr, | ||
| 410 | ✗ | .allocFlags = {}}; | |
| 411 | ✗ | std::size_t granule = 0; | |
| 412 | ✗ | if (memGetAllocationGranularity(&granule, &properties, kMinimumGranule) != 0) { | |
| 413 | ✗ | return std::nullopt; | |
| 414 | } | ||
| 415 | ✗ | Shareable allocation; | |
| 416 | ✗ | allocation.size = (std::max<std::size_t>(bytes, 1) + granule - 1) / granule * granule; | |
| 417 | ✗ | if (memCreate(&allocation.handle, allocation.size, &properties, 0) != 0) { | |
| 418 | ✗ | return std::nullopt; | |
| 419 | } | ||
| 420 | ✗ | if (memAddressReserve(&allocation.pointer, allocation.size, 0, 0, 0) != 0) { | |
| 421 | ✗ | memRelease(allocation.handle); | |
| 422 | ✗ | return std::nullopt; | |
| 423 | } | ||
| 424 | ✗ | const AccessDescription access{.location = {kOnDevice, ordinal}, .flags = kReadWrite}; | |
| 425 | ✗ | if (memMap(allocation.pointer, allocation.size, 0, allocation.handle, 0) != 0 || | |
| 426 | ✗ | memSetAccess(allocation.pointer, allocation.size, &access, 1) != 0) { | |
| 427 | ✗ | FreeShareable(allocation); | |
| 428 | ✗ | return std::nullopt; | |
| 429 | } | ||
| 430 | ✗ | return allocation; | |
| 431 | } | ||
| 432 | |||
| 433 | ✗ | void CudaDriver::FreeShareable(const Shareable& allocation) const { | |
| 434 | ✗ | memUnmap(allocation.pointer, allocation.size); | |
| 435 | ✗ | memAddressFree(allocation.pointer, allocation.size); | |
| 436 | ✗ | memRelease(allocation.handle); | |
| 437 | ✗ | } | |
| 438 | |||
| 439 | ✗ | CUdeviceptr ShareableAllocations::Allocate(int ordinal, std::size_t bytes) { | |
| 440 | ✗ | const std::optional<CudaDriver::Shareable> allocation = | |
| 441 | ✗ | driver_.AllocateShareable(ordinal, bytes); | |
| 442 | ✗ | if (!allocation) { | |
| 443 | return 0; | ||
| 444 | } | ||
| 445 | ✗ | const std::lock_guard lock(mutex_); | |
| 446 | ✗ | live_.emplace(allocation->pointer, *allocation); | |
| 447 | ✗ | return allocation->pointer; | |
| 448 | ✗ | } | |
| 449 | |||
| 450 | ✗ | void ShareableAllocations::Free(CUdeviceptr pointer) { | |
| 451 | ✗ | CudaDriver::Shareable allocation; | |
| 452 | ✗ | { | |
| 453 | ✗ | const std::lock_guard lock(mutex_); | |
| 454 | ✗ | const auto found = live_.find(pointer); | |
| 455 | ✗ | if (found == live_.end()) { | |
| 456 | ✗ | return; | |
| 457 | } | ||
| 458 | ✗ | allocation = found->second; | |
| 459 | ✗ | live_.erase(found); | |
| 460 | ✗ | } | |
| 461 | ✗ | driver_.FreeShareable(allocation); | |
| 462 | } | ||
| 463 | |||
| 464 | ✗ | void ShareableAllocations::FreeAll() { | |
| 465 | ✗ | const std::lock_guard lock(mutex_); | |
| 466 | ✗ | for (const auto& [pointer, allocation] : live_) { | |
| 467 | ✗ | driver_.FreeShareable(allocation); | |
| 468 | } | ||
| 469 | ✗ | live_.clear(); | |
| 470 | ✗ | } | |
| 471 | |||
| 472 | ✗ | CUdeviceptr CudaInterop::Allocate(std::size_t bytes) { | |
| 473 | ✗ | const Current current(*this); | |
| 474 | ✗ | return allocations_.Allocate(ordinal_, bytes); | |
| 475 | ✗ | } | |
| 476 | |||
| 477 | ✗ | void CudaInterop::Free(CUdeviceptr pointer) { | |
| 478 | ✗ | const Current current(*this); | |
| 479 | ✗ | allocations_.Free(pointer); | |
| 480 | ✗ | } | |
| 481 | |||
| 482 | } // namespace ernie | ||
| 483 |