cranelift-wasm: Only allocate if vectors need bitcasts (#4543)

For wasm programs using SIMD vector types, the type known at function
entry or exit may not match the type used within the body of the
function, so we have to bitcast them. We have to check all calls and
returns for this condition, which involves comparing a subset of a
function's signature with the CLIF types we're trying to use. Currently,
this check heap-allocates a short-lived Vec for the appropriate subset
of the signature.

But most of the time none of the values need a bitcast. So this patch
avoids allocating unless at least one bitcast is actually required, and
only saves the pointers to values that need fixing up. I leaned heavily
on iterators to keep space usage constant until we discover allocation
is necessary after all.

I don't think it's possible to eliminate the allocation entirely,
because the function signature we're examining is borrowed from the
FuncBuilder, but we need to mutably borrow the FuncBuilder to insert the
bitcast instructions. Fortunately, the FromIterator implementation for
Vec doesn't reserve any heap space if the iterator is empty, so we can
unconditionally collect into a Vec and still avoid unnecessary
allocations.

Since the relationship between a function signature and a list of CLIF
values is somewhat complicated, I refactored all the uses of
`bitcast_arguments` and `wasm_param_types`. Instead there's
`bitcast_wasm_params` and `bitcast_wasm_returns` which share a helper
that combines the previous pair of functions into one.

According to DHAT, when compiling the Sightglass Spidermonkey benchmark,
this avoids 52k allocations averaging about 9 bytes each, which are
freed on average within 2k instructions.

Most Sightglass benchmarks, including Spidermonkey, show no performance
difference with this change. Only one has a slowdown, and it's small:

compilation :: nanoseconds :: benchmarks/shootout-matrix/benchmark.wasm

  Δ = 689373.34 ± 593720.78 (confidence = 99%)

  lazy-bitcast.so is 0.94x to 1.00x faster than main-e121c209f.so!
  main-e121c209f.so is 1.00x to 1.06x faster than lazy-bitcast.so!

  [19741713 21375844.46 32976047] lazy-bitcast.so
  [19345471 20686471.12 30872267] main-e121c209f.so

But several Sightglass benchmarks have notable speed-ups, with smaller
improvements for shootout-ed25519, meshoptimizer, and pulldown-cmark,
and larger ones as follows:

compilation :: nanoseconds :: benchmarks/bz2/benchmark.wasm

  Δ = 20071545.47 ± 2950014.92 (confidence = 99%)

  lazy-bitcast.so is 1.26x to 1.36x faster than main-e121c209f.so!
  main-e121c209f.so is 0.73x to 0.80x faster than lazy-bitcast.so!

  [55995164 64849257.68 89083031] lazy-bitcast.so
  [79382460 84920803.15 98016185] main-e121c209f.so

compilation :: nanoseconds :: benchmarks/blake3-scalar/benchmark.wasm

  Δ = 16620780.61 ± 5395162.13 (confidence = 99%)

  lazy-bitcast.so is 1.14x to 1.28x faster than main-e121c209f.so!
  main-e121c209f.so is 0.77x to 0.88x faster than lazy-bitcast.so!

  [54604352 79877776.35 103666598] lazy-bitcast.so
  [94011835 96498556.96 106200091] main-e121c209f.so

compilation :: nanoseconds :: benchmarks/intgemm-simd/benchmark.wasm

  Δ = 36891254.34 ± 12403663.50 (confidence = 99%)

  lazy-bitcast.so is 1.12x to 1.24x faster than main-e121c209f.so!
  main-e121c209f.so is 0.79x to 0.90x faster than lazy-bitcast.so!

  [131610845 201289587.88 247341883] lazy-bitcast.so
  [232065032 238180842.22 250957563] main-e121c209f.so

cranelift/wasm/src/code_translator.rs

@@ -89,6 +89,7 @@ use cranelift_codegen::ir::{
 };
 use cranelift_codegen::packed_option::ReservedValue;
 use cranelift_frontend::{FunctionBuilder, Variable};
+use itertools::Itertools;
 use smallvec::SmallVec;
 use std::cmp;
 use std::convert::TryFrom;
@@ -540,10 +541,7 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             };
             {
                 let return_args = state.peekn_mut(return_count);
-                let return_types = wasm_param_types(&builder.func.signature.returns, |i| {
-                    environ.is_wasm_return(&builder.func.signature, i)
-                });
-                bitcast_arguments(return_args, &return_types, builder);
+                bitcast_wasm_returns(environ, return_args, builder);
                 match environ.return_mode() {
                     ReturnMode::NormalReturns => builder.ins().return_(return_args),
                     ReturnMode::FallthroughReturn => {
@@ -575,13 +573,13 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let (fref, num_args) = state.get_direct_func(builder.func, *function_index, environ)?;
 
             // Bitcast any vector arguments to their default type, I8X16, before calling.
-            let callee_signature =
-                &builder.func.dfg.signatures[builder.func.dfg.ext_funcs[fref].signature];
             let args = state.peekn_mut(num_args);
-            let types = wasm_param_types(&callee_signature.params, |i| {
-                environ.is_wasm_parameter(&callee_signature, i)
-            });
-            bitcast_arguments(args, &types, builder);
+            bitcast_wasm_params(
+                environ,
+                builder.func.dfg.ext_funcs[fref].signature,
+                args,
+                builder,
+            );
 
             let call = environ.translate_call(
                 builder.cursor(),
@@ -612,12 +610,8 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             let callee = state.pop1();
 
             // Bitcast any vector arguments to their default type, I8X16, before calling.
-            let callee_signature = &builder.func.dfg.signatures[sigref];
             let args = state.peekn_mut(num_args);
-            let types = wasm_param_types(&callee_signature.params, |i| {
-                environ.is_wasm_parameter(&callee_signature, i)
-            });
-            bitcast_arguments(args, &types, builder);
+            bitcast_wasm_params(environ, sigref, args, builder);
 
             let call = environ.translate_call_indirect(
                 builder,
@@ -3018,40 +3012,75 @@ fn pop2_with_bitcast(
     (bitcast_a, bitcast_b)
 }
 
-/// A helper for bitcasting a sequence of values (e.g. function arguments). If a value is a
-/// vector type that does not match its expected type, this will modify the value in place to point
-/// to the result of a `raw_bitcast`. This conversion is necessary to translate Wasm code that
-/// uses `V128` as function parameters (or implicitly in block parameters) and still use specific
-/// CLIF types (e.g. `I32X4`) in the function body.
-pub fn bitcast_arguments(
-    arguments: &mut [Value],
-    expected_types: &[Type],
-    builder: &mut FunctionBuilder,
-) {
-    assert_eq!(arguments.len(), expected_types.len());
-    for (i, t) in expected_types.iter().enumerate() {
-        if t.is_vector() {
+fn bitcast_arguments<'a>(
+    builder: &FunctionBuilder,
+    arguments: &'a mut [Value],
+    params: &[ir::AbiParam],
+    param_predicate: impl Fn(usize) -> bool,
+) -> Vec<(Type, &'a mut Value)> {
+    let filtered_param_types = params
+        .iter()
+        .enumerate()
+        .filter(|(i, _)| param_predicate(*i))
+        .map(|(_, param)| param.value_type);
+
+    // zip_eq, from the itertools::Itertools trait, is like Iterator::zip but panics if one
+    // iterator ends before the other. The `param_predicate` is required to select exactly as many
+    // elements of `params` as there are elements in `arguments`.
+    let pairs = filtered_param_types.zip_eq(arguments.iter_mut());
+
+    // The arguments which need to be bitcasted are those which have some vector type but the type
+    // expected by the parameter is not the same vector type as that of the provided argument.
+    pairs
+        .filter(|(param_type, _)| param_type.is_vector())
+        .filter(|(param_type, arg)| {
+            let arg_type = builder.func.dfg.value_type(**arg);
             assert!(
-                builder.func.dfg.value_type(arguments[i]).is_vector(),
+                arg_type.is_vector(),
                 "unexpected type mismatch: expected {}, argument {} was actually of type {}",
-                t,
-                arguments[i],
-                builder.func.dfg.value_type(arguments[i])
+                param_type,
+                *arg,
+                arg_type
             );
-            arguments[i] = optionally_bitcast_vector(arguments[i], *t, builder)
-        }
+
+            // This is the same check that would be done by `optionally_bitcast_vector`, except we
+            // can't take a mutable borrow of the FunctionBuilder here, so we defer inserting the
+            // bitcast instruction to the caller.
+            arg_type != *param_type
+        })
+        .collect()
+}
+
+/// A helper for bitcasting a sequence of return values for the function currently being built. If
+/// a value is a vector type that does not match its expected type, this will modify the value in
+/// place to point to the result of a `raw_bitcast`. This conversion is necessary to translate Wasm
+/// code that uses `V128` as function parameters (or implicitly in block parameters) and still use
+/// specific CLIF types (e.g. `I32X4`) in the function body.
+pub fn bitcast_wasm_returns<FE: FuncEnvironment + ?Sized>(
+    environ: &mut FE,
+    arguments: &mut [Value],
+    builder: &mut FunctionBuilder,
+) {
+    let changes = bitcast_arguments(builder, arguments, &builder.func.signature.returns, |i| {
+        environ.is_wasm_return(&builder.func.signature, i)
+    });
+    for (t, arg) in changes {
+        *arg = builder.ins().raw_bitcast(t, *arg);
     }
 }
 
-/// A helper to extract all the `Type` listings of each variable in `params`
-/// for only parameters the return true for `is_wasm`, typically paired with
-/// `is_wasm_return` or `is_wasm_parameter`.
-pub fn wasm_param_types(params: &[ir::AbiParam], is_wasm: impl Fn(usize) -> bool) -> Vec<Type> {
-    let mut ret = Vec::with_capacity(params.len());
-    for (i, param) in params.iter().enumerate() {
-        if is_wasm(i) {
-            ret.push(param.value_type);
-        }
+/// Like `bitcast_wasm_returns`, but for the parameters being passed to a specified callee.
+fn bitcast_wasm_params<FE: FuncEnvironment + ?Sized>(
+    environ: &mut FE,
+    callee_signature: ir::SigRef,
+    arguments: &mut [Value],
+    builder: &mut FunctionBuilder,
+) {
+    let callee_signature = &builder.func.dfg.signatures[callee_signature];
+    let changes = bitcast_arguments(builder, arguments, &callee_signature.params, |i| {
+        environ.is_wasm_parameter(&callee_signature, i)
+    });
+    for (t, arg) in changes {
+        *arg = builder.ins().raw_bitcast(t, *arg);
     }
-    ret
 }

cranelift/wasm/src/func_translator.rs

@@ -4,7 +4,7 @@
 //! function to Cranelift IR guided by a `FuncEnvironment` which provides information about the
 //! WebAssembly module and the runtime environment.
 
-use crate::code_translator::{bitcast_arguments, translate_operator, wasm_param_types};
+use crate::code_translator::{bitcast_wasm_returns, translate_operator};
 use crate::environ::{FuncEnvironment, ReturnMode};
 use crate::state::FuncTranslationState;
 use crate::translation_utils::get_vmctx_value_label;
@@ -255,10 +255,7 @@ fn parse_function_body<FE: FuncEnvironment + ?Sized>(
         if !builder.is_unreachable() {
             match environ.return_mode() {
                 ReturnMode::NormalReturns => {
-                    let return_types = wasm_param_types(&builder.func.signature.returns, |i| {
-                        environ.is_wasm_return(&builder.func.signature, i)
-                    });
-                    bitcast_arguments(&mut state.stack, &return_types, builder);
+                    bitcast_wasm_returns(environ, &mut state.stack, builder);
                     builder.ins().return_(&state.stack)
                 }
                 ReturnMode::FallthroughReturn => builder.ins().fallthrough_return(&state.stack),