Skip to content

Commit d0bce1a

Browse files
committed
fix(cuda.core): free a new pool or VA reservation if its host box allocation throws
1 parent 222e988 commit d0bce1a

2 files changed

Lines changed: 33 additions & 0 deletions

File tree

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,24 @@ static void clear_mempool_peer_access(CUmemoryPool pool, int owner_device) noexc
7373
// owner_device is the ordinal of a device-located pool, or -1 when there is no
7474
// device to skip (host pools) or the location is unknown (imported pools).
7575
static MemoryPoolHandle wrap_mempool_owned(CUmemoryPool pool, int owner_device) {
76+
// new can throw before the shared_ptr deleter exists. This guard destroys
77+
// the pool on that path and is disarmed once the box is allocated; a later
78+
// control-block failure is the deleter's.
79+
struct RawPoolGuard {
80+
CUmemoryPool resource;
81+
int owner_device;
82+
83+
~RawPoolGuard() noexcept {
84+
if (resource) {
85+
GILReleaseGuard gil;
86+
clear_mempool_peer_access(resource, owner_device);
87+
pw_cuMemPoolDestroy(resource);
88+
}
89+
}
90+
} guard{pool, owner_device};
91+
7692
auto* raw = new MemoryPoolBox{pool};
93+
guard.resource = nullptr;
7794
// Count before constructing the shared_ptr: if its control-block allocation
7895
// throws, it runs the deleter on raw, whose decrement then balances this.
7996
g_live_owned_mempools.fetch_add(1, std::memory_order_relaxed);

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,23 @@ VaReservationHandle create_va_reservation_handle(size_t size, size_t alignment,
101101
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemAddressReserve, &ptr, size, alignment, hint, 0))) {
102102
return {};
103103
}
104+
// new can throw before the shared_ptr deleter exists. This guard frees the
105+
// reservation on that path and is disarmed once the box is allocated; a
106+
// later control-block failure is the deleter's.
107+
struct RawReservationGuard {
108+
CUdeviceptr ptr;
109+
size_t size;
110+
111+
~RawReservationGuard() noexcept {
112+
if (ptr) {
113+
GILReleaseGuard gil;
114+
pw_cuMemAddressFree(ptr, size);
115+
}
116+
}
117+
} guard{ptr, size};
118+
104119
auto* raw = new VaReservationBox{{ptr}, size};
120+
guard.ptr = 0;
105121
// Count before constructing the shared_ptr: if its control-block allocation
106122
// throws, it runs the deleter on raw, whose decrement then balances this.
107123
g_live_va_reservations.fetch_add(1, std::memory_order_relaxed);

0 commit comments

Comments
 (0)