Add error for insts not allowed outside fragment (#821)

* Add error for insts not allowed outside fragment

* make get_name return entry point names too

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -81,7 +81,17 @@ fn apply_rewrite_rules(rewrite_rules: &FxHashMap<Word, Word>, blocks: &mut [Bloc
 }
 
 fn get_names(module: &Module) -> FxHashMap<Word, &str> {
-    module
+    let entry_names = module
+        .entry_points
+        .iter()
+        .filter(|i| i.class.opcode == Op::EntryPoint)
+        .map(|i| {
+            (
+                i.operands[1].unwrap_id_ref(),
+                i.operands[2].unwrap_literal_string(),
+            )
+        });
+    let debug_names = module
         .debug_names
         .iter()
         .filter(|i| i.class.opcode == Op::Name)
@@ -90,8 +100,9 @@ fn get_names(module: &Module) -> FxHashMap<Word, &str> {
                 i.operands[0].unwrap_id_ref(),
                 i.operands[1].unwrap_literal_string(),
             )
-        })
-        .collect()
+        });
+    // items later on take priority
+    entry_names.chain(debug_names).collect()
 }
 
 fn get_name<'a>(names: &FxHashMap<Word, &'a str>, id: Word) -> Cow<'a, str> {
@@ -157,6 +168,11 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         import_export_link::run(sess, &mut output)?;
     }
 
+    {
+        let _timer = sess.timer("link_fragment_inst_check");
+        simple_passes::check_fragment_insts(sess, &output)?;
+    }
+
     // HACK(eddyb) this has to run before the `remove_zombies` pass, so that any
     // zombies that are passed as call arguments, but eventually unused, won't
     // be (incorrectly) considered used.

crates/rustc_codegen_spirv/src/linker/simple_passes.rs

@@ -1,6 +1,10 @@
+use super::{get_name, get_names, Result};
 use rspirv::dr::{Block, Function, Module};
-use rspirv::spirv::{Op, Word};
+use rspirv::spirv::{ExecutionModel, Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_errors::ErrorReported;
+use rustc_session::Session;
+use std::iter::once;
 use std::mem::take;
 
 pub fn shift_ids(module: &mut Module, add: u32) {
@@ -151,3 +155,107 @@ pub fn name_variables_pass(module: &mut Module) {
         }
     }
 }
+
+// Some instructions are only valid in fragment shaders. Check them.
+pub fn check_fragment_insts(sess: &Session, module: &Module) -> Result<()> {
+    let mut visited = vec![false; module.functions.len()];
+    let mut stack = Vec::new();
+    let mut names = None;
+    let func_id_to_idx: FxHashMap<Word, usize> = module
+        .functions
+        .iter()
+        .enumerate()
+        .map(|(index, func)| (func.def_id().unwrap(), index))
+        .collect();
+    let entries = module
+        .entry_points
+        .iter()
+        .filter(|i| i.operands[0].unwrap_execution_model() != ExecutionModel::Fragment)
+        .map(|i| func_id_to_idx[&i.operands[1].unwrap_id_ref()]);
+    let mut okay = true;
+    for entry in entries {
+        okay &= visit(
+            sess,
+            module,
+            &mut visited,
+            &mut stack,
+            &mut names,
+            entry,
+            &func_id_to_idx,
+        );
+    }
+    return if okay { Ok(()) } else { Err(ErrorReported) };
+
+    // returns false if error
+    fn visit<'m>(
+        sess: &Session,
+        module: &'m Module,
+        visited: &mut Vec<bool>,
+        stack: &mut Vec<Word>,
+        names: &mut Option<FxHashMap<Word, &'m str>>,
+        index: usize,
+        func_id_to_idx: &FxHashMap<Word, usize>,
+    ) -> bool {
+        if visited[index] {
+            return true;
+        }
+        visited[index] = true;
+        stack.push(module.functions[index].def_id().unwrap());
+        let mut okay = true;
+        for inst in module.functions[index].all_inst_iter() {
+            if inst.class.opcode == Op::FunctionCall {
+                let called_func = func_id_to_idx[&inst.operands[0].unwrap_id_ref()];
+                okay &= visit(
+                    sess,
+                    module,
+                    visited,
+                    stack,
+                    names,
+                    called_func,
+                    func_id_to_idx,
+                );
+            }
+            if matches!(
+                inst.class.opcode,
+                Op::ImageSampleImplicitLod
+                    | Op::ImageSampleDrefImplicitLod
+                    | Op::ImageSampleProjImplicitLod
+                    | Op::ImageSampleProjDrefImplicitLod
+                    | Op::ImageQueryLod
+                    | Op::ImageSparseSampleImplicitLod
+                    | Op::ImageSparseSampleDrefImplicitLod
+                    | Op::DPdx
+                    | Op::DPdy
+                    | Op::Fwidth
+                    | Op::DPdxFine
+                    | Op::DPdyFine
+                    | Op::FwidthFine
+                    | Op::DPdxCoarse
+                    | Op::DPdyCoarse
+                    | Op::FwidthCoarse
+                    | Op::Kill
+            ) {
+                // These instructions are (usually) in system functions - if we get an error, allow
+                // the system function to be visited again from elsewhere to emit another error
+                // from another callsite.
+                visited[index] = false;
+
+                let names = names.get_or_insert_with(|| get_names(module));
+                let stack = stack.iter().rev().map(|&s| get_name(names, s).into_owned());
+                let note = once("Stack:".to_string())
+                    .chain(stack)
+                    .collect::<Vec<_>>()
+                    .join("\n");
+                sess.struct_err(&format!(
+                    "{} cannot be used outside a fragment shader",
+                    inst.class.opname
+                ))
+                .note(&note)
+                .emit();
+                okay = false;
+            }
+        }
+        stack.pop();
+        okay
+    }
+}

tests/ui/image/implicit_not_in_fragment.rs

@@ -0,0 +1,23 @@
+// build-fail
+
+use spirv_std::{arch, Image, Sampler};
+
+fn deeper_stack(image2d: &Image!(2D, type=f32, sampled), sampler: &Sampler) -> glam::Vec4 {
+    let v2 = glam::Vec2::new(0.0, 1.0);
+    image2d.sample(*sampler, v2)
+}
+fn deep_stack(image2d: &Image!(2D, type=f32, sampled), sampler: &Sampler) -> glam::Vec4 {
+    deeper_stack(image2d, sampler)
+}
+
+#[spirv(vertex)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] image2d: &Image!(2D, type=f32, sampled),
+    #[spirv(descriptor_set = 0, binding = 1)] sampler: &Sampler,
+    output: &mut glam::Vec4,
+) {
+    let v2 = glam::Vec2::new(0.0, 1.0);
+    let r0 = deep_stack(image2d, sampler);
+    let r1: glam::Vec4 = image2d.sample(*sampler, v2);
+    *output = r0 + r1;
+}

tests/ui/image/implicit_not_in_fragment.stderr

@@ -0,0 +1,18 @@
+error: ImageSampleImplicitLod cannot be used outside a fragment shader
+  |
+  = note: Stack:
+          <spirv_std::image::Image<f32, 1, 2, 0, 0, 1, 0>>::sample::<f32, glam::vec4::Vec4, glam::vec2::Vec2>
+          implicit_not_in_fragment::deeper_stack
+          implicit_not_in_fragment::deep_stack
+          implicit_not_in_fragment::main
+          main
+
+error: ImageSampleImplicitLod cannot be used outside a fragment shader
+  |
+  = note: Stack:
+          <spirv_std::image::Image<f32, 1, 2, 0, 0, 1, 0>>::sample::<f32, glam::vec4::Vec4, glam::vec2::Vec2>
+          implicit_not_in_fragment::main
+          main
+
+error: aborting due to 2 previous errors
+

tests/ui/lang/consts/nested-ref-in-composite.stderr

@@ -6,7 +6,7 @@ error: constant arrays/structs cannot contain pointers to other constants
    |
    = note: Stack:
            nested_ref_in_composite::main_pair
-           Unnamed function ID %24
+           main_pair
 
 error: constant arrays/structs cannot contain pointers to other constants
   --> $DIR/nested-ref-in-composite.rs:27:19
@@ -16,7 +16,7 @@ error: constant arrays/structs cannot contain pointers to other constants
    |
    = note: Stack:
            nested_ref_in_composite::main_array3
-           Unnamed function ID %33
+           main_array3
 
 error: error:0:0 - No OpEntryPoint instruction was found. This is only allowed if the Linkage capability is being used.
   |

tests/ui/lang/core/ptr/allocate_const_scalar.stderr

@@ -2,7 +2,7 @@ error: pointer has non-null integer address
   |
   = note: Stack:
           allocate_const_scalar::main
-          Unnamed function ID %4
+          main
 
 error: error:0:0 - No OpEntryPoint instruction was found. This is only allowed if the Linkage capability is being used.
   |