# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import subprocess import sys import numpy as np import pytest import tvm import tvm.testing from tvm import te from tvm.support import cc, popen_pool, utils from tvm.testing import env runtime_py = """ import os import sys os.environ["TVM_USE_RUNTIME_LIB"] = "1" import tvm from tvm import te import numpy as np path_dso = sys.argv[1] dtype = sys.argv[2] ff = tvm.runtime.load_module(path_dso) a = tvm.runtime.tensor(np.zeros(10, dtype=dtype)) ff(a) np.testing.assert_equal(a.numpy(), np.arange(a.shape[0])) print("Finish runtime checking...") """ @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.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 @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): 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 run_and_check(): dev = tvm.device(device, 0) 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) worker = popen_pool.PopenWorker() try: worker.send(popen_check) worker.recv() finally: worker.kill() tvm.testing.run_with_gpu_lock(run_and_check) def check_c(device): 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")) def run_and_check(): dev = tvm.device(device, 0) 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) tvm.testing.run_with_gpu_lock(run_and_check) for device in ["cuda", "vulkan", "opencl", "metal"]: check_device(device) check_c(device) @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") def test_combine_module_llvm(): """Test combine multiple module into one shared lib.""" pytest.importorskip("cloudpickle") # needed by popen_pool.PopenWorker # graph nn = 12 n = tvm.runtime.convert(nn) A = te.placeholder((n,), name="A") B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B") mod1 = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd1")) mod2 = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd2")) def check_llvm(): dev = tvm.cpu(0) temp = utils.tempdir() fadd1 = tvm.tirx.build(mod1, "llvm") fadd2 = tvm.tirx.build(mod2, "llvm") path1 = temp.relpath("myadd1.o") path2 = temp.relpath("myadd2.o") path_dso = temp.relpath("mylib.so") fadd1.write_to_file(path1) fadd2.write_to_file(path2) # create shared library with multiple functions cc.create_shared(path_dso, [path1, path2]) m = tvm.runtime.load_module(path_dso) fadd1 = m["myadd1"] fadd2 = m["myadd2"] a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(A.dtype), dev) b = tvm.runtime.tensor(np.zeros(nn, dtype=A.dtype), dev) fadd1(a, b) np.testing.assert_equal(b.numpy(), a.numpy() + 1) fadd2(a, b) np.testing.assert_equal(b.numpy(), a.numpy() + 1) def check_system_lib(): dev = tvm.cpu(0) if not tvm.testing.device_enabled("llvm"): print("Skip because llvm is not enabled") return temp = utils.tempdir() print("Running popen check") fadd1 = tvm.tirx.build(mod1.with_attr("system_lib_prefix", ""), "llvm") fadd2 = tvm.tirx.build(mod2.with_attr("system_lib_prefix", ""), "llvm") path1 = temp.relpath("myadd1.o") path2 = temp.relpath("myadd2.o") path_dso = temp.relpath("mylib.so") fadd1.write_to_file(path1) fadd2.write_to_file(path2) cc.create_shared(path_dso, [path1, path2]) def popen_check(): import ctypes import tvm.runtime # Load dll, will trigger system library registration ctypes.CDLL(path_dso) # Load the system wide library mm = tvm.runtime.system_lib() a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(A.dtype), dev) b = tvm.runtime.tensor(np.zeros(nn, dtype=A.dtype), dev) mm["myadd1"](a, b) np.testing.assert_equal(b.numpy(), a.numpy() + 1) mm["myadd2"](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() if sys.platform != "win32": check_system_lib() check_llvm() if __name__ == "__main__": test_combine_module_llvm()