GPU offload for Nassau's Milnor multiply (p = 2)#271
Conversation
Add a flat, arc-swapped `seqno` index for the Milnor basis: an O(#p-part-entries) rank computed from a precomputed `g` table, with no hash lookup. This is a general (GPU-agnostic) data structure — it is the uploadable/on-device index primitive, and an independently benchmarkable alternative to the basis hashmap. - `SeqnoTables` + `compute_seqno_tables` (idempotent, concurrency-safe via `ArcSwapOption::rcu`) + `seqno` in `milnor_algebra.rs`, applicable at p = 2 with a trivial profile and stable ordering. - The CPU basis index deliberately still uses the hashmap (the tables lose to it on the CPU); `compute_basis` does not build them. Documented at `try_basis_element_to_index`. - `benches/seqno.rs`: A/B of the table index against the hashmap. - Covered by the `seqno_matches_enumeration_order` test. Split out of the Nassau GPU work as standalone infrastructure. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPYvsLdEfitgCbAiPxxx3U
|
Warning Review limit reached
Next review available in: 48 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (16)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
- `seqno`: add a `debug_assert` that the element's degree is within the loaded table's `max_degree`, so a not-built-far-enough call fails with a clear diagnostic instead of a raw slice out-of-bounds panic. - `seqno_matches_enumeration_order`: build the tables partially, rebuild identically (no-op), grow, then request a smaller degree — exercising the idempotent, monotonic (non-shrinking) publish documented on `compute_seqno_tables`, not just a single build. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPYvsLdEfitgCbAiPxxx3U
Offload the admissible-matrix Milnor multiply at the heart of Nassau's `get_partial_matrix` to the GPU via CubeCL, for the p = 2 / trivial-profile regime. Gated behind the `gpu` feature; `cargo check`/`build` need no CUDA toolkit (cudarc dlopens at runtime). - `milnor_algebra.rs`: an admissible-matrix multiply (`multiply_basis_element_by_element_2` + `AdmissibleMatrix`) as the CPU reference model for the kernel, plus the `#[cfg(feature = "gpu")]` upload helpers (`admissible_matrices`, `seqno_table_u32`, `gpu_multiply_applicable`). Covered by `admissible_multiply_agrees_with_reference`. - `milnor_gpu.rs`: the batched CubeCL kernel (`multiply_batch_kernel` / `multiply_batch_on_gpu`), fusing every `(R, s)` product of one matrix build into a single launch, with dedup-by-`R`, device-resident admissible buffers, and per-launch memory release pinned to one CUDA stream. The single-kernel building blocks used only to validate it live under `#[cfg(test)]`, so the production surface is just `multiply_batch_on_gpu` + `GpuProduct`. - `ext::nassau_gpu` + the `build_partial_matrix` hook in `nassau.rs`: route the differential-matrix build through the GPU path when compiled in, opted into (`NASSAU_GPU`), applicable, and large enough to amortise launch overhead; otherwise the CPU path is used unchanged. - `free_module.rs`: skip zero generator blocks in `act` (a general CPU win). - Supporting `nassau_e2e` example, GPU/CPU consistency tests, the `milnor_gpu_ab` A/B bench, and a `gpu` dev shell in `flake.nix`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPYvsLdEfitgCbAiPxxx3U
7824e7d to
31aced0
Compare
…Sequences#272's restricted graph PR SpectralSequences#272's restricted_partial_matrix (degree-bounded columns) had displaced PR SpectralSequences#271's GPU partial-matrix offload and cross-signature row-reuse. Add a restricted GPU variant (get_partial_matrix_restricted[_verified]) that keeps the batched Milnor multiply but sizes output to the frozen prefix and masks bits beyond it, dispatch to it from restricted_partial_matrix_maybe_gpu, and restore row-reuse (one full restricted matrix per bidegree, select_rows per signature). Verified GPU==CPU on H200 NVL (CUDA 12.4): nassau_gpu (14230 rows / 527 bidegrees) and milnor_vs_nassau (9 tests) pass with NASSAU_GPU_VERIFY=1. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2 participants
Reorganized successor to #264: the GPU Milnor-multiply offload, with the benchmark suite (#269, merged) and the
seqnoindex (#270) split out.What
Offloads the admissible-matrix Milnor multiply at the heart of Nassau's
get_partial_matrixto the GPU via CubeCL, for thep = 2/ trivial-profile regime. Gated behind thegpufeature;cargo check/buildneed no CUDA toolkit (cudarc dlopens at runtime).milnor_algebra.rs: an admissible-matrix multiply (multiply_basis_element_by_element_2+AdmissibleMatrix) as the CPU reference model for the kernel, plus the#[cfg(feature = "gpu")]upload helpers. Covered byadmissible_multiply_agrees_with_reference.milnor_gpu.rs: the batched CubeCL kernel (multiply_batch_kernel/multiply_batch_on_gpu), fusing every(R, s)product of one matrix build into a single launch, with dedup-by-R, device-resident admissible buffers, and per-launch memory release pinned to one CUDA stream. Single-kernel building blocks used only to validate it live under#[cfg(test)], so the production surface is justmultiply_batch_on_gpu+GpuProduct.ext::nassau_gpu+ thebuild_partial_matrixhook innassau.rs: route the differential-matrix build through the GPU path when compiled in, opted into (NASSAU_GPU), applicable, and large enough to amortise launch overhead; otherwise the CPU path is unchanged.free_module.rs: skip zero generator blocks inact(a general CPU win).nassau_e2eexample, GPU/CPU consistency tests, themilnor_gpu_abA/B bench, and agpudev shell inflake.nix.Notes vs #264
Dead code and dev-only profiling scaffolding removed; the validation-only kernels moved behind
#[cfg(test)]. The env-var dispatch and the process-global GPU context are unchanged (deferred as larger follow-ups).Test plan
cargo test -p algebra(incl.admissible_multiply_agrees_with_reference)cargo test -p ext --features nassau --test milnor_vs_nassau(CPU Nassau correctness)just lint(nightly fmt + clippy,-D warnings) — incl.--all-features(compiles thegpupath via cubecl)tests/nassau_gpu*.rs) require a live CUDA device — not run in CI🤖 Generated with Claude Code
Generated by Claude Code