Update blake commit_layer to work with DeviceVecs

crates/prover/Cargo.toml

@@ -28,10 +28,10 @@ tracing.workspace = true
 rayon = { version = "1.10.0", optional = true }
 serde = { version = "1.0", features = ["derive"] }
 
-icicle-cuda-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="ff80aea90686d2001989f0cb5c7a0ed652c395ae"}
-icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="ff80aea90686d2001989f0cb5c7a0ed652c395ae"}
-icicle-m31 = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="ff80aea90686d2001989f0cb5c7a0ed652c395ae"}
-icicle-hash = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="ff80aea90686d2001989f0cb5c7a0ed652c395ae"}
+icicle-cuda-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="6c74c54997a104702debe24fac68ab51ec3fe154"}
+icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="6c74c54997a104702debe24fac68ab51ec3fe154"}
+icicle-m31 = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="6c74c54997a104702debe24fac68ab51ec3fe154"}
+icicle-hash = { git = "https://github.com/ingonyama-zk/icicle.git", optional = true, rev="6c74c54997a104702debe24fac68ab51ec3fe154"}
 
 nvtx = { version = "*", optional = true } 
 

crates/prover/src/core/backend/icicle/blake2s.rs

@@ -1,10 +1,14 @@
 use std::cmp::Reverse;
+use std::mem::transmute;
+use std::fmt::Debug;
+use std::ops::Deref;
 
 use icicle_core::tree::{merkle_tree_digests_len, TreeBuilderConfig};
 use icicle_core::Matrix;
-use icicle_cuda_runtime::memory::HostSlice;
-use icicle_hash::blake2s::build_blake2s_mmcs;
+use icicle_hash::blake2s::blake2s_commit_layer;
+use icicle_cuda_runtime::memory::{HostSlice, DeviceVec, DeviceSlice, HostOrDeviceSlice};
 use itertools::Itertools;
+use icicle_core::vec_ops::{are_bytes_equal, VecOpsConfig};
 
 use super::IcicleBackend;
 use crate::core::backend::{BackendForChannel, Col, Column, ColumnOps, CpuBackend};
@@ -13,84 +17,169 @@ use crate::core::vcs::blake2_hash::Blake2sHash;
 use crate::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
 use crate::core::vcs::ops::{MerkleHasher, MerkleOps};
 
+impl BackendForChannel<Blake2sMerkleChannel> for IcicleBackend {}
+
 impl ColumnOps<Blake2sHash> for IcicleBackend {
-    type Column = Vec<Blake2sHash>;
+    type Column = DeviceColumnBlake;
 
     fn bit_reverse_column(_column: &mut Self::Column) {
         unimplemented!()
     }
 }
 
-impl MerkleOps<Blake2sMerkleHasher> for IcicleBackend {
-    const COMMIT_IMPLEMENTED: bool = true;
-
-    fn commit_columns(
-        columns: Vec<&Col<Self, BaseField>>,
-    ) -> Vec<Col<Self, <Blake2sMerkleHasher as MerkleHasher>::Hash>> {
-        let mut config = TreeBuilderConfig::default();
-        config.arity = 2;
-        config.digest_elements = 32;
-        config.sort_inputs = false;
-
-        nvtx::range_push!("[ICICLE] log_max");
-        let log_max = columns
-            .iter()
-            .sorted_by_key(|c| Reverse(c.len()))
-            .next()
-            .unwrap()
-            .len()
-            .ilog2();
-        nvtx::range_pop!();
-        let mut matrices = vec![];
-        nvtx::range_push!("[ICICLE] create matrix");
-        for col in columns.into_iter().sorted_by_key(|c| Reverse(c.len())) {
-            matrices.push(Matrix::from_slice(col, 4, col.len()));
+pub struct DeviceColumnBlake {
+    pub data: DeviceVec<Blake2sHash>,
+    pub length: usize,
+}
+
+impl PartialEq for DeviceColumnBlake {
+    fn eq(&self, other: &Self) -> bool {
+        if self.length != other.length {
+            return false;
         }
-        nvtx::range_pop!();
-        nvtx::range_push!("[ICICLE] merkle_tree_digests_len");
-        let digests_len = merkle_tree_digests_len(log_max as u32, 2, 32);
-        nvtx::range_pop!();
-        let mut digests = vec![0u8; digests_len];
-        let digests_slice = HostSlice::from_mut_slice(&mut digests);
+        let cfg = VecOpsConfig::default();
+        are_bytes_equal::<Blake2sHash>(self.data.deref(), other.data.deref(), &cfg)
+    }
+
+    fn ne(&self, other: &Self) -> bool {
+        !self.eq(other)
+    }
+}
 
-        nvtx::range_push!("[ICICLE] build_blake2s_mmcs");
-        build_blake2s_mmcs(&matrices, digests_slice, &config).unwrap();
-        nvtx::range_pop!();
+impl Clone for DeviceColumnBlake {
+    fn clone(&self) -> Self {
+        let mut data = DeviceVec::cuda_malloc(self.length).unwrap();
+        data.copy_from_device(&self.data);
+        Self{data, length: self.length}
+    }
+}
+
+impl Debug for DeviceColumnBlake {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let data = self.to_cpu();
+        f.debug_struct("DeviceColumnBlake").field("data", &data.as_slice()).field("length", &self.length).finish()
+    }
+}
 
-        let mut digests: &[<Blake2sMerkleHasher as MerkleHasher>::Hash] =
-            unsafe { std::mem::transmute(digests.as_mut_slice()) };
-        // Transmute digests into stwo format
-        let mut layers = vec![];
-        let mut offset = 0usize;
-        nvtx::range_push!("[ICICLE] convert to CPU layer");
-        for log in 0..=log_max {
-            let inv_log = log_max - log;
-            let number_of_rows = 1 << inv_log;
-
-            let mut layer = vec![];
-            layer.extend_from_slice(&digests[offset..offset + number_of_rows]);
-            layers.push(layer);
-
-            if log != log_max {
-                offset += number_of_rows;
-            }
+impl DeviceColumnBlake {
+    pub fn from_cpu(values: &[Blake2sHash]) -> Self {
+        let length = values.len();
+        let mut data: DeviceVec<Blake2sHash> = DeviceVec::cuda_malloc(length).unwrap();
+        data.copy_from_host(HostSlice::from_slice(&values));
+        Self{data, length}
+    }
+
+    pub fn len(&self) -> usize {
+        self.length
+    }
+}
+
+impl Column<Blake2sHash> for DeviceColumnBlake {
+    fn zeros(length: usize) -> Self {
+        let mut data = DeviceVec::cuda_malloc(length).unwrap();
+
+        let host_data = vec![Blake2sHash::default(); length];
+        data.copy_from_host(HostSlice::from_slice(&host_data));
+
+        Self { data, length }
+    }
+
+    #[allow(clippy::uninit_vec)]
+    unsafe fn uninitialized(length: usize) -> Self {
+        let mut data = DeviceVec::cuda_malloc(length).unwrap();
+        Self { data, length }
+    }
+
+    fn to_cpu(&self) -> Vec<Blake2sHash> {
+        let mut host_data = Vec::<Blake2sHash>::with_capacity(self.length);
+        self.data
+            .copy_to_host(HostSlice::from_mut_slice(&mut host_data));
+        host_data
+    }
+
+    fn len(&self) -> usize {
+        self.length
+    }
+
+    fn at(&self, index: usize) -> Blake2sHash {
+        let mut host_vec = vec![Blake2sHash::default(); 1];
+        unsafe {
+            DeviceSlice::from_slice(std::slice::from_raw_parts(self.data.as_ptr().add(index), 1))
+                .copy_to_host(HostSlice::from_mut_slice(&mut host_vec))
+                .unwrap();
         }
+        host_vec[0]
+    }
 
-        layers.reverse();
-        nvtx::range_pop!();
-        layers
+    fn set(&mut self, index: usize, value: Blake2sHash) {
+        let host_vec = vec![value; 1];
+        unsafe {
+            DeviceSlice::from_mut_slice(std::slice::from_raw_parts_mut(
+                self.data.as_mut_ptr().add(index),
+                1,
+            ))
+            .copy_from_host(HostSlice::from_slice(&host_vec))
+            .unwrap();
+        }
+    }
+}
+
+impl FromIterator<Blake2sHash> for DeviceColumnBlake {
+    fn from_iter<I: IntoIterator<Item = Blake2sHash>>(iter: I) -> Self {
+        let host_data = iter.into_iter().collect_vec();
+        let length = host_data.len();
+        let mut data = DeviceVec::cuda_malloc(length).unwrap();
+        data.copy_from_host(HostSlice::from_slice(&host_data))
+            .unwrap();
+
+        Self { data, length }
     }
+}
 
+impl MerkleOps<Blake2sMerkleHasher> for IcicleBackend {
     fn commit_on_layer(
         log_size: u32,
         prev_layer: Option<&Col<Self, <Blake2sMerkleHasher as MerkleHasher>::Hash>>,
         columns: &[&Col<Self, BaseField>],
     ) -> Col<Self, <Blake2sMerkleHasher as MerkleHasher>::Hash> {
-        // todo!()
-        <CpuBackend as MerkleOps<Blake2sMerkleHasher>>::commit_on_layer(
-            log_size, prev_layer, columns,
-        )
+        nvtx::range_push!("[ICICLE] Extract prev_layer");
+        let prev_layer = match prev_layer {
+            Some(layer) => layer,
+            // Hacky, since creating a DeviceVec of size 0 seems to not work
+            // NOTE: blake2s_commit_layer uses a length of 1 as an indicator that
+            // the prev_layer does not exist
+            None => unsafe { &<Col<Self, <Blake2sMerkleHasher as MerkleHasher>::Hash> as Column<Blake2sHash>>::uninitialized(1) },
+        };
+        nvtx::range_pop!();
+
+        nvtx::range_push!("[ICICLE] Create matrices");
+        let mut columns_as_matrices = vec![];
+        for &col in columns {
+          let col_as_slice = col.data[..].as_slice();
+          columns_as_matrices.push(Matrix::from_slice(&col_as_slice, 4, col.len()));
+        }
+        nvtx::range_pop!();
+
+        nvtx::range_push!("[ICICLE] Cuda malloc digests");
+        let digests_bytes = (1 << log_size) * 32;
+        let mut d_digests_slice = DeviceVec::cuda_malloc(digests_bytes).unwrap();
+        nvtx::range_pop!();
+
+        nvtx::range_push!("[ICICLE] cuda commit layer");
+        blake2s_commit_layer(
+            &(unsafe { transmute::<&DeviceVec<Blake2sHash>, &DeviceVec<u8>>(&prev_layer.data) })[..],
+            false,
+            &columns_as_matrices,
+            false,
+            columns.len() as u32,
+            1 << log_size,
+            &mut d_digests_slice[..],
+        ).unwrap();
+        nvtx::range_pop!();
+
+        DeviceColumnBlake {
+            data: unsafe { transmute(d_digests_slice) },
+            length: 1 << log_size,
+        }
     }
 }
-
-impl BackendForChannel<Blake2sMerkleChannel> for IcicleBackend {}

crates/prover/src/core/backend/icicle/poseidon252.rs

@@ -15,8 +15,6 @@ impl ColumnOps<FieldElement> for IcicleBackend {
 }
 
 impl MerkleOps<Poseidon252MerkleHasher> for IcicleBackend {
-    const COMMIT_IMPLEMENTED: bool = false;
-
     fn commit_on_layer(
         log_size: u32,
         prev_layer: Option<&Col<Self, <Poseidon252MerkleHasher as MerkleHasher>::Hash>>,