diff options
Diffstat (limited to 'crates/parser/src')
| -rw-r--r-- | crates/parser/src/conversion.rs | 6 | ||||
| -rw-r--r-- | crates/parser/src/lib.rs | 121 | ||||
| -rw-r--r-- | crates/parser/src/module.rs | 258 | ||||
| -rw-r--r-- | crates/parser/src/optimize.rs | 46 | ||||
| -rw-r--r-- | crates/parser/src/parallel.rs | 174 | ||||
| -rw-r--r-- | crates/parser/src/visit.rs | 10 |
6 files changed, 529 insertions, 86 deletions
diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs index f7fcb34..487f992 100644 --- a/crates/parser/src/conversion.rs +++ b/crates/parser/src/conversion.rs @@ -1,6 +1,6 @@ use alloc::sync::Arc; -use crate::{Result, module::Code, visit::process_operators_and_validate}; +use crate::{Result, module::FunctionCode, visit::process_operators_and_validate}; use alloc::{boxed::Box, format, string::ToString, vec::Vec}; use tinywasm_types::*; use wasmparser::{FuncValidator, FuncValidatorAllocations, OperatorsReader, ValidatorResources}; @@ -163,7 +163,7 @@ pub(crate) fn convert_module_export(export: wasmparser::Export<'_>) -> Result<Ex pub(crate) fn convert_module_code( func: wasmparser::FunctionBody<'_>, mut validator: FuncValidator<ValidatorResources>, -) -> Result<(Code, FuncValidatorAllocations)> { +) -> Result<(FunctionCode, FuncValidatorAllocations)> { let locals_reader = func.get_locals_reader()?; let count = locals_reader.get_count(); let pos = locals_reader.original_position(); @@ -200,7 +200,7 @@ pub(crate) fn convert_module_code( } let (body, data, allocations) = process_operators_and_validate(validator, func, local_addr_map)?; - Ok(((body, data, local_counts), allocations)) + Ok((FunctionCode { instructions: body, data, locals: local_counts, uses_local_memory: false }, allocations)) } pub(crate) fn convert_module_type(ty: wasmparser::RecGroup) -> Result<Arc<FuncType>> { diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs index 4eda84b..c4f6ff7 100644 --- a/crates/parser/src/lib.rs +++ b/crates/parser/src/lib.rs @@ -35,6 +35,10 @@ mod macros; mod module; mod optimize; mod visit; + +#[cfg(parallel_parser)] +mod parallel; + pub use error::*; use module::ModuleReader; use wasmparser::{Validator, WasmFeatures}; @@ -51,11 +55,27 @@ pub struct ParserOptions { pub optimize_rewrite: bool, /// Whether to remove `Nop` and `MergeBarrier` instructions after rewriting. pub optimize_remove_nop: bool, + + #[cfg(parallel_parser)] + /// Number of threads to use for parallel parsing. + /// + /// Requires the `parallel` feature. Ignored when the feature is disabled. + /// + /// - `None`: auto-detect based on available parallelism + /// - `Some(1)`: force single-threaded + /// - `Some(n)`: use up to `n` workers + pub parser_threads: Option<usize>, } impl Default for ParserOptions { fn default() -> Self { - Self { optimize_local_memory_allocation: true, optimize_rewrite: true, optimize_remove_nop: true } + Self { + optimize_local_memory_allocation: true, + optimize_rewrite: true, + optimize_remove_nop: true, + #[cfg(parallel_parser)] + parser_threads: None, + } } } @@ -92,6 +112,21 @@ impl ParserOptions { pub const fn optimize_remove_nop(&self) -> bool { self.optimize_remove_nop } + + #[cfg(parallel_parser)] + /// Set the number of threads for parallel parsing. + /// + /// Requires the `parallel` feature to have any effect. + pub const fn with_parser_threads(mut self, threads: usize) -> Self { + self.parser_threads = Some(threads); + self + } + + #[cfg(parallel_parser)] + /// Returns the configured parser thread count, or `None` for auto-detect. + pub const fn parser_threads(&self) -> Option<usize> { + self.parser_threads + } } /// A WebAssembly parser @@ -139,6 +174,17 @@ impl Parser { Validator::new_with_features(features) } + #[cfg(feature = "std")] + fn read_more(stream: &mut impl std::io::Read, buffer: &mut alloc::vec::Vec<u8>, hint: usize) -> Result<usize> { + let len = buffer.len(); + buffer.extend((0..hint).map(|_| 0u8)); + let read_bytes = stream + .read(&mut buffer[len..]) + .map_err(|e| ParseError::Other(alloc::format!("Error reading from stream: {e}")))?; + buffer.truncate(len + read_bytes); + Ok(read_bytes) + } + /// Parse a [`Module`] from bytes pub fn parse_module_bytes(&self, wasm: impl AsRef<[u8]>) -> Result<Module> { let wasm = wasm.as_ref(); @@ -146,13 +192,22 @@ impl Parser { let mut reader = ModuleReader::default(); for payload in wasmparser::Parser::new(0).parse_all(wasm) { - reader.process_payload(payload?, &mut validator)?; + match payload? { + wasmparser::Payload::CodeSectionStart { count, range, size } => { + reader.begin_code_section(count, range, size, &mut validator, &self.options)?; + } + wasmparser::Payload::CodeSectionEntry(function) => { + reader.process_borrowed_code_section_entry(function, &mut validator, &self.options)?; + } + payload => reader.process_payload(payload, &mut validator)?, + } } if !reader.end_reached { return Err(ParseError::EndNotReached); } + reader.process_pending_functions(&self.options)?; reader.into_module(&self.options) } @@ -176,18 +231,60 @@ impl Parser { loop { match parser.parse(&buffer, eof)? { wasmparser::Chunk::NeedMoreData(hint) => { - let len = buffer.len(); - buffer.extend((0..hint).map(|_| 0u8)); - let read_bytes = stream - .read(&mut buffer[len..]) - .map_err(|e| ParseError::Other(alloc::format!("Error reading from stream: {e}")))?; - buffer.truncate(len + read_bytes); + let read_bytes = Self::read_more(&mut stream, &mut buffer, hint as usize)?; eof = read_bytes == 0; } wasmparser::Chunk::Parsed { consumed, payload } => { - reader.process_payload(payload, &mut validator)?; - buffer.drain(..consumed); + #[cfg(parallel_parser)] + let mut deferred_code_section = None; + + match payload { + wasmparser::Payload::CodeSectionStart { count, range, size } => { + let defer = + reader.begin_code_section(count, range.clone(), size, &mut validator, &self.options)?; + + #[cfg(parallel_parser)] + if defer { + deferred_code_section = Some((count, range.end - size as usize, size as usize)); + } + + #[cfg(not(parallel_parser))] + let _ = defer; + + buffer.drain(..consumed); + } + wasmparser::Payload::CodeSectionEntry(function) => { + reader.process_inline_code_section_entry(function, &mut validator, &self.options)?; + buffer.drain(..consumed); + } + payload => { + reader.process_payload(payload, &mut validator)?; + buffer.drain(..consumed); + } + } + + #[cfg(parallel_parser)] + if let Some((count, body_offset, section_size)) = deferred_code_section { + while buffer.len() < section_size { + let remaining = section_size - buffer.len(); + let read_bytes = Self::read_more(&mut stream, &mut buffer, remaining)?; + if read_bytes == 0 { + return Err(ParseError::ParseError { + message: "unexpected end-of-file".into(), + offset: body_offset + buffer.len(), + }); + } + } + + let section_bytes = alloc::sync::Arc::<[u8]>::from(buffer[..section_size].to_vec()); + reader.queue_owned_code_section(count, body_offset, section_bytes, &mut validator)?; + parser.skip_section(); + buffer.drain(..section_size); + continue; + } + if eof || reader.end_reached { + reader.process_pending_functions(&self.options)?; return reader.into_module(&self.options); } } @@ -196,10 +293,10 @@ impl Parser { } } -impl TryFrom<ModuleReader> for Module { +impl TryFrom<ModuleReader<'_>> for Module { type Error = ParseError; - fn try_from(reader: ModuleReader) -> Result<Self> { + fn try_from(reader: ModuleReader<'_>) -> Result<Self> { reader.into_module(&ParserOptions::default()) } } diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs index b8dd9e7..02a3297 100644 --- a/crates/parser/src/module.rs +++ b/crates/parser/src/module.rs @@ -2,21 +2,59 @@ use crate::log::debug; use crate::{ParseError, ParserOptions, Result, conversion, optimize}; use alloc::sync::Arc; use alloc::{format, string::ToString, vec::Vec}; +use core::marker::PhantomData; +use core::ops::Range; use tinywasm_types::*; use wasmparser::{FuncValidatorAllocations, Payload, Validator}; -pub(crate) type Code = (Vec<Instruction>, WasmFunctionData, ValueCounts); +pub(crate) struct FunctionCode { + pub instructions: Vec<Instruction>, + pub data: WasmFunctionData, + pub locals: ValueCounts, + pub uses_local_memory: bool, +} + +pub(crate) fn imported_func_count(imports: &[Import]) -> usize { + imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count() +} + +pub(crate) fn imported_memory_count(imports: &[Import]) -> u32 { + imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32 +} + +pub(crate) fn optimize_function_code( + mut code: FunctionCode, + options: &ParserOptions, + function_results: ValueCounts, + self_func_addr: u32, + imported_memory_count: u32, +) -> FunctionCode { + let optimized = optimize::optimize_instructions( + code.instructions, + &mut code.data, + options, + function_results, + self_func_addr, + imported_memory_count, + ); + + code.instructions = optimized.instructions; + code.uses_local_memory = optimized.uses_local_memory; + code +} #[derive(Default)] -pub(crate) struct ModuleReader { +pub(crate) struct ModuleReader<'a> { func_validator_allocations: Option<FuncValidatorAllocations>, + has_code_section: bool, + marker: PhantomData<&'a [u8]>, pub(crate) version: Option<u16>, pub(crate) start_func: Option<u32>, pub(crate) func_types: Vec<Arc<FuncType>>, pub(crate) code_type_addrs: Vec<u32>, pub(crate) exports: Vec<Export>, - pub(crate) code: Vec<Code>, + pub(crate) code: Vec<FunctionCode>, pub(crate) globals: Vec<Global>, pub(crate) table_types: Vec<TableType>, pub(crate) memory_types: Vec<MemoryType>, @@ -24,9 +62,18 @@ pub(crate) struct ModuleReader { pub(crate) data: Vec<Data>, pub(crate) elements: Vec<Element>, pub(crate) end_reached: bool, + + #[cfg(parallel_parser)] + pending_functions: Option<Vec<crate::parallel::PendingFunction<'a>>>, } -impl ModuleReader { +impl<'a> ModuleReader<'a> { + fn function_results(&self, ordinal: usize) -> ValueCounts { + let ty_idx = self.code_type_addrs[ordinal]; + let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug"); + ValueCounts::from_iter(ty.results()) + } + pub(crate) fn process_payload(&mut self, payload: Payload<'_>, validator: &mut Validator) -> Result<()> { match payload { Payload::Version { num, encoding, range } => { @@ -55,7 +102,6 @@ impl ModuleReader { self.func_types = reader.into_iter().map(|t| conversion::convert_module_type(t?)).collect::<Result<Vec<_>>>()?; } - Payload::GlobalSection(reader) => { if !self.globals.is_empty() { return Err(ParseError::DuplicateSection("Global section".into())); @@ -69,6 +115,7 @@ impl ModuleReader { if !self.table_types.is_empty() { return Err(ParseError::DuplicateSection("Table section".into())); } + debug!("Found table section"); validator.table_section(&reader)?; self.table_types = conversion::convert_module_tables(reader)?; @@ -112,22 +159,6 @@ impl ModuleReader { validator.function_section(&reader)?; self.code_type_addrs = reader.into_iter().map(|f| Ok(f?)).collect::<Result<Vec<_>>>()?; } - Payload::CodeSectionStart { count, range, .. } => { - debug!("Found code section ({count} functions)"); - if !self.code.is_empty() { - return Err(ParseError::DuplicateSection("Code section".into())); - } - self.code.reserve(count as usize); - validator.code_section_start(&range)?; - } - Payload::CodeSectionEntry(function) => { - debug!("Found code section entry"); - let v = validator.code_section_entry(&function)?; - let func_validator = v.into_validator(self.func_validator_allocations.take().unwrap_or_default()); - let (code, allocations) = conversion::convert_module_code(function, func_validator)?; - self.code.push(code); - self.func_validator_allocations = Some(allocations); - } Payload::ImportSection(reader) => { if !self.imports.is_empty() { return Err(ParseError::DuplicateSection("Import section".into())); @@ -160,9 +191,161 @@ impl ModuleReader { debug!("Found custom section"); debug!("Skipping custom section: {:?}", _reader.name()); } + Payload::CodeSectionStart { .. } | Payload::CodeSectionEntry(_) => { + return Err(ParseError::Other("code section payload handled separately".into())); + } Payload::UnknownSection { .. } => return Err(ParseError::UnsupportedSection("Unknown section".into())), section => return Err(ParseError::UnsupportedSection(format!("Unsupported section: {section:?}"))), + } + + Ok(()) + } + + pub(crate) fn begin_code_section( + &mut self, + count: u32, + range: Range<usize>, + size: u32, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<bool> { + debug!("Found code section ({count} functions)"); + if self.has_code_section { + return Err(ParseError::DuplicateSection("Code section".into())); + } + + self.has_code_section = true; + self.code.reserve(count as usize); + validator.code_section_start(&range)?; + + #[cfg(parallel_parser)] + { + let defer = crate::parallel::should_use_parallel(options, count as usize, size as usize); + if defer { + debug!("Queuing {count} functions from {size} byte code section"); + self.pending_functions = Some(Vec::with_capacity(count as usize)); + } + Ok(defer) + } + + #[cfg(not(parallel_parser))] + { + let _ = (size, options); + Ok(false) + } + } + + pub(crate) fn process_inline_code_section_entry( + &mut self, + function: wasmparser::FunctionBody<'_>, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<()> { + debug!("Found code section entry"); + let ordinal = self.code.len(); + let func_to_validate = validator.code_section_entry(&function)?; + let func_validator = + func_to_validate.into_validator(self.func_validator_allocations.take().unwrap_or_default()); + let (code, allocations) = conversion::convert_module_code(function, func_validator)?; + let code = optimize_function_code( + code, + options, + self.function_results(ordinal), + (imported_func_count(&self.imports) + ordinal) as u32, + imported_memory_count(&self.imports), + ); + self.code.push(code); + self.func_validator_allocations = Some(allocations); + Ok(()) + } + + pub(crate) fn process_borrowed_code_section_entry( + &mut self, + function: wasmparser::FunctionBody<'a>, + validator: &mut Validator, + options: &ParserOptions, + ) -> Result<()> { + debug!("Found code section entry"); + + #[cfg(parallel_parser)] + if let Some(pending) = self.pending_functions.as_mut() { + let func_to_validate = validator.code_section_entry(&function)?; + let ordinal = self.code.len() + pending.len(); + let ty_idx = self.code_type_addrs[ordinal]; + pending.push(crate::parallel::PendingFunction { + ordinal, + ty_idx, + func_to_validate, + body: crate::parallel::FunctionBodyInput::Borrowed(function), + }); + return Ok(()); + } + + self.process_inline_code_section_entry(function, validator, options) + } + + #[cfg(parallel_parser)] + pub(crate) fn queue_owned_code_section( + &mut self, + count: u32, + body_offset: usize, + section_bytes: Arc<[u8]>, + validator: &mut Validator, + ) -> Result<()> { + let code_len = self.code.len(); + let pending = self + .pending_functions + .as_mut() + .ok_or_else(|| ParseError::Other("owned code section queued without pending storage".into()))?; + + let mut reader = wasmparser::BinaryReader::new(§ion_bytes, body_offset); + for _ in 0..count { + let body_reader = reader.read_reader()?; + let body_range = body_reader.range(); + let function = wasmparser::FunctionBody::new(body_reader); + let func_to_validate = validator.code_section_entry(&function)?; + let ordinal = code_len + pending.len(); + let ty_idx = self.code_type_addrs[ordinal]; + pending.push(crate::parallel::PendingFunction { + ordinal, + ty_idx, + func_to_validate, + body: crate::parallel::FunctionBodyInput::Owned(crate::parallel::OwnedFunctionBody { + section_bytes: section_bytes.clone(), + body_range: (body_range.start - body_offset)..(body_range.end - body_offset), + body_offset: body_range.start, + }), + }); + } + + if reader.bytes_remaining() != 0 { + return Err(ParseError::ParseError { + message: "trailing bytes at end of section".into(), + offset: reader.original_position(), + }); + } + + Ok(()) + } + + #[cfg(parallel_parser)] + pub(crate) fn process_pending_functions(&mut self, options: &ParserOptions) -> Result<()> { + let Some(pending) = self.pending_functions.take().filter(|pending| !pending.is_empty()) else { + return Ok(()); }; + + self.code.extend(crate::parallel::process_pending( + pending, + options, + &self.func_types, + imported_func_count(&self.imports), + imported_memory_count(&self.imports), + )?); + Ok(()) + } + + #[cfg(not(parallel_parser))] + pub(crate) fn process_pending_functions(&mut self, _options: &ParserOptions) -> Result<()> { Ok(()) } @@ -175,8 +358,7 @@ impl ModuleReader { return Err(ParseError::Other("Code and code type address count mismatch".to_string())); } - let imported_func_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count(); - let import_mem_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32; + let import_mem_count = imported_memory_count(&self.imports); let has_local_mem_export = self.exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= import_mem_count); let has_active_data_segment_on_local_memory = self.data.iter().any(|data| match &data.kind { @@ -193,33 +375,27 @@ impl ModuleReader { } else { LocalMemoryAllocation::Skip }; + let mut funcs = Vec::with_capacity(self.code.len()); - for (func_idx, ((instructions, mut data, locals), ty_idx)) in - self.code.into_iter().zip(self.code_type_addrs).enumerate() - { + for (code, ty_idx) in self.code.into_iter().zip(self.code_type_addrs) { let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone(); let params = ValueCounts::from_iter(ty.params()); let results = ValueCounts::from_iter(ty.results()); - let self_func = (imported_func_count + func_idx) as u32; - let local_mem_alloc = - optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager; - let optimized = optimize::optimize_instructions( - instructions, - &mut data, - options, - results, - self_func, - import_mem_count, - local_mem_alloc, - ); - - if optimized.uses_local_memory { + if code.uses_local_memory { local_memory_allocation = LocalMemoryAllocation::Eager; } funcs.push( - WasmFunction { instructions: optimized.instructions.into(), data, locals, params, results, ty }.into(), + WasmFunction { + instructions: code.instructions.into(), + data: code.data, + locals: code.locals, + params, + results, + ty, + } + .into(), ); } diff --git a/crates/parser/src/optimize.rs b/crates/parser/src/optimize.rs index 5882fea..aad5b1a 100644 --- a/crates/parser/src/optimize.rs +++ b/crates/parser/src/optimize.rs @@ -15,13 +15,11 @@ pub(crate) fn optimize_instructions( function_results: ValueCounts, self_func_addr: u32, imported_memory_count: u32, - track_local_memory_usage: bool, ) -> OptimizeResult { let uses_local_memory = if options.optimize_rewrite() { - rewrite(&mut instructions, function_results, self_func_addr, imported_memory_count, track_local_memory_usage) + rewrite(&mut instructions, function_results, self_func_addr, imported_memory_count) } else { - track_local_memory_usage - && instructions.iter().any(|instr| instr.memory_addr().is_some_and(|mem| mem >= imported_memory_count)) + instructions.iter().any(|instr| instr.memory_addr().is_some_and(|mem| mem >= imported_memory_count)) }; if options.optimize_remove_nop() { @@ -35,7 +33,6 @@ fn rewrite( function_results: ValueCounts, self_func_addr: u32, imported_memory_count: u32, - track_local_memory_usage: bool, ) -> bool { use Instruction::*; let mut uses_local_memory = false; @@ -380,9 +377,24 @@ fn rewrite( ), Jump(ip) => { let target = resolve_jump_target(instrs, ip); - canonicalize_jump_like_with_target(instrs, i, target); + let exit = next_non_nop(instrs, i + 1) as u32; + let body = next_non_nop(instrs, target as usize + 1) as u32; + + match instrs[target as usize] { + JumpCmpLocalLocal32 { target_ip, left, right, op } + if resolve_jump_target(instrs, target_ip) == exit && body > target => + { + instrs[i] = JumpCmpLocalLocal32 { target_ip: body, left, right, op: inverse_cmp_op(op) }; + } + JumpCmpLocalLocal64 { target_ip, left, right, op } + if resolve_jump_target(instrs, target_ip) == exit && body > target => + { + instrs[i] = JumpCmpLocalLocal64 { target_ip: body, left, right, op: inverse_cmp_op(op) }; + } + _ => canonicalize_jump_like_with_target(instrs, i, target), + } } - JumpIfZero(ip) => { + JumpIfZero32(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet32(local), I32Eqz] => { replace!(instrs, i, 2 => [Nop, Nop, JumpIfLocalNonZero32 { target_ip: target, local }]); @@ -438,7 +450,7 @@ fn rewrite( }); canonicalize_jump_like_with_target(instrs, i, target); } - JumpIfNonZero(ip) => { + JumpIfNonZero32(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet32(local), I32Eqz] => { replace!(instrs, i, 2 => [Nop, Nop, JumpIfLocalZero32 { target_ip: target, local }]); @@ -494,16 +506,6 @@ fn rewrite( }); canonicalize_jump_like_with_target(instrs, i, target); } - JumpIfZero32(ip) => { - let target = resolve_jump_target(instrs, ip); - rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalZero32 { target_ip: target, local }); - canonicalize_jump_like_with_target(instrs, i, target); - } - JumpIfNonZero32(ip) => { - let target = resolve_jump_target(instrs, ip); - rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalNonZero32 { target_ip: target, local }); - canonicalize_jump_like_with_target(instrs, i, target); - } JumpIfZero64(ip) => { let target = resolve_jump_target(instrs, ip); rewrite!(instrs, i, [LocalGet64(local)] => JumpIfLocalZero64 { target_ip: target, local }); @@ -561,7 +563,7 @@ fn rewrite( _ => {} } - if track_local_memory_usage && !uses_local_memory { + if !uses_local_memory { uses_local_memory = instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count); } } @@ -807,8 +809,6 @@ fn resolve_jump_target(instrs: &[Instruction], target: u32) -> u32 { fn jump_target(instr: Instruction) -> Option<u32> { Some(match instr { Instruction::Jump(ip) - | Instruction::JumpIfZero(ip) - | Instruction::JumpIfNonZero(ip) | Instruction::JumpIfZero32(ip) | Instruction::JumpIfNonZero32(ip) | Instruction::JumpIfZero64(ip) @@ -830,8 +830,6 @@ fn jump_target(instr: Instruction) -> Option<u32> { fn set_jump_target(instr: &mut Instruction, target: u32) { match instr { Instruction::Jump(ip) - | Instruction::JumpIfZero(ip) - | Instruction::JumpIfNonZero(ip) | Instruction::JumpIfZero32(ip) | Instruction::JumpIfNonZero32(ip) | Instruction::JumpIfZero64(ip) @@ -898,8 +896,6 @@ fn remove_nop(instructions: &mut Vec<Instruction>, function_data: &mut WasmFunct instructions.retain_mut(|instr| { let ip = match instr { Instruction::Jump(ip) - | Instruction::JumpIfZero(ip) - | Instruction::JumpIfNonZero(ip) | Instruction::JumpIfZero32(ip) | Instruction::JumpIfNonZero32(ip) | Instruction::JumpIfZero64(ip) diff --git a/crates/parser/src/parallel.rs b/crates/parser/src/parallel.rs new file mode 100644 index 0000000..778e4ad --- /dev/null +++ b/crates/parser/src/parallel.rs @@ -0,0 +1,174 @@ +use crate::module::{FunctionCode, optimize_function_code}; +use crate::{ParseError, ParserOptions, Result, conversion}; +use alloc::sync::Arc; +use alloc::vec::Vec; +use core::ops::Range; +use tinywasm_types::{FuncType, ValueCounts}; +use wasmparser::{FuncValidatorAllocations, ValidatorResources}; + +pub(crate) enum FunctionBodyInput<'a> { + Borrowed(wasmparser::FunctionBody<'a>), + Owned(OwnedFunctionBody), +} + +pub(crate) struct OwnedFunctionBody { + // A deferred stream code section is copied once, then shared by all queued + // function jobs from that section. + pub section_bytes: Arc<[u8]>, + pub body_range: Range<usize>, + pub body_offset: usize, +} + +pub(crate) struct PendingFunction<'a> { + pub ordinal: usize, + pub ty_idx: u32, + pub func_to_validate: wasmparser::FuncToValidate<ValidatorResources>, + pub body: FunctionBodyInput<'a>, +} + +pub(crate) const MIN_FUNCTIONS: usize = 8; +const MIN_CODE_SECTION_BYTES: usize = 32 * 1024; +const MIN_FUNCTION_BODY_BYTES: usize = 4; + +pub(crate) fn should_parallelize_function(body_len: usize) -> bool { + body_len >= MIN_FUNCTION_BODY_BYTES +} + +pub(crate) fn should_use_parallel(options: &ParserOptions, num_functions: usize, code_section_bytes: usize) -> bool { + if num_functions < MIN_FUNCTIONS || code_section_bytes < MIN_CODE_SECTION_BYTES { + return false; + } + + worker_count(options, num_functions) > 1 +} + +fn worker_count(options: &ParserOptions, num_functions: usize) -> usize { + let requested = options + .parser_threads() + .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1)) + .max(1); + + requested.min(num_functions).max(1) +} + +fn body_len(body: &FunctionBodyInput<'_>) -> usize { + match body { + FunctionBodyInput::Borrowed(func) => func.as_bytes().len(), + FunctionBodyInput::Owned(body) => body.body_range.len(), + } +} + +fn process_function_job( + job: PendingFunction<'_>, + options: &ParserOptions, + func_types: &[Arc<FuncType>], + imported_func_count: usize, + imported_memory_count: u32, +) -> Result<(usize, FunctionCode)> { + let validator = job.func_to_validate.into_validator(FuncValidatorAllocations::default()); + let (code, _allocations) = match job.body { + FunctionBodyInput::Borrowed(func) => conversion::convert_module_code(func, validator)?, + FunctionBodyInput::Owned(body) => { + let reader = wasmparser::BinaryReader::new(&body.section_bytes[body.body_range], body.body_offset); + let func = wasmparser::FunctionBody::new(reader); + conversion::convert_module_code(func, validator)? + } + }; + + let ty = func_types.get(job.ty_idx as usize).expect("No func type for func, this is a bug"); + let code = optimize_function_code( + code, + options, + ValueCounts::from_iter(ty.results()), + (imported_func_count + job.ordinal) as u32, + imported_memory_count, + ); + + Ok((job.ordinal, code)) +} + +pub(crate) fn process_pending( + pending: Vec<PendingFunction<'_>>, + options: &ParserOptions, + func_types: &[Arc<FuncType>], + imported_func_count: usize, + imported_memory_count: u32, +) -> Result<Vec<FunctionCode>> { + if pending.is_empty() { + return Ok(Vec::new()); + } + + let (small_jobs, large_jobs): (Vec<_>, Vec<_>) = + pending.into_iter().partition(|job| !should_parallelize_function(body_len(&job.body))); + + let mut codes = small_jobs + .into_iter() + .map(|job| process_function_job(job, options, func_types, imported_func_count, imported_memory_count)) + .collect::<Result<Vec<_>>>()?; + + if large_jobs.is_empty() { + codes.sort_by_key(|(ordinal, _)| *ordinal); + return Ok(codes.into_iter().map(|(_, code)| code).collect()); + } + + let num_workers = worker_count(options, large_jobs.len()); + if num_workers == 1 { + codes.extend( + large_jobs + .into_iter() + .map(|job| process_function_job(job, options, func_types, imported_func_count, imported_memory_count)) + .collect::<Result<Vec<_>>>()?, + ); + codes.sort_by_key(|(ordinal, _)| *ordinal); + return Ok(codes.into_iter().map(|(_, code)| code).collect()); + } + + let chunk_size = large_jobs.len().div_ceil(num_workers); + let chunks = { + let mut chunks = Vec::with_capacity(num_workers); + let mut iter = large_jobs.into_iter(); + while let Some(first) = iter.next() { + let mut chunk = alloc::vec![first]; + for _ in 1..chunk_size { + match iter.next() { + Some(job) => chunk.push(job), + None => break, + } + } + chunks.push(chunk); + } + chunks + }; + + let results: Vec<Result<(usize, FunctionCode)>> = std::thread::scope(|s| { + let handles: Vec<_> = chunks + .into_iter() + .map(|chunk| { + s.spawn(move || { + chunk + .into_iter() + .map(|job| { + process_function_job(job, options, func_types, imported_func_count, imported_memory_count) + }) + .collect::<Vec<_>>() + }) + }) + .collect(); + + handles + .into_iter() + .flat_map(|handle| match handle.join() { + Ok(results) => results, + Err(_) => alloc::vec![Err(ParseError::Other("worker thread panicked".into()))], + }) + .collect() + }); + + for result in results { + let (ordinal, code) = result?; + codes.push((ordinal, code)); + } + + codes.sort_by_key(|(ordinal, _)| *ordinal); + Ok(codes.into_iter().map(|(_, code)| code).collect()) +} diff --git a/crates/parser/src/visit.rs b/crates/parser/src/visit.rs index 64821b4..f8d5cfd 100644 --- a/crates/parser/src/visit.rs +++ b/crates/parser/src/visit.rs @@ -282,7 +282,7 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild } fn visit_if(&mut self, _ty: wasmparser::BlockType) -> Self::Output { - self.instructions.push(Instruction::JumpIfZero(0)); + self.instructions.push(Instruction::JumpIfZero32(0)); self.ctx_stack.push(LoweringCtx { kind: BlockKind::If, has_else: false, @@ -325,7 +325,7 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild fn visit_br_if(&mut self, depth: u32) -> Self::Output { let cond_jump_ip = self.instructions.len(); - self.instructions.push(Instruction::JumpIfZero(0)); + self.instructions.push(Instruction::JumpIfZero32(0)); let branch_side_start = self.instructions.len(); self.emit_dropkeep_to_label(depth); @@ -333,7 +333,7 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild if self.instructions.len() == branch_side_start && let Some(ctx_idx) = self.get_ctx_idx(depth) { - self.instructions[cond_jump_ip] = Instruction::JumpIfNonZero(0); + self.instructions[cond_jump_ip] = Instruction::JumpIfNonZero32(0); self.ctx_stack[ctx_idx].branch_jumps.push(cond_jump_ip); return; } @@ -608,7 +608,7 @@ impl<R: WasmModuleResources> FunctionBuilder<R> { fn patch_jump(&mut self, jump_ip: usize, target: usize) { match &mut self.instructions[jump_ip] { - Instruction::Jump(ip) | Instruction::JumpIfNonZero(ip) => { + Instruction::Jump(ip) | Instruction::JumpIfNonZero32(ip) => { *ip = target as u32; } _ => {} @@ -616,7 +616,7 @@ impl<R: WasmModuleResources> FunctionBuilder<R> { } fn patch_jump_if_zero(&mut self, jump_ip: usize, target: usize) { - if let Instruction::JumpIfZero(ip) = &mut self.instructions[jump_ip] { + if let Instruction::JumpIfZero32(ip) = &mut self.instructions[jump_ip] { *ip = target as u32; } } |