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