Image atomics support

CHANGELOG.md

@@ -105,6 +105,8 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148]
 
 ### New Features
 
+Image atomic support in shaders. By @atlv24 in [#6706](https://github.com/gfx-rs/wgpu/pull/6706)
+
 #### Naga
 
 - Clean up tests for atomic operations support in SPIR-V frontend. By @schell in [#6692](https://github.com/gfx-rs/wgpu/pull/6692)

naga/src/back/dot/mod.rs

@@ -254,6 +254,21 @@ impl StatementGraph {
                     }
                     "Atomic"
                 }
+                S::ImageAtomic {
+                    image,
+                    coordinate,
+                    array_index,
+                    fun: _,
+                    value,
+                } => {
+                    self.dependencies.push((id, image, "image"));
+                    self.dependencies.push((id, coordinate, "coordinate"));
+                    if let Some(expr) = array_index {
+                        self.dependencies.push((id, expr, "array_index"));
+                    }
+                    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/mod.rs

@@ -2472,6 +2472,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,
+                array_index,
+                fun,
+                value,
+            } => {
+                write!(self.out, "{level}")?;
+                self.write_image_atomic(ctx, image, coordinate, array_index, fun, value)?
+            }
             Statement::RayQuery { .. } => unreachable!(),
             Statement::SubgroupBallot { result, predicate } => {
                 write!(self.out, "{level}")?;
@@ -4134,6 +4145,56 @@ impl<'a, W: Write> Writer<'a, W> {
         Ok(())
     }
 
+    /// Helper method to write the `ImageAtomic` statement
+    fn write_image_atomic(
+        &mut self,
+        ctx: &back::FunctionCtx,
+        image: Handle<crate::Expression>,
+        coordinate: Handle<crate::Expression>,
+        array_index: Option<Handle<crate::Expression>>,
+        fun: crate::AtomicFunction,
+        value: Handle<crate::Expression>,
+    ) -> Result<(), Error> {
+        use crate::ImageDimension as IDim;
+
+        // NOTE: openGL requires that `imageAtomic`s have no effects 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 `imageAtomic`
+        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,
+            array_index,
+            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 `imageAtomic` and the statement.
+        writeln!(self.out, ");")?;
+
+        Ok(())
+    }
+
     /// Helper method for writing an `ImageLoad` expression.
     #[allow(clippy::too_many_arguments)]
     fn write_image_load(
@@ -4530,6 +4591,9 @@ impl<'a, W: Write> Writer<'a, W> {
     /// they can only be used to query information about the resource which isn't what
     /// we want here so when storage access is both `LOAD` and `STORE` add no modifiers
     fn write_storage_access(&mut self, storage_access: crate::StorageAccess) -> BackendResult {
+        if storage_access.contains(crate::StorageAccess::ATOMIC) {
+            return Ok(());
+        }
         if !storage_access.contains(crate::StorageAccess::STORE) {
             write!(self.out, "readonly ")?;
         }

naga/src/back/hlsl/writer.rs

@@ -2210,6 +2210,32 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
 
                 writeln!(self.out, ");")?;
             }
+            Statement::ImageAtomic {
+                image,
+                coordinate,
+                array_index,
+                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_texture_coordinates(
+                    "int",
+                    coordinate,
+                    array_index,
+                    None,
+                    module,
+                    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

@@ -343,7 +343,9 @@ impl TypedGlobalVariable<'_> {
         let (space, access, reference) = match var.space.to_msl_name() {
             Some(space) if self.reference => {
                 let access = if var.space.needs_access_qualifier()
-                    && !self.usage.contains(valid::GlobalUse::WRITE)
+                    && !self
+                        .usage
+                        .intersects(valid::GlobalUse::WRITE | valid::GlobalUse::ATOMIC)
                 {
                     "const"
                 } else {
@@ -1198,6 +1200,28 @@ 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)?;
+        let op = fun.to_msl();
+        write!(self.out, ".atomic_{}(", op)?;
+        // 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,
@@ -3236,6 +3260,21 @@ impl<W: Write> Writer<W> {
                     // Done
                     writeln!(self.out, ";")?;
                 }
+                crate::Statement::ImageAtomic {
+                    image,
+                    coordinate,
+                    array_index,
+                    fun,
+                    value,
+                } => {
+                    let address = TexelAddress {
+                        coordinate,
+                        array_index,
+                        sample: None,
+                        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

@@ -736,6 +736,20 @@ 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 array_index,
+            fun: _,
+            ref mut value,
+        } => {
+            adjust(image);
+            adjust(coordinate);
+            if let Some(ref mut array_index) = *array_index {
+                adjust(array_index);
+            }
+            adjust(value);
+        }
         Statement::WorkGroupUniformLoad {
             ref mut pointer,
             ref mut result,

naga/src/back/spv/block.rs

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

naga/src/back/spv/image.rs

@@ -1225,4 +1225,77 @@ impl BlockContext<'_> {
 
         Ok(())
     }
+
+    pub(super) fn write_image_atomic(
+        &mut self,
+        image: Handle<crate::Expression>,
+        coordinate: Handle<crate::Expression>,
+        array_index: Option<Handle<crate::Expression>>,
+        fun: crate::AtomicFunction,
+        value: Handle<crate::Expression>,
+        block: &mut Block,
+    ) -> Result<(), Error> {
+        let image_id = match self.ir_function.originating_global(image) {
+            Some(handle) => self.writer.global_variables[handle].var_id,
+            _ => return Err(Error::Validation("Unexpected image type")),
+        };
+        let crate::TypeInner::Image { class, .. } =
+            *self.fun_info[image].ty.inner_with(&self.ir_module.types)
+        else {
+            return Err(Error::Validation("Invalid image type"));
+        };
+        let crate::ImageClass::Storage { format, .. } = class else {
+            return Err(Error::Validation("Invalid image class"));
+        };
+        let scalar = format.into();
+        let pointer_type_id = self.get_type_id(LookupType::Local(LocalType::LocalPointer {
+            base: NumericType::Scalar(scalar),
+            class: spirv::StorageClass::Image,
+        }));
+        let signed = scalar.kind == crate::ScalarKind::Sint;
+        let pointer_id = self.gen_id();
+        let coordinates = self.write_image_coordinates(coordinate, array_index, block)?;
+        let sample_id = self.writer.get_constant_scalar(crate::Literal::U32(0));
+        block.body.push(Instruction::image_texel_pointer(
+            pointer_type_id,
+            pointer_id,
+            image_id,
+            coordinates.value_id,
+            sample_id,
+        ));
+
+        let op = match fun {
+            crate::AtomicFunction::Add => spirv::Op::AtomicIAdd,
+            crate::AtomicFunction::Subtract => spirv::Op::AtomicISub,
+            crate::AtomicFunction::And => spirv::Op::AtomicAnd,
+            crate::AtomicFunction::ExclusiveOr => spirv::Op::AtomicXor,
+            crate::AtomicFunction::InclusiveOr => spirv::Op::AtomicOr,
+            crate::AtomicFunction::Min if signed => spirv::Op::AtomicSMin,
+            crate::AtomicFunction::Min => spirv::Op::AtomicUMin,
+            crate::AtomicFunction::Max if signed => spirv::Op::AtomicSMax,
+            crate::AtomicFunction::Max => spirv::Op::AtomicUMax,
+            crate::AtomicFunction::Exchange { .. } => {
+                return Err(Error::Validation("Exchange atomics are not supported yet"))
+            }
+        };
+        let result_type_id = self.get_expression_type_id(&self.fun_info[value].ty);
+        let id = self.gen_id();
+        let space = crate::AddressSpace::Handle;
+        let (semantics, scope) = space.to_spirv_semantics_and_scope();
+        let scope_constant_id = self.get_scope_constant(scope as u32);
+        let semantics_id = self.get_index_constant(semantics.bits());
+        let value_id = self.cached[value];
+
+        block.body.push(Instruction::image_atomic(
+            op,
+            result_type_id,
+            id,
+            pointer_id,
+            scope_constant_id,
+            semantics_id,
+            value_id,
+        ));
+
+        Ok(())
+    }
 }

naga/src/back/spv/instructions.rs

@@ -702,6 +702,41 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn image_texel_pointer(
+        result_type_id: Word,
+        id: Word,
+        image: Word,
+        coordinates: Word,
+        sample: Word,
+    ) -> Self {
+        let mut instruction = Self::new(Op::ImageTexelPointer);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(image);
+        instruction.add_operand(coordinates);
+        instruction.add_operand(sample);
+        instruction
+    }
+
+    pub(super) fn image_atomic(
+        op: Op,
+        result_type_id: Word,
+        id: Word,
+        pointer: Word,
+        scope_id: Word,
+        semantics_id: Word,
+        value: Word,
+    ) -> Self {
+        let mut instruction = Self::new(op);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(pointer);
+        instruction.add_operand(scope_id);
+        instruction.add_operand(semantics_id);
+        instruction.add_operand(value);
+        instruction
+    }
+
     pub(super) fn image_query(op: Op, result_type_id: Word, id: Word, image: Word) -> Self {
         let mut instruction = Self::new(op);
         instruction.set_type(result_type_id);