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.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "PassDetail.h"
14 #include "mlir/Analysis/Utils.h"
15 #include "mlir/Dialect/Affine/Passes.h"
16 #include "mlir/IR/IntegerSet.h"
17 #include "mlir/Transforms/Utils.h"
18 
19 #define DEBUG_TYPE "simplify-affine-structure"
20 
21 using namespace mlir;
22 
23 namespace {
24 
25 /// Simplifies affine maps and sets appearing in the operations of the Function.
26 /// This part is mainly to test the simplifyAffineExpr method. In addition,
27 /// all memrefs with non-trivial layout maps are converted to ones with trivial
28 /// identity layout ones.
29 struct SimplifyAffineStructures
30     : public SimplifyAffineStructuresBase<SimplifyAffineStructures> {
31   void runOnFunction() override;
32 
33   /// Utility to simplify an affine attribute and update its entry in the parent
34   /// operation if necessary.
35   template <typename AttributeT>
36   void simplifyAndUpdateAttribute(Operation *op, Identifier name,
37                                   AttributeT attr) {
38     auto &simplified = simplifiedAttributes[attr];
39     if (simplified == attr)
40       return;
41 
42     // This is a newly encountered attribute.
43     if (!simplified) {
44       // Try to simplify the value of the attribute.
45       auto value = attr.getValue();
46       auto simplifiedValue = simplify(value);
47       if (simplifiedValue == value) {
48         simplified = attr;
49         return;
50       }
51       simplified = AttributeT::get(simplifiedValue);
52     }
53 
54     // Simplification was successful, so update the attribute.
55     op->setAttr(name, simplified);
56   }
57 
58   IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
59 
60   /// Performs basic affine map simplifications.
61   AffineMap simplify(AffineMap map) {
62     MutableAffineMap mMap(map);
63     mMap.simplify();
64     return mMap.getAffineMap();
65   }
66 
67   DenseMap<Attribute, Attribute> simplifiedAttributes;
68 };
69 
70 } // end anonymous namespace
71 
72 std::unique_ptr<OperationPass<FuncOp>>
73 mlir::createSimplifyAffineStructuresPass() {
74   return std::make_unique<SimplifyAffineStructures>();
75 }
76 
77 void SimplifyAffineStructures::runOnFunction() {
78   auto func = getFunction();
79   simplifiedAttributes.clear();
80   func.walk([&](Operation *opInst) {
81     for (auto attr : opInst->getAttrs()) {
82       if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
83         simplifyAndUpdateAttribute(opInst, attr.first, mapAttr);
84       else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
85         simplifyAndUpdateAttribute(opInst, attr.first, setAttr);
86     }
87   });
88 
89   // Turn memrefs' non-identity layouts maps into ones with identity. Collect
90   // alloc ops first and then process since normalizeMemRef replaces/erases ops
91   // during memref rewriting.
92   SmallVector<AllocOp, 4> allocOps;
93   func.walk([&](AllocOp op) { allocOps.push_back(op); });
94   for (auto allocOp : allocOps) {
95     normalizeMemRef(allocOp);
96   }
97 }
98