[Mlir-commits] [mlir] 3dff8c9 - [MLIR] Fix affine loop tiling utility upper bound bug

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2020-04-21T00:54:01+05:30
New Revision: 3dff8c9109a73b208ef7625c06e0d507ffb02614

URL: https://github.com/llvm/llvm-project/commit/3dff8c9109a73b208ef7625c06e0d507ffb02614
DIFF: https://github.com/llvm/llvm-project/commit/3dff8c9109a73b208ef7625c06e0d507ffb02614.diff

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

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

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
index ce692e405ae1..c411a6a548ff 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
@@ -110,9 +110,11 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
 
     // 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

diff  --git a/mlir/test/Dialect/Affine/loop-tiling.mlir b/mlir/test/Dialect/Affine/loop-tiling.mlir
index 075ab0c230ac..380e82c3a326 100644
--- a/mlir/test/Dialect/Affine/loop-tiling.mlir
+++ b/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_with_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: i
 
 // -----
 
+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>)