summaryrefslogtreecommitdiff
path: root/mlir/test/python/ir/dialects.py
blob: 65e81e84354c22c29315d207459af73b747def5a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# RUN: %PYTHON %s | FileCheck %s

import gc
from mlir.ir import *


def run(f):
  print("\nTEST:", f.__name__)
  f()
  gc.collect()
  assert Context._get_live_count() == 0
  return f


# CHECK-LABEL: TEST: testDialectDescriptor
@run
def testDialectDescriptor():
  ctx = Context()
  d = ctx.get_dialect_descriptor("func")
  # CHECK: <DialectDescriptor func>
  print(d)
  # CHECK: func
  print(d.namespace)
  try:
    _ = ctx.get_dialect_descriptor("not_existing")
  except ValueError:
    pass
  else:
    assert False, "Expected exception"


# CHECK-LABEL: TEST: testUserDialectClass
@run
def testUserDialectClass():
  ctx = Context()
  # Access using attribute.
  d = ctx.dialects.func
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)
  try:
    _ = ctx.dialects.not_existing
  except AttributeError:
    pass
  else:
    assert False, "Expected exception"

  # Access using index.
  d = ctx.dialects["func"]
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)
  try:
    _ = ctx.dialects["not_existing"]
  except IndexError:
    pass
  else:
    assert False, "Expected exception"

  # Using the 'd' alias.
  d = ctx.d["func"]
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)


# CHECK-LABEL: TEST: testCustomOpView
# This test uses the standard dialect AddFOp as an example of a user op.
# TODO: Op creation and access is still quite verbose: simplify this test as
# additional capabilities come online.
@run
def testCustomOpView():

  def createInput():
    op = Operation.create("pytest_dummy.intinput", results=[f32])
    # TODO: Auto result cast from operation
    return op.results[0]

  with Context() as ctx, Location.unknown():
    ctx.allow_unregistered_dialects = True
    m = Module.create()

    with InsertionPoint(m.body):
      f32 = F32Type.get()
      # Create via dialects context collection.
      input1 = createInput()
      input2 = createInput()
      op1 = ctx.dialects.arith.AddFOp(input1, input2)

      # Create via an import
      from mlir.dialects.arith import AddFOp
      AddFOp(input1, op1.result)

  # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[R0:.*]] = arith.addf %[[INPUT0]], %[[INPUT1]] : f32
  # CHECK: %[[R1:.*]] = arith.addf %[[INPUT0]], %[[R0]] : f32
  m.operation.print()


# CHECK-LABEL: TEST: testIsRegisteredOperation
@run
def testIsRegisteredOperation():
  ctx = Context()

  # CHECK: cf.cond_br: True
  print(f"cf.cond_br: {ctx.is_registered_operation('cf.cond_br')}")
  # CHECK: func.not_existing: False
  print(f"func.not_existing: {ctx.is_registered_operation('func.not_existing')}")