fix: handle load mem instructions that write to reg x0 (#1409)

crates/core/machine/src/cpu/air/memory.rs

@@ -1,9 +1,6 @@
 use p3_air::AirBuilder;
 use p3_field::AbstractField;
-use sp1_stark::{
-    air::{BaseAirBuilder, SP1AirBuilder},
-    Word,
-};
+use sp1_stark::{air::SP1AirBuilder, Word};
 
 use crate::{
     air::{SP1CoreAirBuilder, WordAirBuilder},
@@ -155,26 +152,23 @@ impl CpuChip {
         // Get the memory specific columns.
         let memory_columns = local.opcode_specific_columns.memory();
 
-        // Compute whether this is a load instruction.
-        let is_load = self.is_load_instruction::<AB>(&local.selectors);
-
         // Verify the unsigned_mem_value column.
         self.eval_unsigned_mem_value(builder, memory_columns, local);
 
         // If it's a signed operation (such as LB or LH), then we need verify the bit decomposition
         // of the most significant byte to get it's sign.
         self.eval_most_sig_byte_bit_decomp(builder, memory_columns, local, &local.unsigned_mem_val);
 
-        // Assert that if `is_lb` and `is_lh` are both true, then the most significant byte
-        // matches the value of `local.mem_value_is_neg`.
+        // Assert that correct value of `mem_value_is_neg_not_x0`.
         builder.assert_eq(
-            local.mem_value_is_neg,
+            local.mem_value_is_neg_not_x0,
             (local.selectors.is_lb + local.selectors.is_lh)
-                * memory_columns.most_sig_byte_decomp[7],
+                * memory_columns.most_sig_byte_decomp[7]
+                * (AB::Expr::one() - local.instruction.op_a_0),
         );
 
-        // When the memory value is negative, use the SUB opcode to compute the signed value of
-        // the memory value and verify that the op_a value is correct.
+        // When the memory value is negative and not writing to x0, use the SUB opcode to compute
+        // the signed value of the memory value and verify that the op_a value is correct.
         let signed_value = Word([
             AB::Expr::zero(),
             AB::Expr::one() * local.selectors.is_lb,
@@ -189,14 +183,24 @@ impl CpuChip {
             local.shard,
             local.channel,
             local.unsigned_mem_val_nonce,
-            local.mem_value_is_neg,
+            local.mem_value_is_neg_not_x0,
+        );
+
+        // Assert that correct value of `mem_value_is_pos_not_x0`.
+        let mem_value_is_pos = (local.selectors.is_lb + local.selectors.is_lh)
+            * (AB::Expr::one() - memory_columns.most_sig_byte_decomp[7])
+            + local.selectors.is_lbu
+            + local.selectors.is_lhu
+            + local.selectors.is_lw;
+        builder.assert_eq(
+            local.mem_value_is_pos_not_x0,
+            mem_value_is_pos * (AB::Expr::one() - local.instruction.op_a_0),
         );
 
-        // When the memory value is not negaitve, assert that op_a value is equal to the unsigned
-        // memory value.
+        // When the memory value is not positive and not writing to x0, assert that op_a value is
+        // equal to the unsigned memory value.
         builder
-            .when(is_load)
-            .when_not(local.mem_value_is_neg)
+            .when(local.mem_value_is_pos_not_x0)
             .assert_word_eq(local.unsigned_mem_val, local.op_a_val());
     }
 

crates/core/machine/src/cpu/columns/mod.rs

@@ -88,10 +88,21 @@ pub struct CpuCols<T: Copy> {
     /// > (is_bge | is_bgeu) & !(a_eq_b | a_gt_b)
     pub not_branching: T,
 
-    /// The memory value is negative column is equal to:
+    /// Flag for load mem instructions where the value is negative and not writing to x0.
+    /// More formally, it is
     ///
-    /// > (is_lbu | is_lhu) & (most_sig_byte_decomp[7] == 1)
-    pub mem_value_is_neg: T,
+    /// > (is_lb | is_lh) & (most_sig_byte_decomp[7] == 1) & (not writing to x0)
+    pub mem_value_is_neg_not_x0: T,
+
+    /// Flag for load mem instructions where the value is positive and not writing to x0.
+    /// More formally, it is
+    ///
+    /// (
+    ///     ((is_lb | is_lh) & (most_sig_byte_decomp[7] == 0)) |
+    ///     is_lbu | is_lhu | is_lw
+    /// ) &
+    /// (not writing to x0)
+    pub mem_value_is_pos_not_x0: T,
 
     /// The unsigned memory value is the value after the offset logic is applied. Used for the load
     /// memory opcodes (i.e. LB, LH, LW, LBU, and LHU).

crates/core/machine/src/cpu/trace.rs

@@ -8,6 +8,7 @@ use sp1_core_executor::{
     ByteOpcode,
     ByteOpcode::U16Range,
     ExecutionRecord, Opcode, Program,
+    Register::X0,
 };
 use sp1_primitives::consts::WORD_SIZE;
 use sp1_stark::{air::MachineAir, Word};
@@ -351,7 +352,8 @@ impl CpuChip {
                         F::from_canonical_u8(most_sig_mem_value_byte >> i & 0x01);
                 }
                 if memory_columns.most_sig_byte_decomp[7] == F::one() {
-                    cols.mem_value_is_neg = F::one();
+                    cols.mem_value_is_neg_not_x0 =
+                        F::from_bool(event.instruction.op_a != (X0 as u32));
                     let sub_event = AluEvent {
                         lookup_id: event.memory_sub_lookup_id,
                         channel: event.channel,
@@ -373,6 +375,14 @@ impl CpuChip {
                         .or_insert(vec![sub_event]);
                 }
             }
+
+            // Set the `mem_value_is_pos_not_x0` composite flag.
+            cols.mem_value_is_pos_not_x0 = F::from_bool(
+                ((matches!(event.instruction.opcode, Opcode::LB | Opcode::LH)
+                    && (memory_columns.most_sig_byte_decomp[7] == F::zero()))
+                    || matches!(event.instruction.opcode, Opcode::LBU | Opcode::LHU | Opcode::LW))
+                    && event.instruction.op_a != (X0 as u32),
+            );
         }
 
         // Add event to byte lookup for byte range checking each byte in the memory addr