chore: allow optional proof shape in core record and program

crates/core/executor/src/lib.rs

@@ -33,6 +33,7 @@ pub mod programs;
 mod record;
 mod register;
 mod report;
+mod shape;
 mod state;
 pub mod subproof;
 pub mod syscalls;
@@ -47,4 +48,5 @@ pub use program::*;
 pub use record::*;
 pub use register::*;
 pub use report::*;
+pub use shape::*;
 pub use state::*;

crates/core/executor/src/program.rs

@@ -4,11 +4,12 @@ use std::{collections::BTreeMap, fs::File, io::Read};
 
 use p3_field::Field;
 use serde::{Deserialize, Serialize};
-use sp1_stark::air::MachineProgram;
+use sp1_stark::air::{MachineAir, MachineProgram};
 
 use crate::{
     disassembler::{transpile, Elf},
     instruction::Instruction,
+    CoreShape,
 };
 
 /// A program that can be executed by the SP1 zkVM.
@@ -25,13 +26,21 @@ pub struct Program {
     pub pc_base: u32,
     /// The initial memory image, useful for global constants.
     pub memory_image: BTreeMap<u32, u32>,
+    /// The shape for the preprocessed tables.
+    pub preprocessed_shape: Option<CoreShape>,
 }
 
 impl Program {
     /// Create a new [Program].
     #[must_use]
     pub const fn new(instructions: Vec<Instruction>, pc_start: u32, pc_base: u32) -> Self {
-        Self { instructions, pc_start, pc_base, memory_image: BTreeMap::new() }
+        Self {
+            instructions,
+            pc_start,
+            pc_base,
+            memory_image: BTreeMap::new(),
+            preprocessed_shape: None,
+        }
     }
 
     /// Disassemble a RV32IM ELF to a program that be executed by the VM.
@@ -52,6 +61,7 @@ impl Program {
             pc_start: elf.pc_start,
             pc_base: elf.pc_base,
             memory_image: elf.memory_image,
+            preprocessed_shape: None,
         })
     }
 
@@ -65,6 +75,19 @@ impl Program {
         File::open(path)?.read_to_end(&mut elf_code)?;
         Program::from(&elf_code)
     }
+
+    /// Custom logic for padding the trace to a power of two according to the proof shape.
+    pub fn fixed_log2_rows<F: Field, A: MachineAir<F>>(&self, air: &A) -> Option<usize> {
+        self.preprocessed_shape
+            .as_ref()
+            .map(|shape| {
+                shape
+                    .shape
+                    .get(&air.name())
+                    .unwrap_or_else(|| panic!("Chip {} not found in specified shape", air.name()))
+            })
+            .copied()
+    }
 }
 
 impl<F: Field> MachineProgram<F> for Program {

crates/core/executor/src/record.rs

@@ -1,18 +1,24 @@
 use hashbrown::HashMap;
 use itertools::{EitherOrBoth, Itertools};
-use p3_field::AbstractField;
-use sp1_stark::{air::PublicValues, MachineRecord, SP1CoreOpts, SplitOpts};
+use p3_field::{AbstractField, PrimeField};
+use sp1_stark::{
+    air::{MachineAir, PublicValues},
+    MachineRecord, SP1CoreOpts, SplitOpts,
+};
 use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
 
 use super::{program::Program, Opcode};
-use crate::events::{
-    add_sharded_byte_lookup_events, AluEvent, ByteLookupEvent, ByteRecord, CpuEvent,
-    EdDecompressEvent, EllipticCurveAddEvent, EllipticCurveDecompressEvent,
-    EllipticCurveDoubleEvent, Fp2AddSubEvent, Fp2MulEvent, FpOpEvent, KeccakPermuteEvent, LookupId,
-    MemoryInitializeFinalizeEvent, MemoryRecordEnum, ShaCompressEvent, ShaExtendEvent,
-    Uint256MulEvent,
+use crate::{
+    events::{
+        add_sharded_byte_lookup_events, AluEvent, ByteLookupEvent, ByteRecord, CpuEvent,
+        EdDecompressEvent, EllipticCurveAddEvent, EllipticCurveDecompressEvent,
+        EllipticCurveDoubleEvent, Fp2AddSubEvent, Fp2MulEvent, FpOpEvent, KeccakPermuteEvent,
+        LookupId, MemoryInitializeFinalizeEvent, MemoryRecordEnum, ShaCompressEvent,
+        ShaExtendEvent, Uint256MulEvent,
+    },
+    shape::CoreShape,
 };
 
 /// A record of the execution of a program.
@@ -90,6 +96,8 @@ pub struct ExecutionRecord {
     pub public_values: PublicValues<u32, u32>,
     /// The nonce lookup.
     pub nonce_lookup: HashMap<LookupId, u32>,
+    /// The shape of the proof.
+    pub shape: Option<CoreShape>,
 }
 
 impl ExecutionRecord {
@@ -273,6 +281,20 @@ impl ExecutionRecord {
 
         shards
     }
+
+    /// Return the number of rows needed for a chip, according to the proof shape specified in the
+    /// struct.
+    pub fn fixed_log2_rows<F: PrimeField, A: MachineAir<F>>(&self, air: &A) -> Option<usize> {
+        self.shape
+            .as_ref()
+            .map(|shape| {
+                shape
+                    .shape
+                    .get(&air.name())
+                    .unwrap_or_else(|| panic!("Chip {} not found in specified shape", air.name()))
+            })
+            .copied()
+    }
 }
 
 /// A memory access record.

crates/core/executor/src/shape.rs

@@ -0,0 +1,11 @@
+use hashbrown::HashMap;
+use serde::{Deserialize, Serialize};
+
+/// The shape of a core proof.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct CoreShape {
+    /// The id of the shape. Used for enumeration of the possible proof shapes.
+    pub id: usize,
+    /// The shape of the proof. Keys are the chip names and values are the log-heights of the chips.
+    pub shape: HashMap<String, usize>,
+}

crates/core/machine/src/alu/add_sub/mod.rs

@@ -19,7 +19,7 @@ use sp1_stark::{
     Word,
 };
 
-use crate::{operations::AddOperation, utils::pad_to_power_of_two};
+use crate::{operations::AddOperation, utils::pad_rows_fixed};
 
 /// The number of main trace columns for `AddSubChip`.
 pub const NUM_ADD_SUB_COLS: usize = size_of::<AddSubCols<u8>>();
@@ -105,13 +105,15 @@ impl<F: PrimeField> MachineAir<F> for AddSubChip {
             rows.extend(row_batch);
         }
 
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_ADD_SUB_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
         // Convert the trace to a row major matrix.
         let mut trace =
             RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_ADD_SUB_COLS);
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_ADD_SUB_COLS, F>(&mut trace.values);
-
         // Write the nonces to the trace.
         for i in 0..trace.height() {
             let cols: &mut AddSubCols<F> =

crates/core/machine/src/alu/bitwise/mod.rs

@@ -19,7 +19,7 @@ use sp1_stark::{
     Word,
 };
 
-use crate::utils::pad_to_power_of_two;
+use crate::utils::pad_rows_fixed;
 
 /// The number of main trace columns for `BitwiseChip`.
 pub const NUM_BITWISE_COLS: usize = size_of::<BitwiseCols<u8>>();
@@ -74,7 +74,7 @@ impl<F: PrimeField> MachineAir<F> for BitwiseChip {
         input: &ExecutionRecord,
         _: &mut ExecutionRecord,
     ) -> RowMajorMatrix<F> {
-        let rows = input
+        let mut rows = input
             .bitwise_events
             .par_iter()
             .map(|event| {
@@ -86,13 +86,17 @@ impl<F: PrimeField> MachineAir<F> for BitwiseChip {
             })
             .collect::<Vec<_>>();
 
+        // Pad the trace to a power of two.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_BITWISE_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace =
             RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_BITWISE_COLS);
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_BITWISE_COLS, F>(&mut trace.values);
-
         for i in 0..trace.height() {
             let cols: &mut BitwiseCols<F> =
                 trace.values[i * NUM_BITWISE_COLS..(i + 1) * NUM_BITWISE_COLS].borrow_mut();

crates/core/machine/src/alu/divrem/mod.rs

@@ -83,7 +83,7 @@ use crate::{
     air::SP1CoreAirBuilder,
     alu::divrem::utils::{get_msb, get_quotient_and_remainder, is_signed_operation},
     operations::{IsEqualWordOperation, IsZeroWordOperation},
-    utils::pad_to_power_of_two,
+    utils::pad_rows_fixed,
 };
 
 /// The number of main trace columns for `DivRemChip`.
@@ -496,13 +496,17 @@ impl<F: PrimeField> MachineAir<F> for DivRemChip {
             rows.push(row);
         }
 
+        // Pad the trace to a power of two depending on the proof shape in `input`.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_DIVREM_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace =
             RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_DIVREM_COLS);
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_DIVREM_COLS, F>(&mut trace.values);
-
         // Create the template for the padded rows. These are fake rows that don't fail on some
         // sanity checks.
         let padded_row_template = {

crates/core/machine/src/alu/lt/mod.rs

@@ -19,7 +19,7 @@ use sp1_stark::{
     Word,
 };
 
-use crate::utils::pad_to_power_of_two;
+use crate::utils::pad_rows_fixed;
 
 /// The number of main trace columns for `LtChip`.
 pub const NUM_LT_COLS: usize = size_of::<LtCols<u8>>();
@@ -110,7 +110,7 @@ impl<F: PrimeField32> MachineAir<F> for LtChip {
         _: &mut ExecutionRecord,
     ) -> RowMajorMatrix<F> {
         // Generate the trace rows for each event.
-        let rows = input
+        let mut rows = input
             .lt_events
             .par_iter()
             .map(|event| {
@@ -123,13 +123,17 @@ impl<F: PrimeField32> MachineAir<F> for LtChip {
             })
             .collect::<Vec<_>>();
 
+        // Pad the trace to a power of two depending on the proof shape in `input`.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_LT_COLS],
+            input.fixed_log2_rows::<F, Self>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace =
             RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_LT_COLS);
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_LT_COLS, F>(&mut trace.values);
-
         // Write the nonces to the trace.
         for i in 0..trace.height() {
             let cols: &mut LtCols<F> =

crates/core/machine/src/alu/mul/mod.rs

@@ -47,7 +47,7 @@ use sp1_derive::AlignedBorrow;
 use sp1_primitives::consts::WORD_SIZE;
 use sp1_stark::{air::MachineAir, MachineRecord, Word};
 
-use crate::{air::SP1CoreAirBuilder, alu::mul::utils::get_msb, utils::pad_to_power_of_two};
+use crate::{air::SP1CoreAirBuilder, alu::mul::utils::get_msb, utils::pad_rows_fixed};
 
 /// The number of main trace columns for `MulChip`.
 pub const NUM_MUL_COLS: usize = size_of::<MulCols<u8>>();
@@ -267,13 +267,17 @@ impl<F: PrimeField> MachineAir<F> for MulChip {
             output.append(&mut row_and_record.1);
         }
 
+        // Pad the trace to a power of two depending on the proof shape in `input`.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_MUL_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace =
             RowMajorMatrix::new(rows.into_iter().flatten().collect::<Vec<_>>(), NUM_MUL_COLS);
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_MUL_COLS, F>(&mut trace.values);
-
         // Write the nonces to the trace.
         for i in 0..trace.height() {
             let cols: &mut MulCols<F> =

crates/core/machine/src/alu/sll/mod.rs

@@ -49,7 +49,7 @@ use sp1_derive::AlignedBorrow;
 use sp1_primitives::consts::WORD_SIZE;
 use sp1_stark::{air::MachineAir, Word};
 
-use crate::{air::SP1CoreAirBuilder, utils::pad_to_power_of_two};
+use crate::{air::SP1CoreAirBuilder, utils::pad_rows_fixed};
 
 /// The number of main trace columns for `ShiftLeft`.
 pub const NUM_SHIFT_LEFT_COLS: usize = size_of::<ShiftLeftCols<u8>>();
@@ -129,15 +129,19 @@ impl<F: PrimeField> MachineAir<F> for ShiftLeft {
             rows.push(row);
         }
 
+        // Pad the trace to a power of two depending on the proof shape in `input`.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_SHIFT_LEFT_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace = RowMajorMatrix::new(
             rows.into_iter().flatten().collect::<Vec<_>>(),
             NUM_SHIFT_LEFT_COLS,
         );
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_SHIFT_LEFT_COLS, F>(&mut trace.values);
-
         // Create the template for the padded rows. These are fake rows that don't fail on some
         // sanity checks.
         let padded_row_template = {

crates/core/machine/src/alu/sr/mod.rs

@@ -65,7 +65,7 @@ use crate::{
     air::SP1CoreAirBuilder,
     alu::sr::utils::{nb_bits_to_shift, nb_bytes_to_shift},
     bytes::utils::shr_carry,
-    utils::pad_to_power_of_two,
+    utils::pad_rows_fixed,
 };
 
 /// The number of main trace columns for `ShiftRightChip`.
@@ -163,15 +163,19 @@ impl<F: PrimeField> MachineAir<F> for ShiftRightChip {
             rows.push(row);
         }
 
+        // Pad the trace to a power of two depending on the proof shape in `input`.
+        pad_rows_fixed(
+            &mut rows,
+            || [F::zero(); NUM_SHIFT_RIGHT_COLS],
+            input.fixed_log2_rows::<F, _>(self),
+        );
+
         // Convert the trace to a row major matrix.
         let mut trace = RowMajorMatrix::new(
             rows.into_iter().flatten().collect::<Vec<_>>(),
             NUM_SHIFT_RIGHT_COLS,
         );
 
-        // Pad the trace to a power of two.
-        pad_to_power_of_two::<NUM_SHIFT_RIGHT_COLS, F>(&mut trace.values);
-
         // Create the template for the padded rows. These are fake rows that don't fail on some
         // sanity checks.
         let padded_row_template = {