winch: Add full support for integer sub and add instructions (#5737)

This patch adds complete support for the `sub` and `add` WebAssembly instructions
for x64, and complete support for the `add` WebAssembly instruction for aarch64.

This patch also refactors how the binary operations get constructed within the
`VisitOperator` trait implementation. The refactor adds methods in the
`CodeGenContext` to abstract all the common steps to emit binary operations,
making this process less repetitive and less brittle (e.g. omitting to push the resulting value
to the stack, or omitting to free registers after used).

This patch also improves test coverage and refactors the filetests directory to make it
easier to add tests for other instructions.
This commit is contained in:
Saúl Cabrera
2023-02-08 20:01:44 -05:00
committed by GitHub
parent 9637840b4b
commit 7c5c7e4b6d
59 changed files with 992 additions and 120 deletions

View File

@@ -1,9 +1,9 @@
use crate::{
abi::{ABISig, ABI},
frame::Frame,
masm::{MacroAssembler, OperandSize},
masm::{MacroAssembler, OperandSize, RegImm},
regalloc::RegAlloc,
stack::Stack,
stack::{Stack, Val},
};
use anyhow::Result;
use wasmparser::{BinaryReader, FuncValidator, ValType, ValidatorResources, VisitOperator};
@@ -22,9 +22,67 @@ impl<'a, M> CodeGenContext<'a, M>
where
M: MacroAssembler,
{
/// Create a new code generation context.
pub fn new(masm: &'a mut M, stack: Stack, frame: &'a Frame) -> Self {
Self { masm, stack, frame }
}
/// Prepares arguments for emitting an i32 binary operation.
pub fn i32_binop<F>(&mut self, regalloc: &mut RegAlloc, emit: &mut F)
where
F: FnMut(&mut M, RegImm, RegImm, OperandSize),
{
let top = self.stack.peek().expect("value at stack top");
if top.is_i32_const() {
let val = self
.stack
.pop_i32_const()
.expect("i32 const value at stack top");
let reg = regalloc.pop_to_reg(self, OperandSize::S32);
emit(
&mut self.masm,
RegImm::reg(reg),
RegImm::imm(val as i64),
OperandSize::S32,
);
self.stack.push(Val::reg(reg));
} else {
let src = regalloc.pop_to_reg(self, OperandSize::S32);
let dst = regalloc.pop_to_reg(self, OperandSize::S32);
emit(&mut self.masm, dst.into(), src.into(), OperandSize::S32);
regalloc.free_gpr(src);
self.stack.push(Val::reg(dst));
}
}
/// Prepares arguments for emitting an i64 binary operation.
pub fn i64_binop<F>(&mut self, regalloc: &mut RegAlloc, emit: &mut F)
where
F: FnMut(&mut M, RegImm, RegImm, OperandSize),
{
let top = self.stack.peek().expect("value at stack top");
if top.is_i64_const() {
let val = self
.stack
.pop_i64_const()
.expect("i64 const value at stack top");
let reg = regalloc.pop_to_reg(self, OperandSize::S64);
emit(
&mut self.masm,
RegImm::reg(reg),
RegImm::imm(val),
OperandSize::S64,
);
self.stack.push(Val::reg(reg));
} else {
let src = regalloc.pop_to_reg(self, OperandSize::S64);
let dst = regalloc.pop_to_reg(self, OperandSize::S64);
emit(&mut self.masm, dst.into(), src.into(), OperandSize::S64);
regalloc.free_gpr(src);
self.stack.push(Val::reg(dst));
}
}
}
/// The code generation abstraction.

View File

@@ -19,8 +19,8 @@ pub(crate) enum Operand {
Reg(Reg),
/// Memory address.
Mem(Address),
/// Immediate.
Imm(i32),
/// 64-bit signed immediate.
Imm(i64),
}
// Conversions between winch-codegen aarch64 types and cranelift-codegen

View File

@@ -143,6 +143,10 @@ impl Masm for MacroAssembler {
self.asm.add(rhs.into(), lhs.into(), dst.into(), size);
}
fn sub(&mut self, _dst: RegImm, _lhs: RegImm, _rhs: RegImm, _size: OperandSize) {
todo!()
}
fn zero(&mut self, reg: Reg) {
self.asm.load_constant(0, reg);
}

View File

@@ -12,7 +12,7 @@ use cranelift_codegen::{
settings, Final, MachBuffer, MachBufferFinalized, MachInstEmit, Writable,
};
use super::address::Address;
use super::{address::Address, regs};
/// A x64 instruction operand.
#[derive(Debug, Copy, Clone)]
@@ -21,8 +21,8 @@ pub(crate) enum Operand {
Reg(Reg),
/// Memory address.
Mem(Address),
/// Immediate.
Imm(i32),
/// Signed 64-bit immediate.
Imm(i64),
}
// Conversions between winch-codegen x64 types and cranelift-codegen x64 types.
@@ -198,6 +198,37 @@ impl Assembler {
}
}
/// Subtract instruction variants.
pub fn sub(&mut self, src: Operand, dst: Operand, size: OperandSize) {
match &(src, dst) {
(Operand::Imm(imm), Operand::Reg(dst)) => {
if let Ok(val) = i32::try_from(*imm) {
self.sub_ir(val, *dst, size)
} else {
let scratch = regs::scratch();
self.mov_ir(*imm as u64, scratch, size);
self.sub_rr(scratch, *dst, size);
}
}
(Operand::Reg(src), Operand::Reg(dst)) => self.sub_rr(*src, *dst, size),
_ => panic!(
"Invalid operand combination for sub; src = {:?} dst = {:?}",
src, dst
),
}
}
/// Subtract register and register
pub fn sub_rr(&mut self, src: Reg, dst: Reg, size: OperandSize) {
self.emit(Inst::AluRmiR {
size: size.into(),
op: AluRmiROpcode::Sub,
src1: dst.into(),
src2: src.into(),
dst: dst.into(),
});
}
/// Subtact immediate register.
pub fn sub_ir(&mut self, imm: i32, dst: Reg, size: OperandSize) {
let imm = RegMemImm::imm(imm as u32);
@@ -214,7 +245,15 @@ impl Assembler {
/// Add instruction variants.
pub fn add(&mut self, src: Operand, dst: Operand, size: OperandSize) {
match &(src, dst) {
(Operand::Imm(imm), Operand::Reg(dst)) => self.add_ir(*imm, *dst, size),
(Operand::Imm(imm), Operand::Reg(dst)) => {
if let Ok(val) = i32::try_from(*imm) {
self.add_ir(val, *dst, size)
} else {
let scratch = regs::scratch();
self.mov_ir(*imm as u64, scratch, size);
self.add_rr(scratch, *dst, size);
}
}
(Operand::Reg(src), Operand::Reg(dst)) => self.add_rr(*src, *dst, size),
_ => panic!(
"Invalid operand combination for add; src = {:?} dst = {:?}",

View File

@@ -126,6 +126,19 @@ impl Masm for MacroAssembler {
self.asm.add(src, dst, size);
}
fn sub(&mut self, dst: RegImm, lhs: RegImm, rhs: RegImm, size: OperandSize) {
let (src, dst): (Operand, Operand) = if dst == lhs {
(rhs.into(), dst.into())
} else {
panic!(
"the destination and first source argument must be the same, dst={:?}, lhs={:?}",
dst, lhs
);
};
self.asm.sub(src, dst, size);
}
fn epilogue(&mut self, locals_size: u32) {
assert!(self.sp_offset == locals_size);

View File

@@ -18,8 +18,8 @@ pub(crate) enum OperandSize {
pub(crate) enum RegImm {
/// A register.
Reg(Reg),
/// An immediate.
Imm(i32),
/// 64-bit signed immediate.
Imm(i64),
}
impl RegImm {
@@ -29,7 +29,7 @@ impl RegImm {
}
/// Immediate constructor.
pub fn imm(imm: i32) -> Self {
pub fn imm(imm: i64) -> Self {
RegImm::Imm(imm)
}
}
@@ -88,6 +88,9 @@ pub(crate) trait MacroAssembler {
/// Perform add operation.
fn add(&mut self, dst: RegImm, lhs: RegImm, rhs: RegImm, size: OperandSize);
/// Perform subtraction operation.
fn sub(&mut self, dst: RegImm, lhs: RegImm, rhs: RegImm, size: OperandSize);
/// Push the register to the stack, returning the offset.
fn push(&mut self, src: Reg) -> u32;

View File

@@ -76,7 +76,8 @@ impl RegAlloc {
) {
match src {
Val::Reg(src) => masm.mov(RegImm::reg(src), RegImm::reg(dst), size),
Val::I32(imm) => masm.mov(RegImm::imm(imm), RegImm::reg(dst), size),
Val::I32(imm) => masm.mov(RegImm::imm(imm.into()), RegImm::reg(dst), size),
Val::I64(imm) => masm.mov(RegImm::imm(imm), RegImm::reg(dst), size),
Val::Local(index) => {
let slot = frame
.get_local(index)

View File

@@ -6,6 +6,8 @@ use std::collections::VecDeque;
pub(crate) enum Val {
/// I32 Constant.
I32(i32),
/// I64 Constant.
I64(i64),
/// A register.
Reg(Reg),
/// A local slot.
@@ -20,6 +22,11 @@ impl Val {
Self::I32(v)
}
/// Create a new I64 constant value.
pub fn i64(v: i64) -> Self {
Self::I64(v)
}
/// Create a new Reg value.
pub fn reg(r: Reg) -> Self {
Self::Reg(r)
@@ -60,6 +67,17 @@ impl Val {
}
}
/// Get the integer representation of the value.
///
/// # Panics
/// This method will panic if the value is not an i64.
pub fn get_i64(&self) -> i64 {
match self {
Self::I64(v) => *v,
v => panic!("expected value {:?} to be i64", v),
}
}
/// Check whether the value is an i32 constant.
pub fn is_i32_const(&self) -> bool {
match *self {
@@ -67,6 +85,14 @@ impl Val {
_ => false,
}
}
/// Check whether the value is an i64 constant.
pub fn is_i64_const(&self) -> bool {
match *self {
Self::I64(_) => true,
_ => false,
}
}
}
/// The shadow stack used for compilation.
@@ -89,7 +115,7 @@ impl Stack {
}
/// Peek into the top in the stack.
pub fn peek(&mut self) -> Option<&Val> {
pub fn peek(&self) -> Option<&Val> {
self.inner.back()
}
@@ -98,7 +124,7 @@ impl Stack {
self.inner.pop_back()
}
/// Pops the element at the top of the stack if it is a const;
/// Pops the element at the top of the stack if it is an i32 const;
/// returns `None` otherwise.
pub fn pop_i32_const(&mut self) -> Option<i32> {
match self.peek() {
@@ -107,6 +133,15 @@ impl Stack {
}
}
/// Pops the element at the top of the stack if it is an i64 const;
/// returns `None` otherwise.
pub fn pop_i64_const(&mut self) -> Option<i64> {
match self.peek() {
Some(v) => v.is_i64_const().then(|| self.pop().unwrap().get_i64()),
_ => None,
}
}
/// Pops the element at the top of the stack if it is a register;
/// returns `None` otherwise.
pub fn pop_reg(&mut self) -> Option<Reg> {

View File

@@ -10,45 +10,6 @@ use crate::stack::Val;
use wasmparser::ValType;
use wasmparser::VisitOperator;
impl<'a, M> CodeGen<'a, M>
where
M: MacroAssembler,
{
fn add_imm_i32(&mut self) {
let val = self
.context
.stack
.pop_i32_const()
.expect("i32 constant at stack top");
let reg = self
.regalloc
.pop_to_reg(&mut self.context, OperandSize::S32);
let dst = RegImm::reg(reg);
self.context
.masm
.add(dst, dst, RegImm::imm(val), OperandSize::S32);
self.context.stack.push(Val::reg(reg));
}
fn add_i32(&mut self) {
let src = self
.regalloc
.pop_to_reg(&mut self.context, OperandSize::S32);
let dst = self
.regalloc
.pop_to_reg(&mut self.context, OperandSize::S32);
let lhs = RegImm::reg(dst);
self.context
.masm
.add(lhs, lhs, RegImm::reg(src), OperandSize::S32);
self.regalloc.free_gpr(src);
self.context.stack.push(Val::reg(dst));
}
}
/// A macro to define unsupported WebAssembly operators.
///
/// This macro calls itself recursively;
@@ -71,7 +32,11 @@ macro_rules! def_unsupported {
};
(emit I32Const $($rest:tt)*) => {};
(emit I64Const $($rest:tt)*) => {};
(emit I32Add $($rest:tt)*) => {};
(emit I64Add $($rest:tt)*) => {};
(emit I32Sub $($rest:tt)*) => {};
(emit I64Sub $($rest:tt)*) => {};
(emit LocalGet $($rest:tt)*) => {};
(emit LocalSet $($rest:tt)*) => {};
(emit End $($rest:tt)*) => {};
@@ -89,19 +54,36 @@ where
self.context.stack.push(Val::i32(val));
}
fn visit_i32_add(&mut self) {
let is_const = self
.context
.stack
.peek()
.expect("value at stack top")
.is_i32_const();
fn visit_i64_const(&mut self, val: i64) {
self.context.stack.push(Val::i64(val));
}
if is_const {
self.add_imm_i32();
} else {
self.add_i32();
}
fn visit_i32_add(&mut self) {
self.context
.i32_binop(&mut self.regalloc, &mut |masm: &mut M, dst, src, size| {
masm.add(dst, dst, src, size);
});
}
fn visit_i64_add(&mut self) {
self.context
.i64_binop(&mut self.regalloc, &mut |masm: &mut M, dst, src, size| {
masm.add(dst, dst, src, size);
});
}
fn visit_i32_sub(&mut self) {
self.context
.i32_binop(&mut self.regalloc, &mut |masm: &mut M, dst, src, size| {
masm.sub(dst, dst, src, size);
});
}
fn visit_i64_sub(&mut self) {
self.context
.i64_binop(&mut self.regalloc, &mut |masm: &mut M, dst, src, size| {
masm.sub(dst, dst, src, size);
});
}
fn visit_end(&mut self) {}