Shankara/picus backend

Cargo.toml

@@ -16,6 +16,7 @@ members = [
     "crates/cuda",
     "crates/curves",
     "crates/derive",
+    "crates/picus",
     "crates/primitives",
     "crates/prover",
     "crates/recursion/circuit",
@@ -115,6 +116,7 @@ zkm-build = { path = "crates/build" }
 zkm-sdk = { path = "crates/sdk" }
 zkm-cuda = { path = "crates/cuda" }
 zkm-verifier = { path = "crates/verifier" }
+zkm-picus = {path = "crates/picus"}
 
 zkm-lib = { path = "crates/zkvm/lib", default-features = false }
 zkm-zkvm = { path = "crates/zkvm/entrypoint", default-features = false }

crates/core/executor/Cargo.toml

@@ -46,6 +46,9 @@ vec_map = { version = "0.8.2", features = ["serde"] }
 enum-map = { version = "2.7.3", features = ["serde"] }
 sha2 = { workspace = true }
 anyhow = { workspace = true }
+tracing-subscriber = "0.3.19"
+env_logger = "0.11.6"
+num_enum = "0.7.5"
 
 [dev-dependencies]
 test-artifacts = { path = "../../test-artifacts" }

crates/core/executor/src/opcode.rs

@@ -1,14 +1,26 @@
 //! Opcodes for ZKM.
 
 use enum_map::Enum;
+use num_enum::TryFromPrimitive;
 use p3_field::Field;
 use serde::{Deserialize, Serialize};
 use std::fmt::Display;
 
 /// An opcode (short for "operation code") specifies the operation to be performed by the processor.
 #[allow(non_camel_case_types)]
 #[derive(
-    Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord, Enum,
+    TryFromPrimitive,
+    Debug,
+    Clone,
+    Copy,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+    PartialOrd,
+    Ord,
+    Enum,
 )]
 #[repr(u8)]
 pub enum Opcode {

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

@@ -13,9 +13,9 @@ use zkm_core_executor::{
     events::{AluEvent, ByteLookupEvent, ByteRecord},
     ExecutionRecord, Opcode, Program,
 };
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
 use zkm_stark::{
-    air::{MachineAir, ZKMAirBuilder},
+    air::{MachineAir, PicusInfo, ZKMAirBuilder},
     Word,
 };
 
@@ -38,7 +38,7 @@ pub const NUM_ADD_SUB_COLS: usize = size_of::<AddSubCols<u8>>();
 pub struct AddSubChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)]
 #[repr(C)]
 pub struct AddSubCols<T> {
     /// The current/next pc, used for instruction lookup table.
@@ -56,9 +56,11 @@ pub struct AddSubCols<T> {
     pub operand_2: Word<T>,
 
     /// Flag indicating whether the opcode is `ADD`.
+    #[picus(selector)]
     pub is_add: T,
 
     /// Flag indicating whether the opcode is `SUB`.
+    #[picus(selector)]
     pub is_sub: T,
 }
 
@@ -78,6 +80,9 @@ impl<F: PrimeField32> MachineAir<F> for AddSubChip {
             next_power_of_two(input.add_sub_events.len(), input.fixed_log2_rows::<F, _>(self));
         Some(nb_rows)
     }
+    fn picus_info(&self) -> PicusInfo {
+        AddSubCols::<u8>::picus_info()
+    }
 
     fn generate_trace(
         &self,

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

@@ -13,9 +13,9 @@ use zkm_core_executor::{
     events::{AluEvent, ByteLookupEvent, ByteRecord},
     ByteOpcode, ExecutionRecord, Opcode, Program,
 };
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
 use zkm_stark::{
-    air::{MachineAir, ZKMAirBuilder},
+    air::{MachineAir, PicusInfo, ZKMAirBuilder},
     Word,
 };
 
@@ -29,7 +29,7 @@ pub const NUM_BITWISE_COLS: usize = size_of::<BitwiseCols<u8>>();
 pub struct BitwiseChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)]
 #[repr(C)]
 pub struct BitwiseCols<T> {
     /// The current/next pc, used for instruction lookup table.
@@ -46,15 +46,19 @@ pub struct BitwiseCols<T> {
     pub c: Word<T>,
 
     /// If the opcode is NOR.
+    #[picus(selector)]
     pub is_nor: T,
 
     /// If the opcode is XOR.
+    #[picus(selector)]
     pub is_xor: T,
 
     // If the opcode is OR.
+    #[picus(selector)]
     pub is_or: T,
 
     /// If the opcode is AND.
+    #[picus(selector)]
     pub is_and: T,
 }
 

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

@@ -22,8 +22,8 @@ use zkm_core_executor::{
     events::{ByteLookupEvent, ByteRecord},
     ByteOpcode, ExecutionRecord, Opcode, Program,
 };
-use zkm_derive::AlignedBorrow;
-use zkm_stark::{air::MachineAir, Word};
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
+use zkm_stark::{air::MachineAir, PicusInfo, Word};
 
 use crate::{air::ZKMCoreAirBuilder, utils::pad_rows_fixed, CoreChipError};
 
@@ -39,7 +39,7 @@ const BYTE_SIZE: usize = 8;
 pub struct CloClzChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct CloClzCols<T> {
     /// The current/next pc, used for instruction lookup table.
@@ -63,9 +63,11 @@ pub struct CloClzCols<T> {
     pub sr1: Word<T>,
 
     /// Flag to indicate whether the opcode is CLZ.
+    #[picus(selector)]
     pub is_clz: T,
 
     /// Flag to indicate whether the opcode is CLO.
+    #[picus(selector)]
     pub is_clo: T,
 
     /// Selector to know whether this row is enabled.
@@ -83,6 +85,10 @@ impl<F: PrimeField32> MachineAir<F> for CloClzChip {
         "CloClz".to_string()
     }
 
+    fn picus_info(&self) -> PicusInfo {
+        CloClzCols::<u8>::picus_info()
+    }
+
     fn generate_trace(
         &self,
         input: &ExecutionRecord,

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

@@ -76,9 +76,12 @@ use zkm_core_executor::{
 };
 
 use crate::{memory::MemoryReadWriteCols, CoreChipError};
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
 use zkm_primitives::consts::WORD_SIZE;
-use zkm_stark::{air::MachineAir, Word};
+use zkm_stark::{
+    air::{MachineAir, PicusInfo},
+    Word,
+};
 
 use crate::{
     air::{WordAirBuilder, ZKMCoreAirBuilder},
@@ -101,7 +104,7 @@ const LONG_WORD_SIZE: usize = 2 * WORD_SIZE;
 pub struct DivRemChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct DivRemCols<T> {
     /// The current/next pc, used for instruction lookup table.
@@ -139,15 +142,19 @@ pub struct DivRemCols<T> {
     pub is_c_0: IsZeroWordOperation<T>,
 
     /// Flag to indicate whether the opcode is DIV.
+    #[picus(selector)]
     pub is_div: T,
 
     /// Flag to indicate whether the opcode is DIVU.
+    #[picus(selector)]
     pub is_divu: T,
 
     /// Flag to indicate whether the opcode is MOD.
+    #[picus(selector)]
     pub is_mod: T,
 
     /// Flag to indicate whether the opcode is MODU.
+    #[picus(selector)]
     pub is_modu: T,
 
     /// Flag to indicate whether the division operation overflows.

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

@@ -13,10 +13,10 @@ use zkm_core_executor::{
     events::{AluEvent, ByteLookupEvent, ByteRecord},
     ByteOpcode, ExecutionRecord, Opcode, Program,
 };
-use zkm_derive::AlignedBorrow;
+use zkm_derive::{AlignedBorrow, PicusAnnotations};
 use zkm_stark::{
     air::{BaseAirBuilder, MachineAir, ZKMAirBuilder},
-    Word,
+    PicusInfo, Word,
 };
 
 use crate::{
@@ -32,17 +32,19 @@ pub const NUM_LT_COLS: usize = size_of::<LtCols<u8>>();
 pub struct LtChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Clone, Copy)]
 #[repr(C)]
 pub struct LtCols<T> {
     /// The current/next pc, used for instruction lookup table.
     pub pc: T,
     pub next_pc: T,
 
     /// If the opcode is SLT.
+    #[picus(selector)]
     pub is_slt: T,
 
     /// If the opcode is SLTU.
+    #[picus(selector)]
     pub is_sltu: T,
 
     /// The output operand.
@@ -104,6 +106,10 @@ impl<F: PrimeField32> MachineAir<F> for LtChip {
         "Lt".to_string()
     }
 
+    fn picus_info(&self) -> PicusInfo {
+        LtCols::<u8>::picus_info()
+    }
+
     fn generate_trace(
         &self,
         input: &ExecutionRecord,