Fix location assignment

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) {

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -710,6 +710,41 @@ impl<'tcx> CodegenCx<'tcx> {
                 .name(var_id.or(spec_const_id).unwrap(), ident.to_string());
         }
 
+        // location assignment
+        // Note(@firestar99): UniformConstant are things like `SampledImage`, `StorageImage`, `Sampler` and
+        // `Acceleration structure`. Almost always they are assigned a `descriptor_set` and binding, thus never end up
+        // here being assigned locations. I think this is one of those occasions where spirv allows us to assign
+        // locations, but the "client API" Vulkan doesn't describe any use-case for them, or at least none I'm aware of.
+        // A quick scour through the spec revealed that `VK_KHR_dynamic_rendering_local_read` may need this, and while
+        // we don't support it yet (I assume), I'll just keep it here in case it becomes useful in the future.
+        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)),
+            );
+            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.
         let mut decoration_supersedes_location = false;
         if let Some(builtin) = attrs.builtin {
@@ -757,6 +792,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 +931,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 += 1;
+        if !decoration_supersedes_location && has_location {
+            assign_location(var_id, None);
         }
 
         match storage_class {

crates/rustc_codegen_spirv/src/spirv_type.rs

@@ -276,6 +276,34 @@ 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)
+    #[allow(clippy::match_same_arms)]
+    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 +315,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 { .. }

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"),

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

@@ -7,18 +7,11 @@ LL | #![feature(ptr_internals)]
    = note: using it is strongly discouraged
    = note: `#[warn(internal_features)]` on by default
 
-error: pointer has non-null integer address
-   |
-note: used from within `allocate_const_scalar::main`
-  --> $DIR/allocate_const_scalar.rs:16:5
-   |
-LL |     *output = POINTER;
-   |     ^^^^^^^^^^^^^^^^^
-note: called by `main`
-  --> $DIR/allocate_const_scalar.rs:15:8
+error: Type not supported in Input or Output declarations
+  --> $DIR/allocate_const_scalar.rs:15:21
    |
 LL | pub fn main(output: &mut Unique<()>) {
-   |        ^^^^
+   |                     ^^^^^^^^^^^^^^^
 
 error: aborting due to 1 previous error; 1 warning emitted
 

tests/compiletests/ui/spirv-attr/bad-deduce-storage-class.rs

@@ -5,11 +5,11 @@ use spirv_std::{Image, spirv};
 
 #[spirv(vertex)]
 pub fn main(
-    #[spirv(uniform)] error: &Image!(2D, type=f32),
-    #[spirv(uniform_constant)] warning: &Image!(2D, type=f32),
+    #[spirv(descriptor_set = 0, binding = 0, uniform)] error: &Image!(2D, type=f32),
+    #[spirv(descriptor_set = 0, binding = 1, uniform_constant)] warning: &Image!(2D, type=f32),
 ) {
 }
 
 // https://github.com/EmbarkStudios/rust-gpu/issues/585
 #[spirv(vertex)]
-pub fn issue_585(invalid: Image!(2D, type=f32)) {}
+pub fn issue_585(#[spirv(descriptor_set = 0, binding = 0)] invalid: Image!(2D, type=f32)) {}

tests/compiletests/ui/spirv-attr/bad-deduce-storage-class.stderr

@@ -1,25 +1,25 @@
 error: storage class mismatch
   --> $DIR/bad-deduce-storage-class.rs:8:5
    |
-LL |     #[spirv(uniform)] error: &Image!(2D, type=f32),
-   |     ^^^^^^^^-------^^^^^^^^^^---------------------
-   |             |                |
-   |             |                `UniformConstant` deduced from type
-   |             `Uniform` specified in attribute
+LL |     #[spirv(descriptor_set = 0, binding = 0, uniform)] error: &Image!(2D, type=f32),
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^-------^^^^^^^^^^---------------------
+   |                                              |                |
+   |                                              |                `UniformConstant` deduced from type
+   |                                              `Uniform` specified in attribute
    |
    = help: remove storage class attribute to use `UniformConstant` as storage class
 
 warning: redundant storage class attribute, storage class is deduced from type
-  --> $DIR/bad-deduce-storage-class.rs:9:13
+  --> $DIR/bad-deduce-storage-class.rs:9:46
    |
-LL |     #[spirv(uniform_constant)] warning: &Image!(2D, type=f32),
-   |             ^^^^^^^^^^^^^^^^
+LL |     #[spirv(descriptor_set = 0, binding = 1, uniform_constant)] warning: &Image!(2D, type=f32),
+   |                                              ^^^^^^^^^^^^^^^^
 
 error: entry parameter type must be by-reference: `&spirv_std::image::Image<f32, 1, 2, 0, 0, 0, 0, 4>`
-  --> $DIR/bad-deduce-storage-class.rs:15:27
+  --> $DIR/bad-deduce-storage-class.rs:15:69
    |
-LL | pub fn issue_585(invalid: Image!(2D, type=f32)) {}
-   |                           ^^^^^^^^^^^^^^^^^^^^
+LL | pub fn issue_585(#[spirv(descriptor_set = 0, binding = 0)] invalid: Image!(2D, type=f32)) {}
+   |                                                                     ^^^^^^^^^^^^^^^^^^^^
    |
    = note: this error originates in the macro `Image` (in Nightly builds, run with -Z macro-backtrace for more info)
 

tests/compiletests/ui/spirv-attr/bool-inputs-err.stderr

@@ -4,12 +4,24 @@ error: entry-point parameter cannot contain `bool`s
 LL |     input: bool,
    |            ^^^^
 
+error: Type not supported in Input or Output declarations
+  --> $DIR/bool-inputs-err.rs:13:12
+   |
+LL |     input: bool,
+   |            ^^^^
+
 error: entry-point parameter cannot contain `bool`s
   --> $DIR/bool-inputs-err.rs:14:13
    |
 LL |     output: &mut bool,
    |             ^^^^^^^^^
 
+error: Type not supported in Input or Output declarations
+  --> $DIR/bool-inputs-err.rs:14:13
+   |
+LL |     output: &mut bool,
+   |             ^^^^^^^^^
+
 error: entry-point parameter cannot contain `bool`s
   --> $DIR/bool-inputs-err.rs:15:35
    |
@@ -22,5 +34,5 @@ error: entry-point parameter cannot contain `bool`s
 LL |     #[spirv(uniform)] uniform: &Boolthing,
    |                                ^^^^^^^^^^
 
-error: aborting due to 4 previous errors
+error: aborting due to 6 previous errors
 

tests/compiletests/ui/spirv-attr/location_assignment.rs

@@ -0,0 +1,31 @@
+// 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(out1: &mut LargerThanVec4, out2: &mut Vec2, out3: &mut Mat4, out4: &mut f32) {
+    *out1 = Default::default();
+    *out2 = Default::default();
+    *out3 = Default::default();
+    *out4 = Default::default();
+}