NanoVDB Python: GPU/device bindings (nanovdb.cuda)

.github/workflows/nanovdb.yml

@@ -74,7 +74,11 @@ jobs:
           echo "/usr/local/cuda-12/bin" >> $GITHUB_PATH
           echo "LD_LIBRARY_PATH=/usr/local/cuda-12/lib64:$LD_LIBRARY_PATH" >> $GITHUB_ENV
       - name: nanobind
-        run: ./ci/install_nanobind.sh 2.5.0
+        # The NanoVDB Python bindings (CUDA interop, DLPack) track a newer
+        # nanobind than the repo-wide 2.5.0 floor. Bump here only; the shared
+        # build.yml still co-compiles the bindings against 2.5.0, so the
+        # binding code must stay >=2.5.0 source-compatible.
+        run: ./ci/install_nanobind.sh 2.12.0
       - name: build
       # NOTE: CMAKE_POSITION_INDEPENDENT_CODE set to fix default behaviour change in clang 14
       #   https://github.com/AcademySoftwareFoundation/aswf-docker/issues/228

nanovdb/nanovdb/python/CMakeLists.txt

@@ -32,13 +32,36 @@ nanobind_add_module(nanovdb_python NB_STATIC
     PyTree.cc
     PyVoxelBlockManager.cc
     cuda/PyDeviceBuffer.cc
+    cuda/PyUnifiedBuffer.cu
+    cuda/PyUnifiedGridHandle.cu
+    cuda/PyDeviceMesh.cu
+    cuda/PyDeviceStreamMap.cu
+    cuda/PyTempPool.cu
+    cuda/PyDistributedPointsToGrid.cu
     cuda/PyDeviceGridHandle.cu
+    cuda/PyDeviceNodeManager.cu
+    cuda/PyDeviceVoxelBlockManager.cu
     cuda/PyPointsToGrid.cu
     cuda/PySampleFromVoxels.cu
     cuda/PySignedFloodFill.cu
+    cuda/PyDilateGrid.cu
+    cuda/PyCoarsenGrid.cu
+    cuda/PyRefineGrid.cu
+    cuda/PyPruneGrid.cu
+    cuda/PyMergeGrids.cu
+    cuda/PyInjectData.cu
+    cuda/PyIndexToGrid.cu
+    cuda/PyMeshToGrid.cu
+    cuda/PyAddBlindData.cu
+    cuda/PyDeviceGridStats.cu
+    cuda/PyDeviceGridValidator.cu
+    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")
@@ -99,6 +122,14 @@ if(NANOVDB_BUILD_PYTHON_UNITTESTS)
     COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test/TestNanoVDB.py -v
     WORKING_DIRECTORY "${NANOVDB_PYTHON_WORKING_DIR}")
 
+  # GPU interop suite. Self-skips when the build has no CUDA or no GPU is
+  # present (so it is harmless to register unconditionally), but is a
+  # separate test file that the pytest_nanovdb command does not discover,
+  # so register it as its own ctest test mirroring the block above.
+  add_test(NAME pytest_nanovdb_gpu_interop
+    COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test/TestGpuInterop.py -v
+    WORKING_DIRECTORY "${NANOVDB_PYTHON_WORKING_DIR}")
+
   # Smoke-run every example script (they self-skip when optional
   # dependencies such as NumPy or OpenVDB are unavailable).
   add_test(NAME pytest_nanovdb_examples
@@ -110,8 +141,10 @@ if(NANOVDB_BUILD_PYTHON_UNITTESTS)
     set(PYTHONPATH "$ENV{PYTHONPATH};${NANOVDB_PYTHON_WORKING_DIR}")
     string(REPLACE "\\;" ";" PYTHONPATH "${PYTHONPATH}")
     string(REPLACE ";" "\\;" PYTHONPATH "${PYTHONPATH}")
-    set_tests_properties(pytest_nanovdb pytest_nanovdb_examples PROPERTIES ENVIRONMENT "PYTHONPATH=${PYTHONPATH}")
+    set_tests_properties(pytest_nanovdb pytest_nanovdb_gpu_interop pytest_nanovdb_examples
+      PROPERTIES ENVIRONMENT "PYTHONPATH=${PYTHONPATH}")
   else()
-    set_tests_properties(pytest_nanovdb pytest_nanovdb_examples PROPERTIES ENVIRONMENT "PYTHONPATH=$ENV{PYTHONPATH}:${NANOVDB_PYTHON_WORKING_DIR}")
+    set_tests_properties(pytest_nanovdb pytest_nanovdb_gpu_interop pytest_nanovdb_examples
+      PROPERTIES ENVIRONMENT "PYTHONPATH=$ENV{PYTHONPATH}:${NANOVDB_PYTHON_WORKING_DIR}")
   endif()
 endif()

nanovdb/nanovdb/python/NanoVDBModule.cc

@@ -8,15 +8,25 @@
 
 #include <nanovdb/NanoVDB.h>
 #include <nanovdb/io/IO.h> // for __repr__
+#include <nanovdb/tools/GridChecksum.h> // host updateChecksum for the header setters
 #ifdef NANOVDB_USE_CUDA
 #include <nanovdb/cuda/DeviceBuffer.h>
 #endif
 
+#include <cstdint>
 #include <sstream>
 
 #include "cuda/PyDeviceBuffer.h"
+#ifdef NANOVDB_USE_CUDA
+#include "cuda/PyUnifiedBuffer.h"
+#include "cuda/PyUnifiedGridHandle.h"
+#include "cuda/PyDeviceMesh.h"
+#include "cuda/PyDeviceStreamMap.h"
+#include "cuda/PyTempPool.h"
+#endif
 #include "PyBuildGrid.h"
 #include "PyGridHandle.h"
+#include "PyValidate.h"
 #include "PyHostBuffer.h"
 #include "PyIO.h"
 #include "PyMath.h"
@@ -236,21 +246,63 @@ void defineGrid(nb::module_& m)
         "the affine transform, grid flags and blind-data accessors. Concrete "
         "BuildT-typed subclasses (FloatGrid, Vec3fGrid, ...) add tree access "
         "and per-voxel queries.")
-        // Validation and flag mutators (already member functions on GridData).
+        // Validation and header mutators. Every mutator refreshes the grid's
+        // checksum afterwards (preserving its mode; a no-op when the checksum is
+        // disabled), so editing a header field never silently desyncs a
+        // checksummed grid -- a validating reader would otherwise reject it.
         .def("isValid", &GridData::isValid,
              "True iff the grid header looks consistent (magic / version / class tags).")
-        .def("setMinMaxOn", &GridData::setMinMaxOn, "on"_a = true,
-             "Toggle the HasMinMax grid flag.")
-        .def("setBBoxOn", &GridData::setBBoxOn, "on"_a = true,
-             "Toggle the HasBBox grid flag.")
-        .def("setLongGridNameOn", &GridData::setLongGridNameOn, "on"_a = true,
-             "Toggle the HasLongGridName grid flag.")
-        .def("setAverageOn", &GridData::setAverageOn, "on"_a = true,
-             "Toggle the HasAverage grid flag.")
-        .def("setStdDeviationOn", &GridData::setStdDeviationOn, "on"_a = true,
-             "Toggle the HasStdDeviation grid flag.")
-        .def("setGridName", &GridData::setGridName, "src"_a,
-             "Overwrite the grid's short name (truncated to the in-header buffer).")
+        .def("setMinMaxOn",
+             [](GridData& g, bool on) { g.setMinMaxOn(on); nanovdb::tools::updateChecksum(&g); },
+             "on"_a = true, "Toggle the HasMinMax grid flag (refreshes the checksum).")
+        .def("setBBoxOn",
+             [](GridData& g, bool on) { g.setBBoxOn(on); nanovdb::tools::updateChecksum(&g); },
+             "on"_a = true, "Toggle the HasBBox grid flag (refreshes the checksum).")
+        .def("setLongGridNameOn",
+             [](GridData& g, bool on) { g.setLongGridNameOn(on); nanovdb::tools::updateChecksum(&g); },
+             "on"_a = true, "Toggle the HasLongGridName grid flag (refreshes the checksum).")
+        .def("setAverageOn",
+             [](GridData& g, bool on) { g.setAverageOn(on); nanovdb::tools::updateChecksum(&g); },
+             "on"_a = true, "Toggle the HasAverage grid flag (refreshes the checksum).")
+        .def("setStdDeviationOn",
+             [](GridData& g, bool on) { g.setStdDeviationOn(on); nanovdb::tools::updateChecksum(&g); },
+             "on"_a = true, "Toggle the HasStdDeviation grid flag (refreshes the checksum).")
+        .def("setGridName",
+             [](GridData& g, const char* src) {
+                 const bool ok = g.setGridName(src);
+                 nanovdb::tools::updateChecksum(&g);
+                 return ok;
+             },
+             "src"_a,
+             "Overwrite the grid's short name (truncated to the in-header buffer; "
+             "refreshes the checksum). Returns False if the name had to be truncated.")
+        .def("setGridClass",
+             [](GridData& g, GridClass gridClass) {
+                 g.mGridClass = gridClass;
+                 nanovdb::tools::updateChecksum(&g);
+             },
+             "gridClass"_a,
+             "Set the grid's GridClass (e.g. GridClass.LevelSet) in place and refresh "
+             "the checksum. Host twin of tools.cuda.setGridClass for device grids.")
+        .def("setTransform",
+             [](GridData& g, double voxelSize, const Vec3d& translation) {
+                 // Map::set only debug-asserts positivity, so validate here to
+                 // keep a singular / non-finite transform out of the header.
+                 requirePositiveFinite(voxelSize, "setTransform", "voxelSize");
+                 // Uniform scale + translation index->world map, then recompute the
+                 // world-space AABB from the index bbox under the new map (mirrors
+                 // how GridStats sets mWorldBBox).
+                 g.mMap.set(voxelSize, translation, 1.0);
+                 g.mVoxelSize = g.mMap.getVoxelSize();
+                 const CoordBBox& ib = g.indexBBox();
+                 g.mWorldBBox = CoordBBox(ib[0], ib[1].offsetBy(1)).transform(g.mMap);
+                 nanovdb::tools::updateChecksum(&g);
+             },
+             "voxelSize"_a, "translation"_a = Vec3d(0.0, 0.0, 0.0),
+             "Set a uniform-scale + translation index->world transform (voxelSize is "
+             "world units per voxel; translation is in world space), recompute the "
+             "world bounding box, and refresh the checksum. For non-uniform or rotated "
+             "maps, build the grid with the desired transform instead.")
         // Affine transforms (already member functions on GridData).
         .def("applyMap", nb::overload_cast<const Vec3f&>(&GridData::template applyMap<Vec3f>, nb::const_), "xyz"_a,
              "Transform an index-space point to world space using this grid's Map.")
@@ -464,7 +516,16 @@ template<typename BuildT> void defineNanoGrid(nb::module_& m, const char* name)
              nb::overload_cast<>(&NanoGrid<BuildT>::tree, nb::const_),
              nb::rv_policy::reference_internal,
              "Return the tree associated with this grid. Lifetime is "
-             "anchored to the grid (and therefore to the GridHandle).");
+             "anchored to the grid (and therefore to the GridHandle).")
+        .def("data_ptr",
+             [](const NanoGrid<BuildT>& grid) {
+                 return reinterpret_cast<uintptr_t>(&grid);
+             },
+             "Raw pointer to this grid as a Python int. The address is a HOST "
+             "pointer when the grid came from handle.grid(n) and a DEVICE "
+             "pointer when it came from handle.deviceGrid(n) — provenance is "
+             "the caller's responsibility, the grid object itself cannot tell "
+             "host from device.");
     // Grid<TreeT>::valueCount / pointCount are SFINAE-gated in C++ to the
     // index and Point BuildTs respectively — mirror that gating here.
     if constexpr (BuildTraits<BuildT>::is_index) {
@@ -1218,8 +1279,25 @@ NB_MODULE(nanovdb, m)
     defineHostGridHandle(m);
 
 #ifdef NANOVDB_USE_CUDA
-    defineDeviceBuffer(m);
-    defineDeviceGridHandle(m);
+    nb::module_ cudaModule = m.def_submodule("cuda");
+    cudaModule.doc() = "CUDA device buffers, the device GridHandle, and device infrastructure (mirrors nanovdb::cuda).";
+    defineDeviceBuffer(cudaModule);
+    // UnifiedBuffer (CUDA managed memory) + its GridHandle. The unified handle
+    // is the type returned by nanovdb.tools.cuda.DistributedPointsToGrid.
+    defineUnifiedBuffer(cudaModule);
+    defineUnifiedGridHandle(cudaModule);
+    defineDeviceGridHandle(cudaModule);
+    // Device NodeManager (DeviceNodeManagerHandle + createDeviceNodeManager)
+    // on nanovdb.cuda, alongside the device GridHandle.
+    defineDeviceNodeManager(cudaModule);
+    // Multi-GPU / device infrastructure (mirrors nanovdb::cuda): the device
+    // mesh, the device->stream map, and the CUB temp-storage pool primitives.
+    defineDeviceMesh(cudaModule);
+    defineDeviceStreamMap(cudaModule);
+    defineTempPool(cudaModule);
+    // Device VoxelBlockManager (mirrors nanovdb::tools::cuda) is registered on
+    // the existing nanovdb.tools.cuda submodule in PyTools.cc — NOT here — so
+    // the Python layout matches the C++ namespaces.
 #endif
 
     nb::module_ toolsModule = m.def_submodule("tools");

nanovdb/nanovdb/python/PyMath.cc

@@ -486,7 +486,7 @@ template<typename CoordT> void defineBBoxInteger(nb::module_& m, const char* nam
                     "Construct a cube of side dim voxels anchored at the min Coord.")
         .def_static("createCube", nb::overload_cast<ValueType, ValueType>(&math::BBox<CoordT>::createCube), "min"_a, "max"_a,
                     "Construct a cube spanning [min, max] in every axis.")
-        .def("is_divisible", &math::BBox<CoordT>::is_divisible,
+        .def("isDivisible", &math::BBox<CoordT>::is_divisible,
              "True iff this CoordBBox has more than one voxel in every axis.")
         .def("empty", &math::BBox<CoordT>::empty,
              "True iff this CoordBBox is empty (any min component > the matching max).")