From 64cc1bed4690077e88b6cda7ebdb4095a7401b16 Mon Sep 17 00:00:00 2001 From: Lun-Kai Hsu Date: Mon, 22 Jul 2024 16:42:02 -0700 Subject: [PATCH] wip, on hold for now --- Cargo.toml | 1 + chips/src/bits/bit_decompose/air.rs | 56 ++++++++++ chips/src/bits/bit_decompose/columns.rs | 14 +++ chips/src/bits/bit_decompose/mod.rs | 2 + chips/src/lib.rs | 3 + chips/src/utils.rs | 40 ++++++++ sha256-air/Cargo.toml | 13 +++ sha256-air/src/air.rs | 131 ++++++++++++++++++++++++ sha256-air/src/columns.rs | 66 ++++++++++++ sha256-air/src/lib.rs | 2 + 10 files changed, 328 insertions(+) create mode 100644 chips/src/bits/bit_decompose/air.rs create mode 100644 chips/src/bits/bit_decompose/columns.rs create mode 100644 chips/src/bits/bit_decompose/mod.rs create mode 100644 sha256-air/Cargo.toml create mode 100644 sha256-air/src/air.rs create mode 100644 sha256-air/src/columns.rs create mode 100644 sha256-air/src/lib.rs diff --git a/Cargo.toml b/Cargo.toml index f428d4ae50..dd916422dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "vm", "vm/bin", "poseidon2-air", + "sha256-air", ] resolver = "2" diff --git a/chips/src/bits/bit_decompose/air.rs b/chips/src/bits/bit_decompose/air.rs new file mode 100644 index 0000000000..a796f3603e --- /dev/null +++ b/chips/src/bits/bit_decompose/air.rs @@ -0,0 +1,56 @@ +use super::columns::BitDecomposeCols; +use crate::{ + sub_chip::{AirConfig, SubAir}, + utils::Word32, +}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{AbstractField, Field}; +use p3_matrix::Matrix; +use std::borrow::Borrow; + +pub struct BitDecomposeAir { + // TOOD: do we need this? + bus_index: usize, +} + +impl BitDecomposeAir { + pub fn new(bus_index: usize) -> Self { + Self { bus_index } + } +} + +impl AirConfig for BitDecomposeAir { + type Cols = BitDecomposeCols; +} + +impl BaseAir for BitDecomposeAir { + fn width(&self) -> usize { + N + 1 + } +} + +impl Air for BitDecomposeAir { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0); + let local: &[AB::Var] = (*local).borrow(); + let cols = BitDecomposeCols::::from_slice(local); + + SubAir::eval(self, builder, cols, ()); + } +} + +impl SubAir for BitDecomposeAir { + type IoView = BitDecomposeCols; + type AuxView = (); + + fn eval(&self, builder: &mut AB, io: Self::IoView, _aux: Self::AuxView) { + let mut from_bits = AB::Expr::zero(); + for (i, &bit) in io.x_bits.iter().enumerate() { + from_bits += bit * AB::Expr::from_canonical_u32(1 << i); + } + // How to make this a function of Word32? + let from_x = io.x.0[0] * AB::Expr::from_canonical_u32(1 << 16) + io.x.0[1]; + builder.assert_eq(from_bits, from_x); + } +} diff --git a/chips/src/bits/bit_decompose/columns.rs b/chips/src/bits/bit_decompose/columns.rs new file mode 100644 index 0000000000..a7768eaa34 --- /dev/null +++ b/chips/src/bits/bit_decompose/columns.rs @@ -0,0 +1,14 @@ +use crate::utils::Word32; + +pub struct BitDecomposeCols { + pub x: Word32, + pub x_bits: Vec, +} + +impl BitDecomposeCols { + pub fn from_slice(slice: &[T]) -> Self { + let x = Word32([slice[0].clone(), slice[1].clone()]); + let x_bits = slice[2..].to_vec(); + Self { x, x_bits } + } +} diff --git a/chips/src/bits/bit_decompose/mod.rs b/chips/src/bits/bit_decompose/mod.rs new file mode 100644 index 0000000000..602b9b608e --- /dev/null +++ b/chips/src/bits/bit_decompose/mod.rs @@ -0,0 +1,2 @@ +pub mod air; +pub mod columns; diff --git a/chips/src/lib.rs b/chips/src/lib.rs index 4c9778974d..c49fd9a785 100644 --- a/chips/src/lib.rs +++ b/chips/src/lib.rs @@ -1,4 +1,7 @@ pub mod assert_sorted; +pub mod bits { + pub mod bit_decompose; +} pub mod common; pub mod execution_air; pub mod group_by; diff --git a/chips/src/utils.rs b/chips/src/utils.rs index 212984c116..e028a597e7 100644 --- a/chips/src/utils.rs +++ b/chips/src/utils.rs @@ -1,5 +1,8 @@ +use afs_derive::AlignedBorrow; +use p3_air::Air; use p3_air::{AirBuilder, VirtualPairCol}; use p3_field::Field; +use p3_field::PrimeField32; // TODO: Ideally upstream PrimeField implements From pub trait FieldFrom { @@ -45,3 +48,40 @@ pub fn or(a: AB::Expr, b: AB::Expr) -> AB::Expr { pub fn implies(a: AB::Expr, b: AB::Expr) -> AB::Expr { or::(AB::Expr::one() - a, b) } + +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +pub struct Word32(pub [T; 2]); +impl Word32 +where + AB: AirBuilder, + T: AB::Var, +{ + pub fn get_value(&self) -> AB::Expr { + let upper = AB::Expr::from(self.0[0]); + let lower = AB::Expr::from(self.0[1]); + lower + (upper * AB::Expr::from_canonical_u64(1 << 16)) + } +} + +// impl Word32 { +// pub fn get_value>(&self) -> AB::Expr { +// let lower = AB::Expr::from(self.0[1]); +// let upper = AB::Expr::from(self.0[0]); +// lower + (upper * AB::Expr::from_canonical_u64(1 << 16)) +// // AB::Expr::from_canonical_u16(0) +// } +// } + +// impl Word32 { +// fn concat_u8(x: u8, y: u8) -> u16 { +// ((x as u16) << 8) | (y as u16) +// } +// // TODO: verify from and to 32. probably wrong now. Just make it compile first. +// pub fn from_u32(x: u32) -> Self { +// let bytes = x.to_le_bytes(); // TODO: should this be le or be? +// Word32([ +// T::from_canonical_u16(Self::concat_u8(bytes[0], bytes[1])), +// T::from_canonical_u16(Self::concat_u8(bytes[2], bytes[3])), +// ]) +// } +// } diff --git a/sha256-air/Cargo.toml b/sha256-air/Cargo.toml new file mode 100644 index 0000000000..9d80b3151f --- /dev/null +++ b/sha256-air/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "sha256-air" +version.workspace = true +authors.workspace = true +edition.workspace = true + +[dependencies] +p3-air = { workspace = true } +p3-field = { workspace = true } +p3-matrix = { workspace = true } + +afs-chips = { path = "../chips" } +afs-derive = { path = "../derive" } \ No newline at end of file diff --git a/sha256-air/src/air.rs b/sha256-air/src/air.rs new file mode 100644 index 0000000000..5f002e2d31 --- /dev/null +++ b/sha256-air/src/air.rs @@ -0,0 +1,131 @@ +use crate::columns::{Sha256AuxCols, Sha256IoCols, Word}; + +use super::columns::Sha256Cols; +use afs_chips::bits::bit_decompose::air::BitDecomposeAir; +use afs_chips::bits::bit_decompose::columns::BitDecomposeCols; +use afs_chips::sub_chip::{AirConfig, SubAir}; +use p3_air::{Air, AirBuilder, BaseAir, FilteredAirBuilder}; +use p3_field::{AbstractField, Field}; +use p3_matrix::Matrix; +use std::borrow::Borrow; + +// TODO: update +pub const NUM_COLUMNS: usize = 100; + +pub const H: [u32; 8] = [ + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19, +]; + +pub struct Sha256Air { + bus_index: usize, +} + +impl AirConfig for Sha256Air { + type Cols = Sha256Cols; +} + +impl BaseAir for Sha256Air { + fn width(&self) -> usize { + NUM_COLUMNS + } +} + +impl Air for Sha256Air { + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let (local, next) = (main.row_slice(0), main.row_slice(1)); + let local: &Sha256Cols = (*local).borrow(); + let next: &Sha256Cols = (*next).borrow(); + + SubAir::::eval(self, builder, local.io, (local.aux, next.aux)); + } +} + +impl SubAir for Sha256Air { + type IoView = Sha256IoCols; + // local, next + type AuxView = (Sha256AuxCols, Sha256AuxCols); + + fn eval(&self, builder: &mut AB, io: Self::IoView, aux: Self::AuxView) { + self.eval_init(builder, &aux.0); + self.eval_compression(builder, &io, &aux.0, &aux.1); + } +} + +impl Sha256Air { + fn eval_io( + &self, + builder: &mut AB, + io: &Sha256IoCols, + aux: &Sha256AuxCols, + ) { + todo!() + } + + fn eval_init(&self, builder: &mut AB, aux: &Sha256AuxCols) { + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h0, H[0]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h1, H[1]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h2, H[2]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h3, H[3]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h4, H[4]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h5, H[5]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h6, H[6]); + self.assert_word_equal_u32(&mut builder.when_first_row(), &aux.h7, H[7]); + } + + fn eval_compression( + &self, + builder: &mut AB, + local_io: &Sha256IoCols, + local_aux: &Sha256AuxCols, + next_aux: &Sha256AuxCols, + ) { + // Working variables set to hashes at the start of each block. + let mut block_start = builder.when(local_aux.is_block_start); + self.assert_words_equal(&mut block_start, &local_aux.a, &next_aux.h0); + self.assert_words_equal(&mut block_start, &local_aux.b, &next_aux.h1); + self.assert_words_equal(&mut block_start, &local_aux.c, &next_aux.h2); + self.assert_words_equal(&mut block_start, &local_aux.d, &next_aux.h3); + self.assert_words_equal(&mut block_start, &local_aux.e, &next_aux.h4); + self.assert_words_equal(&mut block_start, &local_aux.f, &next_aux.h5); + self.assert_words_equal(&mut block_start, &local_aux.g, &next_aux.h6); + self.assert_words_equal(&mut block_start, &local_aux.h, &next_aux.h7); + + let mut is_main_loop = builder.when_ne(local_aux.row_idx, AB::F::zero()); + let bit_decompose_air = BitDecomposeAir::<32>::new(self.bus_index); + + let x = std::iter::once(local_aux.e).chain(local_aux.e_bits); + SubAir::eval( + &bit_decompose_air, + builder, + BitDecomposeCols::from_slice(std::iter::once(local_aux.e).chain(local_aux.e_bits)), + (), + ); + + // Update the working variables for next round. + + // Update hashes at the end of each block. + } + + fn assert_word_equal_u32( + &self, + builder: &mut FilteredAirBuilder<'_, AB>, + word: &Word, + x: u32, + ) { + let high = AB::F::from_canonical_u16((x >> 16) as u16); + let low = AB::F::from_canonical_u16((x & 0xffff) as u16); + builder.assert_eq(word.0[0], high); + builder.assert_eq(word.0[1], low); + } + + fn assert_words_equal( + &self, + builder: &mut FilteredAirBuilder<'_, AB>, + w1: &Word, + w2: &Word, + ) { + builder.assert_eq(w1.0[0], w2.0[0]); + builder.assert_eq(w1.0[1], w2.0[1]); + } +} diff --git a/sha256-air/src/columns.rs b/sha256-air/src/columns.rs new file mode 100644 index 0000000000..1f1e61c4cd --- /dev/null +++ b/sha256-air/src/columns.rs @@ -0,0 +1,66 @@ +use afs_chips::utils::Word32; +use afs_derive::AlignedBorrow; + +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +pub struct Sha256Cols { + pub io: Sha256IoCols, + pub aux: Sha256AuxCols, +} + +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +pub struct Sha256IoCols { + // Two field elements to hold one word (32bits) of input + pub input: Word32, + // 16 field elements to hold one 256bit output + pub output: [T; 16], +} + +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +pub struct Sha256AuxCols { + // The message schedule at the current round. + pub w: Word32, + // Sliding window of the previous 16 w. + pub w_prev: [T; 32], + // The round constant. + pub k: Word32, + // The working variables a~h. + pub a: Word32, + pub b: Word32, + pub c: Word32, + pub d: Word32, + pub e: Word32, + pub f: Word32, + pub g: Word32, + pub h: Word32, + // The bit operation variables to compute w. + pub s0: Word32, + pub s1: Word32, + // The variables to update working variables each round. + pub cs1: Word32, // S1 in pseudocode + pub ch: Word32, + pub temp1: Word32, + pub cs0: Word32, // S0 in pseudocode + pub maj: Word32, + pub temp2: Word32, + // TODO: use Vec? Fix for now so I can use aligned borrow and skip implementing from slice. + pub e_bits: [T; 32], + // Hash values. + pub h0: Word32, + pub h1: Word32, + pub h2: Word32, + pub h3: Word32, + pub h4: Word32, + pub h5: Word32, + pub h6: Word32, + pub h7: Word32, + // Control flow columns. + pub idx: T, + pub row_idx: T, + pub is_block_start: T, // Bool, basically whether row_idx == 0, but equality is hard to condition on. +} + +impl Sha256Cols { + pub fn flatten(&self) -> Vec { + todo!() + } +} diff --git a/sha256-air/src/lib.rs b/sha256-air/src/lib.rs new file mode 100644 index 0000000000..602b9b608e --- /dev/null +++ b/sha256-air/src/lib.rs @@ -0,0 +1,2 @@ +pub mod air; +pub mod columns;