[Mlir-commits] [mlir] 6ee1aba - [mlir][bytecode] Fix lazy loading of non-isolated regions

[River Riddle]

Author: River Riddle
Date: 2023-06-26T16:33:20-07:00
New Revision: 6ee1aba8acc1255a6eb9e34788a2376c6357458b

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

LOG: [mlir][bytecode] Fix lazy loading of non-isolated regions

The bytecode reader currently assumes all regions can be lazy
loaded, which breaks reading any non-isolated region. This patch
fixes that by properly handling nested non-lazy regions, and only
considers isolated regions as lazy.

Differential Revision: https://reviews.llvm.org/D153795

Added: 
    

Modified: 
    mlir/lib/Bytecode/Reader/BytecodeReader.cpp
    mlir/test/Bytecode/bytecode-lazy-loading.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
index 5aa24ba873ec2..5b313af8cee33 100644
--- a/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
+++ b/mlir/lib/Bytecode/Reader/BytecodeReader.cpp
@@ -1317,11 +1317,11 @@ class mlir::BytecodeReader::Impl {
     regionStack.push_back(std::move(it->getSecond()->second));
     lazyLoadableOps.erase(it->getSecond());
     lazyLoadableOpsMap.erase(it);
-    auto result = parseRegions(regionStack, regionStack.back());
-    assert((regionStack.empty() || failed(result)) &&
-           "broken invariant: regionStack should be empty when parseRegions "
-           "succeeds");
-    return result;
+
+    while (!regionStack.empty())
+      if (failed(parseRegions(regionStack, regionStack.back())))
+        return failure();
+    return success();
   }
 
   /// Return the context for this config.
@@ -2094,14 +2094,11 @@ BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
             childState.owningReader =
                 std::make_unique<EncodingReader>(sectionData, fileLoc);
             childState.reader = childState.owningReader.get();
-          }
 
-          if (lazyLoading) {
-            // If the user has a callback set, they have the opportunity
-            // to control lazyloading as we go.
-            if (!lazyOpsCallback || !lazyOpsCallback(*op)) {
-              lazyLoadableOps.push_back(
-                  std::make_pair(*op, std::move(childState)));
+            // If the user has a callback set, they have the opportunity to
+            // control lazyloading as we go.
+            if (lazyLoading && (!lazyOpsCallback || !lazyOpsCallback(*op))) {
+              lazyLoadableOps.emplace_back(*op, std::move(childState));
               lazyLoadableOpsMap.try_emplace(*op,
                                              std::prev(lazyLoadableOps.end()));
               continue;

diff  --git a/mlir/test/Bytecode/bytecode-lazy-loading.mlir b/mlir/test/Bytecode/bytecode-lazy-loading.mlir
index a4f7974b0b690..731d0e865ded3 100644
--- a/mlir/test/Bytecode/bytecode-lazy-loading.mlir
+++ b/mlir/test/Bytecode/bytecode-lazy-loading.mlir
@@ -3,6 +3,9 @@
 
 
 func.func @op_with_passthrough_region_args() {
+  // Ensure we can handle nested non-isolated/non-lazy regions.
+  "test.one_region_op"() ({}) : () -> ()
+
   %0 = arith.constant 10 : index
   test.isolated_region %0 {
     "test.consumer"(%0) : (index) -> ()
@@ -36,6 +39,7 @@ func.func @op_with_passthrough_region_args() {
 // CHECK-NOT: arith
 // CHECK: Materializing...
 // CHECK: "func.func"() <{function_type = () -> (), sym_name = "op_with_passthrough_region_args"}> ({
+// CHECK: one_region_op
 // CHECK: arith
 // CHECK: isolated_region
 // CHECK-NOT: test.consumer