Winch: Remove specific scratch registers from ABI trait (#8999)

winch/codegen/src/abi/mod.rs

@@ -51,7 +51,7 @@ use crate::masm::SPOffset;
 use smallvec::SmallVec;
 use std::collections::HashSet;
 use std::ops::{Add, BitAnd, Not, Sub};
-use wasmtime_environ::{WasmFuncType, WasmHeapType, WasmRefType, WasmValType};
+use wasmtime_environ::{WasmFuncType, WasmValType};
 
 pub(crate) mod local;
 pub(crate) use local::*;
@@ -75,7 +75,7 @@ macro_rules! vmctx {
 /// designated scratch register for the given type.
 macro_rules! scratch {
     ($m:ident) => {
-        <$m::ABI as $crate::abi::ABI>::scratch_reg()
+        <$m::ABI as $crate::abi::ABI>::scratch_for(&wasmtime_environ::WasmValType::I64)
     };
     ($m:ident, $wasm_type:expr) => {
         <$m::ABI as $crate::abi::ABI>::scratch_for($wasm_type)
@@ -138,25 +138,8 @@ pub(crate) trait ABI {
         Self::word_bits() / 8
     }
 
-    /// Returns the designated general purpose scratch register.
-    fn scratch_reg() -> Reg;
-
-    /// 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: &WasmValType) -> Reg {
-        match ty {
-            WasmValType::I32
-            | WasmValType::I64
-            | WasmValType::Ref(WasmRefType {
-                heap_type: WasmHeapType::Func,
-                ..
-            }) => Self::scratch_reg(),
-            WasmValType::F32 | WasmValType::F64 => Self::float_scratch_reg(),
-            _ => unimplemented!(),
-        }
-    }
+    fn scratch_for(ty: &WasmValType) -> Reg;
 
     /// Returns the pinned register used to hold
     /// the `VMContext`.

winch/codegen/src/isa/aarch64/abi.rs

@@ -1,7 +1,7 @@
 use super::regs;
 use crate::abi::{align_to, ABIOperand, ABIParams, ABIResults, ABISig, ParamsOrReturns, ABI};
 use crate::isa::{reg::Reg, CallingConvention};
-use wasmtime_environ::{WasmHeapType, WasmValType};
+use wasmtime_environ::{WasmHeapType, WasmRefType, WasmValType};
 
 #[derive(Default)]
 pub(crate) struct Aarch64ABI;
@@ -113,12 +113,17 @@ impl ABI for Aarch64ABI {
         })
     }
 
-    fn scratch_reg() -> Reg {
-        regs::scratch()
-    }
-
-    fn float_scratch_reg() -> Reg {
-        regs::float_scratch()
+    fn scratch_for(ty: &WasmValType) -> Reg {
+        match ty {
+            WasmValType::I32
+            | WasmValType::I64
+            | WasmValType::Ref(WasmRefType {
+                heap_type: WasmHeapType::Func,
+                ..
+            }) => regs::scratch(),
+            WasmValType::F32 | WasmValType::F64 => regs::float_scratch(),
+            _ => unimplemented!(),
+        }
     }
 
     fn vmctx_reg() -> Reg {

winch/codegen/src/isa/x64/abi.rs

@@ -3,7 +3,7 @@ use crate::{
     abi::{align_to, ABIOperand, ABIParams, ABIResults, ABISig, ParamsOrReturns, ABI},
     isa::{reg::Reg, CallingConvention},
 };
-use wasmtime_environ::{WasmHeapType, WasmValType};
+use wasmtime_environ::{WasmHeapType, WasmRefType, WasmValType};
 
 /// Helper environment to track argument-register
 /// assignment in x64.
@@ -141,12 +141,17 @@ impl ABI for X64ABI {
         })
     }
 
-    fn scratch_reg() -> Reg {
-        regs::scratch()
-    }
-
-    fn float_scratch_reg() -> Reg {
-        regs::scratch_xmm()
+    fn scratch_for(ty: &WasmValType) -> Reg {
+        match ty {
+            WasmValType::I32
+            | WasmValType::I64
+            | WasmValType::Ref(WasmRefType {
+                heap_type: WasmHeapType::Func,
+                ..
+            }) => regs::scratch(),
+            WasmValType::F32 | WasmValType::F64 => regs::scratch_xmm(),
+            _ => unimplemented!(),
+        }
     }
 
     fn vmctx_reg() -> Reg {

winch/codegen/src/isa/x64/masm.rs

@@ -1041,8 +1041,6 @@ impl MacroAssembler {
 
     /// A common implementation for stack stores.
     fn store_impl(&mut self, src: RegImm, dst: Address, size: OperandSize, flags: MemFlags) {
-        let scratch = <Self as Masm>::ABI::scratch_reg();
-        let float_scratch = <Self as Masm>::ABI::float_scratch_reg();
         match src {
             RegImm::Imm(imm) => match imm {
                 I::I32(v) => self.asm.mov_im(v as i32, &dst, size, flags),
@@ -1051,12 +1049,14 @@ impl MacroAssembler {
                     Err(_) => {
                         // If the immediate doesn't sign extend, use a scratch
                         // register.
+                        let scratch = regs::scratch();
                         self.asm.mov_ir(v, scratch, size);
                         self.asm.mov_rm(scratch, &dst, size, flags);
                     }
                 },
                 I::F32(v) => {
                     let addr = self.asm.add_constant(v.to_le_bytes().as_slice());
+                    let float_scratch = regs::scratch_xmm();
                     // Always trusted, since we are loading the constant from
                     // the constant pool.
                     self.asm
@@ -1065,6 +1065,7 @@ impl MacroAssembler {
                 }
                 I::F64(v) => {
                     let addr = self.asm.add_constant(v.to_le_bytes().as_slice());
+                    let float_scratch = regs::scratch_xmm();
                     // Similar to above, always trusted since we are loading the
                     // constant from the constant pool.
                     self.asm