GPU offload for Nassau's Milnor multiply (p = 2)

[JoeyBF]

Offloads the admissible-matrix Milnor multiplication 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).

Highlights:

• Hash-free seqno index tables (flat, arc-swapped) — the on-device output-indexing primitive.
• Staged CubeCL kernels (F₂ xor → on-device seqno → single-R → batched), with u16-packed transfers, dedup-by-R + prefix-sum pair decoding, and parallel host marshalling.
• Launch-level optimizations: differential matrix computed once per bidegree, admissible-enumeration cache, device-resident admissible buffers, and per-launch memory release pinned to one CUDA stream. ~1.19× end-to-end at stem 100 on an RTX 3050 Ti; the design is transfer-bound and expected to widen on datacenter GPUs.
• Supporting infrastructure: criterion benchmark suite, an end-to-end Nassau timing example, and a compile-time-gated multiply profiler.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added optional CUDA GPU acceleration for Nassau resolution and Milnor algebra calculations.
  • Added GPU-enabled development environment support and configurable CPU/GPU verification paths.
  • Added runtime profiling and flamegraph generation options.

• Performance

  • Improved differential-matrix computation through batching, reusable results, and faster indexing.
  • Added benchmarks covering multiplication, module actions, resolution performance, and GPU comparisons.

• Bug Fixes

  • Skipped empty module input blocks to avoid unnecessary calculations.

• Tests

  • Added CPU/GPU consistency checks for matrix and resolution results.

[coderabbitai]

Warning

Review limit reached

@JoeyBF, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 28 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: b2def9b5-2ea1-4abc-b95f-7d24b96910d6

📥 Commits

Reviewing files that changed from the base of the PR and between 97c5381 and 7589013.

📒 Files selected for processing (7)

• ext/Cargo.toml
• ext/benches/milnor_gpu_ab.rs
• ext/crates/algebra/src/algebra/milnor_algebra.rs
• ext/crates/algebra/src/algebra/milnor_gpu.rs
• ext/examples/nassau_e2e.rs
• ext/src/nassau_gpu.rs
• ext/tests/nassau_gpu_timing.rs

📝 Walkthrough

Walkthrough

Adds a CUDA-backed Milnor multiplication path with seqno indexing and profiling, integrates batched products into Nassau differential construction with matrix reuse, and introduces Criterion benchmarks, CUDA development configuration, end-to-end correctness tests, and timing harnesses.

Changes

GPU Milnor multiplication and Nassau resolution

Layer / File(s)	Summary
Benchmark and profiling harness ext/Cargo.toml, ext/crates/algebra/Cargo.toml, ext/crates/algebra/benches/*, ext/benches/milnor_gpu_ab.rs, ext/examples/nassau_e2e.rs, ext/flake.nix	Adds GPU and flamegraph features, Criterion benchmarks, shared benchmark helpers, profiling controls, CUDA shell configuration, and Nassau timing tools.
CPU indexing and multiplication foundation ext/crates/algebra/src/algebra/milnor_algebra.rs, ext/crates/algebra/src/module/*, ext/crates/algebra/build.rs	Adds optional profiling, seqno-table indexing, admissible-matrix multiplication, zero-input skipping, optimized bounds handling, and CPU equivalence tests.
CUDA Milnor backend ext/crates/algebra/src/algebra/milnor_gpu.rs, ext/crates/algebra/src/algebra/mod.rs	Implements cached GPU data, CUDA kernels, host/device wrappers, single-product execution, batched multiplication, and GPU-specific correctness tests.
Nassau GPU differential integration ext/src/lib.rs, ext/src/nassau.rs, ext/src/nassau_gpu.rs	Adds feature-gated Nassau GPU support, conditional dispatch, batched product extraction, full-matrix reuse, row slicing, and CPU verification paths.
GPU end-to-end validation ext/tests/nassau_gpu*.rs	Checks GPU partial matrices against CPU results, validates reuse behavior, and provides ignored GPU timing comparisons.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Nassau as Nassau resolution
  participant Differential as build_partial_matrix
  participant GPU as nassau_gpu
  participant CUDA as Milnor GPU kernels
  participant Matrix as Differential matrix
  Nassau->>Differential: request partial differential
  Differential->>GPU: dispatch applicable products
  GPU->>CUDA: submit batched GpuProduct records
  CUDA-->>GPU: return packed F₂ output limbs
  GPU->>Matrix: decode and add output rows
  Matrix-->>Nassau: return differential matrix

Loading

Poem

I’m a bunny with benchmarks, hopping through the code,
CUDA carries Milnor terms along a faster road.
Seqno tables guide the way,
Flamegraphs bloom in bright display,
Nassau matrices match at play!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: GPU offload for Nassau's Milnor multiply at p=2.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 6

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@ext/crates/algebra/src/algebra/milnor_algebra.rs`:
- Around line 968-986: The profiling hooks in multiply_basis_element_by_element
run for all primes but are reported as p=2-only metrics. Restrict
profile::record_call, profile::identity, profile::perterm, and
profile::scope_record to the p=2 path, or separate odd-prime data into distinct
counters; update profile::report output and counter field documentation to match
the chosen behavior.
- Around line 1955-1961: Update AdmissibleMatrix::new to handle an empty ps
slice without calling unwrap on the max result, returning a valid zero-column
matrix (or otherwise applying the type’s empty-input invariant). Adjust related
row/column initialization as needed and add a regression test covering
AdmissibleMatrix::new(&[]).
- Around line 1358-1423: Prevent concurrent compute_seqno_tables calls from
replacing a larger cached table with a smaller one. Keep the existing fast-path
check, then publish the newly built SeqnoTables using an atomic CAS/RCU-style
conditional update that only stores it when the currently cached max_degree is
absent or lower; otherwise retain the existing larger table.

In `@ext/crates/algebra/src/algebra/milnor_gpu.rs`:
- Around line 139-141: Add explicit checked conversions for all narrowing from
u32/PPartEntry to u16 in the relevant matrix/term-building code, including the
col_sums, masks, and term_pparts paths around the shown self.col_sums,
self.masks, and index.insert operations. Prefer u16::try_from(...).expect(...)
consistent with seqno_table_u32, or add equivalent bounds assertions, so values
exceeding u16::MAX fail instead of silently wrapping.
- Around line 148-908: Configure CI to run the ignored GPU hardware tests
xor_f2_matches_host, seqno_matches_index_on_gpu,
multiply_single_r_matches_reference, and multiply_batch_matches_reference on a
runner with a live CUDA device. Add the required CUDA setup and invoke the
appropriate ignored cargo tests for ext/crates/algebra, ensuring this job is
part of CI validation and clearly marked as GPU-dependent.

In `@ext/tests/nassau_gpu_timing.rs`:
- Around line 23-44: Serialize the ignored GPU timing tests so they cannot
concurrently mutate NASSAU_GPU: configure the test binary or its documented
invocation to run with a single test thread, updating the module-level
documentation command accordingly. Ensure both gpu_only_resolution_time and the
other ignored test covering lines 46–76 use this serialized execution path, and
update the safety comments to reflect the enforced single-threaded harness.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 0d35a31a-47ce-4a3a-92d9-9e6a0da42cce

📥 Commits

Reviewing files that changed from the base of the PR and between a3a63cb and 97c5381.

📒 Files selected for processing (23)

• ext/Cargo.toml
• ext/benches/milnor_gpu_ab.rs
• ext/crates/algebra/Cargo.toml
• ext/crates/algebra/benches/common/mod.rs
• ext/crates/algebra/benches/milnor.rs
• ext/crates/algebra/benches/module_action.rs
• ext/crates/algebra/benches/multiplication.rs
• ext/crates/algebra/benches/nassau_milnor.rs
• ext/crates/algebra/benches/seqno.rs
• ext/crates/algebra/build.rs
• ext/crates/algebra/src/algebra/milnor_algebra.rs
• ext/crates/algebra/src/algebra/milnor_gpu.rs
• ext/crates/algebra/src/algebra/mod.rs
• ext/crates/algebra/src/module/free_module.rs
• ext/crates/algebra/src/module/homomorphism/mod.rs
• ext/examples/nassau_e2e.rs
• ext/flake.nix
• ext/src/lib.rs
• ext/src/nassau.rs
• ext/src/nassau_gpu.rs
• ext/tests/nassau_gpu.rs
• ext/tests/nassau_gpu_reuse.rs
• ext/tests/nassau_gpu_timing.rs