foxhunt/crates/ml-alpha/cuda/rl_per_tree_rebuild.cu

/* =====================================================================
 * rl_per_tree_rebuild.cu — GPU-resident PER: bottom-up sum-tree rebuild
 *
 * Grid=(1), Block=(1024). Single-block deterministic rebuild.
 *
 * Rebuilds the entire sum-tree from leaf priorities (at indices
 * [capacity..2*capacity)) up to the root (index 1). All levels are
 * processed within ONE block: __syncthreads() between levels is a
 * proper barrier guaranteeing every thread in the block sees every
 * other thread's writes from the previous level.
 *
 * Determinism rationale (Phase 2 §2.F fix — Option C):
 *   The previous Grid=(128) launch used __threadfence() between levels.
 *   __threadfence() is a memory-ordering primitive, NOT a grid-wide
 *   barrier — it orders THIS thread's writes globally but does not
 *   wait for OTHER blocks' writes to be visible.  Under the previous
 *   geometry, block N could read level-L nodes while block M's
 *   level-L writes were still in flight, producing different
 *   addition orderings across same-seed runs.  Same-seed runs of the
 *   determinism diagnostic at b=128 saw root priority diverge
 *   ~592.7 vs ~695.7 at step 2 (Phase 2 sub-investigation dumps).
 *
 *   By collapsing to a single block, we lose grid-level parallelism
 *   but gain bit-deterministic execution: __syncthreads() is a hard
 *   intra-block barrier, and every internal node is the deterministic
 *   sum of its two specific children written in a definite order
 *   (sequential grid-stride within one block, fixed thread→node
 *   mapping via i = tid + k*blockDim.x).
 *
 *   Speed cost: capacity=32768 → ~65k floating-point adds total, 15
 *   levels = 15 syncs.  On RTX 3050 the kernel runs ~30-80 μs which
 *   is a ~10× slowdown vs the parallel version but well within the
 *   per-step budget (the step itself is ~30-50 ms at b=128).
 *
 * Launch: Grid=(1,1,1), Block=(1024,1,1).  Block size 1024 is the
 *   CUDA hard maximum; works on all foxhunt GPUs (sm_86 / sm_89 /
 *   sm_90).  If capacity grows past 2^21 in the future, switch to
 *   cooperative_groups::grid().sync() for a multi-block deterministic
 *   path.
 *
 * No atomicAdd — each internal node is written by exactly one thread.
 * No __threadfence() — __syncthreads() is the only inter-level barrier.
 * ===================================================================== */

extern "C" __global__ void rl_per_tree_rebuild(
    float* __restrict__ priority_tree, /* [2 * capacity] */
    int capacity
)
{
    /* Single block, single-block-stride loop per level.  blockIdx.x is
     * guaranteed 0 by the launch geometry, but guard anyway so a
     * future caller passing Grid>(1,1,1) is a silent no-op rather than
     * a corrupted tree. */
    if (blockIdx.x != 0) return;

    const int tid = threadIdx.x;
    const int block_size = blockDim.x;

    /* Bottom-up: level 0 = parents of leaves, last level = root.
     * At level L, there are (capacity >> (L+1)) internal nodes,
     * with indices [capacity >> (L+1) .. capacity >> L).            */
    int nodes_at_level = capacity >> 1; /* level 0: cap/2 nodes */
    int start = nodes_at_level;         /* first node index at this level */

    while (nodes_at_level >= 1) {
        /* Each thread processes multiple nodes via block-stride loop.
         * Mapping `i = tid + k*block_size` is FIXED across runs:
         * thread t always writes nodes (start + t), (start + t + B),
         * (start + t + 2B), ... ensuring a deterministic write order.
         */
        for (int i = tid; i < nodes_at_level; i += block_size) {
            const int node = start + i;
            priority_tree[node] = priority_tree[2 * node]
                                + priority_tree[2 * node + 1];
        }

        /* Intra-block barrier — proper synchronisation, unlike the
         * previous __threadfence() which did NOT wait for other
         * blocks' writes. */
        __syncthreads();

        /* Move up one level */
        nodes_at_level >>= 1;
        start >>= 1;
    }
}