[mlir] Use rewriting infrastructure in AsyncToLLVM

[tpopp]

Summary

This is needed so a listener hears all changes during the dialect
conversion to allow correct rollbacks upon failure.

[tpopp]

Pass an OpBuilder instead of an optional Listener.

[ezhulenev]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

389–390

This is a bit misleading after the change. It would be nice to explain what exactly is the op and how splitBlock and resume were constructed (also I'd rename splitBlock to suspended).

Something along this line: "Add a LLVM coroutine suspension point to the end of suspended block, to resume execution in resume block. These blocks are the result of splitting blah blah"

[ezhulenev]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

395–412

This example also will be out of date.

[tpopp]

Update comments.

[ftynse]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

518

Builder will not be aware of this operation.

[tpopp]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

518

That's okay because this function is not using the DialectConversion infrastructure. This function uses normal OpBuilders and is an irreversible transformation.

[ftynse]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

518

Maybe add a note to the doc that indicates this. It could end up confusing somebody that some functions are rewriter-compatible and some others are not.

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

static CoroMachinery setupCoroMachinery(FuncOp func) {
 
   // `async.await` op lowering will create resume blocks for async
   // continuations, and will conditionally branch to cleanup or suspend blocks.
 
   return {createToken.getResult(0), coroHdl.getResult(0), cleanupBlock,
           suspendBlock};
 }
 
 // Adds a suspension point before the `op`, and moves `op` and all operations
 // after it into the resume block. Returns a pointer to the resume block.
 //
 // `coroState` must be a value returned from the call to @llvm.coro.save(...)
 // intrinsic (saved coroutine state).
 //
 // Before:
 //
 //   ^bb0:
 //     "opBefore"(...)
 //     "op"(...)
 //   ^cleanup: ...
 //   ^suspend: ...
 //
 // After:
 //
 //   ^bb0:
 //     "opBefore"(...)
 //     %suspend = llmv.call @llvm.coro.suspend(...)
 //     switch %suspend [-1: ^suspend, 0: ^resume, 1: ^cleanup]
 //   ^resume:
 //     "op"(...)
 //   ^cleanup: ...
 //   ^suspend: ...
 //
-static Block *addSuspensionPoint(CoroMachinery coro, Value coroState,
-                                 Operation *op) {
+static void addSuspensionPoint(CoroMachinery coro, Value coroState,
+                               Operation *op, Block *splitBlock, Block *resume,
+                               OpBuilder &builder) {
+  Location loc = op->getLoc();
   MLIRContext *ctx = op->getContext();
   auto i1 = LLVM::LLVMType::getInt1Ty(ctx);
   auto i8 = LLVM::LLVMType::getInt8Ty(ctx);
 
-  Location loc = op->getLoc();
-  Block *splitBlock = op->getBlock();
-
-  // Split the block before `op`, newly added block is the resume block.
-  Block *resume = splitBlock->splitBlock(op);
-
   // Add a coroutine suspension in place of original `op` in the split block.
-  OpBuilder builder = OpBuilder::atBlockEnd(splitBlock);
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToEnd(splitBlock);
 
   auto constFalse =
       builder.create<LLVM::ConstantOp>(loc, i1, builder.getBoolAttr(false));
 
   // Suspend a coroutine: @llvm.coro.suspend
   auto coroSuspend = builder.create<LLVM::CallOp>(
       loc, i8, builder.getSymbolRefAttr(kCoroSuspend),
       ValueRange({coroState, constFalse}));
builder.create<LLVM::CondBrOp>(loc, isNegOne, /*trueDest=*/coro.suspend,
                                  /*falseDest=*/resumeOrCleanup);
 
   // ... or resume or cleanup the coroutine?
   builder.setInsertionPointToStart(resumeOrCleanup);
   auto isZero = builder.create<LLVM::ICmpOp>(
       loc, LLVM::ICmpPredicate::eq, coroSuspend.getResult(0), constZero);
   builder.create<LLVM::CondBrOp>(loc, isZero, /*trueDest=*/resume,
                                  /*falseDest=*/coro.cleanup);
-
-  return resume;
 }
 
 // Outline the body region attached to the `async.execute` op into a standalone
 // function.
 static std::pair<FuncOp, CoroMachinery>
 outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   ModuleOp module = execute.getParentOfType<ModuleOp>();
 
auto coroSave = builder.create<LLVM::CallOp>(
       loc, LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
       ValueRange({coro.coroHandle}));
 
   // Call async runtime API to execute a coroutine in the managed thread.
   SmallVector<Value, 2> executeArgs = {coro.coroHandle, resumePtr.res()};
   builder.create<CallOp>(loc, TypeRange(), kExecute, executeArgs);
 
   // Split the entry block before the terminator.
-  Block *resume = addSuspensionPoint(coro, coroSave.getResult(0),
-                                     entryBlock->getTerminator());
+  auto *terminatorOp = entryBlock->getTerminator();
+  Block *splitBlock = terminatorOp->getBlock();
+  Block *resume = splitBlock->splitBlock(terminatorOp);
+  addSuspensionPoint(coro, coroSave.getResult(0), terminatorOp, splitBlock,
+                     resume, builder);
 
   // Await on all dependencies before starting to execute the body region.
   builder.setInsertionPointToStart(resume);
   for (size_t i = 0; i < execute.dependencies().size(); ++i)
     builder.create<AwaitOp>(loc, func.getArgument(i));
 
   // Map from function inputs defined above the execute op to the function
   // arguments.
if (!isInCoroutine)
       rewriter.create<CallOp>(loc, TypeRange(), blockingAwaitFuncName,
                               ValueRange(operands[0]));
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       const CoroMachinery &coro = outlined->getSecond();
 
-      OpBuilder builder(op);
+      OpBuilder builder(op, rewriter.getListener());
       MLIRContext *ctx = op->getContext();
 
       // A pointer to coroutine resume intrinsic wrapper.
       auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
       auto resumePtr = builder.create<LLVM::AddressOfOp>(
           loc, resumeFnTy.getPointerTo(), kResume);
 
       // Save the coroutine state: @llvm.coro.save
       auto coroSave = builder.create<LLVM::CallOp>(
           loc, LLVM::LLVMTokenType::get(ctx),
           builder.getSymbolRefAttr(kCoroSave), ValueRange(coro.coroHandle));
 
       // Call async runtime API to resume a coroutine in the managed thread when
       // the async await argument becomes ready.
       SmallVector<Value, 3> awaitAndExecuteArgs = {operands[0], coro.coroHandle,
                                                    resumePtr.res()};
       builder.create<CallOp>(loc, TypeRange(), coroAwaitFuncName,
                              awaitAndExecuteArgs);
 
+      Block *splitBlock = op->getBlock();
+
       // Split the entry block before the await operation.
-      addSuspensionPoint(coro, coroSave.getResult(0), op);
+      Block *resume = rewriter.splitBlock(splitBlock, Block::iterator(op));
+      addSuspensionPoint(coro, coroSave.getResult(0), op, splitBlock, resume,
+                         builder);
     }
 
     // Original operation was replaced by function call or suspension point.
     rewriter.eraseOp(op);
 
     return success();
   }