171 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4from tempfile import NamedTemporaryFile5from mlir.ir import *6 7 8def run(f):9 print("\nTEST:", f.__name__)10 f()11 gc.collect()12 assert Context._get_live_count() == 013 return f14 15 16# Verify successful parse.17# CHECK-LABEL: TEST: testParseSuccess18# CHECK: module @successfulParse19@run20def testParseSuccess():21 ctx = Context()22 module = Module.parse(r"""module @successfulParse {}""", ctx)23 assert module.context is ctx24 print("CLEAR CONTEXT")25 ctx = None # Ensure that module captures the context.26 gc.collect()27 module.dump() # Just outputs to stderr. Verifies that it functions.28 print(str(module))29 30 31# Verify successful parse from file.32# CHECK-LABEL: TEST: testParseFromFileSuccess33# CHECK: module @successfulParse34@run35def testParseFromFileSuccess():36 ctx = Context()37 with NamedTemporaryFile(mode="w") as tmp_file:38 tmp_file.write(r"""module @successfulParse {}""")39 tmp_file.flush()40 module = Module.parseFile(tmp_file.name, ctx)41 assert module.context is ctx42 print("CLEAR CONTEXT")43 ctx = None # Ensure that module captures the context.44 gc.collect()45 module.operation.verify()46 print(str(module))47 48 49# Verify parse error.50# CHECK-LABEL: TEST: testParseError51# CHECK: testParseError: <52# CHECK: Unable to parse module assembly:53# CHECK: error: "-":1:1: expected operation name in quotes54# CHECK: >55@run56def testParseError():57 ctx = Context()58 try:59 module = Module.parse(r"""}SYNTAX ERROR{""", ctx)60 except MLIRError as e:61 print(f"testParseError: <{e}>")62 else:63 print("Exception not produced")64 65 66# Verify successful parse.67# CHECK-LABEL: TEST: testCreateEmpty68# CHECK: module {69@run70def testCreateEmpty():71 ctx = Context()72 loc = Location.unknown(ctx)73 module = Module.create(loc)74 print("CLEAR CONTEXT")75 ctx = None # Ensure that module captures the context.76 gc.collect()77 print(str(module))78 79 80# Verify round-trip of ASM that contains unicode.81# Note that this does not test that the print path converts unicode properly82# because MLIR asm always normalizes it to the hex encoding.83# CHECK-LABEL: TEST: testRoundtripUnicode84# CHECK: func private @roundtripUnicode()85# CHECK: foo = "\F0\9F\98\8A"86@run87def testRoundtripUnicode():88 ctx = Context()89 module = Module.parse(90 r"""91 func.func private @roundtripUnicode() attributes { foo = "😊" }92 """,93 ctx,94 )95 print(str(module))96 97 98# Verify round-trip of ASM that contains unicode.99# Note that this does not test that the print path converts unicode properly100# because MLIR asm always normalizes it to the hex encoding.101# CHECK-LABEL: TEST: testRoundtripBinary102# CHECK: func private @roundtripUnicode()103# CHECK: foo = "\F0\9F\98\8A"104@run105def testRoundtripBinary():106 with Context():107 module = Module.parse(108 r"""109 func.func private @roundtripUnicode() attributes { foo = "😊" }110 """111 )112 binary_asm = module.operation.get_asm(binary=True)113 assert isinstance(binary_asm, bytes)114 module = Module.parse(binary_asm)115 print(module)116 117 118# Tests that module.operation works and correctly interns instances.119# CHECK-LABEL: TEST: testModuleOperation120@run121def testModuleOperation():122 ctx = Context()123 module = Module.parse(r"""module @successfulParse {}""", ctx)124 assert ctx._get_live_module_count() == 1125 op1 = module.operation126 # CHECK: module @successfulParse127 print(op1)128 129 # Ensure that operations are the same on multiple calls.130 op2 = module.operation131 assert op1 is not op2132 assert op1 == op2133 134 # Test live operation clearing.135 op1 = module.operation136 op1 = None137 gc.collect()138 op1 = module.operation139 140 # Ensure that if module is de-referenced, the operations are still valid.141 module = None142 gc.collect()143 print(op1)144 145 # Collect and verify lifetime.146 op1 = None147 op2 = None148 gc.collect()149 assert ctx._get_live_module_count() == 0150 151 152# CHECK-LABEL: TEST: testModuleCapsule153@run154def testModuleCapsule():155 ctx = Context()156 module = Module.parse(r"""module @successfulParse {}""", ctx)157 assert ctx._get_live_module_count() == 1158 # CHECK: "mlir.ir.Module._CAPIPtr"159 module_capsule = module._CAPIPtr160 print(module_capsule)161 module_dup = Module._CAPICreate(module_capsule)162 assert module is module_dup163 assert module == module_dup164 assert module_dup.context is ctx165 # Gc and verify destructed.166 module = None167 module_capsule = None168 module_dup = None169 gc.collect()170 assert ctx._get_live_module_count() == 0171