[PATCH] D78505: [MLIR] Fix affine loop tiling utility upper bound bug

[Uday Bondhugula via Phabricator via llvm-commits]

bondhugula updated this revision to Diff 258781.
bondhugula added a comment.

fix comment + cleanup test case


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D78505/new/

https://reviews.llvm.org/D78505

Files:
  mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
  mlir/test/Dialect/Affine/loop-tiling.mlir


Index: mlir/test/Dialect/Affine/loop-tiling.mlir
===================================================================
--- mlir/test/Dialect/Affine/loop-tiling.mlir
+++ mlir/test/Dialect/Affine/loop-tiling.mlir
@@ -170,6 +170,29 @@
 
 // -----
 
+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>)
Index: mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
===================================================================
--- mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
@@ -110,9 +110,11 @@
 
     // 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


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D78505.258781.patch
Type: text/x-patch
Size: 2378 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200420/a8eb38a6/attachment-0001.bin>