261 lines
10 KiB
Python
261 lines
10 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import arith
|
|
from mlir.dialects import func
|
|
from mlir.dialects import memref
|
|
from mlir.dialects import scf
|
|
from mlir.passmanager import PassManager
|
|
|
|
|
|
def constructAndPrintInModule(f):
|
|
print("\nTEST:", f.__name__)
|
|
with Context(), Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
f()
|
|
print(module)
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testSimpleLoop
|
|
@constructAndPrintInModule
|
|
def testSimpleLoop():
|
|
index_type = IndexType.get()
|
|
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type)
|
|
def simple_loop(lb, ub, step):
|
|
loop = scf.ForOp(lb, ub, step, [lb, lb])
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp(loop.inner_iter_args)
|
|
return
|
|
|
|
|
|
# CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
|
|
# CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
|
|
# CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
|
|
# CHECK: scf.yield %[[I1]], %[[I2]]
|
|
|
|
|
|
# CHECK-LABEL: TEST: testInductionVar
|
|
@constructAndPrintInModule
|
|
def testInductionVar():
|
|
index_type = IndexType.get()
|
|
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type)
|
|
def induction_var(lb, ub, step):
|
|
loop = scf.ForOp(lb, ub, step, [lb])
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp([loop.induction_variable])
|
|
return
|
|
|
|
|
|
# CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
|
|
# CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
|
|
# CHECK: scf.yield %[[IV]]
|
|
|
|
|
|
# CHECK-LABEL: TEST: testForSugar
|
|
@constructAndPrintInModule
|
|
def testForSugar():
|
|
index_type = IndexType.get()
|
|
memref_t = MemRefType.get([10], index_type)
|
|
range = scf.for_
|
|
|
|
# CHECK: func.func @range_loop_1(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: scf.for %[[VAL_4:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[VAL_2]] {
|
|
# CHECK: %[[VAL_5:.*]] = arith.addi %[[VAL_4]], %[[VAL_4]] : index
|
|
# CHECK: memref.store %[[VAL_5]], %[[VAL_3]]{{\[}}%[[VAL_4]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_1(lb, ub, step, memref_v):
|
|
for i in range(lb, ub, step):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_2(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 10 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 1 : index
|
|
# CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_0]] to %[[VAL_4]] step %[[VAL_5]] {
|
|
# CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index
|
|
# CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_2(lb, ub, step, memref_v):
|
|
for i in range(lb, 10, 1):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_3(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 1 : index
|
|
# CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_4]] to %[[VAL_1]] step %[[VAL_5]] {
|
|
# CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index
|
|
# CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_3(lb, ub, step, memref_v):
|
|
for i in range(0, ub, 1):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_4(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 10 : index
|
|
# CHECK: scf.for %[[VAL_6:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_2]] {
|
|
# CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_6]] : index
|
|
# CHECK: memref.store %[[VAL_7]], %[[VAL_3]]{{\[}}%[[VAL_6]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_4(lb, ub, step, memref_v):
|
|
for i in range(0, 10, step):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_5(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 10 : index
|
|
# CHECK: %[[VAL_6:.*]] = arith.constant 1 : index
|
|
# CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {
|
|
# CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index
|
|
# CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_5(lb, ub, step, memref_v):
|
|
for i in range(0, 10, 1):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_6(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 10 : index
|
|
# CHECK: %[[VAL_6:.*]] = arith.constant 1 : index
|
|
# CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {
|
|
# CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index
|
|
# CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_6(lb, ub, step, memref_v):
|
|
for i in range(0, 10):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
# CHECK: func.func @range_loop_7(%[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index, %[[VAL_3:.*]]: memref<10xindex>) {
|
|
# CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
|
|
# CHECK: %[[VAL_5:.*]] = arith.constant 10 : index
|
|
# CHECK: %[[VAL_6:.*]] = arith.constant 1 : index
|
|
# CHECK: scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_6]] {
|
|
# CHECK: %[[VAL_8:.*]] = arith.addi %[[VAL_7]], %[[VAL_7]] : index
|
|
# CHECK: memref.store %[[VAL_8]], %[[VAL_3]]{{\[}}%[[VAL_7]]] : memref<10xindex>
|
|
# CHECK: }
|
|
# CHECK: return
|
|
# CHECK: }
|
|
@func.FuncOp.from_py_func(index_type, index_type, index_type, memref_t)
|
|
def range_loop_7(lb, ub, step, memref_v):
|
|
for i in range(10):
|
|
add = arith.addi(i, i)
|
|
memref.store(add, memref_v, [i])
|
|
scf.yield_([])
|
|
|
|
|
|
@constructAndPrintInModule
|
|
def testOpsAsArguments():
|
|
index_type = IndexType.get()
|
|
callee = func.FuncOp("callee", ([], [index_type, index_type]), visibility="private")
|
|
f = func.FuncOp("ops_as_arguments", ([], []))
|
|
with InsertionPoint(f.add_entry_block()):
|
|
lb = arith.ConstantOp.create_index(0)
|
|
ub = arith.ConstantOp.create_index(42)
|
|
step = arith.ConstantOp.create_index(2)
|
|
iter_args = func.CallOp(callee, [])
|
|
loop = scf.ForOp(lb, ub, step, iter_args)
|
|
with InsertionPoint(loop.body):
|
|
scf.YieldOp(loop.inner_iter_args)
|
|
func.ReturnOp([])
|
|
|
|
|
|
# CHECK-LABEL: TEST: testOpsAsArguments
|
|
# CHECK: func private @callee() -> (index, index)
|
|
# CHECK: func @ops_as_arguments() {
|
|
# CHECK: %[[LB:.*]] = arith.constant 0
|
|
# CHECK: %[[UB:.*]] = arith.constant 42
|
|
# CHECK: %[[STEP:.*]] = arith.constant 2
|
|
# CHECK: %[[ARGS:.*]]:2 = call @callee()
|
|
# CHECK: scf.for %arg0 = %c0 to %c42 step %c2
|
|
# CHECK: iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
|
|
# CHECK: scf.yield %{{.*}}, %{{.*}}
|
|
# CHECK: return
|
|
|
|
|
|
@constructAndPrintInModule
|
|
def testIfWithoutElse():
|
|
bool = IntegerType.get_signless(1)
|
|
i32 = IntegerType.get_signless(32)
|
|
|
|
@func.FuncOp.from_py_func(bool)
|
|
def simple_if(cond):
|
|
if_op = scf.IfOp(cond)
|
|
with InsertionPoint(if_op.then_block):
|
|
one = arith.ConstantOp(i32, 1)
|
|
add = arith.AddIOp(one, one)
|
|
scf.YieldOp([])
|
|
return
|
|
|
|
|
|
# CHECK: func @simple_if(%[[ARG0:.*]]: i1)
|
|
# CHECK: scf.if %[[ARG0:.*]]
|
|
# CHECK: %[[ONE:.*]] = arith.constant 1
|
|
# CHECK: %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
|
|
# CHECK: return
|
|
|
|
|
|
@constructAndPrintInModule
|
|
def testIfWithElse():
|
|
bool = IntegerType.get_signless(1)
|
|
i32 = IntegerType.get_signless(32)
|
|
|
|
@func.FuncOp.from_py_func(bool)
|
|
def simple_if_else(cond):
|
|
if_op = scf.IfOp(cond, [i32, i32], hasElse=True)
|
|
with InsertionPoint(if_op.then_block):
|
|
x_true = arith.ConstantOp(i32, 0)
|
|
y_true = arith.ConstantOp(i32, 1)
|
|
scf.YieldOp([x_true, y_true])
|
|
with InsertionPoint(if_op.else_block):
|
|
x_false = arith.ConstantOp(i32, 2)
|
|
y_false = arith.ConstantOp(i32, 3)
|
|
scf.YieldOp([x_false, y_false])
|
|
add = arith.AddIOp(if_op.results[0], if_op.results[1])
|
|
return
|
|
|
|
|
|
# CHECK: func @simple_if_else(%[[ARG0:.*]]: i1)
|
|
# CHECK: %[[RET:.*]]:2 = scf.if %[[ARG0:.*]]
|
|
# CHECK: %[[ZERO:.*]] = arith.constant 0
|
|
# CHECK: %[[ONE:.*]] = arith.constant 1
|
|
# CHECK: scf.yield %[[ZERO]], %[[ONE]]
|
|
# CHECK: } else {
|
|
# CHECK: %[[TWO:.*]] = arith.constant 2
|
|
# CHECK: %[[THREE:.*]] = arith.constant 3
|
|
# CHECK: scf.yield %[[TWO]], %[[THREE]]
|
|
# CHECK: arith.addi %[[RET]]#0, %[[RET]]#1
|
|
# CHECK: return
|