Function normalize_prim_func

python/tvm/s_tir/dlight/analysis/common_analysis.py:290–323

Normalize the primfunc to normal form

(sch: s_tir.Schedule)

def normalize_prim_func(sch: s_tir.Schedule) -> list[SBlockInfo] | None:
    """Normalize the primfunc to normal form"""
    try:
        result = _normalize_prim_func(sch)
        if result is None:
            return None
    except Exception:  # pylint: disable=broad-except
        return None

    def _iter_kind(i: tirx.IterVar) -> str:
        return {
            tirx.IterVar.DataPar: "S",
            tirx.IterVar.CommReduce: "R",
        }.get(i.iter_type, "O")

    blocks: list[SBlockInfo] = []
    for block, loops, iters, is_reduction in zip(*result):
        blocks.append(
            SBlockInfo(
                name=sch.get(block).name_hint,
                iters=[
                    IterInfo(
                        kind=_iter_kind(iter),  # type: ignore
                        var=iter.var,
                        dom=iter.dom.extent,
                        loop_rv=loop,
                    )
                    for loop, iter in zip(loops, iters)
                ],
                block_rv=block,
                reduction_block=is_reduction,
            )
        )
    return blocks


def get_sblock_info(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV) -> SBlockInfo: