OSchip
/
llvm-project
mirror of https://github.com/THU-DSP-LAB/llvm-project.git
13
0
Fork 6
Code Issues Packages Projects Releases Wiki Activity
llvm-project/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

979 lines
41 KiB
C++
Raw Blame History

//===- Fusion.cpp - Implementation of linalg 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 the linalg dialect Fusion pass.
//
//===----------------------------------------------------------------------===//
#include "PassDetail.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include <set>
#define DEBUG_TYPE "linalg-fusion"
using namespace mlir;
using namespace mlir::linalg;
using llvm::dbgs;
/// Implements a simple high-level fusion pass on linalg structured operations.
///
/// In each block, linalg ops are processed in reverse textual order.
/// Given a linalg op `O`, fusion occurs by:
/// 1. inspecting the linalg ops that write into the views read by `O`. There
/// are 2 cases:
/// a) buffer case: use the SSA value of the views and a simple alias
/// analysis on subview ops to determine producer-consumer dependences;
/// b) tensor case: use SSA use-def chains on extract_slice ops;
/// 2. greedily fuse the linalg ops that produce the subview/extract_slice.
/// 3. inspect the fused ops and determine whether they have other remaining
/// LinalgOp uses. If not, then erase the original producing linalg op.
///
/// More advanced use cases, analyses as well as profitability heuristics are
/// left for future work.
struct ShapeDimension {
Value shape;
unsigned dimension;
};
// Given an `op`, returns the first (`shape`, `dimension`) pair that identifies
// the loop range at `loopDepth`. The semantics of the loopToOperandRangesMaps
// guarantees at least one such dimension is found. If multiple candidates exist
// they must agree by construction (i.e. have the same size) and we just return
// the first one.
static ShapeDimension
getShapeDefiningLoopRange(LinalgOp op, unsigned loopDepth,
bool fromSubViewOpOnly = false) {
// Iterate over the inputs and outputs in order.
// Extract the subranges from the linearized ranges.
for (OpOperand *opOperand : op.getInputAndOutputOperands()) {
// The method `getRangeFromOperandShape` requires using SubViewOp or
// ExtractSliceOps. If the value isn't defined from there continue.
// todo: The method should be adapted to get the values from
// `ViewInterface`. The interface needs a `getOrCreateRanges` method which
// currently returns a `linalg.range`. The fix here is to move this op to
// `std` dialect and add the method to `ViewInterface`.
if (fromSubViewOpOnly &&
!isa_and_nonnull<memref::SubViewOp, tensor::ExtractSliceOp>(
opOperand->get().getDefiningOp()))
continue;
AffineMap map = op.getTiedIndexingMap(opOperand);
LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange I/O idx: "
<< opOperand->getOperandNumber() << "\n");
LLVM_DEBUG(llvm::dbgs()
<< "getShapeDefiningLoopRange map: " << map << "\n");
SmallVector<Value, 8> shapeRanges(map.getNumResults(), nullptr);
for (auto en : llvm::enumerate(map.getResults())) {
auto dimExpr = en.value().dyn_cast<AffineDimExpr>();
if (!dimExpr)
continue;
if (loopDepth == en.value().cast<AffineDimExpr>().getPosition()) {
LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange loopDepth: "
<< loopDepth << "\n");
LLVM_DEBUG(llvm::dbgs() << "getShapeDefiningLoopRange shape: "
<< opOperand->get() << "\n");
return ShapeDimension{opOperand->get(),
static_cast<unsigned>(en.index())};
}
}
}
llvm_unreachable("Expect to be able to extract a shape defining loop range");
}
// Return tiled operands for the fused producer op. When fusing into
// `linalg.tiled_loop` one has to update `input` and `output` arguments of the
// loop correspondingly.
// Each input tensor of the producer op has to be added to `inputs` of the
// `tiled_loop` if it is not present there already. Each output tensor has to
// be added either to `inputs` or to `outputs` of `linalg.tiled_loop` depending
// on whether the correponding result is an input or an output to the loop.
//
// NOTE: This way of updating the arguments of the `tiled_loop` assumes that the
// intermediate result is not used by any other operation but the consumer. A
// more generic way is to append all missing output tensors of the producer to
// the tiled loop outputs and hence modify the number of the results, since we
// would need to add the intermediate results to `linalg.yield`. After that a
// canonicalization pass would move the unused output args of the `tiled_loop`
// to the `input` section.
static SmallVector<Value> getTiledOperands(OpBuilder &b, LinalgOp producer) {
auto tiledLoop = dyn_cast<TiledLoopOp>(b.getBlock()->getParentOp());
if (!tiledLoop)
return producer.getInputAndOutputOperands();
SmallVector<Value> tiledOperands;
assert(producer.hasTensorSemantics() &&
"only fusion on tensors is currently supported for TiledLinalgOp");
for (OpOperand *producerInput : producer.getInputOperands()) {
OpOperand *addedInput = tiledLoop.findInputOperand(producerInput->get());
if (addedInput == nullptr)
addedInput = &tiledLoop.appendInputOperand(b, producerInput->get());
BlockArgument addedBlockArg = tiledLoop.getTiedBlockArgument(*addedInput);
tiledOperands.push_back(addedBlockArg);
}
for (OpOperand *producerOutput : producer.getOutputOperands()) {
OpResult result = producer.getTiedOpResult(producerOutput);
OpOperand *resultInputOperand = tiledLoop.findInputOperand(result);
OpOperand *resultOutputOperand = tiledLoop.findOutputOperand(result);
assert((resultInputOperand != nullptr) ^ (resultOutputOperand != nullptr) &&
"The result should be present in `input` or `output` args of "
"`tiled_loop");
bool isInput = resultInputOperand;
int opNumber = isInput ? resultInputOperand->getOperandNumber()
: resultOutputOperand->getOperandNumber();
OpOperand *addedOutput = tiledLoop.findOutputOperand(producerOutput->get());
if (addedOutput == nullptr)
addedOutput =
isInput ? &tiledLoop.appendInputOperand(b, producerOutput->get())
: &tiledLoop.appendOutputOperand(b, producerOutput->get());
OpOperand &resultOperand = tiledLoop->getOpOperand(opNumber);
auto addedBlockArg = tiledLoop.getTiedBlockArgument(*addedOutput);
auto resultOperandBlockArg = tiledLoop.getTiedBlockArgument(resultOperand);
resultOperandBlockArg.replaceAllUsesWith(addedBlockArg);
tiledLoop.eraseOperand(b, resultOperand);
tiledOperands.push_back(addedBlockArg);
}
return tiledOperands;
}
/// Fuses the producer by cloning the `producer`. The `fusedLoopsAndRanges`
/// provides the loop range information for the fused loops. The rest are
/// obtained from the producer itself, since they are not tiled + fused.
static LinalgOp fuse(OpBuilder &b, LinalgOp producer,
const DenseMap<unsigned, Range> &fusedLoopsAndRanges) {
SmallVector<Value, 8> ivs, tileSizes, sizeBounds;
SmallVector<Range, 8> loopRanges;
Location loc = producer.getLoc();
auto zero = b.create<ConstantIndexOp>(loc, 0);
auto one = b.create<ConstantIndexOp>(loc, 1);
for (unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
auto it = fusedLoopsAndRanges.find(i);
if (it != fusedLoopsAndRanges.end()) {
ivs.push_back(it->second.offset);
tileSizes.push_back(it->second.size);
sizeBounds.push_back(nullptr);
loopRanges.push_back(it->second);
LLVM_DEBUG(llvm::dbgs() << "tiled loop#" << i << " with LoopRange "
<< loopRanges.back() << "\n");
} else {
auto shapeDim = getShapeDefiningLoopRange(producer, i);
Value dim = createOrFoldDimOp(b, loc, shapeDim.shape, shapeDim.dimension);
tileSizes.push_back(zero);
sizeBounds.push_back(dim);
loopRanges.push_back(Range{zero, dim, one});
LLVM_DEBUG(llvm::dbgs() << "full loop#" << i << " with LoopRange "
<< loopRanges.back() << "\n");
}
}
SmallVector<Value, 8> clonedShapes;
clonedShapes.reserve(producer.getNumInputsAndOutputs());
// Compute subranges for all tensor input/output operands.
clonedShapes.append(makeTiledShapes(b, loc, producer,
getTiledOperands(b, producer), ivs,
tileSizes, sizeBounds));
// Iterate over the results in order.
// Extract the subtensor type from the linearized range.
// Since we do not enforce any canonicalizations on the fly, this is always
// fully dynamic at construction time.
SmallVector<Type, 4> resultTypes;
resultTypes.reserve(producer->getNumResults());
for (RankedTensorType t : producer.getOutputTensorTypes()) {
unsigned rank = t.getRank();
SmallVector<int64_t, 4> staticOffsetsVector(
rank, ShapedType::kDynamicStrideOrOffset);
SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize);
SmallVector<int64_t, 4> staticStridesVector(
rank, ShapedType::kDynamicStrideOrOffset);
resultTypes.push_back(tensor::ExtractSliceOp::inferResultType(
t.cast<RankedTensorType>(), staticOffsetsVector, staticSizesVector,
staticStridesVector));
}
Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes);
// When the producer has index semantics, we have to transform the indices of
// the producer according to the tiling of the consumer, i.e. offset them by
// the values computed in `loopRanges`.
if (producer.hasIndexSemantics()) {
assert(clonedOp->getNumRegions() == 1 &&
clonedOp->getRegion(0).getBlocks().size() == 1 &&
"expected producer to have one block.");
// Shift all indices by the tile offset.
Block &block = clonedOp->getRegion(0).front();
for (IndexOp indexOp : block.getOps<IndexOp>()) {
OpBuilder::InsertionGuard g(b);
b.setInsertionPointAfter(indexOp);
AffineExpr index, offset;
bindDims(b.getContext(), index, offset);
AffineApplyOp applyOp = b.create<AffineApplyOp>(
indexOp.getLoc(), index + offset,
ValueRange{indexOp.getResult(), loopRanges[indexOp.dim()].offset});
indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
}
}
return clonedOp;
}
/// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is
/// expected to be defined by a subview op or an extract_slice op.
static Range getRangeFromOperandShape(OpBuilder &b, Location loc,
Value shapedOperand, unsigned dim) {
Operation *shapeProducingOp = shapedOperand.getDefiningOp();
if (auto subViewOp = dyn_cast<memref::SubViewOp>(shapeProducingOp))
return subViewOp.getOrCreateRanges(b, loc)[dim];
if (auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(shapeProducingOp))
return sliceOp.getOrCreateRanges(b, loc)[dim];
llvm_unreachable("SubviewOp or ExtractSliceOp expected");
}
/// Fuses the producer into the loop immediately enclosing the consumer.
/// This is achieved by "recomputing" the producer at the time it
/// is needed just before the consumer.
static LinalgOp fuse(OpBuilder &b, LinalgOp producerOp, AffineMap producerMap,
OpOperand &consumerOpOperand) {
LLVM_DEBUG(llvm::dbgs() << "Producer map: " << producerMap << "\n");
DenseMap<unsigned, Range> fusedLoopsAndRanges;
Value shapedOperand = consumerOpOperand.get();
for (auto en : llvm::enumerate(producerMap.getResults())) {
unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition();
fusedLoopsAndRanges[posInProducerLoop] = getRangeFromOperandShape(
b, consumerOpOperand.getOwner()->getLoc(), shapedOperand, en.index());
}
return fuse(b, producerOp, fusedLoopsAndRanges);
}
// Encode structural fusion safety preconditions.
// Some of these will be lifted in the future with better analysis.
static bool isStructurallyFusableProducer(LinalgOp producer, Value consumedView,
LinalgOp consumer) {
assert(producer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
assert(consumer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
if (producer.getNumOutputs() != 1) {
LLVM_DEBUG(llvm::dbgs() << "\nNot structurally fusable (multi-output)");
return false;
}
// Only fuse when the producer block dominates.
DominanceInfo dom(producer.getOperation());
if (!dom.dominates(producer->getBlock(), consumer->getBlock())) {
LLVM_DEBUG(
llvm::dbgs()
<< "\nNot structurally fusable (producer block does not dominate)");
return false;
}
return true;
}
bool mlir::linalg::isProducerLastWriteOfView(const LinalgDependenceGraph &graph,
LinalgOp consumer,
Value consumedView,
LinalgOp producer) {
assert(producer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
assert(consumer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
// Make some simple structural checks that alleviate the need for more
// complex analyses.
if (!isStructurallyFusableProducer(producer, consumedView, consumer)) {
LLVM_DEBUG(llvm::dbgs() << "\n***Not static last write due to structure:\t"
<< *producer.getOperation());
return false;
}
// Check for any interleaved write to consumedView.
if (!graph.findCoveringWrites(producer, consumer, consumedView).empty()) {
LLVM_DEBUG(llvm::dbgs() << "\n***Not fusable due to interleaved write:\t"
<< *producer.getOperation());
return false;
}
return true;
}
bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
LinalgOp consumer, Value consumedView,
LinalgOp producer) {
assert(producer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
assert(consumer.hasBufferSemantics() &&
"expected linalg op with buffer semantics");
if (!isProducerLastWriteOfView(graph, consumer, consumedView, producer))
return false;
// Check for any fusion-preventing dependence to any shape read/written that
// would violate dependences.
if (!graph.findCoveringDependences(producer, consumer).empty()) {
LLVM_DEBUG(llvm::dbgs()
<< "\n***Not fusable due to an interleaved dependence:\t"
<< *producer.getOperation());
return false;
}
if (auto convOp = dyn_cast<linalg::ConvOp>(producer.getOperation())) {
// TODO: add a level of indirection to linalg.generic.
if (convOp.padding())
return false;
}
if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {
// TODO: add a level of indirection to linalg.generic.
if (convOp.padding())
return false;
}
return true;
}
/// For `consumer` with buffer semantics, find the Linalg operation on buffers
/// that is the last writer of `consumerOpOperand`. For now the fusable
/// dependence is returned as an instance of the `dependenceGraph`.
static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
findFusableProducer(OpOperand &consumerOpOperand,
const LinalgDependenceGraph &dependenceGraph) {
LLVM_DEBUG(llvm::dbgs() << "findFusableProducer for: "
<< consumerOpOperand.get() << " @"
<< consumerOpOperand.getOperandNumber() << " in "
<< *consumerOpOperand.getOwner() << "\n");
LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
if (!consumerOp)
return {};
// Only consider RAW and WAW atm.
for (auto depType : {
LinalgDependenceGraph::DependenceType::RAW,
LinalgDependenceGraph::DependenceType::WAW,
}) {
LLVM_DEBUG(llvm::dbgs()
<< "Dependencies into: " << *consumerOp.getOperation() << "\n");
for (auto dependence : llvm::make_filter_range(
dependenceGraph.getDependencesInto(consumerOp, depType),
[&](LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
LLVM_DEBUG(llvm::dbgs() << "Inspect dependence btw: "
<< elem.getIndexingValue() << " and "
<< elem.getDependentValue() << "\n");
Value v = elem.getIndexingValue();
Optional<unsigned> operandNum =
elem.getIndexingOpViewOperandNum();
return isa<LinalgOp>(elem.getDependentOp()) &&
v == consumerOpOperand.get() && operandNum &&
operandNum.getValue() ==
consumerOpOperand.getOperandNumber();
})) {
// Consumer consumes this view, `isStructurallyFusableProducer` also
// checks whether it is a strict subview of the producer view.
auto producer = cast<LinalgOp>(dependence.getDependentOp());
LLVM_DEBUG(llvm::dbgs()
<< "\n"
<< LinalgDependenceGraph::getDependenceTypeStr(depType)
<< "producer: " << *dependence.getDependentOp()
<< " view: " << dependence.getDependentValue() << "\n");
// If the producer and consumer have tensor semantics, the only dependence
// between them is through a RAW dependence and they are fusable by
// construction. For buffer semantics need additional checks.
if (producer.hasBufferSemantics() && consumerOp.hasBufferSemantics() &&
isFusableInto(dependenceGraph, consumerOp, consumerOpOperand.get(),
producer))
return dependence;
if (producer.hasTensorSemantics() && consumerOp.hasTensorSemantics()) {
assert(dependence.dependenceType ==
LinalgDependenceGraph::DependenceType::RAW);
return dependence;
}
}
}
return {};
}
Optional<FusionInfo>
mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, OpOperand &consumerOpOperand,
const LinalgDependenceGraph &graph) {
Optional<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependence =
findFusableProducer(consumerOpOperand, graph);
if (!fusableDependence)
return llvm::None;
LinalgOp producerOp = dyn_cast<LinalgOp>(fusableDependence->getDependentOp());
if (!producerOp)
return llvm::None;
// If producer is already in the same block as consumer, we are done.
if (consumerOpOperand.get().getParentBlock() ==
fusableDependence->getDependentValue().getParentBlock())
return llvm::None;
Optional<AffineMap> producerMap =
fusableDependence->getDependentOpViewIndexingMap();
if (!producerMap)
return llvm::None;
// Must be a subview or an extract_slice to guarantee there are loops we can
// fuse into.
auto subView = consumerOpOperand.get().getDefiningOp<memref::SubViewOp>();
if (!subView) {
LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subview)");
return llvm::None;
}
// Fuse `producer` just before `consumer`.
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(consumerOpOperand.getOwner());
LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: "
<< *consumerOpOperand.getOwner() << "\n");
auto fusedProducer = fuse(b, producerOp, *producerMap, consumerOpOperand);
return FusionInfo{producerOp, fusedProducer};
}
/// Walk back use-def chain through scf::For yields.
/// Sets `producer` and `outputIndex` if it finds a producer LinalgOp
// TODO(ravishankarm, ntv): This can be moved into the dependence graphs
// dependence tracking since the dependence tracking is similar to what is done
// w.r.t to buffers.
static void getProducerOfTensor(Value tensor, OpResult &opResult) {
if (!tensor.getType().isa<RankedTensorType>())
return;
while (true) {
LLVM_DEBUG(llvm::dbgs() << "\ngetProducerOfTensor: " << tensor);
if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) {
opResult = tensor.cast<OpResult>();
return;
}
if (auto sliceOp = tensor.getDefiningOp<tensor::ExtractSliceOp>()) {
tensor = sliceOp.source();
continue;
}
if (auto blockArg = tensor.dyn_cast<BlockArgument>()) {
if (auto forOp = blockArg.getDefiningOp<scf::ForOp>()) {
tensor = *(forOp.getIterOperands().begin() + blockArg.getArgNumber());
continue;
}
}
return;
}
}
Optional<FusionInfo>
mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
Value inputTensor = consumerOpOperand.get();
OpResult producerOpResult;
getProducerOfTensor(inputTensor, producerOpResult);
if (!producerOpResult) {
LLVM_DEBUG(llvm::dbgs() << "\nUnable to find producer");
return {};
}
return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
}
Optional<FusionInfo>
mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
OpOperand &consumerOpOperand) {
auto producerOp = dyn_cast<LinalgOp>(producerOpResult.getOwner());
if (!producerOp)
return llvm::None;
LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
if (!consumerOp)
return llvm::None;
Value inputTensor = consumerOpOperand.get();
// Must be an extract_slice op to guarantee there are loops we can fuse into.
auto sliceOp = inputTensor.getDefiningOp<tensor::ExtractSliceOp>();
if (!sliceOp) {
LLVM_DEBUG(llvm::dbgs()
<< "\nNot fusable, not an extract_slice op: " << inputTensor);
return {};
}
// If producer is already in the same block as consumer, we are done.
if (consumerOpOperand.get().getParentBlock() ==
producerOpResult.getParentBlock())
return {};
// Insert fused `producer` just before `consumer`.
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(consumerOp);
LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumerOp << "\n");
OpOperand *opOperand =
producerOp.getOutputOperand(producerOpResult.getResultNumber());
LinalgOp fusedProducer =
fuse(b, producerOp, producerOp.getTiedIndexingMap(opOperand),
consumerOpOperand);
// Replace use.
// Canonicalizations are not guaranteed to have happened before constructing
// `fusedProducer`. In the tensor case this can result in temporary type
// mismatches. Insert a `tensor.cast` op to propagate the transformation
// invariant that types are compatible.
Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
Type consumerType = consumerOpOperand.get().getType();
if (consumerType != def.getType())
def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
consumerOpOperand.set(def);
return FusionInfo{cast<LinalgOp>(producerOpResult.getOwner()), fusedProducer};
}
/// Prune all dimensions that are of reduction iterator type from `map`.
static AffineMap pruneReductionDimsFromMap(ArrayRef<Attribute> iteratorTypes,
AffineMap map) {
llvm::SmallDenseSet<unsigned> projectedDims;
for (auto attr : llvm::enumerate(iteratorTypes)) {
if (!isParallelIterator(attr.value()))
projectedDims.insert(attr.index());
}
return getProjectedMap(map, projectedDims);
}
/// Returns the mapping from iterations in the consumer that write to the same
/// location as the iterations in the producer. To do so use
/// - indexing map of the fused view in the consumer : consumerIndexMap
/// - indexing map of the fused view in the producer : producerIndexMap
/// consumerLoopToProducerLoop =
/// inverse(producerIndexMap).compose(consumerIndexMap)
static Optional<AffineMap> getConsumerLoopToProducerLoopMap(
LinalgDependenceGraph::LinalgDependenceGraphElem dependence) {
auto producer = dyn_cast<LinalgOp>(dependence.getDependentOp());
if (!producer)
return None;
Optional<AffineMap> producerIndexingMap =
dependence.getDependentOpViewIndexingMap();
Optional<AffineMap> consumerIndexingMap =
dependence.getIndexingOpViewIndexingMap();
if (!producerIndexingMap || !consumerIndexingMap)
return None;
AffineMap prunedProducerIndexingMap = pruneReductionDimsFromMap(
producer.iterator_types().getValue(), *producerIndexingMap);
if (!prunedProducerIndexingMap.isPermutation())
return None;
if (consumerIndexingMap->getNumResults() !=
prunedProducerIndexingMap.getNumResults())
return None;
LLVM_DEBUG({
llvm::dbgs() << "\t producerMap : ";
producerIndexingMap->print(llvm::dbgs());
llvm::dbgs() << " pruned : ";
prunedProducerIndexingMap.print(llvm::dbgs());
llvm::dbgs() << "\n";
llvm::dbgs() << "\t consumerMap : ";
consumerIndexingMap->print(llvm::dbgs());
llvm::dbgs() << "\n";
});
AffineMap invProducerIndexMap = inversePermutation(prunedProducerIndexingMap);
if (!invProducerIndexMap)
return None;
return invProducerIndexMap.compose(*consumerIndexingMap);
}
/// Given a projected permutation `map`, returns true if the map changes the
/// order in which the fused loop dimension appear.
static bool doesTransposeAccess(AffineMap map,
const std::set<unsigned> &fusableLoops) {
Optional<unsigned> lastFusableLoop;
for (unsigned pos : llvm::map_range(map.getResults(), [](AffineExpr expr) {
return expr.cast<AffineDimExpr>().getPosition();
})) {
if (!fusableLoops.count(pos))
continue;
if (!lastFusableLoop) {
lastFusableLoop = pos;
continue;
}
if (pos <= lastFusableLoop.getValue())
return true;
lastFusableLoop = pos;
}
return false;
}
/// Returns the positions of the loop in `op` that can be tiled based on the
/// operations that are to be fused with it. For example, in a
///
/// linalg.matmul ins(%a, %b : ...) outs(%c : ...)
///
/// if the producer of %a needs to be fused with this op, only the `i` loop of
/// the matmul can be tiled while fusing. If producer of %a, and %b are to be
/// fused, then no loops can be tiled while fusing. The conditions used are:
/// 1. Only parallel loops can be used for tile + fuse. Find the number of
/// common outer parallel loops between the op and its producers being fused.
/// 2. Of the parallel loops only some can be fused. Only those loops can be
/// fused such where the fusable loops iteration space only touches one tile
/// of the fused operation. This is because the producer (which is writing
/// the fused subview) has update semantics.
///
/// Since an inverse computation is needed, we need to consider the projection
/// of the producerIndexMap w.r.t the parallel loops. The actual fusable loops
/// are the dimensions of the consumerLoopToProducerLoop map that correspond to
/// parallel loops and appear in the result of the map
///
/// Example 1:
/// linalg.fill(%cst, %c)
/// linalg.matmul ins(%a, %b) outs(%c)
/// Number of parallel loops : 2
/// producerIndexMap = affine_map<(i, j) ->(i , j)>
/// consumerIndexMap = affine_map<(i, j, k) -> (i, j)>
/// consumerLoopToProducerLoop = affine_map<(i, j, k) -> (i, j)>
/// Fused dimensions : i, j
///
/// Example 2:
/// linalg.matmul ins(%a, %b) outs(%c)
/// linalg.generic {indexing_maps = [affine_map<(i, j) -> (j, i)>, ...
/// iterator_types = ["parallel", "parallel"]}
/// ins(%c) ...
///
/// Number of parallel loops = 2:
/// producerIndexMap (projected to parallel loops) =
/// affine_map<(i, j) -> (i, j)>
/// consumerLoopToProducerLoop2 = affine_map<(i, j) -> (j, i)>
/// Fused dimensions : i, j
///
/// Example 3:
/// linalg.copy(%s, %b)
/// linalg.matmul ins(%a, %b) outs(%c)
///
/// Number of parallel loops = 2
/// produceIndexMap : affine_map<(i, j) -> (i, j)>
/// consumerLoopToProduceLoops = affine_map<(i, j, k) -> (k, j)>
/// submap with only parallel loops = affine_map<(i, j) -> (j)>
/// Fused dimensions : j
static std::set<unsigned>
collectFusableLoops(ArrayRef<LinalgOp> ops,
const FusableOpDependencesTy &fusableDependences) {
assert(!ops.empty());
auto getNumOuterParallelLoops = [](LinalgOp linalgOp) {
return linalgOp.iterator_types()
.getValue()
.take_while([](Attribute attr) -> bool {
return attr.cast<StringAttr>().getValue() ==
getParallelIteratorTypeName();
})
.size();
};
size_t numOuterParallelLoops = getNumOuterParallelLoops(ops.back());
for (auto op : ops.drop_back()) {
numOuterParallelLoops =
std::min(numOuterParallelLoops, getNumOuterParallelLoops(op));
}
std::set<unsigned> fusableLoops;
auto range = llvm::seq<unsigned>(0, numOuterParallelLoops);
fusableLoops.insert(range.begin(), range.end());
for (auto op : reverse(ops)) {
for (auto dependence : fusableDependences.lookup(op)) {
LLVM_DEBUG({
llvm::dbgs() << "\t fusable :";
for (unsigned i : fusableLoops)
llvm::dbgs() << " " << i;
llvm::dbgs() << "\n";
});
Optional<AffineMap> consumerLoopToProducerLoop =
getConsumerLoopToProducerLoopMap(dependence);
if (!consumerLoopToProducerLoop) {
op.emitRemark("failed to get map from consumer loop to producer loop");
return {};
}
// todo: This condition is only an implementation limitation. When fusing
// the operation, if the accesses in the producer/consumer are transposes
// of each other, the loop bounds for the tiled producer can be
// manipulated accordingly. This requires some additional bookkeeping in
// the implementation of tile+fuse that is deferred to later.
if (doesTransposeAccess(*consumerLoopToProducerLoop, fusableLoops)) {
op.emitRemark("unhandled fusion when fusion requires permutation");
return {};
}
std::set<unsigned> candidates;
for (AffineExpr expr : consumerLoopToProducerLoop->getResults()) {
unsigned position = expr.cast<AffineDimExpr>().getPosition();
if (fusableLoops.count(position))
candidates.insert(position);
}
LLVM_DEBUG({
llvm::dbgs() << "\t candidates :";
for (unsigned i : candidates)
llvm::dbgs() << " " << i;
llvm::dbgs() << "\n";
});
if (candidates.empty())
return {};
std::swap(candidates, fusableLoops);
}
}
return fusableLoops;
}
/// Find all dependences that are fusable.
FusableOpDependencesTy mlir::linalg::findAllFusableDependences(
ArrayRef<LinalgOp> ops, const LinalgDependenceGraph &dependenceGraph) {
FusableOpDependencesTy fusableDependences;
DenseMap<Operation *, SmallVector<AffineMap, 1>> fusedProducerIndexingMap;
for (LinalgOp op : reverse(ops)) {
for (OpOperand *opOperand : op.getInputAndOutputOperands()) {
Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
fusableDependence = findFusableProducer(*opOperand, dependenceGraph);
if (!fusableDependence)
continue;
LinalgOp producerOp =
dyn_cast<LinalgOp>(fusableDependence->getDependentOp());
if (!producerOp)
continue;
// Do not fuse dependences that are to operations not in the same basic
// block. This avoid moving fused operations across loops that might
// themselves carry dependency making the fusion illegal.
if (producerOp->getBlock() != op->getBlock())
continue;
// Make sure that the indexing map of the view used for fusion in the
// producer is a projected permutation.
Optional<AffineMap> producerMap =
fusableDependence->getDependentOpViewIndexingMap();
Optional<AffineMap> consumerMap =
fusableDependence->getIndexingOpViewIndexingMap();
assert(
consumerMap &&
"unable to find indexing map of operand/result of indexing OpView");
fusedProducerIndexingMap[producerOp.getOperation()].push_back(
*consumerMap);
if (!producerMap || !producerMap->isProjectedPermutation() ||
!consumerMap->isProjectedPermutation())
continue;
fusableDependences[producerOp.getOperation()].push_back(
*fusableDependence);
}
}
// TODO: Currently fusion would not be legal if the fusable dependence is to
// the same producer but different indexing map in the consumer. Fix this, but
// in the meanwhile disallow such a fusion.
for (auto useIndexingMapsList : fusedProducerIndexingMap) {
AffineMap map1 = useIndexingMapsList.second.front();
for (AffineMap map2 :
ArrayRef<AffineMap>(useIndexingMapsList.second).drop_front()) {
if (map1 != map2) {
fusableDependences.erase(useIndexingMapsList.first);
break;
}
}
}
return fusableDependences;
}
/// Tile the fused loops in the root operation, by setting the tile sizes for
/// all other loops to zero (those will be tiled later).
static Optional<TiledLinalgOp>
tileRootOperation(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizeVector,
const LinalgTilingOptions &options,
const std::set<unsigned> &fusedLoops) {
SmallVector<Value, 4> tileSizes(tileSizeVector.begin(), tileSizeVector.end());
auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0);
for (unsigned i = 0, e = tileSizes.size(); i != e; ++i)
if (!fusedLoops.count(i))
tileSizes[i] = zero;
LinalgTilingOptions tileFusedLoopsOptions = options;
tileFusedLoopsOptions.setTileSizes(tileSizes);
return tileLinalgOp(b, op, tileFusedLoopsOptions);
}
/// Fuse the operations in `fusionCandidates` with `tiledOp`. Latter is expected
/// to be a tiled operation such that it is valid to fuse all operations in
/// `fusionCandidates`, i.e. move the operation within the inter-tile loops of
/// `tiledOp`.
static SmallVector<LinalgOp, 1>
fuseOperations(OpBuilder &b, LinalgOp rootOp, TiledLinalgOp tiledLinalgOp,
ArrayRef<LinalgOp> fusionCandidates,
const FusableOpDependencesTy &fusableDependences,
const std::set<unsigned> &fusedLoops) {
LinalgOp tiledOp = tiledLinalgOp.op;
OpBuilder::InsertionGuard guard(b);
b.setInsertionPoint(tiledOp);
DenseMap<unsigned, Range> fusedLoopsAndRanges;
for (unsigned loop : fusedLoops) {
ShapeDimension shapeDim = getShapeDefiningLoopRange(tiledOp, loop, true);
fusedLoopsAndRanges[loop] = getRangeFromOperandShape(
b, tiledOp.getLoc(), shapeDim.shape, shapeDim.dimension);
}
SmallVector<LinalgOp, 1> fusedOps(fusionCandidates.size());
DenseMap<Operation *, LinalgOp> origOpToFusedOp;
origOpToFusedOp[rootOp.getOperation()] = tiledOp;
for (auto candidate : enumerate(llvm::reverse(fusionCandidates))) {
LinalgOp origOp = candidate.value();
LinalgOp fusedOp = fuse(b, origOp, fusedLoopsAndRanges);
origOpToFusedOp[origOp.getOperation()] = fusedOp;
fusedOps[fusionCandidates.size() - candidate.index() - 1] = fusedOp;
// Prepare the builder for the next insertion point.
auto guard = llvm::make_scope_exit([&]() { b.setInsertionPoint(fusedOp); });
if (!origOp.hasTensorSemantics())
continue;
// If the producer consumer operations are linalg operations on tensors, the
// dependence is due to value produced (as a return tensor) by the producer
// and used in the consumer. The returned value of the fused op needs to be
// made the operand of the tiled/fused consumer operation. By construction
// the value returned by the producer is the value used by the consumer.
for (auto &dependence : fusableDependences.lookup(origOp.getOperation())) {
if (dependence.dependenceType !=
LinalgDependenceGraph::DependenceType::RAW)
continue;
unsigned resultIndex =
dependence.getDependentOpViewResultNum().getValue();
LinalgOp consumer = origOpToFusedOp.lookup(dependence.getIndexingOp());
if (!consumer)
continue;
Value replacementValue = fusedOp.getOperation()->getResult(resultIndex);
consumer.getOperation()->setOperand(
dependence.getIndexingOpViewOperandNum().getValue(),
replacementValue);
}
// At this point, all Linalg uses of the tensors produced by `origOp` have
// been replaced. However, there may still be "output tensor"-like uses
// coming from WAW dependencies.
// All these uses are iter_args of the outermost loop (TODO: add a check).
// Such iter_args uses serve 2 purposes:
// 1. give a shape to the output
// 2. encode destructive updates that may be inplaceable by bufferization.
// To keep the second type of information while letting the unfused op die
// unused, we need to forward the producer output operand.
if (auto forOp = dyn_cast<scf::ForOp>(tiledLinalgOp.loops.front())) {
for (auto &operand : forOp.getIterOpOperands()) {
if (auto opResult = operand.get().dyn_cast<OpResult>()) {
if (opResult.getOwner() == origOp) {
Value output =
origOp.getOutputOperand(opResult.getResultNumber())->get();
assert(output.getType().isa<RankedTensorType>());
operand.set(output);
}
}
}
}
}
return fusedOps;
}
static Optional<TiledAndFusedLinalgOps>
tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
const LinalgDependenceGraph &dependenceGraph,
const LinalgTilingOptions &tilingOptions) {
if (ops.size() < 2)
return llvm::None;
LinalgOp rootOp = ops.back();
if (!llvm::all_of(
ops,
[](LinalgOp linalgOp) { return linalgOp.hasBufferSemantics(); }) &&
!llvm::all_of(ops, [](LinalgOp linalgOp) {
return linalgOp.hasTensorSemantics();
})) {
rootOp.emitError(
"unable to fuse operations that have tensor semantics with operations "
"that have buffer semantics and viceversa.");
return llvm::None;
}
// TODO: Support interchange with tile + fuse. This might actually help do
// better fusion.
if (!tilingOptions.interchangeVector.empty()) {
rootOp.emitRemark("unable to handle tile and fuse with interchange");
return llvm::None;
}
OpBuilder::InsertionGuard guard(b);
b.setInsertionPoint(rootOp);
// Find all the producers.
LLVM_DEBUG(llvm::dbgs() << "findAllFusableDependences\n");
FusableOpDependencesTy fusableDependences =
findAllFusableDependences(ops, dependenceGraph);
if (fusableDependences.empty()) {
LLVM_DEBUG(llvm::dbgs() << "no fusable dependencies found\n");
return llvm::None;
}
TiledAndFusedLinalgOps ret;
// Find the loops that can be tiled and fused.
LLVM_DEBUG(llvm::dbgs() << "collectFusableLoops\n");
ret.fusedLoopDims = collectFusableLoops(ops, fusableDependences);
// If there are no fusable dependences or there are no tile+fusable loops,
// just return.
if (ret.fusedLoopDims.empty()) {
LLVM_DEBUG(llvm::dbgs() << "no fusable loops found\n");
return llvm::None;
}
// Tile the fused loops in the last operation in the list.
SmallVector<Value, 4> tileSizeVector =
tilingOptions.tileSizeComputationFunction(b, rootOp);
Optional<TiledLinalgOp> tiledRootOp = tileRootOperation(
b, rootOp, tileSizeVector, tilingOptions, ret.fusedLoopDims);
if (!tiledRootOp) {
rootOp.emitRemark("failed to tile the fused loops");
return llvm::None;
}
ret.op = tiledRootOp->op;
ret.fusedLoops.assign(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
// Fuse the other operations into the fused inter-tile loops produced above.
ret.fusedProducers = fuseOperations(b, rootOp, *tiledRootOp, ops.drop_back(),
fusableDependences, ret.fusedLoopDims);
return ret;
}
Optional<TiledAndFusedLinalgOps>
mlir::linalg::tileAndFuseLinalgOps(OpBuilder &b, ArrayRef<LinalgOp> ops,
const LinalgDependenceGraph &dependenceGraph,
const LinalgTilingOptions &tilingOptions) {
switch (tilingOptions.loopType) {
case LinalgTilingLoopType::Loops:
case LinalgTilingLoopType::ParallelLoops:
case LinalgTilingLoopType::TiledLoops:
return tileAndFuseLinalgOpsImpl(b, ops, dependenceGraph, tilingOptions);
default:;
}
return llvm::None;
}
Reference in New Issue View Git Blame Copy Permalink