Namespaces
namespace	cuda_ipc

namespace	vm

Classes
class	DeviceAPI
	TVM Runtime Device API, abstracts the device specific interface for memory management. More...

class	DiscoChannel
	A bi-directional channel for controler-worker communication. This channel is primarily used to transfer control messages but not data. More...

class	DiscoWorker
	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. More...

class	DRef
	Managed reference to DRefObj. More...

class	DRefObj
	An object that exists on all workers. More...

class	Session
	Managed reference to SessionObj. More...

class	SessionObj
	A Disco interactive session. It allows users to interact with the Disco command queue with various ffi::Function calling convention. More...

class	TensorMemoryPool
	A cache of owning device tensors. More...

struct	ThreadLocalDiscoWorker
	A threadlocal wrapper of DiscoWorker. More...

class	Timer
	Timer for a specific device. More...

class	TimerObj
	Base class for all timer implementations. More...

class	WorkerZeroData
	A special communication channel between controler and worker-0, assuming they are always collocated in the same process. More...

Enumerations
enum	TVMDeviceExtType { TVMDeviceExtType_End = 36 }
	Extension device types in TVM. More...

enum class	ReduceKind : int32_t { kSum = 0 , kProd = 1 , kMin = 2 , kMax = 3 , kAvg = 4 }
	Possible kinds of reduction operations. More...

enum class	DiscoAction : int32_t { kShutDown = 0 , kKillReg = 1 , kGetGlobalFunc = 2 , kCallPacked = 3 , kSyncWorker = 4 , kCopyFromWorker0 = 5 , kCopyToWorker0 = 6 , kDebugGetFromRemote = 7 , kDebugSetRegister = 8 }
	All possible kinds of Disco commands. More...

Functions
const char *	DLDeviceType2Str (int type)
	The name of DLDeviceType.

std::ostream &	operator<< (std::ostream &os, DLDevice dev)

TVM_RUNTIME_DLL bool	RuntimeEnabled (const ffi::String &target)
	Check if runtime module is enabled for target.

std::string	ReduceKind2String (ReduceKind kind)
	Converts `ReduceKind` to string.

TVM_RUNTIME_DLL ffi::Module	LoadVMModule (std::string path, ffi::Optional< Device > device)
	Load a runtime Module, then create and initialize a RelaxVM.

TVM_RUNTIME_DLL ffi::Tensor	DiscoEmptyTensor (ffi::Shape shape, DLDataType dtype, ffi::Optional< Device > device)
	Create an uninitialized empty Tensor.

TVM_RUNTIME_DLL void	AllReduce (ffi::Tensor send, ReduceKind reduce_kind, bool in_group, ffi::Tensor recv)
	Perform an allreduce operation using the underlying communication library.

TVM_RUNTIME_DLL void	AllGather (ffi::Tensor send, bool in_group, ffi::Tensor recv)
	Perform an allgather operation using the underlying communication library.

TVM_RUNTIME_DLL void	BroadcastFromWorker0 (ffi::Tensor send, bool in_group, ffi::Tensor recv)
	Perform a broadcast operation from worker-0.

TVM_RUNTIME_DLL void	ScatterFromWorker0 (ffi::Optional< ffi::Tensor > send, bool in_group, ffi::Tensor recv)
	Perform a scatter operation from worker-0, chunking the given buffer into equal parts.

TVM_RUNTIME_DLL void	GatherToWorker0 (ffi::Tensor send, bool in_group, ffi::Optional< ffi::Tensor > recv)
	Perform a gather operation to worker-0.

TVM_RUNTIME_DLL void	RecvFromWorker0 (ffi::Tensor buffer)
	Receive a buffer from worker-0. No-op if the current worker is worker-0.

TVM_RUNTIME_DLL void	SendToNextGroup (ffi::Tensor buffer)
	Send a buffer to the corresponding worker in the next group. An error is thrown if the worker is already in the last group.

TVM_RUNTIME_DLL void	RecvFromPrevGroup (ffi::Tensor buffer)
	Receive a buffer from the corresponding worker in the previous group. An error is thrown if the worker is already in the first group.

TVM_RUNTIME_DLL void	SendToWorker (ffi::Tensor buffer, int receiver_id)
	Send a buffer to the target receiver worker (globally across all groups).

TVM_RUNTIME_DLL void	RecvFromWorker (ffi::Tensor buffer, int sender_id)
	Receive a buffer from the target sender worker (globally across all groups).

TVM_RUNTIME_DLL int	WorkerId ()
	Get the local worker id.

TVM_RUNTIME_DLL void	SyncWorker ()
	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.

std::string	DiscoAction2String (DiscoAction action)
	Converts the enum class `DiscoAction` to string.

TVM_RUNTIME_DLL ffi::Tensor	Empty (ffi::Shape shape, DLDataType dtype, Device dev, ffi::Optional< ffi::String > mem_scope=std::nullopt)
	Allocate an uninitialized tensor through the runtime device allocator.

TVM_RUNTIME_DLL void	CopyFromTo (ffi::TensorView source, ffi::TensorView destination, TVMStreamHandle stream=nullptr)
	Copy between existing tensors with exactly matching byte sizes.

TVM_RUNTIME_DLL ffi::Tensor	CopyTo (ffi::TensorView source, Device dev, ffi::Optional< ffi::String > mem_scope=std::nullopt)
	Allocate a fresh tensor on dev, copy source, and synchronize the copy.

TVM_RUNTIME_DLL ffi::Tensor	CreateView (ffi::Tensor source, ffi::Shape shape, DLDataType dtype, uint64_t relative_byte_offset=0)
	Create a compact view retaining ownership of source.

TVM_RUNTIME_DLL void	CopyFromBytes (ffi::TensorView destination, const void *source_bytes, size_t nbytes, TVMStreamHandle stream=nullptr)
	Copy exactly nbytes from a byte buffer into a contiguous tensor.

TVM_RUNTIME_DLL void	CopyToBytes (ffi::TensorView source, void *destination_bytes, size_t nbytes, TVMStreamHandle stream=nullptr)
	Copy a contiguous tensor into a byte buffer of exactly nbytes.

void	SaveTensor (ffi::TensorView source, support::Stream *strm)
	Save a tensor in the portable CPU tensor format.

ffi::Tensor	LoadTensor (support::Stream *strm)
	Load a tensor from the portable format into CPU storage.

Device	GetPreferredHostDevice (Device device)
	Get the preferred host device from the input device.

TVM_RUNTIME_DLL ffi::ObjectPtr< TensorMemoryPool >	CreateTensorMemoryPool (DLDevice device)
	Create the default pool for a device.

ffi::Function	WrapTimeEvaluator (ffi::Function f, Device dev, int number, int repeat, int min_repeat_ms, int limit_zero_time_iterations, int cooldown_interval_ms, int repeats_to_cooldown, int cache_flush_bytes=0, ffi::Function f_preproc=nullptr)
	Wrap a timer function to measure the time cost of a given packed function.

Variables
constexpr int	kAllocAlignment = 64
	Number of bytes each allocation must align to.

constexpr int	kTempAllocaAlignment = 64
	Number of bytes each allocation must align to in temporary allocation.

constexpr int	kMaxStackAlloca = 1024
	Maximum size that can be allocated on stack.

constexpr int	kDefaultWorkspaceAlignment = 1
	Number of bytes each allocation must align to by default in the workspace buffer to service intermediate tensors.

constexpr int32_t	kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14
	Static FFI type index for `runtime::disco::DRef`.

constexpr uint64_t	kTVMTensorMagic = 0xDD5E40F096B4A13F
	Magic number for Tensor file.

Enumeration Type Documentation

◆ DiscoAction

enum class tvm::runtime::DiscoAction : int32_t

strong

All possible kinds of Disco commands.

Enumerator
kShutDown
kKillReg
kGetGlobalFunc
kCallPacked
kSyncWorker
kCopyFromWorker0
kCopyToWorker0
kDebugGetFromRemote
kDebugSetRegister

◆ ReduceKind

enum class tvm::runtime::ReduceKind : int32_t

strong

Possible kinds of reduction operations.

Enumerator
kSum
kProd
kMin
kMax
kAvg

◆ TVMDeviceExtType

enum tvm::runtime::TVMDeviceExtType

Extension device types in TVM.

Additional enumerators to supplement those provided by DLPack's DLDeviceType enumeration.

MAINTAINERS NOTE #1: We need to ensure that the two devices are identified by the same integer. Currently this requires manual verification. Discussed here: https://github.com/dmlc/dlpack/issues/111 As of DLPack v0.7, the highest-valued enumerator in DLDeviceType is kDLHexagon = 16.

MAINTAINERS NOTE #2: As of DLPack v0.7, the definition for DLDeviceType specifies an underlying storage type of int32_t. That guarantees a variable of type DLDeviceType is capable of holding any integers provided by either of these enumerations.

However, the int32_t specification only applies when the header file is compiled as C++, and this header file is also meant to work as C code. So the unspecified storage type could be a latent bug when compiled as C.

Enumerator
TVMDeviceExtType_End

Function Documentation

◆ AllGather()

TVM_RUNTIME_DLL void tvm::runtime::AllGather	(	ffi::Tensor	send,
		bool	in_group,
		ffi::Tensor	recv
	)

Perform an allgather operation using the underlying communication library.

Parameters

send	The array send to perform allgather on
in_group	Whether the allgather operation performs globally or in group as default.
recv	The array receives the outcome of allgather

◆ AllReduce()

TVM_RUNTIME_DLL void tvm::runtime::AllReduce	(	ffi::Tensor	send,
		ReduceKind	reduce_kind,
		bool	in_group,
		ffi::Tensor	recv
	)

Perform an allreduce operation using the underlying communication library.

Parameters

send	The array send to perform allreduce on
reduce_kind	The kind of reduction operation (e.g. sum, avg, min, max)
in_group	Whether the allreduce operation performs globally or in group as default.
recv	The array receives the outcome of allreduce

◆ BroadcastFromWorker0()

TVM_RUNTIME_DLL void tvm::runtime::BroadcastFromWorker0	(	ffi::Tensor	send,
		bool	in_group,
		ffi::Tensor	recv
	)

Perform a broadcast operation from worker-0.

Parameters

send	The buffer to be broadcasted
in_group	Whether the broadcast operation performs globally or in group as default.
recv	The buffer receives the broadcasted array

◆ CopyFromBytes()

TVM_RUNTIME_DLL void tvm::runtime::CopyFromBytes	(	ffi::TensorView	destination,
		const void *	source_bytes,
		size_t	nbytes,
		TVMStreamHandle	stream = `nullptr`
	)

Copy exactly nbytes from a byte buffer into a contiguous tensor.

Note: Synchronizes stream before returning.

◆ CopyFromTo()

TVM_RUNTIME_DLL void tvm::runtime::CopyFromTo	(	ffi::TensorView	source,
		ffi::TensorView	destination,
		TVMStreamHandle	stream = `nullptr`
	)

Copy between existing tensors with exactly matching byte sizes.

Note: The copy may be asynchronous. The caller owns stream synchronization.

◆ CopyTo()

TVM_RUNTIME_DLL ffi::Tensor tvm::runtime::CopyTo	(	ffi::TensorView	source,
		Device	dev,
		ffi::Optional< ffi::String >	mem_scope = `std::nullopt`
	)

Allocate a fresh tensor on dev, copy source, and synchronize the copy.

Note: A fresh allocation is made even when source already resides on dev.

◆ CopyToBytes()

TVM_RUNTIME_DLL void tvm::runtime::CopyToBytes	(	ffi::TensorView	source,
		void *	destination_bytes,
		size_t	nbytes,
		TVMStreamHandle	stream = `nullptr`
	)

Copy a contiguous tensor into a byte buffer of exactly nbytes.

Note: Synchronizes stream before returning.

◆ CreateTensorMemoryPool()

TVM_RUNTIME_DLL ffi::ObjectPtr< TensorMemoryPool > tvm::runtime::CreateTensorMemoryPool ( DLDevice device )

Create the default pool for a device.

◆ CreateView()

TVM_RUNTIME_DLL ffi::Tensor tvm::runtime::CreateView	(	ffi::Tensor	source,
		ffi::Shape	shape,
		DLDataType	dtype,
		uint64_t	relative_byte_offset = `0`
	)

Create a compact view retaining ownership of source.

Parameters

source	The owning source tensor, which must be compact.
shape	The view shape.
dtype	The view data type.
relative_byte_offset	Byte offset relative to the source's byte offset.

Note: The view must fit within the source's storage extent.

◆ DiscoAction2String()

std::string tvm::runtime::DiscoAction2String ( DiscoAction action )

inline

Converts the enum class DiscoAction to string.

◆ DiscoEmptyTensor()

TVM_RUNTIME_DLL ffi::Tensor tvm::runtime::DiscoEmptyTensor	(	ffi::Shape	shape,
		DLDataType	dtype,
		ffi::Optional< Device >	device
	)

Create an uninitialized empty Tensor.

Parameters

shape	The shape of the Tensor
dtype	The dtype of the Tensor
device	The device the Tensor is created on. If None, use the thread local default device

Returns: The Tensor created

◆ DLDeviceType2Str()

const char * tvm::runtime::DLDeviceType2Str ( int type )

inline

The name of DLDeviceType.

Parameters

type	The device type.

Returns: the device name.

◆ Empty()

TVM_RUNTIME_DLL ffi::Tensor tvm::runtime::Empty	(	ffi::Shape	shape,
		DLDataType	dtype,
		Device	dev,
		ffi::Optional< ffi::String >	mem_scope = `std::nullopt`
	)

Allocate an uninitialized tensor through the runtime device allocator.

Parameters

shape	The tensor shape.
dtype	The tensor data type.
dev	The allocation device.
mem_scope	The optional device memory scope.

◆ GatherToWorker0()

TVM_RUNTIME_DLL void tvm::runtime::GatherToWorker0	(	ffi::Tensor	send,
		bool	in_group,
		ffi::Optional< ffi::Tensor >	recv
	)

Perform a gather operation to worker-0.

Parameters

send	The sending buffer, which must not be None.
in_group	Whether the gather operation performs globally or in group as default.
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.

◆ GetPreferredHostDevice()

Device tvm::runtime::GetPreferredHostDevice ( Device device )

inline

Get the preferred host device from the input device.

For CUDA and ROCm, CUDAHost and ROCMHost will be returned for pinned memory, since pinned memory reduces copy overhead.
For other devices, CPU is returned as a fallback.

◆ LoadTensor()

ffi::Tensor tvm::runtime::LoadTensor ( support::Stream * strm )

inline

Load a tensor from the portable format into CPU storage.

◆ LoadVMModule()

TVM_RUNTIME_DLL ffi::Module tvm::runtime::LoadVMModule	(	std::string	path,
		ffi::Optional< Device >	device
	)

Load a runtime Module, then create and initialize a RelaxVM.

Parameters

path	The path to the runtime Module (a DSO file) to be loaded
device	The default device used to initialize the RelaxVM

Returns: The RelaxVM as a runtime Module

◆ operator<<()

std::ostream & tvm::runtime::operator<<	(	std::ostream &	os,
		DLDevice	dev
	)

inline

◆ RecvFromPrevGroup()

TVM_RUNTIME_DLL void tvm::runtime::RecvFromPrevGroup ( ffi::Tensor buffer )

Receive a buffer from the corresponding worker in the previous group. An error is thrown if the worker is already in the first group.

Parameters

buffer The receiving buffer.

◆ RecvFromWorker()

TVM_RUNTIME_DLL void tvm::runtime::RecvFromWorker	(	ffi::Tensor	buffer,
		int	sender_id
	)

Receive a buffer from the target sender worker (globally across all groups).

Parameters

buffer	The receiving buffer.
sender_id	The global sender worker id.

◆ RecvFromWorker0()

TVM_RUNTIME_DLL void tvm::runtime::RecvFromWorker0 ( ffi::Tensor buffer )

Receive a buffer from worker-0. No-op if the current worker is worker-0.

Parameters

buffer The buffer to be received

◆ ReduceKind2String()

std::string tvm::runtime::ReduceKind2String ( ReduceKind kind )

inline

Converts ReduceKind to string.

◆ RuntimeEnabled()

TVM_RUNTIME_DLL bool tvm::runtime::RuntimeEnabled ( const ffi::String & target )

Check if runtime module is enabled for target.

Parameters

target The target module name.

Returns: Whether runtime is enabled.

◆ SaveTensor()

void tvm::runtime::SaveTensor	(	ffi::TensorView	source,
		support::Stream *	strm
	)

inline

Save a tensor in the portable CPU tensor format.

◆ ScatterFromWorker0()

TVM_RUNTIME_DLL void tvm::runtime::ScatterFromWorker0	(	ffi::Optional< ffi::Tensor >	send,
		bool	in_group,
		ffi::Tensor	recv
	)

Perform a scatter operation from worker-0, chunking the given buffer into equal parts.

Parameters

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.
in_group	Whether the scatter operation performs globally or in group as default.
recv	The receiving buffer, which must not be None.

◆ SendToNextGroup()

TVM_RUNTIME_DLL void tvm::runtime::SendToNextGroup ( ffi::Tensor buffer )

Send a buffer to the corresponding worker in the next group. An error is thrown if the worker is already in the last group.

Parameters

buffer The sending buffer.

◆ SendToWorker()

TVM_RUNTIME_DLL void tvm::runtime::SendToWorker	(	ffi::Tensor	buffer,
		int	receiver_id
	)

Send a buffer to the target receiver worker (globally across all groups).

Parameters

buffer	The sending buffer.
receiver_id	The global receiver worker id.

◆ SyncWorker()

TVM_RUNTIME_DLL void tvm::runtime::SyncWorker ( )

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.

◆ WorkerId()

TVM_RUNTIME_DLL int tvm::runtime::WorkerId ( )

Get the local worker id.

◆ WrapTimeEvaluator()

ffi::Function tvm::runtime::WrapTimeEvaluator	(	ffi::Function	f,
		Device	dev,
		int	number,
		int	repeat,
		int	min_repeat_ms,
		int	limit_zero_time_iterations,
		int	cooldown_interval_ms,
		int	repeats_to_cooldown,
		int	cache_flush_bytes = `0`,
		ffi::Function	f_preproc = `nullptr`
	)

Wrap a timer function to measure the time cost of a given packed function.

Approximate implementation:

f() // warmup
for i in range(repeat)
  f_preproc()
  while True:
    start = time()
    for j in range(number):
      f()
    duration_ms = time() - start
    if duration_ms >= min_repeat_ms:
      break
    else:
       number = (min_repeat_ms / (duration_ms / number) + 1
  if cooldown_interval_ms and i % repeats_to_cooldown == 0:
    sleep(cooldown_interval_ms)

Parameters

f	The function argument.
dev	The device.
number	The number of times to run this function for taking average. We call these runs as one `repeat` of measurement.
repeat	The number of times to repeat the measurement. In total, the function will be invoked (1 + number x repeat) times, where the first one is warm up and will be discarded. The returned result contains `repeat` costs, each of which is an average of `number` costs.
min_repeat_ms	The minimum duration of one `repeat` in milliseconds. By default, one `repeat` contains `number` runs. If this parameter is set, the parameters `number` will be dynamically adjusted to meet the minimum duration requirement of one `repeat`. i.e., When the run time of one `repeat` falls below this time, the `number` parameter will be automatically increased.
limit_zero_time_iterations	The maximum number of repeats when measured time is equal to 0. It helps to avoid hanging during measurements.
cooldown_interval_ms	The cooldown interval in milliseconds between the number of repeats defined by `repeats_to_cooldown`.
repeats_to_cooldown	The number of repeats before the cooldown is activated.
cache_flush_bytes	The number of bytes to flush from cache before
f_preproc	The function to be executed before we execute time evaluator.

Returns: f_timer A timer function.

Variable Documentation

◆ kAllocAlignment

constexpr int tvm::runtime::kAllocAlignment = 64

constexpr

Number of bytes each allocation must align to.

◆ kDefaultWorkspaceAlignment

constexpr int tvm::runtime::kDefaultWorkspaceAlignment = 1

constexpr

Number of bytes each allocation must align to by default in the workspace buffer to service intermediate tensors.

◆ kMaxStackAlloca

constexpr int tvm::runtime::kMaxStackAlloca = 1024

constexpr

Maximum size that can be allocated on stack.

◆ kRuntimeDiscoDRef

constexpr int32_t tvm::runtime::kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14

constexpr

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.

◆ kTempAllocaAlignment

constexpr int tvm::runtime::kTempAllocaAlignment = 64

constexpr

Number of bytes each allocation must align to in temporary allocation.

◆ kTVMTensorMagic

constexpr uint64_t tvm::runtime::kTVMTensorMagic = 0xDD5E40F096B4A13F

constexpr

Magic number for Tensor file.

Namespaces

Classes

Enumerations

Functions

Variables

Enumeration Type Documentation

◆ DiscoAction

◆ ReduceKind

◆ TVMDeviceExtType

Function Documentation

◆ AllGather()

◆ AllReduce()

◆ BroadcastFromWorker0()

◆ CopyFromBytes()

◆ CopyFromTo()

◆ CopyTo()

◆ CopyToBytes()

◆ CreateTensorMemoryPool()

◆ CreateView()

◆ DiscoAction2String()

◆ DiscoEmptyTensor()

◆ DLDeviceType2Str()

◆ Empty()

◆ GatherToWorker0()

◆ GetPreferredHostDevice()

◆ LoadTensor()

◆ LoadVMModule()

◆ operator<<()

◆ RecvFromPrevGroup()

◆ RecvFromWorker()

◆ RecvFromWorker0()

◆ ReduceKind2String()

◆ RuntimeEnabled()

◆ SaveTensor()

◆ ScatterFromWorker0()

◆ SendToNextGroup()

◆ SendToWorker()

◆ SyncWorker()

◆ WorkerId()

◆ WrapTimeEvaluator()

Variable Documentation

◆ kAllocAlignment

◆ kDefaultWorkspaceAlignment

◆ kMaxStackAlloca

◆ kRuntimeDiscoDRef

◆ kTempAllocaAlignment

◆ kTVMTensorMagic