Implement register allocation for inline assembly

This commit is contained in:
Gary Guo 2021-11-22 01:44:32 +00:00
parent fd87c6dab9
commit 042eb37837
1 changed files with 210 additions and 73 deletions

View File

@ -6,6 +6,7 @@ use std::fmt::Write;
use rustc_ast::ast::{InlineAsmOptions, InlineAsmTemplatePiece};
use rustc_middle::mir::InlineAsmOperand;
use rustc_span::Symbol;
use rustc_target::asm::*;
pub(crate) fn codegen_inline_asm<'tcx>(
@ -115,11 +116,21 @@ pub(crate) fn codegen_inline_asm<'tcx>(
offset
};
let mut asm_gen = InlineAssemblyGenerator {
tcx: fx.tcx,
arch: InlineAsmArch::X86_64,
template,
operands,
options,
registers: Vec::new(),
};
asm_gen.allocate_registers();
// FIXME overlap input and output slots to save stack space
for operand in operands {
for (i, operand) in operands.iter().enumerate() {
match *operand {
InlineAsmOperand::In { reg, ref value } => {
let reg = expect_reg(reg);
let reg = asm_gen.registers[i].unwrap();
clobbered_regs.push((reg, new_slot(reg.reg_class())));
inputs.push((
reg,
@ -128,7 +139,7 @@ pub(crate) fn codegen_inline_asm<'tcx>(
));
}
InlineAsmOperand::Out { reg, late: _, place } => {
let reg = expect_reg(reg);
let reg = asm_gen.registers[i].unwrap();
clobbered_regs.push((reg, new_slot(reg.reg_class())));
if let Some(place) = place {
outputs.push((
@ -139,7 +150,7 @@ pub(crate) fn codegen_inline_asm<'tcx>(
}
}
InlineAsmOperand::InOut { reg, late: _, ref in_value, out_place } => {
let reg = expect_reg(reg);
let reg = asm_gen.registers[i].unwrap();
clobbered_regs.push((reg, new_slot(reg.reg_class())));
inputs.push((
reg,
@ -164,25 +175,145 @@ pub(crate) fn codegen_inline_asm<'tcx>(
fx.inline_asm_index += 1;
let asm_name = format!("{}__inline_asm_{}", fx.symbol_name, inline_asm_index);
let generated_asm = generate_asm_wrapper(
&asm_name,
InlineAsmArch::X86_64,
options,
template,
clobbered_regs,
&inputs,
&outputs,
);
let generated_asm = asm_gen.generate_asm_wrapper(&asm_name, clobbered_regs, &inputs, &outputs);
fx.cx.global_asm.push_str(&generated_asm);
call_inline_asm(fx, &asm_name, slot_size, inputs, outputs);
}
fn generate_asm_wrapper(
asm_name: &str,
struct InlineAssemblyGenerator<'a, 'tcx> {
tcx: TyCtxt<'tcx>,
arch: InlineAsmArch,
template: &'a [InlineAsmTemplatePiece],
operands: &'a [InlineAsmOperand<'tcx>],
options: InlineAsmOptions,
template: &[InlineAsmTemplatePiece],
registers: Vec<Option<InlineAsmReg>>,
}
impl<'tcx> InlineAssemblyGenerator<'_, 'tcx> {
fn allocate_registers(&mut self) {
let sess = self.tcx.sess;
let map = allocatable_registers(
self.arch,
|feature| sess.target_features.contains(&Symbol::intern(feature)),
&sess.target,
);
let mut allocated = FxHashMap::<_, (bool, bool)>::default();
let mut regs = vec![None; self.operands.len()];
// Add explicit registers to the allocated set.
for (i, operand) in self.operands.iter().enumerate() {
match *operand {
InlineAsmOperand::In { reg: InlineAsmRegOrRegClass::Reg(reg), .. } => {
regs[i] = Some(reg);
allocated.entry(reg).or_default().0 = true;
}
InlineAsmOperand::Out {
reg: InlineAsmRegOrRegClass::Reg(reg), late: true, ..
} => {
regs[i] = Some(reg);
allocated.entry(reg).or_default().1 = true;
}
InlineAsmOperand::Out { reg: InlineAsmRegOrRegClass::Reg(reg), .. }
| InlineAsmOperand::InOut { reg: InlineAsmRegOrRegClass::Reg(reg), .. } => {
regs[i] = Some(reg);
allocated.insert(reg, (true, true));
}
_ => (),
}
}
// Allocate out/inout/inlateout registers first because they are more constrained.
for (i, operand) in self.operands.iter().enumerate() {
match *operand {
InlineAsmOperand::Out {
reg: InlineAsmRegOrRegClass::RegClass(class),
late: false,
..
}
| InlineAsmOperand::InOut {
reg: InlineAsmRegOrRegClass::RegClass(class), ..
} => {
let mut alloc_reg = None;
for &reg in &map[&class] {
let mut used = false;
reg.overlapping_regs(|r| {
if allocated.contains_key(&r) {
used = true;
}
});
if !used {
alloc_reg = Some(reg);
break;
}
}
let reg = alloc_reg.expect("cannot allocate registers");
regs[i] = Some(reg);
allocated.insert(reg, (true, true));
}
_ => (),
}
}
// Allocate in/lateout.
for (i, operand) in self.operands.iter().enumerate() {
match *operand {
InlineAsmOperand::In { reg: InlineAsmRegOrRegClass::RegClass(class), .. } => {
let mut alloc_reg = None;
for &reg in &map[&class] {
let mut used = false;
reg.overlapping_regs(|r| {
if allocated.get(&r).copied().unwrap_or_default().0 {
used = true;
}
});
if !used {
alloc_reg = Some(reg);
break;
}
}
let reg = alloc_reg.expect("cannot allocate registers");
regs[i] = Some(reg);
allocated.entry(reg).or_default().0 = true;
}
InlineAsmOperand::Out {
reg: InlineAsmRegOrRegClass::RegClass(class),
late: true,
..
} => {
let mut alloc_reg = None;
for &reg in &map[&class] {
let mut used = false;
reg.overlapping_regs(|r| {
if allocated.get(&r).copied().unwrap_or_default().1 {
used = true;
}
});
if !used {
alloc_reg = Some(reg);
break;
}
}
let reg = alloc_reg.expect("cannot allocate registers");
regs[i] = Some(reg);
allocated.entry(reg).or_default().1 = true;
}
_ => (),
}
}
self.registers = regs;
}
fn generate_asm_wrapper(
&self,
asm_name: &str,
clobbered_regs: Vec<(InlineAsmReg, Size)>,
inputs: &[(InlineAsmReg, Size, Value)],
outputs: &[(InlineAsmReg, Size, CPlace<'_>)],
@ -198,46 +329,51 @@ fn generate_asm_wrapper(
generated_asm.push_str(" mov rbp,rdi\n");
// Save clobbered registers
if !options.contains(InlineAsmOptions::NORETURN) {
if !self.options.contains(InlineAsmOptions::NORETURN) {
// FIXME skip registers saved by the calling convention
for &(reg, offset) in &clobbered_regs {
save_register(&mut generated_asm, arch, reg, offset);
save_register(&mut generated_asm, self.arch, reg, offset);
}
}
// Write input registers
for &(reg, offset, _value) in inputs {
restore_register(&mut generated_asm, arch, reg, offset);
restore_register(&mut generated_asm, self.arch, reg, offset);
}
if options.contains(InlineAsmOptions::ATT_SYNTAX) {
if self.options.contains(InlineAsmOptions::ATT_SYNTAX) {
generated_asm.push_str(".att_syntax\n");
}
// The actual inline asm
for piece in template {
for piece in self.template {
match piece {
InlineAsmTemplatePiece::String(s) => {
generated_asm.push_str(s);
}
InlineAsmTemplatePiece::Placeholder { operand_idx: _, modifier: _, span: _ } => todo!(),
InlineAsmTemplatePiece::Placeholder { operand_idx, modifier, span: _ } => {
self.registers[*operand_idx]
.unwrap()
.emit(&mut generated_asm, self.arch, *modifier)
.unwrap();
}
}
}
generated_asm.push('\n');
if options.contains(InlineAsmOptions::ATT_SYNTAX) {
if self.options.contains(InlineAsmOptions::ATT_SYNTAX) {
generated_asm.push_str(".intel_syntax noprefix\n");
}
if !options.contains(InlineAsmOptions::NORETURN) {
if !self.options.contains(InlineAsmOptions::NORETURN) {
// Read output registers
for &(reg, offset, _place) in outputs {
save_register(&mut generated_asm, arch, reg, offset);
save_register(&mut generated_asm, self.arch, reg, offset);
}
// Restore clobbered registers
for &(reg, offset) in clobbered_regs.iter().rev() {
restore_register(&mut generated_asm, arch, reg, offset);
restore_register(&mut generated_asm, self.arch, reg, offset);
}
generated_asm.push_str(" pop rbp\n");
@ -253,6 +389,7 @@ fn generate_asm_wrapper(
generated_asm
}
}
fn call_inline_asm<'tcx>(
fx: &mut FunctionCx<'_, '_, 'tcx>,