diff --git a/CLAUDE.md b/CLAUDE.md index 65212c28..6bec963c 100644 --- a/CLAUDE.md +++ b/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` 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 diff --git a/src/routing/mmc_op.rs b/src/routing/mmc_op.rs index 4d4b5cbd..a9e71fa6 100644 --- a/src/routing/mmc_op.rs +++ b/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 = primitive_to_vec::(gb_rhs_prim.clone()); - let x_host: Vec = primitive_to_vec::(state.x_sol.clone()); - let pattern = &state.pattern; - let grada: Vec = 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::( + &state.pattern, + gb_rhs_prim, + sparse::SavedX::Cuda(state.x_sol.clone()), + &device, + state.use_cuda, + ); // =========================================================== // B26. gc1 = assemble_backward(gA_values) diff --git a/src/sparse/mod.rs b/src/sparse/mod.rs index 5ddf4c8a..cbd393d1 100644 --- a/src/sparse/mod.rs +++ b/src/sparse/mod.rs @@ -385,9 +385,29 @@ pub(crate) fn cpu_forward_primitive( (out, x) } +thread_local! { + /// Number of `primitive_to_vec` calls made on this thread. Test aid only. + static HOST_READS: std::cell::Cell = 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`. pub(crate) fn primitive_to_vec(prim: B::FloatTensorPrimitive) -> Vec { + HOST_READS.with(|c| c.set(c.get() + 1)); let t: Tensor = Tensor::from_primitive(TensorPrimitive::Float(prim)); t.to_data().to_vec::().expect("expected f32 tensor") } diff --git a/tests/cuda_backward_parity.rs b/tests/cuda_backward_parity.rs index d1adf21c..cd77a3d5 100644 --- a/tests/cuda_backward_parity.rs +++ b/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` and `Cuda`. 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, +} + +/// 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::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::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::( + 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::() + .into_vec::() + .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(cfg: &Config, parent: Parent, weights: &[f32]) -> Vec +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let f = grad_fixture(); + let adj = confluence(&f); + let dev = device::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &dev); + let (loss, parents) = run_forward_loss::( + 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::().into_vec::().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 = (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::(); + let pat = Arc::new(CsrPattern::from_sparse(&adj)); + let asm = AValuesAssembler::::new(&pat, &dev); + let (loss, _) = run_forward_loss::( + &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::().into_vec::().unwrap() + }; + let weights: Vec = 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_cfg, p, &weights); + let gpu = confluence_grad::(&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}" + ); + } +}