diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp --- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp +++ b/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 diff --git a/mlir/test/Dialect/Affine/loop-tiling.mlir b/mlir/test/Dialect/Affine/loop-tiling.mlir --- 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_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>)