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

The device segment boundaries produced by the initial even split were
rebalanced by adjusting only adjacent pairs of boundaries. That cannot
consolidate an upper-node tile spanning three or more GPUs, since a
fully-interior device lies entirely within the tile. Downstream
construction assumes each tile - and therefore each lower node, leaf node
and voxel - is owned by exactly one device, so multiple devices built the
same leaf concurrently and raced on its value mask, silently dropping
active voxels.

Replace the pairwise GPU rebalance with a global, monotonic snap performed
on the host over the (small) set of device boundaries. The keys are
globally sorted at that point, so a tile boundary is simply a position
where the key changes; fully-interior devices are left empty, which the
rest of the pipeline already handles.

Add two multi-GPU regression tests: SingleUpperNode, which places all
input in one upper-node tile and asserts the exact unique active-voxel
count, and MatchesSingleGpu, which cross-checks topology and occupancy
against the single-GPU builder on the same input.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>

Author: swahtz

nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh

@@ -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__
@@ -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
@@ -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());
@@ -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);
@@ -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));
@@ -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;
@@ -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));
     }
 
-    // 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];
 
@@ -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]);

nanovdb/nanovdb/unittest/TestMultiGPU.cu

@@ -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>
@@ -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));
@@ -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

pendingchanges/distributedtileownership.txt

@@ -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.