[MLIR][IR] Fix InProgressAliasInfo init for non-alias (#109013)

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.

2 files changed, 25 insertions, 12 deletions

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index c7ed158..32182c0 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -584,9 +584,8 @@ private:
   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 27c5a75..e878d86 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])