[mlir] Async reference counting for block successors with divergent reference counted liveness

[ezhulenev]

Summary

Support reference counted values implicitly passed (live) only to some of the successors.

Example: if branched to ^bb2 token will leak, unless drop_ref operation is properly created

^entry:
  %token = async.runtime.create : !async.token
   cond_br %cond, ^bb1, ^bb2
^bb1:
  async.runtime.await %token
  async.runtime.drop_ref %token
  br ^bb2
^bb2:
  return

[rriddle]

mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp

261

nit: SmallPtrSet?

301–308

This is only used once, can you just inline it and stream the args in?

emitOpError() << ".." << value;

346

Don't recompute the end iterator every loop

[ezhulenev]

Address comments

mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp

301–308

I removed value printing because in practice it will be too much details, e.g. it can print all the async region there

346

replaced with llvm::enumerate

mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp

private:
   ///      use.
   LogicalResult addDropRefAfterLastUse(Value value);
 
   /// (#2) Adds the `add_ref` operation before the function call taking `value`
   /// operand to ensure that the value passed to the function entry block
   /// has a `+1` reference count.
   LogicalResult addAddRefBeforeFunctionCall(Value value);
 
-  /// (#3) Verifies that if a block has a value in the `liveOut` set, then the
-  /// value is in `liveIn` set in all successors.
+  /// (#3) Adds the `drop_ref` operation to account for successor blocks with
+  /// divergent `liveIn` property: `value` is not in the `liveIn` set of all
+  /// successor blocks.
   ///
   /// Example:
   ///
   ///   ^entry:
   ///     %token = async.runtime.create : !async.token
   ///     cond_br %cond, ^bb1, ^bb2
   ///   ^bb1:
   ///     async.runtime.await %token
-  ///     return
+  ///     async.runtime.drop_ref %token
+  ///     br ^bb2
   ///   ^bb2:
   ///     return
   ///
-  /// This CFG will be rejected because ^bb2 does not have `value` in the
-  /// `liveIn` set, and it will leak a reference counted object.
+  /// In this example ^bb2 does not have `value` in the `liveIn` set, so we have
+  /// to branch into a special "reference counting block" from the ^entry that
+  /// will have a `drop_ref` operation, and then branch into the ^bb2.
+  ///
+  /// After transformation:
+  ///
+  ///   ^entry:
+  ///     %token = async.runtime.create : !async.token
+  ///     cond_br %cond, ^bb1, ^reference_counting
+  ///   ^bb1:
+  ///     async.runtime.await %token
+  ///     async.runtime.drop_ref %token
+  ///     br ^bb2
+  ///   ^reference_counting:
+  ///     async.runtime.drop_ref %token
+  ///     br ^bb2
+  ///   ^bb2:
+  ///     return
   ///
   /// An exception to this rule are blocks with `async.coro.suspend` terminator,
   /// because in Async to LLVM lowering it is guaranteed that the control flow
   /// will jump into the resume block, and then follow into the cleanup and
   /// suspend blocks.
   ///
   /// Example:
   ///
   ///  ^entry(%value: !async.value<f32>):
   ///     async.runtime.await_and_resume %value, %hdl : !async.value<f32>
   ///     async.coro.suspend %ret, ^suspend, ^resume, ^cleanup
   ///   ^resume:
   ///     %0 = async.runtime.load %value
   ///     br ^cleanup
   ///   ^cleanup:
   ///     ...
   ///   ^suspend:
   ///     ...
   ///
   /// Although cleanup and suspend blocks do not have the `value` in the
   /// `liveIn` set, it is guaranteed that execution will eventually continue in
   /// the resume block (we never explicitly destroy coroutines).
-  LogicalResult verifySuccessors(Value value);
+  LogicalResult addDropRefInDivergentLivenessSuccessor(Value value);
 };
 
 } // namespace
 
 LogicalResult AsyncRuntimeRefCountingPass::addDropRefAfterLastUse(Value value) {
   OpBuilder builder(value.getContext());
   Location loc = value.getLoc();
 
for (Operation *user : value.getUsers()) {
     // reference to the function entry block.
     builder.setInsertionPoint(user);
     builder.create<RuntimeAddRefOp>(loc, value, builder.getI32IntegerAttr(1));
   }
 
   return success();
 }
 
-LogicalResult AsyncRuntimeRefCountingPass::verifySuccessors(Value value) {
+LogicalResult
+AsyncRuntimeRefCountingPass::addDropRefInDivergentLivenessSuccessor(
+    Value value) {
+  using BlockSet = llvm::SmallPtrSet<Block *, 4>;
+
   OpBuilder builder(value.getContext());
 
-  // Blocks with successfors with different `liveIn` properties of the `value`.
-  llvm::SmallSet<Block *, 4> divergentLivenessBlocks;
+  // If a block has successors with different `liveIn` property of the `value`,
+  // record block successors that do not thave the `value` in the `liveIn` set.
+  llvm::SmallDenseMap<Block *, BlockSet> divergentLivenessBlocks;
 
   // Use liveness analysis to find the placement of `drop_ref`operation.
   auto &liveness = getAnalysis<Liveness>();
 
   // Because we only add `drop_ref` operations to the region that defines the
   // `value` we can only process CFG for the same region.
   Region *definingRegion = value.getParentRegion();
 
   // Collect blocks with successors with mismatching `liveIn` sets.
   for (Block &block : definingRegion->getBlocks()) {
     const LivenessBlockInfo *blockLiveness = liveness.getLiveness(&block);
 
     // Skip the block if value is not in the `liveOut` set.
     if (!blockLiveness->isLiveOut(value))
       continue;
 
-    // Sucessors with value in `liveIn` set and not value in `liveIn` set.
-    llvm::SmallSet<Block *, 4> liveInSuccessors;
-    llvm::SmallSet<Block *, 4> noLiveInSuccessors;
+    BlockSet liveInSuccessors;   // `value` is in `liveIn` set
+    BlockSet noLiveInSuccessors; // `value` is not in the `liveIn` set
 
     // Collect successors that do not have `value` in the `liveIn` set.
     for (Block *successor : block.getSuccessors()) {
       const LivenessBlockInfo *succLiveness = liveness.getLiveness(successor);
       if (succLiveness->isLiveIn(value))
         liveInSuccessors.insert(successor);
       else
         noLiveInSuccessors.insert(successor);
     }
 
     // Block has successors with different `liveIn` property of the `value`.
     if (!liveInSuccessors.empty() && !noLiveInSuccessors.empty())
-      divergentLivenessBlocks.insert(&block);
+      divergentLivenessBlocks.try_emplace(&block, noLiveInSuccessors);
   }
 
-  // Verify that divergent `liveIn` property only present in blocks with
-  // async.coro.suspend terminator.
-  for (Block *block : divergentLivenessBlocks) {
+  // Try to insert `dropRef` operations to handle blocks with divergent liveness
+  // in successors blocks.
+  for (auto kv : divergentLivenessBlocks) {
+    Block *block = kv.getFirst();
+    BlockSet &successors = kv.getSecond();
+
+    // Coroutine suspension is a special case terminator for wich we do not
+    // need to create additional reference counting (see details above).
     Operation *terminator = block->getTerminator();
     if (isa<CoroSuspendOp>(terminator))
       continue;
 
-    return terminator->emitOpError("successor have different `liveIn` property "
-                                   "of the reference counted value: ");
+    // We only support successor blocks with empty block argument list.
+    auto hasArgs = [](Block *block) { return !block->getArguments().empty(); };
+    if (llvm::any_of(successors, hasArgs))
+      return terminator->emitOpError()
+             << "successor have different `liveIn` property of the reference "
+                "counted value";
+
+    // Make sure that `dropRef` operation is called when branched into the
+    // successor block without `value` in the `liveIn` set.
+    for (Block *successor : successors) {
+      // If successor has a unique predecessor, it is safe to create `dropRef`
+      // operations directly in the successor block.
+      //
+      // Otherwise we need to create a special block for reference counting
+      // operations, and branch from it to the original successor block.
+      Block *refCountingBlock = nullptr;
+
+      if (successor->getUniquePredecessor() == block) {
+        refCountingBlock = successor;
+      } else {
+        refCountingBlock = &successor->getParent()->emplaceBlock();
+        refCountingBlock->moveBefore(successor);
+        OpBuilder builder = OpBuilder::atBlockEnd(refCountingBlock);
+        builder.create<BranchOp>(value.getLoc(), successor);
+      }
+
+      OpBuilder builder = OpBuilder::atBlockBegin(refCountingBlock);
+      builder.create<RuntimeDropRefOp>(value.getLoc(), value,
+                                       builder.getI32IntegerAttr(1));
+
+      // No need to update the terminator operation.
+      if (successor == refCountingBlock)
+        continue;
+
+      // Update terminator `successor` block to `refCountingBlock`.
+      for (auto pair : llvm::enumerate(terminator->getSuccessors()))
+        if (pair.value() == successor)
+          terminator->setSuccessor(refCountingBlock, pair.index());
+    }
   }
 
   return success();
 }
 
 LogicalResult
 AsyncRuntimeRefCountingPass::addAutomaticRefCounting(Value value) {
   OpBuilder builder(value.getContext());
AsyncRuntimeRefCountingPass::addAutomaticRefCounting(Value value) {
   // Add `drop_ref` operations based on the liveness analysis.
   if (failed(addDropRefAfterLastUse(value)))
     return failure();
 
   // Add `add_ref` operations before function calls.
   if (failed(addAddRefBeforeFunctionCall(value)))
     return failure();
 
-  // Verify that the `value` is in `liveIn` set of all successors.
-  if (failed(verifySuccessors(value)))
+  // Add `drop_ref` operations to successors with divergent `value` liveness.
+  if (failed(addDropRefInDivergentLivenessSuccessor(value)))
     return failure();
 
   return success();
 }
 
 void AsyncRuntimeRefCountingPass::runOnOperation() {
   Operation *op = getOperation();
 

mlir/test/Dialect/Async/async-runtime-ref-counting.mlir

^bb1:
   async.runtime.await %token : !async.token
   %0 = call @cond(): () -> (i1)
   cond_br %0, ^bb1, ^bb2
 ^bb2:
   // CHECK: ^[[BB2:.*]]:
   // CHECK:   return
   return
 }
+
+// CHECK-LABEL: @divergent_liveness_one_token
+func @divergent_liveness_one_token(%arg0 : i1) {
+  // CHECK: %[[TOKEN:.*]] = call @token()
+  %token = call @token() : () -> !async.token
+  // CHECK: cond_br %arg0, ^[[LIVE_IN:.*]], ^[[REF_COUNTING:.*]]
+  cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  // CHECK: ^[[LIVE_IN]]:
+  // CHECK:   async.runtime.await %[[TOKEN]]
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN:.*]]
+  async.runtime.await %token : !async.token
+  br ^bb2
+  // CHECK: ^[[REF_COUNTING:.*]]:
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN:.*]]
+^bb2:
+  // CHECK: ^[[RETURN]]:
+  // CHECK:   return
+  return
+}
+
+// CHECK-LABEL: @divergent_liveness_unique_predecessor
+func @divergent_liveness_unique_predecessor(%arg0 : i1) {
+  // CHECK: %[[TOKEN:.*]] = call @token()
+  %token = call @token() : () -> !async.token
+  // CHECK: cond_br %arg0, ^[[LIVE_IN:.*]], ^[[NO_LIVE_IN:.*]]
+  cond_br %arg0, ^bb2, ^bb1
+^bb1:
+  // CHECK: ^[[NO_LIVE_IN]]:
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN:.*]]
+  br ^bb3
+^bb2:
+  // CHECK: ^[[LIVE_IN]]:
+  // CHECK:   async.runtime.await %[[TOKEN]]
+  // CHECK:   async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN]]
+  async.runtime.await %token : !async.token
+  br ^bb3
+^bb3:
+  // CHECK: ^[[RETURN]]:
+  // CHECK:  return
+  return
+}
+
+// CHECK-LABEL: @divergent_liveness_two_tokens
+func @divergent_liveness_two_tokens(%arg0 : i1) {
+  // CHECK: %[[TOKEN0:.*]] = call @token()
+  // CHECK: %[[TOKEN1:.*]] = call @token()
+  %token0 = call @token() : () -> !async.token
+  %token1 = call @token() : () -> !async.token
+  // CHECK: cond_br %arg0, ^[[AWAIT0:.*]], ^[[AWAIT1:.*]]
+  cond_br %arg0, ^await0, ^await1
+^await0:
+  // CHECK: ^[[AWAIT0]]:
+  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i32}
+  // CHECK:   async.runtime.await %[[TOKEN0]]
+  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN:.*]]
+  async.runtime.await %token0 : !async.token
+  br ^ret
+^await1:
+  // CHECK: ^[[AWAIT1]]:
+  // CHECK:   async.runtime.drop_ref %[[TOKEN0]] {count = 1 : i32}
+  // CHECK:   async.runtime.await %[[TOKEN1]]
+  // CHECK:   async.runtime.drop_ref %[[TOKEN1]] {count = 1 : i32}
+  // CHECK:   br ^[[RETURN]]
+  async.runtime.await %token1 : !async.token
+  br ^ret
+^ret:
+  // CHECK: ^[[RETURN]]:
+  // CHECK:   return
+  return
+}