[MLIR] Add method to drop duplicate result exprs from AffineMap

Add a method that given an affine map returns another with just its unique
results. Use this to drop redundant bounds in max/min for affine.for. Update
affine.for's canonicalization pattern and createCanonicalizedForOp to use
this.

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

mlir/include/mlir/IR/AffineMap.h

@@ -215,9 +215,13 @@ private:
   MLIRContext *context;
 };
 
-/// Simplify an affine map by simplifying its underlying AffineExpr results.
+/// Simplifies an affine map by simplifying its underlying AffineExpr results.
 AffineMap simplifyAffineMap(AffineMap map);
 
+/// Returns a map with the same dimension and symbol count as `map`, but whose
+/// results are the unique affine expressions of `map`.
+AffineMap removeDuplicateExprs(AffineMap map);
+
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
 /// Returns an empty map if the input map is empty or if `map` is not invertible

mlir/include/mlir/Transforms/LoopUtils.h

@@ -265,8 +265,8 @@ void gatherLoops(FuncOp func,
                  std::vector<SmallVector<AffineForOp, 2>> &depthToLoops);
 
 /// Creates an AffineForOp while ensuring that the lower and upper bounds are
-/// canonicalized, i.e., unused and duplicate operands are removed, and any
-/// constant operands propagated/folded in.
+/// canonicalized, i.e., unused and duplicate operands are removed, any constant
+/// operands propagated/folded in, and duplicate bound maps dropped.
 AffineForOp createCanonicalizedAffineForOp(OpBuilder b, Location loc,
                                            ValueRange lbOperands,
                                            AffineMap lbMap,

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

@@ -1383,7 +1383,10 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
   auto prevUbMap = ubMap;
 
   canonicalizeMapAndOperands(&lbMap, &lbOperands);
+  lbMap = removeDuplicateExprs(lbMap);
+
   canonicalizeMapAndOperands(&ubMap, &ubOperands);
+  ubMap = removeDuplicateExprs(ubMap);
 
   // Any canonicalization change always leads to updated map(s).
   if (lbMap == prevLbMap && ubMap == prevUbMap)

mlir/lib/IR/AffineMap.cpp

@@ -325,6 +325,15 @@ AffineMap mlir::simplifyAffineMap(AffineMap map) {
   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), exprs);
 }
 
+AffineMap mlir::removeDuplicateExprs(AffineMap map) {
+  auto results = map.getResults();
+  SmallVector<AffineExpr, 4> uniqueExprs(results.begin(), results.end());
+  uniqueExprs.erase(std::unique(uniqueExprs.begin(), uniqueExprs.end()),
+                    uniqueExprs.end());
+  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), uniqueExprs,
+                        map.getContext());
+}
+
 AffineMap mlir::inversePermutation(AffineMap map) {
   if (map.isEmpty())
     return map;

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -2132,8 +2132,10 @@ AffineForOp mlir::createCanonicalizedAffineForOp(
 
   fullyComposeAffineMapAndOperands(&lbMap, &lowerOperands);
   canonicalizeMapAndOperands(&lbMap, &lowerOperands);
+  lbMap = removeDuplicateExprs(lbMap);
   fullyComposeAffineMapAndOperands(&ubMap, &upperOperands);
   canonicalizeMapAndOperands(&ubMap, &upperOperands);
+  ubMap = removeDuplicateExprs(ubMap);
 
   return b.create<AffineForOp>(loc, lowerOperands, lbMap, upperOperands, ubMap,
                                step);

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

@@ -344,7 +344,7 @@ func @loop_min_max(%N : index) {
   return
 }
 
-#map_7_values = affine_map<(i) -> (i, i, i, i, i, i, i)>
+#map_7_values = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5, d6)>
 
 // Check that the "min" (cmpi "slt" + select) reduction sequence is emitted
 // correctly for a an affine map with 7 results.
@@ -369,8 +369,8 @@ func @loop_min_max(%N : index) {
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
-func @min_reduction_tree(%v : index) {
-  affine.for %i = 0 to min #map_7_values(%v)[] {
+func @min_reduction_tree(%v1 : index, %v2 : index, %v3 : index, %v4 : index, %v5 : index, %v6 : index, %v7 : index) {
+  affine.for %i = 0 to min #map_7_values(%v1, %v2, %v3, %v4, %v5, %v6, %v7)[] {
     call @body(%i) : (index) -> ()
   }
   return
@@ -619,4 +619,4 @@ func @affine_max(%arg0: index, %arg1: index) -> index{
   // CHECK: select %[[cmp]], %[[first]], %[[second]]
   %0 = affine.max affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1)
   return %0 : index
-}
+}

mlir/test/Dialect/Affine/canonicalize.mlir

@@ -594,3 +594,15 @@ func @rep(%arg0 : index, %arg1 : index) -> index {
   %1 = affine.min #map2(%0)[%arg1]
   return %1 : index
 }
+
+// -----
+// CHECK-DAG: #[[lb:.*]] = affine_map<()[s0] -> (s0)>
+// CHECK-DAG: #[[ub:.*]] = affine_map<()[s0] -> (s0 + 2)>
+
+func @drop_duplicate_bounds(%N : index) {
+  // affine.for %i = max #lb(%arg0) to min #ub(%arg0)
+  affine.for %i = max affine_map<(d0) -> (d0, d0)>(%N) to min affine_map<(d0) -> (d0 + 2, d0 + 2)>(%N) {
+    "foo"() : () -> ()
+  }
+  return
+}