1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4from mlir.ir import *
5
6def run(f):
7  print("\nTEST:", f.__name__)
8  f()
9  gc.collect()
10  assert Context._get_live_count() == 0
11
12
13# Verify successful parse.
14# CHECK-LABEL: TEST: testParseSuccess
15# CHECK: module @successfulParse
16def testParseSuccess():
17  ctx = Context()
18  module = Module.parse(r"""module @successfulParse {}""", ctx)
19  assert module.context is ctx
20  print("CLEAR CONTEXT")
21  ctx = None  # Ensure that module captures the context.
22  gc.collect()
23  module.dump()  # Just outputs to stderr. Verifies that it functions.
24  print(str(module))
25
26run(testParseSuccess)
27
28
29# Verify parse error.
30# CHECK-LABEL: TEST: testParseError
31# CHECK: testParseError: Unable to parse module assembly (see diagnostics)
32def testParseError():
33  ctx = Context()
34  try:
35    module = Module.parse(r"""}SYNTAX ERROR{""", ctx)
36  except ValueError as e:
37    print("testParseError:", e)
38  else:
39    print("Exception not produced")
40
41run(testParseError)
42
43
44# Verify successful parse.
45# CHECK-LABEL: TEST: testCreateEmpty
46# CHECK: module {
47def testCreateEmpty():
48  ctx = Context()
49  loc = Location.unknown(ctx)
50  module = Module.create(loc)
51  print("CLEAR CONTEXT")
52  ctx = None  # Ensure that module captures the context.
53  gc.collect()
54  print(str(module))
55
56run(testCreateEmpty)
57
58
59# Verify round-trip of ASM that contains unicode.
60# Note that this does not test that the print path converts unicode properly
61# because MLIR asm always normalizes it to the hex encoding.
62# CHECK-LABEL: TEST: testRoundtripUnicode
63# CHECK: func private @roundtripUnicode()
64# CHECK: foo = "\F0\9F\98\8A"
65def testRoundtripUnicode():
66  ctx = Context()
67  module = Module.parse(r"""
68    func private @roundtripUnicode() attributes { foo = "��" }
69  """, ctx)
70  print(str(module))
71
72run(testRoundtripUnicode)
73
74
75# Tests that module.operation works and correctly interns instances.
76# CHECK-LABEL: TEST: testModuleOperation
77def testModuleOperation():
78  ctx = Context()
79  module = Module.parse(r"""module @successfulParse {}""", ctx)
80  assert ctx._get_live_module_count() == 1
81  op1 = module.operation
82  assert ctx._get_live_operation_count() == 1
83  # CHECK: module @successfulParse
84  print(op1)
85
86  # Ensure that operations are the same on multiple calls.
87  op2 = module.operation
88  assert ctx._get_live_operation_count() == 1
89  assert op1 is op2
90
91  # Ensure that if module is de-referenced, the operations are still valid.
92  module = None
93  gc.collect()
94  print(op1)
95
96  # Collect and verify lifetime.
97  op1 = None
98  op2 = None
99  gc.collect()
100  print("LIVE OPERATIONS:", ctx._get_live_operation_count())
101  assert ctx._get_live_operation_count() == 0
102  assert ctx._get_live_module_count() == 0
103
104run(testModuleOperation)
105
106
107# CHECK-LABEL: TEST: testModuleCapsule
108def testModuleCapsule():
109  ctx = Context()
110  module = Module.parse(r"""module @successfulParse {}""", ctx)
111  assert ctx._get_live_module_count() == 1
112  # CHECK: "mlir.ir.Module._CAPIPtr"
113  module_capsule = module._CAPIPtr
114  print(module_capsule)
115  module_dup = Module._CAPICreate(module_capsule)
116  assert module is module_dup
117  assert module_dup.context is ctx
118  # Gc and verify destructed.
119  module = None
120  module_capsule = None
121  module_dup = None
122  gc.collect()
123  assert ctx._get_live_module_count() == 0
124
125
126run(testModuleCapsule)
127