Add support for sampled images (#320)

* Add support for sampled images

Combined image samplers are allowed as resources and require a generic parameter indicating the underlying image type.

* fix formatting issues

* address review comments

* fix formatting issues

crates/rustc_codegen_spirv/src/abi.rs

@@ -7,6 +7,7 @@ use crate::symbols::{parse_attrs, SpirvAttribute};
 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_span::Span;
 use rustc_target::abi::call::{CastTarget, FnAbi, PassMode, Reg, RegKind};
@@ -336,7 +337,7 @@ fn trans_type_impl<'tcx>(
     ty: TyAndLayout<'tcx>,
     is_immediate: bool,
 ) -> Word {
-    if let TyKind::Adt(adt, _) = *ty.ty.kind() {
+    if let TyKind::Adt(adt, substs) = *ty.ty.kind() {
         for attr in parse_attrs(cx, cx.tcx.get_attrs(adt.did)) {
             if matches!(attr, SpirvAttribute::Block) {
                 if !adt.is_struct() {
@@ -367,7 +368,7 @@ fn trans_type_impl<'tcx>(
                 return trans_struct(cx, span, ty, true);
             }
 
-            if let Some(image) = trans_image(cx, span, ty, attr) {
+            if let Some(image) = trans_image(cx, span, ty, substs, attr) {
                 return image;
             }
         }
@@ -792,6 +793,7 @@ fn trans_image<'tcx>(
     cx: &CodegenCx<'tcx>,
     span: Span,
     ty: TyAndLayout<'tcx>,
+    substs: SubstsRef<'tcx>,
     attr: SpirvAttribute,
 ) -> Option<Word> {
     match attr {
@@ -831,6 +833,28 @@ fn trans_image<'tcx>(
             }
             Some(SpirvType::Sampler.def(span, cx))
         }
+        SpirvAttribute::SampledImage => {
+            // see SpirvType::sizeof
+            if ty.size != Size::from_bytes(4) {
+                cx.tcx
+                    .sess
+                    .err("#[spirv(sampled_image)] type must have size 4");
+                return None;
+            }
+
+            // We use a generic to indicate the underlying image type of the sampled image.
+            // The spirv type of it will be generated by querying the type of the first generic.
+            if let Some(image_ty) = substs.types().next() {
+                // TODO: enforce that the generic param is an image type?
+                let image_type = trans_type_impl(cx, span, cx.layout_of(image_ty), false);
+                Some(SpirvType::SampledImage { image_type }.def(span, cx))
+            } else {
+                cx.tcx
+                    .sess
+                    .err("#[spirv(sampled_image)] type must have a generic image type");
+                None
+            }
+        }
         _ => None,
     }
 }

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -198,6 +198,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             SpirvType::Function { .. } => self.fatal("memset on functions not implemented yet"),
             SpirvType::Image { .. } => self.fatal("cannot memset image"),
             SpirvType::Sampler => self.fatal("cannot memset sampler"),
+            SpirvType::SampledImage { .. } => self.fatal("cannot memset sampled image"),
         }
     }
 
@@ -253,6 +254,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             SpirvType::Function { .. } => self.fatal("memset on functions not implemented yet"),
             SpirvType::Image { .. } => self.fatal("cannot memset image"),
             SpirvType::Sampler => self.fatal("cannot memset sampler"),
+            SpirvType::SampledImage { .. } => self.fatal("cannot memset sampled image"),
         }
     }
 

crates/rustc_codegen_spirv/src/codegen_cx/constant.rs

@@ -501,6 +501,10 @@ impl<'tcx> CodegenCx<'tcx> {
                 .tcx
                 .sess
                 .fatal("Cannot create a constant sampler value"),
+            SpirvType::SampledImage { .. } => self
+                .tcx
+                .sess
+                .fatal("Cannot create a constant sampled image value"),
         }
     }
 

crates/rustc_codegen_spirv/src/codegen_cx/type_.rs

@@ -177,6 +177,7 @@ impl<'tcx> BaseTypeMethods<'tcx> for CodegenCx<'tcx> {
             SpirvType::Function { .. } => TypeKind::Function,
             SpirvType::Image { .. } => TypeKind::Integer,
             SpirvType::Sampler => TypeKind::Integer,
+            SpirvType::SampledImage { .. } => TypeKind::Integer,
         }
     }
     fn type_ptr_to(&self, ty: Self::Type) -> Self::Type {

crates/rustc_codegen_spirv/src/spirv_type.rs

@@ -71,6 +71,9 @@ pub enum SpirvType {
         access_qualifier: Option<AccessQualifier>,
     },
     Sampler,
+    SampledImage {
+        image_type: Word,
+    },
 }
 
 impl SpirvType {
@@ -226,6 +229,7 @@ impl SpirvType {
                 access_qualifier,
             ),
             Self::Sampler => cx.emit_global().type_sampler(),
+            Self::SampledImage { image_type } => cx.emit_global().type_sampled_image(image_type),
         };
         cx.type_cache.def(result, self);
         result
@@ -293,6 +297,7 @@ impl SpirvType {
             Self::Function { .. } => cx.tcx.data_layout.pointer_size,
             Self::Image { .. } => Size::from_bytes(4),
             Self::Sampler => Size::from_bytes(4),
+            Self::SampledImage { .. } => Size::from_bytes(4),
         };
         Some(result)
     }
@@ -318,6 +323,7 @@ impl SpirvType {
             Self::Function { .. } => cx.tcx.data_layout.pointer_align.abi,
             Self::Image { .. } => Align::from_bytes(4).unwrap(),
             Self::Sampler => Align::from_bytes(4).unwrap(),
+            Self::SampledImage { .. } => Align::from_bytes(4).unwrap(),
         }
     }
 }
@@ -455,6 +461,11 @@ impl fmt::Debug for SpirvTypePrinter<'_, '_> {
                 .field("access_qualifier", &access_qualifier)
                 .finish(),
             SpirvType::Sampler => f.debug_struct("Sampler").field("id", &self.id).finish(),
+            SpirvType::SampledImage { image_type } => f
+                .debug_struct("SampledImage")
+                .field("id", &self.id)
+                .field("image_type", &self.cx.debug_type(image_type))
+                .finish(),
         };
         {
             let mut debug_stack = DEBUG_STACK.lock().unwrap();
@@ -589,6 +600,10 @@ impl SpirvTypePrinter<'_, '_> {
                 .field("access_qualifier", &access_qualifier)
                 .finish(),
             SpirvType::Sampler => f.write_str("Sampler"),
+            SpirvType::SampledImage { image_type } => f
+                .debug_struct("SampledImage")
+                .field("image_type", &self.cx.debug_type(image_type))
+                .finish(),
         }
     }
 }

crates/rustc_codegen_spirv/src/symbols.rs

@@ -338,6 +338,7 @@ impl Symbols {
             ("sampler", SpirvAttribute::Sampler),
             ("block", SpirvAttribute::Block),
             ("flat", SpirvAttribute::Flat),
+            ("sampled_image", SpirvAttribute::SampledImage),
         ]
         .iter()
         .cloned();
@@ -453,6 +454,7 @@ pub enum SpirvAttribute {
         access_qualifier: Option<AccessQualifier>,
     },
     Sampler,
+    SampledImage,
     Block,
     Flat,
 }

crates/spirv-std/src/textures.rs

@@ -94,3 +94,34 @@ impl Image2dArray {
         }
     }
 }
+
+#[allow(unused_attributes)]
+#[spirv(sampled_image)]
+#[derive(Copy, Clone)]
+pub struct SampledImage<I> {
+    _image: I,
+}
+
+impl SampledImage<Image2d> {
+    pub fn sample(&self, coord: Vec3A) -> Vec4 {
+        #[cfg(not(target_arch = "spirv"))]
+        {
+            let _ = coord;
+            panic!("Image sampling not supported on CPU");
+        }
+        #[cfg(target_arch = "spirv")]
+        unsafe {
+            let mut result = Default::default();
+            asm!(
+                "%sampledImage = OpLoad typeof*{1} {1}",
+                "%coord = OpLoad typeof*{2} {2}",
+                "%result = OpImageSampleImplicitLod typeof*{0} %sampledImage %coord",
+                "OpStore {0} %result",
+                in(reg) &mut result,
+                in(reg) self,
+                in(reg) &coord
+            );
+            result
+        }
+    }
+}