379 lines · cpp
1//===- Loops.cpp - conversion from Linalg named and generic ops to loops --===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Linalg/Passes.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Dialect/Arith/IR/Arith.h"13#include "mlir/Dialect/Func/IR/FuncOps.h"14#include "mlir/Dialect/Linalg/IR/Linalg.h"15#include "mlir/Dialect/Linalg/Transforms/Transforms.h"16#include "mlir/Dialect/Linalg/Utils/Utils.h"17#include "mlir/Dialect/SCF/Transforms/Transforms.h"18#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"19#include "mlir/IR/AffineExpr.h"20#include "mlir/IR/AffineMap.h"21#include "mlir/IR/IRMapping.h"22#include "mlir/Support/LLVM.h"23#include "mlir/Transforms/DialectConversion.h"24#include "mlir/Transforms/FoldUtils.h"25#include "mlir/Transforms/GreedyPatternRewriteDriver.h"26#include "llvm/ADT/TypeSwitch.h"27 28namespace mlir {29#define GEN_PASS_DEF_CONVERTLINALGTOAFFINELOOPSPASS30#define GEN_PASS_DEF_CONVERTLINALGTOLOOPSPASS31#define GEN_PASS_DEF_CONVERTLINALGTOPARALLELLOOPSPASS32#include "mlir/Dialect/Linalg/Passes.h.inc"33} // namespace mlir34 35using namespace mlir;36using namespace mlir::linalg;37 38static SmallVector<Value> makeCanonicalAffineApplies(OpBuilder &b, Location loc,39 AffineMap map,40 ArrayRef<Value> vals) {41 if (map.isEmpty())42 return {};43 44 assert(map.getNumInputs() == vals.size());45 SmallVector<Value> res;46 res.reserve(map.getNumResults());47 auto dims = map.getNumDims();48 for (auto e : map.getResults()) {49 auto exprMap = AffineMap::get(dims, map.getNumSymbols(), e);50 SmallVector<Value> operands(vals);51 affine::canonicalizeMapAndOperands(&exprMap, &operands);52 res.push_back(affine::AffineApplyOp::create(b, loc, exprMap, operands));53 }54 return res;55}56 57template <typename LoadOpTy, typename StoreOpTy, typename OpType>58static void inlineRegionAndEmitStore(OpBuilder &b, Location loc, OpType op,59 ArrayRef<Value> indexedValues,60 ArrayRef<SmallVector<Value>> indexing,61 ArrayRef<Value> outputBuffers) {62 auto &block = op->getRegion(0).front();63 IRMapping map;64 map.map(block.getArguments(), indexedValues);65 for (auto &op : block.without_terminator()) {66 auto *newOp = b.clone(op, map);67 map.map(op.getResults(), newOp->getResults());68 }69 70 Operation *terminator = block.getTerminator();71 for (OpOperand &operand : terminator->getOpOperands()) {72 Value toStore = map.lookupOrDefault(operand.get());73 StoreOpTy::create(b, loc, toStore,74 outputBuffers[operand.getOperandNumber()],75 indexing[operand.getOperandNumber()]);76 }77}78 79// Returns a pair that contains input indices and output indices of a80// SingleInputPoolingOp `op`.81struct InputAndOutputIndices {82 SmallVector<Value> inputs;83 SmallVector<Value> outputs;84};85template <typename SingleInputPoolingOp>86static InputAndOutputIndices87getInputAndOutputIndices(OpBuilder &b, Location loc, ArrayRef<Value> allIvs,88 SingleInputPoolingOp op) {89 auto mapsRange = op.getIndexingMapsArray();90 auto maps = llvm::to_vector<8>(91 llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));92 return InputAndOutputIndices{93 makeCanonicalAffineApplies(b, loc, maps[0], allIvs),94 makeCanonicalAffineApplies(b, loc, maps[2], allIvs)};95}96 97/// Emits the MLIR for the scalar part of the generic op by:98/// 1. Emitting load ops for each input and output view in order. This is99/// achieved by applying the appropriate input or output map to the100/// enclosing induction variables.101/// 2. Emitting a call to `op.fun()` that takes as arguments the scalars102/// from point 1. above.103/// 3. Emitting store ops to store the results of 2. to the output104/// views.105///106/// An example output may resemble:107///108/// ```109/// scf.for %i = %c0 to %0 step %c1 {110/// scf.for %j = %c0 to %1 step %c1 {111/// scf.for %k = %c0 to %4 step %c1 {112/// %11 = load %arg0[%i, %j] :113/// memref<?x?xf32, stride_specification>114/// %12 = load %arg1[%i, %j, %k] :115/// memref<?x?x?xf32, stride_specification>116/// %13 = load %arg2[%i, %k, %j] :117/// memref<?x?x?xf32, stride_specification>118/// %14:2 = call @foo(%11, %12, %13) : (f32, f32, f32) -> (f32, f32)119/// store %14#0, %arg1[%i, %j, %k] :120/// memref<?x?x?Xf32, stride_specification>121/// store %14#1, %arg2[%i, %k, %j] :122/// memref<?x?x?Xf32, stride_specification>123/// }124/// }125/// }126/// ```127template <typename LoadOpTy, typename StoreOpTy>128static void emitScalarImplementation(OpBuilder &b, Location loc,129 ArrayRef<Value> allIvs,130 LinalgOp linalgOp) {131 assert(linalgOp.hasPureBufferSemantics() &&132 "expected linalg op with buffer semantics");133 SmallVector<Value> indexedValues;134 indexedValues.reserve(linalgOp->getNumOperands());135 136 auto allIvsPlusDims = SmallVector<Value>(allIvs);137 138 // TODO: Avoid the loads if the corresponding argument of the139 // region has no uses.140 // 1.a. Emit load from input operand or for scalars access the operand itself.141 for (OpOperand *inputOperand : linalgOp.getDpsInputOperands()) {142 if (linalgOp.isScalar(inputOperand)) {143 indexedValues.push_back(inputOperand->get());144 continue;145 }146 auto indexing = makeCanonicalAffineApplies(147 b, loc, linalgOp.getMatchingIndexingMap(inputOperand), allIvsPlusDims);148 indexedValues.push_back(149 LoadOpTy::create(b, loc, inputOperand->get(), indexing));150 }151 // 1.b. Emit load from output views.152 for (OpOperand &outputOperand : linalgOp.getDpsInitsMutable()) {153 SmallVector<Value> indexing = makeCanonicalAffineApplies(154 b, loc, linalgOp.getMatchingIndexingMap(&outputOperand),155 allIvsPlusDims);156 indexedValues.push_back(157 LoadOpTy::create(b, loc, outputOperand.get(), indexing));158 }159 160 // TODO: When a region inliner exists, use it.161 // 2. Inline region, currently only works for a single basic block.162 // 3. Emit store.163 SmallVector<SmallVector<Value>, 8> indexing;164 SmallVector<Value> outputBuffers;165 for (OpOperand &outputOperand : linalgOp.getDpsInitsMutable()) {166 if (!isa<MemRefType>(outputOperand.get().getType()))167 continue;168 indexing.push_back(makeCanonicalAffineApplies(169 b, loc, linalgOp.getMatchingIndexingMap(&outputOperand),170 allIvsPlusDims));171 outputBuffers.push_back(outputOperand.get());172 }173 inlineRegionAndEmitStore<LoadOpTy, StoreOpTy>(b, loc, linalgOp, indexedValues,174 indexing, outputBuffers);175}176 177/// Replace the index operations in the body of the loop nest by the matching178/// induction variables.179static void replaceIndexOpsByInductionVariables(RewriterBase &rewriter,180 LinalgOp linalgOp,181 ArrayRef<Operation *> loopOps) {182 // Extract the induction variables of the loop nest from outer to inner.183 SmallVector<Value> allIvs;184 for (Operation *loopOp : loopOps) {185 llvm::TypeSwitch<Operation *>(loopOp)186 .Case([&](scf::ParallelOp parallelOp) {187 allIvs.append(parallelOp.getInductionVars());188 })189 .Case([&](scf::ForOp forOp) {190 allIvs.push_back(forOp.getInductionVar());191 })192 .Case([&](affine::AffineForOp affineForOp) {193 allIvs.push_back(affineForOp.getInductionVar());194 })195 .DefaultUnreachable("unexpected op");196 }197 assert(linalgOp.getNumLoops() == allIvs.size() &&198 "expected the number of loops and induction variables to match");199 // Replace the index operations in the body of the innermost loop op.200 if (!loopOps.empty()) {201 auto loopOp = cast<LoopLikeOpInterface>(loopOps.back());202 for (Region *r : loopOp.getLoopRegions())203 for (IndexOp indexOp : llvm::make_early_inc_range(r->getOps<IndexOp>()))204 rewriter.replaceOp(indexOp, allIvs[indexOp.getDim()]);205 }206}207 208template <typename LoopTy>209static FailureOr<LinalgLoops> linalgOpToLoopsImpl(RewriterBase &rewriter,210 LinalgOp linalgOp) {211 using LoadOpTy =212 std::conditional_t<std::is_same<LoopTy, affine::AffineForOp>::value,213 affine::AffineLoadOp, memref::LoadOp>;214 using StoreOpTy =215 std::conditional_t<std::is_same<LoopTy, affine::AffineForOp>::value,216 affine::AffineStoreOp, memref::StoreOp>;217 218 // The flattened loopToOperandRangesMaps is expected to be an invertible219 // permutation map (which is asserted in the inverse calculation).220 assert(linalgOp.hasPureBufferSemantics() &&221 "expected linalg op with buffer semantics");222 223 auto loopRanges = linalgOp.createLoopRanges(rewriter, linalgOp.getLoc());224 auto iteratorTypes = linalgOp.getIteratorTypesArray();225 226 SmallVector<Value> allIvs;227 GenerateLoopNest<LoopTy>::doit(228 rewriter, linalgOp.getLoc(), loopRanges, linalgOp, iteratorTypes,229 [&](OpBuilder &b, Location loc, ValueRange ivs,230 ValueRange operandValuesToUse) -> scf::ValueVector {231 assert(operandValuesToUse == linalgOp->getOperands() &&232 "expect operands are captured and not passed by loop argument");233 allIvs.append(ivs.begin(), ivs.end());234 emitScalarImplementation<LoadOpTy, StoreOpTy>(b, loc, allIvs, linalgOp);235 return scf::ValueVector{};236 });237 // Number of loop ops might be different from the number of ivs since some238 // loops like affine.parallel and scf.parallel have multiple ivs.239 SetVector<Operation *> loopSet;240 for (Value iv : allIvs) {241 if (!iv)242 return failure();243 // The induction variable is a block argument of the entry block of the244 // loop operation.245 BlockArgument ivVal = dyn_cast<BlockArgument>(iv);246 if (!ivVal)247 return failure();248 loopSet.insert(ivVal.getOwner()->getParentOp());249 }250 LinalgLoops loops(loopSet.begin(), loopSet.end());251 // Replace all index operations in the loop body.252 replaceIndexOpsByInductionVariables(rewriter, linalgOp, loops);253 return loops;254}255 256namespace {257template <typename LoopType>258class LinalgRewritePattern : public RewritePattern {259public:260 LinalgRewritePattern(MLIRContext *context)261 : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}262 263 LogicalResult matchAndRewrite(Operation *op,264 PatternRewriter &rewriter) const override {265 auto linalgOp = dyn_cast<LinalgOp>(op);266 if (!isa<LinalgOp>(op) || !linalgOp.hasPureBufferSemantics()) {267 return rewriter.notifyMatchFailure(268 op, "expected linalg op with buffer semantics");269 }270 if (failed(linalgOpToLoopsImpl<LoopType>(rewriter, linalgOp)))271 return failure();272 rewriter.eraseOp(op);273 return success();274 }275};276 277/// Local folding pattern for AffineApplyOp that we can apply greedily.278/// This replaces AffineApplyOp by the proper value in cases where the279/// associated map is trivial.280/// A trivial map here is defined as a map with a single result and either:281/// 1. Zero operand + returns a single AffineConstantExpr282/// 2. One operand + returns a single AffineDimExpr283/// 3. One operand + returns a single AffineSymbolExpr284//285/// In the first case, the AffineApplyOp is replaced by a new constant. In the286/// other cases, it is replaced by its unique operand.287struct FoldAffineOp : public RewritePattern {288 FoldAffineOp(MLIRContext *context)289 : RewritePattern(affine::AffineApplyOp::getOperationName(), 0, context) {}290 291 LogicalResult matchAndRewrite(Operation *op,292 PatternRewriter &rewriter) const override {293 auto affineApplyOp = cast<affine::AffineApplyOp>(op);294 auto map = affineApplyOp.getAffineMap();295 if (map.getNumResults() != 1 || map.getNumInputs() > 1)296 return failure();297 298 AffineExpr expr = map.getResult(0);299 if (map.getNumInputs() == 0) {300 if (auto val = dyn_cast<AffineConstantExpr>(expr)) {301 rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op, val.getValue());302 return success();303 }304 return failure();305 }306 if (isa<AffineDimExpr, AffineSymbolExpr>(expr)) {307 rewriter.replaceOp(op, op->getOperand(0));308 return success();309 }310 return failure();311 }312};313 314template <typename LoopType>315static void lowerLinalgToLoopsImpl(Operation *enclosingOp) {316 MLIRContext *context = enclosingOp->getContext();317 RewritePatternSet patterns(context);318 patterns.add<LinalgRewritePattern<LoopType>>(context);319 memref::DimOp::getCanonicalizationPatterns(patterns, context);320 tensor::DimOp::getCanonicalizationPatterns(patterns, context);321 affine::AffineApplyOp::getCanonicalizationPatterns(patterns, context);322 patterns.add<FoldAffineOp>(context);323 // Just apply the patterns greedily.324 (void)applyPatternsGreedily(enclosingOp, std::move(patterns));325}326 327struct LowerToAffineLoops328 : public impl::ConvertLinalgToAffineLoopsPassBase<LowerToAffineLoops> {329 using impl::ConvertLinalgToAffineLoopsPassBase<330 LowerToAffineLoops>::ConvertLinalgToAffineLoopsPassBase;331 void getDependentDialects(DialectRegistry ®istry) const override {332 registry.insert<memref::MemRefDialect>();333 }334 void runOnOperation() override {335 lowerLinalgToLoopsImpl<affine::AffineForOp>(getOperation());336 }337};338 339struct LowerToLoops : public impl::ConvertLinalgToLoopsPassBase<LowerToLoops> {340 using impl::ConvertLinalgToLoopsPassBase<341 LowerToLoops>::ConvertLinalgToLoopsPassBase;342 void getDependentDialects(DialectRegistry ®istry) const override {343 registry.insert<memref::MemRefDialect, scf::SCFDialect>();344 }345 void runOnOperation() override {346 lowerLinalgToLoopsImpl<scf::ForOp>(getOperation());347 }348};349 350struct LowerToParallelLoops351 : public impl::ConvertLinalgToParallelLoopsPassBase<LowerToParallelLoops> {352 using impl::ConvertLinalgToParallelLoopsPassBase<353 LowerToParallelLoops>::ConvertLinalgToParallelLoopsPassBase;354 void runOnOperation() override {355 lowerLinalgToLoopsImpl<scf::ParallelOp>(getOperation());356 }357};358 359} // namespace360 361/// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.362FailureOr<LinalgLoops>363mlir::linalg::linalgOpToAffineLoops(RewriterBase &rewriter, LinalgOp linalgOp) {364 return linalgOpToLoopsImpl<affine::AffineForOp>(rewriter, linalgOp);365}366 367/// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.368FailureOr<LinalgLoops> mlir::linalg::linalgOpToLoops(RewriterBase &rewriter,369 LinalgOp linalgOp) {370 return linalgOpToLoopsImpl<scf::ForOp>(rewriter, linalgOp);371}372 373/// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.374FailureOr<LinalgLoops>375mlir::linalg::linalgOpToParallelLoops(RewriterBase &rewriter,376 LinalgOp linalgOp) {377 return linalgOpToLoopsImpl<scf::ParallelOp>(rewriter, linalgOp);378}379