64 bit image atomics

CHANGELOG.md

@@ -77,6 +77,8 @@ By @bradwerth [#6216](https://github.com/gfx-rs/wgpu/pull/6216).
 
 ### New Features
 
+64 bit image atomic support in shaders. By @atlv24 in [#5537](https://github.com/gfx-rs/wgpu/pull/5537)
+
 #### Naga
 
 - Support constant evaluation for `firstLeadingBit` and `firstTrailingBit` numeric built-ins in WGSL. Front-ends that translate to these built-ins also benefit from constant evaluation. By @ErichDonGubler in [#5101](https://github.com/gfx-rs/wgpu/pull/5101).

Cargo.toml

@@ -136,7 +136,7 @@ winit = { version = "0.29", features = ["android-native-activity"] }
 # Metal dependencies
 block = "0.1"
 core-graphics-types = "0.1"
-metal = { version = "0.29.0" }
+metal = { version = "0.29.0", git = "https://github.com/gfx-rs/metal-rs.git", rev = "ae7030be0edff4cda88ece74137e5bcd28ef48fa" }
 objc = "0.2.5"
 
 # Vulkan dependencies

naga/src/back/dot/mod.rs

@@ -254,6 +254,19 @@ impl StatementGraph {
                     }
                     "Atomic"
                 }
+                S::ImageAtomic {
+                    image,
+                    coordinate,
+                    sample,
+                    fun: _,
+                    value,
+                } => {
+                    self.dependencies.push((id, image, "image"));
+                    self.dependencies.push((id, coordinate, "coordinate"));
+                    self.dependencies.push((id, sample, "sample"));
+                    self.dependencies.push((id, value, "value"));
+                    "ImageAtomic"
+                }
                 S::WorkGroupUniformLoad { pointer, result } => {
                     self.emits.push((id, result));
                     self.dependencies.push((id, pointer, "pointer"));

naga/src/back/glsl/features.rs

@@ -400,6 +400,7 @@ impl<'a, W> Writer<'a, W> {
                             | StorageFormat::Rgb10a2Uint
                             | StorageFormat::Rgb10a2Unorm
                             | StorageFormat::Rg11b10Ufloat
+                            | StorageFormat::R64Uint
                             | StorageFormat::Rg32Uint
                             | StorageFormat::Rg32Sint
                             | StorageFormat::Rg32Float => {

naga/src/back/glsl/mod.rs

@@ -1095,12 +1095,16 @@ impl<'a, W: Write> Writer<'a, W> {
         // - Array - used if it's an image array
         // - Shadow - used if it's a depth image
         use crate::ImageClass as Ic;
-
-        let (base, kind, ms, comparison) = match class {
-            Ic::Sampled { kind, multi: true } => ("sampler", kind, "MS", ""),
-            Ic::Sampled { kind, multi: false } => ("sampler", kind, "", ""),
-            Ic::Depth { multi: true } => ("sampler", crate::ScalarKind::Float, "MS", ""),
-            Ic::Depth { multi: false } => ("sampler", crate::ScalarKind::Float, "", "Shadow"),
+        use crate::Scalar as S;
+        let float = S {
+            kind: crate::ScalarKind::Float,
+            width: 4,
+        };
+        let (base, scalar, ms, comparison) = match class {
+            Ic::Sampled { kind, multi: true } => ("sampler", S { kind, width: 4 }, "MS", ""),
+            Ic::Sampled { kind, multi: false } => ("sampler", S { kind, width: 4 }, "", ""),
+            Ic::Depth { multi: true } => ("sampler", float, "MS", ""),
+            Ic::Depth { multi: false } => ("sampler", float, "", "Shadow"),
             Ic::Storage { format, .. } => ("image", format.into(), "", ""),
         };
 
@@ -1114,7 +1118,7 @@ impl<'a, W: Write> Writer<'a, W> {
             self.out,
             "{}{}{}{}{}{}{}",
             precision,
-            glsl_scalar(crate::Scalar { kind, width: 4 })?.prefix,
+            glsl_scalar(scalar)?.prefix,
             base,
             glsl_dimension(dim),
             ms,
@@ -2465,6 +2469,17 @@ impl<'a, W: Write> Writer<'a, W> {
                 self.write_expr(value, ctx)?;
                 writeln!(self.out, ");")?;
             }
+            // Stores a value into an image.
+            Statement::ImageAtomic {
+                image,
+                coordinate,
+                sample,
+                fun,
+                value,
+            } => {
+                write!(self.out, "{level}")?;
+                self.write_image_atomic(ctx, image, coordinate, sample, fun, value)?
+            }
             Statement::RayQuery { .. } => unreachable!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
@@ -4085,6 +4100,56 @@ impl<'a, W: Write> Writer<'a, W> {
         Ok(())
     }
 
+    /// Helper method to write the `ImageStore` statement
+    fn write_image_atomic(
+        &mut self,
+        ctx: &back::FunctionCtx,
+        image: Handle<crate::Expression>,
+        coordinate: Handle<crate::Expression>,
+        _sample: Handle<crate::Expression>,
+        fun: crate::AtomicFunction,
+        value: Handle<crate::Expression>,
+    ) -> Result<(), Error> {
+        use crate::ImageDimension as IDim;
+
+        // NOTE: openGL requires that `imageStore`s have no effets when the texel is invalid
+        // so we don't need to generate bounds checks (OpenGL 4.2 Core §3.9.20)
+
+        // This will only panic if the module is invalid
+        let dim = match *ctx.resolve_type(image, &self.module.types) {
+            TypeInner::Image { dim, .. } => dim,
+            _ => unreachable!(),
+        };
+
+        // Begin our call to `imageStore`
+        let fun_str = fun.to_glsl();
+        write!(self.out, "imageAtomic{fun_str}(")?;
+        self.write_expr(image, ctx)?;
+        // Separate the image argument from the coordinates
+        write!(self.out, ", ")?;
+
+        // openGL es doesn't have 1D images so we need workaround it
+        let tex_1d_hack = dim == IDim::D1 && self.options.version.is_es();
+        // Write the coordinate vector
+        self.write_texture_coord(
+            ctx,
+            // Get the size of the coordinate vector
+            self.get_coordinate_vector_size(dim, false),
+            coordinate,
+            None,
+            tex_1d_hack,
+        )?;
+
+        // Separate the coordinate from the value to write and write the expression
+        // of the value to write.
+        write!(self.out, ", ")?;
+        self.write_expr(value, ctx)?;
+        // End the call to `imageStore` and the statement.
+        writeln!(self.out, ");")?;
+
+        Ok(())
+    }
+
     /// Helper method for writing an `ImageLoad` expression.
     #[allow(clippy::too_many_arguments)]
     fn write_image_load(
@@ -4827,6 +4892,7 @@ fn glsl_storage_format(format: crate::StorageFormat) -> Result<&'static str, Err
         Sf::Rgb10a2Uint => "rgb10_a2ui",
         Sf::Rgb10a2Unorm => "rgb10_a2",
         Sf::Rg11b10Ufloat => "r11f_g11f_b10f",
+        Sf::R64Uint => "r64ui",
         Sf::Rg32Uint => "rg32ui",
         Sf::Rg32Sint => "rg32i",
         Sf::Rg32Float => "rg32f",

naga/src/back/hlsl/conv.rs

@@ -122,7 +122,7 @@ impl crate::StorageFormat {
             Self::R16Float | Self::R32Float => "float",
             Self::R8Unorm | Self::R16Unorm => "unorm float",
             Self::R8Snorm | Self::R16Snorm => "snorm float",
-            Self::R8Uint | Self::R16Uint | Self::R32Uint => "uint",
+            Self::R8Uint | Self::R16Uint | Self::R32Uint | Self::R64Uint => "uint",
             Self::R8Sint | Self::R16Sint | Self::R32Sint => "int",
 
             Self::Rg16Float | Self::Rg32Float => "float2",

naga/src/back/hlsl/writer.rs

@@ -2209,6 +2209,25 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
                 writeln!(self.out, ");")?;
             }
+            crate::Statement::ImageAtomic {
+                image,
+                coordinate,
+                sample: _,
+                fun,
+                value,
+            } => {
+                write!(self.out, "{level}")?;
+
+                let fun_str = fun.to_hlsl_suffix();
+                write!(self.out, "Interlocked{fun_str}(")?;
+                self.write_expr(module, image, func_ctx)?;
+                write!(self.out, "[")?;
+                self.write_expr(module, coordinate, func_ctx)?;
+                write!(self.out, "],")?;
+
+                self.write_expr(module, value, func_ctx)?;
+                writeln!(self.out, ");")?;
+            }
             Statement::WorkGroupUniformLoad { pointer, result } => {
                 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
                 write!(self.out, "{level}")?;

naga/src/back/msl/writer.rs

@@ -215,22 +215,27 @@ impl<'a> Display for TypeContext<'a> {
                     crate::ImageDimension::D3 => "3d",
                     crate::ImageDimension::Cube => "cube",
                 };
-                let (texture_str, msaa_str, kind, access) = match class {
+                let (texture_str, msaa_str, scalar, access) = match class {
                     crate::ImageClass::Sampled { kind, multi } => {
                         let (msaa_str, access) = if multi {
                             ("_ms", "read")
                         } else {
                             ("", "sample")
                         };
-                        ("texture", msaa_str, kind, access)
+                        let scalar = crate::Scalar { kind, width: 4 };
+                        ("texture", msaa_str, scalar, access)
                     }
                     crate::ImageClass::Depth { multi } => {
                         let (msaa_str, access) = if multi {
                             ("_ms", "read")
                         } else {
                             ("", "sample")
                         };
-                        ("depth", msaa_str, crate::ScalarKind::Float, access)
+                        let scalar = crate::Scalar {
+                            kind: crate::ScalarKind::Float,
+                            width: 4,
+                        };
+                        ("depth", msaa_str, scalar, access)
                     }
                     crate::ImageClass::Storage { format, .. } => {
                         let access = if self
@@ -254,7 +259,7 @@ impl<'a> Display for TypeContext<'a> {
                         ("texture", "", format.into(), access)
                     }
                 };
-                let base_name = crate::Scalar { kind, width: 4 }.to_msl_name();
+                let base_name = scalar.to_msl_name();
                 let array_str = if arrayed { "_array" } else { "" };
                 write!(
                     out,
@@ -1187,6 +1192,27 @@ impl<W: Write> Writer<W> {
         Ok(())
     }
 
+    fn put_image_atomic(
+        &mut self,
+        level: back::Level,
+        image: Handle<crate::Expression>,
+        address: &TexelAddress,
+        fun: crate::AtomicFunction,
+        value: Handle<crate::Expression>,
+        context: &StatementContext,
+    ) -> BackendResult {
+        write!(self.out, "{level}")?;
+        self.put_expression(image, &context.expression, false)?;
+        write!(self.out, ".atomic_{}(", fun.to_msl_64_bit()?)?;
+        // coordinates in IR are int, but Metal expects uint
+        self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
+        write!(self.out, ", ")?;
+        self.put_expression(value, &context.expression, true)?;
+        writeln!(self.out, ");")?;
+
+        Ok(())
+    }
+
     fn put_image_store(
         &mut self,
         level: back::Level,
@@ -3208,6 +3234,21 @@ impl<W: Write> Writer<W> {
                     // Done
                     writeln!(self.out, ";")?;
                 }
+                crate::Statement::ImageAtomic {
+                    image,
+                    coordinate,
+                    sample,
+                    fun,
+                    value,
+                } => {
+                    let address = TexelAddress {
+                        coordinate,
+                        array_index: None,
+                        sample: Some(sample),
+                        level: None,
+                    };
+                    self.put_image_atomic(level, image, &address, fun, value, context)?
+                }
                 crate::Statement::WorkGroupUniformLoad { pointer, result } => {
                     self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
 

naga/src/back/pipeline_constants.rs

@@ -638,6 +638,18 @@ fn adjust_stmt(new_pos: &HandleVec<Expression, Handle<Expression>>, stmt: &mut S
                 | crate::AtomicFunction::Exchange { compare: None } => {}
             }
         }
+        Statement::ImageAtomic {
+            ref mut image,
+            ref mut coordinate,
+            ref mut sample,
+            fun: _,
+            ref mut value,
+        } => {
+            adjust(image);
+            adjust(coordinate);
+            adjust(sample);
+            adjust(value);
+        }
         Statement::WorkGroupUniformLoad {
             ref mut pointer,
             ref mut result,

naga/src/back/spv/block.rs

@@ -2867,6 +2867,15 @@ impl<'w> BlockContext<'w> {
 
                     block.body.push(instruction);
                 }
+                Statement::ImageAtomic {
+                    image,
+                    coordinate,
+                    sample,
+                    fun,
+                    value,
+                } => {
+                    self.write_image_atomic(image, coordinate, sample, fun, value, &mut block)?;
+                }
                 Statement::WorkGroupUniformLoad { pointer, result } => {
                     self.writer
                         .write_barrier(crate::Barrier::WORK_GROUP, &mut block);