Class ActiveSet

python/tvm/tirx/exec_context.py:125–184

Active thread set represented by a TileLayout.

@dataclass(frozen=True)
class ActiveSet:
    """Active thread set represented by a TileLayout."""

    layout: TileLayout

    @staticmethod
    def from_axes(axes: list[tuple[str, AxisRange]]) -> ActiveSet:
        shard = [Iter(axis_range.extent, axis_range.stride, name) for name, axis_range in axes]
        offset = {
            Axis.get(name): axis_range.offset for name, axis_range in axes if axis_range.offset != 0
        }
        return ActiveSet(TileLayout.from_iters(shard, [], offset))

    @property
    def size(self) -> int:
        result = 1
        for it in self.layout.shard:
            result *= int(it.extent)
        return result

    @property
    def axis_names(self) -> list[str]:
        return [str(it.axis.name) for it in self.layout.shard]

    def axis(self, name: str) -> AxisRange:
        for it in self.layout.shard:
            if str(it.axis.name) != name:
                continue
            offset = 0
            for axis, value in self.layout.offset.items():
                if str(axis.name) == name:
                    offset = int(value)
                    break
            return AxisRange(int(it.extent), offset, int(it.stride))
        raise ValueError(f"unknown active-set axis: {name!r}")

    def replace_axis(self, axis: str, axis_range: AxisRange) -> ActiveSet:
        axes: list[tuple[str, AxisRange]] = []
        found = False
        for name in self.axis_names:
            if name == axis:
                axes.append((name, axis_range))
                found = True
            else:
                axes.append((name, self.axis(name)))
        if not found:
            raise ValueError(f"unknown active-set axis: {axis!r}")
        return ActiveSet.from_axes(axes)

    @property
    def laneid(self) -> AxisRange:
        return self.axis("laneid")

    @property
    def warpid(self) -> AxisRange:
        return self.axis("warpid")

    @property