Repository: bytecodealliance/wasmtime
Pull request #7501 (2023-11-08 15:57:51)

winch: Properly derive a scratch register for arg assignment
This commit properly derives a scratch register for a particular WebAssembly type. The included spec test uncovered that the previous implementation used a int scratch register to assign float stack arguments, which resulted in a panic.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->


github-actions[bot]: #### Subscribe to Label Action

cc @saulecabrera

<details>
This issue or pull request has been labeled: "winch"

Thus the following users have been cc'd because of the following labels:

* saulecabrera: winch

To subscribe or unsubscribe from this label, edit the <code>.github/subscribe-to-label.json</code> configuration file.

[Learn more.](https://github.com/bytecodealliance/subscribe-to-label-action)
</details>

Patch:
diff --git a/cranelift/codegen/src/egraph.rs b/cranelift/codegen/src/egraph.rs
index dab240c3a..2d5802572 100644
--- a/cranelift/codegen/src/egraph.rs
+++ b/cranelift/codegen/src/egraph.rs
@@ -424,7 +424,7 @@ impl<'a> EgraphPass<'a> {
                 }
             }
         }
-        trace!("stats: {:?}", self.stats);
+        trace!("stats: {:#?}", self.stats);
         self.elaborate();
     }
 
diff --git a/cranelift/codegen/src/opts/arithmetic.isle b/cranelift/codegen/src/opts/arithmetic.isle
index 00c56c240..6553b8674 100644
--- a/cranelift/codegen/src/opts/arithmetic.isle
+++ b/cranelift/codegen/src/opts/arithmetic.isle
@@ -126,3 +126,41 @@
 ;; both of them since they cancel out.
 (rule (simplify (fmul ty (fneg ty x) (fneg ty y)))
       (fmul ty x y))
+
+;; (a op (b op (c op d))) ==> ((a op b) op (c op d))
+;;
+;; and
+;;
+;; (((a op b) op c) op d) ==> ((a op b) op (c op d))
+;;
+;; where `op` is an associative operation: `iadd`, `imul`, `band`, or `bxor`.
+;;
+;; This increases instruction-level parallelism and shrinks live ranges. It also
+;; canonicalizes into the shallow-and-wide form for reassociating constants
+;; together for cprop.
+;;
+;; NB: We subsume to avoid exponential e-node blow up due to reassociating very
+;; large chains of operations.
+;;
+;; TODO: We should add `bor` rules for this as well. Unfortunately, they
+;; conflict with our `bswap` recognizing rules when we `subsume`.
+
+(rule (simplify (iadd ty a (iadd ty b (iadd ty c d))))
+      (subsume (iadd ty (iadd ty a b) (iadd ty c d))))
+(rule (simplify (iadd ty (iadd ty (iadd ty a b) c) d))
+      (subsume (iadd ty (iadd ty a b) (iadd ty c d))))
+
+(rule (simplify (imul ty a (imul ty b (imul ty c d))))
+      (subsume (imul ty (imul ty a b) (imul ty c d))))
+(rule (simplify (imul ty (imul ty (imul ty a b) c) d))
+      (subsume (imul ty (imul ty a b) (imul ty c d))))
+
+(rule (simplify (band ty a (band ty b (band ty c d))))
+      (subsume (band ty (band ty a b) (band ty c d))))
+(rule (simplify (band ty (band ty (band ty a b) c) d))
+      (subsume (band ty (band ty a b) (band ty c d))))
+
+(rule (simplify (bxor ty a (bxor ty b (bxor ty c d))))
+      (subsume (bxor ty (bxor ty a b) (bxor ty c d))))
+(rule (simplify (bxor ty (bxor ty (bxor ty a b) c) d))
+      (subsume (bxor ty (bxor ty a b) (bxor ty c d))))
diff --git a/cranelift/codegen/src/opts/cprop.isle b/cranelift/codegen/src/opts/cprop.isle
index 26580a7ea..e44404e05 100644
--- a/cranelift/codegen/src/opts/cprop.isle
+++ b/cranelift/codegen/src/opts/cprop.isle
@@ -208,3 +208,30 @@
       (ushr ty (ushr ty a c) b))
 (rule (simplify (sshr ty (sshr ty a@(iconst _ _) b) c@(iconst _ _)))
       (sshr ty (sshr ty a c) b))
+
+;; When we operations that are both commutative and associative, reassociate
+;; constants together for cprop:
+;;
+;; ((a op B) op (c op D)) ==> ((a op c) op (B op D))
+;;
+;; Where `op` is one of: `iadd`, `imul`, `band`, `bor`, or `bxor`.
+(rule (simplify (iadd ty
+                      (iadd ty a b@(iconst _ _))
+                      (iadd ty c d@(iconst _ _))))
+      (iadd ty (iadd ty a c) (iadd ty b d)))
+(rule (simplify (imul ty
+                      (imul ty a b@(iconst _ _))
+                      (imul ty c d@(iconst _ _))))
+      (imul ty (imul ty a c) (imul ty b d)))
+(rule (simplify (band ty
+                      (band ty a b@(iconst _ _))
+                      (band ty c d@(iconst _ _))))
+      (band ty (band ty a c) (band ty b d)))
+(rule (simplify (bor ty
+                      (bor ty a b@(iconst _ _))
+                      (bor ty c d@(iconst _ _))))
+      (bor ty (bor ty a c) (bor ty b d)))
+(rule (simplify (bxor ty
+                      (bxor ty a b@(iconst _ _))
+                      (bxor ty c d@(iconst _ _))))
+      (bxor ty (bxor ty a c) (bxor ty b d)))
diff --git a/cranelift/codegen/src/unionfind.rs b/cranelift/codegen/src/unionfind.rs
index a51f305c2..8e2be3ba3 100644
--- a/cranelift/codegen/src/unionfind.rs
+++ b/cranelift/codegen/src/unionfind.rs
@@ -13,6 +13,7 @@ pub struct UnionFind<Idx: EntityRef> {
 
 #[derive(Clone, Debug, PartialEq)]
 struct Val<Idx>(Idx);
+
 impl<Idx: EntityRef + ReservedValue> Default for Val<Idx> {
     fn default() -> Self {
         Self(Idx::reserved_value())
diff --git a/winch/codegen/src/abi/mod.rs b/winch/codegen/src/abi/mod.rs
index b4edd3a69..aafaf5ce6 100644
--- a/winch/codegen/src/abi/mod.rs
+++ b/winch/codegen/src/abi/mod.rs
@@ -95,6 +95,15 @@ pub(crate) trait ABI {
     /// Returns the designated floating point scratch register.
     fn float_scratch_reg() -> Reg;
 
+    /// Returns the designated scratch register for the given [WasmType].
+    fn scratch_for(ty: &WasmType) -> Reg {
+        match ty {
+            WasmType::I32 | WasmType::I64 => Self::scratch_reg(),
+            WasmType::F32 | WasmType::F64 => Self::float_scratch_reg(),
+            _ => unimplemented!(),
+        }
+    }
+
     /// Returns the frame pointer register.
     fn fp_reg() -> Reg;
 
diff --git a/winch/codegen/src/codegen/call.rs b/winch/codegen/src/codegen/call.rs
index 977b4757e..981efc139 100644
--- a/winch/codegen/src/codegen/call.rs
+++ b/winch/codegen/src/codegen/call.rs
@@ -99,8 +99,7 @@ impl FnCall {
         let call_stack_space = Self::save(context, masm, &sig);
 
         let reserved_stack = masm.call(arg_stack_space, |masm| {
-            let scratch = <M::ABI as ABI>::scratch_reg();
-            Self::assign(sig, context, masm, scratch);
+            Self::assign(sig, context, masm);
             kind
         });
 
@@ -247,15 +246,11 @@ impl FnCall {
     }
 
     /// Assign arguments for the function call.
-    fn assign<M: MacroAssembler>(
-        sig: &ABISig,
-        context: &mut CodeGenContext,
-        masm: &mut M,
-        scratch: Reg,
-    ) {
+    fn assign<M: MacroAssembler>(sig: &ABISig, context: &mut CodeGenContext, masm: &mut M) {
         let arg_count = sig.params.len();
         let stack = &context.stack;
         let mut stack_values = stack.peekn(arg_count);
+
         for arg in &sig.params {
             let val = stack_values
                 .next()
@@ -267,6 +262,7 @@ impl FnCall {
                 &ABIArg::Stack { ty, offset } => {
                     let addr = masm.address_at_sp(*offset);
                     let size: OperandSize = (*ty).into();
+                    let scratch = <M::ABI as ABI>::scratch_for(&ty);
                     context.move_val_to_reg(val, scratch, masm);
                     masm.store(scratch.into(), addr, size);
                 }
diff --git a/winch/codegen/src/codegen/context.rs b/winch/codegen/src/codegen/context.rs
index ce7551a69..70d4c583e 100644
--- a/winch/codegen/src/codegen/context.rs
+++ b/winch/codegen/src/codegen/context.rs
@@ -454,11 +454,7 @@ impl<'a, 'builtins> CodeGenContext<'a, 'builtins> {
             Val::Local(local) => {
                 let slot = frame.get_local(local.index).expect("valid local at slot");
                 let addr = masm.local_address(&slot);
-                let scratch = match slot.ty {
-                    WasmType::I32 | WasmType::I64 => <M::ABI as ABI>::scratch_reg(),
-                    WasmType::F32 | WasmType::F64 => <M::ABI as ABI>::float_scratch_reg(),
-                    WasmType::V128 | WasmType::Ref(_) => unimplemented!(),
-                };
+                let scratch = <M::ABI as ABI>::scratch_for(&slot.ty);
                 masm.load(addr, scratch, slot.ty.into());
                 let stack_slot = masm.push(scratch, slot.ty.into());
                 *v = Val::mem(slot.ty, stack_slot);


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/associative-and-commutative.clif b/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
new file mode 100644
index 000000000..37bfa8da5
--- /dev/null
+++ b/cranelift/filetests/filetests/egraph/associative-and-commutative.clif
@@ -0,0 +1,81 @@
+test optimize
+set opt_level=speed
+target x86_64
+
+function %iadd_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = iadd v0, v1
+    v5 = iadd v4, v2
+    v6 = iadd v5, v3
+    return v6
+; check:  v4 = iadd v0, v1
+; nextln: v7 = iadd v2, v3
+; nextln: v8 = iadd v4, v7
+; check: return v8
+}
+
+function %iadd_shallow_and_wide_twice(i32, i32, i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32, v4: i32, v5: i32):
+    v6 = iadd v0, v1
+    v7 = iadd v6, v2
+    v8 = iadd v7, v3
+    v9 = iadd v8, v4
+    v10 = iadd v9, v5
+    return v10
+; check:  v6 = iadd v0, v1
+; nextln: v11 = iadd v2, v3
+; nextln: v12 = iadd v6, v11
+; nextln: v16 = iadd v4, v5
+; nextln: v23 = iadd v12, v16
+; nextln: return v23
+}
+
+function %iadd_reassociate_then_cprop(i32, i32) -> i32 {
+block0(v0: i32, v1: i32):
+    v2 = iconst.i32 42
+    v3 = iconst.i32 36
+    v4 = iadd v0, v2
+    v5 = iadd v1, v3
+    v6 = iadd v4, v5
+    return v6
+; check:  v7 = iadd v0, v1
+; nextln: iconst.i32 78
+; nextln: v16 = iadd v7, v14  ; v14 = 78
+; check:  return v16
+}
+
+function %imul_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = imul v0, v1
+    v5 = imul v4, v2
+    v6 = imul v5, v3
+    return v6
+; check:  v4 = imul v0, v1
+; nextln: v7 = imul v2, v3
+; nextln: v8 = imul v4, v7
+; check: return v8
+}
+
+function %band_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = band v0, v1
+    v5 = band v4, v2
+    v6 = band v5, v3
+    return v6
+; check:  v4 = band v0, v1
+; nextln: v7 = band v2, v3
+; nextln: v8 = band v4, v7
+; check: return v8
+}
+
+function %bxor_shallow_and_wide(i32, i32, i32, i32) -> i32 {
+block0(v0: i32, v1: i32, v2: i32, v3: i32):
+    v4 = bxor v0, v1
+    v5 = bxor v4, v2
+    v6 = bxor v5, v3
+    return v6
+; check:  v4 = bxor v0, v1
+; nextln: v7 = bxor v2, v3
+; nextln: v8 = bxor v4, v7
+; check: return v8
+}
diff --git a/cranelift/filetests/filetests/egraph/icmp.clif b/cranelift/filetests/filetests/egraph/icmp.clif
index 1a8874845..ef60cbf95 100644
--- a/cranelift/filetests/filetests/egraph/icmp.clif
+++ b/cranelift/filetests/filetests/egraph/icmp.clif
@@ -53,4 +53,3 @@ block0:
 ;     v5 = icmp ne v2, v3  ; v3 = 0
 ;     return v5
 ; }
-
diff --git a/tests/misc_testsuite/winch/call.wast b/tests/misc_testsuite/winch/call.wast
index 84e92519a..c76b20387 100644
--- a/tests/misc_testsuite/winch/call.wast
+++ b/tests/misc_testsuite/winch/call.wast
@@ -144,6 +144,13 @@
     (block (result i32) (i32.const 2) (call $const-i32) (br_table 0 0))
   )
 
+  (func $return-from-long-argument-list-helper (param f32 i32 i32 f64 f32 f32 f32 f64 f32 i32 i32 f32 f64 i64 i64 i32 i64 i64 f32 i64 i64 i64 i32 f32 f32 f32 f64 f32 i32 i64 f32 f64 f64 f32 i32 f32 f32 f64 i64 f64 i32 i64 f32 f64 i32 i32 i32 i64 f64 i32 i64 i64 f64 f64 f64 f64 f64 f64 i32 f32 f64 f64 i32 i64 f32 f32 f32 i32 f64 f64 f64 f64 f64 f32 i64 i64 i32 i32 i32 f32 f64 i32 i64 f32 f32 f32 i32 i32 f32 f64 i64 f32 f64 f32 f32 f32 i32 f32 i64 i32) (result i32)
+    (local.get 99)
+  )
+
+  (func (export "return-from-long-argument-list") (param i32) (result i32)
+    (call $return-from-long-argument-list-helper (f32.const 0) (i32.const 0) (i32.const 0) (f64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (i64.const 0) (i64.const 0) (f32.const 0) (i64.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (f64.const 0) (f32.const 0) (i32.const 0) (f32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (i32.const 0) (i32.const 0) (i32.const 0) (i64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (i64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f64.const 0) (f32.const 0) (i64.const 0) (i64.const 0) (i32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i32.const 0) (i64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (i32.const 0) (f32.const 0) (f64.const 0) (i64.const 0) (f32.const 0) (f64.const 0) (f32.const 0) (f32.const 0) (f32.const 0) (i32.const 0) (f32.const 0) (i64.const 0) (local.get 0))
+  )
 )
 
 (assert_return (invoke "type-i32") (i32.const 0x132))
@@ -193,3 +200,4 @@
 
 (assert_return (invoke "as-br_table-first") (i32.const 0x132))
 (assert_return (invoke "as-br_table-last") (i32.const 2))
+(assert_return (invoke "return-from-long-argument-list" (i32.const 42)) (i32.const 42))