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 "mlir/Analysis/Utils.h"
14 #include "mlir/Dialect/Affine/Passes.h"
15 #include "mlir/IR/IntegerSet.h"
16 #include "mlir/Pass/Pass.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 PassWrapper<SimplifyAffineStructures, FunctionPass> {
31 /// Include the generated pass utilities.
32 #define GEN_PASS_SimplifyAffineStructures
33 #include "mlir/Dialect/Affine/Passes.h.inc"
34 
35   void runOnFunction() override;
36 
37   /// Utility to simplify an affine attribute and update its entry in the parent
38   /// operation if necessary.
39   template <typename AttributeT>
40   void simplifyAndUpdateAttribute(Operation *op, Identifier name,
41                                   AttributeT attr) {
42     auto &simplified = simplifiedAttributes[attr];
43     if (simplified == attr)
44       return;
45 
46     // This is a newly encountered attribute.
47     if (!simplified) {
48       // Try to simplify the value of the attribute.
49       auto value = attr.getValue();
50       auto simplifiedValue = simplify(value);
51       if (simplifiedValue == value) {
52         simplified = attr;
53         return;
54       }
55       simplified = AttributeT::get(simplifiedValue);
56     }
57 
58     // Simplification was successful, so update the attribute.
59     op->setAttr(name, simplified);
60   }
61 
62   IntegerSet simplify(IntegerSet set) { return simplifyIntegerSet(set); }
63 
64   /// Performs basic affine map simplifications.
65   AffineMap simplify(AffineMap map) {
66     MutableAffineMap mMap(map);
67     mMap.simplify();
68     return mMap.getAffineMap();
69   }
70 
71   DenseMap<Attribute, Attribute> simplifiedAttributes;
72 };
73 
74 } // end anonymous namespace
75 
76 std::unique_ptr<OperationPass<FuncOp>>
77 mlir::createSimplifyAffineStructuresPass() {
78   return std::make_unique<SimplifyAffineStructures>();
79 }
80 
81 void SimplifyAffineStructures::runOnFunction() {
82   auto func = getFunction();
83   simplifiedAttributes.clear();
84   func.walk([&](Operation *opInst) {
85     for (auto attr : opInst->getAttrs()) {
86       if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
87         simplifyAndUpdateAttribute(opInst, attr.first, mapAttr);
88       else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
89         simplifyAndUpdateAttribute(opInst, attr.first, setAttr);
90     }
91   });
92 
93   // Turn memrefs' non-identity layouts maps into ones with identity. Collect
94   // alloc ops first and then process since normalizeMemRef replaces/erases ops
95   // during memref rewriting.
96   SmallVector<AllocOp, 4> allocOps;
97   func.walk([&](AllocOp op) { allocOps.push_back(op); });
98   for (auto allocOp : allocOps) {
99     normalizeMemRef(allocOp);
100   }
101 }
102