188 lines · plain
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-scf-uplift-while-to-for))' -split-input-file -allow-unregistered-dialect | FileCheck %s2 3func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {4 %0 = scf.while (%arg3 = %arg0) : (index) -> (index) {5 %1 = arith.cmpi slt, %arg3, %arg1 : index6 scf.condition(%1) %arg3 : index7 } do {8 ^bb0(%arg3: index):9 "test.test1"(%arg3) : (index) -> ()10 %added = arith.addi %arg3, %arg2 : index11 "test.test2"(%added) : (index) -> ()12 scf.yield %added : index13 }14 return %0 : index15}16 17// CHECK-LABEL: func @uplift_while18// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index19// CHECK: %[[C1:.*]] = arith.constant 1 : index20// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {21// CHECK: "test.test1"(%[[I]]) : (index) -> ()22// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index23// CHECK: "test.test2"(%[[INC]]) : (index) -> ()24// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index25// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index26// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index27// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index28// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index29// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index30// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index31// CHECK: return %[[R7]] : index32 33// -----34 35func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {36 %0 = scf.while (%arg3 = %arg0) : (index) -> (index) {37 %1 = arith.cmpi sgt, %arg1, %arg3 : index38 scf.condition(%1) %arg3 : index39 } do {40 ^bb0(%arg3: index):41 "test.test1"(%arg3) : (index) -> ()42 %added = arith.addi %arg3, %arg2 : index43 "test.test2"(%added) : (index) -> ()44 scf.yield %added : index45 }46 return %0 : index47}48 49// CHECK-LABEL: func @uplift_while50// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index51// CHECK: %[[C1:.*]] = arith.constant 1 : index52// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {53// CHECK: "test.test1"(%[[I]]) : (index) -> ()54// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index55// CHECK: "test.test2"(%[[INC]]) : (index) -> ()56// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index57// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index58// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index59// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index60// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index61// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index62// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index63// CHECK: return %[[R7]] : index64 65// -----66 67func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {68 %0 = scf.while (%arg3 = %arg0) : (index) -> (index) {69 %1 = arith.cmpi slt, %arg3, %arg1 : index70 scf.condition(%1) %arg3 : index71 } do {72 ^bb0(%arg3: index):73 "test.test1"(%arg3) : (index) -> ()74 %added = arith.addi %arg2, %arg3 : index75 "test.test2"(%added) : (index) -> ()76 scf.yield %added : index77 }78 return %0 : index79}80 81// CHECK-LABEL: func @uplift_while82// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index83// CHECK: %[[C1:.*]] = arith.constant 1 : index84// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {85// CHECK: "test.test1"(%[[I]]) : (index) -> ()86// CHECK: %[[INC:.*]] = arith.addi %[[STEP]], %[[I]] : index87// CHECK: "test.test2"(%[[INC]]) : (index) -> ()88// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index89// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index90// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index91// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index92// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index93// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index94// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index95// CHECK: return %[[R7]] : index96 97 98// -----99 100func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> (i32, f32) {101 %c1 = arith.constant 1 : i32102 %c2 = arith.constant 2.0 : f32103 %0:3 = scf.while (%arg4 = %c1, %arg3 = %arg0, %arg5 = %c2) : (i32, index, f32) -> (i32, index, f32) {104 %1 = arith.cmpi slt, %arg3, %arg1 : index105 scf.condition(%1) %arg4, %arg3, %arg5 : i32, index, f32106 } do {107 ^bb0(%arg4: i32, %arg3: index, %arg5: f32):108 %1 = "test.test1"(%arg4) : (i32) -> i32109 %added = arith.addi %arg3, %arg2 : index110 %2 = "test.test2"(%arg5) : (f32) -> f32111 scf.yield %1, %added, %2 : i32, index, f32112 }113 return %0#0, %0#2 : i32, f32114}115 116// CHECK-LABEL: func @uplift_while117// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> (i32, f32)118// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32119// CHECK-DAG: %[[C2:.*]] = arith.constant 2.000000e+00 : f32120// CHECK: %[[RES:.*]]:2 = scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]]121// CHECK-SAME: iter_args(%[[ARG1:.*]] = %[[C1]], %[[ARG2:.*]] = %[[C2]]) -> (i32, f32) {122// CHECK: %[[T1:.*]] = "test.test1"(%[[ARG1]]) : (i32) -> i32123// CHECK: %[[T2:.*]] = "test.test2"(%[[ARG2]]) : (f32) -> f32124// CHECK: scf.yield %[[T1]], %[[T2]] : i32, f32125// CHECK: return %[[RES]]#0, %[[RES]]#1 : i32, f32126 127// -----128 129func.func @uplift_while(%arg0: i64, %arg1: i64, %arg2: i64) -> i64 {130 %0 = scf.while (%arg3 = %arg0) : (i64) -> (i64) {131 %1 = arith.cmpi slt, %arg3, %arg1 : i64132 scf.condition(%1) %arg3 : i64133 } do {134 ^bb0(%arg3: i64):135 "test.test1"(%arg3) : (i64) -> ()136 %added = arith.addi %arg3, %arg2 : i64137 "test.test2"(%added) : (i64) -> ()138 scf.yield %added : i64139 }140 return %0 : i64141}142 143// CHECK-LABEL: func @uplift_while144// CHECK-SAME: (%[[BEGIN:.*]]: i64, %[[END:.*]]: i64, %[[STEP:.*]]: i64) -> i64145// CHECK: %[[C1:.*]] = arith.constant 1 : i64146// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] : i64 {147// CHECK: "test.test1"(%[[I]]) : (i64) -> ()148// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : i64149// CHECK: "test.test2"(%[[INC]]) : (i64) -> ()150// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : i64151// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : i64152// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : i64153// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : i64154// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : i64155// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : i64156// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : i64157// CHECK: return %[[R7]] : i64158 159// -----160 161// A case where all 'before' arguments are forwarded but reordered.162func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> (i32, f32) {163 %c1 = arith.constant 1 : i32164 %c2 = arith.constant 2.0 : f32165 %0:3 = scf.while (%arg4 = %c1, %arg3 = %arg0, %arg5 = %c2) : (i32, index, f32) -> (index, i32, f32) {166 %1 = arith.cmpi slt, %arg3, %arg1 : index167 scf.condition(%1) %arg3, %arg4, %arg5 : index, i32, f32168 } do {169 ^bb0(%arg3: index, %arg4: i32, %arg5: f32):170 %1 = "test.test1"(%arg4) : (i32) -> i32171 %added = arith.addi %arg3, %arg2 : index172 %2 = "test.test2"(%arg5) : (f32) -> f32173 scf.yield %1, %added, %2 : i32, index, f32174 }175 return %0#1, %0#2 : i32, f32176}177 178// CHECK-LABEL: func @uplift_while179// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> (i32, f32)180// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32181// CHECK-DAG: %[[C2:.*]] = arith.constant 2.000000e+00 : f32182// CHECK: %[[RES:.*]]:2 = scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]]183// CHECK-SAME: iter_args(%[[ARG1:.*]] = %[[C1]], %[[ARG2:.*]] = %[[C2]]) -> (i32, f32) {184// CHECK: %[[T1:.*]] = "test.test1"(%[[ARG1]]) : (i32) -> i32185// CHECK: %[[T2:.*]] = "test.test2"(%[[ARG2]]) : (f32) -> f32186// CHECK: scf.yield %[[T1]], %[[T2]] : i32, f32187// CHECK: return %[[RES]]#0, %[[RES]]#1 : i32, f32188