Validate address space for metal

gfx-rs · Apr 3, 2024 · 820dce0 · 820dce0
1 parent 8f15a3a
commit 820dce0
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/naga/src/back/spv/writer.rs b/naga/src/back/spv/writer.rs
@@ -878,10 +878,8 @@ impl Writer {
             crate::TypeInner::RayQuery => {
                 self.require_any("Ray Query", &[spirv::Capability::RayQueryKHR])?;
             }
-            crate::TypeInner::Atomic(crate::Scalar { width, kind: _ }) => {
-                if width == 8 {
-                    self.require_any("64 bit integer atomics", &[spirv::Capability::Int64Atomics])?;
-                }
+            crate::TypeInner::Atomic(crate::Scalar { width: 8, kind: _ }) => {
+                self.require_any("64 bit integer atomics", &[spirv::Capability::Int64Atomics])?;
             }
             _ => {}
         }

diff --git a/naga/src/valid/function.rs b/naga/src/valid/function.rs
@@ -41,6 +41,8 @@ pub enum CallError {
 pub enum AtomicError {
     #[error("Pointer {0:?} to atomic is invalid.")]
     InvalidPointer(Handle<crate::Expression>),
+    #[error("Address space {0:?} does not support 64bit atomics.")]
+    InvalidAddressSpace(crate::AddressSpace),
     #[error("Operand {0:?} has invalid type.")]
     InvalidOperand(Handle<crate::Expression>),
     #[error("Result type for {0:?} doesn't match the statement")]
@@ -424,8 +426,32 @@ impl super::Validator {
     ) -> Result<(), WithSpan<FunctionError>> {
         let pointer_inner = context.resolve_type(pointer, &self.valid_expression_set)?;
         let ptr_scalar = match *pointer_inner {
-            crate::TypeInner::Pointer { base, .. } => match context.types[base].inner {
-                crate::TypeInner::Atomic(scalar) => scalar,
+            crate::TypeInner::Pointer { base, space } => match context.types[base].inner {
+                crate::TypeInner::Atomic(scalar) => match space {
+                    crate::AddressSpace::Function
+                    | crate::AddressSpace::Private
+                    | crate::AddressSpace::WorkGroup => {
+                        if self
+                            .capabilities
+                            .contains(crate::valid::Capabilities::SHADER_INT64_ATOMIC_ALL_OPS)
+                        {
+                            scalar
+                        } else {
+                            // Metal does not support atomics on non-device space pointers. Other platforms can just use normal atomics with returns.
+                            log::error!(
+                                "64 bit atomic no-return function on an invalid address-space {:?}",
+                                space
+                            );
+                            return Err(AtomicError::InvalidAddressSpace(space)
+                                .with_span_handle(pointer, context.expressions)
+                                .into_other());
+                        }
+                    }
+                    crate::AddressSpace::Uniform
+                    | crate::AddressSpace::Storage { .. }
+                    | crate::AddressSpace::Handle
+                    | crate::AddressSpace::PushConstant => scalar,
+                },
                 ref other => {
                     log::error!("Atomic pointer to type {:?}", other);
                     return Err(AtomicError::InvalidPointer(pointer)