Feat/Aggregation VM STARK

Cargo.toml

@@ -5,6 +5,7 @@ members = [
     "test-utils",
     "derive",
     "logical-interface",
+    "vm",
 ]
 resolver = "2"
 

chips/src/is_equal/columns.rs

@@ -21,6 +21,17 @@ pub struct IsEqualAuxCols<T> {
     pub inv: T,
 }
 
+impl<T: Clone> IsEqualAuxCols<T> {
+    pub fn from_slice(slc: &[T]) -> IsEqualAuxCols<T> {
+        IsEqualAuxCols {
+            inv: slc[0].clone(),
+        }
+    }
+
+    pub fn flatten(&self) -> Vec<T> {
+        vec![self.inv.clone()]
+    }
+}
 impl<T: Clone> IsEqualCols<T> {
     pub const fn new(x: T, y: T, is_equal: T, inv: T) -> IsEqualCols<T> {
         IsEqualCols {
@@ -38,6 +49,15 @@ impl<T: Clone> IsEqualCols<T> {
         IsEqualCols::new(x, y, is_equal, inv)
     }
 
+    pub fn flatten(&self) -> Vec<T> {
+        vec![
+            self.io.x.clone(),
+            self.io.y.clone(),
+            self.io.is_equal.clone(),
+            self.aux.inv.clone(),
+        ]
+    }
+
     pub fn get_width() -> usize {
         NUM_COLS
     }

vm/Cargo.toml

@@ -0,0 +1,49 @@
+[package]
+name = "stark-vm"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+description = "Aggregation VM STARK"
+
+[dependencies]
+p3-air = { workspace = true }
+p3-baby-bear = { workspace = true, optional = true }
+p3-commit = { workspace = true }
+p3-field = { workspace = true }
+p3-keccak = { workspace = true }
+p3-keccak-air = { workspace = true }
+p3-matrix = { workspace = true }
+p3-maybe-rayon = { workspace = true }
+p3-symmetric = { workspace = true }
+p3-util = { workspace = true }
+p3-uni-stark = { workspace = true }
+rand = "0.8"
+
+afs-chips = { path = "../chips" }
+afs-stark-backend = { path = "../stark-backend" }
+afs-derive = { path = "../derive" }
+afs-test-utils = { path = "../test-utils" }
+parking_lot = "0.12.2"
+tracing = "0.1.40"
+itertools = "0.13.0"
+getset = "0.1.2"
+derive-new = "0.5.1"
+
+[dev-dependencies]
+p3-uni-stark = { workspace = true }
+p3-baby-bear = { workspace = true }
+p3-challenger = { workspace = true }
+p3-dft = { workspace = true }
+p3-fri = { workspace = true }
+p3-goldilocks = { workspace = true }
+p3-keccak = { workspace = true }
+p3-merkle-tree = { workspace = true }
+p3-poseidon2 = { workspace = true }
+p3-symmetric = { workspace = true }
+tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
+tracing-forest = { version = "0.1.6", features = ["ansi", "smallvec"] }
+test-case = "3.3.1"
+
+[features]
+default = ["test-traits"]
+test-traits = ["p3-baby-bear"]

vm/src/au/air.rs

@@ -0,0 +1,59 @@
+use std::borrow::Borrow;
+
+use p3_air::{Air, AirBuilder, BaseAir};
+use p3_field::{AbstractField, Field};
+use p3_matrix::Matrix;
+
+use super::{columns::FieldArithmeticCols, FieldArithmeticAir};
+use afs_chips::sub_chip::AirConfig;
+
+impl AirConfig for FieldArithmeticAir {
+    type Cols<T> = FieldArithmeticCols<T>;
+}
+
+impl<F: Field> BaseAir<F> for FieldArithmeticAir {
+    fn width(&self) -> usize {
+        FieldArithmeticCols::<F>::NUM_COLS
+    }
+}
+
+impl<AB: AirBuilder> Air<AB> for FieldArithmeticAir {
+    fn eval(&self, builder: &mut AB) {
+        let main = builder.main();
+
+        let local = main.row_slice(0);
+        let au_cols: &FieldArithmeticCols<_> = (*local).borrow();
+
+        let FieldArithmeticCols { io, aux } = au_cols;
+
+        builder.assert_bool(aux.opcode_lo);
+        builder.assert_bool(aux.opcode_hi);
+
+        builder.assert_eq(
+            io.opcode,
+            aux.opcode_lo
+                + aux.opcode_hi * AB::Expr::two()
+                + AB::F::from_canonical_u8(FieldArithmeticAir::BASE_OP),
+        );
+
+        builder.assert_eq(
+            aux.is_mul,
+            aux.opcode_hi * (AB::Expr::one() - aux.opcode_lo),
+        );
+        builder.assert_eq(aux.is_div, aux.opcode_hi * aux.opcode_lo);
+
+        builder.assert_eq(aux.product, io.x * io.y);
+        builder.assert_eq(aux.quotient * io.y, io.x);
+        builder.assert_eq(
+            au_cols.aux.sum_or_diff,
+            io.x + io.y - AB::Expr::two() * aux.opcode_lo * io.y,
+        );
+
+        builder.assert_eq(
+            io.z,
+            aux.is_mul * aux.product
+                + aux.is_div * aux.quotient
+                + aux.sum_or_diff * (AB::Expr::one() - aux.opcode_hi),
+        );
+    }
+}

vm/src/au/bridge.rs

@@ -0,0 +1,18 @@
+use afs_stark_backend::interaction::{AirBridge, Interaction};
+use p3_air::VirtualPairCol;
+use p3_field::Field;
+
+use crate::au::{columns::FieldArithmeticCols, FieldArithmeticAir};
+
+/// Receives all IO columns from another chip on bus 2 (FieldArithmeticAir::BUS_INDEX).
+impl<T: Field> AirBridge<T> for FieldArithmeticAir {
+    fn receives(&self) -> Vec<Interaction<T>> {
+        vec![Interaction {
+            fields: (0..FieldArithmeticCols::<T>::NUM_IO_COLS)
+                .map(VirtualPairCol::single_main)
+                .collect(),
+            count: VirtualPairCol::one(),
+            argument_index: Self::BUS_INDEX,
+        }]
+    }
+}

vm/src/au/columns.rs

@@ -0,0 +1,117 @@
+use crate::cpu::OpCode;
+use afs_derive::AlignedBorrow;
+use p3_field::Field;
+
+use crate::au::FieldArithmeticAir;
+
+/// Columns for field arithmetic chip.
+///
+/// Four IO columns for opcode, x, y, result.
+/// Seven aux columns for interpreting opcode, evaluating indicators, and explicit computations.
+#[derive(AlignedBorrow)]
+pub struct FieldArithmeticCols<T> {
+    pub io: FieldArithmeticIOCols<T>,
+    pub aux: FieldArithmeticAuxCols<T>,
+}
+
+pub struct FieldArithmeticIOCols<T> {
+    pub opcode: T,
+    pub x: T,
+    pub y: T,
+    pub z: T,
+}
+
+pub struct FieldArithmeticAuxCols<T> {
+    pub opcode_lo: T,
+    pub opcode_hi: T,
+    pub is_mul: T,
+    pub is_div: T,
+    pub sum_or_diff: T,
+    pub product: T,
+    pub quotient: T,
+}
+
+impl<T> FieldArithmeticCols<T>
+where
+    T: Field,
+{
+    pub const NUM_COLS: usize = 11;
+    pub const NUM_IO_COLS: usize = 4;
+    pub const NUM_AUX_COLS: usize = 6;
+
+    /// Constructs a new set of columns (including auxiliary columns) given inputs.
+    pub fn new(op: OpCode, x: T, y: T) -> Self {
+        let opcode = op as u32;
+        let opcode_value = opcode - FieldArithmeticAir::BASE_OP as u32;
+        let opcode_lo_u32 = opcode_value % 2;
+        let opcode_hi_u32 = opcode_value / 2;
+        let opcode_lo = T::from_canonical_u32(opcode_lo_u32);
+        let opcode_hi = T::from_canonical_u32(opcode_hi_u32);
+        let is_div = T::from_bool(op == OpCode::FDIV);
+        let is_mul = T::from_bool(op == OpCode::FMUL);
+        let sum_or_diff = x + y - T::two() * opcode_lo * y;
+        let product = x * y;
+        let quotient = if y == T::zero() {
+            T::zero()
+        } else {
+            x * y.inverse()
+        };
+        let z = is_mul * product + is_div * quotient + (T::one() - opcode_hi) * sum_or_diff;
+
+        Self {
+            io: FieldArithmeticIOCols {
+                opcode: T::from_canonical_u32(opcode),
+                x,
+                y,
+                z,
+            },
+            aux: FieldArithmeticAuxCols {
+                opcode_lo,
+                opcode_hi,
+                is_mul,
+                is_div,
+                sum_or_diff,
+                product,
+                quotient,
+            },
+        }
+    }
+
+    pub fn get_width() -> usize {
+        FieldArithmeticIOCols::<T>::get_width() + FieldArithmeticAuxCols::<T>::get_width()
+    }
+
+    pub fn flatten(&self) -> Vec<T> {
+        let mut result = self.io.flatten();
+        result.extend(self.aux.flatten());
+        result
+    }
+}
+
+impl<T: Field> FieldArithmeticIOCols<T> {
+    pub fn get_width() -> usize {
+        4
+    }
+
+    pub fn flatten(&self) -> Vec<T> {
+        vec![self.opcode, self.x, self.y, self.z]
+    }
+}
+
+impl<T: Field> FieldArithmeticAuxCols<T> {
+    pub fn get_width() -> usize {
+        7
+    }
+
+    pub fn flatten(&self) -> Vec<T> {
+        vec![
+            self.opcode_lo,
+            self.opcode_hi,
+            self.is_mul,
+            self.is_div,
+            self.sum_or_diff,
+            self.product,
+            self.quotient,
+        ]
+    }
+}

vm/src/au/mod.rs

@@ -0,0 +1,69 @@
+use super::cpu::trace::ArithmeticOperation;
+use crate::cpu::OpCode;
+use p3_field::Field;
+
+#[cfg(test)]
+pub mod tests;
+
+pub mod air;
+pub mod bridge;
+pub mod columns;
+pub mod trace;
+
+/// Field arithmetic chip.
+///
+/// Carries information about opcodes (currently 5..=8) and bus index (currently 2).
+#[derive(Default, Clone, Copy)]
+pub struct FieldArithmeticAir {}
+
+impl FieldArithmeticAir {
+    pub const BASE_OP: u8 = 6;
+    pub const BUS_INDEX: usize = 2;
+
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    /// Evaluates given opcode using given operands.
+    ///
+    /// Returns None for non-arithmetic operations.
+    pub fn solve<T: Field>(op: OpCode, operands: (T, T)) -> Option<T> {
+        match op {
+            OpCode::LOADW => None,
+            OpCode::STOREW => None,
+            OpCode::JAL => None,
+            OpCode::BEQ => None,
+            OpCode::BNE => None,
+            OpCode::TERMINATE => None,
+
+            OpCode::FADD => Some(operands.0 + operands.1),
+            OpCode::FSUB => Some(operands.0 - operands.1),
+            OpCode::FMUL => Some(operands.0 * operands.1),
+            OpCode::FDIV => Some(operands.0 / operands.1),
+        }
+    }
+
+    /// Vectorized solve<>
+    pub fn solve_all<T: Field>(ops: Vec<OpCode>, operands: Vec<(T, T)>) -> Vec<T> {
+        ops.iter()
+            .zip(operands.iter())
+            .filter_map(|(op, operand)| Self::solve::<T>(*op, *operand))
+            .collect()
+    }
+
+    /// Converts vectorized opcodes and operands into vectorized ArithmeticOperations.
+    pub fn request<T: Field>(
+        ops: Vec<OpCode>,
+        operands: Vec<(T, T)>,
+    ) -> Vec<ArithmeticOperation<T>> {
+        ops.iter()
+            .zip(operands.iter())
+            .map(|(op, operand)| ArithmeticOperation {
+                opcode: *op,
+                operand1: operand.0,
+                operand2: operand.1,
+                result: Self::solve::<T>(*op, *operand).unwrap(),
+            })
+            .collect()
+    }
+}