444 lines · plain
1// RUN: mlir-opt %s | FileCheck %s2// Verify the printed output can be parsed.3// RUN: mlir-opt %s | mlir-opt | FileCheck %s4// Verify the generic form can be parsed.5// RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s6 7func.func @std_for(%arg0 : index, %arg1 : index, %arg2 : index) {8 scf.for %i0 = %arg0 to %arg1 step %arg2 {9 scf.for %i1 = %arg0 to %arg1 step %arg2 {10 %min_cmp = arith.cmpi slt, %i0, %i1 : index11 %min = arith.select %min_cmp, %i0, %i1 : index12 %max_cmp = arith.cmpi sge, %i0, %i1 : index13 %max = arith.select %max_cmp, %i0, %i1 : index14 scf.for %i2 = %min to %max step %i1 {15 }16 }17 }18 return19}20// CHECK-LABEL: func @std_for(21// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {22// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {23// CHECK-NEXT: %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}} : index24// CHECK-NEXT: %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : index25// CHECK-NEXT: %{{.*}} = arith.cmpi sge, %{{.*}}, %{{.*}} : index26// CHECK-NEXT: %{{.*}} = arith.select %{{.*}}, %{{.*}}, %{{.*}} : index27// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {28 29func.func @std_for_i32(%arg0 : i32, %arg1 : i32, %arg2 : i32) {30 scf.for %i0 = %arg0 to %arg1 step %arg2 : i32 {31 scf.for unsigned %i1 = %arg0 to %arg1 step %arg2 : i32 {32 }33 }34 return35}36// CHECK-LABEL: func @std_for_i32(37// CHECK-NEXT: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i32 {38// CHECK-NEXT: scf.for unsigned %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i32 {39 40func.func @scf_for_i64_iter(%arg1: i64, %arg2: i64) {41 %c1_i64 = arith.constant 1 : i6442 %c0_i64 = arith.constant 0 : i6443 %0 = scf.for %arg3 = %arg1 to %arg2 step %c1_i64 iter_args(%arg4 = %c0_i64) -> (i64) : i64 {44 %1 = arith.addi %arg4, %arg3 : i6445 scf.yield %1 : i6446 }47 return48}49// CHECK-LABEL: scf_for_i64_iter50// CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}} -> (i64) : i64 {51 52func.func @std_if(%arg0: i1, %arg1: f32) {53 scf.if %arg0 {54 %0 = arith.addf %arg1, %arg1 : f3255 }56 return57}58// CHECK-LABEL: func @std_if(59// CHECK-NEXT: scf.if %{{.*}} {60// CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f3261 62func.func @std_if_else(%arg0: i1, %arg1: f32) {63 scf.if %arg0 {64 %0 = arith.addf %arg1, %arg1 : f3265 } else {66 %1 = arith.addf %arg1, %arg1 : f3267 }68 return69}70// CHECK-LABEL: func @std_if_else(71// CHECK-NEXT: scf.if %{{.*}} {72// CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f3273// CHECK-NEXT: } else {74// CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f3275 76func.func @std_parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,77 %arg3 : index, %arg4 : index) {78 %step = arith.constant 1 : index79 scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)80 step (%arg4, %step) {81 %min_cmp = arith.cmpi slt, %i0, %i1 : index82 %min = arith.select %min_cmp, %i0, %i1 : index83 %max_cmp = arith.cmpi sge, %i0, %i1 : index84 %max = arith.select %max_cmp, %i0, %i1 : index85 %zero = arith.constant 0.0 : f3286 %int_zero = arith.constant 0 : i3287 %red:2 = scf.parallel (%i2) = (%min) to (%max) step (%i1)88 init (%zero, %int_zero) -> (f32, i32) {89 %one = arith.constant 1.0 : f3290 %int_one = arith.constant 1 : i3291 scf.reduce(%one, %int_one : f32, i32) {92 ^bb0(%lhs : f32, %rhs: f32):93 %res = arith.addf %lhs, %rhs : f3294 scf.reduce.return %res : f3295 }, {96 ^bb0(%lhs : i32, %rhs: i32):97 %res = arith.muli %lhs, %rhs : i3298 scf.reduce.return %res : i3299 }100 }101 scf.reduce102 }103 return104}105// CHECK-LABEL: func @std_parallel_loop(106// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:107// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:108// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:109// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]:110// CHECK-SAME: %[[ARG4:[A-Za-z0-9]+]]:111// CHECK: %[[STEP:.*]] = arith.constant 1 : index112// CHECK-NEXT: scf.parallel (%[[I0:.*]], %[[I1:.*]]) = (%[[ARG0]], %[[ARG1]]) to113// CHECK: (%[[ARG2]], %[[ARG3]]) step (%[[ARG4]], %[[STEP]]) {114// CHECK-NEXT: %[[MIN_CMP:.*]] = arith.cmpi slt, %[[I0]], %[[I1]] : index115// CHECK-NEXT: %[[MIN:.*]] = arith.select %[[MIN_CMP]], %[[I0]], %[[I1]] : index116// CHECK-NEXT: %[[MAX_CMP:.*]] = arith.cmpi sge, %[[I0]], %[[I1]] : index117// CHECK-NEXT: %[[MAX:.*]] = arith.select %[[MAX_CMP]], %[[I0]], %[[I1]] : index118// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32119// CHECK-NEXT: %[[INT_ZERO:.*]] = arith.constant 0 : i32120// CHECK-NEXT: scf.parallel (%{{.*}}) = (%[[MIN]]) to (%[[MAX]])121// CHECK-SAME: step (%[[I1]])122// CHECK-SAME: init (%[[ZERO]], %[[INT_ZERO]]) -> (f32, i32) {123// CHECK-NEXT: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32124// CHECK-NEXT: %[[INT_ONE:.*]] = arith.constant 1 : i32125// CHECK-NEXT: scf.reduce(%[[ONE]], %[[INT_ONE]] : f32, i32) {126// CHECK-NEXT: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):127// CHECK-NEXT: %[[RES:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32128// CHECK-NEXT: scf.reduce.return %[[RES]] : f32129// CHECK-NEXT: }, {130// CHECK-NEXT: ^bb0(%[[LHS:.*]]: i32, %[[RHS:.*]]: i32):131// CHECK-NEXT: %[[RES:.*]] = arith.muli %[[LHS]], %[[RHS]] : i32132// CHECK-NEXT: scf.reduce.return %[[RES]] : i32133// CHECK-NEXT: }134// CHECK-NEXT: }135// CHECK-NEXT: scf.reduce136 137func.func @parallel_explicit_yield(138 %arg0: index, %arg1: index, %arg2: index) {139 scf.parallel (%i0) = (%arg0) to (%arg1) step (%arg2) {140 scf.reduce141 }142 return143}144 145// CHECK-LABEL: func @parallel_explicit_yield(146// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:147// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:148// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:149// CHECK-NEXT: scf.parallel (%{{.*}}) = (%[[ARG0]]) to (%[[ARG1]]) step (%[[ARG2]])150// CHECK-NEXT: scf.reduce151// CHECK-NEXT: }152// CHECK-NEXT: return153// CHECK-NEXT: }154 155func.func @std_if_yield(%arg0: i1, %arg1: f32)156{157 %x, %y = scf.if %arg0 -> (f32, f32) {158 %0 = arith.addf %arg1, %arg1 : f32159 %1 = arith.subf %arg1, %arg1 : f32160 scf.yield %0, %1 : f32, f32161 } else {162 %0 = arith.subf %arg1, %arg1 : f32163 %1 = arith.addf %arg1, %arg1 : f32164 scf.yield %0, %1 : f32, f32165 }166 return167}168// CHECK-LABEL: func @std_if_yield(169// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:170// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:171// CHECK-NEXT: %{{.*}}:2 = scf.if %[[ARG0]] -> (f32, f32) {172// CHECK-NEXT: %[[T1:.*]] = arith.addf %[[ARG1]], %[[ARG1]]173// CHECK-NEXT: %[[T2:.*]] = arith.subf %[[ARG1]], %[[ARG1]]174// CHECK-NEXT: scf.yield %[[T1]], %[[T2]] : f32, f32175// CHECK-NEXT: } else {176// CHECK-NEXT: %[[T3:.*]] = arith.subf %[[ARG1]], %[[ARG1]]177// CHECK-NEXT: %[[T4:.*]] = arith.addf %[[ARG1]], %[[ARG1]]178// CHECK-NEXT: scf.yield %[[T3]], %[[T4]] : f32, f32179// CHECK-NEXT: }180 181func.func @std_for_yield(%arg0 : index, %arg1 : index, %arg2 : index) {182 %s0 = arith.constant 0.0 : f32183 %result = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0) -> (f32) {184 %sn = arith.addf %si, %si : f32185 scf.yield %sn : f32186 }187 return188}189// CHECK-LABEL: func @std_for_yield(190// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:191// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:192// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:193// CHECK-NEXT: %[[INIT:.*]] = arith.constant194// CHECK-NEXT: %{{.*}} = scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]195// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]]) -> (f32) {196// CHECK-NEXT: %[[NEXT:.*]] = arith.addf %[[ITER]], %[[ITER]] : f32197// CHECK-NEXT: scf.yield %[[NEXT]] : f32198// CHECK-NEXT: }199 200 201func.func @std_for_yield_multi(%arg0 : index, %arg1 : index, %arg2 : index) {202 %s0 = arith.constant 0.0 : f32203 %t0 = arith.constant 1 : i32204 %u0 = arith.constant 1.0 : f32205 %result1:3 = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0, %ti = %t0, %ui = %u0) -> (f32, i32, f32) {206 %sn = arith.addf %si, %si : f32207 %tn = arith.addi %ti, %ti : i32208 %un = arith.subf %ui, %ui : f32209 scf.yield %sn, %tn, %un : f32, i32, f32210 }211 return212}213// CHECK-LABEL: func @std_for_yield_multi(214// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]:215// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]:216// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]:217// CHECK-NEXT: %[[INIT1:.*]] = arith.constant218// CHECK-NEXT: %[[INIT2:.*]] = arith.constant219// CHECK-NEXT: %[[INIT3:.*]] = arith.constant220// CHECK-NEXT: %{{.*}}:3 = scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]221// CHECK-SAME: iter_args(%[[ITER1:.*]] = %[[INIT1]], %[[ITER2:.*]] = %[[INIT2]], %[[ITER3:.*]] = %[[INIT3]]) -> (f32, i32, f32) {222// CHECK-NEXT: %[[NEXT1:.*]] = arith.addf %[[ITER1]], %[[ITER1]] : f32223// CHECK-NEXT: %[[NEXT2:.*]] = arith.addi %[[ITER2]], %[[ITER2]] : i32224// CHECK-NEXT: %[[NEXT3:.*]] = arith.subf %[[ITER3]], %[[ITER3]] : f32225// CHECK-NEXT: scf.yield %[[NEXT1]], %[[NEXT2]], %[[NEXT3]] : f32, i32, f32226 227 228func.func @conditional_reduce(%buffer: memref<1024xf32>, %lb: index, %ub: index, %step: index) -> (f32) {229 %sum_0 = arith.constant 0.0 : f32230 %c0 = arith.constant 0.0 : f32231 %sum = scf.for %iv = %lb to %ub step %step iter_args(%sum_iter = %sum_0) -> (f32) {232 %t = memref.load %buffer[%iv] : memref<1024xf32>233 %cond = arith.cmpf ugt, %t, %c0 : f32234 %sum_next = scf.if %cond -> (f32) {235 %new_sum = arith.addf %sum_iter, %t : f32236 scf.yield %new_sum : f32237 } else {238 scf.yield %sum_iter : f32239 }240 scf.yield %sum_next : f32241 }242 return %sum : f32243}244// CHECK-LABEL: func @conditional_reduce(245// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]246// CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]247// CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]248// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]249// CHECK-NEXT: %[[INIT:.*]] = arith.constant250// CHECK-NEXT: %[[ZERO:.*]] = arith.constant251// CHECK-NEXT: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[ARG1]] to %[[ARG2]] step %[[ARG3]]252// CHECK-SAME: iter_args(%[[ITER:.*]] = %[[INIT]]) -> (f32) {253// CHECK-NEXT: %[[T:.*]] = memref.load %[[ARG0]][%[[IV]]]254// CHECK-NEXT: %[[COND:.*]] = arith.cmpf ugt, %[[T]], %[[ZERO]]255// CHECK-NEXT: %[[IFRES:.*]] = scf.if %[[COND]] -> (f32) {256// CHECK-NEXT: %[[THENRES:.*]] = arith.addf %[[ITER]], %[[T]]257// CHECK-NEXT: scf.yield %[[THENRES]] : f32258// CHECK-NEXT: } else {259// CHECK-NEXT: scf.yield %[[ITER]] : f32260// CHECK-NEXT: }261// CHECK-NEXT: scf.yield %[[IFRES]] : f32262// CHECK-NEXT: }263// CHECK-NEXT: return %[[RESULT]]264 265// CHECK-LABEL: @while266func.func @while() {267 %0 = "test.get_some_value"() : () -> i32268 %1 = "test.get_some_value"() : () -> f32269 270 // CHECK: = scf.while (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) : (i32, f32) -> (i64, f64) {271 %2:2 = scf.while (%arg0 = %0, %arg1 = %1) : (i32, f32) -> (i64, f64) {272 %3:2 = "test.some_operation"(%arg0, %arg1) : (i32, f32) -> (i64, f64)273 %4 = "test.some_condition"(%arg0, %arg1) : (i32, f32) -> i1274 // CHECK: scf.condition(%{{.*}}) %{{.*}}, %{{.*}} : i64, f64275 scf.condition(%4) %3#0, %3#1 : i64, f64276 // CHECK: } do {277 } do {278 // CHECK: ^{{.*}}(%{{.*}}: i64, %{{.*}}: f64):279 ^bb0(%arg2: i64, %arg3: f64):280 %5:2 = "test.some_operation"(%arg2, %arg3): (i64, f64) -> (i32, f32)281 // CHECK: scf.yield %{{.*}}, %{{.*}} : i32, f32282 scf.yield %5#0, %5#1 : i32, f32283 // CHECK: attributes {foo = "bar"}284 } attributes {foo="bar"}285 return286}287 288// CHECK-LABEL: @infinite_while289func.func @infinite_while() {290 %true = arith.constant true291 292 // CHECK: scf.while : () -> () {293 scf.while : () -> () {294 // CHECK: scf.condition(%{{.*}})295 scf.condition(%true)296 // CHECK: } do {297 } do {298 // CHECK: scf.yield299 scf.yield300 }301 return302}303 304// CHECK-LABEL: func @execute_region305func.func @execute_region() -> i64 {306 // CHECK: scf.execute_region -> i64 {307 // CHECK-NEXT: arith.constant308 // CHECK-NEXT: scf.yield309 // CHECK-NEXT: }310 %res = scf.execute_region -> i64 {311 %c1 = arith.constant 1 : i64312 scf.yield %c1 : i64313 }314 315 // CHECK: scf.execute_region -> (i64, i64) {316 %res2:2 = scf.execute_region -> (i64, i64) {317 %c1 = arith.constant 1 : i64318 scf.yield %c1, %c1 : i64, i64319 }320 321 // CHECK: scf.execute_region {322 // CHECK-NEXT: cf.br ^bb1323 // CHECK-NEXT: ^bb1:324 // CHECK-NEXT: scf.yield325 // CHECK-NEXT: }326 "scf.execute_region"() ({327 ^bb0:328 cf.br ^bb1329 ^bb1:330 scf.yield331 }) : () -> ()332 return %res : i64333}334 335// CHECK-LABEL: func.func @normalized_forall336func.func @normalized_forall(%in: tensor<100xf32>, %out: tensor<100xf32>) {337 %c1 = arith.constant 1 : index338 %num_threads = arith.constant 100 : index339 340 // CHECK: scf.forall341 // CHECK-NEXT: tensor.extract_slice342 // CHECK-NEXT: scf.forall.in_parallel343 // CHECK-NEXT: tensor.parallel_insert_slice344 // CHECK-NEXT: }345 // CHECK-NEXT: }346 // CHECK-NEXT: return347 %result = scf.forall (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> {348 %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>349 scf.forall.in_parallel {350 tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :351 tensor<1xf32> into tensor<100xf32>352 }353 }354 return355}356 357// CHECK-LABEL: func.func @explicit_loop_bounds_forall358func.func @explicit_loop_bounds_forall(%in: tensor<100xf32>,359 %out: tensor<100xf32>) {360 %c0 = arith.constant 0 : index361 %c1 = arith.constant 1 : index362 %num_threads = arith.constant 100 : index363 364 // CHECK: scf.forall365 // CHECK-NEXT: tensor.extract_slice366 // CHECK-NEXT: scf.forall.in_parallel367 // CHECK-NEXT: tensor.parallel_insert_slice368 // CHECK-NEXT: }369 // CHECK-NEXT: }370 // CHECK-NEXT: return371 %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> tensor<100xf32> {372 %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>373 scf.forall.in_parallel {374 tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :375 tensor<1xf32> into tensor<100xf32>376 }377 }378 return379}380 381// CHECK-LABEL: func.func @normalized_forall_elide_terminator382func.func @normalized_forall_elide_terminator() -> () {383 %num_threads = arith.constant 100 : index384 385 // CHECK: scf.forall386 // CHECK-NEXT: } {mapping = [#gpu.thread<x>]}387 // CHECK-NEXT: return388 scf.forall (%thread_idx) in (%num_threads) {389 scf.forall.in_parallel {390 }391 } {mapping = [#gpu.thread<x>]}392 return393 394}395 396// CHECK-LABEL: func.func @explicit_loop_bounds_forall_elide_terminator397func.func @explicit_loop_bounds_forall_elide_terminator() -> () {398 %c0 = arith.constant 0 : index399 %c1 = arith.constant 1 : index400 %num_threads = arith.constant 100 : index401 402 // CHECK: scf.forall403 // CHECK-NEXT: } {mapping = [#gpu.thread<x>]}404 // CHECK-NEXT: return405 scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) {406 scf.forall.in_parallel {407 }408 } {mapping = [#gpu.thread<x>]}409 return410}411 412// CHECK-LABEL: @switch413func.func @switch(%arg0: index) -> i32 {414 // CHECK: %{{.*}} = scf.index_switch %arg0 -> i32415 %0 = scf.index_switch %arg0 -> i32416 // CHECK-NEXT: case 2 {417 case 2 {418 // CHECK-NEXT: arith.constant419 %c10_i32 = arith.constant 10 : i32420 // CHECK-NEXT: scf.yield %{{.*}} : i32421 scf.yield %c10_i32 : i32422 // CHECK-NEXT: }423 }424 // CHECK-NEXT: case 5 {425 case 5 {426 %c20_i32 = arith.constant 20 : i32427 scf.yield %c20_i32 : i32428 }429 // CHECK: default {430 default {431 %c30_i32 = arith.constant 30 : i32432 scf.yield %c30_i32 : i32433 }434 435 // CHECK: scf.index_switch %arg0436 scf.index_switch %arg0437 // CHECK-NEXT: default {438 default {439 scf.yield440 }441 442 return %0 : i32443}444