1428 lines
61 KiB
C++
1428 lines
61 KiB
C++
//===- LoopFusion.cpp - Code to perform loop fusion -----------------------===//
|
|
//
|
|
// 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 affine fusion.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/Passes.h"
|
|
|
|
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
|
|
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
|
|
#include "mlir/Dialect/Affine/Analysis/Utils.h"
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Affine/LoopFusionUtils.h"
|
|
#include "mlir/Dialect/Affine/LoopUtils.h"
|
|
#include "mlir/Dialect/Affine/Utils.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/Transforms/Passes.h"
|
|
#include "llvm/ADT/DenseMap.h"
|
|
#include "llvm/ADT/DenseSet.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/Support/CommandLine.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
#include <iomanip>
|
|
#include <optional>
|
|
#include <sstream>
|
|
|
|
namespace mlir {
|
|
namespace affine {
|
|
#define GEN_PASS_DEF_AFFINELOOPFUSION
|
|
#include "mlir/Dialect/Affine/Passes.h.inc"
|
|
} // namespace affine
|
|
} // namespace mlir
|
|
|
|
#define DEBUG_TYPE "affine-loop-fusion"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::affine;
|
|
|
|
namespace {
|
|
/// Loop fusion pass. This pass currently supports a greedy fusion policy,
|
|
/// which fuses loop nests with single-writer/single-reader memref dependences
|
|
/// with the goal of improving locality.
|
|
// TODO: Support fusion of source loop nests which write to multiple
|
|
// memrefs, where each memref can have multiple users (if profitable).
|
|
struct LoopFusion : public affine::impl::AffineLoopFusionBase<LoopFusion> {
|
|
LoopFusion() = default;
|
|
LoopFusion(unsigned fastMemorySpace, uint64_t localBufSizeThresholdBytes,
|
|
bool maximalFusion, enum FusionMode affineFusionMode) {
|
|
this->fastMemorySpace = fastMemorySpace;
|
|
this->localBufSizeThreshold = localBufSizeThresholdBytes / 1024;
|
|
this->maximalFusion = maximalFusion;
|
|
this->affineFusionMode = affineFusionMode;
|
|
}
|
|
|
|
void runOnBlock(Block *block);
|
|
void runOnOperation() override;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
/// Returns true if node 'srcId' can be removed after fusing it with node
|
|
/// 'dstId'. The node can be removed if any of the following conditions are met:
|
|
/// 1. 'srcId' has no output dependences after fusion and no escaping memrefs.
|
|
/// 2. 'srcId' has no output dependences after fusion, has escaping memrefs
|
|
/// and the fusion slice is maximal.
|
|
/// 3. 'srcId' has output dependences after fusion, the fusion slice is
|
|
/// maximal and the fusion insertion point dominates all the dependences.
|
|
static bool canRemoveSrcNodeAfterFusion(
|
|
unsigned srcId, unsigned dstId, const ComputationSliceState &fusionSlice,
|
|
Operation *fusedLoopInsPoint, const DenseSet<Value> &escapingMemRefs,
|
|
MemRefDependenceGraph *mdg) {
|
|
|
|
Operation *dstNodeOp = mdg->getNode(dstId)->op;
|
|
bool hasOutDepsAfterFusion = false;
|
|
|
|
for (auto &outEdge : mdg->outEdges[srcId]) {
|
|
Operation *depNodeOp = mdg->getNode(outEdge.id)->op;
|
|
// Skip dependence with dstOp since it will be removed after fusion.
|
|
if (depNodeOp == dstNodeOp)
|
|
continue;
|
|
|
|
// Only fusion within the same block is supported. Use domination analysis
|
|
// when needed.
|
|
if (depNodeOp->getBlock() != dstNodeOp->getBlock())
|
|
return false;
|
|
|
|
// Check if the insertion point of the fused loop dominates the dependence.
|
|
// Otherwise, the src loop can't be removed.
|
|
if (fusedLoopInsPoint != depNodeOp &&
|
|
!fusedLoopInsPoint->isBeforeInBlock(depNodeOp)) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Src loop can't be removed: dst loop doesn't "
|
|
"dominate dependence\n");
|
|
return false;
|
|
}
|
|
|
|
hasOutDepsAfterFusion = true;
|
|
}
|
|
|
|
// If src loop has dependences after fusion or it writes to an live-out or
|
|
// escaping memref, we can only remove it if the fusion slice is maximal so
|
|
// that all the dependences are preserved.
|
|
if (hasOutDepsAfterFusion || !escapingMemRefs.empty()) {
|
|
std::optional<bool> isMaximal = fusionSlice.isMaximal();
|
|
if (!isMaximal) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Src loop can't be removed: can't determine "
|
|
"if fusion is maximal\n");
|
|
return false;
|
|
}
|
|
|
|
if (!*isMaximal) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Src loop can't be removed: fusion is not maximal\n");
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Returns in 'srcIdCandidates' the producer fusion candidates for consumer
|
|
/// 'dstId'. Candidates are sorted by node id order. This order corresponds to
|
|
/// the program order when the 'mdg' is created. However, program order is not
|
|
/// guaranteed and must not be required by the client. Program order won't be
|
|
/// held if the 'mdg' is reused from a previous fusion step or if the node
|
|
/// creation order changes in the future to support more advance cases.
|
|
// TODO: Move this to a loop fusion utility once 'mdg' is also moved.
|
|
static void getProducerCandidates(unsigned dstId, MemRefDependenceGraph *mdg,
|
|
SmallVectorImpl<unsigned> &srcIdCandidates) {
|
|
// Skip if no input edges along which to fuse.
|
|
if (mdg->inEdges.count(dstId) == 0)
|
|
return;
|
|
|
|
// Gather memrefs from loads in 'dstId'.
|
|
auto *dstNode = mdg->getNode(dstId);
|
|
DenseSet<Value> consumedMemrefs;
|
|
for (Operation *load : dstNode->loads)
|
|
consumedMemrefs.insert(cast<AffineReadOpInterface>(load).getMemRef());
|
|
|
|
// Traverse 'dstId' incoming edges and gather the nodes that contain a store
|
|
// to one of the consumed memrefs.
|
|
for (auto &srcEdge : mdg->inEdges[dstId]) {
|
|
auto *srcNode = mdg->getNode(srcEdge.id);
|
|
// Skip if 'srcNode' is not a loop nest.
|
|
if (!isa<AffineForOp>(srcNode->op))
|
|
continue;
|
|
|
|
if (any_of(srcNode->stores, [&](Operation *op) {
|
|
auto storeOp = cast<AffineWriteOpInterface>(op);
|
|
return consumedMemrefs.count(storeOp.getMemRef()) > 0;
|
|
}))
|
|
srcIdCandidates.push_back(srcNode->id);
|
|
}
|
|
|
|
llvm::sort(srcIdCandidates);
|
|
srcIdCandidates.erase(
|
|
std::unique(srcIdCandidates.begin(), srcIdCandidates.end()),
|
|
srcIdCandidates.end());
|
|
}
|
|
|
|
/// Returns in 'producerConsumerMemrefs' the memrefs involved in a
|
|
/// producer-consumer dependence between 'srcId' and 'dstId'.
|
|
static void
|
|
gatherProducerConsumerMemrefs(unsigned srcId, unsigned dstId,
|
|
MemRefDependenceGraph *mdg,
|
|
DenseSet<Value> &producerConsumerMemrefs) {
|
|
auto *dstNode = mdg->getNode(dstId);
|
|
auto *srcNode = mdg->getNode(srcId);
|
|
gatherProducerConsumerMemrefs(srcNode->stores, dstNode->loads,
|
|
producerConsumerMemrefs);
|
|
}
|
|
|
|
/// A memref escapes in the context of the fusion pass if either:
|
|
/// 1. it (or its alias) is a block argument, or
|
|
/// 2. created by an op not known to guarantee alias freedom,
|
|
/// 3. it (or its alias) are used by ops other than affine dereferencing ops
|
|
/// (e.g., by call op, memref load/store ops, alias creating ops, unknown ops,
|
|
/// terminator ops, etc.); such ops do not deference the memref in an affine
|
|
/// way.
|
|
static bool isEscapingMemref(Value memref, Block *block) {
|
|
Operation *defOp = memref.getDefiningOp();
|
|
// Check if 'memref' is a block argument.
|
|
if (!defOp)
|
|
return true;
|
|
|
|
// Check if this is defined to be an alias of another memref.
|
|
if (auto viewOp = dyn_cast<mlir::ViewLikeOpInterface>(defOp))
|
|
if (isEscapingMemref(viewOp.getViewSource(), block))
|
|
return true;
|
|
|
|
// Any op besides allocating ops wouldn't guarantee alias freedom
|
|
if (!hasSingleEffect<mlir::MemoryEffects::Allocate>(defOp, memref))
|
|
return true;
|
|
|
|
// Check if 'memref' is used by a non-deferencing op (including unknown ones)
|
|
// (e.g., call ops, alias creating ops, etc.).
|
|
return llvm::any_of(memref.getUsers(), [&](Operation *user) {
|
|
// Ignore users outside of `block`.
|
|
Operation *ancestorOp = block->getParent()->findAncestorOpInRegion(*user);
|
|
if (!ancestorOp)
|
|
return true;
|
|
if (ancestorOp->getBlock() != block)
|
|
return false;
|
|
return !isa<AffineMapAccessInterface>(*user);
|
|
});
|
|
}
|
|
|
|
/// Returns in 'escapingMemRefs' the memrefs from affine store ops in node 'id'
|
|
/// that escape the block or are accessed in a non-affine way.
|
|
static void gatherEscapingMemrefs(unsigned id, MemRefDependenceGraph *mdg,
|
|
DenseSet<Value> &escapingMemRefs) {
|
|
auto *node = mdg->getNode(id);
|
|
for (Operation *storeOp : node->stores) {
|
|
auto memref = cast<AffineWriteOpInterface>(storeOp).getMemRef();
|
|
if (escapingMemRefs.count(memref))
|
|
continue;
|
|
if (isEscapingMemref(memref, &mdg->block))
|
|
escapingMemRefs.insert(memref);
|
|
}
|
|
}
|
|
|
|
// Sinks all sequential loops to the innermost levels (while preserving
|
|
// relative order among them) and moves all parallel loops to the
|
|
// outermost (while again preserving relative order among them).
|
|
// This can increase the loop depth at which we can fuse a slice, since we are
|
|
// pushing loop carried dependence to a greater depth in the loop nest.
|
|
static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
|
|
assert(isa<AffineForOp>(node->op));
|
|
AffineForOp newRootForOp = sinkSequentialLoops(cast<AffineForOp>(node->op));
|
|
node->op = newRootForOp;
|
|
}
|
|
|
|
// Creates and returns a private (single-user) memref for fused loop rooted
|
|
// at 'forOp', with (potentially reduced) memref size based on the
|
|
// MemRefRegion written to by 'srcStoreOpInst' at depth 'dstLoopDepth'.
|
|
// TODO: consider refactoring the common code from generateDma and
|
|
// this one.
|
|
static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
|
|
unsigned dstLoopDepth,
|
|
std::optional<unsigned> fastMemorySpace,
|
|
uint64_t localBufSizeThreshold) {
|
|
Operation *forInst = forOp.getOperation();
|
|
|
|
// Create builder to insert alloc op just before 'forOp'.
|
|
OpBuilder b(forInst);
|
|
// Builder to create constants at the top level.
|
|
OpBuilder top(forInst->getParentRegion());
|
|
// Create new memref type based on slice bounds.
|
|
auto oldMemRef = cast<AffineWriteOpInterface>(srcStoreOpInst).getMemRef();
|
|
auto oldMemRefType = cast<MemRefType>(oldMemRef.getType());
|
|
unsigned rank = oldMemRefType.getRank();
|
|
|
|
// Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
|
|
MemRefRegion region(srcStoreOpInst->getLoc());
|
|
bool validRegion = succeeded(region.compute(srcStoreOpInst, dstLoopDepth));
|
|
(void)validRegion;
|
|
assert(validRegion && "unexpected memref region failure");
|
|
SmallVector<int64_t, 4> newShape;
|
|
std::vector<SmallVector<int64_t, 4>> lbs;
|
|
SmallVector<int64_t, 8> lbDivisors;
|
|
lbs.reserve(rank);
|
|
// Query 'region' for 'newShape' and lower bounds of MemRefRegion accessed
|
|
// by 'srcStoreOpInst' at depth 'dstLoopDepth'.
|
|
std::optional<int64_t> numElements =
|
|
region.getConstantBoundingSizeAndShape(&newShape, &lbs, &lbDivisors);
|
|
assert(numElements && "non-constant number of elts in local buffer");
|
|
|
|
const FlatAffineValueConstraints *cst = region.getConstraints();
|
|
// 'outerIVs' holds the values that this memory region is symbolic/parametric
|
|
// on; this would correspond to loop IVs surrounding the level at which the
|
|
// slice is being materialized.
|
|
SmallVector<Value, 8> outerIVs;
|
|
cst->getValues(rank, cst->getNumVars(), &outerIVs);
|
|
|
|
// Build 'rank' AffineExprs from MemRefRegion 'lbs'
|
|
SmallVector<AffineExpr, 4> offsets;
|
|
offsets.reserve(rank);
|
|
for (unsigned d = 0; d < rank; ++d) {
|
|
assert(lbs[d].size() == cst->getNumCols() - rank && "incorrect bound size");
|
|
|
|
AffineExpr offset = top.getAffineConstantExpr(0);
|
|
for (unsigned j = 0, e = cst->getNumCols() - rank - 1; j < e; j++) {
|
|
offset = offset + lbs[d][j] * top.getAffineDimExpr(j);
|
|
}
|
|
assert(lbDivisors[d] > 0);
|
|
offset =
|
|
(offset + lbs[d][cst->getNumCols() - 1 - rank]).floorDiv(lbDivisors[d]);
|
|
offsets.push_back(offset);
|
|
}
|
|
|
|
// Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
|
|
// by 'srcStoreOpInst'.
|
|
auto eltSize = getMemRefIntOrFloatEltSizeInBytes(oldMemRefType);
|
|
assert(eltSize && "memrefs with size elt types expected");
|
|
uint64_t bufSize = *eltSize * *numElements;
|
|
unsigned newMemSpace;
|
|
if (bufSize <= localBufSizeThreshold && fastMemorySpace.has_value()) {
|
|
newMemSpace = *fastMemorySpace;
|
|
} else {
|
|
newMemSpace = oldMemRefType.getMemorySpaceAsInt();
|
|
}
|
|
auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
|
|
{}, newMemSpace);
|
|
|
|
// Create new private memref for fused loop 'forOp'. 'newShape' is always
|
|
// a constant shape.
|
|
// TODO: Create/move alloc ops for private memrefs closer to their
|
|
// consumer loop nests to reduce their live range. Currently they are added
|
|
// at the beginning of the block, because loop nests can be reordered
|
|
// during the fusion pass.
|
|
Value newMemRef = top.create<memref::AllocOp>(forOp.getLoc(), newMemRefType);
|
|
|
|
// Build an AffineMap to remap access functions based on lower bound offsets.
|
|
SmallVector<AffineExpr, 4> remapExprs;
|
|
remapExprs.reserve(rank);
|
|
for (unsigned i = 0; i < rank; i++) {
|
|
auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
|
|
|
|
auto remapExpr =
|
|
simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
|
|
remapExprs.push_back(remapExpr);
|
|
}
|
|
|
|
auto indexRemap =
|
|
AffineMap::get(outerIVs.size() + rank, 0, remapExprs, forOp.getContext());
|
|
|
|
// Replace all users of 'oldMemRef' with 'newMemRef'.
|
|
LogicalResult res =
|
|
replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
|
|
/*extraOperands=*/outerIVs,
|
|
/*symbolOperands=*/{},
|
|
/*domOpFilter=*/&*forOp.getBody()->begin());
|
|
assert(succeeded(res) &&
|
|
"replaceAllMemrefUsesWith should always succeed here");
|
|
(void)res;
|
|
return newMemRef;
|
|
}
|
|
|
|
/// Walking from node 'srcId' to node 'dstId' (exclusive of 'srcId' and
|
|
/// 'dstId'), if there is any non-affine operation accessing 'memref', return
|
|
/// true. Otherwise, return false.
|
|
static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
|
|
Value memref,
|
|
MemRefDependenceGraph *mdg) {
|
|
auto *srcNode = mdg->getNode(srcId);
|
|
auto *dstNode = mdg->getNode(dstId);
|
|
Value::user_range users = memref.getUsers();
|
|
// For each MemRefDependenceGraph's node that is between 'srcNode' and
|
|
// 'dstNode' (exclusive of 'srcNodes' and 'dstNode'), check whether any
|
|
// non-affine operation in the node accesses the 'memref'.
|
|
for (auto &idAndNode : mdg->nodes) {
|
|
Operation *op = idAndNode.second.op;
|
|
// Take care of operations between 'srcNode' and 'dstNode'.
|
|
if (srcNode->op->isBeforeInBlock(op) && op->isBeforeInBlock(dstNode->op)) {
|
|
// Walk inside the operation to find any use of the memref.
|
|
// Interrupt the walk if found.
|
|
auto walkResult = op->walk([&](Operation *user) {
|
|
// Skip affine ops.
|
|
if (isa<AffineMapAccessInterface>(*user))
|
|
return WalkResult::advance();
|
|
// Find a non-affine op that uses the memref.
|
|
if (llvm::is_contained(users, user))
|
|
return WalkResult::interrupt();
|
|
return WalkResult::advance();
|
|
});
|
|
if (walkResult.wasInterrupted())
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Check whether a memref value in node 'srcId' has a non-affine that
|
|
/// is between node 'srcId' and node 'dstId' (exclusive of 'srcNode' and
|
|
/// 'dstNode').
|
|
static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
|
|
MemRefDependenceGraph *mdg) {
|
|
// Collect memref values in node 'srcId'.
|
|
auto *srcNode = mdg->getNode(srcId);
|
|
llvm::SmallDenseSet<Value, 2> memRefValues;
|
|
srcNode->op->walk([&](Operation *op) {
|
|
// Skip affine ops.
|
|
if (isa<AffineForOp>(op))
|
|
return WalkResult::advance();
|
|
for (Value v : op->getOperands())
|
|
// Collect memref values only.
|
|
if (isa<MemRefType>(v.getType()))
|
|
memRefValues.insert(v);
|
|
return WalkResult::advance();
|
|
});
|
|
// Looking for users between node 'srcId' and node 'dstId'.
|
|
return llvm::any_of(memRefValues, [&](Value memref) {
|
|
return hasNonAffineUsersOnThePath(srcId, dstId, memref, mdg);
|
|
});
|
|
}
|
|
|
|
// Checks the profitability of fusing a backwards slice of the loop nest
|
|
// surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
|
|
// The argument 'srcStoreOpInst' is used to calculate the storage reduction on
|
|
// the memref being produced and consumed, which is an input to the cost model.
|
|
// For producer-consumer fusion, 'srcStoreOpInst' will be the same as
|
|
// 'srcOpInst', as we are slicing w.r.t to that producer. For input-reuse
|
|
// fusion, 'srcOpInst' will be the src loop nest LoadOp which reads from the
|
|
// same memref as dst loop nest load ops, and 'srcStoreOpInst' will be the
|
|
// unique store op in the src node, which will be used to check that the write
|
|
// region is the same after input-reuse fusion. Computation slices are provided
|
|
// in 'depthSliceUnions' for each legal fusion depth. The maximal depth at which
|
|
// fusion is legal is provided in 'maxLegalFusionDepth'. Returns true if it is
|
|
// profitable to fuse the candidate loop nests. Returns false otherwise.
|
|
// `dstLoopDepth` is set to the most profitable depth at which to materialize
|
|
// the source loop nest slice.
|
|
// The profitability model executes the following steps:
|
|
// *) Computes the backward computation slice at 'srcOpInst'. This
|
|
// computation slice of the loop nest surrounding 'srcOpInst' is
|
|
// represented by modified src loop bounds in 'sliceState', which are
|
|
// functions of loop IVs in the loop nest surrounding 'srcOpInst'.
|
|
// *) Computes the cost of unfused src/dst loop nests (currently the cost of a
|
|
// loop nest is the total number of dynamic operation instances in the loop
|
|
// nest).
|
|
// *) Computes the cost of fusing a slice of the src loop nest into the dst
|
|
// loop nest at various values of dst loop depth, attempting to fuse
|
|
// the largest computation slice at the maximal dst loop depth (closest to
|
|
// the load) to minimize reuse distance and potentially enable subsequent
|
|
// load/store forwarding.
|
|
// NOTE: 'dstLoopDepth' refers to the loop depth within the destination loop
|
|
// nest, at which the src computation slice is inserted/fused.
|
|
// NOTE: We attempt to maximize the dst loop depth, but there are cases
|
|
// where a particular setting for 'dstLoopNest' might fuse an unsliced
|
|
// loop (within the src computation slice) at a depth which results in
|
|
// excessive recomputation (see unit tests for examples).
|
|
// *) Compares the total cost of the unfused loop nests to the min cost fused
|
|
// loop nest computed in the previous step, and returns true if the latter
|
|
// is lower.
|
|
// TODO: Extend profitability analysis to support scenarios with multiple
|
|
// stores.
|
|
static bool isFusionProfitable(Operation *srcOpInst, Operation *srcStoreOpInst,
|
|
AffineForOp dstForOp,
|
|
ArrayRef<ComputationSliceState> depthSliceUnions,
|
|
unsigned maxLegalFusionDepth,
|
|
unsigned *dstLoopDepth,
|
|
double computeToleranceThreshold) {
|
|
LLVM_DEBUG({
|
|
llvm::dbgs() << "Checking whether fusion is profitable between src op:\n";
|
|
llvm::dbgs() << ' ' << *srcOpInst << " and destination loop:\n";
|
|
llvm::dbgs() << dstForOp << "\n";
|
|
});
|
|
|
|
if (maxLegalFusionDepth == 0) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Can't fuse: maxLegalFusionDepth is 0\n");
|
|
return false;
|
|
}
|
|
|
|
// Compute cost of sliced and unsliced src loop nest.
|
|
SmallVector<AffineForOp, 4> srcLoopIVs;
|
|
getAffineForIVs(*srcOpInst, &srcLoopIVs);
|
|
|
|
// Walk src loop nest and collect stats.
|
|
LoopNestStats srcLoopNestStats;
|
|
if (!getLoopNestStats(srcLoopIVs[0], &srcLoopNestStats))
|
|
return false;
|
|
|
|
// Compute cost of dst loop nest.
|
|
LoopNestStats dstLoopNestStats;
|
|
if (!getLoopNestStats(dstForOp, &dstLoopNestStats))
|
|
return false;
|
|
|
|
// Search for min cost value for 'dstLoopDepth'. At each value of
|
|
// 'dstLoopDepth' from 'maxLegalLoopDepth' to '1', compute computation slice
|
|
// bounds between 'srcOpInst' and each op in 'dstOpinsts' (taking the union
|
|
// of these bounds). Next the union slice bounds are used to calculate
|
|
// the cost of the slice and the cost of the slice inserted into the dst
|
|
// loop nest at 'dstLoopDepth'.
|
|
uint64_t minFusedLoopNestComputeCost = std::numeric_limits<uint64_t>::max();
|
|
double maxStorageReduction = 0.0;
|
|
std::optional<uint64_t> sliceMemEstimate;
|
|
|
|
// The best loop depth at which to materialize the slice.
|
|
std::optional<unsigned> bestDstLoopDepth;
|
|
|
|
// Compute op instance count for the src loop nest without iteration slicing.
|
|
uint64_t srcLoopNestCost = getComputeCost(srcLoopIVs[0], srcLoopNestStats);
|
|
|
|
// Compute src loop nest write region size.
|
|
MemRefRegion srcWriteRegion(srcStoreOpInst->getLoc());
|
|
if (failed(srcWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0))) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Unable to compute MemRefRegion for source operation\n");
|
|
return false;
|
|
}
|
|
|
|
std::optional<int64_t> maybeSrcWriteRegionSizeBytes =
|
|
srcWriteRegion.getRegionSize();
|
|
if (!maybeSrcWriteRegionSizeBytes.has_value())
|
|
return false;
|
|
int64_t srcWriteRegionSizeBytes = *maybeSrcWriteRegionSizeBytes;
|
|
|
|
// Compute op instance count for the src loop nest.
|
|
uint64_t dstLoopNestCost = getComputeCost(dstForOp, dstLoopNestStats);
|
|
|
|
// Evaluate all depth choices for materializing the slice in the destination
|
|
// loop nest.
|
|
for (unsigned i = maxLegalFusionDepth; i >= 1; --i) {
|
|
const ComputationSliceState &slice = depthSliceUnions[i - 1];
|
|
// Skip slice union if it wasn't computed for this depth.
|
|
if (slice.isEmpty())
|
|
continue;
|
|
|
|
int64_t fusedLoopNestComputeCost;
|
|
if (!getFusionComputeCost(srcLoopIVs[0], srcLoopNestStats, dstForOp,
|
|
dstLoopNestStats, slice,
|
|
&fusedLoopNestComputeCost)) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Unable to compute fusion compute cost\n");
|
|
continue;
|
|
}
|
|
|
|
double additionalComputeFraction =
|
|
fusedLoopNestComputeCost /
|
|
(static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
|
|
1;
|
|
|
|
// Determine what the slice write MemRefRegion would be, if the src loop
|
|
// nest slice 'slice' were to be inserted into the dst loop nest at loop
|
|
// depth 'i'.
|
|
MemRefRegion sliceWriteRegion(srcStoreOpInst->getLoc());
|
|
if (failed(sliceWriteRegion.compute(srcStoreOpInst, /*loopDepth=*/0,
|
|
&slice))) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Failed to compute slice write region at loopDepth: " << i
|
|
<< "\n");
|
|
continue;
|
|
}
|
|
|
|
std::optional<int64_t> maybeSliceWriteRegionSizeBytes =
|
|
sliceWriteRegion.getRegionSize();
|
|
if (!maybeSliceWriteRegionSizeBytes.has_value() ||
|
|
*maybeSliceWriteRegionSizeBytes == 0) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Failed to get slice write region size at loopDepth: " << i
|
|
<< "\n");
|
|
continue;
|
|
}
|
|
int64_t sliceWriteRegionSizeBytes = *maybeSliceWriteRegionSizeBytes;
|
|
|
|
// If we are fusing for reuse, check that write regions remain the same.
|
|
// TODO: Write region check should check sizes and offsets in
|
|
// each dimension, so that we are sure they are covering the same memref
|
|
// region. Also, move this out to a isMemRefRegionSuperSet helper function.
|
|
if (srcOpInst != srcStoreOpInst &&
|
|
sliceWriteRegionSizeBytes != srcWriteRegionSizeBytes)
|
|
continue;
|
|
|
|
double storageReduction = static_cast<double>(srcWriteRegionSizeBytes) /
|
|
static_cast<double>(sliceWriteRegionSizeBytes);
|
|
|
|
LLVM_DEBUG({
|
|
std::stringstream msg;
|
|
msg << " evaluating fusion profitability at depth : " << i << "\n"
|
|
<< std::fixed << std::setprecision(2)
|
|
<< " additional compute fraction: "
|
|
<< 100.0 * additionalComputeFraction << "%\n"
|
|
<< " storage reduction factor: " << storageReduction << "x\n"
|
|
<< " fused nest cost: " << fusedLoopNestComputeCost << "\n"
|
|
<< " src write region size: " << srcWriteRegionSizeBytes << "\n"
|
|
<< " slice write region size: " << sliceWriteRegionSizeBytes
|
|
<< "\n";
|
|
llvm::dbgs() << msg.str();
|
|
});
|
|
|
|
// TODO: This is a placeholder cost model.
|
|
// Among all choices that add an acceptable amount of redundant computation
|
|
// (as per computeToleranceThreshold), we will simply pick the one that
|
|
// reduces the intermediary size the most.
|
|
if ((storageReduction > maxStorageReduction) &&
|
|
(additionalComputeFraction < computeToleranceThreshold)) {
|
|
maxStorageReduction = storageReduction;
|
|
bestDstLoopDepth = i;
|
|
minFusedLoopNestComputeCost = fusedLoopNestComputeCost;
|
|
sliceMemEstimate = sliceWriteRegionSizeBytes;
|
|
}
|
|
}
|
|
|
|
// A simple cost model: fuse if it reduces the memory footprint.
|
|
|
|
if (!bestDstLoopDepth) {
|
|
LLVM_DEBUG(
|
|
llvm::dbgs()
|
|
<< "All fusion choices involve more than the threshold amount of "
|
|
"redundant computation; NOT fusing.\n");
|
|
return false;
|
|
}
|
|
|
|
if (!bestDstLoopDepth) {
|
|
LLVM_DEBUG(llvm::dbgs() << "no fusion depth could be evaluated.\n");
|
|
return false;
|
|
}
|
|
|
|
// Set dstLoopDepth based on best values from search.
|
|
*dstLoopDepth = *bestDstLoopDepth;
|
|
|
|
LLVM_DEBUG(
|
|
llvm::dbgs() << " LoopFusion fusion stats:"
|
|
<< "\n best loop depth: " << bestDstLoopDepth
|
|
<< "\n src loop nest compute cost: " << srcLoopNestCost
|
|
<< "\n dst loop nest compute cost: " << dstLoopNestCost
|
|
<< "\n fused loop nest compute cost: "
|
|
<< minFusedLoopNestComputeCost << "\n");
|
|
|
|
auto dstMemSize = getMemoryFootprintBytes(dstForOp);
|
|
auto srcMemSize = getMemoryFootprintBytes(srcLoopIVs[0]);
|
|
|
|
std::optional<double> storageReduction;
|
|
|
|
if (!dstMemSize || !srcMemSize) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< " fusion memory benefit cannot be evaluated; NOT fusing.\n");
|
|
return false;
|
|
}
|
|
|
|
auto srcMemSizeVal = *srcMemSize;
|
|
auto dstMemSizeVal = *dstMemSize;
|
|
|
|
assert(sliceMemEstimate && "expected value");
|
|
auto fusedMem = dstMemSizeVal + *sliceMemEstimate;
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << " src mem: " << srcMemSizeVal << "\n"
|
|
<< " dst mem: " << dstMemSizeVal << "\n"
|
|
<< " fused mem: " << fusedMem << "\n"
|
|
<< " slice mem: " << sliceMemEstimate << "\n");
|
|
|
|
if (static_cast<long>(fusedMem) > srcMemSizeVal + dstMemSizeVal) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Fusion is not profitable; NOT fusing.\n");
|
|
return false;
|
|
}
|
|
storageReduction =
|
|
100.0 *
|
|
(1.0 - fusedMem / (static_cast<double>(srcMemSizeVal) + dstMemSizeVal));
|
|
|
|
double additionalComputeFraction =
|
|
100.0 * (minFusedLoopNestComputeCost /
|
|
(static_cast<double>(srcLoopNestCost) + dstLoopNestCost) -
|
|
1);
|
|
(void)additionalComputeFraction;
|
|
LLVM_DEBUG({
|
|
std::stringstream msg;
|
|
msg << " fusion is most profitable at depth " << *dstLoopDepth << " with "
|
|
<< std::setprecision(2) << additionalComputeFraction
|
|
<< "% redundant computation and a ";
|
|
msg << (storageReduction ? std::to_string(*storageReduction) : "<unknown>");
|
|
msg << "% storage reduction.\n";
|
|
llvm::dbgs() << msg.str();
|
|
});
|
|
|
|
return true;
|
|
}
|
|
|
|
namespace {
|
|
|
|
// GreedyFusion greedily fuses loop nests which have a producer/consumer or
|
|
// input-reuse relationship on a memref, with the goal of improving locality.
|
|
//
|
|
// The steps of the producer-consumer fusion algorithm are as follows:
|
|
//
|
|
// *) A worklist is initialized with node ids from the dependence graph.
|
|
// *) For each node id in the worklist:
|
|
// *) Pop an AffineForOp of the worklist. This 'dstAffineForOp' will be a
|
|
// candidate destination AffineForOp into which fusion will be attempted.
|
|
// *) Add each LoadOp currently in 'dstAffineForOp' into list 'dstLoadOps'.
|
|
// *) For each LoadOp in 'dstLoadOps' do:
|
|
// *) Look up dependent loop nests which have a single store op to the same
|
|
// memref.
|
|
// *) Check if dependences would be violated by the fusion.
|
|
// *) Get a computation slice of 'srcLoopNest', which adjusts its loop
|
|
// bounds to be functions of 'dstLoopNest' IVs and symbols.
|
|
// *) Fuse the 'srcLoopNest' computation slice into the 'dstLoopNest',
|
|
// at a loop depth determined by the cost model in 'isFusionProfitable'.
|
|
// *) Add the newly fused load/store operations to the state,
|
|
// and also add newly fused load ops to 'dstLoopOps' to be considered
|
|
// as fusion dst load ops in another iteration.
|
|
// *) Remove old src loop nest and its associated state.
|
|
//
|
|
// The steps of the input-reuse fusion algorithm are as follows:
|
|
//
|
|
// *) Initialize 'worklist' with node ids from the dependence graph.
|
|
// *) For each 'dstNode' in the worklist:
|
|
// *) Find a candidate sibling node 'sibNode' to fuse with 'dstNode' which
|
|
// loads from the same memref, but which has no dependence paths to/from.
|
|
// *) Get a computation slice of 'sibLoopNest', which adjusts its loop
|
|
// bounds to be functions of 'dstLoopNest' IVs and symbols.
|
|
// *) Fuse the 'sibLoopNest' computation slice into the 'dstLoopNest',
|
|
// at a loop depth determined by the cost model in 'isFusionProfitable'.
|
|
// This function also checks that the memref write region of 'sibLoopNest',
|
|
// is preserved in the fused loop nest.
|
|
// *) Update graph state to reflect the fusion of 'sibNode' into 'dstNode'.
|
|
//
|
|
// Given a graph where top-level operations are vertices in the set 'V' and
|
|
// edges in the set 'E' are dependences between vertices, this algorithm
|
|
// takes O(V) time for initialization, and has runtime O(V + E).
|
|
//
|
|
// This greedy algorithm is not 'maximal' due to the current restriction of
|
|
// fusing along single producer consumer edges, but there is a TODO: to fix
|
|
// this.
|
|
//
|
|
// TODO: Experiment with other fusion policies.
|
|
struct GreedyFusion {
|
|
public:
|
|
// The data dependence graph to traverse during fusion.
|
|
MemRefDependenceGraph *mdg;
|
|
// Worklist of graph nodes visited during the fusion pass.
|
|
SmallVector<unsigned, 8> worklist;
|
|
// Parameter for local buffer size threshold.
|
|
unsigned localBufSizeThreshold;
|
|
// Parameter for fast memory space.
|
|
std::optional<unsigned> fastMemorySpace;
|
|
// If true, ignore any additional (redundant) computation tolerance threshold
|
|
// that would have prevented fusion.
|
|
bool maximalFusion;
|
|
// The amount of additional computation that is tolerated while fusing
|
|
// pair-wise as a fraction of the total computation.
|
|
double computeToleranceThreshold;
|
|
|
|
using Node = MemRefDependenceGraph::Node;
|
|
|
|
GreedyFusion(MemRefDependenceGraph *mdg, unsigned localBufSizeThreshold,
|
|
std::optional<unsigned> fastMemorySpace, bool maximalFusion,
|
|
double computeToleranceThreshold)
|
|
: mdg(mdg), localBufSizeThreshold(localBufSizeThreshold),
|
|
fastMemorySpace(fastMemorySpace), maximalFusion(maximalFusion),
|
|
computeToleranceThreshold(computeToleranceThreshold) {}
|
|
|
|
/// Initializes 'worklist' with nodes from 'mdg'.
|
|
void init() {
|
|
// TODO: Add a priority queue for prioritizing nodes by different
|
|
// metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
|
|
worklist.clear();
|
|
for (auto &idAndNode : mdg->nodes) {
|
|
const Node &node = idAndNode.second;
|
|
worklist.push_back(node.id);
|
|
}
|
|
}
|
|
/// Run only sibling fusion on the `mdg`.
|
|
void runSiblingFusionOnly() {
|
|
fuseSiblingNodes();
|
|
eraseUnusedMemRefAllocations();
|
|
}
|
|
|
|
/// Run only producer/consumer fusion on the `mdg`.
|
|
void runProducerConsumerFusionOnly() {
|
|
fuseProducerConsumerNodes(
|
|
/*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
|
|
eraseUnusedMemRefAllocations();
|
|
}
|
|
|
|
// Run the GreedyFusion pass.
|
|
// *) First pass through the nodes fuses single-use producer nodes into their
|
|
// unique consumer.
|
|
// *) Second pass fuses sibling nodes which share no dependence edges.
|
|
// *) Third pass fuses any remaining producer nodes into their users.
|
|
void runGreedyFusion() {
|
|
// TODO: Run this repeatedly until a fixed-point is reached.
|
|
fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
|
|
fuseSiblingNodes();
|
|
fuseProducerConsumerNodes(
|
|
/*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
|
|
eraseUnusedMemRefAllocations();
|
|
}
|
|
|
|
/// Returns true if a private memref can be created for `memref` given
|
|
/// the fusion scenario reflected by the other arguments.
|
|
bool canCreatePrivateMemRef(Value memref,
|
|
const DenseSet<Value> &srcEscapingMemRefs,
|
|
unsigned producerId, unsigned consumerId,
|
|
bool removeSrcNode) {
|
|
const Node *consumerNode = mdg->getNode(consumerId);
|
|
// If `memref` is an escaping one, do not create a private memref
|
|
// for the below scenarios, since doing so will leave the escaping
|
|
// memref unmodified as all the writes originally meant for the
|
|
// escaping memref would be performed on the private memref:
|
|
// 1. The source is to be removed after fusion,
|
|
// OR
|
|
// 2. The destination writes to `memref`.
|
|
if (srcEscapingMemRefs.count(memref) > 0 &&
|
|
(removeSrcNode || consumerNode->getStoreOpCount(memref) > 0))
|
|
return false;
|
|
|
|
// Don't create a private memref if 'srcNode' has in edges on
|
|
// 'memref' or 'dstNode' has out edges on 'memref'.
|
|
if (mdg->getIncomingMemRefAccesses(producerId, memref) > 0 ||
|
|
mdg->getOutEdgeCount(consumerId, memref) > 0)
|
|
return false;
|
|
|
|
// If 'srcNode' will be removed but it has out edges on 'memref' to
|
|
// nodes other than 'dstNode', we have to preserve dependences and
|
|
// cannot create a private memref.
|
|
if (removeSrcNode &&
|
|
any_of(mdg->outEdges[producerId], [&](const auto &edge) {
|
|
return edge.value == memref && edge.id != consumerId;
|
|
}))
|
|
return false;
|
|
|
|
return true;
|
|
}
|
|
|
|
/// Perform fusions with node `dstId` as the destination of fusion, with
|
|
/// No fusion is performed when producers with a user count greater than
|
|
/// `maxSrcUserCount` for any of the memrefs involved.
|
|
void performFusionsIntoDest(unsigned dstId, unsigned maxSrcUserCount) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Evaluating dst loop " << dstId << "\n");
|
|
// Skip if this node was removed (fused into another node).
|
|
if (mdg->nodes.count(dstId) == 0)
|
|
return;
|
|
// Get 'dstNode' into which to attempt fusion.
|
|
auto *dstNode = mdg->getNode(dstId);
|
|
// Skip if 'dstNode' is not a loop nest.
|
|
if (!isa<AffineForOp>(dstNode->op))
|
|
return;
|
|
// Skip if 'dstNode' is a loop nest returning values.
|
|
// TODO: support loop nests that return values.
|
|
if (dstNode->op->getNumResults() > 0)
|
|
return;
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << "Evaluating dst loop " << dstId << "\n");
|
|
|
|
// Sink sequential loops in 'dstNode' (and thus raise parallel loops)
|
|
// while preserving relative order. This can increase the maximum loop
|
|
// depth at which we can fuse a slice of a producer loop nest into a
|
|
// consumer loop nest.
|
|
sinkSequentialLoops(dstNode);
|
|
auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
|
|
|
|
// Try to fuse 'dstNode' with candidate producer loops until a fixed point
|
|
// is reached. Fusing two loops may expose new fusion opportunities.
|
|
bool dstNodeChanged;
|
|
do {
|
|
// Gather src loop candidates for 'dstNode' and visit them in "quasi"
|
|
// reverse program order to minimize the number of iterations needed to
|
|
// reach the fixed point. Note that this is a best effort approach since
|
|
// 'getProducerCandidates' does not always guarantee that program order
|
|
// in 'srcIdCandidates'.
|
|
dstNodeChanged = false;
|
|
SmallVector<unsigned, 16> srcIdCandidates;
|
|
getProducerCandidates(dstId, mdg, srcIdCandidates);
|
|
|
|
for (unsigned srcId : llvm::reverse(srcIdCandidates)) {
|
|
// Get 'srcNode' from which to attempt fusion into 'dstNode'.
|
|
auto *srcNode = mdg->getNode(srcId);
|
|
auto srcAffineForOp = cast<AffineForOp>(srcNode->op);
|
|
LLVM_DEBUG(llvm::dbgs() << "Evaluating src loop " << srcId
|
|
<< " for dst loop " << dstId << "\n");
|
|
|
|
// Skip if 'srcNode' is a loop nest returning values.
|
|
// TODO: support loop nests that return values.
|
|
if (isa<AffineForOp>(srcNode->op) && srcNode->op->getNumResults() > 0)
|
|
continue;
|
|
|
|
DenseSet<Value> producerConsumerMemrefs;
|
|
gatherProducerConsumerMemrefs(srcId, dstId, mdg,
|
|
producerConsumerMemrefs);
|
|
|
|
// Skip if 'srcNode' out edge count on any memref is greater than
|
|
// 'maxSrcUserCount'.
|
|
if (any_of(producerConsumerMemrefs, [&](Value memref) {
|
|
return mdg->getOutEdgeCount(srcNode->id, memref) >
|
|
maxSrcUserCount;
|
|
}))
|
|
continue;
|
|
|
|
// Gather memrefs in 'srcNode' that are written and escape out of the
|
|
// block (e.g., memref block arguments, returned memrefs,
|
|
// memrefs passed to function calls, etc.).
|
|
DenseSet<Value> srcEscapingMemRefs;
|
|
gatherEscapingMemrefs(srcNode->id, mdg, srcEscapingMemRefs);
|
|
|
|
// Skip if there are non-affine operations in between the 'srcNode'
|
|
// and 'dstNode' using their memrefs. If so, we wouldn't be able to
|
|
// compute a legal insertion point for now. 'srcNode' and 'dstNode'
|
|
// memrefs with non-affine operation users would be considered
|
|
// escaping memrefs so we can limit this check to only scenarios with
|
|
// escaping memrefs.
|
|
if (!srcEscapingMemRefs.empty() &&
|
|
hasNonAffineUsersOnThePath(srcId, dstId, mdg)) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Can't fuse: non-affine users in between the loops\n");
|
|
continue;
|
|
}
|
|
|
|
// Compute an operation list insertion point for the fused loop
|
|
// nest which preserves dependences.
|
|
Operation *fusedLoopInsPoint =
|
|
mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
|
|
if (fusedLoopInsPoint == nullptr)
|
|
continue;
|
|
|
|
// It's possible this fusion is at an inner depth (i.e., there are
|
|
// common surrounding affine loops for the source and destination for
|
|
// ops). We need to get this number because the call to canFuseLoops
|
|
// needs to be passed the absolute depth. The max legal depth and the
|
|
// depths we try below are however *relative* and as such don't include
|
|
// the common depth.
|
|
SmallVector<AffineForOp, 4> surroundingLoops;
|
|
getAffineForIVs(*dstAffineForOp, &surroundingLoops);
|
|
unsigned numSurroundingLoops = surroundingLoops.size();
|
|
|
|
// Compute the innermost common loop depth for dstNode
|
|
// producer-consumer loads/stores.
|
|
SmallVector<Operation *, 2> dstMemrefOps;
|
|
for (Operation *op : dstNode->loads)
|
|
if (producerConsumerMemrefs.count(
|
|
cast<AffineReadOpInterface>(op).getMemRef()) > 0)
|
|
dstMemrefOps.push_back(op);
|
|
for (Operation *op : dstNode->stores)
|
|
if (producerConsumerMemrefs.count(
|
|
cast<AffineWriteOpInterface>(op).getMemRef()))
|
|
dstMemrefOps.push_back(op);
|
|
unsigned dstLoopDepthTest =
|
|
getInnermostCommonLoopDepth(dstMemrefOps) - numSurroundingLoops;
|
|
|
|
// Check the feasibility of fusing src loop nest into dst loop nest
|
|
// at loop depths in range [1, dstLoopDepthTest].
|
|
unsigned maxLegalFusionDepth = 0;
|
|
SmallVector<ComputationSliceState, 8> depthSliceUnions;
|
|
depthSliceUnions.resize(dstLoopDepthTest);
|
|
FusionStrategy strategy(FusionStrategy::ProducerConsumer);
|
|
for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
|
|
FusionResult result =
|
|
affine::canFuseLoops(srcAffineForOp, dstAffineForOp,
|
|
/*dstLoopDepth=*/i + numSurroundingLoops,
|
|
&depthSliceUnions[i - 1], strategy);
|
|
|
|
if (result.value == FusionResult::Success)
|
|
maxLegalFusionDepth = i;
|
|
}
|
|
|
|
if (maxLegalFusionDepth == 0) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Can't fuse: fusion is not legal at any depth\n");
|
|
continue;
|
|
}
|
|
|
|
// Check if fusion would be profitable. We skip profitability analysis
|
|
// for maximal fusion since we already know the maximal legal depth to
|
|
// fuse.
|
|
unsigned bestDstLoopDepth = maxLegalFusionDepth;
|
|
if (!maximalFusion) {
|
|
// Retrieve producer stores from the src loop.
|
|
SmallVector<Operation *, 2> producerStores;
|
|
for (Operation *op : srcNode->stores)
|
|
if (producerConsumerMemrefs.count(
|
|
cast<AffineWriteOpInterface>(op).getMemRef()))
|
|
producerStores.push_back(op);
|
|
|
|
// TODO: Suppport multiple producer stores in profitability
|
|
// analysis. We limit profitability analysis to only scenarios with
|
|
// a single producer store for now. Note that some multi-store
|
|
// producer scenarios will still go through profitability analysis
|
|
// if only one of the stores is involved the producer-consumer
|
|
// relationship of the candidate loops.
|
|
assert(!producerStores.empty() && "Expected producer store");
|
|
if (producerStores.size() > 1)
|
|
LLVM_DEBUG(llvm::dbgs() << "Skipping profitability analysis. Not "
|
|
"supported for this case\n");
|
|
else if (!isFusionProfitable(producerStores[0], producerStores[0],
|
|
dstAffineForOp, depthSliceUnions,
|
|
maxLegalFusionDepth, &bestDstLoopDepth,
|
|
computeToleranceThreshold))
|
|
continue;
|
|
}
|
|
|
|
assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
|
|
ComputationSliceState &bestSlice =
|
|
depthSliceUnions[bestDstLoopDepth - 1];
|
|
assert(!bestSlice.isEmpty() && "Missing slice union for depth");
|
|
|
|
// Determine if 'srcId' can be removed after fusion, taking into
|
|
// account remaining dependences, escaping memrefs and the fusion
|
|
// insertion point.
|
|
bool removeSrcNode = canRemoveSrcNodeAfterFusion(
|
|
srcId, dstId, bestSlice, fusedLoopInsPoint, srcEscapingMemRefs,
|
|
mdg);
|
|
|
|
DenseSet<Value> privateMemrefs;
|
|
for (Value memref : producerConsumerMemrefs) {
|
|
if (canCreatePrivateMemRef(memref, srcEscapingMemRefs, srcId, dstId,
|
|
removeSrcNode)) {
|
|
// Create a private version of this memref.
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Creating private memref for " << memref << '\n');
|
|
// Create a private version of this memref.
|
|
privateMemrefs.insert(memref);
|
|
}
|
|
}
|
|
|
|
// Fuse computation slice of 'srcLoopNest' into 'dstLoopNest'.
|
|
fuseLoops(srcAffineForOp, dstAffineForOp, bestSlice);
|
|
dstNodeChanged = true;
|
|
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Fused src loop " << srcId << " into dst loop " << dstId
|
|
<< " at depth " << bestDstLoopDepth << ":\n"
|
|
<< dstAffineForOp << "\n");
|
|
|
|
// Move 'dstAffineForOp' before 'insertPointInst' if needed.
|
|
if (fusedLoopInsPoint != dstAffineForOp)
|
|
dstAffineForOp->moveBefore(fusedLoopInsPoint);
|
|
|
|
// Update edges between 'srcNode' and 'dstNode'.
|
|
mdg->updateEdges(srcNode->id, dstNode->id, privateMemrefs,
|
|
removeSrcNode);
|
|
|
|
// Create private memrefs.
|
|
if (!privateMemrefs.empty()) {
|
|
// Gather stores for all the private-to-be memrefs.
|
|
DenseMap<Value, SmallVector<Operation *, 4>> privateMemRefToStores;
|
|
dstAffineForOp.walk([&](AffineWriteOpInterface storeOp) {
|
|
Value storeMemRef = storeOp.getMemRef();
|
|
if (privateMemrefs.count(storeMemRef) > 0)
|
|
privateMemRefToStores[storeMemRef].push_back(storeOp);
|
|
});
|
|
|
|
// Replace original memrefs with private memrefs. Note that all the
|
|
// loads and stores on these memrefs will be replaced with a new
|
|
// loads and stores. Any reference to the original ones becomes
|
|
// invalid after this point.
|
|
for (auto &memrefToStoresPair : privateMemRefToStores) {
|
|
// TODO: Use union of memref write regions to compute
|
|
// private memref footprint.
|
|
SmallVector<Operation *, 4> &storesForMemref =
|
|
memrefToStoresPair.second;
|
|
Value newMemRef = createPrivateMemRef(
|
|
dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
|
|
fastMemorySpace, localBufSizeThreshold);
|
|
// Create new node in dependence graph for 'newMemRef' alloc op.
|
|
unsigned newMemRefNodeId = mdg->addNode(newMemRef.getDefiningOp());
|
|
// Add edge from 'newMemRef' node to dstNode.
|
|
mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
|
|
}
|
|
// One or more entries for 'newMemRef' alloc op are inserted into
|
|
// the DenseMap mdg->nodes. Since an insertion may cause DenseMap to
|
|
// reallocate, update dstNode.
|
|
dstNode = mdg->getNode(dstId);
|
|
}
|
|
|
|
// Collect dst loop stats after memref privatization transformation.
|
|
LoopNestStateCollector dstLoopCollector;
|
|
dstLoopCollector.collect(dstAffineForOp);
|
|
|
|
// Clear and add back loads and stores.
|
|
mdg->clearNodeLoadAndStores(dstNode->id);
|
|
mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
|
|
dstLoopCollector.storeOpInsts);
|
|
|
|
if (removeSrcNode) {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Removing src loop " << srcId << " after fusion\n");
|
|
// srcNode is no longer valid after it is removed from mdg.
|
|
srcAffineForOp.erase();
|
|
mdg->removeNode(srcId);
|
|
srcNode = nullptr;
|
|
}
|
|
}
|
|
} while (dstNodeChanged);
|
|
}
|
|
|
|
/// Visit each node in the graph, and for each node, attempt to fuse it with
|
|
/// producer-consumer candidates. No fusion is performed when producers with a
|
|
/// user count greater than `maxSrcUserCount` for any of the memrefs involved
|
|
/// are encountered.
|
|
void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
|
|
LLVM_DEBUG(llvm::dbgs() << "--- Producer/Consumer Fusion ---\n");
|
|
init();
|
|
while (!worklist.empty()) {
|
|
unsigned dstId = worklist.back();
|
|
worklist.pop_back();
|
|
performFusionsIntoDest(dstId, maxSrcUserCount);
|
|
}
|
|
}
|
|
|
|
// Visits each node in the graph, and for each node, attempts to fuse it with
|
|
// its sibling nodes (nodes which share a parent, but no dependence edges).
|
|
void fuseSiblingNodes() {
|
|
LLVM_DEBUG(llvm::dbgs() << "--- Sibling Fusion ---\n");
|
|
init();
|
|
while (!worklist.empty()) {
|
|
unsigned dstId = worklist.back();
|
|
worklist.pop_back();
|
|
|
|
// Skip if this node was removed (fused into another node).
|
|
if (mdg->nodes.count(dstId) == 0)
|
|
continue;
|
|
// Get 'dstNode' into which to attempt fusion.
|
|
auto *dstNode = mdg->getNode(dstId);
|
|
// Skip if 'dstNode' is not a loop nest.
|
|
if (!isa<AffineForOp>(dstNode->op))
|
|
continue;
|
|
// Attempt to fuse 'dstNode' with its sibling nodes in the graph.
|
|
fuseWithSiblingNodes(dstNode);
|
|
}
|
|
}
|
|
|
|
// Attempt to fuse 'dstNode' with sibling nodes in the graph.
|
|
void fuseWithSiblingNodes(Node *dstNode) {
|
|
DenseSet<unsigned> visitedSibNodeIds;
|
|
std::pair<unsigned, Value> idAndMemref;
|
|
auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
|
|
|
|
while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
|
|
unsigned sibId = idAndMemref.first;
|
|
Value memref = idAndMemref.second;
|
|
// TODO: Check that 'sibStoreOpInst' post-dominates all other
|
|
// stores to the same memref in 'sibNode' loop nest.
|
|
auto *sibNode = mdg->getNode(sibId);
|
|
// Compute an operation list insertion point for the fused loop
|
|
// nest which preserves dependences.
|
|
assert(sibNode->op->getBlock() == dstNode->op->getBlock());
|
|
Operation *insertPointInst =
|
|
sibNode->op->isBeforeInBlock(dstNode->op)
|
|
? mdg->getFusedLoopNestInsertionPoint(sibNode->id, dstNode->id)
|
|
: mdg->getFusedLoopNestInsertionPoint(dstNode->id, sibNode->id);
|
|
if (insertPointInst == nullptr)
|
|
continue;
|
|
|
|
// Check if fusion would be profitable and at what depth.
|
|
|
|
// Get unique 'sibNode' load op to 'memref'.
|
|
SmallVector<Operation *, 2> sibLoadOpInsts;
|
|
sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
|
|
// Currently findSiblingNodeToFuse searches for siblings with one load.
|
|
assert(sibLoadOpInsts.size() == 1);
|
|
Operation *sibLoadOpInst = sibLoadOpInsts[0];
|
|
|
|
// Gather 'dstNode' load ops to 'memref'.
|
|
SmallVector<Operation *, 2> dstLoadOpInsts;
|
|
dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
|
|
|
|
// It's possible this fusion is at an inner depth (i.e., there are common
|
|
// surrounding affine loops for the source and destination for ops). We
|
|
// need to get this number because the call to canFuseLoops needs to be
|
|
// passed the absolute depth. The max legal depth and the depths we try
|
|
// below are however *relative* and as such don't include the common
|
|
// depth.
|
|
SmallVector<AffineForOp, 4> surroundingLoops;
|
|
getAffineForIVs(*dstAffineForOp, &surroundingLoops);
|
|
unsigned numSurroundingLoops = surroundingLoops.size();
|
|
SmallVector<AffineForOp, 4> dstLoopIVs;
|
|
getAffineForIVs(*dstLoadOpInsts[0], &dstLoopIVs);
|
|
unsigned dstLoopDepthTest = dstLoopIVs.size() - numSurroundingLoops;
|
|
auto sibAffineForOp = cast<AffineForOp>(sibNode->op);
|
|
|
|
// Compute loop depth and slice union for fusion.
|
|
SmallVector<ComputationSliceState, 8> depthSliceUnions;
|
|
depthSliceUnions.resize(dstLoopDepthTest);
|
|
unsigned maxLegalFusionDepth = 0;
|
|
FusionStrategy strategy(memref);
|
|
for (unsigned i = 1; i <= dstLoopDepthTest; ++i) {
|
|
FusionResult result =
|
|
affine::canFuseLoops(sibAffineForOp, dstAffineForOp,
|
|
/*dstLoopDepth=*/i + numSurroundingLoops,
|
|
&depthSliceUnions[i - 1], strategy);
|
|
|
|
if (result.value == FusionResult::Success)
|
|
maxLegalFusionDepth = i;
|
|
}
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << "Max legal depth for fusion: "
|
|
<< maxLegalFusionDepth << '\n');
|
|
|
|
// Skip if fusion is not feasible at any loop depths.
|
|
if (maxLegalFusionDepth == 0)
|
|
continue;
|
|
|
|
unsigned bestDstLoopDepth = maxLegalFusionDepth;
|
|
if (!maximalFusion) {
|
|
// Check if fusion would be profitable. For sibling fusion, the sibling
|
|
// load op is treated as the src "store" op for fusion profitability
|
|
// purposes. The footprint of the load in the slice relative to the
|
|
// unfused source's determines reuse.
|
|
if (!isFusionProfitable(sibLoadOpInst, sibLoadOpInst, dstAffineForOp,
|
|
depthSliceUnions, maxLegalFusionDepth,
|
|
&bestDstLoopDepth, computeToleranceThreshold))
|
|
continue;
|
|
}
|
|
|
|
assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
|
|
assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
|
|
"Fusion depth has no computed slice union");
|
|
// Check if source loop is being inserted in the innermost
|
|
// destination loop. Based on this, the fused loop may be optimized
|
|
// further inside `fuseLoops`.
|
|
bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest);
|
|
// Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
|
|
affine::fuseLoops(sibAffineForOp, dstAffineForOp,
|
|
depthSliceUnions[bestDstLoopDepth - 1],
|
|
isInnermostInsertion);
|
|
|
|
auto dstForInst = cast<AffineForOp>(dstNode->op);
|
|
// Update operation position of fused loop nest (if needed).
|
|
if (insertPointInst != dstForInst) {
|
|
dstForInst->moveBefore(insertPointInst);
|
|
}
|
|
// Update data dependence graph state post fusion.
|
|
updateStateAfterSiblingFusion(sibNode, dstNode);
|
|
}
|
|
}
|
|
|
|
// Searches block argument uses and the graph from 'dstNode' looking for a
|
|
// fusion candidate sibling node which shares no dependences with 'dstNode'
|
|
// but which loads from the same memref. Returns true and sets
|
|
// 'idAndMemrefToFuse' on success. Returns false otherwise.
|
|
bool findSiblingNodeToFuse(Node *dstNode,
|
|
DenseSet<unsigned> *visitedSibNodeIds,
|
|
std::pair<unsigned, Value> *idAndMemrefToFuse) {
|
|
// Returns true if 'sibNode' can be fused with 'dstNode' for input reuse
|
|
// on 'memref'.
|
|
auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
|
|
// Skip if 'outEdge' is not a read-after-write dependence.
|
|
// TODO: Remove restrict to single load op restriction.
|
|
if (sibNode->getLoadOpCount(memref) != 1)
|
|
return false;
|
|
// Skip if there exists a path of dependent edges between
|
|
// 'sibNode' and 'dstNode'.
|
|
if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
|
|
mdg->hasDependencePath(dstNode->id, sibNode->id))
|
|
return false;
|
|
// Skip sib node if it loads to (and stores from) the same memref on
|
|
// which it also has an input dependence edge.
|
|
DenseSet<Value> loadAndStoreMemrefSet;
|
|
sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
|
|
if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
|
|
return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
|
|
}))
|
|
return false;
|
|
|
|
// Check that all stores are to the same memref if any.
|
|
DenseSet<Value> storeMemrefs;
|
|
for (auto *storeOpInst : sibNode->stores) {
|
|
storeMemrefs.insert(
|
|
cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
|
|
}
|
|
if (storeMemrefs.size() > 1)
|
|
return false;
|
|
|
|
// Skip if a memref value in one node is used by a non-affine memref
|
|
// access that lies between 'dstNode' and 'sibNode'.
|
|
if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
|
|
hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
|
|
return false;
|
|
return true;
|
|
};
|
|
|
|
// Search for siblings which load the same memref block argument.
|
|
Block *block = dstNode->op->getBlock();
|
|
for (unsigned i = 0, e = block->getNumArguments(); i != e; ++i) {
|
|
for (Operation *user : block->getArgument(i).getUsers()) {
|
|
auto loadOp = dyn_cast<AffineReadOpInterface>(user);
|
|
if (!loadOp)
|
|
continue;
|
|
// Gather loops surrounding 'use'.
|
|
SmallVector<AffineForOp, 4> loops;
|
|
getAffineForIVs(*user, &loops);
|
|
// Skip 'use' if it is not within a loop nest.
|
|
// Find the surrounding affine.for nested immediately within the
|
|
// block.
|
|
auto *it = llvm::find_if(loops, [&](AffineForOp loop) {
|
|
return loop->getBlock() == &mdg->block;
|
|
});
|
|
// Skip 'use' if it is not within a loop nest in `block`.
|
|
if (it == loops.end())
|
|
continue;
|
|
Node *sibNode = mdg->getForOpNode(*it);
|
|
assert(sibNode != nullptr);
|
|
// Skip 'use' if it not a sibling to 'dstNode'.
|
|
if (sibNode->id == dstNode->id)
|
|
continue;
|
|
// Skip 'use' if it has been visited.
|
|
if (visitedSibNodeIds->count(sibNode->id) > 0)
|
|
continue;
|
|
// Skip 'use' if it does not load from the same memref as 'dstNode'.
|
|
auto memref = loadOp.getMemRef();
|
|
if (dstNode->getLoadOpCount(memref) == 0)
|
|
continue;
|
|
// Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
|
|
if (canFuseWithSibNode(sibNode, memref)) {
|
|
visitedSibNodeIds->insert(sibNode->id);
|
|
idAndMemrefToFuse->first = sibNode->id;
|
|
idAndMemrefToFuse->second = memref;
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Search for siblings by following edges through an intermediate src node.
|
|
// Collect candidate 'dstNode' input edges in 'inEdges'.
|
|
SmallVector<MemRefDependenceGraph::Edge, 2> inEdges;
|
|
mdg->forEachMemRefInputEdge(
|
|
dstNode->id, [&](MemRefDependenceGraph::Edge inEdge) {
|
|
// Add 'inEdge' if it is a read-after-write dependence.
|
|
if (dstNode->getLoadOpCount(inEdge.value) > 0 &&
|
|
mdg->getNode(inEdge.id)->getStoreOpCount(inEdge.value) > 0)
|
|
inEdges.push_back(inEdge);
|
|
});
|
|
|
|
// Search for sibling nodes to fuse by visiting output edges from each input
|
|
// edge in 'inEdges'.
|
|
for (auto &inEdge : inEdges) {
|
|
// Collect candidate output edges from each node 'inEdge.id' in 'inEdges'.
|
|
SmallVector<MemRefDependenceGraph::Edge, 2> outEdges;
|
|
mdg->forEachMemRefOutputEdge(
|
|
inEdge.id, [&](MemRefDependenceGraph::Edge outEdge) {
|
|
unsigned sibNodeId = outEdge.id;
|
|
if (visitedSibNodeIds->count(sibNodeId) > 0)
|
|
return;
|
|
// Skip output edge if not a sibling using the same memref.
|
|
if (outEdge.id == dstNode->id || outEdge.value != inEdge.value)
|
|
return;
|
|
auto *sibNode = mdg->getNode(sibNodeId);
|
|
if (!isa<AffineForOp>(sibNode->op))
|
|
return;
|
|
// Check if 'sibNode/dstNode' can be input-reuse fused on 'memref'.
|
|
if (canFuseWithSibNode(sibNode, outEdge.value)) {
|
|
// Add candidate 'outEdge' to sibling node.
|
|
outEdges.push_back(outEdge);
|
|
}
|
|
});
|
|
|
|
// Add first candidate if any were returned.
|
|
if (!outEdges.empty()) {
|
|
visitedSibNodeIds->insert(outEdges[0].id);
|
|
idAndMemrefToFuse->first = outEdges[0].id;
|
|
idAndMemrefToFuse->second = outEdges[0].value;
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/// Update data dependence graph state to reflect sibling fusion of 'sibNode'
|
|
/// into 'dstNode'.
|
|
void updateStateAfterSiblingFusion(Node *sibNode, Node *dstNode) {
|
|
// Update 'sibNode' and 'dstNode' input/output edges to reflect fusion.
|
|
mdg->updateEdges(sibNode->id, dstNode->id);
|
|
|
|
// Collect dst loop stats after memref privatization transformation.
|
|
auto dstForInst = cast<AffineForOp>(dstNode->op);
|
|
LoopNestStateCollector dstLoopCollector;
|
|
dstLoopCollector.collect(dstForInst);
|
|
// Clear and add back loads and stores
|
|
mdg->clearNodeLoadAndStores(dstNode->id);
|
|
mdg->addToNode(dstNode->id, dstLoopCollector.loadOpInsts,
|
|
dstLoopCollector.storeOpInsts);
|
|
// Remove old sibling loop nest if it no longer has outgoing dependence
|
|
// edges, and it does not write to a memref which escapes the block.
|
|
if (mdg->getOutEdgeCount(sibNode->id) == 0) {
|
|
Operation *op = sibNode->op;
|
|
mdg->removeNode(sibNode->id);
|
|
op->erase();
|
|
}
|
|
}
|
|
|
|
// Clean up any allocs with no users.
|
|
void eraseUnusedMemRefAllocations() {
|
|
for (auto &pair : mdg->memrefEdgeCount) {
|
|
if (pair.second > 0)
|
|
continue;
|
|
auto memref = pair.first;
|
|
// Skip if there exist other uses (return operation or function calls).
|
|
if (!memref.use_empty())
|
|
continue;
|
|
// Use list expected to match the dep graph info.
|
|
auto *op = memref.getDefiningOp();
|
|
if (isa_and_nonnull<memref::AllocOp>(op))
|
|
op->erase();
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
/// Run fusion on `block`.
|
|
void LoopFusion::runOnBlock(Block *block) {
|
|
MemRefDependenceGraph g(*block);
|
|
if (!g.init()) {
|
|
LLVM_DEBUG(llvm::dbgs() << "MDG init failed\n");
|
|
return;
|
|
}
|
|
|
|
std::optional<unsigned> fastMemorySpaceOpt;
|
|
if (fastMemorySpace.hasValue())
|
|
fastMemorySpaceOpt = fastMemorySpace;
|
|
unsigned localBufSizeThresholdBytes = localBufSizeThreshold * 1024;
|
|
GreedyFusion fusion(&g, localBufSizeThresholdBytes, fastMemorySpaceOpt,
|
|
maximalFusion, computeToleranceThreshold);
|
|
|
|
if (affineFusionMode == FusionMode::ProducerConsumer)
|
|
fusion.runProducerConsumerFusionOnly();
|
|
else if (affineFusionMode == FusionMode::Sibling)
|
|
fusion.runSiblingFusionOnly();
|
|
else
|
|
fusion.runGreedyFusion();
|
|
}
|
|
|
|
void LoopFusion::runOnOperation() {
|
|
// Call fusion on every op that has at least two affine.for nests (in post
|
|
// order).
|
|
getOperation()->walk([&](Operation *op) {
|
|
for (Region ®ion : op->getRegions()) {
|
|
for (Block &block : region.getBlocks()) {
|
|
auto affineFors = block.getOps<AffineForOp>();
|
|
if (!affineFors.empty() && !llvm::hasSingleElement(affineFors))
|
|
runOnBlock(&block);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
std::unique_ptr<Pass> mlir::affine::createLoopFusionPass(
|
|
unsigned fastMemorySpace, uint64_t localBufSizeThreshold,
|
|
bool maximalFusion, enum FusionMode affineFusionMode) {
|
|
return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
|
|
maximalFusion, affineFusionMode);
|
|
}
|