Lower aborts (incl. panics) to "return from entry-point", instead of infinite loops.

This commit is contained in:
Eduard-Mihai Burtescu 2023-06-05 23:39:29 +03:00 committed by Eduard-Mihai Burtescu
parent b2e5eb7595
commit ce8c3f8f4c
14 changed files with 272 additions and 77 deletions

View File

@ -30,6 +30,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
### Changed 🛠
- [PR#1070](https://github.com/EmbarkStudios/rust-gpu/pull/1070) made panics (via the `abort` intrinsic)
early-exit (i.e. `return` from) the shader entry-point, instead of looping infinitely
- [PR#1071](https://github.com/EmbarkStudios/rust-gpu/pull/1071) updated toolchain to `nightly-2023-05-27`
## [0.8.0]

View File

@ -2457,8 +2457,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
let is_standard_debug = [Op::Line, Op::NoLine].contains(&inst.class.opcode);
let is_custom_debug = inst.class.opcode == Op::ExtInst
&& inst.operands[0].unwrap_id_ref() == custom_ext_inst_set_import
&& [CustomOp::SetDebugSrcLoc, CustomOp::ClearDebugSrcLoc]
.contains(&CustomOp::decode_from_ext_inst(inst));
&& CustomOp::decode_from_ext_inst(inst).is_debuginfo();
!(is_standard_debug || is_custom_debug)
});

View File

@ -2,6 +2,7 @@ use super::Builder;
use crate::abi::ConvSpirvType;
use crate::builder_spirv::{SpirvValue, SpirvValueExt};
use crate::codegen_cx::CodegenCx;
use crate::custom_insts::CustomInst;
use crate::spirv_type::SpirvType;
use rspirv::spirv::GLOp;
use rustc_codegen_ssa::mir::operand::OperandRef;
@ -337,17 +338,18 @@ impl<'a, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'tcx> {
}
fn abort(&mut self) {
// FIXME(eddyb) this should be cached more efficiently.
let void_ty = SpirvType::Void.def(rustc_span::DUMMY_SP, self);
// HACK(eddyb) there is no `abort` or `trap` instruction in SPIR-V,
// so the best thing we can do is inject an infinite loop.
// (While there is `OpKill`, it doesn't really have the right semantics)
let abort_loop_bb = self.append_sibling_block("abort_loop");
let abort_continue_bb = self.append_sibling_block("abort_continue");
self.br(abort_loop_bb);
// so the best thing we can do is use our own custom instruction.
self.custom_inst(void_ty, CustomInst::Abort);
self.unreachable();
self.switch_to_block(abort_loop_bb);
self.br(abort_loop_bb);
self.switch_to_block(abort_continue_bb);
// HACK(eddyb) we still need an active block in case the user of this
// `Builder` will continue to emit instructions after the `.abort()`.
let post_abort_dead_bb = self.append_sibling_block("post_abort_dead");
self.switch_to_block(post_abort_dead_bb);
}
fn assume(&mut self, _val: Self::Value) {

View File

@ -62,7 +62,7 @@ pub struct CodegenCx<'tcx> {
/// All `panic!(...)`s and builtin panics (from MIR `Assert`s) call into one
/// of these lang items, which we always replace with an "abort", erasing
/// anything passed in (and that "abort" is just an infinite loop for now).
/// anything passed in.
//
// FIXME(eddyb) we should not erase anywhere near as much, but `format_args!`
// is not representable due to containg Rust slices, and Rust 2021 has made

View File

@ -125,4 +125,55 @@ def_custom_insts! {
// Leave the most recent inlined call frame entered by a `PushInlinedCallFrame`
// (i.e. the inlined call frames form a virtual call stack in debuginfo).
3 => PopInlinedCallFrame,
// [Semantic] Similar to some proposed `OpAbort`, but without any ability to
// indicate abnormal termination (so it's closer to `OpTerminateInvocation`,
// which we could theoretically use, but that's limited to fragment shaders).
//
// Lowering takes advantage of inlining happening before CFG structurization
// (by forcing inlining of `Abort`s all the way up to entry-points, as to be
// able to turn the `Abort`s into regular `OpReturn`s, from an entry-point),
// but if/when inlining works on structured SPIR-T instead, it's not much
// harder to make any call to a "may (transitively) abort" function branch on
// an additional returned `bool`, instead (i.e. a form of emulated unwinding).
//
// As this is a custom terminator, it must only appear before `OpUnreachable`,
// *without* any instructions in between (not even debuginfo ones).
//
// FIXME(eddyb) long-term this kind of custom control-flow could be generalized
// to fully emulate unwinding (resulting in codegen similar to `?` in functions
// returning `Option` or `Result`), to e.g. run destructors, or even allow
// users to do `catch_unwind` at the top-level of their shader to handle
// panics specially (e.g. by appending to a custom buffer, or using some
// specific color in a fragment shader, to indicate a panic happened).
4 => Abort,
}
impl CustomOp {
/// Returns `true` iff this `CustomOp` is a custom debuginfo instruction,
/// i.e. non-semantic (can/must be ignored wherever `OpLine`/`OpNoLine` are).
pub fn is_debuginfo(self) -> bool {
match self {
CustomOp::SetDebugSrcLoc
| CustomOp::ClearDebugSrcLoc
| CustomOp::PushInlinedCallFrame
| CustomOp::PopInlinedCallFrame => true,
CustomOp::Abort => false,
}
}
/// Returns `true` iff this `CustomOp` is a custom terminator instruction,
/// i.e. semantic and must always appear just before an `OpUnreachable`
/// standard terminator (without even debuginfo in between the two).
pub fn is_terminator(self) -> bool {
match self {
CustomOp::SetDebugSrcLoc
| CustomOp::ClearDebugSrcLoc
| CustomOp::PushInlinedCallFrame
| CustomOp::PopInlinedCallFrame => false,
CustomOp::Abort => true,
}
}
}

View File

@ -5,6 +5,7 @@
//! run mem2reg (see mem2reg.rs) on the result to "unwrap" the Function pointer.
use super::apply_rewrite_rules;
use super::ipo::CallGraph;
use super::simple_passes::outgoing_edges;
use super::{get_name, get_names};
use crate::custom_insts::{self, CustomInst, CustomOp};
@ -30,6 +31,64 @@ pub fn inline(sess: &Session, module: &mut Module) -> super::Result<()> {
// This algorithm gets real sad if there's recursion - but, good news, SPIR-V bans recursion
deny_recursion_in_module(sess, module)?;
let custom_ext_inst_set_import = module
.ext_inst_imports
.iter()
.find(|inst| {
assert_eq!(inst.class.opcode, Op::ExtInstImport);
inst.operands[0].unwrap_literal_string() == &custom_insts::CUSTOM_EXT_INST_SET[..]
})
.map(|inst| inst.result_id.unwrap());
// HACK(eddyb) compute the set of functions that may `Abort` *transitively*,
// which is only needed because of how we inline (sometimes it's outside-in,
// aka top-down, instead of always being inside-out, aka bottom-up).
//
// (inlining is needed in the first place because our custom `Abort`
// instructions get lowered to a simple `OpReturn` in entry-points, but
// that requires that they get inlined all the way up to the entry-points)
let functions_that_may_abort = custom_ext_inst_set_import
.map(|custom_ext_inst_set_import| {
let mut may_abort_by_id = FxHashSet::default();
// FIXME(eddyb) use this `CallGraph` abstraction more during inlining.
let call_graph = CallGraph::collect(module);
for func_idx in call_graph.post_order() {
let func_id = module.functions[func_idx].def_id().unwrap();
let any_callee_may_abort = call_graph.callees[func_idx].iter().any(|&callee_idx| {
may_abort_by_id.contains(&module.functions[callee_idx].def_id().unwrap())
});
if any_callee_may_abort {
may_abort_by_id.insert(func_id);
continue;
}
let may_abort_directly = module.functions[func_idx].blocks.iter().any(|block| {
match &block.instructions[..] {
[.., last_normal_inst, terminator_inst]
if last_normal_inst.class.opcode == Op::ExtInst
&& last_normal_inst.operands[0].unwrap_id_ref()
== custom_ext_inst_set_import
&& CustomOp::decode_from_ext_inst(last_normal_inst)
== CustomOp::Abort =>
{
assert_eq!(terminator_inst.class.opcode, Op::Unreachable);
true
}
_ => false,
}
});
if may_abort_directly {
may_abort_by_id.insert(func_id);
}
}
may_abort_by_id
})
.unwrap_or_default();
let functions = module
.functions
.iter()
@ -42,7 +101,7 @@ pub fn inline(sess: &Session, module: &mut Module) -> super::Result<()> {
let mut dropped_ids = FxHashSet::default();
let mut inlined_to_legalize_dont_inlines = Vec::new();
module.functions.retain(|f| {
let should_inline_f = should_inline(&legal_globals, f, None);
let should_inline_f = should_inline(&legal_globals, &functions_that_may_abort, f, None);
if should_inline_f != Ok(false) {
if should_inline_f == Err(MustInlineToLegalize) && has_dont_inline(f) {
inlined_to_legalize_dont_inlines.push(f.def_id().unwrap());
@ -82,15 +141,7 @@ pub fn inline(sess: &Session, module: &mut Module) -> super::Result<()> {
id
}),
custom_ext_inst_set_import: module
.ext_inst_imports
.iter()
.find(|inst| {
assert_eq!(inst.class.opcode, Op::ExtInstImport);
inst.operands[0].unwrap_literal_string() == &custom_insts::CUSTOM_EXT_INST_SET[..]
})
.map(|inst| inst.result_id.unwrap())
.unwrap_or_else(|| {
custom_ext_inst_set_import: custom_ext_inst_set_import.unwrap_or_else(|| {
let id = next_id(header);
let inst = Instruction::new(
Op::ExtInstImport,
@ -125,6 +176,7 @@ pub fn inline(sess: &Session, module: &mut Module) -> super::Result<()> {
functions: &functions,
legal_globals: &legal_globals,
functions_that_may_abort: &functions_that_may_abort,
};
for function in &mut module.functions {
inliner.inline_fn(function);
@ -329,12 +381,20 @@ struct MustInlineToLegalize;
/// and inlining is *mandatory* due to an illegal signature/arguments.
fn should_inline(
legal_globals: &FxHashMap<Word, LegalGlobal>,
functions_that_may_abort: &FxHashSet<Word>,
callee: &Function,
call_site: Option<CallSite<'_>>,
) -> Result<bool, MustInlineToLegalize> {
let callee_def = callee.def.as_ref().unwrap();
let callee_control = callee_def.operands[0].unwrap_function_control();
// HACK(eddyb) this "has a call-site" check ensures entry-points don't get
// accidentally removed as "must inline to legalize" function, but can still
// be inlined into other entry-points (if such an unusual situation arises).
if call_site.is_some() && functions_that_may_abort.contains(&callee.def_id().unwrap()) {
return Err(MustInlineToLegalize);
}
let ret_ty = legal_globals
.get(&callee_def.result_type.unwrap())
.ok_or(MustInlineToLegalize)?;
@ -428,6 +488,7 @@ struct Inliner<'m, 'map> {
functions: &'map FunctionMap,
legal_globals: &'map FxHashMap<Word, LegalGlobal>,
functions_that_may_abort: &'map FxHashSet<Word>,
// rewrite_rules: FxHashMap<Word, Word>,
}
@ -509,7 +570,12 @@ impl Inliner<'_, '_> {
caller,
call_inst: inst,
};
match should_inline(self.legal_globals, f, Some(call_site)) {
match should_inline(
self.legal_globals,
self.functions_that_may_abort,
f,
Some(call_site),
) {
Ok(inline) => inline,
Err(MustInlineToLegalize) => true,
}
@ -655,16 +721,9 @@ impl Inliner<'_, '_> {
insts.retain_mut(|inst| {
let is_debuginfo = match inst.class.opcode {
Op::Line | Op::NoLine => true,
Op::ExtInst
if inst.operands[0].unwrap_id_ref()
== self.custom_ext_inst_set_import =>
{
match CustomOp::decode_from_ext_inst(inst) {
CustomOp::SetDebugSrcLoc
| CustomOp::ClearDebugSrcLoc
| CustomOp::PushInlinedCallFrame
| CustomOp::PopInlinedCallFrame => true,
}
Op::ExtInst => {
inst.operands[0].unwrap_id_ref() == self.custom_ext_inst_set_import
&& CustomOp::decode_from_ext_inst(inst).is_debuginfo()
}
_ => false,
};
@ -737,6 +796,12 @@ impl Inliner<'_, '_> {
return_variable: Option<Word>,
return_jump: Word,
) -> Vec<Block> {
let Self {
custom_ext_inst_set_import,
op_type_void_id,
..
} = *self;
// Prepare the debuginfo insts to prepend/append to every block.
// FIXME(eddyb) this could be more efficient if we only used one pair of
// `{Push,Pop}InlinedCallFrame` for the whole inlined callee, but there
@ -774,10 +839,10 @@ impl Inliner<'_, '_> {
let mut custom_inst_to_inst = |inst: CustomInst<_>| {
Instruction::new(
Op::ExtInst,
Some(self.op_type_void_id),
Some(op_type_void_id),
Some(self.id()),
[
Operand::IdRef(self.custom_ext_inst_set_import),
Operand::IdRef(custom_ext_inst_set_import),
Operand::LiteralExtInstInteger(inst.op() as u32),
]
.into_iter()
@ -837,7 +902,21 @@ impl Inliner<'_, '_> {
// Insert the suffix debuginfo instructions before the terminator,
// which sadly can't be covered by them.
{
let i = block.instructions.len() - 1;
let last_non_terminator = block.instructions.iter().rposition(|inst| {
let is_standard_terminator =
rspirv::grammar::reflect::is_block_terminator(inst.class.opcode);
let is_custom_terminator = match inst.class.opcode {
Op::ExtInst
if inst.operands[0].unwrap_id_ref()
== custom_ext_inst_set_import =>
{
CustomOp::decode_from_ext_inst(inst).is_terminator()
}
_ => false,
};
!(is_standard_terminator || is_custom_terminator)
});
let i = last_non_terminator.map_or(0, |x| x + 1);
i..i
},
debuginfo_suffix,

View File

@ -405,6 +405,12 @@ pub fn link(
};
after_pass("lower_from_spv", &module);
// NOTE(eddyb) this *must* run on unstructured CFGs, to do its job.
{
let _timer = sess.timer("spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points");
spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points(&mut module);
}
if opts.structurize {
{
let _timer = sess.timer("spirt::legalize::structurize_func_cfgs");
@ -473,7 +479,10 @@ pub fn link(
report_diagnostics_result?;
// Replace our custom debuginfo instructions just before lifting to SPIR-V.
{
let _timer = sess.timer("spirt_passes::debuginfo::convert_custom_debuginfo_to_spv");
spirt_passes::debuginfo::convert_custom_debuginfo_to_spv(&mut module);
}
let spv_words = {
let _timer = sess.timer("spirt::Module::lift_to_spv_module_emitter");

View File

@ -0,0 +1,69 @@
//! SPIR-T passes related to control-flow.
use crate::custom_insts::{self, CustomOp};
use spirt::{cfg, ControlNodeKind, DataInstKind, DeclDef, ExportKey, Exportee, Module, TypeCtor};
/// Replace our custom extended instruction `Abort`s with standard `OpReturn`s,
/// but only in entry-points (and only before CFG structurization).
pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(module: &mut Module) {
let cx = &module.cx();
let wk = &super::SpvSpecWithExtras::get().well_known;
let custom_ext_inst_set = cx.intern(&custom_insts::CUSTOM_EXT_INST_SET[..]);
for (export_key, exportee) in &module.exports {
let func = match (export_key, exportee) {
(ExportKey::SpvEntryPoint { .. }, &Exportee::Func(func)) => func,
_ => continue,
};
let func_decl = &mut module.funcs[func];
assert!(match &cx[func_decl.ret_type].ctor {
TypeCtor::SpvInst(spv_inst) => spv_inst.opcode == wk.OpTypeVoid,
_ => false,
});
let func_def_body = match &mut func_decl.def {
DeclDef::Present(def) => def,
DeclDef::Imported(_) => continue,
};
let rpo_regions = func_def_body
.unstructured_cfg
.as_ref()
.expect("Abort->OpReturn can only be done on unstructured CFGs")
.rev_post_order(func_def_body);
for region in rpo_regions {
let region_def = &func_def_body.control_regions[region];
let control_node_def = match region_def.children.iter().last {
Some(last_node) => &mut func_def_body.control_nodes[last_node],
_ => continue,
};
let block_insts = match &mut control_node_def.kind {
ControlNodeKind::Block { insts } => insts,
_ => continue,
};
let terminator = &mut func_def_body
.unstructured_cfg
.as_mut()
.unwrap()
.control_inst_on_exit_from[region];
if let cfg::ControlInstKind::Unreachable = terminator.kind {
let abort_inst = block_insts.iter().last.filter(|&last_inst| {
match func_def_body.data_insts[last_inst].kind {
DataInstKind::SpvExtInst { ext_set, inst } => {
ext_set == custom_ext_inst_set
&& CustomOp::decode(inst) == CustomOp::Abort
}
_ => false,
}
});
if let Some(abort_inst) = abort_inst {
block_insts.remove(abort_inst, &mut func_def_body.data_insts);
terminator.kind = cfg::ControlInstKind::Return;
}
}
}
}
}

View File

@ -85,7 +85,8 @@ impl Transformer for CustomDebuginfoToSpv<'_> {
} = data_inst_def.kind
{
if ext_set == self.custom_ext_inst_set {
match CustomOp::decode(ext_inst).with_operands(&data_inst_def.inputs) {
let custom_op = CustomOp::decode(ext_inst);
match custom_op.with_operands(&data_inst_def.inputs) {
CustomInst::SetDebugSrcLoc {
file,
line_start: line,
@ -126,6 +127,12 @@ impl Transformer for CustomDebuginfoToSpv<'_> {
insts_to_remove.push(inst);
continue;
}
CustomInst::Abort => {
assert!(
!custom_op.is_debuginfo(),
"`CustomOp::{custom_op:?}` debuginfo not lowered"
);
}
}
}
}

View File

@ -677,6 +677,7 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
_ => unreachable!(),
}
}
CustomInst::Abort => {}
},
}
}

View File

@ -1,5 +1,6 @@
//! SPIR-T pass infrastructure and supporting utilities.
pub(crate) mod controlflow;
pub(crate) mod debuginfo;
pub(crate) mod diagnostics;
mod fuse_selects;
@ -92,6 +93,8 @@ macro_rules! def_spv_spec_with_extra_well_known {
}
def_spv_spec_with_extra_well_known! {
opcode: spv::spec::Opcode = [
OpTypeVoid,
OpConstantComposite,
OpBitcast,

View File

@ -17,15 +17,6 @@ OpStore %21 %20
OpNoLine
OpBranch %13
%15 = OpLabel
OpBranch %22
%22 = OpLabel
OpLoopMerge %23 %24 None
OpBranch %25
%25 = OpLabel
OpBranch %24
%24 = OpLabel
OpBranch %22
%23 = OpLabel
OpBranch %13
%13 = OpLabel
OpReturn

View File

@ -15,15 +15,6 @@ OpLine %5 10 21
OpNoLine
OpBranch %14
%16 = OpLabel
OpBranch %21
%21 = OpLabel
OpLoopMerge %22 %23 None
OpBranch %24
%24 = OpLabel
OpBranch %23
%23 = OpLabel
OpBranch %21
%22 = OpLabel
OpBranch %14
%14 = OpLabel
OpReturn

View File

@ -17,15 +17,6 @@ OpStore %21 %20
OpNoLine
OpBranch %13
%15 = OpLabel
OpBranch %22
%22 = OpLabel
OpLoopMerge %23 %24 None
OpBranch %25
%25 = OpLabel
OpBranch %24
%24 = OpLabel
OpBranch %22
%23 = OpLabel
OpBranch %13
%13 = OpLabel
OpReturn