Rework br_table to use BlockCall (#5731)

Rework br_table to use BlockCall, allowing us to avoid adding new nodes during ssa construction to hold block arguments. Additionally, many places where we previously matched on InstructionData to extract branch destinations can be replaced with a use of branch_destination or branch_destination_mut.
This commit is contained in:
Trevor Elliott
2023-02-16 09:23:27 -08:00
committed by GitHub
parent c3c16eb207
commit 80c147d9c0
26 changed files with 475 additions and 269 deletions

View File

@@ -191,16 +191,23 @@ pub(crate) fn visit_block_succs<F: FnMut(Inst, Block, bool)>(
}
ir::InstructionData::BranchTable { table, .. } => {
let pool = &f.dfg.value_lists;
let table = &f.stencil.dfg.jump_tables[*table];
// The default block is reached via a direct conditional branch,
// so it is not part of the table. We visit the default block first
// explicitly, as some callers of visit_block_succs depend on that
// ordering.
visit(inst, table.default_block(), false);
// so it is not part of the table. We visit the default block
// first explicitly, to mirror the traversal order of
// `JumpTableData::all_branches`, and transitively the order of
// `InstructionData::branch_destination`.
//
// Additionally, this case is why we are unable to replace this
// whole function with a loop over `branch_destination`: we need
// to report which branch targets come from the table vs the
// default.
visit(inst, table.default_block().block(pool), false);
for &dest in table.as_slice() {
visit(inst, dest, true);
for dest in table.as_slice() {
visit(inst, dest.block(pool), true);
}
}

View File

@@ -718,7 +718,7 @@ impl DataFlowGraph {
.iter()
.chain(
self.insts[inst]
.branch_destination()
.branch_destination(&self.jump_tables)
.into_iter()
.flat_map(|branch| branch.args_slice(&self.value_lists).iter()),
)
@@ -735,10 +735,10 @@ impl DataFlowGraph {
self.inst_args_mut(inst)[i] = body(self, arg);
}
for block_ix in 0..self.insts[inst].branch_destination().len() {
for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
// We aren't changing the size of the args list, so we won't need to write the branch
// back to the instruction.
let mut block = self.insts[inst].branch_destination()[block_ix];
let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
for i in 0..block.args_slice(&self.value_lists).len() {
let arg = block.args_slice(&self.value_lists)[i];
block.args_slice_mut(&mut self.value_lists)[i] = body(self, arg);
@@ -757,8 +757,8 @@ impl DataFlowGraph {
*arg = values.next().unwrap();
}
for block_ix in 0..self.insts[inst].branch_destination().len() {
let mut block = self.insts[inst].branch_destination()[block_ix];
for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
for arg in block.args_slice_mut(&mut self.value_lists) {
*arg = values.next().unwrap();
}

View File

@@ -6,9 +6,9 @@
use crate::entity::{PrimaryMap, SecondaryMap};
use crate::ir::{
self, Block, DataFlowGraph, DynamicStackSlot, DynamicStackSlotData, DynamicStackSlots,
DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, InstructionData,
JumpTable, JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot,
StackSlotData, StackSlots, Table, TableData, Type,
DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, JumpTable,
JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot, StackSlotData,
StackSlots, Table, TableData, Type,
};
use crate::isa::CallConv;
use crate::value_label::ValueLabelsRanges;
@@ -273,37 +273,10 @@ impl FunctionStencil {
/// Rewrite the branch destination to `new_dest` if the destination matches `old_dest`.
/// Does nothing if called with a non-jump or non-branch instruction.
pub fn rewrite_branch_destination(&mut self, inst: Inst, old_dest: Block, new_dest: Block) {
match &mut self.dfg.insts[inst] {
InstructionData::Jump {
destination: dest, ..
} => {
if dest.block(&self.dfg.value_lists) == old_dest {
dest.set_block(new_dest, &mut self.dfg.value_lists)
}
for dest in self.dfg.insts[inst].branch_destination_mut(&mut self.dfg.jump_tables) {
if dest.block(&self.dfg.value_lists) == old_dest {
dest.set_block(new_dest, &mut self.dfg.value_lists)
}
InstructionData::Brif {
blocks: [block_then, block_else],
..
} => {
if block_then.block(&self.dfg.value_lists) == old_dest {
block_then.set_block(new_dest, &mut self.dfg.value_lists);
}
if block_else.block(&self.dfg.value_lists) == old_dest {
block_else.set_block(new_dest, &mut self.dfg.value_lists);
}
}
InstructionData::BranchTable { table, .. } => {
for entry in self.dfg.jump_tables[*table].all_branches_mut() {
if *entry == old_dest {
*entry = new_dest;
}
}
}
inst => debug_assert!(!inst.opcode().is_branch()),
}
}

View File

@@ -304,13 +304,13 @@ impl InstructionData {
/// Get the destinations of this instruction, if it's a branch.
///
/// `br_table` returns the empty slice.
pub fn branch_destination(&self) -> &[BlockCall] {
pub fn branch_destination<'a>(&'a self, jump_tables: &'a ir::JumpTables) -> &[BlockCall] {
match self {
Self::Jump {
ref destination, ..
} => std::slice::from_ref(destination),
Self::Brif { blocks, .. } => blocks,
Self::BranchTable { .. } => &[],
Self::Brif { blocks, .. } => blocks.as_slice(),
Self::BranchTable { table, .. } => jump_tables.get(*table).unwrap().all_branches(),
_ => {
debug_assert!(!self.opcode().is_branch());
&[]
@@ -321,14 +321,19 @@ impl InstructionData {
/// Get a mutable slice of the destinations of this instruction, if it's a branch.
///
/// `br_table` returns the empty slice.
pub fn branch_destination_mut(&mut self) -> &mut [BlockCall] {
pub fn branch_destination_mut<'a>(
&'a mut self,
jump_tables: &'a mut ir::JumpTables,
) -> &mut [BlockCall] {
match self {
Self::Jump {
ref mut destination,
..
} => std::slice::from_mut(destination),
Self::Brif { blocks, .. } => blocks,
Self::BranchTable { .. } => &mut [],
Self::Brif { blocks, .. } => blocks.as_mut_slice(),
Self::BranchTable { table, .. } => {
jump_tables.get_mut(*table).unwrap().all_branches_mut()
}
_ => {
debug_assert!(!self.opcode().is_branch());
&mut []

View File

@@ -3,7 +3,8 @@
//! Jump tables are declared in the preamble and assigned an `ir::entities::JumpTable` reference.
//! The actual table of destinations is stored in a `JumpTableData` struct defined in this module.
use crate::ir::entities::Block;
use crate::ir::instructions::ValueListPool;
use crate::ir::BlockCall;
use alloc::vec::Vec;
use core::fmt::{self, Display, Formatter};
use core::slice::{Iter, IterMut};
@@ -23,57 +24,57 @@ use serde::{Deserialize, Serialize};
#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
pub struct JumpTableData {
// Table entries.
table: Vec<Block>,
table: Vec<BlockCall>,
}
impl JumpTableData {
/// Create a new jump table with the provided blocks
pub fn new(def: Block, table: &[Block]) -> Self {
pub fn new(def: BlockCall, table: &[BlockCall]) -> Self {
Self {
table: std::iter::once(def).chain(table.iter().copied()).collect(),
}
}
/// Fetch the default block for this jump table.
pub fn default_block(&self) -> Block {
pub fn default_block(&self) -> BlockCall {
*self.table.first().unwrap()
}
/// Mutable access to the default block of this jump table.
pub fn default_block_mut(&mut self) -> &mut Block {
pub fn default_block_mut(&mut self) -> &mut BlockCall {
self.table.first_mut().unwrap()
}
/// The jump table and default block as a single slice. The default block will always be first.
pub fn all_branches(&self) -> &[Block] {
pub fn all_branches(&self) -> &[BlockCall] {
self.table.as_slice()
}
/// The jump table and default block as a single mutable slice. The default block will always
/// be first.
pub fn all_branches_mut(&mut self) -> &mut [Block] {
pub fn all_branches_mut(&mut self) -> &mut [BlockCall] {
self.table.as_mut_slice()
}
/// Access the jump table as a slice. This excludes the default block.
pub fn as_slice(&self) -> &[Block] {
pub fn as_slice(&self) -> &[BlockCall] {
&self.table.as_slice()[1..]
}
/// Access the jump table as a mutable slice. This excludes the default block.
pub fn as_mut_slice(&mut self) -> &mut [Block] {
pub fn as_mut_slice(&mut self) -> &mut [BlockCall] {
&mut self.table.as_mut_slice()[1..]
}
/// Returns an iterator to the jump table, excluding the default block.
#[deprecated(since = "7.0.0", note = "please use `.as_slice()` instead")]
pub fn iter(&self) -> Iter<Block> {
pub fn iter(&self) -> Iter<BlockCall> {
self.as_slice().iter()
}
/// Returns an iterator that allows modifying each value, excluding the default block.
#[deprecated(since = "7.0.0", note = "please use `.as_mut_slice()` instead")]
pub fn iter_mut(&mut self) -> IterMut<Block> {
pub fn iter_mut(&mut self) -> IterMut<BlockCall> {
self.as_mut_slice().iter_mut()
}
@@ -81,15 +82,26 @@ impl JumpTableData {
pub fn clear(&mut self) {
self.table.drain(1..);
}
/// Return a value that can display the contents of this jump table.
pub fn display<'a>(&'a self, pool: &'a ValueListPool) -> DisplayJumpTable<'a> {
DisplayJumpTable { jt: self, pool }
}
}
impl Display for JumpTableData {
fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
write!(fmt, "{}, [", self.default_block())?;
if let Some((first, rest)) = self.as_slice().split_first() {
write!(fmt, "{}", first)?;
/// A wrapper for the context required to display a [JumpTableData].
pub struct DisplayJumpTable<'a> {
jt: &'a JumpTableData,
pool: &'a ValueListPool,
}
impl<'a> Display for DisplayJumpTable<'a> {
fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
write!(fmt, "{}, [", self.jt.default_block().display(self.pool))?;
if let Some((first, rest)) = self.jt.as_slice().split_first() {
write!(fmt, "{}", first.display(self.pool))?;
for block in rest {
write!(fmt, ", {}", block)?;
write!(fmt, ", {}", block.display(self.pool))?;
}
}
write!(fmt, "]")
@@ -100,12 +112,14 @@ impl Display for JumpTableData {
mod tests {
use super::JumpTableData;
use crate::entity::EntityRef;
use crate::ir::Block;
use alloc::string::ToString;
use crate::ir::instructions::ValueListPool;
use crate::ir::{Block, BlockCall, Value};
use std::string::ToString;
#[test]
fn empty() {
let def = Block::new(0);
let mut pool = ValueListPool::default();
let def = BlockCall::new(Block::new(0), &[], &mut pool);
let jt = JumpTableData::new(def, &[]);
@@ -114,7 +128,7 @@ mod tests {
assert_eq!(jt.as_slice().get(0), None);
assert_eq!(jt.as_slice().get(10), None);
assert_eq!(jt.to_string(), "block0, []");
assert_eq!(jt.display(&pool).to_string(), "block0, []");
assert_eq!(jt.all_branches(), [def]);
assert_eq!(jt.as_slice(), []);
@@ -122,16 +136,33 @@ mod tests {
#[test]
fn insert() {
let def = Block::new(0);
let mut pool = ValueListPool::default();
let v0 = Value::new(0);
let v1 = Value::new(1);
let e0 = Block::new(0);
let e1 = Block::new(1);
let e2 = Block::new(2);
let jt = JumpTableData::new(def, &[e1, e2, e1]);
let def = BlockCall::new(e0, &[], &mut pool);
let b1 = BlockCall::new(e1, &[v0], &mut pool);
let b2 = BlockCall::new(e2, &[], &mut pool);
let b3 = BlockCall::new(e1, &[v1], &mut pool);
let jt = JumpTableData::new(def, &[b1, b2, b3]);
assert_eq!(jt.default_block(), def);
assert_eq!(jt.to_string(), "block0, [block1, block2, block1]");
assert_eq!(
jt.display(&pool).to_string(),
"block0, [block1(v0), block2, block1(v1)]"
);
assert_eq!(jt.all_branches(), [def, e1, e2, e1]);
assert_eq!(jt.as_slice(), [e1, e2, e1]);
assert_eq!(jt.all_branches(), [def, b1, b2, b3]);
assert_eq!(jt.as_slice(), [b1, b2, b3]);
assert_eq!(jt.as_slice()[0].args_slice(&pool), [v0]);
assert_eq!(jt.as_slice()[1].args_slice(&pool), []);
assert_eq!(jt.as_slice()[2].args_slice(&pool), [v1]);
}
}

View File

@@ -377,9 +377,14 @@ mod test {
let mut pos = FuncCursor::new(&mut func);
pos.insert_block(bb0);
let jt = pos
.func
.create_jump_table(JumpTableData::new(bb3, &[bb1, bb2]));
let jt_data = JumpTableData::new(
pos.func.dfg.block_call(bb3, &[]),
&[
pos.func.dfg.block_call(bb1, &[]),
pos.func.dfg.block_call(bb2, &[]),
],
);
let jt = pos.func.create_jump_table(jt_data);
pos.ins().br_table(arg0, jt);
pos.insert_block(bb1);

View File

@@ -418,9 +418,14 @@ mod test {
let mut pos = FuncCursor::new(&mut func);
pos.insert_block(bb0);
let jt = pos
.func
.create_jump_table(JumpTableData::new(bb3, &[bb1, bb2]));
let jt_data = JumpTableData::new(
pos.func.dfg.block_call(bb3, &[]),
&[
pos.func.dfg.block_call(bb1, &[]),
pos.func.dfg.block_call(bb2, &[]),
],
);
let jt = pos.func.create_jump_table(jt_data);
pos.ins().br_table(arg0, jt);
pos.insert_block(bb1);

View File

@@ -942,29 +942,13 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
// Avoid immutable borrow by explicitly indexing.
let (inst, succ) = self.vcode.block_order().succ_indices(block)[succ_idx];
// Get branch args and convert to Regs.
let branch_args = match &self.f.dfg.insts[inst] {
InstructionData::Jump {
destination: block, ..
} => block.args_slice(&self.f.dfg.value_lists),
InstructionData::Brif {
blocks: [then_block, else_block],
..
} => {
// NOTE: `succ_idx == 0` implying that we're traversing the `then_block` is
// enforced by the traversal order defined in `visit_block_succs`. Eventually
// we should traverse the `branch_destination` slice there, which would
// simplify computing the branch args significantly.
if succ_idx == 0 {
then_block.args_slice(&self.f.dfg.value_lists)
} else {
assert!(succ_idx == 1);
else_block.args_slice(&self.f.dfg.value_lists)
}
}
InstructionData::BranchTable { .. } => &[],
_ => unreachable!(),
};
// The use of `succ_idx` to index `branch_destination` is valid on the assumption that
// the traversal order defined in `visit_block_succs` mirrors the order returned by
// `branch_destination`. If that assumption is violated, the branch targets returned
// here will not match the clif.
let branches = self.f.dfg.insts[inst].branch_destination(&self.f.dfg.jump_tables);
let branch_args = branches[succ_idx].args_slice(&self.f.dfg.value_lists);
let mut branch_arg_vregs: SmallVec<[Reg; 16]> = smallvec![];
for &arg in branch_args {
let arg = self.f.dfg.resolve_aliases(arg);

View File

@@ -7,7 +7,6 @@ use crate::ir;
use crate::ir::Function;
use crate::ir::{Block, BlockCall, Inst, Value};
use crate::timing;
use arrayvec::ArrayVec;
use bumpalo::Bump;
use cranelift_entity::SecondaryMap;
use smallvec::SmallVec;
@@ -171,9 +170,10 @@ struct BlockSummary<'a> {
/// We don't bother to include transfers that pass zero parameters
/// since that makes more work for the solver for no purpose.
///
/// Note that, because blocks used with `br_table`s cannot have block
/// arguments, there are at most two outgoing edges from these blocks.
dests: ArrayVec<OutEdge<'a>, 2>,
/// We optimize for the case where a branch instruction has up to two
/// outgoing edges, as unconditional jumps and conditional branches are
/// more prominent than br_table.
dests: SmallVec<[OutEdge<'a>; 2]>,
}
impl<'a> BlockSummary<'a> {
@@ -239,7 +239,11 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
let mut summary = BlockSummary::new(&bump, formals);
for inst in func.layout.block_insts(b) {
for (ix, dest) in func.dfg.insts[inst].branch_destination().iter().enumerate() {
for (ix, dest) in func.dfg.insts[inst]
.branch_destination(&func.dfg.jump_tables)
.iter()
.enumerate()
{
if let Some(edge) = OutEdge::new(&bump, &func.dfg, inst, ix, *dest) {
summary.dests.push(edge);
}
@@ -382,8 +386,8 @@ pub fn do_remove_constant_phis(func: &mut Function, domtree: &mut DominatorTree)
}
let dfg = &mut func.dfg;
let block =
&mut dfg.insts[edge.inst].branch_destination_mut()[edge.branch_index as usize];
let dests = dfg.insts[edge.inst].branch_destination_mut(&mut dfg.jump_tables);
let block = &mut dests[edge.branch_index as usize];
old_actuals.extend(block.args_slice(&dfg.value_lists));

View File

@@ -849,8 +849,9 @@ impl<'a> Verifier<'a> {
format!("invalid jump table reference {}", j),
))
} else {
for &block in self.func.stencil.dfg.jump_tables[j].all_branches() {
self.verify_block(inst, block, errors)?;
let pool = &self.func.stencil.dfg.value_lists;
for block in self.func.stencil.dfg.jump_tables[j].all_branches() {
self.verify_block(inst, block.block(pool), errors)?;
}
Ok(())
}
@@ -1285,53 +1286,19 @@ impl<'a> Verifier<'a> {
errors: &mut VerifierErrors,
) -> VerifierStepResult<()> {
match &self.func.dfg.insts[inst] {
ir::InstructionData::Jump {
destination: block, ..
} => {
let iter = self
.func
.dfg
.block_params(block.block(&self.func.dfg.value_lists))
.iter()
.map(|&v| self.func.dfg.value_type(v));
let args = block.args_slice(&self.func.dfg.value_lists);
self.typecheck_variable_args_iterator(inst, iter, args, errors)?;
ir::InstructionData::Jump { destination, .. } => {
self.typecheck_block_call(inst, destination, errors)?;
}
ir::InstructionData::Brif {
blocks: [block_then, block_else],
..
} => {
let iter = self
.func
.dfg
.block_params(block_then.block(&self.func.dfg.value_lists))
.iter()
.map(|&v| self.func.dfg.value_type(v));
let args_then = block_then.args_slice(&self.func.dfg.value_lists);
self.typecheck_variable_args_iterator(inst, iter, args_then, errors)?;
let iter = self
.func
.dfg
.block_params(block_else.block(&self.func.dfg.value_lists))
.iter()
.map(|&v| self.func.dfg.value_type(v));
let args_else = block_else.args_slice(&self.func.dfg.value_lists);
self.typecheck_variable_args_iterator(inst, iter, args_else, errors)?;
self.typecheck_block_call(inst, block_then, errors)?;
self.typecheck_block_call(inst, block_else, errors)?;
}
ir::InstructionData::BranchTable { table, .. } => {
for block in self.func.stencil.dfg.jump_tables[*table].all_branches() {
let arg_count = self.func.dfg.num_block_params(*block);
if arg_count != 0 {
return errors.nonfatal((
inst,
self.context(inst),
format!(
"takes no arguments, but had target {} with {} arguments",
block, arg_count,
),
));
}
self.typecheck_block_call(inst, block, errors)?;
}
}
inst => debug_assert!(!inst.opcode().is_branch()),
@@ -1358,6 +1325,23 @@ impl<'a> Verifier<'a> {
Ok(())
}
fn typecheck_block_call(
&self,
inst: Inst,
block: &ir::BlockCall,
errors: &mut VerifierErrors,
) -> VerifierStepResult<()> {
let pool = &self.func.dfg.value_lists;
let iter = self
.func
.dfg
.block_params(block.block(pool))
.iter()
.map(|&v| self.func.dfg.value_type(v));
let args = block.args_slice(pool);
self.typecheck_variable_args_iterator(inst, iter, args, errors)
}
fn typecheck_variable_args_iterator<I: Iterator<Item = Type>>(
&self,
inst: Inst,

View File

@@ -421,7 +421,9 @@ pub fn write_operands(w: &mut dyn Write, dfg: &DataFlowGraph, inst: Inst) -> fmt
write!(w, " {}, {}", arg, block_then.display(pool))?;
write!(w, ", {}", block_else.display(pool))
}
BranchTable { arg, table, .. } => write!(w, " {}, {}", arg, jump_tables[table]),
BranchTable { arg, table, .. } => {
write!(w, " {}, {}", arg, jump_tables[table].display(pool))
}
Call {
func_ref, ref args, ..
} => write!(w, " {}({})", func_ref, DisplayValues(args.as_slice(pool))),