llvm/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp

//===- BufferOptimizations.cpp - pre-pass optimizations for bufferization -===//
//
// 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 logic for three optimization passes. The first two
// passes try to move alloc nodes out of blocks to reduce the number of
// allocations and copies during buffer deallocation. The third pass tries to
// convert heap-based allocations to stack-based allocations, if possible.

#include "mlir/Dialect/Bufferization/Transforms/Passes.h"

#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Pass/Pass.h"

namespace mlir {
namespace bufferization {
#define GEN_PASS_DEF_BUFFERHOISTING
#define GEN_PASS_DEF_BUFFERLOOPHOISTING
#define GEN_PASS_DEF_PROMOTEBUFFERSTOSTACK
#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
} // namespace bufferization
} // namespace mlir

usingnamespacemlir;
usingnamespacemlir::bufferization;

/// Returns true if the given operation implements a known high-level region-
/// based control-flow interface.
static bool isKnownControlFlowInterface(Operation *op) { … }

/// Returns true if the given operation represents a loop by testing whether it
/// implements the `LoopLikeOpInterface` or the `RegionBranchOpInterface`. In
/// the case of a `RegionBranchOpInterface`, it checks all region-based control-
/// flow edges for cycles.
static bool isLoop(Operation *op) { … }

/// Return whether the given operation is a loop with sequential execution
/// semantics.
static bool isSequentialLoop(Operation *op) { … }

/// Returns true if the given operation implements the AllocationOpInterface
/// and it supports the dominate block hoisting.
static bool allowAllocDominateBlockHoisting(Operation *op) { … }

/// Returns true if the given operation implements the AllocationOpInterface
/// and it supports the loop hoisting.
static bool allowAllocLoopHoisting(Operation *op) { … }

/// Check if the size of the allocation is less than the given size. The
/// transformation is only applied to small buffers since large buffers could
/// exceed the stack space.
static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
                                unsigned maxRankOfAllocatedMemRef) { … }

/// Checks whether the given aliases leave the allocation scope.
static bool
leavesAllocationScope(Region *parentRegion,
                      const BufferViewFlowAnalysis::ValueSetT &aliases) { … }

/// Checks, if an automated allocation scope for a given alloc value exists.
static bool hasAllocationScope(Value alloc,
                               const BufferViewFlowAnalysis &aliasAnalysis) { … }

namespace {

//===----------------------------------------------------------------------===//
// BufferAllocationHoisting
//===----------------------------------------------------------------------===//

/// A base implementation compatible with the `BufferAllocationHoisting` class.
struct BufferAllocationHoistingStateBase { … };

/// Implements the actual hoisting logic for allocation nodes.
template <typename StateT>
class BufferAllocationHoisting : public BufferPlacementTransformationBase { … };

/// A state implementation compatible with the `BufferAllocationHoisting` class
/// that hoists allocations into dominator blocks while keeping them inside of
/// loops.
struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase { … };

/// A state implementation compatible with the `BufferAllocationHoisting` class
/// that hoists allocations out of loops.
struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase { … };

//===----------------------------------------------------------------------===//
// BufferPlacementPromotion
//===----------------------------------------------------------------------===//

/// Promotes heap-based allocations to stack-based allocations (if possible).
class BufferPlacementPromotion : BufferPlacementTransformationBase { … };

//===----------------------------------------------------------------------===//
// BufferOptimizationPasses
//===----------------------------------------------------------------------===//

/// The buffer hoisting pass that hoists allocation nodes into dominating
/// blocks.
struct BufferHoistingPass
    : public bufferization::impl::BufferHoistingBase<BufferHoistingPass> { … };

/// The buffer loop hoisting pass that hoists allocation nodes out of loops.
struct BufferLoopHoistingPass
    : public bufferization::impl::BufferLoopHoistingBase<
          BufferLoopHoistingPass> { … };

/// The promote buffer to stack pass that tries to convert alloc nodes into
/// alloca nodes.
class PromoteBuffersToStackPass
    : public bufferization::impl::PromoteBuffersToStackBase<
          PromoteBuffersToStackPass> { … };

} // namespace

void mlir::bufferization::hoistBuffersFromLoops(Operation *op) { … }

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

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

std::unique_ptr<Pass> mlir::bufferization::createPromoteBuffersToStackPass(
    unsigned maxAllocSizeInBytes, unsigned maxRankOfAllocatedMemRef) { … }

std::unique_ptr<Pass> mlir::bufferization::createPromoteBuffersToStackPass(
    std::function<bool(Value)> isSmallAlloc) { … }