feat: 2^32 memory address

crates/circuits/primitives/cuda/include/primitives/constants.h

@@ -6,11 +6,14 @@
 namespace openvm {
 inline constexpr size_t BYTE_BITS = 8;
 inline constexpr size_t U16_BITS = 16;
+// log2 of the byte width of one u16-celled storage cell.
+inline constexpr size_t U16_CELL_SIZE_BITS = 1;
 } // namespace openvm
 
 namespace riscv {
 using openvm::BYTE_BITS;
 using openvm::U16_BITS;
+using openvm::U16_CELL_SIZE_BITS;
 
 inline constexpr size_t RV64_REGISTER_NUM_LIMBS = 8;
 inline constexpr size_t RV64_WORD_NUM_LIMBS = 4;

crates/vm/cuda/include/system/memory/address.cuh

@@ -1,6 +1,8 @@
 #pragma once
 
+#include "system/memory/params.cuh"
+
 template <typename T> struct MemoryAddress {
     T address_space;
-    T pointer;
+    T pointer_limbs[POINTER_LIMBS];
 };

crates/vm/cuda/include/system/memory/params.cuh

@@ -43,6 +43,8 @@
 // merkle leaf l starts at ptr l * DIGEST_WIDTH.
 
 #include "poseidon2.cuh" // brings in CELLS / CELLS_OUT from stark-backend
+#include <cstddef>
+#include <cstdint>
 
 // Cells per Poseidon2 half (and per merkle leaf).
 inline constexpr size_t DIGEST_WIDTH = CELLS_OUT;
@@ -60,6 +62,18 @@ inline constexpr size_t BLOCK_FE_WIDTH = MEMORY_BLOCK_BYTES / U16_CELL_SIZE;
 // Blocks per merkle leaf.
 inline constexpr size_t BLOCKS_PER_LEAF = DIGEST_WIDTH / BLOCK_FE_WIDTH;
 
+// Number of little-endian 16-bit limbs used to represent an AS-native memory pointer on the
+// memory bus.
+inline constexpr size_t POINTER_LIMBS = 2;
+inline constexpr size_t POINTER_LIMB_BITS = 16;
+
+// log2 of DIGEST_WIDTH (= CELLS_OUT = 8).
+inline constexpr size_t DIGEST_WIDTH_BITS = 3;
+
+// Number of low bits of a leaf label kept in the low limb so that
+// `low * DIGEST_WIDTH` fits in one 16-bit memory-bus pointer limb.
+inline constexpr size_t LOW_LEAF_BITS = POINTER_LIMB_BITS - DIGEST_WIDTH_BITS;
+
 // Number of bottom merkle levels omitted from large GPU subtree buffers.
 inline constexpr size_t OMITTED_BOTTOM_LEVELS = 3;
 

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

@@ -16,7 +16,7 @@ template <size_t CHUNK, size_t BLOCKS> struct BoundaryRecord {
 template <typename T> struct PersistentBoundaryCols {
     T expand_direction;
     T address_space;
-    T leaf_label;
+    T leaf_label_limbs[POINTER_LIMBS];
     T values[DIGEST_WIDTH];
     T hash[DIGEST_WIDTH];
     T timestamps[BLOCKS_PER_LEAF];
@@ -40,13 +40,13 @@ __global__ void cukernel_persistent_boundary_tracegen(
     if (record_idx < num_records) {
         BoundaryRecord<DIGEST_WIDTH, BLOCKS_PER_LEAF> record = records[record_idx];
         Poseidon2Buffer poseidon2(poseidon2_buffer, poseidon2_buffer_idx, poseidon2_capacity);
+        uint32_t const leaf_label = record.ptr / DIGEST_WIDTH;
+        uint32_t const leaf_label_limbs[POINTER_LIMBS] = {
+            leaf_label & ((uint32_t(1) << LOW_LEAF_BITS) - 1),
+            leaf_label >> LOW_LEAF_BITS,
+        };
         COL_WRITE_VALUE(row, PersistentBoundaryCols, address_space, record.address_space);
-        COL_WRITE_VALUE(
-            row,
-            PersistentBoundaryCols,
-            leaf_label,
-            record.ptr / DIGEST_WIDTH
-        );
+        COL_WRITE_ARRAY(row, PersistentBoundaryCols, leaf_label_limbs, leaf_label_limbs);
         if (row_idx % 2 == 0) {
             FpArray<DIGEST_WIDTH> init_values;
             uint32_t addr_space_idx = record.address_space - 1;

crates/vm/cuda/src/testing/memory.cu

@@ -2,6 +2,9 @@
 #include "primitives/trace_access.h"
 #include "system/memory/address.cuh"
 
+// AS-native memory-bus pointers are two little-endian 16-bit cell-pointer limbs `[lo16, hi16]`, so
+// the dummy interaction row is `[address_space, ptr_lo, ptr_hi, data[BLOCK_SIZE], timestamp,
+// count]` (width `BLOCK_SIZE + 5`), matching the host `MemoryDummyAir`.
 template <typename T, size_t BLOCK_SIZE> struct DummyMemoryInteractionCols {
     MemoryAddress<T> address;
     T data[BLOCK_SIZE];
@@ -15,11 +18,11 @@ __global__ void memory_testing_tracegen(Fp *trace, size_t height, Fp *records, s
     RowSlice row(trace + idx, height);
     if (idx < num_records) {
         auto record = reinterpret_cast<DummyMemoryInteractionCols<Fp, BLOCK_SIZE> *>(records)[idx];
-        COL_WRITE_VALUE(row, MemoryAddress, address_space, record.address.address_space);
-        COL_WRITE_VALUE(row, MemoryAddress, pointer, record.address.pointer);
-        row.write_array(2, BLOCK_SIZE, record.data);
-        row.write(2 + BLOCK_SIZE, record.timestamp);
-        row.write(2 + BLOCK_SIZE + 1, record.count);
+        row.write(0, record.address.address_space);
+        row.write_array(1, POINTER_LIMBS, record.address.pointer_limbs);
+        row.write_array(1 + POINTER_LIMBS, BLOCK_SIZE, record.data);
+        row.write(1 + POINTER_LIMBS + BLOCK_SIZE, record.timestamp);
+        row.write(1 + POINTER_LIMBS + BLOCK_SIZE + 1, record.count);
     } else if (idx < height) {
 #pragma unroll
         for (size_t i = 0; i < sizeof(DummyMemoryInteractionCols<uint8_t, BLOCK_SIZE>); i++) {

crates/vm/src/arch/config.rs

@@ -86,13 +86,21 @@ pub const BLOCK_FE_WIDTH: usize = 4;
 /// Bytes per memory-bus block.
 pub const MEMORY_BLOCK_BYTES: usize = BLOCK_FE_WIDTH * U16_CELL_SIZE;
 
-/// Default byte-pointer bit width.
+/// Default RV64 *byte*-pointer bit width.
+///
+/// RV64 byte addresses span 2^32 bytes, so byte pointers are 32 bits wide. This is distinct from
+/// [`POINTER_MAX_BITS`] (31), the AS-native u16-*cell* pointer width: 2^32 bytes = 2^31 u16 cells.
 pub const BYTE_POINTER_MAX_BITS: usize = to_byte_ptr_bits(POINTER_MAX_BITS);
 
+const _: () = assert!(BYTE_POINTER_MAX_BITS == 32);
+
+/// Default RV64 byte-addressable memory capacity (2^32 bytes).
+pub const DEFAULT_RV64_MEMORY_BYTE_CAPACITY: usize = 1usize << BYTE_POINTER_MAX_BITS;
+
 /// Byte count for `RV64_MEMORY_AS`.
 // TODO: make executor debug bounds use `MemoryConfig::pointer_max_bits` once
 // execution state carries the memory config.
-pub const RV64_MEMORY_BYTES: usize = 1 << BYTE_POINTER_MAX_BITS;
+pub const RV64_MEMORY_BYTES: usize = DEFAULT_RV64_MEMORY_BYTE_CAPACITY;
 
 /// Number of registers in the RV64 register file.
 pub const NUM_RV64_REGISTERS: usize = 32;

crates/vm/src/arch/execution_mode/metered/memory_ctx.rs

@@ -369,14 +369,17 @@ impl MemoryCtx {
         ptr: u32,
         size: u32,
     ) {
-        let end_ptr = ptr + size - 1;
+        // The inclusive end pointer may be up to 2^32 - 1 for u16-celled byte ranges, so compute
+        // it in u64 to avoid overflowing u32. (`size >= 1`.)
+        let end_ptr = u64::from(ptr) + u64::from(size) - 1;
         let leaf_bits = if address_space == DEFERRAL_AS {
             DEFERRAL_PTRS_PER_LEAF_BITS
         } else {
             BYTE_PTRS_PER_LEAF_BITS
         };
         let leaf_label = ptr >> leaf_bits;
-        let end_leaf_label = end_ptr >> leaf_bits;
+        // `end_leaf_label < 2^address_height` since `end_ptr < 2^BYTE_POINTER_MAX_BITS`.
+        let end_leaf_label = (end_ptr >> leaf_bits) as u32;
         let num_leaves = end_leaf_label - leaf_label + 1;
         debug_assert!(
             leaf_label < (1 << self.memory_dimensions.address_height),
@@ -591,7 +594,22 @@ impl MemoryCtx {
 
 #[cfg(test)]
 mod tests {
+    use openvm_instructions::riscv::RV64_MEMORY_AS;
+
     use super::*;
+    use crate::arch::SystemConfig;
+
+    /// A u16-celled byte range ending at the top of the 2^32 address space (inclusive end
+    /// `= 2^32 - 1`) must not overflow when computing the inclusive end pointer. Here `ptr + size`
+    /// alone overflows `u32`, so the accounting must widen to `u64` (regression for the 2^32
+    /// memory-address change).
+    #[test]
+    fn test_update_boundary_merkle_heights_at_top_of_address_space() {
+        let config = SystemConfig::default();
+        let mut ctx = MemoryCtx::new(&config, 1);
+        // Byte range [u32::MAX - 7, u32::MAX] (8 bytes ending at 2^32 - 1).
+        ctx.update_boundary_merkle_heights(RV64_MEMORY_AS, u32::MAX - 7, 8);
+    }
 
     fn reference_parent_mask(mut mask: u64) -> u64 {
         let mut parents = 0;

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

@@ -8,11 +8,7 @@ use openvm_circuit_primitives::{
     Chip,
 };
 use openvm_cpu_backend::{CpuBackend, CpuDevice, CpuProverError};
-use openvm_instructions::{
-    instruction::Instruction,
-    riscv::{RV64_REGISTER_AS, RV64_REGISTER_NUM_LIMBS},
-    DEFERRAL_AS,
-};
+use openvm_instructions::{instruction::Instruction, riscv::RV64_REGISTER_NUM_LIMBS, DEFERRAL_AS};
 use openvm_poseidon2_air::Poseidon2SubAir;
 use openvm_stark_backend::{
     interaction::{LookupBus, PermutationCheckBus},
@@ -39,7 +35,7 @@ use crate::{
         },
         to_byte_ptr_bits, vm_poseidon2_config, Arena, ExecutionBridge, ExecutionBus,
         ExecutionState, MatrixRecordArena, MemoryConfig, PreflightExecutor, Streams, VmField,
-        VmStateMut, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES,
+        VmStateMut, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, NUM_RV64_REGISTERS,
     },
     system::{
         memory::{
@@ -175,8 +171,13 @@ where
     }
 
     fn get_default_register(&mut self, increment: usize) -> usize {
+        const REGISTER_FILE_BYTES: usize = NUM_RV64_REGISTERS * size_of::<u64>();
+        if self.default_register + increment > REGISTER_FILE_BYTES {
+            self.default_register = 0;
+        }
+        let register = self.default_register;
         self.default_register += increment;
-        self.default_register - increment
+        register
     }
 
     fn get_default_pointer(&mut self, increment: usize) -> usize {
@@ -367,9 +368,6 @@ impl<F: VmField> VmChipTestBuilder<F> {
 impl<F: VmField> Default for VmChipTestBuilder<F> {
     fn default() -> Self {
         let mut mem_config = MemoryConfig::default();
-        // TODO[jpw]: this is because old tests use `gen_pointer` on address space 1; this can be
-        // removed when tests are updated.
-        mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29;
         mem_config.addr_spaces[DEFERRAL_AS as usize].num_cells = 0;
         Self::from_config(mem_config)
     }

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

@@ -25,7 +25,7 @@ use openvm_cuda_common::{
     common::get_device,
     stream::{device_synchronize, CudaStream, GpuDeviceCtx, StreamGuard},
 };
-use openvm_instructions::{program::PC_BITS, riscv::RV64_REGISTER_AS};
+use openvm_instructions::program::PC_BITS;
 use openvm_poseidon2_air::{Poseidon2Config, Poseidon2SubAir};
 use openvm_stark_backend::{
     interaction::{LookupBus, PermutationCheckBus},
@@ -55,7 +55,7 @@ use crate::{
         },
         to_byte_ptr_bits, Arena, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState,
         MatrixRecordArena, MemoryConfig, PreflightExecutor, Streams, VmStateMut, BLOCK_FE_WIDTH,
-        MEMORY_BLOCK_BYTES,
+        MEMORY_BLOCK_BYTES, NUM_RV64_REGISTERS,
     },
     system::{
         cuda::poseidon2::Poseidon2PeripheryChipGPU,
@@ -213,8 +213,13 @@ impl TestBuilder<F> for GpuChipTestBuilder {
     }
 
     fn get_default_register(&mut self, increment: usize) -> usize {
+        const REGISTER_FILE_BYTES: usize = NUM_RV64_REGISTERS * size_of::<u64>();
+        if self.default_register + increment > REGISTER_FILE_BYTES {
+            self.default_register = 0;
+        }
+        let register = self.default_register;
         self.default_register += increment;
-        self.default_register - increment
+        register
     }
 
     fn get_default_pointer(&mut self, increment: usize) -> usize {
@@ -269,11 +274,7 @@ pub struct GpuChipTestBuilder {
 
 impl Default for GpuChipTestBuilder {
     fn default() -> Self {
-        let mut mem_config = MemoryConfig::default();
-        // Tests generate register pointers across the full AS-native pointer range.
-        mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells =
-            1 << mem_config.pointer_max_bits;
-        Self::new(mem_config, default_var_range_checker_bus())
+        Self::new(MemoryConfig::default(), default_var_range_checker_bus())
     }
 }
 

crates/vm/src/arch/testing/memory/air.rs

@@ -41,7 +41,7 @@ impl<'a, T> DummyMemoryInteractionColsRef<'a, T> {
         let (count, slice) = slice.split_last().unwrap();
         let (timestamp, data) = slice.split_last().unwrap();
         Self {
-            address: MemoryAddress::new(&address[0], &address[1]),
+            address: MemoryAddress::new(&address[0], [&address[1], &address[2]]),
             data,
             timestamp,
             count,
@@ -52,19 +52,21 @@ impl<'a, T> DummyMemoryInteractionColsRef<'a, T> {
 impl<'a, T> DummyMemoryInteractionColsMut<'a, T> {
     pub fn from_mut_slice(slice: &'a mut [T]) -> Self {
         let (addr_space, slice) = slice.split_first_mut().unwrap();
-        let (ptr, slice) = slice.split_first_mut().unwrap();
+        let (ptr_lo, slice) = slice.split_first_mut().unwrap();
+        let (ptr_hi, slice) = slice.split_first_mut().unwrap();
         let (count, slice) = slice.split_last_mut().unwrap();
         let (timestamp, data) = slice.split_last_mut().unwrap();
         Self {
-            address: MemoryAddress::new(addr_space, ptr),
+            address: MemoryAddress::new(addr_space, [ptr_lo, ptr_hi]),
             data,
             timestamp,
             count,
         }
     }
 }
 
-/// AIR width = BLOCK_FE_WIDTH + 4 (addr_space, ptr, data[BLOCK_FE_WIDTH], timestamp, count)
+/// AIR width = BLOCK_FE_WIDTH + 5
+/// (addr_space, ptr_lo, ptr_hi, data[BLOCK_FE_WIDTH], timestamp, count)
 #[derive(Clone, Copy, Debug, derive_new::new)]
 pub struct MemoryDummyAir {
     pub bus: MemoryBus,
@@ -77,7 +79,7 @@ impl<F> PartitionedBaseAir<F> for MemoryDummyAir {}
 impl ColumnsAir for MemoryDummyAir {}
 impl<F> BaseAir<F> for MemoryDummyAir {
     fn width(&self) -> usize {
-        BLOCK_FE_WIDTH + 4
+        BLOCK_FE_WIDTH + 5
     }
 }
 
@@ -89,7 +91,10 @@ impl<AB: InteractionBuilder> Air<AB> for MemoryDummyAir {
 
         self.bus
             .send(
-                MemoryAddress::new(*local.address.address_space, *local.address.pointer),
+                MemoryAddress::new(
+                    *local.address.address_space,
+                    local.address.pointer_limbs.map(|p| *p),
+                ),
                 local.data.to_vec(),
                 *local.timestamp,
             )
@@ -124,7 +129,9 @@ impl<F: PrimeField32> MemoryDummyChip<F> {
     pub fn push(&mut self, addr_space: u32, ptr: u32, data: &[F], timestamp: u32, count: F) {
         assert_eq!(data.len(), BLOCK_FE_WIDTH);
         self.trace.push(F::from_u32(addr_space));
-        self.trace.push(F::from_u32(ptr));
+        // Pointer as little-endian 16-bit limbs: [lo16, hi16].
+        self.trace.push(F::from_u32(ptr & 0xffff));
+        self.trace.push(F::from_u32(ptr >> 16));
         self.trace.extend_from_slice(data);
         self.trace.push(F::from_u32(timestamp));
         self.trace.push(count);

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

@@ -55,6 +55,7 @@ impl DeviceMemoryTester {
         ));
         let mut inventory = MemoryInventoryGPU::new(
             mem_config.clone(),
+            range_checker.cpu_chip.clone().unwrap(),
             poseidon2_periphery.clone(),
             device_ctx.clone(),
         );

crates/vm/src/arch/testing/memory/mod.rs

@@ -2,7 +2,7 @@ use air::{MemoryDummyAir, MemoryDummyChip};
 use rand::Rng;
 
 use crate::{
-    arch::{MemoryCellType, VmField, BLOCK_FE_WIDTH, U16_CELL_SIZE},
+    arch::{MemoryCellType, VmField, BLOCK_FE_WIDTH, NUM_RV64_REGISTERS, U16_CELL_SIZE},
     system::memory::{
         offline_checker::pack_u8_block_value, online::TracingMemory, MemoryController,
     },
@@ -149,6 +149,32 @@ pub fn gen_pointer<R>(rng: &mut R, len: usize) -> usize
 where
     R: Rng + ?Sized,
 {
-    const MAX_MEMORY: usize = 1 << 29;
+    const MAX_MEMORY: usize = 1 << 31;
     rng.random_range(0..MAX_MEMORY - len) / len * len
 }
+
+pub fn gen_register_pointer<R>(rng: &mut R, len: usize) -> usize
+where
+    R: Rng + ?Sized,
+{
+    rng.random_range(0..NUM_RV64_REGISTERS * size_of::<u64>() - len) / len * len
+}
+
+/// Generates `N` pairwise-distinct `len`-aligned register pointers. Use this when a test writes
+/// several register operands independently: the register file only has 32 slots, so independent
+/// draws collide often enough to corrupt the expected values.
+pub fn gen_distinct_register_pointers<R, const N: usize>(rng: &mut R, len: usize) -> [usize; N]
+where
+    R: Rng + ?Sized,
+{
+    let mut ptrs = [0usize; N];
+    for i in 0..N {
+        ptrs[i] = loop {
+            let ptr = gen_register_pointer(rng, len);
+            if !ptrs[..i].contains(&ptr) {
+                break ptr;
+            }
+        };
+    }
+    ptrs
+}