chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,247 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user