tvm.tir.analysis

tvm.tir.analysis#

Wrapping existing analysis utils.

tvm.tir.analysis.analysis.OOBChecker()[源代码]#

Detect out of bounds memory access in arrays.

Returns#

fpasstvm.transform.Pass

The result pass

tvm.tir.analysis.analysis.apply_prim_func_arg_and_result_memory_constraints(func, relay_func_type, arg_and_result_memory_scopes)[源代码]#

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: tvm.tir.PrimFunc

The rewritten func.

参数:

  • func (PrimFunc)

  • relay_func_type (Object)

  • arg_and_result_memory_scopes (List[str])

返回类型:

PrimFunc

tvm.tir.analysis.analysis.assert_pure_function(func)[源代码]#

Asserts that the function is a pure function

参数:

func (PrimFunc)

返回类型:

bool

tvm.tir.analysis.analysis.calculate_allocated_bytes(func_or_mod)[源代码]#

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#

resultUnion[Dict[str, int], Dict[str, Dict[str, int]]]

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"

参数:

func_or_mod (PrimFunc | IRModule)

返回类型:

Dict[str, int] | Dict[str, Dict[str, int]]

tvm.tir.analysis.analysis.calculate_constant_bytes(func, constant_byte_alignment)[源代码]#

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_alignmentint

The byte alignment required for each tensor

Returns#

resultint

Workspace size in bytes.

参数:
返回类型:

int

tvm.tir.analysis.analysis.calculate_workspace_bytes(func, workspace_byte_alignment)[源代码]#

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_alignmentint

The byte alignment required for each tensor

Returns#

resultint

Workspace size in bytes.

参数:
返回类型:

int

tvm.tir.analysis.analysis.detect_buffer_access_lca(func)[源代码]#

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#

resultDict[Buffer, Stmt]

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

参数:

func (PrimFunc)

返回类型:

Dict[Buffer, Stmt]

tvm.tir.analysis.analysis.estimate_tir_flops(stmt_or_mod)[源代码]#

Estimate the FLOPs of a TIR fragment.

Parameters#

stmt_or_mod: Union[Stmt, IRModule]

The TIR fragment or IRModule to be estimated.

Returns#

flops: float

The estimated FLOPs.

参数:

stmt_or_mod (Stmt | IRModule)

返回类型:

float

tvm.tir.analysis.analysis.expr_deep_equal(lhs, rhs)[源代码]#

Deeply compare two nested expressions.

Parameters#

lhsPrimExpr

The left operand.

rhsPrimExpr

The right operand.

Returns#

resultbool

The comparison result

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

参数:

  • lhs (PrimExpr)

  • rhs (PrimExpr)

返回类型:

bool

tvm.tir.analysis.analysis.find_anchor_block(mod)[源代码]#

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: Block

The anchor block if found, None otherwise.

参数:

mod (IRModule)

返回类型:

Block

tvm.tir.analysis.analysis.get_block_access_region(block, buffer_var_map)[源代码]#
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_mapDict[Var, Buffer]

The outside buffers which may access the block. Mapping from buffer var to the buffer

Returns#

resultList[List[BufferRegion]]
Array of access regions. There are three arrays of BufferRegion:

  • first: read regions

  • second: write regions

  • third: opaque regions

参数:
返回类型:

List[List[BufferRegion]]

tvm.tir.analysis.analysis.get_block_read_write_region(block, buffer_var_map)[源代码]#
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_mapDict[Var, Buffer]

The outside buffers which may access the block. Mapping from buffer var to the buffer

Returns#

resultList[List[BufferRegion]]

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

参数:
返回类型:

List[List[BufferRegion]]

tvm.tir.analysis.analysis.get_prim_func_arg_and_result_memory_constraints(func, relay_func_type)[源代码]#

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: List[AnyStr]

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

参数:

  • func (PrimFunc)

  • relay_func_type (Object)

返回类型:

List[str]

tvm.tir.analysis.analysis.get_vtcm_compaction_passes()[源代码]#

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

Returns#

resultList[tvm.transform.Pass]

returns list of passes

返回类型:

List[Pass]

tvm.tir.analysis.analysis.is_pure_function(func)[源代码]#

Checks if the function is a pure function

参数:

func (PrimFunc)

返回类型:

bool

tvm.tir.analysis.analysis.undefined_vars(node, defs=None)[源代码]#

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[Var]]

The vars that is defined

Returns#

resultList[Var]

The undefined vars.

参数:
返回类型:

List[Var]

tvm.tir.analysis.analysis.verify_gpu_code(func, constraints)[源代码]#

Verify if module contains illegal host side direct memory access.

Parameters#

func: tvm.tir.PrimFunc

The module to be verified.

constraintsDict[str, int]

The attribute constraints.

Returns#

resultbool

The result of verification.

参数:
返回类型:

None

tvm.tir.analysis.analysis.verify_memory(func)[源代码]#

Verify if func contains illegal host side direct memory access.

Parameters#

func: tvm.tir.PrimFunc

The module to be verified.

Returns#

resultbool

The result of verification.

参数:

func (PrimFunc)

返回类型:

bool

tvm.tir.analysis.analysis.verify_ssa(func)[源代码]#

Verify if the func is in SSA form.

Parameters#

func: tvm.tir.PrimFunc

The module to be verified.

Returns#

resultbool

The result of verification.

参数:

func (PrimFunc)

返回类型:

bool

tvm.tir.analysis.analysis.verify_well_formed(obj, assert_mode=True)[源代码]#
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: bool

Whether it is a well-formed TIR function.

参数:
返回类型:

bool