Add runtime array type (#596)

crates/rustc_codegen_spirv/src/abi.rs

@@ -912,5 +912,25 @@ fn trans_intrinsic_type<'tcx>(
                 Err(ErrorReported)
             }
         }
+        IntrinsicType::RuntimeArray => {
+            if ty.size != Size::from_bytes(4) {
+                cx.tcx
+                    .sess
+                    .err("#[spirv(runtime_array)] type must have size 4");
+                return Err(ErrorReported);
+            }
+
+            // 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) = substs.types().next() {
+                let element = trans_type_impl(cx, span, cx.layout_of(elem_ty), false);
+                Ok(SpirvType::RuntimeArray { element }.def(span, cx))
+            } else {
+                cx.tcx
+                    .sess
+                    .err("#[spirv(runtime_array)] type must have a generic element type");
+                Err(ErrorReported)
+            }
+        }
     }
 }

crates/rustc_codegen_spirv/src/attr.rs

@@ -75,6 +75,7 @@ pub enum IntrinsicType {
     AccelerationStructureKhr,
     SampledImage,
     RayQueryKhr,
+    RuntimeArray,
 }
 
 // NOTE(eddyb) when adding new `#[spirv(...)]` attributes, the tests found inside

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -356,11 +356,17 @@ impl<'tcx> CodegenCx<'tcx> {
         };
         let spirv_ty = self.layout_of(value_ty).spirv_type(hir_param.ty_span, self);
         // Some types automatically specify a storage class. Compute that here.
-        let inferred_storage_class_from_ty = match self.lookup_type(spirv_ty) {
+        let element_ty = match self.lookup_type(spirv_ty) {
+            SpirvType::Array { element, .. } | SpirvType::RuntimeArray { element } => {
+                self.lookup_type(element)
+            }
+            ty => ty,
+        };
+        let inferred_storage_class_from_ty = match element_ty {
             SpirvType::Image { .. }
             | SpirvType::Sampler
             | SpirvType::SampledImage { .. }
-            | SpirvType::AccelerationStructureKhr => {
+            | SpirvType::AccelerationStructureKhr { .. } => {
                 if is_ref {
                     Some(StorageClass::UniformConstant)
                 } else {
@@ -471,6 +477,16 @@ impl<'tcx> CodegenCx<'tcx> {
         // which we represent with `SpirvType::InterfaceBlock` (see its doc comment).
         // This "interface block" construct is also required for "runtime arrays".
         let is_unsized = self.lookup_type(value_spirv_type).sizeof(self).is_none();
+        let is_pair = matches!(entry_arg_abi.mode, PassMode::Pair(..));
+        let is_unsized_with_len = is_pair && is_unsized;
+        if is_pair && !is_unsized {
+            // If PassMode is Pair, then we need to fill in the second part of the pair with a
+            // value. We currently only do that with unsized types, so if a type is a pair for some
+            // other reason (e.g. a tuple), we bail.
+            self.tcx
+                .sess
+                .span_fatal(hir_param.ty_span, "pair type not supported yet")
+        }
         let var_ptr_spirv_type;
         let (value_ptr, value_len) = match storage_class {
             StorageClass::PushConstant | StorageClass::Uniform | StorageClass::StorageBuffer => {
@@ -483,7 +499,7 @@ impl<'tcx> CodegenCx<'tcx> {
 
                 let value_ptr = bx.struct_gep(var.with_type(var_ptr_spirv_type), 0);
 
-                let value_len = if is_unsized {
+                let value_len = if is_unsized_with_len {
                     match self.lookup_type(value_spirv_type) {
                         SpirvType::RuntimeArray { .. } => {}
                         _ => self.tcx.sess.span_err(
@@ -501,11 +517,49 @@ impl<'tcx> CodegenCx<'tcx> {
 
                     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
                 };
 
                 (value_ptr, value_len)
             }
+            StorageClass::UniformConstant => {
+                var_ptr_spirv_type = self.type_ptr_to(value_spirv_type);
+
+                match self.lookup_type(value_spirv_type) {
+                    SpirvType::RuntimeArray { .. } => {
+                        if is_unsized_with_len {
+                            self.tcx.sess.span_err(
+                                hir_param.ty_span,
+                                "uniform_constant must use &RuntimeArray<T>, not &[T]",
+                            )
+                        }
+                    }
+                    _ => {
+                        if is_unsized {
+                            self.tcx.sess.span_err(
+                                hir_param.ty_span,
+                                "only plain slices are supported as unsized types",
+                            )
+                        }
+                    }
+                }
+
+                let value_len = if is_pair {
+                    // We've already emitted an error, fill in a placeholder value
+                    Some(bx.undef(self.type_isize()))
+                } else {
+                    None
+                };
+
+                (var.with_type(var_ptr_spirv_type), value_len)
+            }
             _ => {
                 var_ptr_spirv_type = self.type_ptr_to(value_spirv_type);
 

crates/rustc_codegen_spirv/src/symbols.rs

@@ -338,6 +338,10 @@ impl Symbols {
                 "sampled_image",
                 SpirvAttribute::IntrinsicType(IntrinsicType::SampledImage),
             ),
+            (
+                "runtime_array",
+                SpirvAttribute::IntrinsicType(IntrinsicType::RuntimeArray),
+            ),
             ("unroll_loops", SpirvAttribute::UnrollLoops),
             ("internal_buffer_load", SpirvAttribute::InternalBufferLoad),
             ("internal_buffer_store", SpirvAttribute::InternalBufferStore),

crates/spirv-std/src/lib.rs

@@ -90,6 +90,7 @@ pub mod image;
 pub mod integer;
 pub mod memory;
 pub mod ray_tracing;
+mod runtime_array;
 mod sampler;
 pub mod scalar;
 pub(crate) mod sealed;
@@ -99,6 +100,7 @@ pub mod vector;
 pub use self::sampler::Sampler;
 pub use crate::macros::Image;
 pub use num_traits;
+pub use runtime_array::*;
 pub use textures::*;
 
 #[cfg(all(not(test), target_arch = "spirv"))]

crates/spirv-std/src/runtime_array.rs

@@ -0,0 +1,39 @@
+use core::marker::PhantomData;
+
+#[spirv(runtime_array)]
+pub struct RuntimeArray<T> {
+    // spooky! this field does not exist, so if it's referenced in rust code, things will explode
+    _do_not_touch: u32,
+    _phantom: PhantomData<T>,
+}
+
+// It would be nice to use the Index/IndexMut traits here, but because we do not have the length of
+// the array, it's impossible to make them be safe operations (indexing out of bounds), and
+// Index/IndexMut are marked as safe functions.
+impl<T> RuntimeArray<T> {
+    #[spirv_std_macros::gpu_only]
+    #[allow(clippy::empty_loop)]
+    pub unsafe fn index(&self, index: usize) -> &T {
+        asm! {
+            "%result = OpAccessChain _ {arr} {index}",
+            "OpReturnValue %result",
+            "%unused = OpLabel",
+            arr = in(reg) self,
+            index = in(reg) index,
+        }
+        loop {}
+    }
+
+    #[spirv_std_macros::gpu_only]
+    #[allow(clippy::empty_loop)]
+    pub unsafe fn index_mut(&mut self, index: usize) -> &mut T {
+        asm! {
+            "%result = OpAccessChain _ {arr} {index}",
+            "OpReturnValue %result",
+            "%unused = OpLabel",
+            arr = in(reg) self,
+            index = in(reg) index,
+        }
+        loop {}
+    }
+}

tests/ui/storage_class/runtime_descriptor_array.rs

@@ -0,0 +1,23 @@
+// build-pass
+use spirv_std::{Image, RuntimeArray, Sampler};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] sampler: &Sampler,
+    #[spirv(descriptor_set = 0, binding = 1)] slice: &RuntimeArray<Image!(2D, type=f32, sampled)>,
+    #[spirv(descriptor_set = 0, binding = 2)] sized_slice: &[Image!(2D, type=f32, sampled); 5],
+    output: &mut glam::Vec4,
+) {
+    unsafe {
+        asm!(
+            "OpCapability RuntimeDescriptorArray",
+            "OpExtension \"SPV_EXT_descriptor_indexing\""
+        )
+    }
+    let img = unsafe { slice.index(5) };
+    let v2 = glam::Vec2::new(0.0, 1.0);
+    *output = img.sample(*sampler, v2);
+
+    let img_2 = &sized_slice[2];
+    *output += img_2.sample(*sampler, v2);
+}

tests/ui/storage_class/runtime_descriptor_array_error.rs

@@ -0,0 +1,10 @@
+// build-fail
+
+use spirv_std::{Image, RuntimeArray};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] one: &[Image!(2D, type=f32, sampled)],
+    #[spirv(uniform, descriptor_set = 0, binding = 0)] two: &RuntimeArray<u32>,
+) {
+}

tests/ui/storage_class/runtime_descriptor_array_error.stderr

@@ -0,0 +1,14 @@
+error: uniform_constant must use &RuntimeArray<T>, not &[T]
+ --> $DIR/runtime_descriptor_array_error.rs:7:52
+  |
+7 |     #[spirv(descriptor_set = 0, binding = 0)] one: &[Image!(2D, type=f32, sampled)],
+  |                                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+warning: use &[T] instead of &RuntimeArray<T>
+ --> $DIR/runtime_descriptor_array_error.rs:8:61
+  |
+8 |     #[spirv(uniform, descriptor_set = 0, binding = 0)] two: &RuntimeArray<u32>,
+  |                                                             ^^^^^^^^^^^^^^^^^^
+
+error: aborting due to previous error; 1 warning emitted
+