Shankara/picus frontend updates

crates/core/machine/src/air/memory.rs

@@ -4,7 +4,7 @@ use p3_air::AirBuilder;
 use p3_field::FieldAlgebra;
 use zkm_core_executor::ByteOpcode;
 use zkm_stark::{
-    air::{AirLookup, BaseAirBuilder, ByteAirBuilder, LookupScope},
+    air::{AirLookup, BaseAirBuilder, ByteAirBuilder, LookupScope, OperationSummaryAirBuilder},
     LookupKind,
 };
 
@@ -22,7 +22,9 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
         addr: impl Into<Self::Expr>,
         memory_access: &impl MemoryCols<E>,
         do_check: impl Into<Self::Expr>,
-    ) {
+    ) where
+        Self: OperationSummaryAirBuilder,
+    {
         let do_check: Self::Expr = do_check.into();
         let shard: Self::Expr = shard.into();
         let clk: Self::Expr = clk.into();
@@ -69,7 +71,9 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
         initial_addr: impl Into<Self::Expr> + Clone,
         memory_access_slice: &[impl MemoryCols<E>],
         verify_memory_access: impl Into<Self::Expr> + Copy,
-    ) {
+    ) where
+        Self: OperationSummaryAirBuilder,
+    {
         for (i, access_slice) in memory_access_slice.iter().enumerate() {
             self.eval_memory_access(
                 shard,
@@ -93,24 +97,40 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
         do_check: impl Into<Self::Expr>,
         shard: impl Into<Self::Expr> + Clone,
         clk: impl Into<Self::Expr>,
-    ) {
+    ) where
+        Self: OperationSummaryAirBuilder,
+    {
         let do_check: Self::Expr = do_check.into();
         let compare_clk: Self::Expr = mem_access.compare_clk.clone().into();
         let shard: Self::Expr = shard.clone().into();
         let prev_shard: Self::Expr = mem_access.prev_shard.clone().into();
+        let prev_clk: Self::Expr = mem_access.prev_clk.clone().into();
+        let clk: Self::Expr = clk.into();
+        let diff_16bit_limb: Self::Expr = mem_access.diff_16bit_limb.clone().into();
+        let diff_8bit_limb: Self::Expr = mem_access.diff_8bit_limb.clone().into();
+
+        if self.try_emit_memory_timestamp_summary(
+            do_check.clone(),
+            shard.clone(),
+            clk.clone(),
+            prev_shard.clone(),
+            prev_clk.clone(),
+            compare_clk.clone(),
+            diff_16bit_limb.clone(),
+            diff_8bit_limb.clone(),
+        ) {
+            return;
+        }
 
         // First verify that compare_clk's value is correct.
         self.when(do_check.clone()).assert_bool(compare_clk.clone());
         self.when(do_check.clone()).when(compare_clk.clone()).assert_eq(shard.clone(), prev_shard);
 
         // Get the comparison timestamp values for the current and previous memory access.
-        let prev_comp_value = self.if_else(
-            mem_access.compare_clk.clone(),
-            mem_access.prev_clk.clone(),
-            mem_access.prev_shard.clone(),
-        );
+        let prev_comp_value =
+            self.if_else(mem_access.compare_clk.clone(), prev_clk, mem_access.prev_shard.clone());
 
-        let current_comp_val = self.if_else(compare_clk.clone(), clk.into(), shard.clone());
+        let current_comp_val = self.if_else(compare_clk.clone(), clk, shard.clone());
 
         // Assert `current_comp_val > prev_comp_val`. We check this by asserting that
         // `0 <= current_comp_val-prev_comp_val-1 < 2^24`.
@@ -125,12 +145,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
 
         // Verify that mem_access.ts_diff = mem_access.ts_diff_16bit_limb
         // + mem_access.ts_diff_8bit_limb * 2^16.
-        self.eval_range_check_24bits(
-            diff_minus_one,
-            mem_access.diff_16bit_limb.clone(),
-            mem_access.diff_8bit_limb.clone(),
-            do_check,
-        );
+        self.eval_range_check_24bits(diff_minus_one, diff_16bit_limb, diff_8bit_limb, do_check);
     }
 
     /// Verifies the inputted value is within 24 bits.

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

@@ -81,6 +81,10 @@ impl<F: PrimeField32> MachineAir<F> for CloClzChip {
         "CloClz".to_string()
     }
 
+    fn local_only(&self) -> bool {
+        true
+    }
+
     fn picus_info(&self) -> PicusInfo {
         CloClzCols::<u8>::picus_info()
     }

crates/core/machine/src/cpu/columns/mod.rs

@@ -3,8 +3,8 @@ pub use instruction::*;
 
 use p3_util::indices_arr;
 use std::mem::{size_of, transmute};
-use zkm_derive::AlignedBorrow;
-use zkm_stark::Word;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
+use zkm_stark::{PicusInfo, Word};
 
 use crate::memory::{MemoryCols, MemoryReadCols, MemoryReadWriteCols};
 
@@ -13,7 +13,7 @@ pub const NUM_CPU_COLS: usize = size_of::<CpuCols<u8>>();
 pub const CPU_COL_MAP: CpuCols<usize> = make_col_map();
 
 /// The column layout for the CPU.
-#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[derive(AlignedBorrow, Default, PicusAnnotations, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct CpuCols<T: Copy> {
     /// The current shard.

crates/core/machine/src/cpu/trace.rs

@@ -28,6 +28,10 @@ impl<F: PrimeField32> MachineAir<F> for CpuChip {
         self.id().to_string()
     }
 
+    fn picus_info(&self) -> zkm_stark::PicusInfo {
+        CpuCols::<u8>::picus_info()
+    }
+
     fn num_rows(&self, input: &Self::Record) -> Option<usize> {
         let n_real_rows = input.cpu_events.len();
         let padded_nb_rows = if let Some(shape) = &input.shape {

crates/core/machine/src/memory/global.rs

@@ -10,13 +10,13 @@ use p3_matrix::{dense::RowMajorMatrix, Matrix};
 use p3_maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use zkm_core_executor::events::{GlobalLookupEvent, MemoryInitializeFinalizeEvent};
 use zkm_core_executor::{ExecutionRecord, Program};
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
 use zkm_stark::{
     air::{
         AirLookup, BaseAirBuilder, LookupScope, MachineAir, PublicValues, ZKMAirBuilder,
         ZKM_PROOF_NUM_PV_ELTS,
     },
-    LookupKind, Word,
+    LookupKind, PicusInfo, Word,
 };
 
 use crate::{
@@ -59,6 +59,10 @@ impl<F: PrimeField32> MachineAir<F> for MemoryGlobalChip {
         }
     }
 
+    fn picus_info(&self) -> zkm_stark::PicusInfo {
+        MemoryInitCols::<u8>::picus_info()
+    }
+
     fn generate_dependencies(
         &self,
         input: &ExecutionRecord,
@@ -207,7 +211,7 @@ impl<F: PrimeField32> MachineAir<F> for MemoryGlobalChip {
     }
 }
 
-#[derive(AlignedBorrow, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Clone, Copy)]
 #[repr(C)]
 pub struct MemoryInitCols<T: Copy> {
     /// The shard number of the memory access.
@@ -279,8 +283,8 @@ where
             builder.send(
                 AirLookup::new(
                     vec![
-                        AB::Expr::zero(),
-                        AB::Expr::zero(),
+                        AB::Expr::zero(), // shard
+                        AB::Expr::zero(), // timestamp
                         local.addr.into(),
                         value[0].clone(),
                         value[1].clone(),
@@ -404,6 +408,7 @@ where
 
         if self.kind == MemoryChipType::Initialize {
             builder.when(local.is_real).assert_eq(local.timestamp, AB::F::ONE);
+            builder.when(local.is_real).assert_eq(local.shard, AB::F::ONE);
         }
 
         // Constraints related to register %x0.
@@ -431,6 +436,8 @@ where
         // - The flag `is_real` is set to one and the next `is_real` is set to zero.
 
         // Constrain the `is_last_addr` flag.
+        builder.assert_bool(local.is_last_addr);
+        builder.when_last_row().assert_eq(local.is_last_addr, local.is_real);
         builder
             .when_transition()
             .assert_eq(local.is_last_addr, local.is_real * (AB::Expr::one() - next.is_real));

crates/core/machine/src/misc/others/trace.rs

@@ -36,6 +36,10 @@ impl<F: PrimeField32> MachineAir<F> for MiscInstrsChip {
         MiscInstrColumns::<u8>::picus_info()
     }
 
+    fn local_only(&self) -> bool {
+        true
+    }
+
     fn num_rows(&self, input: &Self::Record) -> Option<usize> {
         let nb_rows = next_power_of_two(
             input.misc_events.len(),

crates/core/machine/src/operations/add.rs

@@ -1,11 +1,18 @@
+use core::{
+    borrow::Borrow,
+    mem::{size_of, transmute},
+};
+
 use zkm_core_executor::events::ByteRecord;
+use zkm_primitives::consts::WORD_SIZE;
 use zkm_stark::{air::ZKMAirBuilder, Word};
 
 use p3_air::AirBuilder;
 use p3_field::{Field, FieldAlgebra};
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusProjection};
 
 use crate::air::WordAirBuilder;
+use crate::utils::indices_arr;
 
 /// A set of columns needed to compute the add of two words.
 #[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
@@ -18,6 +25,46 @@ pub struct AddOperation<T> {
     pub carry: [T; 3],
 }
 
+const NUM_ADD_OPERATION_SUMMARY_COLS: usize = size_of::<AddOperationSummaryCols<u8>>();
+
+const ADD_OPERATION_SUMMARY_COL_MAP: AddOperationSummaryCols<usize> =
+    make_add_operation_summary_col_map();
+
+const fn make_add_operation_summary_col_map() -> AddOperationSummaryCols<usize> {
+    let indices_arr = indices_arr::<NUM_ADD_OPERATION_SUMMARY_COLS>();
+    unsafe {
+        transmute::<[usize; NUM_ADD_OPERATION_SUMMARY_COLS], AddOperationSummaryCols<usize>>(
+            indices_arr,
+        )
+    }
+}
+
+/// Hidden witness layout used when Picus emits the exact two-word add AIR as a
+/// local auxiliary module.
+#[derive(AlignedBorrow, Clone, Copy)]
+#[repr(C)]
+struct AddOperationSummaryCols<T> {
+    pub a: Word<T>,
+    pub b: Word<T>,
+    pub is_real: T,
+    pub cols: AddOperation<T>,
+}
+
+#[derive(PicusProjection)]
+#[picus_projection(
+    source = AddOperationSummaryCols<u8>,
+    col_map = ADD_OPERATION_SUMMARY_COL_MAP
+)]
+#[allow(dead_code)]
+struct AddOperationSummaryProjection {
+    #[picus(input, path = a)]
+    pub a: Word<u8>,
+    #[picus(input, path = b)]
+    pub b: Word<u8>,
+    #[picus(output, path = cols.value)]
+    pub value: Word<u8>,
+}
+
 impl<F: Field> AddOperation<F> {
     #[allow(unused_assignments)]
     pub fn populate(&mut self, record: &mut impl ByteRecord, a_u32: u32, b_u32: u32) -> u32 {
@@ -54,7 +101,7 @@ impl<F: Field> AddOperation<F> {
         expected
     }
 
-    pub fn eval<AB: ZKMAirBuilder>(
+    fn eval_exact<AB: ZKMAirBuilder>(
         builder: &mut AB,
         a: Word<AB::Var>,
         b: Word<AB::Var>,
@@ -98,4 +145,48 @@ impl<F: Field> AddOperation<F> {
             builder.slice_range_check_u8(&cols.value.0, is_real);
         }
     }
+
+    pub fn eval<AB: ZKMAirBuilder>(
+        builder: &mut AB,
+        a: Word<AB::Var>,
+        b: Word<AB::Var>,
+        cols: AddOperation<AB::Var>,
+        is_real: AB::Expr,
+    ) {
+        let mut current_inputs: Vec<AB::Expr> = Vec::with_capacity(WORD_SIZE * 2);
+        for limb in a.0 {
+            current_inputs.push(limb.into());
+        }
+        for limb in b.0 {
+            current_inputs.push(limb.into());
+        }
+
+        let current_outputs: Vec<AB::Expr> =
+            cols.value.0.iter().map(|limb| (*limb).into()).collect();
+
+        if builder.is_known_one(&is_real)
+            && builder.try_emit_projected_summary_with_hidden_consts(
+                "AddOperation",
+                &AddOperationSummaryProjection::picus_projection_info(),
+                &current_inputs,
+                &current_outputs,
+                size_of::<AddOperationSummaryCols<u8>>(),
+                &[(ADD_OPERATION_SUMMARY_COL_MAP.is_real, 1)],
+                |builder, source_row| {
+                    let source: &AddOperationSummaryCols<AB::Var> = (*source_row).borrow();
+                    Self::eval_exact(
+                        builder,
+                        source.a,
+                        source.b,
+                        source.cols,
+                        source.is_real.into(),
+                    );
+                },
+            )
+        {
+            return;
+        }
+
+        Self::eval_exact(builder, a, b, cols, is_real);
+    }
 }