[msl-out] Fix ReadZeroSkipWrite bounds check mode for pointer arguments

Fixes #4541

--
Co-authored-by: Liam Murphy <[email protected]>
Co-Authored-By: Erich Gubler <[email protected]>

Author: andyleiserson

naga/src/back/msl/mod.rs

@@ -29,6 +29,20 @@ holding the result.
 [msl]: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
 [all-atom]: crate::valid::Capabilities::SHADER_INT64_ATOMIC_ALL_OPS
 
+## Pointer-typed bounds-checked expressions and OOB locals
+
+MSL (unlike HLSL and GLSL) has native support for pointer-typed function
+arguments. When the [`BoundsCheckPolicy`] is `ReadZeroSkipWrite` and an
+out-of-bounds index expression is used for such an argument, our strategy is to
+pass a pointer to a dummy variable. These dummy variables are called "OOB
+locals". We emit at most one OOB local per function for each type, since all
+expressions producing a result of that type can share the same OOB local. (Note
+that the OOB local mechanism is not actually implementing "skip write", nor even
+"read zero" in some cases of read-after-write, but doing so would require
+additional effort and the difference is unlikely to matter.)
+
+[`BoundsCheckPolicy`]: crate::proc::BoundsCheckPolicy
+
 */
 
 use alloc::{

naga/src/back/msl/writer.rs

@@ -612,6 +612,17 @@ trait NameKeyExt {
             FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointLocal(idx, local_handle),
         }
     }
+
+    /// Return the name key for a local variable used by ReadZeroSkipWrite bounds-check
+    /// policy when it needs to produce a pointer-typed result for an OOB access. These
+    /// are unique per accessed type, so the second argument is a type handle. See docs
+    /// for [`crate::back::msl`].
+    fn oob_local_for_type(origin: FunctionOrigin, ty: Handle<crate::Type>) -> NameKey {
+        match origin {
+            FunctionOrigin::Handle(handle) => NameKey::FunctionOobLocal(handle, ty),
+            FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointOobLocal(idx, ty),
+        }
+    }
 }
 
 impl NameKeyExt for NameKey {}
@@ -722,6 +733,11 @@ impl<'a> ExpressionContext<'a> {
         index::bounds_check_iter(chain, self.module, self.function, self.info)
     }
 
+    /// See docs for [`proc::index::oob_local_types`].
+    fn oob_local_types(&self) -> FastHashSet<Handle<crate::Type>> {
+        index::oob_local_types(self.module, self.function, self.info, self.policies)
+    }
+
     fn get_packed_vec_kind(&self, expr_handle: Handle<crate::Expression>) -> Option<crate::Scalar> {
         match self.function.expressions[expr_handle] {
             crate::Expression::AccessIndex { base, index } => {
@@ -929,8 +945,18 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
-    /// Writes the local variables of the given function.
+    /// Writes the local variables of the given function, as well as any extra
+    /// out-of-bounds locals that are needed.
+    ///
+    /// The names of the OOB locals are also added to `self.names` at the same
+    /// time.
     fn put_locals(&mut self, context: &ExpressionContext) -> BackendResult {
+        let oob_local_types = context.oob_local_types();
+        for &ty in oob_local_types.iter() {
+            let name_key = NameKey::oob_local_for_type(context.origin, ty);
+            self.names.insert(name_key, self.namer.call("oob"));
+        }
+
         for (name_key, ty, init) in context
             .function
             .local_variables
@@ -939,6 +965,10 @@ impl<W: Write> Writer<W> {
                 let name_key = NameKey::local(context.origin, local_handle);
                 (name_key, local.ty, local.init)
             })
+            .chain(oob_local_types.iter().map(|&ty| {
+                let name_key = NameKey::oob_local_for_type(context.origin, ty);
+                (name_key, ty, None)
+            }))
         {
             let ty_name = TypeContext {
                 handle: ty,
@@ -1761,7 +1791,42 @@ impl<W: Write> Writer<W> {
                 {
                     write!(self.out, " ? ")?;
                     self.put_access_chain(expr_handle, policy, context)?;
-                    write!(self.out, " : DefaultConstructible()")?;
+                    write!(self.out, " : ")?;
+
+                    if context.resolve_type(base).pointer_space().is_some() {
+                        // We can't just use `DefaultConstructible` if this is a pointer.
+                        // Instead, we create a dummy local variable to serve as pointer
+                        // target if the access is out of bounds.
+                        let result_ty = context.info[expr_handle]
+                            .ty
+                            .inner_with(&context.module.types)
+                            .pointer_base_type();
+                        let result_ty_handle = match result_ty {
+                            Some(TypeResolution::Handle(handle)) => handle,
+                            Some(TypeResolution::Value(_)) => {
+                                // As long as the result of a pointer access expression is
+                                // passed to a function or stored in a let binding, the
+                                // type will be in the arena. If additional uses of
+                                // pointers become valid, this assumption might no longer
+                                // hold. Note that the LHS of a load or store doesn't
+                                // take this path -- there is dedicated code in `put_load`
+                                // and `put_store`.
+                                unreachable!(
+                                    "Expected type {result_ty:?} of access through pointer type {base:?} to be in the arena",
+                                );
+                            }
+                            None => {
+                                unreachable!(
+                                    "Expected access through pointer type {base:?} to return a pointer, but got {result_ty:?}",
+                                )
+                            }
+                        };
+                        let name_key =
+                            NameKey::oob_local_for_type(context.origin, result_ty_handle);
+                        self.out.write_str(&self.names[&name_key])?;
+                    } else {
+                        write!(self.out, "DefaultConstructible()")?;
+                    }
 
                     if !is_scoped {
                         write!(self.out, ")")?;

naga/src/proc/index.rs

@@ -2,10 +2,10 @@
 Definitions for index bounds checking.
 */
 
-use core::iter;
+use core::iter::{self, zip};
 
 use crate::arena::{Handle, HandleSet, UniqueArena};
-use crate::valid;
+use crate::{valid, FastHashSet};
 
 /// How should code generated by Naga do bounds checks?
 ///
@@ -389,6 +389,61 @@ pub(crate) fn bounds_check_iter<'a>(
     })
 }
 
+/// Returns all the types which we need out-of-bounds locals for; that is,
+/// all of the types which the code might attempt to get an out-of-bounds
+/// pointer to, in which case we yield a pointer to the out-of-bounds local
+/// of the correct type.
+pub fn oob_local_types(
+    module: &crate::Module,
+    function: &crate::Function,
+    info: &valid::FunctionInfo,
+    policies: BoundsCheckPolicies,
+) -> FastHashSet<Handle<crate::Type>> {
+    let mut result = FastHashSet::default();
+
+    if policies.index != BoundsCheckPolicy::ReadZeroSkipWrite {
+        return result;
+    }
+
+    for statement in &function.body {
+        // The only situation in which we end up actually needing to create an
+        // out-of-bounds pointer is when passing one to a function.
+        //
+        // This is because pointers are never baked; they're just inlined everywhere
+        // they're used. That means that loads can just return 0, and stores can just do
+        // nothing; functions are the only case where you actually *have* to produce a
+        // pointer.
+        if let crate::Statement::Call {
+            function: callee,
+            ref arguments,
+            ..
+        } = *statement
+        {
+            // Now go through the arguments of the function looking for pointers which need bounds checks.
+            for (arg_info, &arg) in zip(&module.functions[callee].arguments, arguments) {
+                match module.types[arg_info.ty].inner {
+                    crate::TypeInner::ValuePointer { .. } => {
+                        // `ValuePointer`s should only ever be used when resolving the types of
+                        // expressions, since the arena can no longer be modified at that point; things
+                        // in the arena should always use proper `Pointer`s.
+                        unreachable!("`ValuePointer` found in arena")
+                    }
+                    crate::TypeInner::Pointer { base, .. } => {
+                        if bounds_check_iter(arg, module, function, info)
+                            .next()
+                            .is_some()
+                        {
+                            result.insert(base);
+                        }
+                    }
+                    _ => continue,
+                };
+            }
+        }
+    }
+    result
+}
+
 impl GuardedIndex {
     /// Make a `GuardedIndex::Known` from a `GuardedIndex::Expression` if possible.
     ///

naga/src/proc/namer.rs

@@ -21,9 +21,19 @@ pub enum NameKey {
     Function(Handle<crate::Function>),
     FunctionArgument(Handle<crate::Function>, u32),
     FunctionLocal(Handle<crate::Function>, Handle<crate::LocalVariable>),
+
+    /// A local variable used by ReadZeroSkipWrite bounds-check policy
+    /// when it needs to produce a pointer-typed result for an OOB access.
+    /// These are unique per accessed type, so the second element is a
+    /// type handle. See docs for [`crate::back::msl`].
+    FunctionOobLocal(Handle<crate::Function>, Handle<crate::Type>),
+
     EntryPoint(EntryPointIndex),
     EntryPointLocal(EntryPointIndex, Handle<crate::LocalVariable>),
     EntryPointArgument(EntryPointIndex, u32),
+
+    /// Entry point version of `FunctionOobLocal`.
+    EntryPointOobLocal(EntryPointIndex, Handle<crate::Type>),
 }
 
 /// This processor assigns names to all the things in a module

naga/tests/in/wgsl/pointer-function-arg-restrict.toml

@@ -0,0 +1,4 @@
+targets = "METAL"
+
+[bounds_check_policies]
+index = "Restrict"

naga/tests/in/wgsl/pointer-function-arg-restrict.wgsl

@@ -0,0 +1,61 @@
+fn takes_ptr(p: ptr<function, i32>) {}
+fn takes_array_ptr(p: ptr<function, array<i32, 4>>) {}
+fn takes_vec_ptr(p: ptr<function, vec2<i32>>) {}
+fn takes_mat_ptr(p: ptr<function, mat2x2<f32>>) {}
+
+fn local_var(i: u32) {
+    var arr = array(1, 2, 3, 4);
+    takes_ptr(&arr[i]);
+    takes_array_ptr(&arr);
+
+}
+
+fn mat_vec_ptrs(
+    pv: ptr<function, array<vec2<i32>, 4>>,
+    pm: ptr<function, array<mat2x2<f32>, 4>>,
+    i: u32,
+) {
+    takes_vec_ptr(&pv[i]);
+    takes_mat_ptr(&pm[i]);
+}
+
+fn argument(v: ptr<function, array<i32, 4>>, i: u32) {
+    takes_ptr(&v[i]);
+}
+
+fn argument_nested_x2(v: ptr<function, array<array<i32, 4>, 4>>, i: u32, j: u32) {
+    takes_ptr(&v[i][j]);
+
+    // Mixing compile and runtime bounds checks
+    takes_ptr(&v[i][0]);
+    takes_ptr(&v[0][j]);
+
+    takes_array_ptr(&v[i]);
+}
+
+fn argument_nested_x3(v: ptr<function, array<array<array<i32, 4>, 4>, 4>>, i: u32, j: u32) {
+    takes_ptr(&v[i][0][j]);
+    takes_ptr(&v[i][j][0]);
+    takes_ptr(&v[0][i][j]);
+}
+
+fn index_from_self(v: ptr<function, array<i32, 4>>, i: u32) {
+    takes_ptr(&v[v[i]]);
+}
+
+fn local_var_from_arg(a: array<i32, 4>, i: u32) {
+    var b = a;
+    takes_ptr(&b[i]);
+}
+
+fn let_binding(a: ptr<function, array<i32, 4>>, i: u32) {
+    let p0 = &a[i];
+    takes_ptr(p0);
+
+    let p1 = &a[0];
+    takes_ptr(p1);
+}
+
+// Runtime-sized arrays can only appear in storage buffers, while (in the base
+// language) pointers can only appear in function or private space, so there
+// is no interaction to test.

naga/tests/in/wgsl/pointer-function-arg-rzsw.toml

@@ -0,0 +1,4 @@
+targets = "METAL"
+
+[bounds_check_policies]
+index = "ReadZeroSkipWrite"

naga/tests/in/wgsl/pointer-function-arg-rzsw.wgsl

@@ -0,0 +1,61 @@
+fn takes_ptr(p: ptr<function, i32>) {}
+fn takes_array_ptr(p: ptr<function, array<i32, 4>>) {}
+fn takes_vec_ptr(p: ptr<function, vec2<i32>>) {}
+fn takes_mat_ptr(p: ptr<function, mat2x2<f32>>) {}
+
+fn local_var(i: u32) {
+    var arr = array(1, 2, 3, 4);
+    takes_ptr(&arr[i]);
+    takes_array_ptr(&arr);
+
+}
+
+fn mat_vec_ptrs(
+    pv: ptr<function, array<vec2<i32>, 4>>,
+    pm: ptr<function, array<mat2x2<f32>, 4>>,
+    i: u32,
+) {
+    takes_vec_ptr(&pv[i]);
+    takes_mat_ptr(&pm[i]);
+}
+
+fn argument(v: ptr<function, array<i32, 4>>, i: u32) {
+    takes_ptr(&v[i]);
+}
+
+fn argument_nested_x2(v: ptr<function, array<array<i32, 4>, 4>>, i: u32, j: u32) {
+    takes_ptr(&v[i][j]);
+
+    // Mixing compile and runtime bounds checks
+    takes_ptr(&v[i][0]);
+    takes_ptr(&v[0][j]);
+
+    takes_array_ptr(&v[i]);
+}
+
+fn argument_nested_x3(v: ptr<function, array<array<array<i32, 4>, 4>, 4>>, i: u32, j: u32) {
+    takes_ptr(&v[i][0][j]);
+    takes_ptr(&v[i][j][0]);
+    takes_ptr(&v[0][i][j]);
+}
+
+fn index_from_self(v: ptr<function, array<i32, 4>>, i: u32) {
+    takes_ptr(&v[v[i]]);
+}
+
+fn local_var_from_arg(a: array<i32, 4>, i: u32) {
+    var b = a;
+    takes_ptr(&b[i]);
+}
+
+fn let_binding(a: ptr<function, array<i32, 4>>, i: u32) {
+    let p0 = &a[i];
+    takes_ptr(p0);
+
+    let p1 = &a[0];
+    takes_ptr(p1);
+}
+
+// Runtime-sized arrays can only appear in storage buffers, while (in the base
+// language) pointers can only appear in function or private space, so there
+// is no interaction to test.

naga/tests/in/wgsl/pointer-function-arg.toml

@@ -0,0 +1 @@
+targets = "METAL | GLSL | HLSL | WGSL"