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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Hideto Ueno (uenoku)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Bytecode/Reader/BytecodeReader.cpp (+37-22) 
- (modified) mlir/test/Bytecode/op_with_properties_deeply_nested_attr.mlir (+10-2) 


``````````diff
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
+}

``````````

</details>


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