Skip to content

[gpu] Scope temporary device buffers with DeviceBuffer so a failed kernel launch cannot leak them #1176

Description

@yingjerkao

Problem

#1175 added CYTNX_CHECK_CUDA_LAUNCH after every kernel launch in the GPU
backend, so a launch that never runs now raises instead of leaving the output
buffer full of zeros. That is the right trade, but it introduces a new exit path
through functions that clean up by hand.

The pattern in utils_internal_gpu/ is: cudaMalloc a few temporaries, launch,
then cudaFree them at the tail of the function. The new check sits between the
launch and the frees, so a failed launch now throws past them and leaks the
device allocations. Previously the function ran to completion and freed them —
silently returning wrong numbers, but not leaking.

An exception plus a leak beats silently wrong numbers, which is why #1175 shipped
as-is. This issue is the follow-up to remove the leak.

Scope

30 launch sites across 7 files, holding 114 cudaFree calls between them:

File Launch sites cudaFree calls
utils_internal_gpu/cuGetElems_gpu.cu 11 44
utils_internal_gpu/cuGetElems_contiguous_gpu.cu 11 44
utils_internal_gpu/cuSetElems_gpu.cu 2 8
utils_internal_gpu/cuSetElems_contiguous_gpu.cu 2 8
linalg_internal_gpu/cuMaxMin_internal.cu 2 4
utils_internal_gpu/cuMovemem_gpu.cu 1 4
utils_internal_gpu/cuReduce_gpu.cu 1 2
Total 30 114

The other 183 checked launch sites allocate no temporaries and need no change.

The type already exists

utils_internal::DeviceBuffer<T> in
src/backend/utils_internal_gpu/cuScopedResource_gpu.hpp
was added by #1146 for precisely this failure mode — cuSOLVER wrappers raised
their LAPACK-info error before their cleanup block and leaked every
allocation in the function. Its constructor does checkCudaErrors(cudaMalloc(...)),
which is exactly what these call sites already do, and it handles the
zero-element case that bit Storage in #1126.

It is currently used in 5 files under linalg_internal_gpu/ (cuSvd, cuEigh,
cuInvM_inplace, cuGeSvd, cuDet) and in none of the 7 above. So this is
reuse, not a new abstraction.

Sketch

Current, in cuGetElems_gpu.cu:

checkCudaErrors(cudaMalloc((void**)&d_offj, sizeof(cytnx_uint64) * offj.size()));
// ... three more, plus cudaMemcpy of each
cuGetElems_kernel<<<NBlocks, 256>>>(new_elem_ptr_, elem_ptr_, d_offj, d_new_offj, ...);
CYTNX_CHECK_CUDA_LAUNCH(cuGetElems_kernel);   // <-- throws past the frees below
cudaFree(d_offj);
cudaFree(d_new_offj);
cudaFree(d_locators);
cudaFree(d_picksize);

After:

DeviceBuffer<cytnx_uint64> d_offj(offj.size());
// ... three more, cudaMemcpy into <buf>.get()
cuGetElems_kernel<<<NBlocks, 256>>>(new_elem_ptr_, elem_ptr_, d_offj.get(), ...);
CYTNX_CHECK_CUDA_LAUNCH(cuGetElems_kernel);   // now unwinds through the destructors
// no manual cudaFree

Notes for whoever picks this up:

  • Use the constructor, not DeviceBuffer::managed(). These are device-only
    scratch buffers; nothing reads them from the host, so plain cudaMalloc is
    correct. managed() exists for cuDet_internal, which does read from the host.
  • cuGetElems_gpu.cu and cuGetElems_contiguous_gpu.cu are 22 of the 30 sites
    and are highly repetitive (the same 4-buffer pattern per dtype), so the bulk of
    this is mechanical.
  • Keep it to the leak fix. Per the AGENTS.md guardrails, do not fold in
    unrelated refactoring of these files.

Acceptance criteria

  • No cudaFree remains between a kernel launch and the end of a function in the
    7 files above; ownership is scope-bound.
  • ctest -L gpu stays green (841/841 at the time of writing).
  • Ideally a test that forces a launch failure and shows the allocation is
    released — see the caveat below.

Caveat on testing

There is no easy in-suite way to force a launch failure: the realistic trigger is
a library built for an architecture the running device cannot execute (#1171),
which the test suite cannot set up. A targeted alternative is a unit test that
drives one of these functions with a deliberately invalid launch configuration,
or simply checking cudaMemGetInfo around a forced-throw path. Reviewers should
not expect a natural regression test here.

Links

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions