Add debugPrintf-based panic reporting, controlled via spirv_builder::ShaderPanicStrategy.

CHANGELOG.md

@@ -29,6 +29,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added ⭐
+- [PR#1080](https://github.com/EmbarkStudios/rust-gpu/pull/1080) added `debugPrintf`-based
+  panic reporting, with the desired behavior selected via `spirv_builder::ShaderPanicStrategy`
+  (see its documentation for more details about each available panic handling strategy)
+
 ### Changed 🛠
 - [PR#1079](https://github.com/EmbarkStudios/rust-gpu/pull/1079) revised `spirv-builder`'s `README.md`,
   and added a way for `docs.rs` to be able to build it (via `cargo +stable doc --no-default-features`)

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -13,7 +13,7 @@ use rustc_codegen_ssa::common::{
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::{
-    BackendTypes, BuilderMethods, ConstMethods, IntrinsicCallMethods, LayoutTypeMethods, OverflowOp,
+    BackendTypes, BuilderMethods, ConstMethods, LayoutTypeMethods, OverflowOp,
 };
 use rustc_codegen_ssa::MemFlags;
 use rustc_data_structures::fx::FxHashSet;
@@ -2647,7 +2647,8 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
 
             // HACK(eddyb) redirect any possible panic call to an abort, to avoid
             // needing to materialize `&core::panic::Location` or `format_args!`.
-            self.abort();
+            // FIXME(eddyb) find a way to extract the original message.
+            self.abort_with_message("panic!(...)".into());
             self.undef(result_type)
         } else if let Some(mode) = buffer_load_intrinsic {
             self.codegen_buffer_load_intrinsic(result_type, args, mode)

crates/rustc_codegen_spirv/src/builder/intrinsics.rs

@@ -4,6 +4,7 @@ use crate::builder_spirv::{SpirvValue, SpirvValueExt};
 use crate::codegen_cx::CodegenCx;
 use crate::custom_insts::CustomInst;
 use crate::spirv_type::SpirvType;
+use rspirv::dr::Operand;
 use rspirv::spirv::GLOp;
 use rustc_codegen_ssa::mir::operand::OperandRef;
 use rustc_codegen_ssa::mir::place::PlaceRef;
@@ -338,18 +339,7 @@ impl<'a, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'tcx> {
     }
 
     fn abort(&mut self) {
-        // FIXME(eddyb) this should be cached more efficiently.
+        self.abort_with_message("intrinsics::abort()".into());
-        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 use our own custom instruction.
-        self.custom_inst(void_ty, CustomInst::Abort);
-        self.unreachable();
-
-        // 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) {
@@ -382,3 +372,26 @@ impl<'a, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'tcx> {
         todo!()
     }
 }
+
+impl Builder<'_, '_> {
+    pub fn abort_with_message(&mut self, message: String) {
+        // 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 use our own custom instruction.
+        let message_id = self.emit().string(message);
+        self.custom_inst(
+            void_ty,
+            CustomInst::Abort {
+                message: Operand::IdRef(message_id),
+            },
+        );
+        self.unreachable();
+
+        // 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);
+    }
+}

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -348,6 +348,12 @@ impl CodegenArgs {
                 "enable additional SPIR-T passes (comma-separated)",
                 "PASSES",
             );
+            opts.optopt(
+                "",
+                "abort-strategy",
+                "select a non-default abort (i.e. panic) strategy - see `spirv-builder` docs",
+                "STRATEGY",
+            );
 
             // NOTE(eddyb) these are debugging options that used to be env vars
             // (for more information see `docs/src/codegen-args.md`).
@@ -529,6 +535,8 @@ impl CodegenArgs {
                 .map(|s| s.to_string())
                 .collect(),
 
+            abort_strategy: matches.opt_str("abort-strategy"),
+
             // FIXME(eddyb) deduplicate between `CodegenArgs` and `linker::Options`.
             emit_multiple_modules: module_output_type == ModuleOutputType::Multiple,
             spirv_metadata,

crates/rustc_codegen_spirv/src/custom_insts.rs

@@ -138,7 +138,7 @@ def_custom_insts! {
     // 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).
+    // with at most debuginfo instructions (standard or custom), between the two.
     //
     // FIXME(eddyb) long-term this kind of custom control-flow could be generalized
     // to fully emulate unwinding (resulting in codegen similar to `?` in functions
@@ -146,7 +146,7 @@ def_custom_insts! {
     // 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,
+    4 => Abort { message },
 }
 
 impl CustomOp {
@@ -164,8 +164,8 @@ impl CustomOp {
     }
 
     /// Returns `true` iff this `CustomOp` is a custom terminator instruction,
-    /// i.e. semantic and must always appear just before an `OpUnreachable`
+    /// i.e. semantic and must precede an `OpUnreachable` standard terminator,
-    /// standard terminator (without even debuginfo in between the two).
+    /// with at most debuginfo instructions (standard or custom), between the two.
     pub fn is_terminator(self) -> bool {
         match self {
             CustomOp::SetDebugSrcLoc

crates/rustc_codegen_spirv/src/linker/inline.rs

@@ -594,19 +594,14 @@ impl Inliner<'_, '_> {
         };
         let call_result_id = call_inst.result_id.unwrap();
 
-        // Get the debuginfo source location instruction that applies to the call.
+        // Get the debuginfo instructions that apply to the call.
-        let call_debug_src_loc_inst = caller.blocks[block_idx].instructions[..call_index]
+        let custom_ext_inst_set_import = self.custom_ext_inst_set_import;
+        let call_debug_insts = caller.blocks[block_idx].instructions[..call_index]
             .iter()
-            .rev()
+            .filter(|inst| match inst.class.opcode {
-            .find(|inst| match inst.class.opcode {
                 Op::Line | Op::NoLine => true,
-                Op::ExtInst
+                Op::ExtInst if inst.operands[0].unwrap_id_ref() == custom_ext_inst_set_import => {
-                    if inst.operands[0].unwrap_id_ref() == self.custom_ext_inst_set_import =>
+                    CustomOp::decode_from_ext_inst(inst).is_debuginfo()
-                {
-                    matches!(
-                        CustomOp::decode_from_ext_inst(inst),
-                        CustomOp::SetDebugSrcLoc | CustomOp::ClearDebugSrcLoc
-                    )
                 }
                 _ => false,
             });
@@ -630,12 +625,8 @@ impl Inliner<'_, '_> {
         };
         let return_jump = self.id();
         // Rewrite OpReturns of the callee.
-        let mut inlined_callee_blocks = self.get_inlined_blocks(
+        let (mut inlined_callee_blocks, extra_debug_insts_pre_call, extra_debug_insts_post_call) =
-            callee,
+            self.get_inlined_blocks(callee, call_debug_insts, return_variable, return_jump);
-            call_debug_src_loc_inst,
-            return_variable,
-            return_jump,
-        );
         // Clone the IDs of the callee, because otherwise they'd be defined multiple times if the
         // fn is inlined multiple times.
         self.add_clone_id_rules(&mut rewrite_rules, &inlined_callee_blocks);
@@ -649,6 +640,7 @@ impl Inliner<'_, '_> {
         let mut post_call_block_insts = caller.blocks[pre_call_block_idx]
             .instructions
             .split_off(call_index + 1);
+
         // pop off OpFunctionCall
         let call = caller.blocks[pre_call_block_idx]
             .instructions
@@ -656,6 +648,13 @@ impl Inliner<'_, '_> {
             .unwrap();
         assert!(call.class.opcode == Op::FunctionCall);
 
+        // HACK(eddyb) inject the additional debuginfo instructions generated by
+        // `get_inlined_blocks`, so the inlined call frame "stack" isn't corrupted.
+        caller.blocks[pre_call_block_idx]
+            .instructions
+            .extend(extra_debug_insts_pre_call);
+        post_call_block_insts.splice(0..0, extra_debug_insts_post_call);
+
         if let Some(call_result_type) = call_result_type {
             // Generate the storage space for the return value: Do this *after* the split above,
             // because if block_idx=0, inserting a variable here shifts call_index.
@@ -789,19 +788,58 @@ impl Inliner<'_, '_> {
         }
     }
 
-    fn get_inlined_blocks(
+    // HACK(eddyb) the second and third return values are additional debuginfo
+    // instructions that need to be inserted just before/after the callsite.
+    fn get_inlined_blocks<'a>(
         &mut self,
         callee: &Function,
-        call_debug_src_loc_inst: Option<&Instruction>,
+        call_debug_insts: impl Iterator<Item = &'a Instruction>,
         return_variable: Option<Word>,
         return_jump: Word,
-    ) -> Vec<Block> {
+    ) -> (
+        Vec<Block>,
+        SmallVec<[Instruction; 8]>,
+        SmallVec<[Instruction; 8]>,
+    ) {
         let Self {
             custom_ext_inst_set_import,
             op_type_void_id,
             ..
         } = *self;
 
+        // HACK(eddyb) this is terrible, but we have to deal with it becasue of
+        // how this inliner is outside-in, instead of inside-out, meaning that
+        // context builds up "outside" of the callee blocks, inside the caller.
+        let mut enclosing_inlined_frames = SmallVec::<[_; 8]>::new();
+        let mut current_debug_src_loc_inst = None;
+        for inst in call_debug_insts {
+            match inst.class.opcode {
+                Op::Line => current_debug_src_loc_inst = Some(inst),
+                Op::NoLine => current_debug_src_loc_inst = None,
+                Op::ExtInst
+                    if inst.operands[0].unwrap_id_ref() == self.custom_ext_inst_set_import =>
+                {
+                    match CustomOp::decode_from_ext_inst(inst) {
+                        CustomOp::SetDebugSrcLoc => current_debug_src_loc_inst = Some(inst),
+                        CustomOp::ClearDebugSrcLoc => current_debug_src_loc_inst = None,
+                        CustomOp::PushInlinedCallFrame => {
+                            enclosing_inlined_frames
+                                .push((current_debug_src_loc_inst.take(), inst));
+                        }
+                        CustomOp::PopInlinedCallFrame => {
+                            if let Some((callsite_debug_src_loc_inst, _)) =
+                                enclosing_inlined_frames.pop()
+                            {
+                                current_debug_src_loc_inst = callsite_debug_src_loc_inst;
+                            }
+                        }
+                        CustomOp::Abort => {}
+                    }
+                }
+                _ => {}
+            }
+        }
+
         // 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
@@ -826,21 +864,21 @@ impl Inliner<'_, '_> {
                 ));
                 id
             });
-        let mut mk_debuginfo_prefix_and_suffix = || {
+        let mut mk_debuginfo_prefix_and_suffix = |include_callee_frame| {
             // NOTE(eddyb) `OpExtInst`s have a result ID, even if unused, and
             // it has to be unique (same goes for the other instructions below).
-            let instantiated_src_loc_inst = call_debug_src_loc_inst.map(|inst| {
+            let instantiate_debuginfo = |this: &mut Self, inst: &Instruction| {
                 let mut inst = inst.clone();
                 if let Some(id) = &mut inst.result_id {
-                    *id = self.id();
+                    *id = this.id();
                 }
                 inst
-            });
+            };
-            let mut custom_inst_to_inst = |inst: CustomInst<_>| {
+            let custom_inst_to_inst = |this: &mut Self, inst: CustomInst<_>| {
                 Instruction::new(
                     Op::ExtInst,
                     Some(op_type_void_id),
-                    Some(self.id()),
+                    Some(this.id()),
                     [
                         Operand::IdRef(custom_ext_inst_set_import),
                         Operand::LiteralExtInstInteger(inst.op() as u32),
@@ -850,14 +888,33 @@ impl Inliner<'_, '_> {
                     .collect(),
                 )
             };
-            (
+            // FIXME(eddyb) this only allocates to avoid borrow conflicts.
-                instantiated_src_loc_inst.into_iter().chain([{
+            let mut prefix = SmallVec::<[_; 8]>::new();
-                    custom_inst_to_inst(CustomInst::PushInlinedCallFrame {
+            let mut suffix = SmallVec::<[_; 8]>::new();
+            for &(callsite_debug_src_loc_inst, push_inlined_call_frame_inst) in
+                &enclosing_inlined_frames
+            {
+                prefix.extend(
+                    callsite_debug_src_loc_inst
+                        .into_iter()
+                        .chain([push_inlined_call_frame_inst])
+                        .map(|inst| instantiate_debuginfo(self, inst)),
+                );
+                suffix.push(custom_inst_to_inst(self, CustomInst::PopInlinedCallFrame));
+            }
+            prefix.extend(current_debug_src_loc_inst.map(|inst| instantiate_debuginfo(self, inst)));
+
+            if include_callee_frame {
+                prefix.push(custom_inst_to_inst(
+                    self,
+                    CustomInst::PushInlinedCallFrame {
                         callee_name: Operand::IdRef(callee_name_id),
-                    })
+                    },
-                }]),
+                ));
-                [custom_inst_to_inst(CustomInst::PopInlinedCallFrame)].into_iter(),
+                suffix.push(custom_inst_to_inst(self, CustomInst::PopInlinedCallFrame));
-            )
+            }
+
+            (prefix, suffix)
         };
 
         let mut blocks = callee.blocks.clone();
@@ -883,47 +940,43 @@ impl Inliner<'_, '_> {
                 }
                 *block.instructions.last_mut().unwrap() =
                     Instruction::new(Op::Branch, None, None, vec![Operand::IdRef(return_jump)]);
-
-                let (debuginfo_prefix, debuginfo_suffix) = mk_debuginfo_prefix_and_suffix();
-                block.instructions.splice(
-                    // Insert the prefix debuginfo instructions after `OpPhi`s,
-                    // which sadly can't be covered by them.
-                    {
-                        let i = block
-                            .instructions
-                            .iter()
-                            .position(|inst| inst.class.opcode != Op::Phi)
-                            .unwrap();
-                        i..i
-                    },
-                    debuginfo_prefix,
-                );
-                block.instructions.splice(
-                    // Insert the suffix debuginfo instructions before the terminator,
-                    // which sadly can't be covered by them.
-                    {
-                        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,
-                );
             }
+
+            let (debuginfo_prefix, debuginfo_suffix) = mk_debuginfo_prefix_and_suffix(true);
+            block.instructions.splice(
+                // Insert the prefix debuginfo instructions after `OpPhi`s,
+                // which sadly can't be covered by them.
+                {
+                    let i = block
+                        .instructions
+                        .iter()
+                        .position(|inst| inst.class.opcode != Op::Phi)
+                        .unwrap();
+                    i..i
+                },
+                debuginfo_prefix,
+            );
+            block.instructions.splice(
+                // Insert the suffix debuginfo instructions before the terminator,
+                // which sadly can't be covered by them.
+                {
+                    let last_non_terminator = block.instructions.iter().rposition(|inst| {
+                        !rspirv::grammar::reflect::is_block_terminator(inst.class.opcode)
+                    });
+                    let i = last_non_terminator.map_or(0, |x| x + 1);
+                    i..i
+                },
+                debuginfo_suffix,
+            );
         }
-        blocks
+
+        let (caller_restore_debuginfo_after_call, calleer_reset_debuginfo_before_call) =
+            mk_debuginfo_prefix_and_suffix(false);
+        (
+            blocks,
+            calleer_reset_debuginfo_before_call,
+            caller_restore_debuginfo_after_call,
+        )
     }
 }
 

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -43,6 +43,8 @@ pub struct Options {
     pub structurize: bool,
     pub spirt_passes: Vec<String>,
 
+    pub abort_strategy: Option<String>,
+
     pub emit_multiple_modules: bool,
     pub spirv_metadata: SpirvMetadata,
 
@@ -408,7 +410,7 @@ pub fn link(
         // 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);
+            spirt_passes::controlflow::convert_custom_aborts_to_unstructured_returns_in_entry_points(opts, &mut module);
         }
 
         if opts.structurize {

crates/rustc_codegen_spirv/src/linker/spirt_passes/controlflow.rs

@@ -1,19 +1,67 @@
 //! SPIR-T passes related to control-flow.
 
-use crate::custom_insts::{self, CustomOp};
+use crate::custom_insts::{self, CustomInst, CustomOp};
-use spirt::{cfg, ControlNodeKind, DataInstKind, DeclDef, ExportKey, Exportee, Module, TypeCtor};
+use smallvec::SmallVec;
+use spirt::func_at::FuncAt;
+use spirt::{
+    cfg, spv, Attr, AttrSet, ConstCtor, ConstDef, ControlNodeKind, DataInstKind, DeclDef,
+    EntityDefs, ExportKey, Exportee, Module, Type, TypeCtor, TypeCtorArg, TypeDef, Value,
+};
+use std::fmt::Write as _;
 
 /// 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) {
+pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(
+    linker_options: &crate::linker::Options,
+    module: &mut Module,
+) {
+    // HACK(eddyb) this shouldn't be the place to parse `abort_strategy`.
+    enum Strategy {
+        Unreachable,
+        DebugPrintf { inputs: bool, backtrace: bool },
+    }
+    let abort_strategy = linker_options.abort_strategy.as_ref().map(|s| {
+        if s == "unreachable" {
+            return Strategy::Unreachable;
+        }
+        if let Some(s) = s.strip_prefix("debug-printf") {
+            let (inputs, s) = s.strip_prefix("+inputs").map_or((false, s), |s| (true, s));
+            let (backtrace, s) = s
+                .strip_prefix("+backtrace")
+                .map_or((false, s), |s| (true, s));
+            if s.is_empty() {
+                return Strategy::DebugPrintf { inputs, backtrace };
+            }
+        }
+        panic!("unknown `--abort-strategy={s}");
+    });
+
     let cx = &module.cx();
     let wk = &super::SpvSpecWithExtras::get().well_known;
 
+    // HACK(eddyb) deduplicate with `diagnostics`.
+    let name_from_attrs = |attrs: AttrSet| {
+        cx[attrs].attrs.iter().find_map(|attr| match attr {
+            Attr::SpvAnnotation(spv_inst) if spv_inst.opcode == wk.OpName => Some(
+                super::diagnostics::decode_spv_lit_str_with(&spv_inst.imms, |name| {
+                    name.to_string()
+                }),
+            ),
+            _ => None,
+        })
+    };
+
     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) {
+        let (entry_point_imms, interface_global_vars, func) = match (export_key, exportee) {
-            (ExportKey::SpvEntryPoint { .. }, &Exportee::Func(func)) => func,
+            (
+                ExportKey::SpvEntryPoint {
+                    imms,
+                    interface_global_vars,
+                },
+                &Exportee::Func(func),
+            ) => (imms, interface_global_vars, func),
             _ => continue,
         };
 
@@ -28,6 +76,114 @@ pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(module: &mu
             DeclDef::Imported(_) => continue,
         };
 
+        let debug_printf_context_fmt_str;
+        let mut debug_printf_context_inputs = SmallVec::<[_; 4]>::new();
+        if let Some(Strategy::DebugPrintf { inputs, .. }) = abort_strategy {
+            let mut fmt = String::new();
+
+            match entry_point_imms[..] {
+                [spv::Imm::Short(em_kind, _), ref name_imms @ ..] => {
+                    assert_eq!(em_kind, wk.ExecutionModel);
+                    super::diagnostics::decode_spv_lit_str_with(name_imms, |name| {
+                        fmt += &name.replace('%', "%%");
+                    });
+                }
+                _ => unreachable!(),
+            }
+            fmt += "(";
+
+            // Collect entry-point inputs `OpLoad`ed by the entry block.
+            // HACK(eddyb) this relies on Rust-GPU always eagerly loading inputs.
+            let loaded_inputs = func_def_body
+                .at(func_def_body
+                    .at_body()
+                    .at_children()
+                    .into_iter()
+                    .next()
+                    .and_then(|func_at_first_node| match func_at_first_node.def().kind {
+                        ControlNodeKind::Block { insts } => Some(insts),
+                        _ => None,
+                    })
+                    .unwrap_or_default())
+                .into_iter()
+                .filter_map(|func_at_inst| {
+                    let data_inst_def = func_at_inst.def();
+                    if let DataInstKind::SpvInst(spv_inst) = &data_inst_def.kind {
+                        if spv_inst.opcode == wk.OpLoad {
+                            if let Value::Const(ct) = data_inst_def.inputs[0] {
+                                if let ConstCtor::PtrToGlobalVar(gv) = cx[ct].ctor {
+                                    if interface_global_vars.contains(&gv) {
+                                        return Some((
+                                            gv,
+                                            data_inst_def.output_type.unwrap(),
+                                            Value::DataInstOutput(func_at_inst.position),
+                                        ));
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    None
+                });
+            if inputs {
+                let mut first_input = true;
+                for (gv, ty, value) in loaded_inputs {
+                    let scalar_type = |ty: Type| match &cx[ty].ctor {
+                        TypeCtor::SpvInst(spv_inst) => match spv_inst.imms[..] {
+                            [spv::Imm::Short(_, 32), spv::Imm::Short(_, signedness)]
+                                if spv_inst.opcode == wk.OpTypeInt =>
+                            {
+                                Some(if signedness != 0 { "i" } else { "u" })
+                            }
+                            [spv::Imm::Short(_, 32)] if spv_inst.opcode == wk.OpTypeFloat => {
+                                Some("f")
+                            }
+                            _ => None,
+                        },
+                        _ => None,
+                    };
+                    let vector_or_scalar_type = |ty: Type| {
+                        let ty_def = &cx[ty];
+                        match (&ty_def.ctor, &ty_def.ctor_args[..]) {
+                            (TypeCtor::SpvInst(spv_inst), &[TypeCtorArg::Type(elem)])
+                                if spv_inst.opcode == wk.OpTypeVector =>
+                            {
+                                match spv_inst.imms[..] {
+                                    [spv::Imm::Short(_, vlen @ 2..=4)] => {
+                                        Some((scalar_type(elem)?, Some(vlen)))
+                                    }
+                                    _ => None,
+                                }
+                            }
+                            _ => Some((scalar_type(ty)?, None)),
+                        }
+                    };
+                    if let Some((scalar_fmt, vlen)) = vector_or_scalar_type(ty) {
+                        if !first_input {
+                            fmt += ", ";
+                        }
+                        first_input = false;
+
+                        if let Some(name) = name_from_attrs(module.global_vars[gv].attrs) {
+                            fmt += &name.replace('%', "%%");
+                            fmt += " = ";
+                        }
+                        match vlen {
+                            Some(vlen) => write!(fmt, "vec{vlen}(%v{vlen}{scalar_fmt})").unwrap(),
+                            None => write!(fmt, "%{scalar_fmt}").unwrap(),
+                        }
+                        debug_printf_context_inputs.push(value);
+                    }
+                }
+            }
+
+            fmt += ")";
+
+            debug_printf_context_fmt_str = fmt;
+        } else {
+            debug_printf_context_fmt_str = String::new();
+        }
+
         let rpo_regions = func_def_body
             .unstructured_cfg
             .as_ref()
@@ -49,20 +205,177 @@ pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(module: &mu
                 .as_mut()
                 .unwrap()
                 .control_inst_on_exit_from[region];
-            if let cfg::ControlInstKind::Unreachable = terminator.kind {
+            match terminator.kind {
-                let abort_inst = block_insts.iter().last.filter(|&last_inst| {
+                cfg::ControlInstKind::Unreachable => {}
-                    match func_def_body.data_insts[last_inst].kind {
+                _ => continue,
-                        DataInstKind::SpvExtInst { ext_set, inst } => {
+            }
-                            ext_set == custom_ext_inst_set
+
-                                && CustomOp::decode(inst) == CustomOp::Abort
+            // HACK(eddyb) this allows accessing the `DataInst` iterator while
+            // mutably borrowing other parts of `FuncDefBody`.
+            let func_at_block_insts = FuncAt {
+                control_nodes: &EntityDefs::new(),
+                control_regions: &EntityDefs::new(),
+                data_insts: &func_def_body.data_insts,
+
+                position: *block_insts,
+            };
+            let block_insts_maybe_custom = func_at_block_insts.into_iter().map(|func_at_inst| {
+                let data_inst_def = func_at_inst.def();
+                (
+                    func_at_inst,
+                    match data_inst_def.kind {
+                        DataInstKind::SpvExtInst { ext_set, inst }
+                            if ext_set == custom_ext_inst_set =>
+                        {
+                            Some(CustomOp::decode(inst).with_operands(&data_inst_def.inputs))
                         }
-                        _ => false,
+                        _ => None,
+                    },
+                )
+            });
+            let custom_terminator_inst = block_insts_maybe_custom
+                .clone()
+                .rev()
+                .take_while(|(_, custom)| custom.is_some())
+                .map(|(func_at_inst, custom)| (func_at_inst, custom.unwrap()))
+                .find(|(_, custom)| !custom.op().is_debuginfo())
+                .filter(|(_, custom)| custom.op().is_terminator());
+            if let Some((func_at_abort_inst, CustomInst::Abort { message })) =
+                custom_terminator_inst
+            {
+                let abort_inst = func_at_abort_inst.position;
+                terminator.kind = cfg::ControlInstKind::Return;
+
+                match abort_strategy {
+                    Some(Strategy::Unreachable) => {
+                        terminator.kind = cfg::ControlInstKind::Unreachable;
                     }
-                });
+                    Some(Strategy::DebugPrintf {
-                if let Some(abort_inst) = abort_inst {
+                        inputs: _,
-                    block_insts.remove(abort_inst, &mut func_def_body.data_insts);
+                        backtrace,
-                    terminator.kind = cfg::ControlInstKind::Return;
+                    }) => {
+                        let const_ctor = |v: Value| match v {
+                            Value::Const(ct) => &cx[ct].ctor,
+                            _ => unreachable!(),
+                        };
+                        let const_str = |v: Value| match const_ctor(v) {
+                            &ConstCtor::SpvStringLiteralForExtInst(s) => s,
+                            _ => unreachable!(),
+                        };
+                        let const_u32 = |v: Value| match const_ctor(v) {
+                            ConstCtor::SpvInst(spv_inst) => {
+                                assert!(spv_inst.opcode == wk.OpConstant);
+                                match spv_inst.imms[..] {
+                                    [spv::Imm::Short(_, x)] => x,
+                                    _ => unreachable!(),
+                                }
+                            }
+                            _ => unreachable!(),
+                        };
+                        let mk_const_str = |s| {
+                            cx.intern(ConstDef {
+                                attrs: Default::default(),
+                                ty: cx.intern(TypeDef {
+                                    attrs: Default::default(),
+                                    ctor: TypeCtor::SpvStringLiteralForExtInst,
+                                    ctor_args: Default::default(),
+                                }),
+                                ctor: ConstCtor::SpvStringLiteralForExtInst(s),
+                                ctor_args: Default::default(),
+                            })
+                        };
+
+                        let mut current_debug_src_loc = None;
+                        let mut call_stack = SmallVec::<[_; 8]>::new();
+                        let block_insts_custom = block_insts_maybe_custom
+                            .filter_map(|(func_at_inst, custom)| Some((func_at_inst, custom?)));
+                        for (func_at_inst, custom) in block_insts_custom {
+                            // Stop at the abort, that we don't undo its debug context.
+                            if func_at_inst.position == abort_inst {
+                                break;
+                            }
+
+                            match custom {
+                                CustomInst::SetDebugSrcLoc {
+                                    file,
+                                    line_start,
+                                    line_end: _,
+                                    col_start,
+                                    col_end: _,
+                                } => {
+                                    current_debug_src_loc = Some((
+                                        &cx[const_str(file)],
+                                        const_u32(line_start),
+                                        const_u32(col_start),
+                                    ));
+                                }
+                                CustomInst::ClearDebugSrcLoc => current_debug_src_loc = None,
+                                CustomInst::PushInlinedCallFrame { callee_name } => {
+                                    if backtrace {
+                                        call_stack.push((
+                                            current_debug_src_loc.take(),
+                                            const_str(callee_name),
+                                        ));
+                                    }
+                                }
+                                CustomInst::PopInlinedCallFrame => {
+                                    if let Some((callsite_debug_src_loc, _)) = call_stack.pop() {
+                                        current_debug_src_loc = callsite_debug_src_loc;
+                                    }
+                                }
+                                CustomInst::Abort { .. } => {}
+                            }
+                        }
+
+                        let mut fmt = String::new();
+
+                        // HACK(eddyb) this improves readability w/ very verbose Vulkan loggers.
+                        fmt += "\n  ";
+
+                        if let Some((file, line, col)) = current_debug_src_loc.take() {
+                            fmt += &format!("{file}:{line}:{col}: ").replace('%', "%%");
+                        }
+                        fmt += &cx[const_str(message)].replace('%', "%%");
+
+                        let mut innermost = true;
+                        let mut append_call = |callsite_debug_src_loc, callee: &str| {
+                            if innermost {
+                                innermost = false;
+                                fmt += "\n    in ";
+                            } else if current_debug_src_loc.is_some() {
+                                fmt += "\n    by ";
+                            } else {
+                                // HACK(eddyb) previous call didn't have a `called at` line.
+                                fmt += "\n    called by ";
+                            }
+                            fmt += callee;
+                            if let Some((file, line, col)) = callsite_debug_src_loc {
+                                fmt += &format!("\n      called at {file}:{line}:{col}")
+                                    .replace('%', "%%");
+                            }
+                            current_debug_src_loc = callsite_debug_src_loc;
+                        };
+                        while let Some((callsite_debug_src_loc, callee)) = call_stack.pop() {
+                            append_call(callsite_debug_src_loc, &cx[callee].replace('%', "%%"));
+                        }
+                        append_call(None, &debug_printf_context_fmt_str);
+
+                        let abort_inst_def = &mut func_def_body.data_insts[abort_inst];
+                        abort_inst_def.kind = DataInstKind::SpvExtInst {
+                            ext_set: cx.intern("NonSemantic.DebugPrintf"),
+                            inst: 1,
+                        };
+                        abort_inst_def.inputs = [Value::Const(mk_const_str(cx.intern(fmt)))]
+                            .into_iter()
+                            .chain(debug_printf_context_inputs.iter().copied())
+                            .collect();
+
+                        // Avoid removing the instruction we just replaced.
+                        continue;
+                    }
+                    None => {}
                 }
+                block_insts.remove(abort_inst, &mut func_def_body.data_insts);
             }
         }
     }

crates/rustc_codegen_spirv/src/linker/spirt_passes/debuginfo.rs

@@ -127,7 +127,7 @@ impl Transformer for CustomDebuginfoToSpv<'_> {
                                 insts_to_remove.push(inst);
                                 continue;
                             }
-                            CustomInst::Abort => {
+                            CustomInst::Abort { .. } => {
                                 assert!(
                                     !custom_op.is_debuginfo(),
                                     "`CustomOp::{custom_op:?}` debuginfo not lowered"

crates/rustc_codegen_spirv/src/linker/spirt_passes/diagnostics.rs

@@ -107,7 +107,7 @@ impl<D> LazilyDecoded<D> {
     }
 }
 
-fn decode_spv_lit_str_with<R>(imms: &[spv::Imm], f: impl FnOnce(&str) -> R) -> R {
+pub(super) fn decode_spv_lit_str_with<R>(imms: &[spv::Imm], f: impl FnOnce(&str) -> R) -> R {
     let wk = &super::SpvSpecWithExtras::get().well_known;
 
     // FIXME(eddyb) deduplicate with `spirt::spv::extract_literal_string`.
@@ -578,8 +578,14 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
     }
 
     fn visit_control_node_def(&mut self, func_at_control_node: FuncAt<'a, ControlNode>) {
+        let original_use_stack_len = self.use_stack.len();
+
         func_at_control_node.inner_visit_with(self);
 
+        // HACK(eddyb) avoid `use_stack` from growing due to having some
+        // `PushInlinedCallFrame` without matching `PopInlinedCallFrame`.
+        self.use_stack.truncate(original_use_stack_len);
+
         // HACK(eddyb) this relies on the fact that `ControlNodeKind::Block` maps
         // to one original SPIR-V block, which may not necessarily be true, and
         // steps should be taken elsewhere to explicitly unset debuginfo, instead
@@ -677,7 +683,7 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
                                         _ => unreachable!(),
                                     }
                                 }
-                                CustomInst::Abort => {}
+                                CustomInst::Abort { .. } => {}
                             },
                         }
                     }

crates/spirv-builder/src/lib.rs

@@ -174,6 +174,62 @@ pub enum SpirvMetadata {
     Full,
 }
 
+/// Strategy used to handle Rust `panic!`s in shaders compiled to SPIR-V.
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum ShaderPanicStrategy {
+    /// Return from shader entry-point with no side-effects **(default)**.
+    ///
+    /// While similar to the standard SPIR-V `OpTerminateInvocation`, this is
+    /// *not* limited to fragment shaders, and instead supports all shaders
+    /// (as it's handled via control-flow rewriting, instead of SPIR-V features).
+    SilentExit,
+
+    /// Like `SilentExit`, but also using `debugPrintf` to report the panic in
+    /// a way that can reach the user, before returning from the entry-point.
+    ///
+    /// Will automatically require the `SPV_KHR_non_semantic_info` extension,
+    /// as `debugPrintf` uses a "non-semantic extended instruction set".
+    ///
+    /// If you have multiple entry-points, you *may* need to also enable the
+    /// `multimodule` node (see <https://github.com/KhronosGroup/SPIRV-Tools/issues/4892>).
+    ///
+    /// **Note**: actually obtaining the `debugPrintf` output requires enabling:
+    /// * `VK_KHR_shader_non_semantic_info` Vulkan *Device* extension
+    /// * Vulkan Validation Layers (which contain the `debugPrintf` implementation)
+    /// * `VK_VALIDATION_FEATURE_ENABLE_DEBUG_PRINTF_EXT` in Validation Layers,
+    ///   either by using `VkValidationFeaturesEXT` during instance creating,
+    ///   setting the `VK_LAYER_ENABLES` environment variable to that value,
+    ///   or adding it to `khronos_validation.enables` in `vk_layer_settings.txt`
+    ///
+    /// See also: <https://github.com/KhronosGroup/Vulkan-ValidationLayers/blob/main/docs/debug_printf.md>.
+    DebugPrintfThenExit {
+        /// Whether to also print the entry-point inputs (excluding buffers/resources),
+        /// which should uniquely identify the panicking shader invocation.
+        print_inputs: bool,
+
+        /// Whether to also print a "backtrace" (i.e. the chain of function calls
+        /// that led to the `panic!).
+        ///
+        /// As there is no way to dynamically compute this information, the string
+        /// containing the full backtrace of each `panic!` is statically generated,
+        /// meaning this option could significantly increase binary size.
+        print_backtrace: bool,
+    },
+
+    /// **Warning**: this is _**unsound**_ (i.e. adds Undefined Behavior to *safe* Rust code)
+    ///
+    /// This option only exists for testing (hence the unfriendly name it has),
+    /// and more specifically testing whether conditional panics are responsible
+    /// for performance differences when upgrading from older Rust-GPU versions
+    /// (which used infinite loops for panics, that `spirv-opt`/drivers could've
+    /// sometimes treated as UB, and optimized as if they were impossible to reach).
+    ///
+    /// Unlike those infinite loops, however, this uses `OpUnreachable`, so it
+    /// forces the old worst-case (all `panic!`s become UB and are optimized out).
+    #[allow(non_camel_case_types)]
+    UNSOUND_DO_NOT_USE_UndefinedBehaviorViaUnreachable,
+}
+
 pub struct SpirvBuilder {
     path_to_crate: PathBuf,
     print_metadata: MetadataPrintout,
@@ -186,6 +242,9 @@ pub struct SpirvBuilder {
     extensions: Vec<String>,
     extra_args: Vec<String>,
 
+    // `rustc_codegen_spirv::linker` codegen args
+    pub shader_panic_strategy: ShaderPanicStrategy,
+
     // spirv-val flags
     pub relax_struct_store: bool,
     pub relax_logical_pointer: bool,
@@ -212,6 +271,8 @@ impl SpirvBuilder {
             extensions: Vec::new(),
             extra_args: Vec::new(),
 
+            shader_panic_strategy: ShaderPanicStrategy::SilentExit,
+
             relax_struct_store: false,
             relax_logical_pointer: false,
             relax_block_layout: false,
@@ -276,6 +337,13 @@ impl SpirvBuilder {
         self
     }
 
+    /// Change the shader `panic!` handling strategy (see [`ShaderPanicStrategy`]).
+    #[must_use]
+    pub fn shader_panic_strategy(mut self, shader_panic_strategy: ShaderPanicStrategy) -> Self {
+        self.shader_panic_strategy = shader_panic_strategy;
+        self
+    }
+
     /// Allow store from one struct type to a different type with compatible layout and members.
     #[must_use]
     pub fn relax_struct_store(mut self, v: bool) -> Self {
@@ -511,6 +579,25 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
     if builder.preserve_bindings {
         llvm_args.push("--preserve-bindings".to_string());
     }
+    let mut target_features = vec![];
+    let abort_strategy = match builder.shader_panic_strategy {
+        ShaderPanicStrategy::SilentExit => None,
+        ShaderPanicStrategy::DebugPrintfThenExit {
+            print_inputs,
+            print_backtrace,
+        } => {
+            target_features.push("+ext:SPV_KHR_non_semantic_info".into());
+            Some(format!(
+                "debug-printf{}{}",
+                if print_inputs { "+inputs" } else { "" },
+                if print_backtrace { "+backtrace" } else { "" }
+            ))
+        }
+        ShaderPanicStrategy::UNSOUND_DO_NOT_USE_UndefinedBehaviorViaUnreachable => {
+            Some("unreachable".into())
+        }
+    };
+    llvm_args.extend(abort_strategy.map(|strategy| format!("--abort-strategy={strategy}")));
 
     if let Ok(extra_codegen_args) = tracked_env_var_get("RUSTGPU_CODEGEN_ARGS") {
         llvm_args.extend(extra_codegen_args.split_whitespace().map(|s| s.to_string()));
@@ -523,7 +610,6 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
         rustflags.push(["-Cllvm-args=", &llvm_args].concat());
     }
 
-    let mut target_features = vec![];
     target_features.extend(builder.capabilities.iter().map(|cap| format!("+{cap:?}")));
     target_features.extend(builder.extensions.iter().map(|ext| format!("+ext:{ext}")));
     let target_features = join_checking_for_separators(target_features, ",");

crates/spirv-std/macros/src/lib.rs

@@ -374,7 +374,7 @@ fn path_from_ident(ident: Ident) -> syn::Type {
 /// debug_printfln!("pos.x: %f, pos.z: %f, int: %i", pos.x, pos.z, int);
 /// ```
 ///
-/// See <https://github.com/KhronosGroup/Vulkan-ValidationLayers/blob/master/docs/debug_printf.md#debug-printf-format-string> for formatting rules.
+/// See <https://github.com/KhronosGroup/Vulkan-ValidationLayers/blob/main/docs/debug_printf.md#debug-printf-format-string> for formatting rules.
 #[proc_macro]
 pub fn debug_printf(input: TokenStream) -> TokenStream {
     debug_printf_inner(syn::parse_macro_input!(input as DebugPrintfInput))

examples/runners/ash/src/main.rs

@@ -130,12 +130,14 @@ pub fn main() {
     ctx.build_pipelines(
         vk::PipelineCache::null(),
         vec![(
+            // HACK(eddyb) used to be `module: "sky_shader"` but we need `multimodule`
+            // for `debugPrintf` instrumentation to work (see `compile_shaders`).
             VertexShaderEntryPoint {
-                module: "sky_shader".into(),
+                module: "sky_shader::main_vs".into(),
                 entry_point: "main_vs".into(),
             },
             FragmentShaderEntryPoint {
-                module: "sky_shader".into(),
+                module: "sky_shader::main_fs".into(),
                 entry_point: "main_fs".into(),
             },
         )],
@@ -207,19 +209,27 @@ pub fn compile_shaders() -> Vec<SpvFile> {
     std::env::set_var("OUT_DIR", env!("OUT_DIR"));
     std::env::set_var("PROFILE", env!("PROFILE"));
 
-    let sky_shader_path =
+    SpirvBuilder::new("examples/shaders/sky-shader", "spirv-unknown-vulkan1.1")
-        SpirvBuilder::new("examples/shaders/sky-shader", "spirv-unknown-vulkan1.1")
+        .print_metadata(MetadataPrintout::None)
-            .print_metadata(MetadataPrintout::None)
+        // HACK(eddyb) having the `ash` runner do this is the easiest way I found
-            .build()
+        // to test this `panic!` feature with actual `debugPrintf` support.
-            .unwrap()
+        .shader_panic_strategy(spirv_builder::ShaderPanicStrategy::DebugPrintfThenExit {
-            .module
+            print_inputs: true,
-            .unwrap_single()
+            print_backtrace: true,
-            .to_path_buf();
+        })
-    let sky_shader = SpvFile {
+        // HACK(eddyb) needed because of `debugPrintf` instrumentation limitations
-        name: "sky_shader".to_string(),
+        // (see https://github.com/KhronosGroup/SPIRV-Tools/issues/4892).
-        data: read_spv(&mut File::open(sky_shader_path).unwrap()).unwrap(),
+        .multimodule(true)
-    };
+        .build()
-    vec![sky_shader]
+        .unwrap()
+        .module
+        .unwrap_multi()
+        .iter()
+        .map(|(name, path)| SpvFile {
+            name: format!("sky_shader::{name}"),
+            data: read_spv(&mut File::open(path).unwrap()).unwrap(),
+        })
+        .collect()
 }
 
 #[derive(Debug)]
@@ -370,7 +380,10 @@ impl RenderBase {
         };
 
         let device: ash::Device = {
-            let device_extension_names_raw = [khr::Swapchain::name().as_ptr()];
+            let mut device_extension_names_raw = vec![khr::Swapchain::name().as_ptr()];
+            if options.debug_layer {
+                device_extension_names_raw.push(vk::KhrShaderNonSemanticInfoFn::name().as_ptr());
+            }
             let features = vk::PhysicalDeviceFeatures {
                 shader_clip_distance: 1,
                 ..Default::default()