## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
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.
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
This PR renames tir.Block to SBlock. This clearly indicate the
scheduable property of the block and is a prereq for followup stir
passes refactor.
Main changes:
- Data structure change from Block to SBlock
- Syntax change from T.block to T.sblock
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.
This PR hides StringObj/BytesObj into details and bring
implementations to directly focus on the String/Bytes.
This change will prepare us for future changes such as SmallStr support.
Also moves more ObjectRef into Any in RPC.
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.
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.
ShapeTuple is an essential item used in Relax stack as the runtime
representation of shapes. It is also used as a current workaround to
represent integers (1-d shape) given standalone non-constant integers
are currently absent in Relax.
This PR introduces formal support for ShapeTuple in Disco's
communication protocol based on the recent enhancement of TVM RPC system
to support TVM Objects: https://github.com/apache/tvm/pull/15631.
* [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>
Tianqi Chen