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

[Laiyyi]

Summary

Currently, TVM Disco supports multi-node distributed execution through NCCL and RCCL. However, these backends require all participating devices to have GPUs. In some environments where only CPU resources are available, existing collective communication backends cannot be directly used.

Therefore, this PR introduces CPUCCL, enabling TVM Disco to perform multi-node distributed execution in CPU-only environments.

• Added cpuccl.cc to implement the CPU-based collective communication backend.
• Added upload_vm_module to support VM module transfer across nodes.
  • In real distributed environments, only the host node usually contains the compiled .so module.
  • The host node can upload the module to remote nodes through upload_vm_module.
  • Remote nodes can then load the module through load_vm_module and participate in distributed execution.
• Implemented allreduce using a ring-based algorithm. These functions interface follows the implementation style of nccl.cc.
• Added build_ring option to ThreadedSession, ProcessSession,SocketSession and process_pool.py to enable optional ring construction for CPUCCL communication.
• Updated disco_worker.py with corresponding changes to support CPUCCL communication.
• Added gather_to_worker0 support in Relax.

SocketSession Cross-node Communication

In SocketSession, CPUCCL cross-node communication is implemented using local rings combined with TCP forwarding.

Assume two nodes, A and B, each containing two workers. Each node first creates its own local ring:

Node A:
W0 ─A0→ W1 ─A1→ W0

Node B:
W0 ─B0→ W1 ─B1→ W0

Each session additionally creates two threads responsible for cross-node TCP communication.

When Node A's W1 needs to send data to Node B's W0, it only writes data into the local ring channel A1. The send thread monitors A1 and forwards the data to Node B through TCP.

On Node B, SocketSession creates a proxy channel. The recv thread receives data from Node A and writes it into proxy_out_to_tcp_. Since Node B W0's ring_in is rerouted through RerouteRingIn, W0 on node B can correctly receive data from Node A W1.

Bug Fix

• Fixed the garbage collection issue in test_session.py caused by duplicated shutdown messages during session destruction. Added shutdown state checking in bcast_session.cc to prevent repeated shutdown handling.

Build Configuration

• Added USE_CPU_CCL option to allow users to choose whether to compile cpuccl.cc.
• Updated related CMakeLists.txt files to include CPUCCL in the build process.

Tests

• Added test_cpuccl.py, reusing NCCL test cases to validate CPUCCL collective communication.
• Added build-ring and non-build-ring test cases in test_session.py.
• Added test_socket_upload to validate the upload_vm_module functionality.
• Added tests in test_op_ccl.py and test_transform_legalize_ops_ccl.py to verify Relax gather_to_worker0 functionality.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.