fix(gpu): survive transient VRAM pressure — recover resident-table declines, close an R2 corruption race

crypto/stark/src/gpu_lde.rs

@@ -54,6 +54,24 @@ fn gpu_lde_threshold() -> usize {
     })
 }
 
+/// Serialize the device R2 window (constraint eval + decompose) across
+/// tables. Concurrent R2 windows under VRAM pressure can transiently corrupt
+/// a whole H buffer (root mechanism unidentified; reruns on the same resident
+/// inputs come out correct), yielding a proof that fails verification.
+/// Serializing only this window eliminates it at negligible cost — the
+/// windows rarely overlap. `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock
+/// (e.g. to bisect or once the underlying race is fixed).
+pub(crate) fn r2_serialize_guard() -> Option<std::sync::MutexGuard<'static, ()>> {
+    static ENABLED: OnceLock<bool> = OnceLock::new();
+    static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+    if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0"))
+    {
+        Some(LOCK.lock().unwrap())
+    } else {
+        None
+    }
+}
+
 /// Incremented by the `try_expand_*` functions per base-field column handed to
 /// the GPU dispatch (an ext3 column counts as 3, one per base component),
 /// before the GPU call. A failed call returns without decrementing it, so it
@@ -1413,6 +1431,224 @@ pub fn gpu_fri_calls() -> u64 {
 /// are counted here, so a single failed dispatch does not necessarily lower
 /// the total; R3's fallbacks are CPU-only, so a failure there does.
 pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0);
+/// Times a device-only table had to be downgraded back to a host trace
+/// because a downstream device path missed at runtime (see
+/// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate
+/// admitted a table some dispatch later declined — correct but slower, and
+/// worth mirroring the missing condition into the gate.
+pub(crate) static GPU_DEVICE_ONLY_DOWNGRADES: AtomicU64 = AtomicU64::new(0);
+pub fn gpu_device_only_downgrades() -> u64 {
+    GPU_DEVICE_ONLY_DOWNGRADES.load(Ordering::Relaxed)
+}
+
+/// Recover a device-only table for the host path: download the resident main
+/// and aux LDEs from their device handles into the host buffers and clear the
+/// device-only flag. A side whose host buffer is already populated (a mixed
+/// state: one commit fell back to CPU while the other stayed device-only) is
+/// kept as is — only the missing side is downloaded. The class-level safety
+/// net under the device-only gate — a static predicate can never mirror every
+/// reason a dynamic dispatch might decline (kernel eligibility, transient
+/// errors, shapes a new workload brings), so any miss lands here and degrades
+/// to a slower-but-correct CPU round instead of a hard abort. Returns false
+/// (→ the caller's abort) only when a missing side has no handle or a
+/// download fails.
+pub(crate) fn materialize_lde_trace_host<F, E>(
+    lde_trace: &mut crate::trace::LDETraceTable<F, E>,
+) -> bool
+where
+    F: IsField + IsSubFieldOf<E> + 'static,
+    E: IsField + 'static,
+{
+    if !lde_trace.host_trace_empty() {
+        return true;
+    }
+    if !is_goldilocks_ext3_tower::<F, E>() {
+        return false;
+    }
+    let Some(stream) = lde_trace.bound_stream() else {
+        return false;
+    };
+
+    // Main: column-major device buf -> row-major host Vec. An empty Vec tells
+    // `set_host_data` to keep the buffer that is already there.
+    let main_data: Vec<FieldElement<F>> =
+        if lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty() {
+            Vec::new()
+        } else {
+            let Some(h) = lde_trace.gpu_main() else {
+                return false;
+            };
+            if h.m != lde_trace.num_main_cols() || h.lde_size != lde_trace.num_rows() {
+                return false;
+            }
+            let Some(data) = download_main_lde_row_major::<F>(h, &stream) else {
+                return false;
+            };
+            data
+        };
+
+    // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec.
+    let aux_data: Vec<FieldElement<E>> =
+        if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() {
+            Vec::new()
+        } else {
+            let Some(h) = lde_trace.gpu_aux() else {
+                return false;
+            };
+            if h.m != lde_trace.num_aux_cols() || h.lde_size != lde_trace.num_rows() {
+                return false;
+            }
+            if h.wait_ready_on(&stream).is_err() {
+                return false;
+            }
+            let Ok(slabs) = stream.clone_dtoh(h.buf.as_ref()) else {
+                return false;
+            };
+            if stream.synchronize().is_err() {
+                return false;
+            }
+            let (m, lde) = (h.m, h.lde_size);
+            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];
+                    }
+                }
+            }
+            // SAFETY: E == Ext3 per the tower check; FieldElement<Ext3> backing
+            // is [u64; 3].
+            unsafe {
+                let mut v = std::mem::ManuallyDrop::new(interleaved);
+                Vec::from_raw_parts(
+                    v.as_mut_ptr() as *mut FieldElement<E>,
+                    v.len() / 3,
+                    v.capacity() / 3,
+                )
+            }
+        };
+
+    lde_trace.set_host_data(main_data, aux_data);
+    GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed);
+    true
+}
+
+/// Download a resident main LDE (column-major device buf) into the row-major
+/// host Vec the CPU rounds read. Shared by the R1 and R2 downgrade paths.
+pub(crate) fn download_main_lde_row_major<F>(
+    h: &math_cuda::lde::GpuLdeBase,
+    stream: &std::sync::Arc<math_cuda::CudaStream>,
+) -> Option<Vec<FieldElement<F>>>
+where
+    F: IsField + 'static,
+{
+    if TypeId::of::<F>() != TypeId::of::<GoldilocksField>() {
+        return None;
+    }
+    h.wait_ready_on(stream).ok()?;
+    let col_major = stream.clone_dtoh(h.buf.as_ref()).ok()?;
+    stream.synchronize().ok()?;
+    let (m, lde) = (h.m, h.lde_size);
+    if col_major.len() != m * lde {
+        return None;
+    }
+    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];
+        }
+    }
+    // SAFETY: F == Goldilocks (gated above); FieldElement<Gl> is
+    // #[repr(transparent)] over u64.
+    Some(unsafe {
+        let mut v = std::mem::ManuallyDrop::new(row_major);
+        Vec::from_raw_parts(
+            v.as_mut_ptr() as *mut FieldElement<F>,
+            v.len(),
+            v.capacity(),
+        )
+    })
+}
+
+/// R1 counterpart of [`materialize_lde_trace_host`]: download the resident
+/// aux trace (already row-major ext3, matching the host layout) into the
+/// trace's aux table, so the aux commit continues on the host arms when the
+/// device aux LDE declines at runtime.
+pub(crate) fn materialize_aux_trace_host<F, E>(trace: &mut crate::trace::TraceTable<F, E>) -> bool
+where
+    F: IsField + IsSubFieldOf<E> + 'static,
+    E: IsField + 'static,
+{
+    if !is_goldilocks_ext3_tower::<F, E>() {
+        return false;
+    }
+    let (buf, rows, cols) = match trace.aux_resident.as_ref() {
+        Some(ra) => (ra.buf.clone(), ra.num_rows, ra.num_aux_cols),
+        None => return false,
+    };
+    let Ok(be) = math_cuda::device::backend() else {
+        return false;
+    };
+    let stream = be.next_stream();
+    let Ok(raw) = stream.clone_dtoh(buf.as_ref()) else {
+        return false;
+    };
+    if stream.synchronize().is_err() || raw.len() != rows * cols * 3 {
+        return false;
+    }
+    let data = u64_to_ext3_vec::<E>(&raw);
+    trace.aux_table = crate::table::Table::new(data, cols);
+    trace.num_aux_columns = cols;
+    // The declined device LDE attempt can leave kernels enqueued on another
+    // stream still reading this buffer; its owning stream is long idle, so
+    // dropping here would complete the stream-ordered free immediately and
+    // the pool could hand the memory to a concurrent table's allocation
+    // while those kernels run. Drain the device before the drop — this is a
+    // rare recovery path.
+    if be.ctx.synchronize().is_err() {
+        return false;
+    }
+    trace.aux_resident = None;
+    GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed);
+    true
+}
+
+/// Diagnostic: download a resident ext3 handle (3-slab layout) as per-column
+/// host Vecs. Used by the xcheck post-mortem to compare the committed R2
+/// parts against a host recompute.
+pub(crate) fn download_ext3_columns<E>(
+    h: &math_cuda::lde::GpuLdeExt3,
+) -> Option<Vec<Vec<FieldElement<E>>>>
+where
+    E: IsField + 'static,
+{
+    if TypeId::of::<E>() != TypeId::of::<Degree3GoldilocksExtensionField>() {
+        return None;
+    }
+    let be = math_cuda::device::backend().ok()?;
+    let stream = be.next_stream();
+    h.wait_ready_on(&stream).ok()?;
+    let slabs = stream.clone_dtoh(h.buf.as_ref()).ok()?;
+    stream.synchronize().ok()?;
+    let (m, lde) = (h.m, h.lde_size);
+    if slabs.len() != m * lde * 3 {
+        return None;
+    }
+    let mut cols = Vec::with_capacity(m);
+    for c in 0..m {
+        let mut interleaved = vec![0u64; lde * 3];
+        for k in 0..3 {
+            let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde];
+            for r in 0..lde {
+                interleaved[r * 3 + k] = slab[r];
+            }
+        }
+        cols.push(u64_to_ext3_vec::<E>(&interleaved));
+    }
+    Some(cols)
+}
+
 pub fn gpu_batch_invert_calls() -> u64 {
     GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed)
 }