Refactor Image to accept const generic parameters

Author: XAMPPRocky

Cargo.lock

@@ -29,6 +29,7 @@ use-compiled-tools = ["spirv-tools/use-compiled-tools"]
 [dependencies]
 bimap = "0.5"
 indexmap = "1.6.0"
+num-traits = "0.2.14"
 rspirv = { git = "https://github.com/gfx-rs/rspirv.git", rev = "01ca0d2e5b667a0e4ff1bc1804511e38f9a08759" }
 rustc-demangle = "0.1.18"
 serde = { version = "1.0", features = ["derive"] }

crates/rustc_codegen_spirv/Cargo.toml

@@ -8,7 +8,7 @@ use rspirv::spirv::{Capability, StorageClass, Word};
 use rustc_middle::bug;
 use rustc_middle::ty::layout::{FnAbiExt, TyAndLayout};
 use rustc_middle::ty::subst::SubstsRef;
-use rustc_middle::ty::{GeneratorSubsts, PolyFnSig, Ty, TyKind, TypeAndMut};
+use rustc_middle::ty::{GeneratorSubsts, ParamEnv, PolyFnSig, Ty, TyKind, TypeAndMut};
 use rustc_span::Span;
 use rustc_target::abi::call::{CastTarget, FnAbi, PassMode, Reg, RegKind};
 use rustc_target::abi::{
@@ -20,6 +20,9 @@ use std::collections::HashMap;
 use std::fmt;
 use std::fmt::Write;
 
+use num_traits::FromPrimitive;
+use rspirv::spirv;
+
 /// If a struct contains a pointer to itself, even indirectly, then doing a naiive recursive walk
 /// of the fields will result in an infinite loop. Because pointers are the only thing that are
 /// allowed to be recursive, keep track of what pointers we've translated, or are currently in the
@@ -797,28 +800,77 @@ fn trans_image<'tcx>(
     attr: SpirvAttribute,
 ) -> Option<Word> {
     match attr {
-        SpirvAttribute::Image {
-            dim,
-            depth,
-            arrayed,
-            multisampled,
-            sampled,
-            image_format,
-            access_qualifier,
-        } => {
+        SpirvAttribute::Image => {
             // see SpirvType::sizeof
             if ty.size != Size::from_bytes(4) {
                 cx.tcx.sess.err("#[spirv(image)] type must have size 4");
                 return None;
             }
+
             // Hardcode to float for now
             let sampled_type = SpirvType::Float(32).def(span, cx);
+
+            macro_rules! type_from_variant_index {
+                ($(let $name:ident : $ty:ty = $exp:expr);+ ;) => {
+                    $(
+                        let $name = <$ty>::from_u64(
+                            cx.tcx
+                            .destructure_const(ParamEnv::reveal_all().and($exp))
+                            .variant
+                            .unwrap()
+                            .as_u32() as u64,
+                        )
+                        .unwrap();
+                    )+
+                }
+            }
+
+            type_from_variant_index! {
+                let dim: spirv::Dim = substs.const_at(0);
+                let depth: u32 = substs.const_at(1);
+                let sampled: u32 = substs.const_at(4);
+                let image_format: spirv::ImageFormat = substs.const_at(5);
+            }
+
+            let arrayed: bool = substs
+                .const_at(2)
+                .val
+                .try_to_value()
+                .unwrap()
+                .try_to_bool()
+                .unwrap();
+            let multisampled: bool = substs
+                .const_at(3)
+                .val
+                .try_to_value()
+                .unwrap()
+                .try_to_bool()
+                .unwrap();
+            let access_qualifier = {
+                let option = cx
+                    .tcx
+                    .destructure_const(ParamEnv::reveal_all().and(substs.const_at(6)));
+
+                match option.variant.map(|i| i.as_u32()).unwrap_or(0) {
+                    0 => None,
+                    1 => spirv::AccessQualifier::from_u64(
+                        option.fields[0]
+                            .val
+                            .try_to_scalar()
+                            .unwrap()
+                            .to_u64()
+                            .unwrap(),
+                    ),
+                    _ => unreachable!(),
+                }
+            };
+
             let ty = SpirvType::Image {
                 sampled_type,
                 dim,
                 depth,
-                arrayed,
-                multisampled,
+                arrayed: arrayed as u32,
+                multisampled: multisampled as u32,
                 sampled,
                 image_format,
                 access_qualifier,

crates/rustc_codegen_spirv/src/abi.rs

@@ -1,8 +1,6 @@
 use crate::builder::libm_intrinsics;
 use crate::codegen_cx::CodegenCx;
-use rspirv::spirv::{
-    AccessQualifier, BuiltIn, Dim, ExecutionMode, ExecutionModel, ImageFormat, StorageClass,
-};
+use rspirv::spirv::{BuiltIn, ExecutionMode, ExecutionModel, StorageClass};
 use rustc_ast::ast::{AttrKind, Attribute, Lit, LitIntType, LitKind, NestedMetaItem};
 use rustc_span::symbol::{Ident, Symbol};
 use std::collections::HashMap;
@@ -33,13 +31,6 @@ pub struct Symbols {
     descriptor_set: Symbol,
     binding: Symbol,
     image: Symbol,
-    dim: Symbol,
-    depth: Symbol,
-    arrayed: Symbol,
-    multisampled: Symbol,
-    sampled: Symbol,
-    image_format: Symbol,
-    access_qualifier: Symbol,
     attributes: HashMap<Symbol, SpirvAttribute>,
     execution_modes: HashMap<Symbol, (ExecutionMode, ExecutionModeExtraDim)>,
     pub libm_intrinsics: HashMap<Symbol, libm_intrinsics::LibmIntrinsic>,
@@ -386,13 +377,6 @@ impl Symbols {
             descriptor_set: Symbol::intern("descriptor_set"),
             binding: Symbol::intern("binding"),
             image: Symbol::intern("image"),
-            dim: Symbol::intern("dim"),
-            depth: Symbol::intern("depth"),
-            arrayed: Symbol::intern("arrayed"),
-            multisampled: Symbol::intern("multisampled"),
-            sampled: Symbol::intern("sampled"),
-            image_format: Symbol::intern("image_format"),
-            access_qualifier: Symbol::intern("access_qualifier"),
             attributes,
             execution_modes,
             libm_intrinsics,
@@ -445,15 +429,7 @@ pub enum SpirvAttribute {
     DescriptorSet(u32),
     Binding(u32),
     ReallyUnsafeIgnoreBitcasts,
-    Image {
-        dim: Dim,
-        depth: u32,
-        arrayed: u32,
-        multisampled: u32,
-        sampled: u32,
-        image_format: ImageFormat,
-        access_qualifier: Option<AccessQualifier>,
-    },
+    Image,
     Sampler,
     SampledImage,
     Block,
@@ -493,7 +469,7 @@ pub fn parse_attrs(
             };
             args.into_iter().filter_map(move |ref arg| {
                 if arg.has_name(cx.sym.image) {
-                    parse_image(cx, arg)
+                    Some(SpirvAttribute::Image)
                 } else if arg.has_name(cx.sym.descriptor_set) {
                     match parse_attr_int_value(cx, arg) {
                         Some(x) => Some(SpirvAttribute::DescriptorSet(x)),
@@ -537,136 +513,6 @@ pub fn parse_attrs(
     result.collect::<Vec<_>>().into_iter()
 }
 
-fn parse_image(cx: &CodegenCx<'_>, attr: &NestedMetaItem) -> Option<SpirvAttribute> {
-    let args = match attr.meta_item_list() {
-        Some(args) => args,
-        None => {
-            cx.tcx
-                .sess
-                .span_err(attr.span(), "image attribute must have arguments");
-            return None;
-        }
-    };
-    if args.len() != 6 && args.len() != 7 {
-        cx.tcx
-            .sess
-            .span_err(attr.span(), "image attribute must have 6 or 7 arguments");
-        return None;
-    }
-    let check = |idx: usize, sym: Symbol| -> bool {
-        if args[idx].has_name(sym) {
-            false
-        } else {
-            cx.tcx.sess.span_err(
-                args[idx].span(),
-                &format!("image attribute argument {} must be {}=...", idx + 1, sym),
-            );
-            true
-        }
-    };
-    if check(0, cx.sym.dim)
-        | check(1, cx.sym.depth)
-        | check(2, cx.sym.arrayed)
-        | check(3, cx.sym.multisampled)
-        | check(4, cx.sym.sampled)
-        | check(5, cx.sym.image_format)
-        | (args.len() == 7 && check(6, cx.sym.access_qualifier))
-    {
-        return None;
-    }
-    let arg_values = args
-        .iter()
-        .map(
-            |arg| match arg.meta_item().and_then(|arg| arg.name_value_literal()) {
-                Some(arg) => Some(arg),
-                None => {
-                    cx.tcx
-                        .sess
-                        .span_err(arg.span(), "image attribute must be name=value");
-                    None
-                }
-            },
-        )
-        .collect::<Option<Vec<_>>>()?;
-    let dim = match arg_values[0].kind {
-        LitKind::Str(dim, _) => match dim.with(|s| s.parse()) {
-            Ok(dim) => dim,
-            Err(()) => {
-                cx.tcx.sess.span_err(args[0].span(), "invalid dim value");
-                return None;
-            }
-        },
-        _ => {
-            cx.tcx
-                .sess
-                .span_err(args[0].span(), "dim value must be str");
-            return None;
-        }
-    };
-    let parse_lit = |idx: usize, name: &str| -> Option<u32> {
-        match arg_values[idx].kind {
-            LitKind::Int(v, _) => Some(v as u32),
-            _ => {
-                cx.tcx
-                    .sess
-                    .span_err(args[idx].span(), &format!("{} value must be int", name));
-                None
-            }
-        }
-    };
-    let depth = parse_lit(1, "depth")?;
-    let arrayed = parse_lit(2, "arrayed")?;
-    let multisampled = parse_lit(3, "multisampled")?;
-    let sampled = parse_lit(4, "sampled")?;
-    let image_format = match arg_values[5].kind {
-        LitKind::Str(dim, _) => match dim.with(|s| s.parse()) {
-            Ok(dim) => dim,
-            Err(()) => {
-                cx.tcx
-                    .sess
-                    .span_err(args[5].span(), "invalid image_format value");
-                return None;
-            }
-        },
-        _ => {
-            cx.tcx
-                .sess
-                .span_err(args[5].span(), "image_format value must be str");
-            return None;
-        }
-    };
-    let access_qualifier = if args.len() == 7 {
-        Some(match arg_values[6].kind {
-            LitKind::Str(dim, _) => match dim.with(|s| s.parse()) {
-                Ok(dim) => dim,
-                Err(()) => {
-                    cx.tcx
-                        .sess
-                        .span_err(args[6].span(), "invalid access_qualifier value");
-                    return None;
-                }
-            },
-            _ => {
-                cx.tcx
-                    .sess
-                    .span_err(args[6].span(), "access_qualifier value must be str");
-                return None;
-            }
-        })
-    } else {
-        None
-    };
-    Some(SpirvAttribute::Image {
-        dim,
-        depth,
-        arrayed,
-        multisampled,
-        sampled,
-        image_format,
-        access_qualifier,
-    })
-}
-
 fn parse_attr_int_value(cx: &CodegenCx<'_>, arg: &NestedMetaItem) -> Option<u32> {
     let arg = match arg.meta_item() {
         Some(arg) => arg,

crates/rustc_codegen_spirv/src/symbols.rs

@@ -143,7 +143,7 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
         format!(" -C target-feature={}", target_features.join(","))
     };
     let rustflags = format!(
-        "-Z codegen-backend={} -Z symbol-mangling-version=v0{}",
+        "-Z codegen-backend={} {}",
         rustc_codegen_spirv.display(),
         feature_flag,
     );

crates/spirv-builder/src/lib.rs

@@ -7,6 +7,12 @@ license = "MIT OR Apache-2.0"
 repository = "https://github.com/EmbarkStudios/rust-gpu"
 description = "Standard functions and types for SPIR-V"
 
+[profile.dev]
+incremental = false
+
+[profile.release]
+incremental = false
+
 [dependencies]
 glam = { version = "0.11.3", default-features = false, features = ["libm", "scalar-math"] }
 num-traits = { version = "0.2.14", default-features = false }

crates/spirv-std/Cargo.toml

@@ -38,18 +38,31 @@
     future_incompatible,
     nonstandard_style
 )]
+#![feature(const_generics)]
+#![allow(incomplete_features)]
 
 #[cfg(not(target_arch = "spirv"))]
 #[macro_use]
 pub extern crate spirv_std_macros;
 
+pub(crate) mod sealed;
 pub mod storage_class;
 mod textures;
 
 pub use glam;
 pub use num_traits;
 pub use textures::*;
 
+/// Marker trait for arguments that accept single scalar values or vectors
+/// of scalars.
+pub trait ScalarOrVector: sealed::Sealed {}
+
+impl ScalarOrVector for f32 {}
+impl ScalarOrVector for glam::Vec2 {}
+impl ScalarOrVector for glam::Vec3 {}
+impl ScalarOrVector for glam::Vec3A {}
+impl ScalarOrVector for glam::Vec4 {}
+
 #[cfg(all(not(test), target_arch = "spirv"))]
 #[panic_handler]
 fn panic(_: &core::panic::PanicInfo<'_>) -> ! {

crates/spirv-std/src/lib.rs

@@ -0,0 +1,9 @@
+/// Sealed trait to ensure certain traits can't be implemented outside
+/// of `spirv-std`.
+pub trait Sealed {}
+
+impl Sealed for f32 {}
+impl Sealed for glam::Vec2 {}
+impl Sealed for glam::Vec3 {}
+impl Sealed for glam::Vec3A {}
+impl Sealed for glam::Vec4 {}