diff --git a/.gitignore b/.gitignore index 827a41cf5..e5b76a0c3 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,7 @@ benchmark.csv # Build Artifacts recursion/gnark-ffi/build prover/build -prover/*.tar.gz \ No newline at end of file +prover/*.tar.gz + +# IDE Conf +.idea \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 4d1515393..6de1c30c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2015,6 +2015,7 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", "allocator-api2", + "serde", ] [[package]] @@ -4395,6 +4396,7 @@ dependencies = [ "curve25519-dalek", "elf", "elliptic-curve", + "hashbrown 0.14.5", "hex", "hybrid-array", "itertools 0.12.1", @@ -4574,6 +4576,7 @@ name = "sphinx-recursion-compiler" version = "0.1.0" dependencies = [ "backtrace", + "hashbrown 0.14.5", "itertools 0.12.1", "p3-air", "p3-baby-bear", diff --git a/Cargo.toml b/Cargo.toml index 4e0b52406..e8ab4882b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -74,7 +74,7 @@ ff = "0.13" futures = "0.3.30" futures-util = "0.3.14" getrandom = "=0.2.14" # 0.2.15 depends on yanked libc 0.2.154 -hashbrown = "0.14.5" +hashbrown = { version = "0.14.5", features = ["serde"] } hex = "0.4.3" home = "0.5.9" hybrid-array = "0.2.0-rc" diff --git a/core/Cargo.toml b/core/Cargo.toml index 45e6a93b6..d52358b41 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -42,6 +42,7 @@ bls12_381 = { workspace = true } cfg-if = { workspace = true } curve25519-dalek = { workspace = true } elliptic-curve = { workspace = true } +hashbrown = { workspace = true } hex = { workspace = true } hybrid-array = { workspace = true } k256 = { workspace = true, features = ["expose-field"] } diff --git a/core/benches/fibonacci.rs b/core/benches/fibonacci.rs index ed09024c4..09c2e0994 100644 --- a/core/benches/fibonacci.rs +++ b/core/benches/fibonacci.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use hashbrown::HashMap; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use sphinx_core::{ diff --git a/core/src/alu/divrem/mod.rs b/core/src/alu/divrem/mod.rs index ebcbd8c4a..dcfef8492 100644 --- a/core/src/alu/divrem/mod.rs +++ b/core/src/alu/divrem/mod.rs @@ -64,7 +64,7 @@ mod utils; use core::borrow::{Borrow, BorrowMut}; use core::mem::size_of; -use std::collections::HashMap; +use hashbrown::HashMap; use p3_air::{Air, AirBuilder, BaseAir}; use p3_field::AbstractField; @@ -394,7 +394,7 @@ impl MachineAir for DivRemChip { } let mut alu_events = HashMap::new(); alu_events.insert(Opcode::ADD, add_events); - output.add_alu_events(&alu_events); + output.add_alu_events(&mut alu_events); } let mut lower_word = 0; diff --git a/core/src/bytes/event.rs b/core/src/bytes/event.rs index acb0fc673..362c1876c 100644 --- a/core/src/bytes/event.rs +++ b/core/src/bytes/event.rs @@ -1,12 +1,10 @@ -use std::collections::BTreeMap; - use p3_field::PrimeField32; use serde::{Deserialize, Serialize}; use super::ByteOpcode; /// A byte lookup event. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] pub struct ByteLookupEvent { /// The shard number, used for byte lookup table. pub shard: u32, @@ -36,9 +34,10 @@ pub trait ByteRecord { fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent); /// Adds a list of `ByteLookupEvent`s to the record. + #[inline] fn add_byte_lookup_events(&mut self, blu_events: Vec) { - for blu_event in blu_events.iter() { - self.add_byte_lookup_event(*blu_event); + for blu_event in blu_events { + self.add_byte_lookup_event(blu_event); } } @@ -121,6 +120,7 @@ pub trait ByteRecord { impl ByteLookupEvent { /// Creates a new `ByteLookupEvent`. + #[inline(always)] pub fn new( shard: u32, channel: u32, @@ -147,13 +147,3 @@ impl ByteRecord for Vec { self.push(blu_event); } } - -impl ByteRecord for BTreeMap> { - fn add_byte_lookup_event(&mut self, blu_event: ByteLookupEvent) { - *self - .entry(blu_event.shard) - .or_default() - .entry(blu_event) - .or_insert(0) += 1 - } -} diff --git a/core/src/bytes/mod.rs b/core/src/bytes/mod.rs index 21e5731bb..388e6910b 100644 --- a/core/src/bytes/mod.rs +++ b/core/src/bytes/mod.rs @@ -5,7 +5,6 @@ pub mod opcode; pub mod trace; pub mod utils; -use alloc::collections::BTreeMap; use core::borrow::BorrowMut; use std::marker::PhantomData; @@ -35,19 +34,10 @@ pub const NUM_BYTE_LOOKUP_CHANNELS: u32 = 16; pub struct ByteChip(PhantomData); impl ByteChip { - /// Creates the preprocessed byte trace and event map. + /// Creates the preprocessed byte trace. /// - /// This function returns a pair `(trace, map)`, where: - /// - `trace` is a matrix containing all possible byte operations. - /// - `map` is a map from a byte lookup to the corresponding row it appears in the table and - /// the index of the result in the array of multiplicities. - pub fn trace_and_map( - shard: u32, - ) -> (RowMajorMatrix, BTreeMap) { - // A map from a byte lookup to its corresponding row in the table and index in the array of - // multiplicities. - let mut event_map = BTreeMap::new(); - + /// This function returns a `trace` which is a matrix containing all possible byte operations. + pub fn trace() -> RowMajorMatrix { // The trace containing all values, with all multiplicities set to zero. let mut initial_trace = RowMajorMatrix::new( vec![F::zero(); NUM_ROWS * NUM_BYTE_PREPROCESSED_COLS], @@ -65,10 +55,11 @@ impl ByteChip { col.b = F::from_canonical_u8(b); col.c = F::from_canonical_u8(c); + let shard = 0; // Iterate over all operations for results and updating the table map. for channel in 0..NUM_BYTE_LOOKUP_CHANNELS { - for (i, opcode) in opcodes.iter().enumerate() { - let event = match opcode { + for opcode in opcodes.iter() { + match opcode { ByteOpcode::AND => { let and = b & c; col.and = F::from_canonical_u8(and); @@ -176,11 +167,25 @@ impl ByteChip { ByteLookupEvent::new(shard, channel, *opcode, v, 0, 0, 0) } }; - event_map.insert(event, (row_index, i)); } } } - (initial_trace, event_map) + initial_trace + } +} + +#[cfg(test)] +mod tests { + use p3_baby_bear::BabyBear; + use std::time::Instant; + + use super::*; + + #[test] + fn test_trace_and_map() { + let start = Instant::now(); + ByteChip::::trace(); + println!("trace and map: {:?}", start.elapsed()); } } diff --git a/core/src/bytes/trace.rs b/core/src/bytes/trace.rs index 100c038c3..0d9cbfbfa 100644 --- a/core/src/bytes/trace.rs +++ b/core/src/bytes/trace.rs @@ -1,5 +1,6 @@ -use std::{borrow::BorrowMut, collections::BTreeMap}; +use std::borrow::BorrowMut; +use hashbrown::HashMap; use p3_field::Field; use p3_matrix::dense::RowMajorMatrix; @@ -9,6 +10,7 @@ use super::{ }; use crate::{ air::{EventLens, MachineAir, WithEvents}, + bytes::ByteOpcode, runtime::{ExecutionRecord, Program}, }; @@ -16,7 +18,7 @@ pub const NUM_ROWS: usize = 1 << 16; impl<'a, F: Field> WithEvents<'a> for ByteChip { // the byte lookups - type Events = &'a BTreeMap>; + type Events = &'a HashMap>; } impl MachineAir for ByteChip { @@ -33,10 +35,7 @@ impl MachineAir for ByteChip { } fn generate_preprocessed_trace(&self, _program: &Self::Program) -> Option> { - // OPT: We should be able to make this a constant. Also, trace / map should be separate. - // Since we only need the trace and not the map, we can just pass 0 as the shard. - let (trace, _) = Self::trace_and_map(0); - + let trace = Self::trace(); Some(trace) } @@ -53,23 +52,23 @@ impl MachineAir for ByteChip { input: &EL, _output: &mut ExecutionRecord, ) -> RowMajorMatrix { - let shard = input.index(); - let (_, event_map) = Self::trace_and_map(shard); - let mut trace = RowMajorMatrix::new( vec![F::zero(); NUM_BYTE_MULT_COLS * NUM_ROWS], NUM_BYTE_MULT_COLS, ); + let shard = input.index(); for (lookup, mult) in input.events()[&shard].iter() { - let (row, index) = event_map[lookup]; + let row = if lookup.opcode != ByteOpcode::U16Range { + ((lookup.b << 8) + lookup.c) as usize + } else { + lookup.a1 as usize + }; + let index = lookup.opcode as usize; let channel = lookup.channel as usize; - let cols: &mut ByteMultCols = trace.row_mut(row).borrow_mut(); - // Update the trace multiplicity + let cols: &mut ByteMultCols = trace.row_mut(row).borrow_mut(); cols.mult_channels[channel].multiplicities[index] += F::from_canonical_usize(*mult); - - // Set the shard column as the current shard. cols.shard = F::from_canonical_u32(shard); } diff --git a/core/src/cpu/columns/channel.rs b/core/src/cpu/columns/channel.rs index c6f1355be..dab7048de 100644 --- a/core/src/cpu/columns/channel.rs +++ b/core/src/cpu/columns/channel.rs @@ -13,6 +13,7 @@ pub struct ChannelSelectorCols { } impl ChannelSelectorCols { + #[inline(always)] pub fn populate(&mut self, channel: u32) { self.channel_selectors = [F::zero(); NUM_BYTE_LOOKUP_CHANNELS as usize]; self.channel_selectors[channel as usize] = F::one(); diff --git a/core/src/cpu/trace.rs b/core/src/cpu/trace.rs index db9e7f5ec..53b2584db 100644 --- a/core/src/cpu/trace.rs +++ b/core/src/cpu/trace.rs @@ -1,10 +1,12 @@ +use hashbrown::HashMap; use std::array; use std::borrow::BorrowMut; -use std::collections::HashMap; use p3_field::{PrimeField, PrimeField32}; use p3_matrix::dense::RowMajorMatrix; -use p3_maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator, ParallelSlice}; +use p3_maybe_rayon::prelude::{ + IntoParallelRefMutIterator, ParallelBridge, ParallelIterator, ParallelSlice, ParallelSliceMut, +}; use tracing::instrument; use super::columns::{CPU_COL_MAP, NUM_CPU_COLS}; @@ -38,47 +40,28 @@ impl MachineAir for CpuChip { fn generate_trace>( &self, input: &EL, - output: &mut ExecutionRecord, + _: &mut ExecutionRecord, ) -> RowMajorMatrix { - let mut new_alu_events = HashMap::new(); - let mut new_blu_events = Vec::new(); - let (events, nonce_lookup) = input.events(); - // Generate the trace rows for each event. - let mut rows_with_events = events - .par_iter() - .map(|op: &CpuEvent| self.event_to_row::(*op, nonce_lookup)) - .collect::>(); - - // No need to sort by the shard, since the cpu events are already partitioned by that. - rows_with_events.sort_unstable_by_key(|(event, _, _)| event[CPU_COL_MAP.clk]); - - let mut rows = Vec::::new(); - for row_with_events in rows_with_events { - let (row, alu_events, blu_events) = row_with_events; - rows.extend(row); - for (key, value) in alu_events { - new_alu_events - .entry(key) - .and_modify(|op_new_events: &mut Vec| { - op_new_events.extend(value.clone()) - }) - .or_insert(value); - } - new_blu_events.extend(blu_events); - } - - // Add the dependency events to the shard. - for (_, value) in new_alu_events.iter_mut() { - value.sort_unstable_by_key(|event| event.clk); - } - new_blu_events.sort_unstable_by_key(|event| event.a1); - output.add_alu_events(&new_alu_events); - output.add_byte_lookup_events(new_blu_events); + let mut values = vec![F::zero(); events.len() * NUM_CPU_COLS]; + let chunk_size = std::cmp::max(events.len() / num_cpus::get(), 1); + values + .chunks_mut(chunk_size * NUM_CPU_COLS) + .enumerate() + .par_bridge() + .for_each(|(i, rows)| { + rows.chunks_mut(NUM_CPU_COLS) + .enumerate() + .for_each(|(j, row)| { + let idx = i * chunk_size + j; + let cols: &mut CpuCols = row.borrow_mut(); + self.event_to_row(&events[idx], nonce_lookup, cols); + }); + }); // Convert the trace to a row major matrix. - let mut trace = RowMajorMatrix::new(rows, NUM_CPU_COLS); + let mut trace = RowMajorMatrix::new(values, NUM_CPU_COLS); // Pad the trace to a power of two. Self::pad_to_power_of_two::(&mut trace.values); @@ -90,15 +73,18 @@ impl MachineAir for CpuChip { fn generate_dependencies>(&self, input: &EL, output: &mut ExecutionRecord) { // Generate the trace rows for each event. let chunk_size = std::cmp::max(input.events().0.len() / num_cpus::get(), 1); - let events = input + let (alu_events, blu_events): (Vec<_>, Vec<_>) = input .events() .0 .par_chunks(chunk_size) .map(|ops: &[CpuEvent]| { let mut alu = HashMap::new(); - let mut blu: Vec<_> = Vec::default(); + let mut blu: Vec<_> = Vec::with_capacity(ops.len() * 8); for op in ops.iter() { - let (_, alu_events, blu_events) = self.event_to_row::(*op, &HashMap::new()); + let mut row = [F::zero(); NUM_CPU_COLS]; + let cols: &mut CpuCols = row.as_mut_slice().borrow_mut(); + let (alu_events, blu_events) = + self.event_to_row::(op, &HashMap::new(), cols); for (key, value) in alu_events { alu.entry(key).or_insert(Vec::default()).extend(value); } @@ -106,16 +92,17 @@ impl MachineAir for CpuChip { } (alu, blu) }) - .collect::>(); + .unzip(); - for (mut alu_events, mut blu_events) in events { - for (_, value) in alu_events.iter_mut() { - value.sort_unstable_by_key(|event| event.clk); - } - // Add the dependency events to the shard. - output.add_alu_events(&alu_events); - blu_events.sort_unstable_by_key(|event| event.a1); - output.add_byte_lookup_events(blu_events); + for mut alu_events_chunk in alu_events { + output.add_alu_events(&mut alu_events_chunk); + } + + let mut blu_events = blu_events.into_iter().flatten().collect::>(); + blu_events.par_sort_unstable_by_key(|event| (event.shard, event.opcode)); + + for blu_event in blu_events { + output.add_byte_lookup_event(blu_event); } } @@ -128,19 +115,13 @@ impl CpuChip { /// Create a row from an event. fn event_to_row( &self, - event: CpuEvent, + event: &CpuEvent, nonce_lookup: &HashMap, - ) -> ( - [F; NUM_CPU_COLS], - HashMap>, - Vec, - ) { + cols: &mut CpuCols, + ) -> (HashMap>, Vec) { let mut new_alu_events = HashMap::new(); let mut new_blu_events = Vec::new(); - let mut row = [F::zero(); NUM_CPU_COLS]; - let cols: &mut CpuCols = row.as_mut_slice().borrow_mut(); - // Populate shard and clk columns. self.populate_shard_clk(cols, event, &mut new_blu_events); @@ -234,19 +215,27 @@ impl CpuChip { // Assert that the instruction is not a no-op. cols.is_real = F::one(); - (row, new_alu_events, new_blu_events) + (new_alu_events, new_blu_events) } /// Populates the shard, channel, and clk related rows. fn populate_shard_clk( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, new_blu_events: &mut Vec, ) { cols.shard = F::from_canonical_u32(event.shard); cols.channel = F::from_canonical_u32(event.channel); + cols.clk = F::from_canonical_u32(event.clk); + + let clk_16bit_limb = event.clk & 0xffff; + let clk_8bit_limb = (event.clk >> 16) & 0xff; + cols.clk_16bit_limb = F::from_canonical_u32(clk_16bit_limb); + cols.clk_8bit_limb = F::from_canonical_u32(clk_8bit_limb); + cols.channel_selectors.populate(event.channel); + new_blu_events.push(ByteLookupEvent::new( event.shard, event.channel, @@ -256,12 +245,6 @@ impl CpuChip { 0, 0, )); - - cols.clk = F::from_canonical_u32(event.clk); - let clk_16bit_limb = event.clk & 0xffff; - cols.clk_16bit_limb = F::from_canonical_u32(clk_16bit_limb); - let clk_8bit_limb = (event.clk >> 16) & 0xff; - cols.clk_8bit_limb = F::from_canonical_u32(clk_8bit_limb); new_blu_events.push(ByteLookupEvent::new( event.shard, event.channel, @@ -286,7 +269,7 @@ impl CpuChip { fn populate_memory( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, new_alu_events: &mut HashMap>, new_blu_events: &mut Vec, nonce_lookup: &HashMap, @@ -438,7 +421,7 @@ impl CpuChip { fn populate_branch( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, alu_events: &mut HashMap>, nonce_lookup: &HashMap, ) { @@ -567,7 +550,7 @@ impl CpuChip { fn populate_jump( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, alu_events: &mut HashMap>, nonce_lookup: &HashMap, ) { @@ -644,7 +627,7 @@ impl CpuChip { fn populate_auipc( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, alu_events: &mut HashMap>, nonce_lookup: &HashMap, ) { @@ -683,7 +666,7 @@ impl CpuChip { fn populate_ecall( &self, cols: &mut CpuCols, - event: CpuEvent, + event: &CpuEvent, nonce_lookup: &HashMap, ) -> bool { let mut is_halt = false; @@ -766,7 +749,7 @@ impl CpuChip { ) }; - rows[n_real_rows..].iter_mut().for_each(|padded_row| { + rows[n_real_rows..].par_iter_mut().for_each(|padded_row| { padded_row[CPU_COL_MAP.selectors.imm_b] = F::one(); padded_row[CPU_COL_MAP.selectors.imm_c] = F::one(); }); diff --git a/core/src/lib.rs b/core/src/lib.rs index 73a81f941..74025dcf9 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -6,8 +6,6 @@ )] #![warn(unused_extern_crates)] -extern crate alloc; - pub mod air; pub mod alu; pub mod bytes; diff --git a/core/src/memory/global.rs b/core/src/memory/global.rs index 5bad9d087..ae11bad09 100644 --- a/core/src/memory/global.rs +++ b/core/src/memory/global.rs @@ -5,6 +5,7 @@ use std::array; use p3_air::{Air, AirBuilder, BaseAir}; use p3_field::{AbstractField, PrimeField}; use p3_matrix::{dense::RowMajorMatrix, Matrix}; +use p3_maybe_rayon::prelude::{ParallelBridge, ParallelIterator}; use sphinx_derive::AlignedBorrow; use super::MemoryInitializeFinalizeEvent; @@ -70,55 +71,67 @@ impl MachineAir for MemoryChip { MemoryChipType::Finalize => input.events().1, } .to_vec(); + memory_events.sort_by_key(|event| event.addr); - let rows: Vec<[F; NUM_MEMORY_INIT_COLS]> = (0..memory_events.len()) // OPT: change this to par_iter - .map(|i| { - let MemoryInitializeFinalizeEvent { - addr, - value, - shard, - timestamp, - used, - } = memory_events[i]; - - let mut row = [F::zero(); NUM_MEMORY_INIT_COLS]; - let cols: &mut MemoryInitCols = row.as_mut_slice().borrow_mut(); - cols.addr = F::from_canonical_u32(addr); - cols.addr_bits.populate(addr); - cols.shard = F::from_canonical_u32(shard); - cols.timestamp = F::from_canonical_u32(timestamp); - cols.value = array::from_fn(|i| F::from_canonical_u32((value >> i) & 1)); - cols.is_real = F::from_canonical_u32(used); - - if i != memory_events.len() - 1 { - let next_addr = memory_events[i + 1].addr; - assert_ne!(next_addr, addr); - - cols.addr_bits.populate(addr); - - cols.seen_diff_bits[0] = F::zero(); - for j in 0..32 { - let rev_j = 32 - j - 1; - let next_bit = ((next_addr >> rev_j) & 1) == 1; - let local_bit = ((addr >> rev_j) & 1) == 1; - cols.match_bits[j] = - F::from_bool((local_bit && next_bit) || (!local_bit && !next_bit)); - cols.seen_diff_bits[j + 1] = cols.seen_diff_bits[j] - + (F::one() - cols.seen_diff_bits[j]) * (F::one() - cols.match_bits[j]); - cols.not_match_and_not_seen_diff_bits[j] = - (F::one() - cols.match_bits[j]) * (F::one() - cols.seen_diff_bits[j]); - } - assert_eq!(cols.seen_diff_bits[cols.seen_diff_bits.len() - 1], F::one()); - } - - row - }) - .collect::>(); - - let mut trace = RowMajorMatrix::new( - rows.into_iter().flatten().collect::>(), - NUM_MEMORY_INIT_COLS, - ); + + let mut rows = vec![F::zero(); memory_events.len() * NUM_MEMORY_INIT_COLS]; + let chunk_size = std::cmp::max(memory_events.len() / num_cpus::get(), 1); + + rows.chunks_mut(chunk_size * NUM_MEMORY_INIT_COLS) + .enumerate() + .par_bridge() + .for_each(|(i, rows)| { + rows.chunks_mut(NUM_MEMORY_INIT_COLS) + .enumerate() + .for_each(|(j, row)| { + let idx = i * chunk_size + j; + + let MemoryInitializeFinalizeEvent { + addr, + value, + shard, + timestamp, + used, + } = memory_events[idx]; + + let cols: &mut MemoryInitCols = row.borrow_mut(); + cols.addr = F::from_canonical_u32(addr); + cols.addr_bits.populate(addr); + cols.shard = F::from_canonical_u32(shard); + cols.timestamp = F::from_canonical_u32(timestamp); + cols.value = array::from_fn(|i| F::from_canonical_u32((value >> i) & 1)); + cols.is_real = F::from_canonical_u32(used); + + if idx != memory_events.len() - 1 { + let next_addr = memory_events[idx + 1].addr; + assert_ne!(next_addr, addr); + + cols.addr_bits.populate(addr); + + cols.seen_diff_bits[0] = F::zero(); + for k in 0..32 { + let rev_k = 32 - k - 1; + let next_bit = ((next_addr >> rev_k) & 1) == 1; + let local_bit = ((addr >> rev_k) & 1) == 1; + cols.match_bits[k] = F::from_bool( + (local_bit && next_bit) || (!local_bit && !next_bit), + ); + cols.seen_diff_bits[k + 1] = cols.seen_diff_bits[k] + + (F::one() - cols.seen_diff_bits[k]) + * (F::one() - cols.match_bits[k]); + cols.not_match_and_not_seen_diff_bits[k] = (F::one() + - cols.match_bits[k]) + * (F::one() - cols.seen_diff_bits[k]); + } + assert_eq!( + cols.seen_diff_bits[cols.seen_diff_bits.len() - 1], + F::one() + ); + } + }); + }); + + let mut trace = RowMajorMatrix::new(rows, NUM_MEMORY_INIT_COLS); pad_to_power_of_two::(&mut trace.values); diff --git a/core/src/program/mod.rs b/core/src/program/mod.rs index 21d06c021..db54778c3 100644 --- a/core/src/program/mod.rs +++ b/core/src/program/mod.rs @@ -2,7 +2,7 @@ use core::{ borrow::{Borrow, BorrowMut}, mem::size_of, }; -use std::collections::HashMap; +use hashbrown::HashMap; use p3_air::{Air, BaseAir, PairBuilder}; use p3_field::PrimeField; diff --git a/core/src/runtime/mod.rs b/core/src/runtime/mod.rs index f453e205e..c5292ca56 100644 --- a/core/src/runtime/mod.rs +++ b/core/src/runtime/mod.rs @@ -20,8 +20,8 @@ pub(crate) use state::*; pub use syscall::*; pub use utils::*; -use std::collections::hash_map::Entry; -use std::collections::HashMap; +use hashbrown::hash_map::Entry; +use hashbrown::HashMap; use std::fmt::{Display, Formatter, Result as FmtResult}; use std::fs::File; use std::io::BufWriter; diff --git a/core/src/runtime/record.rs b/core/src/runtime/record.rs index 2f689f068..e1af8197e 100644 --- a/core/src/runtime/record.rs +++ b/core/src/runtime/record.rs @@ -1,10 +1,6 @@ -use std::{ - collections::{BTreeMap, HashMap}, - mem::take, - sync::Arc, -}; +use std::{mem::take, sync::Arc}; -use itertools::Itertools; +use hashbrown::HashMap; use p3_field::{AbstractField, Field}; use serde::{Deserialize, Serialize}; @@ -94,9 +90,11 @@ pub struct ExecutionRecord { /// A trace of the SLT, SLTI, SLTU, and SLTIU events. pub lt_events: Vec, - /// All byte lookups that are needed. The layout is shard -> (event -> count). Byte lookups are - /// sharded to prevent the multiplicities from overflowing. - pub byte_lookups: BTreeMap>, + /// All byte lookups that are needed. + /// + /// The layout is shard -> (event -> count). Byte lookups are sharded to prevent the + /// multiplicities from overflowing. + pub byte_lookups: HashMap>, pub sha_extend_events: Vec, @@ -583,7 +581,7 @@ impl MachineRecord for ExecutionRecord { Some(existing) => { // If there's already a map for this shard, update counts for each event. for (event, count) in events_map.iter() { - *existing.entry(*event).or_insert(0) += count; + *existing.entry(event.clone()).or_insert(0) += count; } } None => { @@ -989,32 +987,29 @@ impl ExecutionRecord { self.lt_events.push(lt_event); } - pub fn add_alu_events(&mut self, alu_events: &HashMap>) { - let keys = alu_events.keys().sorted(); - for opcode in keys { + pub fn add_alu_events(&mut self, alu_events: &mut HashMap>) { + for (opcode, value) in alu_events.iter_mut() { match opcode { Opcode::ADD => { - self.add_events.extend_from_slice(&alu_events[opcode]); + self.add_events.append(value); } Opcode::MUL | Opcode::MULH | Opcode::MULHU | Opcode::MULHSU => { - self.mul_events.extend_from_slice(&alu_events[opcode]); + self.mul_events.append(value); } Opcode::SUB => { - self.sub_events.extend_from_slice(&alu_events[opcode]); + self.sub_events.append(value); } Opcode::XOR | Opcode::OR | Opcode::AND => { - self.bitwise_events.extend_from_slice(&alu_events[opcode]); + self.bitwise_events.append(value); } Opcode::SLL => { - self.shift_left_events - .extend_from_slice(&alu_events[opcode]); + self.shift_left_events.append(value); } Opcode::SRL | Opcode::SRA => { - self.shift_right_events - .extend_from_slice(&alu_events[opcode]); + self.shift_right_events.append(value); } Opcode::SLT | Opcode::SLTU => { - self.lt_events.extend_from_slice(&alu_events[opcode]); + self.lt_events.append(value); } _ => { panic!("Invalid opcode: {:?}", opcode); diff --git a/core/src/runtime/state.rs b/core/src/runtime/state.rs index 7171e1449..11cf7ff3b 100644 --- a/core/src/runtime/state.rs +++ b/core/src/runtime/state.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use hashbrown::HashMap; use nohash_hasher::BuildNoHashHasher; use serde::{Deserialize, Serialize}; diff --git a/core/src/runtime/syscall.rs b/core/src/runtime/syscall.rs index a77d53b5e..13393e426 100644 --- a/core/src/runtime/syscall.rs +++ b/core/src/runtime/syscall.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use hashbrown::HashMap; use std::fmt; use std::sync::Arc; diff --git a/core/src/stark/config.rs b/core/src/stark/config.rs index 2caf8c2e7..8f87b534d 100644 --- a/core/src/stark/config.rs +++ b/core/src/stark/config.rs @@ -49,7 +49,7 @@ pub type Challenger = ::Challenger; pub trait StarkGenericConfig: Send + Sync + Serialize + DeserializeOwned + Clone { type Val: PrimeField; - type Domain: PolynomialSpace + Sync; + type Domain: PolynomialSpace + Sync + Send; /// The PCS used to commit to trace polynomials. type Pcs: Pcs + Sync; diff --git a/core/src/stark/machine.rs b/core/src/stark/machine.rs index baf51264b..219f47a56 100644 --- a/core/src/stark/machine.rs +++ b/core/src/stark/machine.rs @@ -1,3 +1,4 @@ +use hashbrown::HashMap; use itertools::Itertools; use p3_air::Air; use p3_challenger::CanObserve; @@ -14,7 +15,6 @@ use serde::de::DeserializeOwned; use serde::Deserialize; use serde::Serialize; use std::cmp::Reverse; -use std::collections::HashMap; use std::fmt::Debug; use tracing::instrument; diff --git a/core/src/stark/prover.rs b/core/src/stark/prover.rs index a6c7f9c35..067f94edd 100644 --- a/core/src/stark/prover.rs +++ b/core/src/stark/prover.rs @@ -321,8 +321,8 @@ where let domains_and_perm_traces = tracing::debug_span!("flatten permutation traces and collect domains").in_scope(|| { permutation_traces - .into_iter() - .zip(trace_domains.iter()) + .into_par_iter() + .zip(trace_domains.par_iter()) .map(|(perm_trace, domain)| { let trace = perm_trace.flatten_to_base(); (*domain, trace.clone()) @@ -354,7 +354,7 @@ where let quotient_values = parent_span.in_scope(|| { quotient_domains - .into_par_iter() + .par_iter() .enumerate() .map(|(i, quotient_domain)| { tracing::debug_span!(parent: &parent_span, "compute quotient values for domain") diff --git a/core/src/stark/record.rs b/core/src/stark/record.rs index 10e035eaa..00f14f5cb 100644 --- a/core/src/stark/record.rs +++ b/core/src/stark/record.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use hashbrown::HashMap; use p3_field::AbstractField; diff --git a/core/src/stark/types.rs b/core/src/stark/types.rs index 67ef8d403..2ccd32e19 100644 --- a/core/src/stark/types.rs +++ b/core/src/stark/types.rs @@ -1,5 +1,5 @@ +use hashbrown::HashMap; use std::{ - collections::HashMap, fmt::Debug, fs::File, io::{BufReader, BufWriter, Seek}, diff --git a/core/src/syscall/unconstrained.rs b/core/src/syscall/unconstrained.rs index 65f59a334..291a9d610 100644 --- a/core/src/syscall/unconstrained.rs +++ b/core/src/syscall/unconstrained.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use hashbrown::HashMap; use crate::runtime::{ForkState, Syscall, SyscallContext}; diff --git a/core/src/utils/prove.rs b/core/src/utils/prove.rs index 36f456fd2..c5a4a5bc6 100644 --- a/core/src/utils/prove.rs +++ b/core/src/utils/prove.rs @@ -1,5 +1,5 @@ +use hashbrown::HashMap; use nohash_hasher::BuildNoHashHasher; -use std::collections::HashMap; use std::fs::File; use std::io; use std::io::{Seek, Write}; diff --git a/recursion/compiler/Cargo.toml b/recursion/compiler/Cargo.toml index 55eb92275..a93459d75 100644 --- a/recursion/compiler/Cargo.toml +++ b/recursion/compiler/Cargo.toml @@ -23,6 +23,7 @@ p3-baby-bear = { workspace = true } p3-poseidon2 = { workspace = true } backtrace = { workspace = true } tracing = { workspace = true } +hashbrown = { workspace = true } [dev-dependencies] p3-challenger = { workspace = true } diff --git a/recursion/compiler/src/ir/types.rs b/recursion/compiler/src/ir/types.rs index 148caa6bb..beaa1f9ed 100644 --- a/recursion/compiler/src/ir/types.rs +++ b/recursion/compiler/src/ir/types.rs @@ -1,6 +1,7 @@ use alloc::format; use core::marker::PhantomData; -use std::{collections::HashMap, hash::Hash}; +use hashbrown::HashMap; +use std::hash::Hash; use p3_field::{AbstractField, ExtensionField, Field}; use serde::{Deserialize, Serialize}; diff --git a/recursion/core/src/range_check/trace.rs b/recursion/core/src/range_check/trace.rs index 21005d212..6363c80ed 100644 --- a/recursion/core/src/range_check/trace.rs +++ b/recursion/core/src/range_check/trace.rs @@ -1,5 +1,6 @@ -use std::{borrow::BorrowMut, collections::BTreeMap}; +use std::borrow::BorrowMut; +use hashbrown::HashMap; use p3_field::{Field, PrimeField32}; use p3_matrix::dense::RowMajorMatrix; use sphinx_core::air::{EventLens, MachineAir, WithEvents}; @@ -13,7 +14,7 @@ use crate::runtime::{ExecutionRecord, RecursionProgram}; pub const NUM_ROWS: usize = 1 << 16; impl<'a, F: Field> WithEvents<'a> for RangeCheckChip { - type Events = &'a BTreeMap; + type Events = &'a HashMap; } impl MachineAir for RangeCheckChip { diff --git a/recursion/core/src/runtime/record.rs b/recursion/core/src/runtime/record.rs index 01377f719..9279ad8c2 100644 --- a/recursion/core/src/runtime/record.rs +++ b/recursion/core/src/runtime/record.rs @@ -1,7 +1,7 @@ use std::array; -use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use hashbrown::HashMap; use p3_field::{AbstractField, PrimeField32}; use sphinx_core::air::EventLens; use sphinx_core::stark::{Indexed, MachineRecord, PROOF_MAX_NUM_PVS}; @@ -23,7 +23,7 @@ pub struct ExecutionRecord { pub cpu_events: Vec>, pub poseidon2_events: Vec>, pub fri_fold_events: Vec>, - pub range_check_events: BTreeMap, + pub range_check_events: HashMap, // (address, value) pub first_memory_record: Vec<(F, Block)>,