asm: add support for noreturn option (#717)

msiglreith · web-flow · commit acda7716db2a · 2021-08-13T08:50:44.000+02:00
* 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
diff --git a/crates/rustc_codegen_spirv/src/builder/spirv_asm.rs b/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(
diff --git a/crates/spirv-std/src/arch.rs b/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)) }
 }
diff --git a/crates/spirv-std/src/arch/ray_tracing.rs b/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.
diff --git a/crates/spirv-std/src/ray_tracing.rs b/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]
diff --git a/crates/spirv-std/src/runtime_array.rs b/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 {}
     }
 }
diff --git a/tests/ui/lang/asm/block_tracking_fail.rs b/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);
+}
diff --git a/tests/ui/lang/asm/block_tracking_fail.stderr b/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
+
diff --git a/tests/ui/lang/asm/block_tracking_pass.rs b/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();
+}

Original file line number	Diff line number	Diff line change
`@@ -25,10 +25,9 @@ impl AccelerationStructure {`
`25`	`25`	`"%ret = OpTypeAccelerationStructureKHR",`
`26`	`26`	`"%result = OpConvertUToAccelerationStructureKHR %ret {id}",`
`27`	`27`	`"OpReturnValue %result",`
`28`		`- "%blah = OpLabel",`
`29`	`28`	`id = in(reg) id,`
	`29`	`+ options(noreturn)`
`30`	`30`	`}`
`31`		`- loop {}`
`32`	`31`	`}`
`33`	`32`
`34`	`33`	/// Converts a vector of two 32 bit integers into an [`AccelerationStructure`].
`@@ -47,10 +46,9 @@ impl AccelerationStructure {`
`47`	`46`	`"%id = OpLoad _ {id}",`
`48`	`47`	`"%result = OpConvertUToAccelerationStructureKHR %ret %id",`
`49`	`48`	`"OpReturnValue %result",`
`50`		`- "%blah = OpLabel",`
`51`	`49`	`id = in(reg) &id,`
	`50`	`+ options(noreturn),`
`52`	`51`	`}`
`53`		`- loop {}`
`54`	`52`	`}`
`55`	`53`
`56`	`54`	`#[spirv_std_macros::gpu_only]`
Original file line number	Diff line number	Diff line change
`@@ -17,11 +17,10 @@ impl<T> RuntimeArray<T> {`
`17`	`17`	`asm! {`
`18`	`18`	`"%result = OpAccessChain _ {arr} {index}",`
`19`	`19`	`"OpReturnValue %result",`
`20`		`- "%unused = OpLabel",`
`21`	`20`	`arr = in(reg) self,`
`22`	`21`	`index = in(reg) index,`
	`22`	`+ options(noreturn),`
`23`	`23`	`}`
`24`		`- loop {}`
`25`	`24`	`}`
`26`	`25`
`27`	`26`	`#[spirv_std_macros::gpu_only]`
`@@ -30,10 +29,9 @@ impl<T> RuntimeArray<T> {`
`30`	`29`	`asm! {`
`31`	`30`	`"%result = OpAccessChain _ {arr} {index}",`
`32`	`31`	`"OpReturnValue %result",`
`33`		`- "%unused = OpLabel",`
`34`	`32`	`arr = in(reg) self,`
`35`	`33`	`index = in(reg) index,`
	`34`	`+ options(noreturn),`
`36`	`35`	`}`
`37`		`- loop {}`
`38`	`36`	`}`
`39`	`37`	`}`