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.Server(host, port=9091, port_end=9199, is_proxy=False, use_popen=False, tracker_addr=None, key='', load_library=None, custom_addr=None, silent=False, utvm_dev_id=None, utvm_dev_config_args=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
  • host (str) – The host url of the server.

  • port (int) – The port to be bind to

  • port_end (int, optional) – The end port to search

  • is_proxy (bool, 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.

  • use_popen (bool, optional) – Whether to use Popen to start a fresh new process instead of fork. This is recommended to switch on if we want to do local RPC demonstration for GPU devices to avoid fork safety issues.

  • 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.

  • key (str, optional) – The key used to identify the device type in tracker.

  • load_library (str, 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.

Methods:

terminate()

Terminate the server process

terminate()

Terminate the server process

tvm.rpc.connect(url, port, key='', session_timeout=0, session_constructor_args=None)

Connect to RPC Server

Parameters
  • url (str) – The url of the host

  • port (int) – The port to connect to

  • key (str, optional) – Additional key to match server

  • session_timeout (float, 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.

  • 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.

Returns

sess – The connected session.

Return type

RPCSession

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,
    session_constructor_args=[
        "rpc.Connect", internal_url, internal_port, internal_key])
tvm.rpc.connect_tracker(url, port)

Connect to a RPC tracker

Parameters
  • url (str) – The url of the host

  • port (int) – The port to connect to

Returns

sess – The connected tracker session.

Return type

TrackerSession

class tvm.rpc.RPCSession(sess)

RPC Client session module

Do not directly create the obhect, call connect

Methods:

cl([dev_id])

Construct OpenCL device.

context(dev_type[, dev_id])

Construct a remote context.

cpu([dev_id])

Construct CPU 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.

gpu([dev_id])

Construct GPU device.

hexagon([dev_id])

Construct Hexagon device.

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.

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.

system_lib()

Get system-wide library module.

Returns

module – The system-wide library module.

Return type

runtime.Module

get_function(name)

Get function from the session.

Parameters

name (str) – The name of the function

Returns

f – The result function.

Return type

Function

context(dev_type, dev_id=0)

Construct a remote context.

Parameters
  • dev_type (int or str) –

  • dev_id (int, optional) –

Returns

ctx – The corresponding encoded remote context.

Return type

TVMContext

upload(data, target=None)

Upload file to remote runtime temp folder

Parameters
  • data (str or bytearray) – The file name or binary in local to upload.

  • target (str, optional) – The path in remote

download(path)

Download file from remote temp folder.

Parameters

path (str) – The relative location to remote temp folder.

Returns

blob – The result blob from the file.

Return type

bytearray

remove(path)

Remove file from remote temp folder.

Parameters

path (str) – The relative location to remote temp folder.

load_module(path)

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

Parameters

path (str) – The relative location to remote temp folder.

Returns

m – The remote module containing remote function.

Return type

Module

download_linked_module(path)

Link a module in the remote and download it.

Parameters

path (str) – The relative location to remote temp folder.

Returns

blob – The result blob from the file.

Return type

bytearray

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"))
cpu(dev_id=0)

Construct CPU device.

gpu(dev_id=0)

Construct GPU device.

cl(dev_id=0)

Construct OpenCL device.

vulkan(dev_id=0)

Construct Vulkan device.

metal(dev_id=0)

Construct Metal device.

ext_dev(dev_id=0)

Construct extension device.

hexagon(dev_id=0)

Construct Hexagon device.

webgpu(dev_id=0)

Construct WebGPU device.

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

binary (List[Union[str, bytes]]) – The binary to be executed.

class tvm.rpc.TrackerSession(addr)

Tracker client session.

Parameters

addr (tuple) – 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.

summary()

Get the summary dict of the tracker.

text_summary()

Get a text summary of the tracker.

request(key, priority=1, session_timeout=0, max_retry=5)

Request a new connection from the tracker.

Parameters
  • key (str) – The type key of the device.

  • priority (int, optional) – The priority of the request.

  • session_timeout (float, 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_retry (int, optional) – Maximum number of times to retry before give up.

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
  • key (str) – The type key of the device.

  • func (function of session -> value) – A stateless function

  • priority (int, optional) – The priority of the request.

  • session_timeout (float, 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_retry (int, optional) – Maximum number of times to retry the function before give up.

tvm.rpc.with_minrpc(compile_func, server='posix_popen_server', runtime='libtvm')

Attach the compiler function with minrpc related options.

Parameters
  • compile_func (Union[str, Callable[[str, str, Optional[str]], None]]) – The compilation function to decorate.

  • server (str) – The server type.

  • runtime (str) – The runtime library.

Returns

fcompile – The return compilation.

Return type

function