| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=missing-module-docstring, missing-function-docstring, missing-class-docstring |
| import functools |
| import itertools |
| import operator |
| |
| import pytest |
| |
| import tvm |
| from tvm.arith import Analyzer |
| from tvm.ir import assert_structural_equal |
| from tvm.ir.type import PointerType, PrimType |
| from tvm.script import tirx as T |
| from tvm.script.ir_builder import IRBuilder |
| from tvm.script.ir_builder import tirx as Tx_builder |
| from tvm.tirx import Var |
| from tvm.tirx.cuda.tile_primitive.tma_utils import ( |
| SwizzleMode, |
| mma_shared_layout, |
| ) |
| from tvm.tirx.layout import ( |
| Axis, |
| ComposeLayout, |
| F, |
| Iter, |
| P, |
| R, |
| S, |
| TCol, |
| TileLayout, |
| TLane, |
| laneid, |
| m, |
| tid_in_wg, |
| tmem_mma_operand_layout, |
| tx, |
| warpid, |
| wg_local_layout, |
| wgid, |
| wid_in_wg, |
| ) |
| |
| |
| def test_axis(): |
| assert Axis.bx == Axis.get("bx") |
| assert Axis.by == Axis.get("by") |
| assert Axis.bz == Axis.get("bz") |
| assert Axis.cbx == Axis.get("cbx") |
| assert Axis.cby == Axis.get("cby") |
| assert Axis.cbz == Axis.get("cbz") |
| assert Axis.tx == Axis.get("tx") |
| assert Axis.warpid == Axis.get("warpid") |
| assert Axis.laneid == Axis.get("laneid") |
| assert Axis.wgid == Axis.get("wgid") |
| assert Axis.tid_in_wg == Axis.get("tid_in_wg") |
| assert Axis.wid_in_wg == Axis.get("wid_in_wg") |
| assert Axis.m == Axis.get("m") |
| assert Axis.P == Axis.get("P") |
| assert Axis.F == Axis.get("F") |
| assert Axis.TCol == Axis.get("TCol") |
| assert Axis.TLane == Axis.get("TLane") |
| |
| assert Axis.bx.is_thread() |
| assert Axis.by.is_thread() |
| assert Axis.bz.is_thread() |
| assert Axis.cbx.is_thread() |
| assert Axis.cby.is_thread() |
| assert Axis.cbz.is_thread() |
| assert Axis.tx.is_thread() |
| assert Axis.warpid.is_thread() |
| assert Axis.laneid.is_thread() |
| assert Axis.wgid.is_thread() |
| assert Axis.tid_in_wg.is_thread() |
| assert Axis.wid_in_wg.is_thread() |
| assert Axis.m.is_memory() |
| assert Axis.P.is_memory() |
| assert Axis.F.is_memory() |
| assert Axis.TCol.is_memory() |
| assert Axis.TLane.is_memory() |
| |
| assert Axis.bx.get_scope().name == "thread" |
| assert Axis.bx.get_subscope().name == "cta" |
| |
| |
| def test_constructor(): |
| def assert_tile_layout(layout, shard, replica=None, offset=None): |
| expected = TileLayout.from_iters(shard, replica or [], offset or {}) |
| assert_structural_equal(layout, expected) |
| |
| layout = TileLayout(S[2, 3, 4]) |
| assert_tile_layout(layout, [Iter(2, 12, "m"), Iter(3, 4, "m"), Iter(4, 1, "m")]) |
| |
| layout = TileLayout(S[(2, 3, 4) : (12, 4, 1)]) |
| assert_tile_layout(layout, [Iter(2, 12, "m"), Iter(3, 4, "m"), Iter(4, 1, "m")]) |
| |
| layout = TileLayout(S[(2, 3, 4) : (12 @ m, 4 @ m, 1 @ m)]) |
| assert_tile_layout(layout, [Iter(2, 12, "m"), Iter(3, 4, "m"), Iter(4, 1, "m")]) |
| |
| layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| assert_tile_layout(layout, [Iter(8, 4, "laneid"), Iter(4, 1, "laneid"), Iter(2, 1, "m")]) |
| |
| layout = TileLayout(S[8 : 4 @ laneid] + R[4 : 1 @ laneid]) |
| assert_tile_layout(layout, [Iter(8, 4, "laneid")], replica=[Iter(4, 1, "laneid")]) |
| |
| layout = TileLayout(S[8 : 4 @ laneid] + 1 @ laneid) |
| assert_tile_layout(layout, [Iter(8, 4, "laneid")], offset={laneid: 1}) |
| |
| |
| def test_constructor_multi_term_offset(): |
| """Multiple offset terms can be chained with `+` without parens. |
| |
| `_LayoutSpec.__add__` previously overwrote `self.offset` on each call, |
| silently dropping all but the last axis term in |
| `S[..] + 1 @ a + 2 @ b + 64`. Verify the merge happens for every entry |
| point: `_LayoutSpec + _OnAxis`, `_LayoutSpec + int`, |
| `_LayoutSpec + _OffsetExpr`, and the parenthesised form (which already |
| worked) producing the same result. |
| """ |
| |
| # Chained, no parens: must merge into all three axes. |
| layout = TileLayout(S[8 : 4 @ laneid] + 1 @ laneid + 2 @ warpid + 64) |
| assert dict(layout.offset) == {laneid: 1, warpid: 2, m: 64} |
| |
| # Parenthesised form must produce the same offset. |
| parens = TileLayout(S[8 : 4 @ laneid] + (1 @ laneid + 2 @ warpid + 64)) |
| assert_structural_equal(layout, parens) |
| |
| # Single-axis offset still works (regression sanity). |
| single = TileLayout(S[8 : 4 @ laneid] + 1 @ laneid) |
| assert dict(single.offset) == {laneid: 1} |
| |
| # Bare-int offset alone still routes to `m`. |
| bare = TileLayout(S[8 : 4 @ laneid] + 64) |
| assert dict(bare.offset) == {m: 64} |
| |
| # `_LayoutSpec + _LayoutSpec` where both carry an offset must also merge. |
| a = S[8 : 4 @ laneid] + 1 @ laneid |
| b = R[4 : 1 @ laneid] + 2 @ warpid |
| combined = TileLayout(a + b) |
| assert dict(combined.offset) == {laneid: 1, warpid: 2} |
| |
| # `int + _LayoutSpec` reaches `_LayoutSpec.__radd__` (Python's `int.__add__` |
| # returns NotImplemented for `_LayoutSpec`); verify it merges through the |
| # same path as `__add__`. |
| radd = TileLayout(64 + S[8 : 4 @ laneid] + 1 @ laneid) |
| assert dict(radd.offset) == {laneid: 1, m: 64} |
| |
| |
| def test_wg_local_layout_helper(): |
| layout = wg_local_layout(16) |
| expected = TileLayout(S[(128, 16) : (1 @ tid_in_wg, 1)]) |
| assert_structural_equal(layout.canonicalize(), expected.canonicalize()) |
| |
| layout_rows = wg_local_layout(8, rows=64) |
| expected_rows = TileLayout(S[(64, 8) : (1 @ tid_in_wg, 1)]) |
| assert_structural_equal(layout_rows.canonicalize(), expected_rows.canonicalize()) |
| |
| |
| def test_spec_builder(): |
| """Test S[shape:stride] + R[shape:stride] + offset combinator API.""" |
| |
| # --- S[shape:stride] shard only --- |
| new = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| old = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| assert str(new) == str(old) |
| |
| # --- 1D (no inner parens) --- |
| new = TileLayout(S[128 : 1 @ laneid]) |
| old = TileLayout(S[128 : 1 @ laneid]) |
| assert str(new) == str(old) |
| |
| # --- Extents only --- |
| new = TileLayout(S[8, 4, 2]) |
| old = TileLayout(S[8, 4, 2]) |
| assert str(new) == str(old) |
| |
| # --- S + R (shard + replica) --- |
| new = TileLayout(S[(8,) : (4 @ laneid,)] + R[4 : 1 @ laneid]) |
| old = TileLayout(S[8 : 4 @ laneid] + R[4 : 1 @ laneid]) |
| assert str(new) == str(old) |
| |
| # --- S + offset --- |
| new = TileLayout(S[8 : 4 @ laneid] + 1 @ laneid) |
| old = TileLayout(S[8 : 4 @ laneid] + 1 @ laneid) |
| assert str(new) == str(old) |
| |
| # --- S + R + offset --- |
| new = TileLayout(S[(1,) : (1,)] + R[(8, 4) : (4 @ laneid, 1 @ laneid)] + 2 @ warpid) |
| old = TileLayout(S[1:1] + R[(8, 4) : (4 @ laneid, 1 @ laneid)] + 2 @ warpid) |
| assert str(new) == str(old) |
| |
| # --- Memory axes --- |
| new = TileLayout(S[(2, 3, 4) : (12 @ m, 4 @ m, 1 @ m)]) |
| old = TileLayout(S[(2, 3, 4) : (12 @ m, 4 @ m, 1 @ m)]) |
| assert str(new) == str(old) |
| |
| # --- String axis names (no import needed) --- |
| # stride=1 shorthand |
| assert str(TileLayout(S[8:"laneid"])) == str(TileLayout(S[8 : 1 @ laneid])) |
| assert str(TileLayout(S[32:"warpid"])) == str(TileLayout(S[32 : 1 @ warpid])) |
| # multi-dim with string |
| assert str(TileLayout(S[(8, 4) : ("laneid", 1)])) == str( |
| TileLayout(S[(8, 4) : (1 @ laneid, 1)]) |
| ) |
| # non-unit stride via tuple |
| assert str(TileLayout(S[(8,) : ((4, "laneid"),)])) == str(TileLayout(S[8 : 4 @ laneid])) |
| # string in R |
| assert str(TileLayout(S[1:1] + R[4:"laneid"])) == str(TileLayout(S[1:1] + R[4 : 1 @ laneid])) |
| |
| |
| def test_verify_well_formed(): |
| def test_scope_connected(): |
| layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| res = layout.get_scope() |
| assert res is not None |
| assert res[0].name == "thread" |
| assert res[1].name == "warp" |
| assert layout.verify_well_formed() |
| |
| layout = TileLayout(S[8 : 4 @ laneid] + R[4 : 1 @ laneid]) |
| res = layout.get_scope() |
| assert res is not None |
| assert res[0].name == "thread" |
| assert res[1].name == "warp" |
| assert layout.verify_well_formed() |
| |
| layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| res = layout.get_scope() |
| assert res is not None |
| assert res[0].name == "thread" |
| assert res[1].name == "warp" |
| assert layout.verify_well_formed() |
| |
| layout = TileLayout( |
| S[(2, 8, 2, 4, 2) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid, 1)] |
| ) |
| res = layout.get_scope() |
| assert res is not None |
| assert res[0].name == "thread" |
| assert res[1].name == "cta" |
| assert layout.verify_well_formed() |
| |
| layout = TileLayout( |
| S[(2, 8, 2, 4, 2) : (2 @ wid_in_wg, 4 @ laneid, 1 @ wid_in_wg, 1 @ laneid, 1)] |
| ) |
| res = layout.get_scope() |
| assert res is not None |
| assert res[0].name == "thread" |
| assert res[1].name == "warpgroup" |
| assert layout.verify_well_formed() |
| |
| layout = TileLayout(S[(2, 8, 2, 4, 2) : (2 @ wgid, 4 @ laneid, 1 @ wgid, 1 @ laneid, 1)]) |
| with pytest.raises(Exception): |
| layout.verify_well_formed() |
| |
| test_scope_connected() |
| |
| |
| def test_normalize_tile_layout(): |
| def case1(): |
| layout = TileLayout(S[(8, 8, 8, 4, 2) : (512, 64, 8, 2, 1)]) |
| layout_expected = TileLayout(S[4096:1]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case1() |
| |
| def case2(): |
| layout = TileLayout(S[(8, 8, 1, 8, 4, 2) : (512, 64, 160, 8, 2, 1)]) |
| layout_expected = TileLayout(S[4096:1]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case2() |
| |
| def case3(): |
| layout = TileLayout(S[(8, 8, 8, 4, 1, 1) : (512, 64, 8, 2, 1, 1)]) |
| layout_expected = TileLayout(S[2048:2]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case3() |
| |
| def case4(): |
| layout = TileLayout(S[(8, 8, 1, 1, 1, 4, 1, 1) : (512, 64, 1, 1, 1, 2, 1, 1)]) |
| layout_expected = TileLayout(S[(64, 4) : (64, 2)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case4() |
| |
| def case5(): |
| layout = TileLayout(S[(2, 3, 6) : (18, 6, 1)]) |
| layout_expected = TileLayout(S[36:1]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case5() |
| |
| def case6(): |
| layout = TileLayout(S[(8, 2, 3, 6) : (6, 18, 6, 1)]) |
| layout_expected = TileLayout(S[(8, 36) : (6, 1)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case6() |
| |
| def case7(): |
| layout = TileLayout(S[(8, 2, 3, 6) : (6, 24, 6, 1)]) |
| layout_expected = TileLayout(S[(8, 2, 18) : (6, 24, 1)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case7() |
| |
| def case8(): |
| layout = TileLayout(S[(8, 2, 4, 2, 3, 6) : (2, 1, 4, 24, 6, 1)]) |
| layout_expected = TileLayout(S[(16, 4, 2, 18) : (1, 4, 24, 1)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case8() |
| |
| def case9(): |
| layout = TileLayout(S[(3, 4, 5, 2) : (20, 5, 1, 60)]) |
| layout_expected = TileLayout(S[(60, 2) : (1, 60)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case9() |
| |
| def case10(): |
| layout = TileLayout(S[(18, 8, 2, 4, 2, 3, 6) : (4, 2, 1, 4, 24, 6, 1)]) |
| layout_expected = TileLayout(S[(18, 16, 4, 2, 18) : (4, 1, 4, 24, 1)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case10() |
| |
| def case11(): |
| layout = TileLayout(S[(3, 4, 5, 2, 3, 4) : (20, 5, 1, 60, 20, 5)]) |
| layout_expected = TileLayout(S[(60, 24) : (1, 5)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case11() |
| |
| def case_no_norm(): |
| layout_normalized = TileLayout(S[(8, 8, 8, 4, 2) : (16, 4 @ laneid, 2, 1 @ laneid, 1)]) |
| assert_structural_equal(layout_normalized, layout_normalized.canonicalize()) |
| |
| case_no_norm() |
| |
| def case_both_data_device1(): |
| layout = TileLayout(S[(8, 8, 8, 1, 4, 2, 1) : (16, 4 @ laneid, 2, 1, 1 @ laneid, 1, 1)]) |
| layout_normalized = TileLayout(S[(8, 8, 8, 4, 2) : (16, 4 @ laneid, 2, 1 @ laneid, 1)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device1() |
| |
| def case_both_data_device2(): |
| layout = TileLayout( |
| S[(8, 8, 8, 1, 4, 2, 1) : (16, 4 @ laneid, 2, 1, 1 @ laneid, 1, 4 @ laneid)] |
| ) |
| layout_normalized = TileLayout(S[(8, 8, 8, 4, 2) : (16, 4 @ laneid, 2, 1 @ laneid, 1)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device2() |
| |
| def case_both_data_device3(): |
| layout = TileLayout( |
| S[(8, 8, 8, 1, 1, 2, 1) : (16, 4 @ laneid, 2, 1, 4 @ laneid, 1, 1)] + 0 @ laneid |
| ) |
| layout_normalized = TileLayout(S[(8, 8, 16) : (16, 4 @ laneid, 1)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device3() |
| |
| def case_both_data_device4(): |
| layout = TileLayout(S[(8, 4, 8, 8, 16) : (4 @ laneid, 1 @ laneid, 4, 2, 4)]) |
| layout_normalized = TileLayout(S[(32, 8, 8, 16) : (1 @ laneid, 4, 2, 4)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device4() |
| |
| def case_both_data_device6(): |
| layout = TileLayout(S[(8, 4, 8, 16) : (4 @ laneid, 1 @ laneid, 2, 4)]) |
| layout_normalized = TileLayout(S[(32, 8, 16) : (1 @ laneid, 2, 4)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device6() |
| |
| def case_both_data_device7(): |
| layout = TileLayout(S[(8, 4, 8) : (4 @ laneid, 1 @ laneid, 8)]) |
| layout_normalized = TileLayout(S[(32, 8) : (1 @ laneid, 8)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device7() |
| |
| def case_both_data_device8(): |
| # Fuse-Case 1 |
| layout = TileLayout(S[(8, 4, 8) : (4 @ laneid, 1 @ laneid, 4)]) |
| layout_normalized = TileLayout(S[(32, 8) : (1 @ laneid, 4)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device8() |
| |
| def case_both_data_device9(): |
| # Fuse-Case 2 |
| layout = TileLayout(S[(8, 4) : (4 @ laneid, 1 @ laneid)]) |
| layout_normalized = TileLayout(S[32 : 1 @ laneid]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device9() |
| |
| def case_both_data_device12(): |
| # Fuse-mixed |
| layout = TileLayout(S[(8, 4, 4, 8, 8, 8) : (4 @ laneid, 1 @ laneid, 4, 8, 8, 8)]) |
| layout_normalized = TileLayout(S[(32, 4, 8, 8, 8) : (1 @ laneid, 4, 8, 8, 8)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device12() |
| |
| def case_both_data_device13(): |
| # Fuse-mixed with partial |
| layout = TileLayout(S[(8, 4, 4, 8, 8, 8) : (4 @ laneid, 1 @ laneid, 16, 2, 8, 8)]) |
| layout_normalized = TileLayout(S[(32, 32, 8, 8) : (1 @ laneid, 2, 8, 8)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device13() |
| |
| def case_both_data_device14(): |
| # Fuse-mixed with partial (another case) |
| layout = TileLayout( |
| S[(8, 4, 4, 8, 8, 4, 4, 16, 8) : (4 @ laneid, 1 @ laneid, 16, 2, 8, 2, 16, 1, 4)] |
| ) |
| layout_normalized = TileLayout(S[(32, 32, 32, 64, 8) : (1 @ laneid, 2, 2, 1, 4)]) |
| assert_structural_equal(layout_normalized, layout.canonicalize()) |
| |
| case_both_data_device14() |
| |
| def case15(): |
| # Only data tree (partial norm - middle) #15 |
| layout = TileLayout(S[(32, 3, 4, 5, 2, 3, 4) : (1 @ laneid, 20, 5, 1, 60, 20, 5)]) |
| layout_expected = TileLayout(S[(32, 60, 24) : (1 @ laneid, 1, 5)]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case15() |
| |
| def unit_layout_case1(): |
| layout = TileLayout(S[(1, 1, 1, 1, 1) : (1, 1, 1, 1, 1)]) |
| layout_unit = TileLayout(S[1:1]) |
| assert_structural_equal(layout_unit, layout.canonicalize()) |
| |
| unit_layout_case1() |
| |
| def case_fuse_axis(): |
| with tvm.target.Target("cuda"): |
| layout = TileLayout(S[(2, 8, 2, 4) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid)]) |
| layout_expected = TileLayout(S[(2, 8, 2, 4) : (64 @ tx, 4 @ tx, 32 @ tx, 1 @ tx)]) |
| assert layout.verify_well_formed() |
| assert layout_expected.verify_well_formed() |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| layout = TileLayout(S[(2, 2, 8, 4) : (2 @ warpid, 1 @ warpid, 4 @ laneid, 1 @ laneid)]) |
| layout_expected = TileLayout(S[128 : 1 @ tx]) |
| assert layout.verify_well_formed() |
| assert layout_expected.verify_well_formed() |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| layout = TileLayout( |
| S[ |
| (2, 2, 8, 2, 2, 4) : ( |
| 2 @ wgid, |
| 2 @ wid_in_wg, |
| 4 @ laneid, |
| 1 @ wgid, |
| 1 @ wid_in_wg, |
| 1 @ laneid, |
| ) |
| ] |
| ) |
| layout_expected = TileLayout( |
| S[(2, 2, 8, 2, 2, 4) : (256 @ tx, 64 @ tx, 4 @ tx, 128 @ tx, 32 @ tx, 1 @ tx)] |
| ) |
| assert layout.verify_well_formed() |
| assert layout_expected.verify_well_formed() |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| layout = TileLayout( |
| S[(2, 8, 2, 4) : (2 @ wid_in_wg, 4 @ laneid, 1 @ wid_in_wg, 1 @ laneid)] |
| ) |
| layout_expected = TileLayout( |
| S[(2, 8, 2, 4) : (64 @ tid_in_wg, 4 @ tid_in_wg, 32 @ tid_in_wg, 1 @ tid_in_wg)] |
| ) |
| assert layout.verify_well_formed() |
| assert layout_expected.verify_well_formed() |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| layout = TileLayout( |
| S[(2, 2, 4, 32) : (2 @ wgid, 1 @ wgid, 32 @ tid_in_wg, 1 @ tid_in_wg)] |
| ) |
| layout_expected = TileLayout(S[512 : 1 @ tx]) |
| assert layout.verify_well_formed() |
| assert layout_expected.verify_well_formed() |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case_fuse_axis() |
| |
| def case_sort_replicate_exclude_iters(): |
| layout1 = TileLayout(S[1:1] + R[(8, 4) : (4 @ laneid, 1 @ laneid)] + 2 @ warpid) |
| layout2 = TileLayout(S[1:1] + R[(4, 8) : (1 @ laneid, 4 @ laneid)] + 2 @ warpid) |
| assert_structural_equal(layout1.canonicalize(), layout2.canonicalize()) |
| |
| case_sort_replicate_exclude_iters() |
| |
| def case_empty_shard_canonicalize(): |
| """Regression test for F6: canonicalize must not crash when layout->shard is empty.""" |
| layout = TileLayout(R[32 : 1 @ laneid]) |
| canon = layout.canonicalize() |
| assert canon is not None |
| |
| case_empty_shard_canonicalize() |
| |
| |
| def test_tile_layout(): |
| def case1(): |
| # (8):(1)x(8):(1) -> (64):(1) |
| inner = TileLayout(S[8:1]) |
| outer = inner |
| layout_tile = TileLayout(S[64:1]) |
| assert_structural_equal(layout_tile, inner.tile(outer, [8], [8])) |
| |
| outer_res = inner.is_tile_inner(layout_tile, [64], [8]) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, [64], [8]) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| case1() |
| |
| def case2(): |
| # (8,8):(8,1)x(8,8):(8,1) -> (8,8,8,8):(512,8,64,1) |
| inner = TileLayout(S[(8, 8) : (8, 1)]) |
| outer = inner |
| layout_tile = TileLayout(S[(8, 8, 8, 8) : (512, 8, 64, 1)]) |
| assert_structural_equal(layout_tile, inner.tile(outer, [8, 8], [8, 8])) |
| |
| outer_res = inner.is_tile_inner(layout_tile, [64, 64], [8, 8]) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, [64, 64], [8, 8]) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| case2() |
| |
| def case3(): |
| # (2,4):(1,2)x(8,8):(8,1) -> (8,2,8,4):(64,1,8,2) |
| inner = TileLayout(S[(2, 4) : (1, 2)]) |
| outer = TileLayout(S[(8, 8) : (8, 1)]) |
| layout_tile = TileLayout(S[(8, 2, 32) : (64, 1, 2)]) |
| assert_structural_equal(layout_tile, inner.tile(outer, [8, 8], [2, 4])) |
| |
| outer_res = inner.is_tile_inner(layout_tile, [16, 32], [2, 4]) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, [16, 32], [8, 8]) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, [16, 32], [8, 8]) is None |
| assert inner.is_tile_outer(layout_tile, [16, 32], [2, 4]) is None |
| |
| case3() |
| |
| def case4(): |
| # ((4,2),(2,4)):((16,8),(1,2))x(8,8):(8,1) -> (8,4,2,8,2,4):(512,16,8,64,1,2) |
| inner = TileLayout(S[(4, 2, 2, 4) : (16, 8, 1, 2)]) |
| outer = TileLayout(S[(8, 8) : (8, 1)]) |
| layout_tile = TileLayout(S[(8, 4, 2, 8, 2, 4) : (512, 16, 8, 64, 1, 2)]) |
| assert_structural_equal(layout_tile.canonicalize(), inner.tile(outer, (8, 8), (8, 8))) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (64, 64), (8, 8)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (64, 64), (8, 8)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, (64, 64), (8, 8)) is None |
| assert inner.is_tile_outer(layout_tile, (64, 64), (8, 8)) is None |
| |
| case4() |
| |
| def case5_sharded1(): |
| # Tile over a sharded layout - 1 |
| layout = TileLayout(S[(8, 1, 4, 2) : (4 @ laneid, 2, 1 @ laneid, 1)]) |
| outer = TileLayout(S[(8, 8) : (8, 1)]) |
| layout_tile = layout.tile(outer=outer, outer_shape=(8, 8), inner_shape=(8, 8)) |
| layout_expected = TileLayout(S[(8, 8, 1, 8, 4, 2) : (16, 4 @ laneid, 2, 2, 1 @ laneid, 1)]) |
| assert_structural_equal(layout_expected.canonicalize(), layout_tile) |
| |
| outer_res = layout.is_tile_inner(layout_tile, (64, 64), (8, 8)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (64, 64), (8, 8)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), layout.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, (64, 64), (8, 8)) is None |
| assert layout.is_tile_outer(layout_tile, (64, 64), (8, 8)) is None |
| |
| case5_sharded1() |
| |
| def case6_sharded2(): |
| # Tile over a sharded layout - 2 |
| inner = TileLayout(S[(8, 4) : (4 @ laneid, 1 @ laneid)]) |
| outer = TileLayout(S[(8, 8) : (8, 1)]) |
| layout_tile = inner.tile(outer=outer, outer_shape=(8, 8), inner_shape=(8, 4)) |
| layout_expected = TileLayout(S[(8, 8, 8, 4) : (8, 4 @ laneid, 1, 1 @ laneid)]) |
| assert_structural_equal(layout_expected, layout_tile) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (64, 32), (8, 4)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (64, 32), (8, 8)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, (64, 32), (8, 8)) is None |
| assert inner.is_tile_outer(layout_tile, (64, 32), (8, 4)) is None |
| |
| case6_sharded2() |
| |
| def case7_normalized4(): |
| # Normalized Tile Layout Test - 4 (tile < inner) |
| outer = TileLayout(S[(4, 2, 1) : (2, 1, 1)]) |
| inner = TileLayout(S[(2, 4, 1) : (2, 3, 1)]) |
| layout_tile = inner.tile(outer, outer_shape=(4, 2), inner_shape=(2, 4)) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (8, 8), (4, 2)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (8, 8), (2, 4)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, (8, 8), (4, 2)) is None |
| assert inner.is_tile_outer(layout_tile, (8, 8), (2, 4)) is None |
| |
| case7_normalized4() |
| |
| def case8_normalized5(): |
| # Normalized Tile Layout Test - 5 (tile = inner) |
| outer = TileLayout(S[(8, 2) : (2, 1)]) |
| inner = TileLayout(S[(2, 4) : (4, 1)]) |
| layout_tile = inner.tile(outer, (8, 2), (2, 4)) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (16, 8), (2, 4)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (16, 8), (8, 2)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| assert outer.is_tile_inner(layout_tile, (16, 8), (8, 2)) is None |
| assert inner.is_tile_outer(layout_tile, (16, 8), (2, 4)) is None |
| |
| case8_normalized5() |
| |
| def case9_normalized6(): |
| # Normalized Tile Layout Test - 6 (tile < inner) |
| outer = TileLayout(S[(8, 4, 1) : (4, 1, 4)]) |
| inner = TileLayout(S[(2, 1, 1) : (4, 3, 1)]) |
| TileLayout(S[(8, 2, 2) : (4, 2, 2)]) |
| layout_tile = inner.tile(outer, (8, 4), (2, 1)) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (16, 4), (2, 1)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| inner_res = outer.is_tile_outer(layout_tile, (16, 4), (8, 4)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), inner.canonicalize()) |
| |
| case9_normalized6() |
| |
| def case10_normalized7(): |
| # Normalized Tile Layout Test - 7 (tile = inner) |
| outer = TileLayout(S[(8, 8, 4) : (32, 4, 1)]) |
| inner = TileLayout(S[(1, 2, 1) : (4, 3, 1)]) |
| inner_tmp = TileLayout(S[(1, 2, 2) : (8, 4, 3)]) |
| layout_tile = inner.tile(outer, (8, 8, 4), (1, 2, 1)) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (8, 16, 4), (1, 2, 1)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| |
| assert inner.is_tile_inner(layout_tile.canonicalize(), (8, 16, 4), (1, 2, 1)) |
| |
| assert outer.is_tile_inner(layout_tile, (8, 16, 4), (8, 8, 4)) is None |
| assert inner_tmp.is_tile_inner(layout_tile, (8, 16, 4), (1, 2, 2)) is None |
| |
| case10_normalized7() |
| |
| def case11_normalized8(): |
| # Normalized Tile Layout Test - 8 (tile = inner w/ device) |
| outer = TileLayout(S[(8, 8, 4) : (32, 4, 1)]) |
| inner = TileLayout(S[(8, 8, 1, 4, 2) : (4, 4 @ laneid, 2, 1 @ laneid, 1)]) |
| layout_tile = inner.tile(outer, (8, 8, 4), (8, 8, 8)) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (64, 64, 32), (8, 8, 8)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| assert inner.is_tile_inner(layout_tile.canonicalize(), (64, 64, 32), (8, 8, 8)) |
| assert not outer.canonicalize().is_tile_inner( |
| layout_tile.canonicalize(), (64, 64, 32), (8, 8, 4) |
| ) |
| |
| case11_normalized8() |
| |
| def case12_normalized9(): |
| # Normalized Tile Layout Test - 9 (tile = inner w/ device + diff major-dim) |
| outer = TileLayout(S[(16, 8, 4) : (1, 64, 16)]) |
| inner = TileLayout(S[(2, 4, 2, 2) : (4, 1, 4, 3)]) |
| layout_tile = inner.tile(outer, (16, 8, 4), (8, 2, 2)) |
| |
| outer_res = inner.is_tile_inner(layout_tile, (128, 16, 8), (8, 2, 2)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), outer.canonicalize()) |
| assert inner.is_tile_inner(layout_tile.canonicalize(), (128, 16, 8), (8, 2, 2)) |
| assert not outer.canonicalize().is_tile_inner( |
| layout_tile.canonicalize(), (128, 16, 8), (16, 8, 4) |
| ) |
| |
| case12_normalized9() |
| |
| def case_dims_mismatch(): |
| with pytest.raises(Exception): |
| layout = TileLayout(S[8:1]) |
| layout2 = TileLayout(S[(2, 4) : (1, 2)]) |
| layout2.tile(layout, [8], [2, 4]) |
| |
| case_dims_mismatch() |
| |
| def case_tile_compose_layout(): |
| # tile(TileLayout, ComposeLayout) |
| compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) |
| layout = TileLayout(S[(8, 1) : (1, 1)]) |
| layout_tile = compose.tile(layout, (8, 1), (8, 64)) |
| layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[4096:1])) |
| assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) |
| |
| outer_res = compose.is_tile_inner(layout_tile, (4096,), (512,)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), layout.canonicalize()) |
| |
| inner_res = layout.is_tile_outer(layout_tile, (4096,), (8,)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), compose.canonicalize()) |
| |
| assert layout.is_tile_inner(layout_tile, (4096,), (512,)) is None |
| assert compose.is_tile_outer(layout_tile, (4096,), (8,)) is None |
| |
| case_tile_compose_layout() |
| |
| def case_tile_swizzle_layout(): |
| # swizzle_128B_atom |
| swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layout = TileLayout(S[(8, 4) : (1, 8)]) |
| layout_tile = swizzle.tile(layout, (8, 4), (8, 64)) |
| layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(64, 4, 64) : (64, 4096, 1)])) |
| assert_structural_equal(layout_tile.canonicalize(), layout_expected) |
| |
| outer_res = swizzle.is_tile_inner(layout_tile, (64, 256), (8, 64)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), layout.canonicalize()) |
| |
| inner_res = layout.is_tile_outer(layout_tile, (64, 256), (8, 4)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), swizzle.canonicalize()) |
| |
| case_tile_swizzle_layout() |
| |
| def case_tile_swizzle_layout2(): |
| # swizzle_128B_atom |
| swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| tile = TileLayout(S[(3, 8, 4) : (8 * 4, 1, 8)]) |
| layout_tile = swizzle.tile(tile, (3, 8, 4), (1, 8, 64)) |
| layout_expected = ComposeLayout( |
| swizzle.per_element, |
| swizzle.swizzle_len, |
| swizzle.atom_len, |
| TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]), |
| swizzle.swizzle_inner, |
| ) |
| assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) |
| |
| outer_res = swizzle.is_tile_inner(layout_tile, (3, 64, 256), (1, 8, 64)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), tile.canonicalize()) |
| |
| inner_res = tile.is_tile_outer(layout_tile, (3, 64, 256), (3, 8, 4)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), swizzle.canonicalize()) |
| |
| case_tile_swizzle_layout2() |
| |
| def case_tile_swizzle_layout3(): |
| # swizzle_64B_atom |
| swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) |
| tile = TileLayout(S[(8, 8) : (1, 8)]) |
| layout_tile = swizzle.tile(tile, (8, 8), (8, 32)) |
| layout_expected = ComposeLayout( |
| swizzle.per_element, |
| swizzle.swizzle_len, |
| swizzle.atom_len, |
| TileLayout(S[(64, 8, 32) : (32, 2048, 1)]), |
| swizzle.swizzle_inner, |
| ) |
| assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) |
| |
| outer_res = swizzle.is_tile_inner(layout_tile, (64, 256), (8, 32)) |
| assert outer_res is not None |
| assert_structural_equal(outer_res.canonicalize(), tile.canonicalize()) |
| |
| inner_res = tile.is_tile_outer(layout_tile, (64, 256), (8, 8)) |
| assert inner_res is not None |
| assert_structural_equal(inner_res.canonicalize(), swizzle.canonicalize()) |
| |
| case_tile_swizzle_layout3() |
| |
| def case_tile_swizzle_layout4(): |
| # swizzle_64B_atom |
| swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) |
| outer = swizzle.is_tile_inner(swizzle, (64, 256), (8, 32)) |
| assert outer is None |
| |
| outer = swizzle.is_tile_inner(swizzle, (64, 32), (8, 32)) |
| assert outer is not None |
| outer_expected = TileLayout(S[(8, 1) : (1, 0)]) |
| assert_structural_equal(outer.canonicalize(), outer_expected.canonicalize()) |
| |
| case_tile_swizzle_layout4() |
| |
| def case_tile_swizzle_layout5(): |
| # swizzle_128B_atom |
| swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) |
| tile1 = TileLayout(S[(8, 8) : (1, 8)]) |
| tile2 = TileLayout(S[(2, 2) : (1, 2)]) |
| layout_tile = swizzle.tile(tile1, (8, 8), (8, 32)) |
| layout_tile = layout_tile.tile(tile2, (2, 2), (64, 256)) |
| |
| outer = swizzle.is_tile_inner(layout_tile, (128, 512), (8, 32)) |
| assert outer is not None |
| outer_expected = tile1.tile(tile2, (2, 2), (8, 8)) |
| assert_structural_equal(outer.canonicalize(), outer_expected.canonicalize()) |
| |
| case_tile_swizzle_layout5() |
| |
| |
| def test_shard_layout(): |
| """In the current layout design, shard is just a special case of tile, where the outer tile has thread axes.""" # noqa: E501 |
| |
| def case_mma_layout(): |
| layout = TileLayout(S[(1, 2) : (2, 1)]) |
| layout_warp = TileLayout(S[(8, 4) : (4 @ laneid, 1 @ laneid)]) |
| res = layout.tile(layout_warp, [8, 4], [1, 2]) |
| layout_expected = TileLayout(S[(32, 2) : (1 @ laneid, 1)]) |
| assert_structural_equal(res.canonicalize(), layout_expected.canonicalize()) |
| |
| outer = layout.is_tile_inner(res, [8, 8], [1, 2]) |
| assert outer is not None |
| assert_structural_equal(outer.canonicalize(), layout_warp.canonicalize()) |
| |
| inner = layout_warp.is_tile_outer(res, [8, 8], [8, 4]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), layout.canonicalize()) |
| |
| case_mma_layout() |
| |
| def case_cta_layout(): |
| layout = TileLayout(S[(1, 2) : (2, 1)]) |
| layout_warp = TileLayout(S[(8, 4) : (4 @ laneid, 1 @ laneid)]) |
| layout_cta = TileLayout(S[(2, 2) : (2 @ warpid, 1 @ warpid)]) |
| |
| res_warp = layout.tile(layout_warp, [8, 4], [1, 2]) |
| res = res_warp.tile(layout_cta, [2, 2], [8, 8]) |
| layout_expected = TileLayout( |
| S[(2, 8, 2, 4, 2) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid, 1)] |
| ) |
| assert_structural_equal(res.canonicalize(), layout_expected.canonicalize()) |
| |
| outer = layout.is_tile_inner(res, [16, 16], [1, 2]) |
| outer_expected = TileLayout( |
| S[(2, 8, 2, 4) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid)] |
| ) |
| assert outer is not None |
| assert_structural_equal(outer, outer_expected) |
| |
| inner = layout_cta.is_tile_outer(res, [16, 16], [2, 2]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), res_warp.canonicalize()) |
| |
| case_cta_layout() |
| |
| def case_cta_layout2(): |
| with tvm.target.Target("cuda"): |
| tiled = TileLayout(S[(2, 8, 2, 4, 2) : (64 @ tx, 4 @ tx, 32 @ tx, 1 @ tx, 1)]) |
| # local is inner of cta |
| layout = TileLayout(S[2:1]) |
| outer = layout.is_tile_inner(tiled, [16, 16], [1, 2]) |
| assert outer is not None |
| outer_expected = TileLayout(S[(2, 8, 2, 4) : (64 @ tx, 4 @ tx, 32 @ tx, 1 @ tx)]) |
| assert_structural_equal(outer.canonicalize(), outer_expected.canonicalize()) |
| |
| layout = TileLayout(S[(2, 8, 2, 4) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid)]) |
| inner = layout.is_tile_outer(tiled, [16, 16], [16, 8]) |
| inner_expected = TileLayout(S[2:1]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), inner_expected.canonicalize()) |
| |
| # warp view is inner of cta |
| layout = TileLayout(S[(8, 1, 4, 2) : (4 @ laneid, 2, 1 @ laneid, 1)]) |
| outer = layout.is_tile_inner(tiled, [16, 16], [8, 8]) |
| assert outer is not None |
| outer_expected = TileLayout(S[(2, 2) : (2 @ warpid, 1 @ warpid)]) |
| assert_structural_equal(outer.canonicalize(), outer_expected.canonicalize()) |
| |
| layout = TileLayout(S[(2, 2) : (2 @ warpid, 1 @ warpid)]) |
| inner = layout.is_tile_outer(tiled, [16, 16], [2, 2]) |
| inner_expected = TileLayout(S[(32, 2) : (1 @ laneid, 1)]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), inner_expected.canonicalize()) |
| |
| case_cta_layout2() |
| |
| def case_quad_shuffle(): |
| layout = TileLayout(S[(1, 2) : (2, 1)]) |
| layout_warp = TileLayout(S[8 : 4 @ laneid]) |
| res = layout.tile(layout_warp, [8, 1], [1, 2]) |
| layout_expected = TileLayout(S[(8, 2) : (4 @ laneid, 1)]) |
| assert_structural_equal(res.canonicalize(), layout_expected.canonicalize()) |
| |
| outer = layout.is_tile_inner(res, [8, 2], [1, 2]) |
| assert outer is not None |
| assert_structural_equal(outer.canonicalize(), layout_warp.canonicalize()) |
| |
| inner = layout_warp.is_tile_outer(res, [8, 2], [8, 1]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), layout.canonicalize()) |
| |
| case_quad_shuffle() |
| |
| def case_replicate(): |
| layout = TileLayout(S[(64, 128) : (128, 1)]) |
| layout_rep = TileLayout(S[2 : 2 @ warpid] + R[2 : 1 @ warpid]) |
| res = layout.tile(layout_rep, [2, 1], [64, 128]) |
| layout_expected = TileLayout(S[(2, 8192) : (2 @ warpid, 1)] + R[2 : 1 @ warpid]) |
| assert_structural_equal(res.canonicalize(), layout_expected.canonicalize()) |
| |
| outer = layout.is_tile_inner(res, [128, 128], [64, 128]) |
| assert outer is not None |
| assert_structural_equal(outer.canonicalize(), layout_rep.canonicalize()) |
| |
| inner = layout_rep.is_tile_outer(res, [128, 128], [2, 1]) |
| assert inner is not None |
| assert_structural_equal(inner.canonicalize(), layout.canonicalize()) |
| |
| case_replicate() |
| |
| |
| 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 = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| assert layout.size() == 512 |
| layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) |
| assert layout.size() == 1024 |
| |
| swizzle_layout_size() |
| |
| def compose_layout_size(): |
| layout = ComposeLayout(3, 3, 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 = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| assert layout.span() == 512 |
| layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) |
| assert layout.span() == 1024 |
| |
| swizzle_layout_span() |
| |
| def compose_layout_span(): |
| layout = ComposeLayout(3, 3, 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)]) |
| assert layout.size("P") == 8 |
| assert layout.size("F") == 64 |
| assert layout.span("F") == 456 |
| |
| layout_partition = TileLayout(S[8 : 1 @ P]) |
| assert layout_partition.size("P") == 8 and layout_partition.size("F") == 1 |
| |
| layout_free = TileLayout(S[8 : 1 @ F]) |
| assert layout_free.size("P") == 1 and layout_free.size("F") == 8 |
| |
| layout = TileLayout.trainium("PF", (128, 128)) |
| assert layout.size("P") == 128 and layout.size("F") == 128 |
| |
| layout = TileLayout.trainium("FPF", (32, 512, 512)) |
| assert_structural_equal( |
| layout, TileLayout(S[(32, 4, 128, 512) : (512 @ F, (512 * 32) @ F, 1 @ P, 1 @ F)]) |
| ) |
| |
| layout = TileLayout.trainium("FPPF", (2, 4, 32, 512)) |
| assert_structural_equal( |
| layout, TileLayout(S[(2, 4, 32, 512) : (512 @ F, 32 @ P, 1 @ P, 1 @ F)]) |
| ) |
| |
| trainium_layout_tests() |
| |
| |
| def test_apply(): |
| ################ TileLayout |
| def test_tile_layout_0(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i * 8 + j)["m"] == i * 8 + j * 1 |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i, j, shape=(8, 8))["m"] == i * 8 + j * 1 |
| # # apply can accept coord larger than size |
| # for p in range(1024): |
| # outer = p // 64 |
| # inner = p % 64 |
| # i, j = inner // 8, inner % 8 |
| # assert layout.apply(p)["m"] == outer * 64 + i * 8 + j * 1 |
| with pytest.raises(Exception): |
| layout.apply(1, 1, 1) |
| |
| test_tile_layout_0() |
| |
| def test_tile_layout_1(): |
| layout = TileLayout(S[(8, 8) : (10, 1)]) |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i * 8 + j)["m"] == i * 10 + j * 1 |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i, j, shape=(8, 8))["m"] == i * 10 + j * 1 |
| |
| # # apply can accept coord larger than size |
| # for p in range(1024): |
| # outer = p // 64 |
| # inner = p % 64 |
| # i, j = inner // 8, inner % 8 |
| # assert ( |
| # layout.apply( |
| # p, |
| # )[0] |
| # == outer * 78 + i * 10 + j * 1 |
| # ) |
| |
| test_tile_layout_1() |
| |
| def test_tile_layout_2(): |
| layout = TileLayout(S[(2, 3, 4, 2, 2) : (1, 2, 12, 6, 48)]) |
| |
| def f(i0, i1): |
| leaf1 = i0 // 3 |
| leaf2 = i0 % 3 |
| leaf3 = i1 // 4 |
| leaf4 = (i1 % 4) // 2 |
| leaf5 = i1 % 2 |
| assert ( |
| layout.apply(i0, i1, shape=(6, 16))["m"] |
| == leaf1 * 1 + leaf2 * 2 + leaf3 * 12 + leaf4 * 6 + leaf5 * 48 |
| ) |
| |
| for i0, i1 in itertools.product(range(6), range(16)): |
| f(i0, i1) |
| for i in range(6 * 16): |
| f(i // 16, i % 16) |
| |
| test_tile_layout_2() |
| |
| def test_tile_layout_3(): |
| layout = TileLayout(S[(8, 1, 4, 2) : (4 @ laneid, 2, 1 @ laneid, 1)]) |
| for i0, i1 in itertools.product(range(8), range(8)): |
| res = layout.apply(i0, i1, shape=(8, 8)) |
| assert res["m"] == i1 % 2 |
| assert res["laneid"] == i0 * 4 + i1 // 2 |
| |
| test_tile_layout_3() |
| |
| def test_tile_layout_4(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| v = tvm.tirx.Var("v", ty="int32") |
| res = layout.apply(v) |
| assert res["m"] == v |
| |
| test_tile_layout_4() |
| |
| ################ Swizzle Layout |
| def test_swizzle_layout_0(): |
| layout = ComposeLayout(0, 3, 3, TileLayout(S[(64,)])) |
| # assert layout.size == 64 |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i * 8 + j)["m"] == i * 8 + i ^ j |
| |
| test_swizzle_layout_0() |
| |
| def test_swizzle_layout_1(): |
| layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| assert layout.size() == 512 |
| for i, j, k in itertools.product(range(8), range(8), range(8)): |
| assert layout.apply((i * 8 + j) * 8 + k)["m"] == (i * 8 + (i ^ j)) * 8 + k |
| # apply can accept coord larger than size |
| for p in range(4096): |
| outer = p // 512 |
| inner = p % 512 |
| i, j, k = inner // 64, (inner % 64) // 8, inner % 8 |
| assert layout.apply(p)["m"] == outer * 512 + (i * 8 + (i ^ j)) * 8 + k |
| |
| test_swizzle_layout_1() |
| |
| def test_swizzle_layout_2(): |
| layout = ComposeLayout(0, 3, 3, TileLayout(S[(64,)]), swizzle_inner=False) |
| assert layout.size() == 64 |
| for i, j in itertools.product(range(8), range(8)): |
| assert layout.apply(i * 8 + j)["m"] == (i ^ j) * 8 + j |
| |
| test_swizzle_layout_2() |
| |
| def test_swizzle_layout_3(): |
| layout = ComposeLayout(0, 2, 3, TileLayout(S[(32,)])) |
| for i, j in itertools.product(range(8), range(8)): |
| _outer_i, inner_i = i // 4, i % 4 |
| outer_j, inner_j = j // 4, j % 4 |
| assert layout.apply(i * 8 + j)["m"] == i * 8 + outer_j * 4 + (inner_i ^ inner_j) |
| |
| test_swizzle_layout_3() |
| |
| ################ Compose Layout |
| def test_compose_layout_0(): |
| layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layoutB = TileLayout(S[(8, 64) : (64, 1)]) |
| layout = ComposeLayout( |
| layoutA.per_element, |
| layoutA.swizzle_len, |
| layoutA.atom_len, |
| layoutB, |
| layoutA.swizzle_inner, |
| ) |
| assert layout.size() == 512 |
| assert layout.span() == 512 |
| for i, j in itertools.product(range(8), range(64)): |
| assert ( |
| layout.apply(i * 64 + j)["m"] == layoutA.apply(layoutB.apply(i * 64 + j)["m"])["m"] |
| ) |
| |
| test_compose_layout_0() |
| |
| def test_compose_layout_1(): |
| layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layoutB = TileLayout(S[(16, 64, 8) : (64, 1, 1024)]) |
| layout = ComposeLayout( |
| layoutA.per_element, |
| layoutA.swizzle_len, |
| layoutA.atom_len, |
| layoutB, |
| layoutA.swizzle_inner, |
| ) |
| assert layout.size() == 16 * 64 * 8 |
| assert layout.span() == 16 * 64 * 8 |
| for i, j, k in itertools.product(range(16), range(64), range(8)): |
| assert ( |
| layout.apply(i * 64 * 8 + j * 8 + k)["m"] |
| == layoutA.apply(layoutB.apply(i * 64 * 8 + j * 8 + k)["m"])["m"] |
| ) |
| |
| test_compose_layout_1() |
| |
| ################ Trainium Layout |
| def test_trainium_layout_0(): |
| layout = TileLayout(S[(8, 8) : (8 @ F, 1 @ P)]) |
| for i, j in itertools.product(range(8), range(8)): |
| coord = layout.apply(i, j, shape=(8, 8)) |
| assert coord["P"] == j |
| assert coord["F"] == i * 8 |
| |
| test_trainium_layout_0() |
| |
| def test_trainium_layout_1(): |
| layout = TileLayout(S[(2, 6, 4, 2, 2) : (1 @ F, 1 @ P, 12 @ F, 6 @ P, 48 @ F)]) |
| |
| def f(i0, i1): |
| leaf1 = i0 // 6 |
| leaf2 = i0 % 6 |
| leaf3 = i1 // 4 |
| leaf4 = (i1 % 4) // 2 |
| leaf5 = i1 % 2 |
| coord = layout.apply(i0, i1, shape=(12, 16)) |
| assert coord["P"] == leaf2 + leaf4 * 6 |
| assert coord["F"] == leaf1 * 1 + leaf3 * 12 + leaf5 * 48 |
| |
| for i0, i1 in itertools.product(range(6), range(16)): |
| f(i0, i1) |
| for i in range(6 * 16): |
| f(i // 16, i % 16) |
| |
| test_trainium_layout_1() |
| |
| ################ Trainium PSUM Layout |
| def test_trainium_psum_layout_0(): |
| layout = TileLayout(S[(1024, 8) : (1 @ F, 1 @ P)]).to_psum() |
| for i, j in itertools.product(range(1024), range(8)): |
| coord = layout.apply(i, j, shape=(1024, 8)) |
| assert coord["Bank"] == i // 512 |
| assert coord["P"] == j |
| assert coord["F"] == i % 512 |
| |
| test_trainium_psum_layout_0() |
| |
| |
| def test_normalize_compose_layout(): |
| def case1(): |
| layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layoutB = TileLayout(S[(8, 64) : (64, 1)]) |
| layout = ComposeLayout( |
| layoutA.per_element, |
| layoutA.swizzle_len, |
| layoutA.atom_len, |
| layoutB.canonicalize(), |
| layoutA.swizzle_inner, |
| ) |
| assert_structural_equal(layout.canonicalize(), layoutA) |
| |
| case1() |
| |
| def case2(): |
| layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layoutB = TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) |
| layout = ComposeLayout( |
| layoutA.per_element, |
| layoutA.swizzle_len, |
| layoutA.atom_len, |
| layoutB.canonicalize(), |
| layoutA.swizzle_inner, |
| ) |
| assert_structural_equal(layout.canonicalize(), layout) |
| |
| case2() |
| |
| |
| def test_normalize_trainium_layout(): |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8 @ P, 1 @ F)]) |
| assert_structural_equal(layout, layout.canonicalize()) |
| |
| case1() |
| |
| def case2(): |
| layout = TileLayout(S[(8, 1, 8) : (8 @ F, 1 @ P, 1 @ F)]) |
| layout_expected = TileLayout(S[64 : 1 @ F]) |
| assert_structural_equal(layout_expected, layout.canonicalize()) |
| |
| case2() |
| |
| def case3(): |
| layout = TileLayout(S[(8, 8, 8) : (8 @ F, 1 @ P, 1 @ F)]) |
| assert_structural_equal(layout, layout.canonicalize()) |
| |
| case3() |
| |
| |
| def test_direct_sum(): |
| def case1(): |
| # Example from the appendix: A + B yields contiguous (16):(1) |
| # B = (2,2):(4,1), A = (2,2):(8,2) |
| B = TileLayout(S[(2, 2) : (4, 1)]) |
| A = TileLayout(S[(2, 2) : (8, 2)]) |
| |
| # Compute direct sum on tiling domain S_A ⊗ S_B with shapes (2,2) and (2,2) |
| sum_layout = B.direct_sum(A, [2, 2], [2, 2]).canonicalize() |
| expected = TileLayout(S[16:1]) |
| assert_structural_equal(expected, sum_layout) |
| |
| # Verify Apply equality: 8p + 2q + 4i + j |
| print(f"sum_layout: {sum_layout}") |
| an = Analyzer() |
| for p in [0, 1]: |
| for q in [0, 1]: |
| for i in [0, 1]: |
| for j in [0, 1]: |
| m = sum_layout.apply(p, q, i, j, shape=(2, 2, 2, 2))["m"] |
| m_left = A.apply(p, i, shape=(2, 2))["m"] |
| m_right = B.apply(q, j, shape=(2, 2))["m"] |
| assert an.can_prove(m == m_left + m_right) |
| |
| # Recognition: recover A given B and sum, and recover B given A and sum |
| interleaved_shape = [2, 2, 2, 2] # [A0, B0, A1, B1] |
| A_rec = B.is_direct_sum_right(sum_layout, interleaved_shape, [2, 2]) |
| assert A_rec is not None |
| assert_structural_equal(A.canonicalize(), A_rec.canonicalize()) |
| |
| B_rec = A.is_direct_sum_left(sum_layout, interleaved_shape, [2, 2]) |
| assert B_rec is not None |
| assert_structural_equal(B.canonicalize(), B_rec.canonicalize()) |
| |
| case1() |
| |
| |
| def test_group_by_logical_shape(): |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| layout = layout.tile(layout, outer_shape=[8, 8], inner_shape=[8, 8]) |
| outer, seps = layout.group([64, 64]) |
| assert_structural_equal(outer, layout) |
| assert seps[0] == 0 |
| assert seps[1] == 2 |
| assert seps[2] == 4 |
| |
| case1() |
| |
| |
| def test_group_many_uses_minimal_common_refinement(): |
| layout = TileLayout(S[30:1]) |
| grouped, separators = layout.group_many(([6, 5], [2, 15])) |
| |
| expected = TileLayout(S[(2, 3, 5) : (15, 5, 1)]) |
| assert_structural_equal(grouped, expected) |
| assert [list(seps) for seps in separators] == [[0, 2, 3], [0, 1, 3]] |
| |
| |
| def test_group_many_preserves_repeated_boundaries_as_unit_iters(): |
| layout = TileLayout(S[30:1]) |
| grouped, separators = layout.group_many( |
| ( |
| [1, 6, 1, 5, 1], |
| [2, 1, 3, 5], |
| ) |
| ) |
| |
| expected = TileLayout(S[(1, 2, 1, 3, 1, 5, 1) : (30, 15, 15, 5, 5, 1, 1)]) |
| assert_structural_equal(grouped, expected) |
| assert [list(seps) for seps in separators] == [ |
| [0, 1, 4, 5, 6, 7], |
| [0, 2, 3, 4, 7], |
| ] |
| |
| |
| def test_group_many_preserves_existing_unit_iter_stride(): |
| layout = TileLayout(S[(6, 1, 5) : (5, 97, 1)]) |
| grouped, separators = layout.group_many(([6, 1, 5], [30])) |
| |
| assert_structural_equal(grouped, layout) |
| assert [list(seps) for seps in separators] == [[0, 1, 2, 3], [0, 3]] |
| |
| |
| def test_group_many_rejects_incompatible_boundaries(): |
| layout = TileLayout(S[12:1]) |
| with pytest.raises(tvm.error.InternalError, match="incompatible cumulative boundaries"): |
| layout.group_many(([3, 4], [2, 6])) |
| |
| |
| def test_group_many_accepts_only_provably_ordered_symbolic_boundaries(): |
| n = Var("n", "int32") |
| extent = tvm.tirx.floormod(n, 16) + 2 |
| layout = TileLayout(S[2 * extent : 1]) |
| |
| grouped, separators = layout.group_many(([2, extent], [1, 2, extent])) |
| expected = TileLayout(S[(1, 2, extent) : (Analyzer().simplify(2 * extent), extent, 1)]) |
| assert_structural_equal(grouped, expected) |
| assert [list(seps) for seps in separators] == [[0, 2, 3], [0, 1, 2, 3]] |
| |
| with pytest.raises( |
| tvm.error.InternalError, |
| match="cannot prove cumulative boundary divisibility|order or equality", |
| ): |
| layout.group_many(([2, extent], [extent, 2])) |
| |
| |
| def test_permute_by_groups(): |
| def case_swap_two_groups(): |
| # Two groups, each with 2 shard iters: swap them. |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| layout = layout.tile(layout, outer_shape=[8, 8], inner_shape=[8, 8]) |
| grouped, seps = layout.group([64, 64]) |
| # seps == [0, 2, 4] |
| permuted = grouped.permute_by_groups(seps, [1, 0]) |
| # Expected: shard reordered as [g1[0], g1[1], g0[0], g0[1]] |
| expected = grouped.permute_dims([2, 3, 0, 1]) |
| assert_structural_equal(permuted, expected) |
| |
| def case_identity(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| layout = layout.tile(layout, outer_shape=[8, 8], inner_shape=[8, 8]) |
| grouped, seps = layout.group([64, 64]) |
| permuted = grouped.permute_by_groups(seps, [0, 1]) |
| assert_structural_equal(permuted, grouped) |
| |
| def case_invalid_perm(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| layout = layout.tile(layout, outer_shape=[8, 8], inner_shape=[8, 8]) |
| grouped, seps = layout.group([64, 64]) |
| with pytest.raises(AssertionError): |
| grouped.permute_by_groups(seps, [0, 0]) |
| |
| case_swap_two_groups() |
| case_identity() |
| case_invalid_perm() |
| |
| |
| def test_tile_to(): |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| tiled = layout.tile_to([64, 64], [8, 8]) |
| tiled_expected = layout.tile(layout, [8, 8], [8, 8]) |
| assert_structural_equal(tiled, tiled_expected) |
| |
| case1() |
| |
| |
| def test_mma_shared_layout(): |
| def case1(): |
| layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (64, 256)) |
| layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(64, 4, 64) : (64, 4096, 1)])) |
| assert_structural_equal(layout, layout_expected) |
| |
| case1() |
| |
| def case2(): |
| layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (3, 64, 256)) |
| layout_expected = ComposeLayout( |
| 3, 3, 3, TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]) |
| ) |
| assert_structural_equal(layout, layout_expected) |
| |
| case2() |
| |
| def case3(): |
| layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_64B_ATOM, (3, 64, 256)) |
| layout_expected = ComposeLayout( |
| 3, 2, 3, TileLayout(S[(3, 64, 8, 32) : (16384, 32, 2048, 1)]) |
| ) |
| assert_structural_equal(layout, layout_expected) |
| |
| case3() |
| |
| |
| def test_pool_allocator_alloc_mma(): |
| def alloc_layout(shape, dtype, swizzle_mode="auto"): |
| with IRBuilder(): |
| with Tx_builder.prim_func(): |
| pool = T.SMEMPool(Var("smem_ptr", PointerType(PrimType("uint8")))) |
| buf = pool.alloc_tcgen05_mma_AB(shape, dtype, swizzle_mode=swizzle_mode) |
| return buf.layout |
| |
| cases = [ |
| ("uint8", (3, 64, 256)), |
| ("float16", (3, 64, 256)), |
| ("bfloat16", (3, 64, 256)), |
| ("float32", (3, 64, 256)), |
| ("float4_e2m1fn", (3, 64, 256)), |
| ] |
| for dtype, shape in cases: |
| layout = alloc_layout(shape, dtype) |
| expected = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, shape) |
| assert_structural_equal(layout, expected) |
| |
| shape = (3, 64, 256) |
| layout_64b = alloc_layout(shape, "float32", SwizzleMode.SWIZZLE_64B_ATOM) |
| expected_64b = mma_shared_layout("float32", SwizzleMode.SWIZZLE_64B_ATOM, shape) |
| assert_structural_equal(layout_64b, expected_64b) |
| |
| layout_none = alloc_layout(shape, "float16", "none") |
| expected_none = mma_shared_layout("float16", SwizzleMode.SWIZZLE_NONE, shape) |
| assert_structural_equal(layout_none, expected_none) |
| |
| |
| def test_tmem_mma_operand_layout_grouped_d(): |
| head64_o = tmem_mma_operand_layout( |
| "D", (64, 512), "float32", M=64, cta_group=1, ws=True, group=(2, 2, 128) |
| ) |
| expected_head64_o = TileLayout( |
| S[(64, 2, 2, 128) : (1 @ TLane, 128 @ TCol, 64 @ TLane, 1 @ TCol)] |
| ).canonicalize() |
| assert_structural_equal(head64_o, expected_head64_o) |
| |
| head128_o = tmem_mma_operand_layout( |
| "D", (64, 512), "float32", M=128, cta_group=2, group=(2, 2, 128) |
| ) |
| assert_structural_equal(head128_o, expected_head64_o) |
| |
| with pytest.raises(ValueError, match="ws=True only valid"): |
| tmem_mma_operand_layout("D", (128, 256), "float32", M=128, cta_group=1, ws=True) |
| |
| with pytest.raises(ValueError, match="group must be"): |
| tmem_mma_operand_layout( |
| "D", (64, 512), "float32", M=64, cta_group=1, ws=True, group=(4, 128) |
| ) |
| |
| with pytest.raises(ValueError, match="identity needs 2D"): |
| tmem_mma_operand_layout("A", (1, 128, 64), "bfloat16", M=128, cta_group=1) |
| |
| # M=64 non-.ws A occupies lanes 0..63 identically (Layout F datapath is the |
| # accumulator's scatter, not the A operand's). |
| a_f = tmem_mma_operand_layout("A", (64, 64), "float32", M=64, cta_group=1) |
| assert_structural_equal(a_f, TileLayout(S[(64, 64) : (1 @ TLane, 1 @ TCol)]).canonicalize()) |
| |
| |
| def test_storage(): |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| assert_structural_equal(layout.storage(), layout) |
| |
| case1() |
| |
| def case2(): |
| layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)]) |
| layout_stroage = TileLayout(S[2:1]) |
| assert_structural_equal(layout.storage(), layout_stroage) |
| |
| case2() |
| |
| def case3(): |
| layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| assert_structural_equal(layout.storage(), layout) |
| |
| case3() |
| |
| def case4(): |
| layout = ( |
| TileLayout(S[2:1]) |
| .tile(TileLayout(S[(8, 4) : (4 @ laneid, 1 @ laneid)]), (8, 4), (1, 2)) |
| .tile(TileLayout(S[(2, 1) : (1, 2)]), (2, 1), (8, 8)) |
| .tile(TileLayout(S[(1, 8) : (8, 1)]), (1, 8), (16, 8)) |
| ) |
| layout_stroage = ( |
| TileLayout(S[2:1]) |
| .tile(TileLayout(S[(2, 1) : (1, 2)]), (2, 1), (1, 2)) |
| .tile(TileLayout(S[(1, 8) : (8, 1)]), (1, 8), (2, 2)) |
| ) |
| assert_structural_equal(layout.storage().canonicalize(), layout_stroage.canonicalize()) |
| |
| case4() |
| |
| |
| def test_unpack(): |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| layout_expected = TileLayout(S[(8, 16) : (16, 1)]) |
| assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case1() |
| |
| def case2(): |
| layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| layout_expected = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) |
| assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case2() |
| |
| def case3(): |
| layout = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) |
| layout_expected = ComposeLayout(4, 3, 3, TileLayout(S[(8, 128) : (128, 1)])) |
| assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case3() |
| |
| |
| def test_pack(): |
| def case1(): |
| layout = TileLayout(S[(8, 16) : (16, 1)]) |
| layout_expected = TileLayout(S[(8, 8) : (8, 1)]) |
| assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case1() |
| |
| def case2(): |
| layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) |
| layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case2() |
| |
| def case3(): |
| layout = ComposeLayout(4, 3, 3, TileLayout(S[(8, 128) : (128, 1)])) |
| layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) |
| assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) |
| |
| case3() |
| |
| |
| def test_slice(): |
| def verify_slice(layout, shape, region, sliced): |
| r_shape = [r[1] - r[0] for r in region] |
| r_size = functools.reduce(operator.mul, [r[1] - r[0] for r in region]) |
| |
| def get_region_coord(u): |
| coord = [] |
| for r in reversed(region): |
| coord.append(u % (r[1] - r[0])) |
| u //= r[1] - r[0] |
| return coord[::-1] |
| |
| def get_shape_coord(r_coord, region): |
| return [region[i][0] + r_coord[i] for i in range(len(region))] |
| |
| analyzer = Analyzer() |
| |
| for u in range(r_size): |
| r_coord = get_region_coord(u) |
| s_coord = get_shape_coord(r_coord, region) |
| a = layout.apply(*s_coord, shape=shape)["m"] |
| b = sliced.apply(*r_coord, shape=r_shape)["m"] |
| assert analyzer.simplify(a == b) |
| |
| def case1(): |
| layout = TileLayout(S[(8, 8) : (8, 1)]) |
| shape = [64] |
| region = [(5, 8)] |
| sliced = layout.slice(shape, region).canonicalize() |
| assert sliced is not None |
| verify_slice(layout, shape, region, sliced) |
| |
| region = [tvm.ir.Range(5, 8)] |
| sliced_2 = layout.slice(shape, region).canonicalize() |
| assert sliced_2 is not None |
| assert_structural_equal(sliced, sliced_2) |
| |
| case1() |
| |
| def case2(): |
| # Choose begin and extent to satisfy midpoint condition |
| layout = TileLayout(S[(4, 4, 4, 4) : (64, 4, 16, 1)]) |
| shape = [16, 16] |
| region = [(2, 3), (6, 10)] |
| sliced = layout.slice(shape, region).canonicalize() |
| assert sliced is not None |
| verify_slice(layout, shape, region, sliced) |
| |
| case2() |
| |
| def case3(): |
| layout = TileLayout(S[(2, 8, 3, 8) : (192, 8, 64, 1)]) |
| shape = [16, 24] |
| region = [(2, 6), (4, 12)] |
| sliced = layout.slice(shape, region).canonicalize() |
| assert sliced is not None |
| verify_slice(layout, shape, region, sliced) |
| |
| case3() |
| |
| def case4(): |
| layout = TileLayout(S[(128, 2, 64) : (64, 128 * 64, 1)]) |
| shape = [128, 128] |
| region = [(0, 128), (32, 96)] |
| sliced = layout.slice(shape, region).canonicalize() |
| assert sliced is not None |
| verify_slice(layout, shape, region, sliced) |
| |
| case4() |
| |
| def case_adds_region_offset_to_base(): |
| warp_offset = Var("warp_offset", "int32") |
| layout = TileLayout(S[(4, 64)] + warp_offset * 64) |
| shape = [4, 64] |
| region = [(2, 4), (0, 64)] |
| sliced = layout.slice(shape, region).canonicalize() |
| assert sliced is not None |
| assert Analyzer().can_prove_equal(sliced.offset[m], warp_offset * 64 + 128) |
| |
| case_adds_region_offset_to_base() |
| |
| def case_swizzle_slice(): |
| # bare-swizzle ComposeLayout slice - delegates to ComposeLayout |
| swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) |
| shape = [512] |
| region = [(64, 128)] |
| sliced = swizzle.slice(shape, region) |
| assert sliced is not None |
| verify_slice(swizzle, shape, region, sliced) |
| |
| case_swizzle_slice() |
| |
| def case_compose_slice(): |
| # ComposeLayout slice |
| compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) |
| shape = [512] |
| region = [(64, 128)] |
| sliced = compose.slice(shape, region) |
| assert sliced is not None |
| verify_slice(compose, shape, region, sliced) |
| |
| case_compose_slice() |
| |
| def case_compose_slice_2d(): |
| # ComposeLayout slice with 2D shape |
| compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) |
| shape = [8, 64] |
| region = [(2, 4), (0, 64)] |
| sliced = compose.slice(shape, region) |
| assert sliced is not None |
| verify_slice(compose, shape, region, sliced) |
| |
| case_compose_slice_2d() |
| |
| |
| def test_cuda_copy_extract_swizzle_tile_simplifies_constant_region_extents(): |
| from tvm.tirx.cuda.tile_primitive.copy._common import _extract_tile |
| |
| zero = tvm.tirx.Mul(tvm.tirx.IntImm("int32", 0), tvm.tirx.IntImm("int32", 64)) |
| end = tvm.tirx.Add(zero, tvm.tirx.IntImm("int32", 64)) |
| |
| tile = _extract_tile(ComposeLayout(3, 3, 3, TileLayout(S[(512,)])), [(zero, end)]) |
| |
| assert [int(it.extent) for it in tile.shard] == [64] |
| |
| |
| def _compose_swizzle(value, per_element, swizzle_len, atom_len, swizzle_inner=True): |
| chunk = 1 << per_element |
| index = value // chunk |
| mask = (1 << swizzle_len) - 1 |
| if swizzle_inner: |
| index ^= (index & (mask << atom_len)) >> atom_len |
| else: |
| index ^= (index & mask) << atom_len |
| return index * chunk + value % chunk |
| |
| |
| def _evaluate_layout_expr(expr, values): |
| node_type = type(expr).__name__ |
| if node_type == "IntImm": |
| return int(expr.value) |
| if node_type == "Var": |
| return values[expr] |
| if node_type == "Let": |
| bindings = dict(values) |
| bindings[expr.var] = _evaluate_layout_expr(expr.value, bindings) |
| return _evaluate_layout_expr(expr.body, bindings) |
| if node_type in ("Add", "Sub", "Mul", "FloorDiv", "FloorMod"): |
| lhs = _evaluate_layout_expr(expr.a, values) |
| rhs = _evaluate_layout_expr(expr.b, values) |
| if node_type == "Add": |
| return lhs + rhs |
| if node_type == "Sub": |
| return lhs - rhs |
| if node_type == "Mul": |
| return lhs * rhs |
| if node_type == "FloorDiv": |
| return lhs // rhs |
| return lhs % rhs |
| if node_type == "Cast": |
| return _evaluate_layout_expr(expr.value, values) |
| if node_type == "Call": |
| args = [_evaluate_layout_expr(arg, values) for arg in expr.args] |
| op_name = str(expr.op.name) |
| if op_name == "tirx.bitwise_xor": |
| return args[0] ^ args[1] |
| if op_name == "tirx.bitwise_and": |
| return args[0] & args[1] |
| if op_name == "tirx.shift_left": |
| return args[0] << args[1] |
| if op_name == "tirx.shift_right": |
| return args[0] >> args[1] |
| raise AssertionError(f"Cannot evaluate call {op_name}") |
| raise AssertionError(f"Cannot evaluate node type {node_type}") |
| |
| |
| def _naive_compose_expr(tile, coords, per_element=3, swizzle_len=3, atom_len=3, inner=True): |
| period = 1 << (per_element + swizzle_len + atom_len) |
| swizzle = ComposeLayout( |
| per_element, |
| swizzle_len, |
| atom_len, |
| TileLayout(S[(period,)]), |
| inner, |
| ) |
| return swizzle.apply(tile.apply(*coords)["m"])["m"] |
| |
| |
| @pytest.mark.parametrize("swizzle_len", [3, 2, 1, 0]) |
| def test_compose_structured_apply_uses_bounded_normal_form(swizzle_len): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| tile = TileLayout(S[(8, 8) : (64, 1)] + 72) |
| layout = ComposeLayout(3, swizzle_len, 3, tile) |
| |
| actual = layout.apply(x, y)["m"] |
| shaped = layout.apply(x, y, shape=(8, 8))["m"] |
| if swizzle_len == 0: |
| expected = x * 64 + y + 72 |
| else: |
| mask = (1 << swizzle_len) - 1 |
| base = x * 64 + (y ^ (((x + 1) & mask) << 3)) + 64 |
| expected = base ^ 8 |
| assert_structural_equal(actual, expected) |
| assert_structural_equal(shaped, expected) |
| |
| for xv, yv in itertools.product(range(8), repeat=2): |
| logical = xv * 64 + yv + 72 |
| assert _evaluate_layout_expr(actual, {x: xv, y: yv}) == _compose_swizzle( |
| logical, 3, swizzle_len, 3 |
| ) |
| |
| |
| def test_compose_rank_one_shape_aware_apply_is_structured(): |
| flat = Var("flat", ty="int32") |
| tile = TileLayout(S[(8, 8) : (64, 1)] + 72) |
| layout = ComposeLayout(3, 3, 3, tile) |
| |
| actual = layout.apply(flat, shape=(64,))["m"] |
| x = flat // 8 |
| y = flat % 8 |
| expected = (x * 64 + (y ^ (((x + 1) & 7) << 3)) + 64) ^ 8 |
| assert_structural_equal(actual, expected) |
| assert type(actual).__name__ != "Let" |
| for value in range(64): |
| logical = value // 8 * 64 + value % 8 + 72 |
| assert _evaluate_layout_expr(actual, {flat: value}) == _compose_swizzle(logical, 3, 3, 3) |
| |
| |
| def test_compose_scalar_apply_keeps_full_swizzle_semantics(): |
| coord = Var("coord", ty="int32") |
| tile = TileLayout(S[(8, 8) : (64, 1)] + 72) |
| layout = ComposeLayout(3, 3, 3, tile) |
| actual = layout.apply(coord)["m"] |
| |
| assert "//" in str(actual) |
| assert "%" in str(actual) |
| for value in range(256): |
| logical = value // 8 * 64 + value % 8 + 72 |
| assert _evaluate_layout_expr(actual, {coord: value}) == _compose_swizzle(logical, 3, 3, 3) |
| |
| |
| def test_compose_shape_aware_apply_decomposes_each_logical_coordinate(): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| tile = TileLayout(S[(2, 4, 8, 8) : (256, 64, 8, 1)]) |
| layout = ComposeLayout(3, 3, 3, tile) |
| |
| actual = layout.apply(x, y, shape=(8, 64))["m"] |
| expected = x * 64 + (y ^ ((x & 7) << 3)) |
| assert_structural_equal(actual, expected) |
| for xv, yv in itertools.product(range(8), range(64)): |
| assert _evaluate_layout_expr(actual, {x: xv, y: yv}) == _compose_swizzle( |
| xv * 64 + yv, 3, 3, 3 |
| ) |
| |
| |
| def test_compose_structured_apply_splits_one_cross_atom_term(): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| tile = TileLayout(S[(128, 8) : (8, 1)]) |
| layout = ComposeLayout(3, 3, 3, tile) |
| |
| actual = layout.apply(x, y)["m"] |
| high = x // 8 * 64 |
| low = x % 8 * 8 + y |
| expected = high + (low ^ (((x // 8) & 7) << 3)) |
| assert_structural_equal(actual, expected) |
| assert type(actual).__name__ != "Let" |
| for xv, yv in itertools.product(range(128), range(8)): |
| logical = xv * 8 + yv |
| assert _evaluate_layout_expr(actual, {x: xv, y: yv}) == _compose_swizzle(logical, 3, 3, 3) |
| |
| |
| def test_compose_structured_apply_handles_uint32_casts_per_term(): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| x_u32 = tvm.tirx.Cast("uint32", x) |
| y_u32 = tvm.tirx.Cast("uint32", y) |
| layout = ComposeLayout(3, 3, 3, TileLayout(S[(8, 8) : (64, 1)] + 72)) |
| |
| actual = layout.apply(x_u32, y_u32, shape=(8, 8))["m"] |
| high = x_u32 * 64 + 64 |
| expected = (high + (y_u32 ^ (((x_u32 + 1) & 7) << 3))) ^ 8 |
| assert_structural_equal(actual, expected) |
| assert "// T.uint32(64)" not in str(actual) |
| |
| for xv, yv in itertools.product(range(8), repeat=2): |
| logical = xv * 64 + yv + 72 |
| assert _evaluate_layout_expr(actual, {x: xv, y: yv}) == _compose_swizzle(logical, 3, 3, 3) |
| |
| |
| def test_compose_structured_apply_classifies_dynamic_offsets(): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| high = Var("high", ty="int32") |
| low = Var("low", ty="int32") |
| tile = TileLayout(S[(8, 8) : (64, 1)] + high * 64 + low % 8 + 16) |
| layout = ComposeLayout(3, 3, 3, tile) |
| |
| actual = layout.apply(x, y)["m"] |
| expected = (x * 64 + high * 64 + ((y + low % 8) ^ (((x + high) & 7) << 3))) ^ 16 |
| assert_structural_equal(actual, expected) |
| for xv, yv, hv, lv in itertools.product(range(2), range(8), range(2), range(8)): |
| logical = xv * 64 + yv + hv * 64 + lv % 8 + 16 |
| assert _evaluate_layout_expr(actual, {x: xv, y: yv, high: hv, low: lv}) == ( |
| _compose_swizzle(logical, 3, 3, 3) |
| ) |
| |
| |
| def test_compose_extent_one_coordinate_is_zero_before_classification(): |
| unused = Var("unused", ty="int32") |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| tile = TileLayout(S[(1, 8, 8) : (-1, 64, 1)] + 8) |
| layout = ComposeLayout(3, 3, 3, tile) |
| |
| actual = layout.apply(unused, x, y)["m"] |
| expected = (x * 64 + (y ^ ((x & 7) << 3))) ^ 8 |
| assert_structural_equal(actual, expected) |
| assert "unused" not in str(actual) |
| |
| |
| @pytest.mark.parametrize( |
| "case", |
| ["outer_swizzle", "negative_stride", "unbounded_low", "carry", "symbolic_stride"], |
| ) |
| def test_compose_structured_apply_falls_back_to_full_swizzle(case): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| value = Var("value", ty="int32") |
| inner = True |
| if case == "outer_swizzle": |
| tile = TileLayout(S[(8, 8) : (64, 1)] + 8) |
| inner = False |
| elif case == "negative_stride": |
| tile = TileLayout(S[(8, 8) : (-64, 1)] + 8) |
| elif case == "unbounded_low": |
| tile = TileLayout(S[(8, 8) : (64, 1)] + value) |
| elif case == "carry": |
| tile = TileLayout(S[(8, 8) : (64, 1)] + 4) |
| else: |
| tile = TileLayout(S[(8, 8) : (value, 1)] + 8) |
| |
| layout = ComposeLayout(3, 3, 3, tile, inner) |
| actual = layout.apply(x, y)["m"] |
| expected = _naive_compose_expr(tile, (x, y), inner=inner) |
| |
| assert type(actual).__name__ == "Let" |
| assert_structural_equal(actual.value, tile.apply(x, y)["m"]) |
| quotient = actual.body |
| assert type(quotient).__name__ == "Let" |
| assert_structural_equal(quotient.value, actual.var // 8) |
| body_vars = tvm.tirx.analysis.undefined_vars(quotient.body) |
| assert all(var.same_as(actual.var) or var.same_as(quotient.var) for var in body_vars) |
| |
| for xv, yv, vv in itertools.product(range(3), range(8), [-3, 0, 5, 64]): |
| values = {x: xv, y: yv, value: vv} |
| assert _evaluate_layout_expr(actual, values) == _evaluate_layout_expr(expected, values) |
| |
| |
| def test_compose_structured_apply_emits_shared_base_address_chain(): |
| x = Var("x", ty="int32") |
| y = Var("y", ty="int32") |
| constants = [0, 8, 16, 24, 512, 520, 528, 536] |
| addresses = [ |
| ComposeLayout(3, 3, 3, TileLayout(S[(8, 8) : (64, 1)] + constant)).apply(x, y)["m"] |
| for constant in constants |
| ] |
| |
| base = addresses[0] |
| high_base = base + 512 |
| expected = [ |
| base, |
| base ^ 8, |
| base ^ 16, |
| base ^ 24, |
| high_base, |
| high_base ^ 8, |
| high_base ^ 16, |
| high_base ^ 24, |
| ] |
| for actual, expected_expr in zip(addresses, expected): |
| assert_structural_equal(actual, expected_expr) |
| |
| for xv, yv in itertools.product(range(8), repeat=2): |
| for constant, address in zip(constants, addresses): |
| logical = xv * 64 + yv + constant |
| assert _evaluate_layout_expr(address, {x: xv, y: yv}) == _compose_swizzle( |
| logical, 3, 3, 3 |
| ) |
| |
| |
| def test_apply_to_shape(): |
| """``apply_to_shape`` should give per-shard coord, preferring per-dim |
| split when the input shape aligns with the layout's grouping.""" |
| |
| from tvm.tirx.layout import Iter, TileLayout |
| |
| # 1 shard per dim — coord[d] passes through unchanged. |
| lay = TileLayout(S[16, 16]) |
| assert [int(x) for x in lay.apply_to_shape([5, 7], [16, 16])] == [5, 7] |
| |
| # Dim 1 split into (4, 4) factors — per-dim mixed-radix within dim 1, |
| # no cross-dim flatten needed. |
| lay2 = TileLayout.from_iters([Iter(16, 16, "m"), Iter(4, 4, "m"), Iter(4, 1, "m")]) |
| assert [int(x) for x in lay2.apply_to_shape([5, 7], [16, 16])] == [5, 7 // 4, 7 % 4] |
| |
| # Both dims split — verifies split stays local to each dim. |
| lay3 = TileLayout.from_iters( |
| [Iter(4, 64, "m"), Iter(4, 16, "m"), Iter(4, 4, "m"), Iter(4, 1, "m")] |
| ) |
| r = lay3.apply_to_shape([13, 9], [16, 16]) |
| assert [int(x) for x in r] == [13 // 4, 13 % 4, 9 // 4, 9 % 4] |
| |
| |
| def test_slice_single_shard_skips_defensive_floormod(): |
| """Regression: ``Layout.slice`` must not emit ``floormod(begin, Ek)`` on |
| single-shard groups whose caller-contract guarantees ``begin + extent |
| <= Ek``. |
| |
| Background: ``SlicePerGroup`` in ``src/tirx/ir/layout/tile_slice.cc`` |
| decomposes ``begin`` into per-shard coordinates via |
| ``floormod(floordiv(begin, B[k]), Ek)``. When ``m == 1`` (single shard |
| in the group) and ``begin`` is a runtime expression (e.g. a pipeline |
| stage ``BufferLoad``), the analyzer cannot prove ``begin < Ek`` so the |
| defensive ``floormod`` survives codegen. |
| |
| Concretely, fa4's K_smem with shape ``(SMEM_PIPE_DEPTH_KV=3, 128, 128)`` |
| sliced by ``[stage:stage+1, :, :]`` would emit |
| ``floormod(stage, 3) * 16384`` in every per-MMA SMEM-descriptor offset |
| (72 sites at s1024_kv4) — even though ``PipelineState`` already keeps |
| ``stage`` in ``[0, 3)``. |
| |
| The fix relies on the existing single-shard caller contract noted in |
| the function: |
| ``the slice is valid as long as the caller guarantees |
| begin + slice_extent <= extent (which is assumed)`` |
| |
| With the contract the mod is provably a no-op; this test asserts the |
| sliced layout's ``offset`` is the bare ``stage * stride`` form for |
| runtime ``begin``. |
| """ |
| # Single-shard outer-axis slice with a runtime stage variable. |
| layout = TileLayout(S[(3, 128, 128) : (16384, 128, 1)]) |
| shape = [3, 128, 128] |
| stage = Var("stage", "int32") |
| region = [tvm.ir.Range(stage, stage + 1), tvm.ir.Range(0, 128), tvm.ir.Range(0, 128)] |
| sliced = layout.slice(shape, region) |
| assert sliced is not None |
| offset_strs = [str(off) for _, off in sliced.offset.items()] |
| full = " | ".join(offset_strs) |
| # No defensive floormod-by-extent should remain on the stage axis. |
| assert "FloorMod" not in full and "floormod" not in full and "% 3" not in full, ( |
| f"single-shard slice with runtime begin must not emit defensive floormod, got offset={full}" |
| ) |
| |
| # Multi-shard groups (e.g. row dim with swizzle interleaving |
| # ``(128, 2):(64, 8192)``) still need the floormod for correct |
| # decomposition; verify we did not over-aggressively strip it. |
| multi_shard = TileLayout.from_iters( |
| [Iter(2, 8192, "m"), Iter(128, 64, "m")] # outer (extent=2), inner (extent=128) |
| ) |
| multi_shape = [256] |
| multi_region = [tvm.ir.Range(96, 96 + 32)] |
| multi_sliced = multi_shard.slice(multi_shape, multi_region) |
| assert multi_sliced is not None |
| # Constants — analyzer simplifies floormod(96, 128) to 96 internally; |
| # we just assert offset is non-empty and structurally sane (not None). |
| |
| |
| def test_slice_tcgen05_frag_layout_scope_consistent(): |
| """Slicing a wid_in_wg+laneid frag layout (tcgen05 16x256b) must stay |
| scope-consistent: the sliced result canonicalizes to a single tid_in_wg |
| chain over the full 128 threads (regression for the per-group-fusion bug). |
| """ |
| frag = TileLayout( |
| S[(4, 2, 2, 8, 4, 4, 2) : (1 @ wid_in_wg, 16, 2, 4 @ laneid, 4, 1 @ laneid, 1)] |
| ) |
| |
| def thread_chain(layout): |
| canon = layout.canonicalize() |
| names = {it.axis.name for it in canon.shard if it.axis.is_thread()} |
| titers = sorted( |
| ((int(it.stride), int(it.extent)) for it in canon.shard if it.axis.is_thread()), |
| ) |
| running = 1 |
| for stride, extent in titers: |
| assert stride == running, f"non-contiguous thread chain: {titers}" |
| running *= extent |
| return names, running |
| |
| with tvm.target.Target("cuda"): |
| # Full-region slice and a column sub-slice must both canonicalize to a |
| # single tid_in_wg chain covering all 128 warpgroup threads. |
| full = frag.slice([128, 32], [(0, 128), (0, 32)]) |
| names, total = thread_chain(full) |
| assert names == {"tid_in_wg"}, names |
| assert total == 128, total |
| |
| col = frag.slice([128, 32], [(0, 128), (16, 32)]) |
| names_c, total_c = thread_chain(col) |
| assert names_c == {"tid_in_wg"}, names_c |
| assert total_c == 128, total_c |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
