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

[bondhugula]

Summary

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.

[bondhugula]

Update to use recently added method for 0 result affine maps as well.

mlir/include/mlir/IR/AffineMap.h

private:
   SmallVector<AffineExpr, 8> results;
   unsigned numDims;
   unsigned numSymbols;
   /// A pointer to the IR's context to store all newly created
   /// AffineExprStorage's.
   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
 /// (i.e. `map` does not contain a subset that is a permutation of full domain
 /// rank).
 ///
 /// Prerequisites:
 ///   1. `map` has no symbols.

mlir/include/mlir/Transforms/LoopUtils.h

 void mapLoopToProcessorIds(loop::ForOp forOp, ArrayRef<Value> processorId,
                            ArrayRef<Value> numProcessors);
 
 /// Gathers all AffineForOps in 'func' grouped by loop depth.
 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,
                                            ValueRange ubOperands,
                                            AffineMap ubMap, int64_t step = 1);
 
 /// Separates full tiles from partial tiles for a perfect nest `nest` by
 /// generating a conditional guard that selects between the full tile version

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

static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
   SmallVector<Value, 4> ubOperands(forOp.getUpperBoundOperands());
 
   auto lbMap = forOp.getLowerBoundMap();
   auto ubMap = forOp.getUpperBoundMap();
   auto prevLbMap = lbMap;
   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)
     return failure();
 
   if (lbMap != prevLbMap)
     forOp.setLowerBound(lbOperands, lbMap);
   if (ubMap != prevUbMap)

mlir/lib/IR/AffineMap.cpp

AffineMap mlir::simplifyAffineMap(AffineMap map) {
   SmallVector<AffineExpr, 8> exprs;
   for (auto e : map.getResults()) {
     exprs.push_back(
         simplifyAffineExpr(e, map.getNumDims(), map.getNumSymbols()));
   }
   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;
   assert(map.getNumSymbols() == 0 && "expected map without symbols");
   SmallVector<AffineExpr, 4> exprs(map.getNumDims());
   for (auto en : llvm::enumerate(map.getResults())) {
     auto expr = en.value();
     // Skip non-permutations.

mlir/lib/Transforms/Utils/LoopUtils.cpp

 AffineForOp mlir::createCanonicalizedAffineForOp(
     OpBuilder b, Location loc, ValueRange lbOperands, AffineMap lbMap,
     ValueRange ubOperands, AffineMap ubMap, int64_t step) {
   SmallVector<Value, 4> lowerOperands(lbOperands);
   SmallVector<Value, 4> upperOperands(ubOperands);
 
   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);
 }
 
 /// Creates an AffineIfOp that encodes the conditional to choose between
 /// the constant trip count version and an unknown trip count version of this
 /// nest of loops. This is used to separate partial and full tiles if `loops`

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

func @loop_min_max(%N : index) {
   affine.for %i = 0 to 42 {
     affine.for %j = max #lbMultiMap(%i)[%N] to min #ubMultiMap(%i)[%N] {
       call @body2(%i, %j) : (index, 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.
 
 // CHECK-LABEL: func @min_reduction_tree
 // CHECK-NEXT:   %[[c0:.*]] = constant 0 : index
 // CHECK-NEXT:   %[[c01:.+]] = cmpi "slt", %{{.*}}, %{{.*}} : index
 // CHECK-NEXT:   %[[r01:.+]] = select %[[c01]], %{{.*}}, %{{.*}} : index
 // CHECK-NEXT:   %[[c012:.+]] = cmpi "slt", %[[r01]], %{{.*}} : index
 // CHECK-NEXT:   %[[r012:.+]] = select %[[c012]], %[[r01]], %{{.*}} : index
 // CHECK-NEXT:   %[[c0123:.+]] = cmpi "slt", %[[r012]], %{{.*}} : index
 // CHECK-NEXT:   %[[r0123:.+]] = select %[[c0123]], %[[r012]], %{{.*}} : index
 // CHECK-NEXT:   %[[c01234:.+]] = cmpi "slt", %[[r0123]], %{{.*}} : index
 // CHECK-NEXT:   %[[r01234:.+]] = select %[[c01234]], %[[r0123]], %{{.*}} : index
 // CHECK-NEXT:   %[[c012345:.+]] = cmpi "slt", %[[r01234]], %{{.*}} : index
 // CHECK-NEXT:   %[[r012345:.+]] = select %[[c012345]], %[[r01234]], %{{.*}} : index
 // CHECK-NEXT:   %[[c0123456:.+]] = cmpi "slt", %[[r012345]], %{{.*}} : index
 // CHECK-NEXT:   %[[r0123456:.+]] = select %[[c0123456]], %[[r012345]], %{{.*}} : index
 // CHECK-NEXT:   %[[c1:.*]] = constant 1 : index
 // CHECK-NEXT:   for %{{.*}} = %[[c0]] to %[[v11]] step %[[c1]] {
 // CHECK-NEXT:     call @body(%{{.*}}) : (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
 }
 
 /////////////////////////////////////////////////////////////////////
 
 #map0 = affine_map<() -> (0)>
func @affine_max(%arg0: index, %arg1: index) -> index{
   // CHECK: %[[first:.*]] = addi %[[ARG0]], %[[neg1]]
   // CHECK: %[[Cm2:.*]] = constant -1
   // CHECK: %[[neg2:.*]] = muli %[[ARG0]], %[[Cm2:.*]]
   // CHECK: %[[second:.*]] = addi %[[ARG1]], %[[neg2]]
   // CHECK: %[[cmp:.*]] = cmpi "sgt", %[[first]], %[[second]]
   // 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

func @rep(%arg0 : index, %arg1 : index) -> index {
   %c1024 = constant 1024 : index
   // CHECK-NOT: affine.apply
   %0 = affine.apply #map1(%arg0)[%c1024, %c0]
 
   // CHECK: affine.min #[[MAP]]()[%[[ARG1]], %[[ARG0]]]
   %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
+}