feat: update gpu tracegen for boundary chip

Author: 876pol

crates/vm/cuda/src/system/boundary.cu

@@ -6,12 +6,13 @@
 #include <cassert>
 
 inline constexpr size_t PERSISTENT_CHUNK = 8;
+inline constexpr size_t BLOCKS_PER_CHUNK = 2;
 inline constexpr size_t VOLATILE_CHUNK = 1;
 
-template <size_t CHUNK> struct BoundaryRecord {
+template <size_t CHUNK, size_t BLOCKS> struct BoundaryRecord {
     uint32_t address_space;
     uint32_t ptr;
-    uint32_t timestamp;
+    uint32_t timestamps[BLOCKS];
     uint32_t values[CHUNK];
 };
 
@@ -21,7 +22,7 @@ template <typename T> struct PersistentBoundaryCols {
     T leaf_label;
     T values[PERSISTENT_CHUNK];
     T hash[PERSISTENT_CHUNK];
-    T timestamp;
+    T timestamps[BLOCKS_PER_CHUNK];
 };
 
 inline constexpr size_t ADDR_ELTS = 2;
@@ -42,7 +43,7 @@ __global__ void cukernel_persistent_boundary_tracegen(
     size_t height,
     size_t width,
     uint8_t const *const *initial_mem,
-    BoundaryRecord<PERSISTENT_CHUNK> *records,
+    BoundaryRecord<PERSISTENT_CHUNK, BLOCKS_PER_CHUNK> *records,
     size_t num_records,
     FpArray<16> *poseidon2_buffer,
     uint32_t *poseidon2_buffer_idx,
@@ -53,7 +54,7 @@ __global__ void cukernel_persistent_boundary_tracegen(
     RowSlice row = RowSlice(trace + row_idx, height);
 
     if (record_idx < num_records) {
-        BoundaryRecord<PERSISTENT_CHUNK> record = records[record_idx];
+        BoundaryRecord<PERSISTENT_CHUNK, BLOCKS_PER_CHUNK> record = records[record_idx];
         Poseidon2Buffer poseidon2(poseidon2_buffer, poseidon2_buffer_idx, poseidon2_capacity);
         COL_WRITE_VALUE(row, PersistentBoundaryCols, address_space, record.address_space);
         COL_WRITE_VALUE(row, PersistentBoundaryCols, leaf_label, record.ptr / PERSISTENT_CHUNK);
@@ -77,24 +78,32 @@ __global__ void cukernel_persistent_boundary_tracegen(
             }
             FpArray<8> init_hash = poseidon2.hash_and_record(init_values);
             COL_WRITE_VALUE(row, PersistentBoundaryCols, expand_direction, Fp::one());
-            COL_WRITE_VALUE(row, PersistentBoundaryCols, timestamp, Fp::zero());
             COL_WRITE_ARRAY(
                 row, PersistentBoundaryCols, values, reinterpret_cast<Fp const *>(init_values.v)
             );
             COL_WRITE_ARRAY(
                 row, PersistentBoundaryCols, hash, reinterpret_cast<Fp const *>(init_hash.v)
             );
+            Fp ts_values[BLOCKS_PER_CHUNK];
+            for (int i = 0; i < BLOCKS_PER_CHUNK; ++i) {
+                ts_values[i] = Fp::zero();
+            }
+            COL_WRITE_ARRAY(row, PersistentBoundaryCols, timestamps, ts_values);
         } else {
             FpArray<8> final_values = FpArray<8>::from_raw_array(record.values);
             FpArray<8> final_hash = poseidon2.hash_and_record(final_values);
             COL_WRITE_VALUE(row, PersistentBoundaryCols, expand_direction, Fp::neg_one());
-            COL_WRITE_VALUE(row, PersistentBoundaryCols, timestamp, record.timestamp);
             COL_WRITE_ARRAY(
                 row, PersistentBoundaryCols, values, reinterpret_cast<Fp const *>(final_values.v)
             );
             COL_WRITE_ARRAY(
                 row, PersistentBoundaryCols, hash, reinterpret_cast<Fp const *>(final_hash.v)
             );
+            Fp ts_values[BLOCKS_PER_CHUNK];
+            for (int i = 0; i < BLOCKS_PER_CHUNK; ++i) {
+                ts_values[i] = Fp(record.timestamps[i]);
+            }
+            COL_WRITE_ARRAY(row, PersistentBoundaryCols, timestamps, ts_values);
         }
     } else {
         row.fill_zero(0, width);
@@ -105,7 +114,7 @@ __global__ void cukernel_volatile_boundary_tracegen(
     Fp *trace,
     size_t height,
     size_t width,
-    BoundaryRecord<VOLATILE_CHUNK> const *records,
+    BoundaryRecord<VOLATILE_CHUNK, 1> const *records,
     size_t num_records,
     uint32_t *range_checker,
     size_t range_checker_num_bins,
@@ -122,7 +131,7 @@ __global__ void cukernel_volatile_boundary_tracegen(
             // For the sake of always filling `addr_lt_aux`
             row.fill_zero(0, width);
         }
-        BoundaryRecord<VOLATILE_CHUNK> record = records[idx];
+        BoundaryRecord<VOLATILE_CHUNK, 1> record = records[idx];
         rc.decompose(
             record.address_space,
             as_max_bits,
@@ -137,11 +146,11 @@ __global__ void cukernel_volatile_boundary_tracegen(
         );
         COL_WRITE_VALUE(row, VolatileBoundaryCols, initial_data, Fp::zero());
         COL_WRITE_VALUE(row, VolatileBoundaryCols, final_data, record.values[0]);
-        COL_WRITE_VALUE(row, VolatileBoundaryCols, final_timestamp, record.timestamp);
+        COL_WRITE_VALUE(row, VolatileBoundaryCols, final_timestamp, record.timestamps[0]);
         COL_WRITE_VALUE(row, VolatileBoundaryCols, is_valid, Fp::one());
 
         if (idx != num_records - 1) {
-            BoundaryRecord<VOLATILE_CHUNK> next_record = records[idx + 1];
+            BoundaryRecord<VOLATILE_CHUNK, 1> next_record = records[idx + 1];
             uint32_t curr[ADDR_ELTS] = {record.address_space, record.ptr};
             uint32_t next[ADDR_ELTS] = {next_record.address_space, next_record.ptr};
             IsLessThanArray::generate_subrow(
@@ -189,8 +198,8 @@ extern "C" int _persistent_boundary_tracegen(
     size_t poseidon2_capacity
 ) {
     auto [grid, block] = kernel_launch_params(height);
-    BoundaryRecord<PERSISTENT_CHUNK> *d_records =
-        reinterpret_cast<BoundaryRecord<PERSISTENT_CHUNK> *>(d_raw_records);
+    BoundaryRecord<PERSISTENT_CHUNK, BLOCKS_PER_CHUNK> *d_records =
+        reinterpret_cast<BoundaryRecord<PERSISTENT_CHUNK, BLOCKS_PER_CHUNK> *>(d_raw_records);
     FpArray<16> *d_poseidon2_buffer = reinterpret_cast<FpArray<16> *>(d_poseidon2_raw_buffer);
     cukernel_persistent_boundary_tracegen<<<grid, block>>>(
         d_trace,
@@ -218,7 +227,7 @@ extern "C" int _volatile_boundary_tracegen(
     size_t ptr_max_bits
 ) {
     auto [grid, block] = kernel_launch_params(height, 512);
-    auto d_records = reinterpret_cast<BoundaryRecord<VOLATILE_CHUNK> const *>(d_raw_records);
+    auto d_records = reinterpret_cast<BoundaryRecord<VOLATILE_CHUNK, 1> const *>(d_raw_records);
     cukernel_volatile_boundary_tracegen<<<grid, block>>>(
         d_trace,
         height,

crates/vm/src/arch/testing/cuda.rs

@@ -54,10 +54,10 @@ use crate::{
             POSEIDON2_DIRECT_BUS, READ_INSTRUCTION_BUS,
         },
         Arena, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState, MatrixRecordArena,
-        MemoryConfig, PreflightExecutor, Streams, VmStateMut,
+        MemoryConfig, PreflightExecutor, Streams, VmStateMut, CONST_BLOCK_SIZE,
     },
     system::{
-        cuda::{poseidon2::Poseidon2PeripheryChipGPU, DIGEST_WIDTH},
+        cuda::poseidon2::Poseidon2PeripheryChipGPU,
         memory::{
             offline_checker::{MemoryBridge, MemoryBus},
             MemoryAirInventory, SharedMemoryHelper,
@@ -314,7 +314,7 @@ impl GpuChipTestBuilder {
         )));
         Self {
             memory: DeviceMemoryTester::persistent(
-                default_tracing_memory(&mem_config, DIGEST_WIDTH),
+                default_tracing_memory(&mem_config, CONST_BLOCK_SIZE),
                 mem_bus,
                 mem_config,
                 range_checker.clone(),

crates/vm/src/system/cuda/boundary.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use openvm_circuit::{
+    arch::CONST_BLOCK_SIZE,
     system::memory::{
         persistent::PersistentBoundaryCols, volatile::VolatileBoundaryCols,
         TimestampedEquipartition, TimestampedValues,
@@ -19,7 +20,7 @@ use openvm_stark_backend::{
     Chip,
 };
 
-use super::{merkle_tree::TIMESTAMPED_BLOCK_WIDTH, poseidon2::SharedBuffer};
+use super::{poseidon2::SharedBuffer, DIGEST_WIDTH};
 use crate::cuda_abi::boundary::{persistent_boundary_tracegen, volatile_boundary_tracegen};
 
 pub struct PersistentBoundary {
@@ -28,7 +29,7 @@ pub struct PersistentBoundary {
     /// This struct cannot own the device memory, hence we take extra care not to use memory we
     /// don't own. TODO: use `Arc<DeviceBuffer>` instead?
     pub initial_leaves: Vec<*const std::ffi::c_void>,
-    pub touched_blocks: Option<DeviceBuffer<u32>>,
+    pub records: Option<DeviceBuffer<u32>>,
 }
 
 pub struct VolatileBoundary {
@@ -49,13 +50,24 @@ pub struct BoundaryChipGPU {
     pub trace_width: Option<usize>,
 }
 
+const BLOCKS_PER_CHUNK: usize = DIGEST_WIDTH / CONST_BLOCK_SIZE;
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct PersistentBoundaryRecord {
+    pub address_space: u32,
+    pub ptr: u32,
+    pub timestamps: [u32; BLOCKS_PER_CHUNK],
+    pub values: [F; DIGEST_WIDTH],
+}
+
 impl BoundaryChipGPU {
     pub fn persistent(poseidon2_buffer: SharedBuffer<F>) -> Self {
         Self {
             fields: BoundaryFields::Persistent(PersistentBoundary {
                 poseidon2_buffer,
                 initial_leaves: Vec::new(),
-                touched_blocks: None,
+                records: None,
             }),
             num_records: None,
             trace_width: None,
@@ -106,14 +118,18 @@ impl BoundaryChipGPU {
 
     pub fn finalize_records_persistent<const CHUNK: usize>(
         &mut self,
-        touched_blocks: DeviceBuffer<u32>,
+        records: Vec<PersistentBoundaryRecord>,
     ) {
         match &mut self.fields {
             BoundaryFields::Volatile(_) => panic!("call `finalize_records_volatile`"),
             BoundaryFields::Persistent(fields) => {
-                self.num_records = Some(touched_blocks.len() / TIMESTAMPED_BLOCK_WIDTH);
+                self.num_records = Some(records.len());
                 self.trace_width = Some(PersistentBoundaryCols::<F, CHUNK>::width());
-                fields.touched_blocks = Some(touched_blocks);
+                fields.records = Some(if records.is_empty() {
+                    DeviceBuffer::new()
+                } else {
+                    records.to_device().unwrap().as_buffer::<u32>()
+                });
             }
         }
     }
@@ -144,7 +160,7 @@ impl<RA> Chip<RA, GpuBackend> for BoundaryChipGPU {
                         trace.height(),
                         trace.width(),
                         &mem_ptrs,
-                        boundary.touched_blocks.as_ref().unwrap(),
+                        boundary.records.as_ref().unwrap(),
                         num_records,
                         &boundary.poseidon2_buffer.buffer,
                         &boundary.poseidon2_buffer.idx,

crates/vm/src/system/cuda/memory.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{collections::BTreeMap, sync::Arc};
 
 use openvm_circuit::{
     arch::{
@@ -12,7 +12,7 @@ use openvm_circuit::{
 use openvm_circuit_primitives::var_range::VariableRangeCheckerChipGPU;
 use openvm_cuda_backend::{prover_backend::GpuBackend, types::F};
 use openvm_cuda_common::{
-    copy::{cuda_memcpy, MemCopyD2D, MemCopyH2D},
+    copy::{cuda_memcpy, MemCopyH2D},
     d_buffer::DeviceBuffer,
     memory_manager::MemTracker,
 };
@@ -22,8 +22,8 @@ use openvm_stark_backend::{
 
 use super::{
     access_adapters::AccessAdapterInventoryGPU,
-    boundary::{BoundaryChipGPU, BoundaryFields},
-    merkle_tree::{MemoryMerkleTree, TIMESTAMPED_BLOCK_WIDTH},
+    boundary::{BoundaryChipGPU, BoundaryFields, PersistentBoundaryRecord},
+    merkle_tree::{MemoryMerkleTree, MERKLE_TOUCHED_BLOCK_WIDTH, TIMESTAMPED_BLOCK_WIDTH},
     Poseidon2PeripheryChipGPU, DIGEST_WIDTH,
 };
 
@@ -151,7 +151,34 @@ impl MemoryInventoryGPU {
                 }
 
                 mem.tracing_info("boundary finalize");
-                let (touched_memory, empty) = if partition.is_empty() {
+                let read_chunk_from_device_raw =
+                    |addr_space: u32, chunk_ptr: u32| -> [F; DIGEST_WIDTH] {
+                        let mut res = [F::ZERO; DIGEST_WIDTH];
+                        let addr_space_idx = addr_space as usize;
+                        let d_mem = &persistent.initial_memory[addr_space_idx];
+                        if d_mem.is_empty() {
+                            return res;
+                        }
+                        let layout = &persistent.merkle_tree.mem_config().addr_spaces[addr_space_idx]
+                            .layout;
+                        let one_cell_size = layout.size();
+                        let offset = chunk_ptr as usize * one_cell_size;
+                        let mut values = vec![0u8; one_cell_size * DIGEST_WIDTH];
+                        unsafe {
+                            cuda_memcpy::<true, false>(
+                                values.as_mut_ptr() as *mut std::ffi::c_void,
+                                d_mem.as_ptr().add(offset) as *const std::ffi::c_void,
+                                values.len(),
+                            )
+                            .unwrap();
+                            for i in 0..DIGEST_WIDTH {
+                                res[i] = layout.to_field::<F>(&values[i * one_cell_size..]);
+                            }
+                        }
+                        res
+                    };
+                let partition_is_empty = partition.is_empty();
+                let (touched_memory, empty) = if partition_is_empty {
                     let leftmost_values = 'left: {
                         let mut res = [F::ZERO; CONST_BLOCK_SIZE];
                         if persistent.initial_memory[ADDR_SPACE_OFFSET as usize].is_empty() {
@@ -190,24 +217,87 @@ impl MemoryInventoryGPU {
                 } else {
                     (partition, false)
                 };
+                let merkle_touched_memory = {
+                    let mut chunk_map: BTreeMap<(u32, u32), (u32, [F; DIGEST_WIDTH])> =
+                        BTreeMap::new();
+                    for &((addr_space, ptr), ts_values) in touched_memory.iter() {
+                        let chunk_ptr = (ptr / DIGEST_WIDTH as u32) * DIGEST_WIDTH as u32;
+                        let block_idx_in_chunk =
+                            ((ptr % DIGEST_WIDTH as u32) / CONST_BLOCK_SIZE as u32) as usize;
+                        let entry = chunk_map.entry((addr_space, chunk_ptr)).or_insert_with(|| {
+                            (
+                                0u32,
+                                read_chunk_from_device_raw(addr_space, chunk_ptr),
+                            )
+                        });
+                        entry.0 = entry.0.max(ts_values.timestamp);
+                        for (i, val) in ts_values.values.iter().copied().enumerate() {
+                            entry.1[block_idx_in_chunk * CONST_BLOCK_SIZE + i] = val;
+                        }
+                    }
+                    chunk_map
+                        .into_iter()
+                        .map(|((addr_space, ptr), (timestamp, values))| {
+                            (
+                                (addr_space, ptr),
+                                TimestampedValues { timestamp, values },
+                            )
+                        })
+                        .collect::<Vec<_>>()
+                };
+                let boundary_records = if partition_is_empty {
+                    Vec::new()
+                } else {
+                    let mut chunk_map: BTreeMap<
+                        (u32, u32),
+                        ([u32; DIGEST_WIDTH / CONST_BLOCK_SIZE], [F; DIGEST_WIDTH]),
+                    > = BTreeMap::new();
+                    for &((addr_space, ptr), ts_values) in touched_memory.iter() {
+                        let chunk_ptr = (ptr / DIGEST_WIDTH as u32) * DIGEST_WIDTH as u32;
+                        let block_idx_in_chunk =
+                            ((ptr % DIGEST_WIDTH as u32) / CONST_BLOCK_SIZE as u32) as usize;
+                        let entry = chunk_map.entry((addr_space, chunk_ptr)).or_insert_with(|| {
+                            (
+                                [0u32; DIGEST_WIDTH / CONST_BLOCK_SIZE],
+                                read_chunk_from_device_raw(addr_space, chunk_ptr),
+                            )
+                        });
+                        entry.0[block_idx_in_chunk] = ts_values.timestamp;
+                        for (i, val) in ts_values.values.iter().copied().enumerate() {
+                            entry.1[block_idx_in_chunk * CONST_BLOCK_SIZE + i] = val;
+                        }
+                    }
+                    chunk_map
+                        .into_iter()
+                        .map(|((addr_space, ptr), (timestamps, values))| {
+                            PersistentBoundaryRecord {
+                                address_space: addr_space,
+                                ptr,
+                                timestamps,
+                                values,
+                            }
+                        })
+                        .collect::<Vec<_>>()
+                };
                 debug_assert_eq!(
                     size_of_val(&touched_memory[0]),
                     TIMESTAMPED_BLOCK_WIDTH * size_of::<u32>()
                 );
-                let d_touched_memory = touched_memory.to_device().unwrap().as_buffer::<u32>();
-                if empty {
-                    self.boundary
-                        .finalize_records_persistent::<DIGEST_WIDTH>(DeviceBuffer::new());
-                } else {
-                    self.boundary.finalize_records_persistent::<DIGEST_WIDTH>(
-                        d_touched_memory.device_copy().unwrap().as_buffer::<u32>(),
-                    ); // TODO do not copy
-                }
+                let d_merkle_touched_memory = merkle_touched_memory
+                    .to_device()
+                    .unwrap()
+                    .as_buffer::<u32>();
+                debug_assert_eq!(
+                    size_of_val(&merkle_touched_memory[0]),
+                    MERKLE_TOUCHED_BLOCK_WIDTH * size_of::<u32>()
+                );
+                self.boundary
+                    .finalize_records_persistent::<DIGEST_WIDTH>(boundary_records);
                 mem.tracing_info("merkle update");
                 persistent.merkle_tree.finalize();
                 let merkle_tree_ctx = persistent.merkle_tree.update_with_touched_blocks(
                     unpadded_merkle_height,
-                    &d_touched_memory,
+                    &d_merkle_touched_memory,
                     empty,
                 );
                 Some(merkle_tree_ctx)