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 22 namespace mlir { 23 namespace 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 31 using namespace mlir; 32 using namespace mlir::affine; 33 34 namespace { 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. 40 struct 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 84 std::unique_ptr<OperationPass<func::FuncOp>> 85 mlir::affine::createSimplifyAffineStructuresPass() { 86 return std::make_unique<SimplifyAffineStructures>(); 87 } 88 89 void 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>(op)) 110 opsToSimplify.push_back(op); 111 }); 112 GreedyRewriteConfig config; 113 config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps; 114 (void)applyOpPatternsGreedily(opsToSimplify, frozenPatterns, config); 115 } 116