tilelang.carver.roller.node

PrimFunc Wrapper and Block information Analaysis

Classes

BlockAnalyzer
Edge
Node
PlaceHolderNode
PrimFuncNode
OutputNode

Functions

pre_order_traverse(block_analyzer, blocks, func)
topo_order(list_of_nodes)
find_topo_sort_priority(output_node_list)
find_topo_sort(output_node_list)

Module Contents

tilelang.carver.roller.node.pre_order_traverse(block_analyzer, blocks, func)

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

Bases: object

sch: tvm.tir.Schedule

block_infos: List[tilelang.carver.analysis.BlockInfo] = None

get_block_name(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

str

get_block_info(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

tilelang.carver.analysis.BlockInfo

get_spatial_axis(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.IterVar]

get_reduce_axis(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.IterVar]

get_input_buffers(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.Buffer]

get_output_buffers(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.Buffer]

get_buffers(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.Buffer]

get_producer_blocks(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.schedule.schedule.BlockRV]

get_consumer_blocks(block)

Parameters:

block (tvm.tir.schedule.schedule.BlockRV)

Return type:

List[tvm.tir.schedule.schedule.BlockRV]

class tilelang.carver.roller.node.Edge

src_node: Node

dst_node: Node

src_id: int

dst_id: int

class tilelang.carver.roller.node.Node(tags=None, name='Node')

Bases: object

Parameters:

• tags (Optional[Dict])

• name (str)

name = 'Node'

update_tags(tags)

Parameters:

tags (Dict)

Return type:

None

set_tag(k, v=True)

Parameters:

• k (str)

• v (Any)

Return type:

None

add_tag(k, v=True)

Parameters:

• k (str)

• v (Any)

Return type:

None

get_tag(k)

Parameters:

k (str)

Return type:

Any

is_placeholder()

is_output()

property inputs: List[Edge]

Return type:

List[Edge]

property outputs: List[Edge]

Return type:

List[Edge]

set_inputs(i, edge)

Parameters:

• i (int)

• edge (Edge)

set_outputs(i, edge)

Parameters:

• i (int)

• edge (Edge)

get_dtype(id=0)

Return type:

tvm.DataType

set_dtype(dtype, id=0)

Parameters:

dtype (tvm.DataType)

Return type:

None

get_shape(id=0)

Parameters:

id (int)

Return type:

List[int]

set_shape(shape, id=0, overwrite=False)

Parameters:

shape (List[int])

Return type:

None

num_outputs()

Return type:

int

abstract get_ir()

Return type:

str

__repr__()

Return type:

str

class tilelang.carver.roller.node.PlaceHolderNode(name='')

Bases: Node

is_placeholder()

get_ir()

Return type:

str

class tilelang.carver.roller.node.PrimFuncNode(prim_func, tags=None, name='PrimFuncNode')

Bases: Node

Parameters:

• prim_func (tvm.tir.PrimFunc)

• tags (Optional[Dict])

• name (str)

prim_func

sch: tvm.tir.Schedule

block_analyzer: BlockAnalyzer

schedule_stages: List[tvm.tir.schedule.schedule.BlockRV] = []

blocks: List[tvm.tir.schedule.schedule.BlockRV] = []

output_blocks: List[tvm.tir.schedule.schedule.BlockRV] = None

reduction_block: tvm.tir.schedule.schedule.BlockRV = None

raxis = []

input_buffers = []

output_buffers = []

buffers = []

args = []

ana

get_opt_shape(name)

Return type:

int

extent_wrapper(value)

Return type:

int

get_space_dim()

Return type:

List[int]

set_dtype(dtype, id=0)

Parameters:

dtype (tvm.DataType)

Return type:

None

get_buffer_dtype(buffer)

Parameters:

buffer (tvm.tir.Buffer)

Return type:

tvm.DataType

propagate(tile, rstep=None, targets=None)

Parameters:

rstep (Optional[Dict])

propagate_inputs(tile, rstep=None)

Parameters:

rstep (Optional[Dict])

Return type:

List[List[int]]

propagate_inputs_on_reduction(tile, rstep=None)

Parameters:

rstep (Optional[Dict])

Return type:

List[List[int]]

propagate_outputs(tile, rstep=None)

Parameters:

rstep (Optional[Dict])

Return type:

List[List[int]]

propagate_reduction_inputs(shape, rstep=None)

Parameters:

rstep (Optional[Dict])

Return type:

Dict[str, List[int]]

get_reduce_inputs_dtype()

infer_tensorcore_axis()

Return type:

Tuple[int]

footprint(shape, rstep, stride_map=None)

Parameters:

stride_map (Optional[Dict])

Return type:

int

get_input_buffers()

Return type:

List[tvm.tir.Buffer]

class tilelang.carver.roller.node.OutputNode(node, id=0)

Bases: Node

is_output()

get_ir()

Return type:

str

tilelang.carver.roller.node.topo_order(list_of_nodes)

Return type:

List[Node]

tilelang.carver.roller.node.find_topo_sort_priority(output_node_list)

Return type:

List[Node]

tilelang.carver.roller.node.find_topo_sort(output_node_list)

Return type:

List[Node]