Add query_size_lod and query_size methods to SampledImage

crates/spirv-std/src/image.rs

@@ -10,7 +10,7 @@ mod params;
 /// Contains extra image operands
 pub mod sample_with;
 
-pub use self::params::{ImageCoordinate, ImageCoordinateSubpassData, SampleType};
+pub use self::params::{ImageCoordinate, ImageCoordinateSubpassData, ImageSizeQuery, SampleType};
 pub use crate::macros::Image;
 pub use spirv_std_types::image_params::{
     AccessQualifier, Arrayed, Dimensionality, ImageDepth, ImageFormat, Multisampled, Sampled,
@@ -933,7 +933,7 @@ impl<
     /// Query the dimensions of Image, with no level of detail.
     #[crate::macros::gpu_only]
     #[doc(alias = "OpImageQuerySize")]
-    pub fn query_size<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(&self) -> Size
+    pub fn query_size<Size: ImageSizeQuery<u32, DIM, ARRAYED> + Default>(&self) -> Size
     where
         Self: HasQuerySize,
     {
@@ -971,10 +971,10 @@ impl<
         COMPONENTS,
     >
 {
-    /// Query the dimensions of Image, with no level of detail.
+    /// Query the dimensions of Image at a specific level of detail.
     #[crate::macros::gpu_only]
     #[doc(alias = "OpImageQuerySizeLod")]
-    pub fn query_size_lod<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(
+    pub fn query_size_lod<Size: ImageSizeQuery<u32, DIM, ARRAYED> + Default>(
         &self,
         lod: u32,
     ) -> Size
@@ -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: ImageSizeQuery<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: ImageSizeQuery<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,

crates/spirv-std/src/image/params.rs

@@ -194,3 +194,68 @@ impl<V: Vector<S, 4>, S: Scalar>
 pub trait ImageCoordinateSubpassData<T, const ARRAYED: u32> {}
 impl<V: Vector<I, 2>, I: Integer> ImageCoordinateSubpassData<I, { Arrayed::False as u32 }> for V {}
 impl<V: Vector<I, 3>, I: Integer> ImageCoordinateSubpassData<I, { Arrayed::True as u32 }> for V {}
+
+/// Marker trait for query size results based on image dimension and arraying.
+///
+/// This trait represents the SPIR-V size query results:
+/// - 1D images return a scalar
+/// - 2D/Cube/Rect images return 2 components (Cube returns face width/height)
+/// - 3D images return 3 components
+/// - Arrayed images add one component for the array size
+pub trait ImageSizeQuery<T, const DIM: u32, const ARRAYED: u32> {}
+
+// 1D images
+impl<T: Scalar> ImageSizeQuery<T, { Dimensionality::OneD as u32 }, { Arrayed::False as u32 }>
+    for T
+{
+}
+impl<V: Vector<T, 2>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::OneD as u32 }, { Arrayed::True as u32 }> for V
+{
+}
+
+// 2D images
+impl<V: Vector<T, 2>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::TwoD as u32 }, { Arrayed::False as u32 }> for V
+{
+}
+impl<V: Vector<T, 3>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::TwoD as u32 }, { Arrayed::True as u32 }> for V
+{
+}
+
+// 3D images
+impl<V: Vector<T, 3>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::ThreeD as u32 }, { Arrayed::False as u32 }> for V
+{
+}
+impl<V: Vector<T, 4>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::ThreeD as u32 }, { Arrayed::True as u32 }> for V
+{
+}
+
+// Cube images - returns 2D size (width/height of face)
+impl<V: Vector<T, 2>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::Cube as u32 }, { Arrayed::False as u32 }> for V
+{
+}
+impl<V: Vector<T, 3>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::Cube as u32 }, { Arrayed::True as u32 }> for V
+{
+}
+
+// Rect images
+impl<V: Vector<T, 2>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::Rect as u32 }, { Arrayed::False as u32 }> for V
+{
+}
+impl<V: Vector<T, 3>, T: Scalar>
+    ImageSizeQuery<T, { Dimensionality::Rect as u32 }, { Arrayed::True as u32 }> for V
+{
+}
+
+// Buffer images
+impl<T: Scalar> ImageSizeQuery<T, { Dimensionality::Buffer as u32 }, { Arrayed::False as u32 }>
+    for T
+{
+}

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

@@ -0,0 +1,24 @@
+// build-pass
+// compile-flags: -C target-feature=+ImageQuery
+
+use spirv_std::spirv;
+use spirv_std::{Image, arch, image::Cubemap};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] cubemap: &Cubemap,
+    #[spirv(descriptor_set = 1, binding = 1)] cubemap_array: &Image!(cube, type=f32, sampled, arrayed),
+    #[spirv(descriptor_set = 2, binding = 2)] storage_cubemap: &Image!(cube, type=f32, sampled=false),
+    output: &mut glam::UVec3,
+) {
+    // Cubemaps return 2D size (width, height of one face)
+    let size: glam::UVec2 = cubemap.query_size_lod(0);
+
+    // Arrayed cubemaps return 3D size (width, height, array_layers)
+    let size_array: glam::UVec3 = cubemap_array.query_size_lod(0);
+
+    // Storage cubemaps can use query_size directly
+    let storage_size: glam::UVec2 = storage_cubemap.query_size();
+
+    *output = glam::UVec3::new(size.x, size_array.z, storage_size.x);
+}

tests/compiletests/ui/image/query/query_size_err.stderr

@@ -17,7 +17,7 @@ error[E0277]: the trait bound `Image<f32, 1, 2, 0, 0, 1, 0, 4>: HasQuerySize` is
 note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, MULTISAMPLED, SAMPLED, FORMAT, COMPONENTS>::query_size`
    --> $SPIRV_STD_SRC/image.rs:938:15
     |
-936 |     pub fn query_size<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(&self) -> Size
+936 |     pub fn query_size<Size: ImageSizeQuery<u32, DIM, ARRAYED> + Default>(&self) -> Size
     |            ---------- required by a bound in this associated function
 937 |     where
 938 |         Self: HasQuerySize,

tests/compiletests/ui/image/query/query_size_lod_err.stderr

@@ -12,7 +12,7 @@ error[E0277]: the trait bound `Image<f32, 4, 2, 0, 0, 1, 0, 4>: HasQuerySizeLod`
 note: required by a bound in `Image::<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#7}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>::query_size_lod`
    --> $SPIRV_STD_SRC/image.rs:982:15
     |
-977 |     pub fn query_size_lod<Size: ImageCoordinate<u32, DIM, ARRAYED> + Default>(
+977 |     pub fn query_size_lod<Size: ImageSizeQuery<u32, DIM, ARRAYED> + Default>(
     |            -------------- required by a bound in this associated function
 ...
 982 |         Self: HasQuerySizeLod,

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

@@ -0,0 +1,24 @@
+// build-pass
+// compile-flags: -C target-feature=+ImageQuery,+SampledRect
+// ignore-vulkan1.0
+// ignore-vulkan1.1
+// ignore-vulkan1.1spv1.4
+// ignore-vulkan1.2
+
+use spirv_std::spirv;
+use spirv_std::{Image, arch};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] rect_storage: &Image!(rect, type=f32, sampled=false),
+    #[spirv(descriptor_set = 1, binding = 1)] rect_storage_array: &Image!(rect, type=f32, sampled=false, arrayed),
+    output: &mut glam::UVec3,
+) {
+    // Rect images only support query_size (not query_size_lod)
+    let rect_size: glam::UVec2 = rect_storage.query_size();
+
+    // Arrayed rect images return 3D size (width, height, array_layers)
+    let rect_array_size: glam::UVec3 = rect_storage_array.query_size();
+
+    *output = glam::UVec3::new(rect_size.x, rect_size.y, rect_array_size.z);
+}

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

@@ -0,0 +1,16 @@
+// 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,53 @@
+// build-pass
+// compile-flags: -C target-feature=+ImageQuery,+Sampled1D
+
+use spirv_std::spirv;
+use spirv_std::{Image, arch, image::SampledImage};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] sampled_image1d: &SampledImage<
+        Image!(1D, type=f32, sampled),
+    >,
+    #[spirv(descriptor_set = 1, binding = 1)] sampled_image1d_array: &SampledImage<
+        Image!(1D, type=f32, arrayed, sampled),
+    >,
+    #[spirv(descriptor_set = 2, binding = 2)] sampled_image2d: &SampledImage<
+        Image!(2D, type=f32, sampled),
+    >,
+    #[spirv(descriptor_set = 3, binding = 3)] sampled_image2d_array: &SampledImage<
+        Image!(2D, type=f32, arrayed, sampled),
+    >,
+    #[spirv(descriptor_set = 4, binding = 4)] sampled_image3d: &SampledImage<
+        Image!(3D, type=f32, sampled),
+    >,
+    #[spirv(descriptor_set = 5, binding = 5)] sampled_image3d_array: &SampledImage<
+        Image!(3D, type=f32, arrayed, sampled),
+    >,
+    output: &mut glam::UVec4,
+) {
+    // 1D images return scalar
+    let size_1d: u32 = sampled_image1d.query_size_lod(0);
+
+    // 1D arrayed images return 2 components (width, array_layers)
+    let size_1d_array: glam::UVec2 = sampled_image1d_array.query_size_lod(0);
+
+    // 2D images return 2 components
+    let size_2d: glam::UVec2 = sampled_image2d.query_size_lod(0);
+
+    // 2D arrayed images return 3 components
+    let size_2d_array: glam::UVec3 = sampled_image2d_array.query_size_lod(0);
+
+    // 3D images return 3 components
+    let size_3d: glam::UVec3 = sampled_image3d.query_size_lod(0);
+
+    // 3D arrayed images return 4 components (width, height, depth, array_layers)
+    let size_3d_array: glam::UVec4 = sampled_image3d_array.query_size_lod(0);
+
+    *output = glam::UVec4::new(
+        size_1d,
+        size_1d_array.y,
+        size_2d.x + size_3d.z,
+        size_3d_array.w,
+    );
+}

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

@@ -0,0 +1,20 @@
+// build-fail
+// normalize-stderr-test "\S*/crates/spirv-std/src/" -> "$SPIRV_STD_SRC/"
+// compile-flags: -C target-feature=+ImageQuery,+SampledRect
+// ignore-vulkan1.0
+// ignore-vulkan1.1
+// ignore-vulkan1.1spv1.4
+// ignore-vulkan1.2
+
+use spirv_std::{Image, arch, image::SampledImage, spirv};
+
+#[spirv(fragment)]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0)] rect_sampled: &SampledImage<
+        Image!(rect, type=f32, sampled),
+    >,
+    output: &mut glam::UVec2,
+) {
+    // This should fail because rect images don't support query_size_lod
+    *output = rect_sampled.query_size_lod(0);
+}

tests/compiletests/ui/image/query/sampled_image_rect_query_size_lod_err.stderr

@@ -0,0 +1,23 @@
+error[E0277]: the trait bound `Image<f32, 4, 2, 0, 0, 1, 0, 4>: HasQuerySizeLod` is not satisfied
+    --> $DIR/sampled_image_rect_query_size_lod_err.rs:19:28
+     |
+19   |     *output = rect_sampled.query_size_lod(0);
+     |                            ^^^^^^^^^^^^^^ the trait `HasQuerySizeLod` is not implemented for `Image<f32, 4, 2, 0, 0, 1, 0, 4>`
+     |
+     = help: the following other types implement trait `HasQuerySizeLod`:
+               Image<SampledType, 0, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
+               Image<SampledType, 1, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
+               Image<SampledType, 2, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
+               Image<SampledType, 3, DEPTH, ARRAYED, 0, SAMPLED, FORMAT, COMPONENTS>
+note: required by a bound in `SampledImage::<Image<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#9}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>>::query_size_lod`
+    --> /image.rs:1138:12
+     |
+1124 |     pub fn query_size_lod<Size: ImageSizeQuery<u32, DIM, ARRAYED> + Default>(
+     |            -------------- required by a bound in this associated function
+...
+1138 |         >: HasQuerySizeLod,
+     |            ^^^^^^^^^^^^^^^ required by this bound in `SampledImage::<Image<SampledType, DIM, DEPTH, ARRAYED, spirv_std::::image::{impl#9}::{constant#0}, SAMPLED, FORMAT, COMPONENTS>>::query_size_lod`
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0277`.