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>
38   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 
60   IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
61 
62   /// Performs basic affine map simplifications.
63   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>>
75 mlir::createSimplifyAffineStructuresPass() {
76   return std::make_unique<SimplifyAffineStructures>();
77 }
78 
79 void SimplifyAffineStructures::runOnFunction() {
80   auto func = getFunction();
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