chore: constrain halt, is_real, and pc (#424)

Co-authored-by: Chris Tian <[email protected]>

core/src/air/builder.rs

@@ -168,6 +168,13 @@ pub trait WordAirBuilder: ByteAirBuilder {
         }
     }
 
+    /// Asserts that the word is zero.
+    fn assert_word_zero<I: Into<Self::Expr>>(&mut self, word: Word<I>) {
+        for limb in word.0 {
+            self.assert_zero(limb);
+        }
+    }
+
     /// Check that each limb of the given slice is a u8.
     fn slice_range_check_u8<EWord: Into<Self::Expr> + Clone, EMult: Into<Self::Expr> + Clone>(
         &mut self,

core/src/cpu/air/mod.rs

@@ -71,15 +71,20 @@ where
 
         // Write the `a` or the result to the first register described in the instruction unless
         // we are performing a branch or a store.
+        // If we are writing to register 0, then the new value should be zero.
+        builder
+            .when(local.instruction.op_a_0)
+            .assert_word_zero(*local.op_a_access.value());
         builder.constraint_memory_access(
             local.shard,
             local.clk + AB::F::from_canonical_u32(MemoryAccessPosition::A as u32),
             local.instruction.op_a[0],
             &local.op_a_access,
-            AB::Expr::one() - local.selectors.is_noop - local.selectors.reg_0_write,
+            local.is_real,
         );
 
         // If we are performing a branch or a store, then the value of `a` is the previous value.
+        // Also, if op_a is register 0, then ensure that its value is 0.
         builder
             .when(is_branch_instruction.clone() + self.is_store_instruction::<AB>(&local.selectors))
             .assert_word_eq(local.op_a_val(), local.op_a_access.prev_value);
@@ -129,7 +134,7 @@ where
         self.auipc_eval(builder, local);
 
         // ECALL instruction.
-        let (num_cycles, _is_halt) = self.ecall_eval(builder, local, next);
+        let (num_cycles, is_halt) = self.ecall_eval(builder, local);
 
         builder.send_alu(
             local.instruction.opcode,
@@ -139,19 +144,20 @@ where
             is_alu_instruction,
         );
 
-        // TODO: update the PC.
-        // Verify that the pc increments by 4 for all instructions except branch, jump and halt instructions.
-        // The other case is handled by eval_jump, eval_branch and eval_ecall.
-        // builder
-        //     .when_not(
-        //         is_branch_instruction + local.selectors.is_jal + local.selectors.is_jalr + is_halt,
-        //     )
-        //     .assert_eq(local.pc + AB::Expr::from_canonical_u8(4), next.pc);
-
         self.shard_clk_eval(builder, local, next, num_cycles);
 
-        // Range checks.
+        self.pc_eval(builder, local, next, is_branch_instruction.clone());
+
+        self.halt_unimpl_eval(builder, local, next, is_halt);
+
+        // Check the is_real flag.  It should be 1 for the first row.  Once its 0, it should never
+        // change value.
         builder.assert_bool(local.is_real);
+        builder.when_first_row().assert_one(local.is_real);
+        builder
+            .when_transition()
+            .when_not(local.is_real)
+            .assert_zero(next.is_real);
 
         // Dummy constraint of degree 3.
         builder.assert_eq(
@@ -189,8 +195,12 @@ impl CpuChip {
         let jump_columns = local.opcode_specific_columns.jump();
 
         // Verify that the local.pc + 4 is saved in op_a for both jump instructions.
+        // When op_a is set to register X0, the RISC-V spec states that the jump instruction will
+        // not have a return destination address (it is effectively a GOTO command).  In this case,
+        // we shouldn't verify the return address.
         builder
             .when(local.selectors.is_jal + local.selectors.is_jalr)
+            .when_not(local.instruction.op_a_0)
             .assert_eq(
                 local.op_a_val().reduce::<AB>(),
                 local.pc + AB::F::from_canonical_u8(4),
@@ -252,7 +262,6 @@ impl CpuChip {
         &self,
         builder: &mut AB,
         local: &CpuCols<AB::Var>,
-        _next: &CpuCols<AB::Var>,
     ) -> (AB::Expr, AB::Expr) {
         let ecall_cols = local.opcode_specific_columns.ecall();
         let is_ecall_instruction = self.is_ecall_instruction::<AB>(&local.selectors);
@@ -326,19 +335,6 @@ impl CpuChip {
             .when_not(is_enter_unconstrained + is_lwa + is_halt)
             .assert_word_eq(local.op_a_val(), local.op_a_access.prev_value);
 
-        // // For halt instructions, the next pc is 0.
-        // builder
-        //     .when_transition()
-        //     .when(is_lwa)
-        //     .assert_eq(next.pc, AB::Expr::from_canonical_u16(0));
-
-        // // If we're halting and it's a transition, then the next.is_real should be 0.
-        // builder
-        //     .when_transition()
-        //     .when(is_halt)
-        //     .assert_eq(next.is_real, AB::Expr::zero());
-        // builder.when_first_row().assert_one(local.is_real);
-        // We probably need a "halted" flag, this can be "is_noop" that turns on to control "is_real".
         (
             num_cycles * is_ecall_instruction.clone(),
             is_halt * is_ecall_instruction,
@@ -359,10 +355,7 @@ impl CpuChip {
         num_cycles: AB::Expr,
     ) {
         // Verify that all shard values are the same.
-        builder
-            .when_transition()
-            .when(next.is_real)
-            .assert_eq(local.shard, next.shard);
+        builder.when_transition().assert_eq(local.shard, next.shard);
 
         // Verify that the shard value is within 16 bits.
         builder.send_byte(
@@ -383,6 +376,12 @@ impl CpuChip {
             .when(next.is_real)
             .assert_eq(local.clk + clk_increment, next.clk);
 
+        // The clk value is carried down to the last row for non-real rows.
+        builder
+            .when_transition()
+            .when_not(next.is_real)
+            .assert_eq(local.clk, next.clk);
+
         // Range check that the clk is within 24 bits using it's limb values.
         // First verify that the limb values are correct.
         builder.verify_range_24bits(
@@ -392,6 +391,49 @@ impl CpuChip {
             local.is_real,
         );
     }
+
+    /// Constraints related to the pc for non jump, branch, and halt instructions.
+    ///
+    /// The function will verify that the pc increments by 4 for all instructions except branch, jump
+    /// and halt instructions. Also, it ensures that the pc is carried down to the last row for non-real rows.
+    pub(crate) fn pc_eval<AB: SP1AirBuilder>(
+        &self,
+        builder: &mut AB,
+        local: &CpuCols<AB::Var>,
+        next: &CpuCols<AB::Var>,
+        is_branch_instruction: AB::Expr,
+    ) {
+        // Verify that the pc increments by 4 for all instructions except branch, jump and halt instructions.
+        // The other case is handled by eval_jump, eval_branch and eval_ecall (for halt).
+        // Note that when the instruction is halt, we already contrain that the next new is not real,
+        // so the `when(next.is_real)` condition implies that the instruction is not halt.
+        builder
+            .when_transition()
+            .when(next.is_real)
+            .when_not(is_branch_instruction + local.selectors.is_jal + local.selectors.is_jalr)
+            .assert_eq(local.pc + AB::Expr::from_canonical_u8(4), next.pc);
+
+        // The pc value is carried down to the last row for non-real rows.
+        builder
+            .when_transition()
+            .when_not(next.is_real)
+            .assert_eq(local.pc, next.pc);
+    }
+
+    /// Constraint related to the halt and unimpl instruction.
+    pub(crate) fn halt_unimpl_eval<AB: SP1AirBuilder>(
+        &self,
+        builder: &mut AB,
+        local: &CpuCols<AB::Var>,
+        next: &CpuCols<AB::Var>,
+        is_halt: AB::Expr,
+    ) {
+        // If we're halting and it's a transition, then the next.is_real should be 0.
+        builder
+            .when_transition()
+            .when(is_halt + local.selectors.is_unimpl)
+            .assert_zero(next.is_real);
+    }
 }
 
 impl<F> BaseAir<F> for CpuChip {

core/src/cpu/columns/instruction.rs

@@ -3,6 +3,7 @@ use sp1_derive::AlignedBorrow;
 use std::mem::size_of;
 use std::{iter::once, vec::IntoIter};
 
+use crate::runtime::Register;
 use crate::{air::Word, runtime::Instruction};
 
 pub const NUM_INSTRUCTION_COLS: usize = size_of::<InstructionCols<u8>>();
@@ -22,6 +23,9 @@ pub struct InstructionCols<T> {
 
     /// The third operand for this instruction.
     pub op_c: Word<T>,
+
+    /// Flags to indicate if op_a is register 0.
+    pub op_a_0: T,
 }
 
 impl<F: PrimeField> InstructionCols<F> {
@@ -30,6 +34,8 @@ impl<F: PrimeField> InstructionCols<F> {
         self.op_a = instruction.op_a.into();
         self.op_b = instruction.op_b.into();
         self.op_c = instruction.op_c.into();
+
+        self.op_a_0 = F::from_bool(instruction.op_a == Register::X0 as u32);
     }
 }
 
@@ -42,6 +48,7 @@ impl<T> IntoIterator for InstructionCols<T> {
             .chain(self.op_a)
             .chain(self.op_b)
             .chain(self.op_c)
+            .chain(once(self.op_a_0))
             .collect::<Vec<_>>()
             .into_iter()
     }

core/src/cpu/columns/opcode.rs

@@ -47,8 +47,7 @@ pub struct OpcodeSelectorCols<T> {
 
     /// Miscellaneous.
     pub is_auipc: T,
-    pub is_noop: T,
-    pub reg_0_write: T,
+    pub is_unimpl: T,
 }
 
 impl<F: PrimeField> OpcodeSelectorCols<F> {
@@ -89,17 +88,7 @@ impl<F: PrimeField> OpcodeSelectorCols<F> {
         } else if instruction.opcode == Opcode::AUIPC {
             self.is_auipc = F::one();
         } else if instruction.opcode == Opcode::UNIMP {
-            self.is_noop = F::one();
-        }
-
-        // If op_a is 0 and we're writing to the register, then we don't do a write. We are always
-        // writing to the first register UNLESS it is a branch, is_store.
-        if instruction.op_a == 0
-            && !(instruction.is_branch_instruction()
-                || (self.is_sb + self.is_sh + self.is_sw) == F::one()
-                || self.is_noop == F::one())
-        {
-            self.reg_0_write = F::one();
+            self.is_unimpl = F::one();
         }
     }
 }
@@ -131,8 +120,7 @@ impl<T> IntoIterator for OpcodeSelectorCols<T> {
             self.is_jalr,
             self.is_jal,
             self.is_auipc,
-            self.is_noop,
-            self.reg_0_write,
+            self.is_unimpl,
         ]
         .into_iter()
     }

core/src/cpu/trace.rs

@@ -528,6 +528,7 @@ impl CpuChip {
         let n_real_rows = values.len() / NUM_CPU_COLS;
         let last_row = &values[len - NUM_CPU_COLS..];
         let pc = last_row[CPU_COL_MAP.pc];
+        let shard = last_row[CPU_COL_MAP.shard];
         let clk = last_row[CPU_COL_MAP.clk];
 
         values.resize(n_real_rows.next_power_of_two() * NUM_CPU_COLS, F::zero());
@@ -540,16 +541,13 @@ impl CpuChip {
             )
         };
 
-        rows[n_real_rows..]
-            .iter_mut()
-            .enumerate()
-            .for_each(|(n, padded_row)| {
-                padded_row[CPU_COL_MAP.pc] = pc;
-                padded_row[CPU_COL_MAP.clk] = clk + F::from_canonical_u32((n as u32 + 1) * 4);
-                padded_row[CPU_COL_MAP.selectors.is_noop] = F::one();
-                padded_row[CPU_COL_MAP.selectors.imm_b] = F::one();
-                padded_row[CPU_COL_MAP.selectors.imm_c] = F::one();
-            });
+        rows[n_real_rows..].iter_mut().for_each(|padded_row| {
+            padded_row[CPU_COL_MAP.pc] = pc;
+            padded_row[CPU_COL_MAP.shard] = shard;
+            padded_row[CPU_COL_MAP.clk] = clk;
+            padded_row[CPU_COL_MAP.selectors.imm_b] = F::one();
+            padded_row[CPU_COL_MAP.selectors.imm_c] = F::one();
+        });
     }
 }