location assignment: explicit location assignment with #[spirv(location = 0)]

Firestar99 · Firestar99 · commit aacd739e902d · 2025-12-08T12:11:07.000+01:00
diff --git a/crates/rustc_codegen_spirv/src/attr.rs b/crates/rustc_codegen_spirv/src/attr.rs
@@ -94,6 +94,7 @@ pub enum SpirvAttribute {
     Builtin(BuiltIn),
     DescriptorSet(u32),
     Binding(u32),
+    Location(u32),
     Flat,
     PerPrimitiveExt,
     Invariant,
@@ -130,6 +131,7 @@ pub struct AggregatedSpirvAttributes {
     pub builtin: Option<Spanned<BuiltIn>>,
     pub descriptor_set: Option<Spanned<u32>>,
     pub binding: Option<Spanned<u32>>,
+    pub location: Option<Spanned<u32>>,
     pub flat: Option<Spanned<()>>,
     pub invariant: Option<Spanned<()>>,
     pub per_primitive_ext: Option<Spanned<()>>,
@@ -216,6 +218,7 @@ impl AggregatedSpirvAttributes {
                 "#[spirv(descriptor_set)]",
             ),
             Binding(value) => try_insert(&mut self.binding, value, span, "#[spirv(binding)]"),
+            Location(value) => try_insert(&mut self.location, value, span, "#[spirv(location)]"),
             Flat => try_insert(&mut self.flat, (), span, "#[spirv(flat)]"),
             Invariant => try_insert(&mut self.invariant, (), span, "#[spirv(invariant)]"),
             PerPrimitiveExt => try_insert(
@@ -323,6 +326,7 @@ impl CheckSpirvAttrVisitor<'_> {
                 | SpirvAttribute::Builtin(_)
                 | SpirvAttribute::DescriptorSet(_)
                 | SpirvAttribute::Binding(_)
+                | SpirvAttribute::Location(_)
                 | SpirvAttribute::Flat
                 | SpirvAttribute::Invariant
                 | SpirvAttribute::PerPrimitiveExt
@@ -602,6 +606,8 @@ fn parse_spirv_attr<'a>(
                 SpirvAttribute::DescriptorSet(parse_attr_int_value(arg)?)
             } else if arg.has_name(sym.binding) {
                 SpirvAttribute::Binding(parse_attr_int_value(arg)?)
+            } else if arg.has_name(sym.location) {
+                SpirvAttribute::Location(parse_attr_int_value(arg)?)
             } else if arg.has_name(sym.input_attachment_index) {
                 SpirvAttribute::InputAttachmentIndex(parse_attr_int_value(arg)?)
             } else if arg.has_name(sym.spec_constant) {
diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs b/crates/rustc_codegen_spirv/src/codegen_cx/entry.rs
@@ -710,6 +710,26 @@ impl<'tcx> CodegenCx<'tcx> {
                 .name(var_id.or(spec_const_id).unwrap(), ident.to_string());
         }
 
+        // location assignment
+        let has_location = matches!(
+            storage_class,
+            Ok(StorageClass::Input | StorageClass::Output | StorageClass::UniformConstant)
+        );
+        let mut assign_location = |var_id: Result<Word, &str>, explicit: Option<u32>| {
+            let location = decoration_locations
+                .entry(storage_class.unwrap())
+                .or_insert_with(|| 0);
+            if let Some(explicit) = explicit {
+                *location = explicit;
+            }
+            self.emit_global().decorate(
+                var_id.unwrap(),
+                Decoration::Location,
+                std::iter::once(Operand::LiteralBit32(*location)),
+            );
+            *location += value_layout.size.bytes().div_ceil(16) as u32;
+        };
+
         // Emit `OpDecorate`s based on attributes.
         let mut decoration_supersedes_location = false;
         if let Some(builtin) = attrs.builtin {
@@ -757,6 +777,35 @@ impl<'tcx> CodegenCx<'tcx> {
             );
             decoration_supersedes_location = true;
         }
+        if let Some(location) = attrs.location {
+            if let Err(SpecConstant { .. }) = storage_class {
+                self.tcx.dcx().span_fatal(
+                    location.span,
+                    "`#[spirv(location = ...)]` cannot apply to `#[spirv(spec_constant)]`",
+                );
+            }
+            if attrs.descriptor_set.is_some() {
+                self.tcx.dcx().span_fatal(
+                    location.span,
+                    "`#[spirv(location = ...)]` cannot be combined with `#[spirv(descriptor_set = ...)]`",
+                );
+            }
+            if attrs.binding.is_some() {
+                self.tcx.dcx().span_fatal(
+                    location.span,
+                    "`#[spirv(location = ...)]` cannot be combined with `#[spirv(binding = ...)]`",
+                );
+            }
+            if !has_location {
+                self.tcx.dcx().span_fatal(
+                    location.span,
+                    "`#[spirv(location = ...)]` can only be used on Inputs (declared as plain values, eg. `Vec4`)\
+                     or Outputs (declared as mut ref, eg. `&mut Vec4`)",
+                );
+            }
+            assign_location(var_id, Some(location.value));
+            decoration_supersedes_location = true;
+        }
         if let Some(flat) = attrs.flat {
             if let Err(SpecConstant { .. }) = storage_class {
                 self.tcx.dcx().span_fatal(
@@ -867,21 +916,8 @@ impl<'tcx> CodegenCx<'tcx> {
         // individually.
         // TODO: Is this right for UniformConstant? Do they share locations with
         // input/outpus?
-        let has_location = !decoration_supersedes_location
-            && matches!(
-                storage_class,
-                Ok(StorageClass::Input | StorageClass::Output | StorageClass::UniformConstant)
-            );
-        if has_location {
-            let location = decoration_locations
-                .entry(storage_class.unwrap())
-                .or_insert_with(|| 0);
-            self.emit_global().decorate(
-                var_id.unwrap(),
-                Decoration::Location,
-                std::iter::once(Operand::LiteralBit32(*location)),
-            );
-            *location += value_layout.size.bytes().div_ceil(16) as u32;
+        if !decoration_supersedes_location && has_location {
+            assign_location(var_id, None);
         }
 
         match storage_class {
diff --git a/crates/rustc_codegen_spirv/src/symbols.rs b/crates/rustc_codegen_spirv/src/symbols.rs
@@ -23,6 +23,7 @@ pub struct Symbols {
 
     pub descriptor_set: Symbol,
     pub binding: Symbol,
+    pub location: Symbol,
     pub input_attachment_index: Symbol,
 
     pub spec_constant: Symbol,
@@ -420,6 +421,7 @@ impl Symbols {
 
             descriptor_set: Symbol::intern("descriptor_set"),
             binding: Symbol::intern("binding"),
+            location: Symbol::intern("location"),
             input_attachment_index: Symbol::intern("input_attachment_index"),
 
             spec_constant: Symbol::intern("spec_constant"),
diff --git a/tests/compiletests/ui/spirv-attr/location_assignment_explicit.rs b/tests/compiletests/ui/spirv-attr/location_assignment_explicit.rs
@@ -0,0 +1,37 @@
+// build-pass
+// compile-flags: -C llvm-args=--disassemble
+// normalize-stderr-test "OpSource .*\n" -> ""
+// normalize-stderr-test "OpLine .*\n" -> ""
+// normalize-stderr-test "%\d+ = OpString .*\n" -> ""
+// normalize-stderr-test "; .*\n" -> ""
+// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
+// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
+// ignore-spv1.0
+// ignore-spv1.1
+// ignore-spv1.2
+// ignore-spv1.3
+// ignore-vulkan1.0
+// ignore-vulkan1.1
+
+use spirv_std::glam::*;
+use spirv_std::{Image, spirv};
+
+#[derive(Copy, Clone, Default)]
+pub struct LargerThanVec4 {
+    a: Vec4,
+    b: Vec2,
+}
+
+#[spirv(vertex)]
+pub fn main(
+    #[spirv(location = 4)] out1: &mut LargerThanVec4,
+    out2: &mut Vec2, // should be 6
+    #[spirv(location = 0)] out3: &mut Mat4,
+    // 8 to 11 are unused, that's fine
+    #[spirv(location = 12)] out4: &mut f32,
+) {
+    *out1 = Default::default();
+    *out2 = Default::default();
+    *out3 = Default::default();
+    *out4 = Default::default();
+}
diff --git a/tests/compiletests/ui/spirv-attr/location_assignment_explicit.stderr b/tests/compiletests/ui/spirv-attr/location_assignment_explicit.stderr
@@ -0,0 +1,70 @@
+OpCapability Shader
+OpMemoryModel Logical Simple
+OpEntryPoint Vertex %1 "main" %2 %3 %4 %5
+OpName %7 "LargerThanVec4"
+OpMemberName %7 0 "a"
+OpMemberName %7 1 "b"
+OpName %8 "spirv_std::glam::Mat4"
+OpMemberName %8 0 "x_axis"
+OpMemberName %8 1 "y_axis"
+OpMemberName %8 2 "z_axis"
+OpMemberName %8 3 "w_axis"
+OpName %9 "LargerThanVec4"
+OpMemberName %9 0 "a"
+OpMemberName %9 1 "b"
+OpName %2 "out1"
+OpName %3 "out2"
+OpName %4 "out3"
+OpName %5 "out4"
+OpMemberDecorate %9 0 Offset 0
+OpMemberDecorate %9 1 Offset 16
+OpDecorate %2 Location 4
+OpDecorate %3 Location 6
+OpDecorate %4 Location 0
+OpDecorate %5 Location 12
+%10 = OpTypeFloat 32
+%11 = OpTypeVector %10 4
+%12 = OpTypeVector %10 2
+%7 = OpTypeStruct %11 %12
+%13 = OpTypePointer Output %7
+%14 = OpTypePointer Output %12
+%8 = OpTypeStruct %11 %11 %11 %11
+%15 = OpTypePointer Output %8
+%16 = OpTypePointer Output %10
+%17 = OpTypeVoid
+%18 = OpTypeFunction %17
+%9 = OpTypeStruct %11 %12
+%19 = OpConstant  %10  0
+%20 = OpConstantComposite  %11  %19 %19 %19 %19
+%21 = OpUndef  %9
+%2 = OpVariable  %13  Output
+%3 = OpVariable  %14  Output
+%22 = OpTypeInt 32 0
+%23 = OpConstant  %22  0
+%24 = OpConstant  %22  1
+%4 = OpVariable  %15  Output
+%25 = OpConstant  %10  1
+%26 = OpConstantComposite  %11  %25 %19 %19 %19
+%27 = OpConstantComposite  %11  %19 %25 %19 %19
+%28 = OpConstantComposite  %11  %19 %19 %25 %19
+%29 = OpConstantComposite  %11  %19 %19 %19 %25
+%30 = OpConstantComposite  %8  %26 %27 %28 %29
+%5 = OpVariable  %16  Output
+%1 = OpFunction  %17  None %18
+%31 = OpLabel
+%32 = OpCompositeInsert  %9  %20 %21 0
+%33 = OpCompositeInsert  %9  %19 %32 1 0
+%34 = OpCompositeInsert  %9  %19 %33 1 1
+%35 = OpCompositeExtract  %11  %34 0
+%36 = OpCompositeExtract  %12  %34 1
+%37 = OpCompositeConstruct  %7  %35 %36
+OpStore %2 %37
+%38 = OpInBoundsAccessChain  %16  %3 %23
+OpStore %38 %19
+%39 = OpInBoundsAccessChain  %16  %3 %24
+OpStore %39 %19
+OpStore %4 %30
+OpStore %5 %19
+OpNoLine
+OpReturn
+OpFunctionEnd
diff --git a/tests/compiletests/ui/spirv-attr/location_assignment_explicit_overlap.rs b/tests/compiletests/ui/spirv-attr/location_assignment_explicit_overlap.rs
@@ -0,0 +1,27 @@
+// build-fail
+// compile-flags: -C llvm-args=--disassemble
+// normalize-stderr-test "OpSource .*\n" -> ""
+// normalize-stderr-test "OpLine .*\n" -> ""
+// normalize-stderr-test "%\d+ = OpString .*\n" -> ""
+// normalize-stderr-test "; .*\n" -> ""
+// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
+// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
+// normalize-stderr-test "= note: module `.*`" -> "= note: module `<normalized>`"
+// ignore-spv1.0
+// ignore-spv1.1
+// ignore-spv1.2
+// ignore-spv1.3
+// ignore-spv1.4
+// ignore-spv1.5
+// ignore-spv1.6
+// ignore-vulkan1.0
+// ignore-vulkan1.1
+
+use spirv_std::glam::*;
+use spirv_std::{Image, spirv};
+
+#[spirv(vertex)]
+pub fn main(#[spirv(location = 0)] out1: &mut Mat4, #[spirv(location = 1)] out2: &mut Vec2) {
+    *out1 = Default::default();
+    *out2 = Default::default();
+}
diff --git a/tests/compiletests/ui/spirv-attr/location_assignment_explicit_overlap.stderr b/tests/compiletests/ui/spirv-attr/location_assignment_explicit_overlap.stderr
@@ -0,0 +1,51 @@
+OpCapability Shader
+OpMemoryModel Logical Simple
+OpEntryPoint Vertex %1 "main" %2 %3
+OpName %5 "spirv_std::glam::Mat4"
+OpMemberName %5 0 "x_axis"
+OpMemberName %5 1 "y_axis"
+OpMemberName %5 2 "z_axis"
+OpMemberName %5 3 "w_axis"
+OpName %2 "out1"
+OpName %3 "out2"
+OpDecorate %2 Location 0
+OpDecorate %3 Location 1
+%6 = OpTypeFloat 32
+%7 = OpTypeVector %6 4
+%5 = OpTypeStruct %7 %7 %7 %7
+%8 = OpTypePointer Output %5
+%9 = OpTypeVector %6 2
+%10 = OpTypePointer Output %9
+%11 = OpTypeVoid
+%12 = OpTypeFunction %11
+%2 = OpVariable  %8  Output
+%13 = OpConstant  %6  1
+%14 = OpConstant  %6  0
+%15 = OpConstantComposite  %7  %13 %14 %14 %14
+%16 = OpConstantComposite  %7  %14 %13 %14 %14
+%17 = OpConstantComposite  %7  %14 %14 %13 %14
+%18 = OpConstantComposite  %7  %14 %14 %14 %13
+%19 = OpConstantComposite  %5  %15 %16 %17 %18
+%20 = OpTypePointer Output %6
+%3 = OpVariable  %10  Output
+%21 = OpTypeInt 32 0
+%22 = OpConstant  %21  0
+%23 = OpConstant  %21  1
+%1 = OpFunction  %11  None %12
+%24 = OpLabel
+OpStore %2 %19
+%25 = OpInBoundsAccessChain  %20  %3 %22
+OpStore %25 %14
+%26 = OpInBoundsAccessChain  %20  %3 %23
+OpStore %26 %14
+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 `<normalized>`
+
+error: aborting due to 1 previous error
+
