Merge pull request #2815 from o1-labs/o1vm/decrease-scratch-size

o1vm/mips: decrease scratch size by 40%

o1vm/src/interpreters/mips/column.rs

@@ -1,40 +1,44 @@
-use crate::interpreters::mips::{
-    witness::SCRATCH_SIZE,
-    Instruction::{self, IType, JType, RType},
-};
+use crate::interpreters::mips::Instruction::{self, IType, JType, RType};
 use kimchi_msm::{
     columns::{Column, ColumnIndexer},
     witness::Witness,
 };
 use std::ops::{Index, IndexMut};
 use strum::EnumCount;
 
-pub use super::{
-    witness::SCRATCH_SIZE_INVERSE, ITypeInstruction, JTypeInstruction, RTypeInstruction,
-};
+use super::{ITypeInstruction, JTypeInstruction, RTypeInstruction};
 
+pub(crate) const SCRATCH_SIZE_WITHOUT_KECCAK: usize = 45;
 /// The number of hashes performed so far in the block
-pub(crate) const MIPS_HASH_COUNTER_OFF: usize = 80;
+pub(crate) const MIPS_HASH_COUNTER_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK;
 /// The number of bytes of the preimage that have been read so far in this hash
-pub(crate) const MIPS_BYTE_COUNTER_OFF: usize = 81;
+pub(crate) const MIPS_BYTE_COUNTER_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 1;
 /// A flag indicating whether the preimage has been read fully or not
-pub(crate) const MIPS_END_OF_PREIMAGE_OFF: usize = 82;
+pub(crate) const MIPS_END_OF_PREIMAGE_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 2;
 /// The number of preimage bytes processed in this step
-pub(crate) const MIPS_NUM_BYTES_READ_OFF: usize = 83;
+pub(crate) const MIPS_NUM_BYTES_READ_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 3;
 /// The at most 4-byte chunk of the preimage that has been read in this step.
 /// Contains a field element of at most 4 bytes.
-pub(crate) const MIPS_PREIMAGE_CHUNK_OFF: usize = 84;
+pub(crate) const MIPS_PREIMAGE_CHUNK_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 4;
 /// The at most 4-bytes of the preimage that are currently being processed
 /// Consists of 4 field elements of at most 1 byte each.
-pub(crate) const MIPS_PREIMAGE_BYTES_OFF: usize = 85;
+pub(crate) const MIPS_PREIMAGE_BYTES_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 5;
 /// The at most 4-bytes of the length that are currently being processed
-pub(crate) const MIPS_LENGTH_BYTES_OFF: usize = 89;
+pub(crate) const MIPS_LENGTH_BYTES_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 5 + 4;
 /// Flags indicating whether at least N bytes have been processed in this step
-pub(crate) const MIPS_HAS_N_BYTES_OFF: usize = 93;
+pub(crate) const MIPS_HAS_N_BYTES_OFF: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 5 + 4 + 4;
 /// The maximum size of a chunk (4 bytes)
 pub(crate) const MIPS_CHUNK_BYTES_LEN: usize = 4;
 /// The location of the preimage key as a field element of 248bits
-pub(crate) const MIPS_PREIMAGE_KEY: usize = 97;
+pub(crate) const MIPS_PREIMAGE_KEY: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 5 + 4 + 4 + 4;
+
+// MIPS + hash_counter + byte_counter + eof + num_bytes_read + chunk + bytes
+// + length + has_n_bytes + chunk_bytes + preimage
+pub const SCRATCH_SIZE: usize = SCRATCH_SIZE_WITHOUT_KECCAK + 5 + 4 + 4 + 4 + 1;
+
+/// Number of columns used by the MIPS interpreter to keep values to be
+/// inverted.
+pub const SCRATCH_SIZE_INVERSE: usize = 12;
 
 /// The number of columns used for relation witness in the MIPS circuit
 pub const N_MIPS_REL_COLS: usize = SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2;

o1vm/src/interpreters/mips/constraints.rs

@@ -54,10 +54,10 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
 
     // Use one of the available columns. It won't create a new column every time
     // this function is called. The number of columns is defined upfront by
-    // crate::mips::witness::SCRATCH_SIZE.
+    // crate::interpreters::mips::column::SCRATCH_SIZE.
     fn alloc_scratch(&mut self) -> Self::Position {
         // All columns are implemented using a simple index, and a name is given
-        // to the index. See crate::SCRATCH_SIZE for the maximum number of
+        // to the index. See crate::interpreters::mips::column::SCRATCH_SIZE for the maximum number of
         // columns the circuit can use.
         let scratch_idx = self.scratch_state_idx;
         self.scratch_state_idx += 1;

o1vm/src/interpreters/mips/interpreter.rs

@@ -166,10 +166,10 @@ pub trait InterpreterEnv {
     /// The variables are "freed" after each step/instruction.
     /// The variable allocation can be seen as an allocation on a stack that is
     /// popped after each step execution.
-    /// At the moment, [crate::interpreters::mips::witness::SCRATCH_SIZE - 46]
+    /// At the moment, [crate::interpreters::mips::column::SCRATCH_SIZE - 46]
     /// elements can be allocated. If more temporary variables are required for
     /// an instruction, increase the value
-    /// [crate::interpreters::mips::witness::SCRATCH_SIZE]
+    /// [crate::interpreters::mips::column::SCRATCH_SIZE]
     fn alloc_scratch(&mut self) -> Self::Position;
 
     fn alloc_scratch_inverse(&mut self) -> Self::Position;

o1vm/src/interpreters/mips/tests_helpers.rs

@@ -3,7 +3,7 @@ use crate::{
     interpreters::mips::{
         interpreter::{debugging::InstructionParts, InterpreterEnv},
         registers::Registers,
-        witness::{Env as WEnv, SyscallEnv, SCRATCH_SIZE},
+        witness::{Env as WEnv, SyscallEnv},
     },
     preimage_oracle::PreImageOracleT,
 };
@@ -13,7 +13,7 @@ use std::{fs, path::PathBuf};
 // FIXME: we should parametrize the tests with different fields.
 use ark_bn254::Fr as Fp;
 
-use super::witness::SCRATCH_SIZE_INVERSE;
+use super::column::{SCRATCH_SIZE, SCRATCH_SIZE_INVERSE};
 
 const PAGE_INDEX_EXECUTABLE_MEMORY: u32 = 1;
 

o1vm/src/interpreters/mips/witness.rs

@@ -1,4 +1,4 @@
-use super::column::N_MIPS_SEL_COLS;
+use super::column::{N_MIPS_SEL_COLS, SCRATCH_SIZE, SCRATCH_SIZE_INVERSE};
 use crate::{
     cannon::{
         Hint, Meta, Page, Start, State, StepFrequency, VmConfiguration, PAGE_ADDRESS_MASK,
@@ -50,16 +50,6 @@ pub const NUM_INSTRUCTION_LOOKUP_TERMS: usize = 5;
 pub const NUM_LOOKUP_TERMS: usize =
     NUM_GLOBAL_LOOKUP_TERMS + NUM_DECODING_LOOKUP_TERMS + NUM_INSTRUCTION_LOOKUP_TERMS;
 // TODO: Delete and use a vector instead
-// FIXME: since the introduction of the scratch size inverse, the value below
-// can be decreased. It implies to change the offsets defined in [column]. At
-// the moment, it incurs an overhead we could avoid as some columns are zeroes.
-// MIPS + hash_counter + byte_counter + eof + num_bytes_read + chunk + bytes
-// + length + has_n_bytes + chunk_bytes + preimage
-pub const SCRATCH_SIZE: usize = 98;
-
-/// Number of columns used by the MIPS interpreter to keep values to be
-/// inverted.
-pub const SCRATCH_SIZE_INVERSE: usize = 12;
 
 #[derive(Clone, Default)]
 pub struct SyscallEnv {

o1vm/src/legacy/main.rs

@@ -12,10 +12,10 @@ use o1vm::{
             environment::KeccakEnv,
         },
         mips::{
-            column::{N_MIPS_COLS, N_MIPS_REL_COLS, N_MIPS_SEL_COLS},
+            column::{N_MIPS_COLS, N_MIPS_REL_COLS, N_MIPS_SEL_COLS, SCRATCH_SIZE},
             constraints as mips_constraints,
             interpreter::Instruction,
-            witness::{self as mips_witness, SCRATCH_SIZE},
+            witness::{self as mips_witness},
         },
     },
     legacy::{

o1vm/src/legacy/tests.rs

@@ -1,9 +1,8 @@
 use crate::{
     interpreters::mips::{
-        column::N_MIPS_REL_COLS,
+        column::{N_MIPS_REL_COLS, SCRATCH_SIZE},
         constraints::Env as CEnv,
         interpreter::{debugging::InstructionParts, interpret_itype},
-        witness::SCRATCH_SIZE,
         ITypeInstruction,
     },
     legacy::{

o1vm/src/pickles/column_env.rs

@@ -3,10 +3,7 @@ use ark_poly::{Evaluations, Radix2EvaluationDomain};
 use kimchi_msm::columns::Column;
 
 use crate::{
-    interpreters::mips::{
-        column::N_MIPS_SEL_COLS,
-        witness::{SCRATCH_SIZE, SCRATCH_SIZE_INVERSE},
-    },
+    interpreters::mips::column::{N_MIPS_SEL_COLS, SCRATCH_SIZE, SCRATCH_SIZE_INVERSE},
     pickles::proof::WitnessColumns,
 };
 use kimchi::circuits::{

o1vm/src/pickles/proof.rs

@@ -1,10 +1,7 @@
 use kimchi::{curve::KimchiCurve, proof::PointEvaluations};
 use poly_commitment::{ipa::OpeningProof, PolyComm};
 
-use crate::interpreters::mips::{
-    column::N_MIPS_SEL_COLS,
-    witness::{SCRATCH_SIZE, SCRATCH_SIZE_INVERSE},
-};
+use crate::interpreters::mips::column::{N_MIPS_SEL_COLS, SCRATCH_SIZE, SCRATCH_SIZE_INVERSE};
 
 pub struct WitnessColumns<G, S> {
     pub scratch: [G; SCRATCH_SIZE],

o1vm/src/pickles/tests.rs

@@ -1,7 +1,7 @@
 use std::time::Instant;
 
 use super::{
-    super::interpreters::mips::witness::SCRATCH_SIZE,
+    super::interpreters::mips::column::SCRATCH_SIZE,
     proof::{ProofInputs, WitnessColumns},
     prover::prove,
 };