tvm.relax.frontend¶

Frontends for constructing Relax programs, with the model importers

tvm.relax.frontend.detach_params(mod: tvm.ir.module.IRModule) → Tuple[tvm.ir.module.IRModule, Dict[str, List[tvm.runtime.ndarray.NDArray]]]¶

Detach the attribute “params” in the functions of the input IRModule as separate dictionary of params.

Parameters

mod (tvm.IRModule) – The IRModule whose functions’ “param” attribute is going to be detached.

Returns

detached_mod (tvm.IRModule) – The IRModule after the detachment.
params_dict (Dict[str, List[tvm.nd.NDArray]]) – The detached params. The dict keys corresponds to the names of the functions in the input IRModule that have attribute “params”.

tvm.relax.frontend.nn¶

A PyTorch-like API to build IRModules.

class tvm.relax.frontend.nn.Effect

Effect is a special non-user facing type that is used to represent operations with side effects, for example, print. It is used to represent the output of a computation.

emit_init(name_hint: str, builder: tvm.relax.block_builder.BlockBuilder) → List[tvm.relax.expr.DataflowVar]: Emit the initialization of the effect. This method is called by the compiler to initialize the effect.

create(name_hint: str) → List[tvm.relax.expr.Var]: Create the implicit inputs to a relax.Function that represents the side effect

set_state(state_vars: List[tvm.relax.expr.Var]) → None: Set the variables that represents the effect

finalize() → List[tvm.relax.expr.Var]: finalize the effect as the implicit return value of a relax.Function

to(dtype: Optional[str] = None) → None: Convert the effect to specific dtype. Usually it is no-op for most of the effects

class tvm.relax.frontend.nn.Module

Base class for neural network components. Subclass it to build your models. Modules can nest within each other in a tree structure using regular attribute assignment.

named_parameters(prefix: str = '') → Iterator[Tuple[str, tvm.relax.frontend.nn.core.Parameter]]

This method provides an iterator over module parameters, yielding both the parameter name and its corresponding value.

Parameters: prefix (str) – Prefix to prepend to all parameter names.
Yields: (str, Parameter) - Tuple containing the name and parameter

parameters() → Iterator[tvm.relax.frontend.nn.core.Parameter]

This method provides an iterator over module parameters, yielding only the Parameter value.

Yields: Parameter - The module’s parameter

state_dict(*, prefix: str = '', destination: Optional[Dict[str, tvm.relax.frontend.nn.core.Parameter]] = None) → Dict[str, tvm.relax.frontend.nn.core.Parameter]

Returns a dictionary containing references to the whole state of the module.

Parameters

prefix (str) – Prefix to prepend to all parameter names.
destination (Optional[Dict[str, Parameter]]) – Dictionary to which state will be saved. If None, a new dictionary is created.

Returns

dict – a dictionary containing a whole state of the module

Return type

Dict[str, Parameter]

load_state_dict(state_dict: Dict[str, tvm.relax.frontend.nn.core.Parameter], strict: bool = True) → Tuple[List[str], List[str]]

This function copies parameters and buffers from the state_dict into the current module and its descendants. If strict is set to True, the keys in the state_dict must exactly match the keys returned by the state_dict() function of this module.

Parameters

state_dict (Dict[str, Parameter]) – A dictionary containing a whole state of the module
strict (bool = True) – Whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function.

Returns

(missing_keys, unexpected_keys) – A tuple of two lists: the missing keys and the unexpected keys.

Return type

Tuple[List[str], List[str]]

to(dtype: Optional[str] = None) → None: Convert the module to specific dtype recursively

export_tvm(spec: _spec.ModuleSpecType, debug: bool = False, allow_extern: bool = False) → Union[Tuple[tvm.ir.module.IRModule, List[Tuple[str, tvm.relax.frontend.nn.core.Parameter]]], Tuple[tvm.ir.module.IRModule, List[Tuple[str, tvm.relax.frontend.nn.core.Parameter]], List[ExternModule]]]

Export the module to TVM IRModule and parameters

Parameters

spec (_spec.ModuleSpecType) – A dictionary mapping each input name to a specification that defines the inputs shape and dtype.
debug (bool) – If set to True, then the exported module will support effects. This enables things like printing in the graph.

Returns

irmodule (tvm.ir.IRModule) – The converted tvm IR representation of the model.
params (List[Tuple[str, Parameter]]) – A list of Parameters corresponding to the weights of the model.
ext_mods (List[nn.ExternModule]) – A list of ExternModules that are used in the model.

jit(spec: _spec.ModuleSpec, device: Union[str, tvm._ffi.runtime_ctypes.Device] = 'cpu', pipeline: Union[None, str, tvm.ir.transform.Pass] = 'default_build', out_format: str = 'torch', debug: bool = False) → Any: Just-in-time compilation of a nn.model to an executable

class tvm.relax.frontend.nn.ModuleList(modules: List[tvm.relax.frontend.nn.core.Module])

Holds submodules in a list.

append(module: tvm.relax.frontend.nn.core.Module): Add a module to the end of the ModuleList

to(dtype: Optional[str] = None) → None: Convert the module to specific dtype recursively

forward(x): Feed-forward pass of the module

class tvm.relax.frontend.nn.Object(*, _expr: tvm.ir.expr.RelayExpr, _name: str): A wrapper on top of relax.Expr whose struct_info is the base ObjectStructInfo (rather than any its subclass). Object effectively represents non-tensor frontend components such as KV caches.

class tvm.relax.frontend.nn.Parameter(shape: Sequence[Union[int, str, tvm.ir.expr.PrimExpr]], dtype: Optional[str] = None)

A parameter represents the weight of a neural network layer. It is a special tensor which could be bound or not bound to concrete values. If a parameter is bound to a concrete value, it is called a bound parameter, otherwise it is called an unbound parameter.

property data: Optional[tvm.runtime.ndarray.NDArray]: Returns the concrete value of the parameter if it is bound to a concrete value, otherwise returns None. The returned value is a tvm.runtime.NDArray.

to(dtype: Optional[str] = None) → None: Change the dtype of the parameter if it is not bound to any concrete data

class tvm.relax.frontend.nn.Tensor(*, _expr: tvm.ir.expr.RelayExpr)

A wrapper on top of relax.Expr whose struct_info is a TensorStructInfo, providing more convenient access shape and dtype information. Tensor is always symbolc and not bound to any concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when operators are applied on tensors, the shape and dtype information is already available.

static from_const(data) → tvm.relax.frontend.nn.core.Tensor: Construct a tensor from numpy constants.

static from_scalar(data: Union[int, float], dtype: str) → tvm.relax.frontend.nn.core.Tensor: Construct a tensor from a scalar with dtype specified.

static from_struct_info(struct_info: tvm.relax.struct_info.TensorStructInfo, name: str = 'tensor') → tvm.relax.frontend.nn.core.Tensor: Construct a nn.Tensor from relax TensorStructInfo

static placeholder(shape: Sequence[Union[int, str, tvm.ir.expr.PrimExpr]], dtype: str, name: str = 'tensor') → tvm.relax.frontend.nn.core.Tensor

Create a placeholder tensor with given shape and dtype. A placeholder tensor should never be created directly by users in usual cases, and the only exception is to indicate the shape/dtype of return values of an external function.

If shape is a string name, we create a symbolic shape tvm.tir.Var(name, “int64”).

property shape: List[Union[int, tvm.ir.expr.PrimExpr]]

Returns the shape of the tensor as a list of integers.

An integer can be a python int or tvm.tir.PrimExpr, depending on whether the shape is fully static, for example, [1, 2, tvm.tir.Var(“n”)] is a valid shape where the last dimension is dynamic while the first two dimensions are always static constants.

Returns: shape – The shape of the tensor
Return type: List[Union[int, tir.PrimExpr]]

property ndim: int

Returns the number of dimensions of the tensor.

Returns: ndim – The number of dimensions of the tensor
Return type: int

property dtype: str

Returns the data type of the tensor.

Returns: dtype – The data type of the tensor
Return type: str

tvm.relax.frontend.nn.add_extern(mod: tvm.relax.frontend.nn.extern.ExternModule) → None: Add an external module to the exporter.

class tvm.relax.frontend.nn.ExternModule(symbols: Dict[str, Callable])

The abstract base class for external modules. External modules are designed to help incorporate user-provided handcrafted kernels into the exported TVM IRModule.

load() → tvm.runtime.module.Module: Loads the external module into a TVM runtime module.

class tvm.relax.frontend.nn.ObjectModule(symbols: Dict[str, Callable], filepath: pathlib.Path)

A subclass of nn.ExternModule, which allows users to provide an object .o file to be linked into compiled artifact;

load() → tvm.runtime.module.Module: Loads the external module into a TVM runtime module.

class tvm.relax.frontend.nn.SourceModule(symbols: Dict[str, Callable], source_code: Union[str, pathlib.Path], source_format: str, compile_options: Optional[List[str]] = None, compiler: Optional[str] = None, output_format: str = 'obj')

A subclass of nn.ExternModule. It compiles C++/CUDA source code and link them into the eventual IRModule.

Shape/dtype inference. The nn.ExternModule system requires users to provide additional information to work, namely, symbols. It is a dictionary that maps each symbol in the external object file to its shape/dtype inference function. Consider a case where function my_func accepts two tensors, a of shape (x, y, 1), and b of shape (y, z, 5), and produces a tensor c of shape (x, y, z, 9), the shape/dtype inference function should look like:

def shape_dtype_inference(a, b):
    x, y, _ = a.shape
    _, z, _ = b.shape
    return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")

and the symbols dictionary should be provided as:

symbols={
    "my_func": shape_dtype_inference,
}

Calling convention. All external modules now follows “destination-passing-style” (DPS) calling convention, which means the returned tensors are pre-allocated by the system already and passed in as an argument of the external function.

Reuse the example above, the implementation of my_func should include three parameters in its signature, where tensors are represented using DLTensor from DLPack, the de facto standard of in-memory representation of tensors. More details: https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206.

To expose the symbol, TVM_DLL_EXPORT_TYPED_FUNC(symbol, function) is guaranteed available:

// those headers are guaranteed to be available
#include <dlpack/dlpack.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/packed_func.h>

namespace {
// anonymous namespace hides the symbol `_my_func_impl` from other translation units
int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
    // `a` and `b` are inputs, and `c` is the output
}
}
// expose symbol `my_func` instead of `_my_func_impl`
TVM_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);

A compiler pass `AttachExternModules`. It is introduced to attach a list of nn.ExternModule`s into an IRModule at any stage of the compilation pipeline, and attach the compiled external modules as `runtime.Module`s into IRModule’s `external_mods attribute. It is required by linking in relax.build, but with the existence of this pass, source compilation can be deferred to arbitrary stage of TVM compilation.

Caveats. It is required to call nn.add_extern to register external modules exactly once during export_tvm. Each symbol should be registered exactly once to avoid potential conflicts, and otherwise an error will be raised.

static tvm_home() → pathlib.Path

Find TVM’s home directory. If TVM_HOME environment variable is set, use it. Otherwise, use the directory where the tvm Python package is installed. As a sanity check, it is required to have include and 3rdparty as direct subdirectories.

Returns: tvm_home – The TVM home directory, and it is guaranteed to have include and 3rdparty as direct subdirectories.
Return type: pathlib.Path

static get_includes(tvm_pkg: Optional[List[str]] = None) → List[pathlib.Path]

Returns the default include paths according to tvm_home(). By default, it includes TVM, DLPack, and DMLC-Core. With tvm_pkg provided, it also includes the specified package under tvm_home/3rdparty.

Parameters: tvm_pkg (Optional[List[str]]) – The list of packages to be included under tvm_home/3rdparty. Each element should be a relative path to tvm_home/3rdparty.
Returns: includes – The list of include paths.
Return type: List[pathlib.Path]

static get_compile_options(source_format: str, tvm_pkg: Optional[List[str]] = None) → List[str]

Returns the default compile options depending on source_format, including the default inlcude paths w.r.t. tvm_home(), default flags to configure DMLC-Core, and by default, it uses “-O3” and “-std=c++17”.

Parameters

source_format (str) – The source code format. It can be either “cpp” or “cu”.
tvm_pkg (Optional[List[str]]) – The list of packages to be included under tvm_home/3rdparty. Each element should be a relative path to tvm_home/3rdparty.

Returns

compile_options – The list of compilation flags.

Return type

List[str]

compile(output_path: pathlib.Path) → None: Compiles the source code in a provided directory and returns the compiled artifact.

load() → tvm.runtime.module.Module: Loads the external module into a TVM runtime module.

class tvm.relax.frontend.nn.GELU: relax.frontend.nn.Module for GELU activation layer.

class tvm.relax.frontend.nn.Conv1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: Optional[str] = None)

relax.frontend.nn.Module for conv1d layer.

forward(x: tvm.relax.frontend.nn.core.Tensor) → tvm.relax.frontend.nn.core.Tensor

Forward method for conv1d layer.

Parameters: x (Tensor) – The input tensor.
Returns: ret – The output tensor for the conv1d layer.
Return type: Tensor

class tvm.relax.frontend.nn.Conv2D(in_channels: int, out_channels: int, kernel_size: Union[List[int], int], stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: Optional[str] = None, data_layout: str = 'NCHW')

relax.frontend.nn.Module for conv2d layer.

forward(x: tvm.relax.frontend.nn.core.Tensor) → tvm.relax.frontend.nn.core.Tensor

Forward method for conv2d layer.

Parameters: x (Tensor) – The input tensor.
Returns: ret – The output tensor for the conv2d layer.
Return type: Tensor

class tvm.relax.frontend.nn.Conv3D(in_channels: int, out_channels: int, kernel_size: Union[List[int], int], stride: Union[List[int], int] = 1, padding: Union[List[int], int] = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: Optional[str] = None, data_layout: str = 'NCDHW')

relax.frontend.nn.Module for conv3d layer.

forward(x: tvm.relax.frontend.nn.core.Tensor) → tvm.relax.frontend.nn.core.Tensor

Forward method for conv3d layer.

Parameters: x (Tensor) – The input tensor.
Returns: ret – The output tensor for the conv3d layer.
Return type: Tensor

class tvm.relax.frontend.nn.ConvTranspose1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: Optional[str] = None)

relax.frontend.nn.Module for ConvTranspose1D layer.

forward(x: tvm.relax.frontend.nn.core.Tensor) → tvm.relax.frontend.nn.core.Tensor

Forward method for conv transpose 1d layer.

Parameters: x (Tensor) – The input tensor.
Returns: ret – The output tensor for the conv transpose 1d layer.
Return type: Tensor

class tvm.relax.frontend.nn.Embedding(num: Union[int, str, tvm.ir.expr.PrimExpr], dim: Union[int, str, tvm.ir.expr.PrimExpr], dtype: Optional[str] = None)

relax.frontend.nn.Module for embedding layer.

forward(x: tvm.relax.frontend.nn.core.Tensor)

Forward method for embedding layer.

Parameters: x (Tensor) – The input tensor.
Returns: ret – The output tensor for the embedding layer.
Return type: Tensor

class tvm.relax.frontend.nn.GroupNorm(num_groups: int, num_channels: int, eps: float = 1e-05, affine: bool = True, dtype: Optional[str] = None)

relax.frontend.nn.Module for group norm layer.

forward(x: tvm.relax.frontend.nn.core.Tensor, channel_axis: int = 1, axes: Optional[List[int]] = None)

Forward method for group norm layer.

Parameters

x (Tensor) – The input tensor.
channel_axis (int) – Channel axis of the input data.
axes (Optional[List[int]]) – Optional list of axes to compute norm over, if not specified, assumes that the first two axes should be left alone.

Returns

ret – The output tensor for the group norm layer.

Return type

tvm.relax.frontend¶

tvm.relax.frontend.nn¶

tvm.relax.frontend.onnx¶

tvm.relax.frontend.stablehlo¶

tvm.relax.frontend.torch¶