NanoVDB: handle empty partitions in the distributed merge path search (#2248)

* NanoVDB: handle empty partitions in the distributed merge path search

mergePath's bounds checks compare against (keys1Count - 1) and
(keys2Count - 1) on size_t counts, so a zero count underflows to
SIZE_MAX and the guards never fire; the binary search then reads the
empty array (and past either array) and produces wrong or garbage
intervals. An empty partition is reachable from radixSortAsync's
ceil-split whenever a device receives no items, e.g. sorting a single
element across two devices.

Return the trivial split before the search: with an empty side, every
element up to the diagonal comes from the non-empty side.

New single-GPU test MergePathEmptyPartition covers the empty-right,
empty-left, both-empty, and non-empty control cases; the empty cases
fail deterministically without the fix (inverted split / underflow
garbage) and run clean under compute-sanitizer with it.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Mark Harris <mharris@nvidia.com>

* NanoVDB: compute the empty-partition merge split on the host

An empty side needs no merge path search, so skip the kernel launch:
the trivial split (everything up to the diagonal comes from the
non-empty side) is computed on the host by mergePathTrivial at the
launch site. mergePath reverts to the unguarded search and documents
that both inputs must be non-empty.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Mark Harris <mharris@nvidia.com>

---------

Signed-off-by: Mark Harris <mharris@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: harrism

nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh

@@ -55,7 +55,21 @@ private:
     const T* mValue;
 };
 
+/// @brief Computes the trivial merge path split for the case where either
+///        input is empty: every element up to the diagonal comes from the
+///        non-empty side. Host-callable so callers can skip the kernel launch.
+inline void mergePathTrivial(size_t keys1Count, size_t keys2Count, ptrdiff_t* key1Intervals, ptrdiff_t* key2Intervals, int intervalIndex)
+{
+    const size_t combinedIndex = intervalIndex * (keys1Count + keys2Count) / 2;
+    *key1Intervals = static_cast<ptrdiff_t>(combinedIndex < keys1Count ? combinedIndex : keys1Count);
+    *key2Intervals = static_cast<ptrdiff_t>(combinedIndex < keys2Count ? combinedIndex : keys2Count);
+}
+
 /// @brief Implements the merge path binary search algorithm in order to find the median across two sorted input key arrays
+/// @warning Both inputs must be non-empty: the binary search reads both
+///          arrays, and its unsigned (keysNCount - 1) bounds checks underflow
+///          when a count is zero. Callers handle an empty side with
+///          mergePathTrivial instead of launching this search.
 template<typename KeyIteratorIn>
 __device__
 void mergePath(KeyIteratorIn keys1, size_t keys1Count, KeyIteratorIn keys2, size_t keys2Count, ptrdiff_t* key1Intervals, ptrdiff_t* key2Intervals, int intervalIndex)
@@ -226,7 +240,14 @@ void radixSortAsync(const nanovdb::cuda::DeviceMesh& deviceMesh, nanovdb::cuda::
                     cudaCheck(cudaSetDevice(deviceId));
 
                     cudaCheck(cudaStreamWaitEvent(deviceMesh[deviceId].stream, postEvents[otherDeviceId]));
-                    kernels::mergePathKernel<<<1, 1, 0, deviceMesh[deviceId].stream>>>(leftDeviceKeysIn, leftDeviceItemCount, rightDeviceKeysIn, rightDeviceItemCount, leftIntervals + deviceId, rightIntervals + deviceId, intervalIndex);
+                    if (leftDeviceItemCount == 0 || rightDeviceItemCount == 0) {
+                        // An empty side makes the split trivial; no kernel launch
+                        // is needed and the search itself requires non-empty inputs.
+                        mergePathTrivial(leftDeviceItemCount, rightDeviceItemCount, leftIntervals + deviceId, rightIntervals + deviceId, intervalIndex);
+                    }
+                    else {
+                        kernels::mergePathKernel<<<1, 1, 0, deviceMesh[deviceId].stream>>>(leftDeviceKeysIn, leftDeviceItemCount, rightDeviceKeysIn, rightDeviceItemCount, leftIntervals + deviceId, rightIntervals + deviceId, intervalIndex);
+                    }
                     cudaCheck(cudaEventRecord(postEvents[deviceId], deviceMesh[deviceId].stream));
                 };
                 mergePathSubfunc(leftDeviceId, rightDeviceId, 0);

nanovdb/nanovdb/unittest/TestMultiGPU.cu

@@ -71,6 +71,51 @@ TEST(TestNanoVDBMultiGPU, RadixSort)
     }
 }
 
+/// @brief Tests the merge path split for empty partitions. An empty side
+///        yields a trivial split (everything up to the diagonal comes from
+///        the non-empty side), computed on the host via mergePathTrivial —
+///        the device search requires non-empty inputs. Runs on a single GPU.
+TEST(TestNanoVDBMultiGPU, MergePathEmptyPartition)
+{
+    using nanovdb::tools::cuda::mergePathTrivial;
+    constexpr size_t count = 4;
+    ptrdiff_t k1 = -1, k2 = -1;
+
+    // Median diagonal with an empty right partition: all elements from keys1.
+    mergePathTrivial(count, 0, &k1, &k2, 1);
+    EXPECT_EQ(k1, ptrdiff_t(count / 2));
+    EXPECT_EQ(k2, ptrdiff_t(0));
+
+    // Median diagonal with an empty left partition: all elements from keys2.
+    mergePathTrivial(0, count, &k1, &k2, 1);
+    EXPECT_EQ(k1, ptrdiff_t(0));
+    EXPECT_EQ(k2, ptrdiff_t(count / 2));
+
+    // Both partitions empty.
+    mergePathTrivial(0, 0, &k1, &k2, 0);
+    EXPECT_EQ(k1, ptrdiff_t(0));
+    EXPECT_EQ(k2, ptrdiff_t(0));
+
+    // Control: the device search on two non-empty interleaved partitions
+    // splits both arrays at the median diagonal.
+    using KeyT = int;
+    KeyT* keys = nullptr;
+    cudaCheck(cudaMallocManaged(&keys, 2 * count * sizeof(KeyT)));
+    for (size_t i = 0; i < count; ++i) keys[i] = 10 * (int(i) + 1);              // 10 20 30 40
+    for (size_t i = 0; i < count; ++i) keys[count + i] = 10 * (int(i) + 1) + 5;  // 15 25 35 45
+    ptrdiff_t* intervals = nullptr;
+    cudaCheck(cudaMallocManaged(&intervals, 2 * sizeof(ptrdiff_t)));
+    intervals[0] = intervals[1] = -1;
+    nanovdb::tools::cuda::kernels::mergePathKernel<<<1, 1>>>(keys, count, keys + count, count, intervals, intervals + 1, size_t(1));
+    cudaCheck(cudaDeviceSynchronize());
+    // Merged head is 10 15 20 25, so the median splits both arrays at 2.
+    EXPECT_EQ(intervals[0], ptrdiff_t(2));
+    EXPECT_EQ(intervals[1], ptrdiff_t(2));
+
+    cudaCheck(cudaFree(intervals));
+    cudaCheck(cudaFree(keys));
+}
+
 /// @brief Tests the correctness of multi-GPU exclusive sums against an equivalent CPU implementation
 TEST(TestNanoVDBMultiGPU, ExclusiveSum)
 {