1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-scf-uplift-while-to-for))' -split-input-file -allow-unregistered-dialect | FileCheck %s 2 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 : index 6 scf.condition(%1) %arg3 : index 7 } do { 8 ^bb0(%arg3: index): 9 "test.test1"(%arg3) : (index) -> () 10 %added = arith.addi %arg3, %arg2 : index 11 "test.test2"(%added) : (index) -> () 12 scf.yield %added : index 13 } 14 return %0 : index 15} 16 17// CHECK-LABEL: func @uplift_while 18// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index 19// CHECK: %[[C1:.*]] = arith.constant 1 : index 20// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] { 21// CHECK: "test.test1"(%[[I]]) : (index) -> () 22// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index 23// CHECK: "test.test2"(%[[INC]]) : (index) -> () 24// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index 25// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index 26// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index 27// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index 28// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index 29// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index 30// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index 31// CHECK: return %[[R7]] : index 32 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 : index 38 scf.condition(%1) %arg3 : index 39 } do { 40 ^bb0(%arg3: index): 41 "test.test1"(%arg3) : (index) -> () 42 %added = arith.addi %arg3, %arg2 : index 43 "test.test2"(%added) : (index) -> () 44 scf.yield %added : index 45 } 46 return %0 : index 47} 48 49// CHECK-LABEL: func @uplift_while 50// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index 51// CHECK: %[[C1:.*]] = arith.constant 1 : index 52// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] { 53// CHECK: "test.test1"(%[[I]]) : (index) -> () 54// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index 55// CHECK: "test.test2"(%[[INC]]) : (index) -> () 56// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index 57// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index 58// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index 59// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index 60// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index 61// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index 62// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index 63// CHECK: return %[[R7]] : index 64 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 : index 70 scf.condition(%1) %arg3 : index 71 } do { 72 ^bb0(%arg3: index): 73 "test.test1"(%arg3) : (index) -> () 74 %added = arith.addi %arg2, %arg3 : index 75 "test.test2"(%added) : (index) -> () 76 scf.yield %added : index 77 } 78 return %0 : index 79} 80 81// CHECK-LABEL: func @uplift_while 82// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index 83// CHECK: %[[C1:.*]] = arith.constant 1 : index 84// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] { 85// CHECK: "test.test1"(%[[I]]) : (index) -> () 86// CHECK: %[[INC:.*]] = arith.addi %[[STEP]], %[[I]] : index 87// CHECK: "test.test2"(%[[INC]]) : (index) -> () 88// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index 89// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index 90// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index 91// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index 92// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index 93// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index 94// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index 95// CHECK: return %[[R7]] : index 96 97 98// ----- 99 100func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> (i32, f32) { 101 %c1 = arith.constant 1 : i32 102 %c2 = arith.constant 2.0 : f32 103 %0:3 = scf.while (%arg4 = %c1, %arg3 = %arg0, %arg5 = %c2) : (i32, index, f32) -> (i32, index, f32) { 104 %1 = arith.cmpi slt, %arg3, %arg1 : index 105 scf.condition(%1) %arg4, %arg3, %arg5 : i32, index, f32 106 } do { 107 ^bb0(%arg4: i32, %arg3: index, %arg5: f32): 108 %1 = "test.test1"(%arg4) : (i32) -> i32 109 %added = arith.addi %arg3, %arg2 : index 110 %2 = "test.test2"(%arg5) : (f32) -> f32 111 scf.yield %1, %added, %2 : i32, index, f32 112 } 113 return %0#0, %0#2 : i32, f32 114} 115 116// CHECK-LABEL: func @uplift_while 117// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> (i32, f32) 118// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 119// CHECK-DAG: %[[C2:.*]] = arith.constant 2.000000e+00 : f32 120// 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) -> i32 123// CHECK: %[[T2:.*]] = "test.test2"(%[[ARG2]]) : (f32) -> f32 124// CHECK: scf.yield %[[T1]], %[[T2]] : i32, f32 125// CHECK: return %[[RES]]#0, %[[RES]]#1 : i32, f32 126 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 : i64 132 scf.condition(%1) %arg3 : i64 133 } do { 134 ^bb0(%arg3: i64): 135 "test.test1"(%arg3) : (i64) -> () 136 %added = arith.addi %arg3, %arg2 : i64 137 "test.test2"(%added) : (i64) -> () 138 scf.yield %added : i64 139 } 140 return %0 : i64 141} 142 143// CHECK-LABEL: func @uplift_while 144// CHECK-SAME: (%[[BEGIN:.*]]: i64, %[[END:.*]]: i64, %[[STEP:.*]]: i64) -> i64 145// CHECK: %[[C1:.*]] = arith.constant 1 : i64 146// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] : i64 { 147// CHECK: "test.test1"(%[[I]]) : (i64) -> () 148// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : i64 149// CHECK: "test.test2"(%[[INC]]) : (i64) -> () 150// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : i64 151// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : i64 152// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : i64 153// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : i64 154// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : i64 155// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : i64 156// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : i64 157// CHECK: return %[[R7]] : i64 158