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

[Alex Beloi]

================
@@ -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">,
----------------
alexbeloi wrote:

For some reason using `AllTypesMatch` with `PointeeTypeMatch` isn't propagating the type inference, I tried looking through [Operator.cpp](https://github.com/llvm/llvm-project/blob/155d5849da2b2bfa2da918923d8f148a96c03e72/mlir/lib/TableGen/Operator.cpp#L363) which looks like should be handling this but couldn't spot any obvious issue.

I tried debugging with gdb to see which variables are fully inferred but I can't get any local variables (or line numbers) from gdb even with `-DCMAKE_BUILD_TYPE=Debug`.

Adding something like `AllTypesMatchPointeeType<"pointer", [...]>` would need its own handling in `Operator::populateTypeInferenceInfo`.

Do you mind if we leave resolving this as a follow-up item?

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