Update docs

Author: github-actions[bot]

_sources/autoapi/tilelang/index.rst.txt

@@ -27,6 +27,7 @@ Submodules
    /autoapi/tilelang/profiler/index
    /autoapi/tilelang/quantize/index
    /autoapi/tilelang/testing/index
+   /autoapi/tilelang/tileop/index
    /autoapi/tilelang/tools/index
    /autoapi/tilelang/transform/index
    /autoapi/tilelang/utils/index

_sources/autoapi/tilelang/intrinsics/mma_layout/index.rst.txt

@@ -9,27 +9,49 @@ Attributes
 
 .. autoapisummary::
 
-   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout
-   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_32x4_layout_sr_a
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_32x4_layout_sr_b
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_32x4_layout_rs_a
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_32x4_layout_rs_b
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_sr_a
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_sr_b
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_rs_a
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_rs_b
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_layout_sr_a
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_layout_sr_b
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_layout_rs_a
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_layout_rs_b
 
 
 Functions
 ---------
 
 .. autoapisummary::
 
+   tilelang.intrinsics.mma_layout.ldmatrix_32x4_to_shared_16x8_layout_a
+   tilelang.intrinsics.mma_layout.ldmatrix_32x4_to_shared_16x8_layout_b
    tilelang.intrinsics.mma_layout.ldmatrix_32x8_to_shared_16x16_layout
    tilelang.intrinsics.mma_layout.ldmatrix_trans_32x8_to_shared_16x16_layout
-   tilelang.intrinsics.mma_layout.ldmatrix_16x32_to_shared_16x32_layout_a
-   tilelang.intrinsics.mma_layout.ldmatrix_16x32_to_shared_16x32_layout_b
    tilelang.intrinsics.mma_layout.ldmatrix_32x16_to_shared_16x32_layout_a
    tilelang.intrinsics.mma_layout.ldmatrix_32x16_to_shared_16x32_layout_b
    tilelang.intrinsics.mma_layout.mma_store_32x8_to_shared_16x16_layout
-   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_sr
-   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_layout_rs
-   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_layout
-   tilelang.intrinsics.mma_layout.shared_32x16_to_mma_32x16_layout
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_a_32x4_layout
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_a_32x4_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_b_32x4_layout
+   tilelang.intrinsics.mma_layout.shared_16x8_to_mma_b_32x4_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_a_32x8_layout
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_a_32x8_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_b_32x8_layout
+   tilelang.intrinsics.mma_layout.shared_16x16_to_mma_b_32x8_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_a_32x16_layout
+   tilelang.intrinsics.mma_layout.shared_32x16_to_mma_a_32x16_layout_trans
+   tilelang.intrinsics.mma_layout.shared_16x32_to_mma_b_32x16_layout
+   tilelang.intrinsics.mma_layout.shared_32x16_to_mma_b_32x16_layout_trans
    tilelang.intrinsics.mma_layout.mma_32x8_to_shared_16x16_layout
+   tilelang.intrinsics.mma_layout.mma_load_a_32x4_to_shared_16x8_layout
+   tilelang.intrinsics.mma_layout.mma_load_b_32x4_to_shared_16x8_layout
+   tilelang.intrinsics.mma_layout.mma_load_a_32x16_to_shared_16x32_layout
+   tilelang.intrinsics.mma_layout.mma_load_b_32x16_to_shared_16x32_layout
    tilelang.intrinsics.mma_layout.shared_16x16_to_mma_32x8_smoothlayout
    tilelang.intrinsics.mma_layout.shared_16x32_to_mma_32x16_smoothlayout
    tilelang.intrinsics.mma_layout.shared_32x16_to_mma_32x16_smoothlayout
@@ -40,34 +62,78 @@ Functions
 Module Contents
 ---------------
 
-.. py:function:: ldmatrix_32x8_to_shared_16x16_layout(thread_id, local_id)
+.. py:function:: ldmatrix_32x4_to_shared_16x8_layout_a(thread_id, local_id)
 
-.. py:function:: ldmatrix_trans_32x8_to_shared_16x16_layout(thread_id, local_id)
+.. py:function:: ldmatrix_32x4_to_shared_16x8_layout_b(thread_id, local_id)
 
-.. py:function:: ldmatrix_16x32_to_shared_16x32_layout_a(thread_id, local_id)
+.. py:function:: ldmatrix_32x8_to_shared_16x16_layout(thread_id, local_id)
 
-.. py:function:: ldmatrix_16x32_to_shared_16x32_layout_b(thread_id, local_id)
+.. py:function:: ldmatrix_trans_32x8_to_shared_16x16_layout(thread_id, local_id)
 
 .. py:function:: ldmatrix_32x16_to_shared_16x32_layout_a(thread_id, local_id)
 
 .. py:function:: ldmatrix_32x16_to_shared_16x32_layout_b(thread_id, local_id)
 
 .. py:function:: mma_store_32x8_to_shared_16x16_layout(thread_id, local_id)
 
-.. py:function:: shared_16x16_to_mma_32x8_layout_sr(i, j)
+.. py:function:: shared_16x8_to_mma_a_32x4_layout(i, j)
+
+.. py:function:: shared_16x8_to_mma_a_32x4_layout_trans(i, j)
+
+.. py:function:: shared_16x8_to_mma_b_32x4_layout(i, j)
+
+.. py:function:: shared_16x8_to_mma_b_32x4_layout_trans(i, j)
+
+.. py:data:: shared_16x8_to_mma_32x4_layout_sr_a
+
+.. py:data:: shared_16x8_to_mma_32x4_layout_sr_b
+
+.. py:data:: shared_16x8_to_mma_32x4_layout_rs_a
+
+.. py:data:: shared_16x8_to_mma_32x4_layout_rs_b
+
+.. py:function:: shared_16x16_to_mma_a_32x8_layout(i, j)
+
+.. py:function:: shared_16x16_to_mma_a_32x8_layout_trans(i, j)
 
-.. py:function:: shared_16x16_to_mma_32x8_layout_rs(i, j)
+.. py:function:: shared_16x16_to_mma_b_32x8_layout(i, j)
 
-.. py:data:: shared_16x16_to_mma_32x8_layout
+.. py:function:: shared_16x16_to_mma_b_32x8_layout_trans(i, j)
 
-.. py:data:: shared_16x16_to_mma_32x8_layout_trans
+.. py:data:: shared_16x16_to_mma_32x8_layout_sr_a
 
-.. py:function:: shared_16x32_to_mma_32x16_layout(i, j)
+.. py:data:: shared_16x16_to_mma_32x8_layout_sr_b
 
-.. py:function:: shared_32x16_to_mma_32x16_layout(i, j)
+.. py:data:: shared_16x16_to_mma_32x8_layout_rs_a
+
+.. py:data:: shared_16x16_to_mma_32x8_layout_rs_b
+
+.. py:function:: shared_16x32_to_mma_a_32x16_layout(i, j)
+
+.. py:function:: shared_32x16_to_mma_a_32x16_layout_trans(i, j)
+
+.. py:function:: shared_16x32_to_mma_b_32x16_layout(i, j)
+
+.. py:function:: shared_32x16_to_mma_b_32x16_layout_trans(i, j)
+
+.. py:data:: shared_16x32_to_mma_32x16_layout_sr_a
+
+.. py:data:: shared_16x32_to_mma_32x16_layout_sr_b
+
+.. py:data:: shared_16x32_to_mma_32x16_layout_rs_a
+
+.. py:data:: shared_16x32_to_mma_32x16_layout_rs_b
 
 .. py:function:: mma_32x8_to_shared_16x16_layout(thread_id, local_id)
 
+.. py:function:: mma_load_a_32x4_to_shared_16x8_layout(thread_id, local_id)
+
+.. py:function:: mma_load_b_32x4_to_shared_16x8_layout(thread_id, local_id)
+
+.. py:function:: mma_load_a_32x16_to_shared_16x32_layout(thread_id, local_id)
+
+.. py:function:: mma_load_b_32x16_to_shared_16x32_layout(thread_id, local_id)
+
 .. py:function:: shared_16x16_to_mma_32x8_smoothlayout(i, j)
 
 .. py:function:: shared_16x32_to_mma_32x16_smoothlayout(i, j)

_sources/autoapi/tilelang/intrinsics/mma_macro_generator/index.rst.txt

@@ -28,7 +28,7 @@ Module Contents
 
 .. py:data:: lift
 
-.. py:class:: TensorCoreIntrinEmitter(a_dtype = 'float16', b_dtype = 'float16', accum_dtype = 'float16', a_transposed = False, b_transposed = False, block_row_warps = 2, block_col_warps = 2, warp_row_tiles = 8, warp_col_tiles = 8, chunk = 16, reduce_k = 1, num_elems_per_byte = 1, is_m_first = False)
+.. py:class:: TensorCoreIntrinEmitter(a_dtype = 'float16', b_dtype = 'float16', accum_dtype = 'float16', a_transposed = False, b_transposed = False, block_row_warps = 2, block_col_warps = 2, warp_row_tiles = 8, warp_col_tiles = 8, chunk = 16, reduce_k = 1, num_elems_per_byte = 1, is_m_first = False, thread_var = None)
 
    Bases: :py:obj:`object`
 
@@ -41,7 +41,7 @@ Module Contents
 
 
 
-   .. py:attribute:: N_DIM
+   .. py:attribute:: n_dim
       :value: 16
 
 
@@ -109,16 +109,6 @@ Module Contents
 
 
 
-   .. py:attribute:: warp_rows
-      :value: 0
-
-
-
-   .. py:attribute:: warp_cols
-      :value: 0
-
-
-
    .. py:attribute:: reduce_k
       :value: 1
 
@@ -134,6 +124,14 @@ Module Contents
 
 
 
+   .. py:attribute:: thread_var
+      :value: None
+
+
+
+   .. py:method:: get_thread_binding()
+
+
    .. py:method:: get_store_index_map(inverse = False)
 
 
@@ -209,7 +207,7 @@ Module Contents
    .. py:method:: mma(A_local_buf, B_local_buf, C_local_buf)
 
 
-.. py:class:: INT4TensorCoreIntrinEmitter(a_dtype = 'float16', b_dtype = 'float16', accum_dtype = 'float16', a_transposed = False, b_transposed = False, block_row_warps = 2, block_col_warps = 2, warp_row_tiles = 8, warp_col_tiles = 8, chunk = 16, reduce_k = 1, num_elems_per_byte = 1, is_m_first = False)
+.. py:class:: INT4TensorCoreIntrinEmitter(a_dtype = 'float16', b_dtype = 'float16', accum_dtype = 'float16', a_transposed = False, b_transposed = False, block_row_warps = 2, block_col_warps = 2, warp_row_tiles = 8, warp_col_tiles = 8, chunk = 16, reduce_k = 1, num_elems_per_byte = 1, is_m_first = False, thread_var = None)
 
    Bases: :py:obj:`TensorCoreIntrinEmitter`
 

_sources/autoapi/tilelang/ir/index.rst.txt

@@ -52,6 +52,21 @@ Module Contents
    Bases: :py:obj:`tvm.ir.base.Node`, :py:obj:`tvm.runtime.Scriptable`
 
 
+   .. py:attribute:: policy_type
+      :type:  int
+
+
+   .. py:attribute:: m_warp
+      :type:  int
+
+
+   .. py:attribute:: n_warp
+      :type:  int
+
+
+   .. py:method:: compute_warp_partition(M, N, block_size, target, is_wgmma)
+
+
 .. py:class:: Gemm
 
    Bases: :py:obj:`tvm.ir.base.Node`, :py:obj:`tvm.runtime.Scriptable`

_sources/autoapi/tilelang/language/gemm/index.rst.txt

@@ -15,6 +15,7 @@ Functions
 .. autoapisummary::
 
    tilelang.language.gemm.gemm
+   tilelang.language.gemm.gemm_v2
 
 
 Module Contents
@@ -52,3 +53,35 @@ Module Contents
    :raises AssertionError: If the K dimensions of matrices A and B don't match
 
 
+.. py:function:: gemm_v2(A, B, C, transpose_A = False, transpose_B = False, policy = GemmWarpPolicy.Square, clear_accum = False, k_pack = 1, wg_wait = 0)
+
+   Perform a General Matrix Multiplication (GEMM) operation.
+
+   This function computes C = A @ B where A and B can optionally be transposed.
+   The operation supports various warp policies and accumulation modes.
+
+   :param A: First input matrix
+   :type A: Union[tir.Buffer, tir.Var]
+   :param B: Second input matrix
+   :type B: Union[tir.Buffer, tir.Var]
+   :param C: Output matrix for results
+   :type C: Union[tir.Buffer, tir.Var]
+   :param transpose_A: Whether to transpose matrix A. Defaults to False.
+   :type transpose_A: bool, optional
+   :param transpose_B: Whether to transpose matrix B. Defaults to False.
+   :type transpose_B: bool, optional
+   :param policy: Warp execution policy. Defaults to GemmWarpPolicy.Square.
+   :type policy: GemmWarpPolicy, optional
+   :param clear_accum: Whether to clear accumulator before computation. Defaults to False.
+   :type clear_accum: bool, optional
+   :param k_pack: Number of k dimensions packed into a single warp. Defaults to 1.
+   :type k_pack: int, optional
+   :param wg_wait: Warp group wait count. Defaults to 0.
+   :type wg_wait: int, optional
+
+   :returns: A handle to the GEMM operation
+   :rtype: tir.Call
+
+   :raises AssertionError: If the K dimensions of matrices A and B don't match
+
+