tvm.relay.analysis

The Relay IR namespace containing the analysis passes.

Classes:

AnnotatedRegionSet(expr, region_begin_op, ...)

Class to represent a relay expression split into regions.

CallGraph(module)

Class to represent a call graph.

Feature(value)

The features a program might contain.

Functions:

count_layers(expr, valid_ops)

Determine the number of layers of specified ops in a graph.

all_dtypes(expr)

Collect set of all data types used in expr.

all_type_vars(expr[, mod])

Get all type variables from expression/type e

all_vars(expr)

Get all vars from expression expr in post-DFS order.

bound_type_vars(expr[, mod])

Get bound type variables from expression/type e

bound_vars(expr)

Get bound vars from expression expr in post-DFS order.

check_basic_block_normal_form(expr)

Check whether an expression is in the basic block form

check_constant(expr)

Check whether an expression is constant

check_kind(t[, mod])

Check that the type is well kinded and return the kind.

cpu([dev_id])

Construct a CPU device

detect_feature(a[, b])

Detect the feature used in a relay program.

extract_fused_functions(mod)

Pass to extract IRModule of only fused primitive functions.

extract_intermdeiate_expr(mod, expr_id)

Extract Relay Expr by its expression ID

free_type_vars(expr[, mod])

Get free type variables from expression/type e

free_vars(expr)

Get free Vars from expression expr in Post DFS order.

get_calibration_data(mod, data)

Get the calibration data of a given relay graph

get_total_mac_number(expr)

Count the number of MACs (multiply-accumulate) of a model

list_fake_quantized_op_freqs(mod)

Pass to extract fake quantized op names and the frequency that they appear in fake quantized regions of an IRModule.

list_op_freqs(mod)

Pass to extract unique operator names and how frequently they appear in an IRModule.

post_order_visit(expr, fvisit)

Recursively visit the ir in post DFS order node, apply fvisit.

search_fc_transpose(expr)

Search fc weight name in the patten: y = nn.dense(x, transpose(w, [1, 0]))

unmatched_cases(match[, mod])

Finds cases that the match expression does not catch, if any.

well_formed(expr)

Check that each Var is only bound once (well formed).

class tvm.relay.analysis.AnnotatedRegionSet(expr, region_begin_op, region_end_op)

Class to represent a relay expression split into regions.

Methods:

get_region(expr)

Get the region an expression belongs to.

get_region(expr)

Get the region an expression belongs to.

Parameters:

expr (tvm.relay.Expr) – The expression.

Returns:

The region containing the expression. None if not found.

Return type:

region

class tvm.relay.analysis.CallGraph(module)

Class to represent a call graph.

Attributes:

module

Return the contained Relay IR module.

Methods:

ref_count(var)

Return the number of references to the global var

global_call_count(var)

Return the number of global function calls from a given global var.

is_recursive(var)

Return if the function corresponding to a var is a recursive function.

print_var(var)

Print a call graph of a global function by name or by variable.

property module

Return the contained Relay IR module.

Parameters:

None

Returns:

ret – The contained IRModule

Return type:

tvm.ir.IRModule

ref_count(var)

Return the number of references to the global var

Parameters:

var (Union[String, tvm.relay.GlobalVar])

Returns:

ret – The number reference to the global var

Return type:

int

global_call_count(var)

Return the number of global function calls from a given global var.

Parameters:

var (Union[String, tvm.relay.GlobalVar])

Returns:

ret – The number of global function calls from the given var.

Return type:

int

is_recursive(var)

Return if the function corresponding to a var is a recursive function.

Parameters:

var (Union[String, tvm.relay.GlobalVar])

Returns:

ret – If the function corresponding to var is recurisve.

Return type:

Boolean

print_var(var)

Print a call graph of a global function by name or by variable.

Parameters:

var (Union[String, tvm.relay.GlobalVar]) – The name or global variable.

Returns:

ret – The call graph represented in string.

Return type:

String

tvm.relay.analysis.count_layers(expr, valid_ops)

Determine the number of layers of specified ops in a graph. This pass computes only the deepest chain of ops rather than the total number of ops in a graph. Thus, if there are two parallel convolutions (for example), they would be considered a single layer.

Parameters:
  • expr (tvm.relay.Expr, tvm.relay.Function, or tvm.ir.IRModule.) – The input expression.

  • valid_ops (List[str]) – A list of the operations that should be included in the count.

Returns:

layer_count – The number of layers of the specified operations found in the graph.

Return type:

int

class tvm.relay.analysis.Feature(value)

The features a program might contain.

Attributes:

fMatch

Whether any non-atom fragment of the program is shared, making the program a graph.

fGraph

Whether there is local fixpoint in the program.

fMatch = 14

Whether any non-atom fragment of the program is shared, making the program a graph.

fGraph = 15

Whether there is local fixpoint in the program.

tvm.relay.analysis.all_dtypes(expr)

Collect set of all data types used in expr.

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

ret – Set of data types used in the expression (e.g., {‘int8’, ‘int32’})

Return type:

Set[String]

tvm.relay.analysis.all_type_vars(expr, mod=None)

Get all type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type

  • mod (Optional[tvm.IRModule]) – The global module

Returns:

free – The list of all type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.analysis.all_vars(expr)

Get all vars from expression expr in post-DFS order.

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

free – The list of all variables in post-DFS order.

Return type:

List[tvm.relay.Var]

tvm.relay.analysis.bound_type_vars(expr, mod=None)

Get bound type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type

  • mod (Optional[tvm.IRModule]) – The global module

Returns:

free – The list of bound type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.analysis.bound_vars(expr)

Get bound vars from expression expr in post-DFS order.

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

free – The list of bound variables in post-DFS order.

Return type:

List[tvm.relay.Var]

tvm.relay.analysis.check_basic_block_normal_form(expr)

Check whether an expression is in the basic block form

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

result – Whether the expression is in the basic block form.

Return type:

bool

tvm.relay.analysis.check_constant(expr)

Check whether an expression is constant

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

result – Whether the expression is constant.

Return type:

bool

tvm.relay.analysis.check_kind(t, mod=None)

Check that the type is well kinded and return the kind. For example, this mean type cannot has tensor of tensor, or is a tuple type of 2 shapes.

Parameters:
  • t (tvm.relay.Type) – The type to check

  • mod (Optional[tvm.IRModule]) – The global module.

Returns:

kind – the kind of t

Return type:

Kind

Examples

assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Shape)])) == Shape
assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Type)])) == Type
tvm.relay.analysis.cpu(dev_id=0)

Construct a CPU device

Parameters:

dev_id (int, optional) – The integer device id

Returns:

dev – The created device

Return type:

Device

tvm.relay.analysis.detect_feature(a, b=None)

Detect the feature used in a relay program.

Parameters:
  • a (Union[tvm.relay.Expr, tvm.IRModule]) – The input expression or module.

  • b (Optional[Union[tvm.relay.Expr, tvm.IRModule]]) – The input expression or module. The two arguments cannot both be expression or module.

Returns:

features – Features used in the program.

Return type:

Set[Feature]

tvm.relay.analysis.extract_fused_functions(mod)

Pass to extract IRModule of only fused primitive functions.

The ExtractFusedFunctions pass invokes SimplifyInference, FuseOps(3), and ExtractFusedFunctions in that order

Parameters:

mod (tvm.IRModule)

Returns:

ret – A module containing only fused primitive functions

Return type:

Dict[int, tvm.relay.function.Function]

tvm.relay.analysis.extract_intermdeiate_expr(mod, expr_id)

Extract Relay Expr by its expression ID

This function is used for extracting Relay Expr by its expression ID of the main function that we can see in print(mod[“main”]).

Parameters:
  • mod (tvm.IRModule)

  • expr_id (the Expr ID that we want to extract)

Returns:

ret

Return type:

Extracted IRModule

Examples

# Suppose our module is printed like this:
# def @main(%x: Tensor[(1, 1, 5, 1), float32], %w1, %w2) {
#   %0 = nn.conv2d(%x, %w1, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3]);
#   %1 = nn.conv2d(%0, %w2, padding=[1, 1, 1, 1], channels=1, kernel_size=[3, 3]);
#   %2 = add(%0, %1);
#   %3 = split(%2, indices_or_sections=1);
#   %4 = %3.0;
#   add(%4, 1f)
# }
# if we want to extract `%1 = nn.conv2d`
from tvm import relay

relay.analysis.extract_intermdeiate_expr(mod, 1)
tvm.relay.analysis.free_type_vars(expr, mod=None)

Get free type variables from expression/type e

Parameters:
  • expr (Union[tvm.relay.Expr,tvm.relay.Type]) – The input expression/type

  • mod (Optional[tvm.IRModule]) – The global module

Returns:

free – The list of free type variables in post-DFS order

Return type:

List[tvm.relay.TypeVar]

tvm.relay.analysis.free_vars(expr)

Get free Vars from expression expr in Post DFS order.

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

free – The list of free variables in post DFS order.

Return type:

List[tvm.relay.Var]

Note

The fact that Vars are post-DFS ordred are useful in neural networks: usually this means weights of previous are ordered first.

tvm.relay.analysis.get_calibration_data(mod, data)

Get the calibration data of a given relay graph

This pass uses the graph executor to get the calibration data of a module, which includes the input and output values of each function. The returned data uses the GlobalVar of each function as a key. Users can further access the inputs and outputs by using inputs or outputs as the key.

Following are some limitations: 1. The input module (graph) cannot have control flows. 2. The input arguments of each function cannot be tuples (outputs can be tuples). 3. We only handle top-level functions (i.e., nested function is not handled). 4. We only handle functions with Compiler attribute being set.

Parameters:
  • mod (tvm.IRModule) – The input module for collecting the calibration data

  • data (Dict[str, NDArray]) – The input data for running the module

Returns:

data

Return type:

Dict[tvm.relay.GlobalVar, Dict[str, NDArray]]

tvm.relay.analysis.get_total_mac_number(expr)

Count the number of MACs (multiply-accumulate) of a model

Parameters:

expr (tvm.relay.Expr) – The input expression.

Returns:

result – The number of MACs (multiply-accumulate) of a model

Return type:

int64

tvm.relay.analysis.list_fake_quantized_op_freqs(mod)

Pass to extract fake quantized op names and the frequency that they appear in fake quantized regions of an IRModule.

Parameters:

mod (tvm.IRModule)

Returns:

ret – Dict of fake quantized operator names to frequency

Return type:

Dict[str, int]

tvm.relay.analysis.list_op_freqs(mod)

Pass to extract unique operator names and how frequently they appear in an IRModule. Fused functions are traversed to count the operators that compose them.

Parameters:

mod (tvm.IRModule)

Returns:

ret – Dict of unique operator names to frequency

Return type:

Dict[str, int]

tvm.relay.analysis.post_order_visit(expr, fvisit)

Recursively visit the ir in post DFS order node, apply fvisit. Each node is guaranteed to be visited only once.

Parameters:
  • expr (tvm.relay.Expr) – The input expression.

  • fvisit (function) – The visitor function to be applied.

tvm.relay.analysis.search_fc_transpose(expr)

Search fc weight name in the patten: y = nn.dense(x, transpose(w, [1, 0]))

This function is used in the data_dep_optimization.simplify_fc_transpose method

Parameters:

expr (tvm.relay.Expr)

Returns:

ret – Array of weight variable name in pattern y = nn.dense(x, transpose(w, [1, 0]))

Return type:

Array[String]

tvm.relay.analysis.unmatched_cases(match, mod=None)

Finds cases that the match expression does not catch, if any.

Parameters:
  • match (tvm.relay.Match) – The match expression

  • mod (Optional[tvm.IRModule]) – The module (defaults to an empty module)

Returns:

missing_patterns – Patterns that the match expression does not catch.

Return type:

[tvm.relay.Pattern]

tvm.relay.analysis.well_formed(expr)

Check that each Var is only bound once (well formed).

Parameters:

expr (tvm.relay.Expr) – The input expression

Returns:

well_form – Whether the input expression is well formed

Return type:

bool