xref: /llvm-project/mlir/test/python/dialects/async_dialect.py (revision 37fe3c6788a152dd88a54e2f22db05d9c7e53468)
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4from mlir.dialects import arith
5import mlir.dialects.async_dialect as async_dialect
6import mlir.dialects.async_dialect.passes
7from mlir.passmanager import *
8
9
10def run(f):
11    print("\nTEST:", f.__name__)
12    f()
13
14
15# CHECK-LABEL: TEST: testCreateGroupOp
16@run
17def testCreateGroupOp():
18    with Context() as ctx, Location.unknown():
19        module = Module.create()
20        with InsertionPoint(module.body):
21            i32 = IntegerType.get_signless(32)
22            group_size = arith.ConstantOp(i32, 4)
23            async_dialect.create_group(group_size)
24        # CHECK:         %0 = "arith.constant"() <{value = 4 : i32}> : () -> i32
25        # CHECK:         %1 = "async.create_group"(%0) : (i32) -> !async.group
26        print(module)
27
28def testAsyncPass():
29    with Context() as context:
30        PassManager.parse("any(async-to-async-runtime)")
31    print("SUCCESS")
32
33
34# CHECK-LABEL: testAsyncPass
35#       CHECK: SUCCESS
36run(testAsyncPass)
37