mlir/test/python/dialects/scf.py

   1 # RUN: %PYTHON %s | FileCheck %s
   2
   3 from mlir.ir import *
   4 from mlir.dialects import arith
   5 from mlir.dialects import scf
   6 from mlir.dialects import std
   7 from mlir.dialects import builtin
   8
   9
  10 def 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
  22 def testSimpleLoop():
  23   index_type = IndexType.get()
  24
  25   @builtin.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
  41 def testInductionVar():
  42   index_type = IndexType.get()
  43
  44   @builtin.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
  58 def testOpsAsArguments():
  59   index_type = IndexType.get()
  60   callee = builtin.FuncOp(
  61       "callee", ([], [index_type, index_type]), visibility="private")
  62   func = builtin.FuncOp("ops_as_arguments", ([], []))
  63   with InsertionPoint(func.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 = std.CallOp(callee, [])
  68     loop = scf.ForOp(lb, ub, step, iter_args)
  69     with InsertionPoint(loop.body):
  70       scf.YieldOp(loop.inner_iter_args)
  71     std.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