Skip to content

Commit

Permalink
feat: fixed shapes in precompiles (#1476)
Browse files Browse the repository at this point in the history
  • Loading branch information
tamirhemo authored Sep 4, 2024
1 parent 9d4c565 commit 055d84c
Show file tree
Hide file tree
Showing 11 changed files with 178 additions and 145 deletions.
38 changes: 21 additions & 17 deletions crates/core/machine/src/syscall/precompiles/edwards/ed_add.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use crate::{
operations::field::{
field_den::FieldDenCols, field_inner_product::FieldInnerProductCols, field_op::FieldOpCols,
},
utils::{limbs_from_prev_access, pad_rows},
utils::{limbs_from_prev_access, pad_rows_fixed},
};

pub const NUM_ED_ADD_COLS: usize = size_of::<EdAddAssignCols<u8>>();
Expand Down Expand Up @@ -137,22 +137,26 @@ impl<F: PrimeField32, E: EllipticCurve + EdwardsParameters> MachineAir<F> for Ed
})
.collect::<Vec<_>>();

pad_rows(&mut rows, || {
let mut row = [F::zero(); NUM_ED_ADD_COLS];
let cols: &mut EdAddAssignCols<F> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = [F::zero(); NUM_ED_ADD_COLS];
let cols: &mut EdAddAssignCols<F> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ use typenum::U32;
use crate::{
memory::{MemoryReadCols, MemoryWriteCols},
operations::field::{field_op::FieldOpCols, field_sqrt::FieldSqrtCols, range::FieldLtCols},
utils::{limbs_from_access, limbs_from_prev_access, pad_rows},
utils::{limbs_from_access, limbs_from_prev_access, pad_rows_fixed},
};

pub const NUM_ED_DECOMPRESS_COLS: usize = size_of::<EdDecompressCols<u8>>();
Expand Down Expand Up @@ -273,13 +273,17 @@ impl<F: PrimeField32, E: EdwardsParameters> MachineAir<F> for EdDecompressChip<E
rows.push(row);
}

pad_rows(&mut rows, || {
let mut row = [F::zero(); NUM_ED_DECOMPRESS_COLS];
let cols: &mut EdDecompressCols<F> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
cols.populate_field_ops::<E>(&mut vec![], 0, 0, &zero);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = [F::zero(); NUM_ED_DECOMPRESS_COLS];
let cols: &mut EdDecompressCols<F> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
cols.populate_field_ops::<E>(&mut vec![], 0, 0, &zero);
row
},
input.fixed_log2_rows::<F, _>(self),
);

let mut trace = RowMajorMatrix::new(
rows.into_iter().flatten().collect::<Vec<_>>(),
Expand Down
38 changes: 21 additions & 17 deletions crates/core/machine/src/syscall/precompiles/fptower/fp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use sp1_stark::air::{BaseAirBuilder, MachineAir, Polynomial, SP1AirBuilder};
use crate::{
memory::{value_as_limbs, MemoryReadCols, MemoryWriteCols},
operations::field::field_op::FieldOpCols,
utils::{limbs_from_prev_access, pad_rows, words_to_bytes_le_vec},
utils::{limbs_from_prev_access, pad_rows_fixed, words_to_bytes_le_vec},
};

pub const fn num_fp_cols<P: FpOpField>() -> usize {
Expand Down Expand Up @@ -148,22 +148,26 @@ impl<F: PrimeField32, P: FpOpField> MachineAir<F> for FpOpChip<P> {

output.add_byte_lookup_events(new_byte_lookup_events);

pad_rows(&mut rows, || {
let mut row = vec![F::zero(); num_fp_cols::<P>()];
let cols: &mut FpOpCols<F, P> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
cols.is_add = F::from_canonical_u8(1);
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero,
FieldOperation::Add,
);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = vec![F::zero(); num_fp_cols::<P>()];
let cols: &mut FpOpCols<F, P> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
cols.is_add = F::from_canonical_u8(1);
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero,
FieldOperation::Add,
);
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace =
Expand Down
42 changes: 23 additions & 19 deletions crates/core/machine/src/syscall/precompiles/fptower/fp2_addsub.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use typenum::Unsigned;
use crate::{
memory::{value_as_limbs, MemoryReadCols, MemoryWriteCols},
operations::field::field_op::FieldOpCols,
utils::{limbs_from_prev_access, pad_rows, words_to_bytes_le_vec},
utils::{limbs_from_prev_access, pad_rows_fixed, words_to_bytes_le_vec},
};

pub const fn num_fp2_addsub_cols<P: FpOpField>() -> usize {
Expand Down Expand Up @@ -154,24 +154,28 @@ impl<F: PrimeField32, P: FpOpField> MachineAir<F> for Fp2AddSubAssignChip<P> {

output.add_byte_lookup_events(new_byte_lookup_events);

pad_rows(&mut rows, || {
let mut row = vec![F::zero(); num_fp2_addsub_cols::<P>()];
let cols: &mut Fp2AddSubAssignCols<F, P> = row.as_mut_slice().borrow_mut();
cols.is_add = F::one();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
FieldOperation::Add,
);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = vec![F::zero(); num_fp2_addsub_cols::<P>()];
let cols: &mut Fp2AddSubAssignCols<F, P> = row.as_mut_slice().borrow_mut();
cols.is_add = F::one();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
FieldOperation::Add,
);
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace = RowMajorMatrix::new(
Expand Down
38 changes: 21 additions & 17 deletions crates/core/machine/src/syscall/precompiles/fptower/fp2_mul.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use typenum::Unsigned;
use crate::{
memory::{value_as_limbs, MemoryReadCols, MemoryWriteCols},
operations::field::field_op::FieldOpCols,
utils::{limbs_from_prev_access, pad_rows, words_to_bytes_le_vec},
utils::{limbs_from_prev_access, pad_rows_fixed, words_to_bytes_le_vec},
};

pub const fn num_fp2_mul_cols<P: FieldParameters + NumWords>() -> usize {
Expand Down Expand Up @@ -206,22 +206,26 @@ impl<F: PrimeField32, P: FpOpField> MachineAir<F> for Fp2MulAssignChip<P> {

output.add_byte_lookup_events(new_byte_lookup_events);

pad_rows(&mut rows, || {
let mut row = vec![F::zero(); num_fp2_mul_cols::<P>()];
let cols: &mut Fp2MulAssignCols<F, P> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = vec![F::zero(); num_fp2_mul_cols::<P>()];
let cols: &mut Fp2MulAssignCols<F, P> = row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace = RowMajorMatrix::new(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use super::{
columns::{ShaCompressCols, NUM_SHA_COMPRESS_COLS},
ShaCompressChip, SHA_COMPRESS_K,
};
use crate::utils::pad_rows;
use crate::utils::pad_rows_fixed;

impl<F: PrimeField32> MachineAir<F> for ShaCompressChip {
type Record = ExecutionRecord;
Expand All @@ -42,7 +42,11 @@ impl<F: PrimeField32> MachineAir<F> for ShaCompressChip {

let num_real_rows = rows.len();

pad_rows(&mut rows, || [F::zero(); NUM_SHA_COMPRESS_COLS]);
pad_rows_fixed(
&mut rows,
|| [F::zero(); NUM_SHA_COMPRESS_COLS],
input.fixed_log2_rows::<F, _>(self),
);

// Set the octet_num and octect columns for the padded rows.
let mut octet_num = 0;
Expand Down
22 changes: 13 additions & 9 deletions crates/core/machine/src/syscall/precompiles/uint256/air.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::{
air::MemoryAirBuilder,
operations::{field::range::FieldLtCols, IsZeroOperation},
utils::{
limbs_from_access, limbs_from_prev_access, pad_rows, words_to_bytes_le,
limbs_from_access, limbs_from_prev_access, pad_rows_fixed, words_to_bytes_le,
words_to_bytes_le_vec,
},
};
Expand Down Expand Up @@ -199,16 +199,20 @@ impl<F: PrimeField32> MachineAir<F> for Uint256MulChip {
output.append(&mut record);
}

pad_rows(&mut rows, || {
let mut row: [F; NUM_COLS] = [F::zero(); NUM_COLS];
let cols: &mut Uint256MulCols<F> = row.as_mut_slice().borrow_mut();
pad_rows_fixed(
&mut rows,
|| {
let mut row: [F; NUM_COLS] = [F::zero(); NUM_COLS];
let cols: &mut Uint256MulCols<F> = row.as_mut_slice().borrow_mut();

let x = BigUint::zero();
let y = BigUint::zero();
cols.output.populate(&mut vec![], 0, 0, &x, &y, FieldOperation::Mul);
let x = BigUint::zero();
let y = BigUint::zero();
cols.output.populate(&mut vec![], 0, 0, &x, &y, FieldOperation::Mul);

row
});
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use typenum::Unsigned;
use crate::{
memory::{MemoryCols, MemoryReadCols, MemoryWriteCols},
operations::field::field_op::FieldOpCols,
utils::{limbs_from_prev_access, pad_rows},
utils::{limbs_from_prev_access, pad_rows_fixed},
};

pub const fn num_weierstrass_add_cols<P: FieldParameters + NumWords>() -> usize {
Expand Down Expand Up @@ -257,23 +257,27 @@ impl<F: PrimeField32, E: EllipticCurve + WeierstrassParameters> MachineAir<F>
}
output.add_byte_lookup_events(new_byte_lookup_events);

pad_rows(&mut rows, || {
let mut row = vec![F::zero(); num_weierstrass_add_cols::<E::BaseField>()];
let cols: &mut WeierstrassAddAssignCols<F, E::BaseField> =
row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
});
pad_rows_fixed(
&mut rows,
|| {
let mut row = vec![F::zero(); num_weierstrass_add_cols::<E::BaseField>()];
let cols: &mut WeierstrassAddAssignCols<F, E::BaseField> =
row.as_mut_slice().borrow_mut();
let zero = BigUint::zero();
Self::populate_field_ops(
&mut vec![],
0,
0,
cols,
zero.clone(),
zero.clone(),
zero.clone(),
zero,
);
row
},
input.fixed_log2_rows::<F, _>(self),
);

// Convert the trace to a row major matrix.
let mut trace = RowMajorMatrix::new(
Expand Down
Loading

0 comments on commit 055d84c

Please sign in to comment.