chore: import upstream snapshot with attribution

python/tvm/exec/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+"""Namespace of executables python files that directly run throw cmd"""

python/tvm/exec/disco_remote_socket_session.py

@@ -0,0 +1,36 @@
+# 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
+# ruff: noqa: F401
+"""Launch disco session in the remote node and connect to the server."""
+
+import sys
+
+import tvm
+
+from . import disco_worker as _  # pylint: disable=unused-import
+
+if __name__ == "__main__":
+    if len(sys.argv) != 4:
+        print("Usage: <server_host> <server_port> <num_workers>")
+        sys.exit(1)
+
+    server_host = sys.argv[1]
+    server_port = int(sys.argv[2])
+    num_workers = int(sys.argv[3])
+    func = tvm.get_global_func("runtime.disco.RemoteSocketSession")
+    func(server_host, server_port, num_workers)

python/tvm/exec/disco_worker.py

@@ -0,0 +1,127 @@
+# 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
+# ruff: noqa: RUF005
+"""Internal DiscoWorker for Disco ProcessSession."""
+
+import os
+import sys
+from collections.abc import Callable
+
+from tvm_ffi import Shape, get_global_func, register_global_func
+from tvm_ffi.core import String
+
+import tvm
+from tvm.runtime import Tensor, tensor
+
+
+@register_global_func("tests.disco.add_one", override=True)
+def _add_one(x: int) -> int:
+    return x + 1
+
+
+@register_global_func("tests.disco.add_one_float", override=True)
+def _add_one_float(x: float):
+    return x + 0.5
+
+
+@register_global_func("tests.disco.add_one_tensor", override=True)
+def _add_one_tensor(x: Tensor) -> Tensor:
+    return tensor(x.numpy() + 1)
+
+
+@register_global_func("tests.disco.str", override=True)
+def _str_func(x: str):
+    return x + "_suffix"
+
+
+@register_global_func("tests.disco.str_obj", override=True)
+def _str_obj_func(x: str):
+    assert isinstance(x, str)
+    return String(x + "_suffix")
+
+
+@register_global_func("tests.disco.shape_tuple", override=True)
+def _shape_tuple_func(x: Shape):
+    assert isinstance(x, Shape)
+    return Shape(list(x) + [4, 5])
+
+
+@register_global_func("tests.disco.test_callback", override=True)
+def _make_callback(device: tvm.runtime.Device) -> Callable[[str, int], Tensor]:
+    """For use in tests/python/disco/test_callback.py
+
+    This function simulates a callback to be used for lazy parameter
+    loading.
+
+    Parameters
+    ----------
+    device: tvm.runtime.Device
+
+        The device on which parameters should be located, when
+        returned by the callback function.
+
+    Returns
+    -------
+    fget_item: Callable[[str,int], Tensor]
+
+        A callback function that accepts a parameter's name and index,
+        and returns the specified parameter.
+
+    """
+    import numpy as np  # pylint: disable=import-outside-toplevel
+
+    def fget_item(param_name: str, param_index: int) -> Tensor:
+        if param_index == 0:
+            assert param_name == "A"
+            arr = np.arange(16).reshape([4, 4]).astype("int32")
+        elif param_index == 1:
+            assert param_name == "B"
+            arr = np.arange(4).reshape([2, 2]).astype("float32")
+        else:
+            raise ValueError(f"Unexpected index {param_index}")
+        return tvm.runtime.tensor(arr, device=device)
+
+    return fget_item
+
+
+def main():
+    """Main worker function"""
+    if len(sys.argv) != 6:
+        print("Usage: <worker_id> <num_workers> <num_groups> <read_fd> <write_fd>")
+        return
+    worker_id = int(sys.argv[1])
+    num_workers = int(sys.argv[2])
+    num_groups = int(sys.argv[3])
+    if sys.platform == "win32":
+        import msvcrt  # pylint: disable=import-outside-toplevel,import-error
+
+        reader = msvcrt.open_osfhandle(int(sys.argv[4]), os.O_BINARY)
+        writer = msvcrt.open_osfhandle(int(sys.argv[5]), os.O_BINARY)
+    else:
+        reader = int(sys.argv[4])
+        writer = int(sys.argv[5])
+
+    worker_func = get_global_func("runtime.disco.WorkerProcess")
+    worker_func(worker_id, num_workers, num_groups, reader, writer)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (OSError, KeyboardInterrupt):
+        pass

python/tvm/exec/gpu_memory_bandwidth.py

@@ -0,0 +1,262 @@
+# 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: E501, F821
+"""A script to measure GPU memory bandwidth"""
+
+import argparse
+import itertools
+
+import numpy as np
+
+import tvm
+from tvm import te
+from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig
+from tvm.testing import local_run, rpc_run
+
+
+def _parse_args() -> argparse.Namespace:
+    def _parse_list_int(source: str):
+        return [int(i) for i in source.split(",")]
+
+    parser = argparse.ArgumentParser(
+        prog="GPU memory bandwidth testing",
+        description="""Example for host GPU:
+    python -m tvm.exec.gpu_memory_bandwidth "nvidia/geforce-rtx-3090-ti" \
+        --dtype "float32"
+        --bx "8,16,32,64,128,256"      \
+        --tx "32,64,128,256,512,1024"  \
+        --vec "1,2,4" \
+
+    Example for Android GPU: \
+    python -m tvm.exec.gpu_memory_bandwidth "opencl" --target_host '{"kind": "llvm", "mtriple": "arm64-linux-android"}' \
+        --rpc_host "127.0.0.1" \
+        --rpc_port 9190 \
+        --rpc_key "android" \
+        --export_func "ndk" \
+        --dtype "float32" \
+        --bx "8,16,32,64,128,256"      \
+        --tx "32,64,128,256,512,1024"  \
+        --vec "1,2,4" \
+""",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "target",
+        type=str,
+        help="The target to be benchmarked",
+    )
+    parser.add_argument(
+        "--target_host",
+        type=str,
+        default=None,
+        help="The target host for build",
+    )
+    parser.add_argument(
+        "--xo",
+        type=int,
+        default=1024,
+        help="The value of `XO` in [XO, K, XI] => [XO, XI] reduction",
+    )
+    parser.add_argument(
+        "--k",
+        type=int,
+        default=64,
+        help="The value of `K` in [XO, K, XI] => [XO, XI] reduction",
+    )
+    parser.add_argument(
+        "--xi",
+        type=int,
+        default=4096,
+        help="The value of `XI` in [XO, K, XI] -> [XO, XI] reduction",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="float32",
+        help="The data type to be used in the workload",
+    )
+    parser.add_argument(
+        "--bx",
+        type=_parse_list_int,
+        default=[8, 16, 32, 64, 128, 256],
+        help="The value to be used to split `XO` into [BX, _]",
+    )
+    parser.add_argument(
+        "--tx",
+        type=_parse_list_int,
+        default=[32, 64, 128, 256, 512, 1024],
+        help="Number of threads to be used",
+    )
+    parser.add_argument(
+        "--vec",
+        type=_parse_list_int,
+        default=[1, 2, 4],
+        help="Vector length to be used in vectorized load",
+    )
+    parser.add_argument(
+        "--rpc_host",
+        type=str,
+        default=None,
+        help="The address of RPC host (default: None, that means that RPC is not used)",
+    )
+    parser.add_argument(
+        "--rpc_port",
+        type=int,
+        default=None,
+        help="The port of RPC connection (default: None, that means that RPC is not used)",
+    )
+    parser.add_argument(
+        "--rpc_key",
+        type=str,
+        default=None,
+        help="The device key in RPC tracker (default: None, that means that RPC is not used)",
+    )
+    parser.add_argument(
+        "--export_func",
+        type=str,
+        default="tar",
+        help="Export function, actual only for RPC",
+        choices=["tar", "ndk"],
+    )
+    return parser.parse_args()
+
+
+def _workload(
+    len_xo: int,
+    len_k: int,
+    len_xi: int,
+    dtype: str,
+):
+    # pylint: disable=invalid-name
+    A = te.placeholder((len_xo, len_k, len_xi), dtype=dtype, name="A")
+    k = te.reduce_axis((0, len_k), "k")
+    B = te.compute(
+        (len_xo, len_xi),
+        lambda i, j: te.sum(A[i, k, j], axis=k),
+        name="B",
+    )
+    # pylint: enable=invalid-name
+    return te.create_prim_func([A, B])
+
+
+def _schedule(
+    sch: s_tir.Schedule,
+    len_bx: int,
+    len_tx: int,
+    len_vec: int,
+):
+    # pylint: disable=invalid-name
+    block = sch.get_sblock("B")
+    xo, xi, k = sch.get_loops(block)
+    bx, xo = sch.split(xo, factors=[len_bx, None])
+    xi, tx, vec = sch.split(xi, factors=[None, len_tx, len_vec])
+    sch.reorder(bx, xi, tx, xo, k, vec)
+    bx = sch.fuse(bx, xi)
+    sch.bind(bx, "blockIdx.x")
+    sch.bind(tx, "threadIdx.x")
+    ldg = sch.cache_read(block, 0, "local")
+    sch.compute_at(ldg, k, preserve_unit_loops=True)
+    sch.vectorize(sch.get_loops(ldg)[-1])
+    sch.decompose_reduction(block, k)
+    # pylint: enable=invalid-name
+
+
+def main():  # pylint: disable=too-many-locals
+    """Entry point"""
+    args = _parse_args()
+    # pylint: disable=invalid-name
+    target = tvm.target.Target(args.target)
+    if args.target_host is not None:
+        target = tvm.target.Target(args.target, host=args.target_host)
+    dtype = args.dtype
+    rpcConfig = None
+    if args.rpc_host is not None and args.rpc_port is not None and args.rpc_key is not None:
+        rpcConfig = RPCConfig(
+            tracker_host=args.rpc_host,
+            tracker_port=args.rpc_port,
+            tracker_key=args.rpc_key,
+            session_priority=1,
+            session_timeout_sec=600,
+        )
+
+    a = np.random.uniform(-1, 1, (args.xo, args.k, args.xi)).astype(dtype)
+    b = np.zeros((args.xo, args.xi), dtype=dtype)
+    num_bytes = a.size * a.itemsize + b.size * b.itemsize
+    print("###### Bandwidth Test ######")
+    print(
+        f"Workload [XO, K, XI] => [XO, XI]. "
+        f"[{args.xo}, {args.k}, {args.xi}] => [{args.xo}, {args.xi}]"
+    )
+    print(f"Input size: {num_bytes / 1048576} MB")
+    print(f"Target: {target}")
+
+    # pylint: enable=invalid-name
+    best_bandwidth = -1
+    for len_bx, len_tx, len_vec in itertools.product(
+        args.bx,
+        args.tx,
+        args.vec,
+    ):
+        func = _workload(
+            len_xo=args.xo,
+            len_k=args.k,
+            len_xi=args.xi,
+            dtype=dtype,
+        )
+        sch = tvm.s_tir.Schedule(func)
+        _schedule(sch, len_bx, len_tx, len_vec)
+
+        if rpcConfig is None:
+            _, profile_result = local_run(
+                tvm.compile(sch.mod, target=target),
+                target.kind.name,
+                [a, b],
+                evaluator_config=EvaluatorConfig(
+                    number=10,
+                    repeat=1,
+                    min_repeat_ms=100,
+                    enable_cpu_cache_flush=False,
+                ),
+            )
+        else:
+            _, profile_result = rpc_run(
+                tvm.compile(sch.mod, target=target),
+                target.kind.name,
+                [a, b],
+                evaluator_config=EvaluatorConfig(
+                    number=10,
+                    repeat=1,
+                    min_repeat_ms=100,
+                    enable_cpu_cache_flush=False,
+                ),
+                rpc_config=rpcConfig,
+                export_func=args.export_func,
+            )
+        bandwidth = num_bytes / profile_result.mean / (1024**3)
+        bx = len_bx * args.xi // (len_tx * len_vec)  # pylint: disable=invalid-name
+        mbs = num_bytes / 1024 / 1024
+        print(
+            f"bandwidth = {bandwidth:.3f} GB/s, bx = {bx}, tx = {len_tx}, "
+            f"len_vec = {len_vec}, bytes = {mbs} MB"
+        )
+        if bandwidth > best_bandwidth:
+            best_bandwidth = bandwidth
+    print(f"peak bandwidth: {best_bandwidth:.3f} GB/s")
+
+
+if __name__ == "__main__":
+    main()

python/tvm/exec/popen_worker.py

@@ -0,0 +1,126 @@
+# 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.
+
+# 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
+# ruff: noqa: RUF012
+"""Internal PopenWorker for PopenPool."""
+
+import logging
+import os
+import pickle
+import struct
+import sys
+import threading
+import traceback
+
+import cloudpickle
+
+from tvm.support.popen_pool import StatusKind
+
+
+class TimeoutStatus:
+    __slot__ = ["status"]
+
+    def __init__(self):
+        self.status = StatusKind.RUNNING
+
+
+def main():
+    """Main worker function"""
+    if len(sys.argv) != 3:
+        print("Usage: <read_fd> <write_fd>")
+        return
+    if sys.platform == "win32":
+        # pylint: disable=import-outside-toplevel
+        import msvcrt
+
+        reader = os.fdopen(msvcrt.open_osfhandle(int(sys.argv[1]), os.O_BINARY), "rb")
+        writer = os.fdopen(msvcrt.open_osfhandle(int(sys.argv[2]), os.O_BINARY), "wb")
+    else:
+        reader = os.fdopen(int(sys.argv[1]), "rb")
+        writer = os.fdopen(int(sys.argv[2]), "wb")
+
+    logging.basicConfig(level=logging.INFO)
+
+    lock = threading.Lock()
+
+    def _respond(ret_value):
+        """Send data back to the client."""
+        data = cloudpickle.dumps(ret_value, protocol=pickle.HIGHEST_PROTOCOL)
+        writer.write(struct.pack("<i", len(data)))
+        writer.write(data)
+        writer.flush()
+
+    def _cancel_run(status):
+        lock.acquire()
+        if status.status == StatusKind.RUNNING:
+            _respond((StatusKind.TIMEOUT, TimeoutError()))
+            status.status = StatusKind.TIMEOUT
+        lock.release()
+
+    while True:
+        raw_bytes_size = reader.read(4)
+        if len(raw_bytes_size) != 4:
+            # the parent exited
+            return
+        bytes_size = struct.unpack("<i", raw_bytes_size)[0]
+        fn, args, kwargs, timeout = cloudpickle.loads(reader.read(bytes_size))
+        status = TimeoutStatus()
+
+        if timeout is not None:
+            watcher = threading.Timer(timeout, _cancel_run, [status])
+            watcher.daemon = True
+            watcher.start()
+
+        # pylint: disable=broad-except
+        try:
+            result = fn(*args, **kwargs)
+            ret_value = (StatusKind.COMPLETE, result)
+        except Exception as exception:
+            msg = traceback.format_exc()
+            ret_value = (StatusKind.EXCEPTION, type(exception)(msg))
+
+        if timeout is not None:
+            watcher.cancel()
+
+        lock.acquire()
+        if status.status == StatusKind.RUNNING:
+            _respond(ret_value)
+            status.status = StatusKind.COMPLETE
+        lock.release()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (OSError, KeyboardInterrupt):
+        pass

python/tvm/exec/query_rpc_tracker.py

@@ -0,0 +1,48 @@
+# 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.
+"""Tool to query RPC tracker status"""
+
+import argparse
+import logging
+import os
+
+from .. import rpc
+
+
+def main():
+    """Main function"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="", help="the hostname of the tracker")
+    parser.add_argument("--port", type=int, default=None, help="The port of the RPC")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+
+    # default to local host or environment variable
+    if not args.host:
+        args.host = os.environ.get("TVM_TRACKER_HOST", "127.0.0.1")
+
+    if not args.port:
+        args.port = int(os.environ.get("TVM_TRACKER_PORT", "9190"))
+
+    conn = rpc.connect_tracker(args.host, args.port)
+    # pylint: disable=superfluous-parens
+    print(f"Tracker address {args.host}:{args.port}\n")
+    print(f"{conn.text_summary()}")
+
+
+if __name__ == "__main__":
+    main()

python/tvm/exec/rpc_proxy.py

@@ -0,0 +1,113 @@
+# 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=redefined-outer-name, invalid-name
+"""RPC web proxy, allows redirect to websocket based RPC servers(browsers)"""
+
+import argparse
+import glob
+import logging
+import os
+
+from tvm.rpc.proxy import Proxy
+
+
+def find_example_resource():
+    """Find resource examples."""
+    curr_path = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
+    base_path = os.path.abspath(os.path.join(curr_path, "..", "..", ".."))
+    index_page = os.path.join(base_path, "web", "apps", "browser", "rpc_server.html")
+    default_plugin_page = os.path.join(base_path, "web", "apps", "browser", "rpc_plugin.html")
+
+    resource_files = [
+        ("/", os.path.join(base_path, "web", "dist", "tvmjs.bundle.js")),
+        ("/", os.path.join(base_path, "web", "dist", "wasm", "tvmjs_runtime.wasi.js")),
+        ("/", index_page),
+    ]
+    allow_format = ("json", "bin", "js", "wasm", "html", "css", "model")
+
+    # recursively apend things in www, up to two levels
+    resource_bases = [
+        os.path.join(base_path, "web", "dist", "www"),
+        os.path.join(base_path, "web", ".tensor_cache"),
+    ]
+    for base in resource_bases:
+        if not os.path.isdir(base):
+            continue
+        for full_name in glob.glob(f"{base}/**", recursive=True):
+            fname = os.path.relpath(full_name, base)
+            dirname = os.path.dirname(fname)
+            fmt = fname.rsplit(".", 1)[-1]
+            if os.path.isfile(full_name) and fmt in allow_format:
+                resource_files.append((dirname, full_name))
+
+    for item in resource_files:
+        fname = item[-1]
+        if not os.path.exists(fname):
+            raise RuntimeError(f"Cannot find {fname}")
+
+    if not any(item[-1].endswith("rpc_plugin.html") for item in resource_files):
+        resource_files.append(("/", default_plugin_page))
+
+    return index_page, resource_files
+
+
+def main(args):
+    """Main function"""
+    if args.tracker:
+        url, port = args.tracker.split(":")
+        port = int(port)
+        tracker_addr = (url, port)
+    else:
+        tracker_addr = None
+
+    if args.example_rpc:
+        index, js_files = find_example_resource()
+        prox = Proxy(
+            args.host,
+            port=args.port,
+            port_end=args.port_end,
+            web_port=args.web_port,
+            index_page=index,
+            resource_files=js_files,
+            tracker_addr=tracker_addr,
+        )
+    else:
+        prox = Proxy(
+            args.host,
+            port=args.port,
+            port_end=args.port_end,
+            web_port=args.web_port,
+            tracker_addr=tracker_addr,
+        )
+    prox.proc.join()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="the hostname of the server")
+    parser.add_argument("--port", type=int, default=9090, help="The port of the RPC")
+    parser.add_argument("--port-end", type=int, default=9199, help="The end search port of the RPC")
+    parser.add_argument(
+        "--web-port", type=int, default=8888, help="The port of the http/websocket server"
+    )
+    parser.add_argument(
+        "--example-rpc", type=bool, default=False, help="Whether to switch on example rpc mode"
+    )
+    parser.add_argument("--tracker", type=str, default="", help="Report to RPC tracker")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    main(args)

python/tvm/exec/rpc_server.py

@@ -0,0 +1,104 @@
+# 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=redefined-outer-name, invalid-name
+"""Start an RPC server"""
+
+import argparse
+import logging
+
+from .. import rpc
+
+
+def main(args):
+    """Main function
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        parsed args from command-line invocation
+    """
+    if args.tracker:
+        url, port = args.tracker.rsplit(":", 1)
+        port = int(port)
+        tracker_addr = (url, port)
+        if not args.key:
+            raise RuntimeError("Need key to present type of resource when tracker is available")
+    else:
+        tracker_addr = None
+
+    server = rpc.Server(
+        args.host,
+        args.port,
+        args.port_end,
+        is_proxy=args.through_proxy,
+        key=args.key,
+        tracker_addr=tracker_addr,
+        load_library=args.load_library,
+        custom_addr=args.custom_addr,
+        silent=args.silent,
+        no_fork=not args.fork,
+    )
+    server.proc.join()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="The host IP address the tracker binds to"
+    )
+    parser.add_argument("--port", type=int, default=9090, help="The port of the RPC")
+    parser.add_argument(
+        "--through-proxy",
+        dest="through_proxy",
+        action="store_true",
+        help=(
+            "Whether this server provide service through a proxy. If this is true, the host and"
+            "port actually is the address of the proxy."
+        ),
+    )
+    parser.add_argument("--port-end", type=int, default=9199, help="The end search port of the RPC")
+    parser.add_argument(
+        "--tracker",
+        type=str,
+        help=("The address of RPC tracker in host:port format. e.g. (10.77.1.234:9190)"),
+    )
+    parser.add_argument(
+        "--key", type=str, default="", help="The key used to identify the device type in tracker."
+    )
+    parser.add_argument("--silent", action="store_true", help="Whether run in silent mode.")
+    parser.add_argument("--load-library", type=str, help="Additional library to load")
+    parser.add_argument(
+        "--no-fork",
+        dest="fork",
+        action="store_false",
+        help="Use spawn mode to avoid fork. This option \
+                        is able to avoid potential fork problems with Metal, OpenCL \
+                        and ROCM compilers.",
+    )
+    parser.add_argument(
+        "--custom-addr", type=str, help="Custom IP Address to Report to RPC Tracker"
+    )
+
+    parser.set_defaults(fork=True)
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    if args.fork is not False and not args.silent:
+        logging.info(
+            "If you are running ROCM/Metal, fork will cause "
+            "compiler internal error. Try to launch with arg ```--no-fork```"
+        )
+    main(args)

python/tvm/exec/rpc_tracker.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.
+# pylint: disable=redefined-outer-name, invalid-name
+"""Tool to start RPC tracker"""
+
+import argparse
+import logging
+
+from ..rpc.tracker import Tracker
+
+
+def main(args):
+    """Main function"""
+    tracker = Tracker(args.host, port=args.port, port_end=args.port_end, silent=args.silent)
+    tracker.proc.join()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="The host IP address the tracker binds to"
+    )
+    parser.add_argument("--port", type=int, default=9190, help="The port of the RPC")
+    parser.add_argument("--port-end", type=int, default=9199, help="The end search port of the RPC")
+    parser.add_argument("--silent", action="store_true", help="Whether run in silent mode.")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    main(args)