(
tile,
data_type,
layout_a,
layout_b,
layout_c,
min_cc,
align_a=32,
align_b=32,
align_c=32,
required_cuda_ver_major=9,
required_cuda_ver_minor=2,
)
| 322 | |
| 323 | # |
| 324 | def GeneratesGemm( |
| 325 | tile, |
| 326 | data_type, |
| 327 | layout_a, |
| 328 | layout_b, |
| 329 | layout_c, |
| 330 | min_cc, |
| 331 | align_a=32, |
| 332 | align_b=32, |
| 333 | align_c=32, |
| 334 | required_cuda_ver_major=9, |
| 335 | required_cuda_ver_minor=2, |
| 336 | ): |
| 337 | operations = [] |
| 338 | swizzling_functor = SwizzlingFunctor.Identity1 |
| 339 | |
| 340 | element_a, element_b, element_c, element_epilogue = data_type |
| 341 | |
| 342 | if tile.math_instruction.element_accumulator == DataType.s32: |
| 343 | epilogues = [EpilogueFunctor.LinearCombinationClamp] |
| 344 | else: |
| 345 | assert ( |
| 346 | tile.math_instruction.element_accumulator == DataType.f32 |
| 347 | or tile.math_instruction.element_accumulator == DataType.f16 |
| 348 | ) |
| 349 | epilogues = [EpilogueFunctor.LinearCombination] |
| 350 | |
| 351 | for epilogue in epilogues: |
| 352 | A = TensorDescription( |
| 353 | element_a, layout_a, int(align_a // DataTypeSize[element_a]) |
| 354 | ) |
| 355 | B = TensorDescription( |
| 356 | element_b, layout_b, int(align_b // DataTypeSize[element_b]) |
| 357 | ) |
| 358 | C = TensorDescription( |
| 359 | element_c, layout_c, int(align_c // DataTypeSize[element_c]) |
| 360 | ) |
| 361 | operations.append( |
| 362 | GemmOperation( |
| 363 | GemmKind.Gemm, |
| 364 | min_cc, |
| 365 | tile, |
| 366 | A, |
| 367 | B, |
| 368 | C, |
| 369 | element_epilogue, |
| 370 | epilogue, |
| 371 | swizzling_functor, |
| 372 | required_cuda_ver_major, |
| 373 | required_cuda_ver_minor, |
| 374 | ) |
| 375 | ) |
| 376 | operations.append( |
| 377 | GemmOperation( |
| 378 | GemmKind.SplitKParallel, |
| 379 | min_cc, |
| 380 | tile, |
| 381 | A, |
no test coverage detected