winch: respect the enable_nan_canonicalization setting

tests/disas/winch/x64/f32_add/nan_canon.wat

@@ -0,0 +1,40 @@
+;;! target = "x86_64"
+;;! test = "winch"
+;;! flags = "-Wnan-canonicalization"
+
+(module
+    (func (param f32 f32) (result f32)
+        local.get 0
+        local.get 1
+        f32.add
+    )
+)
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x18(%r11), %r11
+;;       addq    $0x20, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x69
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x20, %rsp
+;;       movq    %rdi, 0x18(%rsp)
+;;       movq    %rsi, 0x10(%rsp)
+;;       movss   %xmm0, 0xc(%rsp)
+;;       movss   %xmm1, 8(%rsp)
+;;       movss   8(%rsp), %xmm0
+;;       movss   0xc(%rsp), %xmm1
+;;       addss   %xmm0, %xmm1
+;;       ucomiss %xmm1, %xmm1
+;;       jnp     0x5d
+;;   55: movss   0x13(%rip), %xmm1
+;;       movaps  %xmm1, %xmm0
+;;       addq    $0x20, %rsp
+;;       popq    %rbp
+;;       retq
+;;   69: ud2
+;;   6b: addb    %al, (%rax)
+;;   6d: addb    %al, (%rax)
+;;   6f: addb    %al, (%rax)
+;;   71: addb    %al, %al

tests/disas/winch/x64/f64_div/nan_canon.wat

@@ -0,0 +1,43 @@
+;;! target = "x86_64"
+;;! test = "winch"
+;;! flags = "-Wnan-canonicalization"
+
+(module
+    (func (param f64 f64) (result f64)
+        local.get 0
+        local.get 1
+        f64.div
+    )
+)
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r11
+;;       movq    0x18(%r11), %r11
+;;       addq    $0x20, %r11
+;;       cmpq    %rsp, %r11
+;;       ja      0x68
+;;   1c: movq    %rdi, %r14
+;;       subq    $0x20, %rsp
+;;       movq    %rdi, 0x18(%rsp)
+;;       movq    %rsi, 0x10(%rsp)
+;;       movsd   %xmm0, 8(%rsp)
+;;       movsd   %xmm1, (%rsp)
+;;       movsd   (%rsp), %xmm0
+;;       movsd   8(%rsp), %xmm1
+;;       divsd   %xmm0, %xmm1
+;;       ucomisd %xmm1, %xmm1
+;;       jnp     0x5c
+;;   54: movsd   0x14(%rip), %xmm1
+;;       movaps  %xmm1, %xmm0
+;;       addq    $0x20, %rsp
+;;       popq    %rbp
+;;       retq
+;;   68: ud2
+;;   6a: addb    %al, (%rax)
+;;   6c: addb    %al, (%rax)
+;;   6e: addb    %al, (%rax)
+;;   70: addb    %al, (%rax)
+;;   72: addb    %al, (%rax)
+;;   74: addb    %al, (%rax)
+;;   76: clc

tests/misc_testsuite/canonicalize-nan-scalar.wast

@@ -0,0 +1,153 @@
+;;! nan_canonicalization = true
+
+;; Scalar counterpart to simd/canonicalize-nan.wast.
+
+(module
+  (func (export "f32.add") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.add)
+  (func (export "f32.sub") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.sub)
+  (func (export "f32.mul") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.mul)
+  (func (export "f32.div") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.div)
+  (func (export "f32.min") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.min)
+  (func (export "f32.max") (param f32 f32) (result f32)
+    local.get 0
+    local.get 1
+    f32.max)
+  (func (export "f32.sqrt") (param f32) (result f32)
+    local.get 0
+    f32.sqrt)
+  (func (export "f32.ceil") (param f32) (result f32)
+    local.get 0
+    f32.ceil)
+  (func (export "f32.floor") (param f32) (result f32)
+    local.get 0
+    f32.floor)
+  (func (export "f32.trunc") (param f32) (result f32)
+    local.get 0
+    f32.trunc)
+  (func (export "f32.nearest") (param f32) (result f32)
+    local.get 0
+    f32.nearest)
+
+  (func (export "f64.add") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.add)
+  (func (export "f64.sub") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.sub)
+  (func (export "f64.mul") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.mul)
+  (func (export "f64.div") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.div)
+  (func (export "f64.min") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.min)
+  (func (export "f64.max") (param f64 f64) (result f64)
+    local.get 0
+    local.get 1
+    f64.max)
+  (func (export "f64.sqrt") (param f64) (result f64)
+    local.get 0
+    f64.sqrt)
+  (func (export "f64.ceil") (param f64) (result f64)
+    local.get 0
+    f64.ceil)
+  (func (export "f64.floor") (param f64) (result f64)
+    local.get 0
+    f64.floor)
+  (func (export "f64.trunc") (param f64) (result f64)
+    local.get 0
+    f64.trunc)
+  (func (export "f64.nearest") (param f64) (result f64)
+    local.get 0
+    f64.nearest)
+
+  (func (export "reinterpret-and-demote") (param i64) (result i32)
+    local.get 0
+    f64.reinterpret_i64
+    f32.demote_f64
+    i32.reinterpret_f32)
+  (func (export "reinterpret-and-promote") (param i32) (result i64)
+    local.get 0
+    f32.reinterpret_i32
+    f64.promote_f32
+    i64.reinterpret_f64)
+
+  ;; Expose raw bits of 0/0 to verify exact canonical NaN bit patterns.
+  (func (export "f32.div-nan-bits") (result i32)
+    f32.const 0
+    f32.const 0
+    f32.div
+    i32.reinterpret_f32)
+  (func (export "f64.div-nan-bits") (result i64)
+    f64.const 0
+    f64.const 0
+    f64.div
+    i64.reinterpret_f64)
+)
+
+;; Exact bit patterns: canonical f32 NaN = 0x7fc00000, f64 = 0x7ff8000000000000
+(assert_return (invoke "f32.div-nan-bits") (i32.const 0x7fc00000))
+(assert_return (invoke "f64.div-nan-bits") (i64.const 0x7ff8000000000000))
+
+;; NaN-producing operations
+(assert_return (invoke "f32.div" (f32.const 0) (f32.const 0)) (f32.const nan:0x400000))
+(assert_return (invoke "f64.div" (f64.const 0) (f64.const 0)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f32.sqrt" (f32.const -1)) (f32.const nan:0x400000))
+(assert_return (invoke "f64.sqrt" (f64.const -1)) (f64.const nan:0x8000000000000))
+
+;; NaN propagation through f32 arithmetic
+(assert_return (invoke "f32.add" (f32.const nan) (f32.const 1)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.sub" (f32.const nan) (f32.const 1)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.mul" (f32.const nan) (f32.const 1)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.min" (f32.const nan) (f32.const 1)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.max" (f32.const nan) (f32.const 1)) (f32.const nan:0x400000))
+
+;; NaN propagation through f64 arithmetic
+(assert_return (invoke "f64.add" (f64.const nan) (f64.const 1)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.sub" (f64.const nan) (f64.const 1)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.mul" (f64.const nan) (f64.const 1)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.min" (f64.const nan) (f64.const 1)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.max" (f64.const nan) (f64.const 1)) (f64.const nan:0x8000000000000))
+
+;; Rounding NaN (f32)
+(assert_return (invoke "f32.ceil" (f32.const nan)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.floor" (f32.const nan)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.trunc" (f32.const nan)) (f32.const nan:0x400000))
+(assert_return (invoke "f32.nearest" (f32.const nan)) (f32.const nan:0x400000))
+
+;; Rounding NaN (f64)
+(assert_return (invoke "f64.ceil" (f64.const nan)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.floor" (f64.const nan)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.trunc" (f64.const nan)) (f64.const nan:0x8000000000000))
+(assert_return (invoke "f64.nearest" (f64.const nan)) (f64.const nan:0x8000000000000))
+
+;; Demote/promote with non-canonical NaN bit patterns
+(assert_return (invoke "reinterpret-and-demote" (i64.const 0xfffefdfccccdcecf)) (i32.const 0x7fc00000))
+(assert_return (invoke "reinterpret-and-promote" (i32.const 0xfffefdfc)) (i64.const 0x7ff8000000000000))
+
+;; Normal values pass through unchanged
+(assert_return (invoke "f32.add" (f32.const 1) (f32.const 2)) (f32.const 3))
+(assert_return (invoke "f64.div" (f64.const 10) (f64.const 2)) (f64.const 5))
+(assert_return (invoke "f32.sqrt" (f32.const 4)) (f32.const 2))

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

@@ -59,6 +59,8 @@ pub(crate) struct MacroAssembler {
     ptr_size: OperandSize,
     /// Scratch register scope.
     scratch_scope: RegAlloc,
+    /// Shared flags.
+    shared_flags: settings::Flags,
 }
 
 impl MacroAssembler {
@@ -71,10 +73,11 @@ impl MacroAssembler {
         Ok(Self {
             sp_max: 0,
             stack_max_use_add: None,
-            asm: Assembler::new(shared_flags, isa_flags),
+            asm: Assembler::new(shared_flags.clone(), isa_flags),
             sp_offset: 0u32,
             ptr_size: ptr_type_from_ptr_size(ptr_size.size()).try_into()?,
             scratch_scope: RegAlloc::from(scratch_gpr_bitset(), scratch_fpr_bitset()),
+            shared_flags,
         })
     }
 
@@ -713,6 +716,35 @@ impl Masm for MacroAssembler {
         Ok(())
     }
 
+    fn canonicalize_nan(&mut self, reg: WritableReg, size: OperandSize) -> Result<()> {
+        if !self.shared_flags.enable_nan_canonicalization() {
+            return Ok(());
+        }
+
+        let done_label = self.asm.buffer_mut().get_label();
+
+        self.asm.fcmp(reg.to_reg(), reg.to_reg(), size);
+        self.asm.jmp_if(Cond::Vc, done_label);
+
+        let canonical_nan: &[u8] = match size {
+            OperandSize::S32 => &0x7FC00000u32.to_le_bytes(),
+            OperandSize::S64 => &0x7FF8000000000000u64.to_le_bytes(),
+            _ => bail!(CodeGenError::unexpected_operand_size()),
+        };
+        let constant = self.asm.add_constant(canonical_nan);
+        self.asm.uload(
+            inst::AMode::Const { addr: constant },
+            reg,
+            size,
+            TRUSTED_FLAGS,
+        );
+
+        self.asm
+            .buffer_mut()
+            .bind_label(done_label, &mut Default::default());
+        Ok(())
+    }
+
     fn and(&mut self, dst: WritableReg, lhs: Reg, rhs: RegImm, size: OperandSize) -> Result<()> {
         match (rhs, lhs, dst) {
             (RegImm::Imm(v), rn, rd) => {

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

@@ -686,6 +686,29 @@ impl Masm for MacroAssembler {
         Ok(())
     }
 
+    fn canonicalize_nan(&mut self, reg: WritableReg, size: OperandSize) -> Result<()> {
+        if !self.shared_flags.enable_nan_canonicalization() {
+            return Ok(());
+        }
+
+        let done_label = self.asm.buffer_mut().get_label();
+
+        self.asm.ucomis(reg.to_reg(), reg.to_reg(), size);
+        self.asm.jmp_if(CC::NP, done_label);
+
+        let canonical_nan: &[u8] = match size {
+            OperandSize::S32 => &0x7FC00000u32.to_le_bytes(),
+            OperandSize::S64 => &0x7FF8000000000000u64.to_le_bytes(),
+            _ => bail!(CodeGenError::unexpected_operand_size()),
+        };
+        self.asm.load_fp_const(reg, canonical_nan, size);
+
+        self.asm
+            .buffer_mut()
+            .bind_label(done_label, &mut Default::default());
+        Ok(())
+    }
+
     fn and(&mut self, dst: WritableReg, lhs: Reg, rhs: RegImm, size: OperandSize) -> Result<()> {
         Self::ensure_two_argument_form(&dst.to_reg(), &lhs)?;
         match (rhs, dst) {

winch/codegen/src/masm.rs

@@ -1692,6 +1692,9 @@ pub(crate) trait MacroAssembler {
     /// Perform a floating point square root operation.
     fn float_sqrt(&mut self, dst: WritableReg, src: Reg, size: OperandSize) -> Result<()>;
 
+    /// Canonicalize NaN values in `reg` if the setting is enabled.
+    fn canonicalize_nan(&mut self, reg: WritableReg, size: OperandSize) -> Result<()>;
+
     /// Perform logical and operation.
     fn and(&mut self, dst: WritableReg, lhs: Reg, rhs: RegImm, size: OperandSize) -> Result<()>;