tvm.contrib

TVM Python包的贡献API。

Contrib API提供了许多有用的非核心功能。其中一些是与第三方库和工具交互的实用工具。

tvm.contrib.cblas

BLAS库的外部函数接口。

tvm.contrib.cblas.matmul(lhs, rhs, transa=False, transb=False, **kwargs)

创建一个外部运算符，使用CrhsLAS计算矩阵A与rhs的乘积该函数展示了如何调用外部库的示例。

Parameters:

lhs (Tensor) – 左侧矩阵操作数
rhs (Tensor) – 右侧矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.cblas.batch_matmul(lhs, rhs, transa=False, transb=False, iterative=False, **kwargs)

创建一个外部操作符，用于计算A和rhs的批量矩阵乘法，使用CBLAS库此函数展示了如何调用外部库的示例。

Parameters:

lhs (Tensor) – 左侧矩阵操作数
rhs (Tensor) – 右侧矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.clang

用于调用系统中clang的工具。

tvm.contrib.clang.find_clang(required=True)

在系统中查找clang。

Parameters:: 必填 (bool) - 是否为必填项，如果编译器要求必填但未填写，将引发运行时错误。
Returns:: valid_list – 有效路径列表。
Return type:: list 的 str

注意

该函数会首先搜索与构建tvm所用的llvm主版本匹配的clang

tvm.contrib.clang.create_llvm(inputs, output=None, options=None, cc=None)

创建llvm文本IR。

Parameters:

inputs (list of str) - 输入文件名或代码源的列表。
output (str, optional) – 输出文件，如果为none则创建一个临时文件
options (list) – 附加选项字符串的列表。
cc (str, optional) – clang编译器，如果未指定，我们将尝试猜测匹配的clang版本。

Returns:

code – 生成的LLVM文本IR。

Return type:

str

tvm.contrib.cc

用于调用系统中C/C++编译器的工具。

tvm.contrib.cc.get_cc()

返回默认C/C++编译器的路径。

Returns:: out – 默认C/C++编译器的路径，如果未找到则为None。
Return type:: 可选[str]

tvm.contrib.cc.create_shared(output, objects, options=None, cc=None, cwd=None, ccache_env=None)

创建共享库。

Parameters:

output (str) – 目标共享库。
objects (List[str]) – 对象文件列表。
options (List[str]) – 附加选项字符串的列表。
cc (可选[str]) – 编译器命令。
cwd (可选[str]) – 当前工作目录。
ccache_env (可选[字典[str, str]]) - ccache的环境变量。默认设置为None来禁用ccache。

tvm.contrib.cc.create_staticlib(output, inputs, ar=None)

创建静态库。

Parameters:

output (str) – 目标共享库。
inputs (List[str]) – 输入文件列表。每个输入文件可以是对象的压缩包或对象文件。
ar (可选[str]) – 要使用的ar命令路径

tvm.contrib.cc.create_executable(output, objects, options=None, cc=None, cwd=None, ccache_env=None)

创建可执行二进制文件。

Parameters:

output (str) – 目标可执行文件。
objects (List[str]) – 对象文件列表。
options (List[str]) - 附加选项字符串的列表。
cc (可选[str]) – 编译器命令。
cwd (可选[str]) – 当前工作目录。
ccache_env (Optional[Dict[str, str]]) – ccache的环境变量。默认设置为None来禁用ccache。

tvm.contrib.cc.get_global_symbol_section_map(path, *, nm=None) → Dict[str, str]

通过nm -g从库中获取全局符号

Parameters:

path (str) – 库路径
nm (str) – nm命令的路径

Returns:

symbol_section_map – 从已定义的全局符号到其所属分区的映射表

Return type:

字典[str, str]

tvm.contrib.cc.get_target_by_dump_machine(compiler)

可以获取目标三元组的仿函数，该目标三元组使用编译器获取。

Parameters:: compiler (可选参数[str]) – 编译器。
Returns:: out – 一个能根据编译器的dumpmachine选项获取目标三元组的函数。
Return type:: 可调用

tvm.contrib.cc.cross_compiler(compile_func, options=None, output_format=None, get_target_triple=None, add_files=None)

通过使用选项特化compile_func来创建一个交叉编译器函数。

此函数可用于构建编译函数，这些函数可以传递给AutoTVM的measure或export_library。

Parameters:

compile_func (Union[str, Callable[[str, str, Optional[str]], None]]) – 执行实际编译工作的函数
options (可选[列表[str]]) - 额外的可选字符串列表。
output_format (可选[str]) – 库的输出格式。
get_target_triple (可选[Callable]) - 根据编译器的dumpmachine选项获取目标三元组的函数。
add_files (Optional[List[str]]) – 作为编译部分传递的额外对象、源文件和库文件路径列表。

Returns:

fcompile – 一个可传递给export_library的编译函数。

Return type:

可调用[[str, str, 可选[str]], 无]

示例

from tvm.contrib import cc, ndk
# export using arm gcc
mod = build_runtime_module()
mod.export_library(path_dso,
                   fcompile=cc.cross_compiler("arm-linux-gnueabihf-gcc"))
# specialize ndk compilation options.
specialized_ndk = cc.cross_compiler(
    ndk.create_shared,
    ["--sysroot=/path/to/sysroot", "-shared", "-fPIC", "-lm"])
mod.export_library(path_dso, fcompile=specialized_ndk)

tvm.contrib.cublas

cuBLAS库的外部函数接口。

tvm.contrib.cublas.matmul(lhs, rhs, transa=False, transb=False, dtype=None)

创建一个外部操作，使用cuBLAS计算矩阵A和rhs的乘积

Parameters:

lhs (Tensor) – 左侧矩阵操作数
rhs (Tensor) – 右侧矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.cublas.batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None)

创建一个外部操作，使用cuBLAS计算A和rhs的批量矩阵乘法

Parameters:

lhs (Tensor) – 左侧矩阵操作数
rhs (Tensor) – 右侧矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.dlpack

封装函数以将有DLPack支持的框架桥接到TVM

tvm.contrib.dlpack.convert_func(tvm_func, tensor_type, to_dlpack_func)

Convert a tvm function into one that accepts a tensor from another: 框架，只要另一个框架支持DLPACK

Parameters:

tvm_func (Function) – 内置的tvm函数，用于操作数组
tensor_type (Type) – 目标框架张量的类型
to_dlpack_func (Function) – 将源张量转换为DLPACK格式的函数

tvm.contrib.dlpack.to_pytorch_func(tvm_func)

将tvm函数转换为可接受PyTorch张量的函数

Parameters:: tvm_func (Function) – 内置的tvm数组操作函数
Returns:: wrapped_func – 操作PyTorch张量的封装tvm函数
Return type:: Function

tvm.contrib.emcc

用于调用系统中emscripten编译器的工具。

tvm.contrib.emcc.create_tvmjs_wasm(output, objects, options=None, cc='emcc', libs=None)

创建预期与tvmjs一起运行的wasm。

Parameters:

output (str) – 目标共享库。
objects (list) – 对象文件列表。
options (str) – 附加选项。
cc (str, optional) – 编译字符串。
libs (list) – 用户自定义库文件列表（例如.bc文件），这些文件将被添加到wasm中。

tvm.contrib.miopen

MIOpen库的外部函数接口。

tvm.contrib.miopen.conv2d_forward(x, w, stride_h=1, stride_w=1, pad_h=0, pad_w=0, dilation_h=1, dilation_w=1, conv_mode=0, data_type=1, group_count=1)

创建一个使用MIOpen计算2D卷积的外部算子

Parameters:

x (Tensor) – 输入特征图
w (Tensor) – 卷积权重
stride_h (int) – 高度步长
stride_w (int) – 宽度步长
pad_h (int) – 高度填充
pad_w (int) – 宽度填充
dilation_h (int) – 高度扩张
dilation_w (int) – 宽度扩张系数
conv_mode (int) – 0: miopenConvolution 1: miopenTranspose
data_type (int) – 0: miopenHalf (fp16) 1: miopenFloat (fp32)
group_count (int) – 组数

Returns:

y – 结果张量

Return type:

Tensor

tvm.contrib.miopen.softmax(x, axis=-1)

使用MIOpen计算softmax

Parameters:

x (tvm.te.Tensor) – 输入张量
axis (int) – 计算softmax的轴

Returns:

ret – 结果张量

Return type:

tvm.te.Tensor

tvm.contrib.miopen.log_softmax(x, axis=-1)

使用MIOpen计算对数softmax

Parameters:

x (tvm.te.Tensor) – 输入张量
axis (int) – 计算log softmax的轴

Returns:

ret – 结果张量

Return type:

tvm.te.Tensor

tvm.contrib.ndk

用于调用NDK编译器工具链的工具。

tvm.contrib.ndk.create_shared(output, objects, options=None)

创建共享库。

Parameters:

output (str) – 目标共享库。
objects (list) – 对象文件列表。
options (list of str, optional) - 额外的选项。

tvm.contrib.ndk.create_staticlib(output, inputs)

创建静态库：

Parameters:

output (str) – 目标静态库。
inputs (list) – 对象文件或tar文件的列表

tvm.contrib.ndk.get_global_symbol_section_map(path, *, nm=None) → Dict[str, str]

通过NDK中的nm -gU命令从库中获取全局符号

Parameters:

path (str) – 库路径
nm (str) – nm命令的路径

Returns:

symbol_section_map – 从已定义的全局符号到其所属分区的映射表

Return type:

字典[str, str]

tvm.contrib.nnpack

NNPACK库的外部函数接口。

tvm.contrib.nnpack.is_available(): 检查NNPACK是否可用，即nnp_initialize() 返回nnp_status_success。

tvm.contrib.nnpack.fully_connected_inference(lhs, rhs, nthreads=1)

创建一个外部操作，使用nnpack计算一维张量lhs和二维张量rhs的全连接。

Parameters:

lhs (Tensor) – lhs 一维数组输入[input_channels]，包含FP32元素
rhs (Tensor) – lhs 2D矩阵核[output_channels][input_channels]，元素类型为FP32

Returns:

C – 左侧一维数组 out[output_channels]，包含 FP32 元素。

Return type:

Tensor

tvm.contrib.nnpack.convolution_inference(data, kernel, bias, padding, stride, nthreads=1, algorithm=0)

创建一个外部操作，使用nnpack对4D张量数据和4D张量核以及1D张量偏置进行推理卷积。

Parameters:

data (Tensor) – 数据4D张量输入[batch][input_channels][input_height][input_width]，元素类型为FP32。
kernel (Tensor) – 4D张量核，格式为kernel[输出通道数][输入通道数][核高度][核宽度]，由FP32元素组成。
bias (Tensor) – 偏置 1D 数组 bias[output_channels][input_channels][kernel_height][kernel_width]，由 FP32 元素组成。
padding (list) – 填充一个4维列表[pad_top, pad_bottom, pad_left, pad_right]，表示特征图周围的填充。
stride (list) – 步长一个二维列表[stride_height, stride_width]，表示步长。

Returns:

output – 输出4D张量 output[batch][output_channels][output_height][output_width]，元素类型为FP32。

Return type:

Tensor

tvm.contrib.nnpack.convolution_inference_without_weight_transform(data, transformed_kernel, bias, padding, stride, nthreads=1, algorithm=0)

创建一个外部操作(extern op)，使用nnpack对4D张量数据和4D预变换张量核以及1D张量偏置进行推理卷积运算。

Parameters:

data (Tensor) – 数据4D张量输入[batch][input_channels][input_height][input_width]，元素类型为FP32。
transformed_kernel (Tensor) – 转换后的4D张量核[output_channels][input_channels][tile][tile]，由FP32元素组成。
bias (Tensor) – 偏置项，1维数组 bias[output_channels][input_channels][kernel_height][kernel_width]，由FP32元素组成。
padding (list) – 填充一个4维列表[pad_top, pad_bottom, pad_left, pad_right]，表示特征图周围的填充。
stride (list) – 步长一个二维列表[stride_height, stride_width]，表示步长。

Returns:

output – 输出4D张量 output[batch][output_channels][output_height][output_width]，元素类型为FP32。

Return type:

Tensor

tvm.contrib.nnpack.convolution_inference_weight_transform(kernel, nthreads=1, algorithm=0, dtype='float32')

创建一个外部操作(extern op)，使用nnpack对3D张量数据、4D张量核和1D张量偏置进行推理卷积运算。

Parameters:: kernel (Tensor) – 内核4D张量 kernel[output_channels][input_channels][kernel_height][kernel_width]，由FP32元素组成。
Returns:: output – 输出4D张量 output[output_channels][input_channels][tile][tile]，元素类型为FP32。
Return type:: Tensor

tvm.contrib.nvcc

用于在系统中调用nvcc编译器的实用工具

tvm.contrib.nvcc.compile_cuda(code, target_format='ptx', arch=None, options=None, path_target=None)

使用环境中的NVCC编译CUDA代码。

Parameters:

code (str) – CUDA代码。
target_format (str) – nvcc编译器的目标格式。
arch (str) – CUDA架构。
options (str 或 list 的 str) - 附加选项。
path_target (str, optional) – 输出文件。

Returns:

cubin – cubin的字节数组

Return type:

bytearray

tvm.contrib.nvcc.find_cuda_path()

用于查找cuda路径的实用函数

Returns:: path – CUDA根目录的路径。
Return type:: str

tvm.contrib.nvcc.get_cuda_version(cuda_path=None)

用于获取CUDA版本的实用函数

Parameters:: cuda_path (可选[str]) – CUDA根目录路径。如果传入None，将默认使用find_cuda_path()。
Returns:: version – CUDA版本号
Return type:: float

tvm.contrib.nvcc.parse_compute_version(compute_version)

解析计算能力字符串以划分主版本和次版本

Parameters:

compute_version (str) – GPU的计算能力版本（例如"6.0"）

Returns:

major (int) – 主版本号
minor (int) – 次版本号

tvm.contrib.nvcc.have_fp16(compute_version)

无论计算能力是否支持fp16

Parameters:: compute_version (str) – GPU的计算能力版本（例如"6.0"）

tvm.contrib.nvcc.have_int8(compute_version)

无论计算能力是否支持int8类型

Parameters:: compute_version (str) – GPU的计算能力版本（例如"6.1"）

tvm.contrib.nvcc.have_tensorcore(compute_version=None, target=None)

无论计算能力是否支持TensorCore

Parameters:

compute_version (str, optional) – GPU的计算能力版本号（例如"7.0"）。
target (tvm.target.Target, 可选参数) – 编译目标，当未指定compute_version时将用于确定架构。

tvm.contrib.nvcc.have_cudagraph(): 提供了CUDA图支持

tvm.contrib.pickle_memoize

通过pickle缓存函数结果，用于测试用例的缓存。

class tvm.contrib.pickle_memoize.Cache(key, save_at_exit)

用于结果缓存的缓存对象。

Parameters:

key (str) – 该函数的文件键
save_at_exit (bool) – 是否在程序退出时将缓存保存到文件

property cache: 返回缓存，首次使用时初始化。

tvm.contrib.pickle_memoize.memoize(key, save_at_exit=False)

将函数结果缓存起来以便多次复用。

Parameters:

key (str) – 文件的唯一标识键
save_at_exit (bool) – 程序退出时是否将缓存保存到文件

Returns:

fmemoize – 用于执行记忆化的装饰器函数。

Return type:

函数

tvm.contrib.random

随机库的外部函数接口。

tvm.contrib.random.randint(low, high, size, dtype='int32')

返回从low（包含）到high（不包含）范围内的随机整数。返回指定dtype在"半开"区间[low, high)内"离散均匀"分布的随机整数。

Parameters:

low (int) – 从分布中抽取的最小(有符号)整数值
high (int) – 从分布中抽取的最大（有符号）整数加一的值

Returns:

out – 具有指定大小和数据类型（dtype）的张量

Return type:

Tensor

tvm.contrib.random.uniform(low, high, size)

从均匀分布中抽取样本。

样本均匀分布在半开区间[low, high)内（包含low但不包含high）。换句话说，在给定区间内的任何值被均匀分布抽中的概率都是相等的。

Parameters:

low (float) – 输出区间的下限。所有生成的值都将大于或等于该下限值。
high (float) – 输出区间的上限。所有生成的值都将小于该上限。
size (tuple of ints) – 输出形状。如果给定的形状是例如(m, n, k)，则会抽取m * n * k个样本。

Returns:

out – 具有指定大小和数据类型（dtype）的张量。

Return type:

Tensor

tvm.contrib.random.normal(loc, scale, size)

从正态分布中抽取样本。

从正态分布中返回随机样本。

Parameters:

loc (float) – 分布的loc参数。
scale (float) – 分布的标准差。
size (tuple of ints) – 输出形状。如果给定的形状是例如(m, n, k)，则会抽取m * n * k个样本。

Returns:

out – 具有指定大小和数据类型（dtype）的张量

Return type:

Tensor

tvm.contrib.rocblas

rocBLAS库的外部函数接口。

tvm.contrib.rocblas.matmul(lhs, rhs, transa=False, transb=False)

创建一个外部操作符，使用rocBLAS计算矩阵A与rhs的乘积

Parameters:

lhs (Tensor) – 左侧矩阵操作数
rhs (Tensor) – 右侧矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.rocblas.batch_matmul(lhs, rhs, transa=False, transb=False)

创建一个外部运算符，使用rocBLAS计算矩阵A与rhs的乘法

Parameters:

lhs (Tensor) – 左侧批量矩阵操作数
rhs (Tensor) – 右侧批量矩阵操作数
transa (bool) – 是否转置左侧矩阵
transb (bool) – 是否转置右侧矩阵

Returns:

C – 结果张量。

Return type:

Tensor

tvm.contrib.rocm

ROCm后端的实用工具

tvm.contrib.rocm.find_lld(required=True)

在系统中查找ld.lld。

Parameters:: 必填 (bool) - 是否为必填项，如果编译器要求必填但未填写，将引发运行时错误。
Returns:: valid_list – 有效路径列表。
Return type:: list 的 str

注意

该函数将首先搜索与构建tvm所用的主要llvm版本相匹配的ld.lld

tvm.contrib.rocm.rocm_link(in_file, out_file, lld=None)

使用 lld 将可重定位 ELF 对象链接为共享 ELF 对象

Parameters:

in_file (str) – 输入文件名（可重定位的ELF对象文件）
out_file (str) – 输出文件名（共享ELF对象文件）
lld (str, optional) – lld链接器，如果未指定，我们将尝试匹配对应的clang版本。

tvm.contrib.rocm.parse_compute_version(compute_version)

解析计算能力字符串以划分主版本和次版本

Parameters:

compute_version (str) – GPU的计算能力版本（例如"6.0"）

Returns:

major (int) – 主版本号
minor (int) – 次版本号

tvm.contrib.rocm.have_matrixcore(compute_version=None)

无论计算能力是否支持MatrixCore

Parameters:: compute_version (str, optional) – GPU的计算能力（例如"7.0"）。
Returns:: have_matrixcore – 如果支持MatrixCore则为True，否则为False
Return type:: bool

tvm.contrib.rocm.find_rocm_path()

用于查找ROCm路径的实用函数

Returns:: path – ROCm根目录的路径。
Return type:: str

tvm.contrib.spirv

用于与SPIRV工具交互的实用程序

tvm.contrib.spirv.optimize(spv_bin)

通过命令行界面使用spirv-opt优化SPIRV

请注意，spirv-opt仍处于实验阶段。

Parameters:: spv_bin (bytearray) – SPIRV文件
Returns:: cobj_bin – HSA代码对象
Return type:: bytearray

tvm.contrib.tar

用于在系统中调用tarball的工具。

tvm.contrib.tar.tar(output, files)

创建包含根目录下所有文件的压缩包。

Parameters:

output (str) – 目标共享库。
files (list) – 需要打包的文件列表。

tvm.contrib.tar.untar(tar_file, directory)

将所有tar文件解压到目录中

Parameters:

tar_file (str) – 源tar文件。
directory (str) – 目标目录

tvm.contrib.tar.normalize_file_list_by_unpacking_tars(temp, file_list)

通过解压列表中的tar文件来规范化文件列表。

当文件名是一个tar文件时，它会将其解压到临时目录中的一个唯一文件夹，并返回tar文件中的文件列表。当文件名是一个普通文件时，它会被简单地添加到列表中。

这对于解压tar包中的对象并将其转换为库非常有用。

Parameters:

temp (tvm.contrib.utils.TempDirectory) – 用于存放解压后文件的临时目录。
file_list (List[str]) – 路径列表

Returns:

ret_list – 更新后的文件列表

Return type:

列表[str]

tvm.contrib.utils

常用系统工具

exception tvm.contrib.utils.DirectoryCreatedPastAtExit: 当在atexit钩子运行后创建TempDirectory时引发此错误。

class tvm.contrib.utils.TempDirectory(custom_path=None, keep_for_debug=None)

用于在测试期间管理临时目录的辅助对象。

当超出作用域范围时自动删除目录。

classmethod set_keep_for_debug(set_to=True): 保留程序退出后的临时目录以便调试。

remove(): 移除临时目录

relpath(name)

临时目录中的相对路径

Parameters:: name (str) – 文件名。
Returns:: path – 拼接后的路径。
Return type:: str

listdir()

列出目录中的内容。

Returns:: names – 目录内容
Return type:: list

tvm.contrib.utils.tempdir(custom_path=None, keep_for_debug=None)

创建一个临时目录，退出时删除其内容。

Parameters:

custom_path (str, optional) – 手动指定临时目录的精确路径
keep_for_debug (bool) – 保留临时目录用于调试目的

Returns:

temp – 临时目录对象

Return type:

TempDirectory

class tvm.contrib.utils.FileLock(path)

文件锁对象

Parameters:: path (str) – 锁文件的路径

release(): 释放锁

tvm.contrib.utils.filelock(path)

创建一个文件锁，用于锁定路径

Parameters:: path (str) – 锁文件的路径
Returns:: 锁定
Return type:: 文件锁对象

tvm.contrib.utils.is_source_path(path)

检查路径是否为源代码路径。

Parameters:: path (str) – 可能的路径
Returns:: valid – 路径是否为可能的源路径
Return type:: bool

tvm.contrib.utils.which(exec_name)

尝试查找exec_name的完整路径

Parameters:: exec_name (str) – 可执行文件名
Returns:: path – 如果找到可执行文件的完整路径则返回该路径，否则返回None
Return type:: str

tvm.contrib.xcode

调用Xcode编译器工具链的实用工具

tvm.contrib.xcode.xcrun(cmd)

运行xcrun并返回输出结果。

Parameters:: cmd (list of str) – 命令序列。
Returns:: out – 输出字符串。
Return type:: str

tvm.contrib.xcode.create_dylib(output, objects, arch, sdk='macosx', min_os_version=None)

创建动态库。

Parameters:

output (str) – 目标共享库。
objects (list) – 对象文件列表。
options (str) – 附加选项。
arch (str) – 目标主要架构
sdk (str) – 要使用的SDK。

tvm.contrib.xcode.compile_metal(code, path_target=None, sdk='macosx', min_os_version=None)

使用环境中的CLI工具编译metal。

Parameters:

code (str) – CUDA代码。
path_target (str, optional) – 输出文件。
sdk (str, optional) – 目标平台SDK。

Returns:

metallib – metallib的字节数组

Return type:

bytearray

tvm.contrib.xcode.compile_coreml(model, model_name='main', out_dir='.'): 编译coreml模型并返回编译后的模型路径。