Contrib APIs

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

Utility to invoke nvcc compiler in the system

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

Compile cuda code with NVCC from env.

Parameters:
  • code (str) – The cuda code.
  • target (str) – The target format
  • arch (str) – The architecture
  • options (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.cc

Util to invoke c++ compilers in the system.

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

Create shared library.

Parameters:
  • output (str) – The target shared library.
  • objects (list) – List of object files.
  • options (list) – The list of additional options string.
  • cc (str, optional) – The compile string.

tvm.contrib.xcode

Utility to invoke Xcode compiler toolchain

class tvm.contrib.xcode.XCodeRPCServer(cmd, lock)

Wrapper for RPC server

Parameters:
  • cmd (list of str) – The command to run
  • lock (FileLock) – Lock on the path
join()

Wait server to finish and release its resource

tvm.contrib.xcode.codesign(lib)

Codesign the shared libary

This is an required step for library to be loaded in the app.

Parameters:lib (The path to the library.) –
tvm.contrib.xcode.compile_metal(code, path_target=None, sdk='macosx')

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.create_dylib(output, objects, arch, sdk='macosx')

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.popen_test_rpc(host, port, key, destination, libs=None, options=None)

Launch rpc server via xcodebuild test through another process.

Parameters:
  • host (str) – The address of RPC proxy host.
  • port (int) – The port of RPC proxy host
  • key (str) – The key of the RPC server
  • destination (str) – Destination device of deployment, as in xcodebuild
  • libs (list of str) – List of files to be packed into app/Frameworks/tvm These can be dylibs that can be loaed remoted by RPC.
  • options (list of str) – Additional options to xcodebuild
Returns:

proc – The test rpc server process. Don’t do wait() on proc, since it can terminate normally.

Return type:

Popen

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

Lightweight TVM RPC module.

RPC enables connect to a remote server, upload and launch functions. This is useful to for cross-compile and remote testing, The compiler stack runs on local server, while we use RPC server to run on remote runtime which don’t have a compiler available.

The test program compiles the program on local server, upload and run remote RPC server, get the result back to verify correctness.

tvm.contrib.graph_runtime

Minimum graph runtime that executes graph containing TVM PackedFunc.

class tvm.contrib.graph_runtime.GraphModule(module, ctx)

Wrapper runtime module.

This is a thin wrapper of the underlying TVM module. you can also directly call set_input, run, and get_output of underlying module functions

Parameters:
  • module (Module) – The interal tvm module that holds the actual graph functions.
  • ctx (TVMContext) – The context this module is under
module

Module – The interal tvm module that holds the actual graph functions.

ctx

TVMContext – The context this module is under

debug_get_output(node, out)

Run graph upto node and get the output to out

Parameters:
  • node (int / str) – The node index or name
  • out (NDArray) – The output array container
get_input(index, out)

Get index-th input to out

Parameters:
  • index (int) – The input index
  • out (NDArray) – The output array container
get_output(index, out)

Get index-th output to out

Parameters:
  • index (int) – The input index
  • out (NDArray) – The output array container
load_params(params_bytes)

Load parameters from serialized byte array of parameter dict.

Parameters:params_bytes (bytearray) – The serialized parameter dict.
run(**input_dict)

Run forward execution of the graph

Parameters:input_dict (dict of str to NDArray) – List of input values to be feed to
set_input(key=None, value=None, **params)

Set inputs to the module via kwargs

Parameters:
  • key (int or str) – The input key
  • value (the input value.) – The input key
  • params (dict of str to NDArray) – Additonal arguments
tvm.contrib.graph_runtime.create(graph_json_str, libmod, ctx)

Create a runtime executor module given a graph and module.

Parameters:
  • graph_json_str (str or graph class) – The graph to be deployed in json format output by nnvm graph. The graph can only contain one operator(tvm_op) that points to the name of PackedFunc in the libmod.
  • libmod (tvm.Module) – The module of the corresponding function
  • ctx (TVMContext) – The context to deploy the module, can be local or remote.
Returns:

graph_module – Runtime graph module that can be used to execute the graph.

Return type:

GraphModule

tvm.contrib.util

Common system utilities

class tvm.contrib.util.FileLock(path)

File lock object

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

Release the lock

class tvm.contrib.util.TempDirectory

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:names – The content of directory
Return type:list
relpath(name)

Relative path in temp dir

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

Remote the tmp dir

tvm.contrib.util.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.util.tempdir()

Create temp dir which deletes the contents when exit.

Returns:temp – The temp directory object
Return type:TempDirectory

tvm.contrib.cblas

External function interface to BLAS libraries.

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

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