tvm.tir.analysis

Wrapping existing analysis utils.

tvm.tir.analysis.analysis.expr_deep_equal(lhs: tvm.ir.expr.PrimExpr, rhs: tvm.ir.expr.PrimExpr) → bool

Deeply compare two nested expressions.

Parameters

• lhs (PrimExpr) – The left operand.

• rhs (PrimExpr) – The right operand.

Returns

result – The comparison result

Return type

bool

Note

This function does not remap variable bindings, it will not return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless x.same_as(y). Use py:func:tvm.ir.structural_equal to handle structural variable remapping.

Due to the restriction of not remapping variables, this function can run faster than StructuralEqual and can be used as a utility function during arithmetic simplifications.

Always consider py:func:tvm.ir.structural_equal first, which handles the structural remapping.

See also

tvm.ir.structural_equal

tvm.tir.analysis.analysis.verify_ssa(func: tvm.tir.function.PrimFunc) → bool

Verify if the func is in SSA form.

Parameters

func (tvm.tir.PrimFunc) – The module to be verified.

Returns

result – The result of verification.

Return type

bool

tvm.tir.analysis.analysis.verify_memory(func: tvm.tir.function.PrimFunc) → bool

Verify if func contains illegal host side direct memory access.

Parameters

func (tvm.tir.PrimFunc) – The module to be verified.

Returns

result – The result of verification.

Return type

bool

tvm.tir.analysis.analysis.verify_gpu_code(func: tvm.tir.function.PrimFunc, constraints: Dict[str, int]) → None

Verify if module contains illegal host side direct memory access.

Parameters

• func (tvm.tir.PrimFunc) – The module to be verified.

• constraints (Dict[str, int]) – The attribute constraints.

Returns

result – The result of verification.

Return type

bool

tvm.tir.analysis.analysis.get_block_access_region(block: tvm.tir.stmt.Block, buffer_var_map: Dict[tvm.tir.expr.Var, tvm.tir.buffer.Buffer]) → List[List[tvm.tir.stmt.BufferRegion]]

Detect which regions of tensors in this block are read or written to.

Regions are sorted by order of appearance in the AST.

Parameters

• block (tvm.tir.Block) – The block in which we are detecting read/write regions.

• buffer_var_map (Dict[tir.Var, Buffer]) – The outside buffers which may access the block. Mapping from buffer var to the buffer

Returns

result –

Array of access regions. There are three arrays of BufferRegion:

• first: read regions

• second: write regions

• third: opaque regions

Return type

List[List[BufferRegion]]

tvm.tir.analysis.analysis.get_block_read_write_region(block: tvm.tir.stmt.Block, buffer_var_map: Dict[tvm.tir.expr.Var, tvm.tir.buffer.Buffer]) → List[List[tvm.tir.stmt.BufferRegion]]

Auto detect the block read/write region according to its body stmt.

An opaque access will be counted as both a read and a write access

Parameters

• block (tvm.tir.Block) – The block in which we are detecting read/write regions.

• buffer_var_map (Dict[tir.Var, Buffer]) – The outside buffers which may access the block. Mapping from buffer var to the buffer

Returns

result – An array only consisting of the read regions and write regions of the input block

Return type

List[List[BufferRegion]]

tvm.tir.analysis.analysis.calculate_workspace_bytes(func: tvm.tir.function.PrimFunc, workspace_byte_alignment: int) → int

Calculate the workspace size in bytes needed by the TIR allocates inside the TIR PrimFunc.

Parameters

• func (tvm.tir.PrimFunc) – The function to be detected.

• workspace_byte_alignment (int) – The byte alignment required for each tensor

Returns

result – Workspace size in bytes.

Return type

int

tvm.tir.analysis.analysis.calculate_constant_bytes(func: tvm.tir.function.PrimFunc, constant_byte_alignment: int) → int

Calculate the constant size in bytes needed by the TIR allocates inside the TIR PrimFunc.

Parameters

• func (tvm.tir.PrimFunc) – The function to be detected.

• constant_byte_alignment (int) – The byte alignment required for each tensor

Returns

result – Workspace size in bytes.

Return type

int

tvm.tir.analysis.analysis.calculate_allocated_bytes(func_or_mod: Union[tvm.tir.function.PrimFunc, tvm.ir.module.IRModule]) → Union[Dict[str, int], Dict[str, Dict[str, int]]]

Calculate allocated memory per memory scope required by TIR PrimFuncs.

Parameters

func_or_mod (Union[PrimFunc, IRModule]) – The function or module to be detected. If a module is passed, allocated memory is calculated for all PrimFuncs inside the module

Returns

result – Allocated memory size per scope in bytes for each function in the IRModule returned as a dict with function names as keys and a dict of allocated sizes as values. If a single PrimFunc is passed, the function name is returned as “main”

Return type

Union[Dict[str, int], Dict[str, Dict[str, int]]]

tvm.tir.analysis.analysis.detect_buffer_access_lca(func: tvm.tir.function.PrimFunc) → Dict[tvm.tir.buffer.Buffer, tvm.tir.stmt.Stmt]

Detect the lowest common ancestor(LCA) of buffer access, including both high-level access (BufferLoad, BufferStore) and low-level access (BufferLoad, BufferStore and opaque access). The LCA may be a For loop or a Block.

Parameters

func (tvm.tir.PrimFunc) – The function to be detected.

Returns

result – Map from buffer to the LCA of all access to it.

Return type

Dict[Buffer, Stmt]

tvm.tir.analysis.analysis.estimate_tir_flops(stmt_or_mod: Union[tvm.tir.stmt.Stmt, tvm.ir.module.IRModule]) → float

Estimate the FLOPs of a TIR fragment.

Parameters

stmt_or_mod (Union[Stmt, IRModule]) – The TIR fragment or IRModule to be estimated.

Returns

flops – The estimated FLOPs.

Return type

float

tvm.tir.analysis.analysis.undefined_vars(node: Union[tvm.tir.stmt.Stmt, tvm.ir.expr.PrimExpr], defs: Optional[List[tvm.tir.expr.Var]] = None) → List[tvm.tir.expr.Var]

Find undefined vars in a TIR statement or expression.

Parameters

• node (Union[Stmt, PrimExpr]) – The TIR statement or expression to be checked.

• defs (Optional[List[tir.Var]]) – The vars that is defined

Returns

result – The undefined vars.

Return type

List[tir.Var]

tvm.tir.analysis.analysis.get_prim_func_arg_and_result_memory_constraints(func: tvm.tir.function.PrimFunc, relay_func_type: tvm.runtime.object.Object) → List[str]

Returns the memory (aka storage) scope constraints for all the arguments and result of func. However the result will be w.r.t. the func’s representation as a Relay Function of relay_func_type before lowering and conversion to DPS.

Visible for testing.

Parameters

• func (tvm.tir.PrimFunc) – The function to retrieve constraints from.

• relay_func_type (tvm.relay.FuncType) – The type of the Relay Function from which the func was derived.

Returns

result – Memory scope constraints for funcs args and result in Relay form. The empty string denotes ‘no constraint’.

Return type

List[AnyStr]

tvm.tir.analysis.analysis.apply_prim_func_arg_and_result_memory_constraints(func: tvm.tir.function.PrimFunc, relay_func_type: tvm.runtime.object.Object, arg_and_result_memory_scopes: List[str]) → tvm.tir.function.PrimFunc

Returns func written to capture the memory (aka storage) scope constraints for each of the func’s parameters given by arg_and_result_memory_scopes. However, arg_and_result_memory_scopes should be w.r.t. the func’s representation as a Relay Function of relay_func_type before lowering and conversion to DPS.

Visible for testing.

CAUTION: This is experimental. The resulting PrimFunc may not have fully accounted for all new memory scopes.

Parameters

• func (tvm.tir.PrimFunc) – The function to retrieve constraints from.

• relay_func_type (tvm.relay.FuncType) – The type of the Relay Function from which the func was derived.

• arg_and_result_memory_scopes (Array[AnyStr]) – Memory constraints for funcs args and result in Relay form. The empty string denotes ‘no constraint’.

Returns

result – The rewritten func.

Return type

tvm.tir.PrimFunc

tvm.tir.analysis.analysis.verify_well_formed(obj: Union[tvm.tir.function.PrimFunc, tvm.ir.module.IRModule], assert_mode: bool = True) → bool

Verify if the given TIR is well-formed. The verification includes:

• Check if expressions not contain vars that is defined outside the block.

Parameters

• obj (Union[tvm.tir.PrimFunc, tvm.ir.IRModule]) – The function or module to be verified.

• assert_mode (bool) – The indicator if it raises an error when the function is not well-formed.

Returns

result – Whether it is a well-formed TIR function.

Return type

bool

tvm.tir.analysis.analysis.OOBChecker()

Detect out of bounds memory access in arrays.

Returns

fpass – The result pass

Return type

tvm.transform.Pass

tvm.tir.analysis.analysis.find_anchor_block(mod: tvm.ir.module.IRModule) → tvm.tir.stmt.Block

Find the “anchor block” of the given module.

We define the anchor block to be the block with (1) an init statement and (2) having the biggest flops count. The latter condition is only used when there are multiple blocks with an init statement.

For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block. The input module may not contain more than one such block. For example, a module having two conv2d is not allowed as an input.

However, a module created from winograd convolution has multiple blocks with an init statement (input transform, batched GEMM, and output transform). We use the second condition, the flops count, to determine that the batched GEMM block is the anchor block.

Parameters

mod (tvm.ir.IRModule) – The input TIR module.

Returns

anchor_block – The anchor block if found, None otherwise.

Return type

Block

tvm.tir.analysis.analysis.get_vtcm_compaction_passes() → List[tvm.ir.transform.Pass]

Utility function to get the list of lowering passes to be applied to calculate the compacted VTCM allocation size

Returns

result – returns list of passes

Return type

List[tvm.transform.Pass]

tvm.tir.analysis.analysis.is_pure_function(func: tvm.tir.function.PrimFunc) → bool

Checks if the function is a pure function

tvm.tir.analysis.analysis.assert_pure_function(func: tvm.tir.function.PrimFunc) → bool

Asserts that the function is a pure function