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

43–44

Lower all hight level ...

50

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

51

Lower async.func

[yijia1212]

Address reviews

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

 
 #ifndef MLIR_DIALECT_ASYNC_PASSES_H_
 #define MLIR_DIALECT_ASYNC_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 class ModuleOp;
+class ConversionTarget;
 
 #define GEN_PASS_DECL
 #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);
 
+void populateAsyncFuncToAsyncRuntimeConversionPatterns(
+    RewritePatternSet &patterns, ConversionTarget &target);
+
+std::unique_ptr<OperationPass<ModuleOp>> createAsyncFuncToAsyncRuntimePass();
+
 std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
 std::unique_ptr<Pass> createAsyncRuntimePolicyBasedRefCountingPass();
 

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

def AsyncParallelFor : Pass<"async-parallel-for", "ModuleOp"> {
 
   let dependentDialects = [
     "arith::ArithDialect",
     "async::AsyncDialect",
     "scf::SCFDialect"
   ];
 }
 
 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 summary = "Lower all high level async operations (e.g. async.execute) to"
+                "the explicit async.runtime and async.coro operations";
   let constructor = "mlir::createAsyncToAsyncRuntimePass()";
   let dependentDialects = ["async::AsyncDialect", "func::FuncDialect"];
 }
 
+def AsyncFuncToAsyncRuntime : Pass<"async-func-to-async-runtime", "ModuleOp"> {
+  let summary = "Lower async.func operations to the explicit async.runtime and"
+                "async.coro operations";
+  let constructor = "mlir::createAsyncFuncToAsyncRuntimePass()";
+  let dependentDialects = ["async::AsyncDialect", "func::FuncDialect"];
+}
+
 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
     `async.execute` and `async.await` operations are lowered to the async
     runtime API calls, and async coroutine operations.
 
     It relies on the LLVM coroutines switched-resume lowering semantics for

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

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_ASYNCTOASYNCRUNTIME
+#define GEN_PASS_DEF_ASYNCFUNCTOASYNCRUNTIME
 #include "mlir/Dialect/Async/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::async;
 
 #define DEBUG_TYPE "async-to-async-runtime"
 // 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;
   void runOnOperation() override;
 };
 
 } // namespace
 
+namespace {
+
+class AsyncFuncToAsyncRuntimePass
+    : public impl::AsyncFuncToAsyncRuntimeBase<AsyncFuncToAsyncRuntimePass> {
+public:
+  AsyncFuncToAsyncRuntimePass() = default;
+  void runOnOperation() override;
+};
+
+} // 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
 /// operation to enable non-blocking waiting via coroutine suspension.
 namespace {
 struct CoroMachinery {
   func::FuncOp func;
struct CoroMachinery {
   Value coroHandle; // coroutine handle (!async.coro.getHandle value)
   Block *entry;     // coroutine entry block
   Optional<Block *> setError; // set returned values to error state
   Block *cleanup;   // coroutine cleanup block
   Block *suspend;   // coroutine suspension block
 };
 } // namespace
 
+using FuncCoroMapPtr =
+    std::shared_ptr<llvm::DenseMap<func::FuncOp, CoroMachinery>>;
+
 /// 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. Adds a new entry block
 /// that branches into preexisting entry block. Also inserts trailing blocks.
 ///
 /// The result types of the passed `func` start with an optional `async.token`
 /// and be continued with some number of `async.value`s.
 ///

 namespace {
 
 //===----------------------------------------------------------------------===//
 // Convert async.func operation to func.func
 //===----------------------------------------------------------------------===//
 
 class AsyncFuncOpLowering : public OpConversionPattern<async::FuncOp> {
 public:
-  AsyncFuncOpLowering(MLIRContext *ctx,
-                      llvm::DenseMap<func::FuncOp, CoroMachinery> &coros)
-      : OpConversionPattern<async::FuncOp>(ctx), coros(coros) {}
+  AsyncFuncOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
+      : OpConversionPattern<async::FuncOp>(ctx), coros_(coros) {}
 
   LogicalResult
   matchAndRewrite(async::FuncOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
 
     auto newFuncOp =
         rewriter.create<func::FuncOp>(loc, op.getName(), op.getFunctionType());
 
     SymbolTable::setSymbolVisibility(newFuncOp,
                                      SymbolTable::getSymbolVisibility(op));
     // Copy over all attributes other than the name.
     for (const auto &namedAttr : op->getAttrs()) {
       if (namedAttr.getName() != SymbolTable::getSymbolAttrName())
         newFuncOp->setAttr(namedAttr.getName(), namedAttr.getValue());
     }
 
     rewriter.inlineRegionBefore(op.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
 
     CoroMachinery coro = setupCoroMachinery(newFuncOp);
-    coros[newFuncOp] = coro;
+    (*coros_)[newFuncOp] = coro;
     // no initial suspend, we should hot-start
 
     rewriter.eraseOp(op);
     return success();
   }
 
 private:
-  llvm::DenseMap<func::FuncOp, CoroMachinery> &coros;
+  FuncCoroMapPtr coros_;
 };
 
 //===----------------------------------------------------------------------===//
 // Convert async.call operation to func.call
 //===----------------------------------------------------------------------===//
 
 class AsyncCallOpLowering : public OpConversionPattern<async::CallOp> {
 public:

 };
 
 //===----------------------------------------------------------------------===//
 // Convert async.return operation to async.runtime operations.
 //===----------------------------------------------------------------------===//
 
 class AsyncReturnOpLowering : public OpConversionPattern<async::ReturnOp> {
 public:
-  AsyncReturnOpLowering(MLIRContext *ctx,
-                        llvm::DenseMap<func::FuncOp, CoroMachinery> &coros)
-      : OpConversionPattern<async::ReturnOp>(ctx), coros(coros) {}
+  AsyncReturnOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
+      : OpConversionPattern<async::ReturnOp>(ctx), coros_(coros) {}
 
   LogicalResult
   matchAndRewrite(async::ReturnOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto func = op->template getParentOfType<func::FuncOp>();
-    auto funcCoro = coros.find(func);
-    if (funcCoro == coros.end())
+    auto funcCoro = coros_->find(func);
+    if (funcCoro == coros_->end())
       return rewriter.notifyMatchFailure(
           op, "operation is not inside the async coroutine function");
 
     Location loc = op->getLoc();
     const CoroMachinery &coro = funcCoro->getSecond();
     rewriter.setInsertionPointAfter(op);
 
     // Store return values into the async values storage and switch async
if (coro.asyncToken)
       rewriter.create<RuntimeSetAvailableOp>(loc, *coro.asyncToken);
 
     rewriter.eraseOp(op);
     rewriter.create<cf::BranchOp>(loc, coro.cleanup);
     return success();
   }
 
 private:
-  llvm::DenseMap<func::FuncOp, CoroMachinery> &coros;
+  FuncCoroMapPtr coros_;
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.await and async.await_all operations to the async.runtime.await
 // or async.runtime.await_and_resume operations.
 //===----------------------------------------------------------------------===//
 
 namespace {
 template <typename AwaitType, typename AwaitableType>
 class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> {
   using AwaitAdaptor = typename AwaitType::Adaptor;
 
 public:
-  AwaitOpLoweringBase(MLIRContext *ctx,
-                      llvm::DenseMap<func::FuncOp, CoroMachinery> &coros)
-      : OpConversionPattern<AwaitType>(ctx), coros(coros) {}
+  AwaitOpLoweringBase(MLIRContext *ctx, FuncCoroMapPtr coros,
+                      bool should_lower_blocking_wait)
+      : OpConversionPattern<AwaitType>(ctx), coros_(coros),
+        should_lower_blocking_wait_(should_lower_blocking_wait) {}
 
   LogicalResult
   matchAndRewrite(AwaitType op, typename AwaitType::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // We can only await on one the `AwaitableType` (for `await` it can be
     // a `token` or a `value`, for `await_all` it must be a `group`).
     if (!op.getOperand().getType().template isa<AwaitableType>())
       return rewriter.notifyMatchFailure(op, "unsupported awaitable type");
 
-    // Check if await operation is inside the outlined coroutine function.
+    // Check if await operation is inside the coroutine function.
     auto func = op->template getParentOfType<func::FuncOp>();
-    auto funcCoro = coros.find(func);
-    const bool isInCoroutine = funcCoro != coros.end();
+    auto funcCoro = coros_->find(func);
+    const bool isInCoroutine = funcCoro != coros_->end();
 
     Location loc = op->getLoc();
     Value operand = adaptor.getOperand();
 
     Type i1 = rewriter.getI1Type();
 
+    // Delay lowering to block wait in case await op is inside async.execute
+    if (!isInCoroutine && !should_lower_blocking_wait_)
+      return failure();
+
     // Inside regular functions we use the blocking wait operation to wait for
     // the async object (token, value or group) to become available.
     if (!isInCoroutine) {
       ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
       builder.create<RuntimeAwaitOp>(loc, operand);
 
       // Assert that the awaited operands is not in the error state.
       Value isError = builder.create<RuntimeIsErrorOp>(i1, operand);
public:
   }
 
   virtual Value getReplacementValue(AwaitType op, Value operand,
                                     ConversionPatternRewriter &rewriter) const {
     return Value();
   }
 
 private:
-  llvm::DenseMap<func::FuncOp, CoroMachinery> &coros;
+  FuncCoroMapPtr coros_;
+  bool should_lower_blocking_wait_;
 };
 
 /// Lowering for `async.await` with a token operand.
 class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
   using Base = AwaitOpLoweringBase<AwaitOp, TokenType>;
 
 public:
   using Base::Base;

 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Convert async.yield operation to async.runtime operations.
 //===----------------------------------------------------------------------===//
 
 class YieldOpLowering : public OpConversionPattern<async::YieldOp> {
 public:
-  YieldOpLowering(MLIRContext *ctx,
-                  const llvm::DenseMap<func::FuncOp, CoroMachinery> &coros)
-      : OpConversionPattern<async::YieldOp>(ctx), coros(coros) {}
+  YieldOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
+      : OpConversionPattern<async::YieldOp>(ctx), coros_(coros) {}
 
   LogicalResult
   matchAndRewrite(async::YieldOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Check if yield operation is inside the async coroutine function.
     auto func = op->template getParentOfType<func::FuncOp>();
-    auto funcCoro = coros.find(func);
-    if (funcCoro == coros.end())
+    auto funcCoro = coros_->find(func);
+    if (funcCoro == coros_->end())
       return rewriter.notifyMatchFailure(
           op, "operation is not inside the async coroutine function");
 
     Location loc = op->getLoc();
     const CoroMachinery &coro = funcCoro->getSecond();
 
     // Store yielded values into the async values storage and switch async
     // values state to available.
matchAndRewrite(async::YieldOp op, OpAdaptor adaptor,
 
     rewriter.eraseOp(op);
     rewriter.create<cf::BranchOp>(loc, coro.cleanup);
 
     return success();
   }
 
 private:
-  const llvm::DenseMap<func::FuncOp, CoroMachinery> &coros;
+  FuncCoroMapPtr coros_;
 };
 
 //===----------------------------------------------------------------------===//
 // Convert cf.assert operation to cf.cond_br into `set_error` block.
 //===----------------------------------------------------------------------===//
 
 class AssertOpLowering : public OpConversionPattern<cf::AssertOp> {
 public:
-  AssertOpLowering(MLIRContext *ctx,
-                   llvm::DenseMap<func::FuncOp, CoroMachinery> &coros)
-      : OpConversionPattern<cf::AssertOp>(ctx), coros(coros) {}
+  AssertOpLowering(MLIRContext *ctx, FuncCoroMapPtr coros)
+      : OpConversionPattern<cf::AssertOp>(ctx), coros_(coros) {}
 
   LogicalResult
   matchAndRewrite(cf::AssertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Check if assert operation is inside the async coroutine function.
     auto func = op->template getParentOfType<func::FuncOp>();
-    auto funcCoro = coros.find(func);
-    if (funcCoro == coros.end())
+    auto funcCoro = coros_->find(func);
+    if (funcCoro == coros_->end())
       return rewriter.notifyMatchFailure(
           op, "operation is not inside the async coroutine function");
 
     Location loc = op->getLoc();
     CoroMachinery &coro = funcCoro->getSecond();
 
     Block *cont = rewriter.splitBlock(op->getBlock(), Block::iterator(op));
     rewriter.setInsertionPointToEnd(cont->getPrevNode());
     rewriter.create<cf::CondBranchOp>(loc, adaptor.getArg(),
                                       /*trueDest=*/cont,
                                       /*trueArgs=*/ArrayRef<Value>(),
                                       /*falseDest=*/setupSetErrorBlock(coro),
                                       /*falseArgs=*/ArrayRef<Value>());
     rewriter.eraseOp(op);
 
     return success();
   }
 
 private:
-  llvm::DenseMap<func::FuncOp, CoroMachinery> &coros;
+  FuncCoroMapPtr coros_;
 };
 
 //===----------------------------------------------------------------------===//
 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;
+  // `async.execute` body regions
+  FuncCoroMapPtr coros =
+      std::make_shared<llvm::DenseMap<func::FuncOp, CoroMachinery>>();
 
   module.walk([&](ExecuteOp execute) {
-    coros.insert(outlineExecuteOp(symbolTable, execute));
+    coros->insert(outlineExecuteOp(symbolTable, execute));
   });
 
   LLVM_DEBUG({
-    llvm::dbgs() << "Outlined " << coros.size()
+    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();
+    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.
   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);
+  asyncPatterns
+      .add<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
+          ctx, coros, /*should_lower_blocking_wait=*/true);
 
   // Lower assertions to conditional branches into error blocks.
-  asyncPatterns.add<AssertOpLowering>(ctx, coros);
+  asyncPatterns.add<YieldOpLowering, AssertOpLowering>(ctx, coros);
 
   // All high level async operations must be lowered to the runtime operations.
   ConversionTarget runtimeTarget(*ctx);
   runtimeTarget.addLegalDialect<AsyncDialect, func::FuncDialect>();
   runtimeTarget.addIllegalOp<CreateGroupOp, AddToGroupOp>();
-  runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp,
-                             async::FuncOp, async::CallOp, async::ReturnOp>();
+  runtimeTarget.addIllegalOp<ExecuteOp, AwaitOp, AwaitAllOp, async::YieldOp>();
 
   // 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();
   });
   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 coros->find(func) == coros->end();
       });
 
   if (failed(applyPartialConversion(module, runtimeTarget,
                                     std::move(asyncPatterns)))) {
     signalPassFailure();
     return;
   }
 }
 
+//===----------------------------------------------------------------------===//
+void mlir::populateAsyncFuncToAsyncRuntimeConversionPatterns(
+    RewritePatternSet &patterns, ConversionTarget &target) {
+  // Functions with coroutine CFG setups, which are results of converting
+  // async.func.
+  FuncCoroMapPtr coros =
+      std::make_shared<llvm::DenseMap<func::FuncOp, CoroMachinery>>();
+  MLIRContext *ctx = patterns.getContext();
+  // Lower async.func to func.func with coroutine cfg.
+  patterns.add<AsyncCallOpLowering>(ctx);
+  patterns.add<AsyncFuncOpLowering, AsyncReturnOpLowering>(ctx, coros);
+
+  patterns.add<AwaitTokenOpLowering, AwaitValueOpLowering, AwaitAllOpLowering>(
+      ctx, coros, /*should_lower_blocking_wait=*/false);
+  patterns.add<YieldOpLowering, AssertOpLowering>(ctx, coros);
+
+  target.addDynamicallyLegalOp<AwaitOp, AwaitAllOp, YieldOp, cf::AssertOp>(
+      [coros](Operation *op) {
+        auto func = op->getParentOfType<func::FuncOp>();
+        return coros->find(func) == coros->end();
+      });
+}
+
+void AsyncFuncToAsyncRuntimePass::runOnOperation() {
+  ModuleOp module = getOperation();
+
+  // Lower async operations to async.runtime operations.
+  MLIRContext *ctx = module->getContext();
+  RewritePatternSet asyncPatterns(ctx);
+  ConversionTarget runtimeTarget(*ctx);
+
+  // Lower async.func to func.func with coroutine cfg.
+  populateAsyncFuncToAsyncRuntimeConversionPatterns(asyncPatterns,
+                                                    runtimeTarget);
+
+  runtimeTarget.addLegalDialect<AsyncDialect, func::FuncDialect>();
+  runtimeTarget.addIllegalOp<async::FuncOp, async::CallOp, async::ReturnOp>();
+
+  runtimeTarget.addLegalOp<arith::XOrIOp, arith::ConstantOp, func::ConstantOp,
+                           cf::BranchOp, cf::CondBranchOp>();
+
+  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::createAsyncFuncToAsyncRuntimePass() {
+  return std::make_unique<AsyncFuncToAsyncRuntimePass>();
+}

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

-// RUN: mlir-opt %s -split-input-file -async-to-async-runtime                  \
-// RUN:   | FileCheck %s --dump-input=always
+// RUN: mlir-opt %s -split-input-file -async-func-to-async-runtime             \
+// RUN:   -async-to-async-runtime | FileCheck %s --dump-input=always
 
 // CHECK-LABEL: @execute_no_async_args
 func.func @execute_no_async_args(%arg0: f32, %arg1: memref<1xf32>) {
   %token = async.execute {
     %c0 = arith.constant 0 : index
     memref.store %arg0, %arg1[%c0] : memref<1xf32>
     async.yield
   }

mlir/test/mlir-cpu-runner/async-func.mlir

-// RUN:   mlir-opt %s -pass-pipeline="builtin.module(async-to-async-runtime,func.func(async-runtime-ref-counting,async-runtime-ref-counting-opt),convert-async-to-llvm,func.func(convert-linalg-to-loops,convert-scf-to-cf),convert-linalg-to-llvm,convert-vector-to-llvm,func.func(convert-arith-to-llvm),convert-func-to-llvm,reconcile-unrealized-casts)" \
+// RUN:   mlir-opt %s -pass-pipeline="builtin.module(async-func-to-async-runtime,async-to-async-runtime,func.func(async-runtime-ref-counting,async-runtime-ref-counting-opt),convert-async-to-llvm,func.func(convert-linalg-to-loops,convert-scf-to-cf),convert-linalg-to-llvm,convert-vector-to-llvm,func.func(convert-arith-to-llvm),convert-func-to-llvm,reconcile-unrealized-casts)" \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:     -e main -entry-point-result=void -O0                               \
 // RUN:     -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext  \
 // RUN:     -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext    \
 // RUN:     -shared-libs=%mlir_lib_dir/libmlir_async_runtime%shlibext   \
 // RUN: | FileCheck %s --dump-input=always
 
 // FIXME: https://github.com/llvm/llvm-project/issues/57231