Translate Wasm block params into Variables instead of CLIF block params

crates/cranelift/src/func_environ/stack_switching/instructions.rs

@@ -1,12 +1,13 @@
-use cranelift_codegen::ir::BlockArg;
-use itertools::{Either, Itertools};
-
+use crate::translate::set_block_params;
 use crate::trap::TranslateTrap;
+use cranelift_codegen::ir::BlockArg;
 use cranelift_codegen::ir::condcodes::*;
 use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::{self, MemFlagsData};
 use cranelift_codegen::ir::{Block, BlockCall, InstBuilder, JumpTableData};
 use cranelift_frontend::FunctionBuilder;
+use itertools::{Either, Itertools};
+use smallvec::SmallVec;
 use wasmtime_environ::{PtrSize, TagIndex, TypeIndex, WasmResult, WasmValType, wasm_unsupported};
 
 fn control_context_size(triple: &target_lexicon::Triple) -> WasmResult<u8> {
@@ -1522,23 +1523,22 @@ pub(crate) fn translate_resume<'a>(
                 .collect();
 
             let values = suspended_contref.values(env, builder);
-            let mut suspend_args: Vec<BlockArg> = values
-                .load_data_entries(env, builder, &param_types)
-                .into_iter()
-                .map(|v| BlockArg::Value(v))
-                .collect();
+            let mut suspend_args: Vec<ir::Value> =
+                values.load_data_entries(env, builder, &param_types);
 
             // At the suspend site, we store the suspend args in the the
             // `values` buffer of the VMContRef that was active at the time that
             // the suspend instruction was performed.
-            suspend_args.push(BlockArg::Value(suspended_contobj));
+            suspend_args.push(suspended_contobj);
 
             // We clear the suspend args. This is mostly for consistency. Note
             // that we don't zero out the data buffer, we still need it for the
 
             values.clear(env, builder, false);
 
-            builder.ins().jump(target_block, &suspend_args);
+            let mut tmp = SmallVec::new();
+            let args = set_block_params(env, builder, &mut tmp, target_block, &suspend_args);
+            builder.ins().jump(target_block, args);
         }
 
         preamble_blocks

crates/cranelift/src/translate/stack.rs

@@ -5,7 +5,9 @@
 //! a single function.
 
 use cranelift_codegen::ir::{self, Block, ExceptionTag, Inst, Value};
-use cranelift_frontend::FunctionBuilder;
+use cranelift_frontend::{FunctionBuilder, Variable};
+use smallvec::SmallVec;
+use std::collections::HashMap;
 use std::vec::Vec;
 use wasmtime_environ::FrameStackShape;
 
@@ -267,6 +269,25 @@ pub struct FuncTranslationStacks {
     pub(crate) stack_shape: Vec<FrameStackShape>,
     /// A stack of active control flow operations at this point in the input wasm function.
     pub(crate) control_stack: Vec<ControlStackFrame>,
+    /// Maps a CLIF block representing a Wasm control-flow target to the
+    /// `Variable`s that hold its Wasm stack parameters.
+    ///
+    /// Rather than giving these blocks CLIF block parameters and passing the
+    /// Wasm operand stack values as block arguments when branching to them, we
+    /// represent each Wasm stack parameter as a `Variable`. When branching to
+    /// such a block we `def_var` the variables and emit an argument-less
+    /// branch; when we begin translating the block we `use_var` each variable
+    /// and push the results onto the operand stack. This lets
+    /// `cranelift-frontend`'s SSA construction decide whether a real block
+    /// parameter is actually needed (i.e. only when multiple predecessors pass
+    /// differing values) instead of pessimistically creating one for every
+    /// Wasm block.
+    ///
+    /// Note that not every CLIF block is present in this map: the function's
+    /// exit block keeps real CLIF block parameters (for the function returns),
+    /// and various internal blocks (e.g. an `if`'s consequent, or the
+    /// fall-through after a `br_if`) have no parameters at all.
+    pub(crate) block_param_vars: HashMap<Block, SmallVec<[Variable; 6]>>,
     /// Exception handler state, updated as we enter and exit
     /// `try_table` scopes and attached to each call that we make.
     pub(crate) handlers: HandlerState,
@@ -291,6 +312,7 @@ impl FuncTranslationStacks {
             stack: Vec::new(),
             stack_shape: Vec::new(),
             control_stack: Vec::new(),
+            block_param_vars: HashMap::new(),
             handlers: HandlerState::default(),
             reachable: true,
         }
@@ -301,6 +323,7 @@ impl FuncTranslationStacks {
         debug_assert!(self.stack_shape.is_empty());
         debug_assert!(self.control_stack.is_empty());
         debug_assert!(self.handlers.is_empty());
+        self.block_param_vars.clear();
         self.reachable = true;
     }
 
@@ -432,13 +455,6 @@ impl FuncTranslationStacks {
         &self.stack[self.stack.len() - n..]
     }
 
-    /// Peek at the top `n` values on the stack in the order they were pushed.
-    pub(crate) fn peekn_mut(&mut self, n: usize) -> &mut [Value] {
-        self.ensure_length_is_at_least(n);
-        let len = self.stack.len();
-        &mut self.stack[len - n..]
-    }
-
     fn push_block_impl(
         &mut self,
         following_code: Block,

crates/cranelift/src/translate/translation_utils.rs

@@ -1,10 +1,12 @@
 //! Helper functions and structures for the translation.
+use crate::func_environ::FuncEnvironment;
 use crate::translate::environ::TargetEnvironment;
 use core::u32;
 use cranelift_codegen::ir;
 use cranelift_frontend::FunctionBuilder;
+use smallvec::SmallVec;
 use wasmparser::{FuncValidator, WasmModuleResources};
-use wasmtime_environ::WasmResult;
+use wasmtime_environ::{TypeConvert, WasmResult};
 
 /// Get the parameter and result types for the given Wasm blocktype.
 pub fn blocktype_params_results<'a, T>(
@@ -41,40 +43,70 @@ where
     });
 }
 
-/// Create a `Block` with the given Wasm parameters.
-pub fn block_with_params<PE: TargetEnvironment + ?Sized>(
+/// Set up the Wasm stack parameters of `destination` to `values` ahead of an
+/// argument-less branch to that block.
+///
+/// If `destination` is a Wasm control-flow target (i.e. it has associated
+/// parameter `Variable`s) then we `def_var` each of those variables and the
+/// branch carries no block arguments. Otherwise `destination` has real CLIF
+/// block parameters (e.g. the function exit block), so we return the values as
+/// block arguments to pass directly.
+pub fn set_block_params<'a>(
+    environ: &FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder,
+    tmp: &'a mut SmallVec<[ir::BlockArg; 16]>,
+    destination: ir::Block,
+    values: &[ir::Value],
+) -> &'a [ir::BlockArg] {
+    debug_assert!(tmp.is_empty());
+    match environ.stacks.block_param_vars.get(&destination) {
+        Some(vars) => {
+            debug_assert_eq!(vars.len(), values.len());
+            for (var, val) in vars.iter().zip(values) {
+                builder.def_var(*var, *val);
+            }
+            &[]
+        }
+        None => {
+            tmp.extend(values.iter().map(|v| ir::BlockArg::from(*v)));
+            tmp.as_slice()
+        }
+    }
+}
+
+/// Create a `Block` representing a Wasm control-flow target with the given Wasm
+/// stack parameters.
+///
+/// Rather than giving the block CLIF block parameters, we create a
+/// `cranelift_frontend::Variable` for each Wasm stack parameter and record the
+/// block-to-variables mapping in `environ.stacks.block_param_vars`. See the
+/// `block_param_vars` docs for more details.
+pub fn block_with_params(
     builder: &mut FunctionBuilder,
     params: impl IntoIterator<Item = wasmparser::ValType>,
-    environ: &PE,
+    environ: &mut FuncEnvironment<'_>,
 ) -> WasmResult<ir::Block> {
     let block = builder.create_block();
+    let mut vars = SmallVec::<[_; 6]>::new();
     for ty in params {
-        match ty {
-            wasmparser::ValType::I32 => {
-                builder.append_block_param(block, ir::types::I32);
-            }
-            wasmparser::ValType::I64 => {
-                builder.append_block_param(block, ir::types::I64);
-            }
-            wasmparser::ValType::F32 => {
-                builder.append_block_param(block, ir::types::F32);
-            }
-            wasmparser::ValType::F64 => {
-                builder.append_block_param(block, ir::types::F64);
-            }
+        let (clif_ty, needs_stack_map) = match ty {
+            wasmparser::ValType::I32 => (ir::types::I32, false),
+            wasmparser::ValType::I64 => (ir::types::I64, false),
+            wasmparser::ValType::F32 => (ir::types::F32, false),
+            wasmparser::ValType::F64 => (ir::types::F64, false),
             wasmparser::ValType::Ref(rt) => {
                 let hty = environ.convert_heap_type(rt.heap_type())?;
-                let (ty, needs_stack_map) = environ.reference_type(hty);
-                let val = builder.append_block_param(block, ty);
-                if needs_stack_map {
-                    builder.declare_value_needs_stack_map(val);
-                }
-            }
-            wasmparser::ValType::V128 => {
-                builder.append_block_param(block, ir::types::I8X16);
+                environ.reference_type(hty)
             }
+            wasmparser::ValType::V128 => (ir::types::I8X16, false),
+        };
+        let var = builder.declare_var(clif_ty);
+        if needs_stack_map {
+            builder.declare_var_needs_stack_map(var);
         }
+        vars.push(var);
     }
+    environ.stacks.block_param_vars.insert(block, vars);
     Ok(block)
 }
 

tests/disas/block-params-br_if-single-pred.wat

@@ -0,0 +1,61 @@
+;;! target = "x86_64"
+
+;; A `br_if` to a block whose only reachable predecessor is that `br_if` (the
+;; fall-through path ends in `unreachable`) does not need block parameters for
+;; the branched-with values.
+
+(module
+  (import "" "f" (func $f (result i32)))
+  (import "" "g" (func $g (param i32 i32)))
+  (import "" "h" (func $h (param i32 i32)))
+  (func
+    block (result i32 i32)
+      i32.const 1
+      i32.const 2
+      call $f
+      br_if 0
+      call $g
+      unreachable
+    end
+    call $h
+  )
+)
+;; function u0:0(i64 vmctx, i64) tail {
+;;     region0 = 8 "VMContext+0x8"
+;;     region1 = 268435480 "VMStoreContext+0x18"
+;;     region2 = 72 "VMContext+0x48"
+;;     region3 = 56 "VMContext+0x38"
+;;     region4 = 104 "VMContext+0x68"
+;;     region5 = 88 "VMContext+0x58"
+;;     region6 = 136 "VMContext+0x88"
+;;     region7 = 120 "VMContext+0x78"
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly can_move region0 gv0+8
+;;     gv2 = load.i64 notrap aligned region1 gv1+24
+;;     sig0 = (i64 vmctx, i64) -> i32 tail
+;;     sig1 = (i64 vmctx, i64, i32, i32) tail
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64):
+;; @0039                               v2 = iconst.i32 1
+;; @003b                               v3 = iconst.i32 2
+;; @003d                               v4 = load.i64 notrap aligned readonly can_move region2 v0+72
+;; @003d                               v5 = load.i64 notrap aligned readonly can_move region3 v0+56
+;; @003d                               v6 = call_indirect sig0, v5(v4, v0)
+;; @003f                               brif v6, block2, block3
+;;
+;;                                 block3:
+;; @0041                               v7 = load.i64 notrap aligned readonly can_move region4 v0+104
+;; @0041                               v8 = load.i64 notrap aligned readonly can_move region5 v0+88
+;; @0041                               call_indirect sig1, v8(v7, v0, v2, v3)  ; v2 = 1, v3 = 2
+;; @0043                               trap user12
+;;
+;;                                 block2:
+;; @0045                               v9 = load.i64 notrap aligned readonly can_move region6 v0+136
+;; @0045                               v10 = load.i64 notrap aligned readonly can_move region7 v0+120
+;; @0045                               call_indirect sig1, v10(v9, v0, v2, v3)  ; v2 = 1, v3 = 2
+;; @0047                               jump block1
+;;
+;;                                 block1:
+;; @0047                               return
+;; }

tests/disas/block-params-if-else-same-values.wat

@@ -0,0 +1,56 @@
+;;! target = "x86_64"
+
+;; The control-flow join after an if/else diamond has two predecessors, but when
+;; both the consequent and the alternative pass the same values, no block
+;; parameter is needed.
+
+(module
+  (import "" "f" (func $f))
+  (import "" "g" (func $g))
+  (func (param i32 i32 i32) (result i32 i32)
+    local.get 0
+    if (result i32 i32)
+      call $f
+      local.get 1
+      local.get 2
+    else
+      call $g
+      local.get 1
+      local.get 2
+    end
+  )
+)
+;; function u0:0(i64 vmctx, i64, i32, i32, i32) -> i32, i32 tail {
+;;     region0 = 8 "VMContext+0x8"
+;;     region1 = 268435480 "VMStoreContext+0x18"
+;;     region2 = 72 "VMContext+0x48"
+;;     region3 = 56 "VMContext+0x38"
+;;     region4 = 104 "VMContext+0x68"
+;;     region5 = 88 "VMContext+0x58"
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly can_move region0 gv0+8
+;;     gv2 = load.i64 notrap aligned region1 gv1+24
+;;     sig0 = (i64 vmctx, i64) tail
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i32, v4: i32):
+;; @0033                               brif v2, block2, block4
+;;
+;;                                 block2:
+;; @0035                               v7 = load.i64 notrap aligned readonly can_move region2 v0+72
+;; @0035                               v8 = load.i64 notrap aligned readonly can_move region3 v0+56
+;; @0035                               call_indirect sig0, v8(v7, v0)
+;; @003b                               jump block3
+;;
+;;                                 block4:
+;; @003c                               v9 = load.i64 notrap aligned readonly can_move region4 v0+104
+;; @003c                               v10 = load.i64 notrap aligned readonly can_move region5 v0+88
+;; @003c                               call_indirect sig0, v10(v9, v0)
+;; @0042                               jump block3
+;;
+;;                                 block3:
+;; @0043                               jump block1
+;;
+;;                                 block1:
+;; @0043                               return v3, v4
+;; }

tests/disas/block-params-loop-single-iteration.wat

@@ -0,0 +1,37 @@
+;;! target = "x86_64"
+
+;; A `loop` that never branches back to the loop header only ever runs a single
+;; iteration, so its header block has just one predecessor and should have no
+;; block parameters.
+
+(module
+  (func (param i32) (result i32)
+    local.get 0
+    (loop (param i32) (result i32)
+      i32.const 1
+      i32.add
+    )
+  )
+)
+;; function u0:0(i64 vmctx, i64, i32) -> i32 tail {
+;;     region0 = 8 "VMContext+0x8"
+;;     region1 = 268435480 "VMStoreContext+0x18"
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly can_move region0 gv0+8
+;;     gv2 = load.i64 notrap aligned region1 gv1+24
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64, v2: i32):
+;; @001b                               jump block2
+;;
+;;                                 block2:
+;; @001d                               v5 = iconst.i32 1
+;; @001f                               v6 = iadd.i32 v2, v5  ; v5 = 1
+;; @0020                               jump block3
+;;
+;;                                 block3:
+;; @0021                               jump block1
+;;
+;;                                 block1:
+;; @0021                               return v6
+;; }