Refactor AsyncToAsyncRuntime pass to boost understandability.

[bakhtiyarneyman]

Summary

Depends On D106730

[ezhulenev]

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

295–297

Don't see why nested blocks are needed here

296

coro.entry should always have BranchOp terminator, just a cast should be ok here

323

is it needed here?

[ezhulenev]

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

190

As Mehdi mentioned in the previous PR it's better to verify preconditions early, also CondBr is another valid terminator, here it's enough to rewrite only YieldOp, and assume that every other terminator is already verified

[bakhtiyarneyman]

Addressing comments.

[bakhtiyarneyman]

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

295–297

Just so that it's clear to what range of lines the comments apply.

[bakhtiyarneyman]

Address comments.

[bakhtiyarneyman]

Superflous update.

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

struct CoroMachinery {
   //   %token, %result = async.execute -> !async.value<T> {
   //     %0 = constant ... : T
   //     async.yield %0 : T
   //   }
   Value asyncToken; // token representing completion of the async region
   llvm::SmallVector<Value, 4> returnValues; // returned async values
 
   Value coroHandle; // coroutine handle (!async.coro.handle value)
+  Block *entry;     // coroutine entry block
   Block *setError;  // switch completion token and all values to error state
   Block *cleanup;   // coroutine cleanup block
   Block *suspend;   // coroutine suspension block
 };
 } // namespace
 
 /// Utility to partially update the regular function CFG to the coroutine CFG
 /// compatible with LLVM coroutines switched-resume lowering using
-/// `async.runtime.*` and `async.coro.*` operations. Modifies the entry block
-/// by prepending its ops with coroutine setup. Also inserts trailing blocks.
+/// `async.runtime.*` and `async.coro.*` operations. Adds a new entry block
+/// that branches into preexisting entry block. Also inserts trailing blocks.
 ///
 /// The result types of the passed `func` must start with an `async.token`
 /// and be continued with some number of `async.value`s.
 ///
-/// It's up to the caller of this function to fix up the terminators of the
-/// preexisting blocks of the passed func op. If the passed `func` is legal,
-/// this typically means rewriting every return op as a yield op and a branch op
-/// to the suspend block.
+/// The func given to this function needs to have been preprocessed to have
+/// either branch or yield ops as terminators. Branches to the cleanup block are
+/// inserted after each yield.
 ///
 /// See LLVM coroutines documentation: https://llvm.org/docs/Coroutines.html
 ///
 ///  - `entry` block sets up the coroutine.
 ///  - `set_error` block sets completion token and async values state to error.
 ///  - `cleanup` block cleans up the coroutine state.
 ///  - `suspend block after the @llvm.coro.end() defines what value will be
 ///    returned to the initial caller of a coroutine. Everything before the
 ///    @llvm.coro.end() will be executed at every suspension point.
 ///
 /// Coroutine structure (only the important bits):
 ///
 ///   func @some_fn(<function-arguments>) -> (!async.token, !async.value<T>)
 ///   {
 ///     ^entry(<function-arguments>):
 ///       %token = <async token> : !async.token    // create async runtime token
 ///       %value = <async value> : !async.value<T> // create async value
 ///       %id = async.coro.id                      // create a coroutine id
 ///       %hdl = async.coro.begin %id              // create a coroutine handle
-///       /* other ops of the preexisting entry block */
+///       br ^preexisting_entry_block
 ///
-///     /* other preexisting blocks */
+///     /*  preexisting blocks modified to branch to the cleanup block */
 ///
 ///     ^set_error: // this block created lazily only if needed (see code below)
 ///       async.runtime.set_error %token : !async.token
 ///       async.runtime.set_error %value : !async.value<T>
 ///       br ^cleanup
 ///
 ///     ^cleanup:
 ///       async.coro.free %hdl // delete the coroutine state
 ///       br ^suspend
 ///
 ///     ^suspend:
 ///       async.coro.end %hdl // marks the end of a coroutine
 ///       return %token, %value : !async.token, !async.value<T>
 ///   }
 ///
 static CoroMachinery setupCoroMachinery(FuncOp func) {
   assert(!func.getBlocks().empty() && "Function must have an entry block");
 
   MLIRContext *ctx = func.getContext();
   Block *entryBlock = &func.getBlocks().front();
+  Block *originalEntryBlock =
+      entryBlock->splitBlock(entryBlock->getOperations().begin());
   auto builder = ImplicitLocOpBuilder::atBlockBegin(func->getLoc(), entryBlock);
 
   // ------------------------------------------------------------------------ //
   // Allocate async token/values that we will return from a ramp function.
   // ------------------------------------------------------------------------ //
   auto retToken = builder.create<RuntimeCreateOp>(TokenType::get(ctx)).result();
 
   llvm::SmallVector<Value, 4> retValues;
   for (auto resType : func.getCallableResults().drop_front())
     retValues.emplace_back(builder.create<RuntimeCreateOp>(resType).result());
 
   // ------------------------------------------------------------------------ //
   // Initialize coroutine: get coroutine id and coroutine handle.
   // ------------------------------------------------------------------------ //
   auto coroIdOp = builder.create<CoroIdOp>(CoroIdType::get(ctx));
   auto coroHdlOp =
       builder.create<CoroBeginOp>(CoroHandleType::get(ctx), coroIdOp.id());
+  builder.create<BranchOp>(originalEntryBlock);
 
   Block *cleanupBlock = func.addBlock();
   Block *suspendBlock = func.addBlock();
 
   // ------------------------------------------------------------------------ //
   // Coroutine cleanup block: deallocate coroutine frame, free the memory.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(cleanupBlock);
static CoroMachinery setupCoroMachinery(FuncOp func) {
   // This will be the return value of a coroutine ramp function.
   SmallVector<Value, 4> ret{retToken};
   ret.insert(ret.end(), retValues.begin(), retValues.end());
   builder.create<ReturnOp>(ret);
 
   // `async.await` op lowering will create resume blocks for async
   // continuations, and will conditionally branch to cleanup or suspend blocks.
 
+  for (Block &block : func.body().getBlocks()) {
+    if (&block == entryBlock || &block == cleanupBlock ||
+        &block == suspendBlock)
+      continue;
+    Operation *terminator = block.getTerminator();
+    if (auto yield = dyn_cast<YieldOp>(terminator)) {
+      builder.setInsertionPointToEnd(&block);
+      builder.create<BranchOp>(cleanupBlock);
+    }
+  }
+
   CoroMachinery machinery;
   machinery.func = func;
   machinery.asyncToken = retToken;
   machinery.returnValues = retValues;
   machinery.coroHandle = coroHdlOp.handle();
+  machinery.entry = entryBlock;
   machinery.setError = nullptr; // created lazily only if needed
   machinery.cleanup = cleanupBlock;
   machinery.suspend = suspendBlock;
   return machinery;
 }
 
 // Lazily creates `set_error` block only if it is required for lowering to the
 // runtime operations (see for example lowering of assert operation).
outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   auto funcAttrs = ArrayRef<NamedAttribute>();
 
   // TODO: Derive outlined function name from the parent FuncOp (support
   // multiple nested async.execute operations).
   FuncOp func = FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs);
   symbolTable.insert(func);
 
   SymbolTable::setSymbolVisibility(func, SymbolTable::Visibility::Private);
+  auto builder = ImplicitLocOpBuilder::atBlockBegin(loc, func.addEntryBlock());
 
-  // Prepare a function for coroutine lowering by adding entry/cleanup/suspend
-  // blocks, adding async.coro operations and setting up control flow.
-  func.addEntryBlock();
-  CoroMachinery coro = setupCoroMachinery(func);
-
-  // Suspend async function at the end of an entry block, and resume it using
-  // Async resume operation (execution will be resumed in a thread managed by
-  // the async runtime).
-  Block *entryBlock = &func.getBlocks().front();
-  auto builder = ImplicitLocOpBuilder::atBlockEnd(loc, entryBlock);
-
-  // Save the coroutine state: async.coro.save
-  auto coroSaveOp =
-      builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
-
-  // Pass coroutine to the runtime to be resumed on a runtime managed thread.
-  builder.create<RuntimeResumeOp>(coro.coroHandle);
-  builder.create<BranchOp>(coro.cleanup);
-
-  // Split the entry block before the terminator (branch to suspend block).
-  auto *terminatorOp = entryBlock->getTerminator();
-  Block *suspended = terminatorOp->getBlock();
-  Block *resume = suspended->splitBlock(terminatorOp);
-
-  // Add async.coro.suspend as a suspended block terminator.
-  builder.setInsertionPointToEnd(suspended);
-  builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend, resume,
-                                coro.cleanup);
-
+  // Prepare for coroutine conversion by creating the body of the function.
+  {
     size_t numDependencies = execute.dependencies().size();
     size_t numOperands = execute.operands().size();
 
     // Await on all dependencies before starting to execute the body region.
-  builder.setInsertionPointToStart(resume);
     for (size_t i = 0; i < numDependencies; ++i)
       builder.create<AwaitOp>(func.getArgument(i));
 
     // Await on all async value operands and unwrap the payload.
     SmallVector<Value, 4> unwrappedOperands(numOperands);
     for (size_t i = 0; i < numOperands; ++i) {
       Value operand = func.getArgument(numDependencies + i);
       unwrappedOperands[i] = builder.create<AwaitOp>(loc, operand).result();
     }
 
     // Map from function inputs defined above the execute op to the function
     // arguments.
     BlockAndValueMapping valueMapping;
     valueMapping.map(functionInputs, func.getArguments());
     valueMapping.map(execute.body().getArguments(), unwrappedOperands);
 
     // Clone all operations from the execute operation body into the outlined
     // function body.
     for (Operation &op : execute.body().getOps())
       builder.clone(op, valueMapping);
+  }
+
+  // Adding entry/cleanup/suspend blocks.
+  CoroMachinery coro = setupCoroMachinery(func);
+
+  // Suspend async function at the end of an entry block, and resume it using
+  // Async resume operation (execution will be resumed in a thread managed by
+  // the async runtime).
+  {
+    BranchOp branch = cast<BranchOp>(coro.entry->getTerminator());
+    builder.setInsertionPointToEnd(coro.entry);
+
+    // Save the coroutine state: async.coro.save
+    auto coroSaveOp =
+        builder.create<CoroSaveOp>(CoroStateType::get(ctx), coro.coroHandle);
+
+    // Pass coroutine to the runtime to be resumed on a runtime managed
+    // thread.
+    builder.create<RuntimeResumeOp>(coro.coroHandle);
+
+    // Add async.coro.suspend as a suspended block terminator.
+    builder.create<CoroSuspendOp>(coroSaveOp.state(), coro.suspend,
+                                  branch.getDest(), coro.cleanup);
+
+    branch.erase();
+  }
 
   // Replace the original `async.execute` with a call to outlined function.
+  {
     ImplicitLocOpBuilder callBuilder(loc, execute);
     auto callOutlinedFunc = callBuilder.create<CallOp>(
         func.getName(), execute.getResultTypes(), functionInputs.getArrayRef());
     execute.replaceAllUsesWith(callOutlinedFunc.getResults());
     execute.erase();
+  }
 
   return {func, coro};
 }
 
 //===----------------------------------------------------------------------===//
 // Convert async.create_group operation to async.runtime.create_group
 //===----------------------------------------------------------------------===//
 
 namespace {
 class CreateGroupOpLowering : public OpConversionPattern<CreateGroupOp> {
static CoroMachinery rewriteFuncAsCoroutine(FuncOp func) {
   auto *ctx = func->getContext();
   auto loc = func.getLoc();
   SmallVector<Type> resultTypes;
   resultTypes.reserve(func.getCallableResults().size());
   llvm::transform(func.getCallableResults(), std::back_inserter(resultTypes),
                   [](Type type) { return ValueType::get(type); });
   func.setType(FunctionType::get(ctx, func.getType().getInputs(), resultTypes));
   func.insertResult(0, TokenType::get(ctx), {});
-  CoroMachinery coro = setupCoroMachinery(func);
   for (Block &block : func.getBlocks()) {
-    if (&block == coro.suspend)
-      continue;
-
     Operation *terminator = block.getTerminator();
     if (auto returnOp = dyn_cast<ReturnOp>(*terminator)) {
       ImplicitLocOpBuilder builder(loc, returnOp);
       builder.create<YieldOp>(returnOp.getOperands());
-      builder.create<BranchOp>(coro.cleanup);
       returnOp.erase();
     }
   }
-  return coro;
+  return setupCoroMachinery(func);
 }
 
 /// Rewrites a call into a function that has been rewritten as a coroutine.
 ///
 /// The invocation of this function is safe only when call ops are traversed in
 /// reverse order of how they appear in a single block. See `funcsToCoroutines`.
 static void rewriteCallsiteForCoroutine(CallOp oldCall, FuncOp func) {
   auto loc = func.getLoc();

mlir/test/Dialect/Async/async-to-async-runtime-eliminate-blocking.mlir

 // RUN: mlir-opt %s -split-input-file \
 // RUN:    -async-to-async-runtime="eliminate-blocking-await-ops=true" \
 // RUN:   | FileCheck %s --dump-input=always
 
 // CHECK-LABEL: func @simple_callee
 // CHECK-SAME: (%[[ARG:.*]]: f32)
 // CHECK-SAME: -> (!async.token, !async.value<f32> {builtin.foo = "bar"})
 func @simple_callee(%arg0: f32) -> (f32 {builtin.foo = "bar"}) {
 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
 // CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
 // CHECK: %[[ID:.*]] = async.coro.id
 // CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
-
-// CHECK: %[[VAL:.*]] = addf %[[ARG]], %[[ARG]] : f32
+// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
+// CHECK  ^[[ORIGINAL_ENTRY]]:
+// CHECK:   %[[VAL:.*]] = addf %[[ARG]], %[[ARG]] : f32
   %0 = addf %arg0, %arg0 : f32
-// CHECK: %[[VAL_STORAGE:.*]] = async.runtime.create : !async.value<f32>
+// CHECK:   %[[VAL_STORAGE:.*]] = async.runtime.create : !async.value<f32>
   %1 = async.runtime.create: !async.value<f32>
-// CHECK: async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : !async.value<f32>
+// CHECK:   async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : !async.value<f32>
   async.runtime.store %0, %1: !async.value<f32>
-// CHECK: async.runtime.set_available %[[VAL_STORAGE]] : !async.value<f32>
+// CHECK:   async.runtime.set_available %[[VAL_STORAGE]] : !async.value<f32>
   async.runtime.set_available %1: !async.value<f32>
 
-// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]]
-// CHECK: async.runtime.await_and_resume %[[VAL_STORAGE]], %[[HDL]]
-// CHECK: async.coro.suspend %[[SAVED]]
+// CHECK:   %[[SAVED:.*]] = async.coro.save %[[HDL]]
+// CHECK:   async.runtime.await_and_resume %[[VAL_STORAGE]], %[[HDL]]
+// CHECK:   async.coro.suspend %[[SAVED]]
 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
   %2 = async.await %1 : !async.value<f32>
 
 // CHECK: ^[[RESUME]]:
 // CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[VAL_STORAGE]] : !async.value<f32>
 // CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]
 
 // CHECK: ^[[BRANCH_OK]]:

 
 // CHECK-LABEL: func @simple_caller()
 // CHECK-SAME: -> (!async.token, !async.value<f32>)
 func @simple_caller() -> f32 {
 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
 // CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
 // CHECK: %[[ID:.*]] = async.coro.id
 // CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
+// CHECK  ^[[ORIGINAL_ENTRY]]:
 
-// CHECK: %[[CONSTANT:.*]] = constant
+// CHECK:   %[[CONSTANT:.*]] = constant
   %c = constant 1.0 : f32
-// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
-// CHECK: %[[SAVED:.*]] = async.coro.save %[[HDL]]
-// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]]#0, %[[HDL]]
-// CHECK: async.coro.suspend %[[SAVED]]
+// CHECK:   %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK:   %[[SAVED:.*]] = async.coro.save %[[HDL]]
+// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER]]#0, %[[HDL]]
+// CHECK:   async.coro.suspend %[[SAVED]]
 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
   %r = call @simple_callee(%c): (f32) -> f32
 
 // CHECK: ^[[RESUME]]:
 // CHECK:   %[[IS_TOKEN_ERROR:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER]]#0 : !async.token
 // CHECK:   cond_br %[[IS_TOKEN_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK:.*]]
 
 // CHECK: ^[[BRANCH_TOKEN_OK]]:

 
 // CHECK-LABEL: func @double_caller()
 // CHECK-SAME: -> (!async.token, !async.value<f32>)
 func @double_caller() -> f32 {
 // CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
 // CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
 // CHECK: %[[ID:.*]] = async.coro.id
 // CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
+// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
+// CHECK  ^[[ORIGINAL_ENTRY]]:
 
-// CHECK: %[[CONSTANT:.*]] = constant
+// CHECK:   %[[CONSTANT:.*]] = constant
   %c = constant 1.0 : f32
-// CHECK: %[[RETURNED_TO_CALLER_1:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
-// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
-// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER_1]]#0, %[[HDL]]
-// CHECK: async.coro.suspend %[[SAVED_1]]
+// CHECK:   %[[RETURNED_TO_CALLER_1:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK:   %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
+// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER_1]]#0, %[[HDL]]
+// CHECK:   async.coro.suspend %[[SAVED_1]]
 // CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
   %r = call @simple_callee(%c): (f32) -> f32
 
 // CHECK: ^[[RESUME_1]]:
 // CHECK:   %[[IS_TOKEN_ERROR_1:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_1]]#0 : !async.token
 // CHECK:   cond_br %[[IS_TOKEN_ERROR_1]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK_1:.*]]
 
 // CHECK: ^[[BRANCH_TOKEN_OK_1]]:

mlir/test/Dialect/Async/async-to-async-runtime.mlir

 // CHECK:   addf %[[LOADED]], %[[LOADED]] : f32
 // CHECK:   async.runtime.set_available %[[TOKEN]]
 
 // CHECK: ^[[CLEANUP]]:
 // CHECK: ^[[SUSPEND]]:
 
 // -----
 
-// CHECK-LABEL: @execute_asserttion
-func @execute_asserttion(%arg0: i1) {
+// CHECK-LABEL: @execute_assertion
+func @execute_assertion(%arg0: i1) {
   %token = async.execute {
     assert %arg0, "error"
     async.yield
   }
   async.await %token : !async.token
   return
 }