| 90 | |
| 91 | class MMA_Traits: |
| 92 | def __init__(self, mma_op: MMA_OP) -> None: |
| 93 | self.mma_op = mma_op |
| 94 | self.D_dtype = mma_op.Accum_dtype |
| 95 | self.A_dtype = mma_op.A_dtype |
| 96 | self.B_dtype = mma_op.B_dtype |
| 97 | self.C_dtype = mma_op.Accum_dtype |
| 98 | |
| 99 | A_major = GmmaMajor.MajorK if not mma_op.A_transpose else GmmaMajor.MajorMN |
| 100 | B_major = GmmaMajor.MajorK if not mma_op.B_transpose else GmmaMajor.MajorMN |
| 101 | self.A_frag_type = SmemDesc(A_major) |
| 102 | self.B_frag_type = SmemDesc(B_major) |
| 103 | |
| 104 | self.MNK_shape = HyperCube(3, [mma_op.M_tile, mma_op.N_tile, mma_op.K_tile]) |
| 105 | |
| 106 | |
| 107 | def gmma_selector( |