37 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import arith5import mlir.dialects.async_dialect as async_dialect6import mlir.dialects.async_dialect.passes7from mlir.passmanager import *8 9 10def run(f):11 print("\nTEST:", f.__name__)12 f()13 14 15# CHECK-LABEL: TEST: testCreateGroupOp16@run17def 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}> : () -> i3225 # CHECK: %1 = "async.create_group"(%0) : (i32) -> !async.group26 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: testAsyncPass35# CHECK: SUCCESS36run(testAsyncPass)37