Add support for sampled images

crates/rustc_codegen_spirv/src/abi.rs

@@ -8,6 +8,7 @@ use rspirv::spirv::{StorageClass, Word};
 use rustc_middle::bug;
 use rustc_middle::ty::layout::{FnAbiExt, TyAndLayout};
 use rustc_middle::ty::{GeneratorSubsts, PolyFnSig, Ty, TyKind, TypeAndMut};
+use rustc_middle::ty::subst::SubstsRef;
 use rustc_span::Span;
 use rustc_target::abi::call::{CastTarget, FnAbi, PassMode, Reg, RegKind};
 use rustc_target::abi::{
@@ -323,7 +324,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() {
@@ -354,7 +355,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;
  }
  }
@@ -779,6 +780,7 @@ fn trans_image<'tcx>(
  cx: &CodegenCx<'tcx>,
  span: Span,
  ty: TyAndLayout<'tcx>,
+ substs: SubstsRef<'tcx>,
  attr: SpirvAttribute,
 ) -> Option<Word> {
  match attr {
@@ -818,6 +820,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 {
@@ -222,6 +225,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
@@ -278,6 +282,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)
  }
@@ -303,6 +308,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(),
  }
  }
 }
@@ -440,6 +446,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();
@@ -574,6 +585,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
+ }
+ }
+}