[mlir] Allow to override type/attr aliases from various hooks

[vinograd47]

Summary

Use new return type for OpAsmDialectInterface::getAlias:

• AliasResult::NoAlias if an alias was not provided.
• AliasResult::OverridableAlias if an alias was provided, but it might be overriden by other hook.
• AliasResult::FinalAlias if an alias was provided and it should be used (no other hooks will be checked).

In that case AsmPrinter will use either the first alias with FinalAlias result or
the last alias with OverridableAlias result (it depends on dialect array order).

Used OverridableAlias result for BuiltinOpAsmDialectInterface.

Use case: provide more informative alias for built-in attributes like AffineMapAttr
instead of generic "map<N>".

[rriddle]

If this is just to avoid the generic alias in the "parent" dialect (in this case, the builtin dialect), can we just flip the iteration order of the dialect interfaces when building the alias set? I'm not sure when Optional would ever really be used outside of that use case.

[vinograd47]

In D107437#2925696, @rriddle wrote:

If this is just to avoid the generic alias in the "parent" dialect (in this case, the builtin dialect), can we just flip the iteration order of the dialect interfaces when building the alias set? I'm not sure when Optional would ever really be used outside of that use case.

The problem is, that the order of dialect interfaces depends on the order of dialects stored in MLIRContext object. The MLIRContext stores dialects in DenseMap<StringRef, std::unique_ptr<Dialect>> with dialect name as key. And it seems that the issue we faced was caused by adding explicit name (builtin) to BuiltinDialect (originally it has an empty name). I didn't find other way to somehow specify the order of alias checks.

Optional might be used for the cases, when some hook provides just general name per object type (like map or set) without extra information from the object properties (like special names for specific affine maps). For example, in our internal project we provide such Optional alias for quantized types from quant dialect, but allows to override this alias with more specific if needed.

[rriddle]

In D107437#2927666, @vinograd47 wrote:

In D107437#2925696, @rriddle wrote:

If this is just to avoid the generic alias in the "parent" dialect (in this case, the builtin dialect), can we just flip the iteration order of the dialect interfaces when building the alias set? I'm not sure when Optional would ever really be used outside of that use case.

The problem is, that the order of dialect interfaces depends on the order of dialects stored in MLIRContext object. The MLIRContext stores dialects in DenseMap<StringRef, std::unique_ptr<Dialect>> with dialect name as key. And it seems that the issue we faced was caused by adding explicit name (builtin) to BuiltinDialect (originally it has an empty name). I didn't find other way to somehow specify the order of alias checks.

Optional might be used for the cases, when some hook provides just general name per object type (like map or set) without extra information from the object properties (like special names for specific affine maps). For example, in our internal project we provide such Optional alias for quantized types from quant dialect, but allows to override this alias with more specific if needed.

The second part makes sense.

mlir/include/mlir/IR/OpImplementation.h

931

What about OverridableAlias instead of OptionalAlias?

And, what about NoAlias as opposed to NotSupported? (It follows the tone of the others results more)

Also, can you document this enum?

mlir/lib/IR/AsmPrinter.cpp

657–659

nit: Spell out auto here.

658

nit: Drop the newline here.

663

nit: Drop the newline here.

mlir/test/lib/Dialect/Test/TestDialect.cpp

83

Can you use a test type instead of index? Less likely to cause a problem in other tests.

[vinograd47]

Apply review remarks

[vinograd47]

@rriddle I've applied your suggestions.

[rriddle]

LG, thanks!

mlir/include/mlir/IR/OpImplementation.h

 
 /// A functor used to set the name of the start of a result group of an
 /// operation. See 'getAsmResultNames' below for more details.
 using OpAsmSetValueNameFn = function_ref<void(Value, StringRef)>;
 
 class OpAsmDialectInterface
     : public DialectInterface::Base<OpAsmDialectInterface> {
 public:
+  /// Holds the result of `getAlias` hook call.
+  enum class AliasResult {
+    /// The object (type or attribute) is not supported by the hook
+    /// and an alias was not provided.
+    NoAlias,
+    /// An alias was provided, but it might be overriden by other hook.
+    OverridableAlias,
+    /// An alias was provided and it should be used
+    /// (no other hooks will be checked).
+    FinalAlias
+  };
+
   OpAsmDialectInterface(Dialect *dialect) : Base(dialect) {}
 
   /// Hooks for getting an alias identifier alias for a given symbol, that is
   /// not necessarily a part of this dialect. The identifier is used in place of
   /// the symbol when printing textual IR. These aliases must not contain `.` or
-  /// end with a numeric digit([0-9]+). Returns success if an alias was
-  /// provided, failure otherwise.
-  virtual LogicalResult getAlias(Attribute attr, raw_ostream &os) const {
-    return failure();
+  /// end with a numeric digit([0-9]+).
+  virtual AliasResult getAlias(Attribute attr, raw_ostream &os) const {
+    return AliasResult::NoAlias;
   }
-  virtual LogicalResult getAlias(Type type, raw_ostream &os) const {
-    return failure();
+  virtual AliasResult getAlias(Type type, raw_ostream &os) const {
+    return AliasResult::NoAlias;
   }
 
   /// Get a special name to use when printing the given operation. See
   /// OpAsmInterface.td#getAsmResultNames for usage details and documentation.
   virtual void getAsmResultNames(Operation *op,
                                  OpAsmSetValueNameFn setNameFn) const {}
 
   /// Get a special name to use when printing the entry block arguments of the

mlir/lib/IR/AsmPrinter.cpp

if (auto subElementInterface = type.dyn_cast<SubElementTypeInterface>()) {
     subElementInterface.walkSubElements([&](Attribute attr) { visit(attr); },
                                         [&](Type type) { visit(type); });
   }
 }
 
 template <typename T>
 LogicalResult AliasInitializer::generateAlias(
     T symbol, llvm::MapVector<StringRef, std::vector<T>> &aliasToSymbol) {
-  SmallString<16> tempBuffer;
+  SmallString<32> nameBuffer;
   for (const auto &interface : interfaces) {
-    if (failed(interface.getAlias(symbol, aliasOS)))
+    OpAsmDialectInterface::AliasResult result =
+        interface.getAlias(symbol, aliasOS);
+    if (result == OpAsmDialectInterface::AliasResult::NoAlias)
       continue;
-    StringRef name = aliasOS.str();
-    assert(!name.empty() && "expected valid alias name");
-    name = sanitizeIdentifier(name, tempBuffer, /*allowedPunctChars=*/"$_-",
-                              /*allowTrailingDigit=*/false);
-    name = name.copy(aliasAllocator);
-
-    aliasToSymbol[name].push_back(symbol);
-    aliasBuffer.clear();
-    return success();
-  }
-  return failure();
-}
+    nameBuffer = std::move(aliasBuffer);
+    assert(!nameBuffer.empty() && "expected valid alias name");
+    if (result == OpAsmDialectInterface::AliasResult::FinalAlias)
+      break;
+  }
+
+  if (nameBuffer.empty())
+    return failure();
+
+  SmallString<16> tempBuffer;
+  StringRef name =
+      sanitizeIdentifier(nameBuffer, tempBuffer, /*allowedPunctChars=*/"$_-",
 
 //===----------------------------------------------------------------------===//
 // AliasState
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class manages the state for type and attribute aliases.
 class AliasState {

mlir/lib/IR/BuiltinDialect.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinDialect.cpp.inc"
 
 namespace {
 struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface {
   using OpAsmDialectInterface::OpAsmDialectInterface;
 
-  LogicalResult getAlias(Attribute attr, raw_ostream &os) const override {
+  AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     if (attr.isa<AffineMapAttr>()) {
       os << "map";
-      return success();
+      return AliasResult::OverridableAlias;
     }
     if (attr.isa<IntegerSetAttr>()) {
       os << "set";
-      return success();
+      return AliasResult::OverridableAlias;
     }
     if (attr.isa<LocationAttr>()) {
       os << "loc";
-      return success();
+      return AliasResult::OverridableAlias;
     }
-    return failure();
+    return AliasResult::NoAlias;
   }
 
-  LogicalResult getAlias(Type type, raw_ostream &os) const final {
+  AliasResult getAlias(Type type, raw_ostream &os) const final {
     if (auto tupleType = type.dyn_cast<TupleType>()) {
       if (tupleType.size() > 16) {
         os << "tuple";
-        return success();
+        return AliasResult::OverridableAlias;
       }
     }
-    return failure();
+    return AliasResult::NoAlias;
   }
 };
 } // end anonymous namespace.
 
 void BuiltinDialect::initialize() {
   registerTypes();
   registerAttributes();
   registerLocationAttributes();

mlir/test/IR/print-attr-type-aliases.mlir

 
 // CHECK-DAG: #test_alias_conflict0_0 = "alias_test:sanitize_conflict_a"
 // CHECK-DAG: #test_alias_conflict0_1 = "alias_test:sanitize_conflict_b"
 "test.op"() {alias_test = ["alias_test:sanitize_conflict_a", "alias_test:sanitize_conflict_b"]} : () -> ()
 
 // CHECK-DAG: !tuple = type tuple<i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32>
 "test.op"() {alias_test = "alias_test:large_tuple"} : () -> (tuple<i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32>)
 
+// CHECK-DAG: !test_tuple = type tuple<!test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla>
+"test.op"() {alias_test = "alias_test:large_tuple"} : () -> (tuple<!test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla, !test.smpla>)
+
 // CHECK-DAG: #test_encoding = "alias_test:tensor_encoding"
 // CHECK-DAG: tensor<32xf32, #test_encoding>
 "test.op"() : () -> tensor<32xf32, "alias_test:tensor_encoding">

mlir/test/lib/Dialect/Test/TestDialect.cpp

static_assert(OpTrait::hasSingleBlockImplicitTerminator<
                   SingleBlockImplicitTerminatorOp>::value,
               "hasSingleBlockImplicitTerminator does not match "
               "SingleBlockImplicitTerminatorOp");
 
 // Test support for interacting with the AsmPrinter.
 struct TestOpAsmInterface : public OpAsmDialectInterface {
   using OpAsmDialectInterface::OpAsmDialectInterface;
 
-  LogicalResult getAlias(Attribute attr, raw_ostream &os) const final {
+  AliasResult getAlias(Attribute attr, raw_ostream &os) const final {
     StringAttr strAttr = attr.dyn_cast<StringAttr>();
     if (!strAttr)
-      return failure();
+      return AliasResult::NoAlias;
 
     // Check the contents of the string attribute to see what the test alias
     // should be named.
     Optional<StringRef> aliasName =
         StringSwitch<Optional<StringRef>>(strAttr.getValue())
             .Case("alias_test:dot_in_name", StringRef("test.alias"))
             .Case("alias_test:trailing_digit", StringRef("test_alias0"))
             .Case("alias_test:prefixed_digit", StringRef("0_test_alias"))
             .Case("alias_test:sanitize_conflict_a",
                   StringRef("test_alias_conflict0"))
             .Case("alias_test:sanitize_conflict_b",
                   StringRef("test_alias_conflict0_"))
             .Case("alias_test:tensor_encoding", StringRef("test_encoding"))
             .Default(llvm::None);
     if (!aliasName)
-      return failure();
+      return AliasResult::NoAlias;
 
     os << *aliasName;
-    return success();
+    return AliasResult::FinalAlias;
+  }
+
+  AliasResult getAlias(Type type, raw_ostream &os) const final {
+    if (auto tupleType = type.dyn_cast<TupleType>()) {
+      if (tupleType.size() > 0 &&
+          llvm::all_of(tupleType.getTypes(), [](Type elemType) {
+            return elemType.isa<SimpleAType>();
+          })) {
+        os << "test_tuple";
+        return AliasResult::FinalAlias;
+      }
+    }
+    return AliasResult::NoAlias;
   }
 
   void getAsmResultNames(Operation *op,
                          OpAsmSetValueNameFn setNameFn) const final {
     if (auto asmOp = dyn_cast<AsmDialectInterfaceOp>(op))
       setNameFn(asmOp, "result");
   }