[mlir] Fix infinite recursion in alias initializer

[zero9178]

Summary

The alias initializer keeps a list of child indices around. When an alias is then marked as non-deferrable, all children are also marked non-deferrable.

This is currently done naively which leads to an infinite recursion if using mutable types or attributes containing a cycle.

This patch fixes this by adding an early return if the alias is already marked non-deferrable. Since this function is the only way to mark an alias as non-deferrable, it is guaranteed that if it is marked non-deferrable, all its children are as well, and it is not required to walk all the children.
This incidentally makes the non-deferrable marking also O(n) instead of O(n^2) (although not performance sensitive obviously).

[mehdi_amini]

mlir/test/lib/Dialect/Test/TestTypes.cpp

493

I really don't like using global state: do we have an alternative?

[zero9178]

mlir/test/lib/Dialect/Test/TestTypes.cpp

493

There currently is not and all upstream mutable attributes and types that may create cycles suffer from this problem.
I am actually currently working on patches that should fix this.
For the time being, this is the best we got.

[mehdi_amini]

mlir/test/lib/Dialect/Test/TestTypes.cpp

493

Adding some tracking to the parser itself instead?

Looking forward to your patch!

mlir/lib/IR/AsmPrinter.cpp

if (maxAliasDepth)
     it->second.aliasDepth = maxAliasDepth + 1;
 
   // Propagate the alias depth of the value.
   return {(size_t)it->second.aliasDepth, aliasIndex};
 }
 
 void AliasInitializer::markAliasNonDeferrable(size_t aliasIndex) {
   auto it = std::next(aliases.begin(), aliasIndex);
+
+  // If already marked non-deferrable stop the recursion.
+  // All children should already be marked non-deferrable as well.
+  if (!it->second.canBeDeferred)
+    return;
+
   it->second.canBeDeferred = false;
 
   // Propagate the non-deferrable flag to any child aliases.
   for (size_t childIndex : it->second.childIndices)
     markAliasNonDeferrable(childIndex);
 }
 
 template <typename T>

mlir/test/IR/recursive-type.mlir

 // RUN: mlir-opt %s -test-recursive-types | FileCheck %s
 
 // CHECK: !testrec = !test.test_rec<type_to_alias, test_rec<type_to_alias>>
+// CHECK: ![[$NAME:.*]] = !test.test_rec_alias<name, !test.test_rec_alias<name>>
+// CHECK: ![[$NAME2:.*]] = !test.test_rec_alias<name2, tuple<!test.test_rec_alias<name2>, i32>>
 
 // CHECK-LABEL: @roundtrip
 func.func @roundtrip() {
   // CHECK: !test.test_rec<a, test_rec<b, test_type>>
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<a, test_rec<b, test_type>>
   // CHECK: !test.test_rec<c, test_rec<c>>
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<c, test_rec<c>>
   // Make sure walkSubElementType, which is used to generate aliases, doesn't go
   // into inifinite recursion.
   // CHECK: !testrec
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<type_to_alias, test_rec<type_to_alias>>
+
+  // CHECK: () -> ![[$NAME]]
+  // CHECK: () -> ![[$NAME]]
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec_alias<name, !test.test_rec_alias<name>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec_alias<name, !test.test_rec_alias<name>>
+
+  // CHECK: () -> ![[$NAME2]]
+  // CHECK: () -> ![[$NAME2]]
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec_alias<name2, tuple<!test.test_rec_alias<name2>, i32>>
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec_alias<name2, tuple<!test.test_rec_alias<name2>, i32>>
   return
 }
 
 // CHECK-LABEL: @create
 func.func @create() {
   // CHECK: !test.test_rec<some_long_and_unique_name, test_rec<some_long_and_unique_name>>
   return
 }

mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp

AliasResult getAlias(Type type, raw_ostream &os) const final {
     }
     if (auto recType = dyn_cast<TestRecursiveType>(type)) {
       if (recType.getName() == "type_to_alias") {
         // We only make alias for a specific recursive type.
         os << "testrec";
         return AliasResult::FinalAlias;
       }
     }
+    if (auto recAliasType = dyn_cast<TestRecursiveAliasType>(type)) {
+      os << recAliasType.getName();
+      return AliasResult::FinalAlias;
+    }
     return AliasResult::NoAlias;
   }
 
   //===------------------------------------------------------------------===//
   // Resources
   //===------------------------------------------------------------------===//
 
   std::string

mlir/test/lib/Dialect/Test/TestTypeDefs.td

def TestTypeElseAnchorStruct : Test_Type<"TestTypeElseAnchorStruct"> {
   let mnemonic = "else_anchor_struct";
   let assemblyFormat = "`<` (`?`) : (struct($a, $b)^)? `>`";
 }
 
 def TestI32 : Test_Type<"TestI32"> {
   let mnemonic = "i32";
 }
 
+def TestRecursiveAlias
+    : Test_Type<"TestRecursiveAlias", [NativeTypeTrait<"IsMutable">]> {
+  let mnemonic = "test_rec_alias";
+  let storageClass = "TestRecursiveTypeStorage";
+  let storageNamespace = "test";
+  let genStorageClass = 0;
+
+  let parameters = (ins "llvm::StringRef":$name);
+
+  let hasCustomAssemblyFormat = 1;
+
+  let extraClassDeclaration = [{
+    Type getBody() const;
+
+    void setBody(Type type);
+  }];
+}
+
 #endif // TEST_TYPEDEFS

mlir/test/lib/Dialect/Test/TestTypes.h

static FailureOr<std::optional<int>> parse(AsmParser &parser) {
     value.reset();
     return value;
   }
 };
 } // namespace mlir
 
 #include "TestTypeInterfaces.h.inc"
 
-#define GET_TYPEDEF_CLASSES
-#include "TestTypeDefs.h.inc"
-
 namespace test {
 
 /// Storage for simple named recursive types, where the type is identified by
 /// its name and can "contain" another type, including itself.
 struct TestRecursiveTypeStorage : public ::mlir::TypeStorage {
   using KeyTy = ::llvm::StringRef;
 
   explicit TestRecursiveTypeStorage(::llvm::StringRef key) : name(key) {}
public:
   ::mlir::Type getBody() const { return getImpl()->body; }
 
   /// Name/key getter.
   ::llvm::StringRef getName() { return getImpl()->name; }
 };
 
 } // namespace test
 
+#define GET_TYPEDEF_CLASSES
+#include "TestTypeDefs.h.inc"
+
 #endif // MLIR_TESTTYPES_H

mlir/test/lib/Dialect/Test/TestTypes.cpp

void TestDialect::printTestType(Type type, AsmPrinter &printer,
   }
   printer << ">";
 }
 
 void TestDialect::printType(Type type, DialectAsmPrinter &printer) const {
   SetVector<Type> stack;
   printTestType(type, printer, stack);
 }
+
+Type TestRecursiveAliasType::getBody() const { return getImpl()->body; }
+
+void TestRecursiveAliasType::setBody(Type type) { (void)Base::mutate(type); }
+
+StringRef TestRecursiveAliasType::getName() const { return getImpl()->name; }
+
+Type TestRecursiveAliasType::parse(AsmParser &parser) {
+  thread_local static SetVector<Type> stack;
+
+  StringRef name;
+  if (parser.parseLess() || parser.parseKeyword(&name))
+    return Type();
+  auto rec = TestRecursiveAliasType::get(parser.getContext(), name);
+
+  // If this type already has been parsed above in the stack, expect just the
+  // name.
+  if (stack.contains(rec)) {
+    if (failed(parser.parseGreater()))
+      return Type();
+    return rec;
+  }
+
+  // Otherwise, parse the body and update the type.
+  if (failed(parser.parseComma()))
+    return Type();
+  stack.insert(rec);
+  Type subtype;
+  if (parser.parseType(subtype))
+    return nullptr;
+  stack.pop_back();
+  if (!subtype || failed(parser.parseGreater()))
+    return Type();
+
+  rec.setBody(subtype);
+
+  return rec;
+}
+
+void TestRecursiveAliasType::print(AsmPrinter &printer) const {
+  thread_local static SetVector<Type> stack;
+
+  printer << "<" << getName();
+  if (!stack.contains(*this)) {
+    printer << ", ";
+    stack.insert(*this);
+    printer << getBody();
+    stack.pop_back();
+  }
+  printer << ">";
+}