Replace spirv_std::storage_class::X<T> with T/&mut T and #[spirv(x)] &T/&mut T.

README.md

@@ -19,12 +19,13 @@ However, many things aren't implemented yet: for example, loops and switches are
 ![Sky shader](docs/assets/sky.jpg)
 
 ```rust
-use spirv_std::{glam::{Vec3, Vec4, vec2, vec3}, storage_class::{Input, Output}};
+use glam::{Vec3, Vec4, vec2, vec3};
 
 #[spirv(fragment)]
 pub fn main(
-    #[spirv(frag_coord)] in_frag_coord: Input<Vec4>,
-    mut output: Output<Vec4>,
+    #[spirv(frag_coord)] in_frag_coord: &Vec4,
+    #[spirv(push_constant)] constants: &ShaderConstants,
+    output: &mut Vec4,
 ) {
     let frag_coord = vec2(in_frag_coord.x, in_frag_coord.y);
     let mut uv = (frag_coord - 0.5 * vec2(constants.width as f32, constants.height as f32))
@@ -61,7 +62,7 @@ Our hope with this project is that we push the industry forward by bringing an e
 
 ### Why Embark?
 
-At Embark, we've been building our own new game engine from the ground up in Rust. We have previous experience in-house developing the [RLSL](https://github.com/MaikKlein/rlsl) prototype, and we have a team of excellent rendering engineers that are familiar with the problems in current shading languages both from games, game engines and other industries. So, we believe we are uniquely positioned to attempt solving this problem. 
+At Embark, we've been building our own new game engine from the ground up in Rust. We have previous experience in-house developing the [RLSL](https://github.com/MaikKlein/rlsl) prototype, and we have a team of excellent rendering engineers that are familiar with the problems in current shading languages both from games, game engines and other industries. So, we believe we are uniquely positioned to attempt solving this problem.
 
 We want to streamline our own internal development with a single great language, build an open source graphics ecosystem and community, facilitate code-sharing between GPU and CPU, and most importantly: to enable our (future) users, and fellow developers, to more rapidly build great looking and engaging experiences.
 
@@ -71,7 +72,7 @@ If we do this project right, one wouldn't necessarily need an entire team of ren
 
 The scope of this overall project is quite broad, but is in multiple stages
 
-- `rustc` compiler backend to generate [SPIR-V], plugging in via `-Z codegen-backend`. 
+- `rustc` compiler backend to generate [SPIR-V], plugging in via `-Z codegen-backend`.
   - This is the same mechanism that [rustc_codegen_cranelift](https://github.com/bjorn3/rustc_codegen_cranelift) and [rustc_codegen_gcc](https://github.com/antoyo/rustc_codegen_gcc) use.
   - Currently only [SPIR-V] support is planned, [Vulkan](https://en.wikipedia.org/wiki/Vulkan_(API))'s open compiler target
   - Possible a future version could support [DXIL](https://github.com/microsoft/DirectXShaderCompiler/blob/master/docs/DXIL.rst) (the target for DirectX) or [WGSL](https://github.com/gpuweb/gpuweb/tree/main/wgsl) (the WebGPU shading language that's bijective with SPIR-V)

crates/rustc_codegen_spirv/src/abi.rs

@@ -585,22 +585,6 @@ fn dig_scalar_pointee_adt<'tcx>(
     }
 }
 
-/// Handles `#[spirv(storage_class="blah")]`. Note this is only called in the entry interface variables code, because this is only
-/// used for spooky builtin stuff, and we pinky promise to never have more than one pointer field in one of these.
-// TODO: Enforce this is only used in spirv-std.
-pub(crate) fn get_storage_class<'tcx>(
-    cx: &CodegenCx<'tcx>,
-    ty: TyAndLayout<'tcx>,
-) -> Option<StorageClass> {
-    if let TyKind::Adt(adt, _substs) = ty.ty.kind() {
-        AggregatedSpirvAttributes::parse(cx, cx.tcx.get_attrs(adt.did))
-            .storage_class
-            .map(|attr| attr.value)
-    } else {
-        None
-    }
-}
-
 fn trans_aggregate<'tcx>(cx: &CodegenCx<'tcx>, span: Span, ty: TyAndLayout<'tcx>) -> Word {
     match ty.fields {
         FieldsShape::Primitive => cx.tcx.sess.fatal(&format!(

crates/rustc_codegen_spirv/src/attr.rs

@@ -79,14 +79,14 @@ pub enum IntrinsicType {
 #[derive(Debug, Clone)]
 pub enum SpirvAttribute {
     // `struct` attributes:
-    StorageClass(StorageClass),
     IntrinsicType(IntrinsicType),
     Block,
 
     // `fn` attributes:
     Entry(Entry),
 
     // (entry) `fn` parameter attributes:
+    StorageClass(StorageClass),
     Builtin(BuiltIn),
     DescriptorSet(u32),
     Binding(u32),
@@ -110,14 +110,14 @@ pub struct Spanned<T> {
 #[derive(Default)]
 pub struct AggregatedSpirvAttributes {
     // `struct` attributes:
-    pub storage_class: Option<Spanned<StorageClass>>,
     pub intrinsic_type: Option<Spanned<IntrinsicType>>,
     pub block: Option<Spanned<()>>,
 
     // `fn` attributes:
     pub entry: Option<Spanned<Entry>>,
 
     // (entry) `fn` parameter attributes:
+    pub storage_class: Option<Spanned<StorageClass>>,
     pub builtin: Option<Spanned<BuiltIn>>,
     pub descriptor_set: Option<Spanned<u32>>,
     pub binding: Option<Spanned<u32>>,
@@ -187,14 +187,14 @@ impl AggregatedSpirvAttributes {
 
         use SpirvAttribute::*;
         match attr {
-            StorageClass(value) => {
-                try_insert(&mut self.storage_class, value, span, "storage class")
-            }
             IntrinsicType(value) => {
                 try_insert(&mut self.intrinsic_type, value, span, "intrinsic type")
             }
             Block => try_insert(&mut self.block, (), span, "#[spirv(block)]"),
             Entry(value) => try_insert(&mut self.entry, value, span, "entry-point"),
+            StorageClass(value) => {
+                try_insert(&mut self.storage_class, value, span, "storage class")
+            }
             Builtin(value) => try_insert(&mut self.builtin, value, span, "builtin"),
             DescriptorSet(value) => try_insert(
                 &mut self.descriptor_set,
@@ -259,9 +259,7 @@ impl CheckSpirvAttrVisitor<'_> {
             struct Expected<T>(T);
 
             let valid_target = match parsed_attr {
-                SpirvAttribute::StorageClass(_)
-                | SpirvAttribute::IntrinsicType(_)
-                | SpirvAttribute::Block => match target {
+                SpirvAttribute::IntrinsicType(_) | SpirvAttribute::Block => match target {
                     Target::Struct => {
                         // FIXME(eddyb) further check type attribute validity,
                         // e.g. layout, generics, other attributes, etc.
@@ -283,7 +281,8 @@ impl CheckSpirvAttrVisitor<'_> {
                     _ => Err(Expected("function")),
                 },
 
-                SpirvAttribute::Builtin(_)
+                SpirvAttribute::StorageClass(_)
+                | SpirvAttribute::Builtin(_)
                 | SpirvAttribute::DescriptorSet(_)
                 | SpirvAttribute::Binding(_)
                 | SpirvAttribute::Flat => match target {
@@ -298,6 +297,31 @@ impl CheckSpirvAttrVisitor<'_> {
                                 span,
                                 "attribute is only valid on a parameter of an entry-point function",
                             );
+                        } else {
+                            // FIXME(eddyb) should we just remove all 5 of these storage class
+                            // attributes, instead of disallowing them here?
+                            if let SpirvAttribute::StorageClass(storage_class) = parsed_attr {
+                                let valid = match storage_class {
+                                    StorageClass::Input | StorageClass::Output => {
+                                        Err("is the default and should not be explicitly specified")
+                                    }
+
+                                    StorageClass::Private
+                                    | StorageClass::Function
+                                    | StorageClass::Generic => {
+                                        Err("can not be used as part of an entry's interface")
+                                    }
+
+                                    _ => Ok(()),
+                                };
+
+                                if let Err(msg) = valid {
+                                    self.tcx.sess.span_err(
+                                        span,
+                                        &format!("`{:?}` storage class {}", storage_class, msg),
+                                    );
+                                }
+                            }
                         }
                         Ok(())
                     }

crates/rustc_codegen_spirv/src/codegen_cx/entry.rs

@@ -7,9 +7,10 @@ use rspirv::dr::Operand;
 use rspirv::spirv::{Decoration, ExecutionModel, FunctionControl, StorageClass, Word};
 use rustc_hir as hir;
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::ty::{Instance, Ty};
+use rustc_middle::ty::{Instance, Ty, TyKind};
 use rustc_span::Span;
 use rustc_target::abi::call::{FnAbi, PassMode};
+use rustc_target::abi::LayoutOf;
 use std::collections::HashMap;
 
 impl<'tcx> CodegenCx<'tcx> {
@@ -37,12 +38,12 @@ impl<'tcx> CodegenCx<'tcx> {
         let fn_hir_id = self.tcx.hir().local_def_id_to_hir_id(local_id);
         let body = self.tcx.hir().body(self.tcx.hir().body_owned_by(fn_hir_id));
         for (abi, arg) in fn_abi.args.iter().zip(body.params) {
-            if let PassMode::Direct(_) = abi.mode {
-            } else {
-                self.tcx.sess.span_err(
+            match abi.mode {
+                PassMode::Direct(_) | PassMode::Indirect { .. } => {}
+                _ => self.tcx.sess.span_err(
                     arg.span,
                     &format!("PassMode {:?} invalid for entry point parameter", abi.mode),
-                )
+                ),
             }
         }
         if let PassMode::Ignore = fn_abi.ret.mode {
@@ -104,7 +105,7 @@ impl<'tcx> CodegenCx<'tcx> {
         };
         let mut decoration_locations = HashMap::new();
         // Create OpVariables before OpFunction so they're global instead of local vars.
-        let arguments = entry_fn_abi
+        let declared_params = entry_fn_abi
             .args
             .iter()
             .zip(hir_params)
@@ -117,25 +118,52 @@ impl<'tcx> CodegenCx<'tcx> {
             .begin_function(void, None, FunctionControl::NONE, fn_void_void)
             .unwrap();
         emit.begin_block(None).unwrap();
+        // Adjust any global `OpVariable`s as needed (e.g. loading from `Input`s).
+        let arguments: Vec<_> = declared_params
+            .iter()
+            .zip(&entry_fn_abi.args)
+            .zip(hir_params)
+            .map(|((&(var, storage_class), entry_fn_arg), hir_param)| {
+                match entry_fn_arg.layout.ty.kind() {
+                    TyKind::Ref(..) => var,
+
+                    _ => match entry_fn_arg.mode {
+                        PassMode::Indirect { .. } => var,
+                        PassMode::Direct(_) => {
+                            assert_eq!(storage_class, StorageClass::Input);
+
+                            // NOTE(eddyb) this should never fail as it has to have
+                            // been already computed earlier by `declare_parameter`.
+                            let value_spirv_type =
+                                entry_fn_arg.layout.spirv_type(hir_param.span, self);
+
+                            emit.load(value_spirv_type, None, var, None, std::iter::empty())
+                                .unwrap()
+                        }
+                        _ => unreachable!(),
+                    },
+                }
+            })
+            .collect();
         emit.function_call(
             entry_func_return_type,
             None,
             entry_func.def_cx(self),
-            arguments.iter().map(|&(a, _)| a),
+            arguments,
         )
         .unwrap();
         emit.ret().unwrap();
         emit.end_function().unwrap();
 
         let interface: Vec<_> = if emit.version().unwrap() > (1, 3) {
             // SPIR-V >= v1.4 includes all OpVariables in the interface.
-            arguments.into_iter().map(|(a, _)| a).collect()
+            declared_params.into_iter().map(|(var, _)| var).collect()
         } else {
             // SPIR-V <= v1.3 only includes Input and Output in the interface.
-            arguments
+            declared_params
                 .into_iter()
                 .filter(|&(_, s)| s == StorageClass::Input || s == StorageClass::Output)
-                .map(|(a, _)| a)
+                .map(|(var, _)| var)
                 .collect()
         };
         emit.entry_point(execution_model, fn_id, name, interface);
@@ -148,59 +176,105 @@ impl<'tcx> CodegenCx<'tcx> {
         hir_param: &hir::Param<'tcx>,
         decoration_locations: &mut HashMap<StorageClass, u32>,
     ) -> (Word, StorageClass) {
-        let storage_class = crate::abi::get_storage_class(self, layout).unwrap_or_else(|| {
-            self.tcx.sess.span_fatal(
-                hir_param.span,
-                &format!("invalid entry param type `{}`", layout.ty),
-            );
-        });
-        let mut has_location = matches!(
-            storage_class,
-            StorageClass::Input | StorageClass::Output | StorageClass::UniformConstant
-        );
-        // Note: this *declares* the variable too.
-        let spirv_type = layout.spirv_type(hir_param.span, self);
-        let variable = self
-            .emit_global()
-            .variable(spirv_type, None, storage_class, None);
+        let attrs = AggregatedSpirvAttributes::parse(self, self.tcx.hir().attrs(hir_param.hir_id));
+
+        // FIXME(eddyb) attribute validation should be done ahead of time.
+        // FIXME(eddyb) also take into account `&T` interior mutability,
+        // i.e. it's only immutable if `T: Freeze`, which we should check.
+        // FIXME(eddyb) also check the type for compatibility with being
+        // part of the interface, including potentially `Sync`ness etc.
+        let (value_ty, storage_class) = if let Some(storage_class_attr) = attrs.storage_class {
+            let storage_class = storage_class_attr.value;
+            let expected_mutbl = match storage_class {
+                StorageClass::UniformConstant
+                | StorageClass::Input
+                | StorageClass::PushConstant => hir::Mutability::Not,
+
+                _ => hir::Mutability::Mut,
+            };
+
+            match *layout.ty.kind() {
+                TyKind::Ref(_, pointee_ty, m) if m == expected_mutbl => (pointee_ty, storage_class),
+
+                _ => self.tcx.sess.span_fatal(
+                    hir_param.span,
+                    &format!(
+                        "invalid entry param type `{}` for storage class `{:?}` \
+                         (expected `&{}T`)",
+                        layout.ty,
+                        storage_class,
+                        expected_mutbl.prefix_str()
+                    ),
+                ),
+            }
+        } else {
+            match *layout.ty.kind() {
+                TyKind::Ref(_, pointee_ty, hir::Mutability::Mut) => {
+                    (pointee_ty, StorageClass::Output)
+                }
+
+                TyKind::Ref(_, pointee_ty, hir::Mutability::Not) => self.tcx.sess.span_fatal(
+                    hir_param.span,
+                    &format!(
+                        "invalid entry param type `{}` (expected `{}` or `&mut {1}`)",
+                        layout.ty, pointee_ty
+                    ),
+                ),
+
+                _ => (layout.ty, StorageClass::Input),
+            }
+        };
+
+        // Pre-allocate the module-scoped `OpVariable`'s *Result* ID.
+        let variable = self.emit_global().id();
+
         if let hir::PatKind::Binding(_, _, ident, _) = &hir_param.pat.kind {
             self.emit_global().name(variable, ident.to_string());
         }
 
-        let attrs = AggregatedSpirvAttributes::parse(self, self.tcx.hir().attrs(hir_param.hir_id));
+        let mut decoration_supersedes_location = false;
         if let Some(builtin) = attrs.builtin.map(|attr| attr.value) {
             self.emit_global().decorate(
                 variable,
                 Decoration::BuiltIn,
                 std::iter::once(Operand::BuiltIn(builtin)),
             );
-            has_location = false;
+            decoration_supersedes_location = true;
         }
         if let Some(index) = attrs.descriptor_set.map(|attr| attr.value) {
             self.emit_global().decorate(
                 variable,
                 Decoration::DescriptorSet,
                 std::iter::once(Operand::LiteralInt32(index)),
             );
-            has_location = false;
+            decoration_supersedes_location = true;
         }
         if let Some(index) = attrs.binding.map(|attr| attr.value) {
             self.emit_global().decorate(
                 variable,
                 Decoration::Binding,
                 std::iter::once(Operand::LiteralInt32(index)),
             );
-            has_location = false;
+            decoration_supersedes_location = true;
         }
         if attrs.flat.is_some() {
             self.emit_global()
                 .decorate(variable, Decoration::Flat, std::iter::empty());
         }
 
+        // FIXME(eddyb) check whether the storage class is compatible with the
+        // specific shader stage of this entry-point, and any decorations
+        // (e.g. Vulkan has specific rules for builtin storage classes).
+
         // Assign locations from left to right, incrementing each storage class
         // individually.
         // TODO: Is this right for UniformConstant? Do they share locations with
         // input/outpus?
+        let has_location = !decoration_supersedes_location
+            && matches!(
+                storage_class,
+                StorageClass::Input | StorageClass::Output | StorageClass::UniformConstant
+            );
         if has_location {
             let location = decoration_locations
                 .entry(storage_class)
@@ -212,6 +286,15 @@ impl<'tcx> CodegenCx<'tcx> {
             );
             *location += 1;
         }
+
+        // Emit the `OpVariable` with its *Result* ID set to `variable`.
+        let var_spirv_type = SpirvType::Pointer {
+            pointee: self.layout_of(value_ty).spirv_type(hir_param.span, self),
+        }
+        .def(hir_param.span, self);
+        self.emit_global()
+            .variable(var_spirv_type, Some(variable), storage_class, None);
+
         (variable, storage_class)
     }