Add shader image query instructions

crates/rustc_codegen_spirv/src/builder/spirv_asm.rs

@@ -6,7 +6,7 @@ use crate::codegen_cx::CodegenCx;
 use rspirv::dr;
 use rspirv::grammar::{LogicalOperand, OperandKind, OperandQuantifier};
 use rspirv::spirv::{
-    FPFastMathMode, FragmentShadingRate, FunctionControl, ImageOperands, KernelProfilingInfo,
+    Dim, FPFastMathMode, FragmentShadingRate, FunctionControl, ImageOperands, KernelProfilingInfo,
     LoopControl, MemoryAccess, MemorySemantics, Op, RayFlags, SelectionControl, StorageClass, Word,
 };
 use rustc_ast::ast::{InlineAsmOptions, InlineAsmTemplatePiece};
@@ -321,6 +321,7 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
                 return;
             }
             _ => {
+                self.validate_instruction(&inst);
                 self.emit()
                     .insert_into_block(dr::InsertPoint::End, inst)
                     .unwrap();
@@ -1301,6 +1302,131 @@ impl<'cx, 'tcx> Builder<'cx, 'tcx> {
         }
         true
     }
+
+    pub fn validate_instruction(&mut self, inst: &dr::Instruction) {
+        fn find_image_ty<'cx, 'tcx>(
+            builder: &mut Builder<'cx, 'tcx>,
+            inst: &dr::Instruction,
+        ) -> Option<SpirvType> {
+            // Assumes the image parameter is the first operand
+            let image_obj = inst.operands[0].unwrap_id_ref();
+            let emit = builder.emit();
+            // Assumes the image's value definition is in the current block
+            let block = &emit.module_ref().functions[emit.selected_function().unwrap()].blocks
+                [emit.selected_block().unwrap()];
+            // Loop through the block to find the defining instruction
+            let defining_inst = match block
+                .instructions
+                .iter()
+                .find(|inst| inst.result_id == Some(image_obj))
+            {
+                Some(defining_inst) => defining_inst,
+                None => {
+                    // Something has gone wrong. All the asm! blocks using these instructions
+                    // should produce the image value in their own basic blocks (usually with
+                    // an OpLoad), so there's probably some typo somewhere with an error
+                    // already emitted, so just skip validation. If there truly is something
+                    // bad going on, spirv-val will catch it.
+                    return None;
+                }
+            };
+            match builder.lookup_type(defining_inst.result_type.unwrap()) {
+                SpirvType::SampledImage { image_type } => Some(builder.lookup_type(image_type)),
+                ty => Some(ty),
+            }
+        }
+
+        fn is_valid_query_size(ty: &SpirvType) -> bool {
+            match *ty {
+                SpirvType::Image {
+                    dim,
+                    multisampled,
+                    sampled,
+                    ..
+                } => match dim {
+                    Dim::Dim1D | Dim::Dim2D | Dim::Dim3D | Dim::DimCube => {
+                        multisampled == 1 || sampled == 0 || sampled == 2
+                    }
+                    Dim::DimBuffer | Dim::DimRect => true,
+                    Dim::DimSubpassData => false,
+                },
+                _ => true,
+            }
+        }
+
+        fn is_valid_query_size_lod(ty: &SpirvType) -> bool {
+            match *ty {
+                SpirvType::Image {
+                    dim, multisampled, ..
+                } => match dim {
+                    Dim::Dim1D | Dim::Dim2D | Dim::Dim3D | Dim::DimCube => multisampled == 0,
+                    _ => false,
+                },
+                _ => true,
+            }
+        }
+
+        match inst.class.opcode {
+            Op::ImageQueryLevels | Op::ImageQueryLod => {
+                let image_ty = match find_image_ty(self, inst) {
+                    Some(ty) => ty,
+                    None => return,
+                };
+                if let SpirvType::Image { dim, .. } = image_ty {
+                    match dim {
+                        Dim::Dim1D | Dim::Dim2D | Dim::Dim3D | Dim::DimCube => {}
+                        bad => self
+                            .struct_err(&format!(
+                                "Op{}'s image has a dimension of {:?}",
+                                inst.class.opname, bad
+                            ))
+                            .note("Allowed dimensions are 1D, 2D, 3D, and Cube")
+                            .emit(),
+                    }
+                }
+                // If the type isn't an image, something has gone wrong. The functions in image.rs
+                // shouldn't allow it, so the user is doing something weird. Let spirv-val handle
+                // the error later on.
+            }
+            Op::ImageQuerySize => {
+                let image_ty = match find_image_ty(self, inst) {
+                    Some(ty) => ty,
+                    None => return,
+                };
+                if !is_valid_query_size(&image_ty) {
+                    let mut err =
+                        self.struct_err("OpImageQuerySize is invalid for this image type");
+                    err.note(
+                        "allowed dimensions are 1D, 2D, 3D, Buffer, Rect, or Cube. \
+                              if dimension is 1D, 2D, 3D, or Cube, it must have either \
+                              multisampled be true, *or* sampled of Unknown or No",
+                    );
+                    if is_valid_query_size_lod(&image_ty) {
+                        err.note("query_size_lod is valid for this image, did you mean to use it instead?");
+                    }
+                    err.emit();
+                }
+            }
+            Op::ImageQuerySizeLod => {
+                let image_ty = match find_image_ty(self, inst) {
+                    Some(ty) => ty,
+                    None => return,
+                };
+                if !is_valid_query_size_lod(&image_ty) {
+                    let mut err =
+                        self.struct_err("OpImageQuerySizeLod is invalid for this image type");
+                    err.note("The image's dimension must be 1D, 2D, 3D, or Cube. Multisampled must be false.");
+                    if is_valid_query_size(&image_ty) {
+                        err.note(
+                            "query_size is valid for this image, did you mean to use it instead?",
+                        );
+                    }
+                    err.emit();
+                }
+            }
+            _ => {}
+        }
+    }
 }
 
 pub const IMAGE_OPERANDS: &[(&str, ImageOperands)] = &[

crates/spirv-std/src/image.rs

@@ -659,6 +659,185 @@ impl<
     }
 }
 
+impl<
+        SampledType: SampleType<FORMAT>,
+        const DIM: Dimensionality,
+        const DEPTH: ImageDepth,
+        const ARRAYED: Arrayed,
+        const MULTISAMPLED: Multisampled,
+        const SAMPLED: Sampled,
+        const FORMAT: ImageFormat,
+        const ACCESS_QUALIFIER: Option<AccessQualifier>,
+    > Image<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, ACCESS_QUALIFIER>
+{
+    /// Query the number of mipmap levels.
+    ///
+    /// Note: Const generics aren't able to reason about the constraints on this function (yet),
+    /// and so are enforced by the compiler. The constraints are:
+    ///
+    /// The image's dimension must be 1D, 2D, 3D, or Cube.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQueryLevels")]
+    pub fn query_levels(&self) -> u32 {
+        let result: u32;
+        unsafe {
+            asm! {
+                "%image = OpLoad _ {this}",
+                "{result} = OpImageQueryLevels typeof{result} %image",
+                this = in(reg) self,
+                result = out(reg) result,
+            }
+        }
+        result
+    }
+
+    /// Query the mipmap level and the level of detail for a hypothetical sampling of Image at
+    /// Coordinate using an implicit level of detail. The first component of the result contains
+    /// the mipmap array layer. The second component of the result contains the implicit level of
+    /// detail relative to the base level.
+    ///
+    /// Note: Const generics aren't able to reason about the constraints on this function (yet),
+    /// and so are enforced by the compiler. The constraints are:
+    ///
+    /// The image's dimension must be 1D, 2D, 3D, or Cube.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQueryLod")]
+    pub fn query_lod<V: Vector<f32, 2>>(
+        &self,
+        sampler: Sampler,
+        coord: impl ImageCoordinate<f32, DIM, { Arrayed::False }>,
+    ) -> V {
+        // Note: Arrayed::False isn't a typo in the ImageCoordinate, the spec states:
+        // Coordinate must be a scalar or vector of floating-point type or integer type. It
+        // contains (u[, v] ... ) as needed by the definition of Sampled Image, **not including any
+        // array layer index**. Unless the Kernel capability is being used, it must be floating
+        // point.
+        let mut result = Default::default();
+        unsafe {
+            asm! {
+                "%typeSampledImage = OpTypeSampledImage typeof*{this}",
+                "%image = OpLoad _ {this}",
+                "%sampler = OpLoad _ {sampler}",
+                "%coord = OpLoad _ {coord}",
+                "%sampledImage = OpSampledImage %typeSampledImage %image %sampler",
+                "%result = OpImageQueryLod typeof*{result} %sampledImage %coord",
+                "OpStore {result} %result",
+                result = in(reg) &mut result,
+                this = in(reg) self,
+                sampler = in(reg) &sampler,
+                coord = in(reg) &coord
+            }
+        }
+        result
+    }
+
+    /// Query the dimensions of Image, with no level of detail.
+    ///
+    /// Note: Const generics aren't able to reason about the constraints on this function (yet),
+    /// and so are enforced by the compiler. The constraints are:
+    ///
+    /// The image's dimension must be 1D, 2D, 3D, Buffer, Rect, or Cube. If dimension is 1D, 2D,
+    /// 3D, or Cube, it must have *either* multisampled be true, *or* sampled of Unknown or No.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQuerySize")]
+    pub fn query_size<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(&self) -> Size {
+        let mut result: Size = Default::default();
+        unsafe {
+            asm! {
+                "%image = OpLoad _ {this}",
+                "%result = OpImageQuerySize typeof*{result} %image",
+                "OpStore {result} %result",
+                this = in(reg) self,
+                result = in(reg) &mut result,
+            }
+        }
+        result
+    }
+}
+
+impl<
+        SampledType: SampleType<FORMAT>,
+        const DIM: Dimensionality,
+        const DEPTH: ImageDepth,
+        const ARRAYED: Arrayed,
+        const SAMPLED: Sampled,
+        const FORMAT: ImageFormat,
+        const ACCESS_QUALIFIER: Option<AccessQualifier>,
+    >
+    Image<
+        SampledType,
+        DIM,
+        DEPTH,
+        ARRAYED,
+        { Multisampled::False },
+        SAMPLED,
+        FORMAT,
+        ACCESS_QUALIFIER,
+    >
+{
+    /// Query the dimensions of Image, with no level of detail.
+    ///
+    /// Note: Const generics aren't able to reason about the constraints on this function (yet),
+    /// and so are enforced by the compiler. The constraints are:
+    ///
+    /// The image's dimension must be 1D, 2D, 3D, or Cube. Multisampled must be false.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQuerySizeLod")]
+    pub fn query_size_lod<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(
+        &self,
+        lod: u32,
+    ) -> Size {
+        let mut result: Size = Default::default();
+        unsafe {
+            asm! {
+                "%image = OpLoad _ {this}",
+                "%result = OpImageQuerySizeLod typeof*{result} %image {lod}",
+                "OpStore {result} %result",
+                this = in(reg) self,
+                lod = in(reg) lod,
+                result = in(reg) &mut result,
+            }
+        }
+        result
+    }
+}
+
+impl<
+        SampledType: SampleType<FORMAT>,
+        const DEPTH: ImageDepth,
+        const ARRAYED: Arrayed,
+        const SAMPLED: Sampled,
+        const FORMAT: ImageFormat,
+        const ACCESS_QUALIFIER: Option<AccessQualifier>,
+    >
+    Image<
+        SampledType,
+        { Dimensionality::TwoD },
+        DEPTH,
+        ARRAYED,
+        { Multisampled::True },
+        SAMPLED,
+        FORMAT,
+        ACCESS_QUALIFIER,
+    >
+{
+    /// Query the number of samples available per texel fetch in a multisample image.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQuerySamples")]
+    pub fn query_samples(&self) -> u32 {
+        let result: u32;
+        unsafe {
+            asm! {
+                "%image = OpLoad _ {this}",
+                "{result} = OpImageQuerySamples typeof{result} %image",
+                this = in(reg) self,
+                result = out(reg) result,
+            }
+        }
+        result
+    }
+}
+
 /// An image combined with a sampler, enabling filtered accesses of the
 /// image’s contents.
 #[spirv(sampled_image)]

tests/ui/image/query/query_levels.rs

@@ -0,0 +1,12 @@
+// build-pass
+
+use spirv_std::{arch, Image};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled),
+    output: &mut u32,
+) {
+    unsafe { asm!("OpCapability ImageQuery") };
+    *output = image.query_levels();
+}

tests/ui/image/query/query_levels_err.rs

@@ -0,0 +1,13 @@
+// build-fail
+// normalize-stderr-test "\S*/crates/spirv-std/src/" -> "$$SPIRV_STD_SRC/"
+
+use spirv_std::{arch, Image};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(rect, type=f32, sampled),
+    output: &mut u32,
+) {
+    unsafe { asm!("OpCapability ImageQuery") };
+    *output = image.query_levels();
+}

tests/ui/image/query/query_levels_err.stderr

@@ -0,0 +1,15 @@
+error: OpImageQueryLevels's image has a dimension of DimRect
+   --> $SPIRV_STD_SRC/image.rs:684:13
+    |
+684 | /             asm! {
+685 | |                 "%image = OpLoad _ {this}",
+686 | |                 "{result} = OpImageQueryLevels typeof{result} %image",
+687 | |                 this = in(reg) self,
+688 | |                 result = out(reg) result,
+689 | |             }
+    | |_____________^
+    |
+    = note: Allowed dimensions are 1D, 2D, 3D, and Cube
+
+error: aborting due to previous error
+

tests/ui/image/query/query_lod.rs

@@ -0,0 +1,13 @@
+// build-pass
+
+use spirv_std::{arch, Image, Sampler};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] image: &Image!(2D, type=f32, sampled),
+    #[spirv(descriptor_set = 0, binding = 1)] sampler: &Sampler,
+    output: &mut glam::Vec2,
+) {
+    unsafe { asm!("OpCapability ImageQuery") };
+    *output = image.query_lod(*sampler, glam::Vec2::new(0.0, 1.0));
+}