09dfc5713d
The greedy rewriter is used in many different flows and it has a lot of
convenience (work list management, debugging actions, tracing, etc). But
it combines two kinds of greedy behavior 1) how ops are matched, 2)
folding wherever it can.
These are independent forms of greedy and leads to inefficiency. E.g.,
cases where one need to create different phases in lowering and is
required to applying patterns in specific order split across different
passes. Using the driver one ends up needlessly retrying folding/having
multiple rounds of folding attempts, where one final run would have
sufficed.
Of course folks can locally avoid this behavior by just building their
own, but this is also a common requested feature that folks keep on
working around locally in suboptimal ways.
For downstream users, there should be no behavioral change. Updating
from the deprecated should just be a find and replace (e.g., `find ./
-type f -exec sed -i
's|applyPatternsAndFoldGreedily|applyPatternsGreedily|g' {} \;` variety)
as the API arguments hasn't changed between the two.
116 lines
4.0 KiB
C++
116 lines
4.0 KiB
C++
//===- SimplifyAffineStructures.cpp ---------------------------------------===//
|
|
//
|
|
// 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 a pass to simplify affine structures in operations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/Passes.h"
|
|
|
|
#include "mlir/Dialect/Affine/Analysis/Utils.h"
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Affine/Utils.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/IR/IntegerSet.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
namespace mlir {
|
|
namespace affine {
|
|
#define GEN_PASS_DEF_SIMPLIFYAFFINESTRUCTURES
|
|
#include "mlir/Dialect/Affine/Passes.h.inc"
|
|
} // namespace affine
|
|
} // namespace mlir
|
|
|
|
#define DEBUG_TYPE "simplify-affine-structure"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::affine;
|
|
|
|
namespace {
|
|
|
|
/// Simplifies affine maps and sets appearing in the operations of the Function.
|
|
/// This part is mainly to test the simplifyAffineExpr method. In addition,
|
|
/// all memrefs with non-trivial layout maps are converted to ones with trivial
|
|
/// identity layout ones.
|
|
struct SimplifyAffineStructures
|
|
: public affine::impl::SimplifyAffineStructuresBase<
|
|
SimplifyAffineStructures> {
|
|
void runOnOperation() override;
|
|
|
|
/// Utility to simplify an affine attribute and update its entry in the parent
|
|
/// operation if necessary.
|
|
template <typename AttributeT>
|
|
void simplifyAndUpdateAttribute(Operation *op, StringAttr name,
|
|
AttributeT attr) {
|
|
auto &simplified = simplifiedAttributes[attr];
|
|
if (simplified == attr)
|
|
return;
|
|
|
|
// This is a newly encountered attribute.
|
|
if (!simplified) {
|
|
// Try to simplify the value of the attribute.
|
|
auto value = attr.getValue();
|
|
auto simplifiedValue = simplify(value);
|
|
if (simplifiedValue == value) {
|
|
simplified = attr;
|
|
return;
|
|
}
|
|
simplified = AttributeT::get(simplifiedValue);
|
|
}
|
|
|
|
// Simplification was successful, so update the attribute.
|
|
op->setAttr(name, simplified);
|
|
}
|
|
|
|
IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
|
|
|
|
/// Performs basic affine map simplifications.
|
|
AffineMap simplify(AffineMap map) {
|
|
MutableAffineMap mMap(map);
|
|
mMap.simplify();
|
|
return mMap.getAffineMap();
|
|
}
|
|
|
|
DenseMap<Attribute, Attribute> simplifiedAttributes;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<OperationPass<func::FuncOp>>
|
|
mlir::affine::createSimplifyAffineStructuresPass() {
|
|
return std::make_unique<SimplifyAffineStructures>();
|
|
}
|
|
|
|
void SimplifyAffineStructures::runOnOperation() {
|
|
auto func = getOperation();
|
|
simplifiedAttributes.clear();
|
|
RewritePatternSet patterns(func.getContext());
|
|
AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
|
|
AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
|
|
AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
|
|
FrozenRewritePatternSet frozenPatterns(std::move(patterns));
|
|
|
|
// The simplification of affine attributes will likely simplify the op. Try to
|
|
// fold/apply canonicalization patterns when we have affine dialect ops.
|
|
SmallVector<Operation *> opsToSimplify;
|
|
func.walk([&](Operation *op) {
|
|
for (auto attr : op->getAttrs()) {
|
|
if (auto mapAttr = dyn_cast<AffineMapAttr>(attr.getValue()))
|
|
simplifyAndUpdateAttribute(op, attr.getName(), mapAttr);
|
|
else if (auto setAttr = dyn_cast<IntegerSetAttr>(attr.getValue()))
|
|
simplifyAndUpdateAttribute(op, attr.getName(), setAttr);
|
|
}
|
|
|
|
if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
|
|
opsToSimplify.push_back(op);
|
|
});
|
|
GreedyRewriteConfig config;
|
|
config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
|
|
(void)applyOpPatternsGreedily(opsToSimplify, frozenPatterns, config);
|
|
}
|