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 "mlir/Dialect/Affine/Passes.h"
14
15#include "mlir/Dialect/Affine/Analysis/Utils.h"
16#include "mlir/Dialect/Affine/IR/AffineOps.h"
17#include "mlir/Dialect/Affine/Utils.h"
18#include "mlir/Dialect/Func/IR/FuncOps.h"
19#include "mlir/IR/IntegerSet.h"
20#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
21
22namespace mlir {
23namespace affine {
24#define GEN_PASS_DEF_SIMPLIFYAFFINESTRUCTURES
25#include "mlir/Dialect/Affine/Passes.h.inc"
26} // namespace affine
27} // namespace mlir
28
29#define DEBUG_TYPE "simplify-affine-structure"
30
31using namespace mlir;
32using namespace mlir::affine;
33
34namespace {
35
36/// Simplifies affine maps and sets appearing in the operations of the Function.
37/// This part is mainly to test the simplifyAffineExpr method. In addition,
38/// all memrefs with non-trivial layout maps are converted to ones with trivial
39/// identity layout ones.
40struct SimplifyAffineStructures
41 : public affine::impl::SimplifyAffineStructuresBase<
42 SimplifyAffineStructures> {
43 void runOnOperation() override;
44
45 /// Utility to simplify an affine attribute and update its entry in the parent
46 /// operation if necessary.
47 template <typename AttributeT>
48 void simplifyAndUpdateAttribute(Operation *op, StringAttr name,
49 AttributeT attr) {
50 auto &simplified = simplifiedAttributes[attr];
51 if (simplified == attr)
52 return;
53
54 // This is a newly encountered attribute.
55 if (!simplified) {
56 // Try to simplify the value of the attribute.
57 auto value = attr.getValue();
58 auto simplifiedValue = simplify(value);
59 if (simplifiedValue == value) {
60 simplified = attr;
61 return;
62 }
63 simplified = AttributeT::get(simplifiedValue);
64 }
65
66 // Simplification was successful, so update the attribute.
67 op->setAttr(name, simplified);
68 }
69
70 IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
71
72 /// Performs basic affine map simplifications.
73 AffineMap simplify(AffineMap map) {
74 MutableAffineMap mMap(map);
75 mMap.simplify();
76 return mMap.getAffineMap();
77 }
78
79 DenseMap<Attribute, Attribute> simplifiedAttributes;
80};
81
82} // namespace
83
84std::unique_ptr<OperationPass<func::FuncOp>>
85mlir::affine::createSimplifyAffineStructuresPass() {
86 return std::make_unique<SimplifyAffineStructures>();
87}
88
89void SimplifyAffineStructures::runOnOperation() {
90 auto func = getOperation();
91 simplifiedAttributes.clear();
92 RewritePatternSet patterns(func.getContext());
93 AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
94 AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
95 AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
96 FrozenRewritePatternSet frozenPatterns(std::move(patterns));
97
98 // The simplification of affine attributes will likely simplify the op. Try to
99 // fold/apply canonicalization patterns when we have affine dialect ops.
100 SmallVector<Operation *> opsToSimplify;
101 func.walk([&](Operation *op) {
102 for (auto attr : op->getAttrs()) {
103 if (auto mapAttr = dyn_cast<AffineMapAttr>(attr.getValue()))
104 simplifyAndUpdateAttribute(op, attr.getName(), mapAttr);
105 else if (auto setAttr = dyn_cast<IntegerSetAttr>(attr.getValue()))
106 simplifyAndUpdateAttribute(op, attr.getName(), setAttr);
107 }
108
109 if (isa<AffineForOp, AffineIfOp, AffineApplyOp>(Val: op))
110 opsToSimplify.push_back(Elt: op);
111 });
112 GreedyRewriteConfig config;
113 config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
114 (void)applyOpPatternsAndFold(ops: opsToSimplify, patterns: frozenPatterns, config);
115}
116

source code of mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp