[msl-out] Bounds checks for ImageLoad and ImageStore. (#1730)

src/back/msl/keywords.rs

@@ -213,4 +213,5 @@ pub const RESERVED: &[&str] = &[
     "M_SQRT1_2",
     // Naga utilities
     "DefaultConstructible",
+    "clamped_lod_e",
 ];

src/back/msl/writer.rs

@@ -65,6 +65,9 @@ fn put_numeric_type(
     }
 }
 
+/// Prefix for cached clamped level-of-detail values for `ImageLoad` expressions.
+const CLAMPED_LOD_LOAD_PREFIX: &str = "clamped_lod_e";
+
 struct TypeContext<'a> {
     handle: Handle<crate::Type>,
     arena: &'a crate::UniqueArena<crate::Type>,
@@ -455,6 +458,37 @@ enum FunctionOrigin {
     EntryPoint(proc::EntryPointIndex),
 }
 
+/// A level of detail argument.
+///
+/// When [`BoundsCheckPolicy::Restrict`] applies to an [`ImageLoad`] access, we
+/// save the clamped level of detail in a temporary variable whose name is based
+/// on the handle of the `ImageLoad` expression. But for other policies, we just
+/// use the expression directly.
+///
+/// [`BoundsCheckPolicy::Restrict`]: index::BoundsCheckPolicy::Restrict
+/// [`ImageLoad`]: crate::Expression::ImageLoad
+#[derive(Clone, Copy)]
+enum LevelOfDetail {
+    Direct(Handle<crate::Expression>),
+    Restricted(Handle<crate::Expression>),
+}
+
+/// Values needed to select a particular texel for [`ImageLoad`] and [`ImageStore`].
+///
+/// When this is used in code paths unconcerned with the `Restrict` bounds check
+/// policy, the `LevelOfDetail` enum introduces an unneeded match, since `level`
+/// will always be either `None` or `Some(Direct(_))`. But this turns out not to
+/// be too awkward. If that changes, we can revisit.
+///
+/// [`ImageLoad`]: crate::Expression::ImageLoad
+/// [`ImageStore`]: crate::Statement::ImageStore
+struct TexelAddress {
+    coordinate: Handle<crate::Expression>,
+    array_index: Option<Handle<crate::Expression>>,
+    sample: Option<Handle<crate::Expression>>,
+    level: Option<LevelOfDetail>,
+}
+
 struct ExpressionContext<'a> {
     function: &'a crate::Function,
     origin: FunctionOrigin,
@@ -475,6 +509,21 @@ impl<'a> ExpressionContext<'a> {
         self.info[handle].ty.inner_with(&self.module.types)
     }
 
+    /// Return true if calls to `image`'s `read` and `write` methods should supply a level of detail.
+    ///
+    /// Only mipmapped images need to specify a level of detail. Since 1D
+    /// textures cannot have mipmaps, MSL requires that the level argument to
+    /// texture1d queries and accesses must be a constexpr 0. It's easiest
+    /// just to omit the level entirely for 1D textures.
+    fn image_needs_lod(&self, image: Handle<crate::Expression>) -> bool {
+        let image_ty = self.resolve_type(image);
+        if let crate::TypeInner::Image { dim, class, .. } = *image_ty {
+            class.is_mipmapped() && dim != crate::ImageDimension::D1
+        } else {
+            false
+        }
+    }
+
     fn choose_bounds_check_policy(
         &self,
         pointer: Handle<crate::Expression>,
@@ -559,17 +608,31 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    fn put_level_of_detail(
+        &mut self,
+        level: LevelOfDetail,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        match level {
+            LevelOfDetail::Direct(expr) => self.put_expression(expr, context, true)?,
+            LevelOfDetail::Restricted(load) => {
+                write!(self.out, "{}{}", CLAMPED_LOD_LOAD_PREFIX, load.index())?
+            }
+        }
+        Ok(())
+    }
+
     fn put_image_query(
         &mut self,
         image: Handle<crate::Expression>,
         query: &str,
-        level: Option<Handle<crate::Expression>>,
+        level: Option<LevelOfDetail>,
         context: &ExpressionContext,
     ) -> BackendResult {
         self.put_expression(image, context, false)?;
         write!(self.out, ".get_{}(", query)?;
-        if let Some(expr) = level {
-            self.put_expression(expr, context, true)?;
+        if let Some(level) = level {
+            self.put_level_of_detail(level, context)?;
         }
         write!(self.out, ")")?;
         Ok(())
@@ -578,7 +641,8 @@ impl<W: Write> Writer<W> {
     fn put_image_size_query(
         &mut self,
         image: Handle<crate::Expression>,
-        level: Option<Handle<crate::Expression>>,
+        level: Option<LevelOfDetail>,
+        kind: crate::ScalarKind,
         context: &ExpressionContext,
     ) -> BackendResult {
         //Note: MSL only has separate width/height/depth queries,
@@ -587,24 +651,31 @@ impl<W: Write> Writer<W> {
             crate::TypeInner::Image { dim, .. } => dim,
             ref other => unreachable!("Unexpected type {:?}", other),
         };
+        let coordinate_type = kind.to_msl_name();
         match dim {
             crate::ImageDimension::D1 => {
-                write!(self.out, "int(")?;
                 // Since 1D textures never have mipmaps, MSL requires that the
                 // `level` argument be a constexpr 0. It's simplest for us just
-                // to omit the level entirely.
-                self.put_image_query(image, "width", None, context)?;
-                write!(self.out, ")")?;
+                // to pass `None` and omit the level entirely.
+                if kind == crate::ScalarKind::Uint {
+                    // No need to construct a vector. No cast needed.
+                    self.put_image_query(image, "width", None, context)?;
+                } else {
+                    // There's no definition for `int` in the `metal` namespace.
+                    write!(self.out, "int(")?;
+                    self.put_image_query(image, "width", None, context)?;
+                    write!(self.out, ")")?;
+                }
             }
             crate::ImageDimension::D2 => {
-                write!(self.out, "int2(")?;
+                write!(self.out, "{}::{}2(", NAMESPACE, coordinate_type)?;
                 self.put_image_query(image, "width", level, context)?;
                 write!(self.out, ", ")?;
                 self.put_image_query(image, "height", level, context)?;
                 write!(self.out, ")")?;
             }
             crate::ImageDimension::D3 => {
-                write!(self.out, "int3(")?;
+                write!(self.out, "{}::{}3(", NAMESPACE, coordinate_type)?;
                 self.put_image_query(image, "width", level, context)?;
                 write!(self.out, ", ")?;
                 self.put_image_query(image, "height", level, context)?;
@@ -613,7 +684,7 @@ impl<W: Write> Writer<W> {
                 write!(self.out, ")")?;
             }
             crate::ImageDimension::Cube => {
-                write!(self.out, "int2(")?;
+                write!(self.out, "{}::{}2(", NAMESPACE, coordinate_type)?;
                 self.put_image_query(image, "width", level, context)?;
                 write!(self.out, ")")?;
             }
@@ -621,7 +692,7 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
-    fn put_storage_image_coordinate(
+    fn put_cast_to_uint_scalar_or_vector(
         &mut self,
         expr: Handle<crate::Expression>,
         context: &ExpressionContext,
@@ -649,13 +720,7 @@ impl<W: Write> Writer<W> {
         level: crate::SampleLevel,
         context: &ExpressionContext,
     ) -> BackendResult {
-        let has_levels = match *context.resolve_type(image) {
-            crate::TypeInner::Image {
-                dim: crate::ImageDimension::D1,
-                ..
-            } => false,
-            _ => true,
-        };
+        let has_levels = context.image_needs_lod(image);
         match level {
             crate::SampleLevel::Auto => {}
             crate::SampleLevel::Zero => {
@@ -685,6 +750,275 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    fn put_image_coordinate_limits(
+        &mut self,
+        image: Handle<crate::Expression>,
+        level: Option<LevelOfDetail>,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
+        write!(self.out, " - 1")?;
+        Ok(())
+    }
+
+    /// General function for writing restricted image indexes.
+    ///
+    /// This is used to produce restricted mip levels, array indices, and sample
+    /// indices for [`ImageLoad`] and [`ImageStore`] accesses under the
+    /// [`Restrict`] bounds check policy.
+    ///
+    /// This function writes an expression of the form:
+    ///
+    /// ```ignore
+    ///
+    ///     metal::min(uint(INDEX), IMAGE.LIMIT_METHOD() - 1)
+    ///
+    /// ```
+    ///
+    /// [`ImageLoad`]: crate::Expression::ImageLoad
+    /// [`ImageStore`]: crate::Statement::ImageStore
+    /// [`Restrict`]: index::BoundsCheckPolicy::Restrict
+    fn put_restricted_scalar_image_index(
+        &mut self,
+        image: Handle<crate::Expression>,
+        index: Handle<crate::Expression>,
+        limit_method: &str,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        write!(self.out, "{}::min(uint(", NAMESPACE)?;
+        self.put_expression(index, context, true)?;
+        write!(self.out, "), ")?;
+        self.put_expression(image, context, false)?;
+        write!(self.out, ".{}() - 1)", limit_method)?;
+        Ok(())
+    }
+
+    fn put_restricted_texel_address(
+        &mut self,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        // Write the coordinate.
+        write!(self.out, "{}::min(", NAMESPACE)?;
+        self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
+        write!(self.out, ", ")?;
+        self.put_image_coordinate_limits(image, address.level, context)?;
+        write!(self.out, ")")?;
+
+        // Write the array index, if present.
+        if let Some(array_index) = address.array_index {
+            write!(self.out, ", ")?;
+            self.put_restricted_scalar_image_index(image, array_index, "get_array_size", context)?;
+        }
+
+        // Write the sample index, if present.
+        if let Some(sample) = address.sample {
+            write!(self.out, ", ")?;
+            self.put_restricted_scalar_image_index(image, sample, "get_num_samples", context)?;
+        }
+
+        // The level of detail should be clamped and cached by
+        // `put_cache_restricted_level`, so we don't need to clamp it here.
+        if let Some(level) = address.level {
+            write!(self.out, ", ")?;
+            self.put_level_of_detail(level, context)?;
+        }
+
+        Ok(())
+    }
+
+    /// Write an expression that is true if the given image access is in bounds.
+    fn put_image_access_bounds_check(
+        &mut self,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        let mut conjunction = "";
+
+        // First, check the level of detail. Only if that is in bounds can we
+        // use it to find the appropriate bounds for the coordinates.
+        let level = if let Some(level) = address.level {
+            write!(self.out, "uint(")?;
+            self.put_level_of_detail(level, context)?;
+            write!(self.out, ") < ")?;
+            self.put_expression(image, context, true)?;
+            write!(self.out, ".get_num_mip_levels()")?;
+            conjunction = " && ";
+            Some(level)
+        } else {
+            None
+        };
+
+        // Check sample index, if present.
+        if let Some(sample) = address.sample {
+            write!(self.out, "uint(")?;
+            self.put_expression(sample, context, true)?;
+            write!(self.out, ") < ")?;
+            self.put_expression(image, context, true)?;
+            write!(self.out, ".get_num_samples()")?;
+            conjunction = " && ";
+        }
+
+        // Check array index, if present.
+        if let Some(array_index) = address.array_index {
+            write!(self.out, "{}uint(", conjunction)?;
+            self.put_expression(array_index, context, true)?;
+            write!(self.out, ") < ")?;
+            self.put_expression(image, context, true)?;
+            write!(self.out, ".get_array_size()")?;
+            conjunction = " && ";
+        }
+
+        // Finally, check if the coordinates are within bounds.
+        let coord_is_vector = match *context.resolve_type(address.coordinate) {
+            crate::TypeInner::Vector { .. } => true,
+            _ => false,
+        };
+        write!(self.out, "{}", conjunction)?;
+        if coord_is_vector {
+            write!(self.out, "{}::all(", NAMESPACE)?;
+        }
+        self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
+        write!(self.out, " < ")?;
+        self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
+        if coord_is_vector {
+            write!(self.out, ")")?;
+        }
+
+        Ok(())
+    }
+
+    fn put_image_load(
+        &mut self,
+        load: Handle<crate::Expression>,
+        image: Handle<crate::Expression>,
+        mut address: TexelAddress,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        match context.policies.image {
+            proc::BoundsCheckPolicy::Restrict => {
+                // Use the cached restricted level of detail, if any. Omit the
+                // level altogether for 1D textures.
+                if address.level.is_some() {
+                    address.level = if context.image_needs_lod(image) {
+                        Some(LevelOfDetail::Restricted(load))
+                    } else {
+                        None
+                    }
+                }
+
+                self.put_expression(image, context, false)?;
+                write!(self.out, ".read(")?;
+                self.put_restricted_texel_address(image, &address, context)?;
+                write!(self.out, ")")?;
+            }
+            proc::BoundsCheckPolicy::ReadZeroSkipWrite => {
+                write!(self.out, "(")?;
+                self.put_image_access_bounds_check(image, &address, context)?;
+                write!(self.out, " ? ")?;
+                self.put_unchecked_image_load(image, &address, context)?;
+                write!(self.out, ": DefaultConstructible())")?;
+            }
+            proc::BoundsCheckPolicy::Unchecked => {
+                self.put_unchecked_image_load(image, &address, context)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn put_unchecked_image_load(
+        &mut self,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        context: &ExpressionContext,
+    ) -> BackendResult {
+        self.put_expression(image, context, false)?;
+        write!(self.out, ".read(")?;
+        // coordinates in IR are int, but Metal expects uint
+        self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
+        if let Some(expr) = address.array_index {
+            write!(self.out, ", ")?;
+            self.put_expression(expr, context, true)?;
+        }
+        if let Some(sample) = address.sample {
+            write!(self.out, ", ")?;
+            self.put_expression(sample, context, true)?;
+        }
+        if let Some(level) = address.level {
+            if context.image_needs_lod(image) {
+                write!(self.out, ", ")?;
+                self.put_level_of_detail(level, context)?;
+            }
+        }
+        write!(self.out, ")")?;
+
+        Ok(())
+    }
+
+    fn put_image_store(
+        &mut self,
+        level: back::Level,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        value: Handle<crate::Expression>,
+        context: &StatementContext,
+    ) -> BackendResult {
+        match context.expression.policies.image {
+            proc::BoundsCheckPolicy::Restrict => {
+                // We don't have a restricted level value, because we don't
+                // support writes to mipmapped textures.
+                debug_assert!(address.level.is_none());
+
+                write!(self.out, "{}", level)?;
+                self.put_expression(image, &context.expression, false)?;
+                write!(self.out, ".write(")?;
+                self.put_expression(value, &context.expression, true)?;
+                write!(self.out, ", ")?;
+                self.put_restricted_texel_address(image, address, &context.expression)?;
+                writeln!(self.out, ");")?;
+            }
+            proc::BoundsCheckPolicy::ReadZeroSkipWrite => {
+                write!(self.out, "{}if (", level)?;
+                self.put_image_access_bounds_check(image, address, &context.expression)?;
+                writeln!(self.out, ") {{")?;
+                self.put_unchecked_image_store(level.next(), image, address, value, context)?;
+                writeln!(self.out, "{}}}", level)?;
+            }
+            proc::BoundsCheckPolicy::Unchecked => {
+                self.put_unchecked_image_store(level, image, address, value, context)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn put_unchecked_image_store(
+        &mut self,
+        level: back::Level,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        value: Handle<crate::Expression>,
+        context: &StatementContext,
+    ) -> BackendResult {
+        write!(self.out, "{}", level)?;
+        self.put_expression(image, &context.expression, false)?;
+        write!(self.out, ".write(")?;
+        self.put_expression(value, &context.expression, true)?;
+        write!(self.out, ", ")?;
+        // coordinates in IR are int, but Metal expects uint
+        self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
+        if let Some(expr) = address.array_index {
+            write!(self.out, ", ")?;
+            self.put_expression(expr, &context.expression, true)?;
+        }
+        writeln!(self.out, ");")?;
+
+        Ok(())
+    }
+
     fn put_compose(
         &mut self,
         ty: Handle<crate::Type>,
@@ -1036,38 +1370,24 @@ impl<W: Write> Writer<W> {
                 sample,
                 level,
             } => {
-                self.put_expression(image, context, false)?;
-                write!(self.out, ".read(")?;
-                self.put_storage_image_coordinate(coordinate, context)?;
-                if let Some(expr) = array_index {
-                    write!(self.out, ", ")?;
-                    self.put_expression(expr, context, true)?;
-                }
-                if let Some(sample) = sample {
-                    write!(self.out, ", ")?;
-                    self.put_expression(sample, context, true)?
-                }
-                if let Some(level) = level {
-                    // Metal requires that the `level` argument to
-                    // `texture1d::read` be a constexpr equal to zero.
-                    if let crate::TypeInner::Image {
-                        dim: crate::ImageDimension::D1,
-                        ..
-                    } = *context.resolve_type(image)
-                    {
-                        // The argument defaults to zero.
-                    } else {
-                        write!(self.out, ", ")?;
-                        self.put_expression(level, context, true)?
-                    }
-                }
-                write!(self.out, ")")?;
+                let address = TexelAddress {
+                    coordinate,
+                    array_index,
+                    sample,
+                    level: level.map(LevelOfDetail::Direct),
+                };
+                self.put_image_load(expr_handle, image, address, context)?;
             }
             //Note: for all the queries, the signed integers are expected,
             // so a conversion is needed.
             crate::Expression::ImageQuery { image, query } => match query {
                 crate::ImageQuery::Size { level } => {
-                    self.put_image_size_query(image, level, context)?;
+                    self.put_image_size_query(
+                        image,
+                        level.map(LevelOfDetail::Direct),
+                        crate::ScalarKind::Sint,
+                        context,
+                    )?;
                 }
                 crate::ImageQuery::NumLevels => {
                     write!(self.out, "int(")?;
@@ -1928,6 +2248,57 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    /// Cache a clamped level of detail value, if necessary.
+    ///
+    /// [`ImageLoad`] accesses covered by [`BoundsCheckPolicy::Restrict`] use a
+    /// properly clamped level of detail value both in the access itself, and
+    /// for fetching the size of the requested MIP level, needed to clamp the
+    /// coordinates. To avoid recomputing this clamped level of detail, we cache
+    /// it in a temporary variable, as part of the [`Emit`] statement covering
+    /// the [`ImageLoad`] expression.
+    ///
+    /// [`ImageLoad`]: crate::Expression::ImageLoad
+    /// [`BoundsCheckPolicy::Restrict`]: index::BoundsCheckPolicy::Restrict
+    /// [`Emit`]: crate::Statement::Emit
+    fn put_cache_restricted_level(
+        &mut self,
+        load: Handle<crate::Expression>,
+        image: Handle<crate::Expression>,
+        mip_level: Option<Handle<crate::Expression>>,
+        indent: back::Level,
+        context: &StatementContext,
+    ) -> BackendResult {
+        // Does this image access actually require (or even permit) a
+        // level-of-detail, and does the policy require us to restrict it?
+        let level_of_detail = match mip_level {
+            Some(level) => level,
+            None => return Ok(()),
+        };
+
+        if context.expression.policies.image != index::BoundsCheckPolicy::Restrict
+            || !context.expression.image_needs_lod(image)
+        {
+            return Ok(());
+        }
+
+        write!(
+            self.out,
+            "{}uint {}{} = ",
+            indent,
+            CLAMPED_LOD_LOAD_PREFIX,
+            load.index(),
+        )?;
+        self.put_restricted_scalar_image_index(
+            image,
+            level_of_detail,
+            "get_num_mip_levels",
+            &context.expression,
+        )?;
+        writeln!(self.out, ";")?;
+
+        Ok(())
+    }
+
     fn put_block(
         &mut self,
         level: back::Level,
@@ -1945,6 +2316,19 @@ impl<W: Write> Writer<W> {
             match *statement {
                 crate::Statement::Emit(ref range) => {
                     for handle in range.clone() {
+                        // `ImageLoad` expressions covered by the `Restrict` bounds check policy
+                        // may need to cache a clamped version of their level-of-detail argument.
+                        if let crate::Expression::ImageLoad {
+                            image,
+                            level: mip_level,
+                            ..
+                        } = context.expression.function.expressions[handle]
+                        {
+                            self.put_cache_restricted_level(
+                                handle, image, mip_level, level, context,
+                            )?;
+                        }
+
                         let info = &context.expression.info[handle];
                         let ptr_class = info
                             .ty
@@ -2123,17 +2507,13 @@ impl<W: Write> Writer<W> {
                     array_index,
                     value,
                 } => {
-                    write!(self.out, "{}", level)?;
-                    self.put_expression(image, &context.expression, false)?;
-                    write!(self.out, ".write(")?;
-                    self.put_expression(value, &context.expression, true)?;
-                    write!(self.out, ", ")?;
-                    self.put_storage_image_coordinate(coordinate, &context.expression)?;
-                    if let Some(expr) = array_index {
-                        write!(self.out, ", ")?;
-                        self.put_expression(expr, &context.expression, true)?;
-                    }
-                    writeln!(self.out, ");")?;
+                    let address = TexelAddress {
+                        coordinate,
+                        array_index,
+                        sample: None,
+                        level: None,
+                    };
+                    self.put_image_store(level, image, &address, value, context)?
                 }
                 crate::Statement::Call {
                     function,

src/proc/index.rs

@@ -178,26 +178,39 @@ pub enum GuardedIndex {
 ///
 /// Such index expressions will be used twice in the generated code: first for the
 /// comparison to see if the index is in bounds, and then for the access itself, should
-/// the comparison succeed. To avoid computing the expressions twice, they should be
-/// cached in temporary variables.
+/// the comparison succeed. To avoid computing the expressions twice, the generated code
+/// should cache them in temporary variables.
 ///
-/// Why do we need to build such a set before processing a function's statements? Whether
-/// an expression needs to be cached depends on whether it appears as the [`index`]
-/// operand of any [`Access`] expression, and on the index bounds check policies that
-/// apply to those accesses. But [`Emit`] statements just identify a range of expressions
-/// by index; there's no good way to tell what an expression is used for. The only way to
-/// do it is to just iterate over all the expressions looking for relevant `Access`
-/// expressions --- which is what this function does.
+/// Why do we need to build such a set in advance, instead of just processing access
+/// expressions as we encounter them? Whether an expression needs to be cached depends on
+/// whether it appears as something like the [`index`] operand of an [`Access`] expression
+/// or the [`level`] operand of an [`ImageLoad`] expression, and on the index bounds check
+/// policies that apply to those accesses. But [`Emit`] statements just identify a range
+/// of expressions by index; there's no good way to tell what an expression is used
+/// for. The only way to do it is to just iterate over all the expressions looking for
+/// relevant `Access` expressions --- which is what this function does.
 ///
 /// Simple expressions like variable loads and constants don't make sense to cache: it's
 /// no better than just re-evaluating them. But constants are not covered by `Emit`
 /// statements, and `Load`s are always cached to ensure they occur at the right time, so
 /// we don't bother filtering them out from this set.
 ///
+/// Fortunately, we don't need to deal with [`ImageStore`] statements here. When we emit
+/// code for a statement, the writer isn't in the middle of an expression, so we can just
+/// emit declarations for temporaries, initialized appropriately.
+///
+/// None of these concerns apply for SPIR-V output, since it's easy to just reuse an
+/// instruction ID in two places; that has the same semantics as a temporary variable, and
+/// it's inherent in the design of SPIR-V. This function is more useful for text-based
+/// back ends.
+///
 /// [`ReadZeroSkipWrite`]: BoundsCheckPolicy::ReadZeroSkipWrite
 /// [`index`]: crate::Expression::Access::index
 /// [`Access`]: crate::Expression::Access
+/// [`level`]: crate::Expression::ImageLoad::level
+/// [`ImageLoad`]: crate::Expression::ImageLoad
 /// [`Emit`]: crate::Statement::Emit
+/// [`ImageStore`]: crate::Statement::ImageStore
 pub fn find_checked_indexes(
     module: &crate::Module,
     function: &crate::Function,
@@ -213,20 +226,43 @@ pub fn find_checked_indexes(
         for (_handle, expr) in function.expressions.iter() {
             // There's no need to handle `AccessIndex` expressions, as their
             // indices never need to be cached.
-            if let Ex::Access { base, index } = *expr {
-                if policies.choose_policy(base, &module.types, info)
-                    == BoundsCheckPolicy::ReadZeroSkipWrite
-                    && access_needs_check(
-                        base,
-                        GuardedIndex::Expression(index),
-                        module,
-                        function,
-                        info,
-                    )
-                    .is_some()
-                {
-                    guarded_indices.insert(index.index());
+            match *expr {
+                Ex::Access { base, index } => {
+                    if policies.choose_policy(base, &module.types, info)
+                        == BoundsCheckPolicy::ReadZeroSkipWrite
+                        && access_needs_check(
+                            base,
+                            GuardedIndex::Expression(index),
+                            module,
+                            function,
+                            info,
+                        )
+                        .is_some()
+                    {
+                        guarded_indices.insert(index.index());
+                    }
                 }
+                Ex::ImageLoad {
+                    coordinate,
+                    array_index,
+                    sample,
+                    level,
+                    ..
+                } => {
+                    if policies.image == BoundsCheckPolicy::ReadZeroSkipWrite {
+                        guarded_indices.insert(coordinate.index());
+                        if let Some(array_index) = array_index {
+                            guarded_indices.insert(array_index.index());
+                        }
+                        if let Some(sample) = sample {
+                            guarded_indices.insert(sample.index());
+                        }
+                        if let Some(level) = level {
+                            guarded_indices.insert(level.index());
+                        }
+                    }
+                }
+                _ => {}
             }
         }
     }

tests/out/msl/bounds-check-image-restrict.msl

@@ -0,0 +1,122 @@
+// language: metal2.0
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using metal::uint;
+
+
+metal::float4 test_textureLoad_1d(
+    int coords,
+    int level,
+    metal::texture1d<float, metal::access::sample> image_1d
+) {
+    metal::float4 _e3 = image_1d.read(metal::min(uint(coords), image_1d.get_width() - 1));
+    return _e3;
+}
+
+metal::float4 test_textureLoad_2d(
+    metal::int2 coords_1,
+    int level_1,
+    metal::texture2d<float, metal::access::sample> image_2d
+) {
+    uint clamped_lod_e4 = metal::min(uint(level_1), image_2d.get_num_mip_levels() - 1);
+    metal::float4 _e4 = image_2d.read(metal::min(metal::uint2(coords_1), metal::uint2(image_2d.get_width(clamped_lod_e4), image_2d.get_height(clamped_lod_e4)) - 1), clamped_lod_e4);
+    return _e4;
+}
+
+metal::float4 test_textureLoad_2d_array(
+    metal::int2 coords_2,
+    int index,
+    int level_2,
+    metal::texture2d_array<float, metal::access::sample> image_2d_array
+) {
+    uint clamped_lod_e6 = metal::min(uint(level_2), image_2d_array.get_num_mip_levels() - 1);
+    metal::float4 _e6 = image_2d_array.read(metal::min(metal::uint2(coords_2), metal::uint2(image_2d_array.get_width(clamped_lod_e6), image_2d_array.get_height(clamped_lod_e6)) - 1), metal::min(uint(index), image_2d_array.get_array_size() - 1), clamped_lod_e6);
+    return _e6;
+}
+
+metal::float4 test_textureLoad_3d(
+    metal::int3 coords_3,
+    int level_3,
+    metal::texture3d<float, metal::access::sample> image_3d
+) {
+    uint clamped_lod_e6 = metal::min(uint(level_3), image_3d.get_num_mip_levels() - 1);
+    metal::float4 _e6 = image_3d.read(metal::min(metal::uint3(coords_3), metal::uint3(image_3d.get_width(clamped_lod_e6), image_3d.get_height(clamped_lod_e6), image_3d.get_depth(clamped_lod_e6)) - 1), clamped_lod_e6);
+    return _e6;
+}
+
+metal::float4 test_textureLoad_multisampled_2d(
+    metal::int2 coords_4,
+    int sample,
+    metal::texture2d_ms<float, metal::access::read> image_multisampled_2d
+) {
+    metal::float4 _e7 = image_multisampled_2d.read(metal::min(metal::uint2(coords_4), metal::uint2(image_multisampled_2d.get_width(), image_multisampled_2d.get_height()) - 1), metal::min(uint(sample), image_multisampled_2d.get_num_samples() - 1));
+    return _e7;
+}
+
+float test_textureLoad_depth_2d(
+    metal::int2 coords_5,
+    int level_4,
+    metal::depth2d<float, metal::access::sample> image_depth_2d
+) {
+    uint clamped_lod_e8 = metal::min(uint(level_4), image_depth_2d.get_num_mip_levels() - 1);
+    float _e8 = image_depth_2d.read(metal::min(metal::uint2(coords_5), metal::uint2(image_depth_2d.get_width(clamped_lod_e8), image_depth_2d.get_height(clamped_lod_e8)) - 1), clamped_lod_e8);
+    return _e8;
+}
+
+float test_textureLoad_depth_2d_array(
+    metal::int2 coords_6,
+    int index_1,
+    int level_5,
+    metal::depth2d_array<float, metal::access::sample> image_depth_2d_array
+) {
+    uint clamped_lod_e10 = metal::min(uint(level_5), image_depth_2d_array.get_num_mip_levels() - 1);
+    float _e10 = image_depth_2d_array.read(metal::min(metal::uint2(coords_6), metal::uint2(image_depth_2d_array.get_width(clamped_lod_e10), image_depth_2d_array.get_height(clamped_lod_e10)) - 1), metal::min(uint(index_1), image_depth_2d_array.get_array_size() - 1), clamped_lod_e10);
+    return _e10;
+}
+
+float test_textureLoad_depth_multisampled_2d(
+    metal::int2 coords_7,
+    int sample_1,
+    metal::depth2d_ms<float, metal::access::read> image_depth_multisampled_2d
+) {
+    float _e10 = image_depth_multisampled_2d.read(metal::min(metal::uint2(coords_7), metal::uint2(image_depth_multisampled_2d.get_width(), image_depth_multisampled_2d.get_height()) - 1), metal::min(uint(sample_1), image_depth_multisampled_2d.get_num_samples() - 1));
+    return _e10;
+}
+
+void test_textureStore_1d(
+    int coords_8,
+    metal::float4 value,
+    metal::texture1d<float, metal::access::write> image_storage_1d
+) {
+    image_storage_1d.write(value, metal::min(uint(coords_8), image_storage_1d.get_width() - 1));
+    return;
+}
+
+void test_textureStore_2d(
+    metal::int2 coords_9,
+    metal::float4 value_1,
+    metal::texture2d<float, metal::access::write> image_storage_2d
+) {
+    image_storage_2d.write(value_1, metal::min(metal::uint2(coords_9), metal::uint2(image_storage_2d.get_width(), image_storage_2d.get_height()) - 1));
+    return;
+}
+
+void test_textureStore_2d_array(
+    metal::int2 coords_10,
+    int array_index,
+    metal::float4 value_2,
+    metal::texture2d_array<float, metal::access::write> image_storage_2d_array
+) {
+    image_storage_2d_array.write(value_2, metal::min(metal::uint2(coords_10), metal::uint2(image_storage_2d_array.get_width(), image_storage_2d_array.get_height()) - 1), metal::min(uint(array_index), image_storage_2d_array.get_array_size() - 1));
+    return;
+}
+
+void test_textureStore_3d(
+    metal::int3 coords_11,
+    metal::float4 value_3,
+    metal::texture3d<float, metal::access::write> image_storage_3d
+) {
+    image_storage_3d.write(value_3, metal::min(metal::uint3(coords_11), metal::uint3(image_storage_3d.get_width(), image_storage_3d.get_height(), image_storage_3d.get_depth()) - 1));
+    return;
+}

tests/out/msl/bounds-check-image-rzsw.msl

@@ -0,0 +1,131 @@
+// language: metal2.0
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using metal::uint;
+
+struct DefaultConstructible {
+    template<typename T>
+    operator T() && {
+        return T {};
+    }
+};
+
+metal::float4 test_textureLoad_1d(
+    int coords,
+    int level,
+    metal::texture1d<float, metal::access::sample> image_1d
+) {
+    metal::float4 _e3 = (uint(level) < image_1d.get_num_mip_levels() && uint(coords) < image_1d.get_width() ? image_1d.read(uint(coords)): DefaultConstructible());
+    return _e3;
+}
+
+metal::float4 test_textureLoad_2d(
+    metal::int2 coords_1,
+    int level_1,
+    metal::texture2d<float, metal::access::sample> image_2d
+) {
+    metal::float4 _e4 = (uint(level_1) < image_2d.get_num_mip_levels() && metal::all(metal::uint2(coords_1) < metal::uint2(image_2d.get_width(level_1), image_2d.get_height(level_1))) ? image_2d.read(metal::uint2(coords_1), level_1): DefaultConstructible());
+    return _e4;
+}
+
+metal::float4 test_textureLoad_2d_array(
+    metal::int2 coords_2,
+    int index,
+    int level_2,
+    metal::texture2d_array<float, metal::access::sample> image_2d_array
+) {
+    metal::float4 _e6 = (uint(level_2) < image_2d_array.get_num_mip_levels() && uint(index) < image_2d_array.get_array_size() && metal::all(metal::uint2(coords_2) < metal::uint2(image_2d_array.get_width(level_2), image_2d_array.get_height(level_2))) ? image_2d_array.read(metal::uint2(coords_2), index, level_2): DefaultConstructible());
+    return _e6;
+}
+
+metal::float4 test_textureLoad_3d(
+    metal::int3 coords_3,
+    int level_3,
+    metal::texture3d<float, metal::access::sample> image_3d
+) {
+    metal::float4 _e6 = (uint(level_3) < image_3d.get_num_mip_levels() && metal::all(metal::uint3(coords_3) < metal::uint3(image_3d.get_width(level_3), image_3d.get_height(level_3), image_3d.get_depth(level_3))) ? image_3d.read(metal::uint3(coords_3), level_3): DefaultConstructible());
+    return _e6;
+}
+
+metal::float4 test_textureLoad_multisampled_2d(
+    metal::int2 coords_4,
+    int sample,
+    metal::texture2d_ms<float, metal::access::read> image_multisampled_2d
+) {
+    metal::float4 _e7 = (uint(sample) < image_multisampled_2d.get_num_samples() && metal::all(metal::uint2(coords_4) < metal::uint2(image_multisampled_2d.get_width(), image_multisampled_2d.get_height())) ? image_multisampled_2d.read(metal::uint2(coords_4), sample): DefaultConstructible());
+    return _e7;
+}
+
+float test_textureLoad_depth_2d(
+    metal::int2 coords_5,
+    int level_4,
+    metal::depth2d<float, metal::access::sample> image_depth_2d
+) {
+    float _e8 = (uint(level_4) < image_depth_2d.get_num_mip_levels() && metal::all(metal::uint2(coords_5) < metal::uint2(image_depth_2d.get_width(level_4), image_depth_2d.get_height(level_4))) ? image_depth_2d.read(metal::uint2(coords_5), level_4): DefaultConstructible());
+    return _e8;
+}
+
+float test_textureLoad_depth_2d_array(
+    metal::int2 coords_6,
+    int index_1,
+    int level_5,
+    metal::depth2d_array<float, metal::access::sample> image_depth_2d_array
+) {
+    float _e10 = (uint(level_5) < image_depth_2d_array.get_num_mip_levels() && uint(index_1) < image_depth_2d_array.get_array_size() && metal::all(metal::uint2(coords_6) < metal::uint2(image_depth_2d_array.get_width(level_5), image_depth_2d_array.get_height(level_5))) ? image_depth_2d_array.read(metal::uint2(coords_6), index_1, level_5): DefaultConstructible());
+    return _e10;
+}
+
+float test_textureLoad_depth_multisampled_2d(
+    metal::int2 coords_7,
+    int sample_1,
+    metal::depth2d_ms<float, metal::access::read> image_depth_multisampled_2d
+) {
+    float _e10 = (uint(sample_1) < image_depth_multisampled_2d.get_num_samples() && metal::all(metal::uint2(coords_7) < metal::uint2(image_depth_multisampled_2d.get_width(), image_depth_multisampled_2d.get_height())) ? image_depth_multisampled_2d.read(metal::uint2(coords_7), sample_1): DefaultConstructible());
+    return _e10;
+}
+
+void test_textureStore_1d(
+    int coords_8,
+    metal::float4 value,
+    metal::texture1d<float, metal::access::write> image_storage_1d
+) {
+    if (uint(coords_8) < image_storage_1d.get_width()) {
+        image_storage_1d.write(value, uint(coords_8));
+    }
+    return;
+}
+
+void test_textureStore_2d(
+    metal::int2 coords_9,
+    metal::float4 value_1,
+    metal::texture2d<float, metal::access::write> image_storage_2d
+) {
+    if (metal::all(metal::uint2(coords_9) < metal::uint2(image_storage_2d.get_width(), image_storage_2d.get_height()))) {
+        image_storage_2d.write(value_1, metal::uint2(coords_9));
+    }
+    return;
+}
+
+void test_textureStore_2d_array(
+    metal::int2 coords_10,
+    int array_index,
+    metal::float4 value_2,
+    metal::texture2d_array<float, metal::access::write> image_storage_2d_array
+) {
+    if (uint(array_index) < image_storage_2d_array.get_array_size() && metal::all(metal::uint2(coords_10) < metal::uint2(image_storage_2d_array.get_width(), image_storage_2d_array.get_height()))) {
+        image_storage_2d_array.write(value_2, metal::uint2(coords_10), array_index);
+    }
+    return;
+}
+
+void test_textureStore_3d(
+    metal::int3 coords_11,
+    metal::float4 value_3,
+    metal::texture3d<float, metal::access::write> image_storage_3d
+) {
+    if (metal::all(metal::uint3(coords_11) < metal::uint3(image_storage_3d.get_width(), image_storage_3d.get_height(), image_storage_3d.get_depth()))) {
+        image_storage_3d.write(value_3, metal::uint3(coords_11));
+    }
+    return;
+}

tests/out/msl/image.msl

@@ -17,7 +17,7 @@ kernel void main_(
 , metal::texture1d<uint, metal::access::sample> image_1d_src [[user(fake0)]]
 , metal::texture1d<uint, metal::access::write> image_dst [[user(fake0)]]
 ) {
-    metal::int2 dim = int2(image_storage_src.get_width(), image_storage_src.get_height());
+    metal::int2 dim = metal::int2(image_storage_src.get_width(), image_storage_src.get_height());
     metal::int2 itc = (dim * static_cast<metal::int2>(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 value2_ = image_multisampled_src.read(metal::uint2(itc), static_cast<int>(local_id.z));
@@ -37,7 +37,7 @@ kernel void depth_load(
 , metal::texture2d<uint, metal::access::read> image_storage_src [[user(fake0)]]
 , metal::texture1d<uint, metal::access::write> image_dst [[user(fake0)]]
 ) {
-    metal::int2 dim_1 = int2(image_storage_src.get_width(), image_storage_src.get_height());
+    metal::int2 dim_1 = metal::int2(image_storage_src.get_width(), image_storage_src.get_height());
     metal::int2 itc_1 = (dim_1 * static_cast<metal::int2>(local_id_1.xy)) % metal::int2(10, 20);
     float val = image_depth_multisampled_src.read(metal::uint2(itc_1), static_cast<int>(local_id_1.z));
     image_dst.write(metal::uint4(static_cast<uint>(val)), uint(itc_1.x));
@@ -58,16 +58,16 @@ vertex queriesOutput queries(
 ) {
     int dim_1d = int(image_1d.get_width());
     int dim_1d_lod = int(image_1d.get_width());
-    metal::int2 dim_2d = int2(image_2d.get_width(), image_2d.get_height());
-    metal::int2 dim_2d_lod = int2(image_2d.get_width(1), image_2d.get_height(1));
-    metal::int2 dim_2d_array = int2(image_2d_array.get_width(), image_2d_array.get_height());
-    metal::int2 dim_2d_array_lod = int2(image_2d_array.get_width(1), image_2d_array.get_height(1));
-    metal::int2 dim_cube = int2(image_cube.get_width());
-    metal::int2 dim_cube_lod = int2(image_cube.get_width(1));
-    metal::int2 dim_cube_array = int2(image_cube_array.get_width());
-    metal::int2 dim_cube_array_lod = int2(image_cube_array.get_width(1));
-    metal::int3 dim_3d = int3(image_3d.get_width(), image_3d.get_height(), image_3d.get_depth());
-    metal::int3 dim_3d_lod = int3(image_3d.get_width(1), image_3d.get_height(1), image_3d.get_depth(1));
+    metal::int2 dim_2d = metal::int2(image_2d.get_width(), image_2d.get_height());
+    metal::int2 dim_2d_lod = metal::int2(image_2d.get_width(1), image_2d.get_height(1));
+    metal::int2 dim_2d_array = metal::int2(image_2d_array.get_width(), image_2d_array.get_height());
+    metal::int2 dim_2d_array_lod = metal::int2(image_2d_array.get_width(1), image_2d_array.get_height(1));
+    metal::int2 dim_cube = metal::int2(image_cube.get_width());
+    metal::int2 dim_cube_lod = metal::int2(image_cube.get_width(1));
+    metal::int2 dim_cube_array = metal::int2(image_cube_array.get_width());
+    metal::int2 dim_cube_array_lod = metal::int2(image_cube_array.get_width(1));
+    metal::int3 dim_3d = metal::int3(image_3d.get_width(), image_3d.get_height(), image_3d.get_depth());
+    metal::int3 dim_3d_lod = metal::int3(image_3d.get_width(1), image_3d.get_height(1), image_3d.get_depth(1));
     int sum = (((((((((dim_1d + dim_2d.y) + dim_2d_lod.y) + dim_2d_array.y) + dim_2d_array_lod.y) + dim_cube.y) + dim_cube_lod.y) + dim_cube_array.y) + dim_cube_array_lod.y) + dim_3d.z) + dim_3d_lod.z;
     return queriesOutput { metal::float4(static_cast<float>(sum)) };
 }

tests/out/msl/policy-mix.msl

@@ -46,7 +46,7 @@ metal::float4 mock_function(
     for(int _i=0; _i<2; ++_i) in_function.inner[_i] = type_9 {metal::float4(0.7070000171661377, 0.0, 0.0, 1.0), metal::float4(0.0, 0.7070000171661377, 0.0, 1.0)}.inner[_i];
     metal::float4 _e22 = in_storage.a.inner[i];
     metal::float4 _e25 = in_uniform.a.inner[i];
-    metal::float4 _e27 = image_2d_array.read(metal::uint2(c), i, l);
+    metal::float4 _e27 = (uint(l) < image_2d_array.get_num_mip_levels() && uint(i) < image_2d_array.get_array_size() && metal::all(metal::uint2(c) < metal::uint2(image_2d_array.get_width(l), image_2d_array.get_height(l))) ? image_2d_array.read(metal::uint2(c), i, l): DefaultConstructible());
     float _e30 = in_workgroup.inner[metal::min(unsigned(i), 29u)];
     float _e34 = in_private.inner[metal::min(unsigned(i), 39u)];
     metal::float4 _e38 = in_function.inner[metal::min(unsigned(i), 1u)];

tests/snapshots.rs

@@ -479,8 +479,11 @@ fn convert_wgsl() {
         ("bounds-check-zero", Targets::SPIRV | Targets::METAL),
         ("bounds-check-zero-atomic", Targets::METAL),
         ("bounds-check-restrict", Targets::SPIRV | Targets::METAL),
-        ("bounds-check-image-restrict", Targets::SPIRV),
-        ("bounds-check-image-rzsw", Targets::SPIRV),
+        (
+            "bounds-check-image-restrict",
+            Targets::SPIRV | Targets::METAL,
+        ),
+        ("bounds-check-image-rzsw", Targets::SPIRV | Targets::METAL),
         ("policy-mix", Targets::SPIRV | Targets::METAL),
         (
             "texture-arg",