Repository: bytecodealliance/wasmtime
Pull request #7934 (2024-02-14 02:00:00)

Restructure the MacroAssembler interface to clobbers
To prepare for moving the stack check into the function prologue, rework the macro assembler's interface to clobbers to ensure that the unwind info defines the prologue correctly on windows. Though clobbers are only used in two trampolines, it's a nice interface to use the save/restore points to bracket the code between the prologue and epilogue.

I made an additional change that's not strictly necessary in this PR: reserving the space for clobbers with `reserve_stack`, and then storing values, rather than calling push for each value. I opted to maintain 16-byte stack alignment in this case, despite winch fixing stack alignment up around calls, and this paid off in trampolines that require clobbers as they do one fewer push/pop combination.

<!--
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/opts/arithmetic.isle b/cranelift/codegen/src/opts/arithmetic.isle
index 400da6bab..2a775d93c 100644
--- a/cranelift/codegen/src/opts/arithmetic.isle
+++ b/cranelift/codegen/src/opts/arithmetic.isle
@@ -167,3 +167,13 @@
       (if-let $true (ty_equal half_ty (ty_half_width ty)))
       (if-let $true (u64_eq k (ty_bits_u64 half_ty)))
       (uextend ty (umulhi half_ty x y)))
+
+;; Cranelift's `fcvt_from_{u,s}int` instructions are polymorphic over the input
+;; type so remove any unnecessary `uextend` or `sextend` to give backends
+;; the chance to convert from the smallest integral type to the float. This
+;; can help lowerings on x64 for example which has a less efficient u64-to-float
+;; conversion than other bit widths.
+(rule (simplify (fcvt_from_uint ty (uextend _ val)))
+      (fcvt_from_uint ty val))
+(rule (simplify (fcvt_from_sint ty (sextend _ val)))
+      (fcvt_from_sint ty val))
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 4aac8b6e1..3230151cf 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -91,6 +91,10 @@ where
         // of the stack pointer. This assumes that no writes to the stack occur in `reserve_stack`.
         self.masm.check_stack();
 
+        // We don't have any callee save registers in the winch calling convention, but
+        // `save_clobbers` does some useful work for setting up unwinding state.
+        self.masm.save_clobbers(&[]);
+
         // Once we have emitted the epilogue and reserved stack space for the locals, we push the
         // base control flow block.
         self.control_frames.push(ControlStackFrame::block(
@@ -309,6 +313,7 @@ where
             self.masm.reset_stack_pointer(base);
         }
         debug_assert_eq!(self.context.stack.len(), 0);
+        self.masm.restore_clobbers(&[]);
         self.masm.epilogue(self.context.frame.locals_size);
         Ok(())
     }
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index 9e95ba848..fe919f0f5 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -68,16 +68,6 @@ impl Masm for MacroAssembler {
 
         self.asm
             .mov_rr(stack_pointer, frame_pointer, OperandSize::S64);
-
-        if self.shared_flags.unwind_info() {
-            self.asm.emit_unwind_inst(UnwindInst::DefineNewFrame {
-                offset_upward_to_caller_sp: Self::ABI::arg_base_offset().try_into().unwrap(),
-
-                // Clobbers appear directly after the RET and FP if they're present. As we just
-                // pushed the frame pointer, the offset to the clobbers will be `0`.
-                offset_downward_to_clobbers: 0,
-            })
-        }
     }
 
     fn check_stack(&mut self) {
@@ -98,6 +88,73 @@ impl Masm for MacroAssembler {
         self.asm.trapif(IntCmpKind::GtU, TrapCode::StackOverflow);
     }
 
+    fn save_clobbers(&mut self, clobbers: &[(Reg, OperandSize)]) {
+        let int_bytes: u32 = Self::ABI::word_bytes().try_into().unwrap();
+        let float_bytes = int_bytes * 2;
+
+        // Determine how much space we need for the clobbers
+        let clobbered_size = align_to(
+            clobbers
+                .iter()
+                .fold(0u32, |total, (reg, _)| match reg.class() {
+                    RegClass::Int => total + int_bytes,
+                    RegClass::Float => align_to(total, float_bytes) + float_bytes,
+                    RegClass::Vector => unimplemented!(),
+                }),
+            16,
+        );
+
+        // Emit unwind info.
+        if self.shared_flags.unwind_info() {
+            self.asm.emit_unwind_inst(UnwindInst::DefineNewFrame {
+                offset_upward_to_caller_sp: Self::ABI::arg_base_offset().try_into().unwrap(),
+                offset_downward_to_clobbers: clobbered_size,
+            })
+        }
+
+        self.reserve_stack(clobbered_size);
+
+        let mut off = 0;
+        for &(reg, size) in clobbers {
+            // Align the current offset
+            off = align_to(off, size.bytes());
+
+            // Emit the store
+            let addr = self.address_at_sp(SPOffset::from_u32(off));
+            self.store(RegImm::Reg(reg), addr, size);
+
+            // Emit unwinding info, if necessary
+            if self.shared_flags.unwind_info() {
+                self.asm.emit_unwind_inst(UnwindInst::SaveReg {
+                    clobber_offset: off,
+                    reg: reg.into(),
+                });
+            }
+
+            // Increment the offset
+            off += size.bytes();
+        }
+
+        debug_assert_eq!(align_to(off, float_bytes), clobbered_size);
+    }
+
+    fn restore_clobbers(&mut self, clobbers: &[(Reg, OperandSize)]) {
+        let mut off = 0;
+        for &(reg, size) in clobbers {
+            // Align the current offset
+            off = align_to(off, size.bytes());
+
+            // Emit the load
+            let addr = self.address_at_sp(SPOffset::from_u32(off));
+            self.load(addr, reg, size);
+
+            // Increment the offset
+            off += size.bytes();
+        }
+
+        self.free_stack(align_to(off, 16));
+    }
+
     fn push(&mut self, reg: Reg, size: OperandSize) -> StackSlot {
         let bytes = match (reg.class(), size) {
             (RegClass::Int, OperandSize::S64) => {
@@ -131,19 +188,6 @@ impl Masm for MacroAssembler {
         }
     }
 
-    fn save(&mut self, clobber_offset: u32, reg: Reg, size: OperandSize) -> StackSlot {
-        let slot = self.push(reg, size);
-
-        if self.shared_flags.unwind_info() {
-            self.asm.emit_unwind_inst(UnwindInst::SaveReg {
-                clobber_offset,
-                reg: reg.into(),
-            });
-        }
-
-        slot
-    }
-
     fn reserve_stack(&mut self, bytes: u32) {
         if bytes == 0 {
             return;
@@ -751,9 +795,8 @@ impl Masm for MacroAssembler {
     fn epilogue(&mut self, locals_size: u32) {
         assert_eq!(self.sp_offset, locals_size);
 
-        let rsp = rsp();
         if locals_size > 0 {
-            self.asm.add_ir(locals_size as i32, rsp, OperandSize::S64);
+            self.free_stack(locals_size);
         }
         self.asm.pop_r(rbp());
         self.asm.ret();
diff --git a/winch/codegen/src/isa/x64/mod.rs b/winch/codegen/src/isa/x64/mod.rs
index f63c51d5a..bcb263a5c 100644
--- a/winch/codegen/src/isa/x64/mod.rs
+++ b/winch/codegen/src/isa/x64/mod.rs
@@ -153,7 +153,7 @@ impl TargetIsa for X64 {
         );
         let call_conv = self.wasmtime_call_conv();
 
-        let mut trampoline = Trampoline::new(
+        let trampoline = Trampoline::new(
             &mut masm,
             regs::scratch(),
             regs::argv(),
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index 877d33d9d..428dd112b 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -397,6 +397,23 @@ pub(crate) trait MacroAssembler {
     /// Emit the function prologue.
     fn prologue(&mut self);
 
+    /// Save all the given clobbered registers to the stack. By default this is the same as pushing
+    /// the registers, however it's present in the [`MacroAssembler`] trait to ensure that it's
+    /// possible to add unwind info for register saves in backends.
+    fn save_clobbers(&mut self, clobbers: &[(Reg, OperandSize)]) {
+        for &(reg, size) in clobbers {
+            self.push(reg, size);
+        }
+    }
+
+    /// Restore all clobbered registers, assumed to be passed in the same order as to
+    /// [`save_clobbers`].
+    fn restore_clobbers(&mut self, clobbers: &[(Reg, OperandSize)]) {
+        for &(reg, size) in clobbers.iter().rev() {
+            self.pop(reg, size);
+        }
+    }
+
     /// Emit a stack check.
     fn check_stack(&mut self);
 
@@ -849,11 +866,4 @@ pub(crate) trait MacroAssembler {
             self.free_stack(bytes);
         }
     }
-
-    /// Save the value of this register to the stack. By default this is the same as pushing the
-    /// register, however it's present in the [`MacroAssembler`] trait to ensure that it's possible
-    /// to add unwind info for register saves in backends.
-    fn save(&mut self, _off: u32, src: Reg, size: OperandSize) -> StackSlot {
-        self.push(src, size)
-    }
 }
diff --git a/winch/codegen/src/trampoline.rs b/winch/codegen/src/trampoline.rs
index 515373415..abde67c67 100644
--- a/winch/codegen/src/trampoline.rs
+++ b/winch/codegen/src/trampoline.rs
@@ -88,7 +88,7 @@ where
     }
 
     /// Emit an array-to-wasm trampoline.
-    pub fn emit_array_to_wasm(&mut self, ty: &WasmFuncType, callee_index: FuncIndex) -> Result<()> {
+    pub fn emit_array_to_wasm(mut self, ty: &WasmFuncType, callee_index: FuncIndex) -> Result<()> {
         let array_sig = self.array_sig();
         let wasm_sig = self.wasm_sig(ty);
 
@@ -190,11 +190,7 @@ where
     }
 
     /// Emit a native-to-wasm trampoline.
-    pub fn emit_native_to_wasm(
-        &mut self,
-        ty: &WasmFuncType,
-        callee_index: FuncIndex,
-    ) -> Result<()> {
+    pub fn emit_native_to_wasm(mut self, ty: &WasmFuncType, callee_index: FuncIndex) -> Result<()> {
         let native_sig = self.native_sig(&ty);
         let wasm_sig = self.wasm_sig(&ty);
         let (vmctx, caller_vmctx) = Self::callee_and_caller_vmctx(&native_sig.params)?;
@@ -365,7 +361,7 @@ where
     }
 
     /// Emit a wasm-to-native trampoline.
-    pub fn emit_wasm_to_native(&mut self, ty: &WasmFuncType) -> Result<()> {
+    pub fn emit_wasm_to_native(mut self, ty: &WasmFuncType) -> Result<()> {
         let mut params = self.callee_and_caller_vmctx_types();
         params.extend_from_slice(ty.params());
 
@@ -698,6 +694,7 @@ where
     /// The trampoline's prologue.
     fn prologue(&mut self) {
         self.masm.prologue();
+        self.masm.save_clobbers(&[]);
     }
 
     /// Similar to [Trampoline::prologue], but saves
@@ -705,11 +702,7 @@ where
     fn prologue_with_callee_saved(&mut self) {
         self.masm.prologue();
         // Save any callee-saved registers.
-        let mut off = 0;
-        for (r, s) in &self.callee_saved_regs {
-            let slot = self.masm.save(off, *r, *s);
-            off += slot.size;
-        }
+        self.masm.save_clobbers(&self.callee_saved_regs);
     }
 
     /// Similar to [Trampoline::epilogue], but restores
@@ -718,9 +711,7 @@ where
         // Free the stack space allocated by pushing the trampoline arguments.
         self.masm.free_stack(arg_size);
         // Restore the callee-saved registers.
-        for (r, s) in self.callee_saved_regs.iter().rev() {
-            self.masm.pop(*r, *s);
-        }
+        self.masm.restore_clobbers(&self.callee_saved_regs);
         self.masm.epilogue(0);
     }
 
@@ -728,6 +719,7 @@ where
     fn epilogue(&mut self, arg_size: u32) {
         // Free the stack space allocated by pushing the trampoline arguments.
         self.masm.free_stack(arg_size);
+        self.masm.restore_clobbers(&[]);
         self.masm.epilogue(0);
     }
 }


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/fcvt.clif b/cranelift/filetests/filetests/egraph/fcvt.clif
new file mode 100644
index 000000000..2aacfbe3b
--- /dev/null
+++ b/cranelift/filetests/filetests/egraph/fcvt.clif
@@ -0,0 +1,43 @@
+test optimize
+set opt_level=speed
+target x86_64
+
+function %u32_ext_to_f32(i32) -> f32 {
+block0(v0: i32):
+  v1 = uextend.i64 v0
+  v2 = fcvt_from_uint.f32 v1
+  return v2
+
+  ; check: v3 = fcvt_from_uint.f32 v0
+  ; check: return v3
+}
+
+function %u32_ext_to_f64(i32) -> f64 {
+block0(v0: i32):
+  v1 = uextend.i64 v0
+  v2 = fcvt_from_uint.f64 v1
+  return v2
+
+  ; check: v3 = fcvt_from_uint.f64 v0
+  ; check: return v3
+}
+
+function %i32_ext_to_f32(i32) -> f32 {
+block0(v0: i32):
+  v1 = sextend.i64 v0
+  v2 = fcvt_from_sint.f32 v1
+  return v2
+
+  ; check: v3 = fcvt_from_sint.f32 v0
+  ; check: return v3
+}
+
+function %i32_ext_to_f64(i32) -> f64 {
+block0(v0: i32):
+  v1 = sextend.i64 v0
+  v2 = fcvt_from_sint.f64 v1
+  return v2
+
+  ; check: v3 = fcvt_from_sint.f64 v0
+  ; check: return v3
+}
diff --git a/cranelift/filetests/filetests/egraph/vector.clif b/cranelift/filetests/filetests/egraph/vector.clif
index 4ecbaa83b..a9cb2cbe7 100644
--- a/cranelift/filetests/filetests/egraph/vector.clif
+++ b/cranelift/filetests/filetests/egraph/vector.clif
@@ -18,10 +18,9 @@ block0(v0: i32):
   v2 = splat.i64x2 v1
   v3 = fcvt_from_uint.f64x2 v2
   return v3
-  ; check: v1 = uextend.i64 v0
-  ; check: v4 = fcvt_from_uint.f64 v1
-  ; check: v5 = splat.f64x2 v4
-  ; check: return v5
+  ; check: v5 = fcvt_from_uint.f64 v0
+  ; check: v7 = splat.f64x2 v5
+  ; check: return v7
 }
 
 function %f3(i32) -> f64x2 {