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

Parameters:

lhs (Tensor) – The left matrix operand
rhs (Tensor) – The right matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

Tensor

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.

Parameters:

lhs (Tensor) – The left matrix operand
rhs (Tensor) – The right matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

Tensor

tvm.contrib.clang

Util to invoke clang in the system.

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

Find clang in system.

Parameters:: required (bool) – Whether it is required, runtime error will be raised if the compiler is required.
Returns:: valid_list – List of possible paths.
Return type:: list of str

Note

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

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

Create llvm text ir.

Parameters:

inputs (list of str) – List of input files name or code source.
output (str, optional) – Output file, if it is none a temporary file is created
options (list) – The list of additional options string.
cc (str, optional) – The clang compiler, if not specified, we will try to guess the matched clang version.

Returns:

code – The generated llvm text IR.

Return type:

str

tvm.contrib.cc

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

tvm.contrib.cc.get_cc()

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

Returns:: out – The path to the default C/C++ compiler, or None if none was found.
Return type:: Optional[str]

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

Create shared library.

Parameters:

output (str) – The target shared library.
objects (List[str]) – List of object files.
options (List[str]) – The list of additional options string.
cc (Optional[str]) – The compiler command.
cwd (Optional[str]) – The current working directory.
ccache_env (Optional[Dict[str, str]]) – The environment variable for ccache. Set None to disable ccache by default.

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

Create static library.

Parameters:

output (str) – The target shared library.
inputs (List[str]) – List of inputs files. Each input file can be a tarball of objects or an object file.
ar (Optional[str]) – Path to the ar command to be used

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

Create executable binary.

Parameters:

output (str) – The target executable.
objects (List[str]) – List of object files.
options (List[str]) – The list of additional options string.
cc (Optional[str]) – The compiler command.
cwd (Optional[str]) – The urrent working directory.
ccache_env (Optional[Dict[str, str]]) – The environment variable for ccache. Set None to disable ccache by default.

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

Get global symbols from a library via nm -g

Parameters:

path (str) – The library path
nm (str) – The path to nm command

Returns:

symbol_section_map – A map from defined global symbol to their sections

Return type:

Dict[str, str]

tvm.contrib.cc.get_target_by_dump_machine(compiler)

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

Parameters:: compiler (Optional[str]) – The compiler.
Returns:: out – A function that can get target triple according to dumpmachine option of compiler.
Return type:: Callable

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.

Parameters:

compile_func (Union[str, Callable[[str, str, Optional[str]], None]]) – Function that performs the actual compilation
options (Optional[List[str]]) – List of additional optional string.
output_format (Optional[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.

Returns:

fcompile – A compilation function that can be passed to export_library.

Return type:

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

Examples

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

External function interface to cuBLAS libraries.

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

Parameters:

lhs (Tensor) – The left matrix operand
rhs (Tensor) – The right matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

Tensor

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

Parameters:

lhs (Tensor) – The left matrix operand
rhs (Tensor) – The right matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

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

Parameters:

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

Parameters:: tvm_func (Function) – Built tvm function operating on arrays
Returns:: wrapped_func – Wrapped tvm function that operates on PyTorch tensors
Return type:: Function

tvm.contrib.emcc

Util to invoke emscripten compilers in the system.

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

Create wasm that is supposed to run with the tvmjs.

Parameters:

output (str) – The target shared library.
objects (list) – List of object files.
options (str) – The additional options.
cc (str, optional) – The compile string.
libs (list) – List of user-defined library files (e.g. .bc files) to add into the wasm.

tvm.contrib.miopen

External function interface to MIOpen library.

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

Parameters:

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

Returns:

y – The result tensor

Return type:

Tensor

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

Compute softmax with MIOpen

Parameters:

x (tvm.te.Tensor) – The input tensor
axis (int) – The axis to compute softmax over

Returns:

ret – The result tensor

Return type:

tvm.te.Tensor

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

Compute log softmax with MIOpen

Parameters:

x (tvm.te.Tensor) – The input tensor
axis (int) – The axis to compute log softmax over

Returns:

ret – The result tensor

Return type:

tvm.te.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.

Parameters:

func (Function) – A TVM function that can take positional arguments
const_loc (list 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.

Returns:

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

Return type:

Function

tvm.contrib.ndk

Util to invoke NDK compiler toolchain.

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

Create shared library.

Parameters:

output (str) – The target shared library.
objects (list) – List of object files.
options (list of str, optional) – The additional options.

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

Create static library:

Parameters:

output (str) – The target static library.
inputs (list) – List of object files or tar files

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

Get global symbols from a library via nm -gU in NDK

Parameters:

path (str) – The library path
nm (str) – The path to nm command

Returns:

symbol_section_map – A map from defined global symbol to their sections

Return type:

Dict[str, str]

tvm.contrib.nnpack

External function interface to NNPACK libraries.

tvm.contrib.nnpack.is_available(): Check whether NNPACK is available, that is, nnp_initialize() returns nnp_status_success.

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.

Parameters:

lhs (Tensor) – lhs 1D array input[input_channels] of FP32 elements
rhs (Tensor) – lhs 2D matrix kernel[output_channels][input_channels] of FP32 elements

Returns:

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

Return type:

Tensor

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.

Parameters:

data (Tensor) – data 4D tensor input[batch][input_channels][input_height][input_width] of FP32 elements.
kernel (Tensor) – kernel 4D tensor kernel[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.
bias (Tensor) – bias 1D array bias[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.
padding (list) – padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right], which indicates the padding around the feature map.
stride (list) – stride A 2-dim list of [stride_height, stride_width], which indicates the stride.

Returns:

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

Return type:

Tensor

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.

Parameters:

data (Tensor) – data 4D tensor input[batch][input_channels][input_height][input_width] of FP32 elements.
transformed_kernel (Tensor) – transformed_kernel 4D tensor kernel[output_channels][input_channels][tile] [tile] of FP32 elements.
bias (Tensor) – bias 1D array bias[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.
padding (list) – padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right], which indicates the padding around the feature map.
stride (list) – stride A 2-dim list of [stride_height, stride_width], which indicates the stride.

Returns:

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

Return type:

Tensor

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.

Parameters:: kernel (Tensor) – kernel 4D tensor kernel[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.
Returns:: output – output 4D tensor output[output_channels][input_channels][tile][tile] of FP32 elements.
Return type:: Tensor

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.

Parameters:

code (str) – The cuda code.
target_format (str) – The target format of nvcc compiler.
arch (str) – The cuda architecture.
options (str or list of str) – The additional options.
path_target (str, optional) – Output file.

Returns:

cubin – The bytearray of the cubin

Return type:

bytearray

tvm.contrib.nvcc.find_cuda_path()

Utility function to find cuda path

Returns:: path – Path to cuda root.
Return type:: str

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

Utility function to get cuda version

Parameters:: cuda_path (Optional[str]) – Path to cuda root. If None is passed, will use find_cuda_path() as default.
Returns:: version – The cuda version
Return type:: float

tvm.contrib.nvcc.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

Parameters:

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

Returns:

major (int) – major version number
minor (int) – minor version number

tvm.contrib.nvcc.have_fp16(compute_version)

Either fp16 support is provided in the compute capability or not

Parameters:: 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

Parameters:: compute_version (str) – 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

Parameters:

compute_version (str, optional) – compute capability of a GPU (e.g. “7.0”).
target (tvm.target.Target, optional) – The compilation target, will be used to determine arch if compute_version isn’t specified.

tvm.contrib.nvcc.have_cudagraph(): Either CUDA Graph support is provided

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.

Parameters:

key (str) – The file key to the function
save_at_exit (bool) – Whether save the cache to file when the program exits

property cache: Return the cache, initializing on first use.

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

Memoize the result of function and reuse multiple times.

Parameters:

key (str) – The unique key to the file
save_at_exit (bool) – Whether save the cache to file when the program exits

Returns:

fmemoize – The decorator function to perform memoization.

Return type:

function

tvm.contrib.random

External function interface to random library.

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

Parameters:

low (int) – Lowest (signed) integer to be drawn from the distribution
high (int) – One above the largest (signed) integer to be drawn from the distribution

Returns:

out – A tensor with specified size and dtype

Return type:

Tensor

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.

Parameters:

low (float) – Lower boundary of the output interval. All values generated will be greater than or equal to low.
high (float) – Upper boundary of the output interval. All values generated will be less than high.
size (tuple of ints) – Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.

Returns:

out – A tensor with specified size and dtype.

Return type:

Tensor

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

Draw samples from a normal distribution.

Return random samples from a normal distribution.

Parameters:

loc (float) – loc of the distribution.
scale (float) – Standard deviation of the distribution.
size (tuple of ints) – Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn.

Returns:

out – A tensor with specified size and dtype

Return type:

Tensor

tvm.contrib.relay_viz

Relay IR Visualizer

class tvm.contrib.relay_viz.RelayVisualizer(relay_mod: IRModule, relay_param: Dict[str, NDArray] | None = None, plotter: Plotter | None = None, parser: VizParser | None = None)

Relay IR Visualizer

Parameters:

relay_mod (tvm.IRModule) – Relay IR module.
relay_param (None | Dict[str, tvm.runtime.NDArray]) – Relay parameter dictionary. Default None.
plotter (Plotter) – An instance of class inheriting from Plotter interface. Default is an instance of terminal.TermPlotter.
parser (VizParser) – An instance of class inheriting from VizParser interface. Default is an instance of terminal.TermVizParser.

Visualize Relay IR by Graphviz DOT language.

class tvm.contrib.relay_viz.dot.DotGraph(name: str, graph_attr: Dict[str, str] | None = None, node_attr: Dict[str, str] | None = None, edge_attr: Dict[str, str] | None = None, get_node_attr: Callable[[VizNode], Dict[str, str]] | None = None)

DOT graph for relay IR.

Parameters:

name (str) – name of this graph.
graph_attr (Optional[Dict[str, str]]) – key-value pairs for the graph.
node_attr (Optional[Dict[str, str]]) – key-value pairs for all nodes.
edge_attr (Optional[Dict[str, str]]) – key-value pairs for all edges.
get_node_attr (Optional[Callable[[VizNode], Dict[str, str]]]) – A callable returning attributes for the node.

node(viz_node: VizNode) → None

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

Parameters:: viz_node (VizNode) – A VizNode instance.

edge(viz_edge: VizEdge) → None

Add an edge to the underlying graph.

Parameters:: viz_edge (VizEdge) – A VizEdge instance.

class tvm.contrib.relay_viz.dot.DotPlotter(graph_attr: Dict[str, str] | None = None, node_attr: Dict[str, str] | None = None, edge_attr: Dict[str, str] | None = None, get_node_attr: Callable[[VizNode], Dict[str, str]] | None = None, render_kwargs: Dict[str, Any] | None = None)

DOT language graph plotter

The plotter accepts various graphviz attributes for graphs, nodes, and edges. Please refer to https://graphviz.org/doc/info/attrs.html for available attributes.

Parameters:

graph_attr (Optional[Dict[str, str]]) – key-value pairs for all graphs.
node_attr (Optional[Dict[str, str]]) – key-value pairs for all nodes.
edge_attr (Optional[Dict[str, str]]) – key-value pairs for all edges.
get_node_attr (Optional[Callable[[VizNode], Dict[str, str]]]) – A callable returning attributes for a specific node.
render_kwargs (Optional[Dict[str, Any]]) – keyword arguments directly passed to graphviz.Digraph.render().

Examples

from tvm.contrib import relay_viz
from tvm.relay.testing import resnet

mod, param = resnet.get_workload(num_layers=18)
# graphviz attributes
graph_attr = {"color": "red"}
node_attr = {"color": "blue"}
edge_attr = {"color": "black"}

# VizNode is passed to the callback.
# We want to color NCHW conv2d nodes. Also give Var a different shape.
def get_node_attr(node):
    if "nn.conv2d" in node.type_name and "NCHW" in node.detail:
        return {
            "fillcolor": "green",
            "style": "filled",
            "shape": "box",
        }
    if "Var" in node.type_name:
        return {"shape": "ellipse"}
    return {"shape": "box"}

# Create plotter and pass it to viz. Then render the graph.
dot_plotter = relay_viz.DotPlotter(
    graph_attr=graph_attr,
    node_attr=node_attr,
    edge_attr=edge_attr,
    get_node_attr=get_node_attr)

viz = relay_viz.RelayVisualizer(
    mod,
    relay_param=param,
    plotter=dot_plotter,
    parser=relay_viz.DotVizParser())
viz.render("hello")

create_graph(name)

Create a VizGraph

Parameters:: name (str) – the name of the graph
Returns:: rv1
Return type:: an instance of class inheriting from VizGraph interface.

render(filename: str | None = None)

render the graph generated from the Relay IR module.

This function is a thin wrapper of graphviz.Digraph.render().

Visualize Relay IR in AST text-form.

class tvm.contrib.relay_viz.terminal.TermVizParser

TermVizParser parse nodes and edges for TermPlotter.

get_node_edges(node: RelayExpr, relay_param: Dict[str, NDArray], node_to_id: Dict[RelayExpr, str]) → Tuple[VizNode | None, List[VizEdge]]: Parse a node and edges from a relay.Expr.

class tvm.contrib.relay_viz.terminal.TermNode(viz_node: VizNode): TermNode is aimed to generate text more suitable for terminal visualization.

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

Terminal graph for a relay IR Module

Parameters:: name (str) – name of this graph.

node(viz_node: VizNode) → None

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

Parameters:: viz_node (VizNode) – A VizNode instance.

edge(viz_edge: VizEdge) → None

Add an edge to the terminal graph.

Parameters:: viz_edge (VizEdge) – A VizEdge instance.

render() → str

Draw a terminal graph

Returns:: rv1 – text representing a graph.
Return type:: str

class tvm.contrib.relay_viz.terminal.TermPlotter

Terminal plotter

create_graph(name)

Create a VizGraph

Parameters:: name (str) – the name of the graph
Returns:: rv1
Return type:: an instance of class inheriting from VizGraph interface.

render(filename): If filename is None, print to stdio. Otherwise, write to the filename.

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

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

VizNode carry node information for VizGraph interface.

Parameters:

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.VizEdge(start_node: str, end_node: str)

VizEdge connect two VizNode.

Parameters:

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.VizParser

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

abstract get_node_edges(node: RelayExpr, relay_param: Dict[str, NDArray], node_to_id: Dict[RelayExpr, str]) → Tuple[VizNode | None, List[VizEdge]]

Get VizNode and VizEdges for a relay.Expr.

Parameters:

node (relay.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_id (Dict[relay.Expr, str]) – This is a mapping from relay.Expr to a unique id, generated by RelayVisualizer.

Returns:

rv1 (Union[VizNode, None]) – VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node to the graph, return None.
rv2 (List[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.VizGraph

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

abstract node(viz_node: VizNode) → None

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

Parameters:: viz_node (VizNode) – A VizNode instance.

abstract edge(viz_edge: VizEdge) → None

Add an edge to the underlying graph.

Parameters:: viz_edge (VizEdge) – A VizEdge instance.

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

get_node_edges(node: RelayExpr, relay_param: Dict[str, NDArray], node_to_id: Dict[RelayExpr, str]) → Tuple[VizNode | None, List[VizEdge]]

Get VizNode and VizEdges for a relay.Expr.

Parameters:

node (relay.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_id (Dict[relay.Expr, str]) – This is a mapping from relay.Expr to a unique id, generated by RelayVisualizer.

Returns:

rv1 (Union[VizNode, None]) – VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node to the graph, return None.
rv2 (List[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) → VizGraph

Create a VizGraph

Parameters:: name (str) – the name of the graph
Returns:: rv1
Return type:: an instance of class inheriting from VizGraph interface.

abstract render(filename: str) → None

Render the graph as a file.

Parameters:: filename (str) – see the definition of implemented class.

tvm.contrib.rocblas

External function interface to rocBLAS libraries.

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

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

Parameters:

lhs (Tensor) – The left matrix operand
rhs (Tensor) – The right matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

Tensor

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

Parameters:

lhs (Tensor) – The left batched matrix operand
rhs (Tensor) – The right batched matrix operand
transa (bool) – Whether transpose lhs
transb (bool) – Whether transpose rhs

Returns:

C – The result tensor.

Return type:

Tensor

tvm.contrib.rocm

Utility for ROCm backend

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

Find ld.lld in system.

Parameters:: required (bool) – Whether it is required, runtime error will be raised if the compiler is required.
Returns:: valid_list – List of possible paths.
Return type:: list of str

Note

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

Parameters:

in_file (str) – Input file name (relocatable ELF object file)
out_file (str) – Output file name (shared ELF object file)
lld (str, optional) – The lld linker, if not specified, we will try to guess the matched clang version.

tvm.contrib.rocm.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

Parameters:

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

Returns:

major (int) – major version number
minor (int) – minor version number

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

Either MatrixCore support is provided in the compute capability or not

Parameters:: compute_version (str, optional) – compute capability of a GPU (e.g. “7.0”).
Returns:: have_matrixcore – True if MatrixCore support is provided, False otherwise
Return type:: bool

tvm.contrib.rocm.find_rocm_path()

Utility function to find ROCm path

Returns:: path – Path to ROCm root.
Return type:: str

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.

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.

tvm.contrib.sparse.array(source_array, device=None, shape=None, stype='csr'): Construct a sparse NDArray from numpy.ndarray

class tvm.contrib.sparse.SparsePlaceholderOp(shape, nonzeros, dtype, name): Placeholder class for sparse tensor representations.

class tvm.contrib.sparse.CSRPlaceholderOp(shape, nonzeros, dtype, name): Placeholder class for CSR based sparse tensor representation.

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

Construct an empty sparse tensor object.

Parameters:

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)

Returns:

tensor – The created sparse tensor placeholder

Return type:

SparsePlaceholderOp

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.

Parameters:: spv_bin (bytearray) – The spirv file
Returns:: cobj_bin – The HSA Code Object
Return type:: bytearray

tvm.contrib.tar

Util to invoke tarball in the system.

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

Create tarball containing all files in root.

Parameters:

output (str) – The target shared library.
files (list) – List of files to be bundled.

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

Unpack all tar files into the directory

Parameters:

tar_file (str) – The source tar file.
directory (str) – The target directory

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

Normalize the file list by unpacking tars in list.

When a filename is a tar, it will untar it into an unique dir in temp and return the list of files in the tar. When a filename is a normal file, it will be simply added to the list.

This is useful to untar objects in tar and then turn them into a library.

Parameters:

temp (tvm.contrib.utils.TempDirectory) – A temp dir to hold the untared files.
file_list (List[str]) – List of path

Returns:

ret_list – An updated list of files

Return type:

List[str]

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

classmethod set_keep_for_debug(set_to=True): Keep temporary directories past program exit for debugging.

remove(): Remove the tmp dir

relpath(name)

Relative path in temp dir

Parameters:: name (str) – The name of the file.
Returns:: path – The concatenated path.
Return type:: str

listdir()

List contents in the dir.

Returns:: names – The content of directory
Return type:: list

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

Create temp dir which deletes the contents when exit.

Parameters:

custom_path (str, optional) – Manually specify the exact temp dir path
keep_for_debug (bool) – Keep temp directory for debugging purposes

Returns:

temp – The temp directory object

Return type:

TempDirectory

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

File lock object

Parameters:: path (str) – The path to the lock

release(): Release the lock

tvm.contrib.utils.filelock(path)

Create a file lock which locks on path

Parameters:: path (str) – The path to the lock
Returns:: lock
Return type:: File lock object

tvm.contrib.utils.is_source_path(path)

Check if path is source code path.

Parameters:: path (str) – A possible path
Returns:: valid – Whether path is a possible source path
Return type:: bool

tvm.contrib.utils.which(exec_name)

Try to find full path of exec_name

Parameters:: exec_name (str) – The executable name
Returns:: path – The full path of executable if found, otherwise returns None
Return type:: str

tvm.contrib.xcode

Utility to invoke Xcode compiler toolchain

tvm.contrib.xcode.xcrun(cmd)

Run xcrun and return the output.

Parameters:: cmd (list of str) – The command sequence.
Returns:: out – The output string.
Return type:: str

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

Create dynamic library.

Parameters:

output (str) – The target shared library.
objects (list) – List of object files.
options (str) – The additional options.
arch (str) – Target major architectures
sdk (str) – The sdk to be used.

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

Compile metal with CLI tool from env.

Parameters:

code (str) – The cuda code.
path_target (str, optional) – Output file.
sdk (str, optional) – The target platform SDK.

Returns:

metallib – The bytearray of the metallib

Return type:

bytearray

tvm.contrib.xcode.compile_coreml(model, model_name='main', out_dir='.'): Compile coreml model and return the compiled model path.