Class AxisRange

python/tvm/tirx/exec_context.py:79–121

An active slice offset + stride * [0, extent) on one TileLayout axis.

@dataclass(frozen=True)
class AxisRange:
    """An active slice offset + stride * [0, extent) on one TileLayout axis."""

    extent: int
    offset: int = 0
    stride: int = 1

    def intersect(self, lo: int, hi: int) -> AxisRange:
        i_lo = max(0, _ceildiv(lo - self.offset, self.stride))
        i_hi = min(self.extent, (hi - 1 - self.offset) // self.stride + 1)
        if i_hi <= i_lo:
            raise ExecContextError(
                f"filter produces empty range: current=[{self.offset},"
                f" {self.offset + self.extent}) ∩ [{lo}, {hi})"
            )
        return AxisRange(
            extent=i_hi - i_lo, offset=self.offset + self.stride * i_lo, stride=self.stride
        )

    def modulo(self, modulus: int, residue: int) -> AxisRange:
        residue %= modulus
        rhs = (residue - self.offset) % modulus
        g = _gcd(self.stride, modulus)
        if rhs % g != 0:
            raise ExecContextError(
                f"modulo filter produces empty range: {self.offset} + {self.stride} * i"
                f" == {residue} mod {modulus}"
            )
        reduced_stride = self.stride // g
        reduced_rhs = rhs // g
        reduced_modulus = modulus // g
        period = reduced_modulus
        i0 = (reduced_rhs * _mod_inverse(reduced_stride, reduced_modulus)) % reduced_modulus
        if i0 >= self.extent:
            raise ExecContextError(
                f"modulo filter produces empty range: {self.offset} + {self.stride} * i"
                f" == {residue} mod {modulus}"
            )
        return AxisRange(
            extent=(self.extent - 1 - i0) // period + 1,
            offset=self.offset + self.stride * i0,
            stride=self.stride * period,
        )


@dataclass(frozen=True)