nanovdb python: bind DeviceBuffer/DeviceGridHandle.recordUse for external streams

The bindings' interop surface hands raw device pointers to CuPy / PyTorch /
Numba via device_ptr(), __cuda_array_interface__ and __dlpack__, and those
frameworks launch kernels on their own non-blocking streams. Such work is
invisible to the buffer's automatic upload/download use tracking, so the
cudaFreeAsync issued when the buffer is cleared or destroyed could race it —
the C++ side grew DeviceBuffer::recordUse(device, stream) for exactly this
caller, but Python had no way to reach it.

Bind recordUse(stream, device=-1) on DeviceBuffer and on DeviceGridHandle
(which owns its buffer internally and is where every real workflow's owning
buffer lives). Both forward to a shared recordUseChecked helper that
validates the device id — the C++ method indexes per-device tracking state
unchecked — raising IndexError on an out-of-range id, with -1 selecting the
current CUDA device. The device_ptr and __cuda_array_interface__ docstrings
now point callers at recordUse. Non-owning (from_external) buffers accept
the call as a no-op, matching the C++ behavior.

Add TestRecordUse to TestGpuInterop.py: a CuPy reduction enqueued on a
non-blocking stream against a zero-copy CAI view, recorded, and validated
after the handle is destroyed while the work may still be in flight; plus
default-stream/explicit-device acceptance, IndexError on a bad device id,
and the non-owning no-op.

Also normalize the cuda/ sources' includes of shared binding headers: add
the python source dir to the target's private include paths so they use
plain names ("PyGridHandle.h", "BuildTypes.def") instead of "../"-relative
paths, which the rest of the codebase does not use.

Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>

Author: swahtz

nanovdb/nanovdb/python/CMakeLists.txt

@@ -58,7 +58,10 @@ nanobind_add_module(nanovdb_python NB_STATIC
     cuda/PyDeviceGridChecksum.cu
 )
 
-target_include_directories(nanovdb_python PRIVATE ${CUDA_INCLUDE_DIRECTORY})
+# CMAKE_CURRENT_SOURCE_DIR lets sources in cuda/ include the shared binding
+# headers by their plain names ("PyGridHandle.h", "BuildTypes.def") instead of
+# "../"-relative paths, matching the include style used across the codebase.
+target_include_directories(nanovdb_python PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${CUDA_INCLUDE_DIRECTORY})
 target_link_libraries(nanovdb_python PRIVATE nanovdb ${CUDA_LIBRARIES} ${NANOVDB_BLOSC} ${NANOVDB_ZLIB} ${NANOVDB_OPENVDB} ${NANOVDB_TBB})
 target_compile_definitions(nanovdb_python PRIVATE ${NANOVDB_USE_CUDA_FLAG} ${NANOVDB_USE_BLOSC_FLAG} ${NANOVDB_USE_ZLIB_FLAG} ${NANOVDB_USE_OPENVDB_FLAG} ${NANOVDB_USE_TBB_FLAG})
 set_target_properties(nanovdb_python PROPERTIES OUTPUT_NAME "nanovdb")

nanovdb/nanovdb/python/cuda/PyDeviceBuffer.cc

@@ -46,7 +46,9 @@ void defineDeviceBuffer(nb::module_& m)
             "does NOT take ownership: it will never free either pointer, so "
             "the caller must keep both allocations alive for the buffer's "
             "lifetime. The device pointer is associated with the current CUDA "
-            "device.");
+            "device.")
+        .def("recordUse", &recordUseChecked, "stream"_a, "device"_a = -1,
+             kRecordUseDoc);
 }
 
 } // namespace pynanovdb

nanovdb/nanovdb/python/cuda/PyDeviceBuffer.h

@@ -9,8 +9,11 @@
 #include <nanobind/ndarray.h>
 
 #include <cstdint>
+#include <string>
 
 #include <cuda_runtime.h>
+
+#include <nanovdb/cuda/DeviceBuffer.h>
 #endif
 
 namespace nb = nanobind;
@@ -19,6 +22,39 @@ namespace pynanovdb {
 
 #ifdef NANOVDB_USE_CUDA
 
+/// @brief Bounds-checked forwarder for DeviceBuffer::recordUse, shared by the
+///        DeviceBuffer and DeviceGridHandle bindings. device == -1 selects the
+///        current CUDA device. recordUse itself indexes per-device tracking
+///        state unchecked, so validate here where we can raise a Python
+///        exception instead.
+inline void recordUseChecked(nanovdb::cuda::DeviceBuffer& buf, uintptr_t stream, int device)
+{
+    int count = 0;
+    cudaCheck(cudaGetDeviceCount(&count));
+    if (device < 0) cudaCheck(cudaGetDevice(&device));
+    if (device >= count) {
+        const std::string msg = "recordUse: device id " + std::to_string(device) +
+                                " out of range [0, " + std::to_string(count) + ").";
+        throw nb::index_error(msg.c_str());
+    }
+    nb::gil_scoped_release release;
+    buf.recordUse(device, reinterpret_cast<cudaStream_t>(stream));
+}
+
+/// @brief Docstring shared by the DeviceBuffer and DeviceGridHandle recordUse
+///        bindings (the two forward to the same underlying buffer method).
+inline constexpr char kRecordUseDoc[] =
+    "Record that this buffer's device data was just used on `stream` (a raw "
+    "cudaStream_t as a Python int, e.g. cupy.cuda.Stream.ptr), so the device "
+    "free issued when the buffer is cleared or destroyed is ordered after "
+    "that work. Uploads/downloads record themselves automatically; call this "
+    "after enqueuing your own kernels or copies against device_ptr() / "
+    "__cuda_array_interface__ / __dlpack__ on a non-blocking stream — "
+    "without it such work is only safe if you synchronize before dropping "
+    "the buffer. device selects which device's buffer was used (-1 = the "
+    "current CUDA device). No-op on non-owning (from_external) buffers, "
+    "which never free their pointers.";
+
 /// @brief Bind the device-interop surface (CUDA Array Interface / DLPack, raw
 ///        device/host pointers, streams) onto a device-buffer-like class.
 ///
@@ -40,7 +76,11 @@ void addDeviceInterop(nb::class_<BufferT>& cls)
             return reinterpret_cast<uintptr_t>(buf.deviceData());
         },
         "Raw device pointer to the current device's buffer as a Python int "
-        "(0 if no device allocation exists yet).");
+        "(0 if no device allocation exists yet). Work you enqueue against "
+        "this pointer on a non-blocking stream is invisible to the buffer's "
+        "lifetime tracking: record it afterwards where the buffer supports "
+        "it (DeviceBuffer.recordUse / DeviceGridHandle.recordUse), or "
+        "synchronize before the buffer is destroyed.");
 
     cls.def(
         "host_ptr",

nanovdb/nanovdb/python/cuda/PyDeviceGridHandle.cu

@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 #ifdef NANOVDB_USE_CUDA
 
-#include "../PyGridHandle.h"
+#include "PyGridHandle.h"
+#include "PyDeviceBuffer.h"  // for recordUseChecked / kRecordUseDoc
 #include <nanobind/ndarray.h>
 
 #include <cstdint>
@@ -54,7 +55,7 @@ static nb::object pyDeviceGrid(nb::handle py_handle, uint32_t n)
             return grid ? nb::cast(grid, nb::rv_policy::reference, py_handle)   \
                         : nb::none();                                           \
         }
-#include "../BuildTypes.def"
+#include "BuildTypes.def"
         default:
             return nb::none();
     }
@@ -114,7 +115,17 @@ void defineDeviceGridHandle(nb::module_& m)
             },
             "Raw device pointer to the base of the whole device buffer as a "
             "Python int (0 if the handle has not been uploaded to the device "
-            "yet).")
+            "yet). Work you enqueue against this pointer on a non-blocking "
+            "stream is invisible to the buffer's lifetime tracking: call "
+            "recordUse(stream) afterwards, or synchronize before the handle "
+            "is destroyed.")
+        .def(
+            "recordUse",
+            [](GridHandle<BufferT>& handle, uintptr_t stream, int device) {
+                recordUseChecked(handle.buffer(), stream, device);
+            },
+            "stream"_a, "device"_a = -1,
+            kRecordUseDoc)
         .def_prop_ro(
             "__cuda_array_interface__",
             [](GridHandle<BufferT>& handle) {
@@ -133,7 +144,10 @@ void defineDeviceGridHandle(nb::module_& m)
             },
             "CUDA Array Interface (v3) view of the whole device buffer as 1-D "
             "uint8 — lets CuPy / Numba / PyTorch consume the serialized grid "
-            "bytes zero-copy. Returns a null data pointer until deviceUpload.")
+            "bytes zero-copy. Returns a null data pointer until deviceUpload. "
+            "After enqueuing work on this view from a non-blocking stream, "
+            "call recordUse(stream) so the buffer's device free is ordered "
+            "after it.")
         .def(
             "__dlpack_device__",
             [](GridHandle<BufferT>&) {

nanovdb/nanovdb/python/cuda/PyDeviceNodeManager.cu

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 #ifdef NANOVDB_USE_CUDA
 
-#include "../PyTree.h"
+#include "PyTree.h"
 
 #include <cstdint>
 
@@ -56,7 +56,7 @@ static nb::object pyDeviceNodeMgr(nb::handle py_self)
     if (auto* m = handle.template deviceMgr<T>()) {                           \
         return nb::cast(m, nb::rv_policy::reference, py_self);                 \
     }
-#include "../BuildTypes.def"
+#include "BuildTypes.def"
     return nb::none();
 }
 
@@ -141,7 +141,7 @@ static void defineCreateDeviceNodeManager(nb::module_& m)
             if (auto obj = tryCreateDeviceNodeManager<T>(py_grid, s); obj.is_valid()) { \
                 return obj;                                                    \
             }
-#include "../BuildTypes.def"
+#include "BuildTypes.def"
             throw nb::type_error(
                 "createDeviceNodeManager: argument is not a NanoVDB device "
                 "grid of any bound BuildT. Pass a device grid obtained from "

nanovdb/nanovdb/python/cuda/PyUnifiedGridHandle.cu

@@ -3,7 +3,7 @@
 #ifdef NANOVDB_USE_CUDA
 
 #include "PyUnifiedGridHandle.h"
-#include "../PyGridHandle.h"
+#include "PyGridHandle.h"
 
 #include <cstdint>
 
@@ -51,7 +51,7 @@ static nb::object pyUnifiedDeviceGrid(nb::handle py_handle, uint32_t n)
             return grid ? nb::cast(grid, nb::rv_policy::reference, py_handle)   \
                         : nb::none();                                           \
         }
-#include "../BuildTypes.def"
+#include "BuildTypes.def"
         default:
             return nb::none();
     }

nanovdb/nanovdb/python/test/TestGpuInterop.py

@@ -295,6 +295,62 @@ def test_from_buffer_rejects_garbage(self):
             cp.cuda.set_allocator(prev)
 
 
+@unittest.skipIf(
+    not nanovdb.isCudaAvailable(), "nanovdb module was compiled without CUDA support"
+)
+@unittest.skipIf(
+    not nanovdb.isGpuAvailable(), "No CUDA-capable GPU available"
+)
+class TestRecordUse(unittest.TestCase):
+    """DeviceBuffer/DeviceGridHandle.recordUse: order the buffer's device free
+    after work enqueued on external non-blocking streams (the CAI/DLPack
+    interop pattern, where kernels launched against device_ptr() are invisible
+    to the buffer's automatic upload/download tracking)."""
+
+    def test_free_ordered_after_external_stream_work(self):
+        cp = _require_cupy(self)
+        dh, _ = _build_device_onindex_grid(20.0)
+        # Baseline: reduce the serialized grid bytes while the handle is alive.
+        view = cp.asarray(dh)  # zero-copy CAI view of the device buffer
+        expected = int(view.sum())
+        # Enqueue the same reduction on a NON-BLOCKING stream, record the use,
+        # and drop the handle while the reduction may still be in flight. The
+        # buffer's cudaFreeAsync must be ordered after the recorded event; if
+        # it were not, the reduction would race the free and read freed memory.
+        s = cp.cuda.Stream(non_blocking=True)
+        with s:
+            pending = view.sum()  # device scalar; no host sync yet
+        dh.recordUse(s.ptr)
+        del dh, view
+        s.synchronize()
+        self.assertEqual(int(pending), expected)
+
+    def test_default_stream_and_explicit_device(self):
+        dh, _ = _build_device_onindex_grid(10.0)
+        # Default stream (0) and the current device, both explicit and implied.
+        dh.recordUse(0)
+        dh.recordUse(0, device=0)
+
+    def test_rejects_out_of_range_device(self):
+        dh, _ = _build_device_onindex_grid(10.0)
+        with self.assertRaises(IndexError):
+            dh.recordUse(0, device=1_000_000)
+
+    def test_noop_on_non_owning_buffer(self):
+        cp = _require_cupy(self)
+        prev = cp.cuda.get_allocator()
+        cp.cuda.set_allocator(cp.cuda.malloc_managed)
+        try:
+            mbuf = cp.zeros(256, dtype=cp.uint8)
+            ptr = int(mbuf.data.ptr)
+            ext = nanovdb.cuda.DeviceBuffer.from_external(256, ptr, ptr)
+            # Non-owning buffers never free their pointers, so recordUse has
+            # nothing to order — it must be accepted and be a no-op.
+            ext.recordUse(0)
+        finally:
+            cp.cuda.set_allocator(prev)
+
+
 @unittest.skipIf(
     not nanovdb.isCudaAvailable(), "nanovdb module was compiled without CUDA support"
 )