chore: import upstream snapshot with attribution

tests/python/runtime/test_executable.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.
+"""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 tirx as T
+
+
+@tvm.script.ir_module
+class MyModule:
+    @T.prim_func(s_tir=True)
+    def add(
+        A: T.Buffer((10,), "float32"),
+        B: T.Buffer((10,), "float32"),
+        C: T.Buffer((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.build(MyModule, target="llvm")
+    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.build(MyModule, target="llvm")
+    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.build(MyModule, target="llvm")
+    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.build(MyModule, target="llvm")
+    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.build(MyModule, target="llvm")
+    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.build(MyModule, target=target)
+
+    # 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.build(MyModule, target=target)
+
+    # 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()

tests/python/runtime/test_rpc_base.py

@@ -0,0 +1,43 @@
+# 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 random
+
+import pytest
+
+from tvm.rpc import base
+
+
+@pytest.mark.parametrize("device_key", ["16e995b6", "127.0.0.1:5555"])
+def test_rpc_base_random_key(device_key):
+    random.seed(0)
+    key = base.random_key(device_key)
+    assert key.startswith(device_key)
+    res_device_key, _ = base.split_random_key(key)
+    assert device_key == res_device_key
+    # start with seed 0 as well, but use cmap arg(a conflict map)
+    # to generate another unique random key
+    random.seed(0)
+    new_key = base.random_key(device_key, cmap={key})
+    assert key != new_key
+    assert new_key.startswith(device_key)
+    res_device_key2, _ = base.split_random_key(new_key)
+    assert device_key == res_device_key2
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/runtime/test_runtime_device_api.py

@@ -0,0 +1,52 @@
+# 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 os
+import subprocess
+import sys
+
+import tvm
+import tvm.testing
+
+
+def test_check_if_device_exists():
+    """kExist can be checked when no devices are present
+
+    This test uses `CUDA_VISIBLE_DEVICES` to disable any CUDA-capable
+    GPUs from being accessed by the subprocess.  Within the
+    subprocess, the CUDA driver cannot be initialized.  While most
+    functionality of CUDADeviceAPI would raise an exception, the
+    `kExist` property can still be checked.
+
+    """
+
+    cmd = [
+        sys.executable,
+        "-c",
+        "import tvm; tvm.device('cuda').exist",
+    ]
+    subprocess.check_call(
+        cmd,
+        env={
+            **os.environ,
+            "CUDA_VISIBLE_DEVICES": "",
+        },
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/runtime/test_runtime_dlpack.py

@@ -0,0 +1,67 @@
+# 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_prim_func([A, B, C]), target=tgt)
+
+    dev = tvm.device(tgt.kind.name, 0)
+
+    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()

tests/python/runtime/test_runtime_error.py

@@ -0,0 +1,32 @@
+# 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()

tests/python/runtime/test_runtime_extension.py

@@ -0,0 +1,42 @@
+# 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 tirx as T
+
+
+def test_dltensor_compatible():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def arange(A: T.handle):
+            n = T.int32()
+            Ab = T.match_buffer(A, (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()

tests/python/runtime/test_runtime_measure.py

@@ -0,0 +1,71 @@
+# 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.build(te.create_prim_func([X]))
+
+    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()

tests/python/runtime/test_runtime_module_export.py

@@ -0,0 +1,73 @@
+# 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
+
+
+@pytest.mark.skipif(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_prim_func([A, B]).with_attr("global_symbol", "myadd0")
+    )
+    irmod1 = tvm.IRModule.from_expr(
+        te.create_prim_func([A, B]).with_attr("global_symbol", "myadd1")
+    )
+
+    mod0 = tvm.tirx.build(irmod0, target="llvm")
+    mod1 = tvm.tirx.build(irmod1, target="llvm")
+
+    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()

tests/python/runtime/test_runtime_module_load.py

@@ -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()

tests/python/runtime/test_runtime_module_property.py

@@ -0,0 +1,60 @@
+# 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.build(te.create_prim_func([A, B]), target="llvm")
+
+
+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()

tests/python/runtime/test_runtime_nd_array.py

@@ -0,0 +1,254 @@
+# 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()

tests/python/runtime/test_runtime_rpc.py

@@ -0,0 +1,704 @@
+# 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: E712, F841
+
+import gc
+import multiprocessing
+import os
+import stat
+import sys
+import tempfile
+import time
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+pytest.importorskip("tornado")  # tvm.rpc.proxy and tvm.rpc.tracker require tornado
+
+import tvm
+import tvm.testing
+from tvm import rpc, te
+from tvm.rpc.proxy import Proxy
+from tvm.rpc.tracker import Tracker
+from tvm.script import ir as I
+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.
+    tvm.testing.main()
+
+
+# tkonolige: The issue as I understand it is this: multiprocessing's spawn
+# method launches a new process and then imports the relevant modules. This
+# means that all registered functions must exist at the top level scope. In
+# this file they are, so all is well when we run this file directly.
+# However, when run under pytest, the functions aren't registered on the
+# server. I believe this is because pytest is also using multiprocessing to
+# run individual functions. Somewhere along the way, the imports are being
+# lost, so the server ends up not registering the functions.
+pytestmark = pytest.mark.skipif(
+    # Windows does not support fork so we can enable Windows for testing
+    sys.platform.startswith("win") == False and multiprocessing.get_start_method() != "fork",
+    reason=(
+        "pytest + multiprocessing spawn method causes tvm.register_global_func to "
+        "not work on the rpc.Server."
+    ),
+)
+
+# 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(not env.build_flag_enabled("USE_RPC"), 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_prim_func([A, B]), target=target)
+
+        dev = remote.cpu(0)
+        a = tvm.runtime.tensor(np.random.randint(0, 256, size=shape).astype(A.dtype), device=dev)
+        b = tvm.runtime.tensor(np.zeros(shape).astype(A.dtype), device=dev)
+        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)
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_simple():
+    server = rpc.Server(key="x1")
+    client = rpc.connect("127.0.0.1", server.port, key="x1")
+
+    def check_remote():
+        f1 = client.get_function("rpc.test.addone")
+        assert f1(10) == 11
+        f3 = client.get_function("rpc.test.except")
+
+        with pytest.raises(RuntimeError):
+            f3("abc")
+
+        f2 = client.get_function("rpc.test.strcat")
+        assert f2("abc", 11) == "abc:11"
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_runtime_string():
+    server = rpc.Server(key="x1")
+    client = rpc.connect("127.0.0.1", server.port, key="x1")
+
+    def check_remote():
+        func = client.get_function("rpc.test.runtime_str_concat")
+        x = tvm_ffi.core.String("abc")
+        y = tvm_ffi.core.String("def")
+        assert str(func(x, y)) == "abcdef"
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_array():
+    server = rpc.Server()
+    remote = rpc.connect("127.0.0.1", server.port)
+
+    def check_remote():
+        x = np.ones((3, 4))
+        r_cpu = tvm.runtime.tensor(x, remote.cpu(0))
+        assert str(r_cpu.device).startswith("remote")
+        np.testing.assert_equal(r_cpu.numpy(), x)
+        fremote = remote.get_function("rpc.test.remote_tensor_func")
+        fremote(r_cpu)
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_large_array():
+    # testcase of large array creation
+    server = rpc.Server()
+    remote = rpc.connect("127.0.0.1", server.port)
+
+    def check_remote():
+        dev = remote.cpu(0)
+        a_np = np.ones((5041, 720)).astype("float32")
+        b_np = np.ones((720, 192)).astype("float32")
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        np.testing.assert_equal(a.numpy(), a_np)
+        np.testing.assert_equal(b.numpy(), b_np)
+
+    check_remote()
+
+
+@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_echo():
+    def check(remote, local_session):
+        fecho = remote.get_function("testing.echo")
+        assert fecho(1, 2, 3) == 1
+        assert fecho(100, 2, 3) == 100
+        assert fecho("xyz") == "xyz"
+        assert bytes(fecho(bytearray(b"123"))) == b"123"
+
+        with pytest.raises(RuntimeError):
+            raise_err = remote.get_function("testing.test_raise_error")
+            raise_err("RuntimeError", "msg")
+
+        remote.cpu().sync()
+        # tests around system lib are not threadsafe by design
+        # and do not work well with multithread pytest
+        # skip local session as they are being tested elsewhere
+        if not local_session:
+            with pytest.raises(AttributeError):
+                f3 = remote.system_lib()["notexist"]
+
+    temp = rpc.server._server_env([])
+    server = rpc.Server()
+    client = rpc.connect("127.0.0.1", server.port)
+    check(rpc.LocalSession(), True)
+
+    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(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_file_exchange():
+    server = rpc.Server()
+    remote = rpc.connect("127.0.0.1", server.port)
+
+    def check_remote():
+        blob = bytearray(np.random.randint(0, 10, size=(10)))
+        remote.upload(blob, "dat.bin")
+        rev = remote.download("dat.bin")
+        assert rev == blob
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+@pytest.mark.skipif(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_prim_func([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 = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
+        b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
+        time_f = f1.time_evaluator(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(A.dtype), tvm.cpu(0))
+        b = tvm.runtime.tensor(np.zeros(102, dtype=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)
+        dev = remote.cpu(0)
+        f1 = remote.system_lib()
+
+        a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
+        b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
+        time_f = f1.time_evaluator("myadd", remote.cpu(0), number=1)
+        cost = time_f(a, b).mean
+        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 = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype), dev)
+        b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
+        fhost(a, b)
+        np.testing.assert_equal(b.numpy(), a.numpy() + 1)
+
+    check_remote(rpc.LocalSession())
+    check_remote(client)
+    check_minrpc()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_return_func():
+    server = rpc.Server(key="x1")
+    client = rpc.connect("127.0.0.1", server.port, key="x1")
+
+    def check_remote():
+        f1 = client.get_function("rpc.test.add_to_lhs")
+        fadd = f1(10)
+        assert fadd(12) == 22
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_session_constructor_args():
+    # start server
+    server0 = rpc.Server(key="x0")
+    server1 = rpc.Server(key="x1")
+
+    def check_multi_hop():
+        # use server0 as proxy to connect to server1
+        client = rpc.connect(
+            "127.0.0.1",
+            server0.port,
+            key="x0",
+            session_constructor_args=["rpc.Connect", "127.0.0.1", server1.port, "x1", False],
+        )
+
+        fecho = client.get_function("testing.echo")
+        assert fecho(1, 2, 3) == 1
+        assert fecho(100, 2, 3) == 100
+        assert fecho("xyz") == "xyz"
+        assert bytes(fecho(bytearray(b"123"))) == b"123"
+
+        nd = tvm.runtime.tensor([1, 2, 3], device=client.cpu(0))
+        assert nd.numpy()[1] == 2
+
+    def check_error_handling():
+        with pytest.raises(tvm.error.RPCError):
+            client = rpc.connect(
+                "127.0.0.1",
+                server0.port,
+                key="x0",
+                session_constructor_args=["rpc.NonExistingConstructor"],
+            )
+
+    check_multi_hop()
+    check_error_handling()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_return_tensor():
+    def run_arr_test():
+        server = rpc.Server(key="x1")
+        client = rpc.connect("127.0.0.1", server.port, key="x1")
+        m = client.get_function("rpc.test.remote_return_nd")
+        get_arr = m("get_arr")
+        get_elem = m("get_elem")
+        get_arr_elem = m("get_arr_elem")
+
+        arr = get_arr()
+        assert get_elem(0) == 0.0
+        assert get_arr_elem(arr, 0) == 0.0
+
+        del arr
+        gc.collect()
+        assert get_elem(0) == 0.0
+
+    run_arr_test()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_rpc_return_remote_object():
+    def check(client, is_local):
+        make_shape = client.get_function("ffi.Shape")
+        get_elem = client.get_function("rpc.testing.GetShapeElem")
+        get_size = client.get_function("rpc.testing.GetShapeSize")
+        shape = make_shape(2, 3)
+        assert get_elem(shape, 0) == 2
+        assert get_elem(shape, 1) == 3
+        assert get_size(shape) == 2
+        # Test free object by assigning to the same variable
+        shape = make_shape(0)
+        assert get_size(shape) == 1
+        assert get_elem(shape, 0) == 0
+
+    # start server
+
+    check(rpc.LocalSession(), True)
+
+    def check_remote():
+        server = rpc.Server(key="x1")
+        client = rpc.connect("127.0.0.1", server.port, key="x1")
+        check(client, False)
+
+    check_remote()
+
+    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)
+
+    check_minrpc()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+def test_local_func():
+    client = rpc.LocalSession()
+
+    def check_remote():
+        f1 = client.get_function("rpc.test.add_to_lhs")
+        fadd = f1(10)
+        assert fadd(12) == 22
+
+        blob = bytearray(np.random.randint(0, 10, size=(10)))
+        client.upload(blob, "dat.bin")
+        rev = client.download("dat.bin")
+        assert rev == blob
+
+    check_remote()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
+def test_rpc_tracker_register(device_key):
+    # test registration
+    tracker = Tracker(port=9000, port_end=10000)
+    server1 = rpc.Server(
+        host="127.0.0.1",
+        port=9000,
+        port_end=10000,
+        key=device_key,
+        tracker_addr=("127.0.0.1", tracker.port),
+    )
+    server2 = rpc.Server(
+        host="127.0.0.1",
+        port=9000,
+        port_end=10000,
+        key=device_key,
+        tracker_addr=("127.0.0.1", tracker.port),
+        custom_addr="test_addr",  # this is a test address, which is unable to connect
+    )
+    time.sleep(1)
+    client = rpc.connect_tracker("127.0.0.1", tracker.port)
+
+    def exist_address(summary, key, host, port):
+        server_info = summary["server_info"]
+        for device in server_info:
+            if device["key"] == f"server:{key}":
+                addr = device["addr"]
+                if (host is None or host == addr[0]) and port == addr[1]:
+                    return True
+        return False
+
+    summary = client.summary()
+    assert summary["queue_info"][device_key]["free"] == 2
+    assert exist_address(summary, device_key, "127.0.0.1", server1.port)
+    assert exist_address(summary, device_key, "test_addr", server2.port)
+
+    remote = client.request(device_key)
+    summary = client.summary()
+    assert summary["queue_info"][device_key]["free"] == 1
+
+    del remote
+    time.sleep(1)
+
+    summary = client.summary()
+    assert summary["queue_info"][device_key]["free"] == 2
+
+    server1.terminate()
+    time.sleep(1)
+
+    summary = client.summary()
+    assert summary["queue_info"][device_key]["free"] == 1
+    assert not exist_address(summary, device_key, "127.0.0.1", server1.port)
+    assert exist_address(summary, device_key, "test_addr", server2.port)
+
+    server2.terminate()
+    time.sleep(1)
+
+    summary = client.summary()
+    assert summary["queue_info"][device_key]["free"] == 0
+    assert not exist_address(summary, device_key, "test_addr", server2.port)
+
+    tracker.terminate()
+
+
+def _target(host, port, device_key, timeout):
+    client = rpc.connect_tracker(host, port)
+    remote = client.request(device_key, session_timeout=timeout)
+    while True:
+        pass
+    remote.cpu()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+@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,
+        port_end=10000,
+        key=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))
+    proc2 = multiprocessing.Process(
+        target=_target, args=("127.0.0.1", tracker.port, device_key, 200)
+    )
+    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()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
+def test_rpc_tracker_via_proxy(device_key):
+    """
+         tracker
+         /     \
+    Host   --   Proxy -- RPC server
+    """
+
+    tracker_server = Tracker(port=9000, port_end=9100)
+    proxy_server = Proxy(
+        host=tracker_server.host,
+        port=8888,
+        port_end=8988,
+        tracker_addr=(tracker_server.host, tracker_server.port),
+    )
+
+    server1 = rpc.Server(
+        host=proxy_server.host,
+        port=proxy_server.port,
+        key=device_key,
+        tracker_addr=(tracker_server.host, tracker_server.port),
+        is_proxy=True,
+    )
+    server2 = rpc.Server(
+        host=proxy_server.host,
+        port=proxy_server.port,
+        key=device_key,
+        tracker_addr=(tracker_server.host, tracker_server.port),
+        is_proxy=True,
+    )
+
+    client = rpc.connect_tracker(tracker_server.host, tracker_server.port)
+    remote1 = client.request(device_key, session_timeout=30)  # pylint: disable=unused-variable
+    remote2 = client.request(device_key, session_timeout=30)  # pylint: disable=unused-variable
+
+    server2.terminate()
+    server1.terminate()
+    proxy_server.terminate()
+    tracker_server.terminate()
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_RPC"), reason="need rpc")
+@pytest.mark.parametrize("with_proxy", (True, False))
+def test_rpc_session_timeout_error(with_proxy):
+    port = 9000
+    port_end = 10000
+
+    tracker = Tracker(port=port, port_end=port_end)
+    time.sleep(0.5)
+    tracker_addr = (tracker.host, tracker.port)
+
+    if with_proxy:
+        proxy = Proxy(host="0.0.0.0", port=port, port_end=port_end, tracker_addr=tracker_addr)
+        time.sleep(0.5)
+        server = rpc.Server(host=proxy.host, port=proxy.port, is_proxy=True, key="x1")
+    else:
+        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):
+        f1 = rpc_sess.get_function("rpc.test.addone")
+        time.sleep(2)
+        f1(10)
+
+    server.terminate()
+    if with_proxy:
+        proxy.terminate()
+    tracker.terminate()
+
+
+@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:
+        @T.prim_func(s_tir=True)
+        def func_without_arg() -> T.int64:
+            return T.int64(42)
+
+        @T.prim_func(s_tir=True)
+        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

tests/python/runtime/test_runtime_trace.py

@@ -0,0 +1,220 @@
+# 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 tvm
+from tvm import te
+
+
+def test_trace_default_action():
+    n = 2
+    x = te.placeholder((n, n, n), name="X", dtype="float32")
+    y = te.compute(x.shape, lambda i, j, k: tvm.tirx.trace([i, j, k, x[i][j][k]]))
+    f = tvm.compile(te.create_prim_func([x, y]), target="llvm")
+    xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype))
+    ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype))
+    f(xnd, ynd)
+
+
+def test_trace_expr_assign():
+    @tvm.register_global_func("tvm.tirx.trace_callback2")
+    def trace_buffer(x):
+        return
+
+    def check_assign(dtype):
+        n = 4
+        x = te.placeholder((n, n, n), name="X", dtype=dtype)
+        y = te.compute(
+            x.shape, lambda i, j, k: tvm.tirx.trace([x[i][j][k]], "tvm.tirx.trace_callback2")
+        )
+        z = te.compute(
+            x.shape, lambda i, j, k: tvm.tirx.trace([y[i][j][k]], "tvm.tirx.trace_callback2")
+        )
+        f = tvm.compile(te.create_prim_func([x, y, z]), "llvm")
+
+        xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype))
+        ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype))
+        znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=z.dtype))
+        f(xnd, ynd, znd)
+
+        assert np.array_equal(xnd.numpy(), np.ones((n, n, n)))
+        assert np.array_equal(ynd.numpy(), np.ones((n, n, n)))
+        assert np.array_equal(znd.numpy(), np.ones((n, n, n)))
+
+    for t in ["float64", "float32", "int64", "int32"]:
+        check_assign(t)
+
+
+def test_trace_expr_sum_generated():
+    @tvm.register_global_func("tvm.tirx.trace_callback3")
+    def trace_buffer(x):
+        return
+
+    def check_expr_sum(dtype):
+        n = 4
+        a = te.placeholder((n, n, n), name="a", dtype=dtype)
+        b = te.placeholder((n, n, n), name="b", dtype=dtype)
+        c = te.compute(
+            a.shape,
+            lambda i, j, k: (
+                tvm.tirx.trace([a[i][j][k]], "tvm.tirx.trace_callback3")
+                + tvm.tirx.trace([b[i][j][k]], "tvm.tirx.trace_callback3")
+            ),
+        )
+        f = tvm.compile(te.create_prim_func([a, b, c]))
+        xnd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=a.dtype)))
+        ynd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=b.dtype)))
+        znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype))
+        f(xnd, ynd, znd)
+        assert np.array_equal(znd.numpy(), xnd.numpy() + ynd.numpy())
+
+    for t in ["float64", "float32", "int64", "int32"]:
+        check_expr_sum(t)
+
+
+def test_trace_expr_sum_args():
+    @tvm.register_global_func("tvm.tirx.trace_silent")
+    def silent(*args):
+        return
+
+    def check_expr_sum(dtype):
+        n = 4
+        a = te.placeholder((n, n, n), name="a", dtype=dtype)
+        b = te.placeholder((n, n, n), name="b", dtype=dtype)
+        e = te.placeholder((n, n, n), name="e", dtype=dtype)
+        d = te.placeholder((n, n, n), name="d", dtype=dtype)
+
+        c = te.compute(
+            a.shape,
+            lambda i, j, k: (
+                tvm.tirx.trace([i, j, k, a[i][j][k]], "tvm.tirx.trace_silent")
+                + tvm.tirx.trace([i, j, k, b[i][j][k]], "tvm.tirx.trace_silent")
+                + tvm.tirx.trace([i, j, k, d[i][j][k]], "tvm.tirx.trace_silent")
+                + tvm.tirx.trace([i, j, k, e[i][j][k]], "tvm.tirx.trace_silent")
+            ),
+        )
+        f = tvm.compile(te.create_prim_func([a, b, d, e, c]))
+        a_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=a.dtype)))
+        b_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=b.dtype)))
+        d_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=d.dtype)))
+        e_nd = tvm.runtime.tensor(np.array(np.ones((n, n, n), dtype=e.dtype)))
+        c_nd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=c.dtype))
+        f(a_nd, b_nd, d_nd, e_nd, c_nd)
+        assert np.array_equal(
+            c_nd.numpy(), a_nd.numpy() + b_nd.numpy() + d_nd.numpy() + e_nd.numpy()
+        )
+
+    for t in ["float64", "float32", "int64", "int32"]:
+        check_expr_sum(t)
+
+
+def test_trace_expr_sum_custom():
+    @tvm.register_global_func("tvm.tirx.trace_callback4")
+    def trace_buffer(x):
+        return
+
+    def check_expr_sum_custom(dtype):
+        n = 4
+        a = te.placeholder((n, n), name="a", dtype=dtype)
+        b = te.placeholder((n, n), name="b", dtype=dtype)
+        c = te.compute(
+            a.shape,
+            lambda i, j: (
+                tvm.tirx.trace([a[i][j]], "tvm.tirx.trace_callback4")
+                + tvm.tirx.trace([b[i][j]], "tvm.tirx.trace_callback4")
+            ),
+        )
+        f = tvm.compile(te.create_prim_func([a, b, c]))
+        npa = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype)
+        npb = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=a.dtype)
+        xnd = tvm.runtime.tensor(npa)
+        ynd = tvm.runtime.tensor(npb)
+        znd = tvm.runtime.tensor(np.zeros((n, n), dtype=c.dtype))
+        f(xnd, ynd, znd)
+        assert np.array_equal(znd.numpy(), npa + npb)
+
+    for t in ["float64", "float32", "int64", "int32"]:
+        check_expr_sum_custom(t)
+
+
+def test_trace_can_change_traced_value_int():
+    @tvm.register_global_func("tvm.tirx.trace_change_int_first")
+    def trace_buffer(x):
+        return 13
+
+    @tvm.register_global_func("tvm.tirx.trace_change_int_second")
+    def trace_buffer(x):
+        return 14
+
+    def check_assign(dtype):
+        n = 4
+        x = te.placeholder((n,), name="X", dtype=dtype)
+        y = te.compute(x.shape, lambda i: tvm.tirx.trace([x[i]], "tvm.tirx.trace_change_int_first"))
+        z = te.compute(
+            x.shape, lambda i: tvm.tirx.trace([y[i]], "tvm.tirx.trace_change_int_second")
+        )
+        f = tvm.compile(te.create_prim_func([x, y, z]))
+
+        xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype))
+        ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype))
+        znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype))
+        f(xnd, ynd, znd)
+        check_array_first = np.array([13, 13, 13, 13])
+        check_array_second = np.array([14, 14, 14, 14])
+        assert np.array_equal(ynd.numpy(), check_array_first)
+        assert np.array_equal(znd.numpy(), check_array_second)
+
+    for t in ["int64", "int32"]:
+        check_assign(t)
+
+
+def test_trace_can_change_traced_value_float():
+    @tvm.register_global_func("tvm.tirx.trace_change_float_first")
+    def trace_buffer(x):
+        return 13.0
+
+    @tvm.register_global_func("tvm.tirx.trace_change_float_second")
+    def trace_buffer(x):
+        return 14.0
+
+    def check_assign(dtype):
+        n = 4
+        x = te.placeholder((n,), name="X", dtype=dtype)
+        y = te.compute(
+            x.shape, lambda i: tvm.tirx.trace([x[i]], "tvm.tirx.trace_change_float_first")
+        )
+        z = te.compute(
+            x.shape, lambda i: tvm.tirx.trace([y[i]], "tvm.tirx.trace_change_float_second")
+        )
+        f = tvm.compile(te.create_prim_func([x, y, z]), target="llvm")
+
+        xnd = tvm.runtime.tensor(np.ones((n,), dtype=x.dtype))
+        ynd = tvm.runtime.tensor(np.zeros((n,), dtype=y.dtype))
+        znd = tvm.runtime.tensor(np.zeros((n,), dtype=z.dtype))
+        f(xnd, ynd, znd)
+        check_array_first = np.array([13.0, 13.0, 13.0, 13.0])
+        check_array_second = np.array([14.0, 14.0, 14.0, 14.0])
+        assert np.array_equal(ynd.numpy(), check_array_first)
+        assert np.array_equal(znd.numpy(), check_array_second)
+
+    for t in ["float64", "float32"]:
+        check_assign(t)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()