Test linalg tiling

test_linalg_tiling

TileLinalgGenericFromAttributePattern

Bases: RewritePattern

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

Source code in xdsl/transforms/test_linalg_tiling.py

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

    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: linalg.ops.GenericOp, 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 = tuple(tile_sizes_attr.get_values())
        tile_linalg_generic(rewriter, op, tile_sizes)

match_and_rewrite(op: linalg.ops.GenericOp, rewriter: PatternRewriter) -> None

Source code in xdsl/transforms/test_linalg_tiling.py

@op_type_rewrite_pattern
def match_and_rewrite(
    self, op: linalg.ops.GenericOp, 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 = tuple(tile_sizes_attr.get_values())
    tile_linalg_generic(rewriter, op, tile_sizes)

TestLinalgTilingPass dataclass

Bases: ModulePass

Tile supported memref-based linalg.generic ops annotated with test_tile_sizes.

Source code in xdsl/transforms/test_linalg_tiling.py

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

    name = "test-linalg-tiling"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        PatternRewriteWalker(
            TileLinalgGenericFromAttributePattern(),
            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(
        TileLinalgGenericFromAttributePattern(),
        apply_recursively=False,
    ).rewrite_module(op)