476 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.extras import types as T5from mlir.dialects import (6 arith,7 func,8 memref,9 scf,10 cf,11)12from mlir.passmanager import PassManager13 14 15def constructAndPrintInModule(f):16 print("\nTEST:", f.__name__)17 with Context(), Location.unknown():18 module = Module.create()19 with InsertionPoint(module.body):20 f()21 print(module)22 return f23 24 25# CHECK-LABEL: TEST: testSimpleForall26# CHECK: scf.forall (%[[IV0:.*]], %[[IV1:.*]]) in (4, 8) shared_outs(%[[BOUND_ARG:.*]] = %{{.*}}) -> (tensor<4x8xf32>)27# CHECK: arith.addi %[[IV0]], %[[IV1]]28# CHECK: scf.forall.in_parallel29@constructAndPrintInModule30def testSimpleForall():31 f32 = F32Type.get()32 tensor_type = RankedTensorType.get([4, 8], f32)33 34 @func.FuncOp.from_py_func(tensor_type)35 def forall_loop(tensor):36 loop = scf.ForallOp([0, 0], [4, 8], [1, 1], [tensor])37 with InsertionPoint(loop.body):38 i, j = loop.induction_variables39 arith.addi(i, j)40 loop.terminator()41 # The verifier will check that the regions have been created properly.42 assert loop.verify()43 44 45# CHECK-LABEL: TEST: testSimpleLoop46@constructAndPrintInModule47def testSimpleLoop():48 index_type = IndexType.get()49 50 @func.FuncOp.from_py_func(index_type, index_type, index_type)51 def simple_loop(lb, ub, step):52 loop = scf.ForOp(lb, ub, step, [lb, lb])53 with InsertionPoint(loop.body):54 scf.YieldOp(loop.inner_iter_args)55 return56 57 58# CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)59# CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]60# CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])61# CHECK: scf.yield %[[I1]], %[[I2]]62 63 64# CHECK-LABEL: TEST: testInductionVar65@constructAndPrintInModule66def testInductionVar():67 index_type = IndexType.get()68 69 @func.FuncOp.from_py_func(index_type, index_type, index_type)70 def induction_var(lb, ub, step):71 loop = scf.ForOp(lb, ub, step, [lb])72 with InsertionPoint(loop.body):73 scf.YieldOp([loop.induction_variable])74 return75 76 77# CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)78# CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]79# CHECK: scf.yield %[[IV]]80 81 82# CHECK-LABEL: TEST: testForSugar83@constructAndPrintInModule84def testForSugar():85 index_type = IndexType.get()86 memref_t = MemRefType.get([10], index_type)87 range = scf.for_88 89 # CHECK: func.func @range_loop_1(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {90 # CHECK: scf.for %[[VAL_4:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[VAL_2]] {91 # CHECK: %[[VAL_5:.*]] = arith.addi %[[VAL_4]], %[[VAL_4]] : index92 # CHECK: memref.store %[[VAL_5]], %[[VAL_3]]{{\[}}%[[VAL_4]]] : memref<10xindex>93 # CHECK: }94 # CHECK: return95 # CHECK: }96 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)97 def range_loop_1(lb, ub, step, memref_v):98 for i in range(lb, ub, step):99 add = arith.addi(i, i)100 memref.store(add, memref_v, [i])101 102 scf.yield_([])103 104 # CHECK: func.func @range_loop_2(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {105 # CHECK: %[[VAL_4:.*]] = arith.constant 10 : index106 # CHECK: %[[VAL_5:.*]] = arith.constant 1 : index107 # CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_0]] to %[[VAL_4]] step %[[VAL_5]] {108 # CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index109 # CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>110 # CHECK: }111 # CHECK: return112 # CHECK: }113 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)114 def range_loop_2(lb, ub, step, memref_v):115 for i in range(lb, 10, 1):116 add = arith.addi(i, i)117 memref.store(add, memref_v, [i])118 scf.yield_([])119 120 # CHECK: func.func @range_loop_3(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {121 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index122 # CHECK: %[[VAL_5:.*]] = arith.constant 1 : index123 # CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_4]] to %[[VAL_1]] step %[[VAL_5]] {124 # CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index125 # CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>126 # CHECK: }127 # CHECK: return128 # CHECK: }129 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)130 def range_loop_3(lb, ub, step, memref_v):131 for i in range(0, ub, 1):132 add = arith.addi(i, i)133 memref.store(add, memref_v, [i])134 scf.yield_([])135 136 # CHECK: func.func @range_loop_4(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {137 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index138 # CHECK: %[[VAL_5:.*]] = arith.constant 10 : index139 # CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_2]] {140 # CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index141 # CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>142 # CHECK: }143 # CHECK: return144 # CHECK: }145 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)146 def range_loop_4(lb, ub, step, memref_v):147 for i in range(0, 10, step):148 add = arith.addi(i, i)149 memref.store(add, memref_v, [i])150 scf.yield_([])151 152 # CHECK: func.func @range_loop_5(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {153 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index154 # CHECK: %[[VAL_5:.*]] = arith.constant 10 : index155 # CHECK: %[[VAL_6:.*]] = arith.constant 1 : index156 # CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {157 # CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index158 # CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>159 # CHECK: }160 # CHECK: return161 # CHECK: }162 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)163 def range_loop_5(lb, ub, step, memref_v):164 for i in range(0, 10, 1):165 add = arith.addi(i, i)166 memref.store(add, memref_v, [i])167 scf.yield_([])168 169 # CHECK: func.func @range_loop_6(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {170 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index171 # CHECK: %[[VAL_5:.*]] = arith.constant 10 : index172 # CHECK: %[[VAL_6:.*]] = arith.constant 1 : index173 # CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {174 # CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index175 # CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>176 # CHECK: }177 # CHECK: return178 # CHECK: }179 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)180 def range_loop_6(lb, ub, step, memref_v):181 for i in range(0, 10):182 add = arith.addi(i, i)183 memref.store(add, memref_v, [i])184 scf.yield_([])185 186 # CHECK: func.func @range_loop_7(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {187 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index188 # CHECK: %[[VAL_5:.*]] = arith.constant 10 : index189 # CHECK: %[[VAL_6:.*]] = arith.constant 1 : index190 # CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {191 # CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index192 # CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>193 # CHECK: }194 # CHECK: return195 # CHECK: }196 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)197 def range_loop_7(lb, ub, step, memref_v):198 for i in range(10):199 add = arith.addi(i, i)200 memref.store(add, memref_v, [i])201 scf.yield_([])202 203 # CHECK: func.func @loop_yield_1(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {204 # CHECK: %[[VAL_4:.*]] = arith.constant 0 : index205 # CHECK: %[[VAL_5:.*]] = arith.constant 0 : index206 # CHECK: %[[VAL_6:.*]] = arith.constant 0 : index207 # CHECK: %[[VAL_7:.*]] = arith.constant 100 : index208 # CHECK: %[[VAL_8:.*]] = arith.constant 1 : index209 # CHECK: %[[VAL_10:.*]] = scf.for %[[IV:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_8]] iter_args(%[[ITER:.*]] = %[[VAL_4]]) -> (index) {210 # CHECK: %[[VAL_9:.*]] = arith.addi %[[ITER]], %[[IV]] : index211 # CHECK: scf.yield %[[VAL_9]] : index212 # CHECK: }213 # CHECK: memref.store %[[VAL_10]], %[[VAL_3]]{{\[}}%[[VAL_5]]] : memref<10xindex>214 # CHECK: return215 # CHECK: }216 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)217 def loop_yield_1(lb, ub, step, memref_v):218 sum = arith.ConstantOp.create_index(0)219 c0 = arith.ConstantOp.create_index(0)220 for i, loc_sum, sum in scf.for_(0, 100, 1, [sum]):221 loc_sum = arith.addi(loc_sum, i)222 scf.yield_([loc_sum])223 memref.store(sum, memref_v, [c0])224 225 # CHECK: func.func @loop_yield_2(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {226 # CHECK: %[[c0:.*]] = arith.constant 0 : index227 # CHECK: %[[c2:.*]] = arith.constant 2 : index228 # CHECK: %[[REF1:.*]] = arith.constant 0 : index229 # CHECK: %[[REF2:.*]] = arith.constant 1 : index230 # CHECK: %[[VAL_6:.*]] = arith.constant 0 : index231 # CHECK: %[[VAL_7:.*]] = arith.constant 100 : index232 # CHECK: %[[VAL_8:.*]] = arith.constant 1 : index233 # CHECK: %[[RES:.*]] = scf.for %[[IV:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_8]] iter_args(%[[ITER1:.*]] = %[[c0]], %[[ITER2:.*]] = %[[c2]]) -> (index, index) {234 # CHECK: %[[VAL_9:.*]] = arith.addi %[[ITER1]], %[[IV]] : index235 # CHECK: %[[VAL_10:.*]] = arith.addi %[[ITER2]], %[[IV]] : index236 # CHECK: scf.yield %[[VAL_9]], %[[VAL_10]] : index, index237 # CHECK: }238 # CHECK: return239 # CHECK: }240 @func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)241 def loop_yield_2(lb, ub, step, memref_v):242 sum1 = arith.ConstantOp.create_index(0)243 sum2 = arith.ConstantOp.create_index(2)244 c0 = arith.ConstantOp.create_index(0)245 c1 = arith.ConstantOp.create_index(1)246 for i, [loc_sum1, loc_sum2], [sum1, sum2] in scf.for_(0, 100, 1, [sum1, sum2]):247 loc_sum1 = arith.addi(loc_sum1, i)248 loc_sum2 = arith.addi(loc_sum2, i)249 scf.yield_([loc_sum1, loc_sum2])250 memref.store(sum1, memref_v, [c0])251 memref.store(sum2, memref_v, [c1])252 253 254@constructAndPrintInModule255def testOpsAsArguments():256 index_type = IndexType.get()257 callee = func.FuncOp("callee", ([], [index_type, index_type]), visibility="private")258 f = func.FuncOp("ops_as_arguments", ([], []))259 with InsertionPoint(f.add_entry_block()):260 lb = arith.ConstantOp.create_index(0)261 ub = arith.ConstantOp.create_index(42)262 step = arith.ConstantOp.create_index(2)263 iter_args = func.CallOp(callee, [])264 loop = scf.ForOp(lb, ub, step, iter_args)265 with InsertionPoint(loop.body):266 scf.YieldOp(loop.inner_iter_args)267 func.ReturnOp([])268 269 270# CHECK-LABEL: TEST: testOpsAsArguments271# CHECK: func private @callee() -> (index, index)272# CHECK: func @ops_as_arguments() {273# CHECK: %[[LB:.*]] = arith.constant 0274# CHECK: %[[UB:.*]] = arith.constant 42275# CHECK: %[[STEP:.*]] = arith.constant 2276# CHECK: %[[ARGS:.*]]:2 = call @callee()277# CHECK: scf.for %arg0 = %c0 to %c42 step %c2278# CHECK: iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)279# CHECK: scf.yield %{{.*}}, %{{.*}}280# CHECK: return281 282 283@constructAndPrintInModule284def testIfWithoutElse():285 bool = IntegerType.get_signless(1)286 i32 = IntegerType.get_signless(32)287 288 @func.FuncOp.from_py_func(bool)289 def simple_if(cond):290 if_op = scf.IfOp(cond)291 with InsertionPoint(if_op.then_block):292 one = arith.ConstantOp(i32, 1)293 add = arith.AddIOp(one, one)294 scf.YieldOp([])295 return296 297 298# CHECK: func @simple_if(%[[ARG0:.*]]: i1)299# CHECK: scf.if %[[ARG0:.*]]300# CHECK: %[[ONE:.*]] = arith.constant 1301# CHECK: %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]302# CHECK: return303 304 305@constructAndPrintInModule306def testNestedIf():307 bool = IntegerType.get_signless(1)308 i32 = IntegerType.get_signless(32)309 310 @func.FuncOp.from_py_func(bool, bool)311 def nested_if(b, c):312 if_op = scf.IfOp(b)313 with InsertionPoint(if_op.then_block) as ip:314 if_op = scf.IfOp(c, ip=ip)315 with InsertionPoint(if_op.then_block):316 one = arith.ConstantOp(i32, 1)317 add = arith.AddIOp(one, one)318 scf.YieldOp([])319 scf.YieldOp([])320 return321 322 323# CHECK: func @nested_if(%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)324# CHECK: scf.if %[[ARG0:.*]]325# CHECK: scf.if %[[ARG1:.*]]326# CHECK: %[[ONE:.*]] = arith.constant 1327# CHECK: %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]328# CHECK: return329 330 331@constructAndPrintInModule332def testIfWithElse():333 bool = IntegerType.get_signless(1)334 i32 = IntegerType.get_signless(32)335 336 @func.FuncOp.from_py_func(bool)337 def simple_if_else(cond):338 if_op = scf.IfOp(cond, [i32, i32], hasElse=True)339 with InsertionPoint(if_op.then_block):340 x_true = arith.ConstantOp(i32, 0)341 y_true = arith.ConstantOp(i32, 1)342 scf.YieldOp([x_true, y_true])343 with InsertionPoint(if_op.else_block):344 x_false = arith.ConstantOp(i32, 2)345 y_false = arith.ConstantOp(i32, 3)346 scf.YieldOp([x_false, y_false])347 add = arith.AddIOp(if_op.results[0], if_op.results[1])348 return349 350 351# CHECK: func @simple_if_else(%[[ARG0:.*]]: i1)352# CHECK: %[[RET:.*]]:2 = scf.if %[[ARG0:.*]]353# CHECK: %[[ZERO:.*]] = arith.constant 0354# CHECK: %[[ONE:.*]] = arith.constant 1355# CHECK: scf.yield %[[ZERO]], %[[ONE]]356# CHECK: } else {357# CHECK: %[[TWO:.*]] = arith.constant 2358# CHECK: %[[THREE:.*]] = arith.constant 3359# CHECK: scf.yield %[[TWO]], %[[THREE]]360# CHECK: arith.addi %[[RET]]#0, %[[RET]]#1361# CHECK: return362 363 364@constructAndPrintInModule365def testIndexSwitch():366 i32 = T.i32()367 368 @func.FuncOp.from_py_func(T.index(), results=[i32])369 def index_switch(index):370 c1 = arith.constant(i32, 1)371 c0 = arith.constant(i32, 0)372 value = arith.constant(i32, 5)373 switch_op = scf.IndexSwitchOp([i32], index, range(3))374 375 assert switch_op.regions[0] == switch_op.default_region376 assert switch_op.regions[1] == switch_op.case_regions[0]377 assert switch_op.regions[1] == switch_op.case_region(0)378 assert len(switch_op.case_regions) == 3379 assert len(switch_op.regions) == 4380 381 with InsertionPoint(switch_op.default_block):382 cf.assert_(arith.constant(T.bool(), 0), "Whoops!")383 scf.yield_([c1])384 385 for i, block in enumerate(switch_op.case_blocks):386 with InsertionPoint(block):387 scf.yield_([arith.constant(i32, i)])388 389 func.return_([switch_op.results[0]])390 391 return index_switch392 393 394# CHECK-LABEL: func.func @index_switch(395# CHECK-SAME: %[[ARG0:.*]]: index) -> i32 {396# CHECK: %[[CONSTANT_0:.*]] = arith.constant 1 : i32397# CHECK: %[[CONSTANT_1:.*]] = arith.constant 0 : i32398# CHECK: %[[CONSTANT_2:.*]] = arith.constant 5 : i32399# CHECK: %[[INDEX_SWITCH_0:.*]] = scf.index_switch %[[ARG0]] -> i32400# CHECK: case 0 {401# CHECK: %[[CONSTANT_3:.*]] = arith.constant 0 : i32402# CHECK: scf.yield %[[CONSTANT_3]] : i32403# CHECK: }404# CHECK: case 1 {405# CHECK: %[[CONSTANT_4:.*]] = arith.constant 1 : i32406# CHECK: scf.yield %[[CONSTANT_4]] : i32407# CHECK: }408# CHECK: case 2 {409# CHECK: %[[CONSTANT_5:.*]] = arith.constant 2 : i32410# CHECK: scf.yield %[[CONSTANT_5]] : i32411# CHECK: }412# CHECK: default {413# CHECK: %[[CONSTANT_6:.*]] = arith.constant false414# CHECK: cf.assert %[[CONSTANT_6]], "Whoops!"415# CHECK: scf.yield %[[CONSTANT_0]] : i32416# CHECK: }417# CHECK: return %[[INDEX_SWITCH_0]] : i32418# CHECK: }419 420 421@constructAndPrintInModule422def testIndexSwitchWithBodyBuilders():423 i32 = T.i32()424 425 @func.FuncOp.from_py_func(T.index(), results=[i32])426 def index_switch(index):427 c1 = arith.constant(i32, 1)428 c0 = arith.constant(i32, 0)429 value = arith.constant(i32, 5)430 431 def default_body_builder(switch_op):432 cf.assert_(arith.constant(T.bool(), 0), "Whoops!")433 scf.yield_([c1])434 435 def case_body_builder(switch_op, case_index: int, case_value: int):436 scf.yield_([arith.constant(i32, case_value)])437 438 result = scf.index_switch(439 results=[i32],440 arg=index,441 cases=range(3),442 case_body_builder=case_body_builder,443 default_body_builder=default_body_builder,444 )445 446 func.return_([result])447 448 return index_switch449 450 451# CHECK-LABEL: func.func @index_switch(452# CHECK-SAME: %[[ARG0:.*]]: index) -> i32 {453# CHECK: %[[CONSTANT_0:.*]] = arith.constant 1 : i32454# CHECK: %[[CONSTANT_1:.*]] = arith.constant 0 : i32455# CHECK: %[[CONSTANT_2:.*]] = arith.constant 5 : i32456# CHECK: %[[INDEX_SWITCH_0:.*]] = scf.index_switch %[[ARG0]] -> i32457# CHECK: case 0 {458# CHECK: %[[CONSTANT_3:.*]] = arith.constant 0 : i32459# CHECK: scf.yield %[[CONSTANT_3]] : i32460# CHECK: }461# CHECK: case 1 {462# CHECK: %[[CONSTANT_4:.*]] = arith.constant 1 : i32463# CHECK: scf.yield %[[CONSTANT_4]] : i32464# CHECK: }465# CHECK: case 2 {466# CHECK: %[[CONSTANT_5:.*]] = arith.constant 2 : i32467# CHECK: scf.yield %[[CONSTANT_5]] : i32468# CHECK: }469# CHECK: default {470# CHECK: %[[CONSTANT_6:.*]] = arith.constant false471# CHECK: cf.assert %[[CONSTANT_6]], "Whoops!"472# CHECK: scf.yield %[[CONSTANT_0]] : i32473# CHECK: }474# CHECK: return %[[INDEX_SWITCH_0]] : i32475# CHECK: }476