The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`,
so the Disco Python wrappers can no longer store instance attributes and
every session construction fails with AttributeError. Declare the
attributes each
wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu`
so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the
disco runtime being present so they skip cleanly on builds (e.g. the pip
wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it.
### Session attribute storage
- `DPackedFunc` / `DModule`: `__slots__ = ("session",)`.
- `Session`: `__slots__ = ("_cache", "_import_python_module")`
## Summary
This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.
The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.
## Main Changes
- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.
## Validation
- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
- `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
- `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
- `664 passed, 25 skipped, 9 xfailed, 1 xpassed`
## Local CI Notes
Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
## Summary
TVM-side cleanup that drops the `python/tvm/runtime/object.py` shim and
routes `tvm.runtime.Object` directly to `tvm_ffi.Object`. The
`tvm.runtime.Object` re-export is preserved (now a re-export of
`tvm_ffi.Object`) so external callers keep working.
The load-bearing `__object_repr__` install — which wires TVM IR objects
up to the rich C++ `ReprPrinter` registered through
`init_ffi_api("node", ...)` — moves into
`python/tvm/runtime/_ffi_node_api.py`.
That module is already imported as a side-effect-only module from
`python/tvm/runtime/__init__.py`, so the override fires at the right
time (after `init_ffi_api` registers the C++ printer).
`_ffi_node_api.AsRepr` itself is **kept**: `tvm_ffi`'s default repr is
primitive (`ClassName(ptr)`); TVM IR objects need the rich printer
registered via `init_ffi_api("node", ...)`. `AsRepr` is what bridges
that printer back into Python `repr(obj)` and is also the runtime-only
fallback when `libtvm.so` is unavailable.
The 7 in-tree importers of the deleted shim (plus one straggler in
`runtime/disco/session.py`) are switched to either
`from tvm.runtime import Object` or `from tvm_ffi import Object`,
depending on which pattern the file already uses.
## Test plan
- [x] `python -c "import tvm; print(repr(tvm.IRModule({})))"` produces
TVMScript-style output (rich repr preserved).
- [x] `pytest tests/python/all-platform-minimal-test/ -x` — 75 passed,
77 skipped (matches baseline).
- [x] `pytest tests/python/tirx-base/ -x` — 273 passed, 2 skipped.
- [x] `pre-commit run --files <changed files>` — all hooks pass.
- [ ] CI green.
This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
This PR cleans up the python API to make things more consistent
with existing python array api and torch.
Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device
API updates:
- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func
This PR Updates the NDArray => Tensor.
Both tensor and ndarray are commonly used terms.
Because the term Tensor is getting more common in the context of ML,
we do the rename to stay more aligned with torch.Tensor and DLTensor.
* [FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This PR establishes tvm_ffi as a standalone python module.
The ffi is structured as a minimal pip module that can be
directly install by path or url.
examples/get_started provided a minimal example.
This is a major change as we are decoupling tvm_ffi as a
separate package, users need to install tvm_ffi separately.
Thanks to its minimal dependency, tvm_ffi can be easily installed
even just from the source by pip install ./ffi
This change would enable future improvement for library plugins
to have lightweight dependencies by just working on top of
the tvm_ffi, while the main compiler toolchain and runtime
can be layered on top.
* [FFI] Improve traceback setups
This PR improves traceback related setups
This PR formalizes original runtime::Module into ffi
as ffi.Module and cleans the APIs around it.
The goal is to stablize the Module API as extra API that can benefit the overall
ffi interactions. We also refactors the c++ code that depends on the Module.
This PR phases out tvm._ffi redirections in favor of new FFI
new functions are now called via tvm.ffi.
We also enabled limited API support for python 3.12+
so the compiled binary can be forward compatible to future
python versions.
This PR cleans up the container redirections and headers
so the files directly points to new ones in ffi folder
- runtime/shape_tuple.h => ffi/container/shape.h
- for IntTuple alias, introduce runtime/int_tuple.h for now
- runtime/container/array.h => ffi/container/array.h
- runtime/container/map.h => ffi/container/map.h
- runtime/container/optional.h => ffi/optional.h
- runtime/container/string.h => ffi/string.h
- runtime/container/variant.h => ffi/container/variant.h
- runtime/container/tuple.h => ffi/container/tuple.h
We also introduce limited number of tvm::ffi classes into tvm namespace,
when they are commonly used and their is no ambiguity.
This PR modernizes the FFI foundation of the project and introduce
a new minimal and lightweight module [tvm ffi](https://github.com/apache/tvm/tree/refactor-s3/ffi)
based on our lessons in the past few years. It implements a modern
version of the [Unified Packed and Object RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0097-unify-packed-and-object.md)
that unifies the packed function call and object systems.
Summary of the change:
- A dedicated clean Any/AnyView that can store strong and weak
references of items
- Function(previously PackedFunc) system built on top of the Any/AnyView
- A minimal C API that backs the overall calls. We are stabilizing the
API with a goal to bring clean, stable FFI conventions for both compiled
and registered code
- A rewrite of core python binding and generated code based on the module
- Update existing code and test cases to the new module
- Latest dlpack support
The new module brings many benefits thanks to the cleaner design,
to name a few:
- Any can support both POD types(int) and object types.
- Containers (e.g. Array) can now also contain Any value, e.g. now
`Array<int>` is supported, no need for boxed types
- Error handling now upgrades to object-based, allowing cleaner
traceback across languages
- Map now preserves insertion orders
- Path toward isolated stabilize minimum core ABI/API foundation module
- Type traits based design that cleanly defines how values interact
with Any system
- Automatic conversion of different types based on traits if needed
Because FFI upgrade is at heart of the project, the change touches every
component of the system. Importantly, this is an upgrade of the ABI so the
change is not backward compatible. The code compiled under the old
FFI won't work under the new one. We did provide example ABI translation
(e.g. LegacyTVMArgValueToFFIAny) functions for compatibility.
The PR tries to leave files in their old places while creating redirections.
The goal is to have the first milestone landed and infrastructure in place,
so we can do further refactors to complete features and cleanup legacy code
as trackable PRs. As of now, python binding and compiled code are under the
new convention while RPC and some other bindings still relies on legacy ABI
translation. We will work on upgrades in the coming PRs, including areas such
as reflection, phasing out legacy redirections etc.
* [Cython][FFI] Fix crash when call del operator for handle
In case of cython when we create a set function for property then the
following code will be generated:
```
static int __pyx_setprop_4test_9TestClass_handle(PyObject *o, PyObject *v, CYTHON_UNUSED void *x) {
if (v) {
return __pyx_pw_4test_9TestClass_6handle_3__set__(o, v);
}
else {
PyErr_SetString(PyExc_NotImplementedError, "__del__");
return -1;
}
}
```
And when we call operator `del` for this handler, then the memory will
be released and operator `__set__` will be called for NULL object. In
this case an exception that operator `__del__` is not implemented will
be generated. To avoid this problem we need to declare `__del__`
function for each property where we define operator `__set__`.
* Apply comments
* Set dref.handle to None instead of using __del__ functions
* [Disco][QoL] Implement broadcast/scatter methods for Session
Prior to this commit, use of the `disco.Session` API to broadcast or
scatter an array required several steps from the caller.
1. Allocate memory on worker0
2. Transfer data from the controller to worker0
3. Allocate memory on each worker
4. Broadcast/scatter data from worker0 to all workers
While exposing these steps is necessary for performance, especially
when used repeatedly, it can be tedious/error-prone to use for
initialization that is only performed once.
This commit adds utility methods `Session.broadcast` and
`Session.scatter`, which are implemented in terms of the existing
lower-level methods `Session.broadcast_from_worker0` and
`Session.scatter_from_worker0`. These methods perform the transfer
from the controller to worker0, and from worker0 to all other
workers.
* lint fix
Prior to this commit, while the `num_workers` argument was provided to
the `disco.Session` object, it could not be determined from an
existing `disco.Session` object. As a result, functions that
interacted with a multi-GPU setup frequently required separate
`num_workers` and `disco_session` argument, which could erroneously be
out-of-sync (e.g. passing the incorrect `num_workers`, or omitting the
`disco_session` argument when `num_workers>1`).
To remove this class of errors, this commit adds a
`disco.Session.num_workers` property. The separate `num_workers`
argument is no longer necessary, as it can be determined from the
`disco.Session` instance.
The `disco.Session.scatter_from_worker0` function expects a `DRef`
which an `NDArray` on worker 0, and `NullOpt` on all other workers.
Prior to this commit, there was no method in the `disco.Session` that
could be used to make such a `DRef`. As a result, every use of
`scatter_from_worker0` generated an error, stating that non-zero
workers should have `NullOpt` as their `send` argument.
This commit adds a `worker0_only: bool` argument to
`disco.Session.empty`. This can be used to generate an allocation
that only exists on worker zero, suitable for use in
`scatter_from_worker0`.
Prior to this commit, the `SessionObj::Shutdown` method could be
called from the C++ API, but could not be called through the Python
API. While it is implicitly called when the `SessionObj` is
destructed, Python's garbage collection may result in the destruction
occurring later than expected.
This commit exposes `SessionObj::Shutdown` through the Python API as
`disco.Session.shutdown`, allowing it to be closed cleanly.
This is a follow-up to #16618, which propagates the `structlog`
configuration to disco worker processes. For configurations that
use `structlog.stdlib` to integrate `structlog` with the stdlib
`logging` module, this integration must also be forwarded.
This PR introduces the CUDA IPC memory support in TVM runtime.
IPC memory allows multiple distribtued workers accessing the GPU
memory of each other directly. This functionality is helpful for
implementing customzied communication primitives across distributed
workers.
In this PR, we bring the customized all-reduce implementation
from TensorRT-LLM into 3rdparty. This all-reduce implementation
makes use of the CUDA IPC memory. We expose the all-reduce function
in global function under namespace `tvm::runtime::disco::cuda_ipc`.
One unit test for the customized all-reduce kernel over two workers
is added.
---
Co-authored-by: Hongyi Jin <hongyij@andrew.cmu.edu>
Prior to this PR, module "psutil" is imported at the top level
of the disco process pool. The pool will try to kill all the processes
at the time of destruction (when `__del__` is implicitly invoked).
The `__del__` function eventually calls into a function that
uses `pstuil`. But it is possible that the top-level `psutil`
has already been released by Python, which leads to a KeyError
as follows:
```
Exception ignored in: <function DiscoPopenWorker.__del__ at 0x7f2c922bfe20>
Traceback (most recent call last):
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 67, in __del__
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 81, in kill
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 162, in _kill_child_processes
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 323, in __init__
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 353, in _init
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_pslinux.py", line 1738, in __init__
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_common.py", line 864, in get_procfs_path
KeyError: 'psutil'
```
This PR fixes the issue by lazily importing `psutil` when needed.
Prior to this commit, while `structlog.configure(...)` would only
impact log statements generated in the main process. Any workers
started with `disco.session.ProcessSession` do not inherit the
`structlog` configuration. While `disco.session.ThreadedSession` would
inherit the `structlog` configuration, it would also inherit
process-specific CUDA variables.
This commit updates `disco.session.ProcessSession` to explicitly
propagate any `structlog` configuration to child processes. This
implementation intentionally avoids introducing a new dependency for
TVM. If the `structlog` package is not available, the config
propagation is skipped.
In addition to the PackedFunc `"runtime.disco.worker_id"`, which
returns the worker ID wrapped in a `ShapeTuple`, this commit adds
`"runtime.disco.worker_rank"`, which returns the worker ID without
wrapping, and `"runtime.disco.device"`, which returns the device for
each worker.
The unit test added in this commit simulates loading of model weights
through a parameter transformation function.
This PR changes how to get the `Session` from `DPackedFunc` or `DModule`. Without this, there will be data corruption in the message channel, if Python gc happens in a different thread than the thread that owns the Disco session, which is typical in a multi-thread environment like LLM inference server.
An explanation on how such data corruption could occur:
1. Every time dref.session gets called. A new Session Python object will be created, due to how FFI works https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/python/tvm/runtime/disco/session.py#L42
2. When we want to call the model through disco, we need to first obtain the DPackedFunc through DModule . This process calls the _get_cached_method on a new Session object (follows (1)) https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/python/tvm/runtime/disco/session.py#L97
3. In hasattr(self, "_cache") , a call is made to the tvm::ReflectionVTable::GetAttr to check if it exists in the TVM object https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/python/tvm/runtime/disco/session.py#L107
4. This call will fail and raise error, which gets caught by https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/python/tvm/_ffi/_ctypes/packed_func.py#L239. hasattr will swallow this error and return False.
5. During the creation of the FFI error in (4), traceback objects are created to help trace into C++ function. https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/python/tvm/_ffi/base.py#L373 Holding a reference to the frame object creates circular references that goes 'frame -> parent frames -> frame that holds the traceback object -> traceback -> frame'. All local variables in the call stack will be indirectly referenced and cannot be freed by ref counting. This is okay because Python gc will collect these objects once they are unreachable.
6. (5) prevents the DPackedFunc returned from _get_cached_method to be freed through Python ref counting, because it's in the frame that's indirectly referred by the circular reference in (5).
7. If Python gc happens in the main thread, the DPackedFunc from (6) will be collected, because there is no references to the newly-created Session, other than the reference path rooted from the circular reference between traceback and frame (from (5)).
8. (7) will result in writing to the pipe in the main thread from the DRef destructor. https://github.com/apache/tvm/blob/ebbe38f3281776cfda4fce0b188892ab1c5c7572/include/tvm/runtime/disco/session.h#L301 If other disco activity happens in another thread, there will be multiple threads writing to the same pipe without synchronization. Data will be corrupted.
This PR fixes (1) and (3), so that the data corruption will not occur anymore.
Removed instances of accidentally repeated words from comments. There
are cases where duplicated words appear legitimately, those cases remain
unmodified.
In our previous implementation, parameter sharding relies on pre-quantization weight processing,
meaning each set of quantized weights corresponds strictly to a hardcoded constant `num_shards`,
and re-quantization is strictly required upon each change of #GPUs, e.g. from 4-GPU to 8-GPU
setting. This PR makes it possible to move parameter sharding to post-quantization loading-time.
During loading, we iterate over all parameters and apply the sharding operation based on the
provided sharding information.
To make this happen, this PR makes an enhancement to the existing `shard_info.json` to include the
sharding function being used at loading time. Each parameter is attached to a list of loading-time
preprocessing methods that are serially applied to it to transform this parameter to the desired
shape, as shown in the example below:
```python
shard_info = {
"x_0": [ # name of the parameter
[ # a list of preprocessing functions to be applied
"tests.disco.shard_dim_1", # name of the sharding function
[(num_shards, 64, 64), "float16"], # output shape/dtype of `tests.disco.shard_dim_1`
num_shards, # extra inputs to `tests.disco.shard_dim_1`
],
],
"x_1": [...],
}
```
To parameter `x_0`, it means we will call method `tests.disco.shard_dim_1` which has the signature:
```python
def shard_dim_1(
input: NDArray,
num_shards, # extra inputs
output: NDArray, # and its shape is (num_shards, 64, 64), and dtype is "float16"
) -> None: ...
```
This approach simplifies parameter sharding for users and ensures correctness.
* [Disco] Add AllGather
* update to allgather to be compatiple with #15766
* lint fix
* add num workers to all gather
* change num_gpus to prim value instead of int
* fix typo
* remove all_gather attrs and minor improvment
* Update allgather doc
Co-authored-by: Lesheng Jin <34279105+LeshengJin@users.noreply.github.com>
---------
Co-authored-by: Lesheng Jin <34279105+LeshengJin@users.noreply.github.com>
This PR introduces `ProcessSession`, a new session implementation based
on multi-processing.
`ProcessSession` shares exactly the same communication protocol with
`ThreadedSession`, but all workers except for worker 0 are launched in a
separate process than thread. Workers communicate with the controller
via pipe provided by the OS, rather than SPSC message queue between
threads.
In our implementation, Python's `subproces.popen` is used to create
subprocesses, and the Python executable, or more specifically,
`sys.executable` calls into `tvm.exec.disco_worker` as the entrypoint.
Besides the launching logic that is only executed once in the very
beginning, the rest of the implementation resides in a C++-only
environment, including reads/writes to pipe file descriptors,
serialization and deserialization of messages, worker interpretation of
each message, etc.
Detailed engineering elements included in this PR:
- Refactors the MinRPC-based communication protocol out to be shared by
`ProcessSession` and `ThreadedSession` as `protocol.h`;
- Refactors a controller-side worker thread into `DiscoWorkerThread`,
which is shared by both session implementation to launch worker-0;
- Added two instructions `kDebugGetFromRemote` and `kDebugSetRegister`,
which are used to communicate with workers other than worker-0 in
debug mode;
- Introduces multi-processing infra including: `tvm.exec.disco_worker`
serving as the entrypoint that launches workers, and
`tvm/runtime/disco/process_pool.py` that exposes APIs to launch worker
processes. `tvm.exec.disco_worker` calls into a global function
`runtime.disco.WorkerProcess` that executes the worker main loop in
pure C++;
- Introduces `src/support/process_id.h` that provides cross-platform pid
and tid printing utilities;
- Refactors Disco's NCCL integration that get rids of initialized-once
global NCCL context, and switches to broadcasting `ncclUniqueId` from
controller to all workers, and then create NCCL communicators in each
worker thread/process accordingly. This is a thread/process-agnostic
way of using NCCL.
* [Runtime] ShapeTuple.Product and ShapeTuple Printing (#15652)
This PR adds two convenient methods for `ShapeTuple`.
```C++
// Returns the number of elements in the shape,
// i.e. the product of all dimensions
ShapeTupleObj::index_type ShapeTupleObj::Product();
// Printing method for shape
std::ostream& operator<<(std::ostream& os, const ShapeTuple& shape);
```
* [Disco] Add `Scatter-From-Worker0`
This PR introduces an API:
```C++
void ScatterFromWorker0(Optional<NDArray> send, NDArray recv);
```
This mimics the MPI-style scattering primitive that divides NDArray
`send` into chunks, and send them to each worker's NDArray `recv`.
* [CMake] Add NCCL to TVM and TVM Runtime (#15605)
This PR introduces NCCL in the cmake system.
NCCL is NVIDIA's library for distributed communication.
* [Runtime] Expose ModuleGetFunction as PackedFunc (#15623)
This PR exposes `Module.GetFunction` as a global PackedFunc.
Previously, the only way to access this method is via TVM's
C API, but the C++ PackedFunc API is missing. This PR patches
this issue.
* [Runtime] Utils to Stringify Device (#15630)
There exist some basic functionality to convert Device and DLDeviceType
to std::string, but they are not following the common naming convention
in TVM, and thus less discoverable. This commit makes changes
accordingly:
- `runtime::DeviceName` to `runtime::DLDeviceType2Str`
- move declaration of `operator << (std::ostream&, Device)` from
`runtime/device_api.h` to `runtime/packed_func.h`
* [RPC] Enhance RPC Protocol to support TVM Object (#15631)
This PR introduces object support in TVM RPC protocol by introducing three
new interfaces in `rpc_reference.h`:
- `uint64_t GetObjectBytes(Object* obj)`, which is a required
implementation that returns the length of the object during serialization;
- `void WriteObject(Object* obj)` used to serialize an object to a
writable channel;
- `void ReadObject(int* type_code, TVMValue* value)`, which deserializes
a TVM Object from a channel.
To serialize an object, a recommended paradigm is to write its
`type_index` first, and then its content. For example, `ShapeTuple` can
be serialized as:
```C++
// pseudocode
void WriteObject(Object* obj) {
if (obj is ShapeTuple) {
this->Write<uint32_t>(type_index of ShapeTuple);
this->Write<int32_t>(obj->ndim);
this->WriteArray<int64_t>(obj->shape);
} else {
throw Unsupported;
}
}
uint64_t GetObjectBytes(Object* obj) {
uint64_t result = 0;
if (obj is ShapeTuple) {
result += sizeof(uint32_t); # for `type_index`
result += sizeof(int32_t); # for `ndim`
result += sizeof(int64_t) * obj->ndim; # for content of the shape
} else {
throw Unsupported;
}
return result;
}
```
To deserialize an object, similar to serialization, the recommended
approach paradigm is to read `type_index` and disptch based on it.
Caveat on deserialization: RPC Reference itself does not own or allocate
any memory to store objects, meaning extra logic is usually required in
`ReadObject` to keep their liveness.
* [Unity] Disco: A Framework-Agnostic SPMD Runtime for Distributed Inference/Training
Disco is a distributed runtime that consists of a controler and a cluster of workers. The
controler is responsible for managing the workers by broadcasting commands to all the workers
together, and the workers are responsible for executing the commands and. The controler and
workers communicate with each other through a bi-directional channel.
Different from a generic system, Disco is designed to as "single-program-multiple-data" (SPMD)
runtime, which means that all the workers execute the same instruction at the same time, but the
data they are working on may be different. For example, in data parallelism, each worker may
work on a different batches of the data, but they all execute the same set of instructions.
Therefore, imagine there is a virtual machine that executes the program, the structures of
workers' register files could be considered as "identical" (single program) although the values
may differ (multiple data).
**DRef.** Following the design above, consider the program in SPMD in a virtual ISA, then each
worker is a virtual machine instance to execute the ISA maintaining its own register file.
The controler denotes each of their register files with a unique integer "register id",
and the workers use this id to refer to the register file that resides on itself.
DRef is a control-side object backed by such a register id. The data it contains is not assumed
to be directly accessible by the controler, with an exception for worker-0, which is a special
worker that is always co-located with the controler.
**Worker-0.** Worker-0 is a special worker that is always co-located with the controler.
It is assumed that the controler can synchronize with and access the registers of worker-0.
The Disco session provides multiple APIs to interact specifically with the worker-0.
To shared data with other workers, a common paradigm in Disco is to copy data from the
controler-side NDArray to the worker-0, and then copy it to other workers using primitives on
the data plane, for example, `broadcast` and `send`.
**Control plane.** The controler broadcasts commands to all the workers as control signals.
For example, the control may ask all workers to load a library or call a function respectively.
Common control signals include: shutdown, retrievel a global PackedFunc, call packed function,
etc. The controler is assumed to keep a message channel to each worker to implement the broadcast
behavior, and the message channel may vary depends on usecases.
**Data plane.** The data channel is usually used to exchange data between workers, especially for
tensor data which is usually large. For example, performing an allreduce operator for sharded
matrix multiplication, or broadcasting for an input tensor. For efficiency, the data channel is
usually backed by NCCL on NVIDIA GPUs, RCCL on AMD GPUs, or MPI on CPUs.
**Session.** A Disco session is a primary interface to interact with the Disco runtime, serving
as a global context that manages the control and workers. It could be implemented as a
multi-threaded with a pool of workers for single-node multi-gpu scenarios, or TCP sockets for
workloads that span over a cluster of nodes.
**Channel.** Disco channel is a bi-directional communication channel between the controler and
workers for exchanging control signals. It is no different from a generic RPC channel, but
adopts TVM's PackedFunc calling convention to support polymorphic and variadic arguments.
Co-Authored-by: Lesheng Jin <34279105+LeshengJin@users.noreply.github.com>
---------
Co-authored-by: Lesheng Jin <34279105+LeshengJin@users.noreply.github.com>
Shushi Hong