Add async_funcs_only option to AsyncToAsyncRuntime pass

[yijia1212]

Summary

This change adds async_funcs_only option to AsyncToAsyncRuntimePass. The goal is to convert async functions to regular functions in early stages of compilation pipeline.

[ezhulenev]

Instead of adding an option to a pass, can you add a populate... function (similar to https://github.com/llvm/llvm-project/blob/2e999b7dd1934a44d38c3a753460f1e5a217e9a5/mlir/include/mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h#L22-L25) that will add all patterns to the rewriter.

Users that need to lower just the async.func can write a custom pass with their own ConversionTarget that will only mark async.func as illegal. That's a more idiomatic MLIR way to parametrize passes.

Also having populateSCFToControlFlowConversionPatterns here is a mistake. If some users require mixing lowering of async and lowering to CFG, again they should use populate... and a custom pass. Let's try to remove it from this pass (in future PRs), I don't think that any tests in MLIR itself depend on this behavior.

[yijia1212]

Handle reviews. Creating a separate pass to lower async.func to coroutine

[ezhulenev]

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

44 ↗

(On Diff #478765)

Lower all hight level ...

50 ↗

(On Diff #478765)

nit: async-func-to-async-runtime for consistency

51 ↗

(On Diff #478765)

Lower async.func

[yijia1212]

Address reviews

mlir/include/mlir/Dialect/Async/Passes.h

 #include "mlir/Dialect/Async/Passes.h.inc"
 
 std::unique_ptr<Pass> createAsyncParallelForPass();
 
 std::unique_ptr<Pass> createAsyncParallelForPass(bool asyncDispatch,
                                                  int32_t numWorkerThreads,
                                                  int32_t minTaskSize);
 
-std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
+std::unique_ptr<OperationPass<ModuleOp>>
+createAsyncToAsyncRuntimePass(bool async_funcs_only = false);
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
 std::unique_ptr<Pass> createAsyncRuntimePolicyBasedRefCountingPass();
 
 //===----------------------------------------------------------------------===//

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

 // Prefix for functions outlined from `async.execute` op regions.
 static constexpr const char kAsyncFnPrefix[] = "async_execute_fn";
 
 namespace {
 
 class AsyncToAsyncRuntimePass
     : public impl::AsyncToAsyncRuntimeBase<AsyncToAsyncRuntimePass> {
 public:
-  AsyncToAsyncRuntimePass() = default;
+  AsyncToAsyncRuntimePass(bool async_funcs_only)
+      : async_funcs_only_(async_funcs_only) {}
   void runOnOperation() override;
+
+private:
+  bool async_funcs_only_;
 };
 
 } // namespace
 
 /// Function targeted for coroutine transformation has two additional blocks at
 /// the end: coroutine cleanup and coroutine suspension.
 ///
 /// async.await op lowering additionaly creates a resume block for each

 void AsyncToAsyncRuntimePass::runOnOperation() {
   ModuleOp module = getOperation();
   SymbolTable symbolTable(module);
 
   // Functions with coroutine CFG setups, which are results of outlining
   // `async.execute` body regions and converting async.func.
   llvm::DenseMap<func::FuncOp, CoroMachinery> coros;
 
+  if (!async_funcs_only_)
     module.walk([&](ExecuteOp execute) {
       coros.insert(outlineExecuteOp(symbolTable, execute));
     });
 
   LLVM_DEBUG({
     llvm::dbgs() << "Outlined " << coros.size()
                  << " functions built from async.execute operations\n";
   });
 
   // Returns true if operation is inside the coroutine.
   auto isInCoroutine = [&](Operation *op) -> bool {
     auto parentFunc = op->getParentOfType<func::FuncOp>();
     return coros.find(parentFunc) != coros.end();
   };
 
   // Lower async operations to async.runtime operations.
   MLIRContext *ctx = module->getContext();
   RewritePatternSet asyncPatterns(ctx);
 
   // Conversion to async runtime augments original CFG with the coroutine CFG,
   // and we have to make sure that structured control flow operations with async
   // operations in nested regions will be converted to branch-based control flow
   // before we add the coroutine basic blocks.
   populateSCFToControlFlowConversionPatterns(asyncPatterns);
 
   // Async lowering does not use type converter because it must preserve all
   // types for async.runtime operations.
+  if (!async_funcs_only_) {
     asyncPatterns.add<CreateGroupOpLowering, AddToGroupOpLowering>(ctx);
-
-  // Lower async.func to func.func with coroutine cfg.
-  asyncPatterns.add<AsyncCallOpLowering>(ctx);
-  asyncPatterns.add<AsyncFuncOpLowering, AsyncReturnOpLowering>(ctx, coros);
-
     asyncPatterns.add<AwaitTokenOpLowering, AwaitValueOpLowering,
                       AwaitAllOpLowering, YieldOpLowering>(ctx, coros);
 
     // Lower assertions to conditional branches into error blocks.
     asyncPatterns.add<AssertOpLowering>(ctx, coros);
+  }
+  // Lower async.func to func.func with coroutine cfg.
+  asyncPatterns.add<AsyncCallOpLowering>(ctx);
+  asyncPatterns.add<AsyncFuncOpLowering, AsyncReturnOpLowering>(ctx, coros);
 
   // All high level async operations must be lowered to the runtime operations.
   ConversionTarget runtimeTarget(*ctx);
   runtimeTarget.addLegalDialect<AsyncDialect, func::FuncDialect>();
+  if (!async_funcs_only_) {
     runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>();
-  runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp,
-                             async::FuncOp, async::CallOp, async::ReturnOp>();
+    runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, YieldOp>();
+  }
+
+  runtimeTarget.addIllegalOp<async::FuncOp, async::CallOp, async::ReturnOp>();
 
   // Decide if structured control flow has to be lowered to branch-based CFG.
   runtimeTarget.addDynamicallyLegalDialect<scf::SCFDialect>([&](Operation *op) {
     auto walkResult = op->walk([&](Operation *nested) {
       bool isAsync = isa<async::AsyncDialect>(nested->getDialect());
       return isAsync && isInCoroutine(nested) ? WalkResult::interrupt()
                                               : WalkResult::advance();
     });
-    return !walkResult.wasInterrupted();
+    return async_funcs_only_ || !walkResult.wasInterrupted();
   });
   runtimeTarget.addLegalOp<cf::AssertOp, arith::XOrIOp, arith::ConstantOp,
                            func::ConstantOp, cf::BranchOp, cf::CondBranchOp>();
 
   // Assertions must be converted to runtime errors inside async functions.
   runtimeTarget.addDynamicallyLegalOp<cf::AssertOp>(
       [&](cf::AssertOp op) -> bool {
         auto func = op->getParentOfType<func::FuncOp>();
-        return coros.find(func) == coros.end();
+        return async_funcs_only_ || coros.find(func) == coros.end();
       });
 
   if (failed(applyPartialConversion(module, runtimeTarget,
                                     std::move(asyncPatterns)))) {
     signalPassFailure();
     return;
   }
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createAsyncToAsyncRuntimePass() {
-  return std::make_unique<AsyncToAsyncRuntimePass>();
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createAsyncToAsyncRuntimePass(bool async_funcs_only) {
+  return std::make_unique<AsyncToAsyncRuntimePass>(async_funcs_only);
 }