Add GridStorage, TorchDeviceResource and device-handle utilities (#770 step 3) - #786
Merged
Merged
Conversation
…nvdb#770 step 3) The primitives the grid-storage flip (openvdb#770 steps 4-5) is built on. Nothing uses them yet; GridBatchData still holds GridHandle<TorchDeviceBuffer>. - TorchDeviceResource: a storage resource bound to one torch device (CUDA:i or PrivateUse1). NanoVDB picks host vs. device storage per buffer *type*, so the device has to travel in the resource: a prototype buffer's resource() is what createDeviceStorage and cuda::copyTo allocate from, which makes every allocation device-correct without a caller-side guard. The block is keyed to the buffer's retained stream (the TorchResource policy), not to torch's current stream as openvdb#773's TorchStorageResource does for the legacy dual-space buffer: a cuda::Buffer orders its writes and its free on the retained stream, so keying the block there makes torch's reuse rule and NanoVDB's ordering model agree with no events. Codex found the alternative wrong: a buffer retained on B but keyed to the current stream A is freed to A, which nothing that ordered B was protecting. The obligation this leaves is that the retained stream outlive the block (torch pool streams and the legacy stream do); DistributedPointsToGrid's mesh streams do not, which step 5 must handle at that builder. And the retained stream must not change after allocation (no set_stream/resize to another stream), since torch frees to the keyed stream whatever stream the free names; GridStorage copies instead. The synchronous allocate/deallocate are spelled out so their synchronization runs under the bound device. PrivateUse1 allocates through its registered allocator. DeviceGridBuffer = cuda::Buffer<std::byte, TorchDeviceResource>. - GridStorage: variant<GridHandle<HostBuffer>, GridHandle<DeviceGridBuffer>> plus a torch::Device. Everything about the grids (count, sizes, offsets, types, typed pointers) is answered from adopted metadata without touching device memory. to(device, stream) is cuda::copyTo with the ordering the design review asked for: the copy stream waits on the source's retained stream, and for a device destination the source's stream waits on the copy, so the source may be destroyed as soon as to() returns even with allocator caching handing its block straight to the next tenant. Across CUDA devices the copy is a peer copy (torch's expandable-segment and cudaMallocAsync backends do not support a plain memcpy between devices without peer access), and every event is recorded under its stream's device. Empty storage moves without allocating and keeps the requested device, in its buffer's resource as well as in device(), where copyTo alone would default-construct a handle on the current device. deviceProto(device, stream) is the zero-byte prototype that replaces the zero-size TorchDeviceBuffer "guide". The header is host-includable. - GridHeaderUtils.h: normalizeStandaloneGridHeader, the one spelling of the header writes every hand-wrapped grid needs now that GridHandle validates the chain at construction (index 0 of 1, size, no blind data with offset = size per GridData::init, checksum disabled). The six existing sites migrate in step 4. - DeviceGridHandleUtils.cuh: the updateGridCountAndZeroChecksum kernel hoisted out of MakeContiguous/ConcatenateGrids; makeDeviceHandleFromLayout for callers that know their layout; mergeDeviceGridHandles, the single-space replacement for nanovdb::cuda::mergeGridHandles with one allocation, one memcpy and one header kernel per grid, host-assembled metadata and no per-grid synchronization, and per-source stream ordering before and after the copy. - GridStorageTest: twelve gtests, including the three contract tests from the openvdb#770 design review (source destroyed right after an asynchronous device copy, with caching, held deterministic by a spin kernel; empty transfer preserving the destination device; explicit stream zero while a non-default torch stream is current), a byte-for-byte comparison of mergeDeviceGridHandles against nanovdb::mergeGrids, and a check that a freed block is handed back on its retained stream and not on torch's current one. The lifetime test uses a source retained on a stream that is not the current one, the case the first draft got wrong. Reviewed with codex before opening; its first pass found the keying error above, an event recorded under the wrong device for cross-device copies, the missing peer copy, empty handles carrying the current device's resource, an overflowable layout check, and the inherited synchronous allocate synchronizing the wrong device. Its second pass found the cross-device wait enqueued under the signaler's device (a null waiter then named the wrong GPU's default stream), the PrivateUse1 synchronize fallback with the same null-stream problem, empty transfers dropping the requested stream, and the set_stream/resize constraint above. Its third pass found the PrivateUse1 fallback skipped by the same-stream shortcut and the merge's copies and kernels running without the prototype's device current. All folded in. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Mark Harris <mharris@nvidia.com>
swahtz
reviewed
Sep 15, 2026
Contributor
There was a problem hiding this comment.
Reviewed at f28cfb3 against the #770 plan and its comment thread. Two items look like they should land before step 4 builds on these primitives (the PrivateUse1 host-source sync and the 2-arg constructor). The rest are nits or optional cleanups, marked as such. Points that belong to later steps rather than this PR (checksum policy for the merge helper, cross-stream readers, PrivateUse1 ordering cost, the long-name flag in the normalizer) are on #770 instead.
…ng, strict layout swahtz's seven threads, all taken. - GridStorage::to, host to PrivateUse1: cuda::copyTo does not synchronize for a device destination and PrivateUse1 storage hands its unified pointer to the host immediately, so hostGridAt() could read pre-copy bytes. The copy stream is synchronized for that destination; the comment is narrowed to the CUDA case. - GridStorage(DeviceHandle&&, device): an empty handle carries the resource it was default-constructed with (the current device, e.g. cuda::copyTo's result for a zero-byte source). It is now re-homed onto the named device so the buffer agrees with device() for anyone who re-wraps it or uses it as a prototype, the same rule empty() and to() already applied. - makeDeviceHandleFromLayout requires what the rest of the stack can consume: grids packed end to end from offset 0 at NANOVDB_DATA_ALIGNMENT, and no metadata iff the buffer is empty. A gap, a misaligned offset or metadata-less bytes now throw instead of yielding a handle the chain parse rejects or the header kernel faults on. Tests cover the gap and the empty-metadata cases. - mergeDeviceGridHandles's comment now says exactly who does what with checksums: the CUDA MakeContiguous/ConcatenateGrids kernel disables them; nanovdb::mergeGrids, the CPU sides of those ops and the replaced nanovdb::cuda::mergeGridHandles update an existing one. It no longer claims parity with the replaced path; the step-5 choice is on openvdb#770. - The block-reuse test skips unless the native caching allocator is active, since only it keys free blocks by stream; under backend:cudaMallocAsync opportunistic reuse could hand the block to another stream and fail the test for a correct resource. - orderStreamAfter was tried on c10::cuda::CUDAEvent as suggested and reverted: c10's stream conversion rejects the cudaStreamLegacy and cudaStreamPerThread sentinels, which raw CUDA accepts. The raw version keeps the same three guards (record under the signaler's device, wait under the waiter's, destroy under the signaler's). - hostGridAt/deviceGridAt delegate to the handles' own grid<T>() / deviceGrid<T>() per variant arm, and one metaAt(i) replaces the three duplicated range checks. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Mark Harris <mharris@nvidia.com>
swahtz
approved these changes
Sep 15, 2026
swahtz
left a comment
swahtz
left a comment
Contributor
There was a problem hiding this comment.
Looks good, thanks for addressing my comments
swahtz
added
the
core library
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Step 3 of #770: the primitives the grid-storage flip is built on. Nothing uses them yet.
GridBatchDatastill holdsGridHandle<TorchDeviceBuffer>; steps 4 and 5 move storage and the builders onto these. Reviewing this on its own keeps the flip PR to mechanical call-site changes.What is added
TorchDeviceResource.hCUDA:ior PrivateUse1). NanoVDB picks host vs. device storage per buffer type, so the device travels in the resource: a prototype buffer'sresource()is whatcreateDeviceStorageandcuda::copyToallocate from. The block is keyed to the buffer's retained stream (theTorchResourcepolicy): acuda::Bufferorders its writes and its free on that stream, so torch's reuse rule and NanoVDB's ordering model agree with no events. This is deliberately notTorchStorageResource's current-stream policy from #773, which is right for the dual-spaceTorchDeviceBuffer(no retained stream, unordered free) and wrong here; see the header comment. Obligations: the retained stream must outlive the block and must not be changed after allocation (set_stream/resizeto another stream);GridStoragecopies instead.DeviceGridBuffer = cuda::Buffer<std::byte, TorchDeviceResource>.GridStorage.{h,cu}variant<GridHandle<HostBuffer>, GridHandle<DeviceGridBuffer>>+torch::Device. Count, sizes, offsets, types and typed grid pointers come from adopted metadata; no device reads.to(device, stream)iscuda::copyTowith explicit ordering (below).empty(device),deviceProto(device, stream)(replaces the zero-sizeTorchDeviceBufferguide). Host-includable.GridHeaderUtils.hnormalizeStandaloneGridHeader: the one spelling of the header writes every hand-wrapped grid needs now thatGridHandlevalidates the chain (index 0 of 1, size, no blind data with offset = size, checksum disabled). The six existing sites migrate in step 4.DeviceGridHandleUtils.cuhupdateGridCountAndZeroChecksumhoisted fromMakeContiguous/ConcatenateGrids;makeDeviceHandleFromLayout;mergeDeviceGridHandles, the single-space replacement fornanovdb::cuda::mergeGridHandles: one allocation, one memcpy + one header kernel per grid, host-assembled metadata, no per-grid synchronization.GridStorageTest.cuStream contract of
GridStorage::tostream. If the source is device storage,streamfirst waits on the source's retained stream, so the source's writers are complete.to()returns is therefore safe even when torch's caching allocator hands its block straight to the next allocation on that stream. This is the failure the Finish adopting NanoVDB's injectable CUDA memory resources (openvdb #2232): move grid storage off the dual-space TorchDeviceBuffer #770 design review reproduced with a probe;SourceMayBeDestroyedRightAfterAsyncCopyreproduces it deterministically (the copy stream is held open by a spin kernel, the source is destroyed, its block is re-taken and overwritten on the source stream) and passes.copyTosynchronizes), so no post-copy ordering is needed there.cudaMemcpyPeerAsync(torch's expandable-segment and cudaMallocAsync backends do not support a plain memcpy between devices without peer access). Events are recorded under their stream's device.copyToalone returns a default-constructed handle, whose resource is the current device (design review item 2).Tests
HostAccessorsAnswerFromMetadata,EmptyStorageKeepsItsDevice,DeviceProtoCarriesDeviceAndStream,RoundTripThroughCudaPreservesBytesAndMetadata,ExplicitStreamZeroWhileNonDefaultStreamIsCurrent,SourceMayBeDestroyedRightAfterAsyncCopy,MergeDeviceGridHandlesMatchesHostMerge(byte-for-byte againstnanovdb::mergeGrids, including an empty source in the middle, then re-parsed from raw bytes),MakeDeviceHandleFromLayoutChecksExtent,NormalizeStandaloneGridHeaderMakesAGridWrappable(the un-normalized copy is rejected byGridHandle; the normalized one wraps),StorageBlockReturnsToItsRetainedStream,SynchronousResourceApiWorksUnderAnyCurrentDevice. The lifetime test uses a source retained on a non-current stream, the case that distinguishes the two keying policies.Verification
./build.sh install gtestswith-Werror=all-warnings: clean. ctest 49/49 (the new suite 12/12 on three consecutive runs). The nine pytest suites touched by the epic: unchanged. clang-format: clean. Codex-reviewed before opening; three passes found twelve real defects, all folded in and listed in the commit message.For #770
Step 5's
DistributedPointsToGridpath creates storage on aDeviceMeshstream that dies with the mesh. Under retained-stream keying that is an orphaned block; it has to be solved at that builder (a long-lived mesh, or re-homing the storage after the builder's own synchronization), not in the resource.🤖 Generated with Claude Code