maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
Uday Bondhugula 2020-04-20 22:29:21 +05:30
parent 85cca945b4
commit 3dff8c9109
2 changed files with 34 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
View File

@ -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

36
mlir/test/Dialect/Affine/loop-tiling.mlir
View File

@ -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>)