feat: remove memory access adapters: CPU Boundary AIR, initialize/finalize memory

crates/vm/src/arch/config.rs

@@ -26,9 +26,7 @@ use crate::{
         Arena, ChipInventoryError, ExecutorInventory, ExecutorInventoryError,
     },
     system::{
-        memory::{
-            merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, CHUNK, POINTER_MAX_BITS,
-        },
+        memory::{merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, POINTER_MAX_BITS},
         SystemChipComplex,
     },
 };
@@ -123,6 +121,11 @@ pub const OPENVM_DEFAULT_INIT_FILE_NAME: &str = "openvm_init.rs";
 const DEFAULT_U8_BLOCK_SIZE: usize = 4;
 const DEFAULT_NATIVE_BLOCK_SIZE: usize = 1;
 
+/// The constant block size used for memory accesses when access adapters are disabled.
+/// All memory accesses for address spaces 1-3 must use this block size.
+/// This is also the block size used by the Boundary AIR for memory bus interactions.
+pub const CONST_BLOCK_SIZE: usize = DEFAULT_U8_BLOCK_SIZE;
+
 /// Trait for generating a init.rs file that contains a call to moduli_init!,
 /// complex_init!, sw_init! with the supported moduli and curves.
 /// Should be implemented by all VM config structs.
@@ -183,6 +186,10 @@ pub struct MemoryConfig {
     pub decomp: usize,
     /// Maximum N AccessAdapter AIR to support.
     pub max_access_adapter_n: usize,
+    /// Whether access adapters are enabled. When disabled, all memory accesses must be of the
+    /// standard block size (ie, 4 for address spaces 1-3).
+    #[new(value = "true")]
+    pub access_adapters_enabled: bool,
 }
 
 impl Default for MemoryConfig {
@@ -194,7 +201,15 @@ impl Default for MemoryConfig {
         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)
+        Self {
+            addr_space_height: 3,
+            addr_spaces,
+            pointer_max_bits: POINTER_MAX_BITS,
+            timestamp_max_bits: 29,
+            decomp: 17,
+            max_access_adapter_n: 32,
+            access_adapters_enabled: true,
+        }
     }
 }
 
@@ -245,6 +260,36 @@ impl MemoryConfig {
             .map(|addr_sp| log2_strict_usize(addr_sp.min_block_size) as u8)
             .collect()
     }
+
+    /// Returns true if the Native address space (AS 4) is used
+    /// Native AS is considered "used" if it has any allocated cells
+    pub fn is_native_as_used(&self) -> bool {
+        self.addr_spaces
+            .get(NATIVE_AS as usize)
+            .is_some_and(|config| config.num_cells > 0)
+    }
+
+    /// Disables access adapters. When disabled, all memory accesses for address spaces 1-3
+    /// must use the constant block size (4). Access adapters will only be used for
+    /// address space 4 (Native) if it is enabled.
+    pub fn without_access_adapters(mut self) -> Self {
+        self.access_adapters_enabled = false;
+        self
+    }
+
+    /// Enables access adapters. This is the default behavior
+    pub fn with_access_adapters(mut self) -> Self {
+        self.access_adapters_enabled = true;
+        self
+    }
+
+    /// Automatically sets `access_adapters_enabled` based on whether Native AS is used.
+    /// If Native AS is not used, access adapters are disabled since all other address spaces
+    /// use a fixed block size of 4
+    pub fn with_auto_access_adapters(mut self) -> Self {
+        self.access_adapters_enabled = self.is_native_as_used();
+        self
+    }
 }
 
 /// System-level configuration for the virtual machine. Contains all configuration parameters that
@@ -375,15 +420,42 @@ impl SystemConfig {
             + num_memory_airs(
                 self.continuation_enabled,
                 self.memory_config.max_access_adapter_n,
+                self.memory_config.access_adapters_enabled,
             )
     }
 
     pub fn initial_block_size(&self) -> usize {
         match self.continuation_enabled {
-            true => CHUNK,
+            true => CONST_BLOCK_SIZE,
             false => 1,
         }
     }
+
+    /// Disables access adapters. When disabled, all memory accesses for address spaces 1-3
+    /// must use the constant block size (4)
+    pub fn without_access_adapters(mut self) -> Self {
+        self.memory_config.access_adapters_enabled = false;
+        self
+    }
+
+    /// Enables access adapters. This is the default behavior.
+    pub fn with_access_adapters(mut self) -> Self {
+        self.memory_config.access_adapters_enabled = true;
+        self
+    }
+
+    /// Automatically sets `access_adapters_enabled` based on whether Native AS is used.
+    /// If Native AS is not used, access adapters are disabled since all other address spaces
+    /// use a fixed block size of 4.
+    pub fn with_auto_access_adapters(mut self) -> Self {
+        self.memory_config = self.memory_config.with_auto_access_adapters();
+        self
+    }
+
+    /// Returns true if access adapters are enabled.
+    pub fn access_adapters_enabled(&self) -> bool {
+        self.memory_config.access_adapters_enabled
+    }
 }
 
 impl Default for SystemConfig {

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

@@ -64,11 +64,13 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
                 air_names[merkle_tree_index]
             );
         }
-        debug_assert!(
-            air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"),
-            "air_name={}",
-            air_names[memory_ctx.adapter_offset]
-        );
+        if memory_ctx.access_adapters_enabled {
+            debug_assert!(
+                air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"),
+                "air_name={}",
+                air_names[memory_ctx.adapter_offset]
+            );
+        }
 
         let segmentation_ctx =
             SegmentationCtx::new(air_names, widths, interactions, config.segmentation_limits);

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

@@ -105,6 +105,7 @@ pub struct MemoryCtx<const PAGE_BITS: usize> {
     pub boundary_idx: usize,
     pub merkle_tree_index: Option<usize>,
     pub adapter_offset: usize,
+    pub access_adapters_enabled: bool,
     continuations_enabled: bool,
     chunk: u32,
     chunk_bits: u32,
@@ -128,6 +129,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
             boundary_idx: config.memory_boundary_air_id(),
             merkle_tree_index: config.memory_merkle_air_id(),
             adapter_offset: config.access_adapter_air_id_offset(),
+            access_adapters_enabled: config.memory_config.access_adapters_enabled,
             chunk,
             chunk_bits,
             memory_dimensions,
@@ -210,6 +212,10 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
         size_bits: u32,
         num: u32,
     ) {
+        if !self.access_adapters_enabled {
+            return;
+        }
+
         debug_assert!((address_space as usize) < self.min_block_size_bits.len());
 
         // SAFETY: address_space passed is usually a hardcoded constant or derived from an

crates/vm/src/arch/execution_mode/metered_cost.rs

@@ -18,13 +18,15 @@ pub const DEFAULT_MAX_COST: u64 = DEFAULT_MAX_SEGMENTS * DEFAULT_SEGMENT_MAX_CEL
 pub struct AccessAdapterCtx {
     min_block_size_bits: Vec<u8>,
     idx_offset: usize,
+    enabled: bool,
 }
 
 impl AccessAdapterCtx {
     pub fn new(config: &SystemConfig) -> Self {
         Self {
             min_block_size_bits: config.memory_config.min_block_size_bits(),
             idx_offset: config.access_adapter_air_id_offset(),
+            enabled: config.memory_config.access_adapters_enabled,
         }
     }
 
@@ -36,6 +38,10 @@ impl AccessAdapterCtx {
         size_bits: u32,
         widths: &[usize],
     ) {
+        if !self.enabled {
+            return;
+        }
+
         debug_assert!((address_space as usize) < self.min_block_size_bits.len());
 
         // SAFETY: address_space passed is usually a hardcoded constant or derived from an

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

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use itertools::zip_eq;
+use openvm_circuit::arch::CONST_BLOCK_SIZE;
 use openvm_circuit_primitives::var_range::{
     SharedVariableRangeCheckerChip, VariableRangeCheckerBus, VariableRangeCheckerChip,
 };
@@ -48,7 +49,7 @@ use crate::{
             adapter::records::arena_size_bound,
             offline_checker::{MemoryBridge, MemoryBus},
             online::TracingMemory,
-            MemoryAirInventory, MemoryController, SharedMemoryHelper, CHUNK,
+            MemoryAirInventory, MemoryController, SharedMemoryHelper,
         },
         poseidon2::Poseidon2PeripheryChip,
         program::ProgramBus,
@@ -323,6 +324,7 @@ impl<F: PrimeField32> VmChipTestBuilder<F> {
         let mut mem_config = MemoryConfig::default();
         mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29;
         mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0;
+        // TODO: Check if need to revert to volatile memory, after access adapters are removed
         Self::persistent(mem_config)
     }
 
@@ -347,7 +349,7 @@ impl<F: PrimeField32> VmChipTestBuilder<F> {
 
     pub fn persistent(mem_config: MemoryConfig) -> Self {
         setup_tracing_with_log_level(Level::INFO);
-        let (range_checker, memory) = Self::range_checker_and_memory(&mem_config, CHUNK);
+        let (range_checker, memory) = Self::range_checker_and_memory(&mem_config, CONST_BLOCK_SIZE);
         let hasher_chip = Arc::new(Poseidon2PeripheryChip::new(
             vm_poseidon2_config(),
             POSEIDON2_DIRECT_BUS,
@@ -403,7 +405,7 @@ impl<F: PrimeField32> Default for VmChipTestBuilder<F> {
         // 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)
+        Self::persistent(mem_config)
     }
 }
 

crates/vm/src/arch/vm.rs

@@ -622,7 +622,13 @@ where
         let system_config: &SystemConfig = self.config().as_ref();
         let adapter_offset = system_config.access_adapter_air_id_offset();
         // ATTENTION: this must agree with `num_memory_airs`
-        let num_adapters = log2_strict_usize(system_config.memory_config.max_access_adapter_n);
+
+        let num_adapters = if system_config.memory_config.access_adapters_enabled {
+            log2_strict_usize(system_config.memory_config.max_access_adapter_n)
+        } else {
+            0
+        };
+
         assert_eq!(adapter_offset + num_adapters, system_config.num_airs());
         let access_adapter_arena_size_bound = records::arena_size_bound(
             &trace_heights[adapter_offset..adapter_offset + num_adapters],

crates/vm/src/system/memory/adapter/mod.rs

@@ -58,21 +58,26 @@ impl<F: Clone + Send + Sync> AccessAdapterInventory<F> {
         memory_bus: MemoryBus,
         memory_config: MemoryConfig,
     ) -> Self {
-        let rc = range_checker;
-        let mb = memory_bus;
-        let tmb = memory_config.timestamp_max_bits;
-        let maan = memory_config.max_access_adapter_n;
-        assert!(matches!(maan, 2 | 4 | 8 | 16 | 32));
-        let chips: Vec<_> = [
-            Self::create_access_adapter_chip::<2>(rc.clone(), mb, tmb, maan),
-            Self::create_access_adapter_chip::<4>(rc.clone(), mb, tmb, maan),
-            Self::create_access_adapter_chip::<8>(rc.clone(), mb, tmb, maan),
-            Self::create_access_adapter_chip::<16>(rc.clone(), mb, tmb, maan),
-            Self::create_access_adapter_chip::<32>(rc.clone(), mb, tmb, maan),
-        ]
-        .into_iter()
-        .flatten()
-        .collect();
+        // Only create adapter chips if access adapters are enabled
+        let chips: Vec<_> = if memory_config.access_adapters_enabled {
+            let rc = range_checker;
+            let mb = memory_bus;
+            let tmb = memory_config.timestamp_max_bits;
+            let maan = memory_config.max_access_adapter_n;
+            assert!(matches!(maan, 2 | 4 | 8 | 16 | 32));
+            [
+                Self::create_access_adapter_chip::<2>(rc.clone(), mb, tmb, maan),
+                Self::create_access_adapter_chip::<4>(rc.clone(), mb, tmb, maan),
+                Self::create_access_adapter_chip::<8>(rc.clone(), mb, tmb, maan),
+                Self::create_access_adapter_chip::<16>(rc.clone(), mb, tmb, maan),
+                Self::create_access_adapter_chip::<32>(rc.clone(), mb, tmb, maan),
+            ]
+            .into_iter()
+            .flatten()
+            .collect()
+        } else {
+            Vec::new()
+        };
         Self {
             memory_config,
             chips,

crates/vm/src/system/memory/controller/mod.rs

@@ -14,7 +14,6 @@ use openvm_stark_backend::{
     interaction::PermutationCheckBus,
     p3_commit::PolynomialSpace,
     p3_field::{Field, PrimeField32},
-    p3_maybe_rayon::prelude::{IntoParallelIterator, ParallelIterator},
     p3_util::{log2_ceil_usize, log2_strict_usize},
     prover::{cpu::CpuBackend, types::AirProvingContext},
     Chip,
@@ -24,7 +23,7 @@ use serde::{Deserialize, Serialize};
 use self::interface::MemoryInterface;
 use super::{volatile::VolatileBoundaryChip, AddressMap};
 use crate::{
-    arch::{DenseRecordArena, MemoryConfig, ADDR_SPACE_OFFSET},
+    arch::{DenseRecordArena, MemoryConfig, ADDR_SPACE_OFFSET, CONST_BLOCK_SIZE},
     system::{
         memory::{
             adapter::AccessAdapterInventory,
@@ -291,10 +290,30 @@ impl<F: PrimeField32> MemoryController<F> {
             ) => {
                 let hasher = self.hasher_chip.as_ref().unwrap();
                 boundary_chip.finalize(initial_memory, &final_memory, hasher.as_ref());
-                let final_memory_values = final_memory
-                    .into_par_iter()
-                    .map(|(key, value)| (key, value.values))
-                    .collect();
+
+                // Rechunk CONST_BLOCK_SIZE blocks into CHUNK-sized blocks for merkle_chip
+                // Note: Equipartition key is (addr_space, ptr) where ptr is the starting pointer
+                let final_memory_values: Equipartition<F, CHUNK> = {
+                    use std::collections::BTreeMap;
+                    let mut chunk_map: BTreeMap<(u32, u32), [F; CHUNK]> = BTreeMap::new();
+                    for ((addr_space, ptr), ts_values) in final_memory.into_iter() {
+                        // Align to CHUNK boundary to get the chunk's starting pointer
+                        let chunk_ptr = (ptr / CHUNK as u32) * CHUNK as u32;
+                        let block_idx_in_chunk =
+                            ((ptr % CHUNK as u32) / CONST_BLOCK_SIZE as u32) as usize;
+                        let entry = chunk_map.entry((addr_space, chunk_ptr)).or_insert_with(|| {
+                            // Initialize with values from initial memory
+                            std::array::from_fn(|i| unsafe {
+                                initial_memory.get_f::<F>(addr_space, chunk_ptr + i as u32)
+                            })
+                        });
+                        // Copy values for this block
+                        for (i, val) in ts_values.values.into_iter().enumerate() {
+                            entry[block_idx_in_chunk * CONST_BLOCK_SIZE + i] = val;
+                        }
+                    }
+                    chunk_map
+                };
                 merkle_chip.finalize(initial_memory, &final_memory_values, hasher.as_ref());
             }
             _ => panic!("TouchedMemory incorrect type"),