Function test_executable_integration

tests/python/runtime/test_executable.py:175–228

Integration test for Executable with a simple TVM module.

()

def test_executable_integration():
    """Integration test for Executable with a simple TVM module."""
    # Create target and build
    target = tvm.target.Target("llvm")
    lib = tvm.tirx.build(MyModule, target=target)

    # Create an executable
    executable = Executable(lib)

    # Test jit
    jitted_mod = executable.jit()
    assert jitted_mod is not None

    # Test __getitem__
    add_func = executable["add"]
    assert add_func is not None

    # Test the function works
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    add_func(a, b, c)

    # Check results
    tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))

    # Test export_library
    temp_dir = tempfile.mkdtemp()
    try:
        lib_path = os.path.join(temp_dir, "test_lib.so")
        executable.export_library(lib_path)

        # Verify the library was created
        assert os.path.exists(lib_path)

        # Load the library back
        loaded_mod = tvm.runtime.load_module(lib_path)
        assert loaded_mod is not None

        # Test the loaded module
        loaded_add = loaded_mod["add"]
        c_loaded = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
        loaded_add(a, b, c_loaded)

        # Check results
        tvm.testing.assert_allclose(c_loaded.numpy(), np.array([3.0] * 10, dtype="float32"))

    finally:
        # Clean up
        if os.path.exists(temp_dir):
            import shutil

            shutil.rmtree(temp_dir)


def test_executable_jit_force_recompile():