tilelang.carver.roller.node module#

PrimFunc Wrapper and Block information Analaysis

class tilelang.carver.roller.node.BlockAnalyzer(sch)#

Bases: object

get_block_info(block: BlockRV) → BlockInfo#
get_block_name(block: BlockRV) → str#
get_buffers(block: BlockRV) → List[Buffer]#
get_consumer_blocks(block: BlockRV) → List[BlockRV]#
get_input_buffers(block: BlockRV) → List[Buffer]#
get_output_buffers(block: BlockRV) → List[Buffer]#
get_producer_blocks(block: BlockRV) → List[BlockRV]#
get_reduce_axis(block: BlockRV) → List[IterVar]#
get_spatial_axis(block: BlockRV) → List[IterVar]#
class tilelang.carver.roller.node.Edge(src_node: 'Node', dst_node: 'Node', src_id: int, dst_id: int)#

Bases: object

dst_id: int#
dst_node: Node#
src_id: int#
src_node: Node#
class tilelang.carver.roller.node.Node(tags: Optional[Dict] = None, name: str = 'Node')#

Bases: object

add_tag(k: str, v: Any = True) → None#
get_dtype(id=0) → DataType#
get_ir() → str#
get_shape(id: int = 0) → List[int]#
get_tag(k: str) → Any#
property inputs: List[Edge]#
is_output()#
is_placeholder()#
num_outputs() → int#
property outputs: List[Edge]#
set_dtype(dtype: DataType, id=0) → None#
set_inputs(i: int, edge: Edge)#
set_outputs(i: int, edge: Edge)#
set_shape(shape: List[int], id=0, overwrite=False) → None#
set_tag(k: str, v: Any = True) → None#
update_tags(tags: Dict) → None#
class tilelang.carver.roller.node.OutputNode(node, id=0)#

Bases: Node

get_ir() → str#
is_output()#
class tilelang.carver.roller.node.PlaceHolderNode(name='')#

Bases: Node

get_ir() → str#
is_placeholder()#
class tilelang.carver.roller.node.PrimFuncNode(prim_func: PrimFunc, tags: Optional[Dict] = None, name: str = 'PrimFuncNode')#

Bases: Node

extent_wrapper(value) → int#
footprint(shape, rstep, stride_map: Optional[Dict] = None) → int#
get_buffer_dtype(buffer: Buffer) → DataType#
get_input_buffers() → List[Buffer]#
get_opt_shape(name) → int#
get_reduce_inputs_dtype()#
get_space_dim() → List[int]#
infer_tensorcore_axis() → Tuple[int]#
propagate(tile, rstep: Optional[Dict] = None, targets=None)#
propagate_inputs(tile, rstep: Optional[Dict] = None) → List[List[int]]#
propagate_inputs_on_reduction(tile, rstep: Optional[Dict] = None) → List[List[int]]#
propagate_outputs(tile, rstep: Optional[Dict] = None) → List[List[int]]#
propagate_reduction_inputs(shape, rstep: Optional[Dict] = None) → Dict[str, List[int]]#
set_dtype(dtype: DataType, id=0) → None#
tilelang.carver.roller.node.find_topo_sort(output_node_list) → List[Node]#
tilelang.carver.roller.node.find_topo_sort_priority(output_node_list) → List[Node]#
tilelang.carver.roller.node.pre_order_traverse(block_analyzer, blocks, func)#
tilelang.carver.roller.node.topo_order(list_of_nodes) → List[Node]#