tvm.contrib

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.

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

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

Parameters#

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.

Returns#

codestr: The generated llvm text IR.

tvm.contrib.clang.find_clang(required=True)[源代码]#

Find clang in system.

Parameters#

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

Returns#

valid_listlist of str: List of possible paths.

Note#

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.

Parameters#

outputstr: The target executable.
objectsList[str]: List of object files.
optionsList[str]: The list of additional options string.
ccOptional[str]: The compiler command.
cwdOptional[str]: The urrent working directory.
ccache_envOptional[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.

Parameters#

outputstr: The target shared library.
objectsList[str]: List of object files.
optionsList[str]: The list of additional options string.
ccOptional[str]: The compiler command.
cwdOptional[str]: The current working directory.
ccache_envOptional[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#

outputstr: The target shared library.
inputsList[str]: List of inputs files. Each input file can be a tarball of objects or an object file.
arOptional[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.

Parameters#

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.

Returns#

fcompileCallable[[str, str, Optional[str]], None]: A compilation function that can be passed to export_library.

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.cc.get_cc()[源代码]#

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

Returns#

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

tvm.contrib.cc.get_global_symbol_section_map(path, *, nm=None)[源代码]#

Get global symbols from a library via nm -g

Parameters#

pathstr: The library path
nm: str: The path to nm command

Returns#

symbol_section_map: Dict[str, str]: 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.

Parameters#

compilerOptional[str]: The compiler.

Returns#

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

Parameters#

lhsTensor: The left matrix operand
rhsTensor: The right matrix operand
transabool: Whether transpose lhs
transbbool: Whether transpose rhs

Returns#

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

Parameters#

lhsTensor: The left matrix operand
rhsTensor: The right matrix operand
transabool: Whether transpose lhs
transbbool: Whether transpose rhs

Returns#

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

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: 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', libs=None)[源代码]#

Create wasm that is supposed to run with the tvmjs.

Parameters#

outputstr: The target shared library.
objectslist: List of object files.
optionsstr: The additional options.
ccstr, optional: The compile string.
libslist: 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._get_np_int32_array_handle(arr)[源代码]#

Return a void_p handle for a numpy array

Parameters#

arr: numpy.NDArray: source numpy array

Returns#

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

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: Tensor: The result tensor

tvm.contrib.miopen.log_softmax(x, axis=-1)[源代码]#

Compute log softmax with MIOpen

Parameters#

xtvm.te.Tensor: The input tensor
axisint: The axis to compute log softmax over

Returns#

rettvm.te.Tensor: The result tensor

tvm.contrib.miopen.softmax(x, axis=-1)[源代码]#

Compute softmax with MIOpen

Parameters#

xtvm.te.Tensor: The input tensor
axisint: The axis to compute softmax over

Returns#

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.

Parameters#

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.

Returns#

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.

Parameters#

outputstr: The target shared library.
objectslist: List of object files.
optionslist of str, optional: The additional options.

tvm.contrib.ndk.create_staticlib(output, inputs)[源代码]#

Create static library:

Parameters#

outputstr: The target static library.
inputslist: List of object files or tar files

tvm.contrib.ndk.get_global_symbol_section_map(path, *, nm=None)[源代码]#

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

Parameters#

pathstr: The library path
nm: str: The path to nm command

Returns#

symbol_section_map: Dict[str, str]: 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.

Parameters#

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.

Returns#

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.

Parameters#

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

Returns#

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.

Parameters#

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.

Returns#

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.

Parameters#

lhsTensor: lhs 1D array input[input_channels] of FP32 elements
rhsTensor: lhs 2D matrix kernel[output_channels][input_channels] of FP32 elements

Returns#

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.

Parameters#

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.

Return#

cubinbytearray: The bytearray of the cubin

tvm.contrib.nvcc.find_cuda_path()[源代码]#

Utility function to find cuda path

Returns#

pathstr: Path to cuda root.

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#

versionfloat: The cuda version

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

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

Parameters#

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

Parameters#

compute_versionstr: compute capability of a GPU (e.g. "6.0")

Returns#

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.

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._atexit()[源代码]#: Save handler.

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#

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.

Parameters#

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.

Returns#

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

Parameters#

lowint: Lowest (signed) integer to be drawn from the distribution
highint: One above the largest (signed) integer to be drawn from the distribution

Returns#

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.

Parameters#

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.

Returns#

outTensor: A tensor with specified size and dtype.

tvm.contrib.relay_viz#

Relay IR Visualizer

class tvm.contrib.relay_viz.RelayVisualizer(relay_mod, relay_param=None, plotter=None, parser=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.

_add_nodes(graph, node_to_id)[源代码]#

add nodes and to the graph.

Parameters#

graphVizGraph: a VizGraph for nodes to be added to.
node_to_idDict[relay.expr, str]: a mapping from nodes to an unique ID.
relay_paramDict[str, tvm.runtime.NDarray]: relay parameter dictionary.

参数:

graph (VizGraph)
node_to_id (Dict[RelayExpr, str])

参数:

relay_mod (IRModule)
relay_param (Dict[str, NDArray])
plotter (Plotter)
parser (VizParser)

Visualize Relay IR by Graphviz DOT language.

class tvm.contrib.relay_viz.dot.DotGraph(name, graph_attr=None, node_attr=None, edge_attr=None, get_node_attr=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.

edge(viz_edge)[源代码]#

Add an edge to the underlying graph.

Parameters#

viz_edgeVizEdge: A VizEdge instance.

参数:: viz_edge (VizEdge)
返回类型:: None

node(viz_node)[源代码]#

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

Parameters#

viz_nodeVizNode: A VizNode instance.

参数:: viz_node (VizNode)
返回类型:: None

参数:

name (str)
graph_attr (Dict[str, str])
node_attr (Dict[str, str])
edge_attr (Dict[str, str])
get_node_attr (Callable[[VizNode], Dict[str, str]])

class tvm.contrib.relay_viz.dot.DotPlotter(graph_attr=None, node_attr=None, edge_attr=None, get_node_attr=None, render_kwargs=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#

namestr: the name of the graph

Return#

rv1: an instance of class inheriting from VizGraph interface.

render(filename=None)[源代码]#

render the graph generated from the Relay IR module.

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

参数:: filename (str)

参数:

graph_attr (Dict[str, str])
node_attr (Dict[str, str])
edge_attr (Dict[str, str])
get_node_attr (Callable[[VizNode], Dict[str, str]])
render_kwargs (Dict[str, Any])

Visualize Relay IR in AST text-form.

class tvm.contrib.relay_viz.terminal.TermGraph(name)[源代码]#

Terminal graph for a relay IR Module

Parameters#

name: str: name of this graph.

edge(viz_edge)[源代码]#

Add an edge to the terminal graph.

Parameters#

viz_edgeVizEdge: A VizEdge instance.

参数:: viz_edge (VizEdge)
返回类型:: None

node(viz_node)[源代码]#

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

Parameters#

viz_nodeVizNode: A VizNode instance.

参数:: viz_node (VizNode)
返回类型:: None

render()[源代码]#

Draw a terminal graph

Returns#

rv1: str: text representing a graph.

返回类型:: str

参数:: name (str)

class tvm.contrib.relay_viz.terminal.TermNode(viz_node)[源代码]#

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

参数:: viz_node (VizNode)

class tvm.contrib.relay_viz.terminal.TermPlotter[源代码]#

Terminal plotter

create_graph(name)[源代码]#

Create a VizGraph

Parameters#

namestr: the name of the graph

Return#

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, relay_param, node_to_id)[源代码]#

Parse a node and edges from a relay.Expr.

参数:

node (RelayExpr)
relay_param (Dict[str, NDArray])
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

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, node_to_id)[源代码]#

Render rule for a relay call node

参数:

node (RelayExpr)
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

_function(node, node_to_id)[源代码]#

Render rule for a relay function node

参数:

node (RelayExpr)
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

_var(node, relay_param, node_to_id)[源代码]#

Render rule for a relay var node

参数:

node (RelayExpr)
relay_param (Dict[str, NDArray])
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

get_node_edges(node, relay_param, node_to_id)[源代码]#

Get VizNode and VizEdges for a relay.Expr.

Parameters#

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.

Returns#

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.

参数:

node (RelayExpr)
relay_param (Dict[str, NDArray])
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

class tvm.contrib.relay_viz.interface.Plotter[源代码]#

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

abstract create_graph(name)[源代码]#

Create a VizGraph

Parameters#

namestr: the name of the graph

Return#

rv1: an instance of class inheriting from VizGraph interface.

参数:: name (str)
返回类型:: VizGraph

abstract render(filename)[源代码]#

Render the graph as a file.

Parameters#

filenamestr: see the definition of implemented class.

参数:: filename (str)
返回类型:: None

class tvm.contrib.relay_viz.interface.VizEdge(start_node, end_node)[源代码]#

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.

参数:

start_node (str)
end_node (str)

class tvm.contrib.relay_viz.interface.VizGraph[源代码]#

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

abstract edge(viz_edge)[源代码]#

Add an edge to the underlying graph.

Parameters#

viz_edgeVizEdge: A VizEdge instance.

参数:: viz_edge (VizEdge)
返回类型:: None

abstract node(viz_node)[源代码]#

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

Parameters#

viz_nodeVizNode: A VizNode instance.

参数:: viz_node (VizNode)
返回类型:: None

class tvm.contrib.relay_viz.interface.VizNode(node_id, node_type, node_detail)[源代码]#

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.

参数:

node_id (str)
node_type (str)
node_detail (str)

class tvm.contrib.relay_viz.interface.VizParser[源代码]#

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

abstract get_node_edges(node, relay_param, node_to_id)[源代码]#

Get VizNode and VizEdges for a relay.Expr.

Parameters#

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.

Returns#

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.

参数:

node (RelayExpr)
relay_param (Dict[str, NDArray])
node_to_id (Dict[RelayExpr, str])

返回类型:

Tuple[VizNode | None, List[VizEdge]]

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

Parameters#

lhsTensor: The left batched matrix operand
rhsTensor: The right batched matrix operand
transabool: Whether transpose lhs
transbbool: Whether transpose rhs

Returns#

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

Parameters#

lhsTensor: The left matrix operand
rhsTensor: The right matrix operand
transabool: Whether transpose lhs
transbbool: Whether transpose rhs

Returns#

CTensor: The result tensor.

tvm.contrib.rocm#

Utility for ROCm backend

tvm.contrib.rocm.find_lld(required=True)[源代码]#

Find ld.lld in system.

Parameters#

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

Returns#

valid_listlist of str: List of possible paths.

Note#

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

Returns#

pathstr: Path to ROCm root.

tvm.contrib.rocm.have_matrixcore(compute_version=None)[源代码]#

Either MatrixCore support is provided in the compute capability or not

Parameters#

compute_versionstr, optional: compute capability of a GPU (e.g. "7.0").

Returns#

have_matrixcorebool: True if MatrixCore support is provided, False otherwise

tvm.contrib.rocm.parse_compute_version(compute_version)[源代码]#

Parse compute capability string to divide major and minor version

Parameters#

compute_versionstr: compute capability of a GPU (e.g. "6.0")

Returns#

majorint: major version number
minorint: minor version number

tvm.contrib.rocm.rocm_link(in_file, out_file, lld=None)[源代码]#

Link relocatable ELF object to shared ELF object using lld

Parameters#

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.

Parameters#

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

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

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

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

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.

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

Parameters#

spv_binbytearray: The spirv file

Return#

cobj_binbytearray: The HSA Code Object

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.

Parameters#

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

Returns#

ret_list: List[str]: An updated list of files

tvm.contrib.tar.tar(output, files)[源代码]#

Create tarball containing all files in root.

Parameters#

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

Parameters#

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

Parameters#

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.

Returns#

nameslist: The content of directory

relpath(name)[源代码]#

Relative path in temp dir

Parameters#

namestr: The name of the file.

Returns#

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

Parameters#

pathstr: The path to the lock

Returns#

lock : File lock object

tvm.contrib.utils.is_source_path(path)[源代码]#

Check if path is source code path.

Parameters#

pathstr: A possible path

Returns#

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.

Parameters#

custom_pathstr, optional: Manually specify the exact temp dir path
keep_for_debugbool: Keep temp directory for debugging purposes

Returns#

tempTempDirectory: The temp directory object

tvm.contrib.utils.which(exec_name)[源代码]#

Try to find full path of exec_name

Parameters#

exec_namestr: The executable name

Returns#

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.

Parameters#

codestr: The cuda code.
path_targetstr, optional: Output file.
sdkstr, optional: The target platform SDK.

Return#

metallibbytearray: The bytearray of the metallib

tvm.contrib.xcode.create_dylib(output, objects, arch, sdk='macosx', min_os_version=None)[源代码]#

Create dynamic library.

Parameters#

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.

Parameters#

cmdlist of str: The command sequence.

Returns#

outstr: The output string.

tvm.contrib

目录

tvm.contrib#

tvm.contrib.cblas#

Parameters#

Returns#

Parameters#

Returns#

tvm.contrib.clang#

Parameters#

Returns#

Parameters#

Returns#

Note#

tvm.contrib.cc#

Parameters#

Parameters#

Parameters#

Parameters#

Returns#

Examples#

Returns#

Parameters#

Returns#

Parameters#

Returns#

tvm.contrib.cublas#

Parameters#

Returns#

Parameters#

Returns#

tvm.contrib.dlpack#

Parameters#

Parameters#

Returns#

tvm.contrib.emcc#

Parameters#

tvm.contrib.miopen#

Parameters#

Returns#

Parameters#

Returns#

Parameters#

Returns#

Parameters#

Returns#

tvm.contrib.mxnet#

Parameters#

Returns#

tvm.contrib.ndk#

Parameters#

Parameters#

Parameters#

Returns#

tvm.contrib.nnpack#

Parameters#

Returns#

Parameters#

Returns#

Parameters#

Returns#

Parameters#

Returns#

tvm.contrib.nvcc#

Parameters#

Return#

Returns#

Parameters#

Returns#

Parameters#

Parameters#

Parameters#

Parameters#

Returns#

tvm.contrib.pickle_memoize#

Parameters#

Parameters#

Returns#

tvm.contrib.random#

Parameters#