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

//===- AsyncRuntimeRefCounting.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 automatic reference counting for Async runtime
// operations and types.
//
//===----------------------------------------------------------------------===//

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

#include "mlir/Analysis/Liveness.h"
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/SmallSet.h"

namespace mlir {
#define GEN_PASS_DEF_ASYNCRUNTIMEREFCOUNTING
#define GEN_PASS_DEF_ASYNCRUNTIMEPOLICYBASEDREFCOUNTING
#include "mlir/Dialect/Async/Passes.h.inc"
} // namespace mlir

#define DEBUG_TYPE …

usingnamespacemlir;
usingnamespacemlir::async;

//===----------------------------------------------------------------------===//
// Utility functions shared by reference counting passes.
//===----------------------------------------------------------------------===//

// Drop the reference count immediately if the value has no uses.
static LogicalResult dropRefIfNoUses(Value value, unsigned count = 1) { … }

// Calls `addRefCounting` for every reference counted value defined by the
// operation `op` (block arguments and values defined in nested regions).
static LogicalResult walkReferenceCountedValues(
    Operation *op, llvm::function_ref<LogicalResult(Value)> addRefCounting) { … }

//===----------------------------------------------------------------------===//
// Automatic reference counting based on the liveness analysis.
//===----------------------------------------------------------------------===//

namespace {

class AsyncRuntimeRefCountingPass
    : public impl::AsyncRuntimeRefCountingBase<AsyncRuntimeRefCountingPass> { … };

} // namespace

LogicalResult AsyncRuntimeRefCountingPass::addDropRefAfterLastUse(Value value) { … }

LogicalResult
AsyncRuntimeRefCountingPass::addAddRefBeforeFunctionCall(Value value) { … }

LogicalResult
AsyncRuntimeRefCountingPass::addDropRefInDivergentLivenessSuccessor(
    Value value) { … }

LogicalResult
AsyncRuntimeRefCountingPass::addAutomaticRefCounting(Value value) { … }

void AsyncRuntimeRefCountingPass::runOnOperation() { … }

//===----------------------------------------------------------------------===//
// Reference counting based on the user defined policy.
//===----------------------------------------------------------------------===//

namespace {

class AsyncRuntimePolicyBasedRefCountingPass
    : public impl::AsyncRuntimePolicyBasedRefCountingBase<
          AsyncRuntimePolicyBasedRefCountingPass> { … };

} // namespace

LogicalResult
AsyncRuntimePolicyBasedRefCountingPass::addRefCounting(Value value) { … }

void AsyncRuntimePolicyBasedRefCountingPass::initializeDefaultPolicy() { … }

void AsyncRuntimePolicyBasedRefCountingPass::runOnOperation() { … }

//----------------------------------------------------------------------------//

std::unique_ptr<Pass> mlir::createAsyncRuntimeRefCountingPass() { … }

std::unique_ptr<Pass> mlir::createAsyncRuntimePolicyBasedRefCountingPass() { … }