1 //===- SimplifyAffineStructures.cpp ---------------------------------------===//
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements a pass to simplify affine structures in operations.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "PassDetail.h"
14 #include "mlir/Analysis/Utils.h"
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/Dialect/Affine/Passes.h"
17 #include "mlir/IR/IntegerSet.h"
18 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
19 #include "mlir/Transforms/Utils.h"
20
21 #define DEBUG_TYPE "simplify-affine-structure"
22
23 using namespace mlir;
24
25 namespace {
26
27 /// Simplifies affine maps and sets appearing in the operations of the Function.
28 /// This part is mainly to test the simplifyAffineExpr method. In addition,
29 /// all memrefs with non-trivial layout maps are converted to ones with trivial
30 /// identity layout ones.
31 struct SimplifyAffineStructures
32 : public SimplifyAffineStructuresBase<SimplifyAffineStructures> {
33 void runOnFunction() override;
34
35 /// Utility to simplify an affine attribute and update its entry in the parent
36 /// operation if necessary.
37 template <typename AttributeT>
simplifyAndUpdateAttribute__anonf5985f950111::SimplifyAffineStructures38 void simplifyAndUpdateAttribute(Operation *op, Identifier name,
39 AttributeT attr) {
40 auto &simplified = simplifiedAttributes[attr];
41 if (simplified == attr)
42 return;
43
44 // This is a newly encountered attribute.
45 if (!simplified) {
46 // Try to simplify the value of the attribute.
47 auto value = attr.getValue();
48 auto simplifiedValue = simplify(value);
49 if (simplifiedValue == value) {
50 simplified = attr;
51 return;
52 }
53 simplified = AttributeT::get(simplifiedValue);
54 }
55
56 // Simplification was successful, so update the attribute.
57 op->setAttr(name, simplified);
58 }
59
simplify__anonf5985f950111::SimplifyAffineStructures60 IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
61
62 /// Performs basic affine map simplifications.
simplify__anonf5985f950111::SimplifyAffineStructures63 AffineMap simplify(AffineMap map) {
64 MutableAffineMap mMap(map);
65 mMap.simplify();
66 return mMap.getAffineMap();
67 }
68
69 DenseMap<Attribute, Attribute> simplifiedAttributes;
70 };
71
72 } // end anonymous namespace
73
74 std::unique_ptr<OperationPass<FuncOp>>
createSimplifyAffineStructuresPass()75 mlir::createSimplifyAffineStructuresPass() {
76 return std::make_unique<SimplifyAffineStructures>();
77 }
78
runOnFunction()79 void SimplifyAffineStructures::runOnFunction() {
80 auto func = getFunction();
81 simplifiedAttributes.clear();
82 OwningRewritePatternList patterns;
83 AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
84 AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
85 AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
86 FrozenRewritePatternList frozenPatterns(std::move(patterns));
87 func.walk([&](Operation *op) {
88 for (auto attr : op->getAttrs()) {
89 if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
90 simplifyAndUpdateAttribute(op, attr.first, mapAttr);
91 else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
92 simplifyAndUpdateAttribute(op, attr.first, setAttr);
93 }
94
95 // The simplification of the attribute will likely simplify the op. Try to
96 // fold / apply canonicalization patterns when we have affine dialect ops.
97 if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
98 applyOpPatternsAndFold(op, frozenPatterns);
99 });
100 }
101