Function test_multi_systemlib

tests/python/relax/test_vm_build.py:992–1060

(exec_mode)

def test_multi_systemlib(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by popen_pool.PopenWorker

    @tvm.script.ir_module
    class ModA:
        I.module_attrs({"system_lib_prefix": "libA_"})

        @T.prim_func(s_tir=True)
        def tir_init(x_handle: T.handle):
            N = T.int64()
            x = T.match_buffer(x_handle, [N], "float32")
            for i in range(N):
                x[i] = T.float32(0)

        @R.function
        def main(s: R.Shape(["m"])) -> R.Tensor:
            m = T.int64()
            gv0 = R.call_tir(ModA.tir_init, (), R.Tensor((m + 1,), dtype="float32"))
            return gv0

    @tvm.script.ir_module
    class ModB:
        I.module_attrs({"system_lib_prefix": "libB_"})

        @T.prim_func(s_tir=True)
        def tir_init(x_handle: T.handle):
            N = T.int64()
            x = T.match_buffer(x_handle, [N], "float32")
            for i in range(N):
                x[i] = T.float32(1)

        @R.function
        def main(s: R.Shape(["m"])) -> R.Tensor:
            m = T.int64()
            gv0 = R.call_tir(ModB.tir_init, (), R.Tensor((m,), dtype="float32"))
            return gv0

    target = tvm.target.Target("llvm", host="llvm")
    libA = relax.build(ModA, target, exec_mode=exec_mode)
    libB = relax.build(ModB, target, exec_mode=exec_mode)

    temp = utils.tempdir()
    pathA = temp.relpath("libA.a")
    pathB = temp.relpath("libB.a")
    path_dso = temp.relpath("mylibAll.so")
    libA.export_library(pathA, fcompile=cc.create_staticlib)
    libB.export_library(pathB, fcompile=cc.create_staticlib)

    # package two static libs together
    # check that they do not interfere with each other
    # even though they have shared global var names
    # intentionally craft same gvar function with different behaviors
    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathA, pathB, "-Wl,--no-whole-archive"])

    def popen_check():
        # Load dll, will trigger system library registration
        ctypes.CDLL(path_dso)
        # Load the system wide library