fix: race-free segment tree propagation via atomicAdd deltas
seg_tree_update and seg_tree_insert used non-atomic tree[node] = tree[2*node] + tree[2*node+1] with 8192 threads racing on shared internal nodes. On H100 (132 SMs), all threads execute simultaneously causing data races and hangs. Replace with atomicAdd delta propagation: each thread computes delta = new_leaf - old_leaf, then atomicAdd to every ancestor. Commutative + associative = race-free. O(log n) per thread, hardware-accelerated on SM 9.0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -7,10 +7,11 @@
|
||||
// Three kernels:
|
||||
// 1. seg_tree_update — priority update: compute (|td|^alpha + eps),
|
||||
// write to priorities[], compute priority^alpha,
|
||||
// write to tree leaf, propagate sums to root.
|
||||
// write to tree leaf, propagate DELTA to root
|
||||
// via atomicAdd (race-free, O(log n) per thread).
|
||||
// 2. seg_tree_insert — insert: take raw priorities (max_priority fill),
|
||||
// compute priority^alpha, write to tree leaf,
|
||||
// propagate sums to root.
|
||||
// propagate DELTA to root via atomicAdd.
|
||||
// 3. seg_tree_sample — proportional sampling: parallel root-to-leaf
|
||||
// traversal with Philox RNG. Output i64 indices.
|
||||
|
||||
@@ -50,13 +51,18 @@ extern "C" __global__ void seg_tree_update(
|
||||
// Compute priority^alpha for tree leaf (same as pow_alpha_f32)
|
||||
float pa = powf(new_prio, alpha);
|
||||
|
||||
// Write leaf and propagate sums to root
|
||||
// Write leaf and propagate delta to root via atomicAdd.
|
||||
// Each thread computes delta = new_leaf - old_leaf and adds it to every
|
||||
// ancestor. atomicAdd is commutative+associative, so concurrent threads
|
||||
// produce correct sums without synchronization barriers.
|
||||
int leaf = capacity + (int)idx;
|
||||
float old_pa = tree[leaf];
|
||||
tree[leaf] = pa;
|
||||
float delta = pa - old_pa;
|
||||
|
||||
int node = leaf >> 1;
|
||||
while (node >= 1) {
|
||||
tree[node] = tree[2 * node] + tree[2 * node + 1];
|
||||
atomicAdd(&tree[node], delta);
|
||||
node >>= 1;
|
||||
}
|
||||
}
|
||||
@@ -82,13 +88,15 @@ extern "C" __global__ void seg_tree_insert(
|
||||
// Compute priority^alpha for tree leaf
|
||||
float pa = powf(priorities[i], alpha);
|
||||
|
||||
// Write leaf and propagate sums to root
|
||||
// Write leaf and propagate delta to root via atomicAdd (race-free).
|
||||
int leaf = capacity + (int)idx;
|
||||
float old_pa = tree[leaf];
|
||||
tree[leaf] = pa;
|
||||
float delta = pa - old_pa;
|
||||
|
||||
int node = leaf >> 1;
|
||||
while (node >= 1) {
|
||||
tree[node] = tree[2 * node] + tree[2 * node + 1];
|
||||
atomicAdd(&tree[node], delta);
|
||||
node >>= 1;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user