tilelang.carver.analysis¶

Analysis on TIR blocks, loops and functions.

Classes¶

IterInfo

Information about a loop/iter var.

BlockInfo

Information about a TIR block.

Functions¶

normalize_prim_func(sch)

Normalize the primfunc to normal form

find_var_from_func(func, var)

check_func_with_dynamic(func)

get_max_threads_per_block(target)

get_max_shared_memory_per_block(target)

get_root_block(sch[, func_name])

collect_block_iter_vars_used_in_access_region(block, ...)

Collect the block iter variables used in the access region of a buffer region.

collect_vars_used_in_prim_expr(expr)

Collect the variables used in the PrimExpr.

detect_dominant_read(block)

Detect the dominant read indices in the block.

is_broadcast_epilogue(sch, block, epilogue)

Check if the epilogue block is a broadcast pattern

get_reduction_blocks(sch, blocks)

get_coalesced_veclen(block_stmt[, target_bits])

Module Contents¶

class tilelang.carver.analysis.IterInfo(kind, var, dom, loop_rv)¶

Information about a loop/iter var.

Parameters:

  • kind (typing_extensions.Literal[S, R, O])

  • var (tvm.tir.Var)

  • dom (tvm.tir.PrimExpr)

  • loop_rv (tvm.tir.schedule.LoopRV)

kind: typing_extensions.Literal[S, R, O]¶
var: tvm.tir.Var¶
loop_rv: tvm.tir.schedule.LoopRV¶
property dom: int | tvm.tir.PrimExpr¶

The iteration domain of the loop.

Return type:

Union[int, tvm.tir.PrimExpr]

__str__()¶
Return type:

str

__repr__()¶
Return type:

str

class tilelang.carver.analysis.BlockInfo(name, iters, block_rv, reduction_block=False)¶

Information about a TIR block.

Parameters:

  • name (str)

  • iters (List[IterInfo])

  • block_rv (tvm.tir.schedule.BlockRV)

  • reduction_block (bool)

name: str¶
iters: List[IterInfo]¶
block_rv: tvm.tir.schedule.BlockRV¶
dom()¶

The iteration domain of the block.

Return type:

List[Union[int, tvm.tir.PrimExpr]]

dom_kind()¶

The iteration domain kind of the block, for example, SSSS, SSSR.

Return type:

str

is_injective()¶

Whether the block is injective, i.e. all its iteration domains are injective.

Return type:

bool

is_elementwise(sch)¶

Whether the block is elementwise, i.e. trivial mapping between read/write region

Parameters:

sch (tvm.tir.Schedule)

Return type:

bool

is_reduction()¶

Whether the block is a reduction workload.

Return type:

bool

abstract is_gemv()¶

Whether the block is a GEMV workload.

Return type:

bool

abstract is_gemm()¶

Whether the block is a GEMM workload.

Return type:

bool

__str__()¶
Return type:

str

__repr__()¶
Return type:

str

tilelang.carver.analysis.normalize_prim_func(sch)¶

Normalize the primfunc to normal form

Parameters:

sch (tvm.tir.Schedule)

Return type:

Optional[List[BlockInfo]]

tilelang.carver.analysis.find_var_from_func(func, var)¶
Parameters:

var (str)

tilelang.carver.analysis.check_func_with_dynamic(func)¶
tilelang.carver.analysis.get_max_threads_per_block(target)¶
Parameters:

target (tvm.target.target.Target)

Return type:

int

tilelang.carver.analysis.get_max_shared_memory_per_block(target)¶
Parameters:

target (tvm.target.target.Target)

Return type:

int

tilelang.carver.analysis.get_root_block(sch, func_name='main')¶
Parameters:

  • sch (tvm.tir.Schedule)

  • func_name (str)

Return type:

tvm.tir.schedule.BlockRV

tilelang.carver.analysis.collect_block_iter_vars_used_in_access_region(block, region)¶

Collect the block iter variables used in the access region of a buffer region.

Parameters:

  • block (tvm.tir.Block)

  • region (List[tvm.ir.Range])

Return type:

Set[tvm.tir.Var]

tilelang.carver.analysis.collect_vars_used_in_prim_expr(expr)¶

Collect the variables used in the PrimExpr.

Parameters:

expr (tvm.tir.PrimExpr)

Return type:

Set[tvm.tir.Var]

tilelang.carver.analysis.detect_dominant_read(block)¶

Detect the dominant read indices in the block.

Parameters:

block (tvm.tir.Block)

Return type:

tvm.tir.PrimExpr

tilelang.carver.analysis.is_broadcast_epilogue(sch, block, epilogue)¶

Check if the epilogue block is a broadcast pattern

Parameters:

  • sch (tvm.tir.Schedule)

  • block (tvm.tir.schedule.BlockRV)

  • epilogue (tvm.tir.schedule.BlockRV)

Return type:

bool

tilelang.carver.analysis.get_reduction_blocks(sch, blocks)¶
Parameters:

  • sch (tvm.tir.Schedule)

  • blocks (List[tvm.tir.schedule.BlockRV])

Return type:

List[tvm.tir.schedule.BlockRV]

tilelang.carver.analysis.get_coalesced_veclen(block_stmt, target_bits=128)¶
Parameters:

  • block_stmt (tvm.tir.Block)

  • target_bits (int)

Return type:

int