3fffffa882
This class represents a rewrite pattern list that has been frozen, and thus immutable. This replaces the uses of OwningRewritePatternList in pattern driver related API, such as dialect conversion. When PDL becomes more prevalent, this API will allow for optimizing a set of patterns once without the need to do this per run of a pass. Differential Revision: https://reviews.llvm.org/D89104
463 lines
19 KiB
C++
463 lines
19 KiB
C++
//===- LinalgTransforms.cpp - Linalg transformations as patterns ----------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements logic and helpers to expose Linalg transforms as rewrite
|
|
// patterns.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
|
|
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
|
|
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
|
|
#include "mlir/Dialect/Linalg/Utils/Utils.h"
|
|
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
|
|
#include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Vector/VectorOps.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
#include <type_traits>
|
|
|
|
#define DEBUG_TYPE "linalg-transforms"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::edsc;
|
|
using namespace mlir::edsc::intrinsics;
|
|
using namespace mlir::linalg;
|
|
|
|
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Transformations exposed as rewrite patterns.
|
|
//===----------------------------------------------------------------------===//
|
|
// Marker used as attribute name in generated Linalg rewriting transformations.
|
|
const StringLiteral mlir::linalg::LinalgTransforms::kLinalgTransformMarker =
|
|
"__internal_linalg_transform__";
|
|
|
|
mlir::linalg::LinalgMarker::LinalgMarker(ArrayRef<Identifier> matchDisjunction,
|
|
Optional<Identifier> replacement)
|
|
: matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
|
|
replacement(replacement) {}
|
|
|
|
LogicalResult
|
|
mlir::linalg::LinalgMarker::checkAndNotify(PatternRewriter &rewriter,
|
|
Operation *op) const {
|
|
auto attr = op->template getAttrOfType<StringAttr>(
|
|
LinalgTransforms::kLinalgTransformMarker);
|
|
|
|
if (!attr) {
|
|
// 1. Has no marker case and matchDisjunction is empty.
|
|
if (matchDisjunction.empty())
|
|
return success();
|
|
|
|
// 2. Has no marker but was expecting a marker.
|
|
return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
|
|
diag << " does not have any marker from list: ";
|
|
interleaveComma(matchDisjunction, diag);
|
|
});
|
|
}
|
|
|
|
// 4. Match explicit marker.
|
|
for (auto marker : matchDisjunction)
|
|
if (attr.getValue() == marker)
|
|
return success();
|
|
|
|
// 5. Fail to match.
|
|
return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
|
|
diag << " does not have any marker from list: ";
|
|
interleaveComma(matchDisjunction, diag);
|
|
});
|
|
}
|
|
|
|
void mlir::linalg::LinalgMarker::replaceLinalgMarker(PatternRewriter &rewriter,
|
|
Operation *op) const {
|
|
if (replacement.hasValue())
|
|
op->setAttr(LinalgTransforms::kLinalgTransformMarker,
|
|
rewriter.getStringAttr(replacement.getValue()));
|
|
else
|
|
op->removeAttr(Identifier::get(LinalgTransforms::kLinalgTransformMarker,
|
|
rewriter.getContext()));
|
|
}
|
|
|
|
LinalgTilingOptions &
|
|
mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {
|
|
SmallVector<int64_t, 4> tileSizes(ts.begin(), ts.end());
|
|
tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
|
|
OpBuilder::InsertionGuard guard(b);
|
|
b.setInsertionPointToStart(
|
|
&op->getParentOfType<FuncOp>().getBody().front());
|
|
return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
|
|
Value v = b.create<ConstantIndexOp>(op->getLoc(), s);
|
|
return v;
|
|
}));
|
|
};
|
|
return *this;
|
|
}
|
|
|
|
/// Linalg base tiling pattern.
|
|
mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
|
|
StringRef opName, MLIRContext *context, LinalgTilingOptions options,
|
|
LinalgMarker marker, PatternBenefit benefit)
|
|
: RewritePattern(opName, {}, benefit, context), marker(marker),
|
|
options(options) {}
|
|
|
|
LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
|
|
Operation *op, PatternRewriter &rewriter,
|
|
SmallVectorImpl<Value> &tensorResults) const {
|
|
LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
|
|
if (!linalgOp)
|
|
return failure();
|
|
if (failed(marker.checkAndNotify(rewriter, linalgOp)))
|
|
return failure();
|
|
|
|
// If LinalgOp has results, they must all be tied to init tensors.
|
|
// We enforce this to ensure all tiled ops have been rewritten in
|
|
// "init tensor" form. This ensures tiling has anchor values into which to
|
|
// subtensor / subtensor_insert. Otherwise tiling would need to allocate which
|
|
// is not acceptable.
|
|
// This would not be the case with a special terminator op that generates the
|
|
// whole tensor (instead of inserting a subtensor). But the generator-based
|
|
// abstraction has other issues.
|
|
if (linalgOp.getNumInitTensors() != linalgOp.getOperation()->getNumResults())
|
|
return failure();
|
|
|
|
Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
|
|
|
|
if (!res)
|
|
return failure();
|
|
|
|
// Return relevant information to derived pattern.
|
|
tensorResults = res->tensorResults;
|
|
|
|
// New marker if specified.
|
|
marker.replaceLinalgMarker(rewriter, res->op.getOperation());
|
|
return success();
|
|
}
|
|
|
|
mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
|
|
StringRef opName, MLIRContext *context,
|
|
const LinalgDependenceGraph &dependenceGraph,
|
|
LinalgTilingOptions tilingOptions, LinalgFusionOptions fusionOptions,
|
|
LinalgMarker marker, LinalgMarker fusedOpMarker,
|
|
LinalgMarker originalOpMarker, PatternBenefit benefit)
|
|
: RewritePattern(opName, {}, benefit, context),
|
|
dependenceGraph(dependenceGraph), tilingOptions(tilingOptions),
|
|
fusionOptions(fusionOptions), marker(marker),
|
|
fusedOpMarker(fusedOpMarker), originalOpMarker(originalOpMarker) {}
|
|
|
|
LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
|
|
if (!linalgOp)
|
|
return failure();
|
|
if (failed(marker.checkAndNotify(rewriter, linalgOp)))
|
|
return failure();
|
|
if (!linalgOp.hasBufferSemantics())
|
|
return failure();
|
|
|
|
Optional<TiledAndFusedLinalgOps> tiledAndFusedOps = tileAndFuseLinalgOps(
|
|
rewriter, op, dependenceGraph, tilingOptions, fusionOptions);
|
|
if (!tiledAndFusedOps)
|
|
return failure();
|
|
marker.replaceLinalgMarker(rewriter, tiledAndFusedOps->op.getOperation());
|
|
for (auto fusedOp : tiledAndFusedOps->fusedProducers) {
|
|
fusedOpMarker.replaceLinalgMarker(rewriter, fusedOp.getOperation());
|
|
}
|
|
for (auto origProducerOp : tiledAndFusedOps->originalProducers)
|
|
originalOpMarker.replaceLinalgMarker(rewriter,
|
|
origProducerOp.getOperation());
|
|
rewriter.updateRootInPlace(
|
|
op, [&]() { originalOpMarker.replaceLinalgMarker(rewriter, op); });
|
|
return success();
|
|
}
|
|
|
|
/// Linalg base interchange pattern.
|
|
mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
|
|
StringRef opName, MLIRContext *context,
|
|
ArrayRef<unsigned> interchangeVector, LinalgMarker marker,
|
|
PatternBenefit benefit)
|
|
: RewritePattern(opName, {}, benefit, context), marker(marker),
|
|
interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
|
|
|
|
LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
|
|
if (!linalgOp)
|
|
return failure();
|
|
if (failed(marker.checkAndNotify(rewriter, linalgOp)))
|
|
return failure();
|
|
if (failed(interchangeGenericLinalgOpPrecondition(op, interchangeVector)))
|
|
return failure();
|
|
|
|
// TODO: figure out how this interplays with named ops. In particular this
|
|
// should break the named op property.
|
|
rewriter.updateRootInPlace(op, [&]() {
|
|
interchange(linalgOp, interchangeVector);
|
|
// New marker if specified.
|
|
marker.replaceLinalgMarker(rewriter, op);
|
|
});
|
|
return success();
|
|
}
|
|
|
|
mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
|
|
StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
|
|
LinalgMarker marker, PatternBenefit benefit)
|
|
: RewritePattern(opName, {}, benefit, context), marker(marker),
|
|
options(options) {}
|
|
|
|
LogicalResult mlir::linalg::LinalgBasePromotionPattern::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
if (failed(marker.checkAndNotify(rewriter, op)))
|
|
return failure();
|
|
if (failed(promoteSubviewsPrecondition(op, options)))
|
|
return failure();
|
|
|
|
// TODO: We cannot use root update here. This pattern is creating other ops,
|
|
// so if the promotion fails, those need to be cleaned up, which doesnt seem
|
|
// to be happening here. So to fail properly, we should be cloning the op and
|
|
// deleting the previous op. This needs more investigation.
|
|
rewriter.startRootUpdate(op);
|
|
Optional<LinalgOp> promotedOp = promoteSubViews(rewriter, op, options);
|
|
if (!promotedOp) {
|
|
rewriter.cancelRootUpdate(op);
|
|
return op->emitError("subview promotion failed");
|
|
}
|
|
rewriter.finalizeRootUpdate(op);
|
|
marker.replaceLinalgMarker(rewriter, op);
|
|
return success();
|
|
}
|
|
|
|
mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
|
|
StringRef opName, MLIRContext *context, LinalgMarker marker,
|
|
PatternBenefit benefit)
|
|
: RewritePattern(opName, {}, benefit, context), marker(marker) {}
|
|
|
|
LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
|
|
if (!linalgOp)
|
|
return failure();
|
|
if (failed(marker.checkAndNotify(rewriter, linalgOp)))
|
|
return failure();
|
|
if (failed(vectorizeLinalgOpPrecondition(op)))
|
|
return failure();
|
|
vectorizeLinalgOp(rewriter, op);
|
|
rewriter.eraseOp(op);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult mlir::linalg::applyStagedPatterns(
|
|
Operation *op, ArrayRef<FrozenRewritePatternList> stage1Patterns,
|
|
const FrozenRewritePatternList &stage2Patterns,
|
|
function_ref<LogicalResult(Operation *)> stage3Lambda) {
|
|
unsigned iteration = 0;
|
|
(void)iteration;
|
|
for (const auto &patterns : stage1Patterns) {
|
|
LLVM_DEBUG(DBGS() << "Before 1st stage, iter: " << ++iteration << "\n"
|
|
<< *op);
|
|
if (failed(applyPatternsAndFoldGreedily(op, patterns))) {
|
|
LLVM_DEBUG(DBGS() << "Underlying first stage rewrite did not converge");
|
|
return failure();
|
|
}
|
|
LLVM_DEBUG(DBGS() << "After 1st stage, iter: " << ++iteration << "\n"
|
|
<< *op);
|
|
if (failed(applyPatternsAndFoldGreedily(op, stage2Patterns))) {
|
|
LLVM_DEBUG(DBGS() << "Underlying 2nd stage rewrite did not converge");
|
|
return failure();
|
|
}
|
|
LLVM_DEBUG(DBGS() << "After 2nd stage, iter : " << iteration << "\n"
|
|
<< *op);
|
|
if (stage3Lambda) {
|
|
if (failed(stage3Lambda(op)))
|
|
return failure();
|
|
LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
|
|
<< *op);
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
|
|
/// been replaced by either:
|
|
/// - `min` if `positivePath` is true when we reach an occurrence of `dim`
|
|
/// - `max` if `positivePath` is true when we reach an occurrence of `dim`
|
|
/// `positivePath` is negated each time we hit a multiplicative or divisive
|
|
/// binary op with a constant negative coefficient.
|
|
static AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
|
|
AffineExpr max, bool positivePath = true) {
|
|
if (e == dim)
|
|
return positivePath ? min : max;
|
|
if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
|
|
AffineExpr lhs = bin.getLHS();
|
|
AffineExpr rhs = bin.getRHS();
|
|
if (bin.getKind() == mlir::AffineExprKind::Add)
|
|
return substWithMin(lhs, dim, min, max, positivePath) +
|
|
substWithMin(rhs, dim, min, max, positivePath);
|
|
|
|
auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
|
|
auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
|
|
if (c1 && c1.getValue() < 0)
|
|
return getAffineBinaryOpExpr(
|
|
bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
|
|
if (c2 && c2.getValue() < 0)
|
|
return getAffineBinaryOpExpr(
|
|
bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
|
|
return getAffineBinaryOpExpr(
|
|
bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
|
|
substWithMin(rhs, dim, min, max, positivePath));
|
|
}
|
|
return e;
|
|
}
|
|
|
|
/// Given the `lbVal`, `ubVal` and `stepVal` of a loop, append `lbVal` and
|
|
/// `ubVal` to `dims` and `stepVal` to `symbols`.
|
|
/// Create new AffineDimExpr (`%lb` and `%ub`) and AffineSymbolExpr (`%step`)
|
|
/// with positions matching the newly appended values. Substitute occurrences of
|
|
/// `dimExpr` by either the min expression (i.e. `%lb`) or the max expression
|
|
/// (i.e. `%lb + %step * floordiv(%ub -1 - %lb, %step)`), depending on whether
|
|
/// the induction variable is used with a positive or negative coefficient.
|
|
static AffineExpr substituteLoopInExpr(AffineExpr expr, AffineExpr dimExpr,
|
|
Value lbVal, Value ubVal, Value stepVal,
|
|
SmallVectorImpl<Value> &dims,
|
|
SmallVectorImpl<Value> &symbols) {
|
|
MLIRContext *ctx = lbVal.getContext();
|
|
AffineExpr lb = getAffineDimExpr(dims.size(), ctx);
|
|
dims.push_back(lbVal);
|
|
AffineExpr ub = getAffineDimExpr(dims.size(), ctx);
|
|
dims.push_back(ubVal);
|
|
AffineExpr step = getAffineSymbolExpr(symbols.size(), ctx);
|
|
symbols.push_back(stepVal);
|
|
LLVM_DEBUG(DBGS() << "Before: " << expr << "\n");
|
|
AffineExpr ee = substWithMin(expr, dimExpr, lb,
|
|
lb + step * ((ub - 1) - lb).floorDiv(step));
|
|
LLVM_DEBUG(DBGS() << "After: " << expr << "\n");
|
|
return ee;
|
|
}
|
|
|
|
/// Traverse the `dims` and substitute known min or max expressions in place of
|
|
/// induction variables in `exprs`.
|
|
static AffineMap substitute(AffineMap map, SmallVectorImpl<Value> &dims,
|
|
SmallVectorImpl<Value> &symbols) {
|
|
auto exprs = llvm::to_vector<4>(map.getResults());
|
|
for (AffineExpr &expr : exprs) {
|
|
bool substituted = true;
|
|
while (substituted) {
|
|
substituted = false;
|
|
for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
|
|
Value dim = dims[dimIdx];
|
|
AffineExpr dimExpr = getAffineDimExpr(dimIdx, expr.getContext());
|
|
LLVM_DEBUG(DBGS() << "Subst: " << dim << " @ " << dimExpr << "\n");
|
|
AffineExpr substitutedExpr;
|
|
if (auto forOp = scf::getForInductionVarOwner(dim))
|
|
substitutedExpr = substituteLoopInExpr(
|
|
expr, dimExpr, forOp.lowerBound(), forOp.upperBound(),
|
|
forOp.step(), dims, symbols);
|
|
|
|
if (auto parallelForOp = scf::getParallelForInductionVarOwner(dim))
|
|
for (unsigned idx = 0, e = parallelForOp.getNumLoops(); idx < e;
|
|
++idx)
|
|
substitutedExpr = substituteLoopInExpr(
|
|
expr, dimExpr, parallelForOp.lowerBound()[idx],
|
|
parallelForOp.upperBound()[idx], parallelForOp.step()[idx],
|
|
dims, symbols);
|
|
|
|
if (!substitutedExpr)
|
|
continue;
|
|
|
|
substituted = (substitutedExpr != expr);
|
|
expr = substitutedExpr;
|
|
}
|
|
}
|
|
|
|
// Cleanup and simplify the results.
|
|
// This needs to happen outside of the loop iterating on dims.size() since
|
|
// it modifies dims.
|
|
SmallVector<Value, 4> operands(dims.begin(), dims.end());
|
|
operands.append(symbols.begin(), symbols.end());
|
|
auto map = AffineMap::get(dims.size(), symbols.size(), exprs,
|
|
exprs.front().getContext());
|
|
|
|
LLVM_DEBUG(DBGS() << "Map to simplify: " << map << "\n");
|
|
|
|
// Pull in affine.apply operations and compose them fully into the
|
|
// result.
|
|
fullyComposeAffineMapAndOperands(&map, &operands);
|
|
canonicalizeMapAndOperands(&map, &operands);
|
|
map = simplifyAffineMap(map);
|
|
// Assign the results.
|
|
exprs.assign(map.getResults().begin(), map.getResults().end());
|
|
dims.assign(operands.begin(), operands.begin() + map.getNumDims());
|
|
symbols.assign(operands.begin() + map.getNumDims(), operands.end());
|
|
|
|
LLVM_DEBUG(DBGS() << "Map simplified: " << map << "\n");
|
|
}
|
|
|
|
assert(!exprs.empty() && "Unexpected empty exprs");
|
|
return AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
|
|
}
|
|
|
|
LogicalResult AffineMinSCFCanonicalizationPattern::matchAndRewrite(
|
|
AffineMinOp minOp, PatternRewriter &rewriter) const {
|
|
LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
|
|
<< "\n");
|
|
|
|
SmallVector<Value, 4> dims(minOp.getDimOperands()),
|
|
symbols(minOp.getSymbolOperands());
|
|
AffineMap map = substitute(minOp.getAffineMap(), dims, symbols);
|
|
|
|
LLVM_DEBUG(DBGS() << "Resulting map: " << map << "\n");
|
|
|
|
// Check whether any of the expressions, when subtracted from all other
|
|
// expressions, produces only >= 0 constants. If so, it is the min.
|
|
for (auto e : minOp.getAffineMap().getResults()) {
|
|
LLVM_DEBUG(DBGS() << "Candidate min: " << e << "\n");
|
|
if (!e.isSymbolicOrConstant())
|
|
continue;
|
|
|
|
auto isNonPositive = [](AffineExpr e) {
|
|
if (auto cst = e.dyn_cast<AffineConstantExpr>())
|
|
return cst.getValue() < 0;
|
|
return true;
|
|
};
|
|
|
|
// Build the subMap and check everything is statically known to be
|
|
// positive.
|
|
SmallVector<AffineExpr, 4> subExprs;
|
|
subExprs.reserve(map.getNumResults());
|
|
for (auto ee : map.getResults())
|
|
subExprs.push_back(ee - e);
|
|
MLIRContext *ctx = minOp.getContext();
|
|
AffineMap subMap = simplifyAffineMap(
|
|
AffineMap::get(map.getNumDims(), map.getNumSymbols(), subExprs, ctx));
|
|
LLVM_DEBUG(DBGS() << "simplified subMap: " << subMap << "\n");
|
|
if (llvm::any_of(subMap.getResults(), isNonPositive))
|
|
continue;
|
|
|
|
// Static min found.
|
|
if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
|
|
rewriter.replaceOpWithNewOp<ConstantIndexOp>(minOp, cst.getValue());
|
|
} else {
|
|
auto resultMap = AffineMap::get(0, map.getNumSymbols(), {e}, ctx);
|
|
SmallVector<Value, 4> resultOperands = dims;
|
|
resultOperands.append(symbols.begin(), symbols.end());
|
|
canonicalizeMapAndOperands(&resultMap, &resultOperands);
|
|
resultMap = simplifyAffineMap(resultMap);
|
|
rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
|
|
resultOperands);
|
|
}
|
|
return success();
|
|
}
|
|
|
|
return failure();
|
|
}
|