MCPcopy Create free account
hub / github.com/NVIDIA/cutlass / run

Method run

python/cutlass_cppgen/op/gemm.py:632–725  ·  view source on GitHub ↗

Runs the kernel currently specified. If it has not already been, the kernel is emitted and compiled. Tensors holding operands and outputs of the kernel are sourced either from the ``A``, ``B``, ``C``, ``D``, ``alpha``, and ``beta`` parameters provided in this call, o

(self, A=None, B=None, C=None, D=None,
            alpha=None, beta=None, sync: bool = True, print_module: bool = False, visitor_args: dict = None,
            stream: Optional[cuda.CUstream] = None)

Source from the content-addressed store, hash-verified

630 f'layout of ({ref_type}, {ref_layout}) and transpose failed.')
631
632 def run(self, A=None, B=None, C=None, D=None,
633 alpha=None, beta=None, sync: bool = True, print_module: bool = False, visitor_args: dict = None,
634 stream: Optional[cuda.CUstream] = None) -> GemmArguments:
635 """
636 Runs the kernel currently specified. If it has not already been, the kernel is emitted and
637 compiled. Tensors holding operands and outputs of the kernel are sourced either from the
638 ``A``, ``B``, ``C``, ``D``, ``alpha``, and ``beta``
639 parameters provided in this call, or from those
640 passed in on the construction of this object -- one of the two must be specified.
641
642 By default, this call returns only once the kernel has completed. To launch the kernel
643 and immediately return, set ``sync=False``. In this case, it is the responsibility of the
644 caller to syncrhonize the results of the kernel before attempting to access outputs
645 by calling ``sync()`` on the arguments returned from this call.
646
647 :param A: tensor representing data type and layout of operand A
648 :param B: tensor representing data type and layout of operand B
649 :param C: tensor representing data type and layout of operand C
650 :param D: tensor representing data type and layout of operand D
651 :param alpha: scalar paramter alpha from GEMM computation that scales the product of operands A and B
652 :param beta: scalar parameter beta from GEMM operation that scales operand C
653 :param sync: whether the call should wait for the kernel to complete before returning
654 :type sync: bool
655 :param print_module: whether to print the emitted C++ code
656 :type print_module: bool
657 :param stream: cuda stream, defaults to cuda.cuda.CUstream(0)
658 :type stream: :class:`cuda.cuda.CUstream`
659
660 :return: arguments passed in to the kernel
661 :rtype: cutlass_cppgen.backend.GemmArguments
662 """
663 if not stream:
664 stream = cuda.CUstream(0)
665 super().run_setup()
666 A = self._verify_tensor(A, self.A, self._element_a, self._layout_a, "A")
667 B = self._verify_tensor(B, self.B, self._element_b, self._layout_b, "B")
668 C = self._verify_tensor(C, self.C, self._element_c, self._layout_c, "C")
669 D = self._verify_tensor(D, self.D, self._element_d, self._layout_d, "D")
670 alpha = self._verify_scalar(alpha, self.alpha, self._element_c, "alpha")
671 beta = self._verify_scalar(beta, self.beta, self._element_c, "beta")
672
673 is_void_c = self._element_c == DataType.void
674
675 self._verify_rank(A)
676 self._verify_rank(B)
677 if not is_void_c:
678 self._verify_rank(C)
679 self._verify_rank(D)
680
681 alignment_a = self.possible_operations.find_alignment(A.shape, self._layout_a, operand="A")
682 alignment_b = self.possible_operations.find_alignment(B.shape, self._layout_b, operand="B")
683
684 # Set C alignment based on D.shape so as to correctly get an alignment with void-C
685 # kernels, for which `C` is None.
686 alignment_c = self.possible_operations.find_alignment(D.shape, self._layout_c, operand="C")
687 self.compile(self._tile_description, alignment_A=alignment_a, alignment_B=alignment_b,
688 alignment_C=alignment_c, print_module=print_module)
689

Callers 6

run_batchedMethod · 0.95
nvcc_versionFunction · 0.45
__init__Method · 0.45
__call__Method · 0.45
compileMethod · 0.45

Calls 10

_verify_rankMethod · 0.95
compileMethod · 0.95
_get_problem_argsMethod · 0.95
_get_batch_strideMethod · 0.95
GemmArgumentsFunction · 0.90
run_setupMethod · 0.80
_verify_tensorMethod · 0.80
_verify_scalarMethod · 0.80
find_alignmentMethod · 0.80
syncMethod · 0.45

Tested by

no test coverage detected