|
| 1 | +use std::{cell::Cell, collections::HashMap}; |
| 2 | + |
| 3 | +use wasmparser::Operator; |
| 4 | + |
| 5 | +use crate::loader::{ |
| 6 | + BlockKind, |
| 7 | + dag::{self, Dag, NodeInput, Operation, WasmValue}, |
| 8 | + settings::{MaybeConstant, Settings}, |
| 9 | +}; |
| 10 | + |
| 11 | +/// This is an optional optimization pass that collapses constants into |
| 12 | +/// the node that uses them, if the ISA supports it. |
| 13 | +/// |
| 14 | +/// For example, if the output ISA is RISC-V, in an `i32.add` instruction that |
| 15 | +/// has a constant node `5` as input, the value `5` could be collapsed into an |
| 16 | +/// immediate for `addi` instruction. |
| 17 | +/// |
| 18 | +/// This pass only happens if `Settings::get_const_collapse_processor()` is |
| 19 | +/// implemented by the user. |
| 20 | +/// |
| 21 | +/// Returns the number of constants collapsed. |
| 22 | +pub fn constant_collapse<'a, S: Settings<'a>>(settings: &S, dag: &mut Dag) -> usize { |
| 23 | + if let Some(processor) = settings.get_const_collapse_processor() { |
| 24 | + let mut finder = ConstantInputFinder::new(); |
| 25 | + recursive_constant_collapse(&mut finder, &processor, &mut dag.nodes, HashMap::new()) |
| 26 | + } else { |
| 27 | + 0 |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +fn recursive_constant_collapse( |
| 32 | + finder: &mut ConstantInputFinder, |
| 33 | + processor: &impl Fn(&Operator, &[MaybeConstant]), |
| 34 | + nodes: &mut [dag::Node], |
| 35 | + const_block_inputs: HashMap<u32, WasmValue>, |
| 36 | +) -> usize { |
| 37 | + let mut num_collapsed = 0; |
| 38 | + |
| 39 | + for i in 0..nodes.len() { |
| 40 | + // We split the nodes so that previous nodes can be borrowed (immutably) independently |
| 41 | + // from the currently processed node (mutably borrowed). |
| 42 | + let (prev_nodes, node) = nodes.split_at_mut(i); |
| 43 | + let node = &mut node[0]; |
| 44 | + |
| 45 | + match &mut node.operation { |
| 46 | + dag::Operation::WASMOp(operator) => { |
| 47 | + let (has_constant, input_constants) = finder.find_constant_inputs( |
| 48 | + &node.inputs, |
| 49 | + prev_nodes, |
| 50 | + const_block_inputs.clone(), |
| 51 | + ); |
| 52 | + if has_constant { |
| 53 | + // Call the user-provided processor to potentially mark constants to be collapsed. |
| 54 | + processor(operator, input_constants); |
| 55 | + |
| 56 | + // Collapse any inputs that were marked. |
| 57 | + for (input, maybe_const) in node.inputs.iter_mut().zip(input_constants) { |
| 58 | + if let MaybeConstant::ReferenceConstant { |
| 59 | + value, |
| 60 | + must_collapse, |
| 61 | + } = maybe_const |
| 62 | + && must_collapse.get() |
| 63 | + { |
| 64 | + *input = NodeInput::Constant(value.clone()); |
| 65 | + num_collapsed += 1; |
| 66 | + } |
| 67 | + } |
| 68 | + } |
| 69 | + } |
| 70 | + dag::Operation::Block { sub_dag, kind } => { |
| 71 | + // Constant input collapse to the block node itself is not supported, |
| 72 | + // but we can recurse into the block's sub-DAG. |
| 73 | + |
| 74 | + // In order to collapse constants that come from the block inputs, |
| 75 | + // we need to know which block inputs are constants. |
| 76 | + let const_inputs = if matches!(kind, BlockKind::Block) { |
| 77 | + finder |
| 78 | + .find_constant_inputs(&node.inputs, prev_nodes, const_block_inputs.clone()) |
| 79 | + .1 |
| 80 | + .iter() |
| 81 | + .enumerate() |
| 82 | + .filter_map(|(idx, maybe_const)| match maybe_const { |
| 83 | + MaybeConstant::ReferenceConstant { value, .. } => { |
| 84 | + Some((idx as u32, value.clone())) |
| 85 | + } |
| 86 | + MaybeConstant::CollapsedConstant(value) => { |
| 87 | + Some((idx as u32, value.clone())) |
| 88 | + } |
| 89 | + MaybeConstant::NonConstant => None, |
| 90 | + }) |
| 91 | + .collect() |
| 92 | + } else { |
| 93 | + // Loops have multiple entry points, so we can't be sure a constant input |
| 94 | + // at entry will be always constant. Besides, an optimized WASM wouldn't |
| 95 | + // have constant inputs to loops anyway. |
| 96 | + HashMap::new() |
| 97 | + }; |
| 98 | + |
| 99 | + num_collapsed += recursive_constant_collapse( |
| 100 | + finder, |
| 101 | + processor, |
| 102 | + &mut sub_dag.nodes, |
| 103 | + const_inputs, |
| 104 | + ); |
| 105 | + } |
| 106 | + _ => { |
| 107 | + // Constant collapse is not supported for other node types. |
| 108 | + } |
| 109 | + } |
| 110 | + } |
| 111 | + |
| 112 | + num_collapsed |
| 113 | +} |
| 114 | + |
| 115 | +struct ConstantInputFinder { |
| 116 | + /// This buffer will be reused to avoid allocations on each node. |
| 117 | + buffer: Vec<MaybeConstant>, |
| 118 | +} |
| 119 | + |
| 120 | +impl ConstantInputFinder { |
| 121 | + fn new() -> Self { |
| 122 | + Self { buffer: Vec::new() } |
| 123 | + } |
| 124 | + |
| 125 | + /// Finds which of the given inputs are constants. |
| 126 | + fn find_constant_inputs<'a>( |
| 127 | + &'a mut self, |
| 128 | + inputs: &[NodeInput], |
| 129 | + nodes: &mut [dag::Node], |
| 130 | + const_block_inputs: HashMap<u32, WasmValue>, |
| 131 | + ) -> (bool, &'a [MaybeConstant]) { |
| 132 | + let mut has_constant = false; |
| 133 | + |
| 134 | + // Check if any of the inputs are constants. |
| 135 | + self.buffer.clear(); |
| 136 | + self.buffer.extend(inputs.iter().map(|input| { |
| 137 | + let origin = match input { |
| 138 | + NodeInput::Reference(value_origin) => value_origin, |
| 139 | + NodeInput::Constant(wasm_value) => { |
| 140 | + // In default Womir pipeline, there shouldn't be any constant |
| 141 | + // before this pass. But we handle it anyway, in case this pass |
| 142 | + // is used in a different context by the user. |
| 143 | + return MaybeConstant::CollapsedConstant(wasm_value.clone()); |
| 144 | + } |
| 145 | + }; |
| 146 | + |
| 147 | + match &nodes[origin.node].operation { |
| 148 | + Operation::Inputs => { |
| 149 | + // This refers to the input of the block, we need to check if it's a constant. |
| 150 | + const_block_inputs.get(&origin.output_idx).cloned() |
| 151 | + } |
| 152 | + Operation::WASMOp(op) => WasmValue::try_from(op).ok(), |
| 153 | + _ => None, |
| 154 | + } |
| 155 | + .map_or(MaybeConstant::NonConstant, |value| { |
| 156 | + has_constant = true; |
| 157 | + MaybeConstant::ReferenceConstant { |
| 158 | + value, |
| 159 | + must_collapse: Cell::new(false), |
| 160 | + } |
| 161 | + }) |
| 162 | + })); |
| 163 | + |
| 164 | + (has_constant, &self.buffer) |
| 165 | + } |
| 166 | +} |
0 commit comments