[Mlir-commits] [mlir] [mlir][SPIRV] update SPIRV Atomic Ops to assemblyFormat (PR #76323)

[Jakub Kuderski]

================
@@ -14,45 +14,55 @@
 #ifndef MLIR_DIALECT_SPIRV_IR_ATOMIC_OPS
 #define MLIR_DIALECT_SPIRV_IR_ATOMIC_OPS
 
-class SPIRV_AtomicUpdateOp<string mnemonic, list<Trait> traits = []> :
-  SPIRV_Op<mnemonic, traits> {
-  let arguments = (ins
-    SPIRV_AnyPtr:$pointer,
-    SPIRV_ScopeAttr:$memory_scope,
-    SPIRV_MemorySemanticsAttr:$semantics
-  );
-
-  let results = (outs
-    SPIRV_Integer:$result
-  );
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+class PointeeTypeMatchTrait<string pointer, string name>
+    : TypesMatchWith<"$" # name # " type matches pointee type of " # "$" # pointer, pointer,
+                     name, "$_self.cast<PointerType>().getPointeeType()">;
+
+// -----
+
+class SPIRV_AtomicUpdateOp<string mnemonic, list<Trait> traits = []>
+    : SPIRV_Op<mnemonic,
+               !listconcat(traits,
+                           [PointeeTypeMatchTrait<"pointer", "result">])> {
+  let arguments = (ins SPIRV_AnyPtr
+                   : $pointer, SPIRV_ScopeAttr
+                   : $memory_scope, SPIRV_MemorySemanticsAttr
+                   : $semantics);
+
+  let results = (outs SPIRV_Integer : $result);
+
+  let assemblyFormat = [{
+    $memory_scope $semantics operands attr-dict `:` type($pointer)
+  }];
 }
 
-class SPIRV_AtomicUpdateWithValueOp<string mnemonic, list<Trait> traits = []> :
-  SPIRV_Op<mnemonic, traits> {
-  let arguments = (ins
-    SPIRV_AnyPtr:$pointer,
-    SPIRV_ScopeAttr:$memory_scope,
-    SPIRV_MemorySemanticsAttr:$semantics,
-    SPIRV_Integer:$value
-  );
-
-  let results = (outs
-    SPIRV_Integer:$result
-  );
-
-  let builders = [
-    OpBuilder<(ins "Value":$pointer, "::mlir::spirv::Scope":$scope,
-      "::mlir::spirv::MemorySemantics":$memory, "Value":$value),
-    [{build($_builder, $_state, value.getType(), pointer, scope, memory, value);}]>
-  ];
+class SPIRV_AtomicUpdateWithValueOp<string mnemonic, list<Trait> traits = []>
+    : SPIRV_Op<mnemonic, !listconcat(traits, [
+                 PointeeTypeMatchTrait<"pointer", "value">,
+                 PointeeTypeMatchTrait<"pointer", "result">,
+               ])> {
+  let arguments = (ins SPIRV_AnyPtr
+                   : $pointer, SPIRV_ScopeAttr
+                   : $memory_scope, SPIRV_MemorySemanticsAttr
+                   : $semantics, SPIRV_Integer
+                   : $value);
+
+  let results = (outs SPIRV_Integer : $result);
+
+  let assemblyFormat = [{
+    $memory_scope $semantics operands attr-dict `:` type($pointer)
+  }];
 }
 
 // -----
 
 def SPIRV_AtomicAndOp : SPIRV_AtomicUpdateWithValueOp<"AtomicAnd", []> {
   let summary = [{
     Perform the following steps atomically with respect to any other atomic
-    accesses within Scope to the same location:
+        accesses within Scope to the same location:
----------------
kuhar wrote:

nit: accidental change?

https://github.com/llvm/llvm-project/pull/76323