github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/mlir/lib/Transforms/SimplifyAffineStructures.cpp (about)

     1  //===- SimplifyAffineStructures.cpp ---------------------------------------===//
     2  //
     3  // Copyright 2019 The MLIR Authors.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //   http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  // =============================================================================
    17  //
    18  // This file implements a pass to simplify affine structures.
    19  //
    20  //===----------------------------------------------------------------------===//
    21  
    22  #include "mlir/Analysis/AffineStructures.h"
    23  #include "mlir/IR/AffineMap.h"
    24  #include "mlir/IR/Function.h"
    25  #include "mlir/IR/IntegerSet.h"
    26  #include "mlir/IR/Operation.h"
    27  #include "mlir/Pass/Pass.h"
    28  #include "mlir/Transforms/Passes.h"
    29  
    30  #define DEBUG_TYPE "simplify-affine-structure"
    31  
    32  using namespace mlir;
    33  
    34  namespace {
    35  
    36  /// Simplifies all affine expressions appearing in the operations of
    37  /// the Function. This is mainly to test the simplifyAffineExpr method.
    38  /// TODO(someone): This should just be defined as a canonicalization pattern
    39  /// on AffineMap and driven from the existing canonicalization pass.
    40  struct SimplifyAffineStructures
    41      : public FunctionPass<SimplifyAffineStructures> {
    42    void runOnFunction() override;
    43  
    44    /// Utility to simplify an affine attribute and update its entry in the parent
    45    /// operation if necessary.
    46    template <typename AttributeT>
    47    void simplifyAndUpdateAttribute(Operation *op, Identifier name,
    48                                    AttributeT attr) {
    49      auto &simplified = simplifiedAttributes[attr];
    50      if (simplified == attr)
    51        return;
    52  
    53      // This is a newly encountered attribute.
    54      if (!simplified) {
    55        // Try to simplify the value of the attribute.
    56        auto value = attr.getValue();
    57        auto simplifiedValue = simplify(value);
    58        if (simplifiedValue == value) {
    59          simplified = attr;
    60          return;
    61        }
    62        simplified = AttributeT::get(simplifiedValue);
    63      }
    64  
    65      // Simplification was successful, so update the attribute.
    66      op->setAttr(name, simplified);
    67    }
    68  
    69    /// Performs basic integer set simplifications. Checks if it's empty, and
    70    /// replaces it with the canonical empty set if it is.
    71    IntegerSet simplify(IntegerSet set) {
    72      FlatAffineConstraints fac(set);
    73      if (fac.isEmpty())
    74        return IntegerSet::getEmptySet(set.getNumDims(), set.getNumSymbols(),
    75                                       &getContext());
    76      return set;
    77    }
    78  
    79    /// Performs basic affine map simplifications.
    80    AffineMap simplify(AffineMap map) {
    81      MutableAffineMap mMap(map);
    82      mMap.simplify();
    83      return mMap.getAffineMap();
    84    }
    85  
    86    DenseMap<Attribute, Attribute> simplifiedAttributes;
    87  };
    88  
    89  } // end anonymous namespace
    90  
    91  std::unique_ptr<FunctionPassBase> mlir::createSimplifyAffineStructuresPass() {
    92    return std::make_unique<SimplifyAffineStructures>();
    93  }
    94  
    95  void SimplifyAffineStructures::runOnFunction() {
    96    simplifiedAttributes.clear();
    97    getFunction().walk([&](Operation *opInst) {
    98      for (auto attr : opInst->getAttrs()) {
    99        if (auto mapAttr = attr.second.dyn_cast<AffineMapAttr>())
   100          simplifyAndUpdateAttribute(opInst, attr.first, mapAttr);
   101        else if (auto setAttr = attr.second.dyn_cast<IntegerSetAttr>())
   102          simplifyAndUpdateAttribute(opInst, attr.first, setAttr);
   103      }
   104    });
   105  }
   106  
   107  static PassRegistration<SimplifyAffineStructures>
   108      pass("simplify-affine-structures", "Simplify affine expressions");