NVIDIA / NVIDIA/cuvs

Possible memory leak: Tensor.Close returns early on the CUDA path and skips two frees

Open Beginner friendly
#2,558 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Cuda
Stars
854
Forks
236
Avg merge
3d 3h
Merged PRs (30d)
62

Description

Possible memory leak: Tensor.Close returns early on the CUDA path and skips two frees

Close handles the device data in the kDLCUDA branch and returns from inside
it, so the shared cleanup below the branch never runs for GPU tensors.

go/dlpack.go:184

func (t *Tensor[T]) Close() error {
	if t.C_tensor.dl_tensor.device.device_type == C.kDLCUDA {
		bytes := t.sizeInBytes()
		res, err := NewResource(nil)
		if err != nil {
			return err
		}
		err = CheckCuvs(CuvsError(C.cuvsRMMFree(res.Resource, t.C_tensor.dl_tensor.data, C.size_t(bytes))))

		return err
	} else if t.C_tensor.dl_tensor.device.device_type == C.kDLCPU {
		if t.C_tensor.dl_tensor.data != nil {
			C.free(t.C_tensor.dl_tensor.data)
			t.C_tensor.dl_tensor.data = nil
		}
	}

	if t.C_tensor.dl_tensor.shape != nil {
		C.free(unsafe.Pointer(t.C_tensor.dl_tensor.shape))
		t.C_tensor.dl_tensor.shape = nil
	}

	if t.C_tensor != nil {
		C.free(unsafe.Pointer(t.C_tensor))
		t.C_tensor = nil
	}

Both skipped frees release host memory that NewTensorOnDevice allocated with
C.malloc:

	shapePtr := C.malloc(C.size_t(len(shape) * int(unsafe.Sizeof(C.int64_t(0)))))
	...
	dlm := (*C.DLManagedTensor)(C.malloc(C.size_t(unsafe.Sizeof(C.DLManagedTensor{}))))

The CPU branch falls through and frees all three, which is what makes the CUDA
path stand out. Every GPU tensor therefore leaks len(shape)*8 bytes for the
shape array plus one DLManagedTensor on the host, and these are created per
index build and per search.

Fix: replace the return err in the CUDA branch with an assignment to a named
error, and let control fall through to the shape and tensor frees.

If you could credit me as a reporter for my contributions to security advisory I will be thankful.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Read go/dlpack.go around line 184 and trace Tensor.Close alongside the allocations in NewTensorOnDevice. Verify that the CUDA path still releases the shape array and DLManagedTensor after device cleanup, then exercise index-build and search paths to confirm GPU tensors no longer retain those host allocations.

Written by the indexing model from the issue text.

Assessment

Tech stack
c, go
Domain
performance
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
84/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.