NanoVDB: injectable CUDA memory resources — design & roadmap

[harrism]

Summary

NanoVDB's CUDA path hard-wires its allocations — grid storage via GridHandle<BufferT>, scratch via cuda::TempPool/cuda::DeviceResource, and raw cudaMallocAsync/cudaMallocHost in builders. This tracks a layered design to let downstream inject its own allocator wherever NanoVDB allocates, aligned with CCCL's cuda::mr / cuda::buffer model — without a hard CCCL dependency. New APIs land alongside the existing ones; the legacy dual buffer + handle surface are deprecated and removed a release or two later, so there is no hard break at introduction.

Why

• Inject your own allocator without forking headers — downstream currently forks NanoVDB headers to route through their pool (e.g. PyTorch's c10::cuda::CUDACachingAllocator; cf. Fork a small set of nanoVDB headers to share PyTorch's CUDA allocator openvdb/fvdb-core#655, which vendored three headers).
• Fix the pinned-host stall — the synchronous cudaMallocHost inside cuda::DeviceBuffer serializes streams in builder-heavy workloads.
• Share one pool across tensors, scratch, and grid buffers (caching; fewer fragmentation OOMs).
• Familiar + future-proof — method/type shapes match CCCL, so a resource_ref / cuda::buffer adapter is a thin shim later, while NanoVDB keeps its low CUDA floor.

Design — three pieces, each a minimal NanoVDB analog of a CCCL standard

1. Resource (allocator): two tiers matching CCCL's refinement. A synchronous resource exposes allocate(bytes, alignment)/deallocate(...) — no stream (cuda::PinnedResource is one). A stream-ordered resource adds allocate_async/deallocate_async(bytes, alignment, stream) and still provides the sync pair (typically thin delegates through the null stream), so is_async_resource<R> implies is_resource<R> — exactly cuda::mr's shape. cuda::DeviceResource provides all four; its legacy static allocateAsync/deallocateAsync are [[deprecated]]. Plus default_resource<R>() and the two detection traits. A custom allocator is a ~25-line struct.
2. cuda::Buffer<T,R> (storage) + cuda::BufferView<T> (view): one single-space, resource-aware, stream-ordered container modelled on CCCL's cuda::buffer<T> (the RMM device_buffer/device_uvector lineage). Buffer<T> typed, Buffer<std::byte> raw. Constructor order and shape follow cuda::buffer exactly — (stream, resource, count, noInit), the tag mandatory: there is no implicitly initializing count constructor (implicit fill of fresh memory is a hidden cost the allocate-then-overwrite pattern wastes). size() counts elements, size_bytes() bytes; a synchronous R yields a buffer with no stream API at all. Alongside it, BufferView<T>: a typed non-owning view with span semantics (element constness in T, trivially copyable, no resource, no stream) that satisfies the same static interface — so a GridHandle can wrap externally owned memory (an ONNX Runtime or Torch tensor) with zero copies. Owning and viewing do not share a type: the legacy mManaged owns-or-wraps flag is deleted, not re-spelled.
3. Single-space GridHandle (surface): one buffer per handle; grid() returns its pointer; host↔device is an explicit stream-carrying copyTo.

Rules that matter

• Stream ownership: a buffer retains its allocation stream and frees on it (not the null stream); the stream must outlive the buffer or be reset via a non-synchronizing setStream (the RMM cuda_stream_view contract). resize orders everything — including the free of the old block, whose last use is the prefix copy — on the passed stream.
• Resource ownership: cuda::Buffer holds the resource by value. A concrete resource that dispatches internally at runtime is a first-class fixed R (no CCCL needed); type-erased resource_ref comes via the optional adapter.
• Synchronous arenas (e.g. ONNX Runtime's) are wrapped as synchronous resources — never as an allocate_async facade. Contract: deallocate implies the memory is quiescent.
• Naming: types are CamelCase per OpenVDB style (cuda::Buffer, cuda::BufferView) — type names are aliasable, so the opt-in using Buffer = ::cuda::buffer<T> at CUDA ≥ 13.2 is preserved. Member names keep the standard spelling (data, size, size_bytes, allocate_async) because member matching is structural and cannot be aliased.
• CUDA graph capture: the async resource tier is capture-safe — stream-ordered allocation records as graph allocation/free nodes, and Buffer's async path performs no hidden synchronization or initialization (verified by a capture → instantiate → relaunch test). The sync tier can never be captured (it synchronizes), visibly: a sync-R buffer has no stream API.
• No hard CCCL dependency; an optional adapter bridges to cuda::mr / cuda::buffer / cuda::std::span.

Roadmap

• Step 1 — resource concept (NanoVDB: add stream-ordered async memory-resource seam (CUDA) #2231, merged): DeviceResource instance methods + default_resource + detection traits; new PinnedResource; TempPool routes through a resource instance and frees on its retained stream; PointsToGrid instance-injection seam. Additive. Follow-up in review: NanoVDB: encode points for any resource in PointsToGrid (CUDA) #2244 (point encoding for any resource).
• Step 2 — cuda::Buffer<T,R> + BufferView<T> + scratch retrofit (in progress): ship the container, the view, the synchronous is_resource trait and the sync/async dispatch (PR A); then convert the raw alloc/free in TopologyBuilder/PointsToGrid (and TempPool's bytes) onto the container — the "no raw allocation" cleanup — with a builder-coverage audit (PR B). Additive.
• Step 3 — single-space grid storage, via deprecation (end state: GridHandle<cuda::Buffer<std::byte,R>> directly — one buffer per handle, no dual surface, cross-space via copyTo; GridHandle over BufferView for externally owned blobs):
  • Introduce: ship the single-space handle path; re-implement the legacy cuda::DeviceBuffer internally as a composition of two cuda::Buffers (transparent — same API, gains the pinned-host fix); migrate NanoVDB's own internal uses + tests/examples off the DeviceBuffer name so the later deprecation fires only externally.
  • Deprecate: using DeviceBuffer [[deprecated]] = <impl> + [[deprecated]] on the dual GridHandle/NodeManager methods (deviceUpload/deviceDownload/deviceGrid/deviceData). Both still compile; warnings external only. Soak one or two releases.
  • Remove: delete the dual buffer + dual surface + hasDeviceDual; flip entry-point defaults to single-space. The break lands here, after the window, only for code that didn't migrate.
• Candidate — capacity-bounded, sync-free (graph-capturable) build path: graph capture prohibits host logic on device-computed values, and PointsToGrid reads count reductions back to size allocations and launch dimensions — so building cannot be captured regardless of allocator (cf. NVIDIA/warp Remove Boost UUID #1606, which reimplemented grid building privately for exactly this reason). Sketch: caller-supplied capacity bounds; one up-front allocation through the injected resource or into caller-owned memory via BufferView; launch dimensions from capacity; counts consumed on device; optional deferred readback. Requirements gathered from the Warp team: an overflow-clamped grid must remain safe to traverse from the root (orphaned leaves are acceptable — safety, not full well-formedness); point-mask support is required; a CPU counterpart is a plus (single-source maintenance); build performance must not regress; capacity-growth policy stays with applications. The memory seam is the prerequisite infrastructure; this is its own item.
• Step 4 — optional CCCL adapter (gated on availability): native cuda::mr / cuda::buffer interop, cuda::std::span conversions, and runtime (resource_ref) selection.

Migration note — custom ResourceT contract

Releases v12.1.0–v13.0.0 accepted a custom resource with static allocateAsync(bytes, alignment, stream) / deallocateAsync(...). Since #2231, builders call instance methods, and the concept now matches CCCL's refinement as described above. The static methods on cuda::DeviceResource are deprecated and will be removed after a deprecation window. Migrating a v13-era custom resource is mechanical: drop static, rename to the snake_case instance forms, add the two-line sync delegates (~8 lines total — see cuda/DeviceResource.h for the reference shape).

Downstream payoff (fvdb-core)

fvdb deletes its three forked headers and writes a small TorchAllocatorResource (forwarding to c10::cuda::CUDACachingAllocator), then uses PointsToGrid<…, TorchAllocatorResource> or GridHandle<cuda::Buffer<std::byte, TorchAllocatorResource>>; ONNX Runtime kernels wrap ORT-owned grid blobs zero-copy via GridHandle<BufferView<std::byte>>. No fork, no patch fragility.