81 lines · plain
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(scf-parallel-for-to-nested-fors))' -split-input-file -verify-diagnostics | FileCheck %s2 3func.func private @callee(%i: index, %j: index)4 5func.func @two_iters(%lb1: index, %lb2: index, %ub1: index, %ub2: index, %step1: index, %step2: index) {6 scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {7 func.call @callee(%i, %j) : (index, index) -> ()8 }9 // CHECK: scf.for %[[VAL_0:.*]] = %[[ARG0:.*]] to %[[ARG2:.*]] step %[[ARG4:.*]] {10 // CHECK: scf.for %[[VAL_1:.*]] = %[[ARG1:.*]] to %[[ARG3:.*]] step %[[ARG5:.*]] {11 // CHECK: func.call @callee(%[[VAL_0]], %[[VAL_1]]) : (index, index) -> ()12 // CHECK: }13 // CHECK: }14 return15}16 17// -----18 19func.func private @callee(%i: index, %j: index)20 21func.func @repeated(%lb1: index, %lb2: index, %ub1: index, %ub2: index, %step1: index, %step2: index) {22 scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {23 func.call @callee(%i, %j) : (index, index) -> ()24 }25 26 scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {27 func.call @callee(%i, %j) : (index, index) -> ()28 }29 // CHECK: scf.for %[[VAL_0:.*]] = %[[ARG0:.*]] to %[[ARG2:.*]] step %[[ARG4:.*]] {30 // CHECK: scf.for %[[VAL_1:.*]] = %[[ARG1:.*]] to %[[ARG3:.*]] step %[[ARG5:.*]] {31 // CHECK: func.call @callee(%[[VAL_0]], %[[VAL_1]]) : (index, index) -> ()32 // CHECK: }33 // CHECK: }34 // CHECK: scf.for %[[VAL_2:.*]] = %[[ARG0]] to %[[ARG2]] step %[[ARG4]] {35 // CHECK: scf.for %[[VAL_3:.*]] = %[[ARG1]] to %[[ARG3]] step %[[ARG5]] {36 // CHECK: func.call @callee(%[[VAL_2]], %[[VAL_3]]) : (index, index) -> ()37 // CHECK: }38 // CHECK: }39 40 return41}42 43// -----44 45func.func private @callee(%i: index, %j: index, %k: index, %l: index)46 47func.func @nested(%lb1: index, %lb2: index, %lb3: index, %lb4: index, %ub1: index, %ub2: index, %ub3: index, %ub4: index, %step1: index, %step2: index, %step3: index, %step4: index) {48 scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {49 scf.parallel (%k, %l) = (%lb3, %lb4) to (%ub3, %ub4) step (%step3, %step4) {50 func.call @callee(%i, %j, %k, %l) : (index, index, index, index) -> ()51 }52 }53 // CHECK: scf.for %[[VAL_0:.*]] = %[[ARG0:.*]] to %[[ARG4:.*]] step %[[ARG8:.*]] {54 // CHECK: scf.for %[[VAL_1:.*]] = %[[ARG1:.*]] to %[[ARG5:.*]] step %[[ARG9:.*]] {55 // CHECK: scf.for %[[VAL_2:.*]] = %[[ARG2:.*]] to %[[ARG6:.*]] step %[[ARG10:.*]] {56 // CHECK: scf.for %[[VAL_3:.*]] = %[[ARG3:.*]] to %[[ARG7:.*]] step %[[ARG11:.*]] {57 // CHECK: func.call @callee(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) : (index, index, index, index) -> ()58 // CHECK: }59 // CHECK: }60 // CHECK: }61 // CHECK: }62 return63}64 65// -----66func.func private @callee(%i: index, %j: index) -> i3267 68func.func @two_iters_with_reduce(%lb1: index, %lb2: index, %ub1: index, %ub2: index, %step1: index, %step2: index) -> i32 {69 %c0 = arith.constant 0 : i3270 // CHECK: scf.parallel71 %0 = scf.parallel (%i, %j) = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) init (%c0) -> i32 {72 %curr = func.call @callee(%i, %j) : (index, index) -> i3273 scf.reduce(%curr : i32) {74 ^bb0(%arg3: i32, %arg4: i32):75 %3 = arith.addi %arg3, %arg4 : i3276 scf.reduce.return %3 : i3277 }78 }79 return %0 : i3280}81