diff --git a/DEBUGGING.md b/DEBUGGING.md new file mode 100644 index 00000000..aae4556b --- /dev/null +++ b/DEBUGGING.md @@ -0,0 +1,48 @@ +# Debugging Guide + +This document outlines the workflow and tools for debugging the Webschembly compiler and JIT. + +## Running Tests with Logs + +To debug specific Scheme files or tests, use `just run` within the `webschembly-js` directory. + +### Environment Variables + +- `LOG_STDOUT=1`: Enables `log::debug!` output from the runtime and compiler (if configured). Use this to see runtime execution trace, JIT instantiation events, and error messages. +- `LOG=1`: Dumps the generated Intermediate Representation (IR) and Wasm binaries to the `webschembly-js/log/` directory. + +### Command Examples + +```bash +# Run a specific fixture with runtime logs +cd webschembly-js +just LOG_STDOUT=1 run ./fixtures/rec.scm + +# Run with IR dumping +just LOG=1 run ./fixtures/rec.scm +``` + +## Analyzing Generated IR + +When `LOG=1` is used, the `webschembly-js/log/` directory will contain files named with the timestamp and content description. + +- `checks-TIMESTAMP-filename-0.ir`: Typically the main module IR (initial compilation). +- `checks-TIMESTAMP-filename-N.ir`: IR for JIT-compiled functions or stubs. `instantiate_func` or `instantiate_bb` calls in the runtime logs (seen with `LOG_STDOUT=1`) will reference `module_id` and `func_id`, which correlate to these files (though the mapping requires checking the `instantiate` log id vs file index). + +**Tip**: Look for "instantiate: id:X" in the runtime logs. The corresponding IR file is often suffix `-X.ir`. + +## Common Issues & Fixes + +### "call target is not a closure" + +This error occurs when the compiled code attempts to invoke a value as a closure, but the compile-time or run-time check fails. + +- **Runtime**: The value on the stack is not a closure object (e.g., encoded as `val_type` mismatch). +- **Compile-time (Optimization)**: If the error appears unconditionally in the IR (e.g., `error "call target is not a closure"`), it means `constant_folding` or another pass determined the check `Is(Closure(None), target)` is false. + - _Watch out for_: Type mismatches in `InstrKind::Is`. Ensure `Closure(None)` (generic check) correctly matches specialized types like `Closure(Some(C))` (constant closure). Using `.remove_constant()` on types before comparison is crucial in `ssa_optimizer.rs`. + +## JIT Optimization Logic + +- **`propagate_types` Pass**: Analyzes dataflow to identify constant closures (`Closure(Some(C))`). Logic in `src/ir_processor/propagate_types.rs`. +- **Specialization**: `jit_func.rs` uses available type info (from `locals`) to unbox closure entries or arguments. +- **Constant Folding**: `ssa_optimizer.rs` folds constants. Be careful with strict equality checks on specialized types. diff --git a/webschembly-compiler-crates/ir/src/id.rs b/webschembly-compiler-crates/ir/src/id.rs index 277f2549..0c5c771e 100644 --- a/webschembly-compiler-crates/ir/src/id.rs +++ b/webschembly-compiler-crates/ir/src/id.rs @@ -161,3 +161,73 @@ impl fmt::Display for Display<'_, JitBasicBlockId> { Ok(()) } } + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct ClosureEnvIndex(pub usize); + +impl ClosureEnvIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ClosureEnvIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ClosureEnvIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "closure_env{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct ClosureArgIndex(pub usize); + +impl ClosureArgIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ClosureArgIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ClosureArgIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "closure_arg{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, From, Into)] +pub struct BBIndex(pub usize); + +impl BBIndex { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, BBIndex> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, BBIndex> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "bb_index{}", self.value.0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ConstantClosure { + pub module_id: JitModuleId, + pub func_id: JitFuncId, + pub env_index: ClosureEnvIndex, +} + +impl ConstantClosure { + pub fn display<'a>(&self, meta: &'a Meta) -> Display<'a, ConstantClosure> { + Display { value: *self, meta } + } +} + +impl fmt::Display for Display<'_, ConstantClosure> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "constant_closure({}, {}, {})", + self.value.module_id.display(self.meta), + self.value.func_id.display(self.meta), + self.value.env_index.display(self.meta) + ) + } +} diff --git a/webschembly-compiler-crates/ir/src/typ.rs b/webschembly-compiler-crates/ir/src/typ.rs index 7535a73f..9091b111 100644 --- a/webschembly-compiler-crates/ir/src/typ.rs +++ b/webschembly-compiler-crates/ir/src/typ.rs @@ -3,6 +3,8 @@ TypeもしくはmutableなTypeを表す Type自体にRefを含めて再帰的にしてしまうと無限種類の型を作れるようになってしまうので、IRではそれを避けるためこのような構造になっている TODO: LocalTypeという名前は適切ではない */ +use crate::id::ConstantClosure; + #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] pub enum LocalType { #[display("ref<{}>", _0)] @@ -38,6 +40,14 @@ impl LocalType { _ => None, } } + + pub fn remove_constant(self) -> Self { + match self { + LocalType::Ref(typ) => LocalType::Ref(typ.remove_constant()), + LocalType::Type(typ) => LocalType::Type(typ.remove_constant()), + _ => self, + } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] @@ -55,6 +65,13 @@ impl Type { Type::Obj => None, } } + + pub fn remove_constant(self) -> Self { + match self { + Type::Val(val_type) => Type::Val(val_type.remove_constant()), + Type::Obj => Type::Obj, + } + } } impl From for Type { @@ -88,7 +105,16 @@ pub enum ValType { #[display("uvector<{0}>", _0)] UVector(UVectorKind), #[display("closure")] - Closure, + Closure(Option), +} + +impl ValType { + pub fn remove_constant(self) -> Self { + match self { + ValType::Closure(_) => ValType::Closure(None), + _ => self, + } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, derive_more::Display)] diff --git a/webschembly-compiler/src/compiler.rs b/webschembly-compiler/src/compiler.rs index f7b94e7f..1121fc2e 100644 --- a/webschembly-compiler/src/compiler.rs +++ b/webschembly-compiler/src/compiler.rs @@ -8,9 +8,7 @@ use crate::ir_processor::optimizer::remove_unused_local; use crate::ir_processor::register_allocation::register_allocation; use crate::ir_processor::ssa::split_critical_edges; use crate::ir_processor::ssa::{debug_assert_ssa, remove_phi}; -use crate::ir_processor::ssa_optimizer::ModuleInliner; use crate::ir_processor::ssa_optimizer::SsaOptimizerConfig; -use crate::ir_processor::ssa_optimizer::inlining; use crate::ir_processor::ssa_optimizer::ssa_optimize; use crate::jit::{Jit, JitConfig}; use crate::lexer; @@ -133,8 +131,8 @@ impl Compiler { &mut self.global_manager, module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), ); preprocess_module(&mut module); @@ -167,10 +165,10 @@ impl Compiler { let mut module = jit.instantiate_bb( module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), bb_id, - crate::jit::bb_index_manager::BBIndex(index), + ir::BBIndex(index), &mut self.global_manager, ); preprocess_module(&mut module); @@ -211,12 +209,12 @@ impl Compiler { &mut self.global_manager, module_id, func_id, - crate::jit::env_index_manager::EnvIndex(env_index), - crate::jit::closure_global_layout::ClosureIndex(func_index), + ir::ClosureEnvIndex(env_index), + ir::ClosureArgIndex(func_index), bb_id, kind, ir::BasicBlockId::from(source_bb_id), - crate::jit::bb_index_manager::BBIndex(source_index), + ir::BBIndex(source_index), ) .map(|mut module| { preprocess_module(&mut module); @@ -243,13 +241,8 @@ fn preprocess_module(module: &mut ir::Module) { } fn optimize_module(module: &mut ir::Module, config: SsaOptimizerConfig) { - let mut module_inliner = ModuleInliner::new(module); - let n = 5; - for i in 0..n { - if config.enable_inlining { - // inliningはInstrKind::Closureのfunc_idに依存しているので、JIT後のモジュールには使えない - inlining(module, &mut module_inliner, i == n - 1); - } + let n = 10; + for _ in 0..n { for func in module.funcs.values_mut() { ssa_optimize( func, @@ -259,6 +252,9 @@ fn optimize_module(module: &mut ir::Module, config: SsaOptimizerConfig) { }, ); } + if config.enable_inlining { + crate::ir_processor::inline::inline_module(module); + } } } @@ -273,6 +269,8 @@ fn postprocess(module: &mut ir::Module, global_manager: &mut GlobalManager) { register_allocation(func); remove_unused_local(func); + + crate::ir_processor::remove_constant::remove_constant(func); } // モジュールごとにグローバルを真面目に管理するのは大変なのでここで計算 diff --git a/webschembly-compiler/src/ir_generator/module_generator.rs b/webschembly-compiler/src/ir_generator/module_generator.rs index f01608e2..a745c7d2 100644 --- a/webschembly-compiler/src/ir_generator/module_generator.rs +++ b/webschembly-compiler/src/ir_generator/module_generator.rs @@ -266,7 +266,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { let bb_entry = self.builder.bbs.allocate_key(); self.builder.current_bb_id = Some(bb_entry); - let self_closure = self.builder.local(Type::Val(ValType::Closure)); + let self_closure = self.builder.local(Type::Val(ValType::Closure(None))); let args = self.builder.local(LocalType::VariadicArgs); let args_len_local = self.builder.local(Type::Val(ValType::Int)); let expected_args_len_local = self.builder.local(Type::Val(ValType::Int)); @@ -559,7 +559,13 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ast::Expr::Lambda(x, lambda) => { let func_id = self.module_generator.gen_func(x, ast.span, lambda); let func_local = self.builder.local(LocalType::FuncRef); - let val_type_local = self.builder.local(Type::Val(ValType::Closure)); + let val_type_local = + self.builder + .local(Type::Val(ValType::Closure(Some(ConstantClosure { + module_id: self.module_generator.id, + func_id: JitFuncId::from(func_id), + env_index: ClosureEnvIndex(0), + })))); self.builder.exprs.push(Instr { local: Some(func_local), kind: InstrKind::FuncRef(func_id), @@ -617,7 +623,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { }); self.builder.exprs.push(Instr { local: result, - kind: InstrKind::ToObj(ValType::Closure, val_type_local), + kind: InstrKind::ToObj(ValType::Closure(None), val_type_local), }); } ast::Expr::If(_, ast::If { cond, then, els }) => { @@ -887,7 +893,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { let is_closure_local = self.builder.local(Type::Val(ValType::Bool)); self.builder.exprs.push(Instr { local: Some(is_closure_local), - kind: InstrKind::Is(ValType::Closure, obj_func_local), + kind: InstrKind::Is(ValType::Closure(None), obj_func_local), }); let then_bb_id = self.builder.bbs.allocate_key(); @@ -912,10 +918,10 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { self.builder.current_bb_id = Some(then_bb_id); - let closure_local = self.builder.local(ValType::Closure); + let closure_local = self.builder.local(ValType::Closure(None)); self.builder.exprs.push(Instr { local: Some(closure_local), - kind: InstrKind::FromObj(ValType::Closure, obj_func_local), + kind: InstrKind::FromObj(ValType::Closure(None), obj_func_local), }); let args_local = self.builder.local(LocalType::VariadicArgs); @@ -1449,7 +1455,7 @@ impl BuiltinConversionRule { ir_gen: |ctx, arg1| { ctx.builder.exprs.push(Instr { local: Some(ctx.dest), - kind: InstrKind::Is(ValType::Closure, arg1), + kind: InstrKind::Is(ValType::Closure(None), arg1), }); }, }], diff --git a/webschembly-compiler/src/ir_processor/desugar.rs b/webschembly-compiler/src/ir_processor/desugar.rs index 9ea73f51..2219bd54 100644 --- a/webschembly-compiler/src/ir_processor/desugar.rs +++ b/webschembly-compiler/src/ir_processor/desugar.rs @@ -92,7 +92,7 @@ fn desugar_call_closure( func_type: FuncType { args: { let mut args = Vec::new(); - args.push(ValType::Closure.into()); + args.push(ValType::Closure(None).into()); args.extend(call_closure.arg_types); args }, diff --git a/webschembly-compiler/src/ir_processor/inline.rs b/webschembly-compiler/src/ir_processor/inline.rs new file mode 100644 index 00000000..a422bd85 --- /dev/null +++ b/webschembly-compiler/src/ir_processor/inline.rs @@ -0,0 +1,528 @@ +use rustc_hash::FxHashMap; +use vec_map::VecMap; +use webschembly_compiler_ir::*; + +use crate::ir_processor::optimizer::remove_unreachable_bb; + +const LIMIT_BB_COUNT: usize = 100; + +pub fn inline_module(module: &mut Module) { + let mut global_map = FxHashMap::default(); + let entry_func = &module.funcs[module.entry]; + // Scan entry function for GlobalSet with ConstantClosure + for bb in entry_func.bbs.values() { + for instr in &bb.instrs { + if let InstrKind::GlobalSet(global_id, val_local) = instr.kind + && let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + entry_func.locals[val_local].typ + { + global_map.insert(global_id, constant); + } + } + } + + let mut new_funcs = VecMap::new(); + + for (func_id_usize, func) in module.funcs.iter() { + let _func_id = func_id_usize; + let mut new_func = func.clone(); + run_inlining(&mut new_func, module, &global_map); + new_funcs.insert(func_id_usize, new_func); + } + + module.funcs = new_funcs; +} + +struct InlineContext<'a> { + module: &'a Module, + tail_instances: FxHashMap, +} + +#[derive(Debug, Clone)] +struct TailCallInfo { + entry_bb: BasicBlockId, + arg_phis: Vec, +} + +fn run_inlining( + func: &mut Func, + module: &Module, + global_map: &FxHashMap, +) { + let mut ctx = InlineContext { + module, + tail_instances: FxHashMap::default(), + }; + + let mut worklist: Vec = func.bbs.keys().collect(); + + while let Some(bb_id) = worklist.pop() { + if func.bbs.iter().count() > LIMIT_BB_COUNT { + log::debug!("BB limit reached for func {:?}", func.id); + break; + } + if !func.bbs.contains_key(bb_id) { + continue; + } + + // 1. Check Non-Tail Calls (Instrs) + let mut call_found = None; + { + let bb = &func.bbs[bb_id]; + for (idx, instr) in bb.instrs.iter().enumerate() { + if let InstrKind::CallClosure(call) = &instr.kind { + let closure_local = call.closure; + let mut constant_opt = None; + + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + func.locals[closure_local].typ + { + constant_opt = Some(constant); + } else if let Some(def_instr) = find_local_def(func, closure_local) + && let InstrKind::GlobalGet(global_id) = def_instr.kind + && let Some(&constant) = global_map.get(&global_id) + { + constant_opt = Some(constant); + } + + if let Some(constant) = constant_opt { + call_found = Some((idx, constant, call.clone())); + break; + } + } + } + } + + if let Some((idx, constant, call)) = call_found { + log::debug!( + "Non-tail call candidate found in BB {:?} to func {:?}", + bb_id, + constant.func_id + ); + let result_local = func.bbs[bb_id].instrs[idx].local; + let continuation_bb_id = inline_non_tail( + func, + &mut ctx, + &mut worklist, + bb_id, + idx, + constant, + &call, + result_local, + ); + worklist.push(continuation_bb_id); + continue; + } + + // 2. Check Tail Call (Terminator) + let mut tail_call_found = None; + { + let bb = &func.bbs[bb_id]; + if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = bb.terminator() { + let closure_local = call.closure; + let mut constant_opt = None; + + if let LocalType::Type(Type::Val(ValType::Closure(Some(constant)))) = + func.locals[closure_local].typ + { + constant_opt = Some(constant); + } else if let Some(def_instr) = find_local_def(func, closure_local) + && let InstrKind::GlobalGet(global_id) = def_instr.kind + && let Some(&constant) = global_map.get(&global_id) + { + constant_opt = Some(constant); + } + + if let Some(constant) = constant_opt { + tail_call_found = Some((constant, call.clone())); + } + } + } + + if let Some((constant, call)) = tail_call_found { + log::debug!( + "Tail call candidate found in BB {:?} to func {:?}", + bb_id, + constant.func_id + ); + inline_tail(func, &mut ctx, &mut worklist, bb_id, constant, &call); + continue; + } + } + + log::debug!( + "Inlining finished for func {:?}. Final BB count: {}", + func.id, + func.bbs.iter().count() + ); + remove_unreachable_bb(func); + + // Verify Phi position + for (id, bb) in func.bbs.iter() { + let mut phi_mode = true; + for (idx, instr) in bb.instrs.iter().enumerate() { + if let InstrKind::Phi { .. } = instr.kind { + if !phi_mode { + panic!( + "INLINE_MODULE_END: Phi at non-start! BB: {:?}, Index: {}, Instrs: {:#?}", + id, idx, bb.instrs + ); + } + } else if !matches!(instr.kind, InstrKind::Nop) { + phi_mode = false; + } + } + } +} + +fn find_local_def(func: &Func, local: LocalId) -> Option<&Instr> { + // Simple scan for single definition (SSA-like), but inefficient. + // However, for GlobalGet, it's usually near the top. + // Optimally, use DefUseChain or similar, but we don't have it here easily. + // We scan all blocks? No, that's too slow. + // But typically instructions are defined before use in the same block or dominator. + // For now, scan all instructions in all blocks (very slow!). + // BETTER: Build a def map at start of inlining? + // OR: Just scan the current block backwards? GlobalGet is usually in the same block for simple code. + // Let's scan ALL blocks for now as a quick fix, optimizing later if needed. + // Actually, `webschembly_compiler_ir` might have a helper? + // I'll implement a simple full scan. + for bb in func.bbs.values() { + for instr in &bb.instrs { + if instr.local == Some(local) { + return Some(instr); + } + } + } + None +} + +fn inline_non_tail( + func: &mut Func, + ctx: &mut InlineContext, + worklist: &mut Vec, + caller_bb_id: BasicBlockId, + instr_idx: usize, + constant: ConstantClosure, + call: &InstrCallClosure, + result_local: Option, +) -> BasicBlockId { + let continuation_bb_id = func.bbs.allocate_key(); + let continuation_bb_id = continuation_bb_id; + let original_terminator = func.bbs[caller_bb_id].terminator().clone(); + + let instrs_after = func.bbs[caller_bb_id].instrs.split_off(instr_idx + 1); + func.bbs[caller_bb_id].instrs.pop(); + + func.bbs.insert_node(BasicBlock { + id: continuation_bb_id, + instrs: instrs_after, + }); + *func.bbs[continuation_bb_id].terminator_mut() = original_terminator; + + let callee_id = FuncId::from(constant.func_id); + let callee = &ctx.module.funcs[callee_id]; + + let mut local_map = FxHashMap::default(); + let mut bb_map = FxHashMap::default(); + + for old_bb_id_usize in callee.bbs.keys() { + let old_bb_id = old_bb_id_usize; + let new_bb_id_usize = func.bbs.allocate_key(); + let new_bb_id = new_bb_id_usize; + bb_map.insert(old_bb_id, new_bb_id); + worklist.push(new_bb_id); + } + + for (i, &arg_local) in callee.args.iter().enumerate() { + let replacement = if i == 0 { + call.closure + } else if i - 1 < call.args.len() { + call.args[i - 1] + } else { + call.args[call.args.len() - 1] + }; + local_map.insert(arg_local, replacement); + } + + for (local_id_usize, local) in callee.locals.iter() { + let local_id = local_id_usize; + if let std::collections::hash_map::Entry::Vacant(e) = local_map.entry(local_id) { + let new_id_usize = func.locals.push_with(|id| Local { + id, + typ: local.typ, + ..*local + }); + let new_id = new_id_usize; + e.insert(new_id); + } + } + + let mut phi_incomings = Vec::new(); + + for (old_bb_id_usize, old_bb) in callee.bbs.iter() { + let old_bb_id = old_bb_id_usize; + let new_bb_id = bb_map[&old_bb_id]; + let mut new_instrs = Vec::new(); + + for instr in &old_bb.instrs { + let mut new_instr = instr.clone(); + if let Some(local) = new_instr.local + && let Some(&mapped) = local_map.get(&local) + { + new_instr.local = Some(mapped); + } + rewrite_usages(&mut new_instr.kind, &local_map); + for bb_ref in new_instr.kind.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + new_instrs.push(new_instr); + } + + new_instrs.pop(); + + let mut new_terminator = old_bb.terminator().clone(); + rewrite_terminator_usages(&mut new_terminator, &local_map); + for bb_ref in new_terminator.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + + if let TerminatorInstr::Exit(ExitInstr::Return(val)) = &new_terminator { + // Note: ExitInstr::Return contains LocalId (not Option) + phi_incomings.push(PhiIncomingValue { + local: *val, + bb: new_bb_id, + }); + new_terminator = TerminatorInstr::Jump(continuation_bb_id); + } else if let TerminatorInstr::Exit(ExitInstr::TailCallClosure(call)) = &new_terminator { + if let Some(dst) = result_local { + let dst_typ = func.locals[dst].typ; + let temp_local_usize = func.locals.push_with(|id| Local { id, typ: dst_typ }); + let temp_local = temp_local_usize; + + new_instrs.push(Instr { + local: Some(temp_local), + kind: InstrKind::CallClosure(call.clone()), + }); + phi_incomings.push(PhiIncomingValue { + local: temp_local, + bb: new_bb_id, + }); + } else { + new_instrs.push(Instr { + local: None, + kind: InstrKind::CallClosure(call.clone()), + }); + } + new_terminator = TerminatorInstr::Jump(continuation_bb_id); + } + + new_instrs.push(Instr { + local: None, + kind: InstrKind::Terminator(new_terminator), + }); + + func.bbs.insert_node(BasicBlock { + id: new_bb_id, + instrs: new_instrs, + }); + } + + if let Some(dst) = result_local + && !phi_incomings.is_empty() + { + func.bbs[continuation_bb_id].instrs.insert( + 0, + Instr { + local: Some(dst), + kind: InstrKind::Phi { + incomings: phi_incomings, + non_exhaustive: false, + }, + }, + ); + } + + let new_entry_id = bb_map[&callee.bb_entry]; + + log::debug!( + "Inlining non-tail: Caller {:?} -> Cont {:?}, Entry {:?}", + caller_bb_id, + continuation_bb_id, + new_entry_id + ); + + func.bbs[caller_bb_id].instrs.push(Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(new_entry_id)), + }); + + continuation_bb_id +} + +fn inline_tail( + func: &mut Func, + ctx: &mut InlineContext, + worklist: &mut Vec, + caller_bb_id: BasicBlockId, + constant: ConstantClosure, + call: &InstrCallClosure, +) { + if let Some(info) = ctx.tail_instances.get(&constant) { + log::debug!("Reusing tail instance at entry BB {:?}", info.entry_bb); + for (i, &phi_local) in info.arg_phis.iter().enumerate() { + let val = if i == 0 { + call.closure + } else { + call.args[i - 1] + }; + let entry_bb = &mut func.bbs[info.entry_bb]; + for instr in &mut entry_bb.instrs { + if let InstrKind::Phi { incomings, .. } = &mut instr.kind + && instr.local == Some(phi_local) + { + incomings.push(PhiIncomingValue { + local: val, + bb: caller_bb_id, + }); + break; + } + } + } + *func.bbs[caller_bb_id].terminator_mut() = TerminatorInstr::Jump(info.entry_bb); + return; + } + + let callee_id = FuncId::from(constant.func_id); + let callee = &ctx.module.funcs[callee_id]; + + let mut local_map = FxHashMap::default(); + let mut bb_map = FxHashMap::default(); + + for old_bb_id_usize in callee.bbs.keys() { + let old_bb_id = old_bb_id_usize; + let new_bb_id_usize = func.bbs.allocate_key(); + let new_bb_id = new_bb_id_usize; + bb_map.insert(old_bb_id, new_bb_id); + worklist.push(new_bb_id); + } + + let new_entry_id = bb_map[&callee.bb_entry]; + let mut arg_phis = Vec::new(); + let mut entry_phi_instrs = Vec::new(); + + for (i, &arg_local) in callee.args.iter().enumerate() { + let new_arg_local_usize = func.locals.push_with(|id| Local { + id, + typ: callee.locals[arg_local].typ, + ..callee.locals[arg_local] + }); + let new_arg_local = new_arg_local_usize; + local_map.insert(arg_local, new_arg_local); + arg_phis.push(new_arg_local); + + let val = if i == 0 { + call.closure + } else { + call.args[i - 1] + }; + entry_phi_instrs.push(Instr { + local: Some(new_arg_local), + kind: InstrKind::Phi { + incomings: vec![PhiIncomingValue { + local: val, + bb: caller_bb_id, + }], + non_exhaustive: false, + }, + }); + } + + for (local_id_usize, local) in callee.locals.iter() { + let local_id = local_id_usize; + if let std::collections::hash_map::Entry::Vacant(e) = local_map.entry(local_id) { + let new_id_usize = func.locals.push_with(|id| Local { + id, + typ: local.typ, + ..*local + }); + let new_id = new_id_usize; + e.insert(new_id); + } + } + + for (old_bb_id_usize, old_bb) in callee.bbs.iter() { + let old_bb_id = old_bb_id_usize; + let new_bb_id = bb_map[&old_bb_id]; + let mut new_instrs = Vec::new(); + + if old_bb_id == callee.bb_entry { + new_instrs.extend(entry_phi_instrs.clone()); + } + + for instr in &old_bb.instrs { + let mut new_instr = instr.clone(); + if let Some(local) = new_instr.local + && let Some(&mapped) = local_map.get(&local) + { + new_instr.local = Some(mapped); + } + rewrite_usages(&mut new_instr.kind, &local_map); + for bb_ref in new_instr.kind.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + new_instrs.push(new_instr); + } + + let mut new_terminator = old_bb.terminator().clone(); + rewrite_terminator_usages(&mut new_terminator, &local_map); + for bb_ref in new_terminator.bb_ids_mut() { + if let Some(&mapped) = bb_map.get(bb_ref) { + *bb_ref = mapped; + } + } + + func.bbs.insert_node(BasicBlock { + id: new_bb_id, + instrs: new_instrs, + }); + *func.bbs[new_bb_id].terminator_mut() = new_terminator; + } + + *func.bbs[caller_bb_id].terminator_mut() = TerminatorInstr::Jump(new_entry_id); + + log::debug!( + "Created new tail instance for func {:?} at entry BB {:?}", + constant.func_id, + new_entry_id + ); + ctx.tail_instances.insert( + constant, + TailCallInfo { + entry_bb: new_entry_id, + arg_phis, + }, + ); +} + +fn rewrite_usages(kind: &mut InstrKind, map: &FxHashMap) { + for (local, _) in kind.local_usages_mut() { + if let Some(&mapped) = map.get(local) { + *local = mapped; + } + } +} + +fn rewrite_terminator_usages(term: &mut TerminatorInstr, map: &FxHashMap) { + for local in term.local_ids_mut() { + if let Some(&mapped) = map.get(local) { + *local = mapped; + } + } +} diff --git a/webschembly-compiler/src/ir_processor/mod.rs b/webschembly-compiler/src/ir_processor/mod.rs index c9a2d875..7d9e70c3 100644 --- a/webschembly-compiler/src/ir_processor/mod.rs +++ b/webschembly-compiler/src/ir_processor/mod.rs @@ -1,7 +1,10 @@ pub mod cfg_analyzer; pub mod dataflow; pub mod desugar; +pub mod inline; pub mod optimizer; +pub mod propagate_types; pub mod register_allocation; +pub mod remove_constant; pub mod ssa; pub mod ssa_optimizer; diff --git a/webschembly-compiler/src/ir_processor/propagate_types.rs b/webschembly-compiler/src/ir_processor/propagate_types.rs new file mode 100644 index 00000000..61d597dc --- /dev/null +++ b/webschembly-compiler/src/ir_processor/propagate_types.rs @@ -0,0 +1,547 @@ +use crate::ir_processor::cfg_analyzer::calculate_rpo; +use vec_map::VecMap; +use webschembly_compiler_ir::*; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LatticeValue { + Top, + Constant(ConstantClosure), + Bottom, +} + +impl LatticeValue { + fn meet(self, other: LatticeValue) -> LatticeValue { + match (self, other) { + (LatticeValue::Top, x) | (x, LatticeValue::Top) => x, + (LatticeValue::Constant(a), LatticeValue::Constant(b)) => { + if a == b { + LatticeValue::Constant(a) + } else { + LatticeValue::Bottom + } + } + (LatticeValue::Bottom, _) | (_, LatticeValue::Bottom) => LatticeValue::Bottom, + } + } +} + +pub fn propagate_types(func: &mut Func) { + let mut lattice = VecMap::new(); + // Initialize args to Bottom (we don't know anything about them) + // Other locals are implicitly Top (not in map or handled as default) + for &arg in &func.args { + lattice.insert(arg, LatticeValue::Bottom); + } + + // Initialize all other locals to Top + for local_id in func.locals.keys() { + if !lattice.contains_key(local_id) { + lattice.insert(local_id, LatticeValue::Top); + } + } + + let rpo = calculate_rpo(&func.bbs, func.bb_entry); + + // Convert RPO map to a list of BBs sorted by RPO index + let mut sorted_bbs: Vec = rpo.keys().cloned().collect(); + sorted_bbs.sort_by_key(|bb| rpo.get(bb).unwrap()); + + let mut changed = true; + while changed { + changed = false; + + for &bb_id in &sorted_bbs { + let bb = &func.bbs[bb_id]; + for instr in &bb.instrs { + if let Some(dest) = instr.local { + let old_val = *lattice.get(dest).unwrap_or(&LatticeValue::Top); + let new_val = match &instr.kind { + /* + InstrKind::Closure { + func_id, env_index, .. + } => LatticeValue::Constant(ConstantClosure { + func_id: *func_id, + env_index: ClosureEnvIndex(*env_index), + }), + */ + InstrKind::Move(src) => *lattice.get(*src).unwrap_or(&LatticeValue::Top), + InstrKind::Phi { + incomings, + non_exhaustive, + } => { + if *non_exhaustive { + LatticeValue::Bottom + } else { + let mut val = LatticeValue::Top; + for incoming in incomings { + let incoming_val = + *lattice.get(incoming.local).unwrap_or(&LatticeValue::Top); + val = val.meet(incoming_val); + } + val + } + } + // If we have a direct reference to a known constant logic in other instructions + // we might want to propagate that, but for now we only care about propagating definitions. + // Any other instruction that defines a value produces Bottom for that value (unknown type/value) + _ => LatticeValue::Bottom, + }; + + if old_val != new_val { + lattice.insert(dest, new_val); + changed = true; + } + } + } + } + } + + // Update types based on lattice results + for (local_id, val) in lattice { + match val { + LatticeValue::Constant(constant) => { + let local = &mut func.locals[local_id]; + match local.typ { + LocalType::Type(Type::Val(ValType::Closure(_))) => { + local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); + } + // If it was Type::Obj (upcasted), we can refine it to Closure + LocalType::Type(Type::Obj) => { + local.typ = LocalType::Type(Type::Val(ValType::Closure(Some(constant)))); + } + _ => {} + } + } + // If Bottom, we might need to revert to generic Closure(None) if it was previously set to something specific + // but here we are primarily taking generic Code (Closure(None)) and refining it. + // If the code was *already* specialized/refined, we should be careful not to overwrite it with less specific info + // unless we are sure. But this pass is usually run to refine. + LatticeValue::Bottom => { + // Ensure if we failed to prove constant, it remains as generic closure if it was a closure + /* + Note: We don't generally downgrade types here explicitly because the lattice implies + we only UPGRADE to Constant or stay Bottom. + However, if this pass is run multiple times or if we have partial information, + we might want to be safe. But currently the IR starts with Closure(None) usually. + */ + } + LatticeValue::Top => { + // Dead code or uninitialized. Leave as is or set to Bottom equivalent? + // Usually safe to ignore. + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use webschembly_compiler_ir::{ + BasicBlock, ExitInstr, Func, Instr, InstrKind, Local, TerminatorInstr, Type, ValType, + }; + + fn create_dummy_func() -> Func { + Func { + id: FuncId::from(0), + args: vec![], + locals: VecMap::new(), + bbs: VecMap::new(), + bb_entry: BasicBlockId::from(0), + ret_type: LocalType::Type(Type::Val(ValType::Nil)), + closure_meta: None, + } + } + + #[test] + fn test_propagate_simple_move() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // Target ID + entrypoint_table: l1, // Dummy + original_entrypoint_table: l1, // Dummy + }, + }, + Instr { + local: Some(l1), + kind: InstrKind::Move(l0), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l1))), + }, + ], + }); + func.bb_entry = bb0; + + propagate_types(&mut func); + + let t1 = func.locals[l1].typ; + if let LocalType::Type(Type::Val(ValType::Closure(Some(c)))) = t1 { + assert_eq!(c.func_id, JitFuncId::from(10)); + assert_eq!(c.env_index, ClosureEnvIndex(0)); + } else { + panic!("Expected ConstantClosure, got {:?}", t1); + } + } + + #[test] + fn test_propagate_phi_same() { + let mut func = create_dummy_func(); + // l0 and l1 are same closure constant + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); // Phi result + + /* + bb0: + l0 = closure(10, 0) + jump bb1 + bb1: + l1 = closure(10, 0) + jump bb2 + bb2: + l2 = phi(l0: bb0, l1: bb1) + */ + + let bb0 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + let bb1 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + let bb2 = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![], // Fill later + }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, + ]; + + // entry -> bb0 (cond jump is hard so just make one path reachable for simple test structure or use Jump from entry) + // Let's make entry jump to bb0 and bb1? No, just linear flow is fine? + // Wait, for Phi to work we need valid preds. + // Let's make entry conditional jump to bb0 and bb1. + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + if let LocalType::Type(Type::Val(ValType::Closure(Some(c)))) = t2 { + assert_eq!(c.func_id, JitFuncId::from(10)); + } else { + panic!("Expected ConstantClosure for l2, got {:?}", t2); + } + } + + #[test] + fn test_propagate_phi_different() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(11), // ID 11 (Different) + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + non_exhaustive: false, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, + ]; + + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + // Should remain generic Closure(None) because input constants are different + if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { + // OK + } else { + panic!( + "Expected generic Closure(None) (Bottom) for l2, got {:?}", + t2 + ); + } + } + + #[test] + fn test_propagate_phi_non_exhaustive() { + let mut func = create_dummy_func(); + let l0 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l1 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + let l2 = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Closure(None))), + }); + + let bb0 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb1 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + let bb2 = func.bbs.push_with(|id| BasicBlock { id, instrs: vec![] }); + + func.bbs[bb0].instrs = vec![ + Instr { + local: Some(l0), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb1].instrs = vec![ + Instr { + local: Some(l1), + kind: InstrKind::Closure { + envs: vec![], + env_types: vec![], + env_index: 0, + module_id: JitModuleId::from(0), + func_id: JitFuncId::from(10), // ID 10 (Same as l0) + entrypoint_table: l0, + original_entrypoint_table: l0, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Jump(bb2)), + }, + ]; + + func.bbs[bb2].instrs = vec![ + Instr { + local: Some(l2), + kind: InstrKind::Phi { + incomings: vec![ + PhiIncomingValue { bb: bb0, local: l0 }, + PhiIncomingValue { bb: bb1, local: l1 }, + ], + // Even though l0 and l1 are same constant, non_exhaustive should make result Bottom + non_exhaustive: true, + }, + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::Exit(ExitInstr::Return(l2))), + }, + ]; + + let l_cond = func.locals.push_with(|id| Local { + id, + typ: LocalType::Type(Type::Val(ValType::Bool)), + }); + func.bb_entry = func.bbs.push_with(|id| BasicBlock { + id, + instrs: vec![ + Instr { + local: Some(l_cond), + kind: InstrKind::Bool(true), + }, + Instr { + local: None, + kind: InstrKind::Terminator(TerminatorInstr::If(l_cond, bb0, bb1)), + }, + ], + }); + + propagate_types(&mut func); + + let t2 = func.locals[l2].typ; + // Should remain generic Closure(None) because non_exhaustive is true + if let LocalType::Type(Type::Val(ValType::Closure(None))) = t2 { + // OK + } else { + panic!( + "Expected generic Closure(None) (Bottom) for l2 due to non_exhaustive, got {:?}", + t2 + ); + } + } +} diff --git a/webschembly-compiler/src/ir_processor/remove_constant.rs b/webschembly-compiler/src/ir_processor/remove_constant.rs new file mode 100644 index 00000000..584a5483 --- /dev/null +++ b/webschembly-compiler/src/ir_processor/remove_constant.rs @@ -0,0 +1,9 @@ +use webschembly_compiler_ir::Func; + +pub fn remove_constant(func: &mut Func) { + for local in func.locals.values_mut() { + local.typ = local.typ.remove_constant(); + } + + func.ret_type = func.ret_type.remove_constant(); +} diff --git a/webschembly-compiler/src/ir_processor/ssa.rs b/webschembly-compiler/src/ir_processor/ssa.rs index cf8767e4..7c449137 100644 --- a/webschembly-compiler/src/ir_processor/ssa.rs +++ b/webschembly-compiler/src/ir_processor/ssa.rs @@ -210,7 +210,10 @@ fn assert_ssa(func: &Func) { phi_area = false; } } else if let InstrKind::Phi { .. } = expr.kind { - panic!("phi instruction must be at the beginning of a basic block"); + panic!( + "phi instruction must be at the beginning of a basic block. BB: {:?}, Instrs: {:#?}", + bb.id, bb.instrs + ); } } } diff --git a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs index 0b6f2996..a9bfd4d8 100644 --- a/webschembly-compiler/src/ir_processor/ssa_optimizer.rs +++ b/webschembly-compiler/src/ir_processor/ssa_optimizer.rs @@ -3,6 +3,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; use crate::ir_processor::{ cfg_analyzer::{DomTreeNode, build_dom_tree, calc_doms, calc_predecessors, calculate_rpo}, optimizer::remove_unreachable_bb, + propagate_types::propagate_types, ssa::{DefUseChain, debug_assert_ssa}, }; use vec_map::VecMap; @@ -370,7 +371,8 @@ pub fn constant_folding( if let Some(&InstrKind::ToObj(typ2, _)) = def_use.get_def_non_move_expr(&func.bbs, src) => { - func.bbs[*bb_id].instrs[expr_idx].kind = InstrKind::Bool(typ1 == typ2); + func.bbs[*bb_id].instrs[expr_idx].kind = + InstrKind::Bool(typ1.remove_constant() == typ2.remove_constant()); } InstrKind::ClosureEnv(_, closure, index) if let Some(InstrKind::Closure { envs, .. }) = @@ -447,7 +449,7 @@ impl Default for SsaOptimizerConfig { SsaOptimizerConfig { enable_cse: true, enable_dce: true, - enable_inlining: false, // true, + enable_inlining: true, iterations: 5, } } @@ -462,6 +464,8 @@ pub fn ssa_optimize(func: &mut Func, config: SsaOptimizerConfig) { let doms = calc_doms(&func.bbs, &rpo, func.bb_entry, &predecessors); let dom_tree = build_dom_tree(&func.bbs, &rpo, func.bb_entry, &doms); + propagate_types(func); + for _ in 0..config.iterations { debug_assert_ssa(func); copy_propagation(func, &rpo); diff --git a/webschembly-compiler/src/jit/bb_index_manager.rs b/webschembly-compiler/src/jit/bb_index_manager.rs index 4b75d2e6..88ebd683 100644 --- a/webschembly-compiler/src/jit/bb_index_manager.rs +++ b/webschembly-compiler/src/jit/bb_index_manager.rs @@ -10,9 +10,6 @@ use super::index_flag::IndexFlag; pub const BB_LAYOUT_MAX_SIZE: usize = 32; pub const BB_LAYOUT_DEFAULT_INDEX: BBIndex = BBIndex(0); -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct BBIndex(pub usize); - #[derive(Debug)] pub struct BBIndexManager { type_params_to_index: FxBiHashMap, BBIndex>, diff --git a/webschembly-compiler/src/jit/closure_global_layout.rs b/webschembly-compiler/src/jit/closure_global_layout.rs index 759d9727..d8cfea6c 100644 --- a/webschembly-compiler/src/jit/closure_global_layout.rs +++ b/webschembly-compiler/src/jit/closure_global_layout.rs @@ -6,10 +6,7 @@ use crate::fxbihashmap::FxBiHashMap; use super::index_flag::IndexFlag; pub const CLOSURE_LAYOUT_MAX_SIZE: usize = 32; -pub const CLOSURE_LAYOUT_DEFAULT_INDEX: ClosureIndex = ClosureIndex(0); - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct ClosureIndex(pub usize); +pub const CLOSURE_LAYOUT_DEFAULT_INDEX: ClosureArgIndex = ClosureArgIndex(0); #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ClosureArgs { @@ -19,8 +16,8 @@ pub enum ClosureArgs { #[derive(Debug)] pub struct ClosureGlobalLayout { - args_to_index: FxBiHashMap, - instantiated_idx: FxHashSet, + args_to_index: FxBiHashMap, + instantiated_idx: FxHashSet, } impl Default for ClosureGlobalLayout { @@ -39,12 +36,12 @@ impl ClosureGlobalLayout { } } - pub fn idx(&mut self, args: &ClosureArgs) -> Option<(ClosureIndex, IndexFlag)> { + pub fn idx(&mut self, args: &ClosureArgs) -> Option<(ClosureArgIndex, IndexFlag)> { // TODO: argsの長さに上限を設定 let index = if let Some(&index) = self.args_to_index.get_by_left(args) { index } else if self.args_to_index.len() < CLOSURE_LAYOUT_MAX_SIZE { - let index = ClosureIndex(self.args_to_index.len()); + let index = ClosureArgIndex(self.args_to_index.len()); self.args_to_index.insert(args.clone(), index); index } else { @@ -58,7 +55,7 @@ impl ClosureGlobalLayout { Some((index, flag)) } - pub fn arg_types(&self, index: ClosureIndex) -> &ClosureArgs { + pub fn arg_types(&self, index: ClosureArgIndex) -> &ClosureArgs { self.args_to_index.get_by_right(&index).unwrap() } } diff --git a/webschembly-compiler/src/jit/env_index_manager.rs b/webschembly-compiler/src/jit/env_index_manager.rs index 9d39e31d..8e924aca 100644 --- a/webschembly-compiler/src/jit/env_index_manager.rs +++ b/webschembly-compiler/src/jit/env_index_manager.rs @@ -8,15 +8,12 @@ use crate::ir_generator::GlobalManager; use super::index_flag::IndexFlag; pub const ENV_LAYOUT_MAX_SIZE: usize = 32; -pub const ENV_LAYOUT_DEFAULT_INDEX: EnvIndex = EnvIndex(0); - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct EnvIndex(pub usize); +pub const ENV_LAYOUT_DEFAULT_INDEX: ClosureEnvIndex = ClosureEnvIndex(0); #[derive(Debug)] pub struct EnvIndexManager { - env_types_to_index: FxBiHashMap, EnvIndex>, - index_to_table_global: FxHashMap, + env_types_to_index: FxBiHashMap, ClosureEnvIndex>, + index_to_table_global: FxHashMap, } impl Default for EnvIndexManager { @@ -40,7 +37,7 @@ impl EnvIndexManager { &mut self, env_types: &VecMap, global_manager: &mut GlobalManager, - ) -> Option<(Option, EnvIndex, IndexFlag)> { + ) -> Option<(Option, ClosureEnvIndex, IndexFlag)> { if let Some(&index) = self .env_types_to_index .get_by_left(VecMapEq::from_ref(env_types)) @@ -49,7 +46,7 @@ impl EnvIndexManager { debug_assert!(index == ENV_LAYOUT_DEFAULT_INDEX || global.is_some()); Some((global.copied(), index, IndexFlag::ExistingInstance)) } else if self.env_types_to_index.len() < ENV_LAYOUT_MAX_SIZE { - let index = EnvIndex(self.env_types_to_index.len()); + let index = ClosureEnvIndex(self.env_types_to_index.len()); self.env_types_to_index .insert(env_types.clone().into(), index); let global = global_manager.gen_global(LocalType::EntrypointTable); @@ -60,7 +57,7 @@ impl EnvIndexManager { } } - pub fn env_types(&self, index: EnvIndex) -> (&VecMap, Option) { + pub fn env_types(&self, index: ClosureEnvIndex) -> (&VecMap, Option) { debug_assert!( index == ENV_LAYOUT_DEFAULT_INDEX || self.index_to_table_global.contains_key(&index) ); diff --git a/webschembly-compiler/src/jit/jit_ctx.rs b/webschembly-compiler/src/jit/jit_ctx.rs index 37f12cba..f556bfd9 100644 --- a/webschembly-compiler/src/jit/jit_ctx.rs +++ b/webschembly-compiler/src/jit/jit_ctx.rs @@ -1,6 +1,6 @@ use rustc_hash::FxHashMap; -use super::closure_global_layout::{ClosureGlobalLayout, ClosureIndex}; +use super::closure_global_layout::ClosureGlobalLayout; use super::jit_config::JitConfig; use webschembly_compiler_ir::*; @@ -12,7 +12,7 @@ pub struct JitCtx { is_instantiated: bool, // 0..CLOSURE_LAYOUT_MAX_SIZEまでのindexに対応する関数のスタブが入ったMutFuncRef // func_indexがインスタンス化されるときにMutFuncRefにFuncRefがセットされる - stub_globals: FxHashMap, + stub_globals: FxHashMap, // instantiate_funcの結果を保存するグローバル instantiate_func_global: Option, } @@ -32,7 +32,7 @@ impl JitCtx { self.config } - pub fn stub_global(&self, index: ClosureIndex) -> Global { + pub fn stub_global(&self, index: ClosureArgIndex) -> Global { debug_assert!(self.is_instantiated); self.stub_globals[&index] } @@ -48,7 +48,7 @@ impl JitCtx { pub fn init_instantiated( &mut self, - stub_globals: FxHashMap, + stub_globals: FxHashMap, instantiate_func_global: Global, ) { debug_assert!(!self.is_instantiated); diff --git a/webschembly-compiler/src/jit/jit_func.rs b/webschembly-compiler/src/jit/jit_func.rs index a985c3ec..70e769a6 100644 --- a/webschembly-compiler/src/jit/jit_func.rs +++ b/webschembly-compiler/src/jit/jit_func.rs @@ -1,11 +1,10 @@ use rustc_hash::{FxHashMap, FxHashSet}; -use super::bb_index_manager::{BB_LAYOUT_DEFAULT_INDEX, BBIndex, BBIndexManager}; +use super::bb_index_manager::{BB_LAYOUT_DEFAULT_INDEX, BBIndexManager}; use super::closure_global_layout::{ CLOSURE_LAYOUT_DEFAULT_INDEX, CLOSURE_LAYOUT_MAX_SIZE, ClosureArgs, ClosureGlobalLayout, - ClosureIndex, }; -use super::env_index_manager::{ENV_LAYOUT_DEFAULT_INDEX, EnvIndex, EnvIndexManager}; +use super::env_index_manager::{ENV_LAYOUT_DEFAULT_INDEX, EnvIndexManager}; use super::index_flag::IndexFlag; use super::jit_ctx::JitCtx; use crate::fxbihashmap::FxBiHashMap; @@ -20,7 +19,7 @@ use webschembly_compiler_ir::*; #[derive(Debug)] pub struct JitFunc { - pub jit_specialized_env_funcs: FxHashMap, + pub jit_specialized_env_funcs: FxHashMap, } impl JitFunc { @@ -49,7 +48,7 @@ impl JitFunc { #[derive(Debug)] pub struct JitSpecializedEnvFunc { - pub jit_specialized_arg_funcs: FxHashMap, + pub jit_specialized_arg_funcs: FxHashMap, pub func: Func, } @@ -58,7 +57,7 @@ impl JitSpecializedEnvFunc { module_id: JitModuleId, global_manager: &mut GlobalManager, func: &Func, - env_index: EnvIndex, + env_index: ClosureEnvIndex, env_index_manager: &EnvIndexManager, jit_ctx: &mut JitCtx, ) -> Self { @@ -84,8 +83,8 @@ impl JitSpecializedEnvFunc { #[derive(Debug)] pub struct JitSpecializedArgFunc { module_id: JitModuleId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, func: Func, jit_bbs: VecMap, } @@ -95,8 +94,8 @@ impl JitSpecializedArgFunc { module_id: JitModuleId, global_manager: &mut GlobalManager, func: &Func, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, jit_ctx: &mut JitCtx, ) -> Self { let mut func = func.clone(); @@ -613,6 +612,10 @@ impl JitSpecializedArgFunc { if let Some(&InstrKind::ToObj(typ, val_local)) = def_use_chain.get_def_non_move_expr(&body_func.bbs, obj_local) { + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = body_func.locals[val_local].typ { + typ = t; + } Some(TypedObj { typ, val_local }) } else { typed_objs.get(&obj_local).copied() @@ -722,7 +725,7 @@ impl JitSpecializedArgFunc { instrs.push(Instr { local: Some(stub), kind: InstrKind::GlobalGet( - jit_ctx.stub_global(ClosureIndex(index)).id, + jit_ctx.stub_global(ClosureArgIndex(index)).id, ), }); locals.push(stub); @@ -763,6 +766,7 @@ impl JitSpecializedArgFunc { call_closure, &def_use_chain, &body_func.bbs, + &body_func.locals, jit_ctx.closure_global_layout(), &mut required_closure_idx, ) => @@ -780,6 +784,7 @@ impl JitSpecializedArgFunc { call_closure, &def_use_chain, &body_func.bbs, + &body_func.locals, jit_ctx.closure_global_layout(), &mut required_closure_idx, ) => @@ -823,6 +828,12 @@ impl JitSpecializedArgFunc { def_use_chain .get_def_non_move_expr(&body_func.bbs, value_local) { + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = + body_func.locals[val_local].typ + { + typ = t; + } Some(TypedObj { typ, val_local }) } else { None @@ -869,6 +880,15 @@ impl JitSpecializedArgFunc { local: Some(entrypoint_table_local), kind: InstrKind::GlobalGet(entrypoint_table_global.id), }); + if let Some(local) = body_func.bbs[bb_id].instrs[instr_idx].local { + body_func.locals[local].typ = + (ValType::Closure(Some(ConstantClosure { + module_id: *module_id, + func_id: *func_id, + env_index: ClosureEnvIndex(env_index.0), + }))) + .into(); + } body_func.bbs[bb_id].instrs[instr_idx].kind = InstrKind::Closure { envs: new_envs, env_types: new_env_types.clone(), @@ -916,7 +936,7 @@ impl JitSpecializedArgFunc { let mut args = Vec::new(); let closure_local = locals.push_with(|id| Local { id, - typ: ValType::Closure.into(), + typ: ValType::Closure(None).into(), }); let mut arg_locals = Vec::new(); args.push(closure_local); @@ -1074,7 +1094,9 @@ impl JitSpecializedArgFunc { }); instrs.push(Instr { local: Some(stub), - kind: InstrKind::GlobalGet(jit_ctx.stub_global(ClosureIndex(index)).id), + kind: InstrKind::GlobalGet( + jit_ctx.stub_global(ClosureArgIndex(index)).id, + ), }); entrypoint_table_locals.push(stub); } @@ -1198,8 +1220,9 @@ fn specialize_call_closure( call_closure: &InstrCallClosure, def_use_chain: &DefUseChain, bbs: &VecMap, + locals: &VecMap, closure_global_layout: &mut ClosureGlobalLayout, - required_closure_idx: &mut Vec, + required_closure_idx: &mut Vec, ) -> Option { if call_closure.func_index != CLOSURE_LAYOUT_DEFAULT_INDEX.0 { return None; @@ -1219,6 +1242,10 @@ fn specialize_call_closure( def_use_chain.get_def_non_move_expr(bbs, obj_arg) { fixed_args.push(val_local); + let mut typ = typ; + if let LocalType::Type(Type::Val(t)) = locals[val_local].typ { + typ = t; + } fixed_arg_types.push(Type::from(typ)); } else { fixed_args.push(obj_arg); @@ -1433,7 +1460,7 @@ fn calculate_args_to_pass( fn closure_func_assign_env_types( func: &mut Func, - env_index: EnvIndex, + env_index: ClosureEnvIndex, env_index_manager: &EnvIndexManager, ) { if env_index == ENV_LAYOUT_DEFAULT_INDEX { @@ -1463,7 +1490,7 @@ fn closure_func_assign_env_types( } let new_closure_arg = func.locals.push_with(|id| Local { id, - typ: LocalType::Type(Type::Val(ValType::Closure)), + typ: LocalType::Type(Type::Val(ValType::Closure(None))), }); let prev_closure_arg = func.args[0]; func.args[0] = new_closure_arg; @@ -1527,7 +1554,7 @@ fn closure_func_assign_env_types( fn closure_func_assign_types( func: &mut Func, - func_index: ClosureIndex, + func_index: ClosureArgIndex, closure_global_layout: &ClosureGlobalLayout, ) { let ClosureArgs::Specified(args_type) = closure_global_layout.arg_types(func_index) else { @@ -1541,7 +1568,7 @@ fn closure_func_assign_types( .map(|&arg| func.locals[arg].typ) .collect::>(), vec![ - LocalType::Type(Type::Val(ValType::Closure)), + LocalType::Type(Type::Val(ValType::Closure(None))), LocalType::VariadicArgs ] ); diff --git a/webschembly-compiler/src/jit/jit_module.rs b/webschembly-compiler/src/jit/jit_module.rs index ded0d1f8..c32c7dee 100644 --- a/webschembly-compiler/src/jit/jit_module.rs +++ b/webschembly-compiler/src/jit/jit_module.rs @@ -1,14 +1,14 @@ use rustc_hash::FxHashMap; -use super::bb_index_manager::BBIndex; -use super::closure_global_layout::{CLOSURE_LAYOUT_MAX_SIZE, ClosureIndex}; -use super::env_index_manager::{EnvIndex, EnvIndexManager}; +use super::closure_global_layout::CLOSURE_LAYOUT_MAX_SIZE; +use super::env_index_manager::EnvIndexManager; use super::jit_ctx::JitCtx; use super::jit_func::{JitFunc, JitSpecializedArgFunc}; use crate::ir_generator::GlobalManager; use crate::jit::jit_func::JitSpecializedEnvFunc; use vec_map::{HasId, VecMap}; use webschembly_compiler_ir::*; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; #[derive(Debug)] pub struct JitModule { module_id: JitModuleId, @@ -170,7 +170,7 @@ impl JitModule { if !jit_ctx.is_instantiated() { let mut stub_globals = FxHashMap::default(); for func_index in 0..CLOSURE_LAYOUT_MAX_SIZE { - let func_index = ClosureIndex(func_index); + let func_index = ClosureArgIndex(func_index); let stub_global = global_manager.gen_global(LocalType::MutFuncRef); stub_globals.insert(func_index, stub_global); let stub_local = entry_func.locals.push_with(|id| Local { @@ -217,8 +217,8 @@ impl JitModule { &mut self, global_manager: &mut GlobalManager, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, jit_ctx: &mut JitCtx, ) -> Module { let jit_func_entry = self.jit_funcs.get_mut(&func_id).unwrap(); @@ -264,8 +264,8 @@ impl JitModule { pub fn instantiate_bb( &mut self, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, index: BBIndex, global_manager: &mut GlobalManager, @@ -297,8 +297,8 @@ impl JitModule { global_manager: &mut GlobalManager, jit_ctx: &mut JitCtx, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, kind: BranchKind, source_bb_id: BasicBlockId, diff --git a/webschembly-compiler/src/jit/mod.rs b/webschembly-compiler/src/jit/mod.rs index a0879b2a..e445b693 100644 --- a/webschembly-compiler/src/jit/mod.rs +++ b/webschembly-compiler/src/jit/mod.rs @@ -1,9 +1,6 @@ use crate::ir_generator::GlobalManager; use vec_map::VecMap; -use bb_index_manager::BBIndex; -use closure_global_layout::ClosureIndex; -use env_index_manager::EnvIndex; mod jit_config; mod jit_module; pub use jit_config::JitConfig; @@ -11,7 +8,9 @@ use jit_module::JitModule; mod jit_ctx; use jit_ctx::JitCtx; use webschembly_compiler_ir::*; +use webschembly_compiler_ir::{BBIndex, ClosureArgIndex, ClosureEnvIndex}; pub mod bb_index_manager; + pub mod closure_global_layout; pub mod env_index_manager; pub mod index_flag; @@ -51,8 +50,8 @@ impl Jit { global_manager: &mut GlobalManager, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, ) -> Module { self.jit_module[module_id].instantiate_func( global_manager, @@ -67,8 +66,8 @@ impl Jit { &mut self, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, index: BBIndex, global_manager: &mut GlobalManager, @@ -89,8 +88,8 @@ impl Jit { global_manager: &mut GlobalManager, module_id: JitModuleId, func_id: FuncId, - env_index: EnvIndex, - func_index: ClosureIndex, + env_index: ClosureEnvIndex, + func_index: ClosureArgIndex, bb_id: BasicBlockId, kind: BranchKind, source_bb_id: BasicBlockId, diff --git a/webschembly-compiler/src/wasm_generator/module_generator.rs b/webschembly-compiler/src/wasm_generator/module_generator.rs index a91cb29b..f3be567a 100644 --- a/webschembly-compiler/src/wasm_generator/module_generator.rs +++ b/webschembly-compiler/src/wasm_generator/module_generator.rs @@ -931,7 +931,7 @@ impl<'a> ModuleGenerator<'a> { nullable: true, heap_type: HeapType::Concrete(self.cons_type), }), - ir::ValType::Closure => ValType::Ref(RefType { + ir::ValType::Closure(_) => ValType::Ref(RefType { nullable: true, heap_type: HeapType::Concrete(self.closure_type), }), @@ -979,7 +979,7 @@ impl<'a> ModuleGenerator<'a> { ir::LocalType::Type(ir::Type::Val(ir::ValType::Cons)) => { Instruction::RefNull(HeapType::Concrete(self.cons_type)) } - ir::LocalType::Type(ir::Type::Val(ir::ValType::Closure)) => { + ir::LocalType::Type(ir::Type::Val(ir::ValType::Closure(_))) => { Instruction::RefNull(HeapType::Concrete(self.closure_type)) } ir::LocalType::Type(ir::Type::Val(ir::ValType::Vector)) => { @@ -1081,6 +1081,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { if let ir::InstrKind::Terminator(..) = expr.kind { debug_assert!(expr.local.is_none()); debug_assert_eq!(i, bb.instrs.len() - 1); + // debug_assert_eq!(i, bb.instrs.len() - 1); } else { self.gen_assign(function, expr); } @@ -1488,7 +1489,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { self.module_generator.cons_type, ))); } - ir::ValType::Closure => { + ir::ValType::Closure(_) => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); function.instruction(&Instruction::RefCastNonNull(HeapType::Concrete( self.module_generator.closure_type, @@ -1551,7 +1552,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ir::ValType::Cons => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); } - ir::ValType::Closure => { + ir::ValType::Closure(_) => { function.instruction(&Instruction::LocalGet(self.local_id_to_idx(*val))); } ir::ValType::Vector => { @@ -1699,7 +1700,7 @@ impl<'a, 'b> FuncGenerator<'a, 'b> { ir::ValType::Symbol => self.module_generator.symbol_type, ir::ValType::Nil => self.module_generator.nil_type, ir::ValType::Cons => self.module_generator.cons_type, - ir::ValType::Closure => self.module_generator.closure_type, + ir::ValType::Closure(_) => self.module_generator.closure_type, ir::ValType::Vector => self.module_generator.vector_type, ir::ValType::UVector(kind) => { self.module_generator.uvector_kind_to_type_idx(*kind)