NanoVDB: parallel checksum, cooperative GridStats, coalesced IndexToGrid optimizations

nanovdb/nanovdb/tools/cuda/GridChecksum.cuh

@@ -109,6 +109,89 @@ inline unique_ptr<uint32_t> createCrc32Lut(size_t extra = 0, cudaStream_t stream
     return lut;
 }
 
+/// @brief Cuda kernel computing per-block CRC32 checksums via slicing-by-4.
+/// The 256-entry base LUT is staged into shared memory and three derived
+/// slice tables are built in place, so the four divergent lookups per 4-byte
+/// step hit shared memory and the dependent update chain advances four bytes
+/// per step instead of one. Bit-identical to the byte-serial crc32().
+/// The final block absorbs any remainder of @c totalSize.
+__global__ inline void crc32SlicedKernel(const void *d_data, uint32_t* d_blockCRC, uint64_t blockCount, uint32_t log2BlockSize, uint64_t totalSize, const uint32_t *d_lut)
+{
+    __shared__ uint32_t sLut[4][256];
+    for (uint32_t i = threadIdx.x; i < 256; i += blockDim.x) sLut[0][i] = d_lut[i];
+    __syncthreads();
+    for (int k = 1; k < 4; ++k) {
+        for (uint32_t i = threadIdx.x; i < 256; i += blockDim.x) {
+            const uint32_t c = sLut[k-1][i];
+            sLut[k][i] = (c >> 8) ^ sLut[0][c & 0xffu];
+        }
+        __syncthreads();
+    }
+    const uint64_t tid = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x;
+    if (tid >= blockCount) return;
+    const uint8_t *p = (const uint8_t*)d_data + (tid << log2BlockSize);
+    uint64_t n = uint64_t(1) << log2BlockSize;
+    if (tid + 1 == blockCount) n += totalSize - (blockCount << log2BlockSize);
+    uint32_t crc = ~0u;
+    // blocks start at power-of-two offsets of a 32B-aligned buffer -> 4B aligned
+    const uint32_t *w = (const uint32_t*)p;
+    for (uint64_t i = 0, nw = n >> 2; i < nw; ++i) {
+        const uint32_t x = crc ^ w[i];
+        crc = sLut[3][x & 0xffu] ^ sLut[2][(x >> 8) & 0xffu] ^ sLut[1][(x >> 16) & 0xffu] ^ sLut[0][x >> 24];
+    }
+    for (uint64_t i = n & ~3ull; i < n; ++i) {
+        crc ^= p[i];
+        for (int j = 0; j < 8; ++j) crc = (crc >> 1) ^ (0xEDB88320u & (-(crc & 1u)));
+    }
+    d_blockCRC[tid] = ~crc;
+}
+
+/// @brief y = M x over GF(2), M given as 32 column words
+__host__ __device__ inline uint32_t crc32Gf2MatTimes(const uint32_t *mat, uint32_t vec)
+{
+    uint32_t sum = 0;
+    while (vec) {
+        if (vec & 1u) sum ^= *mat;
+        vec >>= 1;
+        ++mat;
+    }
+    return sum;
+}
+
+/// @brief Build into @c dst[32] the GF(2) operator that advances a CRC past
+/// @c bits zero bits, by binary exponentiation of the single-bit operator
+/// (zlib crc32_combine construction). O(log bits x 32) - a few microseconds on
+/// the host, which is where it is called so the device combine stays a cheap fold.
+__host__ __device__ inline void crc32BuildShiftOp(uint32_t *dst, uint64_t bits)
+{
+    uint32_t oddBuf[32], evenBuf[32];
+    uint32_t *cur = oddBuf, *nxt = evenBuf;
+    cur[0] = 0xEDB88320u;// operator for a single zero bit
+    for (int n = 1; n < 32; ++n) cur[n] = 1u << (n - 1);
+    for (int n = 0; n < 32; ++n) dst[n] = 1u << n;// identity
+    while (bits) {
+        if (bits & 1ull) for (int n = 0; n < 32; ++n) dst[n] = crc32Gf2MatTimes(cur, dst[n]);
+        bits >>= 1;
+        if (bits) { for (int n = 0; n < 32; ++n) nxt[n] = crc32Gf2MatTimes(cur, cur[n]); uint32_t *t = cur; cur = nxt; nxt = t; }
+    }
+}
+
+/// @brief Single-thread kernel folding per-chunk CRCs into the CRC of the whole
+/// stream, using crc(A||B) = shift(crc(A), len(B)) ^ crc(B). The two GF(2) shift
+/// operators - @c d_acc for a full chunk and @c d_accLast for the final,
+/// possibly shorter, chunk - are precomputed on the host, so this is just an
+/// O(chunkCount) fold. Bit-identical to a serial crc32 over the concatenated
+/// stream. (@c d_accLast equals @c d_acc when the last chunk is full.)
+__global__ inline void crc32CombineKernel(const uint32_t *d_chunkCRC, uint64_t chunkCount, const uint32_t *d_acc, const uint32_t *d_accLast, uint32_t *d_crc)
+{
+    uint32_t crc = d_chunkCRC[0];
+    for (uint64_t i = 1; i < chunkCount; ++i) {
+        const uint32_t *op = (i + 1 == chunkCount) ? d_accLast : d_acc;
+        crc = crc32Gf2MatTimes(op, crc) ^ d_chunkCRC[i];
+    }
+    *d_crc = crc;
+}
+
 /// @brief Compute CRC32 checksum of 4K block
 /// @param d_data device pointer to start of data
 /// @param size number of bytes
@@ -121,14 +204,39 @@ inline void blockedCRC32(const void *d_data, size_t size, const uint32_t *d_lut,
     const uint64_t checksumCount = size >> NANOVDB_CRC32_LOG2_BLOCK_SIZE;// 4 KB (4096 byte)
     unique_ptr<uint32_t> buffer(checksumCount, stream);// for checksums of 4 KB blocks
     uint32_t *d_checksums = buffer.get();
-    lambdaKernel<<<blocksPerGrid(checksumCount, threadsPerBlock), threadsPerBlock, 0, stream>>>(checksumCount, [=] __device__(size_t tid) {
-        uint32_t blockSize = 1 << NANOVDB_CRC32_LOG2_BLOCK_SIZE;
-        if (tid+1 == checksumCount) blockSize += size - (checksumCount<<NANOVDB_CRC32_LOG2_BLOCK_SIZE);
-        d_checksums[tid] = crc32((const uint8_t*)d_data + (tid<<NANOVDB_CRC32_LOG2_BLOCK_SIZE), blockSize, d_lut);
-    }); cudaCheckError();
-    lambdaKernel<<<1, 1, 0, stream>>>(1, [=] __device__(size_t) {// Compute CRC32 of all the 4K blocks
-        *d_crc = crc32((const uint8_t*)d_checksums, checksumCount*sizeof(uint32_t), d_lut);
-    }); cudaCheckError();
+    crc32SlicedKernel<<<blocksPerGrid(checksumCount, threadsPerBlock), threadsPerBlock, 0, stream>>>(
+        d_data, d_checksums, checksumCount, NANOVDB_CRC32_LOG2_BLOCK_SIZE, size, d_lut);
+    cudaCheckError();
+    // CRC of the block-checksum array itself. The former single-thread pass
+    // over checksumCount*4 bytes (megabytes for multi-GB grids) is replaced
+    // by parallel per-chunk CRCs plus a GF(2) combine - bit-identical result.
+    const uint64_t checksumBytes = checksumCount*sizeof(uint32_t);
+    constexpr uint64_t log2ChunkSize = 12, chunkSize = uint64_t(1) << log2ChunkSize;
+    if (checksumBytes <= 2*chunkSize) {// small: single-thread CRC is fine
+        lambdaKernel<<<1, 1, 0, stream>>>(1, [=] __device__(size_t) {
+            *d_crc = crc32((const uint8_t*)d_checksums, checksumBytes, d_lut);
+        }); cudaCheckError();
+    } else {
+        const uint64_t chunkCount = checksumBytes >> log2ChunkSize;// final chunk absorbs the remainder
+        const uint64_t lastChunkBytes = chunkSize + (checksumBytes - (chunkCount << log2ChunkSize));
+        unique_ptr<uint32_t> chunkBuffer(chunkCount, stream);
+        uint32_t *d_chunkCRC = chunkBuffer.get();
+        crc32SlicedKernel<<<blocksPerGrid(chunkCount, threadsPerBlock), threadsPerBlock, 0, stream>>>(
+            d_checksums, d_chunkCRC, chunkCount, log2ChunkSize, checksumBytes, d_lut);
+        cudaCheckError();
+        // Precompute the two GF(2) shift operators on the host (data-independent,
+        // ~microseconds) and upload them, instead of rebuilding them on a single
+        // device thread inside the combine - that build dominated the checksum
+        // for small and medium grids.
+        uint32_t hOps[64];
+        crc32BuildShiftOp(hOps,      chunkSize      * 8ull);// operator for a full chunk
+        crc32BuildShiftOp(hOps + 32, lastChunkBytes * 8ull);// operator for the last chunk
+        unique_ptr<uint32_t> opBuffer(64, stream);
+        uint32_t *d_ops = opBuffer.get();
+        cudaCheck(cudaMemcpyAsync(d_ops, hOps, 64*sizeof(uint32_t), cudaMemcpyHostToDevice, stream));
+        crc32CombineKernel<<<1, 1, 0, stream>>>(d_chunkCRC, chunkCount, d_ops, d_ops + 32, d_crc);
+        cudaCheckError();
+    }
 }// void cudaBlockedCRC32(const void *d_data, size_t size, const uint32_t *d_lut, uint32_t *d_crc, cudaStream_t stream)
 
 /// @brief Compute CRC32 checksum of 4K block

nanovdb/nanovdb/tools/cuda/GridStats.cuh

@@ -58,65 +58,156 @@ public:
 
 namespace {// define cuda kernels in an unnamed namespace
 
+// One warp per leaf: lanes stride the 512 voxel slots (mask-gated) and merge
+// their partial statistics through shared memory; lane 0 handles the bbox
+// update and the final store. Launch with 128 threads (4 warps) per block.
 template<typename BuildT, typename StatsT>
 __global__ void processLeaf(NodeManager<BuildT> *d_nodeMgr, StatsT *d_stats)
 {
-    const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+    constexpr uint32_t WarpsPerBlock = 4;
+    __shared__ StatsT sStats[WarpsPerBlock * 32];
+
+    const uint32_t warpID = threadIdx.x >> 5, lane = threadIdx.x & 31u;
+    const uint32_t tid = blockIdx.x * WarpsPerBlock + warpID;// leaf index
     if (tid >= d_nodeMgr->leafCount()) return;
     auto &d_leaf = d_nodeMgr->leaf(tid);
 
-    if (d_leaf.updateBBox()) {// updates active bounding box (also updates data->mFlags) and return true if non-empty
+    bool nonEmpty = false;
+    // updates leaf's active bounding box (also updates data->mFlags), only run on lane 0
+    if (lane == 0) nonEmpty = d_leaf.updateBBox();
+    // broadcast nonEmpty value to all lanes; register-to-register warp shuffle with all lanes participating
+    nonEmpty = __shfl_sync(0xffffffffu, nonEmpty, 0); // 0xffffffffu is "all 32 lanes" mask
+
+    if (nonEmpty) {
         if constexpr(StatsT::hasStats()) {
+            // 1) Per-lane partial: each of the 32 lanes folds its strided share of the
+            //    512 voxel slots (lane L visits L, L+32, L+64, ...) into a local StatsT,
+            //    skipping inactive voxels via the value mask. This is the coalesced pass -
+            //    consecutive lanes touch consecutive slots on each stride.
             StatsT stats;
-            for (auto it = d_leaf.cbeginValueOn(); it; ++it) stats.add(*it);
-            if constexpr(StatsT::hasAverage()) {
-                d_stats[tid] = stats;
-                *reinterpret_cast<uint32_t*>(&d_leaf.mMinimum) = tid;
-            } else {
-                stats.setStats(d_leaf);
+            const auto &mask = d_leaf.valueMask();
+            for (uint32_t i = lane; i < NanoLeaf<BuildT>::SIZE; i += 32)
+                if (mask.isOn(i)) stats.add(d_leaf.getValue(i));
+            // 2) Warp reduction: stage the 32 partials in this warp's slice of shared
+            //    memory, then merge them pairwise in log2(32)=5 steps (16->8->4->2->1).
+            //    StatsT::add is an associative merge (exact for min/max; a Welford
+            //    combine for mean/variance), so the tree order is safe. __syncwarp
+            //    after each step orders the shared-memory writes within the warp.
+            StatsT *sWarp = sStats + (warpID << 5);// this warp's 32-slot scratch region
+            sWarp[lane] = stats;
+            __syncwarp();
+            // pairwise reduction
+            for (uint32_t d = 16; d; d >>= 1) {
+                if (lane < d) sWarp[lane].add(sWarp[lane + d]);
+                __syncwarp();
+            }
+            // 3) Publish: lane 0 now holds the leaf's fully merged stats in sWarp[0].
+            if (lane == 0) {
+                if constexpr(StatsT::hasAverage()) {
+                    // mean/variance: the parent lower node must Welford-merge this leaf's
+                    // full accumulator (count+mean+M2), not just its extrema, so publish
+                    // the whole StatsT into the per-node scratch array at this leaf's slot
+                    // and stash that slot index in mMinimum (reused as a uint32 handle) so
+                    // the parent kernel can locate it.
+                    d_stats[tid] = sWarp[0];
+                    *reinterpret_cast<uint32_t*>(&d_leaf.mMinimum) = tid;
+                } else {
+                    // min/max only: extrema compose directly from child to parent, so write
+                    // them straight into the leaf - no scratch slot needed.
+                    sWarp[0].setStats(d_leaf);
+                }
             }
         }
     }
-    d_leaf.mFlags &= ~uint8_t(1u);// enable rendering
+    if (lane == 0) d_leaf.mFlags &= ~uint8_t(1u);// enable rendering
 }// processLeaf
 
+// One block per internal node: threads stride the node's child table (4096 or
+// 32768 entries) and merge partial bboxes/statistics through shared memory -
+// the former one-thread-per-node kernel serialized up to 32768 child visits
+// per thread and left the device nearly empty at typical node counts.
 template<typename BuildT, typename StatsT, int LEVEL>
 __global__ void processInternal(NodeManager<BuildT> *d_nodeMgr, StatsT *d_stats)
 {
     using ChildT = typename NanoNode<BuildT,LEVEL-1>::type;
-    uint32_t nodeID = blockIdx.x * blockDim.x + threadIdx.x;// thread id (reused below to avoid compiler warning)
+    using NodeT = typename NanoNode<BuildT,LEVEL>::type;
+    constexpr uint32_t Threads = 128;
+    __shared__ StatsT sStats[Threads];
+    __shared__ CoordBBox sBBox[Threads];
+
+    const uint32_t nodeID = blockIdx.x;
+    const uint32_t tID = threadIdx.x;
     if (nodeID >= d_nodeMgr->nodeCount(LEVEL)) return;
     auto &d_node = d_nodeMgr->template node<LEVEL>(nodeID);
-    auto &bbox   = d_node.mBBox;
-    bbox         = CoordBBox();// empty bbox
-    StatsT stats;
 
-    for (auto it = d_node.beginChild(); it; ++it) {
-        auto &child = *it;
-        bbox.expand( child.bbox() );
-        if constexpr(StatsT::hasAverage()) {
-            nodeID = *reinterpret_cast<uint32_t*>(&child.mMinimum);
-            StatsT &s = d_stats[nodeID];
-            s.setStats(child);
-            stats.add(s);
-        } else if constexpr(StatsT::hasMinMax()) {
-            stats.add(child.minimum());
-            stats.add(child.maximum());
+    // 1) Per-thread partial: each thread folds its strided share of the child table
+    //    (NodeT::SIZE = 4096 for lower, 32768 for upper) into a local bbox + stats.
+    //    Children were finalized by an earlier launch (leaves, then lower, then upper),
+    //    so a child's bbox()/stats are valid to read here. Each entry is one of three:
+    CoordBBox bbox;// empty
+    StatsT stats;
+    for (uint32_t i = tID; i < NodeT::SIZE; i += Threads) {
+        if (d_node.childMask().isOn(i)) {
+            // (a) a child node: union its finalized bbox and merge its statistics.
+            auto &child = *d_node.getChild(i);
+            bbox.expand( child.bbox() );
+            if constexpr(StatsT::hasAverage()) {
+                // The child published its full accumulator into d_stats and stashed the
+                // slot index in its mMinimum. We must read that handle FIRST (to index
+                // d_stats), then setStats commits the child's real min/max/avg/std into
+                // the child node - which also overwrites the mMinimum-as-slot handle with
+                // the true minimum. Finally merge the child's accumulator into ours.
+                StatsT &s = d_stats[*reinterpret_cast<const uint32_t*>(&child.mMinimum)];
+                s.setStats(child);
+                stats.add(s);
+            } else if constexpr(StatsT::hasMinMax()) {
+                // min/max compose directly - the child already holds its final extrema.
+                stats.add(child.minimum());
+                stats.add(child.maximum());
+            }
+        } else if (d_node.valueMask().isOn(i)) {
+            // (b) an active tile: one constant value covering the whole child region
+            //     (no child node). Grow the bbox to span the tile's extent, and add the
+            //     value with multiplicity NUM_VALUES (the voxel count it stands in for)
+            //     so the mean/variance are weighted correctly.
+            const Coord ijk = d_node.offsetToGlobalCoord(i);
+            bbox[0].minComponent(ijk);
+            bbox[1].maxComponent(ijk + Coord(ChildT::DIM - 1));
+            if constexpr(StatsT::hasStats()) stats.add(d_node.data()->getValue(i), ChildT::NUM_VALUES);
         }
+        // (c) otherwise inactive - contributes nothing.
     }
-    for (auto it = d_node.cbeginValueOn(); it; ++it) {
-        const Coord ijk = it.getCoord();
-        bbox[0].minComponent(ijk);
-        bbox[1].maxComponent(ijk + Coord(ChildT::DIM - 1));
-        if constexpr(StatsT::hasStats()) stats.add(*it, ChildT::NUM_VALUES);
+    // 2) Block reduction: merge the 128 per-thread partials (both stats and bbox) via a
+    //    shared-memory tree in log2(128)=7 steps. This spans 4 warps, so it needs
+    //    __syncthreads (block-wide barrier), not the __syncwarp used in processLeaf.
+    sStats[tID] = stats;
+    sBBox[tID] = bbox;
+    __syncthreads();
+    for (uint32_t d = Threads >> 1; d; d >>= 1) {
+        if (tID < d) {
+            sStats[tID].add(sStats[tID + d]);
+            sBBox[tID].expand(sBBox[tID + d]);
+        }
+        __syncthreads();
     }
-    if constexpr(StatsT::hasAverage()) {
-        d_stats[nodeID] = stats;
-        *reinterpret_cast<uint32_t*>(&d_node.mMinimum) = nodeID;
-    } else if constexpr(StatsT::hasMinMax()) {
-        stats.setStats(d_node);
+    // 3) Publish: thread 0 now holds the node's merged bbox + stats.
+    if (tID == 0) {
+        d_node.mBBox = sBBox[0];
+        if constexpr(StatsT::hasAverage()) {
+            // Each node writes its OWN unique d_stats slot (leaves occupy
+            // [0, leafCount), then lower nodes, then upper nodes), so a node
+            // with no children still has a valid slot. The previous scheme
+            // reused a child's slot, writing d_stats[-1] for a childless
+            // (fully-tiled) internal node - an out-of-bounds write.
+            const uint32_t slot = d_nodeMgr->leafCount()
+                                + (LEVEL == 2 ? d_nodeMgr->nodeCount(1) : 0u) + nodeID;
+            d_stats[slot] = sStats[0];
+            *reinterpret_cast<uint32_t*>(&d_node.mMinimum) = slot;
+        } else if constexpr(StatsT::hasMinMax()) {
+            sStats[0].setStats(d_node);
+        }
+        d_node.mFlags &= ~uint64_t(1u);// enable rendering
     }
-    d_node.mFlags &= ~uint64_t(1u);// enable rendering
 }// processInternal
 
 template<typename BuildT, typename StatsT>
@@ -199,13 +290,16 @@ void GridStats<BuildT, StatsT>::update(NanoGrid<BuildT> *d_grid, cudaStream_t st
 
     StatsT *d_stats = nullptr;
 
-    if constexpr(StatsT::hasAverage()) cudaCheck(util::cuda::mallocAsync((void**)&d_stats, nodeCount[0]*sizeof(StatsT), stream));
+    // One d_stats slot per node (leaves, then lower, then upper) so every node
+    // has its own slot - see processInternal.
+    if constexpr(StatsT::hasAverage()) cudaCheck(util::cuda::mallocAsync((void**)&d_stats, (nodeCount[0]+nodeCount[1]+nodeCount[2])*sizeof(StatsT), stream));
 
-    processLeaf<BuildT><<<blocksPerGrid(nodeCount[0]), threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
+    // warp per leaf (4 warps per 128-thread block); block per internal node
+    if (nodeCount[0]) processLeaf<BuildT><<<blocksPerGrid(nodeCount[0]*32), threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
 
-    processInternal<BuildT, StatsT, 1><<<blocksPerGrid(nodeCount[1]), threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
+    if (nodeCount[1]) processInternal<BuildT, StatsT, 1><<<nodeCount[1], threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
 
-    processInternal<BuildT, StatsT, 2><<<blocksPerGrid(nodeCount[2]), threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
+    if (nodeCount[2]) processInternal<BuildT, StatsT, 2><<<nodeCount[2], threadsPerBlock, 0, stream>>>(d_nodeMgr, d_stats);
 
     processRootAndGrid<BuildT><<<1, 1, 0, stream>>>(d_nodeMgr, d_stats);