Image atomics support (#6706)

* Image atomics support

* Address feedback

* fix merge

* Fixes

* Add a couple tests

* Update wgpu-types/src/lib.rs

Co-authored-by: Teodor Tanasoaia <28601907+teoxoy@users.noreply.github.com>

* feedback

* feedback

* glsl

* glsl fix

* fix glsl

* fix fix

* fix fix fic

* fix?

* fix

---------

Co-authored-by: Connor Fitzgerald <connorwadefitzgerald@gmail.com>
Co-authored-by: Teodor Tanasoaia <28601907+teoxoy@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -164,6 +164,7 @@ By @wumpf in [#6849](https://github.com/gfx-rs/wgpu/pull/6849).
 - `DeviceType` and `AdapterInfo` now impl `Hash` by @cwfitzgerald in [#6868](https://github.com/gfx-rs/wgpu/pull/6868)
 - Add build support for Apple Vision Pro. By @guusw in [#6611](https://github.com/gfx-rs/wgpu/pull/6611).
 - Add `wgsl_language_features` for obtaining available WGSL language feature by @sagudev in [#6814](https://github.com/gfx-rs/wgpu/pull/6814)
+- Image atomic support in shaders. By @atlv24 in [#6706](https://github.com/gfx-rs/wgpu/pull/6706)
 - Add `no_std` support to `wgpu-types`. By @bushrat011899 in [#6892](https://github.com/gfx-rs/wgpu/pull/6892).
 
 ##### Vulkan

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/features.rs

@@ -52,6 +52,8 @@ bitflags::bitflags! {
         const TEXTURE_SHADOW_LOD = 1 << 23;
         /// Subgroup operations
         const SUBGROUP_OPERATIONS = 1 << 24;
+        /// Image atomics
+        const TEXTURE_ATOMICS = 1 << 25;
     }
 }
 
@@ -120,6 +122,7 @@ impl FeaturesManager {
         check_feature!(DYNAMIC_ARRAY_SIZE, 430, 310);
         check_feature!(DUAL_SOURCE_BLENDING, 330, 300 /* with extension */);
         check_feature!(SUBGROUP_OPERATIONS, 430, 310);
+        check_feature!(TEXTURE_ATOMICS, 420, 310);
         match version {
             Version::Embedded { is_webgl: true, .. } => check_feature!(MULTI_VIEW, 140, 300),
             _ => check_feature!(MULTI_VIEW, 140, 310),
@@ -278,6 +281,11 @@ impl FeaturesManager {
             )?;
         }
 
+        if self.0.contains(Features::TEXTURE_ATOMICS) {
+            // https://www.khronos.org/registry/OpenGL/extensions/OES/OES_shader_image_atomic.txt
+            writeln!(out, "#extension GL_OES_shader_image_atomic : require")?;
+        }
+
         Ok(())
     }
 }
@@ -546,6 +554,22 @@ impl<W> Writer<'_, W> {
             }
         }
 
+        for blocks in module
+            .functions
+            .iter()
+            .map(|(_, f)| &f.body)
+            .chain(std::iter::once(&entry_point.function.body))
+        {
+            for (stmt, _) in blocks.span_iter() {
+                match *stmt {
+                    crate::Statement::ImageAtomic { .. } => {
+                        features.request(Features::TEXTURE_ATOMICS)
+                    }
+                    _ => {}
+                }
+            }
+        }
+
         self.features.check_availability(self.options.version)
     }
 

naga/src/back/glsl/mod.rs

@@ -2475,6 +2475,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}")?;
@@ -4137,6 +4148,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(
@@ -4533,6 +4594,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

@@ -344,7 +344,7 @@ 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)
                 {
                     "const"
                 } else {
@@ -1201,6 +1201,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,
@@ -3248,6 +3270,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

@@ -2921,6 +2921,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(())
+    }
 }