tvm.rpc#

Lightweight TVM RPC module.

RPC enables connect to a remote server, upload and launch functions. This is useful to for cross-compile and remote testing, The compiler stack runs on local server, while we use RPC server to run on remote runtime which don’t have a compiler available.

The test program compiles the program on local server, upload and run remote RPC server, get the result back to verify correctness.

Classes:

LocalSession()

RPCSession interface backed by local environment.

PopenSession(binary)

RPCSession interface backed by popen.

RPCSession(sess)

RPC Client session module

Server([host, port, port_end, is_proxy, ...])

Start RPC server on a separate process.

TrackerSession(addr)

Tracker client session.

Functions:

connect(url, port[, key, session_timeout, ...])

Connect to RPC Server

connect_tracker(url, port)

Connect to a RPC tracker

with_minrpc(compile_func[, server, runtime])

Attach the compiler function with minrpc related options.

class tvm.rpc.LocalSession[源代码]#

RPCSession interface backed by local environment.

This class can be used to implement functions that need to be ran both locally and remotely.

class tvm.rpc.PopenSession(binary)[源代码]#

RPCSession interface backed by popen.

Parameters#

binaryList[Union[str, bytes]]

The binary to be executed.

class tvm.rpc.RPCSession(sess)[源代码]#

RPC Client session module

Do not directly create the object, call connect

Methods:

cl([dev_id])

Construct OpenCL device.

cpu([dev_id])

Construct CPU device.

cuda([dev_id])

Construct CUDA GPU device.

device(dev_type[, dev_id])

Construct a remote device.

download(path)

Download file from remote temp folder.

download_linked_module(path)

Link a module in the remote and download it.

ext_dev([dev_id])

Construct extension device.

get_function(name)

Get function from the session.

hexagon([dev_id])

Construct Hexagon device.

listdir(path)

ls files from remote temp folder.

load_module(path)

Load a remote module, the file need to be uploaded first.

metal([dev_id])

Construct Metal device.

remove(path)

Remove file from remote temp folder.

rocm([dev_id])

Construct ROCm device.

system_lib()

Get system-wide library module.

upload(data[, target])

Upload file to remote runtime temp folder

vulkan([dev_id])

Construct Vulkan device.

webgpu([dev_id])

Construct WebGPU device.

cl(dev_id=0)[源代码]#

Construct OpenCL device.

cpu(dev_id=0)[源代码]#

Construct CPU device.

cuda(dev_id=0)[源代码]#

Construct CUDA GPU device.

device(dev_type, dev_id=0)[源代码]#

Construct a remote device.

Parameters#

dev_type: int or str

dev_id: int, optional

Returns#

dev: Device

The corresponding encoded remote device.

download(path)[源代码]#

Download file from remote temp folder.

Parameters#

pathstr

The relative location to remote temp folder.

Returns#

blobbytearray

The result blob from the file.

download_linked_module(path)[源代码]#

Link a module in the remote and download it.

Parameters#

pathstr

The relative location to remote temp folder.

Returns#

blobbytearray

The result blob from the file.

Note#

This function can be helpful when a linker is not available on the local client.

Examples#

mod = build_module_with_cross_compilation()
# export the module as tar because a local linker is not available
mod.export_library("lib.tar")
remote.upload("lib.tar")
# invoke the linker on the remote to link the module as a library
# note that the library can only run on the same env as the remote
with open("lib.so", "wb") as file:
    file.write(remote.download_linked_module("lib.tar"))
ext_dev(dev_id=0)[源代码]#

Construct extension device.

get_function(name)[源代码]#

Get function from the session.

Parameters#

namestr

The name of the function

Returns#

fFunction

The result function.

hexagon(dev_id=0)[源代码]#

Construct Hexagon device.

listdir(path)[源代码]#

ls files from remote temp folder.

Parameters#

path: str

The relative location to remote temp folder.

Returns#

dirs: str

The files in the given directory with split token ‘,’.

load_module(path)[源代码]#

Load a remote module, the file need to be uploaded first.

Parameters#

pathstr

The relative location to remote temp folder.

Returns#

mModule

The remote module containing remote function.

metal(dev_id=0)[源代码]#

Construct Metal device.

remove(path)[源代码]#

Remove file from remote temp folder.

Parameters#

path: str

The relative location to remote temp folder.

rocm(dev_id=0)[源代码]#

Construct ROCm device.

system_lib()[源代码]#

Get system-wide library module.

Returns#

moduleruntime.Module

The system-wide library module.

See Also#

tvm.runtime.system_lib

upload(data, target=None)[源代码]#

Upload file to remote runtime temp folder

Parameters#

datastr or bytearray

The file name or binary in local to upload.

targetstr, optional

The path in remote

vulkan(dev_id=0)[源代码]#

Construct Vulkan device.

webgpu(dev_id=0)[源代码]#

Construct WebGPU device.

class tvm.rpc.Server(host='0.0.0.0', port=9091, port_end=9199, is_proxy=False, tracker_addr=None, key='', load_library=None, custom_addr=None, silent=False, no_fork=False, server_init_callback=None, reuse_addr=True, timeout=None)[源代码]#

Start RPC server on a separate process.

This is a simple python implementation based on multi-processing. It is also possible to implement a similar C based server with TVM runtime which does not depend on the python.

Parameters#

hoststr

The host url of the server.

portint

The port to be bind to

port_endint, optional

The end port to search

is_proxybool, optional

Whether the address specified is a proxy. If this is true, the host and port actually corresponds to the address of the proxy server.

tracker_addr: Tuple (str, int) , optional

The address of RPC Tracker in tuple(host, ip) format. If is not None, the server will register itself to the tracker.

keystr, optional

The key used to identify the device type in tracker.

load_librarystr, optional

List of additional libraries to be loaded during execution.

custom_addr: str, optional

Custom IP Address to Report to RPC Tracker

silent: bool, optional

Whether run this server in silent mode.

no_fork: bool, optional

Whether forbid fork in multiprocessing.

server_init_callback: Callable, optional

Additional initialization function when starting the server.

reuse_addr: bool, optional

Allows the kernel to reuse a local socket in TIME_WAIT state.

timeout: float, optional

set a timeout for all operations on the socket

Note#

The RPC server only sees functions in the tvm namespace. To bring additional custom functions to the server env, you can use server_init_callback.

def server_init_callback():
    import tvm
    # must import mypackage here
    import mypackage

    tvm.register_func("function", mypackage.func)

server = rpc.Server(host, server_init_callback=server_init_callback)

Methods:

terminate()

Terminate the server process

terminate()[源代码]#

Terminate the server process

class tvm.rpc.TrackerSession(addr)[源代码]#

Tracker client session.

Parameters#

addrtuple

The address tuple

Methods:

close()

Close the tracker connection.

request(key[, priority, session_timeout, ...])

Request a new connection from the tracker.

request_and_run(key, func[, priority, ...])

Request a resource from tracker and run the func.

summary()

Get the summary dict of the tracker.

text_summary()

Get a text summary of the tracker.

close()[源代码]#

Close the tracker connection.

request(key, priority=1, session_timeout=0, max_retry=5, session_constructor_args=None)[源代码]#

Request a new connection from the tracker.

Parameters#

keystr

The type key of the device.

priorityint, optional

The priority of the request.

session_timeoutfloat, optional

The duration of the session, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.

max_retryint, optional

Maximum number of times to retry before give up.

session_constructor_argslist, optional

List of additional arguments to passed as the remote session constructor. The first element of the list is always a string specifying the name of the session constructor, the following args are the positional args to that function.

request_and_run(key, func, priority=1, session_timeout=0, max_retry=2)[源代码]#

Request a resource from tracker and run the func.

This function safe-guard rare server node dropout during execution. In such case, a new resource will be requested and func will be ran again.

Parameters#

keystr

The type key of the device.

funcfunction of session -> value

A stateless function

priorityint, optional

The priority of the request.

session_timeoutfloat, optional

The duration of the session, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.

max_retryint, optional

Maximum number of times to retry the function before give up.

summary()[源代码]#

Get the summary dict of the tracker.

text_summary()[源代码]#

Get a text summary of the tracker.

tvm.rpc.connect(url, port, key='', session_timeout=0, session_constructor_args=None, enable_logging=False)[源代码]#

Connect to RPC Server

Parameters#

urlstr

The url of the host

portint

The port to connect to

keystr, optional

Additional key to match server

session_timeoutfloat, optional

The duration of the session in seconds, allows server to kill the connection when duration is longer than this value. When duration is zero, it means the request must always be kept alive.

session_constructor_args: List

List of additional arguments to passed as the remote session constructor. The first element of the list is always a string specifying the name of the session constructor, the following args are the positional args to that function.

enable_logging: boolean

flag to enable/disable logging. Logging is disabled by default.

Returns#

sessRPCSession

The connected session.

Examples#

Normal usage .. code-block:: python

client = rpc.connect(server_url, server_port, server_key)

Session_constructor can be used to customize the session in the remote The following code connects to a remote internal server via a proxy by constructing another RPCClientSession on the proxy machine and use that as the serving session of the proxy endpoint.

client_via_proxy = rpc.connect(
    proxy_server_url, proxy_server_port, proxy_server_key, enable_logging
    session_constructor_args=[
        "rpc.Connect", internal_url, internal_port, internal_key, internal_logging])
tvm.rpc.connect_tracker(url, port)[源代码]#

Connect to a RPC tracker

Parameters#

urlstr

The url of the host

portint

The port to connect to

Returns#

sessTrackerSession

The connected tracker session.

tvm.rpc.with_minrpc(compile_func, server='posix_popen_server', runtime='libtvm')[源代码]#

Attach the compiler function with minrpc related options.

Parameters#

compile_funcUnion[str, Callable[[str, str, Optional[str]], None]]

The compilation function to decorate.

serverstr

The server type.

runtimestr

The runtime library.

Returns#

fcompilefunction

The return compilation.