[mlir:Async] Fix a bug in automatic refence counting around function calls

[ezhulenev]

Summary

Depends On D104998

Function calls "transfer ownership" to the callee and it puts additional constraints on the reference counting optimization pass

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

 //===- AsyncRuntimeRefCountingOpt.cpp - Async 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Optimize Async dialect reference counting operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::async;
 
 #define DEBUG_TYPE "async-ref-counting"
 
for (auto &kv : blockUsers) {
 
     for (RuntimeAddRefOp addRef : info.addRefs) {
       for (RuntimeDropRefOp dropRef : info.dropRefs) {
         // `drop_ref` operation after the `add_ref` with matching count.
         if (dropRef.count() != addRef.count() ||
             dropRef->isBeforeInBlock(addRef.getOperation()))
           continue;
 
+        // When reference counted value passed to a function as an argument,
+        // function takes ownership of +1 reference and it will drop it before
+        // returning.
+        //
+        // Example:
+        //
+        //   %token = ... : !async.token
+        //
+        //   async.runtime.add_ref %token {count = 1 : i32} : !async.token
+        //   call @pass_token(%token: !async.token, ...)
+        //
+        //   async.await %token : !async.token
+        //   async.runtime.drop_ref %token {count = 1 : i32} : !async.token
+        //
+        // In this example if we'll cancel a pair of reference counting
+        // operations we might end up with a deallocated token when we'll
+        // reach `async.await` operation.
+        Operation *firstFunctionCallUser = nullptr;
+        Operation *lastNonFunctionCallUser = nullptr;
+
+        for (Operation *user : info.users) {
+          // `user` operation lies after `addRef` ...
+          if (user == addRef || user->isBeforeInBlock(addRef))
+            continue;
+          // ... and before `dropRef`.
+          if (user == dropRef || dropRef->isBeforeInBlock(user))
+            break;
+
+          // Find the first function call user of the reference counted value.
+          Operation *functionCall = dyn_cast<CallOp>(user);
+          if (functionCall &&
+              (!firstFunctionCallUser ||
+               functionCall->isBeforeInBlock(firstFunctionCallUser))) {
+            firstFunctionCallUser = functionCall;
+            continue;
+          }
+
+          // Find the last regular user of the reference counted value.
+          if (!functionCall &&
+              (!lastNonFunctionCallUser ||
+               lastNonFunctionCallUser->isBeforeInBlock(user))) {
+            lastNonFunctionCallUser = user;
+            continue;
+          }
+        }
+
+        // Non function call user after the function call user of the reference
+        // counted value.
+        if (firstFunctionCallUser && lastNonFunctionCallUser &&
+            firstFunctionCallUser->isBeforeInBlock(lastNonFunctionCallUser))
+          continue;
+
         // Try to cancel the pair of `add_ref` and `drop_ref` operations.
         auto emplaced = cancellable.try_emplace(dropRef.getOperation(),
                                                 addRef.getOperation());
 
         if (!emplaced.second) // `drop_ref` was already marked for removal
           continue;           // go to the next `drop_ref`
 
         if (emplaced.second) // successfully cancelled `add_ref` <-> `drop_ref`

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

 add_mlir_dialect_library(MLIRAsyncTransforms
   AsyncParallelFor.cpp
   AsyncRuntimeRefCounting.cpp
   AsyncRuntimeRefCountingOpt.cpp
   AsyncToAsyncRuntime.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Async
 
   DEPENDS
   MLIRAsyncPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRAsync
-  MLIRSCF
   MLIRPass
+  MLIRSCF
+  MLIRStandard
   MLIRTransforms
   MLIRTransformUtils
 )

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

func @cancellable_operations_3(%arg0: !async.token) {
   call @consume_token(%arg0): (!async.token) -> ()
   // CHECK-NOT: async.runtime.drop_ref
   async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
   // CHECK: async.runtime.await
   async.runtime.await %arg0 : !async.token
   // CHECK: return
   return
 }
+
+// CHECK-LABEL: @not_cancellable_operations_0
+func @not_cancellable_operations_0(%arg0: !async.token) {
+  // CHECK: add_ref
+  async.runtime.add_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: call @consume_toke
+  call @consume_token(%arg0): (!async.token) -> ()
+  // CHECK: async.runtime.await
+  async.runtime.await %arg0 : !async.token
+  // CHECK: async.runtime.drop_ref
+  async.runtime.drop_ref %arg0 {count = 1 : i32} : !async.token
+  // CHECK: return
+  return
+}

mlir/test/Integration/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

 // RUN:   mlir-opt %s                                                          \
 // RUN:               -convert-linalg-to-parallel-loops                        \
 // RUN:               -async-parallel-for                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -std-expand                                              \
 // RUN:               -convert-vector-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN: -e entry -entry-point-result=void -O3                                  \
 // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext  \

mlir/test/Integration/Dialect/Async/CPU/microbench-scf-async-parallel-for.mlir

 // RUN:   mlir-opt %s                                                          \
 // RUN:               -async-parallel-for                                      \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-linalg-to-loops                                 \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -std-expand                                              \
 // RUN:               -convert-vector-to-llvm                                  \
-// RUN:               -convert-std-to-llvm                                     \
+// RUN:               -convert-std-to-llvm -print-ir-after-all                                    \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN: -e entry -entry-point-result=void -O3                                  \
 // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext  \
 // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext\
 // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_async_runtime%shlibext \
 // RUN: | FileCheck %s --dump-input=always
 
 // RUN:   mlir-opt %s                                                          \
 // RUN:               -async-parallel-for=async-dispatch=false                 \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// RUN:               -async-runtime-ref-counting-opt                          \
 // RUN:               -convert-async-to-llvm                                   \
 // RUN:               -convert-linalg-to-loops                                 \
 // RUN:               -convert-scf-to-std                                      \
 // RUN:               -std-expand                                              \
 // RUN:               -convert-vector-to-llvm                                  \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN: -e entry -entry-point-result=void -O3                                  \

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                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// 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="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
 // RUN:                                    target-block-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// 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

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                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// 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="async-dispatch=false                \
 // RUN:                                    num-workers=20                      \
 // RUN:                                    target-block-size=1"                \
 // RUN:               -async-to-async-runtime                                  \
 // RUN:               -async-runtime-ref-counting                              \
-// FIXME:             -async-runtime-ref-counting-opt                          \
+// 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