[mlir] Enable setting alias on AsmState.

[jpienaar]

This is useful for application output where a more useful name can be given to a type or where type is rather long and cumbersome.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Jacques Pienaar (jpienaar)

Changes

This is useful for application output where a more useful name can be given to a type or where type is rather long and cumbersome.

---

Full diff: https://github.com/llvm/llvm-project/pull/153776.diff

3 Files Affected:

• (modified) mlir/include/mlir/IR/AsmState.h (+9)
• (modified) mlir/lib/IR/AsmPrinter.cpp (+34)
• (modified) mlir/unittests/IR/TypeTest.cpp (+25)

diff --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h
index 5e9311742bd94..bc25d9e05d73c 100644
--- a/mlir/include/mlir/IR/AsmState.h
+++ b/mlir/include/mlir/IR/AsmState.h
@@ -560,6 +560,15 @@ class AsmState {
            FallbackAsmResourceMap *map = nullptr);
   ~AsmState();
 
+  /// Add an alias for the given attribute. The alias will be used as a
+  /// suggestion when printing. The final alias may be modified to resolve
+  /// conflicts.
+  void addAlias(Attribute attr, StringRef alias);
+
+  /// Add an alias for the given type. The alias will be used as a suggestion
+  /// when printing. The final alias may be modified to resolve conflicts.
+  void addAlias(Type type, StringRef alias);
+
   /// Get the printer flags.
   const OpPrintingFlags &getPrinterFlags() const;
 
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index de52fbd3f215c..775afc02427d5 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -558,6 +558,9 @@ class SymbolAlias {
     }
   }
 
+  /// Returns the alias name.
+  StringRef getAliasName() const { return name; }
+
   /// Returns true if this is a type alias.
   bool isTypeAlias() const { return isType; }
 
@@ -1139,6 +1142,17 @@ void AliasInitializer::initializeAliases(
 void AliasInitializer::initialize(
     Operation *op, const OpPrintingFlags &printerFlags,
     llvm::MapVector<const void *, SymbolAlias> &attrTypeToAlias) {
+  // Pre-populate the alias list with the user-provided aliases.
+  for (auto &it : attrTypeToAlias) {
+    SymbolAlias &symbolAlias = it.second;
+    InProgressAliasInfo info(symbolAlias.getAliasName());
+    info.isType = symbolAlias.isTypeAlias();
+    info.canBeDeferred = symbolAlias.canBeDeferred();
+    info.aliasDepth = 0;
+    aliases.insert({it.first, info});
+  }
+  attrTypeToAlias.clear();
+
   // Use a dummy printer when walking the IR so that we can collect the
   // attributes/types that will actually be used during printing when
   // considering aliases.
@@ -1272,6 +1286,9 @@ class AliasState {
     printAliases(p, newLine, /*isDeferred=*/true);
   }
 
+  /// Add an alias for the given symbol. All aliases are non-deferred.
+  void addAlias(const void *symbol, bool isType, StringRef alias);
+
 private:
   /// Print all of the referenced aliases that support the provided resolution
   /// behavior.
@@ -1286,6 +1303,14 @@ class AliasState {
 };
 } // namespace
 
+void AliasState::addAlias(const void *symbol, bool isType, StringRef alias) {
+  StringRef name = alias.copy(aliasAllocator);
+  // We don't know if this alias is unique. Set suffix index to 0 and defer
+  // conflict resolution to aliases initialization.
+  attrTypeToAlias.insert(
+      {symbol, SymbolAlias(name, 0, isType, /*isDeferrable=*/false)});
+}
+
 void AliasState::initialize(
     Operation *op, const OpPrintingFlags &printerFlags,
     DialectInterfaceCollection<OpAsmDialectInterface> &interfaces) {
@@ -2093,6 +2118,15 @@ AsmState::AsmState(MLIRContext *ctx, const OpPrintingFlags &printerFlags,
 }
 AsmState::~AsmState() = default;
 
+void AsmState::addAlias(Attribute attr, StringRef alias) {
+  impl->getAliasState().addAlias(attr.getAsOpaquePointer(), /*isType=*/false,
+                                 alias);
+}
+void AsmState::addAlias(Type type, StringRef alias) {
+  impl->getAliasState().addAlias(type.getAsOpaquePointer(), /*isType=*/true,
+                                 alias);
+}
+
 const OpPrintingFlags &AsmState::getPrinterFlags() const {
   return impl->getPrinterFlags();
 }
diff --git a/mlir/unittests/IR/TypeTest.cpp b/mlir/unittests/IR/TypeTest.cpp
index 30f6642a9ca71..4222c18f30742 100644
--- a/mlir/unittests/IR/TypeTest.cpp
+++ b/mlir/unittests/IR/TypeTest.cpp
@@ -6,14 +6,19 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 
+using testing::HasSubstr;
+
 /// Mock implementations of a Type hierarchy
 struct LeafType;
 
@@ -69,3 +74,23 @@ TEST(Type, Casting) {
 
   EXPECT_EQ(8u, cast<IntegerType>(intTy).getWidth());
 }
+
+TEST(Type, UserAlias) {
+  MLIRContext ctx;
+  ctx.allowUnregisteredDialects();
+
+  Type intTy = IntegerType::get(&ctx, 8);
+  AsmState state(&ctx);
+  state.addAlias(intTy, "test.alias");
+  Operation *op = Operation::create(
+      UnknownLoc::get(&ctx), OperationName("test.op", &ctx), TypeRange(intTy),
+      ValueRange(), NamedAttrList(), OpaqueProperties(nullptr), BlockRange(),
+      /*numRegions=*/0);
+
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  op->print(os, state);
+  EXPECT_THAT(str, HasSubstr("!test.alias = i8"));
+  EXPECT_THAT(str, HasSubstr("\"test.op\"() : () -> !test.alias\n"));
+  op->erase();
+}

[j2kun]

I don't see any problems here, but I'm curious about the use case. We already have the ability to set asm names for types, which are used as the alias (I'm on mobile but I can find the link later if needed). So when would someone call this new addAlias method? Is something else populating the input map attrTypeToAlias to initialize?

[joker-eph]

How does it play with the existing interfaces that can already alter this? What is the order of priority?

[jpienaar]

I wasn't sure if we want a contract on ordering vs it being a hint, so left it unspecified. Currently this will override - but it is in order of which alias was first specified (e.g, two addAlias calls for same Attribute will result in the first being used). Which corresponds to the user/application being able to dictate alias first and then goes down the regular alias path.

[joker-eph]

Another thing is what is the behavior of repeated calls to the API with the same attribute and different alias names?

[jpienaar]

I mention that above, first one is set, successive calls ignored. Or do you mean I should add to docstring?

[joker-eph]

I'm also curious about the use-case: you can't know all the possible attributes that will be created by the compiler ahead of time, making this feature quite limited in practice.
Usually one would want to customize a specific kind of attribute, not a particular instance somehow.

[jpienaar]

The application here in fact does. E.g., the type itself & dialect doesn't know the use, but the application does. So, as an upstream example, you have a tensor of size of quantized type with loads of params. The dialect and generic hooks can't do much as this is just regular tensor from builtin dialect. But in the application, it knows the user considers it the "jtensor" (to pick a name) that is used everywhere. The builtin dialect assigns no special meaning to it, its only in the application/user context where it has one. So enables giving concise & application relevant aliases even where no generic ones apply. And here it is often the mapping from entry function signature that captures this mapping, so it is easy to populate mapping.

Its not something generic for the type that is, so a type attribute interface would not match. That being said, I did consider creating a aliasinghelper dialect with interface and injecting, but this would then require mutating map before printing (in a thread safe manner) and I don't think one can override that of another dialect like this (could be wrong, this may only be true for those with final alias set).

[joker-eph]

I can see where this is coming from. One annoying thing I can already see is that none of the tooling would round-trip this, there is no way to configure a tool properly for this, and even when you have an application that wants to set this kind of things. The fact that it is an API on the AsmState makes it so that it wouldn't compose with a lot of things (if you interact with other components, the AsmState is not something in general part of the API boundaries)

[jpienaar]

Correct. Now I can make it so :) (which would remove it from API here). And that is by way of adding something like a alias attribute to the top level op - which is just a dictionary attribute of name & attribute. So the printer can query if attribute on op being printed and then populate as here. This would most likely be a discardable attribute though and (downside) would fix a naming convention. E.g., if op has certain named attribute ("attribute_type_alias_map" ? something not likely to be in use already and self-documenting) of dictionary attr kind, then use it. Probably nicer, more composable but at cost of convention.