This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c46d112866 [REFACTOR][TEST] Consolidate runtime tests and isolate RPC
(#20716)
c46d112866 is described below
commit c46d112866cfdd2f2ca2cf601c2fbf4207da3623
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 09:05:12 2026 -0400
[REFACTOR][TEST] Consolidate runtime tests and isolate RPC (#20716)
Consolidate runtime module loading, exports, properties, timing, and
Executable wrappers in one suite, and group RPC tests under runtime/rpc.
Each suite owns its small native fixtures, replacing compiler-generated
kernels in runtime checks.
Combine tensor-view coverage while preserving offsets and
parent-relative bounds, retain useful imported-tensor checks, and move
compiler and dialect cases to their owning suites. Separate JSON
serialization from script roundtrips, with explicit checks of rendered
script. Remove obsolete integration scaffolding and make RPC process
readiness and cleanup explicit.
---
.../codegen/test_target_codegen_cross_llvm.py | 36 ++
.../codegen/test_target_codegen_gpu_common.py | 64 ++++
tests/python/codegen/test_target_codegen_llvm.py | 22 ++
tests/python/ir/test_ir_type.py | 13 +
.../relax/test_vm_alloc_storage_with_scope.py | 20 +
tests/python/runtime/{ => rpc}/test_rpc_base.py | 2 +-
tests/python/runtime/{ => rpc}/test_runtime_rpc.py | 420 +++++++++------------
tests/python/runtime/test_executable.py | 264 -------------
tests/python/runtime/test_runtime_dlpack.py | 67 ----
tests/python/runtime/test_runtime_error.py | 32 --
tests/python/runtime/test_runtime_extension.py | 43 ---
tests/python/runtime/test_runtime_measure.py | 71 ----
tests/python/runtime/test_runtime_memory_scope.py | 116 +-----
tests/python/runtime/test_runtime_module.py | 263 +++++++++++++
tests/python/runtime/test_runtime_module_export.py | 69 ----
tests/python/runtime/test_runtime_module_load.py | 246 ------------
.../python/runtime/test_runtime_module_property.py | 60 ---
tests/python/runtime/test_runtime_nd_array.py | 254 -------------
tests/python/runtime/test_runtime_tensor.py | 84 +++++
tests/python/tirx-base/test_tir_tensor.py | 45 +++
.../python/tirx/script/test_tirx_script_tensors.py | 57 +++
21 files changed, 791 insertions(+), 1457 deletions(-)
diff --git a/tests/python/codegen/test_target_codegen_cross_llvm.py
b/tests/python/codegen/test_target_codegen_cross_llvm.py
index 00f33dd124..3865a48e40 100644
--- a/tests/python/codegen/test_target_codegen_cross_llvm.py
+++ b/tests/python/codegen/test_target_codegen_cross_llvm.py
@@ -95,5 +95,41 @@ def test_llvm_add_pipeline():
build_arm()
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
+def test_bigendian_rpc():
+ """Test big endian rpc when there is a PowerPC RPC server available"""
+ host = os.environ.get("TVM_POWERPC_TEST_HOST", None)
+ port = os.environ.get("TVM_POWERPC_TEST_PORT", 9090)
+ if host is None:
+ pytest.skip("PowerPC RPC host is not configured")
+
+ def verify_rpc(remote, target, shape, dtype):
+ @T.function
+ def main(A: T.Tensor(shape, dtype), B: T.Tensor(shape, dtype)):
+ for i in range(shape[0]):
+ B[i] = A[i] + tvm.tirx.const(1, dtype)
+
+ f = tvm.tirx.compile(main, target=target).mod
+
+ dev = remote.cpu(0)
+ a = dev.tensor(np.random.randint(0, 256, size=shape).astype(dtype))
+ b = dev.tensor(np.zeros(shape).astype(dtype))
+ temp = utils.tempdir()
+ path_dso = temp.relpath("dev_lib.o")
+ f.write_to_file(path_dso)
+ remote.upload(path_dso)
+ f = remote.load_module("dev_lib.o")
+ f(a, b)
+ tvm.testing.assert_allclose(a.numpy() + 1, b.numpy())
+
+ print("Test RPC connection to PowerPC...")
+ remote = rpc.connect(host, port)
+ target = {"kind": "llvm", "mtriple": "powerpc-linux-gnu"}
+ for dtype in ["float32", "float64", "int32", "int8"]:
+ verify_rpc(remote, target, (10,), dtype)
+
+
if __name__ == "__main__":
test_llvm_add_pipeline()
diff --git a/tests/python/codegen/test_target_codegen_gpu_common.py
b/tests/python/codegen/test_target_codegen_gpu_common.py
index 26a4f9a189..fb5d3417e4 100644
--- a/tests/python/codegen/test_target_codegen_gpu_common.py
+++ b/tests/python/codegen/test_target_codegen_gpu_common.py
@@ -23,6 +23,7 @@ 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, utils
from tvm.testing import env
@@ -72,5 +73,68 @@ def test_int_intrin(target, dtype):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_gpu(), reason="need gpu")
+def test_device_module_dump():
+ pytest.importorskip("cloudpickle") # needed by popen_pool.PopenWorker
+
+ @T.function
+ def main(A: T.Tensor((1024,), "float32"), B: T.Tensor((1024,), "float32")):
+ with T.launch_thread("blockIdx.x", 128) as bx:
+ with T.launch_thread("threadIdx.x", 8) as tx:
+ B[bx * 8 + tx] = A[bx * 8 + tx] + T.float32(1)
+
+ 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.tirx.compile(main, 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("float32"), dev)
+ b = tvm.runtime.tensor(np.zeros(1024, dtype="float32"), 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.tirx.compile(main, 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("float32"), dev)
+ b = tvm.runtime.tensor(np.zeros(1024, dtype="float32"), 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)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index 28dd942dd9..f3f2cbfc09 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -1227,5 +1227,27 @@ def test_bool_return_value():
assert not built(15)
[email protected](not env.has_llvm(), reason="need llvm")
+def test_llvm_module_properties_and_formats(tmp_path):
+ @T.function
+ def main(A: T.Tensor((10,), "int64")):
+ for i in range(9):
+ A[i + 1] = A[i] + T.int64(1)
+
+ module = tvm.tirx.compile(main, target="llvm").mod
+ assert not module.is_binary_serializable()
+ assert module.is_runnable()
+ assert module.is_compilation_exportable()
+ for extension in ("o", "ll", "bc"):
+ path = tmp_path / ("module." + extension)
+ module.write_to_file(str(path))
+ assert path.is_file()
+ if extension == "ll":
+ loaded = tvm.runtime.load_module(str(path))
+ data = tvm.runtime.tensor(np.zeros(10, dtype="int64"))
+ loaded(data)
+ np.testing.assert_equal(data.numpy(), np.arange(10))
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/ir/test_ir_type.py b/tests/python/ir/test_ir_type.py
index df19f4ad92..b0012b2a10 100644
--- a/tests/python/ir/test_ir_type.py
+++ b/tests/python/ir/test_ir_type.py
@@ -17,9 +17,12 @@
# ruff: noqa: F401, F821
"""Test type nodes in the IR"""
+import uuid
+
import pytest
import tvm
+from tvm.runtime import MemoryScope
from tvm.script import tirx as T
@@ -69,6 +72,16 @@ def test_tuple_type():
check_json_roundtrip(tup_ty)
+def test_memory_scope_json_roundtrip():
+ custom = MemoryScope.register("test.ir.scope." + uuid.uuid4().hex)
+ for scope in (MemoryScope("global"), MemoryScope("shared.dyn"), custom):
+ assert tvm.ir.load_json(tvm.ir.save_json(scope)).same_as(scope)
+ ty = tvm.tirx.TensorType((4,), "float32", scope=scope)
+ restored = tvm.ir.load_json(tvm.ir.save_json(ty))
+ assert restored.scope.same_as(scope)
+ tvm.ir.assert_structural_equal(ty, restored)
+
+
if __name__ == "__main__":
test_tensor_type_bad_constructor()
test_func_type()
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py
b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index 3d3fd97028..06ea6febe2 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -17,10 +17,12 @@
"""Test tensor allocation with memory scope for Relax VM"""
import numpy as np
+import pytest
import tvm
import tvm.testing
from tvm import relax
+from tvm.runtime import MemoryScope
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import s_tir as Ts
@@ -74,5 +76,23 @@ def test_alloc_storage_with_scope_global():
tvm.testing.assert_allclose(output_ref, output)
+def test_relax_allocation_scope_operands():
+ scope = MemoryScope("global")
+ for make_call in (
+ lambda value: tvm.relax.op.builtin.alloc_tensor((4,), "float32", 0,
value),
+ lambda value: tvm.relax.op.memory.alloc_storage((16,), 0, value,
"float32"),
+ lambda value: tvm.relax.op.vm.alloc_storage((16,), 0, "float32",
value),
+ ):
+ tvm.ir.assert_structural_equal(make_call(scope), make_call(scope.name))
+
+
+def test_vdevice_memory_scope():
+ global_scope = MemoryScope("global")
+ for scope in ("global", global_scope):
+ assert tvm.relax.VDevice("llvm",
memory_scope=scope).memory_scope.same_as(global_scope)
+ with pytest.raises(ValueError):
+ tvm.relax.VDevice("llvm", memory_scope="")
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/runtime/test_rpc_base.py
b/tests/python/runtime/rpc/test_rpc_base.py
similarity index 98%
rename from tests/python/runtime/test_rpc_base.py
rename to tests/python/runtime/rpc/test_rpc_base.py
index 3d54730084..ab2789c309 100644
--- a/tests/python/runtime/test_rpc_base.py
+++ b/tests/python/runtime/rpc/test_rpc_base.py
@@ -14,11 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F821
import random
import pytest
+import tvm.testing
from tvm.rpc import base
diff --git a/tests/python/runtime/test_runtime_rpc.py
b/tests/python/runtime/rpc/test_runtime_rpc.py
similarity index 66%
rename from tests/python/runtime/test_runtime_rpc.py
rename to tests/python/runtime/rpc/test_runtime_rpc.py
index 0bfbcfe36e..0d96670b1e 100644
--- a/tests/python/runtime/test_runtime_rpc.py
+++ b/tests/python/runtime/rpc/test_runtime_rpc.py
@@ -22,7 +22,6 @@ import os
import platform
import stat
import sys
-import tempfile
import time
import numpy as np
@@ -33,14 +32,10 @@ pytest.importorskip("tornado") # tvm.rpc.proxy and
tvm.rpc.tracker require torn
import tvm
import tvm.testing
-from tvm import rpc, te
+from tvm import rpc
from tvm.rpc.proxy import Proxy
from tvm.rpc.tracker import Tracker
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
from tvm.support import cc, utils
-from tvm.testing import env
if __name__ == "__main__":
# NOTE: must live here to avoid registering PackedFunc with
libtvm_compiler.so twice.
@@ -64,41 +59,70 @@ pytestmark = pytest.mark.skipif(
),
)
-# NOTE: When writing tests, wrap remote related checking in a sub-function
-# to ensure all the remote resources destructs before the server terminates
[email protected](scope="module")
+def native_source():
+ return r"""
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/function.h>
+void add_one(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
+ for (int64_t i = 0; i < a.size(0); ++i) {
+ static_cast<float*>(b.data_ptr())[i] =
static_cast<float*>(a.data_ptr())[i] + 1;
+ }
+}
+int64_t answer() { return 42; }
+int64_t with_arg(int64_t) { return 42; }
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(main, add_one);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(answer, answer);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(with_arg, with_arg);
+"""
+
+
[email protected](scope="module")
+def native_compile_options():
+ return [
+ "-std=c++17",
+ "-I" + tvm_ffi.libinfo.find_include_path(),
+ "-I" + tvm_ffi.libinfo.find_dlpack_include_path(),
+ ]
+
+
[email protected](scope="module")
+def native_object(tmp_path_factory, native_source, native_compile_options):
+ directory = tmp_path_factory.mktemp("runtime_module")
+ source = directory / "module.cc"
+ source.write_text(native_source)
+ path = directory / "module.o"
+ cc.create_shared(str(path), [str(source)], options=["-c",
*native_compile_options])
+ return path
+
+
[email protected](scope="module")
+def native_library(native_object):
+ path = native_object.with_suffix(".so")
+ cc.create_shared(str(path), [str(native_object)])
+ return path
+
+
[email protected](scope="module")
+def system_library_source(native_source):
+ return (
+ native_source
+ + r"""
+#include <tvm/ffi/extra/c_env_api.h>
+struct RegisterSystemLibrary {
+ RegisterSystemLibrary() {
+ TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_main",
(void*)__tvm_ffi_main);
+ TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_answer",
(void*)__tvm_ffi_answer);
+ }
+};
+static RegisterSystemLibrary register_system_library;
+"""
+ )
[email protected](
- tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
-)
-def test_bigendian_rpc():
- """Test big endian rpc when there is a PowerPC RPC server available"""
- host = os.environ.get("TVM_POWERPC_TEST_HOST", None)
- port = os.environ.get("TVM_POWERPC_TEST_PORT", 9090)
- if host is None:
- return
-
- def verify_rpc(remote, target, shape, dtype):
- A = te.placeholder(shape, dtype=dtype)
- B = te.compute(A.shape, lambda i: A[i] + tvm.tirx.const(1, A.dtype))
- f = tvm.compile(te.create_function([A, B]), target=target)
-
- dev = remote.cpu(0)
- a = dev.tensor(np.random.randint(0, 256,
size=shape).astype(str(A.dtype)))
- b = dev.tensor(np.zeros(shape).astype(str(A.dtype)))
- temp = utils.tempdir()
- path_dso = temp.relpath("dev_lib.o")
- f.write_to_file(path_dso)
- remote.upload(path_dso)
- f = remote.load_module("dev_lib.o")
- f(a, b)
- tvm.testing.assert_allclose(a.numpy() + 1, b.numpy())
- print("Test RPC connection to PowerPC...")
- remote = rpc.connect(host, port)
- target = {"kind": "llvm", "mtriple": "powerpc-linux-gnu"}
- for dtype in ["float32", "float64", "int32", "int8"]:
- verify_rpc(remote, target, (10,), dtype)
+# NOTE: When writing tests, wrap remote related checking in a sub-function
+# to ensure all the remote resources destructs before the server terminates
@pytest.mark.skipif(
@@ -207,26 +231,6 @@ def test_rpc_echo():
check(client, False)
- def check_minrpc():
- if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is
None:
- return
- # Test minrpc server.
- temp = utils.tempdir()
- minrpc_exec = temp.relpath("minrpc")
- tvm.rpc.with_minrpc(cc.create_executable)(minrpc_exec, [])
- check(rpc.PopenSession(minrpc_exec), False)
- # minrpc on the remote
- server = rpc.Server()
- client = rpc.connect(
- "127.0.0.1",
- server.port,
- session_constructor_args=["rpc.PopenSession", open(minrpc_exec,
"rb").read()],
- )
- check(client, False)
-
- # skip for now until we upgrade to new FFI
- # check_minrpc()
-
@pytest.mark.skipif(
tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
@@ -247,114 +251,66 @@ def test_rpc_file_exchange():
@pytest.mark.skipif(
tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
)
[email protected](not env.has_llvm(), reason="need llvm")
-def test_rpc_remote_module():
- # graph
- n = tvm.runtime.convert(102)
- A = te.placeholder((n,), name="A")
- B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
- mod = tvm.ir.IRModule.from_expr(te.create_function([A,
B]).with_attr("global_symbol", "myadd"))
-
- server0 = rpc.Server(key="x0")
- server1 = rpc.Server(key="x1")
-
- client = rpc.connect(
- "127.0.0.1",
- server0.port,
- key="x0",
- session_constructor_args=["rpc.Connect", "127.0.0.1", server1.port,
"x1", False],
- )
-
- def check_remote(remote):
- temp = utils.tempdir()
- dev = remote.cpu(0)
- f = tvm.compile(mod, "llvm")
- path_dso = temp.relpath("dev_lib.so")
- f.export_library(path_dso)
- remote.upload(path_dso)
- f1 = remote.load_module("dev_lib.so")
- a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
- b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
- time_f = rpc.time_evaluator(f1, f1.entry_name, remote.cpu(0),
number=10)
- cost = time_f(a, b).mean
- print(f"{cost:g} secs/op")
- np.testing.assert_equal(b.numpy(), a.numpy() + 1)
-
- # Download the file from the remote
- path_tar = temp.relpath("dev_lib.tar")
- f.export_library(path_tar)
- remote.upload(path_tar)
- local_download_path = temp.relpath("dev_lib.download.so")
- with open(local_download_path, "wb") as fo:
- fo.write(remote.download_linked_module("dev_lib.tar"))
- fupdated = tvm.runtime.load_module(local_download_path)
- a =
tvm.runtime.tensor(np.random.uniform(size=102).astype(str(A.dtype)), tvm.cpu(0))
- b = tvm.runtime.tensor(np.zeros(102, dtype=str(A.dtype)), tvm.cpu(0))
- fupdated(a, b)
- np.testing.assert_equal(b.numpy(), a.numpy() + 1)
-
- def check_minrpc():
- if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is
None:
- return
- # export to minrpc
- temp = utils.tempdir()
- # system lib prefix will trigger system lib build
- f = tvm.compile(mod.with_attr("system_lib_prefix", ""), "llvm")
- path_minrpc = temp.relpath("dev_lib.minrpc")
- f.export_library(path_minrpc,
fcompile=rpc.with_minrpc(cc.create_executable))
-
- with pytest.raises(RuntimeError):
- rpc.PopenSession("filenotexist")
-
- # statrt the minrpc session.
- remote = tvm.rpc.PopenSession(path_minrpc)
+def test_rpc_remote_module(native_object, system_library_source, tmp_path):
+ module = tvm.runtime.load_static_library(str(native_object), ["main",
"answer", "with_arg"])
+ path_dso = tmp_path / "dev_lib.so"
+ path_tar = tmp_path / "dev_lib.tar"
+ module.export_library(str(path_dso))
+ module.export_library(str(path_tar))
+
+ def check_remote(remote, name):
dev = remote.cpu(0)
- f1 = remote.system_lib()
-
- a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
- b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
- time_f = rpc.time_evaluator(f1, "myadd", remote.cpu(0), number=1)
- cost = time_f(a, b).mean
+ remote.upload(str(path_dso))
+ loaded = remote.load_module("dev_lib.so")
+ a = dev.tensor(np.random.uniform(size=102).astype("float32"))
+ b = dev.tensor(np.zeros(102, dtype="float32"))
+ rpc.time_evaluator(loaded, "main", dev, number=10)(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
- # change to not executable
- os.chmod(path_minrpc, stat.S_IRUSR)
- with pytest.raises(RuntimeError):
- rpc.PopenSession(path_minrpc)
-
- def check_remote_link_cl(remote):
- """Test function to run remote code such as cl
-
- This is not enabled because there is forking issue
- of TVM runtime when server launches after OpenCL
- runtime initializes. We leave it as an example
- on how to do rpc when we want to do linking on remote.
- """
- if not tvm.testing.device_enabled("opencl"):
- print("Skip because opencl is not enabled")
- return
- temp = utils.tempdir()
- dev = remote.cl(0)
-
- s = tvm.s_tir.Schedule(mod)
-
- x = s.get_loops(s.get_sblock("B"))
- xo, xi = s.split(x, factors=[None, 32])
- s.bind(xo, "blockIdx.x")
- s.bind(xi, "threadIdx.x")
- f = tvm.compile(s.mod, tvm.target.Target("opencl", host="llvm"))
- path_tar = temp.relpath("myadd.tar")
- f.export_library(path_tar)
- remote.upload(path_tar)
- fhost = remote.load_module("myadd.tar")
- a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
- b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
- fhost(a, b)
+ remote.upload(str(path_tar))
+ path = tmp_path / f"download-{name}.so"
+ path.write_bytes(remote.download_linked_module("dev_lib.tar"))
+ downloaded = tvm.runtime.load_module(str(path))
+ a = tvm.runtime.tensor(np.random.uniform(size=102).astype("float32"))
+ b = tvm.runtime.tensor(np.zeros(102, dtype="float32"))
+ downloaded(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
- check_remote(rpc.LocalSession())
- check_remote(client)
- check_minrpc()
+ server0 = rpc.Server(key="x0")
+ server1 = rpc.Server(key="x1")
+ try:
+ client = rpc.connect(
+ "127.0.0.1",
+ server0.port,
+ key="x0",
+ session_constructor_args=["rpc.Connect", "127.0.0.1",
server1.port, "x1", False],
+ )
+ check_remote(rpc.LocalSession(), "local")
+ check_remote(client, "tcp")
+ del client
+ finally:
+ server0.terminate()
+ server1.terminate()
+
+ if tvm.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None:
+ return
+ source = tmp_path / "system.cc"
+ source.write_text(system_library_source)
+ path_minrpc = tmp_path / "dev_lib.minrpc"
+ rpc.with_minrpc(cc.create_executable)(str(path_minrpc), [str(source)])
+ with pytest.raises(RuntimeError):
+ rpc.PopenSession("filenotexist")
+ remote = rpc.PopenSession(str(path_minrpc))
+ dev = remote.cpu(0)
+ loaded = remote.system_lib()
+ a = dev.tensor(np.random.uniform(size=102).astype("float32"))
+ b = dev.tensor(np.zeros(102, dtype="float32"))
+ rpc.time_evaluator(loaded, "main", dev, number=1)(a, b)
+ np.testing.assert_equal(b.numpy(), a.numpy() + 1)
+ del a, b, loaded, dev, remote
+ os.chmod(path_minrpc, stat.S_IRUSR)
+ with pytest.raises(RuntimeError):
+ rpc.PopenSession(str(path_minrpc))
@pytest.mark.skipif(
@@ -568,11 +524,11 @@ def test_rpc_tracker_register(device_key):
tracker.terminate()
-def _target(host, port, device_key, timeout):
+def _target(host, port, device_key, timeout, acquired):
client = rpc.connect_tracker(host, port)
remote = client.request(device_key, session_timeout=timeout)
- while True:
- pass
+ acquired.set()
+ multiprocessing.Event().wait()
remote.cpu()
@@ -581,7 +537,6 @@ def _target(host, port, device_key, timeout):
)
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
def test_rpc_tracker_request(device_key):
- # test concurrent request
tracker = Tracker(port=9000, port_end=10000)
server = rpc.Server(
port=9000,
@@ -590,33 +545,39 @@ def test_rpc_tracker_request(device_key):
tracker_addr=("127.0.0.1", tracker.port),
)
client = rpc.connect_tracker("127.0.0.1", tracker.port)
-
- proc1 = multiprocessing.Process(target=_target, args=("127.0.0.1",
tracker.port, device_key, 4))
+ acquired1 = multiprocessing.Event()
+ acquired2 = multiprocessing.Event()
+ proc1 = multiprocessing.Process(
+ target=_target, args=("127.0.0.1", tracker.port, device_key, 4,
acquired1)
+ )
proc2 = multiprocessing.Process(
- target=_target, args=("127.0.0.1", tracker.port, device_key, 200)
+ target=_target, args=("127.0.0.1", tracker.port, device_key, 200,
acquired2)
)
- proc1.start()
- time.sleep(0.5)
- proc2.start()
- time.sleep(0.5)
-
- summary = client.summary()
-
- assert summary["queue_info"][device_key]["free"] == 0
- assert summary["queue_info"][device_key]["pending"] == 1
-
- proc1.terminate()
- proc1.join()
- time.sleep(0.5)
-
- summary = client.summary()
- assert summary["queue_info"][device_key]["free"] == 0
- assert summary["queue_info"][device_key]["pending"] == 0
-
- proc2.terminate()
- proc2.join()
- server.terminate()
- tracker.terminate()
+ try:
+ proc1.start()
+ assert acquired1.wait(20), "first RPC session was not acquired"
+ proc2.start()
+ deadline = time.monotonic() + 10
+ while True:
+ queue = client.summary()["queue_info"][device_key]
+ if queue["pending"] == 1 or time.monotonic() >= deadline:
+ break
+ time.sleep(0.05)
+ assert queue["free"] == 0
+ assert queue["pending"] == 1
+ proc1.terminate()
+ proc1.join()
+ assert acquired2.wait(20), "pending RPC request did not acquire the
released session"
+ queue = client.summary()["queue_info"][device_key]
+ assert queue["free"] == 0
+ assert queue["pending"] == 0
+ finally:
+ for process in (proc1, proc2):
+ if process.pid is not None:
+ process.terminate()
+ process.join()
+ server.terminate()
+ tracker.terminate()
@pytest.mark.skipif(
@@ -683,58 +644,35 @@ def test_rpc_session_timeout_error(with_proxy):
server = rpc.Server(port=port, port_end=port_end,
tracker_addr=tracker_addr, key="x1")
time.sleep(0.5)
- rpc_sess = rpc.connect_tracker(*tracker_addr).request(key="x1",
session_timeout=1)
-
- with pytest.raises(tvm.error.RPCSessionTimeoutError):
+ try:
+ # Leave room for server startup before testing expiration of an
established session.
+ rpc_sess = rpc.connect_tracker(*tracker_addr).request(key="x1",
session_timeout=5)
f1 = rpc_sess.get_function("rpc.test.addone")
- time.sleep(2)
- f1(10)
-
- server.terminate()
- if with_proxy:
- proxy.terminate()
- tracker.terminate()
+ time.sleep(6)
+ with pytest.raises(tvm.error.RPCSessionTimeoutError):
+ f1(10)
+ finally:
+ server.terminate()
+ if with_proxy:
+ proxy.terminate()
+ tracker.terminate()
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.parametrize("call_with_unused_argument", [True, False])
-def test_compiled_function_with_zero_arguments(call_with_unused_argument):
- """RPC functions do not require an argument
-
- This is a regression test. When no arguments are provided, RPC
- provides NULL as the `TVMFFIAny* args` argument to a PackedFunc.
- However, previous implementations of `MakePackedAPI`
- unconditionally asserted that the `args` pointer was non-null.
- This assertion is now generated only when the function accepts
- a non-zero number of arguments.
-
- """
-
- @I.ir_module
- class Module:
- @Ts.function
- def func_without_arg() -> T.int64:
- return T.int64(42)
-
- @Ts.function
- def func_with_arg(unused: T.int64) -> T.int64:
- return T.int64(42)
-
- built = tvm.compile(Module, target="llvm")
-
- server = tvm.rpc.Server(key="x1")
- client = tvm.rpc.connect("127.0.0.1", server.port, key="x1")
-
- libname = "libbuilt.so"
- with tempfile.TemporaryDirectory(prefix="tvm_rpc_testing_") as temp_dir:
- local_path = os.path.join(temp_dir, libname)
- built.export_library(local_path)
- client.upload(local_path)
-
- remote_mod = client.load_module(libname)
-
- if call_with_unused_argument:
- res = remote_mod["func_with_arg"](0)
- else:
- res = remote_mod["func_without_arg"]()
-
- assert res == 42
+def test_remote_function_with_zero_arguments(native_library,
call_with_unused_argument):
+ server = rpc.Server(key="x1")
+ try:
+ client = rpc.connect("127.0.0.1", server.port, key="x1")
+ client.upload(str(native_library))
+ module = client.load_module(native_library.name)
+ if call_with_unused_argument:
+ result = module["with_arg"](0)
+ else:
+ result = module["answer"]()
+ assert result == 42
+ del module, client
+ finally:
+ server.terminate()
diff --git a/tests/python/runtime/test_executable.py
b/tests/python/runtime/test_executable.py
deleted file mode 100644
index 512b830fb9..0000000000
--- a/tests/python/runtime/test_executable.py
+++ /dev/null
@@ -1,264 +0,0 @@
-# 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.
-"""Tests for the Executable class."""
-
-import os
-import tempfile
-
-import numpy as np
-
-import tvm
-import tvm.testing
-from tvm.runtime import Executable
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
[email protected]_module
-class MyModule:
- @Ts.function
- def add(
- A: T.Tensor((10,), "float32"),
- B: T.Tensor((10,), "float32"),
- C: T.Tensor((10,), "float32"),
- ):
- for i in range(10):
- C[i] = A[i] + B[i]
-
-
-def test_executable_init():
- """Test initialization of Executable class."""
- lib = tvm.tirx.compile(MyModule, target="llvm").mod
- executable = Executable(lib)
-
- assert executable.mod is lib
- assert executable._jitted_mod is None
-
-
-def test_executable_getitem():
- """Test __getitem__ method of Executable class."""
- lib = tvm.tirx.compile(MyModule, target="llvm").mod
- executable = Executable(lib)
-
- # Jit the module first
- executable.jit()
-
- # Test __getitem__
- add_func = executable["add"]
-
- # Verify the function works
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- add_func(a, b, c)
-
- # Check results
- tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10,
dtype="float32"))
-
-
-def test_executable_jit_already_jitted():
- """Test jit method when module is already jitted."""
- lib = tvm.tirx.compile(MyModule, target="llvm").mod
- executable = Executable(lib)
-
- # First jit call
- jitted_mod1 = executable.jit()
-
- # Second jit call should return the cached jitted module
- jitted_mod2 = executable.jit()
- assert jitted_mod2 is jitted_mod1
-
- # Test with force_recompile
- jitted_mod3 = executable.jit(force_recompile=True)
- # The module might be different after force recompilation
-
- # Verify both modules work correctly
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c1 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
- c2 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- jitted_mod1["add"](a, b, c1)
- jitted_mod3["add"](a, b, c2)
-
- tvm.testing.assert_allclose(c1.numpy(), np.array([3.0] * 10,
dtype="float32"))
- tvm.testing.assert_allclose(c2.numpy(), np.array([3.0] * 10,
dtype="float32"))
-
-
-def test_executable_export_library():
- """Test export_library method."""
- lib = tvm.tirx.compile(MyModule, target="llvm").mod
- executable = Executable(lib)
-
- # Create a temporary directory for the library
- temp_dir = tempfile.mkdtemp()
- try:
- lib_path = os.path.join(temp_dir, "test_lib.so")
- executable.export_library(lib_path)
-
- # Verify the library was created
- assert os.path.exists(lib_path)
-
- # Load the library back
- loaded_mod = tvm.runtime.load_module(lib_path)
- assert loaded_mod is not None
-
- # Test the loaded module
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- loaded_mod["add"](a, b, c)
-
- # Check results
- tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10,
dtype="float32"))
- finally:
- # Clean up
- if os.path.exists(temp_dir):
- import shutil
-
- shutil.rmtree(temp_dir)
-
-
-def test_executable_export_library_with_workspace():
- """Test export_library method with workspace_dir."""
- lib = tvm.tirx.compile(MyModule, target="llvm").mod
- executable = Executable(lib)
-
- # Create temporary directories
- temp_dir = tempfile.mkdtemp()
- workspace_dir = tempfile.mkdtemp()
-
- try:
- lib_path = os.path.join(temp_dir, "test_lib.so")
- executable.export_library(lib_path, workspace_dir=workspace_dir)
-
- # Verify the library was created
- assert os.path.exists(lib_path)
-
- # Load the library back
- loaded_mod = tvm.runtime.load_module(lib_path)
- assert loaded_mod is not None
-
- # Test the loaded module
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- loaded_mod["add"](a, b, c)
-
- # Check results
- tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10,
dtype="float32"))
- finally:
- # Clean up
- for directory in [temp_dir, workspace_dir]:
- if os.path.exists(directory):
- import shutil
-
- shutil.rmtree(directory)
-
-
-def test_executable_integration():
- """Integration test for Executable with a simple TVM module."""
- # Create target and build
- target = tvm.target.Target("llvm")
- lib = tvm.tirx.compile(MyModule, target=target).mod
-
- # Create an executable
- executable = Executable(lib)
-
- # Test jit
- jitted_mod = executable.jit()
- assert jitted_mod is not None
-
- # Test __getitem__
- add_func = executable["add"]
- assert add_func is not None
-
- # Test the function works
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- add_func(a, b, c)
-
- # Check results
- tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10,
dtype="float32"))
-
- # Test export_library
- temp_dir = tempfile.mkdtemp()
- try:
- lib_path = os.path.join(temp_dir, "test_lib.so")
- executable.export_library(lib_path)
-
- # Verify the library was created
- assert os.path.exists(lib_path)
-
- # Load the library back
- loaded_mod = tvm.runtime.load_module(lib_path)
- assert loaded_mod is not None
-
- # Test the loaded module
- loaded_add = loaded_mod["add"]
- c_loaded = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
- loaded_add(a, b, c_loaded)
-
- # Check results
- tvm.testing.assert_allclose(c_loaded.numpy(), np.array([3.0] * 10,
dtype="float32"))
-
- finally:
- # Clean up
- if os.path.exists(temp_dir):
- import shutil
-
- shutil.rmtree(temp_dir)
-
-
-def test_executable_jit_force_recompile():
- """Test jit method with force_recompile=True."""
- # Create target and build
- target = tvm.target.Target("c")
- lib = tvm.tirx.compile(MyModule, target=target).mod
-
- # Create an executable
- executable = Executable(lib)
-
- # First jit call
- jitted_mod1 = executable.jit()
-
- # Second jit call without force_recompile should return the same module
- jitted_mod2 = executable.jit()
- assert jitted_mod1 is jitted_mod2
-
- # Third jit call with force_recompile should return a new module
- jitted_mod3 = executable.jit(force_recompile=True)
- assert jitted_mod3 is not jitted_mod1
-
- # Test the function works
- a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
- b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
- c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
-
- jitted_mod3["add"](a, b, c)
-
- # Check results
- tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10,
dtype="float32"))
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_dlpack.py
b/tests/python/runtime/test_runtime_dlpack.py
deleted file mode 100644
index 508bab6fba..0000000000
--- a/tests/python/runtime/test_runtime_dlpack.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# 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 numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import te
-
-# These tests exercise the PyTorch DLPack interop path; skip the whole module
-# when torch is unavailable.
-pytest.importorskip("torch")
-
-
-def test_from_dlpack_shape_one():
- # A test case for the issue https://github.com/pytorch/pytorch/issues/99803
- import torch
- from torch.utils.dlpack import to_dlpack
-
- tgt = tvm.target.Target(target="llvm", host="llvm")
-
- rows = 1
- a = tvm.runtime.from_dlpack(to_dlpack(torch.randn(rows, 16)))
-
- A = te.placeholder((rows, 16), name="A")
- B = te.placeholder((rows, 16), name="B")
- C = te.compute(A.shape, lambda i, j: A[i, j] + B[i, j], name="C")
-
- fadd = tvm.compile(te.create_function([A, B, C]), target=tgt)
-
- dev = tvm.cpu()
-
- b = tvm.runtime.tensor(np.random.uniform(size=(rows, 16)).astype(B.dtype),
dev)
- c = tvm.runtime.tensor(np.zeros((rows, 16), dtype=C.dtype), dev)
- fadd(a, b, c)
-
- tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
-
-
-def test_from_dlpack_strided():
- import torch
- from torch.utils.dlpack import to_dlpack
-
- rows = 1
- inp = torch.randn(rows, 16)
- a = tvm.runtime.from_dlpack(to_dlpack(inp))
- view = a._create_view((2, 8))
-
- np.testing.assert_equal(inp.numpy().reshape(2, 8), view.numpy())
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_error.py
b/tests/python/runtime/test_runtime_error.py
deleted file mode 100644
index 2a23a04ffe..0000000000
--- a/tests/python/runtime/test_runtime_error.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# 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.
-"""Test runtime error handling"""
-
-import tvm
-import tvm.testing
-
-
-def test_op_translation_to_not_implemented():
- try:
- tvm.testing.test_raise_error("OpNotImplemented", "myop")
- assert False
- except tvm.error.OpNotImplemented as e:
- assert isinstance(e, NotImplementedError)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_extension.py
b/tests/python/runtime/test_runtime_extension.py
deleted file mode 100644
index 8856edb167..0000000000
--- a/tests/python/runtime/test_runtime_extension.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# 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 numpy as np
-
-import tvm
-from tvm.script import ir as I
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def test_dltensor_compatible():
- n = T.dynamic("n", "int32")
-
- @I.ir_module
- class Module:
- @Ts.function
- def arange(Ab: T.Tensor((n,), "int64")):
- for i in T.serial(n - 1):
- Ab[i + 1] = Ab[i] + T.int64(1)
-
- mod = Module
- f = tvm.compile(mod, target="llvm")
- a = tvm.runtime.tensor(np.zeros(10, dtype="int64"))
- f(a)
- np.testing.assert_equal(a.numpy(), np.arange(a.shape[0]))
-
-
-if __name__ == "__main__":
- test_dltensor_compatible()
diff --git a/tests/python/runtime/test_runtime_measure.py
b/tests/python/runtime/test_runtime_measure.py
deleted file mode 100644
index 29d2832e99..0000000000
--- a/tests/python/runtime/test_runtime_measure.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# 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: F401
-import ctypes
-import time
-
-import tvm
-from tvm import te
-from tvm.runtime.module import BenchmarkResult
-from tvm.support.utils import tempdir
-
-
-def test_min_repeat_ms():
- tmp = tempdir()
- filename = tmp.relpath("log")
-
- @tvm.register_global_func
- def my_debug(filename):
- """one call lasts for 100 ms and writes one character to a file"""
- time.sleep(0.1)
- with open(filename, "a") as fout:
- fout.write("c")
-
- X = te.compute((), lambda: tvm.tirx.call_packed("my_debug", filename))
- func = tvm.tirx.compile(te.create_function([X])).mod
-
- x = tvm.runtime.empty((), dtype="int32")
- ftimer = func.time_evaluator(func.entry_name, tvm.cpu(), number=1,
repeat=1)
- ftimer(x)
-
- with open(filename) as fin:
- ct = len(fin.readline())
-
- assert ct == 2
-
- ftimer = func.time_evaluator(func.entry_name, tvm.cpu(), number=1,
repeat=1, min_repeat_ms=1000)
- ftimer(x)
-
- # make sure we get more than 10 calls
- with open(filename) as fin:
- ct = len(fin.readline())
-
- assert ct > 10 + 2
-
-
-def test_benchmark_result():
- r = BenchmarkResult([1, 2, 2, 5])
- assert r.mean == 2.5
- assert r.median == 2.0
- assert r.min == 1
- assert r.max == 5
- assert r.std == 1.5
-
-
-if __name__ == "__main__":
- test_min_repeat_ms()
- test_benchmark_result()
diff --git a/tests/python/runtime/test_runtime_memory_scope.py
b/tests/python/runtime/test_runtime_memory_scope.py
index 006bbe7f99..c21c2c8f4c 100644
--- a/tests/python/runtime/test_runtime_memory_scope.py
+++ b/tests/python/runtime/test_runtime_memory_scope.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-"""Canonical memory scopes cross the Python, compiler, and runtime
boundaries."""
+"""Canonical runtime memory-scope identity, allocation, and registration."""
import os
import subprocess
@@ -29,8 +29,6 @@ import tvm_ffi
import tvm
from tvm.runtime import MemoryScope
-from tvm.script import ir as I
-from tvm.script import tirx as T
def new_scope():
@@ -58,6 +56,8 @@ def test_identity_and_registration():
MemoryScope.register("")
with pytest.raises(TypeError):
MemoryScope.register(123)
+ with pytest.raises(ValueError):
+ MemoryScope("")
prefix = "test.memory_scope.nul." + uuid.uuid4().hex
with pytest.raises(ValueError):
MemoryScope.register(prefix + "\x00scope")
@@ -65,106 +65,13 @@ def test_identity_and_registration():
MemoryScope(prefix)
-def test_shared_cpp_python_registry_and_typed_conversion():
+def test_shared_cpp_python_registry():
name = "test.memory_scope.cpp." + uuid.uuid4().hex
cpp_scope = tvm_ffi.get_global_func("runtime.MemoryScopeRegister")(name)
assert MemoryScope(name).same_as(cpp_scope)
with pytest.raises(ValueError):
tvm_ffi.get_global_func("runtime.MemoryScopeRegister")(name)
- custom = new_scope()
- for scope in (custom, custom.name):
- ptr_ty = tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope)
- tensor = tvm.tirx.decl_tensor((4,), "float32", scope=scope)
- tensor_ty = tvm.tirx.TensorType((4,), "float32", scope)
- assert tensor_ty.scope.same_as(custom)
- assert ptr_ty.scope.same_as(custom)
- assert tensor.ty.scope.same_as(custom)
- assert tensor.scope().same_as(custom)
- assert tensor.data.ty.scope.same_as(custom)
-
- unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
- for construct in (
- lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), unknown),
- lambda: tvm.tirx.decl_tensor((4,), "float32", scope=unknown),
- lambda: T.Tensor((4,), "float32", scope=unknown),
- ):
- with pytest.raises(ValueError):
- construct()
- with pytest.raises(ValueError):
- MemoryScope(unknown)
-
-
-def test_global_defaults():
- global_scope = MemoryScope("global")
- assert
tvm.ir.PtrType(tvm.ir.PrimType("float32")).scope.same_as(global_scope)
- assert tvm.tirx.decl_tensor((4,), "float32").scope().same_as(global_scope)
- for scope in ("global", global_scope):
- assert tvm.ir.PtrType(tvm.ir.PrimType("float32"),
scope).scope.same_as(global_scope)
- assert tvm.tirx.decl_tensor((4,), "float32",
scope=scope).scope().same_as(global_scope)
- assert T.Tensor((4,), "float32",
scope=scope).scope.same_as(global_scope)
- assert tvm.relax.VDevice("llvm",
memory_scope=scope).memory_scope.same_as(global_scope)
- for construct in (
- lambda: MemoryScope(""),
- lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), ""),
- lambda: tvm.tirx.decl_tensor((4,), "float32", scope=""),
- lambda: T.Tensor((4,), "float32", scope=""),
- lambda: tvm.relax.VDevice("llvm", memory_scope=""),
- lambda: tvm.runtime.empty((4,), mem_scope=""),
- ):
- with pytest.raises(ValueError):
- construct()
-
-
-def test_json_and_script_roundtrip():
- custom = new_scope()
- for scope in (MemoryScope("global"), MemoryScope("shared.dyn"), custom):
- assert tvm.ir.load_json(tvm.ir.save_json(scope)).same_as(scope)
- ty = T.Tensor((4,), "float32", scope=scope)
- restored = tvm.ir.load_json(tvm.ir.save_json(ty))
- assert restored.scope.same_as(scope)
- tvm.ir.assert_structural_equal(ty, restored)
-
- @T.function
- def func(A: T.Tensor((4,), "float32", scope=scope)):
- A[0] = T.float32(1)
-
- script = func.script()
- if scope != MemoryScope("global"):
- assert f'scope="{scope.name}"' in script
- restored_func = tvm.script.from_source(script, extra_vars={"T": T})
- tvm.ir.assert_structural_equal(func, restored_func)
- assert restored_func.params[0].scope().same_as(scope)
-
-
-def test_canonical_constructor_aliases_and_reflected_scopes():
- assert T.Ptr is tvm.ir.PtrType
- assert T.Tensor is tvm.tirx.TensorType
- custom = new_scope()
- for scope in (MemoryScope("global"), custom):
- pointer = T.Ptr(T.float32, scope)
- source = pointer.script()
- expected = (
- "T.Ptr(T.float32)"
- if scope == MemoryScope("global")
- else f'T.Ptr(T.float32, "{scope.name}")'
- )
- assert source == expected
- restored = eval(source, {"T": T, "I": I})
- tvm.ir.assert_structural_equal(pointer, restored)
- assert restored.scope.same_as(scope)
- tensor = T.Tensor((4,), T.float32, scope, byte_offset=8, align=128,
layout="default")
- assert tensor.scope.same_as(scope)
- assert tensor.elem_offset.value == 2
- assert tensor.data_alignment == 128
- restored_tensor = eval(tensor.script(), {"T": T, "I": I})
- tvm.ir.assert_structural_equal(tensor, restored_tensor)
- for scope in (custom, custom.name):
- tensor = T.Tensor((4,), T.float32, scope=scope)
- assert tensor.scope.same_as(custom)
- with pytest.raises(ValueError):
- T.Ptr(T.float32, "")
-
def test_runtime_tensor_allocation():
expected = np.arange(4, dtype="float32")
@@ -172,8 +79,9 @@ def test_runtime_tensor_allocation():
tensor = tvm.runtime.tensor(expected, tvm.cpu(), mem_scope=scope)
np.testing.assert_array_equal(tensor.numpy(), expected)
np.testing.assert_array_equal(tensor.copyto(tvm.cpu(),
mem_scope=scope).numpy(), expected)
- with pytest.raises(ValueError):
- tvm.runtime.empty((4,), mem_scope="test.memory_scope.unregistered")
+ for scope in ("", "test.memory_scope.unregistered"):
+ with pytest.raises(ValueError):
+ tvm.runtime.empty((4,), mem_scope=scope)
with pytest.raises(ValueError, match="Unsupported memory scope"):
tvm.runtime.empty((4,), mem_scope=new_scope())
@@ -277,13 +185,3 @@ for scope in scopes:
check=True,
timeout=60,
)
-
-
-def test_relax_allocation_scope_operands():
- scope = MemoryScope("global")
- for make_call in (
- lambda value: tvm.relax.op.builtin.alloc_tensor((4,), "float32", 0,
value),
- lambda value: tvm.relax.op.memory.alloc_storage((16,), 0, value,
"float32"),
- lambda value: tvm.relax.op.vm.alloc_storage((16,), 0, "float32",
value),
- ):
- tvm.ir.assert_structural_equal(make_call(scope), make_call(scope.name))
diff --git a/tests/python/runtime/test_runtime_module.py
b/tests/python/runtime/test_runtime_module.py
new file mode 100644
index 0000000000..443e45511b
--- /dev/null
+++ b/tests/python/runtime/test_runtime_module.py
@@ -0,0 +1,263 @@
+# 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.
+"""Runtime module loading, properties, exports, timing, and Executable
wrappers."""
+
+import os
+import subprocess
+import sys
+import time
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.runtime import Executable
+from tvm.runtime.module import BenchmarkResult
+from tvm.support import cc
+
+
[email protected](scope="module")
+def native_source():
+ return r"""
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/function.h>
+void add_one(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
+ for (int64_t i = 0; i < a.size(0); ++i) {
+ static_cast<float*>(b.data_ptr())[i] =
static_cast<float*>(a.data_ptr())[i] + 1;
+ }
+}
+int64_t answer() { return 42; }
+void callback(tvm::ffi::Function func) { func(); }
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(main, add_one);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(answer, answer);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(callback, callback);
+"""
+
+
[email protected](scope="module")
+def native_compile_options():
+ return [
+ "-std=c++17",
+ "-I" + tvm_ffi.libinfo.find_include_path(),
+ "-I" + tvm_ffi.libinfo.find_dlpack_include_path(),
+ ]
+
+
[email protected]
+def source_module(native_source):
+ return tvm.runtime._ffi_api.CSourceModuleCreate(
+ native_source, "cc", ["main", "answer", "callback"], None
+ )
+
+
[email protected](scope="module")
+def native_object(tmp_path_factory, native_source, native_compile_options):
+ directory = tmp_path_factory.mktemp("runtime_module")
+ source = directory / "module.cc"
+ source.write_text(native_source)
+ path = directory / "module.o"
+ cc.create_shared(str(path), [str(source)], options=["-c",
*native_compile_options])
+ return path
+
+
[email protected](scope="module")
+def native_library(native_object):
+ path = native_object.with_suffix(".so")
+ cc.create_shared(str(path), [str(native_object)])
+ return path
+
+
[email protected](scope="module")
+def system_library_source(native_source):
+ return (
+ native_source
+ + r"""
+#include <tvm/ffi/extra/c_env_api.h>
+struct RegisterSystemLibrary {
+ RegisterSystemLibrary() {
+ TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_main",
(void*)__tvm_ffi_main);
+ TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_answer",
(void*)__tvm_ffi_answer);
+ }
+};
+static RegisterSystemLibrary register_system_library;
+"""
+ )
+
+
+def test_module_properties(source_module, native_library):
+ assert source_module.is_binary_serializable()
+ assert not source_module.is_runnable()
+ assert source_module.is_compilation_exportable()
+ loaded = tvm.runtime.load_module(str(native_library))
+ assert loaded.kind == "library"
+ assert loaded.is_runnable()
+ assert not loaded.is_compilation_exportable()
+ assert loaded.implements_function("answer")
+ assert loaded.get_function("answer")() == 42
+
+
+def test_dso_module_load(native_library):
+ module = tvm.runtime.load_module(str(native_library))
+ a = tvm.runtime.tensor(np.arange(10, dtype="float32"))
+ b = tvm.runtime.empty((10,), "float32")
+ module(a, b)
+ np.testing.assert_equal(b.numpy(), a.numpy() + 1)
+ # Verify loading and execution in a fresh process using only the runtime
library.
+ subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ """
+import sys
+import numpy as np
+import tvm
+assert tvm.base._RUNTIME_ONLY
+module = tvm.runtime.load_module(sys.argv[1])
+a = tvm.runtime.tensor(np.arange(10, dtype="float32"))
+b = tvm.runtime.empty((10,), "float32")
+module(a, b)
+np.testing.assert_equal(b.numpy(), a.numpy() + 1)
+""",
+ str(native_library),
+ ],
+ env={**os.environ, "TVM_USE_RUNTIME_LIB": "1"},
+ check=True,
+ )
+
+
+def test_import_static_library(source_module, native_compile_options,
tmp_path):
+ source = tmp_path / "other.cc"
+ source.write_text(
+ "#include <tvm/ffi/function.h>\n"
+ "int other() { return 17; }\n"
+ "TVM_FFI_DLL_EXPORT_TYPED_FUNC(other, other);\n"
+ )
+ obj = tmp_path / "other.o"
+ cc.create_shared(str(obj), [str(source)], options=["-c",
*native_compile_options])
+ imported = tvm.runtime.load_static_library(str(obj), ["other"])
+ assert imported.implements_function("other")
+ assert imported.is_compilation_exportable()
+ source_module.import_module(imported)
+ path = tmp_path / "combined.so"
+ source_module.export_library(str(path), options=native_compile_options)
+ loaded = tvm.runtime.load_module(str(path))
+ assert loaded.kind == "library"
+ assert len(loaded.imports) == 0
+ for name, expected in (("answer", 42), ("other", 17)):
+ assert loaded.implements_function(name)
+ assert loaded.get_function(name)() == expected
+ assert not loaded.is_compilation_exportable()
+
+
[email protected](sys.platform == "win32", reason="system library
registration on Unix")
+def test_system_library(system_library_source, native_compile_options,
tmp_path):
+ source = tmp_path / "system.cc"
+ source.write_text(system_library_source)
+ path = tmp_path / "system.so"
+ cc.create_shared(str(path), [str(source)], options=native_compile_options)
+ subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ """
+import ctypes
+import sys
+import numpy as np
+import tvm
+handle = ctypes.CDLL(sys.argv[1])
+module = tvm.runtime.system_lib()
+assert module["answer"]() == 42
+a = tvm.runtime.tensor(np.arange(10, dtype="float32"))
+b = tvm.runtime.empty((10,), "float32")
+module["main"](a, b)
+np.testing.assert_equal(b.numpy(), a.numpy() + 1)
+""",
+ str(path),
+ ],
+ check=True,
+ )
+
+
+def test_min_repeat_ms(native_library):
+ calls = []
+
+ def callback():
+ time.sleep(0.1)
+ calls.append(1)
+
+ module = tvm.runtime.load_module(str(native_library))
+ module.time_evaluator("callback", tvm.cpu(), number=1, repeat=1)(callback)
+ assert len(calls) == 2
+ module.time_evaluator("callback", tvm.cpu(), number=1, repeat=1,
min_repeat_ms=1000)(callback)
+ assert len(calls) > 12
+
+
+def test_benchmark_result():
+ result = BenchmarkResult([1, 2, 2, 5])
+ assert result.mean == 2.5
+ assert result.median == 2.0
+ assert result.min == 1
+ assert result.max == 5
+ assert result.std == 1.5
+
+
+def test_executable_init(source_module):
+ executable = Executable(source_module)
+ assert executable.mod is source_module
+ assert executable._jitted_mod is None
+
+
+def test_executable_getitem(source_module, native_compile_options):
+ executable = Executable(source_module)
+ executable.jit(options=native_compile_options)
+ assert executable["answer"]() == 42
+
+
+def test_executable_jit_already_runnable(native_library):
+ module = tvm.runtime.load_module(str(native_library))
+ executable = Executable(module)
+ assert executable.jit() is module
+ assert executable.jit(force_recompile=True) is module
+ assert executable["answer"]() == 42
+
+
+def test_executable_jit_force_recompile(source_module, native_compile_options):
+ executable = Executable(source_module)
+ first = executable.jit(options=native_compile_options)
+ assert executable.jit() is first
+ rebuilt = executable.jit(force_recompile=True,
options=native_compile_options)
+ assert rebuilt is not first
+ assert first["answer"]() == rebuilt["answer"]() == 42
+
+
[email protected]("use_workspace", [False, True])
+def test_executable_export_library(source_module, native_compile_options,
tmp_path, use_workspace):
+ kwargs = {"options": native_compile_options}
+ if use_workspace:
+ workspace = tmp_path / "workspace"
+ workspace.mkdir()
+ kwargs["workspace_dir"] = str(workspace)
+ path = tmp_path / "module.so"
+ Executable(source_module).export_library(str(path), **kwargs)
+ assert path.is_file()
+ assert tvm.runtime.load_module(str(path))["answer"]() == 42
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_module_export.py
b/tests/python/runtime/test_runtime_module_export.py
deleted file mode 100644
index 65117e958c..0000000000
--- a/tests/python/runtime/test_runtime_module_export.py
+++ /dev/null
@@ -1,69 +0,0 @@
-# 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 pytest
-
-import tvm
-import tvm.testing
-from tvm.support import utils
-from tvm.testing import env
-
-
[email protected](not env.has_llvm(), reason="need llvm")
-def test_import_static_library():
- from tvm import te
-
- # Generate two LLVM modules.
- A = te.placeholder((1024,), name="A")
- B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
- irmod0 = tvm.IRModule.from_expr(te.create_function([A,
B]).with_attr("global_symbol", "myadd0"))
- irmod1 = tvm.IRModule.from_expr(te.create_function([A,
B]).with_attr("global_symbol", "myadd1"))
-
- mod0 = tvm.tirx.compile(irmod0, target="llvm").mod
- mod1 = tvm.tirx.compile(irmod1, target="llvm").mod
-
- assert mod0.implements_function("myadd0")
- assert mod1.implements_function("myadd1")
- assert mod1.is_compilation_exportable()
-
- # mod1 is currently an 'llvm' module.
- # Save and reload it as a vanilla 'static_library'.
- temp = utils.tempdir()
- mod1_o_path = temp.relpath("mod1.o")
- mod1.write_to_file(mod1_o_path)
- mod1_o = tvm.runtime.load_static_library(mod1_o_path, ["myadd1"])
- assert mod1_o.implements_function("myadd1")
- assert mod1_o.is_compilation_exportable()
-
- # Import mod1 as a static library into mod0 and compile to its own DSO.
- mod0.import_module(mod1_o)
- mod0_dso_path = temp.relpath("mod0.so")
- tvm.runtime.Executable(mod0).export_library(mod0_dso_path)
-
- # The imported mod1 is statically linked into mod0.
- loaded_lib = tvm.runtime.load_module(mod0_dso_path)
- assert loaded_lib.kind == "library"
- assert len(loaded_lib.imports) == 0
- assert loaded_lib.implements_function("myadd0")
- assert loaded_lib.get_function("myadd0")
- assert loaded_lib.implements_function("myadd1")
- assert loaded_lib.get_function("myadd1")
- assert not loaded_lib.is_compilation_exportable()
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_module_load.py
b/tests/python/runtime/test_runtime_module_load.py
deleted file mode 100644
index dd09aa1fad..0000000000
--- a/tests/python/runtime/test_runtime_module_load.py
+++ /dev/null
@@ -1,246 +0,0 @@
-# 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...")
-"""
-
-
[email protected](not env.has_llvm(), reason="need llvm")
-def test_dso_module_load():
- dtype = "int64"
- temp = utils.tempdir()
-
- def save_object(names):
- n = te.var("n")
- Ab = tvm.tirx.decl_tensor((n,), dtype)
- i = te.var("i")
- # for i in 0 to n-1:
- stmt = tvm.ir.For(
- i,
- 0,
- n - 1,
- tvm.ir.ForKind.DEFAULT,
- tvm.ir.TensorStore(Ab, [i + 1], tvm.tirx.TensorLoad(Ab, [i]) + 1),
- )
- mod = tvm.IRModule.from_expr(
- tvm.tirx.Function([Ab], stmt).with_attr("global_symbol", "main")
- )
- m = tvm.tirx.compile(mod, target="llvm").mod
- 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}"
-
-
[email protected]
[email protected](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_function([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(str(A.dtype)), dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=str(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(str(A.dtype)), dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=str(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)
-
-
[email protected](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_function([A,
B]).with_attr("global_symbol", "myadd1"))
- mod2 = tvm.IRModule.from_expr(te.create_function([A,
B]).with_attr("global_symbol", "myadd2"))
-
- def check_llvm():
- dev = tvm.cpu(0)
- temp = utils.tempdir()
- fadd1 = tvm.tirx.compile(mod1, "llvm").mod
- fadd2 = tvm.tirx.compile(mod2, "llvm").mod
- 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(str(A.dtype)), dev)
- b = tvm.runtime.tensor(np.zeros(nn, dtype=str(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.compile(mod1.with_attr("system_lib_prefix", ""),
"llvm").mod
- fadd2 = tvm.tirx.compile(mod2.with_attr("system_lib_prefix", ""),
"llvm").mod
- 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(str(A.dtype)), dev)
- b = tvm.runtime.tensor(np.zeros(nn, dtype=str(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()
diff --git a/tests/python/runtime/test_runtime_module_property.py
b/tests/python/runtime/test_runtime_module_property.py
deleted file mode 100644
index 1f1ea4be97..0000000000
--- a/tests/python/runtime/test_runtime_module_property.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# 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 tvm
-import tvm.runtime._ffi_api
-import tvm.target._ffi_api
-from tvm import te
-
-
-def checker(mod, expected):
- assert mod.is_binary_serializable() == expected["is_binary_serializable()"]
- assert mod.is_runnable() == expected["is_runnable"]
- assert mod.is_compilation_exportable() ==
expected["is_compilation_exportable()"]
-
-
-def create_csource_module():
- return tvm.runtime._ffi_api.CSourceModuleCreate("", "cc", [], None)
-
-
-def create_llvm_module():
- A = te.placeholder((1024,), name="A")
- B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
- return tvm.tirx.compile(te.create_function([A, B]), target="llvm").mod
-
-
-def test_property():
- checker(
- create_csource_module(),
- expected={
- "is_binary_serializable()": True,
- "is_runnable": False,
- "is_compilation_exportable()": True,
- },
- )
-
- checker(
- create_llvm_module(),
- expected={
- "is_binary_serializable()": False,
- "is_runnable": True,
- "is_compilation_exportable()": True,
- },
- )
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_nd_array.py
b/tests/python/runtime/test_runtime_nd_array.py
deleted file mode 100644
index e85c89fb22..0000000000
--- a/tests/python/runtime/test_runtime_nd_array.py
+++ /dev/null
@@ -1,254 +0,0 @@
-# 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: F811
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-
-
-def test_1d_full_view_of_1d_arr():
- """Tensor::CreateView may return the same array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([1024])
- np_expected = np_input
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_1d_view_of_first_half_of_1d_arr():
- """Tensor::CreateView may return a subset of an array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([512])
- np_expected = np_input[0:512]
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_1d_view_of_first_half_of_1d_arr():
- """Subset returned by Tensor::CreateView may have a byte offset"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([512], relative_byte_offset=512 * 4)
- np_expected = np_input[512:1024]
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_view_larger_than_original_is_invalid():
- """Subset may not be larger than the original array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- with pytest.raises(ValueError, match="the Tensor being viewed only
contains 4096 bytes"):
- tvm_input._create_view([2048])
-
-
-def test_view_entirely_outside_bounds_of_original_is_invalid():
- """The byte_offset may not place a view outside the original array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- with pytest.raises(ValueError, match="would occupy bytes 8192 <= i_byte <
12288"):
- tvm_input._create_view([1024], relative_byte_offset=2048 * 4)
-
-
-def test_view_partially_outside_bounds_of_original_is_invalid():
- """The byte_offset may not place any elements of a view outside the
original array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- with pytest.raises(ValueError, match="would occupy bytes 2048 <= i_byte <
6144"):
- tvm_input._create_view([1024], relative_byte_offset=512 * 4)
-
-
-def test_subview_first_half_of_first_half():
- """Tensor::CreateView be applied to a view
-
- The first view is at element offset 0 (byte offset 0). The second
- view is at element offset 0 (byte offset 0) relative to the first
- view, or element offset 0 (byte offset 0) relative to the original
- array.
-
- """
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_view = tvm_input._create_view(
- [512],
- relative_byte_offset=0,
- )
- tvm_subview = tvm_view._create_view(
- [256],
- relative_byte_offset=0,
- )
- np_expected = np_input[0:512][0:256]
-
- np.testing.assert_equal(tvm_subview.numpy(), np_expected)
-
-
-def test_subview_first_half_of_second_half():
- """Tensor::CreateView be applied to a view
-
- The first view is at element offset 512 (byte offset 2048). The
- second view is at element offset 0 (byte offset 0) relative to the
- first view, or element offset 512 (byte offset 2048) relative to
- the original array.
-
- """
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_view = tvm_input._create_view(
- [512],
- relative_byte_offset=512 * 4,
- )
- tvm_subview = tvm_view._create_view(
- [256],
- relative_byte_offset=0,
- )
- np_expected = np_input[512:1024][0:256]
-
- np.testing.assert_equal(tvm_subview.numpy(), np_expected)
-
-
-def test_subview_second_half_of_first_half():
- """Tensor::CreateView be applied to a view
-
- The first view is at element offset 0 (byte offset 0). The second
- view is at element offset 256 (byte offset 1024) relative to the
- first view, or element offset 256 (byte offset 1024) relative to
- the original array.
-
- """
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_view = tvm_input._create_view(
- [512],
- relative_byte_offset=0,
- )
- tvm_subview = tvm_view._create_view(
- [256],
- relative_byte_offset=256 * 4,
- )
- np_expected = np_input[0:512][256:512]
-
- np.testing.assert_equal(tvm_subview.numpy(), np_expected)
-
-
-def test_subview_second_half_of_second_half():
- """Tensor::CreateView be applied to a view
-
- The first view is at element offset 512 (byte offset 2048). The
- second view is at element offset 256 (byte offset 1024) relative
- to the first view, or element offset 768 (byte offset 3072)
- relative to the original array.
-
- """
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_view = tvm_input._create_view(
- [512],
- relative_byte_offset=512 * 4,
- )
- tvm_subview = tvm_view._create_view(
- [256],
- relative_byte_offset=256 * 4,
- )
- np_expected = np_input[512:1024][256:512]
-
- np.testing.assert_equal(tvm_subview.numpy(), np_expected)
-
-
-def test_subview_must_be_in_range_of_immediate_parent():
- """Bounds-checking is applied relative to the Tensor
-
- The first view is at location and covers bytes [0,2048). The
- subview would occupy bytes [2048, 4096), and raises an error as
- this is outside the range of the view.
-
- """
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_view = tvm_input._create_view(
- [512],
- relative_byte_offset=0,
- )
-
- with pytest.raises(ValueError, match="would occupy bytes 2048 <= i_byte <
4096"):
- tvm_view._create_view(
- [512],
- relative_byte_offset=512 * 4,
- )
-
-
-def test_2d_view_into_1d_arr():
- """Tensor::CreateView may change the dimensionality of an array"""
- np_input = np.arange(1024, dtype="int32")
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([32, 32])
- np_expected = np_input.reshape(32, 32)
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_2d_full_view_into_2d_arr():
- """Tensor::CreateView may change the shape of an array"""
- np_input = np.arange(1024, dtype="int32").reshape(32, 32)
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([16, 64])
- np_expected = np_input.reshape(16, 64)
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_2d_view_of_first_half_of_2d_arr():
- """Tensor::CreateView may return a multi-dimensional view"""
- np_input = np.arange(1024, dtype="int32").reshape(32, 32)
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([16, 32])
- np_expected = np_input[0:16, :]
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-def test_2d_view_of_second_half_of_2d_arr():
- """Tensor::CreateView may return a multi-dimensional view with byte
offset"""
- np_input = np.arange(1024, dtype="int32").reshape(32, 32)
- tvm_input = tvm.runtime.tensor(np_input)
-
- tvm_output = tvm_input._create_view([16, 32], relative_byte_offset=32 * 16
* 4)
- np_expected = np_input[16:32, :]
-
- np.testing.assert_equal(tvm_output.numpy(), np_expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/runtime/test_runtime_tensor.py
b/tests/python/runtime/test_runtime_tensor.py
new file mode 100644
index 0000000000..de2264700d
--- /dev/null
+++ b/tests/python/runtime/test_runtime_tensor.py
@@ -0,0 +1,84 @@
+# 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.
+"""Runtime tensor views and imported storage."""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+
+
[email protected](
+ "input_shape,view_shape,offset",
+ [
+ ((1024,), (32, 32), 0),
+ ((1024,), (512,), 0),
+ ((32, 32), (16, 32), 512),
+ ],
+)
+def test_tensor_view(input_shape, view_shape, offset):
+ values = np.arange(1024, dtype="int32").reshape(input_shape)
+ tensor = tvm.runtime.tensor(values)
+ view = tensor._create_view(view_shape, relative_byte_offset=offset * 4)
+ expected = values.ravel()[offset : offset +
np.prod(view_shape)].reshape(view_shape)
+ np.testing.assert_equal(view.numpy(), expected)
+
+
[email protected]("parent_offset", [0, 512])
[email protected]("child_offset", [0, 256])
+def test_nested_view(parent_offset, child_offset):
+ values = np.arange(1024, dtype="int32")
+ tensor = tvm.runtime.tensor(values)
+ parent = tensor._create_view((512,), relative_byte_offset=parent_offset *
4)
+ child = parent._create_view((256,), relative_byte_offset=child_offset * 4)
+ expected = values[parent_offset : parent_offset + 512][child_offset :
child_offset + 256]
+ np.testing.assert_equal(child.numpy(), expected)
+
+
[email protected](
+ "size,offset,message",
+ [
+ (2048, 0, "the Tensor being viewed only contains 4096 bytes"),
+ (1024, 2048, "would occupy bytes 8192 <= i_byte < 12288"),
+ (1024, 512, "would occupy bytes 2048 <= i_byte < 6144"),
+ ],
+)
+def test_view_bounds(size, offset, message):
+ tensor = tvm.runtime.tensor(np.arange(1024, dtype="int32"))
+ with pytest.raises(ValueError, match=message):
+ tensor._create_view((size,), relative_byte_offset=offset * 4)
+
+
+def test_subview_must_be_in_range_of_immediate_parent():
+ tensor = tvm.runtime.tensor(np.arange(1024, dtype="int32"))
+ parent = tensor._create_view((512,))
+ with pytest.raises(ValueError, match="would occupy bytes 2048 <= i_byte <
4096"):
+ parent._create_view((512,), relative_byte_offset=512 * 4)
+
+
+def test_from_dlpack_shape_one():
+ torch = pytest.importorskip("torch")
+ values = torch.randn(1, 16)
+ tensor = tvm.runtime.from_dlpack(torch.utils.dlpack.to_dlpack(values))
+ np.testing.assert_equal(tensor.numpy(), values.numpy())
+ view = tensor._create_view((2, 8))
+ np.testing.assert_equal(view.numpy(), values.numpy().reshape(2, 8))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_tensor.py
b/tests/python/tirx-base/test_tir_tensor.py
index 70d5dd8e04..4c2eb95b4a 100644
--- a/tests/python/tirx-base/test_tir_tensor.py
+++ b/tests/python/tirx-base/test_tir_tensor.py
@@ -16,11 +16,14 @@
# under the License.
# ruff: noqa: E741, F401, F841
+import uuid
+
import numpy as np
import pytest
import tvm
import tvm.testing
+from tvm.runtime import MemoryScope
from tvm.script import tirx as T
@@ -258,5 +261,47 @@ def test_tensor_flatten(shape, strides, expected):
tvm.ir.assert_structural_equal(flat.ty.shape, [T.int32(expected)])
+def test_memory_scope_typed_conversion():
+ custom = MemoryScope.register("test.tensor.scope." + uuid.uuid4().hex)
+ for scope in (custom, custom.name):
+ ptr_ty = tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope)
+ tensor = tvm.tirx.decl_tensor((4,), "float32", scope=scope)
+ tensor_ty = tvm.tirx.TensorType((4,), "float32", scope)
+ assert tensor_ty.scope.same_as(custom)
+ assert ptr_ty.scope.same_as(custom)
+ assert tensor.ty.scope.same_as(custom)
+ assert tensor.scope().same_as(custom)
+ assert tensor.data.ty.scope.same_as(custom)
+
+ unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
+ for construct in (
+ lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), unknown),
+ lambda: tvm.tirx.decl_tensor((4,), "float32", scope=unknown),
+ lambda: T.Tensor((4,), "float32", scope=unknown),
+ ):
+ with pytest.raises(ValueError):
+ construct()
+ with pytest.raises(ValueError):
+ MemoryScope(unknown)
+
+
+def test_memory_scope_global_defaults():
+ global_scope = MemoryScope("global")
+ assert
tvm.ir.PtrType(tvm.ir.PrimType("float32")).scope.same_as(global_scope)
+ assert tvm.tirx.decl_tensor((4,), "float32").scope().same_as(global_scope)
+ for scope in ("global", global_scope):
+ assert tvm.ir.PtrType(tvm.ir.PrimType("float32"),
scope).scope.same_as(global_scope)
+ assert tvm.tirx.decl_tensor((4,), "float32",
scope=scope).scope().same_as(global_scope)
+ assert T.Tensor((4,), "float32",
scope=scope).scope.same_as(global_scope)
+ for construct in (
+ lambda: MemoryScope(""),
+ lambda: tvm.ir.PtrType(tvm.ir.PrimType("float32"), ""),
+ lambda: tvm.tirx.decl_tensor((4,), "float32", scope=""),
+ lambda: T.Tensor((4,), "float32", scope=""),
+ ):
+ with pytest.raises(ValueError):
+ construct()
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx/script/test_tirx_script_tensors.py
b/tests/python/tirx/script/test_tirx_script_tensors.py
index 1a16265c0a..b36f91299a 100644
--- a/tests/python/tirx/script/test_tirx_script_tensors.py
+++ b/tests/python/tirx/script/test_tirx_script_tensors.py
@@ -19,6 +19,7 @@
import math
import sys
+import uuid
import pytest
import tvm_ffi
@@ -29,6 +30,7 @@ import tvm.testing
from tirx.layout_fixtures import tcgen05_register_layout
from tvm.backend.cuda.axis import TCol, TLane, laneid, warpid
from tvm.ir import TensorRegion, assert_structural_equal
+from tvm.runtime import MemoryScope
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
@@ -1435,3 +1437,58 @@ def test_roundtrip_tmem_decl_tensor():
tmem_decl = next(decl for decl in decls if decl.var.scope() ==
tvm.runtime.MemoryScope("tmem"))
assert isinstance(tmem_decl.value.args[0], tvm.ir.TensorLoad)
assert len(tmem_decl.value.ty_args) == 1
+
+
+def test_memory_scope_script_roundtrip():
+ custom = MemoryScope.register("test.script.scope." + uuid.uuid4().hex)
+ for scope in (MemoryScope("global"), MemoryScope("shared.dyn"), custom):
+
+ @T.function
+ def func(A: T.Tensor((4,), T.float32, scope=scope)):
+ A[0] = T.float32(1)
+
+ script = func.script()
+ if scope == MemoryScope("global"):
+ assert "A: T.Tensor((4,), T.float32)" in script
+ else:
+ assert f'A: T.Tensor((4,), T.float32, scope="{scope.name}")' in
script
+ restored_func = tvm.script.from_source(script, extra_vars={"T": T})
+ tvm.ir.assert_structural_equal(func, restored_func)
+ assert restored_func.params[0].scope().same_as(scope)
+
+
+def test_memory_scope_canonical_constructor_aliases_and_reflected_scopes():
+ custom = MemoryScope.register("test.script.scope." + uuid.uuid4().hex)
+ for scope in (MemoryScope("global"), custom, custom.name):
+
+ @T.function
+ def func(
+ pointer: T.Ptr(T.float32, scope),
+ tensor: T.Tensor((4,), T.float32, scope, byte_offset=8, align=128,
layout="default"),
+ ):
+ tensor[0] = T.float32(1)
+
+ assert T.Ptr is tvm.ir.PtrType
+ assert T.Tensor is tvm.tirx.TensorType
+ expected_scope = MemoryScope(scope) if isinstance(scope, str) else
scope
+ pointer_type = func.params[0].ty
+ tensor_type = func.params[1].ty
+ assert pointer_type.scope.same_as(expected_scope)
+ assert tensor_type.scope.same_as(expected_scope)
+ assert tensor_type.elem_offset.value == 2
+ assert tensor_type.data_alignment == 128
+ script = func.script()
+ if expected_scope == MemoryScope("global"):
+ assert "pointer: T.Ptr(T.float32)" in script
+ assert "tensor: T.Tensor((4,), T.float32, elem_offset=" in script
+ else:
+ assert f'pointer: T.Ptr(T.float32, "{expected_scope.name}")' in
script
+ assert f'tensor: T.Tensor((4,), T.float32,
scope="{expected_scope.name}", ' in script
+ assert "elem_offset=T.int64(2), data_alignment=128" in script
+ assert 'layout=T.TileLayout.from_iters(shard=[T.Iter(4, 1, "m")],
replica=[])' in script
+ restored = tvm.script.from_source(script, extra_vars={"T": T})
+ tvm.ir.assert_structural_equal(func, restored)
+ assert restored.params[0].ty.scope.same_as(expected_scope)
+ assert restored.params[1].scope().same_as(expected_scope)
+ with pytest.raises(ValueError):
+ T.Ptr(T.float32, "")
