| # Copyright 2022 The IREE Authors | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| import gc | |
| import unittest | |
| import iree.runtime as rt | |
| NONE_CTOR = lambda iface: None | |
| class PyModuleInterfaceTest(unittest.TestCase): | |
| def setUp(self): | |
| self._instance = rt.VmInstance() | |
| def testEmptyModuleLifecycle(self): | |
| iface = rt.PyModuleInterface("test1", NONE_CTOR) | |
| print(iface) | |
| self.assertFalse(iface.initialized) | |
| m = iface.create() | |
| print(iface) | |
| self.assertTrue(iface.initialized) | |
| print(m) | |
| m = None | |
| gc.collect() | |
| print(iface) | |
| self.assertTrue(iface.destroyed) | |
| def testEmptyModuleInstance(self): | |
| iface = rt.PyModuleInterface("test1", NONE_CTOR) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| self.assertTrue(iface.initialized) | |
| print(context) | |
| # Make sure no circular refs and that everything frees. | |
| context = None | |
| m = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testMultiModuleInstance(self): | |
| calls = [] | |
| def ctor(iface): | |
| calls.append(iface) | |
| return None | |
| iface = rt.PyModuleInterface("test1", ctor) | |
| m = iface.create() | |
| context1 = rt.VmContext(self._instance, modules=(m,)) | |
| self.assertTrue(iface.initialized) | |
| context2 = rt.VmContext(self._instance, modules=(m,)) | |
| self.assertTrue(iface.initialized) | |
| self.assertEqual(2, len(calls)) | |
| # Make sure no circular refs and that everything frees. | |
| calls = None | |
| context1 = None | |
| m = None | |
| context2 = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testVoidFunctionExport(self): | |
| messages = [] | |
| class Methods: | |
| def __init__(self, iface): | |
| self.iface = iface | |
| self.counter = 0 | |
| def say_hello(self): | |
| messages.append(f"Hello! Your number is {self.counter}") | |
| print(messages[-1]) | |
| self.counter += 1 | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("say_hello", "0v", Methods.say_hello) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| f = m.lookup_function("say_hello") | |
| self.assertIsNotNone(f) | |
| args = rt.VmVariantList(0) | |
| results = rt.VmVariantList(0) | |
| # Invoke twice - should produce two messages. | |
| context.invoke(f, args, results) | |
| context.invoke(f, args, results) | |
| self.assertListEqual( | |
| messages, | |
| [ | |
| "Hello! Your number is 0", | |
| "Hello! Your number is 1", | |
| ], | |
| ) | |
| # Make sure no circular refs and that everything frees. | |
| context = None | |
| m = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testPythonException(self): | |
| messages = [] | |
| class Methods: | |
| def __init__(self, iface): | |
| pass | |
| def do_it(self): | |
| raise ValueError("This is from Python") | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("do_it", "0v", Methods.do_it) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| f = m.lookup_function("do_it") | |
| self.assertIsNotNone(f) | |
| args = rt.VmVariantList(0) | |
| results = rt.VmVariantList(0) | |
| # We are testing here that the Python level exception is caught and | |
| # translated to an IREE status (surfacing as a RuntimeError) vs percolating | |
| # through the C call stack. | |
| with self.assertRaisesRegex(RuntimeError, "ValueError: This is from Python"): | |
| context.invoke(f, args, results) | |
| # Make sure no circular refs and that everything frees. | |
| context = None | |
| m = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testPrimitiveArguments(self): | |
| values = [] | |
| class Methods: | |
| def __init__(self, iface): | |
| pass | |
| def do_it(self, a, b): | |
| values.append((a, b)) | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("do_it_i32", "0ii", Methods.do_it) | |
| iface.export("do_it_i64", "0II", Methods.do_it) | |
| iface.export("do_it_f32", "0ff", Methods.do_it) | |
| iface.export("do_it_f64", "0FF", Methods.do_it) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| args = rt.VmVariantList(2) | |
| results = rt.VmVariantList(0) | |
| args.push_int(42) | |
| args.push_int(43) | |
| context.invoke(m.lookup_function("do_it_i32"), args, results) | |
| context.invoke(m.lookup_function("do_it_i64"), args, results) | |
| args = rt.VmVariantList(2) | |
| args.push_float(2.0) | |
| args.push_float(4.0) | |
| # TODO: Python doesn't have 32bit floats, so we are populating f64 args. | |
| # These are coming back as zeros, and I expected something to be | |
| # doing a conversion? The same is being done with i64 above but is | |
| # working there. | |
| context.invoke(m.lookup_function("do_it_f32"), args, results) | |
| context.invoke(m.lookup_function("do_it_f64"), args, results) | |
| print(values) | |
| self.assertEqual(repr(values), "[(42, 43), (42, 43), (0.0, 0.0), (2.0, 4.0)]") | |
| # Make sure no circular refs and that everything frees. | |
| context = None | |
| m = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testPrimitiveResults(self): | |
| next_results = None | |
| class Methods: | |
| def __init__(self, iface): | |
| pass | |
| def do_it(self): | |
| return next_results | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("do_it_i32", "0v_ii", Methods.do_it) | |
| iface.export("do_it_i64", "0v_II", Methods.do_it) | |
| iface.export("do_it_f32", "0v_ff", Methods.do_it) | |
| iface.export("do_it_f64", "0v_FF", Methods.do_it) | |
| iface.export("do_it_unary_i32", "0v_i", Methods.do_it) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| args = rt.VmVariantList(0) | |
| # i32 | |
| results = rt.VmVariantList(2) | |
| next_results = (42, 43) | |
| context.invoke(m.lookup_function("do_it_i32"), args, results) | |
| self.assertEqual(repr(results), "<VmVariantList(2): [42, 43]>") | |
| # i64 | |
| results = rt.VmVariantList(2) | |
| next_results = (42, 43) | |
| context.invoke(m.lookup_function("do_it_i64"), args, results) | |
| self.assertEqual(repr(results), "<VmVariantList(2): [42, 43]>") | |
| # f32 | |
| results = rt.VmVariantList(2) | |
| next_results = (2.0, 4.0) | |
| context.invoke(m.lookup_function("do_it_f32"), args, results) | |
| self.assertEqual(repr(results), "<VmVariantList(2): [2.000000, 4.000000]>") | |
| # f64 | |
| results = rt.VmVariantList(2) | |
| next_results = (2.0, 4.0) | |
| context.invoke(m.lookup_function("do_it_f64"), args, results) | |
| self.assertEqual(repr(results), "<VmVariantList(2): [2.000000, 4.000000]>") | |
| # Unary special case. | |
| results = rt.VmVariantList(1) | |
| next_results = 42 | |
| context.invoke(m.lookup_function("do_it_unary_i32"), args, results) | |
| self.assertEqual(repr(results), "<VmVariantList(1): [42]>") | |
| # Make sure no circular refs and that everything frees. | |
| context = None | |
| m = None | |
| gc.collect() | |
| self.assertTrue(iface.destroyed) | |
| def testRefArguments(self): | |
| values = [] | |
| class Methods: | |
| def __init__(self, iface): | |
| pass | |
| def do_it(self, a, b): | |
| values.append((a.deref(rt.VmVariantList), b.deref(rt.VmVariantList))) | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("do_it", "0rr", Methods.do_it) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| # These lists just happen to be reference objects we know how to | |
| # create. | |
| arg0 = rt.VmVariantList(1) | |
| arg0.push_int(42) | |
| arg1 = rt.VmVariantList(1) | |
| arg1.push_int(84) | |
| args = rt.VmVariantList(2) | |
| args.push_list(arg0) | |
| args.push_list(arg1) | |
| results = rt.VmVariantList(2) | |
| context.invoke(m.lookup_function("do_it"), args, results) | |
| print("REF VALUES:", values) | |
| self.assertEqual( | |
| repr(values), "[(<VmVariantList(1): [42]>, <VmVariantList(1): [84]>)]" | |
| ) | |
| def testRefResults(self): | |
| class Methods: | |
| def __init__(self, iface): | |
| pass | |
| def do_it(self): | |
| # These lists just happen to be reference objects we know how to | |
| # create. | |
| r0 = rt.VmVariantList(1) | |
| r0.push_int(42) | |
| r1 = rt.VmVariantList(1) | |
| r1.push_int(84) | |
| return r0.ref, r1.ref | |
| iface = rt.PyModuleInterface("test1", Methods) | |
| iface.export("do_it", "0v_rr", Methods.do_it) | |
| m = iface.create() | |
| context = rt.VmContext(self._instance, modules=(m,)) | |
| args = rt.VmVariantList(0) | |
| results = rt.VmVariantList(2) | |
| context.invoke(m.lookup_function("do_it"), args, results) | |
| print("REF RESULTS:", results) | |
| self.assertEqual(repr(results), "<VmVariantList(2): [List[42], List[84]]>") | |
| if __name__ == "__main__": | |
| unittest.main() |