wehub-resource-sync
248 lines
8.3 KiB
Python
248 lines
8.3 KiB
Python
# 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()
|