From ac980a856e81588e15beb4c6e548e7e27c11f142 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:15:42 +0000 Subject: [PATCH 01/34] cuda: hash single-chunk BLAKE3 rows with one thread per message --- cuda/blake3_rows_bench.cu | 78 +++++++++++++++++++++++++++++ cuda/kernels.cu | 33 +++++++++++++ docs/cuda-blake3-short-rows.md | 85 ++++++++++++++++++++++++++++++++ examples/cuda_commit_bench.rs | 90 ++++++++++++++++++++++++++++++++++ src/cuda/mod.rs | 68 ++++++++++++++----------- 5 files changed, 326 insertions(+), 28 deletions(-) create mode 100644 cuda/blake3_rows_bench.cu create mode 100644 docs/cuda-blake3-short-rows.md create mode 100644 examples/cuda_commit_bench.rs diff --git a/cuda/blake3_rows_bench.cu b/cuda/blake3_rows_bench.cu new file mode 100644 index 0000000..fef03d6 --- /dev/null +++ b/cuda/blake3_rows_bench.cu @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// Standalone resident-kernel comparison; see docs/cuda-blake3-short-rows.md. +// Both kernels run against the same resident bytes, with alternating order. +#include +#include +#include +#include +#include "kernels.cu" + +void checked(cudaError_t error) { + if (error != cudaSuccess) { + fprintf(stderr, "CUDA: %s\n", cudaGetErrorString(error)); + exit(1); + } +} +void launch(bool old, uint8_t* output, const uint8_t* input, + size_t bytes, size_t rows, bool force_stride) { + if (old) { + size_t blocks = std::min((rows + THREADS / 32 - 1) / (THREADS / 32), MAX_BLOCKS); + blake3_hash_rows_kernel<<>>(output, input, bytes, rows); + checked(cudaGetLastError()); + } else if (force_stride && bytes <= BLAKE3_CHUNK_BYTES) { + // Exercise multiple grid-stride iterations without a 16M-row fixture. + blake3_hash_short_rows_kernel<<<2, THREADS, 0, cudaStreamPerThread>>>(output, input, bytes, rows); + checked(cudaGetLastError()); + } else { + checked(launch_blake3_rows(output, input, bytes, rows)); + } +} +int main(int argc, char** argv) { + if (argc > 2 || (argc == 2 && std::strcmp(argv[1], "--check") != 0)) { + fprintf(stderr, "usage: %s [--check]\n", argv[0]); + return 1; + } + const bool small = argc == 2; + checked(cudaSetDevice(0)); + printf("row_bytes,rows,iteration,warp_ms,dispatch_ms,speedup\n"); + for (size_t bytes : {size_t(1), size_t(8), size_t(16), size_t(32), size_t(64), + size_t(128), size_t(256), size_t(320), size_t(512), + size_t(1023), size_t(1024), size_t(1025), size_t(4264), size_t(7400)}) { + size_t rows = small ? 777 : std::min(1 << 20, (128 << 20) / bytes); + std::vector input(bytes * rows); + uint64_t state = 0x81726ab567846217ULL; + for (auto& byte : input) { + state ^= state << 13; state ^= state >> 7; state ^= state << 17; + byte = state; + } + DeviceBuffer source, first, second; + checked(source.allocate((input.size() + 7) / 8)); + checked(first.allocate(rows * 4)); + checked(second.allocate(rows * 4)); + auto* d_source = reinterpret_cast(source.get()); + auto* d_first = reinterpret_cast(first.get()); + auto* d_second = reinterpret_cast(second.get()); + checked(cudaMemcpy(d_source, input.data(), input.size(), cudaMemcpyHostToDevice)); + launch(true, d_first, d_source, bytes, rows, false); + launch(false, d_second, d_source, bytes, rows, small); + std::vector expected(rows * 32), actual(rows * 32); + checked(cudaMemcpy(expected.data(), d_first, expected.size(), cudaMemcpyDeviceToHost)); + checked(cudaMemcpy(actual.data(), d_second, actual.size(), cudaMemcpyDeviceToHost)); + if (expected != actual) { fprintf(stderr, "mismatch for %zu-byte rows\n", bytes); return 2; } + cudaEvent_t start, stop; + checked(cudaEventCreate(&start)); checked(cudaEventCreate(&stop)); + for (int iteration = 0; iteration < (small ? 1 : 7); ++iteration) { + float ms[2]; + for (int order = 0; order < 2; ++order) { + int which = (order + iteration) % 2; + checked(cudaEventRecord(start, cudaStreamPerThread)); + launch(which == 0, which == 0 ? d_first : d_second, d_source, bytes, rows, small); + checked(cudaEventRecord(stop, cudaStreamPerThread)); + checked(cudaEventSynchronize(stop)); + checked(cudaEventElapsedTime(&ms[which], start, stop)); + } + printf("%zu,%zu,%d,%.6f,%.6f,%.3f\n", bytes, rows, iteration, ms[0], ms[1], ms[0]/ms[1]); + } + checked(cudaEventDestroy(start)); checked(cudaEventDestroy(stop)); + } +} diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 9ccdf37..96d90dc 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -23,6 +23,7 @@ constexpr uint64_t GOLDILOCKS_P = 0xffffffff00000001ULL; constexpr uint64_t GOLDILOCKS_EPSILON = 0x00000000ffffffffULL; constexpr unsigned int THREADS = 256; constexpr unsigned int MAX_BLOCKS = 65535; +constexpr size_t BLAKE3_CHUNK_BYTES = 1024; constexpr uint32_t BLAKE3_CHUNK_START = 1U << 0; constexpr uint32_t BLAKE3_CHUNK_END = 1U << 1; constexpr uint32_t BLAKE3_PARENT = 1U << 2; @@ -1793,6 +1794,32 @@ __global__ void goldilocks_ops_kernel(uint64_t* sums, uint64_t* differences, } } +// A single-chunk message has no chunk tree to reduce. Assigning one thread +// per row keeps every lane doing useful compression instead of leaving 31 +// lanes idle in the warp-per-row kernel. Reuse the same chunk primitive and +// little-endian digest encoding; commitments and BLAKE3 flags are unchanged. +__global__ void blake3_hash_short_rows_kernel(uint8_t* digests, + const uint8_t* messages, + size_t message_bytes, + size_t message_count) { + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < message_count; index += grid_stride) { + uint32_t output[8]; + blake3_chunk(messages + index * message_bytes, message_bytes, 0, true, + output); + uint8_t* digest = digests + index * 32; +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + const uint32_t value = output[word]; + digest[word * 4] = static_cast(value); + digest[word * 4 + 1] = static_cast(value >> 8); + digest[word * 4 + 2] = static_cast(value >> 16); + digest[word * 4 + 3] = static_cast(value >> 24); + } + } +} + // One warp owns one message. Chunk compression is distributed across lanes; // lane zero then reduces the (at most 32) chunk chaining values into the root. // This matches the BLAKE3 tree shape while exposing parallelism within wide @@ -1890,6 +1917,12 @@ __global__ void blake3_hash_digest_pairs_kernel( cudaError_t launch_blake3_rows(uint8_t* digests, const uint8_t* messages, size_t message_bytes, size_t message_count) { + if (message_bytes <= BLAKE3_CHUNK_BYTES) { + blake3_hash_short_rows_kernel<<>>( + digests, messages, message_bytes, message_count); + return cudaGetLastError(); + } constexpr unsigned int WARPS_PER_BLOCK = THREADS / 32; const size_t required = (message_count + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; diff --git a/docs/cuda-blake3-short-rows.md b/docs/cuda-blake3-short-rows.md new file mode 100644 index 0000000..ec0eefd --- /dev/null +++ b/docs/cuda-blake3-short-rows.md @@ -0,0 +1,85 @@ +# Short-row BLAKE3 dispatch: 2026-09-11 + +Base: multi-stark `231942a` on `sb/trace-sharding-gpu`. +Isolated branch: `codex/cuda-blake3-short-rows`. +GPU: RTX PRO 6000 Blackwell Server Edition (96 GB), CUDA 13.3, native `sm_120`. +CPU-side benchmarks used `RAYON_NUM_THREADS=8`. + +## Change + +Messages of 1–1024 bytes use one thread per row, reusing our existing +BLAKE3 chunk compression and raw little-endian digest encoding. Longer +messages keep the existing warp-per-row kernel. The dispatch boundary is +BLAKE3's chunk size; there is no machine-dependent threshold or new runtime +flag. No upstream implementation was copied and no dependency was added. + +The new kernel uses 64 registers/thread versus 96 for the old kernel, with +no register spills in either on `sm_120`. It also avoids the old kernel's +8 KiB shared chunk-value array. Both report a 64-byte local stack frame. + +## Measurements + +The resident benchmark runs both kernels against the same uploaded bytes, +alternates their order, checks all output digests, and times with CUDA events. +Medians of seven samples; at most 1,048,576 rows / 128 MiB input per shape: + +| Row bytes | Existing kernel | New dispatch | Speedup | +| ---: | ---: | ---: | ---: | +| 8 | 1.321 ms | 0.384 ms | 3.44× | +| 64 | 1.353 ms | 0.333 ms | 4.07× | +| 320 | 2.351 ms | 0.413 ms | 5.70× | +| 1023 | 2.488 ms | 0.347 ms | 7.16× | +| 1024 | 2.480 ms | 0.472 ms | 5.26× | +| 1025 | 2.632 ms | 2.629 ms | 1.00× | +| 4264 | 0.801 ms | 0.800 ms | 1.00× | +| 7400 | 0.678 ms | 0.677 ms | 1.00× | + +Complete PCS commitments include host upload, LDEs, row gathering/hashing, +Merkle construction, and commitment return. Allocation/filling of the input +witness precedes timing. Three baseline/candidate processes per variant, +five warm samples each, with later process order reversed; first iterations +are excluded. Commitments matched byte-for-byte for every shape/run. + +| Input shape (height × field columns) | Existing | New | Wall reduction | +| --- | ---: | ---: | ---: | +| 2^20 × 8 (64-byte rows) | 11.403 ms | 9.011 ms | 21.0% | +| 2^18 × 40 (320-byte rows) | 12.871 ms | 10.594 ms | 17.7% | +| 2^18 × 128 (1024-byte rows) | 43.977 ms | 35.930 ms | 18.3% | +| 2^18 × 129 (1032-byte rows) | 44.759 ms | 43.986 ms | 1.7% | +| 2^16 × 925 (7400-byte rows) | 94.049 ms | 91.704 ms | 2.5% | +| Mixed heights/widths | 71.689 ms | 68.297 ms | 4.7% | + +The mixed case combines `(2^18,4)`, `(2^18,12)`, `(2^18,24)`, `(2^17,129)`, +and `(2^16,533)`. Equal-height rows concatenate before hashing. The small +wide-only commitment differences are timing variation: their row kernels +are unchanged. These synthetic measurements establish a useful local win; +they do not establish an Init/Mathlib or Stage 2 wall-time improvement. + +## Validation and reproduction + +- All 70 CUDA-enabled library tests passed (5.46 seconds). +- CPU BLAKE3 comparisons cover 25 row lengths, 1–32768 bytes, and eight row + counts across warp/block boundaries, using seeded random messages. +- Merkle roots/openings cover both sides of the 1024-byte boundary; + existing mixed-height, resident, hybrid, and batch-proof tests passed. +- CUDA memcheck reported zero errors on both kernels, with the new kernel + deliberately launched with two blocks to exercise its grid-stride loop. +- CPU, baseline CUDA, and new CUDA `proof_compatibility` outputs were identical + (17,213 bytes), and each proof verified. SHA-256: + `25564a01d1d352b1ec2de56b019b641d24acc81083133e79274a86829b2a5dd5`. +- Clippy completed without warnings for the library and new commit benchmark. + +```sh +MULTI_STARK_CUDA_ARCHS=120 cargo test --release --locked --features parallel,cuda --lib -- --test-threads=1 +nvcc -O3 --std=c++17 --default-stream=per-thread -arch=sm_120 cuda/blake3_rows_bench.cu -o /tmp/blake3-rows-bench +/tmp/blake3-rows-bench +compute-sanitizer --tool memcheck --error-exitcode 1 /tmp/blake3-rows-bench --check +RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_commit_bench +``` + +Select the native architecture appropriate to the test GPU. To compare full +commitments, build the same `cuda_commit_bench.rs` against the base and this +change, preserve both binaries, and alternate runs on an idle GPU. The +benchmark reports cold iteration zero separately and the exact commitment +beside each timing. Existing `cuda_blake3_bench` includes host transfers and +CPU hashing; it is not a resident-kernel timer. diff --git a/examples/cuda_commit_bench.rs b/examples/cuda_commit_bench.rs new file mode 100644 index 0000000..5b22110 --- /dev/null +++ b/examples/cuda_commit_bench.rs @@ -0,0 +1,90 @@ +//! Complete PCS commitment timings: host upload, LDEs, and mixed-height Merkle tree. +//! Run identical binaries before/after a kernel change and compare commitments. +//! Witness allocation is outside the timer; cold iteration zero is reported separately. +//! +//! cargo run --release --locked --features parallel,cuda --example cuda_commit_bench +//! MULTI_STARK_CUDA_BENCH_ITERATIONS controls the number of warm iterations (default 5). + +use std::hint::black_box; +use std::time::Instant; + +use multi_stark::config::StarkGenericConfig; +use multi_stark::types::{ + Challenger, CommitmentParameters, ExtVal, FriParameters, GoldilocksBlake3Config, Pcs, Val, +}; +use p3_commit::Pcs as PcsTrait; +use p3_field::PrimeCharacteristicRing; +use p3_matrix::dense::RowMajorMatrix; + +fn main() { + let iterations: usize = std::env::var("MULTI_STARK_CUDA_BENCH_ITERATIONS") + .map_or(5, |s| s.parse().expect("invalid iteration count")); + let config = GoldilocksBlake3Config::new( + CommitmentParameters { + log_blowup: 1, + cap_height: 0, + }, + FriParameters { + log_final_poly_len: 0, + max_log_arity: 1, + num_queries: 2, + commit_proof_of_work_bits: 0, + query_proof_of_work_bits: 0, + }, + ); + let pcs = config.pcs(); + // Equal-height matrices are hashed as concatenated rows. The mixed case + // covers both dispatches and lower-height row injections in one commitment. + let shapes: &[(&str, &[(usize, usize)])] = &[ + ("narrow", &[(20, 8)]), + ("medium", &[(18, 40)]), + ("chunk_boundary", &[(18, 128)]), + ("above_boundary", &[(18, 129)]), + ("wide", &[(16, 925)]), + ( + "mixed", + &[(18, 4), (18, 12), (18, 24), (17, 129), (16, 533)], + ), + ]; + println!("shape,iteration,seconds,commitment"); + for &(name, shape) in shapes { + let mut expected = None; + for iteration in 0..=iterations { + let inputs = shape + .iter() + .enumerate() + .map(|(matrix_index, &(log_height, width))| { + let height = 1 << log_height; + let values = (0..height * width) + .map(|index| { + let value = (index as u64) + .wrapping_mul(0x9e37_79b9_7f4a_7c15) + .wrapping_add(matrix_index as u64); + Val::from_u64(value) + }) + .collect(); + let domain = >::natural_domain_for_degree( + pcs, height, + ); + (domain, RowMajorMatrix::new(values, width)) + }) + .collect::>(); + let started = Instant::now(); + let (commitment, data) = >::commit(pcs, inputs); + let seconds = started.elapsed().as_secs_f64(); + if let Some(expected) = &expected { + assert_eq!(&commitment, expected); + } else { + expected = Some(commitment.clone()); + } + let bytes = + bincode::serde::encode_to_vec(&commitment, bincode::config::standard()).unwrap(); + let hex: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect(); + println!("{name},{iteration},{seconds:.9},{hex}"); + // Retain all prover data until after timing: committing must not + // silently include materialization or opening work from a consumer. + black_box(&data); + drop(data); + } + } +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index a59619b..638a7aa 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -3527,36 +3527,41 @@ mod tests { } #[test] - fn blake3_rows_match_cpu_across_chunk_boundaries() { + fn blake3_rows_match_cpu_across_chunk_and_launch_boundaries() { use p3_blake3::Blake3; - for message_bytes in [1usize, 63, 64, 65, 1023, 1024, 1025, 4264, 7400] { - let message_count = 17; - let messages: Vec = (0..message_bytes * message_count) - .map(|index| (index as u64).wrapping_mul(0x9e37_79b9).to_le_bytes()[0]) - .collect(); - let mut digests = vec![0u8; 32 * message_count]; - // SAFETY: the input contains `message_count` fixed-size messages, - // the output has one 32-byte digest per message, and the call is - // synchronous. - let status = unsafe { - multi_stark_cuda_blake3_hash_rows( - 0, - digests.as_mut_ptr(), - messages.as_ptr(), - message_bytes, - message_count, - ) - }; - check_cuda(status, "BLAKE3 row hashing contract"); + let mut rng = SmallRng::seed_from_u64(0xb1a3e3); + for message_bytes in [ + 1usize, 3, 4, 7, 8, 16, 31, 32, 63, 64, 65, 127, 128, 129, 511, 512, 513, 1023, 1024, + 1025, 2048, 3072, 4264, 7400, 32768, + ] { + for message_count in [1, 31, 32, 33, 255, 256, 257, 513] { + let messages: Vec = (0..message_bytes * message_count) + .map(|_| rng.random()) + .collect(); + let mut digests = vec![0u8; 32 * message_count]; + // SAFETY: the input contains `message_count` fixed-size messages, + // the output has one 32-byte digest per message, and the call is + // synchronous. + let status = unsafe { + multi_stark_cuda_blake3_hash_rows( + 0, + digests.as_mut_ptr(), + messages.as_ptr(), + message_bytes, + message_count, + ) + }; + check_cuda(status, "BLAKE3 row hashing contract"); - for (index, message) in messages.chunks_exact(message_bytes).enumerate() { - let expected: [u8; 32] = Blake3.hash_iter(message.iter().copied()); - assert_eq!( - &digests[index * 32..(index + 1) * 32], - &expected, - "message_bytes={message_bytes}, index={index}" - ); + for (index, message) in messages.chunks_exact(message_bytes).enumerate() { + let expected: [u8; 32] = Blake3.hash_iter(message.iter().copied()); + assert_eq!( + &digests[index * 32..(index + 1) * 32], + &expected, + "message_bytes={message_bytes}, count={message_count}, index={index}" + ); + } } } } @@ -3565,7 +3570,14 @@ mod tests { fn blake3_merkle_root_matches_cpu() { use p3_blake3::Blake3; - for (row_bytes, row_count) in [(16usize, 1usize), (64, 8), (4264, 1024)] { + for (row_bytes, row_count) in [ + (16usize, 1usize), + (64, 8), + (1023, 512), + (1024, 512), + (1025, 512), + (4264, 1024), + ] { let rows: Vec = (0..row_bytes * row_count) .map(|index| (index as u64).wrapping_mul(0x517c_c1b7).to_le_bytes()[0]) .collect(); From 38f7cc0a51207e19efed500a85f9f61141b995e6 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:33:42 +0000 Subject: [PATCH 02/34] cuda: fuse wide NTTs at any height and trim LDE memory passes --- cuda/kernels.cu | 23 +++---- docs/cuda-benchmarks.md | 6 ++ docs/cuda-lde-pass-reduction.md | 97 +++++++++++++++++++++++++++++ examples/cuda_resident_lde_bench.rs | 52 ++++++++++++++++ src/cuda/mod.rs | 55 ++++++++++++++++ 5 files changed, 222 insertions(+), 11 deletions(-) create mode 100644 docs/cuda-lde-pass-reduction.md create mode 100644 examples/cuda_resident_lde_bench.rs diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 96d90dc..3e745ae 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -1692,9 +1692,10 @@ cudaError_t launch_dif(uint64_t* values, size_t height, size_t width, } const size_t total = (height >> 1) * width; const unsigned int blocks = blocks_for(total); - // Tall, wide row-major batches amortize the heavier fused kernel. Width-2 - // FRI codewords retain the shared-memory tail specialized below. - if (width >= 8 && height >= (size_t(1) << 18)) { + // Wide row-major batches use fused stages regardless of height: short, + // wide traces can contain as many cells as tall traces. Width-2 FRI + // codewords retain the shared-memory tail specialized below. + if (width >= 8) { size_t half = height >> 1; while (half >= 4) { radix8_dif_stage<<> 3) * width), THREADS>>>( @@ -2337,9 +2338,11 @@ extern "C" int multi_stark_cuda_coset_lde_create( // staging pool; they go through the persistent staging slots instead. // Small ones take the direct pageable path. const uint64_t *device_inverse_twiddles=nullptr,*device_shift_powers=nullptr,*device_forward_twiddles=nullptr; - if (status == cudaSuccess) { - status = cudaMemsetAsync(lde->values, 0, - output_elements * sizeof(uint64_t), + // The source prefix is overwritten by the trace copy. Only the padded + // tail needs zeroing before the forward transform. + if (status == cudaSuccess && output_elements > input_elements) { + status = cudaMemsetAsync(lde->values + input_elements, 0, + (output_elements - input_elements) * sizeof(uint64_t), cudaStreamPerThread); } if (status == cudaSuccess) { @@ -2379,11 +2382,9 @@ extern "C" int multi_stark_cuda_coset_lde_create( if (status == cudaSuccess) { status = launch_dif(lde->values, extended_height, width,device_forward_twiddles); } - if (status == cudaSuccess) { - canonicalize_goldilocks<<>>( - lde->values, output_elements); - status = cudaGetLastError(); - } + // Normalization and every DIF butterfly produce canonical field values, + // including the height-one case. A final reduction pass is redundant; + // raw-representation tests protect the Merkle byte contract. if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status != cudaSuccess) { destroy_resident_lde(lde); diff --git a/docs/cuda-benchmarks.md b/docs/cuda-benchmarks.md index ed85386..1af63ef 100644 --- a/docs/cuda-benchmarks.md +++ b/docs/cuda-benchmarks.md @@ -39,3 +39,9 @@ cargo run --release --locked --features parallel,cuda --example cuda_blake3_benc The CSV DFT benchmark labels shapes below the production CUDA thresholds as `cpu-fallback`, warms both implementations, uses the same iteration count, and checks each output against the CPU reference outside the timing window. + +## 2026-09-11: BLAKE3 and LDE kernel experiments + +See [short-row BLAKE3 dispatch](cuda-blake3-short-rows.md) and +[resident LDE/wide NTT optimization](cuda-lde-pass-reduction.md) for isolated +before/after measurements, exact proof compatibility, and reproduction steps. diff --git a/docs/cuda-lde-pass-reduction.md b/docs/cuda-lde-pass-reduction.md new file mode 100644 index 0000000..3390f39 --- /dev/null +++ b/docs/cuda-lde-pass-reduction.md @@ -0,0 +1,97 @@ +# Resident LDE and wide NTT experiment: 2026-09-11 + +Base: `ff3237c` (the short-row BLAKE3 change). Same RTX PRO 6000 Blackwell +Server Edition, CUDA 13.3 / native `sm_120`, and eight Rayon threads as the +[BLAKE3 experiment](cuda-blake3-short-rows.md). Timings below measure the +incremental effect of this change; BLAKE3 is enabled in both binaries. + +## Change retained + +- Use the existing fused radix-8/radix-4 stages for matrices with at least + eight columns regardless of height. Previously a separate `height >= 2^18` + condition left short, wide matrices on individual radix-2 stages even when + their total cell count was large. The existing narrow-column path remains. +- Zero only the padded tail of the resident LDE allocation. The input prefix + is overwritten by the trace copy and does not need a preceding clear. +- Remove the resident LDE's final canonicalization pass. Its normalization + and butterfly arithmetic already produce canonical Goldilocks values, + including height-one transforms. Tests inspect the actual stored u64 + representation, rather than an accessor that canonicalizes values. + +No new NTT implementation, matrix layout, workspace allocation, dependency, +or machine-specific threshold was introduced. Trace retention and output +storage order are unchanged. The NTT dispatch change applies to other callers +of the same transform helper; memory-pass removal is scoped to +`multi_stark_cuda_coset_lde_create`. + +An earlier trial fused bit reversal, coset scaling, virtual zero padding, +and the first forward stage in place, using bit-reversal pairs to avoid +races. It passed correctness checks but did not improve complete commitments +consistently, so that kernel and its transform-continuation helper were +removed. The simpler change above is the measured result. + +## Measurements + +Resident LDE timings include upload, inverse transform, coset handling, +forward transform, and completion, excluding construction of the host input. +Medians of 14 warm samples per variant across two processes; first use of +each shape is reported separately and excluded. Baseline and candidate +processes were interleaved on an idle GPU. + +| Input height | Columns | Blowup | BLAKE3-only | + LDE change | Reduction | +| ---: | ---: | ---: | ---: | ---: | ---: | +| 2^20 | 1 | 2 | 0.992 ms | 0.965 ms | 2.7% | +| 2^20 | 2 | 2 | 1.805 ms | 1.833 ms | -1.6% | +| 2^20 | 8 | 2 | 6.563 ms | 6.435 ms | 1.9% | +| 2^18 | 40 | 2 | 8.582 ms | 8.451 ms | 1.5% | +| 2^18 | 128 | 2 | 30.389 ms | 30.093 ms | 1.0% | +| 2^18 | 129 | 2 | 30.533 ms | 30.099 ms | 1.4% | +| 2^16 | 925 | 2 | 79.048 ms | 57.660 ms | 27.1% | +| 2^18 | 40 | 4 | 10.664 ms | 10.155 ms | 4.8% | + +Complete PCS commitments include LDEs, gathering, hashing, Merkle construction, +and root return. Medians of 15 warm samples per variant across three processes; +all commitment bytes matched for every shape/run. + +| Shape | BLAKE3-only | + LDE change | Reduction | +| --- | ---: | ---: | ---: | +| 2^20 × 8 | 9.011 ms | 8.938 ms | 0.8% | +| 2^18 × 40 | 10.666 ms | 10.478 ms | 1.8% | +| 2^18 × 128 | 36.102 ms | 35.332 ms | 2.1% | +| 2^18 × 129 | 43.792 ms | 43.195 ms | 1.4% | +| 2^16 × 925 | 91.071 ms | 69.281 ms | 23.9% | +| Mixed heights/widths | 68.184 ms | 51.253 ms | 24.8% | + +The mixed shape is the same five-matrix fixture as the BLAKE3 experiment. +The substantial gains occur in shapes affected by the removed NTT height +cutoff. Small percentage differences on unaffected/narrow shapes should not +be treated as established wins; one narrow LDE case was 1.6% slower. These +are synthetic component timings, not an Init/Mathlib or Stage 2 wall-time +measurement, and are not a general 2× NTT claim. + +## Validation and reproduction + +- All 71 CUDA-enabled library tests passed (6.46 seconds). +- CPU comparisons exercise blowups 1/2/4/8, three cosets, height-one inputs, + wide/narrow dispatch, fused-stage residues modulo three, and noncanonical + input representations around the field modulus. +- Raw output field words remain below the modulus. Mixed-height commitments, + resident/hybrid openings, hashing, and batch-proof tests pass. +- CUDA memcheck and initcheck report zero errors for the padding/coset/raw- + representation suite, including the case with no zero-padded tail. +- The complete 17,213-byte compatibility proof verifies and matches both + CPU and BLAKE3-only output exactly. SHA-256: + `25564a01d1d352b1ec2de56b019b641d24acc81083133e79274a86829b2a5dd5`. +- Clippy completed without warnings for the library and both new examples. + +```sh +MULTI_STARK_CUDA_ARCHS=120 RAYON_NUM_THREADS=8 cargo test --release --locked --features parallel,cuda --lib -- --test-threads=1 +RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_resident_lde_bench +RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_commit_bench +``` + +Build the identical benchmark source against `ff3237c` and this commit, +preserve both binaries, and alternate runs. `MULTI_STARK_CUDA_BENCH_ITERATIONS` +controls the warm sample count. The resident benchmark uses the production +resident-LDE interface; the older `cuda_dft_bench` exercises a different, +transfer-inclusive host-return interface and does not isolate these changes. diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs new file mode 100644 index 0000000..84bedfd --- /dev/null +++ b/examples/cuda_resident_lde_bench.rs @@ -0,0 +1,52 @@ +//! Resident LDE timings including upload, excluding host input construction. +//! Run the identical benchmark against both revisions. Correctness is covered +//! by resident_coset_lde_matches_cpu_storage and proof_compatibility. + +#[cfg(not(feature = "cuda"))] +fn main() { + eprintln!("enable --features cuda"); +} + +#[cfg(feature = "cuda")] +fn main() { + use multi_stark::cuda::{CudaDft, pcs::CudaPcsDft}; + use p3_field::{Field, PrimeCharacteristicRing}; + use p3_goldilocks::Goldilocks; + use p3_matrix::dense::RowMajorMatrix; + use std::hint::black_box; + use std::time::Instant; + + let iterations: usize = std::env::var("MULTI_STARK_CUDA_BENCH_ITERATIONS") + .map_or(7, |s| s.parse().expect("invalid iteration count")); + let gpu = CudaDft::default(); + println!("log_height,width,added_bits,iteration,seconds"); + for (log_height, width, added_bits) in [ + (20, 1, 1), + (20, 2, 1), + (20, 8, 1), + (18, 40, 1), + (18, 128, 1), + (18, 129, 1), + (16, 925, 1), + (18, 40, 2), + ] { + let height = 1 << log_height; + let values = (0..height * width) + .map(|index| Goldilocks::from_u64((index as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15))) + .collect(); + let matrix = RowMajorMatrix::new(values, width); + for iteration in 0..=iterations { + let started = Instant::now(); + let lde = CudaPcsDft::coset_lde_batch_resident( + &gpu, + &matrix, + added_bits, + Goldilocks::GENERATOR, + ); + let seconds = started.elapsed().as_secs_f64(); + println!("{log_height},{width},{added_bits},{iteration},{seconds:.9}"); + black_box(&lde); + drop(lde); + } + } +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 638a7aa..db2e8a5 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -3725,6 +3725,55 @@ mod tests { } } + #[test] + fn resident_coset_lde_padding_and_raw_representatives() { + let cpu = Radix2DitParallel::::default(); + let gpu = CudaDft::default(); + let p = Goldilocks::ORDER_U64; + let representatives = [0, 1, p - 1, p, p + 1, u64::MAX]; + // Exercise absent padding, one padded half, multiple padded halves, + // and height-one transforms with lazy input field representatives. + for log_height in [0usize, 1, 2, 5, 8] { + for width in [1usize, 2, 3, 8, 17] { + let height = 1 << log_height; + let matrix = RowMajorMatrix::new( + (0..height * width) + .map(|i| Goldilocks::new(representatives[i % representatives.len()])) + .collect(), + width, + ); + for added_bits in [0usize, 1, 2, 3] { + for shift in [ + Goldilocks::ONE, + Goldilocks::GENERATOR, + Goldilocks::from_u64(11), + ] { + let expected = cpu + .coset_lde_batch(matrix.clone(), added_bits, shift) + .bit_reverse_rows() + .to_row_major_matrix(); + let actual = gpu + .coset_lde_batch_resident(&matrix, added_bits, shift) + .to_row_major_matrix(); + assert_eq!( + actual, expected, + "height=2^{log_height}, width={width}, added_bits={added_bits}, shift={shift}" + ); + assert!( + actual.values.iter().all(|&value| { + // SAFETY: Goldilocks is repr(transparent) over u64. + // Inspect storage without canonicalizing via a field accessor. + let raw = unsafe { core::mem::transmute::(value) }; + raw < p + }), + "LDE digest input must contain canonical field bytes" + ); + } + } + } + } + } + #[test] fn resident_coset_lde_matches_cpu_storage() { let mut rng = SmallRng::seed_from_u64(0x51de); @@ -3736,6 +3785,12 @@ mod tests { (12, 2, 1), (14, 2, 2), (16, 1, 1), + // Fused radix-8 dispatch and all stage-count residues mod 3. + (17, 8, 1), + (18, 8, 1), + (18, 8, 2), + (19, 8, 1), + (18, 3, 1), ] { let height = 1 << log_height; let matrix = From 3c87a593fc29cee53c1836eeb44073270ed026df Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:15:04 +0000 Subject: [PATCH 03/34] cuda: per-device prover binding, host pools per device, GPU trace witnesses Bind a GoldilocksBlake3Config to a CUDA device (GoldilocksBlake3Config::with_device, CudaMmcs::with_device, CudaDft::new taking the device) instead of reading MULTI_STARK_CUDA_DEVICE once per process, so one process can drive one prover per GPU. The host-side pools that were process-wide (the pinned host pool, the upload staging slots, the resident-LDE control-block slab) are now one per device, indexed by the calling thread's current device, which every kernel entry point sets first. Add the witness module (TraceSource, PreparedWitness, TraceGenerator) and the device trace view so a circuit can generate its trace on the GPU, with the lookup-graph LDE writing device rows directly. build.rs: compile the kernels' host code with the POSIX and default feature sets instead of _GNU_SOURCE, so glibc 2.38+ keeps the plain strto*/clock_gettime names; the ix binary is linked by Lean's bundled clang against an older sysroot that lacks the __isoc23_* variants. Claude-Session: https://claude.ai/code/session_01WD1WGwm74eqamotwg34pZt --- build.rs | 8 ++ cuda/kernels.cu | 167 +++++++++++++++++++++------- src/batch.rs | 74 +++++++++++-- src/config.rs | 36 ++++++ src/cuda/mmcs.rs | 22 ++-- src/cuda/mod.rs | 172 +++++++++++++++++++++++++++++ src/cuda/pcs.rs | 2 +- src/cuda/witness.rs | 262 ++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/prover.rs | 16 ++- src/types.rs | 67 ++++++++++- src/witness.rs | 85 ++++++++++++++ 12 files changed, 852 insertions(+), 60 deletions(-) create mode 100644 src/cuda/witness.rs create mode 100644 src/witness.rs diff --git a/build.rs b/build.rs index d3641c7..438365e 100644 --- a/build.rs +++ b/build.rs @@ -45,6 +45,14 @@ fn main() { .arg("-O3") .arg("-lineinfo") .arg("--compiler-options=-fPIC") + // Host code is linked by whatever toolchain links the final binary, + // and Lean's bundled clang links against a sysroot older than + // glibc 2.38. g++ predefines _GNU_SOURCE, under which glibc 2.38+ + // renames strtol and friends to their C23 variants (__isoc23_*), + // which that sysroot lacks. The POSIX and default feature sets keep + // everything the kernels' host code uses (clock_gettime, pthreads) + // under the plain names. + .arg("--compiler-options=-U_GNU_SOURCE,-D_DEFAULT_SOURCE,-D_POSIX_C_SOURCE=200809L") .arg("-o") .arg(&library) .arg("cuda/kernels.cu"); diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 3e745ae..92353e3 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -37,6 +37,18 @@ inline cudaError_t stream_malloc(void** pointer,size_t bytes){return cudaMallocA inline cudaError_t stream_free(void* pointer){return pointer?cudaFreeAsync(pointer,cudaStreamPerThread):cudaSuccess;} #define cudaMalloc stream_malloc #define cudaFree stream_free +// Host-side pools are kept per device, so provers on different devices in +// one process never lease each other's buffers or control blocks. Every +// entry point selects its device first, so the pools index by the calling +// thread's current device. +constexpr int MAX_CUDA_DEVICES = 64; +inline int current_device_index() { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess || device < 0 || device >= MAX_CUDA_DEVICES) { + return 0; + } + return device; +} struct ConstantCacheEntry{int device=-1;size_t count=0;uint64_t kind=0,key0=0,key1=0;uint64_t* values=nullptr;ConstantCacheEntry* next=nullptr;}; ConstantCacheEntry* constant_cache=nullptr;volatile int constant_cache_lock=0; @@ -386,7 +398,7 @@ struct PinnedHostPoolEntry { }; constexpr size_t PINNED_HOST_POOL_SIZE = 2; -PinnedHostPoolEntry pinned_host_pool[PINNED_HOST_POOL_SIZE]; +PinnedHostPoolEntry pinned_host_pool[MAX_CUDA_DEVICES][PINNED_HOST_POOL_SIZE]; pthread_mutex_t pinned_host_pool_mutex = PTHREAD_MUTEX_INITIALIZER; class PinnedHostBuffer { @@ -398,7 +410,7 @@ class PinnedHostBuffer { ~PinnedHostBuffer() { if (pool_index_ < PINNED_HOST_POOL_SIZE) { pthread_mutex_lock(&pinned_host_pool_mutex); - pinned_host_pool[pool_index_].in_use = false; + pinned_host_pool[device_][pool_index_].in_use = false; pthread_mutex_unlock(&pinned_host_pool_mutex); } else if (pointer_ != nullptr) { cudaFreeHost(pointer_); @@ -409,19 +421,21 @@ class PinnedHostBuffer { if (bytes == 0 || pointer_ != nullptr) { return cudaErrorInvalidValue; } + device_ = current_device_index(); + PinnedHostPoolEntry* pool = pinned_host_pool[device_]; pthread_mutex_lock(&pinned_host_pool_mutex); size_t reusable = PINNED_HOST_POOL_SIZE; for (size_t i = 0; i < PINNED_HOST_POOL_SIZE; ++i) { - const auto& entry = pinned_host_pool[i]; + const auto& entry = pool[i]; if (!entry.in_use && entry.pointer != nullptr && entry.capacity >= bytes && (reusable == PINNED_HOST_POOL_SIZE || - entry.capacity < pinned_host_pool[reusable].capacity)) { + entry.capacity < pool[reusable].capacity)) { reusable = i; } } if (reusable < PINNED_HOST_POOL_SIZE) { - auto& entry = pinned_host_pool[reusable]; + auto& entry = pool[reusable]; entry.in_use = true; pointer_ = entry.pointer; pool_index_ = reusable; @@ -431,9 +445,9 @@ class PinnedHostBuffer { size_t available = PINNED_HOST_POOL_SIZE; for (size_t i = 0; i < PINNED_HOST_POOL_SIZE; ++i) { - if (!pinned_host_pool[i].in_use && + if (!pool[i].in_use && (available == PINNED_HOST_POOL_SIZE || - pinned_host_pool[i].capacity < pinned_host_pool[available].capacity)) { + pool[i].capacity < pool[available].capacity)) { available = i; } } @@ -442,7 +456,7 @@ class PinnedHostBuffer { return cudaMallocHost(reinterpret_cast(&pointer_), bytes); } - auto& entry = pinned_host_pool[available]; + auto& entry = pool[available]; entry.in_use = true; uint64_t* old_pointer = entry.pointer; const size_t old_capacity = entry.capacity; @@ -475,6 +489,7 @@ class PinnedHostBuffer { private: uint64_t* pointer_ = nullptr; size_t pool_index_ = PINNED_HOST_POOL_SIZE; + int device_ = 0; }; struct ResidentMerkleTree { @@ -512,19 +527,22 @@ struct ResidentMixedMerkleTree { // count (the PCIe link is shared anyway). constexpr size_t UPLOAD_STAGING_SLOTS = 4; constexpr size_t UPLOAD_STAGING_BYTES = size_t(64) << 20; -uint64_t* upload_staging[UPLOAD_STAGING_SLOTS] = {nullptr}; -bool upload_staging_in_use[UPLOAD_STAGING_SLOTS] = {false}; +uint64_t* upload_staging[MAX_CUDA_DEVICES][UPLOAD_STAGING_SLOTS] = {{nullptr}}; +bool upload_staging_in_use[MAX_CUDA_DEVICES][UPLOAD_STAGING_SLOTS] = {{false}}; pthread_mutex_t upload_staging_mutex = PTHREAD_MUTEX_INITIALIZER; pthread_cond_t upload_staging_cond = PTHREAD_COND_INITIALIZER; // Copies `bytes` of pageable host memory to the device on the per-thread // stream through a leased staging slot; returns with the copy complete. cudaError_t staged_upload(void* device, const void* host, size_t bytes) { + const int device_index = current_device_index(); + uint64_t** slots = upload_staging[device_index]; + bool* in_use = upload_staging_in_use[device_index]; pthread_mutex_lock(&upload_staging_mutex); size_t slot = UPLOAD_STAGING_SLOTS; while (slot == UPLOAD_STAGING_SLOTS) { for (size_t i = 0; i < UPLOAD_STAGING_SLOTS; ++i) { - if (!upload_staging_in_use[i]) { + if (!in_use[i]) { slot = i; break; } @@ -533,15 +551,15 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { pthread_cond_wait(&upload_staging_cond, &upload_staging_mutex); } } - upload_staging_in_use[slot] = true; + in_use[slot] = true; cudaError_t status = cudaSuccess; - if (upload_staging[slot] == nullptr) { - status = cudaMallocHost(reinterpret_cast(&upload_staging[slot]), + if (slots[slot] == nullptr) { + status = cudaMallocHost(reinterpret_cast(&slots[slot]), UPLOAD_STAGING_BYTES); - if (status != cudaSuccess) upload_staging[slot] = nullptr; + if (status != cudaSuccess) slots[slot] = nullptr; } pthread_mutex_unlock(&upload_staging_mutex); - uint64_t* staging = upload_staging[slot]; + uint64_t* staging = slots[slot]; if (status == cudaSuccess) { const auto* source = static_cast(host); auto* target = static_cast(device); @@ -555,12 +573,15 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { } } pthread_mutex_lock(&upload_staging_mutex); - upload_staging_in_use[slot] = false; - pthread_cond_signal(&upload_staging_cond); + in_use[slot] = false; + pthread_cond_broadcast(&upload_staging_cond); pthread_mutex_unlock(&upload_staging_mutex); return status; } +using TraceWriter = int (*)(void*, int, uint64_t*, size_t, size_t); +using TraceDestroy = void (*)(void*); + struct ResidentLde { uint64_t* values = nullptr; // Original row-major evaluations, retained for prover stages (notably @@ -570,12 +591,19 @@ struct ResidentLde { const uint64_t* host_trace_values = nullptr; bool host_trace_registered = false; size_t trace_height = 0; + void* trace_context = nullptr; + TraceWriter trace_writer = nullptr; + TraceDestroy trace_destroy = nullptr; size_t height = 0; size_t width = 0; uint8_t* interpolation_scratch = nullptr; size_t interpolation_scratch_bytes = 0; + // The device whose slab this control block came from, so it returns + // there when destroyed from a thread on another device. + int device = 0; ~ResidentLde() { + if (trace_destroy) trace_destroy(trace_context); if (values != nullptr) { cudaFree(values); } @@ -592,14 +620,17 @@ struct ResidentLde { // whole device, so control blocks come from slabs allocated once per process // and are recycled through a free list rather than allocated per LDE. constexpr size_t RESIDENT_LDE_SLAB = 256; -void* resident_lde_free = nullptr; +// One free list per device: a slab is allocated with its device current, so +// its pages prefer that device and are never read by kernels on another. +void* resident_lde_free[MAX_CUDA_DEVICES] = {nullptr}; volatile int resident_lde_lock = 0; cudaError_t create_resident_lde(ResidentLde** output) { if (output == nullptr) return cudaErrorInvalidValue; *output = nullptr; + const int device = current_device_index(); while (__sync_lock_test_and_set(&resident_lde_lock, 1)) {} - if (resident_lde_free == nullptr) { + if (resident_lde_free[device] == nullptr) { void* slab = nullptr; const cudaError_t status = cudaMallocManaged(&slab, RESIDENT_LDE_SLAB * sizeof(ResidentLde)); @@ -610,23 +641,25 @@ cudaError_t create_resident_lde(ResidentLde** output) { auto* bytes = static_cast(slab); for (size_t i = 0; i < RESIDENT_LDE_SLAB; ++i) { void* slot = bytes + i * sizeof(ResidentLde); - *static_cast(slot) = resident_lde_free; - resident_lde_free = slot; + *static_cast(slot) = resident_lde_free[device]; + resident_lde_free[device] = slot; } } - void* slot = resident_lde_free; - resident_lde_free = *static_cast(slot); + void* slot = resident_lde_free[device]; + resident_lde_free[device] = *static_cast(slot); __sync_lock_release(&resident_lde_lock); *output = new (slot) ResidentLde; + (*output)->device = device; return cudaSuccess; } cudaError_t destroy_resident_lde(ResidentLde* lde) { if (lde == nullptr) return cudaSuccess; + const int device = lde->device; lde->~ResidentLde(); while (__sync_lock_test_and_set(&resident_lde_lock, 1)) {} - *reinterpret_cast(lde) = resident_lde_free; - resident_lde_free = lde; + *reinterpret_cast(lde) = resident_lde_free[device]; + resident_lde_free[device] = lde; __sync_lock_release(&resident_lde_lock); return cudaSuccess; } @@ -2297,12 +2330,12 @@ extern "C" int multi_stark_cuda_coset_lde_batch( return static_cast(status); } -extern "C" int multi_stark_cuda_coset_lde_create( +static int coset_lde_create( int device_id, void** handle, const uint64_t* input, size_t height, size_t width, size_t added_bits, const uint64_t* inverse_twiddles, const uint64_t* shift_powers, const uint64_t* forward_twiddles, - uint64_t height_inverse) { - if (handle == nullptr || input == nullptr || (height > 1 && inverse_twiddles == nullptr) || + uint64_t height_inverse, void* context, TraceWriter writer) { + if (handle == nullptr || (input == nullptr && writer == nullptr) || (height > 1 && inverse_twiddles == nullptr) || shift_powers == nullptr || forward_twiddles == nullptr || !is_power_of_two(height) || width == 0 || added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { @@ -2321,6 +2354,8 @@ extern "C" int multi_stark_cuda_coset_lde_create( ResidentLde* lde = nullptr; status = create_resident_lde(&lde); if (status != cudaSuccess) return static_cast(status); + lde->trace_context = context; + lde->trace_writer = writer; lde->height = extended_height; lde->width = width; const size_t input_elements = height * width; @@ -2346,7 +2381,14 @@ extern "C" int multi_stark_cuda_coset_lde_create( cudaStreamPerThread); } if (status == cudaSuccess) { - if (input_bytes >= (size_t(8) << 20)) { + if (writer) { + constexpr size_t TILE_ROWS = size_t(1) << 16; + for (size_t first = 0; status == cudaSuccess && first < height; first += TILE_ROWS) { + const size_t rows = height - first < TILE_ROWS ? height - first : TILE_ROWS; + status = static_cast(writer(context, device_id, + lde->trace_values + first * width, first, rows)); + } + } else if (input_bytes >= (size_t(8) << 20)) { status = staged_upload(lde->trace_values, input, input_bytes); } else { status = cudaMemcpyAsync(lde->trace_values, input, @@ -2394,6 +2436,33 @@ extern "C" int multi_stark_cuda_coset_lde_create( return static_cast(cudaSuccess); } +extern "C" int multi_stark_cuda_coset_lde_create( + int device_id, void** handle, const uint64_t* input, size_t height, + size_t width, size_t added_bits, const uint64_t* inverse_twiddles, + const uint64_t* shift_powers, const uint64_t* forward_twiddles, + uint64_t height_inverse) { + return coset_lde_create(device_id, handle, input, height, width, added_bits, + inverse_twiddles, shift_powers, forward_twiddles, height_inverse, nullptr, nullptr); +} + +// Context ownership transfers only on success. Its lifetime covers trace +// release and LDE eviction because lookup recovery still needs original rows. +extern "C" int multi_stark_cuda_coset_lde_generate( + int device_id, void** handle, size_t height, size_t width, size_t added_bits, + const uint64_t* inverse_twiddles, const uint64_t* shift_powers, + const uint64_t* forward_twiddles, uint64_t height_inverse, + void* context, TraceWriter writer, TraceDestroy destroy) { + if (!context || !writer || !destroy) return static_cast(cudaErrorInvalidValue); + const int status = coset_lde_create(device_id, handle, nullptr, height, width, added_bits, + inverse_twiddles, shift_powers, forward_twiddles, height_inverse, context, writer); + if (status == 0) static_cast(*handle)->trace_destroy = destroy; + return status; +} + +extern "C" bool multi_stark_cuda_lde_has_generator(const void* handle) { + return handle && static_cast(handle)->trace_writer; +} + extern "C" int multi_stark_cuda_prepare_lde_constants( int device_id,const uint64_t* inverse_twiddles,size_t inverse_count, const uint64_t* shift_powers,size_t height,const uint64_t* forward_twiddles, @@ -3271,7 +3340,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han auto* prep=static_cast(preprocessed_handle); if(!output_handle||!total||!nodes||!node_count||!slot_count||!lookups||!lookup_count|| (lookup_arg_count&&!lookup_args)||!main|| - (!main->trace_values&&!main->host_trace_values)||!main->trace_height|| + (!main->trace_values&&!main->host_trace_values&&!main->trace_writer)||!main->trace_height|| !group_size||!beta||!gamma||!inverse_twiddles||!shift_powers||!forward_twiddles|| !is_power_of_two(main->trace_height)||added_bits>=sizeof(size_t)*8|| main->trace_height>(SIZE_MAX>>added_bits)||(prep&&!prep->trace_values)) { @@ -3280,7 +3349,13 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han *output_handle=nullptr;const size_t height=main->trace_height; const size_t groups=(lookup_count+group_size-1)/group_size,width=2*groups; const size_t count=height*groups,extended_height=height<trace_writer ? getenv("MULTI_STARK_CUDA_LOOKUP_TRACE_TILE_ROWS") : nullptr) { + char* end=nullptr; + const unsigned long rows=strtoul(configured,&end,10); + if (end!=configured && *end=='\0' && rows>0 && rows<=LOOKUP_ROWS_PER_CHUNK) + LOOKUP_ROWS_PER_CHUNK=rows; + } const size_t chunk_rows=height(cudaErrorInvalidValue); @@ -3306,14 +3381,19 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han const size_t budget=48*1024;size_t tile=budget/(slot_count*sizeof(uint64_t));const bool global=tile<32; if(global)tile=128;else if(tile>32)tile=32;size_t blocks=(chunk_rows+tile-1)/tile;const size_t cap=global?256:1024;if(blocks>cap)blocks=cap; if(status==cudaSuccess&&global)status=cudaMalloc(reinterpret_cast(&scratch),blocks*slot_count*tile*sizeof(uint64_t)); - if(status==cudaSuccess&&main->host_trace_values)status=cudaMalloc(reinterpret_cast(&trace_chunk),(chunk_rows+1)*main->width*sizeof(uint64_t)); + const bool generated = !main->trace_values && main->trace_writer; + const bool tiled = main->host_trace_values || generated; + if(status==cudaSuccess&&tiled)status=cudaMalloc(reinterpret_cast(&trace_chunk),(chunk_rows+1)*main->width*sizeof(uint64_t)); for(size_t row_start=0;status==cudaSuccess&&row_starttrace_values; - if(main->host_trace_values){ + if(generated){ + status=static_cast(main->trace_writer(main->trace_context,device_id,trace_chunk,row_start,rows+1)); + active_trace=trace_chunk; + } else if(main->host_trace_values){ status=cudaMemcpy(trace_chunk,main->host_trace_values+row_start*main->width,rows*main->width*sizeof(uint64_t),cudaMemcpyHostToDevice); const size_t next_row=(row_start+rows)&(height-1); if(status==cudaSuccess)status=cudaMemcpy(trace_chunk+rows*main->width,main->host_trace_values+next_row*main->width,main->width*sizeof(uint64_t),cudaMemcpyHostToDevice); @@ -3321,7 +3401,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han } if(status==cudaSuccess){lookup_messages_graph<<(chunk_blocks),static_cast(tile),global?0:slot_count*tile*sizeof(uint64_t)>>>( conjugates,norms,multiplicities,dn,node_count,slot_count,dl,lookup_count,da,prep,active_trace,main->width, - main->host_trace_values!=nullptr,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w,height,row_start,rows,scratch);status=cudaGetLastError();} + tiled,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w,height,row_start,rows,scratch);status=cudaGetLastError();} if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,messages); if(status==cudaSuccess){lookup_group_deltas_batched<<>>( deltas+row_start*groups,multiplicities,conjugates,norm_inverses,rows,lookup_count,group_size,ext_w);status=cudaGetLastError();} @@ -3649,7 +3729,7 @@ extern "C" int multi_stark_cuda_lde_release_trace(int device_id,void* handle){ if(!handle)return static_cast(cudaSuccess);cudaError_t status=cudaSetDevice(device_id); auto* lde=static_cast(handle);if(status==cudaSuccess&&lde->trace_values){status=cudaFree(lde->trace_values);lde->trace_values=nullptr;} if(status==cudaSuccess&&lde->host_trace_registered){status=cudaHostUnregister(const_cast(lde->host_trace_values));lde->host_trace_registered=false;} - if(status==cudaSuccess){lde->host_trace_values=nullptr;lde->trace_height=0;} + if(status==cudaSuccess){lde->host_trace_values=nullptr;if(!lde->trace_writer)lde->trace_height=0;} return static_cast(status); } @@ -3682,7 +3762,7 @@ extern "C" int multi_stark_cuda_lde_release_values(int device_id, void* handle) } if (status == cudaSuccess) { lde->host_trace_values = nullptr; - lde->trace_height = 0; + if (!lde->trace_writer) lde->trace_height = 0; } return static_cast(status); } @@ -4479,3 +4559,16 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, } return static_cast(status); } + +extern "C" int multi_stark_cuda_generate_trace_rows(int device, void* context, + TraceWriter writer, size_t first, size_t rows, size_t width, uint64_t* output) { + if (!context || !writer || !output || !rows || !width || !product_fits(rows,width)) + return static_cast(cudaErrorInvalidValue); + cudaError_t status = cudaSetDevice(device); + uint64_t* tile = nullptr; + if (status == cudaSuccess) status = cudaMalloc(reinterpret_cast(&tile), rows * width * sizeof(uint64_t)); + if (status == cudaSuccess) status = static_cast(writer(context, device, tile, first, rows)); + if (status == cudaSuccess) status = cudaMemcpy(output, tile, rows * width * sizeof(uint64_t), cudaMemcpyDeviceToHost); + if (tile) cudaFree(tile); + return static_cast(status); +} diff --git a/src/batch.rs b/src/batch.rs index abbefa4..b1c29db 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -127,6 +127,9 @@ pub enum Retention { /// plus whatever the witness source holds, at the price of one extra /// witness build and stage-1 commitment per shard. Regenerate, + /// Retain only tree nodes within this per-batch device budget. Every + /// witness and LDE is regenerated; pressure evicts trees before proving. + MerkleTrees { max_bytes: usize }, } /// A public lookup message the verifier contributes to the batch's balance, @@ -199,6 +202,7 @@ pub struct BatchBarrier { claims: Vec>>>, headers: Vec>, retained: Vec>>, + trees: Vec>, } impl BatchBarrier { @@ -206,6 +210,16 @@ impl BatchBarrier { pub fn shards(&self) -> usize { self.headers.len() } + + /// Evict optional trees when a resident session needs their memory. + pub fn trim_tree_cache(&mut self, max_bytes: usize) { + evict_trees(&mut self.trees, max_bytes); + } + + /// Tree nodes and declared metadata retained across this batch's barrier. + pub fn tree_cache_bytes(&self) -> usize { + self.trees.iter().flatten().map(|tree| tree.bytes()).sum() + } } /// A vector of shard proofs under one preamble. @@ -478,7 +492,7 @@ where /// # Panics /// Panics if `claims` is empty, if any shard's traces are all empty, or /// if a regenerated shard does not reproduce its round-one header. - pub fn prove_batch_with( + pub fn prove_batch_with( &self, key: &ProverKey, claims: &[Vec>>], @@ -488,7 +502,8 @@ where ) -> BatchProof where Com: PartialEq, - W: FnMut(usize) -> SystemWitness>, + W: FnMut(usize) -> T, + T: Into>> + Send, Self: Sync, ProverKey: Sync, BatchPreamble: Sync, @@ -518,9 +533,10 @@ where /// the next shard's witness on the calling thread; at most one witness /// is held ahead of the prover. #[tracing::instrument(level = "info", skip_all, name = "stark/batch_round_one")] - pub fn batch_round_one(&self, round_one: I, retention: Retention) -> BatchBarrier + pub fn batch_round_one(&self, round_one: I, retention: Retention) -> BatchBarrier where - I: IntoIterator>>, SystemWitness>)>, + I: IntoIterator>>, T)>, + T: Into>> + Send, Self: Sync, SystemWitness>: Send, Stage1: Send, @@ -530,17 +546,40 @@ where claims: Vec::new(), headers: Vec::new(), retained: Vec::new(), + trees: Vec::new(), }; consume_ahead(round_one.into_iter(), |shard, (claims, witness)| { let _g = tracing::info_span!("stark/batch_round_1", shard).entered(); + let witness = witness.into(); + evict_trees( + &mut barrier.trees, + self.config.tree_cache_headroom(&witness), + ); let stage_1 = self.prove_stage_1(witness); barrier .headers .push(stage_1.header(&claim_slices::(&claims))); + let mut tree = None; barrier.retained.push(match retention { Retention::Retain => Some(stage_1), Retention::Regenerate => None, + Retention::MerkleTrees { max_bytes } => { + let used = barrier + .trees + .iter() + .flatten() + .map(|t| t.bytes()) + .sum::(); + let remaining = max_bytes.saturating_sub(used); + if remaining != 0 { + tree = self + .config + .checkpoint_main(stage_1.stage_1_trace_data, remaining); + } + None + } }); + barrier.trees.push(tree); barrier.claims.push(claims); }); barrier @@ -559,7 +598,7 @@ where /// on the calling thread for every shard round one did not retain; at /// most one witness is held ahead of the prover. #[tracing::instrument(level = "info", skip_all, name = "stark/batch_round_two")] - pub fn batch_round_two( + pub fn batch_round_two( &self, key: &ProverKey, barrier: BatchBarrier, @@ -568,7 +607,8 @@ where ) -> BatchProof where Com: PartialEq, - W: FnMut(usize) -> SystemWitness>, + W: FnMut(usize) -> T, + T: Into>> + Send, Self: Sync, ProverKey: Sync, BatchPreamble: Sync, @@ -581,6 +621,7 @@ where claims, headers, retained, + mut trees, } = barrier; assert!(!claims.is_empty(), "cannot prove an empty batch"); let preamble = BatchPreamble { headers, messages }; @@ -595,7 +636,9 @@ where let stage_1 = match (witness, retained) { (None, Some(stage_1)) => stage_1, (Some(witness), None) => { - let stage_1 = self.prove_stage_1(witness); + let witness = witness.into(); + evict_trees(&mut trees, self.config.tree_cache_headroom(&witness)); + let stage_1 = self.prove_stage_1_restored(witness, trees[shard].take()); assert!( stage_1.header(&claims) == preamble.headers[shard], "shard {shard} did not reproduce its round-one header" @@ -970,3 +1013,20 @@ mod tests { type ExtVal = ::Challenge; } + +fn evict_trees(trees: &mut [Option], headroom: usize) { + let mut used = trees + .iter() + .flatten() + .map(|tree| tree.bytes()) + .sum::(); + for entry in trees.iter_mut().rev() { + if used <= headroom { + break; + } + if let Some(tree) = entry.take() { + used -= tree.bytes(); + tracing::debug!(bytes = tree.bytes(), "evicted stage-one tree checkpoint"); + } + } +} diff --git a/src/config.rs b/src/config.rs index bca90b2..ff5cec6 100644 --- a/src/config.rs +++ b/src/config.rs @@ -175,6 +175,42 @@ pub trait StarkGenericConfig { /// evaluations. fn log_blowup(&self) -> usize; + /// Commit deterministic main-trace sources, optionally restoring tree nodes. + fn commit_main( + &self, + evaluations: Vec<(Domain, crate::witness::TraceSource>)>, + _checkpoint: Option, + ) -> (Com, PcsData) + where + Self: Sized, + { + self.pcs().commit( + evaluations + .into_iter() + .map(|(domain, trace)| (domain, trace.materialize())), + ) + } + + /// Consume all PCS data, keeping at most a bounded tree-only checkpoint. + fn checkpoint_main( + &self, + _data: PcsData, + _max_bytes: usize, + ) -> Option + where + Self: Sized, + { + None + } + + /// Available cache allowance after reserving the incoming proof's workspace. + fn tree_cache_headroom(&self, _witness: &crate::witness::PreparedWitness>) -> usize + where + Self: Sized, + { + 0 + } + /// Normalize any backend-dependent field representatives before proof /// serialization. Most fields have canonical in-memory representations; /// configurations whose field permits lazy reduction can override this. diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index 7b419b6..78c529f 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -175,13 +175,15 @@ impl CudaMmcsData> { retained_traces, .. } => { + let lde = resident.get(index)?.as_ref()?; if !resident_active .get(index)? .load(std::sync::atomic::Ordering::Acquire) + && !lde.has_generator() + && retained_traces.get(index)?.is_none() { return None; } - let lde = resident.get(index)?.as_ref()?; if let Some(trace) = retained_traces.get(index)?.as_ref() { // SAFETY: the trace is owned by this prover data and drops // after the resident LDE which holds the registered pointer. @@ -219,7 +221,7 @@ pub(crate) fn hash_cpu_height_groups( pub(crate) fn hash_host_only_height_groups( matrices: &[Option>], - resident: &[Option], + resident: &[Option<&CudaLde>], deferred_dimensions: &[Option], prehashed_heights: &std::collections::BTreeSet, ) -> Vec<(usize, Vec<[u8; 32]>)> { @@ -242,7 +244,7 @@ pub(crate) fn hash_host_only_height_groups( let height = matrix.as_ref().map_or_else( || { lde.as_ref() - .map_or_else(|| deferred.unwrap().height, CudaLde::height) + .map_or_else(|| deferred.unwrap().height, |lde| lde.height()) }, Matrix::height, ); @@ -389,11 +391,13 @@ pub struct CudaMmcs { } impl CudaMmcs { - pub(crate) fn new(cpu: CpuMmcs) -> Self { - Self { - cpu, - device_id: super::configured_device(), - } + /// A commitment scheme resident on the given CUDA device. Every kernel + /// it launches, allocation it makes and buffer it stages through belongs + /// to that device, so several of these in one process can each own a + /// device of their own. + pub(crate) fn with_device(cpu: CpuMmcs, device_id: i32) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + Self { cpu, device_id } } } @@ -1337,7 +1341,7 @@ mod tests { let host_matrices = vec![None, Some(matrices[1].clone()), Some(matrices[2].clone())]; let host_digest_groups = hash_host_only_height_groups( &host_matrices, - &resident, + &resident.iter().map(Option::as_ref).collect::>(), &[None, None, None], &std::collections::BTreeSet::new(), ); diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index db2e8a5..5c90ad2 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod mmcs; #[doc(hidden)] pub mod pcs; +pub(crate) mod witness; use core::ffi::{CStr, c_char, c_void}; use core::mem::{align_of, size_of}; @@ -242,6 +243,173 @@ impl CudaDft { /// Bit-reversed coset-LDE storage owned by a CUDA device allocation. #[doc(hidden)] +/// A borrowed output tile on the selected CUDA device. The writer must finish +/// using it before returning; rows wrap at the source's padded height. +pub struct DeviceTraceView<'a> { + device_id: i32, + output: *mut u64, + first: usize, + rows: usize, + width: usize, + _borrow: core::marker::PhantomData<&'a mut [u64]>, +} + +impl DeviceTraceView<'_> { + pub fn device_id(&self) -> i32 { + self.device_id + } + pub fn as_mut_ptr(&self) -> *mut u64 { + self.output + } + pub fn first_row(&self) -> usize { + self.first + } + pub fn rows(&self) -> usize { + self.rows + } + pub fn width(&self) -> usize { + self.width + } +} + +type Generator = Arc>; + +unsafe extern "C" fn write_generated_trace( + context: *mut c_void, + device_id: i32, + output: *mut u64, + first: usize, + rows: usize, +) -> i32 { + let generator = unsafe { &*context.cast::() }; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + generator.write_device_rows(DeviceTraceView { + device_id, + output, + first, + rows, + width: generator.width(), + _borrow: core::marker::PhantomData, + }) + })); + match result { + Ok(Ok(())) => 0, + Ok(Err(error)) => { + eprintln!("[multi-stark/cuda] trace generation failed: {error}"); + 999 + } + Err(_) => 999, + } +} + +unsafe extern "C" fn destroy_generated_trace(context: *mut c_void) { + drop(unsafe { Box::from_raw(context.cast::()) }); +} + +impl CudaDft { + /// Materialize a bounded generated tile for diagnostics and reference checks. + pub fn generated_trace_rows( + &self, + source: Generator, + first: usize, + rows: usize, + ) -> RowMajorMatrix { + assert!( + rows > 0 && rows <= (1 << 16) + 1, + "generated tile exceeds the row limit" + ); + let width = source.width(); + let mut output = + vec![Goldilocks::ZERO; rows.checked_mul(width).expect("trace tile overflow")]; + let mut source = source; + let status = unsafe { + multi_stark_cuda_generate_trace_rows( + self.device_id, + (&mut source as *mut Generator).cast(), + write_generated_trace, + first, + rows, + width, + output.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "generated trace tile download"); + RowMajorMatrix::new(output, width) + } + + pub fn generate_coset_lde( + &self, + generator: Generator, + added_bits: usize, + shift: Goldilocks, + ) -> CudaLde { + let height = generator.height(); + let width = generator.width(); + Self::validate_dimensions(height, width); + let extended_height = height + .checked_shl(added_bits.try_into().unwrap()) + .expect("LDE height overflow"); + Self::validate_dimensions(extended_height, width); + let log_height = log2_strict_usize(height); + let inverse_twiddles = self.twiddles(log_height, true); + let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let shift_powers = self.shift_powers(height, shift); + let mut context = Box::new(generator); + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_coset_lde_generate( + self.device_id, + &mut handle, + height, + width, + added_bits, + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(Goldilocks::ONE.div_2exp_u64(log_height as u64)), + (&mut *context as *mut Generator).cast(), + write_generated_trace, + destroy_generated_trace, + ) + }; + check_cuda(status, "generated trace LDE"); + let _ = Box::into_raw(context); + CudaLde { + device_id: self.device_id, + handle: NonNull::new(handle).expect("null generated LDE"), + height: extended_height, + width, + } + } +} + +unsafe extern "C" { + fn multi_stark_cuda_generate_trace_rows( + device: i32, + context: *mut c_void, + writer: unsafe extern "C" fn(*mut c_void, i32, *mut u64, usize, usize) -> i32, + first: usize, + rows: usize, + width: usize, + output: *mut u64, + ) -> i32; + fn multi_stark_cuda_coset_lde_generate( + device: i32, + output: *mut *mut c_void, + height: usize, + width: usize, + added_bits: usize, + inverse: *const u64, + shifts: *const u64, + forward: *const u64, + height_inverse: u64, + context: *mut c_void, + writer: unsafe extern "C" fn(*mut c_void, i32, *mut u64, usize, usize) -> i32, + destroy: unsafe extern "C" fn(*mut c_void), + ) -> i32; + fn multi_stark_cuda_lde_has_generator(handle: *const c_void) -> bool; +} + pub struct CudaLde { device_id: i32, handle: NonNull, @@ -253,6 +421,10 @@ unsafe impl Send for CudaLde {} unsafe impl Sync for CudaLde {} impl CudaLde { + pub(crate) fn has_generator(&self) -> bool { + unsafe { multi_stark_cuda_lde_has_generator(self.raw_handle()) } + } + pub(crate) const fn raw_handle(&self) -> *const c_void { self.handle.as_ptr() } diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 250f24d..c2bf1df 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -891,7 +891,7 @@ where .collect_vec(); host_digest_groups.extend(hash_host_only_height_groups( &host_matrices, - &resident, + &resident.iter().map(Option::as_ref).collect::>(), &deferred_dimensions, &prehashed_heights, )); diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs new file mode 100644 index 0000000..6e0c577 --- /dev/null +++ b/src/cuda/witness.rs @@ -0,0 +1,262 @@ +//! Device main-trace generation and tree-only stage-one checkpoints. + +use std::sync::{OnceLock, atomic::AtomicBool}; + +use p3_commit::Pcs; +use p3_field::{Field, PrimeCharacteristicRing}; +use p3_matrix::{Dimensions, dense::RowMajorMatrix}; +use p3_symmetric::MerkleCap; + +use super::mmcs::CudaMmcsData; +use super::{CudaLde, CudaMixedMerkleTree, device_memory_info}; +use crate::config::{Com, Domain, PcsData}; +use crate::types::{GoldilocksBlake3Config as Config, Pcs as CudaPcs, Val}; +use crate::witness::{PreparedWitness, TraceSource, TreeCheckpoint}; + +struct Checkpoint { + tree: CudaMixedMerkleTree, + dimensions: Vec, +} + +pub(crate) fn checkpoint(data: PcsData, max_bytes: usize) -> Option { + let (tree, dimensions) = match data { + CudaMmcsData::Cpu(_) => return None, + CudaMmcsData::Hybrid { + tree, + dimensions, + deferred_matrices, + .. + } => { + // Joining releases deferred host LDEs before crossing the barrier. + for deferred in deferred_matrices.into_iter().flatten() { + if let Some(worker) = deferred + .into_inner() + .expect("deferred matrix lock poisoned") + { + drop(worker.join().expect("deferred LDE worker panicked")); + } + } + (tree, dimensions) + } + CudaMmcsData::Cuda { tree, resident, .. } => { + let dimensions = resident + .iter() + .map(|m| Dimensions { + height: m.height(), + width: m.width(), + }) + .collect(); + (tree, dimensions) + } + }; + let bytes = tree + .row_count + .saturating_mul(64) + .saturating_add( + dimensions + .capacity() + .saturating_mul(size_of::()), + ) + .saturating_add(size_of::()); + if bytes > max_bytes { + return None; + } + tracing::debug!(bytes, "retained stage-one Merkle tree"); + Some(TreeCheckpoint { + bytes, + data: Box::new(Checkpoint { tree, dimensions }), + }) +} + +fn reserve_bytes(total: usize) -> usize { + std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") + .ok() + .and_then(|n| n.parse().ok()) + .unwrap_or(total / 4) +} + +pub(crate) fn cache_headroom(pcs: &CudaPcs, witness: &PreparedWitness) -> usize { + let (free, total) = device_memory_info(pcs.dft.device_id()); + let main = witness + .traces + .iter() + .map(|m| m.height().saturating_mul(m.width()).saturating_mul(8)) + .sum::(); + let lde = main + .checked_shl(pcs.fri.log_blowup as u32) + .unwrap_or(usize::MAX); + // Cache memory is expendable. Leave the active proof ample room for its + // main and lookup LDEs, quotient, FRI, row tiles and transform constants. + free.saturating_sub( + lde.saturating_mul(4) + .saturating_add(main) + .saturating_add(reserve_bytes(total)) + .saturating_add(256 << 20), + ) +} + +pub(crate) fn commit( + pcs: &CudaPcs, + evaluations: Vec<(Domain, TraceSource)>, + checkpoint: Option, +) -> (Com, PcsData) { + if checkpoint.is_none() + && evaluations + .iter() + .all(|(_, m)| matches!(m, TraceSource::Host(_))) + { + return >::commit( + pcs, + evaluations.into_iter().map(|(d, m)| (d, m.materialize())), + ); + } + let device = pcs.dft.device_id(); + let blowup = pcs.fri.log_blowup; + let dimensions: Vec<_> = evaluations + .iter() + .map(|(_, m)| Dimensions { + width: m.width(), + height: m + .height() + .checked_shl(blowup.try_into().expect("LDE blowup exceeds u32")) + .expect("LDE height overflow"), + }) + .collect(); + let mut cached = checkpoint + .and_then(|c| c.data.downcast::().ok()) + .filter(|c| { + c.tree.device_id == device + && c.dimensions == dimensions + && evaluations.iter().all(|(d, _)| d.shift() == Val::ONE) + }); + for (domain, source) in &evaluations { + assert_eq!( + domain.size(), + source.height(), + "main-trace domain height mismatch" + ); + pcs.dft.prepare_coset_lde_constants( + source.height(), + blowup, + Val::GENERATOR / domain.shift(), + ); + } + let max_height = dimensions.iter().map(|d| d.height).max().unwrap(); + let (_, total) = device_memory_info(device); + let reserve = reserve_bytes(total) + .saturating_add(max_height.saturating_mul(96)) + .saturating_add(128 << 20); + let mut resident: Vec> = Vec::with_capacity(evaluations.len()); + let mut host: Vec>> = Vec::with_capacity(evaluations.len()); + let mut retained = Vec::with_capacity(evaluations.len()); + for (domain, source) in evaluations { + let needed = source + .height() + .saturating_mul(source.width()) + .saturating_mul(8) + .saturating_mul((1 << blowup) + 1) + .saturating_add(reserve); + if device_memory_info(device).0 < needed { + drop(cached.take()); + } + for index in 0..resident.len() { + if device_memory_info(device).0 >= needed { + break; + } + if host[index].is_none() { + let lde = resident[index].as_ref().unwrap(); + host[index] = Some(lde.to_row_major_matrix()); + // Synchronous materialization has finished every device use. + unsafe { lde.release_values() }; + } + } + assert!( + device_memory_info(device).0 >= needed, + "generated trace commitment exceeds device admission; reduce the shard cell budget" + ); + let shift = Val::GENERATOR / domain.shift(); + let (lde, trace) = match source { + TraceSource::Host(matrix) => ( + pcs.dft.coset_lde_batch_resident(&matrix, blowup, shift), + Some(matrix), + ), + TraceSource::Generated(source) => { + (pcs.dft.generate_coset_lde(source, blowup, shift), None) + } + }; + // Both kinds retain a bounded recovery source. Raw device rows need + // not coexist with later lookup/quotient workspace. + unsafe { lde.release_trace() }; + let spilled = if std::env::var("MULTI_STARK_CUDA_TRACE_FORCE_SPILL").is_ok_and(|v| v == "1") + { + let matrix = lde.to_row_major_matrix(); + unsafe { lde.release_values() }; + Some(matrix) + } else { + None + }; + resident.push(Some(lde)); + host.push(spilled); + retained.push(trace); + } + let spilled_heights: std::collections::BTreeSet<_> = host + .iter() + .enumerate() + .filter_map(|(index, h)| h.is_some().then_some(dimensions[index].height)) + .collect(); + for index in 0..host.len() { + if host[index].is_none() && spilled_heights.contains(&dimensions[index].height) { + let lde = resident[index].as_ref().unwrap(); + host[index] = Some(lde.to_row_major_matrix()); + unsafe { lde.release_values() }; + } + } + let resident_refs: Vec<_> = resident + .iter() + .zip(&host) + .map(|(r, h)| if h.is_none() { r.as_ref() } else { None }) + .collect(); + let host_refs: Vec<_> = host.iter().map(Option::as_ref).collect(); + let tree = if let Some(cached) = cached { + tracing::debug!( + bytes = cached.tree.row_count * 64, + "reused stage-one Merkle tree" + ); + cached.tree + } else { + let deferred = vec![None; dimensions.len()]; + let digests = super::mmcs::hash_host_only_height_groups( + &host, + &resident_refs, + &deferred, + &std::collections::BTreeSet::new(), + ); + CudaMixedMerkleTree::from_hybrid(device, &resident_refs, &host_refs, &deferred, &digests) + }; + let commitment = MerkleCap::new(vec![tree.root()]); + let active = host.iter().map(|h| AtomicBool::new(h.is_none())).collect(); + let committed = host + .into_iter() + .map(|h| { + let cell = OnceLock::new(); + if let Some(h) = h { + cell.set(h).unwrap(); + } + cell + }) + .collect(); + let count = dimensions.len(); + ( + commitment, + CudaMmcsData::Hybrid { + resident, + resident_active: active, + materialize: Box::new(CudaLde::to_row_major_matrix), + committed_matrices: committed, + deferred_matrices: (0..count).map(|_| None).collect(), + dimensions, + retained_traces: retained, + tree, + }, + ) +} diff --git a/src/lib.rs b/src/lib.rs index 4172a75..c1f1aed 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub mod system; mod test_circuits; pub mod types; pub mod verifier; +pub mod witness; pub use p3_air; pub use p3_field; diff --git a/src/prover.rs b/src/prover.rs index bbc6d27..ca01124 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -434,7 +434,18 @@ where /// # Panics /// Panics if every circuit's trace is empty (nothing to prove). #[tracing::instrument(level = "info", skip_all, name = "stark/stage1_commit")] - pub fn prove_stage_1(&self, witness: SystemWitness>) -> Stage1 { + pub fn prove_stage_1( + &self, + witness: impl Into>>, + ) -> Stage1 { + self.prove_stage_1_restored(witness.into(), None) + } + + pub(crate) fn prove_stage_1_restored( + &self, + witness: crate::witness::PreparedWitness>, + checkpoint: Option, + ) -> Stage1 { let pcs = self.config.pcs(); // Sparse activation: a circuit whose stage-1 trace is empty is @@ -474,7 +485,8 @@ where log_degrees.push(log_degree); (trace_domain, trace) }); - let (stage_1_trace_commit, stage_1_trace_data) = pcs.commit(evaluations); + let (stage_1_trace_commit, stage_1_trace_data) = + self.config.commit_main(evaluations.collect(), checkpoint); // Only active circuits enter the accumulator chain; the chain (and // `intermediate_accumulators`) is indexed by active position. diff --git a/src/types.rs b/src/types.rs index 09f8e7d..ec7266e 100644 --- a/src/types.rs +++ b/src/types.rs @@ -277,6 +277,18 @@ pub struct GoldilocksBlake3Config { impl GoldilocksBlake3Config { pub fn new(commitment_parameters: CommitmentParameters, fri_parameters: FriParameters) -> Self { + Self::with_device(commitment_parameters, fri_parameters, None) + } + + /// A configuration whose prover lives on the given CUDA device (`None`: + /// the device `MULTI_STARK_CUDA_DEVICE` names, or 0). Without the + /// `cuda` feature only device 0 exists. Configurations on distinct + /// devices prove concurrently in one process. + pub fn with_device( + commitment_parameters: CommitmentParameters, + fri_parameters: FriParameters, + device_id: Option, + ) -> Self { #[cfg(feature = "cuda")] { assert_eq!( @@ -288,7 +300,20 @@ impl GoldilocksBlake3Config { "the CUDA backend currently supports only binary FRI folds" ); } - let (pcs, dft) = new_pcs(commitment_parameters, fri_parameters); + #[cfg(feature = "cuda")] + let (pcs, dft) = new_pcs( + commitment_parameters, + fri_parameters, + device_id.unwrap_or_else(crate::cuda::configured_device), + ); + #[cfg(not(feature = "cuda"))] + let (pcs, dft) = { + assert!( + device_id.is_none_or(|device| device == 0), + "no CUDA backend to place a prover on device {device_id:?}" + ); + new_pcs(commitment_parameters, fri_parameters) + }; // Seed the challenger with a protocol tag for domain separation, // followed by every protocol parameter. Binding the parameters into // the seed means transcripts produced under different parameters @@ -350,6 +375,32 @@ impl StarkGenericConfig for GoldilocksBlake3Config { self.log_blowup } + #[cfg(feature = "cuda")] + fn commit_main( + &self, + evaluations: Vec<( + crate::config::Domain, + crate::witness::TraceSource, + )>, + checkpoint: Option, + ) -> (crate::config::Com, crate::config::PcsData) { + crate::cuda::witness::commit(&self.pcs, evaluations, checkpoint) + } + + #[cfg(feature = "cuda")] + fn checkpoint_main( + &self, + data: crate::config::PcsData, + max_bytes: usize, + ) -> Option { + crate::cuda::witness::checkpoint(data, max_bytes) + } + + #[cfg(feature = "cuda")] + fn tree_cache_headroom(&self, witness: &crate::witness::PreparedWitness) -> usize { + crate::cuda::witness::cache_headroom(&self.pcs, witness) + } + fn canonicalize_proof(proof: &mut crate::prover::Proof) { fn canonical_base(value: &mut Val) { *value = Val::from_u64(value.as_canonical_u64()); @@ -1052,7 +1103,7 @@ type PcsDft = Dft; #[cfg(feature = "cuda")] type PcsDft = CudaDft; -fn new_mmcs(cap_height: usize) -> Mmcs { +fn new_mmcs(cap_height: usize, #[cfg(feature = "cuda")] device_id: i32) -> Mmcs { let byte_hash = Blake3; let field_hash = SerializingHasher::new(byte_hash); let compress = Blake3CompressionFunction::new(byte_hash); @@ -1060,13 +1111,17 @@ fn new_mmcs(cap_height: usize) -> Mmcs { #[cfg(not(feature = "cuda"))] return cpu; #[cfg(feature = "cuda")] - crate::cuda::mmcs::CudaMmcs::new(cpu) + crate::cuda::mmcs::CudaMmcs::with_device(cpu, device_id) } fn new_pcs( commitment_parameters: CommitmentParameters, fri_parameters: FriParameters, + #[cfg(feature = "cuda")] device_id: i32, ) -> (Pcs, Dft) { + #[cfg(feature = "cuda")] + let val_mmcs = new_mmcs(commitment_parameters.cap_height, device_id); + #[cfg(not(feature = "cuda"))] let val_mmcs = new_mmcs(commitment_parameters.cap_height); let mmcs = ExtensionMmcs::new(val_mmcs.clone()); let inner_parameters = InnerFriParameters { @@ -1079,7 +1134,11 @@ fn new_pcs( mmcs, }; let dft = Dft::default(); - let pcs = Pcs::new(PcsDft::default(), val_mmcs, inner_parameters); + #[cfg(feature = "cuda")] + let pcs_dft = CudaDft::new(device_id); + #[cfg(not(feature = "cuda"))] + let pcs_dft = PcsDft::default(); + let pcs = Pcs::new(pcs_dft, val_mmcs, inner_parameters); (pcs, dft) } diff --git a/src/witness.rs b/src/witness.rs new file mode 100644 index 0000000..ca998a5 --- /dev/null +++ b/src/witness.rs @@ -0,0 +1,85 @@ +//! Owned, deterministic main-trace sources. Producers do not allocate device memory. + +use std::sync::Arc; + +use p3_field::Field; +use p3_matrix::{Matrix, dense::RowMajorMatrix}; + +/// A frozen source which reproduces the same canonical cells on every call. +pub trait TraceGenerator: Send + Sync { + fn height(&self) -> usize; + fn width(&self) -> usize; + fn host_bytes(&self) -> usize; + /// Fill contiguous rows, wrapping at the padded height for lookup halos. + fn write_rows(&self, first: usize, output: &mut [F]); + + /// Fill a device tile synchronously on its owning device and calling stream. + #[cfg(feature = "cuda")] + fn write_device_rows(&self, output: crate::cuda::DeviceTraceView<'_>) -> Result<(), String>; +} + +#[derive(Clone)] +pub enum TraceSource { + Host(RowMajorMatrix), + Generated(Arc>), +} + +impl TraceSource { + pub fn height(&self) -> usize { + match self { + Self::Host(m) => m.height(), + Self::Generated(g) => g.height(), + } + } + + pub fn width(&self) -> usize { + match self { + Self::Host(m) => m.width(), + Self::Generated(g) => g.width(), + } + } + + pub fn materialize(self) -> RowMajorMatrix { + match self { + Self::Host(m) => m, + Self::Generated(g) => { + let mut values = vec![ + F::ZERO; + g.height() + .checked_mul(g.width()) + .expect("trace size overflow") + ]; + g.write_rows(0, &mut values); + RowMajorMatrix::new(values, g.width()) + } + } + } +} + +#[derive(Clone)] +pub struct PreparedWitness { + pub traces: Vec>, + pub lookups: Vec>, +} + +impl From> for PreparedWitness { + fn from(witness: crate::system::SystemWitness) -> Self { + Self { + traces: witness.traces.into_iter().map(TraceSource::Host).collect(), + lookups: witness.lookups, + } + } +} + +/// Backend-owned tree nodes and metadata only; never a trace or LDE owner. +pub struct TreeCheckpoint { + pub(crate) bytes: usize, + #[cfg_attr(not(feature = "cuda"), allow(dead_code))] + pub(crate) data: Box, +} + +impl TreeCheckpoint { + pub fn bytes(&self) -> usize { + self.bytes + } +} From 98692347d5ae0966934deb6f3cfb2a7a066519af Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Tue, 15 Sep 2026 21:26:03 +0000 Subject: [PATCH 04/34] cuda: remove the Merkle tree cache, keep per-construction device admission The bounded tree cache (TreeCheckpoint, Retention::MerkleTrees, checkpoint_main, prove_stage_1_restored, trim_tree_cache, tree_cache_bytes and the cache_headroom reservation) is removed. Its headroom rule reserved four times the main LDE plus the main trace plus the minimum-free allowance, about 95 GiB of a 96 GiB device for one 2^20-row BLAKE3 matrix, so it evicted every retained tree before reuse; with the accounting fixed (each trace/LDE construction admitted on its own, existing cache bytes not charged twice, optional trees released before required LDEs spill) a 24 GiB allowance reused all 91 trees on the eight-claim Init fixture and still did not move end-to-end time (290.06 s against a 284.54 s regenerate median), because hashing is a small slice of the proof. Regenerate is the one cross-round policy. What stays: the generated main commitment admits each trace and LDE construction separately from Merkle hashing, and LDE spills are counted per device (record_lde_spill) so a run can report them. Measurements and the archived experiment: ix bench/tree-cache-2026-09-15 and bench/gpu-trace-regenerate-2026-09-15. --- src/batch.rs | 59 +--------------- src/config.rs | 23 +----- src/cuda/mmcs.rs | 16 +++-- src/cuda/pcs.rs | 2 +- src/cuda/witness.rs | 169 +++++++++++++------------------------------- src/prover.rs | 11 +-- src/types.rs | 88 ++++++++++++----------- src/witness.rs | 13 ---- 8 files changed, 113 insertions(+), 268 deletions(-) diff --git a/src/batch.rs b/src/batch.rs index b1c29db..54b6da7 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -127,9 +127,6 @@ pub enum Retention { /// plus whatever the witness source holds, at the price of one extra /// witness build and stage-1 commitment per shard. Regenerate, - /// Retain only tree nodes within this per-batch device budget. Every - /// witness and LDE is regenerated; pressure evicts trees before proving. - MerkleTrees { max_bytes: usize }, } /// A public lookup message the verifier contributes to the batch's balance, @@ -202,7 +199,6 @@ pub struct BatchBarrier { claims: Vec>>>, headers: Vec>, retained: Vec>>, - trees: Vec>, } impl BatchBarrier { @@ -210,16 +206,6 @@ impl BatchBarrier { pub fn shards(&self) -> usize { self.headers.len() } - - /// Evict optional trees when a resident session needs their memory. - pub fn trim_tree_cache(&mut self, max_bytes: usize) { - evict_trees(&mut self.trees, max_bytes); - } - - /// Tree nodes and declared metadata retained across this batch's barrier. - pub fn tree_cache_bytes(&self) -> usize { - self.trees.iter().flatten().map(|tree| tree.bytes()).sum() - } } /// A vector of shard proofs under one preamble. @@ -546,40 +532,17 @@ where claims: Vec::new(), headers: Vec::new(), retained: Vec::new(), - trees: Vec::new(), }; consume_ahead(round_one.into_iter(), |shard, (claims, witness)| { let _g = tracing::info_span!("stark/batch_round_1", shard).entered(); - let witness = witness.into(); - evict_trees( - &mut barrier.trees, - self.config.tree_cache_headroom(&witness), - ); let stage_1 = self.prove_stage_1(witness); barrier .headers .push(stage_1.header(&claim_slices::(&claims))); - let mut tree = None; barrier.retained.push(match retention { Retention::Retain => Some(stage_1), Retention::Regenerate => None, - Retention::MerkleTrees { max_bytes } => { - let used = barrier - .trees - .iter() - .flatten() - .map(|t| t.bytes()) - .sum::(); - let remaining = max_bytes.saturating_sub(used); - if remaining != 0 { - tree = self - .config - .checkpoint_main(stage_1.stage_1_trace_data, remaining); - } - None - } }); - barrier.trees.push(tree); barrier.claims.push(claims); }); barrier @@ -621,7 +584,6 @@ where claims, headers, retained, - mut trees, } = barrier; assert!(!claims.is_empty(), "cannot prove an empty batch"); let preamble = BatchPreamble { headers, messages }; @@ -636,9 +598,7 @@ where let stage_1 = match (witness, retained) { (None, Some(stage_1)) => stage_1, (Some(witness), None) => { - let witness = witness.into(); - evict_trees(&mut trees, self.config.tree_cache_headroom(&witness)); - let stage_1 = self.prove_stage_1_restored(witness, trees[shard].take()); + let stage_1 = self.prove_stage_1(witness); assert!( stage_1.header(&claims) == preamble.headers[shard], "shard {shard} did not reproduce its round-one header" @@ -1013,20 +973,3 @@ mod tests { type ExtVal = ::Challenge; } - -fn evict_trees(trees: &mut [Option], headroom: usize) { - let mut used = trees - .iter() - .flatten() - .map(|tree| tree.bytes()) - .sum::(); - for entry in trees.iter_mut().rev() { - if used <= headroom { - break; - } - if let Some(tree) = entry.take() { - used -= tree.bytes(); - tracing::debug!(bytes = tree.bytes(), "evicted stage-one tree checkpoint"); - } - } -} diff --git a/src/config.rs b/src/config.rs index ff5cec6..6dcd85b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -175,11 +175,10 @@ pub trait StarkGenericConfig { /// evaluations. fn log_blowup(&self) -> usize; - /// Commit deterministic main-trace sources, optionally restoring tree nodes. + /// Commit deterministic main-trace sources. fn commit_main( &self, evaluations: Vec<(Domain, crate::witness::TraceSource>)>, - _checkpoint: Option, ) -> (Com, PcsData) where Self: Sized, @@ -191,26 +190,6 @@ pub trait StarkGenericConfig { ) } - /// Consume all PCS data, keeping at most a bounded tree-only checkpoint. - fn checkpoint_main( - &self, - _data: PcsData, - _max_bytes: usize, - ) -> Option - where - Self: Sized, - { - None - } - - /// Available cache allowance after reserving the incoming proof's workspace. - fn tree_cache_headroom(&self, _witness: &crate::witness::PreparedWitness>) -> usize - where - Self: Sized, - { - 0 - } - /// Normalize any backend-dependent field representatives before proof /// serialization. Most fields have canonical in-memory representations; /// configurations whose field permits lazy reduction can override this. diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index 78c529f..a1f65d3 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -126,6 +126,10 @@ fn evict_hybrid_resident(data: &CudaMmcsData, index: usize) { // SAFETY: admission transitions run between proving stages, when no CUDA // operation can access this LDE or its trace. unsafe { lde.release_values() }; + super::witness::record_lde_spill( + lde.device_id, + lde.height() * lde.width() * size_of::(), + ); resident_active[index].store(false, std::sync::atomic::Ordering::Release); } @@ -434,6 +438,7 @@ pub trait CudaCommitMmcs: Mmcs { data: &Self::ProverData, target_free_bytes: usize, protected_index: Option, + phase: &'static str, ) -> usize; /// Selects spill candidates across all supplied commitments instead of @@ -442,6 +447,7 @@ pub trait CudaCommitMmcs: Mmcs { &self, data: &[&Self::ProverData>], target_free_bytes: usize, + phase: &'static str, ) -> usize; fn matrix_dimensions(&self, data: &Self::ProverData>) -> Vec; @@ -591,11 +597,12 @@ impl CudaCommitMmcs for CudaMmcs { data: &Self::ProverData, target_free_bytes: usize, protected_index: Option, + phase: &'static str, ) -> usize { let (mut free_bytes, _) = super::device_memory_info(self.device_id); if super::memory_diagnostics_enabled() { eprintln!( - "[multi-stark/cuda] admission start: target={} free={}", + "[multi-stark/cuda] {phase} admission start: target={} free={}", target_free_bytes, free_bytes ); } @@ -640,11 +647,12 @@ impl CudaCommitMmcs for CudaMmcs { &self, data: &[&Self::ProverData>], target_free_bytes: usize, + phase: &'static str, ) -> usize { let (mut measured_free_bytes, _) = super::device_memory_info(self.device_id); if super::memory_diagnostics_enabled() { eprintln!( - "[multi-stark/cuda] batch admission start: target={} free={}", + "[multi-stark/cuda] {phase} batch admission start: target={} free={}", target_free_bytes, measured_free_bytes ); } @@ -1332,7 +1340,7 @@ mod tests { let (expected_commitment, expected_data) = cpu.commit(matrices.clone()); let expected_opening = cpu.open_batch(5, &expected_data); - let mmcs = CudaMmcs::new(cpu); + let mmcs = CudaMmcs::with_device(cpu, 0); let resident = vec![ Some(CudaLde::from_row_major_matrix(mmcs.device_id, &matrices[0])), None, @@ -1366,7 +1374,7 @@ mod tests { expected_opening.opening_proof ); - mmcs.ensure_device_headroom(&data, usize::MAX, None); + mmcs.ensure_device_headroom(&data, usize::MAX, None, "test"); assert!(!(0..matrices.len()).any(|index| mmcs.is_matrix_cuda_resident(&data, index))); let spilled_opening = mmcs.open_batch(5, &data); assert_eq!( diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index c2bf1df..8c3fd7e 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -1480,7 +1480,7 @@ where .map(|entry| entry.0) .collect_vec(); self.mmcs - .ensure_device_headroom_batch(&admission_data, fri_workspace_bytes); + .ensure_device_headroom_batch(&admission_data, fri_workspace_bytes, "fri"); if crate::cuda::memory_diagnostics_enabled() { eprintln!( "[multi-stark/cuda] FRI admission: {:.3}s", diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index 6e0c577..fe04d3b 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -1,9 +1,9 @@ -//! Device main-trace generation and tree-only stage-one checkpoints. +//! Device main-trace generation and commitment with bounded construction workspace. use std::sync::{OnceLock, atomic::AtomicBool}; use p3_commit::Pcs; -use p3_field::{Field, PrimeCharacteristicRing}; +use p3_field::Field; use p3_matrix::{Dimensions, dense::RowMajorMatrix}; use p3_symmetric::MerkleCap; @@ -11,61 +11,18 @@ use super::mmcs::CudaMmcsData; use super::{CudaLde, CudaMixedMerkleTree, device_memory_info}; use crate::config::{Com, Domain, PcsData}; use crate::types::{GoldilocksBlake3Config as Config, Pcs as CudaPcs, Val}; -use crate::witness::{PreparedWitness, TraceSource, TreeCheckpoint}; +use crate::witness::TraceSource; -struct Checkpoint { - tree: CudaMixedMerkleTree, - dimensions: Vec, +pub(crate) fn record_lde_spill(device: i32, bytes: usize) { + tracing::debug!(device, bytes, "spilled active LDE"); } -pub(crate) fn checkpoint(data: PcsData, max_bytes: usize) -> Option { - let (tree, dimensions) = match data { - CudaMmcsData::Cpu(_) => return None, - CudaMmcsData::Hybrid { - tree, - dimensions, - deferred_matrices, - .. - } => { - // Joining releases deferred host LDEs before crossing the barrier. - for deferred in deferred_matrices.into_iter().flatten() { - if let Some(worker) = deferred - .into_inner() - .expect("deferred matrix lock poisoned") - { - drop(worker.join().expect("deferred LDE worker panicked")); - } - } - (tree, dimensions) - } - CudaMmcsData::Cuda { tree, resident, .. } => { - let dimensions = resident - .iter() - .map(|m| Dimensions { - height: m.height(), - width: m.width(), - }) - .collect(); - (tree, dimensions) - } - }; - let bytes = tree - .row_count - .saturating_mul(64) - .saturating_add( - dimensions - .capacity() - .saturating_mul(size_of::()), - ) - .saturating_add(size_of::()); - if bytes > max_bytes { - return None; - } - tracing::debug!(bytes, "retained stage-one Merkle tree"); - Some(TreeCheckpoint { - bytes, - data: Box::new(Checkpoint { tree, dimensions }), - }) +fn spill_lde(lde: &CudaLde) -> RowMajorMatrix { + let matrix = lde.to_row_major_matrix(); + // Materialization synchronizes every use of the released values. + unsafe { lde.release_values() }; + record_lde_spill(lde.device_id, lde.height() * lde.width() * 8); + matrix } fn reserve_bytes(total: usize) -> usize { @@ -75,42 +32,26 @@ fn reserve_bytes(total: usize) -> usize { .unwrap_or(total / 4) } -pub(crate) fn cache_headroom(pcs: &CudaPcs, witness: &PreparedWitness) -> usize { - let (free, total) = device_memory_info(pcs.dft.device_id()); - let main = witness - .traces - .iter() - .map(|m| m.height().saturating_mul(m.width()).saturating_mul(8)) - .sum::(); - let lde = main - .checked_shl(pcs.fri.log_blowup as u32) - .unwrap_or(usize::MAX); - // Cache memory is expendable. Leave the active proof ample room for its - // main and lookup LDEs, quotient, FRI, row tiles and transform constants. - free.saturating_sub( - lde.saturating_mul(4) - .saturating_add(main) - .saturating_add(reserve_bytes(total)) - .saturating_add(256 << 20), - ) -} - pub(crate) fn commit( pcs: &CudaPcs, evaluations: Vec<(Domain, TraceSource)>, - checkpoint: Option, ) -> (Com, PcsData) { - if checkpoint.is_none() - && evaluations - .iter() - .all(|(_, m)| matches!(m, TraceSource::Host(_))) + if evaluations + .iter() + .all(|(_, m)| matches!(m, TraceSource::Host(_))) { + tracing::debug!( + device = pcs.dft.device_id(), + path = "host_pcs", + "main commitment path" + ); return >::commit( pcs, evaluations.into_iter().map(|(d, m)| (d, m.materialize())), ); } let device = pcs.dft.device_id(); + tracing::debug!(device, path = "prepared", "main commitment path"); let blowup = pcs.fri.log_blowup; let dimensions: Vec<_> = evaluations .iter() @@ -122,30 +63,16 @@ pub(crate) fn commit( .expect("LDE height overflow"), }) .collect(); - let mut cached = checkpoint - .and_then(|c| c.data.downcast::().ok()) - .filter(|c| { - c.tree.device_id == device - && c.dimensions == dimensions - && evaluations.iter().all(|(d, _)| d.shift() == Val::ONE) - }); for (domain, source) in &evaluations { assert_eq!( domain.size(), source.height(), "main-trace domain height mismatch" ); - pcs.dft.prepare_coset_lde_constants( - source.height(), - blowup, - Val::GENERATOR / domain.shift(), - ); } let max_height = dimensions.iter().map(|d| d.height).max().unwrap(); let (_, total) = device_memory_info(device); - let reserve = reserve_bytes(total) - .saturating_add(max_height.saturating_mul(96)) - .saturating_add(128 << 20); + let reserve = reserve_bytes(total).saturating_add(128 << 20); let mut resident: Vec> = Vec::with_capacity(evaluations.len()); let mut host: Vec>> = Vec::with_capacity(evaluations.len()); let mut retained = Vec::with_capacity(evaluations.len()); @@ -156,18 +83,13 @@ pub(crate) fn commit( .saturating_mul(8) .saturating_mul((1 << blowup) + 1) .saturating_add(reserve); - if device_memory_info(device).0 < needed { - drop(cached.take()); - } for index in 0..resident.len() { if device_memory_info(device).0 >= needed { break; } if host[index].is_none() { let lde = resident[index].as_ref().unwrap(); - host[index] = Some(lde.to_row_major_matrix()); - // Synchronous materialization has finished every device use. - unsafe { lde.release_values() }; + host[index] = Some(spill_lde(lde)); } } assert!( @@ -175,6 +97,8 @@ pub(crate) fn commit( "generated trace commitment exceeds device admission; reduce the shard cell budget" ); let shift = Val::GENERATOR / domain.shift(); + pcs.dft + .prepare_coset_lde_constants(source.height(), blowup, shift); let (lde, trace) = match source { TraceSource::Host(matrix) => ( pcs.dft.coset_lde_batch_resident(&matrix, blowup, shift), @@ -189,9 +113,7 @@ pub(crate) fn commit( unsafe { lde.release_trace() }; let spilled = if std::env::var("MULTI_STARK_CUDA_TRACE_FORCE_SPILL").is_ok_and(|v| v == "1") { - let matrix = lde.to_row_major_matrix(); - unsafe { lde.release_values() }; - Some(matrix) + Some(spill_lde(&lde)) } else { None }; @@ -199,6 +121,19 @@ pub(crate) fn commit( host.push(spilled); retained.push(trace); } + let needed = max_height.saturating_mul(96).saturating_add(reserve); + for index in 0..resident.len() { + if device_memory_info(device).0 >= needed { + break; + } + if host[index].is_none() { + host[index] = Some(spill_lde(resident[index].as_ref().unwrap())); + } + } + assert!( + device_memory_info(device).0 >= needed, + "main Merkle tree exceeds device admission; reduce the shard cell budget" + ); let spilled_heights: std::collections::BTreeSet<_> = host .iter() .enumerate() @@ -207,8 +142,7 @@ pub(crate) fn commit( for index in 0..host.len() { if host[index].is_none() && spilled_heights.contains(&dimensions[index].height) { let lde = resident[index].as_ref().unwrap(); - host[index] = Some(lde.to_row_major_matrix()); - unsafe { lde.release_values() }; + host[index] = Some(spill_lde(lde)); } } let resident_refs: Vec<_> = resident @@ -217,22 +151,15 @@ pub(crate) fn commit( .map(|(r, h)| if h.is_none() { r.as_ref() } else { None }) .collect(); let host_refs: Vec<_> = host.iter().map(Option::as_ref).collect(); - let tree = if let Some(cached) = cached { - tracing::debug!( - bytes = cached.tree.row_count * 64, - "reused stage-one Merkle tree" - ); - cached.tree - } else { - let deferred = vec![None; dimensions.len()]; - let digests = super::mmcs::hash_host_only_height_groups( - &host, - &resident_refs, - &deferred, - &std::collections::BTreeSet::new(), - ); - CudaMixedMerkleTree::from_hybrid(device, &resident_refs, &host_refs, &deferred, &digests) - }; + let deferred = vec![None; dimensions.len()]; + let digests = super::mmcs::hash_host_only_height_groups( + &host, + &resident_refs, + &deferred, + &std::collections::BTreeSet::new(), + ); + let tree = + CudaMixedMerkleTree::from_hybrid(device, &resident_refs, &host_refs, &deferred, &digests); let commitment = MerkleCap::new(vec![tree.root()]); let active = host.iter().map(|h| AtomicBool::new(h.is_none())).collect(); let committed = host diff --git a/src/prover.rs b/src/prover.rs index ca01124..e1adbd1 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -438,14 +438,7 @@ where &self, witness: impl Into>>, ) -> Stage1 { - self.prove_stage_1_restored(witness.into(), None) - } - - pub(crate) fn prove_stage_1_restored( - &self, - witness: crate::witness::PreparedWitness>, - checkpoint: Option, - ) -> Stage1 { + let witness = witness.into(); let pcs = self.config.pcs(); // Sparse activation: a circuit whose stage-1 trace is empty is @@ -486,7 +479,7 @@ where (trace_domain, trace) }); let (stage_1_trace_commit, stage_1_trace_data) = - self.config.commit_main(evaluations.collect(), checkpoint); + self.config.commit_main(evaluations.collect()); // Only active circuits enter the accumulator chain; the chain (and // `intermediate_accumulators`) is indexed by active position. diff --git a/src/types.rs b/src/types.rs index ec7266e..7d2eaeb 100644 --- a/src/types.rs +++ b/src/types.rs @@ -382,23 +382,8 @@ impl StarkGenericConfig for GoldilocksBlake3Config { crate::config::Domain, crate::witness::TraceSource, )>, - checkpoint: Option, ) -> (crate::config::Com, crate::config::PcsData) { - crate::cuda::witness::commit(&self.pcs, evaluations, checkpoint) - } - - #[cfg(feature = "cuda")] - fn checkpoint_main( - &self, - data: crate::config::PcsData, - max_bytes: usize, - ) -> Option { - crate::cuda::witness::checkpoint(data, max_bytes) - } - - #[cfg(feature = "cuda")] - fn tree_cache_headroom(&self, witness: &crate::witness::PreparedWitness) -> usize { - crate::cuda::witness::cache_headroom(&self.pcs, witness) + crate::cuda::witness::commit(&self.pcs, evaluations) } fn canonicalize_proof(proof: &mut crate::prover::Proof) { @@ -637,18 +622,24 @@ impl StarkGenericConfig for GoldilocksBlake3Config { .saturating_add(total_device_bytes / 64) }; let mut target = required(); - self.pcs - .mmcs - .ensure_device_headroom(input.stage_1.0, target, Some(input.stage_1.1)); + self.pcs.mmcs.ensure_device_headroom( + input.stage_1.0, + target, + Some(input.stage_1.1), + "quotient", + ); target = required(); - self.pcs - .mmcs - .ensure_device_headroom(input.stage_2.0, target, Some(input.stage_2.1)); + self.pcs.mmcs.ensure_device_headroom( + input.stage_2.0, + target, + Some(input.stage_2.1), + "quotient", + ); if let Some((data, matrix)) = input.preprocessed { target = required(); self.pcs .mmcs - .ensure_device_headroom(data, target, Some(matrix)); + .ensure_device_headroom(data, target, Some(matrix), "quotient"); } target = required(); let free_bytes = crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()).0; @@ -743,17 +734,23 @@ impl StarkGenericConfig for GoldilocksBlake3Config { .saturating_mul(96) .saturating_add(total_device_bytes / 64); if let Some(input) = inputs.first() { - self.pcs - .mmcs - .ensure_device_headroom(input.stage_1.0, tree_headroom, None); - self.pcs - .mmcs - .ensure_device_headroom(input.stage_2.0, tree_headroom, None); + self.pcs.mmcs.ensure_device_headroom( + input.stage_1.0, + tree_headroom, + None, + "quotient_tree", + ); + self.pcs.mmcs.ensure_device_headroom( + input.stage_2.0, + tree_headroom, + None, + "quotient_tree", + ); } if let Some((data, _)) = inputs.iter().find_map(|input| input.preprocessed) { self.pcs .mmcs - .ensure_device_headroom(data, tree_headroom, None); + .ensure_device_headroom(data, tree_headroom, None, "quotient_tree"); } if crate::cuda::device_memory_info(self.pcs.mmcs.cuda_device_id()).0 < tree_headroom { return None; @@ -993,15 +990,16 @@ impl StarkGenericConfig for GoldilocksBlake3Config { input.stage_1.0, target, Some(input.stage_1.1), + "lookup", ); if free_bytes < target { // This circuit alone does not fit beside its resident trace. // Spill it as a last resort and use the direct lookup-values // path, which remains protocol-identical. - free_bytes = self - .pcs - .mmcs - .ensure_device_headroom(input.stage_1.0, target, None); + free_bytes = + self.pcs + .mmcs + .ensure_device_headroom(input.stage_1.0, target, None, "lookup"); } if free_bytes < target { return None; @@ -1028,10 +1026,12 @@ impl StarkGenericConfig for GoldilocksBlake3Config { .saturating_mul(96) .saturating_add(total_device_bytes / 64); if let Some(input) = inputs.first() { - let free_bytes = - self.pcs - .mmcs - .ensure_device_headroom(input.stage_1.0, tree_headroom, None); + let free_bytes = self.pcs.mmcs.ensure_device_headroom( + input.stage_1.0, + tree_headroom, + None, + "lookup_tree", + ); if free_bytes < tree_headroom { return None; } @@ -1332,7 +1332,11 @@ mod pcs_ref_gen { m1[8] = f(109); // row 2 = [107, 108, 109] let mut m2 = vec![f(0); 2]; m2[1] = f(202); // row 1 = [202] - let mmcs = new_mmcs(0); + let mmcs = new_mmcs( + 0, + #[cfg(feature = "cuda")] + 0, + ); let (commit, pd) = mmcs.commit(vec![ RowMajorMatrix::new(m0.clone(), 2), RowMajorMatrix::new(m1.clone(), 3), @@ -1382,7 +1386,11 @@ mod pcs_ref_gen { ]] ); - let mmcs = new_mmcs(2); + let mmcs = new_mmcs( + 2, + #[cfg(feature = "cuda")] + 0, + ); let (commit, pd) = mmcs.commit(vec![ RowMajorMatrix::new(m0, 2), RowMajorMatrix::new(m1, 3), diff --git a/src/witness.rs b/src/witness.rs index ca998a5..c592ff0 100644 --- a/src/witness.rs +++ b/src/witness.rs @@ -70,16 +70,3 @@ impl From> for PreparedWitness { } } } - -/// Backend-owned tree nodes and metadata only; never a trace or LDE owner. -pub struct TreeCheckpoint { - pub(crate) bytes: usize, - #[cfg_attr(not(feature = "cuda"), allow(dead_code))] - pub(crate) data: Box, -} - -impl TreeCheckpoint { - pub fn bytes(&self) -> usize { - self.bytes - } -} From 9b172ed5392a9a143af4fa059ae353c7b7a35034 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:23:07 +0000 Subject: [PATCH 05/34] cuda: hash height groups wider than 32 KiB per row on the host The BLAKE3 leaf kernel hashes one message of at most 32 KiB per row, and a mixed tree's leaf row at a height is every matrix of that height side by side, so a height group wider than 4,096 Goldilocks columns made multi_stark_cuda_mixed_merkle_create reject the tree ("invalid argument") and the prover panic in stage-one commit. ix's Aiur test programs hit it (a 9,282-column BLAKE3 hashing circuit at height 4); the production IxVM circuits are narrower, so no proof run had. Such a group is now hashed on the host and enters the tree as a digest group, the way a group spilled for memory already does: - commit_cuda_resident, the sink of every all-resident and spillable commit, copies the group's LDEs back to the host and delegates to commit_cuda_hybrid with hash_host_only_height_groups' digests; - the PCS placement under memory pressure never makes such a group durable or transient, so its LDEs are computed on the host directly; - the generated-trace commit spills the group whole, like a memory spill. host_hashed_heights (MAX_DEVICE_LEAF_ROW_BYTES) is the one rule all three share, with a unit test; the CUDA suite (72 tests) passes. --- src/cuda/mmcs.rs | 130 ++++++++++++++++++++++++++++++++++++++++++-- src/cuda/mod.rs | 26 ++++++++- src/cuda/pcs.rs | 16 +++++- src/cuda/witness.rs | 5 +- 4 files changed, 168 insertions(+), 9 deletions(-) diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index a1f65d3..23e8d74 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -223,6 +223,30 @@ pub(crate) fn hash_cpu_height_groups( .collect() } +/// The CUDA BLAKE3 leaf kernel hashes one message of at most 32 KiB per +/// row, and a mixed tree's leaf row at a height is every matrix of that +/// height side by side. A height group whose rows are wider than that is +/// hashed on the host and enters the tree as a digest group; the device +/// still builds the rest of the tree. +pub(crate) const MAX_DEVICE_LEAF_ROW_BYTES: usize = 32 * 1024; + +/// The heights whose concatenated leaf row exceeds +/// [`MAX_DEVICE_LEAF_ROW_BYTES`], and so must be hashed on the host. +pub(crate) fn host_hashed_heights( + dimensions: impl IntoIterator, +) -> std::collections::BTreeSet { + let mut row_bytes = std::collections::BTreeMap::::new(); + for dimensions in dimensions { + *row_bytes.entry(dimensions.height).or_default() += + dimensions.width * size_of::(); + } + row_bytes + .into_iter() + .filter(|(_, bytes)| *bytes > MAX_DEVICE_LEAF_ROW_BYTES) + .map(|(height, _)| height) + .collect() +} + pub(crate) fn hash_host_only_height_groups( matrices: &[Option>], resident: &[Option<&CudaLde>], @@ -946,6 +970,39 @@ impl CudaCommitMmcs for CudaMmcs { Self::Commitment, Self::ProverData>, ) { + // A height group whose leaf row is wider than the device kernel + // hashes cannot be part of an all-resident tree: its LDEs come back + // to the host, its digests are hashed there, and the commitment is + // the hybrid one every other stage already handles. + let wide_heights = host_hashed_heights(ldes.iter().map(|lde| Dimensions { + width: lde.width(), + height: lde.height(), + })); + if !wide_heights.is_empty() { + let mut resident: Vec> = Vec::with_capacity(ldes.len()); + let mut host: Vec>> = + Vec::with_capacity(ldes.len()); + for lde in ldes { + if wide_heights.contains(&lde.height()) { + host.push(Some(lde.to_row_major_matrix())); + resident.push(None); + } else { + host.push(None); + resident.push(Some(lde)); + } + } + let deferred: Vec>> = + (0..resident.len()).map(|_| None).collect(); + let deferred_dimensions = vec![None; resident.len()]; + let digests = hash_host_only_height_groups( + &host, + &resident.iter().map(Option::as_ref).collect::>(), + &deferred_dimensions, + &std::collections::BTreeSet::new(), + ); + let retained = (0..resident.len()).map(|_| None).collect(); + return self.commit_cuda_hybrid(resident, host, deferred, retained, digests); + } let ldes = std::sync::Arc::new(ldes); let tree = CudaMixedMerkleTree::from_ldes(self.device_id, &ldes); let commitment = MerkleCap::new(vec![tree.root()]); @@ -1106,11 +1163,20 @@ impl Mmcs for CudaMmcs { &self, inputs: Vec, ) -> (Self::Commitment, Self::ProverData) { - if inputs - .iter() - .map(|matrix| matrix.height().saturating_mul(matrix.width())) - .sum::() - > 1 << 18 + // Small commitments stay on the CPU, and so does one with a height + // group wider than the device leaf kernel hashes (32 KiB per row); + // the CPU MMCS is the reference the device tree reproduces. + let device_hashable = host_hashed_heights(inputs.iter().map(|matrix| Dimensions { + width: matrix.width(), + height: matrix.height(), + })) + .is_empty(); + if device_hashable + && inputs + .iter() + .map(|matrix| matrix.height().saturating_mul(matrix.width())) + .sum::() + > 1 << 18 { let resident: Vec<_> = inputs .iter() @@ -1329,6 +1395,60 @@ mod tests { ) } + #[test] + fn wide_height_groups_are_hashed_on_the_host() { + let d = |width, height| Dimensions { width, height }; + // 4096 columns of 8 bytes is exactly the 32 KiB the leaf kernel takes. + assert!(host_hashed_heights([d(4096, 8)]).is_empty()); + assert_eq!( + host_hashed_heights([d(4097, 8)]).into_iter().collect::>(), + vec![8] + ); + // Matrices at one height share a leaf row, so their widths add up; + // other heights are judged on their own. + assert_eq!( + host_hashed_heights([d(3000, 16), d(1500, 16), d(9282, 4), d(10, 4), d(14, 4096)]) + .into_iter() + .collect::>(), + vec![4, 16] + ); + } + + /// The dimensions of the Aiur BLAKE3 hashing test's stage-one traces at + /// size 64: the height-4 group is 9,731 columns wide, 78 KB per leaf row. + #[test] + fn resident_commit_with_a_wide_height_group_matches_the_cpu() { + let dims = [ + (4096, 14), (2048, 14), (4096, 17), (4, 37), (4, 35), (4, 371), + (4, 9282), (256, 110), (4, 6), (1024, 3), (262144, 10), + ]; + let matrices: Vec<_> = dims + .iter() + .enumerate() + .map(|(i, &(height, width))| matrix(height, width, 7 * i + 1)) + .collect(); + let cpu = CpuMmcs::new( + SerializingHasher::new(Blake3), + Blake3CompressionFunction::new(Blake3), + 0, + ); + let (expected_commitment, expected_data) = cpu.commit(matrices.clone()); + let indices = [0, 3, 1000, 262143]; + let expected_openings: Vec<_> = + indices.iter().map(|&i| cpu.open_batch(i, &expected_data)).collect(); + let mmcs = CudaMmcs::with_device(cpu, 0); + let ldes = matrices + .iter() + .map(|m| CudaLde::from_row_major_matrix(mmcs.device_id, m)) + .collect(); + let (commitment, data) = mmcs.commit_cuda_resident(ldes); + assert_eq!(commitment, expected_commitment); + for (&index, expected) in indices.iter().zip(&expected_openings) { + let opening = mmcs.open_batch(index, &data); + assert_eq!(opening.opened_values, expected.opened_values); + } + } + #[test] fn hybrid_openings_survive_device_spill() { let matrices = vec![matrix(16, 2, 3), matrix(8, 3, 71), matrix(16, 1, 109)]; diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 5c90ad2..6a8c744 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -2542,7 +2542,14 @@ impl CudaMixedMerkleTree { handles.len(), ) }; - check_cuda(status, "resident LDE Merkle tree creation"); + if status != 0 { + let dimensions: Vec<(usize, usize)> = + ldes.iter().map(|lde| (lde.height(), lde.width())).collect(); + check_cuda( + status, + &format!("resident LDE Merkle tree creation over (height, width) {dimensions:?}"), + ); + } Self { device_id, handle: NonNull::new(handle).expect("CUDA returned a null mixed Merkle handle"), @@ -2658,7 +2665,22 @@ impl CudaMixedMerkleTree { host_digest_groups.len(), ) }; - check_cuda(status, "hybrid CPU/CUDA mixed-height Merkle tree creation"); + if status != 0 { + let dims: Vec> = ldes + .iter() + .zip(host_matrices) + .zip(deferred_dimensions) + .map(|((lde, host), deferred)| { + lde.map(|l| (l.height(), l.width())) + .or_else(|| host.map(|m| (m.height(), m.width()))) + .or_else(|| deferred.map(|d| (d.height, d.width))) + }) + .collect(); + check_cuda( + status, + &format!("hybrid Merkle tree creation over (height, width) {dims:?}, host digest groups at heights {:?}", host_digest_groups.iter().map(|(h, _)| *h).collect::>()), + ); + } let row_count = heights.into_iter().max().unwrap(); Self { device_id, diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 8c3fd7e..56ae00c 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -604,6 +604,18 @@ where // Split allocatable device capacity between persistent LDEs and // later phase workspace. This prevents stage one from filling VRAM // with values which lookup immediately has to copy back and evict. + // Height groups wider than the device leaf kernel can hash are + // neither durable nor transient: their LDEs and digests come + // from the host, whatever the memory budget says. + let wide_heights = super::mmcs::host_hashed_heights( + evaluations.iter().map(|(_, matrix)| p3_matrix::Dimensions { + width: matrix.width(), + height: matrix.height(), + }), + ); + for height in &wide_heights { + height_groups.remove(height); + } let durable_budget = gpu_lde_budget / 2; // A Merkle leaf combines every matrix at a given height. Keeping // height groups intact avoids streaming the CPU half of a split @@ -631,7 +643,9 @@ where let select_transient_plan = |transient_budget| { height_indices .iter() - .filter(|(height, _)| !durable_heights.contains(height)) + .filter(|(height, _)| { + !durable_heights.contains(height) && !wide_heights.contains(height) + }) .filter_map(|(&height, indices)| { let resources = indices .iter() diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index fe04d3b..d13f318 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -134,11 +134,14 @@ pub(crate) fn commit( device_memory_info(device).0 >= needed, "main Merkle tree exceeds device admission; reduce the shard cell budget" ); - let spilled_heights: std::collections::BTreeSet<_> = host + let mut spilled_heights: std::collections::BTreeSet<_> = host .iter() .enumerate() .filter_map(|(index, h)| h.is_some().then_some(dimensions[index].height)) .collect(); + // A height group wider than the device leaf kernel hashes is spilled + // whole and hashed on the host, like a group spilled for memory. + spilled_heights.extend(super::mmcs::host_hashed_heights(dimensions.iter().copied())); for index in 0..host.len() { if host[index].is_none() && spilled_heights.contains(&dimensions[index].height) { let lde = resident[index].as_ref().unwrap(); From 39d59cba8509eaa91a8a39022ca22b3964c6631f Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Wed, 16 Sep 2026 04:38:28 +0000 Subject: [PATCH 06/34] cuda: share Goldilocks arithmetic with generated trace kernels Export the field helpers in a shared CUDA header and expose its include directory through Cargo dependency metadata. Keep backend kernels on the same arithmetic implementation. --- Cargo.toml | 1 + build.rs | 4 +++ cuda/goldilocks.cuh | 80 +++++++++++++++++++++++++++++++++++++++++++++ cuda/kernels.cu | 71 ++-------------------------------------- 4 files changed, 87 insertions(+), 69 deletions(-) create mode 100644 cuda/goldilocks.cuh diff --git a/Cargo.toml b/Cargo.toml index e3aac95..3455901 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" authors = ["Argument Engineering "] license = "MIT OR Apache-2.0" rust-version = "1.98" +links = "multi_stark_cuda" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/build.rs b/build.rs index 438365e..03dc72c 100644 --- a/build.rs +++ b/build.rs @@ -11,6 +11,7 @@ use std::process::Command; fn main() { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=cuda/kernels.cu"); + println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); println!("cargo:rerun-if-env-changed=NVCC"); println!("cargo:rerun-if-env-changed=CUDA_HOME"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); @@ -20,6 +21,9 @@ fn main() { return; } + let include = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").unwrap()).join("cuda"); + println!("cargo:include={}", include.display()); + assert_eq!( env::var("CARGO_CFG_TARGET_OS").as_deref(), Ok("linux"), diff --git a/cuda/goldilocks.cuh b/cuda/goldilocks.cuh new file mode 100644 index 0000000..abebe64 --- /dev/null +++ b/cuda/goldilocks.cuh @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +#pragma once + +#include +#include + +namespace multi_stark_cuda { + +constexpr uint64_t GOLDILOCKS_P = 0xffffffff00000001ULL; +constexpr uint64_t GOLDILOCKS_EPSILON = 0x00000000ffffffffULL; + +__device__ __forceinline__ uint64_t canonicalize(uint64_t value) { + return value >= GOLDILOCKS_P ? value - GOLDILOCKS_P : value; +} + +__device__ __forceinline__ uint64_t goldilocks_add(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + // Written this way to avoid relying on overflow behavior in the source + // language. `GOLDILOCKS_P - right` is always representable. + const uint64_t gap = GOLDILOCKS_P - right; + return left >= gap ? left - gap : left + right; +} + +__device__ __forceinline__ uint64_t goldilocks_sub(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + return left >= right ? left - right : GOLDILOCKS_P - (right - left); +} + +__device__ __forceinline__ uint64_t goldilocks_mul(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + + const uint64_t low = left * right; + const uint64_t high = __umul64hi(left, right); + + // Reduce low + high * 2^64 using 2^64 = 2^32 - 1 (mod p). + const uint64_t high_high = high >> 32; + const uint64_t high_low = high & GOLDILOCKS_EPSILON; + uint64_t reduced_low = low - high_high; + if (low < high_high) { + // The wrapped subtraction added 2^64; replace that with +p. + reduced_low -= GOLDILOCKS_EPSILON; + } + const uint64_t reduced_high = high_low * GOLDILOCKS_EPSILON; + return goldilocks_add(reduced_low, reduced_high); +} + +__device__ __forceinline__ uint64_t goldilocks_pow(uint64_t base, uint64_t exponent) { + uint64_t result = 1; + while (exponent != 0) { + if ((exponent & 1U) != 0) { + result = goldilocks_mul(result, base); + } + base = goldilocks_mul(base, base); + exponent >>= 1; + } + return result; +} + +// The resident PCS canonicalizes every committed LDE, and interpolation +// tables are serialized with `as_canonical_u64`. Restrict the faster +// canonical-input arithmetic to this boundary instead of weakening the +// representation guarantees of lookup/quotient code, whose host inputs may +// legitimately use lazy representatives. +__device__ __forceinline__ uint64_t canonical_add(uint64_t a,uint64_t b){ + const uint64_t sum=a+b; + if(sum=GOLDILOCKS_P?sum-GOLDILOCKS_P:sum; +} +__device__ __forceinline__ uint64_t canonical_mul(uint64_t a,uint64_t b){ + const uint64_t low=a*b,high=__umul64hi(a,b),high_high=high>>32; + uint64_t reduced_low=low-high_high; + if(low +#include "goldilocks.cuh" #include #include @@ -19,8 +20,7 @@ namespace { -constexpr uint64_t GOLDILOCKS_P = 0xffffffff00000001ULL; -constexpr uint64_t GOLDILOCKS_EPSILON = 0x00000000ffffffffULL; +using namespace multi_stark_cuda; constexpr unsigned int THREADS = 256; constexpr unsigned int MAX_BLOCKS = 65535; constexpr size_t BLAKE3_CHUNK_BYTES = 1024; @@ -75,56 +75,6 @@ __device__ __constant__ unsigned int BLAKE3_PERMUTATION[16] = { 2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8, }; -__device__ __forceinline__ uint64_t canonicalize(uint64_t value) { - return value >= GOLDILOCKS_P ? value - GOLDILOCKS_P : value; -} - -__device__ __forceinline__ uint64_t goldilocks_add(uint64_t left, uint64_t right) { - left = canonicalize(left); - right = canonicalize(right); - // Written this way to avoid relying on overflow behavior in the source - // language. `GOLDILOCKS_P - right` is always representable. - const uint64_t gap = GOLDILOCKS_P - right; - return left >= gap ? left - gap : left + right; -} - -__device__ __forceinline__ uint64_t goldilocks_sub(uint64_t left, uint64_t right) { - left = canonicalize(left); - right = canonicalize(right); - return left >= right ? left - right : GOLDILOCKS_P - (right - left); -} - -__device__ __forceinline__ uint64_t goldilocks_mul(uint64_t left, uint64_t right) { - left = canonicalize(left); - right = canonicalize(right); - - const uint64_t low = left * right; - const uint64_t high = __umul64hi(left, right); - - // Reduce low + high * 2^64 using 2^64 = 2^32 - 1 (mod p). - const uint64_t high_high = high >> 32; - const uint64_t high_low = high & GOLDILOCKS_EPSILON; - uint64_t reduced_low = low - high_high; - if (low < high_high) { - // The wrapped subtraction added 2^64; replace that with +p. - reduced_low -= GOLDILOCKS_EPSILON; - } - const uint64_t reduced_high = high_low * GOLDILOCKS_EPSILON; - return goldilocks_add(reduced_low, reduced_high); -} - -__device__ __forceinline__ uint64_t goldilocks_pow(uint64_t base, uint64_t exponent) { - uint64_t result = 1; - while (exponent != 0) { - if ((exponent & 1U) != 0) { - result = goldilocks_mul(result, base); - } - base = goldilocks_mul(base, base); - exponent >>= 1; - } - return result; -} - __device__ __forceinline__ uint32_t rotate_right(uint32_t value, unsigned int count) { return __funnelshift_r(value, value, count); @@ -714,23 +664,6 @@ struct PendingLookupLde { } }; -// The resident PCS canonicalizes every committed LDE, and interpolation -// tables are serialized with `as_canonical_u64`. Restrict the faster -// canonical-input arithmetic to this boundary instead of weakening the -// representation guarantees of lookup/quotient code, whose host inputs may -// legitimately use lazy representatives. -__device__ __forceinline__ uint64_t canonical_add(uint64_t a,uint64_t b){ - const uint64_t sum=a+b; - if(sum=GOLDILOCKS_P?sum-GOLDILOCKS_P:sum; -} -__device__ __forceinline__ uint64_t canonical_mul(uint64_t a,uint64_t b){ - const uint64_t low=a*b,high=__umul64hi(a,b),high_high=high>>32; - uint64_t reduced_low=low-high_high; - if(low Date: Thu, 17 Sep 2026 02:05:48 +0000 Subject: [PATCH 07/34] cuda: admit lookup graphs only for resident LDEs that kept their trace Lookup admission budgeted the graph kernel for every resident LDE while execution required the raw trace as well, so a spilled or trace-released matrix was planned as regenerable and then evaluated on the slow path. `CudaMmcsData::has_resident_with_trace` is now the one predicate behind admission, the spill sort key and the re-budget after a spill, and the chosen path is passed into `evaluate` instead of being recomputed there. Under `MULTI_STARK_CUDA_TRACE_FORCE_SPILL` every lookup job now logs `graph=true`. --- src/cuda/mmcs.rs | 40 +++++++++++++++++++ src/types.rs | 101 +++++++++++++++++++++++++++++++---------------- 2 files changed, 107 insertions(+), 34 deletions(-) diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index 23e8d74..8765a0e 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -157,6 +157,46 @@ impl CudaMmcsData { } impl CudaMmcsData> { + /// Whether [`Self::resident_with_trace`] returns an LDE for this matrix, + /// without attaching anything. A hybrid LDE qualifies while its values are + /// resident, and after a spill as long as it can regenerate its trace from + /// a generator or a retained host trace. Admission decisions that size a + /// budget for the graph kernel must use this predicate, so that the budget + /// and the kernel that then runs agree. + pub(crate) fn has_resident_with_trace(&self, index: usize) -> bool { + match self { + Self::Cuda { + resident, + retained_traces, + .. + } => { + index < resident.len() + && retained_traces + .get() + .is_none_or(|retained| index < retained.len()) + } + Self::Hybrid { + resident, + resident_active, + retained_traces, + .. + } => { + let Some(Some(lde)) = resident.get(index) else { + return false; + }; + let (Some(active), Some(retained)) = + (resident_active.get(index), retained_traces.get(index)) + else { + return false; + }; + active.load(std::sync::atomic::Ordering::Acquire) + || lde.has_generator() + || retained.is_some() + } + Self::Cpu(_) => false, + } + } + pub(crate) fn resident_with_trace(&self, index: usize) -> Option<&CudaLde> { match self { Self::Cuda { diff --git a/src/types.rs b/src/types.rs index 7d2eaeb..266f13a 100644 --- a/src/types.rs +++ b/src/types.rs @@ -15,8 +15,8 @@ use p3_commit::{ExtensionMmcs, Pcs as PcsTrait}; use p3_dft::Radix2DitParallel; use p3_field::BasedVectorSpace; use p3_field::{ - ExtensionField, Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField, - extension::BinomialExtensionField, + extension::BinomialExtensionField, ExtensionField, Field, PrimeCharacteristicRing, + PrimeField64, TwoAdicField, }; use p3_fri::FriParameters as InnerFriParameters; #[cfg(not(feature = "cuda"))] @@ -28,10 +28,10 @@ use p3_maybe_rayon::prelude::*; use p3_merkle_tree::MerkleTreeMmcs; use p3_symmetric::{CompressionFunctionFromHasher, SerializingHasher}; -#[cfg(feature = "cuda")] -use crate::cuda::CudaDft; #[cfg(feature = "cuda")] use crate::cuda::pcs::CudaPcsDft; +#[cfg(feature = "cuda")] +use crate::cuda::CudaDft; pub type Val = Goldilocks; pub type PackedVal = ::Packing; @@ -854,10 +854,10 @@ impl StarkGenericConfig for GoldilocksBlake3Config { // resident, then allow that trace to become an eviction candidate for // later jobs. lookup_jobs.sort_unstable_by_key(|&(index, output, direct, graph, _)| { - let temporary = if self - .pcs - .mmcs - .is_matrix_cuda_resident(inputs[index].stage_1.0, inputs[index].stage_1.1) + let temporary = if inputs[index] + .stage_1 + .0 + .has_resident_with_trace(inputs[index].stage_1.1) { graph.unwrap_or(direct) } else { @@ -894,16 +894,27 @@ impl StarkGenericConfig for GoldilocksBlake3Config { let ext_w = pair(extension_generator * extension_generator)[0]; let lookup_pool = std::sync::Arc::new(cuda_host_pool("lookup-rows", cuda_lookup_worker_count())); - let evaluate = |input: &crate::config::LookupCommitInput<'_, Self>, cooperative: bool| { + // `graph_path` is decided by admission below, which budgets the graph + // kernel from the same residency predicate; the kernel must not be + // re-chosen here or it can run under the other path's budget. + let evaluate = |input: &crate::config::LookupCommitInput<'_, Self>, + graph_path: bool, + cooperative: bool| { let (height, num_lookups, multiplicities, args, arg_offsets) = input.lookup_values.cuda_parts(); let group_size = input.circuit.lookup_group_size.max(1); - let main = input.stage_1.0.resident_with_trace(input.stage_1.1); - let result = if let Some(main) = main.filter(|_| num_lookups != 0) { - let preprocessed = match input.preprocessed { - Some((data, index)) => Some(data.resident(index)?), - None => None, - }; + let main = graph_path.then(|| { + input + .stage_1 + .0 + .resident_with_trace(input.stage_1.1) + .expect("lookup graph admission requires a trace-backed resident LDE") + }); + let result = if let Some(main) = main { + let preprocessed = input.preprocessed.map(|(data, index)| { + data.resident(index) + .expect("lookup graph admission requires a resident preprocessed LDE") + }); crate::cuda::lookup_graph_lde_resident( &self.pcs.dft, &input.circuit.graph, @@ -967,22 +978,32 @@ impl StarkGenericConfig for GoldilocksBlake3Config { for (index, output_bytes, direct_temporary_bytes, graph_temporary_bytes, _) in lookup_jobs { let job_started = std::time::Instant::now(); let input = &inputs[index]; - let graph_path = self - .pcs - .mmcs - .is_matrix_cuda_resident(input.stage_1.0, input.stage_1.1) - && graph_temporary_bytes.is_some() - && input.preprocessed.is_none_or(|(data, matrix)| { - self.pcs.mmcs.is_matrix_cuda_resident(data, matrix) - }); - let temporary_bytes = if graph_path { - graph_temporary_bytes.unwrap() - } else { - direct_temporary_bytes + // The graph kernel reads the trace through `resident_with_trace`, + // which also serves spilled LDEs that regenerate from a generator + // or a retained host trace. Admit it under the same predicate so + // the budget matches the kernel that runs. + let admits_graph = || { + graph_temporary_bytes.is_some() + && input.stage_1.0.has_resident_with_trace(input.stage_1.1) + && input + .preprocessed + .is_none_or(|(data, matrix)| data.resident(matrix).is_some()) + }; + let target_for = |graph_path: bool| { + let temporary_bytes = if graph_path { + graph_temporary_bytes.unwrap() + } else { + direct_temporary_bytes + }; + ( + temporary_bytes, + output_bytes + .saturating_add(temporary_bytes) + .saturating_add(total_device_bytes / 64), + ) }; - let target = output_bytes - .saturating_add(temporary_bytes) - .saturating_add(total_device_bytes / 64); + let mut graph_path = admits_graph(); + let (mut temporary_bytes, mut target) = target_for(graph_path); if target > total_device_bytes { return None; } @@ -994,12 +1015,24 @@ impl StarkGenericConfig for GoldilocksBlake3Config { ); if free_bytes < target { // This circuit alone does not fit beside its resident trace. - // Spill it as a last resort and use the direct lookup-values - // path, which remains protocol-identical. + // Spill it as a last resort. If the spilled LDE can still + // regenerate its trace it stays on the graph path; otherwise + // re-budget for the direct lookup-values path, which remains + // protocol-identical. free_bytes = self.pcs .mmcs .ensure_device_headroom(input.stage_1.0, target, None, "lookup"); + graph_path = admits_graph(); + (temporary_bytes, target) = target_for(graph_path); + if free_bytes < target { + free_bytes = self.pcs.mmcs.ensure_device_headroom( + input.stage_1.0, + target, + None, + "lookup", + ); + } } if free_bytes < target { return None; @@ -1014,7 +1047,7 @@ impl StarkGenericConfig for GoldilocksBlake3Config { && output_bytes >= (8usize << 30) && num_lookups >= 64 && arg_offsets.last().copied().unwrap_or(0) >= 256; - results[index] = Some(evaluate(&inputs[index], cooperative)?); + results[index] = Some(evaluate(&inputs[index], graph_path, cooperative)?); if crate::cuda::memory_diagnostics_enabled() { eprintln!( "[multi-stark/cuda] lookup job {index} complete: {:.3}s", @@ -1147,7 +1180,7 @@ mod pcs_ref_gen { use super::*; use p3_commit::Mmcs as _; use p3_field::{ - BasedVectorSpace, PrimeCharacteristicRing, PrimeField64, batch_multiplicative_inverse, + batch_multiplicative_inverse, BasedVectorSpace, PrimeCharacteristicRing, PrimeField64, }; use p3_matrix::dense::RowMajorMatrix; use p3_symmetric::{CryptographicHasher, PseudoCompressionFunction}; From 89747d3a698be7206a1ebc9383cb393544fa21a9 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:37:27 +0000 Subject: [PATCH 08/34] cuda: let trace generators keep device state, released after the lookup job A generator may hold device memory that serves its tiles faster, such as seeds kept resident between a shard's commitment and the lookup pass that regenerates rows from them. `TraceGenerator::release_device` frees it; the lookup path calls it after each LDE's job, and admission calls it before any LDE is spilled or evicted, in the commit loop and in both headroom paths, since a cache is rebuilt from host seeds on demand while a spilled LDE has to be uploaded again. --- cuda/kernels.cu | 6 +++++ src/cuda/mmcs.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++ src/cuda/mod.rs | 14 ++++++++++++ src/cuda/witness.rs | 8 +++++++ src/types.rs | 1 + src/witness.rs | 6 +++++ 6 files changed, 90 insertions(+) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index e85ca48..a27e85a 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -2396,6 +2396,12 @@ extern "C" bool multi_stark_cuda_lde_has_generator(const void* handle) { return handle && static_cast(handle)->trace_writer; } +extern "C" void* multi_stark_cuda_lde_generator_context(const void* handle) { + if (!handle) return nullptr; + const auto* lde = static_cast(handle); + return lde->trace_writer ? lde->trace_context : nullptr; +} + extern "C" int multi_stark_cuda_prepare_lde_constants( int device_id,const uint64_t* inverse_twiddles,size_t inverse_count, const uint64_t* shift_powers,size_t height,const uint64_t* forward_twiddles, diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index 8765a0e..c871d1f 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -108,6 +108,38 @@ fn hybrid_resident_candidates( .collect() } +/// Frees the device caches of generator-backed resident LDEs, one at a time +/// until `target_free_bytes` is met, and returns the free bytes then. A cache +/// is the cheapest headroom there is: the generator rebuilds it from host +/// seeds on demand, whereas an evicted LDE must be materialized and uploaded +/// again. +fn release_generator_caches( + device_id: i32, + data: &CudaMmcsData, + protected_index: Option, + target_free_bytes: usize, +) -> usize { + let mut free_bytes = super::device_memory_info(device_id).0; + let CudaMmcsData::Hybrid { resident, .. } = data else { + return free_bytes; + }; + for (index, lde) in resident.iter().enumerate() { + if free_bytes >= target_free_bytes { + break; + } + if protected_index == Some(index) { + continue; + } + let Some(lde) = lde else { continue }; + if !lde.has_generator() { + continue; + } + lde.release_generator_device(); + free_bytes = super::device_memory_info(device_id).0; + } + free_bytes +} + fn evict_hybrid_resident(data: &CudaMmcsData, index: usize) { let CudaMmcsData::Hybrid { resident, @@ -126,6 +158,7 @@ fn evict_hybrid_resident(data: &CudaMmcsData, index: usize) { // SAFETY: admission transitions run between proving stages, when no CUDA // operation can access this LDE or its trace. unsafe { lde.release_values() }; + lde.release_generator_device(); super::witness::record_lde_spill( lde.device_id, lde.height() * lde.width() * size_of::(), @@ -673,6 +706,16 @@ impl CudaCommitMmcs for CudaMmcs { if free_bytes >= target_free_bytes { return free_bytes; } + free_bytes = + release_generator_caches(self.device_id, data, protected_index, target_free_bytes); + if free_bytes >= target_free_bytes { + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] {phase} admitted after releasing generator caches: free={free_bytes}" + ); + } + return free_bytes; + } let mut candidates = hybrid_resident_candidates(data, protected_index); while free_bytes < target_free_bytes && !candidates.is_empty() { let deficit_cells = target_free_bytes @@ -723,6 +766,18 @@ impl CudaCommitMmcs for CudaMmcs { if measured_free_bytes >= target_free_bytes { return measured_free_bytes; } + for data in data { + measured_free_bytes = + release_generator_caches(self.device_id, data, None, target_free_bytes); + if measured_free_bytes >= target_free_bytes { + if super::memory_diagnostics_enabled() { + eprintln!( + "[multi-stark/cuda] {phase} batch admitted after releasing generator caches: free={measured_free_bytes}" + ); + } + return measured_free_bytes; + } + } let initial_free_bytes = measured_free_bytes; let mut released_bytes = 0usize; let mut candidates = data diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 6a8c744..4a9b426 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -408,6 +408,7 @@ unsafe extern "C" { destroy: unsafe extern "C" fn(*mut c_void), ) -> i32; fn multi_stark_cuda_lde_has_generator(handle: *const c_void) -> bool; + fn multi_stark_cuda_lde_generator_context(handle: *const c_void) -> *mut c_void; } pub struct CudaLde { @@ -425,6 +426,19 @@ impl CudaLde { unsafe { multi_stark_cuda_lde_has_generator(self.raw_handle()) } } + /// Frees what the trace generator, if any, caches on this device. The + /// generator stays attached and keeps serving tiles from the host. + pub(crate) fn release_generator_device(&self) { + let context = unsafe { multi_stark_cuda_lde_generator_context(self.raw_handle()) }; + if context.is_null() { + return; + } + // SAFETY: the handle owns the boxed generator until it is destroyed, + // and the box holds exactly the `Generator` the LDE was created with. + let generator = unsafe { &*context.cast::() }; + generator.release_device(self.device_id); + } + pub(crate) const fn raw_handle(&self) -> *const c_void { self.handle.as_ptr() } diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index d13f318..d8aa0e8 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -83,6 +83,14 @@ pub(crate) fn commit( .saturating_mul(8) .saturating_mul((1 << blowup) + 1) .saturating_add(reserve); + // Generator caches go before any LDE spills: they are rebuilt from + // host seeds on demand, a spilled LDE is uploaded again. + for index in 0..resident.len() { + if device_memory_info(device).0 >= needed { + break; + } + resident[index].as_ref().unwrap().release_generator_device(); + } for index in 0..resident.len() { if device_memory_info(device).0 >= needed { break; diff --git a/src/types.rs b/src/types.rs index 266f13a..a3786a4 100644 --- a/src/types.rs +++ b/src/types.rs @@ -971,6 +971,7 @@ impl StarkGenericConfig for GoldilocksBlake3Config { // retained matrix remains owned by the prover data. if let Some(main) = main { unsafe { main.release_trace() }; + main.release_generator_device(); } result }; diff --git a/src/witness.rs b/src/witness.rs index c592ff0..98a0183 100644 --- a/src/witness.rs +++ b/src/witness.rs @@ -16,6 +16,12 @@ pub trait TraceGenerator: Send + Sync { /// Fill a device tile synchronously on its owning device and calling stream. #[cfg(feature = "cuda")] fn write_device_rows(&self, output: crate::cuda::DeviceTraceView<'_>) -> Result<(), String>; + + /// Free whatever the source keeps on `device_id` to serve tiles faster, + /// such as seeds left resident between a commitment and the lookup pass + /// that regenerates rows from them. Tiles must still be served afterwards. + #[cfg(feature = "cuda")] + fn release_device(&self, _device_id: i32) {} } #[derive(Clone)] From 2ec82faf12d08fe789e386c06606cf3c7a9f0ae1 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:01:03 +0000 Subject: [PATCH 09/34] cuda: spans for LDE shapes, commit sources and the FRI opening phases Every resident LDE construction records its kind, height, width and blowup; the commit loop records each source and the Merkle build; the streamed FRI opening records preparation, interpolation, observation, reduction, the per-round commitments and the query phase. A CUPTI profile joined to these spans attributes kernel time per transform shape and idle time per phase. --- src/cuda/mod.rs | 10 ++++++++++ src/cuda/pcs.rs | 39 +++++++++++++++++++++++++++------------ src/cuda/witness.rs | 12 ++++++++++++ 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 4a9b426..6b3529b 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -174,6 +174,8 @@ impl CudaDft { ) -> CudaLde { let height = matrix.height(); let width = matrix.width(); + let _span = + tracing::info_span!("cuda/lde", kind = "host", height, width, added_bits).entered(); Self::validate_dimensions(height, width); assert!(width > 0, "resident CUDA LDE requires at least one column"); let extended_height = height @@ -345,6 +347,8 @@ impl CudaDft { ) -> CudaLde { let height = generator.height(); let width = generator.width(); + let _span = + tracing::info_span!("cuda/lde", kind = "generated", height, width, added_bits).entered(); Self::validate_dimensions(height, width); let extended_height = height .checked_shl(added_bits.try_into().unwrap()) @@ -492,6 +496,7 @@ impl CudaLde { g_inv: Goldilocks, ext_w: Goldilocks, ) -> Self { + let _span = tracing::info_span!("cuda/fri_fold").entered(); let mut handle = core::ptr::null_mut(); let status = unsafe { multi_stark_cuda_fri_fold_resident( @@ -1207,6 +1212,7 @@ pub(crate) fn quotient_lde_mixed( quotient_degree: usize, log_blowup: usize, ) -> CudaLde { + let _span = tracing::info_span!("cuda/quotient_lde").entered(); quotient_lde_sources( dft, graph, @@ -1750,6 +1756,7 @@ pub(crate) fn lookup_lde_resident( ext_w: Goldilocks, log_blowup: usize, ) -> (CudaLde, [Goldilocks; 2]) { + let _span = tracing::info_span!("cuda/lookup_lde", path = "direct", height, num_lookups, group_size).entered(); assert!((1..=8).contains(&group_size)); assert_eq!(arg_offsets.len(), num_lookups + 1); assert_eq!(arg_offsets.first(), Some(&0)); @@ -1831,6 +1838,7 @@ pub(crate) fn lookup_lde_resident_partitioned( log_blowup: usize, cpu_deltas: impl Fn(core::ops::Range) -> Vec<[Goldilocks; 2]> + Sync, ) -> (CudaLde, [Goldilocks; 2]) { + let _span = tracing::info_span!("cuda/lookup_lde", path = "partitioned", height, num_lookups, group_size).entered(); assert!((1..=8).contains(&group_size)); assert_eq!(arg_offsets.len(), num_lookups + 1); assert_eq!(arg_offsets.first(), Some(&0)); @@ -2018,6 +2026,8 @@ pub(crate) fn lookup_graph_lde_resident( log_blowup: usize, ) -> Option<(CudaLde, [Goldilocks; 2])> { assert!((1..=8).contains(&group_size)); + let _span = + tracing::info_span!("cuda/lookup_lde", path = "graph", height, group_size).entered(); let (nodes, slot_count, lookups, args) = encode_lookup_nodes(graph)?; let num_lookups = lookups.len(); let groups = num_lookups.div_ceil(group_size.max(1)); diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 56ae00c..518a30c 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -341,10 +341,14 @@ where params.max_log_arity, ); let arity = 1 << log_arity; - let (commitment, round) = params.mmcs.commit_cuda_fri(codeword, log_arity); + let (commitment, round) = tracing::info_span!("stark/fri_round_commit") + .in_scope(|| params.mmcs.commit_cuda_fri(codeword, log_arity)); challenger.observe(commitment.clone()); commits.push(commitment); - commit_pow_witnesses.push(challenger.grind(params.commit_proof_of_work_bits)); + commit_pow_witnesses.push( + tracing::info_span!("stark/fri_commit_grind") + .in_scope(|| challenger.grind(params.commit_proof_of_work_bits)), + ); let beta: Challenge = challenger.sample_algebra_element(); let mut beta_step = beta; let betas = (0..log_arity) @@ -396,10 +400,13 @@ where for &log_arity in &log_arities { challenger.observe(Val::from_usize(log_arity)); } - let query_pow_witness = challenger.grind(params.query_proof_of_work_bits); + let _query_span = tracing::info_span!("stark/fri_queries").entered(); + let query_pow_witness = tracing::info_span!("stark/fri_query_grind") + .in_scope(|| challenger.grind(params.query_proof_of_work_bits)); let query_indices = iter::repeat_with(|| challenger.sample_bits(log_global_max_height)) .take(params.num_queries) .collect_vec(); + let input_openings_span = tracing::info_span!("stark/fri_input_openings").entered(); let input_openings = prover_data_with_opening_points .iter() .map(|(data, _)| { @@ -415,6 +422,7 @@ where } }) .collect_vec(); + drop(input_openings_span); let mut current_indices = query_indices; let commit_phase_openings = rounds .iter() @@ -429,8 +437,8 @@ where .iter() .map(|&index| index >> log_arity) .collect_vec(); - let (opened_rows, opening_proof) = - params.mmcs.open_cuda_fri_batch(round, &group_indices); + let (opened_rows, opening_proof) = tracing::info_span!("stark/fri_commit_phase_opening") + .in_scope(|| params.mmcs.open_cuda_fri_batch(round, &group_indices)); current_indices = group_indices; let sibling_values = positions .into_iter() @@ -1198,7 +1206,7 @@ where .iter() .all(|(data, _)| self.mmcs.is_cuda_resident(data)) { - let resident_rounds = debug_span!("cuda prepare resident rounds").in_scope(|| { + let resident_rounds = tracing::info_span!("stark/fri_prepare_rounds").in_scope(|| { commitment_data_with_opening_points .iter() .map(|(data, points)| (self.mmcs.resident_or_upload(data), points)) @@ -1212,7 +1220,7 @@ where .unwrap_or(0); let final_fri_height = self.fri.blowup() * self.fri.final_poly_len(); if resident_max_height > 1024 && resident_max_height > final_fri_height { - let _resident_guard = debug_span!("cuda resident fri").entered(); + let _resident_guard = tracing::info_span!("stark/fri_resident").entered(); let rounds = resident_rounds; let device_id = self.mmcs.cuda_device_id(); assert_eq!(>::DIMENSION, 2); @@ -1308,7 +1316,7 @@ where .collect_vec() }) .collect_vec(); - let interpolated = debug_span!("cuda interpolate openings") + let interpolated = tracing::info_span!("stark/fri_interpolate") .in_scope(|| workspace.interpolate(&interpolation_tasks, output_count, ext_w)); let all_opened_values = layouts .into_iter() @@ -1366,10 +1374,10 @@ where } } } - debug_span!("cuda reduce openings") + tracing::info_span!("stark/fri_reduce") .in_scope(|| workspace.reduce(&reduction_tasks, &alpha_pairs, ext_w)); let fri_input = reduced.into_iter().rev().flatten().collect_vec(); - let fri_proof = debug_span!("cuda prove fri").in_scope(|| { + let fri_proof = tracing::info_span!("stark/fri_prove").in_scope(|| { prove_fri_cuda_resident( &self.fri, fri_input, @@ -1415,7 +1423,8 @@ where .expect("No Matrices Supplied?"); let final_fri_height = self.fri.blowup() * self.fri.final_poly_len(); if cuda_max_height > 1024 && cuda_max_height > final_fri_height { - let _resident_guard = debug_span!("cuda streamed fri").entered(); + let _resident_guard = tracing::info_span!("stark/fri_streamed").entered(); + let prepare_span = tracing::info_span!("stark/fri_prepare").entered(); let phase_started = std::time::Instant::now(); let device_id = self.mmcs.cuda_device_id(); assert_eq!(>::DIMENSION, 2); @@ -1584,6 +1593,8 @@ where ); } + drop(prepare_span); + let interpolate_span = tracing::info_span!("stark/fri_interpolate").entered(); let interpolation_started = std::time::Instant::now(); let ((cpu_opened, cpu_interpolation_seconds), gpu_opened, gpu_interpolation_seconds) = std::thread::scope(|scope| { @@ -1713,6 +1724,8 @@ where ); } + drop(interpolate_span); + let observe_span = tracing::info_span!("stark/fri_observe_openings").entered(); for round in &all_opened_values { for matrix in round { for values in matrix { @@ -1749,6 +1762,8 @@ where }) .collect_vec(); + drop(observe_span); + let _reduce_span = tracing::info_span!("stark/fri_reduce").entered(); let reduction_started = std::time::Instant::now(); let ((cpu_reduced, cpu_reduction_seconds), mut gpu_reduced, gpu_reduction_seconds) = std::thread::scope(|scope| { @@ -1881,7 +1896,7 @@ where } let fri_input = gpu_reduced.into_iter().rev().flatten().collect_vec(); let folding_started = std::time::Instant::now(); - let fri_proof = debug_span!("cuda prove streamed fri").in_scope(|| { + let fri_proof = tracing::info_span!("stark/fri_prove").in_scope(|| { prove_fri_cuda_resident( &self.fri, fri_input, diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index d8aa0e8..9f282d5 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -105,6 +105,16 @@ pub(crate) fn commit( "generated trace commitment exceeds device admission; reduce the shard cell budget" ); let shift = Val::GENERATOR / domain.shift(); + let source_span = tracing::info_span!( + "stark/commit_source", + kind = match &source { + TraceSource::Host(_) => "host", + TraceSource::Generated(_) => "generated", + }, + height = source.height(), + width = source.width() + ) + .entered(); pcs.dft .prepare_coset_lde_constants(source.height(), blowup, shift); let (lde, trace) = match source { @@ -128,7 +138,9 @@ pub(crate) fn commit( resident.push(Some(lde)); host.push(spilled); retained.push(trace); + drop(source_span); } + let _merkle_span = tracing::info_span!("stark/commit_merkle").entered(); let needed = max_height.saturating_mul(96).saturating_add(reserve); for index in 0..resident.len() { if device_memory_info(device).0 >= needed { From a1fcbf58425a4e0d1836b6e3714278189d5dd8d8 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:57:01 +0000 Subject: [PATCH 10/34] cuda: cache the FRI coset per process; double-buffer staged uploads The streamed FRI opening rebuilt, bit-reversed and converted the coset of the largest LDE height for every shard and uploaded it again for the denominators: 0.32 s of host time per 2^26 shard. The bit-reversed coset of a smaller height is a prefix of a larger one's, so one vector per process serves every height, and the workspace takes it from the device constant cache instead of owning a copy. A staged upload copied each 64 MiB chunk into the pinned slot, transferred it and waited before the next. The slot is now a ring of 16 MiB chunks with per-chunk events, so the host copy of one chunk overlaps the transfers before it, and the host copy of a chunk is split across four pthreads since one thread fills pinned memory well below the PCIe rate. On the profiled Init claim the proof went from 44.9 s to 40.3 s: FRI opening 5.7 s to 2.9 s, host-trace commit idle 5.8 s to 4.2 s. --- cuda/kernels.cu | 87 ++++++++++++++++++++++++++++++++++++++++++------- src/cuda/pcs.rs | 66 +++++++++++++++++++++++++++++++------ 2 files changed, 132 insertions(+), 21 deletions(-) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index a27e85a..b0434e0 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -475,13 +475,61 @@ struct ResidentMixedMerkleTree { // copy costs one host memcpy per chunk. Slots are leased under a mutex and // callers wait for a free one, which bounds concurrent uploads to the slot // count (the PCIe link is shared anyway). +// A slot is a ring of chunks, each with the event of its last transfer, so +// the host copy of one chunk overlaps the DMA of the chunks before it. constexpr size_t UPLOAD_STAGING_SLOTS = 4; constexpr size_t UPLOAD_STAGING_BYTES = size_t(64) << 20; +constexpr size_t UPLOAD_STAGING_RING = 4; +constexpr size_t UPLOAD_STAGING_CHUNK = UPLOAD_STAGING_BYTES / UPLOAD_STAGING_RING; uint64_t* upload_staging[MAX_CUDA_DEVICES][UPLOAD_STAGING_SLOTS] = {{nullptr}}; +cudaEvent_t upload_staging_events[MAX_CUDA_DEVICES][UPLOAD_STAGING_SLOTS][UPLOAD_STAGING_RING] = {{{nullptr}}}; bool upload_staging_in_use[MAX_CUDA_DEVICES][UPLOAD_STAGING_SLOTS] = {{false}}; pthread_mutex_t upload_staging_mutex = PTHREAD_MUTEX_INITIALIZER; pthread_cond_t upload_staging_cond = PTHREAD_COND_INITIALIZER; +// One thread copies into pinned memory at a fraction of the PCIe rate, so a +// serial copy, not the DMA, would bound a staged upload; large chunks are +// copied by a few threads. +struct MemcpyPart { + unsigned char* destination; + const unsigned char* source; + size_t bytes; +}; + +void* memcpy_part(void* argument) { + const auto* part = static_cast(argument); + memcpy(part->destination, part->source, part->bytes); + return nullptr; +} + +// pthreads rather than std::thread: the executable is linked by the Lean +// toolchain without the C++ runtime library. +void parallel_memcpy(void* destination, const void* source, size_t bytes) { + constexpr size_t THREADS = 4; + constexpr size_t MIN_PARALLEL = size_t(8) << 20; + if (bytes < MIN_PARALLEL) { + memcpy(destination, source, bytes); + return; + } + const size_t part = (bytes + THREADS - 1) / THREADS; + auto* dst = static_cast(destination); + const auto* src = static_cast(source); + MemcpyPart parts[THREADS]; + pthread_t workers[THREADS]; + bool started[THREADS] = {false}; + for (size_t t = 1; t < THREADS; ++t) { + const size_t offset = t * part; + if (offset >= bytes) break; + parts[t] = {dst + offset, src + offset, std::min(part, bytes - offset)}; + started[t] = pthread_create(&workers[t], nullptr, memcpy_part, &parts[t]) == 0; + if (!started[t]) memcpy_part(&parts[t]); + } + memcpy(dst, src, std::min(part, bytes)); + for (size_t t = 1; t < THREADS; ++t) { + if (started[t]) pthread_join(workers[t], nullptr); + } +} + // Copies `bytes` of pageable host memory to the device on the per-thread // stream through a leased staging slot; returns with the copy complete. cudaError_t staged_upload(void* device, const void* host, size_t bytes) { @@ -503,24 +551,38 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { } in_use[slot] = true; cudaError_t status = cudaSuccess; + cudaEvent_t* events = upload_staging_events[device_index][slot]; if (slots[slot] == nullptr) { status = cudaMallocHost(reinterpret_cast(&slots[slot]), UPLOAD_STAGING_BYTES); if (status != cudaSuccess) slots[slot] = nullptr; + for (size_t i = 0; status == cudaSuccess && i < UPLOAD_STAGING_RING; ++i) { + status = cudaEventCreateWithFlags(&events[i], cudaEventDisableTiming); + } } pthread_mutex_unlock(&upload_staging_mutex); - uint64_t* staging = slots[slot]; + auto* staging = reinterpret_cast(slots[slot]); if (status == cudaSuccess) { const auto* source = static_cast(host); auto* target = static_cast(device); + size_t index = 0; for (size_t offset = 0; offset < bytes && status == cudaSuccess; - offset += UPLOAD_STAGING_BYTES) { - const size_t chunk = std::min(UPLOAD_STAGING_BYTES, bytes - offset); - memcpy(staging, source + offset, chunk); - status = cudaMemcpyAsync(target + offset, staging, chunk, - cudaMemcpyHostToDevice, cudaStreamPerThread); - if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); + offset += UPLOAD_STAGING_CHUNK, ++index) { + const size_t ring = index % UPLOAD_STAGING_RING; + unsigned char* buffer = staging + ring * UPLOAD_STAGING_CHUNK; + // The buffer is free once its previous transfer has landed. + if (index >= UPLOAD_STAGING_RING) status = cudaEventSynchronize(events[ring]); + const size_t chunk = std::min(UPLOAD_STAGING_CHUNK, bytes - offset); + if (status == cudaSuccess) { + parallel_memcpy(buffer, source + offset, chunk); + status = cudaMemcpyAsync(target + offset, buffer, chunk, + cudaMemcpyHostToDevice, cudaStreamPerThread); + } + if (status == cudaSuccess) status = cudaEventRecord(events[ring], cudaStreamPerThread); } + // The slot is handed to the next caller on release. + const cudaError_t synced = cudaStreamSynchronize(cudaStreamPerThread); + if (status == cudaSuccess) status = synced; } pthread_mutex_lock(&upload_staging_mutex); in_use[slot] = false; @@ -693,11 +755,13 @@ struct ResidentReducedOpening { }; struct ResidentFriWorkspace { - Ext2* inv_denoms=nullptr; uint64_t* coset=nullptr; Ext2* output=nullptr; + // The coset is a process-wide device constant shared by every workspace + // of its size; the workspace does not own it. + Ext2* inv_denoms=nullptr; const uint64_t* coset=nullptr; Ext2* output=nullptr; Ext2* alpha=nullptr; Ext2* interpolation_partials=nullptr; size_t inv_count=0; size_t coset_count=0; size_t output_capacity=0; size_t alpha_capacity=0,partial_capacity=0; - ~ResidentFriWorkspace(){if(inv_denoms)cudaFree(inv_denoms);if(coset)cudaFree(coset);if(output)cudaFree(output);if(alpha)cudaFree(alpha);if(interpolation_partials)cudaFree(interpolation_partials);} + ~ResidentFriWorkspace(){if(inv_denoms)cudaFree(inv_denoms);if(output)cudaFree(output);if(alpha)cudaFree(alpha);if(interpolation_partials)cudaFree(interpolation_partials);} }; struct InterpolationTask { @@ -3102,8 +3166,9 @@ extern "C" int multi_stark_cuda_fri_workspace_create(int device_id,void** handle cudaError_t status=cudaSetDevice(device_id);auto* w=new(std::nothrow) ResidentFriWorkspace;if(!w)return static_cast(cudaErrorMemoryAllocation); w->inv_count=inv_count;w->coset_count=coset_count; if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->inv_denoms),inv_count*sizeof(Ext2)); - if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->coset),coset_count*sizeof(uint64_t)); - if(status==cudaSuccess)status=cudaMemcpy(w->coset,coset,coset_count*sizeof(uint64_t),cudaMemcpyHostToDevice); + // One upload per device and size: every shard's opening indexes the same + // bit-reversed coset, 512 MiB at 2^26 rows. + if(status==cudaSuccess)status=cached_device_constants(device_id,coset,coset_count,5,coset[0],coset_count>1?coset[1]:0,&w->coset); uint64_t *norms=nullptr,*inverses=nullptr; if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norms),max_count*sizeof(uint64_t)); if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&inverses),max_count*sizeof(uint64_t)); diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 518a30c..4f9ae0a 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -263,6 +263,60 @@ pub struct CudaTwoAdicFriPcs { _phantom: PhantomData, } +/// The coset `GENERATOR * ` in bit-reversed order, as the opening's +/// denominators and interpolation index it, with at least `2^log_height` +/// elements. Element `i` of the bit-reversed coset of size `2^k` is the +/// shift times `g_k^{rev_k(i)}`, and `rev_{k+1}(i) = 2 rev_k(i)` for +/// `i < 2^k`, so every smaller coset is a prefix of a larger one: one vector +/// per process serves every height, and grows only when a larger height is +/// opened. Every shard of a proof opens on the same coset; before caching, +/// building it cost 0.32 s per shard at 2^26 rows. +fn bit_reversed_coset( + log_height: usize, +) -> std::sync::Arc> { + static CACHE: std::sync::Mutex>>> = + std::sync::Mutex::new(None); + if let Some(coset) = cache_at_least(&CACHE, 1 << log_height) { + return coset; + } + let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); + let generator: Goldilocks = TwoAdicField::two_adic_generator(log_height); + let shift = ::GENERATOR; + assert_eq!(to_gold(Val::two_adic_generator(log_height)), generator); + assert_eq!(to_gold(Val::GENERATOR), shift); + let size = 1usize << log_height; + let mut coset = vec![Goldilocks::ZERO; size]; + const CHUNK: usize = 1 << 14; + coset + .par_chunks_mut(CHUNK) + .enumerate() + .for_each(|(index, chunk)| { + let mut x: Goldilocks = shift * generator.exp_u64((index * CHUNK) as u64); + for value in chunk { + *value = x; + x *= generator; + } + }); + reverse_slice_index_bits(&mut coset); + let mut slot = CACHE.lock().unwrap(); + match &*slot { + Some(cached) if cached.len() >= size => std::sync::Arc::clone(cached), + _ => std::sync::Arc::clone(slot.insert(std::sync::Arc::new(coset))), + } +} + +fn cache_at_least( + cache: &std::sync::Mutex>>>, + size: usize, +) -> Option>> { + cache + .lock() + .unwrap() + .as_ref() + .filter(|coset| coset.len() >= size) + .map(std::sync::Arc::clone) +} + fn prove_fri_cuda_resident( params: &FriParameters, mut inputs: Vec, @@ -1247,10 +1301,7 @@ where .max() .unwrap(); let log_global_max_height = log2_strict_usize(global_max_height); - let coset_domain = - TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); - let mut coset: Vec = coset_domain.iter().collect(); - reverse_slice_index_bits(&mut coset); + let coset_gold = bit_reversed_coset::(log_global_max_height); let mut max_log: LinearMap = LinearMap::new(); for (ldes, points) in &rounds { for (lde, ps) in ldes.iter().zip(points.iter()) { @@ -1267,7 +1318,6 @@ where Challenge::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE]).unwrap(); let ext_w = to_pair(ext_x * ext_x)[0]; assert_eq!(ext_w, Goldilocks::from_u64(7)); - let coset_gold: Vec<_> = coset.iter().copied().map(to_gold).collect(); let mut inv_offsets = LinearMap::new(); let mut inverse_points = Vec::new(); let mut inverse_counts = Vec::new(); @@ -1442,10 +1492,7 @@ where .expect("quadratic extension element") }; let log_global_max_height = log2_strict_usize(cuda_max_height); - let coset_domain = - TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); - let mut coset: Vec = coset_domain.iter().collect(); - reverse_slice_index_bits(&mut coset); + let coset_gold = bit_reversed_coset::(log_global_max_height); let mut max_log: LinearMap = LinearMap::new(); for ((_, points), round_dimensions) in commitment_data_with_opening_points @@ -1468,7 +1515,6 @@ where .expect("quadratic extension generator"); let ext_w = to_pair(ext_x * ext_x)[0]; assert_eq!(ext_w, Goldilocks::from_u64(7)); - let coset_gold: Vec<_> = coset.iter().copied().map(to_gold).collect(); let mut inv_offsets = LinearMap::new(); let mut inverse_points = Vec::new(); let mut inverse_counts = Vec::new(); From 82c6c9ba6ad090963f41e29dffd7d7a62c82a3b2 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:36:59 +0000 Subject: [PATCH 11/34] cuda: collect lightweight prover operation metrics Add opt-in process/device counters for staged uploads, coset caching, memory samples and NTT shapes. Emit cumulative snapshots at piece boundaries and annotate coarse DFT/LDE operations and admission decisions. Metrics introduce no CUDA events or synchronization. Validated with the focused aiur release check using cuda-trace-codegen and the local backend override. --- build.rs | 1 + cuda/kernels.cu | 67 +++++++++++++++++++-------- cuda/metrics.cuh | 56 +++++++++++++++++++++++ src/batch.rs | 4 ++ src/cuda/metrics.rs | 34 ++++++++++++++ src/cuda/mod.rs | 109 ++++++++++++++++++++++++++++++++++++++++---- src/cuda/pcs.rs | 21 +++++---- src/cuda/witness.rs | 1 + src/types.rs | 16 +++++-- 9 files changed, 268 insertions(+), 41 deletions(-) create mode 100644 cuda/metrics.cuh create mode 100644 src/cuda/metrics.rs diff --git a/build.rs b/build.rs index 03dc72c..b6814f7 100644 --- a/build.rs +++ b/build.rs @@ -9,6 +9,7 @@ use std::path::{Path, PathBuf}; use std::process::Command; fn main() { + println!("cargo:rerun-if-changed=cuda/metrics.cuh"); println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=cuda/kernels.cu"); println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); diff --git a/cuda/kernels.cu b/cuda/kernels.cu index b0434e0..2053469 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -17,6 +17,7 @@ #include #include #include +#include "metrics.cuh" namespace { @@ -57,13 +58,22 @@ cudaError_t cached_device_constants(int device,const uint64_t* host,size_t count while(__sync_lock_test_and_set(&constant_cache_lock,1)){} for(auto* entry=constant_cache;entry;entry=entry->next)if(entry->device==device&& entry->count==count&&entry->kind==kind&&entry->key0==key0&&entry->key1==key1){ + if (kind == 5 && multi_stark_metrics::enabled()) multi_stark_metrics::add(device, multi_stark_metrics::CosetHits, 1); *output=entry->values;__sync_lock_release(&constant_cache_lock);return cudaSuccess;} auto* entry=new(std::nothrow) ConstantCacheEntry; if(!entry){__sync_lock_release(&constant_cache_lock);return cudaErrorMemoryAllocation;} entry->device=device;entry->count=count;entry->kind=kind;entry->key0=key0;entry->key1=key1; cudaError_t status=persistent_malloc(reinterpret_cast(&entry->values),count*sizeof(uint64_t)); if(status==cudaSuccess)status=cudaMemcpy(entry->values,host,count*sizeof(uint64_t),cudaMemcpyHostToDevice); - if(status==cudaSuccess){*output=entry->values;entry->next=constant_cache;constant_cache=entry;} + if(status==cudaSuccess){ + if (multi_stark_metrics::enabled()) { + multi_stark_metrics::add(device, multi_stark_metrics::ConstantBytes, count * sizeof(uint64_t)); + if (kind == 5) { + multi_stark_metrics::add(device, multi_stark_metrics::CosetMisses, 1); + multi_stark_metrics::add(device, multi_stark_metrics::CosetUploadedBytes, count * sizeof(uint64_t)); + } + } + *output=entry->values;entry->next=constant_cache;constant_cache=entry;} else{persistent_free(entry->values);delete entry;} __sync_lock_release(&constant_cache_lock);return status; } @@ -534,6 +544,8 @@ void parallel_memcpy(void* destination, const void* source, size_t bytes) { // stream through a leased staging slot; returns with the copy complete. cudaError_t staged_upload(void* device, const void* host, size_t bytes) { const int device_index = current_device_index(); + cudaError_t status = cudaSuccess; + multi_stark_metrics::Upload metrics(device_index, bytes, status); uint64_t** slots = upload_staging[device_index]; bool* in_use = upload_staging_in_use[device_index]; pthread_mutex_lock(&upload_staging_mutex); @@ -550,7 +562,6 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { } } in_use[slot] = true; - cudaError_t status = cudaSuccess; cudaEvent_t* events = upload_staging_events[device_index][slot]; if (slots[slot] == nullptr) { status = cudaMallocHost(reinterpret_cast(&slots[slot]), @@ -575,6 +586,7 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { const size_t chunk = std::min(UPLOAD_STAGING_CHUNK, bytes - offset); if (status == cudaSuccess) { parallel_memcpy(buffer, source + offset, chunk); + ++metrics.chunks; status = cudaMemcpyAsync(target + offset, buffer, chunk, cudaMemcpyHostToDevice, cudaStreamPerThread); } @@ -1715,8 +1727,14 @@ __global__ void radix2_dif_tail_tiled(uint64_t* values,size_t height,size_t widt } } -cudaError_t launch_dif(uint64_t* values, size_t height, size_t width, +cudaError_t launch_dif(int device_id, uint64_t* values, size_t height, size_t width, const uint64_t* twiddles) { + if (multi_stark_metrics::enabled() && height > 1 && width > 0) { + const unsigned log = strict_log2(height); + const unsigned bucket = width == 1 ? 0 : width == 2 ? 1 : width < 8 ? 2 : 3; + if (log <= 32) multi_stark_metrics::add(device_id, + multi_stark_metrics::NTT_OFFSET + log * 4 + bucket, 1); + } if (height <= 1 || width == 0) { return cudaSuccess; } @@ -2248,7 +2266,7 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, status = copy_to_device(device_twiddles, twiddles, height / 2); } if (status == cudaSuccess) { - status = launch_dif(device_values.get(), height, width, device_twiddles.get()); + status = launch_dif(device_id, device_values.get(), height, width, device_twiddles.get()); } if (status == cudaSuccess) { status = copy_to_host(values, device_values, elements); @@ -2307,7 +2325,7 @@ extern "C" int multi_stark_cuda_coset_lde_batch( extended_height / 2); } if (status == cudaSuccess) { - status = launch_dif(device_values.get(), height, width, + status = launch_dif(device_id, device_values.get(), height, width, device_inverse_twiddles.get()); } if (status == cudaSuccess) { @@ -2318,7 +2336,7 @@ extern "C" int multi_stark_cuda_coset_lde_batch( status = cudaGetLastError(); } if (status == cudaSuccess) { - status = launch_dif(device_values.get(), extended_height, width, + status = launch_dif(device_id, device_values.get(), extended_height, width, device_forward_twiddles.get()); } if (status == cudaSuccess) { @@ -2410,7 +2428,7 @@ static int coset_lde_create( status = cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&device_forward_twiddles); } if (status == cudaSuccess) { - status = launch_dif(lde->values, height, width,device_inverse_twiddles); + status = launch_dif(device_id, lde->values, height, width,device_inverse_twiddles); } if (status == cudaSuccess) { bit_reverse_scale_and_shift<<>>( @@ -2419,7 +2437,7 @@ static int coset_lde_create( status = cudaGetLastError(); } if (status == cudaSuccess) { - status = launch_dif(lde->values, extended_height, width,device_forward_twiddles); + status = launch_dif(device_id, lde->values, extended_height, width,device_forward_twiddles); } // Normalization and every DIF butterfly produce canonical field values, // including the height-one case. A final reduction pass is redundant; @@ -2888,7 +2906,7 @@ extern "C" int multi_stark_cuda_quotient_lde( dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=launch_dif(device_id, quotient,quotient_size,2,device_quotient_twiddles); ResidentLde* lde = nullptr; if(status==cudaSuccess) { @@ -2905,7 +2923,7 @@ extern "C" int multi_stark_cuda_quotient_lde( quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,lde_height,width,device_lde_twiddles); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess) { *output_handle=lde; @@ -3117,7 +3135,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( } for(size_t i=0;i<2;++i)if(status==cudaSuccess&&stream_busy[i])status=cudaStreamSynchronize(streams[i]); - if(status==cudaSuccess)status=launch_dif(quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=launch_dif(device_id, quotient,quotient_size,2,device_quotient_twiddles); ResidentLde* lde=nullptr; if(status==cudaSuccess)status=create_resident_lde(&lde); if(status==cudaSuccess){lde->height=lde_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t));} @@ -3127,7 +3145,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( lde->values,quotient,device_weights,quotient_size,trace_height,quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,lde_height,width,device_lde_twiddles); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else if(lde)destroy_resident_lde(lde); for(size_t i=0;i<2;++i){ @@ -3417,9 +3435,9 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); - if(status==cudaSuccess)status=launch_dif(lde->values,height,width,dit); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,height,width,dit); if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} - if(status==cudaSuccess)status=launch_dif(lde->values,extended_height,width,dft); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,extended_height,width,dft); if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); @@ -3483,9 +3501,9 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); - if(status==cudaSuccess)status=launch_dif(lde->values,height,width,dit); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,height,width,dit); if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} - if(status==cudaSuccess)status=launch_dif(lde->values,extended_height,width,dft); + if(status==cudaSuccess)status=launch_dif(device_id, lde->values,extended_height,width,dft); if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, @@ -3675,7 +3693,7 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( &forward); const size_t width = 2 * pending->slots; if (status == cudaSuccess) - status = launch_dif(pending->lde->values, pending->height, width, inverse); + status = launch_dif(device_id, pending->lde->values, pending->height, width, inverse); if (status == cudaSuccess) { bit_reverse_scale_and_shift<<height * width), THREADS>>>( pending->lde->values, pending->height, width, @@ -3683,7 +3701,7 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( status = cudaGetLastError(); } if (status == cudaSuccess) - status = launch_dif(pending->lde->values, pending->extended_height, + status = launch_dif(device_id, pending->lde->values, pending->extended_height, width, forward); if (status == cudaSuccess) { canonicalize_goldilocks<<extended_height * width), THREADS>>>( @@ -4540,6 +4558,11 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, cudaError_t status = cudaSetDevice(device_id); if (status == cudaSuccess) retain_default_pool(device_id); if (status == cudaSuccess) status = cudaMemGetInfo(free_bytes, total_bytes); + if (status == cudaSuccess && multi_stark_metrics::enabled()) { + multi_stark_metrics::sample(device_id, multi_stark_metrics::DriverFreeBytes, *free_bytes); + multi_stark_metrics::sample(device_id, multi_stark_metrics::TotalBytes, *total_bytes); + multi_stark_metrics::add(device_id, multi_stark_metrics::MemorySamples, 1); + } // cudaMemGetInfo excludes pages retained by cudaMallocAsync's default // pool, even though subsequent stream allocations can reuse them. Treat // the unused part of that pool as available for admission decisions; using @@ -4576,3 +4599,11 @@ extern "C" int multi_stark_cuda_generate_trace_rows(int device, void* context, if (tile) cudaFree(tile); return static_cast(status); } + +extern "C" void multi_stark_cuda_metrics_snapshot(uint64_t* output, size_t count) { + if (!output || count != multi_stark_metrics::DEVICES * multi_stark_metrics::WORDS) return; + for (size_t device = 0; device < multi_stark_metrics::DEVICES; ++device) + for (size_t key = 0; key < multi_stark_metrics::WORDS; ++key) + output[device * multi_stark_metrics::WORDS + key] = + multi_stark_metrics::counters[device][key].load(std::memory_order_relaxed); +} diff --git a/cuda/metrics.cuh b/cuda/metrics.cuh new file mode 100644 index 0000000..57719f2 --- /dev/null +++ b/cuda/metrics.cuh @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +#pragma once +#include +#include +#include +#include +#include +#include + +namespace multi_stark_metrics { +constexpr size_t DEVICES = 64; +constexpr size_t NTT_OFFSET = 16; +constexpr size_t WORDS = NTT_OFFSET + 33 * 4; +enum Counter : size_t { + UploadCalls, UploadRequestedBytes, UploadChunks, UploadFailures, UploadHostNs, + CosetHits, CosetMisses, CosetUploadedBytes, ConstantBytes, + DriverFreeBytes, TotalBytes, MemorySamples +}; +std::atomic counters[DEVICES][WORDS]{}; +inline bool enabled() { + static const bool value = std::getenv("AIUR_METRICS") != nullptr; + return value; +} +inline void add(int device, size_t key, uint64_t value) { + if (device >= 0 && device < int(DEVICES)) + counters[device][key].fetch_add(value, std::memory_order_relaxed); +} +inline void sample(int device, size_t key, uint64_t value) { + if (device >= 0 && device < int(DEVICES)) + counters[device][key].store(value, std::memory_order_relaxed); +} +inline uint64_t host_ns() { + timespec ts{}; + clock_gettime(CLOCK_MONOTONIC, &ts); + return uint64_t(ts.tv_sec) * 1000000000 + uint64_t(ts.tv_nsec); +} +struct Upload { + bool active = enabled(); + int device; + size_t bytes; + uint64_t chunks = 0; + cudaError_t& status; + uint64_t start = 0; + Upload(int d, size_t b, cudaError_t& s) : device(d), bytes(b), status(s) { + if (active) start = host_ns(); + } + ~Upload() { + if (!active) return; + add(device, UploadCalls, 1); + add(device, UploadRequestedBytes, bytes); + add(device, UploadChunks, chunks); + add(device, UploadFailures, status != cudaSuccess); + add(device, UploadHostNs, host_ns() - start); + } +}; +} diff --git a/src/batch.rs b/src/batch.rs index 54b6da7..88504ec 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -544,6 +544,8 @@ where Retention::Regenerate => None, }); barrier.claims.push(claims); + #[cfg(feature = "cuda")] + crate::cuda::metrics::emit_snapshot(); }); barrier } @@ -608,6 +610,8 @@ where _ => unreachable!("a shard is either retained or rebuilt"), }; proofs.push(self.prove_batch_shard(key, stage_1, &claims, &preamble, shard)); + #[cfg(feature = "cuda")] + crate::cuda::metrics::emit_snapshot(); }); BatchProof { preamble, proofs } } diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs new file mode 100644 index 0000000..56a7dea --- /dev/null +++ b/src/cuda/metrics.rs @@ -0,0 +1,34 @@ +//! CPU-only snapshots of counters accumulated at CUDA operation boundaries. + +pub(crate) fn emit_snapshot() { + if !tracing::enabled!(target: "prover_metrics", tracing::Level::INFO) { + return; + } + const WORDS: usize = 16 + 33 * 4; + let mut values = vec![0u64; 64 * WORDS]; + unsafe { + multi_stark_cuda_metrics_snapshot(values.as_mut_ptr(), values.len()); + } + for (device, c) in values.chunks_exact(WORDS).enumerate() { + if c.iter().all(|&v| v == 0) { + continue; + } + tracing::info!(target: "prover_metrics", metric = "cuda_device_snapshot", device, + scope = "process_device_cumulative", upload_calls = c[0], upload_requested_bytes = c[1], + upload_chunks = c[2], upload_failures = c[3], upload_host_ns = c[4], + coset_hits = c[5], coset_misses = c[6], coset_uploaded_bytes = c[7], constant_bytes = c[8], + last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11]); + for (shape, &count) in c[16..].iter().enumerate() { + if count == 0 { + continue; + } + tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, + scope = "process_device_cumulative", backend = "legacy", log_height = shape / 4, + width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); + } + } +} + +unsafe extern "C" { + fn multi_stark_cuda_metrics_snapshot(output: *mut u64, count: usize); +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 6b3529b..8d16ea2 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -5,6 +5,7 @@ //! lookup traces, quotient evaluations, and FRI codewords on the selected GPU. //! All public protocol types and serialized proofs remain unchanged. +pub(crate) mod metrics; pub(crate) mod mmcs; #[doc(hidden)] pub mod pcs; @@ -174,8 +175,16 @@ impl CudaDft { ) -> CudaLde { let height = matrix.height(); let width = matrix.width(); - let _span = - tracing::info_span!("cuda/lde", kind = "host", height, width, added_bits).entered(); + let _span = tracing::info_span!( + "cuda/lde", + kind = "host", + backend = "legacy", + device = self.device_id, + height, + width, + added_bits + ) + .entered(); Self::validate_dimensions(height, width); assert!(width > 0, "resident CUDA LDE requires at least one column"); let extended_height = height @@ -347,8 +356,16 @@ impl CudaDft { ) -> CudaLde { let height = generator.height(); let width = generator.width(); - let _span = - tracing::info_span!("cuda/lde", kind = "generated", height, width, added_bits).entered(); + let _span = tracing::info_span!( + "cuda/lde", + kind = "generated", + backend = "legacy", + device = self.device_id, + height, + width, + added_bits + ) + .entered(); Self::validate_dimensions(height, width); let extended_height = height .checked_shl(added_bits.try_into().unwrap()) @@ -1212,7 +1229,15 @@ pub(crate) fn quotient_lde_mixed( quotient_degree: usize, log_blowup: usize, ) -> CudaLde { - let _span = tracing::info_span!("cuda/quotient_lde").entered(); + let _span = tracing::info_span!( + "cuda/quotient_lde", + backend = "legacy", + device = dft.device_id, + quotient_size, + quotient_degree, + added_bits = log_blowup + ) + .entered(); quotient_lde_sources( dft, graph, @@ -1756,7 +1781,17 @@ pub(crate) fn lookup_lde_resident( ext_w: Goldilocks, log_blowup: usize, ) -> (CudaLde, [Goldilocks; 2]) { - let _span = tracing::info_span!("cuda/lookup_lde", path = "direct", height, num_lookups, group_size).entered(); + let _span = tracing::info_span!( + "cuda/lookup_lde", + path = "direct", + backend = "legacy", + device = dft.device_id, + height, + num_lookups, + group_size, + added_bits = log_blowup + ) + .entered(); assert!((1..=8).contains(&group_size)); assert_eq!(arg_offsets.len(), num_lookups + 1); assert_eq!(arg_offsets.first(), Some(&0)); @@ -1838,7 +1873,17 @@ pub(crate) fn lookup_lde_resident_partitioned( log_blowup: usize, cpu_deltas: impl Fn(core::ops::Range) -> Vec<[Goldilocks; 2]> + Sync, ) -> (CudaLde, [Goldilocks; 2]) { - let _span = tracing::info_span!("cuda/lookup_lde", path = "partitioned", height, num_lookups, group_size).entered(); + let _span = tracing::info_span!( + "cuda/lookup_lde", + path = "partitioned", + backend = "legacy", + device = dft.device_id, + height, + num_lookups, + group_size, + added_bits = log_blowup + ) + .entered(); assert!((1..=8).contains(&group_size)); assert_eq!(arg_offsets.len(), num_lookups + 1); assert_eq!(arg_offsets.first(), Some(&0)); @@ -2026,8 +2071,17 @@ pub(crate) fn lookup_graph_lde_resident( log_blowup: usize, ) -> Option<(CudaLde, [Goldilocks; 2])> { assert!((1..=8).contains(&group_size)); - let _span = - tracing::info_span!("cuda/lookup_lde", path = "graph", height, group_size).entered(); + let _span = tracing::info_span!( + "cuda/lookup_lde", + path = "graph", + backend = "legacy", + device = dft.device_id, + height, + num_lookups = graph.lookups.len(), + group_size, + added_bits = log_blowup + ) + .entered(); let (nodes, slot_count, lookups, args) = encode_lookup_nodes(graph)?; let num_lookups = lookups.len(); let groups = num_lookups.div_ceil(group_size.max(1)); @@ -2130,6 +2184,20 @@ impl TwoAdicSubgroupDft for CudaDft { let height = matrix.height(); let width = matrix.width(); Self::validate_dimensions(height, width); + let _span = tracing::info_span!( + "cuda/dft_batch", + device = self.device_id, + height, + width, + backend = if height == 1 || width == 0 { + "noop" + } else if Self::use_cuda_dft(height, width) { + "legacy" + } else { + "cpu" + } + ) + .entered(); if height == 1 || width == 0 { return BitReversalPerm::new_view(matrix); } @@ -2171,6 +2239,21 @@ impl TwoAdicSubgroupDft for CudaDft { .checked_shl(u32::try_from(added_bits).expect("LDE blowup exceeds u32")) .expect("LDE height overflows usize"); Self::validate_dimensions(extended_height, width); + let _span = tracing::info_span!( + "cuda/coset_lde_batch", + device = self.device_id, + height, + width, + added_bits, + backend = if width == 0 { + "noop" + } else if height > 1 && Self::use_cuda_coset_lde(extended_height, width) { + "legacy" + } else { + "cpu" + } + ) + .entered(); if width == 0 { return BitReversalPerm::new_view(RowMajorMatrix::new(Vec::new(), width)); @@ -2702,7 +2785,13 @@ impl CudaMixedMerkleTree { .collect(); check_cuda( status, - &format!("hybrid Merkle tree creation over (height, width) {dims:?}, host digest groups at heights {:?}", host_digest_groups.iter().map(|(h, _)| *h).collect::>()), + &format!( + "hybrid Merkle tree creation over (height, width) {dims:?}, host digest groups at heights {:?}", + host_digest_groups + .iter() + .map(|(h, _)| *h) + .collect::>() + ), ); } let row_count = heights.into_iter().max().unwrap(); diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 4f9ae0a..7f64d94 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -277,8 +277,11 @@ fn bit_reversed_coset( static CACHE: std::sync::Mutex>>> = std::sync::Mutex::new(None); if let Some(coset) = cache_at_least(&CACHE, 1 << log_height) { + tracing::info!(target: "prover_metrics", metric = "host_coset", action = "hit", + log_height, cached_bytes = coset.len() * 8); return coset; } + tracing::info!(target: "prover_metrics", metric = "host_coset", action = "build", log_height); let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); let generator: Goldilocks = TwoAdicField::two_adic_generator(log_height); let shift = ::GENERATOR; @@ -491,8 +494,9 @@ where .iter() .map(|&index| index >> log_arity) .collect_vec(); - let (opened_rows, opening_proof) = tracing::info_span!("stark/fri_commit_phase_opening") - .in_scope(|| params.mmcs.open_cuda_fri_batch(round, &group_indices)); + let (opened_rows, opening_proof) = + tracing::info_span!("stark/fri_commit_phase_opening") + .in_scope(|| params.mmcs.open_cuda_fri_batch(round, &group_indices)); current_indices = group_indices; let sibling_values = positions .into_iter() @@ -669,12 +673,13 @@ where // Height groups wider than the device leaf kernel can hash are // neither durable nor transient: their LDEs and digests come // from the host, whatever the memory budget says. - let wide_heights = super::mmcs::host_hashed_heights( - evaluations.iter().map(|(_, matrix)| p3_matrix::Dimensions { - width: matrix.width(), - height: matrix.height(), - }), - ); + let wide_heights = + super::mmcs::host_hashed_heights(evaluations.iter().map(|(_, matrix)| { + p3_matrix::Dimensions { + width: matrix.width(), + height: matrix.height(), + } + })); for height in &wide_heights { height_groups.remove(height); } diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index 9f282d5..e1eb4f0 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -15,6 +15,7 @@ use crate::witness::TraceSource; pub(crate) fn record_lde_spill(device: i32, bytes: usize) { tracing::debug!(device, bytes, "spilled active LDE"); + tracing::info!(target: "prover_metrics", metric = "lde_spill", device, bytes); } fn spill_lde(lde: &CudaLde) -> RowMajorMatrix { diff --git a/src/types.rs b/src/types.rs index a3786a4..9b74931 100644 --- a/src/types.rs +++ b/src/types.rs @@ -15,8 +15,8 @@ use p3_commit::{ExtensionMmcs, Pcs as PcsTrait}; use p3_dft::Radix2DitParallel; use p3_field::BasedVectorSpace; use p3_field::{ - extension::BinomialExtensionField, ExtensionField, Field, PrimeCharacteristicRing, - PrimeField64, TwoAdicField, + ExtensionField, Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField, + extension::BinomialExtensionField, }; use p3_fri::FriParameters as InnerFriParameters; #[cfg(not(feature = "cuda"))] @@ -28,10 +28,10 @@ use p3_maybe_rayon::prelude::*; use p3_merkle_tree::MerkleTreeMmcs; use p3_symmetric::{CompressionFunctionFromHasher, SerializingHasher}; -#[cfg(feature = "cuda")] -use crate::cuda::pcs::CudaPcsDft; #[cfg(feature = "cuda")] use crate::cuda::CudaDft; +#[cfg(feature = "cuda")] +use crate::cuda::pcs::CudaPcsDft; pub type Val = Goldilocks; pub type PackedVal = ::Packing; @@ -1006,6 +1006,9 @@ impl StarkGenericConfig for GoldilocksBlake3Config { let mut graph_path = admits_graph(); let (mut temporary_bytes, mut target) = target_for(graph_path); if target > total_device_bytes { + tracing::info!(target: "prover_metrics", metric = "lookup_admission", + job = index, graph_path, admitted = false, reason = "device_capacity", + output_bytes, temporary_bytes, target_bytes = target); return None; } let mut free_bytes = self.pcs.mmcs.ensure_device_headroom( @@ -1035,6 +1038,9 @@ impl StarkGenericConfig for GoldilocksBlake3Config { ); } } + tracing::info!(target: "prover_metrics", metric = "lookup_admission", + job = index, graph_path, admitted = free_bytes >= target, + output_bytes, temporary_bytes, target_bytes = target, free_bytes); if free_bytes < target { return None; } @@ -1181,7 +1187,7 @@ mod pcs_ref_gen { use super::*; use p3_commit::Mmcs as _; use p3_field::{ - batch_multiplicative_inverse, BasedVectorSpace, PrimeCharacteristicRing, PrimeField64, + BasedVectorSpace, PrimeCharacteristicRing, PrimeField64, batch_multiplicative_inverse, }; use p3_matrix::dense::RowMajorMatrix; use p3_symmetric::{CryptographicHasher, PseudoCompressionFunction}; From 1691727ec86fbecbc121a6536db4090bbe7b281c Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:11:42 +0000 Subject: [PATCH 12/34] cuda: sppark's Goldilocks NTT as an optional comparison backend `cuda-sppark` adds the pinned sppark crate and compiles a device-pointer adapter with the upstream GPU runtime into the kernel archive, in their own nvcc step because that runtime needs libstdc++'s ``, which the kernels' host flags would hide. Each transform runs on an upstream stream fenced to the caller's per-thread stream by events, so the completion contract matches the first-party kernels. The contract tests pin what the prover will rely on: the natural forward transform is the CPU DFT, the reversed orders are its bit-reversal, the inverse is normalized upstream so an inverse-forward pair needs no scaling, the coset variant uses the field generator, canonical extremes come back canonical, and the compiled domain limit covers 2^26 rows. Upstream takes canonical words only; representatives at or above the modulus transform to different values. --- Cargo.lock | 139 +++++++++++++++++++++-- Cargo.toml | 12 ++ build.rs | 47 ++++++++ cuda/sppark_ntt.cu | 107 ++++++++++++++++++ src/cuda/mod.rs | 2 + src/cuda/sppark.rs | 267 +++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 566 insertions(+), 8 deletions(-) create mode 100644 cuda/sppark_ntt.cu create mode 100644 src/cuda/sppark.rs diff --git a/Cargo.lock b/Cargo.lock index 7c13591..bb3611f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -267,7 +267,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -320,6 +320,15 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "is-terminal" version = "0.4.17" @@ -328,7 +337,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -392,6 +401,12 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "linux-raw-sys" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -466,6 +481,7 @@ dependencies = [ "rand", "rayon", "serde", + "sppark", "tracing", "tracing-subscriber", "tracing-texray", @@ -477,7 +493,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -968,6 +984,19 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "rustix" +version = "0.38.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys 0.4.15", + "windows-sys 0.59.0", +] + [[package]] name = "rustix" version = "1.1.4" @@ -977,8 +1006,8 @@ dependencies = [ "bitflags", "errno", "libc", - "linux-raw-sys", - "windows-sys", + "linux-raw-sys 0.12.1", + "windows-sys 0.61.2", ] [[package]] @@ -1081,6 +1110,15 @@ dependencies = [ "lock_api", ] +[[package]] +name = "sppark" +version = "0.1.15" +source = "git+https://github.com/supranational/sppark?rev=17278d74295392f9813f009300b257a688422b7a#17278d74295392f9813f009300b257a688422b7a" +dependencies = [ + "cc", + "which", +] + [[package]] name = "syn" version = "2.0.117" @@ -1098,8 +1136,8 @@ version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ - "rustix", - "windows-sys", + "rustix 1.1.4", + "windows-sys 0.61.2", ] [[package]] @@ -1312,13 +1350,25 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix 0.38.44", +] + [[package]] name = "winapi-util" version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1336,6 +1386,15 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -1345,6 +1404,70 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "zerocopy" version = "0.8.48" diff --git a/Cargo.toml b/Cargo.toml index 3455901..aad8291 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,14 @@ p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89 p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } +# sppark's Goldilocks NTT, for the comparison backend behind `cuda-sppark`. +# Pinned; the adapter compiles against the headers its build script exports. +[dependencies.sppark] +git = "https://github.com/supranational/sppark" +rev = "17278d74295392f9813f009300b257a688422b7a" +optional = true +features = ["cuda"] + [dev-dependencies] criterion = "0.5" p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } @@ -50,6 +58,10 @@ parallel = ["p3-maybe-rayon/parallel"] # requires a CUDA toolkit at build time and an NVIDIA GPU at runtime; the # default CPU build never invokes nvcc or links the CUDA runtime. cuda = ["dep:itertools", "dep:rayon"] +# Adds sppark's NTT/LDE engine as a comparison backend for the transforms. +# Requires `cuda`; the first-party kernels remain the default until the +# comparison is accepted. +cuda-sppark = ["cuda", "dep:sppark"] # Similar to `release`, but preserves debug info [profile.dev-ci] diff --git a/build.rs b/build.rs index b6814f7..dfb41c3 100644 --- a/build.rs +++ b/build.rs @@ -13,6 +13,7 @@ fn main() { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=cuda/kernels.cu"); println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); + println!("cargo:rerun-if-changed=cuda/sppark_ntt.cu"); println!("cargo:rerun-if-env-changed=NVCC"); println!("cargo:rerun-if-env-changed=CUDA_HOME"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); @@ -41,6 +42,51 @@ fn main() { let library = out_dir.join("libmulti_stark_cuda.a"); let architectures = cuda_architectures(&nvcc); + // sppark's units are compiled first, on their own, without the + // `_GNU_SOURCE` undefinition below: its runtime includes libstdc++'s + // ``, whose GNU-only pthread functions that flag would hide. Their + // objects then join the archive. + let mut sppark_objects = Vec::new(); + if env::var_os("CARGO_FEATURE_CUDA_SPPARK").is_some() { + let root = PathBuf::from( + env::var_os("DEP_SPPARK_ROOT").expect("sppark's build script exports DEP_SPPARK_ROOT"), + ); + for (source, object) in [ + (PathBuf::from("cuda/sppark_ntt.cu"), "sppark_ntt.o"), + (root.join("util/all_gpus.cpp"), "sppark_all_gpus.o"), + ] { + let object = out_dir.join(object); + let mut compile = Command::new(&nvcc); + compile + .arg("-c") + .arg("--std=c++17") + .arg("--cudart=static") + .arg("--default-stream=per-thread") + .arg("-O3") + .arg("-lineinfo") + .arg("--compiler-options=-fPIC") + .arg(format!("-I{}", root.display())) + .arg("-DFEATURE_GOLDILOCKS") + .arg("-o") + .arg(&object) + .arg(&source); + for architecture in &architectures { + compile.arg(format!( + "-gencode=arch=compute_{architecture},code=sm_{architecture}" + )); + } + let status = compile + .status() + .unwrap_or_else(|error| panic!("failed to execute {:?}: {error}", nvcc)); + assert!( + status.success(), + "nvcc failed on {} with status {status}", + source.display() + ); + sppark_objects.push(object); + } + } + let mut command = Command::new(&nvcc); command .arg("--lib") @@ -61,6 +107,7 @@ fn main() { .arg("-o") .arg(&library) .arg("cuda/kernels.cu"); + command.args(&sppark_objects); for architecture in &architectures { command.arg(format!( diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu new file mode 100644 index 0000000..cecf467 --- /dev/null +++ b/cuda/sppark_ntt.cu @@ -0,0 +1,107 @@ +// sppark's Goldilocks NTT behind a C interface for the comparison backend. +// +// Every transform runs on an upstream stream leased for the call and is +// fenced against the caller's per-thread stream with events, so a caller +// sees the same completion contract as the first-party kernels: work +// enqueued after the call on its own stream follows the transform, and the +// host is not blocked. Upstream's field type is a plain 64-bit word with +// canonical values, the same storage the prover's matrices use. +#include +#include + +#include +#include + +namespace { + +// Status codes above the CUDA range: an upstream exception carries its own +// CUDA error when it has one. +constexpr int SPPARK_EXCEPTION = 20000; + +struct EventPair { + cudaEvent_t before = nullptr; + cudaEvent_t after = nullptr; + ~EventPair() { + if (before) cudaEventDestroy(before); + if (after) cudaEventDestroy(after); + } + cudaError_t create() { + cudaError_t status = cudaEventCreateWithFlags(&before, cudaEventDisableTiming); + if (status == cudaSuccess) + status = cudaEventCreateWithFlags(&after, cudaEventDisableTiming); + return status; + } +}; + +bool valid_arguments(const void* d_inout, uint32_t lg, int order, int direction, int coset) { + return d_inout && lg > 0 && lg <= MAX_LG_DOMAIN_SIZE && order >= 0 && order <= 3 && + direction >= 0 && direction <= 1 && coset >= 0 && coset <= 1; +} + +} // namespace + +extern "C" int multi_stark_sppark_max_lg_domain() { return MAX_LG_DOMAIN_SIZE; } + +// One in-place transform of 2^lg field elements at `d_inout` on `device`. +// `order` is NTT::InputOutputOrder (NN, NR, RN, RR), `direction` 0 forward +// or 1 inverse, `coset` 1 for the multiplicative coset by the field +// generator. Inverse transforms are normalized by 1/2^lg upstream. +extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint32_t lg, + int order, int direction, int coset) { + if (!valid_arguments(d_inout, lg, order, direction, coset)) + return static_cast(cudaErrorInvalidValue); + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) return static_cast(status); + try { + const gpu_t& gpu = select_gpu(device); + // Upstream enumerates the visible devices in CUDA order; the caller's + // ordinal must name the same device. + if (gpu.cid() != device) return static_cast(cudaErrorInvalidDevice); + EventPair events; + status = events.create(); + if (status != cudaSuccess) return static_cast(status); + status = cudaEventRecord(events.before, cudaStreamPerThread); + if (status != cudaSuccess) return static_cast(status); + stream_t stream(gpu.id()); + stream.wait(events.before); + NTT::Base_dev_ptr(stream, reinterpret_cast(d_inout), lg, + static_cast(order), + static_cast(direction), + static_cast(coset)); + stream.record(events.after); + status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); + // The stream is destroyed on return; CUDA defers that until its work + // has drained, and the caller's stream already waits on it. + return static_cast(status); + } catch (const cuda_error& error) { + // Upstream stores the negated CUDA code. + const int code = static_cast(error.code()); + return SPPARK_EXCEPTION + (code < 0 ? -code : code); + } catch (const std::exception&) { + return SPPARK_EXCEPTION; + } +} + +// The same transform on host memory: uploads, transforms, downloads and +// synchronizes. For contract checks and small inputs, not the prover. +extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t lg, int order, + int direction, int coset) { + if (!valid_arguments(inout, lg, order, direction, coset)) + return static_cast(cudaErrorInvalidValue); + const size_t bytes = (size_t(1) << lg) * sizeof(uint64_t); + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) return static_cast(status); + uint64_t* d_inout = nullptr; + status = cudaMalloc(reinterpret_cast(&d_inout), bytes); + if (status != cudaSuccess) return static_cast(status); + status = cudaMemcpyAsync(d_inout, inout, bytes, cudaMemcpyHostToDevice, cudaStreamPerThread); + int result = static_cast(status); + if (result == 0) result = multi_stark_sppark_ntt_device(device, d_inout, lg, order, direction, coset); + if (result == 0) + result = static_cast( + cudaMemcpyAsync(inout, d_inout, bytes, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + const cudaError_t synced = cudaStreamSynchronize(cudaStreamPerThread); + if (result == 0) result = static_cast(synced); + cudaFree(d_inout); + return result; +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 8d16ea2..82331ad 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -7,6 +7,8 @@ pub(crate) mod metrics; pub(crate) mod mmcs; +#[cfg(feature = "cuda-sppark")] +pub mod sppark; #[doc(hidden)] pub mod pcs; pub(crate) mod witness; diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs new file mode 100644 index 0000000..b33bb7e --- /dev/null +++ b/src/cuda/sppark.rs @@ -0,0 +1,267 @@ +//! sppark's Goldilocks NTT through the device-pointer adapter in +//! `cuda/sppark_ntt.cu`: the comparison backend behind `cuda-sppark`. +//! +//! The contracts the prover relies on, checked by the tests below against +//! the CPU reference: a forward transform in natural order equals the DFT, +//! the `R` orders are the bit-reversed permutation of the `N` orders, +//! inverse transforms are normalized upstream, the coset variant is the +//! DFT on the coset shifted by the field generator, and outputs are stored +//! as canonical words. Inputs must be canonical: upstream does not reduce +//! representatives at or above the modulus. + +use core::ffi::c_int; + +use p3_field::PrimeField64; +use p3_goldilocks::Goldilocks; + +use super::check_cuda; + +/// `NTT::InputOutputOrder`: whether the input and the output are in natural +/// (`N`) or bit-reversed (`R`) order. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(i32)] +pub enum Order { + NN = 0, + NR = 1, + RN = 2, + RR = 3, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(i32)] +pub enum Direction { + Forward = 0, + Inverse = 1, +} + +unsafe extern "C" { + fn multi_stark_sppark_max_lg_domain() -> c_int; + fn multi_stark_sppark_ntt_device( + device: c_int, + d_inout: *mut u64, + lg: u32, + order: c_int, + direction: c_int, + coset: c_int, + ) -> c_int; + fn multi_stark_sppark_ntt_host( + device: c_int, + inout: *mut u64, + lg: u32, + order: c_int, + direction: c_int, + coset: c_int, + ) -> c_int; +} + +/// The largest log domain size the compiled upstream parameters support. +pub fn max_log_domain() -> usize { + usize::try_from(unsafe { multi_stark_sppark_max_lg_domain() }).expect("domain limit") +} + +fn log_len(values: usize) -> u32 { + assert!( + values.is_power_of_two() && values > 1, + "sppark transforms take 2^lg elements, lg > 0" + ); + values.trailing_zeros() +} + +/// Transforms `values` in place on `device`: uploaded, transformed and +/// downloaded within the call. `coset` selects the coset by the field +/// generator. +pub fn ntt_host( + device: i32, + values: &mut [Goldilocks], + order: Order, + direction: Direction, + coset: bool, +) { + let lg = log_len(values.len()); + let status = unsafe { + multi_stark_sppark_ntt_host( + device, + values.as_mut_ptr().cast(), + lg, + order as c_int, + direction as c_int, + c_int::from(coset), + ) + }; + check_cuda(status, "sppark host transform"); +} + +/// Transforms `2^lg` field elements at `d_inout` in place on the calling +/// thread's stream. +/// +/// # Safety +/// +/// `d_inout` must be a device allocation of `2^lg` 64-bit words on +/// `device`, and no other stream may access it until work enqueued after +/// this call on the calling thread's stream has completed. +pub unsafe fn ntt_device( + device: i32, + d_inout: *mut u64, + lg: u32, + order: Order, + direction: Direction, + coset: bool, +) { + let status = unsafe { + multi_stark_sppark_ntt_device( + device, + d_inout, + lg, + order as c_int, + direction as c_int, + c_int::from(coset), + ) + }; + check_cuda(status, "sppark device transform"); +} + +/// The raw stored words of `values`; Goldilocks is `repr(transparent)`. +pub fn raw_words(values: &[Goldilocks]) -> &[u64] { + // SAFETY: Goldilocks is a transparent wrapper over u64. + unsafe { core::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } +} + +#[cfg(test)] +mod tests { + use super::*; + use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; + use p3_field::{Field, PrimeCharacteristicRing}; + use p3_matrix::Matrix; + use p3_matrix::dense::RowMajorMatrix; + use p3_util::reverse_slice_index_bits; + use rand::{RngExt, SeedableRng, rngs::SmallRng}; + + const LOGS: [usize; 6] = [1, 5, 10, 11, 16, 20]; + + fn random(lg: usize, seed: u64) -> Vec { + let mut rng = SmallRng::seed_from_u64(seed); + (0..1usize << lg).map(|_| rng.random()).collect() + } + + fn cpu_dft(values: &[Goldilocks]) -> Vec { + Radix2DitParallel::::default() + .dft_batch(RowMajorMatrix::new(values.to_vec(), 1)) + .to_row_major_matrix() + .values + } + + fn cpu_idft(values: &[Goldilocks]) -> Vec { + Radix2DitParallel::::default() + .idft_batch(RowMajorMatrix::new(values.to_vec(), 1)) + .values + } + + fn cpu_coset_dft(values: &[Goldilocks]) -> Vec { + Radix2DitParallel::::default() + .coset_dft_batch( + RowMajorMatrix::new(values.to_vec(), 1), + Goldilocks::GENERATOR, + ) + .to_row_major_matrix() + .values + } + + fn assert_canonical(values: &[Goldilocks], what: &str) { + for (i, &word) in raw_words(values).iter().enumerate() { + assert!( + word < Goldilocks::ORDER_U64, + "{what}: word {i} is not canonical: {word:#x}" + ); + } + } + + #[test] + fn domain_limit_covers_the_trace_cap_and_its_lde() { + assert!( + max_log_domain() >= 26, + "the prover extends 2^24 traces by 4" + ); + } + + #[test] + fn forward_natural_matches_the_cpu_dft() { + for lg in LOGS { + let input = random(lg, 0xf0 + lg as u64); + let expected = cpu_dft(&input); + let mut actual = input.clone(); + ntt_host(0, &mut actual, Order::NN, Direction::Forward, false); + assert_eq!(actual, expected, "lg {lg}"); + assert_canonical(&actual, "forward NN"); + } + } + + #[test] + fn reversed_orders_are_the_bit_reversal_of_natural_ones() { + for lg in LOGS { + let input = random(lg, 0xb1 + lg as u64); + let mut natural = cpu_dft(&input); + let mut nr = input.clone(); + ntt_host(0, &mut nr, Order::NR, Direction::Forward, false); + reverse_slice_index_bits(&mut natural); + assert_eq!(nr, natural, "NR at lg {lg}"); + // RN: bit-reversed input yields the natural-order output. + let mut reversed_input = input.clone(); + reverse_slice_index_bits(&mut reversed_input); + ntt_host(0, &mut reversed_input, Order::RN, Direction::Forward, false); + assert_eq!(reversed_input, cpu_dft(&input), "RN at lg {lg}"); + } + } + + #[test] + fn inverse_is_normalized_upstream() { + for lg in LOGS { + let input = random(lg, 0x1d + lg as u64); + let mut actual = input.clone(); + ntt_host(0, &mut actual, Order::NN, Direction::Inverse, false); + assert_eq!(actual, cpu_idft(&input), "inverse NN at lg {lg}"); + // A forward NR followed by an inverse RN is the identity, so the + // prover's inverse-then-forward pair needs no scaling of its own. + let mut round_trip = input.clone(); + ntt_host(0, &mut round_trip, Order::NR, Direction::Forward, false); + ntt_host(0, &mut round_trip, Order::RN, Direction::Inverse, false); + assert_eq!(round_trip, input, "round trip at lg {lg}"); + } + } + + #[test] + fn coset_transform_uses_the_field_generator() { + for lg in LOGS { + let input = random(lg, 0xc0 + lg as u64); + let mut actual = input.clone(); + ntt_host(0, &mut actual, Order::NN, Direction::Forward, true); + assert_eq!(actual, cpu_coset_dft(&input), "coset forward NN at lg {lg}"); + } + } + + #[test] + fn extreme_canonical_values_transform_and_stay_canonical() { + // Upstream takes canonical words only: raw representatives at or + // above the modulus, which the first-party kernels reduce lazily, + // transform to different values, so the adapter's callers reduce + // first. Canonical extremes must survive unchanged. + let p = Goldilocks::ORDER_U64; + let words: Vec = [0, 1, p - 1, 7, p - 7, 1 << 32, (1 << 32) - 1, p - (1 << 32)] + .into_iter() + .cycle() + .take(1 << 10) + .collect(); + let input: Vec = words.iter().map(|&w| Goldilocks::from_u64(w)).collect(); + assert_eq!( + raw_words(&input), + &words[..], + "the inputs are stored canonically" + ); + let mut actual = input.clone(); + ntt_host(0, &mut actual, Order::NN, Direction::Forward, false); + assert_eq!(actual, cpu_dft(&input)); + assert_canonical(&actual, "forward of canonical extremes"); + ntt_host(0, &mut actual, Order::NN, Direction::Inverse, false); + assert_eq!(actual, input); + assert_canonical(&actual, "inverse of canonical extremes"); + } +} From fa64c28d609a05354cfa315a4489135b6917f286 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:30:26 +0000 Subject: [PATCH 13/34] cuda: build sppark's runtime without exceptions or threads The fork's `SPPARK_NO_CXX_RUNTIME` mode records a failing CUDA call in a per-thread status instead of throwing, and drops the thread pool from the GPU handle, so the sppark objects reference nothing from libstdc++ and the archive links into the Lean executable, which carries libc++. The adapter reads the recorded status after each transform. The dependency now names the organization's fork; the pinned revision moves to the fork's branch once it is pushed. --- Cargo.toml | 8 +++++--- build.rs | 5 +++++ cuda/sppark_ntt.cu | 48 +++++++++++++++++++++------------------------- 3 files changed, 32 insertions(+), 29 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index aad8291..3974cb1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,10 +33,12 @@ p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89 p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } -# sppark's Goldilocks NTT, for the comparison backend behind `cuda-sppark`. -# Pinned; the adapter compiles against the headers its build script exports. +# sppark's Goldilocks NTT, for the comparison backend behind `cuda-sppark`, +# from the organization's fork: its `multi-stark/no-cxx-runtime` branch adds +# the `SPPARK_NO_CXX_RUNTIME` build mode the adapter relies on. Pinned; the +# adapter compiles against the headers its build script exports. [dependencies.sppark] -git = "https://github.com/supranational/sppark" +git = "https://github.com/argumentcomputer/sppark" rev = "17278d74295392f9813f009300b257a688422b7a" optional = true features = ["cuda"] diff --git a/build.rs b/build.rs index dfb41c3..28865bb 100644 --- a/build.rs +++ b/build.rs @@ -67,6 +67,11 @@ fn main() { .arg("--compiler-options=-fPIC") .arg(format!("-I{}", root.display())) .arg("-DFEATURE_GOLDILOCKS") + // The fork's runtime without exceptions or its thread pool: + // no C++ runtime library symbols, so the archive links into + // the Lean executable, which carries libc++ rather than + // libstdc++. + .arg("-DSPPARK_NO_CXX_RUNTIME") .arg("-o") .arg(&object) .arg(&source); diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index cecf467..8a81239 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -5,19 +5,16 @@ // sees the same completion contract as the first-party kernels: work // enqueued after the call on its own stream follows the transform, and the // host is not blocked. Upstream's field type is a plain 64-bit word with -// canonical values, the same storage the prover's matrices use. +// canonical values, the same storage the prover's matrices use. The fork +// is built with SPPARK_NO_CXX_RUNTIME: CUDA failures inside upstream are +// recorded per thread and read back here rather than thrown. #include #include #include -#include namespace { -// Status codes above the CUDA range: an upstream exception carries its own -// CUDA error when it has one. -constexpr int SPPARK_EXCEPTION = 20000; - struct EventPair { cudaEvent_t before = nullptr; cudaEvent_t after = nullptr; @@ -52,16 +49,19 @@ extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint return static_cast(cudaErrorInvalidValue); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - try { - const gpu_t& gpu = select_gpu(device); - // Upstream enumerates the visible devices in CUDA order; the caller's - // ordinal must name the same device. - if (gpu.cid() != device) return static_cast(cudaErrorInvalidDevice); - EventPair events; - status = events.create(); - if (status != cudaSuccess) return static_cast(status); - status = cudaEventRecord(events.before, cudaStreamPerThread); - if (status != cudaSuccess) return static_cast(status); + // Upstream enumerates the visible devices in CUDA order; the caller's + // ordinal must name the same device. + if (device < 0 || static_cast(device) >= ngpus()) + return static_cast(cudaErrorInvalidDevice); + (void)sppark_take_cuda_error(); + const gpu_t& gpu = select_gpu(device); + if (gpu.cid() != device) return static_cast(cudaErrorInvalidDevice); + EventPair events; + status = events.create(); + if (status != cudaSuccess) return static_cast(status); + status = cudaEventRecord(events.before, cudaStreamPerThread); + if (status != cudaSuccess) return static_cast(status); + { stream_t stream(gpu.id()); stream.wait(events.before); NTT::Base_dev_ptr(stream, reinterpret_cast(d_inout), lg, @@ -69,17 +69,13 @@ extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint static_cast(direction), static_cast(coset)); stream.record(events.after); - status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); - // The stream is destroyed on return; CUDA defers that until its work - // has drained, and the caller's stream already waits on it. - return static_cast(status); - } catch (const cuda_error& error) { - // Upstream stores the negated CUDA code. - const int code = static_cast(error.code()); - return SPPARK_EXCEPTION + (code < 0 ? -code : code); - } catch (const std::exception&) { - return SPPARK_EXCEPTION; + // The stream is destroyed here; CUDA defers that until its work has + // drained, and the caller's stream waits on it below. } + // The fork records the first failing CUDA call instead of throwing. + if (const int recorded = sppark_take_cuda_error()) return recorded; + status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); + return static_cast(status); } // The same transform on host memory: uploads, transforms, downloads and From aa714fd6040cead476986e3649299f9118441030 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:37:45 +0000 Subject: [PATCH 14/34] cuda: the sppark fork's dev branch carries the runtime mode --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 3974cb1..318a0ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,7 +34,7 @@ p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a8906 p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } # sppark's Goldilocks NTT, for the comparison backend behind `cuda-sppark`, -# from the organization's fork: its `multi-stark/no-cxx-runtime` branch adds +# from the organization's fork: its `dev` branch adds # the `SPPARK_NO_CXX_RUNTIME` build mode the adapter relies on. Pinned; the # adapter compiles against the headers its build script exports. [dependencies.sppark] From d9645509bf5952e3c2d2df15e258d9b4c58c0d0d Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:42:40 +0000 Subject: [PATCH 15/34] cuda: pin sppark at the fork's dev branch, 6c5d826 --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bb3611f..228dfc4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1113,7 +1113,7 @@ dependencies = [ [[package]] name = "sppark" version = "0.1.15" -source = "git+https://github.com/supranational/sppark?rev=17278d74295392f9813f009300b257a688422b7a#17278d74295392f9813f009300b257a688422b7a" +source = "git+https://github.com/argumentcomputer/sppark?rev=6c5d826de1a36324a6c3e2ce3663139476f9b143#6c5d826de1a36324a6c3e2ce3663139476f9b143" dependencies = [ "cc", "which", diff --git a/Cargo.toml b/Cargo.toml index 318a0ad..1e7c9e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,7 +39,7 @@ p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c837 # adapter compiles against the headers its build script exports. [dependencies.sppark] git = "https://github.com/argumentcomputer/sppark" -rev = "17278d74295392f9813f009300b257a688422b7a" +rev = "6c5d826de1a36324a6c3e2ce3663139476f9b143" optional = true features = ["cuda"] From 44ea38cb46398499be9c75768e39f8d0a290f103 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:59:47 +0000 Subject: [PATCH 16/34] cuda: sppark adapter resolves the device by ordinal and fences before releasing its stream Upstream keeps only the devices it supports, so its logical index can differ from the CUDA ordinal; the adapter now finds the entry by ordinal and rejects an unknown one before touching the buffer, which a test covers. The caller's stream waits on the transform while the private stream is still alive, and if that wait cannot be installed the private stream is drained before returning, so no queued kernel outlives the call. With the fork stopping at the first failing CUDA call, the recorded-status read is gone. The pin moves to the fork's fail-fast commit once pushed. --- cuda/sppark_ntt.cu | 45 +++++++++++++++++++++++---------------------- src/cuda/sppark.rs | 33 +++++++++++++++++++++++++++++---- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 8a81239..47d841c 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -6,8 +6,9 @@ // enqueued after the call on its own stream follows the transform, and the // host is not blocked. Upstream's field type is a plain 64-bit word with // canonical values, the same storage the prover's matrices use. The fork -// is built with SPPARK_NO_CXX_RUNTIME: CUDA failures inside upstream are -// recorded per thread and read back here rather than thrown. +// is built with SPPARK_NO_CXX_RUNTIME: a CUDA failure inside upstream ends +// the process with a message rather than throwing, which matches the +// status checks on the Rust side. #include #include @@ -49,32 +50,32 @@ extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint return static_cast(cudaErrorInvalidValue); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - // Upstream enumerates the visible devices in CUDA order; the caller's - // ordinal must name the same device. - if (device < 0 || static_cast(device) >= ngpus()) - return static_cast(cudaErrorInvalidDevice); - (void)sppark_take_cuda_error(); - const gpu_t& gpu = select_gpu(device); - if (gpu.cid() != device) return static_cast(cudaErrorInvalidDevice); + // Upstream keeps only the devices it supports, so its logical index can + // differ from the CUDA ordinal; find the entry by ordinal. + const gpu_t* found = nullptr; + for (const gpu_t* candidate : all_gpus()) + if (candidate->cid() == device) found = candidate; + if (!found) return static_cast(cudaErrorInvalidDevice); + const gpu_t& gpu = select_gpu(found->id()); EventPair events; status = events.create(); if (status != cudaSuccess) return static_cast(status); status = cudaEventRecord(events.before, cudaStreamPerThread); if (status != cudaSuccess) return static_cast(status); - { - stream_t stream(gpu.id()); - stream.wait(events.before); - NTT::Base_dev_ptr(stream, reinterpret_cast(d_inout), lg, - static_cast(order), - static_cast(direction), - static_cast(coset)); - stream.record(events.after); - // The stream is destroyed here; CUDA defers that until its work has - // drained, and the caller's stream waits on it below. - } - // The fork records the first failing CUDA call instead of throwing. - if (const int recorded = sppark_take_cuda_error()) return recorded; + // Upstream's CUDA failures end the process (the fork's runtime mode), so + // everything past the launch either completes or never returns. + stream_t stream(gpu.id()); + stream.wait(events.before); + NTT::Base_dev_ptr(stream, reinterpret_cast(d_inout), lg, + static_cast(order), + static_cast(direction), + static_cast(coset)); + stream.record(events.after); status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); + // Only once the caller's stream waits on the transform may the private + // stream go; if that wait could not be installed, drain the stream here + // so nothing still runs against the caller's buffer on return. + if (status != cudaSuccess) stream.sync(); return static_cast(status); } diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index b33bb7e..65cde4e 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -69,14 +69,15 @@ fn log_len(values: usize) -> u32 { /// Transforms `values` in place on `device`: uploaded, transformed and /// downloaded within the call. `coset` selects the coset by the field -/// generator. -pub fn ntt_host( +/// generator. An error is the CUDA status the adapter returned; `device` +/// must be the CUDA ordinal of a device upstream supports. +pub fn try_ntt_host( device: i32, values: &mut [Goldilocks], order: Order, direction: Direction, coset: bool, -) { +) -> Result<(), i32> { let lg = log_len(values.len()); let status = unsafe { multi_stark_sppark_ntt_host( @@ -88,7 +89,20 @@ pub fn ntt_host( c_int::from(coset), ) }; - check_cuda(status, "sppark host transform"); + if status == 0 { Ok(()) } else { Err(status) } +} + +/// [`try_ntt_host`], panicking on a CUDA status like the other backends. +pub fn ntt_host( + device: i32, + values: &mut [Goldilocks], + order: Order, + direction: Direction, + coset: bool, +) { + if let Err(status) = try_ntt_host(device, values, order, direction, coset) { + check_cuda(status, "sppark host transform"); + } } /// Transforms `2^lg` field elements at `d_inout` in place on the calling @@ -175,6 +189,17 @@ mod tests { } } + #[test] + fn an_unknown_device_ordinal_is_rejected_before_any_launch() { + // cudaErrorInvalidDevice is 101; the adapter answers it for an + // ordinal upstream did not enumerate, without touching the buffer. + let input = random(8, 0xde); + let mut values = input.clone(); + let status = try_ntt_host(1 << 20, &mut values, Order::NN, Direction::Forward, false); + assert_eq!(status, Err(101)); + assert_eq!(values, input); + } + #[test] fn domain_limit_covers_the_trace_cap_and_its_lde() { assert!( From 9fef49971bad308fc31b8e9ff371b22585ec30fa Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:16:35 +0000 Subject: [PATCH 17/34] cuda: resident coset LDE through sppark behind a runtime switch With `cuda-sppark`, MULTI_STARK_CUDA_NTT=sppark (or a runtime selection for one-process comparisons) routes `coset_lde_create` through a column panel: a gather that canonicalizes into column-major scratch, an inverse transform per column, one pass that restores natural coefficient order and applies the coset powers into a second panel with a zeroed tail, a forward transform per column, and a scatter into the bit-reversed row order the commitment expects. Panels are sized from MULTI_STARK_SPPARK_PANEL_BYTES. The stored words match the first-party kernels bit for bit across the tested heights, widths, blowups and shifts, raw representatives included. Medians of five warm resident LDEs on one RTX PRO 6000: 2^24 x 6 at blowup 4 from 170 to 71 ms, the 2^22 width-2 codeword from 9.7 to 4.5 ms, the 2^20 x 533 BLAKE3 shape from 386 to 315 ms; 2^24 x 17 even; short wide matrices slower, 2^16 x 925 from 25 to 57 ms, on per-column launches. --- build.rs | 4 + cuda/kernels.cu | 36 ++++-- cuda/sppark_ntt.cu | 163 ++++++++++++++++++++++++++++ examples/cuda_resident_lde_bench.rs | 8 ++ src/cuda/sppark.rs | 91 ++++++++++++++++ 5 files changed, 293 insertions(+), 9 deletions(-) diff --git a/build.rs b/build.rs index 28865bb..bb7ccb5 100644 --- a/build.rs +++ b/build.rs @@ -66,6 +66,7 @@ fn main() { .arg("-lineinfo") .arg("--compiler-options=-fPIC") .arg(format!("-I{}", root.display())) + .arg("-Icuda") .arg("-DFEATURE_GOLDILOCKS") // The fork's runtime without exceptions or its thread pool: // no C++ runtime library symbols, so the archive links into @@ -112,6 +113,9 @@ fn main() { .arg("-o") .arg(&library) .arg("cuda/kernels.cu"); + if !sppark_objects.is_empty() { + command.arg("-DMULTI_STARK_SPPARK"); + } command.args(&sppark_objects); for architecture in &architectures { diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 2053469..2c283cc 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -18,6 +18,12 @@ #include #include #include "metrics.cuh" +#ifdef MULTI_STARK_SPPARK +extern "C" int multi_stark_sppark_backend_selected(); +extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, + size_t height, size_t width, size_t added_bits, + const uint64_t* shift_powers); +#endif namespace { @@ -2384,13 +2390,19 @@ static int coset_lde_create( } if (status == cudaSuccess) lde->trace_height = height; +#ifdef MULTI_STARK_SPPARK + // The sppark path writes every output row itself from its own scratch. + const bool sppark = multi_stark_sppark_backend_selected() != 0; +#else + const bool sppark = false; +#endif // Large pageable uploads otherwise serialize through the driver's hidden // staging pool; they go through the persistent staging slots instead. // Small ones take the direct pageable path. const uint64_t *device_inverse_twiddles=nullptr,*device_shift_powers=nullptr,*device_forward_twiddles=nullptr; // The source prefix is overwritten by the trace copy. Only the padded // tail needs zeroing before the forward transform. - if (status == cudaSuccess && output_elements > input_elements) { + if (status == cudaSuccess && !sppark && output_elements > input_elements) { status = cudaMemsetAsync(lde->values + input_elements, 0, (output_elements - input_elements) * sizeof(uint64_t), cudaStreamPerThread); @@ -2413,30 +2425,36 @@ static int coset_lde_create( } } if (status == cudaSuccess) { + status = cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&device_shift_powers); + } +#ifdef MULTI_STARK_SPPARK + if (status == cudaSuccess && sppark) { + status = static_cast(multi_stark_sppark_coset_lde( + device_id, lde->trace_values, lde->values, height, width, added_bits, device_shift_powers)); + } +#endif + if (status == cudaSuccess && !sppark) { status = cudaMemcpyAsync(lde->values, lde->trace_values, input_elements * sizeof(uint64_t), cudaMemcpyDeviceToDevice, cudaStreamPerThread); } - if (status == cudaSuccess && height > 1) { + if (status == cudaSuccess && !sppark && height > 1) { status = cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&device_inverse_twiddles); } - if (status == cudaSuccess) { - status = cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&device_shift_powers); - } - if (status == cudaSuccess) { + if (status == cudaSuccess && !sppark) { status = cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&device_forward_twiddles); } - if (status == cudaSuccess) { + if (status == cudaSuccess && !sppark) { status = launch_dif(device_id, lde->values, height, width,device_inverse_twiddles); } - if (status == cudaSuccess) { + if (status == cudaSuccess && !sppark) { bit_reverse_scale_and_shift<<>>( lde->values, height, width, strict_log2(height), height_inverse, device_shift_powers); status = cudaGetLastError(); } - if (status == cudaSuccess) { + if (status == cudaSuccess && !sppark) { status = launch_dif(device_id, lde->values, extended_height, width,device_forward_twiddles); } // Normalization and every DIF butterfly produce canonical field values, diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 47d841c..a2cf696 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -13,6 +13,8 @@ #include #include +#include +#include namespace { @@ -102,3 +104,164 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t cudaFree(d_inout); return result; } + +// --- Resident coset LDE through sppark --------------------------------- +// +// The prover's matrices are row-major with the transform along the column, +// while upstream transforms one contiguous vector. A panel of columns is +// gathered into column-major scratch, transformed column by column, and +// scattered back in the bit-reversed row order the commitment expects: +// +// A[c][r] = canonical(trace[r][f + c]) r < N +// inverse NR on A[c], normalized upstream (bit-reversed coeffs) +// B[c][i] = A[c][rev(i)] * shift^i i < N, zero beyond +// forward NR on B[c] over M = N << added_bits (bit-reversed evals) +// values[r][f + c] = B[c][r] r < M +// +// The gather reduces representatives at or above the modulus, which upstream +// does not accept. Two panels of C columns cost 16 * M * C bytes; C is sized +// from MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. + +#include "goldilocks.cuh" + +namespace { + +constexpr unsigned PANEL_THREADS = 256; +constexpr size_t MAX_BLOCKS = 65535; + +unsigned blocks_for_total(size_t total) { + const size_t blocks = (total + PANEL_THREADS - 1) / PANEL_THREADS; + return static_cast(blocks < MAX_BLOCKS ? blocks : MAX_BLOCKS); +} + +__device__ __forceinline__ size_t reverse_bits(size_t index, unsigned log) { + return log == 0 ? 0 : static_cast(__brev(static_cast(index)) >> (32 - log)); +} + +__global__ void gather_columns(const uint64_t* __restrict__ trace, size_t height, size_t width, + size_t first, size_t columns, size_t extended_height, + uint64_t* __restrict__ panel) { + const size_t total = height * columns; + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; + index += stride) { + const size_t row = index / columns; + const size_t column = index - row * columns; + panel[column * extended_height + row] = multi_stark_cuda::canonicalize(trace[row * width + first + column]); + } +} + +__global__ void shift_columns(const uint64_t* __restrict__ coefficients, uint64_t* __restrict__ panel, + size_t height, unsigned log_height, size_t extended_height, + size_t columns, const uint64_t* __restrict__ shift_powers) { + const size_t total = extended_height * columns; + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; + index += stride) { + const size_t column = index / extended_height; + const size_t row = index - column * extended_height; + uint64_t value = 0; + if (row < height) { + const uint64_t coefficient = coefficients[column * extended_height + reverse_bits(row, log_height)]; + value = multi_stark_cuda::goldilocks_mul(coefficient, shift_powers[row]); + } + panel[index] = value; + } +} + +__global__ void scatter_columns(const uint64_t* __restrict__ panel, size_t extended_height, size_t width, + size_t first, size_t columns, uint64_t* __restrict__ values) { + const size_t total = extended_height * columns; + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; + index += stride) { + const size_t row = index / columns; + const size_t column = index - row * columns; + values[row * width + first + column] = panel[column * extended_height + row]; + } +} + +unsigned log2_exact(size_t value) { + unsigned log = 0; + while ((size_t(1) << log) < value) ++log; + return log; +} + +int backend_flag = -1; + +// Read per construction: one getenv against a transform of gigabytes, and +// tests vary it within a process. +size_t panel_budget_bytes() { + size_t budget = size_t(4) << 30; + if (const char* configured = getenv("MULTI_STARK_SPPARK_PANEL_BYTES")) { + char* end = nullptr; + const unsigned long long parsed = strtoull(configured, &end, 10); + if (end != configured && *end == '\0' && parsed > 0) budget = parsed; + } + return budget; +} + +} // namespace + +// Whether the prover's transforms take the sppark path: MULTI_STARK_CUDA_NTT=sppark, +// or a runtime selection, which tests use to compare both paths in one process. +extern "C" int multi_stark_sppark_backend_selected() { + if (backend_flag < 0) { + const char* configured = getenv("MULTI_STARK_CUDA_NTT"); + backend_flag = configured && strcmp(configured, "sppark") == 0; + } + return backend_flag; +} + +extern "C" void multi_stark_sppark_select_backend(int selected) { backend_flag = selected ? 1 : 0; } + +// The coset LDE of `trace` (height x width, natural row order, device memory) +// into `values` (extended_height x width, bit-reversed rows), with the coset +// shift powers `shift_powers[i] = shift^i` for i < height. Scratch is +// allocated per call within the panel budget. +extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, + size_t height, size_t width, size_t added_bits, + const uint64_t* shift_powers) { + if (!trace || !values || !shift_powers || height == 0 || width == 0 || (height & (height - 1))) + return static_cast(cudaErrorInvalidValue); + const size_t extended_height = height << added_bits; + const unsigned log_height = log2_exact(height); + const unsigned log_extended = log_height + static_cast(added_bits); + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) return static_cast(status); + const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); + size_t columns = panel_budget_bytes() / column_bytes; + if (columns == 0) columns = 1; + if (columns > width) columns = width; + uint64_t* scratch = nullptr; + status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes); + if (status != cudaSuccess) return static_cast(status); + uint64_t* a = scratch; + uint64_t* b = scratch + columns * extended_height; + int result = 0; + for (size_t first = 0; result == 0 && first < width; first += columns) { + const size_t count = width - first < columns ? width - first : columns; + gather_columns<<>>( + trace, height, width, first, count, extended_height, a); + result = static_cast(cudaGetLastError()); + for (size_t column = 0; result == 0 && column < count && log_height > 0; ++column) + result = multi_stark_sppark_ntt_device(device, a + column * extended_height, log_height, 1, 1, 0); + if (result == 0) { + shift_columns<<>>( + a, b, height, log_height, extended_height, count, shift_powers); + result = static_cast(cudaGetLastError()); + } + for (size_t column = 0; result == 0 && column < count && log_extended > 0; ++column) + result = multi_stark_sppark_ntt_device(device, b + column * extended_height, log_extended, 1, 0, 0); + if (result == 0) { + scatter_columns<<>>( + b, extended_height, width, first, count, values); + result = static_cast(cudaGetLastError()); + } + } + // The scratch outlives every kernel that reads it: the free is ordered + // behind them on the same stream. + const cudaError_t freed = cudaFreeAsync(scratch, cudaStreamPerThread); + if (result == 0) result = static_cast(freed); + return result; +} diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs index 84bedfd..658f9c5 100644 --- a/examples/cuda_resident_lde_bench.rs +++ b/examples/cuda_resident_lde_bench.rs @@ -20,6 +20,9 @@ fn main() { .map_or(7, |s| s.parse().expect("invalid iteration count")); let gpu = CudaDft::default(); println!("log_height,width,added_bits,iteration,seconds"); + // The last five are the shapes the profiled Init proofs commit most: + // BLAKE3 pieces, wide and narrow IxVM circuits at the height cap, and + // the width-2 codewords of the quotient and narrow lookups. for (log_height, width, added_bits) in [ (20, 1, 1), (20, 2, 1), @@ -29,6 +32,11 @@ fn main() { (18, 129, 1), (16, 925, 1), (18, 40, 2), + (20, 533, 2), + (24, 6, 2), + (24, 17, 2), + (22, 2, 2), + (20, 2, 2), ] { let height = 1 << log_height; let values = (0..height * width) diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 65cde4e..bd80b39 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -36,6 +36,8 @@ pub enum Direction { unsafe extern "C" { fn multi_stark_sppark_max_lg_domain() -> c_int; + fn multi_stark_sppark_backend_selected() -> c_int; + fn multi_stark_sppark_select_backend(selected: c_int); fn multi_stark_sppark_ntt_device( device: c_int, d_inout: *mut u64, @@ -54,6 +56,20 @@ unsafe extern "C" { ) -> c_int; } +/// Whether the prover's resident LDEs take the sppark path: selected by +/// `MULTI_STARK_CUDA_NTT=sppark` or by [`select_backend`]. +pub fn backend_selected() -> bool { + unsafe { multi_stark_sppark_backend_selected() != 0 } +} + +/// Routes the prover's resident LDEs through sppark, or back to the +/// first-party kernels, for the rest of the process. Comparisons in one +/// process toggle it between constructions; concurrent constructions all +/// see the latest value. +pub fn select_backend(sppark: bool) { + unsafe { multi_stark_sppark_select_backend(c_int::from(sppark)) } +} + /// The largest log domain size the compiled upstream parameters support. pub fn max_log_domain() -> usize { usize::try_from(unsafe { multi_stark_sppark_max_lg_domain() }).expect("domain limit") @@ -200,6 +216,81 @@ mod tests { assert_eq!(values, input); } + /// One resident LDE both ways, comparing stored words and the Merkle root + /// of a commitment over the matrix. + fn resident_lde_both_ways( + matrix: RowMajorMatrix, + added_bits: usize, + shift: Goldilocks, + ) { + let dft = super::super::CudaDft::new(0); + select_backend(false); + let legacy = dft.coset_lde_batch_resident(&matrix, added_bits, shift); + let legacy_rows = legacy.to_row_major_matrix(); + select_backend(true); + let candidate = dft.coset_lde_batch_resident(&matrix, added_bits, shift); + let candidate_rows = candidate.to_row_major_matrix(); + select_backend(false); + assert_eq!( + raw_words(&candidate_rows.values), + raw_words(&legacy_rows.values), + "height {} width {} blowup {added_bits}", + matrix.height(), + matrix.width() + ); + } + + #[test] + fn resident_lde_matches_the_first_party_kernels_bit_for_bit() { + let mut rng = SmallRng::seed_from_u64(0x1de5); + for log_height in [0usize, 1, 2, 5, 8, 12, 14] { + for added_bits in [0usize, 1, 2, 3] { + for width in [1usize, 2, 3, 7, 8, 33] { + let height = 1 << log_height; + let matrix = RowMajorMatrix::new( + (0..height * width).map(|_| rng.random()).collect(), + width, + ); + resident_lde_both_ways(matrix, added_bits, Goldilocks::GENERATOR); + } + } + } + } + + #[test] + fn resident_lde_reduces_raw_representatives_like_the_first_party_kernels() { + let p = Goldilocks::ORDER_U64; + let words = [0u64, 1, p - 1, p, p + 1, u64::MAX, 7, p + 7]; + let height = 1usize << 10; + let width = 9; + // SAFETY: the test reinterprets raw words as field elements on purpose. + let values: Vec = (0..height * width) + .map(|i| unsafe { core::mem::transmute::(words[i % words.len()]) }) + .collect(); + for shift in [ + Goldilocks::GENERATOR, + Goldilocks::ONE, + Goldilocks::from_u64(11), + ] { + resident_lde_both_ways(RowMajorMatrix::new(values.clone(), width), 2, shift); + } + } + + #[test] + fn resident_lde_panels_narrower_than_the_matrix_cover_every_column() { + // A 2^12 x 33 matrix at blowup 2 needs 16 KiB x 2 per column, so a + // 256 KiB budget forces panels of a few columns. + let mut rng = SmallRng::seed_from_u64(0x9a7e); + let previous = std::env::var("MULTI_STARK_SPPARK_PANEL_BYTES").ok(); + unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", "262144") }; + let matrix = RowMajorMatrix::new((0..(1 << 12) * 33).map(|_| rng.random()).collect(), 33); + resident_lde_both_ways(matrix, 2, Goldilocks::GENERATOR); + match previous { + Some(value) => unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", value) }, + None => unsafe { std::env::remove_var("MULTI_STARK_SPPARK_PANEL_BYTES") }, + } + } + #[test] fn domain_limit_covers_the_trace_cap_and_its_lde() { assert!( From 62cee20cbcb6de1e643f633dd5ed3dc8913e150a Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:23:12 +0000 Subject: [PATCH 18/34] cuda: admit the sppark panel scratch, dispatch by height, count shapes per backend The commit loop adds the two panels the sppark path allocates to a source's admission. Below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (20) a resident LDE stays on the first-party kernels, where the per-column baseline is launch-bound and slower; the selection has a third state that takes sppark at every height for one-process comparisons. The metrics snapshot counts taken and declined dispatches and keeps transform shapes per backend. --- cuda/kernels.cu | 11 ++++++-- cuda/metrics.cuh | 14 ++++++++-- cuda/sppark_ntt.cu | 48 ++++++++++++++++++++++++++++---- src/cuda/metrics.rs | 23 ++++++++++------ src/cuda/sppark.rs | 67 +++++++++++++++++++++++++++++++++++++++------ src/cuda/witness.rs | 5 ++++ 6 files changed, 141 insertions(+), 27 deletions(-) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 2c283cc..8fbf6cb 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -19,7 +19,7 @@ #include #include "metrics.cuh" #ifdef MULTI_STARK_SPPARK -extern "C" int multi_stark_sppark_backend_selected(); +extern "C" int multi_stark_sppark_takes(size_t height); extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, size_t height, size_t width, size_t added_bits, const uint64_t* shift_powers); @@ -2392,7 +2392,14 @@ static int coset_lde_create( #ifdef MULTI_STARK_SPPARK // The sppark path writes every output row itself from its own scratch. - const bool sppark = multi_stark_sppark_backend_selected() != 0; + const bool sppark = multi_stark_sppark_takes(height) != 0; + if (multi_stark_metrics::enabled()) { + multi_stark_metrics::add(device_id, sppark ? multi_stark_metrics::SpparkTaken : multi_stark_metrics::SpparkDeclined, 1); + if (sppark) { + multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); + multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(extended_height, width), 1); + } + } #else const bool sppark = false; #endif diff --git a/cuda/metrics.cuh b/cuda/metrics.cuh index 57719f2..87e0918 100644 --- a/cuda/metrics.cuh +++ b/cuda/metrics.cuh @@ -10,12 +10,22 @@ namespace multi_stark_metrics { constexpr size_t DEVICES = 64; constexpr size_t NTT_OFFSET = 16; -constexpr size_t WORDS = NTT_OFFSET + 33 * 4; +// Transform shape counts, one block per backend: first-party, then sppark. +constexpr size_t NTT_SHAPES = 33 * 4; +constexpr size_t NTT_SPPARK_OFFSET = NTT_OFFSET + NTT_SHAPES; +constexpr size_t WORDS = NTT_SPPARK_OFFSET + NTT_SHAPES; enum Counter : size_t { UploadCalls, UploadRequestedBytes, UploadChunks, UploadFailures, UploadHostNs, CosetHits, CosetMisses, CosetUploadedBytes, ConstantBytes, - DriverFreeBytes, TotalBytes, MemorySamples + DriverFreeBytes, TotalBytes, MemorySamples, + SpparkTaken, SpparkDeclined }; +inline size_t ntt_shape(size_t height, size_t width) { + unsigned log = 0; + while ((size_t(1) << log) < height) ++log; + const size_t bucket = width == 1 ? 0 : width == 2 ? 1 : width < 8 ? 2 : 3; + return log * 4 + bucket; +} std::atomic counters[DEVICES][WORDS]{}; inline bool enabled() { static const bool value = std::getenv("AIUR_METRICS") != nullptr; diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index a2cf696..1ba38fa 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -187,8 +187,6 @@ unsigned log2_exact(size_t value) { return log; } -int backend_flag = -1; - // Read per construction: one getenv against a transform of gigabytes, and // tests vary it within a process. size_t panel_budget_bytes() { @@ -201,6 +199,27 @@ size_t panel_budget_bytes() { return budget; } +// -1 unread, 0 first-party, 1 sppark above the height threshold, 2 sppark +// for every height (tests compare the paths on small shapes). +int backend_flag = -1; + +unsigned min_log_height() { + unsigned log = 20; + if (const char* configured = getenv("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT")) { + char* end = nullptr; + const unsigned long parsed = strtoul(configured, &end, 10); + if (end != configured && *end == '\0') log = static_cast(parsed); + } + return log; +} + +size_t panel_columns(size_t width, size_t extended_height) { + const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); + size_t columns = panel_budget_bytes() / column_bytes; + if (columns == 0) columns = 1; + return columns < width ? columns : width; +} + } // namespace // Whether the prover's transforms take the sppark path: MULTI_STARK_CUDA_NTT=sppark, @@ -213,7 +232,26 @@ extern "C" int multi_stark_sppark_backend_selected() { return backend_flag; } -extern "C" void multi_stark_sppark_select_backend(int selected) { backend_flag = selected ? 1 : 0; } +// 0 first-party, 1 sppark above the height threshold, 2 sppark always. +extern "C" void multi_stark_sppark_select_backend(int selected) { backend_flag = selected; } + +// Whether a resident LDE of `height` input rows takes the sppark path. +// Short transforms are launch-bound on the per-column baseline and stay on +// the first-party kernels below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (20). +extern "C" int multi_stark_sppark_takes(size_t height) { + const int flag = multi_stark_sppark_backend_selected(); + if (flag == 2) return 1; + if (flag != 1) return 0; + return height >= (size_t(1) << min_log_height()); +} + +// The scratch the sppark path allocates for one LDE: two panels of the +// columns the budget admits, sized for the extended height. +extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { + if (!multi_stark_sppark_takes(height) || width == 0) return 0; + const size_t extended_height = height << added_bits; + return panel_columns(width, extended_height) * 2 * extended_height * sizeof(uint64_t); +} // The coset LDE of `trace` (height x width, natural row order, device memory) // into `values` (extended_height x width, bit-reversed rows), with the coset @@ -230,9 +268,7 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); - size_t columns = panel_budget_bytes() / column_bytes; - if (columns == 0) columns = 1; - if (columns > width) columns = width; + const size_t columns = panel_columns(width, extended_height); uint64_t* scratch = nullptr; status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes); if (status != cudaSuccess) return static_cast(status); diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs index 56a7dea..69ca624 100644 --- a/src/cuda/metrics.rs +++ b/src/cuda/metrics.rs @@ -4,7 +4,8 @@ pub(crate) fn emit_snapshot() { if !tracing::enabled!(target: "prover_metrics", tracing::Level::INFO) { return; } - const WORDS: usize = 16 + 33 * 4; + const SHAPES: usize = 33 * 4; + const WORDS: usize = 16 + 2 * SHAPES; let mut values = vec![0u64; 64 * WORDS]; unsafe { multi_stark_cuda_metrics_snapshot(values.as_mut_ptr(), values.len()); @@ -17,14 +18,20 @@ pub(crate) fn emit_snapshot() { scope = "process_device_cumulative", upload_calls = c[0], upload_requested_bytes = c[1], upload_chunks = c[2], upload_failures = c[3], upload_host_ns = c[4], coset_hits = c[5], coset_misses = c[6], coset_uploaded_bytes = c[7], constant_bytes = c[8], - last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11]); - for (shape, &count) in c[16..].iter().enumerate() { - if count == 0 { - continue; + last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11], + sppark_taken = c[12], sppark_declined = c[13]); + for (backend, shapes) in [ + ("legacy", &c[16..16 + SHAPES]), + ("sppark", &c[16 + SHAPES..]), + ] { + for (shape, &count) in shapes.iter().enumerate() { + if count == 0 { + continue; + } + tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, + scope = "process_device_cumulative", backend, log_height = shape / 4, + width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); } - tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, - scope = "process_device_cumulative", backend = "legacy", log_height = shape / 4, - width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); } } } diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index bd80b39..7f48ec9 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -38,6 +38,8 @@ unsafe extern "C" { fn multi_stark_sppark_max_lg_domain() -> c_int; fn multi_stark_sppark_backend_selected() -> c_int; fn multi_stark_sppark_select_backend(selected: c_int); + fn multi_stark_sppark_takes(height: usize) -> c_int; + fn multi_stark_sppark_panel_bytes(height: usize, width: usize, added_bits: usize) -> usize; fn multi_stark_sppark_ntt_device( device: c_int, d_inout: *mut u64, @@ -62,12 +64,40 @@ pub fn backend_selected() -> bool { unsafe { multi_stark_sppark_backend_selected() != 0 } } -/// Routes the prover's resident LDEs through sppark, or back to the -/// first-party kernels, for the rest of the process. Comparisons in one -/// process toggle it between constructions; concurrent constructions all -/// see the latest value. -pub fn select_backend(sppark: bool) { - unsafe { multi_stark_sppark_select_backend(c_int::from(sppark)) } +/// Which transforms take the sppark path. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Backend { + /// The first-party kernels only. + Legacy, + /// sppark for LDEs at or above the height threshold + /// (`MULTI_STARK_SPPARK_MIN_LOG_HEIGHT`, 20), where the per-column + /// baseline wins; the first-party kernels below it. + Sppark, + /// sppark at every height: for comparing the paths on small shapes. + SpparkAllHeights, +} + +/// Selects the backend for the rest of the process. Comparisons in one +/// process switch between constructions; concurrent constructions all see +/// the latest value. +pub fn select_backend(backend: Backend) { + let flag = match backend { + Backend::Legacy => 0, + Backend::Sppark => 1, + Backend::SpparkAllHeights => 2, + }; + unsafe { multi_stark_sppark_select_backend(flag) } +} + +/// Whether a resident LDE of `height` input rows takes the sppark path. +pub fn takes(height: usize) -> bool { + unsafe { multi_stark_sppark_takes(height) != 0 } +} + +/// The scratch the sppark path allocates for that LDE, to admit alongside +/// the trace and the LDE; zero when the first-party kernels take it. +pub fn panel_bytes(height: usize, width: usize, added_bits: usize) -> usize { + unsafe { multi_stark_sppark_panel_bytes(height, width, added_bits) } } /// The largest log domain size the compiled upstream parameters support. @@ -224,13 +254,13 @@ mod tests { shift: Goldilocks, ) { let dft = super::super::CudaDft::new(0); - select_backend(false); + select_backend(Backend::Legacy); let legacy = dft.coset_lde_batch_resident(&matrix, added_bits, shift); let legacy_rows = legacy.to_row_major_matrix(); - select_backend(true); + select_backend(Backend::SpparkAllHeights); let candidate = dft.coset_lde_batch_resident(&matrix, added_bits, shift); let candidate_rows = candidate.to_row_major_matrix(); - select_backend(false); + select_backend(Backend::Legacy); assert_eq!( raw_words(&candidate_rows.values), raw_words(&legacy_rows.values), @@ -291,6 +321,25 @@ mod tests { } } + #[test] + fn the_height_threshold_and_the_panel_budget_decide_dispatch_and_scratch() { + select_backend(Backend::Sppark); + assert!( + !takes(1 << 12), + "short transforms stay on the first-party kernels" + ); + assert!(takes(1 << 20)); + assert_eq!(panel_bytes(1 << 12, 533, 2), 0); + // 2^20 rows, 533 columns, blowup 4: 2 x 2^22 x 8 bytes per column is + // 64 MiB, so a 4 GiB budget admits 64 columns. + assert_eq!(panel_bytes(1 << 20, 533, 2), 64 * 2 * (1 << 22) * 8); + select_backend(Backend::SpparkAllHeights); + assert!(takes(2)); + select_backend(Backend::Legacy); + assert!(!takes(1 << 24)); + assert_eq!(panel_bytes(1 << 24, 6, 2), 0); + } + #[test] fn domain_limit_covers_the_trace_cap_and_its_lde() { assert!( diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index e1eb4f0..3c8e4ce 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -78,11 +78,16 @@ pub(crate) fn commit( let mut host: Vec>> = Vec::with_capacity(evaluations.len()); let mut retained = Vec::with_capacity(evaluations.len()); for (domain, source) in evaluations { + #[cfg(feature = "cuda-sppark")] + let panel = super::sppark::panel_bytes(source.height(), source.width(), blowup); + #[cfg(not(feature = "cuda-sppark"))] + let panel = 0; let needed = source .height() .saturating_mul(source.width()) .saturating_mul(8) .saturating_mul((1 << blowup) + 1) + .saturating_add(panel) .saturating_add(reserve); // Generator caches go before any LDE spills: they are rebuilt from // host seeds on demand, a spilled LDE is uploaded again. From 24ef583ca858653dd048c06d23ed1fdd9bc78e3b Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:37:38 +0000 Subject: [PATCH 19/34] cuda: every transform of a proof through sppark above the height threshold The lookup LDEs (graph, direct and partitioned), the quotient's two forward transforms, the general DFT and the host coset LDE dispatch through two helpers, one for the inverse-shift-forward sequence in place and one for a forward transform in place, so the backend rule lives in one place: sppark's column panel above MULTI_STARK_SPPARK_MIN_LOG_HEIGHT, the first-party stages otherwise. The adapter gains the forward panel entry and an ungated transform counter. A batch proof and the compatibility example produce identical bytes on the first-party kernels, on sppark above the threshold and on sppark at every height; the general DFT and host coset LDE entries are compared bit for bit as well. --- cuda/kernels.cu | 108 +++++++++++++++++++++++++++----------------- cuda/sppark_ntt.cu | 44 ++++++++++++++++++ src/batch.rs | 23 ++++++++++ src/cuda/metrics.rs | 2 +- src/cuda/sppark.rs | 82 +++++++++++++++++++++++++++++---- src/cuda/witness.rs | 4 +- 6 files changed, 211 insertions(+), 52 deletions(-) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 8fbf6cb..7c05439 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -20,6 +20,7 @@ #include "metrics.cuh" #ifdef MULTI_STARK_SPPARK extern "C" int multi_stark_sppark_takes(size_t height); +extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width); extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, size_t height, size_t width, size_t added_bits, const uint64_t* shift_powers); @@ -2251,6 +2252,59 @@ cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, } // namespace +// The inverse-shift-forward sequence in place on `values`, whose first +// `height` rows hold natural-order evaluations and whose tail is zero: the +// sppark panel above its height threshold, the first-party stages +// otherwise. Both leave bit-reversed rows; `canonical_pass` adds the +// reduction pass the lookup paths run after the first-party stages. +static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t height, size_t width, + size_t extended_height, const uint64_t* inverse_twiddles, + const uint64_t* shift_powers, const uint64_t* forward_twiddles, + uint64_t height_inverse, bool canonical_pass) { +#ifdef MULTI_STARK_SPPARK + if (multi_stark_sppark_takes(height)) { + if (multi_stark_metrics::enabled()) { + multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); + multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); + multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(extended_height, width), 1); + } + const size_t added_bits = strict_log2(extended_height) - strict_log2(height); + return static_cast(multi_stark_sppark_coset_lde( + device_id, values, values, height, width, added_bits, shift_powers)); + } + if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); +#endif + cudaError_t status = launch_dif(device_id, values, height, width, inverse_twiddles); + if (status == cudaSuccess) { + bit_reverse_scale_and_shift<<>>( + values, height, width, strict_log2(height), height_inverse, shift_powers); + status = cudaGetLastError(); + } + if (status == cudaSuccess) status = launch_dif(device_id, values, extended_height, width, forward_twiddles); + if (status == cudaSuccess && canonical_pass) { + canonicalize_goldilocks<<>>(values, extended_height * width); + status = cudaGetLastError(); + } + return status; +} + +// A forward transform in place on natural-order rows, leaving bit-reversed +// rows: sppark above its height threshold, the first-party stages otherwise. +static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t height, size_t width, + const uint64_t* twiddles) { +#ifdef MULTI_STARK_SPPARK + if (multi_stark_sppark_takes(height)) { + if (multi_stark_metrics::enabled()) { + multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); + multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); + } + return static_cast(multi_stark_sppark_forward(device_id, values, height, width)); + } + if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); +#endif + return launch_dif(device_id, values, height, width, twiddles); +} + extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, size_t height, size_t width, const uint64_t* twiddles) { @@ -2272,7 +2326,7 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, status = copy_to_device(device_twiddles, twiddles, height / 2); } if (status == cudaSuccess) { - status = launch_dif(device_id, device_values.get(), height, width, device_twiddles.get()); + status = forward_in_place(device_id, device_values.get(), height, width, device_twiddles.get()); } if (status == cudaSuccess) { status = copy_to_host(values, device_values, elements); @@ -2331,19 +2385,9 @@ extern "C" int multi_stark_cuda_coset_lde_batch( extended_height / 2); } if (status == cudaSuccess) { - status = launch_dif(device_id, device_values.get(), height, width, - device_inverse_twiddles.get()); - } - if (status == cudaSuccess) { - const size_t total = input_elements; - bit_reverse_scale_and_shift<<>>( - device_values.get(), height, width, strict_log2(height), - height_inverse, device_shift_powers.get()); - status = cudaGetLastError(); - } - if (status == cudaSuccess) { - status = launch_dif(device_id, device_values.get(), extended_height, width, - device_forward_twiddles.get()); + status = coset_lde_in_place(device_id, device_values.get(), height, width, extended_height, + device_inverse_twiddles.get(), device_shift_powers.get(), + device_forward_twiddles.get(), height_inverse, false); } if (status == cudaSuccess) { status = copy_to_host(output, device_values, output_elements); @@ -2931,7 +2975,7 @@ extern "C" int multi_stark_cuda_quotient_lde( dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(device_id, quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,device_quotient_twiddles); ResidentLde* lde = nullptr; if(status==cudaSuccess) { @@ -2948,7 +2992,7 @@ extern "C" int multi_stark_cuda_quotient_lde( quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,device_lde_twiddles); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess) { *output_handle=lde; @@ -3160,7 +3204,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( } for(size_t i=0;i<2;++i)if(status==cudaSuccess&&stream_busy[i])status=cudaStreamSynchronize(streams[i]); - if(status==cudaSuccess)status=launch_dif(device_id, quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,device_quotient_twiddles); ResidentLde* lde=nullptr; if(status==cudaSuccess)status=create_resident_lde(&lde); if(status==cudaSuccess){lde->height=lde_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t));} @@ -3170,7 +3214,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( lde->values,quotient,device_weights,quotient_size,trace_height,quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,device_lde_twiddles); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else if(lde)destroy_resident_lde(lde); for(size_t i=0;i<2;++i){ @@ -3460,10 +3504,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,height,width,dit); - if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,extended_height,width,dft); - if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} + if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); cudaFree(trace_chunk);cudaFree(scratch);cudaFree(deltas);cudaFree(multiplicities);cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(metadata); @@ -3526,10 +3567,7 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,height,width,dit); - if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} - if(status==cudaSuccess)status=launch_dif(device_id, lde->values,extended_height,width,dft); - if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} + if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, "[multi-stark/cuda] lookup phases: height=%zu lookups=%zu slots=%zu args_width=%zu allocate=%.3fs rows=%.3fs scan=%.3fs dft=%.3fs\n", @@ -3718,21 +3756,9 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( &forward); const size_t width = 2 * pending->slots; if (status == cudaSuccess) - status = launch_dif(device_id, pending->lde->values, pending->height, width, inverse); - if (status == cudaSuccess) { - bit_reverse_scale_and_shift<<height * width), THREADS>>>( - pending->lde->values, pending->height, width, - strict_log2(pending->height), height_inverse, shifts); - status = cudaGetLastError(); - } - if (status == cudaSuccess) - status = launch_dif(device_id, pending->lde->values, pending->extended_height, - width, forward); - if (status == cudaSuccess) { - canonicalize_goldilocks<<extended_height * width), THREADS>>>( - pending->lde->values, pending->extended_height * width); - status = cudaGetLastError(); - } + status = coset_lde_in_place(device_id, pending->lde->values, pending->height, width, + pending->extended_height, inverse, shifts, forward, + height_inverse, true); if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status == cudaSuccess) { *output_handle = pending->lde; diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 1ba38fa..397cf5c 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -257,9 +258,18 @@ extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, si // into `values` (extended_height x width, bit-reversed rows), with the coset // shift powers `shift_powers[i] = shift^i` for i < height. Scratch is // allocated per call within the panel budget. +// Transforms the adapter has run so far, whatever the metrics setting, so a +// test can tell a proof that went through sppark from one that did not. +static std::atomic transforms_run{0}; + +extern "C" uint64_t multi_stark_sppark_transforms_run() { + return transforms_run.load(std::memory_order_relaxed); +} + extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, size_t height, size_t width, size_t added_bits, const uint64_t* shift_powers) { + transforms_run.fetch_add(1, std::memory_order_relaxed); if (!trace || !values || !shift_powers || height == 0 || width == 0 || (height & (height - 1))) return static_cast(cudaErrorInvalidValue); const size_t extended_height = height << added_bits; @@ -301,3 +311,37 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u if (result == 0) result = static_cast(freed); return result; } + +// A forward transform in place on `values` (height x width, natural row +// order) leaving bit-reversed rows: gather, upstream forward per column, +// scatter. The quotient's transforms and the general DFT take this path. +extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width) { + transforms_run.fetch_add(1, std::memory_order_relaxed); + if (!values || height == 0 || width == 0 || (height & (height - 1))) + return static_cast(cudaErrorInvalidValue); + const unsigned log_height = log2_exact(height); + if (log_height == 0) return 0; + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) return static_cast(status); + const size_t columns = panel_columns(width, height); + uint64_t* panel = nullptr; + status = cudaMalloc(reinterpret_cast(&panel), columns * height * sizeof(uint64_t)); + if (status != cudaSuccess) return static_cast(status); + int result = 0; + for (size_t first = 0; result == 0 && first < width; first += columns) { + const size_t count = width - first < columns ? width - first : columns; + gather_columns<<>>( + values, height, width, first, count, height, panel); + result = static_cast(cudaGetLastError()); + for (size_t column = 0; result == 0 && column < count; ++column) + result = multi_stark_sppark_ntt_device(device, panel + column * height, log_height, 1, 0, 0); + if (result == 0) { + scatter_columns<<>>( + panel, height, width, first, count, values); + result = static_cast(cudaGetLastError()); + } + } + const cudaError_t freed = cudaFreeAsync(panel, cudaStreamPerThread); + if (result == 0) result = static_cast(freed); + return result; +} diff --git a/src/batch.rs b/src/batch.rs index 88504ec..0d97a2b 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -773,6 +773,29 @@ mod tests { system.verify_batch(&decoded).unwrap(); } + /// Every transform of a batch proof, main and lookup LDEs and the + /// quotient, produces the same bytes through sppark as through the + /// first-party kernels. + #[cfg(feature = "cuda-sppark")] + #[test] + fn sppark_transforms_prove_the_same_batch() { + use crate::cuda::sppark::{Backend, select_backend, transforms_run}; + let (system, key) = byte_system(config()); + select_backend(Backend::Legacy); + let legacy = system.prove_batch(&key, two_shards(&system), vec![]); + select_backend(Backend::SpparkAllHeights); + let before = transforms_run(); + let candidate = system.prove_batch(&key, two_shards(&system), vec![]); + let through_sppark = transforms_run() - before; + select_backend(Backend::Legacy); + assert!( + through_sppark > 0, + "the proof's transforms ran through sppark" + ); + assert_eq!(candidate.to_bytes().unwrap(), legacy.to_bytes().unwrap()); + system.verify_batch(&candidate).unwrap(); + } + #[test] fn regenerated_shards_prove_the_same_batch() { let (system, key) = byte_system(config()); diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs index 69ca624..5e432ee 100644 --- a/src/cuda/metrics.rs +++ b/src/cuda/metrics.rs @@ -10,7 +10,7 @@ pub(crate) fn emit_snapshot() { unsafe { multi_stark_cuda_metrics_snapshot(values.as_mut_ptr(), values.len()); } - for (device, c) in values.chunks_exact(WORDS).enumerate() { + for (device, c) in values.as_chunks::().0.iter().enumerate() { if c.iter().all(|&v| v == 0) { continue; } diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 7f48ec9..b669e3c 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -11,7 +11,6 @@ use core::ffi::c_int; -use p3_field::PrimeField64; use p3_goldilocks::Goldilocks; use super::check_cuda; @@ -39,6 +38,7 @@ unsafe extern "C" { fn multi_stark_sppark_backend_selected() -> c_int; fn multi_stark_sppark_select_backend(selected: c_int); fn multi_stark_sppark_takes(height: usize) -> c_int; + fn multi_stark_sppark_transforms_run() -> u64; fn multi_stark_sppark_panel_bytes(height: usize, width: usize, added_bits: usize) -> usize; fn multi_stark_sppark_ntt_device( device: c_int, @@ -89,6 +89,11 @@ pub fn select_backend(backend: Backend) { unsafe { multi_stark_sppark_select_backend(flag) } } +/// How many transforms the adapter has run in this process. +pub fn transforms_run() -> u64 { + unsafe { multi_stark_sppark_transforms_run() } +} + /// Whether a resident LDE of `height` input rows takes the sppark path. pub fn takes(height: usize) -> bool { unsafe { multi_stark_sppark_takes(height) != 0 } @@ -190,7 +195,7 @@ pub fn raw_words(values: &[Goldilocks]) -> &[u64] { mod tests { use super::*; use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; - use p3_field::{Field, PrimeCharacteristicRing}; + use p3_field::{Field, PrimeCharacteristicRing, PrimeField64}; use p3_matrix::Matrix; use p3_matrix::dense::RowMajorMatrix; use p3_util::reverse_slice_index_bits; @@ -249,16 +254,16 @@ mod tests { /// One resident LDE both ways, comparing stored words and the Merkle root /// of a commitment over the matrix. fn resident_lde_both_ways( - matrix: RowMajorMatrix, + matrix: &RowMajorMatrix, added_bits: usize, shift: Goldilocks, ) { let dft = super::super::CudaDft::new(0); select_backend(Backend::Legacy); - let legacy = dft.coset_lde_batch_resident(&matrix, added_bits, shift); + let legacy = dft.coset_lde_batch_resident(matrix, added_bits, shift); let legacy_rows = legacy.to_row_major_matrix(); select_backend(Backend::SpparkAllHeights); - let candidate = dft.coset_lde_batch_resident(&matrix, added_bits, shift); + let candidate = dft.coset_lde_batch_resident(matrix, added_bits, shift); let candidate_rows = candidate.to_row_major_matrix(); select_backend(Backend::Legacy); assert_eq!( @@ -281,7 +286,7 @@ mod tests { (0..height * width).map(|_| rng.random()).collect(), width, ); - resident_lde_both_ways(matrix, added_bits, Goldilocks::GENERATOR); + resident_lde_both_ways(&matrix, added_bits, Goldilocks::GENERATOR); } } } @@ -302,7 +307,7 @@ mod tests { Goldilocks::ONE, Goldilocks::from_u64(11), ] { - resident_lde_both_ways(RowMajorMatrix::new(values.clone(), width), 2, shift); + resident_lde_both_ways(&RowMajorMatrix::new(values.clone(), width), 2, shift); } } @@ -314,13 +319,74 @@ mod tests { let previous = std::env::var("MULTI_STARK_SPPARK_PANEL_BYTES").ok(); unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", "262144") }; let matrix = RowMajorMatrix::new((0..(1 << 12) * 33).map(|_| rng.random()).collect(), 33); - resident_lde_both_ways(matrix, 2, Goldilocks::GENERATOR); + resident_lde_both_ways(&matrix, 2, Goldilocks::GENERATOR); match previous { Some(value) => unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", value) }, None => unsafe { std::env::remove_var("MULTI_STARK_SPPARK_PANEL_BYTES") }, } } + #[test] + fn general_dft_matches_the_first_party_kernels_bit_for_bit() { + // Shapes above the CUDA DFT threshold of 2^15 cells. + let mut rng = SmallRng::seed_from_u64(0xdf7); + let dft = super::super::CudaDft::new(0); + for (log_height, width) in [(15usize, 1usize), (12, 8), (10, 33), (16, 2), (11, 129)] { + let matrix = RowMajorMatrix::new( + (0..(1 << log_height) * width) + .map(|_| rng.random()) + .collect(), + width, + ); + select_backend(Backend::Legacy); + let legacy = dft.dft_batch(matrix.clone()).to_row_major_matrix(); + select_backend(Backend::SpparkAllHeights); + let candidate = dft.dft_batch(matrix).to_row_major_matrix(); + select_backend(Backend::Legacy); + assert_eq!( + raw_words(&candidate.values), + raw_words(&legacy.values), + "2^{log_height} x {width}" + ); + } + } + + #[test] + fn host_coset_lde_matches_the_first_party_kernels_bit_for_bit() { + // The host entry takes matrices of width at most two whose extended + // height reaches 2^15. + let mut rng = SmallRng::seed_from_u64(0x1e5); + let dft = super::super::CudaDft::new(0); + let generator = Goldilocks::GENERATOR; + for (log_height, width, added_bits) in [ + (14usize, 1usize, 1usize), + (14, 2, 2), + (15, 2, 3), + (16, 1, 1), + ] { + let matrix = RowMajorMatrix::new( + (0..(1 << log_height) * width) + .map(|_| rng.random()) + .collect(), + width, + ); + select_backend(Backend::Legacy); + let legacy = dft + .coset_lde_batch(matrix.clone(), added_bits, generator) + .to_row_major_matrix(); + select_backend(Backend::SpparkAllHeights); + let candidate = dft + .coset_lde_batch(matrix, added_bits, generator) + .to_row_major_matrix(); + select_backend(Backend::Legacy); + assert_eq!( + raw_words(&candidate.values), + raw_words(&legacy.values), + "2^{log_height} x {width} blowup {added_bits}" + ); + } + } + #[test] fn the_height_threshold_and_the_panel_budget_decide_dispatch_and_scratch() { select_backend(Backend::Sppark); diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index 3c8e4ce..6377826 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -91,11 +91,11 @@ pub(crate) fn commit( .saturating_add(reserve); // Generator caches go before any LDE spills: they are rebuilt from // host seeds on demand, a spilled LDE is uploaded again. - for index in 0..resident.len() { + for slot in resident.iter() { if device_memory_info(device).0 >= needed { break; } - resident[index].as_ref().unwrap().release_generator_device(); + slot.as_ref().unwrap().release_generator_device(); } for index in 0..resident.len() { if device_memory_info(device).0 >= needed { From a15a9795b225478b6a5aa25a36939f21f91b278d Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:43:00 +0000 Subject: [PATCH 20/34] cuda: the sppark adapter parses its settings without strtoul nvcc's host compiler builds the adapter unit without the C standard pin the crate's C units get, so glibc redirected strtoul to a C23 symbol the Lean toolchain's libc lacks and the ix executable no longer linked. --- cuda/sppark_ntt.cu | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 397cf5c..0c9eaba 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -188,16 +188,26 @@ unsigned log2_exact(size_t value) { return log; } +// A decimal setting, or `fallback` when unset or not a plain number. This +// unit is compiled by nvcc's host compiler without the C standard pin the +// crate's C units get, where glibc redirects strtoul to a C23 symbol the +// Lean toolchain's libc does not carry. +unsigned long long decimal_setting(const char* name, unsigned long long fallback) { + const char* configured = getenv(name); + if (!configured || !*configured) return fallback; + unsigned long long value = 0; + for (const char* c = configured; *c; ++c) { + if (*c < '0' || *c > '9' || value > (~0ull - 9) / 10) return fallback; + value = value * 10 + unsigned(*c - '0'); + } + return value; +} + // Read per construction: one getenv against a transform of gigabytes, and // tests vary it within a process. size_t panel_budget_bytes() { - size_t budget = size_t(4) << 30; - if (const char* configured = getenv("MULTI_STARK_SPPARK_PANEL_BYTES")) { - char* end = nullptr; - const unsigned long long parsed = strtoull(configured, &end, 10); - if (end != configured && *end == '\0' && parsed > 0) budget = parsed; - } - return budget; + const unsigned long long budget = decimal_setting("MULTI_STARK_SPPARK_PANEL_BYTES", 0); + return budget ? size_t(budget) : size_t(4) << 30; } // -1 unread, 0 first-party, 1 sppark above the height threshold, 2 sppark @@ -205,13 +215,7 @@ size_t panel_budget_bytes() { int backend_flag = -1; unsigned min_log_height() { - unsigned log = 20; - if (const char* configured = getenv("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT")) { - char* end = nullptr; - const unsigned long parsed = strtoul(configured, &end, 10); - if (end != configured && *end == '\0') log = static_cast(parsed); - } - return log; + return unsigned(decimal_setting("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT", 20)); } size_t panel_columns(size_t width, size_t extended_height) { From 4739cb7dd5525f03bdf067950c232a6dbd281fae Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:30:28 +0000 Subject: [PATCH 21/34] cuda: sppark dispatch by whole shape, admitted everywhere, labelled in spans The dispatch rule takes the complete shape: tall enough, the extended height within upstream's compiled domain, one column's scratch within the panel budget; a shape outside any of these stays on the first-party kernels instead of failing after allocation. Lookup and quotient admission add the sppark panels their transforms allocate, as the commit loop already did. The paths that go through sppark no longer build or upload the first-party twiddle tables, in the prewarm and at the lookup and quotient sites. The backend selector is an atomic with a single publication of the environment setting, the comparison tests serialize on a lock so both constructions of a comparison see the backend they selected, and the LDE, lookup and quotient spans label the backend they ran on. --- cuda/kernels.cu | 60 +++++++++++++++++++++++++---------- cuda/sppark_ntt.cu | 79 +++++++++++++++++++++++++++++++++++----------- src/batch.rs | 3 +- src/cuda/mod.rs | 76 ++++++++++++++++++++++++++++++++++++++------ src/cuda/sppark.rs | 48 +++++++++++++++++++++++++++- src/types.rs | 20 ++++++++++-- 6 files changed, 237 insertions(+), 49 deletions(-) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 7c05439..1099184 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -20,12 +20,16 @@ #include "metrics.cuh" #ifdef MULTI_STARK_SPPARK extern "C" int multi_stark_sppark_takes(size_t height); +extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits); +extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width); extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width); extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, size_t height, size_t width, size_t added_bits, const uint64_t* shift_powers); #endif + + namespace { using namespace multi_stark_cuda; @@ -2252,6 +2256,26 @@ cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, } // namespace +// Whether the shape's inverse-shift-forward sequence, or forward transform, +// runs through sppark; false without the feature. +static bool sppark_takes_lde(size_t height, size_t width, size_t extended_height) { +#ifdef MULTI_STARK_SPPARK + return multi_stark_sppark_takes_lde(height, width, strict_log2(extended_height) - strict_log2(height)) != 0; +#else + (void)height; (void)width; (void)extended_height; + return false; +#endif +} + +static bool sppark_takes_forward(size_t height, size_t width) { +#ifdef MULTI_STARK_SPPARK + return multi_stark_sppark_takes_forward(height, width) != 0; +#else + (void)height; (void)width; + return false; +#endif +} + // The inverse-shift-forward sequence in place on `values`, whose first // `height` rows hold natural-order evaluations and whose tail is zero: the // sppark panel above its height threshold, the first-party stages @@ -2262,7 +2286,7 @@ static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t he const uint64_t* shift_powers, const uint64_t* forward_twiddles, uint64_t height_inverse, bool canonical_pass) { #ifdef MULTI_STARK_SPPARK - if (multi_stark_sppark_takes(height)) { + if (sppark_takes_lde(height, width, extended_height)) { if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); @@ -2293,7 +2317,7 @@ static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t he static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t height, size_t width, const uint64_t* twiddles) { #ifdef MULTI_STARK_SPPARK - if (multi_stark_sppark_takes(height)) { + if (sppark_takes_forward(height, width)) { if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); @@ -2436,7 +2460,7 @@ static int coset_lde_create( #ifdef MULTI_STARK_SPPARK // The sppark path writes every output row itself from its own scratch. - const bool sppark = multi_stark_sppark_takes(height) != 0; + const bool sppark = multi_stark_sppark_takes_lde(height, width, added_bits) != 0; if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, sppark ? multi_stark_metrics::SpparkTaken : multi_stark_metrics::SpparkDeclined, 1); if (sppark) { @@ -2558,13 +2582,13 @@ extern "C" int multi_stark_cuda_prepare_lde_constants( const uint64_t* shift_powers,size_t height,const uint64_t* forward_twiddles, size_t forward_count) { if((inverse_count&&!inverse_twiddles)||!shift_powers||!height|| - !forward_twiddles||!forward_count)return static_cast(cudaErrorInvalidValue); + (forward_count&&!forward_twiddles))return static_cast(cudaErrorInvalidValue); cudaError_t status=cudaSetDevice(device_id);const uint64_t* ignored=nullptr; if(status==cudaSuccess&&inverse_count) status=cached_device_constants(device_id,inverse_twiddles,inverse_count,1,0,0,&ignored); if(status==cudaSuccess) status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&ignored); - if(status==cudaSuccess) + if(status==cudaSuccess&&forward_count) status=cached_device_constants(device_id,forward_twiddles,forward_count,2,0,0,&ignored); return static_cast(status); } @@ -2952,8 +2976,8 @@ extern "C" int multi_stark_cuda_quotient_lde( coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess)status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); - if(status==cudaSuccess)status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess&&!sppark_takes_forward(quotient_size,2))status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); + if(status==cudaSuccess&&!sppark_takes_forward(lde_height,width))status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); @@ -3081,8 +3105,8 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess)status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); - if(status==cudaSuccess)status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess&&!sppark_takes_forward(quotient_size,2))status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); + if(status==cudaSuccess&&!sppark_takes_forward(lde_height,width))status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); const auto* prep_resident=static_cast(preprocessed_handle); @@ -3501,9 +3525,10 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + const bool legacy_tables=!sppark_takes_lde(height,width,extended_height); + if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); @@ -3564,9 +3589,10 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); scan_done=now(); const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + const bool legacy_tables=!sppark_takes_lde(height,width,extended_height); + if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, @@ -3742,7 +3768,10 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( const uint64_t* inverse = nullptr; const uint64_t* shifts = nullptr; const uint64_t* forward = nullptr; - if (status == cudaSuccess) + const size_t width = 2 * pending->slots; + const bool legacy_tables = + !sppark_takes_lde(pending->height, width, pending->extended_height); + if (status == cudaSuccess && legacy_tables) status = cached_device_constants(device_id, inverse_twiddles, pending->height / 2, 1, 0, 0, &inverse); if (status == cudaSuccess) @@ -3750,11 +3779,10 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( pending->height, 3, shift_powers[0], pending->height > 1 ? shift_powers[1] : 0, &shifts); - if (status == cudaSuccess) + if (status == cudaSuccess && legacy_tables) status = cached_device_constants(device_id, forward_twiddles, pending->extended_height / 2, 2, 0, 0, &forward); - const size_t width = 2 * pending->slots; if (status == cudaSuccess) status = coset_lde_in_place(device_id, pending->lde->values, pending->height, width, pending->extended_height, inverse, shifts, forward, diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 0c9eaba..32a2ff8 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -211,38 +211,54 @@ size_t panel_budget_bytes() { } // -1 unread, 0 first-party, 1 sppark above the height threshold, 2 sppark -// for every height (tests compare the paths on small shapes). -int backend_flag = -1; +// for every height (tests compare the paths on small shapes). Read and +// written from concurrent constructions. +std::atomic backend_flag{-1}; unsigned min_log_height() { return unsigned(decimal_setting("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT", 20)); } -size_t panel_columns(size_t width, size_t extended_height) { - const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); - size_t columns = panel_budget_bytes() / column_bytes; - if (columns == 0) columns = 1; +// The columns one panel holds at `column_bytes` each: as many as the +// budget admits, at most the width. Zero when one column does not fit, +// which the dispatch rules decline before any allocation. +size_t panel_columns(size_t width, size_t column_bytes) { + const size_t columns = panel_budget_bytes() / column_bytes; return columns < width ? columns : width; } +// Whether a transform of 2^log rows is within upstream's compiled domain. +bool within_domain(size_t height, size_t added_bits) { + if (height == 0 || (height & (height - 1))) return false; + const unsigned log = log2_exact(height); + return added_bits <= MAX_LG_DOMAIN_SIZE && log + added_bits <= MAX_LG_DOMAIN_SIZE; +} + } // namespace // Whether the prover's transforms take the sppark path: MULTI_STARK_CUDA_NTT=sppark, // or a runtime selection, which tests use to compare both paths in one process. extern "C" int multi_stark_sppark_backend_selected() { - if (backend_flag < 0) { + int flag = backend_flag.load(std::memory_order_acquire); + if (flag < 0) { const char* configured = getenv("MULTI_STARK_CUDA_NTT"); - backend_flag = configured && strcmp(configured, "sppark") == 0; + int expected = -1; + const int read = configured && strcmp(configured, "sppark") == 0; + // The first reader publishes; a concurrent selection wins over it. + flag = backend_flag.compare_exchange_strong(expected, read, std::memory_order_acq_rel) ? read : expected; } - return backend_flag; + return flag; } // 0 first-party, 1 sppark above the height threshold, 2 sppark always. -extern "C" void multi_stark_sppark_select_backend(int selected) { backend_flag = selected; } +extern "C" void multi_stark_sppark_select_backend(int selected) { + backend_flag.store(selected, std::memory_order_release); +} -// Whether a resident LDE of `height` input rows takes the sppark path. -// Short transforms are launch-bound on the per-column baseline and stay on -// the first-party kernels below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (20). +// Whether a transform of `height` input rows is tall enough for the sppark +// path. Short transforms are launch-bound on the per-column baseline and +// stay on the first-party kernels below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT +// (20). extern "C" int multi_stark_sppark_takes(size_t height) { const int flag = multi_stark_sppark_backend_selected(); if (flag == 2) return 1; @@ -250,12 +266,35 @@ extern "C" int multi_stark_sppark_takes(size_t height) { return height >= (size_t(1) << min_log_height()); } +// Whether a resident coset LDE of the shape takes the sppark path: tall +// enough, the extended height within upstream's compiled domain, and one +// column's scratch within the panel budget. +extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits) { + if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, added_bits)) return 0; + const size_t extended_height = height << added_bits; + return 2 * extended_height * sizeof(uint64_t) <= panel_budget_bytes(); +} + +// Whether a forward transform of the shape takes the sppark path. +extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width) { + if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, 0)) return 0; + return height * sizeof(uint64_t) <= panel_budget_bytes(); +} + // The scratch the sppark path allocates for one LDE: two panels of the -// columns the budget admits, sized for the extended height. +// columns the budget admits, sized for the extended height. Zero when the +// shape does not take the path. extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { - if (!multi_stark_sppark_takes(height) || width == 0) return 0; - const size_t extended_height = height << added_bits; - return panel_columns(width, extended_height) * 2 * extended_height * sizeof(uint64_t); + if (!multi_stark_sppark_takes_lde(height, width, added_bits)) return 0; + const size_t column_bytes = 2 * (height << added_bits) * sizeof(uint64_t); + return panel_columns(width, column_bytes) * column_bytes; +} + +// The scratch the sppark path allocates for one forward transform. +extern "C" size_t multi_stark_sppark_forward_panel_bytes(size_t height, size_t width) { + if (!multi_stark_sppark_takes_forward(height, width)) return 0; + const size_t column_bytes = height * sizeof(uint64_t); + return panel_columns(width, column_bytes) * column_bytes; } // The coset LDE of `trace` (height x width, natural row order, device memory) @@ -282,7 +321,8 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); - const size_t columns = panel_columns(width, extended_height); + const size_t columns = panel_columns(width, column_bytes); + if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* scratch = nullptr; status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes); if (status != cudaSuccess) return static_cast(status); @@ -327,7 +367,8 @@ extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t h if (log_height == 0) return 0; cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - const size_t columns = panel_columns(width, height); + const size_t columns = panel_columns(width, height * sizeof(uint64_t)); + if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* panel = nullptr; status = cudaMalloc(reinterpret_cast(&panel), columns * height * sizeof(uint64_t)); if (status != cudaSuccess) return static_cast(status); diff --git a/src/batch.rs b/src/batch.rs index 0d97a2b..2c42e2d 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -779,7 +779,8 @@ mod tests { #[cfg(feature = "cuda-sppark")] #[test] fn sppark_transforms_prove_the_same_batch() { - use crate::cuda::sppark::{Backend, select_backend, transforms_run}; + use crate::cuda::sppark::{Backend, backend_lock, select_backend, transforms_run}; + let _guard = backend_lock(); let (system, key) = byte_system(config()); select_backend(Backend::Legacy); let legacy = system.prove_batch(&key, two_shards(&system), vec![]); diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 82331ad..52a2014 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -7,10 +7,10 @@ pub(crate) mod metrics; pub(crate) mod mmcs; -#[cfg(feature = "cuda-sppark")] -pub mod sppark; #[doc(hidden)] pub mod pcs; +#[cfg(feature = "cuda-sppark")] +pub mod sppark; pub(crate) mod witness; use core::ffi::{CStr, c_char, c_void}; @@ -180,7 +180,7 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "host", - backend = "legacy", + backend = lde_backend(height, width, added_bits), device = self.device_id, height, width, @@ -225,6 +225,9 @@ impl CudaDft { } } + /// Uploads the constants a resident LDE of the shape will need: the + /// coset powers always, the first-party twiddle tables only when the + /// shape stays on the first-party kernels. pub(crate) fn prepare_coset_lde_constants( &self, height: usize, @@ -236,9 +239,18 @@ impl CudaDft { .expect("LDE height overflows usize"); Self::validate_dimensions(height, 1); Self::validate_dimensions(extended_height, 1); - let inverse_twiddles = self.twiddles(log2_strict_usize(height), true); + let legacy_tables = lde_backend(height, 1, added_bits) == "legacy"; + let inverse_twiddles = if legacy_tables { + self.twiddles(log2_strict_usize(height), true) + } else { + Arc::from(Vec::new()) + }; let shift_powers = self.shift_powers(height, shift); - let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let forward_twiddles = if legacy_tables { + self.twiddles(log2_strict_usize(extended_height), false) + } else { + Arc::from(Vec::new()) + }; let status = unsafe { multi_stark_cuda_prepare_lde_constants( self.device_id, @@ -361,7 +373,7 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "generated", - backend = "legacy", + backend = lde_backend(height, width, added_bits), device = self.device_id, height, width, @@ -816,6 +828,50 @@ fn encode_quotient_nodes( /// kernel use the same liveness calculation. The scratch term assumes the /// global-memory path; devices able to fit the slots in shared memory need /// less than this bound. +/// The backend a resident coset LDE of the shape runs on, for spans and +/// for skipping the tables the other backend would need. +pub(crate) fn lde_backend(height: usize, width: usize, added_bits: usize) -> &'static str { + #[cfg(feature = "cuda-sppark")] + if sppark::takes_lde(height, width, added_bits) { + return "sppark"; + } + #[cfg(not(feature = "cuda-sppark"))] + let _ = (height, width, added_bits); + "legacy" +} + +/// The backend the quotient's wide forward transform runs on. +pub(crate) fn quotient_backend( + quotient_size: usize, + quotient_degree: usize, + log_blowup: usize, +) -> &'static str { + #[cfg(feature = "cuda-sppark")] + { + let lde_height = (quotient_size / quotient_degree.max(1)) << log_blowup; + if sppark::takes_forward(quotient_size, 2) + || sppark::takes_forward(lde_height, 2 * quotient_degree) + { + return "sppark"; + } + } + #[cfg(not(feature = "cuda-sppark"))] + let _ = (quotient_size, quotient_degree, log_blowup); + "legacy" +} + +/// The backend a lookup LDE of the shape runs on: its committed width is +/// two columns per lookup group. +pub(crate) fn lookup_backend( + height: usize, + num_lookups: usize, + group_size: usize, + log_blowup: usize, +) -> &'static str { + let groups = num_lookups.div_ceil(group_size.max(1)).max(1); + lde_backend(height, 2 * groups, log_blowup) +} + pub(crate) fn quotient_lde_memory_upper_bound( graph: &ConstraintGraph, public_count: usize, @@ -1233,7 +1289,7 @@ pub(crate) fn quotient_lde_mixed( ) -> CudaLde { let _span = tracing::info_span!( "cuda/quotient_lde", - backend = "legacy", + backend = quotient_backend(quotient_size, quotient_degree, log_blowup), device = dft.device_id, quotient_size, quotient_degree, @@ -1786,7 +1842,7 @@ pub(crate) fn lookup_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "direct", - backend = "legacy", + backend = lookup_backend(height, num_lookups, group_size, log_blowup), device = dft.device_id, height, num_lookups, @@ -1878,7 +1934,7 @@ pub(crate) fn lookup_lde_resident_partitioned( let _span = tracing::info_span!( "cuda/lookup_lde", path = "partitioned", - backend = "legacy", + backend = lookup_backend(height, num_lookups, group_size, log_blowup), device = dft.device_id, height, num_lookups, @@ -2076,7 +2132,7 @@ pub(crate) fn lookup_graph_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "graph", - backend = "legacy", + backend = lookup_backend(height, graph.lookups.len(), group_size, log_blowup), device = dft.device_id, height, num_lookups = graph.lookups.len(), diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index b669e3c..8653458 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -39,7 +39,10 @@ unsafe extern "C" { fn multi_stark_sppark_select_backend(selected: c_int); fn multi_stark_sppark_takes(height: usize) -> c_int; fn multi_stark_sppark_transforms_run() -> u64; + fn multi_stark_sppark_takes_lde(height: usize, width: usize, added_bits: usize) -> c_int; + fn multi_stark_sppark_takes_forward(height: usize, width: usize) -> c_int; fn multi_stark_sppark_panel_bytes(height: usize, width: usize, added_bits: usize) -> usize; + fn multi_stark_sppark_forward_panel_bytes(height: usize, width: usize) -> usize; fn multi_stark_sppark_ntt_device( device: c_int, d_inout: *mut u64, @@ -94,11 +97,38 @@ pub fn transforms_run() -> u64 { unsafe { multi_stark_sppark_transforms_run() } } -/// Whether a resident LDE of `height` input rows takes the sppark path. +/// Whether a transform of `height` input rows is tall enough for the sppark +/// path; the shape rules below decide a dispatch. pub fn takes(height: usize) -> bool { unsafe { multi_stark_sppark_takes(height) != 0 } } +/// Whether a resident coset LDE of the shape takes the sppark path: tall +/// enough, within upstream's compiled domain after expansion, and one +/// column's scratch within the panel budget. +pub fn takes_lde(height: usize, width: usize, added_bits: usize) -> bool { + unsafe { multi_stark_sppark_takes_lde(height, width, added_bits) != 0 } +} + +/// Whether a forward transform of the shape takes the sppark path. +pub fn takes_forward(height: usize, width: usize) -> bool { + unsafe { multi_stark_sppark_takes_forward(height, width) != 0 } +} + +/// The scratch one forward transform of the shape allocates on the sppark +/// path; zero when the shape stays on the first-party kernels. +pub fn forward_panel_bytes(height: usize, width: usize) -> usize { + unsafe { multi_stark_sppark_forward_panel_bytes(height, width) } +} + +/// Serializes tests that switch the process-wide backend, so a comparison +/// sees the backend it selected on both of its constructions. +#[cfg(test)] +pub(crate) fn backend_lock() -> std::sync::MutexGuard<'static, ()> { + static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) +} + /// The scratch the sppark path allocates for that LDE, to admit alongside /// the trace and the LDE; zero when the first-party kernels take it. pub fn panel_bytes(height: usize, width: usize, added_bits: usize) -> usize { @@ -277,6 +307,7 @@ mod tests { #[test] fn resident_lde_matches_the_first_party_kernels_bit_for_bit() { + let _guard = backend_lock(); let mut rng = SmallRng::seed_from_u64(0x1de5); for log_height in [0usize, 1, 2, 5, 8, 12, 14] { for added_bits in [0usize, 1, 2, 3] { @@ -294,6 +325,7 @@ mod tests { #[test] fn resident_lde_reduces_raw_representatives_like_the_first_party_kernels() { + let _guard = backend_lock(); let p = Goldilocks::ORDER_U64; let words = [0u64, 1, p - 1, p, p + 1, u64::MAX, 7, p + 7]; let height = 1usize << 10; @@ -313,6 +345,7 @@ mod tests { #[test] fn resident_lde_panels_narrower_than_the_matrix_cover_every_column() { + let _guard = backend_lock(); // A 2^12 x 33 matrix at blowup 2 needs 16 KiB x 2 per column, so a // 256 KiB budget forces panels of a few columns. let mut rng = SmallRng::seed_from_u64(0x9a7e); @@ -328,6 +361,7 @@ mod tests { #[test] fn general_dft_matches_the_first_party_kernels_bit_for_bit() { + let _guard = backend_lock(); // Shapes above the CUDA DFT threshold of 2^15 cells. let mut rng = SmallRng::seed_from_u64(0xdf7); let dft = super::super::CudaDft::new(0); @@ -353,6 +387,7 @@ mod tests { #[test] fn host_coset_lde_matches_the_first_party_kernels_bit_for_bit() { + let _guard = backend_lock(); // The host entry takes matrices of width at most two whose extended // height reaches 2^15. let mut rng = SmallRng::seed_from_u64(0x1e5); @@ -389,13 +424,24 @@ mod tests { #[test] fn the_height_threshold_and_the_panel_budget_decide_dispatch_and_scratch() { + let _guard = backend_lock(); select_backend(Backend::Sppark); assert!( !takes(1 << 12), "short transforms stay on the first-party kernels" ); assert!(takes(1 << 20)); + assert!(takes_lde(1 << 20, 533, 2)); + assert!(takes_forward(1 << 20, 2)); + assert!( + !takes_lde(1 << 24, 6, 5), + "2^29 output rows are beyond the compiled domain" + ); + assert!(!takes_forward(1 << 29, 2)); assert_eq!(panel_bytes(1 << 12, 533, 2), 0); + assert_eq!(panel_bytes(1 << 24, 6, 5), 0); + // A forward transform's panel is one column set at the height. + assert_eq!(forward_panel_bytes(1 << 22, 2), 2 * (1 << 22) * 8); // 2^20 rows, 533 columns, blowup 4: 2 x 2^22 x 8 bytes per column is // 64 MiB, so a 4 GiB budget admits 64 columns. assert_eq!(panel_bytes(1 << 20, 533, 2), 64 * 2 * (1 << 22) * 8); diff --git a/src/types.rs b/src/types.rs index 9b74931..2052746 100644 --- a/src/types.rs +++ b/src/types.rs @@ -570,6 +570,15 @@ impl StarkGenericConfig for GoldilocksBlake3Config { .saturating_div(2) .saturating_add(quotient_degree) .saturating_mul(size_of::()); + // The sppark path's column panels for the two forward + // transforms, zero when they stay on the first-party kernels. + #[cfg(feature = "cuda-sppark")] + let kernel_workspace = kernel_workspace + .saturating_add(crate::cuda::sppark::forward_panel_bytes(quotient_size, 2)) + .saturating_add(crate::cuda::sppark::forward_panel_bytes( + lde_height, + 2 * quotient_degree, + )); ( index, output_bytes, @@ -840,11 +849,18 @@ impl StarkGenericConfig for GoldilocksBlake3Config { ) }) .flatten(); + // The sppark path's column panels for the lookup LDE, on + // either path; zero when it stays on the first-party kernels. + #[cfg(feature = "cuda-sppark")] + let sppark_panel = + crate::cuda::sppark::panel_bytes(height, 2 * groups, self.log_blowup); + #[cfg(not(feature = "cuda-sppark"))] + let sppark_panel = 0; ( index, output_bytes, - direct_temporary_bytes, - graph_memory.map(|(_, temporary)| temporary), + direct_temporary_bytes.saturating_add(sppark_panel), + graph_memory.map(|(_, temporary)| temporary.saturating_add(sppark_panel)), extended_height, ) }) From 9c517949b0c321d6a11a6f21aaa28ca8402a40a3 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:39:47 +0000 Subject: [PATCH 22/34] cuda: sppark panels transform their columns in batched launch sequences The fork's batched entry runs a group of panel columns through one launch sequence per stage. A group is sized to the device's L2 cache (MULTI_STARK_SPPARK_BATCH_BYTES overrides it, 0 launches every column on its own): short columns share launches, and a group of tall columns keeps the reuse between stages that one column at a time had. The coefficient panel is compact, so an LDE's scratch is (N + M) words per column rather than 2M. The build script reruns when the fork's NTT or runtime sources change. On one RTX PRO 6000 the resident LDE medians against the unbatched baseline: 2^16 x 925 x2 52.4 to 20.8 ms, 2^18 x 128 x2 15.8 to 11.7 ms, 2^20 x 533 x4 343 to 336 ms, the tall narrow shapes unchanged; every shape now at or ahead of the first-party kernels. --- build.rs | 5 ++ cuda/sppark_ntt.cu | 126 ++++++++++++++++++++++++++++++++------------- src/cuda/sppark.rs | 108 ++++++++++++++++++++++++++++++++++++-- 3 files changed, 200 insertions(+), 39 deletions(-) diff --git a/build.rs b/build.rs index bb7ccb5..762b763 100644 --- a/build.rs +++ b/build.rs @@ -14,6 +14,11 @@ fn main() { println!("cargo:rerun-if-changed=cuda/kernels.cu"); println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); println!("cargo:rerun-if-changed=cuda/sppark_ntt.cu"); + if let Some(root) = env::var_os("DEP_SPPARK_ROOT") { + let root = PathBuf::from(root); + println!("cargo:rerun-if-changed={}", root.join("ntt").display()); + println!("cargo:rerun-if-changed={}", root.join("util").display()); + } println!("cargo:rerun-if-env-changed=NVCC"); println!("cargo:rerun-if-env-changed=CUDA_HOME"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 32a2ff8..82a2c15 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -43,13 +43,16 @@ bool valid_arguments(const void* d_inout, uint32_t lg, int order, int direction, extern "C" int multi_stark_sppark_max_lg_domain() { return MAX_LG_DOMAIN_SIZE; } -// One in-place transform of 2^lg field elements at `d_inout` on `device`. +// `batch` in-place transforms of 2^lg field elements each, `stride` +// elements apart from `d_inout` on `device`, in one launch sequence. // `order` is NTT::InputOutputOrder (NN, NR, RN, RR), `direction` 0 forward // or 1 inverse, `coset` 1 for the multiplicative coset by the field // generator. Inverse transforms are normalized by 1/2^lg upstream. -extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint32_t lg, - int order, int direction, int coset) { - if (!valid_arguments(d_inout, lg, order, direction, coset)) +extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* d_inout, uint32_t lg, + int order, int direction, int coset, + uint32_t batch, size_t stride) { + if (!valid_arguments(d_inout, lg, order, direction, coset) || batch == 0 || batch > 65535 || + (batch > 1 && stride < (size_t(1) << lg))) return static_cast(cudaErrorInvalidValue); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); @@ -69,10 +72,10 @@ extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint // everything past the launch either completes or never returns. stream_t stream(gpu.id()); stream.wait(events.before); - NTT::Base_dev_ptr(stream, reinterpret_cast(d_inout), lg, - static_cast(order), - static_cast(direction), - static_cast(coset)); + NTT::Base_dev_ptr_batch(stream, reinterpret_cast(d_inout), lg, + static_cast(order), + static_cast(direction), + static_cast(coset), batch, stride); stream.record(events.after); status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); // Only once the caller's stream waits on the transform may the private @@ -82,13 +85,22 @@ extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint return static_cast(status); } -// The same transform on host memory: uploads, transforms, downloads and -// synchronizes. For contract checks and small inputs, not the prover. -extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t lg, int order, - int direction, int coset) { - if (!valid_arguments(inout, lg, order, direction, coset)) +// One transform: the batch entry with a single vector. +extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint32_t lg, + int order, int direction, int coset) { + return multi_stark_sppark_ntt_batch_device(device, d_inout, lg, order, direction, coset, 1, + size_t(1) << lg); +} + +// The batched transform on host memory: `batch * stride` words uploaded, +// transformed, downloaded and synchronized. For contract checks and small +// inputs, not the prover. +extern "C" int multi_stark_sppark_ntt_batch_host(int device, uint64_t* inout, uint32_t lg, int order, + int direction, int coset, uint32_t batch, + size_t stride) { + if (!valid_arguments(inout, lg, order, direction, coset) || batch == 0 || stride < (size_t(1) << lg)) return static_cast(cudaErrorInvalidValue); - const size_t bytes = (size_t(1) << lg) * sizeof(uint64_t); + const size_t bytes = batch * stride * sizeof(uint64_t); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); uint64_t* d_inout = nullptr; @@ -96,7 +108,8 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t if (status != cudaSuccess) return static_cast(status); status = cudaMemcpyAsync(d_inout, inout, bytes, cudaMemcpyHostToDevice, cudaStreamPerThread); int result = static_cast(status); - if (result == 0) result = multi_stark_sppark_ntt_device(device, d_inout, lg, order, direction, coset); + if (result == 0) + result = multi_stark_sppark_ntt_batch_device(device, d_inout, lg, order, direction, coset, batch, stride); if (result == 0) result = static_cast( cudaMemcpyAsync(inout, d_inout, bytes, cudaMemcpyDeviceToHost, cudaStreamPerThread)); @@ -106,6 +119,12 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t return result; } +// One transform on host memory: the batch entry with a single vector. +extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t lg, int order, + int direction, int coset) { + return multi_stark_sppark_ntt_batch_host(device, inout, lg, order, direction, coset, 1, size_t(1) << lg); +} + // --- Resident coset LDE through sppark --------------------------------- // // The prover's matrices are row-major with the transform along the column, @@ -120,8 +139,12 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t // values[r][f + c] = B[c][r] r < M // // The gather reduces representatives at or above the modulus, which upstream -// does not accept. Two panels of C columns cost 16 * M * C bytes; C is sized -// from MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. +// does not accept. The compact panel A and the extended panel B cost +// 8 * (N + M) * C bytes for C columns; C is sized from +// MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. The columns +// of a panel go through batched launch sequences (the fork's grid +// dimension) in groups sized to the L2 cache; MULTI_STARK_SPPARK_BATCH_BYTES +// sets the group, 0 launches every column on its own. #include "goldilocks.cuh" @@ -152,6 +175,7 @@ __global__ void gather_columns(const uint64_t* __restrict__ trace, size_t height } } +// `coefficients` holds `height` bit-reversed coefficients per column. __global__ void shift_columns(const uint64_t* __restrict__ coefficients, uint64_t* __restrict__ panel, size_t height, unsigned log_height, size_t extended_height, size_t columns, const uint64_t* __restrict__ shift_powers) { @@ -163,7 +187,7 @@ __global__ void shift_columns(const uint64_t* __restrict__ coefficients, uint64_ const size_t row = index - column * extended_height; uint64_t value = 0; if (row < height) { - const uint64_t coefficient = coefficients[column * extended_height + reverse_bits(row, log_height)]; + const uint64_t coefficient = coefficients[column * height + reverse_bits(row, log_height)]; value = multi_stark_cuda::goldilocks_mul(coefficient, shift_powers[row]); } panel[index] = value; @@ -220,10 +244,12 @@ unsigned min_log_height() { } // The columns one panel holds at `column_bytes` each: as many as the -// budget admits, at most the width. Zero when one column does not fit, -// which the dispatch rules decline before any allocation. +// budget admits, at most the width and the batch a launch grid can carry. +// Zero when one column does not fit, which the dispatch rules decline +// before any allocation. size_t panel_columns(size_t width, size_t column_bytes) { - const size_t columns = panel_budget_bytes() / column_bytes; + size_t columns = panel_budget_bytes() / column_bytes; + if (columns > 65535) columns = 65535; return columns < width ? columns : width; } @@ -234,6 +260,37 @@ bool within_domain(size_t height, size_t added_bits) { return added_bits <= MAX_LG_DOMAIN_SIZE && log + added_bits <= MAX_LG_DOMAIN_SIZE; } +// The bytes one batched launch sequence keeps in flight: the columns of a +// group go through every stage together, so a group that fits the L2 cache +// keeps the reuse between stages that one column at a time had, while +// short columns still share launches. MULTI_STARK_SPPARK_BATCH_BYTES +// overrides the device's L2 size; 0 launches every column on its own. +size_t batch_group_bytes(int device) { + const unsigned long long setting = decimal_setting("MULTI_STARK_SPPARK_BATCH_BYTES", ~0ull); + if (setting != ~0ull) return static_cast(setting); + int l2 = 0; + if (cudaDeviceGetAttribute(&l2, cudaDevAttrL2CacheSize, device) != cudaSuccess || l2 <= 0) + l2 = 64 << 20; + return static_cast(l2); +} + +// The `count` columns of a panel, `stride` elements apart, through batched +// launch sequences of as many columns as the group budget holds. +int transform_columns(int device, uint64_t* panel, uint32_t lg, int direction, size_t count, + size_t stride) { + if (lg == 0) return 0; + const size_t column_bytes = (size_t(1) << lg) * sizeof(uint64_t); + size_t group = batch_group_bytes(device) / column_bytes; + if (group == 0) group = 1; + int result = 0; + for (size_t first = 0; result == 0 && first < count; first += group) { + const size_t batch = count - first < group ? count - first : group; + result = multi_stark_sppark_ntt_batch_device(device, panel + first * stride, lg, 1, direction, 0, + static_cast(batch), stride); + } + return result; +} + } // namespace // Whether the prover's transforms take the sppark path: MULTI_STARK_CUDA_NTT=sppark, @@ -268,11 +325,11 @@ extern "C" int multi_stark_sppark_takes(size_t height) { // Whether a resident coset LDE of the shape takes the sppark path: tall // enough, the extended height within upstream's compiled domain, and one -// column's scratch within the panel budget. +// column's scratch, the compact and the extended panel, within the budget. extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits) { if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, added_bits)) return 0; const size_t extended_height = height << added_bits; - return 2 * extended_height * sizeof(uint64_t) <= panel_budget_bytes(); + return (height + extended_height) * sizeof(uint64_t) <= panel_budget_bytes(); } // Whether a forward transform of the shape takes the sppark path. @@ -281,12 +338,12 @@ extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width) { return height * sizeof(uint64_t) <= panel_budget_bytes(); } -// The scratch the sppark path allocates for one LDE: two panels of the -// columns the budget admits, sized for the extended height. Zero when the -// shape does not take the path. +// The scratch the sppark path allocates for one LDE: the compact and the +// extended panel of the columns the budget admits. Zero when the shape +// does not take the path. extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { if (!multi_stark_sppark_takes_lde(height, width, added_bits)) return 0; - const size_t column_bytes = 2 * (height << added_bits) * sizeof(uint64_t); + const size_t column_bytes = (height + (height << added_bits)) * sizeof(uint64_t); return panel_columns(width, column_bytes) * column_bytes; } @@ -320,29 +377,27 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u const unsigned log_extended = log_height + static_cast(added_bits); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - const size_t column_bytes = 2 * extended_height * sizeof(uint64_t); + const size_t column_bytes = (height + extended_height) * sizeof(uint64_t); const size_t columns = panel_columns(width, column_bytes); if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* scratch = nullptr; status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes); if (status != cudaSuccess) return static_cast(status); uint64_t* a = scratch; - uint64_t* b = scratch + columns * extended_height; + uint64_t* b = scratch + columns * height; int result = 0; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; gather_columns<<>>( - trace, height, width, first, count, extended_height, a); + trace, height, width, first, count, height, a); result = static_cast(cudaGetLastError()); - for (size_t column = 0; result == 0 && column < count && log_height > 0; ++column) - result = multi_stark_sppark_ntt_device(device, a + column * extended_height, log_height, 1, 1, 0); + if (result == 0) result = transform_columns(device, a, log_height, 1, count, height); if (result == 0) { shift_columns<<>>( a, b, height, log_height, extended_height, count, shift_powers); result = static_cast(cudaGetLastError()); } - for (size_t column = 0; result == 0 && column < count && log_extended > 0; ++column) - result = multi_stark_sppark_ntt_device(device, b + column * extended_height, log_extended, 1, 0, 0); + if (result == 0) result = transform_columns(device, b, log_extended, 0, count, extended_height); if (result == 0) { scatter_columns<<>>( b, extended_height, width, first, count, values); @@ -378,8 +433,7 @@ extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t h gather_columns<<>>( values, height, width, first, count, height, panel); result = static_cast(cudaGetLastError()); - for (size_t column = 0; result == 0 && column < count; ++column) - result = multi_stark_sppark_ntt_device(device, panel + column * height, log_height, 1, 0, 0); + if (result == 0) result = transform_columns(device, panel, log_height, 0, count, height); if (result == 0) { scatter_columns<<>>( panel, height, width, first, count, values); diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 8653458..9ec7238 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -59,6 +59,16 @@ unsafe extern "C" { direction: c_int, coset: c_int, ) -> c_int; + fn multi_stark_sppark_ntt_batch_host( + device: c_int, + inout: *mut u64, + lg: u32, + order: c_int, + direction: c_int, + coset: c_int, + batch: u32, + stride: usize, + ) -> c_int; } /// Whether the prover's resident LDEs take the sppark path: selected by @@ -173,6 +183,41 @@ pub fn try_ntt_host( if status == 0 { Ok(()) } else { Err(status) } } +/// `count` vectors of `2^lg` elements laid out `stride` elements apart. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Batch { + pub lg: u32, + pub count: u32, + pub stride: usize, +} + +/// Transforms the vectors of `batch` in `values` (`count * stride` long) in +/// one batched launch sequence, uploaded, transformed and downloaded within +/// the call. +pub fn try_ntt_batch_host( + device: i32, + values: &mut [Goldilocks], + batch: Batch, + order: Order, + direction: Direction, + coset: bool, +) -> Result<(), i32> { + assert_eq!(values.len(), batch.count as usize * batch.stride); + let status = unsafe { + multi_stark_sppark_ntt_batch_host( + device, + values.as_mut_ptr().cast(), + batch.lg, + order as c_int, + direction as c_int, + c_int::from(coset), + batch.count, + batch.stride, + ) + }; + if status == 0 { Ok(()) } else { Err(status) } +} + /// [`try_ntt_host`], panicking on a CUDA status like the other backends. pub fn ntt_host( device: i32, @@ -422,6 +467,62 @@ mod tests { } } + /// The fork's batched launch sequence transforms every vector of a + /// batch exactly as the single-vector entry does, for each order, + /// direction and coset setting, with padding between the vectors. + #[test] + fn a_batch_of_vectors_matches_the_vectors_one_by_one() { + let mut rng = SmallRng::seed_from_u64(0xba7c); + for (lg, count, padding) in [(4u32, 3u32, 0usize), (10, 7, 8), (12, 5, 0), (16, 3, 16)] { + let length = 1usize << lg; + let stride = length + padding; + let batch = Batch { lg, count, stride }; + let values: Vec = + (0..count as usize * stride).map(|_| rng.random()).collect(); + for order in [Order::NN, Order::NR, Order::RN, Order::RR] { + for direction in [Direction::Forward, Direction::Inverse] { + for coset in [false, true] { + let mut batched = values.clone(); + try_ntt_batch_host(0, &mut batched, batch, order, direction, coset) + .unwrap(); + let mut expected = values.clone(); + for vector in expected.chunks_exact_mut(stride) { + ntt_host(0, &mut vector[..length], order, direction, coset); + } + assert_eq!( + raw_words(&batched), + raw_words(&expected), + "2^{lg} x {count} stride {stride} {order:?} {direction:?} coset {coset}" + ); + } + } + } + } + } + + /// With column batching switched off the panel path launches its + /// columns one by one and still matches the first-party kernels. + #[test] + fn unbatched_columns_match_the_first_party_kernels() { + let _guard = backend_lock(); + let previous = std::env::var("MULTI_STARK_SPPARK_BATCH_BYTES").ok(); + unsafe { std::env::set_var("MULTI_STARK_SPPARK_BATCH_BYTES", "0") }; + let mut rng = SmallRng::seed_from_u64(0x0ff); + for (log_height, width, added_bits) in [(10usize, 5usize, 2usize), (12, 33, 1), (8, 3, 3)] { + let matrix = RowMajorMatrix::new( + (0..(1 << log_height) * width) + .map(|_| rng.random()) + .collect(), + width, + ); + resident_lde_both_ways(&matrix, added_bits, Goldilocks::GENERATOR); + } + match previous { + Some(value) => unsafe { std::env::set_var("MULTI_STARK_SPPARK_BATCH_BYTES", value) }, + None => unsafe { std::env::remove_var("MULTI_STARK_SPPARK_BATCH_BYTES") }, + } + } + #[test] fn the_height_threshold_and_the_panel_budget_decide_dispatch_and_scratch() { let _guard = backend_lock(); @@ -442,9 +543,10 @@ mod tests { assert_eq!(panel_bytes(1 << 24, 6, 5), 0); // A forward transform's panel is one column set at the height. assert_eq!(forward_panel_bytes(1 << 22, 2), 2 * (1 << 22) * 8); - // 2^20 rows, 533 columns, blowup 4: 2 x 2^22 x 8 bytes per column is - // 64 MiB, so a 4 GiB budget admits 64 columns. - assert_eq!(panel_bytes(1 << 20, 533, 2), 64 * 2 * (1 << 22) * 8); + // 2^20 rows, 533 columns, blowup 4: (2^20 + 2^22) x 8 bytes per + // column is 40 MiB, so a 4 GiB budget admits 102 columns. + let column_bytes = ((1 << 20) + (1 << 22)) * 8; + assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes); select_backend(Backend::SpparkAllHeights); assert!(takes(2)); select_backend(Backend::Legacy); From 0ac94c1ac0a6826e4569219976b6877a9860c7fe Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:42:14 +0000 Subject: [PATCH 23/34] cuda: sppark height threshold at 2^18, short shapes in the resident bench With the columns batched, the resident LDE through sppark is ahead of the first-party kernels from 2^18 input rows up (0.97 to 1.19x at 2^18, 1.07 to 1.22x at 2^19) and behind below (0.64 to 0.93x from 2^8 to 2^17), so the default threshold comes down from 2^20. The bench covers those short shapes now. --- cuda/sppark_ntt.cu | 8 ++++---- examples/cuda_resident_lde_bench.rs | 12 ++++++++++++ src/cuda/sppark.rs | 6 ++++-- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 82a2c15..4bfff58 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -240,7 +240,7 @@ size_t panel_budget_bytes() { std::atomic backend_flag{-1}; unsigned min_log_height() { - return unsigned(decimal_setting("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT", 20)); + return unsigned(decimal_setting("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT", 18)); } // The columns one panel holds at `column_bytes` each: as many as the @@ -313,9 +313,9 @@ extern "C" void multi_stark_sppark_select_backend(int selected) { } // Whether a transform of `height` input rows is tall enough for the sppark -// path. Short transforms are launch-bound on the per-column baseline and -// stay on the first-party kernels below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT -// (20). +// path. Below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (18) the panel's gather and +// scatter passes cost more than the first-party kernels save, and the +// first-party kernels stay. extern "C" int multi_stark_sppark_takes(size_t height) { const int flag = multi_stark_sppark_backend_selected(); if (flag == 2) return 1; diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs index 658f9c5..cce1c89 100644 --- a/examples/cuda_resident_lde_bench.rs +++ b/examples/cuda_resident_lde_bench.rs @@ -37,6 +37,18 @@ fn main() { (24, 17, 2), (22, 2, 2), (20, 2, 2), + // Short shapes below the sppark height threshold: the aggregation + // circuits' tables and narrow lookup and quotient codewords. + (8, 20, 2), + (10, 33, 2), + (12, 8, 2), + (14, 40, 2), + (15, 129, 2), + (16, 2, 2), + (17, 16, 2), + (18, 2, 2), + (19, 2, 2), + (19, 49, 2), ] { let height = 1 << log_height; let values = (0..height * width) diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 9ec7238..3565e5f 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -83,8 +83,8 @@ pub enum Backend { /// The first-party kernels only. Legacy, /// sppark for LDEs at or above the height threshold - /// (`MULTI_STARK_SPPARK_MIN_LOG_HEIGHT`, 20), where the per-column - /// baseline wins; the first-party kernels below it. + /// (`MULTI_STARK_SPPARK_MIN_LOG_HEIGHT`, 18), where the panel path + /// wins; the first-party kernels below it. Sppark, /// sppark at every height: for comparing the paths on small shapes. SpparkAllHeights, @@ -531,6 +531,8 @@ mod tests { !takes(1 << 12), "short transforms stay on the first-party kernels" ); + assert!(!takes(1 << 17)); + assert!(takes(1 << 18)); assert!(takes(1 << 20)); assert!(takes_lde(1 << 20, 533, 2)); assert!(takes_forward(1 << 20, 2)); From e7be72c4120f368e67d4731f62595b189813c891 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:03:25 +0000 Subject: [PATCH 24/34] cuda: sppark panel glue as tiled transposes, the fused expansion measured The gather, expansion and scatter passes around the transforms were per-element kernels with a 64-bit division each and strided accesses; they took 227 of the 349 ms of the BLAKE3 shape's LDE against 122 ms of transforms. The gather and scatter are shared-memory tiled transposes now, or row-per-thread kernels for matrices narrower than eight columns, and the expansion runs one column per grid row. The admission counts the reversed coset powers. The fused expansion of the plan (bit-reversed coefficients spread with the coset powers straight into the forward transform in RN order, the row permutation folded into the scatter) is implemented behind MULTI_STARK_SPPARK_FUSED and measured: equal to the restoring path on the wide shape and 15 to 18 percent behind on the tall ones, since its reversing scatter costs what the restoring pass saves, so the restoring path stays the default. MULTI_STARK_SPPARK_STAGE_TIMING prints the per-stage times of every LDE; the resident bench takes a shape list. Resident LDE medians on one RTX PRO 6000 against the first-party kernels: 2^20 x 533 x4 385 to 284 ms, 2^24 x 6 x4 169 to 63 ms, 2^22 x 2 x4 9.5 to 4.3 ms, 2^18 x 128 x2 13.9 to 11.4 ms; the shapes below 2^18 rows stay on the first-party kernels. --- cuda/sppark_ntt.cu | 323 +++++++++++++++++++++------- examples/cuda_resident_lde_bench.rs | 20 ++ src/cuda/sppark.rs | 5 +- 3 files changed, 270 insertions(+), 78 deletions(-) diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 4bfff58..bb23731 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -128,23 +129,28 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t // --- Resident coset LDE through sppark --------------------------------- // // The prover's matrices are row-major with the transform along the column, -// while upstream transforms one contiguous vector. A panel of columns is -// gathered into column-major scratch, transformed column by column, and -// scattered back in the bit-reversed row order the commitment expects: +// while upstream transforms contiguous vectors. A panel of columns is +// gathered into column-major scratch, transformed, and scattered back in +// the bit-reversed row order the commitment expects: // -// A[c][r] = canonical(trace[r][f + c]) r < N -// inverse NR on A[c], normalized upstream (bit-reversed coeffs) -// B[c][i] = A[c][rev(i)] * shift^i i < N, zero beyond -// forward NR on B[c] over M = N << added_bits (bit-reversed evals) -// values[r][f + c] = B[c][r] r < M +// A[c][r] = canonical(trace[r][f + c]) r < N +// inverse NR on A[c], normalized upstream (bit-reversed coeffs) +// B[c][i << added_bits] = A[c][i] * shift^rev(i) i < N, zero elsewhere +// forward RN on B[c] over M = N << added_bits (natural-order evals) +// values[rev(r)][f + c] = B[c][r] r < M // -// The gather reduces representatives at or above the modulus, which upstream -// does not accept. The compact panel A and the extended panel B cost -// 8 * (N + M) * C bytes for C columns; C is sized from -// MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. The columns -// of a panel go through batched launch sequences (the fork's grid -// dimension) in groups sized to the L2 cache; MULTI_STARK_SPPARK_BATCH_BYTES -// sets the group, 0 launches every column on its own. +// That is the fused expansion: it feeds the forward transform in +// bit-reversed order, no pass restores the coefficient order, and the row +// permutation folds into the scatter. The default restores the order +// instead (B[c][i] = A[c][rev(i)] * shift^i), transforms in NR order and +// scatters naturally; `fused_expansion` decides and records why. The gather reduces representatives at or +// above the modulus, which upstream does not accept. The compact panel A +// and the extended panel B cost 8 * (N + M) * C bytes for C columns; C is +// sized from MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. +// The columns of a panel go through batched launch sequences (the fork's +// grid dimension) in groups sized to the L2 cache; +// MULTI_STARK_SPPARK_BATCH_BYTES sets the group, 0 launches every column +// on its own. #include "goldilocks.cuh" @@ -152,57 +158,163 @@ namespace { constexpr unsigned PANEL_THREADS = 256; constexpr size_t MAX_BLOCKS = 65535; +constexpr unsigned TILE = 32; +constexpr unsigned TILE_ROWS = 8; +// Matrices narrower than this take the row-per-thread gather and scatter, +// whose per-thread row segments are then a few contiguous words; from here +// up the tiled transposes win. +constexpr size_t NARROW_MATRIX = 8; unsigned blocks_for_total(size_t total) { const size_t blocks = (total + PANEL_THREADS - 1) / PANEL_THREADS; return static_cast(blocks < MAX_BLOCKS ? blocks : MAX_BLOCKS); } +// Tiles over rows on the grid's first dimension, which is wide enough for +// 2^26 rows, and over columns on the second. +dim3 tile_grid(size_t rows, size_t columns) { + return dim3(static_cast((rows + TILE - 1) / TILE), static_cast((columns + TILE - 1) / TILE)); +} + __device__ __forceinline__ size_t reverse_bits(size_t index, unsigned log) { return log == 0 ? 0 : static_cast(__brev(static_cast(index)) >> (32 - log)); } -__global__ void gather_columns(const uint64_t* __restrict__ trace, size_t height, size_t width, - size_t first, size_t columns, size_t extended_height, - uint64_t* __restrict__ panel) { - const size_t total = height * columns; - const size_t stride = static_cast(blockDim.x) * gridDim.x; - for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; - index += stride) { - const size_t row = index / columns; - const size_t column = index - row * columns; - panel[column * extended_height + row] = multi_stark_cuda::canonicalize(trace[row * width + first + column]); +// Row-major rows [0, height) of columns [first, first + count) into the +// column-major panel (columns `column_stride` apart), canonical, through +// a shared-memory tile so both sides are coalesced. +__global__ void gather_tiles(const uint64_t* __restrict__ source, size_t height, size_t width, size_t first, + size_t count, size_t column_stride, uint64_t* __restrict__ panel) { + __shared__ uint64_t tile[TILE][TILE + 1]; + const size_t row0 = static_cast(blockIdx.x) * TILE; + const size_t col0 = static_cast(blockIdx.y) * TILE; + for (unsigned r = threadIdx.y; r < TILE; r += TILE_ROWS) { + const size_t row = row0 + r, col = col0 + threadIdx.x; + if (row < height && col < count) + tile[r][threadIdx.x] = multi_stark_cuda::canonicalize(source[row * width + first + col]); + } + __syncthreads(); + for (unsigned c = threadIdx.y; c < TILE; c += TILE_ROWS) { + const size_t col = col0 + c, row = row0 + threadIdx.x; + if (row < height && col < count) panel[col * column_stride + row] = tile[threadIdx.x][c]; } } -// `coefficients` holds `height` bit-reversed coefficients per column. -__global__ void shift_columns(const uint64_t* __restrict__ coefficients, uint64_t* __restrict__ panel, +// The panel back into row-major rows of columns [first, first + count); +// with `reverse_rows`, panel row r lands in output row rev(r) over +// 2^log_rows rows, so a natural-order transform result is stored in the +// bit-reversed row order the commitment expects. +__global__ void scatter_tiles(const uint64_t* __restrict__ panel, size_t rows, size_t column_stride, size_t width, + size_t first, size_t count, unsigned log_rows, bool reverse_rows, + uint64_t* __restrict__ values) { + __shared__ uint64_t tile[TILE][TILE + 1]; + const size_t row0 = static_cast(blockIdx.x) * TILE; + const size_t col0 = static_cast(blockIdx.y) * TILE; + for (unsigned c = threadIdx.y; c < TILE; c += TILE_ROWS) { + const size_t col = col0 + c, row = row0 + threadIdx.x; + if (row < rows && col < count) tile[c][threadIdx.x] = panel[col * column_stride + row]; + } + __syncthreads(); + for (unsigned r = threadIdx.y; r < TILE; r += TILE_ROWS) { + const size_t row = row0 + r, col = col0 + threadIdx.x; + if (row < rows && col < count) { + const size_t out = reverse_rows ? reverse_bits(row, log_rows) : row; + values[out * width + first + col] = tile[threadIdx.x][r]; + } + } +} + +// The expansion with the coefficient order restored: B[c][i] = +// A[c][rev(i)] * shift^i for i < N, zero beyond, so the forward transform +// runs in NR order and the scatter stays natural. The permuted reads are +// cheap while the compact panel fits the L2 cache. One column per grid row. +__global__ void shift_columns(const uint64_t* __restrict__ coefficients, const uint64_t* __restrict__ shift_powers, size_t height, unsigned log_height, size_t extended_height, - size_t columns, const uint64_t* __restrict__ shift_powers) { - const size_t total = extended_height * columns; + uint64_t* __restrict__ panel) { + const size_t column = blockIdx.y; const size_t stride = static_cast(blockDim.x) * gridDim.x; - for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; - index += stride) { - const size_t column = index / extended_height; - const size_t row = index - column * extended_height; + for (size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < extended_height; i += stride) { uint64_t value = 0; - if (row < height) { - const uint64_t coefficient = coefficients[column * height + reverse_bits(row, log_height)]; - value = multi_stark_cuda::goldilocks_mul(coefficient, shift_powers[row]); - } - panel[index] = value; + if (i < height) + value = multi_stark_cuda::goldilocks_mul(coefficients[column * height + reverse_bits(i, log_height)], + shift_powers[i]); + panel[column * extended_height + i] = value; + } +} + +// The gather and scatter for a narrow panel: one thread per row moves its +// few words, the panel side coalesced across the block; the tiled +// transpose would leave most of its block idle on such panels. +__global__ void gather_rows(const uint64_t* __restrict__ source, size_t height, size_t width, size_t first, + unsigned count, size_t column_stride, uint64_t* __restrict__ panel) { + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; row < height; row += stride) { + const uint64_t* in = source + row * width + first; + for (unsigned c = 0; c < count; ++c) panel[c * column_stride + row] = multi_stark_cuda::canonicalize(in[c]); } } -__global__ void scatter_columns(const uint64_t* __restrict__ panel, size_t extended_height, size_t width, - size_t first, size_t columns, uint64_t* __restrict__ values) { - const size_t total = extended_height * columns; +__global__ void scatter_rows(const uint64_t* __restrict__ panel, size_t rows, size_t column_stride, size_t width, + size_t first, unsigned count, unsigned log_rows, bool reverse_rows, + uint64_t* __restrict__ values) { + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; row < rows; row += stride) { + const size_t out_row = reverse_rows ? reverse_bits(row, log_rows) : row; + uint64_t* out = values + out_row * width + first; + for (unsigned c = 0; c < count; ++c) out[c] = panel[c * column_stride + row]; + } +} + +// The panel's gather and scatter by the matrix width. +cudaError_t gather_panel(const uint64_t* source, size_t height, size_t width, size_t first, size_t count, + size_t column_stride, uint64_t* panel) { + if (width < NARROW_MATRIX) + gather_rows<<>>( + source, height, width, first, static_cast(count), column_stride, panel); + else + gather_tiles<<>>( + source, height, width, first, count, column_stride, panel); + return cudaGetLastError(); +} + +cudaError_t scatter_panel(const uint64_t* panel, size_t rows, size_t column_stride, size_t width, size_t first, + size_t count, unsigned log_rows, bool reverse_rows, uint64_t* values) { + if (width < NARROW_MATRIX) + scatter_rows<<>>( + panel, rows, column_stride, width, first, static_cast(count), log_rows, reverse_rows, values); + else + scatter_tiles<<>>( + panel, rows, column_stride, width, first, count, log_rows, reverse_rows, values); + return cudaGetLastError(); +} + +// powers[rev(i)] for i < 2^log: the coset powers in the order the +// bit-reversed coefficients meet them. +__global__ void reverse_powers(const uint64_t* __restrict__ powers, size_t count, unsigned log, + uint64_t* __restrict__ reversed) { + const size_t stride = static_cast(blockDim.x) * gridDim.x; + for (size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < count; i += stride) + reversed[i] = powers[reverse_bits(i, log)]; +} + +// The expansion: bit-reversed coefficients A[c][i] (height per column) +// become the bit-reversed input of the forward transform over the extended +// height, B[c][i << added_bits] = A[c][i] * shift^rev(i), zero elsewhere, +// so a forward RN transform yields natural-order evaluations of the coset. +// One column per grid row, rows coalesced along the block. +__global__ void spread_columns(const uint64_t* __restrict__ coefficients, const uint64_t* __restrict__ reversed_powers, + size_t height, size_t extended_height, unsigned added_bits, + uint64_t* __restrict__ panel) { + const size_t column = blockIdx.y; + const size_t mask = (size_t(1) << added_bits) - 1; const size_t stride = static_cast(blockDim.x) * gridDim.x; - for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; - index += stride) { - const size_t row = index / columns; - const size_t column = index - row * columns; - values[row * width + first + column] = panel[column * extended_height + row]; + for (size_t j = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; j < extended_height; j += stride) { + uint64_t value = 0; + if ((j & mask) == 0) { + const size_t i = j >> added_bits; + value = multi_stark_cuda::goldilocks_mul(coefficients[column * height + i], reversed_powers[i]); + } + panel[column * extended_height + j] = value; } } @@ -265,18 +377,68 @@ bool within_domain(size_t height, size_t added_bits) { // keeps the reuse between stages that one column at a time had, while // short columns still share launches. MULTI_STARK_SPPARK_BATCH_BYTES // overrides the device's L2 size; 0 launches every column on its own. -size_t batch_group_bytes(int device) { - const unsigned long long setting = decimal_setting("MULTI_STARK_SPPARK_BATCH_BYTES", ~0ull); - if (setting != ~0ull) return static_cast(setting); +size_t l2_bytes(int device) { int l2 = 0; if (cudaDeviceGetAttribute(&l2, cudaDevAttrL2CacheSize, device) != cudaSuccess || l2 <= 0) l2 = 64 << 20; return static_cast(l2); } +size_t batch_group_bytes(int device) { + const unsigned long long setting = decimal_setting("MULTI_STARK_SPPARK_BATCH_BYTES", ~0ull); + return setting != ~0ull ? static_cast(setting) : l2_bytes(device); +} + +// Whether a panel's expansion feeds the forward transform in bit-reversed +// order with the row permutation in the scatter (fused), or restores the +// coefficient order first and scatters naturally. Measured on the RTX PRO +// 6000, the fused expansion saves the restoring pass but its reversing +// scatter costs as much on wide panels and more on tall narrow ones, so +// restoring is the default. MULTI_STARK_SPPARK_FUSED: 1 always, 2 for +// compact panels beyond the L2 cache, otherwise never. +bool fused_expansion(int device, size_t compact_panel_bytes) { + switch (decimal_setting("MULTI_STARK_SPPARK_FUSED", 0)) { + case 1: return true; + case 2: return compact_panel_bytes > l2_bytes(device); + default: return false; + } +} + +// MULTI_STARK_SPPARK_STAGE_TIMING=1 prints the stage times of every coset +// LDE to stderr: events on the caller's stream around each stage, which +// the transforms are fenced to, and one synchronization per LDE. +struct StageTimer { + static constexpr int STAGES = 5; + bool enabled = decimal_setting("MULTI_STARK_SPPARK_STAGE_TIMING", 0) != 0; + cudaEvent_t marks[STAGES + 1] = {}; + StageTimer() { + if (!enabled) return; + for (auto& mark : marks) + if (cudaEventCreate(&mark) != cudaSuccess) enabled = false; + } + ~StageTimer() { + for (auto mark : marks) + if (mark) cudaEventDestroy(mark); + } + void mark(int stage) { + if (enabled) cudaEventRecord(marks[stage], cudaStreamPerThread); + } + void report(size_t height, size_t width, size_t added_bits) { + if (!enabled || cudaEventSynchronize(marks[STAGES]) != cudaSuccess) return; + static const char* const names[STAGES] = {"gather", "inverse", "spread", "forward", "scatter"}; + fprintf(stderr, "sppark lde height=%zu width=%zu added_bits=%zu", height, width, added_bits); + for (int stage = 0; stage < STAGES; ++stage) { + float ms = 0; + cudaEventElapsedTime(&ms, marks[stage], marks[stage + 1]); + fprintf(stderr, " %s=%.3f", names[stage], ms); + } + fprintf(stderr, "\n"); + } +}; + // The `count` columns of a panel, `stride` elements apart, through batched // launch sequences of as many columns as the group budget holds. -int transform_columns(int device, uint64_t* panel, uint32_t lg, int direction, size_t count, +int transform_columns(int device, uint64_t* panel, uint32_t lg, int order, int direction, size_t count, size_t stride) { if (lg == 0) return 0; const size_t column_bytes = (size_t(1) << lg) * sizeof(uint64_t); @@ -285,7 +447,7 @@ int transform_columns(int device, uint64_t* panel, uint32_t lg, int direction, s int result = 0; for (size_t first = 0; result == 0 && first < count; first += group) { const size_t batch = count - first < group ? count - first : group; - result = multi_stark_sppark_ntt_batch_device(device, panel + first * stride, lg, 1, direction, 0, + result = multi_stark_sppark_ntt_batch_device(device, panel + first * stride, lg, order, direction, 0, static_cast(batch), stride); } return result; @@ -339,12 +501,12 @@ extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width) { } // The scratch the sppark path allocates for one LDE: the compact and the -// extended panel of the columns the budget admits. Zero when the shape -// does not take the path. +// extended panel of the columns the budget admits, and the reversed coset +// powers. Zero when the shape does not take the path. extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { if (!multi_stark_sppark_takes_lde(height, width, added_bits)) return 0; const size_t column_bytes = (height + (height << added_bits)) * sizeof(uint64_t); - return panel_columns(width, column_bytes) * column_bytes; + return panel_columns(width, column_bytes) * column_bytes + height * sizeof(uint64_t); } // The scratch the sppark path allocates for one forward transform. @@ -381,28 +543,42 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u const size_t columns = panel_columns(width, column_bytes); if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* scratch = nullptr; - status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes); + status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes + height * sizeof(uint64_t)); if (status != cudaSuccess) return static_cast(status); uint64_t* a = scratch; uint64_t* b = scratch + columns * height; + uint64_t* powers = b + columns * extended_height; int result = 0; + reverse_powers<<>>( + shift_powers, height, log_height, powers); + result = static_cast(cudaGetLastError()); + StageTimer timer; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; - gather_columns<<>>( - trace, height, width, first, count, height, a); - result = static_cast(cudaGetLastError()); - if (result == 0) result = transform_columns(device, a, log_height, 1, count, height); + timer.mark(0); + result = static_cast(gather_panel(trace, height, width, first, count, height, a)); + timer.mark(1); + if (result == 0) result = transform_columns(device, a, log_height, 1, 1, count, height); + timer.mark(2); + const bool fused = fused_expansion(device, height * count * sizeof(uint64_t)); + const dim3 column_grid(blocks_for_total(extended_height), static_cast(count)); if (result == 0) { - shift_columns<<>>( - a, b, height, log_height, extended_height, count, shift_powers); - result = static_cast(cudaGetLastError()); - } - if (result == 0) result = transform_columns(device, b, log_extended, 0, count, extended_height); - if (result == 0) { - scatter_columns<<>>( - b, extended_height, width, first, count, values); + if (fused) + spread_columns<<>>( + a, powers, height, extended_height, static_cast(added_bits), b); + else + shift_columns<<>>( + a, shift_powers, height, log_height, extended_height, b); result = static_cast(cudaGetLastError()); } + timer.mark(3); + if (result == 0) result = transform_columns(device, b, log_extended, fused ? 2 : 1, 0, count, extended_height); + timer.mark(4); + if (result == 0) + result = static_cast( + scatter_panel(b, extended_height, extended_height, width, first, count, log_extended, fused, values)); + timer.mark(5); + if (result == 0) timer.report(height, count, added_bits); } // The scratch outlives every kernel that reads it: the free is ordered // behind them on the same stream. @@ -430,15 +606,10 @@ extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t h int result = 0; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; - gather_columns<<>>( - values, height, width, first, count, height, panel); - result = static_cast(cudaGetLastError()); - if (result == 0) result = transform_columns(device, panel, log_height, 0, count, height); - if (result == 0) { - scatter_columns<<>>( - panel, height, width, first, count, values); - result = static_cast(cudaGetLastError()); - } + result = static_cast(gather_panel(values, height, width, first, count, height, panel)); + if (result == 0) result = transform_columns(device, panel, log_height, 1, 0, count, height); + if (result == 0) + result = static_cast(scatter_panel(panel, height, height, width, first, count, log_height, false, values)); } const cudaError_t freed = cudaFreeAsync(panel, cudaStreamPerThread); if (result == 0) result = static_cast(freed); diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs index cce1c89..7008ef5 100644 --- a/examples/cuda_resident_lde_bench.rs +++ b/examples/cuda_resident_lde_bench.rs @@ -19,6 +19,23 @@ fn main() { let iterations: usize = std::env::var("MULTI_STARK_CUDA_BENCH_ITERATIONS") .map_or(7, |s| s.parse().expect("invalid iteration count")); let gpu = CudaDft::default(); + // MULTI_STARK_CUDA_BENCH_SHAPES="20,533,2;24,6,2" runs those shapes only. + let only: Vec<(usize, usize, usize)> = std::env::var("MULTI_STARK_CUDA_BENCH_SHAPES") + .map(|list| { + list.split(';') + .map(|shape| { + let mut parts = shape + .split(',') + .map(|part| part.trim().parse().expect("invalid shape")); + ( + parts.next().expect("log height"), + parts.next().expect("width"), + parts.next().expect("added bits"), + ) + }) + .collect() + }) + .unwrap_or_default(); println!("log_height,width,added_bits,iteration,seconds"); // The last five are the shapes the profiled Init proofs commit most: // BLAKE3 pieces, wide and narrow IxVM circuits at the height cap, and @@ -50,6 +67,9 @@ fn main() { (19, 2, 2), (19, 49, 2), ] { + if !only.is_empty() && !only.contains(&(log_height, width, added_bits)) { + continue; + } let height = 1 << log_height; let values = (0..height * width) .map(|index| Goldilocks::from_u64((index as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15))) diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 3565e5f..4b3dd02 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -546,9 +546,10 @@ mod tests { // A forward transform's panel is one column set at the height. assert_eq!(forward_panel_bytes(1 << 22, 2), 2 * (1 << 22) * 8); // 2^20 rows, 533 columns, blowup 4: (2^20 + 2^22) x 8 bytes per - // column is 40 MiB, so a 4 GiB budget admits 102 columns. + // column is 40 MiB, so a 4 GiB budget admits 102 columns, plus the + // reversed coset powers. let column_bytes = ((1 << 20) + (1 << 22)) * 8; - assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes); + assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes + (1 << 20) * 8); select_backend(Backend::SpparkAllHeights); assert!(takes(2)); select_backend(Backend::Legacy); From 1326c657e23c2004ebb130a1bee68bfec1b0d83c Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:15:25 +0000 Subject: [PATCH 25/34] README: the sppark transform backend and its settings --- README.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/README.md b/README.md index 66d2482..6db5c7a 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,29 @@ hardware and downstream Ix measurements live in See [cuda/README.md](cuda/README.md) for architecture, build settings, the NVIDIA correctness harness, platform limitations, and benchmark commands. +### sppark transforms + +The `cuda-sppark` feature adds [sppark](https://github.com/argumentcomputer/sppark)'s +Goldilocks NTT as a second transform backend, pinned to the fork's `dev` +branch and built in its `SPPARK_NO_CXX_RUNTIME` mode, so the archive links +without libstdc++. Proofs are byte-identical on either backend. The backend +is selected at run time: + +| Setting | Effect | +| --- | --- | +| `MULTI_STARK_CUDA_NTT=sppark` | Route transforms of at least the threshold height through sppark; unset, the first-party kernels run | +| `MULTI_STARK_SPPARK_MIN_LOG_HEIGHT=18` | Height threshold; below it the first-party kernels are faster | +| `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Scratch budget per LDE; a shape whose single column exceeds it stays on the first-party kernels | +| `MULTI_STARK_SPPARK_BATCH_BYTES=` | Columns per batched launch sequence, by bytes; defaults to the device's L2 size, 0 launches every column alone | +| `MULTI_STARK_SPPARK_FUSED=0\|1\|2` | The measured alternative expansion: never (default), always, or for coefficient panels beyond the L2 | +| `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each LDE's gather, inverse, expansion, forward and scatter times to stderr | +| `AIUR_METRICS=1` with `RUST_LOG=prover_metrics=info` | The snapshot reports dispatches taken and declined per backend and transform shapes per backend | + +```sh +cargo test --release --features parallel,cuda,cuda-sppark +MULTI_STARK_CUDA_NTT=sppark cargo run --release --features parallel,cuda,cuda-sppark --example cuda_resident_lde_bench +``` + ## License MIT or Apache-2.0 From fde18f8a719b0a3fd3f4fa774a02e353e1337930 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:25:03 +0000 Subject: [PATCH 26/34] README: AIUR_METRICS is a path for Ix, any value for the counters --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6db5c7a..f23f1db 100644 --- a/README.md +++ b/README.md @@ -120,7 +120,7 @@ is selected at run time: | `MULTI_STARK_SPPARK_BATCH_BYTES=` | Columns per batched launch sequence, by bytes; defaults to the device's L2 size, 0 launches every column alone | | `MULTI_STARK_SPPARK_FUSED=0\|1\|2` | The measured alternative expansion: never (default), always, or for coefficient panels beyond the L2 | | `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each LDE's gather, inverse, expansion, forward and scatter times to stderr | -| `AIUR_METRICS=1` with `RUST_LOG=prover_metrics=info` | The snapshot reports dispatches taken and declined per backend and transform shapes per backend | +| `AIUR_METRICS=` with `RUST_LOG=prover_metrics=info` | Any value enables the counters (Ix writes its lightweight metrics to the path); the snapshot reports dispatches taken and declined per backend and transform shapes per backend | ```sh cargo test --release --features parallel,cuda,cuda-sppark From 41fd44c66a1ac49ac1ab0f8ad130b2e9254d2bb2 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:27:22 +0000 Subject: [PATCH 27/34] build: pin sppark at the fork's batched entry and patch record --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 228dfc4..db24f3e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1113,7 +1113,7 @@ dependencies = [ [[package]] name = "sppark" version = "0.1.15" -source = "git+https://github.com/argumentcomputer/sppark?rev=6c5d826de1a36324a6c3e2ce3663139476f9b143#6c5d826de1a36324a6c3e2ce3663139476f9b143" +source = "git+https://github.com/argumentcomputer/sppark?rev=176af27edd74ff51e0e8c601e6aa344b1a854f59#176af27edd74ff51e0e8c601e6aa344b1a854f59" dependencies = [ "cc", "which", diff --git a/Cargo.toml b/Cargo.toml index 1e7c9e4..ea0dcc9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,7 +39,7 @@ p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c837 # adapter compiles against the headers its build script exports. [dependencies.sppark] git = "https://github.com/argumentcomputer/sppark" -rev = "6c5d826de1a36324a6c3e2ce3663139476f9b143" +rev = "176af27edd74ff51e0e8c601e6aa344b1a854f59" optional = true features = ["cuda"] From ce77de2f1f7b014a371063c3226b186aad9ba7e0 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:41:56 +0000 Subject: [PATCH 28/34] cuda: one dispatch decision per transform, tables built for that backend only The Rust side decides a transform's backend once, by the shape, and builds the first-party twiddle tables only for the first-party path; the kernels take the sppark path exactly when no table was handed to them, so a backend switch between the two sides cannot split a transform, and the generic DFT and LDE entries neither build nor upload tables they will not use. Their spans label the backend that ran. The batched host entry checks its word count and stride representability before the FFI, and the adapter rejects a stride beyond upstream's 32-bit index. The bit-reversed feed's reversed coset powers are allocated only when that expansion is selected and count against the panel budget with the panels. --- README.md | 2 +- cuda/kernels.cu | 124 +++++++++++++++++++-------------------- cuda/sppark_ntt.cu | 76 +++++++++++++----------- src/cuda/mod.rs | 143 +++++++++++++++++++++++++++++---------------- src/cuda/sppark.rs | 21 +++++-- 5 files changed, 213 insertions(+), 153 deletions(-) diff --git a/README.md b/README.md index f23f1db..8b41156 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,7 @@ is selected at run time: | `MULTI_STARK_SPPARK_MIN_LOG_HEIGHT=18` | Height threshold; below it the first-party kernels are faster | | `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Scratch budget per LDE; a shape whose single column exceeds it stays on the first-party kernels | | `MULTI_STARK_SPPARK_BATCH_BYTES=` | Columns per batched launch sequence, by bytes; defaults to the device's L2 size, 0 launches every column alone | -| `MULTI_STARK_SPPARK_FUSED=0\|1\|2` | The measured alternative expansion: never (default), always, or for coefficient panels beyond the L2 | +| `MULTI_STARK_SPPARK_FUSED=1` | The measured alternative expansion, feeding the forward transform in bit-reversed order; off by default | | `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each LDE's gather, inverse, expansion, forward and scatter times to stderr | | `AIUR_METRICS=` with `RUST_LOG=prover_metrics=info` | Any value enables the counters (Ix writes its lightweight metrics to the path); the snapshot reports dispatches taken and declined per backend and transform shapes per backend | diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 1099184..6ca288d 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -19,9 +19,6 @@ #include #include "metrics.cuh" #ifdef MULTI_STARK_SPPARK -extern "C" int multi_stark_sppark_takes(size_t height); -extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits); -extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width); extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width); extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, size_t height, size_t width, size_t added_bits, @@ -2256,24 +2253,17 @@ cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, } // namespace -// Whether the shape's inverse-shift-forward sequence, or forward transform, -// runs through sppark; false without the feature. -static bool sppark_takes_lde(size_t height, size_t width, size_t extended_height) { -#ifdef MULTI_STARK_SPPARK - return multi_stark_sppark_takes_lde(height, width, strict_log2(extended_height) - strict_log2(height)) != 0; -#else - (void)height; (void)width; (void)extended_height; - return false; -#endif -} - -static bool sppark_takes_forward(size_t height, size_t width) { -#ifdef MULTI_STARK_SPPARK - return multi_stark_sppark_takes_forward(height, width) != 0; -#else - (void)height; (void)width; - return false; -#endif +// A first-party twiddle table of `count` entries for a transform, cached +// on the device when `legacy`. The Rust side builds no table for a +// transform it sends through sppark, and that absence is the dispatch +// decision on this side; a table required by a first-party transform but +// missing is an argument error rather than a read past an empty buffer. +static cudaError_t legacy_table(int device_id, const uint64_t* host, size_t count, int kind, bool legacy, + const uint64_t** out) { + *out = nullptr; + if (!legacy) return cudaSuccess; + if (!host) return cudaErrorInvalidValue; + return cached_device_constants(device_id, host, count, kind, 0, 0, out); } // The inverse-shift-forward sequence in place on `values`, whose first @@ -2285,8 +2275,12 @@ static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t he size_t extended_height, const uint64_t* inverse_twiddles, const uint64_t* shift_powers, const uint64_t* forward_twiddles, uint64_t height_inverse, bool canonical_pass) { + // The Rust side decides the backend once, by the shape, and hands the + // first-party tables over only for the first-party path; the tables' + // absence is the sppark decision, so a backend switch between the two + // sides cannot split a transform. #ifdef MULTI_STARK_SPPARK - if (sppark_takes_lde(height, width, extended_height)) { + if (forward_twiddles == nullptr) { if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); @@ -2297,6 +2291,8 @@ static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t he device_id, values, values, height, width, added_bits, shift_powers)); } if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); +#else + if (forward_twiddles == nullptr) return cudaErrorInvalidValue; #endif cudaError_t status = launch_dif(device_id, values, height, width, inverse_twiddles); if (status == cudaSuccess) { @@ -2317,7 +2313,7 @@ static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t he static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t height, size_t width, const uint64_t* twiddles) { #ifdef MULTI_STARK_SPPARK - if (sppark_takes_forward(height, width)) { + if (twiddles == nullptr) { if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); @@ -2325,6 +2321,8 @@ static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t heig return static_cast(multi_stark_sppark_forward(device_id, values, height, width)); } if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); +#else + if (twiddles == nullptr) return cudaErrorInvalidValue; #endif return launch_dif(device_id, values, height, width, twiddles); } @@ -2332,10 +2330,11 @@ static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t heig extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, size_t height, size_t width, const uint64_t* twiddles) { - if (values == nullptr || twiddles == nullptr || !is_power_of_two(height) || + if (values == nullptr || !is_power_of_two(height) || width == 0 || !product_fits(height, width)) { return static_cast(cudaErrorInvalidValue); } + const bool legacy = twiddles != nullptr; cudaError_t status = cudaSetDevice(device_id); if (status != cudaSuccess) { return static_cast(status); @@ -2346,7 +2345,7 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, DeviceBuffer device_values; DeviceBuffer device_twiddles; status = copy_to_device(device_values, values, elements); - if (status == cudaSuccess) { + if (status == cudaSuccess && legacy) { status = copy_to_device(device_twiddles, twiddles, height / 2); } if (status == cudaSuccess) { @@ -2363,8 +2362,7 @@ extern "C" int multi_stark_cuda_coset_lde_batch( size_t width, size_t added_bits, const uint64_t* inverse_twiddles, const uint64_t* shift_powers, const uint64_t* forward_twiddles, uint64_t height_inverse) { - if (output == nullptr || input == nullptr || inverse_twiddles == nullptr || - shift_powers == nullptr || forward_twiddles == nullptr || + if (output == nullptr || input == nullptr || shift_powers == nullptr || !is_power_of_two(height) || width == 0 || added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { return static_cast(cudaErrorInvalidValue); @@ -2373,6 +2371,7 @@ extern "C" int multi_stark_cuda_coset_lde_batch( if (!product_fits(extended_height, width)) { return static_cast(cudaErrorInvalidValue); } + const bool legacy = forward_twiddles != nullptr; cudaError_t status = cudaSetDevice(device_id); if (status != cudaSuccess) { @@ -2398,13 +2397,13 @@ extern "C" int multi_stark_cuda_coset_lde_batch( input_elements * sizeof(uint64_t), cudaMemcpyHostToDevice); } - if (status == cudaSuccess) { + if (status == cudaSuccess && legacy) { status = copy_to_device(device_inverse_twiddles, inverse_twiddles, height / 2); } if (status == cudaSuccess) { status = copy_to_device(device_shift_powers, shift_powers, height); } - if (status == cudaSuccess) { + if (status == cudaSuccess && legacy) { status = copy_to_device(device_forward_twiddles, forward_twiddles, extended_height / 2); } @@ -2424,8 +2423,8 @@ static int coset_lde_create( size_t width, size_t added_bits, const uint64_t* inverse_twiddles, const uint64_t* shift_powers, const uint64_t* forward_twiddles, uint64_t height_inverse, void* context, TraceWriter writer) { - if (handle == nullptr || (input == nullptr && writer == nullptr) || (height > 1 && inverse_twiddles == nullptr) || - shift_powers == nullptr || forward_twiddles == nullptr || + if (handle == nullptr || (input == nullptr && writer == nullptr) || + shift_powers == nullptr || !is_power_of_two(height) || width == 0 || added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { return static_cast(cudaErrorInvalidValue); @@ -2460,7 +2459,7 @@ static int coset_lde_create( #ifdef MULTI_STARK_SPPARK // The sppark path writes every output row itself from its own scratch. - const bool sppark = multi_stark_sppark_takes_lde(height, width, added_bits) != 0; + const bool sppark = forward_twiddles == nullptr; if (multi_stark_metrics::enabled()) { multi_stark_metrics::add(device_id, sppark ? multi_stark_metrics::SpparkTaken : multi_stark_metrics::SpparkDeclined, 1); if (sppark) { @@ -2470,6 +2469,10 @@ static int coset_lde_create( } #else const bool sppark = false; + if (forward_twiddles == nullptr) { + destroy_resident_lde(lde); + return static_cast(cudaErrorInvalidValue); + } #endif // Large pageable uploads otherwise serialize through the driver's hidden // staging pool; they go through the persistent staging slots instead. @@ -2514,12 +2517,10 @@ static int coset_lde_create( cudaMemcpyDeviceToDevice, cudaStreamPerThread); } - if (status == cudaSuccess && !sppark && height > 1) { - status = cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&device_inverse_twiddles); - } - if (status == cudaSuccess && !sppark) { - status = cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&device_forward_twiddles); - } + if (status == cudaSuccess) + status = legacy_table(device_id, inverse_twiddles, height / 2, 1, !sppark && height > 1, &device_inverse_twiddles); + if (status == cudaSuccess) + status = legacy_table(device_id, forward_twiddles, extended_height / 2, 2, !sppark, &device_forward_twiddles); if (status == cudaSuccess && !sppark) { status = launch_dif(device_id, lde->values, height, width,device_inverse_twiddles); } @@ -2931,7 +2932,6 @@ extern "C" int multi_stark_cuda_quotient_lde( if (output_handle == nullptr || nodes == nullptr || roots == nullptr || main_handle == nullptr || stage2_handle == nullptr || publics == nullptr || alpha == nullptr || delta == nullptr || - quotient_twiddles == nullptr || lde_twiddles == nullptr || slice_weights == nullptr || node_count == 0 || slot_count == 0 || group_size == 0 || !is_power_of_two(quotient_size) || !is_power_of_two(quotient_degree) || quotient_degree > quotient_size || @@ -2976,8 +2976,8 @@ extern "C" int multi_stark_cuda_quotient_lde( coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess&&!sppark_takes_forward(quotient_size,2))status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); - if(status==cudaSuccess&&!sppark_takes_forward(lde_height,width))status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess)status=legacy_table(device_id,quotient_twiddles,quotient_size/2,2,quotient_twiddles!=nullptr,&device_quotient_twiddles); + if(status==cudaSuccess)status=legacy_table(device_id,lde_twiddles,lde_height/2,2,lde_twiddles!=nullptr,&device_lde_twiddles); if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); @@ -3057,8 +3057,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( if (output_handle == nullptr || nodes == nullptr || roots == nullptr || !one_source(main_handle, main_host) || !one_source(stage2_handle, stage2_host) || publics == nullptr || - alpha == nullptr || delta == nullptr || quotient_twiddles == nullptr || - lde_twiddles == nullptr || slice_weights == nullptr || node_count == 0 || + alpha == nullptr || delta == nullptr || slice_weights == nullptr || node_count == 0 || slot_count == 0 || group_size == 0 || !is_power_of_two(quotient_size) || !is_power_of_two(quotient_degree) || quotient_degree > quotient_size || quotient_size % quotient_degree != 0 || !is_power_of_two(next_step) || @@ -3105,8 +3104,8 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess&&!sppark_takes_forward(quotient_size,2))status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); - if(status==cudaSuccess&&!sppark_takes_forward(lde_height,width))status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess)status=legacy_table(device_id,quotient_twiddles,quotient_size/2,2,quotient_twiddles!=nullptr,&device_quotient_twiddles); + if(status==cudaSuccess)status=legacy_table(device_id,lde_twiddles,lde_height/2,2,lde_twiddles!=nullptr,&device_lde_twiddles); if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); const auto* prep_resident=static_cast(preprocessed_handle); @@ -3456,7 +3455,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(!output_handle||!total||!nodes||!node_count||!slot_count||!lookups||!lookup_count|| (lookup_arg_count&&!lookup_args)||!main|| (!main->trace_values&&!main->host_trace_values&&!main->trace_writer)||!main->trace_height|| - !group_size||!beta||!gamma||!inverse_twiddles||!shift_powers||!forward_twiddles|| + !group_size||!beta||!gamma||!shift_powers|| !is_power_of_two(main->trace_height)||added_bits>=sizeof(size_t)*8|| main->trace_height>(SIZE_MAX>>added_bits)||(prep&&!prep->trace_values)) { return static_cast(cudaErrorInvalidValue); @@ -3525,10 +3524,10 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - const bool legacy_tables=!sppark_takes_lde(height,width,extended_height); - if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + const bool legacy_tables=forward_twiddles!=nullptr; + if(status==cudaSuccess)status=legacy_table(device_id,inverse_twiddles,height/2,1,legacy_tables,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess)status=legacy_table(device_id,forward_twiddles,extended_height/2,2,legacy_tables,&dft); if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); @@ -3543,8 +3542,8 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui const uint64_t* inverse_twiddles,const uint64_t* shift_powers, const uint64_t* forward_twiddles,uint64_t height_inverse){ if(!output_handle||!total||!multiplicities||!arg_offsets||!height||!num_lookups|| - !group_size||!beta||!gamma||!inverse_twiddles||!shift_powers|| - !forward_twiddles||(args_width&&!args)||!is_power_of_two(height)|| + !group_size||!beta||!gamma||!shift_powers|| + (args_width&&!args)||!is_power_of_two(height)|| added_bits>=sizeof(size_t)*8||height>(SIZE_MAX>>added_bits)) return static_cast(cudaErrorInvalidValue); *output_handle=nullptr;const size_t slots=(num_lookups+group_size-1)/group_size; @@ -3589,10 +3588,10 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); scan_done=now(); const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - const bool legacy_tables=!sppark_takes_lde(height,width,extended_height); - if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + const bool legacy_tables=forward_twiddles!=nullptr; + if(status==cudaSuccess)status=legacy_table(device_id,inverse_twiddles,height/2,1,legacy_tables,&dit); if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess&&legacy_tables)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess)status=legacy_table(device_id,forward_twiddles,extended_height/2,2,legacy_tables,&dft); if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, @@ -3748,8 +3747,7 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( const uint64_t* inverse_twiddles, const uint64_t* shift_powers, const uint64_t* forward_twiddles, uint64_t height_inverse) { auto* pending = static_cast(pending_handle); - if (!pending || !output_handle || !total || !inverse_twiddles || - !shift_powers || !forward_twiddles) { + if (!pending || !output_handle || !total || !shift_powers) { return static_cast(cudaErrorInvalidValue); } *output_handle = nullptr; @@ -3769,20 +3767,16 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( const uint64_t* shifts = nullptr; const uint64_t* forward = nullptr; const size_t width = 2 * pending->slots; - const bool legacy_tables = - !sppark_takes_lde(pending->height, width, pending->extended_height); - if (status == cudaSuccess && legacy_tables) - status = cached_device_constants(device_id, inverse_twiddles, - pending->height / 2, 1, 0, 0, &inverse); + const bool legacy_tables = forward_twiddles != nullptr; + if (status == cudaSuccess) + status = legacy_table(device_id, inverse_twiddles, pending->height / 2, 1, legacy_tables, &inverse); if (status == cudaSuccess) status = cached_device_constants(device_id, shift_powers, pending->height, 3, shift_powers[0], pending->height > 1 ? shift_powers[1] : 0, &shifts); - if (status == cudaSuccess && legacy_tables) - status = cached_device_constants(device_id, forward_twiddles, - pending->extended_height / 2, 2, 0, 0, - &forward); + if (status == cudaSuccess) + status = legacy_table(device_id, forward_twiddles, pending->extended_height / 2, 2, legacy_tables, &forward); if (status == cudaSuccess) status = coset_lde_in_place(device_id, pending->lde->values, pending->height, width, pending->extended_height, inverse, shifts, forward, diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index bb23731..a9f3cab 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -17,6 +17,7 @@ #include #include #include +#include namespace { @@ -52,8 +53,9 @@ extern "C" int multi_stark_sppark_max_lg_domain() { return MAX_LG_DOMAIN_SIZE; } extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* d_inout, uint32_t lg, int order, int direction, int coset, uint32_t batch, size_t stride) { + // Upstream indexes with index_t, 32 bits at this domain limit. if (!valid_arguments(d_inout, lg, order, direction, coset) || batch == 0 || batch > 65535 || - (batch > 1 && stride < (size_t(1) << lg))) + (batch > 1 && stride < (size_t(1) << lg)) || stride > size_t(std::numeric_limits::max())) return static_cast(cudaErrorInvalidValue); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); @@ -139,11 +141,12 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t // forward RN on B[c] over M = N << added_bits (natural-order evals) // values[rev(r)][f + c] = B[c][r] r < M // -// That is the fused expansion: it feeds the forward transform in -// bit-reversed order, no pass restores the coefficient order, and the row -// permutation folds into the scatter. The default restores the order -// instead (B[c][i] = A[c][rev(i)] * shift^i), transforms in NR order and -// scatters naturally; `fused_expansion` decides and records why. The gather reduces representatives at or +// That feeds the forward transform in bit-reversed order, so no pass +// restores the coefficient order and the row permutation folds into the +// scatter. The default restores the order instead (B[c][i] = A[c][rev(i)] +// * shift^i), transforms in NR order and scatters naturally; +// `fused_expansion` decides and records why. Neither folds the expansion +// into the transform's first pass; that remains upstream's kernels' job. The gather reduces representatives at or // above the modulus, which upstream does not accept. The compact panel A // and the extended panel B cost 8 * (N + M) * C bytes for C columns; C is // sized from MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. @@ -359,8 +362,8 @@ unsigned min_log_height() { // budget admits, at most the width and the batch a launch grid can carry. // Zero when one column does not fit, which the dispatch rules decline // before any allocation. -size_t panel_columns(size_t width, size_t column_bytes) { - size_t columns = panel_budget_bytes() / column_bytes; +size_t panel_columns(size_t width, size_t column_bytes, size_t budget) { + size_t columns = budget / column_bytes; if (columns > 65535) columns = 65535; return columns < width ? columns : width; } @@ -390,18 +393,17 @@ size_t batch_group_bytes(int device) { } // Whether a panel's expansion feeds the forward transform in bit-reversed -// order with the row permutation in the scatter (fused), or restores the -// coefficient order first and scatters naturally. Measured on the RTX PRO -// 6000, the fused expansion saves the restoring pass but its reversing -// scatter costs as much on wide panels and more on tall narrow ones, so -// restoring is the default. MULTI_STARK_SPPARK_FUSED: 1 always, 2 for -// compact panels beyond the L2 cache, otherwise never. -bool fused_expansion(int device, size_t compact_panel_bytes) { - switch (decimal_setting("MULTI_STARK_SPPARK_FUSED", 0)) { - case 1: return true; - case 2: return compact_panel_bytes > l2_bytes(device); - default: return false; - } +// order with the row permutation in the scatter (MULTI_STARK_SPPARK_FUSED=1), +// or restores the coefficient order first and scatters naturally. Measured +// on the RTX PRO 6000, the bit-reversed feed saves the restoring pass but +// its reversing scatter costs as much on wide panels and more on tall +// narrow ones, so restoring is the default. +bool fused_expansion() { return decimal_setting("MULTI_STARK_SPPARK_FUSED", 0) == 1; } + +// The reversed coset powers the bit-reversed feed reads: one word per input +// row, nothing on the restoring path. +size_t expansion_extra_bytes(size_t height) { + return fused_expansion() ? height * sizeof(uint64_t) : 0; } // MULTI_STARK_SPPARK_STAGE_TIMING=1 prints the stage times of every coset @@ -491,7 +493,7 @@ extern "C" int multi_stark_sppark_takes(size_t height) { extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits) { if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, added_bits)) return 0; const size_t extended_height = height << added_bits; - return (height + extended_height) * sizeof(uint64_t) <= panel_budget_bytes(); + return (height + extended_height) * sizeof(uint64_t) + expansion_extra_bytes(height) <= panel_budget_bytes(); } // Whether a forward transform of the shape takes the sppark path. @@ -500,20 +502,22 @@ extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width) { return height * sizeof(uint64_t) <= panel_budget_bytes(); } -// The scratch the sppark path allocates for one LDE: the compact and the -// extended panel of the columns the budget admits, and the reversed coset -// powers. Zero when the shape does not take the path. +// The scratch the sppark path allocates for one LDE, within the panel +// budget: the compact and the extended panel of the columns the budget +// admits, plus the bit-reversed feed's reversed coset powers when that +// expansion is selected. Zero when the shape does not take the path. extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { if (!multi_stark_sppark_takes_lde(height, width, added_bits)) return 0; const size_t column_bytes = (height + (height << added_bits)) * sizeof(uint64_t); - return panel_columns(width, column_bytes) * column_bytes + height * sizeof(uint64_t); + const size_t extra = expansion_extra_bytes(height); + return panel_columns(width, column_bytes, panel_budget_bytes() - extra) * column_bytes + extra; } // The scratch the sppark path allocates for one forward transform. extern "C" size_t multi_stark_sppark_forward_panel_bytes(size_t height, size_t width) { if (!multi_stark_sppark_takes_forward(height, width)) return 0; const size_t column_bytes = height * sizeof(uint64_t); - return panel_columns(width, column_bytes) * column_bytes; + return panel_columns(width, column_bytes, panel_budget_bytes()) * column_bytes; } // The coset LDE of `trace` (height x width, natural row order, device memory) @@ -539,19 +543,24 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u const unsigned log_extended = log_height + static_cast(added_bits); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); + const bool fused = fused_expansion(); const size_t column_bytes = (height + extended_height) * sizeof(uint64_t); - const size_t columns = panel_columns(width, column_bytes); - if (columns == 0) return static_cast(cudaErrorInvalidValue); + const size_t extra = expansion_extra_bytes(height); + const size_t budget = panel_budget_bytes(); + if (budget < column_bytes + extra) return static_cast(cudaErrorInvalidValue); + const size_t columns = panel_columns(width, column_bytes, budget - extra); uint64_t* scratch = nullptr; - status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes + height * sizeof(uint64_t)); + status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes + extra); if (status != cudaSuccess) return static_cast(status); uint64_t* a = scratch; uint64_t* b = scratch + columns * height; uint64_t* powers = b + columns * extended_height; int result = 0; - reverse_powers<<>>( - shift_powers, height, log_height, powers); - result = static_cast(cudaGetLastError()); + if (fused) { + reverse_powers<<>>( + shift_powers, height, log_height, powers); + result = static_cast(cudaGetLastError()); + } StageTimer timer; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; @@ -560,7 +569,6 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u timer.mark(1); if (result == 0) result = transform_columns(device, a, log_height, 1, 1, count, height); timer.mark(2); - const bool fused = fused_expansion(device, height * count * sizeof(uint64_t)); const dim3 column_grid(blocks_for_total(extended_height), static_cast(count)); if (result == 0) { if (fused) @@ -598,7 +606,7 @@ extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t h if (log_height == 0) return 0; cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - const size_t columns = panel_columns(width, height * sizeof(uint64_t)); + const size_t columns = panel_columns(width, height * sizeof(uint64_t), panel_budget_bytes()); if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* panel = nullptr; status = cudaMalloc(reinterpret_cast(&panel), columns * height * sizeof(uint64_t)); diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 52a2014..89af67a 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -83,6 +83,18 @@ impl CudaDft { self.device_id } + /// The first-party twiddle table for a transform, or none when the + /// transform runs through sppark, which needs no table; the kernels + /// reject a missing table on the first-party path. + fn legacy_twiddles( + &self, + legacy: bool, + log_height: usize, + inverse: bool, + ) -> Option> { + legacy.then(|| self.twiddles(log_height, inverse)) + } + fn twiddles(&self, log_height: usize, inverse: bool) -> Arc<[Goldilocks]> { let key = (log_height, inverse); if let Some(twiddles) = self @@ -195,8 +207,10 @@ impl CudaDft { Self::validate_dimensions(extended_height, width); let log_height = log2_strict_usize(height); - let inverse_twiddles = self.twiddles(log_height, true); - let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let legacy = lde_backend(height, width, added_bits) == "legacy"; + let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); + let forward_twiddles = + self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let shift_powers = self.shift_powers(height, shift); let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); let mut handle = core::ptr::null_mut(); @@ -210,9 +224,9 @@ impl CudaDft { height, width, added_bits, - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(height_inverse), ) }; @@ -239,27 +253,20 @@ impl CudaDft { .expect("LDE height overflows usize"); Self::validate_dimensions(height, 1); Self::validate_dimensions(extended_height, 1); - let legacy_tables = lde_backend(height, 1, added_bits) == "legacy"; - let inverse_twiddles = if legacy_tables { - self.twiddles(log2_strict_usize(height), true) - } else { - Arc::from(Vec::new()) - }; + let legacy = lde_backend(height, 1, added_bits) == "legacy"; + let inverse_twiddles = self.legacy_twiddles(legacy, log2_strict_usize(height), true); let shift_powers = self.shift_powers(height, shift); - let forward_twiddles = if legacy_tables { - self.twiddles(log2_strict_usize(extended_height), false) - } else { - Arc::from(Vec::new()) - }; + let forward_twiddles = + self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let status = unsafe { multi_stark_cuda_prepare_lde_constants( self.device_id, - inverse_twiddles.as_ptr().cast(), - inverse_twiddles.len(), + table_ptr(&inverse_twiddles), + inverse_twiddles.as_ref().map_or(0, |table| table.len()), shift_powers.as_ptr().cast(), height, - forward_twiddles.as_ptr().cast(), - forward_twiddles.len(), + table_ptr(&forward_twiddles), + forward_twiddles.as_ref().map_or(0, |table| table.len()), ) }; check_cuda(status, "prepare resident LDE constants"); @@ -386,8 +393,10 @@ impl CudaDft { .expect("LDE height overflow"); Self::validate_dimensions(extended_height, width); let log_height = log2_strict_usize(height); - let inverse_twiddles = self.twiddles(log_height, true); - let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let legacy = lde_backend(height, width, added_bits) == "legacy"; + let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); + let forward_twiddles = + self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let shift_powers = self.shift_powers(height, shift); let mut context = Box::new(generator); let mut handle = core::ptr::null_mut(); @@ -398,9 +407,9 @@ impl CudaDft { height, width, added_bits, - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(Goldilocks::ONE.div_2exp_u64(log_height as u64)), (&mut *context as *mut Generator).cast(), write_generated_trace, @@ -828,6 +837,25 @@ fn encode_quotient_nodes( /// kernel use the same liveness calculation. The scratch term assumes the /// global-memory path; devices able to fit the slots in shared memory need /// less than this bound. +/// A host table's pointer for the kernels: null when the transform runs +/// through sppark and no table was built. +fn table_ptr(table: &Option>) -> *const u64 { + table + .as_ref() + .map_or(core::ptr::null(), |table| table.as_ptr().cast()) +} + +/// The backend a forward transform of the shape runs on. +pub(crate) fn forward_backend(height: usize, width: usize) -> &'static str { + #[cfg(feature = "cuda-sppark")] + if sppark::takes_forward(height, width) { + return "sppark"; + } + #[cfg(not(feature = "cuda-sppark"))] + let _ = (height, width); + "legacy" +} + /// The backend a resident coset LDE of the shape runs on, for spans and /// for skipping the tables the other backend would need. pub(crate) fn lde_backend(height: usize, width: usize, added_bits: usize) -> &'static str { @@ -1367,8 +1395,16 @@ fn quotient_lde_sources( assert_eq!(alpha.len(), 2 * expected_constraints); let trace_height = quotient_size / quotient_degree; let lde_height = trace_height << log_blowup; - let quotient_twiddles = dft.twiddles(log2_strict_usize(quotient_size), false); - let lde_twiddles = dft.twiddles(log2_strict_usize(lde_height), false); + let quotient_twiddles = dft.legacy_twiddles( + forward_backend(quotient_size, 2) == "legacy", + log2_strict_usize(quotient_size), + false, + ); + let lde_twiddles = dft.legacy_twiddles( + forward_backend(lde_height, 2 * quotient_degree) == "legacy", + log2_strict_usize(lde_height), + false, + ); let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(quotient_size) as u64); let weight_step = Goldilocks::GENERATOR.exp_u64(trace_height as u64).inverse(); let weights: Vec<_> = weight_step @@ -1444,8 +1480,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - quotient_twiddles.as_ptr().cast(), - lde_twiddles.as_ptr().cast(), + table_ptr("ient_twiddles), + table_ptr(&lde_twiddles), weights.as_ptr().cast(), ) } @@ -1484,8 +1520,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - quotient_twiddles.as_ptr().cast(), - lde_twiddles.as_ptr().cast(), + table_ptr("ient_twiddles), + table_ptr(&lde_twiddles), weights.as_ptr().cast(), ) } @@ -1876,9 +1912,10 @@ pub(crate) fn lookup_lde_resident( assert_eq!(multiplicities.len(), height * num_lookups); assert_eq!(args.len(), height * args_width); let extended_height = height << log_blowup; - let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; + let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); let mut tail = [Goldilocks::ZERO; 4]; let mut handle = core::ptr::null_mut(); @@ -1898,9 +1935,9 @@ pub(crate) fn lookup_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(height_inverse), ) }; @@ -1952,9 +1989,10 @@ pub(crate) fn lookup_lde_resident_partitioned( assert_eq!(multiplicities.len(), height * num_lookups); assert_eq!(args.len(), height * args_width); let extended_height = height << log_blowup; - let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; + let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); let total_started = std::time::Instant::now(); let create_started = std::time::Instant::now(); @@ -2085,9 +2123,9 @@ pub(crate) fn lookup_lde_resident_partitioned( pending_handle.as_ptr(), &mut handle, tail.as_mut_ptr().cast(), - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(height_inverse), ) }; @@ -2144,9 +2182,10 @@ pub(crate) fn lookup_graph_lde_resident( let num_lookups = lookups.len(); let groups = num_lookups.div_ceil(group_size.max(1)); let extended_height = height << log_blowup; - let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; + let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); let mut tail = [Goldilocks::ZERO; 4]; let mut handle = core::ptr::null_mut(); @@ -2169,9 +2208,9 @@ pub(crate) fn lookup_graph_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(height_inverse), ) }; @@ -2250,7 +2289,7 @@ impl TwoAdicSubgroupDft for CudaDft { backend = if height == 1 || width == 0 { "noop" } else if Self::use_cuda_dft(height, width) { - "legacy" + forward_backend(height, width) } else { "cpu" } @@ -2263,7 +2302,11 @@ impl TwoAdicSubgroupDft for CudaDft { return self.cpu.dft_batch(matrix); } - let twiddles = self.twiddles(log2_strict_usize(height), false); + let twiddles = self.legacy_twiddles( + forward_backend(height, width) == "legacy", + log2_strict_usize(height), + false, + ); // SAFETY: Goldilocks is repr(transparent) over u64 (asserted above), // every u64 bit pattern is a valid Goldilocks value, all buffers have // the element counts implied by height/width, and the FFI call is @@ -2274,7 +2317,7 @@ impl TwoAdicSubgroupDft for CudaDft { matrix.values.as_mut_ptr().cast(), height, width, - twiddles.as_ptr().cast(), + table_ptr(&twiddles), ) }; check_cuda(status, "batched DFT"); @@ -2306,7 +2349,7 @@ impl TwoAdicSubgroupDft for CudaDft { backend = if width == 0 { "noop" } else if height > 1 && Self::use_cuda_coset_lde(extended_height, width) { - "legacy" + lde_backend(height, width, added_bits) } else { "cpu" } @@ -2328,8 +2371,10 @@ impl TwoAdicSubgroupDft for CudaDft { } let log_height = log2_strict_usize(height); - let inverse_twiddles = self.twiddles(log_height, true); - let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let legacy = lde_backend(height, width, added_bits) == "legacy"; + let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); + let forward_twiddles = + self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); let shift_powers = self.shift_powers(height, shift); let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); let mut output = Goldilocks::zero_vec(extended_height * width); @@ -2345,9 +2390,9 @@ impl TwoAdicSubgroupDft for CudaDft { height, width, added_bits, - inverse_twiddles.as_ptr().cast(), + table_ptr(&inverse_twiddles), shift_powers.as_ptr().cast(), - forward_twiddles.as_ptr().cast(), + table_ptr(&forward_twiddles), raw_u64(height_inverse), ) }; diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 4b3dd02..2f31773 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -202,7 +202,14 @@ pub fn try_ntt_batch_host( direction: Direction, coset: bool, ) -> Result<(), i32> { - assert_eq!(values.len(), batch.count as usize * batch.stride); + let words = (batch.count as usize) + .checked_mul(batch.stride) + .expect("batch words overflow usize"); + assert_eq!(values.len(), words); + assert!( + u32::try_from(batch.stride).is_ok(), + "the adapter indexes vectors with 32 bits" + ); let status = unsafe { multi_stark_sppark_ntt_batch_host( device, @@ -546,10 +553,16 @@ mod tests { // A forward transform's panel is one column set at the height. assert_eq!(forward_panel_bytes(1 << 22, 2), 2 * (1 << 22) * 8); // 2^20 rows, 533 columns, blowup 4: (2^20 + 2^22) x 8 bytes per - // column is 40 MiB, so a 4 GiB budget admits 102 columns, plus the - // reversed coset powers. + // column is 40 MiB, so a 4 GiB budget admits 102 columns, and the + // bit-reversed feed's reversed powers when that expansion is on. let column_bytes = ((1 << 20) + (1 << 22)) * 8; - assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes + (1 << 20) * 8); + let extra = if std::env::var("MULTI_STARK_SPPARK_FUSED").as_deref() == Ok("1") { + (1 << 20) * 8 + } else { + 0 + }; + assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes + extra); + assert!(panel_bytes(1 << 20, 533, 2) <= 4 << 30); select_backend(Backend::SpparkAllHeights); assert!(takes(2)); select_backend(Backend::Legacy); From b6104b6a3d9d581c87a5b3b21ea46161364a57bf Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:12:24 +0000 Subject: [PATCH 29/34] cuda: the sppark panels' scratch comes from the stream's memory pool The panel scratch was a cudaMalloc per LDE freed asynchronously: a device synchronization per transform, and a pairing on which CUPTI 2026.2.1's memory tracking faults (an async free of memory that did not come from a pool dereferences the pool it expects, with the MEMORY2 activity kind on, whatever the stream). Pool allocations on the caller's stream remove both: 2^24 x 6 x4 61 to 55 ms, 2^22 x 2 x4 4.1 to 3.5 ms, and the prover profile collector records the sppark path again. --- cuda/sppark_ntt.cu | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index a9f3cab..035599d 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -549,8 +549,12 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u const size_t budget = panel_budget_bytes(); if (budget < column_bytes + extra) return static_cast(cudaErrorInvalidValue); const size_t columns = panel_columns(width, column_bytes, budget - extra); + // Pool allocations on the caller's stream: no device synchronization + // per LDE, and the async free below returns them to the pool they came + // from. (CUPTI 2026.2.1's memory tracking faults on an async free of + // memory from cudaMalloc, which this pairing also avoids.) uint64_t* scratch = nullptr; - status = cudaMalloc(reinterpret_cast(&scratch), columns * column_bytes + extra); + status = cudaMallocAsync(reinterpret_cast(&scratch), columns * column_bytes + extra, cudaStreamPerThread); if (status != cudaSuccess) return static_cast(status); uint64_t* a = scratch; uint64_t* b = scratch + columns * height; @@ -609,7 +613,7 @@ extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t h const size_t columns = panel_columns(width, height * sizeof(uint64_t), panel_budget_bytes()); if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* panel = nullptr; - status = cudaMalloc(reinterpret_cast(&panel), columns * height * sizeof(uint64_t)); + status = cudaMallocAsync(reinterpret_cast(&panel), columns * height * sizeof(uint64_t), cudaStreamPerThread); if (status != cudaSuccess) return static_cast(status); int result = 0; for (size_t first = 0; result == 0 && first < width; first += columns) { From 9c62a108e11662099193a232e92b84d2dc6cd1bc Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:14:17 +0000 Subject: [PATCH 30/34] Use sppark for all CUDA transforms Share immutable panel plans between admission and execution, and run the fork on borrowed caller streams. Remove the first-party GPU NTT, backend selection, twiddle plumbing, and unused expansion experiment. Keep the generic CPU DFT and compare resident results and proof bytes against CPU references. --- Cargo.lock | 2 +- Cargo.toml | 19 +- README.md | 34 +- build.rs | 6 +- cuda/README.md | 19 +- cuda/kernels.cu | 568 +++------------------------- cuda/metrics.cuh | 8 +- cuda/ntt.cuh | 23 ++ cuda/smoke.sh | 8 +- cuda/sppark_ntt.cu | 372 ++++-------------- examples/cuda_resident_lde_bench.rs | 3 +- src/batch.rs | 24 -- src/cuda/metrics.rs | 23 +- src/cuda/mod.rs | 364 ++++-------------- src/cuda/pcs.rs | 41 +- src/cuda/sppark.rs | 471 +++++++++++------------ src/cuda/witness.rs | 12 +- src/prover.rs | 85 ----- src/types.rs | 59 ++- 19 files changed, 587 insertions(+), 1554 deletions(-) create mode 100644 cuda/ntt.cuh diff --git a/Cargo.lock b/Cargo.lock index db24f3e..f119e7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1113,7 +1113,7 @@ dependencies = [ [[package]] name = "sppark" version = "0.1.15" -source = "git+https://github.com/argumentcomputer/sppark?rev=176af27edd74ff51e0e8c601e6aa344b1a854f59#176af27edd74ff51e0e8c601e6aa344b1a854f59" +source = "git+https://github.com/argumentcomputer/sppark?rev=e10e107673aa22861f0f8b9758fc62169ab919ae#e10e107673aa22861f0f8b9758fc62169ab919ae" dependencies = [ "cc", "which", diff --git a/Cargo.toml b/Cargo.toml index ea0dcc9..7e042d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,13 +33,11 @@ p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89 p3-symmetric = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c83775a8076cc262ffc48a1fd7c" } -# sppark's Goldilocks NTT, for the comparison backend behind `cuda-sppark`, -# from the organization's fork: its `dev` branch adds -# the `SPPARK_NO_CXX_RUNTIME` build mode the adapter relies on. Pinned; the -# adapter compiles against the headers its build script exports. +# The CUDA transform engine. The fork supplies column batching, borrowed +# streams and a runtime compatible with the Lean executable's C++ linkage. [dependencies.sppark] git = "https://github.com/argumentcomputer/sppark" -rev = "176af27edd74ff51e0e8c601e6aa344b1a854f59" +rev = "e10e107673aa22861f0f8b9758fc62169ab919ae" optional = true features = ["cuda"] @@ -56,14 +54,9 @@ harness = false [features] parallel = ["p3-maybe-rayon/parallel"] -# Use the first-party CUDA Goldilocks DFT/LDE backend. Enabling this feature -# requires a CUDA toolkit at build time and an NVIDIA GPU at runtime; the -# default CPU build never invokes nvcc or links the CUDA runtime. -cuda = ["dep:itertools", "dep:rayon"] -# Adds sppark's NTT/LDE engine as a comparison backend for the transforms. -# Requires `cuda`; the first-party kernels remain the default until the -# comparison is accepted. -cuda-sppark = ["cuda", "dep:sppark"] +# CUDA requires the toolkit at build time and an NVIDIA GPU at runtime. +# CPU builds neither compile sppark nor link CUDA. +cuda = ["dep:itertools", "dep:rayon", "dep:sppark"] # Similar to `release`, but preserves debug info [profile.dev-ci] diff --git a/README.md b/README.md index 8b41156..3bbbeca 100644 --- a/README.md +++ b/README.md @@ -106,25 +106,31 @@ correctness harness, platform limitations, and benchmark commands. ### sppark transforms -The `cuda-sppark` feature adds [sppark](https://github.com/argumentcomputer/sppark)'s -Goldilocks NTT as a second transform backend, pinned to the fork's `dev` -branch and built in its `SPPARK_NO_CXX_RUNTIME` mode, so the archive links -without libstdc++. Proofs are byte-identical on either backend. The backend -is selected at run time: +Every CUDA build uses [sppark](https://github.com/argumentcomputer/sppark)'s +Goldilocks NTT, pinned to the fork's `dev` branch. Its borrowed streams keep +panel allocation, transforms and freeing on the caller's stream. Tiny generic +host matrices retain the CPU DFT. CPU builds do not depend on sppark or CUDA. + +An immutable per-device plan fixes each transform's panel size, launch groups +and coset powers. Admission and execution share the plan. Dimensions must fit +the compiled domain (currently 2^28 rows), checked byte arithmetic and at least +one column within the panel budget; invalid shapes fail before allocation. | Setting | Effect | | --- | --- | -| `MULTI_STARK_CUDA_NTT=sppark` | Route transforms of at least the threshold height through sppark; unset, the first-party kernels run | -| `MULTI_STARK_SPPARK_MIN_LOG_HEIGHT=18` | Height threshold; below it the first-party kernels are faster | -| `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Scratch budget per LDE; a shape whose single column exceeds it stays on the first-party kernels | -| `MULTI_STARK_SPPARK_BATCH_BYTES=` | Columns per batched launch sequence, by bytes; defaults to the device's L2 size, 0 launches every column alone | -| `MULTI_STARK_SPPARK_FUSED=1` | The measured alternative expansion, feeding the forward transform in bit-reversed order; off by default | -| `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each LDE's gather, inverse, expansion, forward and scatter times to stderr | -| `AIUR_METRICS=` with `RUST_LOG=prover_metrics=info` | Any value enables the counters (Ix writes its lightweight metrics to the path); the snapshot reports dispatches taken and declined per backend and transform shapes per backend | +| `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Maximum transform scratch; zero selects the default. A budget smaller than one column is a configuration error | +| `MULTI_STARK_SPPARK_BATCH_BYTES=` | Bytes per batched launch group; defaults to the device's L2 size, zero launches one column at a time | +| `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each panel's gather, inverse, expansion, forward and scatter times; synchronizes each measured panel | +| `AIUR_METRICS=` with `RUST_LOG=prover_metrics=info` | Enables lightweight per-device transform-shape counters; Ix writes them to the selected metrics file | + +Panel and batching settings are captured when the device DFT is constructed. +There is no backend selector, height threshold or alternative expansion mode. +The fork's `SPPARK_NO_CXX_RUNTIME` mode aborts on CUDA errors with diagnostics; +normal ix builds also abort on Rust panics. ```sh -cargo test --release --features parallel,cuda,cuda-sppark -MULTI_STARK_CUDA_NTT=sppark cargo run --release --features parallel,cuda,cuda-sppark --example cuda_resident_lde_bench +cargo test --release --features parallel,cuda --lib cuda::sppark::tests:: +MULTI_STARK_CUDA_BENCH_SHAPES="20,533,2" cargo run --release --features parallel,cuda --example cuda_resident_lde_bench ``` ## License diff --git a/build.rs b/build.rs index 762b763..cc4e9f1 100644 --- a/build.rs +++ b/build.rs @@ -14,6 +14,7 @@ fn main() { println!("cargo:rerun-if-changed=cuda/kernels.cu"); println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); println!("cargo:rerun-if-changed=cuda/sppark_ntt.cu"); + println!("cargo:rerun-if-changed=cuda/ntt.cuh"); if let Some(root) = env::var_os("DEP_SPPARK_ROOT") { let root = PathBuf::from(root); println!("cargo:rerun-if-changed={}", root.join("ntt").display()); @@ -52,7 +53,7 @@ fn main() { // ``, whose GNU-only pthread functions that flag would hide. Their // objects then join the archive. let mut sppark_objects = Vec::new(); - if env::var_os("CARGO_FEATURE_CUDA_SPPARK").is_some() { + { let root = PathBuf::from( env::var_os("DEP_SPPARK_ROOT").expect("sppark's build script exports DEP_SPPARK_ROOT"), ); @@ -118,9 +119,6 @@ fn main() { .arg("-o") .arg(&library) .arg("cuda/kernels.cu"); - if !sppark_objects.is_empty() { - command.arg("-DMULTI_STARK_SPPARK"); - } command.args(&sppark_objects); for architecture in &architectures { diff --git a/cuda/README.md b/cuda/README.md index 8d203f1..48d4dc0 100644 --- a/cuda/README.md +++ b/cuda/README.md @@ -1,9 +1,9 @@ # CUDA backend -This directory contains multi-stark's first-party CUDA implementation of -Goldilocks arithmetic, DFTs, coset low-degree extensions, BLAKE3 Merkle -commitments, lookup construction, quotient evaluation, and FRI proving. It -does not use ICICLE or copy code from it. +This directory contains multi-stark's CUDA prover: Goldilocks arithmetic, +BLAKE3 Merkle commitments, lookup construction, quotient evaluation, and FRI +proving. All GPU DFTs and coset low-degree extensions use the pinned sppark +fork through `sppark_ntt.cu`. The `cuda` Cargo feature selects `CudaDft` for the production `GoldilocksBlake3Config`. BabyBear tests remain on their CPU DFT. Without the @@ -15,9 +15,10 @@ independent of CUDA. The backend preserves Plonky3's public PCS interfaces while keeping the hot prover pipeline device-resident: -1. Rust validates dimensions and builds exact P3-compatible twiddle tables. -2. Trace matrices are uploaded and transformed with fused radix-4/radix-8 - DIF kernels into resident coset LDEs. +1. Rust validates dimensions and caches immutable transform plans, including + panel scratch, batch groups and coset powers. Admission uses the same plan. +2. Trace matrices are uploaded and transformed through sppark on the caller's + stream into resident coset LDEs. 3. First-party BLAKE3 kernels commit mixed-height matrices without copying LDEs back to the host. 4. Lookup traces and quotient LDEs are constructed from resident commitments. @@ -97,8 +98,8 @@ Ix measurements are documented in [`docs/cuda-benchmarks.md`](../docs/cuda-bench - CUDA outputs are canonical field representatives. - Rust checks power-of-two sizes, two-adicity, integer overflow, and buffer lengths before each synchronous FFI call. -- The CUDA source is licensed under the repository's MIT/Apache-2.0 terms and - depends only on the CUDA runtime/toolkit when enabled. +- First-party CUDA sources use the repository's MIT/Apache-2.0 terms. The + sppark dependency retains its own license; see the pinned fork's `LICENSE`. - CUDA affects prover execution only. Fields, BLAKE3 hashing, transcripts, proof format, and the CPU verifier are unchanged. Canonical representation changed newly generated proof bytes relative to pre-CUDA revisions as diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 6ca288d..62ed57c 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -18,12 +18,7 @@ #include #include #include "metrics.cuh" -#ifdef MULTI_STARK_SPPARK -extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width); -extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, - size_t height, size_t width, size_t added_bits, - const uint64_t* shift_powers); -#endif +#include "ntt.cuh" @@ -1568,270 +1563,6 @@ __global__ void gather_mixed_lde_rows(uint64_t* output, } } -__global__ void radix2_dif_stage(uint64_t* values, size_t height, size_t width, - size_t half, const uint64_t* twiddles) { - const size_t total = (height >> 1) * width; - const size_t stride = height / (2 * half); - const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; - for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - index < total; index += grid_stride) { - const size_t butterfly = index / width; - const size_t column = index - butterfly * width; - const size_t offset = butterfly % half; - const size_t group = butterfly / half; - const size_t row_0 = group * (2 * half) + offset; - const size_t row_1 = row_0 + half; - const size_t index_0 = row_0 * width + column; - const size_t index_1 = row_1 * width + column; - const uint64_t left = values[index_0]; - const uint64_t right = values[index_1]; - values[index_0] = goldilocks_add(left, right); - values[index_1] = goldilocks_mul( - goldilocks_sub(left, right), twiddles[offset * stride]); - } -} - -// Fuse two consecutive radix-2 DIF stages. Each thread owns one column of -// one four-row butterfly, so row-major accesses remain coalesced while the -// intermediate values never return to global memory. -__global__ void radix4_dif_stage(uint64_t* values,size_t height,size_t width, - size_t half,const uint64_t* twiddles){ - const size_t quarter=half>>1,total=(height>>2)*width,stride=height/(2*half); - const size_t grid_stride=static_cast(blockDim.x)*gridDim.x; - for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index>2,total=(height>>3)*width,stride=height/(2*half); - const size_t grid_stride=static_cast(blockDim.x)*gridDim.x; - for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index>= 1) { - const size_t butterflies = start_half * width; - const size_t stride = height / (2 * half); - for (size_t index = threadIdx.x; index < butterflies; - index += blockDim.x) { - const size_t butterfly = index / width; - const size_t column = index - butterfly * width; - const size_t offset = butterfly % half; - const size_t subgroup = butterfly / half; - const size_t row_0 = subgroup * (2 * half) + offset; - const size_t row_1 = row_0 + half; - const size_t index_0 = row_0 * width + column; - const size_t index_1 = row_1 * width + column; - const uint64_t left = local_values[index_0]; - const uint64_t right = local_values[index_1]; - local_values[index_0] = goldilocks_add(left, right); - local_values[index_1] = goldilocks_mul( - goldilocks_sub(left, right), twiddles[offset * stride]); - } - __syncthreads(); - if (half == 1) { - break; - } - } - - for (size_t index = threadIdx.x; index < elements_per_group; - index += blockDim.x) { - values[global_base + index] = local_values[index]; - } - __syncthreads(); - } -} - -// Wide row-major matrices cannot fit all columns of a row group in shared -// memory. Tile the columns as well, preserving coalesced row-major loads while -// fusing the final DIF stages independently for each column tile. -__global__ void radix2_dif_tail_tiled(uint64_t* values,size_t height,size_t width, - size_t start_half,const uint64_t* twiddles,size_t columns_per_tile){ - extern __shared__ uint64_t local_values[];const size_t rows=2*start_half; - const size_t groups=height/rows,column_tiles=(width+columns_per_tile-1)/columns_per_tile; - const size_t tile_count=groups*column_tiles; - for(size_t tile_index=blockIdx.x;tile_index>=1){const size_t butterflies=start_half*columns,stride=height/(2*half); - for(size_t index=threadIdx.x;index 1 && width > 0) { - const unsigned log = strict_log2(height); - const unsigned bucket = width == 1 ? 0 : width == 2 ? 1 : width < 8 ? 2 : 3; - if (log <= 32) multi_stark_metrics::add(device_id, - multi_stark_metrics::NTT_OFFSET + log * 4 + bucket, 1); - } - if (height <= 1 || width == 0) { - return cudaSuccess; - } - const size_t total = (height >> 1) * width; - const unsigned int blocks = blocks_for(total); - // Wide row-major batches use fused stages regardless of height: short, - // wide traces can contain as many cells as tall traces. Width-2 FRI - // codewords retain the shared-memory tail specialized below. - if (width >= 8) { - size_t half = height >> 1; - while (half >= 4) { - radix8_dif_stage<<> 3) * width), THREADS>>>( - values, height, width, half, twiddles); - const cudaError_t status = cudaGetLastError(); - if (status != cudaSuccess) return status; - half >>= 3; - } - if (half == 2) { - radix4_dif_stage<<> 2) * width), THREADS>>>( - values, height, width, half, twiddles); - return cudaGetLastError(); - } - if (half == 1) { - radix2_dif_stage<<>>(values, height, width, 1, - twiddles); - } - return cudaGetLastError(); - } - if (width > 2) { - for (size_t half = height >> 1;; half >>= 1) { - radix2_dif_stage<<>>(values, height, width, half, - twiddles); - const cudaError_t status = cudaGetLastError(); - if (status != cudaSuccess) return status; - if (half == 1) return cudaSuccess; - } - } - constexpr size_t TAIL_HALF = 128; - for (size_t half = height >> 1; half > TAIL_HALF; half >>= 1) { - radix2_dif_stage<<>>(values, height, width, half, - twiddles); - const cudaError_t status = cudaGetLastError(); - if (status != cudaSuccess) { - return status; - } - } - - const size_t start_half = - height < 2 * TAIL_HALF ? height >> 1 : TAIL_HALF; - const size_t groups = height / (2 * start_half); - const unsigned int tail_blocks = static_cast( - groups < MAX_BLOCKS ? groups : MAX_BLOCKS); - const size_t shared_bytes = 2 * start_half * width * sizeof(uint64_t); - radix2_dif_tail<<>>( - values, height, width, start_half, twiddles); - return cudaGetLastError(); -} - -__global__ void bit_reverse_scale_and_shift(uint64_t* values, size_t height, - size_t width, unsigned int log_height, - uint64_t height_inverse, - const uint64_t* shift_powers) { - const size_t total = height * width; - const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; - for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - index < total; index += grid_stride) { - const size_t row = index / width; - const size_t column = index - row * width; - const size_t reverse_row = reverse_index_bits(row, log_height); - if (row > reverse_row) { - continue; - } - - const size_t reverse_index = reverse_row * width + column; - const uint64_t left = values[index]; - if (row == reverse_row) { - values[index] = goldilocks_mul( - left, goldilocks_mul(height_inverse, shift_powers[row])); - continue; - } - - const uint64_t right = values[reverse_index]; - values[index] = goldilocks_mul( - right, goldilocks_mul(height_inverse, shift_powers[row])); - values[reverse_index] = goldilocks_mul( - left, goldilocks_mul(height_inverse, shift_powers[reverse_row])); - } -} - __global__ void goldilocks_ops_kernel(uint64_t* sums, uint64_t* differences, uint64_t* products, uint64_t* inverses, const uint64_t* left, const uint64_t* right, @@ -2253,88 +1984,41 @@ cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, } // namespace -// A first-party twiddle table of `count` entries for a transform, cached -// on the device when `legacy`. The Rust side builds no table for a -// transform it sends through sppark, and that absence is the dispatch -// decision on this side; a table required by a first-party transform but -// missing is an argument error rather than a read past an empty buffer. -static cudaError_t legacy_table(int device_id, const uint64_t* host, size_t count, int kind, bool legacy, - const uint64_t** out) { - *out = nullptr; - if (!legacy) return cudaSuccess; - if (!host) return cudaErrorInvalidValue; - return cached_device_constants(device_id, host, count, kind, 0, 0, out); -} - -// The inverse-shift-forward sequence in place on `values`, whose first -// `height` rows hold natural-order evaluations and whose tail is zero: the -// sppark panel above its height threshold, the first-party stages -// otherwise. Both leave bit-reversed rows; `canonical_pass` adds the -// reduction pass the lookup paths run after the first-party stages. -static cudaError_t coset_lde_in_place(int device_id, uint64_t* values, size_t height, size_t width, - size_t extended_height, const uint64_t* inverse_twiddles, - const uint64_t* shift_powers, const uint64_t* forward_twiddles, - uint64_t height_inverse, bool canonical_pass) { - // The Rust side decides the backend once, by the shape, and hands the - // first-party tables over only for the first-party path; the tables' - // absence is the sppark decision, so a backend switch between the two - // sides cannot split a transform. -#ifdef MULTI_STARK_SPPARK - if (forward_twiddles == nullptr) { - if (multi_stark_metrics::enabled()) { - multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); - multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); - multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(extended_height, width), 1); - } - const size_t added_bits = strict_log2(extended_height) - strict_log2(height); - return static_cast(multi_stark_sppark_coset_lde( - device_id, values, values, height, width, added_bits, shift_powers)); - } - if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); -#else - if (forward_twiddles == nullptr) return cudaErrorInvalidValue; -#endif - cudaError_t status = launch_dif(device_id, values, height, width, inverse_twiddles); - if (status == cudaSuccess) { - bit_reverse_scale_and_shift<<>>( - values, height, width, strict_log2(height), height_inverse, shift_powers); - status = cudaGetLastError(); - } - if (status == cudaSuccess) status = launch_dif(device_id, values, extended_height, width, forward_twiddles); - if (status == cudaSuccess && canonical_pass) { - canonicalize_goldilocks<<>>(values, extended_height * width); - status = cudaGetLastError(); - } - return status; +static cudaError_t plan_constants(int device, const MultiStarkNttPlan* plan, + const uint64_t** powers) { + const auto* host = plan->shift_powers; + return cached_device_constants(device, host, plan->height, 3, host[0], + plan->height > 1 ? host[1] : 0, powers); } -// A forward transform in place on natural-order rows, leaving bit-reversed -// rows: sppark above its height threshold, the first-party stages otherwise. -static cudaError_t forward_in_place(int device_id, uint64_t* values, size_t height, size_t width, - const uint64_t* twiddles) { -#ifdef MULTI_STARK_SPPARK - if (twiddles == nullptr) { - if (multi_stark_metrics::enabled()) { - multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkTaken, 1); - multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); - } - return static_cast(multi_stark_sppark_forward(device_id, values, height, width)); - } - if (multi_stark_metrics::enabled()) multi_stark_metrics::add(device_id, multi_stark_metrics::SpparkDeclined, 1); -#else - if (twiddles == nullptr) return cudaErrorInvalidValue; -#endif - return launch_dif(device_id, values, height, width, twiddles); +static void record_transform(int device, size_t height, size_t width) { + if (multi_stark_metrics::enabled() && height > 1) + multi_stark_metrics::add(device, multi_stark_metrics::NTT_OFFSET + + multi_stark_metrics::ntt_shape(height, width), 1); +} + +static cudaError_t coset_lde(int device, const uint64_t* trace, uint64_t* values, + const MultiStarkNttPlan* plan) { + const uint64_t* powers = nullptr; + cudaError_t status = plan_constants(device, plan, &powers); + if (status != cudaSuccess) return status; + record_transform(device, plan->height, plan->width); + record_transform(device, plan->extended_height, plan->width); + return static_cast(multi_stark_sppark_coset_lde(device, trace, values, plan, powers)); +} + +static cudaError_t forward_in_place(int device, uint64_t* values, const MultiStarkNttPlan* plan) { + record_transform(device, plan->height, plan->width); + return static_cast(multi_stark_sppark_forward(device, values, plan)); } extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, size_t height, size_t width, - const uint64_t* twiddles) { + const MultiStarkNttPlan* plan) { if (values == nullptr || !is_power_of_two(height) || width == 0 || !product_fits(height, width)) { return static_cast(cudaErrorInvalidValue); } - const bool legacy = twiddles != nullptr; cudaError_t status = cudaSetDevice(device_id); if (status != cudaSuccess) { return static_cast(status); @@ -2343,13 +2027,9 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, const size_t elements = height * width; HostRegistration registered_values(values, elements * sizeof(uint64_t)); DeviceBuffer device_values; - DeviceBuffer device_twiddles; status = copy_to_device(device_values, values, elements); - if (status == cudaSuccess && legacy) { - status = copy_to_device(device_twiddles, twiddles, height / 2); - } if (status == cudaSuccess) { - status = forward_in_place(device_id, device_values.get(), height, width, device_twiddles.get()); + status = forward_in_place(device_id, device_values.get(), plan); } if (status == cudaSuccess) { status = copy_to_host(values, device_values, elements); @@ -2359,10 +2039,8 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, extern "C" int multi_stark_cuda_coset_lde_batch( int device_id, uint64_t* output, const uint64_t* input, size_t height, - size_t width, size_t added_bits, const uint64_t* inverse_twiddles, - const uint64_t* shift_powers, const uint64_t* forward_twiddles, - uint64_t height_inverse) { - if (output == nullptr || input == nullptr || shift_powers == nullptr || + size_t width, size_t added_bits, const MultiStarkNttPlan* plan) { + if (output == nullptr || input == nullptr || plan == nullptr || plan->shift_powers == nullptr || !is_power_of_two(height) || width == 0 || added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { return static_cast(cudaErrorInvalidValue); @@ -2371,7 +2049,6 @@ extern "C" int multi_stark_cuda_coset_lde_batch( if (!product_fits(extended_height, width)) { return static_cast(cudaErrorInvalidValue); } - const bool legacy = forward_twiddles != nullptr; cudaError_t status = cudaSetDevice(device_id); if (status != cudaSuccess) { @@ -2385,32 +2062,14 @@ extern "C" int multi_stark_cuda_coset_lde_batch( HostRegistration registered_output(output, output_elements * sizeof(uint64_t)); DeviceBuffer device_values; - DeviceBuffer device_inverse_twiddles; - DeviceBuffer device_shift_powers; - DeviceBuffer device_forward_twiddles; status = device_values.allocate(output_elements); - if (status == cudaSuccess) { - status = cudaMemset(device_values.get(), 0, output_elements * sizeof(uint64_t)); - } if (status == cudaSuccess) { status = cudaMemcpy(device_values.get(), input, input_elements * sizeof(uint64_t), cudaMemcpyHostToDevice); } - if (status == cudaSuccess && legacy) { - status = copy_to_device(device_inverse_twiddles, inverse_twiddles, height / 2); - } - if (status == cudaSuccess) { - status = copy_to_device(device_shift_powers, shift_powers, height); - } - if (status == cudaSuccess && legacy) { - status = copy_to_device(device_forward_twiddles, forward_twiddles, - extended_height / 2); - } if (status == cudaSuccess) { - status = coset_lde_in_place(device_id, device_values.get(), height, width, extended_height, - device_inverse_twiddles.get(), device_shift_powers.get(), - device_forward_twiddles.get(), height_inverse, false); + status = coset_lde(device_id, device_values.get(), device_values.get(), plan); } if (status == cudaSuccess) { status = copy_to_host(output, device_values, output_elements); @@ -2420,11 +2079,9 @@ extern "C" int multi_stark_cuda_coset_lde_batch( static int coset_lde_create( int device_id, void** handle, const uint64_t* input, size_t height, - size_t width, size_t added_bits, const uint64_t* inverse_twiddles, - const uint64_t* shift_powers, const uint64_t* forward_twiddles, - uint64_t height_inverse, void* context, TraceWriter writer) { + size_t width, size_t added_bits, const MultiStarkNttPlan* plan, void* context, TraceWriter writer) { if (handle == nullptr || (input == nullptr && writer == nullptr) || - shift_powers == nullptr || + plan == nullptr || plan->shift_powers == nullptr || !is_power_of_two(height) || width == 0 || added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { return static_cast(cudaErrorInvalidValue); @@ -2457,34 +2114,6 @@ static int coset_lde_create( } if (status == cudaSuccess) lde->trace_height = height; -#ifdef MULTI_STARK_SPPARK - // The sppark path writes every output row itself from its own scratch. - const bool sppark = forward_twiddles == nullptr; - if (multi_stark_metrics::enabled()) { - multi_stark_metrics::add(device_id, sppark ? multi_stark_metrics::SpparkTaken : multi_stark_metrics::SpparkDeclined, 1); - if (sppark) { - multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(height, width), 1); - multi_stark_metrics::add(device_id, multi_stark_metrics::NTT_SPPARK_OFFSET + multi_stark_metrics::ntt_shape(extended_height, width), 1); - } - } -#else - const bool sppark = false; - if (forward_twiddles == nullptr) { - destroy_resident_lde(lde); - return static_cast(cudaErrorInvalidValue); - } -#endif - // Large pageable uploads otherwise serialize through the driver's hidden - // staging pool; they go through the persistent staging slots instead. - // Small ones take the direct pageable path. - const uint64_t *device_inverse_twiddles=nullptr,*device_shift_powers=nullptr,*device_forward_twiddles=nullptr; - // The source prefix is overwritten by the trace copy. Only the padded - // tail needs zeroing before the forward transform. - if (status == cudaSuccess && !sppark && output_elements > input_elements) { - status = cudaMemsetAsync(lde->values + input_elements, 0, - (output_elements - input_elements) * sizeof(uint64_t), - cudaStreamPerThread); - } if (status == cudaSuccess) { if (writer) { constexpr size_t TILE_ROWS = size_t(1) << 16; @@ -2502,40 +2131,8 @@ static int coset_lde_create( cudaStreamPerThread); } } - if (status == cudaSuccess) { - status = cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&device_shift_powers); - } -#ifdef MULTI_STARK_SPPARK - if (status == cudaSuccess && sppark) { - status = static_cast(multi_stark_sppark_coset_lde( - device_id, lde->trace_values, lde->values, height, width, added_bits, device_shift_powers)); - } -#endif - if (status == cudaSuccess && !sppark) { - status = cudaMemcpyAsync(lde->values, lde->trace_values, - input_elements * sizeof(uint64_t), - cudaMemcpyDeviceToDevice, - cudaStreamPerThread); - } if (status == cudaSuccess) - status = legacy_table(device_id, inverse_twiddles, height / 2, 1, !sppark && height > 1, &device_inverse_twiddles); - if (status == cudaSuccess) - status = legacy_table(device_id, forward_twiddles, extended_height / 2, 2, !sppark, &device_forward_twiddles); - if (status == cudaSuccess && !sppark) { - status = launch_dif(device_id, lde->values, height, width,device_inverse_twiddles); - } - if (status == cudaSuccess && !sppark) { - bit_reverse_scale_and_shift<<>>( - lde->values, height, width, strict_log2(height), height_inverse, - device_shift_powers); - status = cudaGetLastError(); - } - if (status == cudaSuccess && !sppark) { - status = launch_dif(device_id, lde->values, extended_height, width,device_forward_twiddles); - } - // Normalization and every DIF butterfly produce canonical field values, - // including the height-one case. A final reduction pass is redundant; - // raw-representation tests protect the Merkle byte contract. + status = coset_lde(device_id, lde->trace_values, lde->values, plan); if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status != cudaSuccess) { destroy_resident_lde(lde); @@ -2547,23 +2144,20 @@ static int coset_lde_create( extern "C" int multi_stark_cuda_coset_lde_create( int device_id, void** handle, const uint64_t* input, size_t height, - size_t width, size_t added_bits, const uint64_t* inverse_twiddles, - const uint64_t* shift_powers, const uint64_t* forward_twiddles, - uint64_t height_inverse) { + size_t width, size_t added_bits, const MultiStarkNttPlan* plan) { return coset_lde_create(device_id, handle, input, height, width, added_bits, - inverse_twiddles, shift_powers, forward_twiddles, height_inverse, nullptr, nullptr); + plan, nullptr, nullptr); } // Context ownership transfers only on success. Its lifetime covers trace // release and LDE eviction because lookup recovery still needs original rows. extern "C" int multi_stark_cuda_coset_lde_generate( int device_id, void** handle, size_t height, size_t width, size_t added_bits, - const uint64_t* inverse_twiddles, const uint64_t* shift_powers, - const uint64_t* forward_twiddles, uint64_t height_inverse, + const MultiStarkNttPlan* plan, void* context, TraceWriter writer, TraceDestroy destroy) { if (!context || !writer || !destroy) return static_cast(cudaErrorInvalidValue); const int status = coset_lde_create(device_id, handle, nullptr, height, width, added_bits, - inverse_twiddles, shift_powers, forward_twiddles, height_inverse, context, writer); + plan, context, writer); if (status == 0) static_cast(*handle)->trace_destroy = destroy; return status; } @@ -2578,19 +2172,10 @@ extern "C" void* multi_stark_cuda_lde_generator_context(const void* handle) { return lde->trace_writer ? lde->trace_context : nullptr; } -extern "C" int multi_stark_cuda_prepare_lde_constants( - int device_id,const uint64_t* inverse_twiddles,size_t inverse_count, - const uint64_t* shift_powers,size_t height,const uint64_t* forward_twiddles, - size_t forward_count) { - if((inverse_count&&!inverse_twiddles)||!shift_powers||!height|| - (forward_count&&!forward_twiddles))return static_cast(cudaErrorInvalidValue); - cudaError_t status=cudaSetDevice(device_id);const uint64_t* ignored=nullptr; - if(status==cudaSuccess&&inverse_count) - status=cached_device_constants(device_id,inverse_twiddles,inverse_count,1,0,0,&ignored); - if(status==cudaSuccess) - status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&ignored); - if(status==cudaSuccess&&forward_count) - status=cached_device_constants(device_id,forward_twiddles,forward_count,2,0,0,&ignored); +extern "C" int multi_stark_cuda_prepare_lde_constants(int device_id, const MultiStarkNttPlan* plan) { + cudaError_t status = cudaSetDevice(device_id); + const uint64_t* ignored = nullptr; + if (status == cudaSuccess) status = plan_constants(device_id, plan, &ignored); return static_cast(status); } @@ -2927,8 +2512,8 @@ extern "C" int multi_stark_cuda_quotient_lde( const uint64_t* alpha, size_t constraint_count, const uint64_t* delta, uint64_t ext_w, size_t quotient_size, size_t next_step, size_t quotient_degree, - size_t log_blowup, const uint64_t* quotient_twiddles, - const uint64_t* lde_twiddles, const uint64_t* slice_weights) { + size_t log_blowup, const MultiStarkNttPlan* quotient_plan, + const MultiStarkNttPlan* lde_plan, const uint64_t* slice_weights) { if (output_handle == nullptr || nodes == nullptr || roots == nullptr || main_handle == nullptr || stage2_handle == nullptr || publics == nullptr || alpha == nullptr || delta == nullptr || @@ -2975,9 +2560,7 @@ extern "C" int multi_stark_cuda_quotient_lde( if(status==cudaSuccess)status=generate_coset_selectors(ds,quotient_size,next_step, coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); - const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess)status=legacy_table(device_id,quotient_twiddles,quotient_size/2,2,quotient_twiddles!=nullptr,&device_quotient_twiddles); - if(status==cudaSuccess)status=legacy_table(device_id,lde_twiddles,lde_height/2,2,lde_twiddles!=nullptr,&device_lde_twiddles); + const uint64_t* device_weights=nullptr; if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); @@ -2999,7 +2582,7 @@ extern "C" int multi_stark_cuda_quotient_lde( dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_plan); ResidentLde* lde = nullptr; if(status==cudaSuccess) { @@ -3016,7 +2599,7 @@ extern "C" int multi_stark_cuda_quotient_lde( quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess) { *output_handle=lde; @@ -3049,8 +2632,8 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( uint64_t vanishing_start, uint64_t vanishing_step, const uint64_t* alpha, size_t constraint_count, const uint64_t* delta, uint64_t ext_w, size_t quotient_size, size_t next_step, size_t quotient_degree, - size_t log_blowup, const uint64_t* quotient_twiddles, - const uint64_t* lde_twiddles, const uint64_t* slice_weights) { + size_t log_blowup, const MultiStarkNttPlan* quotient_plan, + const MultiStarkNttPlan* lde_plan, const uint64_t* slice_weights) { const auto one_source = [](const void* handle, const uint64_t* host) { return (handle != nullptr) != (host != nullptr); }; @@ -3103,9 +2686,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( if(status==cudaSuccess)status=generate_coset_selectors(ds,quotient_size,next_step, coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); - const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; - if(status==cudaSuccess)status=legacy_table(device_id,quotient_twiddles,quotient_size/2,2,quotient_twiddles!=nullptr,&device_quotient_twiddles); - if(status==cudaSuccess)status=legacy_table(device_id,lde_twiddles,lde_height/2,2,lde_twiddles!=nullptr,&device_lde_twiddles); + const uint64_t* device_weights=nullptr; if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); const auto* prep_resident=static_cast(preprocessed_handle); @@ -3227,7 +2808,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( } for(size_t i=0;i<2;++i)if(status==cudaSuccess&&stream_busy[i])status=cudaStreamSynchronize(streams[i]); - if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,device_quotient_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_plan); ResidentLde* lde=nullptr; if(status==cudaSuccess)status=create_resident_lde(&lde); if(status==cudaSuccess){lde->height=lde_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t));} @@ -3237,7 +2818,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( lde->values,quotient,device_weights,quotient_size,trace_height,quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else if(lde)destroy_resident_lde(lde); for(size_t i=0;i<2;++i){ @@ -3448,14 +3029,13 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han const void* nodes,size_t node_count,size_t slot_count,const void* lookups,size_t lookup_count, const uint32_t* lookup_args,size_t lookup_arg_count,const void* preprocessed_handle, const void* main_handle,size_t group_size,const uint64_t* beta,const uint64_t* gamma, - uint64_t ext_w,size_t added_bits,const uint64_t* inverse_twiddles, - const uint64_t* shift_powers,const uint64_t* forward_twiddles,uint64_t height_inverse){ + uint64_t ext_w,size_t added_bits,const MultiStarkNttPlan* plan){ auto* main=const_cast(static_cast(main_handle)); auto* prep=static_cast(preprocessed_handle); if(!output_handle||!total||!nodes||!node_count||!slot_count||!lookups||!lookup_count|| (lookup_arg_count&&!lookup_args)||!main|| (!main->trace_values&&!main->host_trace_values&&!main->trace_writer)||!main->trace_height|| - !group_size||!beta||!gamma||!shift_powers|| + !group_size||!beta||!gamma||(!plan || !plan->shift_powers)|| !is_power_of_two(main->trace_height)||added_bits>=sizeof(size_t)*8|| main->trace_height>(SIZE_MAX>>added_bits)||(prep&&!prep->trace_values)) { return static_cast(cudaErrorInvalidValue); @@ -3523,12 +3103,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=exclusive_scan_ext2(reinterpret_cast(lde->values),deltas,count); if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); - const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - const bool legacy_tables=forward_twiddles!=nullptr; - if(status==cudaSuccess)status=legacy_table(device_id,inverse_twiddles,height/2,1,legacy_tables,&dit); - if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess)status=legacy_table(device_id,forward_twiddles,extended_height/2,2,legacy_tables,&dft); - if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); + if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); cudaFree(trace_chunk);cudaFree(scratch);cudaFree(deltas);cudaFree(multiplicities);cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(metadata); @@ -3539,10 +3114,9 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui const uint64_t* multiplicities,const uint64_t* args,const size_t* arg_offsets, size_t height,size_t num_lookups,size_t args_width,size_t group_size, const uint64_t* beta,const uint64_t* gamma,uint64_t ext_w,size_t added_bits, - const uint64_t* inverse_twiddles,const uint64_t* shift_powers, - const uint64_t* forward_twiddles,uint64_t height_inverse){ + const MultiStarkNttPlan* plan){ if(!output_handle||!total||!multiplicities||!arg_offsets||!height||!num_lookups|| - !group_size||!beta||!gamma||!shift_powers|| + !group_size||!beta||!gamma||(!plan || !plan->shift_powers)|| (args_width&&!args)||!is_power_of_two(height)|| added_bits>=sizeof(size_t)*8||height>(SIZE_MAX>>added_bits)) return static_cast(cudaErrorInvalidValue); @@ -3587,12 +3161,7 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); scan_done=now(); - const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; - const bool legacy_tables=forward_twiddles!=nullptr; - if(status==cudaSuccess)status=legacy_table(device_id,inverse_twiddles,height/2,1,legacy_tables,&dit); - if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); - if(status==cudaSuccess)status=legacy_table(device_id,forward_twiddles,extended_height/2,2,legacy_tables,&dft); - if(status==cudaSuccess)status=coset_lde_in_place(device_id,lde->values,height,width,extended_height,dit,dshift,dft,height_inverse,true); + if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, "[multi-stark/cuda] lookup phases: height=%zu lookups=%zu slots=%zu args_width=%zu allocate=%.3fs rows=%.3fs scan=%.3fs dft=%.3fs\n", @@ -3744,10 +3313,9 @@ extern "C" int multi_stark_cuda_lookup_lde_cpu_rows_partitioned( extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( int device_id, void* pending_handle, void** output_handle, uint64_t* total, - const uint64_t* inverse_twiddles, const uint64_t* shift_powers, - const uint64_t* forward_twiddles, uint64_t height_inverse) { + const MultiStarkNttPlan* plan) { auto* pending = static_cast(pending_handle); - if (!pending || !output_handle || !total || !shift_powers) { + if (!pending || !output_handle || !total || (!plan || !plan->shift_powers)) { return static_cast(cudaErrorInvalidValue); } *output_handle = nullptr; @@ -3763,24 +3331,8 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( status = cudaMemcpy(total + 2, pending->deltas + count - 1, sizeof(Ext2), cudaMemcpyDeviceToHost); - const uint64_t* inverse = nullptr; - const uint64_t* shifts = nullptr; - const uint64_t* forward = nullptr; - const size_t width = 2 * pending->slots; - const bool legacy_tables = forward_twiddles != nullptr; - if (status == cudaSuccess) - status = legacy_table(device_id, inverse_twiddles, pending->height / 2, 1, legacy_tables, &inverse); - if (status == cudaSuccess) - status = cached_device_constants(device_id, shift_powers, - pending->height, 3, shift_powers[0], - pending->height > 1 ? shift_powers[1] : 0, - &shifts); - if (status == cudaSuccess) - status = legacy_table(device_id, forward_twiddles, pending->extended_height / 2, 2, legacy_tables, &forward); if (status == cudaSuccess) - status = coset_lde_in_place(device_id, pending->lde->values, pending->height, width, - pending->extended_height, inverse, shifts, forward, - height_inverse, true); + status = coset_lde(device_id, pending->lde->values, pending->lde->values, plan); if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status == cudaSuccess) { *output_handle = pending->lde; diff --git a/cuda/metrics.cuh b/cuda/metrics.cuh index 87e0918..bf6f8c1 100644 --- a/cuda/metrics.cuh +++ b/cuda/metrics.cuh @@ -10,15 +10,13 @@ namespace multi_stark_metrics { constexpr size_t DEVICES = 64; constexpr size_t NTT_OFFSET = 16; -// Transform shape counts, one block per backend: first-party, then sppark. +// Transform shape counts for the CUDA NTT. constexpr size_t NTT_SHAPES = 33 * 4; -constexpr size_t NTT_SPPARK_OFFSET = NTT_OFFSET + NTT_SHAPES; -constexpr size_t WORDS = NTT_SPPARK_OFFSET + NTT_SHAPES; +constexpr size_t WORDS = NTT_OFFSET + NTT_SHAPES; enum Counter : size_t { UploadCalls, UploadRequestedBytes, UploadChunks, UploadFailures, UploadHostNs, CosetHits, CosetMisses, CosetUploadedBytes, ConstantBytes, - DriverFreeBytes, TotalBytes, MemorySamples, - SpparkTaken, SpparkDeclined + DriverFreeBytes, TotalBytes, MemorySamples }; inline size_t ntt_shape(size_t height, size_t width) { unsigned log = 0; diff --git a/cuda/ntt.cuh b/cuda/ntt.cuh new file mode 100644 index 0000000..c23654e --- /dev/null +++ b/cuda/ntt.cuh @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +#pragma once +#include +#include + +// Borrowed from the Rust plan; its host constants outlive the synchronous +// prover entry that uploads them. All sizes and groups are fixed at admission. +struct MultiStarkNttPlan { + size_t height; + size_t width; + size_t extended_height; + size_t columns; + size_t inverse_group; + size_t forward_group; + size_t scratch_bytes; + const uint64_t* shift_powers; +}; + +extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, + const MultiStarkNttPlan* plan); +extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, + uint64_t* values, const MultiStarkNttPlan* plan, + const uint64_t* device_shift_powers); diff --git a/cuda/smoke.sh b/cuda/smoke.sh index 0ad061a..ff752b0 100755 --- a/cuda/smoke.sh +++ b/cuda/smoke.sh @@ -17,12 +17,12 @@ echo "MULTI_STARK_CUDA_ARCHS=${MULTI_STARK_CUDA_ARCHS}" cargo clippy --release --locked --all-targets --features parallel,cuda -- -D warnings cargo test --release --locked --features parallel,cuda \ - cuda::tests::goldilocks_field_kernels_match_cpu -- --test-threads=1 + cuda::tests::goldilocks_field_kernels_match_cpu cargo test --release --locked --features parallel,cuda \ - cuda::tests::batched_dft_matches_cpu -- --test-threads=1 + cuda::tests::batched_dft_matches_cpu cargo test --release --locked --features parallel,cuda \ - cuda::tests::coset_lde_matches_cpu_including_storage_layout -- --test-threads=1 -cargo test --release --locked --features parallel,cuda -- --test-threads=1 + cuda::tests::coset_lde_matches_cpu_including_storage_layout +cargo test --release --locked --features parallel,cuda compat_dir="$(mktemp -d)" trap 'rm -rf "$compat_dir"' EXIT diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 035599d..0ac8174 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -1,16 +1,7 @@ -// sppark's Goldilocks NTT behind a C interface for the comparison backend. -// -// Every transform runs on an upstream stream leased for the call and is -// fenced against the caller's per-thread stream with events, so a caller -// sees the same completion contract as the first-party kernels: work -// enqueued after the call on its own stream follows the transform, and the -// host is not blocked. Upstream's field type is a plain 64-bit word with -// canonical values, the same storage the prover's matrices use. The fork -// is built with SPPARK_NO_CXX_RUNTIME: a CUDA failure inside upstream ends -// the process with a message rather than throwing, which matches the -// status checks on the Rust side. +// Goldilocks transforms on the caller's stream, with immutable panel plans. #include #include +#include "ntt.cuh" #include #include @@ -21,21 +12,6 @@ namespace { -struct EventPair { - cudaEvent_t before = nullptr; - cudaEvent_t after = nullptr; - ~EventPair() { - if (before) cudaEventDestroy(before); - if (after) cudaEventDestroy(after); - } - cudaError_t create() { - cudaError_t status = cudaEventCreateWithFlags(&before, cudaEventDisableTiming); - if (status == cudaSuccess) - status = cudaEventCreateWithFlags(&after, cudaEventDisableTiming); - return status; - } -}; - bool valid_arguments(const void* d_inout, uint32_t lg, int order, int direction, int coset) { return d_inout && lg > 0 && lg <= MAX_LG_DOMAIN_SIZE && order >= 0 && order <= 3 && direction >= 0 && direction <= 1 && coset >= 0 && coset <= 1; @@ -50,9 +26,9 @@ extern "C" int multi_stark_sppark_max_lg_domain() { return MAX_LG_DOMAIN_SIZE; } // `order` is NTT::InputOutputOrder (NN, NR, RN, RR), `direction` 0 forward // or 1 inverse, `coset` 1 for the multiplicative coset by the field // generator. Inverse transforms are normalized by 1/2^lg upstream. -extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* d_inout, uint32_t lg, +static int ntt_batch_on_stream(int device, uint64_t* d_inout, uint32_t lg, int order, int direction, int coset, - uint32_t batch, size_t stride) { + uint32_t batch, size_t stride, cudaStream_t caller) { // Upstream indexes with index_t, 32 bits at this domain limit. if (!valid_arguments(d_inout, lg, order, direction, coset) || batch == 0 || batch > 65535 || (batch > 1 && stride < (size_t(1) << lg)) || stride > size_t(std::numeric_limits::max())) @@ -66,26 +42,41 @@ extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* d_inout if (candidate->cid() == device) found = candidate; if (!found) return static_cast(cudaErrorInvalidDevice); const gpu_t& gpu = select_gpu(found->id()); - EventPair events; - status = events.create(); - if (status != cudaSuccess) return static_cast(status); - status = cudaEventRecord(events.before, cudaStreamPerThread); - if (status != cudaSuccess) return static_cast(status); - // Upstream's CUDA failures end the process (the fork's runtime mode), so - // everything past the launch either completes or never returns. - stream_t stream(gpu.id()); - stream.wait(events.before); + stream_t stream(gpu.id(), caller); NTT::Base_dev_ptr_batch(stream, reinterpret_cast(d_inout), lg, static_cast(order), static_cast(direction), static_cast(coset), batch, stride); - stream.record(events.after); - status = cudaStreamWaitEvent(cudaStreamPerThread, events.after, 0); - // Only once the caller's stream waits on the transform may the private - // stream go; if that wait could not be installed, drain the stream here - // so nothing still runs against the caller's buffer on return. - if (status != cudaSuccess) stream.sync(); - return static_cast(status); + return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* data, uint32_t lg, + int order, int direction, int coset, + uint32_t batch, size_t stride) { + return ntt_batch_on_stream(device, data, lg, order, direction, coset, batch, stride, cudaStreamPerThread); +} + +// A fresh caller-owned stream carrying producer, transforms and consumer. +// Reusing it after each borrowed wrapper is dropped checks its ownership. +extern "C" int multi_stark_sppark_borrowed_round_trip(int device, uint64_t* values, uint32_t lg) { + cudaError_t status = cudaSetDevice(device); + if (status != cudaSuccess) return static_cast(status); + cudaStream_t caller = nullptr; + status = cudaStreamCreateWithFlags(&caller, cudaStreamNonBlocking); + if (status != cudaSuccess) return static_cast(status); + uint64_t* data = nullptr; + const size_t count = size_t(1) << lg, bytes = count * sizeof(uint64_t); + status = cudaMallocAsync(reinterpret_cast(&data), bytes, caller); + if (status == cudaSuccess) status = cudaMemcpyAsync(data, values, bytes, cudaMemcpyHostToDevice, caller); + int result = static_cast(status); + if (result == 0) result = ntt_batch_on_stream(device, data, lg, 1, 0, 0, 1, count, caller); + if (result == 0) result = ntt_batch_on_stream(device, data, lg, 2, 1, 0, 1, count, caller); + if (result == 0) result = static_cast(cudaMemcpyAsync(values, data, bytes, cudaMemcpyDeviceToHost, caller)); + if (data) cudaFreeAsync(data, caller); + const auto synced = cudaStreamSynchronize(caller); + if (result == 0) result = static_cast(synced); + cudaStreamDestroy(caller); + return result; } // One transform: the batch entry with a single vector. @@ -128,33 +119,10 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t return multi_stark_sppark_ntt_batch_host(device, inout, lg, order, direction, coset, 1, size_t(1) << lg); } -// --- Resident coset LDE through sppark --------------------------------- -// -// The prover's matrices are row-major with the transform along the column, -// while upstream transforms contiguous vectors. A panel of columns is -// gathered into column-major scratch, transformed, and scattered back in -// the bit-reversed row order the commitment expects: -// -// A[c][r] = canonical(trace[r][f + c]) r < N -// inverse NR on A[c], normalized upstream (bit-reversed coeffs) -// B[c][i << added_bits] = A[c][i] * shift^rev(i) i < N, zero elsewhere -// forward RN on B[c] over M = N << added_bits (natural-order evals) -// values[rev(r)][f + c] = B[c][r] r < M -// -// That feeds the forward transform in bit-reversed order, so no pass -// restores the coefficient order and the row permutation folds into the -// scatter. The default restores the order instead (B[c][i] = A[c][rev(i)] -// * shift^i), transforms in NR order and scatters naturally; -// `fused_expansion` decides and records why. Neither folds the expansion -// into the transform's first pass; that remains upstream's kernels' job. The gather reduces representatives at or -// above the modulus, which upstream does not accept. The compact panel A -// and the extended panel B cost 8 * (N + M) * C bytes for C columns; C is -// sized from MULTI_STARK_SPPARK_PANEL_BYTES (default 4 GiB) and the width. -// The columns of a panel go through batched launch sequences (the fork's -// grid dimension) in groups sized to the L2 cache; -// MULTI_STARK_SPPARK_BATCH_BYTES sets the group, 0 launches every column -// on its own. - +// Row-major matrices are gathered into contiguous column panels for the +// inverse NTT, coefficient-order restoration with coset shift and expansion, +// and forward NR transform. Scatter preserves the bit-reversed row storage +// consumed by commitments. Gather canonicalizes lazy field representatives. #include "goldilocks.cuh" namespace { @@ -203,12 +171,9 @@ __global__ void gather_tiles(const uint64_t* __restrict__ source, size_t height, } } -// The panel back into row-major rows of columns [first, first + count); -// with `reverse_rows`, panel row r lands in output row rev(r) over -// 2^log_rows rows, so a natural-order transform result is stored in the -// bit-reversed row order the commitment expects. +// The column panel back into row-major storage without a row permutation. __global__ void scatter_tiles(const uint64_t* __restrict__ panel, size_t rows, size_t column_stride, size_t width, - size_t first, size_t count, unsigned log_rows, bool reverse_rows, + size_t first, size_t count, uint64_t* __restrict__ values) { __shared__ uint64_t tile[TILE][TILE + 1]; const size_t row0 = static_cast(blockIdx.x) * TILE; @@ -221,8 +186,7 @@ __global__ void scatter_tiles(const uint64_t* __restrict__ panel, size_t rows, s for (unsigned r = threadIdx.y; r < TILE; r += TILE_ROWS) { const size_t row = row0 + r, col = col0 + threadIdx.x; if (row < rows && col < count) { - const size_t out = reverse_rows ? reverse_bits(row, log_rows) : row; - values[out * width + first + col] = tile[threadIdx.x][r]; + values[row * width + first + col] = tile[threadIdx.x][r]; } } } @@ -258,12 +222,11 @@ __global__ void gather_rows(const uint64_t* __restrict__ source, size_t height, } __global__ void scatter_rows(const uint64_t* __restrict__ panel, size_t rows, size_t column_stride, size_t width, - size_t first, unsigned count, unsigned log_rows, bool reverse_rows, + size_t first, unsigned count, uint64_t* __restrict__ values) { const size_t stride = static_cast(blockDim.x) * gridDim.x; for (size_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; row < rows; row += stride) { - const size_t out_row = reverse_rows ? reverse_bits(row, log_rows) : row; - uint64_t* out = values + out_row * width + first; + uint64_t* out = values + row * width + first; for (unsigned c = 0; c < count; ++c) out[c] = panel[c * column_stride + row]; } } @@ -281,46 +244,16 @@ cudaError_t gather_panel(const uint64_t* source, size_t height, size_t width, si } cudaError_t scatter_panel(const uint64_t* panel, size_t rows, size_t column_stride, size_t width, size_t first, - size_t count, unsigned log_rows, bool reverse_rows, uint64_t* values) { + size_t count, uint64_t* values) { if (width < NARROW_MATRIX) scatter_rows<<>>( - panel, rows, column_stride, width, first, static_cast(count), log_rows, reverse_rows, values); + panel, rows, column_stride, width, first, static_cast(count), values); else scatter_tiles<<>>( - panel, rows, column_stride, width, first, count, log_rows, reverse_rows, values); + panel, rows, column_stride, width, first, count, values); return cudaGetLastError(); } -// powers[rev(i)] for i < 2^log: the coset powers in the order the -// bit-reversed coefficients meet them. -__global__ void reverse_powers(const uint64_t* __restrict__ powers, size_t count, unsigned log, - uint64_t* __restrict__ reversed) { - const size_t stride = static_cast(blockDim.x) * gridDim.x; - for (size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < count; i += stride) - reversed[i] = powers[reverse_bits(i, log)]; -} - -// The expansion: bit-reversed coefficients A[c][i] (height per column) -// become the bit-reversed input of the forward transform over the extended -// height, B[c][i << added_bits] = A[c][i] * shift^rev(i), zero elsewhere, -// so a forward RN transform yields natural-order evaluations of the coset. -// One column per grid row, rows coalesced along the block. -__global__ void spread_columns(const uint64_t* __restrict__ coefficients, const uint64_t* __restrict__ reversed_powers, - size_t height, size_t extended_height, unsigned added_bits, - uint64_t* __restrict__ panel) { - const size_t column = blockIdx.y; - const size_t mask = (size_t(1) << added_bits) - 1; - const size_t stride = static_cast(blockDim.x) * gridDim.x; - for (size_t j = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; j < extended_height; j += stride) { - uint64_t value = 0; - if ((j & mask) == 0) { - const size_t i = j >> added_bits; - value = multi_stark_cuda::goldilocks_mul(coefficients[column * height + i], reversed_powers[i]); - } - panel[column * extended_height + j] = value; - } -} - unsigned log2_exact(size_t value) { unsigned log = 0; while ((size_t(1) << log) < value) ++log; @@ -342,73 +275,9 @@ unsigned long long decimal_setting(const char* name, unsigned long long fallback return value; } -// Read per construction: one getenv against a transform of gigabytes, and -// tests vary it within a process. -size_t panel_budget_bytes() { - const unsigned long long budget = decimal_setting("MULTI_STARK_SPPARK_PANEL_BYTES", 0); - return budget ? size_t(budget) : size_t(4) << 30; -} - -// -1 unread, 0 first-party, 1 sppark above the height threshold, 2 sppark -// for every height (tests compare the paths on small shapes). Read and -// written from concurrent constructions. -std::atomic backend_flag{-1}; - -unsigned min_log_height() { - return unsigned(decimal_setting("MULTI_STARK_SPPARK_MIN_LOG_HEIGHT", 18)); -} - -// The columns one panel holds at `column_bytes` each: as many as the -// budget admits, at most the width and the batch a launch grid can carry. -// Zero when one column does not fit, which the dispatch rules decline -// before any allocation. -size_t panel_columns(size_t width, size_t column_bytes, size_t budget) { - size_t columns = budget / column_bytes; - if (columns > 65535) columns = 65535; - return columns < width ? columns : width; -} - -// Whether a transform of 2^log rows is within upstream's compiled domain. -bool within_domain(size_t height, size_t added_bits) { - if (height == 0 || (height & (height - 1))) return false; - const unsigned log = log2_exact(height); - return added_bits <= MAX_LG_DOMAIN_SIZE && log + added_bits <= MAX_LG_DOMAIN_SIZE; -} - -// The bytes one batched launch sequence keeps in flight: the columns of a -// group go through every stage together, so a group that fits the L2 cache -// keeps the reuse between stages that one column at a time had, while -// short columns still share launches. MULTI_STARK_SPPARK_BATCH_BYTES -// overrides the device's L2 size; 0 launches every column on its own. -size_t l2_bytes(int device) { - int l2 = 0; - if (cudaDeviceGetAttribute(&l2, cudaDevAttrL2CacheSize, device) != cudaSuccess || l2 <= 0) - l2 = 64 << 20; - return static_cast(l2); -} - -size_t batch_group_bytes(int device) { - const unsigned long long setting = decimal_setting("MULTI_STARK_SPPARK_BATCH_BYTES", ~0ull); - return setting != ~0ull ? static_cast(setting) : l2_bytes(device); -} - -// Whether a panel's expansion feeds the forward transform in bit-reversed -// order with the row permutation in the scatter (MULTI_STARK_SPPARK_FUSED=1), -// or restores the coefficient order first and scatters naturally. Measured -// on the RTX PRO 6000, the bit-reversed feed saves the restoring pass but -// its reversing scatter costs as much on wide panels and more on tall -// narrow ones, so restoring is the default. -bool fused_expansion() { return decimal_setting("MULTI_STARK_SPPARK_FUSED", 0) == 1; } - -// The reversed coset powers the bit-reversed feed reads: one word per input -// row, nothing on the restoring path. -size_t expansion_extra_bytes(size_t height) { - return fused_expansion() ? height * sizeof(uint64_t) : 0; -} - // MULTI_STARK_SPPARK_STAGE_TIMING=1 prints the stage times of every coset // LDE to stderr: events on the caller's stream around each stage, which -// the transforms are fenced to, and one synchronization per LDE. +// the transforms are fenced to, and one synchronization per panel. struct StageTimer { static constexpr int STAGES = 5; bool enabled = decimal_setting("MULTI_STARK_SPPARK_STAGE_TIMING", 0) != 0; @@ -441,11 +310,8 @@ struct StageTimer { // The `count` columns of a panel, `stride` elements apart, through batched // launch sequences of as many columns as the group budget holds. int transform_columns(int device, uint64_t* panel, uint32_t lg, int order, int direction, size_t count, - size_t stride) { + size_t stride, size_t group) { if (lg == 0) return 0; - const size_t column_bytes = (size_t(1) << lg) * sizeof(uint64_t); - size_t group = batch_group_bytes(device) / column_bytes; - if (group == 0) group = 1; int result = 0; for (size_t first = 0; result == 0 && first < count; first += group) { const size_t batch = count - first < group ? count - first : group; @@ -457,75 +323,13 @@ int transform_columns(int device, uint64_t* panel, uint32_t lg, int order, int d } // namespace -// Whether the prover's transforms take the sppark path: MULTI_STARK_CUDA_NTT=sppark, -// or a runtime selection, which tests use to compare both paths in one process. -extern "C" int multi_stark_sppark_backend_selected() { - int flag = backend_flag.load(std::memory_order_acquire); - if (flag < 0) { - const char* configured = getenv("MULTI_STARK_CUDA_NTT"); - int expected = -1; - const int read = configured && strcmp(configured, "sppark") == 0; - // The first reader publishes; a concurrent selection wins over it. - flag = backend_flag.compare_exchange_strong(expected, read, std::memory_order_acq_rel) ? read : expected; - } - return flag; -} - -// 0 first-party, 1 sppark above the height threshold, 2 sppark always. -extern "C" void multi_stark_sppark_select_backend(int selected) { - backend_flag.store(selected, std::memory_order_release); -} - -// Whether a transform of `height` input rows is tall enough for the sppark -// path. Below MULTI_STARK_SPPARK_MIN_LOG_HEIGHT (18) the panel's gather and -// scatter passes cost more than the first-party kernels save, and the -// first-party kernels stay. -extern "C" int multi_stark_sppark_takes(size_t height) { - const int flag = multi_stark_sppark_backend_selected(); - if (flag == 2) return 1; - if (flag != 1) return 0; - return height >= (size_t(1) << min_log_height()); -} - -// Whether a resident coset LDE of the shape takes the sppark path: tall -// enough, the extended height within upstream's compiled domain, and one -// column's scratch, the compact and the extended panel, within the budget. -extern "C" int multi_stark_sppark_takes_lde(size_t height, size_t width, size_t added_bits) { - if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, added_bits)) return 0; - const size_t extended_height = height << added_bits; - return (height + extended_height) * sizeof(uint64_t) + expansion_extra_bytes(height) <= panel_budget_bytes(); -} - -// Whether a forward transform of the shape takes the sppark path. -extern "C" int multi_stark_sppark_takes_forward(size_t height, size_t width) { - if (width == 0 || !multi_stark_sppark_takes(height) || !within_domain(height, 0)) return 0; - return height * sizeof(uint64_t) <= panel_budget_bytes(); -} - -// The scratch the sppark path allocates for one LDE, within the panel -// budget: the compact and the extended panel of the columns the budget -// admits, plus the bit-reversed feed's reversed coset powers when that -// expansion is selected. Zero when the shape does not take the path. -extern "C" size_t multi_stark_sppark_panel_bytes(size_t height, size_t width, size_t added_bits) { - if (!multi_stark_sppark_takes_lde(height, width, added_bits)) return 0; - const size_t column_bytes = (height + (height << added_bits)) * sizeof(uint64_t); - const size_t extra = expansion_extra_bytes(height); - return panel_columns(width, column_bytes, panel_budget_bytes() - extra) * column_bytes + extra; -} - -// The scratch the sppark path allocates for one forward transform. -extern "C" size_t multi_stark_sppark_forward_panel_bytes(size_t height, size_t width) { - if (!multi_stark_sppark_takes_forward(height, width)) return 0; - const size_t column_bytes = height * sizeof(uint64_t); - return panel_columns(width, column_bytes, panel_budget_bytes()) * column_bytes; +extern "C" int multi_stark_sppark_l2_bytes(int device, size_t* bytes) { + int l2 = 0; + const cudaError_t status = cudaDeviceGetAttribute(&l2, cudaDevAttrL2CacheSize, device); + if (status == cudaSuccess) *bytes = l2 > 0 ? size_t(l2) : size_t(64) << 20; + return static_cast(status); } -// The coset LDE of `trace` (height x width, natural row order, device memory) -// into `values` (extended_height x width, bit-reversed rows), with the coset -// shift powers `shift_powers[i] = shift^i` for i < height. Scratch is -// allocated per call within the panel budget. -// Transforms the adapter has run so far, whatever the metrics setting, so a -// test can tell a proof that went through sppark from one that did not. static std::atomic transforms_run{0}; extern "C" uint64_t multi_stark_sppark_transforms_run() { @@ -533,95 +337,65 @@ extern "C" uint64_t multi_stark_sppark_transforms_run() { } extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, - size_t height, size_t width, size_t added_bits, + const MultiStarkNttPlan* plan, const uint64_t* shift_powers) { transforms_run.fetch_add(1, std::memory_order_relaxed); - if (!trace || !values || !shift_powers || height == 0 || width == 0 || (height & (height - 1))) - return static_cast(cudaErrorInvalidValue); - const size_t extended_height = height << added_bits; - const unsigned log_height = log2_exact(height); - const unsigned log_extended = log_height + static_cast(added_bits); + const size_t height = plan->height, width = plan->width; + const size_t extended_height = plan->extended_height, columns = plan->columns; + const unsigned log_height = log2_exact(height), log_extended = log2_exact(extended_height); cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - const bool fused = fused_expansion(); - const size_t column_bytes = (height + extended_height) * sizeof(uint64_t); - const size_t extra = expansion_extra_bytes(height); - const size_t budget = panel_budget_bytes(); - if (budget < column_bytes + extra) return static_cast(cudaErrorInvalidValue); - const size_t columns = panel_columns(width, column_bytes, budget - extra); - // Pool allocations on the caller's stream: no device synchronization - // per LDE, and the async free below returns them to the pool they came - // from. (CUPTI 2026.2.1's memory tracking faults on an async free of - // memory from cudaMalloc, which this pairing also avoids.) + // Pool allocation and freeing follow all panel kernels on the same stream. + // CUPTI 2026.2.1 faults when tracing an async free of non-pool memory. uint64_t* scratch = nullptr; - status = cudaMallocAsync(reinterpret_cast(&scratch), columns * column_bytes + extra, cudaStreamPerThread); + status = cudaMallocAsync(reinterpret_cast(&scratch), plan->scratch_bytes, cudaStreamPerThread); if (status != cudaSuccess) return static_cast(status); uint64_t* a = scratch; uint64_t* b = scratch + columns * height; - uint64_t* powers = b + columns * extended_height; int result = 0; - if (fused) { - reverse_powers<<>>( - shift_powers, height, log_height, powers); - result = static_cast(cudaGetLastError()); - } StageTimer timer; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; timer.mark(0); result = static_cast(gather_panel(trace, height, width, first, count, height, a)); timer.mark(1); - if (result == 0) result = transform_columns(device, a, log_height, 1, 1, count, height); + if (result == 0) result = transform_columns(device, a, log_height, 1, 1, count, height, plan->inverse_group); timer.mark(2); - const dim3 column_grid(blocks_for_total(extended_height), static_cast(count)); if (result == 0) { - if (fused) - spread_columns<<>>( - a, powers, height, extended_height, static_cast(added_bits), b); - else - shift_columns<<>>( - a, shift_powers, height, log_height, extended_height, b); + const dim3 column_grid(blocks_for_total(extended_height), static_cast(count)); + shift_columns<<>>( + a, shift_powers, height, log_height, extended_height, b); result = static_cast(cudaGetLastError()); } timer.mark(3); - if (result == 0) result = transform_columns(device, b, log_extended, fused ? 2 : 1, 0, count, extended_height); + if (result == 0) result = transform_columns(device, b, log_extended, 1, 0, count, extended_height, plan->forward_group); timer.mark(4); if (result == 0) - result = static_cast( - scatter_panel(b, extended_height, extended_height, width, first, count, log_extended, fused, values)); + result = static_cast(scatter_panel(b, extended_height, extended_height, width, first, count, values)); timer.mark(5); - if (result == 0) timer.report(height, count, added_bits); + if (result == 0) timer.report(height, count, log_extended - log_height); } - // The scratch outlives every kernel that reads it: the free is ordered - // behind them on the same stream. const cudaError_t freed = cudaFreeAsync(scratch, cudaStreamPerThread); if (result == 0) result = static_cast(freed); return result; } -// A forward transform in place on `values` (height x width, natural row -// order) leaving bit-reversed rows: gather, upstream forward per column, -// scatter. The quotient's transforms and the general DFT take this path. -extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, size_t height, size_t width) { +extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, const MultiStarkNttPlan* plan) { transforms_run.fetch_add(1, std::memory_order_relaxed); - if (!values || height == 0 || width == 0 || (height & (height - 1))) - return static_cast(cudaErrorInvalidValue); + const size_t height = plan->height, width = plan->width, columns = plan->columns; const unsigned log_height = log2_exact(height); - if (log_height == 0) return 0; cudaError_t status = cudaSetDevice(device); if (status != cudaSuccess) return static_cast(status); - const size_t columns = panel_columns(width, height * sizeof(uint64_t), panel_budget_bytes()); - if (columns == 0) return static_cast(cudaErrorInvalidValue); uint64_t* panel = nullptr; - status = cudaMallocAsync(reinterpret_cast(&panel), columns * height * sizeof(uint64_t), cudaStreamPerThread); + status = cudaMallocAsync(reinterpret_cast(&panel), plan->scratch_bytes, cudaStreamPerThread); if (status != cudaSuccess) return static_cast(status); int result = 0; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; result = static_cast(gather_panel(values, height, width, first, count, height, panel)); - if (result == 0) result = transform_columns(device, panel, log_height, 1, 0, count, height); + if (result == 0) result = transform_columns(device, panel, log_height, 1, 0, count, height, plan->forward_group); if (result == 0) - result = static_cast(scatter_panel(panel, height, height, width, first, count, log_height, false, values)); + result = static_cast(scatter_panel(panel, height, height, width, first, count, values)); } const cudaError_t freed = cudaFreeAsync(panel, cudaStreamPerThread); if (result == 0) result = static_cast(freed); diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs index 7008ef5..1bbc3cd 100644 --- a/examples/cuda_resident_lde_bench.rs +++ b/examples/cuda_resident_lde_bench.rs @@ -54,8 +54,7 @@ fn main() { (24, 17, 2), (22, 2, 2), (20, 2, 2), - // Short shapes below the sppark height threshold: the aggregation - // circuits' tables and narrow lookup and quotient codewords. + // Short aggregation tables and narrow lookup and quotient codewords. (8, 20, 2), (10, 33, 2), (12, 8, 2), diff --git a/src/batch.rs b/src/batch.rs index 2c42e2d..88504ec 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -773,30 +773,6 @@ mod tests { system.verify_batch(&decoded).unwrap(); } - /// Every transform of a batch proof, main and lookup LDEs and the - /// quotient, produces the same bytes through sppark as through the - /// first-party kernels. - #[cfg(feature = "cuda-sppark")] - #[test] - fn sppark_transforms_prove_the_same_batch() { - use crate::cuda::sppark::{Backend, backend_lock, select_backend, transforms_run}; - let _guard = backend_lock(); - let (system, key) = byte_system(config()); - select_backend(Backend::Legacy); - let legacy = system.prove_batch(&key, two_shards(&system), vec![]); - select_backend(Backend::SpparkAllHeights); - let before = transforms_run(); - let candidate = system.prove_batch(&key, two_shards(&system), vec![]); - let through_sppark = transforms_run() - before; - select_backend(Backend::Legacy); - assert!( - through_sppark > 0, - "the proof's transforms ran through sppark" - ); - assert_eq!(candidate.to_bytes().unwrap(), legacy.to_bytes().unwrap()); - system.verify_batch(&candidate).unwrap(); - } - #[test] fn regenerated_shards_prove_the_same_batch() { let (system, key) = byte_system(config()); diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs index 5e432ee..0bc3ff9 100644 --- a/src/cuda/metrics.rs +++ b/src/cuda/metrics.rs @@ -5,7 +5,7 @@ pub(crate) fn emit_snapshot() { return; } const SHAPES: usize = 33 * 4; - const WORDS: usize = 16 + 2 * SHAPES; + const WORDS: usize = 16 + SHAPES; let mut values = vec![0u64; 64 * WORDS]; unsafe { multi_stark_cuda_metrics_snapshot(values.as_mut_ptr(), values.len()); @@ -18,20 +18,15 @@ pub(crate) fn emit_snapshot() { scope = "process_device_cumulative", upload_calls = c[0], upload_requested_bytes = c[1], upload_chunks = c[2], upload_failures = c[3], upload_host_ns = c[4], coset_hits = c[5], coset_misses = c[6], coset_uploaded_bytes = c[7], constant_bytes = c[8], - last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11], - sppark_taken = c[12], sppark_declined = c[13]); - for (backend, shapes) in [ - ("legacy", &c[16..16 + SHAPES]), - ("sppark", &c[16 + SHAPES..]), - ] { - for (shape, &count) in shapes.iter().enumerate() { - if count == 0 { - continue; - } - tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, - scope = "process_device_cumulative", backend, log_height = shape / 4, - width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); + last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11]); + let shapes = &c[16..]; + for (shape, &count) in shapes.iter().enumerate() { + if count == 0 { + continue; } + tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, + scope = "process_device_cumulative", log_height = shape / 4, + width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); } } } diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 89af67a..da2ffd7 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -9,15 +9,14 @@ pub(crate) mod metrics; pub(crate) mod mmcs; #[doc(hidden)] pub mod pcs; -#[cfg(feature = "cuda-sppark")] pub mod sppark; pub(crate) mod witness; use core::ffi::{CStr, c_char, c_void}; use core::mem::{align_of, size_of}; use core::ptr::NonNull; -use std::collections::BTreeMap; -use std::sync::{Arc, RwLock}; +use sppark::{RawPlan, TransformPlan}; +use std::sync::Arc; use crate::expr::{RowOffset, Source}; use crate::graph::{ConstraintGraph, Node}; @@ -33,20 +32,14 @@ use p3_util::log2_strict_usize; const _: () = assert!(size_of::() == size_of::()); const _: () = assert!(align_of::() == align_of::()); -type CachedPowers = Arc<[Goldilocks]>; -type SharedPowerCache = Arc>>; - /// CUDA-backed batched DFT for the Goldilocks field. /// -/// Clones share the host twiddle cache. This is important because the -/// production configuration gives one clone to the PCS and retains another -/// for quotient transforms. +/// Clones share immutable transform plans and their host coset powers. #[derive(Clone, Debug)] pub struct CudaDft { device_id: i32, cpu: Radix2DitParallel, - twiddles: SharedPowerCache<(usize, bool)>, - shift_powers: SharedPowerCache<(usize, u64)>, + planner: sppark::Planner, } impl Default for CudaDft { @@ -72,8 +65,7 @@ impl CudaDft { Self { device_id, cpu: Radix2DitParallel::default(), - twiddles: Arc::default(), - shift_powers: Arc::default(), + planner: sppark::Planner::new(device_id), } } @@ -83,57 +75,18 @@ impl CudaDft { self.device_id } - /// The first-party twiddle table for a transform, or none when the - /// transform runs through sppark, which needs no table; the kernels - /// reject a missing table on the first-party path. - fn legacy_twiddles( - &self, - legacy: bool, - log_height: usize, - inverse: bool, - ) -> Option> { - legacy.then(|| self.twiddles(log_height, inverse)) - } - - fn twiddles(&self, log_height: usize, inverse: bool) -> Arc<[Goldilocks]> { - let key = (log_height, inverse); - if let Some(twiddles) = self - .twiddles - .read() - .expect("twiddle cache poisoned") - .get(&key) - { - return Arc::clone(twiddles); - } - - let mut cache = self.twiddles.write().expect("twiddle cache poisoned"); - Arc::clone(cache.entry(key).or_insert_with(|| { - let root = Goldilocks::two_adic_generator(log_height); - let root = if inverse { root.inverse() } else { root }; - root.powers().take((1 << log_height) / 2).collect().into() - })) + pub(crate) fn forward_plan(&self, height: usize, width: usize) -> Arc { + self.planner.plan(height, width, 0, None) } - fn shift_powers(&self, height: usize, shift: Goldilocks) -> Arc<[Goldilocks]> { - let key = (height, shift.as_canonical_u64()); - if let Some(powers) = self - .shift_powers - .read() - .expect("shift-power cache poisoned") - .get(&key) - { - return Arc::clone(powers); - } - - let mut cache = self - .shift_powers - .write() - .expect("shift-power cache poisoned"); - Arc::clone( - cache - .entry(key) - .or_insert_with(|| shift.powers().take(height).collect().into()), - ) + pub(crate) fn lde_plan( + &self, + height: usize, + width: usize, + added_bits: usize, + shift: Goldilocks, + ) -> Arc { + self.planner.plan(height, width, added_bits, Some(shift)) } fn validate_dimensions(height: usize, width: usize) { @@ -192,7 +145,7 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "host", - backend = lde_backend(height, width, added_bits), + backend = "sppark", device = self.device_id, height, width, @@ -206,13 +159,7 @@ impl CudaDft { .expect("LDE height overflows usize"); Self::validate_dimensions(extended_height, width); - let log_height = log2_strict_usize(height); - let legacy = lde_backend(height, width, added_bits) == "legacy"; - let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); - let forward_twiddles = - self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let shift_powers = self.shift_powers(height, shift); - let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); + let plan = self.lde_plan(height, width, added_bits, shift); let mut handle = core::ptr::null_mut(); // SAFETY: every host buffer has the exact dimensions validated above; // successful creation transfers the device allocation to `CudaLde`. @@ -224,10 +171,7 @@ impl CudaDft { height, width, added_bits, - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(height_inverse), + &plan.raw, ) }; check_cuda(status, "resident coset LDE"); @@ -239,9 +183,7 @@ impl CudaDft { } } - /// Uploads the constants a resident LDE of the shape will need: the - /// coset powers always, the first-party twiddle tables only when the - /// shape stays on the first-party kernels. + /// Uploads the coset powers before taking a device-memory snapshot. pub(crate) fn prepare_coset_lde_constants( &self, height: usize, @@ -253,22 +195,8 @@ impl CudaDft { .expect("LDE height overflows usize"); Self::validate_dimensions(height, 1); Self::validate_dimensions(extended_height, 1); - let legacy = lde_backend(height, 1, added_bits) == "legacy"; - let inverse_twiddles = self.legacy_twiddles(legacy, log2_strict_usize(height), true); - let shift_powers = self.shift_powers(height, shift); - let forward_twiddles = - self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let status = unsafe { - multi_stark_cuda_prepare_lde_constants( - self.device_id, - table_ptr(&inverse_twiddles), - inverse_twiddles.as_ref().map_or(0, |table| table.len()), - shift_powers.as_ptr().cast(), - height, - table_ptr(&forward_twiddles), - forward_twiddles.as_ref().map_or(0, |table| table.len()), - ) - }; + let plan = self.lde_plan(height, 1, added_bits, shift); + let status = unsafe { multi_stark_cuda_prepare_lde_constants(self.device_id, &plan.raw) }; check_cuda(status, "prepare resident LDE constants"); } } @@ -380,7 +308,7 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "generated", - backend = lde_backend(height, width, added_bits), + backend = "sppark", device = self.device_id, height, width, @@ -392,12 +320,7 @@ impl CudaDft { .checked_shl(added_bits.try_into().unwrap()) .expect("LDE height overflow"); Self::validate_dimensions(extended_height, width); - let log_height = log2_strict_usize(height); - let legacy = lde_backend(height, width, added_bits) == "legacy"; - let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); - let forward_twiddles = - self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let shift_powers = self.shift_powers(height, shift); + let plan = self.lde_plan(height, width, added_bits, shift); let mut context = Box::new(generator); let mut handle = core::ptr::null_mut(); let status = unsafe { @@ -407,10 +330,7 @@ impl CudaDft { height, width, added_bits, - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(Goldilocks::ONE.div_2exp_u64(log_height as u64)), + &plan.raw, (&mut *context as *mut Generator).cast(), write_generated_trace, destroy_generated_trace, @@ -443,10 +363,7 @@ unsafe extern "C" { height: usize, width: usize, added_bits: usize, - inverse: *const u64, - shifts: *const u64, - forward: *const u64, - height_inverse: u64, + plan: *const RawPlan, context: *mut c_void, writer: unsafe extern "C" fn(*mut c_void, i32, *mut u64, usize, usize) -> i32, destroy: unsafe extern "C" fn(*mut c_void), @@ -829,77 +746,6 @@ fn encode_quotient_nodes( (nodes, slots, count as usize) } -/// Conservative device-memory requirement for one fused quotient job. -/// -/// The graph evaluator reuses slots as soon as their final consumer has run, -/// so `graph.nodes.len()` can be orders of magnitude larger than the live -/// device scratch. Keep this estimate beside the encoder so admission and the -/// kernel use the same liveness calculation. The scratch term assumes the -/// global-memory path; devices able to fit the slots in shared memory need -/// less than this bound. -/// A host table's pointer for the kernels: null when the transform runs -/// through sppark and no table was built. -fn table_ptr(table: &Option>) -> *const u64 { - table - .as_ref() - .map_or(core::ptr::null(), |table| table.as_ptr().cast()) -} - -/// The backend a forward transform of the shape runs on. -pub(crate) fn forward_backend(height: usize, width: usize) -> &'static str { - #[cfg(feature = "cuda-sppark")] - if sppark::takes_forward(height, width) { - return "sppark"; - } - #[cfg(not(feature = "cuda-sppark"))] - let _ = (height, width); - "legacy" -} - -/// The backend a resident coset LDE of the shape runs on, for spans and -/// for skipping the tables the other backend would need. -pub(crate) fn lde_backend(height: usize, width: usize, added_bits: usize) -> &'static str { - #[cfg(feature = "cuda-sppark")] - if sppark::takes_lde(height, width, added_bits) { - return "sppark"; - } - #[cfg(not(feature = "cuda-sppark"))] - let _ = (height, width, added_bits); - "legacy" -} - -/// The backend the quotient's wide forward transform runs on. -pub(crate) fn quotient_backend( - quotient_size: usize, - quotient_degree: usize, - log_blowup: usize, -) -> &'static str { - #[cfg(feature = "cuda-sppark")] - { - let lde_height = (quotient_size / quotient_degree.max(1)) << log_blowup; - if sppark::takes_forward(quotient_size, 2) - || sppark::takes_forward(lde_height, 2 * quotient_degree) - { - return "sppark"; - } - } - #[cfg(not(feature = "cuda-sppark"))] - let _ = (quotient_size, quotient_degree, log_blowup); - "legacy" -} - -/// The backend a lookup LDE of the shape runs on: its committed width is -/// two columns per lookup group. -pub(crate) fn lookup_backend( - height: usize, - num_lookups: usize, - group_size: usize, - log_blowup: usize, -) -> &'static str { - let groups = num_lookups.div_ceil(group_size.max(1)).max(1); - lde_backend(height, 2 * groups, log_blowup) -} - pub(crate) fn quotient_lde_memory_upper_bound( graph: &ConstraintGraph, public_count: usize, @@ -1140,20 +986,13 @@ pub(crate) fn lookup_graph_lde_memory_upper_bound( .saturating_add(1) .saturating_mul(main_width) .saturating_mul(size_of::()); - // Device-cached twiddles and shift powers may be cold for this height. - let constant_bytes = height - .saturating_div(2) - .saturating_add(height) - .saturating_add(extended_height / 2) - .saturating_mul(size_of::()); Some(( output_bytes, metadata_bytes .saturating_add(message_bytes) .saturating_add(delta_bytes) .saturating_add(scratch_bytes) - .saturating_add(trace_chunk_bytes) - .saturating_add(constant_bytes), + .saturating_add(trace_chunk_bytes), )) } @@ -1317,7 +1156,7 @@ pub(crate) fn quotient_lde_mixed( ) -> CudaLde { let _span = tracing::info_span!( "cuda/quotient_lde", - backend = quotient_backend(quotient_size, quotient_degree, log_blowup), + backend = "sppark", device = dft.device_id, quotient_size, quotient_degree, @@ -1395,16 +1234,8 @@ fn quotient_lde_sources( assert_eq!(alpha.len(), 2 * expected_constraints); let trace_height = quotient_size / quotient_degree; let lde_height = trace_height << log_blowup; - let quotient_twiddles = dft.legacy_twiddles( - forward_backend(quotient_size, 2) == "legacy", - log2_strict_usize(quotient_size), - false, - ); - let lde_twiddles = dft.legacy_twiddles( - forward_backend(lde_height, 2 * quotient_degree) == "legacy", - log2_strict_usize(lde_height), - false, - ); + let quotient_plan = dft.forward_plan(quotient_size, 2); + let lde_plan = dft.forward_plan(lde_height, 2 * quotient_degree); let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(quotient_size) as u64); let weight_step = Goldilocks::GENERATOR.exp_u64(trace_height as u64).inverse(); let weights: Vec<_> = weight_step @@ -1480,8 +1311,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - table_ptr("ient_twiddles), - table_ptr(&lde_twiddles), + "ient_plan.raw, + &lde_plan.raw, weights.as_ptr().cast(), ) } @@ -1520,8 +1351,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - table_ptr("ient_twiddles), - table_ptr(&lde_twiddles), + "ient_plan.raw, + &lde_plan.raw, weights.as_ptr().cast(), ) } @@ -1878,7 +1709,7 @@ pub(crate) fn lookup_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "direct", - backend = lookup_backend(height, num_lookups, group_size, log_blowup), + backend = "sppark", device = dft.device_id, height, num_lookups, @@ -1912,11 +1743,12 @@ pub(crate) fn lookup_lde_resident( assert_eq!(multiplicities.len(), height * num_lookups); assert_eq!(args.len(), height * args_width); let extended_height = height << log_blowup; - let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; - let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); - let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let plan = dft.lde_plan( + height, + 2 * num_lookups.div_ceil(group_size), + log_blowup, + Goldilocks::GENERATOR, + ); let mut tail = [Goldilocks::ZERO; 4]; let mut handle = core::ptr::null_mut(); let status = unsafe { @@ -1935,10 +1767,7 @@ pub(crate) fn lookup_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(height_inverse), + &plan.raw, ) }; check_cuda(status, "resident CUDA lookup LDE"); @@ -1971,7 +1800,7 @@ pub(crate) fn lookup_lde_resident_partitioned( let _span = tracing::info_span!( "cuda/lookup_lde", path = "partitioned", - backend = lookup_backend(height, num_lookups, group_size, log_blowup), + backend = "sppark", device = dft.device_id, height, num_lookups, @@ -1989,11 +1818,12 @@ pub(crate) fn lookup_lde_resident_partitioned( assert_eq!(multiplicities.len(), height * num_lookups); assert_eq!(args.len(), height * args_width); let extended_height = height << log_blowup; - let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; - let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); - let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let plan = dft.lde_plan( + height, + 2 * num_lookups.div_ceil(group_size), + log_blowup, + Goldilocks::GENERATOR, + ); let total_started = std::time::Instant::now(); let create_started = std::time::Instant::now(); let mut pending_handle = core::ptr::null_mut(); @@ -2123,10 +1953,7 @@ pub(crate) fn lookup_lde_resident_partitioned( pending_handle.as_ptr(), &mut handle, tail.as_mut_ptr().cast(), - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(height_inverse), + &plan.raw, ) }; pending.handle = None; @@ -2170,7 +1997,7 @@ pub(crate) fn lookup_graph_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "graph", - backend = lookup_backend(height, graph.lookups.len(), group_size, log_blowup), + backend = "sppark", device = dft.device_id, height, num_lookups = graph.lookups.len(), @@ -2182,11 +2009,12 @@ pub(crate) fn lookup_graph_lde_resident( let num_lookups = lookups.len(); let groups = num_lookups.div_ceil(group_size.max(1)); let extended_height = height << log_blowup; - let legacy = lookup_backend(height, num_lookups, group_size, log_blowup) == "legacy"; - let inverse_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(height), true); - let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); - let forward_twiddles = dft.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let plan = dft.lde_plan( + height, + 2 * num_lookups.div_ceil(group_size), + log_blowup, + Goldilocks::GENERATOR, + ); let mut tail = [Goldilocks::ZERO; 4]; let mut handle = core::ptr::null_mut(); let status = unsafe { @@ -2208,10 +2036,7 @@ pub(crate) fn lookup_graph_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(height_inverse), + &plan.raw, ) }; check_cuda(status, "resident CUDA graph lookup LDE"); @@ -2289,7 +2114,7 @@ impl TwoAdicSubgroupDft for CudaDft { backend = if height == 1 || width == 0 { "noop" } else if Self::use_cuda_dft(height, width) { - forward_backend(height, width) + "sppark" } else { "cpu" } @@ -2302,27 +2127,23 @@ impl TwoAdicSubgroupDft for CudaDft { return self.cpu.dft_batch(matrix); } - let twiddles = self.legacy_twiddles( - forward_backend(height, width) == "legacy", - log2_strict_usize(height), - false, - ); + let plan = self.forward_plan(height, width); // SAFETY: Goldilocks is repr(transparent) over u64 (asserted above), // every u64 bit pattern is a valid Goldilocks value, all buffers have // the element counts implied by height/width, and the FFI call is - // synchronous so the borrowed twiddle buffer outlives device use. + // synchronous so the matrix and transform plan outlive device use. let status = unsafe { multi_stark_cuda_dft_batch( self.device_id, matrix.values.as_mut_ptr().cast(), height, width, - table_ptr(&twiddles), + &plan.raw, ) }; check_cuda(status, "batched DFT"); - // The CUDA DIF kernel writes bit-reversed rows. Wrap that storage so + // The CUDA transform writes bit-reversed rows. Wrap that storage so // callers observe the natural-order evaluations required by the trait. BitReversalPerm::new_view(matrix) } @@ -2349,7 +2170,7 @@ impl TwoAdicSubgroupDft for CudaDft { backend = if width == 0 { "noop" } else if height > 1 && Self::use_cuda_coset_lde(extended_height, width) { - lde_backend(height, width, added_bits) + "sppark" } else { "cpu" } @@ -2370,13 +2191,7 @@ impl TwoAdicSubgroupDft for CudaDft { return self.cpu.coset_lde_batch(matrix, added_bits, shift); } - let log_height = log2_strict_usize(height); - let legacy = lde_backend(height, width, added_bits) == "legacy"; - let inverse_twiddles = self.legacy_twiddles(legacy, log_height, true); - let forward_twiddles = - self.legacy_twiddles(legacy, log2_strict_usize(extended_height), false); - let shift_powers = self.shift_powers(height, shift); - let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); + let plan = self.lde_plan(height, width, added_bits, shift); let mut output = Goldilocks::zero_vec(extended_height * width); // SAFETY: the input/output and cached tables have the exact lengths @@ -2390,10 +2205,7 @@ impl TwoAdicSubgroupDft for CudaDft { height, width, added_bits, - table_ptr(&inverse_twiddles), - shift_powers.as_ptr().cast(), - table_ptr(&forward_twiddles), - raw_u64(height_inverse), + &plan.raw, ) }; check_cuda(status, "coset LDE"); @@ -3101,7 +2913,7 @@ unsafe extern "C" { values: *mut u64, height: usize, width: usize, - twiddles: *const u64, + plan: *const RawPlan, ) -> i32; fn multi_stark_cuda_coset_lde_batch( @@ -3111,10 +2923,7 @@ unsafe extern "C" { height: usize, width: usize, added_bits: usize, - inverse_twiddles: *const u64, - shift_powers: *const u64, - forward_twiddles: *const u64, - height_inverse: u64, + plan: *const RawPlan, ) -> i32; fn multi_stark_cuda_coset_lde_create( @@ -3124,22 +2933,10 @@ unsafe extern "C" { height: usize, width: usize, added_bits: usize, - inverse_twiddles: *const u64, - shift_powers: *const u64, - forward_twiddles: *const u64, - height_inverse: u64, - ) -> i32; - - fn multi_stark_cuda_prepare_lde_constants( - device_id: i32, - inverse_twiddles: *const u64, - inverse_count: usize, - shift_powers: *const u64, - height: usize, - forward_twiddles: *const u64, - forward_count: usize, + plan: *const RawPlan, ) -> i32; + fn multi_stark_cuda_prepare_lde_constants(device_id: i32, plan: *const RawPlan) -> i32; fn multi_stark_cuda_lde_create_from_host( device_id: i32, handle: *mut *mut c_void, @@ -3265,8 +3062,8 @@ unsafe extern "C" { next_step: usize, quotient_degree: usize, log_blowup: usize, - quotient_twiddles: *const u64, - lde_twiddles: *const u64, + quotient_plan: *const RawPlan, + lde_plan: *const RawPlan, slice_weights: *const u64, ) -> i32; fn multi_stark_cuda_quotient_lde_mixed( @@ -3309,8 +3106,8 @@ unsafe extern "C" { next_step: usize, quotient_degree: usize, log_blowup: usize, - quotient_twiddles: *const u64, - lde_twiddles: *const u64, + quotient_plan: *const RawPlan, + lde_plan: *const RawPlan, slice_weights: *const u64, ) -> i32; fn multi_stark_cuda_mixed_lde_open_row( @@ -3405,10 +3202,7 @@ unsafe extern "C" { gamma: *const u64, ext_w: u64, added_bits: usize, - inverse_twiddles: *const u64, - shift_powers: *const u64, - forward_twiddles: *const u64, - height_inverse: u64, + plan: *const RawPlan, ) -> i32; fn multi_stark_cuda_lookup_lde( device_id: i32, @@ -3425,10 +3219,7 @@ unsafe extern "C" { gamma: *const u64, ext_w: u64, added_bits: usize, - inverse_twiddles: *const u64, - shift_powers: *const u64, - forward_twiddles: *const u64, - height_inverse: u64, + plan: *const RawPlan, ) -> i32; fn multi_stark_cuda_lookup_lde_begin_partitioned( device_id: i32, @@ -3463,10 +3254,7 @@ unsafe extern "C" { pending_handle: *mut c_void, output_handle: *mut *mut c_void, total: *mut u64, - inverse_twiddles: *const u64, - shift_powers: *const u64, - forward_twiddles: *const u64, - height_inverse: u64, + plan: *const RawPlan, ) -> i32; fn multi_stark_cuda_lookup_lde_cancel_partitioned( device_id: i32, diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 7f64d94..e61c7dd 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -72,6 +72,14 @@ fn goldilocks_quadratic_inverse_denominators( } pub trait CudaPcsDft: TwoAdicSubgroupDft { + fn coset_lde_workspace_bytes( + &self, + height: usize, + width: usize, + added_bits: usize, + shift: T, + ) -> usize; + fn prepare_coset_lde_constants(&self, height: usize, added_bits: usize, shift: T); fn coset_lde_batch_resident( @@ -614,6 +622,18 @@ where .sum::(); let dft = &self.dft; let log_blowup = self.fri.log_blowup; + let transform_workspaces: Vec<_> = evaluations + .iter() + .map(|(domain, matrix)| { + dft.coset_lde_workspace_bytes( + matrix.height(), + matrix.width(), + log_blowup, + Val::GENERATOR / domain.shift(), + ) + }) + .collect(); + let max_transform_workspace = transform_workspaces.iter().copied().max().unwrap_or(0); let (initial_free, total_bytes) = crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); let minimum_free = std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") @@ -627,8 +647,13 @@ where let release_traces_during_construction = source_bytes .saturating_add(lde_bytes) .saturating_add(minimum_free) + .saturating_add(max_transform_workspace) > initial_free; - if lde_bytes.saturating_add(minimum_free) > initial_free { + if lde_bytes + .saturating_add(minimum_free) + .saturating_add(max_transform_workspace) + > initial_free + { let max_lde_height = evaluations .iter() .map(|(_, matrix)| matrix.height() << log_blowup) @@ -640,7 +665,8 @@ where let tree_workspace_bytes = max_lde_height.saturating_mul(96); let gpu_lde_budget = initial_free .saturating_sub(minimum_free) - .saturating_sub(tree_workspace_bytes); + .saturating_sub(tree_workspace_bytes) + .saturating_sub(max_transform_workspace); let matrix_resources = evaluations .iter() .map(|(_, matrix)| { @@ -706,7 +732,10 @@ where // GPU lane during commitment. Its selected matrices are transformed // temporarily for hashing, then recomputed by the CPU for durable // host storage while CUDA moves on to the retained groups. - let transient_reserve = max_lde_height.saturating_mul(32).saturating_add(64 << 20); + let transient_reserve = max_lde_height + .saturating_mul(32) + .saturating_add(64 << 20) + .saturating_add(max_transform_workspace); let select_transient_plan = |transient_budget| { height_indices .iter() @@ -1035,7 +1064,11 @@ where let wave_size = if release_traces_during_construction { 1 } else { - CUDA_LDE_WAVE + let workspace_budget = initial_free + .saturating_sub(source_bytes) + .saturating_sub(lde_bytes) + .saturating_sub(minimum_free); + (workspace_budget / max_transform_workspace.max(1)).clamp(1, CUDA_LDE_WAVE) }; for wave in evaluations.chunks(wave_size) { let transform = diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 2f31773..115d9a5 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -1,5 +1,5 @@ //! sppark's Goldilocks NTT through the device-pointer adapter in -//! `cuda/sppark_ntt.cu`: the comparison backend behind `cuda-sppark`. +//! `cuda/sppark_ntt.cu`, with shared immutable transform plans. //! //! The contracts the prover relies on, checked by the tests below against //! the CPU reference: a forward transform in natural order equals the DFT, @@ -15,6 +15,170 @@ use p3_goldilocks::Goldilocks; use super::check_cuda; +/// The allocation and launch shape shared by admission and CUDA execution. +#[derive(Clone, Copy, Debug)] +#[repr(C)] +pub(crate) struct RawPlan { + pub height: usize, + pub width: usize, + pub extended_height: usize, + pub columns: usize, + pub inverse_group: usize, + pub forward_group: usize, + pub scratch_bytes: usize, + pub shift_powers: *const u64, +} + +#[derive(Debug)] +pub(crate) struct TransformPlan { + pub raw: RawPlan, + powers: Option>, +} + +// SAFETY: the pointer references immutable host powers owned by this plan. +// CUDA entry points borrow it only for their synchronous host-side upload. +unsafe impl Send for TransformPlan {} +unsafe impl Sync for TransformPlan {} + +impl TransformPlan { + pub(crate) fn scratch_bytes(&self) -> usize { + self.raw.scratch_bytes + } + + pub(crate) fn constant_bytes(&self) -> usize { + self.powers.as_ref().map_or(0, |powers| powers.len() * 8) + } +} + +type PlanKey = (usize, usize, usize, Option); +type PowerCache = std::collections::BTreeMap<(usize, u64), std::sync::Arc<[Goldilocks]>>; + +#[derive(Default, Debug)] +struct Plans { + shapes: std::collections::BTreeMap>, + powers: PowerCache, +} + +/// Immutable per-device settings and shared plans; environment changes cannot +/// change the workspace between admission and execution of a transform. +#[derive(Clone, Debug)] +pub(crate) struct Planner { + panel_bytes: usize, + batch_bytes: usize, + plans: std::sync::Arc>, +} + +impl Planner { + pub(crate) fn new(device: i32) -> Self { + fn setting(name: &str) -> Option { + std::env::var(name).ok().map(|value| { + value + .parse() + .unwrap_or_else(|_| panic!("{name} must be a non-negative byte count")) + }) + } + let mut l2_bytes = 0; + check_cuda( + unsafe { multi_stark_sppark_l2_bytes(device, &mut l2_bytes) }, + "NTT device properties", + ); + Self::with_budgets( + setting("MULTI_STARK_SPPARK_PANEL_BYTES") + .filter(|&n| n != 0) + .unwrap_or(4usize << 30), + setting("MULTI_STARK_SPPARK_BATCH_BYTES").unwrap_or(l2_bytes), + ) + } + + pub(crate) fn with_budgets(panel_bytes: usize, batch_bytes: usize) -> Self { + Self { + panel_bytes, + batch_bytes, + plans: Default::default(), + } + } + + pub(crate) fn plan( + &self, + height: usize, + width: usize, + added_bits: usize, + shift: Option, + ) -> std::sync::Arc { + use p3_field::{PrimeCharacteristicRing, PrimeField64}; + use std::sync::Arc; + assert!( + height.is_power_of_two(), + "NTT height must be a power of two" + ); + assert!(width > 0, "NTT width must be positive"); + let log = height.trailing_zeros() as usize; + assert!( + added_bits <= max_log_domain() && log + added_bits <= max_log_domain(), + "NTT exceeds sppark's compiled domain" + ); + assert!(shift.is_some() || added_bits == 0); + let extended_height = height << added_bits; + extended_height + .checked_mul(width) + .and_then(|n| n.checked_mul(8)) + .expect("NTT matrix byte count overflows usize"); + let column_bytes = extended_height + .checked_add(if shift.is_some() { height } else { 0 }) + .and_then(|n| n.checked_mul(8)) + .expect("NTT column byte count overflows usize"); + assert!( + column_bytes <= self.panel_bytes, + "MULTI_STARK_SPPARK_PANEL_BYTES cannot hold one NTT column: need {column_bytes} bytes, budget {}", + self.panel_bytes + ); + let key = ( + height, + width, + added_bits, + shift.map(|s| s.as_canonical_u64()), + ); + if let Some(plan) = self + .plans + .read() + .expect("NTT plan cache poisoned") + .shapes + .get(&key) + { + return Arc::clone(plan); + } + let mut cache = self.plans.write().expect("NTT plan cache poisoned"); + if let Some(plan) = cache.shapes.get(&key) { + return Arc::clone(plan); + } + let powers = shift.map(|shift| { + Arc::clone( + cache + .powers + .entry((height, shift.as_canonical_u64())) + .or_insert_with(|| shift.powers().take(height).collect().into()), + ) + }); + let columns = (self.panel_bytes / column_bytes).min(width).min(65535); + let group = |rows: usize| (self.batch_bytes / (rows * 8)).max(1).min(columns); + let raw = RawPlan { + height, + width, + extended_height, + columns, + inverse_group: group(height), + forward_group: group(extended_height), + scratch_bytes: columns * column_bytes, + shift_powers: powers + .as_ref() + .map_or(core::ptr::null(), |p| p.as_ptr().cast()), + }; + let plan = Arc::new(TransformPlan { raw, powers }); + cache.shapes.insert(key, Arc::clone(&plan)); + plan + } +} + /// `NTT::InputOutputOrder`: whether the input and the output are in natural /// (`N`) or bit-reversed (`R`) order. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -35,14 +199,10 @@ pub enum Direction { unsafe extern "C" { fn multi_stark_sppark_max_lg_domain() -> c_int; - fn multi_stark_sppark_backend_selected() -> c_int; - fn multi_stark_sppark_select_backend(selected: c_int); - fn multi_stark_sppark_takes(height: usize) -> c_int; + fn multi_stark_sppark_l2_bytes(device: c_int, bytes: *mut usize) -> c_int; + #[cfg(test)] + fn multi_stark_sppark_borrowed_round_trip(device: c_int, values: *mut u64, lg: u32) -> c_int; fn multi_stark_sppark_transforms_run() -> u64; - fn multi_stark_sppark_takes_lde(height: usize, width: usize, added_bits: usize) -> c_int; - fn multi_stark_sppark_takes_forward(height: usize, width: usize) -> c_int; - fn multi_stark_sppark_panel_bytes(height: usize, width: usize, added_bits: usize) -> usize; - fn multi_stark_sppark_forward_panel_bytes(height: usize, width: usize) -> usize; fn multi_stark_sppark_ntt_device( device: c_int, d_inout: *mut u64, @@ -71,80 +231,11 @@ unsafe extern "C" { ) -> c_int; } -/// Whether the prover's resident LDEs take the sppark path: selected by -/// `MULTI_STARK_CUDA_NTT=sppark` or by [`select_backend`]. -pub fn backend_selected() -> bool { - unsafe { multi_stark_sppark_backend_selected() != 0 } -} - -/// Which transforms take the sppark path. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Backend { - /// The first-party kernels only. - Legacy, - /// sppark for LDEs at or above the height threshold - /// (`MULTI_STARK_SPPARK_MIN_LOG_HEIGHT`, 18), where the panel path - /// wins; the first-party kernels below it. - Sppark, - /// sppark at every height: for comparing the paths on small shapes. - SpparkAllHeights, -} - -/// Selects the backend for the rest of the process. Comparisons in one -/// process switch between constructions; concurrent constructions all see -/// the latest value. -pub fn select_backend(backend: Backend) { - let flag = match backend { - Backend::Legacy => 0, - Backend::Sppark => 1, - Backend::SpparkAllHeights => 2, - }; - unsafe { multi_stark_sppark_select_backend(flag) } -} - -/// How many transforms the adapter has run in this process. +/// Number of transform operations launched, including identity shapes. pub fn transforms_run() -> u64 { unsafe { multi_stark_sppark_transforms_run() } } -/// Whether a transform of `height` input rows is tall enough for the sppark -/// path; the shape rules below decide a dispatch. -pub fn takes(height: usize) -> bool { - unsafe { multi_stark_sppark_takes(height) != 0 } -} - -/// Whether a resident coset LDE of the shape takes the sppark path: tall -/// enough, within upstream's compiled domain after expansion, and one -/// column's scratch within the panel budget. -pub fn takes_lde(height: usize, width: usize, added_bits: usize) -> bool { - unsafe { multi_stark_sppark_takes_lde(height, width, added_bits) != 0 } -} - -/// Whether a forward transform of the shape takes the sppark path. -pub fn takes_forward(height: usize, width: usize) -> bool { - unsafe { multi_stark_sppark_takes_forward(height, width) != 0 } -} - -/// The scratch one forward transform of the shape allocates on the sppark -/// path; zero when the shape stays on the first-party kernels. -pub fn forward_panel_bytes(height: usize, width: usize) -> usize { - unsafe { multi_stark_sppark_forward_panel_bytes(height, width) } -} - -/// Serializes tests that switch the process-wide backend, so a comparison -/// sees the backend it selected on both of its constructions. -#[cfg(test)] -pub(crate) fn backend_lock() -> std::sync::MutexGuard<'static, ()> { - static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -/// The scratch the sppark path allocates for that LDE, to admit alongside -/// the trace and the LDE; zero when the first-party kernels take it. -pub fn panel_bytes(height: usize, width: usize, added_bits: usize) -> usize { - unsafe { multi_stark_sppark_panel_bytes(height, width, added_bits) } -} - /// The largest log domain size the compiled upstream parameters support. pub fn max_log_domain() -> usize { usize::try_from(unsafe { multi_stark_sppark_max_lg_domain() }).expect("domain limit") @@ -279,6 +370,7 @@ mod tests { use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; use p3_field::{Field, PrimeCharacteristicRing, PrimeField64}; use p3_matrix::Matrix; + use p3_matrix::bitrev::BitReversibleMatrix; use p3_matrix::dense::RowMajorMatrix; use p3_util::reverse_slice_index_bits; use rand::{RngExt, SeedableRng, rngs::SmallRng}; @@ -333,51 +425,25 @@ mod tests { assert_eq!(values, input); } - /// One resident LDE both ways, comparing stored words and the Merkle root - /// of a commitment over the matrix. - fn resident_lde_both_ways( + fn resident_matches_cpu( + dft: &super::super::CudaDft, matrix: &RowMajorMatrix, added_bits: usize, shift: Goldilocks, ) { - let dft = super::super::CudaDft::new(0); - select_backend(Backend::Legacy); - let legacy = dft.coset_lde_batch_resident(matrix, added_bits, shift); - let legacy_rows = legacy.to_row_major_matrix(); - select_backend(Backend::SpparkAllHeights); - let candidate = dft.coset_lde_batch_resident(matrix, added_bits, shift); - let candidate_rows = candidate.to_row_major_matrix(); - select_backend(Backend::Legacy); - assert_eq!( - raw_words(&candidate_rows.values), - raw_words(&legacy_rows.values), - "height {} width {} blowup {added_bits}", - matrix.height(), - matrix.width() - ); - } - - #[test] - fn resident_lde_matches_the_first_party_kernels_bit_for_bit() { - let _guard = backend_lock(); - let mut rng = SmallRng::seed_from_u64(0x1de5); - for log_height in [0usize, 1, 2, 5, 8, 12, 14] { - for added_bits in [0usize, 1, 2, 3] { - for width in [1usize, 2, 3, 7, 8, 33] { - let height = 1 << log_height; - let matrix = RowMajorMatrix::new( - (0..height * width).map(|_| rng.random()).collect(), - width, - ); - resident_lde_both_ways(&matrix, added_bits, Goldilocks::GENERATOR); - } - } - } + let expected = Radix2DitParallel::::default() + .coset_lde_batch(matrix.clone(), added_bits, shift) + .bit_reverse_rows() + .to_row_major_matrix(); + let actual = dft + .coset_lde_batch_resident(matrix, added_bits, shift) + .to_row_major_matrix(); + assert_eq!(actual, expected); + assert_canonical(&actual.values, "resident LDE"); } #[test] - fn resident_lde_reduces_raw_representatives_like_the_first_party_kernels() { - let _guard = backend_lock(); + fn resident_lde_reduces_raw_representatives() { let p = Goldilocks::ORDER_U64; let words = [0u64, 1, p - 1, p, p + 1, u64::MAX, 7, p + 7]; let height = 1usize << 10; @@ -391,86 +457,25 @@ mod tests { Goldilocks::ONE, Goldilocks::from_u64(11), ] { - resident_lde_both_ways(&RowMajorMatrix::new(values.clone(), width), 2, shift); + resident_matches_cpu( + &super::super::CudaDft::new(0), + &RowMajorMatrix::new(values.clone(), width), + 2, + shift, + ); } } #[test] - fn resident_lde_panels_narrower_than_the_matrix_cover_every_column() { - let _guard = backend_lock(); - // A 2^12 x 33 matrix at blowup 2 needs 16 KiB x 2 per column, so a - // 256 KiB budget forces panels of a few columns. + fn panel_and_batch_boundaries_match_cpu_without_global_settings() { let mut rng = SmallRng::seed_from_u64(0x9a7e); - let previous = std::env::var("MULTI_STARK_SPPARK_PANEL_BYTES").ok(); - unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", "262144") }; let matrix = RowMajorMatrix::new((0..(1 << 12) * 33).map(|_| rng.random()).collect(), 33); - resident_lde_both_ways(&matrix, 2, Goldilocks::GENERATOR); - match previous { - Some(value) => unsafe { std::env::set_var("MULTI_STARK_SPPARK_PANEL_BYTES", value) }, - None => unsafe { std::env::remove_var("MULTI_STARK_SPPARK_PANEL_BYTES") }, - } - } - - #[test] - fn general_dft_matches_the_first_party_kernels_bit_for_bit() { - let _guard = backend_lock(); - // Shapes above the CUDA DFT threshold of 2^15 cells. - let mut rng = SmallRng::seed_from_u64(0xdf7); - let dft = super::super::CudaDft::new(0); - for (log_height, width) in [(15usize, 1usize), (12, 8), (10, 33), (16, 2), (11, 129)] { - let matrix = RowMajorMatrix::new( - (0..(1 << log_height) * width) - .map(|_| rng.random()) - .collect(), - width, - ); - select_backend(Backend::Legacy); - let legacy = dft.dft_batch(matrix.clone()).to_row_major_matrix(); - select_backend(Backend::SpparkAllHeights); - let candidate = dft.dft_batch(matrix).to_row_major_matrix(); - select_backend(Backend::Legacy); - assert_eq!( - raw_words(&candidate.values), - raw_words(&legacy.values), - "2^{log_height} x {width}" - ); - } - } - - #[test] - fn host_coset_lde_matches_the_first_party_kernels_bit_for_bit() { - let _guard = backend_lock(); - // The host entry takes matrices of width at most two whose extended - // height reaches 2^15. - let mut rng = SmallRng::seed_from_u64(0x1e5); - let dft = super::super::CudaDft::new(0); - let generator = Goldilocks::GENERATOR; - for (log_height, width, added_bits) in [ - (14usize, 1usize, 1usize), - (14, 2, 2), - (15, 2, 3), - (16, 1, 1), - ] { - let matrix = RowMajorMatrix::new( - (0..(1 << log_height) * width) - .map(|_| rng.random()) - .collect(), - width, - ); - select_backend(Backend::Legacy); - let legacy = dft - .coset_lde_batch(matrix.clone(), added_bits, generator) - .to_row_major_matrix(); - select_backend(Backend::SpparkAllHeights); - let candidate = dft - .coset_lde_batch(matrix, added_bits, generator) - .to_row_major_matrix(); - select_backend(Backend::Legacy); - assert_eq!( - raw_words(&candidate.values), - raw_words(&legacy.values), - "2^{log_height} x {width} blowup {added_bits}" - ); + for batch_bytes in [0, 64 << 10, 64 << 20] { + let mut dft = super::super::CudaDft::new(0); + dft.planner = Planner::with_budgets(256 << 10, batch_bytes); + let plan = dft.lde_plan(matrix.height(), matrix.width(), 2, Goldilocks::GENERATOR); + assert!(plan.raw.columns < matrix.width()); + resident_matches_cpu(&dft, &matrix, 2, Goldilocks::GENERATOR); } } @@ -507,67 +512,45 @@ mod tests { } } - /// With column batching switched off the panel path launches its - /// columns one by one and still matches the first-party kernels. #[test] - fn unbatched_columns_match_the_first_party_kernels() { - let _guard = backend_lock(); - let previous = std::env::var("MULTI_STARK_SPPARK_BATCH_BYTES").ok(); - unsafe { std::env::set_var("MULTI_STARK_SPPARK_BATCH_BYTES", "0") }; - let mut rng = SmallRng::seed_from_u64(0x0ff); - for (log_height, width, added_bits) in [(10usize, 5usize, 2usize), (12, 33, 1), (8, 3, 3)] { - let matrix = RowMajorMatrix::new( - (0..(1 << log_height) * width) - .map(|_| rng.random()) - .collect(), - width, + fn plan_reuses_constants_and_rejects_unrepresentable_shapes() { + use std::sync::Arc; + let planner = Planner::with_budgets(256 << 10, 64 << 10); + let plan = planner.plan(1 << 12, 33, 2, Some(Goldilocks::GENERATOR)); + assert!(Arc::ptr_eq( + &plan, + &planner.plan(1 << 12, 33, 2, Some(Goldilocks::GENERATOR)) + )); + assert!(plan.scratch_bytes() <= 256 << 10); + let narrow = planner.plan(1 << 12, 1, 2, Some(Goldilocks::GENERATOR)); + assert_eq!(plan.raw.shift_powers, narrow.raw.shift_powers); + for (height, width, bits) in [ + (1 << 29, 1, 0), + (1 << 24, 1, 5), + (1 << 20, 1, 2), + (2, usize::MAX, 0), + ] { + assert!( + std::panic::catch_unwind(|| planner.plan( + height, + width, + bits, + Some(Goldilocks::ONE) + )) + .is_err() ); - resident_lde_both_ways(&matrix, added_bits, Goldilocks::GENERATOR); - } - match previous { - Some(value) => unsafe { std::env::set_var("MULTI_STARK_SPPARK_BATCH_BYTES", value) }, - None => unsafe { std::env::remove_var("MULTI_STARK_SPPARK_BATCH_BYTES") }, } } #[test] - fn the_height_threshold_and_the_panel_budget_decide_dispatch_and_scratch() { - let _guard = backend_lock(); - select_backend(Backend::Sppark); - assert!( - !takes(1 << 12), - "short transforms stay on the first-party kernels" + fn borrowed_stream_preserves_order_and_caller_ownership() { + let mut values = random(12, 0xb0770); + let expected = values.clone(); + check_cuda( + unsafe { multi_stark_sppark_borrowed_round_trip(0, values.as_mut_ptr().cast(), 12) }, + "borrowed stream round trip", ); - assert!(!takes(1 << 17)); - assert!(takes(1 << 18)); - assert!(takes(1 << 20)); - assert!(takes_lde(1 << 20, 533, 2)); - assert!(takes_forward(1 << 20, 2)); - assert!( - !takes_lde(1 << 24, 6, 5), - "2^29 output rows are beyond the compiled domain" - ); - assert!(!takes_forward(1 << 29, 2)); - assert_eq!(panel_bytes(1 << 12, 533, 2), 0); - assert_eq!(panel_bytes(1 << 24, 6, 5), 0); - // A forward transform's panel is one column set at the height. - assert_eq!(forward_panel_bytes(1 << 22, 2), 2 * (1 << 22) * 8); - // 2^20 rows, 533 columns, blowup 4: (2^20 + 2^22) x 8 bytes per - // column is 40 MiB, so a 4 GiB budget admits 102 columns, and the - // bit-reversed feed's reversed powers when that expansion is on. - let column_bytes = ((1 << 20) + (1 << 22)) * 8; - let extra = if std::env::var("MULTI_STARK_SPPARK_FUSED").as_deref() == Ok("1") { - (1 << 20) * 8 - } else { - 0 - }; - assert_eq!(panel_bytes(1 << 20, 533, 2), 102 * column_bytes + extra); - assert!(panel_bytes(1 << 20, 533, 2) <= 4 << 30); - select_backend(Backend::SpparkAllHeights); - assert!(takes(2)); - select_backend(Backend::Legacy); - assert!(!takes(1 << 24)); - assert_eq!(panel_bytes(1 << 24, 6, 2), 0); + assert_eq!(values, expected); } #[test] diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index 6377826..e3dc72b 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -78,16 +78,17 @@ pub(crate) fn commit( let mut host: Vec>> = Vec::with_capacity(evaluations.len()); let mut retained = Vec::with_capacity(evaluations.len()); for (domain, source) in evaluations { - #[cfg(feature = "cuda-sppark")] - let panel = super::sppark::panel_bytes(source.height(), source.width(), blowup); - #[cfg(not(feature = "cuda-sppark"))] - let panel = 0; + let shift = Val::GENERATOR / domain.shift(); + let plan = pcs + .dft + .lde_plan(source.height(), source.width(), blowup, shift); let needed = source .height() .saturating_mul(source.width()) .saturating_mul(8) .saturating_mul((1 << blowup) + 1) - .saturating_add(panel) + .saturating_add(plan.scratch_bytes()) + .saturating_add(plan.constant_bytes()) .saturating_add(reserve); // Generator caches go before any LDE spills: they are rebuilt from // host seeds on demand, a spilled LDE is uploaded again. @@ -110,7 +111,6 @@ pub(crate) fn commit( device_memory_info(device).0 >= needed, "generated trace commitment exceeds device admission; reduce the shard cell budget" ); - let shift = Val::GENERATOR / domain.shift(); let source_span = tracing::info_span!( "stark/commit_source", kind = match &source { diff --git a/src/prover.rs b/src/prover.rs index e1adbd1..4bf1f01 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -1215,45 +1215,6 @@ mod tests { use p3_dft::{NaiveDft, Radix2DitParallel}; use rand::{RngExt, SeedableRng, rngs::SmallRng}; - /// Host model of the CUDA radix-2 DIF kernel. It intentionally mirrors - /// `cuda/kernels.cu` stage/index/twiddle ordering and returns raw - /// bit-reversed storage. - fn cuda_dif_model(mut matrix: RowMajorMatrix, inverse: bool) -> RowMajorMatrix { - let height = matrix.height(); - if height <= 1 { - return matrix; - } - let width = matrix.width(); - let log_height = log2_strict_usize(height); - let root = Val::two_adic_generator(log_height); - let root = if inverse { root.inverse() } else { root }; - let twiddles: Vec<_> = root.powers().take(height / 2).collect(); - - let mut half = height / 2; - loop { - let stride = height / (2 * half); - for butterfly in 0..height / 2 { - let offset = butterfly % half; - let group = butterfly / half; - let row_0 = group * (2 * half) + offset; - let row_1 = row_0 + half; - for column in 0..width { - let index_0 = row_0 * width + column; - let index_1 = row_1 * width + column; - let left = matrix.values[index_0]; - let right = matrix.values[index_1]; - matrix.values[index_0] = left + right; - matrix.values[index_1] = (left - right) * twiddles[offset * stride]; - } - } - if half == 1 { - break; - } - half >>= 1; - } - matrix - } - /// `lde_from_coefficients` must reproduce, value for value, the matrix /// `TwoAdicFriPcs::commit` stores for the same polynomials given as /// trace-domain evaluations: `coset_lde_batch` with the generator shift @@ -1365,50 +1326,4 @@ mod tests { } } } - - /// Pins the first-party CUDA kernel's DIF ordering and fused coset-LDE - /// pipeline even on machines without a CUDA toolkit or GPU. - #[test] - fn cuda_dif_and_coset_lde_model_match_cpu() { - let mut rng = SmallRng::seed_from_u64(3); - let radix = Radix2DitParallel::::default(); - for log_height in [0usize, 1, 2, 5, 8] { - for width in [1usize, 2, 7] { - let height = 1 << log_height; - let matrix = - RowMajorMatrix::new((0..height * width).map(|_| rng.random()).collect(), width); - let expected_dft = radix - .dft_batch(matrix.clone()) - .bit_reverse_rows() - .to_row_major_matrix(); - assert_eq!(cuda_dif_model(matrix.clone(), false), expected_dft); - - for added_bits in [0usize, 1, 2, 3] { - // CUDA pipeline: inverse DIF (bit-reversed coefficients), - // bit-reverse + normalize + shift, zero-pad, forward DIF. - let mut coefficients = cuda_dif_model(matrix.clone(), true) - .bit_reverse_rows() - .to_row_major_matrix(); - let height_inverse = Val::ONE.div_2exp_u64(log_height as u64); - let mut shift = Val::ONE; - for row in 0..height { - for value in &mut coefficients.values[row * width..(row + 1) * width] { - *value *= height_inverse * shift; - } - shift *= Val::GENERATOR; - } - coefficients.pad_to_height(height << added_bits, Val::ZERO); - let actual = cuda_dif_model(coefficients, false); - let expected = radix - .coset_lde_batch(matrix.clone(), added_bits, Val::GENERATOR) - .bit_reverse_rows() - .to_row_major_matrix(); - assert_eq!( - actual, expected, - "height=2^{log_height}, blowup=2^{added_bits}, width={width}" - ); - } - } - } - } } diff --git a/src/types.rs b/src/types.rs index 2052746..ec1df5e 100644 --- a/src/types.rs +++ b/src/types.rs @@ -565,20 +565,11 @@ impl StarkGenericConfig for GoldilocksBlake3Config { self.log_blowup, ); let current_staging = staging_bytes(input); - let constant_bytes = quotient_size - .saturating_add(lde_height) - .saturating_div(2) - .saturating_add(quotient_degree) - .saturating_mul(size_of::()); - // The sppark path's column panels for the two forward - // transforms, zero when they stay on the first-party kernels. - #[cfg(feature = "cuda-sppark")] + let constant_bytes = quotient_degree.saturating_mul(size_of::()); + let quotient_plan = self.pcs.dft.forward_plan(quotient_size, 2); + let lde_plan = self.pcs.dft.forward_plan(lde_height, 2 * quotient_degree); let kernel_workspace = kernel_workspace - .saturating_add(crate::cuda::sppark::forward_panel_bytes(quotient_size, 2)) - .saturating_add(crate::cuda::sppark::forward_panel_bytes( - lde_height, - 2 * quotient_degree, - )); + .saturating_add(quotient_plan.scratch_bytes().max(lde_plan.scratch_bytes())); ( index, output_bytes, @@ -822,14 +813,6 @@ impl StarkGenericConfig for GoldilocksBlake3Config { .saturating_mul(2 * size_of::()), ) .saturating_add(arg_offsets.len().saturating_mul(size_of::())) - // Device-cached inverse/forward twiddles and coset - // powers may be cold for this height. - .saturating_add( - height - .saturating_add(height / 2) - .saturating_add(extended_height / 2) - .saturating_mul(size_of::()), - ) }; let main_width = self .pcs @@ -849,18 +832,20 @@ impl StarkGenericConfig for GoldilocksBlake3Config { ) }) .flatten(); - // The sppark path's column panels for the lookup LDE, on - // either path; zero when it stays on the first-party kernels. - #[cfg(feature = "cuda-sppark")] - let sppark_panel = - crate::cuda::sppark::panel_bytes(height, 2 * groups, self.log_blowup); - #[cfg(not(feature = "cuda-sppark"))] - let sppark_panel = 0; + let transform_bytes = if num_lookups == 0 { + 0 + } else { + let plan = + self.pcs + .dft + .lde_plan(height, 2 * groups, self.log_blowup, Val::GENERATOR); + plan.scratch_bytes().saturating_add(plan.constant_bytes()) + }; ( index, output_bytes, - direct_temporary_bytes.saturating_add(sppark_panel), - graph_memory.map(|(_, temporary)| temporary.saturating_add(sppark_panel)), + direct_temporary_bytes.saturating_add(transform_bytes), + graph_memory.map(|(_, temporary)| temporary.saturating_add(transform_bytes)), extended_height, ) }) @@ -1139,6 +1124,20 @@ pub(crate) type Blake3CompressionFunction = CompressionFunctionFromHasher for CudaDft { + fn coset_lde_workspace_bytes( + &self, + height: usize, + width: usize, + added_bits: usize, + shift: Val, + ) -> usize { + if width == 0 { + return 0; + } + let plan = self.lde_plan(height, width, added_bits, shift); + plan.scratch_bytes().saturating_add(plan.constant_bytes()) + } + fn prepare_coset_lde_constants(&self, height: usize, added_bits: usize, shift: Val) { self.prepare_coset_lde_constants(height, added_bits, shift); } From bc864c15f4585ba0b186ceae28c23a4d1d227bd1 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:18:24 +0000 Subject: [PATCH 31/34] Enable device snapshots with event-only tracing filters --- README.md | 2 +- src/cuda/metrics.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 3bbbeca..4d98d4f 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ one column within the panel budget; invalid shapes fail before allocation. | `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Maximum transform scratch; zero selects the default. A budget smaller than one column is a configuration error | | `MULTI_STARK_SPPARK_BATCH_BYTES=` | Bytes per batched launch group; defaults to the device's L2 size, zero launches one column at a time | | `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each panel's gather, inverse, expansion, forward and scatter times; synchronizes each measured panel | -| `AIUR_METRICS=` with `RUST_LOG=prover_metrics=info` | Enables lightweight per-device transform-shape counters; Ix writes them to the selected metrics file | +| `AIUR_METRICS=` | Enables lightweight per-device transform-shape counters; Ix writes them to the selected metrics file | Panel and batching settings are captured when the device DFT is constructed. There is no backend selector, height threshold or alternative expansion mode. diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs index 0bc3ff9..24c98f3 100644 --- a/src/cuda/metrics.rs +++ b/src/cuda/metrics.rs @@ -1,7 +1,7 @@ //! CPU-only snapshots of counters accumulated at CUDA operation boundaries. pub(crate) fn emit_snapshot() { - if !tracing::enabled!(target: "prover_metrics", tracing::Level::INFO) { + if !tracing::event_enabled!(target: "prover_metrics", tracing::Level::INFO) { return; } const SHAPES: usize = 33 * 4; From 7eae0793f7cddad67d45fac0b5ad640dbbb5a1bb Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:29:09 +0000 Subject: [PATCH 32/34] Validate transform plans against CUDA buffer dimensions Reject mismatched heights, widths, expansions, and forward/coset plans before the shared NTT helpers launch. Keep raw plans private and non-copyable, borrow them from their owner, and test that rejected calls preserve host output and do not return resident storage. Refresh the backend descriptions and restore-stage timer label. --- README.md | 4 +- cuda/kernels.cu | 35 +++++++++----- cuda/sppark_ntt.cu | 2 +- docs/cuda-lde-pass-reduction.md | 3 ++ src/cuda/mod.rs | 24 +++++----- src/cuda/sppark.rs | 83 +++++++++++++++++++++++++++++---- 6 files changed, 115 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 4d98d4f..aacb827 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ lookup arguments for shared state. a batteries-included Goldilocks/Blake3 instantiation is provided - **Serialization** — `Proof::to_bytes` / `Proof::from_bytes` via bincode - **Parallel proving** — opt-in via the `parallel` feature flag -- **CUDA transforms** — opt-in first-party Goldilocks DFT/LDE backend via the +- **CUDA transforms** — opt-in sppark Goldilocks DFT/LDE backend via the `cuda` feature; normal builds remain independent of CUDA ## Reference configuration @@ -80,7 +80,7 @@ are enabled by default via `.cargo/config.toml`. ## CUDA acceleration The optional `cuda` feature routes the production Goldilocks configuration's -PCS and quotient transforms through first-party CUDA kernels. CUDA and CPU +PCS and quotient transforms through sppark. CUDA and CPU proofs generated by the same revision are byte-identical and use the same CPU verifier. Canonical field serialization changes newly generated CPU proof bytes relative to revisions before this backend; previously generated proofs remain diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 62ed57c..b92bf65 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -1986,6 +1986,7 @@ cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, static cudaError_t plan_constants(int device, const MultiStarkNttPlan* plan, const uint64_t** powers) { + if (!plan || !plan->shift_powers) return cudaErrorInvalidValue; const auto* host = plan->shift_powers; return cached_device_constants(device, host, plan->height, 3, host[0], plan->height > 1 ? host[1] : 0, powers); @@ -1997,8 +1998,16 @@ static void record_transform(int device, size_t height, size_t width) { multi_stark_metrics::ntt_shape(height, width), 1); } +static bool plan_matches(const MultiStarkNttPlan* plan, size_t height, + size_t width, size_t extended_height) { + return plan && plan->height == height && plan->width == width && + plan->extended_height == extended_height; +} + static cudaError_t coset_lde(int device, const uint64_t* trace, uint64_t* values, + size_t height, size_t width, size_t extended_height, const MultiStarkNttPlan* plan) { + if (!plan_matches(plan, height, width, extended_height)) return cudaErrorInvalidValue; const uint64_t* powers = nullptr; cudaError_t status = plan_constants(device, plan, &powers); if (status != cudaSuccess) return status; @@ -2007,7 +2016,10 @@ static cudaError_t coset_lde(int device, const uint64_t* trace, uint64_t* values return static_cast(multi_stark_sppark_coset_lde(device, trace, values, plan, powers)); } -static cudaError_t forward_in_place(int device, uint64_t* values, const MultiStarkNttPlan* plan) { +static cudaError_t forward_in_place(int device, uint64_t* values, size_t height, + size_t width, const MultiStarkNttPlan* plan) { + if (!plan_matches(plan, height, width, height) || plan->shift_powers) + return cudaErrorInvalidValue; record_transform(device, plan->height, plan->width); return static_cast(multi_stark_sppark_forward(device, values, plan)); } @@ -2029,7 +2041,7 @@ extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, DeviceBuffer device_values; status = copy_to_device(device_values, values, elements); if (status == cudaSuccess) { - status = forward_in_place(device_id, device_values.get(), plan); + status = forward_in_place(device_id, device_values.get(), height, width, plan); } if (status == cudaSuccess) { status = copy_to_host(values, device_values, elements); @@ -2069,7 +2081,7 @@ extern "C" int multi_stark_cuda_coset_lde_batch( cudaMemcpyHostToDevice); } if (status == cudaSuccess) { - status = coset_lde(device_id, device_values.get(), device_values.get(), plan); + status = coset_lde(device_id, device_values.get(), device_values.get(), height, width, extended_height, plan); } if (status == cudaSuccess) { status = copy_to_host(output, device_values, output_elements); @@ -2132,7 +2144,7 @@ static int coset_lde_create( } } if (status == cudaSuccess) - status = coset_lde(device_id, lde->trace_values, lde->values, plan); + status = coset_lde(device_id, lde->trace_values, lde->values, height, width, extended_height, plan); if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status != cudaSuccess) { destroy_resident_lde(lde); @@ -2582,7 +2594,7 @@ extern "C" int multi_stark_cuda_quotient_lde( dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch,0,quotient_size,false); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_plan); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,quotient_plan); ResidentLde* lde = nullptr; if(status==cudaSuccess) { @@ -2599,7 +2611,7 @@ extern "C" int multi_stark_cuda_quotient_lde( quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_plan); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,lde_plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess) { *output_handle=lde; @@ -2808,7 +2820,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( } for(size_t i=0;i<2;++i)if(status==cudaSuccess&&stream_busy[i])status=cudaStreamSynchronize(streams[i]); - if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_plan); + if(status==cudaSuccess)status=forward_in_place(device_id, quotient,quotient_size,2,quotient_plan); ResidentLde* lde=nullptr; if(status==cudaSuccess)status=create_resident_lde(&lde); if(status==cudaSuccess){lde->height=lde_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t));} @@ -2818,7 +2830,7 @@ extern "C" int multi_stark_cuda_quotient_lde_mixed( lde->values,quotient,device_weights,quotient_size,trace_height,quotient_degree,2); status=cudaGetLastError(); } - if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_plan); + if(status==cudaSuccess)status=forward_in_place(device_id, lde->values,lde_height,width,lde_plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else if(lde)destroy_resident_lde(lde); for(size_t i=0;i<2;++i){ @@ -3103,7 +3115,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han if(status==cudaSuccess)status=exclusive_scan_ext2(reinterpret_cast(lde->values),deltas,count); if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); - if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,plan); + if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,height,width,extended_height,plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); cudaFree(trace_chunk);cudaFree(scratch);cudaFree(deltas);cudaFree(multiplicities);cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(metadata); @@ -3161,7 +3173,7 @@ extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,ui if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); scan_done=now(); - if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,plan); + if(status==cudaSuccess)status=coset_lde(device_id,lde->values,lde->values,height,width,extended_height,plan); if(status==cudaSuccess)status=cudaStreamSynchronize(0); if(profile){const double finished=now();fprintf(stderr, "[multi-stark/cuda] lookup phases: height=%zu lookups=%zu slots=%zu args_width=%zu allocate=%.3fs rows=%.3fs scan=%.3fs dft=%.3fs\n", @@ -3332,7 +3344,8 @@ extern "C" int multi_stark_cuda_lookup_lde_finish_partitioned( sizeof(Ext2), cudaMemcpyDeviceToHost); if (status == cudaSuccess) - status = coset_lde(device_id, pending->lde->values, pending->lde->values, plan); + status = coset_lde(device_id, pending->lde->values, pending->lde->values, + pending->height, pending->lde->width, pending->lde->height, plan); if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); if (status == cudaSuccess) { *output_handle = pending->lde; diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 0ac8174..24ef2d8 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -296,7 +296,7 @@ struct StageTimer { } void report(size_t height, size_t width, size_t added_bits) { if (!enabled || cudaEventSynchronize(marks[STAGES]) != cudaSuccess) return; - static const char* const names[STAGES] = {"gather", "inverse", "spread", "forward", "scatter"}; + static const char* const names[STAGES] = {"gather", "inverse", "restore", "forward", "scatter"}; fprintf(stderr, "sppark lde height=%zu width=%zu added_bits=%zu", height, width, added_bits); for (int stage = 0; stage < STAGES; ++stage) { float ms = 0; diff --git a/docs/cuda-lde-pass-reduction.md b/docs/cuda-lde-pass-reduction.md index 3390f39..58f62d8 100644 --- a/docs/cuda-lde-pass-reduction.md +++ b/docs/cuda-lde-pass-reduction.md @@ -1,5 +1,8 @@ # Resident LDE and wide NTT experiment: 2026-09-11 +Historical measurements of the first-party NTT removed in September 2026. +Current GPU transforms use sppark; see [CUDA acceleration](../README.md#cuda-acceleration). + Base: `ff3237c` (the short-row BLAKE3 change). Same RTX PRO 6000 Blackwell Server Edition, CUDA 13.3 / native `sm_120`, and eight Rayon threads as the [BLAKE3 experiment](cuda-blake3-short-rows.md). Timings below measure the diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index da2ffd7..ab42f5f 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -171,7 +171,7 @@ impl CudaDft { height, width, added_bits, - &plan.raw, + plan.raw(), ) }; check_cuda(status, "resident coset LDE"); @@ -196,7 +196,7 @@ impl CudaDft { Self::validate_dimensions(height, 1); Self::validate_dimensions(extended_height, 1); let plan = self.lde_plan(height, 1, added_bits, shift); - let status = unsafe { multi_stark_cuda_prepare_lde_constants(self.device_id, &plan.raw) }; + let status = unsafe { multi_stark_cuda_prepare_lde_constants(self.device_id, plan.raw()) }; check_cuda(status, "prepare resident LDE constants"); } } @@ -330,7 +330,7 @@ impl CudaDft { height, width, added_bits, - &plan.raw, + plan.raw(), (&mut *context as *mut Generator).cast(), write_generated_trace, destroy_generated_trace, @@ -1311,8 +1311,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - "ient_plan.raw, - &lde_plan.raw, + quotient_plan.raw(), + lde_plan.raw(), weights.as_ptr().cast(), ) } @@ -1351,8 +1351,8 @@ fn quotient_lde_sources( next_step, quotient_degree, log_blowup, - "ient_plan.raw, - &lde_plan.raw, + quotient_plan.raw(), + lde_plan.raw(), weights.as_ptr().cast(), ) } @@ -1767,7 +1767,7 @@ pub(crate) fn lookup_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - &plan.raw, + plan.raw(), ) }; check_cuda(status, "resident CUDA lookup LDE"); @@ -1953,7 +1953,7 @@ pub(crate) fn lookup_lde_resident_partitioned( pending_handle.as_ptr(), &mut handle, tail.as_mut_ptr().cast(), - &plan.raw, + plan.raw(), ) }; pending.handle = None; @@ -2036,7 +2036,7 @@ pub(crate) fn lookup_graph_lde_resident( gamma.as_ptr().cast(), raw_u64(ext_w), log_blowup, - &plan.raw, + plan.raw(), ) }; check_cuda(status, "resident CUDA graph lookup LDE"); @@ -2138,7 +2138,7 @@ impl TwoAdicSubgroupDft for CudaDft { matrix.values.as_mut_ptr().cast(), height, width, - &plan.raw, + plan.raw(), ) }; check_cuda(status, "batched DFT"); @@ -2205,7 +2205,7 @@ impl TwoAdicSubgroupDft for CudaDft { height, width, added_bits, - &plan.raw, + plan.raw(), ) }; check_cuda(status, "coset LDE"); diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 115d9a5..116a6dd 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -16,22 +16,22 @@ use p3_goldilocks::Goldilocks; use super::check_cuda; /// The allocation and launch shape shared by admission and CUDA execution. -#[derive(Clone, Copy, Debug)] +#[derive(Debug)] #[repr(C)] pub(crate) struct RawPlan { - pub height: usize, - pub width: usize, - pub extended_height: usize, - pub columns: usize, - pub inverse_group: usize, - pub forward_group: usize, - pub scratch_bytes: usize, - pub shift_powers: *const u64, + height: usize, + width: usize, + extended_height: usize, + columns: usize, + inverse_group: usize, + forward_group: usize, + scratch_bytes: usize, + shift_powers: *const u64, } #[derive(Debug)] pub(crate) struct TransformPlan { - pub raw: RawPlan, + raw: RawPlan, powers: Option>, } @@ -41,6 +41,10 @@ unsafe impl Send for TransformPlan {} unsafe impl Sync for TransformPlan {} impl TransformPlan { + pub(crate) fn raw(&self) -> &RawPlan { + &self.raw + } + pub(crate) fn scratch_bytes(&self) -> usize { self.raw.scratch_bytes } @@ -542,6 +546,65 @@ mod tests { } } + #[test] + fn ffi_rejects_mismatched_transform_plans() { + let dft = super::super::CudaDft::new(0); + let input = vec![Goldilocks::ONE; 8 * 3]; + for plan in [ + dft.forward_plan(4, 3), + dft.forward_plan(8, 4), + dft.lde_plan(8, 3, 1, Goldilocks::GENERATOR), + dft.lde_plan(8, 3, 0, Goldilocks::GENERATOR), + ] { + let mut values = input.clone(); + let status = unsafe { + super::super::multi_stark_cuda_dft_batch( + 0, + values.as_mut_ptr().cast(), + 8, + 3, + plan.raw(), + ) + }; + assert_eq!(status, 1, "cudaErrorInvalidValue"); + assert_eq!(values, input); + } + for plan in [ + dft.lde_plan(4, 3, 2, Goldilocks::GENERATOR), + dft.lde_plan(8, 4, 1, Goldilocks::GENERATOR), + dft.lde_plan(8, 3, 2, Goldilocks::GENERATOR), + ] { + let mut output = vec![Goldilocks::TWO; 16 * 3]; + let status = unsafe { + super::super::multi_stark_cuda_coset_lde_batch( + 0, + output.as_mut_ptr().cast(), + input.as_ptr().cast(), + 8, + 3, + 1, + plan.raw(), + ) + }; + assert_eq!(status, 1, "cudaErrorInvalidValue"); + assert!(output.iter().all(|v| *v == Goldilocks::TWO)); + let mut handle = core::ptr::null_mut(); + let status = unsafe { + super::super::multi_stark_cuda_coset_lde_create( + 0, + &mut handle, + input.as_ptr().cast(), + 8, + 3, + 1, + plan.raw(), + ) + }; + assert_eq!(status, 1, "cudaErrorInvalidValue"); + assert!(handle.is_null()); + } + } + #[test] fn borrowed_stream_preserves_order_and_caller_ownership() { let mut values = random(12, 0xb0770); From 6516490741bba25db4cf39e79c375c3f9df07ee1 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Mon, 28 Sep 2026 14:18:52 -0400 Subject: [PATCH 33/34] Format src/cuda/mmcs.rs with rustfmt --- src/cuda/mmcs.rs | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs index c871d1f..41e7d7f 100644 --- a/src/cuda/mmcs.rs +++ b/src/cuda/mmcs.rs @@ -1075,8 +1075,7 @@ impl CudaCommitMmcs for CudaMmcs { })); if !wide_heights.is_empty() { let mut resident: Vec> = Vec::with_capacity(ldes.len()); - let mut host: Vec>> = - Vec::with_capacity(ldes.len()); + let mut host: Vec>> = Vec::with_capacity(ldes.len()); for lde in ldes { if wide_heights.contains(&lde.height()) { host.push(Some(lde.to_row_major_matrix())); @@ -1496,7 +1495,9 @@ mod tests { // 4096 columns of 8 bytes is exactly the 32 KiB the leaf kernel takes. assert!(host_hashed_heights([d(4096, 8)]).is_empty()); assert_eq!( - host_hashed_heights([d(4097, 8)]).into_iter().collect::>(), + host_hashed_heights([d(4097, 8)]) + .into_iter() + .collect::>(), vec![8] ); // Matrices at one height share a leaf row, so their widths add up; @@ -1514,8 +1515,17 @@ mod tests { #[test] fn resident_commit_with_a_wide_height_group_matches_the_cpu() { let dims = [ - (4096, 14), (2048, 14), (4096, 17), (4, 37), (4, 35), (4, 371), - (4, 9282), (256, 110), (4, 6), (1024, 3), (262144, 10), + (4096, 14), + (2048, 14), + (4096, 17), + (4, 37), + (4, 35), + (4, 371), + (4, 9282), + (256, 110), + (4, 6), + (1024, 3), + (262144, 10), ]; let matrices: Vec<_> = dims .iter() @@ -1529,8 +1539,10 @@ mod tests { ); let (expected_commitment, expected_data) = cpu.commit(matrices.clone()); let indices = [0, 3, 1000, 262143]; - let expected_openings: Vec<_> = - indices.iter().map(|&i| cpu.open_batch(i, &expected_data)).collect(); + let expected_openings: Vec<_> = indices + .iter() + .map(|&i| cpu.open_batch(i, &expected_data)) + .collect(); let mmcs = CudaMmcs::with_device(cpu, 0); let ldes = matrices .iter() From c9532e37438f72e17fcb59e78c66c32c4d51e0b7 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Fri, 2 Oct 2026 09:30:54 -0400 Subject: [PATCH 34/34] Drop the unused metrics layer, debug knobs and experiment docs Nothing read the per-device counters, the prover_metrics tracing events or the sppark transform counter, so they go, together with the AIUR_METRICS, MULTI_STARK_SPPARK_BATCH_BYTES, MULTI_STARK_SPPARK_STAGE_TIMING, MULTI_STARK_CUDA_TRACE_FORCE_SPILL and MULTI_STARK_CUDA_LOOKUP_TRACE_TILE_ROWS settings. Launch groups always follow the device's L2 size. The two dated experiment docs described the first-party NTT this branch replaced, and the standalone BLAKE3 row bench no longer linked once kernels.cu called into sppark. The short-row dispatch boundary is noted in cuda/README.md instead. Also removes the unused generated-tile download, the device-pointer NTT entry, a dead bit-reversal helper and the TraceGenerator::host_bytes method, parses MULTI_STARK_CUDA_MIN_FREE_BYTES in one place, and makes the sppark module crate-private with its host contract entries gated to tests. --- README.md | 16 ++--- build.rs | 1 - cuda/README.md | 3 +- cuda/blake3_rows_bench.cu | 78 ---------------------- cuda/kernels.cu | 75 ++------------------- cuda/metrics.cuh | 64 ------------------ cuda/sppark_ntt.cu | 87 ++---------------------- docs/cuda-benchmarks.md | 6 -- docs/cuda-blake3-short-rows.md | 85 ----------------------- docs/cuda-lde-pass-reduction.md | 100 ---------------------------- examples/cuda_resident_lde_bench.rs | 3 +- src/batch.rs | 4 -- src/cuda/metrics.rs | 36 ---------- src/cuda/mod.rs | 71 ++++++-------------- src/cuda/pcs.rs | 11 +-- src/cuda/sppark.rs | 98 +++++++++------------------ src/cuda/witness.rs | 18 +---- src/types.rs | 6 -- src/witness.rs | 1 - 19 files changed, 78 insertions(+), 685 deletions(-) delete mode 100644 cuda/blake3_rows_bench.cu delete mode 100644 cuda/metrics.cuh delete mode 100644 docs/cuda-blake3-short-rows.md delete mode 100644 docs/cuda-lde-pass-reduction.md delete mode 100644 src/cuda/metrics.rs diff --git a/README.md b/README.md index aacb827..641add6 100644 --- a/README.md +++ b/README.md @@ -116,17 +116,11 @@ and coset powers. Admission and execution share the plan. Dimensions must fit the compiled domain (currently 2^28 rows), checked byte arithmetic and at least one column within the panel budget; invalid shapes fail before allocation. -| Setting | Effect | -| --- | --- | -| `MULTI_STARK_SPPARK_PANEL_BYTES=4294967296` | Maximum transform scratch; zero selects the default. A budget smaller than one column is a configuration error | -| `MULTI_STARK_SPPARK_BATCH_BYTES=` | Bytes per batched launch group; defaults to the device's L2 size, zero launches one column at a time | -| `MULTI_STARK_SPPARK_STAGE_TIMING=1` | Print each panel's gather, inverse, expansion, forward and scatter times; synchronizes each measured panel | -| `AIUR_METRICS=` | Enables lightweight per-device transform-shape counters; Ix writes them to the selected metrics file | - -Panel and batching settings are captured when the device DFT is constructed. -There is no backend selector, height threshold or alternative expansion mode. -The fork's `SPPARK_NO_CXX_RUNTIME` mode aborts on CUDA errors with diagnostics; -normal ix builds also abort on Rust panics. +`MULTI_STARK_SPPARK_PANEL_BYTES` caps the transform scratch, 4 GiB by default; +a budget smaller than one column is a configuration error. Columns are +transformed in launch groups sized to the device's L2 cache. Both are captured +when the device DFT is constructed. The fork's `SPPARK_NO_CXX_RUNTIME` mode +aborts on CUDA errors with diagnostics. ```sh cargo test --release --features parallel,cuda --lib cuda::sppark::tests:: diff --git a/build.rs b/build.rs index cc4e9f1..514b734 100644 --- a/build.rs +++ b/build.rs @@ -9,7 +9,6 @@ use std::path::{Path, PathBuf}; use std::process::Command; fn main() { - println!("cargo:rerun-if-changed=cuda/metrics.cuh"); println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=cuda/kernels.cu"); println!("cargo:rerun-if-changed=cuda/goldilocks.cuh"); diff --git a/cuda/README.md b/cuda/README.md index 48d4dc0..4fb5f93 100644 --- a/cuda/README.md +++ b/cuda/README.md @@ -20,7 +20,8 @@ prover pipeline device-resident: 2. Trace matrices are uploaded and transformed through sppark on the caller's stream into resident coset LDEs. 3. First-party BLAKE3 kernels commit mixed-height matrices without copying - LDEs back to the host. + LDEs back to the host. Rows of at most one BLAKE3 chunk (1024 bytes) hash + with one thread per row; longer rows take a warp per row. 4. Lookup traces and quotient LDEs are constructed from resident commitments. 5. Batched openings, reductions, FRI folding, and Merkle authentication remain resident; only protocol-visible openings and proofs return to Rust. diff --git a/cuda/blake3_rows_bench.cu b/cuda/blake3_rows_bench.cu deleted file mode 100644 index fef03d6..0000000 --- a/cuda/blake3_rows_bench.cu +++ /dev/null @@ -1,78 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 -// Standalone resident-kernel comparison; see docs/cuda-blake3-short-rows.md. -// Both kernels run against the same resident bytes, with alternating order. -#include -#include -#include -#include -#include "kernels.cu" - -void checked(cudaError_t error) { - if (error != cudaSuccess) { - fprintf(stderr, "CUDA: %s\n", cudaGetErrorString(error)); - exit(1); - } -} -void launch(bool old, uint8_t* output, const uint8_t* input, - size_t bytes, size_t rows, bool force_stride) { - if (old) { - size_t blocks = std::min((rows + THREADS / 32 - 1) / (THREADS / 32), MAX_BLOCKS); - blake3_hash_rows_kernel<<>>(output, input, bytes, rows); - checked(cudaGetLastError()); - } else if (force_stride && bytes <= BLAKE3_CHUNK_BYTES) { - // Exercise multiple grid-stride iterations without a 16M-row fixture. - blake3_hash_short_rows_kernel<<<2, THREADS, 0, cudaStreamPerThread>>>(output, input, bytes, rows); - checked(cudaGetLastError()); - } else { - checked(launch_blake3_rows(output, input, bytes, rows)); - } -} -int main(int argc, char** argv) { - if (argc > 2 || (argc == 2 && std::strcmp(argv[1], "--check") != 0)) { - fprintf(stderr, "usage: %s [--check]\n", argv[0]); - return 1; - } - const bool small = argc == 2; - checked(cudaSetDevice(0)); - printf("row_bytes,rows,iteration,warp_ms,dispatch_ms,speedup\n"); - for (size_t bytes : {size_t(1), size_t(8), size_t(16), size_t(32), size_t(64), - size_t(128), size_t(256), size_t(320), size_t(512), - size_t(1023), size_t(1024), size_t(1025), size_t(4264), size_t(7400)}) { - size_t rows = small ? 777 : std::min(1 << 20, (128 << 20) / bytes); - std::vector input(bytes * rows); - uint64_t state = 0x81726ab567846217ULL; - for (auto& byte : input) { - state ^= state << 13; state ^= state >> 7; state ^= state << 17; - byte = state; - } - DeviceBuffer source, first, second; - checked(source.allocate((input.size() + 7) / 8)); - checked(first.allocate(rows * 4)); - checked(second.allocate(rows * 4)); - auto* d_source = reinterpret_cast(source.get()); - auto* d_first = reinterpret_cast(first.get()); - auto* d_second = reinterpret_cast(second.get()); - checked(cudaMemcpy(d_source, input.data(), input.size(), cudaMemcpyHostToDevice)); - launch(true, d_first, d_source, bytes, rows, false); - launch(false, d_second, d_source, bytes, rows, small); - std::vector expected(rows * 32), actual(rows * 32); - checked(cudaMemcpy(expected.data(), d_first, expected.size(), cudaMemcpyDeviceToHost)); - checked(cudaMemcpy(actual.data(), d_second, actual.size(), cudaMemcpyDeviceToHost)); - if (expected != actual) { fprintf(stderr, "mismatch for %zu-byte rows\n", bytes); return 2; } - cudaEvent_t start, stop; - checked(cudaEventCreate(&start)); checked(cudaEventCreate(&stop)); - for (int iteration = 0; iteration < (small ? 1 : 7); ++iteration) { - float ms[2]; - for (int order = 0; order < 2; ++order) { - int which = (order + iteration) % 2; - checked(cudaEventRecord(start, cudaStreamPerThread)); - launch(which == 0, which == 0 ? d_first : d_second, d_source, bytes, rows, small); - checked(cudaEventRecord(stop, cudaStreamPerThread)); - checked(cudaEventSynchronize(stop)); - checked(cudaEventElapsedTime(&ms[which], start, stop)); - } - printf("%zu,%zu,%d,%.6f,%.6f,%.3f\n", bytes, rows, iteration, ms[0], ms[1], ms[0]/ms[1]); - } - checked(cudaEventDestroy(start)); checked(cudaEventDestroy(stop)); - } -} diff --git a/cuda/kernels.cu b/cuda/kernels.cu index b92bf65..fe77eac 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -1,8 +1,8 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 // -// First-party CUDA kernels for Goldilocks arithmetic and batched radix-2 -// transforms. The ABI accepts host buffers; device residency is intentionally -// deferred to the later PCS/FRI backend. +// First-party CUDA kernels: Goldilocks arithmetic, BLAKE3 Merkle +// commitments, lookup and quotient construction and FRI over device-resident +// LDEs. Transforms go through the sppark adapter declared in ntt.cuh. #include #include "goldilocks.cuh" @@ -17,11 +17,8 @@ #include #include #include -#include "metrics.cuh" #include "ntt.cuh" - - namespace { using namespace multi_stark_cuda; @@ -61,22 +58,13 @@ cudaError_t cached_device_constants(int device,const uint64_t* host,size_t count while(__sync_lock_test_and_set(&constant_cache_lock,1)){} for(auto* entry=constant_cache;entry;entry=entry->next)if(entry->device==device&& entry->count==count&&entry->kind==kind&&entry->key0==key0&&entry->key1==key1){ - if (kind == 5 && multi_stark_metrics::enabled()) multi_stark_metrics::add(device, multi_stark_metrics::CosetHits, 1); *output=entry->values;__sync_lock_release(&constant_cache_lock);return cudaSuccess;} auto* entry=new(std::nothrow) ConstantCacheEntry; if(!entry){__sync_lock_release(&constant_cache_lock);return cudaErrorMemoryAllocation;} entry->device=device;entry->count=count;entry->kind=kind;entry->key0=key0;entry->key1=key1; cudaError_t status=persistent_malloc(reinterpret_cast(&entry->values),count*sizeof(uint64_t)); if(status==cudaSuccess)status=cudaMemcpy(entry->values,host,count*sizeof(uint64_t),cudaMemcpyHostToDevice); - if(status==cudaSuccess){ - if (multi_stark_metrics::enabled()) { - multi_stark_metrics::add(device, multi_stark_metrics::ConstantBytes, count * sizeof(uint64_t)); - if (kind == 5) { - multi_stark_metrics::add(device, multi_stark_metrics::CosetMisses, 1); - multi_stark_metrics::add(device, multi_stark_metrics::CosetUploadedBytes, count * sizeof(uint64_t)); - } - } - *output=entry->values;entry->next=constant_cache;constant_cache=entry;} + if(status==cudaSuccess){*output=entry->values;entry->next=constant_cache;constant_cache=entry;} else{persistent_free(entry->values);delete entry;} __sync_lock_release(&constant_cache_lock);return status; } @@ -257,13 +245,6 @@ __device__ __forceinline__ void blake3_hash_digest_pair( } } -__device__ __forceinline__ size_t reverse_index_bits(size_t index, unsigned int bits) { - if (bits == 0) { - return 0; - } - return static_cast(__brevll(static_cast(index)) >> (64 - bits)); -} - unsigned int blocks_for(size_t work_items) { if (work_items == 0) { return 0; @@ -548,7 +529,6 @@ void parallel_memcpy(void* destination, const void* source, size_t bytes) { cudaError_t staged_upload(void* device, const void* host, size_t bytes) { const int device_index = current_device_index(); cudaError_t status = cudaSuccess; - multi_stark_metrics::Upload metrics(device_index, bytes, status); uint64_t** slots = upload_staging[device_index]; bool* in_use = upload_staging_in_use[device_index]; pthread_mutex_lock(&upload_staging_mutex); @@ -589,7 +569,6 @@ cudaError_t staged_upload(void* device, const void* host, size_t bytes) { const size_t chunk = std::min(UPLOAD_STAGING_CHUNK, bytes - offset); if (status == cudaSuccess) { parallel_memcpy(buffer, source + offset, chunk); - ++metrics.chunks; status = cudaMemcpyAsync(target + offset, buffer, chunk, cudaMemcpyHostToDevice, cudaStreamPerThread); } @@ -1584,8 +1563,8 @@ __global__ void goldilocks_ops_kernel(uint64_t* sums, uint64_t* differences, // A single-chunk message has no chunk tree to reduce. Assigning one thread // per row keeps every lane doing useful compression instead of leaving 31 -// lanes idle in the warp-per-row kernel. Reuse the same chunk primitive and -// little-endian digest encoding; commitments and BLAKE3 flags are unchanged. +// lanes idle in the warp-per-row kernel, with the same chunk primitive and +// little-endian digest encoding. __global__ void blake3_hash_short_rows_kernel(uint8_t* digests, const uint8_t* messages, size_t message_bytes, @@ -1992,12 +1971,6 @@ static cudaError_t plan_constants(int device, const MultiStarkNttPlan* plan, plan->height > 1 ? host[1] : 0, powers); } -static void record_transform(int device, size_t height, size_t width) { - if (multi_stark_metrics::enabled() && height > 1) - multi_stark_metrics::add(device, multi_stark_metrics::NTT_OFFSET + - multi_stark_metrics::ntt_shape(height, width), 1); -} - static bool plan_matches(const MultiStarkNttPlan* plan, size_t height, size_t width, size_t extended_height) { return plan && plan->height == height && plan->width == width && @@ -2011,8 +1984,6 @@ static cudaError_t coset_lde(int device, const uint64_t* trace, uint64_t* values const uint64_t* powers = nullptr; cudaError_t status = plan_constants(device, plan, &powers); if (status != cudaSuccess) return status; - record_transform(device, plan->height, plan->width); - record_transform(device, plan->extended_height, plan->width); return static_cast(multi_stark_sppark_coset_lde(device, trace, values, plan, powers)); } @@ -2020,7 +1991,6 @@ static cudaError_t forward_in_place(int device, uint64_t* values, size_t height, size_t width, const MultiStarkNttPlan* plan) { if (!plan_matches(plan, height, width, height) || plan->shift_powers) return cudaErrorInvalidValue; - record_transform(device, plan->height, plan->width); return static_cast(multi_stark_sppark_forward(device, values, plan)); } @@ -3055,13 +3025,7 @@ extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_han *output_handle=nullptr;const size_t height=main->trace_height; const size_t groups=(lookup_count+group_size-1)/group_size,width=2*groups; const size_t count=height*groups,extended_height=height<trace_writer ? getenv("MULTI_STARK_CUDA_LOOKUP_TRACE_TILE_ROWS") : nullptr) { - char* end=nullptr; - const unsigned long rows=strtoul(configured,&end,10); - if (end!=configured && *end=='\0' && rows>0 && rows<=LOOKUP_ROWS_PER_CHUNK) - LOOKUP_ROWS_PER_CHUNK=rows; - } + constexpr size_t LOOKUP_ROWS_PER_CHUNK=size_t(1)<<16; const size_t chunk_rows=height(cudaErrorInvalidValue); @@ -4196,11 +4160,6 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, cudaError_t status = cudaSetDevice(device_id); if (status == cudaSuccess) retain_default_pool(device_id); if (status == cudaSuccess) status = cudaMemGetInfo(free_bytes, total_bytes); - if (status == cudaSuccess && multi_stark_metrics::enabled()) { - multi_stark_metrics::sample(device_id, multi_stark_metrics::DriverFreeBytes, *free_bytes); - multi_stark_metrics::sample(device_id, multi_stark_metrics::TotalBytes, *total_bytes); - multi_stark_metrics::add(device_id, multi_stark_metrics::MemorySamples, 1); - } // cudaMemGetInfo excludes pages retained by cudaMallocAsync's default // pool, even though subsequent stream allocations can reuse them. Treat // the unused part of that pool as available for admission decisions; using @@ -4225,23 +4184,3 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, return static_cast(status); } -extern "C" int multi_stark_cuda_generate_trace_rows(int device, void* context, - TraceWriter writer, size_t first, size_t rows, size_t width, uint64_t* output) { - if (!context || !writer || !output || !rows || !width || !product_fits(rows,width)) - return static_cast(cudaErrorInvalidValue); - cudaError_t status = cudaSetDevice(device); - uint64_t* tile = nullptr; - if (status == cudaSuccess) status = cudaMalloc(reinterpret_cast(&tile), rows * width * sizeof(uint64_t)); - if (status == cudaSuccess) status = static_cast(writer(context, device, tile, first, rows)); - if (status == cudaSuccess) status = cudaMemcpy(output, tile, rows * width * sizeof(uint64_t), cudaMemcpyDeviceToHost); - if (tile) cudaFree(tile); - return static_cast(status); -} - -extern "C" void multi_stark_cuda_metrics_snapshot(uint64_t* output, size_t count) { - if (!output || count != multi_stark_metrics::DEVICES * multi_stark_metrics::WORDS) return; - for (size_t device = 0; device < multi_stark_metrics::DEVICES; ++device) - for (size_t key = 0; key < multi_stark_metrics::WORDS; ++key) - output[device * multi_stark_metrics::WORDS + key] = - multi_stark_metrics::counters[device][key].load(std::memory_order_relaxed); -} diff --git a/cuda/metrics.cuh b/cuda/metrics.cuh deleted file mode 100644 index bf6f8c1..0000000 --- a/cuda/metrics.cuh +++ /dev/null @@ -1,64 +0,0 @@ -// SPDX-License-Identifier: MIT OR Apache-2.0 -#pragma once -#include -#include -#include -#include -#include -#include - -namespace multi_stark_metrics { -constexpr size_t DEVICES = 64; -constexpr size_t NTT_OFFSET = 16; -// Transform shape counts for the CUDA NTT. -constexpr size_t NTT_SHAPES = 33 * 4; -constexpr size_t WORDS = NTT_OFFSET + NTT_SHAPES; -enum Counter : size_t { - UploadCalls, UploadRequestedBytes, UploadChunks, UploadFailures, UploadHostNs, - CosetHits, CosetMisses, CosetUploadedBytes, ConstantBytes, - DriverFreeBytes, TotalBytes, MemorySamples -}; -inline size_t ntt_shape(size_t height, size_t width) { - unsigned log = 0; - while ((size_t(1) << log) < height) ++log; - const size_t bucket = width == 1 ? 0 : width == 2 ? 1 : width < 8 ? 2 : 3; - return log * 4 + bucket; -} -std::atomic counters[DEVICES][WORDS]{}; -inline bool enabled() { - static const bool value = std::getenv("AIUR_METRICS") != nullptr; - return value; -} -inline void add(int device, size_t key, uint64_t value) { - if (device >= 0 && device < int(DEVICES)) - counters[device][key].fetch_add(value, std::memory_order_relaxed); -} -inline void sample(int device, size_t key, uint64_t value) { - if (device >= 0 && device < int(DEVICES)) - counters[device][key].store(value, std::memory_order_relaxed); -} -inline uint64_t host_ns() { - timespec ts{}; - clock_gettime(CLOCK_MONOTONIC, &ts); - return uint64_t(ts.tv_sec) * 1000000000 + uint64_t(ts.tv_nsec); -} -struct Upload { - bool active = enabled(); - int device; - size_t bytes; - uint64_t chunks = 0; - cudaError_t& status; - uint64_t start = 0; - Upload(int d, size_t b, cudaError_t& s) : device(d), bytes(b), status(s) { - if (active) start = host_ns(); - } - ~Upload() { - if (!active) return; - add(device, UploadCalls, 1); - add(device, UploadRequestedBytes, bytes); - add(device, UploadChunks, chunks); - add(device, UploadFailures, status != cudaSuccess); - add(device, UploadHostNs, host_ns() - start); - } -}; -} diff --git a/cuda/sppark_ntt.cu b/cuda/sppark_ntt.cu index 24ef2d8..5a8416f 100644 --- a/cuda/sppark_ntt.cu +++ b/cuda/sppark_ntt.cu @@ -1,13 +1,10 @@ // Goldilocks transforms on the caller's stream, with immutable panel plans. #include #include +#include "goldilocks.cuh" #include "ntt.cuh" -#include #include -#include -#include -#include #include namespace { @@ -50,9 +47,8 @@ static int ntt_batch_on_stream(int device, uint64_t* d_inout, uint32_t lg, return static_cast(cudaSuccess); } -extern "C" int multi_stark_sppark_ntt_batch_device(int device, uint64_t* data, uint32_t lg, - int order, int direction, int coset, - uint32_t batch, size_t stride) { +static int ntt_batch_device(int device, uint64_t* data, uint32_t lg, int order, int direction, + int coset, uint32_t batch, size_t stride) { return ntt_batch_on_stream(device, data, lg, order, direction, coset, batch, stride, cudaStreamPerThread); } @@ -79,13 +75,6 @@ extern "C" int multi_stark_sppark_borrowed_round_trip(int device, uint64_t* valu return result; } -// One transform: the batch entry with a single vector. -extern "C" int multi_stark_sppark_ntt_device(int device, uint64_t* d_inout, uint32_t lg, - int order, int direction, int coset) { - return multi_stark_sppark_ntt_batch_device(device, d_inout, lg, order, direction, coset, 1, - size_t(1) << lg); -} - // The batched transform on host memory: `batch * stride` words uploaded, // transformed, downloaded and synchronized. For contract checks and small // inputs, not the prover. @@ -103,7 +92,7 @@ extern "C" int multi_stark_sppark_ntt_batch_host(int device, uint64_t* inout, ui status = cudaMemcpyAsync(d_inout, inout, bytes, cudaMemcpyHostToDevice, cudaStreamPerThread); int result = static_cast(status); if (result == 0) - result = multi_stark_sppark_ntt_batch_device(device, d_inout, lg, order, direction, coset, batch, stride); + result = ntt_batch_device(device, d_inout, lg, order, direction, coset, batch, stride); if (result == 0) result = static_cast( cudaMemcpyAsync(inout, d_inout, bytes, cudaMemcpyDeviceToHost, cudaStreamPerThread)); @@ -123,7 +112,6 @@ extern "C" int multi_stark_sppark_ntt_host(int device, uint64_t* inout, uint32_t // inverse NTT, coefficient-order restoration with coset shift and expansion, // and forward NR transform. Scatter preserves the bit-reversed row storage // consumed by commitments. Gather canonicalizes lazy field representatives. -#include "goldilocks.cuh" namespace { @@ -260,53 +248,6 @@ unsigned log2_exact(size_t value) { return log; } -// A decimal setting, or `fallback` when unset or not a plain number. This -// unit is compiled by nvcc's host compiler without the C standard pin the -// crate's C units get, where glibc redirects strtoul to a C23 symbol the -// Lean toolchain's libc does not carry. -unsigned long long decimal_setting(const char* name, unsigned long long fallback) { - const char* configured = getenv(name); - if (!configured || !*configured) return fallback; - unsigned long long value = 0; - for (const char* c = configured; *c; ++c) { - if (*c < '0' || *c > '9' || value > (~0ull - 9) / 10) return fallback; - value = value * 10 + unsigned(*c - '0'); - } - return value; -} - -// MULTI_STARK_SPPARK_STAGE_TIMING=1 prints the stage times of every coset -// LDE to stderr: events on the caller's stream around each stage, which -// the transforms are fenced to, and one synchronization per panel. -struct StageTimer { - static constexpr int STAGES = 5; - bool enabled = decimal_setting("MULTI_STARK_SPPARK_STAGE_TIMING", 0) != 0; - cudaEvent_t marks[STAGES + 1] = {}; - StageTimer() { - if (!enabled) return; - for (auto& mark : marks) - if (cudaEventCreate(&mark) != cudaSuccess) enabled = false; - } - ~StageTimer() { - for (auto mark : marks) - if (mark) cudaEventDestroy(mark); - } - void mark(int stage) { - if (enabled) cudaEventRecord(marks[stage], cudaStreamPerThread); - } - void report(size_t height, size_t width, size_t added_bits) { - if (!enabled || cudaEventSynchronize(marks[STAGES]) != cudaSuccess) return; - static const char* const names[STAGES] = {"gather", "inverse", "restore", "forward", "scatter"}; - fprintf(stderr, "sppark lde height=%zu width=%zu added_bits=%zu", height, width, added_bits); - for (int stage = 0; stage < STAGES; ++stage) { - float ms = 0; - cudaEventElapsedTime(&ms, marks[stage], marks[stage + 1]); - fprintf(stderr, " %s=%.3f", names[stage], ms); - } - fprintf(stderr, "\n"); - } -}; - // The `count` columns of a panel, `stride` elements apart, through batched // launch sequences of as many columns as the group budget holds. int transform_columns(int device, uint64_t* panel, uint32_t lg, int order, int direction, size_t count, @@ -315,8 +256,8 @@ int transform_columns(int device, uint64_t* panel, uint32_t lg, int order, int d int result = 0; for (size_t first = 0; result == 0 && first < count; first += group) { const size_t batch = count - first < group ? count - first : group; - result = multi_stark_sppark_ntt_batch_device(device, panel + first * stride, lg, order, direction, 0, - static_cast(batch), stride); + result = ntt_batch_device(device, panel + first * stride, lg, order, direction, 0, + static_cast(batch), stride); } return result; } @@ -330,16 +271,9 @@ extern "C" int multi_stark_sppark_l2_bytes(int device, size_t* bytes) { return static_cast(status); } -static std::atomic transforms_run{0}; - -extern "C" uint64_t multi_stark_sppark_transforms_run() { - return transforms_run.load(std::memory_order_relaxed); -} - extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, uint64_t* values, const MultiStarkNttPlan* plan, const uint64_t* shift_powers) { - transforms_run.fetch_add(1, std::memory_order_relaxed); const size_t height = plan->height, width = plan->width; const size_t extended_height = plan->extended_height, columns = plan->columns; const unsigned log_height = log2_exact(height), log_extended = log2_exact(extended_height); @@ -353,27 +287,19 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u uint64_t* a = scratch; uint64_t* b = scratch + columns * height; int result = 0; - StageTimer timer; for (size_t first = 0; result == 0 && first < width; first += columns) { const size_t count = width - first < columns ? width - first : columns; - timer.mark(0); result = static_cast(gather_panel(trace, height, width, first, count, height, a)); - timer.mark(1); if (result == 0) result = transform_columns(device, a, log_height, 1, 1, count, height, plan->inverse_group); - timer.mark(2); if (result == 0) { const dim3 column_grid(blocks_for_total(extended_height), static_cast(count)); shift_columns<<>>( a, shift_powers, height, log_height, extended_height, b); result = static_cast(cudaGetLastError()); } - timer.mark(3); if (result == 0) result = transform_columns(device, b, log_extended, 1, 0, count, extended_height, plan->forward_group); - timer.mark(4); if (result == 0) result = static_cast(scatter_panel(b, extended_height, extended_height, width, first, count, values)); - timer.mark(5); - if (result == 0) timer.report(height, count, log_extended - log_height); } const cudaError_t freed = cudaFreeAsync(scratch, cudaStreamPerThread); if (result == 0) result = static_cast(freed); @@ -381,7 +307,6 @@ extern "C" int multi_stark_sppark_coset_lde(int device, const uint64_t* trace, u } extern "C" int multi_stark_sppark_forward(int device, uint64_t* values, const MultiStarkNttPlan* plan) { - transforms_run.fetch_add(1, std::memory_order_relaxed); const size_t height = plan->height, width = plan->width, columns = plan->columns; const unsigned log_height = log2_exact(height); cudaError_t status = cudaSetDevice(device); diff --git a/docs/cuda-benchmarks.md b/docs/cuda-benchmarks.md index 1af63ef..ed85386 100644 --- a/docs/cuda-benchmarks.md +++ b/docs/cuda-benchmarks.md @@ -39,9 +39,3 @@ cargo run --release --locked --features parallel,cuda --example cuda_blake3_benc The CSV DFT benchmark labels shapes below the production CUDA thresholds as `cpu-fallback`, warms both implementations, uses the same iteration count, and checks each output against the CPU reference outside the timing window. - -## 2026-09-11: BLAKE3 and LDE kernel experiments - -See [short-row BLAKE3 dispatch](cuda-blake3-short-rows.md) and -[resident LDE/wide NTT optimization](cuda-lde-pass-reduction.md) for isolated -before/after measurements, exact proof compatibility, and reproduction steps. diff --git a/docs/cuda-blake3-short-rows.md b/docs/cuda-blake3-short-rows.md deleted file mode 100644 index ec0eefd..0000000 --- a/docs/cuda-blake3-short-rows.md +++ /dev/null @@ -1,85 +0,0 @@ -# Short-row BLAKE3 dispatch: 2026-09-11 - -Base: multi-stark `231942a` on `sb/trace-sharding-gpu`. -Isolated branch: `codex/cuda-blake3-short-rows`. -GPU: RTX PRO 6000 Blackwell Server Edition (96 GB), CUDA 13.3, native `sm_120`. -CPU-side benchmarks used `RAYON_NUM_THREADS=8`. - -## Change - -Messages of 1–1024 bytes use one thread per row, reusing our existing -BLAKE3 chunk compression and raw little-endian digest encoding. Longer -messages keep the existing warp-per-row kernel. The dispatch boundary is -BLAKE3's chunk size; there is no machine-dependent threshold or new runtime -flag. No upstream implementation was copied and no dependency was added. - -The new kernel uses 64 registers/thread versus 96 for the old kernel, with -no register spills in either on `sm_120`. It also avoids the old kernel's -8 KiB shared chunk-value array. Both report a 64-byte local stack frame. - -## Measurements - -The resident benchmark runs both kernels against the same uploaded bytes, -alternates their order, checks all output digests, and times with CUDA events. -Medians of seven samples; at most 1,048,576 rows / 128 MiB input per shape: - -| Row bytes | Existing kernel | New dispatch | Speedup | -| ---: | ---: | ---: | ---: | -| 8 | 1.321 ms | 0.384 ms | 3.44× | -| 64 | 1.353 ms | 0.333 ms | 4.07× | -| 320 | 2.351 ms | 0.413 ms | 5.70× | -| 1023 | 2.488 ms | 0.347 ms | 7.16× | -| 1024 | 2.480 ms | 0.472 ms | 5.26× | -| 1025 | 2.632 ms | 2.629 ms | 1.00× | -| 4264 | 0.801 ms | 0.800 ms | 1.00× | -| 7400 | 0.678 ms | 0.677 ms | 1.00× | - -Complete PCS commitments include host upload, LDEs, row gathering/hashing, -Merkle construction, and commitment return. Allocation/filling of the input -witness precedes timing. Three baseline/candidate processes per variant, -five warm samples each, with later process order reversed; first iterations -are excluded. Commitments matched byte-for-byte for every shape/run. - -| Input shape (height × field columns) | Existing | New | Wall reduction | -| --- | ---: | ---: | ---: | -| 2^20 × 8 (64-byte rows) | 11.403 ms | 9.011 ms | 21.0% | -| 2^18 × 40 (320-byte rows) | 12.871 ms | 10.594 ms | 17.7% | -| 2^18 × 128 (1024-byte rows) | 43.977 ms | 35.930 ms | 18.3% | -| 2^18 × 129 (1032-byte rows) | 44.759 ms | 43.986 ms | 1.7% | -| 2^16 × 925 (7400-byte rows) | 94.049 ms | 91.704 ms | 2.5% | -| Mixed heights/widths | 71.689 ms | 68.297 ms | 4.7% | - -The mixed case combines `(2^18,4)`, `(2^18,12)`, `(2^18,24)`, `(2^17,129)`, -and `(2^16,533)`. Equal-height rows concatenate before hashing. The small -wide-only commitment differences are timing variation: their row kernels -are unchanged. These synthetic measurements establish a useful local win; -they do not establish an Init/Mathlib or Stage 2 wall-time improvement. - -## Validation and reproduction - -- All 70 CUDA-enabled library tests passed (5.46 seconds). -- CPU BLAKE3 comparisons cover 25 row lengths, 1–32768 bytes, and eight row - counts across warp/block boundaries, using seeded random messages. -- Merkle roots/openings cover both sides of the 1024-byte boundary; - existing mixed-height, resident, hybrid, and batch-proof tests passed. -- CUDA memcheck reported zero errors on both kernels, with the new kernel - deliberately launched with two blocks to exercise its grid-stride loop. -- CPU, baseline CUDA, and new CUDA `proof_compatibility` outputs were identical - (17,213 bytes), and each proof verified. SHA-256: - `25564a01d1d352b1ec2de56b019b641d24acc81083133e79274a86829b2a5dd5`. -- Clippy completed without warnings for the library and new commit benchmark. - -```sh -MULTI_STARK_CUDA_ARCHS=120 cargo test --release --locked --features parallel,cuda --lib -- --test-threads=1 -nvcc -O3 --std=c++17 --default-stream=per-thread -arch=sm_120 cuda/blake3_rows_bench.cu -o /tmp/blake3-rows-bench -/tmp/blake3-rows-bench -compute-sanitizer --tool memcheck --error-exitcode 1 /tmp/blake3-rows-bench --check -RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_commit_bench -``` - -Select the native architecture appropriate to the test GPU. To compare full -commitments, build the same `cuda_commit_bench.rs` against the base and this -change, preserve both binaries, and alternate runs on an idle GPU. The -benchmark reports cold iteration zero separately and the exact commitment -beside each timing. Existing `cuda_blake3_bench` includes host transfers and -CPU hashing; it is not a resident-kernel timer. diff --git a/docs/cuda-lde-pass-reduction.md b/docs/cuda-lde-pass-reduction.md deleted file mode 100644 index 58f62d8..0000000 --- a/docs/cuda-lde-pass-reduction.md +++ /dev/null @@ -1,100 +0,0 @@ -# Resident LDE and wide NTT experiment: 2026-09-11 - -Historical measurements of the first-party NTT removed in September 2026. -Current GPU transforms use sppark; see [CUDA acceleration](../README.md#cuda-acceleration). - -Base: `ff3237c` (the short-row BLAKE3 change). Same RTX PRO 6000 Blackwell -Server Edition, CUDA 13.3 / native `sm_120`, and eight Rayon threads as the -[BLAKE3 experiment](cuda-blake3-short-rows.md). Timings below measure the -incremental effect of this change; BLAKE3 is enabled in both binaries. - -## Change retained - -- Use the existing fused radix-8/radix-4 stages for matrices with at least - eight columns regardless of height. Previously a separate `height >= 2^18` - condition left short, wide matrices on individual radix-2 stages even when - their total cell count was large. The existing narrow-column path remains. -- Zero only the padded tail of the resident LDE allocation. The input prefix - is overwritten by the trace copy and does not need a preceding clear. -- Remove the resident LDE's final canonicalization pass. Its normalization - and butterfly arithmetic already produce canonical Goldilocks values, - including height-one transforms. Tests inspect the actual stored u64 - representation, rather than an accessor that canonicalizes values. - -No new NTT implementation, matrix layout, workspace allocation, dependency, -or machine-specific threshold was introduced. Trace retention and output -storage order are unchanged. The NTT dispatch change applies to other callers -of the same transform helper; memory-pass removal is scoped to -`multi_stark_cuda_coset_lde_create`. - -An earlier trial fused bit reversal, coset scaling, virtual zero padding, -and the first forward stage in place, using bit-reversal pairs to avoid -races. It passed correctness checks but did not improve complete commitments -consistently, so that kernel and its transform-continuation helper were -removed. The simpler change above is the measured result. - -## Measurements - -Resident LDE timings include upload, inverse transform, coset handling, -forward transform, and completion, excluding construction of the host input. -Medians of 14 warm samples per variant across two processes; first use of -each shape is reported separately and excluded. Baseline and candidate -processes were interleaved on an idle GPU. - -| Input height | Columns | Blowup | BLAKE3-only | + LDE change | Reduction | -| ---: | ---: | ---: | ---: | ---: | ---: | -| 2^20 | 1 | 2 | 0.992 ms | 0.965 ms | 2.7% | -| 2^20 | 2 | 2 | 1.805 ms | 1.833 ms | -1.6% | -| 2^20 | 8 | 2 | 6.563 ms | 6.435 ms | 1.9% | -| 2^18 | 40 | 2 | 8.582 ms | 8.451 ms | 1.5% | -| 2^18 | 128 | 2 | 30.389 ms | 30.093 ms | 1.0% | -| 2^18 | 129 | 2 | 30.533 ms | 30.099 ms | 1.4% | -| 2^16 | 925 | 2 | 79.048 ms | 57.660 ms | 27.1% | -| 2^18 | 40 | 4 | 10.664 ms | 10.155 ms | 4.8% | - -Complete PCS commitments include LDEs, gathering, hashing, Merkle construction, -and root return. Medians of 15 warm samples per variant across three processes; -all commitment bytes matched for every shape/run. - -| Shape | BLAKE3-only | + LDE change | Reduction | -| --- | ---: | ---: | ---: | -| 2^20 × 8 | 9.011 ms | 8.938 ms | 0.8% | -| 2^18 × 40 | 10.666 ms | 10.478 ms | 1.8% | -| 2^18 × 128 | 36.102 ms | 35.332 ms | 2.1% | -| 2^18 × 129 | 43.792 ms | 43.195 ms | 1.4% | -| 2^16 × 925 | 91.071 ms | 69.281 ms | 23.9% | -| Mixed heights/widths | 68.184 ms | 51.253 ms | 24.8% | - -The mixed shape is the same five-matrix fixture as the BLAKE3 experiment. -The substantial gains occur in shapes affected by the removed NTT height -cutoff. Small percentage differences on unaffected/narrow shapes should not -be treated as established wins; one narrow LDE case was 1.6% slower. These -are synthetic component timings, not an Init/Mathlib or Stage 2 wall-time -measurement, and are not a general 2× NTT claim. - -## Validation and reproduction - -- All 71 CUDA-enabled library tests passed (6.46 seconds). -- CPU comparisons exercise blowups 1/2/4/8, three cosets, height-one inputs, - wide/narrow dispatch, fused-stage residues modulo three, and noncanonical - input representations around the field modulus. -- Raw output field words remain below the modulus. Mixed-height commitments, - resident/hybrid openings, hashing, and batch-proof tests pass. -- CUDA memcheck and initcheck report zero errors for the padding/coset/raw- - representation suite, including the case with no zero-padded tail. -- The complete 17,213-byte compatibility proof verifies and matches both - CPU and BLAKE3-only output exactly. SHA-256: - `25564a01d1d352b1ec2de56b019b641d24acc81083133e79274a86829b2a5dd5`. -- Clippy completed without warnings for the library and both new examples. - -```sh -MULTI_STARK_CUDA_ARCHS=120 RAYON_NUM_THREADS=8 cargo test --release --locked --features parallel,cuda --lib -- --test-threads=1 -RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_resident_lde_bench -RAYON_NUM_THREADS=8 cargo run --release --locked --features parallel,cuda --example cuda_commit_bench -``` - -Build the identical benchmark source against `ff3237c` and this commit, -preserve both binaries, and alternate runs. `MULTI_STARK_CUDA_BENCH_ITERATIONS` -controls the warm sample count. The resident benchmark uses the production -resident-LDE interface; the older `cuda_dft_bench` exercises a different, -transfer-inclusive host-return interface and does not isolate these changes. diff --git a/examples/cuda_resident_lde_bench.rs b/examples/cuda_resident_lde_bench.rs index 1bbc3cd..6e8ec0e 100644 --- a/examples/cuda_resident_lde_bench.rs +++ b/examples/cuda_resident_lde_bench.rs @@ -1,6 +1,5 @@ //! Resident LDE timings including upload, excluding host input construction. -//! Run the identical benchmark against both revisions. Correctness is covered -//! by resident_coset_lde_matches_cpu_storage and proof_compatibility. +//! Correctness is covered by the `cuda::sppark` tests and `proof_compatibility`. #[cfg(not(feature = "cuda"))] fn main() { diff --git a/src/batch.rs b/src/batch.rs index 88504ec..54b6da7 100644 --- a/src/batch.rs +++ b/src/batch.rs @@ -544,8 +544,6 @@ where Retention::Regenerate => None, }); barrier.claims.push(claims); - #[cfg(feature = "cuda")] - crate::cuda::metrics::emit_snapshot(); }); barrier } @@ -610,8 +608,6 @@ where _ => unreachable!("a shard is either retained or rebuilt"), }; proofs.push(self.prove_batch_shard(key, stage_1, &claims, &preamble, shard)); - #[cfg(feature = "cuda")] - crate::cuda::metrics::emit_snapshot(); }); BatchProof { preamble, proofs } } diff --git a/src/cuda/metrics.rs b/src/cuda/metrics.rs deleted file mode 100644 index 24c98f3..0000000 --- a/src/cuda/metrics.rs +++ /dev/null @@ -1,36 +0,0 @@ -//! CPU-only snapshots of counters accumulated at CUDA operation boundaries. - -pub(crate) fn emit_snapshot() { - if !tracing::event_enabled!(target: "prover_metrics", tracing::Level::INFO) { - return; - } - const SHAPES: usize = 33 * 4; - const WORDS: usize = 16 + SHAPES; - let mut values = vec![0u64; 64 * WORDS]; - unsafe { - multi_stark_cuda_metrics_snapshot(values.as_mut_ptr(), values.len()); - } - for (device, c) in values.as_chunks::().0.iter().enumerate() { - if c.iter().all(|&v| v == 0) { - continue; - } - tracing::info!(target: "prover_metrics", metric = "cuda_device_snapshot", device, - scope = "process_device_cumulative", upload_calls = c[0], upload_requested_bytes = c[1], - upload_chunks = c[2], upload_failures = c[3], upload_host_ns = c[4], - coset_hits = c[5], coset_misses = c[6], coset_uploaded_bytes = c[7], constant_bytes = c[8], - last_driver_free_bytes = c[9], total_bytes = c[10], memory_samples = c[11]); - let shapes = &c[16..]; - for (shape, &count) in shapes.iter().enumerate() { - if count == 0 { - continue; - } - tracing::info!(target: "prover_metrics", metric = "ntt_snapshot", device, - scope = "process_device_cumulative", log_height = shape / 4, - width_bucket = ["1", "2", "3-7", "8+"][shape % 4], transforms = count); - } - } -} - -unsafe extern "C" { - fn multi_stark_cuda_metrics_snapshot(output: *mut u64, count: usize); -} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index ab42f5f..4d12da4 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -5,11 +5,10 @@ //! lookup traces, quotient evaluations, and FRI codewords on the selected GPU. //! All public protocol types and serialized proofs remain unchanged. -pub(crate) mod metrics; pub(crate) mod mmcs; #[doc(hidden)] pub mod pcs; -pub mod sppark; +pub(crate) mod sppark; pub(crate) mod witness; use core::ffi::{CStr, c_char, c_void}; @@ -145,7 +144,6 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "host", - backend = "sppark", device = self.device_id, height, width, @@ -201,8 +199,6 @@ impl CudaDft { } } -/// Bit-reversed coset-LDE storage owned by a CUDA device allocation. -#[doc(hidden)] /// A borrowed output tile on the selected CUDA device. The writer must finish /// using it before returning; rows wrap at the source's padded height. pub struct DeviceTraceView<'a> { @@ -267,37 +263,7 @@ unsafe extern "C" fn destroy_generated_trace(context: *mut c_void) { } impl CudaDft { - /// Materialize a bounded generated tile for diagnostics and reference checks. - pub fn generated_trace_rows( - &self, - source: Generator, - first: usize, - rows: usize, - ) -> RowMajorMatrix { - assert!( - rows > 0 && rows <= (1 << 16) + 1, - "generated tile exceeds the row limit" - ); - let width = source.width(); - let mut output = - vec![Goldilocks::ZERO; rows.checked_mul(width).expect("trace tile overflow")]; - let mut source = source; - let status = unsafe { - multi_stark_cuda_generate_trace_rows( - self.device_id, - (&mut source as *mut Generator).cast(), - write_generated_trace, - first, - rows, - width, - output.as_mut_ptr().cast(), - ) - }; - check_cuda(status, "generated trace tile download"); - RowMajorMatrix::new(output, width) - } - - pub fn generate_coset_lde( + pub(crate) fn generate_coset_lde( &self, generator: Generator, added_bits: usize, @@ -308,7 +274,6 @@ impl CudaDft { let _span = tracing::info_span!( "cuda/lde", kind = "generated", - backend = "sppark", device = self.device_id, height, width, @@ -348,15 +313,6 @@ impl CudaDft { } unsafe extern "C" { - fn multi_stark_cuda_generate_trace_rows( - device: i32, - context: *mut c_void, - writer: unsafe extern "C" fn(*mut c_void, i32, *mut u64, usize, usize) -> i32, - first: usize, - rows: usize, - width: usize, - output: *mut u64, - ) -> i32; fn multi_stark_cuda_coset_lde_generate( device: i32, output: *mut *mut c_void, @@ -372,6 +328,8 @@ unsafe extern "C" { fn multi_stark_cuda_lde_generator_context(handle: *const c_void) -> *mut c_void; } +/// Bit-reversed coset-LDE storage owned by a CUDA device allocation. +#[doc(hidden)] pub struct CudaLde { device_id: i32, handle: NonNull, @@ -746,6 +704,14 @@ fn encode_quotient_nodes( (nodes, slots, count as usize) } +/// Conservative device-memory requirement for one fused quotient job. +/// +/// The graph evaluator reuses slots as soon as their final consumer has run, +/// so `graph.nodes.len()` can be orders of magnitude larger than the live +/// device scratch. Keep this estimate beside the encoder so admission and the +/// kernel use the same liveness calculation. The scratch term assumes the +/// global-memory path; devices able to fit the slots in shared memory need +/// less than this bound. pub(crate) fn quotient_lde_memory_upper_bound( graph: &ConstraintGraph, public_count: usize, @@ -1156,7 +1122,6 @@ pub(crate) fn quotient_lde_mixed( ) -> CudaLde { let _span = tracing::info_span!( "cuda/quotient_lde", - backend = "sppark", device = dft.device_id, quotient_size, quotient_degree, @@ -1709,7 +1674,6 @@ pub(crate) fn lookup_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "direct", - backend = "sppark", device = dft.device_id, height, num_lookups, @@ -1800,7 +1764,6 @@ pub(crate) fn lookup_lde_resident_partitioned( let _span = tracing::info_span!( "cuda/lookup_lde", path = "partitioned", - backend = "sppark", device = dft.device_id, height, num_lookups, @@ -1997,7 +1960,6 @@ pub(crate) fn lookup_graph_lde_resident( let _span = tracing::info_span!( "cuda/lookup_lde", path = "graph", - backend = "sppark", device = dft.device_id, height, num_lookups = graph.lookups.len(), @@ -2228,6 +2190,15 @@ pub(crate) fn device_memory_info(device_id: i32) -> (usize, usize) { (free_bytes, total_bytes) } +/// Device bytes the prover keeps free for its later stages: a quarter of the +/// device unless `MULTI_STARK_CUDA_MIN_FREE_BYTES` says otherwise. +pub(crate) fn minimum_free_bytes(total_bytes: usize) -> usize { + std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(total_bytes / 4) +} + pub(crate) fn memory_diagnostics_enabled() -> bool { std::env::var_os("MULTI_STARK_CUDA_MEMORY_LOG").is_some() } diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index e61c7dd..81767dd 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -277,19 +277,15 @@ pub struct CudaTwoAdicFriPcs { /// shift times `g_k^{rev_k(i)}`, and `rev_{k+1}(i) = 2 rev_k(i)` for /// `i < 2^k`, so every smaller coset is a prefix of a larger one: one vector /// per process serves every height, and grows only when a larger height is -/// opened. Every shard of a proof opens on the same coset; before caching, -/// building it cost 0.32 s per shard at 2^26 rows. +/// opened. Every shard of a proof opens on the same coset. fn bit_reversed_coset( log_height: usize, ) -> std::sync::Arc> { static CACHE: std::sync::Mutex>>> = std::sync::Mutex::new(None); if let Some(coset) = cache_at_least(&CACHE, 1 << log_height) { - tracing::info!(target: "prover_metrics", metric = "host_coset", action = "hit", - log_height, cached_bytes = coset.len() * 8); return coset; } - tracing::info!(target: "prover_metrics", metric = "host_coset", action = "build", log_height); let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); let generator: Goldilocks = TwoAdicField::two_adic_generator(log_height); let shift = ::GENERATOR; @@ -636,10 +632,7 @@ where let max_transform_workspace = transform_workspaces.iter().copied().max().unwrap_or(0); let (initial_free, total_bytes) = crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); - let minimum_free = std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(total_bytes / 4); + let minimum_free = crate::cuda::minimum_free_bytes(total_bytes); let source_bytes = source_cells.saturating_mul(size_of::()); let lde_bytes = source_bytes .checked_shl(u32::try_from(log_blowup).expect("LDE blowup exceeds u32")) diff --git a/src/cuda/sppark.rs b/src/cuda/sppark.rs index 116a6dd..8d49b4e 100644 --- a/src/cuda/sppark.rs +++ b/src/cuda/sppark.rs @@ -74,24 +74,23 @@ pub(crate) struct Planner { impl Planner { pub(crate) fn new(device: i32) -> Self { - fn setting(name: &str) -> Option { - std::env::var(name).ok().map(|value| { - value - .parse() - .unwrap_or_else(|_| panic!("{name} must be a non-negative byte count")) + let panel_bytes = std::env::var("MULTI_STARK_SPPARK_PANEL_BYTES") + .ok() + .map(|value| { + value.parse().unwrap_or_else(|_| { + panic!("MULTI_STARK_SPPARK_PANEL_BYTES must be a non-negative byte count") + }) }) - } + .filter(|&n| n != 0) + .unwrap_or(4usize << 30); + // Launch groups sized to the L2 cache keep a group's columns resident + // across the stages of a batched transform. let mut l2_bytes = 0; check_cuda( unsafe { multi_stark_sppark_l2_bytes(device, &mut l2_bytes) }, "NTT device properties", ); - Self::with_budgets( - setting("MULTI_STARK_SPPARK_PANEL_BYTES") - .filter(|&n| n != 0) - .unwrap_or(4usize << 30), - setting("MULTI_STARK_SPPARK_BATCH_BYTES").unwrap_or(l2_bytes), - ) + Self::with_budgets(panel_bytes, l2_bytes) } pub(crate) fn with_budgets(panel_bytes: usize, batch_bytes: usize) -> Self { @@ -185,18 +184,20 @@ impl Planner { /// `NTT::InputOutputOrder`: whether the input and the output are in natural /// (`N`) or bit-reversed (`R`) order. +#[cfg(test)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(i32)] -pub enum Order { +enum Order { NN = 0, NR = 1, RN = 2, RR = 3, } +#[cfg(test)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(i32)] -pub enum Direction { +enum Direction { Forward = 0, Inverse = 1, } @@ -206,15 +207,7 @@ unsafe extern "C" { fn multi_stark_sppark_l2_bytes(device: c_int, bytes: *mut usize) -> c_int; #[cfg(test)] fn multi_stark_sppark_borrowed_round_trip(device: c_int, values: *mut u64, lg: u32) -> c_int; - fn multi_stark_sppark_transforms_run() -> u64; - fn multi_stark_sppark_ntt_device( - device: c_int, - d_inout: *mut u64, - lg: u32, - order: c_int, - direction: c_int, - coset: c_int, - ) -> c_int; + #[cfg(test)] fn multi_stark_sppark_ntt_host( device: c_int, inout: *mut u64, @@ -223,6 +216,7 @@ unsafe extern "C" { direction: c_int, coset: c_int, ) -> c_int; + #[cfg(test)] fn multi_stark_sppark_ntt_batch_host( device: c_int, inout: *mut u64, @@ -235,16 +229,12 @@ unsafe extern "C" { ) -> c_int; } -/// Number of transform operations launched, including identity shapes. -pub fn transforms_run() -> u64 { - unsafe { multi_stark_sppark_transforms_run() } -} - /// The largest log domain size the compiled upstream parameters support. -pub fn max_log_domain() -> usize { +pub(crate) fn max_log_domain() -> usize { usize::try_from(unsafe { multi_stark_sppark_max_lg_domain() }).expect("domain limit") } +#[cfg(test)] fn log_len(values: usize) -> u32 { assert!( values.is_power_of_two() && values > 1, @@ -257,7 +247,8 @@ fn log_len(values: usize) -> u32 { /// downloaded within the call. `coset` selects the coset by the field /// generator. An error is the CUDA status the adapter returned; `device` /// must be the CUDA ordinal of a device upstream supports. -pub fn try_ntt_host( +#[cfg(test)] +fn try_ntt_host( device: i32, values: &mut [Goldilocks], order: Order, @@ -279,17 +270,19 @@ pub fn try_ntt_host( } /// `count` vectors of `2^lg` elements laid out `stride` elements apart. +#[cfg(test)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Batch { - pub lg: u32, - pub count: u32, - pub stride: usize, +struct Batch { + lg: u32, + count: u32, + stride: usize, } /// Transforms the vectors of `batch` in `values` (`count * stride` long) in /// one batched launch sequence, uploaded, transformed and downloaded within /// the call. -pub fn try_ntt_batch_host( +#[cfg(test)] +fn try_ntt_batch_host( device: i32, values: &mut [Goldilocks], batch: Batch, @@ -321,7 +314,8 @@ pub fn try_ntt_batch_host( } /// [`try_ntt_host`], panicking on a CUDA status like the other backends. -pub fn ntt_host( +#[cfg(test)] +fn ntt_host( device: i32, values: &mut [Goldilocks], order: Order, @@ -333,37 +327,9 @@ pub fn ntt_host( } } -/// Transforms `2^lg` field elements at `d_inout` in place on the calling -/// thread's stream. -/// -/// # Safety -/// -/// `d_inout` must be a device allocation of `2^lg` 64-bit words on -/// `device`, and no other stream may access it until work enqueued after -/// this call on the calling thread's stream has completed. -pub unsafe fn ntt_device( - device: i32, - d_inout: *mut u64, - lg: u32, - order: Order, - direction: Direction, - coset: bool, -) { - let status = unsafe { - multi_stark_sppark_ntt_device( - device, - d_inout, - lg, - order as c_int, - direction as c_int, - c_int::from(coset), - ) - }; - check_cuda(status, "sppark device transform"); -} - /// The raw stored words of `values`; Goldilocks is `repr(transparent)`. -pub fn raw_words(values: &[Goldilocks]) -> &[u64] { +#[cfg(test)] +fn raw_words(values: &[Goldilocks]) -> &[u64] { // SAFETY: Goldilocks is a transparent wrapper over u64. unsafe { core::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } } diff --git a/src/cuda/witness.rs b/src/cuda/witness.rs index e3dc72b..badf474 100644 --- a/src/cuda/witness.rs +++ b/src/cuda/witness.rs @@ -15,7 +15,6 @@ use crate::witness::TraceSource; pub(crate) fn record_lde_spill(device: i32, bytes: usize) { tracing::debug!(device, bytes, "spilled active LDE"); - tracing::info!(target: "prover_metrics", metric = "lde_spill", device, bytes); } fn spill_lde(lde: &CudaLde) -> RowMajorMatrix { @@ -26,13 +25,6 @@ fn spill_lde(lde: &CudaLde) -> RowMajorMatrix { matrix } -fn reserve_bytes(total: usize) -> usize { - std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") - .ok() - .and_then(|n| n.parse().ok()) - .unwrap_or(total / 4) -} - pub(crate) fn commit( pcs: &CudaPcs, evaluations: Vec<(Domain, TraceSource)>, @@ -73,7 +65,7 @@ pub(crate) fn commit( } let max_height = dimensions.iter().map(|d| d.height).max().unwrap(); let (_, total) = device_memory_info(device); - let reserve = reserve_bytes(total).saturating_add(128 << 20); + let reserve = super::minimum_free_bytes(total).saturating_add(128 << 20); let mut resident: Vec> = Vec::with_capacity(evaluations.len()); let mut host: Vec>> = Vec::with_capacity(evaluations.len()); let mut retained = Vec::with_capacity(evaluations.len()); @@ -135,14 +127,8 @@ pub(crate) fn commit( // Both kinds retain a bounded recovery source. Raw device rows need // not coexist with later lookup/quotient workspace. unsafe { lde.release_trace() }; - let spilled = if std::env::var("MULTI_STARK_CUDA_TRACE_FORCE_SPILL").is_ok_and(|v| v == "1") - { - Some(spill_lde(&lde)) - } else { - None - }; resident.push(Some(lde)); - host.push(spilled); + host.push(None); retained.push(trace); drop(source_span); } diff --git a/src/types.rs b/src/types.rs index ec1df5e..10ad8bc 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1007,9 +1007,6 @@ impl StarkGenericConfig for GoldilocksBlake3Config { let mut graph_path = admits_graph(); let (mut temporary_bytes, mut target) = target_for(graph_path); if target > total_device_bytes { - tracing::info!(target: "prover_metrics", metric = "lookup_admission", - job = index, graph_path, admitted = false, reason = "device_capacity", - output_bytes, temporary_bytes, target_bytes = target); return None; } let mut free_bytes = self.pcs.mmcs.ensure_device_headroom( @@ -1039,9 +1036,6 @@ impl StarkGenericConfig for GoldilocksBlake3Config { ); } } - tracing::info!(target: "prover_metrics", metric = "lookup_admission", - job = index, graph_path, admitted = free_bytes >= target, - output_bytes, temporary_bytes, target_bytes = target, free_bytes); if free_bytes < target { return None; } diff --git a/src/witness.rs b/src/witness.rs index 98a0183..3f29286 100644 --- a/src/witness.rs +++ b/src/witness.rs @@ -9,7 +9,6 @@ use p3_matrix::{Matrix, dense::RowMajorMatrix}; pub trait TraceGenerator: Send + Sync { fn height(&self) -> usize; fn width(&self) -> usize; - fn host_bytes(&self) -> usize; /// Fill contiguous rows, wrapping at the padded height for lookup halos. fn write_rows(&self, first: usize, output: &mut [F]);