winch: Add support for <i32|i64>.rem_* WebAssembly instructions (#5823)

This commit adds support for i32 and i64 remainder instructions for
x64.

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

@@ -7,7 +7,7 @@ use crate::{
     abi::local::LocalSlot,
     codegen::CodeGenContext,
     isa::reg::Reg,
-    masm::{DivKind, MacroAssembler as Masm, OperandSize, RegImm},
+    masm::{DivKind, MacroAssembler as Masm, OperandSize, RegImm, RemKind},
 };
 use cranelift_codegen::{settings, Final, MachBufferFinalized};
 
@@ -156,6 +156,10 @@ impl Masm for MacroAssembler {
         todo!()
     }
 
+    fn rem(&mut self, _context: &mut CodeGenContext, _kind: RemKind, _size: OperandSize) {
+        todo!()
+    }
+
     fn zero(&mut self, reg: Reg) {
         self.asm.load_constant(0, reg);
     }

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

@@ -2,7 +2,7 @@
 
 use crate::{
     isa::reg::Reg,
-    masm::{DivKind, OperandSize},
+    masm::{DivKind, OperandSize, RemKind},
 };
 use cranelift_codegen::{
     isa::x64::{
@@ -73,6 +73,15 @@ impl From<DivKind> for DivOrRemKind {
     }
 }
 
+impl From<RemKind> for DivOrRemKind {
+    fn from(kind: RemKind) -> Self {
+        match kind {
+            RemKind::Signed => DivOrRemKind::SignedRem,
+            RemKind::Unsigned => DivOrRemKind::UnsignedRem,
+        }
+    }
+}
+
 /// Low level assembler implementation for x64.
 pub(crate) struct Assembler {
     /// The machine instruction buffer.
@@ -299,6 +308,26 @@ impl Assembler {
         });
     }
 
+    /// Signed/unsigned remainder.
+    ///
+    /// Emits a sequence of instructions to ensure the correctness of the
+    /// division invariants and ultimately calculate the remainder.
+    /// This function assumes that the
+    /// caller has correctly allocated the dividend as `(rdx:rax)` and
+    /// accounted for the remainder to be stored in `rdx`.
+    pub fn rem(&mut self, divisor: Reg, dst: (Reg, Reg), kind: RemKind, size: OperandSize) {
+        self.emit(Inst::CheckedDivOrRemSeq {
+            kind: kind.into(),
+            size: size.into(),
+            divisor: divisor.into(),
+            dividend_lo: dst.0.into(),
+            dividend_hi: dst.1.into(),
+            dst_quotient: dst.0.into(),
+            dst_remainder: dst.1.into(),
+            tmp: None,
+        });
+    }
+
     /// Multiply immediate and register.
     pub fn mul_ir(&mut self, imm: i32, dst: Reg, size: OperandSize) {
         let imm = RegMemImm::imm(imm as u32);

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

@@ -4,7 +4,7 @@ use super::{
     regs::{self, rbp, rsp},
 };
 use crate::isa::reg::Reg;
-use crate::masm::{DivKind, MacroAssembler as Masm, OperandSize, RegImm};
+use crate::masm::{DivKind, MacroAssembler as Masm, OperandSize, RegImm, RemKind};
 use crate::{abi::LocalSlot, codegen::CodeGenContext, stack::Val};
 use cranelift_codegen::{isa::x64::settings as x64_settings, settings, Final, MachBufferFinalized};
 
@@ -174,6 +174,28 @@ impl Masm for MacroAssembler {
         context.stack.push(Val::reg(rax));
     }
 
+    fn rem(&mut self, context: &mut CodeGenContext, kind: RemKind, size: OperandSize) {
+        // Allocate rdx:rax.
+        let rdx = context.gpr(regs::rdx(), self);
+        let rax = context.gpr(regs::rax(), self);
+
+        // Allocate the divisor, which can be any gpr.
+        let divisor = context.pop_to_reg(self, size);
+
+        // Mark rax as allocatable.
+        context.regalloc.free_gpr(rax);
+        // Move the top value to rax.
+        let rax = context.pop_to_named_reg(self, rax, size);
+        self.asm.rem(divisor, (rax, rdx), kind, size);
+
+        // Free the divisor and rax.
+        context.free_gpr(divisor);
+        context.free_gpr(rax);
+
+        // Push the remainder.
+        context.stack.push(Val::reg(rdx));
+    }
+
     fn epilogue(&mut self, locals_size: u32) {
         assert!(self.sp_offset == locals_size);
 

winch/codegen/src/masm.rs

@@ -13,6 +13,14 @@ pub(crate) enum DivKind {
     Unsigned,
 }
 
+/// Remainder kind.
+pub(crate) enum RemKind {
+    /// Signed remainder.
+    Signed,
+    /// Unsigned remainder.
+    Unsigned,
+}
+
 /// Operand size, in bits.
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum OperandSize {
@@ -116,6 +124,9 @@ pub(crate) trait MacroAssembler {
     /// functions.
     fn div(&mut self, context: &mut CodeGenContext, kind: DivKind, size: OperandSize);
 
+    /// Calculate remainder.
+    fn rem(&mut self, context: &mut CodeGenContext, kind: RemKind, size: OperandSize);
+
     /// Push the register to the stack, returning the offset.
     fn push(&mut self, src: Reg) -> u32;
 

winch/codegen/src/visitor.rs

@@ -5,7 +5,7 @@
 //! machine code emitter.
 
 use crate::codegen::CodeGen;
-use crate::masm::{DivKind, MacroAssembler, OperandSize, RegImm};
+use crate::masm::{DivKind, MacroAssembler, OperandSize, RegImm, RemKind};
 use crate::stack::Val;
 use wasmparser::ValType;
 use wasmparser::VisitOperator;
@@ -41,6 +41,10 @@ macro_rules! def_unsupported {
     (emit I32DivU $($rest:tt)*) => {};
     (emit I64DivS $($rest:tt)*) => {};
     (emit I64DivU $($rest:tt)*) => {};
+    (emit I64RemU $($rest:tt)*) => {};
+    (emit I64RemS $($rest:tt)*) => {};
+    (emit I32RemU $($rest:tt)*) => {};
+    (emit I32RemS $($rest:tt)*) => {};
     (emit I64Mul $($rest:tt)*) => {};
     (emit I64Sub $($rest:tt)*) => {};
     (emit LocalGet $($rest:tt)*) => {};
@@ -134,6 +138,34 @@ where
         self.masm.div(&mut self.context, Unsigned, S64);
     }
 
+    fn visit_i32_rem_s(&mut self) {
+        use OperandSize::*;
+        use RemKind::*;
+
+        self.masm.rem(&mut self.context, Signed, S32);
+    }
+
+    fn visit_i32_rem_u(&mut self) {
+        use OperandSize::*;
+        use RemKind::*;
+
+        self.masm.rem(&mut self.context, Unsigned, S32);
+    }
+
+    fn visit_i64_rem_s(&mut self) {
+        use OperandSize::*;
+        use RemKind::*;
+
+        self.masm.rem(&mut self.context, Signed, S64);
+    }
+
+    fn visit_i64_rem_u(&mut self) {
+        use OperandSize::*;
+        use RemKind::*;
+
+        self.masm.rem(&mut self.context, Unsigned, S64);
+    }
+
     fn visit_end(&mut self) {}
 
     fn visit_local_get(&mut self, index: u32) {