[mlir:bytecode] Support lazy loading dynamically isolated regions

[rriddle]

Summary

We currently only support lazy loading for regions that
statically implement the IsolatedFromAbove trait, but that
limits the amount of operations that can be lazily loaded. This review
lifts that restriction by computing which operations have isolated
regions when numbering, allowing any operation to be lazily loaded
as long as it doesn't use values defined above.

[mehdi_amini]

mlir/lib/Bytecode/Writer/IRNumbering.cpp

190

Isn't this redundant with the numbering optional isIsolatedFromAbove?
If numbering->isIsolatedFromAbove has a value, it means we already know and shouldCheckIsolation should be false.

194

Some description of what's happening below would be useful, it's a non-trivial algorithm.
I would document the invariants around the stack as well

200

This if the longest, can we put it at the end and check the other conditions with an early return (and reduce indentation)

[rriddle]

mlir/lib/Bytecode/Writer/IRNumbering.cpp

190

Not exactly. isIsolatedFromAbove applies for the current op, but this flag is if the current op or any of its parents still need checking. We can only avoid checking if we know the current region and all of the parents have been resolved, otherwise we might miss checking something like:

top.op {
  %value = ...
  middle.op {
    %value2 = ...
    inner.op {
      // Here we mark inner.op as not isolated (but not middle.op yet)
      use %value2 

      // Here inner.op is already known to be non-isolated, but middle.op is now also
      // discovered to be non-isolated.
      use %value
    }
  }
}

[mehdi_amini]

mlir/lib/Bytecode/Writer/IRNumbering.cpp

190

Tricky... Can you add this comment in the code?

Edit: Oh.. you did already! :)

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

LogicalResult BytecodeWriter::writeOp(EncodingEmitter &emitter, Operation *op) {
   // Update the mask for the operation.
   emitter.patchByte(maskOffset, opEncodingMask);
 
   // With the mask emitted, we can now emit the regions of the operation. We do
   // this after mask emission to avoid offset complications that may arise by
   // emitting the regions first (e.g. if the regions are huge, backpatching the
   // op encoding mask is more annoying).
   if (numRegions) {
-    bool isIsolatedFromAbove = op->hasTrait<OpTrait::IsIsolatedFromAbove>();
+    bool isIsolatedFromAbove = numberingState.isIsolatedFromAbove(op);
     emitter.emitVarIntWithFlag(numRegions, isIsolatedFromAbove);
 
     // If the region is not isolated from above, or we are emitting bytecode
     // targeting version <kLazyLoading, we don't use a section.
     if (isIsolatedFromAbove &&
         config.bytecodeVersion >= bytecode::kLazyLoading) {
       EncodingEmitter regionEmitter;
       if (failed(writeRegions(regionEmitter, op->getRegions())))

mlir/lib/Bytecode/Writer/IRNumbering.h

struct DialectNumbering {
   /// The referenced resources of this dialect.
   SetVector<AsmDialectResourceHandle> resources;
 
   /// A mapping from resource key to the corresponding resource numbering entry.
   llvm::MapVector<StringRef, DialectResourceNumbering *> resourceMap;
 };
 
 //===----------------------------------------------------------------------===//
+// Operation Numbering
+//===----------------------------------------------------------------------===//
+
+/// This class represents the numbering entry of an operation.
+struct OperationNumbering {
+  OperationNumbering(unsigned number) : number(number) {}
+
+  /// The number assigned to this operation.
+  unsigned number;
+
+  /// A flag indicating if this operation's regions are isolated. If unset, the
+  /// operation isn't yet known to be isolated.
+  std::optional<bool> isIsolatedFromAbove;
+};
+
+//===----------------------------------------------------------------------===//
 // IRNumberingState
 //===----------------------------------------------------------------------===//
 
 /// This class manages numbering IR entities in preparation of bytecode
 /// emission.
 class IRNumberingState {
 public:
   IRNumberingState(Operation *op, const BytecodeWriterConfig &config);
unsigned getNumber(Attribute attr) {
     assert(attrs.count(attr) && "attribute not numbered");
     return attrs[attr]->number;
   }
   unsigned getNumber(Block *block) {
     assert(blockIDs.count(block) && "block not numbered");
     return blockIDs[block];
   }
   unsigned getNumber(Operation *op) {
-    assert(operationIDs.count(op) && "operation not numbered");
-    return operationIDs[op];
+    assert(operations.count(op) && "operation not numbered");
+    return operations[op]->number;
   }
   unsigned getNumber(OperationName opName) {
     assert(opNames.count(opName) && "opName not numbered");
     return opNames[opName]->number;
   }
   unsigned getNumber(Type type) {
     assert(types.count(type) && "type not numbered");
     return types[type]->number;
public:
   }
 
   /// Return the number of operations in the given block.
   unsigned getOperationCount(Block *block) {
     assert(blockOperationCounts.count(block) && "block not numbered");
     return blockOperationCounts[block];
   }
 
+  /// Return if the given operation is isolated from above.
+  bool isIsolatedFromAbove(Operation *op) {
+    assert(operations.count(op) && "operation not numbered");
+    return operations[op]->isIsolatedFromAbove.value_or(false);
+  }
+
   /// Get the set desired bytecode version to emit.
   int64_t getDesiredBytecodeVersion() const;
   
 private:
   /// This class is used to provide a fake dialect writer for numbering nested
   /// attributes and types.
   struct NumberingDialectWriter;
 
+  /// Compute the global numbering state for the given root operation.
+  void computeGlobalNumberingState(Operation *rootOp);
+
   /// Number the given IR unit for bytecode emission.
   void number(Attribute attr);
   void number(Block &block);
   DialectNumbering &numberDialect(Dialect *dialect);
   DialectNumbering &numberDialect(StringRef dialect);
   void number(Operation &op);
   void number(OperationName opName);
   void number(Region &region);
   void number(Type type);
 
   /// Number the given dialect resources.
   void number(Dialect *dialect, ArrayRef<AsmDialectResourceHandle> resources);
 
   /// Finalize the numberings of any dialect resources.
   void finalizeDialectResourceNumberings(Operation *rootOp);
 
   /// Mapping from IR to the respective numbering entries.
   DenseMap<Attribute, AttributeNumbering *> attrs;
+  DenseMap<Operation *, OperationNumbering *> operations;
   DenseMap<OperationName, OpNameNumbering *> opNames;
   DenseMap<Type, TypeNumbering *> types;
   DenseMap<Dialect *, DialectNumbering *> registeredDialects;
   llvm::MapVector<StringRef, DialectNumbering *> dialects;
   std::vector<AttributeNumbering *> orderedAttrs;
   std::vector<OpNameNumbering *> orderedOpNames;
   std::vector<TypeNumbering *> orderedTypes;
 
   /// A mapping from dialect resource handle to the numbering for the referenced
   /// resource.
   llvm::DenseMap<AsmDialectResourceHandle, DialectResourceNumbering *>
       dialectResources;
 
   /// Allocators used for the various numbering entries.
   llvm::SpecificBumpPtrAllocator<AttributeNumbering> attrAllocator;
   llvm::SpecificBumpPtrAllocator<DialectNumbering> dialectAllocator;
+  llvm::SpecificBumpPtrAllocator<OperationNumbering> opAllocator;
   llvm::SpecificBumpPtrAllocator<OpNameNumbering> opNameAllocator;
   llvm::SpecificBumpPtrAllocator<DialectResourceNumbering> resourceAllocator;
   llvm::SpecificBumpPtrAllocator<TypeNumbering> typeAllocator;
 
-  /// The value ID for each Operation, Block and Value.
-  DenseMap<Operation *, unsigned> operationIDs;
+  /// The value ID for each Block and Value.
   DenseMap<Block *, unsigned> blockIDs;
   DenseMap<Value, unsigned> valueIDs;
 
   /// The number of operations in each block.
   DenseMap<Block *, unsigned> blockOperationCounts;
 
   /// A map from region to the number of blocks and values within that region.
   DenseMap<Region *, std::pair<unsigned, unsigned>> regionBlockValueCounts;

mlir/lib/Bytecode/Writer/IRNumbering.cpp

static void groupByDialectPerByte(T range) {
   // Assign the entry numbers based on the sort order.
   for (auto [idx, value] : llvm::enumerate(range))
     value->number = idx;
 }
 
 IRNumberingState::IRNumberingState(Operation *op,
                                    const BytecodeWriterConfig &config)
     : config(config) {
-  // Compute a global operation ID numbering according to the pre-order walk of
-  // the IR. This is used as reference to construct use-list orders.
-  unsigned operationID = 0;
-  op->walk<WalkOrder::PreOrder>(
-      [&](Operation *op) { operationIDs.try_emplace(op, operationID++); });
+  computeGlobalNumberingState(op);
 
   // Number the root operation.
   number(*op);
 
-  // Push all of the regions of the root operation onto the worklist.
+  // A worklist of region contexts to number and the next value id before that
+  // region.
   SmallVector<std::pair<Region *, unsigned>, 8> numberContext;
-  for (Region &region : op->getRegions())
-    numberContext.emplace_back(&region, nextValueID);
+
+  // Functor to push the regions of the given operation onto the numbering
+  // context.
+  auto addOpRegionsToNumber = [&](Operation *op) {
+    MutableArrayRef<Region> regions = op->getRegions();
+    if (regions.empty())
+      return;
+
+    // Isolated regions don't share value numbers with their parent, so we can
+    // start numbering these regions at zero.
+    unsigned opFirstValueID = isIsolatedFromAbove(op) ? 0 : nextValueID;
+    for (Region &region : regions)
+      numberContext.emplace_back(&region, opFirstValueID);
+  };
+  addOpRegionsToNumber(op);
 
   // Iteratively process each of the nested regions.
   while (!numberContext.empty()) {
     Region *region;
     std::tie(region, nextValueID) = numberContext.pop_back_val();
     number(*region);
 
     // Traverse into nested regions.
-    for (Operation &op : region->getOps()) {
-      // Isolated regions don't share value numbers with their parent, so we can
-      // start numbering these regions at zero.
-      unsigned opFirstValueID =
-          op.hasTrait<OpTrait::IsIsolatedFromAbove>() ? 0 : nextValueID;
-      for (Region &region : op.getRegions())
-        numberContext.emplace_back(&region, opFirstValueID);
-    }
+    for (Operation &op : region->getOps())
+      addOpRegionsToNumber(&op);
   }
 
   // Number each of the dialects. For now this is just in the order they were
   // found, given that the number of dialects on average is small enough to fit
   // within a singly byte (128). If we ever have real world use cases that have
   // a huge number of dialects, this could be made more intelligent.
   for (auto [idx, dialect] : llvm::enumerate(dialects))
     dialect.second->number = idx;
IRNumberingState::IRNumberingState(Operation *op,
   groupByDialectPerByte(llvm::MutableArrayRef(orderedAttrs));
   groupByDialectPerByte(llvm::MutableArrayRef(orderedOpNames));
   groupByDialectPerByte(llvm::MutableArrayRef(orderedTypes));
 
   // Finalize the numbering of the dialect resources.
   finalizeDialectResourceNumberings(op);
 }
 
+void IRNumberingState::computeGlobalNumberingState(Operation *rootOp) {
+  // A simple state struct tracking data used when walking operations.
+  struct StackState {
+    /// The operation currently being walked.
+    Operation *op;
+
+    /// The numbering of the operation.
+    OperationNumbering *numbering;
+
+    /// A flag indicating if the current state or one of its parents has
+    /// unresolved isolation status. This is tracked separately from the
+    /// isIsolatedFromAbove bit on `numbering` because we need to be able to
+    /// handle the given case:
+    ///   top.op {
+    ///     %value = ...
+    ///     middle.op {
+    ///       %value2 = ...
+    ///       inner.op {
+    ///         // Here we mark `inner.op` as not isolated. Note `middle.op`
+    ///         // isn't known not isolated yet.
+    ///         use.op %value2
+    ///
+    ///         // Here inner.op is already known to be non-isolated, but
+    ///         // `middle.op` is now also discovered to be non-isolated.
+    ///         use.op %value
+    ///       }
+    ///     }
+    ///   }
+    bool hasUnresolvedIsolation;
+  };
+
+  // Compute a global operation ID numbering according to the pre-order walk of
+  // the IR. This is used as reference to construct use-list orders.
+  unsigned operationID = 0;
+
+  // Walk each of the operations within the IR, tracking a stack of operations
+  // as we recurse into nested regions. This walk method hooks in at two stages
+  // during the walk:
+  //
+  //   BeforeAllRegions:
+  //     Here we generate a numbering for the operation and push it onto the
+  //     stack if it has regions. We also compute the isolation status of parent
+  //     regions at this stage. This is done by checking the parent regions of
+  //     operands used by the operation, and marking each region between the
+  //     the operand region and the current as not isolated. See
+  //     StackState::hasUnresolvedIsolation above for an example.
+  //
+  //   AfterAllRegions:
+  //     Here we pop the operation from the stack, and if it hasn't been marked
+  //     as non-isolated, we mark it as so. A non-isolated use would have been
+  //     found while walking the regions, so it is safe to mark the operation at
+  //     this point.
+  //
+  SmallVector<StackState> opStack;
+  rootOp->walk([&](Operation *op, const WalkStage &stage) {
+    // After visiting all nested regions, we pop the operation from the stack.
+    if (stage.isAfterAllRegions()) {
+      // If no non-isolated uses were found, we can safely mark this operation
+      // as isolated from above.
+      OperationNumbering *numbering = opStack.pop_back_val().numbering;
+      if (!numbering->isIsolatedFromAbove.has_value())
+        numbering->isIsolatedFromAbove = true;
+      return;
+    }
+
+    // When visiting before nested regions, we process "IsolatedFromAbove"
+    // checks and compute the number for this operation.
+    if (!stage.isBeforeAllRegions())
+      return;
+    // Update the isolation status of parent regions if any have yet to be
+    // resolved.
+    if (!opStack.empty() && opStack.back().hasUnresolvedIsolation) {
+      Region *parentRegion = op->getParentRegion();
+      for (Value operand : op->getOperands()) {
+        Region *operandRegion = operand.getParentRegion();
+        if (operandRegion == parentRegion)
+          continue;
+        // We've found a use of an operand outside of the current region,
+        // walk the operation stack searching for the parent operation,
+        // marking every region on the way as not isolated.
+        Operation *operandContainerOp = operandRegion->getParentOp();
+        auto it = std::find_if(
+            opStack.rbegin(), opStack.rend(), [=](const StackState &it) {
+              // We only need to mark up to the container region, or the first
+              // that has an unresolved status.
+              return !it.hasUnresolvedIsolation || it.op == operandContainerOp;
+            });
+        assert(it != opStack.rend() && "expected to find the container");
+        for (auto &state : llvm::make_range(opStack.rbegin(), it)) {
+          // If we stopped at a region that knows its isolation status, we can
+          // stop updating the isolation status for the parent regions.
+          state.hasUnresolvedIsolation = it->hasUnresolvedIsolation;
+          state.numbering->isIsolatedFromAbove = false;
+        }
+      }
+    }
+
+    // Compute the number for this op and push it onto the stack.
+    auto *numbering =
+        new (opAllocator.Allocate()) OperationNumbering(operationID++);
+    if (op->hasTrait<OpTrait::IsIsolatedFromAbove>())
+      numbering->isIsolatedFromAbove = true;
+    operations.try_emplace(op, numbering);
+    if (op->getNumRegions()) {
+      opStack.emplace_back(StackState{
+          op, numbering, !numbering->isIsolatedFromAbove.has_value()});
+    }
+  });
+}
+
 void IRNumberingState::number(Attribute attr) {
   auto it = attrs.insert({attr, nullptr});
   if (!it.second) {
     ++it.first->second->refCount;
     return;
   }
   auto *numbering = new (attrAllocator.Allocate()) AttributeNumbering(attr);
   it.first->second = numbering;

mlir/test/Bytecode/bytecode-lazy-loading.mlir

test.isolated_region %result#1 {
     "test.consumer"(%result#1) : (index) -> ()
   }
   
   test.isolated_regions {
     "test.unknown_op"() : () -> ()
   }, {
     "test.unknown_op"() : () -> ()
   }
+  
+  // Ensure operations that aren't tagged as IsolatedFromAbove can
+  // still be lazy loaded if they don't have references to values
+  // defined above.
+  "test.one_region_op"() ({
+    "test.unknown_op"() : () -> ()
+  }) : () -> ()
+
+  // Similar test as above, but check that if one region has a reference
+  // to a value defined above, we don't lazy load the operation.
+  "test.two_region_op"() ({
+    "test.unknown_op"() : () -> ()
+  }, {
+    "test.consumer"(%0) : (index) -> ()
+  }) : () -> ()
   return
 }
 
 // Before version 2, we can't support lazy loading.
 // OLD-BYTECODE-NOT: Has 1 ops to materialize
 // OLD-BYTECODE-NOT: Materializing
 // OLD-BYTECODE: Has 0 ops to materialize
 

 // CHECK-NOT: arith
 // CHECK: Materializing...
 // CHECK: "func.func"() <{function_type = () -> (), sym_name = "op_with_passthrough_region_args"}> ({
 // CHECK: arith
 // CHECK: one_region_op
 // CHECK: test.consumer
 // CHECK: isolated_region
 // CHECK-NOT: test.consumer
-// CHECK: Has 3 ops to materialize
+// CHECK: test.one_region_op
+// CHECK-NOT: test.op
+// CHECK: test.two_region_op
+// CHECK: test.unknown_op
+// CHECK: test.consumer
+// CHECK: Has 4 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_region
 // CHECK-NOT:  test.consumer
 // CHECK: Materializing...
 // CHECK: test.isolated_region
 // CHECK: ^bb0(%arg0: index):
 // CHECK:  test.consumer
-// CHECK: Has 2 ops to materialize
+// CHECK: Has 3 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_region
 // CHECK-NOT: test.consumer
 // CHECK: Materializing...
 // CHECK: test.isolated_region
 // CHECK: test.consumer
-// CHECK: Has 1 ops to materialize
+// CHECK: Has 2 ops to materialize
 
 // CHECK: Before Materializing...
 // CHECK: test.isolated_regions
 // CHECK-NOT: test.unknown_op
 // CHECK: Materializing...
 // CHECK: test.isolated_regions
 // CHECK: test.unknown_op
 // CHECK: test.unknown_op
+// CHECK: Has 1 ops to materialize
+
+// CHECK: Before Materializing...
+// CHECK: test.one_region_op
+// CHECK-NOT: test.unknown_op
+// CHECK: Materializing...
+// CHECK: test.one_region_op
+// CHECK: test.unknown_op
 // CHECK: Has 0 ops to materialize