(
A_dtype: DType,
B_dtype: DType,
C_dtype: DType,
MNK_tiling: HyperCube,
A_major: GmmaMajor,
B_major: GmmaMajor,
)
| 105 | |
| 106 | |
| 107 | def gmma_selector( |
| 108 | A_dtype: DType, |
| 109 | B_dtype: DType, |
| 110 | C_dtype: DType, |
| 111 | MNK_tiling: HyperCube, |
| 112 | A_major: GmmaMajor, |
| 113 | B_major: GmmaMajor, |
| 114 | ): |
| 115 | assert MNK_tiling.ndim == 3 |
| 116 | M_tile = MNK_tiling[0] |
| 117 | N_tile = MNK_tiling[1] |
| 118 | K_tile = MNK_tiling[2] |
| 119 | assert ( |
| 120 | not MNK_tiling.has_dynamic() |
| 121 | ), "Tiling info shouldn't have dynamic at GMMA level." |
| 122 | assert M_tile % 64 == 0, "GMMA tiling M dimension should be multiple of 64." |
| 123 | A_transpose = False if A_major == GmmaMajor.MajorK else True |
| 124 | B_transpose = False if B_major == GmmaMajor.MajorK else True |
| 125 | |
| 126 | # FP32 accumulator |
| 127 | if C_dtype.is_same(float_t()): |
| 128 | # FP16 inputs |
| 129 | if A_dtype.is_same(half_t()): |
| 130 | assert A_dtype.is_same(B_dtype), "A and B should have the same dtype." |
| 131 | assert K_tile % 16 == 0, "GMMA tiling K dimension should be multiple of 16." |
| 132 | assert N_tile % 8 == 0, "GMMA tiling N dimension should be multiple of 8." |
| 133 | switch_table = [ |
| 134 | ( |
| 135 | k, |
| 136 | SM90_SS( |
| 137 | M_tile=64, |
| 138 | N_tile=k, |
| 139 | K_tile=16, |
| 140 | A_dtype=half_t(), |
| 141 | B_dtype=half_t(), |
| 142 | Accum_dtype=float_t(), |
| 143 | A_transpose=A_transpose, |
| 144 | B_transpose=B_transpose, |
| 145 | ), |
| 146 | ) |
| 147 | for k in [256, 192, 128, 96, 64, 32, 16, 8] |
| 148 | ] |
| 149 | for k, v in switch_table: |
| 150 | if N_tile % k == 0: |
| 151 | return v |
| 152 | raise RuntimeError( |
| 153 | f"Selector can't find proper GMMA operations for:\n" |
| 154 | f"A_dtype: {A_dtype}, B_dtype: {B_dtype}, C_dtype: {C_dtype}\n" |
| 155 | f"MNK_tiling: {MNK_tiling}, A_major: {A_major}, B_major: {B_major}\n" |
| 156 | ) |
| 157 | |
| 158 | if __name__ == "__main__": |
| 159 | wgmma_op = SM90_SS( |
nothing calls this directly
no test coverage detected