Add an escape-hatch for conversion of funcs with blocking awaits to coroutines.

[bakhtiyarneyman]

Summary

Currently TFRT does not support top-level coroutines, so this functionality will allow to have a single blocking await at the top level until TFRT implements the necessary functionality.

[mehdi_amini]

I feel that what is missing from these two revisions is a more contextual description of the motivation and what we're trying to achieve and how, I don't quite have a good understanding of the design just yet.

[ezhulenev]

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

635–637

Just & is enough here, unless lambda captures by value don't see a need to specify ref captures explicitly

[bakhtiyarneyman]

Addressing comments.

[mehdi_amini]

Currently TFRT does not support top-level coroutines, so this functionality will allow to have a single blocking await at the top level until TFRT implements the necessary functionality.

I suspect that in general this should be handled using visibility: you should only be able to rewrite the signature of private functions and not public one (which are by definition exported by the module to the outside world): would that be enough to handle the TFRT case?

[ezhulenev]

We'll need to rewrite public functions at some point as well (after resolving all the issues with TFRT threading). Alternatively instead of passing list of function names it can be a boolean flag allow-rewriting-public-functions.

[mehdi_amini]

In D106730#2907908, @ezhulenev wrote:

We'll need to rewrite public functions at some point as well (after resolving all the issues with TFRT threading). Alternatively instead of passing list of function names it can be a boolean flag allow-rewriting-public-functions.

Yeah I'd rather have the boolean flag as an interface here right now, thanks!

[bakhtiyarneyman]

Addressing comments.

[bakhtiyarneyman]

In D106730#2907908, @ezhulenev wrote:

We'll need to rewrite public functions at some point as well (after resolving all the issues with TFRT threading). Alternatively instead of passing list of function names it can be a boolean flag allow-rewriting-public-functions.

I think this creates an unnecessary and rigid coupling between privacy and asynchronicity, so I'm not a fan of this proposal.

[ezhulenev]

Not to mix visibility and asynchronicity I'd prefer to have a special attribute on functions that should not become async, this will give control to the caller who knows best (e.g. in tfrt currently entrypoint will be marked sync-only).

[mehdi_amini]

I like this!

[bakhtiyarneyman]

Use attributes as opposed to command line options.

[mehdi_amini]

mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td

37

FYI: in C++14 (pre 17) you need a definition in the .cpp file as well to avoid link-time failure in some cases.

[bakhtiyarneyman]

Addressing comments.

[bakhtiyarneyman]

Address comments.

[ezhulenev]

Fix compilation

mlir/include/mlir/Dialect/Async/IR/AsyncDialect.td

def AsyncDialect : Dialect {
   let name = "async";
 
   let summary = "Types and operations for async dialect";
   let description = [{
     This dialect contains operations for modeling asynchronous execution.
   }];
 
   let cppNamespace = "::mlir::async";
+
+  let extraClassDeclaration = [{
+    // The name of a unit attribute on funcs that are allowed to have a blocking
+    // async.runtime.await ops. Only useful in combination with
+    // 'eliminate-blocking-await-ops' option, which in absence of this attribute
+    // might convert a func to a coroutine.
+    static constexpr StringRef kAllowedToBlockAttrName = "async.allowed_to_block";
+  }];
+
 }
 
 #endif // ASYNC_DIALECT_TD

mlir/include/mlir/Dialect/Async/Passes.td

def AsyncToAsyncRuntime : Pass<"async-to-async-runtime", "ModuleOp"> {
   let summary = "Lower high level async operations (e.g. async.execute) to the"
                 "explicit async.runtime and async.coro operations";
   let constructor = "mlir::createAsyncToAsyncRuntimePass()";
   let options = [
     // Temporary for bringup, should become the default.
     Option<"eliminateBlockingAwaitOps", "eliminate-blocking-await-ops", "bool",
             /*default=*/"false",
            "Rewrite functions with blocking async.runtime.await as coroutines "
-           "with async.runtime.await_and_resume.">
+           "with async.runtime.await_and_resume.">,
   ];
   let dependentDialects = ["async::AsyncDialect"];
 }
 
 def AsyncRuntimeRefCounting : Pass<"async-runtime-ref-counting"> {
   let summary = "Automatic reference counting for Async runtime operations";
   let description = [{
     This pass works at the async runtime abtraction level, after all

mlir/lib/Dialect/Async/IR/Async.cpp

 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::async;
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.cpp.inc"
 
+constexpr StringRef AsyncDialect::kAllowedToBlockAttrName;
+
 void AsyncDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"
       >();
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/Async/IR/AsyncOpsTypes.cpp.inc"

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

for (Value result : newCall.getResults().drop_front())
     unwrappedResults.push_back(
         callBuilder.create<AwaitOp>(loc, result).result());
   // Careful, when result of a call is piped into another call this could lead
   // to a dangling pointer.
   oldCall.replaceAllUsesWith(unwrappedResults);
   oldCall.erase();
 }
 
+static bool isAllowedToBlock(FuncOp func) {
+  return !!func->getAttrOfType<UnitAttr>(AsyncDialect::kAllowedToBlockAttrName);
+}
+
 static LogicalResult
 funcsToCoroutines(ModuleOp module,
                   llvm::DenseMap<FuncOp, CoroMachinery> &outlinedFunctions) {
   // The following code supports the general case when 2 functions mutually
   // recurse into each other. Because of this and that we are relying on
   // SymbolUserMap to find pointers to calling FuncOps, we cannot simply erase
   // a FuncOp while inserting an equivalent coroutine, because that could lead
   // to dangling pointers.
 
   SmallVector<FuncOp> funcWorklist;
 
   // Careful, it's okay to add a func to the worklist multiple times if and only
   // if the loop processing the worklist will skip the functions that have
   // already been converted to coroutines.
-  auto addToWorklist = [&outlinedFunctions, &funcWorklist](FuncOp func) {
+  auto addToWorklist = [&](FuncOp func) {
+    if (isAllowedToBlock(func))
+      return;
     // N.B. To refactor this code into a separate pass the lookup in
     // outlinedFunctions is the most obvious obstacle. Looking at an arbitrary
     // func and recognizing if it has a coroutine structure is messy. Passing
     // this dict between the passes is ugly.
-    if (outlinedFunctions.find(func) == outlinedFunctions.end()) {
+    if (isAllowedToBlock(func) ||
+        outlinedFunctions.find(func) == outlinedFunctions.end()) {
       for (Operation &op : func.body().getOps()) {
         if (dyn_cast<AwaitOp>(op) || dyn_cast<AwaitAllOp>(op)) {
           funcWorklist.push_back(func);
           break;
         }
       }
     }
   };
void AsyncToAsyncRuntimePass::runOnOperation() {
 
   // Assertions must be converted to runtime errors inside async functions.
   runtimeTarget.addDynamicallyLegalOp<AssertOp>([&](AssertOp op) -> bool {
     auto func = op->getParentOfType<FuncOp>();
     return outlinedFunctions.find(func) == outlinedFunctions.end();
   });
 
   if (eliminateBlockingAwaitOps)
-    runtimeTarget.addIllegalOp<RuntimeAwaitOp>();
+    runtimeTarget.addDynamicallyLegalOp<RuntimeAwaitOp>(
+        [&](RuntimeAwaitOp op) -> bool {
+          return isAllowedToBlock(op->getParentOfType<FuncOp>());
+        });
 
   if (failed(applyPartialConversion(module, runtimeTarget,
                                     std::move(asyncPatterns)))) {
     signalPassFailure();
     return;
   }
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() {
   return std::make_unique<AsyncToAsyncRuntimePass>();
 }

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

 // CHECK: ^[[CLEANUP]]:
 // CHECK:   async.coro.free %[[ID]], %[[HDL]]
 // CHECK:   br ^[[SUSPEND]]
 
 // CHECK: ^[[SUSPEND]]:
 // CHECK:   async.coro.end %[[HDL]]
 // CHECK:   return %[[TOKEN]] : !async.token
 }
+
+// CHECK-LABEL: func @caller_allowed_to_block
+// CHECK-SAME: () -> f32
+func @caller_allowed_to_block() -> f32 attributes { async.allowed_to_block } {
+// CHECK: %[[CONSTANT:.*]] = constant
+  %c = constant 1.0 : f32
+// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
+// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#0
+// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#1
+// CHECK: %[[RETURNED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1
+  %r = call @simple_callee(%c): (f32) -> f32
+
+// CHECK:   return %[[RETURNED]] : f32
+  return %r: f32
+}