Add a OpAtomicIIncrement function to arch (#839)

Author: expenses

crates/spirv-std/src/arch.rs

@@ -226,6 +226,29 @@ pub fn signed_max<T: SignedInteger>(a: T, b: T) -> T {
     unsafe { call_glsl_op_with_ints::<_, 42>(a, b) }
 }
 
+/// Atomically increment an integer and return the old value.
+#[spirv_std_macros::gpu_only]
+#[doc(alias = "OpAtomicIIncrement")]
+pub unsafe fn atomic_i_increment<I: Integer, const SCOPE: u32, const SEMANTICS: u32>(
+    ptr: &mut I,
+) -> I {
+    let mut old = I::default();
+
+    asm! {
+        "%u32 = OpTypeInt 32 0",
+        "%scope = OpConstant %u32 {scope}",
+        "%semantics = OpConstant %u32 {semantics}",
+        "%old = OpAtomicIIncrement _ {ptr} %scope %semantics",
+        "OpStore {old} %old",
+        scope = const SCOPE,
+        semantics = const SEMANTICS,
+        ptr = in(reg) ptr,
+        old = in(reg) &mut old,
+    }
+
+    old
+}
+
 /// Index into an array without bounds checking.
 ///
 /// The main purpose of this trait is to work around the fact that the regular `get_unchecked*`

tests/ui/arch/atomic_i_increment.rs

@@ -0,0 +1,17 @@
+// build-pass
+
+use spirv_std::arch::IndexUnchecked;
+
+#[spirv(compute(threads(64)))]
+pub fn main(#[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buffer: &mut [u32]) {
+    let reference = unsafe { buffer.index_unchecked_mut(0) };
+
+    let old = unsafe {
+        spirv_std::arch::atomic_i_increment::<
+            _,
+            { spirv_std::memory::Scope::Workgroup as u32 },
+            { spirv_std::memory::Semantics::NONE.bits() as u32 },
+        >(reference)
+    };
+    assert!(old == 0);
+}