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:
jgrusewski
2026-04-06 11:15:17 +02:00
parent ef62ee14a0
commit ae65740a3d

View File

@@ -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;
}
}