| 1383 | |
| 1384 | |
| 1385 | def test_mma_shared_layout(): |
| 1386 | def case1(): |
| 1387 | layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (64, 256)) |
| 1388 | layout_expected = ComposeLayout( |
| 1389 | SwizzleLayout(3, 3, 3, swizzle_inner=True), TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) |
| 1390 | ) |
| 1391 | assert_structural_equal(layout, layout_expected) |
| 1392 | |
| 1393 | case1() |
| 1394 | |
| 1395 | def case2(): |
| 1396 | layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (3, 64, 256)) |
| 1397 | layout_expected = ComposeLayout( |
| 1398 | SwizzleLayout(3, 3, 3, swizzle_inner=True), |
| 1399 | TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]), |
| 1400 | ) |
| 1401 | assert_structural_equal(layout, layout_expected) |
| 1402 | |
| 1403 | case2() |
| 1404 | |
| 1405 | def case3(): |
| 1406 | layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_64B_ATOM, (3, 64, 256)) |
| 1407 | layout_expected = ComposeLayout( |
| 1408 | SwizzleLayout(3, 2, 3, swizzle_inner=True), |
| 1409 | TileLayout(S[(3, 64, 8, 32) : (16384, 32, 2048, 1)]), |
| 1410 | ) |
| 1411 | assert_structural_equal(layout, layout_expected) |
| 1412 | |
| 1413 | case3() |
| 1414 | |
| 1415 | |
| 1416 | def test_tma_shared_layout_alias(): |