103 lines · python
1# RUN: %PYTHON %s 2>&1 | FileCheck %s2 3import gc, sys4from mlir.ir import *5from mlir.passmanager import *6from mlir.dialects.builtin import ModuleOp7from mlir.dialects import pdl8from mlir.rewrite import *9 10 11def log(*args):12 print(*args, file=sys.stderr)13 sys.stderr.flush()14 15 16def run(f):17 log("\nTEST:", f.__name__)18 f()19 gc.collect()20 assert Context._get_live_count() == 021 22 23def make_pdl_module():24 with Location.unknown():25 pdl_module = Module.create()26 with InsertionPoint(pdl_module.body):27 # Change all arith.addi with index types to arith.muli.28 @pdl.pattern(benefit=1, sym_name="addi_to_mul")29 def pat():30 # Match arith.addi with index types.31 i64_type = pdl.TypeOp(IntegerType.get_signless(64))32 operand0 = pdl.OperandOp(i64_type)33 operand1 = pdl.OperandOp(i64_type)34 op0 = pdl.OperationOp(35 name="arith.addi", args=[operand0, operand1], types=[i64_type]36 )37 38 # Replace the matched op with arith.muli.39 @pdl.rewrite()40 def rew():41 newOp = pdl.OperationOp(42 name="arith.muli", args=[operand0, operand1], types=[i64_type]43 )44 pdl.ReplaceOp(op0, with_op=newOp)45 46 return pdl_module47 48 49# CHECK-LABEL: TEST: testCustomPass50@run51def testCustomPass():52 with Context():53 pdl_module = make_pdl_module()54 frozen = PDLModule(pdl_module).freeze()55 56 module = ModuleOp.parse(57 r"""58 module {59 func.func @add(%a: i64, %b: i64) -> i64 {60 %sum = arith.addi %a, %b : i6461 return %sum : i6462 }63 }64 """65 )66 67 def custom_pass_1(op, pass_):68 print("hello from pass 1!!!", file=sys.stderr)69 70 class CustomPass2:71 def __call__(self, op, pass_):72 apply_patterns_and_fold_greedily(op, frozen)73 74 custom_pass_2 = CustomPass2()75 76 pm = PassManager("any")77 pm.enable_ir_printing()78 79 # CHECK: hello from pass 1!!!80 # CHECK-LABEL: Dump After custom_pass_181 pm.add(custom_pass_1)82 # CHECK-LABEL: Dump After CustomPass283 # CHECK: arith.muli84 pm.add(custom_pass_2, "CustomPass2")85 # CHECK-LABEL: Dump After ArithToLLVMConversionPass86 # CHECK: llvm.mul87 pm.add("convert-arith-to-llvm")88 pm.run(module)89 90 # test signal_pass_failure91 def custom_pass_that_fails(op, pass_):92 print("hello from pass that fails")93 pass_.signal_pass_failure()94 95 pm = PassManager("any")96 pm.add(custom_pass_that_fails, "CustomPassThatFails")97 # CHECK: hello from pass that fails98 # CHECK: caught exception: Failure while executing pass pipeline99 try:100 pm.run(module)101 except Exception as e:102 print(f"caught exception: {e}")103