Skip to content

Commit

Permalink
chore: make FpOperations use PrimeField32 to clean up type incons…
Browse files Browse the repository at this point in the history 
…istency hacks (#203)
  • Loading branch information
@hidenori-shinohara
hidenori-shinohara authored Feb 13, 2024
1 parent ecc9b92 commit db58ea3
Show file tree
Hide file tree
Showing 11 changed files with 320 additions and 344 deletions.
64 changes: 30 additions & 34 deletions core/src/operations/field/fp_den.rs → core/src/operations/field/field_den.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use super::params::Limbs;
use super::params::NUM_WITNESS_LIMBS;
use super::params::{convert_polynomial, convert_vec, Limbs};
use super::util::{compute_root_quotient_and_shift, split_u16_limbs_to_u8_limbs};
use super::util_air::eval_field_operation;
use crate::air::CurtaAirBuilder;
Expand All @@ -8,38 +8,34 @@ use crate::utils::ec::field::FieldParameters;
use core::borrow::{Borrow, BorrowMut};
use core::mem::size_of;
use num::BigUint;
use p3_baby_bear::BabyBear;
use p3_field::Field;
use p3_field::PrimeField32;
use std::fmt::Debug;
use valida_derive::AlignedBorrow;

// a / (1 + b) if sign
// a/ -b if !sign
/// A set of columns to compute `FpDen(a, b)` where a, b are field elements.
/// A set of columns to compute `FieldDen(a, b)` where `a`, `b` are field elements.
///
/// `a / (1 + b)` if `sign`
/// `a / -b` if `!sign`
///
/// Right now the number of limbs is assumed to be a constant, although this could be macro-ed
/// or made generic in the future.
#[derive(Debug, Clone, AlignedBorrow)]
#[repr(C)]
pub struct FpDenCols<T> {
pub struct FieldDenCols<T> {
/// The result of `a den b`, where a, b are field elements
pub result: Limbs<T>,
pub(crate) carry: Limbs<T>,
pub(crate) witness_low: [T; NUM_WITNESS_LIMBS],
pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
}

impl<F: Field> FpDenCols<F> {
impl<F: PrimeField32> FieldDenCols<F> {
pub fn populate<P: FieldParameters>(
&mut self,
a: &BigUint,
b: &BigUint,
sign: bool,
) -> BigUint {
/// TODO: This operation relies on `F` being a PrimeField32, but our traits do not
/// support that. This is a hack, since we always use BabyBear, to get around that, but
/// all operations using "PF" should use "F" in the future.
type PF = BabyBear;

let p = P::modulus();
let minus_b_int = &p - b;
let b_signed = if sign { b.clone() } else { minus_b_int };
Expand All @@ -59,11 +55,11 @@ impl<F: Field> FpDenCols<F> {
debug_assert!(carry < p);
debug_assert_eq!(&carry * &p, &equation_lhs - &equation_rhs);

let p_a: Polynomial<PF> = P::to_limbs_field::<PF>(a).into();
let p_b: Polynomial<PF> = P::to_limbs_field::<PF>(b).into();
let p_p: Polynomial<PF> = P::to_limbs_field::<PF>(&p).into();
let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(&result).into();
let p_carry: Polynomial<PF> = P::to_limbs_field::<PF>(&carry).into();
let p_a: Polynomial<F> = P::to_limbs_field::<F>(a).into();
let p_b: Polynomial<F> = P::to_limbs_field::<F>(b).into();
let p_p: Polynomial<F> = P::to_limbs_field::<F>(&p).into();
let p_result: Polynomial<F> = P::to_limbs_field::<F>(&result).into();
let p_carry: Polynomial<F> = P::to_limbs_field::<F>(&carry).into();

// Compute the vanishing polynomial.
let vanishing_poly = if sign {
Expand All @@ -80,16 +76,16 @@ impl<F: Field> FpDenCols<F> {
);
let (p_witness_low, p_witness_high) = split_u16_limbs_to_u8_limbs(&p_witness);

self.result = convert_polynomial(p_result);
self.carry = convert_polynomial(p_carry);
self.witness_low = convert_vec(p_witness_low).try_into().unwrap();
self.witness_high = convert_vec(p_witness_high).try_into().unwrap();
self.result = p_result.into();
self.carry = p_carry.into();
self.witness_low = p_witness_low.try_into().unwrap();
self.witness_high = p_witness_high.try_into().unwrap();

result
}
}

impl<V: Copy> FpDenCols<V> {
impl<V: Copy> FieldDenCols<V> {
#[allow(unused_variables)]
pub fn eval<AB: CurtaAirBuilder<Var = V>, P: FieldParameters>(
&self,
Expand Down Expand Up @@ -132,9 +128,9 @@ impl<V: Copy> FpDenCols<V> {
mod tests {
use num::BigUint;
use p3_air::BaseAir;
use p3_field::Field;
use p3_field::{Field, PrimeField32};

use super::{FpDenCols, Limbs};
use super::{FieldDenCols, Limbs};
use crate::air::MachineAir;
use crate::utils::ec::edwards::ed25519::Ed25519BaseField;
use crate::utils::ec::field::FieldParameters;
Expand All @@ -154,17 +150,17 @@ mod tests {
pub struct TestCols<T> {
pub a: Limbs<T>,
pub b: Limbs<T>,
pub a_den_b: FpDenCols<T>,
pub a_den_b: FieldDenCols<T>,
}

pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8>>();

struct FpDenChip<P: FieldParameters> {
struct FieldDenChip<P: FieldParameters> {
pub sign: bool,
pub _phantom: std::marker::PhantomData<P>,
}

impl<P: FieldParameters> FpDenChip<P> {
impl<P: FieldParameters> FieldDenChip<P> {
pub fn new(sign: bool) -> Self {
Self {
sign,
Expand All @@ -173,9 +169,9 @@ mod tests {
}
}

impl<F: Field, P: FieldParameters> MachineAir<F> for FpDenChip<P> {
impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for FieldDenChip<P> {
fn name(&self) -> String {
"FpDen".to_string()
"FieldDen".to_string()
}

fn shard(&self, _: &ExecutionRecord, _: &mut Vec<ExecutionRecord>) {}
Expand Down Expand Up @@ -227,13 +223,13 @@ mod tests {
}
}

impl<F: Field, P: FieldParameters> BaseAir<F> for FpDenChip<P> {
impl<F: Field, P: FieldParameters> BaseAir<F> for FieldDenChip<P> {
fn width(&self) -> usize {
NUM_TEST_COLS
}
}

impl<AB, P: FieldParameters> Air<AB> for FpDenChip<P>
impl<AB, P: FieldParameters> Air<AB> for FieldDenChip<P>
where
AB: CurtaAirBuilder,
{
Expand All @@ -254,7 +250,7 @@ mod tests {
#[test]
fn generate_trace() {
let mut shard = ExecutionRecord::default();
let chip: FpDenChip<Ed25519BaseField> = FpDenChip::new(true);
let chip: FieldDenChip<Ed25519BaseField> = FieldDenChip::new(true);
let trace: RowMajorMatrix<BabyBear> = chip.generate_trace(&mut shard);
println!("{:?}", trace.values)
}
Expand All @@ -266,7 +262,7 @@ mod tests {

let mut shard = ExecutionRecord::default();

let chip: FpDenChip<Ed25519BaseField> = FpDenChip::new(true);
let chip: FieldDenChip<Ed25519BaseField> = FieldDenChip::new(true);
let trace: RowMajorMatrix<BabyBear> = chip.generate_trace(&mut shard);
// This it to test that the proof DOESN'T work if messed up.
// let row = trace.row_mut(0);
Expand Down
69 changes: 29 additions & 40 deletions .../src/operations/field/fp_inner_product.rs → ...c/operations/field/field_inner_product.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use super::params::Limbs;
use super::params::NUM_WITNESS_LIMBS;
use super::params::{convert_polynomial, convert_vec, Limbs};
use super::util::{compute_root_quotient_and_shift, split_u16_limbs_to_u8_limbs};
use super::util_air::eval_field_operation;
use crate::air::CurtaAirBuilder;
Expand All @@ -9,40 +9,29 @@ use core::borrow::{Borrow, BorrowMut};
use core::mem::size_of;
use num::BigUint;
use num::Zero;
use p3_baby_bear::BabyBear;
use p3_field::AbstractField;
use p3_field::Field;
use p3_field::{AbstractField, PrimeField32};
use std::fmt::Debug;
use valida_derive::AlignedBorrow;

/// A set of columns to compute `FpInnerProduct(Vec<a>, Vec<b>)` where a, b are field elements.
/// A set of columns to compute `FieldInnerProduct(Vec<a>, Vec<b>)` where a, b are field elements.
/// Right now the number of limbs is assumed to be a constant, although this could be macro-ed
/// or made generic in the future.
#[derive(Debug, Clone, AlignedBorrow)]
#[repr(C)]
pub struct FpInnerProductCols<T> {
pub struct FieldInnerProductCols<T> {
/// The result of `a inner product b`, where a, b are field elements
pub result: Limbs<T>,
pub(crate) carry: Limbs<T>,
pub(crate) witness_low: [T; NUM_WITNESS_LIMBS],
pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
}

impl<F: Field> FpInnerProductCols<F> {
impl<F: PrimeField32> FieldInnerProductCols<F> {
pub fn populate<P: FieldParameters>(&mut self, a: &[BigUint], b: &[BigUint]) -> BigUint {
/// TODO: This operation relies on `F` being a PrimeField32, but our traits do not
/// support that. This is a hack, since we always use BabyBear, to get around that, but
/// all operations using "PF" should use "F" in the future.
type PF = BabyBear;

let p_a_vec: Vec<Polynomial<PF>> = a
.iter()
.map(|x| P::to_limbs_field::<PF>(x).into())
.collect();
let p_b_vec: Vec<Polynomial<PF>> = b
.iter()
.map(|x| P::to_limbs_field::<PF>(x).into())
.collect();
let p_a_vec: Vec<Polynomial<F>> =
a.iter().map(|x| P::to_limbs_field::<F>(x).into()).collect();
let p_b_vec: Vec<Polynomial<F>> =
b.iter().map(|x| P::to_limbs_field::<F>(x).into()).collect();

let modulus = &P::modulus();
let inner_product = a
Expand All @@ -56,15 +45,15 @@ impl<F: Field> FpInnerProductCols<F> {
assert!(carry < &(2u32 * modulus));
assert_eq!(carry * modulus, inner_product - result);

let p_modulus: Polynomial<PF> = P::to_limbs_field::<PF>(modulus).into();
let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(result).into();
let p_carry: Polynomial<PF> = P::to_limbs_field::<PF>(carry).into();
let p_modulus: Polynomial<F> = P::to_limbs_field::<F>(modulus).into();
let p_result: Polynomial<F> = P::to_limbs_field::<F>(result).into();
let p_carry: Polynomial<F> = P::to_limbs_field::<F>(carry).into();

// Compute the vanishing polynomial.
let p_inner_product = p_a_vec
.into_iter()
.zip(p_b_vec)
.fold(Polynomial::<PF>::new(vec![PF::zero()]), |acc, (c, d)| {
.fold(Polynomial::<F>::new(vec![F::zero()]), |acc, (c, d)| {
acc + &c * &d
});
let p_vanishing = p_inner_product - &p_result - &p_carry * &p_modulus;
Expand All @@ -77,16 +66,16 @@ impl<F: Field> FpInnerProductCols<F> {
);
let (p_witness_low, p_witness_high) = split_u16_limbs_to_u8_limbs(&p_witness);

self.result = convert_polynomial(p_result);
self.carry = convert_polynomial(p_carry);
self.witness_low = convert_vec(p_witness_low).try_into().unwrap();
self.witness_high = convert_vec(p_witness_high).try_into().unwrap();
self.result = p_result.into();
self.carry = p_carry.into();
self.witness_low = p_witness_low.try_into().unwrap();
self.witness_high = p_witness_high.try_into().unwrap();

result.clone()
}
}

impl<V: Copy> FpInnerProductCols<V> {
impl<V: Copy> FieldInnerProductCols<V> {
#[allow(unused_variables)]
pub fn eval<AB: CurtaAirBuilder<Var = V>, P: FieldParameters>(
&self,
Expand Down Expand Up @@ -127,9 +116,9 @@ impl<V: Copy> FpInnerProductCols<V> {
mod tests {
use num::BigUint;
use p3_air::BaseAir;
use p3_field::Field;
use p3_field::{Field, PrimeField32};

use super::{FpInnerProductCols, Limbs};
use super::{FieldInnerProductCols, Limbs};
use crate::air::MachineAir;
use crate::utils::ec::edwards::ed25519::Ed25519BaseField;
use crate::utils::ec::field::FieldParameters;
Expand All @@ -150,26 +139,26 @@ mod tests {
pub struct TestCols<T> {
pub a: [Limbs<T>; 1],
pub b: [Limbs<T>; 1],
pub a_ip_b: FpInnerProductCols<T>,
pub a_ip_b: FieldInnerProductCols<T>,
}

pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8>>();

struct FpIpChip<P: FieldParameters> {
struct FieldIpChip<P: FieldParameters> {
pub _phantom: std::marker::PhantomData<P>,
}

impl<P: FieldParameters> FpIpChip<P> {
impl<P: FieldParameters> FieldIpChip<P> {
pub fn new() -> Self {
Self {
_phantom: std::marker::PhantomData,
}
}
}

impl<F: Field, P: FieldParameters> MachineAir<F> for FpIpChip<P> {
impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for FieldIpChip<P> {
fn name(&self) -> String {
"FpInnerProduct".to_string()
"FieldInnerProduct".to_string()
}

fn shard(&self, _: &ExecutionRecord, _: &mut Vec<ExecutionRecord>) {}
Expand Down Expand Up @@ -219,13 +208,13 @@ mod tests {
}
}

impl<F: Field, P: FieldParameters> BaseAir<F> for FpIpChip<P> {
impl<F: Field, P: FieldParameters> BaseAir<F> for FieldIpChip<P> {
fn width(&self) -> usize {
NUM_TEST_COLS
}
}

impl<AB, P: FieldParameters> Air<AB> for FpIpChip<P>
impl<AB, P: FieldParameters> Air<AB> for FieldIpChip<P>
where
AB: CurtaAirBuilder,
{
Expand All @@ -245,7 +234,7 @@ mod tests {
#[test]
fn generate_trace() {
let mut shard = ExecutionRecord::default();
let chip: FpIpChip<Ed25519BaseField> = FpIpChip::new();
let chip: FieldIpChip<Ed25519BaseField> = FieldIpChip::new();
let trace: RowMajorMatrix<BabyBear> = chip.generate_trace(&mut shard);
println!("{:?}", trace.values)
}
Expand All @@ -257,7 +246,7 @@ mod tests {

let mut shard = ExecutionRecord::default();

let chip: FpIpChip<Ed25519BaseField> = FpIpChip::new();
let chip: FieldIpChip<Ed25519BaseField> = FieldIpChip::new();
let trace: RowMajorMatrix<BabyBear> = chip.generate_trace(&mut shard);
let proof = prove::<BabyBearPoseidon2, _>(&config, &chip, &mut challenger, trace);

Expand Down
Loading

0 comments on commit db58ea3

Please sign in to comment.