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/Dialect/Affine/Analysis/Utils.h"
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/Dialect/Affine/Passes.h"
17 #include "mlir/Dialect/Affine/Utils.h"
18 #include "mlir/IR/IntegerSet.h"
19 #include "mlir/Transforms/GreedyPatternRewriteDriver.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 runOnOperation() 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__anonfe12588e0111::SimplifyAffineStructures38 void simplifyAndUpdateAttribute(Operation *op, StringAttr 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__anonfe12588e0111::SimplifyAffineStructures60 IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
61
62 /// Performs basic affine map simplifications.
simplify__anonfe12588e0111::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 } // namespace
73
74 std::unique_ptr<OperationPass<func::FuncOp>>
createSimplifyAffineStructuresPass()75 mlir::createSimplifyAffineStructuresPass() {
76 return std::make_unique<SimplifyAffineStructures>();
77 }
78
runOnOperation()79 void SimplifyAffineStructures::runOnOperation() {
80 auto func = getOperation();
81 simplifiedAttributes.clear();
82 RewritePatternSet patterns(func.getContext());
83 AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
84 AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
85 AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
86 FrozenRewritePatternSet frozenPatterns(std::move(patterns));
87
88 // The simplification of affine attributes will likely simplify the op. Try to
89 // fold/apply canonicalization patterns when we have affine dialect ops.
90 SmallVector<Operation *> opsToSimplify;
91 func.walk([&](Operation *op) {
92 for (auto attr : op->getAttrs()) {
93 if (auto mapAttr = attr.getValue().dyn_cast<AffineMapAttr>())
94 simplifyAndUpdateAttribute(op, attr.getName(), mapAttr);
95 else if (auto setAttr = attr.getValue().dyn_cast<IntegerSetAttr>())
96 simplifyAndUpdateAttribute(op, attr.getName(), setAttr);
97 }
98
99 if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(op))
100 opsToSimplify.push_back(op);
101 });
102 (void)applyOpPatternsAndFold(opsToSimplify, frozenPatterns, /*strict=*/true);
103 }
104