Feat/multi merkle tree gpu resident

crypto/crypto/src/merkle_tree/backends/field_element_vector.rs

@@ -1,7 +1,7 @@
 use core::marker::PhantomData;
 
 use crate::hash::poseidon::Poseidon;
-use crate::merkle_tree::traits::IsMerkleTreeBackend;
+use crate::merkle_tree::traits::{IsLeafHasher, IsMerkleTreeBackend, IsStreamingLeafBackend};
 use alloc::vec::Vec;
 use digest::{Digest, Output};
 use math::{
@@ -202,6 +202,140 @@ where
     }
 }
 
+/// Exposes the streaming leaf routes to callers that reach this backend through
+/// a commitment configuration rather than by name. Both bodies go through
+/// [`hash_streamed`], which is where the absorbed byte layout is defined, so
+/// they agree with `hash_data` by construction.
+impl<F, D: Digest + Send + 'static, const NUM_BYTES: usize> IsStreamingLeafBackend<F>
+    for FieldElementVectorBackend<F, D, NUM_BYTES>
+where
+    F: IsField,
+    FieldElement<F>: AsBytes,
+    [u8; NUM_BYTES]: From<Output<D>>,
+    Vec<FieldElement<F>>: Sync + Send,
+{
+    fn hash_bytes(data: &[u8]) -> [u8; NUM_BYTES] {
+        hash_streamed::<D, NUM_BYTES>(|sink| sink(data))
+    }
+
+    fn hash_data_from_slices(a: &[FieldElement<F>], b: &[FieldElement<F>]) -> [u8; NUM_BYTES] {
+        // A size threshold below which this streams straight through was tried
+        // and MEASURED NEUTRAL-TO-WORSE (963.56M vs 963.28M cycles on a blowup8
+        // verify, with `verify_fri` unmoved to the cycle). The 6.9M `verify_fri`
+        // rise that this change costs is NOT the staging buffer — gating the
+        // buffer away does not recover it — so it is not worth a branch here.
+        // Do not re-add one without a measurement.
+        hash_streamed::<D, NUM_BYTES>(|sink| {
+            let mut stage = LeafStage::new();
+            for element in a.iter().chain(b.iter()) {
+                element.stream_bytes(&mut |bytes| stage.push(bytes, sink));
+            }
+            stage.flush(sink);
+        })
+    }
+
+    type LeafHasher = DigestLeafHasher<F, D, NUM_BYTES>;
+
+    fn leaf_hasher() -> Self::LeafHasher {
+        DigestLeafHasher {
+            hasher: D::new(),
+            phantom: PhantomData,
+        }
+    }
+}
+
+/// [`IsLeafHasher`] over the same digest the one-shot routes use.
+///
+/// The split-invariance the trait demands is inherited rather than argued:
+/// `hash_streamed` opens a fresh `D`, feeds it every element's `stream_bytes`
+/// and finalizes, with no length prefix, padding or framing of its own — so
+/// absorbing the same elements across several `update` calls presents `D` with
+/// the identical byte stream. There is no place for a split to show.
+///
+/// This is a PROVER-side construct: the guest verifier authenticates leaves it
+/// receives whole, through `hash_data_from_slices`.
+pub struct DigestLeafHasher<F, D: Digest, const NUM_BYTES: usize> {
+    hasher: D,
+    /// `fn() -> F` rather than `F`: the field is a type-level label here, never a
+    /// value, and the function-pointer form is unconditionally `Send`/`Sync`. The
+    /// bare `PhantomData<F>` would make every leaf hasher's thread-safety hinge on
+    /// a marker type nobody ever moves.
+    phantom: PhantomData<fn() -> F>,
+}
+
+impl<F, D: Digest, const NUM_BYTES: usize> IsLeafHasher<F> for DigestLeafHasher<F, D, NUM_BYTES>
+where
+    F: IsField,
+    FieldElement<F>: AsBytes,
+    [u8; NUM_BYTES]: From<Output<D>>,
+{
+    type Node = [u8; NUM_BYTES];
+
+    fn update(&mut self, data: &[FieldElement<F>]) {
+        for element in data {
+            element.stream_bytes(&mut |bytes| self.hasher.update(bytes));
+        }
+    }
+
+    fn finalize(self) -> [u8; NUM_BYTES] {
+        let mut result = [0u8; NUM_BYTES];
+        result.copy_from_slice(&self.hasher.finalize());
+        result
+    }
+}
+
+/// Bytes of the leaf staging buffer. Large enough that the run reaching the
+/// hasher is worth batching — 16 blocks.
+const LEAF_STAGE_BYTES: usize = 1024;
+
+/// Coalesces a leaf's field elements into large aligned runs before they reach
+/// the hasher.
+///
+/// A leaf arrives one field element at a time — eight bytes per `stream_bytes`
+/// call — so without staging the hasher only ever sees eight bytes at a time.
+/// Coalescing presents it with fewer, larger `update` calls.
+///
+/// **This cannot change any digest.** The same bytes reach the hasher in the
+/// same order; only the call boundaries move, and the sponge is split-invariant
+/// by construction, so it simply sees fewer, larger `update` calls.
+#[repr(align(8))]
+struct LeafStage {
+    buf: [u8; LEAF_STAGE_BYTES],
+    len: usize,
+}
+
+impl LeafStage {
+    #[inline]
+    fn new() -> Self {
+        Self {
+            buf: [0u8; LEAF_STAGE_BYTES],
+            len: 0,
+        }
+    }
+
+    #[inline]
+    fn push(&mut self, mut bytes: &[u8], sink: &mut dyn FnMut(&[u8])) {
+        while !bytes.is_empty() {
+            if self.len == LEAF_STAGE_BYTES {
+                sink(&self.buf[..LEAF_STAGE_BYTES]);
+                self.len = 0;
+            }
+            let take = (LEAF_STAGE_BYTES - self.len).min(bytes.len());
+            self.buf[self.len..self.len + take].copy_from_slice(&bytes[..take]);
+            self.len += take;
+            bytes = &bytes[take..];
+        }
+    }
+
+    #[inline]
+    fn flush(&mut self, sink: &mut dyn FnMut(&[u8])) {
+        if self.len > 0 {
+            sink(&self.buf[..self.len]);
+            self.len = 0;
+        }
+    }
+}
+
 #[derive(Clone, Default)]
 pub struct BatchPoseidonTree<P: Poseidon + Default> {
     _poseidon: PhantomData<P>,

crypto/crypto/src/merkle_tree/traits.rs

@@ -1,4 +1,7 @@
 use alloc::vec::Vec;
+use math::field::element::FieldElement;
+use math::field::traits::IsField;
+use math::traits::AsBytes;
 #[cfg(feature = "parallel")]
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 
@@ -27,3 +30,71 @@ pub trait IsMerkleTreeBackend {
     /// It will be used in the construction of the Merkle tree.
     fn hash_new_parent(child_1: &Self::Node, child_2: &Self::Node) -> Self::Node;
 }
+
+/// A leaf backend that can hash a leaf without being handed one.
+///
+/// [`IsMerkleTreeBackend::hash_data`] takes a `&Self::Data`, which for the
+/// batched backends is a `Vec<FieldElement<F>>`. Building one per leaf costs an
+/// allocation per leaf — millions on a real trace — so the prover and verifier
+/// never do: they serialize into a reused buffer, or hold two slices they want
+/// hashed as if concatenated. These are the two shapes they use.
+///
+/// Both must agree with `hash_data` on the bytes they absorb, so a leaf hashed
+/// through either route is the leaf the tree was built from. That is the whole
+/// contract, and it is why these live on a trait rather than staying inherent
+/// methods on one concrete backend: a commitment configuration that names its
+/// leaf backend generically still has to reach them.
+pub trait IsStreamingLeafBackend<F>: IsMerkleTreeBackend
+where
+    F: IsField,
+    FieldElement<F>: AsBytes,
+{
+    /// Hash a pre-serialized leaf buffer. Equals `hash_data` applied to the
+    /// elements `data` encodes, in that order.
+    fn hash_bytes(data: &[u8]) -> Self::Node;
+
+    /// Hash `a ‖ b` without materializing the concatenation. Equals
+    /// `hash_data(&[a, b].concat())`.
+    fn hash_data_from_slices(a: &[FieldElement<F>], b: &[FieldElement<F>]) -> Self::Node;
+
+    /// The incremental form of the same leaf hash. See [`IsLeafHasher`].
+    ///
+    /// `Send` because there is one of these per leaf and the base layer of a real
+    /// epoch has millions: absorbing them is parallel across leaves, exactly as
+    /// the one-shot leaf hashing is.
+    type LeafHasher: IsLeafHasher<F, Node = Self::Node> + Send;
+
+    /// A leaf hasher that has absorbed nothing yet.
+    fn leaf_hasher() -> Self::LeafHasher;
+}
+
+/// One leaf's hash, absorbed in an arbitrary number of updates.
+///
+/// [`IsStreamingLeafBackend::hash_data_from_slices`] covers the two-slice case,
+/// which is every leaf the per-table trees hash. A mixed-height MMCS leaf is
+/// different: it concatenates one row pair per matrix at that height, and a
+/// prover that wants to produce those matrices ONE AT A TIME — absorbing each
+/// into the leaves and dropping its buffer — cannot hand over all the slices at
+/// once. This is the API that lets it, and the memory it costs is one hasher
+/// state per leaf rather than one LDE per matrix.
+///
+/// # Contract
+///
+/// Splitting is free: for any partition of a leaf's elements into consecutive
+/// chunks, updating with each chunk in order and finalizing must equal
+/// [`IsMerkleTreeBackend::hash_data`] over the whole. A backend whose framing
+/// depended on where the updates fell would produce leaves no verifier could
+/// re-derive, since the verifier only ever sees the concatenation.
+pub trait IsLeafHasher<F>
+where
+    F: IsField,
+    FieldElement<F>: AsBytes,
+{
+    type Node;
+
+    /// Absorb the next consecutive run of the leaf's elements.
+    fn update(&mut self, data: &[FieldElement<F>]);
+
+    /// Finish the leaf.
+    fn finalize(self) -> Self::Node;
+}

crypto/math-cuda/kernels/deep.cu

@@ -130,3 +130,24 @@ extern "C" __global__ void bit_reverse_ext3_interleaved(
         out[i * 3 + 2] = in[j * 3 + 2];
     }
 }
+
+// Re-interleave the 3 de-interleaved ext3 slabs of `m` columns into the
+// per-column row-major ext3 layout the host consumes: for column `p` and row
+// `r`, out[(p*lde + r)*3 + k] = slab[(p*3 + k)*lde + r]. One thread per
+// (column, row). Moves the host D2H's de-interleave onto the device (the host
+// side dominates the parts download; the GPU does the permute at device
+// bandwidth). Byte-identical — a pure data permutation.
+extern "C" __global__ void interleave_ext3_slabs(
+    const uint64_t *__restrict__ slab,
+    uint64_t *__restrict__ out,
+    uint64_t m,
+    uint64_t lde) {
+    uint64_t idx = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= m * lde) return;
+    uint64_t p = idx / lde;
+    uint64_t r = idx % lde;
+    uint64_t o = (p * lde + r) * 3;
+    out[o + 0] = slab[(p * 3 + 0) * lde + r];
+    out[o + 1] = slab[(p * 3 + 1) * lde + r];
+    out[o + 2] = slab[(p * 3 + 2) * lde + r];
+}

crypto/math-cuda/kernels/fri.cu

@@ -76,3 +76,23 @@ extern "C" __global__ void gather_ext3_at(
     out[i * 3 + 1] = evals[p * 3 + 1];
     out[i * 3 + 2] = evals[p * 3 + 2];
 }
+
+// Batched-FRI injection: `out[i] += beta_sq * bucket[i]` over ext3 elements,
+// interleaved layout. `out` is the just-folded running codeword and `bucket`
+// the shorter DEEP codeword whose height matches this layer; both hold `n` ext3
+// elements (3*n u64). Matches the host `inject_bucket` (fri/batched.rs).
+extern "C" __global__ void fri_inject_bucket_ext3(
+    uint64_t *out,             // 3 * n u64, modified in place
+    const uint64_t *bucket,    // 3 * n u64
+    const uint64_t *beta_sq,   // 3 u64 (ext3)
+    uint64_t n) {
+    uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n) return;
+    ext3::Fe3 v   = ext3::make(out[i * 3], out[i * 3 + 1], out[i * 3 + 2]);
+    ext3::Fe3 b   = ext3::make(bucket[i * 3], bucket[i * 3 + 1], bucket[i * 3 + 2]);
+    ext3::Fe3 bsq = ext3::make(beta_sq[0], beta_sq[1], beta_sq[2]);
+    ext3::Fe3 res = ext3::add(v, ext3::mul(bsq, b));
+    out[i * 3]     = res.a;
+    out[i * 3 + 1] = res.b;
+    out[i * 3 + 2] = res.c;
+}