Translate Wasm block params into Variables instead of CLIF block params

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

@@ -1,12 +1,12 @@
-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 wasmtime_environ::{PtrSize, TagIndex, TypeIndex, WasmResult, WasmValType, wasm_unsupported};
 
 fn control_context_size(triple: &target_lexicon::Triple) -> WasmResult<u8> {
@@ -1522,23 +1522,21 @@ 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);
+            set_block_params(env, builder, target_block, &suspend_args);
+            builder.ins().jump(target_block, &[]);
         }
 
         preamble_blocks

crates/cranelift/src/translate/func_translator.rs

@@ -11,7 +11,8 @@ use crate::translate::translation_utils::get_vmctx_value_label;
 use cranelift_codegen::entity::EntityRef;
 use cranelift_codegen::ir::{self, Block, InstBuilder, ValueLabel};
 use cranelift_codegen::timing;
-use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
+use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
+use smallvec::SmallVec;
 use wasmparser::{BinaryReader, FuncValidator, FunctionBody, OperatorsReader, WasmModuleResources};
 use wasmtime_environ::{TypeConvert, WasmResult};
 
@@ -84,12 +85,34 @@ impl FuncTranslator {
         // Set up the translation state with a single pushed control block representing the whole
         // function and its return values.
         let exit_block = builder.create_block();
-        builder.append_block_params_for_function_returns(exit_block);
         environ
             .stacks
             .initialize(&builder.func.signature, exit_block);
 
         parse_local_decls(&mut reader, &mut builder, num_params, environ, validator)?;
+
+        // Represent the function's return values as `Variable`s, exactly like
+        // Wasm block parameters, rather than as real CLIF block parameters.
+        //
+        // Note this must happen *after* `parse_local_decls`: the Wasm locals'
+        // `Variable`s must be indexed contiguously after the parameters, so we
+        // can only declare these additional variables once all locals exist.
+        let return_types: SmallVec<[ir::Type; 6]> = builder
+            .func
+            .signature
+            .returns
+            .iter()
+            .map(|ret| ret.value_type)
+            .collect();
+        let return_vars: SmallVec<[Variable; 6]> = return_types
+            .iter()
+            .map(|ty| builder.declare_var(*ty))
+            .collect();
+        environ
+            .stacks
+            .block_param_vars
+            .insert(exit_block, return_vars);
+
         parse_function_body(validator, reader, &mut builder, environ)?;
 
         builder.finalize(environ.target_config());

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_entity::SecondaryMap;
+use cranelift_frontend::{FunctionBuilder, Variable};
+use smallvec::SmallVec;
 use std::vec::Vec;
 use wasmtime_environ::FrameStackShape;
 
@@ -267,6 +269,24 @@ 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.
+    ///
+    /// The only blocks with real CLIF block parameters are the entry block
+    /// (function parameters) and `try_table` catch blocks (the exception
+    /// payload, filled in by the exception ABI).
+    pub(crate) block_param_vars: SecondaryMap<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 +311,7 @@ impl FuncTranslationStacks {
             stack: Vec::new(),
             stack_shape: Vec::new(),
             control_stack: Vec::new(),
+            block_param_vars: SecondaryMap::new(),
             handlers: HandlerState::default(),
             reachable: true,
         }
@@ -301,6 +322,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 +454,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,55 @@ where
     });
 }
 
-/// Create a `Block` with the given Wasm parameters.
-pub fn block_with_params<PE: TargetEnvironment + ?Sized>(
+/// Set the parameter `Variable`s of `destination` to `values` ahead of an
+/// argument-less branch to that block.
+pub fn set_block_params(
+    environ: &FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder,
+    destination: ir::Block,
+    values: &[ir::Value],
+) {
+    let vars = &environ.stacks.block_param_vars[destination];
+    debug_assert_eq!(vars.len(), values.len());
+    for (var, val) in vars.iter().zip(values) {
+        builder.def_var(*var, *val);
+    }
+}
+
+/// 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);
     }
+    let old = environ.stacks.block_param_vars.insert(block, vars);
+    debug_assert!(old.is_none());
     Ok(block)
 }
 

tests/disas/alias-region-globals.wat

@@ -26,8 +26,8 @@
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64, v2: i32):
-;; @0039                               v4 = load.i64 notrap aligned readonly can_move region2 v0+48
-;; @0039                               store notrap aligned region3 v2, v4
+;; @0039                               v3 = load.i64 notrap aligned readonly can_move region2 v0+48
+;; @0039                               store notrap aligned region3 v2, v3
 ;; @003d                               store notrap aligned region4 v2, v0+80
 ;; @0041                               jump block1
 ;;

tests/disas/alias-region-memories.wat

@@ -28,14 +28,14 @@
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i32):
-;; @003b                               v6 = load.i64 notrap aligned readonly can_move region2 v0+48
-;; @003b                               v7 = load.i64 notrap aligned readonly can_move region3 v6
-;; @003b                               v5 = uextend.i64 v2
-;; @003b                               v8 = iadd v7, v5
-;; @003b                               store little region5 v3, v8
-;; @0042                               v10 = load.i64 notrap aligned readonly can_move region3 v0+80
-;; @0042                               v11 = iadd v10, v5
-;; @0042                               store little region6 v3, v11
+;; @003b                               v5 = load.i64 notrap aligned readonly can_move region2 v0+48
+;; @003b                               v6 = load.i64 notrap aligned readonly can_move region3 v5
+;; @003b                               v4 = uextend.i64 v2
+;; @003b                               v7 = iadd v6, v4
+;; @003b                               store little region5 v3, v7
+;; @0042                               v9 = load.i64 notrap aligned readonly can_move region3 v0+80
+;; @0042                               v10 = iadd v9, v4
+;; @0042                               store little region6 v3, v10
 ;; @004b                               jump block1
 ;;
 ;;                                 block1:

tests/disas/alias-region-tables.wat

@@ -33,47 +33,47 @@
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i64):
-;; @0043                               v5 = load.i64 notrap aligned readonly can_move region2 v0+48
-;; @0043                               v6 = load.i64 notrap aligned region4 v5+8
-;; @0043                               v11 = load.i64 notrap aligned region3 v5
-;; @0043                               v17 = iconst.i64 1
-;; @0043                               v18 = bor v3, v17  ; v17 = 1
-;; @0043                               v7 = ireduce.i32 v6
-;; @0043                               v8 = icmp uge v2, v7
-;; @0043                               v15 = iconst.i64 0
-;; @0043                               v9 = uextend.i64 v2
-;; @0043                               v12 = iconst.i64 3
-;; @0043                               v13 = ishl v9, v12  ; v12 = 3
-;; @0043                               v14 = iadd v11, v13
-;; @0043                               v16 = select_spectre_guard v8, v15, v14  ; v15 = 0
-;; @0043                               store user6 aligned region5 v18, v16
-;; @0049                               v19 = load.i64 notrap aligned region4 v0+80
-;; @0049                               v23 = load.i64 notrap aligned region3 v0+72
-;; @0049                               v20 = ireduce.i32 v19
-;; @0049                               v21 = icmp uge v2, v20
-;; @0049                               v26 = iadd v23, v13
-;; @0049                               v28 = select_spectre_guard v21, v15, v26  ; v15 = 0
-;; @0049                               store user6 aligned region6 v18, v28
-;; @004d                               v44 = iconst.i64 -2
-;; @004d                               v45 = band v18, v44  ; v44 = -2
-;; @004d                               brif v18, block3(v45), block2
+;; @0043                               v4 = load.i64 notrap aligned readonly can_move region2 v0+48
+;; @0043                               v5 = load.i64 notrap aligned region4 v4+8
+;; @0043                               v10 = load.i64 notrap aligned region3 v4
+;; @0043                               v16 = iconst.i64 1
+;; @0043                               v17 = bor v3, v16  ; v16 = 1
+;; @0043                               v6 = ireduce.i32 v5
+;; @0043                               v7 = icmp uge v2, v6
+;; @0043                               v14 = iconst.i64 0
+;; @0043                               v8 = uextend.i64 v2
+;; @0043                               v11 = iconst.i64 3
+;; @0043                               v12 = ishl v8, v11  ; v11 = 3
+;; @0043                               v13 = iadd v10, v12
+;; @0043                               v15 = select_spectre_guard v7, v14, v13  ; v14 = 0
+;; @0043                               store user6 aligned region5 v17, v15
+;; @0049                               v18 = load.i64 notrap aligned region4 v0+80
+;; @0049                               v22 = load.i64 notrap aligned region3 v0+72
+;; @0049                               v19 = ireduce.i32 v18
+;; @0049                               v20 = icmp uge v2, v19
+;; @0049                               v25 = iadd v22, v12
+;; @0049                               v27 = select_spectre_guard v20, v14, v25  ; v14 = 0
+;; @0049                               store user6 aligned region6 v17, v27
+;; @004d                               v43 = iconst.i64 -2
+;; @004d                               v44 = band v17, v43  ; v43 = -2
+;; @004d                               brif v17, block3(v44), block2
 ;;
 ;;                                 block2 cold:
-;; @004d                               v47 = iconst.i32 0
-;; @004d                               v49 = call fn0(v0, v47, v9)  ; v47 = 0
-;; @004d                               jump block3(v49)
+;; @004d                               v46 = iconst.i32 0
+;; @004d                               v48 = call fn0(v0, v46, v8)  ; v46 = 0
+;; @004d                               jump block3(v48)
 ;;
-;;                                 block3(v46: i64):
-;; @004d                               v52 = load.i32 user7 aligned readonly v46+16
-;; @004d                               v50 = load.i64 notrap aligned readonly can_move region7 v0+40
-;; @004d                               v51 = load.i32 notrap aligned readonly can_move v50
-;; @004d                               v53 = icmp eq v52, v51
-;; @004d                               trapz v53, user8
-;; @004d                               v55 = load.i64 notrap aligned readonly v46+8
-;; @004d                               v56 = load.i64 notrap aligned readonly v46+24
-;; @004d                               v57 = call_indirect sig0, v55(v56, v0)
+;;                                 block3(v45: i64):
+;; @004d                               v51 = load.i32 user7 aligned readonly v45+16
+;; @004d                               v49 = load.i64 notrap aligned readonly can_move region7 v0+40
+;; @004d                               v50 = load.i32 notrap aligned readonly can_move v49
+;; @004d                               v52 = icmp eq v51, v50
+;; @004d                               trapz v52, user8
+;; @004d                               v54 = load.i64 notrap aligned readonly v45+8
+;; @004d                               v55 = load.i64 notrap aligned readonly v45+24
+;; @004d                               v56 = call_indirect sig0, v54(v55, v0)
 ;; @0050                               jump block1
 ;;
 ;;                                 block1:
-;; @0050                               return v57
+;; @0050                               return v56
 ;; }

tests/disas/basic-wat-test.wat

@@ -21,17 +21,17 @@
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i32):
-;; @0021                               v5 = uextend.i64 v2
-;; @0021                               v6 = load.i64 notrap aligned readonly can_move region2 v0+56
-;; @0021                               v7 = iadd v6, v5
-;; @0021                               v8 = load.i32 little region4 v7
-;; @0026                               v9 = uextend.i64 v3
-;; @0026                               v10 = load.i64 notrap aligned readonly can_move region2 v0+56
-;; @0026                               v11 = iadd v10, v9
-;; @0026                               v12 = load.i32 little region4 v11
-;; @0029                               v13 = iadd v8, v12
+;; @0021                               v4 = uextend.i64 v2
+;; @0021                               v5 = load.i64 notrap aligned readonly can_move region2 v0+56
+;; @0021                               v6 = iadd v5, v4
+;; @0021                               v7 = load.i32 little region4 v6
+;; @0026                               v8 = uextend.i64 v3
+;; @0026                               v9 = load.i64 notrap aligned readonly can_move region2 v0+56
+;; @0026                               v10 = iadd v9, v8
+;; @0026                               v11 = load.i32 little region4 v10
+;; @0029                               v12 = iadd v7, v11
 ;; @002a                               jump block1
 ;;
 ;;                                 block1:
-;; @002a                               return v13
+;; @002a                               return v12
 ;; }