Shankara/picus backend

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "1.2.4"
+version = "1.2.5"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 rust-version = "1.80"
@@ -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 }

README.md

@@ -33,4 +33,4 @@ Learn more about Ziren at [docs.zkm.io](https://docs.zkm.io).
 ## Acknowledgements
 Ziren draws inspiration from the following projects, which represents the cutting-edge zero-knowledge proof systems. 
 - [Plonky3](https://github.com/Plonky3/Plonky3): Ziren proving backend is based on Plonky3.
-- [SP1](https://github.com/succinctlabs/sp1): Ziren circuit builder, recursion compiler, and precompiles originate from SP1.
+- [SP1](https://github.com/succinctlabs/sp1): Ziren circuit builder, recursion compiler, and precompiles originate from SP1.

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/Cargo.toml

@@ -32,6 +32,7 @@ zkm-primitives = { workspace = true }
 zkm-core-executor = { workspace = true, default-features = true }
 zkm-curves = { workspace = true }
 
+crossbeam-channel = "0.5.12"
 rayon = "1.10.0"
 rayon-scan = "0.1.1"
 

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,
@@ -229,7 +237,7 @@ where
         builder.assert_bool(local.is_xor);
         builder.assert_bool(local.is_or);
         builder.assert_bool(local.is_and);
-        builder.assert_bool(local.is_xor);
+        builder.assert_bool(local.is_nor);
         builder.assert_bool(is_real);
     }
 }

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,

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

@@ -44,9 +44,9 @@ use zkm_core_executor::{
     events::{ByteLookupEvent, ByteRecord, CompAluEvent, MemoryAccessPosition, MemoryRecordEnum},
     ByteOpcode, ExecutionRecord, Opcode, Program,
 };
-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},
@@ -74,10 +74,11 @@ const BYTE_MASK: u8 = 0xff;
 pub struct MulChip;
 
 /// The column layout for the chip.
-#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
+#[derive(AlignedBorrow, PicusAnnotations, Default, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct MulCols<T> {
     /// The current/next pc, used for instruction lookup table.
+    #[picus(input)]
     pub pc: T,
     pub next_pc: T,
 
@@ -112,15 +113,17 @@ pub struct MulCols<T> {
     pub c_sign_extend: T,
 
     /// Flag indicating whether the opcode is `MUL`.
+    #[picus(selector)]
     pub is_mul: T,
 
     /// Flag indicating whether the opcode is `MULT`.
+    #[picus(selector)]
     pub is_mult: T,
 
     /// Flag indicating whether the opcode is `MULTU`.
+    #[picus(selector)]
     pub is_multu: T,
 
-    /// Selector to know whether this row is enabled.
     pub is_real: T,
 
     /// Access to hi register
@@ -146,6 +149,10 @@ impl<F: PrimeField32> MachineAir<F> for MulChip {
         "Mul".to_string()
     }
 
+    fn picus_info(&self) -> PicusInfo {
+        MulCols::<u8>::picus_info()
+    }
+
     fn generate_trace(
         &self,
         input: &ExecutionRecord,
@@ -395,11 +402,11 @@ where
         let product = {
             for i in 0..PRODUCT_SIZE {
                 if i == 0 {
-                    builder.assert_eq(local.product[i], m[i].clone() - local.carry[i] * base);
+                    builder.assert_eq(m[i].clone(), local.carry[i] * base + local.product[i]);
                 } else {
                     builder.assert_eq(
-                        local.product[i],
-                        m[i].clone() + local.carry[i - 1] - local.carry[i] * base,
+                        local.product[i] + local.carry[i] * base - local.carry[i - 1],
+                        m[i].clone(),
                     );
                 }
             }