1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4import io
5import itertools
6from mlir.ir import *
7
8def run(f):
9  print("\nTEST:", f.__name__)
10  f()
11  gc.collect()
12  assert Context._get_live_count() == 0
13
14
15# Verify iterator based traversal of the op/region/block hierarchy.
16# CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
17def testTraverseOpRegionBlockIterators():
18  ctx = Context()
19  ctx.allow_unregistered_dialects = True
20  module = Module.parse(r"""
21    func @f1(%arg0: i32) -> i32 {
22      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
23      return %1 : i32
24    }
25  """, ctx)
26  op = module.operation
27  assert op.context is ctx
28  # Get the block using iterators off of the named collections.
29  regions = list(op.regions)
30  blocks = list(regions[0].blocks)
31  # CHECK: MODULE REGIONS=1 BLOCKS=1
32  print(f"MODULE REGIONS={len(regions)} BLOCKS={len(blocks)}")
33
34  # Should verify.
35  # CHECK: .verify = True
36  print(f".verify = {module.operation.verify()}")
37
38  # Get the regions and blocks from the default collections.
39  default_regions = list(op)
40  default_blocks = list(default_regions[0])
41  # They should compare equal regardless of how obtained.
42  assert default_regions == regions
43  assert default_blocks == blocks
44
45  # Should be able to get the operations from either the named collection
46  # or the block.
47  operations = list(blocks[0].operations)
48  default_operations = list(blocks[0])
49  assert default_operations == operations
50
51  def walk_operations(indent, op):
52    for i, region in enumerate(op):
53      print(f"{indent}REGION {i}:")
54      for j, block in enumerate(region):
55        print(f"{indent}  BLOCK {j}:")
56        for k, child_op in enumerate(block):
57          print(f"{indent}    OP {k}: {child_op}")
58          walk_operations(indent + "      ", child_op)
59
60  # CHECK: REGION 0:
61  # CHECK:   BLOCK 0:
62  # CHECK:     OP 0: func
63  # CHECK:       REGION 0:
64  # CHECK:         BLOCK 0:
65  # CHECK:           OP 0: %0 = "custom.addi"
66  # CHECK:           OP 1: return
67  walk_operations("", op)
68
69run(testTraverseOpRegionBlockIterators)
70
71
72# Verify index based traversal of the op/region/block hierarchy.
73# CHECK-LABEL: TEST: testTraverseOpRegionBlockIndices
74def testTraverseOpRegionBlockIndices():
75  ctx = Context()
76  ctx.allow_unregistered_dialects = True
77  module = Module.parse(r"""
78    func @f1(%arg0: i32) -> i32 {
79      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
80      return %1 : i32
81    }
82  """, ctx)
83
84  def walk_operations(indent, op):
85    for i in range(len(op.regions)):
86      region = op.regions[i]
87      print(f"{indent}REGION {i}:")
88      for j in range(len(region.blocks)):
89        block = region.blocks[j]
90        print(f"{indent}  BLOCK {j}:")
91        for k in range(len(block.operations)):
92          child_op = block.operations[k]
93          print(f"{indent}    OP {k}: {child_op}")
94          print(f"{indent}    OP {k}: parent {child_op.operation.parent.name}")
95          walk_operations(indent + "      ", child_op)
96
97  # CHECK: REGION 0:
98  # CHECK:   BLOCK 0:
99  # CHECK:     OP 0: func
100  # CHECK:     OP 0: parent builtin.module
101  # CHECK:       REGION 0:
102  # CHECK:         BLOCK 0:
103  # CHECK:           OP 0: %0 = "custom.addi"
104  # CHECK:           OP 0: parent builtin.func
105  # CHECK:           OP 1: return
106  # CHECK:           OP 1: parent builtin.func
107  walk_operations("", module.operation)
108
109run(testTraverseOpRegionBlockIndices)
110
111
112# CHECK-LABEL: TEST: testBlockArgumentList
113def testBlockArgumentList():
114  with Context() as ctx:
115    module = Module.parse(r"""
116      func @f1(%arg0: i32, %arg1: f64, %arg2: index) {
117        return
118      }
119    """, ctx)
120    func = module.body.operations[0]
121    entry_block = func.regions[0].blocks[0]
122    assert len(entry_block.arguments) == 3
123    # CHECK: Argument 0, type i32
124    # CHECK: Argument 1, type f64
125    # CHECK: Argument 2, type index
126    for arg in entry_block.arguments:
127      print(f"Argument {arg.arg_number}, type {arg.type}")
128      new_type = IntegerType.get_signless(8 * (arg.arg_number + 1))
129      arg.set_type(new_type)
130
131    # CHECK: Argument 0, type i8
132    # CHECK: Argument 1, type i16
133    # CHECK: Argument 2, type i24
134    for arg in entry_block.arguments:
135      print(f"Argument {arg.arg_number}, type {arg.type}")
136
137    # Check that slicing works for block argument lists.
138    # CHECK: Argument 1, type i16
139    # CHECK: Argument 2, type i24
140    for arg in entry_block.arguments[1:]:
141      print(f"Argument {arg.arg_number}, type {arg.type}")
142
143    # Check that we can concatenate slices of argument lists.
144    # CHECK: Length: 4
145    print("Length: ",
146          len(entry_block.arguments[:2] + entry_block.arguments[1:]))
147
148    # CHECK: Type: i8
149    # CHECK: Type: i16
150    # CHECK: Type: i24
151    for t in entry_block.arguments.types:
152      print("Type: ", t)
153
154
155run(testBlockArgumentList)
156
157
158# CHECK-LABEL: TEST: testOperationOperands
159def testOperationOperands():
160  with Context() as ctx:
161    ctx.allow_unregistered_dialects = True
162    module = Module.parse(r"""
163      func @f1(%arg0: i32) {
164        %0 = "test.producer"() : () -> i64
165        "test.consumer"(%arg0, %0) : (i32, i64) -> ()
166        return
167      }""")
168    func = module.body.operations[0]
169    entry_block = func.regions[0].blocks[0]
170    consumer = entry_block.operations[1]
171    assert len(consumer.operands) == 2
172    # CHECK: Operand 0, type i32
173    # CHECK: Operand 1, type i64
174    for i, operand in enumerate(consumer.operands):
175      print(f"Operand {i}, type {operand.type}")
176
177
178run(testOperationOperands)
179
180
181# CHECK-LABEL: TEST: testOperationOperandsSlice
182def testOperationOperandsSlice():
183  with Context() as ctx:
184    ctx.allow_unregistered_dialects = True
185    module = Module.parse(r"""
186      func @f1() {
187        %0 = "test.producer0"() : () -> i64
188        %1 = "test.producer1"() : () -> i64
189        %2 = "test.producer2"() : () -> i64
190        %3 = "test.producer3"() : () -> i64
191        %4 = "test.producer4"() : () -> i64
192        "test.consumer"(%0, %1, %2, %3, %4) : (i64, i64, i64, i64, i64) -> ()
193        return
194      }""")
195    func = module.body.operations[0]
196    entry_block = func.regions[0].blocks[0]
197    consumer = entry_block.operations[5]
198    assert len(consumer.operands) == 5
199    for left, right in zip(consumer.operands, consumer.operands[::-1][::-1]):
200      assert left == right
201
202    # CHECK: test.producer0
203    # CHECK: test.producer1
204    # CHECK: test.producer2
205    # CHECK: test.producer3
206    # CHECK: test.producer4
207    full_slice = consumer.operands[:]
208    for operand in full_slice:
209      print(operand)
210
211    # CHECK: test.producer0
212    # CHECK: test.producer1
213    first_two = consumer.operands[0:2]
214    for operand in first_two:
215      print(operand)
216
217    # CHECK: test.producer3
218    # CHECK: test.producer4
219    last_two = consumer.operands[3:]
220    for operand in last_two:
221      print(operand)
222
223    # CHECK: test.producer0
224    # CHECK: test.producer2
225    # CHECK: test.producer4
226    even = consumer.operands[::2]
227    for operand in even:
228      print(operand)
229
230    # CHECK: test.producer2
231    fourth = consumer.operands[::2][1::2]
232    for operand in fourth:
233      print(operand)
234
235
236run(testOperationOperandsSlice)
237
238
239# CHECK-LABEL: TEST: testOperationOperandsSet
240def testOperationOperandsSet():
241  with Context() as ctx, Location.unknown(ctx):
242    ctx.allow_unregistered_dialects = True
243    module = Module.parse(r"""
244      func @f1() {
245        %0 = "test.producer0"() : () -> i64
246        %1 = "test.producer1"() : () -> i64
247        %2 = "test.producer2"() : () -> i64
248        "test.consumer"(%0) : (i64) -> ()
249        return
250      }""")
251    func = module.body.operations[0]
252    entry_block = func.regions[0].blocks[0]
253    producer1 = entry_block.operations[1]
254    producer2 = entry_block.operations[2]
255    consumer = entry_block.operations[3]
256    assert len(consumer.operands) == 1
257    type = consumer.operands[0].type
258
259    # CHECK: test.producer1
260    consumer.operands[0] = producer1.result
261    print(consumer.operands[0])
262
263    # CHECK: test.producer2
264    consumer.operands[-1] = producer2.result
265    print(consumer.operands[0])
266
267
268run(testOperationOperandsSet)
269
270
271# CHECK-LABEL: TEST: testDetachedOperation
272def testDetachedOperation():
273  ctx = Context()
274  ctx.allow_unregistered_dialects = True
275  with Location.unknown(ctx):
276    i32 = IntegerType.get_signed(32)
277    op1 = Operation.create(
278        "custom.op1", results=[i32, i32], regions=1, attributes={
279            "foo": StringAttr.get("foo_value"),
280            "bar": StringAttr.get("bar_value"),
281        })
282    # CHECK: %0:2 = "custom.op1"() ( {
283    # CHECK: }) {bar = "bar_value", foo = "foo_value"} : () -> (si32, si32)
284    print(op1)
285
286  # TODO: Check successors once enough infra exists to do it properly.
287
288run(testDetachedOperation)
289
290
291# CHECK-LABEL: TEST: testOperationInsertionPoint
292def testOperationInsertionPoint():
293  ctx = Context()
294  ctx.allow_unregistered_dialects = True
295  module = Module.parse(r"""
296    func @f1(%arg0: i32) -> i32 {
297      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
298      return %1 : i32
299    }
300  """, ctx)
301
302  # Create test op.
303  with Location.unknown(ctx):
304    op1 = Operation.create("custom.op1")
305    op2 = Operation.create("custom.op2")
306
307    func = module.body.operations[0]
308    entry_block = func.regions[0].blocks[0]
309    ip = InsertionPoint.at_block_begin(entry_block)
310    ip.insert(op1)
311    ip.insert(op2)
312    # CHECK: func @f1
313    # CHECK: "custom.op1"()
314    # CHECK: "custom.op2"()
315    # CHECK: %0 = "custom.addi"
316    print(module)
317
318  # Trying to add a previously added op should raise.
319  try:
320    ip.insert(op1)
321  except ValueError:
322    pass
323  else:
324    assert False, "expected insert of attached op to raise"
325
326run(testOperationInsertionPoint)
327
328
329# CHECK-LABEL: TEST: testOperationWithRegion
330def testOperationWithRegion():
331  ctx = Context()
332  ctx.allow_unregistered_dialects = True
333  with Location.unknown(ctx):
334    i32 = IntegerType.get_signed(32)
335    op1 = Operation.create("custom.op1", regions=1)
336    block = op1.regions[0].blocks.append(i32, i32)
337    # CHECK: "custom.op1"() ( {
338    # CHECK: ^bb0(%arg0: si32, %arg1: si32):  // no predecessors
339    # CHECK:   "custom.terminator"() : () -> ()
340    # CHECK: }) : () -> ()
341    terminator = Operation.create("custom.terminator")
342    ip = InsertionPoint(block)
343    ip.insert(terminator)
344    print(op1)
345
346    # Now add the whole operation to another op.
347    # TODO: Verify lifetime hazard by nulling out the new owning module and
348    # accessing op1.
349    # TODO: Also verify accessing the terminator once both parents are nulled
350    # out.
351    module = Module.parse(r"""
352      func @f1(%arg0: i32) -> i32 {
353        %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
354        return %1 : i32
355      }
356    """)
357    func = module.body.operations[0]
358    entry_block = func.regions[0].blocks[0]
359    ip = InsertionPoint.at_block_begin(entry_block)
360    ip.insert(op1)
361    # CHECK: func @f1
362    # CHECK: "custom.op1"()
363    # CHECK:   "custom.terminator"
364    # CHECK: %0 = "custom.addi"
365    print(module)
366
367run(testOperationWithRegion)
368
369
370# CHECK-LABEL: TEST: testOperationResultList
371def testOperationResultList():
372  ctx = Context()
373  module = Module.parse(r"""
374    func @f1() {
375      %0:3 = call @f2() : () -> (i32, f64, index)
376      return
377    }
378    func private @f2() -> (i32, f64, index)
379  """, ctx)
380  caller = module.body.operations[0]
381  call = caller.regions[0].blocks[0].operations[0]
382  assert len(call.results) == 3
383  # CHECK: Result 0, type i32
384  # CHECK: Result 1, type f64
385  # CHECK: Result 2, type index
386  for res in call.results:
387    print(f"Result {res.result_number}, type {res.type}")
388
389  # CHECK: Result type i32
390  # CHECK: Result type f64
391  # CHECK: Result type index
392  for t in call.results.types:
393    print(f"Result type {t}")
394
395
396run(testOperationResultList)
397
398
399# CHECK-LABEL: TEST: testOperationResultListSlice
400def testOperationResultListSlice():
401  with Context() as ctx:
402    ctx.allow_unregistered_dialects = True
403    module = Module.parse(r"""
404      func @f1() {
405        "some.op"() : () -> (i1, i2, i3, i4, i5)
406        return
407      }
408    """)
409    func = module.body.operations[0]
410    entry_block = func.regions[0].blocks[0]
411    producer = entry_block.operations[0]
412
413    assert len(producer.results) == 5
414    for left, right in zip(producer.results, producer.results[::-1][::-1]):
415      assert left == right
416      assert left.result_number == right.result_number
417
418    # CHECK: Result 0, type i1
419    # CHECK: Result 1, type i2
420    # CHECK: Result 2, type i3
421    # CHECK: Result 3, type i4
422    # CHECK: Result 4, type i5
423    full_slice = producer.results[:]
424    for res in full_slice:
425      print(f"Result {res.result_number}, type {res.type}")
426
427    # CHECK: Result 1, type i2
428    # CHECK: Result 2, type i3
429    # CHECK: Result 3, type i4
430    middle = producer.results[1:4]
431    for res in middle:
432      print(f"Result {res.result_number}, type {res.type}")
433
434    # CHECK: Result 1, type i2
435    # CHECK: Result 3, type i4
436    odd = producer.results[1::2]
437    for res in odd:
438      print(f"Result {res.result_number}, type {res.type}")
439
440    # CHECK: Result 3, type i4
441    # CHECK: Result 1, type i2
442    inverted_middle = producer.results[-2:0:-2]
443    for res in inverted_middle:
444      print(f"Result {res.result_number}, type {res.type}")
445
446
447run(testOperationResultListSlice)
448
449
450# CHECK-LABEL: TEST: testOperationAttributes
451def testOperationAttributes():
452  ctx = Context()
453  ctx.allow_unregistered_dialects = True
454  module = Module.parse(r"""
455    "some.op"() { some.attribute = 1 : i8,
456                  other.attribute = 3.0,
457                  dependent = "text" } : () -> ()
458  """, ctx)
459  op = module.body.operations[0]
460  assert len(op.attributes) == 3
461  iattr = IntegerAttr(op.attributes["some.attribute"])
462  fattr = FloatAttr(op.attributes["other.attribute"])
463  sattr = StringAttr(op.attributes["dependent"])
464  # CHECK: Attribute type i8, value 1
465  print(f"Attribute type {iattr.type}, value {iattr.value}")
466  # CHECK: Attribute type f64, value 3.0
467  print(f"Attribute type {fattr.type}, value {fattr.value}")
468  # CHECK: Attribute value text
469  print(f"Attribute value {sattr.value}")
470
471  # We don't know in which order the attributes are stored.
472  # CHECK-DAG: NamedAttribute(dependent="text")
473  # CHECK-DAG: NamedAttribute(other.attribute=3.000000e+00 : f64)
474  # CHECK-DAG: NamedAttribute(some.attribute=1 : i8)
475  for attr in op.attributes:
476    print(str(attr))
477
478  # Check that exceptions are raised as expected.
479  try:
480    op.attributes["does_not_exist"]
481  except KeyError:
482    pass
483  else:
484    assert False, "expected KeyError on accessing a non-existent attribute"
485
486  try:
487    op.attributes[42]
488  except IndexError:
489    pass
490  else:
491    assert False, "expected IndexError on accessing an out-of-bounds attribute"
492
493
494run(testOperationAttributes)
495
496
497# CHECK-LABEL: TEST: testOperationPrint
498def testOperationPrint():
499  ctx = Context()
500  module = Module.parse(r"""
501    func @f1(%arg0: i32) -> i32 {
502      %0 = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
503      return %arg0 : i32
504    }
505  """, ctx)
506
507  # Test print to stdout.
508  # CHECK: return %arg0 : i32
509  module.operation.print()
510
511  # Test print to text file.
512  f = io.StringIO()
513  # CHECK: <class 'str'>
514  # CHECK: return %arg0 : i32
515  module.operation.print(file=f)
516  str_value = f.getvalue()
517  print(str_value.__class__)
518  print(f.getvalue())
519
520  # Test print to binary file.
521  f = io.BytesIO()
522  # CHECK: <class 'bytes'>
523  # CHECK: return %arg0 : i32
524  module.operation.print(file=f, binary=True)
525  bytes_value = f.getvalue()
526  print(bytes_value.__class__)
527  print(bytes_value)
528
529  # Test get_asm with options.
530  # CHECK: value = opaque<"_", "0xDEADBEEF"> : tensor<4xi32>
531  # CHECK: "std.return"(%arg0) : (i32) -> () -:4:7
532  module.operation.print(large_elements_limit=2, enable_debug_info=True,
533      pretty_debug_info=True, print_generic_op_form=True, use_local_scope=True)
534
535run(testOperationPrint)
536
537
538# CHECK-LABEL: TEST: testKnownOpView
539def testKnownOpView():
540  with Context(), Location.unknown():
541    Context.current.allow_unregistered_dialects = True
542    module = Module.parse(r"""
543      %1 = "custom.f32"() : () -> f32
544      %2 = "custom.f32"() : () -> f32
545      %3 = addf %1, %2 : f32
546    """)
547    print(module)
548
549    # addf should map to a known OpView class in the std dialect.
550    # We know the OpView for it defines an 'lhs' attribute.
551    addf = module.body.operations[2]
552    # CHECK: <mlir.dialects._std_ops_gen._AddFOp object
553    print(repr(addf))
554    # CHECK: "custom.f32"()
555    print(addf.lhs)
556
557    # One of the custom ops should resolve to the default OpView.
558    custom = module.body.operations[0]
559    # CHECK: OpView object
560    print(repr(custom))
561
562    # Check again to make sure negative caching works.
563    custom = module.body.operations[0]
564    # CHECK: OpView object
565    print(repr(custom))
566
567run(testKnownOpView)
568
569
570# CHECK-LABEL: TEST: testSingleResultProperty
571def testSingleResultProperty():
572  with Context(), Location.unknown():
573    Context.current.allow_unregistered_dialects = True
574    module = Module.parse(r"""
575      "custom.no_result"() : () -> ()
576      %0:2 = "custom.two_result"() : () -> (f32, f32)
577      %1 = "custom.one_result"() : () -> f32
578    """)
579    print(module)
580
581  try:
582    module.body.operations[0].result
583  except ValueError as e:
584    # CHECK: Cannot call .result on operation custom.no_result which has 0 results
585    print(e)
586  else:
587    assert False, "Expected exception"
588
589  try:
590    module.body.operations[1].result
591  except ValueError as e:
592    # CHECK: Cannot call .result on operation custom.two_result which has 2 results
593    print(e)
594  else:
595    assert False, "Expected exception"
596
597  # CHECK: %1 = "custom.one_result"() : () -> f32
598  print(module.body.operations[2])
599
600run(testSingleResultProperty)
601
602# CHECK-LABEL: TEST: testPrintInvalidOperation
603def testPrintInvalidOperation():
604  ctx = Context()
605  with Location.unknown(ctx):
606    module = Operation.create("builtin.module", regions=2)
607    # This module has two region and is invalid verify that we fallback
608    # to the generic printer for safety.
609    block = module.regions[0].blocks.append()
610    # CHECK: // Verification failed, printing generic form
611    # CHECK: "builtin.module"() ( {
612    # CHECK: }) : () -> ()
613    print(module)
614    # CHECK: .verify = False
615    print(f".verify = {module.operation.verify()}")
616run(testPrintInvalidOperation)
617
618
619# CHECK-LABEL: TEST: testCreateWithInvalidAttributes
620def testCreateWithInvalidAttributes():
621  ctx = Context()
622  with Location.unknown(ctx):
623    try:
624      Operation.create(
625          "builtin.module", attributes={None: StringAttr.get("name")})
626    except Exception as e:
627      # CHECK: Invalid attribute key (not a string) when attempting to create the operation "builtin.module"
628      print(e)
629    try:
630      Operation.create(
631          "builtin.module", attributes={42: StringAttr.get("name")})
632    except Exception as e:
633      # CHECK: Invalid attribute key (not a string) when attempting to create the operation "builtin.module"
634      print(e)
635    try:
636      Operation.create("builtin.module", attributes={"some_key": ctx})
637    except Exception as e:
638      # CHECK: Invalid attribute value for the key "some_key" when attempting to create the operation "builtin.module"
639      print(e)
640    try:
641      Operation.create("builtin.module", attributes={"some_key": None})
642    except Exception as e:
643      # CHECK: Found an invalid (`None`?) attribute value for the key "some_key" when attempting to create the operation "builtin.module"
644      print(e)
645run(testCreateWithInvalidAttributes)
646
647
648# CHECK-LABEL: TEST: testOperationName
649def testOperationName():
650  ctx = Context()
651  ctx.allow_unregistered_dialects = True
652  module = Module.parse(r"""
653    %0 = "custom.op1"() : () -> f32
654    %1 = "custom.op2"() : () -> i32
655    %2 = "custom.op1"() : () -> f32
656  """, ctx)
657
658  # CHECK: custom.op1
659  # CHECK: custom.op2
660  # CHECK: custom.op1
661  for op in module.body.operations:
662    print(op.operation.name)
663
664run(testOperationName)
665
666# CHECK-LABEL: TEST: testCapsuleConversions
667def testCapsuleConversions():
668  ctx = Context()
669  ctx.allow_unregistered_dialects = True
670  with Location.unknown(ctx):
671    m = Operation.create("custom.op1").operation
672    m_capsule = m._CAPIPtr
673    assert '"mlir.ir.Operation._CAPIPtr"' in repr(m_capsule)
674    m2 = Operation._CAPICreate(m_capsule)
675    assert m2 is m
676
677run(testCapsuleConversions)
678
679# CHECK-LABEL: TEST: testOperationErase
680def testOperationErase():
681  ctx = Context()
682  ctx.allow_unregistered_dialects = True
683  with Location.unknown(ctx):
684    m = Module.create()
685    with InsertionPoint(m.body):
686      op = Operation.create("custom.op1")
687
688      # CHECK: "custom.op1"
689      print(m)
690
691      op.operation.erase()
692
693      # CHECK-NOT: "custom.op1"
694      print(m)
695
696      # Ensure we can create another operation
697      Operation.create("custom.op2")
698
699run(testOperationErase)
700