Function test_device_module_dump

tests/python/runtime/test_runtime_module_load.py:102–156

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_device_module_dump():
    pytest.importorskip("cloudpickle")  # needed by popen_pool.PopenWorker

    # graph
    n = tvm.runtime.convert(1024)
    A = te.placeholder((n,), name="A")
    B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")

    sch = tvm.s_tir.Schedule(te.create_prim_func([A, B]))
    # create iter var and assign them tags.
    num_thread = 8
    bx, tx = sch.split(sch.get_loops("B")[0], factors=[None, num_thread])
    sch.bind(bx, "blockIdx.x")
    sch.bind(tx, "threadIdx.x")

    def check_device(device):
        dev = tvm.device(device, 0)
        if not tvm.testing.device_enabled(device):
            print(f"Skip because {device} is not enabled")
            return
        temp = utils.tempdir()
        f = tvm.compile(sch.mod, target=device)

        path_dso = temp.relpath("dev_lib.so")
        # test cross compiler function
        f.export_library(path_dso, fcompile=cc.cross_compiler("g++"))

        def popen_check():
            import tvm

            f1 = tvm.runtime.load_module(path_dso)
            a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
            b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
            f1(a, b)
            np.testing.assert_equal(b.numpy(), a.numpy() + 1)

        # system lib should be loaded in different process
        worker = popen_pool.PopenWorker()
        worker.send(popen_check)
        worker.recv()

    def check_c(device):
        dev = tvm.device(device, 0)
        if not tvm.testing.device_enabled(device):
            print(f"Skip because {device} is not enabled")
            return
        f = tvm.compile(sch.mod, target=tvm.target.Target(device, host="c"))
        a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
        b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
        f["main"](a, b)
        np.testing.assert_equal(b.numpy(), a.numpy() + 1)

    for device in ["cuda", "vulkan", "opencl", "metal"]:
        check_device(device)
        check_c(device)


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")