[mlir] Convert async dialect passes from function passes to op agnostic passes

[ezhulenev]

Rebased

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 {
 
-std::unique_ptr<OperationPass<FuncOp>> createAsyncParallelForPass();
+std::unique_ptr<Pass> createAsyncParallelForPass();
 
-std::unique_ptr<OperationPass<FuncOp>>
-createAsyncParallelForPass(int numWorkerThreads);
+std::unique_ptr<Pass> createAsyncParallelForPass(int numWorkerThreads);
 
 std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
 
-std::unique_ptr<OperationPass<FuncOp>> createAsyncRuntimeRefCountingPass();
+std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
-std::unique_ptr<OperationPass<FuncOp>> createAsyncRuntimeRefCountingOptPass();
+std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Async/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ASYNC_PASSES_H_

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

 //===-- Passes.td - Async pass definition file -------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_ASYNC_PASSES
 #define MLIR_DIALECT_ASYNC_PASSES
 
 include "mlir/Pass/PassBase.td"
 
-def AsyncParallelFor : FunctionPass<"async-parallel-for"> {
+def AsyncParallelFor : Pass<"async-parallel-for"> {
   let summary = "Convert scf.parallel operations to multiple async regions "
                 "executed concurrently for non-overlapping iteration ranges";
   let constructor = "mlir::createAsyncParallelForPass()";
   let options = [
     Option<"numConcurrentAsyncExecute", "num-concurrent-async-execute",
       "int32_t", /*default=*/"4",
       "The number of async.execute operations that will be used for concurrent "
       "loop execution.">
   ];
   let dependentDialects = ["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 constructor = "mlir::createAsyncToAsyncRuntimePass()";
   let dependentDialects = ["async::AsyncDialect"];
 }
 
-def AsyncRuntimeRefCounting : FunctionPass<"async-runtime-ref-counting"> {
+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
     the correct placing of the reference counting operations.
 
     See: https://llvm.org/docs/Coroutines.html#switched-resume-lowering
   }];
 
   let constructor = "mlir::createAsyncRuntimeRefCountingPass()";
   let dependentDialects = ["async::AsyncDialect"];
 }
 
-def AsyncRuntimeRefCountingOpt :
-    FunctionPass<"async-runtime-ref-counting-opt"> {
+def AsyncRuntimeRefCountingOpt : Pass<"async-runtime-ref-counting-opt"> {
   let summary = "Optimize automatic reference counting operations for the"
                 "Async runtime by removing redundant operations";
   let constructor = "mlir::createAsyncRuntimeRefCountingOptPass()";
   let dependentDialects = ["async::AsyncDialect"];
 }
 
 #endif // MLIR_DIALECT_ASYNC_PASSES

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

 
 struct AsyncParallelForPass
     : public AsyncParallelForBase<AsyncParallelForPass> {
   AsyncParallelForPass() = default;
   AsyncParallelForPass(int numWorkerThreads) {
     assert(numWorkerThreads >= 1);
     numConcurrentAsyncExecute = numWorkerThreads;
   }
-  void runOnFunction() override;
+  void runOnOperation() override;
 };
 
 } // namespace
 
 LogicalResult
 AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
                                          PatternRewriter &rewriter) const {
   // We do not currently support rewrite for parallel op with reductions.
AsyncParallelForRewrite::matchAndRewrite(scf::ParallelOp op,
   rewriter.create<AwaitAllOp>(loc, group.result());
 
   // Erase the original parallel operation.
   rewriter.eraseOp(op);
 
   return success();
 }
 
-void AsyncParallelForPass::runOnFunction() {
+void AsyncParallelForPass::runOnOperation() {
   MLIRContext *ctx = &getContext();
 
   RewritePatternSet patterns(ctx);
   patterns.add<AsyncParallelForRewrite>(ctx, numConcurrentAsyncExecute);
 
-  if (failed(applyPatternsAndFoldGreedily(getFunction(), std::move(patterns))))
+  if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
     signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<FuncOp>> mlir::createAsyncParallelForPass() {
+std::unique_ptr<Pass> mlir::createAsyncParallelForPass() {
   return std::make_unique<AsyncParallelForPass>();
 }
 
-std::unique_ptr<OperationPass<FuncOp>>
-mlir::createAsyncParallelForPass(int numWorkerThreads) {
+std::unique_ptr<Pass> mlir::createAsyncParallelForPass(int numWorkerThreads) {
   return std::make_unique<AsyncParallelForPass>(numWorkerThreads);
 }

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

 #define DEBUG_TYPE "async-runtime-ref-counting"
 
 namespace {
 
 class AsyncRuntimeRefCountingPass
     : public AsyncRuntimeRefCountingBase<AsyncRuntimeRefCountingPass> {
 public:
   AsyncRuntimeRefCountingPass() = default;
-  void runOnFunction() override;
+  void runOnOperation() override;
 
 private:
   /// Adds an automatic reference counting to the `value`.
   ///
   /// All values (token, group or value) are semantically created with a
   /// reference count of +1 and it is the responsibility of the async value user
   /// to place the `add_ref` and `drop_ref` operations to ensure that the value
   /// is destroyed after the last use.
AsyncRuntimeRefCountingPass::addAutomaticRefCounting(Value value) {
 
   // Verify that the `value` is in `liveIn` set of all successors.
   if (failed(verifySuccessors(value)))
     return failure();
 
   return success();
 }
 
-void AsyncRuntimeRefCountingPass::runOnFunction() {
-  FuncOp func = getFunction();
+void AsyncRuntimeRefCountingPass::runOnOperation() {
+  Operation *op = getOperation();
 
   // Check that we do not have high level async operations in the IR because
   // otherwise automatic reference counting will produce incorrect results after
   // execute operations will be lowered to `async.runtime`
-  WalkResult executeOpWalk = func.walk([&](Operation *op) -> WalkResult {
+  WalkResult executeOpWalk = op->walk([&](Operation *op) -> WalkResult {
     if (!isa<ExecuteOp, AwaitOp, AwaitAllOp, YieldOp>(op))
       return WalkResult::advance();
 
     return op->emitError()
            << "async operations must be lowered to async runtime operations";
   });
 
   if (executeOpWalk.wasInterrupted()) {
     signalPassFailure();
     return;
   }
 
   // Add reference counting to block arguments.
-  WalkResult blockWalk = func.walk([&](Block *block) -> WalkResult {
+  WalkResult blockWalk = op->walk([&](Block *block) -> WalkResult {
     for (BlockArgument arg : block->getArguments())
       if (isRefCounted(arg.getType()))
         if (failed(addAutomaticRefCounting(arg)))
           return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   if (blockWalk.wasInterrupted()) {
     signalPassFailure();
     return;
   }
 
   // Add reference counting to operation results.
-  WalkResult opWalk = func.walk([&](Operation *op) -> WalkResult {
+  WalkResult opWalk = op->walk([&](Operation *op) -> WalkResult {
     for (unsigned i = 0; i < op->getNumResults(); ++i)
       if (isRefCounted(op->getResultTypes()[i]))
         if (failed(addAutomaticRefCounting(op->getResult(i))))
           return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   if (opWalk.wasInterrupted())
     signalPassFailure();
 }
 
-std::unique_ptr<OperationPass<FuncOp>>
-mlir::createAsyncRuntimeRefCountingPass() {
+std::unique_ptr<Pass> mlir::createAsyncRuntimeRefCountingPass() {
   return std::make_unique<AsyncRuntimeRefCountingPass>();
 }

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

 #define DEBUG_TYPE "async-ref-counting"
 
 namespace {
 
 class AsyncRuntimeRefCountingOptPass
     : public AsyncRuntimeRefCountingOptBase<AsyncRuntimeRefCountingOptPass> {
 public:
   AsyncRuntimeRefCountingOptPass() = default;
-  void runOnFunction() override;
+  void runOnOperation() override;
 
 private:
   LogicalResult optimizeReferenceCounting(
       Value value, llvm::SmallDenseMap<Operation *, Operation *> &cancellable);
 };
 
 } // namespace
 
for (RuntimeAddRefOp addRef : info.addRefs) {
           break;             // go to the next `add_ref`
       }
     }
   }
 
   return success();
 }
 
-void AsyncRuntimeRefCountingOptPass::runOnFunction() {
-  FuncOp func = getFunction();
+void AsyncRuntimeRefCountingOptPass::runOnOperation() {
+  Operation *op = getOperation();
 
   // Mapping from `dropRef.getOperation()` to `addRef.getOperation()`.
   //
   // Find all cancellable pairs of operation and erase them in the end to keep
   // all iterators valid while we are walking the function operations.
   llvm::SmallDenseMap<Operation *, Operation *> cancellable;
 
   // Optimize reference counting for values defined by block arguments.
-  WalkResult blockWalk = func.walk([&](Block *block) -> WalkResult {
+  WalkResult blockWalk = op->walk([&](Block *block) -> WalkResult {
     for (BlockArgument arg : block->getArguments())
       if (isRefCounted(arg.getType()))
         if (failed(optimizeReferenceCounting(arg, cancellable)))
           return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   if (blockWalk.wasInterrupted())
     signalPassFailure();
 
   // Optimize reference counting for values defined by operation results.
-  WalkResult opWalk = func.walk([&](Operation *op) -> WalkResult {
+  WalkResult opWalk = op->walk([&](Operation *op) -> WalkResult {
     for (unsigned i = 0; i < op->getNumResults(); ++i)
       if (isRefCounted(op->getResultTypes()[i]))
         if (failed(optimizeReferenceCounting(op->getResult(i), cancellable)))
           return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   if (opWalk.wasInterrupted())
     signalPassFailure();
 
   LLVM_DEBUG({
     llvm::dbgs() << "Found " << cancellable.size()
                  << " cancellable reference counting operations\n";
   });
 
   // Erase all cancellable `add_ref <-> drop_ref` operation pairs.
   for (auto &kv : cancellable) {
     kv.first->erase();
     kv.second->erase();
   }
 }
 
-std::unique_ptr<OperationPass<FuncOp>>
-mlir::createAsyncRuntimeRefCountingOptPass() {
+std::unique_ptr<Pass> mlir::createAsyncRuntimeRefCountingOptPass() {
   return std::make_unique<AsyncRuntimeRefCountingOptPass>();
 }

mlir/test/Integration/GPU/CUDA/async.mlir

 // RUN: mlir-opt %s \
 // RUN:   -gpu-kernel-outlining \
 // RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin)' \
-// RUN:   -gpu-async-region -async-ref-counting -gpu-to-llvm \
-// RUN:   -async-to-async-runtime -convert-async-to-llvm -convert-std-to-llvm \
+// RUN:   -gpu-async-region -gpu-to-llvm \
+// RUN:   -async-to-async-runtime -async-runtime-ref-counting \
+// RUN:   -convert-async-to-llvm -convert-std-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_cuda_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void -O0 \
 // RUN: | FileCheck %s
 
 func @main() {