Skip to content

NanoVDB: assign each upper-node tile to a single GPU in DistributedPointsToGrid #2283

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
swahtz wants to merge 2 commits into
base: master
Choose a base branch
from
+203 −94
Open

NanoVDB: assign each upper-node tile to a single GPU in DistributedPointsToGrid #2283

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8fc1545
NanoVDB: assign each upper-node tile to a single GPU in DistributedPo…
swahtz Aug 18, 2026
b61d568
Added a `snapBoundariesToRunsKernel` and running on device to avoid s…
swahtz Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 50 additions & 93 deletions nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,21 +40,6 @@ private:
uint32_t* mNodeCounts;
};

/// @brief Indicator functor that returns 1 if the input value matches the member value, 0 otherwise
template <typename T, typename std::enable_if<std::is_integral<T>::value>::type* = nullptr>
struct EqualityIndicator
{
EqualityIndicator(const T* value) : mValue(value) {}

__hostdev__
T operator()(const T& x) const
{
return x == (*mValue);
}
private:
const T* mValue;
};

/// @brief Find a partition at an arbitrary diagonal in the conceptual merge of two sorted input arrays.
template<typename KeyIteratorIn>
__device__
Expand Down Expand Up @@ -86,34 +71,6 @@ void mergePathKernel(KeyIteratorIn keys1, size_t keys1Count, KeyIteratorIn keys2
mergePath(keys1, keys1Count, keys2, keys2Count, key1Intervals, key2Intervals, combinedIndex);
}

/// @brief Extends or shortens the left end of an array interval
template<typename DistanceIteratorIn, typename CountIteratorOut, typename OffsetIteratorOut>
__global__
void leftRebalanceKernel(DistanceIteratorIn leftDistance, DistanceIteratorIn rightDistance, CountIteratorOut leftCount, OffsetIteratorOut leftOffset)
{
if (*leftDistance < *rightDistance) {
*leftCount -= *leftDistance;
}
else {
*leftCount += *rightDistance;
}
}

/// @brief Extends or shortens the right end of an array interval
template<typename DistanceIteratorIn, typename CountIteratorOut, typename OffsetIteratorOut>
__global__
void rightRebalanceKernel(DistanceIteratorIn leftDistance, DistanceIteratorIn rightDistance, CountIteratorOut rightCount, OffsetIteratorOut rightOffset)
{
if (*leftDistance < *rightDistance) {
*rightCount += *leftDistance;
*rightOffset -= *leftDistance;
}
else {
*rightCount -= *rightDistance;
*rightOffset += *rightDistance;
}
}

} // namespace kernels

// Define utility macro used to call cub functions that use dynamic temporary storage
Expand Down Expand Up @@ -580,8 +537,6 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord
// to reduce overhead.
std::vector<cudaEvent_t> sortEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> runLengthEncodeEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> transformReduceEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> rebalanceEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> tilePrefixSumEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> voxelCountEvents(mDeviceMesh.deviceCount());
std::vector<cudaEvent_t> leafCountEvents(mDeviceMesh.deviceCount());
Expand All @@ -592,8 +547,6 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord
cudaCheck(cudaSetDevice(deviceId));
cudaEventCreateWithFlags(&sortEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&runLengthEncodeEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&transformReduceEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&rebalanceEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&tilePrefixSumEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&voxelCountEvents[deviceId], cudaEventDisableTiming);
cudaEventCreateWithFlags(&leafCountEvents[deviceId], cudaEventDisableTiming);
Expand All @@ -604,7 +557,6 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord

// Advise per-coord quantities to be split evenly across devices. Clamp each stripe to
// the input range so that inputs smaller than the device count produce valid trailing empty stripes.
std::vector<size_t> deviceStripeCounts(mDeviceMesh.deviceCount());
const size_t deviceStripeSize = ::cuda::ceil_div(coordCount, mDeviceMesh.deviceCount());
for (const auto& [deviceId, stream] : mDeviceMesh) {
cudaCheck(cudaSetDevice(deviceId));
Expand All @@ -614,7 +566,6 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord

mStripeCounts[deviceId] = deviceStripeCount;
mStripeOffsets[deviceId] = deviceStripeOffset;
deviceStripeCounts[deviceId] = deviceStripeCount;

if (deviceStripeCount) {
nanovdb::Coord* deviceCoords = coords + deviceStripeOffset;
Expand Down Expand Up @@ -650,60 +601,68 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord

radixSortAsync(mDeviceMesh, mTempDevicePools, mData->d_keys, mKeys, mData->d_indx, mIndices, coordCount, mIntervals, mStripeOffsets, mStripeCounts, sortEvents.data(), sortEvents.data());

// For each segment of sorted keys on each device, we count how many of the leftmost key occur past the left boundary of the segment. The same is done for the rightmost key with the right boundary of the segment.
auto leftIntervals = mIntervals;
auto rightIntervals = mIntervals + mDeviceMesh.deviceCount() + 1;
// Rebalance the device segments so that a device boundary always coincides
// with a change in key value. Because TileKeyFunctor assigns identical keys
// to every point that falls in the same upper-node "tile", this aligns the
// device ownership boundaries with tile boundaries. Downstream construction
// assumes each tile (and therefore each lower node, leaf node, and voxel) is
// owned by exactly one device; if a tile straddled a boundary, multiple
// devices would concurrently build the same leaf and race on its value mask.
//
// A single tile can span three or more devices (e.g. one dense leaf whose
// points are split evenly across the mesh). Adjusting only adjacent pairs of
// boundaries cannot consolidate such a tile because a fully-interior device
// lies entirely within it, so we compute the boundaries globally and
// monotonically. mKeys is globally sorted at this point, so a tile boundary
// is simply a position where mKeys changes. Snapping is performed on the host
// over the (small) set of device boundaries; fully-interior devices are left
// empty, which the rest of the pipeline already handles.
for (const auto& [deviceId, stream] : mDeviceMesh) {
cudaCheck(cudaSetDevice(deviceId));

auto deviceStripeCount = mStripeCounts[deviceId];
auto deviceStripeOffset = mStripeOffsets[deviceId];
uint64_t* deviceInputKeys = mKeys + deviceStripeOffset;

if (deviceStripeCounts[deviceId] && deviceId > 0 && deviceStripeCounts[deviceId - 1]) {
cudaCheck(cudaStreamWaitEvent(stream, sortEvents[deviceId - 1]));
EqualityIndicator<uint64_t> indicator(deviceInputKeys - 1);
CUB_LAUNCH(DeviceReduce::TransformReduce, mTempDevicePools[deviceId], stream, deviceInputKeys, rightIntervals + deviceId, deviceStripeCount, ::cuda::std::plus(), indicator, 0);
} else {
rightIntervals[deviceId] = 0;
}

if (deviceStripeCounts[deviceId] && deviceId < static_cast<int>(mDeviceMesh.deviceCount() - 1) && deviceStripeCounts[deviceId + 1]) {
cudaCheck(cudaStreamWaitEvent(stream, sortEvents[deviceId + 1]));
EqualityIndicator<uint64_t> indicator(deviceInputKeys + deviceStripeCount);
CUB_LAUNCH(DeviceReduce::TransformReduce, mTempDevicePools[deviceId], stream, deviceInputKeys, leftIntervals + deviceId, deviceStripeCount, ::cuda::std::plus(), indicator, 0);
} else {
leftIntervals[deviceId] = 0;
}
cudaCheck(cudaEventRecord(transformReduceEvents[deviceId], stream));
cudaCheck(cudaStreamSynchronize(stream));

@matthewdcong matthewdcong Aug 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to do this on the device instead of the host in order to avoid the stream sync? Alternatively, would it be possible to run a benchmark on an analytic example (e.g. a sampled torus) to show that performance isn't affected?

@swahtz swahtz Aug 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I added a snapBoundariesToRunsKernel so we run the snapping computation on-device.

}

// Rebalance the segments so that a device segment boundary also corresponds to a change in key value. Effectively, this aligns upper node boundaries with device ownership boundaries.
for (const auto& [deviceId, stream] : mDeviceMesh) {
cudaCheck(cudaSetDevice(deviceId));

if (deviceId > 0 && deviceStripeCounts[deviceId] && deviceStripeCounts[deviceId - 1])
{
cudaCheck(cudaStreamWaitEvent(stream, transformReduceEvents[deviceId - 1]));
kernels::rightRebalanceKernel<<<1, 1, 0, stream>>>(leftIntervals + deviceId - 1, rightIntervals + deviceId, mStripeCounts + deviceId, mStripeOffsets + deviceId);
cudaCheckError();
{
const int deviceCount = static_cast<int>(mDeviceMesh.deviceCount());
const ptrdiff_t keyCount = static_cast<ptrdiff_t>(coordCount);
ptrdiff_t previousBoundary = mStripeOffsets[0]; // device 0 always starts at 0
for (int deviceId = 1; deviceId < deviceCount; ++deviceId) {
ptrdiff_t boundary = mStripeOffsets[deviceId];
if (boundary >= keyCount) {
boundary = keyCount;
} else if (boundary > previousBoundary && mKeys[boundary] == mKeys[boundary - 1]) {
// The even-split boundary falls inside a tile run; find its extent
// and snap to whichever end keeps the boundary closest to the
// even split without crossing the previous boundary.
ptrdiff_t runStart = boundary;
while (runStart > previousBoundary && mKeys[runStart - 1] == mKeys[boundary]) --runStart;
ptrdiff_t runEnd = boundary;
while (runEnd < keyCount && mKeys[runEnd] == mKeys[boundary]) ++runEnd;
if (runStart <= previousBoundary) {
boundary = runEnd; // the run reaches the previous device, give the whole tile away
} else {
boundary = (boundary - runStart <= runEnd - boundary) ? runStart : runEnd;
}
}
if (boundary < previousBoundary) boundary = previousBoundary;
mStripeOffsets[deviceId] = boundary;
previousBoundary = boundary;
}

if (deviceId < static_cast<int>(mDeviceMesh.deviceCount() - 1) && deviceStripeCounts[deviceId] && deviceStripeCounts[deviceId + 1])
{
cudaCheck(cudaStreamWaitEvent(stream, transformReduceEvents[deviceId + 1]));
kernels::leftRebalanceKernel<<<1, 1, 0, stream>>>(leftIntervals + deviceId, rightIntervals + deviceId + 1, mStripeCounts + deviceId, mStripeOffsets + deviceId);
cudaCheckError();
// Recompute the per-device counts from the adjusted, monotonic offsets.
for (int deviceId = 0; deviceId < deviceCount; ++deviceId) {
const ptrdiff_t nextOffset = (deviceId + 1 < deviceCount)
? mStripeOffsets[deviceId + 1]
: keyCount;
mStripeCounts[deviceId] = static_cast<size_t>(nextOffset - mStripeOffsets[deviceId]);
}
cudaCheck(cudaEventRecord(rebalanceEvents[deviceId], stream));
}

// Parallel RLE in order to obtain tiles
// Parallel RLE in order to obtain tiles. The device boundaries were finalized
// synchronously on the host above, so no per-device rebalance event is needed.
for (const auto& [deviceId, stream] : mDeviceMesh) {
cudaCheck(cudaSetDevice(deviceId));

cudaCheck(cudaEventSynchronize(rebalanceEvents[deviceId]));

auto deviceStripeCount = mStripeCounts[deviceId];
auto deviceStripeOffset = mStripeOffsets[deviceId];

Expand Down Expand Up @@ -877,8 +836,6 @@ void DistributedPointsToGrid<BuildT>::countNodes(const PtrT coords, size_t coord
cudaCheck(cudaSetDevice(deviceId));
cudaEventDestroy(sortEvents[deviceId]);
cudaEventDestroy(runLengthEncodeEvents[deviceId]);
cudaEventDestroy(transformReduceEvents[deviceId]);
cudaEventDestroy(rebalanceEvents[deviceId]);
cudaEventDestroy(tilePrefixSumEvents[deviceId]);
cudaEventDestroy(voxelCountEvents[deviceId]);
cudaEventDestroy(leafCountEvents[deviceId]);
Expand Down
120 changes: 119 additions & 1 deletion nanovdb/nanovdb/unittest/TestMultiGPU.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,9 @@
#include <cuda_runtime_api.h>
#include <gtest/gtest.h>
#include <thread> // for std::thread
#include <algorithm> // for std::sort, std::unique
#include <cstdint>
#include <vector>

#include <thrust/fill.h>
#include <thrust/universal_vector.h>
Expand Down Expand Up @@ -328,7 +331,7 @@ TEST(TestNanoVDBMultiGPU, DenseLeaf_DistributedCudaPointsToGrid_UnifiedBuffer)
EXPECT_TRUE(data);
grid = handle.grid<BuildT>();
EXPECT_TRUE(grid);
EXPECT_TRUE(grid->activeVoxelCount() == 512);
EXPECT_EQ(voxelCount, grid->activeVoxelCount());
EXPECT_EQ(nanovdb::Vec3d(1.0), grid->voxelSize());

cudaCheck(cudaFree(voxels));
Expand Down Expand Up @@ -500,3 +503,118 @@ TEST(TestNanoVDBMultiGPU, ManyTiles_DistributedCudaPointsToGrid)
cudaCheck(cudaFree(voxels));
cudaSetDevice(current);
}// ManyTiles_DistributedCudaPointsToGrid

/// @brief Regression test for the multi-GPU shared-tile race. All coordinates
/// fall inside a single upper-node tile (each upper node spans 4096
/// voxels per axis, so ijk in [0, 4095] maps to tile (0,0,0)), yet they
/// populate many lower and leaf nodes. With the default even initial
/// split this one tile is shared across every GPU, so the builder must
/// consolidate it onto a single device. A cross-device race in leaf
/// construction would drop a device's contribution and undercount the
/// active voxels, which the exact-count assertion below catches
/// deterministically.
TEST(TestNanoVDBMultiGPU, SingleUpperNode_DistributedCudaPointsToGrid_UnifiedBuffer)
{
int current = 0;
cudaCheck(cudaGetDevice(&current));

using BufferT = nanovdb::cuda::UnifiedBuffer;
using BuildT = nanovdb::ValueOnIndex;

const size_t inputCount = 1 << 18;// 262144
nanovdb::Coord* voxels = nullptr;
cudaCheck(cudaMallocManaged(&voxels, inputCount * sizeof(nanovdb::Coord)));
std::srand(24680);
auto op = [](){ return rand() % 4096; };// stays within a single upper node
for (size_t i = 0; i < inputCount; ++i)
voxels[i] = nanovdb::Coord(op(), op(), op());

// Deterministic expected active-voxel count (the input may contain duplicates).
std::vector<uint64_t> packed(inputCount);
for (size_t i = 0; i < inputCount; ++i)
packed[i] = (uint64_t(voxels[i][0]) << 24) | (uint64_t(voxels[i][1]) << 12) | uint64_t(voxels[i][2]);
std::sort(packed.begin(), packed.end());
const size_t uniqueCount = static_cast<size_t>(std::unique(packed.begin(), packed.end()) - packed.begin());

nanovdb::cuda::DeviceMesh deviceMesh;
nanovdb::tools::cuda::DistributedPointsToGrid<BuildT> converter(deviceMesh);
auto handle = converter.getHandle(voxels, inputCount);

EXPECT_TRUE(handle.deviceData());
EXPECT_TRUE(handle.deviceGrid<BuildT>());
handle.deviceDownload();
auto *grid = handle.grid<BuildT>();
EXPECT_TRUE(grid);
EXPECT_EQ(nanovdb::Vec3d(1.0), grid->voxelSize());
// The input occupies exactly one upper-node tile, so the shared-tile
// consolidation path is exercised.
EXPECT_EQ(1u, grid->tree().nodeCount(2));
// Every unique input voxel must be active exactly once.
EXPECT_EQ(static_cast<uint64_t>(uniqueCount), grid->activeVoxelCount());

nanovdb::util::forEach(0, inputCount, 1, [&](const nanovdb::util::Range1D &r){
auto acc = grid->getAccessor();
for (size_t i=r.begin(); i!=r.end(); ++i) {
const nanovdb::Coord &ijk = voxels[i];
EXPECT_TRUE(acc.probeLeaf(ijk)!=nullptr);
EXPECT_TRUE(acc.isActive(ijk));
EXPECT_TRUE(acc.getValue(ijk) > 0u);
}
});

cudaCheck(cudaFree(voxels));
cudaSetDevice(current);
}// SingleUpperNode_DistributedCudaPointsToGrid_UnifiedBuffer

/// @brief Cross-checks the distributed builder against the trusted single-GPU
/// PointsToGrid on an input that forces a single tile to be split across
/// devices. Index assignment order may differ between the two builders,
/// so we compare topology (node counts, active-voxel count) and voxel
/// occupancy rather than the ValueOnIndex indices themselves.
TEST(TestNanoVDBMultiGPU, MatchesSingleGpu_DistributedCudaPointsToGrid)
{
int current = 0;
cudaCheck(cudaGetDevice(&current));

using BufferT = nanovdb::cuda::UnifiedBuffer;
using BuildT = nanovdb::ValueOnIndex;

const size_t inputCount = 1 << 17;// 131072, all within a single upper node
nanovdb::Coord* voxels = nullptr;
cudaCheck(cudaMallocManaged(&voxels, inputCount * sizeof(nanovdb::Coord)));
std::srand(1357);
auto op = [](){ return rand() % 4096; };
for (size_t i = 0; i < inputCount; ++i)
voxels[i] = nanovdb::Coord(op(), op(), op());

nanovdb::cuda::DeviceMesh deviceMesh;
nanovdb::tools::cuda::DistributedPointsToGrid<BuildT> converter(deviceMesh);
auto distributedHandle = converter.getHandle(voxels, inputCount);
distributedHandle.deviceDownload();
auto *distributedGrid = distributedHandle.grid<BuildT>();
EXPECT_TRUE(distributedGrid);

cudaSetDevice(current);
auto referenceHandle = nanovdb::tools::cuda::voxelsToGrid<BuildT, nanovdb::Coord*, BufferT>(voxels, inputCount);
referenceHandle.deviceDownload();
auto *referenceGrid = referenceHandle.grid<BuildT>();
EXPECT_TRUE(referenceGrid);

EXPECT_EQ(referenceGrid->activeVoxelCount(), distributedGrid->activeVoxelCount());
EXPECT_EQ(referenceGrid->tree().nodeCount(0), distributedGrid->tree().nodeCount(0));
EXPECT_EQ(referenceGrid->tree().nodeCount(1), distributedGrid->tree().nodeCount(1));
EXPECT_EQ(referenceGrid->tree().nodeCount(2), distributedGrid->tree().nodeCount(2));

nanovdb::util::forEach(0, inputCount, 1, [&](const nanovdb::util::Range1D &r){
auto distributedAcc = distributedGrid->getAccessor();
auto referenceAcc = referenceGrid->getAccessor();
for (size_t i=r.begin(); i!=r.end(); ++i) {
const nanovdb::Coord &ijk = voxels[i];
EXPECT_TRUE(referenceAcc.isActive(ijk));
EXPECT_TRUE(distributedAcc.isActive(ijk));
}
});

cudaCheck(cudaFree(voxels));
cudaSetDevice(current);
}// MatchesSingleGpu_DistributedCudaPointsToGrid
10 changes: 10 additions & 0 deletions pendingchanges/distributedtileownership.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
NanoVDB:
Bug Fixes:
- Fixed a cross-device race in tools::cuda::DistributedPointsToGrid. The
device segment boundaries were rebalanced by adjusting only adjacent pairs
of boundaries, which cannot consolidate an upper-node tile that spans three
or more GPUs, because a fully-interior device lies entirely within the tile.
Multiple devices then built the same leaf concurrently and raced on its
value mask, silently dropping active voxels. The boundaries are now snapped
globally and monotonically to tile boundaries, so every tile - and therefore
every lower node, leaf node and voxel - is owned by exactly one device.
Loading