25 #include "llvm/ADT/TypeSwitch.h" 44 res.push_back(b.
create<AffineApplyOp>(loc, exprMap, operands));
49 template <
typename LoadOpTy,
typename StoreOpTy,
typename OpType>
54 auto &block = op->getRegion(0).front();
56 map.
map(block.getArguments(), indexedValues);
57 for (
auto &op : block.without_terminator()) {
58 auto *newOp = b.
clone(op, map);
59 map.
map(op.getResults(), newOp->getResults());
62 Operation *terminator = block.getTerminator();
65 b.
create<StoreOpTy>(loc, toStore, outputBuffers[operand.getOperandNumber()],
66 indexing[operand.getOperandNumber()]);
76 template <
typename SingleInputPoolingOp>
79 SingleInputPoolingOp op) {
80 auto mapsRange = op.getIndexingMapsArray();
81 auto maps = llvm::to_vector<8>(
82 llvm::map_range(mapsRange, [](AffineMapAttr a) {
return a.getValue(); }));
118 template <
typename LoadOpTy,
typename StoreOpTy>
122 assert(linalgOp.hasBufferSemantics() &&
123 "expected linalg op with buffer semantics");
125 indexedValues.reserve(linalgOp.getNumInputsAndOutputs());
132 for (
OpOperand *inputOperand : linalgOp.getInputOperands()) {
133 if (linalgOp.isScalar(inputOperand)) {
134 indexedValues.push_back(inputOperand->get());
138 b, loc, linalgOp.getTiedIndexingMap(inputOperand), allIvsPlusDims);
139 indexedValues.push_back(
140 b.
create<LoadOpTy>(loc, inputOperand->get(), indexing));
143 for (
OpOperand *outputOperand : linalgOp.getOutputOperands()) {
145 b, loc, linalgOp.getTiedIndexingMap(outputOperand), allIvsPlusDims);
146 indexedValues.push_back(
147 b.
create<LoadOpTy>(loc, outputOperand->get(), indexing));
155 for (
OpOperand *outputOperand : linalgOp.getOutputBufferOperands()) {
157 b, loc, linalgOp.getTiedIndexingMap(outputOperand), allIvsPlusDims));
158 outputBuffers.push_back(outputOperand->get());
160 inlineRegionAndEmitStore<LoadOpTy, StoreOpTy>(b, loc, linalgOp, indexedValues,
161 indexing, outputBuffers);
173 .Case([&](scf::ParallelOp parallelOp) {
174 allIvs.append(parallelOp.getInductionVars().begin(),
175 parallelOp.getInductionVars().end());
177 .Case([&](scf::ForOp forOp) {
178 allIvs.push_back(forOp.getInductionVar());
180 .Case([&](AffineForOp affineForOp) {
181 allIvs.push_back(affineForOp.getInductionVar());
183 .Default([&](
Operation *op) { assert(
false &&
"unexpected op"); });
185 assert(linalgOp.getNumLoops() == allIvs.size() &&
186 "expected the number of loops and induction variables to match");
188 if (!loopOps.empty()) {
189 LoopLikeOpInterface loopOp = loopOps.back();
190 for (IndexOp indexOp :
191 llvm::make_early_inc_range(loopOp.getLoopBody().getOps<IndexOp>()))
192 rewriter.
replaceOp(indexOp, allIvs[indexOp.getDim()]);
196 template <
typename LoopTy>
201 AffineLoadOp, memref::LoadOp>::type;
204 AffineStoreOp, memref::StoreOp>::type;
208 assert(linalgOp.hasBufferSemantics() &&
209 "expected linalg op with buffer semantics");
211 auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc());
212 auto iteratorTypes = llvm::to_vector<4>(linalgOp.iterator_types().getValue());
216 rewriter, linalgOp.getLoc(), loopRanges, linalgOp, iteratorTypes,
219 assert(operandValuesToUse == linalgOp->getOperands() &&
220 "expect operands are captured and not passed by loop argument");
221 allIvs.append(ivs.begin(), ivs.end());
222 emitScalarImplementation<LoadOpTy, StoreOpTy>(b, loc, allIvs, linalgOp);
228 for (
Value iv : allIvs) {
245 template <
typename LoopType>
253 auto linalgOp = dyn_cast<LinalgOp>(op);
254 if (!isa<LinalgOp>(op))
256 if (
failed(linalgOpToLoopsImpl<LoopType>(rewriter, linalgOp)))
275 :
RewritePattern(AffineApplyOp::getOperationName(), 0, context) {}
279 AffineApplyOp affineApplyOp = cast<AffineApplyOp>(op);
280 auto map = affineApplyOp.getAffineMap();
281 if (map.getNumResults() != 1 || map.getNumInputs() > 1)
285 if (map.getNumInputs() == 0) {
300 template <
typename LoopType>
301 static void lowerLinalgToLoopsImpl(func::FuncOp funcOp) {
304 patterns.
add<LinalgRewritePattern<LoopType>>(context);
305 memref::DimOp::getCanonicalizationPatterns(patterns, context);
306 tensor::DimOp::getCanonicalizationPatterns(patterns, context);
307 AffineApplyOp::getCanonicalizationPatterns(patterns, context);
308 patterns.
add<FoldAffineOp>(context);
313 struct LowerToAffineLoops
314 :
public LinalgLowerToAffineLoopsBase<LowerToAffineLoops> {
316 registry.
insert<memref::MemRefDialect>();
318 void runOnOperation()
override {
319 lowerLinalgToLoopsImpl<AffineForOp>(getOperation());
323 struct LowerToLoops :
public LinalgLowerToLoopsBase<LowerToLoops> {
325 registry.
insert<memref::MemRefDialect, scf::SCFDialect>();
327 void runOnOperation()
override {
328 lowerLinalgToLoopsImpl<scf::ForOp>(getOperation());
332 struct LowerToParallelLoops
333 :
public LinalgLowerToParallelLoopsBase<LowerToParallelLoops> {
334 void runOnOperation()
override {
335 lowerLinalgToLoopsImpl<scf::ParallelOp>(getOperation());
341 std::unique_ptr<OperationPass<func::FuncOp>>
343 return std::make_unique<LowerToLoops>();
346 std::unique_ptr<OperationPass<func::FuncOp>>
348 return std::make_unique<LowerToParallelLoops>();
351 std::unique_ptr<OperationPass<func::FuncOp>>
353 return std::make_unique<LowerToAffineLoops>();
360 return linalgOpToLoopsImpl<AffineForOp>(rewriter, linalgOp);
366 return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp);
373 return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp);
Include the generated interface declarations.
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments...
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Operation is a basic unit of execution within MLIR.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
unsigned getNumSymbols() const
static void emitScalarImplementation(OpBuilder &b, Location loc, ArrayRef< Value > allIvs, LinalgOp linalgOp)
Emits the MLIR for the scalar part of the generic op by:
unsigned getNumDims() const
bool isEmpty() const
Returns true if this affine map is an empty map, i.e., () -> ().
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
Operation * clone(Operation &op, BlockAndValueMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Value getOperand(unsigned idx)
static void inlineRegionAndEmitStore(OpBuilder &b, Location loc, OpType op, ArrayRef< Value > indexedValues, ArrayRef< SmallVector< Value >> indexing, ArrayRef< Value > outputBuffers)
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
static FailureOr< LinalgLoops > linalgOpToLoopsImpl(PatternRewriter &rewriter, LinalgOp linalgOp)
std::vector< Value > ValueVector
An owning vector of values, handy to return from functions.
std::unique_ptr< OperationPass< func::FuncOp > > createConvertLinalgToLoopsPass()
Create a pass to convert Linalg operations to scf.for loops and memref.load/memref.store accesses.
FailureOr< LinalgLoops > linalgOpToAffineLoops(PatternRewriter &rewriter, LinalgOp linalgOp)
Emit a loop nest of affine.for with the proper body for linalgOp.
std::unique_ptr< OperationPass< func::FuncOp > > createConvertLinalgToParallelLoopsPass()
Create a pass to convert Linalg operations to scf.parallel loops and memref.load/memref.store accesses.
An integer constant appearing in affine expression.
static constexpr const bool value
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
void map(Block *from, Block *to)
Inserts a new mapping for 'from' to 'to'.
unsigned getNumInputs() const
Block * getOwner() const
Returns the block that owns this argument.
RewritePattern is the common base class for all DAG to DAG replacements.
MutableArrayRef< OpOperand > getOpOperands()
std::unique_ptr< OperationPass< func::FuncOp > > createConvertLinalgToAffineLoopsPass()
Create a pass to convert Linalg operations to affine.for loops and affine_load/affine_store accesses...
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents an efficient way to signal success or failure.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
This class provides support for representing a failure result, or a valid value of type T...
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
static void replaceIndexOpsByInductionVariables(LinalgOp linalgOp, PatternRewriter &rewriter, ArrayRef< Operation *> loopOps)
Replace the index operations in the body of the loop nest by the matching induction variables...
Base type for affine expression.
void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl< Value > *operands)
Modifies both map and operands in-place so as to:
unsigned getNumResults() const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued...
This class represents an argument of a Block.
ArrayRef< AffineExpr > getResults() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
FailureOr< LinalgLoops > linalgOpToLoops(PatternRewriter &rewriter, LinalgOp linalgOp)
Emit a loop nest of scf.for with the proper body for linalgOp.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
static void doit(OpBuilder &b, Location loc, ArrayRef< Range > loopRanges, LinalgOp linalgOp, ArrayRef< Attribute > iteratorTypes, function_ref< scf::ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilderFn, Optional< LinalgLoopDistributionOptions >=None, ArrayRef< StringRef > distributionTypes={})
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
A dimensional identifier appearing in an affine expression.
Specialization of arith.constant op that returns an integer of index type.
MLIRContext is the top-level object for a collection of MLIR operations.
Block * lookupOrDefault(Block *from) const
Lookup a mapped value within the map.
This class represents an operand of an operation.
static SmallVector< Value > makeCanonicalAffineApplies(OpBuilder &b, Location loc, AffineMap map, ArrayRef< Value > vals)
FailureOr< LinalgLoops > linalgOpToParallelLoops(PatternRewriter &rewriter, LinalgOp linalgOp)
Emit a loop nest of scf.parallel with the proper body for linalgOp.
LogicalResult applyPatternsAndFoldGreedily(MutableArrayRef< Region > regions, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig())
Rewrite the regions of the specified operation, which must be isolated from above, by repeatedly applying the highest benefit patterns in a greedy work-list driven manner.
This class helps build Operations.
This class provides an abstraction over the different types of ranges over Values.
static InputAndOutputIndices getInputAndOutputIndices(OpBuilder &b, Location loc, ArrayRef< Value > allIvs, SingleInputPoolingOp op)
A symbolic identifier appearing in an affine expression.