From 346e6670f5ae13398d6b3f3fc5fa9bbf434ff165 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 28 May 2026 01:01:53 +0200
Subject: [PATCH] =?UTF-8?q?perf(rl):=20precompute=20tree=5Frebuild=5Flevel?=
 =?UTF-8?q?s=20=E2=80=94=20eliminate=20host=20loop=20in=20mega-graph?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The per-step while loop computing (grid, start, nodes_at_level) for
each tree level was host-side work inside the mega-graph replay path.
Now precomputed at init: 16 entries for capacity=65536, 192 bytes.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 crates/ml-alpha/src/trainer/integrated.rs | 60 ++++++++++++++---------
 1 file changed, 37 insertions(+), 23 deletions(-)

diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs
index 007e9cf15..2aa07c286 100644
--- a/crates/ml-alpha/src/trainer/integrated.rs
+++ b/crates/ml-alpha/src/trainer/integrated.rs
@@ -1002,6 +1002,10 @@ pub struct IntegratedTrainer {
     rl_per_update_priority_fn: CudaFunction,
     _rl_per_tree_rebuild_module: Arc<CudaModule>,
     rl_per_tree_rebuild_level_fn: CudaFunction,
+    /// Pre-computed per-level (grid, start, nodes_at_level) for the
+    /// bottom-up sum-tree rebuild. Avoids recomputing log2(capacity)
+    /// loop iterations on every step in the non-mega-graph path.
+    tree_rebuild_levels: Vec<(u32, i32, i32)>,
 
     // ── Bidirectional HER kernel handles ───────────────────────────
     pub hindsight: crate::rl::gpu_hindsight::GpuHindsight,
@@ -2136,6 +2140,22 @@ impl IntegratedTrainer {
         )
         .context("GpuReplayBuffer::new")?;
 
+        // Pre-compute per-level launch params for the bottom-up sum-tree
+        // rebuild. Avoids recomputing log2(capacity) iterations per step.
+        let tree_rebuild_levels = {
+            let cap = cfg.per_capacity;
+            let mut levels = Vec::new();
+            let mut nodes_at_level = (cap >> 1) as i32;
+            let mut start = nodes_at_level;
+            while nodes_at_level >= 1 {
+                let grid = ((nodes_at_level as u32) + 255) / 256;
+                levels.push((grid.max(1), start, nodes_at_level));
+                nodes_at_level >>= 1;
+                start >>= 1;
+            }
+            levels
+        };
+
         // SP20 P3 FRD head — construct + own per-step buffers. Seed
         // distinct from dqn/ppo to keep Xavier draws independent.
         let frd_head = crate::rl::frd::FrdHead::new(
@@ -2602,6 +2622,7 @@ impl IntegratedTrainer {
             rl_per_update_priority_fn,
             _rl_per_tree_rebuild_module: rl_per_tree_rebuild_module,
             rl_per_tree_rebuild_level_fn,
+            tree_rebuild_levels,
             hindsight,
             _rl_hindsight_track_module: rl_hindsight_track_module,
             rl_hindsight_track_fn,
@@ -6950,29 +6971,22 @@ impl IntegratedTrainer {
             // barrier that __threadfence() cannot (it is only a memory
             // fence, not a synchronization barrier). Graph-safe: each
             // level is a separate graph node with implicit dependency.
-            {
-                let cap = self.gpu_replay.capacity;
-                let mut nodes_at_level = (cap >> 1) as i32;
-                let mut start = nodes_at_level;
-                while nodes_at_level >= 1 {
-                    let grid = ((nodes_at_level as u32) + 255) / 256;
-                    let mut args = RawArgs::new();
-                    args.push_ptr(self.gpu_replay.priority_tree_d.raw_ptr());
-                    args.push_i32(start);
-                    args.push_i32(nodes_at_level);
-                    let mut ptrs = args.build_arg_ptrs();
-                    unsafe {
-                        raw_launch(
-                            self.rl_per_tree_rebuild_level_fn.cu_function(),
-                            (grid.max(1), 1, 1), (256, 1, 1), 0,
-                            per_stream,
-                            &mut ptrs[..args.len()],
-                        ).map_err(|e| anyhow::anyhow!(
-                            "rl_per_tree_rebuild_level (start={start}, nodes={nodes_at_level}): {:?}", e
-                        ))?;
-                    }
-                    nodes_at_level >>= 1;
-                    start >>= 1;
+            // Per-level (grid, start, nodes) pre-computed at construction.
+            for &(grid, start, nodes_at_level) in &self.tree_rebuild_levels {
+                let mut args = RawArgs::new();
+                args.push_ptr(self.gpu_replay.priority_tree_d.raw_ptr());
+                args.push_i32(start);
+                args.push_i32(nodes_at_level);
+                let mut ptrs = args.build_arg_ptrs();
+                unsafe {
+                    raw_launch(
+                        self.rl_per_tree_rebuild_level_fn.cu_function(),
+                        (grid, 1, 1), (256, 1, 1), 0,
+                        per_stream,
+                        &mut ptrs[..args.len()],
+                    ).map_err(|e| anyhow::anyhow!(
+                        "rl_per_tree_rebuild_level (start={start}, nodes={nodes_at_level}): {:?}", e
+                    ))?;
                 }
             }
         }