[mlir] Automatic reference counting for Async values + runtime support for ref counte...
authorEugene Zhulenev <ezhulenev@google.com>
Fri, 20 Nov 2020 10:42:28 +0000 (02:42 -0800)
committerEugene Zhulenev <ezhulenev@google.com>
Fri, 20 Nov 2020 11:08:44 +0000 (03:08 -0800)
commita86a9b5ef777552d1683e2b9031e2045d39de2f0
tree5d4a3f27203286201835bf28e9d839b5da6bb21b
parent7a18bdb350e4a0318776cf66cf280b386eb9c3e8
[mlir] Automatic reference counting for Async values + runtime support for ref counted objects

Depends On D89963

**Automatic reference counting algorithm outline:**

1. `ReturnLike` operations forward the reference counted values without
    modifying the reference count.
2. Use liveness analysis to find blocks in the CFG where the lifetime of
   reference counted values ends, and insert `drop_ref` operations after
   the last use of the value.
3. Insert `add_ref` before the `async.execute` operation capturing the
   value, and pairing `drop_ref` before the async body region terminator,
   to release the captured reference counted value when execution
   completes.
4. If the reference counted value is passed only to some of the block
   successors, insert `drop_ref` operations in the beginning of the blocks
   that do not have reference coutned value uses.

Reviewed By: silvas

Differential Revision: https://reviews.llvm.org/D90716
19 files changed:
mlir/include/mlir/Dialect/Async/IR/Async.h
mlir/include/mlir/Dialect/Async/IR/AsyncBase.td
mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
mlir/include/mlir/Dialect/Async/Passes.h
mlir/include/mlir/Dialect/Async/Passes.td
mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
mlir/integration_test/Dialect/Async/CPU/test-async-parallel-for-1d.mlir
mlir/integration_test/Dialect/Async/CPU/test-async-parallel-for-2d.mlir
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
mlir/lib/Dialect/Async/Transforms/AsyncRefCounting.cpp [new file with mode: 0644]
mlir/lib/Dialect/Async/Transforms/AsyncRefCountingOptimization.cpp [new file with mode: 0644]
mlir/lib/Dialect/Async/Transforms/CMakeLists.txt
mlir/lib/ExecutionEngine/AsyncRuntime.cpp
mlir/test/Conversion/AsyncToLLVM/convert-to-llvm.mlir
mlir/test/Dialect/Async/async-ref-counting-optimization.mlir [new file with mode: 0644]
mlir/test/Dialect/Async/async-ref-counting.mlir [new file with mode: 0644]
mlir/test/Dialect/Async/ops.mlir
mlir/test/mlir-cpu-runner/async-group.mlir
mlir/test/mlir-cpu-runner/async.mlir