Test linalg tiling

test_linalg_tiling

TileLinalgFromAttributePattern

Bases: RewritePattern

Rewrite supported structured linalg ops annotated with test_tile_sizes into tiled form.

A tile size is normally taken straight from that attribute. Dimensions named by test_dynamic_tile_sizes instead take a tile size that is not known until the op runs, which the pass has no way of writing in an attribute, so one is produced by a test.op for the tiling to consume.

Source code in xdsl/transforms/test_linalg_tiling.py

class TileLinalgFromAttributePattern(RewritePattern):
    """
    Rewrite supported structured linalg ops annotated with `test_tile_sizes` into
    tiled form.

    A tile size is normally taken straight from that attribute. Dimensions named
    by `test_dynamic_tile_sizes` instead take a tile size that is not known until
    the op runs, which the pass has no way of writing in an attribute, so one is
    produced by a `test.op` for the tiling to consume.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(
        self,
        op: linalg.abstract_ops.LinalgStructuredOperation,
        rewriter: PatternRewriter,
        /,
    ) -> None:

        tile_sizes_attr = op.attributes.get("test_tile_sizes")
        if tile_sizes_attr is None:
            return

        assert isa(tile_sizes_attr, DenseArrayBase[IntegerType])
        tile_sizes: list[SSAValue | int] = list(tile_sizes_attr.get_values())

        dynamic_dims_attr = op.attributes.get("test_dynamic_tile_sizes")
        if dynamic_dims_attr is not None:
            assert isa(dynamic_dims_attr, DenseArrayBase[IntegerType])
            for dim in dynamic_dims_attr.get_values():
                tile_size_op = test.TestOp(result_types=[IndexType()])
                rewriter.insert(tile_size_op, InsertPoint.before(op))
                tile_sizes[dim] = tile_size_op.res[0]

        tile_structured_op(rewriter, op, tile_sizes)

match_and_rewrite(op: linalg.abstract_ops.LinalgStructuredOperation, rewriter: PatternRewriter) -> None

Source code in xdsl/transforms/test_linalg_tiling.py

@op_type_rewrite_pattern
def match_and_rewrite(
    self,
    op: linalg.abstract_ops.LinalgStructuredOperation,
    rewriter: PatternRewriter,
    /,
) -> None:

    tile_sizes_attr = op.attributes.get("test_tile_sizes")
    if tile_sizes_attr is None:
        return

    assert isa(tile_sizes_attr, DenseArrayBase[IntegerType])
    tile_sizes: list[SSAValue | int] = list(tile_sizes_attr.get_values())

    dynamic_dims_attr = op.attributes.get("test_dynamic_tile_sizes")
    if dynamic_dims_attr is not None:
        assert isa(dynamic_dims_attr, DenseArrayBase[IntegerType])
        for dim in dynamic_dims_attr.get_values():
            tile_size_op = test.TestOp(result_types=[IndexType()])
            rewriter.insert(tile_size_op, InsertPoint.before(op))
            tile_sizes[dim] = tile_size_op.res[0]

    tile_structured_op(rewriter, op, tile_sizes)

TestLinalgTilingPass dataclass

Bases: ModulePass

Tile supported structured linalg ops annotated with test_tile_sizes.

Source code in xdsl/transforms/test_linalg_tiling.py

@dataclass(frozen=True)
class TestLinalgTilingPass(ModulePass):
    """
    Tile supported structured linalg ops annotated with `test_tile_sizes`.
    """

    name = "test-linalg-tiling"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        PatternRewriteWalker(
            TileLinalgFromAttributePattern(),
            apply_recursively=False,
        ).rewrite_module(op)

name = 'test-linalg-tiling' class-attribute instance-attribute

__init__() -> None

apply(ctx: Context, op: ModuleOp) -> None

Source code in xdsl/transforms/test_linalg_tiling.py

def apply(self, ctx: Context, op: ModuleOp) -> None:
    PatternRewriteWalker(
        TileLinalgFromAttributePattern(),
        apply_recursively=False,
    ).rewrite_module(op)