Keep the gA gradient on the GPU in the timestep backward

CLAUDE.md

@@ -599,6 +599,16 @@ writing an entry: `.claude/skills/ddrs-journal/SKILL.md`.
   unification was explicitly rejected.
 - **The sparse path uses a single `Arc<CsrPattern>` per network**, built once at
   `setup_inputs` and reused for every timestep. Don't rebuild per step.
+- **The timestep backward keeps gA on the device.** With `sparse_solver: cuda`,
+  `src/routing/mmc_op.rs::timestep_backward_core` computes gA with
+  `dispatch::grada_primitive` and makes no host copy. Do not put a host loop
+  back. `src/sparse/mod.rs::host_reads_on_this_thread_for_test` counts the host
+  copies (`primitive_to_vec`) made on the calling thread. Part E of
+  `tests/cuda_backward_parity.rs` uses it.
+- **On a host with CUDA 13.3, set `CUDARC_CUDA_VERSION=13020` at build time.**
+  `cudarc` 0.19.7 panics on `nvcc` 13.3. Put it in the gitignored
+  `.cargo/config.toml` `[env]` block, or export it. A fresh worktree needs that
+  file copied (trap T12 in `.claude/skills/ddrs-dev/references/traps.md`).
 - **Configs are YAML to match DDR's tooling**, deserialized via `serde_yaml`.
 - **Cite symbols, not lines, inside `src/`.** Write
   `src/config.rs::validate_subdivision_reaches_the_builder`, never

src/routing/mmc_op.rs

@@ -535,7 +535,6 @@ where
             state.use_cuda,
         );
 
-        // gA_values via direct gather+multiply on primitives (mirrors dispatch::grada_primitive).
         let gb_rhs_raw = wrap(gb_rhs_prim.clone());
 
         // Split `gb_rhs` for the MASS BOUND. Where the bound binds, the forward
@@ -558,19 +557,19 @@ where
         // the geometry-side grads to inject below. `None` ⇒ pre-leakance math.
         let zeta_geom = zeta_hook(&gb_rhs_for_zeta);
 
-        let g_a_values_prim = {
-            // -gb[row] * x[col]
-            let gradb_host: Vec<f32> = primitive_to_vec::<I>(gb_rhs_prim.clone());
-            let x_host: Vec<f32> = primitive_to_vec::<I>(state.x_sol.clone());
-            let pattern = &state.pattern;
-            let grada: Vec<f32> = pattern
-                .row_for_nnz
-                .iter()
-                .zip(pattern.col.iter())
-                .map(|(&r, &c)| -gradb_host[r as usize] * x_host[c as usize])
-                .collect();
-            I::float_from_data(burn::tensor::TensorData::from(grada.as_slice()), &device)
-        };
+        // gA[k] = -gb_rhs[row[k]] * x_sol[col[k]]. With `sparse_solver: cuda` this
+        // stays on the device (gather, multiply, negate). The CPU branch is the
+        // same host loop as before, so CPU results do not change. `x_sol` is
+        // always a backend primitive here, so it goes in as `SavedX::Cuda`; the
+        // CPU branch reads it to the host. `gb_rhs_prim` is not used after this
+        // line (`gb_rhs_raw` holds its own clone), so it moves in.
+        let g_a_values_prim = dispatch::grada_primitive::<I>(
+            &state.pattern,
+            gb_rhs_prim,
+            sparse::SavedX::Cuda(state.x_sol.clone()),
+            &device,
+            state.use_cuda,
+        );
 
         // ===========================================================
         // B26. gc1 = assemble_backward(gA_values)

src/sparse/mod.rs

@@ -385,9 +385,29 @@ pub(crate) fn cpu_forward_primitive<B: Backend>(
     (out, x)
 }
 
+thread_local! {
+    /// Number of `primitive_to_vec` calls made on this thread. Test aid only.
+    static HOST_READS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
+}
+
+/// Number of host copies (`primitive_to_vec` calls) made on this thread so far.
+///
+/// Test hook, not a stable API. A test reads it before and after a call and
+/// compares the two values. It counts only the copies in `primitive_to_vec`,
+/// the helper that the sparse solver and the timestep backward use to move a
+/// tensor to the host. It does not see other host reads, such as
+/// `Tensor::into_data` in a caller. The counter is per thread: the tests run
+/// in parallel threads, and burn's autodiff backward runs in the calling
+/// thread. The cost is one `Cell` increment per host copy. No result changes.
+#[doc(hidden)]
+pub fn host_reads_on_this_thread_for_test() -> u64 {
+    HOST_READS.with(|c| c.get())
+}
+
 /// Wrap a raw `FloatTensorPrimitive` in a `Tensor` wrapper just long enough to
 /// call `.to_data()` synchronously, then yield a `Vec<f32>`.
 pub(crate) fn primitive_to_vec<B: Backend>(prim: B::FloatTensorPrimitive) -> Vec<f32> {
+    HOST_READS.with(|c| c.set(c.get() + 1));
     let t: Tensor<B, 1> = Tensor::from_primitive(TensorPrimitive::Float(prim));
     t.to_data().to_vec::<f32>().expect("expected f32 tensor")
 }

tests/cuda_backward_parity.rs

@@ -35,6 +35,8 @@
 //! * **Part C** — non-vacuity guards, evaluated *on the CUDA backend*, so a
 //!   fixture that went degenerate only on GPU would be caught.
 //! * **Part D** — the transitive `enforce_positivity ⟹ !use_cuda_graphs` guard.
+//! * **Part E** — the `gA` gradient of the timestep backward stays on the
+//!   device (no host copy), and it is right on a confluence network.
 //!
 //! Every numeric helper below is generic in the inner backend `I` and is called
 //! with BOTH `NdArray<f32>` and `Cuda<f32, i32>`. That is deliberate: a parity
@@ -1247,3 +1249,226 @@ fn enforce_positivity_can_never_reach_the_cuda_graph_path() {
         .expect("plain ddr_match:false must load");
     assert!(!cfg.params.enforce_positivity, "enforce_positivity must default to false");
 }
+
+// ===========================================================================
+// Part E — the gA gradient of the timestep backward stays on the device.
+//
+// `timestep_backward_core` computes `gA[k] = -gb_rhs[row[k]] * x_sol[col[k]]`.
+// With `sparse_solver: cuda` this must run on the device. The old code copied
+// `gb_rhs` and `x_sol` to the host, looped over nnz and uploaded the result.
+// That is 2 device-to-host copies, 1 upload and a stream sync per timestep op.
+//
+// E1 counts the host copies in one backward. E2 checks that the device gA is
+// right on a network with confluences, where some rows have 2 off-diagonal
+// entries and `row_for_nnz` repeats. The chain fixture of Parts A to D has one
+// off-diagonal entry per row, so it cannot show a wrong gather on those rows.
+// ===========================================================================
+
+/// What one forward and one backward of a timestep op cost in host copies.
+struct BackwardTrace {
+    /// Host copies (`primitive_to_vec`) made by the forward.
+    forward_reads: u64,
+    /// Host copies made by `backward()`. This is the number E1 checks.
+    backward_reads: u64,
+    /// Gradient of `n`, to prove the backward really ran through the op.
+    grad_n: Vec<f32>,
+}
+
+/// Run one timestep forward and its backward on `Gpu` with `cfg`, and count the
+/// host copies on this thread. The counter is per thread. Burn's autodiff
+/// backward runs in the calling thread, so the count is exact even when other
+/// tests run at the same time.
+fn trace_backward_reads(cfg: &Config) -> BackwardTrace {
+    let f = grad_fixture();
+    let adj = chain(&f);
+    let dev = device::<Gpu>();
+    let pattern = Arc::new(CsrPattern::from_sparse(&adj));
+    let assembler = AValuesAssembler::<Gpu>::new(&pattern, &dev);
+    // The weights come from a CPU run. Compute them before the first read of
+    // the counter, so their host copies do not count.
+    let weights = conditioning_weights(true);
+
+    let reads_start = ddrs::sparse::host_reads_on_this_thread_for_test();
+    let (loss, parents) = run_forward_loss::<Gpu>(
+        cfg, &pattern, &assembler, &dev, &f.n, &f.qsp, &f.psp, &f.qt, &f.qpt, &f.length, &f.slope,
+        &weights, Some(Parent::N),
+    );
+    let reads_after_forward = ddrs::sparse::host_reads_on_this_thread_for_test();
+    let grads = loss.sum().backward();
+    let reads_after_backward = ddrs::sparse::host_reads_on_this_thread_for_test();
+
+    let grad_n = parents
+        .n
+        .grad(&grads)
+        .expect("grad on n")
+        .into_data()
+        .convert::<f32>()
+        .into_vec::<f32>()
+        .unwrap();
+    BackwardTrace {
+        forward_reads: reads_after_forward - reads_start,
+        backward_reads: reads_after_backward - reads_after_forward,
+        grad_n,
+    }
+}
+
+/// E1. With `sparse_solver: cuda` the timestep backward makes no host copy.
+///
+/// Falsifiability. On the code before the fix this fails: the backward made 2
+/// host copies per timestep op (`gb_rhs` and `x_sol`). The control run below
+/// uses the host solver on the same device and must show copies. If it showed
+/// none, the counter would not see the backward and this test would be vacuous.
+///
+/// Scope. The counter sees only `primitive_to_vec`. On this path that is the
+/// only host-read helper: the cuSPARSE solve, SpMV and assemble kernels of the
+/// backward read no tensor to the host.
+#[test]
+fn cusparse_backward_keeps_ga_on_the_device() {
+    // Control: the host solver copies to the host, and the counter must see it.
+    let host = trace_backward_reads(&stress_cfg(true));
+    println!(
+        "host solver on Gpu: forward reads = {}, backward reads = {}",
+        host.forward_reads, host.backward_reads
+    );
+    assert!(
+        host.backward_reads > 0,
+        "the control run (sparse_solver: cpu) shows no host copy in the backward, \
+         so the counter does not see the backward and E1 would be vacuous"
+    );
+    assert!(
+        host.grad_n.iter().any(|g| *g != 0.0),
+        "the control run made an all-zero gradient: the backward did not run through the op"
+    );
+
+    let mut cfg = stress_cfg(true);
+    cfg.params.sparse_solver = ddrs::config::SparseSolver::Cuda;
+    let cuda = trace_backward_reads(&cfg);
+    println!(
+        "cuSPARSE solver on Gpu: forward reads = {}, backward reads = {}",
+        cuda.forward_reads, cuda.backward_reads
+    );
+    assert!(
+        cuda.grad_n.iter().any(|g| *g != 0.0),
+        "the cuSPARSE run made an all-zero gradient: the backward did not run through the op"
+    );
+    assert_eq!(
+        cuda.backward_reads, 0,
+        "the timestep backward copied to the host {} time(s) with sparse_solver: cuda; \
+         gA must stay on the device",
+        cuda.backward_reads
+    );
+}
+
+/// A river network with two confluences, on the reaches of `grad_fixture()`.
+///
+/// Edges are `(downstream, upstream)`. Reach 4 takes reaches 2 and 3. Reach 7
+/// takes reaches 5 and 6. Reaches 3 and 6 are headwaters. The order is
+/// topological, so the matrix is lower triangular. Rows 4 and 7 have 2
+/// off-diagonal entries each, so `row_for_nnz` repeats 3 times on them.
+fn confluence(f: &Fixture) -> SparseAdjacency {
+    let n = f.length.len();
+    let mut dense = vec![0.0_f32; n * n];
+    for (down, up) in [(1, 0), (2, 1), (4, 2), (4, 3), (5, 4), (7, 5), (7, 6), (8, 7), (9, 8)] {
+        dense[down * n + up] = 1.0;
+    }
+    SparseAdjacency::from_dense(n, &dense, f.length.clone(), f.slope.clone())
+}
+
+/// Analytic gradient of `Σ w·q_next` on the confluence network, for backend `I`
+/// and the solver in `cfg`.
+fn confluence_grad<I: Backend + 'static>(cfg: &Config, parent: Parent, weights: &[f32]) -> Vec<f32>
+where
+    I::FloatTensorPrimitive: 'static,
+    I::Device: 'static,
+{
+    let f = grad_fixture();
+    let adj = confluence(&f);
+    let dev = device::<I>();
+    let pattern = Arc::new(CsrPattern::from_sparse(&adj));
+    let assembler = AValuesAssembler::<I>::new(&pattern, &dev);
+    let (loss, parents) = run_forward_loss::<I>(
+        cfg, &pattern, &assembler, &dev, &f.n, &f.qsp, &f.psp, &f.qt, &f.qpt, &f.length, &f.slope,
+        weights, Some(parent),
+    );
+    let grads = loss.sum().backward();
+    let g = match parent {
+        Parent::N => parents.n.grad(&grads).expect("grad on n"),
+        Parent::QSpatial => parents.qsp.grad(&grads).expect("grad on q_spatial"),
+        Parent::PSpatial => parents.psp.grad(&grads).expect("grad on p_spatial"),
+        Parent::QT => parents.qt.grad(&grads).expect("grad on q_t"),
+    };
+    g.into_data().convert::<f32>().into_vec::<f32>().unwrap()
+}
+
+/// E2. On a confluence network, the gradients with the cuSPARSE solver and the
+/// device gA match the gradients of the CPU backend with the host solver.
+///
+/// The loss weights `1/max(q_next, 1)` come from a CPU forward on the
+/// confluence network. Both sides use the same weights bit for bit, as in
+/// Part B. E2 passes on the code before the fix too: it pins the new device
+/// gather, it does not show the host copies.
+///
+/// Falsifiability (checked by hand, then reverted). In `cusparse_grada`, swap
+/// the row and col indices of the gather: E2 and
+/// `cusparse_solver_path_carries_the_same_gradients` fail. Index `gradb` by
+/// `min(row, col + 1)` instead of `row`: this equals `row` on a chain, and only
+/// E2 fails. The chain fixture alone cannot see that fault.
+#[test]
+fn cusparse_ga_matches_cpu_on_a_confluence() {
+    let f = grad_fixture();
+    let adj = confluence(&f);
+    let pattern = CsrPattern::from_sparse(&adj);
+
+    // Non-vacuity: the network must have rows with 2 off-diagonal entries,
+    // or `row_for_nnz` never repeats on an off-diagonal slot.
+    let off_diag: Vec<i32> = (0..pattern.n)
+        .map(|i| pattern.crow[i + 1] - pattern.crow[i] - 1)
+        .collect();
+    let wide_rows = off_diag.iter().filter(|&&k| k >= 2).count();
+    println!("off-diagonal entries per row: {off_diag:?}");
+    assert!(
+        wide_rows >= 2,
+        "the confluence network needs at least 2 rows with 2 off-diagonal entries, got {wide_rows}"
+    );
+
+    // Both physics configs. Without the clamp every reach carries gradient, so
+    // all 10 components are compared. With the clamp the short reaches are
+    // masked to zero, as in Part B.
+    for enforce_positivity in [false, true] {
+        // Loss weights from the CPU forward on this network.
+        let cpu_cfg = stress_cfg(enforce_positivity);
+        let ones = vec![1.0_f32; f.length.len()];
+        let q_next = {
+            let dev = device::<Cpu>();
+            let pat = Arc::new(CsrPattern::from_sparse(&adj));
+            let asm = AValuesAssembler::<Cpu>::new(&pat, &dev);
+            let (loss, _) = run_forward_loss::<Cpu>(
+                &cpu_cfg, &pat, &asm, &dev, &f.n, &f.qsp, &f.psp, &f.qt, &f.qpt, &f.length,
+                &f.slope, &ones, None,
+            );
+            loss.into_data().convert::<f32>().into_vec::<f32>().unwrap()
+        };
+        let weights: Vec<f32> = q_next.iter().map(|&q| 1.0 / q.max(1.0)).collect();
+
+        let mut gpu_cfg = stress_cfg(enforce_positivity);
+        gpu_cfg.params.sparse_solver = ddrs::config::SparseSolver::Cuda;
+
+        let mut worst = 0.0f32;
+        for p in ALL_PARENTS {
+            let cpu = confluence_grad::<Cpu>(&cpu_cfg, p, &weights);
+            let gpu = confluence_grad::<Gpu>(&gpu_cfg, p, &weights);
+            worst = worst.max(compare_backends(
+                &format!(
+                    "confluence d(loss)/d({}) [enforce_positivity={enforce_positivity}, \
+                     cuSPARSE + device gA vs CPU]",
+                    p.name()
+                ),
+                &cpu,
+                &gpu,
+            ));
+        }
+        println!(
+            "WORST confluence cross-backend rel (enforce_positivity={enforce_positivity}) = {worst:.3e}"
+        );
+    }
+}