tvm.sparse.FormatRewriteRule

class tvm.sparse.FormatRewriteRule(name: str, new_format_desc: tvm.tir.function.PrimFunc, buffers_to_rewrite: List[str], axes_before_rewrite: List[str], axes_after_rewrite: List[str], axis_map: Dict[str, List[str]], idx_map: Union[Callable, tvm.tir.function.IndexMap], inv_idx_map: Union[Callable, tvm.tir.function.IndexMap])

Format rewriting rule.

namestr

Name of the format rewriting rule.

format_descPrimFunc

A TIR script describing the new format.

buffers_to_rewrite: List[str]

The list of sparse buffers we need to rewrite.

axes_before_rewriteList[str]

The list of axes before the rewrite.

axes_after_rewriteList[str]

The list of axes after the rewrite.

axis_mapDict[str, List[str]]

The axis mapping from the old format to the new format.

idx_map_funcUnion[Callable, IndexMap]

A function describing the index mapping from the old format to indices in new format.

inv_idx_map_funcUnion[Callable, IndexMap]

A function describing the coordinate mapping from indices in new format. to indices in old format.

__init__(name: str, new_format_desc: tvm.tir.function.PrimFunc, buffers_to_rewrite: List[str], axes_before_rewrite: List[str], axes_after_rewrite: List[str], axis_map: Dict[str, List[str]], idx_map: Union[Callable, tvm.tir.function.IndexMap], inv_idx_map: Union[Callable, tvm.tir.function.IndexMap]) → None

Methods

__init__(name, new_format_desc, ...)
same_as	Check object identity.

Attributes

handle