327 lines · python
1# RUN: %PYTHON %s 2>&1 | FileCheck %s2 3from mlir.dialects import arith, func, pdl4from mlir.dialects.builtin import module5from mlir.ir import *6from mlir.rewrite import *7 8 9def construct_and_print_in_module(f):10 print("\nTEST:", f.__name__)11 with Context(), Location.unknown():12 module = Module.create()13 with InsertionPoint(module.body):14 module = f(module)15 if module is not None:16 print(module)17 return f18 19 20def get_pdl_patterns():21 # Create a rewrite from add to mul. This will match22 # - operation name is arith.addi23 # - operands are index types.24 # - there are two operands.25 with Location.unknown():26 m = Module.create()27 with InsertionPoint(m.body):28 # Change all arith.addi with index types to arith.muli.29 @pdl.pattern(benefit=1, sym_name="addi_to_mul")30 def pat():31 # Match arith.addi with index types.32 index_type = pdl.TypeOp(IndexType.get())33 operand0 = pdl.OperandOp(index_type)34 operand1 = pdl.OperandOp(index_type)35 op0 = pdl.OperationOp(36 name="arith.addi", args=[operand0, operand1], types=[index_type]37 )38 39 # Replace the matched op with arith.muli.40 @pdl.rewrite()41 def rew():42 newOp = pdl.OperationOp(43 name="arith.muli", args=[operand0, operand1], types=[index_type]44 )45 pdl.ReplaceOp(op0, with_op=newOp)46 47 # Create a PDL module from module and freeze it. At this point the ownership48 # of the module is transferred to the PDL module. This ownership transfer is49 # not yet captured Python side/has sharp edges. So best to construct the50 # module and PDL module in same scope.51 # FIXME: This should be made more robust.52 return PDLModule(m).freeze()53 54 55# CHECK-LABEL: TEST: test_add_to_mul56# CHECK: arith.muli57@construct_and_print_in_module58def test_add_to_mul(module_):59 index_type = IndexType.get()60 61 # Create a test case.62 @module(sym_name="ir")63 def ir():64 @func.func(index_type, index_type)65 def add_func(a, b):66 return arith.addi(a, b)67 68 frozen = get_pdl_patterns()69 # Could apply frozen pattern set multiple times.70 apply_patterns_and_fold_greedily(module_, frozen)71 return module_72 73 74# CHECK-LABEL: TEST: test_add_to_mul_with_op75# CHECK: arith.muli76@construct_and_print_in_module77def test_add_to_mul_with_op(module_):78 index_type = IndexType.get()79 80 # Create a test case.81 @module(sym_name="ir")82 def ir():83 @func.func(index_type, index_type)84 def add_func(a, b):85 return arith.addi(a, b)86 87 frozen = get_pdl_patterns()88 apply_patterns_and_fold_greedily(module_.operation, frozen)89 return module_90 91 92# If we use arith.constant and arith.addi here,93# these C++-defined folding/canonicalization will be applied94# implicitly in the greedy pattern rewrite driver to95# make our Python-defined folding useless,96# so here we define a new dialect to workaround this.97def load_myint_dialect():98 from mlir.dialects import irdl99 100 m = Module.create()101 with InsertionPoint(m.body):102 myint = irdl.dialect("myint")103 with InsertionPoint(myint.body):104 constant = irdl.operation_("constant")105 with InsertionPoint(constant.body):106 iattr = irdl.base(base_name="#builtin.integer")107 i32 = irdl.is_(TypeAttr.get(IntegerType.get_signless(32)))108 irdl.attributes_([iattr], ["value"])109 irdl.results_([i32], ["cst"], [irdl.Variadicity.single])110 add = irdl.operation_("add")111 with InsertionPoint(add.body):112 i32 = irdl.is_(TypeAttr.get(IntegerType.get_signless(32)))113 irdl.operands_(114 [i32, i32],115 ["lhs", "rhs"],116 [irdl.Variadicity.single, irdl.Variadicity.single],117 )118 irdl.results_([i32], ["res"], [irdl.Variadicity.single])119 120 m.operation.verify()121 irdl.load_dialects(m)122 123 124# This PDL pattern is to fold constant additions,125# including two patterns:126# 1. add(constant0, constant1) -> constant2127# where constant2 = constant0 + constant1;128# 2. add(x, 0) or add(0, x) -> x.129def get_pdl_pattern_fold():130 m = Module.create()131 i32 = IntegerType.get_signless(32)132 with InsertionPoint(m.body):133 134 @pdl.pattern(benefit=1, sym_name="myint_add_fold")135 def pat():136 t = pdl.TypeOp(i32)137 a0 = pdl.AttributeOp()138 a1 = pdl.AttributeOp()139 c0 = pdl.OperationOp(140 name="myint.constant", attributes={"value": a0}, types=[t]141 )142 c1 = pdl.OperationOp(143 name="myint.constant", attributes={"value": a1}, types=[t]144 )145 v0 = pdl.ResultOp(c0, 0)146 v1 = pdl.ResultOp(c1, 0)147 op0 = pdl.OperationOp(name="myint.add", args=[v0, v1], types=[t])148 149 @pdl.rewrite()150 def rew():151 sum = pdl.apply_native_rewrite(152 [pdl.AttributeType.get()], "add_fold", [a0, a1]153 )154 newOp = pdl.OperationOp(155 name="myint.constant", attributes={"value": sum}, types=[t]156 )157 pdl.ReplaceOp(op0, with_op=newOp)158 159 @pdl.pattern(benefit=1, sym_name="myint_add_zero_fold")160 def pat():161 t = pdl.TypeOp(i32)162 v0 = pdl.OperandOp()163 v1 = pdl.OperandOp()164 v = pdl.apply_native_constraint([pdl.ValueType.get()], "has_zero", [v0, v1])165 op0 = pdl.OperationOp(name="myint.add", args=[v0, v1], types=[t])166 167 @pdl.rewrite()168 def rew():169 pdl.ReplaceOp(op0, with_values=[v])170 171 def add_fold(rewriter, results, values):172 a0, a1 = values173 results.append(IntegerAttr.get(i32, a0.value + a1.value))174 175 def is_zero(value):176 op = value.owner177 if isinstance(op, Operation):178 return op.name == "myint.constant" and op.attributes["value"].value == 0179 return False180 181 # Check if either operand is a constant zero,182 # and append the other operand to the results if so.183 def has_zero(rewriter, results, values):184 v0, v1 = values185 if is_zero(v0):186 results.append(v1)187 return False188 if is_zero(v1):189 results.append(v0)190 return False191 return True192 193 pdl_module = PDLModule(m)194 pdl_module.register_rewrite_function("add_fold", add_fold)195 pdl_module.register_constraint_function("has_zero", has_zero)196 return pdl_module.freeze()197 198 199# CHECK-LABEL: TEST: test_pdl_register_function200# CHECK: "myint.constant"() {value = 8 : i32} : () -> i32201@construct_and_print_in_module202def test_pdl_register_function(module_):203 load_myint_dialect()204 205 module_ = Module.parse(206 """207 %c0 = "myint.constant"() { value = 2 }: () -> (i32)208 %c1 = "myint.constant"() { value = 3 }: () -> (i32)209 %x = "myint.add"(%c0, %c1): (i32, i32) -> (i32)210 "myint.add"(%x, %c1): (i32, i32) -> (i32)211 """212 )213 214 frozen = get_pdl_pattern_fold()215 apply_patterns_and_fold_greedily(module_, frozen)216 217 return module_218 219 220# CHECK-LABEL: TEST: test_pdl_register_function_constraint221# CHECK: return %arg0 : i32222@construct_and_print_in_module223def test_pdl_register_function_constraint(module_):224 load_myint_dialect()225 226 module_ = Module.parse(227 """228 func.func @f(%x : i32) -> i32 {229 %c0 = "myint.constant"() { value = 1 }: () -> (i32)230 %c1 = "myint.constant"() { value = -1 }: () -> (i32)231 %a = "myint.add"(%c0, %c1): (i32, i32) -> (i32)232 %b = "myint.add"(%a, %x): (i32, i32) -> (i32)233 %c = "myint.add"(%b, %a): (i32, i32) -> (i32)234 func.return %c : i32235 }236 """237 )238 239 frozen = get_pdl_pattern_fold()240 apply_patterns_and_fold_greedily(module_, frozen)241 242 return module_243 244 245# This pattern is to expand constant to additions246# unless the constant is no more than 1,247# e.g. 3 -> 1 + 2 -> 1 + (1 + 1).248def get_pdl_pattern_expand():249 m = Module.create()250 i32 = IntegerType.get_signless(32)251 with InsertionPoint(m.body):252 253 @pdl.pattern(benefit=1, sym_name="myint_constant_expand")254 def pat():255 t = pdl.TypeOp(i32)256 cst = pdl.AttributeOp()257 pdl.apply_native_constraint([], "is_one", [cst])258 op0 = pdl.OperationOp(259 name="myint.constant", attributes={"value": cst}, types=[t]260 )261 262 @pdl.rewrite()263 def rew():264 expanded = pdl.apply_native_rewrite(265 [pdl.OperationType.get()], "expand", [cst]266 )267 pdl.ReplaceOp(op0, with_op=expanded)268 269 def is_one(rewriter, results, values):270 cst = values[0].value271 return cst <= 1272 273 def expand(rewriter, results, values):274 cst = values[0].value275 c1 = cst // 2276 c2 = cst - c1277 with rewriter.ip:278 op1 = Operation.create(279 "myint.constant",280 results=[i32],281 attributes={"value": IntegerAttr.get(i32, c1)},282 )283 op2 = Operation.create(284 "myint.constant",285 results=[i32],286 attributes={"value": IntegerAttr.get(i32, c2)},287 )288 res = Operation.create(289 "myint.add", results=[i32], operands=[op1.result, op2.result]290 )291 results.append(res)292 293 pdl_module = PDLModule(m)294 pdl_module.register_constraint_function("is_one", is_one)295 pdl_module.register_rewrite_function("expand", expand)296 return pdl_module.freeze()297 298 299# CHECK-LABEL: TEST: test_pdl_register_function_expand300# CHECK: %0 = "myint.constant"() {value = 1 : i32} : () -> i32301# CHECK: %1 = "myint.constant"() {value = 1 : i32} : () -> i32302# CHECK: %2 = "myint.add"(%0, %1) : (i32, i32) -> i32303# CHECK: %3 = "myint.constant"() {value = 1 : i32} : () -> i32304# CHECK: %4 = "myint.constant"() {value = 1 : i32} : () -> i32305# CHECK: %5 = "myint.constant"() {value = 1 : i32} : () -> i32306# CHECK: %6 = "myint.add"(%4, %5) : (i32, i32) -> i32307# CHECK: %7 = "myint.add"(%3, %6) : (i32, i32) -> i32308# CHECK: %8 = "myint.add"(%2, %7) : (i32, i32) -> i32309# CHECK: return %8 : i32310@construct_and_print_in_module311def test_pdl_register_function_expand(module_):312 load_myint_dialect()313 314 module_ = Module.parse(315 """316 func.func @f() -> i32 {317 %0 = "myint.constant"() { value = 5 }: () -> (i32)318 return %0 : i32319 }320 """321 )322 323 frozen = get_pdl_pattern_expand()324 apply_patterns_and_fold_greedily(module_, frozen)325 326 return module_327