Skip to content

NanoVDB: injectable CUDA memory resources — design & roadmap #2232

Description

@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 where a dual buffer is actually constructed host-side. Note this is not every DeviceBuffer use: init allocates host or device but never both, so buffers created with a real device id never touch cudaMallocHost (see the correction under B1).
  • 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 (revised — decision of record): cuda::Buffer holds the resource by value, matching cuda::buffer — but ownership semantics are selected by what is placed in that slot, completing the CCCL model rather than adopting half of it (cuda::buffer owns an any_resource; CCCL's borrowing and sharing tiers are resource_ref and shared_resource, and its constructor's static_assert points users at them). The tiers here: a concrete resource in the slot is owned as a copy (so it must be a cheap-to-copy handle, stateless or pointer-to-state); ResourceRef<R> (NanoVDB: SyncFromAsync, ResourceRef, and resource seams for MeshToGrid and TempPool (CUDA) #2269) borrows — a non-owning, dependency-free static analog of cuda::mr::resource_ref that is itself a resource, with enable_if-gated async members so a ref over a synchronous resource does not misreport its tier, and pointer-identity equality; a refcounted SharedResource<R> analog is deferred until the Python-bindings need materializes. Non-owning containers (TempPool, and the builders' ResourceT* members over time) compose with Buffer via ResourceRef. History: an earlier decision ("reference-at-API, pointer-as-member", pre-cuda::buffer reshape) was superseded by the by-value rule without reconciling the artifacts built on it — TempPool's pointer contract and the stateful test resources — which is what the first TempPool conversion attempt tripped over. Type-erased CCCL resource_ref interop remains step 4.
  • Synchronous arenas (e.g. ONNX Runtime's) are wrapped as synchronous resources — never as a silent allocate_async facade (a facade misrepresents its semantics and hands multi-stream callers unexpected serialization). Contract: deallocate implies the memory is quiescent. The explicit lift is AsyncFromSync<R> (NanoVDB: run the builders on a synchronous memory resource (CUDA) #2272), the analog of CCCL's synchronous_resource_adapter: allocate_async forwards (synchronously allocated memory is valid on every stream — stronger than stream-ordering requires) and deallocate_async synchronizes the stream first, making the quiescence contract hold; the serialization cost is documented at the type and chosen by the caller.
  • 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 — NanoVDB: add cuda::Buffer, cuda::BufferView, and the synchronous resource concept (CUDA) #2251, merged); then convert the raw alloc/free onto the container — the "no raw allocation" cleanup — with a builder-coverage audit. The retrofit was originally scoped as a single PR B; a scoping pass found the three targets differ enough in risk that they land separately:

    • B1 — TopologyBuilder (NanoVDB: inject a resource into TopologyBuilder scratch, and align cuda::Buffer with cuda::buffer (CUDA) #2268): 8 of its 10 cuda::DeviceBuffer members are device-only (zero host .data() uses), as are 3 function-local buffers. Convert those 11 to Buffer<std::byte,R> and add a defaulted ResourceT parameter; only DilateGrid and MergeGrids instantiate TopologyBuilder, so existing callers are unaffected. mProcessedRoot and mData are genuinely dual-space and stay until Step 3. Carries the builder-coverage audit table. Correction: an earlier revision of this entry described B1 as fixing the pinned-host stall. It does not. DeviceBuffer::init allocates host memory or device memory, never both, and TopologyBuilder always passed a real device id, so no cudaMallocHost was ever on this path. Measured on dilate at 1k/20k/200k points: no difference beyond run-to-run noise. B1's value is injectability (the last builder without the resource seam) and scope-based ownership, not speed. The pinned-host stall, where it exists, is on GridHandle-side paths that construct host-side buffers, and belongs to Step 3.

    • B2 — resource ergonomics + the remaining seams (NanoVDB: SyncFromAsync, ResourceRef, and resource seams for MeshToGrid and TempPool (CUDA) #2269): SyncFromAsync<Derived> CRTP mixin (a custom resource is two methods rather than four, with the mandatory synchronize in one audited place — its first users were the two test resources in TestMemoryResource, which turned out never to have modelled the concept); a ResourceT seam for MeshToGrid, the last builder on a hard-wired DeviceResource; ResourceRef<R> per the revised ownership rule above; and TempPool's bytes onto Buffer<std::byte, ResourceRef<R>> — same pointer contract and stream retention, block freed by ownership (discard-on-growth via destroy(stream) + move-assignment; the pool keeps a size_t mirror because cub's two-pass API wants a mutable size_t&). Buffer::swap and the destroy/set_stream alignment landed in NanoVDB: inject a resource into TopologyBuilder scratch, and align cuda::Buffer with cuda::buffer (CUDA) #2268.

    • B3 — PointsToGrid: 48 sites, and the only target that is not mechanical. mData.* are device-visible raw pointers uploaded at two separate points; std::swap(d_indx, mData.d_indx) moves ownership between a local and one of those aliased fields after the first upload; a goto retry loop over voxel density makes lifetimes non-lexical; and allocations are freed across three exit paths ~500 lines apart. Wants a written ownership design plus a separate goto-to-while prep commit before the RAII change. Sequenced after NanoVDB: encode points for any resource in PointsToGrid (CUDA) #2244, which edits the same function.

    • B4 — builders on a synchronous resource (NanoVDB: run the builders on a synchronous memory resource (CUDA) #2272): MallocResource (synchronous cudaMalloc/cudaFree, works on pool-less devices) + AsyncFromSync<R> per the rule above. Closes the scratch half of the acceptance criterion as pure library code — PointsToGrid<BuildT, AsyncFromSync<MallocResource>>, no build flags. The grid handle's buffer still needs NANOVDB_USE_SYNC_CUDA_MALLOC until Step 3 delivers GridHandle<cuda::Buffer<std::byte,R>>. Tested by driving PointsToGrid end-to-end on an injected stateful synchronous resource with balanced accounting.

    Splitting matters more than usual here because NanoVDB: please run the CUDA unit tests on a GPU in CI #2264 means no CUDA test runs on a GPU in CI, so a mistake in the intricate conversion would land unverified by review. Additive throughout. Acceptance criterion (from Issue with DeviceResource (cudaMallocAsync) on vGPU without memory pool support #2255): NanoVDB builds and runs on a pool-less vGPU (cudaDevAttrMemoryPoolsSupported == 0, e.g. AWS g6f) via an injected synchronous cudaMalloc-backed resource; interim coverage via the NANOVDB_USE_SYNC_CUDA_MALLOC macro (the caller's explicit opt-in). The util::cuda wrappers deliberately do not silently fall back: on a pool-less device they fail with an actionable diagnostic (NanoVDB: fail with an actionable diagnostic when CUDA memory pools are unavailable #2256) rather than making an async resource misrepresent its semantics; CI detects the capability with a cudaDevAttrMemoryPoolsSupported probe and sets the macro.

  • 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.
    • Direction of adaptation (rule of record): cuda::Buffer is the type the legacy buffers are being replaced by, so it keeps the interface we want to live with — the cuda::buffer-aligned spelling — and GridHandle/NodeManager adapt to it. Not the reverse: legacy names are not accreted onto cuda::Buffer merely because HostBuffer/DeviceBuffer/UnifiedBuffer carry them, since those three are being deleted. Concretely, GridHandle should move to destroy() and empty() rather than cuda::Buffer gaining clear() and isEmpty() permanently. Where a shim is genuinely needed to survive the deprecation window it is marked transitional at the declaration, naming what replaces it and when it goes — as cuda::Buffer::clear is in NanoVDB: inject a resource into TopologyBuilder scratch, and align cuda::Buffer with cuda::buffer (CUDA) #2268. This applies to every spelling the window might otherwise drag across: isEmpty, deviceData, deviceUpload, deviceDownload, create.
    • Known gap: GridHandle::copy() calls mBuffer.isEmpty(), which cuda::Buffer and BufferView do not provide (they have empty()). It is latent today because copy() is only instantiated on use, but it blocks the end state. Per the rule above the fix is in GridHandle, not in cuda::Buffer — and it needs more than a rename regardless, since copy() also calls OtherBufferT::create(...) and performs a host-side std::memcpy over the buffer contents, neither of which is valid for device memory.
  • 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions