Complete migration to exclusive upper bound

cl/220448963 had missed a part of the updates.

- while on this, clean up some of the test cases to use ops' custom forms.

PiperOrigin-RevId: 220675303
This commit is contained in:
Uday Bondhugula 2018-11-08 11:46:18 -08:00 committed by jpienaar
parent 76bbe2cff6
commit 23ddd577ef
2 changed files with 17 additions and 12 deletions

View File

@ -98,24 +98,29 @@ static bool setTiledIndexSetHyperRect(ArrayRef<ForStmt *> origLoops,
// TODO(bondhugula): Keep it simple for now - constant upper bound.
if (!origLoops[i]->hasConstantUpperBound())
return false;
int64_t largestDiv = getLargestDivisorOfTripCount(*origLoops[i]);
auto mayBeConstantCount = getConstantTripCount(*origLoops[i]);
AffineMap lbMap, ubMap;
auto dim = b.getAffineDimExpr(0);
lbMap = b.getAffineMap(1, 0, dim, {});
newLoops[width + i]->setLowerBound(newLoops[i], lbMap);
// Set the upper bound.
if (mayBeConstantCount.hasValue() &&
mayBeConstantCount.getValue() < tileSizes[i]) {
ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue() - 1);
// Trip count is less than tile size; upper bound is the trip count.
ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue());
newLoops[width + i]->setUpperBoundMap(ubMap);
} else if (largestDiv % tileSizes[i] == 0) {
// No need of min.
ubMap = b.getAffineMap(1, 0, dim + tileSizes[i] - 1, {});
newLoops[width + i]->setUpperBound(newLoops[i], ubMap);
} else {
} else if (largestDiv % tileSizes[i] != 0) {
// Intra-tile loop ii goes from i to min(i + tileSize, ub_i).
auto ubMax =
b.getAffineConstantExpr(origLoops[i]->getConstantUpperBound());
ubMap = b.getAffineMap(1, 0, {dim + tileSizes[i] - 1, ubMax}, {});
ubMap = b.getAffineMap(1, 0, {dim + tileSizes[i], ubMax}, {});
newLoops[width + i]->setUpperBound(newLoops[i], ubMap);
} else {
// No need of the min expression.
ubMap = b.getAffineMap(1, 0, dim + tileSizes[i], {});
newLoops[width + i]->setUpperBound(newLoops[i], ubMap);
}
}

View File

@ -1,7 +1,7 @@
// RUN: mlir-opt %s -loop-tile | FileCheck %s
// CHECK: #map0 = (d0) -> (d0 + 31)
// CHECK: #map1 = (d0) -> (d0 + 31, 51)
// CHECK: #map0 = (d0) -> (d0 + 32)
// CHECK: #map1 = (d0) -> (d0 + 32, 50)
// CHECK-LABEL: mlfunc @loop_tiling()
// CHECK-NEXT: for %i0 = 0 to 256 step 32 {
// CHECK-NEXT: for %i1 = 0 to 512 step 32 {
@ -16,13 +16,13 @@
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: for %i6 = 0 to 51 step 32 {
// CHECK-NEXT: for %i6 = 0 to 50 step 32 {
// CHECK-NEXT: for %i7 = (d0) -> (d0)(%i6) to min #map1(%i6) {
// CHECK-NEXT: "bar"(%i7, %i7) : (index, index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: for %i8 = 0 to 21 step 32 {
// CHECK-NEXT: for %i9 = (d0) -> (d0)(%i8) to 20 {
// CHECK-NEXT: for %i9 = (d0) -> (d0)(%i8) to 21 {
// CHECK-NEXT: "foobar"(%i9) : (index) -> ()
// CHECK-NEXT: }
// CHECK-NEXT: }
@ -36,7 +36,7 @@ mlfunc @loop_tiling() {
}
}
for %x = 0 to 51 {
for %x = 0 to 50 {
"bar"(%x, %x) : (index, index) -> ()
}