| 20 | |
| 21 | # |
| 22 | class GemmOperation: |
| 23 | # |
| 24 | def __init__( |
| 25 | self, |
| 26 | gemm_kind, |
| 27 | arch, |
| 28 | tile_description, |
| 29 | A, |
| 30 | B, |
| 31 | C, |
| 32 | element_epilogue, |
| 33 | epilogue_functor=EpilogueFunctor.LinearCombination, |
| 34 | swizzling_functor=SwizzlingFunctor.Identity8, |
| 35 | required_cuda_ver_major=9, |
| 36 | required_cuda_ver_minor=2, |
| 37 | ): |
| 38 | |
| 39 | self.operation_kind = OperationKind.Gemm |
| 40 | self.arch = arch |
| 41 | self.tile_description = tile_description |
| 42 | self.gemm_kind = gemm_kind |
| 43 | self.A = A |
| 44 | self.B = B |
| 45 | self.C = C |
| 46 | self.element_epilogue = element_epilogue |
| 47 | self.epilogue_functor = epilogue_functor |
| 48 | self.swizzling_functor = swizzling_functor |
| 49 | self.required_cuda_ver_major = required_cuda_ver_major |
| 50 | self.required_cuda_ver_minor = required_cuda_ver_minor |
| 51 | |
| 52 | # |
| 53 | def is_complex(self): |
| 54 | complex_operators = [ |
| 55 | MathOperation.multiply_add_complex, |
| 56 | MathOperation.multiply_add_complex_gaussian, |
| 57 | ] |
| 58 | return ( |
| 59 | self.tile_description.math_instruction.math_operation in complex_operators |
| 60 | ) |
| 61 | |
| 62 | # |
| 63 | def is_split_k_parallel(self): |
| 64 | return self.gemm_kind == GemmKind.SplitKParallel |
| 65 | |
| 66 | # |
| 67 | def is_planar_complex(self): |
| 68 | return self.gemm_kind in (GemmKind.PlanarComplex, GemmKind.PlanarComplexArray) |
| 69 | |
| 70 | # |
| 71 | def accumulator_type(self): |
| 72 | accum = self.tile_description.math_instruction.element_accumulator |
| 73 | |
| 74 | if self.is_complex(): |
| 75 | return get_complex_from_real(accum) |
| 76 | |
| 77 | return accum |
| 78 | |
| 79 | # |
no outgoing calls
no test coverage detected