[maxwbuckley]

Description

Fixes several memory-management bugs in the Go bindings' Tensor lifecycle (go/dlpack.go) and in cagra.SearchIndex. The changes are confined to the Go bindings, with no API changes.

Bugs fixed

#	Where	Bug	Effect
1	Tensor.Close()	On a device tensor it returned before freeing the DLPack shape and the DLManagedTensor, and never destroyed the Resource it created to call cuvsRMMFree	Host memory leaked on every device-tensor Close(). A second Close() dereferenced a nil C_tensor and panicked
2	Tensor.ToDevice()	Overwrote the host data pointer without freeing it	Every tensor moved to the GPU leaked its whole host copy
3	Tensor.ToDevice() on a tensor already on the device	Allocated a second device buffer, did an H2D cudaMemcpy from a device pointer (undefined per the CUDA docs), and dropped the old buffer	The full device buffer leaked on each redundant call
4	Tensor.ToHost() on a tensor already on the host	Did a D2H copy from a host pointer and then cuvsRMMFree'd it	Returned invalid argument. Separately, the staging buffer leaked on both ToHost error paths
5	Tensor.Expand()	Pointed dl_tensor.shape at a Go-allocated slice	Breaks the cgo pointer rules (fatal error: unpinned Go pointer stored into non-Go memory under GOEXPERIMENT=cgocheck2). A later Close() then passes a Go heap pointer to free(): in the new TestExpandRepeatedly this aborts with munmap_chunk(): invalid pointer
6	cagra.SearchIndex()	Built, uploaded and freed a device bitset on every search, even with no allow-list	One extra malloc, RMM allocation and H2D copy per unfiltered search, plus the leak from #1
7	NewTensor / NewVector / NewTensorOnDevice	Earlier allocations leaked if the DLManagedTensor allocation failed	Leak on the out-of-memory path

Fixes

• Close() has one cleanup path for both host and device tensors, closes the temporary Resource, and is a no-op on a closed tensor. If the device free fails, the tensor is left intact and the error is returned.
• ToDevice() frees the host buffer once the copy succeeds. ToDevice() and ToHost() are no-ops when the tensor is already where it's asked to go, which is what the existing CAGRA filtering test already assumes (it calls ToDevice on tensors that are on the device).
• Expand() updates the existing C shape array in place, so the DLPack struct no longer points into Go memory.
• cagra.SearchIndex() only builds the bitset when allowList != nil. An empty non-nil allow-list still means "allow nothing", as before.

Measurements

These are end-to-end measurements of the Go bindings on an RTX 5090, linked against the libcuvs 26.12 conda nightly with the CI link flags (-lcudart_static -ldl -lrt -lcuvs -lcuvs_c). Host memory is glibc heap in use (mallinfo2: uordblks + hblkhd). Device memory is from cudaMemGetInfo. Each probe ran 5 times after a warmup, and the before/after runs used the same probe program.

Operation	Before	After
NewTensorOnDevice + Close	2,478 B/call host (5/5 runs identical)	Total growth stays flat as n grows: 72 KB at n=5,000, 72–108 KB at n=50,000, i.e. under 1 B/call
1 MiB host tensor → ToDevice → Close	1,051,071 B/call host	3,744 B total, identical at n=200 and n=2,000 (6/6 runs)
ToDevice on a 16 MiB tensor already on the device	16,777,216 B/call device	0
Unfiltered cagra.SearchIndex	2,496 B/search host	0
Filtered cagra.SearchIndex	~3,000 B/search host	0

Search results are the same before and after: identical neighbour IDs, and data round-trips intact. I'm not claiming any latency change for #6. The saving is the allocations and the copy listed above.

Tests

Three regression tests added to go/dlpack_test.go, each checked to fail on the old code and pass with this change:

• TestRedundantTransfersAreNoOps: before, ToHost on a host tensor returned invalid argument.
• TestCloseIsIdempotent (host and device): before, the second Close() panicked with a nil dereference.
• TestExpandRepeatedly: before, Close() after Expand aborted with munmap_chunk(): invalid pointer, and under cgocheck2 the Expand itself was a fatal error.

go test ./... passes in all five packages, both normally and with GOEXPERIMENT=cgocheck2. On current main, the unmodified suite fails under cgocheck2 (in the existing TestDlPack, which calls Expand). After this change the whole suite is cgocheck2-clean, so running the Go tests under cgocheck2 in CI would stop this class of bug from coming back. I've left that CI change for a separate PR.

Not in this PR

• cagra.SearchIndex sizes the allow-list bitset from the largest allowed ID rather than from the dataset's row count. When the largest ID is below the last row, CAGRA reads past the end of the bitset, because bitset_view::test has no bounds check. That's a separate correctness fix with its own test, so I'll send it separately.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxwbuckley]

@achirkin any chance you can take a look :)

[maxwbuckley]

Gentle bump on this one. go/ has no CODEOWNERS entry, so nobody was auto-requested as a reviewer and CI hasn't been vetted to run yet.

@benfred you reviewed my earlier Go-binding PRs (#1471, #1472). Would you be up for taking a look, or pointing me to whoever owns the Go bindings now? The change is confined to go/dlpack.go and go/cagra/cagra.go (plus regression tests), with no API changes. It fixes the Tensor lifecycle leaks and the Expand() Go-pointer-in-C-memory abort described above.

It also needs labels (bug, non-breaking) and a /ok to test from a vetter. Thanks!

[benfred]

changes look good to me - but I'm no longer at nvidia, so you'll need to get a review from someone else. Maybe @cjnolet, @achirkin or @divyegala will be able to help you out?

[maxwbuckley]

Thanks @benfred! @cjnolet @achirkin @divyegala, benfred suggested one of you might be able to review this. It's confined to go/dlpack.go and go/cagra/cagra.go (plus regression tests), with no API changes. It fixes leaked device/host buffers in the Go Tensor bindings and an invalid free. It also needs labels (bug, non-breaking) and an /ok to test 418d2de before CI can run. Thank you!