feat(stark): full GPU trace residency — skip the round-1 D2H

crypto/math-cuda/kernels/barycentric.cu

@@ -190,3 +190,45 @@ extern "C" __global__ void barycentric_ext3_batched_strided(
         out_ext3_int[col * 3 + 2] = sum.c;
     }
 }
+
+// Gather full rows from a device-resident base-field LDE (`buf[col*col_stride +
+// row]`). One block per gathered row, threads stride over columns. Output is
+// row-major `out[q*num_cols + col]` for gathered-row slot `q` — directly the
+// concatenation of `gather_main_row(rows[q])` for each q. `rows` are the LDE row
+// indices to gather (already the reversed query rows on the host side).
+extern "C" __global__ void gather_rows_base(
+    const uint64_t *__restrict__ columns,
+    uint64_t col_stride,
+    uint64_t num_cols,
+    const uint32_t *__restrict__ rows,
+    uint64_t num_rows,
+    uint64_t *__restrict__ out
+) {
+    uint64_t q = blockIdx.x;
+    if (q >= num_rows) return;
+    uint64_t row = rows[q];
+    for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) {
+        out[q * num_cols + col] = columns[col * col_stride + row];
+    }
+}
+
+// Ext3 variant: M ext3 columns as 3M base slabs, `columns[(col*3+k)*col_stride +
+// row]`. Output interleaved ext3: `out[(q*num_cols + col)*3 + k]`.
+extern "C" __global__ void gather_rows_ext3(
+    const uint64_t *__restrict__ columns,
+    uint64_t col_stride,
+    uint64_t num_cols,
+    const uint32_t *__restrict__ rows,
+    uint64_t num_rows,
+    uint64_t *__restrict__ out
+) {
+    uint64_t q = blockIdx.x;
+    if (q >= num_rows) return;
+    uint64_t row = rows[q];
+    for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) {
+        uint64_t o = (q * num_cols + col) * 3;
+        out[o + 0] = columns[(col * 3 + 0) * col_stride + row];
+        out[o + 1] = columns[(col * 3 + 1) * col_stride + row];
+        out[o + 2] = columns[(col * 3 + 2) * col_stride + row];
+    }
+}

crypto/math-cuda/src/barycentric.rs

@@ -337,3 +337,81 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms(
     stream.synchronize()?;
     Ok(out)
 }
+
+/// Gather full rows from a device-resident base-field LDE handle. `rows` are LDE
+/// row indices; returns their column values row-major (`rows.len() * main.m`
+/// u64, `out[q*num_cols + col]`) — i.e. the concatenation of
+/// `gather_main_row(rows[q])` for each `q`. Runs on the caller's stream.
+pub fn gather_rows_base_on_device(
+    main: &GpuLdeBase,
+    rows: &[u32],
+    stream: &Arc<CudaStream>,
+) -> Result<Vec<u64>> {
+    let num_cols = main.m;
+    if num_cols == 0 || rows.is_empty() {
+        return Ok(Vec::new());
+    }
+    let be = backend()?;
+    let rows_dev = stream.clone_htod(rows)?;
+    let mut out = stream.alloc_zeros::<u64>(rows.len() * num_cols)?;
+    let col_stride = main.lde_size as u64;
+    let num_cols_u64 = num_cols as u64;
+    let num_rows_u64 = rows.len() as u64;
+    let cfg = LaunchConfig {
+        grid_dim: (rows.len() as u32, 1, 1),
+        block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1),
+        shared_mem_bytes: 0,
+    };
+    unsafe {
+        stream
+            .launch_builder(&be.gather_rows_base)
+            .arg(main.buf.as_ref())
+            .arg(&col_stride)
+            .arg(&num_cols_u64)
+            .arg(&rows_dev)
+            .arg(&num_rows_u64)
+            .arg(&mut out)
+            .launch(cfg)?;
+    }
+    let host = stream.clone_dtoh(&out)?;
+    stream.synchronize()?;
+    Ok(host)
+}
+
+/// Ext3 sibling of [`gather_rows_base_on_device`]: returns `rows.len() * aux.m *
+/// 3` u64, interleaved ext3 (`out[(q*num_cols + col)*3 + k]`).
+pub fn gather_rows_ext3_on_device(
+    aux: &GpuLdeExt3,
+    rows: &[u32],
+    stream: &Arc<CudaStream>,
+) -> Result<Vec<u64>> {
+    let num_cols = aux.m;
+    if num_cols == 0 || rows.is_empty() {
+        return Ok(Vec::new());
+    }
+    let be = backend()?;
+    let rows_dev = stream.clone_htod(rows)?;
+    let mut out = stream.alloc_zeros::<u64>(rows.len() * num_cols * 3)?;
+    let col_stride = aux.lde_size as u64;
+    let num_cols_u64 = num_cols as u64;
+    let num_rows_u64 = rows.len() as u64;
+    let cfg = LaunchConfig {
+        grid_dim: (rows.len() as u32, 1, 1),
+        block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1),
+        shared_mem_bytes: 0,
+    };
+    unsafe {
+        stream
+            .launch_builder(&be.gather_rows_ext3)
+            .arg(aux.buf.as_ref())
+            .arg(&col_stride)
+            .arg(&num_cols_u64)
+            .arg(&rows_dev)
+            .arg(&num_rows_u64)
+            .arg(&mut out)
+            .launch(cfg)?;
+    }
+    let host = stream.clone_dtoh(&out)?;
+    stream.synchronize()?;
+    Ok(host)
+}

crypto/math-cuda/src/device.rs

@@ -167,6 +167,8 @@ pub struct Backend {
     pub barycentric_ext3_batched: CudaFunction,
     pub barycentric_base_batched_strided: CudaFunction,
     pub barycentric_ext3_batched_strided: CudaFunction,
+    pub gather_rows_base: CudaFunction,
+    pub gather_rows_ext3: CudaFunction,
 
     // deep.ptx
     pub deep_composition_ext3_row: CudaFunction,
@@ -367,6 +369,8 @@ impl Backend {
                 .load_function("barycentric_base_batched_strided")?,
             barycentric_ext3_batched_strided: bary
                 .load_function("barycentric_ext3_batched_strided")?,
+            gather_rows_base: bary.load_function("gather_rows_base")?,
+            gather_rows_ext3: bary.load_function("gather_rows_ext3")?,
             deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?,
             fri_fold_ext3: fri.load_function("fri_fold_ext3")?,
             fri_update_twiddles: fri.load_function("fri_update_twiddles")?,

crypto/math-cuda/src/lde.rs

@@ -400,6 +400,7 @@ enum InnerInput<'a> {
 /// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in
 /// the appropriate LDE handle.
 #[allow(clippy::type_complexity)]
+#[allow(clippy::too_many_arguments)]
 fn coset_lde_row_major_inner(
     input: InnerInput,
     n: usize,
@@ -408,6 +409,7 @@ fn coset_lde_row_major_inner(
     weights: &[u64],
     what: &str,
     retain_trace_col_major: bool,
+    retain_host_lde: bool,
 ) -> Result<(
     GpuMerkleTree,
     CudaSlice<u64>,
@@ -514,7 +516,10 @@ fn coset_lde_row_major_inner(
 
     // D2H the row-major LDE first (before the handle transpose). Release the
     // staging lock before the Merkle nodes transfer to minimise lock contention.
-    let lde_out = {
+    // Skipped entirely when `retain_host_lde` is false (the caller keeps the LDE
+    // device-only): this is the round-1 trace D2H the full-residency path
+    // eliminates — the big transfer/alloc win — so we return an empty host Vec.
+    let lde_out = if retain_host_lde {
         let staging_slot = be.pinned_staging();
         let mut staging = staging_slot.lock().unwrap();
         staging.ensure_capacity(lde_size * total_cols, &be.ctx)?;
@@ -524,6 +529,8 @@ fn coset_lde_row_major_inner(
         let out = pinned[..lde_size * total_cols].to_vec();
         drop(staging);
         out
+    } else {
+        Vec::new()
     };
 
     // Keep the Merkle tree resident on device; copy only the 32 byte root so the
@@ -562,6 +569,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep(
     m: usize,
     blowup_factor: usize,
     weights: &[u64],
+    retain_host_lde: bool,
 ) -> Result<(GpuLdeBase, Vec<u64>)> {
     let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner(
         InnerInput::Host(row_major),
@@ -571,6 +579,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep(
         weights,
         "coset_lde_row_major lde_size",
         true,
+        retain_host_lde,
     )?;
     let handle = GpuLdeBase {
         buf: Arc::new(col_major_dev),
@@ -599,6 +608,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep(
     m: usize,
     blowup_factor: usize,
     weights: &[u64],
+    retain_host_lde: bool,
 ) -> Result<(GpuLdeExt3, Vec<u64>)> {
     let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner(
         InnerInput::Host(row_major),
@@ -608,6 +618,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep(
         weights,
         "coset_lde_ext3_row_major lde_size",
         false,
+        retain_host_lde,
     )?;
     let handle = GpuLdeExt3 {
         buf: Arc::new(col_major_dev),
@@ -628,6 +639,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev(
     m: usize,
     blowup_factor: usize,
     weights: &[u64],
+    retain_host_lde: bool,
 ) -> Result<(GpuLdeExt3, Vec<u64>)> {
     let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner(
         InnerInput::Dev(input_dev),
@@ -637,6 +649,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev(
         weights,
         "coset_lde_ext3_row_major_dev lde_size",
         false,
+        retain_host_lde,
     )?;
     let handle = GpuLdeExt3 {
         buf: Arc::new(col_major_dev),

crypto/math-cuda/tests/gather_rows.rs

@@ -0,0 +1,95 @@
+//! Parity: the device row-gather (`gather_rows_base` / `gather_rows_ext3`,
+//! used by R4 query openings to read opened column values from the resident LDE
+//! instead of the host trace) returns exactly the column values obtained by
+//! directly indexing the column-major device buffer at each requested row.
+
+use std::sync::Arc;
+
+use math_cuda::barycentric::{gather_rows_base_on_device, gather_rows_ext3_on_device};
+use math_cuda::device::backend;
+use math_cuda::lde::{GpuLdeBase, GpuLdeExt3};
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
+
+fn run_base(lde_size: usize, num_cols: usize, seed: u64) {
+    let mut rng = ChaCha8Rng::seed_from_u64(seed);
+    // Column-major base LDE: buf[col*lde_size + row].
+    let buf: Vec<u64> = (0..num_cols * lde_size)
+        .map(|_| rng.r#gen::<u64>())
+        .collect();
+
+    let be = backend().unwrap();
+    let stream = be.next_stream();
+    let dev = stream.clone_htod(&buf).unwrap();
+    stream.synchronize().unwrap();
+    let handle = GpuLdeBase {
+        buf: Arc::new(dev),
+        m: num_cols,
+        lde_size,
+        tree: None,
+        trace_dev: None,
+        trace_rows: 0,
+    };
+
+    let rows: Vec<u32> = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect();
+    let got = gather_rows_base_on_device(&handle, &rows, &stream).unwrap();
+    assert_eq!(got.len(), rows.len() * num_cols, "base gather shape");
+    for (q, &row) in rows.iter().enumerate() {
+        for col in 0..num_cols {
+            assert_eq!(
+                got[q * num_cols + col],
+                buf[col * lde_size + row as usize],
+                "base gather mismatch: row {row}, col {col}"
+            );
+        }
+    }
+}
+
+fn run_ext3(lde_size: usize, num_cols: usize, seed: u64) {
+    let mut rng = ChaCha8Rng::seed_from_u64(seed);
+    // De-interleaved ext3 LDE: buf[(col*3 + k)*lde_size + row].
+    let buf: Vec<u64> = (0..num_cols * 3 * lde_size)
+        .map(|_| rng.r#gen::<u64>())
+        .collect();
+
+    let be = backend().unwrap();
+    let stream = be.next_stream();
+    let dev = stream.clone_htod(&buf).unwrap();
+    stream.synchronize().unwrap();
+    let handle = GpuLdeExt3 {
+        buf: Arc::new(dev),
+        m: num_cols,
+        lde_size,
+        tree: None,
+    };
+
+    let rows: Vec<u32> = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect();
+    let got = gather_rows_ext3_on_device(&handle, &rows, &stream).unwrap();
+    assert_eq!(got.len(), rows.len() * num_cols * 3, "ext3 gather shape");
+    for (q, &row) in rows.iter().enumerate() {
+        for col in 0..num_cols {
+            let o = (q * num_cols + col) * 3;
+            for k in 0..3 {
+                assert_eq!(
+                    got[o + k],
+                    buf[(col * 3 + k) * lde_size + row as usize],
+                    "ext3 gather mismatch: row {row}, col {col}, comp {k}"
+                );
+            }
+        }
+    }
+}
+
+#[test]
+fn gather_rows_base_matches_direct_indexing() {
+    for (log_size, cols) in [(6u32, 3usize), (12, 20), (16, 8)] {
+        run_base(1usize << log_size, cols, 100 + log_size as u64);
+    }
+}
+
+#[test]
+fn gather_rows_ext3_matches_direct_indexing() {
+    for (log_size, cols) in [(6u32, 2usize), (12, 5), (16, 3)] {
+        run_ext3(1usize << log_size, cols, 200 + log_size as u64);
+    }
+}

crypto/math-cuda/tests/merkle_root_parity.rs

@@ -305,6 +305,7 @@ fn new_row_major_pipeline_base_root_matches_cpu() {
                     num_cols,
                     blowup,
                     &weights_u64,
+                    true,
                 )
                 .expect("new row-major GPU pipeline");
                 let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root;
@@ -368,6 +369,7 @@ fn new_row_major_pipeline_ext3_root_matches_cpu() {
                         num_cols,
                         blowup,
                         &weights_u64,
+                        true,
                     )
                     .expect("new ext3 row-major GPU pipeline");
                 let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root;

crypto/stark/src/constraint_ir/gpu_interp.rs

@@ -65,6 +65,39 @@ unsafe fn base_slice_to_u64<F: IsField>(xs: &[FieldElement<F>]) -> Vec<u64> {
     raw.to_vec()
 }
 
+/// Lift raw base-field limbs (one canonical-Goldilocks `u64` per element, as
+/// produced by the device row-gather kernels) back into owned `FieldElement`s.
+/// Returns `None` unless `F == GoldilocksField`. The inverse of
+/// [`base_slice_to_u64`]; used to feed device-gathered LDE rows into the generic
+/// prover openings.
+pub fn base_u64_to_field<F: IsField + 'static>(raw: &[u64]) -> Option<Vec<FieldElement<F>>> {
+    if TypeId::of::<F>() != TypeId::of::<GoldilocksField>() {
+        return None;
+    }
+    // SAFETY: the gate established `F == GoldilocksField`, whose `FieldElement`
+    // is `#[repr(transparent)]` over `u64`; `raw` (a `*const u64`) is 8-aligned.
+    let fe =
+        unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement<F>, raw.len()) };
+    Some(fe.to_vec())
+}
+
+/// Ext3 sibling of [`base_u64_to_field`]: `raw` holds `3` interleaved limbs per
+/// element (`[c0, c1, c2]`). Returns `None` unless `E` is the degree-3
+/// Goldilocks extension. Inverse of [`ext3_slice_to_u64`].
+pub fn ext3_u64_to_field<E: IsField + 'static>(raw: &[u64]) -> Option<Vec<FieldElement<E>>> {
+    if TypeId::of::<E>() != TypeId::of::<GoldilocksExtension>() {
+        return None;
+    }
+    debug_assert_eq!(raw.len() % 3, 0, "ext3 limbs must come in triples");
+    // SAFETY: the gate established the degree-3 Goldilocks extension, whose
+    // `FieldElement` is `#[repr(transparent)]` over `[u64; 3]`; `raw` (a
+    // `*const u64`) is 8-aligned, matching `[u64; 3]`'s alignment.
+    let fe = unsafe {
+        std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement<E>, raw.len() / 3)
+    };
+    Some(fe.to_vec())
+}
+
 /// Per-proof accumulation inputs (in `FieldElement` form) for
 /// [`try_eval_composition_gpu`], mirroring the CPU accumulation in
 /// `constraints::evaluator`.