cleanup

Author: ColoCarletti

crypto/crypto/src/merkle_tree/merkle.rs

@@ -168,10 +168,9 @@ where
         })
     }
 
-    /// Create a root-only Merkle tree placeholder: stores the commitment root
-    /// but no nodes. Used when the tree's authentication paths are gathered from
-    /// a device-resident copy (GPU) instead of this host tree, so the host nodes
-    /// are never materialised (saving the full-tree Device→Host copy).
+    /// Create a root only Merkle tree placeholder: stores the commitment root
+    /// but no nodes. Used when paths are gathered from a device resident copy
+    /// (GPU) instead of this host tree, so the host nodes are never built.
     /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it.
     pub fn from_root(root: B::Node) -> Self {
         MerkleTree {
@@ -255,9 +254,9 @@ where
     /// For example, give me an inclusion proof for the 3rd element in the
     /// Merkle tree
     pub fn get_proof_by_pos(&self, pos: usize) -> Option<Proof<B::Node>> {
-        // A root-only tree (from `from_root`) has no nodes to walk — callers
-        // must gather paths from the device-resident copy instead. Catch the
-        // misuse early in debug builds rather than returning a misleading `None`.
+        // A root only tree (from `from_root`) has no nodes to walk. Callers must
+        // gather paths from the device resident copy instead. Catch the misuse
+        // early in debug builds rather than returning a misleading None.
         debug_assert!(
             !self.nodes.is_empty(),
             "get_proof_by_pos called on a root-only MerkleTree (no nodes)"

crypto/math-cuda/src/device.rs

@@ -214,13 +214,13 @@ fn retain_default_mempool(ctx: &CudaContext) {
     }
 }
 
-/// Device VRAM budget (bytes) for table-session admission control.
+/// Device VRAM budget in bytes for table session admission control.
 ///
-/// `LAMBDA_VM_VRAM_BUDGET_MB` overrides it explicitly — used to force-exercise
-/// the throttle in tests/benchmarks. Otherwise it is 80% of total device
-/// memory, leaving headroom for the context, module code, and the pool's
-/// retained blocks. On any query failure it returns `u64::MAX`, which disables
-/// budgeting: admission then falls back to the core-bound chunk size alone.
+/// LAMBDA_VM_VRAM_BUDGET_MB overrides it (used to force the throttle in tests).
+/// Otherwise it is 80% of total device memory, leaving headroom for the
+/// context, module code, and retained pool blocks. Returns u64::MAX on any
+/// query failure, which disables budgeting (chunks fall back to the core bound
+/// size alone).
 fn detect_vram_budget_bytes(ctx: &CudaContext) -> u64 {
     if let Ok(mb) = std::env::var("LAMBDA_VM_VRAM_BUDGET_MB")
         && let Ok(mb) = mb.parse::<u64>()
@@ -255,17 +255,15 @@ impl Backend {
         // before returning), so the tracking is pure overhead. Disable it.
         unsafe { ctx.disable_event_tracking() };
 
-        // Retain freed device memory in the stream-ordered pool for reuse.
+        // Retain freed device memory in the stream ordered pool for reuse.
         //
-        // cudarc routes `CudaStream::alloc*` through `cuMemAllocAsync`, which
-        // draws from the device's default memory pool. That pool's release
-        // threshold defaults to 0, so every freed buffer is handed back to the
-        // OS at the next sync — meaning the prover's large, repeatedly-shaped
-        // LDE / FRI buffers are re-malloc'd from scratch each op. Raising the
-        // threshold to "unbounded" keeps freed blocks resident in the pool so
-        // subsequent allocations of the same size are satisfied without a real
-        // driver allocation. Best-effort: on any error (no pool support, sync
-        // allocator) we silently keep the current behaviour.
+        // cudarc routes CudaStream::alloc* through cuMemAllocAsync, drawing from
+        // the device default memory pool. Its release threshold defaults to 0,
+        // so every freed buffer goes back to the OS at the next sync and the
+        // prover's large LDE/FRI buffers are rebuilt from scratch each op.
+        // Raising the threshold keeps freed blocks in the pool so a same size
+        // allocation skips a real driver allocation. Best effort: on any error
+        // we keep the current behaviour.
         retain_default_mempool(&ctx);
 
         let arith = ctx.load_module(Ptx::from_src(ARITH_PTX))?;

crypto/math-cuda/src/fri.rs

@@ -215,7 +215,7 @@ impl FriCommitState {
         };
 
         // Keep the layer tree resident on device; copy only the 32-byte root so
-        // R4 query openings gather paths on device instead of D2H'ing the tree.
+        // R4 query openings gather paths on device instead of copying the tree.
         let mut root = [0u8; 32];
         self.stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?;
         self.stream.synchronize()?;

crypto/math-cuda/src/lde.rs

@@ -220,10 +220,9 @@ fn launch_pointwise_mul_batched(
 /// Layout: `m` columns, each `lde_size` u64s, column `c` at byte offset
 /// `c * lde_size * 8` within `buf`. Freed when `buf` Arc drops.
 ///
-/// `tree` optionally carries the main-trace Merkle tree kept resident on device
-/// (populated by the keep path), so R4 query openings can gather authentication
-/// paths on device instead of D2H'ing the whole tree to host. `None` on the CPU
-/// path or when the tree wasn't retained.
+/// `tree` optionally carries the main trace Merkle tree kept resident on device
+/// (the keep path), so R4 query openings gather paths on device instead of
+/// copying the whole tree to host. None on the CPU path.
 #[derive(Clone)]
 pub struct GpuLdeBase {
     pub buf: Arc<CudaSlice<u64>>,
@@ -240,24 +239,22 @@ pub struct GpuLdeExt3 {
     pub buf: Arc<CudaSlice<u64>>,
     pub m: usize,
     pub lde_size: usize,
-    /// Optionally the aux/composition Merkle tree kept resident on device (the
-    /// keep path), so R4 openings gather paths on device instead of D2H'ing the
-    /// whole tree. `None` on the CPU path or when not retained.
+    /// Optionally the aux or composition Merkle tree kept resident on device
+    /// (the keep path), so R4 openings gather paths on device. None otherwise.
     pub tree: Option<GpuMerkleTree>,
 }
 
-/// Handle to a Merkle tree kept live on device after a commit, so query
-/// openings can gather authentication paths on device instead of D2H'ing the
-/// whole tree to host. Node layout matches the CPU tree
-/// (`crypto/crypto/src/merkle_tree`): `nodes[0..leaves_len-1]` are inner nodes
-/// (root at index 0), `nodes[leaves_len-1..]` are the leaves; each node is 32
+/// Merkle tree kept resident on device after a commit, so query openings gather
+/// paths on device instead of copying the whole tree to host. Node layout
+/// matches the CPU tree (`crypto/crypto/src/merkle_tree`): `nodes[0..leaves_len-1]`
+/// are inner nodes (root at 0), `nodes[leaves_len-1..]` are the leaves, each 32
 /// bytes. Freed when the `nodes` Arc drops.
 #[derive(Clone)]
 pub struct GpuMerkleTree {
     pub nodes: Arc<CudaSlice<u8>>,
     pub leaves_len: usize,
-    /// The Merkle root (node 0), copied to host at build time (32 bytes) so the
-    /// commitment is available without D2H'ing the whole tree.
+    /// The Merkle root (node 0), copied to host at build time so the commitment
+    /// is available without copying the whole tree.
     pub root: [u8; 32],
 }
 
@@ -643,41 +640,11 @@ pub fn coset_lde_batch_base_into_with_leaf_hash(
     .map(|_| ())
 }
 
-/// Like `coset_lde_batch_base_into_with_leaf_hash`, but also builds the full
-/// Merkle tree on device and returns the `2*lde_size - 1` node buffer back
-/// to the caller in `merkle_nodes_out` (byte length `(2*lde_size - 1) * 32`).
-///
-/// The leaf hashes are never exposed to the caller — they stay on device and
-/// feed straight into the pair-hash tree kernel, avoiding the
-/// pinned→pageable→pinned round-trip that the separate-step GPU tree build
-/// would pay.
-pub fn coset_lde_batch_base_into_with_merkle_tree(
-    columns: &[&[u64]],
-    blowup_factor: usize,
-    weights: &[u64],
-    outputs: &mut [&mut [u64]],
-    merkle_nodes_out: &mut [u8],
-) -> Result<()> {
-    coset_lde_batch_base_into_with_merkle_tree_inner(
-        columns,
-        blowup_factor,
-        weights,
-        outputs,
-        Some(merkle_nodes_out),
-        KeccakCommit::FullTree,
-        false,
-        false,
-    )
-    .map(|_| ())
-}
-
-/// Fused LDE + leaf-hash + Merkle tree build, keeping **both** the LDE buffer
-/// and the Merkle tree resident on device (returned in `GpuLdeBase` + its
-/// `tree`, including the root). The host tree is **not** materialised — there is
-/// no `merkle_nodes_out` — so the whole-tree Device→Host copy is eliminated.
-/// Callers gather query authentication paths from the device tree instead
-/// (`crate::merkle::gather_merkle_paths_dev`) and use a root-only host tree for
-/// the commitment.
+/// Fused LDE, leaf hash, and Merkle tree build that keeps both the LDE buffer
+/// and the Merkle tree (with root) resident on device, returned in `GpuLdeBase`.
+/// The host tree is not built (no `merkle_nodes_out`), so the whole tree copy to
+/// host is eliminated. Callers gather query paths from the device tree
+/// (`crate::merkle::gather_merkle_paths_dev`) and use a root only host tree.
 pub fn coset_lde_batch_base_into_with_merkle_tree_keep(
     columns: &[&[u64]],
     blowup_factor: usize,
@@ -840,32 +807,32 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner(
         crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, lde_size)?;
     }
 
-    // D2H the LDE columns via pinned staging (still read on host by constraint
-    // eval / openings). The full tree nodes are D2H'd only when the caller asks
-    // for a host tree (`nodes_out`); the device-resident path passes `None` and
-    // keeps the tree on device, eliminating the whole-tree D2H.
+    // Copy the LDE columns to host via pinned staging (constraint eval and
+    // openings still read them). The full tree is copied only when the caller
+    // wants a host tree (`nodes_out`); the resident path passes None and keeps
+    // the tree on device.
     stream.memcpy_dtoh(&buf, &mut pinned[..m * lde_size])?;
     if let Some(no) = nodes_out {
         d2h_bytes_via_pinned_hashes(&stream, be, &nodes_dev, no)?;
     }
-    // Ensure the (async) LDE D2H above has landed before reading `pinned`.
-    // Previously the tree D2H's `synchronize` covered this; the device-resident
-    // path skips that D2H, so synchronize explicitly.
+    // Make sure the async LDE copy above landed before reading `pinned`. The
+    // tree copy used to provide this sync; the resident path skips it, so sync
+    // here.
     stream.synchronize()?;
 
-    // Copy pinned into caller outputs. Runs under the pinned-staging lock,
+    // Copy pinned into caller outputs. Runs under the pinned staging lock,
     // where rayon can deadlock. See `Backend::pinned_staging`.
     for (c, dst) in outputs.iter_mut().enumerate() {
         dst.copy_from_slice(&pinned[c * lde_size..c * lde_size + lde_size]);
     }
     drop(staging);
 
     if keep_device_buf {
-        // Retain the device tree nodes for on-device opening (FullTree only; in
-        // LeavesOnly mode there are no inner nodes to gather paths from).
+        // Retain the device tree for on device opening (FullTree only; LeavesOnly
+        // has no inner nodes to gather paths from).
         let tree = if keep_tree && commit == KeccakCommit::FullTree {
-            // Copy just the 32-byte root (node 0) so the commitment is available
-            // without D2H'ing the whole tree.
+            // Copy just the 32 byte root (node 0) so the commitment is available
+            // without copying the whole tree.
             let mut root = [0u8; 32];
             stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?;
             stream.synchronize()?;
@@ -891,60 +858,10 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner(
     }
 }
 
-/// Ext3 variant of `coset_lde_batch_base_into_with_leaf_hash`: fused
-/// LDE + Keccak-256 leaf hashing over ext3 columns. Thin wrapper over
-/// `coset_lde_batch_ext3_into_with_merkle_tree_inner` with `LeavesOnly`.
-pub fn coset_lde_batch_ext3_into_with_leaf_hash(
-    columns: &[&[u64]],
-    n: usize,
-    blowup_factor: usize,
-    weights: &[u64],
-    outputs: &mut [&mut [u64]],
-    hashed_leaves_out: &mut [u8],
-) -> Result<()> {
-    coset_lde_batch_ext3_into_with_merkle_tree_inner(
-        columns,
-        n,
-        blowup_factor,
-        weights,
-        outputs,
-        Some(hashed_leaves_out),
-        KeccakCommit::LeavesOnly,
-        false,
-        false,
-    )
-    .map(|_| ())
-}
-
-/// Ext3 variant of the fused `coset_lde_batch_base_into_with_merkle_tree`.
-/// LDE + leaf hashing + inner-tree build, all on device; D2Hs only the LDE
-/// evaluations and the full `2*lde_size - 1` node buffer.
-pub fn coset_lde_batch_ext3_into_with_merkle_tree(
-    columns: &[&[u64]],
-    n: usize,
-    blowup_factor: usize,
-    weights: &[u64],
-    outputs: &mut [&mut [u64]],
-    merkle_nodes_out: &mut [u8],
-) -> Result<()> {
-    coset_lde_batch_ext3_into_with_merkle_tree_inner(
-        columns,
-        n,
-        blowup_factor,
-        weights,
-        outputs,
-        Some(merkle_nodes_out),
-        KeccakCommit::FullTree,
-        false,
-        false,
-    )
-    .map(|_| ())
-}
-
 /// Ext3 variant of [`coset_lde_batch_base_into_with_merkle_tree_keep`]: keeps
-/// both the de-interleaved LDE buffer and the Merkle tree (with root) resident
-/// on device. The host tree is **not** materialised (no `merkle_nodes_out`), so
-/// the whole-tree D2H is eliminated; openings gather paths from the device tree.
+/// both the deinterleaved LDE buffer and the Merkle tree (with root) resident on
+/// device. No host tree is built, so the whole tree copy to host is eliminated;
+/// openings gather paths from the device tree.
 pub fn coset_lde_batch_ext3_into_with_merkle_tree_keep(
     columns: &[&[u64]],
     n: usize,
@@ -1103,15 +1020,15 @@ fn coset_lde_batch_ext3_into_with_merkle_tree_inner(
         crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, lde_size)?;
     }
 
-    // D2H the LDE columns. The full tree nodes are D2H'd only when the caller
-    // asks for a host tree (`nodes_out`); the device-resident path passes `None`
-    // and keeps the tree on device, eliminating the whole-tree D2H.
+    // Copy the LDE columns to host. The full tree is copied only when the caller
+    // wants a host tree (`nodes_out`); the resident path passes None and keeps
+    // the tree on device.
     stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?;
     if let Some(no) = nodes_out {
         d2h_bytes_via_pinned_hashes(&stream, be, &nodes_dev, no)?;
     }
-    // Ensure the (async) LDE D2H has landed before reading `pinned` (the tree
-    // D2H's synchronize used to cover this; the device-resident path skips it).
+    // Make sure the async LDE copy landed before reading `pinned` (the tree copy
+    // used to provide this sync; the resident path skips it).
     stream.synchronize()?;
 
     unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size);

crypto/math-cuda/src/merkle.rs

@@ -221,11 +221,10 @@ pub fn build_merkle_tree_on_device(hashed_leaves: &[u8]) -> Result<Vec<u8>> {
 /// Gather Merkle authentication paths on device for `positions` (leaf indices)
 /// against the resident tree `nodes_dev` (standard layout, `2*leaves_len-1`
 /// nodes of 32 bytes). Returns `positions.len() * depth * 32` bytes, where
-/// `depth = log2(leaves_len)`: query `q`'s path occupies bytes
-/// `[q*depth*32 .. (q+1)*depth*32]`, each 32-byte node a sibling from leaf to
-/// root — byte-for-byte the same nodes the CPU `MerkleTree::get_proof_by_pos`
-/// (`build_merkle_path`) collects. Runs on the caller's `stream` (pass the
-/// table's session stream so it shares the queue with the rest of R4).
+/// `depth = log2(leaves_len)`. Query `q`'s path is `[q*depth*32 ..
+/// (q+1)*depth*32]`, each 32 byte node a sibling from leaf to root. These are
+/// the same nodes the CPU `MerkleTree::get_proof_by_pos` collects. Runs on the
+/// caller's `stream` (pass the table's session stream).
 pub fn gather_merkle_paths_dev(
     nodes_dev: &CudaSlice<u8>,
     leaves_len: usize,
@@ -241,9 +240,9 @@ pub fn gather_merkle_paths_dev(
         "leaves_len must be a power of two >= 2"
     );
     let depth = leaves_len.trailing_zeros() as usize;
-    // Guard the kernel's device reads: an out-of-range position would walk off
-    // the node buffer (OOB device read). Positions are valid by construction;
-    // this catches any caller bug before it becomes UB.
+    // Guard the kernel's device reads: a position past leaves_len would walk
+    // off the node buffer. Positions are valid by construction; this catches a
+    // caller bug before it becomes an out of bounds device read.
     assert!(
         positions.iter().all(|&p| (p as usize) < leaves_len),
         "gather_merkle_paths_dev: leaf position >= leaves_len"
@@ -286,9 +285,9 @@ pub fn gather_merkle_paths_dev(
 ///
 /// Returns `(2*(lde_size/2) - 1) * 32` bytes of tree nodes in the standard
 /// layout (root at byte offset 0, leaves in the tail).
-/// Build the composition-poly Merkle tree on device (leaves hash row-pairs, so
-/// `num_leaves = lde_size / 2`). Returns the device node buffer, the leaf count,
-/// and the stream it was built on. Shared by the host-D2H and device-keep
+/// Build the composition Merkle tree on device. Leaves hash row pairs, so
+/// `num_leaves = lde_size / 2`. Returns the device node buffer, the leaf count,
+/// and the stream it was built on. Shared by the host copy and device keep
 /// wrappers below.
 fn build_comp_poly_tree_nodes_dev(
     parts_interleaved: &[&[u64]],
@@ -366,8 +365,8 @@ pub fn build_comp_poly_tree_from_evals_ext3(parts_interleaved: &[&[u64]]) -> Res
 
 /// Like [`build_comp_poly_tree_from_evals_ext3`] but keeps the tree nodes on
 /// device (returned as a [`crate::lde::GpuMerkleTree`] with its root), so R4
-/// composition openings gather authentication paths on device instead of
-/// D2H'ing the whole tree. `leaves_len = lde_size / 2` (row-pair leaves).
+/// composition openings gather paths on device instead of copying the whole
+/// tree to host. `leaves_len = lde_size / 2` (row pair leaves).
 pub fn build_comp_poly_tree_from_evals_ext3_keep(
     parts_interleaved: &[&[u64]],
 ) -> Result<crate::lde::GpuMerkleTree> {

crypto/math-cuda/tests/merkle_gather.rs

@@ -1,7 +1,7 @@
 //! Parity: GPU `gather_merkle_paths_dev` must produce, for each leaf position,
-//! the exact `merkle_path` the CPU `MerkleTree::get_proof_by_pos` returns —
-//! same sibling-from-leaf-to-root order, byte-for-byte. This is the gate for
-//! gathering R4 query openings on device instead of D2H'ing the whole tree.
+//! the exact `merkle_path` the CPU `MerkleTree::get_proof_by_pos` returns: the
+//! same sibling order from leaf to root, byte for byte. This is the gate for
+//! gathering R4 query openings on device instead of copying the whole tree.
 
 use crypto::merkle_tree::backends::field_element_vector::FieldElementVectorBackend;
 use crypto::merkle_tree::merkle::MerkleTree;

crypto/stark/src/fri/fri_commitment.rs

@@ -15,7 +15,7 @@ where
     pub merkle_tree: MerkleTree<B>,
     /// The layer's Merkle tree kept resident on device (GPU FRI commit path),
     /// so R4 query openings gather authentication paths on device. When set,
-    /// `merkle_tree` is a root-only placeholder. `None` on the CPU path.
+    /// `merkle_tree` is a root only placeholder. `None` on the CPU path.
     #[cfg(feature = "cuda")]
     pub gpu_tree: Option<math_cuda::lde::GpuMerkleTree>,
 }