NanoVDB: chain DeviceBuffer::recordUse and expose orderAfterPriorUses (CUDA)

nanovdb/nanovdb/cuda/DeviceBuffer.h

@@ -47,13 +47,6 @@ class DeviceBuffer
     /// @warning size is expected to be non-zero. Use clear() clear buffer!
     void init(uint64_t size, int device, cudaStream_t stream);
 
-    /// @brief Order work subsequently issued on @a stream after every prior use of this
-    ///        device buffer, whichever stream those uses were issued on.
-    void orderAfterPriorUses(int device, cudaStream_t stream) const
-    {
-        if (mEvents && mEvents[device]) cudaCheck(cudaStreamWaitEvent(stream, mEvents[device], 0));
-    }
-
     /// @brief Free every managed device allocation, each ordered after all tracked uses of the
     ///        buffer, and destroy the tracking events.
     /// @param stream Stream the frees are issued on for allocations owned by the CURRENT device.
@@ -265,6 +258,18 @@ class DeviceBuffer
 
     ///////////////////////////////////////////////////////////////////////
 
+    /// @brief Order work subsequently issued on @a stream after every prior use of this
+    ///        device buffer, whichever stream those uses were issued on. The consume-side
+    ///        companion of recordUse: an external consumer (e.g. a zero-copy array-interface
+    ///        export) calls this with its own stream before reading, so it cannot observe a
+    ///        partially-written buffer after asynchronous uploads or recorded kernels.
+    /// @param device Device whose buffer is about to be read
+    /// @param stream Stream the consumer's work will be issued on
+    void orderAfterPriorUses(int device, cudaStream_t stream) const
+    {
+        if (mEvents && mEvents[device]) cudaCheck(cudaStreamWaitEvent(stream, mEvents[device], 0));
+    }
+
     /// @brief Record that this buffer's device data was just used on @a stream, so that the
     ///        buffer's device frees (destructor, move-assignment, clear) are ordered after that
     ///        work. Uses issued through deviceUpload/deviceDownload are recorded automatically;
@@ -274,6 +279,16 @@ class DeviceBuffer
     ///        implicitly) or if the caller synchronizes before the buffer is cleared/destroyed.
     /// @param device Device whose buffer was used
     /// @param stream Stream the work was issued on
+    /// @note Recording chains across streams: @a stream is first ordered after the previously
+    ///       recorded use (if any) so the single per-device event transitively covers every
+    ///       recorded use, not just the last one. Without this, concurrent uses on streams A
+    ///       then B would leave only B's event, and the device free could run while A's work
+    ///       is still in flight. The side effect is that work subsequently issued on @a stream
+    ///       also waits on the previously recorded use -- acceptable for a shared buffer, where
+    ///       later-recorded consumers observing earlier writes is the expected ordering. Note
+    ///       this also serializes CONCURRENT READERS that record uses (the single event cannot
+    ///       distinguish read-read from write-read); if that ever matters in a profile, the
+    ///       upgrade path is a read/write-separated or per-record event scheme, not a revert.
     void recordUse(int device, cudaStream_t stream)
     {
         if (!mEvents) return;
@@ -283,6 +298,10 @@ class DeviceBuffer
             if (current != device) cudaCheck(cudaSetDevice(device));
             cudaCheck(cudaEventCreateWithFlags(&mEvents[device], cudaEventDisableTiming));
             if (current != device) cudaCheck(cudaSetDevice(current));
+        } else {
+            // Re-recording MOVES the event; chain first so the new capture also covers the
+            // prior recorded use (waiting on a never-recorded or completed event is a no-op).
+            cudaCheck(cudaStreamWaitEvent(stream, mEvents[device], 0));
         }
         cudaCheck(cudaEventRecord(mEvents[device], stream));
     }

nanovdb/nanovdb/unittest/TestNanoVDB.cu

@@ -3890,6 +3890,79 @@ TEST(TestNanoVDBCUDA, DeviceBufferNonBlockingFreeOrdering)
     testDeviceBufferFreeOrdering(/*nonBlockingUser=*/true, /*registerUse=*/true);
 }// DeviceBufferNonBlockingFreeOrdering
 
+TEST(TestNanoVDBCUDA, DeviceBufferChainedRecordUse)
+{
+    // Regression: recordUse re-records the single per-device tracking event, and re-recording
+    // MOVES an event. Two non-blocking streams recording uses in sequence must therefore CHAIN
+    // (the second record first waits on the first capture); otherwise the second record
+    // discards the only coverage of the first stream's in-flight work and the free races it.
+    // Same shape as testDeviceBufferFreeOrdering, with the late writer recorded FIRST and an
+    // idle second stream recorded after it.
+    const size_t N = size_t(64) << 20;
+    const unsigned char LATE = 0xAA, VICTIM = 0x55;
+    const unsigned long long CYCLES = 400000000ull;// parks 'userA' for O(100 ms)
+
+    cudaStream_t userA = nullptr, userB = nullptr, other = nullptr;
+    cudaCheck(cudaStreamCreateWithFlags(&userA, cudaStreamNonBlocking));
+    cudaCheck(cudaStreamCreateWithFlags(&userB, cudaStreamNonBlocking));
+    cudaCheck(cudaStreamCreate(&other));
+    unsigned long long *bad = nullptr;
+    cudaCheck(cudaMallocManaged(&bad, sizeof(*bad)));
+
+    {// warm-up (see testDeviceBufferFreeOrdering)
+        unsigned char *w = nullptr;
+        cudaCheck(cudaMallocAsync((void**)&w, N, other));
+        streamBusyWaitKernel<<<1,1,0,userA>>>(CYCLES/10);
+        deviceBufferFillKernel<<<1024,256,0,other>>>(w, N, 0);
+        deviceBufferCountKernel<<<1024,256,0,other>>>(w, N, 0, bad);
+        cudaCheck(cudaFreeAsync(w, other));
+        cudaCheck(cudaDeviceSynchronize());
+    }
+
+    void *devPtr = nullptr;
+    {
+        auto buf = nanovdb::cuda::DeviceBuffer::create(N, nullptr, 0, other);// device-only
+        devPtr = buf.deviceData(0);
+        ASSERT_TRUE(devPtr);
+        streamBusyWaitKernel<<<1,1,0,userA>>>(CYCLES);// park 'userA'
+        deviceBufferFillKernel<<<1024,256,0,userA>>>((unsigned char*)devPtr, N, LATE);
+        buf.recordUse(0, userA);// covers the in-flight write...
+        buf.recordUse(0, userB);// ...and must NOT be discarded by a later record
+    }// destroyed here; the free must still be ordered after 'userA'
+
+    unsigned char *victim = nullptr;
+    cudaCheck(cudaMallocAsync((void**)&victim, N, other));
+    deviceBufferFillKernel<<<1024,256,0,other>>>(victim, N, VICTIM);
+    cudaCheck(cudaStreamSynchronize(other));
+
+    const bool stillPending = (cudaStreamQuery(userA) == cudaErrorNotReady);
+    cudaGetLastError();// clear the cudaErrorNotReady left by the query above
+    const bool recycled = (victim == devPtr);
+
+    cudaCheck(cudaStreamSynchronize(userA));// let the late write land
+    *bad = 0;
+    deviceBufferCountKernel<<<1024,256>>>(victim, N, VICTIM, bad);
+    cudaCheck(cudaDeviceSynchronize());
+    const unsigned long long clobbered = *bad;
+
+    cudaCheck(cudaFreeAsync(victim, other));
+    cudaCheck(cudaStreamSynchronize(other));
+    cudaCheck(cudaFree(bad));
+    cudaCheck(cudaStreamDestroy(userA));
+    cudaCheck(cudaStreamDestroy(userB));
+    cudaCheck(cudaStreamDestroy(other));
+
+    // Detection relies on the pool recycling the freed block into 'victim' (stream-ordered
+    // pools recycle WITH the dependency attached, so recycling is expected even with a
+    // correctly ordered free). If it did not recycle, the chain was never exercised -- make
+    // that visible instead of a vacuous pass.
+    if (!recycled) GTEST_SKIP() << "allocator did not recycle the block; ordering not exercised";
+
+    EXPECT_EQ(0u, clobbered) << "a later recordUse on another stream discarded the tracking "
+                                "event covering in-flight work (block recycled: " << recycled
+                             << ", work still pending when it was reused: " << stillPending << ")";
+}// DeviceBufferChainedRecordUse
+
 TEST(TestNanoVDBCUDA, RefineCoarsen_ValueOnIndex)
 {
     using BuildT = nanovdb::ValueOnIndex;

pendingchanges/nanovdbrecordusechain.txt

@@ -0,0 +1,15 @@
+NanoVDB:
+
+  Bug fixes:
+  - cuda::DeviceBuffer::recordUse now chains: recording a use on a second
+    stream first orders that stream after the previously recorded use, so the
+    single per-device tracking event transitively covers every recorded use
+    instead of only the most recent one. Previously, concurrent uses recorded
+    on streams A then B left only B's event, and the buffer's device free
+    could run while A's work was still in flight.
+
+  Improvements:
+  - cuda::DeviceBuffer::orderAfterPriorUses is now public — the consume-side
+    companion of recordUse, letting external consumers (e.g. zero-copy
+    array-interface exports) order their own stream after the buffer's
+    tracked uses before reading.