Intercept panic calls and replace them with aborts. (#305)

* new_structurizer: fix infinite loops.

* intrinsics: use an infinite loop for `abort`.

* Don't deduplicate zombie values even with other zombies.

* Bring back `zombie_even_in_user_code` and use it for constants.

* Use global `OpVariable`s instead of undefs for `ConstantPointer`s.

* Intercept panic calls and replace them with aborts.

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -10,7 +10,9 @@ 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::{BuilderMethods, ConstMethods, LayoutTypeMethods, OverflowOp};
+use rustc_codegen_ssa::traits::{
+    BuilderMethods, ConstMethods, IntrinsicCallMethods, LayoutTypeMethods, OverflowOp,
+};
 use rustc_codegen_ssa::MemFlags;
 use rustc_middle::bug;
 use rustc_middle::ty::Ty;
@@ -1995,6 +1997,13 @@ 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(llfn_def))
+        {
+            // HACK(eddyb) redirect builtin panic calls to an abort, to avoid
+            // needing to materialize `&core::panic::Location` or `format_args!`.
+            self.abort();
+            self.undef(result_type)
         } else {
             let args = args.iter().map(|arg| arg.def(self)).collect::<Vec<_>>();
             self.emit()

crates/rustc_codegen_spirv/src/builder/intrinsics.rs

@@ -487,13 +487,14 @@ impl<'a, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'tcx> {
     }
 
     fn abort(&mut self) {
-        // codegen_llvm uses call(llvm.trap) here, so it is not a block terminator
-        if !self.kernel_mode {
-            self.emit().kill().unwrap();
+        // 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 mut abort_loop_bx = self.build_sibling_block("abort_loop");
+        abort_loop_bx.br(abort_loop_bx.llbb());
+        self.br(abort_loop_bx.llbb());
         *self = self.build_sibling_block("abort_continue");
     }
-        // TODO: Figure out an appropriate instruction for kernel mode.
-    }
 
     fn assume(&mut self, _val: Self::Value) {
         // TODO: llvm.assume

crates/rustc_codegen_spirv/src/builder_spirv.rs

@@ -6,6 +6,7 @@ use rspirv::dr::{Block, Builder, Module, Operand};
 use rspirv::spirv::{AddressingModel, Capability, MemoryModel, Op, Word};
 use rspirv::{binary::Assemble, binary::Disassemble};
 use rustc_middle::bug;
+use rustc_span::{Span, DUMMY_SP};
 use std::cell::{RefCell, RefMut};
 use std::{fs::File, io::Write, path::Path};
 
@@ -20,7 +21,13 @@ pub enum SpirvValueKind {
     /// pointers to constants, or function pointers. So, instead, we create this ConstantPointer
     /// "meta-value": directly using it is an error, however, if it is attempted to be
     /// dereferenced, the "load" is instead a no-op that returns the underlying value directly.
-    ConstantPointer(Word),
+    ConstantPointer {
+        initializer: Word,
+
+        /// The global (module-scoped) `OpVariable` (with `initializer` set as
+        /// its initializer) to attach zombies to.
+        global_var: Word,
+    },
 }
 
 #[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
@@ -32,7 +39,10 @@ pub struct SpirvValue {
 impl SpirvValue {
     pub fn const_ptr_val(self, cx: &CodegenCx<'_>) -> Option<Self> {
         match self.kind {
-            SpirvValueKind::ConstantPointer(word) => {
+            SpirvValueKind::ConstantPointer {
+                initializer,
+                global_var: _,
+            } => {
                 let ty = match cx.lookup_type(self.ty) {
                     SpirvType::Pointer {
                         storage_class: _,
@@ -40,7 +50,7 @@ impl SpirvValue {
                     } => pointee,
                     ty => bug!("load called on variable that wasn't a pointer: {:?}", ty),
                 };
-                Some(word.with_type(ty))
+                Some(initializer.with_type(ty))
             }
             SpirvValueKind::Def(_) => None,
         }
@@ -50,43 +60,31 @@ impl SpirvValue {
     // contexts where the emitter is already locked. Doing so may cause subtle
     // rare bugs.
     pub fn def(self, bx: &builder::Builder<'_, '_>) -> Word {
-        match self.kind {
-            SpirvValueKind::Def(word) => word,
-            SpirvValueKind::ConstantPointer(_) => {
-                if bx.is_system_crate() {
-                    *bx.zombie_undefs_for_system_constant_pointers
-                        .borrow()
-                        .get(&self.ty)
-                        .expect("ConstantPointer didn't go through proper undef registration")
-                } else {
-                    bx.err("Cannot use this pointer directly, it must be dereferenced first");
-                    // Because we never get beyond compilation (into e.g. linking),
-                    // emitting an invalid ID reference here is OK.
-                    0
-                }
-            }
-        }
+        self.def_with_span(bx, bx.span())
     }
 
     // def and def_cx are separated, because Builder has a span associated with
     // what it's currently emitting.
     pub fn def_cx(self, cx: &CodegenCx<'_>) -> Word {
+        self.def_with_span(cx, DUMMY_SP)
+    }
+
+    pub fn def_with_span(self, cx: &CodegenCx<'_>, span: Span) -> Word {
         match self.kind {
             SpirvValueKind::Def(word) => word,
-            SpirvValueKind::ConstantPointer(_) => {
-                if cx.is_system_crate() {
-                    *cx.zombie_undefs_for_system_constant_pointers
-                        .borrow()
-                        .get(&self.ty)
-                        .expect("ConstantPointer didn't go through proper undef registration")
-                } else {
-                    cx.tcx
-                        .sess
-                        .err("Cannot use this pointer directly, it must be dereferenced first");
-                    // Because we never get beyond compilation (into e.g. linking),
-                    // emitting an invalid ID reference here is OK.
-                    0
-                }
+            SpirvValueKind::ConstantPointer {
+                initializer: _,
+                global_var,
+            } => {
+                // HACK(eddyb) we don't know whether this constant originated
+                // in a system crate, so it's better to always zombie.
+                cx.zombie_even_in_user_code(
+                    global_var,
+                    span,
+                    "Cannot use this pointer directly, it must be dereferenced first",
+                );
+
+                global_var
             }
         }
     }

crates/rustc_codegen_spirv/src/codegen_cx/constant.rs

@@ -4,7 +4,7 @@ use crate::builder_spirv::{SpirvConst, SpirvValue, SpirvValueExt};
 use crate::spirv_type::SpirvType;
 use rspirv::spirv::Word;
 use rustc_codegen_ssa::mir::place::PlaceRef;
-use rustc_codegen_ssa::traits::{BaseTypeMethods, ConstMethods, MiscMethods, StaticMethods};
+use rustc_codegen_ssa::traits::{ConstMethods, MiscMethods, StaticMethods};
 use rustc_middle::bug;
 use rustc_middle::mir::interpret::{read_target_uint, Allocation, GlobalAlloc, Pointer};
 use rustc_middle::ty::layout::TyAndLayout;
@@ -168,13 +168,14 @@ impl<'tcx> ConstMethods<'tcx> for CodegenCx<'tcx> {
 
     fn const_str(&self, s: Symbol) -> (Self::Value, Self::Value) {
         let len = s.as_str().len();
-        let ty = self.type_ptr_to(
-            self.layout_of(self.tcx.types.str_)
-                .spirv_type(DUMMY_SP, self),
-        );
-        let result = self.undef(ty);
-        self.zombie_no_span(result.def_cx(self), "constant string");
-        (result, self.const_usize(len as u64))
+        let str_ty = self
+            .layout_of(self.tcx.types.str_)
+            .spirv_type(DUMMY_SP, self);
+        // FIXME(eddyb) include the actual byte data.
+        (
+            self.make_constant_pointer(DUMMY_SP, self.undef(str_ty)),
+            self.const_usize(len as u64),
+        )
     }
     fn const_struct(&self, elts: &[Self::Value], _packed: bool) -> Self::Value {
         // Presumably this will get bitcasted to the right type?
@@ -471,7 +472,13 @@ impl<'tcx> CodegenCx<'tcx> {
                 *data = *c + asdf->y[*c];
                 }
                 */
-                self.zombie_no_span(result.def_cx(self), "constant runtime array value");
+                // HACK(eddyb) we don't know whether this constant originated
+                // in a system crate, so it's better to always zombie.
+                self.zombie_even_in_user_code(
+                    result.def_cx(self),
+                    DUMMY_SP,
+                    "constant runtime array value",
+                );
                 result
             }
             SpirvType::Pointer { .. } => {

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -142,6 +142,13 @@ impl<'tcx> CodegenCx<'tcx> {
             }
         }
 
+        if Some(instance_def_id) == self.tcx.lang_items().panic_fn() {
+            self.panic_fn_id.set(Some(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));
+        }
+
         declared
     }
 

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -27,7 +27,7 @@ use rustc_span::{SourceFile, Span, DUMMY_SP};
 use rustc_target::abi::call::FnAbi;
 use rustc_target::abi::{HasDataLayout, TargetDataLayout};
 use rustc_target::spec::{HasTargetSpec, Target};
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::collections::{HashMap, HashSet};
 use std::iter::once;
 
@@ -53,8 +53,13 @@ pub struct CodegenCx<'tcx> {
     pub sym: Box<Symbols>,
     pub instruction_table: InstructionTable,
     pub really_unsafe_ignore_bitcasts: RefCell<HashSet<SpirvValue>>,
-    pub zombie_undefs_for_system_constant_pointers: RefCell<HashMap<Word, Word>>,
     pub libm_intrinsics: RefCell<HashMap<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>>,
+    /// 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.
     pub i8_i16_atomics_allowed: bool,
@@ -105,8 +110,9 @@ impl<'tcx> CodegenCx<'tcx> {
             sym,
             instruction_table: InstructionTable::new(),
             really_unsafe_ignore_bitcasts: Default::default(),
-            zombie_undefs_for_system_constant_pointers: Default::default(),
             libm_intrinsics: Default::default(),
+            panic_fn_id: Default::default(),
+            panic_bounds_check_fn_id: Default::default(),
             i8_i16_atomics_allowed: false,
         }
     }
@@ -163,6 +169,9 @@ impl<'tcx> CodegenCx<'tcx> {
             self.tcx.sess.err(reason);
         }
     }
+    pub fn zombie_even_in_user_code(&self, word: Word, span: Span, reason: &'static str) {
+        self.zombie_values.borrow_mut().insert(word, (reason, span));
+    }
 
     pub fn is_system_crate(&self) -> bool {
         self.tcx
@@ -199,19 +208,19 @@ impl<'tcx> CodegenCx<'tcx> {
             pointee: value.ty,
         }
         .def(span, self);
-        if self.is_system_crate() {
-            // Create these undefs up front instead of on demand in SpirvValue::def because
+        let initializer = value.def_cx(self);
+
+        // Create these up front instead of on demand in SpirvValue::def because
         // SpirvValue::def can't use cx.emit()
-            self.zombie_undefs_for_system_constant_pointers
-                .borrow_mut()
-                .entry(ty)
-                .or_insert_with(|| {
-                    // We want a unique ID for these undefs, so don't use the caching system.
-                    self.emit_global().undef(ty, None)
-                });
-        }
+        let global_var =
+            self.emit_global()
+                .variable(ty, None, StorageClass::Function, Some(initializer));
+
         SpirvValue {
-            kind: SpirvValueKind::ConstantPointer(value.def_cx(self)),
+            kind: SpirvValueKind::ConstantPointer {
+                initializer,
+                global_var,
+            },
             ty,
         }
     }

crates/rustc_codegen_spirv/src/linker/duplicates.rs

@@ -121,7 +121,15 @@ fn make_type_key(
         }
     }
     if let Some(id) = inst.result_id {
-        data.push(if zombies.contains(&id) { 1 } else { 0 });
+        data.push(if zombies.contains(&id) {
+            if inst.result_type.is_some() {
+                id
+            } else {
+                1
+            }
+        } else {
+            0
+        });
         if let Some(annos) = annotations.get(&id) {
             data.extend_from_slice(annos)
         }

crates/rustc_codegen_spirv/src/linker/new_structurizer.rs

@@ -240,8 +240,14 @@ impl Structurizer<'_> {
 
                 // Move all of the contents of the original `block` into the
                 // new loop body, but keep labels and indices intact.
+                // Also update the existing merge if it happens to be the `block`
+                // we just moved (this should only be relevant to infinite loops).
                 self.func.blocks_mut()[while_body_block].instructions =
                     mem::replace(&mut self.func.blocks_mut()[block].instructions, vec![]);
+                if region.merge == block {
+                    region.merge = while_body_block;
+                    region.merge_id = while_body_block_id;
+                }
 
                 // Create a separate merge block for the loop body, as the original
                 // one might be used by an `OpSelectionMerge` and cannot be reused.

crates/spirv-builder/src/test/basic.rs

@@ -110,9 +110,7 @@ pub fn main() {
 }"#);
 }
 
-// TODO: Implement strings to make this compile
 #[test]
-#[ignore]
 fn panic() {
     val(r#"
 #[allow(unused_attributes)]
@@ -123,6 +121,36 @@ pub fn main() {
 "#);
 }
 
+#[test]
+fn panic_builtin() {
+    val(r#"
+fn int_div(x: usize) -> usize {
+    1 / x
+}
+
+#[allow(unused_attributes)]
+#[spirv(fragment)]
+pub fn main() {
+    int_div(0);
+}
+"#);
+}
+
+#[test]
+fn panic_builtin_bounds_check() {
+    val(r#"
+fn array_bounds_check(x: [u32; 4], i: usize) -> u32 {
+    x[i]
+}
+
+#[allow(unused_attributes)]
+#[spirv(fragment)]
+pub fn main() {
+    array_bounds_check([0, 1, 2, 3], 5);
+}
+"#);
+}
+
 // NOTE(eddyb) this won't pass Vulkan validation (see `push_constant_vulkan`),
 // but should still pass the basline SPIR-V validation.
 #[test]
@@ -167,3 +195,13 @@ pub fn main(constants: PushConstant<ShaderConstants>) {
 "#,
     );
 }
+
+#[test]
+fn infinite_loop() {
+    val(r#"
+#[allow(unused_attributes)]
+#[spirv(fragment)]
+pub fn main() {
+    loop {}
+}"#);
+}