chore: import upstream snapshot with attribution

tests/python/contrib/pickle_memoize_script.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+# 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 sys
+
+import tvm
+
+
+@tvm.contrib.pickle_memoize.memoize("test_memoize_save_data", save_at_exit=True)
+def get_data_saved():
+    return 42
+
+
+@tvm.contrib.pickle_memoize.memoize("test_memoize_transient_data", save_at_exit=False)
+def get_data_transient():
+    return 42
+
+
+def main():
+    assert len(sys.argv) == 3, "Expect arguments SCRIPT NUM_SAVED NUM_TRANSIENT"
+
+    num_iter_saved = int(sys.argv[1])
+    num_iter_transient = int(sys.argv[2])
+
+    for _ in range(num_iter_saved):
+        get_data_saved()
+    for _ in range(num_iter_transient):
+        get_data_transient()
+
+
+if __name__ == "__main__":
+    main()

tests/python/contrib/test_android/__init__.py

@@ -0,0 +1,19 @@
+# isort: skip_file
+# 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.
+
+"""Testing infrastructure for Android"""

tests/python/contrib/test_android/infrastructure.py

@@ -0,0 +1,58 @@
+# 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.
+# pylint: disable=invalid-name
+
+"""Android testing infrastructure"""
+
+import os
+
+import tvm
+from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig, RPCRunner
+
+
+def get_rpc_runner() -> tvm.s_tir.meta_schedule.runner.RPCRunner:
+    if (
+        "TVM_TRACKER_HOST" in os.environ
+        and "TVM_TRACKER_PORT" in os.environ
+        and "RPC_DEVICE_KEY" in os.environ
+    ):
+        rpc_host = os.environ["TVM_TRACKER_HOST"]
+        rpc_port = int(os.environ["TVM_TRACKER_PORT"])
+        rpc_key = os.environ["RPC_DEVICE_KEY"]
+    else:
+        raise Exception("Please initialize environment variables for using RPC tracker")
+
+    rpc_config = RPCConfig(
+        tracker_host=rpc_host,
+        tracker_port=rpc_port,
+        tracker_key=rpc_key,
+        session_priority=1,
+        session_timeout_sec=100,
+    )
+    evaluator_config = EvaluatorConfig(
+        number=1,
+        repeat=1,
+        min_repeat_ms=0,
+    )
+    return RPCRunner(rpc_config, evaluator_config)
+
+
+def get_android_gpu_target() -> tvm.target.Target:
+    """Creates a Android GPU target"""
+    target_c = "opencl"
+    target_h = {"kind": "llvm", "mtriple": "arm64-linux-android"}
+    return tvm.target.Target(target_c, host=target_h)

tests/python/contrib/test_android/test_meta_schedule.py

@@ -0,0 +1,78 @@
+# 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
+
+"""Test rpc based launcher for Android"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")  # tvm.topi.testing imports scipy
+
+import tvm.testing
+import tvm.topi.testing
+from tvm.s_tir import meta_schedule as ms
+from tvm.s_tir.meta_schedule.builder import LocalBuilder
+from tvm.script import tirx as T
+
+from .infrastructure import get_android_gpu_target, get_rpc_runner
+
+
+@T.prim_func(s_tir=True)
+def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, [128, 128])
+    B = T.match_buffer(b, [128, 128])
+    C = T.match_buffer(c, [128, 128])
+    for i, j, k in T.grid(128, 128, 128):
+        with T.sblock("update"):
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vi, vj] = 0.0
+            C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
+
+
+@pytest.mark.skip("Integration test")
+def test_tune_tir_on_android():
+    """Test tune_tir on Android through RPC."""
+    max_workers = 4
+    builder = LocalBuilder(
+        f_export="s_tir.meta_schedule.builder.export_ndk", max_workers=max_workers
+    )
+    runner = get_rpc_runner()
+    target = get_android_gpu_target()
+    with tempfile.TemporaryDirectory() as work_dir:
+        database = ms.tir_integration.tune_tir(
+            mod=matmul,
+            target=target,
+            work_dir=work_dir,
+            max_trials_global=32,
+            num_trials_per_iter=16,
+            builder=builder,
+            runner=runner,
+        )
+        sch = ms.tir_integration.compile_tir(database, matmul, target)
+        if sch is None:
+            print("No valid schedule found!")
+        else:
+            sch.mod.show()
+            sch.trace.show()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_cblas.py

@@ -0,0 +1,301 @@
+# 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.
+"""Configure pytest"""
+
+import numpy as np
+import pytest
+
+pytest.importorskip("scipy")  # tvm.topi.testing imports scipy
+
+import tvm
+import tvm.testing
+import tvm.topi.testing
+from tvm import te
+from tvm.contrib import cblas, dnnl, mkl
+
+
+def verify_matmul_add(
+    matrix_m, matrix_l, matrix_n, lib, transa=False, transb=False, dtype="float32"
+):
+    """Tests matmul+add op"""
+    bias = te.var("bias", dtype=dtype)
+    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
+    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
+    matmul_result = lib.matmul(input1_data, input2_data, transa, transb)
+    final_result = te.compute(
+        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
+    )
+
+    def get_numpy(a, b, matrix_bias, transa, transb):
+        if transa:
+            a = a.transpose()
+        if transb:
+            b = b.transpose()
+        return np.dot(a, b) + matrix_bias
+
+    def compiling(f, name="test_matmul_add", ext=".so"):
+        path = name + ext
+        f.export_library(path)
+        mod = tvm.runtime.load_module(path)
+        f = mod[name]
+        return f
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        name = "test_matmul_add"
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result, bias]).with_attr(
+                "global_symbol", name
+            ),
+            target=target,
+        )
+        if target == "c":
+            f = compiling(f, name)
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        matrix_bias = 10.0
+        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), matrix_bias, transa, transb),
+            rtol=1e-5,
+        )
+
+    verify("llvm")
+    verify("c")
+
+
+def test_matmul_add():
+    """Tests of matmul+add op"""
+    verify_matmul_add(235, 128, 1024, cblas)
+    verify_matmul_add(235, 128, 1024, cblas, True, False)
+    verify_matmul_add(235, 128, 1024, cblas, False, True)
+    verify_matmul_add(235, 128, 1024, cblas, True, True)
+    verify_matmul_add(235, 128, 1024, mkl)
+    verify_matmul_add(235, 128, 1024, mkl, True, False)
+    verify_matmul_add(235, 128, 1024, mkl, False, True)
+    verify_matmul_add(235, 128, 1024, mkl, True, True)
+    verify_matmul_add(235, 128, 1024, dnnl)
+    verify_matmul_add(235, 128, 1024, dnnl, True, False)
+    verify_matmul_add(235, 128, 1024, dnnl, False, True)
+    verify_matmul_add(235, 128, 1024, dnnl, True, True)
+    verify_matmul_add(1, 16, 4, cblas)
+    verify_matmul_add(1, 16, 3, cblas, True, False)
+    verify_matmul_add(1, 16, 3, cblas, False, False)
+    verify_matmul_add(1, 16, 3, cblas, True, True)
+    verify_matmul_add(1, 16, 4, mkl)
+    verify_matmul_add(1, 16, 3, mkl, True, False)
+    verify_matmul_add(1, 16, 3, mkl, False, False)
+    verify_matmul_add(1, 16, 3, mkl, True, True)
+    verify_matmul_add(1, 16, 4, dnnl)
+    verify_matmul_add(1, 16, 3, dnnl, True, False)
+    verify_matmul_add(1, 16, 3, dnnl, False, False)
+    verify_matmul_add(1, 16, 3, dnnl, True, True)
+
+
+def verify_quantized_matmul_add(matrix_m, matrix_l, matrix_n, transa=False, transb=False):
+    """Tests quantized matmul+add op"""
+    if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
+        pytest.skip("Quantized dense is supported only for MKL. TVM GPU CI uses openblas")
+    data_dtype = "uint8"
+    kernel_dtype = "int8"
+    out_dtype = "int32"
+    bias = te.var("bias", dtype=out_dtype)
+    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
+    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=data_dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=kernel_dtype)
+    matmul_result = mkl.matmul_u8s8s32(input1_data, input2_data, transa, transb, dtype=out_dtype)
+    final_result = te.compute(
+        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
+    )
+
+    def get_numpy(a, b, matrix_bias, transa, transb):
+        if transa:
+            a = a.transpose()
+        if transb:
+            b = b.transpose()
+        return np.dot(a, b) + matrix_bias
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result, bias]), target=target
+        )
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.randint(low=0, high=50, size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.randint(low=0, high=50, size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        matrix_bias = 10
+        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(
+                matrix_input1.numpy().astype("int32"),
+                matrix_input2.numpy().astype("int32"),
+                matrix_bias,
+                transa,
+                transb,
+            ),
+            rtol=1e-5,
+        )
+
+    verify()
+
+
+def test_quantized_matmul_add():
+    """Tests of quantized matmul+add op"""
+    verify_quantized_matmul_add(235, 128, 1024)
+    verify_quantized_matmul_add(235, 128, 1024, True, False)
+    verify_quantized_matmul_add(235, 128, 1024, False, True)
+    verify_quantized_matmul_add(235, 128, 1024, True, True)
+    verify_quantized_matmul_add(1, 16, 4)
+    verify_quantized_matmul_add(1, 16, 3, True, False)
+    verify_quantized_matmul_add(1, 16, 3, False, True)
+    verify_quantized_matmul_add(1, 16, 3, True, True)
+
+
+def verify_batch_matmul(
+    batch_a,
+    batch_b,
+    matrix_m,
+    matrix_l,
+    matrix_n,
+    lib,
+    transa=False,
+    transb=False,
+    dtype="float32",
+):
+    """Tests matmul op where matrices are in batch"""
+    batch = max(batch_a, batch_b)
+    ashape = (batch_a, matrix_l, matrix_n) if transa else (batch_a, matrix_n, matrix_l)
+    bshape = (batch_b, matrix_m, matrix_l) if transb else (batch_b, matrix_l, matrix_m)
+    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
+    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
+    matmul_result = lib.batch_matmul(input1_data, input2_data, transa, transb)
+    final_result = te.compute(
+        matmul_result.shape, lambda k, i, j: matmul_result[k, i, j], name="final_result"
+    )
+
+    def get_numpy(a, b, transa, transb):
+        if transa:
+            a = a.transpose(0, 2, 1)
+        if not transb:
+            b = b.transpose(0, 2, 1)
+        return tvm.topi.testing.batch_matmul(a, b)
+
+    def compiling(f, name="test_batch_matmul", ext=".so"):
+        path = name + ext
+        f.export_library(path)
+        mod = tvm.runtime.load_module(path)
+        f = mod[name]
+        return f
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        name = "test_batch_matmul"
+        f = tvm.compile(
+            te.create_prim_func([input1_data, input2_data, final_result]), target=target
+        )
+        if target == "c":
+            f = compiling(f, name)
+        matrix_input1 = tvm.runtime.tensor(
+            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
+        )
+        matrix_input2 = tvm.runtime.tensor(
+            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
+        )
+        matrix_result = tvm.runtime.tensor(
+            np.zeros((batch, matrix_n, matrix_m), dtype=final_result.dtype), dev
+        )
+        f(matrix_input1, matrix_input2, matrix_result)
+        tvm.testing.assert_allclose(
+            matrix_result.numpy(),
+            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), transa, transb),
+            rtol=1e-5,
+        )
+
+    verify("llvm")
+    verify("c")
+
+
+def test_batch_matmul():
+    """Tests of matmul op where matrices are in batch"""
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, False)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, False, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, False)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, False, True)
+    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, True)
+    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
+    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
+    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
+    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, False, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, True)
+    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, False, False)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, True)
+    verify_batch_matmul(1, 1, 1, 16, 3, mkl)
+
+
+if __name__ == "__main__":
+    test_matmul_add()
+    test_quantized_matmul_add()
+    test_batch_matmul()

tests/python/contrib/test_coreml_runtime.py

@@ -0,0 +1,108 @@
+# 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 os
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm import rpc, te
+from tvm.contrib import coreml_runtime
+from tvm.support import utils, xcode
+
+proxy_host = os.environ.get("TVM_IOS_RPC_PROXY_HOST", "127.0.0.1")
+proxy_port = os.environ.get("TVM_IOS_RPC_PROXY_PORT", 9090)
+destination = os.environ.get("TVM_IOS_RPC_DESTINATION", "")
+key = "iphone"
+
+
+@pytest.mark.skip("skip because coremltools is not available in CI")
+def test_coreml_runtime():
+    import coremltools
+    from coremltools.models.neural_network import NeuralNetworkBuilder
+
+    def create_coreml_model():
+        shape = (2,)
+        alpha = 2
+
+        inputs = [
+            ("input0", coremltools.models.datatypes.Array(*shape)),
+            ("input1", coremltools.models.datatypes.Array(*shape)),
+        ]
+        outputs = [
+            ("output0", coremltools.models.datatypes.Array(*shape)),
+            ("output1", coremltools.models.datatypes.Array(*shape)),
+        ]
+        builder = NeuralNetworkBuilder(inputs, outputs)
+        builder.add_elementwise(
+            name="Add", input_names=["input0", "input1"], output_name="output0", mode="ADD"
+        )
+        builder.add_elementwise(
+            name="Mul", alpha=alpha, input_names=["input0"], output_name="output1", mode="MULTIPLY"
+        )
+        return coremltools.models.MLModel(builder.spec)
+
+    def verify(coreml_model, model_path, dev):
+        coreml_model = create_coreml_model()
+
+        out_spec = coreml_model.output_description._fd_spec
+        out_names = [spec.name for spec in out_spec]
+
+        # inference via coremltools
+        inputs = {}
+        for in_spec in coreml_model.input_description._fd_spec:
+            name = in_spec.name
+            shape = in_spec.type.multiArrayType.shape
+            inputs[name] = np.random.random_sample(shape)
+
+        coreml_outputs = [coreml_model.predict(inputs)[name] for name in out_names]
+
+        # inference via tvm coreml runtime
+        runtime = coreml_runtime.create("main", model_path, dev)
+        for name in inputs:
+            runtime.set_input(name, tvm.runtime.tensor(inputs[name], dev))
+        runtime.invoke()
+        tvm_outputs = [runtime.get_output(i).numpy() for i in range(runtime.get_num_outputs())]
+
+        for c_out, t_out in zip(coreml_outputs, tvm_outputs):
+            np.testing.assert_almost_equal(c_out, t_out, 3)
+
+    def check_remote(coreml_model):
+        temp = utils.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        xcode.popen_test_rpc(
+            proxy_host, proxy_port, key, destination=destination, libs=[compiled_model]
+        )
+        compiled_model = os.path.basename(compiled_model)
+        remote = rpc.connect(proxy_host, proxy_port, key=key)
+        dev = remote.cpu(0)
+        verify(coreml_model, compiled_model, dev)
+
+    def check_local(coreml_model):
+        temp = utils.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        dev = tvm.cpu(0)
+        verify(coreml_model, compiled_model, dev)
+
+    coreml_model = create_coreml_model()
+    check_remote(coreml_model)
+    check_local(coreml_model)
+
+
+if __name__ == "__main__":
+    test_coreml_runtime()

tests/python/contrib/test_cutlass_gemm.py

@@ -0,0 +1,399 @@
+# 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 ml_dtypes
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.contrib.pickle_memoize import memoize
+from tvm.testing import env
+
+
+def get_random_tensor(shape, dtype):
+    if dtype == "int8":
+        return np.random.randint(-128, 128, shape).astype(dtype)
+    elif dtype == "uint8":
+        return np.random.randint(0, 256, shape).astype(dtype)
+    return np.random.uniform(-1, 1, shape).astype(dtype)
+
+
+def verify_group_gemm(
+    func_name, M, N, K, num_groups, x_dtype, weight_dtype, out_dtype, use_scale, rtol, atol
+):
+    group_gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if group_gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    @memoize("tvm.contrib.cutlass.test_group_gemm_sm90")
+    def get_ref_data():
+        assert M % num_groups == 0
+        M_per_group = M // num_groups
+        a_np = get_random_tensor((M, K), x_dtype)
+        b_np = get_random_tensor((num_groups, N, K), weight_dtype)
+        indptr_np = np.arange(1, num_groups + 1).astype("int64") * M_per_group
+        c_np = np.concatenate(
+            [a_np[i * M_per_group : (i + 1) * M_per_group] @ b_np[i].T for i in range(num_groups)],
+            axis=0,
+        )
+        return a_np, b_np, indptr_np, c_np
+
+    def to_numpy_dtype(dtype):
+        mapping = {"float8_e5m2": ml_dtypes.float8_e5m2, "float8_e4m3fn": ml_dtypes.float8_e4m3fn}
+        return mapping.get(dtype, dtype)
+
+    a_np, b_np, indptr_np, c_np = get_ref_data()
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        a_nd = tvm.runtime.tensor(a_np.astype(to_numpy_dtype(x_dtype)), device=dev)
+        b_nd = tvm.runtime.tensor(b_np.astype(to_numpy_dtype(weight_dtype)), device=dev)
+        c_nd = tvm.runtime.empty(c_np.shape, dtype=out_dtype, device=dev)
+        indptr_nd = tvm.runtime.tensor(indptr_np, device=dev)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=dev)
+        if use_scale:
+            scale = tvm.runtime.tensor(np.array([1.0], dtype="float32"), device=dev)
+            group_gemm_func(a_nd, b_nd, indptr_nd, workspace, scale, c_nd)
+        else:
+            group_gemm_func(a_nd, b_nd, indptr_nd, workspace, c_nd)
+        tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=rtol, atol=atol)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_group_gemm_sm90():
+    verify_group_gemm(
+        "cutlass.group_gemm",
+        8,
+        128,
+        128,
+        4,
+        "float16",
+        "float16",
+        "float16",
+        False,
+        rtol=1e-3,
+        atol=1e-3,
+    )
+    verify_group_gemm(
+        "cutlass.group_gemm_e5m2_e5m2_fp16",
+        8,
+        16,
+        16,
+        4,
+        "float8_e5m2",
+        "float8_e5m2",
+        "float16",
+        True,
+        rtol=1e-1,
+        atol=1,
+    )
+    verify_group_gemm(
+        "cutlass.group_gemm_e4m3_e4m3_fp16",
+        8,
+        16,
+        16,
+        4,
+        "float8_e4m3fn",
+        "float8_e4m3fn",
+        "float16",
+        True,
+        rtol=1e-1,
+        atol=1,
+    )
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_group_gemm_sm100():
+    verify_group_gemm(
+        "cutlass.group_gemm",
+        8,
+        128,
+        128,
+        4,
+        "bfloat16",
+        "bfloat16",
+        "bfloat16",
+        False,
+        rtol=1e-2,
+        atol=1e-3,
+    )
+
+
+def rowwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
+    x_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
+    x_scale_shape = (
+        *shape[:-1],
+        (shape[-1] + block_size[1] - 1) // block_size[1],
+    )
+    # For each (block_size[1]) block, compute the max abs value of `w_full_np`
+    x_max_abs_np = np.zeros(x_scale_shape, dtype="float32")
+    for i in range(x_scale_shape[-1]):
+        x_max_abs_np[..., i] = np.max(
+            np.abs(x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]),
+            axis=-1,
+        )[0]
+    # Scale is the `x_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
+    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
+    x_scale_np = x_max_abs_np / fp8_max
+    # `x_np` is the `x_full_np` divided by the `x_scale_np` (with block awareness),
+    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
+    x_np = np.zeros_like(x_full_np, dtype="float8_e4m3fn")
+    for i in range(x_scale_shape[-1]):
+        x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = np.clip(
+            x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])]
+            / x_scale_np[..., i : i + 1],
+            -fp8_max,
+            fp8_max,
+        )
+
+    x_scale_np = np.random.rand(*x_scale_np.shape).astype("float32") / fp8_max
+    for i in range(x_scale_shape[-1]):
+        x_full_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])] = (
+            x_np[..., i * block_size[1] : min((i + 1) * block_size[1], shape[-1])].astype(
+                x_scale_np.dtype
+            )
+            * x_scale_np[..., i : i + 1]
+        )
+    return x_np, x_scale_np
+
+
+def blockwise_quant_fp8_e4m3(shape: tuple[int, int], block_size: tuple[int, int], dtype: str):
+    w_full_np = (np.random.rand(*shape) * 2 - 1).astype(dtype)
+    w_scale_shape = (
+        *shape[:-2],
+        (shape[-2] + block_size[0] - 1) // block_size[0],
+        (shape[-1] + block_size[1] - 1) // block_size[1],
+    )
+    # For each (block_size[0], block_size[1]) block, compute the max abs value of `w_full_np`
+    w_max_abs_np = np.zeros(w_scale_shape, dtype="float32")
+    for i in range(w_scale_shape[-2]):
+        for j in range(w_scale_shape[-1]):
+            block_shape = (
+                *shape[:-2],
+                min(block_size[0], shape[-2] - i * block_size[0]),
+                min(block_size[1], shape[-1] - j * block_size[1]),
+            )
+            w_max_abs_np[..., i, j] = np.max(
+                np.abs(
+                    w_full_np[
+                        ...,
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ]
+                ).reshape(*shape[:-2], block_shape[-2] * block_shape[-1]),
+                axis=-1,
+            )
+    # Scale is the `w_max_abs_np` divided by the max value of quant_dtype in ml_dtypes
+    fp8_max = float(ml_dtypes.finfo("float8_e4m3fn").max)
+    w_scale_np = w_max_abs_np / fp8_max
+    # `w_np` is the `w_full_np` divided by the `w_scale_np` (with block awareness),
+    # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype`
+    w_np = np.zeros_like(w_full_np, dtype="float8_e4m3fn")
+    if len(w_scale_shape) == 2:
+        for i in range(w_scale_shape[-2]):
+            for j in range(w_scale_shape[-1]):
+                w_np[
+                    i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                    j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                ] = np.clip(
+                    w_full_np[
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ]
+                    / w_scale_np[..., i, j],
+                    -fp8_max,
+                    fp8_max,
+                )
+    else:
+        for e in range(w_scale_shape[0]):
+            for i in range(w_scale_shape[-2]):
+                for j in range(w_scale_shape[-1]):
+                    w_np[
+                        e,
+                        i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                        j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                    ] = np.clip(
+                        w_full_np[
+                            e,
+                            i * block_size[0] : min((i + 1) * block_size[0], shape[-2]),
+                            j * block_size[1] : min((j + 1) * block_size[1], shape[-1]),
+                        ]
+                        / w_scale_np[e, i, j],
+                        -fp8_max,
+                        fp8_max,
+                    )
+
+    w_scale_np = np.random.rand(*w_scale_np.shape).astype("float32") / fp8_max
+    return w_np, w_scale_np
+
+
+def blockwise_matmul(
+    x_fp8_np: np.ndarray,
+    x_scale_np: np.ndarray,
+    w_np: np.ndarray,
+    w_scale_np: np.ndarray,
+    block_size: tuple[int, int],
+    dtype: str,
+):
+    o_np = np.zeros((x_fp8_np.shape[0], w_np.shape[0]), dtype=dtype)
+    for j in range(w_scale_np.shape[0]):
+        for k in range(w_scale_np.shape[1]):
+            o_np[:, j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[0])] += (
+                np.matmul(
+                    x_fp8_np[
+                        :, k * block_size[1] : min((k + 1) * block_size[1], x_fp8_np.shape[1])
+                    ].astype(dtype),
+                    w_np[
+                        j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[0]),
+                        k * block_size[1] : min((k + 1) * block_size[1], w_np.shape[1]),
+                    ].T.astype(dtype),
+                )
+                * x_scale_np[:, k : k + 1]
+                * w_scale_np[j, k]
+            )
+    return o_np
+
+
+def blockwise_bmm(
+    x_fp8_np: np.ndarray,
+    x_scale_np: np.ndarray,
+    w_np: np.ndarray,
+    w_scale_np: np.ndarray,
+    block_size: tuple[int, int],
+    dtype: str,
+):
+    o_np = np.zeros((x_fp8_np.shape[0], x_fp8_np.shape[1], w_np.shape[1]), dtype=dtype)
+    for j in range(w_scale_np.shape[1]):
+        for k in range(w_scale_np.shape[2]):
+            o_np[..., j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1])] += (
+                np.matmul(
+                    x_fp8_np[
+                        ..., k * block_size[1] : min((k + 1) * block_size[1], x_fp8_np.shape[2])
+                    ].astype(dtype),
+                    w_np[
+                        ...,
+                        j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1]),
+                        k * block_size[1] : min((k + 1) * block_size[1], w_np.shape[2]),
+                    ]
+                    .transpose(0, 2, 1)
+                    .astype(dtype),
+                )
+                * x_scale_np[..., k : k + 1]
+                * w_scale_np[..., j : j + 1, k : k + 1]
+            )
+    return o_np
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_fp8_e4m3_groupwise_scaled_gemm():
+    M = 16
+    N = 4608
+    K = 896
+    block_size = (128, 128)
+    assert N % 128 == 0 and K % 128 == 0  # Only support N/K are multiple of 128
+
+    func_name = "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn"
+    gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    dtype = "bfloat16"
+    x_np, x_scale_np = rowwise_quant_fp8_e4m3((M, K), block_size, dtype)
+    w_np, w_scale_np = blockwise_quant_fp8_e4m3((N, K), block_size, dtype)
+    o_np = blockwise_matmul(x_np, x_scale_np, w_np, w_scale_np, block_size, dtype)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_tvm = tvm.runtime.tensor(x_np, device=device)
+        x_scale_tvm = tvm.runtime.tensor(x_scale_np.T, device=device)
+        w_tvm = tvm.runtime.tensor(w_np, device=device)
+        w_scale_tvm = tvm.runtime.tensor(w_scale_np, device=device)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=device)
+        o_tvm = tvm.runtime.empty((M, N), dtype=dtype, device=device)
+        gemm_func(
+            x_tvm,
+            w_tvm,
+            x_scale_tvm,
+            w_scale_tvm,
+            workspace,
+            block_size[0],
+            block_size[1],
+            o_tvm,
+        )
+        tvm.testing.assert_allclose(o_tvm.numpy(), o_np, rtol=1e-4, atol=0.5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.skipif(not env.build_flag_enabled("USE_CUTLASS"), reason="need cutlass")
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_fp8_e4m3_groupwise_scaled_bmm():
+    B = 16
+    M = 40
+    N = 512
+    K = 128
+    block_size = (128, 128)
+    assert N % 128 == 0 and K % 128 == 0  # Only support N/K are multiple of 128
+
+    func_name = "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn"
+    gemm_func = tvm.get_global_func(func_name, allow_missing=True)
+    if gemm_func is None:
+        print(f"Skipped as {func_name} is not available")
+        return
+
+    dtype = "bfloat16"
+    x_np, x_scale_np = rowwise_quant_fp8_e4m3((B, M, K), block_size, dtype)
+    w_np, w_scale_np = blockwise_quant_fp8_e4m3((B, N, K), block_size, dtype)
+    o_np = blockwise_bmm(x_np, x_scale_np, w_np, w_scale_np, block_size, dtype)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_tvm = tvm.runtime.tensor(x_np, device=device)
+        x_scale_tvm = tvm.runtime.tensor(x_scale_np.transpose(0, 2, 1), device=device)
+        w_tvm = tvm.runtime.tensor(w_np, device=device)
+        w_scale_tvm = tvm.runtime.tensor(w_scale_np, device=device)
+        workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=device)
+        o_tvm = tvm.runtime.empty((B, M, N), dtype=dtype, device=device)
+        gemm_func(
+            x_tvm,
+            w_tvm,
+            x_scale_tvm,
+            w_scale_tvm,
+            workspace,
+            block_size[0],
+            block_size[1],
+            o_tvm,
+        )
+        tvm.testing.assert_allclose(o_tvm.numpy(), o_np, rtol=1e-4, atol=0.5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_dlpack.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.
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.contrib.dlpack import to_pytorch_func
+
+
+def verify_torch_dlpack():
+    a = np.random.randn(1337)
+    tvm_a = tvm.runtime.tensor(a)
+    np.testing.assert_equal(tvm.runtime.from_dlpack(tvm_a).numpy(), a)
+
+    try:
+        import torch
+        import torch.utils.dlpack
+
+        x = torch.rand(56, 56)
+        tvm_x = tvm.runtime.from_dlpack(torch.utils.dlpack.to_dlpack(x))
+        np.testing.assert_equal(x.numpy(), tvm_x.numpy())
+        y = tvm.runtime.from_dlpack(tvm_x)
+        np.testing.assert_equal(y.numpy(), tvm_x.numpy())
+        np.testing.assert_equal(torch.utils.dlpack.from_dlpack(y).numpy(), tvm_x.numpy())
+
+        n = tvm.runtime.convert(137)
+        xx = torch.rand(137, 137)
+        yy = torch.rand(137, 137)
+        zz2 = torch.empty(137, 137)
+        zz = xx.mm(yy)
+        XX = te.placeholder((n, n), name="X")
+        YY = te.placeholder((n, n), name="Y")
+
+        k = te.reduce_axis((0, n), name="k")
+        ZZ = te.compute((n, n), lambda i, j: te.sum(XX[i, k] * YY[k, j], axis=k))
+        # No need to speficy target_host if it's llvm
+        # Otherwise you will need to specify the target and target_host
+        f = tvm.compile(te.create_prim_func([XX, YY, ZZ]))
+
+        f_pytorch = to_pytorch_func(f)
+        zz2 = torch.empty(137, 137)
+        f_pytorch(xx, yy, zz2)
+        tvm.testing.assert_allclose(zz.numpy(), zz2.numpy(), rtol=1e-4, atol=1e-4)
+
+    except ImportError:
+        pass
+
+
+def test_torch_dlpack():
+    # Run dlpack interoperability test a few times to make sure it's stable.
+    for i in range(5):
+        verify_torch_dlpack()
+
+
+if __name__ == "__main__":
+    test_torch_dlpack()

tests/python/contrib/test_example_npu.py

@@ -0,0 +1,277 @@
+# 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 Example NPU Backend
+
+This test file demonstrates how to test a custom NPU backend
+implementation using TVM's testing infrastructure.
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
+from tvm.script import relax as R
+
+
+@tvm.script.ir_module
+class MatmulReLU:
+    """Example module with matrix multiplication and ReLU"""
+
+    @R.function
+    def main(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 8), "float32"),
+    ) -> R.Tensor((2, 8), "float32"):
+        with R.dataflow():
+            y = relax.op.matmul(x, w)
+            z = relax.op.nn.relu(y)
+            R.output(z)
+        return z
+
+
+@tvm.script.ir_module
+class Conv2dReLU:
+    """Example module with 2D convolution and ReLU"""
+
+    @R.function
+    def main(
+        x: R.Tensor((1, 3, 32, 32), "float32"),
+        w: R.Tensor((16, 3, 3, 3), "float32"),
+    ) -> R.Tensor((1, 16, 30, 30), "float32"):
+        with R.dataflow():
+            y = relax.op.nn.conv2d(x, w)
+            z = relax.op.nn.relu(y)
+            R.output(z)
+        return z
+
+
+@tvm.script.ir_module
+class MultipleOps:
+    """Example module with multiple operations that can be fused"""
+
+    @R.function
+    def main(
+        x: R.Tensor((1, 16, 32, 32), "float32"),
+    ) -> R.Tensor((1, 16, 16, 16), "float32"):
+        with R.dataflow():
+            # First ReLU
+            y = relax.op.nn.relu(x)
+            # Max pooling
+            z = relax.op.nn.max_pool2d(y, pool_size=(2, 2), strides=(2, 2))
+            # Second ReLU
+            out = relax.op.nn.relu(z)
+            R.output(out)
+        return out
+
+
+@tvm.script.ir_module
+class Softmax:
+    """Example module with softmax"""
+
+    @R.function
+    def main(x: R.Tensor((2, 8), "float32")) -> R.Tensor((2, 8), "float32"):
+        with R.dataflow():
+            z = relax.op.nn.softmax(x)
+            R.output(z)
+        return z
+
+
+# Check if the example NPU runtime is available
+has_example_npu_codegen = tvm.get_global_func("relax.ext.example_npu", True)
+has_example_npu_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
+has_example_npu = has_example_npu_codegen and has_example_npu_runtime
+
+example_npu_enabled = pytest.mark.skipif(
+    not has_example_npu,
+    reason="Example NPU backend not enabled. Compile with the example NPU runtime.",
+)
+
+
+def test_example_npu_patterns_registered():
+    """Test that all expected patterns are registered"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    patterns = get_patterns_with_prefix("example_npu")
+    pattern_names = {p.name for p in patterns}
+
+    # Core patterns that should always be available
+    core_patterns = {
+        "example_npu.dense",
+        "example_npu.matmul",
+        "example_npu.conv1d",
+        "example_npu.conv2d",
+        "example_npu.max_pool2d",
+    }
+
+    assert core_patterns.issubset(pattern_names), (
+        f"Missing core patterns: {core_patterns - pattern_names}"
+    )
+
+    # Check that at least some activation patterns are available
+    activation_patterns = {name for name in pattern_names if "relu" in name or "sigmoid" in name}
+    assert len(activation_patterns) > 0, "No activation patterns found"
+
+
+@example_npu_enabled
+def test_example_npu_matmul_relu_partitioning():
+    """Test graph partitioning for MatMul + ReLU pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    # Verify partitioning happened
+    assert partitioned_mod is not None
+
+    # Check that composite functions were created
+    for gvar, func in partitioned_mod.functions.items():
+        if gvar.name_hint != "main":
+            # This should be a composite function
+            assert "Composite" in str(func)
+
+
+@example_npu_enabled
+def test_example_npu_conv2d_relu_partitioning():
+    """Test graph partitioning for Conv2D + ReLU pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = Conv2dReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+
+@example_npu_enabled
+def test_example_npu_multiple_ops():
+    """Test partitioning with multiple fusable operations"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MultipleOps
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition the graph
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+
+@example_npu_enabled
+def test_example_npu_softmax_partitioning():
+    """Test graph partitioning for softmax pattern"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = Softmax
+    patterns = get_patterns_with_prefix("example_npu")
+
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+    for gvar, func in partitioned_mod.functions.items():
+        if gvar.name_hint != "main":
+            assert "Composite" in str(func)
+
+
+@example_npu_enabled
+def test_example_npu_codegen():
+    """Test code generation for the example NPU backend"""
+    import tvm.relax.backend.contrib.example_npu  # noqa: F401
+
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Partition and generate code
+    partitioned_mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
+    partitioned_mod = MergeCompositeFunctions()(partitioned_mod)
+    partitioned_mod = RunCodegen()(partitioned_mod)
+
+    assert partitioned_mod is not None
+
+    # The module should now contain external function calls
+    main_func = partitioned_mod["main"]
+    assert main_func is not None
+
+
+@example_npu_enabled
+def test_example_npu_runtime_execution():
+    """Test end-to-end execution with the example NPU runtime"""
+    import tvm.relax.backend.contrib.example_npu
+
+    # Create simple test inputs
+    np.random.seed(42)
+    x_np = np.random.randn(2, 4).astype("float32")
+    w_np = np.random.randn(4, 8).astype("float32")
+
+    # Expected output (computed with NumPy)
+    expected = np.maximum(0, np.matmul(x_np, w_np))
+
+    # Build and run with example NPU backend
+    mod = MatmulReLU
+    patterns = get_patterns_with_prefix("example_npu")
+
+    # Apply transformations
+    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    mod = RunCodegen()(mod)
+
+    # Build the module
+    target = tvm.target.Target("llvm")
+    with tvm.transform.PassContext(opt_level=3):
+        built = relax.build(mod, target)
+
+    # Create VM and run
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    x_tvm = tvm.runtime.tensor(x_np, tvm.cpu())
+    w_tvm = tvm.runtime.tensor(w_np, tvm.cpu())
+
+    result = vm["main"](x_tvm, w_tvm)
+
+    # Verify the result shape is correct (the runtime is a stub and does not compute numerically)
+    assert result.numpy().shape == expected.shape
+
+
+if __name__ == "__main__":
+    # Run tests locally for debugging
+    test_example_npu_patterns_registered()
+
+    if has_example_npu:
+        print("Example NPU backend is available, running tests...")
+        test_example_npu_matmul_relu_partitioning()
+        test_example_npu_conv2d_relu_partitioning()
+        test_example_npu_softmax_partitioning()
+        test_example_npu_multiple_ops()
+        test_example_npu_codegen()
+        test_example_npu_runtime_execution()
+        print("All tests passed!")
+    else:
+        print("Example NPU backend not available. Compile with example NPU runtime to run tests.")

tests/python/contrib/test_hipblas.py

@@ -0,0 +1,121 @@
+# 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: E741
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te
+from tvm.contrib import hipblas
+from tvm.testing import env
+
+
+def verify_matmul_add(in_dtype, out_dtype, rtol=1e-5):
+    n = 1024
+    l = 128
+    m = 236
+    A = te.placeholder((n, l), name="A", dtype=in_dtype)
+    B = te.placeholder((l, m), name="B", dtype=in_dtype)
+    C = hipblas.matmul(A, B, dtype=out_dtype)
+
+    def verify(target="rocm"):
+        if not tvm.get_global_func("tvm.contrib.hipblas.matmul", True):
+            print("skip because extern function is not available")
+            return
+        f = tvm.compile(te.create_prim_func([A, B, C]), target=target)
+
+        def run_and_check():
+            dev = tvm.rocm(0)
+            a = tvm.runtime.tensor(np.random.uniform(0, 128, size=(n, l)).astype(A.dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(0, 128, size=(l, m)).astype(B.dtype), dev)
+            c = tvm.runtime.tensor(np.zeros((n, m), dtype=C.dtype), dev)
+            f(a, b, c)
+            tvm.testing.assert_allclose(
+                c.numpy(),
+                np.dot(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)),
+                rtol=rtol,
+            )
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+    verify()
+
+
+def roundoff(v, d):
+    return int(np.floor((v + d - 1) / d) * d)
+
+
+def verify_batch_matmul(Ashape, Bshape, Cshape, in_dtype, out_dtype, rtol=1e-5):
+    A = te.placeholder(Ashape, name="A", dtype=in_dtype)
+    B = te.placeholder(Bshape, name="B", dtype=in_dtype)
+    C = hipblas.batch_matmul(A, B, dtype=out_dtype)
+
+    f = tvm.compile(te.create_prim_func([A, B, C]), target="rocm")
+
+    def run_and_check():
+        dev = tvm.rocm(0)
+        if "int" in in_dtype:
+            a = tvm.runtime.tensor(np.random.uniform(1, 10, size=Ashape).astype(in_dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(1, 10, size=Bshape).astype(in_dtype), dev)
+        else:
+            a = tvm.runtime.tensor(np.random.uniform(size=Ashape).astype(A.dtype), dev)
+            b = tvm.runtime.tensor(np.random.uniform(size=Bshape).astype(B.dtype), dev)
+
+        c = tvm.runtime.tensor(np.zeros(Cshape, dtype=C.dtype), dev)
+        f(a, b, c)
+        tvm.testing.assert_allclose(
+            c.numpy(),
+            np.matmul(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)).astype(C.dtype),
+            rtol=rtol,
+        )
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_matmul_add():
+    verify_matmul_add("float", "float", rtol=1e-3)
+    verify_matmul_add("float16", "float")
+    verify_matmul_add("float16", "float16", rtol=1e-2)
+    verify_matmul_add("int8", "int32")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_rocm(), reason="need rocm")
+def test_batch_matmul():
+    if not tvm.get_global_func("tvm.contrib.hipblas.batch_matmul", True):
+        print("skip because extern function is not available")
+        return
+
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float", "float")
+    verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float", "float")
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float")
+    verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float")
+    verify_batch_matmul(
+        (16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2
+    )
+    verify_batch_matmul(
+        (16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2
+    )
+
+    verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "int8", "int32")
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_memoize.py

@@ -0,0 +1,128 @@
+# 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 tvm.contrib.pickle_memoize"""
+
+import os
+import pathlib
+import subprocess
+import sys
+import tempfile
+
+import tvm.testing
+
+TEST_SCRIPT_FILE = pathlib.Path(__file__).with_name("pickle_memoize_script.py").resolve()
+
+
+def test_cache_dir_not_in_current_working_dir():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+        subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1", "1"], cwd=temp_dir)
+
+        new_files = list(temp_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+
+def test_current_directory_is_not_required_to_be_writable():
+    """TVM may be imported without directory permissions
+
+    This is a regression test.  In previous implementations, the
+    `tvm.contrib.pickle_memoize.memoize` function would write to the
+    current directory when importing TVM.  Import of a Python module
+    should not write to any directory.
+
+    """
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        # User may read/cd into the temp dir, nobody may write to temp
+        # dir.
+        temp_dir.chmod(0o500)
+        subprocess.check_call([sys.executable, "-c", "import tvm"], cwd=temp_dir)
+
+
+def test_cache_dir_defaults_to_home_config_cache():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        temp_dir = pathlib.Path(temp_dir)
+
+        subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1", "0"], cwd=temp_dir)
+
+        new_files = list(temp_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+        cache_dir = pathlib.Path.home().joinpath(".cache", "tvm", "pkl_memoize_py3")
+        assert cache_dir.exists()
+        cache_files = list(cache_dir.iterdir())
+        assert len(cache_files) >= 1
+
+
+def test_cache_dir_respects_xdg_cache_home():
+    with (
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_working_dir,
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_cache_dir,
+    ):
+        temp_cache_dir = pathlib.Path(temp_cache_dir)
+        temp_working_dir = pathlib.Path(temp_working_dir)
+
+        subprocess.check_call(
+            [sys.executable, str(TEST_SCRIPT_FILE), "1", "0"],
+            cwd=temp_working_dir,
+            env={
+                **os.environ,
+                "XDG_CACHE_HOME": temp_cache_dir.as_posix(),
+            },
+        )
+
+        new_files = list(temp_working_dir.iterdir())
+        assert not new_files, (
+            "Use of tvm.contrib.pickle_memorize may not write to current directory."
+        )
+
+        cache_dir = temp_cache_dir.joinpath("tvm", "pkl_memoize_py3")
+        assert cache_dir.exists()
+        cache_files = list(cache_dir.iterdir())
+        assert len(cache_files) == 1
+
+
+def test_cache_dir_only_created_when_used():
+    with (
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_working_dir,
+        tempfile.TemporaryDirectory(prefix="tvm_") as temp_cache_dir,
+    ):
+        temp_cache_dir = pathlib.Path(temp_cache_dir)
+        temp_working_dir = pathlib.Path(temp_working_dir)
+
+        subprocess.check_call(
+            [sys.executable, str(TEST_SCRIPT_FILE), "0", "1"],
+            cwd=temp_working_dir,
+            env={
+                **os.environ,
+                "XDG_CACHE_HOME": temp_cache_dir.as_posix(),
+            },
+        )
+
+        cache_dir = temp_cache_dir.joinpath("tvm", "pkl_memoize_py3")
+        assert not cache_dir.exists()
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/contrib/test_random.py

@@ -0,0 +1,200 @@
+# 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: E722
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import threading
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import rpc, te
+from tvm.contrib import random
+
+
+def test_randint():
+    """Tests randint function"""
+    m = 10240
+    n = 10240
+    A = random.randint(-127, 128, size=(m, n), dtype="int32")
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.randint", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na)) < 0.3
+        assert np.min(na) == -127
+        assert np.max(na) == 127
+
+    verify()
+
+
+def test_uniform():
+    """Tests uniform function"""
+    m = 10240
+    n = 10240
+    A = random.uniform(0, 1, size=(m, n))
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.uniform", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na) - 0.5) < 1e-1
+        assert abs(np.min(na) - 0.0) < 1e-3
+        assert abs(np.max(na) - 1.0) < 1e-3
+
+    verify()
+
+
+def test_normal():
+    """Tests normal function"""
+    m = 10240
+    n = 10240
+    A = random.normal(3, 4, size=(m, n))
+
+    def verify(target="llvm"):
+        if not tvm.testing.device_enabled(target):
+            print(f"skip because {target} is not enabled...")
+            return
+        if not tvm.get_global_func("tvm.contrib.random.normal", True):
+            print("skip because extern function is not available")
+            return
+        dev = tvm.cpu(0)
+        f = tvm.compile(te.create_prim_func([A]), target=target)
+        a = tvm.runtime.tensor(np.zeros((m, n), dtype=A.dtype), dev)
+        f(a)
+        na = a.numpy()
+        assert abs(np.mean(na) - 3) < 1e-1
+        assert abs(np.std(na) - 4) < 1e-2
+
+    verify()
+
+
+@pytest.mark.gpu
+def test_random_fill():
+    """Tests random_fill function"""
+
+    def test_local(dev, dtype):
+        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
+            print("skip because extern function is not available")
+            return
+        value = tvm.runtime.empty((512, 512), dtype, dev)
+        random_fill = tvm.get_global_func("tvm.contrib.random.random_fill")
+        random_fill(value)
+
+        assert np.count_nonzero(value.numpy()) == 512 * 512
+
+        # make sure arithmentic doesn't overflow too
+        np_values = value.numpy()
+        assert np.isfinite(np_values * np_values + np_values).any()
+
+    def test_rpc(dtype):
+        if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
+            print("skip because extern function is not available")
+            return
+        if not tvm.testing.device_enabled("rpc") or not tvm.runtime.enabled("llvm"):
+            return
+
+        def check_remote(server):
+            remote = rpc.connect(server.host, server.port)
+            value = tvm.runtime.empty((512, 512), dtype, remote.cpu())
+            random_fill = remote.get_function("tvm.contrib.random.random_fill")
+            random_fill(value)
+
+            assert np.count_nonzero(value.numpy()) == 512 * 512
+
+            # make sure arithmentic doesn't overflow too
+            np_values = value.numpy()
+            assert np.isfinite(np_values * np_values + np_values).any()
+
+        check_remote(rpc.Server("127.0.0.1"))
+
+    # Packed sub-byte dtypes (e.g. int4) are intentionally unsupported by
+    # random_fill since #19714 and raise an error instead.
+    for dtype in [
+        "bool",
+        "int8",
+        "uint8",
+        "int16",
+        "uint16",
+        "int32",
+        "int32",
+        "int64",
+        "uint64",
+        "float16",
+        "float32",
+        "float64",
+    ]:
+        for target, dev in tvm.testing.enabled_targets():
+            if tvm.target.Target(target).kind.name == "llvm":
+                test_local(dev, dtype)
+            else:
+                tvm.testing.run_with_gpu_lock(test_local, dev, dtype)
+        test_rpc(dtype)
+
+
+def test_random_fill_mt():
+    """Check random filler applicability in case of nontrivial thread pool configuration.
+    Particularly when MaxConcurrency != num_workers_used_ which is actual for big-little systems.
+    """
+    no_exception_happened = True
+
+    def test_body():
+        try:
+            num_thread_used = 1
+            configure_threads = tvm.get_global_func("runtime.config_threadpool")
+            configure_threads(1, num_thread_used)
+
+            test_input = tvm.runtime.empty((10, 10))
+            random_fill = tvm.get_global_func("tvm.contrib.random.random_fill_for_measure")
+            random_fill(test_input)
+        except:  # pylint: disable=bare-except
+            nonlocal no_exception_happened
+            no_exception_happened = False
+
+    # ThreadPool object is thread local. To eliminate effect on other test cases put it into thread
+    x = threading.Thread(target=test_body)
+    x.start()
+    x.join()
+    assert no_exception_happened
+
+
+if __name__ == "__main__":
+    test_randint()
+    test_uniform()
+    test_normal()
+    test_random_fill()
+    test_random_fill_mt()

tests/python/contrib/test_rpc_proxy.py

@@ -0,0 +1,69 @@
+# 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.
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import logging
+import multiprocessing
+import time
+
+import tvm
+from tvm import rpc
+
+
+def rpc_proxy_check():
+    """This is a simple test function for RPC Proxy
+
+    It is not included as pytests, because:
+    - It depends on tornado
+    - It relies on the fact that Proxy starts before client and server connects,
+      which is often the case but not always
+
+    User can directly run this script to verify correctness.
+    """
+
+    try:
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc import proxy
+
+        web_port = 8888
+        prox = proxy.Proxy("127.0.0.1", web_port=web_port)
+
+        def check():
+            if not tvm.runtime.enabled("rpc"):
+                return
+
+            server = multiprocessing.Process(
+                target=proxy.websocket_proxy_server, args=(f"ws://localhost:{web_port}/ws", "x1")
+            )
+            # Need to make sure that the connection start after proxy comes up
+            time.sleep(0.1)
+            server.deamon = True
+            server.start()
+            client = rpc.connect(prox.host, prox.port, key="x1")
+            f1 = client.get_function("testing.echo")
+            assert f1(10) == 10
+            assert f1("xyz") == "xyz"
+
+        check()
+    except ImportError:
+        print("Skipping because tornado is not avaliable...")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    rpc_proxy_check()

tests/python/contrib/test_rpc_tracker.py

@@ -0,0 +1,152 @@
+# 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.
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import logging
+import time
+
+import tvm
+from tvm import rpc
+
+
+def check_server_drop():
+    """test when server drops"""
+    try:
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc import base, proxy, tracker
+
+        # pylint: disable=import-outside-toplevel
+        from tvm.rpc.base import TrackerCode
+
+        @tvm.register_global_func("rpc.test2.addone")
+        def addone(x):
+            return x + 1
+
+        def _put(tclient, value):
+            base.sendjson(tclient._sock, value)
+            base.recvjson(tclient._sock)
+
+        tserver = tracker.Tracker("127.0.0.1", 8888)
+        tproxy = proxy.Proxy("127.0.0.1", 8881, tracker_addr=("127.0.0.1", tserver.port))
+        tclient = rpc.connect_tracker("127.0.0.1", tserver.port)
+
+        server0 = rpc.Server(
+            "127.0.0.1", port=9099, tracker_addr=("127.0.0.1", tserver.port), key="abc"
+        )
+        server1 = rpc.Server(
+            "127.0.0.1", port=9099, tracker_addr=("127.0.0.1", tserver.port), key="xyz"
+        )
+        server2 = rpc.Server("127.0.0.1", tproxy.port, is_proxy=True, key="xyz")
+        server3 = rpc.Server("127.0.0.1", tproxy.port, is_proxy=True, key="xyz1")
+
+        # Fault tolerence to un-handled requested value
+        _put(tclient, [TrackerCode.REQUEST, "abc", "", 1])
+        _put(tclient, [TrackerCode.REQUEST, "xyz1", "", 1])
+
+        # Fault tolerence to stale worker value
+        _put(tclient, [TrackerCode.PUT, "xyz", (server1.port, "abc")])
+        _put(tclient, [TrackerCode.PUT, "xyz", (server1.port, "abcxxx")])
+        _put(tclient, [TrackerCode.PUT, "xyz", (tproxy.port, "abcxxx11")])
+
+        # Fault tolerence server timeout
+        def check_timeout(timeout, sleeptime):
+            def myfunc(remote):
+                time.sleep(sleeptime)
+                f1 = remote.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+
+            try:
+                tclient.request_and_run("xyz", myfunc, session_timeout=timeout)
+            except RuntimeError:
+                pass
+            print(tclient.text_summary())
+            try:
+                remote = tclient.request("xyz", priority=0, session_timeout=timeout)
+                remote2 = tclient.request("xyz", session_timeout=timeout)
+                time.sleep(sleeptime)
+                f1 = remote.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+                f1 = remote2.get_function("rpc.test2.addone")
+                assert f1(10) == 11
+
+            except RuntimeError:
+                pass
+            remote3 = tclient.request("abc")
+            f1 = remote3.get_function("rpc.test2.addone")
+            assert f1(10) == 11
+            remote3 = tclient.request("xyz1")
+            f1 = remote3.get_function("rpc.test2.addone")
+            assert f1(10) == 11
+
+        check_timeout(0.01, 0.1)
+        check_timeout(2, 0)
+        tserver.terminate()
+        server0.terminate()
+        server1.terminate()
+        server2.terminate()
+        server3.terminate()
+        tproxy.terminate()
+    except ImportError:
+        print("Skip because tornado is not available")
+
+
+def check_tracker_rejects_oversized_msg_size():
+    """Tracker must reject an oversized msg_size header and close the connection
+    instead of buffering an unbounded amount of data on a single TCP connection.
+
+    Regression test for the unbounded buffer growth defect in
+    TCPEventHandler.on_message. See MAX_TRACKER_MSG_BYTES in tracker.py.
+    """
+    try:
+        # pylint: disable=import-outside-toplevel
+        import socket
+        import struct
+
+        from tvm.rpc import base, tracker
+
+        tserver = tracker.Tracker(port=9180, port_end=9290, silent=True)
+        try:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(5)
+            sock.connect(("127.0.0.1", tserver.port))
+            # complete the 4-byte magic handshake
+            sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+            magic_reply = sock.recv(4)
+            assert struct.unpack("<i", magic_reply)[0] == base.RPC_TRACKER_MAGIC
+
+            # send an oversized msg_size header (2 GiB)
+            sock.sendall(struct.pack("<i", 0x7FFFFFFF))
+
+            # server must close the connection (no payload buffering)
+            for _ in range(20):
+                chunk = sock.recv(4096)
+                if chunk == b"":
+                    break
+                time.sleep(0.05)
+            else:
+                raise AssertionError("tracker did not close connection after oversized msg_size")
+        finally:
+            tserver.terminate()
+    except ImportError:
+        print("Skip because tornado is not available")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    check_server_drop()
+    check_tracker_rejects_oversized_msg_size()

tests/python/contrib/test_sort.py

@@ -0,0 +1,102 @@
+# 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: E741
+"""Configure pytest"""
+
+# pylint: disable=invalid-name
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import te
+
+
+def test_sort():
+    """Tests sort function"""
+    n = 2
+    l = 5
+    m = 3
+    data = te.placeholder((n, l, m), name="data")
+    sort_num = te.placeholder((n, m), name="sort_num", dtype="int32")
+    axis = 1
+    is_ascend = False
+    out = te.extern(
+        data.shape,
+        [data, sort_num],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
+        ),
+        dtype="int32",
+        name="sort_tensor",
+    )
+    input_data = [
+        [[1, 2, 3], [2, 4.5, 3.5], [1.1, 0.5, 1], [3.2, -5, 0.5], [1.5, 0, 0]],
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]],
+    ]
+    sort_num_input = [[1, 2, 3], [4, 5, 5]]
+    sorted_index = [
+        [[0, 1, 1], [1, 0, 0], [2, 2, 2], [3, 3, 3], [4, 4, 4]],
+        [[3, 4, 4], [2, 3, 3], [1, 2, 2], [0, 1, 1], [4, 0, 0]],
+    ]
+
+    dev = tvm.cpu(0)
+    target = "llvm"
+    f = tvm.compile(te.create_prim_func([data, sort_num, out]), target=target)
+    a = tvm.runtime.tensor(np.array(input_data).astype(data.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.array(sort_num_input).astype(sort_num.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.zeros(a.shape, dtype=out.dtype.dtype), dev)
+    f(a, b, c)
+    tvm.testing.assert_allclose(
+        c.numpy(), np.array(sorted_index).astype(out.dtype.dtype), rtol=1e-5
+    )
+
+
+def test_sort_np():
+    """Tests sort function using numpy"""
+    dshape = (1, 2, 3, 4, 5, 6)
+    axis = 4
+    reduced_shape = (1, 2, 3, 4, 6)
+    is_ascend = True
+    data = te.placeholder(dshape, name="data")
+    sort_num = te.placeholder(reduced_shape, name="sort_num", dtype="int32")
+    out = te.extern(
+        data.shape,
+        [data, sort_num],
+        lambda ins, outs: tvm.tirx.call_packed(
+            "tvm.contrib.sort.argsort_nms", ins[0], ins[1], outs[0], axis, is_ascend
+        ),
+        dtype="int32",
+        name="sort_tensor",
+    )
+
+    dev = tvm.cpu(0)
+    target = "llvm"
+    f = tvm.compile(te.create_prim_func([data, sort_num, out]), target=target)
+
+    np_data = np.random.uniform(size=dshape)
+    np_out = np.argsort(np_data, axis=axis)
+    sort_num_input = np.full(reduced_shape, dshape[axis])
+    a = tvm.runtime.tensor(np.array(np_data).astype(data.dtype.dtype), dev)
+    b = tvm.runtime.tensor(np.array(sort_num_input).astype(sort_num.dtype.dtype), dev)
+    c = tvm.runtime.tensor(np.zeros(a.shape, dtype=out.dtype.dtype), dev)
+    f(a, b, c)
+    tvm.testing.assert_allclose(c.numpy(), np_out, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    test_sort()
+    test_sort_np()

tests/python/contrib/test_tir_triton_integration.py

@@ -0,0 +1,134 @@
+# 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 sys
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.relax.frontend import nn
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+from tvm.testing import env
+
+try:
+    import triton
+    import triton.language as tl
+    from packaging import version
+except ImportError:
+    pytestmark = pytest.skip("Triton is not available", allow_module_level=True)
+else:
+    if version.parse(triton.__version__) < version.parse("3.3.0"):
+        pytestmark = pytest.skip("Triton >= 3.3.0 is required", allow_module_level=True)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_tir_triton_integration():
+    @triton.jit
+    def add_kernel(
+        x_ptr,  # *Pointer* to first input vector.
+        y_ptr,  # *Pointer* to second input vector.
+        output_ptr,  # *Pointer* to output vector.
+        n_elements,  # Size of the vector.
+        BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
+    ):
+        """Triton vector add kernel from its tutorial."""
+        pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+        x = tl.load(x_ptr + offsets, mask=mask)
+        y = tl.load(y_ptr + offsets, mask=mask)
+        output = x + y
+        tl.store(output_ptr + offsets, output, mask=mask)
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def add(x_handle: T.handle, y_handle: T.handle, output_handle: T.handle) -> None:
+            T.func_attr({"global_symbol": "add"})
+            m = T.int64()
+            x = T.match_buffer(x_handle, (m,), "float32")
+            y = T.match_buffer(y_handle, (m,), "float32")
+            output = T.match_buffer(output_handle, (m,), "float32")
+            with T.sblock("root"):
+                T.reads(x[0:m], y[0:m])
+                T.writes(output[0:m])
+                BLOCK_SIZE = T.meta_var(64)
+                T.call_kernel(
+                    add_kernel,
+                    (T.ceildiv(m, BLOCK_SIZE),),
+                    x.data,
+                    y.data,
+                    output.data,
+                    m,
+                    BLOCK_SIZE,
+                    num_warps=8,
+                )
+
+        @R.function
+        def main(x: R.Tensor(("m",), "float32"), y: R.Tensor(("m",), "float32")):
+            m = T.int64()
+            with R.dataflow():
+                output = R.call_tir(Module.add, [x, y], relax.TensorType((m,), "float32"))
+                R.output(output)
+            return output
+
+    # Constexpr parameters (BLOCK_SIZE) stay in the kernel arguments, and the
+    # thread extent is 256 because the kernel is compiled with num_warps=8.
+    @I.ir_module(s_tir=True)
+    class Parsed:
+        @T.prim_func(s_tir=True)
+        def add(x_handle: T.handle, y_handle: T.handle, output_handle: T.handle):
+            m = T.int64()
+            x = T.match_buffer(x_handle, (m,))
+            y = T.match_buffer(y_handle, (m,))
+            output = T.match_buffer(output_handle, (m,))
+            with T.sblock("root"):
+                T.reads(x[0:m], y[0:m])
+                T.writes(output[0:m])
+                T.call_packed(
+                    "add_kernel",
+                    x.data,
+                    y.data,
+                    output.data,
+                    m,
+                    64,
+                    256,
+                    (m + T.int64(64) - T.int64(1)) // T.int64(64),
+                )
+
+    tvm.ir.assert_structural_equal(Module["add"], Parsed["add"])
+    assert len(Module.get_attr("external_mods")) == 1
+
+    with tvm.target.Target("cuda"):
+        lib = tvm.compile(Module)
+
+    def run_and_check():
+        device = tvm.cuda(0)
+        x_nd = tvm.runtime.tensor(np.random.rand(256).astype(np.float32), device)
+        y_nd = tvm.runtime.tensor(np.random.rand(256).astype(np.float32), device)
+        output_np = x_nd.numpy() + y_nd.numpy()
+        output_nd = tvm.runtime.vm.VirtualMachine(lib, device)["main"](x_nd, y_nd)
+        tvm.testing.assert_allclose(output_nd.numpy(), output_np, rtol=1e-5)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)

tests/python/contrib/test_tvmjs.py

@@ -0,0 +1,64 @@
+# 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 contrib.tvmjs"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm.testing
+from tvm.contrib import tvmjs
+
+dtype = tvm.testing.parameter(
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "float16",
+    "float32",
+    "float64",
+    "float8_e4m3fn",
+    "float8_e5m2",
+)
+
+
+def test_save_load_float8(dtype):
+    if "float8" in dtype or "bfloat16" in dtype:
+        ml_dtypes = pytest.importorskip("ml_dtypes")
+        np_dtype = np.dtype(getattr(ml_dtypes, dtype))
+    else:
+        np_dtype = np.dtype(dtype)
+
+    arr = np.arange(16, dtype=np_dtype)
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache({"arr": arr}, temp_dir)
+        cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu())
+
+    after_roundtrip = cache["arr"].numpy()
+
+    np.testing.assert_array_equal(arr, after_roundtrip)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()