diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2a9547153..e4e85792e 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -26,6 +26,8 @@ use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; use math::field::goldilocks::GoldilocksField; use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; use math::traits::AsBytes; +#[cfg(feature = "parallel")] +use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut}; use crate::config::{Commitment, FriLayerMerkleTreeBackend}; use crate::domain::Domain; @@ -1512,12 +1514,37 @@ where return false; } let (m, lde) = (h.m, h.lde_size); + // Short download: degrade like the sibling paths + // (`download_main_lde_row_major`, `materialize_aux_trace_host`) + // rather than panic on the slab slicing below. + if slabs.len() != m * lde * 3 { + return false; + } + // Parallel de-interleaved slabs → row-major interleaved: each row + // chunk gathers from the source slabs independently. let mut interleaved = vec![0u64; m * lde * 3]; - for c in 0..m { - for k in 0..3 { - let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; - for r in 0..lde { - interleaved[(r * m + c) * 3 + k] = slab[r]; + if m > 0 { + #[cfg(feature = "parallel")] + { + interleaved + .par_chunks_exact_mut(m * 3) + .enumerate() + .for_each(|(r, dst)| { + for (c, dst_col) in dst.chunks_exact_mut(3).enumerate() { + for (k, d) in dst_col.iter_mut().enumerate() { + *d = slabs[(c * 3 + k) * lde + r]; + } + } + }); + } + #[cfg(not(feature = "parallel"))] + { + for (r, dst) in interleaved.chunks_exact_mut(m * 3).enumerate() { + for (c, dst_col) in dst.chunks_exact_mut(3).enumerate() { + for (k, d) in dst_col.iter_mut().enumerate() { + *d = slabs[(c * 3 + k) * lde + r]; + } + } } } } @@ -1525,6 +1552,10 @@ where // is [u64; 3]. unsafe { let mut v = std::mem::ManuallyDrop::new(interleaved); + debug_assert!( + v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3), + "interleaved len/capacity must be a multiple of 3 for Fp3 reinterpret" + ); Vec::from_raw_parts( v.as_mut_ptr() as *mut FieldElement, v.len() / 3, @@ -1557,10 +1588,28 @@ where if col_major.len() != m * lde { return None; } + // Parallel col-major → row-major transpose: each row chunk gathers from + // the source columns independently. let mut row_major = vec![0u64; m * lde]; - for c in 0..m { - for r in 0..lde { - row_major[r * m + c] = col_major[c * lde + r]; + if m > 0 { + #[cfg(feature = "parallel")] + { + row_major + .par_chunks_exact_mut(m) + .enumerate() + .for_each(|(r, dst)| { + for (c, d) in dst.iter_mut().enumerate() { + *d = col_major[c * lde + r]; + } + }); + } + #[cfg(not(feature = "parallel"))] + { + for (r, dst) in row_major.chunks_exact_mut(m).enumerate() { + for (c, d) in dst.iter_mut().enumerate() { + *d = col_major[c * lde + r]; + } + } } } // SAFETY: F == Goldilocks (gated above); FieldElement is diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d5cca05e1..b609786e7 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -3478,6 +3478,11 @@ pub trait IsStarkProver< // the main LDE if this table was device-only — and // continue fully host-backed on the arms below. let mut recovered = crate::gpu_lde::materialize_aux_trace_host(*trace); + // Once the aux download lands, the host aux trace is + // populated: a later failure is the main-LDE + // download's, and the error has to name that step + // instead of claiming an empty aux trace. + let aux_recovered = recovered; if recovered && device_only { let mut cell = main_lde_cells[idx].lock().unwrap(); if let Some((data, _)) = cell.as_mut() @@ -3506,7 +3511,14 @@ pub trait IsStarkProver< } if !recovered { return Err(ProvingError::Fft( - "resident aux LDE failed; host aux trace is empty".to_string(), + if aux_recovered { + "resident aux LDE declined; the aux trace was recovered \ + but the main-LDE download failed" + } else { + "resident aux LDE declined and the aux-trace download \ + recovery failed" + } + .to_string(), )); } eprintln!(