Add 64 bit atomics

CHANGELOG.md

@@ -81,8 +81,41 @@ for message in compilation_info
 
 By @stefnotch in [#5410](https://github.com/gfx-rs/wgpu/pull/5410)
 
-### New features
+#### 64 bit integer atomic support in shaders.
+
+Add support for 64 bit integer atomic operations in shaders.
+
+Add the following flags to `wgpu_types::Features`:
+
+- `SHADER_INT64_ATOMIC_ALL_OPS` enables all atomic operations on `atomic<i64>` and
+  `atomic<u64>` values.
+
+- `SHADER_INT64_ATOMIC_MIN_MAX` is a subset of the above, enabling only
+  `AtomicFunction::Min` and `AtomicFunction::Max` operations on `atomic<i64>` and
+  `atomic<u64>` values in the `Storage` address space. These are the only 64-bit
+  atomic operations available on Metal as of 3.1.
+
+Add corresponding flags to `naga::valid::Capabilities`. These are supported by the
+WGSL front end, and all Naga backends.
 
+Platform support:
+
+- On Direct3d 12, in `D3D12_FEATURE_DATA_D3D12_OPTIONS9`, if
+  `AtomicInt64OnTypedResourceSupported` and `AtomicInt64OnGroupSharedSupported` are
+  both available, then both wgpu features described above are available.
+
+- On Metal, `SHADER_INT64_ATOMIC_MIN_MAX` is available on Apple9 hardware, and on
+  hardware that advertises both Apple8 and Mac2 support. This also requires Metal
+  Shading Language 2.4 or later. Metal does not yet support the more general
+  `SHADER_INT64_ATOMIC_ALL_OPS`.
+
+- On Vulkan, if the `VK_KHR_shader_atomic_int64` extension is available with both the
+  `shader_buffer_int64_atomics` and `shader_shared_int64_atomics` features, then both
+  wgpu features described above are available.
+
+By @atlv24 in [#5383](https://github.com/gfx-rs/wgpu/pull/5383)
+
+### New features
 #### Vulkan
 
 - Added a `PipelineCache` resource to allow using Vulkan pipeline caches. By @DJMcNab in [#5319](https://github.com/gfx-rs/wgpu/pull/5319)

naga/src/back/dot/mod.rs

@@ -244,7 +244,9 @@ impl StatementGraph {
                     value,
                     result,
                 } => {
-                    self.emits.push((id, result));
+                    if let Some(result) = result {
+                        self.emits.push((id, result));
+                    }
                     self.dependencies.push((id, pointer, "pointer"));
                     self.dependencies.push((id, value, "value"));
                     if let crate::AtomicFunction::Exchange { compare: Some(cmp) } = *fun {

naga/src/back/glsl/mod.rs

@@ -2368,11 +2368,13 @@ impl<'a, W: Write> Writer<'a, W> {
                 result,
             } => {
                 write!(self.out, "{level}")?;
-                let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
-                let res_ty = ctx.resolve_type(result, &self.module.types);
-                self.write_value_type(res_ty)?;
-                write!(self.out, " {res_name} = ")?;
-                self.named_expressions.insert(result, res_name);
+                if let Some(result) = result {
+                    let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                    let res_ty = ctx.resolve_type(result, &self.module.types);
+                    self.write_value_type(res_ty)?;
+                    write!(self.out, " {res_name} = ")?;
+                    self.named_expressions.insert(result, res_name);
+                }
 
                 let fun_str = fun.to_glsl();
                 write!(self.out, "atomic{fun_str}(")?;

naga/src/back/hlsl/writer.rs

@@ -1919,11 +1919,20 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 result,
             } => {
                 write!(self.out, "{level}")?;
-                let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
-                match func_ctx.info[result].ty {
-                    proc::TypeResolution::Handle(handle) => self.write_type(module, handle)?,
-                    proc::TypeResolution::Value(ref value) => {
-                        self.write_value_type(module, value)?
+                let res_name = match result {
+                    None => None,
+                    Some(result) => {
+                        let name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                        match func_ctx.info[result].ty {
+                            proc::TypeResolution::Handle(handle) => {
+                                self.write_type(module, handle)?
+                            }
+                            proc::TypeResolution::Value(ref value) => {
+                                self.write_value_type(module, value)?
+                            }
+                        };
+                        write!(self.out, " {name}; ")?;
+                        Some((result, name))
                     }
                 };
 
@@ -1934,7 +1943,6 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     .unwrap();
 
                 let fun_str = fun.to_hlsl_suffix();
-                write!(self.out, " {res_name}; ")?;
                 match pointer_space {
                     crate::AddressSpace::WorkGroup => {
                         write!(self.out, "Interlocked{fun_str}(")?;
@@ -1970,8 +1978,16 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                     _ => {}
                 }
                 self.write_expr(module, value, func_ctx)?;
-                writeln!(self.out, ", {res_name});")?;
-                self.named_expressions.insert(result, res_name);
+
+                // The `original_value` out parameter is optional for all the
+                // `Interlocked` functions we generate other than
+                // `InterlockedExchange`.
+                if let Some((result, name)) = res_name {
+                    write!(self.out, ", {name}")?;
+                    self.named_expressions.insert(result, name);
+                }
+
+                writeln!(self.out, ");")?;
             }
             Statement::WorkGroupUniformLoad { pointer, result } => {
                 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;

naga/src/back/msl/mod.rs

@@ -1,6 +1,9 @@
 /*!
 Backend for [MSL][msl] (Metal Shading Language).
 
+This backend does not support the [`SHADER_INT64_ATOMIC_ALL_OPS`][all-atom]
+capability.
+
 ## Binding model
 
 Metal's bindings are flat per resource. Since there isn't an obvious mapping
@@ -24,6 +27,8 @@ For the result type, if it's a structure, we re-compose it with a temporary valu
 holding the result.
 
 [msl]: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
+[all-atom]: crate::valid::Capabilities::SHADER_INT64_ATOMIC_ALL_OPS
+
 */
 
 use crate::{arena::Handle, proc::index, valid::ModuleInfo};
@@ -661,21 +666,3 @@ fn test_error_size() {
     use std::mem::size_of;
     assert_eq!(size_of::<Error>(), 32);
 }
-
-impl crate::AtomicFunction {
-    fn to_msl(self) -> Result<&'static str, Error> {
-        Ok(match self {
-            Self::Add => "fetch_add",
-            Self::Subtract => "fetch_sub",
-            Self::And => "fetch_and",
-            Self::InclusiveOr => "fetch_or",
-            Self::ExclusiveOr => "fetch_xor",
-            Self::Min => "fetch_min",
-            Self::Max => "fetch_max",
-            Self::Exchange { compare: None } => "exchange",
-            Self::Exchange { compare: Some(_) } => Err(Error::FeatureNotImplemented(
-                "atomic CompareExchange".to_string(),
-            ))?,
-        })
-    }
-}

naga/src/back/msl/writer.rs

@@ -3058,11 +3058,22 @@ impl<W: Write> Writer<W> {
                     value,
                     result,
                 } => {
+                    // This backend supports `SHADER_INT64_ATOMIC_MIN_MAX` but not
+                    // `SHADER_INT64_ATOMIC_ALL_OPS`, so we can assume that if `result` is
+                    // `Some`, we are not operating on a 64-bit value, and that if we are
+                    // operating on a 64-bit value, `result` is `None`.
                     write!(self.out, "{level}")?;
-                    let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
-                    self.start_baking_expression(result, &context.expression, &res_name)?;
-                    self.named_expressions.insert(result, res_name);
-                    let fun_str = fun.to_msl()?;
+                    let fun_str = if let Some(result) = result {
+                        let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                        self.start_baking_expression(result, &context.expression, &res_name)?;
+                        self.named_expressions.insert(result, res_name);
+                        fun.to_msl()?
+                    } else if context.expression.resolve_type(value).scalar_width() == Some(8) {
+                        fun.to_msl_64_bit()?
+                    } else {
+                        fun.to_msl()?
+                    };
+
                     self.put_atomic_operation(pointer, fun_str, value, &context.expression)?;
                     // done
                     writeln!(self.out, ";")?;
@@ -5914,3 +5925,31 @@ fn test_stack_size() {
         }
     }
 }
+
+impl crate::AtomicFunction {
+    fn to_msl(self) -> Result<&'static str, Error> {
+        Ok(match self {
+            Self::Add => "fetch_add",
+            Self::Subtract => "fetch_sub",
+            Self::And => "fetch_and",
+            Self::InclusiveOr => "fetch_or",
+            Self::ExclusiveOr => "fetch_xor",
+            Self::Min => "fetch_min",
+            Self::Max => "fetch_max",
+            Self::Exchange { compare: None } => "exchange",
+            Self::Exchange { compare: Some(_) } => Err(Error::FeatureNotImplemented(
+                "atomic CompareExchange".to_string(),
+            ))?,
+        })
+    }
+
+    fn to_msl_64_bit(self) -> Result<&'static str, Error> {
+        Ok(match self {
+            Self::Min => "min",
+            Self::Max => "max",
+            _ => Err(Error::FeatureNotImplemented(
+                "64-bit atomic operation other than min/max".to_string(),
+            ))?,
+        })
+    }
+}

naga/src/back/pipeline_constants.rs

@@ -617,7 +617,9 @@ fn adjust_stmt(new_pos: &[Handle<Expression>], stmt: &mut Statement) {
         } => {
             adjust(pointer);
             adjust(value);
-            adjust(result);
+            if let Some(ref mut result) = *result {
+                adjust(result);
+            }
             match *fun {
                 crate::AtomicFunction::Exchange {
                     compare: Some(ref mut compare),

naga/src/back/spv/block.rs

@@ -2423,9 +2423,15 @@ impl<'w> BlockContext<'w> {
                     result,
                 } => {
                     let id = self.gen_id();
-                    let result_type_id = self.get_expression_type_id(&self.fun_info[result].ty);
-
-                    self.cached[result] = id;
+                    // Compare-and-exchange operations produce a struct result,
+                    // so use `result`'s type if it is available. For no-result
+                    // operations, fall back to `value`'s type.
+                    let result_type_id =
+                        self.get_expression_type_id(&self.fun_info[result.unwrap_or(value)].ty);
+
+                    if let Some(result) = result {
+                        self.cached[result] = id;
+                    }
 
                     let pointer_id =
                         match self.write_expression_pointer(pointer, &mut block, None)? {

naga/src/back/spv/writer.rs

@@ -878,6 +878,9 @@ impl Writer {
             crate::TypeInner::RayQuery => {
                 self.require_any("Ray Query", &[spirv::Capability::RayQueryKHR])?;
             }
+            crate::TypeInner::Atomic(crate::Scalar { width: 8, kind: _ }) => {
+                self.require_any("64 bit integer atomics", &[spirv::Capability::Int64Atomics])?;
+            }
             _ => {}
         }
         Ok(())

naga/src/back/wgsl/writer.rs

@@ -754,9 +754,11 @@ impl<W: Write> Writer<W> {
                 result,
             } => {
                 write!(self.out, "{level}")?;
-                let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
-                self.start_named_expr(module, result, func_ctx, &res_name)?;
-                self.named_expressions.insert(result, res_name);
+                if let Some(result) = result {
+                    let res_name = format!("{}{}", back::BAKE_PREFIX, result.index());
+                    self.start_named_expr(module, result, func_ctx, &res_name)?;
+                    self.named_expressions.insert(result, res_name);
+                }
 
                 let fun_str = fun.to_wgsl();
                 write!(self.out, "atomic{fun_str}(")?;

naga/src/compact/statements.rs

@@ -75,7 +75,9 @@ impl FunctionTracer<'_> {
                         self.expressions_used.insert(pointer);
                         self.trace_atomic_function(fun);
                         self.expressions_used.insert(value);
-                        self.expressions_used.insert(result);
+                        if let Some(result) = result {
+                            self.expressions_used.insert(result);
+                        }
                     }
                     St::WorkGroupUniformLoad { pointer, result } => {
                         self.expressions_used.insert(pointer);
@@ -255,7 +257,9 @@ impl FunctionMap {
                         adjust(pointer);
                         self.adjust_atomic_function(fun);
                         adjust(value);
-                        adjust(result);
+                        if let Some(ref mut result) = *result {
+                            adjust(result);
+                        }
                     }
                     St::WorkGroupUniformLoad {
                         ref mut pointer,

naga/src/front/spv/mod.rs

@@ -63,6 +63,7 @@ pub const SUPPORTED_CAPABILITIES: &[spirv::Capability] = &[
     spirv::Capability::Int8,
     spirv::Capability::Int16,
     spirv::Capability::Int64,
+    spirv::Capability::Int64Atomics,
     spirv::Capability::Float16,
     spirv::Capability::Float64,
     spirv::Capability::Geometry,
@@ -4028,7 +4029,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                         pointer: p_lexp_handle,
                         fun: crate::AtomicFunction::Add,
                         value: one_lexp_handle,
-                        result: r_lexp_handle,
+                        result: Some(r_lexp_handle),
                     };
                     block.push(stmt, span);
                 }

naga/src/front/type_gen.rs

@@ -291,10 +291,10 @@ impl crate::Module {
                                 name: Some("exchanged".to_string()),
                                 ty: bool_ty,
                                 binding: None,
-                                offset: 4,
+                                offset: scalar.width as u32,
                             },
                         ],
-                        span: 8,
+                        span: scalar.width as u32 * 2,
                     },
                 }
             }