location assignment: proper location layout algorithm, as required by spec

https://registry.khronos.org/vulkan/specs/latest/html/vkspec.html#interfaces-iointerfaces-locations

Author: Firestar99

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -727,7 +727,16 @@ impl<'tcx> CodegenCx<'tcx> {
                 Decoration::Location,
                 std::iter::once(Operand::LiteralBit32(*location)),
             );
-            *location += value_layout.size.bytes().div_ceil(16) as u32;
+            let spirv_type = self.lookup_type(value_spirv_type);
+            if let Some(location_size) = spirv_type.location_size(self) {
+                *location += location_size;
+            } else {
+                *location += 1;
+                self.tcx.dcx().span_err(
+                    hir_param.ty_span,
+                    "Type not supported in Input or Output declarations",
+                );
+            }
         };
 
         // Emit `OpDecorate`s based on attributes.

crates/rustc_codegen_spirv/src/spirv_type.rs

@@ -276,6 +276,33 @@ impl SpirvType<'_> {
         id
     }
 
+    /// Returns how many Input / Output `location`s this type occupies, or None if this type is not allowed to be sent.
+    ///
+    /// See [Vulkan Spec 16.1.4. Location and Component Assignment](https://registry.khronos.org/vulkan/specs/latest/html/vkspec.html#interfaces-iointerfaces-locations)
+    pub fn location_size(&self, cx: &CodegenCx<'_>) -> Option<u32> {
+        let result = match *self {
+            // bools cannot be in an Input / Output interface
+            Self::Bool => return None,
+            Self::Integer(_, _) | Self::Float(_) => 1,
+            Self::Vector { .. } => 1,
+            Self::Adt { field_types, .. } => {
+                let mut locations = 0;
+                for f in field_types {
+                    locations += cx.lookup_type(*f).location_size(cx)?
+                }
+                locations
+            }
+            Self::Matrix { element, count } => cx.lookup_type(element).location_size(cx)? * count,
+            Self::Array { element, count } => {
+                let count = cx.builder.lookup_const_scalar(count).unwrap();
+                let count: u32 = count.try_into().unwrap();
+                cx.lookup_type(element).location_size(cx)? * count
+            }
+            _ => return None,
+        };
+        Some(result)
+    }
+
     pub fn sizeof(&self, cx: &CodegenCx<'_>) -> Option<Size> {
         let result = match *self {
             // Types that have a dynamic size, or no concept of size at all.
@@ -287,12 +314,9 @@ impl SpirvType<'_> {
             Self::Vector { size, .. } => size,
             Self::Matrix { element, count } => cx.lookup_type(element).sizeof(cx)? * count as u64,
             Self::Array { element, count } => {
-                cx.lookup_type(element).sizeof(cx)?
-                    * cx.builder
-                        .lookup_const_scalar(count)
-                        .unwrap()
-                        .try_into()
-                        .unwrap()
+                let count = cx.builder.lookup_const_scalar(count).unwrap();
+                let count = count.try_into().unwrap();
+                cx.lookup_type(element).sizeof(cx)? * count
             }
             Self::Pointer { .. } => cx.tcx.data_layout.pointer_size,
             Self::Image { .. }

tests/compiletests/ui/spirv-attr/location_assignment_array_f32.stderr

@@ -5,7 +5,7 @@ OpName %2 "out1"
 OpName %3 "out2"
 OpName %6 "<[f32OpDecorate %7 ArrayStride 4
 OpDecorate %2 Location 0
-OpDecorate %3 Location 1
+OpDecorate %3 Location 3
 %8 = OpTypeFloat 32
 %9 = OpTypeInt 32 0
 %10 = OpConstant  %9  3
@@ -37,11 +37,3 @@ OpFunctionEnd
 OpNoLine
 OpReturnValue %25
 OpFunctionEnd
-error: error:0:0 - [VUID-StandaloneSpirv-OpEntryPoint-08722] Entry-point has conflicting output location assignment at location 1, component 0
-         OpEntryPoint Vertex %1 "main" %out1 %out2
-   |
-   = note: spirv-val failed
-   = note: module `$TEST_BUILD_DIR/spirv-attr/location_assignment_array_f32.vulkan1.2`
-
-error: aborting due to 1 previous error
-

tests/compiletests/ui/spirv-attr/location_assignment_many_f32_struct.stderr

@@ -15,7 +15,7 @@ OpMemberDecorate %6 0 Offset 0
 OpMemberDecorate %6 1 Offset 4
 OpMemberDecorate %6 2 Offset 8
 OpDecorate %2 Location 0
-OpDecorate %3 Location 1
+OpDecorate %3 Location 3
 %7 = OpTypeFloat 32
 %5 = OpTypeStruct %7 %7 %7
 %8 = OpTypePointer Output %5
@@ -38,11 +38,3 @@ OpStore %3 %12
 OpNoLine
 OpReturn
 OpFunctionEnd
-error: error:0:0 - [VUID-StandaloneSpirv-OpEntryPoint-08722] Entry-point has conflicting output location assignment at location 1, component 0
-         OpEntryPoint Vertex %1 "main" %out1 %out2
-   |
-   = note: spirv-val failed
-   = note: module `$TEST_BUILD_DIR/spirv-attr/location_assignment_many_f32_struct.vulkan1.2`
-
-error: aborting due to 1 previous error
-

tests/compiletests/ui/spirv-attr/location_assignment_vec3_f32.stderr

@@ -12,7 +12,7 @@ OpName %3 "out2"
 OpMemberDecorate %6 0 Offset 0
 OpMemberDecorate %6 1 Offset 12
 OpDecorate %2 Location 0
-OpDecorate %3 Location 1
+OpDecorate %3 Location 2
 %7 = OpTypeFloat 32
 %8 = OpTypeVector %7 3
 %5 = OpTypeStruct %8 %7
@@ -36,11 +36,3 @@ OpStore %3 %13
 OpNoLine
 OpReturn
 OpFunctionEnd
-error: error:0:0 - [VUID-StandaloneSpirv-OpEntryPoint-08722] Entry-point has conflicting output location assignment at location 1, component 0
-         OpEntryPoint Vertex %1 "main" %out1 %out2
-   |
-   = note: spirv-val failed
-   = note: module `$TEST_BUILD_DIR/spirv-attr/location_assignment_vec3_f32.vulkan1.2`
-
-error: aborting due to 1 previous error
-