tilelang.cuda.language.gemm_op
==============================

.. py:module:: tilelang.cuda.language.gemm_op

.. autoapi-nested-parse::

   CUDA dialect of ``T.gemm``: the common GEMM plus CUDA-specific knobs.



Functions
---------

.. autoapisummary::

   tilelang.cuda.language.gemm_op.gemm
   tilelang.cuda.language.gemm_op.gemm_sp


Module Contents
---------------

.. py:function:: gemm(A, B, C, transpose_A = False, transpose_B = False, policy = GemmWarpPolicy.Square, clear_accum = False, mbar = None, annotations = None)

   TileLang GEMM operator for CUDA.

   Same semantics as the common :func:`tilelang.language.gemm_op.gemm`: the
   default synchronous GEMM. On Hopper, if the compiler selects WGMMA
   lowering, TileLang inserts the corresponding wait implicitly. On Blackwell
   TCGEN5MMA, TileLang inserts the corresponding
   ``mbarrier_wait_parity(...)`` implicitly after issue.

   For manual asynchronous scheduling, use ``T.wgmma_gemm(...)`` with
   ``T.wait_wgmma(...)`` on Hopper, or ``T.tcgen05_gemm(...)`` with
   ``T.mbarrier_wait_parity(...)`` on Blackwell.

   :param A: Input buffer A.
   :type A: BufferLikeType, i.e. Buffer | BufferLoad | BufferRegion, or Var
   :param B: Input buffer B.
   :type B: BufferLikeType
   :param C: Output buffer C.
   :type C: BufferLikeType
   :param transpose_A: Whether to transpose A. Defaults to False.
   :type transpose_A: bool
   :param transpose_B: Whether to transpose B. Defaults to False.
   :type transpose_B: bool
   :param policy: GEMM warp partition policy.
   :type policy: GemmWarpPolicy
   :param clear_accum: Whether to clear the accumulator.
   :type clear_accum: bool
   :param mbar: Mbarrier in Blackwell.
                Required when this GEMM lowers to TCGEN5MMA. Defaults to None.
   :type mbar: BarrierType, i.e. Buffer | BufferLoad, or Var, optional
   :param annotations: Additional annotations.
   :type annotations: Optional[dict]

   :returns: A handle to the GEMM operation.
   :rtype: tirx.Call


.. py:function:: gemm_sp(A_sparse, E, B, C, transpose_A = False, transpose_E = False, transpose_B = False, policy = GemmWarpPolicy.Square, clear_accum = False, wg_wait = 0, annotations = None)

   Sparse GEMM (2:4 structured sparsity) for CUDA.

   Same semantics as the common :func:`tilelang.language.experimental.gemm_sp_op.gemm_sp`.
   ``wg_wait`` is the Hopper warpgroup wait count consumed when the WGMMA SP
   lowering is selected (``-1`` defers the wait to an explicit
   ``T.wait_wgmma``); it rides in the tile-op annotations.

   :param A_sparse: Compressed sparse matrix containing only non-zero elements.
   :param E: Metadata tensor encoding the sparsity pattern of A.
   :param B: Dense input matrix.
   :param C: Output accumulator matrix.
   :param transpose_A: Whether to transpose A. Defaults to False.
   :param transpose_E: Whether to transpose E. Defaults to False.
   :param transpose_B: Whether to transpose B. Defaults to False.
   :param policy: Warp partition policy. Defaults to GemmWarpPolicy.Square.
   :param clear_accum: Whether to zero the accumulator before computation. Defaults to False.
   :param wg_wait: Warp group wait count. Defaults to 0.
   :param annotations: Additional annotations; values in it take precedence.

   :returns: A handle to the sparse GEMM operation.
   :rtype: tirx.Call