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/air.rs

@@ -99,7 +99,7 @@ pub enum MipsAirId {
     SyscallInstrs = 38,
     /// The MemoryInstructionChip.
     MemoryInstrs = 39,
-    /// The MiscInstructionChip.
+    /// The MiscInstrsChip.
     MiscInstrs = 40,
     /// The memory global init chip.
     MemoryGlobalInit = 41,

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/include/cpu.hpp

@@ -50,9 +50,12 @@ __ZKM_HOSTDEV__ void event_to_row(
 
     populate_instruction<F>(cols.instruction, instruction);
 
+    // TEQ is a read-only instruction: it compares two registers and traps if equal,
+    // but must not modify register A. Mark it immutable to prevent register writes.
     cols.op_a_immutable = F::from_bool(
         is_memory_store_instruction_except_sc(instruction)
             || is_branch_instruction(instruction)
+            || instruction.opcode == Opcode::TEQ
     );
 
     cols.is_rw_a = F::from_bool(is_rw_a_instruction(instruction));

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,
 }
 
@@ -69,6 +73,10 @@ impl<F: PrimeField32> MachineAir<F> for BitwiseChip {
         "Bitwise".to_string()
     }
 
+    fn picus_info(&self) -> PicusInfo {
+        BitwiseCols::<u8>::picus_info()
+    }
+
     fn generate_trace(
         &self,
         input: &ExecutionRecord,

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,