Add arrays of buffer descriptors

crates/rustc_codegen_spirv/src/abi.rs

@@ -933,8 +933,23 @@ fn trans_intrinsic_type<'tcx>(
             // We use a generic to indicate the underlying element type.
             // The spirv type of it will be generated by querying the type of the first generic.
             if let Some(elem_ty) = args.types().next() {
-                let element = cx.layout_of(elem_ty).spirv_type(span, cx);
-                Ok(SpirvType::RuntimeArray { element }.def(span, cx))
+                let layout = cx.layout_of(elem_ty);
+                let element = layout.spirv_type(span, cx);
+                let element_ty = cx.lookup_type(element);
+
+                if element_ty.is_uniform_constant() {
+                    // array of image, sampler, SampledImage etc. descriptors
+                    Ok(SpirvType::RuntimeArray { element }.def(span, cx))
+                } else {
+                    // array of buffer descriptors
+                    Ok(SpirvType::RuntimeArray {
+                        element: SpirvType::InterfaceBlock {
+                            inner_type: element,
+                        }
+                        .def(span, cx),
+                    }
+                    .def(span, cx))
+                }
             } else {
                 Err(cx
                     .tcx

crates/rustc_codegen_spirv/src/attr.rs

@@ -98,6 +98,7 @@ pub enum SpirvAttribute {
     // `fn`/closure attributes:
     BufferLoadIntrinsic,
     BufferStoreIntrinsic,
+    RuntimeArrayIndexIntrinsic,
 }
 
 // HACK(eddyb) this is similar to `rustc_span::Spanned` but with `value` as the
@@ -133,6 +134,7 @@ pub struct AggregatedSpirvAttributes {
     // `fn`/closure attributes:
     pub buffer_load_intrinsic: Option<Spanned<()>>,
     pub buffer_store_intrinsic: Option<Spanned<()>>,
+    pub runtime_array_index_intrinsic: Option<Spanned<()>>,
 }
 
 struct MultipleAttrs {
@@ -237,6 +239,12 @@ impl AggregatedSpirvAttributes {
                 span,
                 "#[spirv(buffer_store_intrinsic)]",
             ),
+            RuntimeArrayIndexIntrinsic => try_insert(
+                &mut self.runtime_array_index_intrinsic,
+                (),
+                span,
+                "#[spirv(runtime_array_index_intrinsic)]",
+            ),
         }
     }
 }
@@ -358,12 +366,12 @@ impl CheckSpirvAttrVisitor<'_> {
 
                     _ => Err(Expected("function parameter")),
                 },
-                SpirvAttribute::BufferLoadIntrinsic | SpirvAttribute::BufferStoreIntrinsic => {
-                    match target {
-                        Target::Fn => Ok(()),
-                        _ => Err(Expected("function")),
-                    }
-                }
+                SpirvAttribute::BufferLoadIntrinsic
+                | SpirvAttribute::BufferStoreIntrinsic
+                | SpirvAttribute::RuntimeArrayIndexIntrinsic => match target {
+                    Target::Fn => Ok(()),
+                    _ => Err(Expected("function")),
+                },
             };
             match valid_target {
                 Err(Expected(expected_target)) => {

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -2531,6 +2531,11 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             .borrow()
             .get(&callee_val)
             .copied();
+        let runtime_array_index_intrinsic = self
+            .runtime_array_index_intrinsic_fn_id
+            .borrow()
+            .get(&callee_val)
+            .copied();
         if let Some(libm_intrinsic) = libm_intrinsic {
             let result = self.call_libm_intrinsic(libm_intrinsic, result_type, args);
             if result_type != result.ty {
@@ -3024,6 +3029,8 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 kind: SpirvValueKind::IllegalTypeUsed(void_ty),
                 ty: void_ty,
             }
+        } else if let Some(mode) = runtime_array_index_intrinsic {
+            self.codegen_runtime_array_index_intrinsic(result_type, args, mode)
         } else {
             let args = args.iter().map(|arg| arg.def(self)).collect::<Vec<_>>();
             self.emit()

crates/rustc_codegen_spirv/src/builder/mod.rs

@@ -3,6 +3,7 @@ mod byte_addressable_buffer;
 mod ext_inst;
 mod intrinsics;
 pub mod libm_intrinsics;
+mod runtime_array;
 mod spirv_asm;
 
 pub use ext_inst::ExtInst;

crates/rustc_codegen_spirv/src/builder/runtime_array.rs

@@ -0,0 +1,130 @@
+use crate::builder::Builder;
+use crate::builder_spirv::{SpirvValue, SpirvValueExt, SpirvValueKind};
+use crate::spirv_type::SpirvType;
+use rspirv::spirv::Word;
+use rustc_codegen_ssa::traits::{BaseTypeMethods, BuilderMethods};
+use rustc_target::abi::call::PassMode;
+
+impl<'a, 'tcx> Builder<'a, 'tcx> {
+    /// Note: DOES NOT do bounds checking! Bounds checking is expected to be done in the caller.
+    pub fn codegen_runtime_array_index_intrinsic(
+        &mut self,
+        result_type: Word,
+        args: &[SpirvValue],
+        pass_mode: &PassMode,
+    ) -> SpirvValue {
+        match pass_mode {
+            PassMode::Ignore => {
+                return SpirvValue {
+                    kind: SpirvValueKind::IllegalTypeUsed(result_type),
+                    ty: result_type,
+                };
+            }
+            // PassMode::Pair is identical to PassMode::Direct - it's returned as a struct
+            PassMode::Direct(_) | PassMode::Pair(_, _) => (),
+            PassMode::Cast { .. } => {
+                self.fatal("PassMode::Cast not supported in codegen_runtime_array_index_intrinsic")
+            }
+            PassMode::Indirect { .. } => self
+                .fatal("PassMode::Indirect not supported in codegen_runtime_array_index_intrinsic"),
+        }
+
+        // Signatures:
+        // fn <T: ?Sized>(runtime_array: &RuntimeArray<T>, index: usize) -> &T
+        // fn <T: ?Sized>(runtime_array: &mut RuntimeArray<T>, index: usize) -> &mut T
+        if args.len() != 2 {
+            self.fatal(format!(
+                "runtime_array_index_intrinsic should have 3 args, it has {}",
+                args.len()
+            ));
+        }
+        let runtime_array = args[0];
+        let index = args[1];
+
+        let runtime_array_type = self.lookup_type(runtime_array.ty);
+        let element_ty = match runtime_array_type {
+			SpirvType::Pointer { pointee } => {
+				match self.lookup_type(pointee) {
+					SpirvType::RuntimeArray { element } => {
+						element
+					}
+					_ => self.fatal(format!(
+						"runtime_array_index_intrinsic args[0] is {:?} and not a Pointer to a RuntimeArray!",
+						runtime_array_type
+					)),
+				}
+			}
+			_ => self.fatal(format!(
+				"runtime_array_index_intrinsic args[0] is {:?} and not a Pointer!",
+				runtime_array_type
+			)),
+		};
+
+        let ptr_element = self.type_ptr_to(element_ty);
+        let element = self
+            .emit()
+            .access_chain(
+                ptr_element,
+                None,
+                runtime_array.def(self),
+                [index.def(self)],
+            )
+            .unwrap()
+            .with_type(ptr_element);
+
+        match self.lookup_type(element_ty) {
+            SpirvType::InterfaceBlock { .. } => {
+                // array of buffer descriptors
+                let inner = self.struct_gep(element_ty, element, 0);
+                match pass_mode {
+                    PassMode::Direct(_) => {
+                        // element is sized
+                        if inner.ty == result_type {
+                            inner
+                        } else {
+                            self.fatal(format!(
+								"runtime_array_index_intrinsic expected result_type to equal RuntimeArray's InterfaceBlock's inner_type: {:?} == {:?}",
+								self.lookup_type(result_type).debug(result_type, self),
+								self.lookup_type(inner.ty).debug(inner.ty, self)
+							))
+                        }
+                    }
+                    PassMode::Pair(_, _) => {
+                        // element is a slice
+                        match self.lookup_type(result_type) {
+                            SpirvType::Adt { field_types, .. } if field_types.len() == 2
+								&& matches!(self.lookup_type(field_types[0]), SpirvType::Pointer {..})
+								&& field_types[1] == self.type_isize() => {
+							}
+							_ => self.fatal(format!(
+								"Expected element of RuntimeArray to be a plain slice, like `&RuntimeArray<[u32]>`, but got {:?}!",
+								self.lookup_type(result_type).debug(result_type, self)
+							))
+                        };
+                        let len = self
+                            .emit()
+                            .array_length(self.type_isize(), None, element.def(self), 0)
+                            .unwrap();
+                        self.emit()
+                            .composite_construct(result_type, None, [inner.def(self), len])
+                            .unwrap()
+                            .with_type(result_type)
+                    }
+                    _ => unreachable!(),
+                }
+            }
+            _ => {
+                // array of UniformConstant (image, sampler, etc.) descriptors
+                if ptr_element == result_type {
+                    element
+                } else {
+                    self.fatal(format!(
+						"runtime_array_index_intrinsic expected result_type to equal RuntimeArray's element_ty: {:?} == {:?}",
+						self.lookup_type(result_type).debug(result_type, self),
+						self.lookup_type(ptr_element).debug(ptr_element, self)
+					))
+                }
+            }
+        }
+    }
+}

crates/rustc_codegen_spirv/src/codegen_cx/declare.rs

@@ -155,6 +155,12 @@ impl<'tcx> CodegenCx<'tcx> {
                 .borrow_mut()
                 .insert(fn_id, mode);
         }
+        if attrs.runtime_array_index_intrinsic.is_some() {
+            let mode = &fn_abi.ret.mode;
+            self.runtime_array_index_intrinsic_fn_id
+                .borrow_mut()
+                .insert(fn_id, mode);
+        }
 
         let instance_def_id = instance.def_id();
 

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -266,25 +266,21 @@ impl<'tcx> CodegenCx<'tcx> {
             }
             ty => ty,
         };
-        let deduced_storage_class_from_ty = match element_ty {
-            SpirvType::Image { .. }
-            | SpirvType::Sampler
-            | SpirvType::SampledImage { .. }
-            | SpirvType::AccelerationStructureKhr { .. } => {
-                if is_ref {
-                    Some(StorageClass::UniformConstant)
-                } else {
-                    self.tcx.sess.span_err(
-                        hir_param.ty_span,
-                        format!(
-                            "entry parameter type must be by-reference: `&{}`",
-                            value_layout.ty,
-                        ),
-                    );
-                    None
-                }
+        let deduced_storage_class_from_ty = if element_ty.is_uniform_constant() {
+            if is_ref {
+                Some(StorageClass::UniformConstant)
+            } else {
+                self.tcx.sess.span_err(
+                    hir_param.ty_span,
+                    format!(
+                        "entry parameter type must be by-reference: `&{}`",
+                        value_layout.ty,
+                    ),
+                );
+                None
             }
-            _ => None,
+        } else {
+            None
         };
         // Storage classes can be specified via attribute. Compute that here, and emit diagnostics.
         let attr_storage_class = attrs.storage_class.map(|storage_class_attr| {
@@ -502,49 +498,56 @@ impl<'tcx> CodegenCx<'tcx> {
             Ok(
                 StorageClass::PushConstant | StorageClass::Uniform | StorageClass::StorageBuffer,
             ) => {
-                let var_spirv_type = SpirvType::InterfaceBlock {
-                    inner_type: value_spirv_type,
-                }
-                .def(hir_param.span, self);
-                var_ptr_spirv_type = self.type_ptr_to(var_spirv_type);
-
-                let value_ptr = bx.struct_gep(
-                    var_spirv_type,
-                    var_id.unwrap().with_type(var_ptr_spirv_type),
-                    0,
-                );
-
-                let value_len = if is_unsized_with_len {
-                    match self.lookup_type(value_spirv_type) {
-                        SpirvType::RuntimeArray { .. } => {}
-                        _ => {
-                            self.tcx.sess.span_err(
-                                hir_param.ty_span,
-                                "only plain slices are supported as unsized types",
-                            );
-                        }
+                match self.lookup_type(value_spirv_type) {
+                    SpirvType::RuntimeArray { element }
+                        if matches!(
+                            self.lookup_type(element),
+                            SpirvType::InterfaceBlock { .. }
+                        ) =>
+                    {
+                        // array of buffer descriptors
+                        var_ptr_spirv_type = self.type_ptr_to(value_spirv_type);
+                        (Ok(var_id.unwrap().with_type(var_ptr_spirv_type)), None)
                     }
+                    _ => {
+                        // single buffer descriptor
+                        let var_spirv_type = SpirvType::InterfaceBlock {
+                            inner_type: value_spirv_type,
+                        }
+                        .def(hir_param.span, self);
+                        var_ptr_spirv_type = self.type_ptr_to(var_spirv_type);
 
-                    // FIXME(eddyb) shouldn't this be `usize`?
-                    let len_spirv_type = self.type_isize();
-                    let len = bx
-                        .emit()
-                        .array_length(len_spirv_type, None, var_id.unwrap(), 0)
-                        .unwrap();
+                        let value_ptr = bx.struct_gep(
+                            var_spirv_type,
+                            var_id.unwrap().with_type(var_ptr_spirv_type),
+                            0,
+                        );
 
-                    Some(len.with_type(len_spirv_type))
-                } else {
-                    if is_unsized {
-                        // It's OK to use a RuntimeArray<u32> and not have a length parameter, but
-                        // it's just nicer ergonomics to use a slice.
-                        self.tcx
-                            .sess
-                            .span_warn(hir_param.ty_span, "use &[T] instead of &RuntimeArray<T>");
-                    }
-                    None
-                };
+                        let value_len = if is_unsized_with_len {
+                            match self.lookup_type(value_spirv_type) {
+                                SpirvType::RuntimeArray { .. } => {}
+                                _ => {
+                                    self.tcx.sess.span_err(
+                                        hir_param.ty_span,
+                                        "only plain slices are supported as unsized types",
+                                    );
+                                }
+                            }
+
+                            // FIXME(eddyb) shouldn't this be `usize`?
+                            let len_spirv_type = self.type_isize();
+                            let len = bx
+                                .emit()
+                                .array_length(len_spirv_type, None, var_id.unwrap(), 0)
+                                .unwrap();
+                            Some(len.with_type(len_spirv_type))
+                        } else {
+                            None
+                        };
 
-                (Ok(value_ptr), value_len)
+                        (Ok(value_ptr), value_len)
+                    }
+                }
             }
             Ok(StorageClass::UniformConstant) => {
                 var_ptr_spirv_type = self.type_ptr_to(value_spirv_type);

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -74,6 +74,8 @@ pub struct CodegenCx<'tcx> {
     pub buffer_load_intrinsic_fn_id: RefCell<FxHashMap<Word, &'tcx PassMode>>,
     /// Intrinsic for storing a <T> into a &[u32]. The PassMode is the mode of the <T>.
     pub buffer_store_intrinsic_fn_id: RefCell<FxHashMap<Word, &'tcx PassMode>>,
+    /// Intrinsic for loading a descriptor from a `RuntimeArray`. The PassMode is the mode of the <T>.
+    pub runtime_array_index_intrinsic_fn_id: RefCell<FxHashMap<Word, &'tcx PassMode>>,
 
     /// Some runtimes (e.g. intel-compute-runtime) disallow atomics on i8 and i16, even though it's allowed by the spec.
     /// This enables/disables them.
@@ -132,6 +134,7 @@ impl<'tcx> CodegenCx<'tcx> {
             fmt_rt_arg_new_fn_ids_to_ty_and_spec: Default::default(),
             buffer_load_intrinsic_fn_id: Default::default(),
             buffer_store_intrinsic_fn_id: Default::default(),
+            runtime_array_index_intrinsic_fn_id: Default::default(),
             i8_i16_atomics_allowed: false,
             codegen_args,
         }