[mlir] Add dialect hooks for registering custom type and attribute alias printers

This patch introduces a mechanism for dialects to register custom alias printers
for types and attributes via the `OpAsmDialectInterface`. This allows dialects
to provide alternative printed representations for types and attributes based
on their TypeID, including types/attributes from other dialects.

The new `registerAttrAliasPrinter` and `registerTypeAliasPrinter` virtual
methods accept callbacks that register printers for specific TypeIDs. When
printing, these custom printers are invoked in registration order, and the
first one to produce output is used.

The precedence for alias resolution is:
1. Explicit type/attribute aliases returned by `getAlias`
2. Dialect-specific alias printers registered via the new hooks
3. Default type/attribute printers

Signed-off-by: Fabian Mora <[email protected]>

Author: fabianmcg

mlir/include/mlir/IR/OpImplementation.h

@@ -17,6 +17,8 @@
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/OpAsmSupport.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/SMLoc.h"
 #include <optional>
@@ -169,6 +171,9 @@ class AsmPrinter {
     if (succeeded(printAlias(attrOrType)))
       return;
 
+    if (succeeded(printDialectAlias(attrOrType, /*printStripped=*/true)))
+      return;
+
     raw_ostream &os = getStream();
     uint64_t posPrior = os.tell();
     attrOrType.print(*this);
@@ -218,6 +223,14 @@ class AsmPrinter {
   /// be printed.
   virtual LogicalResult printAlias(Type type);
 
+  /// Print the alias for the given attribute, return failure if no alias could
+  /// be printed.
+  virtual LogicalResult printDialectAlias(Attribute attr, bool printStripped);
+
+  /// Print the alias for the given type, return failure if no alias could
+  /// be printed.
+  virtual LogicalResult printDialectAlias(Type type, bool printStripped);
+
   /// Print the given string as a keyword, or a quoted and escaped string if it
   /// has any special or non-printable characters in it.
   virtual void printKeywordOrString(StringRef keyword);
@@ -1799,6 +1812,30 @@ class OpAsmDialectInterface
     return AliasResult::NoAlias;
   }
 
+  /// Hooks for registering alias printers for types and attributes. These
+  /// printers are invoked when printing types or attributes of the given
+  /// TypeID. Printers are invoked in the order they are registered, and the
+  /// first one to print an alias is used.
+  /// The precedence of these printers is as follow:
+  /// 1. The type and attribute aliases returned by `getAlias`.
+  /// 2. Dialect-specific alias printers registered here.
+  /// 3. The type and attribute printers.
+  /// The boolean argument to the printer indicates whether the stripped form
+  /// of the type or attribute is being printed.
+  /// NOTE: This mechanism caches the printed object, therefore the printer
+  /// must always produce the same output for the same input.
+  using AttributeAliasPrinter =
+      llvm::function_ref<void(Attribute, AsmPrinter &, bool)>;
+  using InsertAttrAliasPrinter =
+      llvm::function_ref<void(TypeID, AttributeAliasPrinter)>;
+  virtual void registerAttrAliasPrinter(InsertAttrAliasPrinter insertFn) const {
+  }
+  using TypeAliasPrinter = llvm::function_ref<void(Type, AsmPrinter &, bool)>;
+  using InsertTypeAliasPrinter =
+      llvm::function_ref<void(TypeID, TypeAliasPrinter)>;
+  virtual void registerTypeAliasPrinter(InsertTypeAliasPrinter insertFn) const {
+  }
+
   //===--------------------------------------------------------------------===//
   // Resources
   //===--------------------------------------------------------------------===//