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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
# RUN: %PYTHON %s | FileCheck %s
import gc
import sys
from mlir.ir import *
from mlir.dialects._ods_common import _cext
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')}")
# CHECK-LABEL: TEST: testAppendPrefixSearchPath
@run
def testAppendPrefixSearchPath():
ctx = Context()
ctx.allow_unregistered_dialects = True
with Location.unknown(ctx):
assert not _cext.globals._check_dialect_module_loaded("custom")
Operation.create("custom.op")
assert not _cext.globals._check_dialect_module_loaded("custom")
sys.path.append(".")
_cext.globals.append_dialect_search_prefix("custom_dialect")
assert _cext.globals._check_dialect_module_loaded("custom")
|
