404 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.ir import *4from mlir.dialects import transform5from mlir.dialects.transform import pdl as transform_pdl6 7 8def run(f):9 with Context(), Location.unknown():10 module = Module.create()11 with InsertionPoint(module.body):12 print("\nTEST:", f.__name__)13 f(module)14 print(module)15 return f16 17 18@run19def testTypes(module: Module):20 # CHECK-LABEL: TEST: testTypes21 # CHECK: !transform.any_op22 any_op = transform.AnyOpType.get()23 print(any_op)24 25 # CHECK: !transform.any_param26 any_param = transform.AnyParamType.get()27 print(any_param)28 29 # CHECK: !transform.any_value30 any_value = transform.AnyValueType.get()31 print(any_value)32 33 # CHECK: !transform.op<"foo.bar">34 # CHECK: foo.bar35 concrete_op = transform.OperationType.get("foo.bar")36 print(concrete_op)37 print(concrete_op.operation_name)38 39 # CHECK: !transform.param<i32>40 # CHECK: i3241 param = transform.ParamType.get(IntegerType.get_signless(32))42 print(param)43 print(param.type)44 45 46@run47def testSequenceOp(module: Module):48 sequence = transform.SequenceOp(49 transform.FailurePropagationMode.Propagate,50 [transform.AnyOpType.get()],51 transform.AnyOpType.get(),52 )53 with InsertionPoint(sequence.body):54 res = transform.CastOp(transform.AnyOpType.get(), sequence.bodyTarget)55 res2 = transform.cast(transform.any_op_t(), res.result)56 transform.YieldOp([res2])57 # CHECK-LABEL: TEST: testSequenceOp58 # CHECK: transform.sequence59 # CHECK: ^{{.*}}(%[[ARG0:.+]]: !transform.any_op):60 # CHECK: %[[RES:.+]] = cast %[[ARG0]] : !transform.any_op to !transform.any_op61 # CHECK: %[[RES2:.+]] = cast %[[RES]] : !transform.any_op to !transform.any_op62 # CHECK: yield %[[RES2]] : !transform.any_op63 # CHECK: }64 65 66@run67def testSequenceOp(module: Module):68 sequence = transform.SequenceOp(69 transform.FailurePropagationMode.Propagate,70 [transform.AnyOpType.get()],71 transform.AnyOpType.get(),72 )73 with InsertionPoint(sequence.body):74 transform.YieldOp([sequence.bodyTarget])75 # CHECK-LABEL: TEST: testSequenceOp76 # CHECK: = transform.sequence -> !transform.any_op failures(propagate) {77 # CHECK: ^{{.*}}(%[[ARG0:.+]]: !transform.any_op):78 # CHECK: yield %[[ARG0]] : !transform.any_op79 # CHECK: }80 81 82@run83def testNestedSequenceOp(module: Module):84 sequence = transform.SequenceOp(85 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()86 )87 with InsertionPoint(sequence.body):88 nested = transform.SequenceOp(89 transform.FailurePropagationMode.Propagate, [], sequence.bodyTarget90 )91 with InsertionPoint(nested.body):92 doubly_nested = transform.SequenceOp(93 transform.FailurePropagationMode.Propagate,94 [transform.AnyOpType.get()],95 nested.bodyTarget,96 )97 with InsertionPoint(doubly_nested.body):98 transform.YieldOp([doubly_nested.bodyTarget])99 transform.YieldOp()100 transform.YieldOp()101 # CHECK-LABEL: TEST: testNestedSequenceOp102 # CHECK: transform.sequence failures(propagate) {103 # CHECK: ^{{.*}}(%[[ARG0:.+]]: !transform.any_op):104 # CHECK: sequence %[[ARG0]] : !transform.any_op failures(propagate) {105 # CHECK: ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):106 # CHECK: = sequence %[[ARG1]] : !transform.any_op -> !transform.any_op failures(propagate) {107 # CHECK: ^{{.*}}(%[[ARG2:.+]]: !transform.any_op):108 # CHECK: yield %[[ARG2]] : !transform.any_op109 # CHECK: }110 # CHECK: }111 # CHECK: }112 113 114@run115def testSequenceOpWithExtras(module: Module):116 sequence = transform.SequenceOp(117 transform.FailurePropagationMode.Propagate,118 [],119 transform.AnyOpType.get(),120 [transform.AnyOpType.get(), transform.OperationType.get("foo.bar")],121 )122 with InsertionPoint(sequence.body):123 transform.YieldOp()124 # CHECK-LABEL: TEST: testSequenceOpWithExtras125 # CHECK: transform.sequence failures(propagate)126 # CHECK: ^{{.*}}(%{{.*}}: !transform.any_op, %{{.*}}: !transform.any_op, %{{.*}}: !transform.op<"foo.bar">):127 sequence = transform.sequence(128 transform.FailurePropagationMode.Propagate,129 [],130 transform.AnyOpType.get(),131 [transform.AnyOpType.get(), transform.OperationType.get("foo.bar")],132 )133 with InsertionPoint(sequence.body):134 transform.yield_()135 # CHECK: transform.sequence failures(propagate)136 # CHECK: ^{{.*}}(%{{.*}}: !transform.any_op, %{{.*}}: !transform.any_op, %{{.*}}: !transform.op<"foo.bar">):137 138 139@run140def testNestedSequenceOpWithExtras(module: Module):141 sequence = transform.SequenceOp(142 transform.FailurePropagationMode.Propagate,143 [],144 transform.AnyOpType.get(),145 [transform.AnyOpType.get(), transform.OperationType.get("foo.bar")],146 )147 with InsertionPoint(sequence.body):148 nested = transform.SequenceOp(149 transform.FailurePropagationMode.Propagate,150 [],151 sequence.bodyTarget,152 sequence.bodyExtraArgs,153 )154 with InsertionPoint(nested.body):155 transform.YieldOp()156 transform.YieldOp()157 # CHECK-LABEL: TEST: testNestedSequenceOpWithExtras158 # CHECK: transform.sequence failures(propagate)159 # CHECK: ^{{.*}}(%[[ARG0:.*]]: !transform.any_op, %[[ARG1:.*]]: !transform.any_op, %[[ARG2:.*]]: !transform.op<"foo.bar">):160 # CHECK: sequence %[[ARG0]], %[[ARG1]], %[[ARG2]] : (!transform.any_op, !transform.any_op, !transform.op<"foo.bar">)161 162 163@run164def testTransformPDLOps(module: Module):165 withPdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())166 with InsertionPoint(withPdl.body):167 sequence = transform.SequenceOp(168 transform.FailurePropagationMode.Propagate,169 [transform.AnyOpType.get()],170 withPdl.bodyTarget,171 )172 with InsertionPoint(sequence.body):173 match = transform_pdl.PDLMatchOp(174 transform.AnyOpType.get(), sequence.bodyTarget, "pdl_matcher"175 )176 transform.YieldOp(match)177 # CHECK-LABEL: TEST: testTransformPDLOps178 # CHECK: transform.with_pdl_patterns {179 # CHECK: ^{{.*}}(%[[ARG0:.+]]: !transform.any_op):180 # CHECK: = sequence %[[ARG0]] : !transform.any_op -> !transform.any_op failures(propagate) {181 # CHECK: ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):182 # CHECK: %[[RES:.+]] = pdl_match @pdl_matcher in %[[ARG1]]183 # CHECK: yield %[[RES]] : !transform.any_op184 # CHECK: }185 # CHECK: }186 187 188@run189def testNamedSequenceOp(module: Module):190 module.operation.attributes["transform.with_named_sequence"] = UnitAttr.get()191 named_sequence = transform.NamedSequenceOp(192 "__transform_main",193 [transform.AnyOpType.get()],194 [transform.AnyOpType.get()],195 arg_attrs=[{"transform.consumed": UnitAttr.get()}],196 )197 with InsertionPoint(named_sequence.body):198 transform.YieldOp([named_sequence.bodyTarget])199 # CHECK-LABEL: TEST: testNamedSequenceOp200 # CHECK: module attributes {transform.with_named_sequence} {201 # CHECK: transform.named_sequence @__transform_main(%[[ARG0:.+]]: !transform.any_op {transform.consumed}) -> !transform.any_op {202 # CHECK: yield %[[ARG0]] : !transform.any_op203 named_sequence = transform.named_sequence(204 "other_seq",205 [transform.AnyOpType.get()],206 [transform.AnyOpType.get()],207 arg_attrs=[{"transform.consumed": UnitAttr.get()}],208 )209 with InsertionPoint(named_sequence.body):210 transform.yield_([named_sequence.bodyTarget])211 # CHECK: transform.named_sequence @other_seq(%[[ARG1:.+]]: !transform.any_op {transform.consumed}) -> !transform.any_op {212 # CHECK: yield %[[ARG1]] : !transform.any_op213 214 215@run216def testGetParentOp(module: Module):217 sequence = transform.SequenceOp(218 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()219 )220 with InsertionPoint(sequence.body):221 transform.GetParentOp(222 transform.AnyOpType.get(),223 sequence.bodyTarget,224 isolated_from_above=True,225 nth_parent=2,226 )227 transform.get_parent_op(228 transform.AnyOpType.get(),229 sequence.bodyTarget,230 isolated_from_above=True,231 nth_parent=2,232 allow_empty_results=True,233 op_name="func.func",234 deduplicate=True,235 )236 transform.YieldOp()237 # CHECK-LABEL: TEST: testGetParentOp238 # CHECK: transform.sequence239 # CHECK: ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):240 # CHECK: = get_parent_op %[[ARG1]] {isolated_from_above, nth_parent = 2 : i64}241 # CHECK: = get_parent_op %[[ARG1]] {allow_empty_results, deduplicate, isolated_from_above, nth_parent = 2 : i64, op_name = "func.func"}242 243 244@run245def testMergeHandlesOp(module: Module):246 sequence = transform.SequenceOp(247 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()248 )249 with InsertionPoint(sequence.body):250 res = transform.MergeHandlesOp([sequence.bodyTarget])251 transform.merge_handles([res.result], deduplicate=True)252 transform.YieldOp()253 # CHECK-LABEL: TEST: testMergeHandlesOp254 # CHECK: transform.sequence255 # CHECK: ^{{.*}}(%[[ARG1:.+]]: !transform.any_op):256 # CHECK: %[[RES1:.+]] = merge_handles %[[ARG1]] : !transform.any_op257 # CHECK: = merge_handles deduplicate %[[RES1]] : !transform.any_op258 259 260@run261def testApplyPatternsOpCompact(module: Module):262 sequence = transform.SequenceOp(263 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()264 )265 with InsertionPoint(sequence.body):266 with InsertionPoint(transform.ApplyPatternsOp(sequence.bodyTarget).patterns):267 transform.ApplyCanonicalizationPatternsOp()268 with InsertionPoint(269 transform.apply_patterns(270 sequence.bodyTarget,271 apply_cse=True,272 max_iterations=3,273 max_num_rewrites=5,274 ).patterns275 ):276 transform.ApplyCanonicalizationPatternsOp()277 transform.YieldOp()278 # CHECK-LABEL: TEST: testApplyPatternsOpCompact279 # CHECK: apply_patterns to280 # CHECK: transform.apply_patterns.canonicalization281 # CHECK: } : !transform.any_op282 # CHECK: apply_patterns to283 # CHECK: transform.apply_patterns.canonicalization284 # CHECK: } {apply_cse, max_iterations = 3 : i64, max_num_rewrites = 5 : i64} : !transform.any_op285 286 287@run288def testApplyPatternsOpWithType(module: Module):289 sequence = transform.SequenceOp(290 transform.FailurePropagationMode.Propagate,291 [],292 transform.OperationType.get("test.dummy"),293 )294 with InsertionPoint(sequence.body):295 with InsertionPoint(transform.ApplyPatternsOp(sequence.bodyTarget).patterns):296 transform.ApplyCanonicalizationPatternsOp()297 transform.YieldOp()298 # CHECK-LABEL: TEST: testApplyPatternsOp299 # CHECK: apply_patterns to300 # CHECK: transform.apply_patterns.canonicalization301 # CHECK: !transform.op<"test.dummy">302 303 304@run305def testReplicateOp(module: Module):306 with_pdl = transform_pdl.WithPDLPatternsOp(transform.AnyOpType.get())307 with InsertionPoint(with_pdl.body):308 sequence = transform.SequenceOp(309 transform.FailurePropagationMode.Propagate, [], with_pdl.bodyTarget310 )311 with InsertionPoint(sequence.body):312 m1 = transform_pdl.PDLMatchOp(313 transform.AnyOpType.get(), sequence.bodyTarget, "first"314 )315 m2 = transform_pdl.PDLMatchOp(316 transform.AnyOpType.get(), sequence.bodyTarget, "second"317 )318 transform.ReplicateOp(m1, [m2])319 transform.replicate(m1, [m2])320 transform.YieldOp()321 # CHECK-LABEL: TEST: testReplicateOp322 # CHECK: %[[FIRST:.+]] = pdl_match323 # CHECK: %[[SECOND:.+]] = pdl_match324 # CHECK: %{{.*}} = replicate num(%[[FIRST]]) %[[SECOND]]325 # CHECK: %{{.*}} = replicate num(%[[FIRST]]) %[[SECOND]]326 327 328# CHECK-LABEL: TEST: testApplyRegisteredPassOp329@run330def testApplyRegisteredPassOp(module: Module):331 # CHECK: transform.sequence332 sequence = transform.SequenceOp(333 transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()334 )335 with InsertionPoint(sequence.body):336 # CHECK: %{{.*}} = apply_registered_pass "canonicalize" to {{.*}} : (!transform.any_op) -> !transform.any_op337 mod = transform.ApplyRegisteredPassOp(338 transform.AnyOpType.get(), sequence.bodyTarget, "canonicalize"339 )340 # CHECK: %{{.*}} = apply_registered_pass "canonicalize"341 # CHECK-SAME: with options = {"top-down" = false}342 # CHECK-SAME: to {{.*}} : (!transform.any_op) -> !transform.any_op343 mod = transform.ApplyRegisteredPassOp(344 transform.AnyOpType.get(),345 mod.result,346 "canonicalize",347 options={"top-down": BoolAttr.get(False)},348 )349 # CHECK: %[[MAX_ITER:.+]] = transform.param.constant350 max_iter = transform.param_constant(351 transform.AnyParamType.get(),352 IntegerAttr.get(IntegerType.get_signless(64), 10),353 )354 # CHECK: %[[MAX_REWRITE:.+]] = transform.param.constant355 max_rewrites = transform.param_constant(356 transform.AnyParamType.get(),357 IntegerAttr.get(IntegerType.get_signless(64), 1),358 )359 # CHECK: %{{.*}} = apply_registered_pass "canonicalize"360 # NB: MLIR has sorted the dict lexicographically by key:361 # CHECK-SAME: with options = {"max-iterations" = %[[MAX_ITER]],362 # CHECK-SAME: "max-rewrites" = %[[MAX_REWRITE]],363 # CHECK-SAME: "test-convergence" = true,364 # CHECK-SAME: "top-down" = false}365 # CHECK-SAME: to %{{.*}} : (!transform.any_op, !transform.any_param, !transform.any_param) -> !transform.any_op366 mod = transform.apply_registered_pass(367 transform.AnyOpType.get(),368 mod,369 "canonicalize",370 options={371 "top-down": BoolAttr.get(False),372 "max-iterations": max_iter,373 "test-convergence": True,374 "max-rewrites": max_rewrites,375 },376 )377 # CHECK: %{{.*}} = apply_registered_pass "symbol-privatize"378 # CHECK-SAME: with options = {"exclude" = ["a", "b"]}379 # CHECK-SAME: to %{{.*}} : (!transform.any_op) -> !transform.any_op380 mod = transform.apply_registered_pass(381 transform.AnyOpType.get(),382 mod,383 "symbol-privatize",384 options={"exclude": ("a", "b")},385 )386 # CHECK: %[[SYMBOL_A:.+]] = transform.param.constant387 symbol_a = transform.param_constant(388 transform.AnyParamType.get(), StringAttr.get("a")389 )390 # CHECK: %[[SYMBOL_B:.+]] = transform.param.constant391 symbol_b = transform.param_constant(392 transform.AnyParamType.get(), StringAttr.get("b")393 )394 # CHECK: %{{.*}} = apply_registered_pass "symbol-privatize"395 # CHECK-SAME: with options = {"exclude" = [%[[SYMBOL_A]], %[[SYMBOL_B]]]}396 # CHECK-SAME: to %{{.*}} : (!transform.any_op, !transform.any_param, !transform.any_param) -> !transform.any_op397 mod = transform.apply_registered_pass(398 transform.AnyOpType.get(),399 mod,400 "symbol-privatize",401 options={"exclude": (symbol_a, symbol_b)},402 )403 transform.YieldOp()404