Better structure cpu cols (#14)

* Better structure cpu cols

* cleanup

* builds

core/src/air/bool.rs

@@ -1,10 +1,12 @@
+use core::borrow::{Borrow, BorrowMut};
 use p3_air::AirBuilder;
 use p3_field::AbstractField;
 
 use super::AirVariable;
+use valida_derive::AlignedBorrow;
 
 /// An AIR representation of a boolean value.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default, AlignedBorrow)]
 pub struct Bool<T>(pub T);
 
 impl<AB: AirBuilder> AirVariable<AB> for Bool<AB::Var> {

core/src/air/word.rs

@@ -1,14 +1,17 @@
 use std::ops::{Index, IndexMut};
 
+use core::borrow::{Borrow, BorrowMut};
 use p3_air::AirBuilder;
+use p3_field::Field;
+use valida_derive::AlignedBorrow;
 
 use super::AirVariable;
 
 /// Using a 32-bit word size, we use four field elements to represent a 32-bit word.
 const WORD_LEN: usize = 4;
 
 /// An AIR representation of a word in the instruction set.
-#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, AlignedBorrow)]
 pub struct Word<T>(pub [T; WORD_LEN]);
 
 impl<AB: AirBuilder> AirVariable<AB> for Word<AB::Var> {
@@ -34,3 +37,16 @@ impl<T> IndexMut<usize> for Word<T> {
         &mut self.0[index]
     }
 }
+
+impl<F: Field> From<u32> for Word<F> {
+    fn from(value: u32) -> Self {
+        let inner = value
+            .to_le_bytes()
+            .iter()
+            .map(|v| F::from_canonical_u8(*v))
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap();
+        Word(inner)
+    }
+}

core/src/cpu/air.rs

@@ -7,26 +7,12 @@ use p3_field::AbstractField;
 use p3_field::{Field, PrimeField};
 use p3_matrix::MatrixRowSlices;
 use p3_util::indices_arr;
+use valida_derive::AlignedBorrow;
 
-#[derive(Debug, Clone, Copy)]
-pub struct CpuAir;
-
-/// An AIR table for memory accesses.
-#[derive(Debug, Clone)]
-pub struct CpuCols<T> {
-    /// The clock cycle value.
-    pub clk: T,
-    /// The program counter value.
-    pub pc: T,
-    /// The opcode for this cycle.
-    pub opcode: T,
-    /// The first operand for this instruction.
-    pub op_a: T,
-    /// The second operand for this instruction.
-    pub op_b: T,
-    /// The third operand for this instruction.
-    pub op_c: T,
-    // Whether op_b is an immediate value.
+#[derive(AlignedBorrow, Default)]
+#[repr(C)]
+pub struct OpcodeSelectors<T> {
+    // // Whether op_b is an immediate value.
     pub imm_b: T,
     // Whether op_c is an immediate value.
     pub imm_c: T,
@@ -46,24 +32,51 @@ pub struct CpuCols<T> {
     pub system_instruction: T,
     // Whether this is a multiply instruction.
     pub multiply_instruction: T,
-    // Selectors for load/store instructions and their types.
-    pub byte: Bool<T>,
-    pub half: Bool<T>,
-    pub word: Bool<T>,
-    pub unsigned: Bool<T>,
-    // TODO: we might need a selector for "MULSU" since no other instruction has "SU"
+    // // Selectors for load/store instructions and their types.
+    pub byte: T,
+    pub half: T,
+    pub word: T,
+    pub unsigned: T,
+    // // TODO: we might need a selector for "MULSU" since no other instruction has "SU"
     pub JALR: T,
     pub JAL: T,
     pub AUIPC: T,
+    // // Whether this instruction is reading from register A.
+    pub reg_a_read: T,
+}
+
+#[derive(AlignedBorrow, Default)]
+#[repr(C)]
+pub struct InstructionCols<T> {
+    // /// The opcode for this cycle.
+    pub opcode: T,
+    // /// The first operand for this instruction.
+    pub op_a: T,
+    // /// The second operand for this instruction.
+    pub op_b: T,
+    // /// The third operand for this instruction.
+    pub op_c: T,
+}
 
-    // Operand values, either from registers or immediate values.
+/// An AIR table for memory accesses.
+#[derive(AlignedBorrow, Default)]
+#[repr(C)]
+pub struct CpuCols<T> {
+    /// The clock cycle value.
+    pub clk: T,
+    // /// The program counter value.
+    pub pc: T,
+
+    // Columns related to the instruction.
+    pub instruction: InstructionCols<T>,
+    // Selectors for the opcode.
+    pub selectors: OpcodeSelectors<T>,
+
+    // // Operand values, either from registers or immediate values.
     pub op_a_val: Word<T>,
     pub op_b_val: Word<T>,
     pub op_c_val: Word<T>,
 
-    // Whether this instruction is reading from register A.
-    pub reg_a_read: T,
-
     // An addr that we are reading from or writing to.
     pub addr: Word<T>,
     // The associated memory value for `addr`.
@@ -81,28 +94,6 @@ const fn make_col_map() -> CpuCols<usize> {
     unsafe { transmute::<[usize; NUM_CPU_COLS], CpuCols<usize>>(indices_arr) }
 }
 
-impl<T> Borrow<CpuCols<T>> for [T] {
-    fn borrow(&self) -> &CpuCols<T> {
-        // TODO: Double check if this is correct & consider making asserts debug-only.
-        let (prefix, shorts, suffix) = unsafe { self.align_to::<CpuCols<T>>() };
-        assert!(prefix.is_empty(), "Data was not aligned");
-        assert!(suffix.is_empty(), "Data was not aligned");
-        assert_eq!(shorts.len(), 1);
-        &shorts[0]
-    }
-}
-
-impl<T> BorrowMut<CpuCols<T>> for [T] {
-    fn borrow_mut(&mut self) -> &mut CpuCols<T> {
-        // TODO: Double check if this is correct & consider making asserts debug-only.
-        let (prefix, shorts, suffix) = unsafe { self.align_to_mut::<CpuCols<T>>() };
-        assert!(prefix.is_empty(), "Data was not aligned");
-        assert!(suffix.is_empty(), "Data was not aligned");
-        assert_eq!(shorts.len(), 1);
-        &mut shorts[0]
-    }
-}
-
 impl<AB: AirBuilder> AirConstraint<AB> for CpuCols<AB::Var> {
     fn eval(&self, builder: &mut AB) {
         let main = builder.main();
@@ -120,47 +111,49 @@ impl<AB: AirBuilder> AirConstraint<AB> for CpuCols<AB::Var> {
         //// Constraint op_a_val, op_b_val, op_c_val
         // Constraint the op_b_val and op_c_val columns when imm_b and imm_c are true.
         builder
-            .when(local.imm_b)
-            .assert_eq(reduce::<AB>(local.op_b_val), local.op_b);
+            .when(local.selectors.imm_b)
+            .assert_eq(reduce::<AB>(local.op_b_val), local.instruction.op_b);
         builder
-            .when(local.imm_c)
-            .assert_eq(reduce::<AB>(local.op_c_val), local.op_c);
+            .when(local.selectors.imm_c)
+            .assert_eq(reduce::<AB>(local.op_c_val), local.instruction.op_c);
 
         // We only read from the first register if there is a store or branch instruction. In all other cases we write.
-        let reg_a_read =
-            local.store_instruction + local.branch_instruction + local.multiply_instruction;
+        let reg_a_read = local.selectors.store_instruction
+            + local.selectors.branch_instruction
+            + local.selectors.multiply_instruction;
 
         //// For r-type, i-type and multiply instructions, we must constraint by an "opcode-oracle" table
         // TODO: lookup (clk, op_a_val, op_b_val, op_c_val) in the "opcode-oracle" table with multiplicity (register_instruction + immediate_instruction + multiply_instruction)
 
         //// For branch instructions
         // TODO: lookup (clk, branch_cond_val, op_a_val, op_b_val) in the "branch" table with multiplicity branch_instruction
         // Increment the pc by 4 + op_c_val * branch_cond_val where we interpret the first result as a bool that it is.
-        builder.when(local.branch_instruction).assert_eq(
+        builder.when(local.selectors.branch_instruction).assert_eq(
             local.pc
                 + AB::F::from_canonical_u8(4)
                 + reduce::<AB>(local.op_c_val) * local.branch_cond_val.0[0],
             next.pc,
         );
 
         //// For jump instructions
-        builder.when(local.jump_instruction).assert_eq(
+        builder.when(local.selectors.jump_instruction).assert_eq(
             reduce::<AB>(local.op_a_val),
             local.pc + AB::F::from_canonical_u8(4),
         );
-        builder.when(local.JAL).assert_eq(
+        builder.when(local.selectors.JAL).assert_eq(
             local.pc + AB::F::from_canonical_u8(4) + reduce::<AB>(local.op_b_val),
             next.pc,
         );
-        builder
-            .when(local.JALR)
-            .assert_eq(reduce::<AB>(local.op_b_val) + local.op_c, next.pc);
+        builder.when(local.selectors.JALR).assert_eq(
+            reduce::<AB>(local.op_b_val) + local.instruction.op_c,
+            next.pc,
+        );
 
         //// For system instructions
 
         //// Upper immediate instructions
         // lookup(clk, op_c_val, imm, 12) in SLT table with multiplicity AUIPC
-        builder.when(local.AUIPC).assert_eq(
+        builder.when(local.selectors.AUIPC).assert_eq(
             reduce::<AB>(local.op_a_val),
             reduce::<AB>(local.op_c_val) + local.pc,
         );

core/src/cpu/mod.rs

@@ -8,4 +8,7 @@ pub struct CpuEvent {
     pub pc: u32,
     pub instruction: Instruction,
     pub operands: [u32; 3],
+    pub addr: Option<u32>,
+    pub memory_value: Option<u32>,
+    pub branch_condition: Option<bool>,
 }