Skip to content

Commit 8fb56cb

Browse files
mdboommdboom
committed
Address CodeRabbit's comment
1 parent 52ef259 commit 8fb56cb

1 file changed

Lines changed: 67 additions & 20 deletions

File tree

‎cuda_core/cuda/core/_cpp/rt/memory.cpp‎

Lines changed: 67 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include <atomic>
1212
#include <cstddef>
1313
#include <cstring>
14+
#include <exception>
1415
#include <memory>
1516
#include <mutex>
1617
#include <stdexcept>
@@ -183,21 +184,51 @@ CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h
183184
return CUDA_SUCCESS;
184185
}
185186

187+
// Build the box for a freshly acquired device pointer. Binding the deallocation
188+
// stream and allocating the box can fail: make_deallocation_stream returns
189+
// false (err set) when no context is current, and either step can throw
190+
// std::bad_alloc. On any failure run undo() to release the pointer and return
191+
// nullptr; a thrown exception is left in `failure` for the caller to rethrow
192+
// once it holds no lock (undo may not report under a lock, see
193+
// deviceptr_import_ipc). Never throws.
194+
template <typename Undo>
195+
static DevicePtrBox* make_deviceptr_box(
196+
CUdeviceptr ptr, const StreamHandle& h_stream, Undo&& undo,
197+
std::exception_ptr& failure) noexcept {
198+
DevicePtrBox* box = nullptr;
199+
try {
200+
DeallocationStream ds;
201+
if (make_deallocation_stream(h_stream, ds)) {
202+
box = new DevicePtrBox{ptr, std::move(ds)};
203+
}
204+
} catch (...) {
205+
failure = std::current_exception();
206+
}
207+
if (!box) {
208+
undo();
209+
}
210+
return box;
211+
}
212+
186213
DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) {
187214
GILReleaseGuard gil;
188215
CUdeviceptr ptr;
189216
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemAllocFromPoolAsync, &ptr, size, *h_pool, as_cu(h_stream)))) {
190217
return {};
191218
}
192219

193-
DeallocationStream ds;
194-
if (!make_deallocation_stream(h_stream, ds)) {
195-
pw_cuMemFreeAsync(ptr, as_cu(h_stream));
220+
std::exception_ptr failure;
221+
DevicePtrBox* raw = make_deviceptr_box(
222+
ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
223+
if (!raw) {
224+
if (failure) {
225+
std::rethrow_exception(failure);
226+
}
196227
return {};
197228
}
198229

199230
auto box = std::shared_ptr<DevicePtrBox>(
200-
new DevicePtrBox{ptr, std::move(ds)},
231+
raw,
201232
[h_pool](DevicePtrBox* b) {
202233
GILReleaseGuard gil;
203234
const DeallocationStream& stream = b->deallocation;
@@ -220,14 +251,18 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream)
220251
return {};
221252
}
222253

223-
DeallocationStream ds;
224-
if (!make_deallocation_stream(h_stream, ds)) {
225-
pw_cuMemFreeAsync(ptr, as_cu(h_stream));
254+
std::exception_ptr failure;
255+
DevicePtrBox* raw = make_deviceptr_box(
256+
ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
257+
if (!raw) {
258+
if (failure) {
259+
std::rethrow_exception(failure);
260+
}
226261
return {};
227262
}
228263

229264
auto box = std::shared_ptr<DevicePtrBox>(
230-
new DevicePtrBox{ptr, std::move(ds)},
265+
raw,
231266
[](DevicePtrBox* b) {
232267
GILReleaseGuard gil;
233268
const DeallocationStream& stream = b->deallocation;
@@ -488,6 +523,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
488523
GILReleaseGuard gil;
489524
CUresult discard_status = CUDA_SUCCESS;
490525
CUdeviceptr discarded = 0;
526+
std::exception_ptr failure;
491527
{
492528
std::lock_guard<std::mutex> lock(ipc_import_mutex);
493529

@@ -500,10 +536,19 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
500536
return {};
501537
}
502538

503-
DeallocationStream ds;
504-
if (make_deallocation_stream(h_stream, ds)) {
539+
// No deallocation stream could be recorded (or the box could not be
540+
// allocated): discard the import with the raw call (a pw_ report would
541+
// acquire the GIL under the mutex).
542+
DevicePtrBox* raw = make_deviceptr_box(
543+
ptr, h_stream,
544+
[&]() noexcept {
545+
discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream)); // raw: under ipc_import_mutex
546+
discarded = ptr;
547+
},
548+
failure);
549+
if (raw) {
505550
auto box = std::shared_ptr<DevicePtrBox>(
506-
new DevicePtrBox{ptr, std::move(ds)},
551+
raw,
507552
[h_pool, key](DevicePtrBox* b) {
508553
// Release the GIL first (the GIL is the outermost lock), then hold the
509554
// mutex across unregister + free: a concurrent import that finds this
@@ -530,11 +575,6 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
530575
ipc_ptr_cache.register_handle(key, h);
531576
return h;
532577
}
533-
534-
// No deallocation stream could be recorded: discard the import with
535-
// the raw call (a pw_ report would acquire the GIL under the mutex).
536-
discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream)); // raw: under ipc_import_mutex
537-
discarded = ptr;
538578
}
539579
if (discard_status != CUDA_SUCCESS) {
540580
char operation[160];
@@ -543,6 +583,9 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
543583
"failed while discarding an IPC import that could not record a "
544584
"deallocation stream; the mapping leaked");
545585
}
586+
if (failure) {
587+
std::rethrow_exception(failure);
588+
}
546589
return {};
547590

548591
} else {
@@ -552,14 +595,18 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
552595
return {};
553596
}
554597

555-
DeallocationStream ds;
556-
if (!make_deallocation_stream(h_stream, ds)) {
557-
pw_cuMemFreeAsync(ptr, as_cu(h_stream));
598+
std::exception_ptr failure;
599+
DevicePtrBox* raw = make_deviceptr_box(
600+
ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure);
601+
if (!raw) {
602+
if (failure) {
603+
std::rethrow_exception(failure);
604+
}
558605
return {};
559606
}
560607

561608
auto box = std::shared_ptr<DevicePtrBox>(
562-
new DevicePtrBox{ptr, std::move(ds)},
609+
raw,
563610
[h_pool](DevicePtrBox* b) {
564611
GILReleaseGuard gil;
565612
const DeallocationStream& stream = b->deallocation;

0 commit comments

Comments
 (0)