tvm.contrib

Contrib APIs of TVM python package.

Contrib API provides many useful not core features. Some of these are useful utilities to interact with thirdparty libraries and tools.

tvm.contrib.cblas

External function interface to BLAS libraries.

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

Create an extern op that compute batched matrix mult of A and rhs with CBLAS This function serves as an example on how to call external libraries.

lhs: Tensor

The left matrix operand

rhs: Tensor

The right matrix operand

transa: bool

Whether transpose lhs

transb: bool

Whether transpose rhs

C: Tensor

The result tensor.

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

Create an extern op that compute matrix mult of A and rhs with CrhsLAS This function serves as an example on how to call external libraries.

lhs: Tensor

The left matrix operand

rhs: Tensor

The right matrix operand

transa: bool

Whether transpose lhs

transb: bool

Whether transpose rhs

C: Tensor

The result tensor.

tvm.contrib.clang

Util to invoke clang in the system.

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

Create llvm text ir.

inputslist of str

List of input files name or code source.

outputstr, optional

Output file, if it is none a temporary file is created

optionslist

The list of additional options string.

ccstr, optional

The clang compiler, if not specified, we will try to guess the matched clang version.

codestr

The generated llvm text IR.

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

Find clang in system.

requiredbool

Whether it is required, runtime error will be raised if the compiler is required.

valid_listlist of str

List of possible paths.

This function will first search clang that matches the major llvm version that built with tvm

tvm.contrib.cc

Util to invoke C/C++ compilers in the system.

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

Create executable binary.

outputstr

The target executable.

objectsList[str]

List of object files.

optionsList[str]

The list of additional options string.

ccOptional[str]

The compiler command.

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

Create shared library.

outputstr

The target shared library.

objectsList[str]

List of object files.

optionsList[str]

The list of additional options string.

ccOptional[str]

The compiler command.

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

Create a cross compiler function by specializing compile_func with options.

This function can be used to construct compile functions that can be passed to AutoTVM measure or export_library.

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

Function that performs the actual compilation

optionsOptional[List[str]]

List of additional optional string.

output_formatOptional[str]

Library output format.

get_target_triple: Optional[Callable]

Function that can target triple according to dumpmachine option of compiler.

add_files: Optional[List[str]]

List of paths to additional object, source, library files to pass as part of the compilation.

fcompileCallable[[str, str, Optional[str]], None]

A compilation function that can be passed to export_library.

from tvm.contrib import cc, ndk
# export using arm gcc
mod = build_runtime_module()
mod.export_library(path_dso,
                   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, specialized_ndk)

tvm.contrib.cc.get_cc()

Return the path to the default C/C++ compiler.

out: Optional[str]

The path to the default C/C++ compiler, or None if none was found.

tvm.contrib.cc.get_target_by_dump_machine(compiler)

Functor of get_target_triple that can get the target triple using compiler.

compilerOptional[str]

The compiler.

out: Callable

A function that can get target triple according to dumpmachine option of compiler.

tvm.contrib.cublas

External function interface to cuBLAS libraries.

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

Create an extern op that compute batch matrix mult of A and rhs with cuBLAS

lhsTensor

The left matrix operand

rhsTensor

The right matrix operand

transabool

Whether transpose lhs

transbbool

Whether transpose rhs

CTensor

The result tensor.

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

Create an extern op that compute matrix mult of A and rhs with cuBLAS

lhsTensor

The left matrix operand

rhsTensor

The right matrix operand

transabool

Whether transpose lhs

transbbool

Whether transpose rhs

CTensor

The result tensor.

tvm.contrib.dlpack

Wrapping functions to bridge frameworks with DLPack support to 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

framework, provided the other framework supports DLPACK

tvm_func: Function

Built tvm function operating on arrays

tensor_type: Type

Type of the tensors of the target framework

to_dlpack_func: Function

Function to convert the source tensors to DLPACK

tvm.contrib.dlpack.to_pytorch_func(tvm_func)

Convert a tvm function into one that accepts PyTorch tensors

tvm_func: Function

Built tvm function operating on arrays

wrapped_func: Function

Wrapped tvm function that operates on PyTorch tensors

tvm.contrib.emcc

Util to invoke emscripten compilers in the system.

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

Create wasm that is supposed to run with the tvmjs.

outputstr

The target shared library.

objectslist

List of object files.

optionsstr

The additional options.

ccstr, optional

The compile string.

tvm.contrib.miopen

External function interface to MIOpen library.

tvm.contrib.miopen._get_np_int32_array_handle(arr)

Return a void_p handle for a numpy array

arr: numpy.NDArray

source numpy array

ptr: ctypes.c_void_p

pointer to the data

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)

Create an extern op that compute 2D convolution with MIOpen

x: Tensor

input feature map

w: Tensor

convolution weight

stride_h: int

height stride

stride_w: int

width stride

pad_h: int

height pad

pad_w: int

weight pad

dilation_h: int

height dilation

dilation_w: int

width dilation

conv_mode: int

0: miopenConvolution 1: miopenTranspose

data_type: int

0: miopenHalf (fp16) 1: miopenFloat (fp32)

group_count: int

number of groups

y: Tensor

The result tensor

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

Compute log softmax with MIOpen

xtvm.te.Tensor

The input tensor

axisint

The axis to compute log softmax over

rettvm.te.Tensor

The result tensor

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

Compute softmax with MIOpen

xtvm.te.Tensor

The input tensor

axisint

The axis to compute softmax over

rettvm.te.Tensor

The result tensor

tvm.contrib.mxnet

MXNet bridge wrap Function MXNet’s async function.

tvm.contrib.mxnet.to_mxnet_func(func, const_loc=None)

Wrap a TVM function as MXNet function

MXNet function runs asynchrously via its engine.

funcFunction

A TVM function that can take positional arguments

const_loclist of int

List of integers indicating the argument position of read only NDArray argument. The NDArray argument location that are not annotated will be viewed as mutable arrays in MXNet’s engine.

async_funcFunction

A function that can take MXNet NDArray as argument in places that used to expect TVM NDArray. Run asynchrously in MXNet’s async engine.

tvm.contrib.ndk

Util to invoke NDK compiler toolchain.

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

Create shared library.

outputstr

The target shared library.

objectslist

List of object files.

optionslist of str, optional

The additional options.

tvm.contrib.nnpack

External function interface to NNPACK libraries.

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

Create an extern op to do inference convolution of 4D tensor data and 4D tensor kernel and 1D tensor bias with nnpack.

dataTensor

data 4D tensor input[batch][input_channels][input_height][input_width] of FP32 elements.

kernelTensor

kernel 4D tensor kernel[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.

biasTensor

bias 1D array bias[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.

paddinglist

padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right], which indicates the padding around the feature map.

stridelist

stride A 2-dim list of [stride_height, stride_width], which indicates the stride.

outputTensor

output 4D tensor output[batch][output_channels][output_height][output_width] of FP32 elements.

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

Create an extern op to do inference convolution of 3D tensor data and 4D tensor kernel and 1D tensor bias with nnpack.

kernelTensor

kernel 4D tensor kernel[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.

outputTensor

output 4D tensor output[output_channels][input_channels][tile][tile] of FP32 elements.

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

Create an extern op to do inference convolution of 4D tensor data and 4D pre-transformed tensor kernel and 1D tensor bias with nnpack.

dataTensor

data 4D tensor input[batch][input_channels][input_height][input_width] of FP32 elements.

transformed_kernelTensor

transformed_kernel 4D tensor kernel[output_channels][input_channels][tile] [tile] of FP32 elements.

biasTensor

bias 1D array bias[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.

paddinglist

padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right], which indicates the padding around the feature map.

stridelist

stride A 2-dim list of [stride_height, stride_width], which indicates the stride.

outputTensor

output 4D tensor output[batch][output_channels][output_height][output_width] of FP32 elements.

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

Create an extern op that compute fully connected of 1D tensor lhs and 2D tensor rhs with nnpack.

lhsTensor

lhs 1D array input[input_channels] of FP32 elements

rhsTensor

lhs 2D matrix kernel[output_channels][input_channels] of FP32 elements

CTensor

lhs 1D array out[output_channels] of FP32 elements.

tvm.contrib.nnpack.is_available()

Check whether NNPACK is available, that is, nnp_initialize() returns nnp_status_success.

tvm.contrib.nvcc

Utility to invoke nvcc compiler in the system

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

Compile cuda code with NVCC from env.

codestr

The cuda code.

target_formatstr

The target format of nvcc compiler.

archstr

The cuda architecture.

optionsstr or list of str

The additional options.

path_targetstr, optional

Output file.

cubinbytearray

The bytearray of the cubin

tvm.contrib.nvcc.find_cuda_path()

Utility function to find cuda path

pathstr

Path to cuda root.

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

Utility function to get cuda version

cuda_path : Optional[str]

Path to cuda root. If None is passed, will use find_cuda_path() as default.

versionfloat

The cuda version

tvm.contrib.nvcc.get_target_compute_version(target=None)

Utility function to get compute capability of compilation target.

Looks for the target arch in three different places, first in the target input, then the Target.current() scope, and finally the GPU device (if it exists).

targettvm.target.Target, optional

The compilation target

compute_versionstr

compute capability of a GPU (e.g. “8.6”)

tvm.contrib.nvcc.have_bf16(compute_version)

Either bf16 support is provided in the compute capability or not

compute_versionstr

compute capability of a GPU (e.g. “8.0”)

tvm.contrib.nvcc.have_cudagraph()

Either CUDA Graph support is provided

tvm.contrib.nvcc.have_fp16(compute_version)

Either fp16 support is provided in the compute capability or not

compute_version: str

compute capability of a GPU (e.g. “6.0”)

tvm.contrib.nvcc.have_int8(compute_version)

Either int8 support is provided in the compute capability or not

compute_versionstr

compute capability of a GPU (e.g. “6.1”)

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

Either TensorCore support is provided in the compute capability or not

compute_versionstr, optional

compute capability of a GPU (e.g. “7.0”).

targettvm.target.Target, optional

The compilation target, will be used to determine arch if compute_version isn’t specified.

tvm.contrib.nvcc.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

compute_versionstr

compute capability of a GPU (e.g. “6.0”)

majorint

major version number

minorint

minor version number

tvm.contrib.pickle_memoize

Memoize result of function via pickle, used for cache testcases.

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

A cache object for result cache.

key: str

The file key to the function

save_at_exit: bool

Whether save the cache to file when the program exits

tvm.contrib.pickle_memoize._atexit()

Save handler.

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

Memoize the result of function and reuse multiple times.

key: str

The unique key to the file

save_at_exit: bool

Whether save the cache to file when the program exits

fmemoizefunction

The decorator function to perform memoization.

tvm.contrib.random

External function interface to random library.

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

Draw samples from a normal distribution.

Return random samples from a normal distribution.

locfloat

loc of the distribution.

scalefloat

Standard deviation of the distribution.

sizetuple of ints

Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.

outTensor

A tensor with specified size and dtype

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

Return random integers from low (inclusive) to high (exclusive). Return random integers from the “discrete uniform” distribution of the specified dtype in the “half-open” interval [low, high).

lowint

Lowest (signed) integer to be drawn from the distribution

highint

One above the largest (signed) integer to be drawn from the distribution

outTensor

A tensor with specified size and dtype

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

Draw samples from a uniform distribution.

Samples are uniformly distributed over the half-open interval [low, high) (includes low, but excludes high). In other words, any value within the given interval is equally likely to be drawn by uniform.

lowfloat

Lower boundary of the output interval. All values generated will be greater than or equal to low.

highfloat

Upper boundary of the output interval. All values generated will be less than high.

sizetuple of ints

Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.

outTensor

A tensor with specified size and dtype.

tvm.contrib.relay_viz

Visualize Relay IR in AST text-form.

class tvm.contrib.relay_viz.terminal.TermGraph(name: str)

Terminal graph for a relay IR Module

name: str

name of this graph.

edge(viz_edge: tvm.contrib.relay_viz.interface.VizEdge) → None

Add an edge to the terminal graph.

viz_edgeVizEdge

A VizEdge instance.

node(viz_node: tvm.contrib.relay_viz.interface.VizNode) → None

Add a node to the underlying graph. Nodes in a Relay IR Module are expected to be added in the post-order.

viz_nodeVizNode

A VizNode instance.

render() → str

Draw a terminal graph

rv1: str

text representing a graph.

class tvm.contrib.relay_viz.terminal.TermNode(viz_node: tvm.contrib.relay_viz.interface.VizNode)

TermNode is aimed to generate text more suitable for terminal visualization.

class tvm.contrib.relay_viz.terminal.TermPlotter

Terminal plotter

create_graph(name)

Create a VizGraph

namestr

the name of the graph

rv1: an instance of class inheriting from VizGraph interface.

render(filename)

If filename is None, print to stdio. Otherwise, write to the filename.

class tvm.contrib.relay_viz.terminal.TermVizParser

TermVizParser parse nodes and edges for TermPlotter.

get_node_edges(node: tvm.ir.expr.RelayExpr, relay_param: Dict[str, tvm.runtime.ndarray.NDArray], node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Parse a node and edges from a relay.Expr.

Abstract class used by tvm.contrib.relay_viz.RelayVisualizer.

class tvm.contrib.relay_viz.interface.DefaultVizParser

DefaultVizParser provde a set of logics to parse a various relay types. These logics are inspired and heavily based on visualize function in https://tvm.apache.org/2020/07/14/bert-pytorch-tvm

_call(node: tvm.ir.expr.RelayExpr, node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Render rule for a relay call node

_function(node: tvm.ir.expr.RelayExpr, node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Render rule for a relay function node

_var(node: tvm.ir.expr.RelayExpr, relay_param: Dict[str, tvm.runtime.ndarray.NDArray], node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Render rule for a relay var node

get_node_edges(node: tvm.ir.expr.RelayExpr, relay_param: Dict[str, tvm.runtime.ndarray.NDArray], node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Get VizNode and VizEdges for a relay.Expr.

noderelay.Expr

relay.Expr which will be parsed and generate a node and edges.

relay_param: Dict[str, tvm.runtime.NDArray]

relay parameters dictionary.

node_to_idDict[relay.Expr, str]

This is a mapping from relay.Expr to a unique id, generated by RelayVisualizer.

rv1Union[VizNode, None]

VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node to the graph, return None.

rv2List[VizEdge]

a list of VizEdges to describe the connectivity of the relay.Expr. Can be empty list to indicate no connectivity.

class tvm.contrib.relay_viz.interface.Plotter

Plotter can render a collection of Graph interfaces to a file.

abstract create_graph(name: str) → tvm.contrib.relay_viz.interface.VizGraph

Create a VizGraph

namestr

the name of the graph

rv1: an instance of class inheriting from VizGraph interface.

abstract render(filename: str) → None

Render the graph as a file.

filenamestr

see the definition of implemented class.

class tvm.contrib.relay_viz.interface.VizEdge(start_node: str, end_node: str)

VizEdge connect two VizNode.

start_node: str

The identifier of the node starting the edge.

end_node: str

The identifier of the node ending the edge.

class tvm.contrib.relay_viz.interface.VizGraph

Abstract class for graph, which is composed of nodes and edges.

abstract edge(viz_edge: tvm.contrib.relay_viz.interface.VizEdge) → None

Add an edge to the underlying graph.

viz_edgeVizEdge

A VizEdge instance.

abstract node(viz_node: tvm.contrib.relay_viz.interface.VizNode) → None

Add a node to the underlying graph. Nodes in a Relay IR Module are expected to be added in the post-order.

viz_nodeVizNode

A VizNode instance.

class tvm.contrib.relay_viz.interface.VizNode(node_id: str, node_type: str, node_detail: str)

VizNode carry node information for VizGraph interface.

node_id: str

Unique identifier for this node.

node_type: str

Type of this node.

node_detail: str

Any supplement for this node such as attributes.

class tvm.contrib.relay_viz.interface.VizParser

VizParser parses out a VizNode and VizEdges from a relay.Expr.

abstract get_node_edges(node: tvm.ir.expr.RelayExpr, relay_param: Dict[str, tvm.runtime.ndarray.NDArray], node_to_id: Dict[tvm.ir.expr.RelayExpr, str]) → Tuple[Optional[tvm.contrib.relay_viz.interface.VizNode], List[tvm.contrib.relay_viz.interface.VizEdge]]

Get VizNode and VizEdges for a relay.Expr.

noderelay.Expr

relay.Expr which will be parsed and generate a node and edges.

relay_param: Dict[str, tvm.runtime.NDArray]

relay parameters dictionary.

node_to_idDict[relay.Expr, str]

This is a mapping from relay.Expr to a unique id, generated by RelayVisualizer.

rv1Union[VizNode, None]

VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node to the graph, return None.

rv2List[VizEdge]

a list of VizEdges to describe the connectivity of the relay.Expr. Can be empty list to indicate no connectivity.

tvm.contrib.rocblas

External function interface to rocBLAS libraries.

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

Create an extern op that compute matrix mult of A and rhs with rocBLAS

lhsTensor

The left batched matrix operand

rhsTensor

The right batched matrix operand

transabool

Whether transpose lhs

transbbool

Whether transpose rhs

CTensor

The result tensor.

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

Create an extern op that compute matrix mult of A and rhs with rocBLAS

lhsTensor

The left matrix operand

rhsTensor

The right matrix operand

transabool

Whether transpose lhs

transbbool

Whether transpose rhs

CTensor

The result tensor.

tvm.contrib.rocm

Utility for ROCm backend

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

Find ld.lld in system.

requiredbool

Whether it is required, runtime error will be raised if the compiler is required.

valid_listlist of str

List of possible paths.

This function will first search ld.lld that matches the major llvm version that built with tvm

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

Link relocatable ELF object to shared ELF object using lld

in_filestr

Input file name (relocatable ELF object file)

out_filestr

Output file name (shared ELF object file)

lldstr, optional

The lld linker, if not specified, we will try to guess the matched clang version.

tvm.contrib.sparse

Tensor and Operation class for computation declaration.

class tvm.contrib.sparse.CSRNDArray(arg1, device=None, shape=None)

Sparse tensor object in CSR format.

__init__(arg1, device=None, shape=None)

Construct a sparse matrix in CSR format.

arg1numpy.ndarray or a tuple with (data, indices, indptr)

The corresponding a dense numpy array, or a tuple for constructing a sparse matrix directly.

device: Device

The corresponding device.

shapetuple of int

The shape of the array

asnumpy()

Construct a full matrix and convert it to numpy array. This API will be deprecated in TVM v0.8 release. Please use numpy instead.

numpy()

Construct a full matrix and convert it to numpy array.

class tvm.contrib.sparse.CSRPlaceholderOp(shape, nonzeros, dtype, name)

Placeholder class for CSR based sparse tensor representation.

__init__(shape, nonzeros, dtype, name)

Contructing a bare bone structure for a csr_matrix

shape: Tuple of Expr

The shape of the tensor

nonzeros: int

The number of non-zero values

dtype: str, optional

The data type of the tensor

name: str, optional

The name hint of the tensor

class tvm.contrib.sparse.SparsePlaceholderOp(shape, nonzeros, dtype, name)

Placeholder class for sparse tensor representations.

__init__(shape, nonzeros, dtype, name)

Contructing a bare bone structure for a sparse matrix

shape: Tuple of Expr

The shape of the tensor

nonzeros: int

The number of non-zero values

dtype: str, optional

The data type of the tensor

name: str, optional

The name hint of the tensor

tvm.contrib.sparse.array(source_array, device=None, shape=None, stype='csr')

Construct a sparse NDArray from numpy.ndarray

tvm.contrib.sparse.placeholder(shape, nonzeros=None, dtype=None, name='placeholder', stype=None)

Construct an empty sparse tensor object.

shape: Tuple of Expr

The shape of the tensor

nonzeros: int

The number of non-zero values

dtype: str, optional

The data type of the tensor

name: str, optional

The name hint of the tensor

stype: str, optional

The name storage type of the sparse tensor (e.g. csr, coo, ell)

tensor: SparsePlaceholderOp

The created sparse tensor placeholder

tvm.contrib.spirv

Utility for Interacting with SPIRV Tools

tvm.contrib.spirv.optimize(spv_bin)

Optimize SPIRV using spirv-opt via CLI

Note that the spirv-opt is still experimental.

spv_binbytearray

The spirv file

cobj_binbytearray

The HSA Code Object

tvm.contrib.tar

Util to invoke tarball in the system.

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

Create tarball containing all files in root.

outputstr

The target shared library.

fileslist

List of files to be bundled.

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

Unpack all tar files into the directory

tar_filestr

The source tar file.

directorystr

The target directory

tvm.contrib.utils

Common system utilities

exception tvm.contrib.utils.DirectoryCreatedPastAtExit

Raised when a TempDirectory is created after the atexit hook runs.

class tvm.contrib.utils.FileLock(path)

File lock object

pathstr

The path to the lock

release()

Release the lock

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

Helper object to manage temp directory during testing.

Automatically removes the directory when it went out of scope.

listdir()

List contents in the dir.

nameslist

The content of directory

relpath(name)

Relative path in temp dir

namestr

The name of the file.

pathstr

The concatenated path.

remove()

Remove the tmp dir

classmethod set_keep_for_debug(set_to=True)

Keep temporary directories past program exit for debugging.

tvm.contrib.utils.filelock(path)

Create a file lock which locks on path

pathstr

The path to the lock

lock : File lock object

tvm.contrib.utils.is_source_path(path)

Check if path is source code path.

pathstr

A possible path

validbool

Whether path is a possible source path

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

Create temp dir which deletes the contents when exit.

custom_pathstr, optional

Manually specify the exact temp dir path

keep_for_debugbool

Keep temp directory for debugging purposes

tempTempDirectory

The temp directory object

tvm.contrib.utils.which(exec_name)

Try to find full path of exec_name

exec_namestr

The executable name

pathstr

The full path of executable if found, otherwise returns None

tvm.contrib.xcode

Utility to invoke Xcode compiler toolchain

tvm.contrib.xcode.compile_coreml(model, model_name='main', out_dir='.')

Compile coreml model and return the compiled model path.

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

Compile metal with CLI tool from env.

codestr

The cuda code.

path_targetstr, optional

Output file.

sdkstr, optional

The target platform SDK.

metallibbytearray

The bytearray of the metallib

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

Create dynamic library.

outputstr

The target shared library.

objectslist

List of object files.

optionsstr

The additional options.

archstr

Target major architectures

sdkstr

The sdk to be used.

tvm.contrib.xcode.xcrun(cmd)

Run xcrun and return the output.

cmdlist of str

The command sequence.

outstr

The output string.