chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,347 @@
+# 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.
+"""Basic tests for a Disco session"""
+
+# pylint: disable=missing-docstring
+import socket
+import subprocess
+import sys
+import tempfile
+import threading
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape
+from tvm_ffi.core import String
+
+import tvm
+import tvm.testing
+
+# Imported for the side effect of registering the tests.disco.* worker functions.
+from tvm.exec import disco_worker as _  # noqa: F401  # pylint: disable=unused-import
+from tvm.runtime import disco as di
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+if di is None:
+    pytest.skip("disco runtime is not available", allow_module_level=True)
+
+
+_SOCKET_SESSION_TESTER = None
+
+
+def _get_free_port():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+class SocketSessionTester:
+    """Run a disco SocketSession with one local node and remote nodes.
+
+    Each remote node is a `tvm.exec.disco_remote_socket_session` subprocess
+    launched with the current Python interpreter.
+    """
+
+    def __init__(self, num_workers, num_nodes=2, num_groups=1):
+        # Initialize the attributes used by __del__ first, so that teardown is
+        # safe even when __init__ raises below.
+        self.sess = None
+        self.remote_nodes = []
+        assert num_workers % num_nodes == 0
+        num_workers_per_node = num_workers // num_nodes
+        server_host = "localhost"
+        server_port = _get_free_port()
+        server_exc = []
+
+        def start_server():
+            try:
+                self.sess = di.SocketSession(
+                    num_nodes, num_workers_per_node, num_groups, server_host, server_port
+                )
+            except Exception as exc:  # pylint: disable=broad-except
+                server_exc.append(exc)
+
+        thread = threading.Thread(target=start_server)
+        thread.start()
+
+        cmd = "tvm.exec.disco_remote_socket_session"
+        for _i in range(num_nodes - 1):
+            self.remote_nodes.append(
+                subprocess.Popen(
+                    [
+                        sys.executable,
+                        "-m",
+                        cmd,
+                        server_host,
+                        str(server_port),
+                        str(num_workers_per_node),
+                    ],
+                    stdout=sys.stdout,
+                    stderr=sys.stderr,
+                )
+            )
+
+        thread.join()
+        if server_exc:
+            raise server_exc[0]
+
+    # Bound at class creation: module globals may already be cleared when
+    # __del__ runs during interpreter shutdown.
+    _TIMEOUT_EXPIRED = subprocess.TimeoutExpired
+
+    def __del__(self):
+        try:
+            # Shut down the session first so remote nodes can exit gracefully.
+            if self.sess is not None:
+                self.sess.shutdown()
+        finally:
+            for node in self.remote_nodes:
+                try:
+                    node.wait(timeout=10)
+                except self._TIMEOUT_EXPIRED:
+                    node.kill()
+                    node.wait()
+
+
+def create_socket_session(num_workers):
+    """Create a socket session backed by one local and one remote node.
+
+    The tester is kept alive in a module-level global so that the session
+    survives until the next call (or interpreter exit) replaces it.
+    """
+    global _SOCKET_SESSION_TESTER
+    # Rebind (not `del`) so the global stays defined if the constructor raises.
+    _SOCKET_SESSION_TESTER = None
+    _SOCKET_SESSION_TESTER = SocketSessionTester(num_workers)
+    assert _SOCKET_SESSION_TESTER.sess is not None
+    return _SOCKET_SESSION_TESTER.sess
+
+
+def _numpy_to_worker_0(sess: di.Session, np_array: np.array, device):
+    x_array = sess.empty(np_array.shape, "float32", device=device)
+    host_array = tvm.runtime.tensor(np_array, device=device)
+    sess.copy_to_worker_0(host_array, x_array)
+    return x_array
+
+
+def _numpy_from_worker_0(sess: di.Session, remote_array, shape, dtype):
+    host_array = tvm.runtime.empty(shape, dtype, device=tvm.cpu())
+    sess.copy_from_worker_0(host_array, remote_array)
+    sess.sync_worker_0()
+    return host_array.numpy()
+
+
+_all_session_kinds = [di.ThreadedSession, di.ProcessSession, create_socket_session]
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_int(session_kind):  # pylint: disable=invalid-name
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one")
+    result: di.DRef = func(1)
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == 2
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_float(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.add_one_float")
+    result: di.DRef = func(1.5)
+
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == 2.0
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_tensor(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    device = tvm.cpu(0)
+    x_np = np.arange(6).astype("float32").reshape([2, 3])
+    y_np = np.arange(6).astype("float32").reshape([2, 3]) + 1
+    x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+    y_disc = sess.get_global_func("tests.disco.add_one_tensor")(x_disc)
+    y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+    np.testing.assert_equal(y_nd, y_np)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_string(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.str")
+    result: di.DRef = func("hello")
+
+    for i in range(num_workers):
+        assert result.debug_get_from_remote(i) == "hello_suffix"
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_string_obj(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.str_obj")
+    result: di.DRef = func(String("hello"))
+
+    for i in range(num_workers):
+        value = result.debug_get_from_remote(i)
+        assert isinstance(value, str)
+        assert value == "hello_suffix"
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_shape_tuple(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+    func: di.DPackedFunc = sess.get_global_func("tests.disco.shape_tuple")
+    result: di.DRef = func(Shape([1, 2, 3]))
+    for i in range(num_workers):
+        value = result.debug_get_from_remote(i)
+        assert isinstance(value, Shape)
+        assert list(value) == [1, 2, 3, 4, 5]
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_vm_module(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+
+    # pylint: disable=invalid-name
+    @I.ir_module(s_tir=True)
+    class TestMod:
+        @T.prim_func(s_tir=True)
+        def transpose(A: T.Buffer((8, 16), "float32"), B: T.Buffer((16, 8), "float32")):
+            for i, j in T.grid(16, 8):
+                with T.sblock("transpose"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @R.function
+        def main(A: R.Tensor((8, 16), dtype="float32")) -> R.Tensor((16, 8), dtype="float32"):
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.transpose, (A,), out_ty=R.Tensor((16, 8), dtype="float32"))
+                R.output(B)
+            return B
+
+    # pylint: enable=invalid-name
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        device = tvm.cpu()
+        x_np = np.arange(8 * 16).astype("float32").reshape([8, 16])
+        y_np = x_np.transpose()
+
+        tvm.compile(TestMod, target="llvm").export_library(path)
+        mod = sess.load_vm_module(path, device=device)
+
+        x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+        y_disc = mod["main"](x_disc)
+        y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+        np.testing.assert_equal(y_nd, y_np)
+
+        # sync all workers to make sure the temporary files are cleaned up after all workers
+        # finish the execution
+        for i in range(num_workers):
+            sess._sync_worker(i)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+def test_vm_multi_func(session_kind):
+    num_workers = 4
+    sess = session_kind(num_workers=num_workers)
+
+    # pylint: disable=invalid-name
+    @I.ir_module(s_tir=True)
+    class TestMod:
+        @T.prim_func(s_tir=True)
+        def t1(A: T.Buffer((8, 16), "float32"), B: T.Buffer((16, 8), "float32")):
+            for i, j in T.grid(16, 8):
+                with T.sblock("t1"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @T.prim_func(s_tir=True)
+        def t2(A: T.Buffer((16, 8), "float32"), B: T.Buffer((8, 16), "float32")):
+            for i, j in T.grid(8, 16):
+                with T.sblock("t2"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    B[vi, vj] = A[vj, vi]
+
+        @R.function
+        def transpose_1(A: R.Tensor((8, 16), dtype="float32")) -> R.Tensor(
+            (16, 8), dtype="float32"
+        ):
+            R.func_attr({"global_symbol": "transpose_1"})
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.t1, (A,), out_ty=R.Tensor((16, 8), dtype="float32"))
+                R.output(B)
+            return B
+
+        @R.function
+        def transpose_2(A: R.Tensor((16, 8), dtype="float32")) -> R.Tensor(
+            (8, 16), dtype="float32"
+        ):
+            R.func_attr({"global_symbol": "transpose_2"})
+            cls = TestMod
+            with R.dataflow():
+                B = R.call_tir(cls.t2, (A,), out_ty=R.Tensor((8, 16), dtype="float32"))
+                R.output(B)
+            return B
+
+    # pylint: enable=invalid-name
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = tmpdir + "/test.so"
+        device = tvm.cpu()
+        x_np = np.arange(8 * 16).astype("float32").reshape([8, 16])
+        y_np = x_np.transpose()
+
+        tvm.compile(TestMod, target="llvm").export_library(path)
+        mod = sess.load_vm_module(path, device=device)
+
+        x_disc = _numpy_to_worker_0(sess, x_np, device=device)
+        y_disc = mod["transpose_1"](x_disc)
+        z_disc = mod["transpose_2"](y_disc)
+        y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
+        z_nd = _numpy_from_worker_0(sess, z_disc, shape=x_np.shape, dtype=x_np.dtype)
+        np.testing.assert_equal(y_nd, y_np)
+        np.testing.assert_equal(z_nd, x_np)
+
+        # sync all workers to make sure the temporary files are cleaned up after all workers
+        # finish the execution
+        for i in range(num_workers):
+            sess._sync_worker(i)
+
+
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("num_workers", [1, 2, 4])
+def test_num_workers(session_kind, num_workers):
+    if session_kind == create_socket_session and num_workers < 2:
+        return
+    sess = session_kind(num_workers=num_workers)
+    assert sess.num_workers == num_workers
+
+
+if __name__ == "__main__":
+    tvm.testing.main()