[Mlir-commits] [mlir] 23fd266 - [mlir][affine] Deduplicate affine min/max op expressions

[Lei Zhang]

Author: Lei Zhang
Date: 2021-03-24T18:17:57-04:00
New Revision: 23fd26608ca899f13445908e6afd61d88f6aea9c

URL: https://github.com/llvm/llvm-project/commit/23fd26608ca899f13445908e6afd61d88f6aea9c
DIFF: https://github.com/llvm/llvm-project/commit/23fd26608ca899f13445908e6afd61d88f6aea9c.diff

LOG: [mlir][affine] Deduplicate affine min/max op expressions

If there are multiple identical expressions in an affine
min/max op's map, we can just keep one.

Differential Revision: https://reviews.llvm.org/D99015

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 930d0bce96c6..305e39553503 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2327,6 +2327,34 @@ static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) {
   return IntegerAttr::get(IndexType::get(op.getContext()), *resultIt);
 }
 
+/// Remove duplicated expressions in affine min/max ops.
+template <typename T>
+struct DeduplicateAffineMinMaxExpressions : public OpRewritePattern<T> {
+  using OpRewritePattern<T>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(T affineOp,
+                                PatternRewriter &rewriter) const override {
+    AffineMap oldMap = affineOp.getAffineMap();
+
+    SmallVector<AffineExpr, 4> newExprs;
+    for (AffineExpr expr : oldMap.getResults()) {
+      // This is a linear scan over newExprs, but it should be fine given that
+      // we typically just have a few expressions per op.
+      if (!llvm::is_contained(newExprs, expr))
+        newExprs.push_back(expr);
+    }
+
+    if (newExprs.size() == oldMap.getNumResults())
+      return failure();
+
+    auto newMap = AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(),
+                                 newExprs, rewriter.getContext());
+    rewriter.replaceOpWithNewOp<T>(affineOp, newMap, affineOp.getMapOperands());
+
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // AffineMinOp
 //===----------------------------------------------------------------------===//
@@ -2340,7 +2368,8 @@ OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
 
 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
-  patterns.add<SimplifyAffineOp<AffineMinOp>>(context);
+  patterns.add<DeduplicateAffineMinMaxExpressions<AffineMinOp>,
+               SimplifyAffineOp<AffineMinOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
@@ -2356,7 +2385,8 @@ OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
 
 void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
-  patterns.add<SimplifyAffineOp<AffineMaxOp>>(context);
+  patterns.add<DeduplicateAffineMinMaxExpressions<AffineMaxOp>,
+               SimplifyAffineOp<AffineMaxOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1fddf5c882c1..4f20431624b4 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -694,3 +694,27 @@ func @compose_affine_maps_div_symbol(%A : memref<i64>, %i0 : index, %i1 : index)
   }
   return
 }
+
+// -----
+
+// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
+
+// CHECK: func @deduplicate_affine_min_expressions
+// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
+func @deduplicate_affine_min_expressions(%i0: index, %i1: index) -> index {
+  // CHECK:  affine.min #[[MAP]]()[%[[I0]], %[[I1]]]
+  %0 = affine.min affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
+  return %0: index
+}
+
+// -----
+
+// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1, s0 * s1)>
+
+// CHECK: func @deduplicate_affine_max_expressions
+// CHECK-SAME: (%[[I0:.+]]: index, %[[I1:.+]]: index)
+func @deduplicate_affine_max_expressions(%i0: index, %i1: index) -> index {
+  // CHECK:  affine.max #[[MAP]]()[%[[I0]], %[[I1]]]
+  %0 = affine.max affine_map<()[s0, s1] -> (s0 + s1, s0 * s1, s1 + s0, s0 * s1)> ()[%i0, %i1]
+  return %0: index
+}