Cranelift: introduce a function inliner (bytecodealliance#11210)

* Cranelift: introduce a function inliner

This comit adds "inlining as a library" to Cranelift; it does _not_ provide a
complete, off-the-shelf inlining solution. Cranelift's compilation context is
per-function and does not encompass the full call graph. It does not know which
functions are hot and which are cold, which have been marked the equivalent of
`#[inline(always)]` versus `#[inline(never)]`, etc... Only the Cranelift user
can understand these aspects of the full compilation pipeline, and these things
can be very different between (say) Wasmtime and `cg_clif`. Therefore, this
infrastructure does not attempt to define hueristics for when inlining a
particular call is likely beneficial. This module only provides hooks for the
Cranelift user to tell Cranelift whether a given call should be inlined or not,
and the mechanics to inline a callee into a particular call site when the user
directs Cranelift to do so.

This commit also creates a new kind of filetest that will always inline calls to
functions that have already been defined in the file. This lets us exercise the
inliner in filetests.

Fixes bytecodealliance#4127

* Address review feedback

* Require callee bodies are pre-legalized

Author: fitzgen

cranelift/codegen/src/context.rs

@@ -14,6 +14,7 @@ use crate::dominator_tree::DominatorTree;
 use crate::dominator_tree::DominatorTreePreorder;
 use crate::egraph::EgraphPass;
 use crate::flowgraph::ControlFlowGraph;
+use crate::inline::{Inline, do_inlining};
 use crate::ir::Function;
 use crate::isa::TargetIsa;
 use crate::legalizer::simple_legalize;
@@ -193,6 +194,13 @@ impl Context {
         Ok(())
     }
 
+    /// Perform function call inlining.
+    ///
+    /// Returns `true` if any function call was inlined, `false` otherwise.
+    pub fn inline(&mut self, inliner: impl Inline) -> CodegenResult<bool> {
+        do_inlining(&mut self.func, inliner)
+    }
+
     /// Compile the function,
     ///
     /// Run the function through all the passes necessary to generate

cranelift/codegen/src/inline.rs

@@ -65,10 +65,24 @@ impl Blocks {
         self.0.len()
     }
 
+    /// Reserves capacity for at least `additional` more elements to be
+    /// inserted.
+    pub fn reserve(&mut self, additional: usize) {
+        self.0.reserve(additional);
+    }
+
     /// Returns `true` if the given block reference is valid.
     pub fn is_valid(&self, block: Block) -> bool {
         self.0.is_valid(block)
     }
+
+    /// Iterate over all blocks, regardless whether a block is actually inserted
+    /// in the layout or not.
+    ///
+    /// Iterates in creation order, not layout order.
+    pub fn iter(&self) -> impl Iterator<Item = Block> {
+        self.0.keys()
+    }
 }
 
 impl Index<Block> for Blocks {
@@ -107,13 +121,6 @@ pub struct DataFlowGraph {
     results: SecondaryMap<Inst, ValueList>,
 
     /// User-defined stack maps.
-    ///
-    /// Not to be confused with the stack maps that `regalloc2` produces. These
-    /// are defined by the user in `cranelift-frontend`. These will eventually
-    /// replace the stack maps support in `regalloc2`, but in the name of
-    /// incrementalism and avoiding gigantic PRs that completely overhaul
-    /// Cranelift and Wasmtime at the same time, we are allowing them to live in
-    /// parallel for the time being.
     user_stack_maps: alloc::collections::BTreeMap<Inst, UserStackMapEntryVec>,
 
     /// basic blocks in the function and their parameters.
@@ -337,6 +344,11 @@ impl DataFlowGraph {
         self.values.push(data.into())
     }
 
+    /// The number of values defined in this DFG.
+    pub fn len_values(&self) -> usize {
+        self.values.len()
+    }
+
     /// Get an iterator over all values.
     pub fn values<'a>(&'a self) -> Values<'a> {
         Values {
@@ -599,6 +611,30 @@ impl DataFlowGraph {
         assert!(opcode.is_safepoint());
         self.user_stack_maps.entry(inst).or_default().push(entry);
     }
+
+    /// Append multiple stack map entries for the given call instruction.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the given instruction is not a (non-tail) call instruction.
+    pub fn append_user_stack_map_entries(
+        &mut self,
+        inst: Inst,
+        entries: impl IntoIterator<Item = UserStackMapEntry>,
+    ) {
+        for entry in entries {
+            self.append_user_stack_map_entry(inst, entry);
+        }
+    }
+
+    /// Take the stack map entries for a given instruction, leaving the
+    /// instruction without stack maps.
+    pub(crate) fn take_user_stack_map_entries(
+        &mut self,
+        inst: Inst,
+    ) -> Option<UserStackMapEntryVec> {
+        self.user_stack_maps.remove(&inst)
+    }
 }
 
 /// Where did a value come from?

cranelift/codegen/src/ir/dfg.rs

@@ -106,6 +106,15 @@ impl ExceptionTableData {
         DisplayExceptionTable { table: self, pool }
     }
 
+    /// Deep-clone this exception table.
+    pub fn deep_clone(&self, pool: &mut ValueListPool) -> Self {
+        Self {
+            targets: self.targets.iter().map(|b| b.deep_clone(pool)).collect(),
+            tags: self.tags.clone(),
+            sig: self.sig,
+        }
+    }
+
     /// Get the default target for the non-exceptional return case.
     pub fn normal_return(&self) -> &BlockCall {
         self.targets.last().unwrap()
@@ -140,6 +149,18 @@ impl ExceptionTableData {
             .zip(self.targets.iter_mut())
     }
 
+    /// The number of catch edges in this exception table.
+    pub fn len_catches(&self) -> usize {
+        self.tags.len()
+    }
+
+    /// Get the `index`th catch edge from this table.
+    pub fn get_catch(&self, index: usize) -> Option<(Option<ExceptionTag>, &BlockCall)> {
+        let tag = self.tags.get(index)?.expand();
+        let target = &self.targets[index];
+        Some((tag, target))
+    }
+
     /// Get all branch targets.
     pub fn all_branches(&self) -> &[BlockCall] {
         &self.targets[..]
@@ -156,11 +177,34 @@ impl ExceptionTableData {
         self.sig
     }
 
+    /// Get a mutable handle to this exception table's signature.
+    pub(crate) fn signature_mut(&mut self) -> &mut SigRef {
+        &mut self.sig
+    }
+
     /// Clears all entries in this exception table, but leaves the function signature.
     pub fn clear(&mut self) {
         self.tags.clear();
         self.targets.clear();
     }
+
+    /// Push a catch target onto this exception table.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this exception table has been cleared.
+    pub fn push_catch(&mut self, tag: Option<ExceptionTag>, block_call: BlockCall) {
+        assert_eq!(
+            self.tags.len() + 1,
+            self.targets.len(),
+            "cannot push onto an exception table that has been cleared"
+        );
+
+        self.tags.push(tag.into());
+
+        let target_index = self.targets.len() - 1;
+        self.targets.insert(target_index, block_call);
+    }
 }
 
 /// A wrapper for the context required to display a

cranelift/codegen/src/ir/exception_table.rs

@@ -68,6 +68,7 @@ pub use crate::ir::stackslot::{
 };
 pub use crate::ir::trapcode::TrapCode;
 pub use crate::ir::types::Type;
+pub(crate) use crate::ir::user_stack_maps::UserStackMapEntryVec;
 pub use crate::ir::user_stack_maps::{UserStackMap, UserStackMapEntry};
 
 use crate::entity::{PrimaryMap, SecondaryMap, entity_impl};

cranelift/codegen/src/ir/mod.rs

@@ -21,9 +21,9 @@ extern crate alloc;
 extern crate std;
 
 #[cfg(not(feature = "std"))]
-use hashbrown::{HashMap, hash_map};
+use hashbrown::{HashMap, HashSet, hash_map};
 #[cfg(feature = "std")]
-use std::collections::{HashMap, hash_map};
+use std::collections::{HashMap, HashSet, hash_map};
 
 pub use crate::context::Context;
 pub use crate::value_label::{LabelValueLoc, ValueLabelsRanges, ValueLocRange};
@@ -50,6 +50,7 @@ pub mod data_value;
 pub mod dbg;
 pub mod dominator_tree;
 pub mod flowgraph;
+pub mod inline;
 pub mod ir;
 pub mod isa;
 pub mod loop_analysis;

cranelift/codegen/src/lib.rs

@@ -14,7 +14,7 @@ rust-version.workspace = true
 workspace = true
 
 [dependencies]
-cranelift-codegen = { workspace = true, features = ["disas"] }
+cranelift-codegen = { workspace = true, features = ["disas", "timing"] }
 cranelift-frontend = { workspace = true }
 cranelift-interpreter = { workspace = true }
 cranelift-native = { workspace = true }

cranelift/filetests/Cargo.toml

@@ -0,0 +1,35 @@
+test inline precise-output
+target x86_64
+
+function %f0(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+    v2 = iadd v0, v1
+    return v2
+}
+
+; (no functions inlined into %f0)
+
+function %f1() -> i32 {
+    fn0 = %f0(i32, i32) -> i32
+block0():
+    v0 = iconst.i32 10
+    v1 = call fn0(v0, v0)
+    return v1
+}
+
+; function %f1() -> i32 fast {
+;     sig0 = (i32, i32) -> i32 fast
+;     fn0 = %f0 sig0
+;
+; block0:
+;     v0 = iconst.i32 10
+;     jump block1
+;
+; block1:
+;     v3 = iadd.i32 v0, v0  ; v0 = 10, v0 = 10
+;     jump block2(v3)
+;
+; block2(v2: i32):
+;     v1 -> v2
+;     return v1
+; }

cranelift/filetests/filetests/inline/basic.clif

@@ -0,0 +1,57 @@
+test inline precise-output
+target x86_64
+
+function %f0(i32, i32) -> i32 tail {
+    fn0 = %whatever(i32, i32) -> i32 tail
+block0(v0: i32, v1: i32):
+    br_table v0, block3, [block1(v1), block2]
+block1(v2: i32):
+    return v2
+block2:
+    v3 = iconst.i32 0
+    return v3
+block3:
+    return v0
+}
+
+; (no functions inlined into %f0)
+
+function %f1() -> i32 tail {
+    fn0 = %f0(i32, i32) -> i32 tail
+block0:
+    v0 = iconst.i32 42
+    v1 = call fn0(v0, v0)
+    v2 = iconst.i32 1
+    v3 = iadd v1, v2
+    return v3
+}
+
+; function %f1() -> i32 tail {
+;     sig0 = (i32, i32) -> i32 tail
+;     sig1 = (i32, i32) -> i32 tail
+;     fn0 = %f0 sig0
+;     fn1 = %whatever sig1
+;
+; block0:
+;     v0 = iconst.i32 42
+;     jump block1
+;
+; block1:
+;     br_table v0, block4, [block2(v0), block3]  ; v0 = 42, v0 = 42
+;
+; block2(v4: i32):
+;     jump block5(v4)
+;
+; block3:
+;     v6 = iconst.i32 0
+;     jump block5(v6)  ; v6 = 0
+;
+; block4:
+;     jump block5(v0)  ; v0 = 42
+;
+; block5(v5: i32):
+;     v1 -> v5
+;     v2 = iconst.i32 1
+;     v3 = iadd v1, v2  ; v2 = 1
+;     return v3
+; }

cranelift/filetests/filetests/inline/br-table.clif

@@ -0,0 +1,63 @@
+test inline precise-output
+target x86_64
+
+function %f0(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+    v2 = iconst.i32 0
+    jump block1(v2, v2)
+
+block1(v3: i32, v4: i32):
+    v5 = icmp eq v1, v4
+    brif v4, block3, block2
+
+block2:
+    v6 = iadd v3, v0
+    v7 = iconst.i32 1
+    v8 = iadd v4, v7
+    jump block1(v6, v8)
+
+block3:
+    return v3
+}
+
+; (no functions inlined into %f0)
+
+function %f1() -> i32 {
+    fn0 = %f0(i32, i32) -> i32
+block0:
+    v0 = iconst.i32 42
+    v1 = iconst.i32 10
+    v2 = call fn0(v0, v1)
+    return v2
+}
+
+; function %f1() -> i32 fast {
+;     sig0 = (i32, i32) -> i32 fast
+;     fn0 = %f0 sig0
+;
+; block0:
+;     v0 = iconst.i32 42
+;     v1 = iconst.i32 10
+;     jump block1
+;
+; block1:
+;     v6 = iconst.i32 0
+;     jump block2(v6, v6)  ; v6 = 0, v6 = 0
+;
+; block2(v3: i32, v4: i32):
+;     v7 = icmp.i32 eq v1, v4  ; v1 = 10
+;     brif v4, block4, block3
+;
+; block3:
+;     v8 = iadd.i32 v3, v0  ; v0 = 42
+;     v9 = iconst.i32 1
+;     v10 = iadd.i32 v4, v9  ; v9 = 1
+;     jump block2(v8, v10)
+;
+; block4:
+;     jump block5(v3)
+;
+; block5(v5: i32):
+;     v2 -> v5
+;     return v2
+; }