Perf/registers on main

prover/src/lib.rs

@@ -41,8 +41,8 @@ use crate::tables::types::BusId;
 use crate::test_utils::{
     E, F, VmAir, create_bitwise_air, create_branch_air, create_commit_air, create_cpu_air,
     create_decode_air, create_dvrm_air, create_halt_air, create_load_air, create_lt_air,
-    create_memw_air, create_memw_aligned_air, create_memw_register_air, create_mul_air,
-    create_page_air, create_register_air, create_shift_air,
+    create_memw_air, create_memw_aligned_air, create_mul_air, create_page_air, create_register_air,
+    create_register_reload_air, create_shift_air,
 };
 
 use stark::proof::options::{GoldilocksCubicProofOptions, ProofOptions};
@@ -73,7 +73,6 @@ pub struct TableCounts {
     pub dvrm: usize,
     pub shift: usize,
     pub branch: usize,
-    pub memw_register: usize,
 }
 
 impl TableCounts {
@@ -92,7 +91,6 @@ impl TableCounts {
             + self.dvrm
             + self.shift
             + self.branch
-            + self.memw_register
     }
 
     /// Validate that all required tables have at least one chunk.
@@ -110,7 +108,6 @@ impl TableCounts {
             ("dvrm", self.dvrm),
             ("shift", self.shift),
             ("branch", self.branch),
-            ("memw_register", self.memw_register),
         ];
         for (name, count) in checks {
             if count == 0 {
@@ -197,8 +194,8 @@ pub(crate) struct VmAirs {
     pub halt: VmAir,
     pub commit: VmAir,
     pub register: VmAir,
+    pub register_reload: VmAir,
     pub pages: Vec<VmAir>,
-    pub memw_registers: Vec<VmAir>,
 }
 
 impl VmAirs {
@@ -210,6 +207,7 @@ impl VmAirs {
             (&self.halt, &mut traces.halt, &()),
             (&self.commit, &mut traces.commit, &()),
             (&self.register, &mut traces.register, &()),
+            (&self.register_reload, &mut traces.register_reload, &()),
         ];
 
         for (air, trace) in self.cpus.iter().zip(traces.cpus.iter_mut()) {
@@ -246,13 +244,6 @@ impl VmAirs {
         for (air, trace) in self.pages.iter().zip(traces.pages.iter_mut()) {
             pairs.push((air, trace, &()));
         }
-        for (air, trace) in self
-            .memw_registers
-            .iter()
-            .zip(traces.memw_registers.iter_mut())
-        {
-            pairs.push((air, trace, &()));
-        }
 
         pairs
     }
@@ -265,6 +256,7 @@ impl VmAirs {
             &self.halt,
             &self.commit,
             &self.register,
+            &self.register_reload,
         ];
 
         for air in &self.cpus {
@@ -297,9 +289,6 @@ impl VmAirs {
         for air in &self.pages {
             refs.push(air);
         }
-        for air in &self.memw_registers {
-            refs.push(air);
-        }
 
         refs
     }
@@ -363,6 +352,7 @@ impl VmAirs {
             register::preprocessed_commitment(proof_options, elf.entry_point),
             register::NUM_PREPROCESSED_COLS,
         );
+        let register_reload = create_register_reload_air(proof_options);
         let pages: Vec<_> = page_configs
             .iter()
             .map(|config| {
@@ -372,10 +362,6 @@ impl VmAirs {
                 )
             })
             .collect();
-        let memw_registers: Vec<_> = (0..table_counts.memw_register)
-            .map(|i| create_memw_register_air(proof_options).with_name(&format!("MEMW_R[{}]", i)))
-            .collect();
-
         #[cfg(feature = "debug-checks")]
         debug_report::print_bus_legend();
 
@@ -394,8 +380,8 @@ impl VmAirs {
             halt,
             commit,
             register,
+            register_reload,
             pages,
-            memw_registers,
         }
     }
 }
@@ -594,11 +580,11 @@ pub fn verify_with_options(
         Traces::page_configs_from_elf_and_runtime(&program, &vm_proof.runtime_page_ranges);
 
     // Cross-check: table_counts must match the number of sub-proofs.
-    // Fixed tables (bitwise, decode, halt, commit, register) = 5, plus page tables.
-    let expected_proof_count = vm_proof.table_counts.total() + 5 + page_configs.len();
+    // Fixed tables (bitwise, decode, halt, commit, register, register_reload) = 6, plus page tables.
+    let expected_proof_count = vm_proof.table_counts.total() + 6 + page_configs.len();
     if expected_proof_count != vm_proof.proof.proofs.len() {
         return Err(Error::InvalidTableCounts(format!(
-            "table_counts total ({}) + 5 fixed + {} pages = {}, but proof contains {} sub-proofs",
+            "table_counts total ({}) + 6 fixed + {} pages = {}, but proof contains {} sub-proofs",
             vm_proof.table_counts.total(),
             page_configs.len(),
             expected_proof_count,