wehub-resource-sync
109 lines
3.8 KiB
Python
109 lines
3.8 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.
|
|
# ruff: noqa: F821
|
|
|
|
import ctypes
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import ir as I
|
|
from tvm.script import tirx as T
|
|
from tvm.support import cc, popen_pool, tar, utils
|
|
|
|
|
|
@pytest.mark.gpu
|
|
def test_cuda_multi_lib():
|
|
pytest.importorskip("cloudpickle")
|
|
|
|
# test combining two system lib together
|
|
# each contains a fatbin component in cuda
|
|
for device in ["llvm", "cuda"]:
|
|
if not tvm.testing.device_enabled(device):
|
|
print(f"skip because {device} is not enabled...")
|
|
return
|
|
|
|
@tvm.script.ir_module
|
|
class ModA:
|
|
I.module_attrs({"system_lib_prefix": "modA_"})
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
|
|
T.func_attr({"global_symbol": "modA_my_inplace_update"})
|
|
for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
|
for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
|
|
x[tx] = x[tx] + 1
|
|
|
|
@tvm.script.ir_module
|
|
class ModB:
|
|
I.module_attrs({"system_lib_prefix": "modB_"})
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
|
|
T.func_attr({"global_symbol": "modB_my_inplace_update"})
|
|
for bx in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
|
for tx in T.thread_binding(T.int64(12), thread="threadIdx.x"):
|
|
x[tx] = x[tx] + 2
|
|
|
|
temp = utils.tempdir()
|
|
target = tvm.target.Target("cuda", host="llvm")
|
|
libA = tvm.compile(ModA, target=target)
|
|
libB = tvm.compile(ModB, target=target)
|
|
|
|
pathA = temp.relpath("libA.tar")
|
|
pathB = temp.relpath("libB.tar")
|
|
pathAll = temp.relpath("libAll.a")
|
|
|
|
path_dso = temp.relpath("mylib.so")
|
|
libA.export_library(pathA, fcompile=tar.tar)
|
|
libB.export_library(pathB, fcompile=tar.tar)
|
|
cc.create_staticlib(pathAll, [pathA, pathB])
|
|
# package two static libs together
|
|
cc.create_shared(path_dso, ["-Wl,--whole-archive", pathAll, "-Wl,--no-whole-archive"])
|
|
|
|
def popen_check():
|
|
def run_and_check():
|
|
# Load dll, will trigger system library registration
|
|
ctypes.CDLL(path_dso)
|
|
# Load the system wide library
|
|
dev = tvm.cuda()
|
|
a_np = np.random.uniform(size=12).astype("float32")
|
|
a_nd = tvm.runtime.tensor(a_np, dev)
|
|
b_nd = tvm.runtime.tensor(a_np, dev)
|
|
syslibA = tvm.runtime.system_lib("modA_")
|
|
syslibB = tvm.runtime.system_lib("modB_")
|
|
# reload same lib twice
|
|
syslibA = tvm.runtime.system_lib("modA_")
|
|
syslibA["my_inplace_update"](a_nd)
|
|
syslibB["my_inplace_update"](b_nd)
|
|
np.testing.assert_equal(a_nd.numpy(), a_np + 1)
|
|
np.testing.assert_equal(b_nd.numpy(), a_np + 2)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
# system lib should be loaded in different process
|
|
worker = popen_pool.PopenWorker()
|
|
worker.send(popen_check)
|
|
worker.recv()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_synthetic()
|
|
test_cuda_multilib()
|