From 94dafc73e8baec06e389adc4072d4538ca676439 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 12 Aug 2026 13:28:32 -0300 Subject: [PATCH] feat(prover): spill the aux LDE under StorageMode::Disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under `StorageMode::Disk` the aux trace and the aux Merkle tree are both spilled, but the aux LDE itself — `lde_size × aux_cols` ext3 elements at 24 B each, the largest of the three — had no spill path and stayed heap-resident from the aux commit through rounds 2-4. Carry it as a `Table` (the crate's existing mmap-backed row-major container, already used for spilled trace tables) instead of a bare `Vec`, and spill it right after the aux commit. It is write-once at that point and read-only afterwards, so mmap-backing it frees the heap buffer for the whole rounds 2-4 window and lets the OS evict the pages under memory pressure. Behaviour-neutral off `disk-spill` and off Disk mode: the `Table` arm without an mmap backing is the same buffer and the same indexing. Measured with prover/tests/calibration.rs (fib_iterative_372k, 5 runs per arm): peak heap under FORCE_DISK_SPILL=1 goes from 2.615-2.825 GB to 2.290-2.378 GB — non-overlapping ranges, about -403 MB / -14.8%. Ram mode is unchanged. --- crypto/stark/src/gpu_lde.rs | 2 +- crypto/stark/src/prover.rs | 50 +++++++++++++++++++++++++------------ crypto/stark/src/table.rs | 19 ++++++++++++++ crypto/stark/src/trace.rs | 25 +++++++++++-------- 4 files changed, 69 insertions(+), 27 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 4aa756b25..617316d45 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1553,7 +1553,7 @@ where // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. let aux_data: Vec> = - if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() { + if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.row_major_data().is_empty() { Vec::new() } else { let Some(h) = lde_trace.gpu_aux() else { diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 232e1faaf..c5acffe2e 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -275,8 +275,10 @@ where struct Lde { /// Row-major main LDE buffer + its column count. main: (Vec>, usize), - /// Row-major aux LDE buffer + its column count (`(vec![], 0)` if no aux). - aux: (Vec>, usize), + /// Row-major aux LDE, as a `Table` whose `width` is the column count (an + /// empty, zero-width table if no aux). `Table` rather than a bare `Vec` so + /// the buffer can be disk-spilled for the whole rounds 2-4 window. + aux: Table, /// Device-side main LDE buffer, populated only when the R1 GPU fused /// pipeline ran for this table. Kept so R2/R3/R4 GPU paths can read /// the LDE without re-H2D. @@ -302,7 +304,9 @@ where blowup_factor: usize, ) -> Round1 { let (main_data, num_main_cols) = lde.main; - let (aux_data, num_aux_cols) = lde.aux; + let aux_data = lde.aux; + #[cfg(feature = "cuda")] + let num_aux_cols = aux_data.width; // Stage-3 device-only detection, inferred from the ACTUAL buffer state // (not the gate's intent): a table whose round-1 D2H was skipped has an @@ -322,8 +326,8 @@ where #[cfg(feature = "cuda")] let main_empty = num_main_cols > 0 && main_data.is_empty(); #[cfg(feature = "cuda")] - let host_trace_empty = - main_empty || (num_aux_cols > 0 && aux_data.is_empty() && lde.gpu_aux.is_some()); + let host_trace_empty = main_empty + || (num_aux_cols > 0 && aux_data.row_major_data().is_empty() && lde.gpu_aux.is_some()); #[cfg(feature = "cuda")] let device_num_rows = lde .gpu_main @@ -336,7 +340,6 @@ where main_data, num_main_cols, aux_data, - num_aux_cols, step_size, blowup_factor, ); @@ -1407,9 +1410,9 @@ pub trait IsStarkProver< } } } - (aux_data, num_aux_cols) + Table::from_row_major(aux_data, num_aux_cols) } else { - (Vec::new(), 0) + Table::new(Vec::new(), 0) }; Ok(commitment.build_round1( @@ -3323,11 +3326,11 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] type AuxResult = ( Option>, - (Vec>, usize), + Table, Option, ); #[cfg(not(feature = "cuda"))] - type AuxResult = (Option>, (Vec>, usize)); + type AuxResult = (Option>, Table); // R1 aux commit and rounds 2 to 4 share the peak working set: the main // and aux LDEs are co-resident, plus the composition and Merkle // transients (in the scratch factor). The aux width comes from the AIR @@ -3484,7 +3487,7 @@ pub trait IsStarkProver< let root = tree.root; return Ok(( Some(TableCommit::plain(tree, root)), - (aux_data, num_cols), + Table::from_row_major(aux_data, num_cols), Some(handle), )); } @@ -3580,7 +3583,7 @@ pub trait IsStarkProver< crate::instruments::accum_r1_aux(aux_lde_dur, Duration::ZERO); return Ok(( Some(TableCommit::plain(tree, root)), - (aux_data, num_cols), + Table::from_row_major(aux_data, num_cols), Some(handle), )); } @@ -3627,15 +3630,30 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] crate::instruments::accum_r1_aux(aux_lde_dur, t_sub.elapsed()); + // Spill the aux LDE itself (`lde_size × aux_cols` ext3 + // elements, 24 B each — the dominant survivor of this + // stage). It is write-once here and read-only from + // rounds 2-4, so mmap-backing it frees the heap buffer + // for the whole window and lets the OS evict the pages + // under pressure. + #[allow(unused_mut)] + let mut aux_lde = Table::from_row_major(aux_data, total_cols); + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + aux_lde + .spill_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("aux LDE: {e}")))?; + } + #[cfg(feature = "cuda")] - return Ok((Some(commit), (aux_data, total_cols), None)); + return Ok((Some(commit), aux_lde, None)); #[cfg(not(feature = "cuda"))] - Ok((Some(commit), (aux_data, total_cols))) + Ok((Some(commit), aux_lde)) } else { #[cfg(feature = "cuda")] - return Ok((None, (Vec::new(), 0), None)); + return Ok((None, Table::new(Vec::new(), 0), None)); #[cfg(not(feature = "cuda"))] - Ok((None, (Vec::new(), 0))) + Ok((None, Table::new(Vec::new(), 0))) } })()?; // Tuple shape is cfg-gated; `.0` is the optional TableCommit in diff --git a/crypto/stark/src/table.rs b/crypto/stark/src/table.rs index 238c4fcfb..ec3a67ff7 100644 --- a/crypto/stark/src/table.rs +++ b/crypto/stark/src/table.rs @@ -313,6 +313,23 @@ impl Table { } } + /// Wrap a buffer the caller already produced in row-major order, skipping + /// `new`'s 2-D revalidation — that `debug_assert` clones the whole buffer, + /// which is prohibitive for LDE-sized data. + pub(crate) fn from_row_major(data: Vec>, width: usize) -> Self { + if width == 0 { + return Self::new(Vec::new(), 0); + } + let height = data.len() / width; + Self { + data, + width, + height, + #[cfg(feature = "disk-spill")] + mmap_backing: None, + } + } + /// Creates a Table instance from a vector of the intended columns. pub fn from_columns(columns: Vec>>) -> Self { if columns.is_empty() { @@ -336,6 +353,7 @@ impl Table { } /// Given a row index, returns a reference to that row as a slice of field elements. + #[inline] pub fn get_row(&self, row_idx: usize) -> &[FieldElement] { #[cfg(feature = "disk-spill")] if let Some(ref backing) = self.mmap_backing { @@ -424,6 +442,7 @@ impl Table { } /// Given row and column indexes, returns the stored field element in that position of the table. + #[inline] pub fn get(&self, row: usize, col: usize) -> &FieldElement { #[cfg(feature = "disk-spill")] if let Some(ref backing) = self.mmap_backing { diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index ccf35cca5..e245a852e 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -321,7 +321,9 @@ where /// Row-major main-trace buffer of length `num_rows * num_main_cols`. pub(crate) main_data: Vec>, /// Row-major auxiliary-trace buffer of length `num_rows * num_aux_cols`. - pub(crate) aux_data: Vec>, + /// A `Table` rather than a bare `Vec` so it can carry a disk-spilled + /// (mmap-backed) buffer — see the aux-LDE spill in the prover's aux stage. + pub(crate) aux_data: Table, pub(crate) num_main_cols: usize, pub(crate) num_aux_cols: usize, pub(crate) num_rows: usize, @@ -467,7 +469,7 @@ where Self { main_data, - aux_data, + aux_data: Table::from_row_major(aux_data, num_aux_cols), num_main_cols, num_aux_cols, num_rows, @@ -483,21 +485,24 @@ where /// Build an LDETraceTable directly from row-major flat buffers. Skips the /// O(N·M) col→row transpose that `from_columns` pays — the caller produces /// the buffers row-major already (e.g. via `coset_lde_full_expand_row_major`). + /// + /// `aux_data` arrives as a `Table` so an already disk-spilled aux LDE keeps + /// its mmap backing instead of being pulled back onto the heap; its `width` + /// is the aux column count. pub fn from_row_major( main_data: Vec>, num_main_cols: usize, - aux_data: Vec>, - num_aux_cols: usize, + aux_data: Table, trace_step_size: usize, blowup_factor: usize, ) -> Self { let lde_step_size = trace_step_size * blowup_factor; + let num_aux_cols = aux_data.width; let num_rows = if num_main_cols > 0 { debug_assert_eq!(main_data.len() % num_main_cols, 0); main_data.len() / num_main_cols } else if num_aux_cols > 0 { - debug_assert_eq!(aux_data.len() % num_aux_cols, 0); - aux_data.len() / num_aux_cols + aux_data.height } else { 0 }; @@ -566,7 +571,7 @@ where self.main_data = main_data; } if !aux_data.is_empty() { - self.aux_data = aux_data; + self.aux_data = Table::from_row_major(aux_data, self.num_aux_cols); } self.host_trace_empty = false; } @@ -636,7 +641,7 @@ where /// Get a single aux-trace element by (row, col). #[inline] pub fn get_aux(&self, row: usize, col: usize) -> &FieldElement { - &self.aux_data[row * self.num_aux_cols + col] + self.aux_data.get(row, col) } /// Borrow a full main-trace row as a contiguous slice (row-major buffer). @@ -648,7 +653,7 @@ where /// Borrow a full aux-trace row as a contiguous slice (row-major buffer). #[inline] pub fn aux_row(&self, row: usize) -> &[FieldElement] { - &self.aux_data[row * self.num_aux_cols..(row + 1) * self.num_aux_cols] + self.aux_data.get_row(row) } /// Gather a full main-trace row into an owned Vec. @@ -878,7 +883,7 @@ where // buffer-level check as the main arm: mixed states are valid here. #[cfg(feature = "cuda")] assert!( - lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty(), + lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.row_major_data().is_empty(), "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)" ); let inv_denoms_v =