tilelang.carver.roller.policy.tensorcore¶

Policy for tensorcore schedule

Attributes¶

Classes¶

TensorCorePolicy

Module Contents¶

tilelang.carver.roller.policy.tensorcore.logger¶

class tilelang.carver.roller.policy.tensorcore.TensorCorePolicy¶

Bases: tilelang.carver.roller.policy.default.DefaultPolicy

wmma_k: int = 16¶

pipeline_stage: int = 1¶

use_async_copy: bool = False¶

block_reduction_depth: int | None = None¶

infer_node_smem_usage(td, node)¶

Parameters:

td (tilelang.carver.roller.hint.TileDict)
node (tilelang.carver.roller.node.PrimFuncNode)

get_node_reduce_step_candidates(node)¶

check_tile_shape_isvalid(td)¶

Parameters:: td (tilelang.carver.roller.hint.TileDict)

compute_node_stride_map(node, td)¶

Parameters:

node (tilelang.carver.roller.node.PrimFuncNode)
td (tilelang.carver.roller.hint.TileDict)

plan_rasterization(td)¶

Parameters:: td (tilelang.carver.roller.hint.TileDict)