perf: the epoch's columns go to the card once

Four things read the same trace columns — the commitment, the sumcheck's
factors, the evaluation at the reduction point, and the opening's message
— and each uploaded its own copy. Counted on the real block that is 160.9
GiB crossing the bus and 10.21 s of a 44.8 s prove, for forty gigabytes of
data. At thirteen to nineteen gigabytes a second they were already running
at what pageable host memory gives, so there was nothing to make faster:
only three passes to stop making.

The columns now go up once, when the epoch's tables are committed, and
live as long as the tables that read them. A table's columns are a
contiguous run of equal height, which is the layout the factor gather and
the batched evaluation already want, so those two read them where they
lie; the commitment and the opening scatter theirs by offset, which on the
card is a copy at device bandwidth. Every entry point still takes columns
that are only here, which is what runs when there is no device or it will
not promise the room.

Real block, 19 epochs of 2^21: 51.16 -> 45.31 s. Peak VRAM 30223 -> 31055
MiB of 32607 — 832 more, against the 1280 that holding a copy without
removing the uploads cost, because what each site used to allocate to
upload into is gone. Peak host RSS unchanged at 9.3 GB.

The margin this leaves ties the epoch to 2^21: a sweep already measured
32.1 GiB of peak at 2^22 without any of this. Epoch size stops being a
free knob and starts depending on a memory model that does not exist yet.

Six consecutive runs of the GPU suite green, the cross-check sweep green,
the four kernel parity tests green, 274 and 261 on the host.

Author: ColoCarletti

crypto/math-cuda/src/columns.rs

@@ -0,0 +1,148 @@
+//! The epoch's trace columns, on the card once.
+//!
+//! Four things read the same columns — the commitment, the sumcheck's factors,
+//! the evaluation at the reduction point, and the opening's message — and each
+//! used to upload its own copy. That is the same forty gigabytes crossing the
+//! bus four times, and at the speed pageable host memory gives it is a fifth of
+//! the prove. They are put here once and read where they lie.
+//!
+//! A table's columns are a contiguous run of equal height, which is the layout
+//! the factor gather and the batched evaluation already want; the commitment and
+//! the opening scatter theirs, which on the card is a copy at device bandwidth.
+
+use std::sync::Arc;
+
+use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
+
+use crate::Result;
+use crate::device::{DeviceReservation, alloc_or_trim, backend};
+
+/// Where a set of columns is: still here, or already there.
+pub enum Columns<'a> {
+    Host(&'a [&'a [u64]]),
+    /// A run of `width` columns of `rows` each, starting at column `first`.
+    Device {
+        store: &'a DeviceColumns,
+        first: usize,
+        width: usize,
+    },
+}
+
+impl Columns<'_> {
+    pub fn width(&self) -> usize {
+        match self {
+            Self::Host(columns) => columns.len(),
+            Self::Device { width, .. } => *width,
+        }
+    }
+
+    pub fn rows(&self) -> usize {
+        match self {
+            Self::Host(columns) => columns.first().map_or(0, |c| c.len()),
+            Self::Device { store, first, .. } => store.spans[*first].1,
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.width() == 0
+    }
+}
+
+/// The columns themselves, laid end to end in one allocation.
+pub struct DeviceColumns {
+    stream: Arc<CudaStream>,
+    buffer: CudaSlice<u64>,
+    /// `(offset, len)` in elements, per column, in the order uploaded.
+    spans: Vec<(usize, usize)>,
+    _room: DeviceReservation,
+}
+
+impl DeviceColumns {
+    /// `None` when the card will not promise the room, in which case every
+    /// caller uploads its own copy as before.
+    pub fn upload(columns: &[&[u64]]) -> Option<Self> {
+        if columns.is_empty() {
+            return None;
+        }
+        let total: usize = columns.iter().map(|c| c.len()).sum();
+        let be = backend().ok()?;
+        let room = be.reserve(total as u64 * 8)?;
+        let stream = be.next_stream();
+        // SAFETY: every element is written by the copies below.
+        let mut buffer = unsafe { alloc_or_trim::<u64>(&stream, total) }.ok()?;
+        let mut spans = Vec::with_capacity(columns.len());
+        let mut at = 0usize;
+        for column in columns {
+            let mut slab = buffer.slice_mut(at..at + column.len());
+            stream.memcpy_htod(*column, &mut slab).ok()?;
+            spans.push((at, column.len()));
+            at += column.len();
+        }
+        stream.synchronize().ok()?;
+        Some(Self {
+            stream,
+            buffer,
+            spans,
+            _room: room,
+        })
+    }
+
+    pub fn num_columns(&self) -> usize {
+        self.spans.len()
+    }
+
+    /// Whether `width` columns from `first` are a run of equal height — which
+    /// is what the kernels that read a table's columns in place need.
+    pub fn is_run(&self, first: usize, width: usize) -> bool {
+        if width == 0 || first + width > self.spans.len() {
+            return false;
+        }
+        let (start, rows) = self.spans[first];
+        (0..width).all(|k| self.spans[first + k] == (start + k * rows, rows))
+    }
+
+    /// The device address of column `k` and its length in elements.
+    pub fn at(&self, k: usize) -> (u64, usize) {
+        let (offset, len) = self.spans[k];
+        let (base, _guard) = self.buffer.device_ptr(&self.stream);
+        (base + (offset * 8) as u64, len)
+    }
+
+    /// A view over `width` columns from `first`, which must be a run.
+    pub fn view(&self, first: usize, width: usize) -> cudarc::driver::CudaView<'_, u64> {
+        let (offset, rows) = self.spans[first];
+        self.buffer.slice(offset..offset + width * rows)
+    }
+
+    pub fn stream(&self) -> &Arc<CudaStream> {
+        &self.stream
+    }
+
+    /// Copies column `k` into `dst` at `offset`, on the card.
+    ///
+    /// Issued on the **destination's** stream, so whatever reads `dst` next is
+    /// ordered behind it. The store is written once and synchronized at upload
+    /// and read-only after, so another stream reading it races with nothing.
+    pub fn copy_into(
+        &self,
+        k: usize,
+        dst: &mut CudaSlice<u64>,
+        offset: usize,
+        stream: &Arc<CudaStream>,
+    ) -> Result<()> {
+        let (src, len) = self.at(k);
+        let (base, _guard) = dst.device_ptr_mut(stream);
+        // SAFETY: both ranges are inside allocations this call holds a handle
+        // to, and the caller checked `offset + len` against `dst`.
+        unsafe {
+            cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
+                base + (offset * 8) as u64,
+                src,
+                len * 8,
+                stream.cu_stream(),
+            )
+            .result()?;
+        }
+        Ok(())
+    }
+}

crypto/math-cuda/src/lib.rs

@@ -6,6 +6,7 @@
 //! pipelines or used by the parity test suite.
 
 pub mod barycentric;
+pub mod columns;
 pub mod constraint_interp;
 pub mod deep;
 pub mod device;

crypto/math-cuda/src/sumcheck.rs

@@ -553,49 +553,63 @@ pub fn evaluate_mle_base(table: &[u64], point: &[u64]) -> Result<[u64; 3]> {
 ///
 /// `columns` are base-field slices of `2^point.len()/3` values each. Returns
 /// one ext3 value per column, in order.
-pub fn evaluate_many_base(columns: &[&[u64]], point: &[u64]) -> Result<Vec<[u64; 3]>> {
+pub fn evaluate_many_base(
+    columns: crate::columns::Columns<'_>,
+    point: &[u64],
+) -> Result<Vec<[u64; 3]>> {
     assert!(point.len().is_multiple_of(3), "three u64 per coordinate");
     let vars = point.len() / 3;
     let rows = 1usize << vars;
     assert!(vars > 0, "a point with no coordinates is the table itself");
-    assert!(
-        columns.iter().all(|column| column.len() == rows),
-        "every column spans the point"
-    );
     if columns.is_empty() {
         return Ok(Vec::new());
     }
+    assert_eq!(columns.rows(), rows, "every column spans the point");
+    let num_columns = columns.width();
 
     let be = backend()?;
     let stream = be.next_stream();
     // How many go up at a time: the base copy plus the ext3 half it folds to
     // is 20 bytes a row, and this keeps that transient bounded however wide
     // the table is.
     let per_column = rows as u64 * 20;
-    let chunk = (CHUNK_BUDGET_BYTES / per_column.max(1)).clamp(1, columns.len() as u64) as usize;
+    let chunk = (CHUNK_BUDGET_BYTES / per_column.max(1)).clamp(1, num_columns as u64) as usize;
 
-    let mut out = Vec::with_capacity(columns.len());
-    for group in columns.chunks(chunk) {
+    let mut out = Vec::with_capacity(num_columns);
+    for start in (0..num_columns).step_by(chunk) {
+        let group = start..(start + chunk).min(num_columns);
+        let group_len = group.len();
         let half = rows / 2;
         // Promised before it is used, like everything else that takes a slab
         // of the card: the caller's fallback is to evaluate the columns one at
         // a time, which needs almost nothing.
-        let Some(_room) = crate::device::reserve(group.len() as u64 * per_column) else {
+        let Some(_room) = crate::device::reserve(group_len as u64 * per_column) else {
             return Err(cudarc::driver::DriverError(
                 cudarc::driver::sys::CUresult::CUDA_ERROR_OUT_OF_MEMORY,
             ));
         };
-        let mut base = unsafe { alloc_or_trim::<u64>(&stream, group.len() * rows) }?;
-        for (k, column) in group.iter().enumerate() {
-            let at = k * rows;
-            let mut slab = base.slice_mut(at..at + rows);
-            stream.memcpy_htod(*column, &mut slab)?;
-        }
+        // Read where they lie when they are already there; a copy otherwise.
+        let uploaded;
+        let base = match &columns {
+            crate::columns::Columns::Device { store, first, .. } => {
+                store.view(first + group.start, group_len)
+            }
+            crate::columns::Columns::Host(host) => {
+                let mut up = unsafe { alloc_or_trim::<u64>(&stream, group_len * rows) }?;
+                for (k, column) in host[group.clone()].iter().enumerate() {
+                    let at = k * rows;
+                    let mut slab = up.slice_mut(at..at + rows);
+                    stream.memcpy_htod(*column, &mut slab)?;
+                }
+                uploaded = up;
+                uploaded.slice(0..group_len * rows)
+            }
+        };
         let r = crate::device::htod_or_trim(&stream, &point[..3])?;
         // SAFETY: the kernel writes every element of the halves it produces.
-        let mut values = unsafe { alloc_or_trim::<u64>(&stream, group.len() * half * 3) }?;
+        let mut values = unsafe { alloc_or_trim::<u64>(&stream, group_len * half * 3) }?;
         let half_arg = half as u64;
-        let tables = group.len() as u64;
+        let tables = group_len as u64;
         let total = half_arg * tables;
         let grid = total.div_ceil(BLOCK_DIM as u64).clamp(1, MAX_GRID as u64) as u32;
         unsafe {
@@ -612,13 +626,11 @@ pub fn evaluate_many_base(columns: &[&[u64]], point: &[u64]) -> Result<Vec<[u64;
                     shared_mem_bytes: 0,
                 })?;
         }
-        drop(base);
-
         // From here every table is ext3 and they all fold together: one launch
         // per level, over the list of addresses.
         let addresses: Vec<u64> = {
             let (at, _guard) = values.device_ptr(&stream);
-            (0..group.len())
+            (0..group_len)
                 .map(|k| at + (k * half * 3 * 8) as u64)
                 .collect()
         };
@@ -631,7 +643,7 @@ pub fn evaluate_many_base(columns: &[&[u64]], point: &[u64]) -> Result<Vec<[u64;
                 break;
             }
             stream.memcpy_htod(coordinate, &mut r_dev)?;
-            let width = group.len() as u64;
+            let width = group_len as u64;
             let total = width * fold_half;
             let grid = total.div_ceil(BLOCK_DIM as u64).clamp(1, MAX_GRID as u64) as u32;
             unsafe {
@@ -656,7 +668,7 @@ pub fn evaluate_many_base(columns: &[&[u64]], point: &[u64]) -> Result<Vec<[u64;
         // Gathered there and brought back in one copy. A trace has thousands of
         // columns, and reading each one's head on its own is a transfer and a
         // stream synchronize apiece for twenty-four bytes.
-        let heads: Vec<u32> = (0..group.len()).map(|k| (k * half) as u32).collect();
+        let heads: Vec<u32> = (0..group_len).map(|k| (k * half) as u32).collect();
         let packed = crate::fri::gather_ext3_at(&values, &heads, &stream)?;
         for head in packed.chunks_exact(3) {
             out.push([head[0], head[1], head[2]]);
@@ -782,7 +794,7 @@ impl DeviceFactors {
     /// reduced mod `rows`), and the slot it fills; `public` is the extension
     /// tables that are not views of a column, each with the slot it goes to.
     pub fn from_columns(
-        columns: &[&[u64]],
+        columns: crate::columns::Columns<'_>,
         plan: &[u64],
         public: &[(usize, &[u64])],
         rows: usize,
@@ -800,7 +812,7 @@ impl DeviceFactors {
             "every slot is filled once"
         );
         assert!(
-            columns.iter().all(|column| column.len() == rows),
+            columns.is_empty() || columns.rows() == rows,
             "every column spans the cube"
         );
 
@@ -826,13 +838,26 @@ impl DeviceFactors {
         }
 
         if !plan.is_empty() {
-            // SAFETY: every cell is written by the copies below.
-            let mut base = unsafe { alloc_or_trim::<u64>(&stream, columns.len() * rows) }?;
-            for (k, column) in columns.iter().enumerate() {
-                let at = k * rows;
-                let mut slab = base.slice_mut(at..at + rows);
-                stream.memcpy_htod(*column, &mut slab)?;
-            }
+            // Read where they lie when they are already there; a copy otherwise.
+            let uploaded;
+            let base = match &columns {
+                crate::columns::Columns::Device {
+                    store,
+                    first,
+                    width,
+                } => store.view(*first, *width),
+                crate::columns::Columns::Host(host) => {
+                    // SAFETY: every cell is written by the copies below.
+                    let mut up = unsafe { alloc_or_trim::<u64>(&stream, host.len() * rows) }?;
+                    for (k, column) in host.iter().enumerate() {
+                        let at = k * rows;
+                        let mut slab = up.slice_mut(at..at + rows);
+                        stream.memcpy_htod(*column, &mut slab)?;
+                    }
+                    uploaded = up;
+                    uploaded.slice(0..host.len() * rows)
+                }
+            };
             let plan_dev = crate::device::htod_or_trim(&stream, plan)?;
             let num_plan = (plan.len() / 3) as u64;
             let rows_arg = rows as u64;
@@ -853,9 +878,9 @@ impl DeviceFactors {
                     .arg(&mut buffer)
                     .launch(cfg)?;
             }
-            // The columns are spent. Freeing them is stream-ordered, so it
-            // happens behind the kernel that just read them.
-            drop(base);
+            // A copy made here is spent and goes at the end of this block,
+            // which is stream-ordered behind the kernel that just read it; a
+            // view of the epoch's columns frees nothing, because they stay.
         }
 
         let addresses: Vec<u64> = {

crypto/math-cuda/src/whir.rs

@@ -193,6 +193,29 @@ pub fn commit_codeword_parts(
     )
 }
 
+/// The same for parts the card already holds: `(column index, offset)` into the
+/// epoch's columns. The scatter is then a copy at device bandwidth rather than
+/// the trace crossing the bus again.
+pub fn commit_codeword_resident(
+    store: &crate::columns::DeviceColumns,
+    parts: &[(usize, usize)],
+    log_evals: usize,
+    log_blowup: usize,
+    log_folding: usize,
+    transient: bool,
+) -> Result<(DeviceCodeword, [u8; 32])> {
+    commit_from(
+        Source::Resident {
+            store,
+            parts,
+            log_evals,
+        },
+        log_blowup,
+        log_folding,
+        transient,
+    )
+}
+
 /// Where a commit's coefficients come from: one slab the host holds, or the
 /// columns a stacked polynomial is made of.
 enum Source<'a> {
@@ -201,13 +224,18 @@ enum Source<'a> {
         parts: &'a [(&'a [u64], usize)],
         log_evals: usize,
     },
+    Resident {
+        store: &'a crate::columns::DeviceColumns,
+        parts: &'a [(usize, usize)],
+        log_evals: usize,
+    },
 }
 
 impl Source<'_> {
     fn log_evals(&self) -> u64 {
         match self {
             Self::Whole(evals) => evals.len().trailing_zeros() as u64,
-            Self::Parts { log_evals, .. } => *log_evals as u64,
+            Self::Parts { log_evals, .. } | Self::Resident { log_evals, .. } => *log_evals as u64,
         }
     }
 
@@ -225,6 +253,13 @@ impl Source<'_> {
                 }
                 Ok(())
             }
+            Self::Resident { store, parts, .. } => {
+                stream.memset_zeros(coeffs)?;
+                for (column, offset) in *parts {
+                    store.copy_into(*column, coeffs, *offset, stream)?;
+                }
+                Ok(())
+            }
         }
     }
 }