asm: add support for noreturn option (#717)

* asm: add support for noreturn option

OpUnreachable will be appended as terminator at the end of the asm block.

* asm: implicit label after return or abort terminator

* rework handling

* fix tests and add few comments

* fix tests

crates/rustc_codegen_spirv/src/builder/spirv_asm.rs

@@ -3,7 +3,7 @@ use crate::builder_spirv::{BuilderCursor, SpirvValue};
 use crate::codegen_cx::CodegenCx;
 use crate::spirv_type::SpirvType;
 use rspirv::dr;
-use rspirv::grammar::{LogicalOperand, OperandKind, OperandQuantifier};
+use rspirv::grammar::{reflect, LogicalOperand, OperandKind, OperandQuantifier};
 use rspirv::spirv::{
     FPFastMathMode, FragmentShadingRate, FunctionControl, ImageOperands, KernelProfilingInfo,
     LoopControl, MemoryAccess, MemorySemantics, Op, RayFlags, SelectionControl, StorageClass, Word,
@@ -70,8 +70,13 @@ impl<'a, 'tcx> AsmBuilderMethods<'tcx> for Builder<'a, 'tcx> {
         options: InlineAsmOptions,
         _line_spans: &[Span],
     ) {
-        if !options.is_empty() {
-            self.err(&format!("asm flags not supported: {:?}", options));
+        const SUPPORTED_OPTIONS: InlineAsmOptions = InlineAsmOptions::NORETURN;
+        let unsupported_options = options & !SUPPORTED_OPTIONS;
+        if !unsupported_options.is_empty() {
+            self.err(&format!(
+                "asm flags not supported: {:?}",
+                unsupported_options
+            ));
         }
         // vec of lines, and each line is vec of tokens
         let mut tokens = vec![vec![]];
@@ -141,14 +146,41 @@ impl<'a, 'tcx> AsmBuilderMethods<'tcx> for Builder<'a, 'tcx> {
                 id_to_type_map.insert(value.def(self), value.ty);
             }
         }
+
+        let mut asm_block = AsmBlock::Open;
         for line in tokens {
             self.codegen_asm(
                 &mut id_map,
                 &mut defined_ids,
                 &mut id_to_type_map,
+                &mut asm_block,
                 line.into_iter(),
             );
         }
+
+        match (options.contains(InlineAsmOptions::NORETURN), asm_block) {
+            (true, AsmBlock::Open) => {
+                self.err("`noreturn` requires a terminator at the end");
+            }
+            (true, AsmBlock::End(_)) => {
+                // `noreturn` appends an `OpUnreachable` after the asm block.
+                // This requires starting a new block for this.
+                let label = self.emit().id();
+                self.emit()
+                    .insert_into_block(
+                        dr::InsertPoint::End,
+                        dr::Instruction::new(Op::Label, None, Some(label), vec![]),
+                    )
+                    .unwrap();
+            }
+            (false, AsmBlock::Open) => (),
+            (false, AsmBlock::End(terminator)) => {
+                self.err(&format!(
+                    "trailing terminator {:?} requires `options(noreturn)`",
+                    terminator
+                ));
+            }
+        }
         for (id, num) in id_map {
             if !defined_ids.contains(&num) {
                 self.err(&format!("%{} is used but not defined", id));
@@ -178,6 +210,11 @@ enum OutRegister<'a> {
     Place(PlaceRef<'a, SpirvValue>),
 }
 
+enum AsmBlock {
+    Open,
+    End(Op),
+}
+
 impl<'cx, 'tcx> Builder<'cx, 'tcx> {
     fn lex_word<'a>(&self, line: &mut std::str::Chars<'a>) -> Option<Token<'a, 'cx, 'tcx>> {
         loop {
@@ -242,6 +279,7 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
         &mut self,
         id_map: &mut FxHashMap<&str, Word>,
         defined_ids: &mut FxHashSet<Word>,
+        asm_block: &mut AsmBlock,
         inst: dr::Instruction,
     ) {
         // Types declared must be registered in our type system.
@@ -328,10 +366,32 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
                 }
                 return;
             }
-            _ => {
+
+            op => {
                 self.emit()
                     .insert_into_block(dr::InsertPoint::End, inst)
                     .unwrap();
+
+                *asm_block = match *asm_block {
+                    AsmBlock::Open => {
+                        if reflect::is_block_terminator(op) {
+                            AsmBlock::End(op)
+                        } else {
+                            AsmBlock::Open
+                        }
+                    }
+                    AsmBlock::End(terminator) => {
+                        if op != Op::Label {
+                            self.err(&format!(
+                                "expected OpLabel after terminator {:?}",
+                                terminator
+                            ));
+                        }
+
+                        AsmBlock::Open
+                    }
+                };
+
                 return;
             }
         };
@@ -351,6 +411,7 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
         id_map: &mut FxHashMap<&'a str, Word>,
         defined_ids: &mut FxHashSet<Word>,
         id_to_type_map: &mut FxHashMap<Word, Word>,
+        asm_block: &mut AsmBlock,
         mut tokens: impl Iterator<Item = Token<'a, 'cx, 'tcx>>,
     ) where
         'cx: 'a,
@@ -427,7 +488,7 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
         if let Some(result_type) = instruction.result_type {
             id_to_type_map.insert(instruction.result_id.unwrap(), result_type);
         }
-        self.insert_inst(id_map, defined_ids, instruction);
+        self.insert_inst(id_map, defined_ids, asm_block, instruction);
         if let Some(OutRegister::Place(place)) = out_register {
             self.emit()
                 .store(

crates/spirv-std/src/arch.rs

@@ -148,8 +148,5 @@ pub unsafe fn vector_insert_dynamic<T: Scalar, V: Vector<T, N>, const N: usize>(
 #[doc(alias = "OpKill", alias = "discard")]
 #[allow(clippy::empty_loop)]
 pub fn kill() -> ! {
-    unsafe {
-        asm!("OpKill", "%unused = OpLabel");
-    }
-    loop {}
+    unsafe { asm!("OpKill", options(noreturn)) }
 }

crates/spirv-std/src/arch/ray_tracing.rs

@@ -44,8 +44,7 @@ pub unsafe fn report_intersection(hit: f32, hit_kind: u32) -> bool {
 #[inline]
 #[allow(clippy::empty_loop)]
 pub unsafe fn ignore_intersection() -> ! {
-    asm!("OpIgnoreIntersectionKHR", "%unused = OpLabel");
-    loop {}
+    asm!("OpIgnoreIntersectionKHR", options(noreturn));
 }
 
 /// Terminates the invocation that executes it, stops the ray traversal, accepts
@@ -57,8 +56,7 @@ pub unsafe fn ignore_intersection() -> ! {
 #[inline]
 #[allow(clippy::empty_loop)]
 pub unsafe fn terminate_ray() -> ! {
-    asm!("OpTerminateRayKHR", "%unused = OpLabel");
-    loop {}
+    asm!("OpTerminateRayKHR", options(noreturn));
 }
 
 /// Invoke a callable shader.

crates/spirv-std/src/ray_tracing.rs

@@ -25,10 +25,9 @@ impl AccelerationStructure {
             "%ret = OpTypeAccelerationStructureKHR",
             "%result = OpConvertUToAccelerationStructureKHR %ret {id}",
             "OpReturnValue %result",
-            "%blah = OpLabel",
             id = in(reg) id,
+            options(noreturn)
         }
-        loop {}
     }
 
     /// Converts a vector of two 32 bit integers into an [`AccelerationStructure`].
@@ -47,10 +46,9 @@ impl AccelerationStructure {
             "%id = OpLoad _ {id}",
             "%result = OpConvertUToAccelerationStructureKHR %ret %id",
             "OpReturnValue %result",
-            "%blah = OpLabel",
             id = in(reg) &id,
+            options(noreturn),
         }
-        loop {}
     }
 
     #[spirv_std_macros::gpu_only]

crates/spirv-std/src/runtime_array.rs

@@ -17,11 +17,10 @@ impl<T> RuntimeArray<T> {
         asm! {
             "%result = OpAccessChain _ {arr} {index}",
             "OpReturnValue %result",
-            "%unused = OpLabel",
             arr = in(reg) self,
             index = in(reg) index,
+            options(noreturn),
         }
-        loop {}
     }
 
     #[spirv_std_macros::gpu_only]
@@ -30,10 +29,9 @@ impl<T> RuntimeArray<T> {
         asm! {
             "%result = OpAccessChain _ {arr} {index}",
             "OpReturnValue %result",
-            "%unused = OpLabel",
             arr = in(reg) self,
             index = in(reg) index,
+            options(noreturn),
         }
-        loop {}
     }
 }

tests/ui/lang/asm/block_tracking_fail.rs

@@ -0,0 +1,39 @@
+// Tests validating tracking of basic blocks
+// within the `asm!` macro.
+// build-fail
+
+use spirv_std as _;
+
+// Active basic block with `noreturn`.
+fn asm_noreturn_open() {
+    unsafe {
+        asm!("", options(noreturn));
+    }
+}
+
+// No active basic block without `noreturn`.
+fn asm_closed() {
+    unsafe {
+        asm!(
+            "OpUnreachable",
+        );
+    }
+}
+
+// Invalid op after terminator
+fn asm_invalid_op_terminator(x: f32) {
+    unsafe {
+        asm!(
+            "OpKill",
+            "%sum = OpFAdd _ {x} {x}",
+            x = in(reg) x,
+        );
+    }
+}
+
+#[spirv(fragment)]
+pub fn main() {
+    asm_closed();
+    asm_noreturn_open();
+    asm_invalid_op_terminator(1.0);
+}

tests/ui/lang/asm/block_tracking_fail.stderr

@@ -0,0 +1,26 @@
+error: `noreturn` requires a terminator at the end
+  --> $DIR/block_tracking_fail.rs:10:9
+   |
+10 |         asm!("", options(noreturn));
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: trailing terminator Unreachable requires `options(noreturn)`
+  --> $DIR/block_tracking_fail.rs:17:9
+   |
+17 | /         asm!(
+18 | |             "OpUnreachable",
+19 | |         );
+   | |__________^
+
+error: expected OpLabel after terminator Kill
+  --> $DIR/block_tracking_fail.rs:26:9
+   |
+26 | /         asm!(
+27 | |             "OpKill",
+28 | |             "%sum = OpFAdd _ {x} {x}",
+29 | |             x = in(reg) x,
+30 | |         );
+   | |__________^
+
+error: aborting due to 3 previous errors
+

tests/ui/lang/asm/block_tracking_pass.rs

@@ -0,0 +1,29 @@
+// Tests validating tracking of basic blocks
+// within the `asm!` macro.
+// build-pass
+
+use spirv_std as _;
+
+fn asm_label() {
+    unsafe {
+        asm!(
+            "OpReturn", // close active block
+            "%unused = OpLabel", // open new block
+        );
+    }
+}
+
+fn asm_noreturn_single() -> ! {
+    unsafe {
+        asm!(
+            "OpKill", // close active block
+            options(noreturn),
+        );
+    }
+}
+
+#[spirv(fragment)]
+pub fn main() {
+    asm_label();
+    asm_noreturn_single();
+}