Error when int doesn't have spirv(flat) (#815)

Author: khyperia

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -517,11 +517,12 @@ impl<'tcx> CodegenCx<'tcx> {
             );
         }
 
-        self.check_for_bools(
+        self.check_for_bad_types(
             hir_param.ty_span,
             var_ptr_spirv_type,
             storage_class,
             attrs.builtin.is_some(),
+            attrs.flat.is_some(),
         );
 
         // Assign locations from left to right, incrementing each storage class
@@ -563,7 +564,15 @@ impl<'tcx> CodegenCx<'tcx> {
     }
 
     // Booleans are only allowed in some storage classes. Error if they're in others.
-    fn check_for_bools(&self, span: Span, ty: Word, storage_class: StorageClass, is_builtin: bool) {
+    // Integers and f64s must be decorated with `#[spirv(flat)]`.
+    fn check_for_bad_types(
+        &self,
+        span: Span,
+        ty: Word,
+        storage_class: StorageClass,
+        is_builtin: bool,
+        is_flat: bool,
+    ) {
         // private and function are allowed here, but they can't happen.
         // SPIR-V technically allows all input/output variables to be booleans, not just builtins,
         // but has a note:
@@ -578,28 +587,49 @@ impl<'tcx> CodegenCx<'tcx> {
         {
             return;
         }
-        if recurse(self, ty) {
+        let mut has_bool = false;
+        let mut must_be_flat = false;
+        recurse(self, ty, &mut has_bool, &mut must_be_flat);
+        if has_bool {
             self.tcx
                 .sess
                 .span_err(span, "entrypoint parameter cannot contain a boolean");
         }
-        fn recurse(cx: &CodegenCx<'_>, ty: Word) -> bool {
+        if matches!(storage_class, StorageClass::Input | StorageClass::Output)
+            && must_be_flat
+            && !is_flat
+        {
+            self.tcx
+                .sess
+                .span_err(span, "parameter must be decorated with #[spirv(flat)]");
+        }
+        fn recurse(cx: &CodegenCx<'_>, ty: Word, has_bool: &mut bool, must_be_flat: &mut bool) {
             match cx.lookup_type(ty) {
-                SpirvType::Bool => true,
-                SpirvType::Adt { field_types, .. } => field_types.iter().any(|&f| recurse(cx, f)),
+                SpirvType::Bool => *has_bool = true,
+                SpirvType::Integer(_, _) | SpirvType::Float(64) => *must_be_flat = true,
+                SpirvType::Adt { field_types, .. } => {
+                    for f in field_types {
+                        recurse(cx, f, has_bool, must_be_flat);
+                    }
+                }
                 SpirvType::Vector { element, .. }
                 | SpirvType::Matrix { element, .. }
                 | SpirvType::Array { element, .. }
                 | SpirvType::RuntimeArray { element }
                 | SpirvType::Pointer { pointee: element }
                 | SpirvType::InterfaceBlock {
                     inner_type: element,
-                } => recurse(cx, element),
+                } => recurse(cx, element, has_bool, must_be_flat),
                 SpirvType::Function {
                     return_type,
                     arguments,
-                } => recurse(cx, return_type) || arguments.iter().any(|&a| recurse(cx, a)),
-                _ => false,
+                } => {
+                    recurse(cx, return_type, has_bool, must_be_flat);
+                    for a in arguments {
+                        recurse(cx, a, has_bool, must_be_flat);
+                    }
+                }
+                _ => (),
             }
         }
     }

tests/ui/byte_addressable_buffer/arr.rs

@@ -5,7 +5,7 @@ use spirv_std::{glam::Vec4, ByteAddressableBuffer};
 #[spirv(fragment)]
 pub fn load(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
-    out: &mut [i32; 4],
+    #[spirv(flat)] out: &mut [i32; 4],
 ) {
     unsafe {
         let buf = ByteAddressableBuffer::new(buf);
@@ -16,7 +16,7 @@ pub fn load(
 #[spirv(fragment)]
 pub fn store(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
-    val: [i32; 4],
+    #[spirv(flat)] val: [i32; 4],
 ) {
     unsafe {
         let mut buf = ByteAddressableBuffer::new(buf);

tests/ui/byte_addressable_buffer/big_struct.rs

@@ -14,7 +14,7 @@ pub struct BigStruct {
 #[spirv(fragment)]
 pub fn load(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
-    out: &mut BigStruct,
+    #[spirv(flat)] out: &mut BigStruct,
 ) {
     unsafe {
         let buf = ByteAddressableBuffer::new(buf);
@@ -25,7 +25,7 @@ pub fn load(
 #[spirv(fragment)]
 pub fn store(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
-    val: BigStruct,
+    #[spirv(flat)] val: BigStruct,
 ) {
     unsafe {
         let mut buf = ByteAddressableBuffer::new(buf);

tests/ui/byte_addressable_buffer/u32.rs

@@ -5,7 +5,7 @@ use spirv_std::ByteAddressableBuffer;
 #[spirv(fragment)]
 pub fn load(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
-    out: &mut u32,
+    #[spirv(flat)] out: &mut u32,
 ) {
     unsafe {
         let buf = ByteAddressableBuffer::new(buf);
@@ -14,7 +14,10 @@ pub fn load(
 }
 
 #[spirv(fragment)]
-pub fn store(#[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32], val: u32) {
+pub fn store(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
+    #[spirv(flat)] val: u32,
+) {
     unsafe {
         let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);

tests/ui/dis/pass-mode-cast-struct.rs

@@ -24,7 +24,7 @@ impl Foo {
 }
 
 #[spirv(fragment)]
-pub fn main(in_packed: u64, out_sum: &mut u32) {
+pub fn main(#[spirv(flat)] in_packed: u64, #[spirv(flat)] out_sum: &mut u32) {
     let foo = Foo::unpack(in_packed);
     *out_sum = foo.a + (foo.b + foo.c) as u32;
 }

tests/ui/image/query/query_levels.rs

@@ -6,7 +6,7 @@ use spirv_std::{arch, Image};
 #[spirv(fragment)]
 pub fn main(
     #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled),
-    output: &mut u32,
+    #[spirv(flat)] output: &mut u32,
 ) {
     *output = image.query_levels();
 }

tests/ui/image/query/query_samples.rs

@@ -6,7 +6,7 @@ use spirv_std::{arch, Image};
 #[spirv(fragment)]
 pub fn main(
     #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled, multisampled),
-    output: &mut u32,
+    #[spirv(flat)] output: &mut u32,
 ) {
     *output = image.query_samples();
 }

tests/ui/image/query/query_size.rs

@@ -6,7 +6,7 @@ use spirv_std::{arch, Image};
 #[spirv(fragment)]
 pub fn main(
     #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled=false),
-    output: &mut glam::UVec2,
+    #[spirv(flat)] output: &mut glam::UVec2,
 ) {
     *output = image.query_size();
 }

tests/ui/image/query/query_size_lod.rs

@@ -6,7 +6,7 @@ use spirv_std::{arch, Image};
 #[spirv(fragment)]
 pub fn main(
     #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled),
-    output: &mut glam::UVec2,
+    #[spirv(flat)] output: &mut glam::UVec2,
 ) {
     *output = image.query_size_lod(0);
 }

tests/ui/lang/asm/infer-access-chain-array.rs

@@ -8,7 +8,7 @@ use spirv_std as _;
 use glam::Vec4;
 
 #[spirv(fragment)]
-pub fn main(#[spirv(push_constant)] array_in: &[Vec4; 16], i: u32, out: &mut Vec4) {
+pub fn main(#[spirv(push_constant)] array_in: &[Vec4; 16], #[spirv(flat)] i: u32, out: &mut Vec4) {
     unsafe {
         asm!(
             "%val_ptr = OpAccessChain _ {array_ptr} {index}",