@@ -58,65 +58,156 @@ public:
5858
5959namespace {// define cuda kernels in an unnamed namespace
6060
61+ // One warp per leaf: lanes stride the 512 voxel slots (mask-gated) and merge
62+ // their partial statistics through shared memory; lane 0 handles the bbox
63+ // update and the final store. Launch with 128 threads (4 warps) per block.
6164template <typename BuildT, typename StatsT>
6265__global__ void processLeaf (NodeManager<BuildT> *d_nodeMgr, StatsT *d_stats)
6366{
64- const uint32_t tid = blockIdx .x * blockDim .x + threadIdx .x ;
67+ constexpr uint32_t WarpsPerBlock = 4 ;
68+ __shared__ StatsT sStats [WarpsPerBlock * 32 ];
69+
70+ const uint32_t warpID = threadIdx .x >> 5 , lane = threadIdx .x & 31u ;
71+ const uint32_t tid = blockIdx .x * WarpsPerBlock + warpID;// leaf index
6572 if (tid >= d_nodeMgr->leafCount ()) return ;
6673 auto &d_leaf = d_nodeMgr->leaf (tid);
6774
68- if (d_leaf.updateBBox ()) {// updates active bounding box (also updates data->mFlags) and return true if non-empty
75+ bool nonEmpty = false ;
76+ // updates leaf's active bounding box (also updates data->mFlags), only run on lane 0
77+ if (lane == 0 ) nonEmpty = d_leaf.updateBBox ();
78+ // broadcast nonEmpty value to all lanes; register-to-register warp shuffle with all lanes participating
79+ nonEmpty = __shfl_sync (0xffffffffu , nonEmpty, 0 ); // 0xffffffffu is "all 32 lanes" mask
80+
81+ if (nonEmpty) {
6982 if constexpr (StatsT::hasStats ()) {
83+ // 1) Per-lane partial: each of the 32 lanes folds its strided share of the
84+ // 512 voxel slots (lane L visits L, L+32, L+64, ...) into a local StatsT,
85+ // skipping inactive voxels via the value mask. This is the coalesced pass -
86+ // consecutive lanes touch consecutive slots on each stride.
7087 StatsT stats;
71- for (auto it = d_leaf.cbeginValueOn (); it; ++it) stats.add (*it);
72- if constexpr (StatsT::hasAverage ()) {
73- d_stats[tid] = stats;
74- *reinterpret_cast <uint32_t *>(&d_leaf.mMinimum ) = tid;
75- } else {
76- stats.setStats (d_leaf);
88+ const auto &mask = d_leaf.valueMask ();
89+ for (uint32_t i = lane; i < NanoLeaf<BuildT>::SIZE ; i += 32 )
90+ if (mask.isOn (i)) stats.add (d_leaf.getValue (i));
91+ // 2) Warp reduction: stage the 32 partials in this warp's slice of shared
92+ // memory, then merge them pairwise in log2(32)=5 steps (16->8->4->2->1).
93+ // StatsT::add is an associative merge (exact for min/max; a Welford
94+ // combine for mean/variance), so the tree order is safe. __syncwarp
95+ // after each step orders the shared-memory writes within the warp.
96+ StatsT *sWarp = sStats + (warpID << 5 );// this warp's 32-slot scratch region
97+ sWarp [lane] = stats;
98+ __syncwarp ();
99+ // pairwise reduction
100+ for (uint32_t d = 16 ; d; d >>= 1 ) {
101+ if (lane < d) sWarp [lane].add (sWarp [lane + d]);
102+ __syncwarp ();
103+ }
104+ // 3) Publish: lane 0 now holds the leaf's fully merged stats in sWarp[0].
105+ if (lane == 0 ) {
106+ if constexpr (StatsT::hasAverage ()) {
107+ // mean/variance: the parent lower node must Welford-merge this leaf's
108+ // full accumulator (count+mean+M2), not just its extrema, so publish
109+ // the whole StatsT into the per-node scratch array at this leaf's slot
110+ // and stash that slot index in mMinimum (reused as a uint32 handle) so
111+ // the parent kernel can locate it.
112+ d_stats[tid] = sWarp [0 ];
113+ *reinterpret_cast <uint32_t *>(&d_leaf.mMinimum ) = tid;
114+ } else {
115+ // min/max only: extrema compose directly from child to parent, so write
116+ // them straight into the leaf - no scratch slot needed.
117+ sWarp [0 ].setStats (d_leaf);
118+ }
77119 }
78120 }
79121 }
80- d_leaf.mFlags &= ~uint8_t (1u );// enable rendering
122+ if (lane == 0 ) d_leaf.mFlags &= ~uint8_t (1u );// enable rendering
81123}// processLeaf
82124
125+ // One block per internal node: threads stride the node's child table (4096 or
126+ // 32768 entries) and merge partial bboxes/statistics through shared memory -
127+ // the former one-thread-per-node kernel serialized up to 32768 child visits
128+ // per thread and left the device nearly empty at typical node counts.
83129template <typename BuildT, typename StatsT, int LEVEL >
84130__global__ void processInternal (NodeManager<BuildT> *d_nodeMgr, StatsT *d_stats)
85131{
86132 using ChildT = typename NanoNode<BuildT,LEVEL -1 >::type;
87- uint32_t nodeID = blockIdx .x * blockDim .x + threadIdx .x ;// thread id (reused below to avoid compiler warning)
133+ using NodeT = typename NanoNode<BuildT,LEVEL >::type;
134+ constexpr uint32_t Threads = 128 ;
135+ __shared__ StatsT sStats [Threads];
136+ __shared__ CoordBBox sBBox [Threads];
137+
138+ const uint32_t nodeID = blockIdx .x ;
139+ const uint32_t tID = threadIdx .x ;
88140 if (nodeID >= d_nodeMgr->nodeCount (LEVEL )) return ;
89141 auto &d_node = d_nodeMgr->template node <LEVEL >(nodeID);
90- auto &bbox = d_node.mBBox ;
91- bbox = CoordBBox ();// empty bbox
92- StatsT stats;
93142
94- for (auto it = d_node.beginChild (); it; ++it) {
95- auto &child = *it;
96- bbox.expand ( child.bbox () );
97- if constexpr (StatsT::hasAverage ()) {
98- nodeID = *reinterpret_cast <uint32_t *>(&child.mMinimum );
99- StatsT &s = d_stats[nodeID];
100- s.setStats (child);
101- stats.add (s);
102- } else if constexpr (StatsT::hasMinMax ()) {
103- stats.add (child.minimum ());
104- stats.add (child.maximum ());
143+ // 1) Per-thread partial: each thread folds its strided share of the child table
144+ // (NodeT::SIZE = 4096 for lower, 32768 for upper) into a local bbox + stats.
145+ // Children were finalized by an earlier launch (leaves, then lower, then upper),
146+ // so a child's bbox()/stats are valid to read here. Each entry is one of three:
147+ CoordBBox bbox;// empty
148+ StatsT stats;
149+ for (uint32_t i = tID; i < NodeT::SIZE ; i += Threads) {
150+ if (d_node.childMask ().isOn (i)) {
151+ // (a) a child node: union its finalized bbox and merge its statistics.
152+ auto &child = *d_node.getChild (i);
153+ bbox.expand ( child.bbox () );
154+ if constexpr (StatsT::hasAverage ()) {
155+ // The child published its full accumulator into d_stats and stashed the
156+ // slot index in its mMinimum. We must read that handle FIRST (to index
157+ // d_stats), then setStats commits the child's real min/max/avg/std into
158+ // the child node - which also overwrites the mMinimum-as-slot handle with
159+ // the true minimum. Finally merge the child's accumulator into ours.
160+ StatsT &s = d_stats[*reinterpret_cast <const uint32_t *>(&child.mMinimum )];
161+ s.setStats (child);
162+ stats.add (s);
163+ } else if constexpr (StatsT::hasMinMax ()) {
164+ // min/max compose directly - the child already holds its final extrema.
165+ stats.add (child.minimum ());
166+ stats.add (child.maximum ());
167+ }
168+ } else if (d_node.valueMask ().isOn (i)) {
169+ // (b) an active tile: one constant value covering the whole child region
170+ // (no child node). Grow the bbox to span the tile's extent, and add the
171+ // value with multiplicity NUM_VALUES (the voxel count it stands in for)
172+ // so the mean/variance are weighted correctly.
173+ const Coord ijk = d_node.offsetToGlobalCoord (i);
174+ bbox[0 ].minComponent (ijk);
175+ bbox[1 ].maxComponent (ijk + Coord (ChildT::DIM - 1 ));
176+ if constexpr (StatsT::hasStats ()) stats.add (d_node.data ()->getValue (i), ChildT::NUM_VALUES );
105177 }
178+ // (c) otherwise inactive - contributes nothing.
106179 }
107- for (auto it = d_node.cbeginValueOn (); it; ++it) {
108- const Coord ijk = it.getCoord ();
109- bbox[0 ].minComponent (ijk);
110- bbox[1 ].maxComponent (ijk + Coord (ChildT::DIM - 1 ));
111- if constexpr (StatsT::hasStats ()) stats.add (*it, ChildT::NUM_VALUES );
180+ // 2) Block reduction: merge the 128 per-thread partials (both stats and bbox) via a
181+ // shared-memory tree in log2(128)=7 steps. This spans 4 warps, so it needs
182+ // __syncthreads (block-wide barrier), not the __syncwarp used in processLeaf.
183+ sStats [tID] = stats;
184+ sBBox [tID] = bbox;
185+ __syncthreads ();
186+ for (uint32_t d = Threads >> 1 ; d; d >>= 1 ) {
187+ if (tID < d) {
188+ sStats [tID].add (sStats [tID + d]);
189+ sBBox [tID].expand (sBBox [tID + d]);
190+ }
191+ __syncthreads ();
112192 }
113- if constexpr (StatsT::hasAverage ()) {
114- d_stats[nodeID] = stats;
115- *reinterpret_cast <uint32_t *>(&d_node.mMinimum ) = nodeID;
116- } else if constexpr (StatsT::hasMinMax ()) {
117- stats.setStats (d_node);
193+ // 3) Publish: thread 0 now holds the node's merged bbox + stats.
194+ if (tID == 0 ) {
195+ d_node.mBBox = sBBox [0 ];
196+ if constexpr (StatsT::hasAverage ()) {
197+ // Each node writes its OWN unique d_stats slot (leaves occupy
198+ // [0, leafCount), then lower nodes, then upper nodes), so a node
199+ // with no children still has a valid slot. The previous scheme
200+ // reused a child's slot, writing d_stats[-1] for a childless
201+ // (fully-tiled) internal node - an out-of-bounds write.
202+ const uint32_t slot = d_nodeMgr->leafCount ()
203+ + (LEVEL == 2 ? d_nodeMgr->nodeCount (1 ) : 0u ) + nodeID;
204+ d_stats[slot] = sStats [0 ];
205+ *reinterpret_cast <uint32_t *>(&d_node.mMinimum ) = slot;
206+ } else if constexpr (StatsT::hasMinMax ()) {
207+ sStats [0 ].setStats (d_node);
208+ }
209+ d_node.mFlags &= ~uint64_t (1u );// enable rendering
118210 }
119- d_node.mFlags &= ~uint64_t (1u );// enable rendering
120211}// processInternal
121212
122213template <typename BuildT, typename StatsT>
@@ -199,13 +290,16 @@ void GridStats<BuildT, StatsT>::update(NanoGrid<BuildT> *d_grid, cudaStream_t st
199290
200291 StatsT *d_stats = nullptr ;
201292
202- if constexpr (StatsT::hasAverage ()) cudaCheck (util::cuda::mallocAsync ((void **)&d_stats, nodeCount[0 ]*sizeof (StatsT), stream));
293+ // One d_stats slot per node (leaves, then lower, then upper) so every node
294+ // has its own slot - see processInternal.
295+ if constexpr (StatsT::hasAverage ()) cudaCheck (util::cuda::mallocAsync ((void **)&d_stats, (nodeCount[0 ]+nodeCount[1 ]+nodeCount[2 ])*sizeof (StatsT), stream));
203296
204- processLeaf<BuildT><<<blocksPerGrid(nodeCount[0 ]), threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
297+ // warp per leaf (4 warps per 128-thread block); block per internal node
298+ if (nodeCount[0 ]) processLeaf<BuildT><<<blocksPerGrid(nodeCount[0 ]*32 ), threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
205299
206- processInternal<BuildT, StatsT, 1 ><<<blocksPerGrid( nodeCount[1 ]) , threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
300+ if (nodeCount[ 1 ]) processInternal<BuildT, StatsT, 1 ><<<nodeCount[1 ], threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
207301
208- processInternal<BuildT, StatsT, 2 ><<<blocksPerGrid( nodeCount[2 ]) , threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
302+ if (nodeCount[ 2 ]) processInternal<BuildT, StatsT, 2 ><<<nodeCount[2 ], threadsPerBlock, 0 , stream>>> (d_nodeMgr, d_stats);
209303
210304 processRootAndGrid<BuildT><<<1 , 1 , 0 , stream>>> (d_nodeMgr, d_stats);
211305
0 commit comments