tvm.rpc

轻量级 TVM RPC 模块。

RPC 允许连接到远程服务器,上传和启动函数。这对于交叉编译和远程测试非常有用。编译器堆栈在本地服务器上运行,而我们使用 RPC 服务器在没有可用编译器的远程运行时上运行。

测试程序在本地服务器上编译程序,上传并在远程 RPC 服务器上运行,取回结果以验证正确性。

TVM RPC 服务器假定用户是可信任的,需要在可信任的网络环境和加密通道中使用。它允许将任意文件写入服务器,并为任何可以访问此 API 的人提供完整的远程代码执行能力。

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

在单独的进程中启动 RPC 服务器。

RPCSession(sess)

RPC 客户端会话模块

LocalSession()

由本地环境支持的 RPCSession 接口。

PopenSession(binary)

由 popen 支持的 RPCSession 接口。

TrackerSession(addr)

Tracker 客户端会话。

函数

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

连接到 RPC 服务器

connect_tracker(url, port)

连接到 RPC tracker

with_minrpc(compile_func[, server, runtime])

将编译器函数附加到 minrpc 相关选项。

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)

在单独的进程中启动 RPC 服务器。

这是一个基于多处理的简单 Python 实现。也可以使用不依赖 Python 的 TVM 运行时来实现类似的基于 C 的服务器。

参数:
  • host (str) – 服务器的主机 url。

  • port (int) – 要绑定的端口

  • port_end (int, optional) – 要搜索的结束端口

  • is_proxy (bool, optional) – 指定的地址是否为代理。如果为 true,则主机和端口实际上对应于代理服务器的地址。

  • tracker_addr (Tuple (str, int) , optional) – RPC Tracker 的地址,格式为 tuple(host, ip)。如果不是 None,服务器将自己注册到 tracker。

  • key (str, optional) – 用于在 tracker 中标识设备类型的 key。

  • load_library (str, optional) – 执行期间要加载的其他库的列表。

  • custom_addr (str, optional) – 要报告给 RPC Tracker 的自定义 IP 地址

  • silent (bool, optional) – 是否以静默模式运行此服务器。

  • no_fork (bool, optional) – 是否禁止在多处理中 fork。

  • server_init_callback (Callable, optional) – 启动服务器时的其他初始化函数。

  • reuse_addr (bool, optional) – 允许内核重用 TIME_WAIT 状态下的本地套接字。

  • timeout (float, optional) – 设置套接字上所有操作的超时时间

注意

TVM RPC 服务器假定用户是可信任的,需要在可信任的网络环境和加密通道中使用。它允许将任意文件写入服务器,并为任何可以访问此 API 的人提供完整的远程代码执行能力。

RPC 服务器仅查看 tvm 命名空间中的函数。要将其他自定义函数引入服务器环境,可以使用 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)

方法

terminate()

终止服务器进程

terminate()

终止服务器进程

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

连接到 RPC 服务器

参数:
  • url (str) – 主机的 url

  • port (int) – 要连接的端口

  • key (str, optional) – 用于匹配服务器的其他 key

  • session_timeout (float, optional) – 会话的持续时间(秒),允许服务器在持续时间超过此值时终止连接。当持续时间为零时,表示请求必须始终保持活动状态。

  • session_constructor_args (List) – 作为远程会话构造函数的附加参数列表。列表的第一个元素始终是一个字符串,指定会话构造函数的名称,后面的参数是该函数的位置参数。

  • enable_logging (boolean) – 用于启用/禁用日志记录的标志。默认情况下禁用日志记录。

返回:

sess – 连接的会话。

返回类型:

RPCSession

示例

正常用法 .. code-block:: python

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

Session_constructor 可用于自定义远程会话。以下代码通过在代理计算机上构建另一个 RPCClientSession 并将其用作代理端点的服务会话,从而通过代理连接到远程内部服务器。

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)

连接到 RPC tracker

参数:
  • url (str) – 主机的 url

  • port (int) – 要连接的端口

返回:

sess – 连接的 tracker 会话。

返回类型:

TrackerSession

class tvm.rpc.RPCSession(sess)

RPC 客户端会话模块

不要直接创建对象,调用 connect

方法

system_lib()

获取系统范围的库模块。

get_function(name)

从会话中获取函数。

device(dev_type[, dev_id])

构造远程设备。

upload(data[, target])

将文件上传到远程运行时临时文件夹

download(path)

从远程临时文件夹下载文件。

remove(path)

从远程临时文件夹中删除文件。

listdir(path)

ls 来自远程临时文件夹的文件。

load_module(path)

加载远程模块,需要先上传文件。

download_linked_module(path)

链接远程模块并下载它。

cpu([dev_id])

构造 CPU 设备。

cuda([dev_id])

构造 CUDA GPU 设备。

cl([dev_id])

构造 OpenCL 设备。

vulkan([dev_id])

构造 Vulkan 设备。

metal([dev_id])

构造 Metal 设备。

rocm([dev_id])

构造 ROCm 设备。

ext_dev([dev_id])

构造扩展设备。

hexagon([dev_id])

构造 Hexagon 设备。

webgpu([dev_id])

构造 WebGPU 设备。

system_lib()

获取系统范围的库模块。

返回:

module – 系统范围的库模块。

返回类型:

runtime.Module

另请参阅

tvm.runtime.system_lib

get_function(name)

从会话中获取函数。

参数:

name (str) – 函数的名称

返回:

f – 结果函数。

返回类型:

函数

device(dev_type, dev_id=0)

构造远程设备。

参数:
  • dev_type (int or str)

  • dev_id (int, optional)

返回:

dev – 相应的编码远程设备。

返回类型:

设备

upload(data, target=None)

将文件上传到远程运行时临时文件夹

参数:
  • data (str or bytearray) – 要上传的本地文件名或二进制文件。

  • target (str, optional) – 远程路径

download(path)

从远程临时文件夹下载文件。

参数:

path (str) – 相对于远程临时文件夹的位置。

返回:

blob – 来自文件的结果 blob。

返回类型:

bytearray

remove(path)

从远程临时文件夹中删除文件。

参数:

path (str) – 相对于远程临时文件夹的位置。

listdir(path)

ls 来自远程临时文件夹的文件。

参数:

path (str) – 相对于远程临时文件夹的位置。

返回:

dirs – 给定目录中的文件,用分隔符“,”分割。

返回类型:

str

load_module(path)

加载远程模块,需要先上传文件。

参数:

path (str) – 相对于远程临时文件夹的位置。

返回:

m – 包含远程函数的远程模块。

返回类型:

Module

download_linked_module(path)

链接远程模块并下载它。

参数:

path (str) – 相对于远程临时文件夹的位置。

返回:

blob – 来自文件的结果 blob。

返回类型:

bytearray

注意

当本地客户端上没有链接器时,此函数可能很有用。

示例

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)

构造 CPU 设备。

cuda(dev_id=0)

构造 CUDA GPU 设备。

cl(dev_id=0)

构造 OpenCL 设备。

vulkan(dev_id=0)

构造 Vulkan 设备。

metal(dev_id=0)

构造 Metal 设备。

rocm(dev_id=0)

构造 ROCm 设备。

ext_dev(dev_id=0)

构造扩展设备。

hexagon(dev_id=0)

构造 Hexagon 设备。

webgpu(dev_id=0)

构造 WebGPU 设备。

class tvm.rpc.LocalSession

由本地环境支持的 RPCSession 接口。

此类可用于实现需要在本地和远程运行的函数。

class tvm.rpc.PopenSession(binary)

由 popen 支持的 RPCSession 接口。

参数:

binary (List[Union[str, bytes]]) – 要执行的二进制文件。

class tvm.rpc.TrackerSession(addr)

Tracker 客户端会话。

参数:

addr (tuple) – 地址元组

方法

close()

关闭 tracker 连接。

summary()

获取 tracker 的摘要字典。

text_summary()

获取 tracker 的文本摘要。

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

从 tracker 请求新连接。

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

从 tracker 请求资源并运行 func。

close()

关闭 tracker 连接。

summary()

获取 tracker 的摘要字典。

text_summary()

获取 tracker 的文本摘要。

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

从 tracker 请求新连接。

参数:
  • key (str) – 设备的类型 key。

  • priority (int, optional) – 请求的优先级。

  • session_timeout (float, optional) – 会话的持续时间,允许服务器在持续时间超过此值时终止连接。当持续时间为零时,表示请求必须始终保持活动状态。

  • max_retry (int, optional) – 放弃前最大重试次数。

  • session_constructor_args (list, optional) – 作为远程会话构造函数的附加参数列表。列表的第一个元素始终是一个字符串,指定会话构造函数的名称,后面的参数是该函数的位置参数。

request_and_run(key, func, priority=1, session_timeout=0, max_retry=2)

从 tracker 请求资源并运行 func。

此函数可以安全地防止执行期间罕见的服务器节点掉线。在这种情况下,将请求新资源并再次运行 func。

参数:
  • key (str) – 设备的类型 key。

  • func (function of session -> value) – 一个无状态函数

  • priority (int, optional) – 请求的优先级。

  • session_timeout (float, optional) – 会话的持续时间,允许服务器在持续时间超过此值时终止连接。当持续时间为零时,表示请求必须始终保持活动状态。

  • max_retry (int, optional) – 放弃前最大函数重试次数。

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

将编译器函数附加到 minrpc 相关选项。

参数:
  • compile_func (Union[str, Callable[[str, str, Optional[str]], None]]) – 要装饰的编译函数。

  • server (str) – 服务器类型。

  • runtime (str) – 运行时库。

返回:

fcompile – 返回编译。

返回类型:

函数