diff --git a/Cargo.lock b/Cargo.lock index 72a8274551..8c1ee3e44a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6081,6 +6081,7 @@ dependencies = [ "openvm-cuda-backend", "openvm-cuda-builder", "openvm-cuda-common", + "openvm-deferral-guest", "openvm-deferral-transpiler", "openvm-instructions", "openvm-poseidon2-air", diff --git a/crates/circuits/primitives/cuda/include/primitives/constants.h b/crates/circuits/primitives/cuda/include/primitives/constants.h index 12d1745307..e37d90fa70 100644 --- a/crates/circuits/primitives/cuda/include/primitives/constants.h +++ b/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; diff --git a/crates/vm/cuda/include/system/memory/address.cuh b/crates/vm/cuda/include/system/memory/address.cuh index e321d4c91e..244a7bd472 100644 --- a/crates/vm/cuda/include/system/memory/address.cuh +++ b/crates/vm/cuda/include/system/memory/address.cuh @@ -1,6 +1,8 @@ #pragma once +#include "system/memory/params.cuh" + template struct MemoryAddress { T address_space; - T pointer; + T pointer_limbs[POINTER_LIMBS]; }; diff --git a/crates/vm/cuda/include/system/memory/params.cuh b/crates/vm/cuda/include/system/memory/params.cuh index 71d5d63fa9..343ae0d4f2 100644 --- a/crates/vm/cuda/include/system/memory/params.cuh +++ b/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 +#include // 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; diff --git a/crates/vm/cuda/src/system/boundary.cu b/crates/vm/cuda/src/system/boundary.cu index b3b7460f26..741d2fde8d 100644 --- a/crates/vm/cuda/src/system/boundary.cu +++ b/crates/vm/cuda/src/system/boundary.cu @@ -16,7 +16,7 @@ template struct BoundaryRecord { template 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 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 init_values; uint32_t addr_space_idx = record.address_space - 1; diff --git a/crates/vm/cuda/src/testing/memory.cu b/crates/vm/cuda/src/testing/memory.cu index e9caf47a12..6a1eb4121d 100644 --- a/crates/vm/cuda/src/testing/memory.cu +++ b/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 struct DummyMemoryInteractionCols { MemoryAddress 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 *>(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); i++) { diff --git a/crates/vm/src/arch/config.rs b/crates/vm/src/arch/config.rs index 8b34b7dc02..52ea9777cc 100644 --- a/crates/vm/src/arch/config.rs +++ b/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; diff --git a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs index e80961c3f4..987079aff5 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/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; diff --git a/crates/vm/src/arch/testing/cpu.rs b/crates/vm/src/arch/testing/cpu.rs index 16aa633ed4..6f12a56019 100644 --- a/crates/vm/src/arch/testing/cpu.rs +++ b/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::(); + 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 VmChipTestBuilder { impl Default for VmChipTestBuilder { 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) } diff --git a/crates/vm/src/arch/testing/cuda.rs b/crates/vm/src/arch/testing/cuda.rs index f14a95f3ec..83a8053c17 100644 --- a/crates/vm/src/arch/testing/cuda.rs +++ b/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 for GpuChipTestBuilder { } fn get_default_register(&mut self, increment: usize) -> usize { + const REGISTER_FILE_BYTES: usize = NUM_RV64_REGISTERS * size_of::(); + 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()) } } diff --git a/crates/vm/src/arch/testing/memory/air.rs b/crates/vm/src/arch/testing/memory/air.rs index 6baf5568c7..0679e4eae0 100644 --- a/crates/vm/src/arch/testing/memory/air.rs +++ b/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,11 +52,12 @@ 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, @@ -64,7 +65,8 @@ impl<'a, T> DummyMemoryInteractionColsMut<'a, T> { } } -/// 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 PartitionedBaseAir for MemoryDummyAir {} impl ColumnsAir for MemoryDummyAir {} impl BaseAir for MemoryDummyAir { fn width(&self) -> usize { - BLOCK_FE_WIDTH + 4 + BLOCK_FE_WIDTH + 5 } } @@ -89,7 +91,10 @@ impl Air 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 MemoryDummyChip { 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); diff --git a/crates/vm/src/arch/testing/memory/cuda.rs b/crates/vm/src/arch/testing/memory/cuda.rs index fb3af43f8e..80e536cd1c 100644 --- a/crates/vm/src/arch/testing/memory/cuda.rs +++ b/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(), ); diff --git a/crates/vm/src/arch/testing/memory/mod.rs b/crates/vm/src/arch/testing/memory/mod.rs index 77b8e149b9..3a2fd9aefa 100644 --- a/crates/vm/src/arch/testing/memory/mod.rs +++ b/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(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(rng: &mut R, len: usize) -> usize +where + R: Rng + ?Sized, +{ + rng.random_range(0..NUM_RV64_REGISTERS * size_of::() - 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(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 +} diff --git a/crates/vm/src/system/cuda/memory.rs b/crates/vm/src/system/cuda/memory.rs index d2a2bcbb2e..d6e6af8811 100644 --- a/crates/vm/src/system/cuda/memory.rs +++ b/crates/vm/src/system/cuda/memory.rs @@ -3,11 +3,14 @@ use std::sync::Arc; use openvm_circuit::{ arch::{AddressSpaceHostLayout, MemoryConfig, ADDR_SPACE_OFFSET, BLOCK_FE_WIDTH}, system::{ - memory::{persistent::BLOCKS_PER_LEAF, AddressMap}, + memory::{ + persistent::{BLOCKS_PER_LEAF, LOW_LEAF_BITS}, + AddressMap, + }, TouchedMemory, }, }; -use openvm_circuit_primitives::Chip; +use openvm_circuit_primitives::{var_range::SharedVariableRangeCheckerChip, Chip}; use openvm_cuda_backend::{prelude::F, GpuBackend}; use openvm_cuda_common::{ copy::{cuda_memcpy_on, MemCopyD2H, MemCopyH2D}, @@ -37,6 +40,7 @@ pub struct MemoryInventoryGPU { pub device_ctx: GpuDeviceCtx, pub boundary: BoundaryChipGPU, pub merkle_tree: MemoryMerkleTree, + pub range_checker: SharedVariableRangeCheckerChip, pub hasher_chip: Arc, pub initial_memory: Vec>>, pub merkle_records: Option>, @@ -70,6 +74,7 @@ impl MemoryInventoryGPU { pub fn new( config: MemoryConfig, + range_checker: SharedVariableRangeCheckerChip, hasher_chip: Arc, device_ctx: GpuDeviceCtx, ) -> Self { @@ -77,6 +82,7 @@ impl MemoryInventoryGPU { device_ctx: device_ctx.clone(), boundary: BoundaryChipGPU::new(hasher_chip.shared_buffer(), device_ctx.clone()), merkle_tree: MemoryMerkleTree::new(config.clone(), hasher_chip.clone(), device_ctx), + range_checker, hasher_chip, initial_memory: Vec::new(), merkle_records: None, @@ -243,6 +249,13 @@ impl MemoryInventoryGPU { .to_host_on(&self.device_ctx) .unwrap(); let record_words = 2 + BLOCKS_PER_LEAF + DIGEST_WIDTH; + let low_mask = (1u32 << LOW_LEAF_BITS) - 1; + let high_bits = self + .merkle_tree + .mem_config() + .memory_dimensions() + .address_height + .saturating_sub(LOW_LEAF_BITS); let mut merkle_records = Vec::with_capacity(out_num_records); for i in 0..out_num_records { let base = i * record_words; @@ -255,6 +268,13 @@ impl MemoryInventoryGPU { .iter() .max() .unwrap(); + let leaf_label = out_records[base + 1] / DIGEST_WIDTH as u32; + let low = leaf_label & low_mask; + let high = leaf_label >> LOW_LEAF_BITS; + for _ in 0..2 { + self.range_checker.add_count(low, LOW_LEAF_BITS); + self.range_checker.add_count(high, high_bits); + } let record = MemoryMerkleRecord { address_space: out_records[base], ptr: out_records[base + 1], @@ -330,6 +350,7 @@ mod tests { poseidon2::Poseidon2PeripheryChip, }, }; + use openvm_circuit_primitives::var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip}; use openvm_cuda_backend::prelude::F; use openvm_cuda_common::{ common::get_device, @@ -370,8 +391,17 @@ mod tests { stream: StreamGuard::new(CudaStream::new_non_blocking().unwrap()), }; let hasher_chip = Arc::new(Poseidon2PeripheryChipGPU::new(1, device_ctx.clone())); - let mut inventory = - MemoryInventoryGPU::new(mem_config.clone(), hasher_chip, device_ctx.clone()); + let range_max_bits = mem_config.pointer_max_bits.max(LOW_LEAF_BITS); + let range_checker = Arc::new(VariableRangeCheckerChip::new(VariableRangeCheckerBus::new( + 0, + range_max_bits, + ))); + let mut inventory = MemoryInventoryGPU::new( + mem_config.clone(), + range_checker, + hasher_chip, + device_ctx.clone(), + ); inventory.set_initial_memory(&memory.memory); let ctxs = inventory.generate_proving_ctxs(Vec::new()); @@ -445,8 +475,17 @@ mod tests { stream: StreamGuard::new(CudaStream::new_non_blocking().unwrap()), }; let hasher_chip = Arc::new(Poseidon2PeripheryChipGPU::new(1, device_ctx.clone())); - let mut inventory = - MemoryInventoryGPU::new(mem_config.clone(), hasher_chip, device_ctx.clone()); + let range_max_bits = mem_config.pointer_max_bits.max(LOW_LEAF_BITS); + let range_checker = Arc::new(VariableRangeCheckerChip::new(VariableRangeCheckerBus::new( + 0, + range_max_bits, + ))); + let mut inventory = MemoryInventoryGPU::new( + mem_config.clone(), + range_checker, + hasher_chip, + device_ctx.clone(), + ); inventory.set_initial_memory(&memory.memory); let touched_memory = vec![ diff --git a/crates/vm/src/system/cuda/mod.rs b/crates/vm/src/system/cuda/mod.rs index 722098a46f..cba562dffa 100644 --- a/crates/vm/src/system/cuda/mod.rs +++ b/crates/vm/src/system/cuda/mod.rs @@ -54,6 +54,7 @@ impl SystemChipInventoryGPU { let memory_inventory = MemoryInventoryGPU::new( config.memory_config.clone(), + cpu_range_checker, hasher_chip, device_ctx.clone(), ); diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index e0a486b35f..05bc0f8826 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -115,7 +115,13 @@ impl MemoryController { let memory_dims = mem_config.memory_dimensions(); let range_checker_bus = range_checker.bus(); let interface_chip = MemoryInterface { - boundary_chip: PersistentBoundaryChip::new(memory_bus, merkle_bus, compression_bus), + boundary_chip: PersistentBoundaryChip::new( + memory_bus, + merkle_bus, + compression_bus, + range_checker.clone(), + memory_dims, + ), merkle_chip: MemoryMerkleChip::new(memory_dims, merkle_bus, compression_bus), initial_memory: AddressMap::from_mem_config(&mem_config), }; diff --git a/crates/vm/src/system/memory/mod.rs b/crates/vm/src/system/memory/mod.rs index 243e7714a9..8dc0d9864b 100644 --- a/crates/vm/src/system/memory/mod.rs +++ b/crates/vm/src/system/memory/mod.rs @@ -27,8 +27,13 @@ use crate::{ // TODO: Currently this is only used for debug assertions, but we may switch to making it constant // and removing from MemoryConfig. -// This is the max bit width of AS-native OpenVM memory pointers. -pub const POINTER_MAX_BITS: usize = 28; +// This is the max bit width of AS-native OpenVM memory pointers (i.e. pointers measured in +// AS-native storage cells, *not* in guest bytes). +// +// RV64 targets 2^32 byte-addressable memory. Since `RV64_MEMORY_AS` uses u16 storage cells, 2^32 +// bytes equals 2^31 u16 cells, so the AS-native pointer width is 31 bits. See [`crate::arch`]'s +// `BYTE_POINTER_MAX_BITS` for the corresponding *byte*-pointer width (32). +pub const POINTER_MAX_BITS: usize = 31; #[derive(PartialEq, Copy, Clone, Debug, Eq)] pub enum OpType { @@ -36,20 +41,32 @@ pub enum OpType { Write = 1, } +/// Number of little-endian 16-bit limbs used to represent an AS-native memory pointer on the +/// memory bus. +/// +/// AS-native pointers can be up to [`POINTER_MAX_BITS`] (31) bits wide, which exceeds the +/// BabyBear field modulus. To avoid composing a full pointer into a single field element, every +/// memory-bus pointer is carried as two little-endian 16-bit limbs `[lo16, hi16]`. +pub const POINTER_LIMBS: usize = 2; + /// The full pointer to a location in memory consists of an address space and a pointer within /// the address space. +/// +/// The AS-native pointer is stored as little-endian 16-bit limbs `pointer_limbs = [lo16, hi16]` +/// (see [`POINTER_LIMBS`]). These limbs are *AS-native cell* pointer limbs, not guest byte-pointer +/// limbs. #[derive(Clone, Copy, Debug, PartialEq, Eq, AlignedBorrow, StructReflection)] #[repr(C)] pub struct MemoryAddress { pub address_space: S, - pub pointer: T, + pub pointer_limbs: [T; POINTER_LIMBS], } impl MemoryAddress { - pub fn new(address_space: S, pointer: T) -> Self { + pub fn new(address_space: S, pointer_limbs: [T; POINTER_LIMBS]) -> Self { Self { address_space, - pointer, + pointer_limbs, } } @@ -60,11 +77,31 @@ impl MemoryAddress { { Self { address_space: a.address_space.into(), - pointer: a.pointer.into(), + pointer_limbs: a.pointer_limbs.map(Into::into), + } + } +} + +impl MemoryAddress { + /// Build an address from a concrete AS-native `pointer`, split into little-endian 16-bit + /// limbs `[lo16, hi16]`. + #[inline(always)] + pub fn from_u32_pointer(address_space: S, pointer: u32) -> Self { + Self { + address_space, + pointer_limbs: pointer_to_limbs(pointer), } } } +/// Splits a concrete AS-native pointer into little-endian 16-bit limbs `[lo16, hi16]`. +#[inline(always)] +pub fn pointer_to_limbs( + pointer: u32, +) -> [T; POINTER_LIMBS] { + [T::from_u32(pointer & 0xffff), T::from_u32(pointer >> 16)] +} + #[derive(Clone)] pub struct MemoryAirInventory { pub bridge: MemoryBridge, @@ -84,6 +121,8 @@ impl MemoryAirInventory { memory_bus, merkle_bus, compression_bus, + range_bus: bridge.range_bus(), + memory_dimensions: memory_dims, }; let merkle = MemoryMerkleAir:: { memory_dimensions: memory_dims, diff --git a/crates/vm/src/system/memory/offline_checker/bus.rs b/crates/vm/src/system/memory/offline_checker/bus.rs index e41a5eeee9..bbfc451806 100644 --- a/crates/vm/src/system/memory/offline_checker/bus.rs +++ b/crates/vm/src/system/memory/offline_checker/bus.rs @@ -58,7 +58,10 @@ impl MemoryBus { MemoryBusInteraction { bus: self.inner, is_send, - address: MemoryAddress::new(address.address_space.into(), address.pointer.into()), + address: MemoryAddress::new( + address.address_space.into(), + address.pointer_limbs.map(Into::into), + ), data: data.into_iter().map(|item| item.into()).collect(), timestamp: timestamp.into(), } @@ -87,9 +90,10 @@ impl MemoryBusInteraction { where AB: InteractionBuilder, { + // Memory-bus payload order: [address_space, pointer_lo, pointer_hi, data..., timestamp]. let fields = iter::empty() .chain(iter::once(self.address.address_space)) - .chain(iter::once(self.address.pointer)) + .chain(self.address.pointer_limbs) .chain(self.data) .chain(iter::once(self.timestamp)); diff --git a/crates/vm/src/system/memory/persistent.rs b/crates/vm/src/system/memory/persistent.rs index 841425c1e7..f8b8b870d2 100644 --- a/crates/vm/src/system/memory/persistent.rs +++ b/crates/vm/src/system/memory/persistent.rs @@ -4,7 +4,10 @@ use std::{ iter, }; -use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper}; +use openvm_circuit_primitives::{ + var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus}, + ColumnsAir, StructReflection, StructReflectionHelper, U16_BITS, +}; use openvm_circuit_primitives_derive::AlignedBorrow; use openvm_cpu_backend::CpuBackend; use openvm_stark_backend::{ @@ -23,14 +26,25 @@ use crate::{ arch::{hasher::Hasher, ADDR_SPACE_OFFSET, BLOCK_FE_WIDTH}, primitives::Chip, system::memory::{ - controller::DIGEST_WIDTH, offline_checker::MemoryBus, MemoryAddress, MemoryImage, - TimestampedEquipartition, + controller::{dimensions::MemoryDimensions, DIGEST_WIDTH, DIGEST_WIDTH_BITS}, + offline_checker::MemoryBus, + MemoryAddress, MemoryImage, TimestampedEquipartition, }, }; /// Number of memory-bus blocks covered by one merkle leaf. pub const BLOCKS_PER_LEAF: usize = DIGEST_WIDTH / BLOCK_FE_WIDTH; +/// Number of low bits of a leaf label kept in the `low` limb of [`PersistentBoundaryCols:: +/// leaf_label_limbs`]. +/// +/// A merkle leaf spans `DIGEST_WIDTH` AS-native cells, so the leaf's base cell pointer is +/// `leaf_label * DIGEST_WIDTH`. To send that pointer's low 16-bit limb without composing the full +/// (up to 31-bit) pointer into a field element, we keep the low `LOW_LEAF_BITS` bits of the leaf +/// label in `low` so that `low * DIGEST_WIDTH` (plus a small in-leaf block offset) stays within +/// 16 bits: `LOW_LEAF_BITS = U16_BITS - DIGEST_WIDTH_BITS`. +pub const LOW_LEAF_BITS: usize = U16_BITS - DIGEST_WIDTH_BITS; + /// The values describe one merkle leaf (`DIGEST_WIDTH` cells)---the data together with the /// last accessed timestamp---in either the initial or final memory state. #[repr(C)] @@ -41,7 +55,13 @@ pub struct PersistentBoundaryCols { // `expand_direction` = 0 corresponds to irrelevant row (all interactions multiplicity 0) pub expand_direction: T, pub address_space: T, - pub leaf_label: T, + /// Leaf label decomposed into little-endian limbs `[low, high]`: + /// `leaf_label = low + 2^LOW_LEAF_BITS * high`, + /// where `low` is range-checked to [`LOW_LEAF_BITS`] bits and `high` to + /// `address_height - LOW_LEAF_BITS` bits. The decomposition lets us emit the leaf's base + /// AS-native cell pointer as two 16-bit limbs without composing the full pointer into one + /// field element. + pub leaf_label_limbs: [T; 2], pub values: [T; DIGEST_WIDTH], pub hash: [T; DIGEST_WIDTH], /// Per-block timestamps. Each BLOCK_FE_WIDTH block within the leaf has its own timestamp. @@ -63,6 +83,8 @@ pub struct PersistentBoundaryAir { pub memory_bus: MemoryBus, pub merkle_bus: PermutationCheckBus, pub compression_bus: PermutationCheckBus, + pub range_bus: VariableRangeCheckerBus, + pub memory_dimensions: MemoryDimensions, } impl BaseAir for PersistentBoundaryAir { @@ -100,13 +122,35 @@ impl Air when_initial.assert_zero(local.timestamps[i]); } + // Decompose the leaf label into `[low, high]` limbs and reconstruct it for the merkle bus. + // `leaf_label = low + 2^LOW_LEAF_BITS * high`. We range-check the limbs (on active rows + // only) so that the leaf's base cell pointer `leaf_label * DIGEST_WIDTH` splits cleanly + // into two 16-bit limbs below: `low * DIGEST_WIDTH < 2^16` and `high < 2^16`. + let low = local.leaf_label_limbs[0]; + let high = local.leaf_label_limbs[1]; + let leaf_label = low.into() + high.into() * AB::F::from_u32(1u32 << LOW_LEAF_BITS); + + // Active rows have `expand_direction in {1, -1}`, so `expand_direction^2 = 1`; padding rows + // have `expand_direction = 0`. + let is_active = local.expand_direction * local.expand_direction; + let high_bits = self + .memory_dimensions + .address_height + .saturating_sub(LOW_LEAF_BITS); + self.range_bus + .range_check(low, LOW_LEAF_BITS) + .eval(builder, is_active.clone()); + self.range_bus + .range_check(high, high_bits) + .eval(builder, is_active); + let mut expand_fields = vec![ // direction = 1 => is_final = 0 // direction = -1 => is_final = 1 local.expand_direction.into(), AB::Expr::ZERO, local.address_space - AB::F::from_u32(ADDR_SPACE_OFFSET), - local.leaf_label.into(), + leaf_label, ]; expand_fields.extend(local.hash.map(Into::into)); self.merkle_bus @@ -121,13 +165,18 @@ impl Air local.expand_direction * local.expand_direction, ); - let leaf_ptr = local.leaf_label * AB::F::from_usize(DIGEST_WIDTH); for block_idx in 0..BLOCKS_PER_LEAF { - let ptr = leaf_ptr.clone() + AB::F::from_usize(block_idx * BLOCK_FE_WIDTH); + // The leaf's base cell pointer is `leaf_label * DIGEST_WIDTH`; block `block_idx` starts + // at `+ block_idx * BLOCK_FE_WIDTH`. As little-endian 16-bit limbs: + // pointer_lo = low * DIGEST_WIDTH + block_idx * BLOCK_FE_WIDTH (< 2^16) + // pointer_hi = high + let pointer_lo = low.into() * AB::F::from_usize(DIGEST_WIDTH) + + AB::F::from_usize(block_idx * BLOCK_FE_WIDTH); + let pointer_hi = high.into(); // Each block uses its own timestamp; untouched blocks stay at t=0. self.memory_bus .send( - MemoryAddress::new(local.address_space, ptr), + MemoryAddress::new(local.address_space.into(), [pointer_lo, pointer_hi]), local.values[block_idx * BLOCK_FE_WIDTH..(block_idx + 1) * BLOCK_FE_WIDTH] .to_vec(), local.timestamps[block_idx], @@ -139,6 +188,7 @@ impl Air pub struct PersistentBoundaryChip { pub air: PersistentBoundaryAir, + range_checker: SharedVariableRangeCheckerChip, touched_labels: Option>>, overridden_height: Option, } @@ -190,13 +240,18 @@ impl PersistentBoundaryChip Self { Self { air: PersistentBoundaryAir { memory_bus, merkle_bus, compression_bus, + range_bus: range_checker.bus(), + memory_dimensions, }, + range_checker, touched_labels: None, overridden_height: None, } @@ -283,14 +338,32 @@ where } let mut rows = Val::::zero_vec(height * width); + // `leaf_label = low + 2^LOW_LEAF_BITS * high`. + let low_mask = (1u32 << LOW_LEAF_BITS) - 1; + let high_bits = self + .air + .memory_dimensions + .address_height + .saturating_sub(LOW_LEAF_BITS); + rows.par_chunks_mut(2 * width) .zip(touched_labels.par_iter()) .for_each(|(row, touched_label)| { + let low = touched_label.label & low_mask; + let high = touched_label.label >> LOW_LEAF_BITS; + let leaf_label_limbs = [Val::::from_u32(low), Val::::from_u32(high)]; + // Both the initial and final active rows range-check the limbs (the AIR sends + // the range check with multiplicity `expand_direction^2 = 1` on each). + for _ in 0..2 { + self.range_checker.add_count(low, LOW_LEAF_BITS); + self.range_checker.add_count(high, high_bits); + } + let (initial_row, final_row) = row.split_at_mut(width); *initial_row.borrow_mut() = PersistentBoundaryCols { expand_direction: Val::::ONE, address_space: Val::::from_u32(touched_label.address_space), - leaf_label: Val::::from_u32(touched_label.label), + leaf_label_limbs, values: touched_label.init_values, hash: touched_label.init_hash, timestamps: [Val::::from_u32(INITIAL_TIMESTAMP); BLOCKS_PER_LEAF], @@ -299,7 +372,7 @@ where *final_row.borrow_mut() = PersistentBoundaryCols { expand_direction: Val::::NEG_ONE, address_space: Val::::from_u32(touched_label.address_space), - leaf_label: Val::::from_u32(touched_label.label), + leaf_label_limbs, values: touched_label.final_values, hash: touched_label.final_hash, timestamps: touched_label.final_timestamps.map(Val::::from_u32), diff --git a/crates/vm/src/system/memory/tests.rs b/crates/vm/src/system/memory/tests.rs index 75878071e1..9d249c8489 100644 --- a/crates/vm/src/system/memory/tests.rs +++ b/crates/vm/src/system/memory/tests.rs @@ -54,6 +54,25 @@ fn test_memory_write(its: usize) { tester.simple_test().expect("Verification failed"); } +#[test] +fn test_memory_write_max_address() { + let mut rng = create_seeded_rng(); + let mem_config = MemoryConfig::default(); + // The default config gives RV64_MEMORY_AS its full 2^32-byte capacity (2^31 u16 cells). + // Touch the last cell block so the boundary/merkle chips process the top of the address range. + let last_block = mem_config.addr_spaces[RV64_MEMORY_AS as usize].num_cells - BLOCK_FE_WIDTH; + let mut tester = VmChipTestBuilder::::from_config(mem_config); + let values: [F; BLOCK_FE_WIDTH] = + array::from_fn(|_| F::from_u32(rng.random_range(0..u16::MAX as u32 + 1))); + tester.write::(RV64_MEMORY_AS as usize, last_block, values); + assert_eq!( + tester.read::(RV64_MEMORY_AS as usize, last_block), + values + ); + let tester = tester.build().finalize(); + tester.simple_test().expect("Verification failed"); +} + #[cfg(feature = "cuda")] #[test_case(1000)] #[test_case(0)] diff --git a/extensions/algebra/circuit/src/fp2_chip/tests.rs b/extensions/algebra/circuit/src/fp2_chip/tests.rs index 42939f8c93..c2a50f2173 100644 --- a/extensions/algebra/circuit/src/fp2_chip/tests.rs +++ b/extensions/algebra/circuit/src/fp2_chip/tests.rs @@ -7,7 +7,10 @@ use num_bigint::BigUint; use num_traits::Zero; use openvm_algebra_transpiler::Fp2Opcode; use openvm_circuit::arch::{ - testing::{memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder}, + testing::{ + memory::{gen_distinct_register_pointers, gen_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, + }, Arena, PreflightExecutor, MEMORY_BLOCK_BYTES, }; use openvm_circuit_primitives::bigint::utils::secp256k1_coord_prime; @@ -148,9 +151,7 @@ fn set_and_execute_fp2( let ptr_as = RV64_REGISTER_AS as usize; let data_as = RV64_MEMORY_AS as usize; - let rs1_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rd_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1_ptr, rs2_ptr, rd_ptr] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let a_base_addr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS) as u32; let b_base_addr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS) as u32; diff --git a/extensions/deferral/circuit/Cargo.toml b/extensions/deferral/circuit/Cargo.toml index cc84264b41..8ede47cfd4 100644 --- a/extensions/deferral/circuit/Cargo.toml +++ b/extensions/deferral/circuit/Cargo.toml @@ -19,6 +19,7 @@ openvm-circuit-derive = { workspace = true } openvm-instructions = { workspace = true } openvm-riscv-circuit = { workspace = true } openvm-deferral-transpiler = { workspace = true } +openvm-deferral-guest = { workspace = true } rvr-openvm-lift = { workspace = true, optional = true } rvr-openvm-ext-deferral = { workspace = true, optional = true } rvr-openvm-ext-ffi-common = { workspace = true, optional = true } diff --git a/extensions/deferral/circuit/build.rs b/extensions/deferral/circuit/build.rs index ec0fb1787a..e4abc51bd0 100644 --- a/extensions/deferral/circuit/build.rs +++ b/extensions/deferral/circuit/build.rs @@ -13,8 +13,10 @@ fn main() { .include("../../../crates/circuits/primitives/cuda/include") .include("../../../crates/circuits/poseidon2-air/cuda/include") .include("../../../crates/vm/cuda/include") + .include("../../riscv-adapters/cuda/include") .include("cuda/include") .watch("cuda/src") + .watch("../../riscv-adapters/cuda") .library_name("tracegen_gpu_deferral") .files_from_glob("cuda/src/**/*.cu"); diff --git a/extensions/deferral/circuit/cuda/include/def_types.h b/extensions/deferral/circuit/cuda/include/def_types.h index 8406ea8338..ef14ca384e 100644 --- a/extensions/deferral/circuit/cuda/include/def_types.h +++ b/extensions/deferral/circuit/cuda/include/def_types.h @@ -4,6 +4,7 @@ #include #include "fp.h" +#include "primitives/constants.h" #include "system/memory/params.cuh" namespace deferral { @@ -30,4 +31,8 @@ inline constexpr uint8_t BABY_BEAR_ORDER_BE[F_NUM_BYTES] = { static_cast(BABY_BEAR_ORDER & 0xff), }; +// Maximum AS-native cell-pointer width; mirrors +// `openvm_circuit::system::memory::POINTER_MAX_BITS`. +inline constexpr size_t POINTER_MAX_BITS = 31; + } // namespace deferral diff --git a/extensions/deferral/circuit/cuda/src/call.cu b/extensions/deferral/circuit/cuda/src/call.cu index 32059bb37e..fc8b9672cd 100644 --- a/extensions/deferral/circuit/cuda/src/call.cu +++ b/extensions/deferral/circuit/cuda/src/call.cu @@ -12,6 +12,7 @@ #include "primitives/fp_array.cuh" #include "primitives/histogram.cuh" #include "primitives/trace_access.h" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -206,6 +207,14 @@ template struct DeferralCallAdapterCols { MemoryWriteAuxCols output_commit_and_len_aux[OUTPUT_TOTAL_MEMORY_OPS]; MemoryWriteAuxCols new_input_acc_aux[DIGEST_F_MEMORY_OPS]; MemoryWriteAuxCols new_output_acc_aux[DIGEST_F_MEMORY_OPS]; + + // Carries for converting the heap `input`/`output` base byte pointers to AS-native u16 cell + // pointer limbs, plus per-block cell-offset carries. The DEFERRAL_AS accumulator pointers are + // bounded below 2^16 and need no carry/decomposition columns. + T input_cell_carry; + T output_cell_carry; + T input_commit_add_carry[COMMIT_MEMORY_OPS]; + T output_add_carry[OUTPUT_TOTAL_MEMORY_OPS]; }; __device__ __forceinline__ void deferral_call_adapter_tracegen( @@ -213,6 +222,7 @@ __device__ __forceinline__ void deferral_call_adapter_tracegen( const DeferralCallAdapterRecord &record, BitwiseOperationLookup &bitwise_buffer, MemoryAuxColsFactory &mem_helper, + VariableRangeChecker &range_checker, const size_t address_bits ) { const uint32_t limb_shift_bits = RV64_BYTE_BITS * RV64_WORD_NUM_LIMBS - address_bits; @@ -315,6 +325,49 @@ __device__ __forceinline__ void deferral_call_adapter_tracegen( ); mem_helper.fill(aux_row, record.new_output_acc_aux[i].prev_timestamp, timestamp++); } + + // Convert the heap `input` (rs_val) and `output` (rd_val) base byte pointers to AS-native u16 + // cell pointer limbs and emit the matching range-check counts. Mirrors the per-block carry + // computation in the host `DeferralCallAdapterFiller`. + const uint32_t heap_cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + + { + const uint32_t input_ptr = + static_cast(record.rs_val[0]) | + (static_cast(record.rs_val[1]) << RV64_BYTE_BITS) | + (static_cast(record.rs_val[2]) << (2 * RV64_BYTE_BITS)) | + (static_cast(record.rs_val[3]) << (3 * RV64_BYTE_BITS)); + uint32_t add_carries[COMMIT_MEMORY_OPS]; + const uint32_t conv_carry = compute_pointer_carries( + range_checker, + input_ptr, + address_bits, + COMMIT_MEMORY_OPS, + heap_cell_stride, + add_carries + ); + COL_WRITE_VALUE(row, DeferralCallAdapterCols, input_cell_carry, Fp(conv_carry)); + COL_WRITE_ARRAY(row, DeferralCallAdapterCols, input_commit_add_carry, add_carries); + } + + { + const uint32_t output_ptr = + static_cast(record.rd_val[0]) | + (static_cast(record.rd_val[1]) << RV64_BYTE_BITS) | + (static_cast(record.rd_val[2]) << (2 * RV64_BYTE_BITS)) | + (static_cast(record.rd_val[3]) << (3 * RV64_BYTE_BITS)); + uint32_t add_carries[OUTPUT_TOTAL_MEMORY_OPS]; + const uint32_t conv_carry = compute_pointer_carries( + range_checker, + output_ptr, + address_bits, + OUTPUT_TOTAL_MEMORY_OPS, + heap_cell_stride, + add_carries + ); + COL_WRITE_VALUE(row, DeferralCallAdapterCols, output_cell_carry, Fp(conv_carry)); + COL_WRITE_ARRAY(row, DeferralCallAdapterCols, output_add_carry, add_carries); + } } // =========================== COMBINED =========================== @@ -356,6 +409,7 @@ __global__ void deferral_call_tracegen( DeferralCallRecord record = records[row_idx]; Histogram count_buffer(count_ptr, num_def_circuits); + VariableRangeChecker range_checker(range_checker_ptr, range_checker_num_bins); MemoryAuxColsFactory mem_helper( VariableRangeChecker(range_checker_ptr, range_checker_num_bins), timestamp_max_bits ); @@ -364,7 +418,9 @@ __global__ void deferral_call_tracegen( poseidon2_records, poseidon2_counts, poseidon2_idx, poseidon2_capacity ); - deferral_call_adapter_tracegen(row, record.adapter, bitwise_buffer, mem_helper, address_bits); + deferral_call_adapter_tracegen( + row, record.adapter, bitwise_buffer, mem_helper, range_checker, address_bits + ); deferral_call_core_tracegen( row.slice_from(COL_INDEX(DeferralCallCols, core)), record.core, diff --git a/extensions/deferral/circuit/cuda/src/output.cu b/extensions/deferral/circuit/cuda/src/output.cu index 193d6d3da4..a7492b0a9f 100644 --- a/extensions/deferral/circuit/cuda/src/output.cu +++ b/extensions/deferral/circuit/cuda/src/output.cu @@ -12,6 +12,7 @@ #include "primitives/fp_array.cuh" #include "primitives/histogram.cuh" #include "primitives/trace_access.h" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -105,6 +106,15 @@ template struct DeferralOutputCols { // Capacity of the permutation of write_bytes and the previous row's capacity on // non-last rows, compression on the last row. T poseidon2_res[DIGEST_SIZE]; + + // Carry for converting the heap `input` base byte pointer (read on the first row) to AS-native + // u16 cell pointer limbs, plus per-block cell-offset carries. + T input_cell_carry; + T input_add_carry[OUTPUT_TOTAL_MEMORY_OPS]; + + // Per-row output write cell pointer limbs `[lo16, hi16]`. On write rows this equals + // `(output_ptr + (section_idx - 1) * DIGEST_SIZE) / 2`. + T write_cell_ptr[2]; }; __global__ void deferral_output_tracegen( @@ -147,6 +157,7 @@ __global__ void deferral_output_tracegen( const bool is_last = section_idx + 1 == header.num_rows; Histogram count_buffer(count_ptr, num_def_circuits); + VariableRangeChecker range_checker(range_checker_ptr, range_checker_num_bins); MemoryAuxColsFactory mem_helper( VariableRangeChecker(range_checker_ptr, range_checker_num_bins), timestamp_max_bits ); @@ -254,6 +265,31 @@ __global__ void deferral_output_tracegen( COL_WRITE_ARRAY(row, DeferralOutputCols, sponge_inputs, sponge_inputs); COL_FILL_ZERO(row, DeferralOutputCols, write_bytes_aux); + + // Convert the heap `input` (rs_val) base byte pointer to AS-native u16 cell pointer limbs + // and emit the matching range-check counts. Mirrors the first-row branch of the host + // `DeferralOutputFiller`. + const uint32_t heap_cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + const uint32_t input_ptr = + static_cast(header.rs_val[0]) | + (static_cast(header.rs_val[1]) << RV64_BYTE_BITS) | + (static_cast(header.rs_val[2]) << (2 * RV64_BYTE_BITS)) | + (static_cast(header.rs_val[3]) << (3 * RV64_BYTE_BITS)); + uint32_t add_carries[OUTPUT_TOTAL_MEMORY_OPS]; + const uint32_t conv_carry = compute_pointer_carries( + range_checker, + input_ptr, + address_bits, + OUTPUT_TOTAL_MEMORY_OPS, + heap_cell_stride, + add_carries + ); + COL_WRITE_VALUE(row, DeferralOutputCols, input_cell_carry, Fp(conv_carry)); + COL_WRITE_ARRAY(row, DeferralOutputCols, input_add_carry, add_carries); + + // The output write cell pointer is unconstrained on the first row (its constraints are + // gated by `is_write_row`); match the host trace, which leaves it zero. + COL_FILL_ZERO(row, DeferralOutputCols, write_cell_ptr); } else { const uint8_t *header_end = record_start + sizeof(DeferralOutputRecordHeader); const uint8_t *write_bytes_start = header_end + (section_idx - 1) * DIGEST_SIZE; @@ -294,6 +330,27 @@ __global__ void deferral_output_tracegen( static_cast(aux_idx) ); } + + // Input-conversion carries are only populated on the first row; match the host trace, + // which leaves them zero on write rows. + COL_WRITE_VALUE(row, DeferralOutputCols, input_cell_carry, Fp::zero()); + COL_FILL_ZERO(row, DeferralOutputCols, input_add_carry); + + // Output write cell pointer for this row = `(output_ptr + (section_idx - 1) * DIGEST_SIZE) + // / 2`, witnessed as little-endian 16-bit cell-pointer limbs `[lo16, hi16]` and + // range-checked. Mirrors the write-row branch of the host `DeferralOutputFiller`. + const uint32_t rd_val = static_cast(header.rd_val[0]) | + (static_cast(header.rd_val[1]) << RV64_BYTE_BITS) | + (static_cast(header.rd_val[2]) << (2 * RV64_BYTE_BITS)) | + (static_cast(header.rd_val[3]) << (3 * RV64_BYTE_BITS)); + const uint32_t write_byte_ptr = rd_val + (section_idx - 1) * DIGEST_SIZE; + const uint32_t write_cell = write_byte_ptr >> 1; + const uint32_t write_cell_lo = write_cell & 0xffffu; + const uint32_t write_cell_hi = write_cell >> openvm::U16_BITS; + range_checker.add_count(write_cell_lo, openvm::U16_BITS); + range_checker.add_count(write_cell_hi, POINTER_MAX_BITS - openvm::U16_BITS); + const Fp write_cell_limbs[2] = {Fp(write_cell_lo), Fp(write_cell_hi)}; + COL_WRITE_ARRAY(row, DeferralOutputCols, write_cell_ptr, write_cell_limbs); } Fp prev_capacity[DIGEST_SIZE]; diff --git a/extensions/deferral/circuit/src/call/air.rs b/extensions/deferral/circuit/src/call/air.rs index 35a6734b11..4ffdc949fe 100644 --- a/extensions/deferral/circuit/src/call/air.rs +++ b/extensions/deferral/circuit/src/call/air.rs @@ -4,7 +4,7 @@ use itertools::{izip, Itertools as _}; use openvm_circuit::{ arch::{ AdapterAirContext, ExecutionBridge, ExecutionState, ImmInstruction, VmAdapterAir, - VmAdapterInterface, VmCoreAir, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, + VmAdapterInterface, VmCoreAir, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{pack_u8_block, MemoryBridge, MemoryReadAuxCols, MemoryWriteAuxCols}, @@ -12,8 +12,8 @@ use openvm_circuit::{ }, }; use openvm_circuit_primitives::{ - bitwise_op_lookup::BitwiseOperationLookupBus, ColumnsAir, StructReflection, - StructReflectionHelper, + bitwise_op_lookup::BitwiseOperationLookupBus, var_range::VariableRangeCheckerBus, ColumnsAir, + StructReflection, StructReflectionHelper, }; use openvm_circuit_primitives_derive::AlignedBorrow; use openvm_deferral_transpiler::DeferralOpcode; @@ -22,7 +22,10 @@ use openvm_instructions::{ riscv::{RV64_BYTE_BITS, RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_WORD_NUM_LIMBS}, LocalOpcode, DEFERRAL_AS, }; -use openvm_riscv_circuit::adapters::{byte_ptr_to_u16_ptr, expand_to_rv64_register}; +use openvm_riscv_circuit::adapters::{ + eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_register, + pack_u8_pair, reg_byte_ptr_to_cell_ptr_limbs, +}; use openvm_stark_backend::{ interaction::InteractionBuilder, p3_air::BaseAir, @@ -38,7 +41,7 @@ use crate::{ count::DeferralCircuitCountBus, poseidon2::DeferralPoseidon2Bus, utils::{ - byte_commit_to_f, bytes_to_f, combine_output, split_byte_memory_ops, split_f_memory_ops, + byte_commit_to_f, combine_output, split_byte_memory_ops, split_f_memory_ops, COMMIT_MEMORY_OPS, COMMIT_NUM_BYTES, DIGEST_F_MEMORY_OPS, F_NUM_BYTES, OUTPUT_TOTAL_BYTES, OUTPUT_TOTAL_MEMORY_OPS, }, @@ -254,6 +257,18 @@ pub struct DeferralCallAdapterCols { pub output_commit_and_len_aux: [MemoryWriteAuxCols; OUTPUT_TOTAL_MEMORY_OPS], pub new_input_acc_aux: [MemoryWriteAuxCols; DIGEST_F_MEMORY_OPS], pub new_output_acc_aux: [MemoryWriteAuxCols; DIGEST_F_MEMORY_OPS], + + // Carries for converting the heap `input`/`output` base *byte* pointers to AS-native u16 + // *cell* pointer limbs. + pub input_cell_carry: T, + pub output_cell_carry: T, + // Per-block carries for adding the cell offset `chunk_idx * (MEMORY_BLOCK_BYTES / + // U16_CELL_SIZE)` to each base cell pointer. + pub input_commit_add_carry: [T; COMMIT_MEMORY_OPS], + pub output_add_carry: [T; OUTPUT_TOTAL_MEMORY_OPS], + // The DEFERRAL_AS accumulator cell pointers need no limb decomposition or add-carry columns: + // they are bounded below 2^16 (see the static assert in `super`), so the high cell limb is + // identically zero and the low limb is the algebraic cell pointer. } #[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)] @@ -262,6 +277,7 @@ pub struct DeferralCallAdapterAir { pub execution_bridge: ExecutionBridge, pub memory_bridge: MemoryBridge, pub bitwise_bus: BitwiseOperationLookupBus, + pub range_bus: VariableRangeCheckerBus, pub address_bits: usize, } @@ -297,10 +313,11 @@ impl VmAdapterAir for DeferralCallAdapterAir { let rd_full = expand_to_rv64_register(&cols.rd_val); let rs_full = expand_to_rv64_register(&cols.rs_val); - // Heap pointers are first read from their respective registers. + // Heap pointers are first read from their respective registers. Register byte pointers are + // small: `ptr / 2` in the low cell limb, high cell limb zero. self.memory_bridge .read( - MemoryAddress::new(d.clone(), byte_ptr_to_u16_ptr::(cols.rd_ptr)), + MemoryAddress::new(d.clone(), reg_byte_ptr_to_cell_ptr_limbs::(cols.rd_ptr)), pack_u8_block::(&rd_full), timestamp_pp(), &cols.rd_aux, @@ -309,7 +326,7 @@ impl VmAdapterAir for DeferralCallAdapterAir { self.memory_bridge .read( - MemoryAddress::new(d.clone(), byte_ptr_to_u16_ptr::(cols.rs_ptr)), + MemoryAddress::new(d.clone(), reg_byte_ptr_to_cell_ptr_limbs::(cols.rs_ptr)), pack_u8_block::(&rs_full), timestamp_pp(), &cols.rs_aux, @@ -338,20 +355,57 @@ impl VmAdapterAir for DeferralCallAdapterAir { } } + // Convert the heap `input`/`output` base *byte* pointers (read from registers) into + // AS-native u16 *cell* pointer limbs `[cell_lo, cell_hi]`. The byte pointers are + // little-endian 16-bit limbs packed from the low 4 register bytes. + let input_byte_limbs: [AB::Expr; 2] = [ + pack_u8_pair(cols.rs_val[0].into(), cols.rs_val[1].into()), + pack_u8_pair(cols.rs_val[2].into(), cols.rs_val[3].into()), + ]; + let input_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + input_byte_limbs, + cols.input_cell_carry, + self.address_bits, + ctx.instruction.is_valid.clone(), + ); + let output_byte_limbs: [AB::Expr; 2] = [ + pack_u8_pair(cols.rd_val[0].into(), cols.rd_val[1].into()), + pack_u8_pair(cols.rd_val[2].into(), cols.rd_val[3].into()), + ]; + let output_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + output_byte_limbs, + cols.output_cell_carry, + self.address_bits, + ctx.instruction.is_valid.clone(), + ); + + // Cell offset (in u16 cells) between consecutive heap blocks. + let heap_cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + // Accumulators are read then updated in the deferral address space, // using deferral_idx (instruction immediate / operand c) to determine // the accumulator memory address. - let input_ptr = bytes_to_f(&cols.rs_val); - let output_ptr = bytes_to_f(&cols.rd_val); - let deferral_idx = ctx.instruction.immediate; let deferral_as = AB::Expr::from_u32(DEFERRAL_AS); - // Accumulators are consecutive DEFERRAL_AS cell ranges. - let digest_size = AB::F::from_usize(DIGEST_SIZE); - let num_accumulators = AB::F::from_usize(NUM_ACCUMULATORS_PER_IDX); - let input_acc_ptr = deferral_idx.clone() * num_accumulators * digest_size; - let output_acc_ptr = input_acc_ptr.clone() + AB::Expr::from(digest_size); + // Accumulators are consecutive DEFERRAL_AS *cell* ranges. The base accumulator pointer + // `NUM_ACCUMULATORS_PER_IDX * deferral_idx * DIGEST_SIZE` and every accumulator cell access + // are bounded below 2^16 because the count bus constrains `deferral_idx < MAX_DEF_CIRCUITS` + // (see the static assert in `super`). The pointer therefore fits entirely in the low cell + // limb with a high limb of zero, so — unlike the heap pointers — no limb decomposition, + // range checks, or add carries are needed: each cell pointer is just `[base + offset, 0]`. + let acc_base_ptr = + deferral_idx.clone() * AB::Expr::from_usize(NUM_ACCUMULATORS_PER_IDX * DIGEST_SIZE); + let acc_cell_ptr = |offset: usize| -> [AB::Expr; 2] { + [ + acc_base_ptr.clone() + AB::Expr::from_usize(offset), + AB::Expr::ZERO, + ] + }; let DeferralCallReads { input_commit, @@ -383,20 +437,24 @@ impl VmAdapterAir for DeferralCallAdapterAir { let input_commit_chunks = split_byte_memory_ops::<_, COMMIT_NUM_BYTES, COMMIT_MEMORY_OPS>(input_commit); - for (chunk_idx, (data, aux)) in input_commit_chunks - .into_iter() - .zip(&cols.input_commit_aux) - .enumerate() + for (chunk_idx, (data, aux, carry)) in izip!( + input_commit_chunks, + &cols.input_commit_aux, + &cols.input_commit_add_carry + ) + .enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + input_base_cell.clone(), + chunk_idx as u32 * heap_cell_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new( - e.clone(), - byte_ptr_to_u16_ptr::( - input_ptr.clone() - + AB::Expr::from_usize(chunk_idx * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e.clone(), block_cell_ptr), pack_u8_block::(&data), timestamp_pp(), aux, @@ -415,7 +473,7 @@ impl VmAdapterAir for DeferralCallAdapterAir { .read( MemoryAddress::new( deferral_as.clone(), - input_acc_ptr.clone() + AB::Expr::from_usize(chunk_idx * BLOCK_FE_WIDTH), + acc_cell_ptr(chunk_idx * BLOCK_FE_WIDTH), ), data, timestamp_pp(), @@ -435,7 +493,7 @@ impl VmAdapterAir for DeferralCallAdapterAir { .read( MemoryAddress::new( deferral_as.clone(), - output_acc_ptr.clone() + AB::Expr::from_usize(chunk_idx * BLOCK_FE_WIDTH), + acc_cell_ptr(DIGEST_SIZE + chunk_idx * BLOCK_FE_WIDTH), ), data, timestamp_pp(), @@ -449,20 +507,24 @@ impl VmAdapterAir for DeferralCallAdapterAir { split_byte_memory_ops::<_, OUTPUT_TOTAL_BYTES, OUTPUT_TOTAL_MEMORY_OPS>( output_commit_and_len, ); - for (chunk_idx, (data, aux)) in output_commit_and_len_chunks - .into_iter() - .zip(&cols.output_commit_and_len_aux) - .enumerate() + for (chunk_idx, (data, aux, carry)) in izip!( + output_commit_and_len_chunks, + &cols.output_commit_and_len_aux, + &cols.output_add_carry + ) + .enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + output_base_cell.clone(), + chunk_idx as u32 * heap_cell_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .write( - MemoryAddress::new( - e.clone(), - byte_ptr_to_u16_ptr::( - output_ptr.clone() - + AB::Expr::from_usize(chunk_idx * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e.clone(), block_cell_ptr), pack_u8_block::(&data), timestamp_pp(), aux, @@ -481,7 +543,7 @@ impl VmAdapterAir for DeferralCallAdapterAir { .write( MemoryAddress::new( deferral_as.clone(), - input_acc_ptr.clone() + AB::Expr::from_usize(chunk_idx * BLOCK_FE_WIDTH), + acc_cell_ptr(chunk_idx * BLOCK_FE_WIDTH), ), data, timestamp_pp(), @@ -501,7 +563,7 @@ impl VmAdapterAir for DeferralCallAdapterAir { .write( MemoryAddress::new( deferral_as.clone(), - output_acc_ptr.clone() + AB::Expr::from_usize(chunk_idx * BLOCK_FE_WIDTH), + acc_cell_ptr(DIGEST_SIZE + chunk_idx * BLOCK_FE_WIDTH), ), data, timestamp_pp(), diff --git a/extensions/deferral/circuit/src/call/mod.rs b/extensions/deferral/circuit/src/call/mod.rs index 50dbefaff8..69ab32f9de 100644 --- a/extensions/deferral/circuit/src/call/mod.rs +++ b/extensions/deferral/circuit/src/call/mod.rs @@ -1,9 +1,19 @@ use openvm_circuit::arch::{VmAirWrapper, VmChipWrapper}; +use openvm_deferral_guest::MAX_DEF_CIRCUITS; use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE; /// Number of accumulator digests stored for each `deferral_idx`. pub(in crate::call) const NUM_ACCUMULATORS_PER_IDX: usize = 2; +// The accumulator cell pointers live entirely in the low 16-bit pointer limb. The count bus +// constrains `deferral_idx < num_deferral_circuits <= MAX_DEF_CIRCUITS`, so every accumulator cell +// pointer (`NUM_ACCUMULATORS_PER_IDX * deferral_idx * DIGEST_SIZE` plus an offset +// `< NUM_ACCUMULATORS_PER_IDX * DIGEST_SIZE`) is strictly below +// `NUM_ACCUMULATORS_PER_IDX * MAX_DEF_CIRCUITS * DIGEST_SIZE`. As long as that bound fits in 2^16, +// the high pointer limb is identically zero. +const _: () = + assert!(NUM_ACCUMULATORS_PER_IDX * (MAX_DEF_CIRCUITS as usize) * DIGEST_SIZE <= (1 << 16)); + #[inline(always)] pub(in crate::call) const fn accumulator_ptrs(deferral_idx: u32) -> (u32, u32) { let input_acc_ptr = (NUM_ACCUMULATORS_PER_IDX as u32) * deferral_idx * DIGEST_SIZE as u32; diff --git a/extensions/deferral/circuit/src/call/tests.rs b/extensions/deferral/circuit/src/call/tests.rs index 30d94589e8..bc30d722b2 100644 --- a/extensions/deferral/circuit/src/call/tests.rs +++ b/extensions/deferral/circuit/src/call/tests.rs @@ -3,7 +3,8 @@ use std::{array::from_fn, sync::Arc}; use openvm_circuit::arch::{ deferral::{DeferralState, InputMapVal}, testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, + memory::{gen_distinct_register_pointers, gen_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, Arena, MatrixRecordArena, MemoryConfig, PreflightExecutor, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, }; @@ -30,7 +31,9 @@ use { super::{DeferralCallAdapterRecord, DeferralCallChipGpu, DeferralCallCoreRecord}, crate::{count::DeferralCircuitCountChipGpu, poseidon2::DeferralPoseidon2ChipGpu}, openvm_circuit::arch::{ - testing::{default_bitwise_lookup_bus, GpuChipTestBuilder, GpuTestChipHarness}, + testing::{ + default_bitwise_lookup_bus, dummy_range_checker, GpuChipTestBuilder, GpuTestChipHarness, + }, DenseRecordArena, EmptyAdapterCoreLayout, }, openvm_cuda_common::d_buffer::DeviceBuffer, @@ -106,7 +109,6 @@ struct CudaHarnessBundle { fn test_memory_config() -> MemoryConfig { let mut config = MemoryConfig::default(); - config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << config.pointer_max_bits; config.addr_spaces[DEFERRAL_AS as usize].num_cells = 1 << 20; config } @@ -149,8 +151,7 @@ fn set_and_execute_call( RA: Arena, E: PreflightExecutor, { - let rd = gen_pointer(rng, MEMORY_BLOCK_BYTES); - let rs = gen_pointer(rng, MEMORY_BLOCK_BYTES); + let [rd, rs] = gen_distinct_register_pointers(rng, MEMORY_BLOCK_BYTES); let output_ptr = gen_pointer(rng, MEMORY_BLOCK_BYTES); let input_ptr = gen_pointer(rng, MEMORY_BLOCK_BYTES); let deferral_idx = rng.random_range(0..num_deferrals); @@ -278,6 +279,7 @@ fn create_cpu_harness( tester.execution_bridge(), tester.memory_bridge(), bitwise_bus, + tester.range_checker().bus(), tester.address_bits(), ), DeferralCallCoreAir::new(count_bus, poseidon2_bus, bitwise_bus), @@ -285,7 +287,11 @@ fn create_cpu_harness( let executor = DeferralCallExecutor::new(DeferralCallAdapterExecutor, fns); let chip = DeferralCallChip::new( DeferralCallCoreFiller::new( - DeferralCallAdapterFiller::new(bitwise_chip.clone(), tester.address_bits()), + DeferralCallAdapterFiller::new( + bitwise_chip.clone(), + tester.range_checker(), + tester.address_bits(), + ), count_chip.clone(), poseidon2_chip.clone(), bitwise_chip.clone(), @@ -328,6 +334,7 @@ fn create_cuda_harness( tester.execution_bridge(), tester.memory_bridge(), bitwise_bus, + tester.cpu_range_checker().bus(), tester.address_bits(), ), DeferralCallCoreAir::new(count_bus, poseidon2_bus, bitwise_bus), @@ -335,7 +342,13 @@ fn create_cuda_harness( let executor = DeferralCallExecutor::new(DeferralCallAdapterExecutor, fns); let cpu_chip = DeferralCallChip::new( DeferralCallCoreFiller::new( - DeferralCallAdapterFiller::new(dummy_bitwise_chip.clone(), tester.address_bits()), + DeferralCallAdapterFiller::new( + dummy_bitwise_chip.clone(), + // Dummy range checker: the GPU kernel already emits the AS-pointer range-check + // counts; using the real (hybrid) range checker here would double-count them. + dummy_range_checker(tester.cpu_range_checker().bus()), + tester.address_bits(), + ), count_chip_cpu, poseidon2_chip_cpu, dummy_bitwise_chip, diff --git a/extensions/deferral/circuit/src/call/trace.rs b/extensions/deferral/circuit/src/call/trace.rs index b34d5d9673..c26322763c 100644 --- a/extensions/deferral/circuit/src/call/trace.rs +++ b/extensions/deferral/circuit/src/call/trace.rs @@ -5,7 +5,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterTraceExecutor, AdapterTraceFiller, EmptyAdapterCoreLayout, ExecutionError, PreflightExecutor, RecordArena, TraceFiller, VmField, VmStateMut, - BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, + BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{ @@ -17,7 +17,8 @@ use openvm_circuit::{ }, }; use openvm_circuit_primitives::{ - bitwise_op_lookup::SharedBitwiseOperationLookupChip, AlignedBytesBorrow, + bitwise_op_lookup::SharedBitwiseOperationLookupChip, var_range::SharedVariableRangeCheckerChip, + AlignedBytesBorrow, }; use openvm_deferral_transpiler::DeferralOpcode; use openvm_instructions::{ @@ -28,7 +29,9 @@ use openvm_instructions::{ RV64_WORD_NUM_LIMBS, }, }; -use openvm_riscv_circuit::adapters::{rv64_bytes_to_u32, tracing_read, tracing_write}; +use openvm_riscv_circuit::adapters::{ + compute_pointer_carries, rv64_bytes_to_u32, tracing_read, tracing_write, +}; use openvm_stark_backend::p3_field::PrimeField32; use super::accumulator_ptrs; @@ -268,6 +271,7 @@ pub struct DeferralCallAdapterExecutor; #[derive(Clone, derive_new::new)] pub struct DeferralCallAdapterFiller { bitwise_lookup_chip: SharedBitwiseOperationLookupChip, + range_checker_chip: SharedVariableRangeCheckerChip, address_bits: usize, } @@ -432,6 +436,28 @@ impl AdapterTraceFiller for DeferralCallAdapterFiller { } } + // Byte -> cell pointer conversion carries and per-block cell-offset carries for the heap + // `input`/`output` pointers, plus matching range-check counts. Must read the record values + // here before any column overwrites; carry columns are written at the very end. + let heap_cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let (input_conv, input_add) = compute_pointer_carries( + &self.range_checker_chip, + record.rs_val, + COMMIT_MEMORY_OPS, + heap_cell_stride, + self.address_bits, + ); + let (output_conv, output_add) = compute_pointer_carries( + &self.range_checker_chip, + record.rd_val, + OUTPUT_TOTAL_MEMORY_OPS, + heap_cell_stride, + self.address_bits, + ); + + // The DEFERRAL_AS accumulator cell pointers are bounded below 2^16 (see the static assert + // in `super`), so they need no limb decomposition, range checks, or add-carry columns. + // Timestamps in AIR are assigned in strict sequence starting from // `from_state.timestamp`; mirror that exact sequence in reverse here. let timestamp_delta = @@ -511,5 +537,23 @@ impl AdapterTraceFiller for DeferralCallAdapterFiller { adapter_row.rd_ptr = record.rd_ptr; adapter_row.from_state.timestamp = F::from_u32(record.from_timestamp); adapter_row.from_state.pc = F::from_u32(record.from_pc); + + // Pointer-conversion / block-offset carry columns (computed above, after all record reads). + adapter_row.input_cell_carry = F::from_u32(input_conv); + adapter_row.output_cell_carry = F::from_u32(output_conv); + for (col, &c) in adapter_row + .input_commit_add_carry + .iter_mut() + .zip(input_add.iter()) + { + *col = F::from_u32(c); + } + for (col, &c) in adapter_row + .output_add_carry + .iter_mut() + .zip(output_add.iter()) + { + *col = F::from_u32(c); + } } } diff --git a/extensions/deferral/circuit/src/extension/mod.rs b/extensions/deferral/circuit/src/extension/mod.rs index 0e474cabc3..00160e7460 100644 --- a/extensions/deferral/circuit/src/extension/mod.rs +++ b/extensions/deferral/circuit/src/extension/mod.rs @@ -139,6 +139,7 @@ where } }; + let range_bus = inventory.range_checker().bus; let base_num_airs = inventory.num_airs(); let address_bits = to_byte_ptr_bits(inventory.pointer_max_bits()); @@ -149,7 +150,13 @@ where assert_eq!(inventory.num_airs() - base_num_airs, CALL_AIR_REL_IDX); inventory.add_air(DeferralCallAir::new( - DeferralCallAdapterAir::new(execution_bridge, memory_bridge, bitwise_bus, address_bits), + DeferralCallAdapterAir::new( + execution_bridge, + memory_bridge, + bitwise_bus, + range_bus, + address_bits, + ), DeferralCallCoreAir::new(count_bus, poseidon2_bus, bitwise_bus), )); @@ -160,6 +167,7 @@ where count_bus, poseidon2_bus, bitwise_bus, + range_bus, address_bits, )); @@ -211,7 +219,11 @@ where inventory.next_air::()?; inventory.add_executor_chip(DeferralCallChip::new( DeferralCallCoreFiller::new( - DeferralCallAdapterFiller::new(bitwise_lu.clone(), address_bits), + DeferralCallAdapterFiller::new( + bitwise_lu.clone(), + range_checker.clone(), + address_bits, + ), count_chip.clone(), poseidon2_chip.clone(), bitwise_lu.clone(), @@ -226,6 +238,7 @@ where count_chip.clone(), poseidon2_chip.clone(), bitwise_lu, + range_checker.clone(), address_bits, ), mem_helper, diff --git a/extensions/deferral/circuit/src/output/air.rs b/extensions/deferral/circuit/src/output/air.rs index 3e76b8c75d..782300e4fd 100644 --- a/extensions/deferral/circuit/src/output/air.rs +++ b/extensions/deferral/circuit/src/output/air.rs @@ -2,16 +2,17 @@ use std::{array::from_fn, borrow::Borrow}; use itertools::{izip, Itertools}; use openvm_circuit::{ - arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES}, + arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE}, system::memory::{ offline_checker::{pack_u8_block, MemoryBridge, MemoryReadAuxCols, MemoryWriteAuxCols}, - MemoryAddress, + MemoryAddress, POINTER_MAX_BITS, }, }; use openvm_circuit_primitives::{ bitwise_op_lookup::BitwiseOperationLookupBus, utils::{assert_array_eq, not}, - ColumnsAir, StructReflection, StructReflectionHelper, + var_range::VariableRangeCheckerBus, + ColumnsAir, StructReflection, StructReflectionHelper, U16_BITS, }; use openvm_circuit_primitives_derive::AlignedBorrow; use openvm_deferral_transpiler::DeferralOpcode; @@ -20,7 +21,10 @@ use openvm_instructions::{ riscv::{RV64_BYTE_BITS, RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_WORD_NUM_LIMBS}, LocalOpcode, }; -use openvm_riscv_circuit::adapters::{byte_ptr_to_u16_ptr, expand_to_rv64_register}; +use openvm_riscv_circuit::adapters::{ + byte_ptr_to_u16_ptr, eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, + expand_to_rv64_register, pack_u8_pair, reg_byte_ptr_to_cell_ptr_limbs, +}; use openvm_stark_backend::{ interaction::InteractionBuilder, p3_air::{Air, AirBuilder, BaseAir}, @@ -84,6 +88,15 @@ pub struct DeferralOutputCols { // Capacity of the permutation of write_bytes and the previous row's capacity on // non-last rows, compression on the last row. pub poseidon2_res: [T; DIGEST_SIZE], + + // Carry for converting the heap `input` base *byte* pointer (read on the first row) to + // AS-native u16 *cell* pointer limbs, plus per-block cell-offset carries. + pub input_cell_carry: T, + pub input_add_carry: [T; OUTPUT_TOTAL_MEMORY_OPS], + + // Per-row output write *cell* pointer limbs `[lo16, hi16]`. On write rows this equals + // `(output_ptr + (section_idx - 1) * DIGEST_SIZE) / 2` (cell = byte / 2). + pub write_cell_ptr: [T; 2], } #[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)] @@ -94,6 +107,7 @@ pub struct DeferralOutputAir { pub count_bus: DeferralCircuitCountBus, pub poseidon2_bus: DeferralPoseidon2Bus, pub bitwise_bus: BitwiseOperationLookupBus, + pub range_bus: VariableRangeCheckerBus, pub address_bits: usize, } @@ -281,9 +295,13 @@ where let rd_full = expand_to_rv64_register(&local.rd_val); let rs_full = expand_to_rv64_register(&local.rs_val); + // Register byte pointers are small: `ptr / 2` in the low cell limb, high cell limb zero. self.memory_bridge .read( - MemoryAddress::new(d.clone(), byte_ptr_to_u16_ptr::(local.rd_ptr)), + MemoryAddress::new( + d.clone(), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), + ), pack_u8_block::(&rd_full), local.from_state.timestamp, &local.rd_aux, @@ -292,7 +310,10 @@ where self.memory_bridge .read( - MemoryAddress::new(d.clone(), byte_ptr_to_u16_ptr::(local.rs_ptr)), + MemoryAddress::new( + d.clone(), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs_ptr), + ), pack_u8_block::(&rs_full), local.from_state.timestamp + AB::Expr::ONE, &local.rs_aux, @@ -302,7 +323,6 @@ where // Constrain memory reads and writes using the MemoryBridge. a and b are // register pointers whose values are read first, then used as heap // pointers. c carries deferral_idx. - let input_ptr = bytes_to_f(&local.rs_val); let output_ptr = bytes_to_f(&local.rd_val); let output_len_full = from_fn(|i| { if i < F_NUM_BYTES { @@ -318,20 +338,43 @@ where output_commit_and_len, ); - for (chunk_idx, (data, aux)) in output_commit_and_len_chunks - .into_iter() - .zip(&local.output_commit_and_len_aux) - .enumerate() + // Cell offset (in u16 cells) between consecutive heap read blocks. + let heap_cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert the `input` base *byte* pointer (read on the first row) to AS-native u16 *cell* + // pointer limbs. The byte pointer is little-endian 16-bit limbs from the low 4 register + // bytes. + let input_byte_limbs: [AB::Expr; 2] = [ + pack_u8_pair(local.rs_val[0].into(), local.rs_val[1].into()), + pack_u8_pair(local.rs_val[2].into(), local.rs_val[3].into()), + ]; + let input_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + input_byte_limbs, + local.input_cell_carry, + self.address_bits, + local.is_first.into(), + ); + + for (chunk_idx, (data, aux, carry)) in izip!( + output_commit_and_len_chunks, + &local.output_commit_and_len_aux, + &local.input_add_carry + ) + .enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + input_base_cell.clone(), + chunk_idx as u32 * heap_cell_stride, + *carry, + local.is_first.into(), + ); self.memory_bridge .read( - MemoryAddress::new( - e.clone(), - byte_ptr_to_u16_ptr::( - input_ptr.clone() - + AB::Expr::from_usize(chunk_idx * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e.clone(), block_cell_ptr), pack_u8_block::(&data), local.from_state.timestamp + AB::Expr::from_usize(2 + chunk_idx), aux, @@ -342,6 +385,29 @@ where let write_bytes_chunks = split_byte_memory_ops::<_, DIGEST_SIZE, DIGEST_BYTE_MEMORY_OPS>(local.sponge_inputs); let section_idx_minus_one = local.section_idx - AB::Expr::ONE; + let is_write_row = local.is_valid - local.is_first; + + // Output write *cell* pointer for this row: `(output_ptr + (section_idx - 1) * DIGEST_SIZE) + // / 2`. The composed cell-pointer field expression can approach 2^31 (and thus exceed the + // field modulus), so it is witnessed as little-endian 16-bit cell-pointer limbs + // `[lo16, hi16]` and range-checked; the limbs are then constrained to equal the composed + // expression. With `lo16 < 2^16` and `hi16 < 2^(POINTER_MAX_BITS - 16)` the decomposition + // is the unique canonical representative in `[0, 2^31)`. + let write_cell_lo = local.write_cell_ptr[0]; + let write_cell_hi = local.write_cell_ptr[1]; + let composed_write_cell = byte_ptr_to_u16_ptr::( + output_ptr.clone() + section_idx_minus_one.clone() * AB::Expr::from_usize(DIGEST_SIZE), + ); + builder.when(is_write_row.clone()).assert_zero( + write_cell_lo + write_cell_hi * AB::Expr::from_u32(1 << U16_BITS) - composed_write_cell, + ); + self.range_bus + .range_check(write_cell_lo, U16_BITS) + .eval(builder, is_write_row.clone()); + self.range_bus + .range_check(write_cell_hi, POINTER_MAX_BITS - U16_BITS) + .eval(builder, is_write_row.clone()); + let write_cell: [AB::Expr; 2] = [write_cell_lo.into(), write_cell_hi.into()]; for (chunk_idx, (data, aux)) in write_bytes_chunks .into_iter() @@ -351,21 +417,16 @@ where for bytes in data.chunks(2) { self.bitwise_bus .send_range(bytes[0], bytes[1]) - .eval(builder, local.is_valid - local.is_first); + .eval(builder, is_write_row.clone()); } let data_expr: [AB::Expr; MEMORY_BLOCK_BYTES] = from_fn(|i| data[i].into()); + // DIGEST_BYTE_MEMORY_OPS == 1, so there is a single write block per row at + // `write_cell`. + debug_assert_eq!(chunk_idx, 0); self.memory_bridge .write( - MemoryAddress::new( - e.clone(), - byte_ptr_to_u16_ptr::( - output_ptr.clone() - + (section_idx_minus_one.clone() - * AB::Expr::from_usize(DIGEST_SIZE)) - + AB::Expr::from_usize(chunk_idx * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e.clone(), write_cell.clone()), pack_u8_block::(&data_expr), local.from_state.timestamp + AB::Expr::from_usize(2 + OUTPUT_TOTAL_MEMORY_OPS + chunk_idx) @@ -373,7 +434,7 @@ where * AB::Expr::from_usize(DIGEST_BYTE_MEMORY_OPS)), aux, ) - .eval(builder, local.is_valid - local.is_first); + .eval(builder, is_write_row.clone()); } // Evaluate the execution interaction. Because a single opcode spans many diff --git a/extensions/deferral/circuit/src/output/tests.rs b/extensions/deferral/circuit/src/output/tests.rs index 67c41132df..110e58161c 100644 --- a/extensions/deferral/circuit/src/output/tests.rs +++ b/extensions/deferral/circuit/src/output/tests.rs @@ -3,10 +3,10 @@ use std::sync::Arc; use openvm_circuit::arch::{ deferral::{DeferralResult, DeferralState}, testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, + memory::{gen_distinct_register_pointers, gen_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, - to_byte_ptr_bits, Arena, MatrixRecordArena, MemoryConfig, PreflightExecutor, - MEMORY_BLOCK_BYTES, + Arena, MatrixRecordArena, MemoryConfig, PreflightExecutor, MEMORY_BLOCK_BYTES, }; use openvm_circuit_primitives::bitwise_op_lookup::{ BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip, @@ -28,7 +28,9 @@ use { super::{DeferralOutputChipGpu, DeferralOutputRecordMut}, crate::{count::DeferralCircuitCountChipGpu, poseidon2::DeferralPoseidon2ChipGpu}, openvm_circuit::arch::{ - testing::{default_bitwise_lookup_bus, GpuChipTestBuilder, GpuTestChipHarness}, + testing::{ + default_bitwise_lookup_bus, dummy_range_checker, GpuChipTestBuilder, GpuTestChipHarness, + }, DenseRecordArena, }, openvm_cuda_common::d_buffer::DeviceBuffer, @@ -98,18 +100,10 @@ struct CudaHarnessBundle { fn test_memory_config() -> MemoryConfig { let mut config = MemoryConfig::default(); - config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << config.pointer_max_bits; config.addr_spaces[DEFERRAL_AS as usize].num_cells = 1 << 20; config } -fn test_memory_config_cpu() -> MemoryConfig { - let mut config = test_memory_config(); - config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = - 1 << to_byte_ptr_bits(config.pointer_max_bits); - config -} - fn init_streams(tester: &mut impl TestBuilder, num_deferrals: usize) { tester.streams_mut().deferrals = vec![DeferralState::new(vec![]); num_deferrals]; } @@ -155,8 +149,7 @@ fn set_and_execute_output( RA: Arena, E: PreflightExecutor, { - let rd = gen_pointer(rng, MEMORY_BLOCK_BYTES); - let rs = gen_pointer(rng, MEMORY_BLOCK_BYTES); + let [rd, rs] = gen_distinct_register_pointers(rng, MEMORY_BLOCK_BYTES); let output_ptr = gen_pointer(rng, MEMORY_BLOCK_BYTES); let input_ptr = gen_pointer(rng, MEMORY_BLOCK_BYTES); let deferral_idx = rng.random_range(0..num_deferrals); @@ -226,6 +219,7 @@ fn create_cpu_harness(tester: &VmChipTestBuilder, num_deferrals: usize) -> Cp count_bus, poseidon2_bus, bitwise_bus, + tester.range_checker().bus(), tester.address_bits(), ); let executor = DeferralOutputExecutor::new(); @@ -234,6 +228,7 @@ fn create_cpu_harness(tester: &VmChipTestBuilder, num_deferrals: usize) -> Cp count_chip.clone(), poseidon2_chip.clone(), bitwise_chip.clone(), + tester.range_checker(), tester.address_bits(), ), tester.memory_helper(), @@ -269,6 +264,7 @@ fn create_cuda_harness(tester: &GpuChipTestBuilder, num_deferrals: usize) -> Cud count_bus, poseidon2_bus, bitwise_bus, + tester.cpu_range_checker().bus(), tester.address_bits(), ); let executor = DeferralOutputExecutor::new(); @@ -277,6 +273,9 @@ fn create_cuda_harness(tester: &GpuChipTestBuilder, num_deferrals: usize) -> Cud count_chip_cpu, poseidon2_chip_cpu, dummy_bitwise_chip, + // Dummy range checker: the GPU kernel already emits the AS-pointer range-check counts; + // using the real (hybrid) range checker here would double-count them. + dummy_range_checker(tester.cpu_range_checker().bus()), tester.address_bits(), ), tester.dummy_memory_helper(), @@ -319,7 +318,7 @@ fn create_cuda_harness(tester: &GpuChipTestBuilder, num_deferrals: usize) -> Cud #[test] fn rand_deferral_output_test() { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::::from_config(test_memory_config_cpu()); + let mut tester = VmChipTestBuilder::::from_config(test_memory_config()); let CpuHarnessBundle { mut harness, bitwise, @@ -367,7 +366,7 @@ fn rand_deferral_output_test() { #[test] fn deferral_output_non_first_row_canonicity_aux_cleared_test() { let mut rng = create_seeded_rng(); - let mut tester = VmChipTestBuilder::::from_config(test_memory_config_cpu()); + let mut tester = VmChipTestBuilder::::from_config(test_memory_config()); let CpuHarnessBundle { mut harness, bitwise, @@ -382,8 +381,7 @@ fn deferral_output_non_first_row_canonicity_aux_cleared_test() { // columns the filler skips (the bug under test) retain this value. harness.arena.trace_buffer.fill(F::from_u32(0xdead)); - let rd = gen_pointer(&mut rng, MEMORY_BLOCK_BYTES); - let rs = gen_pointer(&mut rng, MEMORY_BLOCK_BYTES); + let [rd, rs] = gen_distinct_register_pointers(&mut rng, MEMORY_BLOCK_BYTES); let output_ptr = gen_pointer(&mut rng, MEMORY_BLOCK_BYTES); let input_ptr = gen_pointer(&mut rng, MEMORY_BLOCK_BYTES); let deferral_idx = 0; diff --git a/extensions/deferral/circuit/src/output/trace.rs b/extensions/deferral/circuit/src/output/trace.rs index e5a5a88735..d78b8c905d 100644 --- a/extensions/deferral/circuit/src/output/trace.rs +++ b/extensions/deferral/circuit/src/output/trace.rs @@ -10,16 +10,17 @@ use openvm_circuit::{ arch::{ get_record_from_slice, CustomBorrow, ExecutionError, MultiRowLayout, MultiRowMetadata, PreflightExecutor, RecordArena, SizedRecord, TraceFiller, VmField, VmStateMut, - MEMORY_BLOCK_BYTES, + MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{pack_u8_block_bytes, MemoryReadAuxRecord, MemoryWriteBytesAuxRecord}, online::TracingMemory, - MemoryAuxColsFactory, + MemoryAuxColsFactory, POINTER_MAX_BITS, }, }; use openvm_circuit_primitives::{ - bitwise_op_lookup::SharedBitwiseOperationLookupChip, AlignedBytesBorrow, + bitwise_op_lookup::SharedBitwiseOperationLookupChip, var_range::SharedVariableRangeCheckerChip, + AlignedBytesBorrow, U16_BITS, }; use openvm_deferral_transpiler::DeferralOpcode; use openvm_instructions::{ @@ -31,7 +32,8 @@ use openvm_instructions::{ }, }; use openvm_riscv_circuit::adapters::{ - memory_read, read_rv64_register_as_u32, rv64_bytes_to_u32, tracing_read, tracing_write, + byte_ptr_to_u16_ptr_value, compute_pointer_carries, memory_read, read_rv64_register_as_u32, + rv64_bytes_to_u32, tracing_read, tracing_write, u32_to_ptr_limbs, }; use openvm_stark_backend::{p3_field::PrimeField32, p3_matrix::dense::RowMajorMatrix}; use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE; @@ -152,6 +154,7 @@ pub struct DeferralOutputFiller { count_chip: Arc, poseidon2_chip: Arc>, bitwise_lookup_chip: SharedBitwiseOperationLookupChip, + range_checker_chip: SharedVariableRangeCheckerChip, address_bits: usize, } @@ -379,6 +382,23 @@ where cols.output_commit_and_len_aux[chunk_idx].as_mut(), ); } + + // Byte -> cell pointer conversion carry and per-block cell-offset carries for + // the `input` base pointer (read on the first row), plus matching range checks. + let heap_cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let (input_conv, input_add) = compute_pointer_carries( + &self.range_checker_chip, + header.rs_val, + OUTPUT_TOTAL_MEMORY_OPS, + heap_cell_stride, + self.address_bits, + ); + cols.input_cell_carry = F::from_u32(input_conv); + for (carry_col, &add_carry) in + cols.input_add_carry.iter_mut().zip(input_add.iter()) + { + *carry_col = F::from_u32(add_carry); + } } else { mem_helper.fill_zero(cols.rd_aux.as_mut()); mem_helper.fill_zero(cols.rs_aux.as_mut()); @@ -432,6 +452,18 @@ where cols.write_bytes_aux[chunk_idx].as_mut(), ); } + + // Output write *cell* pointer for this row = + // `(output_ptr + (section_idx - 1) * DIGEST_SIZE) / 2`, witnessed as + // little-endian 16-bit cell-pointer limbs `[lo16, hi16]` and range-checked. + let write_byte_ptr = header.rd_val + ((row_idx - 1) * DIGEST_SIZE) as u32; + let write_cell_ptr = byte_ptr_to_u16_ptr_value(write_byte_ptr); + let write_cell_limbs = u32_to_ptr_limbs(write_cell_ptr); + self.range_checker_chip + .add_count(write_cell_limbs[0], U16_BITS); + self.range_checker_chip + .add_count(write_cell_limbs[1], POINTER_MAX_BITS - U16_BITS); + cols.write_cell_ptr = write_cell_limbs.map(F::from_u32); } cols.poseidon2_res = current_poseidon2_res; } diff --git a/extensions/ecc/circuit/src/weierstrass_chip/tests.rs b/extensions/ecc/circuit/src/weierstrass_chip/tests.rs index 8291906107..444da4ab74 100644 --- a/extensions/ecc/circuit/src/weierstrass_chip/tests.rs +++ b/extensions/ecc/circuit/src/weierstrass_chip/tests.rs @@ -6,7 +6,10 @@ use halo2curves_axiom::secp256r1; use num_bigint::BigUint; use num_traits::{FromPrimitive, Num, Zero}; use openvm_circuit::arch::{ - testing::{memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder}, + testing::{ + memory::{gen_distinct_register_pointers, gen_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, + }, Arena, MatrixRecordArena, PreflightExecutor, MEMORY_BLOCK_BYTES, }; use openvm_circuit_primitives::bigint::utils::{secp256k1_coord_prime, secp256r1_coord_prime}; @@ -224,9 +227,8 @@ mod ec_addne_tests { let ptr_as = RV64_REGISTER_AS as usize; let data_as = RV64_MEMORY_AS as usize; - let rs1_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rd_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1_ptr, rs2_ptr, rd_ptr] = + gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let p1_base_addr = gen_pointer(rng, MEMORY_BLOCK_BYTES) as u64; let p2_base_addr = gen_pointer(rng, MEMORY_BLOCK_BYTES) as u64; @@ -642,8 +644,7 @@ mod ec_double_tests { let ptr_as = RV64_REGISTER_AS as usize; let data_as = RV64_MEMORY_AS as usize; - let rs1_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rd_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1_ptr, rd_ptr] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let p1_base_addr = gen_pointer(rng, MEMORY_BLOCK_BYTES) as u64; let result_base_addr = gen_pointer(rng, MEMORY_BLOCK_BYTES) as u64; diff --git a/extensions/keccak256/circuit/build.rs b/extensions/keccak256/circuit/build.rs index 7832536924..7ed4c5c384 100644 --- a/extensions/keccak256/circuit/build.rs +++ b/extensions/keccak256/circuit/build.rs @@ -12,9 +12,11 @@ fn main() { .include_from_dep("DEP_CUDA_COMMON_INCLUDE") .include("../../../crates/circuits/primitives/cuda/include") .include("../../../crates/vm/cuda/include") + .include("../../riscv-adapters/cuda/include") .include("cuda/include") .watch("cuda") .watch("../../../crates/circuits/primitives/cuda") + .watch("../../riscv-adapters/cuda") .watch("../../../crates/vm/cuda") .library_name("tracegen_gpu_keccak256") .files_from_glob("cuda/src/*.cu"); diff --git a/extensions/keccak256/circuit/cuda/include/keccakf_op.cuh b/extensions/keccak256/circuit/cuda/include/keccakf_op.cuh index 02b6173dc9..c8b5c40122 100644 --- a/extensions/keccak256/circuit/cuda/include/keccakf_op.cuh +++ b/extensions/keccak256/circuit/cuda/include/keccakf_op.cuh @@ -38,6 +38,10 @@ template struct KeccakfOpCols { T postimage[KECCAK_WIDTH_U16S]; MemoryReadAuxCols rd_aux; MemoryBaseAuxCols buffer_word_aux[KECCAK_WIDTH_MEM_OPS]; // 25 words + // Carry for converting the base buffer byte pointer to AS-native u16 cell pointer limbs. + T buffer_cell_carry; + // Per-block carry for adding the cell offset to the base cell pointer. + T buffer_word_add_carry[KECCAK_WIDTH_MEM_OPS]; }; inline constexpr size_t NUM_KECCAKF_OP_COLS = sizeof(KeccakfOpCols); diff --git a/extensions/keccak256/circuit/cuda/include/xorin.cuh b/extensions/keccak256/circuit/cuda/include/xorin.cuh index 1dcbdb12fd..3198a77de6 100644 --- a/extensions/keccak256/circuit/cuda/include/xorin.cuh +++ b/extensions/keccak256/circuit/cuda/include/xorin.cuh @@ -17,10 +17,8 @@ struct XorinInstructionCols { T buffer_reg_ptr; T input_reg_ptr; T len_reg_ptr; - T buffer_ptr; // Low 32 bits of [buffer_reg_ptr:8]_1 as u16 cells. T buffer_ptr_limbs[RV64_PTR_U16_LIMBS]; - T input_ptr; // Low 32 bits of [input_reg_ptr:8]_1 as u16 cells. T input_ptr_limbs[RV64_PTR_U16_LIMBS]; T len; @@ -43,6 +41,16 @@ struct XorinMemoryCols { MemoryReadAuxCols buffer_bytes_read_aux_cols[keccak256::KECCAK_RATE_MEM_OPS]; MemoryWriteAuxCols buffer_bytes_write_aux_cols[keccak256::KECCAK_RATE_MEM_OPS]; + // Carry for converting the base `buffer`/`input` *byte* pointers to AS-native u16 *cell* + // pointer limbs. + T buffer_cell_carry; + T input_cell_carry; + // Per-block carry for adding the cell offset `i * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to each + // base cell pointer (block `i`'s carry into the high cell limb). One set per heap access group + // (buffer read, input read, buffer write). + T buffer_read_add_carry[keccak256::KECCAK_RATE_MEM_OPS]; + T input_read_add_carry[keccak256::KECCAK_RATE_MEM_OPS]; + T buffer_write_add_carry[keccak256::KECCAK_RATE_MEM_OPS]; }; template diff --git a/extensions/keccak256/circuit/cuda/src/keccakf_op.cu b/extensions/keccak256/circuit/cuda/src/keccakf_op.cu index 0e9008e565..b833677652 100644 --- a/extensions/keccak256/circuit/cuda/src/keccakf_op.cu +++ b/extensions/keccak256/circuit/cuda/src/keccakf_op.cu @@ -6,6 +6,7 @@ #include "primitives/histogram.cuh" #include "primitives/trace_access.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include @@ -15,7 +16,6 @@ using namespace keccakf_op; using namespace riscv; -using openvm::U16_BITS; #define KECCAKF_OP_WRITE(FIELD, VALUE) COL_WRITE_VALUE(row, KeccakfOpCols, FIELD, VALUE) #define KECCAKF_OP_WRITE_ARRAY(FIELD, VALUES) COL_WRITE_ARRAY(row, KeccakfOpCols, FIELD, VALUES) @@ -80,11 +80,20 @@ static __device__ __noinline__ void fill_keccakf_op_row( ts++; } - // Bound the high u16 cell so the low-32-bit buffer pointer fits in pointer_max_bits. - range_checker.add_count( - ptr_bound_from_high_u16(buffer_ptr_limbs[RV64_PTR_U16_LIMBS - 1], pointer_max_bits), - U16_BITS + // Byte -> cell pointer conversion carry and per-block cell-offset carries, plus the matching + // range-check counts (mirrors KeccakfOpChip::fill_trace). + uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + uint32_t add_carries[KECCAK_WIDTH_MEM_OPS]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + rec.buffer_ptr, + pointer_max_bits, + KECCAK_WIDTH_MEM_OPS, + cell_stride, + add_carries ); + KECCAKF_OP_WRITE(buffer_cell_carry, conv_carry); + KECCAKF_OP_WRITE_ARRAY(buffer_word_add_carry, add_carries); } // Main kernel for KeccakfOpChip trace generation diff --git a/extensions/keccak256/circuit/cuda/src/xorin.cu b/extensions/keccak256/circuit/cuda/src/xorin.cu index 81819b8de8..96606da8ad 100644 --- a/extensions/keccak256/circuit/cuda/src/xorin.cu +++ b/extensions/keccak256/circuit/cuda/src/xorin.cu @@ -5,6 +5,7 @@ #include "primitives/histogram.cuh" #include "primitives/trace_access.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "xorin.cuh" #include @@ -15,7 +16,6 @@ using namespace xorin; using namespace riscv; using namespace keccak256; using namespace program; -using openvm::U16_BITS; #define XORIN_WRITE(FIELD, VALUE) COL_WRITE_VALUE(row, XorinVmCols, FIELD, VALUE) #define XORIN_WRITE_ARRAY(FIELD, VALUES) COL_WRITE_ARRAY(row, XorinVmCols, FIELD, VALUES) @@ -46,7 +46,7 @@ __global__ void xorin_tracegen( assert(rec.len <= XORIN_RATE_BYTES); assert((uint64_t)rec.buffer + rec.len <= (1ULL << pointer_max_bits)); assert((uint64_t)rec.input + rec.len <= (1ULL << pointer_max_bits)); - assert(rec.len < (1U << pointer_max_bits)); + assert(rec.len < (1ULL << pointer_max_bits)); VariableRangeChecker range_checker(d_range_checker_ptr, range_checker_num_bins); MemoryAuxColsFactory mem_helper(range_checker, timestamp_max_bits); @@ -62,8 +62,6 @@ __global__ void xorin_tracegen( XORIN_WRITE(instruction.buffer_reg_ptr, rec.rd_ptr); XORIN_WRITE(instruction.input_reg_ptr, rec.rs1_ptr); XORIN_WRITE(instruction.len_reg_ptr, rec.rs2_ptr); - XORIN_WRITE(instruction.buffer_ptr, rec.buffer); - XORIN_WRITE(instruction.input_ptr, rec.input); XORIN_WRITE(instruction.len, rec.len); XORIN_WRITE(instruction.start_timestamp, rec.timestamp); @@ -158,15 +156,48 @@ __global__ void xorin_tracegen( XORIN_FILL_ZERO(mem_oc.buffer_bytes_write_aux_cols[t]); } - // Bound the high u16 cells so the low-32-bit pointers fit in pointer_max_bits. - range_checker.add_count( - ptr_bound_from_high_u16(buffer_ptr_limbs[RV64_PTR_U16_LIMBS - 1], pointer_max_bits), - U16_BITS + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors `xorin/trace.rs`. + uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + + uint32_t buffer_add_carry[KECCAK_RATE_MEM_OPS]; + uint32_t buffer_conv_carry = compute_pointer_carries( + range_checker, + rec.buffer, + pointer_max_bits, + KECCAK_RATE_MEM_OPS, + cell_stride, + buffer_add_carry ); - range_checker.add_count( - ptr_bound_from_high_u16(input_ptr_limbs[RV64_PTR_U16_LIMBS - 1], pointer_max_bits), - U16_BITS + XORIN_WRITE(mem_oc.buffer_cell_carry, buffer_conv_carry); + XORIN_WRITE_ARRAY(mem_oc.buffer_read_add_carry, buffer_add_carry); + + uint32_t input_add_carry[KECCAK_RATE_MEM_OPS]; + uint32_t input_conv_carry = compute_pointer_carries( + range_checker, + rec.input, + pointer_max_bits, + KECCAK_RATE_MEM_OPS, + cell_stride, + input_add_carry ); + XORIN_WRITE(mem_oc.input_cell_carry, input_conv_carry); + XORIN_WRITE_ARRAY(mem_oc.input_read_add_carry, input_add_carry); + + // The write reuses the converted `buffer` base cell pointer; only the per-block write add + // carries (and their range checks) are needed. The base conversion was registered above. + uint32_t buffer_write_add_carry[KECCAK_RATE_MEM_OPS]; + CellPtr buffer_cell = byte_ptr_limbs_to_cell_ptr_limbs_value( + rec.buffer & 0xffffu, rec.buffer >> openvm::U16_BITS + ); + compute_block_add_carries( + range_checker, + buffer_cell.limbs[0], + KECCAK_RATE_MEM_OPS, + cell_stride, + buffer_write_add_carry + ); + XORIN_WRITE_ARRAY(mem_oc.buffer_write_add_carry, buffer_write_add_carry); } else { // Zero-fill padding rows diff --git a/extensions/keccak256/circuit/src/keccakf_op/air.rs b/extensions/keccak256/circuit/src/keccakf_op/air.rs index 3cf10b0472..9491fb1e43 100644 --- a/extensions/keccak256/circuit/src/keccakf_op/air.rs +++ b/extensions/keccak256/circuit/src/keccakf_op/air.rs @@ -2,18 +2,18 @@ use std::{borrow::Borrow, iter}; use itertools::izip; use openvm_circuit::{ - arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES}, + arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE}, system::memory::{ offline_checker::{MemoryBridge, MemoryWriteAuxInput}, MemoryAddress, }, }; -use openvm_circuit_primitives::{var_range::VariableRangeCheckerBus, ColumnsAir, U16_BITS}; +use openvm_circuit_primitives::{var_range::VariableRangeCheckerBus, ColumnsAir}; use openvm_instructions::riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}; use openvm_keccak256_transpiler::KeccakfOpcode; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, u16_limbs_to_ptr, - RV64_PTR_U16_LIMBS, + eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, + reg_byte_ptr_to_cell_ptr_limbs, }; use openvm_stark_backend::{ interaction::{InteractionBuilder, PermutationCheckBus}, @@ -74,7 +74,8 @@ impl Air for KeccakfOpAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(rd_ptr), + // Register byte pointers are small: `rd_ptr / 2` in the low cell limb. + reg_byte_ptr_to_cell_ptr_limbs::(rd_ptr), ), buffer_ptr_data, timestamp_pp(), @@ -82,46 +83,53 @@ impl Air for KeccakfOpAir { ) .eval(builder, is_valid); - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr::( - local.buffer_ptr_limbs[RV64_PTR_U16_LIMBS - 1], - self.ptr_max_bits, - ), - U16_BITS, - ) - .eval(builder, is_valid); - let buffer_ptr = u16_limbs_to_ptr(&local.buffer_ptr_limbs); + // Convert the base `buffer` *byte* pointer to base AS-native u16 *cell* pointer limbs. + let buffer_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.buffer_ptr_limbs[i].into()); + let buffer_base_cell_ptr = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + buffer_byte_limbs, + local.buffer_cell_carry, + self.ptr_max_bits, + is_valid.into(), + ); + // Cell-pointer stride (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; // ======== Constrain new writes of `buffer` to memory ========= // Keccak state and memory both consume these values as packed u16 cells. - for (word_idx, (prev_word, post_word, base_aux)) in izip!( + for (word_idx, (prev_word, post_word, base_aux, add_carry)) in izip!( local.preimage.chunks_exact(BLOCK_FE_WIDTH), local.postimage.chunks_exact(BLOCK_FE_WIDTH), - local.buffer_word_aux + local.buffer_word_aux, + local.buffer_word_add_carry ) .enumerate() { // Safety: - // - we range checked that buffer_ptr < 2^ptr_max_bits but not that buffer_ptr + - // KECCAK_WIDTH_BYTES is in range. - // - the previous range check implies `buffer_ptr + KECCAK_WIDTH_BYTES` does not - // overflow the field `F` hence it is safe to consider `ptr` as a field element. - // - the memory_bridge.write at `ptr` consists of a receive on memory bus at a previous - // timestamp. The only way this bus interaction could balance is if there was already - // a previous valid write at `ptr`. Assuming the invariant that all previous memory - // accesses are valid and timestamp always moves forward, the new write to `ptr` must - // be valid as well. - let ptr = buffer_ptr.clone() + AB::F::from_usize(word_idx * MEMORY_BLOCK_BYTES); + // - `buffer_base_cell_ptr` is range-checked to be a canonical cell pointer below + // `2^cell_max_bits`, and each `eval_add_const_u16_limbs` range-checks the new low + // limb, so the per-block cell pointer is canonical. + // - the memory_bridge.write at this cell pointer consists of a receive on memory bus at + // a previous timestamp. The only way this bus interaction could balance is if there + // was already a previous valid write there. Assuming the invariant that all previous + // memory accesses are valid and timestamp always moves forward, the new write must be + // valid as well. + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + buffer_base_cell_ptr.clone(), + word_idx as u32 * cell_ptr_block_stride, + add_carry, + is_valid.into(), + ); let prev_data: [AB::Expr; BLOCK_FE_WIDTH] = std::array::from_fn(|i| prev_word[i].into()); let data: [AB::Expr; BLOCK_FE_WIDTH] = std::array::from_fn(|i| post_word[i].into()); self.memory_bridge .write( - MemoryAddress::new( - AB::F::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::(ptr), - ), + MemoryAddress::new(AB::F::from_u32(RV64_MEMORY_AS), block_cell_ptr), data, timestamp_pp(), MemoryWriteAuxInput::from_prev_data_exprs(&base_aux, prev_data), diff --git a/extensions/keccak256/circuit/src/keccakf_op/columns.rs b/extensions/keccak256/circuit/src/keccakf_op/columns.rs index ba6bc0be33..6ec83c1c97 100644 --- a/extensions/keccak256/circuit/src/keccakf_op/columns.rs +++ b/extensions/keccak256/circuit/src/keccakf_op/columns.rs @@ -34,6 +34,12 @@ pub struct KeccakfOpCols { /// Auxiliary columns for timestamp checking of the writes to `buffer`. The writes are done one /// word at a time, and each write requires a separate previous timestamp. pub buffer_word_aux: [MemoryBaseAuxCols; KECCAK_WIDTH_MEM_OPS], + /// Carry for converting the base `buffer` *byte* pointer to AS-native u16 *cell* pointer + /// limbs. + pub buffer_cell_carry: T, + /// Per-block carry for adding the cell offset `word_idx * (MEMORY_BLOCK_BYTES / + /// U16_CELL_SIZE)` to the base cell pointer (block's carry into the high cell limb). + pub buffer_word_add_carry: [T; KECCAK_WIDTH_MEM_OPS], } pub const NUM_KECCAKF_OP_COLS: usize = size_of::>(); diff --git a/extensions/keccak256/circuit/src/keccakf_op/tests.rs b/extensions/keccak256/circuit/src/keccakf_op/tests.rs index cd8bec1019..40f2a4af07 100644 --- a/extensions/keccak256/circuit/src/keccakf_op/tests.rs +++ b/extensions/keccak256/circuit/src/keccakf_op/tests.rs @@ -8,8 +8,8 @@ use itertools::Itertools; use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, + memory::{gen_pointer, gen_register_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, Arena, ExecutionBridge, PreflightExecutor, MEMORY_BLOCK_BYTES, }, @@ -126,7 +126,7 @@ fn set_and_execute_single_perm>( let mut rand_buffer_arr = [0u8; MAX_LEN]; rand_buffer_arr.copy_from_slice(&rand_buffer); - let rd = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let rd = gen_register_pointer(rng, RV64_REGISTER_NUM_LIMBS); let buffer_ptr = gen_pointer(rng, MAX_LEN); tester.write_bytes( RV64_REGISTER_AS as usize, @@ -296,11 +296,11 @@ fn cuda_set_and_execute( arena: &mut DenseRecordArena, rng: &mut StdRng, ) { - use openvm_circuit::arch::testing::memory::gen_pointer; + use openvm_circuit::arch::testing::memory::{gen_pointer, gen_register_pointer}; const KECCAK_STATE_BYTES: usize = 200; - let buffer_reg = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let buffer_reg = gen_register_pointer(rng, RV64_REGISTER_NUM_LIMBS); let buffer_ptr = gen_pointer(rng, KECCAK_STATE_BYTES); tester.write_bytes( @@ -406,11 +406,11 @@ fn test_keccakf_cuda_tracegen_zero_state() { let mut harness = create_cuda_harness(&tester); - use openvm_circuit::arch::testing::memory::gen_pointer; + use openvm_circuit::arch::testing::memory::{gen_pointer, gen_register_pointer}; const KECCAK_STATE_BYTES: usize = 200; - let buffer_reg = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); + let buffer_reg = gen_register_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); let buffer_ptr = gen_pointer(&mut rng, KECCAK_STATE_BYTES); tester.write_bytes( diff --git a/extensions/keccak256/circuit/src/keccakf_op/trace.rs b/extensions/keccak256/circuit/src/keccakf_op/trace.rs index a5b6766d03..f1a992ecb2 100644 --- a/extensions/keccak256/circuit/src/keccakf_op/trace.rs +++ b/extensions/keccak256/circuit/src/keccakf_op/trace.rs @@ -13,7 +13,7 @@ use openvm_circuit::{ }, }; use openvm_circuit_primitives::{ - var_range::SharedVariableRangeCheckerChip, AlignedBytesBorrow, Chip, U16_BITS, + var_range::SharedVariableRangeCheckerChip, AlignedBytesBorrow, Chip, }; use openvm_cpu_backend::CpuBackend; use openvm_instructions::{ @@ -23,7 +23,7 @@ use openvm_instructions::{ }; use openvm_keccak256_transpiler::KeccakfOpcode; use openvm_riscv_circuit::adapters::{ - ptr_bound_from_ptr, ptr_to_field_u16_limbs, rv64_bytes_to_u32, timed_write, tracing_read, + compute_pointer_carries, ptr_to_field_u16_limbs, rv64_bytes_to_u32, timed_write, tracing_read, }; use openvm_stark_backend::{ p3_field::PrimeField32, @@ -252,10 +252,25 @@ impl TraceFiller for KeccakfOpChip { timestamp += 1; } - self.range_checker_chip.add_count( - ptr_bound_from_ptr(record.buffer_ptr, self.pointer_max_bits), - U16_BITS, + // Byte -> cell pointer conversion carry and per-block cell-offset carries, plus + // matching range-check counts. `record` is a stable clone, so writing the carry + // columns here does not alias the trace reads above. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let (conv_carry, add_carries) = compute_pointer_carries( + &self.range_checker_chip, + record.buffer_ptr, + KECCAK_WIDTH_MEM_OPS, + cell_stride, + self.pointer_max_bits, ); + local.buffer_cell_carry = F::from_u32(conv_carry); + for (col, &add_carry) in local + .buffer_word_add_carry + .iter_mut() + .zip(add_carries.iter()) + { + *col = F::from_u32(add_carry); + } }); *self.shared_records.lock().unwrap() = records; } diff --git a/extensions/keccak256/circuit/src/xorin/air.rs b/extensions/keccak256/circuit/src/xorin/air.rs index 583f666f6a..493626aae7 100644 --- a/extensions/keccak256/circuit/src/xorin/air.rs +++ b/extensions/keccak256/circuit/src/xorin/air.rs @@ -2,7 +2,7 @@ use std::borrow::Borrow; use itertools::izip; use openvm_circuit::{ - arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES}, + arch::{ExecutionBridge, ExecutionState, BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE}, system::memory::{ offline_checker::{pack_u8_block, MemoryBridge, MemoryReadAuxCols, MemoryWriteAuxCols}, MemoryAddress, @@ -10,13 +10,13 @@ use openvm_circuit::{ }; use openvm_circuit_primitives::{ bitwise_op_lookup::BitwiseOperationLookupBus, utils::not, var_range::VariableRangeCheckerBus, - ColumnsAir, U16_BITS, + ColumnsAir, }; use openvm_instructions::riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}; use openvm_keccak256_transpiler::XorinOpcode; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, u16_limbs_to_ptr, - RV64_PTR_U16_LIMBS, + eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, + reg_byte_ptr_to_cell_ptr_limbs, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -65,7 +65,7 @@ impl Air for XorinVmAir { let start_read_timestamp = self.eval_instruction(builder, local, &mem.register_aux_cols); - let start_write_timestamp = self.constrain_input_read( + let (start_write_timestamp, buffer_base_cell_ptr) = self.constrain_input_read( builder, local, start_read_timestamp, @@ -79,6 +79,7 @@ impl Air for XorinVmAir { builder, local, start_write_timestamp, + buffer_base_cell_ptr, &mem.buffer_bytes_write_aux_cols, ); } @@ -160,7 +161,8 @@ impl XorinVmAir { .read( MemoryAddress::new( AB::Expr::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(ptr), + // Register byte pointers are small: `ptr / 2` in the low cell limb. + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), value, timestamp.clone(), @@ -171,27 +173,6 @@ impl XorinVmAir { timestamp += AB::Expr::ONE; } - for top_cell in [ - instruction.buffer_ptr_limbs[RV64_PTR_U16_LIMBS - 1], - instruction.input_ptr_limbs[RV64_PTR_U16_LIMBS - 1], - ] { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr::(top_cell, self.ptr_max_bits), - U16_BITS, - ) - .eval(builder, is_enabled); - } - - builder.assert_eq( - instruction.buffer_ptr, - u16_limbs_to_ptr(&instruction.buffer_ptr_limbs), - ); - builder.assert_eq( - instruction.input_ptr, - u16_limbs_to_ptr(&instruction.input_ptr_limbs), - ); - builder.assert_eq(instruction.len, instruction.len_limb); timestamp @@ -206,31 +187,54 @@ impl XorinVmAir { start_read_timestamp: AB::Expr, input_bytes_read_aux_cols: &[MemoryReadAuxCols; KECCAK_RATE_MEM_OPS], buffer_bytes_read_aux_cols: &[MemoryReadAuxCols; KECCAK_RATE_MEM_OPS], - ) -> AB::Expr { + ) -> (AB::Expr, [AB::Expr; 2]) { let is_enabled = local.instruction.is_enabled; let mut timestamp = start_read_timestamp; + let mem = &local.mem_oc; + // Cell-pointer stride (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert the base `buffer` *byte* pointer to base AS-native u16 *cell* pointer limbs. + let buffer_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.instruction.buffer_ptr_limbs[i].into()); + let buffer_base_cell_ptr = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + buffer_byte_limbs, + mem.buffer_cell_carry, + self.ptr_max_bits, + is_enabled.into(), + ); // Constrain read of buffer bytes // Timestamp increases by <= (136/8) = 17 - for (i, (input, mem_aux)) in izip!( + for (i, (input, mem_aux, add_carry)) in izip!( local .sponge .preimage_buffer_bytes .chunks_exact(MEMORY_BLOCK_BYTES), - buffer_bytes_read_aux_cols + buffer_bytes_read_aux_cols, + mem.buffer_read_add_carry ) .enumerate() { - let ptr = local.instruction.buffer_ptr + AB::F::from_usize(i * MEMORY_BLOCK_BYTES); let is_padding = local.sponge.is_padding_bytes[i]; let should_read = is_enabled * not(is_padding); + // The cell-offset add (and its range check) is gated by `is_enabled` (degree 1) rather + // than `should_read` (degree 2) to keep the constraint degree within bounds; the filler + // therefore registers an add-carry range check for *every* block, padding or not. + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + buffer_base_cell_ptr.clone(), + i as u32 * cell_ptr_block_stride, + add_carry, + is_enabled.into(), + ); self.memory_bridge .read( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::(ptr), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), pack_u8_block::(&[ input[0].into(), input[1].into(), @@ -249,24 +253,42 @@ impl XorinVmAir { timestamp += not(is_padding); } + // Convert the base `input` *byte* pointer to base AS-native u16 *cell* pointer limbs. + let input_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.instruction.input_ptr_limbs[i].into()); + let input_base_cell_ptr = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + input_byte_limbs, + mem.input_cell_carry, + self.ptr_max_bits, + is_enabled.into(), + ); + // Constrain read of input_bytes // Timestamp increases by at most (136/8) = 17 - for (i, (input, mem_aux)) in izip!( + for (i, (input, mem_aux, add_carry)) in izip!( local.sponge.input_bytes.chunks_exact(MEMORY_BLOCK_BYTES), - input_bytes_read_aux_cols + input_bytes_read_aux_cols, + mem.input_read_add_carry ) .enumerate() { - let ptr = local.instruction.input_ptr + AB::F::from_usize(i * MEMORY_BLOCK_BYTES); let is_padding = local.sponge.is_padding_bytes[i]; let should_read = is_enabled * not(is_padding); + // Gated by `is_enabled` (see buffer-read note); filler range-checks every block. + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + input_base_cell_ptr.clone(), + i as u32 * cell_ptr_block_stride, + add_carry, + is_enabled.into(), + ); self.memory_bridge .read( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::(ptr), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), pack_u8_block::(&[ input[0].into(), input[1].into(), @@ -285,7 +307,7 @@ impl XorinVmAir { timestamp += not(is_padding); } - timestamp + (timestamp, buffer_base_cell_ptr) } #[inline] @@ -321,31 +343,40 @@ impl XorinVmAir { builder: &mut AB, local: &XorinVmCols, start_write_timestamp: AB::Expr, + buffer_base_cell_ptr: [AB::Expr; 2], mem_aux: &[MemoryWriteAuxCols; KECCAK_RATE_MEM_OPS], ) { let mut timestamp = start_write_timestamp; let is_enabled = local.instruction.is_enabled; + // Cell-pointer stride (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; // Constrain write of buffer bytes - for (i, (output, mem_aux)) in izip!( + for (i, (output, mem_aux, add_carry)) in izip!( local .sponge .postimage_buffer_bytes .chunks_exact(MEMORY_BLOCK_BYTES), - mem_aux + mem_aux, + local.mem_oc.buffer_write_add_carry ) .enumerate() { let is_padding = local.sponge.is_padding_bytes[i]; let should_write = is_enabled * not(is_padding); - let ptr = local.instruction.buffer_ptr + AB::F::from_usize(i * MEMORY_BLOCK_BYTES); + // Gated by `is_enabled` (see buffer-read note); filler range-checks every block. + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + buffer_base_cell_ptr.clone(), + i as u32 * cell_ptr_block_stride, + add_carry, + is_enabled.into(), + ); self.memory_bridge .write( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::(ptr), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), pack_u8_block::(&[ output[0].into(), output[1].into(), diff --git a/extensions/keccak256/circuit/src/xorin/columns.rs b/extensions/keccak256/circuit/src/xorin/columns.rs index 971d9d0066..bab7f23946 100644 --- a/extensions/keccak256/circuit/src/xorin/columns.rs +++ b/extensions/keccak256/circuit/src/xorin/columns.rs @@ -25,10 +25,8 @@ pub struct XorinInstructionCols { pub buffer_reg_ptr: T, pub input_reg_ptr: T, pub len_reg_ptr: T, - pub buffer_ptr: T, /// Low 32 bits of the `rs0` register as u16 cells. pub buffer_ptr_limbs: [T; RV64_PTR_U16_LIMBS], - pub input_ptr: T, /// Low 32 bits of the `rs1` register as u16 cells. pub input_ptr_limbs: [T; RV64_PTR_U16_LIMBS], pub len: T, @@ -54,6 +52,16 @@ pub struct XorinMemoryCols { pub input_bytes_read_aux_cols: [MemoryReadAuxCols; KECCAK_RATE_MEM_OPS], pub buffer_bytes_read_aux_cols: [MemoryReadAuxCols; KECCAK_RATE_MEM_OPS], pub buffer_bytes_write_aux_cols: [MemoryWriteAuxCols; KECCAK_RATE_MEM_OPS], + /// Carry for converting the base `buffer`/`input` *byte* pointers to AS-native u16 *cell* + /// pointer limbs. + pub buffer_cell_carry: T, + pub input_cell_carry: T, + /// Per-block carry for adding the cell offset `i * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `i`'s carry into the high cell limb). One set per heap + /// access group (buffer read, input read, buffer write). + pub buffer_read_add_carry: [T; KECCAK_RATE_MEM_OPS], + pub input_read_add_carry: [T; KECCAK_RATE_MEM_OPS], + pub buffer_write_add_carry: [T; KECCAK_RATE_MEM_OPS], } pub const NUM_XORIN_VM_COLS: usize = size_of::>(); diff --git a/extensions/keccak256/circuit/src/xorin/tests.rs b/extensions/keccak256/circuit/src/xorin/tests.rs index ea27350250..ca2eb24aed 100644 --- a/extensions/keccak256/circuit/src/xorin/tests.rs +++ b/extensions/keccak256/circuit/src/xorin/tests.rs @@ -123,10 +123,8 @@ fn set_and_execute>( let mut rand_input_arr = [0u8; MAX_LEN]; rand_input_arr.copy_from_slice(&rand_input); - use openvm_circuit::arch::testing::memory::gen_pointer; - let rd = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + use openvm_circuit::arch::testing::memory::{gen_distinct_register_pointers, gen_pointer}; + let [rd, rs1, rs2] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); // Align buffer/input pointers to MEMORY_BLOCK_BYTES-byte blocks for memory bus compatibility let num_blocks = buffer_length.div_ceil(MEMORY_BLOCK_BYTES); @@ -345,16 +343,15 @@ fn cuda_set_and_execute( rng: &mut StdRng, len: Option, ) { - use openvm_circuit::arch::testing::memory::gen_pointer; + use openvm_circuit::arch::testing::memory::{gen_distinct_register_pointers, gen_pointer}; let len = len.unwrap_or_else(|| rng.random_range(1..=KECCAK_RATE_MEM_OPS) * MEMORY_BLOCK_BYTES); if len == 0 { return; } - let buffer_reg = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let input_reg = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let len_reg = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [buffer_reg, input_reg, len_reg] = + gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let buffer_ptr = gen_pointer(rng, len); let input_ptr = gen_pointer(rng, len); diff --git a/extensions/keccak256/circuit/src/xorin/trace.rs b/extensions/keccak256/circuit/src/xorin/trace.rs index 887c7e2828..3e762b978d 100644 --- a/extensions/keccak256/circuit/src/xorin/trace.rs +++ b/extensions/keccak256/circuit/src/xorin/trace.rs @@ -11,7 +11,7 @@ use openvm_circuit::{ MemoryAuxColsFactory, }, }; -use openvm_circuit_primitives::{AlignedBytesBorrow, U16_BITS}; +use openvm_circuit_primitives::AlignedBytesBorrow; use openvm_instructions::{ instruction::Instruction, program::DEFAULT_PC_STEP, @@ -19,8 +19,9 @@ use openvm_instructions::{ }; use openvm_keccak256_transpiler::XorinOpcode; use openvm_riscv_circuit::adapters::{ - ptr_bound_from_ptr, ptr_to_field_u16_limbs, read_rv64_register_as_u32, rv64_bytes_to_u32, - tracing_read, tracing_write, + byte_ptr_limbs_to_cell_ptr_limbs_value, compute_block_add_carries, compute_pointer_carries, + ptr_to_field_u16_limbs, read_rv64_register_as_u32, rv64_bytes_to_u32, tracing_read, + tracing_write, u32_to_ptr_limbs, }; use openvm_stark_backend::p3_field::PrimeField32; @@ -154,9 +155,13 @@ where ); record.inner.len = rv64_bytes_to_u32(len_val); - debug_assert!(record.inner.buffer as usize + len <= (1 << self.pointer_max_bits)); - debug_assert!(record.inner.input as usize + len < (1 << self.pointer_max_bits)); - debug_assert!(record.inner.len < (1 << self.pointer_max_bits)); + debug_assert!( + (record.inner.buffer as u64) + (len as u64) <= (1u64 << self.pointer_max_bits) + ); + debug_assert!( + (record.inner.input as u64) + (len as u64) <= (1u64 << self.pointer_max_bits) + ); + debug_assert!((record.inner.len as u64) < (1u64 << self.pointer_max_bits)); // read buffer for idx in 0..num_reads { @@ -233,9 +238,7 @@ impl TraceFiller for XorinVmFiller { trace_row.instruction.buffer_reg_ptr = F::from_u32(record.rd_ptr); trace_row.instruction.input_reg_ptr = F::from_u32(record.rs1_ptr); trace_row.instruction.len_reg_ptr = F::from_u32(record.rs2_ptr); - trace_row.instruction.buffer_ptr = F::from_u32(record.buffer); trace_row.instruction.buffer_ptr_limbs = ptr_to_field_u16_limbs(record.buffer); - trace_row.instruction.input_ptr = F::from_u32(record.input); trace_row.instruction.input_ptr_limbs = ptr_to_field_u16_limbs(record.input); trace_row.instruction.len = F::from_u32(record.len); trace_row.instruction.len_limb = F::from_u8(record.len as u8); @@ -309,9 +312,67 @@ impl TraceFiller for XorinVmFiller { timestamp += 1; } - for ptr in [record.buffer, record.input] { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. `record` is a stable clone, so writing the carry columns here does + // not alias the trace reads above. + // + // The AIR gates the per-block cell-offset add by `is_enabled` (degree 1) rather than the + // per-block `should_read`/`should_write` (degree 2) to stay within the max constraint + // degree. So add carries (and their range checks) are computed for *every* block, padding + // or not, matching the AIR's `is_enabled`-gated `eval_add_const_u16_limbs` for all blocks. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let (buffer_conv, buffer_add) = compute_pointer_carries( + &self.range_checker_chip, + record.buffer, + KECCAK_RATE_MEM_OPS, + cell_stride, + self.pointer_max_bits, + ); + trace_row.mem_oc.buffer_cell_carry = F::from_u32(buffer_conv); + for (col, &add_carry) in trace_row + .mem_oc + .buffer_read_add_carry + .iter_mut() + .zip(buffer_add.iter()) + { + *col = F::from_u32(add_carry); + } + let (input_conv, input_add) = compute_pointer_carries( + &self.range_checker_chip, + record.input, + KECCAK_RATE_MEM_OPS, + cell_stride, + self.pointer_max_bits, + ); + trace_row.mem_oc.input_cell_carry = F::from_u32(input_conv); + for (col, &add_carry) in trace_row + .mem_oc + .input_read_add_carry + .iter_mut() + .zip(input_add.iter()) + { + *col = F::from_u32(add_carry); + } + // The write reuses the converted `buffer` base cell pointer; only register the per-block + // write add carries (and their range checks). The base conversion carry is already filled + // above for the buffer read group. + { + let byte_limbs = u32_to_ptr_limbs(record.buffer); + let (_conv_carry, base_cell) = byte_ptr_limbs_to_cell_ptr_limbs_value(byte_limbs); + let buffer_write_add = compute_block_add_carries( + &self.range_checker_chip, + base_cell.map(|limb| limb as u16), + KECCAK_RATE_MEM_OPS, + cell_stride, + ); + for (col, &add_carry) in trace_row + .mem_oc + .buffer_write_add_carry + .iter_mut() + .zip(buffer_write_add.iter()) + { + *col = F::from_u32(add_carry); + } } } } diff --git a/extensions/riscv-adapters/cuda/include/riscv-adapters/pointer_conv.cuh b/extensions/riscv-adapters/cuda/include/riscv-adapters/pointer_conv.cuh new file mode 100644 index 0000000000..bbe39ff94c --- /dev/null +++ b/extensions/riscv-adapters/cuda/include/riscv-adapters/pointer_conv.cuh @@ -0,0 +1,82 @@ +#pragma once + +#include "primitives/constants.h" +#include "primitives/histogram.cuh" +#include "system/memory/params.cuh" + +// CUDA mirrors of the host-side pointer-conversion value helpers in +// `openvm_riscv_circuit::adapters` (see `extensions/riscv/circuit/src/adapters/mod.rs`). +// They convert a guest *byte* pointer into AS-native u16 *cell* pointer limbs and add a small +// per-block cell offset, returning the witness carries that the heap adapters store in their +// `*_cell_carry` / `*_add_carry` columns. + +// Cell high-limb range-check bit width corresponding to a guest `byte_ptr_max_bits`. +__device__ __forceinline__ uint32_t cell_ptr_hi_bits(size_t byte_ptr_max_bits) { + return uint32_t(byte_ptr_max_bits) - openvm::U16_CELL_SIZE_BITS - openvm::U16_BITS; +} + +struct CellPtr { + // Witness boolean carry, equal to `byte_hi & 1`. + uint32_t carry; + // AS-native u16 cell pointer limbs `[cell_lo, cell_hi]`. + uint32_t limbs[2]; +}; + +// Value form of `byte_ptr_limbs_to_cell_ptr_limbs_value`: given an aligned byte pointer's +// little-endian 16-bit limbs, returns `(carry, [cell_lo, cell_hi])`. The caller is responsible for +// range-checking `cell_hi` to `cell_ptr_hi_bits(...)`. +__device__ __forceinline__ CellPtr byte_ptr_limbs_to_cell_ptr_limbs_value( + uint32_t byte_lo, + uint32_t byte_hi +) { + uint32_t carry = byte_hi & 1u; + uint32_t cell_lo = (byte_lo + (carry << openvm::U16_BITS)) >> 1; + uint32_t cell_hi = byte_hi >> 1; + return CellPtr{carry, {cell_lo, cell_hi}}; +} + +// Value form of `add_const_u16_limbs_value`: adds a small `constant` (`< 2^16`) to a pointer given +// as little-endian 16-bit limbs `[lo, hi]`, carrying into the high limb. Returns `(carry, [new_lo, +// new_hi])`. The caller is responsible for range-checking `new_lo` to `U16_BITS`. +__device__ __forceinline__ CellPtr add_const_u16_limbs_value( + uint32_t lo, + uint32_t hi, + uint32_t constant +) { + uint32_t sum_lo = lo + constant; + uint32_t carry = sum_lo >> openvm::U16_BITS; + return CellPtr{carry, {sum_lo & 0xffffu, hi + carry}}; +} + +__device__ __forceinline__ void compute_block_add_carries( + VariableRangeChecker &range_checker, + uint32_t base_cell_lo, + uint32_t num_blocks, + uint32_t cell_stride, + uint32_t *add_carry_out +) { + for (uint32_t i = 0; i < num_blocks; i++) { + uint32_t sum_lo = base_cell_lo + i * cell_stride; + range_checker.add_count(sum_lo & 0xffffu, openvm::U16_BITS); + add_carry_out[i] = sum_lo >> openvm::U16_BITS; + } +} + +// Returns the conversion carry; writes one add-carry per block into add_carry_out. +__device__ __forceinline__ uint32_t compute_pointer_carries( + VariableRangeChecker &range_checker, + uint32_t byte_ptr, + size_t byte_ptr_max_bits, + uint32_t num_blocks, + uint32_t cell_stride, + uint32_t *add_carry_out +) { + CellPtr conv = byte_ptr_limbs_to_cell_ptr_limbs_value( + byte_ptr & 0xffffu, byte_ptr >> openvm::U16_BITS + ); + range_checker.add_count(conv.limbs[1], cell_ptr_hi_bits(byte_ptr_max_bits)); + compute_block_add_carries( + range_checker, conv.limbs[0], num_blocks, cell_stride, add_carry_out + ); + return conv.carry; +} diff --git a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap.cuh b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap.cuh index 7de45fc750..16d4f66f9e 100644 --- a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap.cuh +++ b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap.cuh @@ -3,6 +3,7 @@ #include "primitives/execution.h" #include "primitives/trace_access.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -22,6 +23,14 @@ struct Rv64VecHeapAdapterCols { T rs_val[NUM_READS][RV64_PTR_U16_LIMBS]; T rd_val[RV64_PTR_U16_LIMBS]; + // Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + T rs_cell_carry[NUM_READS]; + T rd_cell_carry; + // Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to each + // base cell pointer (block `j`'s carry into the high cell limb). + T reads_add_carry[NUM_READS][BLOCKS_PER_READ]; + T writes_add_carry[BLOCKS_PER_WRITE]; + MemoryReadAuxCols rs_read_aux[NUM_READS]; MemoryReadAuxCols rd_read_aux; @@ -83,17 +92,43 @@ struct Rv64VecHeapAdapter { ) { static_assert(NUM_READS == 1 || NUM_READS == 2); + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors the host filler in vec_heap.rs. + const uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + #pragma unroll for (size_t i = 0; i < NUM_READS; i++) { - range_checker.add_count( - ptr_bound_from_high_u16(uint16_t(record.rs_vals[i] >> U16_BITS), pointer_max_bits), - U16_BITS + uint32_t add_carries[BLOCKS_PER_READ]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rs_vals[i], + pointer_max_bits, + BLOCKS_PER_READ, + cell_stride, + add_carries ); + COL_WRITE_VALUE(row, Cols, rs_cell_carry[i], conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_READ; j++) { + COL_WRITE_VALUE(row, Cols, reads_add_carry[i][j], add_carries[j]); + } + } + { + uint32_t add_carries[BLOCKS_PER_WRITE]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rd_val, + pointer_max_bits, + BLOCKS_PER_WRITE, + cell_stride, + add_carries + ); + COL_WRITE_VALUE(row, Cols, rd_cell_carry, conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_WRITE; j++) { + COL_WRITE_VALUE(row, Cols, writes_add_carry[j], add_carries[j]); + } } - range_checker.add_count( - ptr_bound_from_high_u16(uint16_t(record.rd_val >> U16_BITS), pointer_max_bits), - U16_BITS - ); uint32_t timestamp = record.from_timestamp + NUM_READS + 1 + NUM_READS * BLOCKS_PER_READ + BLOCKS_PER_WRITE; diff --git a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch.cuh b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch.cuh index ddc0163cdf..81c319125f 100644 --- a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch.cuh +++ b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch.cuh @@ -3,6 +3,7 @@ #include "primitives/execution.h" #include "primitives/trace_access.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -14,9 +15,16 @@ struct Rv64VecHeapBranchAdapterCols { T rs_ptr[NUM_READS]; T rs_val[NUM_READS][RV64_PTR_U16_LIMBS]; + + // Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + T rs_cell_carry[NUM_READS]; + // Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to each + // base cell pointer (block `j`'s carry into the high cell limb). + T reads_add_carry[NUM_READS][BLOCKS_PER_READ]; + MemoryReadAuxCols rs_read_aux[NUM_READS]; - MemoryReadAuxCols heap_read_aux[NUM_READS][BLOCKS_PER_READ]; + MemoryReadAuxCols reads_aux[NUM_READS][BLOCKS_PER_READ]; }; template @@ -28,7 +36,7 @@ struct Rv64VecHeapBranchAdapterRecord { uint32_t rs_vals[NUM_READS]; MemoryReadAuxRecord rs_read_aux[NUM_READS]; - MemoryReadAuxRecord heap_read_aux[NUM_READS][BLOCKS_PER_READ]; + MemoryReadAuxRecord reads_aux[NUM_READS][BLOCKS_PER_READ]; }; template @@ -54,12 +62,26 @@ struct Rv64VecHeapBranchAdapter { ) { static_assert(NUM_READS == 1 || NUM_READS == 2); + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors the host filler in vec_heap_branch.rs. + const uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + #pragma unroll for (size_t i = 0; i < NUM_READS; i++) { - range_checker.add_count( - ptr_bound_from_high_u16(uint16_t(record.rs_vals[i] >> U16_BITS), pointer_max_bits), - U16_BITS + uint32_t add_carries[BLOCKS_PER_READ]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rs_vals[i], + pointer_max_bits, + BLOCKS_PER_READ, + cell_stride, + add_carries ); + COL_WRITE_VALUE(row, Cols, rs_cell_carry[i], conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_READ; j++) { + COL_WRITE_VALUE(row, Cols, reads_add_carry[i][j], add_carries[j]); + } } uint32_t timestamp = record.from_timestamp + NUM_READS + NUM_READS * BLOCKS_PER_READ; @@ -68,8 +90,8 @@ struct Rv64VecHeapBranchAdapter { for (int j = BLOCKS_PER_READ - 1; j >= 0; j--) { timestamp--; mem_helper.fill( - row.slice_from(COL_INDEX(Cols, heap_read_aux[i][j])), - record.heap_read_aux[i][j].prev_timestamp, + row.slice_from(COL_INDEX(Cols, reads_aux[i][j])), + record.reads_aux[i][j].prev_timestamp, timestamp ); } diff --git a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch_u16.cuh b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch_u16.cuh index 4be2fa2010..b5b2bc8577 100644 --- a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch_u16.cuh +++ b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_branch_u16.cuh @@ -4,6 +4,7 @@ #include "primitives/trace_access.h" #include "primitives/constants.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -16,6 +17,13 @@ struct Rv64VecHeapBranchU16AdapterCols { T rs_ptr[NUM_READS]; // Low 32 bits of each source pointer register as u16 limbs. T rs_val[NUM_READS][RV64_PTR_U16_LIMBS]; + + // Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + T rs_cell_carry[NUM_READS]; + // Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to each + // base cell pointer (block `j`'s carry into the high cell limb). + T reads_add_carry[NUM_READS][BLOCKS_PER_READ]; + MemoryReadAuxCols rs_read_aux[NUM_READS]; MemoryReadAuxCols reads_aux[NUM_READS][BLOCKS_PER_READ]; @@ -53,12 +61,26 @@ template struct Rv64VecHeapBranchU16A RowSlice row, Rv64VecHeapBranchU16AdapterRecord record ) { - // Bound each source pointer to pointer_max_bits by narrowing the high u16 limb. - const size_t limb_shift_bits = RV64_PTR_BITS - pointer_max_bits; + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors the host filler in vec_heap_branch_u16.rs. + const uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + +#pragma unroll for (size_t i = 0; i < NUM_READS; i++) { - range_checker.add_count( - (record.rs_vals[i] >> U16_BITS) << limb_shift_bits, U16_BITS + uint32_t add_carries[BLOCKS_PER_READ]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rs_vals[i], + pointer_max_bits, + BLOCKS_PER_READ, + cell_stride, + add_carries ); + COL_WRITE_VALUE(row, Cols, rs_cell_carry[i], conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_READ; j++) { + COL_WRITE_VALUE(row, Cols, reads_add_carry[i][j], add_carries[j]); + } } uint32_t timestamp = record.from_timestamp + NUM_READS + NUM_READS * BLOCKS_PER_READ; diff --git a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_u16.cuh b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_u16.cuh index b3028a8c35..39a89fe790 100644 --- a/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_u16.cuh +++ b/extensions/riscv-adapters/cuda/include/riscv-adapters/vec_heap_u16.cuh @@ -5,6 +5,7 @@ #include "primitives/trace_access.h" #include "primitives/constants.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" @@ -26,6 +27,14 @@ struct Rv64VecHeapU16AdapterCols { // Low 32 bits of rd register as u16 limbs. T rd_val[RV64_PTR_U16_LIMBS]; + // Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + T rs_cell_carry[NUM_READS]; + T rd_cell_carry; + // Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to each + // base cell pointer (block `j`'s carry into the high cell limb). + T reads_add_carry[NUM_READS][BLOCKS_PER_READ]; + T writes_add_carry[BLOCKS_PER_WRITE]; + MemoryReadAuxCols rs_read_aux[NUM_READS]; MemoryReadAuxCols rd_read_aux; @@ -85,16 +94,43 @@ struct Rv64VecHeapU16Adapter { BLOCKS_PER_READ, BLOCKS_PER_WRITE> record ) { - // Bound each register pointer to pointer_max_bits by narrowing the high u16 limb. - const size_t limb_shift_bits = RV64_PTR_BITS - pointer_max_bits; + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors the host filler in vec_heap_u16.rs. + const uint32_t cell_stride = MEMORY_BLOCK_BYTES / U16_CELL_SIZE; + +#pragma unroll for (size_t i = 0; i < NUM_READS; i++) { - range_checker.add_count( - (record.rs_vals[i] >> U16_BITS) << limb_shift_bits, U16_BITS + uint32_t add_carries[BLOCKS_PER_READ]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rs_vals[i], + pointer_max_bits, + BLOCKS_PER_READ, + cell_stride, + add_carries ); + COL_WRITE_VALUE(row, Cols, rs_cell_carry[i], conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_READ; j++) { + COL_WRITE_VALUE(row, Cols, reads_add_carry[i][j], add_carries[j]); + } + } + { + uint32_t add_carries[BLOCKS_PER_WRITE]; + uint32_t conv_carry = compute_pointer_carries( + range_checker, + record.rd_val, + pointer_max_bits, + BLOCKS_PER_WRITE, + cell_stride, + add_carries + ); + COL_WRITE_VALUE(row, Cols, rd_cell_carry, conv_carry); +#pragma unroll + for (size_t j = 0; j < BLOCKS_PER_WRITE; j++) { + COL_WRITE_VALUE(row, Cols, writes_add_carry[j], add_carries[j]); + } } - range_checker.add_count( - (record.rd_val >> U16_BITS) << limb_shift_bits, U16_BITS - ); uint32_t timestamp = record.from_timestamp + NUM_READS + 1 + NUM_READS * BLOCKS_PER_READ + BLOCKS_PER_WRITE; diff --git a/extensions/riscv-adapters/src/eq_mod.rs b/extensions/riscv-adapters/src/eq_mod.rs index 2713ea3001..f1d2041e09 100644 --- a/extensions/riscv-adapters/src/eq_mod.rs +++ b/extensions/riscv-adapters/src/eq_mod.rs @@ -8,7 +8,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, BasicAdapterInterface, ExecutionBridge, ExecutionState, MinimalInstruction, VmAdapterAir, - BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, + BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{ @@ -30,9 +30,10 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, - ptr_to_field_u16_limbs, tracing_read, tracing_read_reg_ptr, tracing_write, u16_limbs_to_ptr, - RV64_PTR_BITS, RV64_PTR_U16_LIMBS, RV64_REGISTER_NUM_LIMBS, U16_BITS, + compute_pointer_carries, eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, + expand_to_rv64_block, ptr_to_field_u16_limbs, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_read_reg_ptr, tracing_write, RV64_PTR_BITS, RV64_PTR_U16_LIMBS, + RV64_REGISTER_NUM_LIMBS, U16_BITS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -57,6 +58,12 @@ pub struct Rv64IsEqualModAdapterCols; NUM_READS], pub heap_read_aux: [[MemoryReadAuxCols; BLOCKS_PER_READ]; NUM_READS], + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub rd_ptr: T, pub writes_aux: MemoryWriteAuxCols, } @@ -127,11 +134,11 @@ impl< let d = AB::F::from_u32(RV64_REGISTER_AS); let e = AB::F::from_u32(RV64_MEMORY_AS); - // Read register values for rs. + // Read register values for rs (register pointers are small). for (ptr, val, aux) in izip!(cols.rs_ptr, cols.rs_val, &cols.rs_read_aux) { self.memory_bridge .read( - MemoryAddress::new(d, byte_ptr_to_u16_ptr::(ptr)), + MemoryAddress::new(d, reg_byte_ptr_to_cell_ptr_limbs::(ptr)), expand_to_rv64_block(&val), timestamp_pp(), aux, @@ -139,34 +146,47 @@ impl< .eval(builder, ctx.instruction.is_valid.clone()); } - // Compose the two u16 cells into low 32-bit heap pointers. - let rs_val_f: [AB::Expr; NUM_READS] = cols.rs_val.map(|limbs| u16_limbs_to_ptr(&limbs)); - - // Each materialized pointer is stored as two u16 cells. Bound the high - // cell against the guest byte-pointer limit. - for val in cols.rs_val.iter() { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr(val[1], self.pointer_max_bits), - U16_BITS, - ) - .eval(builder, ctx.instruction.is_valid.clone()); - } + let byte_ptr_max_bits = self.pointer_max_bits; + // Cell offset (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert each base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let rs_base_cell: [[AB::Expr; 2]; NUM_READS] = from_fn(|i| { + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); - // Reads from heap + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. let read_block_data: [[[_; MEMORY_BLOCK_BYTES]; BLOCKS_PER_READ]; NUM_READS] = ctx.reads.map(|r: [AB::Expr; TOTAL_READ_SIZE]| { let mut r_it = r.into_iter(); from_fn(|_| from_fn(|_| r_it.next().unwrap())) }); - let block_ptr_offset: [_; BLOCKS_PER_READ] = - from_fn(|i| AB::F::from_usize(i * MEMORY_BLOCK_BYTES)); - for (ptr, block_data, block_aux) in izip!(rs_val_f, read_block_data, &cols.heap_read_aux) { - for (offset, data, aux) in izip!(block_ptr_offset, block_data, block_aux) { + for (base_cell, block_data, block_aux, add_carry) in izip!( + rs_base_cell, + read_block_data, + &cols.heap_read_aux, + &cols.reads_add_carry + ) { + for (j, (data, aux, carry)) in izip!(block_data, block_aux, add_carry).enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new(e, byte_ptr_to_u16_ptr::(ptr.clone() + offset)), + MemoryAddress::new(e, block_cell_ptr), pack_u8_block::(&data), timestamp_pp(), aux, @@ -175,10 +195,10 @@ impl< } } - // Write to rd register + // Write to rd register (register pointer is small). self.memory_bridge .write( - MemoryAddress::new(d, byte_ptr_to_u16_ptr::(cols.rd_ptr)), + MemoryAddress::new(d, reg_byte_ptr_to_cell_ptr_limbs::(cols.rd_ptr)), pack_u8_block::(&ctx.writes[0].clone()), timestamp_pp(), &cols.writes_aux, @@ -358,10 +378,19 @@ impl Adap let cols: &mut Rv64IsEqualModAdapterCols = adapter_row.borrow_mut(); - for &ptr in record.rs_val.iter() { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_val[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); let mut timestamp = record.timestamp + (NUM_READS + NUM_READS * BLOCKS_PER_READ) as u32 + 1; let mut timestamp_mm = || { @@ -405,6 +434,14 @@ impl Adap cols.rs_val = record.rs_val.map(ptr_to_field_u16_limbs); cols.rs_ptr = record.rs_ptr.map(|ptr| F::from_u32(ptr)); + // Pointer-conversion / block-offset carry columns (computed above). + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv-adapters/src/eq_mod_u16.rs b/extensions/riscv-adapters/src/eq_mod_u16.rs index e8ac41a27d..fff0d88c81 100644 --- a/extensions/riscv-adapters/src/eq_mod_u16.rs +++ b/extensions/riscv-adapters/src/eq_mod_u16.rs @@ -8,7 +8,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, BasicAdapterInterface, ExecutionBridge, ExecutionState, MinimalInstruction, VmAdapterAir, - BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, + BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{ @@ -30,10 +30,10 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, expand_to_rv64_block, - ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, ptr_to_field_u16_limbs, tracing_read_reg_ptr, - tracing_read_u16, tracing_write_u16, u16_limbs_to_ptr, RV64_PTR_BITS, RV64_PTR_U16_LIMBS, - U16_BITS, + byte_ptr_to_u16_ptr_value, compute_pointer_carries, eval_add_const_u16_limbs, + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, ptr_to_field_u16_limbs, + reg_byte_ptr_to_cell_ptr_limbs, tracing_read_reg_ptr, tracing_read_u16, tracing_write_u16, + RV64_PTR_BITS, RV64_PTR_U16_LIMBS, U16_BITS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -53,6 +53,12 @@ pub struct Rv64IsEqualModU16AdapterCols; NUM_READS], pub heap_read_aux: [[MemoryReadAuxCols; BLOCKS_PER_READ]; NUM_READS], + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub rd_ptr: T, pub writes_aux: MemoryWriteAuxCols, } @@ -124,11 +130,11 @@ impl< let d = AB::F::from_u32(RV64_REGISTER_AS); let e = AB::F::from_u32(RV64_MEMORY_AS); - // Read register values for rs + // Read register values for rs (register pointers are small). for (ptr, val, aux) in izip!(cols.rs_ptr, cols.rs_val, &cols.rs_read_aux) { self.memory_bridge .read( - MemoryAddress::new(d, byte_ptr_to_u16_ptr::(ptr)), + MemoryAddress::new(d, reg_byte_ptr_to_cell_ptr_limbs::(ptr)), expand_to_rv64_block(&val), timestamp_pp(), aux, @@ -136,34 +142,47 @@ impl< .eval(builder, ctx.instruction.is_valid.clone()); } - // Compose the two u16 cells into low 32-bit heap pointers. - let rs_val_f: [AB::Expr; NUM_READS] = cols.rs_val.map(|limbs| u16_limbs_to_ptr(&limbs)); - - // Each materialized pointer is stored as two u16 cells. Bound the high - // cell against the guest byte-pointer limit. - for val in cols.rs_val.iter() { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr(val[1], self.pointer_max_bits), - U16_BITS, - ) - .eval(builder, ctx.instruction.is_valid.clone()); - } + let byte_ptr_max_bits = self.pointer_max_bits; + // Cell offset (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert each base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let rs_base_cell: [[AB::Expr; 2]; NUM_READS] = from_fn(|i| { + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); - // Reads from heap + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. let read_block_data: [[[_; BLOCK_FE_WIDTH]; BLOCKS_PER_READ]; NUM_READS] = ctx.reads.map(|r: [AB::Expr; TOTAL_READ_SIZE]| { let mut r_it = r.into_iter(); from_fn(|_| from_fn(|_| r_it.next().unwrap())) }); - let block_ptr_offset: [_; BLOCKS_PER_READ] = - from_fn(|i| AB::F::from_usize(i * MEMORY_BLOCK_BYTES)); - for (ptr, block_data, block_aux) in izip!(rs_val_f, read_block_data, &cols.heap_read_aux) { - for (offset, data, aux) in izip!(block_ptr_offset, block_data, block_aux) { + for (base_cell, block_data, block_aux, add_carry) in izip!( + rs_base_cell, + read_block_data, + &cols.heap_read_aux, + &cols.reads_add_carry + ) { + for (j, (data, aux, carry)) in izip!(block_data, block_aux, add_carry).enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new(e, byte_ptr_to_u16_ptr::(ptr.clone() + offset)), + MemoryAddress::new(e, block_cell_ptr), data, timestamp_pp(), aux, @@ -172,10 +191,10 @@ impl< } } - // Write to rd register + // Write to rd register (register pointer is small). self.memory_bridge .write( - MemoryAddress::new(d, byte_ptr_to_u16_ptr::(cols.rd_ptr)), + MemoryAddress::new(d, reg_byte_ptr_to_cell_ptr_limbs::(cols.rd_ptr)), ctx.writes[0].clone(), timestamp_pp(), &cols.writes_aux, @@ -355,10 +374,19 @@ impl Adap let cols: &mut Rv64IsEqualModU16AdapterCols = adapter_row.borrow_mut(); - for &ptr in record.rs_val.iter() { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_val[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); let mut timestamp = record.timestamp + (NUM_READS + NUM_READS * BLOCKS_PER_READ) as u32 + 1; let mut timestamp_mm = || { @@ -402,6 +430,14 @@ impl Adap cols.rs_val = record.rs_val.map(ptr_to_field_u16_limbs); cols.rs_ptr = record.rs_ptr.map(F::from_u32); + // Pointer-conversion / block-offset carry columns (computed above). + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv-adapters/src/test_utils.rs b/extensions/riscv-adapters/src/test_utils.rs index 458dcb5ca6..8e709d46ae 100644 --- a/extensions/riscv-adapters/src/test_utils.rs +++ b/extensions/riscv-adapters/src/test_utils.rs @@ -1,5 +1,5 @@ use openvm_circuit::arch::{ - testing::{memory::gen_pointer, TestBuilder}, + testing::{memory::gen_distinct_register_pointers, TestBuilder}, BLOCK_FE_WIDTH, U16_CELL_SIZE, }; use openvm_instructions::{instruction::Instruction, VmOpcode}; @@ -168,9 +168,8 @@ pub fn rv64_rand_write_register_or_imm( ) -> (Instruction, usize) { let rs2_is_imm = imm.is_some(); - let rs1 = gen_pointer(rng, NUM_LIMBS); - let rs2 = imm.unwrap_or_else(|| gen_pointer(rng, NUM_LIMBS)); - let rd = gen_pointer(rng, NUM_LIMBS); + let [rs1, rs2_reg, rd] = gen_distinct_register_pointers(rng, NUM_LIMBS); + let rs2 = imm.unwrap_or(rs2_reg); tester.write::(1, rs1, rs1_writes.map(BabyBear::from_u32)); if !rs2_is_imm { diff --git a/extensions/riscv-adapters/src/vec_heap.rs b/extensions/riscv-adapters/src/vec_heap.rs index d701778454..c5944afc5f 100644 --- a/extensions/riscv-adapters/src/vec_heap.rs +++ b/extensions/riscv-adapters/src/vec_heap.rs @@ -1,7 +1,7 @@ use std::{ array::from_fn, borrow::{Borrow, BorrowMut}, - iter::{once, zip}, + iter::once, }; use itertools::izip; @@ -9,7 +9,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, ExecutionBridge, ExecutionState, VecHeapAdapterInterface, VmAdapterAir, BLOCK_FE_WIDTH, - MEMORY_BLOCK_BYTES, + MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{ @@ -31,9 +31,9 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, - ptr_to_field_u16_limbs, tracing_read, tracing_read_reg_ptr, tracing_write, u16_limbs_to_ptr, - RV64_PTR_U16_LIMBS, U16_BITS, + compute_pointer_carries, eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, + expand_to_rv64_block, ptr_to_field_u16_limbs, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_read_reg_ptr, tracing_write, RV64_PTR_U16_LIMBS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -61,9 +61,19 @@ pub struct Rv64VecHeapAdapterCols< pub rs_ptr: [T; NUM_READS], pub rd_ptr: T, + /// Low 32 bits of rs registers as little-endian 16-bit *byte*-pointer limbs. pub rs_val: [[T; RV64_PTR_U16_LIMBS]; NUM_READS], + /// Low 32 bits of rd register as little-endian 16-bit *byte*-pointer limbs. pub rd_val: [T; RV64_PTR_U16_LIMBS], + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + pub rd_cell_carry: T, + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub writes_add_carry: [T; BLOCKS_PER_WRITE], + pub rs_read_aux: [MemoryReadAuxCols; NUM_READS], pub rd_read_aux: MemoryReadAuxCols, @@ -129,7 +139,7 @@ impl< timestamp + AB::F::from_usize(timestamp_delta - 1) }; - // Read register values for rs, rd + // Read register values for rs, rd (register pointers are small). for (ptr, val, aux) in izip!(cols.rs_ptr, cols.rs_val, &cols.rs_read_aux).chain(once(( cols.rd_ptr, cols.rd_val, @@ -139,7 +149,7 @@ impl< .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(ptr), + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), expand_to_rv64_block(&val), timestamp_pp(), @@ -148,33 +158,50 @@ impl< .eval(builder, ctx.instruction.is_valid.clone()); } - // Each materialized pointer is stored as two u16 cells. Bound the high - // cell against the guest byte-pointer limit. - for val in cols.rs_val.iter().chain(once(&cols.rd_val)) { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr(val[1], self.pointer_max_bits), - U16_BITS, - ) - .eval(builder, ctx.instruction.is_valid.clone()); - } - - // Compose the two u16 cells into low 32-bit heap/register pointers. - let rd_val_f: AB::Expr = u16_limbs_to_ptr(&cols.rd_val); - let rs_val_f: [AB::Expr; NUM_READS] = cols.rs_val.map(|limbs| u16_limbs_to_ptr(&limbs)); - + let byte_ptr_max_bits = self.pointer_max_bits; let e = AB::F::from_u32(RV64_MEMORY_AS); - // Reads from heap - for (address, reads, reads_aux) in izip!(rs_val_f, ctx.reads, &cols.reads_aux,) { - for (i, (read, aux)) in zip(reads, reads_aux).enumerate() { + // Cell offset (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert each base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let rs_base_cell: [[AB::Expr; 2]; NUM_READS] = from_fn(|i| { + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); + let rd_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rd_val.map(Into::into), + cols.rd_cell_carry, + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ); + + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. + for (base_cell, reads, reads_aux, add_carry) in izip!( + rs_base_cell, + ctx.reads, + &cols.reads_aux, + &cols.reads_add_carry + ) { + for (j, (read, aux, carry)) in izip!(reads, reads_aux, add_carry).enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - address.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), pack_u8_block::(&read), timestamp_pp(), aux, @@ -184,15 +211,20 @@ impl< } // Writes to heap - for (i, (write, aux)) in zip(ctx.writes, &cols.writes_aux).enumerate() { + for (j, (write, aux, carry)) in + izip!(ctx.writes, &cols.writes_aux, &cols.writes_add_carry).enumerate() + { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + rd_base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .write( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - rd_val_f.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), pack_u8_block::(&write), timestamp_pp(), aux, @@ -390,10 +422,26 @@ impl< let cols: &mut Rv64VecHeapAdapterCols = adapter_row.borrow_mut(); - for &ptr in record.rs_vals.iter().chain(once(&record.rd_val)) { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_vals[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); + let rd_carries = compute_pointer_carries( + &self.range_checker_chip, + record.rd_val, + BLOCKS_PER_WRITE, + cell_stride, + self.pointer_max_bits, + ); let timestamp_delta = NUM_READS + 1 + NUM_READS * BLOCKS_PER_READ + BLOCKS_PER_WRITE; let mut timestamp = record.from_timestamp + timestamp_delta as u32; @@ -459,6 +507,20 @@ impl< .for_each(|(cols_ptr, ptr)| { *cols_ptr = F::from_u32(*ptr); }); + + // Pointer-conversion / block-offset carry columns (computed above). + let (rd_conv, rd_add) = &rd_carries; + cols.rd_cell_carry = F::from_u32(*rd_conv); + for (col, &c) in cols.writes_add_carry.iter_mut().zip(rd_add.iter()) { + *col = F::from_u32(c); + } + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.from_timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv-adapters/src/vec_heap_branch.rs b/extensions/riscv-adapters/src/vec_heap_branch.rs index ef7717c78a..f6b570030f 100644 --- a/extensions/riscv-adapters/src/vec_heap_branch.rs +++ b/extensions/riscv-adapters/src/vec_heap_branch.rs @@ -1,7 +1,6 @@ use std::{ array::from_fn, borrow::{Borrow, BorrowMut}, - iter::zip, }; use itertools::izip; @@ -9,7 +8,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, ExecutionBridge, ExecutionState, VecHeapBranchAdapterInterface, VmAdapterAir, - MEMORY_BLOCK_BYTES, + MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{pack_u8_block, MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord}, @@ -28,9 +27,9 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, - ptr_to_field_u16_limbs, tracing_read, tracing_read_reg_ptr, u16_limbs_to_ptr, - RV64_PTR_U16_LIMBS, U16_BITS, + compute_pointer_carries, eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, + expand_to_rv64_block, ptr_to_field_u16_limbs, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_read_reg_ptr, RV64_PTR_U16_LIMBS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -50,7 +49,15 @@ pub struct Rv64VecHeapBranchAdapterCols, pub rs_ptr: [T; NUM_READS], + /// Low 32 bits of each source pointer register as little-endian 16-bit *byte*-pointer limbs. pub rs_val: [[T; RV64_PTR_U16_LIMBS]; NUM_READS], + + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub rs_read_aux: [MemoryReadAuxCols; NUM_READS], pub reads_aux: [[MemoryReadAuxCols; BLOCKS_PER_READ]; NUM_READS], @@ -95,13 +102,13 @@ impl(ptr), + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), expand_to_rv64_block(&val), timestamp_pp(), @@ -110,32 +117,42 @@ impl( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); + + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. + for (base_cell, reads, reads_aux, add_carry) in izip!( + rs_base_cell, + ctx.reads, + &cols.reads_aux, + &cols.reads_add_carry + ) { + for (j, (read, aux, carry)) in izip!(reads, reads_aux, add_carry).enumerate() { + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - address.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), pack_u8_block::(&read), timestamp_pp(), aux, @@ -277,10 +294,19 @@ impl Adap let cols: &mut Rv64VecHeapBranchAdapterCols = adapter_row.borrow_mut(); - for &ptr in record.rs_vals.iter() { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_vals[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); let timestamp_delta = NUM_READS + NUM_READS * BLOCKS_PER_READ; let mut timestamp = record.from_timestamp + timestamp_delta as u32; @@ -328,6 +354,15 @@ impl Adap .for_each(|(cols_ptr, ptr)| { *cols_ptr = F::from_u32(*ptr); }); + + // Pointer-conversion / block-offset carry columns (computed above). + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.from_timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv-adapters/src/vec_heap_branch_u16.rs b/extensions/riscv-adapters/src/vec_heap_branch_u16.rs index 3d4d2837e1..8a8e9a2206 100644 --- a/extensions/riscv-adapters/src/vec_heap_branch_u16.rs +++ b/extensions/riscv-adapters/src/vec_heap_branch_u16.rs @@ -1,7 +1,6 @@ use std::{ array::from_fn, borrow::{Borrow, BorrowMut}, - iter::zip, }; use itertools::izip; @@ -9,7 +8,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, ExecutionBridge, ExecutionState, VecHeapBranchAdapterInterface, VmAdapterAir, - BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, + BLOCK_FE_WIDTH, MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord}, @@ -28,9 +27,9 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, expand_to_rv64_block, - ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, ptr_to_u16_limbs, tracing_read_reg_ptr, - tracing_read_u16, u16_limbs_to_ptr, RV64_PTR_U16_LIMBS, U16_BITS, + byte_ptr_to_u16_ptr_value, compute_pointer_carries, eval_add_const_u16_limbs, + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, ptr_to_u16_limbs, + reg_byte_ptr_to_cell_ptr_limbs, tracing_read_reg_ptr, tracing_read_u16, RV64_PTR_U16_LIMBS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -51,8 +50,15 @@ pub struct Rv64VecHeapBranchU16AdapterCols, pub rs_ptr: [T; NUM_READS], - /// Low 32 bits of each source pointer register as u16 limbs. + /// Low 32 bits of each source pointer register as little-endian 16-bit *byte*-pointer limbs. pub rs_val: [[T; RV64_PTR_U16_LIMBS]; NUM_READS], + + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub rs_read_aux: [MemoryReadAuxCols; NUM_READS], pub reads_aux: [[MemoryReadAuxCols; BLOCKS_PER_READ]; NUM_READS], @@ -97,14 +103,14 @@ impl(ptr), + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), bus_payload, timestamp_pp(), @@ -113,33 +119,43 @@ impl( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); + + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. + for (base_cell, reads, reads_aux, add_carry) in izip!( + rs_base_cell, + ctx.reads, + &cols.reads_aux, + &cols.reads_add_carry + ) { + for (j, (read, aux, carry)) in izip!(reads, reads_aux, add_carry).enumerate() { + let read_array: [AB::Expr; BLOCK_FE_WIDTH] = from_fn(|k| read[k].clone()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - address.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), read_array, timestamp_pp(), aux, @@ -282,12 +298,19 @@ impl Adap let cols: &mut Rv64VecHeapBranchU16AdapterCols = adapter_row.borrow_mut(); - // Range checks: - // **NOTE**: Must do the range checks before overwriting the records - for &v in record.rs_vals.iter() { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(v, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_vals[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); let timestamp_delta = NUM_READS + NUM_READS * BLOCKS_PER_READ; let mut timestamp = record.from_timestamp + timestamp_delta as u32; @@ -335,6 +358,15 @@ impl Adap .for_each(|(cols_ptr, ptr)| { *cols_ptr = F::from_u32(*ptr); }); + + // Pointer-conversion / block-offset carry columns (computed above). + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.from_timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv-adapters/src/vec_heap_u16.rs b/extensions/riscv-adapters/src/vec_heap_u16.rs index 37fcfead34..e9e973e759 100644 --- a/extensions/riscv-adapters/src/vec_heap_u16.rs +++ b/extensions/riscv-adapters/src/vec_heap_u16.rs @@ -1,7 +1,7 @@ use std::{ array::from_fn, borrow::{Borrow, BorrowMut}, - iter::{once, zip}, + iter::once, }; use itertools::izip; @@ -9,7 +9,7 @@ use openvm_circuit::{ arch::{ get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller, ExecutionBridge, ExecutionState, VecHeapAdapterInterface, VmAdapterAir, BLOCK_FE_WIDTH, - MEMORY_BLOCK_BYTES, + MEMORY_BLOCK_BYTES, U16_CELL_SIZE, }, system::memory::{ offline_checker::{ @@ -31,9 +31,10 @@ use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, }; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, expand_to_rv64_block, - ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, ptr_to_u16_limbs, tracing_read_reg_ptr, - tracing_read_u16, tracing_write_u16, u16_limbs_to_ptr, RV64_PTR_U16_LIMBS, U16_BITS, + byte_ptr_to_u16_ptr_value, compute_pointer_carries, eval_add_const_u16_limbs, + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, ptr_to_u16_limbs, + reg_byte_ptr_to_cell_ptr_limbs, tracing_read_reg_ptr, tracing_read_u16, tracing_write_u16, + RV64_PTR_U16_LIMBS, }; use openvm_stark_backend::{ interaction::InteractionBuilder, @@ -61,11 +62,19 @@ pub struct Rv64VecHeapU16AdapterCols< pub rs_ptr: [T; NUM_READS], pub rd_ptr: T, - /// Low 32 bits of rs registers as u16 limbs. + /// Low 32 bits of rs registers as little-endian 16-bit *byte*-pointer limbs. pub rs_val: [[T; RV64_PTR_U16_LIMBS]; NUM_READS], - /// Low 32 bits of rd register as u16 limbs. + /// Low 32 bits of rd register as little-endian 16-bit *byte*-pointer limbs. pub rd_val: [T; RV64_PTR_U16_LIMBS], + /// Carry for converting each base byte pointer to AS-native u16 *cell* pointer limbs. + pub rs_cell_carry: [T; NUM_READS], + pub rd_cell_carry: T, + /// Per-block carry for adding the cell offset `j * (MEMORY_BLOCK_BYTES / U16_CELL_SIZE)` to + /// each base cell pointer (block `j`'s carry into the high cell limb). + pub reads_add_carry: [[T; BLOCKS_PER_READ]; NUM_READS], + pub writes_add_carry: [T; BLOCKS_PER_WRITE], + pub rs_read_aux: [MemoryReadAuxCols; NUM_READS], pub rd_read_aux: MemoryReadAuxCols, @@ -131,7 +140,7 @@ impl< timestamp + AB::F::from_usize(timestamp_delta - 1) }; - // Read register values for rs, rd + // Read register values for rs, rd (register pointers are small). for (ptr, val, aux) in izip!(cols.rs_ptr, cols.rs_val, &cols.rs_read_aux).chain(once(( cols.rd_ptr, cols.rd_val, @@ -142,7 +151,7 @@ impl< .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(ptr), + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), bus_payload, timestamp_pp(), @@ -151,34 +160,51 @@ impl< .eval(builder, ctx.instruction.is_valid.clone()); } - // Each materialized pointer is stored as two u16 cells. Bound the high - // cell against the guest byte-pointer limit. - for val in cols.rs_val.iter().chain(once(&cols.rd_val)) { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr(val[1], self.pointer_max_bits), - U16_BITS, - ) - .eval(builder, ctx.instruction.is_valid.clone()); - } - - // Compose the two u16 cells into low 32-bit heap/register pointers. - let rd_val_f: AB::Expr = u16_limbs_to_ptr(&cols.rd_val); - let rs_val_f: [AB::Expr; NUM_READS] = cols.rs_val.map(|limbs| u16_limbs_to_ptr(&limbs)); - + let byte_ptr_max_bits = self.pointer_max_bits; let e = AB::F::from_u32(RV64_MEMORY_AS); - // Reads from heap - for (address, reads, reads_aux) in izip!(rs_val_f, ctx.reads, &cols.reads_aux,) { - for (i, (read, aux)) in zip(reads, reads_aux).enumerate() { - let read_array: [AB::Expr; BLOCK_FE_WIDTH] = from_fn(|j| read[j].clone()); + // Cell offset (in u16 cells) between consecutive heap blocks. + let cell_ptr_block_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + + // Convert each base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let rs_base_cell: [[AB::Expr; 2]; NUM_READS] = from_fn(|i| { + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rs_val[i].map(Into::into), + cols.rs_cell_carry[i], + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ) + }); + let rd_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + cols.rd_val.map(Into::into), + cols.rd_cell_carry, + byte_ptr_max_bits, + ctx.instruction.is_valid.clone(), + ); + + // Reads from heap: block `j` is at base cell pointer + `j * cell_ptr_block_stride`. + for (base_cell, reads, reads_aux, add_carry) in izip!( + rs_base_cell, + ctx.reads, + &cols.reads_aux, + &cols.reads_add_carry + ) { + for (j, (read, aux, carry)) in izip!(reads, reads_aux, add_carry).enumerate() { + let read_array: [AB::Expr; BLOCK_FE_WIDTH] = from_fn(|k| read[k].clone()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .read( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - address.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), read_array, timestamp_pp(), aux, @@ -188,16 +214,21 @@ impl< } // Writes to heap - for (i, (write, aux)) in zip(ctx.writes, &cols.writes_aux).enumerate() { - let write_array: [AB::Expr; BLOCK_FE_WIDTH] = from_fn(|j| write[j].clone()); + for (j, (write, aux, carry)) in + izip!(ctx.writes, &cols.writes_aux, &cols.writes_add_carry).enumerate() + { + let write_array: [AB::Expr; BLOCK_FE_WIDTH] = from_fn(|k| write[k].clone()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + rd_base_cell.clone(), + j as u32 * cell_ptr_block_stride, + *carry, + ctx.instruction.is_valid.clone(), + ); self.memory_bridge .write( - MemoryAddress::new( - e, - byte_ptr_to_u16_ptr::( - rd_val_f.clone() + AB::Expr::from_usize(i * MEMORY_BLOCK_BYTES), - ), - ), + MemoryAddress::new(e, block_cell_ptr), write_array, timestamp_pp(), aux, @@ -391,12 +422,26 @@ impl< let cols: &mut Rv64VecHeapU16AdapterCols = adapter_row.borrow_mut(); - // Range checks: - // **NOTE**: Must do the range checks before overwriting the records - for &v in record.rs_vals.iter().chain(once(&record.rd_val)) { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(v, self.pointer_max_bits), U16_BITS); - } + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. **NOTE**: Must read the record values before overwriting them below. + // `carry` columns are written near the end (after all record reads), so store them here. + let cell_stride = (MEMORY_BLOCK_BYTES / U16_CELL_SIZE) as u32; + let rs_carries: [(u32, Vec); NUM_READS] = from_fn(|i| { + compute_pointer_carries( + &self.range_checker_chip, + record.rs_vals[i], + BLOCKS_PER_READ, + cell_stride, + self.pointer_max_bits, + ) + }); + let rd_carries = compute_pointer_carries( + &self.range_checker_chip, + record.rd_val, + BLOCKS_PER_WRITE, + cell_stride, + self.pointer_max_bits, + ); let timestamp_delta = NUM_READS + 1 + NUM_READS * BLOCKS_PER_READ + BLOCKS_PER_WRITE; let mut timestamp = record.from_timestamp + timestamp_delta as u32; @@ -462,6 +507,20 @@ impl< .for_each(|(cols_ptr, ptr)| { *cols_ptr = F::from_u32(*ptr); }); + + // Pointer-conversion / block-offset carry columns (computed above). + let (rd_conv, rd_add) = &rd_carries; + cols.rd_cell_carry = F::from_u32(*rd_conv); + for (col, &c) in cols.writes_add_carry.iter_mut().zip(rd_add.iter()) { + *col = F::from_u32(c); + } + for (i, (conv, add)) in rs_carries.iter().enumerate() { + cols.rs_cell_carry[i] = F::from_u32(*conv); + for (col, &c) in cols.reads_add_carry[i].iter_mut().zip(add.iter()) { + *col = F::from_u32(c); + } + } + cols.from_state.timestamp = F::from_u32(record.from_timestamp); cols.from_state.pc = F::from_u32(record.from_pc); } diff --git a/extensions/riscv/circuit/cuda/include/riscv/adapters/loadstore.cuh b/extensions/riscv/circuit/cuda/include/riscv/adapters/loadstore.cuh index f69309d789..403fc685e2 100644 --- a/extensions/riscv/circuit/cuda/include/riscv/adapters/loadstore.cuh +++ b/extensions/riscv/circuit/cuda/include/riscv/adapters/loadstore.cuh @@ -19,8 +19,12 @@ template struct Rv64LoadStoreAdapterCols { MemoryReadAuxCols read_data_aux; T imm; T imm_sign; - /// mem_ptr is the intermediate memory pointer limbs, needed to check the correct addition + /// mem_ptr is the intermediate memory pointer limbs, needed to check the correct addition. + /// These are *byte*-pointer limbs of `rs1 + imm` (each < 2^16). T mem_ptr_limbs[2]; + /// Carry bit (`mem_ptr_limbs[1] & 1`) used to convert the aligned heap *byte* pointer into + /// AS-native u16 *cell* pointer limbs. + T mem_ptr_carry; T mem_as; /// Timestamp aux for the write; previous data is provided by the core chip. MemoryBaseAuxCols write_base_aux; @@ -47,6 +51,8 @@ struct Rv64LoadStoreAdapterRecord { bool imm_sign; uint8_t mem_as; + /// Whether this is a load (heap access is the read) vs a store (heap access is the write). + bool is_load; uint32_t write_prev_timestamp; }; @@ -101,10 +107,21 @@ struct Rv64LoadStoreAdapter { uint32_t ptr_limbs[RV64_PTR_U16_LIMBS]; ptr_to_u16_limbs(ptr_limbs, ptr); COL_WRITE_ARRAY(row, Rv64LoadStoreAdapterCols, mem_ptr_limbs, ptr_limbs); + + // Convert the aligned heap byte pointer to AS-native u16 cell pointer limbs. The carry + // depends only on the high byte limb's parity; only `cell_hi` is range-checked (hi_bits). + uint32_t shift = ptr & (uint32_t(RV64_REGISTER_NUM_LIMBS) - 1); + uint32_t aligned_byte_limbs[RV64_PTR_U16_LIMBS]; + ptr_to_u16_limbs(aligned_byte_limbs, ptr - shift); + uint32_t mem_ptr_carry = aligned_byte_limbs[1] & 1u; + uint32_t heap_cell_hi = aligned_byte_limbs[1] >> 1; + COL_WRITE_VALUE(row, Rv64LoadStoreAdapterCols, mem_ptr_carry, mem_ptr_carry); COL_WRITE_VALUE(row, Rv64LoadStoreAdapterCols, mem_as, record.mem_as); + uint32_t hi_bits = uint32_t(pointer_max_bits) - U16_CELL_SIZE_BITS - U16_BITS; + // Alignment check: `(mem_ptr_limbs[0] - shift) / 8 < 2^13`. range_checker.add_count(ptr_limbs[0] >> 3, U16_BITS - 3); - range_checker.add_count(ptr_limbs[1], pointer_max_bits - U16_BITS); + range_checker.add_count(heap_cell_hi, hi_bits); COL_WRITE_VALUE(row, Rv64LoadStoreAdapterCols, needs_write, needs_write); if (needs_write) { diff --git a/extensions/riscv/circuit/cuda/src/hintstore.cu b/extensions/riscv/circuit/cuda/src/hintstore.cu index 073189bbee..6dc505bac1 100644 --- a/extensions/riscv/circuit/cuda/src/hintstore.cu +++ b/extensions/riscv/circuit/cuda/src/hintstore.cu @@ -28,6 +28,11 @@ template struct Rv64HintStoreCols { // Low 32 bits of the 8-byte RV64 register that holds `mem_ptr`; the upper 4 bytes are // known to be zero and are hardcoded in the memory bus interaction. T mem_ptr_limbs[RV64_PTR_U16_LIMBS]; + // Carry (`mem_ptr_limbs[1] & 1`) for converting the byte pointer to AS-native u16 *cell* + // pointer limbs. + T mem_ptr_carry; + // Carry for the per-row `next.mem_ptr = mem_ptr + 8` byte increment. + T mem_ptr_inc_carry; MemoryReadAuxCols mem_ptr_aux_cols; MemoryWriteAuxCols write_aux; @@ -96,6 +101,19 @@ struct Rv64HintStore { COL_WRITE_VALUE(row, Rv64HintStoreCols, mem_ptr_ptr, record.mem_ptr_ptr); COL_WRITE_ARRAY(row, Rv64HintStoreCols, mem_ptr_limbs, mem_ptr_limbs); + // Byte -> cell pointer conversion (heap write) and the per-row range checks: + // cell_hi (hi_bits) and the low byte limb (16 bits, for the limb-wise `+8` increment). + // `mem_ptr_limbs` are the little-endian 16-bit *byte*-pointer limbs `[byte_lo, byte_hi]`. + uint32_t mem_carry = mem_ptr_limbs[1] & 1u; + uint32_t cell_hi = mem_ptr_limbs[1] >> 1; + uint32_t inc_carry = + (mem_ptr_limbs[0] + (uint32_t)RV64_REGISTER_NUM_LIMBS) >> U16_BITS; + COL_WRITE_VALUE(row, Rv64HintStoreCols, mem_ptr_carry, mem_carry); + COL_WRITE_VALUE(row, Rv64HintStoreCols, mem_ptr_inc_carry, inc_carry); + uint32_t hi_bits = (uint32_t)pointer_max_bits - U16_CELL_SIZE_BITS - U16_BITS; + range_checker.add_count(cell_hi, hi_bits); + range_checker.add_count(mem_ptr_limbs[0], U16_BITS); + if (local_idx == 0) { #ifdef CUDA_DEBUG // The overflow check for mem_ptr + num_words * 8 is not needed because @@ -104,11 +122,6 @@ struct Rv64HintStore { assert(record.num_words <= MAX_HINT_BUFFER_DWORDS); #endif - // Range check for mem_ptr (using pointer_max_bits). - uint32_t mem_ptr_shift = RV64_PTR_BITS - (uint32_t)pointer_max_bits; - uint32_t mem_ptr_high_u16 = record.mem_ptr >> U16_BITS; - range_checker.add_count(mem_ptr_high_u16 << mem_ptr_shift, U16_BITS); - // Range check for num_words (using MAX_HINT_BUFFER_DWORDS_BITS). range_checker.add_count(record.num_words << REM_WORDS_SHIFT, U16_BITS); diff --git a/extensions/riscv/circuit/src/adapters/alu.rs b/extensions/riscv/circuit/src/adapters/alu.rs index f5b9355c6d..f19460e375 100644 --- a/extensions/riscv/circuit/src/adapters/alu.rs +++ b/extensions/riscv/circuit/src/adapters/alu.rs @@ -32,7 +32,7 @@ use openvm_stark_backend::{ p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, }; -use super::{byte_ptr_to_u16_ptr, tracing_read, tracing_read_imm, tracing_write}; +use super::{reg_byte_ptr_to_cell_ptr_limbs, tracing_read, tracing_read_imm, tracing_write}; #[repr(C)] #[derive(AlignedBorrow, StructReflection)] @@ -116,7 +116,7 @@ impl VmAdapterAir for Rv64BaseAluAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), pack_u8_block::(&ctx.reads[0].clone()), timestamp_pp(), @@ -130,7 +130,10 @@ impl VmAdapterAir for Rv64BaseAluAdapterAir { .assert_one(ctx.instruction.is_valid.clone()); self.memory_bridge .read( - MemoryAddress::new(local.rs2_as, byte_ptr_to_u16_ptr::(local.rs2)), + MemoryAddress::new( + local.rs2_as, + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2), + ), pack_u8_block::(&ctx.reads[1].clone()), timestamp_pp(), &local.reads_aux[1], @@ -141,7 +144,7 @@ impl VmAdapterAir for Rv64BaseAluAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), pack_u8_block::(&ctx.writes[0].clone()), timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/alu_u16.rs b/extensions/riscv/circuit/src/adapters/alu_u16.rs index 78438505f9..3dff4032fe 100644 --- a/extensions/riscv/circuit/src/adapters/alu_u16.rs +++ b/extensions/riscv/circuit/src/adapters/alu_u16.rs @@ -32,7 +32,9 @@ use openvm_stark_backend::{ p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, }; -use super::{byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, tracing_read_u16, tracing_write_u16}; +use super::{ + byte_ptr_to_u16_ptr_value, reg_byte_ptr_to_cell_ptr_limbs, tracing_read_u16, tracing_write_u16, +}; use crate::adapters::{imm_to_rv64_u64, U16_BITS}; #[repr(C)] @@ -123,7 +125,7 @@ impl VmAdapterAir for Rv64BaseAluU16AdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), ctx.reads[0].clone(), timestamp_pp(), @@ -136,7 +138,10 @@ impl VmAdapterAir for Rv64BaseAluU16AdapterAir { .assert_one(ctx.instruction.is_valid.clone()); self.memory_bridge .read( - MemoryAddress::new(local.rs2_as, byte_ptr_to_u16_ptr::(local.rs2)), + MemoryAddress::new( + local.rs2_as, + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2), + ), ctx.reads[1].clone(), timestamp_pp(), &local.reads_aux[1], @@ -147,7 +152,7 @@ impl VmAdapterAir for Rv64BaseAluU16AdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), ctx.writes[0].clone(), timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/alu_w.rs b/extensions/riscv/circuit/src/adapters/alu_w.rs index d0b6ab9526..ffecdb79c4 100644 --- a/extensions/riscv/circuit/src/adapters/alu_w.rs +++ b/extensions/riscv/circuit/src/adapters/alu_w.rs @@ -38,8 +38,8 @@ use openvm_stark_backend::{ }; use super::{ - byte_ptr_to_u16_ptr, pack_high_u16, pack_rv64_u16_block, tracing_read, tracing_read_imm, - tracing_write, RV64_PTR_U16_LIMBS, + pack_high_u16, pack_rv64_u16_block, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_read_imm, tracing_write, RV64_PTR_U16_LIMBS, }; #[repr(C)] @@ -132,7 +132,7 @@ impl VmAdapterAir for Rv64BaseAluWAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), rs1_data, timestamp_pp(), @@ -148,7 +148,10 @@ impl VmAdapterAir for Rv64BaseAluWAdapterAir { pack_rv64_u16_block(&ctx.reads[1], &local.rs2_high); self.memory_bridge .read( - MemoryAddress::new(local.rs2_as, byte_ptr_to_u16_ptr::(local.rs2)), + MemoryAddress::new( + local.rs2_as, + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2), + ), rs2_data, timestamp_pp(), &local.reads_aux[1], @@ -177,7 +180,7 @@ impl VmAdapterAir for Rv64BaseAluWAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), write_data, timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/alu_w_u16.rs b/extensions/riscv/circuit/src/adapters/alu_w_u16.rs index f9fafbadb8..126e613f68 100644 --- a/extensions/riscv/circuit/src/adapters/alu_w_u16.rs +++ b/extensions/riscv/circuit/src/adapters/alu_w_u16.rs @@ -36,8 +36,9 @@ use openvm_stark_backend::{ }; use super::{ - byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, concat_rv64_u16_block, imm_to_rv64_u64, - tracing_read_u16, tracing_write_u16, RV64_WORD_U16_LIMBS, U16_BITS, + byte_ptr_to_u16_ptr_value, concat_rv64_u16_block, imm_to_rv64_u64, + reg_byte_ptr_to_cell_ptr_limbs, tracing_read_u16, tracing_write_u16, RV64_WORD_U16_LIMBS, + U16_BITS, }; #[repr(C)] @@ -138,7 +139,7 @@ impl VmAdapterAir for Rv64BaseAluWU16AdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), rs1_data, timestamp_pp(), @@ -154,7 +155,10 @@ impl VmAdapterAir for Rv64BaseAluWU16AdapterAir { concat_rv64_u16_block(&ctx.reads[1], &local.rs2_high); self.memory_bridge .read( - MemoryAddress::new(local.rs2_as, byte_ptr_to_u16_ptr::(local.rs2)), + MemoryAddress::new( + local.rs2_as, + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2), + ), rs2_data, timestamp_pp(), &local.reads_aux[1], @@ -180,7 +184,7 @@ impl VmAdapterAir for Rv64BaseAluWU16AdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), write_data, timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/branch.rs b/extensions/riscv/circuit/src/adapters/branch.rs index a10bb69ba3..d0c34e3e79 100644 --- a/extensions/riscv/circuit/src/adapters/branch.rs +++ b/extensions/riscv/circuit/src/adapters/branch.rs @@ -25,7 +25,9 @@ use openvm_stark_backend::{ p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, }; -use crate::adapters::{byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, tracing_read_u16}; +use crate::adapters::{ + byte_ptr_to_u16_ptr_value, reg_byte_ptr_to_cell_ptr_limbs, tracing_read_u16, +}; #[repr(C)] #[derive(AlignedBorrow, StructReflection)] @@ -71,7 +73,7 @@ impl VmAdapterAir for Rv64BranchAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), ctx.reads[0].clone(), timestamp_pp(), @@ -83,7 +85,7 @@ impl VmAdapterAir for Rv64BranchAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs2_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2_ptr), ), ctx.reads[1].clone(), timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/jalr.rs b/extensions/riscv/circuit/src/adapters/jalr.rs index 989b802dd4..ddcc6574ac 100644 --- a/extensions/riscv/circuit/src/adapters/jalr.rs +++ b/extensions/riscv/circuit/src/adapters/jalr.rs @@ -32,7 +32,7 @@ use openvm_stark_backend::{ }; use crate::adapters::{ - byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, tracing_read_u16, tracing_write_u16, + byte_ptr_to_u16_ptr_value, reg_byte_ptr_to_cell_ptr_limbs, tracing_read_u16, tracing_write_u16, }; #[repr(C)] @@ -97,7 +97,7 @@ impl VmAdapterAir for Rv64JalrAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.rs1_ptr), ), ctx.reads[0].clone(), timestamp_pp(), @@ -109,7 +109,7 @@ impl VmAdapterAir for Rv64JalrAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.rd_ptr), ), ctx.writes[0].clone(), timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/loadstore.rs b/extensions/riscv/circuit/src/adapters/loadstore.rs index 7deea67f61..ddc65b1afb 100644 --- a/extensions/riscv/circuit/src/adapters/loadstore.rs +++ b/extensions/riscv/circuit/src/adapters/loadstore.rs @@ -39,9 +39,10 @@ use openvm_stark_backend::{ }; use super::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_to_field_u16_limbs, ptr_to_u16_limbs, - rv64_address_add_imm, sign_extend_imm16, try_rv64_bytes_to_u32, RV64_PTR_BITS, - RV64_PTR_U16_LIMBS, RV64_REGISTER_NUM_LIMBS, U16_BITS, + byte_ptr_limbs_to_cell_ptr_limbs_value, cell_ptr_hi_bits, + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, ptr_to_field_u16_limbs, + ptr_to_u16_limbs, reg_byte_ptr_to_cell_ptr_limbs, rv64_address_add_imm, sign_extend_imm16, + try_rv64_bytes_to_u32, RV64_PTR_BITS, RV64_PTR_U16_LIMBS, RV64_REGISTER_NUM_LIMBS, U16_BITS, }; use crate::adapters::{memory_read, timed_write, tracing_read}; @@ -95,8 +96,12 @@ pub struct Rv64LoadStoreAdapterCols { pub read_data_aux: MemoryReadAuxCols, pub imm: T, pub imm_sign: T, - /// mem_ptr is the intermediate memory pointer limbs, needed to check the correct addition + /// mem_ptr is the intermediate memory pointer limbs, needed to check the correct addition. + /// These are *byte*-pointer limbs of `rs1 + imm` (each < 2^16). pub mem_ptr_limbs: [T; 2], + /// Carry bit (`mem_ptr_limbs[1] & 1`) used to convert the aligned heap *byte* pointer into + /// AS-native u16 *cell* pointer limbs. See `eval_byte_ptr_limbs_to_u16_cell_ptr_limbs`. + pub mem_ptr_carry: T, pub mem_as: T, /// Timestamp aux for the write; previous data is provided by the core chip. pub write_base_aux: MemoryBaseAuxCols, @@ -144,9 +149,9 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { let is_load = ctx.instruction.is_load; let is_valid = ctx.instruction.is_valid; - let load_shift_amount = ctx.instruction.load_shift_amount; - let store_shift_amount = ctx.instruction.store_shift_amount; - let shift_amount = load_shift_amount.clone() + store_shift_amount.clone(); + let load_shift_amount: AB::Expr = ctx.instruction.load_shift_amount; + let store_shift_amount: AB::Expr = ctx.instruction.store_shift_amount; + let shift_amount = load_shift_amount + store_shift_amount; let write_count = local_cols.needs_write; @@ -169,7 +174,7 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.rs1_ptr), ), rs1_data, timestamp_pp(), @@ -196,24 +201,18 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { .when(is_valid.clone()) .assert_eq(carry, local_cols.imm_sign); - // preventing mem_ptr overflow + // Alignment: the (shifted) heap byte pointer is 8-byte aligned, i.e. + // `(mem_ptr_limbs[0] - shift_amount) / 8 < 2^13`, which also implies + // `mem_ptr_limbs[0] - shift_amount < 2^16`. (The high byte limb `mem_ptr_limbs[1]` is + // bounded by the cell-pointer range checks below, via whichever of read/write is the heap + // access.) self.range_bus .range_check( - // (limb[0] - shift_amount) / 8 < 2^13 => limb[0] - shift_amount < 2^16 - (local_cols.mem_ptr_limbs[0] - shift_amount) + (local_cols.mem_ptr_limbs[0] - shift_amount.clone()) * AB::F::from_u32(RV64_REGISTER_NUM_LIMBS as u32).inverse(), U16_BITS - 3, ) .eval(builder, is_valid.clone()); - self.range_bus - .range_check( - local_cols.mem_ptr_limbs[1], - self.pointer_max_bits - U16_BITS, - ) - .eval(builder, is_valid.clone()); - - let mem_ptr = local_cols.mem_ptr_limbs[0] - + local_cols.mem_ptr_limbs[1] * AB::F::from_u32(1u32 << U16_BITS); // Constrain loads to address space 2 and stores to address spaces 2 or 3. let mem_as_minus_two = local_cols.mem_as - AB::Expr::TWO; @@ -225,24 +224,48 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { .when(not::(is_valid.clone())) .assert_zero(local_cols.mem_as); + // Exactly one of {read, write} is a heap access (read for loads, write for stores); the + // other is a register access. The heap access uses the 8-byte-aligned byte pointer + // `mem_ptr - shift_amount`; convert it ONCE to AS-native u16 *cell* limbs (range-checked). + // The register access uses `rd_rs2_ptr` via the small-pointer helper (no range check, since + // register locations are small). Route heap/register limbs to read/write by `is_load`. + // (`shift_amount` is the per-instruction shift: `load_shift` for loads, `store_shift` for + // stores, so the aligned heap pointer is correct for whichever side is the heap access.) + let byte_ptr_max_bits = self.pointer_max_bits; + let heap_cell_limbs = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + [ + local_cols.mem_ptr_limbs[0] - shift_amount, + local_cols.mem_ptr_limbs[1].into(), + ], + local_cols.mem_ptr_carry, + byte_ptr_max_bits, + is_valid.clone(), + ); + let reg_cell_limbs = reg_byte_ptr_to_cell_ptr_limbs::(local_cols.rd_rs2_ptr); + // read_as is [local_cols.mem_as] for loads and 1 for stores let read_as = select::( is_load.clone(), local_cols.mem_as, AB::F::from_u32(RV64_REGISTER_AS), ); - - // read_ptr is mem_ptr for loads and rd_rs2_ptr for stores - // Note: shift_amount is expected to have degree 2, thus we can't put it in the select - // clause since the resulting read_ptr/write_ptr's degree will be 3 which is - // too high. Instead, the solution without using additional columns is to get - // two different shift amounts from core chip - let read_ptr = select::(is_load.clone(), mem_ptr.clone(), local_cols.rd_rs2_ptr) - - load_shift_amount; - + let read_ptr_limbs = [ + select::( + is_load.clone(), + heap_cell_limbs[0].clone(), + reg_cell_limbs[0].clone(), + ), + select::( + is_load.clone(), + heap_cell_limbs[1].clone(), + reg_cell_limbs[1].clone(), + ), + ]; self.memory_bridge .read( - MemoryAddress::new(read_as, byte_ptr_to_u16_ptr::(read_ptr)), + MemoryAddress::new(read_as, read_ptr_limbs), pack_u8_block::(&ctx.reads.1), timestamp_pp(), &local_cols.read_data_aux, @@ -255,17 +278,25 @@ impl VmAdapterAir for Rv64LoadStoreAdapterAir { AB::F::from_u32(RV64_REGISTER_AS), local_cols.mem_as, ); - - // write_ptr is rd_rs2_ptr for loads and mem_ptr for stores - let write_ptr = select::(is_load.clone(), local_cols.rd_rs2_ptr, mem_ptr.clone()) - - store_shift_amount; + let write_ptr_limbs = [ + select::( + is_load.clone(), + reg_cell_limbs[0].clone(), + heap_cell_limbs[0].clone(), + ), + select::( + is_load.clone(), + reg_cell_limbs[1].clone(), + heap_cell_limbs[1].clone(), + ), + ]; // The core supplies the previous write bytes; this adapter stores only the base aux // columns. let prev_data_expr: [AB::Expr; MEMORY_BLOCK_BYTES] = ctx.reads.0.map(Into::into); self.memory_bridge .write( - MemoryAddress::new(write_as, byte_ptr_to_u16_ptr::(write_ptr)), + MemoryAddress::new(write_as, write_ptr_limbs), pack_u8_block::(&ctx.writes[0].clone()), timestamp_pp(), MemoryWriteAuxInput::from_prev_data_exprs( @@ -321,6 +352,8 @@ pub struct Rv64LoadStoreAdapterRecord { pub imm_sign: bool, pub mem_as: u8, + /// Whether this is a load (heap access is the read) vs a store (heap access is the write). + pub is_load: bool, pub write_prev_timestamp: u32, } @@ -411,6 +444,7 @@ where LOADD | LOADW | LOADB | LOADH | LOADBU | LOADHU | LOADWU => { debug_assert_eq!(e, F::from_u32(RV64_MEMORY_AS)); record.mem_as = RV64_MEMORY_AS as u8; + record.is_load = true; let read_data = tracing_read( memory, RV64_MEMORY_AS, @@ -425,6 +459,7 @@ where debug_assert_ne!(e, RV64_IMM_AS); debug_assert_ne!(e, RV64_REGISTER_AS); record.mem_as = e as u8; + record.is_load = false; let read_data = tracing_read( memory, RV64_REGISTER_AS, @@ -523,12 +558,24 @@ impl AdapterTraceFiller for Rv64LoadStoreAdapterFiller { .rs1_val .wrapping_add(sign_extend_imm16(record.imm as u32, record.imm_sign as u32)); let ptr_limbs = ptr_to_u16_limbs(ptr).map(u32::from); + // Alignment check: `(mem_ptr_limbs[0] - shift) / 8 < 2^13`. self.range_checker_chip .add_count(ptr_limbs[0] >> 3, U16_BITS - 3); - self.range_checker_chip - .add_count(ptr_limbs[1], self.pointer_max_bits - U16_BITS); adapter_row.mem_ptr_limbs = ptr_limbs.map(F::from_u32); + // Convert the aligned heap byte pointer to AS-native cell pointer limbs and register the + // single set of range-check counts (the AIR converts the heap pointer once, with + // `enabled = is_valid`). `mem_ptr_carry` depends only on the high byte limb's parity. The + // register-side access uses the small-pointer helper, which is NOT range-checked. + let hi_bits = cell_ptr_hi_bits(self.pointer_max_bits); + let shift = ptr & (RV64_REGISTER_NUM_LIMBS as u32 - 1); + let aligned_byte_limbs = ptr_to_u16_limbs(ptr - shift).map(u32::from); + let (mem_carry, heap_cell_limbs) = + byte_ptr_limbs_to_cell_ptr_limbs_value(aligned_byte_limbs); + adapter_row.mem_ptr_carry = F::from_u32(mem_carry); + self.range_checker_chip + .add_count(heap_cell_limbs[1], hi_bits); + adapter_row.imm_sign = F::from_bool(record.imm_sign); adapter_row.imm = F::from_u16(record.imm); diff --git a/extensions/riscv/circuit/src/adapters/mod.rs b/extensions/riscv/circuit/src/adapters/mod.rs index 035e294d80..3c2ac1eb75 100644 --- a/extensions/riscv/circuit/src/adapters/mod.rs +++ b/extensions/riscv/circuit/src/adapters/mod.rs @@ -1,12 +1,15 @@ use std::ops::Mul; use openvm_circuit::{ - arch::{execution_mode::ExecutionCtxTrait, VmStateMut, BLOCK_FE_WIDTH}, + arch::{execution_mode::ExecutionCtxTrait, VmStateMut, BLOCK_FE_WIDTH, U16_CELL_SIZE_BITS}, system::memory::{ merkle::public_values::PUBLIC_VALUES_AS, online::{GuestMemory, TracingMemory}, }, }; +use openvm_circuit_primitives::var_range::{ + SharedVariableRangeCheckerChip, VariableRangeCheckerBus, +}; pub use openvm_circuit_primitives::U16_BITS; use openvm_instructions::{ riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}, @@ -14,6 +17,7 @@ use openvm_instructions::{ }; use openvm_stark_backend::{ interaction::InteractionBuilder, + p3_air::AirBuilder, p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, }; @@ -296,6 +300,185 @@ where }) } +// ---------------------------------------------------------------------------- +// AS-native pointer-limb helpers. +// +// Every memory-bus pointer is two little-endian 16-bit *AS-native cell* pointer limbs +// `[lo16, hi16]` (see `openvm_circuit::system::memory::MemoryAddress`). These helpers convert +// between RV64 *byte* pointers (read from registers) and AS-native *cell* pointer limbs without +// composing a full (up to 31-bit) pointer into one field element. +// ---------------------------------------------------------------------------- + +/// AS-native memory pointer represented as little-endian 16-bit limbs `[lo16, hi16]`. +pub type PtrLimbs = [T; 2]; + +/// Splits a concrete pointer into little-endian 16-bit limb *values* `[lo16, hi16]`. +#[inline(always)] +pub fn u32_to_ptr_limbs(ptr: u32) -> PtrLimbs { + [ptr & 0xffff, ptr >> U16_BITS] +} + +/// Recomposes little-endian 16-bit limb values into a `u32`. +#[inline(always)] +pub fn ptr_limbs_to_u32(limbs: PtrLimbs) -> u32 { + limbs[0] | (limbs[1] << U16_BITS) +} + +/// AS-native cell-pointer limbs for a byte pointer in the register address space +/// ([`RV64_REGISTER_AS`]). +/// +/// The register file holds at most `NUM_RV64_REGISTERS * 8` bytes, so a register byte pointer's +/// cell pointer `ptr / 2` is far below `2^16`: it fits entirely in the low 16-bit limb and the +/// high limb is always zero. This lets us skip the carry/decomposition columns and range checks +/// that a general (up to `POINTER_MAX_BITS`-bit) memory pointer requires. Only use this for +/// register-AS pointers; for the memory address space use the range-checked decomposition helpers. +#[inline(always)] +pub fn reg_byte_ptr_to_cell_ptr_limbs( + byte_ptr: impl Into, +) -> PtrLimbs { + [byte_ptr_to_u16_ptr::(byte_ptr), AB::Expr::ZERO] +} + +/// Value form of [`reg_byte_ptr_to_cell_ptr_limbs`]. +#[inline(always)] +pub fn reg_byte_ptr_to_cell_ptr_limbs_value(byte_ptr: u32) -> PtrLimbs { + [byte_ptr_to_u16_ptr_value(byte_ptr), 0] +} + +/// Converts an aligned RV64 byte pointer given as little-endian 16-bit limbs `[byte_lo, byte_hi]` +/// into AS-native u16 *cell* pointer limbs `[cell_lo, cell_hi]` (cell = byte / 2). +/// +/// `carry` is a witness boolean intended to equal `byte_hi & 1`. The returned limbs are the +/// expressions +/// cell_lo = (byte_lo + carry * 2^16) / 2, cell_hi = (byte_hi - carry) / 2. +/// The composed cell pointer `cell_lo + 2^16 * cell_hi` equals `byte_ptr / 2` as a field identity +/// for *any* `carry`. The caller must already constrain `byte_lo` to be a canonical 16-bit value +/// divisible by 8; this makes `cell_lo < 2^16` for either boolean carry. This function range-checks +/// only `cell_hi < 2^cell_hi_bits`, where the high-limb bound is derived from `byte_ptr_max_bits` +/// (the guest *byte* pointer width): a u16 cell is two bytes, so +/// `cell_max_bits = byte_ptr_max_bits - U16_CELL_SIZE_BITS` and +/// `cell_hi_bits = cell_max_bits - U16_BITS`. Since `cell_hi` is a bounded integer expression, this +/// also forces `carry = byte_hi & 1`. +#[allow(clippy::too_many_arguments)] +pub fn eval_byte_ptr_limbs_to_u16_cell_ptr_limbs( + builder: &mut AB, + range_bus: VariableRangeCheckerBus, + byte_limbs: [AB::Expr; 2], + carry: impl Into, + byte_ptr_max_bits: usize, + enabled: AB::Expr, +) -> PtrLimbs { + let cell_hi_bits = byte_ptr_max_bits - U16_CELL_SIZE_BITS - U16_BITS; + let carry_e: AB::Expr = carry.into(); + builder.when(enabled.clone()).assert_bool(carry_e.clone()); + let inv2 = AB::F::TWO.inverse(); + let [byte_lo, byte_hi] = byte_limbs; + let cell_lo = (byte_lo + carry_e.clone() * AB::F::from_u32(1 << U16_BITS)) * inv2; + let cell_hi = (byte_hi - carry_e) * inv2; + range_bus + .range_check(cell_hi.clone(), cell_hi_bits) + .eval(builder, enabled); + [cell_lo, cell_hi] +} + +/// Cell high-limb range-check bit width corresponding to a guest `byte_ptr_max_bits`. +#[inline(always)] +pub fn cell_ptr_hi_bits(byte_ptr_max_bits: usize) -> usize { + byte_ptr_max_bits - U16_CELL_SIZE_BITS - U16_BITS +} + +/// Adds a small constant `constant` (`< 2^16`) to a pointer given as little-endian 16-bit limbs +/// `[lo, hi]`, carrying into the high limb: +/// new_lo = lo + constant - carry * 2^16, new_hi = hi + carry. +/// `carry` is a witness boolean. Only `new_lo` is range-checked (to 16 bits): this forces `carry` +/// to be the correct carry bit (given `lo` canonical), so `new_hi = hi + carry` is canonical +/// whenever `hi` is. Use to add a per-block cell offset to an already-converted base cell pointer. +#[allow(clippy::too_many_arguments)] +pub fn eval_add_const_u16_limbs( + builder: &mut AB, + range_bus: VariableRangeCheckerBus, + limbs: [AB::Expr; 2], + constant: u32, + carry: AB::Var, + enabled: AB::Expr, +) -> PtrLimbs { + let carry_e: AB::Expr = carry.into(); + builder.when(enabled.clone()).assert_bool(carry_e.clone()); + let [lo, hi] = limbs; + let new_lo = + lo + AB::Expr::from_u32(constant) - carry_e.clone() * AB::F::from_u32(1 << U16_BITS); + let new_hi = hi + carry_e; + range_bus + .range_check(new_lo.clone(), U16_BITS) + .eval(builder, enabled); + [new_lo, new_hi] +} + +/// Value form of [`eval_add_const_u16_limbs`]: returns `(carry, [new_lo, new_hi])`. +#[inline(always)] +pub fn add_const_u16_limbs_value(limbs: PtrLimbs, constant: u32) -> (u32, PtrLimbs) { + let sum_lo = limbs[0] + constant; + let carry = sum_lo >> U16_BITS; + (carry, [sum_lo & 0xffff, limbs[1] + carry]) +} + +/// Computes one add-carry per memory block from an already-converted base cell pointer, +/// registering the matching range checks for each block's new low limb. +pub fn compute_block_add_carries( + range_checker: &SharedVariableRangeCheckerChip, + base_cell: [u16; 2], + num_blocks: usize, + cell_stride: u32, +) -> Vec { + let base_cell = base_cell.map(u32::from); + (0..num_blocks) + .map(|i| { + let (add_carry, block_cell_ptr) = + add_const_u16_limbs_value(base_cell, i as u32 * cell_stride); + range_checker.add_count(block_cell_ptr[0], U16_BITS); + add_carry + }) + .collect() +} + +/// Value form of [`eval_byte_ptr_limbs_to_u16_cell_ptr_limbs`]. Returns +/// `(carry, [cell_lo, cell_hi])` for an aligned byte pointer given as little-endian 16-bit limb +/// values. The caller is responsible for registering the matching range-check for `cell_hi` +/// to `hi_bits`. +#[inline(always)] +pub fn byte_ptr_limbs_to_cell_ptr_limbs_value(byte_limbs: PtrLimbs) -> (u32, PtrLimbs) { + let carry = byte_limbs[1] & 1; + let cell_lo = (byte_limbs[0] + (carry << U16_BITS)) >> 1; + let cell_hi = byte_limbs[1] >> 1; + (carry, [cell_lo, cell_hi]) +} + +/// Computes the byte->cell conversion carry and one add-carry per block for a heap +/// access group, registering the matching range checks. +/// +/// Returns `(conv_carry, add_carries)`. +/// +/// Column writes are left to the caller because vec_heap-family fillers must buffer +/// carries before overwriting their records. +pub fn compute_pointer_carries( + range_checker: &SharedVariableRangeCheckerChip, + byte_ptr: u32, + num_blocks: usize, + cell_stride: u32, + byte_ptr_max_bits: usize, +) -> (u32, Vec) { + let byte_limbs = u32_to_ptr_limbs(byte_ptr); + let (conv_carry, base_cell) = byte_ptr_limbs_to_cell_ptr_limbs_value(byte_limbs); + range_checker.add_count(base_cell[1], cell_ptr_hi_bits(byte_ptr_max_bits)); + let add_carries = compute_block_add_carries( + range_checker, + base_cell.map(|limb| limb as u16), + num_blocks, + cell_stride, + ); + (conv_carry, add_carries) +} + /// Expand `N` limbs to `RV64_REGISTER_NUM_LIMBS` (8) by zero-padding the upper limbs. Used for /// register bus reads where the register holds a value in fewer than 8 bytes. pub fn expand_to_rv64_register, T: PrimeCharacteristicRing, const N: usize>( diff --git a/extensions/riscv/circuit/src/adapters/mul.rs b/extensions/riscv/circuit/src/adapters/mul.rs index 96a4e8bb5d..7168b2013f 100644 --- a/extensions/riscv/circuit/src/adapters/mul.rs +++ b/extensions/riscv/circuit/src/adapters/mul.rs @@ -29,7 +29,7 @@ use openvm_stark_backend::{ }; use super::{tracing_write, RV64_REGISTER_NUM_LIMBS}; -use crate::adapters::{byte_ptr_to_u16_ptr, tracing_read}; +use crate::adapters::{reg_byte_ptr_to_cell_ptr_limbs, tracing_read}; #[repr(C)] #[derive(AlignedBorrow, StructReflection)] @@ -85,7 +85,7 @@ impl VmAdapterAir for Rv64MultAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), pack_u8_block::(&ctx.reads[0].clone()), timestamp_pp(), @@ -97,7 +97,7 @@ impl VmAdapterAir for Rv64MultAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs2_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2_ptr), ), pack_u8_block::(&ctx.reads[1].clone()), timestamp_pp(), @@ -109,7 +109,7 @@ impl VmAdapterAir for Rv64MultAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), pack_u8_block::(&ctx.writes[0].clone()), timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/mul_w.rs b/extensions/riscv/circuit/src/adapters/mul_w.rs index 4c545fd547..890c186841 100644 --- a/extensions/riscv/circuit/src/adapters/mul_w.rs +++ b/extensions/riscv/circuit/src/adapters/mul_w.rs @@ -32,8 +32,8 @@ use openvm_stark_backend::{ }; use super::{ - byte_ptr_to_u16_ptr, pack_high_u16, pack_rv64_u16_block, tracing_read, tracing_write, - RV64_PTR_U16_LIMBS, + pack_high_u16, pack_rv64_u16_block, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_write, RV64_PTR_U16_LIMBS, }; #[repr(C)] @@ -102,7 +102,7 @@ impl VmAdapterAir for Rv64MultWAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs1_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs1_ptr), ), rs1_data, timestamp_pp(), @@ -116,7 +116,7 @@ impl VmAdapterAir for Rv64MultWAdapterAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rs2_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rs2_ptr), ), rs2_data, timestamp_pp(), @@ -144,7 +144,7 @@ impl VmAdapterAir for Rv64MultWAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local.rd_ptr), ), write_data, timestamp_pp(), diff --git a/extensions/riscv/circuit/src/adapters/rdwrite.rs b/extensions/riscv/circuit/src/adapters/rdwrite.rs index dba9d9b9b9..02c099d447 100644 --- a/extensions/riscv/circuit/src/adapters/rdwrite.rs +++ b/extensions/riscv/circuit/src/adapters/rdwrite.rs @@ -28,7 +28,9 @@ use openvm_stark_backend::{ p3_field::{Field, PrimeCharacteristicRing, PrimeField32}, }; -use crate::adapters::{byte_ptr_to_u16_ptr, byte_ptr_to_u16_ptr_value, tracing_write_u16}; +use crate::adapters::{ + byte_ptr_to_u16_ptr_value, reg_byte_ptr_to_cell_ptr_limbs, tracing_write_u16, +}; #[repr(C)] #[derive(Debug, Clone, AlignedBorrow, StructReflection)] @@ -103,7 +105,7 @@ impl Rv64RdWriteAdapterAir { .write( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.rd_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.rd_ptr), ), ctx.writes[0].clone(), timestamp, diff --git a/extensions/riscv/circuit/src/branch_eq/tests.rs b/extensions/riscv/circuit/src/branch_eq/tests.rs index c72bfd0ee1..e9b04dafea 100644 --- a/extensions/riscv/circuit/src/branch_eq/tests.rs +++ b/extensions/riscv/circuit/src/branch_eq/tests.rs @@ -2,7 +2,9 @@ use std::{array, borrow::BorrowMut}; use openvm_circuit::{ arch::{ - testing::{memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder}, + testing::{ + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, VmChipTestBuilder, + }, Arena, ExecutionBridge, PreflightExecutor, BLOCK_FE_WIDTH, }, system::memory::{offline_checker::MemoryBridge, SharedMemoryHelper}, @@ -108,8 +110,7 @@ fn set_and_execute>( }); let imm = imm.unwrap_or(rng.random_range((-ABS_MAX_IMM)..ABS_MAX_IMM)); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1, rs2] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes::(1, rs1, a.map(F::from_u8)); tester.write_bytes::(1, rs2, b.map(F::from_u8)); diff --git a/extensions/riscv/circuit/src/branch_lt/tests.rs b/extensions/riscv/circuit/src/branch_lt/tests.rs index 4b194a8e90..6df728c13b 100644 --- a/extensions/riscv/circuit/src/branch_lt/tests.rs +++ b/extensions/riscv/circuit/src/branch_lt/tests.rs @@ -4,7 +4,9 @@ use std::{array, borrow::BorrowMut}; use openvm_circuit::{ arch::{ - testing::{memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder}, + testing::{ + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, VmChipTestBuilder, + }, Arena, ExecutionBridge, PreflightExecutor, BLOCK_FE_WIDTH, }, system::memory::{offline_checker::MemoryBridge, SharedMemoryHelper}, @@ -119,8 +121,7 @@ fn set_and_execute>( }); let imm = imm.unwrap_or(rng.random_range((-ABS_MAX_IMM)..ABS_MAX_IMM)); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1, rs2] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let a_bytes: [F; RV64_REGISTER_NUM_LIMBS] = rv64_u16_block_to_bytes(a).map(F::from_u8); let b_bytes: [F; RV64_REGISTER_NUM_LIMBS] = rv64_u16_block_to_bytes(b).map(F::from_u8); tester.write_bytes::(1, rs1, a_bytes); diff --git a/extensions/riscv/circuit/src/divrem/tests.rs b/extensions/riscv/circuit/src/divrem/tests.rs index e6a39d5a10..a7e0241680 100644 --- a/extensions/riscv/circuit/src/divrem/tests.rs +++ b/extensions/riscv/circuit/src/divrem/tests.rs @@ -3,8 +3,8 @@ use std::{array, borrow::BorrowMut, sync::Arc}; use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, + VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, }, Arena, ExecutionBridge, PreflightExecutor, }, @@ -157,9 +157,7 @@ fn set_and_execute>( rng.random_range(0..(RV64_REGISTER_NUM_LIMBS - 1)), )); - let rs1 = gen_pointer(rng, 8); - let rs2 = gen_pointer(rng, 8); - let rd = gen_pointer(rng, 8); + let [rs1, rs2, rd] = gen_distinct_register_pointers(rng, 8); tester.write_bytes::(1, rs1, b.map(F::from_u32)); tester.write_bytes::(1, rs2, c.map(F::from_u32)); diff --git a/extensions/riscv/circuit/src/divrem_w/tests.rs b/extensions/riscv/circuit/src/divrem_w/tests.rs index a16f6c6273..d74ee98c76 100644 --- a/extensions/riscv/circuit/src/divrem_w/tests.rs +++ b/extensions/riscv/circuit/src/divrem_w/tests.rs @@ -3,8 +3,8 @@ use std::{array, borrow::BorrowMut, sync::Arc}; use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, + VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, }, Arena, ExecutionBridge, PreflightExecutor, }, @@ -179,9 +179,7 @@ fn set_and_execute>( }); // Write full 8-byte registers. Upper bytes are arbitrary and remain adapter-constrained. - let rs1_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rd_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1_ptr, rs2_ptr, rd_ptr] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes::(1, rs1_ptr, b.map(F::from_u32)); tester.write_bytes::(1, rs2_ptr, c.map(F::from_u32)); diff --git a/extensions/riscv/circuit/src/hintstore/mod.rs b/extensions/riscv/circuit/src/hintstore/mod.rs index 32169815e3..7ddb493fbe 100644 --- a/extensions/riscv/circuit/src/hintstore/mod.rs +++ b/extensions/riscv/circuit/src/hintstore/mod.rs @@ -38,9 +38,10 @@ use openvm_stark_backend::{ }; use crate::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, ptr_bound_from_ptr, - ptr_to_field_u16_limbs, read_rv64_register_as_u32, tracing_read, tracing_read_reg_ptr, - tracing_write, u16_limbs_to_ptr, RV64_PTR_BITS, RV64_PTR_U16_LIMBS, U16_BITS, + byte_ptr_limbs_to_cell_ptr_limbs_value, cell_ptr_hi_bits, + eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, ptr_to_field_u16_limbs, + ptr_to_u16_limbs, read_rv64_register_as_u32, reg_byte_ptr_to_cell_ptr_limbs, tracing_read, + tracing_read_reg_ptr, tracing_write, RV64_PTR_U16_LIMBS, U16_BITS, }; mod execution; @@ -76,10 +77,16 @@ pub struct Rv64HintStoreCols { pub from_state: ExecutionState, pub mem_ptr_ptr: T, - /// Low 32 bits of the 8-byte RV64 register that holds `mem_ptr`. `mem_ptr` is a - /// u32 memory address, so the upper 4 bytes are known to be zero and are hardcoded - /// in the memory bus interaction rather than materialized as columns. + /// Low 32 bits of the 8-byte RV64 register that holds `mem_ptr`, as little-endian 16-bit + /// *byte*-pointer limbs `[lo16, hi16]`. The byte pointer may span the full 2^32 byte address + /// space. pub mem_ptr_limbs: [T; RV64_PTR_U16_LIMBS], + /// Carry (`mem_ptr_limbs[1] & 1`) for converting the byte pointer to AS-native u16 *cell* + /// pointer limbs. See `eval_byte_ptr_limbs_to_u16_cell_ptr_limbs`. + pub mem_ptr_carry: T, + /// Carry for the per-row `next.mem_ptr = mem_ptr + 8` byte increment (computed limb-wise to + /// avoid composing a 32-bit pointer into one field element). + pub mem_ptr_inc_carry: T, pub mem_ptr_aux_cols: MemoryReadAuxCols, pub write_aux: MemoryWriteAuxCols, @@ -144,9 +151,6 @@ impl Air for Rv64HintStoreAir { let rem_words: AB::Expr = local_cols.rem_words.into(); let next_rem_words: AB::Expr = next_cols.rem_words.into(); - let mem_ptr: AB::Expr = u16_limbs_to_ptr(&local_cols.mem_ptr_limbs); - let next_mem_ptr: AB::Expr = u16_limbs_to_ptr(&next_cols.mem_ptr_limbs); - // Constrain that if local is invalid, then the next state is invalid as well builder .when_transition() @@ -168,7 +172,7 @@ impl Air for Rv64HintStoreAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.mem_ptr_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.mem_ptr_ptr), ), mem_ptr_data, timestamp_pp(), @@ -187,7 +191,7 @@ impl Air for Rv64HintStoreAir { .read( MemoryAddress::new( AB::F::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(local_cols.num_words_ptr), + reg_byte_ptr_to_cell_ptr_limbs::(local_cols.num_words_ptr), ), num_words_data, timestamp_pp(), @@ -195,13 +199,19 @@ impl Air for Rv64HintStoreAir { ) .eval(builder, local_cols.is_buffer_start); - // write hint + // write hint: convert the (aligned) heap byte pointer `mem_ptr_limbs` to AS-native cell + // pointer limbs without composing the full byte pointer into one field element. + let mem_ptr_cell_limbs = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + local_cols.mem_ptr_limbs.map(Into::into), + local_cols.mem_ptr_carry, + self.pointer_max_bits, + is_valid.clone(), + ); self.memory_bridge .write( - MemoryAddress::new( - AB::F::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::(mem_ptr.clone()), - ), + MemoryAddress::new(AB::F::from_u32(RV64_MEMORY_AS), mem_ptr_cell_limbs), local_cols.data.map(Into::into), timestamp_pp(), &local_cols.write_aux, @@ -226,21 +236,13 @@ impl Air for Rv64HintStoreAir { ) .eval(builder, is_start.clone()); - assert!( - (U16_BITS..=RV64_PTR_BITS).contains(&self.pointer_max_bits), - "pointer_max_bits must fit in the low 32-bit mem_ptr view" - ); - - // Preventing mem_ptr overflow: mem_ptr < 2^pointer_max_bits. + // Range-check the low byte-pointer limb to 16 bits so the limb-wise `+8` increment below + // is sound. The high byte limb is bounded by the cell-pointer range checks in + // `eval_byte_ptr_limbs_to_u16_cell_ptr_limbs` above (`cell_hi < 2^(pointer_max_bits - 16)` + // implies `byte_hi < 2^16`, i.e. byte pointer `< 2^32`). self.range_bus - .range_check( - ptr_bound_from_high_u16_expr( - local_cols.mem_ptr_limbs[RV64_PTR_U16_LIMBS - 1], - self.pointer_max_bits, - ), - U16_BITS, - ) - .eval(builder, is_start.clone()); + .range_check(local_cols.mem_ptr_limbs[0], U16_BITS) + .eval(builder, is_valid.clone()); // Preventing rem_words overflow: rem_words < 2^MAX_HINT_BUFFER_DWORDS_BITS. self.range_bus .range_check( @@ -269,15 +271,22 @@ impl Air for Rv64HintStoreAir { // additional `buffer` rows we will always increment `mem_ptr` to an illegal memory address // at some point, which prevents this exploit. when_buffer_transition.assert_one(rem_words.clone() - next_rem_words.clone()); - // Note: we only care about the composed `next_mem_ptr` and not the individual limbs: - // the limbs do not need to be in the range, they can be anything that makes - // `next_mem_ptr` correct -- this is just a way to avoid another column for `mem_ptr`. - // The constraint we care about is `next.mem_ptr == local.mem_ptr + 8`. Since we increment - // by `8` each time, any out of bounds memory access will be rejected by the memory bus - // before we overflow the field. + // `next.mem_ptr == local.mem_ptr + 8`, computed limb-wise to avoid composing a 32-bit byte + // pointer into one field element: + // next_lo = lo + 8 - inc_carry * 2^16 + // next_hi = hi + inc_carry + // with `inc_carry` boolean. The byte limbs are canonical (each `< 2^16`, range-checked), + // so this pins the increment exactly. + let inc_carry = local_cols.mem_ptr_inc_carry; + when_buffer_transition.assert_bool(inc_carry); when_buffer_transition.assert_eq( - next_mem_ptr.clone() - mem_ptr.clone(), - AB::F::from_usize(RV64_REGISTER_NUM_LIMBS), + next_cols.mem_ptr_limbs[0], + local_cols.mem_ptr_limbs[0] + AB::F::from_usize(RV64_REGISTER_NUM_LIMBS) + - inc_carry * AB::F::from_u32(1u32 << U16_BITS), + ); + when_buffer_transition.assert_eq( + next_cols.mem_ptr_limbs[1], + local_cols.mem_ptr_limbs[1] + inc_carry, ); when_buffer_transition.assert_eq( timestamp + AB::F::from_usize(timestamp_delta), @@ -457,7 +466,7 @@ where ); debug_assert_ne!(num_words, 0); - debug_assert!(num_words <= (1 << self.pointer_max_bits)); + debug_assert!(u64::from(num_words) <= (1u64 << self.pointer_max_bits)); record.inner.num_words = num_words; if local_opcode == HINT_STORED { @@ -557,15 +566,12 @@ impl TraceFiller for Rv64HintStoreFiller { num_words <= MAX_HINT_BUFFER_DWORDS as u32, "num_words must be <= MAX_HINT_BUFFER_DWORDS" ); - // Range check for mem_ptr (using pointer_max_bits) and num_words (using - // MAX_HINT_BUFFER_DWORDS_BITS). - self.range_checker_chip.add_count( - ptr_bound_from_ptr(record.inner.mem_ptr, self.pointer_max_bits), - U16_BITS, - ); + // Range check num_words (using MAX_HINT_BUFFER_DWORDS_BITS). Per-row pointer-limb + // range checks are added in the row loop below. self.range_checker_chip .add_count(num_words << REM_WORDS_SHIFT, U16_BITS); + let hi_bits = cell_ptr_hi_bits(self.pointer_max_bits); let mut timestamp = record.inner.timestamp + num_words * 3; let mut mem_ptr = record.inner.mem_ptr + num_words * RV64_REGISTER_NUM_LIMBS as u32; @@ -618,6 +624,20 @@ impl TraceFiller for Rv64HintStoreFiller { cols.mem_ptr_limbs = ptr_to_field_u16_limbs(mem_ptr); cols.mem_ptr_ptr = F::from_u32(record.inner.mem_ptr_ptr); + // Byte -> cell pointer conversion (heap write) and the per-row range + // checks: cell_hi (hi_bits) and the low byte limb + // (16 bits, for the limb-wise `+8` increment). + let byte_limbs = ptr_to_u16_limbs(mem_ptr).map(u32::from); + let (mem_carry, cell_limbs) = + byte_ptr_limbs_to_cell_ptr_limbs_value(byte_limbs); + cols.mem_ptr_carry = F::from_u32(mem_carry); + // `+8` carry from this row's low byte limb into the high limb. + cols.mem_ptr_inc_carry = F::from_u32( + (byte_limbs[0] + RV64_REGISTER_NUM_LIMBS as u32) >> U16_BITS, + ); + self.range_checker_chip.add_count(cell_limbs[1], hi_bits); + self.range_checker_chip.add_count(byte_limbs[0], U16_BITS); + cols.from_state.timestamp = F::from_u32(timestamp); cols.from_state.pc = F::from_u32(record.inner.from_pc); diff --git a/extensions/riscv/circuit/src/hintstore/tests.rs b/extensions/riscv/circuit/src/hintstore/tests.rs index 11a9335f37..75f57a5ac1 100644 --- a/extensions/riscv/circuit/src/hintstore/tests.rs +++ b/extensions/riscv/circuit/src/hintstore/tests.rs @@ -4,7 +4,10 @@ use std::sync::Arc; use openvm_circuit::{ arch::{ - testing::{memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder}, + testing::{ + memory::{gen_distinct_register_pointers, gen_pointer, gen_register_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, + }, Arena, ExecutionBridge, MatrixRecordArena, PreflightExecutor, BLOCK_FE_WIDTH, }, system::memory::{ @@ -99,20 +102,19 @@ fn set_and_execute>( HINT_BUFFER => rng.random_range(1..28), } as u32; + let [a_reg, b] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let a = if opcode == HINT_BUFFER { - let a = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, - a, + a_reg, u64_to_rv64_limbs(num_words.into()), ); - a + a_reg } else { 0 }; let mem_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS) as u32; - let b = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, b, @@ -199,7 +201,7 @@ fn test_hint_buffer_exceeds_max_words() { let num_words = (MAX_HINT_BUFFER_DWORDS + 1) as u32; - let a = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); + let [a, b] = gen_distinct_register_pointers(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, a, @@ -207,7 +209,6 @@ fn test_hint_buffer_exceeds_max_words() { ); let mem_ptr = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS) as u32; - let b = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, b, @@ -237,7 +238,7 @@ fn test_hint_buffer_rem_words_range_check() { let mut harness = create_harness(&mut tester); let num_words: u32 = 1; - let a = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); + let [a, b] = gen_distinct_register_pointers(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, a, @@ -245,7 +246,6 @@ fn test_hint_buffer_rem_words_range_check() { ); let mem_ptr = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS) as u32; - let b = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, b, @@ -294,7 +294,7 @@ fn test_hint_buffer_mem_ptr_range_check() { let mut harness = create_harness(&mut tester); let num_words: u32 = 1; - let a = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); + let [a, b] = gen_distinct_register_pointers(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, a, @@ -302,7 +302,6 @@ fn test_hint_buffer_mem_ptr_range_check() { ); let mem_ptr = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS) as u32; - let b = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); tester.write_bytes( RV64_REGISTER_AS as usize, b, @@ -353,7 +352,7 @@ fn test_hintstore_rs1_upper_bytes_non_zero() { // Write b with a non-zero byte in the upper half; `mem_ptr_u64 >> 32` is then non-zero, // so the preflight executor must panic before it reaches the data write. - let b = gen_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); + let b = gen_register_pointer(&mut rng, RV64_REGISTER_NUM_LIMBS); let mut mem_ptr_limbs = [F::ZERO; RV64_REGISTER_NUM_LIMBS]; mem_ptr_limbs[4] = F::from_u8(1); tester.write_bytes(RV64_REGISTER_AS as usize, b, mem_ptr_limbs); diff --git a/extensions/riscv/circuit/src/load_sign_extend/tests.rs b/extensions/riscv/circuit/src/load_sign_extend/tests.rs index 1c19f4b6d3..beb55909c9 100644 --- a/extensions/riscv/circuit/src/load_sign_extend/tests.rs +++ b/extensions/riscv/circuit/src/load_sign_extend/tests.rs @@ -3,8 +3,8 @@ use std::{array, borrow::BorrowMut, sync::Arc}; use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, + VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, Arena, ExecutionBridge, PreflightExecutor, }, @@ -153,15 +153,15 @@ fn set_and_execute>( _ => unreachable!(), }; - let ptr_val: u32 = rng.random_range(0..(1 << (tester.address_bits() - alignment))) << alignment; + let ptr_val: u32 = + (rng.random_range(0..(1u64 << (tester.address_bits() - alignment))) << alignment) as u32; // rs1 is 8 bytes, but only low 4 bytes used for address let rs1 = rs1.unwrap_or_else(|| { let low4 = ptr_val.wrapping_sub(imm_ext).to_le_bytes(); [low4[0], low4[1], low4[2], low4[3], 0, 0, 0, 0] }); let ptr_val = imm_ext.wrapping_add(rv64_bytes_to_u32(rs1)); - let a = gen_pointer(rng, 8); - let b = gen_pointer(rng, 8); + let [a, b] = gen_distinct_register_pointers(rng, 8); let shift_amount = ptr_val % 8; tester.write_bytes(1, b, rs1.map(F::from_u8)); diff --git a/extensions/riscv/circuit/src/loadstore/tests.rs b/extensions/riscv/circuit/src/loadstore/tests.rs index ab93679f2c..43ba24a893 100644 --- a/extensions/riscv/circuit/src/loadstore/tests.rs +++ b/extensions/riscv/circuit/src/loadstore/tests.rs @@ -2,7 +2,10 @@ use std::{array, borrow::BorrowMut, sync::Arc}; use openvm_circuit::{ arch::{ - testing::{TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS}, + testing::{ + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, + VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, + }, Arena, ExecutionBridge, MemoryConfig, PreflightExecutor, }, system::memory::{ @@ -16,7 +19,7 @@ use openvm_circuit_primitives::{ }, var_range::VariableRangeCheckerChip, }; -use openvm_instructions::{instruction::Instruction, riscv::RV64_REGISTER_AS, LocalOpcode}; +use openvm_instructions::{instruction::Instruction, LocalOpcode}; use openvm_riscv_transpiler::Rv64LoadStoreOpcode::{self, *}; use openvm_stark_backend::{ p3_air::BaseAir, @@ -149,18 +152,16 @@ fn set_and_execute>( _ => unreachable!("loadstore tests should not handle sign-extension load opcodes"), }; - let ptr_val: u32 = rng.random_range(0..(1 << (tester.address_bits() - alignment))) << alignment; + let ptr_val: u32 = + (rng.random_range(0..(1u64 << (tester.address_bits() - alignment))) << alignment) as u32; let ptr = ptr_val.wrapping_sub(imm_ext).to_le_bytes(); let rs1 = rs1.unwrap_or([ptr[0], ptr[1], ptr[2], ptr[3], 0, 0, 0, 0]); let rs1_low = rv64_bytes_to_u32(rs1); let ptr_val = imm_ext.wrapping_add(rs1_low); let shift_amount = (ptr_val as usize) & (RV64_REGISTER_NUM_LIMBS - 1); - let max_addr = 1usize << tester.address_bits(); - let a = rng.random_range(0..(max_addr - RV64_REGISTER_NUM_LIMBS)) / RV64_REGISTER_NUM_LIMBS - * RV64_REGISTER_NUM_LIMBS; - let b = rng.random_range(0..(max_addr - RV64_REGISTER_NUM_LIMBS)) / RV64_REGISTER_NUM_LIMBS - * RV64_REGISTER_NUM_LIMBS; + // `a`/`b` are register *locations* (instruction operands). + let [a, b] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let is_load = [LOADD, LOADWU, LOADHU, LOADBU].contains(&opcode); // Store tests choose writable u16-celled address spaces. @@ -254,9 +255,8 @@ fn set_and_execute>( fn rand_loadstore_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { let mut rng = create_seeded_rng(); let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; if [STORED, STOREW, STOREB, STOREH].contains(&opcode) { - mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; + mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 31; } let mut tester = VmChipTestBuilder::from_config(mem_config); let (mut harness, bitwise) = create_harness(&mut tester); @@ -286,9 +286,7 @@ fn rand_loadstore_test(opcode: Rv64LoadStoreOpcode, num_ops: usize) { #[test] fn positive_loadwu_shift4_test() { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - let mut tester = VmChipTestBuilder::from_config(mem_config); + let mut tester = VmChipTestBuilder::from_config(MemoryConfig::default()); let (mut harness, bitwise) = create_harness(&mut tester); set_and_execute( @@ -314,9 +312,7 @@ fn positive_loadwu_shift4_test() { #[test] fn positive_loadhu_shift6_test() { let mut rng = create_seeded_rng(); - let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - let mut tester = VmChipTestBuilder::from_config(mem_config); + let mut tester = VmChipTestBuilder::from_config(MemoryConfig::default()); let (mut harness, bitwise) = create_harness(&mut tester); set_and_execute( @@ -339,12 +335,47 @@ fn positive_loadhu_shift6_test() { tester.simple_test().expect("Verification failed"); } +#[test] +fn positive_loadstore_max_address_test() { + let mut rng = create_seeded_rng(); + let mut tester = VmChipTestBuilder::from_config(MemoryConfig::default()); + let (mut harness, bitwise) = create_harness(&mut tester); + + let top = 1u64 << tester.address_bits(); + let imm = 8u32; + for (opcode, byte_addr) in [ + (LOADD, top - 8), + (LOADBU, top - 1), + (STORED, top - 8), + (STOREB, top - 1), + ] { + let rs1 = (byte_addr as u32).wrapping_sub(imm).to_le_bytes(); + set_and_execute( + &mut tester, + &mut harness.executor, + &mut harness.arena, + &mut rng, + opcode, + Some([rs1[0], rs1[1], rs1[2], rs1[3], 0, 0, 0, 0]), + Some(imm), + Some(0), + Some(2), + ); + } + + let tester = tester + .build() + .load(harness) + .load_periphery(bitwise) + .finalize(); + tester.simple_test().expect("Verification failed"); +} + #[test] fn positive_storew_public_values_test() { let mut rng = create_seeded_rng(); let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; + mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 31; let mut tester = VmChipTestBuilder::from_config(mem_config); let (mut harness, bitwise) = create_harness(&mut tester); @@ -397,8 +428,7 @@ fn run_negative_loadstore_test( ) { let mut rng = create_seeded_rng(); let mut mem_config = MemoryConfig::default(); - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 29; - mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; + mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 31; let mut tester = VmChipTestBuilder::from_config(mem_config); let (mut harness, bitwise) = create_harness(&mut tester); @@ -858,7 +888,6 @@ fn test_cuda_rand_load_store_tracegen(opcode: Rv64LoadStoreOpcode, num_ops: usiz pointer_max_bits: 20, ..Default::default() }; - mem_config.addr_spaces[RV64_REGISTER_AS as usize].num_cells = 1 << 20; mem_config.addr_spaces[RV64_MEMORY_AS as usize].num_cells = 1 << 20; if [STORED, STOREW, STOREB, STOREH].contains(&opcode) { mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 20; diff --git a/extensions/riscv/circuit/src/mulh/tests.rs b/extensions/riscv/circuit/src/mulh/tests.rs index 4ad7a06fd6..98ba517782 100644 --- a/extensions/riscv/circuit/src/mulh/tests.rs +++ b/extensions/riscv/circuit/src/mulh/tests.rs @@ -9,8 +9,8 @@ use openvm_circuit::arch::{ use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, + memory::gen_distinct_register_pointers, TestBuilder, TestChipHarness, + VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, RANGE_TUPLE_CHECKER_BUS, }, Arena, ExecutionBridge, PreflightExecutor, }, @@ -162,9 +162,7 @@ fn set_and_execute>( RV64_BYTE_BITS, >(rng)); - let rs1 = gen_pointer(rng, 8); - let rs2 = gen_pointer(rng, 8); - let rd = gen_pointer(rng, 8); + let [rs1, rs2, rd] = gen_distinct_register_pointers(rng, 8); tester.write_bytes::(1, rs1, b.map(F::from_u32)); tester.write_bytes::(1, rs2, c.map(F::from_u32)); diff --git a/extensions/riscv/circuit/src/test_utils.rs b/extensions/riscv/circuit/src/test_utils.rs index 71c3966c5b..1a169a19c9 100644 --- a/extensions/riscv/circuit/src/test_utils.rs +++ b/extensions/riscv/circuit/src/test_utils.rs @@ -1,5 +1,5 @@ use openvm_circuit::arch::{ - testing::{memory::gen_pointer, TestBuilder}, + testing::{memory::gen_distinct_register_pointers, TestBuilder}, BLOCK_FE_WIDTH, }; use openvm_instructions::{instruction::Instruction, VmOpcode}; @@ -21,9 +21,8 @@ pub fn rv64_rand_write_register_or_imm( ) -> (Instruction, usize) { let rs2_is_imm = imm.is_some(); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = imm.unwrap_or_else(|| gen_pointer(rng, RV64_REGISTER_NUM_LIMBS)); - let rd = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rs1, rs2_reg, rd] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); + let rs2 = imm.unwrap_or(rs2_reg); tester.write_bytes::(1, rs1, rs1_writes.map(BabyBear::from_u8)); if !rs2_is_imm { diff --git a/extensions/sha2/circuit/build.rs b/extensions/sha2/circuit/build.rs index ea99d52b66..dda5843dd1 100644 --- a/extensions/sha2/circuit/build.rs +++ b/extensions/sha2/circuit/build.rs @@ -12,8 +12,10 @@ fn main() { .include("cuda/include") .include("../../../crates/circuits/primitives/cuda/include") .include("../../../crates/vm/cuda/include") + .include("../../riscv-adapters/cuda/include") .watch("cuda") .watch("../../../crates/circuits/primitives/cuda") + .watch("../../riscv-adapters/cuda") .watch("../../../crates/vm/cuda") .library_name("tracegen_gpu_sha2") .files_from_glob("cuda/src/*.cu"); diff --git a/extensions/sha2/circuit/cuda/include/main/columns.cuh b/extensions/sha2/circuit/cuda/include/main/columns.cuh index cb2dd37bf4..8eed33a708 100644 --- a/extensions/sha2/circuit/cuda/include/main/columns.cuh +++ b/extensions/sha2/circuit/cuda/include/main/columns.cuh @@ -36,6 +36,16 @@ template struct Sha2MainMemoryCols { MemoryReadAuxCols input_reads[V::BLOCK_READS]; MemoryReadAuxCols state_reads[V::STATE_READS]; MemoryWriteAuxCols write_aux[V::STATE_WRITES]; + // Carry for converting each base heap *byte* pointer (`input`, `state`, `dst`) to AS-native + // u16 *cell* pointer limbs. + T input_cell_carry; + T state_cell_carry; + T dst_cell_carry; + // Per-block carry for adding the cell offset `i * (SHA2_READ_SIZE / U16_CELL_SIZE)` to each + // base cell pointer (block `i`'s carry into the high cell limb). + T input_add_carry[V::BLOCK_READS]; + T state_add_carry[V::STATE_READS]; + T write_add_carry[V::STATE_WRITES]; }; template struct Sha2MainCols { @@ -114,5 +124,13 @@ template struct Sha2MainLayout { #define SHA2_MAIN_SLICE_MEM(V, ROW, FIELD) \ (ROW).slice_from(SHA2_MAIN_COL_INDEX_MEM_V(V, FIELD)) +#define SHA2_MAIN_WRITE_MEM(V, ROW, FIELD, VALUE) \ + (ROW).write(SHA2_MAIN_COL_INDEX_MEM_V(V, FIELD), VALUE) +#define SHA2_MAIN_WRITE_ARRAY_MEM(V, ROW, FIELD, VALUES) \ + (ROW).write_array( \ + SHA2_MAIN_COL_INDEX_MEM_V(V, FIELD), \ + SHA2_MAIN_COL_ARRAY_LEN_V(V, Sha2MainMemoryCols, FIELD), \ + VALUES \ + ) } // namespace sha2 diff --git a/extensions/sha2/circuit/cuda/src/sha2_main.cu b/extensions/sha2/circuit/cuda/src/sha2_main.cu index 4dd4c27cc7..1e726f805e 100644 --- a/extensions/sha2/circuit/cuda/src/sha2_main.cu +++ b/extensions/sha2/circuit/cuda/src/sha2_main.cu @@ -7,6 +7,7 @@ #include "primitives/histogram.cuh" #include "primitives/trace_access.h" #include "primitives/utils.cuh" +#include "riscv-adapters/pointer_conv.cuh" #include "system/memory/controller.cuh" #include "system/memory/offline_checker.cuh" #include @@ -15,7 +16,6 @@ using namespace riscv; using namespace sha2; -using openvm::U16_BITS; // Body shared by both the inlined (SHA-256) and outlined (SHA-512) paths. // @@ -71,20 +71,46 @@ static __device__ __forceinline__ void sha2_main_row_body( SHA2_MAIN_WRITE_ARRAY_INSTR(V, row, state_ptr_limbs, state_ptr_u16s); SHA2_MAIN_WRITE_ARRAY_INSTR(V, row, input_ptr_limbs, input_ptr_u16s); - // Range-check the high u16 of each pointer via the variable range checker. - range_checker.add_count( - ptr_bound_from_high_u16(dst_ptr_u16s[RV64_PTR_U16_LIMBS - 1], ptr_max_bits), - U16_BITS + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. Mirrors `compute_pointer_carries` in `main_chip/trace.rs`. + uint32_t read_cell_stride = SHA2_READ_SIZE / 2; + uint32_t write_cell_stride = SHA2_WRITE_SIZE / 2; + + uint32_t input_add_carry[V::BLOCK_READS]; + uint32_t state_add_carry[V::STATE_READS]; + uint32_t write_add_carry[V::STATE_WRITES]; + uint32_t input_conv_carry = compute_pointer_carries( + range_checker, + header->input_ptr, + ptr_max_bits, + V::BLOCK_READS, + read_cell_stride, + input_add_carry ); - range_checker.add_count( - ptr_bound_from_high_u16(state_ptr_u16s[RV64_PTR_U16_LIMBS - 1], ptr_max_bits), - U16_BITS + uint32_t state_conv_carry = compute_pointer_carries( + range_checker, + header->state_ptr, + ptr_max_bits, + V::STATE_READS, + read_cell_stride, + state_add_carry ); - range_checker.add_count( - ptr_bound_from_high_u16(input_ptr_u16s[RV64_PTR_U16_LIMBS - 1], ptr_max_bits), - U16_BITS + uint32_t dst_conv_carry = compute_pointer_carries( + range_checker, + header->dst_ptr, + ptr_max_bits, + V::STATE_WRITES, + write_cell_stride, + write_add_carry ); + SHA2_MAIN_WRITE_MEM(V, row, input_cell_carry, Fp(input_conv_carry)); + SHA2_MAIN_WRITE_MEM(V, row, state_cell_carry, Fp(state_conv_carry)); + SHA2_MAIN_WRITE_MEM(V, row, dst_cell_carry, Fp(dst_conv_carry)); + SHA2_MAIN_WRITE_ARRAY_MEM(V, row, input_add_carry, input_add_carry); + SHA2_MAIN_WRITE_ARRAY_MEM(V, row, state_add_carry, state_add_carry); + SHA2_MAIN_WRITE_ARRAY_MEM(V, row, write_add_carry, write_add_carry); + // Memory aux uint32_t timestamp = header->timestamp; for (int i = 0; i < static_cast(sha2::SHA2_REGISTER_READS); i++) { diff --git a/extensions/sha2/circuit/src/sha2_chips/config.rs b/extensions/sha2/circuit/src/sha2_chips/config.rs index aff63b230a..5498626c1b 100644 --- a/extensions/sha2/circuit/src/sha2_chips/config.rs +++ b/extensions/sha2/circuit/src/sha2_chips/config.rs @@ -63,9 +63,11 @@ impl Sha2Config for Sha512Config { debug_assert!(state.len() >= Sha512Config::STATE_BYTES); debug_assert!(input.len() == Sha512Config::BLOCK_BYTES); - // `state` may only be 4-byte aligned, so copy through an aligned `u64` buffer. - // We store the state as bytes in the record because the word size is - // not known at compile time (u32 for Sha256, u64 for Sha512). + // `state` is a byte sub-slice of a record and is not guaranteed to be 8-byte aligned, so we + // cannot reinterpret it in place as `&mut [u64; 8]`. Copy through an aligned buffer (using + // native-endian bytes, matching the in-place reinterpretation this replaces). We store the + // state as bytes in the record because the word size is not known at compile time (u32 for + // Sha256, u64 for Sha512). let mut state_u64s = [0u64; 8]; for (w, chunk) in state_u64s.iter_mut().zip(state.chunks_exact(8)) { *w = u64::from_ne_bytes(chunk.try_into().unwrap()); diff --git a/extensions/sha2/circuit/src/sha2_chips/main_chip/air.rs b/extensions/sha2/circuit/src/sha2_chips/main_chip/air.rs index d9917322bb..5c9820d76f 100644 --- a/extensions/sha2/circuit/src/sha2_chips/main_chip/air.rs +++ b/extensions/sha2/circuit/src/sha2_chips/main_chip/air.rs @@ -8,11 +8,11 @@ use openvm_circuit::{ SystemPort, }, }; -use openvm_circuit_primitives::{var_range::VariableRangeCheckerBus, ColumnsAir, U16_BITS}; +use openvm_circuit_primitives::{var_range::VariableRangeCheckerBus, ColumnsAir}; use openvm_instructions::riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS}; use openvm_riscv_circuit::adapters::{ - byte_ptr_to_u16_ptr, expand_to_rv64_block, ptr_bound_from_high_u16_expr, u16_limbs_to_ptr, - RV64_PTR_U16_LIMBS, + eval_add_const_u16_limbs, eval_byte_ptr_limbs_to_u16_cell_ptr_limbs, expand_to_rv64_block, + reg_byte_ptr_to_cell_ptr_limbs, }; use openvm_sha2_air::Sha2BlockHasherSubairConfig; use openvm_stark_backend::{ @@ -218,7 +218,8 @@ impl Sha2MainAir { .read( MemoryAddress::new( AB::Expr::from_u32(RV64_REGISTER_AS), - byte_ptr_to_u16_ptr::(ptr), + // Register byte pointers are small: `ptr / 2` in the low cell limb. + reg_byte_ptr_to_cell_ptr_limbs::(ptr), ), bus_payload, timestamp_pp(), @@ -227,21 +228,6 @@ impl Sha2MainAir { .eval(builder, *local.instruction.is_enabled); } - for limbs in [ - local.instruction.dst_ptr_limbs, - local.instruction.state_ptr_limbs, - local.instruction.input_ptr_limbs, - ] { - self.range_bus - .range_check( - ptr_bound_from_high_u16_expr::( - limbs[RV64_PTR_U16_LIMBS - 1], - self.ptr_max_bits, - ), - U16_BITS, - ) - .eval(builder, *local.instruction.is_enabled); - } self.execution_bridge .execute_and_increment_pc( AB::Expr::from_usize(C::OPCODE as usize + self.offset), @@ -264,19 +250,34 @@ impl Sha2MainAir { local: &Sha2ColsRef, timestamp_pp: &mut impl FnMut() -> AB::Expr, ) { - let input_ptr_limbs = std::array::from_fn(|i| local.instruction.input_ptr_limbs[i]); - let input_ptr_val = u16_limbs_to_ptr(&input_ptr_limbs); + // Cell offset (in u16 cells) between consecutive heap read blocks. + let cell_ptr_block_stride = (SHA2_READ_SIZE / 2) as u32; + + // Convert the `input` base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let input_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.instruction.input_ptr_limbs[i].into()); + let input_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + input_byte_limbs, + *local.mem.input_cell_carry, + self.ptr_max_bits, + (*local.instruction.is_enabled).into(), + ); for i in 0..C::BLOCK_READS { let chunk: [AB::Expr; BLOCK_FE_WIDTH] = std::array::from_fn(|j| local.block.message_u16s[i * BLOCK_FE_WIDTH + j].into()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + input_base_cell.clone(), + i as u32 * cell_ptr_block_stride, + local.mem.input_add_carry[i], + (*local.instruction.is_enabled).into(), + ); self.memory_bridge .read( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::( - input_ptr_val.clone() + AB::F::from_usize(i * SHA2_READ_SIZE), - ), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), chunk, timestamp_pp(), &local.mem.input_reads[i], @@ -284,19 +285,31 @@ impl Sha2MainAir { .eval(builder, *local.instruction.is_enabled); } - let state_ptr_limbs = std::array::from_fn(|i| local.instruction.state_ptr_limbs[i]); - let state_ptr_val = u16_limbs_to_ptr(&state_ptr_limbs); + // Convert the `state` base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let state_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.instruction.state_ptr_limbs[i].into()); + let state_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + state_byte_limbs, + *local.mem.state_cell_carry, + self.ptr_max_bits, + (*local.instruction.is_enabled).into(), + ); for i in 0..C::STATE_READS { let chunk: [AB::Expr; BLOCK_FE_WIDTH] = std::array::from_fn(|j| local.block.prev_state[i * BLOCK_FE_WIDTH + j].into()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + state_base_cell.clone(), + i as u32 * cell_ptr_block_stride, + local.mem.state_add_carry[i], + (*local.instruction.is_enabled).into(), + ); self.memory_bridge .read( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::( - state_ptr_val.clone() + AB::F::from_usize(i * SHA2_READ_SIZE), - ), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), chunk, timestamp_pp(), &local.mem.state_reads[i], @@ -311,19 +324,34 @@ impl Sha2MainAir { local: &Sha2ColsRef, timestamp_pp: &mut impl FnMut() -> AB::Expr, ) { - let dst_ptr_limbs = std::array::from_fn(|i| local.instruction.dst_ptr_limbs[i]); - let dst_ptr_val = u16_limbs_to_ptr(&dst_ptr_limbs); + // Cell offset (in u16 cells) between consecutive heap write blocks. + let cell_ptr_block_stride = (SHA2_WRITE_SIZE / 2) as u32; + + // Convert the `dst` base *byte* pointer to base AS-native u16 *cell* pointer limbs. + let dst_byte_limbs: [AB::Expr; 2] = + std::array::from_fn(|i| local.instruction.dst_ptr_limbs[i].into()); + let dst_base_cell = eval_byte_ptr_limbs_to_u16_cell_ptr_limbs::( + builder, + self.range_bus, + dst_byte_limbs, + *local.mem.dst_cell_carry, + self.ptr_max_bits, + (*local.instruction.is_enabled).into(), + ); for i in 0..C::STATE_WRITES { let chunk: [AB::Expr; BLOCK_FE_WIDTH] = std::array::from_fn(|j| local.block.new_state[i * BLOCK_FE_WIDTH + j].into()); + let block_cell_ptr = eval_add_const_u16_limbs::( + builder, + self.range_bus, + dst_base_cell.clone(), + i as u32 * cell_ptr_block_stride, + local.mem.write_add_carry[i], + (*local.instruction.is_enabled).into(), + ); self.memory_bridge .write( - MemoryAddress::new( - AB::Expr::from_u32(RV64_MEMORY_AS), - byte_ptr_to_u16_ptr::( - dst_ptr_val.clone() + AB::F::from_usize(i * SHA2_WRITE_SIZE), - ), - ), + MemoryAddress::new(AB::Expr::from_u32(RV64_MEMORY_AS), block_cell_ptr), chunk, timestamp_pp(), &local.mem.write_aux[i], diff --git a/extensions/sha2/circuit/src/sha2_chips/main_chip/columns.rs b/extensions/sha2/circuit/src/sha2_chips/main_chip/columns.rs index 810e2b04df..f998d47368 100644 --- a/extensions/sha2/circuit/src/sha2_chips/main_chip/columns.rs +++ b/extensions/sha2/circuit/src/sha2_chips/main_chip/columns.rs @@ -76,4 +76,14 @@ pub struct Sha2MemoryCols< pub state_reads: [MemoryReadAuxCols; STATE_READS], #[aligned_borrow] pub write_aux: [MemoryWriteAuxCols; STATE_WRITES], + /// Carry for converting each base heap *byte* pointer (`input`, `state`, `dst`) to AS-native + /// u16 *cell* pointer limbs. + pub input_cell_carry: T, + pub state_cell_carry: T, + pub dst_cell_carry: T, + /// Per-block carry for adding the cell offset `i * (SHA2_READ_SIZE / U16_CELL_SIZE)` to each + /// base cell pointer (block `i`'s carry into the high cell limb). + pub input_add_carry: [T; BLOCK_READS], + pub state_add_carry: [T; STATE_READS], + pub write_add_carry: [T; STATE_WRITES], } diff --git a/extensions/sha2/circuit/src/sha2_chips/main_chip/trace.rs b/extensions/sha2/circuit/src/sha2_chips/main_chip/trace.rs index d5b56de13a..74fe84dd0f 100644 --- a/extensions/sha2/circuit/src/sha2_chips/main_chip/trace.rs +++ b/extensions/sha2/circuit/src/sha2_chips/main_chip/trace.rs @@ -6,9 +6,9 @@ use openvm_circuit::{ }, utils::next_power_of_two_or_zero, }; -use openvm_circuit_primitives::{Chip, U16_BITS}; +use openvm_circuit_primitives::Chip; use openvm_cpu_backend::CpuBackend; -use openvm_riscv_circuit::adapters::{ptr_bound_from_ptr, ptr_to_u16_limbs}; +use openvm_riscv_circuit::adapters::{compute_pointer_carries, ptr_to_u16_limbs}; use openvm_sha2_air::{set_arrayview_from_u16_le_bytes, set_arrayview_from_u16_slice}; use openvm_stark_backend::{ p3_field::{PrimeCharacteristicRing, PrimeField32}, @@ -20,7 +20,7 @@ use openvm_stark_backend::{ use crate::{ Sha2ColsRefMut, Sha2Config, Sha2MainChip, Sha2Metadata, Sha2RecordLayout, Sha2RecordMut, - Sha2SharedRecords, SHA2_WRITE_SIZE, + Sha2SharedRecords, SHA2_READ_SIZE, SHA2_WRITE_SIZE, }; // We will allocate a new trace matrix instead of using the record arena directly, @@ -169,9 +169,44 @@ impl Sha2MainChip { ptr_to_u16_limbs(vm_record.input_ptr), ); - for ptr in [vm_record.dst_ptr, vm_record.state_ptr, vm_record.input_ptr] { - self.range_checker_chip - .add_count(ptr_bound_from_ptr(ptr, self.pointer_max_bits), U16_BITS); + // Byte -> cell pointer conversion carries and per-block cell-offset carries, plus matching + // range-check counts. `vm_record` holds stable copies of the pointer values, so computing + // and writing the carry columns here does not alias the in-place record reads above. + let read_cell_stride = (SHA2_READ_SIZE / 2) as u32; + let write_cell_stride = (SHA2_WRITE_SIZE / 2) as u32; + let (input_conv, input_add) = compute_pointer_carries( + &self.range_checker_chip, + vm_record.input_ptr, + C::BLOCK_READS, + read_cell_stride, + self.pointer_max_bits, + ); + let (state_conv, state_add) = compute_pointer_carries( + &self.range_checker_chip, + vm_record.state_ptr, + C::STATE_READS, + read_cell_stride, + self.pointer_max_bits, + ); + let (dst_conv, dst_add) = compute_pointer_carries( + &self.range_checker_chip, + vm_record.dst_ptr, + C::STATE_WRITES, + write_cell_stride, + self.pointer_max_bits, + ); + + *cols.mem.input_cell_carry = F::from_u32(input_conv); + *cols.mem.state_cell_carry = F::from_u32(state_conv); + *cols.mem.dst_cell_carry = F::from_u32(dst_conv); + for (col, &c) in cols.mem.input_add_carry.iter_mut().zip(input_add.iter()) { + *col = F::from_u32(c); + } + for (col, &c) in cols.mem.state_add_carry.iter_mut().zip(state_add.iter()) { + *col = F::from_u32(c); + } + for (col, &c) in cols.mem.write_add_carry.iter_mut().zip(dst_add.iter()) { + *col = F::from_u32(c); } // fill in the register reads aux diff --git a/extensions/sha2/circuit/src/sha2_chips/tests.rs b/extensions/sha2/circuit/src/sha2_chips/tests.rs index 4571fbcf7c..91a4707beb 100644 --- a/extensions/sha2/circuit/src/sha2_chips/tests.rs +++ b/extensions/sha2/circuit/src/sha2_chips/tests.rs @@ -5,8 +5,8 @@ use itertools::Itertools; use openvm_circuit::{ arch::{ testing::{ - memory::gen_pointer, TestBuilder, TestChipHarness, VmChipTestBuilder, - BITWISE_OP_LOOKUP_BUS, + memory::{gen_distinct_register_pointers, gen_pointer}, + TestBuilder, TestChipHarness, VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS, }, Arena, MatrixRecordArena, PreflightExecutor, }, @@ -150,9 +150,7 @@ fn set_and_execute_single_block, prev_state: Option<&[u8]>, ) { - let rd = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rd, rs1, rs2] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let dst_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); let state_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); @@ -255,9 +253,7 @@ fn set_and_execute_full_message, len: Option, ) { - let rd = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs1 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); - let rs2 = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); + let [rd, rs1, rs2] = gen_distinct_register_pointers(rng, RV64_REGISTER_NUM_LIMBS); let state_ptr = gen_pointer(rng, RV64_REGISTER_NUM_LIMBS); let dst_ptr = state_ptr;