Skip to content

Commit

Permalink
perf: optimize assignment (#263)
Browse files Browse the repository at this point in the history
* perf: minor optimiazation for assignment

* perf: minor optimiazation for image table assignment

* perf: optimize memory_table

* opt post image table assignment

* opt mtable assignment

* perf: optimize etable assignment

* perf: remove assertion for perf

* perf: optimize image table assignment

* perf: optimize image table assignment

* perf: optimize etable assignment

* minor optimization for etable and image table

* remove mimalloc

* remove config in test_cli script
  • Loading branch information
junyu0312 authored Jun 14, 2024
1 parent b3eefcd commit f855b26
Show file tree
Hide file tree
Showing 54 changed files with 598 additions and 395 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,11 @@ anyhow = { version = "1.0.68", features = ["backtrace"] }
bincode = "1.3.3"
cfg-if = "1.0.0"
halo2_proofs = { git = "https://github.com/DelphinusLab/halo2-gpu-specific.git", default-features = true }
num-traits = "0.2.15"
parity-wasm = { version = "0.42.0", features = ["sign_ext"] }
rayon = "1.8.0"
regex = "1.10.2"
static_assertions = "1.1.0"
wasmi = { path = "third-party/wasmi" }
zkwasm-host-circuits = { git = "https://github.com/DelphinusLab/zkWasm-host-circuits.git" }
circuits-batcher = { git = "https://github.com/DelphinusLab/continuation-batcher.git" }
Expand Down
2 changes: 1 addition & 1 deletion crates/host/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ hex = "0.4.3"
log = "0.4.17"
num-integer = "0.1"
num-bigint = { version = "0.4", features = ["rand"] }
num-traits = "0.2.15"
wabt = "0.10.0"
lazy_static = "1.4.0"
rand = "0.8.4"
Expand All @@ -27,6 +26,7 @@ ff = "0.12"
sha2 = "0.10.6"
anyhow.workspace = true
halo2_proofs.workspace = true
num-traits.workspace = true
parity-wasm.workspace = true
poseidon.workspace = true
wasmi.workspace = true
Expand Down
4 changes: 3 additions & 1 deletion crates/specs/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,11 @@ strum_macros = "0.24.1"
bincode.workspace = true
cfg-if.workspace = true
halo2_proofs.workspace = true
parity-wasm.workspace = true
hex = "0.4.3"
num-traits.workspace = true
parity-wasm.workspace = true
rayon.workspace = true
static_assertions.workspace = true

[features]
default = []
Expand Down
34 changes: 25 additions & 9 deletions crates/specs/src/encode/image_table.rs
Original file line number Diff line number Diff line change
@@ -1,28 +1,44 @@
use num_bigint::BigUint;
use num_bigint::ToBigUint;
use static_assertions::const_assert;

use crate::encode::br_table::BR_TABLE_ENCODE_BOUNDARY;
use crate::encode::init_memory_table::INIT_MEMORY_ENCODE_BOUNDARY;
use crate::encode::instruction_table::INSTRUCTION_ENCODE_BOUNDARY;

use super::FromBn;

const CLASS_SHIFT: u32 = 224;

#[derive(Clone, Copy)]
pub enum ImageTableEncoder {
Instruction = 1,
BrTable = 2,
InitMemory = 3,
}

impl ImageTableEncoder {
pub fn encode<T: FromBn>(&self, data: T) -> T {
const CLASS_SHIFT: u32 = 224;
const_assert!(INSTRUCTION_ENCODE_BOUNDARY <= CLASS_SHIFT);
const_assert!(BR_TABLE_ENCODE_BOUNDARY <= CLASS_SHIFT);
const_assert!(INIT_MEMORY_ENCODE_BOUNDARY <= CLASS_SHIFT);

assert!(INSTRUCTION_ENCODE_BOUNDARY <= CLASS_SHIFT);
assert!(BR_TABLE_ENCODE_BOUNDARY <= CLASS_SHIFT);
assert!(INIT_MEMORY_ENCODE_BOUNDARY <= CLASS_SHIFT);
lazy_static! {
static ref INSTRUCTION_TAG: BigUint = (ImageTableEncoder::Instruction as u64)
.to_biguint()
.unwrap()
<< CLASS_SHIFT;
static ref BR_TABLE_TAG: BigUint =
(ImageTableEncoder::BrTable as u64).to_biguint().unwrap() << CLASS_SHIFT;
static ref INIT_MEMORY_TAG: BigUint =
(ImageTableEncoder::InitMemory as u64).to_biguint().unwrap() << CLASS_SHIFT;
}

T::from_bn(&(*self as u64).to_biguint().unwrap())
* T::from_bn(&(1u64.to_biguint().unwrap() << CLASS_SHIFT))
+ data
impl ImageTableEncoder {
#[inline(always)]
pub fn encode<T: FromBn>(&self, data: T) -> T {
match self {
ImageTableEncoder::Instruction => T::from_bn(&INSTRUCTION_TAG) + data,
ImageTableEncoder::BrTable => T::from_bn(&BR_TABLE_TAG) + data,
ImageTableEncoder::InitMemory => T::from_bn(&INIT_MEMORY_TAG) + data,
}
}
}
53 changes: 37 additions & 16 deletions crates/specs/src/encode/init_memory_table.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,36 @@
use halo2_proofs::arithmetic::FieldExt;
use num_bigint::BigUint;
use num_bigint::ToBigUint;
use num_traits::identities::One;
use static_assertions::const_assert;
use static_assertions::const_assert_eq;

use super::bn_to_field;
use super::FromBn;
use crate::imtable::InitMemoryTableEntry;

pub(crate) const INIT_MEMORY_ENCODE_BOUNDARY: u32 = 224;
pub const INIT_MEMORY_ENCODE_BOUNDARY: u32 = 224;
pub const MEMORY_ADDRESS_OFFSET: u32 = 97;

const LTYPE_SHIFT: u32 = OFFSET_SHIFT + u32::BITS;
const OFFSET_SHIFT: u32 = IS_MUTABLE_SHIFT + 1;
const IS_MUTABLE_SHIFT: u32 = EID_OFFSET_SHIFT + u32::BITS;
const EID_OFFSET_SHIFT: u32 = VALUE_SHIFT + u64::BITS;
const VALUE_SHIFT: u32 = 0;

const_assert_eq!(OFFSET_SHIFT, MEMORY_ADDRESS_OFFSET);
const_assert!(LTYPE_SHIFT + 8 <= INIT_MEMORY_ENCODE_BOUNDARY);

lazy_static! {
static ref IS_MUTABLE_OFFSET: BigUint = BigUint::one() << IS_MUTABLE_SHIFT;
static ref EID_OFFSET: BigUint = BigUint::one() << EID_OFFSET_SHIFT;
static ref LTYPE_OFFSET: BigUint = BigUint::one() << LTYPE_SHIFT;
static ref MEMORY_OFFSET_OFFSET: BigUint = BigUint::one() << OFFSET_SHIFT;
static ref MEMORY_LOCATION_OFFSET: BigUint = BigUint::one() << 32;
}

pub fn encode_init_memory_table_address<T: FromBn>(location_type: T, offset: T) -> T {
location_type * T::from_bn(&(1u64.to_biguint().unwrap() << 32)) + offset
location_type * T::from_bn(&MEMORY_LOCATION_OFFSET) + offset
}

pub fn encode_init_memory_table_entry<T: FromBn>(
Expand All @@ -18,25 +40,24 @@ pub fn encode_init_memory_table_entry<T: FromBn>(
eid: T,
value: T,
) -> T {
const LTYPE_SHIFT: u32 = OFFSET_SHIFT + u32::BITS;
const OFFSET_SHIFT: u32 = IS_MUTABLE_SHIFT + 1;
const IS_MUTABLE_SHIFT: u32 = EID_OFFSET_SHIFT + u32::BITS;
const EID_OFFSET_SHIFT: u32 = VALUE_SHIFT + u64::BITS;
const VALUE_SHIFT: u32 = 0;

assert_eq!(OFFSET_SHIFT, MEMORY_ADDRESS_OFFSET);
assert!(LTYPE_SHIFT + 8 <= INIT_MEMORY_ENCODE_BOUNDARY);
let encode =
is_mutable * T::from_bn(&IS_MUTABLE_OFFSET) + eid * T::from_bn(&EID_OFFSET) + value;

let encode = is_mutable * T::from_bn(&(1u64.to_biguint().unwrap() << IS_MUTABLE_SHIFT))
+ eid * T::from_bn(&(1u64.to_biguint().unwrap() << EID_OFFSET_SHIFT))
+ value;
if cfg!(feature = "continuation") {
encode
} else {
ltype * T::from_bn(&LTYPE_OFFSET) + offset * T::from_bn(&MEMORY_OFFSET_OFFSET) + encode
}
}

pub fn init_memory_table_entry_encode_update_offset<T: FieldExt>(
encode: T,
offset: impl FnOnce() -> T,
) -> T {
if cfg!(feature = "continuation") {
encode
} else {
ltype * T::from_bn(&(1u64.to_biguint().unwrap() << LTYPE_SHIFT))
+ offset * T::from_bn(&(1u64.to_biguint().unwrap() << OFFSET_SHIFT))
+ encode
encode + offset() * bn_to_field::<T>(&MEMORY_OFFSET_OFFSET)
}
}

Expand Down
29 changes: 19 additions & 10 deletions crates/specs/src/encode/memory_table.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,26 @@
use num_bigint::ToBigUint;
use num_bigint::BigUint;
use num_traits::One;
use static_assertions::const_assert;

use super::FromBn;
use crate::encode::COMMON_RANGE_OFFSET;

pub fn encode_memory_table_entry<T: FromBn>(offset: T, location_type: T, is_i32: T) -> T {
const END_SHIFT: u32 = OFFSET_SHIFT + COMMON_RANGE_OFFSET;
const OFFSET_SHIFT: u32 = LOCATION_TYPE_SHIFT + COMMON_RANGE_OFFSET;
const LOCATION_TYPE_SHIFT: u32 = IS_I32_SHIFT + 1;
const IS_I32_SHIFT: u32 = 0;
const _END_SHIFT: u32 = OFFSET_SHIFT + COMMON_RANGE_OFFSET;
const OFFSET_SHIFT: u32 = LOCATION_TYPE_SHIFT + COMMON_RANGE_OFFSET;
const LOCATION_TYPE_SHIFT: u32 = IS_I32_SHIFT + 1;
const IS_I32_SHIFT: u32 = 0;

const_assert!(_END_SHIFT < 240);

assert!(END_SHIFT < 240);
lazy_static! {
pub static ref MEMORY_TABLE_ENTRY_OFFSET: BigUint = BigUint::one() << OFFSET_SHIFT;
pub static ref MEMORY_TABLE_ENTRY_LOCATION_TYPE: BigUint =
BigUint::one() << LOCATION_TYPE_SHIFT;
pub static ref MEMORY_TABLE_ENTRY_IS_I32: BigUint = BigUint::one() << IS_I32_SHIFT;
}

offset * T::from_bn(&(1u64.to_biguint().unwrap() << OFFSET_SHIFT))
+ location_type * T::from_bn(&(1u64.to_biguint().unwrap() << LOCATION_TYPE_SHIFT))
+ is_i32 * T::from_bn(&(1u64.to_biguint().unwrap() << IS_I32_SHIFT))
pub fn encode_memory_table_entry<T: FromBn>(offset: T, location_type: T, is_i32: T) -> T {
offset * T::from_bn(&MEMORY_TABLE_ENTRY_OFFSET)
+ location_type * T::from_bn(&MEMORY_TABLE_ENTRY_LOCATION_TYPE)
+ is_i32 * T::from_bn(&MEMORY_TABLE_ENTRY_IS_I32)
}
12 changes: 4 additions & 8 deletions crates/specs/src/imtable.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::collections::BTreeMap;
use std::collections::HashMap;

use crate::mtable::LocationType;
use crate::mtable::VarType;
Expand All @@ -16,12 +16,12 @@ pub struct InitMemoryTableEntry {
pub eid: u32,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct InitMemoryTable(pub BTreeMap<(LocationType, u32), InitMemoryTableEntry>);
#[derive(Default, Debug)]
pub struct InitMemoryTable(pub HashMap<(LocationType, u32), InitMemoryTableEntry>);

impl InitMemoryTable {
pub fn new(entries: Vec<InitMemoryTableEntry>) -> Self {
let mut map = BTreeMap::new();
let mut map = HashMap::new();

entries.into_iter().for_each(|entry| {
map.insert((entry.ltype, entry.offset), entry);
Expand All @@ -30,10 +30,6 @@ impl InitMemoryTable {
Self(map)
}

pub fn entries(&self) -> &BTreeMap<(LocationType, u32), InitMemoryTableEntry> {
&self.0
}

pub fn to_string(&self) -> String {
serde_json::to_string(&self.0).unwrap()
}
Expand Down
6 changes: 4 additions & 2 deletions crates/specs/src/slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use std::sync::Arc;

use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;
use rayon::prelude::ParallelSliceMut;

use crate::brtable::BrTable;
use crate::brtable::ElemTable;
Expand Down Expand Up @@ -39,7 +40,7 @@ impl From<FrameTable> for FrameTableSlice {

impl FrameTableSlice {
pub fn build_returned_lookup_mapping(&self) -> HashMap<(u32, u32), bool> {
let mut lookup_table = HashMap::new();
let mut lookup_table = HashMap::with_capacity(self.called.len() + self.inherited.0.len());
for entry in self.called.iter() {
lookup_table.insert((entry.0.frame_id, entry.0.callee_fid), entry.0.returned);
}
Expand Down Expand Up @@ -114,6 +115,7 @@ impl Slice {
.collect::<Vec<Vec<_>>>()
.concat();

// Use a set to deduplicate
let mut set = HashSet::<MemoryTableEntry>::default();

memory_entries.iter().for_each(|entry| {
Expand Down Expand Up @@ -145,7 +147,7 @@ impl Slice {

memory_entries.append(&mut set.into_iter().collect());

memory_entries.sort_by_key(|item| (item.ltype, item.offset, item.eid));
memory_entries.par_sort_unstable_by_key(|item| (item.ltype, item.offset, item.eid));

MTable::new(memory_entries)
}
Expand Down
2 changes: 1 addition & 1 deletion crates/zkwasm/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ hex = "0.4.3"
log = "0.4.17"
num-integer = "0.1"
num-bigint = { version = "0.4", features = ["rand"] }
num-traits = "0.2.15"
wabt = "0.10.0"
lazy_static = "1.4.0"
rand = "0.8.4"
Expand All @@ -28,6 +27,7 @@ sha2 = "0.10.6"
anyhow.workspace = true
cfg-if.workspace = true
halo2_proofs.workspace = true
num-traits.workspace = true
parity-wasm.workspace = true
rayon.workspace = true
regex.workspace = true
Expand Down
15 changes: 0 additions & 15 deletions crates/zkwasm/src/circuits/cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,6 @@ macro_rules! define_cell {
ctx: &mut Context<'_, F>,
value: F,
) -> Result<AssignedCell<F, F>, Error> {
assert!(
value <= $limit,
"assigned value {:?} exceeds the limit {:?}",
value,
$limit
);

self.cell.assign(ctx, value)
}
}
Expand All @@ -145,7 +138,6 @@ define_cell!(AllocatedUnlimitedCell, -F::one());
#[derive(Debug, Clone, Copy)]
pub(crate) struct AllocatedCommonRangeCell<F: FieldExt> {
pub(crate) cell: AllocatedCell<F>,
pub(crate) upper_bound: F,
}

impl<F: FieldExt> CellExpression<F> for AllocatedCommonRangeCell<F> {
Expand All @@ -154,13 +146,6 @@ impl<F: FieldExt> CellExpression<F> for AllocatedCommonRangeCell<F> {
}

fn assign(&self, ctx: &mut Context<'_, F>, value: F) -> Result<AssignedCell<F, F>, Error> {
assert!(
value <= self.upper_bound,
"assigned value {:?} exceeds the limit {:?}",
value,
self.upper_bound
);

self.cell.assign(ctx, value)
}
}
Expand Down
8 changes: 1 addition & 7 deletions crates/zkwasm/src/circuits/etable/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use super::AllocatedU32StateCell;
use super::EVENT_TABLE_ENTRY_ROWS;
use crate::circuits::bit_table::BitTableOp;
use crate::circuits::cell::*;
use crate::circuits::config::common_range_max;
use crate::circuits::etable::ConstraintBuilder;
use crate::circuits::rtable::RangeTableConfig;
use crate::circuits::traits::ConfigureLookupTable;
Expand Down Expand Up @@ -269,7 +268,6 @@ impl AllocatorFreeCellsProfiler {

#[derive(Debug, Clone)]
pub(crate) struct EventTableCellAllocator<F: FieldExt> {
k: u32,
pub(crate) free_cells: BTreeMap<EventTableCellType, (usize, u32)>,
all_cols: BTreeMap<EventTableCellType, Vec<Vec<Column<Advice>>>>,
free_u32_cells: Vec<AllocatedU32Cell<F>>,
Expand Down Expand Up @@ -348,13 +346,12 @@ impl<F: FieldExt> EventTableCellAllocator<F> {

pub(super) fn new(
meta: &mut ConstraintSystem<F>,
k: u32,
sel: Column<Fixed>,
rtable: &RangeTableConfig<F>,
mtable: &impl ConfigureLookupTable<F>,
cols: &mut impl Iterator<Item = Column<Advice>>,
) -> Self {
let mut allocator = Self::_new(meta, k, sel, rtable, mtable, cols);
let mut allocator = Self::_new(meta, sel, rtable, mtable, cols);
for _ in 0..U32_CELLS {
let cell = allocator.prepare_alloc_u32_cell();
allocator.free_u32_cells.push(cell);
Expand All @@ -373,7 +370,6 @@ impl<F: FieldExt> EventTableCellAllocator<F> {

fn _new(
meta: &mut ConstraintSystem<F>,
k: u32,
sel: Column<Fixed>,
rtable: &RangeTableConfig<F>,
mtable: &impl ConfigureLookupTable<F>,
Expand Down Expand Up @@ -442,7 +438,6 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
);

Self {
k,
all_cols,
free_cells: BTreeMap::from_iter(
vec![
Expand Down Expand Up @@ -510,7 +505,6 @@ impl<F: FieldExt> EventTableCellAllocator<F> {
pub(crate) fn alloc_common_range_cell(&mut self) -> AllocatedCommonRangeCell<F> {
AllocatedCommonRangeCell {
cell: self.alloc(&EventTableCellType::CommonRange),
upper_bound: F::from(common_range_max(self.k) as u64),
}
}

Expand Down
Loading

0 comments on commit f855b26

Please sign in to comment.