tvm.relay.transform#

The Relay IR namespace containing transformations.

Classes:

ChangeBatch(data[, batch_size])

Change the batch size.

FlexibleShapeDispatch(buckets[, axis, ...])

Enable inference of multiple shaped inputs in one module.

FunctionPass()

A pass that works on each tvm.relay.Function in a module.

LayoutConfig([skip_layers])

A structure for customizing the ConvertLayout pass.

Functions:

AlterOpLayout()

Alternate the layouts of operators or replace primitive operators with other expressions.

AnnotateSpans()

Annotate a program with span information by first generating its textual representation and then parsing it back into a Relay AST annotated with span information.

AnnotateTarget(targets[, include_non_call_ops])

Annotate ops in an experession with a provied compiler/target and then use it for codegen.

BackwardFoldScaleAxis()

Backward fold axis scaling into weights of conv2d/dense.

BatchingOps()

Batching parallel operators into one for Conv2D, Dense and BatchMatmul.

CanonicalizeCast()

Canonicalize cast expressions to make operator fusion more efficient.

CanonicalizeOps()

Canonicalize special operators to basic operators.

CapturePostDfsIndexInSpans()

Captures the post-dfs index and dominator post-dfs index of (most) expression nodes in their span, in the form "index:<post-dfs index>:<dominator post-dfs index>".

CollagePartition(config[, cost_estimator])

Partition the bodies of all functions according to the available targets so as to minimize model latency.

CombineParallelBatchMatmul([min_num_branches])

Combine multiple batch matmul operators into one.

CombineParallelConv2D([min_num_branches])

Combine multiple conv2d operators into one.

CombineParallelDense([min_num_branches, ...])

Combine multiple dense operators into one.

Conv2dToSparse(weight_name, weight_shape, ...)

Rewrite qualified `nn.conv2d operation` to `nn.sparse_conv2d`

Conv2dToSparse2(layout, kernel_size, ...)

Rewrite freezed `nn.conv2d` operation to `nn.sparse_conv2d`

ConvertLayout(desired_layouts)

Given a dest layout, this pass transforms the expr such that most of the ops input data layout is changed to the dest layout.

DeadCodeElimination([inline_once, ...])

Remove expressions that do not have any users (dead code).

Defunctionalization(func, mod)

Performs defunctionalization on func, transforming func from a higher-order program to a first-order program.

DefuseOps()

The inverse operation of FuseOps.

DenseToSparse(weight_name, weight_shape)

Rewrite qualified `nn.dense operation` to `nn.sparse_dense` This pass is used in `data_dep_optimization.bsr_dense` Parameters of this pass is generated by `analysis.sparse_dense.process_params`

DivToMul()

Transform division by a constant to multiplication by the inverse of the constant

DynamicToStatic()

If possible, convert tvm.relay.dynamic* ops to static versions

EliminateCommonSubexpr([fskip])

Eliminate common subexpressions.

EtaExpand([expand_constructor, ...])

Add abstraction over a constructor or global variable bound to a function

FakeQuantizationToInteger([hard_fail, ...])

Find regions of the graph of the form

FastMath()

Converts the expensive non linear functions to their fast but approximate counterparts.

FirstOrderGradient()

Transforms all global functions in the module to return the original result, paired with the gradients of the inputs.

FlattenAtrousConv()

The purpose of this pass is to find a sequence of space_to_batch_nd-conv2d-batch_to_space_nd operations:

FoldConstant([fold_qnn])

Fold the constant expressions in a Relay program.

FoldConstantExpr(expr, mod[, fold_qnn])

Fold the constant expressions in a Relay program. Parameters ---------- expr: Expr The expression to fold mod: IRModule The module the expr lives in (for global calls) fold_qnn: bool Whether to fold constants for QNN operations.

FoldExplicitPadding()

FoldExplicitPadding finds explict padding before an op that can support implicit padding and fuses them.

FoldScaleAxis()

Fold the scaling of axis into weights of conv2d/dense.

ForwardFoldScaleAxis()

Fold the scaling of axis into weights of conv2d/dense.

FuseOps([fuse_opt_level])

Fuse operators in an expr to a larger operator according to some rules.

InferType()

Infer the type of an expr.

InferTypeLocal(expr)

Infer the type of a single expr, reusing type information to do so.

Inline()

Perform inlining on the given Relay IR module.

InlineCompilerFunctionsBoundTo(global_vars)

Inlines all global functions bound to a global var in global_vars.

LambdaLift()

Lift the closure to global function.

LazyGradientInit()

Reduces memory usage of gradient tensors

Legalize([legalize_map_attr_name])

Legalizes an expression with another expression.

ManifestLifetimes()

Manifest the lifetimes of variables after allocations have been manifested, by inserting kill operations once variables become dead.

MarkCompilerFunctionsAsExtern([compiler_filter])

Marks all global functions which have a "Compiler" attribute matching compiler_filter as 'extern'.

MergeCompilerRegions()

Merge together compiler regions.

MergeComposite(pattern_table)

Merge multiple operators into a single composite relay function.

OutlineCompilerFunctionsWithExistingGlobalSymbols([...])

Outlines all literal functions in direct call positions which have a "Compiler" attribute.

PartialEvaluate()

Evaluate the static fragment of the code.

PartitionGraph([mod_name, bind_constants])

Partition a Relay program into regions that can be executed on different backends.

PlanDevices(config)

Uses existing "on_device" and "device_copy" calls to infer the virtual device on which every Relay sub-expression should run and the result stored.

RemoveUnusedFunctions([entry_functions])

Remove unused global relay functions in a relay module.

SimplifyExpr()

Simplify the Relay expression, including merging consecutive reshapes.

SimplifyFCTranspose(target_weight_name)

Rewrite `y = nn.dense(x, transpose(w, [1, 0]))` to `y = nn.dense(x, wt)` This pass is used in `data_dep_optimization.simplify_fc_transpose`

SimplifyInference()

Simplify the data-flow graph for inference phase.

SplitArgs(max_function_args)

Split function with huge number of arguments to smaller pieces.

ToANormalForm()

Turn Graph Normal Form expression into A Normal Form Expression.

ToANormalFormExpr(e)

ToANormalForm, but on expression level.

ToBasicBlockNormalForm()

Turn an expression to Basic Block Normal Form.

ToCPS(expr[, mod])

Turn expression into continuation passing style(CPS).

ToGraphNormalForm()

Turn a Relay program in A Normal Form into Graph Normal Form

ToMixedPrecision([mixed_precision_type, ...])

Automatic mixed precision rewriter.

build_config([opt_level, required_pass, ...])

Configure the build behavior by setting config variables.

function_pass([pass_func, opt_level, name, ...])

Decorate a function pass.

gradient(expr[, mod, mode])

Transform the input function, returning a function that calculate the original result, paired with gradient of the input.

recast(expr, dtype, out_dtype[, ops, ...])

Convert the types of operations in a graph to a new value.

to_cps(func[, mod])

Turn expression into CPS expression.

un_cps(func)

Turn an cps function into a Function without the continuation argument.

class tvm.relay.transform.ChangeBatch(data, batch_size=16)[源代码]#

Change the batch size.

Parameters#

data: Dict[relay.Var, int]

A dictionary of all the params to change. The keys are all params, and the values are which dimension hold the batch.

batch_size: int

The batch size to change to.

Returns#

pass: FunctionPass

The pass.

class tvm.relay.transform.FlexibleShapeDispatch(buckets, axis=0, auto_pad=False, pad_value=0, input_indices=None, affects_output=True)[源代码]#

Enable inference of multiple shaped inputs in one module.

This transformation adds a handler around a module that checks input shapes and dispatches to a subgraph specialized to handle the specific shapes of that input. If no exactly matching subgraph is available, the input will be run using full dynamism. For best performance, specify all the sizes the module will be likely to see using the buckets argument.

By default, this pass will dispatch shapes that exactly match one of the buckets to a corresponding subgraph. All non-matching shapes use the same fully dynamic fallback. This can be detrimental to performance for those non-matching shapes. Setting auto_pad to True causes this pass to round-up the shape of non-matching inputs to the closest bucket. This allows them to use the tuned kernels of bucket shapes which can improve performance.

Models that have multiple inputs sharing a dynamic axis, which is common for batch size or sequence length dynamism, are supported through the input_indices argument.

Many types of dynamism such as batching affect both the input and output shape, however this is not always the case. If the output shape is independent of the input, the affects_output argument of this pass must be set to False.

Parameters#

buckets: list[int]

The sizes of the input dimension that should be explicitly handled. Each value in buckets will have a corresponding subgraph constructed to handle it.

axis: int

The dimension of the input that should be made flexible. This will most often be used for the batch dimension.

auto_pad: Optional[bool]

If True, then padding will be inserted to values that don’t match one of the provided buckets.

pad_value: Optional[float]

When auto_pad is true, padding will be done with this value.

input_indices: Optional[List[int]]

Which inputs should be dispatched dynamically, provided by index. All inputs must share the same dynamic axis.

affects_output: Optional[bool]

Whether the change in input shape has a corresponding effect on the output shape. Batching for example effects both the input and output whereas changing sequence length in an NLP model typically does not.

Returns#

retFlexibleShapeDispatch

A pass that can be applied to a module to add flexible shape handling.

class tvm.relay.transform.FunctionPass[源代码]#

A pass that works on each tvm.relay.Function in a module. A function pass class should be created through function_pass.

class tvm.relay.transform.LayoutConfig(skip_layers=None)[源代码]#

A structure for customizing the ConvertLayout pass.

tvm.relay.transform.AlterOpLayout()[源代码]#

Alternate the layouts of operators or replace primitive operators with other expressions. This pass can be used for computing convolution in custom layouts or other general weight pre-transformation.

Returns#

rettvm.transform.Pass

The registered pass that alters the layout of operators.

tvm.relay.transform.AnnotateSpans()[源代码]#

Annotate a program with span information by first generating its textual representation and then parsing it back into a Relay AST annotated with span information.

Returns#

rettvm.transform.Pass

The registered AnnotateSpans pass.

tvm.relay.transform.AnnotateTarget(targets, include_non_call_ops=True)[源代码]#

Annotate ops in an experession with a provied compiler/target and then use it for codegen.

Parameters#

targetsstr or List[str]

The list of target compilers used for codegen.

include_non_call_opsboolean

If True then non-call ops also will be annotated with targets If False then non-call ops will not be processed

Returns#

rettvm.transform.Pass

The annotated pass that wrapps ops with subgraph_start and subgraph_end.

tvm.relay.transform.BackwardFoldScaleAxis()[源代码]#

Backward fold axis scaling into weights of conv2d/dense.

Returns#

rettvm.transform.Pass

The registered pass to backward fold expressions.

Note#

It is recommended to call backward_fold_scale_axis before using forward_fold_scale_axis as backward folding targets the common conv->bn pattern.

tvm.relay.transform.BatchingOps()[源代码]#

Batching parallel operators into one for Conv2D, Dense and BatchMatmul.

Returns#

ret: tvm.transform.Pass

The sequential pass which apply batching for different operator types.

tvm.relay.transform.CanonicalizeCast()[源代码]#

Canonicalize cast expressions to make operator fusion more efficient.

Returns#

rettvm.transform.Pass

The registered pass that canonicalizes cast expression.

tvm.relay.transform.CanonicalizeOps()[源代码]#

Canonicalize special operators to basic operators. This can simplify followed analysis, e.g. expanding bias_add to expand_dims and broadcast_add.

Returns#

ret: tvm.transform.Pass

The registered pass performing the canonicalization.

tvm.relay.transform.CapturePostDfsIndexInSpans()[源代码]#

Captures the post-dfs index and dominator post-dfs index of (most) expression nodes in their span, in the form “index:<post-dfs index>:<dominator post-dfs index>”.

This is useful for debugging since a) it helps identify pretty-printed sub-expressions within the overall model and b) the indexes are heavily used by Collage for its compact representation of sub-graphs.

Note that Op and Constructor nodes are not changed even though they are assigned an post-dfs index.

Returns#

rettvm.transform.Pass

The pass.

tvm.relay.transform.CollagePartition(config, cost_estimator=None)[源代码]#

Partition the bodies of all functions according to the available targets so as to minimize model latency. See apache/tvm-rfcs.

Parameters#

configCompilationConfig

The available targets.

cost_estimatorCostEstimator, optional

The custom cost estimator to use for costing each candidate partition.

Returns#

rettvm.transform.Pass

The pass.

tvm.relay.transform.CombineParallelBatchMatmul(min_num_branches=3)[源代码]#

Combine multiple batch matmul operators into one. For example:

Would become:

Parameters#

min_num_branchesint

The minimum number of required parallel branches for performing this optimization.

Returns#

ret: tvm.transform.Pass

The registered pass that combines parallel dense operators.

tvm.relay.transform.CombineParallelConv2D(min_num_branches=3)[源代码]#

Combine multiple conv2d operators into one.

Parameters#

min_num_branchesint

The minimum number of required parallel branches for performing this optimization.

Returns#

ret: tvm.transform.Pass

The registered pass that combines parallel conv2d operators.

tvm.relay.transform.CombineParallelDense(min_num_branches=3, to_batch=True)[源代码]#

Combine multiple dense operators into one. For example:

Would become:

or (if to_batch=False)

Parameters#

min_num_branchesint

The minimum number of required parallel branches for performing this optimization.

to_batch_matmulbool

If True, combine parallel dense ops into batch_matmul op. If False, combine parallel dense ops into dense op.

Returns#

ret: tvm.transform.Pass

The registered pass that combines parallel dense operators.

tvm.relay.transform.Conv2dToSparse(weight_name, weight_shape, layout, kernel_size)[源代码]#

Rewrite qualified `nn.conv2d operation` to `nn.sparse_conv2d`

Parameters#

weight_name: Array[String]

Names of weights which qualified sparse contrains

weight_shape: Array[Array[IntImm]]

Weights shape in BSR format.

layoutstr

layout of data

Returns#

rettvm.transform.Pass

The registered DenseToSparse pass.

tvm.relay.transform.Conv2dToSparse2(layout, kernel_size, blocksize, sparsity_threshold)[源代码]#

Rewrite freezed `nn.conv2d` operation to `nn.sparse_conv2d`

Parameters#

layoutstr

layout of data

kernel_sizeint

kernel size of conv2d

Returns#

rettvm.transform.Pass

The registered DenseToSparse pass.

tvm.relay.transform.ConvertLayout(desired_layouts)[源代码]#

Given a dest layout, this pass transforms the expr such that most of the ops input data layout is changed to the dest layout. In ideal situation, there are only 2 layout transforms, one at the start and one at the end.

This pass is not a part of relay.build and is expected to be called between framework-relay parser and relay.build call. This is very helpful for hardware backends that support/prefer only type of data layout.

RFC - https://discuss.tvm.apache.org/t/layout-conversion-pass/4009

This pass uses most of the AlterOpLayout and InferCorrectLayout infrastructure. We can define new layouts for conv2d ops for now. Most of the other operators try to adapt to their input layout using the InferCorrectLayout infrastructure.

Parameters#

desired_layoutsmap of op_name to list of layouts

Specify a mapping of operator names to a list of layouts to convert to, in the order defined by the operator. An example for nn.conv2d could be: {“nn.conv2d”, [“NHWC”, “OHWI]}, where the first item in the list specifies the data layout and the second specifies the kernel layout.

Returns#

pass: FunctionPass

The pass.

tvm.relay.transform.DeadCodeElimination(inline_once=False, ignore_impurity=False)[源代码]#

Remove expressions that do not have any users (dead code).

Parameters#

inline_once: Optional[Bool]

Whether to inline a binding that is referenced exactly once.

ignore_impurity: Optional[Bool]

Whether to ignore possible side-effects in let-bound expressions.

Returns#

ret: tvm.transform.Pass

The registered pass that eliminates the dead code in a Relay program.

tvm.relay.transform.Defunctionalization(func, mod)[源代码]#

Performs defunctionalization on func, transforming func from a higher-order program to a first-order program.

At each call site, the function is cloned and type parameters are substituted in. Function arguments are encoded as datatypes and additional apply functions are used for application.

Parameters#

functvm.relay.Function

The input function, which should not be polymorphic or be higher-order. This is because all types must be known and we can’t encode function arguments to the program itself.

modtvm.IRModule

The IRModule containing function and type definitions, which is also mutated during this pass.

Returns#

exprtvm.relay.Function

The output function.

tvm.relay.transform.DefuseOps()[源代码]#

The inverse operation of FuseOps. It transforms a fused program returned by FuseOps into the program before FuseOps. (i.e., x == DefuseOps(FuseOps(x)))

Returns#

rettvm.transform.Pass

The registered pass for operator defusion.

tvm.relay.transform.DenseToSparse(weight_name, weight_shape)[源代码]#

Rewrite qualified `nn.dense operation` to `nn.sparse_dense` This pass is used in `data_dep_optimization.bsr_dense` Parameters of this pass is generated by `analysis.sparse_dense.process_params`

Parameters#

weight_name: Array[String]

Names of weights which qualified sparse contrains

weight_shape: Array[Array[IntImm]]

Weights shape in BSR format.

Returns#

rettvm.transform.Pass

The registered DenseToSparse pass.

tvm.relay.transform.DivToMul()[源代码]#

Transform division by a constant to multiplication by the inverse of the constant

tvm.relay.transform.DynamicToStatic()[源代码]#

If possible, convert tvm.relay.dynamic* ops to static versions

Returns#

rettvm.transform.Pass

The registered pass for dynamic->static conversion.

tvm.relay.transform.EliminateCommonSubexpr(fskip=None)[源代码]#

Eliminate common subexpressions.

Parameters#

fskip: Callable

The callback function that decides whether an expression should be skipped.

Returns#

rettvm.transform.Pass

The registered pass that eliminates common subexpressions.

tvm.relay.transform.EtaExpand(expand_constructor=False, expand_global_var=False)[源代码]#

Add abstraction over a constructor or global variable bound to a function

Parameters#

expand_constructor: bool

Whether to expand constructors.

expand_global_var: bool

Whether to expand global variables.

Returns#

ret: tvm.transform.Pass

The registered pass that eta expands an expression.

tvm.relay.transform.FakeQuantizationToInteger(hard_fail=False, use_qat=False, optional_qnn_ops=None)[源代码]#

Find regions of the graph of the form

x    w
|    |
dq   dq
 \   /
  op1
   |
  op2
   |
   q

where q == qnn.quantize and dq = qnn.dequantize and rewrite them into integer versions of op1 and op2

Rules for rewriting indivdual ops are in fake_quantization_to_integer.py

Parameters#

hard_failboolean

How do deal with errors during graph rewriting. If true, raise an error. If false, skip rewriting the subgraph.

use_qatboolean

To perform an additional QAT pass - convert enabled operations with dequantized inputs. Example: in the graph above op2 is not registered with the FakeQuantizationToInteger attribute, op1 operation can still be converted. Converted pattern below:

x    w
|    |
\   /
  op1
  |
  dq
  |
  op2
  |
  q
optional_qnn_opsList[str]

Specify a list of operator names to explicitly enable conversion for specific ops disabled by default. Example: [‘nn.softmax’]

Returns#

rettvm.transform.Pass

The registered FakeQuantizationToInteger pass.

tvm.relay.transform.FastMath()[源代码]#

Converts the expensive non linear functions to their fast but approximate counterparts.

Returns#

ret: tvm.transform.Pass

The registered pass to perform fast math operations.

tvm.relay.transform.FirstOrderGradient()[源代码]#

Transforms all global functions in the module to return the original result, paired with the gradients of the inputs. This pass transforms each global function independently and does not support interprocedural AD. Additionally, this pass does not support any control-flow or references, and should only be used on pure data-flow graphs.

Returns#

rettvm.transform.Pass

The registered FirstOrderGradient pass.

tvm.relay.transform.FlattenAtrousConv()[源代码]#

The purpose of this pass is to find a sequence of space_to_batch_nd-conv2d-batch_to_space_nd operations:

x     w
|     |
s2b   |
 \   /
  conv2d
   |
   b2s

and convert them into subgraphs with a convolution with the modified “dilation” and recalculated “padding” parameters.

Returns#

rettvm.transform.Pass

The registered FlattenAtrousConv pass.

tvm.relay.transform.FoldConstant(fold_qnn=False)[源代码]#

Fold the constant expressions in a Relay program.

Because of backward compatibility reason it skips QNN primitives from folding by default. There are some transformation passes like FakeQuantizationToInteger, which requires to keep QNN primitives for constant subgraphs. Uncontrolled constant folding of QNN primitives may break applicability of FakeQuantizationToInteger. We suggest to use FoldConstant pass with none default fold_qnn=True value only when all other QNN sensitive passes were already applied.

Parameters#

fold_qnn: bool

Whether to fold constants for QNN operations.

Returns#

rettvm.transform.Pass

The registered pass for constant folding.

tvm.relay.transform.FoldConstantExpr(expr, mod, fold_qnn=False)[源代码]#

Fold the constant expressions in a Relay program. Parameters ———- expr: Expr

The expression to fold

mod: IRModule

The module the expr lives in (for global calls)

fold_qnn: bool

Whether to fold constants for QNN operations.

Returns#

new_expr: Expr

The expr after Constant Folding

tvm.relay.transform.FoldExplicitPadding()[源代码]#

FoldExplicitPadding finds explict padding before an op that can support implicit padding and fuses them.

Returns#

rettvm.transform.Pass

The registered ImplicitPadding pass.

tvm.relay.transform.FoldScaleAxis()[源代码]#

Fold the scaling of axis into weights of conv2d/dense. This pass will invoke both forward and backward scale folding.

Returns#

rettvm.transform.Pass

The registered pass to fold expressions.

Note#

Internally, we will call backward_fold_scale_axis before using forward_fold_scale_axis as backward folding targets the common conv->bn pattern.

tvm.relay.transform.ForwardFoldScaleAxis()[源代码]#

Fold the scaling of axis into weights of conv2d/dense.

Returns#

rettvm.transform.Pass

The registered pass to forward fold expressions.

Note#

It is recommended to call backward_fold_scale_axis before using forward_fold_scale_axis, as backward folding targets the common conv->bn pattern.

tvm.relay.transform.FuseOps(fuse_opt_level=-1)[源代码]#

Fuse operators in an expr to a larger operator according to some rules.

Parameters#

fuse_opt_levelint

The level of fuse optimization. -1 indicates that the level will be inferred from pass context.

Returns#

rettvm.transform.Pass

The registered pass for operator fusion.

tvm.relay.transform.InferType()[源代码]#

Infer the type of an expr.

Returns#

rettvm.transform.Pass

The registered type inference pass.

tvm.relay.transform.InferTypeLocal(expr)[源代码]#

Infer the type of a single expr, reusing type information to do so.

This populates the checked_type field in expr. We assume existing type information in the graph is correct!

Parameters#

expr: relay.Expr

The expression we want to know the type of

Returns#

type: relay.Type

The type of the expression

tvm.relay.transform.Inline()[源代码]#

Perform inlining on the given Relay IR module. The global functions that are marked as inline should be always inlined. A cost model will be needed in the future to decide if it is profitable to inline the function.

Returns#

ret: tvm.transform.Pass

The registered pass that performs inlining for a Relay IR module.

tvm.relay.transform.InlineCompilerFunctionsBoundTo(global_vars)[源代码]#

Inlines all global functions bound to a global var in global_vars.

Both the global “Compiler” attributed function, and any calls to “Composite” functions it its body are inlined.

This pass may be useful for external codegen which needs to undo partitioning based on properties of the entire partition.

Parameters#

global_varsArray[tvm.relay.GlobalVar]

The global vars of all ‘Compiler’ functions to inline.

Returns#

rettvm.transform.Pass

The pass.

tvm.relay.transform.LambdaLift()[源代码]#

Lift the closure to global function.

Returns#

rettvm.transform.Pass

The registered pass that lifts the lambda function.

tvm.relay.transform.LazyGradientInit()[源代码]#

Reduces memory usage of gradient tensors

Parameters#

Returns#

ret: tvm.transform.Pass

A pass which delays and/or reduces memory allocation, by lazily allocating 0 or one filled tensors.

tvm.relay.transform.Legalize(legalize_map_attr_name='FTVMLegalize')[源代码]#

Legalizes an expression with another expression. This pass can be used to replace an expr with another expr for target dependent optimizations. For example, one expr, though semnatically equivalent to the other, can have better performance on a target. This pass can be used to legalize the expr in a target-dependent manner.

Parameters#

legalize_map_attr_namestr

The Op’s attr name which corresponds to the legalize rule function.

Returns#

rettvm.transform.Pass

The registered pass that rewrites an expr.

tvm.relay.transform.ManifestLifetimes()[源代码]#

Manifest the lifetimes of variables after allocations have been manifested, by inserting kill operations once variables become dead.

tvm.relay.transform.MarkCompilerFunctionsAsExtern(compiler_filter='')[源代码]#

Marks all global functions which have a “Compiler” attribute matching compiler_filter as ‘extern’.

The function’s attributes are replaced with a single “Extern” attribute, and all calls to the function are switched to use the ‘call_lowered’ calling convention.

If compiler_filter is non-empty only functions with that as their attribute value are outlined.

This pass may be useful for external codegen using the “RelayToTIR” custom pass mechanism to cleanup the IRModule after custom lowering.

Parameters#

compiler_filterString

If non-empty, the “Compiler” attribute to filter on.

Returns#

rettvm.transform.Pass

The pass.

tvm.relay.transform.MergeCompilerRegions()[源代码]#

Merge together compiler regions.

Returns#

rettvm.transform.Pass

The registered pass that merges compiler regions.

tvm.relay.transform.MergeComposite(pattern_table)[源代码]#

Merge multiple operators into a single composite relay function.

Parameters#

pattern_tableList[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Function]]

A list of (pattern_name, pattern, check) tuples. The order of the patterns in the list will determine the order of priority in which they are matched. ‘check’ is a function to check whether an extracted pattern matches. It can be implemented by pattern writer but if not specified it will always return True.

Returns#

rettvm.transform.Pass

The registered pass that merges operators into a single composite relay function.

tvm.relay.transform.OutlineCompilerFunctionsWithExistingGlobalSymbols(compiler_filter='')[源代码]#

Outlines all literal functions in direct call positions which have a “Compiler” attribute.

The outlined functions are bound to unique global vars according to their existing “global_symbol” attribute. At most one function with the same global symbol is outlined.

If compiler_filter is non-empty only functions with that as their attribute value are outlined.

This pass may be useful for external codegen using the “RelayToTIR” custom pass mechanism to prepare the IRModule before custom lowering.

Parameters#

compiler_filterString

If non-empty, the “Compiler” attribute to filter on.

Returns#

rettvm.transform.Pass

The pass.

tvm.relay.transform.PartialEvaluate()[源代码]#

Evaluate the static fragment of the code.

Note#

This transformation could be either Module -> Module or Expr -> Expr. It will directly transform the input expression to a new one if the target expression is provided. Otherwise, it will rely on the pass manager to carry out transformation.

Returns#

ret: tvm.transform.Pass

The registered pass that performs partial evaluation on an expression.

tvm.relay.transform.PartitionGraph(mod_name='default', bind_constants=True)[源代码]#

Partition a Relay program into regions that can be executed on different backends.

Parameters#

mod_namestring

Controls the prefix of the name of each partitioned subraph. If mod_name is None, then tvmgen_ prefix is used. Otherwise, tvmgen_mod_name_ prefix is used.

bind_constants: bool

Whether or not to bind constants in partitioned subgraphs. Note that the codegen needs to maintain the bound constants; Otherwise the constants will be maintained by the metadata module. So it is recommended for C-source based codegens to set bind_constants=False to avoid embedding large constants in a C source file.

Returns#

ret: tvm.transform.Pass

The registered pass that partitions the Relay program.

tvm.relay.transform.PlanDevices(config)[源代码]#

Uses existing “on_device” and “device_copy” calls to infer the virtual device on which every Relay sub-expression should run and the result stored. Captures the result of that analysis using new “on_device” and “device_copy” calls. Sub-expressions which are not otherwise constrained are assigned to the default primitive virtual device describe by config. However data and computations which must be hosted on a CPU (such as shapes and shape functions) use the host virtual device of the config.

Parameters#

configtvm.CompilationConfig

The compilation configuration, specifying available targets and default devices.

Returns#

rettvm.transforms.Pass

The pass.

tvm.relay.transform.RemoveUnusedFunctions(entry_functions=None)[源代码]#

Remove unused global relay functions in a relay module.

Parameters#

entry_functions: list[string]

The set of entry functions to start from.

Returns#

rettvm.transform.Pass

The registered pass to remove unused functions.

tvm.relay.transform.SimplifyExpr()[源代码]#

Simplify the Relay expression, including merging consecutive reshapes.

Returns#

rettvm.transform.Pass

The registered SimplifyExpr pass.

tvm.relay.transform.SimplifyFCTranspose(target_weight_name)[源代码]#

Rewrite `y = nn.dense(x, transpose(w, [1, 0]))` to `y = nn.dense(x, wt)` This pass is used in `data_dep_optimization.simplify_fc_transpose`

Parameters#

weight_name: Array[String]

Names of weights which qualified `y = nn.dense(x, transpose(w, [1, 0]))` This parameter is generated by `analysis.search_fc_transpose` function

Returns#

rettvm.transform.Pass

The registered SimplifyFCTranspose pass.

tvm.relay.transform.SimplifyInference()[源代码]#

Simplify the data-flow graph for inference phase. An simplified expression which is semantically equal to the input expression will be returned.

Note that batch norms will only be simplified if their result is indexed at tuple index 0.

Returns#

ret: tvm.transform.Pass

The registered pass to perform operator simplification.

tvm.relay.transform.SplitArgs(max_function_args)[源代码]#

Split function with huge number of arguments to smaller pieces.

Parameters#

max_function_args: int

Maximum number of function arguments. If it equals 0 then SplitArgs shouldn’t split the function.

Returns#

rettvm.transform.Pass

The registered pass.

tvm.relay.transform.ToANormalForm()[源代码]#

Turn Graph Normal Form expression into A Normal Form Expression. The scope of the root expression is the global scope. The scope of any non root expression is the least common ancestor of all it’s scope. Values are ordered by post-DFS order in each scope.

Returns#

retUnion[tvm.transform.Pass, tvm.relay.Expr]

The registered pass that transforms an expression into A Normal Form.

tvm.relay.transform.ToANormalFormExpr(e)[源代码]#

ToANormalForm, but on expression level.

Parameters#

eExpr

The graph expression.

Returns#

retExpr

The transformed expresion.

tvm.relay.transform.ToBasicBlockNormalForm()[源代码]#

Turn an expression to Basic Block Normal Form. We define a block as a group of expressions implied by the scope structure. Each graph node can only belong to a single block. For any value that is being used in multiple blocks, it has to be referred by a Var which is defined in a block, whose scope is the least common ancestor of blocks this value is used.

Returns#

ret: tvm.transform.Pass

The registered pass that transforms an expression into Basic Block Normal Form.

tvm.relay.transform.ToCPS(expr, mod=None)[源代码]#

Turn expression into continuation passing style(CPS).

Every intermediate compute will be passed to a continuation.

Returns#

result: tvm.transform.Pass

The registered pass that transforms an expression into CPS.

tvm.relay.transform.ToGraphNormalForm()[源代码]#

Turn a Relay program in A Normal Form into Graph Normal Form

Returns#

rettvm.transform.Pass

The registered pass that transforms an expression into Graph Normal Form.

tvm.relay.transform.ToMixedPrecision(mixed_precision_type='float16', missing_op_mode=1)[源代码]#

Automatic mixed precision rewriter. Rewrite an FP32 relay graph into a version where as many operations as possible are in the target mixed_precision_type.

Parameters#

mixed_precision_type: str

The target datatype to transform operations in the graph to use.

missing_op_mode: int
Determines how to handle ops not registered with FTVMMixedPrecisionConversionType

0: Does not allow any missing ops. Will throw errors when encountering any. 1: Allow missing ops but emit warnings. 2: Allow missing ops and silently ignore them.

relay.ToMixedPrecision.keep_orig_output_dtype: boolean

Defines if outputs should be retained in original data type or convert to mixed_precision_type. By default this parameter is False and transformation modifies the data types of outputs to mixed_precision_type. This parameter is not part of explicit arguments of the transformation, but should be passed through tvm.transform.PassContext.

Returns#

rettvm.transform.Pass

The registered pass.

tvm.relay.transform.build_config(opt_level=2, required_pass=None, disabled_pass=None, trace=None)[源代码]#

Configure the build behavior by setting config variables. This function will be deprecated in TVM v0.7. Instead, we should directly use tvm.transform.PassContext.

Parameters#

opt_level: int, optional

Optimization level. The optimization pass name and level are as the following:

OPT_PASS_LEVEL = {
    "SimplifyInference": 0,
    "OpFusion": 1,
    "FoldConstant": 2,
    "FoldScaleAxis": 3,
    "AlterOpLayout": 3,
    "CanonicalizeOps": 3,
    "CanonicalizeCast": 3,
    "EliminateCommonSubexpr": 3,
    "CombineParallelConv2D": 4,
    "CombineParallelDense": 4,
    "CombineParallelBatchMatmul": 4,
    "FastMath": 4
}
required_pass: set of str, optional

Optimization passes that are required regardless of optimization level.

disabled_pass: set of str, optional

Optimization passes to be disabled during optimization.

trace: Callable[[IRModule, PassInfo, bool], None]

A tracing function for debugging or introspection.

Returns#

pass_context: PassContext

The pass context for optimizations.

tvm.relay.transform.function_pass(pass_func=None, opt_level=None, name=None, required=None)[源代码]#

Decorate a function pass.

This function returns a callback when pass_func is provided. Otherwise, it returns the created function pass using the given optimization function.

Parameters#

pass_funcOptional[Callable[(Function, Module, PassContext) -> Function]]

The transformation function or class.

opt_levelint

The optimization level of this module pass.

nameOptional[str]

The name of the function pass. The name could be empty. In this case, the name of the optimization function will be used as the pass name.

requiredOptional[List[str]]

The list of passes that the module pass is dependent on.

Returns#

create_function_pass : Union[Callable, FunctionPass]

A decorator will be returned if pass_func is not provided, otherwise return the decorated result. The returned decorator has two behaviors depending on the input: A new FunctionPass will be returned when we decorate a pass function. A new FunctionPass class will be returned when we decorate a class type.

Examples#

The following code block decorates a function pass class.

@relay.transform.function_pass(opt_level=1)
class TestReplaceFunc:
    def __init__(self, new_func):
        self.new_func = new_func

    def transform_function(self, func, mod, ctx):
        # just for demo purposes
        # transform func to new_func
        return self.new_func

x = relay.var("x", shape=(10, 20))
f1 = relay.Function([x], x)
f2 = relay.Function([x], relay.log(x))
# fpass is now a special pass that replaces every
# function to f1
fpass = TestReplaceFunc(f1)
# now every function in input_mod is replaced by f1
res_mod = fpass(input_mod)

The following code creates a function pass by decorating a user defined transform function.

@relay.transform.function_pass(opt_level=2)
def transform(func, mod, ctx):
    # my transformations here.
    return func

function_pass = transform
assert isinstance(function_pass, transform.FunctionPass)
assert function_pass.info.opt_level == 2

# Given a module m, the optimization could be invoked as the follwoing:
updated_mod = function_pass(m)
# Now constant folding should have been applied to every function in
# the provided module m. And the updated module will be returned.
tvm.relay.transform.gradient(expr, mod=None, mode='higher_order')[源代码]#

Transform the input function, returning a function that calculate the original result, paired with gradient of the input.

Parameters#

exprtvm.relay.Expr

The input expression, which is a Function or a GlobalVar.

mod : Optional[tvm.IRModule]

modeOptional[String]

The mode of the automatic differentiation algorithm. ‘first_order’ only works on first order code, but will not produce reference nor closure. ‘higher_order’ works on all code using reference and closure.

Returns#

exprtvm.relay.Expr

The transformed expression.

tvm.relay.transform.recast(expr, dtype, out_dtype, ops=None, skip_layers=None)[源代码]#

Convert the types of operations in a graph to a new value. Note that this is primarily useful for testing performance of individual operations at the new datatype. In a real setting, this pass will almost certainly do a poor job converting from one datatype to another as it just applies hard casting. For example, when recasting from float to integer, many small values will simply be set to 0. Although this will allow autotuning and benchmarking to produce proper timings at the new data type, the output of the model will of course be heavily impacted.

Parameters#

expr: tvm.relay.Expr, tvm.relay.Function, or tvm.ir.IRModule

The original function that will have its type changed.

dtype: str

The target type to cast to.

out_dtype: str

The output type to cast to.

ops: List[str]

A list of operations that should have their type changed, others will be left as is.

skip_layers: List[int]

A list of integers indicating operations that should not have their type changed, counted starting with the first valid operation encountered. Negative indices are allowed and indicate starting at the last layer.

Returns#

output_exprtvm.relay.Expr, tvm.relay.Function, or tvm.ir.IRModule

The graph after recasting to the specified datatype.

tvm.relay.transform.to_cps(func, mod=None)[源代码]#

Turn expression into CPS expression.

Every intermediate compute will be passed to a continuation.

Parameters#

func: tvm.relay.Function

The input function.

mod: Optional[tvm.IRModule]

The global module.

Returns#

result: tvm.relay.Function

The output function.

tvm.relay.transform.un_cps(func)[源代码]#

Turn an cps function into a Function without the continuation argument.

Note that this will not give the exact same interface as before cps:

If the input/output is higher order, they will still be in cps form.

Parameters#

func: tvm.relay.Function

The input function

Returns#

result: tvm.relay.Function

The output function