285 lines · plain
1// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(func.func(test-affine-loop-unswitch))" | FileCheck %s2 3// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>4 5// CHECK-LABEL: func @if_else_imperfect6func.func @if_else_imperfect(%A : memref<100xi32>, %B : memref<100xi32>, %v : i32) {7// CHECK: %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref8 affine.for %i = 0 to 100 {9 affine.store %v, %A[%i] : memref<100xi32>10 affine.for %j = 0 to 100 {11 affine.store %v, %A[%j] : memref<100xi32>12 affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {13 affine.store %v, %B[%j] : memref<100xi32>14 }15 func.call @external() : () -> ()16 }17 affine.store %v, %A[%i] : memref<100xi32>18 }19 return20}21func.func private @external()22 23// CHECK: affine.for %[[I:.*]] = 0 to 100 {24// CHECK-NEXT: affine.store %{{.*}}, %[[A]][%[[I]]]25// CHECK-NEXT: affine.if #[[$SET]](%[[I]]) {26// CHECK-NEXT: affine.for %[[J:.*]] = 0 to 100 {27// CHECK-NEXT: affine.store %{{.*}}, %[[A]][%[[J]]]28// CHECK-NEXT: affine.store %{{.*}}, %[[B]][%[[J]]]29// CHECK-NEXT: call30// CHECK-NEXT: }31// CHECK-NEXT: } else {32// CHECK-NEXT: affine.for %[[JJ:.*]] = 0 to 100 {33// CHECK-NEXT: affine.store %{{.*}}, %[[A]][%[[J]]]34// CHECK-NEXT: call35// CHECK-NEXT: }36// CHECK-NEXT: }37// CHECK-NEXT: affine.store %{{.*}}, %[[A]][%[[I]]]38// CHECK-NEXT: }39// CHECK-NEXT: return40 41// -----42 43func.func private @foo()44func.func private @bar()45func.func private @abc()46func.func private @xyz()47 48// CHECK-LABEL: func @if_then_perfect49func.func @if_then_perfect(%A : memref<100xi32>, %v : i32) {50 affine.for %i = 0 to 100 {51 affine.for %j = 0 to 100 {52 affine.for %k = 0 to 100 {53 affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {54 affine.store %v, %A[%i] : memref<100xi32>55 }56 }57 }58 }59 return60}61// CHECK: affine.for62// CHECK-NEXT: affine.if63// CHECK-NEXT: affine.for64// CHECK-NEXT: affine.for65// CHECK-NOT: else66 67 68// CHECK-LABEL: func @if_else_perfect69func.func @if_else_perfect(%A : memref<100xi32>, %v : i32) {70 affine.for %i = 0 to 99 {71 affine.for %j = 0 to 100 {72 affine.for %k = 0 to 100 {73 func.call @foo() : () -> ()74 affine.if affine_set<(d0, d1) : (d0 - 2 >= 0, -d1 + 80 >= 0)>(%i, %j) {75 affine.store %v, %A[%i] : memref<100xi32>76 func.call @abc() : () -> ()77 } else {78 affine.store %v, %A[%i + 1] : memref<100xi32>79 func.call @xyz() : () -> ()80 }81 func.call @bar() : () -> ()82 }83 }84 }85 return86}87// CHECK: affine.for88// CHECK-NEXT: affine.for89// CHECK-NEXT: affine.if90// CHECK-NEXT: affine.for91// CHECK-NEXT: call @foo92// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}]93// CHECK-NEXT: call @abc94// CHECK-NEXT: call @bar95// CHECK-NEXT: }96// CHECK-NEXT: else97// CHECK-NEXT: affine.for98// CHECK-NEXT: call @foo99// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} + 1]100// CHECK-NEXT: call @xyz101// CHECK-NEXT: call @bar102// CHECK-NEXT: }103// CHECK-NEXT: }104// CHECK-NEXT: }105// CHECK-NEXT: }106 107// CHECK-LABEL: func @if_then_imperfect108func.func @if_then_imperfect(%A : memref<100xi32>, %N : index, %v: i32) {109 affine.for %i = 0 to 100 {110 affine.store %v, %A[0] : memref<100xi32>111 affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%N) {112 affine.store %v, %A[%i] : memref<100xi32>113 }114 }115 return116}117// CHECK: affine.if118// CHECK-NEXT: affine.for119// CHECK-NEXT: affine.store120// CHECK-NEXT: affine.store121// CHECK-NEXT: }122// CHECK-NEXT: } else {123// CHECK-NEXT: affine.for124// CHECK-NEXT: affine.store125// CHECK-NEXT: }126// CHECK-NEXT: }127// CHECK-NEXT: return128 129// Check if unused operands are dropped: hence, hoisting is possible.130// CHECK-LABEL: func @hoist_after_canonicalize131func.func @hoist_after_canonicalize() {132 affine.for %i = 0 to 100 {133 affine.for %j = 0 to 100 {134 affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%j) {135 affine.if affine_set<(d0, d1) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%i, %j) {136 // The call to external is to avoid DCE on affine.if.137 func.call @foo() : () -> ()138 }139 }140 }141 }142 return143}144// CHECK: affine.for145// CHECK-NEXT: affine.if146// CHECK-NEXT: affine.for147// CHECK-NEXT: affine.if148// CHECK-NEXT: call149// CHECK-NEXT: }150// CHECK-NEXT: }151// CHECK-NEXT: }152// CHECK-NEXT: }153// CHECK-NEXT: return154 155// CHECK-LABEL: func @handle_dead_if156func.func @handle_dead_if(%N : index) {157 affine.for %i = 0 to 100 {158 affine.if affine_set<(d0) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%N) {159 }160 }161 return162}163// CHECK-NEXT: affine.for164// CHECK-NEXT: }165// CHECK-NEXT: return166 167// -----168 169// A test case with affine.parallel.170 171#flb1 = affine_map<(d0) -> (d0 * 3)>172#fub1 = affine_map<(d0) -> (d0 * 3 + 3)>173#flb0 = affine_map<(d0) -> (d0 * 16)>174#fub0 = affine_map<(d0) -> (d0 * 16 + 16)>175#pub1 = affine_map<(d0)[s0] -> (s0, d0 * 3 + 3)>176#pub0 = affine_map<(d0)[s0] -> (s0, d0 * 16 + 16)>177#lb1 = affine_map<(d0) -> (d0 * 480)>178#ub1 = affine_map<(d0)[s0] -> (s0, d0 * 480 + 480)>179#lb0 = affine_map<(d0) -> (d0 * 110)>180#ub0 = affine_map<(d0)[s0] -> (d0 * 110 + 110, s0 floordiv 3)>181 182#set0 = affine_set<(d0, d1)[s0, s1] : (d0 * -16 + s0 - 16 >= 0, d1 * -3 + s1 - 3 >= 0)>183 184// CHECK-LABEL: func @perfect_if_else185func.func @perfect_if_else(%arg0 : memref<?x?xf64>, %arg1 : memref<?x?xf64>, %v : f64,186 %arg4 : index, %arg5 : index, %arg6 : index, %sym : index) {187 affine.for %arg7 = #lb0(%arg5) to min #ub0(%arg5)[%sym] {188 affine.parallel (%i0, %j0) = (0, 0) to (symbol(%sym), 100) step (10, 10) {189 affine.for %arg8 = #lb1(%arg4) to min #ub1(%arg4)[%sym] {190 affine.if #set0(%arg6, %arg7)[%sym, %sym] {191 affine.for %arg9 = #flb0(%arg6) to #fub0(%arg6) {192 affine.for %arg10 = #flb1(%arg7) to #fub1(%arg7) {193 affine.store %v, %arg0[0, 0] : memref<?x?xf64>194 }195 }196 } else {197 affine.for %arg9 = #lb0(%arg6) to min #pub0(%arg6)[%sym] {198 affine.for %arg10 = #lb1(%arg7) to min #pub1(%arg7)[%sym] {199 affine.store %v, %arg0[0, 0] : memref<?x?xf64>200 }201 }202 }203 }204 }205 }206 return207}208 209// CHECK: affine.for210// CHECK-NEXT: affine.if211// CHECK-NEXT: affine.parallel212// CHECK-NEXT: affine.for213// CHECK-NEXT: affine.for214// CHECK-NEXT: affine.for215// CHECK-NEXT: affine.store216// CHECK-NEXT: }217// CHECK-NEXT: }218// CHECK-NEXT: }219// CHECK-NEXT: }220// CHECK-NEXT: } else {221// CHECK-NEXT: affine.parallel222// CHECK-NEXT: affine.for223// CHECK-NEXT: affine.for224// CHECK-NEXT: affine.for225// CHECK-NEXT: affine.store226// CHECK-NEXT: }227// CHECK-NEXT: }228// CHECK-NEXT: }229// CHECK-NEXT: }230// CHECK-NEXT: }231// CHECK-NEXT: }232 233// With multiple if ops in a function, the test pass just looks for the first if234// op that it is able to successfully hoist.235 236// CHECK-LABEL: func @multiple_if237func.func @multiple_if(%N : index) {238 affine.if affine_set<() : (0 == 0)>() {239 func.call @external() : () -> ()240 }241 affine.for %i = 0 to 100 {242 affine.if affine_set<()[s0] : (s0 >= 0)>()[%N] {243 func.call @external() : () -> ()244 }245 }246 return247}248// CHECK: call249// CHECK-NEXT: affine.if250// CHECK-NEXT: affine.for251// CHECK-NEXT: call252// CHECK-NEXT: }253// CHECK-NEXT: }254// CHECK-NEXT: return255 256func.func private @external()257 258// Check to ensure affine.parallel ops are handled as well.259 260#set = affine_set<(d0) : (-d0 + 3 >= 0)>261// CHECK-LABEL: affine_parallel262func.func @affine_parallel(%arg0: memref<35xf32>) -> memref<35xf32> {263 %0 = llvm.mlir.constant(1.000000e+00 : f32) : f32264 %alloc = memref.alloc() {alignment = 64 : i64} : memref<35xf32>265 // CHECK: affine.parallel266 affine.parallel (%arg1) = (0) to (35) step (32) {267 // This can't be hoisted further.268 // CHECK-NEXT: affine.if269 affine.if #set(%arg1) {270 affine.parallel (%arg2) = (%arg1) to (%arg1 + 32) {271 %1 = affine.load %arg0[%arg2] : memref<35xf32>272 %2 = llvm.fdiv %0, %1 : f32273 affine.store %2, %alloc[%arg2] : memref<35xf32>274 }275 } else {276 affine.parallel (%arg2) = (%arg1) to (min(%arg1 + 32, 35)) {277 %1 = affine.load %arg0[%arg2] : memref<35xf32>278 %2 = llvm.fdiv %0, %1 : f32279 affine.store %2, %alloc[%arg2] : memref<35xf32>280 }281 }282 }283 return %alloc : memref<35xf32>284}285