1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4from mlir.ir import *
5
6
7def run(f):
8  print("\nTEST:", f.__name__)
9  f()
10  gc.collect()
11  assert Context._get_live_count() == 0
12  return f
13
14
15# CHECK-LABEL: TEST: testDialectDescriptor
16@run
17def testDialectDescriptor():
18  ctx = Context()
19  d = ctx.get_dialect_descriptor("func")
20  # CHECK: <DialectDescriptor func>
21  print(d)
22  # CHECK: func
23  print(d.namespace)
24  try:
25    _ = ctx.get_dialect_descriptor("not_existing")
26  except ValueError:
27    pass
28  else:
29    assert False, "Expected exception"
30
31
32# CHECK-LABEL: TEST: testUserDialectClass
33@run
34def testUserDialectClass():
35  ctx = Context()
36  # Access using attribute.
37  d = ctx.dialects.func
38  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
39  print(d)
40  try:
41    _ = ctx.dialects.not_existing
42  except AttributeError:
43    pass
44  else:
45    assert False, "Expected exception"
46
47  # Access using index.
48  d = ctx.dialects["func"]
49  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
50  print(d)
51  try:
52    _ = ctx.dialects["not_existing"]
53  except IndexError:
54    pass
55  else:
56    assert False, "Expected exception"
57
58  # Using the 'd' alias.
59  d = ctx.d["func"]
60  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
61  print(d)
62
63
64# CHECK-LABEL: TEST: testCustomOpView
65# This test uses the standard dialect AddFOp as an example of a user op.
66# TODO: Op creation and access is still quite verbose: simplify this test as
67# additional capabilities come online.
68@run
69def testCustomOpView():
70
71  def createInput():
72    op = Operation.create("pytest_dummy.intinput", results=[f32])
73    # TODO: Auto result cast from operation
74    return op.results[0]
75
76  with Context() as ctx, Location.unknown():
77    ctx.allow_unregistered_dialects = True
78    m = Module.create()
79
80    with InsertionPoint(m.body):
81      f32 = F32Type.get()
82      # Create via dialects context collection.
83      input1 = createInput()
84      input2 = createInput()
85      op1 = ctx.dialects.arith.AddFOp(input1, input2)
86
87      # Create via an import
88      from mlir.dialects.arith import AddFOp
89      AddFOp(input1, op1.result)
90
91  # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
92  # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
93  # CHECK: %[[R0:.*]] = arith.addf %[[INPUT0]], %[[INPUT1]] : f32
94  # CHECK: %[[R1:.*]] = arith.addf %[[INPUT0]], %[[R0]] : f32
95  m.operation.print()
96
97
98# CHECK-LABEL: TEST: testIsRegisteredOperation
99@run
100def testIsRegisteredOperation():
101  ctx = Context()
102
103  # CHECK: cf.cond_br: True
104  print(f"cf.cond_br: {ctx.is_registered_operation('cf.cond_br')}")
105  # CHECK: func.not_existing: False
106  print(f"func.not_existing: {ctx.is_registered_operation('func.not_existing')}")
107