[Mlir-commits] [mlir] d540855 - [mlir][Bytecode] Fix stale deferred worklist entries in attribute callback fallthrough (#186150)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-13T14:17:19+01:00
New Revision: d540855723a8afb2e13fe0c3dd37efd929ab0c1a

URL: https://github.com/llvm/llvm-project/commit/d540855723a8afb2e13fe0c3dd37efd929ab0c1a
DIFF: https://github.com/llvm/llvm-project/commit/d540855723a8afb2e13fe0c3dd37efd929ab0c1a.diff

LOG: [mlir][Bytecode] Fix stale deferred worklist entries in attribute callback fallthrough (#186150)

When parseCustomEntry() calls a user attribute/type callback that
internally reads sub-attributes/types via the bytecode reader, the
reader may add entries to the deferredWorklist if the depth limit is
exceeded. If the callback then returns success with an empty entry
(falling through to the regular dialect reader), the reader position is
reset but deferredWorklist retains stale entries from the failed partial
read.

This causes an assert(deferredWorklist.empty()) failure in debug builds
when the fallback dialect reader successfully parses the attribute.

Fix by saving and restoring deferredWorklist.size() around each callback
invocation, discarding any stale entries added during a callback's
partial read when the reader position is rolled back.

Fixes #163337

Assisted-by: Claude Code

Added: 
    mlir/test/Bytecode/bytecode_callback_with_spirv_and_custom_attr.mlir

Modified: 
    mlir/lib/Bytecode/Reader/BytecodeReader.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 07a9266dcd1fa..caa17ea57b950 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -1503,6 +1503,7 @@ LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
     // Try parsing with callbacks first if available.
     for (const auto &callback :
          parserConfig.getBytecodeReaderConfig().getTypeCallbacks()) {
+      size_t savedWorklistSize = deferredWorklist.size();
       if (failed(
               callback->read(dialectReader, entry.dialect->name, entry.entry)))
         return failure();
@@ -1510,14 +1511,17 @@ LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
       if (!!entry.entry)
         return success();
 
-      // Reset the reader if we failed to parse, so we can fall through the
-      // other parsing functions.
+      // The callback fell through without consuming the encoding. Reset the
+      // reader and restore the deferred worklist: any entries added during the
+      // callback's partial read are stale and must not persist.
+      deferredWorklist.resize(savedWorklistSize);
       reader = EncodingReader(entry.data, reader.getLoc());
     }
   } else {
     // Try parsing with callbacks first if available.
     for (const auto &callback :
          parserConfig.getBytecodeReaderConfig().getAttributeCallbacks()) {
+      size_t savedWorklistSize = deferredWorklist.size();
       if (failed(
               callback->read(dialectReader, entry.dialect->name, entry.entry)))
         return failure();
@@ -1525,8 +1529,10 @@ LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
       if (!!entry.entry)
         return success();
 
-      // Reset the reader if we failed to parse, so we can fall through the
-      // other parsing functions.
+      // The callback fell through without consuming the encoding. Reset the
+      // reader and restore the deferred worklist: any entries added during the
+      // callback's partial read are stale and must not persist.
+      deferredWorklist.resize(savedWorklistSize);
       reader = EncodingReader(entry.data, reader.getLoc());
     }
   }

diff  --git a/mlir/test/Bytecode/bytecode_callback_with_spirv_and_custom_attr.mlir b/mlir/test/Bytecode/bytecode_callback_with_spirv_and_custom_attr.mlir
new file mode 100644
index 0000000000000..e2bf0453f3f45
--- /dev/null
+++ b/mlir/test/Bytecode/bytecode_callback_with_spirv_and_custom_attr.mlir
@@ -0,0 +1,30 @@
+// RUN: mlir-opt %s --test-bytecode-roundtrip="test-kind=4" | FileCheck %s
+
+// Regression test for https://github.com/llvm/llvm-project/issues/163337
+//
+// When using test-kind=4, the attribute callback calls the builtin dialect
+// reader for each attribute. For dialect-specific attributes (e.g., spirv.*),
+// the builtin reader fails and the callback falls through to the regular reader.
+// During the failing read, the callback may add entries to the deferred parsing
+// worklist. These stale entries must be discarded when the reader position is
+// reset, otherwise the `deferredWorklist.empty()` assertion fires in debug
+// builds and may corrupt subsequent attribute resolution in release builds.
+
+spirv.module Logical GLSL450 {
+  spirv.func @callee() -> () "None" {
+    spirv.Kill
+  }
+  spirv.func @caller() -> () "None" {
+    spirv.FunctionCall @callee() : () -> ()
+    spirv.Return
+  }
+}
+
+// CHECK-LABEL: spirv.module Logical GLSL450
+// CHECK: spirv.func @callee
+// CHECK: spirv.func @caller
+
+// Verify the custom test attribute roundtrips correctly through the callback.
+"test.versionedC"() <{attribute = #test.attr_params<42, 24>}> : () -> ()
+
+// CHECK: "test.versionedC"() <{attribute = #test.attr_params<42, 24>}> : () -> ()