cranelift-frontend: Propagate needs-stack-map from variables to values during finalization (bytecodealliance#10468)

* cranelift-frontend: Propagate needs-stack-map from variables to values during finalization

Rather than trying to propagate the needs-stack-map bit from variables to values
online, while we are building the IR and defining and using variables, wait
until the function is being finalized, and then propagate everything all at
once. This avoids potential repeated work and is easier to ensure that it is
complete and covers all values associated with a variable, since by the time we
are finalizing the function, we won't add any new values for a variable that we
will need to keep track of and propagate this information to.

This also means that we can remove the `params_added_to_blocks` vector from the
SSA side effects structure, since it was only used to online-update the
`stack_map_values` set.

* Initialize the env-logger in `#[wasmtime_test]`

* Fix needs-stack-map set iteration

For reasons I do not understand, the `EntitySet::keys` method includes keys that
are not in the set, and we have unit tests asserting this bizarre behavior. Very
perplexing. So I added a new method to iterate over just the elements of the
set.

Author: fitzgen

cranelift/entity/src/lib.rs

@@ -287,7 +287,7 @@ pub use self::keys::Keys;
 pub use self::list::{EntityList, ListPool};
 pub use self::map::SecondaryMap;
 pub use self::primary::PrimaryMap;
-pub use self::set::EntitySet;
+pub use self::set::{EntitySet, SetIter};
 pub use self::signed::Signed;
 pub use self::sparse::{SparseMap, SparseMapValue, SparseSet};
 pub use self::unsigned::Unsigned;

cranelift/entity/src/set.rs

@@ -89,11 +89,56 @@ where
         self.bitset.clear()
     }
 
-    /// Iterate over all the keys in this set.
+    /// Iterate over all the keys up to the maximum in this set.
+    ///
+    /// This will yield intermediate keys on the way up to the max key, even if
+    /// they are not contained within the set.
+    ///
+    /// ```
+    /// use cranelift_entity::{entity_impl, EntityRef, EntitySet};
+    ///
+    /// #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+    /// struct Entity(u32);
+    /// entity_impl!(Entity);
+    ///
+    /// let mut set = EntitySet::new();
+    /// set.insert(Entity::new(2));
+    ///
+    /// let mut keys = set.keys();
+    /// assert_eq!(keys.next(), Some(Entity::new(0)));
+    /// assert_eq!(keys.next(), Some(Entity::new(1)));
+    /// assert_eq!(keys.next(), Some(Entity::new(2)));
+    /// assert!(keys.next().is_none());
+    /// ```
     pub fn keys(&self) -> Keys<K> {
         Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
     }
 
+    /// Iterate over the elements of this set.
+    ///
+    /// ```
+    /// use cranelift_entity::{entity_impl, EntityRef, EntitySet};
+    ///
+    /// #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+    /// struct Entity(u32);
+    /// entity_impl!(Entity);
+    ///
+    /// let mut set = EntitySet::new();
+    /// set.insert(Entity::new(2));
+    /// set.insert(Entity::new(3));
+    ///
+    /// let mut iter = set.iter();
+    /// assert_eq!(iter.next(), Some(Entity::new(2)));
+    /// assert_eq!(iter.next(), Some(Entity::new(3)));
+    /// assert!(iter.next().is_none());
+    /// ```
+    pub fn iter(&self) -> SetIter<K> {
+        SetIter {
+            inner: self.bitset.iter(),
+            _phantom: PhantomData,
+        }
+    }
+
     /// Insert the element at `k`.
     ///
     /// Returns `true` if `k` was not present in the set, i.e. this is a
@@ -110,6 +155,25 @@ where
     }
 }
 
+/// An iterator over the elements in an `EntitySet`.
+pub struct SetIter<'a, K> {
+    inner: cranelift_bitset::compound::Iter<'a>,
+    _phantom: PhantomData<K>,
+}
+
+impl<K> Iterator for SetIter<'_, K>
+where
+    K: EntityRef,
+{
+    type Item = K;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        let k = self.inner.next()?;
+        Some(K::new(k))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

cranelift/frontend/src/frontend.rs

@@ -473,12 +473,6 @@ impl<'a> FunctionBuilder<'a> {
         };
         self.handle_ssa_side_effects(side_effects);
 
-        // If the variable was declared as needing stack maps, then we should
-        // have propagated that to the value as well.
-        if cfg!(debug_assertions) && self.func_ctx.stack_map_vars.contains(var) {
-            assert!(self.func_ctx.stack_map_values.contains(val));
-        }
-
         Ok(val)
     }
 
@@ -494,6 +488,8 @@ impl<'a> FunctionBuilder<'a> {
     /// an error if the value supplied does not match the type the variable was
     /// declared to have.
     pub fn try_def_var(&mut self, var: Variable, val: Value) -> Result<(), DefVariableError> {
+        log::trace!("try_def_var: {var:?} = {val:?}");
+
         let var_ty = *self
             .func_ctx
             .types
@@ -503,11 +499,6 @@ impl<'a> FunctionBuilder<'a> {
             return Err(DefVariableError::TypeMismatch(var, val));
         }
 
-        // If `var` needs inclusion in stack maps, then `val` does too.
-        if self.func_ctx.stack_map_vars.contains(var) {
-            self.declare_value_needs_stack_map(val);
-        }
-
         self.func_ctx.ssa.def_var(var, val, self.position.unwrap());
         Ok(())
     }
@@ -720,6 +711,19 @@ impl<'a> FunctionBuilder<'a> {
             }
         }
 
+        // Propagate the needs-stack-map bit from variables to each of their
+        // associated values.
+        for var in self.func_ctx.stack_map_vars.iter() {
+            for val in self.func_ctx.ssa.values_for_var(var) {
+                log::trace!("propagating needs-stack-map from {var:?} to {val:?}");
+                debug_assert_eq!(self.func.dfg.value_type(val), self.func_ctx.types[var]);
+                self.func_ctx.stack_map_values.insert(val);
+            }
+        }
+
+        // If we have any values that need inclusion in stack maps, then we need
+        // to run our pass to spill those values to the stack at safepoints and
+        // generate stack maps.
         if !self.func_ctx.stack_map_values.is_empty() {
             self.func_ctx
                 .safepoints
@@ -1180,24 +1184,13 @@ impl<'a> FunctionBuilder<'a> {
     fn handle_ssa_side_effects(&mut self, side_effects: SideEffects) {
         let SideEffects {
             instructions_added_to_blocks,
-            params_added_to_blocks,
         } = side_effects;
 
         for modified_block in instructions_added_to_blocks {
             if self.is_pristine(modified_block) {
                 self.func_ctx.status[modified_block] = BlockStatus::Partial;
             }
         }
-
-        // Propagate needs-inclusion-in-stack-maps metadata from variables to
-        // their block parameter SSA values.
-        if !self.func_ctx.stack_map_vars.is_empty() {
-            for (val, var) in params_added_to_blocks {
-                if self.func_ctx.stack_map_vars.contains(var) {
-                    self.func_ctx.stack_map_values.insert(val);
-                }
-            }
-        }
     }
 }
 

cranelift/frontend/src/ssa.rs

@@ -67,21 +67,14 @@ pub struct SideEffects {
     /// This field signals if it is the case and return the `Block` to which the initialization has
     /// been added.
     pub instructions_added_to_blocks: Vec<Block>,
-
-    /// When a variable does not have a single, dominating definition of a use,
-    /// we must insert block parameters to control-flow join points. This field
-    /// contains all of the block parameters we inserted, along with the
-    /// variable that they are associated with.
-    pub params_added_to_blocks: Vec<(Value, Variable)>,
 }
 
 impl SideEffects {
     fn is_empty(&self) -> bool {
         let Self {
             instructions_added_to_blocks,
-            params_added_to_blocks,
         } = self;
-        instructions_added_to_blocks.is_empty() && params_added_to_blocks.is_empty()
+        instructions_added_to_blocks.is_empty()
     }
 }
 
@@ -195,6 +188,12 @@ fn emit_zero(ty: Type, mut cur: FuncCursor) -> Value {
 /// Phi functions.
 ///
 impl SSABuilder {
+    /// Get all of the values associated with the given variable that we have
+    /// inserted in the function thus far.
+    pub fn values_for_var(&self, var: Variable) -> impl Iterator<Item = Value> + '_ {
+        self.variables[var].values().filter_map(|v| v.expand())
+    }
+
     /// Declares a new definition of a variable in a given basic block.
     /// The SSA value is passed as an argument because it should be created with
     /// `ir::DataFlowGraph::append_result`.
@@ -318,7 +317,6 @@ impl SSABuilder {
         // find a usable definition. So create one.
         let val = func.dfg.append_block_param(block, ty);
         var_defs[block] = PackedOption::from(val);
-        self.side_effects.params_added_to_blocks.push((val, var));
 
         // Now every predecessor needs to pass its definition of this variable to the newly added
         // block parameter. To do that we have to "recursively" call `use_var`, but there are two

crates/test-macros/src/wasmtime_test.rs

@@ -279,6 +279,7 @@ fn expand(test_config: &TestConfig, func: Fn) -> Result<TokenStream> {
             #should_panic
             #(#attrs)*
             #asyncness fn #test_name() {
+                let _ = env_logger::try_init();
                 let mut config = Config::new();
                 wasmtime_test_util::wasmtime_wast::apply_test_config(
                     &mut config,

fuzz/fuzz_targets/cranelift-icache.rs

@@ -43,6 +43,8 @@ pub struct FunctionWithIsa {
 
 impl FunctionWithIsa {
     pub fn generate(u: &mut Unstructured) -> anyhow::Result<Self> {
+        let _ = env_logger::try_init();
+
         // We filter out targets that aren't supported in the current build
         // configuration after randomly choosing one, instead of randomly choosing
         // a supported one, so that the same fuzz input works across different build