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 -- The result tensor.

返回类型:

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 -- The result tensor.

返回类型:

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.

参数:

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

返回:

code -- The generated llvm text IR.

返回类型:

str

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

Find clang in system.

参数:

required (bool) -- Whether it is required, runtime error will be raised if the compiler is required.

返回:

valid_list -- List of possible paths.

返回类型:

list of str

备注

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, cwd=None, ccache_env=None)

Create executable binary.

参数:

• 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.create_shared(output, objects, options=None, cc=None, cwd=None, ccache_env=None)

Create shared library.

参数:

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

参数:

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

返回:

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

返回类型:

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

示例

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.cc.get_cc()

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

返回:

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

返回类型:

Optional[str]

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

Get global symbols from a library via nm -g

参数:

• path (str) -- The library path

• nm (str) -- The path to nm command

返回:

symbol_section_map -- A map from defined global symbol to their sections

返回类型:

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.

参数:

compiler (Optional[str]) -- The compiler.

返回:

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

返回类型:

Callable

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

参数:

• lhs (Tensor) -- The left matrix operand

• rhs (Tensor) -- The right matrix operand

• transa (bool) -- Whether transpose lhs

• transb (bool) -- Whether transpose rhs

返回:

C -- The result tensor.

返回类型:

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

参数:

• lhs (Tensor) -- The left matrix operand

• rhs (Tensor) -- The right matrix operand

• transa (bool) -- Whether transpose lhs

• transb (bool) -- Whether transpose rhs

返回:

C -- The result tensor.

返回类型:

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 -- Wrapped tvm function that operates on PyTorch tensors

返回类型:

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.

参数:

• 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

参数:

• 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 -- The result tensor

返回类型:

Tensor

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

Compute log softmax with MIOpen

参数:

• x (tvm.te.Tensor) -- The input tensor

• axis (int) -- The axis to compute log softmax over

返回:

ret -- The result tensor

返回类型:

tvm.te.Tensor

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

Compute softmax with MIOpen

参数:

• x (tvm.te.Tensor) -- The input tensor

• axis (int) -- The axis to compute softmax over

返回:

ret -- The result tensor

返回类型:

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.

参数:

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

返回:

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.

返回类型:

Function

tvm.contrib.ndk

Util to invoke NDK compiler toolchain.

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

Create shared library.

参数:

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

参数:

• 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

参数:

• path (str) -- The library path

• nm (str) -- The path to nm command

返回:

symbol_section_map -- A map from defined global symbol to their sections

返回类型:

Dict[str, str]

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.

参数:

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

返回:

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

返回类型:

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.

参数:

kernel (Tensor) -- kernel 4D tensor kernel[output_channels][input_channels][kernel_height] [kernel_width] of FP32 elements.

返回:

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

返回类型:

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.

参数:

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

返回:

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

返回类型:

Tensor

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.

参数:

• lhs (Tensor) -- lhs 1D array input[input_channels] of FP32 elements

• rhs (Tensor) -- lhs 2D matrix kernel[output_channels][input_channels] of FP32 elements

返回:

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

返回类型:

Tensor

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.

参数:

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

返回:

cubin -- The bytearray of the cubin

返回类型:

bytearray

tvm.contrib.nvcc.find_cuda_path()

Utility function to find cuda path

返回:

path -- Path to cuda root.

返回类型:

str

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.

返回:

version -- The cuda version

返回类型:

float

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

参数:

• 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.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

参数:

compute_version (str) -- compute capability of a GPU (e.g. "6.0")

返回:

• major (int) -- major version number

• minor (int) -- 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

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.

参数:

• key (str) -- The unique key to the file

• save_at_exit (bool) -- Whether save the cache to file when the program exits

返回:

fmemoize -- The decorator function to perform memoization.

返回类型:

function

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.

参数:

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

返回:

out -- A tensor with specified size and dtype

返回类型:

Tensor

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

参数:

• 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

返回:

out -- A tensor with specified size and dtype

返回类型:

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.

参数:

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

返回:

out -- A tensor with specified size and dtype.

返回类型:

Tensor

tvm.contrib.relay_viz

Relay IR Visualizer

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

Relay IR Visualizer

参数:

• 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, node_attr: Dict[str, str] = None, edge_attr: Dict[str, str] = None, get_node_attr: Callable[[VizNode], Dict[str, str]] = None)

DOT graph for relay IR.

See also tvm.contrib.relay_viz.dot.DotPlotter

参数:

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

edge(viz_edge: VizEdge) → None

Add an edge to the underlying graph.

参数:

viz_edge (VizEdge) -- A VizEdge instance.

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.

参数:

viz_node (VizNode) -- A VizNode instance.

class tvm.contrib.relay_viz.dot.DotPlotter(graph_attr: Dict[str, str] = None, node_attr: Dict[str, str] = None, edge_attr: Dict[str, str] = None, get_node_attr: Callable[[VizNode], Dict[str, str]] = None, render_kwargs: Dict[str, Any] = 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.

参数:

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

示例

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

参数:

name (str) -- the name of the graph

返回:

rv1

返回类型:

an instance of class inheriting from VizGraph interface.

render(filename: str = 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.TermGraph(name: str)

Terminal graph for a relay IR Module

参数:

name (str) -- name of this graph.

edge(viz_edge: VizEdge) → None

Add an edge to the terminal graph.

参数:

viz_edge (VizEdge) -- A VizEdge instance.

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.

参数:

viz_node (VizNode) -- A VizNode instance.

render() → str

Draw a terminal graph

返回:

rv1 -- text representing a graph.

返回类型:

str

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

Terminal plotter

create_graph(name)

Create a VizGraph

参数:

name (str) -- 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: 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.

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

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.

参数:

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

返回:

• 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

参数:

name (str) -- 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.

参数:

filename (str) -- 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: VizEdge) → None

Add an edge to the underlying graph.

参数:

viz_edge (VizEdge) -- A VizEdge instance.

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.

参数:

viz_node (VizNode) -- 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: 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.

参数:

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

返回:

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

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

参数:

• lhs (Tensor) -- The left batched matrix operand

• rhs (Tensor) -- The right batched matrix operand

• transa (bool) -- Whether transpose lhs

• transb (bool) -- Whether transpose rhs

返回:

C -- The result tensor.

返回类型:

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

参数:

• lhs (Tensor) -- The left matrix operand

• rhs (Tensor) -- The right matrix operand

• transa (bool) -- Whether transpose lhs

• transb (bool) -- Whether transpose rhs

返回:

C -- The result tensor.

返回类型:

Tensor

tvm.contrib.rocm

Utility for ROCm backend

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

Find ld.lld in system.

参数:

required (bool) -- Whether it is required, runtime error will be raised if the compiler is required.

返回:

valid_list -- List of possible paths.

返回类型:

list of str

备注

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

tvm.contrib.rocm.find_rocm_path()

Utility function to find ROCm path

返回:

path -- Path to ROCm root.

返回类型:

str

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

Either MatrixCore support is provided in the compute capability or not

参数:

compute_version (str, optional) -- compute capability of a GPU (e.g. "7.0").

返回:

have_matrixcore -- True if MatrixCore support is provided, False otherwise

返回类型:

bool

tvm.contrib.rocm.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

参数:

compute_version (str) -- compute capability of a GPU (e.g. "6.0")

返回:

• major (int) -- major version number

• minor (int) -- minor version number

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

Link relocatable ELF object to shared ELF object using lld

参数:

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

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

Placeholder class for CSR based sparse tensor representation.

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

Placeholder class for sparse tensor representations.

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 -- The created sparse tensor placeholder

返回类型:

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.

参数:

spv_bin (bytearray) -- The spirv file

返回:

cobj_bin -- The HSA Code Object

返回类型:

bytearray

tvm.contrib.tar

Util to invoke tarball in the system.

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.

参数:

• temp (tvm.contrib.utils.TempDirectory) -- A temp dir to hold the untared files.

• file_list (List[str]) -- List of path

返回:

ret_list -- An updated list of files

返回类型:

List[str]

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

Create tarball containing all files in root.

参数:

• 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

参数:

• tar_file (str) -- The source tar file.

• directory (str) -- 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

参数:

path (str) -- 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.

返回:

names -- The content of directory

返回类型:

list

relpath(name)

Relative path in temp dir

参数:

name (str) -- The name of the file.

返回:

path -- The concatenated path.

返回类型:

str

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

参数:

path (str) -- The path to the lock

返回:

lock

返回类型:

File lock object

tvm.contrib.utils.is_source_path(path)

Check if path is source code path.

参数:

path (str) -- A possible path

返回:

valid -- Whether path is a possible source path

返回类型:

bool

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

Create temp dir which deletes the contents when exit.

参数:

• custom_path (str, optional) -- Manually specify the exact temp dir path

• keep_for_debug (bool) -- Keep temp directory for debugging purposes

返回:

temp -- The temp directory object

返回类型:

TempDirectory

tvm.contrib.utils.which(exec_name)

Try to find full path of exec_name

参数:

exec_name (str) -- The executable name

返回:

path -- The full path of executable if found, otherwise returns None

返回类型:

str

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.

参数:

• code (str) -- The cuda code.

• path_target (str, optional) -- Output file.

• sdk (str, optional) -- The target platform SDK.

返回:

metallib -- The bytearray of the metallib

返回类型:

bytearray

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

Create dynamic library.

参数:

• 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.xcrun(cmd)

Run xcrun and return the output.

参数:

cmd (list of str) -- The command sequence.

返回:

out -- The output string.

返回类型:

str