[MLIR] Fix affine loop tiling utility upper bound bug

[bondhugula]

Summary

Fix intra-tile upper bound setting in a scenario where the tile size was
larger than the trip count.

[bondhugula]

fix comment + cleanup test case

[bondhugula]

Drop white space in test case

[bondhugula]

Use an alternative equivalent method that is better.

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

for (unsigned i = 0; i < width; i++) {
     auto mayBeConstantCount = getConstantTripCount(origLoops[i]);
     // The lower bound is just the tile-space loop.
     AffineMap lbMap = b.getDimIdentityMap();
     newLoops[width + i].setLowerBound(
         /*operands=*/newLoops[i].getInductionVar(), lbMap);
 
     // Set the upper bound.
     if (mayBeConstantCount && mayBeConstantCount.getValue() < tileSizes[i]) {
-      // Trip count is less than tile size; upper bound is the trip count.
-      auto ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue());
-      newLoops[width + i].setUpperBoundMap(ubMap);
+      // Trip count is less than the tile size; upper bound is the tile space
+      // loop's upper bound.
+      newLoops[width + i].setUpperBound(
+          /*operands=*/newLoops[i].getUpperBoundOperands(),
+          newLoops[i].getUpperBoundMap());
     } else if (largestDiv % tileSizes[i] != 0) {
       // Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
       // Construct the upper bound map; the operands are the original operands
       // with 'i' (tile-space loop) appended to it. The new upper bound map is
       // the original one with an additional expression i + tileSize appended.
 
       // Add dim operands from original upper bound.
       SmallVector<Value, 4> ubOperands;

mlir/test/Dialect/Affine/loop-tiling.mlir

 // CHECK-NEXT:  affine.for %{{.*}} = 0 to [[MAP1]]()[%{{.*}}, %{{.*}}] step 32 {
 // CHECK-NEXT:    affine.for %{{.*}} = [[MAP0]](%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}, %{{.*}}] {
 // CHECK-NEXT:      %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<?xf32>
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 
 // -----
 
+func @tile_size_larger_than_trip_count_symbolic_bound(%M: index, %N :  index) {
+  affine.for %i = affine_map<(d0) -> (d0)>(%M) to affine_map<(d0) -> (d0 + 2)>(%M) {
+    affine.for %j = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 + 4)>(%N) {
+      "test.foo" () : () -> ()
+    }
+  }
+  return
+}
+
+// CHECK-DAG: #[[ID:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[ID_PLUS_2:.*]] = affine_map<(d0) -> (d0 + 2)>
+// CHECK-DAG: #[[ID_PLUS_4:.*]] = affine_map<(d0) -> (d0 + 4)>
+
+// CHECK: %[[M:.*]]: index, %[[N:.*]]: index
+
+// CHECK:      affine.for %[[I:.*]] = #[[ID]](%[[M]]) to #[[ID_PLUS_2]](%[[M]]) step 32 {
+// CHECK-NEXT:   affine.for %[[J:.*]] = #[[ID]](%[[N]]) to #[[ID_PLUS_4]](%[[N]]) step 32 {
+// CHECK-NEXT:     affine.for %arg4 = #[[ID]](%[[I]]) to #[[ID_PLUS_2]](%[[M]]) {
+// CHECK-NEXT:       affine.for %arg5 = #[[ID]](%[[J]]) to #[[ID_PLUS_4]](%[[N]]) {
+// CHECK-NEXT:         "test.foo"() : () -> ()
+
+// -----
+
 // CHECK-LABEL: func @trip_count_1
 // SEPARATE-LABEL: func @trip_count_1
 func @trip_count_1(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)
     -> memref<196608x1xf32> {
   affine.for %i1 = 0 to 196608 {
     affine.for %i3 = 0 to 1 {
       %4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32>
       affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32>