GPU trace-gen (part 1): device-resident fills for 7 tables

crypto/math-cuda/build.rs

@@ -121,5 +121,6 @@ fn main() {
     compile_ptx("fri.cu", "fri.ptx", have_nvcc);
     compile_ptx("inverse.cu", "inverse.ptx", have_nvcc);
     compile_ptx("logup.cu", "logup.ptx", have_nvcc);
+    compile_ptx("trace_cpu.cu", "trace_cpu.ptx", have_nvcc);
     compile_ptx("constraint_interp.cu", "constraint_interp.ptx", have_nvcc);
 }

crypto/math-cuda/src/device.rs

@@ -98,6 +98,7 @@ const DEEP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/deep.ptx"));
 const FRI_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/fri.ptx"));
 const INVERSE_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/inverse.ptx"));
 const LOGUP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/logup.ptx"));
+const TRACE_CPU_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/trace_cpu.ptx"));
 const CONSTRAINT_INTERP_PTX: &str =
     include_str!(concat!(env!("OUT_DIR"), "/constraint_interp.ptx"));
 
@@ -190,6 +191,14 @@ pub struct Backend {
     pub logup_finalize_accum_ext3: CudaFunction,
     pub logup_assemble_aux_ext3: CudaFunction,
 
+    // On-GPU trace generation.
+    pub trace_cpu_fill: CudaFunction,
+    pub memw_aligned_fill: CudaFunction,
+    pub load_fill: CudaFunction,
+    pub store_fill: CudaFunction,
+    pub shift_fill: CudaFunction,
+    pub lt_fill: CudaFunction,
+    pub memw_register_fill: CudaFunction,
     // constraint_interp.ptx
     pub constraint_interp_kernel: CudaFunction,
     pub constraint_composition_kernel: CudaFunction,
@@ -297,6 +306,7 @@ impl Backend {
         let fri = ctx.load_module(Ptx::from_src(FRI_PTX))?;
         let inverse = ctx.load_module(Ptx::from_src(INVERSE_PTX))?;
         let logup = ctx.load_module(Ptx::from_src(LOGUP_PTX))?;
+        let trace_cpu = ctx.load_module(Ptx::from_src(TRACE_CPU_PTX))?;
         let constraint_interp = ctx.load_module(Ptx::from_src(CONSTRAINT_INTERP_PTX))?;
 
         let mut streams = Vec::with_capacity(STREAM_POOL_SIZE);
@@ -385,6 +395,13 @@ impl Backend {
             logup_apply_offsets_add_ext3: logup.load_function("logup_apply_offsets_add_ext3")?,
             logup_finalize_accum_ext3: logup.load_function("logup_finalize_accum_ext3")?,
             logup_assemble_aux_ext3: logup.load_function("logup_assemble_aux_ext3")?,
+            trace_cpu_fill: trace_cpu.load_function("trace_cpu_fill")?,
+            memw_aligned_fill: trace_cpu.load_function("memw_aligned_fill")?,
+            load_fill: trace_cpu.load_function("load_fill")?,
+            store_fill: trace_cpu.load_function("store_fill")?,
+            shift_fill: trace_cpu.load_function("shift_fill")?,
+            lt_fill: trace_cpu.load_function("lt_fill")?,
+            memw_register_fill: trace_cpu.load_function("memw_register_fill")?,
             constraint_interp_kernel: constraint_interp
                 .load_function("constraint_interp_kernel")?,
             constraint_composition_kernel: constraint_interp

crypto/math-cuda/src/inverse.rs

@@ -84,12 +84,15 @@ pub fn batch_inverse_ext3_dev(
     stream: &Arc<CudaStream>,
 ) -> Result<CudaSlice<u64>> {
     assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1");
-    // Runtime guard (not debug_assert): a u32 grid_dim is truncated past
-    // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks
-    // and leave a tail uninverted. Reachable on LDE size 2^23+ × multi-
-    // eval-point R4. Returning Err lets the dispatcher's Err(_) => None
-    // route the caller to the CPU `inplace_batch_inverse` fallback.
-    if n > u32::MAX as usize / BLOCK_SIZE as usize {
+    // Runtime guard (not debug_assert): `n` must fit in u32 so neither `n as u32`
+    // below nor the kernel's flat index truncates. The block count
+    // `n.div_ceil(BLOCK_SIZE)` is well within CUDA's grid.x bound (2^31-1) for any
+    // such `n` — grid.x is NOT the 65535 cap that applies to grid.y/z, so the old
+    // `u32::MAX / BLOCK_SIZE` bound was 256× too strict and spuriously declined
+    // legitimate launches (e.g. logup batch-invert of 18 interactions × 2^20 rows
+    // = ~18.9M), silently routing to a wrong fallback. Returning Err lets the
+    // dispatcher's Err(_) => None route the caller to the CPU fallback.
+    if n > u32::MAX as usize {
         return Err(cudarc::driver::DriverError(
             cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE,
         ));
@@ -182,10 +185,11 @@ pub fn compute_and_invert_denoms_ext3_dev(
     let total = k_scalars
         .checked_mul(n)
         .expect("compute_and_invert_denoms_ext3_dev: k_scalars * n overflow");
-    // See `batch_inverse_ext3_dev` for the rationale: runtime Err, not
-    // debug_assert, so release builds also route past the silent-truncation
-    // hazard via the caller's CPU fallback.
-    if total > u32::MAX as usize / BLOCK_SIZE as usize {
+    // See `batch_inverse_ext3_dev` for the rationale: guard on `u32::MAX` (the
+    // cast/index bound), NOT `u32::MAX / BLOCK_SIZE` (which is the grid.y/z cap, not
+    // grid.x, and 256× too strict). Runtime Err, not debug_assert, so release
+    // builds route past the truncation hazard via the caller's CPU fallback.
+    if total > u32::MAX as usize {
         return Err(cudarc::driver::DriverError(
             cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE,
         ));

crypto/math-cuda/src/lde.rs

@@ -583,6 +583,40 @@ pub fn coset_lde_row_major_with_merkle_tree_keep(
     Ok((handle, lde_out))
 }
 
+/// Like [`coset_lde_row_major_with_merkle_tree_keep`] but the input is an
+/// already-resident device buffer (`n * m` base-field elements, row-major
+/// `[row*m + col]`). No PCIe upload: the buffer is copied device-to-device into
+/// the LDE scratch. Used by the on-GPU trace generator so a device-born main
+/// trace feeds the LDE without a host round-trip. Retains the trace-domain
+/// column-major snapshot exactly like the host keep path, so the resident LogUp
+/// aux fingerprint (#762) still works unchanged.
+pub fn coset_lde_row_major_with_merkle_tree_keep_dev(
+    input_dev: &CudaSlice<u64>,
+    n: usize,
+    m: usize,
+    blowup_factor: usize,
+    weights: &[u64],
+) -> Result<(GpuLdeBase, Vec<u64>)> {
+    let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner(
+        InnerInput::Dev(input_dev),
+        n,
+        m,
+        blowup_factor,
+        weights,
+        "coset_lde_row_major_dev lde_size",
+        true,
+    )?;
+    let handle = GpuLdeBase {
+        buf: Arc::new(col_major_dev),
+        m,
+        lde_size: n * blowup_factor,
+        tree: Some(tree),
+        trace_dev: trace_col_major.map(Arc::new),
+        trace_rows: n,
+    };
+    Ok((handle, lde_out))
+}
+
 /// Row-major ext3 LDE + Keccak + Merkle, all on-device.
 ///
 /// `Fp3` is `[u64; 3]` in memory, so row-major ext3 with `m` ext3 columns is

crypto/math-cuda/src/lib.rs

@@ -15,6 +15,7 @@ pub mod lde;
 pub mod logup;
 pub mod merkle;
 pub mod ntt;
+pub mod trace_cpu;
 
 // Re-exported for downstream crates so they can refer to CUDA primitive
 // types without depending on cudarc directly.

crypto/math-cuda/src/logup.rs

@@ -44,12 +44,16 @@ pub struct LogupDescriptor<'a> {
 }
 
 fn cfg(total: usize) -> Result<LaunchConfig> {
-    // See `batch_inverse_ext3_dev` for the rationale: a u32 grid_dim is
-    // truncated past u32::MAX / BLOCK_SIZE, which would silently launch too
-    // few blocks and leave a tail of the (uninitialized) output unwritten.
-    // Runtime Err, not debug_assert, so release builds also route to the
-    // caller's CPU fallback.
-    if total > u32::MAX as usize / BLOCK_SIZE as usize {
+    // One thread per element: `total` must fit in u32 so neither the `total as u32`
+    // cast below nor the kernel's `blockIdx.x*blockDim.x + threadIdx.x` index
+    // truncates. The block count `total.div_ceil(BLOCK_SIZE)` is well within CUDA's
+    // grid.x bound (2^31-1) for any such `total` — grid.x is NOT capped at the
+    // 65535 that applies to grid.y/z. (The previous bound `u32::MAX / BLOCK_SIZE`
+    // conflated the two and spuriously declined large single-chunk tables — e.g.
+    // SHIFT at 2^20 rows × 18 interactions = ~18.9M > u32::MAX/256 — silently
+    // falling back to a wrong, zeroed-host aux.) Runtime Err (not debug_assert) so
+    // release builds route to the caller's CPU fallback rather than truncate.
+    if total > u32::MAX as usize {
         return Err(cudarc::driver::DriverError(
             cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE,
         ));