23 #include "llvm/ADT/SmallBitVector.h"
24 #include "llvm/Support/Debug.h"
26 #define DEBUG_TYPE "linalg-fusion"
59 bool fromSubViewOpOnly =
false) {
62 for (
OpOperand &opOperand : op->getOpOperands()) {
69 if (fromSubViewOpOnly &&
70 !isa_and_nonnull<memref::SubViewOp, tensor::ExtractSliceOp>(
71 opOperand.get().getDefiningOp()))
74 AffineMap map = op.getMatchingIndexingMap(&opOperand);
75 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange I/O idx: "
76 << opOperand.getOperandNumber() <<
"\n");
77 LLVM_DEBUG(llvm::dbgs()
78 <<
"getShapeDefiningLoopRange map: " << map <<
"\n");
80 auto dimExpr = dyn_cast<AffineDimExpr>(en.value());
83 if (loopDepth == cast<AffineDimExpr>(en.value()).getPosition()) {
84 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange loopDepth: "
85 << loopDepth <<
"\n");
86 LLVM_DEBUG(llvm::dbgs() <<
"getShapeDefiningLoopRange shape: "
87 << opOperand.get() <<
"\n");
89 static_cast<unsigned>(en.index())};
93 llvm_unreachable(
"Expect to be able to extract a shape defining loop range");
97 return producer->getOperands();
109 for (
unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
113 sizeBounds.push_back(dim);
114 auto it = fusedLoopsAndRanges.find(i);
115 if (it != fusedLoopsAndRanges.end()) {
116 ivs.push_back(it->second.offset);
117 tileSizes.push_back(it->second.size);
118 loopRanges.push_back(it->second);
119 LLVM_DEBUG(llvm::dbgs() <<
"tiled loop#" << i <<
" with LoopRange "
120 << loopRanges.back() <<
"\n");
124 LLVM_DEBUG(llvm::dbgs() <<
"full loop#" << i <<
" with LoopRange "
125 << loopRanges.back() <<
"\n");
130 clonedShapes.reserve(producer->getNumOperands());
140 resultTypes.reserve(producer->getNumResults());
141 int64_t firstInitOperandIdx =
143 for (int64_t i = 0, e = producer->getNumResults(); i < e; ++i) {
144 resultTypes.push_back(clonedShapes[firstInitOperandIdx + i].
getType());
148 LinalgOp clonedOp =
clone(b, producer, resultTypes, clonedShapes);
152 llvm::map_range(loopRanges, [&](
Range range) {
return range.
offset; }));
161 Value shapedOperand,
unsigned dim) {
163 if (
auto subViewOp = dyn_cast<memref::SubViewOp>(shapeProducingOp))
164 return subViewOp.getOrCreateRanges(b, loc)[dim];
165 if (
auto sliceOp = dyn_cast<tensor::ExtractSliceOp>(shapeProducingOp))
166 return sliceOp.getOrCreateRanges(b, loc)[dim];
167 llvm_unreachable(
"SubviewOp or ExtractSliceOp expected");
175 LLVM_DEBUG(llvm::dbgs() <<
"Producer map: " << producerMap <<
"\n");
177 Value shapedOperand = consumerOpOperand.
get();
179 unsigned posInProducerLoop = cast<AffineDimExpr>(en.value()).getPosition();
181 b, consumerOpOperand.
getOwner()->
getLoc(), shapedOperand, en.index());
183 return fuse(b, producerOp, fusedLoopsAndRanges);
193 if (!isa<RankedTensorType>(tensor.
getType()))
197 LLVM_DEBUG(llvm::dbgs() <<
"\ngetProducerOfTensor: " << tensor);
199 opResult = cast<OpResult>(tensor);
202 if (
auto sliceOp = tensor.
getDefiningOp<tensor::ExtractSliceOp>()) {
203 tensor = sliceOp.getSource();
206 if (
auto blockArg = dyn_cast<BlockArgument>(tensor)) {
207 if (
auto forOp = blockArg.getDefiningOp<scf::ForOp>()) {
208 tensor = forOp.getInitArgs()[blockArg.getArgNumber()];
216 FailureOr<FusionInfo>
218 Value inputTensor = consumerOpOperand.
get();
221 if (!producerOpResult) {
222 LLVM_DEBUG(llvm::dbgs() <<
"\nUnable to find producer");
228 FailureOr<FusionInfo>
231 auto producerOp = dyn_cast<LinalgOp>(producerOpResult.
getOwner());
235 LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.
getOwner());
239 Value inputTensor = consumerOpOperand.
get();
242 auto sliceOp = inputTensor.
getDefiningOp<tensor::ExtractSliceOp>();
244 LLVM_DEBUG(llvm::dbgs()
245 <<
"\nNot fusable, not an extract_slice op: " << inputTensor);
257 LLVM_DEBUG(llvm::dbgs() <<
"Fuse into consumer: " << *consumerOp <<
"\n");
260 LinalgOp fusedProducer =
261 fuse(b, producerOp, producerOp.getMatchingIndexingMap(opOperand),
269 if (cast<ShapedType>(consumerType).getRank() !=
270 cast<ShapedType>(def.
getType()).getRank()) {
271 llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
279 if (consumerType != def.
getType())
280 def = tensor::CastOp::create(b, fusedProducer.getLoc(), consumerType, def);
281 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.
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.
CollapseShapeOp dropGivenUnitDims(OpBuilder &b, Location loc, Value src, const llvm::SmallBitVector &dropDims)
Create tensor.collapse_shape to drop unit dimensions in dropDims in tensor src.
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.