fix(gpu): harden the downgrade downloads, parallelize the recovery transposes

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,19 +1514,48 @@ 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];
+                            }
+                        }
                     }
                 }
             }
             // SAFETY: E == Ext3 per the tower check; FieldElement<Ext3> backing
             // 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<E>,
                     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<Gl> is

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!(