From 684d5a04a928ea4132d0c5e61bb4086fac9feb22 Mon Sep 17 00:00:00 2001 From: Henry Date: Sun, 19 Apr 2026 15:12:25 +0200 Subject: feat: lazy memory allocation, fix instruction rewrite order Signed-off-by: Henry --- crates/parser/src/lib.rs | 26 ++++++++++++- crates/parser/src/module.rs | 87 ++++++++++++++++++++++++++++++++----------- crates/parser/src/optimize.rs | 74 +++++++++++++++++++++++------------- 3 files changed, 138 insertions(+), 49 deletions(-) (limited to 'crates/parser') diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs index ad58cc2..0b5af2e 100644 --- a/crates/parser/src/lib.rs +++ b/crates/parser/src/lib.rs @@ -43,8 +43,30 @@ pub use tinywasm_types::TinyWasmModule; /// Parser optimization and lowering options. #[non_exhaustive] -#[derive(Debug, Clone, Default)] -pub struct ParserOptions {} +#[derive(Debug, Clone)] +pub struct ParserOptions { + /// Whether to optimize local memory allocation by skipping allocation of unused local memories. + pub optimize_local_memory_allocation: bool, +} + +impl Default for ParserOptions { + fn default() -> Self { + Self { optimize_local_memory_allocation: true } + } +} + +impl ParserOptions { + /// Enable or disable the optimization that skips allocating unused local memories. + pub const fn with_local_memory_allocation_optimization(mut self, enabled: bool) -> Self { + self.optimize_local_memory_allocation = enabled; + self + } + + /// Returns whether unused local memory allocation optimization is enabled. + pub const fn optimize_local_memory_allocation(&self) -> bool { + self.optimize_local_memory_allocation + } +} /// A WebAssembly parser #[derive(Debug, Default)] diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs index 461f79f..9c4d6b4 100644 --- a/crates/parser/src/module.rs +++ b/crates/parser/src/module.rs @@ -167,7 +167,7 @@ impl ModuleReader { Ok(()) } - pub(crate) fn into_module(self, _options: &ParserOptions) -> Result { + pub(crate) fn into_module(self, options: &ParserOptions) -> Result { if !self.end_reached { return Err(ParseError::EndNotReached); } @@ -176,28 +176,73 @@ 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 funcs = self.code.into_iter().zip(self.code_type_addrs).enumerate().map( - |(func_idx, ((instructions, mut data, locals), ty_idx))| { - 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 self_func = (imported_func_count + func_idx) as u32; - let instructions = optimize::optimize_instructions(instructions, &mut data, self_func); - WasmFunction { instructions: ArcSlice::from(instructions), data, locals, params, ty } - }, - ); + let Self { + start_func, + func_types, + code_type_addrs, + exports, + code, + globals, + table_types, + memory_types, + imports, + data, + elements, + .. + } = self; + + let imported_func_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count(); + let imported_memory_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32; + let has_local_memory_export = + exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= imported_memory_count); + let has_active_data_segment_on_local_memory = data.iter().any(|data| match &data.kind { + DataKind::Active { mem, .. } => *mem >= imported_memory_count, + DataKind::Passive => false, + }); + let optimize_local_memory_allocation = options.optimize_local_memory_allocation(); + let mut local_memory_allocation = if memory_types.is_empty() { + LocalMemoryAllocation::Skip + } else if !optimize_local_memory_allocation || has_active_data_segment_on_local_memory { + LocalMemoryAllocation::Eager + } else if has_local_memory_export { + LocalMemoryAllocation::Lazy + } else { + LocalMemoryAllocation::Skip + }; + let mut funcs = Vec::with_capacity(code.len()); + + for (func_idx, ((instructions, mut data, locals), ty_idx)) in code.into_iter().zip(code_type_addrs).enumerate() + { + let ty = 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 self_func = (imported_func_count + func_idx) as u32; + let optimized = optimize::optimize_instructions( + instructions, + &mut data, + self_func, + imported_memory_count, + optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager, + ); + + if optimized.uses_local_memory { + local_memory_allocation = LocalMemoryAllocation::Eager; + } + + funcs.push(WasmFunction { instructions: ArcSlice::from(optimized.instructions), data, locals, params, ty }); + } Ok(TinyWasmModule { - funcs: funcs.collect(), - func_types: self.func_types.into(), - globals: self.globals.into(), - table_types: self.table_types.into(), - imports: self.imports.into(), - start_func: self.start_func, - data: self.data.into(), - exports: self.exports.into(), - elements: self.elements.into(), - memory_types: self.memory_types.into(), + funcs: funcs.into(), + func_types: func_types.into(), + globals: globals.into(), + table_types: table_types.into(), + imports: imports.into(), + start_func, + data: data.into(), + exports: exports.into(), + elements: elements.into(), + memory_types: memory_types.into(), + local_memory_allocation, }) } } diff --git a/crates/parser/src/optimize.rs b/crates/parser/src/optimize.rs index 0539b61..beb944d 100644 --- a/crates/parser/src/optimize.rs +++ b/crates/parser/src/optimize.rs @@ -2,18 +2,32 @@ use crate::macros::optimize::*; use alloc::vec::Vec; use tinywasm_types::{CmpOp, Instruction, WasmFunctionData}; +pub(crate) struct OptimizeResult { + pub(crate) instructions: Vec, + pub(crate) uses_local_memory: bool, +} + pub(crate) fn optimize_instructions( mut instructions: Vec, function_data: &mut WasmFunctionData, self_func_addr: u32, -) -> Vec { - rewrite(&mut instructions, self_func_addr); + imported_memory_count: u32, + track_local_memory_usage: bool, +) -> OptimizeResult { + let uses_local_memory = rewrite(&mut instructions, self_func_addr, imported_memory_count, track_local_memory_usage); remove_nop(&mut instructions, function_data); - instructions + OptimizeResult { instructions, uses_local_memory } } -fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) { +fn rewrite( + instrs: &mut [Instruction], + self_func_addr: u32, + imported_memory_count: u32, + track_local_memory_usage: bool, +) -> bool { use Instruction::*; + let mut uses_local_memory = false; + for i in 0..instrs.len() { match instrs[i] { LocalCopy32(a, b) if a == b => instrs[i] = Nop, @@ -22,14 +36,14 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) { Call(addr) if addr == self_func_addr => instrs[i] = CallSelf, ReturnCall(addr) if addr == self_func_addr => instrs[i] = ReturnCallSelf, I32Add => { - rewrite!(instrs, i, [I32Const(c)] => AddConst32(c)); rewrite!(instrs, i, [LocalGet32(a), LocalGet32(b)] => AddLocalLocal32(a, b)); rewrite!(instrs, i, [LocalGet32(local), I32Const(c)] => [ Nop, LocalGet32(local), AddConst32(c)]); + rewrite!(instrs, i, [I32Const(c)] => AddConst32(c)); } I64Add => { - rewrite!(instrs, i, [I64Const(c)] => AddConst64(c)); rewrite!(instrs, i, [LocalGet64(a), LocalGet64(b)] => AddLocalLocal64(a, b)); rewrite!(instrs, i, [LocalGet64(local), I64Const(c)] => [ Nop, LocalGet64(local), AddConst64(c)]); + rewrite!(instrs, i, [I64Const(c)] => AddConst64(c)); } I64Rotl => rewrite!(instrs, i, [I64Xor, I64Const(c)] => XorRotlConst64(c)), I32Store(memarg) => { @@ -69,6 +83,7 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) { rewrite!(instrs, i, [LocalGet32(src)] => if src == dst { Nop } else { LocalCopy32(src, dst) }); rewrite!(instrs, i, [I32Const(c)] => SetLocalConst32(dst, c)); rewrite!(instrs, i, [F32Const(c)] => SetLocalConst32(dst, i32::from_ne_bytes(c.to_bits().to_ne_bytes()))); + rewrite!(instrs, i, [AddLocalLocal32(a, b)] => AddLocalLocalSet32(a, b, dst)); rewrite!(instrs, i, [LocalGet32(src), AddConst32(c)] if (src == dst) => AddLocalConst32(dst, c)); rewrite!(instrs, i, [LoadLocal32(memarg, addr)] if (let Ok(dst) = u8::try_from(dst)) => LoadLocalSet32(memarg, addr, dst)); rewrite!(instrs, i, @@ -81,6 +96,7 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) { rewrite!(instrs, i, [LocalGet64(src)] => if src == dst { Nop } else { LocalCopy64(src, dst) }); rewrite!(instrs, i, [I64Const(c)] => SetLocalConst64(dst, c)); rewrite!(instrs, i, [F64Const(c)] => SetLocalConst64(dst, i64::from_ne_bytes(c.to_bits().to_ne_bytes()))); + rewrite!(instrs, i, [AddLocalLocal64(a, b)] => AddLocalLocalSet64(a, b, dst)); rewrite!(instrs, i, [LocalGet64(src), AddConst64(c)] if (src == dst) => AddLocalConst64(dst, c) @@ -130,62 +146,68 @@ fn rewrite(instrs: &mut [Instruction], self_func_addr: u32) { replace!(instrs, i, 1 => [Nop, JumpIfNonZero(ip)]); continue; }); - rewrite!(instrs, i, [cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) => - JumpCmpStackConst32 { target_ip: ip, imm, op: inverse_cmp_op(op) } - ); - rewrite!(instrs, i, [cmp, I64Const(imm)] if (let Some(op) = cmp_op_64(cmp)) => - JumpCmpStackConst64 { target_ip: ip, imm, op: inverse_cmp_op(op) } - ); rewrite!(instrs, i, - [LocalGet32(local), cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) => + [LocalGet32(local), I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) => JumpCmpLocalConst32 { target_ip: ip, local, imm, op: inverse_cmp_op(op) } ); rewrite!(instrs, i, - [LocalGet64(local), cmp, I64Const(imm)] if + [LocalGet64(local), I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp) && let Ok(imm) = i32::try_from(imm)) => JumpCmpLocalConst64 { target_ip: ip, local, imm, op: inverse_cmp_op(op) } ); rewrite!(instrs, i, - [LocalGet32(left), cmp, LocalGet32(right)] if (let Some(op) = cmp_op(cmp)) => + [LocalGet32(left), LocalGet32(right), cmp] if (let Some(op) = cmp_op(cmp)) => JumpCmpLocalLocal32 { target_ip: ip, left, right, op: inverse_cmp_op(op) } ); rewrite!(instrs, i, - [LocalGet64(left), cmp, LocalGet64(right)] if (let Some(op) = cmp_op_64(cmp)) => + [LocalGet64(left), LocalGet64(right), cmp] if (let Some(op) = cmp_op_64(cmp)) => JumpCmpLocalLocal64 { target_ip: ip, left, right, op: inverse_cmp_op(op) } ); + rewrite!(instrs, i, [I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) => + JumpCmpStackConst32 { target_ip: ip, imm, op: inverse_cmp_op(op) } + ); + rewrite!(instrs, i, [I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp)) => + JumpCmpStackConst64 { target_ip: ip, imm, op: inverse_cmp_op(op) } + ); } JumpIfNonZero(ip) => { rewrite!(instrs, i, [I32Eqz] => { replace!(instrs, i, 1 => [Nop, JumpIfZero(ip)]); continue; }); - rewrite!(instrs, i, [cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) => - JumpCmpStackConst32 { target_ip: ip, imm, op } - ); - rewrite!(instrs, i, [cmp, I64Const(imm)] if (let Some(op) = cmp_op_64(cmp)) => - JumpCmpStackConst64 { target_ip: ip, imm, op } - ); rewrite!(instrs, i, - [LocalGet32(local), cmp, I32Const(imm)] if (let Some(op) = cmp_op(cmp)) => + [LocalGet32(local), I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) => JumpCmpLocalConst32 { target_ip: ip, local, imm, op } ); rewrite!(instrs, i, - [LocalGet64(local), cmp, I64Const(imm)] if + [LocalGet64(local), I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp) && let Ok(imm) = i32::try_from(imm)) => JumpCmpLocalConst64 { target_ip: ip, local, imm, op } ); rewrite!(instrs, i, - [LocalGet32(left), cmp, LocalGet32(right)] if (let Some(op) = cmp_op(cmp)) => + [LocalGet32(left), LocalGet32(right), cmp] if (let Some(op) = cmp_op(cmp)) => JumpCmpLocalLocal32 { target_ip: ip, left, right, op } ); rewrite!(instrs, i, - [LocalGet64(left), cmp, LocalGet64(right)] if (let Some(op) = cmp_op_64(cmp)) => + [LocalGet64(left), LocalGet64(right), cmp] if (let Some(op) = cmp_op_64(cmp)) => JumpCmpLocalLocal64 { target_ip: ip, left, right, op } ); + rewrite!(instrs, i, [I32Const(imm), cmp] if (let Some(op) = cmp_op(cmp)) => + JumpCmpStackConst32 { target_ip: ip, imm, op } + ); + rewrite!(instrs, i, [I64Const(imm), cmp] if (let Some(op) = cmp_op_64(cmp)) => + JumpCmpStackConst64 { target_ip: ip, imm, op } + ); } _ => {} } + + if track_local_memory_usage { + uses_local_memory |= instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count); + } } + + uses_local_memory } fn cmp_op(instr: Instruction) -> Option { -- cgit v1.3.1