Support all panic entry-points equally - mainly for Rust 2021's panic!("...").

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -2335,10 +2335,99 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 );
             }
             result
-        } else if [self.panic_fn_id.get(), self.panic_bounds_check_fn_id.get()]
-            .contains(&Some(callee_val))
-        {
-            // HACK(eddyb) redirect builtin panic calls to an abort, to avoid
+        } else if self.panic_entry_point_ids.borrow().contains(&callee_val) {
+            // HACK(eddyb) Rust 2021 `panic!` always uses `format_args!`, even
+            // in the simple case that used to pass a `&str` constant, which
+            // would not remain reachable in the SPIR-V - but `format_args!` is
+            // more complex and neither immediate (`fmt::Arguments` is too big)
+            // nor simplified in MIR (e.g. promoted to a constant) in any way,
+            // so we have to try and remove the `fmt::Arguments::new` call here.
+            // HACK(eddyb) this is basically a `try` block.
+            let remove_simple_format_args_if_possible = || -> Option<()> {
+                let format_args_id = match args {
+                    &[SpirvValue {
+                        kind: SpirvValueKind::Def(format_args_id),
+                        ..
+                    }, SpirvValue {
+                        kind: SpirvValueKind::IllegalConst(_panic_location_id),
+                        ..
+                    }] => format_args_id,
+
+                    _ => return None,
+                };
+
+                // HACK(eddyb) we can remove SSA instructions even when they have
+                // side-effects, *as long as* they are "local" enough and cannot
+                // be observed from outside this current invocation - because the
+                // the abort, any SSA definitions or local variable writes can't
+                // be actually used anywhere else (other than *before* the abort).
+                let mut builder = self.emit();
+                let func_idx = builder.selected_function().unwrap();
+                let block_idx = builder.selected_block().unwrap();
+                let func = &mut builder.module_mut().functions[func_idx];
+                let mut non_debug_insts = func.blocks[block_idx]
+                    .instructions
+                    .iter()
+                    .enumerate()
+                    .filter(|(_, inst)| ![Op::Line, Op::NoLine].contains(&inst.class.opcode));
+                let mut relevant_insts_next_back = |expected_op| {
+                    non_debug_insts
+                        .next_back()
+                        .filter(|(_, inst)| inst.class.opcode == expected_op)
+                        .map(|(i, inst)| {
+                            (
+                                i,
+                                inst.result_id,
+                                inst.operands.iter().map(|operand| operand.unwrap_id_ref()),
+                            )
+                        })
+                };
+                let (_, load_src_id) = relevant_insts_next_back(Op::Load)
+                    .map(|(_, result_id, mut operands)| {
+                        (result_id.unwrap(), operands.next().unwrap())
+                    })
+                    .filter(|&(result_id, _)| result_id == format_args_id)?;
+                let (_, store_val_id) = relevant_insts_next_back(Op::Store)
+                    .map(|(_, _, mut operands)| {
+                        (operands.next().unwrap(), operands.next().unwrap())
+                    })
+                    .filter(|&(store_dst_id, _)| store_dst_id == load_src_id)?;
+                let (call_fmt_args_new_idx, _) = relevant_insts_next_back(Op::FunctionCall)
+                    .filter(|&(_, result_id, _)| result_id == Some(store_val_id))
+                    .map(|(i, _, mut operands)| (i, operands.next().unwrap(), operands))
+                    .filter(|&(_, callee, _)| self.fmt_args_new_fn_ids.borrow().contains(&callee))
+                    .map(|(i, _, mut call_args)| {
+                        assert_eq!(call_args.len(), 4);
+                        let mut arg = || call_args.next().unwrap();
+                        (i, [arg(), arg(), arg(), arg()])
+                    })
+                    .filter(|&(_, [_, _, _, fmt_args_len_id])| {
+                        // Only ever remove `fmt::Arguments` with no runtime values.
+                        matches!(
+                            self.builder.lookup_const_by_id(fmt_args_len_id),
+                            Some(SpirvConst::U32(0))
+                        )
+                    })?;
+
+                // Lastly, ensure that the `Op{Store,Load}` pair operates on
+                // a local `OpVariable`, i.e. is not externally observable.
+                let store_load_local_var = func.blocks[0]
+                    .instructions
+                    .iter()
+                    .take_while(|inst| inst.class.opcode == Op::Variable)
+                    .find(|inst| inst.result_id == Some(load_src_id));
+                if store_load_local_var.is_some() {
+                    // Keep all instructions up to (but not including) the call.
+                    func.blocks[block_idx]
+                        .instructions
+                        .truncate(call_fmt_args_new_idx);
+                }
+
+                None
+            };
+            remove_simple_format_args_if_possible();
+
+            // HACK(eddyb) redirect any possible panic call to an abort, to avoid
             // needing to materialize `&core::panic::Location` or `format_args!`.
             self.abort();
             self.undef(result_type)

crates/rustc_codegen_spirv/src/builder_spirv.rs

@@ -616,10 +616,14 @@ impl<'tcx> BuilderSpirv<'tcx> {
         SpirvValue { kind, ty }
     }
 
+    pub fn lookup_const_by_id(&self, id: Word) -> Option<SpirvConst<'tcx>> {
+        Some(self.id_to_const.borrow().get(&id)?.val)
+    }
+
     pub fn lookup_const(&self, def: SpirvValue) -> Option<SpirvConst<'tcx>> {
         match def.kind {
             SpirvValueKind::Def(id) | SpirvValueKind::IllegalConst(id) => {
-                Some(self.id_to_const.borrow().get(&id)?.val)
+                self.lookup_const_by_id(id)
             }
             _ => None,
         }

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -149,11 +149,21 @@ impl<'tcx> CodegenCx<'tcx> {
             }
         }
 
-        if Some(instance_def_id) == self.tcx.lang_items().panic_fn() {
-            self.panic_fn_id.set(Some(fn_id));
+        if [
+            self.tcx.lang_items().panic_fn(),
+            self.tcx.lang_items().panic_fmt(),
+            self.tcx.lang_items().panic_display(),
+            self.tcx.lang_items().panic_bounds_check_fn(),
+        ]
+        .contains(&Some(instance_def_id))
+        {
+            self.panic_entry_point_ids.borrow_mut().insert(fn_id);
         }
-        if Some(instance_def_id) == self.tcx.lang_items().panic_bounds_check_fn() {
-            self.panic_bounds_check_fn_id.set(Some(fn_id));
+
+        // HACK(eddyb) there is no good way to identify this definition
+        // (e.g. no `#[lang = "..."]` attribute), but this works well enough.
+        if demangled_symbol_name == "<core::fmt::Arguments>::new_v1" {
+            self.fmt_args_new_fn_ids.borrow_mut().insert(fn_id);
         }
 
         declared

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -18,7 +18,7 @@ use rustc_codegen_ssa::traits::{
     AsmMethods, BackendTypes, CoverageInfoMethods, DebugInfoMethods, GlobalAsmOperandRef,
     MiscMethods,
 };
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::mir::mono::CodegenUnit;
 use rustc_middle::mir::Body;
 use rustc_middle::ty::layout::{HasParamEnv, HasTyCtxt};
@@ -30,7 +30,7 @@ use rustc_span::{SourceFile, Span, DUMMY_SP};
 use rustc_target::abi::call::{FnAbi, PassMode};
 use rustc_target::abi::{HasDataLayout, TargetDataLayout};
 use rustc_target::spec::{HasTargetSpec, Target};
-use std::cell::{Cell, RefCell};
+use std::cell::RefCell;
 use std::collections::BTreeSet;
 use std::iter::once;
 use std::path::{Path, PathBuf};
@@ -59,14 +59,22 @@ pub struct CodegenCx<'tcx> {
     pub instruction_table: InstructionTable,
     pub libm_intrinsics: RefCell<FxHashMap<Word, super::builder::libm_intrinsics::LibmIntrinsic>>,
 
-    /// Simple `panic!("...")` and builtin panics (from MIR `Assert`s) call `#[lang = "panic"]`.
-    pub panic_fn_id: Cell<Option<Word>>,
+    /// 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).
+    //
+    // 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
+    // it mandatory even for `panic!("...")` (that were previously separate).
+    pub panic_entry_point_ids: RefCell<FxHashSet<Word>>,
+
+    /// `core::fmt::Arguments::new_v1` instances (for Rust 2021 panics).
+    pub fmt_args_new_fn_ids: RefCell<FxHashSet<Word>>,
+
     /// Intrinsic for loading a <T> from a &[u32]. The PassMode is the mode of the <T>.
     pub buffer_load_intrinsic_fn_id: RefCell<FxHashMap<Word, &'tcx PassMode>>,
     /// Intrinsic for storing a <T> into a &[u32]. The PassMode is the mode of the <T>.
     pub buffer_store_intrinsic_fn_id: RefCell<FxHashMap<Word, &'tcx PassMode>>,
-    /// Builtin bounds-checking panics (from MIR `Assert`s) call `#[lang = "panic_bounds_check"]`.
-    pub panic_bounds_check_fn_id: Cell<Option<Word>>,
 
     /// Some runtimes (e.g. intel-compute-runtime) disallow atomics on i8 and i16, even though it's allowed by the spec.
     /// This enables/disables them.
@@ -120,10 +128,10 @@ impl<'tcx> CodegenCx<'tcx> {
             sym,
             instruction_table: InstructionTable::new(),
             libm_intrinsics: Default::default(),
-            panic_fn_id: Default::default(),
+            panic_entry_point_ids: Default::default(),
+            fmt_args_new_fn_ids: Default::default(),
             buffer_load_intrinsic_fn_id: Default::default(),
             buffer_store_intrinsic_fn_id: Default::default(),
-            panic_bounds_check_fn_id: Default::default(),
             i8_i16_atomics_allowed: false,
             codegen_args,
         }

crates/rustc_codegen_spirv/src/linker/zombies.rs

@@ -3,7 +3,7 @@
 use super::{get_name, get_names};
 use crate::decorations::{CustomDecoration, ZombieDecoration};
 use rspirv::dr::{Instruction, Module};
-use rspirv::spirv::Word;
+use rspirv::spirv::{Op, Word};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_session::Session;
 use rustc_span::{Span, DUMMY_SP};
@@ -84,6 +84,8 @@ fn spread_zombie(module: &mut Module, zombie: &mut FxHashMap<Word, ZombieInfo<'_
     // No need to zombie defs within a function: If any def within a function is zombied, then the
     // whole function is zombied. But, we don't have to mark the defs within a function as zombie,
     // because the defs can't escape the function.
+    // HACK(eddyb) one exception to this is function-local variables, which may
+    // be unused and as such cannot be allowed to always zombie the function.
     for func in &module.functions {
         let func_id = func.def_id().unwrap();
         // Can't use zombie.entry() here, due to using the map in contains_zombie
@@ -91,13 +93,23 @@ fn spread_zombie(module: &mut Module, zombie: &mut FxHashMap<Word, ZombieInfo<'_
             // Func is already zombie, no need to scan it again.
             continue;
         }
-        let func_is_zombie = func
-            .all_inst_iter()
-            .find_map(|inst| is_or_contains_zombie(inst, zombie));
-        if let Some(reason) = func_is_zombie {
-            let pushed_reason = reason.push_stack(func_id);
-            zombie.insert(func_id, pushed_reason);
-            any = true;
+        for inst in func.all_inst_iter() {
+            if inst.class.opcode == Op::Variable {
+                let result_id = inst.result_id.unwrap();
+                if let Some(reason) = contains_zombie(inst, zombie) {
+                    if zombie.contains_key(&result_id) {
+                        continue;
+                    }
+                    let reason = reason.clone();
+                    zombie.insert(result_id, reason);
+                    any = true;
+                }
+            } else if let Some(reason) = is_or_contains_zombie(inst, zombie) {
+                let pushed_reason = reason.push_stack(func_id);
+                zombie.insert(func_id, pushed_reason);
+                any = true;
+                break;
+            }
         }
     }
     any