winch: Add support for the <i32|i64>.div_* instructions (#5807)

* Refactor the structure and responsibilities of `CodeGenContext`

This commit refactors how the `CodeGenContext` is used throughout the code
generation process, making it easier to pass it around when more flexibility is
desired in the MacroAssembler to perform the lowering of certain instructions.

As of this change, the responsibility of the `CodeGenContext` is to provide an
interface for operations that require an orchestration between the register
allocator, the value stack and function's frame. The MacroAssembler is removed
from the CodeGenContext as is passed as a dependency where needed, effectly
using it as an independent code generation interface only.

By giving more responsibilities to the `CodeGenContext` we can clearly separate
the concerns of the register allocator, which previously did more than it
should (e.g. popping values and spilling).

This change ultimately allows passing in the `CodeGenContext` to the
`MacroAssembler` when a given instruction cannot be generically described
through a common interface. Allowing each implementation to decide the best way
to lower a particular instruction.

* winch: Add support for the WebAssembly `<i32|i64>.div_*` instructions

Given that some architectures have very specific requirements on how to handle
division, this change uses `CodeGenContext` as a dependency to the `div`
MacroAssembler instruction to ensure that each implementation can decide on how to lower the
division. This approach also allows -- in architectures where division can be
expressed as an ordinary binary operation -- to rely on the
`CodeGenContext::i32_binop` or `CodeGenContext::i64_binop` helpers.

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

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

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

@@ -82,8 +82,8 @@ impl TargetIsa for Aarch64 {
         let frame = Frame::new(&abi_sig, &mut body, &mut validator, &abi)?;
         // TODO: Add floating point bitmask
         let regalloc = RegAlloc::new(RegSet::new(ALL_GPR, 0), scratch());
-        let codegen_context = CodeGenContext::new(&mut masm, stack, &frame);
-        let mut codegen = CodeGen::new::<abi::Aarch64ABI>(codegen_context, abi_sig, regalloc);
+        let codegen_context = CodeGenContext::new(regalloc, stack, &frame);
+        let mut codegen = CodeGen::new::<abi::Aarch64ABI>(&mut masm, codegen_context, abi_sig);
 
         codegen.emit(&mut body, validator)?;
         Ok(masm.finalize())

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

@@ -1,11 +1,14 @@
 //! Assembler library implementation for x64.
 
-use crate::{isa::reg::Reg, masm::OperandSize};
+use crate::{
+    isa::reg::Reg,
+    masm::{DivKind, OperandSize},
+};
 use cranelift_codegen::{
     isa::x64::{
         args::{
-            self, AluRmiROpcode, Amode, ExtMode, FromWritableReg, Gpr, GprMem, GprMemImm, RegMem,
-            RegMemImm, SyntheticAmode, WritableGpr,
+            self, AluRmiROpcode, Amode, DivOrRemKind, ExtMode, FromWritableReg, Gpr, GprMem,
+            GprMemImm, RegMem, RegMemImm, SyntheticAmode, WritableGpr,
         },
         settings as x64_settings, EmitInfo, EmitState, Inst,
     },
@@ -61,6 +64,15 @@ impl From<OperandSize> for args::OperandSize {
     }
 }
 
+impl From<DivKind> for DivOrRemKind {
+    fn from(kind: DivKind) -> Self {
+        match kind {
+            DivKind::Signed => DivOrRemKind::SignedDiv,
+            DivKind::Unsigned => DivOrRemKind::UnsignedDiv,
+        }
+    }
+}
+
 /// Low level assembler implementation for x64.
 pub(crate) struct Assembler {
     /// The machine instruction buffer.
@@ -262,6 +274,31 @@ impl Assembler {
         }
     }
 
+    /// Signed/unsigned division.
+    ///
+    /// Emits a sequence of instructions to ensure the correctness of
+    /// the division invariants.  This function assumes that the
+    /// caller has correctly allocated the dividend as `(rdx:rax)` and
+    /// accounted for the quotient to be stored in `rax`.
+    pub fn div(&mut self, divisor: Reg, dst: (Reg, Reg), kind: DivKind, size: OperandSize) {
+        let tmp = if size == OperandSize::S64 && kind == DivKind::Signed {
+            Some(regs::scratch())
+        } else {
+            None
+        };
+
+        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: tmp.map(|reg| reg.into()),
+        });
+    }
+
     /// 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

@@ -1,11 +1,11 @@
 use super::{
     address::Address,
     asm::{Assembler, Operand},
-    regs::{rbp, rsp},
+    regs::{self, rbp, rsp},
 };
-use crate::abi::LocalSlot;
 use crate::isa::reg::Reg;
-use crate::masm::{MacroAssembler as Masm, OperandSize, RegImm};
+use crate::masm::{DivKind, MacroAssembler as Masm, OperandSize, RegImm};
+use crate::{abi::LocalSlot, codegen::CodeGenContext, stack::Val};
 use cranelift_codegen::{isa::x64::settings as x64_settings, settings, Final, MachBufferFinalized};
 
 /// x64 MacroAssembler.
@@ -152,6 +152,28 @@ impl Masm for MacroAssembler {
         self.asm.mul(src, dst, size);
     }
 
+    fn div(&mut self, context: &mut CodeGenContext, kind: DivKind, 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.div(divisor, (rax, rdx), kind, size);
+
+        // Free the divisor and rdx.
+        context.free_gpr(divisor);
+        context.free_gpr(rdx);
+
+        // Push the quotient.
+        context.stack.push(Val::reg(rax));
+    }
+
     fn epilogue(&mut self, locals_size: u32) {
         assert!(self.sp_offset == locals_size);
 

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

@@ -87,8 +87,8 @@ impl TargetIsa for X64 {
         let frame = Frame::new(&abi_sig, &mut body, &mut validator, &abi)?;
         // TODO Add in floating point bitmask
         let regalloc = RegAlloc::new(RegSet::new(ALL_GPR, 0), regs::scratch());
-        let codegen_context = CodeGenContext::new(&mut masm, stack, &frame);
-        let mut codegen = CodeGen::new::<abi::X64ABI>(codegen_context, abi_sig, regalloc);
+        let codegen_context = CodeGenContext::new(regalloc, stack, &frame);
+        let mut codegen = CodeGen::new::<abi::X64ABI>(&mut masm, codegen_context, abi_sig);
 
         codegen.emit(&mut body, validator)?;