Repository: bytecodealliance/wasmtime
Pull request #7230 (2023-10-12 17:54:32)

Add support for float sqrt operators to Winch
<!--
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
-->
Adding support for `f32.sqrt` and `f64.sqrt` Wasm instructions to Winch. Part of #6528.


Patch:
diff --git a/fuzz/fuzz_targets/differential.rs b/fuzz/fuzz_targets/differential.rs
index d4bdd3bdd..61d31f6a8 100644
--- a/fuzz/fuzz_targets/differential.rs
+++ b/fuzz/fuzz_targets/differential.rs
@@ -375,6 +375,8 @@ fn winch_supports_module(module: &[u8]) -> bool {
                         | F64Abs { .. }
                         | F32Neg { .. }
                         | F64Neg { .. }
+                        | F32Sqrt { .. }
+                        | F64Sqrt { .. }
                         | CallIndirect { .. }
                         | ElemDrop { .. }
                         | TableCopy { .. }
diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
index f6032a946..e6fb506c1 100644
--- a/winch/codegen/src/isa/aarch64/masm.rs
+++ b/winch/codegen/src/isa/aarch64/masm.rs
@@ -262,6 +262,10 @@ impl Masm for MacroAssembler {
         todo!();
     }
 
+    fn float_sqrt(&mut self, _dst: Reg, _src: Reg, _size: OperandSize) {
+        todo!()
+    }
+
     fn and(&mut self, _dst: Reg, _lhs: Reg, _rhs: RegImm, _size: OperandSize) {
         todo!()
     }
diff --git a/winch/codegen/src/isa/x64/asm.rs b/winch/codegen/src/isa/x64/asm.rs
index 6b94c5ad7..afd1a2a71 100644
--- a/winch/codegen/src/isa/x64/asm.rs
+++ b/winch/codegen/src/isa/x64/asm.rs
@@ -845,6 +845,20 @@ impl Assembler {
         })
     }
 
+    pub fn sqrt(&mut self, src: Reg, dst: Reg, size: OperandSize) {
+        let op = match size {
+            OperandSize::S32 => SseOpcode::Sqrtss,
+            OperandSize::S64 => SseOpcode::Sqrtsd,
+            OperandSize::S128 => unreachable!(),
+        };
+
+        self.emit(Inst::XmmUnaryRmR {
+            op,
+            src: Xmm::from(src).into(),
+            dst: dst.into(),
+        })
+    }
+
     /// Emit a call to an unknown location through a register.
     pub fn call_with_reg(&mut self, callee: Reg) {
         self.emit(Inst::CallUnknown {
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index 25264de83..182d8f902 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -404,6 +404,10 @@ impl Masm for MacroAssembler {
         }
     }
 
+    fn float_sqrt(&mut self, dst: Reg, src: Reg, size: OperandSize) {
+        self.asm.sqrt(src, dst, size);
+    }
+
     fn and(&mut self, dst: Reg, lhs: Reg, rhs: RegImm, size: OperandSize) {
         Self::ensure_two_argument_form(&dst, &lhs);
         match (rhs, dst) {
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index ab032cfb7..3c597e709 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -383,6 +383,9 @@ pub(crate) trait MacroAssembler {
     /// Perform a floating point floor operation.
     fn float_round(&mut self, mode: RoundingMode, dst: Reg, src: RegImm, size: OperandSize);
 
+    /// Perform a floating point square root operation.
+    fn float_sqrt(&mut self, dst: Reg, src: Reg, size: OperandSize);
+
     /// Perform logical and operation.
     fn and(&mut self, dst: Reg, lhs: Reg, rhs: RegImm, size: OperandSize);
 
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 40a290666..53143f4b4 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -56,6 +56,8 @@ macro_rules! def_unsupported {
     (emit F64Nearest $($rest:tt)*) => {};
     (emit F32Trunc $($rest:tt)*) => {};
     (emit F64Trunc $($rest:tt)*) => {};
+    (emit F32Sqrt $($rest:tt)*) => {};
+    (emit F64Sqrt $($rest:tt)*) => {};
     (emit I32Add $($rest:tt)*) => {};
     (emit I64Add $($rest:tt)*) => {};
     (emit I32Sub $($rest:tt)*) => {};
@@ -252,6 +254,20 @@ where
             });
     }
 
+    fn visit_f32_sqrt(&mut self) {
+        self.context
+            .unop(self.masm, OperandSize::S32, &mut |masm, reg, size| {
+                masm.float_sqrt(reg, reg, size);
+            });
+    }
+
+    fn visit_f64_sqrt(&mut self) {
+        self.context
+            .unop(self.masm, OperandSize::S64, &mut |masm, reg, size| {
+                masm.float_sqrt(reg, reg, size);
+            });
+    }
+
     fn visit_i32_add(&mut self) {
         self.context.i32_binop(self.masm, |masm, dst, src, size| {
             masm.add(dst, dst, src, size);


Test patch:
diff --git a/tests/misc_testsuite/winch/f32.wast b/tests/misc_testsuite/winch/f32.wast
index 5d8497461..96ab0d8d2 100644
--- a/tests/misc_testsuite/winch/f32.wast
+++ b/tests/misc_testsuite/winch/f32.wast
@@ -306,12 +306,33 @@
 
 
 (module
+  (func (export "sqrt") (param $x f32) (result f32) (f32.sqrt (local.get $x)))
   (func (export "ceil") (param $x f32) (result f32) (f32.ceil (local.get $x)))
   (func (export "floor") (param $x f32) (result f32) (f32.floor (local.get $x)))
   (func (export "trunc") (param $x f32) (result f32) (f32.trunc (local.get $x)))
   (func (export "nearest") (param $x f32) (result f32) (f32.nearest (local.get $x)))
 )
 
+(assert_return (invoke "sqrt" (f32.const -0x0p+0)) (f32.const -0x0p+0))
+(assert_return (invoke "sqrt" (f32.const 0x0p+0)) (f32.const 0x0p+0))
+(assert_return (invoke "sqrt" (f32.const -0x1p-149)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1p-149)) (f32.const 0x1.6a09e6p-75))
+(assert_return (invoke "sqrt" (f32.const -0x1p-126)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1p-126)) (f32.const 0x1p-63))
+(assert_return (invoke "sqrt" (f32.const -0x1p-1)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1p-1)) (f32.const 0x1.6a09e6p-1))
+(assert_return (invoke "sqrt" (f32.const -0x1p+0)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1p+0)) (f32.const 0x1p+0))
+(assert_return (invoke "sqrt" (f32.const -0x1.921fb6p+2)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1.921fb6p+2)) (f32.const 0x1.40d932p+1))
+(assert_return (invoke "sqrt" (f32.const -0x1.fffffep+127)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const 0x1.fffffep+127)) (f32.const 0x1.fffffep+63))
+(assert_return (invoke "sqrt" (f32.const -inf)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const inf)) (f32.const inf))
+(assert_return (invoke "sqrt" (f32.const -nan)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const -nan:0x200000)) (f32.const nan:arithmetic))
+(assert_return (invoke "sqrt" (f32.const nan)) (f32.const nan:canonical))
+(assert_return (invoke "sqrt" (f32.const nan:0x200000)) (f32.const nan:arithmetic))
 (assert_return (invoke "floor" (f32.const -0x0p+0)) (f32.const -0x0p+0))
 (assert_return (invoke "floor" (f32.const 0x0p+0)) (f32.const 0x0p+0))
 (assert_return (invoke "floor" (f32.const -0x1p-149)) (f32.const -0x1p+0))
diff --git a/tests/misc_testsuite/winch/f64.wast b/tests/misc_testsuite/winch/f64.wast
index 7acb99be6..9792f6d45 100644
--- a/tests/misc_testsuite/winch/f64.wast
+++ b/tests/misc_testsuite/winch/f64.wast
@@ -329,12 +329,33 @@
 (assert_return (invoke "type-first-f64") (f64.const 1.64))
 
 (module
+  (func (export "sqrt") (param $x f64) (result f64) (f64.sqrt (local.get $x)))
   (func (export "ceil") (param $x f64) (result f64) (f64.ceil (local.get $x)))
   (func (export "floor") (param $x f64) (result f64) (f64.floor (local.get $x)))
   (func (export "trunc") (param $x f64) (result f64) (f64.trunc (local.get $x)))
   (func (export "nearest") (param $x f64) (result f64) (f64.nearest (local.get $x)))
 )
 
+(assert_return (invoke "sqrt" (f64.const -0x0p+0)) (f64.const -0x0p+0))
+(assert_return (invoke "sqrt" (f64.const 0x0p+0)) (f64.const 0x0p+0))
+(assert_return (invoke "sqrt" (f64.const -0x0.0000000000001p-1022)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x0.0000000000001p-1022)) (f64.const 0x1p-537))
+(assert_return (invoke "sqrt" (f64.const -0x1p-1022)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x1p-1022)) (f64.const 0x1p-511))
+(assert_return (invoke "sqrt" (f64.const -0x1p-1)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x1p-1)) (f64.const 0x1.6a09e667f3bcdp-1))
+(assert_return (invoke "sqrt" (f64.const -0x1p+0)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x1p+0)) (f64.const 0x1p+0))
+(assert_return (invoke "sqrt" (f64.const -0x1.921fb54442d18p+2)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x1.921fb54442d18p+2)) (f64.const 0x1.40d931ff62705p+1))
+(assert_return (invoke "sqrt" (f64.const -0x1.fffffffffffffp+1023)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const 0x1.fffffffffffffp+1023)) (f64.const 0x1.fffffffffffffp+511))
+(assert_return (invoke "sqrt" (f64.const -inf)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const inf)) (f64.const inf))
+(assert_return (invoke "sqrt" (f64.const -nan)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const -nan:0x4000000000000)) (f64.const nan:arithmetic))
+(assert_return (invoke "sqrt" (f64.const nan)) (f64.const nan:canonical))
+(assert_return (invoke "sqrt" (f64.const nan:0x4000000000000)) (f64.const nan:arithmetic))
 (assert_return (invoke "floor" (f64.const -0x0p+0)) (f64.const -0x0p+0))
 (assert_return (invoke "floor" (f64.const 0x0p+0)) (f64.const 0x0p+0))
 (assert_return (invoke "floor" (f64.const -0x0.0000000000001p-1022)) (f64.const -0x1p+0))
diff --git a/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_const.wat b/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_const.wat
new file mode 100644
index 000000000..0cc61434b
--- /dev/null
+++ b/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_const.wat
@@ -0,0 +1,21 @@
+;;! target = "x86_64"
+
+(module
+    (func (result f32)
+        (f32.const 1.32)
+        (f32.sqrt)
+    )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 f30f10050c000000     	movss	xmm0, dword ptr [rip + 0xc]
+;;   14:	 f30f51c0             	sqrtss	xmm0, xmm0
+;;   18:	 4883c408             	add	rsp, 8
+;;   1c:	 5d                   	pop	rbp
+;;   1d:	 c3                   	ret	
+;;   1e:	 0000                 	add	byte ptr [rax], al
+;;   20:	 c3                   	ret	
+;;   21:	 f5                   	cmc	
+;;   22:	 a83f                 	test	al, 0x3f
diff --git a/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_param.wat b/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_param.wat
new file mode 100644
index 000000000..6e096d088
--- /dev/null
+++ b/winch/filetests/filetests/x64/f32_sqrt/f32_sqrt_param.wat
@@ -0,0 +1,18 @@
+;;! target = "x86_64"
+
+(module
+    (func (param f32) (result f32)
+        (local.get 0)
+        (f32.sqrt)
+    )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec10             	sub	rsp, 0x10
+;;    8:	 f30f1144240c         	movss	dword ptr [rsp + 0xc], xmm0
+;;    e:	 4c89742404           	mov	qword ptr [rsp + 4], r14
+;;   13:	 f30f1044240c         	movss	xmm0, dword ptr [rsp + 0xc]
+;;   19:	 f30f51c0             	sqrtss	xmm0, xmm0
+;;   1d:	 4883c410             	add	rsp, 0x10
+;;   21:	 5d                   	pop	rbp
+;;   22:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_const.wat b/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_const.wat
new file mode 100644
index 000000000..0b70b0885
--- /dev/null
+++ b/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_const.wat
@@ -0,0 +1,18 @@
+;;! target = "x86_64"
+
+(module
+    (func (result f64)
+        (f64.const 1.32)
+        (f64.sqrt)
+    )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 f20f10050c000000     	movsd	xmm0, qword ptr [rip + 0xc]
+;;   14:	 f20f51c0             	sqrtsd	xmm0, xmm0
+;;   18:	 4883c408             	add	rsp, 8
+;;   1c:	 5d                   	pop	rbp
+;;   1d:	 c3                   	ret	
+;;   1e:	 0000                 	add	byte ptr [rax], al
diff --git a/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_param.wat b/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_param.wat
new file mode 100644
index 000000000..9b06a9637
--- /dev/null
+++ b/winch/filetests/filetests/x64/f64_sqrt/f64_sqrt_param.wat
@@ -0,0 +1,18 @@
+;;! target = "x86_64"
+
+(module
+    (func (param f64) (result f64)
+        (local.get 0)
+        (f64.sqrt)
+    )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec10             	sub	rsp, 0x10
+;;    8:	 f20f11442408         	movsd	qword ptr [rsp + 8], xmm0
+;;    e:	 4c893424             	mov	qword ptr [rsp], r14
+;;   12:	 f20f10442408         	movsd	xmm0, qword ptr [rsp + 8]
+;;   18:	 f20f51c0             	sqrtsd	xmm0, xmm0
+;;   1c:	 4883c410             	add	rsp, 0x10
+;;   20:	 5d                   	pop	rbp
+;;   21:	 c3                   	ret