27 #include "llvm/ADT/MapVector.h"
28 #include "llvm/ADT/ScopeExit.h"
29 #include "llvm/Support/CommandLine.h"
30 #include "llvm/Support/Debug.h"
35 #define DEBUG_TYPE "linalg-fusion"
68 bool fromSubViewOpOnly =
false) {
71 for (
OpOperand &opOperand : op->getOpOperands()) {
78 if (fromSubViewOpOnly &&
79 !isa_and_nonnull<memref::SubViewOp, tensor::ExtractSliceOp>(
80 opOperand.get().getDefiningOp()))
83 AffineMap map = op.getMatchingIndexingMap(&opOperand);
84 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange I/O idx: "
85 << opOperand.getOperandNumber() <<
"\n");
86 LLVM_DEBUG(llvm::dbgs()
87 <<
"getShapeDefiningLoopRange map: " << map <<
"\n");
89 auto dimExpr = dyn_cast<AffineDimExpr>(en.value());
92 if (loopDepth == cast<AffineDimExpr>(en.value()).getPosition()) {
93 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange loopDepth: "
94 << loopDepth <<
"\n");
95 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange shape: "
96 << opOperand.get() <<
"\n");
98 static_cast<unsigned>(en.index())};
102 llvm_unreachable(
"Expect to be able to extract a shape defining loop range");
106 return producer->getOperands();
118 for (
unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
122 sizeBounds.push_back(dim);
123 auto it = fusedLoopsAndRanges.find(i);
124 if (it != fusedLoopsAndRanges.end()) {
125 ivs.push_back(it->second.offset);
126 tileSizes.push_back(it->second.size);
127 loopRanges.push_back(it->second);
128 LLVM_DEBUG(llvm::dbgs() <<
"tiled loop#" << i <<
" with LoopRange "
129 << loopRanges.back() <<
"\n");
133 LLVM_DEBUG(llvm::dbgs() <<
"full loop#" << i <<
" with LoopRange "
134 << loopRanges.back() <<
"\n");
139 clonedShapes.reserve(producer->getNumOperands());
149 resultTypes.reserve(producer->getNumResults());
150 int64_t firstInitOperandIdx =
152 for (int64_t i = 0, e = producer->getNumResults(); i < e; ++i) {
153 resultTypes.push_back(clonedShapes[firstInitOperandIdx + i].
getType());
157 LinalgOp clonedOp =
clone(b, producer, resultTypes, clonedShapes);
161 llvm::map_range(loopRanges, [&](
Range range) {
return range.
offset; }));
170 Value shapedOperand,
unsigned dim) {
172 if (
auto subViewOp = dyn_cast<memref::SubViewOp>(shapeProducingOp))
173 return subViewOp.getOrCreateRanges(b, loc)[dim];
174 if (
auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(shapeProducingOp))
175 return sliceOp.getOrCreateRanges(b, loc)[dim];
176 llvm_unreachable(
"SubviewOp or ExtractSliceOp expected");
184 LLVM_DEBUG(llvm::dbgs() <<
"Producer map: " << producerMap <<
"\n");
186 Value shapedOperand = consumerOpOperand.
get();
188 unsigned posInProducerLoop = cast<AffineDimExpr>(en.value()).getPosition();
190 b, consumerOpOperand.
getOwner()->
getLoc(), shapedOperand, en.index());
192 return fuse(b, producerOp, fusedLoopsAndRanges);
202 if (!isa<RankedTensorType>(tensor.
getType()))
206 LLVM_DEBUG(llvm::dbgs() <<
"\ngetProducerOfTensor: " << tensor);
208 opResult = cast<OpResult>(tensor);
211 if (
auto sliceOp = tensor.
getDefiningOp<tensor::ExtractSliceOp>()) {
212 tensor = sliceOp.getSource();
215 if (
auto blockArg = dyn_cast<BlockArgument>(tensor)) {
216 if (
auto forOp = blockArg.getDefiningOp<scf::ForOp>()) {
217 tensor = forOp.getInitArgs()[blockArg.getArgNumber()];
225 FailureOr<FusionInfo>
227 Value inputTensor = consumerOpOperand.
get();
230 if (!producerOpResult) {
231 LLVM_DEBUG(llvm::dbgs() <<
"\nUnable to find producer");
237 FailureOr<FusionInfo>
240 auto producerOp = dyn_cast<LinalgOp>(producerOpResult.
getOwner());
244 LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.
getOwner());
248 Value inputTensor = consumerOpOperand.
get();
251 auto sliceOp = inputTensor.
getDefiningOp<tensor::ExtractSliceOp>();
253 LLVM_DEBUG(llvm::dbgs()
254 <<
"\nNot fusable, not an extract_slice op: " << inputTensor);
266 LLVM_DEBUG(llvm::dbgs() <<
"Fuse into consumer: " << *consumerOp <<
"\n");
269 LinalgOp fusedProducer =
270 fuse(b, producerOp, producerOp.getMatchingIndexingMap(opOperand),
280 if (consumerType != def.
getType())
281 def = b.
create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
282 consumerOpOperand.
set(def);
static LinalgOp fuse(OpBuilder &b, LinalgOp producer, const DenseMap< unsigned, Range > &fusedLoopsAndRanges)
Fuses the producer by cloning the producer.
static void getProducerOfTensor(Value tensor, OpResult &opResult)
Walk back use-def chain through scf::For yields.
static SmallVector< Value > getTiledOperands(LinalgOp producer)
static ShapeDimension getShapeDefiningLoopRange(LinalgOp op, unsigned loopDepth, bool fromSubViewOpOnly=false)
static Range getRangeFromOperandShape(OpBuilder &b, Location loc, Value shapedOperand, unsigned dim)
Get the loop range for a dimension dim based on the shapedOperand.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
ArrayRef< AffineExpr > getResults() const
IntegerAttr getIndexAttr(int64_t value)
IRValueT get() const
Return the current value being used by this operand.
void set(IRValueT newValue)
Set the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
This class provides a mutable adaptor for a range of operands.
OperandRange getAsOperandRange() const
Explicit conversion to an OperandRange.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents a single result from folding an operation.
This class represents an operand of an operation.
This is a value defined by a result of an operation.
Operation * getOwner() const
Returns the operation that owns this result.
unsigned getResultNumber() const
Returns the number of this result.
unsigned getBeginOperandIndex() const
Return the operand index of the first element of this range.
Operation is the basic unit of execution within MLIR.
Location getLoc()
The source location the operation was defined or derived from.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
Block * getParentBlock()
Return the Block in which this Value is defined.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Operation * getOwner() const
Return the owner of this operand.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
FailureOr< FusionInfo > fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand)
This implements the fusion part of the "tileAndFuse on tensors" transformation and thus requires the ...
SmallVector< Value > makeTiledShapes(OpBuilder &builder, Location loc, LinalgOp linalgOp, ValueRange valuesToTile, ArrayRef< OpFoldResult > ivs, ArrayRef< OpFoldResult > tileSizes, ArrayRef< OpFoldResult > sizeBounds, bool omitPartialTileCheck)
Creates extract_slice/subview ops for all valuesToTile of the given linalgOp with builder,...
OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value val, int64_t dim)
Create one memref::DimOp or tensor::DimOp depending on the type of val.
void offsetIndices(OpBuilder &b, LinalgOp linalgOp, ArrayRef< OpFoldResult > offests)
Add the specified offsets to any linalg.index ops contained in the given linalgOp.
Include the generated interface declarations.
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
Implements a simple high-level fusion pass on linalg structured operations.
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...
A struct containing the Linalg producer before and after fusion.