feat: add host configurable memory cell type to MemoryConfig (#1903)

Add trait `AddressSpaceHostLayout` with concrete implementation as enum
`MemoryCellType`.

Add `AddressSpaceHostConfig` struct and now `MemoryConfig` stores a vec
of these, one per address space.
Now all handling of conversions from raw bytes to `F` is done through
this config.

Comparison:
https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/16532679208
Memory finalize time seems to have increased significantly
percentage-wise, but it is still on the order of milliseconds so I think
negligible?

---------

Co-authored-by: Golovanov399 <[email protected]>

Author: jonathanpwang

CHANGELOG.md

@@ -7,8 +7,13 @@ and this project follows a versioning principles documented in [VERSIONING.md](.
 
 ## [Unreleased]
 
+### Added
+- (Config) Added `addr_spaces` vector of `AddressSpaceHostConfig` to `MemoryConfig`.
+
 ### Changed
 - (Toolchain) Removed `step` from `Program` struct because `DEFAULT_PC_STEP = 4` is always used.
+- (Config) The `clk_max_bits` field in `MemoryConfig` has been renamed to `timestamp_max_bits`.
+- (Prover) Guest memory is stored on host with address space-specified memory layouts. In particular address space `1` through `3` are now represented in bytes instead of field elements.
 
 ## v1.3.0 (2025-07-15)
 

crates/sdk/src/config/global.rs

@@ -412,7 +412,7 @@ impl From<SdkVmConfigWithDefaultDeser> for SdkVmConfig {
         if config.native.is_none() && config.castf.is_none() {
             // There should be no need to write to native address space if Native extension and
             // CastF extension are not enabled.
-            system.config.memory_config.addr_space_sizes[NATIVE_AS as usize] = 0;
+            system.config.memory_config.addr_spaces[NATIVE_AS as usize].num_cells = 0;
         }
         Self {
             system,

crates/vm/src/arch/config.rs

@@ -2,7 +2,10 @@ use std::{fs::File, io::Write, path::Path};
 
 use derive_new::new;
 use getset::{Setters, WithSetters};
-use openvm_instructions::NATIVE_AS;
+use openvm_instructions::{
+    riscv::{RV32_IMM_AS, RV32_MEMORY_AS, RV32_REGISTER_AS},
+    NATIVE_AS,
+};
 use openvm_poseidon2_air::Poseidon2Config;
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
@@ -18,7 +21,9 @@ use crate::{
         Arena, ChipInventoryError, ExecutorInventory, ExecutorInventoryError,
     },
     system::{
-        memory::{merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, CHUNK},
+        memory::{
+            merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, CHUNK, POINTER_MAX_BITS,
+        },
         SystemChipComplex,
     },
 };
@@ -105,6 +110,10 @@ where
 
 pub const OPENVM_DEFAULT_INIT_FILE_BASENAME: &str = "openvm_init";
 pub const OPENVM_DEFAULT_INIT_FILE_NAME: &str = "openvm_init.rs";
+/// The minimum block size is 4, but RISC-V `lb` only requires alignment of 1 and `lh` only requires
+/// alignment of 2 because the instructions are implemented by doing an access of block size 4.
+const DEFAULT_U8_BLOCK_SIZE: usize = 4;
+const DEFAULT_NATIVE_BLOCK_SIZE: usize = 1;
 
 /// Trait for generating a init.rs file that contains a call to moduli_init!,
 /// complex_init!, sw_init! with the supported moduli and curves.
@@ -132,18 +141,36 @@ pub trait InitFileGenerator {
     }
 }
 
+/// Each address space in guest memory may be configured with a different type `T` to represent a
+/// memory cell in the address space. On host, the address space will be mapped to linear host
+/// memory in bytes. The type `T` must be plain old data (POD) and be safely transmutable from a
+/// fixed size array of bytes. Moreover, each type `T` must be convertible to a field element `F`.
+///
+/// We currently implement this trait on the enum [MemoryCellType], which includes all cell types
+/// that we expect to be used in the VM context.
+pub trait AddressSpaceHostLayout {
+    /// Size in bytes of the memory cell type.
+    fn size(&self) -> usize;
+
+    /// # Safety
+    /// - This function must only be called when `value` is guaranteed to be of size `self.size()`.
+    /// - Alignment of `value` must be a multiple of the alignment of `F`.
+    /// - The field type `F` must be plain old data.
+    unsafe fn to_field<F: Field>(&self, value: &[u8]) -> F;
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone, new)]
 pub struct MemoryConfig {
     /// The maximum height of the address space. This means the trie has `addr_space_height` layers
     /// for searching the address space. The allowed address spaces are those in the range `[1,
     /// 1 + 2^addr_space_height)` where it starts from 1 to not allow address space 0 in memory.
     pub addr_space_height: usize,
-    /// The number of cells in each address space. It is expected that the size of the list is
-    /// `1 << addr_space_height + 1` and the first element is 0, which means no address space.
-    pub addr_space_sizes: Vec<usize>,
+    /// It is expected that the size of the list is `(1 << addr_space_height) + 1` and the first
+    /// element is 0, which means no address space.
+    pub addr_spaces: Vec<AddressSpaceHostConfig>,
     pub pointer_max_bits: usize,
-    /// All timestamps must be in the range `[0, 2^clk_max_bits)`. Maximum allowed: 29.
-    pub clk_max_bits: usize,
+    /// All timestamps must be in the range `[0, 2^timestamp_max_bits)`. Maximum allowed: 29.
+    pub timestamp_max_bits: usize,
     /// Limb size used by the range checker
     pub decomp: usize,
     /// Maximum N AccessAdapter AIR to support.
@@ -152,19 +179,46 @@ pub struct MemoryConfig {
 
 impl Default for MemoryConfig {
     fn default() -> Self {
-        let mut addr_space_sizes = vec![0; (1 << 3) + ADDR_SPACE_OFFSET as usize];
-        addr_space_sizes[ADDR_SPACE_OFFSET as usize..=NATIVE_AS as usize].fill(1 << 29);
-        addr_space_sizes[PUBLIC_VALUES_AS as usize] = DEFAULT_MAX_NUM_PUBLIC_VALUES;
-        Self::new(3, addr_space_sizes, 29, 29, 17, 32)
+        let mut addr_spaces =
+            Self::empty_address_space_configs((1 << 3) + ADDR_SPACE_OFFSET as usize);
+        const MAX_CELLS: usize = 1 << 29;
+        addr_spaces[RV32_REGISTER_AS as usize].num_cells = 32 * size_of::<u32>();
+        addr_spaces[RV32_MEMORY_AS as usize].num_cells = MAX_CELLS;
+        addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = DEFAULT_MAX_NUM_PUBLIC_VALUES;
+        addr_spaces[NATIVE_AS as usize].num_cells = MAX_CELLS;
+        Self::new(3, addr_spaces, POINTER_MAX_BITS, 29, 17, 32)
     }
 }
 
 impl MemoryConfig {
+    pub fn empty_address_space_configs(num_addr_spaces: usize) -> Vec<AddressSpaceHostConfig> {
+        // All except address spaces 0..4 default to native 32-bit field.
+        // By default only address spaces 1..=4 have non-empty cell counts.
+        let mut addr_spaces = vec![
+            AddressSpaceHostConfig::new(
+                0,
+                DEFAULT_NATIVE_BLOCK_SIZE,
+                MemoryCellType::native32()
+            );
+            num_addr_spaces
+        ];
+        addr_spaces[RV32_IMM_AS as usize] = AddressSpaceHostConfig::new(0, 1, MemoryCellType::Null);
+        addr_spaces[RV32_REGISTER_AS as usize] =
+            AddressSpaceHostConfig::new(0, DEFAULT_U8_BLOCK_SIZE, MemoryCellType::U8);
+        addr_spaces[RV32_MEMORY_AS as usize] =
+            AddressSpaceHostConfig::new(0, DEFAULT_U8_BLOCK_SIZE, MemoryCellType::U8);
+        addr_spaces[PUBLIC_VALUES_AS as usize] =
+            AddressSpaceHostConfig::new(0, DEFAULT_U8_BLOCK_SIZE, MemoryCellType::U8);
+
+        addr_spaces
+    }
+
     /// Config for aggregation usage with only native address space.
     pub fn aggregation() -> Self {
-        let mut addr_space_sizes = vec![0; (1 << 3) + ADDR_SPACE_OFFSET as usize];
-        addr_space_sizes[NATIVE_AS as usize] = 1 << 29;
-        Self::new(3, addr_space_sizes, 29, 29, 17, 8)
+        let mut addr_spaces =
+            Self::empty_address_space_configs((1 << 3) + ADDR_SPACE_OFFSET as usize);
+        addr_spaces[NATIVE_AS as usize].num_cells = 1 << 29;
+        Self::new(3, addr_spaces, POINTER_MAX_BITS, 29, 17, 8)
     }
 }
 
@@ -209,10 +263,10 @@ impl SystemConfig {
         num_public_values: usize,
     ) -> Self {
         assert!(
-            memory_config.clk_max_bits <= 29,
+            memory_config.timestamp_max_bits <= 29,
             "Timestamp max bits must be <= 29 for LessThan to work in 31-bit field"
         );
-        memory_config.addr_space_sizes[PUBLIC_VALUES_AS as usize] = num_public_values;
+        memory_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = num_public_values;
         Self {
             max_constraint_degree,
             continuation_enabled: false,
@@ -243,7 +297,7 @@ impl SystemConfig {
 
     pub fn with_public_values(mut self, num_public_values: usize) -> Self {
         self.num_public_values = num_public_values;
-        self.memory_config.addr_space_sizes[PUBLIC_VALUES_AS as usize] = num_public_values;
+        self.memory_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = num_public_values;
         self
     }
 
@@ -316,3 +370,67 @@ impl AsMut<SystemConfig> for SystemConfig {
 
 // Default implementation uses no init file
 impl InitFileGenerator for SystemConfig {}
+
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, new)]
+pub struct AddressSpaceHostConfig {
+    /// The number of cells in each address space.
+    pub num_cells: usize,
+    /// Minimum block size for memory accesses supported. This is a property of the address space
+    /// that is determined by the ISA.
+    ///
+    /// **Note**: Block size is in terms of memory cells.
+    pub min_block_size: usize,
+    pub layout: MemoryCellType,
+}
+
+pub(crate) const MAX_CELL_BYTE_SIZE: usize = 8;
+
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
+pub enum MemoryCellType {
+    Null,
+    U8,
+    U16,
+    /// Represented in little-endian format.
+    U32,
+    /// `size` is the size in bytes of the native field type. This should not exceed 8.
+    Native {
+        size: u8,
+    },
+}
+
+impl MemoryCellType {
+    pub fn native32() -> Self {
+        Self::Native {
+            size: size_of::<u32>() as u8,
+        }
+    }
+}
+
+impl AddressSpaceHostLayout for MemoryCellType {
+    fn size(&self) -> usize {
+        match self {
+            Self::Null => 1, // to avoid divide by zero
+            Self::U8 => size_of::<u8>(),
+            Self::U16 => size_of::<u16>(),
+            Self::U32 => size_of::<u32>(),
+            Self::Native { size } => *size as usize,
+        }
+    }
+
+    /// # Safety
+    /// - This function must only be called when `value` is guaranteed to be of size `self.size()`.
+    /// - Alignment of `value` must be a multiple of the alignment of `F`.
+    /// - The field type `F` must be plain old data.
+    ///
+    /// # Panics
+    /// If the value is of integer type and overflows the field.
+    unsafe fn to_field<F: Field>(&self, value: &[u8]) -> F {
+        match self {
+            Self::Null => unreachable!(),
+            Self::U8 => F::from_canonical_u8(*value.get_unchecked(0)),
+            Self::U16 => F::from_canonical_u16(core::ptr::read(value.as_ptr() as *const u16)),
+            Self::U32 => F::from_canonical_u32(core::ptr::read(value.as_ptr() as *const u32)),
+            Self::Native { .. } => core::ptr::read(value.as_ptr() as *const F),
+        }
+    }
+}

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

@@ -25,7 +25,6 @@ pub const DEFAULT_SEGMENT_CHECK_INSNS: u64 = 1000;
 #[derive(Clone, Debug, WithSetters)]
 pub struct MeteredCtx<const PAGE_BITS: usize = DEFAULT_PAGE_BITS> {
     pub trace_heights: Vec<u32>,
-    // TODO[jpw]: should this be in Ctrl?
     pub is_trace_height_constant: Vec<bool>,
 
     pub memory_ctx: MemoryCtx<PAGE_BITS>,

crates/vm/src/arch/extensions.rs

@@ -583,7 +583,7 @@ where
     }
 
     pub fn timestamp_max_bits(&self) -> usize {
-        self.airs.config().memory_config.clk_max_bits
+        self.airs.config().memory_config.timestamp_max_bits
     }
 }
 

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

@@ -35,11 +35,10 @@ impl<F: PrimeField32> MemoryTester<F> {
         }
     }
 
-    // TODO: change interface by implementing GuestMemory trait after everything works
     pub fn read<const N: usize>(&mut self, addr_space: usize, ptr: usize) -> [F; N] {
         let memory = &mut self.memory;
         let t = memory.timestamp();
-        // TODO: hack
+        // TODO: this could be improved if we added a TracingMemory::get_f function
         let (t_prev, data) = if addr_space <= 3 {
             let (t_prev, data) = unsafe { memory.read::<u8, N, 4>(addr_space as u32, ptr as u32) };
             (t_prev, data.map(F::from_canonical_u8))
@@ -60,11 +59,10 @@ impl<F: PrimeField32> MemoryTester<F> {
         data
     }
 
-    // TODO: see read
     pub fn write<const N: usize>(&mut self, addr_space: usize, ptr: usize, data: [F; N]) {
         let memory = &mut self.memory;
         let t = memory.timestamp();
-        // TODO: hack
+        // TODO: this could be improved if we added a TracingMemory::write_f function
         let (t_prev, data_prev) = if addr_space <= 3 {
             let (t_prev, data_prev) = unsafe {
                 memory.write::<u8, N, 4>(

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

@@ -7,7 +7,7 @@ use openvm_circuit_primitives::{
         SharedVariableRangeCheckerChip, VariableRangeCheckerBus, VariableRangeCheckerChip,
     },
 };
-use openvm_instructions::{instruction::Instruction, NATIVE_AS};
+use openvm_instructions::{instruction::Instruction, riscv::RV32_REGISTER_AS, NATIVE_AS};
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
     engine::VerificationData,
@@ -272,7 +272,7 @@ impl<F: PrimeField32> VmChipTestBuilder<F> {
     }
 
     pub fn address_bits(&self) -> usize {
-        self.memory.controller.mem_config.pointer_max_bits
+        self.memory.controller.memory_config().pointer_max_bits
     }
 
     pub fn get_default_register(&mut self, increment: usize) -> usize {
@@ -336,7 +336,8 @@ impl VmChipTestBuilder<BabyBear> {
 impl<F: PrimeField32> VmChipTestBuilder<F> {
     pub fn default_persistent() -> Self {
         let mut mem_config = MemoryConfig::default();
-        mem_config.addr_space_sizes[NATIVE_AS as usize] = 0;
+        mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29;
+        mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0;
         Self::persistent(mem_config)
     }
 
@@ -411,7 +412,10 @@ impl<F: PrimeField32> VmChipTestBuilder<F> {
 impl<F: PrimeField32> Default for VmChipTestBuilder<F> {
     fn default() -> Self {
         let mut mem_config = MemoryConfig::default();
-        mem_config.addr_space_sizes[NATIVE_AS as usize] = 0;
+        // 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[RV32_REGISTER_AS as usize].num_cells = 1 << 29;
+        mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0;
         Self::volatile(mem_config)
     }
 }
@@ -487,7 +491,7 @@ where
             }
             let mem_inventory = MemoryAirInventory::new(
                 memory_controller.memory_bridge(),
-                &memory_controller.mem_config,
+                memory_controller.memory_config(),
                 range_checker.bus(),
                 is_persistent.then_some((
                     PermutationCheckBus::new(MEMORY_MERKLE_BUS),

crates/vm/src/arch/vm.rs

@@ -237,16 +237,11 @@ where
         interactions: &[usize],
     ) -> MeteredCtx {
         let system_config = self.config.as_ref();
-        let num_addr_sp = 1 + (1 << system_config.memory_config.addr_space_height);
-        let mut min_block_size = vec![1; num_addr_sp];
-        // TMP: hardcoding for now
-        // TODO[jpw]: move to mem_config
-        min_block_size[1] = 4;
-        min_block_size[2] = 4;
-        min_block_size[3] = 4;
-        let as_byte_alignment_bits = min_block_size
+        let as_byte_alignment_bits = system_config
+            .memory_config
+            .addr_spaces
             .iter()
-            .map(|&x| log2_strict_usize(x as usize) as u8)
+            .map(|addr_sp| log2_strict_usize(addr_sp.min_block_size) as u8)
             .collect();
 
         MeteredCtx::new(
@@ -543,7 +538,6 @@ where
         );
         let memory = TracingMemory::from_image(
             state.memory,
-            &system_config.memory_config,
             system_config.initial_block_size(),
             access_adapter_arena_size_bound,
         );
@@ -1201,7 +1195,7 @@ pub(super) fn create_memory_image(
     init_memory: SparseMemoryImage,
 ) -> GuestMemory {
     GuestMemory::new(AddressMap::from_sparse(
-        memory_config.addr_space_sizes.clone(),
+        memory_config.addr_spaces.clone(),
         init_memory,
     ))
 }