From 8fc1545f3500e630021a94239d055178342bdf21 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Tue, 18 Aug 2026 17:13:16 +1200 Subject: [PATCH 1/2] 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) Signed-off-by: Jonathan Swartz --- .../tools/cuda/DistributedPointsToGrid.cuh | 143 ++++++------------ nanovdb/nanovdb/unittest/TestMultiGPU.cu | 120 ++++++++++++++- pendingchanges/distributedtileownership.txt | 10 ++ 3 files changed, 179 insertions(+), 94 deletions(-) create mode 100644 pendingchanges/distributedtileownership.txt diff --git a/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh b/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh index da098c98f7..2c750b91d0 100644 --- a/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh +++ b/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 ::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 __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 -__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 -__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::countNodes(const PtrT coords, size_t coord // to reduce overhead. std::vector sortEvents(mDeviceMesh.deviceCount()); std::vector runLengthEncodeEvents(mDeviceMesh.deviceCount()); - std::vector transformReduceEvents(mDeviceMesh.deviceCount()); - std::vector rebalanceEvents(mDeviceMesh.deviceCount()); std::vector tilePrefixSumEvents(mDeviceMesh.deviceCount()); std::vector voxelCountEvents(mDeviceMesh.deviceCount()); std::vector leafCountEvents(mDeviceMesh.deviceCount()); @@ -592,8 +547,6 @@ void DistributedPointsToGrid::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::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 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::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::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 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(mDeviceMesh.deviceCount() - 1) && deviceStripeCounts[deviceId + 1]) { - cudaCheck(cudaStreamWaitEvent(stream, sortEvents[deviceId + 1])); - EqualityIndicator 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(mDeviceMesh.deviceCount()); + const ptrdiff_t keyCount = static_cast(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(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(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::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]); diff --git a/nanovdb/nanovdb/unittest/TestMultiGPU.cu b/nanovdb/nanovdb/unittest/TestMultiGPU.cu index dd8526417b..5930a62aa6 100644 --- a/nanovdb/nanovdb/unittest/TestMultiGPU.cu +++ b/nanovdb/nanovdb/unittest/TestMultiGPU.cu @@ -10,6 +10,9 @@ #include #include #include // for std::thread +#include // for std::sort, std::unique +#include +#include #include #include @@ -328,7 +331,7 @@ TEST(TestNanoVDBMultiGPU, DenseLeaf_DistributedCudaPointsToGrid_UnifiedBuffer) EXPECT_TRUE(data); grid = handle.grid(); 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(¤t)); + + 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 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(std::unique(packed.begin(), packed.end()) - packed.begin()); + + nanovdb::cuda::DeviceMesh deviceMesh; + nanovdb::tools::cuda::DistributedPointsToGrid converter(deviceMesh); + auto handle = converter.getHandle(voxels, inputCount); + + EXPECT_TRUE(handle.deviceData()); + EXPECT_TRUE(handle.deviceGrid()); + handle.deviceDownload(); + auto *grid = handle.grid(); + 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(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(¤t)); + + 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 converter(deviceMesh); + auto distributedHandle = converter.getHandle(voxels, inputCount); + distributedHandle.deviceDownload(); + auto *distributedGrid = distributedHandle.grid(); + EXPECT_TRUE(distributedGrid); + + cudaSetDevice(current); + auto referenceHandle = nanovdb::tools::cuda::voxelsToGrid(voxels, inputCount); + referenceHandle.deviceDownload(); + auto *referenceGrid = referenceHandle.grid(); + 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 diff --git a/pendingchanges/distributedtileownership.txt b/pendingchanges/distributedtileownership.txt new file mode 100644 index 0000000000..19cd5016bb --- /dev/null +++ b/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. From b61d56831c72199f3be8a1fd27cf6a8e994d0655 Mon Sep 17 00:00:00 2001 From: Jonathan Swartz Date: Thu, 20 Aug 2026 16:07:49 +1200 Subject: [PATCH 2/2] Added a `snapBoundariesToRunsKernel` and running on device to avoid stream sync Signed-off-by: Jonathan Swartz --- .../tools/cuda/DistributedPointsToGrid.cuh | 108 +++++++++++------- 1 file changed, 66 insertions(+), 42 deletions(-) diff --git a/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh b/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh index 2c750b91d0..fbeb37b9d1 100644 --- a/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh +++ b/nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh @@ -71,6 +71,56 @@ void mergePathKernel(KeyIteratorIn keys1, size_t keys1Count, KeyIteratorIn keys2 mergePath(keys1, keys1Count, keys2, keys2Count, key1Intervals, key2Intervals, combinedIndex); } +/// @brief Snap each device boundary in stripeOffsets to the nearest edge of the run of equal +/// keys containing it, so that no run (i.e. upper-node tile) straddles two devices. The keys +/// must be globally sorted. The boundaries are adjusted monotonically from left to right, so a +/// run spanning several stripes is consolidated onto one device and the fully-interior devices +/// are left with empty stripes. Runs sequentially on a single thread: the boundary chain is a +/// sequential dependence over deviceCount entries and each run extent is found by binary search. +template +__global__ +void snapBoundariesToRunsKernel(const KeyT* keys, ptrdiff_t keyCount, int deviceCount, ptrdiff_t* stripeOffsets, size_t* stripeCounts) +{ + ptrdiff_t previousBoundary = stripeOffsets[0]; // device 0 always starts at 0 + for (int deviceId = 1; deviceId < deviceCount; ++deviceId) { + ptrdiff_t boundary = stripeOffsets[deviceId]; + if (boundary >= keyCount) { + boundary = keyCount; + } else if (boundary > previousBoundary && keys[boundary] == keys[boundary - 1]) { + // The even-split boundary falls inside a run; binary search for the run's extent + // and snap to whichever end keeps the boundary closest to the even split without + // crossing the previous boundary. + const KeyT key = keys[boundary]; + ptrdiff_t lo = previousBoundary, hi = boundary; + while (lo < hi) { + const ptrdiff_t mid = lo + (hi - lo) / 2; + if (keys[mid] < key) lo = mid + 1; else hi = mid; + } + const ptrdiff_t runStart = lo; + lo = boundary; hi = keyCount; + while (lo < hi) { + const ptrdiff_t mid = lo + (hi - lo) / 2; + if (keys[mid] <= key) lo = mid + 1; else hi = mid; + } + const ptrdiff_t runEnd = lo; + 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; + stripeOffsets[deviceId] = boundary; + previousBoundary = boundary; + } + + // Recompute the per-device counts from the adjusted, monotonic offsets. + for (int deviceId = 0; deviceId < deviceCount; ++deviceId) { + const ptrdiff_t nextOffset = (deviceId + 1 < deviceCount) ? stripeOffsets[deviceId + 1] : keyCount; + stripeCounts[deviceId] = static_cast(nextOffset - stripeOffsets[deviceId]); + } +} + } // namespace kernels // Define utility macro used to call cub functions that use dynamic temporary storage @@ -614,52 +664,26 @@ void DistributedPointsToGrid::countNodes(const PtrT coords, size_t coord // 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)); - cudaCheck(cudaStreamSynchronize(stream)); - } - + // is simply a position where mKeys changes. Snapping runs in a single-thread + // kernel on one device (every device stream is already ordered after the + // whole sort by radixSortAsync's final stream merge) with binary searches for + // the run extents; the host only waits on its completion event before reading + // back the small offset/count arrays, mirroring how the sort's merge-path + // partitions are synchronized. Fully-interior devices are left with empty + // stripes, which the rest of the pipeline already handles. { - const int deviceCount = static_cast(mDeviceMesh.deviceCount()); - const ptrdiff_t keyCount = static_cast(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; - } - - // 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(nextOffset - mStripeOffsets[deviceId]); - } + static constexpr int snapDeviceId = 0; + cudaCheck(cudaSetDevice(snapDeviceId)); + const auto snapStream = mDeviceMesh[snapDeviceId].stream; + kernels::snapBoundariesToRunsKernel<<<1, 1, 0, snapStream>>>(mKeys, static_cast(coordCount), static_cast(mDeviceMesh.deviceCount()), mStripeOffsets, mStripeCounts); + cudaCheckError(); + cudaCheck(cudaEventRecord(sortEvents[snapDeviceId], snapStream)); + // CUB needs the rebalanced partition sizes on the host. Synchronize only the tiny snap kernel. + cudaCheck(cudaEventSynchronize(sortEvents[snapDeviceId])); } // 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. + // above before the host read them back, so no per-device rebalance event is needed. for (const auto& [deviceId, stream] : mDeviceMesh) { cudaCheck(cudaSetDevice(deviceId));