Stratification of call graphs for parallel bottom-up inlining (bytecodealliance#11269)

* Stratification of call graphs for parallel bottom-up inlining

This commit takes a call graph and constructs a strata, which is essentially a
parallel execution plan. A strata consists of an ordered sequence of layers, and
a layer of an unordered set of functions. The `i`th layer must be processed
before the `i + 1`th layer, but functions within the same layer may be processed
in any order (and in parallel).

For example, when given the following tree-like call graph:

    +---+   +---+   +---+
    | a |-->| b |-->| c |
    +---+   +---+   +---+
      |       |
      |       |     +---+
      |       '---->| d |
      |             +---+
      |
      |     +---+   +---+
      '---->| e |-->| f |
            +---+   +---+
              |
              |     +---+
              '---->| g |
                    +---+

then stratification will produce these layers:

    [
        {c, d, f, g},
        {b, e},
        {a},
    ]

Our goal in constructing the layers is to maximize potential parallelism at each
layer. Logically, we do this by finding the strongly-connected components of the
input call graph and peeling off all of the leaves of SCCs'
condensation (i.e. the DAG that the SCCs form; see the documentation for the
`StronglyConnectedComponents::evaporation` method for details). These leaves
become the strata's first layer. The layer's components are removed from the
condensation graph, and we repeat the process, so that the condensation's new
leaves become the strata's second layer, and etc... until the condensation graph
is empty and all components have been processed. In practice we don't actually
mutate the condensation graph or remove its nodes but instead count how many
unprocessed dependencies each component has, and a component is ready for
inclusion in a layer once its unprocessed-dependencies count reaches zero.

This commit also renames the entity type for strongly-connected components from
`Component` to `Scc`, as I felt the former was a bit ambiguous given Wasm
components.

The next PR will extend Wasmtime's compilation driver code to actually make use
of this new infrastructure.

* Address review feedback

Author: fitzgen

cranelift/entity/src/map.rs

@@ -149,6 +149,23 @@ where
     }
 }
 
+impl<K, V> FromIterator<(K, V)> for SecondaryMap<K, V>
+where
+    K: EntityRef,
+    V: Clone + Default,
+{
+    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+        let (min, max) = iter.size_hint();
+        let cap = max.unwrap_or_else(|| 2 * min);
+        let mut map = Self::with_capacity(cap);
+        for (k, v) in iter {
+            map[k] = v;
+        }
+        map
+    }
+}
+
 /// Immutable indexing into an `SecondaryMap`.
 ///
 /// All keys are permitted. Untouched entries have the default value.

crates/wasmtime/src/compile.rs

@@ -31,6 +31,7 @@ use std::{
     borrow::Cow,
     collections::{BTreeMap, BTreeSet, btree_map},
     mem,
+    ops::Range,
 };
 
 use wasmtime_environ::CompiledFunctionBody;
@@ -43,7 +44,9 @@ use wasmtime_environ::{
     StaticModuleIndex,
 };
 
+mod call_graph;
 mod scc;
+mod stratify;
 
 mod code_builder;
 pub use self::code_builder::{CodeBuilder, CodeHint, HashedEngineCompileEnv};
@@ -1017,3 +1020,17 @@ impl Artifacts {
         self.modules.into_iter().next().unwrap().1
     }
 }
+
+/// Extend `dest` with `items` and return the range of indices in `dest` where
+/// they ended up.
+fn extend_with_range<T>(dest: &mut Vec<T>, items: impl IntoIterator<Item = T>) -> Range<u32> {
+    let start = dest.len();
+    let start = u32::try_from(start).unwrap();
+
+    dest.extend(items);
+
+    let end = dest.len();
+    let end = u32::try_from(end).unwrap();
+
+    start..end
+}

crates/wasmtime/src/compile/call_graph.rs

@@ -0,0 +1,107 @@
+//! Construction of the call-graph, for the purposes of inlining.
+//!
+//! These call graphs are not necessarily complete or accurate, and Wasmtime's
+//! soundness does not rely on those properties. First off, we do not attempt to
+//! understand indirect calls, which at their worst must force any call analysis
+//! give up and say "the callee could be absolutely any function". More
+//! interestingly, these call graphs are only used for scheduling bottom-up
+//! inlining, so the worst that inaccurate information can do is cause us to
+//! miss inlining opportunities or lose potential parallelism in our
+//! schedule. For best results, however, every direct call that is potentially
+//! inlinable should be reported when constructing these call graphs.
+
+#![cfg_attr(not(test), expect(dead_code, reason = "used in upcoming PRs"))]
+
+use super::*;
+use core::{
+    fmt::{self, Debug},
+    ops::Range,
+};
+use wasmtime_environ::{EntityRef, SecondaryMap};
+
+/// A call graph reified into a densely packed and quickly accessible
+/// representation.
+///
+/// In a call graph, nodes are functions, and an edge `f --> g` means that the
+/// function `f` calls the function `g`.
+pub struct CallGraph<Node>
+where
+    Node: EntityRef + Debug,
+{
+    /// A map from each node to the subslice of `self.edge_elems` that are its
+    /// edges.
+    edges: SecondaryMap<Node, Range<u32>>,
+
+    /// Densely packed edge elements for `self.edges`.
+    edge_elems: Vec<Node>,
+}
+
+impl<Node> Debug for CallGraph<Node>
+where
+    Node: EntityRef + Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        struct Edges<'a, Node: EntityRef + Debug>(&'a CallGraph<Node>);
+
+        impl<'a, Node: EntityRef + Debug> Debug for Edges<'a, Node> {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.debug_map()
+                    .entries(self.0.nodes().map(|n| (n, self.0.edges(n))))
+                    .finish()
+            }
+        }
+
+        f.debug_struct("CallGraph")
+            .field("edges", &Edges(self))
+            .finish()
+    }
+}
+
+impl<Node> CallGraph<Node>
+where
+    Node: EntityRef + Debug,
+{
+    /// Construct a new call graph.
+    ///
+    /// `funcs` should be an iterator over all function nodes in this call
+    /// graph's translation unit.
+    ///
+    /// The `get_calls` function should yield (by pushing onto the given `Vec`)
+    /// all of the callee function nodes that the given caller function node
+    /// calls.
+    pub fn new(
+        funcs: impl IntoIterator<Item = Node>,
+        get_calls: impl Fn(Node, &mut Vec<Node>) -> Result<()>,
+    ) -> Result<Self> {
+        let funcs = funcs.into_iter();
+
+        let (min, max) = funcs.size_hint();
+        let capacity = max.unwrap_or_else(|| 2 * min);
+        let mut edges = SecondaryMap::with_capacity(capacity);
+        let mut edge_elems = vec![];
+
+        let mut calls = vec![];
+        for caller in funcs {
+            debug_assert!(calls.is_empty());
+            get_calls(caller, &mut calls)?;
+
+            debug_assert_eq!(edges[caller], Range::default());
+            edges[caller] = extend_with_range(&mut edge_elems, calls.drain(..));
+        }
+
+        Ok(CallGraph { edges, edge_elems })
+    }
+
+    /// Get the function nodes in this call graph.
+    pub fn nodes(&self) -> impl ExactSizeIterator<Item = Node> {
+        self.edges.keys()
+    }
+
+    /// Get the callee function nodes that the given caller function node calls.
+    pub fn edges(&self, node: Node) -> &[Node] {
+        let Range { start, end } = self.edges[node].clone();
+        let start = usize::try_from(start).unwrap();
+        let end = usize::try_from(end).unwrap();
+        &self.edge_elems[start..end]
+    }
+}