From d14fbc103fad7074a52ee85bbd379bd51493bc34 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Sun, 18 Oct 2026 10:34:45 -0400 Subject: [PATCH 1/5] Add a host-copy counter and Part E tests for the gA backward on CUDA Add a thread-local counter in src/sparse/mod.rs. `primitive_to_vec` increments it. `host_reads_on_this_thread()` reads it. It costs one `Cell` increment per host copy and changes no result. Add Part E to tests/cuda_backward_parity.rs: - E1 `cusparse_backward_keeps_ga_on_the_device` counts the host copies in one timestep backward with `sparse_solver: cuda`. It expects 0. A control run with the host solver must show copies, so the counter is not blind. E1 fails on the current code: 2 copies per timestep op. - E2 `cusparse_ga_matches_cpu_on_a_confluence` compares the gradients of CUDA + cuSPARSE with the CPU backend on a network with two confluences. Some rows there have 2 off-diagonal entries, so `row_for_nnz` repeats. E2 passes before and after the fix. Co-Authored-By: Claude Opus 5.5 --- src/sparse/mod.rs | 19 +++ tests/cuda_backward_parity.rs | 215 ++++++++++++++++++++++++++++++++++ 2 files changed, 234 insertions(+) diff --git a/src/sparse/mod.rs b/src/sparse/mod.rs index 5ddf4c8a..40c65e60 100644 --- a/src/sparse/mod.rs +++ b/src/sparse/mod.rs @@ -385,9 +385,28 @@ 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 aid. 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. 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() -> 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..307062d6 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,216 @@ 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(); + 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(); + let grads = loss.sum().backward(); + let reads_after_backward = ddrs::sparse::host_reads_on_this_thread(); + + 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. +#[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. +#[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}" + ); + } +} From bcc419e3bd8c928b840aa82901dbef943243356c Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Sun, 18 Oct 2026 10:36:19 -0400 Subject: [PATCH 2/5] Keep the gA gradient on the GPU in the timestep backward `timestep_backward_core` copied `gb_rhs` and `x_sol` to the host, looped over nnz to build gA, and uploaded the result. With `sparse_solver: cuda` that was 2 device-to-host copies, 1 upload and a stream sync per timestep op, in every backward (all timestep op variants share this function). Call `dispatch::grada_primitive` instead. `CsrSolveOp::backward` already uses it. On CUDA it gathers, multiplies and negates on the device. On the CPU it runs the same host loop as before, in the same order, so CPU results do not change. E1 (`cusparse_backward_keeps_ga_on_the_device`) now passes: 0 host copies in the backward. It made 2 before this change. Co-Authored-By: Claude Opus 5.5 --- src/routing/mmc_op.rs | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) 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) From 486262be9bc48eb5a7548e929ef837b125fea9a0 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Sun, 18 Oct 2026 12:13:20 -0400 Subject: [PATCH 3/5] Note the hand mutation checks in the E2 test comment Swapping the row and col indices in `cusparse_grada` fails E2 and `cusparse_solver_path_carries_the_same_gradients`. Indexing `gradb` by `min(row, col + 1)` equals `row` on a chain, and only E2 fails. Both edits were reverted. No code change. Co-Authored-By: Claude Opus 5.5 --- tests/cuda_backward_parity.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/cuda_backward_parity.rs b/tests/cuda_backward_parity.rs index 307062d6..686b3a1a 100644 --- a/tests/cuda_backward_parity.rs +++ b/tests/cuda_backward_parity.rs @@ -1403,6 +1403,12 @@ where /// 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(); From ee9e0039ade3706f8108c367674d73ef88119008 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Sun, 18 Oct 2026 12:21:47 -0400 Subject: [PATCH 4/5] Name the host-copy counter as a test hook and state its scope Rename `host_reads_on_this_thread` to `host_reads_on_this_thread_for_test`, as the other `#[doc(hidden)]` test hooks in ddrs (`content_key_for_test`, `run_smoke_for_test`). Its doc now says that it sees only `primitive_to_vec`, not other host reads such as `Tensor::into_data`. The E1 doc now says why that is enough: on the timestep backward path, `primitive_to_vec` is the only host-read helper. The cuSPARSE solve, SpMV and assemble kernels read no tensor to the host. No code behaviour changes. Co-Authored-By: Claude Opus 5.5 --- src/sparse/mod.rs | 15 ++++++++------- tests/cuda_backward_parity.rs | 10 +++++++--- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/sparse/mod.rs b/src/sparse/mod.rs index 40c65e60..cbd393d1 100644 --- a/src/sparse/mod.rs +++ b/src/sparse/mod.rs @@ -392,14 +392,15 @@ thread_local! { /// Number of host copies (`primitive_to_vec` calls) made on this thread so far. /// -/// Test aid. 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. 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. +/// 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() -> u64 { +pub fn host_reads_on_this_thread_for_test() -> u64 { HOST_READS.with(|c| c.get()) } diff --git a/tests/cuda_backward_parity.rs b/tests/cuda_backward_parity.rs index 686b3a1a..cd77a3d5 100644 --- a/tests/cuda_backward_parity.rs +++ b/tests/cuda_backward_parity.rs @@ -1288,14 +1288,14 @@ fn trace_backward_reads(cfg: &Config) -> BackwardTrace { // the counter, so their host copies do not count. let weights = conditioning_weights(true); - let reads_start = ddrs::sparse::host_reads_on_this_thread(); + 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(); + 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(); + let reads_after_backward = ddrs::sparse::host_reads_on_this_thread_for_test(); let grad_n = parents .n @@ -1318,6 +1318,10 @@ fn trace_backward_reads(cfg: &Config) -> BackwardTrace { /// 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. From 31672343223bfeefc14dbbb0c9b43043bf78d5e6 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Sun, 18 Oct 2026 13:41:06 -0400 Subject: [PATCH 5/5] CLAUDE.md: note the on-device gA backward and the CUDARC_CUDA_VERSION build need The timestep backward keeps gA on the device with sparse_solver: cuda. host_reads_on_this_thread_for_test counts the host copies (tests/cuda_backward_parity.rs Part E). On a host with CUDA 13.3, cudarc 0.19.7 needs CUDARC_CUDA_VERSION=13020 at build time. Co-Authored-By: Claude Sonnet 5.5 --- CLAUDE.md | 10 ++++++++++ 1 file changed, 10 insertions(+) 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