tvm.tir.analysis

Wrapping existing analysis utils.

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.apply_prim_func_arg_and_result_memory_constraints(func: PrimFunc, relay_func_type: Object, arg_and_result_memory_scopes: List[str]) → 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.calculate_allocated_bytes(func_or_mod: PrimFunc | IRModule) → 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.calculate_constant_bytes(func: 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_workspace_bytes(func: 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.detect_buffer_access_lca(func: PrimFunc) → Dict[Buffer, 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: Stmt | 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.expr_deep_equal(lhs: PrimExpr, rhs: PrimExpr) → bool

Deeply compare two nested expressions.

Parameters:
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.find_anchor_block(mod: IRModule) → 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_block_access_region(block: Block, buffer_var_map: Dict[Var, Buffer]) → List[List[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: Block, buffer_var_map: Dict[Var, Buffer]) → List[List[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.get_prim_func_arg_and_result_memory_constraints(func: PrimFunc, relay_func_type: 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.get_vtcm_compaction_passes() → List[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.undefined_vars(node: Stmt | PrimExpr, defs: List[Var] | None = None) → List[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.verify_gpu_code(func: 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.verify_memory(func: 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_ssa(func: 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_well_formed(obj: PrimFunc | 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