Function test_dso_module_load

tests/python/runtime/test_runtime_module_load.py:49–97

(target)

@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@pytest.mark.parametrize("target", ["llvm", {"kind": "llvm", "jit": "mcjit"}])
def test_dso_module_load(target):
    dtype = "int64"
    temp = utils.tempdir()

    def save_object(names):
        n = te.size_var("n")
        Ab = tvm.tirx.decl_buffer((n,), dtype)
        i = te.var("i")
        # for i in 0 to n-1:
        stmt = tvm.tirx.For(
            i,
            0,
            n - 1,
            tvm.tirx.ForKind.SERIAL,
            tvm.tirx.BufferStore(Ab, tvm.tirx.BufferLoad(Ab, [i]) + 1, [i + 1]),
        )
        mod = tvm.IRModule.from_expr(
            tvm.tirx.PrimFunc([Ab], stmt).with_attr("global_symbol", "main")
        )
        m = tvm.tirx.build(mod, target=target)
        for name in names:
            m.write_to_file(name)

    path_obj = temp.relpath("test.o")
    path_ll = temp.relpath("test.ll")
    path_bc = temp.relpath("test.bc")
    path_dso = temp.relpath("test.so")
    save_object([path_obj, path_ll, path_bc])
    cc.create_shared(path_dso, [path_obj])

    f1 = tvm.runtime.load_module(path_dso)
    f2 = tvm.runtime.load_module(path_ll)
    a = tvm.runtime.tensor(np.zeros(10, dtype=dtype))
    f1(a)
    np.testing.assert_equal(a.numpy(), np.arange(a.shape[0]))
    a = tvm.runtime.tensor(np.zeros(10, dtype=dtype))
    f2(a)
    np.testing.assert_equal(a.numpy(), np.arange(a.shape[0]))

    path_runtime_py = temp.relpath("runtime.py")
    with open(path_runtime_py, "w") as fo:
        fo.write(runtime_py)

    proc = subprocess.run(
        [sys.executable, path_runtime_py, path_dso, dtype],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    assert proc.returncode == 0, f"{proc.args} exited with {proc.returncode}: {proc.stdout}"


@pytest.mark.gpu