Function test_size_span

tests/python/tirx/test_layout.py:983–1064

()

def test_size_span():
    def tile_layout_size():
        layout = TileLayout(S[(8, 8) : (8, 1)])
        assert layout.size() == 64

    tile_layout_size()

    def swizzle_layout_size():
        layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
        assert layout.size() == 512
        layout = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3)
        assert layout.size() == 1024

    swizzle_layout_size()

    def compose_layout_size():
        layout = ComposeLayout(
            SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3),
            TileLayout(S[(8, 64) : (64, 1)]),
        )
        assert layout.size() == 512

    compose_layout_size()

    def tile_layout_span():
        layout = TileLayout(S[(8, 8) : (8, 1)])
        assert layout.span() == 64
        layout = TileLayout(S[(8, 6) : (8, 1)])
        assert layout.span() == 62
        layout = TileLayout(S[(8, 1, 4, 2) : (4 @ laneid, 2, 1 @ laneid, 1)])
        assert layout.span() == 2

    tile_layout_span()

    def swizzle_layout_span():
        layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
        assert layout.span() == 512
        layout = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3)
        assert layout.span() == 1024

    swizzle_layout_span()

    def compose_layout_span():
        layout = ComposeLayout(
            SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3),
            TileLayout(S[(8, 64) : (64, 1)]),
        )
        assert layout.span() == 512

    compose_layout_span()

    def trainium_layout_tests():
        # TrainiumLayout tests
        layout = TileLayout(S[(8, 8) : (1 @ P, 1 @ F)])
        assert layout.size("P") == 8
        assert layout.size("F") == 8

        layout = TileLayout(S[(8, 8, 8) : (64 @ F, 1 @ P, 1 @ F)])