170 lines · cpp
1//===- DecomposeAffineOps.cpp - Decompose affine ops into finer-grained ---===//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// This file implements functionality to progressively decompose coarse-grained10// affine ops into finer-grained ops.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Affine/Transforms/Transforms.h"16#include "mlir/IR/PatternMatch.h"17#include "llvm/Support/Debug.h"18#include "llvm/Support/DebugLog.h"19#include "llvm/Support/InterleavedRange.h"20 21using namespace mlir;22using namespace mlir::affine;23 24#define DEBUG_TYPE "decompose-affine-ops"25 26/// Count the number of loops surrounding `operand` such that operand could be27/// hoisted above.28/// Stop counting at the first loop over which the operand cannot be hoisted.29static int64_t numEnclosingInvariantLoops(OpOperand &operand) {30 int64_t count = 0;31 Operation *currentOp = operand.getOwner();32 while (auto loopOp = currentOp->getParentOfType<LoopLikeOpInterface>()) {33 if (!loopOp.isDefinedOutsideOfLoop(operand.get()))34 break;35 currentOp = loopOp;36 count++;37 }38 return count;39}40 41void mlir::affine::reorderOperandsByHoistability(RewriterBase &rewriter,42 AffineApplyOp op) {43 SmallVector<int64_t> numInvariant = llvm::to_vector(44 llvm::map_range(op->getOpOperands(), [&](OpOperand &operand) {45 return numEnclosingInvariantLoops(operand);46 }));47 48 int64_t numOperands = op.getNumOperands();49 SmallVector<int64_t> operandPositions =50 llvm::to_vector(llvm::seq<int64_t>(0, numOperands));51 llvm::stable_sort(operandPositions, [&numInvariant](size_t i1, size_t i2) {52 return numInvariant[i1] > numInvariant[i2];53 });54 55 SmallVector<AffineExpr> replacements(numOperands);56 SmallVector<Value> operands(numOperands);57 for (int64_t i = 0; i < numOperands; ++i) {58 operands[i] = op.getOperand(operandPositions[i]);59 replacements[operandPositions[i]] = getAffineSymbolExpr(i, op.getContext());60 }61 62 AffineMap map = op.getAffineMap();63 ArrayRef<AffineExpr> repls{replacements};64 map = map.replaceDimsAndSymbols(repls.take_front(map.getNumDims()),65 repls.drop_front(map.getNumDims()),66 /*numResultDims=*/0,67 /*numResultSyms=*/numOperands);68 map = AffineMap::get(0, numOperands,69 simplifyAffineExpr(map.getResult(0), 0, numOperands),70 op->getContext());71 canonicalizeMapAndOperands(&map, &operands);72 73 rewriter.startOpModification(op);74 op.setMap(map);75 op->setOperands(operands);76 rewriter.finalizeOpModification(op);77}78 79/// Build an affine.apply that is a subexpression `expr` of `originalOp`s affine80/// map and with the same operands.81/// Canonicalize the map and operands to deduplicate and drop dead operands82/// before returning but do not perform maximal composition of AffineApplyOp83/// which would defeat the purpose.84static AffineApplyOp createSubApply(RewriterBase &rewriter,85 AffineApplyOp originalOp, AffineExpr expr) {86 MLIRContext *ctx = originalOp->getContext();87 AffineMap m = originalOp.getAffineMap();88 auto rhsMap = AffineMap::get(m.getNumDims(), m.getNumSymbols(), expr, ctx);89 SmallVector<Value> rhsOperands = originalOp->getOperands();90 canonicalizeMapAndOperands(&rhsMap, &rhsOperands);91 return AffineApplyOp::create(rewriter, originalOp.getLoc(), rhsMap,92 rhsOperands);93}94 95FailureOr<AffineApplyOp> mlir::affine::decompose(RewriterBase &rewriter,96 AffineApplyOp op) {97 // 1. Preconditions: only handle dimensionless AffineApplyOp maps with a98 // top-level binary expression that we can reassociate (i.e. add or mul).99 AffineMap m = op.getAffineMap();100 if (m.getNumDims() > 0)101 return rewriter.notifyMatchFailure(op, "expected no dims");102 103 AffineExpr remainingExp = m.getResult(0);104 auto binExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);105 if (!binExpr)106 return rewriter.notifyMatchFailure(op, "terminal affine.apply");107 108 if (!isa<AffineBinaryOpExpr>(binExpr.getLHS()) &&109 !isa<AffineBinaryOpExpr>(binExpr.getRHS()))110 return rewriter.notifyMatchFailure(op, "terminal affine.apply");111 112 bool supportedKind = ((binExpr.getKind() == AffineExprKind::Add) ||113 (binExpr.getKind() == AffineExprKind::Mul));114 if (!supportedKind)115 return rewriter.notifyMatchFailure(116 op, "only add or mul binary expr can be reassociated");117 118 LDBG() << "Start decomposeIntoFinerGrainedOps: " << op;119 120 // 2. Iteratively extract the RHS subexpressions while the top-level binary121 // expr kind remains the same.122 MLIRContext *ctx = op->getContext();123 SmallVector<AffineExpr> subExpressions;124 while (true) {125 auto currentBinExpr = dyn_cast<AffineBinaryOpExpr>(remainingExp);126 if (!currentBinExpr || currentBinExpr.getKind() != binExpr.getKind()) {127 subExpressions.push_back(remainingExp);128 LDBG() << "--terminal: " << subExpressions.back();129 break;130 }131 subExpressions.push_back(currentBinExpr.getRHS());132 LDBG() << "--subExpr: " << subExpressions.back();133 remainingExp = currentBinExpr.getLHS();134 }135 136 // 3. Reorder subExpressions by the min symbol they are a function of.137 // This also takes care of properly reordering local variables.138 // This however won't be able to split expression that cannot be reassociated139 // such as ones that involve divs and multiple symbols.140 auto getMaxSymbol = [&](AffineExpr e) -> int64_t {141 for (int64_t i = m.getNumSymbols(); i >= 0; --i)142 if (e.isFunctionOfSymbol(i))143 return i;144 return -1;145 };146 llvm::stable_sort(subExpressions, [&](AffineExpr e1, AffineExpr e2) {147 return getMaxSymbol(e1) < getMaxSymbol(e2);148 });149 LDBG() << "--sorted subexprs: " << llvm::interleaved(subExpressions);150 151 // 4. Merge sorted subExpressions iteratively, thus achieving reassociation.152 auto s0 = getAffineSymbolExpr(0, ctx);153 auto s1 = getAffineSymbolExpr(1, ctx);154 AffineMap binMap = AffineMap::get(155 /*dimCount=*/0, /*symbolCount=*/2,156 getAffineBinaryOpExpr(binExpr.getKind(), s0, s1), ctx);157 158 auto current = createSubApply(rewriter, op, subExpressions[0]);159 for (int64_t i = 1, e = subExpressions.size(); i < e; ++i) {160 Value tmp = createSubApply(rewriter, op, subExpressions[i]);161 current = AffineApplyOp::create(rewriter, op.getLoc(), binMap,162 ValueRange{current, tmp});163 LDBG() << "--reassociate into: " << current;164 }165 166 // 5. Replace original op.167 rewriter.replaceOp(op, current.getResult());168 return current;169}170