Class SBlockInfo

python/tvm/s_tir/dlight/analysis/common_analysis.py:171–284

Information about a TIR block.

class SBlockInfo:
    """Information about a TIR block."""

    name: str
    iters: list[IterInfo]
    block_rv: s_tir.schedule.SBlockRV
    _reduction_block: bool

    def __init__(
        self,
        name: str,
        iters: list[IterInfo],
        block_rv: s_tir.schedule.SBlockRV,
        reduction_block: bool = False,
    ):
        """Construct a SBlockInfo object."""
        self.name = name
        self.block_rv = block_rv
        self.iters = iters
        self._reduction_block = reduction_block

    def dom(self) -> list[int | tirx.PrimExpr]:
        """The iteration domain of the block."""
        return [i.dom for i in self.iters]

    def read_bufs(self, sch: s_tir.Schedule) -> list[BufferInfo]:
        block_stmt = sch.get(self.block_rv)
        lps = sch.get_loops(self.block_rv)
        return [BufferInfo(sch, self.block_rv, buf, lps) for buf in block_stmt.reads]

    def write_bufs(self, sch: s_tir.Schedule) -> list[BufferInfo]:
        block_stmt = sch.get(self.block_rv)
        lps = sch.get_loops(self.block_rv)
        return [BufferInfo(sch, self.block_rv, buf, lps) for buf in block_stmt.writes]

    def dom_kind(self) -> str:
        """The iteration domain kind of the block, for example, SSSS, SSSR."""
        return "".join(i.kind for i in self.iters)

    def is_injective(self) -> bool:
        """Whether the SBlock is injective, i.e. all its iteration domains are injective."""
        return all(k == "S" for k in self.dom_kind())

    def is_elementwise(self, sch: s_tir.Schedule) -> bool:
        """Whether the SBlock is elementwise, i.e. trivial mapping between read/write region"""

        def _check_unit_var_range(dom: ir.Range, var: tirx.Var) -> bool:
            return dom.min.same_as(var) and dom.extent == 1

        if not self.is_injective():
            return False
        block = sch.get(self.block_rv)
        if len(block.reads) != 1 or len(block.writes) != 1:
            return False
        r_region = block.reads[0].region
        w_region = block.writes[0].region
        if len(r_region) != len(w_region):
            return False