1# RUN: %PYTHON %s | FileCheck %s
2# This is just a smoke test that the dialect is functional.
3
4from mlir.ir import *
5from mlir.dialects import ml_program
6
7
8def constructAndPrintInModule(f):
9  print("\nTEST:", f.__name__)
10  with Context(), Location.unknown():
11    module = Module.create()
12    with InsertionPoint(module.body):
13      f()
14    print(module)
15  return f
16
17
18# CHECK-LABEL: testFuncOp
19@constructAndPrintInModule
20def testFuncOp():
21  # CHECK: ml_program.func @foobar(%arg0: si32) -> si32
22  f = ml_program.FuncOp(
23      name="foobar",
24      type=([IntegerType.get_signed(32)], [IntegerType.get_signed(32)]))
25  block = f.add_entry_block()
26  with InsertionPoint(block):
27    # CHECK: ml_program.return
28    ml_program.ReturnOp([block.arguments[0]])
29