# 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.operator.tile_primitive.tma_utils import ( SwizzleMode, mma_shared_layout, tma_shared_layout, ) from tvm.tirx.layout import ( Axis, ComposeLayout, F, Iter, P, R, S, SwizzleLayout, TileLayout, laneid, m, tid_in_wg, 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( layout_A=SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), layout_B=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( SwizzleLayout(3, 3, 3, swizzle_inner=True), 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layout = TileLayout(S[(8, 4) : (1, 8)]) layout_tile = swizzle.tile(layout, (8, 4), (8, 64)) layout_expected = ComposeLayout( SwizzleLayout(3, 3, 3, swizzle_inner=True), 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) 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, TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]) ) 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 = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) tile = TileLayout(S[(8, 8) : (1, 8)]) layout_tile = swizzle.tile(tile, (8, 8), (8, 32)) layout_expected = ComposeLayout(swizzle, TileLayout(S[(64, 8, 32) : (32, 2048, 1)])) 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 = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) 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 = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) 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 = 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)]) 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", dtype="int32") res = layout.apply(v) assert res["m"] == v test_tile_layout_4() ################ Swizzle Layout def test_swizzle_layout_0(): layout = SwizzleLayout(per_element=0, swizzle_len=3, atom_len=3) # 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) 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 = SwizzleLayout(per_element=0, swizzle_len=3, atom_len=3, 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 = SwizzleLayout(per_element=0, swizzle_len=2, atom_len=3) 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layoutB = TileLayout(S[(8, 64) : (64, 1)]) layout = ComposeLayout(layoutA, layoutB) 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layoutB = TileLayout(S[(16, 64, 8) : (64, 1, 1024)]) layout = ComposeLayout(layoutA, layoutB) 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layoutB = TileLayout(S[(8, 64) : (64, 1)]) layout = ComposeLayout(layoutA, layoutB.canonicalize()) assert_structural_equal(layout.canonicalize(), layoutA) case1() def case2(): layoutA = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layoutB = TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) layout = ComposeLayout(layoutA, layoutB.canonicalize()) 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_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( SwizzleLayout(3, 3, 3, swizzle_inner=True), 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( SwizzleLayout(3, 3, 3, swizzle_inner=True), 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( SwizzleLayout(3, 2, 3, swizzle_inner=True), TileLayout(S[(3, 64, 8, 32) : (16384, 32, 2048, 1)]), ) assert_structural_equal(layout, layout_expected) case3() def test_tma_shared_layout_alias(): shape = (3, 64, 256) layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, shape) alias_layout = tma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, shape) assert_structural_equal(alias_layout, layout) 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_mma(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_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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) 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 = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) layout_expected = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) case2() def case3(): layout = ComposeLayout( SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), TileLayout(S[(8, 64) : (64, 1)]), ) layout_expected = ComposeLayout( SwizzleLayout(per_element=4, swizzle_len=3, atom_len=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 = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) layout_expected = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) case2() def case3(): layout = ComposeLayout( SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3), TileLayout(S[(8, 128) : (128, 1)]), ) layout_expected = ComposeLayout( SwizzleLayout(per_element=3, swizzle_len=3, atom_len=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_swizzle_slice(): # SwizzleLayout slice - delegates to ComposeLayout swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) 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( SwizzleLayout(per_element=3, swizzle_len=3, atom_len=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( SwizzleLayout(per_element=3, swizzle_len=3, atom_len=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_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()