1239 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4import io5from tempfile import NamedTemporaryFile6from mlir.ir import *7from mlir.dialects.builtin import ModuleOp8from mlir.dialects import arith, func, scf, shape9from mlir.dialects._ods_common import _cext10from mlir.extras import types as T11 12 13def run(f):14 print("\nTEST:", f.__name__)15 f()16 gc.collect()17 assert Context._get_live_count() == 018 return f19 20 21def expect_index_error(callback):22 try:23 _ = callback()24 raise RuntimeError("Expected IndexError")25 except IndexError:26 pass27 28 29# Verify iterator based traversal of the op/region/block hierarchy.30# CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators31@run32def testTraverseOpRegionBlockIterators():33 ctx = Context()34 ctx.allow_unregistered_dialects = True35 module = Module.parse(36 r"""37 func.func @f1(%arg0: i32) -> i32 {38 %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i3239 return %1 : i3240 }41 """,42 ctx,43 )44 op = module.operation45 assert op.context is ctx46 # Get the block using iterators off of the named collections.47 regions = list(op.regions[:])48 blocks = list(regions[0].blocks)49 # CHECK: MODULE REGIONS=1 BLOCKS=150 print(f"MODULE REGIONS={len(regions)} BLOCKS={len(blocks)}")51 52 # Should verify.53 # CHECK: .verify = True54 print(f".verify = {module.operation.verify()}")55 56 # Get the blocks from the default collection.57 default_blocks = list(regions[0])58 # They should compare equal regardless of how obtained.59 assert default_blocks == blocks60 61 # Should be able to get the operations from either the named collection62 # or the block.63 operations = list(blocks[0].operations)64 default_operations = list(blocks[0])65 assert default_operations == operations66 67 def walk_operations(indent, op):68 for i, region in enumerate(op.regions):69 print(f"{indent}REGION {i}:")70 for j, block in enumerate(region):71 print(f"{indent} BLOCK {j}:")72 for k, child_op in enumerate(block):73 print(f"{indent} OP {k}: {child_op}")74 walk_operations(indent + " ", child_op)75 76 # CHECK: REGION 0:77 # CHECK: BLOCK 0:78 # CHECK: OP 0: func79 # CHECK: REGION 0:80 # CHECK: BLOCK 0:81 # CHECK: OP 0: %0 = "custom.addi"82 # CHECK: OP 1: func.return83 walk_operations("", op)84 85 # CHECK: Region iter: <mlir.{{.+}}.RegionIterator86 # CHECK: Block iter: <mlir.{{.+}}.BlockIterator87 # CHECK: Operation iter: <mlir.{{.+}}.OperationIterator88 print(" Region iter:", iter(op.regions))89 print(" Block iter:", iter(op.regions[-1]))90 print("Operation iter:", iter(op.regions[-1].blocks[-1]))91 92 try:93 op.regions[-42]94 except IndexError as e:95 # CHECK: Region OOB: index out of range96 print("Region OOB:", e)97 try:98 op.regions[0].blocks[-42]99 except IndexError as e:100 # CHECK: attempt to access out of bounds block101 print(e)102 try:103 op.regions[0].blocks[0].operations[-42]104 except IndexError as e:105 # CHECK: attempt to access out of bounds operation106 print(e)107 108 109# Verify index based traversal of the op/region/block hierarchy.110# CHECK-LABEL: TEST: testTraverseOpRegionBlockIndices111@run112def testTraverseOpRegionBlockIndices():113 ctx = Context()114 ctx.allow_unregistered_dialects = True115 module = Module.parse(116 r"""117 func.func @f1(%arg0: i32) -> i32 {118 %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32119 return %1 : i32120 }121 """,122 ctx,123 )124 125 def walk_operations(indent, op):126 for i in range(len(op.regions)):127 region = op.regions[i]128 print(f"{indent}REGION {i}:")129 for j in range(len(region.blocks)):130 block = region.blocks[j]131 print(f"{indent} BLOCK {j}:")132 for k in range(len(block.operations)):133 child_op = block.operations[k]134 print(f"{indent} OP {k}: {child_op}")135 print(136 f"{indent} OP {k}: parent {child_op.operation.parent.name}"137 )138 walk_operations(indent + " ", child_op)139 140 # CHECK: REGION 0:141 # CHECK: BLOCK 0:142 # CHECK: OP 0: func143 # CHECK: OP 0: parent builtin.module144 # CHECK: REGION 0:145 # CHECK: BLOCK 0:146 # CHECK: OP 0: %0 = "custom.addi"147 # CHECK: OP 0: parent func.func148 # CHECK: OP 1: func.return149 # CHECK: OP 1: parent func.func150 walk_operations("", module.operation)151 152 153# CHECK-LABEL: TEST: testBlockAndRegionOwners154@run155def testBlockAndRegionOwners():156 ctx = Context()157 ctx.allow_unregistered_dialects = True158 module = Module.parse(159 r"""160 builtin.module {161 func.func @f() {162 func.return163 }164 }165 """,166 ctx,167 )168 169 assert module.operation.regions[0].owner == module.operation170 assert module.operation.regions[0].blocks[0].owner == module.operation171 172 func = module.body.operations[0]173 assert func.operation.regions[0].owner == func174 assert func.operation.regions[0].blocks[0].owner == func175 176 177# CHECK-LABEL: TEST: testBlockArgumentList178@run179def testBlockArgumentList():180 with Context() as ctx:181 module = Module.parse(182 r"""183 func.func @f1(%arg0: i32, %arg1: f64, %arg2: index) {184 return185 }186 """,187 ctx,188 )189 func = module.body.operations[0]190 entry_block = func.regions[0].blocks[0]191 assert len(entry_block.arguments) == 3192 # CHECK: Argument 0, type i32193 # CHECK: Argument 1, type f64194 # CHECK: Argument 2, type index195 for arg in entry_block.arguments:196 print(f"Argument {arg.arg_number}, type {arg.type}")197 new_type = IntegerType.get_signless(8 * (arg.arg_number + 1))198 arg.set_type(new_type)199 200 # CHECK: Argument 0, type i8201 # CHECK: Argument 1, type i16202 # CHECK: Argument 2, type i24203 for arg in entry_block.arguments:204 print(f"Argument {arg.arg_number}, type {arg.type}")205 206 # Check that slicing works for block argument lists.207 # CHECK: Argument 1, type i16208 # CHECK: Argument 2, type i24209 for arg in entry_block.arguments[1:]:210 print(f"Argument {arg.arg_number}, type {arg.type}")211 212 # Check that we can concatenate slices of argument lists.213 # CHECK: Length: 4214 print("Length: ", len(entry_block.arguments[:2] + entry_block.arguments[1:]))215 216 # CHECK: Type: i8217 # CHECK: Type: i16218 # CHECK: Type: i24219 for t in entry_block.arguments.types:220 print("Type: ", t)221 222 # Check that slicing and type access compose.223 # CHECK: Sliced type: i16224 # CHECK: Sliced type: i24225 for t in entry_block.arguments[1:].types:226 print("Sliced type: ", t)227 228 # Check that slice addition works as expected.229 # CHECK: Argument 2, type i24230 # CHECK: Argument 0, type i8231 restructured = entry_block.arguments[-1:] + entry_block.arguments[:1]232 for arg in restructured:233 print(f"Argument {arg.arg_number}, type {arg.type}")234 235 236# CHECK-LABEL: TEST: testOperationOperands237@run238def testOperationOperands():239 with Context() as ctx:240 ctx.allow_unregistered_dialects = True241 module = Module.parse(242 r"""243 func.func @f1(%arg0: i32) {244 %0 = "test.producer"() : () -> i64245 "test.consumer"(%arg0, %0) : (i32, i64) -> ()246 return247 }"""248 )249 func = module.body.operations[0]250 entry_block = func.regions[0].blocks[0]251 consumer = entry_block.operations[1]252 assert len(consumer.operands) == 2253 # CHECK: Operand 0, type i32254 # CHECK: Operand 1, type i64255 for i, operand in enumerate(consumer.operands):256 print(f"Operand {i}, type {operand.type}")257 258 259# CHECK-LABEL: TEST: testOperationOperandsSlice260@run261def testOperationOperandsSlice():262 with Context() as ctx:263 ctx.allow_unregistered_dialects = True264 module = Module.parse(265 r"""266 func.func @f1() {267 %0 = "test.producer0"() : () -> i64268 %1 = "test.producer1"() : () -> i64269 %2 = "test.producer2"() : () -> i64270 %3 = "test.producer3"() : () -> i64271 %4 = "test.producer4"() : () -> i64272 "test.consumer"(%0, %1, %2, %3, %4) : (i64, i64, i64, i64, i64) -> ()273 return274 }"""275 )276 func = module.body.operations[0]277 entry_block = func.regions[0].blocks[0]278 consumer = entry_block.operations[5]279 assert len(consumer.operands) == 5280 for left, right in zip(consumer.operands, consumer.operands[::-1][::-1]):281 assert left == right282 283 # CHECK: test.producer0284 # CHECK: test.producer1285 # CHECK: test.producer2286 # CHECK: test.producer3287 # CHECK: test.producer4288 full_slice = consumer.operands[:]289 for operand in full_slice:290 print(operand)291 292 # CHECK: test.producer0293 # CHECK: test.producer1294 first_two = consumer.operands[0:2]295 for operand in first_two:296 print(operand)297 298 # CHECK: test.producer3299 # CHECK: test.producer4300 last_two = consumer.operands[3:]301 for operand in last_two:302 print(operand)303 304 # CHECK: test.producer0305 # CHECK: test.producer2306 # CHECK: test.producer4307 even = consumer.operands[::2]308 for operand in even:309 print(operand)310 311 # CHECK: test.producer2312 fourth = consumer.operands[::2][1::2]313 for operand in fourth:314 print(operand)315 316 317# CHECK-LABEL: TEST: testOperationOperandsSet318@run319def testOperationOperandsSet():320 with Context() as ctx, Location.unknown(ctx):321 ctx.allow_unregistered_dialects = True322 module = Module.parse(323 r"""324 func.func @f1() {325 %0 = "test.producer0"() : () -> i64326 %1 = "test.producer1"() : () -> i64327 %2 = "test.producer2"() : () -> i64328 "test.consumer"(%0) : (i64) -> ()329 return330 }"""331 )332 func = module.body.operations[0]333 entry_block = func.regions[0].blocks[0]334 producer1 = entry_block.operations[1]335 producer2 = entry_block.operations[2]336 consumer = entry_block.operations[3]337 assert len(consumer.operands) == 1338 type = consumer.operands[0].type339 340 # CHECK: test.producer1341 consumer.operands[0] = producer1.result342 print(consumer.operands[0])343 344 # CHECK: test.producer2345 consumer.operands[-1] = producer2.result346 print(consumer.operands[0])347 348 349# CHECK-LABEL: TEST: testDetachedOperation350@run351def testDetachedOperation():352 ctx = Context()353 ctx.allow_unregistered_dialects = True354 with Location.unknown(ctx):355 i32 = IntegerType.get_signed(32)356 op1 = Operation.create(357 "custom.op1",358 results=[i32, i32],359 regions=1,360 attributes={361 "foo": StringAttr.get("foo_value"),362 "bar": StringAttr.get("bar_value"),363 },364 )365 # CHECK: %0:2 = "custom.op1"() ({366 # CHECK: }) {bar = "bar_value", foo = "foo_value"} : () -> (si32, si32)367 print(op1)368 369 # TODO: Check successors once enough infra exists to do it properly.370 371 372# CHECK-LABEL: TEST: testOperationInsertionPoint373@run374def testOperationInsertionPoint():375 ctx = Context()376 ctx.allow_unregistered_dialects = True377 module = Module.parse(378 r"""379 func.func @f1(%arg0: i32) -> i32 {380 %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32381 return %1 : i32382 }383 """,384 ctx,385 )386 387 # Create test op.388 with Location.unknown(ctx):389 op1 = Operation.create("custom.op1")390 op2 = Operation.create("custom.op2")391 392 func = module.body.operations[0]393 entry_block = func.regions[0].blocks[0]394 ip = InsertionPoint.at_block_begin(entry_block)395 ip.insert(op1)396 ip.insert(op2)397 # CHECK: func @f1398 # CHECK: "custom.op1"()399 # CHECK: "custom.op2"()400 # CHECK: %0 = "custom.addi"401 print(module)402 403 # Trying to add a previously added op should raise.404 try:405 ip.insert(op1)406 except ValueError:407 pass408 else:409 assert False, "expected insert of attached op to raise"410 411 412# CHECK-LABEL: TEST: testOperationWithRegion413@run414def testOperationWithRegion():415 ctx = Context()416 ctx.allow_unregistered_dialects = True417 with Location.unknown(ctx):418 i32 = IntegerType.get_signed(32)419 op1 = Operation.create("custom.op1", regions=1)420 block = op1.regions[0].blocks.append(i32, i32)421 # CHECK: "custom.op1"() ({422 # CHECK: ^bb0(%arg0: si32, %arg1: si32):423 # CHECK: "custom.terminator"() : () -> ()424 # CHECK: }) : () -> ()425 terminator = Operation.create("custom.terminator")426 ip = InsertionPoint(block)427 ip.insert(terminator)428 print(op1)429 430 # Now add the whole operation to another op.431 # TODO: Verify lifetime hazard by nulling out the new owning module and432 # accessing op1.433 # TODO: Also verify accessing the terminator once both parents are nulled434 # out.435 module = Module.parse(436 r"""437 func.func @f1(%arg0: i32) -> i32 {438 %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32439 return %1 : i32440 }441 """442 )443 func = module.body.operations[0]444 entry_block = func.regions[0].blocks[0]445 ip = InsertionPoint.at_block_begin(entry_block)446 ip.insert(op1)447 # CHECK: func @f1448 # CHECK: "custom.op1"()449 # CHECK: "custom.terminator"450 # CHECK: %0 = "custom.addi"451 print(module)452 453 454# CHECK-LABEL: TEST: testOperationResultList455@run456def testOperationResultList():457 ctx = Context()458 module = Module.parse(459 r"""460 func.func @f1() {461 %0:3 = call @f2() : () -> (i32, f64, index)462 call @f3() : () -> ()463 return464 }465 func.func private @f2() -> (i32, f64, index)466 func.func private @f3() -> ()467 """,468 ctx,469 )470 caller = module.body.operations[0]471 call = caller.regions[0].blocks[0].operations[0]472 assert len(call.results) == 3473 # CHECK: Result 0, type i32474 # CHECK: Result 1, type f64475 # CHECK: Result 2, type index476 for res in call.results:477 print(f"Result {res.result_number}, type {res.type}")478 479 # CHECK: Result type i32480 # CHECK: Result type f64481 # CHECK: Result type index482 for t in call.results.types:483 print(f"Result type {t}")484 485 # Out of range486 expect_index_error(lambda: call.results[3])487 expect_index_error(lambda: call.results[-4])488 489 no_results_call = caller.regions[0].blocks[0].operations[1]490 assert len(no_results_call.results) == 0491 assert no_results_call.results.owner == no_results_call492 493 494# CHECK-LABEL: TEST: testOperationResultListSlice495@run496def testOperationResultListSlice():497 with Context() as ctx:498 ctx.allow_unregistered_dialects = True499 module = Module.parse(500 r"""501 func.func @f1() {502 "some.op"() : () -> (i1, i2, i3, i4, i5)503 return504 }505 """506 )507 func = module.body.operations[0]508 entry_block = func.regions[0].blocks[0]509 producer = entry_block.operations[0]510 511 assert len(producer.results) == 5512 for left, right in zip(producer.results, producer.results[::-1][::-1]):513 assert left == right514 assert left.result_number == right.result_number515 516 # CHECK: Result 0, type i1517 # CHECK: Result 1, type i2518 # CHECK: Result 2, type i3519 # CHECK: Result 3, type i4520 # CHECK: Result 4, type i5521 full_slice = producer.results[:]522 for res in full_slice:523 print(f"Result {res.result_number}, type {res.type}")524 525 # CHECK: Result 1, type i2526 # CHECK: Result 2, type i3527 # CHECK: Result 3, type i4528 middle = producer.results[1:4]529 for res in middle:530 print(f"Result {res.result_number}, type {res.type}")531 532 # CHECK: Result 1, type i2533 # CHECK: Result 3, type i4534 odd = producer.results[1::2]535 for res in odd:536 print(f"Result {res.result_number}, type {res.type}")537 538 # CHECK: Result 3, type i4539 # CHECK: Result 1, type i2540 inverted_middle = producer.results[-2:0:-2]541 for res in inverted_middle:542 print(f"Result {res.result_number}, type {res.type}")543 544 545# CHECK-LABEL: TEST: testOperationAttributes546@run547def testOperationAttributes():548 ctx = Context()549 ctx.allow_unregistered_dialects = True550 module = Module.parse(551 r"""552 "some.op"() { some.attribute = 1 : i8,553 other.attribute = 3.0,554 dependent = "text" } : () -> ()555 """,556 ctx,557 )558 op = module.body.operations[0]559 assert len(op.attributes) == 3560 iattr = op.attributes["some.attribute"]561 fattr = op.attributes["other.attribute"]562 sattr = op.attributes["dependent"]563 # CHECK: Attribute type i8, value 1564 print(f"Attribute type {iattr.type}, value {iattr.value}")565 # CHECK: Attribute type f64, value 3.0566 print(f"Attribute type {fattr.type}, value {fattr.value}")567 # CHECK: Attribute value text568 print(f"Attribute value {sattr.value}")569 # CHECK: Attribute value b'text'570 print(f"Attribute value {sattr.value_bytes}")571 572 # Python dict-style iteration573 # We don't know in which order the attributes are stored.574 # CHECK-DAG: dependent575 # CHECK-DAG: other.attribute576 # CHECK-DAG: some.attribute577 for name in op.attributes:578 print(name)579 580 # Basic dict-like introspection581 # CHECK: True582 print("some.attribute" in op.attributes)583 # CHECK: False584 print("missing" in op.attributes)585 # CHECK: Keys: ['dependent', 'other.attribute', 'some.attribute']586 print("Keys:", sorted(op.attributes.keys()))587 # CHECK: Values count 3588 print("Values count", len(op.attributes.values()))589 # CHECK: Items count 3590 print("Items count", len(op.attributes.items()))591 592 # Dict() conversion test593 d = {k: v.value for k, v in dict(op.attributes).items()}594 # CHECK: Dict mapping {'dependent': 'text', 'other.attribute': 3.0, 'some.attribute': 1}595 print("Dict mapping", d)596 597 # Check that exceptions are raised as expected.598 try:599 op.attributes["does_not_exist"]600 except KeyError:601 pass602 else:603 assert False, "expected KeyError on accessing a non-existent attribute"604 605 try:606 op.attributes[42]607 except IndexError:608 pass609 else:610 assert False, "expected IndexError on accessing an out-of-bounds attribute"611 612 613# CHECK-LABEL: TEST: testOperationPrint614@run615def testOperationPrint():616 ctx = Context()617 module = Module.parse(618 r"""619 func.func @f1(%arg0: i32) -> i32 {620 %0 = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> loc("nom")621 %1 = arith.constant dense_resource<resource1> : tensor<3xi64>622 return %arg0 : i32623 }624 625 {-#626 dialect_resources: {627 builtin: {628 resource1: "0x08000000010000000000000002000000000000000300000000000000"629 }630 }631 #-}632 """,633 ctx,634 )635 636 # Test print to stdout.637 # CHECK: return %arg0 : i32638 # CHECK: resource1: "0x08639 module.operation.print()640 641 # Test print to text file.642 f = io.StringIO()643 # CHECK: <class 'str'>644 # CHECK: return %arg0 : i32645 module.operation.print(file=f)646 str_value = f.getvalue()647 print(str_value.__class__)648 print(f.getvalue())649 650 # Test roundtrip to bytecode.651 bytecode_stream = io.BytesIO()652 module.operation.write_bytecode(bytecode_stream, desired_version=1)653 bytecode = bytecode_stream.getvalue()654 assert bytecode.startswith(b"ML\xefR"), "Expected bytecode to start with MLïR"655 with NamedTemporaryFile() as tmpfile:656 module.operation.write_bytecode(str(tmpfile.name), desired_version=1)657 tmpfile.seek(0)658 assert tmpfile.read().startswith(659 b"ML\xefR"660 ), "Expected bytecode to start with MLïR"661 ctx2 = Context()662 module_roundtrip = Module.parse(bytecode, ctx2)663 f = io.StringIO()664 module_roundtrip.operation.print(file=f)665 roundtrip_value = f.getvalue()666 assert str_value == roundtrip_value, "Mismatch after roundtrip bytecode"667 668 # Test print to binary file.669 f = io.BytesIO()670 # CHECK: <class 'bytes'>671 # CHECK: return %arg0 : i32672 module.operation.print(file=f, binary=True)673 bytes_value = f.getvalue()674 print(bytes_value.__class__)675 print(bytes_value)676 677 # Test print local_scope.678 # CHECK: constant dense<[1, 2, 3, 4]> : tensor<4xi32> loc("nom")679 module.operation.print(enable_debug_info=True, use_local_scope=True)680 # CHECK: %nom = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>681 module.operation.print(use_name_loc_as_prefix=True, use_local_scope=True)682 683 # Test printing using state.684 state = AsmState(module.operation)685 # CHECK: constant dense<[1, 2, 3, 4]> : tensor<4xi32>686 module.operation.print(state)687 688 # Test print with options.689 # CHECK: value = dense_resource<__elided__> : tensor<4xi32>690 # CHECK: "func.return"(%arg0) : (i32) -> () -:5:7691 # CHECK-NOT: resource1: "0x08692 module.operation.print(693 large_elements_limit=2,694 enable_debug_info=True,695 pretty_debug_info=True,696 print_generic_op_form=True,697 use_local_scope=True,698 )699 700 # Test print with skip_regions option701 # CHECK: func.func @f1(%arg0: i32) -> i32702 # CHECK-NOT: func.return703 module.body.operations[0].print(704 skip_regions=True,705 )706 707 # Test print with large_resource_limit.708 # CHECK: func.func @f1(%arg0: i32) -> i32709 # CHECK-NOT: resource1: "0x08710 module.operation.print(large_resource_limit=2)711 712 # Test large_elements_limit has no effect on resource string713 # CHECK: func.func @f1(%arg0: i32) -> i32714 # CHECK: resource1: "0x08715 module.operation.print(large_elements_limit=2)716 717 718# CHECK-LABEL: TEST: testKnownOpView719@run720def testKnownOpView():721 with Context(), Location.unknown():722 Context.current.allow_unregistered_dialects = True723 module = Module.parse(724 r"""725 %1 = "custom.f32"() : () -> f32726 %2 = "custom.f32"() : () -> f32727 %3 = arith.addf %1, %2 : f32728 %4 = arith.constant 0 : i32729 """730 )731 print(module)732 733 # addf should map to a known OpView class in the arithmetic dialect.734 # We know the OpView for it defines an 'lhs' attribute.735 addf = module.body.operations[2]736 # CHECK: <mlir.dialects._arith_ops_gen.AddFOp object737 print(repr(addf))738 # CHECK: "custom.f32"()739 print(addf.lhs)740 741 # One of the custom ops should resolve to the default OpView.742 custom = module.body.operations[0]743 # CHECK: OpView object744 print(repr(custom))745 746 # Check again to make sure negative caching works.747 custom = module.body.operations[0]748 # CHECK: OpView object749 print(repr(custom))750 751 # constant should map to an extension OpView class in the arithmetic dialect.752 constant = module.body.operations[3]753 # CHECK: <mlir.dialects.arith.ConstantOp object754 print(repr(constant))755 # Checks that the arith extension is being registered successfully756 # (literal_value is a property on the extension class but not on the default OpView).757 # CHECK: literal value 0758 print("literal value", constant.literal_value)759 760 # Checks that "late" registration/replacement (i.e., post all module loading/initialization)761 # is working correctly.762 @_cext.register_operation(arith._Dialect, replace=True)763 class ConstantOp(arith.ConstantOp):764 def __init__(self, result, value, *, loc=None, ip=None):765 if isinstance(value, int):766 super().__init__(IntegerAttr.get(result, value), loc=loc, ip=ip)767 elif isinstance(value, float):768 super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip)769 else:770 super().__init__(value, loc=loc, ip=ip)771 772 constant = module.body.operations[3]773 # CHECK: <__main__.testKnownOpView.<locals>.ConstantOp object774 print(repr(constant))775 776 777# CHECK-LABEL: TEST: testFailedGenericOperationCreationReportsError778@run779def testFailedGenericOperationCreationReportsError():780 with Context(), Location.unknown():781 c0 = shape.const_shape([])782 c1 = shape.const_shape([1, 2, 3])783 try:784 shape.MeetOp.build_generic(operands=[c0, c1])785 except MLIRError as e:786 # CHECK: unequal shape cardinality787 print(e)788 else:789 assert False, "Expected exception"790 791 792# CHECK-LABEL: TEST: testSingleResultProperty793@run794def testSingleResultProperty():795 with Context(), Location.unknown():796 Context.current.allow_unregistered_dialects = True797 module = Module.parse(798 r"""799 "custom.no_result"() : () -> ()800 %0:2 = "custom.two_result"() : () -> (f32, f32)801 %1 = "custom.one_result"() : () -> f32802 """803 )804 print(module)805 806 try:807 module.body.operations[0].result808 except ValueError as e:809 # CHECK: Cannot call .result on operation custom.no_result which has 0 results810 print(e)811 else:812 assert False, "Expected exception"813 814 try:815 module.body.operations[1].result816 except ValueError as e:817 # CHECK: Cannot call .result on operation custom.two_result which has 2 results818 print(e)819 else:820 assert False, "Expected exception"821 822 # CHECK: %1 = "custom.one_result"() : () -> f32823 print(module.body.operations[2])824 825 826def create_invalid_operation():827 # This module has two region and is invalid verify that we fallback828 # to the generic printer for safety.829 op = Operation.create("builtin.module", regions=2)830 op.regions[0].blocks.append()831 return op832 833 834# CHECK-LABEL: TEST: testInvalidOperationStrSoftFails835@run836def testInvalidOperationStrSoftFails():837 ctx = Context()838 with Location.unknown(ctx):839 invalid_op = create_invalid_operation()840 # Verify that we fallback to the generic printer for safety.841 # CHECK: "builtin.module"() ({842 # CHECK: }) : () -> ()843 print(invalid_op)844 try:845 invalid_op.verify()846 except MLIRError as e:847 # CHECK: Exception: <848 # CHECK: Verification failed:849 # CHECK: error: unknown: 'builtin.module' op requires one region850 # CHECK: note: unknown: see current operation:851 # CHECK: "builtin.module"() ({852 # CHECK: ^bb0:853 # CHECK: }, {854 # CHECK: }) : () -> ()855 # CHECK: >856 print(f"Exception: <{e}>")857 858 859# CHECK-LABEL: TEST: testInvalidModuleStrSoftFails860@run861def testInvalidModuleStrSoftFails():862 ctx = Context()863 with Location.unknown(ctx):864 module = Module.create()865 with InsertionPoint(module.body):866 invalid_op = create_invalid_operation()867 # Verify that we fallback to the generic printer for safety.868 # CHECK: "builtin.module"() ({869 # CHECK: }) : () -> ()870 print(module)871 872 873# CHECK-LABEL: TEST: testInvalidOperationGetAsmBinarySoftFails874@run875def testInvalidOperationGetAsmBinarySoftFails():876 ctx = Context()877 with Location.unknown(ctx):878 invalid_op = create_invalid_operation()879 # Verify that we fallback to the generic printer for safety.880 # CHECK: b'"builtin.module"() ({\n^bb0:\n}, {\n}) : () -> ()\n'881 print(invalid_op.get_asm(binary=True))882 883 884# CHECK-LABEL: TEST: testCreateWithInvalidAttributes885@run886def testCreateWithInvalidAttributes():887 ctx = Context()888 with Location.unknown(ctx):889 try:890 Operation.create(891 "builtin.module", attributes={None: StringAttr.get("name")}892 )893 except Exception as e:894 # CHECK: Invalid attribute key (not a string) when attempting to create the operation "builtin.module"895 print(e)896 try:897 Operation.create("builtin.module", attributes={42: StringAttr.get("name")})898 except Exception as e:899 # CHECK: Invalid attribute key (not a string) when attempting to create the operation "builtin.module"900 print(e)901 try:902 Operation.create("builtin.module", attributes={"some_key": ctx})903 except Exception as e:904 # CHECK: Invalid attribute value for the key "some_key" when attempting to create the operation "builtin.module"905 print(e)906 try:907 Operation.create("builtin.module", attributes={"some_key": None})908 except Exception as e:909 # CHECK: Found an invalid (`None`?) attribute value for the key "some_key" when attempting to create the operation "builtin.module"910 print(e)911 912 913# CHECK-LABEL: TEST: testOperationName914@run915def testOperationName():916 ctx = Context()917 ctx.allow_unregistered_dialects = True918 module = Module.parse(919 r"""920 %0 = "custom.op1"() : () -> f32921 %1 = "custom.op2"() : () -> i32922 %2 = "custom.op1"() : () -> f32923 """,924 ctx,925 )926 927 # CHECK: custom.op1928 # CHECK: custom.op2929 # CHECK: custom.op1930 for op in module.body.operations:931 print(op.operation.name)932 933 934# CHECK-LABEL: TEST: testCapsuleConversions935@run936def testCapsuleConversions():937 ctx = Context()938 ctx.allow_unregistered_dialects = True939 with Location.unknown(ctx):940 m = Operation.create("custom.op1").operation941 m_capsule = m._CAPIPtr942 assert '"mlir.ir.Operation._CAPIPtr"' in repr(m_capsule)943 m2 = Operation._CAPICreate(m_capsule)944 assert m2 is not m945 assert m2 == m946 # Gc and verify destructed.947 m = None948 m_capsule = None949 m2 = None950 gc.collect()951 952 953# CHECK-LABEL: TEST: testOperationErase954@run955def testOperationErase():956 ctx = Context()957 ctx.allow_unregistered_dialects = True958 with Location.unknown(ctx):959 m = Module.create()960 with InsertionPoint(m.body):961 op = Operation.create("custom.op1")962 963 # CHECK: "custom.op1"964 print(m)965 966 op.operation.erase()967 968 # CHECK-NOT: "custom.op1"969 print(m)970 971 # Ensure we can create another operation972 Operation.create("custom.op2")973 974 975# CHECK-LABEL: TEST: testOperationClone976@run977def testOperationClone():978 ctx = Context()979 ctx.allow_unregistered_dialects = True980 with Location.unknown(ctx):981 m = Module.create()982 with InsertionPoint(m.body):983 op = Operation.create("custom.op1")984 985 # CHECK: "custom.op1"986 print(m)987 988 clone = op.operation.clone()989 op.operation.erase()990 991 # CHECK: "custom.op1"992 print(m)993 994 995# CHECK-LABEL: TEST: testOperationLoc996@run997def testOperationLoc():998 ctx = Context()999 ctx.allow_unregistered_dialects = True1000 with ctx:1001 loc = Location.name("loc")1002 op = Operation.create("custom.op", loc=loc)1003 assert op.location == loc1004 assert op.operation.location == loc1005 1006 another_loc = Location.name("another_loc")1007 op.location = another_loc1008 assert op.location == another_loc1009 assert op.operation.location == another_loc1010 # CHECK: loc("another_loc")1011 print(op.location)1012 1013 1014# CHECK-LABEL: TEST: testModuleMerge1015@run1016def testModuleMerge():1017 with Context():1018 m1 = Module.parse("func.func private @foo()")1019 m2 = Module.parse(1020 """1021 func.func private @bar()1022 func.func private @qux()1023 """1024 )1025 foo = m1.body.operations[0]1026 bar = m2.body.operations[0]1027 qux = m2.body.operations[1]1028 assert bar.is_before_in_block(qux)1029 bar.move_before(foo)1030 assert bar.is_before_in_block(foo)1031 qux.move_after(foo)1032 assert bar.is_before_in_block(qux)1033 assert foo.is_before_in_block(qux)1034 1035 # CHECK: module1036 # CHECK: func private @bar1037 # CHECK: func private @foo1038 # CHECK: func private @qux1039 print(m1)1040 1041 # CHECK: module {1042 # CHECK-NEXT: }1043 print(m2)1044 1045 1046# CHECK-LABEL: TEST: testAppendMoveFromAnotherBlock1047@run1048def testAppendMoveFromAnotherBlock():1049 with Context():1050 m1 = Module.parse("func.func private @foo()")1051 m2 = Module.parse("func.func private @bar()")1052 func = m1.body.operations[0]1053 m2.body.append(func)1054 1055 # CHECK: module1056 # CHECK: func private @bar1057 # CHECK: func private @foo1058 1059 print(m2)1060 # CHECK: module {1061 # CHECK-NEXT: }1062 print(m1)1063 1064 1065# CHECK-LABEL: TEST: testDetachFromParent1066@run1067def testDetachFromParent():1068 with Context():1069 m1 = Module.parse("func.func private @foo()")1070 func = m1.body.operations[0].detach_from_parent()1071 # CHECK: func.attached=False1072 print(f"{func.attached=}")1073 1074 try:1075 func.detach_from_parent()1076 except ValueError as e:1077 if "has no parent" not in str(e):1078 raise1079 else:1080 assert False, "expected ValueError when detaching a detached operation"1081 1082 print(m1)1083 # CHECK-NOT: func private @foo1084 1085 1086# CHECK-LABEL: TEST: testOperationHash1087@run1088def testOperationHash():1089 ctx = Context()1090 ctx.allow_unregistered_dialects = True1091 with ctx, Location.unknown():1092 op = Operation.create("custom.op1")1093 assert hash(op) == hash(op.operation)1094 1095 module = Module.create()1096 with InsertionPoint(module.body):1097 op2 = Operation.create("custom.op2")1098 custom_op2 = module.body.operations[0]1099 assert hash(op2) == hash(custom_op2)1100 1101 1102# CHECK-LABEL: TEST: testOperationParse1103@run1104def testOperationParse():1105 with Context() as ctx:1106 ctx.allow_unregistered_dialects = True1107 1108 # Generic operation parsing.1109 m = Operation.parse("module {}")1110 o = Operation.parse('"test.foo"() : () -> ()')1111 assert isinstance(m, ModuleOp)1112 assert type(o) is OpView1113 1114 # Parsing specific operation.1115 m = ModuleOp.parse("module {}")1116 assert isinstance(m, ModuleOp)1117 try:1118 ModuleOp.parse('"test.foo"() : () -> ()')1119 except MLIRError as e:1120 # CHECK: error: Expected a 'builtin.module' op, got: 'test.foo'1121 print(f"error: {e}")1122 else:1123 assert False, "expected error"1124 1125 o = Operation.parse('"test.foo"() : () -> ()', source_name="my-source-string")1126 # CHECK: op_with_source_name: "test.foo"() : () -> () loc("my-source-string":1:1)1127 print(1128 f"op_with_source_name: {o.get_asm(enable_debug_info=True, use_local_scope=True)}"1129 )1130 1131 1132# CHECK-LABEL: TEST: testOpWalk1133@run1134def testOpWalk():1135 ctx = Context()1136 ctx.allow_unregistered_dialects = True1137 module = Module.parse(1138 r"""1139 builtin.module {1140 func.func @f() {1141 func.return1142 }1143 }1144 """,1145 ctx,1146 )1147 1148 def callback(op):1149 print(op.name)1150 return WalkResult.ADVANCE1151 1152 # Test post-order walk (default).1153 # CHECK-NEXT: Post-order1154 # CHECK-NEXT: func.return1155 # CHECK-NEXT: func.func1156 # CHECK-NEXT: builtin.module1157 print("Post-order")1158 module.operation.walk(callback)1159 1160 # Test pre-order walk.1161 # CHECK-NEXT: Pre-order1162 # CHECK-NEXT: builtin.module1163 # CHECK-NEXT: func.fun1164 # CHECK-NEXT: func.return1165 print("Pre-order")1166 module.operation.walk(callback, WalkOrder.PRE_ORDER)1167 1168 # Test interrput.1169 # CHECK-NEXT: Interrupt post-order1170 # CHECK-NEXT: func.return1171 print("Interrupt post-order")1172 1173 def callback(op):1174 print(op.name)1175 return WalkResult.INTERRUPT1176 1177 module.operation.walk(callback)1178 1179 # Test skip.1180 # CHECK-NEXT: Skip pre-order1181 # CHECK-NEXT: builtin.module1182 print("Skip pre-order")1183 1184 def callback(op):1185 print(op.name)1186 return WalkResult.SKIP1187 1188 module.operation.walk(callback, WalkOrder.PRE_ORDER)1189 1190 # Test exception.1191 # CHECK: Exception1192 # CHECK-NEXT: func.return1193 # CHECK-NEXT: Exception raised1194 print("Exception")1195 1196 def callback(op):1197 print(op.name)1198 raise ValueError1199 return WalkResult.ADVANCE1200 1201 try:1202 module.operation.walk(callback)1203 except RuntimeError:1204 print("Exception raised")1205 1206 1207# CHECK-LABEL: TEST: testGetOwnerConcreteOpview1208@run1209def testGetOwnerConcreteOpview():1210 with Context() as ctx, Location.unknown():1211 module = Module.create()1212 with InsertionPoint(module.body):1213 a = arith.ConstantOp(value=42, result=IntegerType.get_signless(32))1214 r = arith.AddIOp(a, a, overflowFlags=arith.IntegerOverflowFlags.nsw)1215 for u in a.result.uses:1216 assert isinstance(u.owner, arith.AddIOp)1217 1218 1219# CHECK-LABEL: TEST: testIndexSwitch1220@run1221def testIndexSwitch():1222 with Context() as ctx, Location.unknown():1223 i32 = T.i32()1224 module = Module.create()1225 with InsertionPoint(module.body):1226 1227 @func.FuncOp.from_py_func(T.index())1228 def index_switch(index):1229 c1 = arith.constant(i32, 1)1230 switch_op = scf.IndexSwitchOp(results=[i32], arg=index, cases=range(3))1231 1232 assert len(switch_op.regions) == 41233 assert len(switch_op.regions[2:]) == 21234 assert len([i for i in switch_op.regions[2:]]) == 21235 assert len(switch_op.caseRegions) == 31236 assert len([i for i in switch_op.caseRegions]) == 31237 assert len(switch_op.caseRegions[1:]) == 21238 assert len([i for i in switch_op.caseRegions[1:]]) == 21239