opt: reduce unnecessary cols in jump and cloclz

crates/core/executor/src/artifacts/mips_costs.json

@@ -1,26 +1,23 @@
 {
   "U256XU2048Mul": 5849,
   "Bn254Fp2MulAssign": 2885,
-  "CloClz": 101,
+  "CloClz": 47,
   "Bls12831Fp2AddSubAssign": 2070,
+  "SyscallInstrs": 89,
   "DivRem": 144,
   "ShiftRight": 136,
   "Secp256r1Decompress": 2686,
   "Secp256k1Decompress": 2686,
   "KeccakSponge": 102216,
-  "Bitwise": 43,
   "Bn254AddAssign": 4013,
+  "Bitwise": 43,
   "ShiftLeft": 69,
   "MemoryLocal": 100,
   "Program": 31,
   "Global": 428,
   "Secp256k1AddAssign": 4013,
   "AddSub": 48,
-  "Jump": 87,
-  "Branch": 99,
-  "MemoryInstrs": 120,
-  "SyscallInstrs": 89,
-  "MiscInstrs": 115,
+  "Jump": 83,
   "Bn254FpOpAssign": 704,
   "Mul": 87,
   "ShaExtend": 15936,
@@ -39,10 +36,13 @@
   "Bls12381Decompress": 4149,
   "Lt": 57,
   "Secp256r1DoubleAssign": 4492,
+  "Branch": 99,
   "SyscallCore": 22,
   "Bn254Fp2AddSubAssign": 1382,
   "Bls12381FpOpAssign": 1048,
   "Cpu": 137,
   "ShaCompress": 40480,
+  "MemoryInstrs": 120,
+  "MiscInstrs": 115,
   "Secp256k1DoubleAssign": 4492
 }

crates/core/executor/src/executor.rs

@@ -1061,16 +1061,21 @@ impl<'a> Executor<'a> {
                     self.local_counts.event_counts[Opcode::ADD] += 1;
                     self.local_counts.event_counts[Opcode::SLT] += 2;
                 }
-                Opcode::DIVU | Opcode::DIV => {
-                    self.local_counts.event_counts[Opcode::MUL] += 2;
+                Opcode::DIV => {
+                    self.local_counts.event_counts[Opcode::MULT] += 2;
+                    self.local_counts.event_counts[Opcode::ADD] += 2;
+                    self.local_counts.event_counts[Opcode::SLTU] += 1;
+                }
+                Opcode::DIVU => {
+                    self.local_counts.event_counts[Opcode::MULTU] += 2;
                     self.local_counts.event_counts[Opcode::ADD] += 2;
                     self.local_counts.event_counts[Opcode::SLTU] += 1;
                 }
                 Opcode::CLZ | Opcode::CLO => {
                     self.local_counts.event_counts[Opcode::SRL] += 1;
                 }
                 Opcode::MADDU | Opcode::MSUBU => {
-                    self.local_counts.event_counts[Opcode::MULT] += 1;
+                    self.local_counts.event_counts[Opcode::MULTU] += 1;
                 }
                 Opcode::EXT => {
                     self.local_counts.event_counts[Opcode::SLL] += 1;

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

@@ -25,8 +25,7 @@ use zkm_core_executor::{
 use zkm_derive::AlignedBorrow;
 use zkm_stark::{air::MachineAir, Word};
 
-use crate::operations::{FixedShiftRightOperation, IsEqualWordOperation, IsZeroOperation};
-use crate::{air::ZKMCoreAirBuilder, operations::IsZeroWordOperation, utils::pad_rows_fixed};
+use crate::{air::ZKMCoreAirBuilder, utils::pad_rows_fixed};
 
 /// The number of main trace columns for `CloClzChip`.
 pub const NUM_CLOCLZ_COLS: usize = size_of::<CloClzCols<u8>>();
@@ -60,25 +59,12 @@ pub struct CloClzCols<T> {
     /// if clz, bb == 0xffffffff - b
     pub bb: Word<T>,
 
-    /// whether the result is 32.
-    pub is_result_32: IsZeroOperation<T>,
-
     /// whether the `bb` is zero.
-    pub is_bb_zero: IsZeroWordOperation<T>,
-
-    /// bb shift right by `32 - result`.
-    /// Use right-shifting sr1 by 1 to obtain sr0.
-    pub sr0: FixedShiftRightOperation<T>,
+    pub is_bb_zero: T,
 
     /// bb shift right by `32 - (result + 1)`.
     pub sr1: Word<T>,
 
-    /// sr0 == 0
-    pub is_sr0_zero: IsZeroWordOperation<T>,
-
-    /// sr1 == 1
-    pub is_sr1_one: IsEqualWordOperation<T>,
-
     /// Flag to indicate whether the opcode is CLZ.
     pub is_clz: T,
 
@@ -124,16 +110,11 @@ impl<F: PrimeField32> MachineAir<F> for CloClzChip {
             cols.bb = Word::from(bb);
 
             // if bb == 0, then result is 32.
-            cols.is_result_32.populate(32 - event.a);
-            cols.is_bb_zero.populate(bb);
+            cols.is_bb_zero = F::from_bool(bb == 0);
 
             if bb != 0 {
                 let sr1_val = bb >> (31 - event.a);
                 cols.sr1 = Word::from(sr1_val);
-                cols.sr0.populate(output, sr1_val, 1);
-
-                cols.is_sr0_zero.populate(sr1_val >> 1);
-                cols.is_sr1_one.populate(sr1_val, 1);
             }
 
             // Range check.
@@ -168,8 +149,7 @@ impl<F: PrimeField32> MachineAir<F> for CloClzChip {
             // clz(0) = 32
             cols.a = Word::from(32);
             cols.is_clz = F::ONE;
-            cols.is_bb_zero.populate(0);
-            cols.is_result_32.populate(0);
+            cols.is_bb_zero = F::ONE;
 
             row
         };
@@ -222,7 +202,7 @@ where
         builder.send_byte(
             ByteOpcode::LTU.as_field::<AB::F>(),
             AB::F::ONE,
-            local.a[0],
+            local.a.reduce::<AB>(),
             AB::Expr::from_canonical_u8(33),
             local.is_real,
         );
@@ -251,24 +231,13 @@ where
             local.is_real,
         );
 
-        // if bb == 0, result is 32
-        let is_bb_zero: AB::Expr = local.is_bb_zero.result.into();
+        // if is_bb_zero == 1, bb == 0, and result is 32
         {
-            IsZeroWordOperation::<AB::F>::eval(
-                builder,
-                local.bb.map(|x| x.into()),
-                local.is_bb_zero,
-                local.is_real.into(),
-            );
+            builder.assert_bool(local.is_bb_zero);
 
-            IsZeroOperation::<AB::F>::eval(
-                builder,
-                AB::Expr::from_canonical_u32(32) - local.a[0],
-                local.is_result_32,
-                local.is_real.into(),
-            );
+            builder.when(local.is_bb_zero).assert_zero(local.bb.reduce::<AB>());
 
-            builder.when(is_bb_zero.clone()).assert_one(local.is_result_32.result);
+            builder.when(local.is_bb_zero).assert_eq(local.a[0], AB::Expr::from_canonical_u32(32));
         }
 
         {
@@ -283,35 +252,18 @@ where
                     zero.clone(),
                     zero.clone(),
                 ]),
-                one.clone() - is_bb_zero.clone(),
-            );
-
-            FixedShiftRightOperation::<AB::F>::eval(
-                builder,
-                local.sr1,
-                1,
-                local.sr0,
-                one.clone() - is_bb_zero.clone(),
+                one.clone() - local.is_bb_zero,
             );
         }
 
-        // if bb!=0, check sr == 0 and sr1 == 1
+        // if bb!=0, check sr1 == 1
         {
-            builder.when_not(is_bb_zero.clone()).assert_one(local.sr1[0]);
-            local
-                .sr0
-                .value
-                .into_iter()
-                .chain(local.sr1.into_iter().skip(1))
-                .for_each(|x| builder.when_not(is_bb_zero.clone()).assert_zero(x));
+            builder.when_not(local.is_bb_zero).assert_one(local.sr1.reduce::<AB>());
         }
 
         builder.assert_bool(local.is_clo);
         builder.assert_bool(local.is_clz);
         builder.assert_one(local.is_clo + local.is_clz);
-        builder.assert_zero(local.a[1]);
-        builder.assert_zero(local.a[2]);
-        builder.assert_zero(local.a[3]);
     }
 }
 

crates/core/machine/src/control_flow/jump/air.rs

@@ -9,7 +9,6 @@ use zkm_stark::{
     Word,
 };
 
-use crate::air::WordAirBuilder;
 use crate::operations::KoalaBearWordRangeChecker;
 
 use super::{JumpChip, JumpColumns};
@@ -96,8 +95,8 @@ where
         );
         KoalaBearWordRangeChecker::<AB::F>::range_check(
             builder,
-            local.target_pc,
-            local.target_pc_range_checker,
+            local.next_next_pc,
+            local.next_next_pc_range_checker,
             is_real.clone(),
         );
 
@@ -107,13 +106,10 @@ where
         // SAFETY: `is_jal` is boolean, and zero for padding rows.
         builder.send_alu(
             AB::Expr::from_canonical_u32(Opcode::ADD as u32),
-            local.target_pc,
+            local.next_next_pc,
             local.next_pc,
             local.op_b_value,
             local.is_jumpdirect,
         );
-
-        // Assert that local.next_next_pc <==> local.target_pc.
-        builder.when(is_real.clone()).assert_word_eq(local.target_pc, local.next_next_pc);
     }
 }

crates/core/machine/src/control_flow/jump/columns.rs

@@ -16,12 +16,9 @@ pub struct JumpColumns<T> {
     pub next_pc: Word<T>,
     pub next_pc_range_checker: KoalaBearWordRangeChecker<T>,
 
-    /// The target program counter.
-    pub target_pc: Word<T>,
-    pub target_pc_range_checker: KoalaBearWordRangeChecker<T>,
-
     /// The next program counter.
     pub next_next_pc: Word<T>,
+    pub next_next_pc_range_checker: KoalaBearWordRangeChecker<T>,
 
     /// The value of the first operand.
     pub op_a_value: Word<T>,

crates/core/machine/src/control_flow/jump/trace.rs

@@ -90,26 +90,10 @@ impl JumpChip {
         cols.op_b_value = event.b.into();
         cols.op_c_value = event.c.into();
         cols.op_a_0 = F::from_bool(event.op_a_0);
+        cols.op_a_range_checker.populate(event.a);
+        cols.next_pc = Word::from(event.next_pc);
+        cols.next_pc_range_checker.populate(event.next_pc);
         cols.next_next_pc = Word::from(event.next_next_pc);
-
-        match event.opcode {
-            Opcode::Jump | Opcode::Jumpi => {
-                let target_pc = event.b;
-                cols.op_a_range_checker.populate(event.a);
-                cols.target_pc = Word::from(target_pc);
-                cols.next_pc = Word::from(event.next_pc);
-                cols.next_pc_range_checker.populate(event.next_pc);
-                cols.target_pc_range_checker.populate(target_pc);
-            }
-            Opcode::JumpDirect => {
-                let target_pc = event.next_pc.wrapping_add(event.b);
-                cols.op_a_range_checker.populate(event.a);
-                cols.next_pc = Word::from(event.next_pc);
-                cols.next_pc_range_checker.populate(event.next_pc);
-                cols.target_pc = Word::from(target_pc);
-                cols.target_pc_range_checker.populate(target_pc);
-            }
-            _ => unreachable!(),
-        }
+        cols.next_next_pc_range_checker.populate(event.next_next_pc);
     }
 }

crates/core/machine/src/mips/mod.rs

@@ -723,7 +723,6 @@ pub mod tests {
     }
 
     #[test]
-    #[ignore]
     fn write_core_air_costs() {
         let costs = MipsAir::<KoalaBear>::costs();
         println!("{costs:?}");