Skip to content

Commit 0adc072

Browse files
swahtzclaude
swahtz
and
claude
committed
NanoVDB CUDA: consolidate VBM decode to a thread-local per-slot API
Replace the cooperative shared-memory decodeInverseMaps with decodeInverseMap: a per-slot decode into registers (no shared memory or synchronization; callable from divergent threads), factored into jumpMapRank + selectVoxelInLeaf helpers. computeBoxStencil now takes the decoded (leafIndex, voxelOffset) by value. Block-level facts formerly read from the materialized maps derive directly from the VBM metadata: the block's first leaf is firstLeafID, slot p starts a new leaf iff jumpMap bit p is set, and the spanned-leaf count is 1 + the jumpMap popcount. Decode maps and all 27 box-stencil taps verified byte-exact against the prior select decode at widths 64/128/512 across 16-86% leaf occupancy. Decode-only +14..17% at width 128 and +8..13% at width 512 over the prior select, neutral at width 64; vs master's sweep: 1.56-2.08x (w7) and 1.24-1.68x (w9) across the sparsity range. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>
1 parent ab56089 commit 0adc072

3 files changed

Lines changed: 132 additions & 103 deletions

File tree

‎nanovdb/nanovdb/tools/cuda/VoxelBlockManager.cuh‎

Lines changed: 117 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,14 @@
1717
of an OnIndexGrid, independent of occupancy. This file provides:
1818
- buildVoxelBlockManager (device): constructs the firstLeafID array and
1919
jumpMap on the GPU from a device-resident NanoGrid.
20-
- decodeInverseMaps (device): per-block SIMT decode of the inverse maps
21-
(sequential active-voxel index -> leaf ID + intra-leaf voxel offset),
22-
with one thread per output slot across a CUDA thread block.
20+
- decodeInverseMap (device): per-slot, thread-local decode of one inverse
21+
map entry (sequential active-voxel index -> leaf ID + intra-leaf voxel
22+
offset) into registers; no shared memory or synchronization, callable
23+
from divergent threads. Block-level facts a cooperative consumer might
24+
otherwise read from a materialized map are derivable directly from the
25+
VBM metadata: the block's first leaf is firstLeafID, slot p starts a
26+
new leaf iff jumpMap bit p is set, and the number of leaves spanned by
27+
the block is 1 + the jumpMap popcount.
2328
*/
2429

2530
#ifndef NANOVDB_VOXELBLOCKMANAGER_CUH_HAS_BEEN_INCLUDED
@@ -48,138 +53,153 @@ struct VoxelBlockManager : nanovdb::tools::VoxelBlockManagerBase<Log2BlockWidth>
4853
using Base::UnusedLeafIndex;
4954
using Base::UnusedVoxelOffset;
5055

51-
// The efficiency of the functions in this class are contingent on
52-
// threadblock-level coordination, which manifests either as using shared
53-
// memory for synchronization, or warp-level shift operations.
56+
// The decode is a rank + select over the VBM's bit-vectors and is
57+
// thread-local per output slot: decodeInverseMap requires no threadblock
58+
// coordination. Consumers that need cross-slot information derive it from
59+
// the VBM metadata (firstLeafID + jumpMap) rather than from a materialized
60+
// per-block map.
5461

55-
/// @brief Decode the inverse maps for a single voxel block on the device.
62+
/// @brief Rank one output slot into its leaf: the number of leaves that
63+
/// begin at in-block positions [1, blockOffset] (bit 0 is never set),
64+
/// counted via popcounts over the block's jumpMap.
65+
/// Thread-local; no synchronization.
66+
__device__
67+
static uint32_t jumpMapRank(const uint64_t *jumpMap, const int blockOffset)
68+
{
69+
uint32_t leafRank = 0;
70+
const int jumpWord = blockOffset >> 6; // index into jumpMap
71+
#pragma unroll
72+
for (int i = 0; i < JumpMapLength; ++i) {
73+
if (i < jumpWord) leafRank += util::countOn(jumpMap[i]); // count leaves before the current jump word
74+
else if (i == jumpWord) {
75+
// count leaves in the current jump word, masking those that are at or before blockOffset
76+
leafRank += util::countOn(jumpMap[i] & ((uint64_t(2) << (blockOffset & 63)) - 1u));
77+
}
78+
}
79+
return leafRank;
80+
}
81+
82+
/// @brief Select the voxel with sequential index @c globalOffset inside leaf
83+
/// @c leafID: find its 64-bit mask word via the leaf's precomputed 9-bit
84+
/// prefix sums (mPrefixSum), then its bit within that word via __fns.
85+
/// Writes the sentinels if globalOffset lies beyond the leaf's last active
86+
/// voxel (i.e. beyond the last active voxel of the grid).
87+
/// Thread-local; no synchronization.
88+
template <class BuildT>
89+
__device__
90+
static typename util::enable_if<BuildTraits<BuildT>::is_index, void>::type
91+
selectVoxelInLeaf(
92+
const NanoGrid<BuildT> *grid,
93+
const uint32_t leafID,
94+
const uint64_t globalOffset,
95+
uint32_t &leafIndex,
96+
uint16_t &voxelOffset)
97+
{
98+
const auto& leafData = *grid->tree().template getFirstNode<0>()[leafID].data();
99+
// 0-based rank among the leaf's actives
100+
const uint64_t rankInLeaf = globalOffset - leafData.firstOffset();
101+
if (rankInLeaf < leafData.valueCount()) { // if the rank is within the leaf's active voxels (bounds check)
102+
int wordID = 0;
103+
uint32_t activesBeforeWord = 0;
104+
#pragma unroll
105+
for (int candidateWord = 1; candidateWord < 8; ++candidateWord) {
106+
// the number of active voxels before the candidateWord's mask word (& 0x1ffu masks to 9 bits)
107+
const uint32_t cumulative = uint32_t(leafData.mPrefixSum >> (9*(candidateWord-1))) & 0x1ffu;
108+
if (cumulative <= rankInLeaf) {
109+
wordID = candidateWord; // word ID of the mask word that contains the rankInLeaf-th active voxel
110+
activesBeforeWord = cumulative; // the number of active voxels before the wordID's mask word
111+
}
112+
}
113+
uint32_t rankInWord = uint32_t(rankInLeaf) - activesBeforeWord; // active voxel's rank within the mask word (0-based)
114+
// select the in-word bit position of the voxel using __fns (find n-th set bit)
115+
// __fns(mask, base, k) is the hardware find-nth-set-bit intrinsic - the k-th set bit (1-based) at/after base
116+
// but it's a 32-bit op while `maskWord` is 64-bit, so we need to split it into two 32-bit halves
117+
const uint64_t maskWord = leafData.mValueMask.words()[wordID];
118+
const uint32_t lowHalf = uint32_t(maskWord); // low 32 bits of the mask word
119+
const uint32_t lowHalfCount = __popc(lowHalf);
120+
int bit;
121+
// if rank is less than the number of active voxels in the low half, __fns finds the bit in the lower half
122+
// otherwise, shift maskWord by 32 bits and __fns finds the bit in the upper half
123+
if (rankInWord < lowHalfCount) bit = __fns(lowHalf, 0, rankInWord + 1);
124+
else bit = 32 + __fns(uint32_t(maskWord >> 32), 0, rankInWord - lowHalfCount + 1);
125+
leafIndex = leafID;
126+
voxelOffset = uint16_t((wordID << 6) + bit);
127+
} else { // beyond the last active voxel in the grid
128+
leafIndex = UnusedLeafIndex;
129+
voxelOffset = UnusedVoxelOffset;
130+
}
131+
}
132+
133+
/// @brief Decode a single inverse-map entry into registers on the device.
56134
///
57135
/// Given the VBM metadata for one block (firstLeafID and the block's slice of
58-
/// the jumpMap) and the block's base sequential offset, fills smem_leafIndex[]
59-
/// and smem_voxelOffset[] in shared memory so that for each position p in
60-
/// [0, BlockWidth):
61-
/// - smem_leafIndex[p] = index of the leaf node containing sequential voxel
62-
/// (blockFirstOffset + p), or UnusedLeafIndex if that
63-
/// index is beyond the last active voxel.
64-
/// - smem_voxelOffset[p] = local (0..511) offset of that voxel within its leaf,
65-
/// or UnusedVoxelOffset.
136+
/// the jumpMap), the block's base sequential offset, and a slot position
137+
/// blockOffset in [0, BlockWidth), computes:
138+
/// - leafIndex = index of the leaf node containing sequential voxel
139+
/// (blockFirstOffset + blockOffset), or UnusedLeafIndex if
140+
/// that index is beyond the last active voxel.
141+
/// - voxelOffset = local (0..511) offset of that voxel within its leaf,
142+
/// or UnusedVoxelOffset.
66143
///
67-
/// Must be called by all threads in the block (uses __syncthreads internally).
68-
/// Do not call from divergent threads within a thread block.
144+
/// No shared memory or synchronization; may be called from divergent threads.
69145
///
70146
/// @tparam BuildT Build type of the grid (must be an index type)
71147
/// @param grid Device-accessible OnIndex grid
72148
/// @param firstLeafID Index of the first leaf overlapping this block
73149
/// @param jumpMap Pointer to the JumpMapLength words for this block
74150
/// @param blockFirstOffset Sequential index of the first voxel in this block
75-
/// @param smem_leafIndex Output array of length BlockWidth in shared memory
76-
/// @param smem_voxelOffset Output array of length BlockWidth in shared memory
151+
/// @param blockOffset Slot position within the block, in [0, BlockWidth)
152+
/// @param leafIndex Output leaf index (register)
153+
/// @param voxelOffset Output intra-leaf voxel offset (register)
77154
template <class BuildT>
78155
__device__
79156
static typename util::enable_if<BuildTraits<BuildT>::is_index, void>::type
80-
decodeInverseMaps(
157+
decodeInverseMap(
81158
const NanoGrid<BuildT> *grid,
82159
const uint32_t firstLeafID,
83160
const uint64_t *jumpMap,
84161
const uint64_t blockFirstOffset,
85-
uint32_t *smem_leafIndex,
86-
uint16_t *smem_voxelOffset)
162+
const int blockOffset,
163+
uint32_t &leafIndex,
164+
uint16_t &voxelOffset)
87165
{
88166
// Verify that the nodes can be accessed linearly
89167
NANOVDB_ASSERT(grid->isSequential());
90-
NANOVDB_ASSERT(blockDim.x <= 512);
91-
92-
// Select-based decode: one thread per output slot. Here each
93-
// slot ranks itself into its leaf via the jumpMap popcount,
94-
// then locates its voxel with the leaf's 9-bit prefix sums
95-
// plus an in-word bit select - O(1) per slot.
96-
const int tID = threadIdx.x;
97-
const auto *leaf0 = grid->tree().template getFirstNode<0>();
98-
for (int blockOffset = tID; blockOffset < BlockWidth; blockOffset += blockDim.x) {
99-
// rank this slot into its leaf: count leaves beginning at in-block positions
100-
// [1, blockOffset] (bit 0 is never set) via the jumpMap popcount
101-
uint32_t leafRank = 0;
102-
const int jumpWord = blockOffset >> 6; // index into jumpMap
103-
// count the number of leaves that begin before blockOffset
104-
#pragma unroll
105-
for (int i = 0; i < JumpMapLength; ++i) {
106-
if (i < jumpWord) leafRank += util::countOn(jumpMap[i]); // count leaves before the current jump word
107-
else if (i == jumpWord) {
108-
// count leaves in the current jump word, masking those that are before blockOffset
109-
leafRank += util::countOn(jumpMap[i] & ((uint64_t(2) << (blockOffset & 63)) - 1u));
110-
}
168+
NANOVDB_ASSERT(blockOffset >= 0 && blockOffset < BlockWidth);
111169

112-
}
113-
const uint32_t leafID = firstLeafID + leafRank;
114-
const auto& leafData = *leaf0[leafID].data();
115-
// 0-based rank among the leaf's actives
116-
const uint64_t rankInLeaf = (blockFirstOffset + blockOffset) - leafData.firstOffset();
117-
if (rankInLeaf < leafData.valueCount()) { // if the rank is within the leaf's active voxels (bounds check)
118-
// select the rankInLeaf-th active voxel: find its 64-bit mask word via the
119-
// leaf's 9-bit prefix sums, then its bit within that word
120-
int wordID = 0;
121-
uint32_t activesBeforeWord = 0;
122-
#pragma unroll
123-
for (int candidateWord = 1; candidateWord < 8; ++candidateWord) {
124-
// the number of active voxels before the candidateWord's mask word (& 0x1ffu masks to 9 bits)
125-
const uint32_t cumulative = uint32_t(leafData.mPrefixSum >> (9*(candidateWord-1))) & 0x1ffu;
126-
if (cumulative <= rankInLeaf) {
127-
wordID = candidateWord; // word ID of the mask word that contains the rankInLeaf-th active voxel
128-
activesBeforeWord = cumulative; // the number of active voxels before the wordID's mask word
129-
}
130-
}
131-
132-
uint32_t rankInWord = uint32_t(rankInLeaf) - activesBeforeWord; // active voxel's rank within the mask word (0-based)
133-
// select the in-word bit position of the voxel using __fns (find n-th set bit)
134-
// /__fns(mask, base, k) is the hardware find-nth-set-bit intrinsic - the k-th set bit (1-based) at/after base
135-
// but it's a 32-bit op while `maskWord` is 64-bit, so we need to split it into two 32-bit halves
136-
const uint64_t maskWord = leafData.mValueMask.words()[wordID];
137-
const uint32_t lowHalf = uint32_t(maskWord); // low 32 bits of the mask word
138-
const uint32_t lowHalfCount = util::countOn(uint64_t(lowHalf));
139-
int bit;
140-
// if rank is less than the number of active voxels in the low half, __fns finds the bit in the lower half
141-
// otherwise, shift maskWord by 32 bits and __fns finds the bit in the upper half
142-
if (rankInWord < lowHalfCount) bit = __fns(lowHalf, 0, rankInWord + 1);
143-
else bit = 32 + __fns(uint32_t(maskWord >> 32), 0, rankInWord - lowHalfCount + 1);
144-
smem_leafIndex[blockOffset] = leafID;
145-
smem_voxelOffset[blockOffset] = uint16_t((wordID << 6) + bit);
146-
} else {// beyond the last active voxel in the grid
147-
smem_leafIndex[blockOffset] = UnusedLeafIndex;
148-
smem_voxelOffset[blockOffset] = UnusedVoxelOffset;
149-
}
150-
}
151-
__syncthreads();
170+
const uint32_t leafRank = jumpMapRank(jumpMap, blockOffset);
171+
selectVoxelInLeaf(grid, firstLeafID + leafRank,
172+
blockFirstOffset + blockOffset, leafIndex, voxelOffset);
152173
}
153174

154-
/// @brief Given a grid and its decoded voxel map, compute the stencil.
155-
/// This function accesses shared memory but does not synchronize threads
156-
/// so it may be called from divergent threads within a thread block.
157-
/// offsets for a 3x3x3 box stencil.
175+
/// @brief Given a grid and one decoded inverse-map entry (from
176+
/// decodeInverseMap), compute the stencil indices for a 3x3x3 box stencil.
177+
/// Thread-local: no shared memory, no synchronization; may be called from
178+
/// divergent threads within a thread block. Leaves stencilIndices untouched
179+
/// when leafIndex is UnusedLeafIndex (a slot beyond the last active voxel).
158180
/// @tparam BuildT Build type of the grid
159-
/// @param grid
160-
/// @param smem_leafIndex Leaf indices stored in shared memory
161-
/// @param smem_voxelOffset Voxel offsets stored in shared memory
181+
/// @param grid Device-accessible OnIndex grid
182+
/// @param leafIndex This thread's decoded leaf index (or UnusedLeafIndex)
183+
/// @param voxelOffset This thread's decoded intra-leaf voxel offset
162184
/// @param stencilIndices Pointer to output stencil indices. Must have
163185
/// length of at least 27 (corresponding to the 3x3x3 stencil)
164186
template <class BuildT>
165187
__device__
166188
static typename util::enable_if<BuildTraits<BuildT>::is_index, void>::type
167189
computeBoxStencil(
168190
const NanoGrid<BuildT> *grid,
169-
const uint32_t *smem_leafIndex,
170-
const uint16_t *smem_voxelOffset,
191+
const uint32_t leafIndex,
192+
const uint16_t voxelOffset,
171193
uint64_t *stencilIndices)
172194
{
173195
// Verify that the nodes can be accessed linearly
174196
NANOVDB_ASSERT(grid->isSequential());
175197

176-
int tID = threadIdx.x;
177198
const auto& tree = grid->tree();
178-
if (smem_leafIndex[tID] != UnusedLeafIndex) {
199+
if (leafIndex != UnusedLeafIndex) {
179200
// This presumes that leaf nodes are fixed-size and sequentially accessible in memory
180-
const auto& leaf = tree.template getFirstNode<0>()[ smem_leafIndex[tID] ];
181-
const Coord coord = leaf.offsetToGlobalCoord( smem_voxelOffset[tID] );
182-
const auto index = leaf.getValue( smem_voxelOffset[tID] );
201+
const auto& leaf = tree.template getFirstNode<0>()[ leafIndex ];
202+
const Coord coord = leaf.offsetToGlobalCoord( voxelOffset );
183203
for (int di = -1; di <= 1; di++)
184204
for (int dj = -1; dj <= 1; dj++)
185205
for (int dk = -1; dk <= 1; dk++) {

‎nanovdb/nanovdb/unittest/TestNanoVDB.cu‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3465,16 +3465,15 @@ __global__ void testComputeStencilNeighborsKernel(
34653465
uint64_t *jumpMap = jumpMapArray + JumpMapLength * bID;
34663466
int firstOffset = 1;
34673467
int blockFirstOffset = firstOffset + bID * BlockWidth;
3468-
__shared__ uint32_t leafIndex[BlockWidth];
3469-
__shared__ uint16_t voxelOffset[BlockWidth];
3468+
uint32_t leafIndex;
3469+
uint16_t voxelOffset;
34703470

3471-
nanovdb::tools::cuda::VoxelBlockManager<Log2BlockWidth>::decodeInverseMaps(
3472-
grid, firstLeafID, jumpMap, blockFirstOffset, &leafIndex[0], &voxelOffset[0]);
3471+
nanovdb::tools::cuda::VoxelBlockManager<Log2BlockWidth>::decodeInverseMap(
3472+
grid, firstLeafID, jumpMap, blockFirstOffset, tID, leafIndex, voxelOffset);
34733473

34743474
uint64_t localNeighbors[27] = {};
34753475
nanovdb::tools::cuda::VoxelBlockManager<Log2BlockWidth>::computeBoxStencil(
3476-
grid, &leafIndex[0], &voxelOffset[0], localNeighbors);
3477-
__syncthreads();
3476+
grid, leafIndex, voxelOffset, localNeighbors);
34783477

34793478
using StencilNeighborsType = uint64_t (*)[27];
34803479
auto stencilNeighbors = reinterpret_cast<StencilNeighborsType>(stencilNeighborsArray+27*BlockWidth*bID);

‎pendingchanges/nanovdbvbmdecode.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,13 @@ NanoVDB:
33
- Replaced the CUDA VoxelBlockManager inverse-map decode's O(nLeaves x 512) per-block
44
sweep with an O(1)-per-slot select (jumpMap popcount rank + the leaf's 9-bit prefix
55
sums + an in-word bit select); byte-identical decode maps and ~2x faster decode.
6+
API Changes:
7+
- Consolidated the CUDA VoxelBlockManager decode to a single thread-local API:
8+
decodeInverseMap decodes one slot into registers (no shared memory or
9+
synchronization; callable from divergent threads) and replaces the cooperative
10+
shared-memory decodeInverseMaps, which has been removed. computeBoxStencil now
11+
takes the decoded (leafIndex, voxelOffset) by value instead of shared-memory
12+
arrays. Block-level facts formerly read from the materialized maps derive
13+
directly from the VBM metadata: the block's first leaf is firstLeafID, slot p
14+
starts a new leaf iff jumpMap bit p is set, and the spanned-leaf count is
15+
1 + the jumpMap popcount.

0 commit comments

Comments
 (0)