Add persistent CUDA ray tracing path for iso-surface example

nanovdb/nanovdb/examples/ex_raytrace_iso_surface/common.h

@@ -10,22 +10,9 @@
 #include <nanovdb/NanoVDB.h>
 #include "ComputePrimitives.h"
 
-template<typename RenderFn, typename GridT>
-inline float renderImage(bool useCuda, const RenderFn renderOp, int width, int height, float* image, const GridT* grid)
-{
-    using ClockT = std::chrono::high_resolution_clock;
-    auto t0 = ClockT::now();
-
-    computeForEach(
-        useCuda, width * height, 512, __FILE__, __LINE__, [renderOp, image, grid] __hostdev__(int start, int end) {
-            renderOp(start, end, image, grid);
-        });
-    computeSync(useCuda, __FILE__, __LINE__);
-
-    auto t1 = ClockT::now();
-    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count() / 1000.f;
-    return duration;
-}
+struct RenderOp;
+template<typename GridT>
+__global__ void renderIsoSurfacePersistentKernel(RenderOp renderOp, float* image, const GridT* grid, int numPixels, int* nextPixel);
 
 struct RenderOp
 {
@@ -69,18 +56,57 @@ struct RenderOp
     {
         static_assert(nanovdb::util::is_same<typename GridT::BuildType, float, nanovdb::ValueOnIndex, nanovdb::ValueIndex>::value, "only works for float and OnIndex grids");
         auto acc = nanovdb::getAccessor<GridT, float>(*grid);
-        float  t0, v;
-        nanovdb::Coord ijk;
         for (int i = start; i < end; ++i) {
-            RayT iRay = this->getIndexRay(i, grid);   
-            if (nanovdb::math::isoCrossing(iRay, acc, ijk, v, t0, mIso)) {// intersect...
-                this->composite(image, i, (t0 * mDx) / (mWBBoxDimZ * 2), 1.0f);
-            } else {
-                this->composite(image, i, 0.0f, 0.0f);// write background value.
-            }
+            this->renderPixel(i, image, grid, acc);
         }
     }
 
+    template <typename GridT, typename AccT>
+    inline __hostdev__ void renderPixel(int i, float* image, const GridT* grid, AccT& acc) const
+    {
+        float          t0, v;
+        nanovdb::Coord ijk;
+        RayT           iRay = this->getIndexRay(i, grid);
+        if (nanovdb::math::isoCrossing(iRay, acc, ijk, v, t0, mIso)) {// intersect...
+            this->composite(image, i, (t0 * mDx) / (mWBBoxDimZ * 2), 1.0f);
+        } else {
+            this->composite(image, i, 0.0f, 0.0f);// write background value.
+        }
+    }
+
+    template<typename GridT>
+    inline float renderImagePersistent(float* image, const GridT* grid, int* nextPixel) const
+    {
+        int device = 0;
+        NANOVDB_CUDA_CHECK_ERROR(cudaGetDevice(&device), __FILE__, __LINE__);
+
+        cudaDeviceProp properties;
+        NANOVDB_CUDA_CHECK_ERROR(cudaGetDeviceProperties(&properties, device), __FILE__, __LINE__);
+
+        constexpr int blockSize = 256;
+        // Launch a small, fixed pool of blocks that persists on the GPU and
+        // pulls pixel work from a global counter instead of launching one
+        // logical thread per pixel up front.
+        int           blockCount = properties.multiProcessorCount * 4;
+        if (blockCount < 1) blockCount = 1;
+
+        // Reset the work queue before each timed render. The kernel advances
+        // this counter by one warp of pixels at a time.
+        NANOVDB_CUDA_CHECK_ERROR(cudaMemset(nextPixel, 0, sizeof(int)), __FILE__, __LINE__);
+
+        using ClockT = std::chrono::high_resolution_clock;
+        auto t0 = ClockT::now();
+
+        const int numPixels = mWidth * mHeight;
+        renderIsoSurfacePersistentKernel<GridT><<<blockCount, blockSize>>>(*this, image, grid, numPixels, nextPixel);
+        NANOVDB_CUDA_CHECK_ERROR(cudaGetLastError(), __FILE__, __LINE__);
+        NANOVDB_CUDA_CHECK_ERROR(cudaDeviceSynchronize(), __FILE__, __LINE__);
+
+        auto t1 = ClockT::now();
+        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count() / 1000.f;
+        return duration;
+    }
+
     template <typename GridT>  
     inline __hostdev__ RayT getIndexRay(int i, const GridT *grid) const
     {
@@ -134,3 +160,28 @@ struct RenderOp
         }
     }
 };
+
+template<typename GridT>
+__global__ void renderIsoSurfacePersistentKernel(RenderOp renderOp, float* image, const GridT* grid, int numPixels, int* nextPixel)
+{
+    static_assert(nanovdb::util::is_same<typename GridT::BuildType, float, nanovdb::ValueOnIndex, nanovdb::ValueIndex>::value, "only works for float and OnIndex grids");
+    auto acc = nanovdb::getAccessor<GridT, float>(*grid);
+    const unsigned int lane = threadIdx.x & 31u;
+
+    // Keep the fixed set of launched threads busy until all pixels have been assigned.
+    while (true) {
+        int base = 0;
+        // Each warp asks the shared counter for the next batch of 32 pixels.
+        // Only lane 0 updates the counter; __shfl_sync copies lane 0's result
+        // to the other lanes in the warp.
+        if (lane == 0) base = atomicAdd(nextPixel, 32);
+        base = __shfl_sync(0xFFFFFFFFu, base, 0);
+
+        // Each lane renders one pixel from the batch: lane 0 renders base,
+        // lane 1 renders base + 1, and so on.
+        const int i = base + int(lane);
+        if (i >= numPixels) break;
+
+        renderOp.renderPixel(i, image, grid, acc);
+    }
+}

nanovdb/nanovdb/examples/ex_raytrace_iso_surface/main.cc

@@ -13,14 +13,25 @@ using BufferT = nanovdb::cuda::DeviceBuffer;
 using BufferT = nanovdb::HostBuffer;
 #endif
 
-extern void runNanoVDB(nanovdb::GridHandle<BufferT>& handle, int numIterations, int width, int height, BufferT& imageBuffer);
+extern void runNanoVDB(nanovdb::GridHandle<BufferT>& handle, int numIterations, int width, int height, BufferT& imageBuffer, bool usePersistentThreads);
 
 int main(int ac, char** av)
 {
     try {
+        bool        usePersistentThreads = false;
+        const char* gridName = nullptr;
+        for (int i = 1; i < ac; ++i) {
+            if (std::strcmp(av[i], "--persistent") == 0) {
+                usePersistentThreads = true;
+            } else if (!gridName) {
+                gridName = av[i];
+            } else {
+                throw std::runtime_error("Usage: ex_raytrace_iso_surface [--persistent] [grid.nvdb]");
+            }
+        }
         nanovdb::GridHandle<BufferT> handle;
-        if (ac > 1) {
-            handle = nanovdb::io::readGrid<BufferT>(av[1]);
+        if (gridName) {
+            handle = nanovdb::io::readGrid<BufferT>(gridName);
             std::cout << "Loaded NanoVDB grid[" << handle.gridMetaData()->shortGridName() << "]...\n";
         } else {
             handle = nanovdb::tools::createLevelSetSphere<float, BufferT>(100.0f, nanovdb::Vec3d(-20, 0, 0), 1.0, 3.0, nanovdb::Vec3d(0), "sphere");
@@ -31,7 +42,7 @@ int main(int ac, char** av)
         const int height = 4096;
         BufferT imageBuffer(width * height * sizeof(float));
 
-        runNanoVDB(handle, numIterations, width, height, imageBuffer);
+        runNanoVDB(handle, numIterations, width, height, imageBuffer, usePersistentThreads);
     }
     catch (const std::exception& e) {
         std::cerr << "An exception occurred: \"" << e.what() << "\"" << std::endl;

nanovdb/nanovdb/examples/ex_raytrace_iso_surface/nanovdb.cu

@@ -18,7 +18,7 @@ using BufferT = nanovdb::HostBuffer;
 
 #include "common.h"
 
-void runNanoVDB(nanovdb::GridHandle<BufferT>& handle, int numIterations, int width, int height, BufferT& imageBuffer)
+void runNanoVDB(nanovdb::GridHandle<BufferT>& handle, int numIterations, int width, int height, BufferT& imageBuffer, bool usePersistentThreads)
 {
     float *h_outImage = reinterpret_cast<float*>(imageBuffer.data());
     RenderOp renderOp(handle, width, height);
@@ -35,11 +35,22 @@ void runNanoVDB(nanovdb::GridHandle<BufferT>& handle, int numIterations, int wid
         auto* d_grid = handle.deviceGrid<BuildT>();
         if (!d_grid) throw std::runtime_error("GridHandle does not contain a valid device grid");
         imageBuffer.deviceUpload();
+        float* d_outImage = reinterpret_cast<float*>(imageBuffer.deviceData());
         sum = 0;
-        for (int i = 0; i < numIterations; ++i, sum += renderOp.renderImage(true/*useCuda*/, (float*)imageBuffer.deviceData(), d_grid));
-        std::cout << "Average of " << numIterations << " renderings (NanoVDB-Cuda) = " << (sum/numIterations) << " ms " << std::endl;
-        imageBuffer.deviceDownload();
-        renderOp.saveImage("raytrace_iso_surface-nanovdb-cuda.pfm", (float*)imageBuffer.data());
+        if (usePersistentThreads) {
+            int* d_nextPixel = nullptr;
+            NANOVDB_CUDA_CHECK_ERROR(cudaMalloc(&d_nextPixel, sizeof(int)), __FILE__, __LINE__);
+            for (int i = 0; i < numIterations; ++i, sum += renderOp.renderImagePersistent(d_outImage, d_grid, d_nextPixel));
+            NANOVDB_CUDA_CHECK_ERROR(cudaFree(d_nextPixel), __FILE__, __LINE__);
+            std::cout << "Average of " << numIterations << " renderings (NanoVDB-Cuda-Persistent) = " << (sum/numIterations) << " ms " << std::endl;
+            imageBuffer.deviceDownload();
+            renderOp.saveImage("raytrace_iso_surface-nanovdb-cuda-persistent.pfm", (float*)imageBuffer.data());
+        } else {
+            for (int i = 0; i < numIterations; ++i, sum += renderOp.renderImage(true/*useCuda*/, d_outImage, d_grid));
+            std::cout << "Average of " << numIterations << " renderings (NanoVDB-Cuda) = " << (sum/numIterations) << " ms " << std::endl;
+            imageBuffer.deviceDownload();
+            renderOp.saveImage("raytrace_iso_surface-nanovdb-cuda.pfm", (float*)imageBuffer.data());
+        }
 #endif
     };// kernel