[mlir] Add getAlias for OpAsmTypeInterface

Author: ZenithalHourlyRate

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;"
+    >,
   ];
 }
 

mlir/lib/IR/AsmPrinter.cpp

@@ -1159,15 +1159,31 @@ template <typename T>
 void AliasInitializer::generateAlias(T symbol, InProgressAliasInfo &alias,
                                      bool canBeDeferred) {
   SmallString<32> nameBuffer;
-  for (const auto &interface : interfaces) {
-    OpAsmDialectInterface::AliasResult result =
-        interface.getAlias(symbol, aliasOS);
-    if (result == OpAsmDialectInterface::AliasResult::NoAlias)
-      continue;
-    nameBuffer = std::move(aliasBuffer);
-    assert(!nameBuffer.empty() && "expected valid alias name");
-    if (result == OpAsmDialectInterface::AliasResult::FinalAlias)
-      break;
+
+  OpAsmDialectInterface::AliasResult symbolInterfaceResult =
+      OpAsmDialectInterface::AliasResult::NoAlias;
+  if constexpr (std::is_base_of_v<Type, T>) {
+    if (auto symbolInterface = mlir::dyn_cast<OpAsmTypeInterface>(symbol)) {
+      symbolInterfaceResult = symbolInterface.getAlias(aliasOS);
+      if (symbolInterfaceResult !=
+          OpAsmDialectInterface::AliasResult::NoAlias) {
+        nameBuffer = std::move(aliasBuffer);
+        assert(!nameBuffer.empty() && "expected valid alias name");
+      }
+    }
+  }
+
+  if (symbolInterfaceResult != OpAsmDialectInterface::AliasResult::FinalAlias) {
+    for (const auto &interface : interfaces) {
+      OpAsmDialectInterface::AliasResult result =
+          interface.getAlias(symbol, aliasOS);
+      if (result == OpAsmDialectInterface::AliasResult::NoAlias)
+        continue;
+      nameBuffer = std::move(aliasBuffer);
+      assert(!nameBuffer.empty() && "expected valid alias name");
+      if (result == OpAsmDialectInterface::AliasResult::FinalAlias)
+        break;
+    }
   }
 
   if (nameBuffer.empty())

mlir/test/IR/op-asm-interface.mlir

@@ -22,3 +22,14 @@ func.func @block_argument_name_from_op_asm_type_interface() {
   }
   return
 }
+
+// -----
+
+// CHECK: !op_asm_type_interface_type =
+!type = !test.op_asm_type_interface
+
+func.func @alias_from_op_asm_type_interface() {
+  // CHECK-LABEL: @alias_from_op_asm_type_interface
+  %0 = "test.result_name_from_type"() : () -> !type
+  return
+}

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";
 }
 

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;
+}