chore: cleanup tasks for core

core/src/alu/mul/mod.rs

@@ -412,14 +412,6 @@ where
             .when(local.c_sign_extend)
             .assert_eq(local.c_msb, one.clone());
 
-        // If the opcode doesn't allow sign extension for an operand, we must not extend their sign.
-        builder
-            .when(local.is_mul + local.is_mulhu)
-            .assert_zero(local.b_sign_extend + local.c_sign_extend);
-        builder
-            .when(local.is_mul + local.is_mulhsu + local.is_mulhsu)
-            .assert_zero(local.c_sign_extend);
-
         // Calculate the opcode.
         let opcode = {
             // Exactly one of the op codes must be on.

core/src/alu/sr/mod.rs

@@ -348,6 +348,16 @@ where
                     .assert_eq(num_bits_to_shift.clone(), AB::F::from_canonical_usize(i));
             }
 
+            // Bool check c_least_sig_bytes elements.
+            for bit in local.c_least_sig_byte.iter() {
+                builder.assert_bool(*bit);
+            }
+
+            // Bool check shift_by_n_bits elements.
+            for shift in local.shift_by_n_bits.iter() {
+                builder.assert_bool(*shift);
+            }
+
             // Exactly one of the shift_by_n_bits must be 1.
             builder.assert_eq(
                 local
@@ -485,6 +495,8 @@ where
         builder.assert_bool(local.is_sra);
         builder.assert_bool(local.is_real);
 
+        builder.assert_eq(local.is_srl + local.is_sra, local.is_real);
+
         // Receive the arguments.
         builder.receive_alu(
             local.is_srl * AB::F::from_canonical_u32(Opcode::SRL as u32)

core/src/cpu/air/branch.rs

@@ -3,6 +3,7 @@ use p3_field::AbstractField;
 
 use crate::air::{BaseAirBuilder, SP1AirBuilder, Word, WordAirBuilder};
 use crate::cpu::columns::{CpuCols, OpcodeSelectorCols};
+use crate::operations::BabyBearWord;
 use crate::{cpu::CpuChip, runtime::Opcode};
 
 impl CpuChip {
@@ -42,26 +43,38 @@ impl CpuChip {
             // When we are branching, assert local.pc <==> branch_cols.pc as Word.
             builder
                 .when(local.branching)
-                .assert_eq(branch_cols.pc.reduce::<AB>(), local.pc);
+                .assert_eq(branch_cols.pc.value.reduce::<AB>(), local.pc);
 
             // When we are branching, assert that next.pc <==> branch_columns.next_pc as Word.
             builder
                 .when_transition()
                 .when(next.is_real)
                 .when(local.branching)
-                .assert_eq(branch_cols.next_pc.reduce::<AB>(), next.pc);
+                .assert_eq(branch_cols.next_pc.value.reduce::<AB>(), next.pc);
 
             // When the current row is real and local.branching, assert that local.next_pc <==> branch_columns.next_pc as Word.
             builder
                 .when(local.is_real)
                 .when(local.branching)
-                .assert_eq(branch_cols.next_pc.reduce::<AB>(), local.next_pc);
+                .assert_eq(branch_cols.next_pc.value.reduce::<AB>(), local.next_pc);
+
+            // Range check branch_cols.pc and branch_cols.next_pc
+            BabyBearWord::<AB::F>::range_check(
+                builder,
+                branch_cols.pc,
+                is_branch_instruction.clone(),
+            );
+            BabyBearWord::<AB::F>::range_check(
+                builder,
+                branch_cols.next_pc,
+                is_branch_instruction.clone(),
+            );
 
             // When we are branching, calculate branch_cols.next_pc <==> branch_cols.pc + c.
             builder.send_alu(
                 Opcode::ADD.as_field::<AB::F>(),
-                branch_cols.next_pc,
-                branch_cols.pc,
+                branch_cols.next_pc.value,
+                branch_cols.pc.value,
                 local.op_c_val(),
                 local.shard,
                 local.channel,

core/src/cpu/air/ecall.rs

@@ -35,10 +35,14 @@ impl CpuChip {
         let syscall_id = syscall_code[0];
         let send_to_table = syscall_code[1];
 
-        // When is_ecall_instruction == true AND sent_to_table == true, ecall_mul_send_to_table should be true.
+        // When is_ecall_instruction == true AND send_to_table == true, ecall_mul_send_to_table should be true.
+        builder.assert_bool(local.ecall_mul_send_to_table);
         builder
             .when(is_ecall_instruction.clone())
-            .assert_eq(send_to_table, local.ecall_mul_send_to_table);
+            .assert_bool(send_to_table);
+        builder
+            .when(local.ecall_mul_send_to_table)
+            .assert_one(send_to_table * is_ecall_instruction.clone());
         builder.send_syscall(
             local.shard,
             local.channel,

core/src/cpu/air/memory.rs

@@ -88,6 +88,35 @@ impl CpuChip {
                 memory_columns.addr_word.reduce::<AB>(),
             );
 
+        // Verify that the least significant byte of addr_word - addr_offset is divisible by 4.
+        let offset = [
+            memory_columns.offset_is_one,
+            memory_columns.offset_is_two,
+            memory_columns.offset_is_three,
+        ]
+        .iter()
+        .enumerate()
+        .fold(AB::Expr::zero(), |acc, (index, &value)| {
+            acc + AB::Expr::from_canonical_usize(index + 1) * value
+        });
+        let mut recomposed_byte = AB::Expr::zero();
+        memory_columns
+            .aa_least_sig_byte_decomp
+            .iter()
+            .enumerate()
+            .for_each(|(i, value)| {
+                builder
+                    .when(is_memory_instruction.clone())
+                    .assert_bool(*value);
+
+                recomposed_byte =
+                    recomposed_byte.clone() + AB::Expr::from_canonical_usize(1 << (i + 2)) * *value;
+            });
+
+        builder
+            .when(is_memory_instruction.clone())
+            .assert_eq(memory_columns.addr_word[0] - offset, recomposed_byte);
+
         // For operations that require reading from memory (not registers), we need to read the
         // value into the memory columns.
         builder.eval_memory_access(
@@ -98,6 +127,14 @@ impl CpuChip {
             &memory_columns.memory_access,
             is_memory_instruction.clone(),
         );
+
+        // On memory load instructions, make sure that the memory value is not changed.
+        builder
+            .when(self.is_load_instruction::<AB>(&local.selectors))
+            .assert_word_eq(
+                *memory_columns.memory_access.value(),
+                *memory_columns.memory_access.prev_value(),
+            );
     }
 
     /// Evaluates constraints related to loading from memory.
@@ -146,7 +183,7 @@ impl CpuChip {
             local.mem_value_is_neg,
         );
 
-        // When the memory value is not negaitve, assert that op_a value is equal to the unsigned
+        // When the memory value is not negative, assert that op_a value is equal to the unsigned
         // memory value.
         builder
             .when(is_load)
@@ -195,6 +232,11 @@ impl CpuChip {
             .when(local.selectors.is_sh)
             .assert_zero(memory_columns.offset_is_one + memory_columns.offset_is_three);
 
+        // When the instruction is SW, ensure that the offset is 0.
+        builder
+            .when(local.selectors.is_sw)
+            .assert_one(offset_is_zero.clone());
+
         // Compute the expected stored value for a SH instruction.
         let a_is_lower_half = offset_is_zero;
         let a_is_upper_half = memory_columns.offset_is_two;
@@ -247,6 +289,12 @@ impl CpuChip {
         builder
             .when(local.selectors.is_lh + local.selectors.is_lhu)
             .assert_zero(memory_columns.offset_is_one + memory_columns.offset_is_three);
+
+        // When the instruction is LW, ensure that the offset is zero.
+        builder
+            .when(local.selectors.is_lw)
+            .assert_one(offset_is_zero.clone());
+
         let use_lower_half = offset_is_zero;
         let use_upper_half = memory_columns.offset_is_two;
         let half_value = Word([

core/src/cpu/air/mod.rs

@@ -22,6 +22,7 @@ use crate::bytes::ByteOpcode;
 use crate::cpu::columns::OpcodeSelectorCols;
 use crate::cpu::columns::{CpuCols, NUM_CPU_COLS};
 use crate::cpu::CpuChip;
+use crate::operations::BabyBearWord;
 use crate::runtime::Opcode;
 
 use super::columns::eval_channel_selectors;
@@ -160,26 +161,34 @@ impl CpuChip {
         // Verify that the word form of local.pc is correct for JAL instructions.
         builder
             .when(local.selectors.is_jal)
-            .assert_eq(jump_columns.pc.reduce::<AB>(), local.pc);
+            .assert_eq(jump_columns.pc.value.reduce::<AB>(), local.pc);
 
         // Verify that the word form of next.pc is correct for both jump instructions.
         builder
             .when_transition()
             .when(next.is_real)
             .when(is_jump_instruction.clone())
-            .assert_eq(jump_columns.next_pc.reduce::<AB>(), next.pc);
+            .assert_eq(jump_columns.next_pc.value.reduce::<AB>(), next.pc);
 
         // When the last row is real and it's a jump instruction, assert that local.next_pc <==> jump_column.next_pc
         builder
             .when(local.is_real)
             .when(is_jump_instruction.clone())
-            .assert_eq(jump_columns.next_pc.reduce::<AB>(), local.next_pc);
+            .assert_eq(jump_columns.next_pc.value.reduce::<AB>(), local.next_pc);
+
+        // Range check pc and next_pc
+        BabyBearWord::<AB::F>::range_check(builder, jump_columns.pc, local.selectors.is_jal.into());
+        BabyBearWord::<AB::F>::range_check(
+            builder,
+            jump_columns.next_pc,
+            is_jump_instruction.clone(),
+        );
 
         // Verify that the new pc is calculated correctly for JAL instructions.
         builder.send_alu(
             AB::Expr::from_canonical_u32(Opcode::ADD as u32),
-            jump_columns.next_pc,
-            jump_columns.pc,
+            jump_columns.next_pc.value,
+            jump_columns.pc.value,
             local.op_b_val(),
             local.shard,
             local.channel,
@@ -189,7 +198,7 @@ impl CpuChip {
         // Verify that the new pc is calculated correctly for JALR instructions.
         builder.send_alu(
             AB::Expr::from_canonical_u32(Opcode::ADD as u32),
-            jump_columns.next_pc,
+            jump_columns.next_pc.value,
             local.op_b_val(),
             local.op_c_val(),
             local.shard,
@@ -206,13 +215,20 @@ impl CpuChip {
         // Verify that the word form of local.pc is correct.
         builder
             .when(local.selectors.is_auipc)
-            .assert_eq(auipc_columns.pc.reduce::<AB>(), local.pc);
+            .assert_eq(auipc_columns.pc.value.reduce::<AB>(), local.pc);
+
+        // Range check the pc.
+        BabyBearWord::<AB::F>::range_check(
+            builder,
+            auipc_columns.pc,
+            local.selectors.is_auipc.into(),
+        );
 
         // Verify that op_a == pc + op_b.
         builder.send_alu(
             AB::Expr::from_canonical_u32(Opcode::ADD as u32),
             local.op_a_val(),
-            auipc_columns.pc,
+            auipc_columns.pc.value,
             local.op_b_val(),
             local.shard,
             local.channel,
@@ -288,17 +304,15 @@ impl CpuChip {
         next: &CpuCols<AB::Var>,
         is_branch_instruction: AB::Expr,
     ) {
-        // Verify that if is_sequential_instr is true, assert that local.is_real is true.
-        // This is needed for the following constraint, which is already degree 3.
-        builder
-            .when(local.is_sequential_instr)
-            .assert_one(local.is_real);
-
-        // When is_sequential_instr is true, assert that instruction is not branch, jump, or halt.
-        // Note that the condition `when(local_is_real)` is implied from the previous constraint.
+        // Verify the is_sequential_instr flag.
         let is_halt = self.get_is_halt_syscall::<AB>(builder, local);
-        builder.when(local.is_sequential_instr).assert_zero(
-            is_branch_instruction + local.selectors.is_jal + local.selectors.is_jalr + is_halt,
+        builder.when(local.is_real).assert_eq(
+            local.is_sequential_instr,
+            AB::Expr::one()
+                - (is_branch_instruction
+                    + local.selectors.is_jal
+                    + local.selectors.is_jalr
+                    + is_halt),
         );
 
         // Verify that the pc increments by 4 for all instructions except branch, jump and halt instructions.

core/src/cpu/columns/auipc.rs

@@ -1,13 +1,13 @@
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 
-use crate::air::Word;
+use crate::operations::BabyBearWord;
 
 pub const NUM_AUIPC_COLS: usize = size_of::<AuipcCols<u8>>();
 
 #[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct AuipcCols<T> {
     /// The current program counter.
-    pub pc: Word<T>,
+    pub pc: BabyBearWord<T>,
 }

core/src/cpu/columns/branch.rs

@@ -1,7 +1,7 @@
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 
-use crate::air::Word;
+use crate::operations::BabyBearWord;
 
 pub const NUM_BRANCH_COLS: usize = size_of::<BranchCols<u8>>();
 
@@ -10,10 +10,10 @@ pub const NUM_BRANCH_COLS: usize = size_of::<BranchCols<u8>>();
 #[repr(C)]
 pub struct BranchCols<T> {
     /// The current program counter.
-    pub pc: Word<T>,
+    pub pc: BabyBearWord<T>,
 
     /// The next program counter.
-    pub next_pc: Word<T>,
+    pub next_pc: BabyBearWord<T>,
 
     /// Whether a equals b.
     pub a_eq_b: T,

core/src/cpu/columns/jump.rs

@@ -1,16 +1,16 @@
 use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 
-use crate::air::Word;
+use crate::operations::BabyBearWord;
 
 pub const NUM_JUMP_COLS: usize = size_of::<JumpCols<u8>>();
 
 #[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
 #[repr(C)]
 pub struct JumpCols<T> {
     /// The current program counter.
-    pub pc: Word<T>,
+    pub pc: BabyBearWord<T>,
 
     /// THe next program counter.
-    pub next_pc: Word<T>,
+    pub next_pc: BabyBearWord<T>,
 }

core/src/cpu/columns/memory.rs

@@ -18,6 +18,8 @@ pub struct MemoryColumns<T> {
     // Note that this all needs to be verified in the AIR
     pub addr_word: Word<T>,
     pub addr_aligned: T,
+    /// The LE bit decomp of the least significant byte of address aligned.
+    pub aa_least_sig_byte_decomp: [T; 6],
     pub addr_offset: T,
     pub memory_access: MemoryReadWriteCols<T>,