Skip to content

Commit

Permalink
refactor: free some unlimited cells related to u32 cells
Browse files Browse the repository at this point in the history
  • Loading branch information
@junyu0312
junyu0312 committed Apr 14, 2024
1 parent 61416fd commit 3a4e8c0
Show file tree
Hide file tree
Showing 10 changed files with 144 additions and 127 deletions.
37 changes: 36 additions & 1 deletion crates/zkwasm/src/circuits/cell.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@ use halo2_proofs::plonk::Expression;
use halo2_proofs::plonk::VirtualCells;
use num_bigint::BigUint;

use crate::circuits::etable::EVENT_TABLE_ENTRY_ROWS;
use crate::circuits::utils::bn_to_field;
use crate::circuits::utils::Context;
use crate::constant_from;
use crate::nextn;

#[derive(Debug, Clone, Copy)]
Expand Down Expand Up @@ -79,6 +81,11 @@ impl<F: FieldExt> AllocatedU64Cell<F> {
#[derive(Debug, Clone, Copy)]
pub(crate) struct AllocatedU32Cell<F: FieldExt> {
pub(crate) u16_cells_le: [AllocatedU16Cell<F>; 2],
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct AllocatedU32PermutationCell<F: FieldExt> {
pub(crate) u16_cells_le: [AllocatedU16Cell<F>; 2],
pub(crate) u32_cell: AllocatedUnlimitedCell<F>,
}

Expand Down Expand Up @@ -160,13 +167,41 @@ impl<F: FieldExt> CellExpression<F> for AllocatedCommonRangeCell<F> {

impl<F: FieldExt> AllocatedU32Cell<F> {
pub(crate) fn expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.u32_cell.curr_expr(meta)
self.u16_cells_le[0].curr_expr(meta)
+ (self.u16_cells_le[1].curr_expr(meta) * constant_from!(1 << 16))
}

pub(crate) fn curr_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.expr(meta)
}

pub(crate) fn assign(&self, ctx: &mut Context<'_, F>, value: u32) -> Result<(), Error> {
for i in 0..2 {
self.u16_cells_le[i].assign(ctx, (((value >> (i * 16)) & 0xffffu32) as u64).into())?;
}

Ok(())
}
}

#[allow(dead_code)]
impl<F: FieldExt> AllocatedU32PermutationCell<F> {
pub(crate) fn expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.curr_expr(meta)
}

pub(crate) fn curr_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
self.u32_cell.expr(meta)
}

pub(crate) fn next_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
nextn!(
meta,
self.u32_cell.cell.col,
self.u32_cell.cell.rot + EVENT_TABLE_ENTRY_ROWS
)
}

pub(crate) fn assign(
&self,
ctx: &mut Context<'_, F>,
Expand Down
93 changes: 60 additions & 33 deletions crates/zkwasm/src/circuits/etable/allocator.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use super::AllocatedU32StateCell;
use super::EVENT_TABLE_ENTRY_ROWS;
use crate::circuits::bit_table::BitTableOp;
use crate::circuits::cell::*;
Expand All @@ -9,7 +10,6 @@ use crate::circuits::traits::ConfigureLookupTable;
use crate::circuits::utils::bit::BitColumn;
use crate::circuits::utils::common_range::CommonRangeColumn;
use crate::circuits::utils::u16::U16Column;
use crate::circuits::utils::u32_state::AllocatedU32StateCell;
use crate::circuits::utils::u8::U8Column;
use crate::circuits::Context;
use crate::circuits::Lookup;
Expand Down Expand Up @@ -66,24 +66,6 @@ impl_cell!(AllocatedCommonRangeCell);
impl_cell!(AllocatedUnlimitedCell);
impl_cell!(AllocatedJumpTableLookupCell);

impl<F: FieldExt> EventTableCellExpression<F> for AllocatedU32Cell<F> {
fn next_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
nextn!(
meta,
self.u32_cell.cell.col,
self.u32_cell.cell.rot + EVENT_TABLE_ENTRY_ROWS as i32
)
}

fn prev_expr(&self, meta: &mut VirtualCells<'_, F>) -> Expression<F> {
nextn!(
meta,
self.u32_cell.cell.col,
self.u32_cell.cell.rot - EVENT_TABLE_ENTRY_ROWS as i32
)
}
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct AllocatedJumpTableLookupCell<F: FieldExt> {
pub(crate) cell: AllocatedCell<F>,
Expand Down Expand Up @@ -184,18 +166,20 @@ pub(crate) enum EventTableCellType {

const BIT_COLUMNS: usize = 12;
const U8_COLUMNS: usize = 1;
const U32_CELLS: usize = if cfg!(feature = "continuation") {
12
const U32_CELLS: usize = 2;
const U32_PERMUTATION_CELLS: usize = if cfg!(feature = "continuation") {
10
} else {
2
0
};
const U64_CELLS: usize = 5;
const U16_COLUMNS: usize = U64_CELLS + (U32_CELLS / 2);
const U16_COLUMNS: usize =
U64_CELLS + ((U32_CELLS + U32_PERMUTATION_CELLS).next_multiple_of(2) / 2);
const COMMON_RANGE_COLUMNS: usize = if cfg!(feature = "continuation") { 4 } else { 6 };
const UNLIMITED_COLUMNS: usize = if cfg!(feature = "continuation") {
10
} else {
8
7
};
const MEMORY_TABLE_LOOKUP_COLUMNS: usize = 2;
const JUMP_TABLE_LOOKUP_COLUMNS: usize = 1;
Expand Down Expand Up @@ -229,6 +213,7 @@ impl<F: FieldExt> AllocatedBitTableLookupCells<F> {
pub(crate) struct AllocatorFreeCellsProfiler {
free_cells: BTreeMap<EventTableCellType, (usize, u32)>,
free_u32_cells: usize,
free_u32_permutation_cells: usize,
free_u64_cells: usize,
}

Expand All @@ -237,6 +222,7 @@ impl AllocatorFreeCellsProfiler {
Self {
free_cells: allocator.free_cells.clone(),
free_u32_cells: allocator.free_u32_cells.len(),
free_u32_permutation_cells: allocator.free_u32_permutation_cells.len(),
free_u64_cells: allocator.free_u64_cells.len(),
}
}
Expand All @@ -254,18 +240,40 @@ impl AllocatorFreeCellsProfiler {
}

self.free_u32_cells = usize::min(self.free_u32_cells, allocator.free_u32_cells.len());
self.free_u32_permutation_cells = usize::min(
self.free_u32_permutation_cells,
allocator.free_u32_permutation_cells.len(),
);
self.free_u64_cells = usize::min(self.free_u64_cells, allocator.free_u64_cells.len());
}

pub(crate) fn assert_no_free_cells<F: FieldExt>(&self, allocator: &EventTableCellAllocator<F>) {
for (t, (i, j)) in &self.free_cells {
let cols = allocator.all_cols.get(t).unwrap();

assert!(*i == cols.len() || (*i == cols.len() - 1 && *j > 0));
assert!(
*i == cols.len() || (*i == cols.len() - 1 && *j > 0),
"unused {:?} col should be removed: {}.",
t,
cols.len() - *i - (*j != 0) as usize
);
}

assert!(self.free_u32_cells == 0);
assert!(self.free_u64_cells == 0);
assert!(
self.free_u32_cells == 0,
"unused u32 cells should be removed: {:?}.",
self.free_u32_cells
);
assert!(
self.free_u32_permutation_cells == 0,
"unused u32 permutation cells should be removed: {:?}.",
self.free_u32_permutation_cells
);
assert!(
self.free_u64_cells == 0,
"unused u64 cells should be removed: {:?}.",
self.free_u64_cells
);
}
}

Expand All @@ -275,6 +283,7 @@ pub(crate) struct EventTableCellAllocator<F: FieldExt> {
pub(crate) free_cells: BTreeMap<EventTableCellType, (usize, u32)>,
all_cols: BTreeMap<EventTableCellType, Vec<Vec<Column<Advice>>>>,
free_u32_cells: Vec<AllocatedU32Cell<F>>,
free_u32_permutation_cells: Vec<AllocatedU32PermutationCell<F>>,
free_u64_cells: Vec<AllocatedU64Cell<F>>,
_mark: PhantomData<F>,
}
Expand All @@ -293,11 +302,17 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
}
}

pub(super) fn prepare_alloc_u32_cell(
pub(super) fn prepare_alloc_u32_cell(&mut self) -> AllocatedU32Cell<F> {
let u16_cells_le = [0; 2].map(|_| self.alloc_u16_cell());

AllocatedU32Cell { u16_cells_le }
}

pub(super) fn prepare_alloc_u32_permutation_cell(
&mut self,
meta: &mut ConstraintSystem<F>,
enable: impl Fn(&mut VirtualCells<'_, F>) -> Expression<F>,
) -> AllocatedU32Cell<F> {
) -> AllocatedU32PermutationCell<F> {
let u16_cells_le = [0; 2].map(|_| self.alloc_u16_cell());
let u32_cell = self.alloc_unlimited_cell();
meta.create_gate("c9. u32 decompose", |meta| {
Expand All @@ -312,7 +327,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
});
meta.enable_equality(u32_cell.cell.col);

AllocatedU32Cell {
AllocatedU32PermutationCell {
u16_cells_le,
u32_cell,
}
Expand Down Expand Up @@ -352,9 +367,14 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
) -> Self {
let mut allocator = Self::_new(meta, k, sel, rtable, mtable, jtable, cols);
for _ in 0..U32_CELLS {
let cell = allocator.prepare_alloc_u32_cell(meta, |meta| fixed_curr!(meta, sel));
let cell = allocator.prepare_alloc_u32_cell();
allocator.free_u32_cells.push(cell);
}
for _ in 0..U32_PERMUTATION_CELLS {
let cell =
allocator.prepare_alloc_u32_permutation_cell(meta, |meta| fixed_curr!(meta, sel));
allocator.free_u32_permutation_cells.push(cell);
}
for _ in 0..U64_CELLS {
let cell = allocator.prepare_alloc_u64_cell(meta, |meta| fixed_curr!(meta, sel));
allocator.free_u64_cells.push(cell);
Expand Down Expand Up @@ -462,6 +482,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
.into_iter(),
),
free_u32_cells: vec![],
free_u32_permutation_cells: vec![],
free_u64_cells: vec![],
_mark: PhantomData,
}
Expand Down Expand Up @@ -522,7 +543,7 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
pub(crate) fn alloc_u32_state_cell(&mut self) -> AllocatedU32StateCell<F> {
cfg_if::cfg_if! {
if #[cfg(feature = "continuation")] {
self.alloc_u32_cell()
self.alloc_u32_permutation_cell()
} else {
self.alloc_common_range_cell()
}
Expand Down Expand Up @@ -731,11 +752,17 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
cell
}

#[allow(dead_code)]
pub(crate) fn alloc_u32_cell(&mut self) -> AllocatedU32Cell<F> {
self.free_u32_cells.pop().expect("no more free u32 cells")
}

#[allow(dead_code)]
pub(crate) fn alloc_u32_permutation_cell(&mut self) -> AllocatedU32PermutationCell<F> {
self.free_u32_permutation_cells
.pop()
.expect("no more free u32 permutation cells")
}

pub(crate) fn alloc_u64_cell(&mut self) -> AllocatedU64Cell<F> {
self.free_u64_cells.pop().expect("no more free u64 cells")
}
Expand Down
29 changes: 0 additions & 29 deletions crates/zkwasm/src/circuits/etable/assign.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,35 +124,6 @@ impl<F: FieldExt> EventTableChip<F> {
#[cfg(feature = "continuation")]
let jops = assign_common_range_advice!(jops_cell, state.jops);

// if let Some(assigned_pre_initialization_state) = assigned_pre_initialization_state {
// macro_rules! constrain_equal {
// ($field:ident) => {
// ctx.region.constrain_equal(
// $field.cell(),
// assigned_pre_initialization_state.$field.cell(),
// )?;
// };
// }

// constrain_equal!(eid);
// constrain_equal!(fid);
// constrain_equal!(iid);
// constrain_equal!(sp);
// constrain_equal!(frame_id);

// constrain_equal!(host_public_inputs);
// constrain_equal!(context_in_index);
// constrain_equal!(context_out_index);
// constrain_equal!(external_host_call_call_index);

// constrain_equal!(initial_memory_pages);
// constrain_equal!(maximal_memory_pages);

// #[cfg(feature = "continuation")]
// constrain_equal!(jops);
// }

// The context will be stepped by EVENT_TABLE_ENTRY_ROWS.
ctx.step(EVENT_TABLE_ENTRY_ROWS as usize);

Ok(InitializationState {
Expand Down
6 changes: 5 additions & 1 deletion crates/zkwasm/src/circuits/etable/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ use super::rtable::RangeTableConfig;
use super::traits::ConfigureLookupTable;
use super::utils::step_status::StepStatus;
use super::utils::table_entry::EventTableEntryWithMemoryInfo;
use super::utils::u32_state::AllocatedU32StateCell;
use super::utils::Context;
use crate::circuits::etable::op_configure::op_bin::BinConfigBuilder;
use crate::circuits::etable::op_configure::op_bin_bit::BinBitConfigBuilder;
Expand Down Expand Up @@ -68,6 +67,11 @@ mod op_configure;
pub(crate) mod allocator;
pub(crate) mod constraint_builder;

#[cfg(feature = "continuation")]
type AllocatedU32StateCell<F> = AllocatedU32PermutationCell<F>;
#[cfg(not(feature = "continuation"))]
type AllocatedU32StateCell<F> = AllocatedCommonRangeCell<F>;

pub(crate) const EVENT_TABLE_ENTRY_ROWS: i32 = 4;
pub(crate) const OP_CAPABILITY: usize = 32;

Expand Down
Loading

0 comments on commit 3a4e8c0

Please sign in to comment.