Skip to content

[mlir] Enable setting alias on AsmState. #153776

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir] Enable setting alias on AsmState. #153776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions mlir/include/mlir/IR/AsmState.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Copy link
Collaborator

@joker-eph joker-eph Aug 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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);

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

/// Get the printer flags.
const OpPrintingFlags &getPrinterFlags() const;

Expand Down
34 changes: 34 additions & 0 deletions mlir/lib/IR/AsmPrinter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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; }

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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) {
Expand Down Expand Up @@ -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();
}
Expand Down
25 changes: 25 additions & 0 deletions mlir/unittests/IR/TypeTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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();
}