1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import arith
5from mlir.dialects import func
6from mlir.dialects import scf
7from mlir.dialects import builtin
8
9
10def constructAndPrintInModule(f):
11  print("\nTEST:", f.__name__)
12  with Context(), Location.unknown():
13    module = Module.create()
14    with InsertionPoint(module.body):
15      f()
16    print(module)
17  return f
18
19
20# CHECK-LABEL: TEST: testSimpleLoop
21@constructAndPrintInModule
22def testSimpleLoop():
23  index_type = IndexType.get()
24
25  @func.FuncOp.from_py_func(index_type, index_type, index_type)
26  def simple_loop(lb, ub, step):
27    loop = scf.ForOp(lb, ub, step, [lb, lb])
28    with InsertionPoint(loop.body):
29      scf.YieldOp(loop.inner_iter_args)
30    return
31
32
33# CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
34# CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
35# CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
36# CHECK: scf.yield %[[I1]], %[[I2]]
37
38
39# CHECK-LABEL: TEST: testInductionVar
40@constructAndPrintInModule
41def testInductionVar():
42  index_type = IndexType.get()
43
44  @func.FuncOp.from_py_func(index_type, index_type, index_type)
45  def induction_var(lb, ub, step):
46    loop = scf.ForOp(lb, ub, step, [lb])
47    with InsertionPoint(loop.body):
48      scf.YieldOp([loop.induction_variable])
49    return
50
51
52# CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
53# CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
54# CHECK: scf.yield %[[IV]]
55
56
57@constructAndPrintInModule
58def testOpsAsArguments():
59  index_type = IndexType.get()
60  callee = func.FuncOp(
61      "callee", ([], [index_type, index_type]), visibility="private")
62  f = func.FuncOp("ops_as_arguments", ([], []))
63  with InsertionPoint(f.add_entry_block()):
64    lb = arith.ConstantOp.create_index(0)
65    ub = arith.ConstantOp.create_index(42)
66    step = arith.ConstantOp.create_index(2)
67    iter_args = func.CallOp(callee, [])
68    loop = scf.ForOp(lb, ub, step, iter_args)
69    with InsertionPoint(loop.body):
70      scf.YieldOp(loop.inner_iter_args)
71    func.ReturnOp([])
72
73
74# CHECK-LABEL: TEST: testOpsAsArguments
75# CHECK: func private @callee() -> (index, index)
76# CHECK: func @ops_as_arguments() {
77# CHECK:   %[[LB:.*]] = arith.constant 0
78# CHECK:   %[[UB:.*]] = arith.constant 42
79# CHECK:   %[[STEP:.*]] = arith.constant 2
80# CHECK:   %[[ARGS:.*]]:2 = call @callee()
81# CHECK:   scf.for %arg0 = %c0 to %c42 step %c2
82# CHECK:   iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
83# CHECK:     scf.yield %{{.*}}, %{{.*}}
84# CHECK:   return
85
86
87@constructAndPrintInModule
88def testIfWithoutElse():
89  bool = IntegerType.get_signless(1)
90  i32 = IntegerType.get_signless(32)
91
92  @func.FuncOp.from_py_func(bool)
93  def simple_if(cond):
94    if_op = scf.IfOp(cond)
95    with InsertionPoint(if_op.then_block):
96      one = arith.ConstantOp(i32, 1)
97      add = arith.AddIOp(one, one)
98      scf.YieldOp([])
99    return
100
101
102# CHECK: func @simple_if(%[[ARG0:.*]]: i1)
103# CHECK: scf.if %[[ARG0:.*]]
104# CHECK:   %[[ONE:.*]] = arith.constant 1
105# CHECK:   %[[ADD:.*]] = arith.addi %[[ONE]], %[[ONE]]
106# CHECK: return
107
108
109@constructAndPrintInModule
110def testIfWithElse():
111  bool = IntegerType.get_signless(1)
112  i32 = IntegerType.get_signless(32)
113
114  @func.FuncOp.from_py_func(bool)
115  def simple_if_else(cond):
116    if_op = scf.IfOp(cond, [i32, i32], hasElse=True)
117    with InsertionPoint(if_op.then_block):
118      x_true = arith.ConstantOp(i32, 0)
119      y_true = arith.ConstantOp(i32, 1)
120      scf.YieldOp([x_true, y_true])
121    with InsertionPoint(if_op.else_block):
122      x_false = arith.ConstantOp(i32, 2)
123      y_false = arith.ConstantOp(i32, 3)
124      scf.YieldOp([x_false, y_false])
125    add = arith.AddIOp(if_op.results[0], if_op.results[1])
126    return
127
128
129# CHECK: func @simple_if_else(%[[ARG0:.*]]: i1)
130# CHECK: %[[RET:.*]]:2 = scf.if %[[ARG0:.*]]
131# CHECK:   %[[ZERO:.*]] = arith.constant 0
132# CHECK:   %[[ONE:.*]] = arith.constant 1
133# CHECK:   scf.yield %[[ZERO]], %[[ONE]]
134# CHECK: } else {
135# CHECK:   %[[TWO:.*]] = arith.constant 2
136# CHECK:   %[[THREE:.*]] = arith.constant 3
137# CHECK:   scf.yield %[[TWO]], %[[THREE]]
138# CHECK: arith.addi %[[RET]]#0, %[[RET]]#1
139# CHECK: return
140