[mlir:Async] Add an async reference counting pass based on the user defined policy

[ezhulenev]

Summary

Depends On D104999

Automatic reference counting based on the liveness analysis can add a lot of reference counting overhead at runtime. If the IR is known to be constrained to few particular "shapes", it's much more efficient to provide a custom reference counting policy that will specify where it is required to update the async value reference count.

[ezhulenev]

Fix a bug in ref counting

[ezhulenev]

Put all reference counting passes into the single file

[mehdi_amini]

I suspect the title should be "based" instead of "baded"?

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

32

I don't understand the method name here?

39

Nit: three ///.

Also, if the first sentence could be along the line of "Create a pass that does XYZ".

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

79

I am concerned about this: this will break reproducer. Can we look into solutions for this?

[mehdi_amini]

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

79

One option would be to keep the pass without option, and expose an entry point that isn't a pass. The user would create their own pass and inject the policy into the utility entry point,

[ezhulenev]

Rename pass to AsyncRuntimePolicyBasedRefCountingPass

[mehdi_amini]

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

79

Thinking more about it: isn't the "policy" just implementing a calling convention of some sort?
If so, why isn't it an op or dialect interface?

[ezhulenev]

Update documentation

[ezhulenev]

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

32

Renamed to AsyncRuntimePolicyBasedRefCountingPass

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

79

Yeah, agree that this is weird, but I don't actually have examples of policy beyond default. I'll take a look at it later. For now I just added a little bit more documentation.

[mehdi_amini]

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

79

If we don't have other policy, can we just not introduce the ability to inject behavior here for now?
When such injection will be needed, it'll be easier to look into interfaces.

[ezhulenev]

Remove user defined policies from policy based ref counting pass

[ezhulenev]

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

79

Yeah, good point. Removed it for now.

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

std::unique_ptr<Pass> createAsyncParallelForPass(bool asyncDispatch,
                                                  int32_t targetBlockSize);
 
 std::unique_ptr<OperationPass<ModuleOp>> createAsyncToAsyncRuntimePass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingPass();
 
 std::unique_ptr<Pass> createAsyncRuntimeRefCountingOptPass();
 
+std::unique_ptr<Pass> createAsyncRuntimeBespokeRefCountingPass();
+
+// Each policy function must return a number of references that should be added
+// or dropped for the given operand:
+//
+//  0 - no new reference counting operations are required
+// +N - `add_ref` operation will be created before the operand owner
+// -N - `drop_ref` operation will be created after the operand owner
+std::unique_ptr<Pass> createAsyncRuntimeBespokeRefCountingPass(
+    llvm::SmallVector<std::function<FailureOr<int>(OpOperand &)>> policy);
+
 //===----------------------------------------------------------------------===//
 // 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

 
 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"];
 }
 
+def AsyncRuntimeBespokeRefCounting
+    : Pass<"async-runtime-bespoke-ref-counting"> {
+  let summary = "Policy based 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.
+
+    This pass takes user defined policy that specifies where to put reference
+    counting operations. Currently there is no way to specify the policy using
+    command line flags.
+
+    The default reference counting policy makes few assumptions:
+      1. Async token can be awaited or added to the group only once.
+      2. Async value or group can be awaited only once.
+
+    Under these assumptions reference counting only needs to drop reference:
+      1. After `async.runtime.await` operation for async tokens and groups
+         (until error handling is not implemented for the sync await).
+      2. After `async.runtime.is_error` operation for async tokens and groups
+         (this is the last operation in the coroutine resume function).
+      3. After `async.runtime.load` operation for async values.
+
+    This pass introduces significanly less runtime overhead compared to the
+    automatic reference counting.
+  }];
+
+  let constructor = "mlir::createAsyncRuntimeBespokeRefCountingPass()";
+  let dependentDialects = ["async::AsyncDialect"];
+}
+
 #endif // MLIR_DIALECT_ASYNC_PASSES

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

+//===- AsyncRuntimeBespokeRefCounting.cpp - Async Runtime Ref Counting ----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements policy based reference counting for Async runtime
+// operations and types.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Analysis/Liveness.h"
+#include "mlir/Dialect/Async/IR/Async.h"
+#include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallSet.h"
+
+using namespace mlir;
+using namespace mlir::async;
+
+#define DEBUG_TYPE "async-runtime-ref-counting"
+
+namespace {
+
+class AsyncRuntimeBespokeRefCountingPass
+    : public AsyncRuntimeBespokeRefCountingBase<
+          AsyncRuntimeBespokeRefCountingPass> {
+public:
+  AsyncRuntimeBespokeRefCountingPass() { initializeDefaultPolicy(); }
+
+  AsyncRuntimeBespokeRefCountingPass(
+      llvm::SmallVector<std::function<FailureOr<int>(OpOperand &)>> policy)
+      : policy(policy) {}
+
+  void runOnOperation() override;
+
+private:
+  // Adds a reference counting operations for all uses of the `value` according
+  // to the reference counting policy.
+  LogicalResult addRefCounting(Value value);
+
+  void initializeDefaultPolicy();
+
+  llvm::SmallVector<std::function<FailureOr<int>(OpOperand &)>> policy;
+};
+
+} // namespace
+
+LogicalResult AsyncRuntimeBespokeRefCountingPass::addRefCounting(Value value) {
+  OpBuilder b(value.getContext());
+
+  for (OpOperand &operand : value.getUses()) {
+    Location loc = operand.getOwner()->getLoc();
+
+    for (auto &func : policy) {
+      FailureOr<int> refCount = func(operand);
+      if (failed(refCount))
+        return failure();
+
+      int cnt = refCount.getValue();
+
+      // Create `add_ref` operation before the operand owner.
+      if (cnt > 0) {
+        b.setInsertionPoint(operand.getOwner());
+        b.create<RuntimeAddRefOp>(loc, value, b.getI32IntegerAttr(cnt));
+      }
+
+      // Create `drop_ref` operation after the operand owner.
+      if (cnt < 0) {
+        b.setInsertionPointAfter(operand.getOwner());
+        b.create<RuntimeDropRefOp>(loc, value, b.getI32IntegerAttr(-cnt));
+      }
+    }
+  }
+
+  return success();
+}
+
+void AsyncRuntimeBespokeRefCountingPass::initializeDefaultPolicy() {
+  policy.push_back([](OpOperand &operand) -> FailureOr<int> {
+    Operation *op = operand.getOwner();
+    Type type = operand.get().getType();
+
+    bool isToken = type.isa<TokenType>();
+    bool isGroup = type.isa<GroupType>();
+    bool isValue = type.isa<ValueType>();
+
+    // Drop reference after async token or group await (sync await)
+    if (auto await = dyn_cast<RuntimeAwaitOp>(op))
+      return (isToken || isGroup) ? -1 : 0;
+
+    // Drop reference after async token or group error check (coro await).
+    if (auto await = dyn_cast<RuntimeIsErrorOp>(op))
+      return (isToken || isGroup) ? -1 : 0;
+
+    // Drop reference after async value load.
+    if (auto load = dyn_cast<RuntimeLoadOp>(op))
+      return isValue ? -1 : 0;
+
+    // Drop reference after async token added to the group.
+    if (auto add = dyn_cast<RuntimeAddToGroupOp>(op))
+      return isToken ? -1 : 0;
+
+    return 0;
+  });
+}
+
+void AsyncRuntimeBespokeRefCountingPass::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 = 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())
+    return signalPassFailure();
+
+  // Add reference counting to block arguments.
+  WalkResult blockWalk = op->walk([&](Block *block) -> WalkResult {
+    for (BlockArgument arg : block->getArguments())
+      if (isRefCounted(arg.getType()))
+        if (failed(addRefCounting(arg)))
+          return WalkResult::interrupt();
+
+    return WalkResult::advance();
+  });
+
+  if (blockWalk.wasInterrupted()) {
+    signalPassFailure();
+    return;
+  }
+
+  // Add reference counting to operation results.
+  WalkResult opWalk = op->walk([&](Operation *op) -> WalkResult {
+    for (unsigned i = 0; i < op->getNumResults(); ++i)
+      if (isRefCounted(op->getResultTypes()[i]))
+        if (failed(addRefCounting(op->getResult(i))))
+          return WalkResult::interrupt();
+
+    return WalkResult::advance();
+  });
+
+  if (opWalk.wasInterrupted())
+    signalPassFailure();
+}
+
+std::unique_ptr<Pass> mlir::createAsyncRuntimeBespokeRefCountingPass() {
+  return std::make_unique<AsyncRuntimeBespokeRefCountingPass>();
+}
+
+std::unique_ptr<Pass> createAsyncRuntimeBespokeRefCountingPass(
+    llvm::SmallVector<std::function<FailureOr<int>(OpOperand &)>> policy) {
+  return std::make_unique<AsyncRuntimeBespokeRefCountingPass>(
+      std::move(policy));
+}

mlir/lib/Dialect/Async/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRAsyncTransforms
   AsyncParallelFor.cpp
+  AsyncRuntimeBespokeRefCounting.cpp
   AsyncRuntimeRefCounting.cpp
   AsyncRuntimeRefCountingOpt.cpp
   AsyncToAsyncRuntime.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Async
 
   DEPENDS

mlir/test/Dialect/Async/async-runtime-bespoke-ref-counting.mlir

+// RUN: mlir-opt %s -async-runtime-bespoke-ref-counting | FileCheck %s
+
+// CHECK-LABEL: @token_await
+// CHECK:         %[[TOKEN:.*]]: !async.token
+func @token_await(%arg0: !async.token) {
+  // CHECK: async.runtime.await %[[TOKEN]]
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  async.runtime.await %arg0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @group_await
+// CHECK:         %[[GROUP:.*]]: !async.group
+func @group_await(%arg0: !async.group) {
+  // CHECK: async.runtime.await %[[GROUP]]
+  // CHECK: async.runtime.drop_ref %[[GROUP]] {count = 1 : i32}
+  async.runtime.await %arg0 : !async.group
+  return
+}
+
+// CHECK-LABEL: @add_token_to_group
+// CHECK:         %[[GROUP:.*]]: !async.group
+// CHECK:         %[[TOKEN:.*]]: !async.token
+func @add_token_to_group(%arg0: !async.group, %arg1: !async.token) {
+  // CHECK: async.runtime.add_to_group %[[TOKEN]], %[[GROUP]]
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  async.runtime.add_to_group %arg1, %arg0 : !async.token
+  return
+}
+
+// CHECK-LABEL: @value_load
+// CHECK:         %[[VALUE:.*]]: !async.value<f32>
+func @value_load(%arg0: !async.value<f32>) {
+  // CHECK: async.runtime.load %[[VALUE]]
+  // CHECK: async.runtime.drop_ref %[[VALUE]] {count = 1 : i32}
+  %0 = async.runtime.load %arg0 : !async.value<f32>
+  return
+}
+
+// CHECK-LABEL: @error_check
+// CHECK:         %[[TOKEN:.*]]: !async.token
+func @error_check(%arg0: !async.token) {
+  // CHECK: async.runtime.is_error %[[TOKEN]]
+  // CHECK: async.runtime.drop_ref %[[TOKEN]] {count = 1 : i32}
+  %0 = async.runtime.is_error %arg0 : !async.token
+  return
+}

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-1d.mlir

 // RUN:   mlir-opt %s -async-parallel-for                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:  -e entry -entry-point-result=void -O0                                 \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
 // RUN: | FileCheck %s --dump-input=always
 
+// RUN:   mlir-opt %s -async-parallel-for                                      \
+// RUN:               -async-to-async-runtime                                  \
+// RUN:               -async-runtime-bespoke-ref-counting                      \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:  -e entry -entry-point-result=void -O0                                 \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
+// RUN: | FileCheck %s --dump-input=always
+
 // RUN:   mlir-opt %s -async-parallel-for="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
 // RUN:                                    target-block-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

 // RUN:   mlir-opt %s -async-parallel-for                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:  -e entry -entry-point-result=void -O0                                 \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
 // RUN: | FileCheck %s --dump-input=always
 
+// RUN:   mlir-opt %s -async-parallel-for                                      \
+// RUN:               -async-to-async-runtime                                  \
+// RUN:               -async-runtime-bespoke-ref-counting                      \
+// RUN:               -convert-async-to-llvm                                   \
+// RUN:               -convert-scf-to-std                                      \
+// RUN:               -convert-std-to-llvm                                     \
+// RUN: | mlir-cpu-runner                                                      \
+// RUN:  -e entry -entry-point-result=void -O0                                 \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext\
+// RUN: | FileCheck %s --dump-input=always
+
 // RUN:   mlir-opt %s -async-parallel-for="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
 // RUN:                                    target-block-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
 // RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \