Update docs

Author: github-actions[bot]

_images/sparse_mma_storage_example.png

@@ -16,6 +16,8 @@ Submodules
    /autoapi/tilelang/intrinsics/mma_macro_generator/index
    /autoapi/tilelang/intrinsics/mma_sm70_layout/index
    /autoapi/tilelang/intrinsics/mma_sm70_macro_generator/index
+   /autoapi/tilelang/intrinsics/mma_sp_layout/index
+   /autoapi/tilelang/intrinsics/mma_sp_macro_generator/index
    /autoapi/tilelang/intrinsics/tcgen05_macro_generator/index
    /autoapi/tilelang/intrinsics/utils/index
    /autoapi/tilelang/intrinsics/wgmma_macro_generator/index

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

@@ -52,7 +52,9 @@ Functions
    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_a_32x8_to_shared_16x16_layout
    tilelang.intrinsics.mma_layout.mma_load_b_32x16_to_shared_16x32_layout
+   tilelang.intrinsics.mma_layout.mma_load_b_32x8_to_shared_16x16_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
@@ -135,8 +137,31 @@ Module Contents
 
 .. py:function:: mma_load_a_32x16_to_shared_16x32_layout(thread_id, local_id)
 
+.. py:function:: mma_load_a_32x8_to_shared_16x16_layout(thread_id, local_id)
+
+   groupID           = %laneid >> 2
+   threadID_in_group = %laneid % 4
+
+   row =      groupID            for ai where  0 <= i < 2 || 4 <= i < 6
+           groupID + 8         Otherwise
+
+   col =  (threadID_in_group * 2) + (i & 0x1)          for ai where i <  4
+   (threadID_in_group * 2) + (i & 0x1) + 8      for ai where i >= 4
+
+
 .. py:function:: mma_load_b_32x16_to_shared_16x32_layout(thread_id, local_id)
 
+.. py:function:: mma_load_b_32x8_to_shared_16x16_layout(thread_id, local_id)
+
+   groupID           = %laneid >> 2
+   threadID_in_group = %laneid % 4
+
+   row =  (threadID_in_group * 2) + (i & 0x1)           for bi where i <  2
+       (threadID_in_group * 2) + (i & 0x1) + 8       for bi where i >= 2
+
+   col = groupID
+
+
 .. 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_layout/index.rst.txt

@@ -0,0 +1,94 @@
+tilelang.intrinsics.mma_sp_layout
+=================================
+
+.. py:module:: tilelang.intrinsics.mma_sp_layout
+
+
+Functions
+---------
+
+.. autoapisummary::
+
+   tilelang.intrinsics.mma_sp_layout.shared_16x16_to_mma_sp_layout_sr_a
+   tilelang.intrinsics.mma_sp_layout.shared_16x16_to_mma_sp_layout_sr_b
+   tilelang.intrinsics.mma_sp_layout.shared_16x32_to_mma_sp_layout_sr_a
+   tilelang.intrinsics.mma_sp_layout.shared_16x32_to_mma_sp_layout_sr_b
+   tilelang.intrinsics.mma_sp_layout.shared_16x64_to_mma_sp_layout_sr_a
+   tilelang.intrinsics.mma_sp_layout.shared_16x64_to_mma_sp_layout_sr_b
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_a_32x4_to_shared_16x16_layout
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_a_32x8_to_shared_16x32_layout
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_a_32x16_to_shared_16x64_layout
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_b_32x8_to_shared_16x16_layout
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_b_32x16_to_shared_16x32_layout
+   tilelang.intrinsics.mma_sp_layout.mma_sp_load_b_32x32_to_shared_16x64_layout
+   tilelang.intrinsics.mma_sp_layout.get_logical_id_32bit
+   tilelang.intrinsics.mma_sp_layout.metadata_8bit_load_32x4_to_shared_16x4_layout_32bit
+   tilelang.intrinsics.mma_sp_layout.metadata_16bit_load_32x2_to_shared_16x2_layout_32bit
+   tilelang.intrinsics.mma_sp_layout.metadata_8bit_load_32x4_to_shared_16x4_layout_16bit
+   tilelang.intrinsics.mma_sp_layout.metadata_16bit_load_32x2_to_shared_16x2_layout_16bit
+   tilelang.intrinsics.mma_sp_layout.get_logical_id_8bit
+   tilelang.intrinsics.mma_sp_layout.metadata_8bit_load_32x4_to_shared_16x4_layout_8bit
+   tilelang.intrinsics.mma_sp_layout.metadata_16bit_load_32x2_to_shared_16x4_layout_8bit
+   tilelang.intrinsics.mma_sp_layout.metadata_32bit_load_32x1_to_shared_16x2_layout_8bit
+   tilelang.intrinsics.mma_sp_layout.ldmatrix_trans_32x8_to_shared_16x16_layout
+   tilelang.intrinsics.mma_sp_layout.ldmatrix_32x16_to_shared_32x16_layout
+   tilelang.intrinsics.mma_sp_layout.ldmatrix_trans_32x16_to_shared_16x32_layout
+   tilelang.intrinsics.mma_sp_layout.ldmatrix_trans_32x32_to_shared_shared_16x64_layout
+   tilelang.intrinsics.mma_sp_layout.get_ldmatrix_offset_b
+
+
+Module Contents
+---------------
+
+.. py:function:: shared_16x16_to_mma_sp_layout_sr_a(i, j)
+
+.. py:function:: shared_16x16_to_mma_sp_layout_sr_b(i, j)
+
+.. py:function:: shared_16x32_to_mma_sp_layout_sr_a(i, j)
+
+.. py:function:: shared_16x32_to_mma_sp_layout_sr_b(i, j)
+
+.. py:function:: shared_16x64_to_mma_sp_layout_sr_a(i, j)
+
+.. py:function:: shared_16x64_to_mma_sp_layout_sr_b(i, j)
+
+.. py:function:: mma_sp_load_a_32x4_to_shared_16x16_layout(thread_id, local_id)
+
+.. py:function:: mma_sp_load_a_32x8_to_shared_16x32_layout(thread_id, local_id)
+
+.. py:function:: mma_sp_load_a_32x16_to_shared_16x64_layout(thread_id, local_id)
+
+.. py:function:: mma_sp_load_b_32x8_to_shared_16x16_layout(thread_id, local_id)
+
+.. py:function:: mma_sp_load_b_32x16_to_shared_16x32_layout(thread_id, local_id)
+
+.. py:function:: mma_sp_load_b_32x32_to_shared_16x64_layout(thread_id, local_id)
+
+.. py:function:: get_logical_id_32bit(thread_id)
+
+.. py:function:: metadata_8bit_load_32x4_to_shared_16x4_layout_32bit(thread_id, local_id)
+
+.. py:function:: metadata_16bit_load_32x2_to_shared_16x2_layout_32bit(thread_id, local_id)
+
+.. py:function:: metadata_8bit_load_32x4_to_shared_16x4_layout_16bit(thread_id, local_id)
+
+.. py:function:: metadata_16bit_load_32x2_to_shared_16x2_layout_16bit(thread_id, local_id)
+
+.. py:function:: get_logical_id_8bit(thread_id)
+
+.. py:function:: metadata_8bit_load_32x4_to_shared_16x4_layout_8bit(thread_id, local_id)
+
+.. py:function:: metadata_16bit_load_32x2_to_shared_16x4_layout_8bit(thread_id, local_id)
+
+.. py:function:: metadata_32bit_load_32x1_to_shared_16x2_layout_8bit(thread_id, local_id)
+
+.. py:function:: ldmatrix_trans_32x8_to_shared_16x16_layout(thread_id, local_id)
+
+.. py:function:: ldmatrix_32x16_to_shared_32x16_layout(thread_id, local_id)
+
+.. py:function:: ldmatrix_trans_32x16_to_shared_16x32_layout(thread_id, local_id)
+
+.. py:function:: ldmatrix_trans_32x32_to_shared_shared_16x64_layout(thread_id, local_id)
+
+.. py:function:: get_ldmatrix_offset_b(matrix, row_idx, col_idx, stride, dtype = 'float16', transposed = False)
+

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

@@ -0,0 +1,219 @@
+tilelang.intrinsics.mma_sp_macro_generator
+==========================================
+
+.. py:module:: tilelang.intrinsics.mma_sp_macro_generator
+
+
+Attributes
+----------
+
+.. autoapisummary::
+
+   tilelang.intrinsics.mma_sp_macro_generator.lift
+
+
+Classes
+-------
+
+.. autoapisummary::
+
+   tilelang.intrinsics.mma_sp_macro_generator.SparseTensorCoreIntrinEmitter
+
+
+Module Contents
+---------------
+
+.. py:data:: lift
+
+.. py:class:: SparseTensorCoreIntrinEmitter(a_dtype = 'float16', e_dtype = 'uint8', b_dtype = 'float16', accum_dtype = 'float16', a_transposed = False, b_transposed = False, e_transposed = False, block_row_warps = 2, block_col_warps = 2, warp_row_tiles = 8, warp_col_tiles = 8, warp_k = 16, reduce_k = 1, num_elems_per_byte = 1, is_m_first = False, thread_var = None)
+
+   To eliminate Python syntax within TIR Macro.
+
+
+   .. py:attribute:: M_DIM
+      :value: 16
+
+
+
+   .. py:attribute:: SPARSE_FACTOR
+      :value: 2
+
+
+
+   .. py:attribute:: SPARSE_SELECTOR
+      :value: 0
+
+
+
+   .. py:attribute:: n_dim
+      :value: 16
+
+
+
+   .. py:attribute:: WARP_SIZE
+      :value: 32
+
+
+
+   .. py:attribute:: dtype_abbrv
+
+
+   .. py:attribute:: E_FACTOR_MAP
+
+
+   .. py:attribute:: E_REPLICATE_FACTOR
+
+
+   .. py:attribute:: is_m_first
+      :value: False
+
+
+
+   .. py:attribute:: a_dtype
+      :value: 'float16'
+
+
+
+   .. py:attribute:: e_dtype
+      :value: 'uint8'
+
+
+
+   .. py:attribute:: b_dtype
+      :value: 'float16'
+
+
+
+   .. py:attribute:: accum_dtype
+      :value: 'float16'
+
+
+
+   .. py:attribute:: a_transposed
+      :value: False
+
+
+
+   .. py:attribute:: b_transposed
+      :value: False
+
+
+
+   .. py:attribute:: e_transposed
+      :value: False
+
+
+
+   .. py:attribute:: block_row_warps
+      :value: 2
+
+
+
+   .. py:attribute:: block_col_warps
+      :value: 2
+
+
+
+   .. py:attribute:: warp_row_tiles
+      :value: 8
+
+
+
+   .. py:attribute:: warp_col_tiles
+      :value: 8
+
+
+
+   .. py:attribute:: warp_k
+      :value: 16
+
+
+
+   .. py:attribute:: e_factor
+      :value: 8
+
+
+
+   .. py:attribute:: reduce_k
+      :value: 1
+
+
+
+   .. py:attribute:: threads
+      :value: 128
+
+
+
+   .. py:attribute:: num_elems_per_byte
+      :value: 1
+
+
+
+   .. py:attribute:: thread_var
+      :value: None
+
+
+
+   .. py:method:: get_thread_binding()
+
+
+   .. py:method:: get_store_index_map(inverse = False)
+
+
+   .. py:method:: extract_thread_binding(thread_id, is_m_first = None)
+
+      is_m_first: True if the thread binding is in the form of (tx, warp_n, warp_m)
+      which represents [warp_size, block_row_warps (split n), block_col_warps (split m)]
+      Otherwise, it is in the form of [warp_size, block_col_warps (split m), block_row_warps (split n)]
+
+
+
+   .. py:method:: ldmatrix_a(A_local_buf, A_shared_buf, ki, rk = 0)
+
+
+   .. py:method:: ldmatrix_e(E_local_buf, E_shared_buf, ki, rk = 0)
+
+
+   .. py:method:: ldmatrix_b(B_local_buf, B_shared_buf, ki, rk = 0)
+
+
+   .. py:method:: mma_sp(A_local_buf, E_local_buf, B_local_buf, C_local_buf, k_inner = 0)
+
+
+   .. py:method:: stmatrix(C_local_buf, C_buf, pid_m=None, pid_n=None)
+
+
+   .. py:method:: make_mma_load_layout(local_buf, matrix = 'A')
+
+      Create a layout function for storing MMA results into a fragment buffer.
+      This layout is used in conjunction with `inverse_mma_store_layout` to
+      map fragment indices to threads and local indices.
+
+      :param local_buf: The local buffer representing a fragment of a matrix.
+      :type local_buf: tir.Buffer
+
+      :returns: A fragment object that describes how threads and indices
+                in `local_buf` are laid out.
+      :rtype: T.Fragment
+
+      :raises AssertionError: If `local_buf` is not detected to be a fragment buffer.
+
+
+
+   .. py:method:: make_mma_store_layout(local_buf)
+
+      Create a layout function for storing MMA results into a fragment buffer.
+      This layout is used in conjunction with `inverse_mma_store_layout` to
+      map fragment indices to threads and local indices.
+
+      :param local_buf: The local buffer representing a fragment of a matrix.
+      :type local_buf: tir.Buffer
+
+      :returns: A fragment object that describes how threads and indices
+                in `local_buf` are laid out.
+      :rtype: T.Fragment
+
+      :raises AssertionError: If `local_buf` is not detected to be a fragment buffer.
+
+
+