[Mlir-commits] [mlir] [mlir][Bytecode] Fix infinite loop by tracking type/attribute in deferred worklist (PR #174874)

[Hideto Ueno]

https://github.com/uenoku created https://github.com/llvm/llvm-project/pull/174874

The bytecode reader could enter an infinite loop when parsing deeply nested attributes containing type references. The deferred worklist stored only indices without distinguishing between attributes and types, causing type indexes to be misinterpreted as attributes.

This patch changes the deferred worklist to store pairs of (index, isType) to track whether each deferred entry is a type or attribute. The worklist processing logic is updated to resolve the correct entry type.

https://github.com/llvm/llvm-project/pull/170993#issuecomment-3717224879 is the context.

>From f858ec6ec5392218cfd9aa156e38bddb54864155 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko at gmail.com>
Date: Tue, 6 Jan 2026 18:14:29 -0800
Subject: [PATCH] [mlir][Bytecode] Fix infinite loop by tracking type/attribute
 in deferred worklist

The bytecode reader entered an infinite loop when parsing deeply nested
attributes containing type references. The deferred worklist stored only
indices without distinguishing between attributes and types, causing
type index 0 to be misinterpreted as attribute index 0.

This patch changes the deferred worklist to store pairs of (index, isType)
to track whether each deferred entry is a type or attribute. The worklist
processing logic is updated to resolve the correct entry type.
---
 mlir/lib/Bytecode/Reader/BytecodeReader.cpp   | 59 ++++++++++++-------
 ...op_with_properties_deeply_nested_attr.mlir | 12 +++-
 2 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 6159e0f037370..6bd4c91f3590b 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -950,7 +950,9 @@ class AttrTypeReader {
   }
 
   /// Add an index to the deferred worklist for re-parsing.
-  void addDeferredParsing(uint64_t index) { deferredWorklist.push_back(index); }
+  void addDeferredParsing(uint64_t index, char isType) {
+    deferredWorklist.emplace_back(index, isType);
+  }
 
   /// Whether currently resolving.
   bool isResolving() const { return resolving; }
@@ -1007,7 +1009,9 @@ class AttrTypeReader {
 
   /// Worklist for deferred attribute/type parsing. This is used to handle
   /// deeply nested structures like CallSiteLoc iteratively.
-  std::vector<uint64_t> deferredWorklist;
+  /// - The first element is the index of the attribute/type to parse.
+  /// - The second element is a flag indicating if the entry is a type.
+  std::vector<std::pair<uint64_t, char>> deferredWorklist;
 
   /// Flag indicating if we are currently resolving an attribute or type.
   bool resolving = false;
@@ -1084,7 +1088,7 @@ class DialectReader : public DialectBytecodeReader {
         result = attr;
         return success();
       }
-      attrTypeReader.addDeferredParsing(index);
+      attrTypeReader.addDeferredParsing(index, /*isType=*/false);
       return failure();
     }
     return attrTypeReader.readAttribute(index, result, depth + 1);
@@ -1113,7 +1117,7 @@ class DialectReader : public DialectBytecodeReader {
         result = type;
         return success();
       }
-      attrTypeReader.addDeferredParsing(index);
+      attrTypeReader.addDeferredParsing(index, /*isType=*/true);
       return failure();
     }
     return attrTypeReader.readType(index, result, depth + 1);
@@ -1380,28 +1384,40 @@ T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries,
   // - Pop from front to process
   // - Push new dependencies to front (depth-first)
   // - Move failed entries to back (retry after dependencies)
-  std::deque<size_t> worklist;
-  llvm::DenseSet<size_t> inWorklist;
+  std::deque<std::pair<uint64_t, char>> worklist;
+  llvm::DenseSet<std::pair<uint64_t, char>> inWorklist;
+
+  char isTypeEntry = std::is_same_v<T, Type>;
+  auto addToWorklistFront = [&](std::pair<uint64_t, char> entry) {
+    if (inWorklist.insert(entry).second)
+      worklist.push_front(entry);
+  };
 
   // Add the original index and any dependencies from the fast path attempt.
-  worklist.push_back(index);
-  inWorklist.insert(index);
-  for (uint64_t idx : llvm::reverse(deferredWorklist)) {
-    if (inWorklist.insert(idx).second)
-      worklist.push_front(idx);
-  }
+  worklist.emplace_back(index, isTypeEntry);
+  inWorklist.insert({index, isTypeEntry});
+  for (auto entry : llvm::reverse(deferredWorklist))
+    addToWorklistFront(entry);
 
   while (!worklist.empty()) {
-    size_t currentIndex = worklist.front();
+    auto [currentIndex, isType] = worklist.front();
     worklist.pop_front();
 
     // Clear the deferred worklist before parsing to capture any new entries.
     deferredWorklist.clear();
 
-    T result;
-    if (succeeded(readEntry(entries, currentIndex, result, entryType, depth))) {
-      inWorklist.erase(currentIndex);
-      continue;
+    if (isType) {
+      Type result;
+      if (succeeded(readType(currentIndex, result, depth))) {
+        inWorklist.erase({currentIndex, isType});
+        continue;
+      }
+    } else {
+      Attribute result;
+      if (succeeded(readAttribute(currentIndex, result, depth))) {
+        inWorklist.erase({currentIndex, isType});
+        continue;
+      }
     }
 
     if (deferredWorklist.empty()) {
@@ -1410,13 +1426,12 @@ T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries,
     }
 
     // Move this entry to the back to retry after dependencies.
-    worklist.push_back(currentIndex);
+    worklist.emplace_back(currentIndex, isType);
 
     // Add dependencies to the front (in reverse so they maintain order).
-    for (uint64_t idx : llvm::reverse(deferredWorklist)) {
-      if (inWorklist.insert(idx).second)
-        worklist.push_front(idx);
-    }
+    for (auto entry : llvm::reverse(deferredWorklist))
+      addToWorklistFront(entry);
+
     deferredWorklist.clear();
   }
   return entries[index].entry;
diff --git a/mlir/test/Bytecode/op_with_properties_deeply_nested_attr.mlir b/mlir/test/Bytecode/op_with_properties_deeply_nested_attr.mlir
index a7d964927a6ec..9ee3fa6e12ccf 100644
--- a/mlir/test/Bytecode/op_with_properties_deeply_nested_attr.mlir
+++ b/mlir/test/Bytecode/op_with_properties_deeply_nested_attr.mlir
@@ -1,6 +1,14 @@
-// RUN: mlir-opt %s  --verify-roundtrip
+// RUN: mlir-opt %s --verify-roundtrip  --split-input-file
 
 func.func @deeply_nested_attribute() -> i32 {
   %0 = "test.constant"() {value = [[[[[[[[[[[[1]]]]]]]]]]]]} : () -> i32
   return %0 : i32
-}
\ No newline at end of file
+}
+
+// -----
+// Make sure bytecode reader doesn't get stuck in an infinite loop when
+// parsing a specific deeply nested attribute structure.
+func.func @deeply_nested_infinite_loop() -> i32 {
+  %0 = "test.constant"() {value = [[[[[[1]]]]]]} : () -> i32
+  return %0 : i32
+}