NanoVDB: support single-space device buffers in GridHandle (CUDA)

GridHandle can now own a nanovdb::cuda::Buffer<std::byte, R>: a new
hasDeviceSingle buffer trait selects a constructor (implemented in
cuda/GridHandle.cuh) that parses the grid metadata through the device,
with every operation -- the validated header walk, metadata scratch
allocation through the buffer's resource, the copy kernel and the
readback -- ordered on the buffer's retained stream.
deviceData()/deviceGrid() work on such handles. The host accessors
(data, grid, gridData, gridMetaData) are SFINAE-removed for them,
since the handle owns no host-readable bytes; the read/write I/O
members and splitGrids/mergeGrids stay addressable (the python
bindings take write's address via overload_cast) and fail instead
with an explanatory static_assert when instantiated for a
single-space handle.

The device parse validates the whole grid chain (per-header bounds
checks against the allocation) before launching the metadata kernel,
so truncated buffers and forged mGridCount/mGridSize headers are
rejected with an exception instead of an out-of-bounds device read;
the pre-existing dual-space parse had the same hole and now shares the
validation helper. Its metadata scratch uses MallocResource so that
long-standing path keeps working on devices without memory-pool
support, and the dirty-flag scratch in splitGridHandles and
mergeGridHandles is a resource-aware buffer as well -- together
replacing the file's six raw cudaMalloc/mallocAsync sites, one of
which was missing its cudaCheck.

copy() dispatches on the traits: host buffers keep the memcpy path,
single-space buffers deep-copy device-to-device through their own
resource on the retained stream and reuse the host-resident metadata
instead of re-parsing. Buffer gains ElementType/ResourceType typedefs,
a resource() accessor, and a no-argument copy() for stream-ordered
resources that orders the copy on the retained stream.

Host accessibility is a property of the resource, not the element
type: the new is_host_accessible_resource trait detects a
HOST_ACCESSIBLE marker (declared by PinnedResource and forwarded by
ResourceRef and AsyncFromSync). A cuda::Buffer over a host-accessible
resource is rejected at handle scope with a named error for now --
GridHandle's host paths require an allocation interface cuda::Buffer
does not yet provide -- and non-byte element types fail loudly in the
single-space constructor. Both hasDeviceSingle and hasHostSingle are
detected rather than required, so pre-existing BufferTraits
specializations in and out of tree compile unchanged.

Part of #2232 (step 3).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Mark Harris <mharris@nvidia.com>

Author: harrism

nanovdb/nanovdb/GridHandle.h

@@ -98,7 +98,8 @@ namespace nanovdb {
 template<typename BufferT>
 struct BufferTraits
 {
-    static constexpr bool hasDeviceDual = false;
+    static constexpr bool hasDeviceDual   = false;
+    static constexpr bool hasDeviceSingle = false;
 };
 
 // ----------------------------> HostBuffer <--------------------------------------

nanovdb/nanovdb/HostBuffer.h

@@ -62,6 +62,13 @@ class Buffer : private detail::StreamHolder<is_async_resource<R>::value>
 
     static constexpr bool IsAsync = is_async_resource<R>::value;
 
+public:
+    /// @brief Element and resource types, for generic code that rebinds one
+    ///        or constructs sibling buffers over the same resource.
+    using ElementType  = T;
+    using ResourceType = R;
+
+private:
     R      mResource;
     T*     mData = nullptr;
     size_t mSize = 0; // element count
@@ -156,6 +163,11 @@ class Buffer : private detail::StreamHolder<is_async_resource<R>::value>
         return out;
     }
 
+    /// @brief Returns a deep copy of this buffer ordered on the retained
+    ///        stream, i.e. copy(this->stream()).
+    template<typename S = R, std::enable_if_t<is_async_resource<S>::value, int> = 0>
+    Buffer copy() const { return this->copy(this->stream()); }
+
     /// @brief Returns a deep copy of this buffer, allocated from a copy of the
     ///        synchronous resource.
     template<typename S = R, std::enable_if_t<!is_async_resource<S>::value && is_resource<S>::value, int> = 0>
@@ -255,6 +267,13 @@ class Buffer : private detail::StreamHolder<is_async_resource<R>::value>
     T*       data()       { return mData; }
     const T* data() const { return mData; }
 
+    /// @brief Returns a copy of the resource; for a ResourceRef this refers
+    ///        to the same underlying instance.
+    /// @note Requires R to be copy-constructible (the cuda::mr convention:
+    ///       resources are cheap handles). A resource that owns its pool by
+    ///       value hands the caller an independent copy of that pool.
+    R resource() const { return mResource; }
+
     /// @brief Returns the number of elements.
     size_t size() const { return mSize; }
 
@@ -402,6 +421,36 @@ class BufferView
 
 } // namespace cuda
 
+// Primary template defined in HostBuffer.h; declared here so this header
+// stays self-contained without pulling in the host-buffer machinery.
+template<typename BufferT>
+struct BufferTraits;
+
+/// @brief GridHandle support for the single-space cuda::Buffer: the buffer
+///        owns exactly one allocation, resident on the device, so the handle
+///        parses metadata through a device read and exposes only the device
+///        accessors. Requires byte-addressed storage.
+/// @note A buffer whose trait sets hasDeviceSingle must provide the interface
+///       the single-space GridHandle constructor consumes: ElementType and
+///       ResourceType typedefs, data(), size_bytes(), resource(), stream()
+///       (stream-ordered resources), and the (stream, resource, count, noInit)
+///       constructor shape.
+template<typename T, typename R>
+struct BufferTraits<cuda::Buffer<T, R>>
+{
+    static constexpr bool hasDeviceDual   = false;
+    // Device-resident storage; the byte-addressed requirement is enforced by
+    // the single-space GridHandle constructor, so trait queries stay
+    // answerable for any element type.
+    static constexpr bool hasDeviceSingle = !cuda::is_host_accessible_resource<R>::value;
+    // A buffer over a host-accessible resource (e.g. PinnedResource) is
+    // host-readable, but GridHandle's host paths also require the create()
+    // static interface and byte-count size semantics that cuda::Buffer does
+    // not provide -- GridHandle rejects such buffers with a named error until
+    // that adaptation lands.
+    static constexpr bool hasHostSingle   = cuda::is_host_accessible_resource<R>::value;
+};
+
 } // namespace nanovdb
 
 #endif // end of NANOVDB_CUDA_BUFFER_H_HAS_BEEN_INCLUDED

nanovdb/nanovdb/cuda/Buffer.h

@@ -123,6 +123,17 @@ struct is_async_resource<R, std::void_t<
     decltype(std::declval<R&>().deallocate(std::declval<void*>(), size_t{0}, size_t{0}))>>
     : std::true_type {};
 
+/// @brief Detection trait: @c is_host_accessible_resource<R>::value is true
+///        iff @c R declares `static constexpr bool HOST_ACCESSIBLE = true`,
+///        i.e. its allocations are mapped into the host address space (e.g.
+///        PinnedResource). Defaults to false: allocations are device-resident.
+/// @note Unlike the void_t detections above, this checks the member's VALUE:
+///       a resource declaring HOST_ACCESSIBLE = false stays device-resident.
+template<typename R, typename = void>
+struct is_host_accessible_resource : std::false_type {};
+template<typename R>
+struct is_host_accessible_resource<R, typename std::enable_if<bool(R::HOST_ACCESSIBLE)>::type> : std::true_type {};
+
 /// @brief Detection trait: @c is_resource<R>::value is true iff @c R models
 ///        the synchronous Resource concept, i.e. exposes
 ///        allocate(size_t, size_t) and deallocate(void*, size_t, size_t).
@@ -244,6 +255,9 @@ struct AsyncFromSync
 
     static constexpr size_t DEFAULT_ALIGNMENT = R::DEFAULT_ALIGNMENT;
 
+    /// @brief The adapter is host-accessible iff the adapted resource is.
+    static constexpr bool HOST_ACCESSIBLE = is_host_accessible_resource<R>::value;
+
     R resource;
 
     /// @brief Allocates through the synchronous resource; the result is valid
@@ -286,6 +300,9 @@ struct ResourceRef
 
     static constexpr size_t DEFAULT_ALIGNMENT = R::DEFAULT_ALIGNMENT;
 
+    /// @brief A reference is host-accessible iff the referenced resource is.
+    static constexpr bool HOST_ACCESSIBLE = is_host_accessible_resource<R>::value;
+
     /// @brief Constructs a ref borrowing @c resource.
     /// @param resource resource to allocate from; must outlive this ref
     ResourceRef(R& resource) : mResource(&resource) {}

nanovdb/nanovdb/cuda/DeviceResource.h

@@ -17,6 +17,7 @@
 #ifndef NANOVDB_CUDA_GRIDHANDLE_CUH_HAS_BEEN_INCLUDED
 #define NANOVDB_CUDA_GRIDHANDLE_CUH_HAS_BEEN_INCLUDED
 
+#include <nanovdb/cuda/Buffer.h>// for the resource-aware scratch buffers below
 #include <nanovdb/cuda/DeviceBuffer.h>// required for instantiation of move c-tor of GridHandle
 #include <nanovdb/tools/cuda/GridChecksum.cuh>// for cuda::updateChecksum
 #include <nanovdb/GridHandle.h>
@@ -43,6 +44,35 @@ static __global__ void updateGridCount(GridData *d_data, uint32_t gridIndex, uin
     }
 }
 
+/// @brief Walks the grid chain with per-header device-to-host reads, checking
+///        that every header and every grid span lies inside the allocation, so
+///        the device-side metadata walk that follows can never read out of
+///        bounds from a truncated buffer or a forged header.
+/// @return the validated grid count
+inline uint32_t validGridChainCount(const GridData *d_head, uint64_t bytes, cudaStream_t stream)
+{
+    uint64_t offset = 0;
+    uint32_t count = 0, expected = 0;
+    GridData tmp;
+    do {
+        if (offset + sizeof(GridData) > bytes)
+            throw std::runtime_error("GridHandle: grid chain exceeds the device buffer (truncated or corrupt grid data)");
+        cudaCheck(cudaMemcpyAsync(&tmp, util::PtrAdd<GridData>(d_head, offset), sizeof(GridData), cudaMemcpyDeviceToHost, stream));
+        cudaCheck(cudaStreamSynchronize(stream));
+        if (!tmp.isValid()) throw std::runtime_error("GridHandle was constructed with an invalid device buffer");
+        if (count == 0) {
+            expected = tmp.mGridCount;
+            if (expected == 0) throw std::runtime_error("GridHandle: device buffer contains no grids");
+        }
+        if (tmp.mGridIndex != count || tmp.mGridCount != expected)
+            throw std::runtime_error("GridHandle: inconsistent grid index/count in the device buffer's grid chain");
+        if (tmp.mGridSize < sizeof(GridData) || tmp.mGridSize > bytes - offset)
+            throw std::runtime_error("GridHandle: grid size field exceeds the device buffer (truncated or corrupt grid data)");
+        offset += tmp.mGridSize;
+    } while (++count < expected);
+    return expected;
+}
+
 }// namespace detail
 
 template<typename BufferT, template <class, class...> class VectorT = std::vector>
@@ -52,8 +82,9 @@ splitGridHandles(const GridHandle<BufferT> &handle, const BufferT* other = nullp
     const void *ptr = handle.deviceData();
     if (ptr == nullptr) return VectorT<GridHandle<BufferT>>();
     VectorT<GridHandle<BufferT>> handles(handle.gridCount());
-    bool dirty, *d_dirty;// use this to check if the checksum needs to be recomputed
-    cudaCheck(util::cuda::mallocAsync((void**)&d_dirty, sizeof(bool), stream));
+    bool dirty;// set when the checksum needs to be recomputed
+    Buffer<bool, DeviceResource> dirtyBuf(stream, 1, noInit);
+    bool *d_dirty = dirtyBuf.data();
     int device = util::cuda::currentDevice();
     for (uint32_t n=0; n<handle.gridCount(); ++n) {
         auto buffer = BufferT::create(handle.gridSize(n), other, device, stream);
@@ -68,7 +99,6 @@ splitGridHandles(const GridHandle<BufferT> &handle, const BufferT* other = nullp
         handles[n] = nanovdb::GridHandle<BufferT>(std::move(buffer));
         ptr = util::PtrAdd(ptr, handle.gridSize(n));
     }
-    cudaCheck(util::cuda::freeAsync(d_dirty, stream));
     return handles;
 }// cuda::splitGridHandles
 
@@ -85,8 +115,9 @@ mergeGridHandles(const VectorT<GridHandle<BufferT>> &handles, const BufferT* oth
     int device = util::cuda::currentDevice();
     auto buffer = BufferT::create(size, other, device, stream);
     void *dst = buffer.deviceData();
-    bool dirty, *d_dirty;// use this to check if the checksum needs to be recomputed
-    cudaCheck(util::cuda::mallocAsync((void**)&d_dirty, sizeof(bool), stream));
+    bool dirty;// set when the checksum needs to be recomputed
+    Buffer<bool, DeviceResource> dirtyBuf(stream, 1, noInit);
+    bool *d_dirty = dirtyBuf.data();
     for (auto &h : handles) {
         const void *src = h.deviceData();
         for (uint32_t n=0; n<h.gridCount(); ++n) {
@@ -101,7 +132,6 @@ mergeGridHandles(const VectorT<GridHandle<BufferT>> &handles, const BufferT* oth
             src = util::PtrAdd(src, h.gridSize(n));
         }
     }
-    cudaCheck(util::cuda::freeAsync(d_dirty, stream));
     return GridHandle<BufferT>(std::move(buffer));
 }// cuda::mergeGridHandles
 
@@ -122,28 +152,57 @@ mergeDeviceGrids(const VectorT<GridHandle<BufferT>> &handles, const BufferT* oth
 template<typename BufferT>
 template<typename T, typename util::enable_if<BufferTraits<T>::hasDeviceDual, int>::type>
 GridHandle<BufferT>::GridHandle(T&& buffer)
+    : mBuffer(std::move(buffer))
 {
     static_assert(util::is_same<T,BufferT>::value, "Expected U==BufferT");
-    mBuffer = std::move(buffer);
     if (auto *data = reinterpret_cast<const GridData*>(mBuffer.data())) {
         if (!data->isValid()) throw std::runtime_error("GridHandle was constructed with an invalid host buffer");
         mMetaData.resize(data->mGridCount);
         cpyGridHandleMeta(data, mMetaData.data());
     } else {
         if (auto *d_data = reinterpret_cast<const GridData*>(mBuffer.deviceData())) {
-            GridData tmp;
-            cudaCheck(cudaMemcpy(&tmp, d_data, sizeof(GridData), cudaMemcpyDeviceToHost));
-            if (!tmp.isValid()) throw std::runtime_error("GridHandle was constructed with an invalid device buffer");
-            GridHandleMetaData *d_metaData;
-            cudaMalloc((void**)&d_metaData, tmp.mGridCount*sizeof(GridHandleMetaData));
-            cuda::detail::cpyGridHandleMeta<<<1,1>>>(d_data, d_metaData);
-            mMetaData.resize(tmp.mGridCount);
-            cudaCheck(cudaMemcpy(mMetaData.data(), d_metaData,tmp.mGridCount*sizeof(GridHandleMetaData), cudaMemcpyDeviceToHost));
-            cudaCheck(cudaFree(d_metaData));
+            const uint32_t count = cuda::detail::validGridChainCount(d_data, mBuffer.size(), cudaStream_t(0));
+            // MallocResource: plain cudaMalloc, so this long-standing parse path
+            // keeps working on devices without memory-pool support.
+            cuda::Buffer<GridHandleMetaData, cuda::MallocResource> scratch(count, cuda::noInit);
+            cuda::detail::cpyGridHandleMeta<<<1,1>>>(d_data, scratch.data());
+            cudaCheckError();
+            mMetaData.resize(count);
+            cudaCheck(cudaMemcpy(mMetaData.data(), scratch.data(), count*sizeof(GridHandleMetaData), cudaMemcpyDeviceToHost));
         }
     }
 }// GridHandle(T&& buffer)
 
+// move constructor from a single-space device buffer: all device work runs on
+// the buffer's retained stream (or the default stream for a synchronous
+// resource), and the metadata scratch allocates through the buffer's resource.
+template<typename BufferT>
+template<typename T, typename util::enable_if<BufferHasDeviceSingle<T>::value, int>::type, typename>
+GridHandle<BufferT>::GridHandle(T&& buffer)
+    : mBuffer(std::move(buffer))
+{
+    static_assert(util::is_same<T,BufferT>::value, "Expected U==BufferT");
+    static_assert(sizeof(typename BufferT::ElementType) == 1,
+                  "GridHandle requires byte-addressed single-space storage, e.g. cuda::Buffer<std::byte, R>");
+    using ResourceT = typename BufferT::ResourceType;
+    if (const GridData *d_data = reinterpret_cast<const GridData*>(mBuffer.data())) {
+        constexpr bool isAsync = cuda::is_async_resource<ResourceT>::value;
+        cudaStream_t stream = cudaStream_t(0);
+        if constexpr (isAsync) stream = mBuffer.stream();
+        const uint32_t count = cuda::detail::validGridChainCount(d_data, mBuffer.size_bytes(), stream);
+        using ScratchT = cuda::Buffer<GridHandleMetaData, ResourceT>;
+        ScratchT scratch = [&]{
+            if constexpr (isAsync) return ScratchT(stream, mBuffer.resource(), count, cuda::noInit);
+            else                   return ScratchT(mBuffer.resource(), count, cuda::noInit);
+        }();
+        cuda::detail::cpyGridHandleMeta<<<1, 1, 0, stream>>>(d_data, scratch.data());
+        cudaCheckError();
+        mMetaData.resize(count);
+        cudaCheck(cudaMemcpyAsync(mMetaData.data(), scratch.data(), count*sizeof(GridHandleMetaData), cudaMemcpyDeviceToHost, stream));
+        cudaCheck(cudaStreamSynchronize(stream));
+    }
+}// GridHandle(T&& buffer) for single-space device buffers
+
 // Dummy function that ensures instantiation of the move-constructor above when BufferT=cuda::DeviceBuffer
 namespace {auto __dummy(){return GridHandle<cuda::DeviceBuffer>(std::move(cuda::DeviceBuffer()));}}
 

nanovdb/nanovdb/cuda/GridHandle.cuh

@@ -30,6 +30,11 @@ class PinnedResource
     // alignment is always satisfied; 256 is the nominal default advertised here.
     static constexpr size_t DEFAULT_ALIGNMENT = 256;
 
+    /// @brief Pinned allocations are mapped into the host address space, so
+    ///        containers over this resource remain host-readable (detected by
+    ///        nanovdb::cuda::is_host_accessible_resource).
+    static constexpr bool HOST_ACCESSIBLE = true;
+
     /// @brief Synchronous allocation of page-locked host memory.
     /// @param bytes number of bytes to allocate
     /// @param alignment requested alignment (ignored; cudaMallocHost is page-aligned)