[naga wgsl-in] Add support for unsigned types when calling textureLoad with the level parameter. (#7058)

Author: ygdrasil-io

CHANGELOG.md

@@ -136,6 +136,7 @@ By @cwfitzgerald in [#7030](https://github.com/gfx-rs/wgpu/pull/7030).
 
 #### Naga
 
+- Add support for unsigned types when calling textureLoad with the level parameter. By @ygdrasil-io in [#7058](https://github.com/gfx-rs/wgpu/pull/7058).
 - Support @must_use attribute on function declarations. By @turbocrime in [#6801](https://github.com/gfx-rs/wgpu/pull/6801).
 - Support for generating the candidate intersections from AABB geometry, and confirming the hits. By @kvark in [#7047](https://github.com/gfx-rs/wgpu/pull/7047).
 - Make naga::back::spv::Function::to_words write the OpFunctionEnd instruction in itself, instead of making another call after it. By @junjunjd in [#7156](https://github.com/gfx-rs/wgpu/pull/7156).

naga/src/back/glsl/mod.rs

@@ -4449,10 +4449,28 @@ impl<'a, W: Write> Writer<'a, W> {
                 writeln!(self.out, ") - 1)")?;
             }
         } else if let Some(sample_or_level) = sample.or(level) {
+            // GLSL only support SInt on this field while WGSL support also UInt
+            let cast_to_int = matches!(
+                *ctx.resolve_type(sample_or_level, &self.module.types),
+                TypeInner::Scalar(crate::Scalar {
+                    kind: crate::ScalarKind::Uint,
+                    ..
+                })
+            );
+
             // If no bounds checking is need just add the sample or level argument
             // after the coordinates
             write!(self.out, ", ")?;
+
+            if cast_to_int {
+                write!(self.out, "int(")?;
+            }
+
             self.write_expr(sample_or_level, ctx)?;
+
+            if cast_to_int {
+                write!(self.out, ")")?;
+            }
         }
 
         // Close the image load function.

naga/src/back/hlsl/help.rs

@@ -1611,8 +1611,26 @@ impl<W: Write> super::Writer<'_, W> {
                 self.write_expr(module, expr, func_ctx)?;
             }
             if let Some(expr) = mip_level {
+                // Explicit cast if needed
+                let cast_to_int = matches!(
+                    *func_ctx.resolve_type(expr, &module.types),
+                    crate::TypeInner::Scalar(crate::Scalar {
+                        kind: ScalarKind::Uint,
+                        ..
+                    })
+                );
+
                 write!(self.out, ", ")?;
+
+                if cast_to_int {
+                    write!(self.out, "int(")?;
+                }
+
                 self.write_expr(module, expr, func_ctx)?;
+
+                if cast_to_int {
+                    write!(self.out, ")")?;
+                }
             }
             write!(self.out, ")")?;
         }

naga/src/back/hlsl/writer.rs

@@ -3213,57 +3213,16 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 array_index,
                 sample,
                 level,
-            } => {
-                let mut wrapping_type = None;
-                match *func_ctx.resolve_type(image, &module.types) {
-                    TypeInner::Image {
-                        class: crate::ImageClass::Storage { format, .. },
-                        ..
-                    } => {
-                        if format.single_component() {
-                            wrapping_type = Some(Scalar::from(format));
-                        }
-                    }
-                    _ => {}
-                }
-                if let Some(scalar) = wrapping_type {
-                    write!(
-                        self.out,
-                        "{}{}(",
-                        help::IMAGE_STORAGE_LOAD_SCALAR_WRAPPER,
-                        scalar.to_hlsl_str()?
-                    )?;
-                }
-                // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-to-load
-                self.write_expr(module, image, func_ctx)?;
-                write!(self.out, ".Load(")?;
-
-                self.write_texture_coordinates(
-                    "int",
-                    coordinate,
-                    array_index,
-                    level,
-                    module,
-                    func_ctx,
-                )?;
-
-                if let Some(sample) = sample {
-                    write!(self.out, ", ")?;
-                    self.write_expr(module, sample, func_ctx)?;
-                }
-
-                // close bracket for Load function
-                write!(self.out, ")")?;
-
-                if wrapping_type.is_some() {
-                    write!(self.out, ")")?;
-                }
-
-                // return x component if return type is scalar
-                if let TypeInner::Scalar(_) = *func_ctx.resolve_type(expr, &module.types) {
-                    write!(self.out, ".x")?;
-                }
-            }
+            } => self.write_image_load(
+                &module,
+                expr,
+                func_ctx,
+                image,
+                coordinate,
+                array_index,
+                sample,
+                level,
+            )?,
             Expression::GlobalVariable(handle) => {
                 let global_variable = &module.global_variables[handle];
                 let ty = &module.types[global_variable.ty].inner;
@@ -4003,6 +3962,63 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
         Ok(())
     }
 
+    #[allow(clippy::too_many_arguments)]
+    fn write_image_load(
+        &mut self,
+        module: &&Module,
+        expr: Handle<crate::Expression>,
+        func_ctx: &back::FunctionCtx,
+        image: Handle<crate::Expression>,
+        coordinate: Handle<crate::Expression>,
+        array_index: Option<Handle<crate::Expression>>,
+        sample: Option<Handle<crate::Expression>>,
+        level: Option<Handle<crate::Expression>>,
+    ) -> Result<(), Error> {
+        let mut wrapping_type = None;
+        match *func_ctx.resolve_type(image, &module.types) {
+            TypeInner::Image {
+                class: crate::ImageClass::Storage { format, .. },
+                ..
+            } => {
+                if format.single_component() {
+                    wrapping_type = Some(Scalar::from(format));
+                }
+            }
+            _ => {}
+        }
+        if let Some(scalar) = wrapping_type {
+            write!(
+                self.out,
+                "{}{}(",
+                help::IMAGE_STORAGE_LOAD_SCALAR_WRAPPER,
+                scalar.to_hlsl_str()?
+            )?;
+        }
+        // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-to-load
+        self.write_expr(module, image, func_ctx)?;
+        write!(self.out, ".Load(")?;
+
+        self.write_texture_coordinates("int", coordinate, array_index, level, module, func_ctx)?;
+
+        if let Some(sample) = sample {
+            write!(self.out, ", ")?;
+            self.write_expr(module, sample, func_ctx)?;
+        }
+
+        // close bracket for Load function
+        write!(self.out, ")")?;
+
+        if wrapping_type.is_some() {
+            write!(self.out, ")")?;
+        }
+
+        // return x component if return type is scalar
+        if let TypeInner::Scalar(_) = *func_ctx.resolve_type(expr, &module.types) {
+            write!(self.out, ".x")?;
+        }
+        Ok(())
+    }
+
     /// Find the [`BindingArraySamplerInfo`] from an expression so that such an access
     /// can be generated later.
     fn sampler_binding_array_info_from_expression(

naga/src/valid/expression.rs

@@ -671,11 +671,15 @@ impl super::Validator {
 
                         match (level, class.is_mipmapped()) {
                             (None, false) => {}
-                            (Some(level), true) => {
-                                if resolver[level].scalar_kind() != Some(Sk::Sint) {
-                                    return Err(ExpressionError::InvalidImageOtherIndexType(level));
+                            (Some(level), true) => match resolver[level] {
+                                Ti::Scalar(Sc {
+                                    kind: Sk::Sint | Sk::Uint,
+                                    width: _,
+                                }) => {}
+                                _ => {
+                                    return Err(ExpressionError::InvalidImageArrayIndexType(level))
                                 }
-                            }
+                            },
                             _ => {
                                 return Err(ExpressionError::InvalidImageOtherIndex);
                             }

naga/tests/in/image.wgsl

@@ -21,6 +21,8 @@ fn main(@builtin(local_invocation_id) local_id: vec3<u32>) {
     let itc = vec2<i32>(dim * local_id.xy) % vec2<i32>(10, 20);
     // loads with ivec2 coords.
     let value1 = textureLoad(image_mipmapped_src, itc, i32(local_id.z));
+    // doing the same thing as the line above, but with u32, as textureLoad must also support unsigned integers.
+    let value1_2 = textureLoad(image_mipmapped_src, itc, u32(local_id.z));
     let value2 = textureLoad(image_multisampled_src, itc, i32(local_id.z));
     let value4 = textureLoad(image_storage_src, itc);
     let value5 = textureLoad(image_array_src, itc, local_id.z, i32(local_id.z) + 1);

naga/tests/out/glsl/image.main.Compute.glsl

@@ -20,6 +20,7 @@ void main() {
     uvec2 dim = uvec2(imageSize(_group_0_binding_1_cs).xy);
     ivec2 itc = (ivec2((dim * local_id.xy)) % ivec2(10, 20));
     uvec4 value1_ = texelFetch(_group_0_binding_0_cs, itc, int(local_id.z));
+    uvec4 value1_2_ = texelFetch(_group_0_binding_0_cs, itc, int(uint(local_id.z)));
     uvec4 value2_ = texelFetch(_group_0_binding_3_cs, itc, int(local_id.z));
     uvec4 value4_ = imageLoad(_group_0_binding_1_cs, itc);
     uvec4 value5_ = texelFetch(_group_0_binding_5_cs, ivec3(itc, local_id.z), (int(local_id.z) + 1));

naga/tests/out/hlsl/image.hlsl

@@ -42,6 +42,7 @@ void main(uint3 local_id : SV_GroupThreadID)
     uint2 dim = NagaRWDimensions2D(image_storage_src);
     int2 itc = naga_mod(int2((dim * local_id.xy)), int2(int(10), int(20)));
     uint4 value1_ = image_mipmapped_src.Load(int3(itc, int(local_id.z)));
+    uint4 value1_2_ = image_mipmapped_src.Load(int3(itc, int(uint(local_id.z))));
     uint4 value2_ = image_multisampled_src.Load(itc, int(local_id.z));
     uint4 value4_ = image_storage_src.Load(itc);
     uint4 value5_ = image_array_src.Load(int4(itc, local_id.z, asint(asuint(int(local_id.z)) + asuint(int(1)))));

naga/tests/out/msl/image.msl

@@ -24,6 +24,7 @@ kernel void main_(
     metal::uint2 dim = metal::uint2(image_storage_src.get_width(), image_storage_src.get_height());
     metal::int2 itc = naga_mod(static_cast<metal::int2>(dim * local_id.xy), metal::int2(10, 20));
     metal::uint4 value1_ = image_mipmapped_src.read(metal::uint2(itc), static_cast<int>(local_id.z));
+    metal::uint4 value1_2_ = image_mipmapped_src.read(metal::uint2(itc), static_cast<uint>(local_id.z));
     metal::uint4 value2_ = image_multisampled_src.read(metal::uint2(itc), static_cast<int>(local_id.z));
     metal::uint4 value4_ = image_storage_src.read(metal::uint2(itc));
     metal::uint4 value5_ = image_array_src.read(metal::uint2(itc), local_id.z, as_type<int>(as_type<uint>(static_cast<int>(local_id.z)) + as_type<uint>(1)));