Tiling

@dataclass(frozen=True)
class OperandTileInfo:
    """
    This records how one operand should be sliced when we enter a tile.
    - `source_type` keeps the original type.
    - `loop_dims` the loop dimension each indexing-map result reads, where it
      reads exactly one, and `None` where it reads an expression over several or
      none, which no single loop range can be read back from.
    """

    source_type: MemRefType[Attribute] | TensorType[Attribute]
    loop_dims: tuple[int | None, ...]

    @staticmethod
    def analyze(
        indexing_map: AffineMap,
        source_type: MemRefType[Attribute] | TensorType[Attribute],
    ) -> "OperandTileInfo":
        """
        Analyze how one operand should be sliced for each tile.
        """

        loop_dims = tuple(
            expr.position if isinstance(expr, AffineDimExpr) else None
            for expr in indexing_map.results
        )
        return OperandTileInfo(source_type, loop_dims)

@dataclass(frozen=True)
class TilingPlan:
    """
    This stores the information needed to turn one op into tiled loop and tiled subview.
    - `loop_ranges` are original static loop ranges.
    - `tiled_dims` the dimensions that really get tiled.
    - `partial_tiled_dims` the tiled dimensions whose loop range is not divisible
      by the tile size, so that their last tile is smaller than the rest.
    - `operand_infos` stores one `OperandTileInfo` per operand.
    - `tile_sizes` are the normalized tile sizes, padded to match the op loop
      count. A tile size that is not known until the op runs is a value.
    """

    loop_ranges: tuple[int, ...]
    tiled_dims: tuple[int, ...]
    partial_tiled_dims: frozenset[int]
    operand_infos: tuple[OperandTileInfo, ...]
    tile_sizes: tuple[SSAValue | int, ...]

    @staticmethod
    def analyze(
        op: linalg.abstract_ops.LinalgStructuredOperation,
        tile_sizes: Sequence[SSAValue | int],
    ) -> "TilingPlan":
        """
        Analyze one supported structured linalg op and return a `TilingPlan`.
        """

        num_loops = op.get_num_loops()
        normalized_tile_sizes: tuple[SSAValue | int, ...] = tuple(
            tile_sizes[:num_loops]
        ) + (0,) * (num_loops - len(tile_sizes))

        # Tiling by zero means leaving a dimension alone, so a dimension is tiled
        # unless its tile size is known to be zero. One that is not known until
        # the op runs cannot be, so it is tiled.
        tiled_dims = tuple(
            dim
            for dim, tile_size in enumerate(normalized_tile_sizes)
            if not _is_zero(tile_size)
        )

        if not tiled_dims:
            return TilingPlan(
                loop_ranges=(),
                tiled_dims=(),
                partial_tiled_dims=frozenset(),
                operand_infos=(),
                tile_sizes=normalized_tile_sizes,
            )

        loop_ranges = _verify_is_tileable(
            op,
            normalized_tile_sizes,
            tiled_dims,
        )

        # A range that is not known until the op runs cannot be shown to divide
        # by its tile size, so it is treated as leaving a leftover tile.
        partial_tiled_dims = frozenset(
            dim
            for dim in tiled_dims
            if not _divides(normalized_tile_sizes[dim], loop_ranges[dim])
        )

        operand_infos_list: list[OperandTileInfo] = []
        for operand, indexing_map in zip(
            op.operands, op.get_indexing_maps(), strict=True
        ):
            source_type = operand.type
            assert isa(source_type, MemRefType | TensorType)
            operand_infos_list.append(
                OperandTileInfo.analyze(indexing_map.data, source_type)
            )
        operand_infos = tuple(operand_infos_list)

        return TilingPlan(
            loop_ranges=loop_ranges,
            tiled_dims=tiled_dims,
            partial_tiled_dims=partial_tiled_dims,
            operand_infos=operand_infos,
            tile_sizes=normalized_tile_sizes,
        )

@dataclass(frozen=True)
class SliceParameters:
    """
    Where one operand's tile sits within that operand.

    This is the geometry of the tile, which is the same whether the operand is a
    memref or a tensor, and so does not depend on which op ends up materializing
    the slice.
    """

    offsets: tuple[SSAValue | int, ...]
    sizes: tuple[SSAValue | int, ...]
    strides: tuple[SSAValue | int, ...]

    @staticmethod
    def compute(
        rewriter: PatternRewriter,
        insertion_point: InsertPoint,
        indexing_map: AffineMap,
        operand_info: OperandTileInfo,
        tiled_loop_ivs: dict[int, SSAValue],
        effective_tile_sizes: dict[int, SSAValue | int],
        loop_ranges: Sequence[int],
    ) -> "SliceParameters":
        """
        Compute the offsets and sizes for one operand's `memref.subview` or
        `tensor.extract_slice`.

        There is one offset and one size per operand dimension, and the indexing
        map has one result per operand dimension saying how the loops reach it.

        - A result that no tiled loop appears in is not sliced at all: offset 0,
          and the whole dimension for its size.
        - A result that is just `dN`, for a tiled loop, takes that loop's
          induction variable as its offset and its tile size as its size.
        - Any other result, `d0 + d1` or `d0 * 2`, gets both from an
          `affine.apply`. Its offset evaluates the result with each tiled loop at
          its induction variable and the rest at zero, less what it evaluates to
          with every loop at zero. Its size evaluates it at the last index each
          loop reaches inside the tile, plus one to turn that index back into a
          count.
        """

        source_shape = operand_info.source_type.get_shape()
        num_loops = indexing_map.num_dims

        # The induction variable of each tiled loop, and zero for the rest.
        starts: tuple[SSAValue | int, ...] = tuple(
            tiled_loop_ivs.get(dim, 0) for dim in range(num_loops)
        )
        zeros = (0,) * num_loops
        # The last index each loop reaches inside the tile. This costs an
        # `affine.apply` per loop, so it is built once, and only if some result
        # turns out to need it.
        extents: tuple[SSAValue | int, ...] | None = None

        offsets: list[SSAValue | int] = []
        sizes: list[SSAValue | int] = []
        for result_index, expr in enumerate(indexing_map.results):
            if not (expr.used_dims() & tiled_loop_ivs.keys()):
                offsets.append(0)
                sizes.append(source_shape[result_index])
                continue

            # `dN` on its own: the slice starts at that loop's induction
            # variable and is one tile long. The general case below computes
            # exactly this, so taking it here just avoids emitting an
            # `affine.apply` to say so.
            if isinstance(expr, AffineDimExpr):
                offsets.append(tiled_loop_ivs[expr.position])
                sizes.append(effective_tile_sizes[expr.position])
                continue

            if extents is None:
                extents = tuple(
                    _decrement(
                        rewriter,
                        insertion_point,
                        effective_tile_sizes.get(dim, size),
                    )
                    for dim, size in enumerate(loop_ranges[:num_loops])
                )

            # Take off what the result evaluates to with every loop at zero, so
            # that a constant the map adds does not shift the slice: the tiled op
            # applies the same map inside the slice and adds it back there.
            # Taking it off the expression rather than off the value it produces
            # keeps this to a single `affine.apply`.
            at_zero = expr.eval(zeros, ())
            offsets.append(
                _apply_affine_expr(
                    rewriter,
                    insertion_point,
                    expr if at_zero == 0 else expr - at_zero,
                    starts,
                )
            )
            sizes.append(
                _apply_affine_expr(rewriter, insertion_point, expr + 1, extents)
            )

        return SliceParameters(tuple(offsets), tuple(sizes), (1,) * len(source_shape))

@dataclass(frozen=True)
class TilingResult:
    """The tiled op, enclosing loops (outermost first), and replacement results."""

    tiled_op: linalg.abstract_ops.LinalgStructuredOperation
    loops: tuple[scf.ForOp, ...]
    replacements: tuple[SSAValue, ...]