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

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

External function interface to cuBLAS libraries.

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

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

Util to invoke emscripten compilers in the system.

tvm.contrib.emscripten.create_js(output, objects, options=None, side_module=False, cc='emcc')

Create emscripten javascript library.

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

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)

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

y – The result tensor

Return type:

Tensor

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.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:
  • 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_weight_transform(kernel, nthreads=1, algorithm=0)

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

Utility function to get cuda version

Parameters:cuda_path (str) – Path to cuda root.
Returns:version – The cuda version
Return type:float
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)

Either TensorCore support is provided in the compute capability or not

Parameters:compute_version (str) – compute capability of a GPU (e.g. “7.0”)
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.pickle_memoize

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

class tvm.contrib.pickle_memoize.Cache(key)

A cache object for result cache.

Parameters:key (str) – The file key to the function
tvm.contrib.pickle_memoize.memoize(key)

Memoize the result of function and reuse multiple times.

Parameters:key (str) – The unique key to the file
Returns:fmemoize – The decorator function to perform memoization.
Return type: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.

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

Utility for ROCm backend

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)

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.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.get_lower_ir(s)

Get lower ir code of a schedule. This is useful for debug, since you don’t have to find all inputs/outputs for a schedule in a fused subgraph.

Parameters:s (Schedule) –
Returns:ir – The lower ir
Return type:str
tvm.contrib.util.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.util.tempdir()

Create temp dir which deletes the contents when exit.

Returns:temp – The temp directory object
Return type:TempDirectory
tvm.contrib.util.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

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