wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/runtime/disco/builtin.h
+154
View File
@@ -0,0 +1,154 @@
/*
* 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.
*/
#ifndef TVM_RUNTIME_DISCO_BUILTIN_H_
#define TVM_RUNTIME_DISCO_BUILTIN_H_
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/runtime/tensor.h>
#include <string>
namespace tvm {
namespace runtime {
/*!
* \brief Possible kinds of reduction operations.
*/
enum class ReduceKind : int32_t {
kSum = 0,
kProd = 1,
kMin = 2,
kMax = 3,
kAvg = 4,
};
/*! \brief Converts `ReduceKind` to string */
inline std::string ReduceKind2String(ReduceKind kind) {
switch (kind) {
case ReduceKind::kSum:
return "kSum";
case ReduceKind::kProd:
return "kProd";
case ReduceKind::kMin:
return "kMin";
case ReduceKind::kMax:
return "kMax";
case ReduceKind::kAvg:
return "kAvg";
}
TVM_FFI_THROW(ValueError) << "Unknown ReduceKind: " << static_cast<int>(kind);
}
/*!
* \brief Load a runtime Module, then create and initialize a RelaxVM
* \param path The path to the runtime Module (a DSO file) to be loaded
* \param device The default device used to initialize the RelaxVM
* \return The RelaxVM as a runtime Module
*/
TVM_RUNTIME_DLL ffi::Module LoadVMModule(std::string path, ffi::Optional<Device> device);
/*!
* \brief Create an uninitialized empty Tensor
* \param shape The shape of the Tensor
* \param dtype The dtype of the Tensor
* \param device The device the Tensor is created on. If None, use the thread local default device
* \return The Tensor created
*/
TVM_RUNTIME_DLL Tensor DiscoEmptyTensor(ffi::Shape shape, DLDataType dtype,
ffi::Optional<Device> device);
/*!
* \brief Perform an allreduce operation using the underlying communication library
* \param send The array send to perform allreduce on
* \param reduce_kind The kind of reduction operation (e.g. sum, avg, min, max)
* \param in_group Whether the allreduce operation performs globally or in group as default.
* \param recv The array receives the outcome of allreduce
*/
TVM_RUNTIME_DLL void AllReduce(Tensor send, ReduceKind reduce_kind, bool in_group, Tensor recv);
/*!
* \brief Perform an allgather operation using the underlying communication library
* \param send The array send to perform allgather on
* \param in_group Whether the allgather operation performs globally or in group as default.
* \param recv The array receives the outcome of allgather
*/
TVM_RUNTIME_DLL void AllGather(Tensor send, bool in_group, Tensor recv);
/*!
* \brief Perform a broadcast operation from worker-0
* \param send The buffer to be broadcasted
* \param in_group Whether the broadcast operation performs globally or in group as default.
* \param recv The buffer receives the broadcasted array
*/
TVM_RUNTIME_DLL void BroadcastFromWorker0(Tensor send, bool in_group, Tensor recv);
/*!
* \brief Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
* \param send For worker-0, it must be provided, and otherwise, the buffer must be None.
* The buffer will be divided into equal parts and sent to each worker accordingly.
* \param in_group Whether the scatter operation performs globally or in group as default.
* \param recv The receiving buffer, which must not be None.
*/
TVM_RUNTIME_DLL void ScatterFromWorker0(ffi::Optional<Tensor> send, bool in_group, Tensor recv);
/*!
* \brief Perform a gather operation to worker-0.
* \param send The sending buffer, which must not be None.
* \param in_group Whether the gather operation performs globally or in group as default.
* \param recv For worker-0, it must be provided, and otherwise, the buffer must be None. The
* receiving buffer will be divided into equal parts and receive from each worker accordingly.
*/
TVM_RUNTIME_DLL void GatherToWorker0(Tensor send, bool in_group, ffi::Optional<Tensor> recv);
/*!
* \brief Receive a buffer from worker-0. No-op if the current worker is worker-0.
* \param buffer The buffer to be received
*/
TVM_RUNTIME_DLL void RecvFromWorker0(Tensor buffer);
/*!
* \brief Send a buffer to the corresponding worker in the next group.
* An error is thrown if the worker is already in the last group.
* \param buffer The sending buffer.
*/
TVM_RUNTIME_DLL void SendToNextGroup(Tensor buffer);
/*!
* \brief Receive a buffer from the corresponding worker in the previous group.
* An error is thrown if the worker is already in the first group.
* \param buffer The receiving buffer.
*/
TVM_RUNTIME_DLL void RecvFromPrevGroup(Tensor buffer);
/*!
* \brief Send a buffer to the target receiver worker (globally across all groups).
* \param buffer The sending buffer.
* \param receiver_id The global receiver worker id.
*/
TVM_RUNTIME_DLL void SendToWorker(Tensor buffer, int receiver_id);
/*!
* \brief Receive a buffer from the target sender worker (globally across all groups).
* \param buffer The receiving buffer.
* \param sender_id The global sender worker id.
*/
TVM_RUNTIME_DLL void RecvFromWorker(Tensor buffer, int sender_id);
/*! \brief Get the local worker id */
TVM_RUNTIME_DLL int WorkerId();
/*!
* \brief Called by the worker thread. Waiting until the worker completes all its tasks.
* As a specific example, on a CUDA worker, it blocks until all kernels are launched and
* cudaStreamSynchronize is complete.
*/
TVM_RUNTIME_DLL void SyncWorker();
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_BUILTIN_H_
include/tvm/runtime/disco/cuda_ipc_memory.h
+99
View File
@@ -0,0 +1,99 @@
/*
* 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.
*/
#ifndef TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#include <tvm/runtime/base.h>
#include <tvm/runtime/memory/memory_manager.h>
#include <vector>
namespace tvm {
namespace runtime {
namespace cuda_ipc {
/*!
* \brief The CUDA IPC (interprocess communication) memory object,
* which internally contains data pointers to CUDA IPC memory.
* It is be useful for efficient all-reduce implementation.
* \note Right now the class members are closely tied with customized
* all-reduce kernel. They may also be extended for other uses in
* the future.
*/
class CUDAIPCMemoryObj : public ffi::Object {
public:
/*! \brief The number of GPU workers. */
int num_workers;
/*! \brief The worker id corresponding to this IPC memory object. */
int worker_id;
/*!
* \brief The data pointers of all all-reduce inputs.
* It has "num_workers" pointers. The i-th pointer is the data pointer on worker i.
* If "i != worker_id", the pointer is an IPC data pointer.
* Otherwise, the pointer is a local CUDA data pointer.
*/
std::vector<void*> remote_data;
// We introduce the barrier helper data below per CUDAIPCMemory object
// so that they can be used by custom collective operations and allow
// fine-grained synchronization on each buffer. These barriers have
// low overhead, and can potentially enable concurrent execution of
// kernels in future.
/*!
* \brief The pointers to input barrier signals of all workers for all-reduce.
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
*/
std::vector<void*> barrier_in;
/*!
* \brief The pointers to output barrier signals of all workers for all-reduce.
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
*/
std::vector<void*> barrier_out;
/*! \brief The integer buffer flag for all-reduce. */
int barrier_flag;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("tvm.runtime.disco.cuda_ipc_memory", CUDAIPCMemoryObj, ffi::Object);
};
/*!
* \brief Managed reference to CUDAIPCMemoryObj.
* \sa CUDAIPCMemory
*/
class CUDAIPCMemory : public ffi::ObjectRef {
public:
/*! \brief Get the global singleton CUDAIPCMemory allocator. */
TVM_RUNTIME_DLL static memory::Allocator* GlobalAllocator();
/*!
* \brief Given a local CUDA data pointer, return the CUDAIPCMemory object of the pointer.
* \note The pointer's CUDAIPCMemory is expected to have been allocated
* through global function "cuda_ipc.alloc_storage". Or otherwise this
* function will raise exception.
*/
TVM_RUNTIME_DLL static CUDAIPCMemory GetIPCMemoryFromDevicePtr(void* ptr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CUDAIPCMemory, ffi::ObjectRef, CUDAIPCMemoryObj);
};
} // namespace cuda_ipc
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
include/tvm/runtime/disco/disco_worker.h
+118
View File
@@ -0,0 +1,118 @@
/*
* 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.
*/
/*!
* \file disco_worker.h
* \brief This file defines a worker in Disco. A worker can be launched in a separate thread or
* process as long as the channel supports bi-directional communication in-between the worker and
* the controler.
*/
#ifndef TVM_RUNTIME_DISCO_DISCO_WORKER_H_
#define TVM_RUNTIME_DISCO_DISCO_WORKER_H_
#include <tvm/ffi/function.h>
#include <tvm/runtime/disco/session.h>
#include <vector>
namespace tvm {
namespace runtime {
/*!
* \brief A worker in Disco. It takes a channel to communication with the controler.
* The worker can be run in a separate thread or process as long as the channel supports
* bi-directional communication in-between.
*/
class DiscoWorker {
public:
/*!
* \brief Construct a worker.
* \param worker_id The id of the worker.
* \param num_workers The number of the workers.
* \param num_groups The number of the worker groups.
* \param worker_zero_data The data shared between worker-0 and the controler. It's a nullptr if
* the worker is not worker-0.
* \param channel The communication channel between the worker and the controler.
*/
explicit DiscoWorker(int worker_id, int num_workers, int num_groups,
WorkerZeroData* worker_zero_data, DiscoChannel* channel)
: worker_id(worker_id),
local_worker_id(worker_id),
num_workers(num_workers),
num_groups(num_groups),
default_device(Device{DLDeviceType::kDLCPU, 0}),
worker_zero_data(worker_zero_data),
channel(channel),
register_file{} {}
/*! \brief Main loop of the worker */
void MainLoop();
/*! \brief Get the worker instance on the current thread */
TVM_RUNTIME_DLL static DiscoWorker* ThreadLocal();
/*! \brief Set the specific register to a specific value */
void SetRegister(int reg_id, ffi::AnyView value);
/*! \brief The id of the worker.*/
int worker_id;
/*! \brief The local id of the worker. This can be different from worker_id if the session is
* consisted with multiple sub-sessions. */
int local_worker_id;
/*! \brief Total number of workers */
int num_workers;
/*! \brief Total number of workers */
int num_groups;
/*! \brief The default device to allocate data if not specified */
Device default_device;
/*! \brief The name of the underlying collective communication library. */
ffi::String ccl;
/*!
* \brief The data shared between worker-0 and the controler. It's a nullptr if
* the worker is not worker-0.
* \note This data structure is owned by the controler.
*/
WorkerZeroData* worker_zero_data;
/*!
* \brief The communication channel between the worker and the controler.
* \note This data structure is owned by the controler.
*/
DiscoChannel* channel;
/*! \brief The registers in the worker */
std::vector<ffi::Any> register_file;
struct Impl;
friend struct DiscoWorker::Impl;
};
/*!
* \brief A threadlocal wrapper of DiscoWorker.
*/
struct ThreadLocalDiscoWorker {
/*! \brief The Disco worker */
DiscoWorker* worker;
/*!
* \brief Get the threadlocal Disco worker.
*/
static ThreadLocalDiscoWorker* Get() {
thread_local static ThreadLocalDiscoWorker worker;
return &worker;
}
};
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_DISCO_WORKER_H_
include/tvm/runtime/disco/session.h
+385
View File
@@ -0,0 +1,385 @@
/*
* 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.
*/
/*!
* \file session.h
* \brief This file serves as the entry point of Disco and defines key data structures and
* interfaces.
*
* 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 Tensor 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 ffi::Function, 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 ffi::Function calling convention to support polymorphic and variadic arguments.
*/
#ifndef TVM_RUNTIME_DISCO_SESSION_H_
#define TVM_RUNTIME_DISCO_SESSION_H_
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
#include <mutex>
#include <queue>
#include <string>
#include <utility>
namespace tvm {
namespace runtime {
/*!
* \brief Static FFI type index for `runtime::disco::DRef`.
*
* Allocated within the [kTVMFFIDynObjectBegin - 16, kTVMFFIDynObjectBegin)
* custom-static slot range. The sibling constant `kRuntimeRPCObjectRef`
* lives in `src/runtime/rpc/rpc_session.h` and uses `... - 13`; values must
* remain disjoint across this small reserved block.
*/
constexpr int32_t kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14;
static_assert(kRuntimeDiscoDRef >= TVMFFITypeIndex::kTVMFFIStaticObjectEnd &&
kRuntimeDiscoDRef < TVMFFITypeIndex::kTVMFFIDynObjectBegin,
"kRuntimeDiscoDRef must live in the static custom-index slot range");
/*!
* \brief All possible kinds of Disco commands.
*/
enum class DiscoAction : int32_t {
kShutDown = 0,
kKillReg = 1,
kGetGlobalFunc = 2,
kCallPacked = 3,
kSyncWorker = 4,
kCopyFromWorker0 = 5,
kCopyToWorker0 = 6,
kDebugGetFromRemote = 7,
kDebugSetRegister = 8,
};
/*! \brief Converts the enum class `DiscoAction` to string */
inline std::string DiscoAction2String(DiscoAction action) {
switch (action) {
case DiscoAction::kShutDown:
return "kShutDown";
case DiscoAction::kKillReg:
return "kKillReg";
case DiscoAction::kGetGlobalFunc:
return "kGetGlobalFunc";
case DiscoAction::kCallPacked:
return "kCallPacked";
case DiscoAction::kSyncWorker:
return "kSyncWorker";
case DiscoAction::kCopyFromWorker0:
return "kCopyFromWorker0";
case DiscoAction::kCopyToWorker0:
return "kCopyToWorker0";
case DiscoAction::kDebugGetFromRemote:
return "kDebugGetFromRemote";
case DiscoAction::kDebugSetRegister:
return "kDebugSetRegister";
}
TVM_FFI_THROW(ValueError) << "Unknown DiscoAction: " << static_cast<int>(action);
}
class SessionObj;
/*!
* \brief An object that exists on all workers.
*
* The controler assigns a unique "register id" to each object, and the worker uses this id to
* refer to the object residing on itself.
*/
class DRefObj : public ffi::Object {
public:
/*!\ brief Send dellocation command for `reg_id` */
inline ~DRefObj();
/*!
* \brief Get the value of a DRef from a remote worker.
* \param worker_id The id of the worker to be fetched from.
* \return The value of the register.
*/
inline ffi::Any DebugGetFromRemote(int worker_id);
/*!
* \brief Copy from the Tensor provided to a remote worker.
* \param worker_id The id of the worker to be copied to.
* \param source The Tensor to be copied.
*/
inline void DebugCopyFrom(int worker_id, ffi::AnyView source);
static constexpr const uint32_t _type_index = kRuntimeDiscoDRef;
static const constexpr bool _type_final = true;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_STATIC("runtime.disco.DRef", DRefObj, ffi::Object);
/*! \brief The id of the register */
int64_t reg_id;
/*! \brief Back-pointer to the host controler session */
ffi::ObjectRef session{nullptr};
private:
inline SessionObj* GetSession();
};
/*!
* \brief Managed reference to DRefObj.
* \sa DRefObj
* \note No public constructor is provided as it is not supposed to be directly created by users.
*/
class DRef : public ffi::ObjectRef {
public:
explicit DRef(ffi::ObjectPtr<DRefObj> data) : ffi::ObjectRef(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DRef, ffi::ObjectRef, DRefObj);
};
/*!
* \brief A Disco interactive session. It allows users to interact with the Disco command queue with
* various ffi::Function calling convention.
*/
class SessionObj : public ffi::Object {
public:
virtual ~SessionObj() = default;
/*!
* \brief Call a ffi::Function on workers providing variadic arguments.
* \tparam Args In the variadic arguments, the supported types include:
* - integers and floating point numbers;
* - DataType;
* - Device;
* - std::string;
* - DRef.
* Examples of unsupported types:
* - Tensor, DLTensor;
* - TVM Objects, including ffi::Function, Module and String;
* \param func The function to be called.
* \param args The variadic arguments.
* \return The return value of function call
*/
template <typename... Args>
TVM_FFI_INLINE DRef CallPacked(const DRef& func, Args&&... args);
/*!
* \brief Call packed function on each worker using a packed sequence. The calling convention:
* The first element must be DiscoAction::kCallPacked,
* The second element must be 0, which will later be updated by the session to return reg_id
* The thirtd element is the function to be called.
*/
TVM_RUNTIME_DLL virtual DRef CallWithPacked(const ffi::PackedArgs& args) = 0;
/*! \brief Get the number of workers in the session. */
TVM_RUNTIME_DLL virtual int64_t GetNumWorkers() = 0;
/*! \brief Get a global functions on workers. */
TVM_RUNTIME_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
/*!
* \brief Copy an Tensor from worker-0 to the controler-side Tensor
* \param host_array The array to be copied to worker-0
* \param remote_array The Tensor on worker-0
*/
TVM_RUNTIME_DLL virtual void CopyFromWorker0(const Tensor& host_array,
const DRef& remote_array) = 0;
/*!
* \brief Copy the controler-side Tensor to worker-0
* \param host_array The array to be copied to worker-0
* \param remote_array The Tensor on worker-0
*/
TVM_RUNTIME_DLL virtual void CopyToWorker0(const Tensor& host_array,
const DRef& remote_array) = 0;
/*!
* \brief Synchrnoize the controler with a worker, and it will wait until worker finishes
* executing this instruction.
* \param worker_id The id of the worker to be synced with.
* \note This function is usually used for worker-0, because it is the only worker that is
* assumed to collocate with the controler. Syncing with other workers may not be supported.
*/
TVM_RUNTIME_DLL virtual void SyncWorker(int worker_id) = 0;
/*! \brief Signal all the workers to shutdown */
TVM_RUNTIME_DLL virtual void Shutdown() = 0;
/*!
* \brief Initialize the data plane between workers.
* \param ccl The name of the communication backend, e.g., nccl, rccl, mpi.
* \param device_ids The device ids of the workers.
*/
TVM_RUNTIME_DLL virtual void InitCCL(ffi::String ccl, ffi::Shape device_ids) = 0;
/*!
* \brief Get the value of a register from a remote worker.
* \param reg_id The id of the register to be fetched.
* \param worker_id The id of the worker to be fetched from.
* \return The value of the register.
*/
TVM_RUNTIME_DLL virtual ffi::Any DebugGetFromRemote(int64_t reg_id, int worker_id) = 0;
/*!
* \brief Set the value of a register on a remote worker.
* \param reg_id The id of the register to be set.
* \param value The value to be set.
* \param worker_id The id of the worker to be set.
*/
TVM_RUNTIME_DLL virtual void DebugSetRegister(int64_t reg_id, ffi::AnyView value,
int worker_id) = 0;
struct FFI;
friend struct SessionObj::FFI;
friend class DRefObj;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("runtime.disco.Session", SessionObj, ffi::Object);
protected:
/*! \brief Deallocate a register id, kill it on all workers, and append it to `free_regs_`. */
virtual void DeallocReg(int reg_id) = 0;
};
/*!
* \brief Managed reference to SessionObj
* \sa SessionObj
*/
class Session : public ffi::ObjectRef {
public:
/*!
* \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.
*/
TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups);
/*!
* \brief Create a session backed by pipe-based multiprocessing
* \param num_workers The number of workers.
* \param num_groups The number of worker groups.
* \param process_pool_creator The name of a global function that takes `num_workers` as an input,
* and returns a ffi::Function, which takes an integer `worker_id` as the input and returns None.
* 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.
* \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);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj);
};
/*!
* \brief A bi-directional channel for controler-worker communication.
* This channel is primarily used to transfer control messages but not data.
*/
class DiscoChannel {
public:
virtual ~DiscoChannel() = default;
/*! \brief Send a packed sequence to the receiver */
virtual void Send(const ffi::PackedArgs& args) = 0;
/*! \brief Receive a packed sequence from worker */
virtual ffi::PackedArgs Recv() = 0;
/*! \brief Reply a packed sequence to the sender */
virtual void Reply(const ffi::PackedArgs& args) = 0;
/*! \brief Receive a reply from the worker */
virtual ffi::PackedArgs RecvReply() = 0;
};
/*!
* \brief A special communication channel between controler and worker-0,
* assuming they are always collocated in the same process.
*/
class WorkerZeroData {
public:
/*!
* \brief The host-side arrays to passed to worker-0 for special uses, for example,
* copy-to-worker0 and copy-from-worker0
*/
std::queue<Tensor> host_arrays;
/*! \brief The mutex that guards `host_arrays` */
std::mutex queue_mutex_;
};
// Implementation details
inline SessionObj* DRefObj::GetSession() {
return const_cast<SessionObj*>(static_cast<const SessionObj*>(session.get()));
}
DRefObj::~DRefObj() {
if (this->session.defined()) {
GetSession()->DeallocReg(reg_id);
}
}
ffi::Any DRefObj::DebugGetFromRemote(int worker_id) {
return GetSession()->DebugGetFromRemote(this->reg_id, worker_id);
}
void DRefObj::DebugCopyFrom(int worker_id, ffi::AnyView value) {
return GetSession()->DebugSetRegister(this->reg_id, value, worker_id);
}
template <typename... Args>
DRef SessionObj::CallPacked(const DRef& func, Args&&... args) {
constexpr int offset = 3;
constexpr int kNumArgs = offset + sizeof...(Args);
ffi::AnyView packed_args[kNumArgs];
ffi::PackedArgs::Fill(packed_args,
/*.0=*/static_cast<int>(DiscoAction::kCallPacked), // action
/*.1=*/0, // reg_id, which will be updated by this->CallWithPacked
/*.2=*/func, // the function to be called
std::forward<Args>(args)...);
return this->CallWithPacked(ffi::PackedArgs(packed_args, kNumArgs));
}
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_DISCO_SESSION_H_