117 lines · cpp
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-exception6//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_SIMPLIFYAFFINESTRUCTURES25#include "mlir/Dialect/Affine/Passes.h.inc"26} // namespace affine27} // namespace mlir28 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 trivial39/// identity layout ones.40struct SimplifyAffineStructures41 : public affine::impl::SimplifyAffineStructuresBase<42 SimplifyAffineStructures> {43 void runOnOperation() override;44 45 /// Utility to simplify an affine attribute and update its entry in the parent46 /// 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} // namespace83 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 to99 // 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 (void)applyOpPatternsGreedily(113 opsToSimplify, frozenPatterns,114 GreedyRewriteConfig().setStrictness(115 GreedyRewriteStrictness::ExistingAndNewOps));116}117