Proposal: ⚫ ONNX Runtime support for fVDB operations

[swahtz]

ONNX Runtime Support for fVDB Sparse Voxel Operations

Motivation

fVDB operations currently exist only within the PyTorch ecosystem. Users who want to deploy fVDB-based models (neural radiance fields, 3D reconstruction, sparse convolution networks, etc.) in ONNX Runtime -- for cross-language inference (C++, C#, Java, JS), hardware-specific execution providers, or standardized model interchange -- have no path to do so.

The core challenge is that ONNX graphs transport only tensors between nodes, while fVDB operations pass two non-tensor types: GridBatchData (a frozen C++ struct wrapping NanoVDB OnIndexGrid data in a contiguous byte buffer) and JaggedTensor (variable-length batched data). Grid topology is dynamic -- constructed at inference time from input data, not baked into the model.

Design Overview

The approach is split into two phases: (1) a custom ONNX operator domain with tensor-based representations of fVDB types, and (2) an optional ONNX Runtime Execution Provider plugin for optimized execution.

---

Phase 1: Custom ONNX Operator Domain (fvdb)

1.1 Tensor Representations

GridBatchData as a tensor bundle ("GridBatch Bundle")

GridBatchData (src/fvdb/detail/GridBatchData.h) is an immutable struct holding sparse voxel grid topology. Its NanoVDB data is stored as a single contiguous byte buffer via mGridHdl, a shared_ptr<nanovdb::GridHandle<TorchDeviceBuffer>> backed by a raw uint8_t* + size (src/fvdb/detail/TorchDeviceBuffer.h). The struct is decomposed into a group of tensors that always travel together through the ONNX graph:

Tensor	Dtype	Shape	Source in GridBatchData
grid_blob	uint8	[N] (dynamic)	mGridHdl->data() -- the raw NanoVDB buffer containing all grids packed sequentially
grid_byte_offsets	int64	[B]	Per-grid mCumBytes from GridMetadata
voxel_sizes	float64	[B, 3]	Per-grid mVoxelSize from GridMetadata
origins	float64	[B, 3]	Per-grid voxel origins via GridMetadata::voxelOrigin()
leaf_batch_indices	int32	[L] (dynamic)	mLeafBatchIndices
batch_offsets	int64	[B+1]	mBatchOffsets
list_indices	int32	[M] (dynamic)	mListIndices

Individual grids within the blob are accessed via pointer arithmetic into the contiguous buffer, as in GridBatchData::Accessor::grid():

const nanovdb::OnIndexGrid *grid(int64_t bi) const {
    return reinterpret_cast<const nanovdb::OnIndexGrid *>(
        reinterpret_cast<const char *>(mGridPtr) + mMetadata[bi].mCumBytes);
}

Alignment: NanoVDB requires 32-byte alignment (NanoVDB.h ~L67-78). CUDA allocators satisfy this (typically 256B+). CPU allocators may not -- custom op implementations must validate or enforce alignment.

JaggedTensor as a tensor bundle ("JaggedTensor Bundle")

A JaggedTensor (src/fvdb/JaggedTensor.h) is decomposed into its constituent tensors:

Tensor	Dtype	Shape	Source
jdata	varies	[N, *esizes]	mData -- packed values
joffsets	int64	[T+1]	mOffsets -- CSR boundaries
jidx	int32	[N]	mBatchIdx -- per-element batch index
jlidx	int32	[T, ldim]	mListIdx -- list-of-lists indexing

1.2 Custom Operator Schemas

Each ONNX custom op in the fvdb domain maps 1:1 to a C++ free function bound in GridBatchOps.cpp (src/python/GridBatchOps.cpp). The fvdb.functional module (~120 exports in fvdb/functional/__init__.py) serves as the definitive op catalog.

Grid Construction Ops (produce a GridBatch Bundle from input tensors):

ONNX Custom Op	GridBatchOps.cpp binding	fvdb.functional entry point
fvdb.GridFromPoints	create_from_points -> ops::buildGridFromPoints	gridbatch_from_points / grid_from_points
fvdb.GridFromIJK	create_from_ijk -> ops::createNanoGridFromIJK	gridbatch_from_ijk / grid_from_ijk
fvdb.GridFromMesh	create_from_mesh -> ops::buildGridFromMesh	gridbatch_from_mesh / grid_from_mesh
fvdb.GridFromDense	create_dense -> ops::createNanoGridFromDense	gridbatch_from_dense / grid_from_dense
fvdb.GridFromNanoVDB	New -- wraps makeGridBatchData from GridBatchDataFactory.h	New

The fvdb.GridFromNanoVDB op is for users who provide a pre-built NanoVDB blob as input. It takes the raw bytes + voxel sizes + origins and derives metadata and index tensors via the existing makeGridBatchData factory.

Grid-Consuming Ops (take a GridBatch Bundle + data, produce outputs):

Priority ops for initial implementation:

ONNX Custom Op	GridBatchOps.cpp binding	fvdb.functional
fvdb.SampleTrilinear	sample_trilinear -> ops::sampleTrilinear	sample_trilinear_batch / _single
fvdb.SampleBezier	sample_bezier -> ops::sampleBezier	sample_bezier_batch / _single
fvdb.SplatTrilinear	splat_trilinear -> ops::splatTrilinear	splat_trilinear_batch / _single
fvdb.SplatBezier	splat_bezier -> ops::splatBezier	splat_bezier_batch / _single
fvdb.PointsInGrid	points_in_grid -> ops::pointsInGrid	points_in_grid_batch / _single
fvdb.IJKToIndex	ijk_to_index -> ops::ijkToIndex	ijk_to_index_batch / _single
fvdb.CoordsInGrid	coords_in_grid -> ops::coordsInGrid	coords_in_grid_batch / _single
fvdb.VoxelsAlongRays	voxels_along_rays -> ops::voxelsAlongRays	voxels_along_rays_batch / _single
fvdb.SegmentsAlongRays	segments_along_rays -> ops::segmentsAlongRays	segments_along_rays_batch / _single
fvdb.MaxPool	max_pool -> ops::maxPool	max_pool_batch / _single
fvdb.AvgPool	avg_pool -> ops::avgPool	avg_pool_batch / _single
fvdb.VoxelToWorld	voxel_to_world -> ops::voxelToWorld	voxel_to_world_batch / _single
fvdb.WorldToVoxel	world_to_voxel -> ops::worldToVoxel	world_to_voxel_batch / _single

Grid-to-Grid Ops (produce a new GridBatch Bundle):

ONNX Custom Op	GridBatchOps.cpp binding	fvdb.functional
fvdb.CoarsenedGrid	coarsen_grid -> ops::buildCoarseGridFromFine	coarsened_grid_batch / _single
fvdb.RefinedGrid	upsample_grid -> ops::buildFineGridFromCoarse	refined_grid_batch / _single
fvdb.DilatedGrid	dilate_grid -> ops::dilateGrid	dilated_grid_batch / _single
fvdb.ConvGrid	conv_grid -> ops::buildGridForConv	conv_grid_batch / _single
fvdb.ConvTransposeGrid	conv_transpose_grid -> ops::buildGridForConvTranspose	conv_transpose_grid_batch / _single
fvdb.PrunedGrid	prune_grid -> ops::pruneGrid	pruned_grid_batch / _single
fvdb.MergedGrid	merge_grids -> ops::mergeGrids	merged_grid_batch / _single

The full op list need not all be implemented at once. The above covers the most common inference workloads.

1.3 Custom Op Library Implementation

Build a shared library (libfvdb_onnx_ops.so / fvdb_onnx_ops.dll) that:

1. Exports RegisterCustomOps per the ONNX Runtime custom op library convention
2. Registers each fVDB custom op in the fvdb domain using Ort::CustomOpDomain + Ort::Custom::CreateLiteCustomOp
3. Provides both CPU and CUDA implementations using Ort::Custom::CudaContext for GPU ops

Each custom op is a thin adapter that:

1. Reconstructs a GridBatchData from the tensor bundle using makeGridBatchData (GridBatchDataFactory.h)
2. Reconstructs JaggedTensor views from component tensors
3. Calls the existing C++ ops::* free function (no kernel rewrite needed)
4. Decomposes the results back into output tensors

GridBatchData is immutable after construction and has a single factory entry point (makeGridBatchData in GridBatchDataFactory.h), which makes reconstruction straightforward:

// Reconstruction from tensor bundle -> GridBatchData
auto grid_hdl = wrapBlobAsGridHandle(grid_blob.Data(), grid_blob.NumberOfElement());
auto gbd = makeGridBatchData(
    std::move(grid_hdl),
    extractVoxelSizes(voxel_sizes),
    extractOrigins(origins));
// Call the existing op
auto result = ops::sampleTrilinear(*gbd, jt_points, voxel_data_tensor);
// Decompose result into output tensors...

1.4 ONNX Model Export

fVDB ops are bound to Python via pybind11 (src/python/GridBatchOps.cpp, src/python/Bindings.cpp), not torch.library, so torch.onnx.export() will not trace them automatically. Provide a Python utility for ONNX export:

• Option A: Custom symbolic functions -- Register torch.onnx symbolic handlers for each _fvdb_cpp.* call that emit the corresponding fvdb.* custom ONNX nodes.
• Option B: Manual graph construction -- Provide an fvdb.export_onnx(model, sample_inputs, path) utility that traces the model and builds the ONNX graph programmatically using onnx.helper, decomposing GridBatchData and JaggedTensor into tensor bundles at graph boundaries.

The fvdb.functional module serves as the reference for which ops need export support: every function in fvdb/functional/__init__.py's __all__ list is a candidate.

1.5 User-Provided Grids

Users who construct grids externally and pass them as inference inputs provide:

• The raw NanoVDB bytes as a uint8 input tensor
• Voxel sizes as a float64[B, 3] input tensor
• Origins as a float64[B, 3] input tensor

The fvdb.GridFromNanoVDB op derives the remaining fields by calling makeGridBatchData from GridBatchDataFactory.h, which scans the NanoVDB buffer to compute per-grid metadata, leaf batch indices, batch offsets, and list indices.

---

Phase 2: fVDB Execution Provider Plugin

Build an ONNX Runtime EP plugin (libfvdb_ep.so) using the EP ABI kernel-based plugin API introduced in ORT v1.24.

2.1 Architecture

The EP claims subgraphs of fvdb.* custom ops via OrtEp::GetCapability + EpGraphSupportInfo_LookUpKernel. Inside the EP:

• GridBatchData objects are maintained as native C++ state, not reconstructed from tensor bundles per-op
• Grid construction ops create GridBatchData via makeGridBatchData (GridBatchDataFactory.h) and store the result in EP-managed state
• Grid-consuming ops access the native GridBatchData directly, calling the same ops::* free functions from src/fvdb/detail/ops/
• Only plain tensors (feature data, query points, output samples) cross the EP boundary

The EP kernel registry (OrtKernelRegistry + KernelRegistry_AddKernel) contains entries for each fVDB op, each with a Compute method that:

1. Retrieves the GridBatchData from EP state (not from tensor inputs)
2. Extracts other inputs from OrtKernelContext
3. Calls the existing ops::* free function
4. Writes outputs via OrtKernelContext

2.2 Key ORT APIs

API	Purpose	Reference
OrtEp::GetCapability	Claim fvdb.* subgraph nodes	PR #24887
OrtEp::GetKernelRegistry	Return registered fVDB kernel implementations	PR #26206
OrtKernelImpl::Compute	Per-op execution entry point	PR #26206
OrtEpApi::KernelRegistry_AddKernel	Register each fVDB kernel	PR #26206
Weight pre-packing	Pre-process static data at session init	PR #26754
OrtEpApi::GetEnvConfigEntries	Read fVDB-specific config	v1.24.1 release

2.3 Benefits Over Phase 1 Alone

• Eliminates per-op GridBatchData reconstruction from the tensor bundle (the makeGridBatchData call that scans the NanoVDB buffer)
• Enables cross-op optimizations (e.g., fusing grid construction + immediate sample)
• Grid state never materializes as ONNX tensor edges, reducing memory management overhead
• Full control over CUDA stream and memory management within the fVDB subgraph

---

Implementation Plan

Phase 1 milestones:

1. Define GridBatch Bundle and JaggedTensor Bundle tensor representations; write bundleFromGridBatchData / gridBatchDataFromBundle helper functions
2. Implement fvdb.GridFromPoints and fvdb.GridFromNanoVDB construction ops wrapping ops::buildGridFromPoints and makeGridBatchData
3. Implement fvdb.SampleTrilinear and fvdb.PointsInGrid as initial consuming ops wrapping ops::sampleTrilinear and ops::pointsInGrid
4. Build and test the custom op shared library with ONNX Runtime (CPU + CUDA)
5. Implement Python export utility mapping fvdb.functional calls to ONNX custom nodes
6. Expand op coverage guided by fvdb/functional/__init__.py's __all__ list

Phase 2 milestones:

1. Scaffold EP plugin with ORT EP ABI
2. Implement GetCapability to claim fvdb.* subgraphs
3. Implement native GridBatchData state management within the EP (using makeGridBatchData for construction, storing c10::intrusive_ptr<GridBatchData> as EP state)
4. Register kernel implementations for all Phase 1 ops via KernelRegistry_AddKernel
5. Benchmark against Phase 1 (per-op reconstruction via makeGridBatchData) to quantify improvement

Key Source Files

File	Relevance
src/fvdb/detail/GridBatchData.h	Frozen struct: all public fields define what the tensor bundle must capture. Contains GridMetadata, GridBatchMetadata, Accessor
src/fvdb/detail/GridBatchDataFactory.h	makeGridBatchData -- the single entry point for constructing GridBatchData from a GridHandle + voxel sizes/origins. The ONNX adapter's reconstruction path
src/fvdb/detail/TorchDeviceBuffer.h	NanoVDB buffer wrapper -- the uint8_t* blob that becomes grid_blob
src/python/GridBatchOps.cpp	All C++ ops bound as free functions -- the definitive list of functions that ONNX custom ops will wrap
fvdb/functional/__init__.py	Complete __all__ export list (~120 ops) -- the ONNX custom op domain catalog
fvdb/functional/_interpolation.py	Representative pattern showing how ops call _fvdb_cpp free functions with GridBatchData
fvdb/functional/_constructors.py	Grid construction functions showing how _fvdb_cpp.gridbatch_from_* factories work
src/fvdb/detail/ops/*.h	~55 op implementations (CUDA kernels) that custom ops ultimately call
src/fvdb/JaggedTensor.h	Data members define the JaggedTensor Bundle tensors
nanovdb/nanovdb/NanoVDB.h	NanoVDB memory layout, 32-byte alignment, GridData (672B header), contiguous buffer
nanovdb/nanovdb/GridHandle.h	GridHandle<BufferT> -- owns the buffer, multi-grid metadata

Key External References

• ONNX Runtime Custom Operator docs -- custom op domain registration, CUDA custom ops, shared library convention
• EP ABI: kernel-based EPs (PR #26206) -- OrtKernelImpl, OrtKernelRegistry, KernelRegistry_AddKernel, EpGraphSupportInfo_LookUpKernel
• EP ABI: GetCapability/Compile infrastructure (PR #24887) -- OrtEp::GetCapability, Graph IR APIs
• EP ABI: weight pre-packing (PR #26754) -- pre-process data at session init
• ONNX Runtime v1.24.1 release notes -- EP Plugin API overview, CreateEnvWithOptions, GetEnvConfigEntries