1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import scf
5from mlir.dialects import std
6from mlir.dialects import builtin
7
8
9def constructAndPrintInModule(f):
10  print("\nTEST:", f.__name__)
11  with Context(), Location.unknown():
12    module = Module.create()
13    with InsertionPoint(module.body):
14      f()
15    print(module)
16  return f
17
18
19# CHECK-LABEL: TEST: testSimpleLoop
20@constructAndPrintInModule
21def testSimpleLoop():
22  index_type = IndexType.get()
23
24  @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
25  def simple_loop(lb, ub, step):
26    loop = scf.ForOp(lb, ub, step, [lb, lb])
27    with InsertionPoint(loop.body):
28      scf.YieldOp(loop.inner_iter_args)
29    return
30
31
32# CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
33# CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
34# CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
35# CHECK: scf.yield %[[I1]], %[[I2]]
36
37
38# CHECK-LABEL: TEST: testInductionVar
39@constructAndPrintInModule
40def testInductionVar():
41  index_type = IndexType.get()
42
43  @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
44  def induction_var(lb, ub, step):
45    loop = scf.ForOp(lb, ub, step, [lb])
46    with InsertionPoint(loop.body):
47      scf.YieldOp([loop.induction_variable])
48    return
49
50
51# CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
52# CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
53# CHECK: scf.yield %[[IV]]
54
55
56@constructAndPrintInModule
57def testOpsAsArguments():
58  index_type = IndexType.get()
59  callee = builtin.FuncOp(
60      "callee", ([], [index_type, index_type]), visibility="private")
61  func = builtin.FuncOp("ops_as_arguments", ([], []))
62  with InsertionPoint(func.add_entry_block()):
63    lb = std.ConstantOp.create_index(0)
64    ub = std.ConstantOp.create_index(42)
65    step = std.ConstantOp.create_index(2)
66    iter_args = std.CallOp(callee, [])
67    loop = scf.ForOp(lb, ub, step, iter_args)
68    with InsertionPoint(loop.body):
69      scf.YieldOp(loop.inner_iter_args)
70    std.ReturnOp([])
71
72
73# CHECK-LABEL: TEST: testOpsAsArguments
74# CHECK: func private @callee() -> (index, index)
75# CHECK: func @ops_as_arguments() {
76# CHECK:   %[[LB:.*]] = constant 0
77# CHECK:   %[[UB:.*]] = constant 42
78# CHECK:   %[[STEP:.*]] = constant 2
79# CHECK:   %[[ARGS:.*]]:2 = call @callee()
80# CHECK:   scf.for %arg0 = %c0 to %c42 step %c2
81# CHECK:   iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
82# CHECK:     scf.yield %{{.*}}, %{{.*}}
83# CHECK:   return
84