()
| 1032 | compose_layout_span() |
| 1033 | |
| 1034 | def trainium_layout_tests(): |
| 1035 | # TrainiumLayout tests |
| 1036 | layout = TileLayout(S[(8, 8) : (1 @ P, 1 @ F)]) |
| 1037 | assert layout.size("P") == 8 |
| 1038 | assert layout.size("F") == 8 |
| 1039 | |
| 1040 | layout = TileLayout(S[(8, 8, 8) : (64 @ F, 1 @ P, 1 @ F)]) |
| 1041 | assert layout.size("P") == 8 |
| 1042 | assert layout.size("F") == 64 |
| 1043 | assert layout.span("F") == 456 |
| 1044 | |
| 1045 | layout_partition = TileLayout(S[8 : 1 @ P]) |
| 1046 | assert layout_partition.size("P") == 8 and layout_partition.size("F") == 1 |
| 1047 | |
| 1048 | layout_free = TileLayout(S[8 : 1 @ F]) |
| 1049 | assert layout_free.size("P") == 1 and layout_free.size("F") == 8 |
| 1050 | |
| 1051 | layout = TileLayout.trainium("PF", (128, 128)) |
| 1052 | assert layout.size("P") == 128 and layout.size("F") == 128 |
| 1053 | |
| 1054 | layout = TileLayout.trainium("FPF", (32, 512, 512)) |
| 1055 | assert_structural_equal( |
| 1056 | layout, TileLayout(S[(32, 4, 128, 512) : (512 @ F, (512 * 32) @ F, 1 @ P, 1 @ F)]) |
| 1057 | ) |
| 1058 | |
| 1059 | layout = TileLayout.trainium("FPPF", (2, 4, 32, 512)) |
| 1060 | assert_structural_equal( |
| 1061 | layout, TileLayout(S[(2, 4, 32, 512) : (512 @ F, 32 @ P, 1 @ P, 1 @ F)]) |
| 1062 | ) |
| 1063 | |
| 1064 | trainium_layout_tests() |
| 1065 |
no test coverage detected
searching dependent graphs…