[Mlir-commits] [mlir] [mlir][SPIR-V] Add AtomicLoad and AtomicStore ops (PR #195797)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon May 4 23:14:09 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Arseniy Obolenskiy (aobolensk)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/195797.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td (+89)
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td (+4-1)
- (modified) mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir (+60)
- (modified) mlir/test/Target/SPIRV/atomic-ops.mlir (+8)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td
index 202fa0ad60d2f..d756a923b744f 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td
@@ -231,6 +231,95 @@ def SPIRV_AtomicCompareExchangeWeakOp : SPIRV_Op<"AtomicCompareExchangeWeak", [
// -----
+def SPIRV_AtomicLoadOp : SPIRV_Op<"AtomicLoad", [
+ PointeeTypeMatchTrait<"pointer", "result">,
+]> {
+ let summary = "Atomically load through Pointer using the given Semantics.";
+
+ let description = [{
+ All subparts of the value that is loaded are read atomically with respect
+ to all other atomic accesses to it within Scope.
+
+ Result Type must be a scalar of integer type or floating-point type.
+
+ The type of the value pointed to by Pointer must be the same as Result
+ Type.
+
+ Memory is a memory Scope.
+
+ <!-- End of AutoGen section -->
+
+ #### Example:
+
+ ```mlir
+ %0 = spirv.AtomicLoad <Workgroup> <Acquire> %pointer :
+ !spirv.ptr<i32, Workgroup>
+ ```
+ }];
+
+ let arguments = (ins
+ SPIRV_AnyPtr:$pointer,
+ SPIRV_ScopeAttr:$memory_scope,
+ SPIRV_MemorySemanticsAttr:$semantics
+ );
+
+ let results = (outs
+ SPIRV_Numerical:$result
+ );
+
+ let assemblyFormat = [{
+ $memory_scope $semantics operands attr-dict `:` type($pointer)
+ }];
+
+ let hasCustomAssemblyFormat = 0;
+ let hasVerifier = 0;
+}
+
+// -----
+
+def SPIRV_AtomicStoreOp : SPIRV_Op<"AtomicStore", [
+ PointeeTypeMatchTrait<"pointer", "value">,
+]> {
+ let summary = "Atomically store through Pointer using the given Semantics.";
+
+ let description = [{
+ All subparts of Value are written atomically with respect to all other
+ atomic accesses to it within Scope.
+
+ The type of the value pointed to by Pointer must be the same as the type
+ of Value, and must be a scalar of integer type or floating-point type.
+
+ Memory is a memory Scope.
+
+ <!-- End of AutoGen section -->
+
+ #### Example:
+
+ ```mlir
+ spirv.AtomicStore <Workgroup> <Release> %pointer, %value :
+ !spirv.ptr<i32, Workgroup>
+ ```
+ }];
+
+ let arguments = (ins
+ SPIRV_AnyPtr:$pointer,
+ SPIRV_ScopeAttr:$memory_scope,
+ SPIRV_MemorySemanticsAttr:$semantics,
+ SPIRV_Numerical:$value
+ );
+
+ let results = (outs);
+
+ let assemblyFormat = [{
+ $memory_scope $semantics operands attr-dict `:` type($pointer)
+ }];
+
+ let hasCustomAssemblyFormat = 0;
+ let hasVerifier = 0;
+}
+
+// -----
+
def SPIRV_AtomicExchangeOp : SPIRV_Op<"AtomicExchange", [
PointeeTypeMatchTrait<"pointer", "value">,
PointeeTypeMatchTrait<"pointer", "result">,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index c13d24942a9e9..27d8ec3496d36 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4573,6 +4573,8 @@ def SPIRV_OC_OpEmitVertex : I32EnumAttrCase<"OpEmitVertex",
def SPIRV_OC_OpEndPrimitive : I32EnumAttrCase<"OpEndPrimitive", 219>;
def SPIRV_OC_OpControlBarrier : I32EnumAttrCase<"OpControlBarrier", 224>;
def SPIRV_OC_OpMemoryBarrier : I32EnumAttrCase<"OpMemoryBarrier", 225>;
+def SPIRV_OC_OpAtomicLoad : I32EnumAttrCase<"OpAtomicLoad", 227>;
+def SPIRV_OC_OpAtomicStore : I32EnumAttrCase<"OpAtomicStore", 228>;
def SPIRV_OC_OpAtomicExchange : I32EnumAttrCase<"OpAtomicExchange", 229>;
def SPIRV_OC_OpAtomicCompareExchange : I32EnumAttrCase<"OpAtomicCompareExchange", 230>;
def SPIRV_OC_OpAtomicCompareExchangeWeak : I32EnumAttrCase<"OpAtomicCompareExchangeWeak", 231>;
@@ -4749,7 +4751,8 @@ def SPIRV_OpcodeAttr :
SPIRV_OC_OpNot, SPIRV_OC_OpBitFieldInsert, SPIRV_OC_OpBitFieldSExtract,
SPIRV_OC_OpBitFieldUExtract, SPIRV_OC_OpBitReverse, SPIRV_OC_OpBitCount,
SPIRV_OC_OpEmitVertex, SPIRV_OC_OpEndPrimitive, SPIRV_OC_OpControlBarrier,
- SPIRV_OC_OpMemoryBarrier, SPIRV_OC_OpAtomicExchange,
+ SPIRV_OC_OpMemoryBarrier, SPIRV_OC_OpAtomicLoad, SPIRV_OC_OpAtomicStore,
+ SPIRV_OC_OpAtomicExchange,
SPIRV_OC_OpAtomicCompareExchange, SPIRV_OC_OpAtomicCompareExchangeWeak,
SPIRV_OC_OpAtomicIIncrement, SPIRV_OC_OpAtomicIDecrement,
SPIRV_OC_OpAtomicIAdd, SPIRV_OC_OpAtomicISub, SPIRV_OC_OpAtomicSMin,
diff --git a/mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir b/mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir
index 661497d5fff38..845ea6c98ed67 100644
--- a/mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir
@@ -177,6 +177,36 @@ func.func @atomic_isub(%ptr : !spirv.ptr<i32, StorageBuffer>, %value : i32) -> i
return %0 : i32
}
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.AtomicLoad
+//===----------------------------------------------------------------------===//
+
+func.func @atomic_load(%ptr : !spirv.ptr<i32, Workgroup>) -> i32 {
+ // CHECK: spirv.AtomicLoad <Workgroup> <Acquire> %{{.*}} : !spirv.ptr<i32, Workgroup>
+ %0 = spirv.AtomicLoad <Workgroup> <Acquire> %ptr : !spirv.ptr<i32, Workgroup>
+ return %0 : i32
+}
+
+// -----
+
+func.func @atomic_load_float(%ptr : !spirv.ptr<f32, StorageBuffer>) -> f32 {
+ // CHECK: spirv.AtomicLoad <Device> <None> %{{.*}} : !spirv.ptr<f32, StorageBuffer>
+ %0 = spirv.AtomicLoad <Device> <None> %ptr : !spirv.ptr<f32, StorageBuffer>
+ return %0 : f32
+}
+
+// -----
+
+func.func @atomic_load_mismatch(%ptr : !spirv.ptr<i32, Workgroup>) -> i64 {
+ // expected-error @+1 {{'spirv.AtomicLoad' op failed to verify that `result` type matches pointee type of `pointer`}}
+ %0 = "spirv.AtomicLoad"(%ptr) {memory_scope = #spirv.scope<Workgroup>, semantics = #spirv.memory_semantics<Acquire>} : (!spirv.ptr<i32, Workgroup>) -> (i64)
+ return %0 : i64
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.AtomicOr
//===----------------------------------------------------------------------===//
@@ -207,6 +237,36 @@ func.func @atomic_smin(%ptr : !spirv.ptr<i32, StorageBuffer>, %value : i32) -> i
return %0 : i32
}
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.AtomicStore
+//===----------------------------------------------------------------------===//
+
+func.func @atomic_store(%ptr : !spirv.ptr<i32, Workgroup>, %value : i32) {
+ // CHECK: spirv.AtomicStore <Workgroup> <Release> %{{.*}}, %{{.*}} : !spirv.ptr<i32, Workgroup>
+ spirv.AtomicStore <Workgroup> <Release> %ptr, %value : !spirv.ptr<i32, Workgroup>
+ return
+}
+
+// -----
+
+func.func @atomic_store_float(%ptr : !spirv.ptr<f32, StorageBuffer>, %value : f32) {
+ // CHECK: spirv.AtomicStore <Device> <None> %{{.*}}, %{{.*}} : !spirv.ptr<f32, StorageBuffer>
+ spirv.AtomicStore <Device> <None> %ptr, %value : !spirv.ptr<f32, StorageBuffer>
+ return
+}
+
+// -----
+
+func.func @atomic_store_mismatch(%ptr : !spirv.ptr<i32, Workgroup>, %value : i64) {
+ // expected-error @+1 {{'spirv.AtomicStore' op failed to verify that `value` type matches pointee type of `pointer`}}
+ "spirv.AtomicStore"(%ptr, %value) {memory_scope = #spirv.scope<Workgroup>, semantics = #spirv.memory_semantics<Release>} : (!spirv.ptr<i32, Workgroup>, i64) -> ()
+ return
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.AtomicUMax
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/SPIRV/atomic-ops.mlir b/mlir/test/Target/SPIRV/atomic-ops.mlir
index 3e4908123a345..aecc60cbbc71f 100644
--- a/mlir/test/Target/SPIRV/atomic-ops.mlir
+++ b/mlir/test/Target/SPIRV/atomic-ops.mlir
@@ -36,6 +36,10 @@ spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Linkage, Addre
%12 = spirv.AtomicCompareExchange <Workgroup> <Release> <Acquire> %ptr, %value, %comparator: !spirv.ptr<i32, Workgroup>
// CHECK: spirv.AtomicExchange <Workgroup> <Release> %{{.*}}, %{{.*}} : !spirv.ptr<i32, Workgroup>
%13 = spirv.AtomicExchange <Workgroup> <Release> %ptr, %value: !spirv.ptr<i32, Workgroup>
+ // CHECK: spirv.AtomicLoad <Workgroup> <Acquire> %{{.*}} : !spirv.ptr<i32, Workgroup>
+ %14 = spirv.AtomicLoad <Workgroup> <Acquire> %ptr : !spirv.ptr<i32, Workgroup>
+ // CHECK: spirv.AtomicStore <Workgroup> <Release> %{{.*}}, %{{.*}} : !spirv.ptr<i32, Workgroup>
+ spirv.AtomicStore <Workgroup> <Release> %ptr, %value : !spirv.ptr<i32, Workgroup>
spirv.ReturnValue %0: i32
}
@@ -43,6 +47,10 @@ spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Linkage, Addre
spirv.func @test_float_atomics(%ptr: !spirv.ptr<f32, Workgroup>, %value: f32) -> f32 "None" {
// CHECK: spirv.EXT.AtomicFAdd <Workgroup> <Acquire> %{{.*}}, %{{.*}} : !spirv.ptr<f32, Workgroup>
%0 = spirv.EXT.AtomicFAdd <Workgroup> <Acquire> %ptr, %value : !spirv.ptr<f32, Workgroup>
+ // CHECK: spirv.AtomicLoad <Workgroup> <Acquire> %{{.*}} : !spirv.ptr<f32, Workgroup>
+ %1 = spirv.AtomicLoad <Workgroup> <Acquire> %ptr : !spirv.ptr<f32, Workgroup>
+ // CHECK: spirv.AtomicStore <Workgroup> <Release> %{{.*}}, %{{.*}} : !spirv.ptr<f32, Workgroup>
+ spirv.AtomicStore <Workgroup> <Release> %ptr, %value : !spirv.ptr<f32, Workgroup>
spirv.ReturnValue %0: f32
}
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/195797
More information about the Mlir-commits
mailing list