perf(stark): GPU transfer-residency

crypto/math-cuda/kernels/fri.cu

@@ -47,6 +47,40 @@ extern "C" __global__ void fri_fold_ext3(
     out_p[2] = res.c;
 }
 
+// First DEEP->FRI fold when the DEEP producer hands over its natural-order
+// codeword directly. The legacy host bridge bit-reverses the full codeword and
+// then pairs entries (2*j, 2*j+1). Bit reversal is an involution, so reading
+// natural input at br(2*j) / br(2*j+1) is byte-identical while avoiding both
+// the CPU permutation and a separate device permutation pass. Output j is in
+// the ordinary FRI layer order, so every later fold uses fri_fold_ext3.
+extern "C" __global__ void fri_fold_ext3_from_natural(
+    const uint64_t *__restrict__ in,
+    uint64_t n_out,
+    uint64_t log_n_in,
+    const uint64_t *__restrict__ inv_tw,
+    const uint64_t *__restrict__ zeta,
+    uint64_t *__restrict__ out) {
+    uint64_t j = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (j >= n_out) return;
+
+    uint64_t br_lo = __brevll(2 * j) >> (64 - log_n_in);
+    uint64_t br_hi = __brevll(2 * j + 1) >> (64 - log_n_in);
+    const uint64_t *lo_p = in + br_lo * 3;
+    const uint64_t *hi_p = in + br_hi * 3;
+
+    ext3::Fe3 lo = ext3::make(lo_p[0], lo_p[1], lo_p[2]);
+    ext3::Fe3 hi = ext3::make(hi_p[0], hi_p[1], hi_p[2]);
+    ext3::Fe3 sum = ext3::add(lo, hi);
+    ext3::Fe3 diff = ext3::sub(lo, hi);
+    ext3::Fe3 z = ext3::make(zeta[0], zeta[1], zeta[2]);
+    ext3::Fe3 res = ext3::add(sum, ext3::mul_base(ext3::mul(z, diff), inv_tw[j]));
+
+    uint64_t *out_p = out + j * 3;
+    out_p[0] = res.a;
+    out_p[1] = res.b;
+    out_p[2] = res.c;
+}
+
 // update_twiddles: tw_out[j] = tw_in[2j]^2 for j in 0..n_out.
 // Separate input/output buffers: thread j reads tw_in[2j] while thread 2j
 // writes tw_out[2j], so an in-place version would race across threads.

crypto/math-cuda/src/barycentric.rs

@@ -257,6 +257,11 @@ pub fn barycentric_ext3_on_device(
 
     let be = backend()?;
     let stream = be.next_stream();
+    // Order this stream against the producer's fill of the LDE (no-op unless the
+    // handle carries a ready event, i.e. the composition device-only path).
+    if let Some(ev) = aux_handle.ready.as_deref() {
+        stream.wait(ev)?;
+    }
 
     let points_dev = stream.clone_htod(coset_points)?;
     let inv_dev = stream.clone_htod(inv_denoms_ext3)?;
@@ -390,6 +395,11 @@ pub fn gather_rows_ext3_on_device(
         return Ok(Vec::new());
     }
     let be = backend()?;
+    // Order the caller's stream against the producer's fill (no-op unless the
+    // handle carries a ready event, i.e. the composition device-only path).
+    if let Some(ev) = aux.ready.as_deref() {
+        stream.wait(ev)?;
+    }
     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;

crypto/math-cuda/src/constraint_interp.rs

@@ -274,5 +274,6 @@ pub fn eval_composition_on_device(
     }
     let out = stream.clone_dtoh(&d_h)?;
     stream.synchronize()?;
+    crate::stagebytes::add_comp_dh_d2h(out.len() * 8);
     Ok(out)
 }

crypto/math-cuda/src/deep.rs

@@ -16,6 +16,21 @@ use crate::Result;
 use crate::device::backend;
 use crate::lde::{GpuLdeBase, GpuLdeExt3};
 
+/// Natural-order DEEP codeword kept on the producer stream for direct FRI
+/// consumption. `len` counts ext3 elements; `buf` contains `3 * len` u64s.
+pub struct GpuDeepEvals {
+    pub buf: CudaSlice<u64>,
+    pub len: usize,
+    pub stream: Arc<CudaStream>,
+}
+
+/// Result of the retaining DEEP entry point. `host` is empty when the caller
+/// elects device-only handoff; the device handle is always populated.
+pub struct DeepCompositionKeep {
+    pub host: Vec<u64>,
+    pub device: GpuDeepEvals,
+}
+
 /// Compute deep-composition evaluations on device.
 ///
 /// `num_eval_points = trace_terms_gammas_interleaved.len() / ((num_main +
@@ -88,6 +103,11 @@ pub fn deep_composition_ext3_with_dev_parts(
 ) -> Result<Vec<u64>> {
     let be = backend()?;
     let stream = be.next_stream();
+    // Order this stream against the producer's fill of the composition LDE
+    // (no-op while the producer still `synchronize()`s; real barrier at C4).
+    if let Some(ev) = h_parts_dev.ready.as_deref() {
+        stream.wait(ev)?;
+    }
     deep_composition_ext3_impl(
         &stream,
         main_lde,
@@ -138,6 +158,51 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms(
     row_stride: usize,
     domain_size: usize,
 ) -> Result<Vec<u64>> {
+    Ok(deep_composition_ext3_with_dev_parts_and_inv_denoms_keep(
+        stream,
+        main_lde,
+        aux_lde,
+        h_parts_dev,
+        inv_denoms_dev,
+        h_ood,
+        trace_ood,
+        gammas_h,
+        gammas_tr,
+        num_parts,
+        num_main,
+        num_aux,
+        num_eval_points,
+        row_stride,
+        domain_size,
+        true,
+    )?
+    .host)
+}
+
+/// Retaining variant of
+/// [`deep_composition_ext3_with_dev_parts_and_inv_denoms`]. When
+/// `retain_host` is false it does not synchronize or copy the codeword to the
+/// host; the returned handle remains bound to the producer stream so FRI can
+/// consume it in FIFO order.
+#[allow(clippy::too_many_arguments)]
+pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms_keep(
+    stream: &Arc<CudaStream>,
+    main_lde: &GpuLdeBase,
+    aux_lde: Option<&GpuLdeExt3>,
+    h_parts_dev: &GpuLdeExt3,
+    inv_denoms_dev: &CudaSlice<u64>,
+    h_ood: &[u64],
+    trace_ood: &[u64],
+    gammas_h: &[u64],
+    gammas_tr: &[u64],
+    num_parts: usize,
+    num_main: usize,
+    num_aux: usize,
+    num_eval_points: usize,
+    row_stride: usize,
+    domain_size: usize,
+    retain_host: bool,
+) -> Result<DeepCompositionKeep> {
     assert_eq!(main_lde.m, num_main);
     assert_eq!(h_parts_dev.m, num_parts);
     assert_eq!(h_parts_dev.lde_size, main_lde.lde_size);
@@ -174,6 +239,12 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms(
 
     let be = backend()?;
 
+    // Order the caller's stream against the producer's fill of the composition
+    // LDE (no-op while the producer still `synchronize()`s; real barrier at C4).
+    if let Some(ev) = h_parts_dev.ready.as_deref() {
+        stream.wait(ev)?;
+    }
+
     // H2D only the small scalars on the caller's stream.
     let h_ood_dev = stream.clone_htod(h_ood)?;
     let trace_ood_dev = stream.clone_htod(trace_ood)?;
@@ -232,9 +303,22 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms(
             .launch(cfg)?;
     }
 
-    let out = stream.clone_dtoh(&deep_out)?;
-    stream.synchronize()?;
-    Ok(out)
+    let host = if retain_host {
+        let out = stream.clone_dtoh(&deep_out)?;
+        crate::stagebytes::add_deep_out_d2h(core::mem::size_of_val(out.as_slice()));
+        stream.synchronize()?;
+        out
+    } else {
+        Vec::new()
+    };
+    Ok(DeepCompositionKeep {
+        host,
+        device: GpuDeepEvals {
+            buf: deep_out,
+            len: domain_size,
+            stream: Arc::clone(stream),
+        },
+    })
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -359,6 +443,7 @@ fn deep_composition_ext3_impl(
     }
 
     let out = stream.clone_dtoh(&deep_out)?;
+    crate::stagebytes::add_deep_out_d2h(core::mem::size_of_val(out.as_slice()));
     stream.synchronize()?;
     Ok(out)
 }

crypto/math-cuda/src/device.rs

@@ -180,6 +180,7 @@ pub struct Backend {
 
     // fri.cubin
     pub fri_fold_ext3: CudaFunction,
+    pub fri_fold_ext3_from_natural: CudaFunction,
     pub fri_update_twiddles: CudaFunction,
 
     // inverse.cubin
@@ -379,6 +380,7 @@ impl Backend {
             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_fold_ext3_from_natural: fri.load_function("fri_fold_ext3_from_natural")?,
             fri_update_twiddles: fri.load_function("fri_update_twiddles")?,
             compute_denoms_ext3: inverse.load_function("compute_denoms_ext3")?,
             block_inclusive_scan_fwd_ext3: inverse