[Runtime][Disco][Relax] Add CPUCCL support for distributed execution

CMakeLists.txt

@@ -46,7 +46,7 @@ tvm_option(USE_OPENCL "Build with OpenCL" OFF)
 tvm_option(USE_OPENCL_ENABLE_HOST_PTR "Enable OpenCL memory object access to host" OFF)
 tvm_option(USE_OPENCL_GTEST "Path to OpenCL specific gtest version for runtime cpp tests." /path/to/opencl/gtest)
 tvm_option(USE_VULKAN "Build with Vulkan" OFF)
-
+tvm_option(USE_CPU_CCL "Build with CPU CCL" OFF)
 
 # Whether to use spirv-tools and SPIRV-Headers from Khronos GitHub or GitLab.
 #
@@ -360,7 +360,7 @@ tvm_file_glob(GLOB RUNTIME_SRCS
   src/runtime/vm/*.cc
   src/runtime/memory/*.cc
   src/runtime/rpc/minrpc/*.cc
-)
+  )
 # Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
 # Note: src/backend/*/runtime sources move to per-backend DSOs.
 set(TVM_RUNTIME_EXT_OBJS "")
@@ -524,6 +524,15 @@ if(USE_ROCM AND USE_RCCL)
   target_link_libraries(tvm_runtime_extra PRIVATE tvm_rccl_objs rccl)
 endif()
 
+# CPUCCL.
+if(USE_CPU_CCL)
+  message(STATUS "Build with CPU CCL support...")
+  tvm_file_glob(GLOB _cpu_ccl_srcs src/runtime/extra/disco/cpuccl/*.cc)
+  add_library(tvm_cpu_ccl_objs OBJECT ${_cpu_ccl_srcs})
+  target_link_libraries(tvm_cpu_ccl_objs PRIVATE tvm_runtime_extra_defs)
+  target_link_libraries(tvm_runtime_extra PRIVATE tvm_cpu_ccl_objs)
+endif()
+
 target_link_libraries(tvm_runtime_extra PUBLIC tvm_runtime)
 
 # If disco/cuda_ipc is included, link the CUDA DSO.

cmake/config.cmake

@@ -322,3 +322,8 @@ SET(CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE "x64")
 
 # Enable Qualcomm OpenCL extension support
 set(USE_OPENCL_EXTN_QCOM OFF)
+
+# Whether to enable CPU CCL support:
+# - ON: enable CPU collectives over a pipe ring between disco workers
+# - OFF: disable CPU CCL
+set(USE_CPU_CCL OFF)

include/tvm/relax/attrs/ccl.h

@@ -85,6 +85,23 @@ struct ScatterCollectiveAttrs : public tvm::AttrsNode {
                                     AttrsNode);
 };  // struct ScatterCollectiveAttrs
 
+/*! \brief Attributes used in gather operators */
+struct GatherCollectiveAttrs : public tvm::AttrsNode {
+  int num_workers;
+  bool in_group;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<GatherCollectiveAttrs>()
+        .def_ro("num_workers", &GatherCollectiveAttrs::num_workers,
+                "The number of workers to gather data from.")
+        .def_ro("in_group", &GatherCollectiveAttrs::in_group,
+                "Whether the gather operation performs in group or globally as default.");
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherCollectiveAttrs", GatherCollectiveAttrs,
+                                    AttrsNode);
+};  // struct GatherCollectiveAttrs
+
 }  // namespace relax
 }  // namespace tvm
 

include/tvm/runtime/disco/disco_worker.h

@@ -58,7 +58,9 @@ class DiscoWorker {
         default_device(Device{DLDeviceType::kDLCPU, 0}),
         worker_zero_data(worker_zero_data),
         channel(channel),
-        register_file{} {}
+        register_file{},
+        ring_in(nullptr),
+        ring_out(nullptr) {}
 
   /*! \brief Main loop of the worker */
   void MainLoop();
@@ -96,6 +98,9 @@ class DiscoWorker {
 
   struct Impl;
   friend struct DiscoWorker::Impl;
+  /*! \brief CPU communication ring endpoints assigned by the session during BuildRing(). */
+  DiscoRingChannel* ring_in;
+  DiscoRingChannel* ring_out;
 };
 /*!
  * \brief A threadlocal wrapper of DiscoWorker.

include/tvm/runtime/disco/session.h

@@ -76,6 +76,7 @@
 #include <tvm/ffi/container/shape.h>
 #include <tvm/ffi/function.h>
 #include <tvm/runtime/tensor.h>
+#include <unistd.h>
 
 #include <mutex>
 #include <queue>
@@ -294,8 +295,10 @@ class Session : public ffi::ObjectRef {
    * \brief Create a session backed by a thread pool of workers
    * \param num_workers The number of workers.
    * \param num_groups The number of worker groups.
+   * \param build_ring Whether to wire the workers into a unidirectional pipe ring
+   * (worker `i` -> worker `i + 1`), used as the data plane for CPU collectives.
    */
-  TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups);
+  TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups, bool build_ring);
   /*!
    * \brief Create a session backed by pipe-based multiprocessing
    * \param num_workers The number of workers.
@@ -305,12 +308,14 @@ class Session : public ffi::ObjectRef {
    * When `worker-id` is 0, it shuts down the process pool; Otherwise, it retursn a tuple
    * (read_fd, writefd) used to communicate with the corresponding worker.
    * \param entrypoint The entrypoint of DiscoWorker main worker function.
+   * \param build_ring Whether to wire the workers into a unidirectional pipe ring
+   * (worker `i` -> worker `i + 1`), used as the data plane for CPU collectives.
    * \note Worker-0 is always co-located with the controler as a separate thread, and therefore
    * worker-0 does not exist in the process pool.
    */
   TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups,
                                                 ffi::String process_pool_creator,
-                                                ffi::String entrypoint);
+                                                ffi::String entrypoint, bool build_ring);
 
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj);
 };
@@ -332,6 +337,72 @@ class DiscoChannel {
   virtual ffi::PackedArgs RecvReply() = 0;
 };
 
+class DiscoRingChannel {
+ public:
+  explicit DiscoRingChannel(int fd) : fd_(fd) {
+    TVM_FFI_ICHECK_GE(fd_, 0) << "DiscoRingChannel: invalid fd " << fd_;
+  }
+  ~DiscoRingChannel() { Close(); }
+
+  DiscoRingChannel(const DiscoRingChannel&) = delete;
+  DiscoRingChannel& operator=(const DiscoRingChannel&) = delete;
+  DiscoRingChannel(DiscoRingChannel&& other) noexcept : fd_(other.fd_) { other.fd_ = -1; }
+
+  void Send(const void* data, size_t size) {
+    const char* write_data = static_cast<const char*>(data);
+    while (size > 0) {
+      ssize_t n;
+      do {
+        n = ::write(fd_, write_data, size);
+      } while (n < 0 && errno == EINTR);
+      TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Send failed: " << std::strerror(errno);
+      write_data += n;
+      size -= n;
+    }
+  }
+
+  void Recv(void* data, size_t size) {
+    char* read_data = static_cast<char*>(data);
+    while (size > 0) {
+      ssize_t n;
+      do {
+        n = ::read(fd_, read_data, size);
+      } while (n < 0 && errno == EINTR);
+      TVM_FFI_ICHECK_GT(n, 0) << "DiscoRingChannel::Recv failed: " << std::strerror(errno);
+      read_data += n;
+      size -= n;
+    }
+  }
+
+  ssize_t ReadSome(void* data, size_t max_size) {
+    if (fd_ < 0) return 0;
+    ssize_t n;
+    do {
+      n = ::read(fd_, data, max_size);
+    } while (n < 0 && errno == EINTR);
+    return n;
+  }
+
+  ssize_t WriteSome(const void* data, size_t size) {
+    if (fd_ < 0) return -1;
+    ssize_t n;
+    do {
+      n = ::write(fd_, data, size);
+    } while (n < 0 && errno == EINTR);
+    return n;
+  }
+
+  void Close() {
+    if (fd_ >= 0) {
+      ::close(fd_);
+      fd_ = -1;
+    }
+  }
+
+ private:
+  int fd_;
+};
+
 /*!
  * \brief A special communication channel between controler and worker-0,
  * assuming they are always collocated in the same process.

python/tvm/exec/disco_worker.py

@@ -101,12 +101,21 @@ def fget_item(param_name: str, param_index: int) -> Tensor:
 
 def main():
     """Main worker function"""
-    if len(sys.argv) != 6:
-        print("Usage: <worker_id> <num_workers> <num_groups> <read_fd> <write_fd>")
+    argc = len(sys.argv)
+    if argc not in (6, 8):
+        print(
+            "Usage: <worker_id> <num_workers> <num_groups> "
+            "<read_fd> <write_fd> [<ring_in_fd> <ring_out_fd>]"
+        )
         return
     worker_id = int(sys.argv[1])
     num_workers = int(sys.argv[2])
     num_groups = int(sys.argv[3])
+
+    ring_out_fd = -1
+    ring_in_fd = -1
+    has_ring = argc == 8
+
     if sys.platform == "win32":
         import msvcrt  # pylint: disable=import-outside-toplevel,import-error
 
@@ -115,9 +124,12 @@ def main():
     else:
         reader = int(sys.argv[4])
         writer = int(sys.argv[5])
+        if has_ring:
+            ring_in_fd = int(sys.argv[6])
+            ring_out_fd = int(sys.argv[7])
 
     worker_func = get_global_func("runtime.disco.WorkerProcess")
-    worker_func(worker_id, num_workers, num_groups, reader, writer)
+    worker_func(worker_id, num_workers, num_groups, reader, writer, ring_in_fd, ring_out_fd)
 
 
 if __name__ == "__main__":

python/tvm/relax/op/ccl/__init__.py

@@ -17,4 +17,10 @@
 # under the License.
 """CCL related operators."""
 
-from .ccl import allgather, allreduce, broadcast_from_worker0, scatter_from_worker0
+from .ccl import (
+    allgather,
+    allreduce,
+    broadcast_from_worker0,
+    gather_to_worker0,
+    scatter_from_worker0,
+)

python/tvm/relax/op/ccl/ccl.py

@@ -106,3 +106,25 @@ def scatter_from_worker0(x: Expr, num_workers: int, axis: int = 0) -> Expr:
       Chunked Tensor received by different workers.
     """
     return _ffi_api.scatter_from_worker0(x, num_workers, axis)
+
+
+def gather_to_worker0(x: Expr, num_workers: int, in_group: bool = True) -> Expr:
+    """Gather data from all workers to worker-0.
+
+    Parameters
+    ----------
+    x : relax.Expr
+      The tensor to be gathered from each worker.
+
+    num_workers : int
+      The number of workers to gather data from.
+
+    in_group : bool
+      Whether the gather operation performs globally or in group as default.
+
+    Returns
+    -------
+    result : relax.Expr
+      The concatenated tensor received by worker-0 only.
+    """
+    return _ffi_api.gather_to_worker0(x, num_workers, in_group)

python/tvm/relax/transform/legalize_ops/ccl.py

@@ -123,3 +123,26 @@ def _scatter_from_worker0(_bb: BlockBuilder, call: Call) -> Expr:
             vdevice=call.args[0].ty.vdevice,
         ),
     )
+
+
+@register_legalize("relax.ccl.gather_to_worker0")
+def _gather_to_worker0(_bb: BlockBuilder, call: Call) -> Expr:
+    output_shape = []
+    arg_ty = call.args[0].ty
+    assert isinstance(arg_ty, TensorType), "The input of gather_to_worker0 should be TensorType."
+    assert isinstance(arg_ty.shape.ty, ShapeType)
+    arg_shape = arg_ty.shape.ty
+    for i, shape_value in enumerate(arg_shape.values):
+        if i == 0:
+            output_shape.append(shape_value * call.attrs.num_workers)
+        else:
+            output_shape.append(shape_value)
+    return call_dps_packed(
+        "runtime.disco.gather_to_worker0",
+        [call.args[0], call.attrs.in_group],
+        out_ty=TensorType(
+            shape=output_shape,
+            dtype=arg_ty.dtype,
+            vdevice=arg_ty.vdevice,
+        ),
+    )