2d6478fa92
are a max/min of several expressions. - Extend loop tiling to handle non-constant loop bounds and bounds that are a max/min of several expressions, i.e., bounds using multi-result affine maps - also fix b/120630124 as a result (the IR was in an invalid state when tiled loop generation failed; SSA uses were created that weren't plugged into the IR). PiperOrigin-RevId: 224604460
71 lines
2.4 KiB
MLIR
71 lines
2.4 KiB
MLIR
// RUN: mlir-opt %s -loop-tile -tile-size=32 | FileCheck %s
|
|
|
|
// CHECK-DAG: #map0 = (d0) -> (d0 + 32)
|
|
// CHECK-DAG: #map1 = (d0) -> (d0 + 32, 50)
|
|
// CHECK-DAG: [[IDENTITY:#map[0-9]+]] = (d0) -> (d0)
|
|
// CHECK-DAG: [[LB:#map[0-9]+]] = ()[s0] -> (0, s0)
|
|
// CHECK-DAG: [[UB:#map[0-9]+]] = ()[s0, s1] -> (s0, 4096 floordiv s1)
|
|
// CHECK-DAG: [[UB_INTRA_TILE:#map[0-9]+]] = (d0, d1, d2) -> (d2 + 32, s0, 4096 floordiv s1)
|
|
|
|
// CHECK-LABEL: mlfunc @loop_tiling()
|
|
// CHECK-NEXT: for %i0 = 0 to 256 step 32 {
|
|
// CHECK-NEXT: for %i1 = 0 to 512 step 32 {
|
|
// CHECK-NEXT: for %i2 = 0 to 1024 step 32 {
|
|
// CHECK-NEXT: for %i3 = [[IDENTITY]](%i0) to #map0(%i0) {
|
|
// CHECK-NEXT: for %i4 = [[IDENTITY]](%i1) to #map0(%i1) {
|
|
// CHECK-NEXT: for %i5 = [[IDENTITY]](%i2) to #map0(%i2) {
|
|
// CHECK-NEXT: "foo"(%i3, %i4, %i5) : (index, index, index) -> ()
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: for %i6 = 0 to 50 step 32 {
|
|
// CHECK-NEXT: for %i7 = [[IDENTITY]](%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 = [[IDENTITY]](%i8) to 21 {
|
|
// CHECK-NEXT: "foobar"(%i9) : (index) -> ()
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: return
|
|
mlfunc @loop_tiling() {
|
|
for %i = 0 to 256 {
|
|
for %j = 0 to 512 {
|
|
for %k = 0 to 1024 {
|
|
"foo"(%i, %j, %k) : (index, index, index) -> ()
|
|
}
|
|
}
|
|
}
|
|
|
|
for %x = 0 to 50 {
|
|
"bar"(%x, %x) : (index, index) -> ()
|
|
}
|
|
|
|
// Intra-tile loop won't need a min expression.
|
|
for %y = 0 to 21 {
|
|
"foobar"(%y) : (index) -> ()
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
#lb = ()[s0] -> (0, s0)
|
|
#ub = ()[s0, s1] -> (s0, 4096 floordiv s1)
|
|
// CHECK-LABEL: mlfunc @loop_max_min_bound(%arg0 : memref<?xi32>, %arg1 : index, %arg2 : index) {
|
|
mlfunc @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) {
|
|
%M = dim %A, 0 : memref<? x i32>
|
|
for %iTT = max #lb()[%L] to min #ub()[%M, %U] {
|
|
%out = affine_apply (d0) -> (d0) (%iTT)
|
|
}
|
|
return
|
|
// CHECK: for %i0 = max [[LB]]()[%arg1] to min [[UB]]()[%0, %arg2] step 32 {
|
|
// CHECK-NEXT: for %i1 = [[IDENTITY]](%i0) to min [[UB_INTRA_TILE]](%0, %arg2, %i0) {
|
|
// CHECK-NEXT: %1 = affine_apply [[IDENTITY]](%i1)
|
|
// CHECK-NEXT: }
|
|
// CHECK-NEXT: }
|
|
}
|