71 lines · python
1# RUN: %PYTHON %s 2>&1 | FileCheck %s2 3from mlir.ir import *4from mlir.passmanager import *5from mlir.dialects.builtin import ModuleOp6from mlir.dialects import arith7from mlir.rewrite import *8 9 10def run(f):11 print("\nTEST:", f.__name__)12 f()13 14 15# CHECK-LABEL: TEST: testRewritePattern16@run17def testRewritePattern():18 def to_muli(op, rewriter):19 with rewriter.ip:20 assert isinstance(op, arith.AddIOp)21 new_op = arith.muli(op.lhs, op.rhs, loc=op.location)22 rewriter.replace_op(op, new_op.owner)23 24 def constant_1_to_2(op, rewriter):25 c = op.value.value26 if c != 1:27 return True # failed to match28 with rewriter.ip:29 new_op = arith.constant(op.type, 2, loc=op.location)30 rewriter.replace_op(op, [new_op])31 32 with Context():33 patterns = RewritePatternSet()34 patterns.add(arith.AddIOp, to_muli)35 patterns.add(arith.ConstantOp, constant_1_to_2)36 frozen = patterns.freeze()37 38 module = ModuleOp.parse(39 r"""40 module {41 func.func @add(%a: i64, %b: i64) -> i64 {42 %sum = arith.addi %a, %b : i6443 return %sum : i6444 }45 }46 """47 )48 49 apply_patterns_and_fold_greedily(module, frozen)50 # CHECK: %0 = arith.muli %arg0, %arg1 : i6451 # CHECK: return %0 : i6452 print(module)53 54 module = ModuleOp.parse(55 r"""56 module {57 func.func @const() -> (i64, i64) {58 %0 = arith.constant 1 : i6459 %1 = arith.constant 3 : i6460 return %0, %1 : i64, i6461 }62 }63 """64 )65 66 apply_patterns_and_fold_greedily(module, frozen)67 # CHECK: %c2_i64 = arith.constant 2 : i6468 # CHECK: %c3_i64 = arith.constant 3 : i6469 # CHECK: return %c2_i64, %c3_i64 : i64, i6470 print(module)71