304 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects.pdl import *5 6 7def constructAndPrintInModule(f):8 print("\nTEST:", f.__name__)9 with Context(), Location.unknown():10 module = Module.create()11 with InsertionPoint(module.body):12 f()13 print(module)14 return f15 16 17# CHECK: module {18# CHECK: pdl.pattern @operations : benefit(1) {19# CHECK: %0 = attribute20# CHECK: %1 = type21# CHECK: %2 = operation {"attr" = %0} -> (%1 : !pdl.type)22# CHECK: %3 = result 0 of %223# CHECK: %4 = operand24# CHECK: %5 = operation(%3, %4 : !pdl.value, !pdl.value)25# CHECK: rewrite %5 with "rewriter"26# CHECK: }27# CHECK: }28@constructAndPrintInModule29def test_operations():30 pattern = PatternOp(1, "operations")31 with InsertionPoint(pattern.body):32 attr = AttributeOp()33 ty = TypeOp()34 op0 = OperationOp(attributes={"attr": attr}, types=[ty])35 op0_result = ResultOp(op0, 0)36 input = OperandOp()37 root = OperationOp(args=[op0_result, input])38 RewriteOp(root, "rewriter")39 40 41# CHECK: module {42# CHECK: pdl.pattern @rewrite_with_args : benefit(1) {43# CHECK: %0 = operand44# CHECK: %1 = operation(%0 : !pdl.value)45# CHECK: rewrite %1 with "rewriter"(%0 : !pdl.value)46# CHECK: }47# CHECK: }48@constructAndPrintInModule49def test_rewrite_with_args():50 pattern = PatternOp(1, "rewrite_with_args")51 with InsertionPoint(pattern.body):52 input = OperandOp()53 root = OperationOp(args=[input])54 RewriteOp(root, "rewriter", args=[input])55 56 57# CHECK: module {58# CHECK: pdl.pattern @rewrite_multi_root_optimal : benefit(1) {59# CHECK: %0 = operand60# CHECK: %1 = operand61# CHECK: %2 = type62# CHECK: %3 = operation(%0 : !pdl.value) -> (%2 : !pdl.type)63# CHECK: %4 = result 0 of %364# CHECK: %5 = operation(%4 : !pdl.value)65# CHECK: %6 = operation(%1 : !pdl.value) -> (%2 : !pdl.type)66# CHECK: %7 = result 0 of %667# CHECK: %8 = operation(%4, %7 : !pdl.value, !pdl.value)68# CHECK: rewrite with "rewriter"(%5, %8 : !pdl.operation, !pdl.operation)69# CHECK: }70# CHECK: }71@constructAndPrintInModule72def test_rewrite_multi_root_optimal():73 pattern = PatternOp(1, "rewrite_multi_root_optimal")74 with InsertionPoint(pattern.body):75 input1 = OperandOp()76 input2 = OperandOp()77 ty = TypeOp()78 op1 = OperationOp(args=[input1], types=[ty])79 val1 = ResultOp(op1, 0)80 root1 = OperationOp(args=[val1])81 op2 = OperationOp(args=[input2], types=[ty])82 val2 = ResultOp(op2, 0)83 root2 = OperationOp(args=[val1, val2])84 RewriteOp(name="rewriter", args=[root1, root2])85 86 87# CHECK: module {88# CHECK: pdl.pattern @rewrite_multi_root_forced : benefit(1) {89# CHECK: %0 = operand90# CHECK: %1 = operand91# CHECK: %2 = type92# CHECK: %3 = operation(%0 : !pdl.value) -> (%2 : !pdl.type)93# CHECK: %4 = result 0 of %394# CHECK: %5 = operation(%4 : !pdl.value)95# CHECK: %6 = operation(%1 : !pdl.value) -> (%2 : !pdl.type)96# CHECK: %7 = result 0 of %697# CHECK: %8 = operation(%4, %7 : !pdl.value, !pdl.value)98# CHECK: rewrite %5 with "rewriter"(%8 : !pdl.operation)99# CHECK: }100# CHECK: }101@constructAndPrintInModule102def test_rewrite_multi_root_forced():103 pattern = PatternOp(1, "rewrite_multi_root_forced")104 with InsertionPoint(pattern.body):105 input1 = OperandOp()106 input2 = OperandOp()107 ty = TypeOp()108 op1 = OperationOp(args=[input1], types=[ty])109 val1 = ResultOp(op1, 0)110 root1 = OperationOp(args=[val1])111 op2 = OperationOp(args=[input2], types=[ty])112 val2 = ResultOp(op2, 0)113 root2 = OperationOp(args=[val1, val2])114 RewriteOp(root1, name="rewriter", args=[root2])115 116 117# CHECK: module {118# CHECK: pdl.pattern @rewrite_add_body : benefit(1) {119# CHECK: %0 = type : i32120# CHECK: %1 = type121# CHECK: %2 = operation -> (%0, %1 : !pdl.type, !pdl.type)122# CHECK: rewrite %2 {123# CHECK: %3 = type124# CHECK: %4 = operation "foo.op" -> (%0, %3 : !pdl.type, !pdl.type)125# CHECK: replace %2 with %4126# CHECK: }127# CHECK: }128# CHECK: }129@constructAndPrintInModule130def test_rewrite_add_body():131 pattern = PatternOp(1, "rewrite_add_body")132 with InsertionPoint(pattern.body):133 ty1 = TypeOp(IntegerType.get_signless(32))134 ty2 = TypeOp()135 root = OperationOp(types=[ty1, ty2])136 rewrite = RewriteOp(root)137 with InsertionPoint(rewrite.add_body()):138 ty3 = TypeOp()139 newOp = OperationOp(name="foo.op", types=[ty1, ty3])140 ReplaceOp(root, with_op=newOp)141 142 143# CHECK: module {144# CHECK: pdl.pattern @rewrite_type : benefit(1) {145# CHECK: %0 = type : i32146# CHECK: %1 = type147# CHECK: %2 = operation -> (%0, %1 : !pdl.type, !pdl.type)148# CHECK: rewrite %2 {149# CHECK: %3 = operation "foo.op" -> (%0, %1 : !pdl.type, !pdl.type)150# CHECK: }151# CHECK: }152# CHECK: }153@constructAndPrintInModule154def test_rewrite_type():155 pattern = PatternOp(1, "rewrite_type")156 with InsertionPoint(pattern.body):157 ty1 = TypeOp(IntegerType.get_signless(32))158 ty2 = TypeOp()159 root = OperationOp(types=[ty1, ty2])160 rewrite = RewriteOp(root)161 with InsertionPoint(rewrite.add_body()):162 newOp = OperationOp(name="foo.op", types=[ty1, ty2])163 164 165# CHECK: module {166# CHECK: pdl.pattern @rewrite_types : benefit(1) {167# CHECK: %0 = types168# CHECK: %1 = operation -> (%0 : !pdl.range<type>)169# CHECK: rewrite %1 {170# CHECK: %2 = types : [i32, i64]171# CHECK: %3 = operation "foo.op" -> (%0, %2 : !pdl.range<type>, !pdl.range<type>)172# CHECK: }173# CHECK: }174# CHECK: }175@constructAndPrintInModule176def test_rewrite_types():177 pattern = PatternOp(1, "rewrite_types")178 with InsertionPoint(pattern.body):179 types = TypesOp()180 root = OperationOp(types=[types])181 rewrite = RewriteOp(root)182 with InsertionPoint(rewrite.add_body()):183 otherTypes = TypesOp(184 [IntegerType.get_signless(32), IntegerType.get_signless(64)]185 )186 newOp = OperationOp(name="foo.op", types=[types, otherTypes])187 188 189# CHECK: module {190# CHECK: pdl.pattern @rewrite_operands : benefit(1) {191# CHECK: %0 = types192# CHECK: %1 = operands : %0193# CHECK: %2 = operation(%1 : !pdl.range<value>)194# CHECK: rewrite %2 {195# CHECK: %3 = operation "foo.op" -> (%0 : !pdl.range<type>)196# CHECK: }197# CHECK: }198# CHECK: }199@constructAndPrintInModule200def test_rewrite_operands():201 pattern = PatternOp(1, "rewrite_operands")202 with InsertionPoint(pattern.body):203 types = TypesOp()204 operands = OperandsOp(types)205 root = OperationOp(args=[operands])206 rewrite = RewriteOp(root)207 with InsertionPoint(rewrite.add_body()):208 newOp = OperationOp(name="foo.op", types=[types])209 210 211# CHECK: module {212# CHECK: pdl.pattern @native_rewrite : benefit(1) {213# CHECK: %0 = operation214# CHECK: rewrite %0 {215# CHECK: apply_native_rewrite "NativeRewrite"(%0 : !pdl.operation)216# CHECK: }217# CHECK: }218# CHECK: }219@constructAndPrintInModule220def test_native_rewrite():221 pattern = PatternOp(1, "native_rewrite")222 with InsertionPoint(pattern.body):223 root = OperationOp()224 rewrite = RewriteOp(root)225 with InsertionPoint(rewrite.add_body()):226 ApplyNativeRewriteOp([], "NativeRewrite", args=[root])227 228 229# CHECK: module {230# CHECK: pdl.pattern @attribute_with_value : benefit(1) {231# CHECK: %0 = operation232# CHECK: rewrite %0 {233# CHECK: %1 = attribute = "value"234# CHECK: apply_native_rewrite "NativeRewrite"(%1 : !pdl.attribute)235# CHECK: }236# CHECK: }237# CHECK: }238@constructAndPrintInModule239def test_attribute_with_value():240 pattern = PatternOp(1, "attribute_with_value")241 with InsertionPoint(pattern.body):242 root = OperationOp()243 rewrite = RewriteOp(root)244 with InsertionPoint(rewrite.add_body()):245 attr = AttributeOp(value=Attribute.parse('"value"'))246 ApplyNativeRewriteOp([], "NativeRewrite", args=[attr])247 248 249# CHECK: module {250# CHECK: pdl.pattern @erase : benefit(1) {251# CHECK: %0 = operation252# CHECK: rewrite %0 {253# CHECK: erase %0254# CHECK: }255# CHECK: }256# CHECK: }257@constructAndPrintInModule258def test_erase():259 pattern = PatternOp(1, "erase")260 with InsertionPoint(pattern.body):261 root = OperationOp()262 rewrite = RewriteOp(root)263 with InsertionPoint(rewrite.add_body()):264 EraseOp(root)265 266 267# CHECK: module {268# CHECK: pdl.pattern @operation_results : benefit(1) {269# CHECK: %0 = types270# CHECK: %1 = operation -> (%0 : !pdl.range<type>)271# CHECK: %2 = results of %1272# CHECK: %3 = operation(%2 : !pdl.range<value>)273# CHECK: rewrite %3 with "rewriter"274# CHECK: }275# CHECK: }276@constructAndPrintInModule277def test_operation_results():278 valueRange = RangeType.get(ValueType.get())279 pattern = PatternOp(1, "operation_results")280 with InsertionPoint(pattern.body):281 types = TypesOp()282 inputOp = OperationOp(types=[types])283 results = ResultsOp(valueRange, inputOp)284 root = OperationOp(args=[results])285 RewriteOp(root, name="rewriter")286 287 288# CHECK: module {289# CHECK: pdl.pattern : benefit(1) {290# CHECK: %0 = type291# CHECK: apply_native_constraint "typeConstraint"(%0 : !pdl.type)292# CHECK: %1 = operation -> (%0 : !pdl.type)293# CHECK: rewrite %1 with "rewrite"294# CHECK: }295# CHECK: }296@constructAndPrintInModule297def test_apply_native_constraint():298 pattern = PatternOp(1)299 with InsertionPoint(pattern.body):300 resultType = TypeOp()301 ApplyNativeConstraintOp([], "typeConstraint", args=[resultType])302 root = OperationOp(types=[resultType])303 RewriteOp(root, name="rewrite")304