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

[Uday Bondhugula via Phabricator via llvm-commits]

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

Use an alternative equivalent method that is better.


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
@@ -6,15 +6,16 @@
 
 // CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
 // CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32, 50)>
-// CHECK-DAG: [[IDENTITY:#map[0-9]+]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: [[ID:#map[0-9]+]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: [[ID_PLUS_21:#map[0-9]+]] = affine_map<(d0) -> (d0 + 21)>
 
 // CHECK-LABEL: func @loop_tiling()
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 256 step 32 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 512 step 32 {
 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 step 32 {
-// CHECK-NEXT:         affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
-// CHECK-NEXT:           affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
-// CHECK-NEXT:             affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to [[MAP0]](%{{.*}}) {
+// CHECK-NEXT:         affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
+// CHECK-NEXT:           affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
+// CHECK-NEXT:             affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
 // CHECK-NEXT:               "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> ()
 // CHECK-NEXT:             }
 // CHECK-NEXT:           }
@@ -23,12 +24,12 @@
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 50 step 32 {
-// CHECK-NEXT:     affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to min [[MAP1]](%{{.*}}) {
+// CHECK-NEXT:     affine.for %{{.*}} = [[ID]](%{{.*}}) to min [[MAP1]](%{{.*}}) {
 // CHECK-NEXT:       "bar"(%{{.*}}, %{{.*}}) : (index, index) -> ()
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
-// CHECK-NEXT: affine.for %{{.*}} = 0 to 21 step 32 {
-// CHECK-NEXT:    affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to 21 {
+// CHECK-NEXT: affine.for %[[I:.*]] = 0 to 21 step 32 {
+// CHECK-NEXT:    affine.for %{{.*}} = [[ID]](%[[I]]) to [[ID_PLUS_21]](%[[I]]) {
 // CHECK-NEXT:      "foobar"(%{{.*}}) : (index) -> ()
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
@@ -170,6 +171,27 @@
 
 // -----
 
+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]](%[[I]])
+// CHECK-NEXT:       affine.for %arg5 = #[[ID]](%[[J]]) to #[[ID_PLUS_4]](%[[J]])
+// 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 lower bound +
+      // trip count.
+      auto ubMap = b.getSingleDimShiftAffineMap(mayBeConstantCount.getValue());
+      newLoops[width + i].setUpperBound(
+          /*operands=*/newLoops[i].getInductionVar(), ubMap);
     } 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.258786.patch
Type: text/x-patch
Size: 4409 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200420/5e959f61/attachment.bin>