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

[antiagainst]

Summary

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

Depends On D99014

[mravishankar]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2369–2372

Why not just add it for Max too?

[antiagainst]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2369–2372

Isn't that at L2388? Am I missing something?

[antiagainst]

Rebase

[mravishankar]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2369–2372

Oh right. Misread it.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) {
   auto resultIt = std::is_same<T, AffineMinOp>::value
                       ? std::min_element(results.begin(), results.end())
                       : std::max_element(results.begin(), results.end());
   if (resultIt == results.end())
     return {};
   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
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
 //
 
 OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
   return foldMinMaxOp(*this, operands);
 }
 
 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
-  patterns.add<SimplifyAffineOp<AffineMinOp>>(context);
+  patterns.add<DeduplicateAffineMinMaxExpressions<AffineMinOp>,
+               SimplifyAffineOp<AffineMinOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineMaxOp
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
 //
 
 OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
   return foldMinMaxOp(*this, operands);
 }
 
 void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
-  patterns.add<SimplifyAffineOp<AffineMaxOp>>(context);
+  patterns.add<DeduplicateAffineMinMaxExpressions<AffineMaxOp>,
+               SimplifyAffineOp<AffineMaxOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AffinePrefetchOp
 //===----------------------------------------------------------------------===//
 
 //
 // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>

mlir/test/Dialect/Affine/canonicalize.mlir

func @compose_affine_maps_div_symbol(%A : memref<i64>, %i0 : index, %i1 : index) {
   memref.store %3, %A[]: memref<i64>
   affine.for %i2 = 0 to 3 {
     %4 = affine.apply affine_map<(d0)[s0, s1] -> (d0 ceildiv s1 + s0 + s0 * 3)> (%i2)[%0, %1]
     %5 = index_cast %4: index to i64
     memref.store %5, %A[]: memref<i64>
   }
   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
+}