Skip to content

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

Description

@OvOhao

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions