[Mlir-commits] [mlir] aa16ca3 - [mlir] Add getAlias for OpAsmTypeInterface (#126364)

[llvmlistbot at llvm.org]

Author: Hongren Zheng
Date: 2025-02-19T02:28:49+08:00
New Revision: aa16ca3422d6ed4847bffb469cc5a1db14dca0b3

URL: https://github.com/llvm/llvm-project/commit/aa16ca3422d6ed4847bffb469cc5a1db14dca0b3
DIFF: https://github.com/llvm/llvm-project/commit/aa16ca3422d6ed4847bffb469cc5a1db14dca0b3.diff

LOG: [mlir] Add getAlias for OpAsmTypeInterface (#126364)

See
https://discourse.llvm.org/t/rfc-introduce-opasm-type-attr-interface-for-pretty-print-in-asmprinter/83792
for detailed introduction.

This PR should be rebased once #124721 is merged.

This PR adds

* Definition of `getAlias` for `OpAsmTypeInterface`
* Integration of `OpAsmTypeInterface` with `AsmPrinter` alias handling
part

This is partly in response to
https://github.com/llvm/llvm-project/pull/124721/files#r1940399862

Cc @River707 for review.

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpAsmInterface.td
    mlir/lib/IR/AsmPrinter.cpp
    mlir/test/IR/op-asm-interface.mlir
    mlir/test/lib/Dialect/Test/TestTypeDefs.td
    mlir/test/lib/Dialect/Test/TestTypes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpAsmInterface.td b/mlir/include/mlir/IR/OpAsmInterface.td
index c3e84bccc5dee..1bd8eb04714c5 100644
--- a/mlir/include/mlir/IR/OpAsmInterface.td
+++ b/mlir/include/mlir/IR/OpAsmInterface.td
@@ -127,6 +127,13 @@ def OpAsmTypeInterface : TypeInterface<"OpAsmTypeInterface"> {
       "void", "getAsmName",
       (ins "::mlir::OpAsmSetNameFn":$setNameFn), "", ";"
     >,
+    InterfaceMethod<[{
+        Get a name to use when generating an alias for this type.
+      }],
+      "::mlir::OpAsmDialectInterface::AliasResult", "getAlias",
+      (ins "::llvm::raw_ostream&":$os), "",
+      "return ::mlir::OpAsmDialectInterface::AliasResult::NoAlias;"
+    >,
   ];
 }
 

diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index cc578eae3ee36..1f22d4f37a813 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -1163,14 +1163,13 @@ void AliasInitializer::generateAlias(T symbol, InProgressAliasInfo &alias,
 
   OpAsmDialectInterface::AliasResult symbolInterfaceResult =
       OpAsmDialectInterface::AliasResult::NoAlias;
-  if constexpr (std::is_base_of_v<Attribute, T>) {
-    if (auto symbolInterface = dyn_cast<OpAsmAttrInterface>(symbol)) {
-      symbolInterfaceResult = symbolInterface.getAlias(aliasOS);
-      if (symbolInterfaceResult !=
-          OpAsmDialectInterface::AliasResult::NoAlias) {
-        nameBuffer = std::move(aliasBuffer);
-        assert(!nameBuffer.empty() && "expected valid alias name");
-      }
+  using InterfaceT = std::conditional_t<std::is_base_of_v<Attribute, T>,
+                                        OpAsmAttrInterface, OpAsmTypeInterface>;
+  if (auto symbolInterface = dyn_cast<InterfaceT>(symbol)) {
+    symbolInterfaceResult = symbolInterface.getAlias(aliasOS);
+    if (symbolInterfaceResult != OpAsmDialectInterface::AliasResult::NoAlias) {
+      nameBuffer = std::move(aliasBuffer);
+      assert(!nameBuffer.empty() && "expected valid alias name");
     }
   }
 

diff  --git a/mlir/test/IR/op-asm-interface.mlir b/mlir/test/IR/op-asm-interface.mlir
index 44a6e7afece03..086dc7da421c2 100644
--- a/mlir/test/IR/op-asm-interface.mlir
+++ b/mlir/test/IR/op-asm-interface.mlir
@@ -61,6 +61,16 @@ func.func @block_argument_name_from_op_asm_type_interface_asmprinter() {
 
 // -----
 
+// CHECK: !op_asm_type_interface_type =
+!type = !test.op_asm_type_interface
+
+func.func @alias_from_op_asm_type_interface() {
+  %0 = "test.result_name_from_type"() : () -> !type
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // Test OpAsmAttrInterface
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index 6335701786ecc..c048f8b654ec2 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -399,7 +399,7 @@ def TestTypeVerification : Test_Type<"TestTypeVerification"> {
 }
 
 def TestTypeOpAsmTypeInterface : Test_Type<"TestTypeOpAsmTypeInterface",
-    [DeclareTypeInterfaceMethods<OpAsmTypeInterface, ["getAsmName"]>]> {
+    [DeclareTypeInterfaceMethods<OpAsmTypeInterface, ["getAsmName", "getAlias"]>]> {
   let mnemonic = "op_asm_type_interface";
 }
 

diff  --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index 1ae7ac472d989..0c237440834ef 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -537,3 +537,9 @@ void TestTypeOpAsmTypeInterfaceType::getAsmName(
     OpAsmSetNameFn setNameFn) const {
   setNameFn("op_asm_type_interface");
 }
+
+::mlir::OpAsmDialectInterface::AliasResult
+TestTypeOpAsmTypeInterfaceType::getAlias(::llvm::raw_ostream &os) const {
+  os << "op_asm_type_interface_type";
+  return ::mlir::OpAsmDialectInterface::AliasResult::FinalAlias;
+}