tvm.relax

The Relax IR namespace containing the IR, type, operator, builder, vm, etc.

class tvm.relax.Binding

The base class of a binding in Relax.

class tvm.relax.BindingBlock(bindings, span=None)

base class of binding block, bindings inside can be impure (with side effect or control flow)

参数:

• bindings (List[Binding])

• span (Span | None)

class tvm.relax.Call(op, args, attrs=None, sinfo_args=None, span=None)

Function call node in Relax.

Call node corresponds the operator application node in computational graph terminology.

Parameters

op: tvm.ir.Op or any tvm.relax.Expr with function type.

The operation to be called.

args: Union[List[Expr], typing.Tuple[Expr, ...]]

The arguments to the call.

attrs: Optional[tvm.ir.Attrs]

Attributes to the call, can be None

sinfo_args: Optional[Union[List[StructInfo], typing.Tuple[StructInfo, ...]]]

The structure info arguments of a CallNode. sinfo_args is designed to be non-empty only for intrinsic op (e.g., call_tir, call_builtin_with_ctx, etc.) and calls to ExternFuncs, with the main usage of structure info inference.

span: Optional[Span]

Span that points to original source code

参数:

• op (RelayExpr)

• args (List[RelayExpr])

• attrs (Attrs)

• sinfo_args (List[StructInfo])

• span (Span | None)

class tvm.relax.Constant(data, struct_info=None, span=None)

Constant Tensor

Parameters

data: tvm.nd.NDArray

The data of the constant tensor.

struct_info: Optional[StructInfo]

The struct info of the constant tensor. If not specified, infer it from data.

span: Optional[Span]

Span that points to original source code

Note

Scalar constants are represented by ndim-0 constant tensors.

参数:

• data (NDArray)

• struct_info (StructInfo | None)

• span (Span | None)

class tvm.relax.DataTypeImm(value, span=None)

Represent a data type constant.

参数:

• value (DataType)

• span (Span | None)

class tvm.relax.DataflowBlock(bindings, span=None)

dataflow block, bindings inside are pure (no side effect and no control flow)

参数:

• bindings (List[Binding])

• span (Span | None)

class tvm.relax.DataflowBlockRewrite(dfb, root_fn)

A binding/statement-level dataflow block rewriter.

Notes

Due to the immutable and copy-on-write nature of TVM AST nodes, the rewriting is not done in place. Instead, a new DataflowBlock is created and returned with mutated_dfb. Similarly, its new root Function is created and returned by mutated_root_fn. To apply this change for an IRModule, use mutate_irmodule which rewrites the old function that registered in the constructor.

__init__(dfb, root_fn)

Construct a rewriter with the DataflowBlock to rewrite and its root function.

Parameters

dfbDataflowBlock

The DataflowBlock to rewrite.

root_fnFunction

The root function of the DataflowBlock.

参数:

• dfb (DataflowBlock)

• root_fn (Function)

add(expr, name=None, is_dfvar=False)

Add a new statement to the DataflowBlock with an automatically generated variable name.

Parameters

exprExpr

The expression to add.

nameOptional[str], optional

Variable name, by default None

is_dfvarbool, optional

The variable type, by default False

Notes

If the variable name is not given, it will be automatically generated in a form of "tmp${COUNTER}". The variable type will be DataflowVar if is_dfvar is True, otherwise it will be Var. Being Var means the variables are output variables of the DataflowBlock. While being DataflowVar means the variables are internal variables of the DataflowBlock.

参数:

• expr (RelayExpr)

• name (str | None)

• is_dfvar (bool)

返回类型:

None

mutate_irmodule(irmodule)

Return an updated IRModule by replacing the old function with the mutated root function.

Parameters

irmoduletvm.IRModule

The base IRModule to update.

Returns

tvm.IRModule

The updated IRModule.

参数:

irmodule (IRModule)

返回类型:

IRModule

mutated_dfb()

Returns the mutated DataflowBlock.

返回类型:

DataflowBlock

mutated_root_fn()

Returns the mutated root function.

返回类型:

Function

remove_all_unused()

Remove all unused variables.

Notes

This could remove unused variables in other DataflowBlocks as well.

返回类型:

None

remove_unused(var, allow_undef=False)

Remove a statement by its variable definition if and only if it is unused.

Parameters

varVar

The unused variable definition.

allow_undefbool, optional

Whether to allow var being undefined variable, by default False

Raises

TVMError if the variable is used or undefined (allow_undef=False).

参数:

var (Var)

返回类型:

None

replace_all_uses(old_var, new_var)

Replace all uses of old_var with new_var.

Parameters

old_varVar

The old variable to replace.

new_varVar

The new variable to replace with.

参数:

• old_var (Var)

• new_var (Var)

返回类型:

None

参数:

• dfb (DataflowBlock)

• root_fn (Function)

class tvm.relax.DataflowVar(name_hint, struct_info=None, span=None)

A sub-type of the variable node used to mark dataflow variables from normal visible "function local" bindings.

Parameters

name_hint: Union[str, Id]

The name hint of the variable.

struct_info: Optional[StructInfo]

The struct info annotation of the variable.

span: Optional[Span]

Span that points to original source code

参数:

• name_hint (str | Id)

• struct_info (StructInfo | None)

• span (Span | None)

class tvm.relax.DynTensorType(ndim=-1, dtype='float32', span=None)

A dynamic tensor type in Relax.

This is the type assigned to tensors with a known dtype and unknown shape.

Parameters

ndimOptional[int]

The ndim of the Tensor

dtypeOptional[str]

The content data type.

参数:

span (Span)

class tvm.relax.ExecBuilder

A builder to emit instructions and build executable for the virtual machine.

c(idx)

set instruction's argument as a constant.

参数:

idx (int)

返回类型:

int

declare_function(func_name, kind=VMFuncKind.PACKED_FUNC)

Declare a function

参数:

• func_name (str)

• kind (VMFuncKind)

返回类型:

None

emit_call(name, args=None, dst=None)

emit a call instruction which calls a packed function.

参数:

• name (str)

• args (List[NDArray | DataType] | None)

• dst (int)

返回类型:

None

emit_goto(pc_offset)

emit a goto instruction

emit_if(cond, false_offset)

emit an if instruction

emit_ret(result)

emit a return instruction

参数:

result (int)

返回类型:

None

f(name)

set instruction's argument as a function.

参数:

name (str)

返回类型:

int

function(func_name, num_inputs=0, param_names=None)

annotate a VM function.

参数:

• func_name (str)

• num_inputs (int | None)

• param_names (List[str])

返回类型:

VMFuncScope

get()

return the executable

返回类型:

Executable

imm(value)

set instruction's argument as an immediate.

参数:

value (int)

返回类型:

int

r(idx)

set instruction's argument as a register.

参数:

idx (int)

返回类型:

int

class tvm.relax.Executable(mod)

The executable object emitted by the VM compiler or the ExecBuilder.

参数:

mod (Module)

as_python()

print the instructions as python program.

返回类型:

str

as_text()

print the instructions as text format.

返回类型:

str

export_library(file_name, fcompile=None, workspace_dir=None, **kwargs)

Export the executable to a library which can then be loaded back.

Parameters

file_namestr

The name of the shared library.

fcompilefunction(target, file_list, kwargs), optional

The compilation function to use create the final library object during

workspace_dirstr, optional

The path of the directory used to create the intermediate artifacts when exporting the module. If this is not provided a temporary dir will be created.

kwargsdict, optional

Additional arguments passed to fcompile

Returns

result of fcompile()unknown, optional

If the compilation function returns an artifact it would be returned via export_library, if any.

Examples

ex = relax.build(mod, target)
# export the library
ex.export_library("exported.so")

# load it back for future uses.
rt_mod = tvm.runtime.load_module("exported.so")
vm = tvm.relax.VirtualMachine(rt_mod, tvm.cuda())

参数:

• file_name (str)

• fcompile (str | callable | None)

• workspace_dir (str | None)

返回类型:

Any

jit(fcompile=None, addons=None, **kwargs)

Just-in-time compile and link the modules.

The Executable returned by relax.build may not be directly runnable as they may contain cuda source files and objects that are yet to be compiled and linked. This function helps to create a runtime.Module for these cases.

Parameters

fcompilefunction(target, file_list, kwargs), optional

The compilation function to use create the final library object during

kwargsdict, optional

Additional arguments passed to fcompile

Returns

rt_mod: tvm.runtime.Module

A runnable runtime module that can be passed to VirtualMachine.

Examples

ex = relax.build(mod, target)
# build a runnable module using nvcc to link everything
rt_mod = ex.jit()
vm = tvm.relax.VirtualMachine(rt_mod, tvm.cuda())

返回类型:

Module

stats()

print the detailed statistics of the executable.

返回类型:

str

tvm.relax.Expr

RelayExpr 的别名

class tvm.relax.ExprFunctor

An abstract visitor defined over Expr. Defines the default dispatch over expressions, and implements memoization.

visit_expr(expr)

Apply the visitor to an expression.

参数:

expr (RelayExpr)

返回类型:

RelayExpr

class tvm.relax.ExternFunc(global_symbol, struct_info=None, span=None)

extern function, which represents a PackedFunc.

参数:

• global_symbol (String)

• struct_info (StructInfo | None)

• span (Span | None)

class tvm.relax.FuncStructInfo(params, ret, purity=True, span=None)

StructInfo of a function value.

Parameters

params: List[StructInfo]

The struct info of the fields.

ret: StructInfo

The struct info of return value

purity: bool

Whether the function is pure (has no visible side effects). Note: We consider a function to be pure only if it is pure on all inputs. If a function can have visible side effects only in some cases, we still consider it impure.

static opaque_func(*, ret=None, derive_func=None, purity=False, span=None)

Create an opaque FuncStructInfo.

The opaque function takes either a ret that specificies the struct info of the return value or a derive_func that provides a customized derivation rule.

Parameters

ret: Optional[StructInfo]

The struct info of the function return value.

derive_func: Optional[Union[str,EnvFunc]]

The environment function used for derivation

purity: bool

Whether the function is pure (false by default, as most opaque functions are not pure)

span: Optional[Span]

Optional span information of the ast.

Returns

info: FuncStructInfo

Note

We cannot specify ret and derive_func simultaneously.

参数:

• ret (StructInfo | None)

• derive_func (str | EnvFunc | None)

• purity (bool)

• span (Span)

返回类型:

FuncStructInfo

参数:

• params (List[StructInfo] | None)

• ret (StructInfo)

• purity (bool)

• span (Span)

class tvm.relax.Function(params, body, ret_struct_info=None, is_pure=True, attrs=None, span=None)

A Relax function.

参数:

• params (List[Var])

• body (RelayExpr)

• ret_struct_info (StructInfo)

• is_pure (bool)

• attrs (DictAttrs)

• span (Span | None)

__call__(*args)

Invoke the global function.

Parameters

args: List[relax.Expr]

Arguments.

bind_params(binding_map)

Return a new function with updated symbolic variable

Parameters

binding_map: Mapping[

Union[str, Var], Union[int, float, PrimExpr, tvm.runtime.NDArray, _np.ndarray, Expr],

]

The mapping of values to be replaced.

Keys may be either a relax.Var or a string name of the Relax variable. If the variables are referred to by name, the name must uniquely identify a parameter in the function.

Values must be a relax expression, or a value that is convertible into a relax expression. The value must be compatible with the variable being replaced.

Returns

func: Function

The updated function

参数:

binding_map (Mapping[str | Var, int | float | PrimExpr | NDArray | ndarray | RelayExpr])

返回类型:

Function

bind_symbolic_vars(binding_map)

Return a new function with updated symbolic variable

Parameters

binding_map: Mapping[Union[str, tvm.tir.Var], PrimExpr]

The mapping of values to be replaced. Keys may be either a tir.Var or a string name of the variable. If the variables are referred to by name, the name must uniquely identify a symbolic variable in the function.

Returns

func: Function

The updated function

参数:

binding_map (Mapping[str | Var, PrimExpr])

返回类型:

Function

static create_empty(params, ret_struct_info, is_pure=True, attrs=None, span=None)

Construct a relax.Function but without body

参数:

• params (List[Var])

• ret_struct_info (StructInfo)

• is_pure (bool | None)

• attrs (DictAttrs | None)

• span (Span | None)

class tvm.relax.Id

Unique identifier(name) used in Var. Guaranteed to be stable across all passes.

class tvm.relax.If(cond, true_branch, false_branch, span=None)

A conditional expression in Relax.

Parameters

cond: Expr

The condition.

true_branch: Expr

The expression evaluated when condition is true.

false_branch: Expr

The expression evaluated when condition is false.

span: Optional[Span]

Span that points to original source code

参数:

• cond (RelayExpr)

• true_branch (RelayExpr)

• false_branch (RelayExpr)

• span (Span | None)

class tvm.relax.MatchCast(var, value, struct_info, span=None)

Runtime-match the value to the struct info.

This operation does runtime check, populates the un-defined symbolic shape vars and vars in struct_info in the first occurrence, and insert equality assertions in other cases.

Parameters

var: Var

The return variable that the match cast bind to.

value: Expr

The input value expression.

struct_info: tvm.relax.StructInfo

The struct info to match cast to.

参数:

• var (Var)

• value (RelayExpr)

• struct_info (StructInfo)

• span (Span | None)

class tvm.relax.ObjectStructInfo(span=None)

StructInfo of an Object.

参数:

span (Span)

class tvm.relax.ObjectType(span=None)

A type that corresponds to tvm::runtime::Object, is base of all possible object values in TVM.

参数:

span (Span)

class tvm.relax.PackedFuncType(span=None)

The type of ExternFunc in Relax.

参数:

span (Span)

class tvm.relax.PrimStructInfo(dtype=None, value=None, span=None)

StructInfo of a primitive POD value.

Parameters

dtype_or_expr : Union[str, DataType, PrimExpr]

The data type of the prim value, or a known expression for the prim value.

参数:

• dtype (str)

• value (PrimExpr | None)

• span (Span)

class tvm.relax.PrimValue(value, span=None)

The prim expr representing the value.

参数:

• value (PrimExpr)

• span (Span | None)

class tvm.relax.PyExprMutator(mod=None)

An abstract ExprMutator with customized methods on the python-side. This is the user facing class for method overwriting inheritance. _tvm_metadata discribes the class to inherit("cls"), the methods that users can overwrite("methods"), the constructor's parameters("fields")

Note: @relax.expr_functor.mutator is required for proper usage of any inherited class.

See also: visitor, _PyExprVisitor

Example:

@relax.expr_functor.mutator
def MyExprMutator(PyExprMutator):
    ...

参数:

mod (IRModule | None)

__init__(mod=None)

Constructor

参数:

mod (IRModule | None)

返回类型:

None

get_var_remap(vid)

Remap a var to a new var in use-site.

Parameters

vidId

The vid of the old var

Returns

varVar

The remapped var.

参数:

vid (Id)

返回类型:

Var

lookup_binding(var)

Look up the value bound to a variable. Note: For function parameters, this function returns NullOpt.

Parameters

varVar

The var to be looked up.

Returns

varVar

The value bound to the input var.

参数:

var (Var)

返回类型:

RelayExpr | None

set_var_remap(vid, var)

Remap a var to a new var in use-site.

Parameters

vidId

The vid of the old var.

varVar

The new var.

参数:

• vid (Id)

• var (Var)

返回类型:

None

visit_binding(binding)

Generic dispatcher for Binding. Users can customized this function to overwrite VisitBinding(const Binding& binding) on the C++ side.

Parameters

bindingBinding

The binding to be visited.

参数:

binding (Binding)

返回类型:

None

visit_binding_block(block)

Generic dispatcher for BindingBlock. Users can customized this function to overwrite VisitBindingBlock(const BindingBlock& block) on the C++ side.

Parameters

blockBindingBlock

The block to be visited.

Returns

resultBindingBlock

The binding block after transformation.

参数:

block (BindingBlock)

返回类型:

BindingBlock

visit_binding_block_(block)

Visit BindingBlock. Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode* block) on the C++ side.

Parameters

blockBindingBlock

The BindingBlock to be visited.

Returns

resultBindingBlock

The binding block after transformation

参数:

block (BindingBlock)

返回类型:

BindingBlock

visit_call_(op)

Visit Call. Users can customized this function to overwrite VisitExpr_(const CallNode* op) on the C++ side.

Parameters

opCall

The Call to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Call)

返回类型:

RelayExpr

visit_constant_(op)

Visit Constant. Users can customized this function to overwrite VisitExpr_(const ConstantNode* op) on the C++ side.

Parameters

opConstant

The Constant to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Constant)

返回类型:

RelayExpr

visit_data_type_imm_(op)

Visit DataTypeImm. Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op) on the C++ side.

Parameters

opDataTypeImm

The DataTypeImm to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (DataTypeImm)

返回类型:

RelayExpr

visit_dataflow_block_(block)

Visit DataflowBlock. Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode* block) on the C++ side.

Parameters

blockDataflowBlock

The DataflowBlock to be visited.

Returns

resultBindingBlock

The binding block after transformation

参数:

block (DataflowBlock)

返回类型:

BindingBlock

visit_dataflow_var_(op)

Visit DataflowVar. Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op) on the C++ side.

Parameters

opDataflowVar

The DataflowVar to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (DataflowVar)

返回类型:

RelayExpr

visit_dataflow_var_def_(var)

Visit the DataflowVar definition site. Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var) on the C++ side.

Parameters

varDataflowVar

The DataflowVar to be visited.

Returns

resultVar

The var after post-order rewritten.

参数:

var (DataflowVar)

返回类型:

Var

visit_expr(expr)

Generic dispatcher for Expr. Users can customized this function to overwrite VisitExpr(const Expr& expr) on the C++ side.

Parameters

exprExpr

The expr to be visited.

Returns

resultExpr

The Expr after transformation.

参数:

expr (RelayExpr)

返回类型:

RelayExpr

visit_expr_post_order(expr)

Post-order rewrite an Expr and normalize.

Parameters

exprExpr

The Expr to be rewritten.

Returns

resultExpr

The Expr after post-order rewritten.

参数:

expr (RelayExpr)

返回类型:

RelayExpr

visit_extern_func_(op)

Visit ExternFunc. Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op) on the C++ side.

Parameters

opExternFunc

The ExternFunc to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (ExternFunc)

返回类型:

RelayExpr

visit_function_(op)

Visit Function. Users can customized this function to overwrite VisitExpr_(const FunctionNode* op) on the C++ side.

Parameters

opFunction

The Function to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Function)

返回类型:

RelayExpr

visit_global_var_(op)

Visit GlobalVar. Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op) on the C++ side.

Parameters

opGlobalVar

The GlobalVar to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (GlobalVar)

返回类型:

RelayExpr

visit_if_(op)

Visit If. Users can customized this function to overwrite VisitExpr_(const IfNode* op) on the C++ side.

Parameters

opIf

The If to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (If)

返回类型:

RelayExpr

visit_match_cast_(binding)

Visit MatchCast. Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding) on the C++ side.

Parameters

bindingMatchCast

The MatchCast to be visited.

参数:

binding (MatchCast)

返回类型:

None

visit_op_(op)

Visit Op. Users can customized this function to overwrite VisitExpr_(const OpNode* op) on the C++ side.

Parameters

opOp

The Op to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Op)

返回类型:

RelayExpr

visit_prim_value_(op)

Visit PrimValue. Users can customized this function to overwrite VisitExpr_(const PrimValueNode* op) on the C++ side.

Parameters

opPrimValue

The PrimValue to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (PrimValue)

返回类型:

RelayExpr

visit_seq_expr_(op)

Visit SeqExpr. Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op) on the C++ side.

Parameters

opSeqExpr

The SeqExpr to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (SeqExpr)

返回类型:

RelayExpr

visit_shape_expr_(op)

Visit ShapeExpr. Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op) on the C++ side.

Parameters

opShapeExpr

The ShapeExpr to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (ShapeExpr)

返回类型:

RelayExpr

visit_span(span)

Visit Span. Users can customized this function to overwrite VisitSpan(const Span& span) on the C++ side.

Parameters

spanSpan

The Span to be visited.

Returns

resultSpan

The span after transformation.

参数:

span (Span)

返回类型:

Span

visit_string_imm_(op)

Visit StringImm. Users can customized this function to overwrite VisitExpr_(const StringImmNode* op) on the C++ side.

Parameters

opStringImm

The StringImm to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (StringImm)

返回类型:

RelayExpr

visit_tuple_(op)

Visit Tuple. Users can customized this function to overwrite VisitExpr_(const TupleNode* op) on the C++ side.

Parameters

opTuple

The Tuple to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Tuple)

返回类型:

RelayExpr

visit_tuple_getitem_(op)

Visit TupleGetItem. Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op) on the C++ side.

Parameters

opTupleGetItem

The TupleGetItem to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (TupleGetItem)

返回类型:

RelayExpr

visit_var_(op)

Visit Var. Users can customized this function to overwrite VisitExpr_(const VarNode* op) on the C++ side.

Parameters

opVar

The Var to be visited.

Returns

resultExpr

The Expr after transformation

参数:

op (Var)

返回类型:

RelayExpr

visit_var_binding_(binding)

Visit VarBinding. Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding) on the C++ side.

Parameters

bindingVarBinding

The VarBinding to be visited.

参数:

binding (VarBinding)

返回类型:

None

visit_var_def(var)

Generic dispatcher for visiting the var definition site. Users can customized this function to overwrite VisitVarDef(const Var& var) on the C++ side. Note that visit_var_() will only visit the usage site of an Var.

Parameters

varVar

The var to be visited.

Returns

result: Var

The var after post-order rewritten.

参数:

var (Var)

返回类型:

Var

visit_var_def_(var)

Visit the Var definition site. Users can customized this function to overwrite VisitVarDef_(const VarNode* var) on the C++ side.

Parameters

varVar

The Var to be visited.

Returns

resultVar

The var after post-order rewritten.

参数:

var (Var)

返回类型:

Var

visit_with_new_scope(expr)

Rewrite the expr with a new scope, used in a Function's body and the branches of If.

Parameters

exprExpr

The expr to be visited.

Returns

varVar

The expr after visiting.

参数:

expr (RelayExpr)

返回类型:

RelayExpr

with_struct_info(var, struct_info)

Create a new var with specified shape and type if the original var's shape or type does not match with the specified ones.

Parameters

varVar

The var to be updated.

struct_infoStructInfo

The struct info.

Returns

varVar

The var filled with shape and type.

参数:

• var (Var)

• struct_info (StructInfo)

返回类型:

Var

class tvm.relax.PyExprVisitor

An abstract ExprVisitor with customized methods on the python-side. This is the user facing class for method overwriting inheritance. _tvm_metadata discribes the class to inherit("cls"), the methods that users can overwrite("methods").

Note: @relax.expr_functor.visitor is required for proper usage of any inherited class.

See also: visitor, _PyExprVisitor

Example:

@relax.expr_functor.visitor
def MyExprVisitor(PyExprVisitor):
    ...

visit_binding(binding)

Generic dispatcher for Binding. Users can customized this function to overwrite VisitBinding(const Binding& binding) on the C++ side.

Parameters

bindingBinding

The binding to be visited.

参数:

binding (Binding)

返回类型:

None

visit_binding_block(block)

Generic dispatcher for BindingBlock. Users can customized this function to overwrite VisitBindingBlock(const BindingBlock& block) on the C++ side.

Parameters

blockBindingBlock

The block to be visited.

参数:

block (BindingBlock)

返回类型:

None

visit_binding_block_(block)

Visit BindingBlock. Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode* block) on the C++ side.

Parameters

blockBindingBlock

The BindingBlock to be visited.

参数:

block (BindingBlock)

返回类型:

None

visit_call_(op)

Visit Call. Users can customized this function to overwrite VisitExpr_(const CallNode* op) on the C++ side.

Parameters

opCall

The Call to be visited.

参数:

op (Call)

返回类型:

None

visit_constant_(op)

Visit Constant. Users can customized this function to overwrite VisitExpr_(const ConstantNode* op) on the C++ side.

Parameters

opConstant

The Constant to be visited.

参数:

op (Constant)

返回类型:

None

visit_data_type_imm_(op)

Visit DataTypeImm. Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op) on the C++ side.

Parameters

opDataTypeImm

The DataTypeImm to be visited.

参数:

op (DataTypeImm)

返回类型:

None

visit_dataflow_block_(block)

Visit DataflowBlock. Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode* block) on the C++ side.

Parameters

blockDataflowBlock

The DataflowBlock to be visited.

参数:

block (DataflowBlock)

返回类型:

None

visit_dataflow_var_(op)

Visit DataflowVar. Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op) on the C++ side.

Parameters

opDataflowVar

The DataflowVar to be visited.

参数:

op (DataflowVar)

返回类型:

None

visit_dataflow_var_def_(var)

Visit the DataflowVar definition site. Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var) on the C++ side.

Parameters

varDataflowVar

The DataflowVar to be visited.

参数:

var (DataflowVar)

返回类型:

None

visit_expr(expr)

Generic dispatcher for Expr. Users can customized this function to overwrite VisitExpr(const Expr& expr) on the C++ side.

Parameters

exprExpr

The expr to be visited.

参数:

expr (RelayExpr)

返回类型:

None

visit_extern_func_(op)

Visit ExternFunc. Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op) on the C++ side.

Parameters

opExternFunc

The ExternFunc to be visited.

参数:

op (ExternFunc)

返回类型:

None

visit_function_(op)

Visit Function. Users can customized this function to overwrite VisitExpr_(const FunctionNode* op) on the C++ side.

Parameters

opFunction

The Function to be visited.

参数:

op (Function)

返回类型:

None

visit_global_var_(op)

Visit GlobalVar. Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op) on the C++ side.

Parameters

opGlobalVar

The GlobalVar to be visited.

参数:

op (GlobalVar)

返回类型:

None

visit_if_(op)

Visit If. Users can customized this function to overwrite VisitExpr_(const IfNode* op) on the C++ side.

Parameters

opIf

The If to be visited.

参数:

op (If)

返回类型:

None

visit_match_cast_(binding)

Visit MatchCast. Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding) on the C++ side.

Parameters

bindingMatchCast

The MatchCast to be visited.

参数:

binding (MatchCast)

返回类型:

None

visit_op_(op)

Visit Op. Users can customized this function to overwrite VisitExpr_(const OpNode* op) on the C++ side.

Parameters

opOp

The Op to be visited.

参数:

op (Op)

返回类型:

None

visit_prim_value_(op)

Visit PrimValue. Users can customized this function to overwrite VisitExpr_(const PrimValueNode* op) on the C++ side.

Parameters

opPrimValue

The PrimValue to be visited.

参数:

op (PrimValue)

返回类型:

None

visit_seq_expr_(op)

Visit SeqExpr. Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op) on the C++ side.

Parameters

opSeqExpr

The SeqExpr to be visited.

参数:

op (SeqExpr)

返回类型:

None

visit_shape_expr_(op)

Visit ShapeExpr. Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op) on the C++ side.

Parameters

opShapeExpr

The ShapeExpr to be visited.

参数:

op (ShapeExpr)

返回类型:

None

visit_span(span)

Visit Span. Users can customized this function to overwrite VisitSpan(const Span& span) on the C++ side.

Parameters

spanSpan

The Span to be visited.

参数:

span (Span)

返回类型:

None

visit_string_imm_(op)

Visit StringImm. Users can customized this function to overwrite VisitExpr_(const StringImmNode* op) on the C++ side.

Parameters

opStringImm

The StringImm to be visited.

参数:

op (StringImm)

返回类型:

None

visit_tuple_(op)

Visit Tuple. Users can customized this function to overwrite VisitExpr_(const TupleNode* op) on the C++ side.

Parameters

opTuple

The Tuple to be visited.

参数:

op (Tuple)

返回类型:

None

visit_tuple_getitem_(op)

Visit TupleGetItem. Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op) on the C++ side.

Parameters

opTupleGetItem

The TupleGetItem to be visited.

参数:

op (TupleGetItem)

返回类型:

None

visit_var_(op)

Visit Var. Users can customized this function to overwrite VisitExpr_(const VarNode* op) on the C++ side.

Parameters

opVar

The Var to be visited.

参数:

op (Var)

返回类型:

None

visit_var_binding_(binding)

Visit VarBinding. Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding) on the C++ side.

Parameters

bindingVarBinding

The VarBinding to be visited.

参数:

binding (VarBinding)

返回类型:

None

visit_var_def(var)

Generic dispatcher for visiting the var definition site. Users can customized this function to overwrite VisitVarDef(const Var& var) on the C++ side. Note that visit_var_() will only visit the usage site of an Var.

Parameters

varVar

The var to be visited.

参数:

var (Var)

返回类型:

None

visit_var_def_(var)

Visit the Var definition site. Users can customized this function to overwrite VisitVarDef_(const VarNode* var) on the C++ side.

Parameters

varVar

The Var to be visited.

参数:

var (Var)

返回类型:

None

class tvm.relax.SeqExpr(blocks, body, span=None)

A sequence of binding blocks followed by an expression.

参数:

• blocks (List[BindingBlock])

• body (RelayExpr)

• span (Span | None)

class tvm.relax.ShapeExpr(values, span=None)

A shape expression which allows users to construct a shape containing PrimExpr.

Parameters

values: Union[List[PrimExpr], typing.Tuple[PrimExpr, ...], tvm.ir.Array]

The values of the shape expression.

span: Optional[Span]

Span that points to original source code

参数:

• values (List[PrimExpr])

• span (Span | None)

class tvm.relax.ShapeStructInfo(values=None, ndim=-1, span=None)

StructInfo of a shape value.

Parameters

valuesOptional[List[PrimExpr]]

The symbolic shape values if known.

ndimOptional[int]

The size of the shape.

Note

Do not specify values and ndim at the same time.

参数:

• values (List[PrimExpr] | None)

• ndim (int)

• span (Span)

class tvm.relax.ShapeType(ndim=-1, span=None)

The type of shape in Relax.

Parameters

ndimOptional[int]

The size of the shape.

参数:

• ndim (int)

• span (Span)

class tvm.relax.StringImm(value, span=None)

Represent a string literal constant.

参数:

• value (str)

• span (Span | None)

class tvm.relax.StructInfo

The base class of all StructInfo.

StructInfo contains both the static type and runtime structural information.

__eq__(other)

Compare two struct info for structural equivalence.

is_base_of(derived)

Check if self is base of another derived struct info.

Parameters

derivedStructInfo

The derived struct info to be checked.

Returns

resultbool

The check result.

参数:

derived (StructInfo)

返回类型:

bool

same_as(other)

Overload with structural equality.

class tvm.relax.TensorStructInfo(shape=None, dtype='float32', vdevice=None, ndim=-1, span=None)

StructInfo of a Tensor value.

Parameters

shapeOptional[Expr]

The shape expression.

dtypeOptional[str]

The content data type.

vdeviceOptional[Vdevice]

The virtual device.

ndimOptional[int]

The number of dimensions of the tensor.

Note

Do not specify shape and ndim at the same time.

参数:

• shape (RelayExpr | None)

• dtype (str)

• vdevice (VDevice | None)

• ndim (int)

• span (Span)

class tvm.relax.Tuple(fields, span=None)

Tuple expression that groups several fields together.

Parameters

fieldsUnion[List[Expr], typing.Tuple[Expr, ...]]

The fields in the tuple.

span: Optional[Span]

Span that points to original source code

参数:

• fields (List[RelayExpr])

• span (Span | None)

class tvm.relax.TupleGetItem(tuple_value, index, span=None)

Get index-th item from a tuple.

Parameters

tuple_value: Expr

The input tuple expression.

index: int

The index.

span: Optional[Span]

Span that points to original source code

参数:

• tuple_value (RelayExpr)

• index (int)

• span (Span | None)

class tvm.relax.TupleStructInfo(fields, span=None)

StructInfo of a Tuple value.

Parameters

fields: List[StructInfo]

The struct info of the fields.

参数:

• fields (List[StructInfo])

• span (Span)

class tvm.relax.VMInstrumentReturnKind(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

class tvm.relax.Var(name_hint, struct_info=None, span=None)

The variable class for all Relax bindings.

Parameters

name_hint: Union[str, Id]

The name hint of the variable.

struct_info: Optional[StructInfo]

The struct info annotation of the variable.

span: Optional[Span]

Span that points to original source code

property name_hint: str

Get name hint of the current var.

参数:

• name_hint (str | Id)

• struct_info (StructInfo | None)

• span (Span | None)

class tvm.relax.VarBinding(var, value, span=None)

Variable binding, bind he variable of the lhs with the rhs.

Parameters

var: Var

The return variable that the match cast bind to.

value: Expr

The input value expression.

参数:

• var (Var)

• value (RelayExpr)

• span (Span | None)

class tvm.relax.VirtualMachine(rt_mod, device, memory_cfg=None, profile=False)

Relax VM runtime.

参数:

• rt_mod (Module | tvm.relax.Executable)

• device (Device | List[Device])

• memory_cfg (str | Dict[Device, str] | None)

• profile (bool)

__init__(rt_mod, device, memory_cfg=None, profile=False)

Construct a VirtualMachine wrapper object.

Parameters

rt_mod: Union[tvm.runtime.Module, tvm.relax.Executable]

Runtime module exported by the result of build.

deviceUnion[Device, List[Device]]

The device to deploy the module.

memory_cfgOptional[Union[str, Dict[Device, str]]]

Config the type of memory allocator. The allocator type can be ["naive", "pooled"]. If memory_cfg is None, all devices will use pooled allocator by default. If memory_cfg is string, all devices will use the specified allocator type. If memory_cfg is a dict, each device uses the allocator type specified in the dict, or pooled allocator if not specified in the dict.

profileOptional[bool]

Whether or not to enable profiling.

参数:

• rt_mod (Module | Executable)

• device (Device | List[Device])

• memory_cfg (str | Dict[Device, str] | None)

• profile (bool)

返回类型:

None

_convert(arg, cargs)

helper function to convert arguments to vm function.

参数:

• arg (Any)

• cargs (List)

返回类型:

None

_convert_func_named_args(func_name, args, **kwargs)

Takes named function parameters and returns a list of those needed, in the order they should appear

参数:

• func_name (str)

• args (Any)

• kwargs (Any)

返回类型:

Any

_setup_device(dev, memory_cfg)

init devices and allocators.

参数:

• dev (Device)

• memory_cfg (str | Dict[Device, str])

返回类型:

None

get_outputs(func_name)

Get the value output by the function by the given name after a call of invoke_stateful.

It is an error to call this function without first calling invoke_stateful.

Parameters

func_name: str

The name of the function whose output should be fetched.

Returns

ret: Union[tvm.Object, Tuple[Any]]

The result of the earlier call to the function via invoke_stateful. If the result is a tuple, it returns a list of the fields. The fields are potentially also tuples, so these can be arbitrily nested.

参数:

func_name (str)

返回类型:

Object | Tuple[Any]

invoke_closure(closure, *args)

Invoke a closure.

Parameters

closureObject

The VMClosure Object.

argslist[tvm.runtime.NDArray] or list[np.ndarray]

The arguments to the closure.

Returns

resultObject

The output.

参数:

• closure (Object)

• args (Any)

返回类型:

Object

invoke_stateful(func_name)

Call the named function from the VM module using the arguments set using set_input. It is an error to call invoke_stateful without using set_input first (even if it's to set 0 inputs); conversely, if set_input has been called, it is an error to call the function without using invoke_stateful.

The results of the call can be obtained by calling get_outputs.

Parameters

func_name: str

The name of the function to call.

参数:

func_name (str)

返回类型:

None

profile(func_name, *args)

Profile a function call.

Parameters

func_namestr

The name of the function.

args: List of NDArray or other objects supported by PackedFunc.

The arguments to the function.

Returns

report: tvm.runtime.profiling.Report

The formatted profiling result, showing per-op timing measurements.

参数:

func_name (str)

save_function(func_name, saved_name, *args, include_return=True, **kwargs)

Convenience function. Takes a function from the module and saves a PackedFunc that, when called, will invoke the function with the given arguments. The PackedFunc can be accessed from the module using saved_name. This is included to facilitate timing trials: Invoking the returned PackedFunc will have less overhead from dictionary lookups than normally running through the VM.

If the saved name is taken, it can be overridden, though it cannot override the name of a function defined in the Relax source.

This is really creating a closure, but the function has a different name to avoid confusion with invoke_closure (they are not meant to be used together).

Parameters

func_namestr

The function that should be packaged up.

saved_namestr

The name that the resulting closure should be saved under.

include_returnbool

Whether the saved PackedFunc should return its output. If timing over RPC, it may not be desirable to send output between machines.

argsList[Any]

The arguments to package up with the function.

kwargsDict[str, Any]

Any named arguments to package up with the function

参数:

• func_name (str)

• saved_name (str)

• args (List[Any])

• include_return (bool)

• kwargs (Dict[str, Any])

返回类型:

None

set_input(func_name, *args, **kwargs)

Set the inputs to a function. This interface works when using VM over RPC by internally converting NDArray in the arguments to DLTensor, which is supported in RPC where remote could only have a minimal C runtime.

Note: If set_input is used, the function must be called using invoke_stateful and the results must be obtained using get_outputs.

Parameters

func_namestr

The name of the function.

args: List[tvm.runtime.NDArray] or List[np.ndarray]

The arguments to the function.

kwargs: dict of str to tvm.runtime.NDArray or np.ndarray

Named arguments to the function.

参数:

• func_name (str)

• args (Any)

• kwargs (Any)

返回类型:

None

set_instrument(instrument)

Set an instrumentation function.

If instrument is present, the function will be called before/after each Call instruction. The function have the following signature:

def instrument(
    func: Union[VMClosure, PackedFunc],
    func_symbol: str,
    before_run: bool,
    ret_value: any,
    *args) -> bool:
    pass

The instrument takes the following parameters: - func: function object to be called. - func_symbol: the symbol name of the function. - before_run: whether it is before or after call. - ret_value: the return value of the call, only valid after run. - args: the arguments being passed to call.

The instrument function can choose an integer, which corresponds to action direction for the following run. See VMInstrumentReturnKind for more details.

Parameters

instrument: tvm.runtime.PackedFunc

A instrumentation function that get invoked every VM call instr.

See Also

VMInstrumentReturnKind: the possible return values in VM.

参数:

instrument (PackedFunc)

返回类型:

None

time_evaluator(func_name, dev, number=10, repeat=1, min_repeat_ms=0, cooldown_interval_ms=0, repeats_to_cooldown=1, f_preproc='')

Returns an evaluator that times a function in the module. This follows the same convention as time_evaluator in tvm.runtime.module. This can be used in combination with save_function() so that the timings avoid extra dictionary lookups.

Parameters

func_name: str

The name of the function in the module.

dev: Device

The device we should run this function on.

number: int

The number of times to run this function for taking average. We call these runs as one repeat of measurement.

repeat: int, optional

The number of times to repeat the measurement. In total, the function will be invoked (1 + number x repeat) times, where the first one is warm up and will be discarded. The returned result contains repeat costs, each of which is an average of number costs.

min_repeat_ms: int, optional

The minimum duration of one repeat in milliseconds. By default, one repeat contains number runs. If this parameter is set, the parameters number will be dynamically adjusted to meet the minimum duration requirement of one repeat. i.e., When the run time of one repeat falls below this time, the number parameter will be automatically increased.

cooldown_interval_ms: int, optional

The cooldown interval in milliseconds between the number of repeats defined by repeats_to_cooldown.

repeats_to_cooldown: int, optional

The number of repeats before the cooldown is activated.

f_preproc: str, optional

The preprocess function name we want to execute before executing the time evaluator.

Note

The function will be invoked (1 + number x repeat) times, with the first call discarded in case there is lazy initialization.

Example

Normal use with a VM function (may not work over RPC if the function returns a tuple):

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)

Use with the stateful API:

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.set_input("func_name", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")

With saved closures via save_function (this results in fewer dictionary lookups in the timed portion):

target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(mod, tvm.cpu())
vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn)
timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()

Returns

ftimerfunction

The function that takes same argument as func and returns a BenchmarkResult. The ProfileResult reports repeat time costs in seconds.

参数:

• func_name (str)

• dev (Device)

• number (int)

• repeat (int)

• min_repeat_ms (int)

• cooldown_interval_ms (int)

• repeats_to_cooldown (int)

• f_preproc (str)

返回类型:

Callable[[...], BenchmarkResult]

tvm.relax.build(mod, target=None, params=None, pipeline='default_build', exec_mode='bytecode', *, system_lib=None)

Build an IRModule to VM executable.

Parameters

mod: IRModule

The input IRModule to be built.

targetOptional[Union[str, tvm.target.Target]]

A build target which can have optional host side compilation target.

When TVM compiles device specific program such as CUDA, we also need host(CPU) side code to interact with the driver to setup the dimensions and parameters correctly. host is used to specify the host side codegen target. By default, llvm is used if it is enabled, otherwise a stackvm interpreter is used.

params: Optional[Dict[str, list]]

Parameters for the input IRModule that will be bound.

pipelinestr = "default_build"

The compilation pipeline to use.

exec_mode: {"bytecode", "compiled"}

The execution mode.

system_lib: Optional[bool]

Whether to build system lib that is being packed statically and auto registers generated functions to the system. By default auto detects based on the target.

Returns

ex: tvm.relax.Executable

An executable that can be loaded by virtual machine.

Example

class InputModule:
    @R.function
    def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
        z = R.add(x, y)
        return z

mod = InputModule
target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(mod, target)

参数:

• mod (IRModule)

• target (str | Target | None)

• params (Dict[str, list] | None)

• pipeline (None | str | Pass)

• exec_mode (str)

• system_lib (bool | None)

返回类型:

Executable

tvm.relax.call_dps_packed(func, args, out_sinfo)

Call a destination-passing-style packed function and return the output.

Note: The called function is assumed to be _pure_ (other than modifying the designated output arguments). If the function _does_ result in other side effects, then the compiler may end up removing, reordering, or repeating those effects--no guarantees can be made.

Parameters

funcUnion[str, Expr]

The destination-passing-style function, can be ExternFunc.

argsExpr

The input arguments.

out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]

The structure info of the call_dps_packed output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.

Returns

ret: Call

A call node for the call_dps_packed operator.

参数:

• func (str | RelayExpr)

• args (RelayExpr)

• out_sinfo (TensorStructInfo | List[TensorStructInfo])

返回类型:

Call

tvm.relax.call_pure_packed(func, *args, sinfo_args)

Construct a call to a packed function that should be treated as pure, even though packed calls are normally not treated as pure.

The resulting call will have the same semantics as calling the packed function directly.

Note: This should be used for cases when the user knows that calling the packed function with these arguments will in reality not cause any side effects. If it is used for a call that does result in side effects, then the compiler may end up removing, reordering, or repeating that call, with no guarantees made about any side effects from the callee.

Parameters

funcUnion[str, ExternFunc]

The name (global symbol) for a PackedFunc or an ExternFunc node.

args: Expr

The arguments for the PackedFunc.

sinfo_args: Union[StructInfo, List[StructInfo]]

The list of structure info arguments (giving the structural info for the returned value).

Returns

resultExpr

A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)

参数:

• func (str | ExternFunc | GlobalVar)

• args (RelayExpr)

• sinfo_args (StructInfo | List[StructInfo])

返回类型:

RelayExpr

tvm.relax.call_tir(gvar, args, out_sinfo, tir_vars=None)

Call a tir.prim_func and return the output.

Parameters

gvarGlobalVar

The GlobalVar referring to a tir PrimFunc.

argsExpr

The input arguments.

out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]

The structure info of the call_tir output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.

tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]

ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns

ret: Call

A call node for the call_tir operator.

参数:

• gvar (GlobalVar)

• args (RelayExpr)

• out_sinfo (TensorStructInfo | List[TensorStructInfo])

• tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)

返回类型:

Call

tvm.relax.call_tir_inplace(gvar, args, inplace_indices, out_sinfo, tir_vars=None)

Call a TIR PrimFunc and return the result, doing the specified computations in-place (based on the inplace_indices argument; outputs will alias the inputs selected by in-place indices).

Warning: This operator is considered pure by the type system but actually mutates the arguments specified by inplace_indices. This operator should not be used directly, but rather should be inserted by passes that have checked whether it is safe to perform operations in-place (i.e., none of the arguments specified as an output is aliased or is live after calling call_tir_inplace).

Direct calls to this operator should be done for testing purposes only.

Parameters

gvarGlobalVar

The GlobalVar referring to a TIR PrimFunc.

argsExpr

The input arguments.

inplace_indicesUnion[int, List[int]]

Specify which arguments should be used for in-place computations. If inplace_indices is a single integer, it will be made into a singleton list. Suppose inplace_indices[i] = j, where j >= 0. Then the i`th output will be an alias of `args[j]. If inplace_indices[i] = -1, then the i`th output will be a freshly allocated tensor. At least one member of `inplace_indices must not be -1.

out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]

The structure info of the call_tir_inplace output. It should be a single TensorStructInfo or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor. If a list of TensorStructInfo is given, the result will be a tuple of TensorStructInfo.

tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]

ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns

ret: Call

A call node for the call_tir operator.

参数:

• gvar (GlobalVar)

• args (RelayExpr)

• inplace_indices (int | List[int])

• out_sinfo (TensorStructInfo | List[TensorStructInfo])

• tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)

返回类型:

Call

tvm.relax.call_tir_with_grad(gvar, args, out_sinfo, te_grad_name, te_grad_kwargs=None, tir_vars=None)

Call a tir.prim_func and return the output. This intrinsic will bind a te gradient function (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be called by the Gradient pass.

Parameters

gvarGlobalVar

The GlobalVar referring to a tir PrimFunc.

argsExpr

The input arguments.

out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]

The structure info of the call_tir_with_grad output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.

te_grad_namestr

The registered name of the te gradient function associated with the call_tir_with_grad node. Must be provided as a keyword argument.

te_grad_kwargsDict[str, Object], optional

The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}.

tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]

ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns

ret: Call

A call node for the call_tir_with_grad operator.

参数:

• gvar (GlobalVar)

• args (RelayExpr)

• out_sinfo (TensorStructInfo | List[TensorStructInfo])

• te_grad_name (str)

• te_grad_kwargs (Dict[str, Object])

• tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)

返回类型:

Call

tvm.relax.const(value, dtype=None)

Create a constant value.

Parameters

value: Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]

The constant value.

dtype: Optional[str]

The data type of the resulting constant.

Note

When dtype is None, we use the following rule:

• int maps to "int32"

• float maps to "float32"

• bool maps to "bool"

• other using the same default rule as numpy.

参数:

• value (bool | int | float | ndarray | NDArray)

• dtype (str | None)

返回类型:

Constant

tvm.relax.convert_to_expr(value)

Helper function to convert the input to Expr, which follows the rules: 1. Return the input itself if it's already a relax.Expr; 2. Return relax.PrimValue if the input is a PrimExpr; 3. Return relax.StringImm if the input is tvm.String or str; 4. Return relax.Tuple if the input is a tuple/list of Expr.

Notes

1. tvm.tir.StringImm is not allowed because of ambiguity, which can be either relax.StringImm or relax.PrimValue.

参数:

value (Any)

返回类型:

RelayExpr

tvm.relax.extern(name, struct_info=None, span=None)

Create extern function.

参数:

• name (str)

• struct_info (StructInfo | None)

• span (Span | None)

tvm.relax.get_pipeline(name='zero', **kwargs)

Get pre-build pipeline by name

Parameters

nameOptional[str]

Name of the pipeline

kwargsDict[str, object]

Keyword args for configuring the pipeline.

Returns

pipeline: tvm.transform.Pass

The transformation pipeline.

参数:

name (str)

返回类型:

Pass

tvm.relax.get_shape_of(expr)

Get shape of expr.

Parameters

expr: Expr

The input expr.

Returns

shape: Expr

The shape expression

Note

This function requires expr to be normalized. The function will report an error if expr's StructInfo is not TensorStructInfo. It will try to return symbolic function when possible. If the tensor do not have a compile-time symbolic shape, the function will then choose to return Call(relax.op.shape_of, [expr]).

参数:

expr (RelayExpr)

返回类型:

RelayExpr

tvm.relax.register_pipeline(name)

Register a new pipeline

Parameters

namestr

Name of the pipeline

参数:

name (str)