From bbc412af8f4befbb1a0f7fe31f71c44984f2e5a7 Mon Sep 17 00:00:00 2001 From: Jonas Rickert Date: Thu, 5 Mar 2026 08:37:39 -0700 Subject: [PATCH 1/2] [MLIR] Remove restriction that type/attribute alias names cannot end with a digit. To prevent collisions when multiple aliases are registered with the same name, keep track of all used aliases and skip used ones when generating conflict-resolving suffixes. This does not affect the reading of .mlir files, but will affect the printing of aliases that end in a digit, as they won't have an _ appended anymore. Signed-off-by: Jonas Rickert --- mlir/docs/DefiningDialects/Assembly.md | 3 +- mlir/include/mlir/IR/OpImplementation.h | 6 ++-- mlir/lib/IR/AsmPrinter.cpp | 20 +++++++++++-- mlir/test/IR/print-attr-type-aliases.mlir | 28 ++++++++++++++----- mlir/test/IR/recursive-type.mlir | 6 ++-- .../Dialect/Test/TestDialectInterfaces.cpp | 11 ++++++++ 6 files changed, 59 insertions(+), 15 deletions(-) diff --git a/mlir/docs/DefiningDialects/Assembly.md b/mlir/docs/DefiningDialects/Assembly.md index aa6272e87358..23301e937d9f 100644 --- a/mlir/docs/DefiningDialects/Assembly.md +++ b/mlir/docs/DefiningDialects/Assembly.md @@ -47,8 +47,7 @@ void MyDialect::initialize() { } ``` -* If `getAlias` provides an alias with a trailing digit, `AsmPrinter` appends an underscore to avoid conflicts with autogenerated IDs. -* If multiple types/attributes have the same alias from `getAlias`, a number is appended to the alias to avoid conflicts. +* If multiple types/attributes have the same alias from `getAlias`, a numeric suffix is appended to the alias to disambiguate. The suffix assignment automatically avoids collisions with other registered alias names. ## Suggesting SSA/Block Names diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h index d9c925a9c56e..c32c104ea85f 100644 --- a/mlir/include/mlir/IR/OpImplementation.h +++ b/mlir/include/mlir/IR/OpImplementation.h @@ -4,6 +4,9 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // +// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates +// //===----------------------------------------------------------------------===// // // This classes used by the implementation details of Op types. @@ -1762,8 +1765,7 @@ class OpAsmDialectInterface /// 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]+). + /// the symbol when printing textual IR. These aliases must not contain `.`. virtual AliasResult getAlias(Attribute attr, raw_ostream &os) const { return AliasResult::NoAlias; } diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp index 0f2364bc2170..8318f077e9b2 100644 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -4,6 +4,9 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // +// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates +// //===----------------------------------------------------------------------===// // // This file implements the MLIR AsmPrinter class, which is used to implement @@ -1109,12 +1112,26 @@ void AliasInitializer::initializeAliases( return lhs.second < rhs.second; }); + llvm::StringSet<> usedNames; + for (auto &[symbol, aliasInfo] : unprocessedAliases) + if (aliasInfo.alias) + usedNames.insert(*aliasInfo.alias); + + const auto tryClaimName = [&usedNames](StringRef alias, + unsigned index) -> bool { + return usedNames.insert((alias + Twine(index)).str()).second; + }; + llvm::StringMap nameCounts; for (auto &[symbol, aliasInfo] : unprocessedAliases) { if (!aliasInfo.alias) continue; StringRef alias = *aliasInfo.alias; unsigned nameIndex = nameCounts[alias]++; + + while (nameIndex > 0 && !tryClaimName(alias, nameIndex)) + nameIndex = nameCounts[alias]++; + symbolToAlias.insert( {symbol, SymbolAlias(alias, nameIndex, aliasInfo.isType, aliasInfo.canBeDeferred)}); @@ -1206,8 +1223,7 @@ void AliasInitializer::generateAlias(T symbol, InProgressAliasInfo &alias, SmallString<16> tempBuffer; StringRef name = - sanitizeIdentifier(nameBuffer, tempBuffer, /*allowedPunctChars=*/"$_-", - /*allowTrailingDigit=*/false); + sanitizeIdentifier(nameBuffer, tempBuffer, /*allowedPunctChars=*/"$_-"); name = name.copy(aliasAllocator); alias = InProgressAliasInfo(name); } diff --git a/mlir/test/IR/print-attr-type-aliases.mlir b/mlir/test/IR/print-attr-type-aliases.mlir index e878d862076c..cbf7159c9471 100644 --- a/mlir/test/IR/print-attr-type-aliases.mlir +++ b/mlir/test/IR/print-attr-type-aliases.mlir @@ -1,3 +1,5 @@ +// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates // RUN: mlir-opt %s -split-input-file -mlir-print-debuginfo | FileCheck %s // Verify printer of type & attr aliases. // RUN: mlir-opt %s -split-input-file -mlir-print-debuginfo | mlir-opt -split-input-file -mlir-print-debuginfo | FileCheck %s @@ -5,7 +7,7 @@ // CHECK-DAG: #test2Ealias = "alias_test:dot_in_name" "test.op"() {alias_test = "alias_test:dot_in_name"} : () -> () -// CHECK-DAG: #test_alias0_ = "alias_test:trailing_digit" +// CHECK-DAG: #test_alias0 = "alias_test:trailing_digit" "test.op"() {alias_test = "alias_test:trailing_digit"} : () -> () // CHECK-DAG: #_0_test_alias = "alias_test:prefixed_digit" @@ -14,10 +16,22 @@ // CHECK-DAG: #_25test = "alias_test:prefixed_symbol" "test.op"() {alias_test = "alias_test:prefixed_symbol"} : () -> () -// CHECK-DAG: #test_alias_conflict0_ = "alias_test:sanitize_conflict_a" -// CHECK-DAG: #test_alias_conflict0_1 = "alias_test:sanitize_conflict_b" +// CHECK-DAG: #test_alias_conflict0 = "alias_test:sanitize_conflict_a" +// CHECK-DAG: #test_alias_conflict0_ = "alias_test:sanitize_conflict_b" "test.op"() {alias_test = ["alias_test:sanitize_conflict_a", "alias_test:sanitize_conflict_b"]} : () -> () +// CHECK-DAG: #collide = "alias_test:suffix_collision_a" +// CHECK-DAG: #collide2 = "alias_test:suffix_collision_b" +// CHECK-DAG: #collide1 = "alias_test:suffix_collision_c" +"test.op"() {alias_test = ["alias_test:suffix_collision_a", "alias_test:suffix_collision_b", "alias_test:suffix_collision_c"]} : () -> () + +// CHECK-DAG: #cross = "alias_test:cross_collision_a" +// CHECK-DAG: #cross2 = "alias_test:cross_collision_b" +// CHECK-DAG: #cross3 = "alias_test:cross_collision_c" +// CHECK-DAG: #cross1 = "alias_test:cross_collision_d" +// CHECK-DAG: #cross11 = "alias_test:cross_collision_e" +"test.op"() {alias_test = ["alias_test:cross_collision_a", "alias_test:cross_collision_b", "alias_test:cross_collision_c", "alias_test:cross_collision_d", "alias_test:cross_collision_e"]} : () -> () + // CHECK-DAG: !tuple = tuple "test.op"() {alias_test = "alias_test:large_tuple"} : () -> (tuple) @@ -28,8 +42,8 @@ // CHECK-DAG: tensor<32xf32, #test_encoding> "test.op"() : () -> tensor<32xf32, "alias_test:tensor_encoding"> -// CHECK-DAG: !test_ui8_ = !test.int -// CHECK-DAG: tensor<32x!test_ui8_> +// CHECK-DAG: !test_ui8 = !test.int +// CHECK-DAG: tensor<32x!test_ui8> "test.op"() : () -> tensor<32x!test.int> // CHECK-DAG: #[[LOC_NESTED:.+]] = loc("nested") @@ -47,8 +61,8 @@ // ----- // Ensure self type parameters get considered for aliases. -// CHECK: !test_ui8_ = !test.int -// CHECK: #test.attr_with_self_type_param : !test_ui8_ +// CHECK: !test_ui8 = !test.int +// CHECK: #test.attr_with_self_type_param : !test_ui8 "test.op"() {alias_test = #test.attr_with_self_type_param : !test.int } : () -> () // ----- diff --git a/mlir/test/IR/recursive-type.mlir b/mlir/test/IR/recursive-type.mlir index 42aecb41d998..1f24a5e9d619 100644 --- a/mlir/test/IR/recursive-type.mlir +++ b/mlir/test/IR/recursive-type.mlir @@ -1,11 +1,13 @@ +// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates // RUN: mlir-opt %s -test-recursive-types | FileCheck %s // CHECK: !testrec = !test.test_rec> // CHECK: ![[$NAME:.*]] = !test.test_rec_alias> // CHECK: ![[$NAME5:.*]] = !test.test_rec_alias>>> // CHECK: ![[$NAME2:.*]] = !test.test_rec_alias, i32>> -// CHECK: ![[$NAME4:.*]] = !test.test_rec_alias -// CHECK: ![[$NAME3:.*]] = !test.test_rec_alias +// CHECK: ![[$NAME4:.*]] = !test.test_rec_alias +// CHECK: ![[$NAME3:.*]] = !test.test_rec_alias // CHECK-LABEL: @roundtrip func.func @roundtrip() { diff --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp index 64add8cef369..18c169a979a3 100644 --- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp @@ -4,6 +4,9 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // +// Modifications (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates +// //===----------------------------------------------------------------------===// #include "TestDialect.h" @@ -194,6 +197,14 @@ struct TestOpAsmInterface : public OpAsmDialectInterface { .Case("alias_test:sanitize_conflict_b", StringRef("test_alias_conflict0_")) .Case("alias_test:tensor_encoding", StringRef("test_encoding")) + .Case("alias_test:suffix_collision_a", StringRef("collide")) + .Case("alias_test:suffix_collision_b", StringRef("collide")) + .Case("alias_test:suffix_collision_c", StringRef("collide1")) + .Case("alias_test:cross_collision_a", StringRef("cross")) + .Case("alias_test:cross_collision_b", StringRef("cross")) + .Case("alias_test:cross_collision_c", StringRef("cross")) + .Case("alias_test:cross_collision_d", StringRef("cross1")) + .Case("alias_test:cross_collision_e", StringRef("cross1")) .Default(std::nullopt); if (!aliasName) return AliasResult::NoAlias; From f3caaa509a90b0d50add78cc13f498baa983bb88 Mon Sep 17 00:00:00 2001 From: Jonas Rickert Date: Thu, 26 Feb 2026 10:08:10 -0700 Subject: [PATCH 2/2] [MLIR][Quant] Add alias type printing for BlockFloat Signed-off-by: Jonas Rickert --- .../mlir/Dialect/Quant/IR/QuantTypes.h | 4 +- mlir/lib/Dialect/Quant/IR/QuantOps.cpp | 5 ++ mlir/lib/Dialect/Quant/IR/QuantTypes.cpp | 12 +++-- mlir/lib/Dialect/Quant/IR/TypeParser.cpp | 29 +++++++++++- .../Quant/print-block-float-aliases.mlir | 46 +++++++++++++++++++ 5 files changed, 91 insertions(+), 5 deletions(-) create mode 100644 mlir/test/Dialect/Quant/print-block-float-aliases.mlir diff --git a/mlir/include/mlir/Dialect/Quant/IR/QuantTypes.h b/mlir/include/mlir/Dialect/Quant/IR/QuantTypes.h index 9733efb6d83f..2342dc3be520 100644 --- a/mlir/include/mlir/Dialect/Quant/IR/QuantTypes.h +++ b/mlir/include/mlir/Dialect/Quant/IR/QuantTypes.h @@ -3,7 +3,7 @@ // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its +// Modifications (c) Copyright 2025-2026 Advanced Micro Devices, Inc. or its // affiliates // //===----------------------------------------------------------------------===// @@ -263,6 +263,8 @@ class BlockFloatQuantizedType static std::optional parseBlockMode(StringRef name); static StringRef getBlockModeName(BlockMode blockMode); + /// An empty alias means no alias should be emitted. + static StringRef getBlockModeAlias(BlockMode blockMode); static BlockFloatQuantizedType get(MLIRContext *ctx, BlockMode blockMode, int32_t axis); diff --git a/mlir/lib/Dialect/Quant/IR/QuantOps.cpp b/mlir/lib/Dialect/Quant/IR/QuantOps.cpp index 1237c869fef9..3ac7419a0794 100644 --- a/mlir/lib/Dialect/Quant/IR/QuantOps.cpp +++ b/mlir/lib/Dialect/Quant/IR/QuantOps.cpp @@ -23,6 +23,10 @@ namespace mlir { namespace quant { +namespace detail { +void addAsmInterface(QuantDialect *dialect); +} // namespace detail + namespace { // Verify the integrity of per-axis quantization information, if present. @@ -245,6 +249,7 @@ void QuantDialect::initialize() { >(); detail::addBytecodeInterface(this); addInterfaces(); + detail::addAsmInterface(this); } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp b/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp index 11dc44825d81..57611e6be2cf 100644 --- a/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp +++ b/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp @@ -3,7 +3,7 @@ // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its +// Modifications (c) Copyright 2025-2026 Advanced Micro Devices, Inc. or its // affiliates // //===----------------------------------------------------------------------===// @@ -38,6 +38,8 @@ double getMaxScale(Type expressedType) { struct BlockFloatQuantizedTypeConfig { StringRef mode; + // If empty, no type alias should be emitted for this mode. + StringRef alias; unsigned singleElementBitWidth; unsigned averageBitsPerElement; unsigned blockSize; @@ -48,12 +50,12 @@ getBlockFloatQuantizedTypeConfig(BlockFloatQuantizedType::BlockMode blockMode) { switch (blockMode) { case BlockFloatQuantizedType::BlockMode::BFP16: { static constexpr BlockFloatQuantizedTypeConfig config = { - StringLiteral("BFP16"), 16, 9, 8}; + StringLiteral("BFP16"), StringLiteral("bfp16"), 16, 9, 8}; return config; } case BlockFloatQuantizedType::BlockMode::MX6: { static constexpr BlockFloatQuantizedTypeConfig config = { - StringLiteral("MX6"), 13, 6, 16}; + StringLiteral("MX6"), StringLiteral("mx6"), 13, 6, 16}; return config; } } @@ -463,6 +465,10 @@ StringRef BlockFloatQuantizedType::getBlockModeName(BlockMode blockMode) { return getBlockFloatQuantizedTypeConfig(blockMode).mode; } +StringRef BlockFloatQuantizedType::getBlockModeAlias(BlockMode blockMode) { + return getBlockFloatQuantizedTypeConfig(blockMode).alias; +} + UniformQuantizedPerAxisType UniformQuantizedPerAxisType::get( unsigned flags, Type storageType, Type expressedType, ArrayRef scales, ArrayRef zeroPoints, diff --git a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp index d81cfc226724..0c8d84ba10af 100644 --- a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp +++ b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp @@ -3,7 +3,7 @@ // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its +// Modifications (c) Copyright 2025-2026 Advanced Micro Devices, Inc. or its // affiliates // //===----------------------------------------------------------------------===// @@ -13,6 +13,7 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/DialectImplementation.h" #include "mlir/IR/Location.h" +#include "mlir/IR/OpImplementation.h" #include "mlir/IR/Types.h" #include "llvm/ADT/APFloat.h" #include "llvm/Support/Format.h" @@ -708,3 +709,29 @@ void QuantDialect::printType(Type type, DialectAsmPrinter &os) const { else llvm_unreachable("Unhandled quantized type"); } + +namespace { +struct QuantOpAsmDialectInterface : public OpAsmDialectInterface { + using OpAsmDialectInterface::OpAsmDialectInterface; + + AliasResult getAlias(Type type, raw_ostream &os) const final { + const auto blockType = dyn_cast(type); + if (!blockType) + return AliasResult::NoAlias; + if (blockType.getAxis() != 1) + return AliasResult::NoAlias; + const StringRef alias = + BlockFloatQuantizedType::getBlockModeAlias(blockType.getBlockMode()); + if (alias.empty()) + return AliasResult::NoAlias; + os << alias; + return AliasResult::OverridableAlias; + } +}; +} // namespace + +namespace mlir::quant::detail { +void addAsmInterface(QuantDialect *dialect) { + dialect->addInterfaces(); +} +} // namespace mlir::quant::detail diff --git a/mlir/test/Dialect/Quant/print-block-float-aliases.mlir b/mlir/test/Dialect/Quant/print-block-float-aliases.mlir new file mode 100644 index 000000000000..3cd6dc685fc4 --- /dev/null +++ b/mlir/test/Dialect/Quant/print-block-float-aliases.mlir @@ -0,0 +1,46 @@ +// (c) Copyright 2026 Advanced Micro Devices, Inc. or its +// affiliates + +// RUN: mlir-opt %s -split-input-file | FileCheck %s + +// CHECK: !mx6 = !quant.block_float +// CHECK-LABEL: func.func @alias_positive( +// CHECK: tensor<4x!mx6> +// CHECK: return %arg0 : tensor<4x!mx6> +func.func @alias_positive(%arg0: tensor<4x!quant.block_float>) + -> tensor<4x!quant.block_float> { + return %arg0 : tensor<4x!quant.block_float> +} + +// ----- + +// CHECK-LABEL: func.func @no_alias_mx6_other_axis( +// CHECK-NOT: !mx6 +// CHECK: tensor<4x!quant.block_float> +func.func @no_alias_mx6_other_axis( + %arg0: tensor<4x!quant.block_float>) + -> tensor<4x!quant.block_float> { + return %arg0 : tensor<4x!quant.block_float> +} + +// ----- + +// CHECK: !bfp16 = !quant.block_float +// CHECK-LABEL: func.func @alias_bfp16_axis1( +// CHECK: tensor<4x!bfp16> +func.func @alias_bfp16_axis1( + %arg0: tensor<4x!quant.block_float>) + -> tensor<4x!quant.block_float> { + return %arg0 : tensor<4x!quant.block_float> +} + +// ----- + +// CHECK-LABEL: func.func @no_alias_bfp16_other_axis( +// CHECK-NOT: !bfp16 +// CHECK: tensor<4x!quant.block_float> +func.func @no_alias_bfp16_other_axis( + %arg0: tensor<4x!quant.block_float>) + -> tensor<4x!quant.block_float> { + return %arg0 : tensor<4x!quant.block_float> +}