Perf/metal msm optimizations - #100
Merged
Merged
Conversation
…nges The build script gated test-kernel compilation on this env var but never declared it to cargo, so a cached build-script run could leave the metallib without test kernels.
Previously every stage committed its own command buffer, blocked on wait_until_completed, copied results back to host Vecs and re-uploaded them as fresh buffers for the next stage, and every MSM call rebuilt all pipeline states from the metallib. Now all dispatches are encoded up front against device-resident buffers, committed back-to-back (one command buffer per dispatch so the GPU watchdog never sees a long-running one) with a single host sync at the end; only the final per-subtask points are read back. Zero-filled buffers come from newBufferWithLength instead of host-side staging vectors, and the pipeline object is cached per thread across calls. The old per-stage structs are inlined into execute_pipeline.
The 16-bit x 16 limb representation was inherited from WGSL, which has no 64-bit integers. Metal does: a widening 32x32->64 multiply-add is a single fused operation on Apple GPUs, so 8x32-bit limbs quarter the partial products in Montgomery multiplication (256 -> 64) and remove all mask-and-shift carry handling. Field elements shrink 64 -> 32 bytes and Jacobian points 192 -> 96, halving register pressure in the ILP-1 bucket accumulation chain. - CIOS Montgomery multiplication rewritten with ulong accumulators - coords/scalars buffers are now plain little-endian words; the convert kernel converts to Montgomery form via mont_mul(x, R^2), replacing the Barrett reduction path entirely (deleted along with the 12/15-bit Montgomery variants that only work for narrow limbs) - host: 32-bit-safe to_limbs/from_limbs fast paths; fixed u32 shift overflow in calc_rinv_and_n0 and calc_nsafe for log_limb_size = 32 - constants.metal regenerated (N0 = -p^-1 mod 2^32, R^2, 8-limb tables) - tests swept to the 8/32 config; obsolete Barrett and narrow-limb Montgomery tests removed Measured: 2^16 e2e 180 -> 106 ms, 2^20 1351 -> 949 ms.
SMVP input points are affine (z = 1), but the inner loop paid the full Jacobian-Jacobian addition (11M + 5S) by manufacturing a z = 1 point. jacobian_madd_safe implements madd-2007-bl (7M + 4S, ~31% fewer field muls) with the edge cases the textbook formula silently breaks on: accumulator at infinity, P + P (falls back to doubling) and P + (-P) (returns infinity). The u2/s2 values the guards compare are needed by the formula anyway, and mont_mul outputs are canonical, so the checks are limb-wise equality. Bucket-merge additions (both operands Jacobian) keep the general path. Measured: 2^20 e2e 949 -> 793 ms.
The transpose ran one thread per subtask (~20 threads on the whole GPU), each doing three serial passes over up to 2^20 elements. It is a counting sort, so it splits into three kernels: - histogram: one thread per (element, subtask), relaxed atomic adds on per-column counters - prefix sum: still one thread per subtask, but over 2^w columns instead of n elements - scatter: one thread per (element, subtask); each thread claims a unique slot via an atomic per-column counter Within-column order becomes nondeterministic, which is sound because the only consumer sums each column and point addition commutes. The legacy serial kernel is kept for its unit test. Measured: 2^20 e2e 793 -> 372 ms, overtaking Arkworks CPU (~700 ms).
Implements mont_mul_logjumps (Domb/Koh: replace each Montgomery round's quotient computation with a shift plus t0 * rho, rho = 2^-32 mod p), saving n-1 of n^2+n limb products and one serial dependency per round. Not enabled: for BN254, rho ~ 0.89p, so the exact worst-case recurrence reaches 2.89p before reduction and needs a second conditional subtraction, which together with the in-loop round branch costs more than the saved multiplications. Measured 8-10% slower end-to-end than CIOS on Apple GPUs. Kept for reference and for primes/limb widths where rho is small; call sites now go through a mont_mul alias so the implementation can be swapped in one place.
arkworks stores Fq internally in Montgomery form with R = 2^256 -- the same radix the GPU uses. The old path converted out of Montgomery form on the CPU (one REDC per coordinate via into_bigint) only for the GPU to convert straight back (mont_mul by R^2 per coordinate). Coordinates now flow raw from pt.x.0: no host reduction, no GPU conversion, and SMVP reads the input buffer directly as interleaved affine points (one contiguous 64-byte record per gather instead of two cache lines in two buffers). The convert kernel shrinks to scalar decomposition only (decompose_scalars.metal); scalars still need into_bigint because window decomposition slices the standard bit pattern. The separate point_x/point_y buffers (2 x 32 MB at 2^20) are gone. Peak GPU allocation at 2^20 drops 410 -> 263 MB combined with the earlier removal of over-allocations.
- window 13 for all inputs below 2^19 (was 8 below 2^14): at small n the binding constraint is dispatch overhead and GPU occupancy, not the Pippenger add count; measured 2^14 56 -> 37 ms, 2^12 26 -> 23 ms - halve workgroup scale factors: more, smaller threadgroups pack better around variable-length buckets; 2^16 66 -> 60 ms, 2^20 -4% - MSM_WINDOW_SIZE / MSM_SCALE_FACTOR env overrides for tuning runs Note from the sweep: window sizes 12 and 14 are pathological (254 mod w leaves a 2-bit top window, cramming all points into ~4 buckets whose serial accumulation stalls the dispatch ~17x); 13/15/16 are safe. Window 8 also has a known size-dependent correctness bug at n = 1024 (pre-existing, both old and new code; defaults never select it).
It measured 8-10% slower than CIOS end-to-end (BN254's rho = 2^-32 mod p is ~0.89p, forcing a second conditional subtraction that cancels the saved multiplications) and nothing calls it. The mont_mul alias stays as the single swap point and documents the result; the implementation and the BN254_RHO constant live in the previous commit if rho-friendly parameters ever make it worth revisiting.
…ory tracking - examples/quick_bench: min/median timings per size, correctness asserted against Arkworks, and peak GPU memory via a background thread sampling MTLDevice currentAllocatedSize (relevant for mobile targets; VRAM is now reported on every run) - MSM_PROFILE=1 runs each pipeline dispatch synchronously with wall times and prints per-pipeline occupancy info - benches/shaders: replace benches of deleted kernels (convert/Barrett) with a decompose_scalars bench; port SMVP bench to the interleaved point layout and the 8x32 limb config - examples/debug_w8: stage-by-stage GPU-vs-CPU validation harness used to isolate the pre-existing w=8/n=1024 correctness bug (PBPR stage) - README: optimization notes with per-change measurements, LogJumps and Elastic MSM findings, Metal 4 status, updated future work
The 4-subtask chunking was inherited from the WebGPU implementation (which has per-dispatch workgroup limits Metal doesn't have). Separate dispatches serialize on the queue, so each chunk's slowest bucket stalled the GPU before the next chunk could start; a single dispatch lets the scheduler overlap the load-imbalance tail across all subtasks. Measured: 2^16 60 -> 53 ms, 2^20 358 -> 339 ms, 2^22 1056 -> 977 ms.
XYZZ (X, Y, ZZ = z^2, ZZZ = z^3) trades one extra stored coordinate for cheaper group operations everywhere the pipeline is hot: - mixed addition (SMVP inner loop): 8M+2S with ~5 field add/subs, vs Jacobian madd's 7M+4S with ~13 add/subs - general addition (bucket merge, PBPR running sums): 12M+2S vs 11M+5S with roughly half the add/subs The identity is ZZ = 0, so zero-initialized bucket buffers still encode infinity for never-written slots. The host converts the final partial sums to Jacobian without inversions via the scale-equivalence (X*ZZ*ZZZ^2, Y*ZZ^3*ZZZ^2, ZZ*ZZZ). Jacobian code stays for its unit tests and the scalar-mul kernel. Measured: 2^16 53 -> 47 ms, 2^20 339 -> 305 ms, 2^22 977 -> 945 ms; peak GPU memory +3% (fourth bucket coordinate). Cumulative vs v0.2.0 on this machine: 2^12 2.9x, 2^16 4.4x, 2^20 5.0x.
Raising it to 256 cut bpr_stage_1 from 9 to 2 ms at 2^20 but added ~40 ms of host-side final reduction (8x more partial sums to convert and add on the CPU) -- a net loss at every size. Keep it tied to the scale factor; moving the final reduction onto the GPU is the prerequisite for revisiting this.
…groups Follow-up to the documented b_workgroup_size trade-off: the host used to convert and add num_subtasks * b_workgroup_size partial sums, which made larger PBPR workgroups a net loss (cheap GPU time traded for expensive CPU time). bpr_reduce collapses each subtask's partial sums to a single point on the GPU: one threadgroup per subtask, grid-stride accumulate into 128 threads, then a barrier tree reduction in threadgroup memory (16 KB of XYZZ points). The host now reads back and Horner-combines only num_subtasks points, so b_workgroup_size rises to 256 for free (bpr_stage_1 drops from ~9 to ~2 ms at 2^20). Measured: 2^16 47.5 -> 39.1 ms, 2^20 300 -> 288 ms, 2^22 926 -> 869 ms.
One-thread-per-bucket SMVP finishes when the fattest bucket does; bucket sizes are Poisson-distributed, so SIMD groups idle on the tail, and window sizes whose top digit is narrow (12/14) were catastrophically imbalanced (~17x). smvp_accumulate assigns threads fixed-size slices of the bucket-sorted point stream instead -- uniform work by construction. Because the stream is sorted, a run that starts and ends inside a slice covers its bucket completely and is written exclusively (no atomics); only a slice's first and last runs can be partial, so each thread spills at most two entries. smvp_merge combines spills (a row's spills occupy the contiguous slice range [row_ptr[r]/L, (row_ptr[r+1]-1)/L] -- no searching) and folds mirror rows into signed buckets as before. Slice length L = 128 (MSM_SMVP_L to override); the classic kernel remains available via MSM_LEGACY_SMVP=1 and for its unit test. Measured (medians): 2^12 18.9 -> 13.7 ms, 2^16 39.1 -> 21.8 ms, 2^20 288 -> 175-188 ms, 2^22 869 -> 652 ms. Pathological windows fixed: w=12 at 2^20 6365 -> 238 ms; all window sizes are now load-safe. Peak GPU memory +~100-200 MB at the top sizes (row-sum and spill buffers) -- reducible later by storing row sums in affine form.
Batched-affine bucket accumulation was worked through in detail and rejected on paper for this architecture: with XYZZ mixed addition at 10 mult-ops, batched affine needs its ~380-mul batch inversion amortized over k >= 128 additions to win, and every way of assembling such a batch on Apple GPUs (threadgroup Blelloch scans + ~30 barriers/round, 32-wide SIMD batches, or Yrrid-style per-thread multi-bucket accumulators) costs more than the ~3 mult-ops it would save. CUDA MSMs profit from it via warp-synchronous k = 1024 batches against a 16-mul Jacobian baseline -- neither condition holds here.
The segmented SMVP's row-sum and spill buffers scaled with num_subtasks (+100-200 MB at the top sizes). Subtasks are now processed in chunks of 4 (when num_columns >= 32768) sharing one chunk-sized set of buffers; Metal's hazard tracking serializes the reuse, and stale data is safe because row sums are only read for rows direct-written in the same chunk and every spill slot is rewritten (value or sentinel) each chunk. Note: the earlier idea of affine row storage died with batched-affine -- converting XYZZ to affine at flush time needs the same inversion that made batched affine uneconomic. Chunked reuse saves more anyway. Measured: peak GPU memory 2^20 375 -> 296 MB, 2^22 1176 -> 1017 MB; timings unchanged within noise (2^20 ~186 ms, 2^22 ~650 ms).
Session tooling, not library surface: the Criterion benches (e2e, shaders) remain the supported benchmark path. The debug_w8 harness's finding (pre-existing w=8/n=1024 PBPR bug) stays documented in the tuning commit; both files remain recoverable from git history.
The XYZZ migration added a fourth bucket/partial-sum coordinate to the smvp and pbpr kernel signatures, but benches/shaders.rs still bound the old 3-coordinate buffer sets; the unbound slot made the smvp bench case hang the GPU. Buckets now bind x/y/zz/zzz (pbpr inputs seed ZZ = ZZZ = 1 in Montgomery form for z = 1 points), matching the kernels.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe PR migrates Metal MSM to eight 32-bit limbs and XYZZ coordinates. It adds GPU scalar decomposition, parallel transpose, segmented SMVP, GPU bucket reduction, pipelined execution, updated shader constants, and corresponding benchmarks and tests. ChangesMetal MSM pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant MetalMSM
participant ShaderManager
participant MetalHelper
participant MetalKernels
MetalMSM->>ShaderManager: Load MSM shader pipelines
MetalMSM->>MetalHelper: Allocate GPU-resident buffers
MetalMSM->>MetalHelper: Encode and commit shader operations
MetalHelper->>MetalKernels: Run decomposition, transpose, SMVP, and reduction
MetalKernels-->>MetalMSM: Return per-subtask XYZZ results
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
mopro-msm/src/msm/metal_msm/metal_msm.rs (1)
567-596: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winValidate the
MSM_WINDOW_SIZEandMSM_SCALE_FACTORoverrides.Both overrides accept any value that parses as the target integer type. Out-of-range values fail later in confusing ways:
MSM_WINDOW_SIZE=0makesnum_columns1 andhalf_columns0, sob_workgroup_sizebecomes 0 and the BPR dispatch requests a zero-size threadgroup.MSM_WINDOW_SIZE=64or larger overflows1 << window_sizeand panics.MSM_SCALE_FACTOR=0makesc_workgroup_size0, soinput_size / c_workgroup_sizeat Line 146 divides by zero.Clamp both values to the supported range before use.
🛡️ Proposed fix
let window_size = if let Ok(w) = std::env::var("MSM_WINDOW_SIZE") { - w.parse().unwrap_or(13) + w.parse::<usize>().ok().filter(|&w| (1..=20).contains(&w)).unwrap_or(13) } else if input_size < 524288 {let scale_factor = if let Ok(s) = std::env::var("MSM_SCALE_FACTOR") { - s.parse().unwrap_or(1) + s.parse::<usize>().ok().filter(|&s| s > 0).unwrap_or(1) } else if input_size <= 65536 {🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/metal_msm.rs` around lines 567 - 596, Validate the parsed overrides in the window-size and scale-factor selection logic: clamp MSM_WINDOW_SIZE to a supported positive range below the shift-overflow limit, and clamp MSM_SCALE_FACTOR to a supported positive range that keeps c_workgroup_size nonzero. Preserve the existing size-based defaults when the environment variables are absent or invalid.mopro-msm/src/msm/metal_msm/utils/limbs_conversion.rs (1)
556-580: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winExtend the parallel-packing test to cover the 32-bit zero-copy path.
test_parallel_packing_correctnessonly configureslog_limb_size: 16, num_limbs: 16. The newlog_limb_size == 32branch inpack_affine_and_scalars(Lines 390-439) takes a materially different code path: it copies arkworks' internal Montgomery-form limbs directly (pt.x.0.0), instead of callinginto_bigint()to get the standard-form value used bypack_affine_and_scalars_sequential. No test currently checks that this zero-copy output matches the correct Montgomery-domain reference values.Add a dedicated test (or a separate reference implementation using
into_bigint()plus manual Montgomery conversion viacalc_mont_radix) that exerciseslog_limb_size: 32, num_limbs: 8and validates the packed coordinates and scalars.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/utils/limbs_conversion.rs` around lines 556 - 580, Extend test_parallel_packing_correctness with a dedicated log_limb_size: 32, num_limbs: 8 case that exercises the zero-copy branch in pack_affine_and_scalars. Compare its packed coordinates and scalars against a correct Montgomery-domain reference built from into_bigint() and calc_mont_radix, preserving the existing 16-bit assertions.
🧹 Nitpick comments (14)
mopro-msm/benches/shaders.rs (2)
64-69: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe transpose benchmark still declares 16 limbs of 16 bits.
Every other benchmark block in this file moved to
log_limb_size: 32andnum_limbs: 8. The transpose kernel does not use field limbs, so behavior does not change, but the stale values now contradict the rest of the file. Update them for consistency.♻️ Proposed change
let config = MetalConfig { - log_limb_size: 16, - num_limbs: 16, + log_limb_size: 32, + num_limbs: 8, shader_file: "cuzk/transpose.metal".to_string(),🤖 Prompt for 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. In `@mopro-msm/benches/shaders.rs` around lines 64 - 69, Update the MetalConfig for the transpose benchmark, identified by kernel_name "transpose", to use log_limb_size 32 and num_limbs 8, matching the other benchmark configurations while leaving the shader_file and kernel name unchanged.
402-411: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
smvp_gpuallocates and bindsbucket_zzz_bufbut never returns it.This helper creates four bucket buffers and passes all four to the kernel, then returns only X, Y, and ZZ. The equivalent helper in
mopro-msm/src/msm/metal_msm/tests/cuzk/smvp.rslines 14-112 is now identical except that it returns all four coordinates. Two near-duplicate helpers with different return arities will drift further.For a benchmark, the readback is only there to force completion, so returning four values costs nothing and removes the divergence.
Also applies to: 441-445, 481-485, 496-498
🤖 Prompt for 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. In `@mopro-msm/benches/shaders.rs` around lines 402 - 411, Update smvp_gpu to return all four bucket coordinate buffers, including bucket_zzz_buf, matching the equivalent cuzk smvp helper. Adjust the function’s tuple return type and every return/destructuring site in the benchmark, while preserving the existing readback used to force GPU completion.mopro-msm/src/msm/metal_msm/tests/cuzk/convert_point_coords_and_decompose_scalars.rs (1)
140-143: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
pack_limbsis now an identity function, and the file name no longer matches its contents.With 32-bit limbs the helper just clones the slice. Two call sites (Line 37 and Line 71) would read the same with
scalarsused directly. The file is also still namedconvert_point_coords_and_decompose_scalars.rs, but the point-coordinate conversion test was removed, so only scalar decomposition remains. Rename the file todecompose_scalars.rsand remove the identity helper.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/tests/cuzk/convert_point_coords_and_decompose_scalars.rs` around lines 140 - 143, Remove the identity helper pack_limbs and update both call sites to use the scalar limb slices directly. Rename the test file from convert_point_coords_and_decompose_scalars.rs to decompose_scalars.rs, preserving the existing scalar decomposition behavior.mopro-msm/src/msm/metal_msm/metal_msm.rs (2)
239-262: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider removing the collapsed chunk loop in the legacy SMVP path.
num_subtask_chunk_sizeequalsnum_subtasks, so thestep_byloop always runs exactly one iteration,offsetis always 0, andremaining_subtasksalways equalsnum_subtasks. The two vectorssmvp_params_bufsandsmvp_dimshold one element each. Flattening this to a single dispatch would remove the unused chunking arithmetic and make the "single dispatch" comment self-evident.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/metal_msm.rs` around lines 239 - 262, Remove the collapsed chunk loop around the legacy SMVP setup and construct a single params buffer and dimension entry directly. Replace the derived offset, remaining_subtasks, and chunking arithmetic with the equivalent num_subtasks-based values while preserving the existing smvp_params_bufs and smvp_dims contents and single-dispatch behavior.
540-543: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the thread-local MSM pipeline cache with a shared cache.
metal_variable_base_msmstores oneMetalMSMPipelineper Rust thread inPIPELINE, and the cache is never released. Because each pipeline has a device and compiled shader pipeline states, calls from many short-lived threads can grow process memory with the number of distinct threads. Use a process-wide shared cache instead, ifMetalMSMPipelineisSend + Syncand concurrentexecute_pipelinecalls are safe.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/metal_msm.rs` around lines 540 - 543, Replace the thread-local PIPELINE cache used by metal_variable_base_msm with a process-wide shared OnceCell, verifying that MetalMSMPipeline is Send + Sync and concurrent execute_pipeline calls are safe; otherwise synchronize access appropriately. Preserve lazy initialization and reuse the single pipeline across MSM calls.mopro-msm/src/msm/metal_msm/shader/bigint/bigint.metal (1)
14-22: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winPin the 32-bit limb assumption with a compile-time check.
The new add and subtract paths use
(uint)truncation in place of& MASK. That is exact only whenLOG_LIMB_SIZE == 32.write_constantsstill accepts anylog_limb_size, so a future regeneration with a smaller limb size would produce unmasked limbs and silently wrong field arithmetic.Add a static assertion so the mismatch fails at compile time.
♻️ Proposed addition near the top of the file
`#pragma` once using namespace metal; `#include` "../misc/get_constant.metal" + +// The carry/borrow paths below rely on (uint) truncation acting as the limb mask. +static_assert(LOG_LIMB_SIZE == 32, "bigint.metal requires 32-bit limbs");Also applies to: 49-57
🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/bigint/bigint.metal` around lines 14 - 22, Add a compile-time assertion near the top of bigint.metal that requires LOG_LIMB_SIZE to equal 32, covering both the addition and subtraction paths that rely on uint truncation. Keep the existing limb arithmetic unchanged and ensure generation with any other limb size fails during compilation.mopro-msm/src/msm/metal_msm/host/metal_wrapper.rs (1)
239-244: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCheck the command buffer status before you return.
The function waits on the last command buffer but ignores its status. If the GPU aborts a stage, the caller reads stale buffer contents and treats them as a result. Inspect the status of each committed command buffer, or at minimum of the last one, and surface the error.
♻️ Proposed change
if let Some(last) = command_buffers.last() { last.wait_until_completed(); + for (i, cb) in command_buffers.iter().enumerate() { + if cb.status() == MTLCommandBufferStatus::Error { + panic!("Metal command buffer {} failed", i); + } + } }🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/host/metal_wrapper.rs` around lines 239 - 244, Update the command-buffer completion handling after the commit loop to inspect the status of the completed buffer, at minimum the buffer returned by command_buffers.last() after wait_until_completed(). Propagate a GPU execution failure to the caller instead of continuing with potentially stale output, while preserving successful result handling.mopro-msm/src/msm/metal_msm/shader/curve/jacobian.metal (1)
168-232: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
jacobian_madd_safehelper.
jacobian_madd_safehas no callers, while the XYZZ path already providesxyzz_madd_safe. Keeping this guarded mixed-addition only adds dead shader code.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/curve/jacobian.metal` around lines 168 - 232, Remove the unused jacobian_madd_safe helper and its entire implementation. Leave the existing xyzz_madd_safe path and all other Jacobian helpers unchanged.mopro-msm/src/msm/metal_msm/shader/cuzk/decompose_scalars.metal (2)
57-69: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a comment for the intentional unsigned wraparound.
slice_valcan be negative at Line 68.uint(slice_val) + sthen relies on unsigned wraparound to produce the biased digit in[0, l). The behavior is correct, but a reader can easily read it as a bug. State that the wraparound is intentional.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/cuzk/decompose_scalars.metal` around lines 57 - 69, Add a concise comment immediately before the chunks[id + offset] assignment in the scalar decomposition flow explaining that converting negative slice_val to uint and adding s intentionally uses unsigned wraparound to produce the biased digit in [0, l). Do not alter the existing computation.
50-55: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the assumed
window_sizeandnum_subtaskrelationship.The last-chunk shift is
((num_subtask * window_size - 256u) + 16u) - window_size. It computes in unsigned arithmetic. Ifnum_subtask * window_size < 240u + window_size, the expression underflows and the shift count becomes very large, which is undefined. Thescalar_byteshalfword layout also assumeswindow_size <= 16. Add a comment stating both preconditions, or clamp the shift.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/cuzk/decompose_scalars.metal` around lines 50 - 55, Document the preconditions directly above the last-chunk calculation in the decompose-scalars shader: require window_size <= 16 and num_subtask * window_size >= 240u + window_size to prevent unsigned shift underflow. Keep the existing calculation unchanged unless implementing an equivalent clamp.mopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_wide.rs (1)
48-62: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the wide carry limb.
The buffer holds
num_limbs + 1limbs, but the test reads and compares onlynum_limbs. The overflow case therefore never verifies the top carry limb, which is the only part ofbigint_add_widethat differs frombigint_add_unsafe. Readnum_limbs + 1limbs and assert the last limb equals the expected carry.♻️ Proposed test change
- let result_limbs = helper.read_results(&result_buf, config.num_limbs); + let all_limbs = helper.read_results(&result_buf, config.num_limbs + 1); + let result_limbs = all_limbs[..config.num_limbs].to_vec(); let expected_limbs = expected.to_limbs(config.num_limbs, config.log_limb_size); assert_eq!(result_limbs, expected_limbs); + let expected_carry = if require_overflow { 1u32 } else { 0u32 }; + assert_eq!(all_limbs[config.num_limbs], expected_carry, "carry limb mismatch");
run_bigint_add_testneeds therequire_overflowflag passed in for this assertion.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_wide.rs` around lines 48 - 62, Update run_bigint_add_test to accept and use a require_overflow flag, read num_limbs + 1 limbs from result_buf, and compare the full result against the expected limbs including the carry. When require_overflow is enabled, explicitly assert that the final limb equals the expected carry value.mopro-msm/src/msm/metal_msm/shader/cuzk/transpose.metal (1)
5-8: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe production transpose path has no dedicated test.
The comment states that the legacy single-kernel
transposeis kept for tests while the pipeline uses the three new kernels. Test coverage then applies only to the unused kernel. Add a test that runstranspose_histogram,transpose_prefix_sum, andtranspose_scatter, and that compares the result against the legacy kernel with per-column sorting, since column order is unspecified.I can draft that test or open an issue to track it.
🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/cuzk/transpose.metal` around lines 5 - 8, Add dedicated coverage for the production transpose pipeline by invoking transpose_histogram, transpose_prefix_sum, and transpose_scatter in sequence. Compare its output with the legacy transpose result after sorting values within each column, preserving the unordered-column behavior while ensuring both paths produce equivalent contents.mopro-msm/src/msm/metal_msm/shader/mont_backend/mont.metal (1)
65-96: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider reusing
conditional_reducefor the final subtraction.The file already defines
conditional_reduce, andBigIntprovides>=and-. The inlined comparison and borrow-subtraction duplicate that logic. If the manual version exists only for performance, add a short comment stating that; otherwise build theBigIntfromtand callconditional_reduce.🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/mont_backend/mont.metal` around lines 65 - 96, Replace the duplicated final comparison and borrow-subtraction logic in the surrounding Montgomery reduction function with the existing conditional_reduce helper, constructing a BigInt from t and relying on BigInt’s >= and - operators. If retaining the manual implementation for performance, add a concise comment documenting that rationale.mopro-msm/src/msm/metal_msm/shader/misc/get_constant.metal (1)
14-18: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a compile-time guard for unsupported
NUM_LIMBS.If
NUM_LIMBS != 8,MODULUSbecomes an empty initializer and every limb is zero. Montgomery arithmetic then silently produces wrong results instead of failing the build. Add an#elsebranch with#error.🛡️ Proposed guard
constant BigInt MODULUS = { `#if` (NUM_LIMBS == 8) LIMBS_8 +#else +#error "Unsupported NUM_LIMBS: only 8 x 32-bit limbs are supported" `#endif` };🤖 Prompt for 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. In `@mopro-msm/src/msm/metal_msm/shader/misc/get_constant.metal` around lines 14 - 18, Add an `#else` branch to the NUM_LIMBS preprocessor conditional initializing MODULUS, and emit a compile-time `#error` for every value other than 8. Preserve the existing LIMBS_8 initializer for NUM_LIMBS == 8 so unsupported configurations fail compilation instead of creating a zero modulus.
🤖 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 `@mopro-msm/src/msm/metal_msm/metal_msm.rs`:
- Around line 291-295: Validate the parsed `MSM_SMVP_L` value in the
segment-length initialization before calculating `num_segments`. Ensure zero is
rejected and replaced with the existing default length of 128, so the division
in the `num_segments` calculation always uses a positive `seg_len`.
- Around line 144-164: Update the dispatch setup around c_num_y_workgroups to
use ceiling division by c_workgroup_size * c_num_x_workgroups, ensuring all
input scalars are covered. In decompose_scalars, guard threads whose global
index is >= input_size so they return before any out-of-range reads or writes.
In `@mopro-msm/src/msm/metal_msm/utils/limbs_conversion.rs`:
- Around line 389-440: The 32-bit zero-copy path assumes exactly eight 32-bit
limbs but does not enforce that invariant. Add a debug assertion at the
`msm_config.log_limb_size == 32` guard requiring `msm_config.num_limbs == 8`,
and update public/default API paths that enable this mode to enforce the same
requirement without changing other limb-size behavior.
---
Outside diff comments:
In `@mopro-msm/src/msm/metal_msm/metal_msm.rs`:
- Around line 567-596: Validate the parsed overrides in the window-size and
scale-factor selection logic: clamp MSM_WINDOW_SIZE to a supported positive
range below the shift-overflow limit, and clamp MSM_SCALE_FACTOR to a supported
positive range that keeps c_workgroup_size nonzero. Preserve the existing
size-based defaults when the environment variables are absent or invalid.
In `@mopro-msm/src/msm/metal_msm/utils/limbs_conversion.rs`:
- Around line 556-580: Extend test_parallel_packing_correctness with a dedicated
log_limb_size: 32, num_limbs: 8 case that exercises the zero-copy branch in
pack_affine_and_scalars. Compare its packed coordinates and scalars against a
correct Montgomery-domain reference built from into_bigint() and
calc_mont_radix, preserving the existing 16-bit assertions.
---
Nitpick comments:
In `@mopro-msm/benches/shaders.rs`:
- Around line 64-69: Update the MetalConfig for the transpose benchmark,
identified by kernel_name "transpose", to use log_limb_size 32 and num_limbs 8,
matching the other benchmark configurations while leaving the shader_file and
kernel name unchanged.
- Around line 402-411: Update smvp_gpu to return all four bucket coordinate
buffers, including bucket_zzz_buf, matching the equivalent cuzk smvp helper.
Adjust the function’s tuple return type and every return/destructuring site in
the benchmark, while preserving the existing readback used to force GPU
completion.
In `@mopro-msm/src/msm/metal_msm/host/metal_wrapper.rs`:
- Around line 239-244: Update the command-buffer completion handling after the
commit loop to inspect the status of the completed buffer, at minimum the buffer
returned by command_buffers.last() after wait_until_completed(). Propagate a GPU
execution failure to the caller instead of continuing with potentially stale
output, while preserving successful result handling.
In `@mopro-msm/src/msm/metal_msm/metal_msm.rs`:
- Around line 239-262: Remove the collapsed chunk loop around the legacy SMVP
setup and construct a single params buffer and dimension entry directly. Replace
the derived offset, remaining_subtasks, and chunking arithmetic with the
equivalent num_subtasks-based values while preserving the existing
smvp_params_bufs and smvp_dims contents and single-dispatch behavior.
- Around line 540-543: Replace the thread-local PIPELINE cache used by
metal_variable_base_msm with a process-wide shared OnceCell, verifying that
MetalMSMPipeline is Send + Sync and concurrent execute_pipeline calls are safe;
otherwise synchronize access appropriately. Preserve lazy initialization and
reuse the single pipeline across MSM calls.
In `@mopro-msm/src/msm/metal_msm/shader/bigint/bigint.metal`:
- Around line 14-22: Add a compile-time assertion near the top of bigint.metal
that requires LOG_LIMB_SIZE to equal 32, covering both the addition and
subtraction paths that rely on uint truncation. Keep the existing limb
arithmetic unchanged and ensure generation with any other limb size fails during
compilation.
In `@mopro-msm/src/msm/metal_msm/shader/curve/jacobian.metal`:
- Around line 168-232: Remove the unused jacobian_madd_safe helper and its
entire implementation. Leave the existing xyzz_madd_safe path and all other
Jacobian helpers unchanged.
In `@mopro-msm/src/msm/metal_msm/shader/cuzk/decompose_scalars.metal`:
- Around line 57-69: Add a concise comment immediately before the chunks[id +
offset] assignment in the scalar decomposition flow explaining that converting
negative slice_val to uint and adding s intentionally uses unsigned wraparound
to produce the biased digit in [0, l). Do not alter the existing computation.
- Around line 50-55: Document the preconditions directly above the last-chunk
calculation in the decompose-scalars shader: require window_size <= 16 and
num_subtask * window_size >= 240u + window_size to prevent unsigned shift
underflow. Keep the existing calculation unchanged unless implementing an
equivalent clamp.
In `@mopro-msm/src/msm/metal_msm/shader/cuzk/transpose.metal`:
- Around line 5-8: Add dedicated coverage for the production transpose pipeline
by invoking transpose_histogram, transpose_prefix_sum, and transpose_scatter in
sequence. Compare its output with the legacy transpose result after sorting
values within each column, preserving the unordered-column behavior while
ensuring both paths produce equivalent contents.
In `@mopro-msm/src/msm/metal_msm/shader/misc/get_constant.metal`:
- Around line 14-18: Add an `#else` branch to the NUM_LIMBS preprocessor
conditional initializing MODULUS, and emit a compile-time `#error` for every value
other than 8. Preserve the existing LIMBS_8 initializer for NUM_LIMBS == 8 so
unsupported configurations fail compilation instead of creating a zero modulus.
In `@mopro-msm/src/msm/metal_msm/shader/mont_backend/mont.metal`:
- Around line 65-96: Replace the duplicated final comparison and
borrow-subtraction logic in the surrounding Montgomery reduction function with
the existing conditional_reduce helper, constructing a BigInt from t and relying
on BigInt’s >= and - operators. If retaining the manual implementation for
performance, add a concise comment documenting that rationale.
In `@mopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_wide.rs`:
- Around line 48-62: Update run_bigint_add_test to accept and use a
require_overflow flag, read num_limbs + 1 limbs from result_buf, and compare the
full result against the expected limbs including the carry. When
require_overflow is enabled, explicitly assert that the final limb equals the
expected carry value.
In
`@mopro-msm/src/msm/metal_msm/tests/cuzk/convert_point_coords_and_decompose_scalars.rs`:
- Around line 140-143: Remove the identity helper pack_limbs and update both
call sites to use the scalar limb slices directly. Rename the test file from
convert_point_coords_and_decompose_scalars.rs to decompose_scalars.rs,
preserving the existing scalar decomposition behavior.
🪄 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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 20449c43-863d-4093-9ca6-20c7eb0f22f5
📒 Files selected for processing (51)
mopro-msm/benches/shaders.rsmopro-msm/build.rsmopro-msm/src/msm/metal_msm/host/metal_wrapper.rsmopro-msm/src/msm/metal_msm/host/shader.rsmopro-msm/src/msm/metal_msm/host/shader_manager.rsmopro-msm/src/msm/metal_msm/metal_msm.rsmopro-msm/src/msm/metal_msm/shader/bigint/bigint.metalmopro-msm/src/msm/metal_msm/shader/constants.metalmopro-msm/src/msm/metal_msm/shader/curve/jacobian.metalmopro-msm/src/msm/metal_msm/shader/curve/xyzz.metalmopro-msm/src/msm/metal_msm/shader/cuzk/barrett_reduction.metalmopro-msm/src/msm/metal_msm/shader/cuzk/convert_point_coords_and_decompose_scalars.metalmopro-msm/src/msm/metal_msm/shader/cuzk/decompose_scalars.metalmopro-msm/src/msm/metal_msm/shader/cuzk/kernel_barrett_reduction.metalmopro-msm/src/msm/metal_msm/shader/cuzk/kernel_field_mul.metalmopro-msm/src/msm/metal_msm/shader/cuzk/pbpr.metalmopro-msm/src/msm/metal_msm/shader/cuzk/smvp.metalmopro-msm/src/msm/metal_msm/shader/cuzk/smvp_segmented.metalmopro-msm/src/msm/metal_msm/shader/cuzk/transpose.metalmopro-msm/src/msm/metal_msm/shader/misc/get_constant.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_cios_benchmarks.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_modified.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_modified_benchmarks.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_optimised.metalmopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_optimised_benchmarks.metalmopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_unsafe.rsmopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_wide.rsmopro-msm/src/msm/metal_msm/tests/bigint/bigint_sub.rsmopro-msm/src/msm/metal_msm/tests/curve/jacobian_add_2007_b1.rsmopro-msm/src/msm/metal_msm/tests/curve/jacobian_dbl_2009_l.rsmopro-msm/src/msm/metal_msm/tests/curve/jacobian_madd_2007_bl.rsmopro-msm/src/msm/metal_msm/tests/curve/jacobian_neg.rsmopro-msm/src/msm/metal_msm/tests/curve/jacobian_scalar_mul.rsmopro-msm/src/msm/metal_msm/tests/cuzk/barrett_reduction.rsmopro-msm/src/msm/metal_msm/tests/cuzk/convert_point_coords_and_decompose_scalars.rsmopro-msm/src/msm/metal_msm/tests/cuzk/mod.rsmopro-msm/src/msm/metal_msm/tests/cuzk/pbpr.rsmopro-msm/src/msm/metal_msm/tests/cuzk/smvp.rsmopro-msm/src/msm/metal_msm/tests/cuzk/transpose.rsmopro-msm/src/msm/metal_msm/tests/field/ff_add.rsmopro-msm/src/msm/metal_msm/tests/field/ff_reduce.rsmopro-msm/src/msm/metal_msm/tests/field/ff_sub.rsmopro-msm/src/msm/metal_msm/tests/misc/get_constant.rsmopro-msm/src/msm/metal_msm/tests/mont_backend/mod.rsmopro-msm/src/msm/metal_msm/tests/mont_backend/mont_benchmarks.rsmopro-msm/src/msm/metal_msm/tests/mont_backend/mont_mul_cios.rsmopro-msm/src/msm/metal_msm/tests/mont_backend/mont_mul_modified.rsmopro-msm/src/msm/metal_msm/tests/mont_backend/mont_mul_optimised.rsmopro-msm/src/msm/metal_msm/utils/limbs_conversion.rsmopro-msm/src/msm/metal_msm/utils/mont_params.rs
💤 Files with no reviewable changes (15)
- mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_optimised_benchmarks.metal
- mopro-msm/src/msm/metal_msm/shader/cuzk/kernel_field_mul.metal
- mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_modified.metal
- mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_optimised.metal
- mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_modified_benchmarks.metal
- mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_cios_benchmarks.metal
- mopro-msm/src/msm/metal_msm/shader/cuzk/convert_point_coords_and_decompose_scalars.metal
- mopro-msm/src/msm/metal_msm/shader/cuzk/kernel_barrett_reduction.metal
- mopro-msm/src/msm/metal_msm/tests/mont_backend/mont_benchmarks.rs
- mopro-msm/src/msm/metal_msm/tests/mont_backend/mont_mul_optimised.rs
- mopro-msm/src/msm/metal_msm/tests/mont_backend/mod.rs
- mopro-msm/src/msm/metal_msm/tests/mont_backend/mont_mul_modified.rs
- mopro-msm/src/msm/metal_msm/tests/cuzk/barrett_reduction.rs
- mopro-msm/src/msm/metal_msm/tests/cuzk/mod.rs
- mopro-msm/src/msm/metal_msm/shader/cuzk/barrett_reduction.metal
Comment on lines
144
to
+164
| let c_workgroup_size = self.simd_width * scale_factor; | ||
| let c_num_x_workgroups = c_workgroup_size; | ||
| let c_num_y_workgroups = input_size / c_workgroup_size / c_num_x_workgroups; | ||
| let c_num_z_workgroups = 1; | ||
|
|
||
| let (point_x, point_y, scalar_chunks) = stage1.execute( | ||
| &coords, | ||
| &scals, | ||
| input_size, | ||
| window_size, | ||
| num_columns, | ||
| num_subtasks, | ||
| c_num_x_workgroups, | ||
| c_num_y_workgroups, | ||
| c_num_z_workgroups, | ||
| c_workgroup_size, | ||
| )?; | ||
|
|
||
| // Stage 2: Transpose | ||
| let stage2 = Transpose::new(&self.shader_manager); | ||
|
|
||
| let t_num_x_workgroups = 1; | ||
| let t_num_y_workgroups = 1; | ||
| let t_num_z_workgroups = 1; | ||
| let t_workgroup_size = num_subtasks; | ||
|
|
||
| let (csc_col_ptr, csc_val_idxs) = stage2.execute( | ||
| &scalar_chunks, | ||
| num_subtasks, | ||
| input_size, | ||
| num_columns, | ||
| t_num_x_workgroups, | ||
| t_num_y_workgroups, | ||
| t_num_z_workgroups, | ||
| t_workgroup_size, | ||
| )?; | ||
| let convert_params_buf = helper.create_buffer(&vec![ | ||
| input_size as u32, | ||
| window_size as u32, | ||
| num_columns as u32, | ||
| num_subtasks as u32, | ||
| ]); | ||
|
|
||
| // Stage 3: Sparse Matrix-Vector Multiplication | ||
| // according to the write-up (https://github.com/z-prize/2023-entries/tree/main/prize-2-msm-wasm/webgpu-only/tal-derei-koh-wei-jie#thread-count) | ||
| // the threads is (input size / 2) * num_subtasks | ||
| // after testing, 1D dim config is better than 3D dim config | ||
| let stage3 = SMVP::new(&self.shader_manager); | ||
| ops.push(ShaderOperation { | ||
| pipeline_state: convert_shader.pipeline_state.clone(), | ||
| buffers: vec![&scalars_buf, &scalar_chunks_buf, &convert_params_buf], | ||
| thread_group_count: helper.create_thread_group_size( | ||
| c_num_x_workgroups as u64, | ||
| c_num_y_workgroups.max(1) as u64, | ||
| 1, | ||
| ), | ||
| threads_per_threadgroup: helper.create_thread_group_size(c_workgroup_size as u64, 1, 1), | ||
| }); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate the decompose_scalars kernel and inspect its bounds check.
fd -t f 'decompose_scalars.metal' | while IFS= read -r f; do
echo "=== $f ==="
cat -n "$f"
doneRepository: zkmopro/gpu-acceleration
Length of output: 3090
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate metal_msm.rs:"
fd -t f '^metal_msm\.rs$' .
echo
echo "Find relevant symbols/line ranges:"
rg -n "metal_variable_base_msm|decompose_scalars|create_thread_group_size|simd_width|c_workgroup_size|c_num_x_workgroups|c_num_y_workgroups|scale_factor" mopro-msm/src -S
echo
echo "Inspect metal_msm.rs around relevant functions:"
wc -l mopro-msm/src/msm/metal_msm/metal_msm.rs
sed -n '1,240p' mopro-msm/src/msm/metal_msm/metal_msm.rs
echo
echo "Inspect remaining file around variable_base_msm usages:"
sed -n '240,440p' mopro-msm/src/msm/metal_msm/metal_msm.rsRepository: zkmopro/gpu-acceleration
Length of output: 33529
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Inspect variable_base_msm and remainder of execute_pipeline:"
sed -n '440,720p' mopro-msm/src/msm/metal_msm/metal_msm.rs
echo
echo "Search for scalar decomposition bound checks and tests with non-power-of-two sizes:"
rg -n "input_size|num_subtasks|id <|global index|thread_position_in_grid|metal_variable_base_msm|convert_point|decompose" mopro-msm/src/msm/metal_msm -S
echo
echo "Standalone arithmetic check for dispatched threads and coverage:"
python3 - <<'PY'
cases = [
(32, 2, 100000),
(32, 1, 100000),
(32, 2, 8192),
(32, 2, 8193),
]
for simd_width, scale_factor, input_size in cases:
c_workgroup_size = simd_width * scale_factor
c_num_x_workgroups = c_workgroup_size
c_num_y_workgroups = input_size // c_workgroup_size // c_num_x_workgroups
dispatched = c_workgroup_size * (max(1, c_num_y_workgroups)) * c_num_x_workgroups
covered = min(input_size, dispatched)
rest = input_size - covered
ceiling_y = (input_size + c_workgroup_size * c_num_x_workgroups - 1) // (c_workgroup_size * c_num_x_workgroups)
ceiling_dispatched = c_workgroup_size * ceiling_y * c_num_x_workgroups
print({
"simd_width": simd_width,
"scale_factor": scale_factor,
"input_size": input_size,
"c_workgroup_size": c_workgroup_size,
"current_y": c_num_y_workgroups,
"current_dispatched": dispatched,
"rest": rest,
"ceiling_y": ceiling_y,
"ceiling_dispatched": ceiling_dispatched,
})
PY
echo
echo "Inspect Metal wrapper for thread limit availability:"
sed -n '90,130p' mopro-msm/src/msm/metal_msm/host/metal_wrapper.rsRepository: zkmopro/gpu-acceleration
Length of output: 32840
Dispatch scalar decomposition with ceiling division
metal_variable_base_msm can set scale_factor = 2 for 65,537–1,048,576 entries. With simd_width = 32, that makes c_workgroup_size = 64, and the current floor-based grid only dispatches c_workgroup_size^2 * floor(input_size / c_workgroup_size^2) threads. For 100,000 scalars, 1,696 chunks stay zero, and later stages use that partial result. Compute c_num_y_workgroups with ceiling division over c_workgroup_size * c_num_x_workgroups, and add a guard in decompose_scalars so threads with global index >= input_size do not read/write out of range.
🤖 Prompt for 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.
In `@mopro-msm/src/msm/metal_msm/metal_msm.rs` around lines 144 - 164, Update the
dispatch setup around c_num_y_workgroups to use ceiling division by
c_workgroup_size * c_num_x_workgroups, ensuring all input scalars are covered.
In decompose_scalars, guard threads whose global index is >= input_size so they
return before any out-of-range reads or writes.
Comment on lines
+291
to
+295
| let seg_len: usize = std::env::var("MSM_SMVP_L") | ||
| .ok() | ||
| .and_then(|v| v.parse().ok()) | ||
| .unwrap_or(128); | ||
| let num_segments = (input_size + seg_len - 1) / seg_len; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Guard MSM_SMVP_L against zero.
MSM_SMVP_L=0 parses successfully, so seg_len becomes 0 and num_segments divides by zero. That panics the MSM call.
🛡️ Proposed fix
let seg_len: usize = std::env::var("MSM_SMVP_L")
.ok()
.and_then(|v| v.parse().ok())
+ .filter(|&v| v > 0)
.unwrap_or(128);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| let seg_len: usize = std::env::var("MSM_SMVP_L") | |
| .ok() | |
| .and_then(|v| v.parse().ok()) | |
| .unwrap_or(128); | |
| let num_segments = (input_size + seg_len - 1) / seg_len; | |
| let seg_len: usize = std::env::var("MSM_SMVP_L") | |
| .ok() | |
| .and_then(|v| v.parse().ok()) | |
| .filter(|&v| v > 0) | |
| .unwrap_or(128); | |
| let num_segments = (input_size + seg_len - 1) / seg_len; |
🤖 Prompt for 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.
In `@mopro-msm/src/msm/metal_msm/metal_msm.rs` around lines 291 - 295, Validate
the parsed `MSM_SMVP_L` value in the segment-length initialization before
calculating `num_segments`. Ensure zero is rejected and replaced with the
existing default length of 128, so the division in the `num_segments`
calculation always uses a positive `seg_len`.
Comment on lines
+389
to
+440
|
|
||
| if msm_config.log_limb_size == 32 { | ||
| // Limbs are plain little-endian 32-bit words. | ||
| // | ||
| // Coordinates are copied in arkworks' internal Montgomery representation | ||
| // (R = 2^256, identical to the GPU's Montgomery radix), so no | ||
| // Montgomery reduction happens on the host and no to-Montgomery | ||
| // conversion is needed on the GPU. Scalars go through into_bigint() | ||
| // because window decomposition needs the standard representation. | ||
| let words_per_coord = msm_config.num_limbs; | ||
| let coords_per_point = words_per_coord * 2; | ||
| let mut coords = vec![0u32; num_elements * coords_per_point]; | ||
| let mut scalars_u32 = vec![0u32; num_elements * words_per_coord]; | ||
|
|
||
| let write_words = |v: &[u64; 4], out: &mut [u32]| { | ||
| for (i, w) in v.iter().enumerate() { | ||
| out[2 * i] = *w as u32; | ||
| out[2 * i + 1] = (*w >> 32) as u32; | ||
| } | ||
| }; | ||
|
|
||
| const CHUNK_SIZE: usize = 1024; | ||
| bases | ||
| .par_chunks(CHUNK_SIZE) | ||
| .zip(scalars.par_chunks(CHUNK_SIZE)) | ||
| .zip(coords.par_chunks_mut(CHUNK_SIZE * coords_per_point)) | ||
| .zip(scalars_u32.par_chunks_mut(CHUNK_SIZE * words_per_coord)) | ||
| .for_each( | ||
| |(((base_chunk, scalar_chunk), coord_chunk), scalar_u32_chunk)| { | ||
| for (i, (pt, sc)) in base_chunk.iter().zip(scalar_chunk.iter()).enumerate() { | ||
| let coord_start = i * coords_per_point; | ||
| write_words( | ||
| &pt.x.0 .0, | ||
| &mut coord_chunk[coord_start..coord_start + words_per_coord], | ||
| ); | ||
| write_words( | ||
| &pt.y.0 .0, | ||
| &mut coord_chunk | ||
| [coord_start + words_per_coord..coord_start + coords_per_point], | ||
| ); | ||
| let scalar_start = i * words_per_coord; | ||
| write_words( | ||
| &sc.into_bigint().0, | ||
| &mut scalar_u32_chunk[scalar_start..scalar_start + words_per_coord], | ||
| ); | ||
| } | ||
| }, | ||
| ); | ||
|
|
||
| return (coords, scalars_u32); | ||
| } | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find all MetalConfig construction sites and check log_limb_size/num_limbs pairing.
rg -n -B2 -A6 'MetalConfig\s*\{' --type rustRepository: zkmopro/gpu-acceleration
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Files matching config/limbs names:"
fd -i 'config|limbs|metal|mopro' . | sed -n '1,120p'
echo
echo "Find MetalConfig type/struct definitions and usages:"
rg -n -i 'MetalConfig|log_limb_size|num_limbs|limbs_conversion|write_words' --type rust | sed -n '1,240p'Repository: zkmopro/gpu-acceleration
Length of output: 3153
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Repo files:"
git ls-files | sed -n '1,120p'
echo
echo "Locate limbs_conversion.rs:"
fd -a limbs_conversion.rs . || true
echo
echo "Relevant source around lines 360-450:"
file="$(fd 'limbs_conversion.rs' . | head -n 1)"
if [ -n "${file:-}" ]; then
wc -l "$file"
sed -n '1,520p' "$file"
fiRepository: zkmopro/gpu-acceleration
Length of output: 24818
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find struct/type/defaults definitions:"
rg -n -i 'struct MetalConfig|type MetalConfig|derive\(Default|metal_config|Ms' mopro-msm/src/msm/metal_msm mopro-msm/src/msm/mod.rs mopro-msm/src/lib.rs --type rust
echo
echo "Inspect host files likely containing config/default:"
for f in mopro-msm/src/msm/metal_msm/host/errors.rs \
mopro-msm/src/msm/metal_msm/host/gpu.rs \
mopro-msm/src/msm/metal_msm/host/metal_wrapper.rs \
mopro-msm/src/msm/metal_msm/metal_msm.rs \
mopro-msm/src/msm/metal_msm/mod.rs \
mopro-msm/src/lib.rs \
mopro-msm/src/msm/mod.rs; do
[ -f "$f" ] || continue
echo "--- $f ---"
ast-grep outline "$f" 2>/dev/null | sed -n '1,160p' || true
done
echo
echo "Search for default-like config values and field names:"
rg -n 'Default|Default.*Metal|log_limb_size|num_limbs|const .*MONT|MONT|Mont' mopro-msm/src/msm --type rust | sed -n '1,240p'Repository: zkmopro/gpu-acceleration
Length of output: 50380
Add an invariant check for the zero-copy 32-bit limb path.
write_words only packs a fixed 4-element u64 coordinate/scalar into 8 u32 limbs. If log_limb_size: 32 can be used with num_limbs != 8, this path can panic or truncate the Montgomery value. Add a debug assert at the guard and update any public/default API paths to require num_limbs == 8 for this zero-copy mode.
🤖 Prompt for 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.
In `@mopro-msm/src/msm/metal_msm/utils/limbs_conversion.rs` around lines 389 -
440, The 32-bit zero-copy path assumes exactly eight 32-bit limbs but does not
enforce that invariant. Add a debug assertion at the `msm_config.log_limb_size
== 32` guard requiring `msm_config.num_limbs == 8`, and update public/default
API paths that enable this mode to enforce the same requirement without changing
other limb-size behavior.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
perf: 8-9x faster Metal MSM using bunch of optimizations
Summary
This PR rewrites the hot paths of the Metal MSM pipeline end to end. On
M4 Air,
metal_variable_base_msmis 4.4-9.5x faster than v0.2.0 across 2^12-2^22 and now beats Arkworks CPU
MSM from 2^16 upward. Peak
GPU memory at the top sizes is lower than v0.2.0 despite the pipeline doing
strictly more on-device.
(Medians of 5 runs)
What changed
MSM_COMPILE_ALL_SHADERSchanges, so the test metallib can no longer silently miss test kernels on
cached build-script runs.
device-resident buffers and committed back-to-back with one host sync per
MSM. Before, every stage did a full wait, readback, and re-upload.
ulongfusedmultiply-add CIOS replace the WGSL-inherited 16x16-bit representation. 4x
fewer partial products per Montgomery multiplication and half the register
footprint.
madd-2007-bl formula (11 field muls instead of 16).
additions with far fewer field add/subs; the identity is encoded as ZZ = 0
so zero-initialized buffers still mean infinity. The host converts final
sums to Jacobian without inversions via scale equivalence.
becomes three parallel kernels: histogram with relaxed atomics, prefix sum,
and scatter with per-column atomic slot claims.
the same radix the GPU uses, so coordinates pass through raw, deleting a
CPU Montgomery reduction and a GPU multiplication per coordinate that
converted a number into itself. SMVP reads the input buffer directly as
interleaved affine records.
on the GPU in threadgroup memory, so the host only Horner-combines one
point per subtask. This freed PBPR to use 256-thread workgroups.
four, so the scheduler can overlap load-imbalance tails.
bucket-sorted point stream, giving uniform per-thread work. Sortedness
guarantees a run starting and ending inside a slice covers its bucket
completely and is written exclusively without atomics; only a slice's first
and last runs can be partial, so each thread spills at most two entries,
and a merge kernel folds spills and mirror rows into signed buckets. This
also fixes the window-12/14 load pathology, so every window size is safe.
across subtask chunks of four, keeping the segmented overhead to a few
tens of MB.
scale factors, and env knobs for future machines:
MSM_WINDOW_SIZE,MSM_SCALE_FACTOR,MSM_SMVP_L,MSM_LEGACY_SMVP, andMSM_PROFILE=1for per-dispatch timings and occupancy.
Some Attempts
CIOS for BN254. rho is close to p, forcing a second conditional
subtraction that cancels the saved multiplications(?)
XYZZ, and assembling such batches on Apple GPUs costs more than it saves.
Summary by CodeRabbit
Performance
Reliability
Maintenance