[MLIR][IR] Fix InProgressAliasInfo init for non-alias

[zyx-billy]

When visiting an attr/type that is NoAlias, the created InProgressAliasInfo was not getting its canBeDeferred and isType fields set. Not setting canBeDeferred when it should be true breaks the assumption that all nested elements are also false. This will cause problems when at a later point the attr/type needs to be converted by markAliasNonDeferrable, as recursion will stop when a canBeDeferred=false attr/type is reached, leaving its nested elements not flipped. This causes nested elements to be printed later in the textual IR and cannot be parsed back in.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Billy Zhu (zyx-billy)

Changes

When visiting an attr/type that is NoAlias, the created InProgressAliasInfo was not getting its canBeDeferred and isType fields set. Not setting canBeDeferred when it should be true breaks the assumption that all nested elements are also false. This will cause problems when at a later point the attr/type needs to converted by markAliasNonDeferrable, as recursion will stop when a canBeDeferred=false attr/type is reached, and the nested elements are not flipped. This causes nested elements to be printed later in the textual IR and cannot be loaded back in.

---

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

2 Files Affected:

• (modified) mlir/lib/IR/AsmPrinter.cpp (+5-5)
• (modified) mlir/test/IR/print-attr-type-aliases.mlir (+20-7)

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index c7ed158aabb6e7..32182c083a8a38 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -584,9 +584,8 @@ class AliasInitializer {
   struct InProgressAliasInfo {
     InProgressAliasInfo()
         : aliasDepth(0), isType(false), canBeDeferred(false) {}
-    InProgressAliasInfo(StringRef alias, bool isType, bool canBeDeferred)
-        : alias(alias), aliasDepth(1), isType(isType),
-          canBeDeferred(canBeDeferred) {}
+    InProgressAliasInfo(StringRef alias)
+        : alias(alias), aliasDepth(1), isType(false), canBeDeferred(false) {}
 
     bool operator<(const InProgressAliasInfo &rhs) const {
       // Order first by depth, then by attr/type kind, and then by name.
@@ -1096,6 +1095,8 @@ std::pair<size_t, size_t> AliasInitializer::visitImpl(
 
   // Try to generate an alias for this value.
   generateAlias(value, it->second, canBeDeferred);
+  it->second.isType = std::is_base_of_v<Type, T>;
+  it->second.canBeDeferred = canBeDeferred;
 
   // Print the value, capturing any nested elements that require aliases.
   SmallVector<size_t> childAliases;
@@ -1153,8 +1154,7 @@ void AliasInitializer::generateAlias(T symbol, InProgressAliasInfo &alias,
       sanitizeIdentifier(nameBuffer, tempBuffer, /*allowedPunctChars=*/"$_-",
                          /*allowTrailingDigit=*/false);
   name = name.copy(aliasAllocator);
-  alias = InProgressAliasInfo(name, /*isType=*/std::is_base_of_v<Type, T>,
-                              canBeDeferred);
+  alias = InProgressAliasInfo(name);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/IR/print-attr-type-aliases.mlir b/mlir/test/IR/print-attr-type-aliases.mlir
index 27c5a75addbb59..e878d862076c90 100644
--- a/mlir/test/IR/print-attr-type-aliases.mlir
+++ b/mlir/test/IR/print-attr-type-aliases.mlir
@@ -1,6 +1,6 @@
-// RUN: mlir-opt %s -split-input-file | FileCheck %s
+// 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-opt -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -mlir-print-debuginfo | mlir-opt -split-input-file -mlir-print-debuginfo | FileCheck %s
 
 // CHECK-DAG: #test2Ealias = "alias_test:dot_in_name"
 "test.op"() {alias_test = "alias_test:dot_in_name"} : () -> ()
@@ -32,16 +32,16 @@
 // CHECK-DAG: tensor<32x!test_ui8_>
 "test.op"() : () -> tensor<32x!test.int<unsigned, 8>>
 
-// CHECK-DAG: #loc = loc("nested")
-// CHECK-DAG: #loc1 = loc("test.mlir":10:8)
-// CHECK-DAG: #loc2 = loc(fused<#loc>[#loc1])
+// CHECK-DAG: #[[LOC_NESTED:.+]] = loc("nested")
+// CHECK-DAG: #[[LOC_RAW:.+]] = loc("test.mlir":10:8)
+// CHECK-DAG: = loc(fused<#[[LOC_NESTED]]>[#[[LOC_RAW]]])
 "test.op"() {alias_test = loc(fused<loc("nested")>["test.mlir":10:8])} : () -> ()
 
 // -----
 
 // Check proper ordering of intermixed attribute/type aliases.
 // CHECK: !tuple = tuple<
-// CHECK: #loc1 = loc(fused<!tuple
+// CHECK: = loc(fused<!tuple
 "test.op"() {alias_test = loc(fused<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, i32>>["test.mlir":10:8])} : () -> ()
 
 // -----
@@ -54,7 +54,7 @@
 // -----
 
 // Check that we don't print aliases for things that aren't printed.
-// CHECK: #loc1 = loc(fused<memref<1xi32>
+// CHECK: = loc(fused<memref<1xi32>
 // CHECK-NOT: #map
 "test.op"() {alias_test = loc(fused<memref<1xi32, affine_map<(d0) -> (d0)>>>["test.mlir":10:8])} : () -> ()
 
@@ -71,3 +71,16 @@
 "test.op"() {attr = #test.conditional_alias<#unalias_me>} : () -> ()
 // CHECK-NEXT: #test.conditional_alias<#test2Ealias>
 "test.op"() {attr = #test.conditional_alias<#keep_aliased>} : () -> ()
+
+// -----
+
+// Check that a deferred no_alias attr can be un-deferred.
+
+#keep_aliased = "alias_test:dot_in_name"
+#cond_alias = #test.conditional_alias<#keep_aliased>
+#no_alias = loc(fused<#cond_alias>["test.mlir":1:1])
+
+// CHECK: #[[TEST_ALIAS:.+]] = "alias_test:dot_in_name"
+// CHECK: fused<#test.conditional_alias<#[[TEST_ALIAS]]>
+// CHECK: "test.op"
+"test.op"() {attr = #no_alias} : () -> () loc(fused<#no_alias>["test.mlir":0:0])

[River707]

Nice!