Argument wrapper for GEMM in CUTLASS 2 or 3. It returns either 2x arguments or 3x arguments depending on the `arch` field specified in `operation`. :param operation: the GEMM operation to take the argument :type operation: :class:`cutlass_cppgen.backend.GemmOperationUniversal` |
(operation, problem_size, A, B, C, D, gemm_mode=GemmUniversalMode.Gemm, **kwargs)
| 628 | |
| 629 | |
| 630 | def GemmArguments(operation, problem_size, A, B, C, D, gemm_mode=GemmUniversalMode.Gemm, **kwargs): |
| 631 | """ |
| 632 | Argument wrapper for GEMM in CUTLASS 2 or 3. It returns either 2x arguments |
| 633 | or 3x arguments depending on the `arch` field specified in `operation`. |
| 634 | |
| 635 | :param operation: the GEMM operation to take the argument |
| 636 | :type operation: :class:`cutlass_cppgen.backend.GemmOperationUniversal` | |
| 637 | :class:`cutlass_cppgen.backend.GemmOperationGrouped` |
| 638 | |
| 639 | :param problem_size: GEMM problem size gemm(M, N, K) |
| 640 | :type operation: :class:`cutlass_cppgen.shape.GemmCoord` |
| 641 | |
| 642 | :param A: tensor A |
| 643 | :type A: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray |
| 644 | |
| 645 | :param B: tensor B |
| 646 | :type B: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray |
| 647 | |
| 648 | :param C: tensor C |
| 649 | :type C: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray |
| 650 | |
| 651 | :param D: tensor D |
| 652 | :type D: cuda.CUdeviceptr | numpy.ndarray | torch.Tensor | cupy.ndarray |
| 653 | |
| 654 | :param gemm_mode: GEMM mode |
| 655 | :type gemm_mode: :class:`cutlass_library.GemmUniversalMode` |
| 656 | |
| 657 | :param output_op: output operator, optional |
| 658 | :type output_op: :class:`cutlass_cppgen.backend.LinearCombinationFunctorArguments` |
| 659 | """ |
| 660 | if operation.swizzling_functor == SwizzlingFunctor.StreamK: |
| 661 | if operation.api == ApiVersion.v3x: |
| 662 | raise Exception("Stream K is currently only supported in CUTLASS 2.x") |
| 663 | ArgClass = GemmArguments2xStreamK |
| 664 | else: |
| 665 | ArgClass = GemmArguments3x if operation.api == ApiVersion.v3x else GemmArguments2x |
| 666 | return ArgClass(operation, problem_size, A, B, C, D, gemm_mode, **kwargs) |
| 667 | |
| 668 | |
| 669 | class GemmGroupedArguments: |