Add query_size_lod and query_size methods to SampledImage

- Add query_size_lod for non-multisampled sampled images (Sampled=1)
- Add query_size for multisampled sampled images
- Add missing HasQuerySize trait implementations for Sampled::Yes
- Add tests for both methods

Author: LegNeato

crates/spirv-std/src/image.rs

@@ -1117,6 +1117,94 @@ impl<
         }
         result.truncate_into()
     }
+
+    /// Query the dimensions of the image at the specified level of detail.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQuerySizeLod")]
+    pub fn query_size_lod<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(
+        &self,
+        lod: u32,
+    ) -> Size
+    where
+        Image<
+            SampledType,
+            DIM,
+            DEPTH,
+            ARRAYED,
+            { Multisampled::False as u32 },
+            SAMPLED,
+            FORMAT,
+            COMPONENTS,
+        >: HasQuerySizeLod,
+    {
+        let mut result: Size = Default::default();
+        unsafe {
+            asm! {
+                "%sampledImage = OpLoad _ {this}",
+                "%image = OpImage _ %sampledImage",
+                "%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, COMPONENTS>,
+    const DIM: u32,
+    const DEPTH: u32,
+    const ARRAYED: u32,
+    const SAMPLED: u32,
+    const FORMAT: u32,
+    const COMPONENTS: u32,
+>
+    SampledImage<
+        Image<
+            SampledType,
+            DIM,
+            DEPTH,
+            ARRAYED,
+            { Multisampled::True as u32 },
+            SAMPLED,
+            FORMAT,
+            COMPONENTS,
+        >,
+    >
+{
+    /// Query the dimensions of the image, with no level of detail.
+    /// Available only for multisampled images.
+    #[crate::macros::gpu_only]
+    #[doc(alias = "OpImageQuerySize")]
+    pub fn query_size<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(&self) -> Size
+    where
+        Image<
+            SampledType,
+            DIM,
+            DEPTH,
+            ARRAYED,
+            { Multisampled::True as u32 },
+            SAMPLED,
+            FORMAT,
+            COMPONENTS,
+        >: HasQuerySize,
+    {
+        let mut result: Size = Default::default();
+        unsafe {
+            asm! {
+                "%sampledImage = OpLoad _ {this}",
+                "%image = OpImage _ %sampledImage",
+                "%result = OpImageQuerySize typeof*{result} %image",
+                "OpStore {result} %result",
+                this = in(reg) self,
+                result = in(reg) &mut result,
+            }
+        }
+        result
+    }
 }
 
 /// Helper trait that defines all `*_with` methods on an `Image` that use the extra image operands,
@@ -1633,6 +1721,25 @@ impl<
     >
 {
 }
+impl<
+    SampledType: SampleType<FORMAT, COMPONENTS>,
+    const DEPTH: u32,
+    const FORMAT: u32,
+    const ARRAYED: u32,
+    const COMPONENTS: u32,
+> HasQuerySize
+    for Image<
+        SampledType,
+        { Dimensionality::OneD as u32 },
+        DEPTH,
+        ARRAYED,
+        { Multisampled::False as u32 },
+        { Sampled::Yes as u32 },
+        FORMAT,
+        COMPONENTS,
+    >
+{
+}
 impl<
     SampledType: SampleType<FORMAT, COMPONENTS>,
     const DEPTH: u32,
@@ -1691,6 +1798,25 @@ impl<
     >
 {
 }
+impl<
+    SampledType: SampleType<FORMAT, COMPONENTS>,
+    const DEPTH: u32,
+    const FORMAT: u32,
+    const ARRAYED: u32,
+    const COMPONENTS: u32,
+> HasQuerySize
+    for Image<
+        SampledType,
+        { Dimensionality::TwoD as u32 },
+        DEPTH,
+        ARRAYED,
+        { Multisampled::False as u32 },
+        { Sampled::Yes as u32 },
+        FORMAT,
+        COMPONENTS,
+    >
+{
+}
 impl<
     SampledType: SampleType<FORMAT, COMPONENTS>,
     const DEPTH: u32,
@@ -1749,6 +1875,25 @@ impl<
     >
 {
 }
+impl<
+    SampledType: SampleType<FORMAT, COMPONENTS>,
+    const DEPTH: u32,
+    const FORMAT: u32,
+    const ARRAYED: u32,
+    const COMPONENTS: u32,
+> HasQuerySize
+    for Image<
+        SampledType,
+        { Dimensionality::ThreeD as u32 },
+        DEPTH,
+        ARRAYED,
+        { Multisampled::False as u32 },
+        { Sampled::Yes as u32 },
+        FORMAT,
+        COMPONENTS,
+    >
+{
+}
 impl<
     SampledType: SampleType<FORMAT, COMPONENTS>,
     const DEPTH: u32,
@@ -1807,6 +1952,25 @@ impl<
     >
 {
 }
+impl<
+    SampledType: SampleType<FORMAT, COMPONENTS>,
+    const DEPTH: u32,
+    const FORMAT: u32,
+    const ARRAYED: u32,
+    const COMPONENTS: u32,
+> HasQuerySize
+    for Image<
+        SampledType,
+        { Dimensionality::Cube as u32 },
+        DEPTH,
+        ARRAYED,
+        { Multisampled::False as u32 },
+        { Sampled::Yes as u32 },
+        FORMAT,
+        COMPONENTS,
+    >
+{
+}
 impl<
     SampledType: SampleType<FORMAT, COMPONENTS>,
     const DEPTH: u32,

tests/compiletests/ui/image/query/sampled_image_multisampled_query_size.rs

@@ -0,0 +1,17 @@
+// Test `OpImageQuerySize` on multisampled `SampledImage`
+// build-pass
+// compile-flags: -C target-feature=+ImageQuery
+
+use spirv_std::spirv;
+use spirv_std::{Image, arch, image::SampledImage};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] sampled_image2d_ms: &SampledImage<
+        Image!(2D, type=f32, multisampled, sampled),
+    >,
+    output: &mut glam::UVec2,
+) {
+    // Multisampled sampled images can use query_size directly
+    *output = sampled_image2d_ms.query_size();
+}

tests/compiletests/ui/image/query/sampled_image_query_size.rs

@@ -0,0 +1,27 @@
+// Test `OpImageQuerySizeLod` on `SampledImage`
+// build-pass
+// compile-flags: -C target-feature=+ImageQuery
+
+use spirv_std::spirv;
+use spirv_std::{Image, arch, image::SampledImage};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] sampled_image2d: &SampledImage<
+        Image!(2D, type=f32, sampled),
+    >,
+    #[spirv(descriptor_set = 1, binding = 1)] sampled_image2d_array: &SampledImage<
+        Image!(2D, type=f32, arrayed, sampled),
+    >,
+    #[spirv(descriptor_set = 2, binding = 2)] sampled_image3d: &SampledImage<
+        Image!(3D, type=f32, sampled),
+    >,
+    output: &mut glam::UVec3,
+) {
+    // For sampled images, we need to use query_size_lod
+    let size_2d: glam::UVec2 = sampled_image2d.query_size_lod(0);
+    let size_2d_array: glam::UVec3 = sampled_image2d_array.query_size_lod(0);
+    let size_3d: glam::UVec3 = sampled_image3d.query_size_lod(0);
+
+    *output = glam::UVec3::new(size_2d.x, size_2d_array.y, size_3d.z);
+}