Tiling

tiling

OperandTileInfo dataclass

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 code in xdsl/dialects/linalg/transforms/tiling.py

@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)

source_type: MemRefType[Attribute] | TensorType[Attribute] instance-attribute

loop_dims: tuple[int | None, ...] instance-attribute

__init__(source_type: MemRefType[Attribute] | TensorType[Attribute], loop_dims: tuple[int | None, ...]) -> None

analyze(indexing_map: AffineMap, source_type: MemRefType[Attribute] | TensorType[Attribute]) -> OperandTileInfo staticmethod

Analyze how one operand should be sliced for each tile.

Source code in xdsl/dialects/linalg/transforms/tiling.py

@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)

TilingPlan dataclass

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.

Source code in xdsl/dialects/linalg/transforms/tiling.py

@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,
        )

loop_ranges: tuple[int, ...] instance-attribute

tiled_dims: tuple[int, ...] instance-attribute

partial_tiled_dims: frozenset[int] instance-attribute

operand_infos: tuple[OperandTileInfo, ...] instance-attribute

tile_sizes: tuple[SSAValue | int, ...] instance-attribute

__init__(loop_ranges: tuple[int, ...], tiled_dims: tuple[int, ...], partial_tiled_dims: frozenset[int], operand_infos: tuple[OperandTileInfo, ...], tile_sizes: tuple[SSAValue | int, ...]) -> None

analyze(op: linalg.abstract_ops.LinalgStructuredOperation, tile_sizes: Sequence[SSAValue | int]) -> TilingPlan staticmethod

Analyze one supported structured linalg op and return a TilingPlan.

Source code in xdsl/dialects/linalg/transforms/tiling.py

@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,
    )

SliceParameters dataclass

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.

Source code in xdsl/dialects/linalg/transforms/tiling.py

@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))

offsets: tuple[SSAValue | int, ...] instance-attribute

sizes: tuple[SSAValue | int, ...] instance-attribute

strides: tuple[SSAValue | int, ...] instance-attribute

__init__(offsets: tuple[SSAValue | int, ...], sizes: tuple[SSAValue | int, ...], strides: tuple[SSAValue | int, ...]) -> None

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 staticmethod

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 code in xdsl/dialects/linalg/transforms/tiling.py

@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))

tile_structured_op(rewriter: PatternRewriter, op: linalg.abstract_ops.LinalgStructuredOperation, tile_sizes: Sequence[SSAValue | int]) -> bool

Rewrite supported structured linalg ops into tiled form.

Source code in xdsl/dialects/linalg/transforms/tiling.py

def tile_structured_op(
    rewriter: PatternRewriter,
    op: linalg.abstract_ops.LinalgStructuredOperation,
    tile_sizes: Sequence[SSAValue | int],
) -> bool:
    """
    Rewrite supported structured linalg ops into tiled form.
    """
    try:
        plan = TilingPlan.analyze(op, tile_sizes)
    except (ValueError, NotImplementedError) as e:
        raise PassFailedException(str(e)) from e

    if not plan.tiled_dims:
        return False

    # Outputs with value semantics are threaded through the loops, since each
    # tile produces a new value instead of writing through a view.
    has_tensor_outputs = bool(op.res)
    iter_args = tuple(op.outputs) if has_tensor_outputs else ()

    loops, tiled_loop_ivs, inner_ip = _build_tile_loops(
        rewriter,
        InsertPoint.before(op),
        plan.loop_ranges,
        plan.tile_sizes,
        plan.tiled_dims,
        _loop_range_sources(op, plan),
        iter_args,
    )

    num_inputs = len(op.inputs)
    effective_tile_sizes = _build_effective_tile_sizes(rewriter, inner_ip, plan, loops)

    slice_parameters = tuple(
        SliceParameters.compute(
            rewriter,
            inner_ip,
            indexing_map.data,
            operand_info,
            tiled_loop_ivs,
            effective_tile_sizes,
            plan.loop_ranges,
        )
        for operand_info, indexing_map in zip(
            plan.operand_infos, op.get_indexing_maps(), strict=True
        )
    )

    # Output tensors are sliced from the values the innermost loop carries, so
    # that each tile builds on the tiles the surrounding iterations wrote back.
    # Slicing the originals instead would discard their work.
    carried = loops[-1].body.block.args[1:]
    slice_sources = list(op.operands)
    slice_sources[num_inputs:] = carried if has_tensor_outputs else op.outputs

    tiled_operands = [
        _build_tiled_slice(
            rewriter, inner_ip, source, operand_info.source_type, parameters
        )
        for source, operand_info, parameters in zip(
            slice_sources, plan.operand_infos, slice_parameters, strict=True
        )
    ]

    tiled_outputs = tiled_operands[num_inputs:]
    # A tile is the same op over its slices, whatever op that is, so it is built
    # as one of those rather than as a generic. A named op says what it computes
    # by which op it is, which tiling it has no reason to take away from it.
    tiled_op = type(op).create(
        operands=tiled_operands,
        result_types=(
            tuple(value.type for value in tiled_outputs) if has_tensor_outputs else ()
        ),
        properties=dict(op.properties),
        regions=[op.body.clone()],
    )
    rewriter.insert(tiled_op, inner_ip)
    _offset_tiled_indices(rewriter, tiled_op, tiled_loop_ivs)

    # Memref outputs are written through their subview, so only tensor outputs
    # need their computed tile written back into the value being carried.
    yielded: Sequence[SSAValue] = (
        tuple(
            _build_tiled_insert(
                rewriter, inner_ip, tiled_result, destination, parameters
            )
            for tiled_result, destination, parameters in zip(
                tiled_op.res, carried, slice_parameters[num_inputs:], strict=True
            )
        )
        if has_tensor_outputs
        else ()
    )

    # The innermost loop yields the updated tensors, and each enclosing loop
    # yields the results of the loop nested inside it.
    for loop in reversed(loops):
        rewriter.insert(scf.YieldOp(*yielded), InsertPoint.at_end(loop.body.block))
        yielded = loop.results

    # The outermost loop carries out the fully updated tensors. An op tiling
    # memrefs carries nothing and has no results, so this replaces it with
    # nothing, which is the erase that case needs.
    rewriter.replace(op, [], loops[0].results)
    return True