Comprehensively support WGMMA GEMM SS

Author: LeiWang1999

src/layout/gemm_layouts.cc

@@ -385,6 +385,7 @@ Layout makeQuarterBankSwizzleLayout(int stride, int continuous,
   Var i = InputPlaceholder(0);
   Var j = InputPlaceholder(1);
   int vector_size = 128 / element_size;
+  LOG(INFO) << "makeQuarterBankSwizzleLayout: " << stride << ", " << continuous << ", " << element_size;
   ICHECK(stride % 8 == 0) << "stride=" << stride;
   ICHECK(continuous % (vector_size * 2) == 0)
       << "continuous=" << continuous << ", vector_size=" << vector_size;
@@ -740,6 +741,7 @@ Layout makeGemmABLayout(int mat_stride, int mat_continuous, int continuity,
 
 Layout makeGemmABLayoutHopper(int mat_stride, int mat_continuous,
                               int continuity, int element_size, bool k_inner) {
+  LOG(INFO) << "makeGemmABLayoutHopper: " << mat_stride << ", " << mat_continuous << ", " << continuity << ", " << element_size << ", " << k_inner;
   if (element_size == 64) {
     if (!k_inner && continuity % 16 == 0) // float64 KxN
       return makeGemmABLayoutF64_Kouter(mat_stride, mat_continuous);
@@ -749,6 +751,12 @@ Layout makeGemmABLayoutHopper(int mat_stride, int mat_continuous,
                                         element_size);
   }
   int vector_size = 128 / element_size;
+  LOG(INFO) << "makeGemmABLayoutHopper: mat_continuous: " << mat_continuous << ", mat_stride: " << mat_stride << ", element_size: " << element_size;
+  LOG(INFO) << "vector_size: " << vector_size;
+  LOG(INFO) << "mat_continuous % (vector_size * 8): " << mat_continuous % (vector_size * 8);
+  LOG(INFO) << "mat_continuous % (vector_size * 4): " << mat_continuous % (vector_size * 4);
+  LOG(INFO) << "mat_continuous % (vector_size * 2): " << mat_continuous % (vector_size * 2);
+  LOG(INFO) << "mat_continuous % vector_size: " << mat_continuous % vector_size;
   if (mat_continuous % (vector_size * 8) == 0)
     return makeFullBankSwizzleLayout(mat_stride, mat_continuous, element_size);
   else if (mat_continuous % (vector_size * 4) == 0)

src/layout/layout.cc

@@ -82,7 +82,7 @@ void LayoutNode::RegisterReflection() {
 }
 
 void LayoutNode::UpdateAnalyzer(arith::Analyzer *analyzer) const {
-  for (const auto &[var, dom] : LayoutNode::getVarMap()) {
+  for (const auto &[var, dom] : getVarMap()) {
     analyzer->Bind(var, dom);
   }
 }
@@ -547,6 +547,10 @@ TVM_FFI_STATIC_INIT_BLOCK({
            [](int stride, int continuous, int element_size) {
              return makeQuarterBankSwizzleLayout(stride, continuous,
                                                  element_size);
+           })
+      .def("tl.make_linear_layout",
+           [](int stride, int continuous) {
+             return makeGemmLayoutLinear(stride, continuous);
            });
 });
 

src/op/gemm.cc

@@ -756,6 +756,8 @@ LayoutMap GemmNode::InferLayout(const LayoutInferArgs &T,
       const int64_t mat_continuous = *as_const_int(B->shape[dim_B - 1]);
       const int64_t continuity =
           trans_B ? mat_continuous : mat_continuous / warp_n;
+
+    LOG(INFO) << "gemm_inst: " << (int)gemm_inst << ", trans_B: " << trans_B;
       auto ABLayout =
           gemm_inst == GemmInst::kWGMMA
               ? makeGemmABLayoutHopper(mat_stride, mat_continuous, continuity,

tilelang/intrinsics/wgmma_macro_generator.py

@@ -10,6 +10,7 @@
     make_full_bank_swizzled_layout,
     make_half_bank_swizzled_layout,
     make_quarter_bank_swizzled_layout,
+    make_linear_layout,
 )
 from tvm.runtime import convert
 from tilelang.intrinsics.mma_layout import (shared_16x8_to_mma_32x4_layout_sr_a,
@@ -131,13 +132,20 @@ def _initialize_micro_size(self, m_dim: int = 16, k_dim: int = 16):
         self.micro_size_k = k_dim
 
     def _determinate_swizzle_mode(self, buffer: Buffer, layout: Layout) -> SwizzleMode:
+        # same behavior to src/layout/gemm_layouts.cc::makeGemmABLayoutHopper
+        mat_stride = int(buffer.shape[-2])
+        mat_continuous = int(buffer.shape[-1])
+        element_size = DataType(buffer.dtype).bits
+        print(f"_determinate_swizzle_mode mat_stride: {mat_stride}, mat_continuous: {mat_continuous}, element_size: {element_size}")
         if layout is None:
             return SwizzleMode.NONE
-        elif layout.is_equal(make_quarter_bank_swizzled_layout(buffer)):
+        elif layout.is_equal(make_linear_layout(mat_stride, mat_continuous)):
+            return SwizzleMode.NONE
+        elif layout.is_equal(make_quarter_bank_swizzled_layout(mat_stride, mat_continuous, element_size)):
             return SwizzleMode.SWIZZLE_32B
-        elif layout.is_equal(make_half_bank_swizzled_layout(buffer)):
+        elif layout.is_equal(make_half_bank_swizzled_layout(mat_stride, mat_continuous, element_size)):
             return SwizzleMode.SWIZZLE_64B
-        elif layout.is_equal(make_full_bank_swizzled_layout(buffer)):
+        elif layout.is_equal(make_full_bank_swizzled_layout(mat_stride, mat_continuous, element_size)):
             return SwizzleMode.SWIZZLE_128B
         else:
             raise ValueError(f"Unsupported swizzle mode: {layout}")
@@ -173,7 +181,11 @@ def wgmma(self,
         a_swizzle_mode = self._determinate_swizzle_mode(A_buf, self.a_shared_layout)
         b_swizzle_mode = self._determinate_swizzle_mode(B_buf, self.b_shared_layout)
 
-        elems_in_bytes = DataType(self.a_dtype).bits // 8
+        elems_in_bits = DataType(self.a_dtype).bits
+        elems_in_bytes = elems_in_bits // 8
+        
+        a_swizzle_atom_elems = a_swizzle_mode.swizzle_byte_size() // elems_in_bytes
+        b_swizzle_atom_elems = n_dim if b_swizzle_mode.is_none() else b_swizzle_mode.swizzle_byte_size() // elems_in_bytes
 
         # by default, we utilize non-swizzle layout offset
         a_leading_byte_offset = (8 * 8 * elems_in_bytes) if a_is_k_major else (8 * m_dim *
@@ -186,52 +198,59 @@ def wgmma(self,
             # https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-leading-dimension-byte-offset
             if a_is_k_major:
                 a_leading_byte_offset = 16
+                a_stride_byte_offset = 8 * a_swizzle_mode.swizzle_byte_size()
             else:
                 # MN Major
                 # LBO represents the distance between two atoms along the M dimension
                 # SBO represents the distance between two atoms along the K dimension
-                a_leading_byte_offset = a_swizzle_mode.swizzle_atom_size()
-                a_stride_byte_offset = 8 * 64 * elems_in_bytes
+                a_m_axis_atoms = m_dim // a_swizzle_atom_elems
+                if a_m_axis_atoms <= 1:
+                    a_leading_byte_offset = 0
+                else:
+                    a_leading_byte_offset = 8 * a_swizzle_mode.swizzle_atom_size() * (a_swizzle_mode.swizzle_byte_size() // elems_in_bytes)
+
+                if a_m_axis_atoms <= 1:
+                    a_stride_byte_offset = 8 * elems_in_bytes * m_dim
+                else:
+                    a_stride_byte_offset = 8 * elems_in_bytes * a_swizzle_atom_elems
 
         b_leading_byte_offset = (8 * 8 * elems_in_bytes) if b_is_k_major else (8 * n_dim *
                                                                                elems_in_bytes)
-        b_stride_byte_offset = (8 * k_dim * elems_in_bytes) if b_is_k_major else (8 * 8 *
-                                                                                  elems_in_bytes)
+        b_stride_byte_offset = (8 * k_dim * elems_in_bytes) if b_is_k_major else (
+            0 if n_dim == 8 else (8 * 8 * elems_in_bytes)
+        )
         if not b_swizzle_mode.is_none():
             # swizzle mode doesn't require LBO/SBO to be 1
             # https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-leading-dimension-byte-offset
             if b_is_k_major:
                 b_leading_byte_offset = 16
+                b_stride_byte_offset = 8 * b_swizzle_mode.swizzle_byte_size()
             else:
                 # MN Major, K * N
                 # LBO represents the distance between two atoms along the N dimension
                 # SBO represents the distance between two atoms along the K dimension
-                b_n_axis_atoms = n_dim // (b_swizzle_mode.swizzle_byte_size() // elems_in_bytes)
+                b_n_axis_atoms = n_dim // b_swizzle_atom_elems
                 if b_n_axis_atoms <= 1:
                     b_leading_byte_offset = 0
                 else:
                     b_leading_byte_offset = 8 * 8 * elems_in_bytes * k_dim
-
                 if b_n_axis_atoms <= 1:
                     b_stride_byte_offset = 8 * elems_in_bytes * n_dim
                 else:
-                    b_stride_byte_offset = 8 * elems_in_bytes * (b_swizzle_mode.swizzle_byte_size() // elems_in_bytes)
-
-                
+                    b_stride_byte_offset = 8 * elems_in_bytes * b_swizzle_atom_elems
         print(f"a_leading_byte_offset: {a_leading_byte_offset >> 4}")
         print(f"a_stride_byte_offset: {a_stride_byte_offset >> 4}")
+        print(f"b_leading_byte_offset: {b_leading_byte_offset >> 4}")
+        print(f"b_stride_byte_offset: {b_stride_byte_offset >> 4}")
 
         print(f"b_swizzle_atom_size: {b_swizzle_mode.swizzle_atom_size()}")
         print(f"b_swizzle_byte_size: {b_swizzle_mode.swizzle_byte_size()}")
-        print(f"m_dim: {m_dim}")
-        print(f"n_dim: {n_dim}")
-        print(f"k_dim: {k_dim}")
-        print(f"micro_size_k: {micro_size_k}")
-        print(f"a_leading_byte_offset: {a_leading_byte_offset}")
-        print(f"a_stride_byte_offset: {a_stride_byte_offset}")
-        print(f"b_leading_byte_offset: {b_leading_byte_offset}")
-        print(f"b_stride_byte_offset: {b_stride_byte_offset}")
-        # exit()
+
+        # for example, if [n, k] where k is 128, we should split it into 2 atoms
+        # where max specially handles the case when n_dim is 8.
+        ak_atom_size = max(a_swizzle_atom_elems // micro_size_k, 1)
+        bk_atom_size = max(b_swizzle_atom_elems // micro_size_k, 1)
+
         @T.macro
         def _warp_mma(A_buf, B_buf, C_local_buf):
             desc_a = T.alloc_descriptor()
@@ -242,10 +261,8 @@ def _warp_mma(A_buf, B_buf, C_local_buf):
                                     int(b_leading_byte_offset >> 4), int(b_stride_byte_offset >> 4))
             for ki in T.serial(0, (k_dim // micro_size_k)):
                 for i in T.serial(m_dim // 64):
-                    k_dim_offset = ki * micro_size_k
-                    A_offset = i * 64 * A_buf.shape[
-                        -1] + k_dim_offset if a_is_k_major else ki * micro_size_k * 64 + i * 64 * k_dim
-                    B_offset = k_dim_offset if b_is_k_major else k_dim_offset * (b_swizzle_mode.swizzle_byte_size() // elems_in_bytes)
+                    A_offset = (ki % ak_atom_size) * micro_size_k + i * 64 * a_swizzle_atom_elems + (ki // ak_atom_size) * m_dim * a_swizzle_atom_elems if a_is_k_major else i * 64 * k_dim + ki * a_swizzle_atom_elems * micro_size_k
+                    B_offset = (ki // bk_atom_size) * n_dim * b_swizzle_atom_elems + (ki % bk_atom_size) * micro_size_k if b_is_k_major else ki * b_swizzle_atom_elems * micro_size_k
                     C_offset = i * warp_cols * local_size_out  # 4 warps as an unit
                     T.ptx_wgmma_ss(accum_dtype, wgmma_prefix, a_is_k_major,
                                    b_is_k_major, a_dtype_abbrv, b_dtype_abbrv,
@@ -300,7 +317,7 @@ def wgmma_rs(self,
                 if b_n_axis_atoms <= 1:
                     b_leading_byte_offset = 0
                 else:
-                    b_leading_byte_offset = 8 * 8 * elems_in_bytes * k_dim
+                    b_leading_byte_offset = 8 * b_swizzle_mode.swizzle_atom_size() * (b_swizzle_mode.swizzle_byte_size() // elems_in_bytes)
 
                 if b_n_axis_atoms <= 1:
                     b_stride_byte_offset = 8 * elems_in_bytes * n_dim

tilelang/layout/__init__.py

@@ -9,5 +9,6 @@
     make_full_bank_swizzled_layout,  # noqa: F401
     make_half_bank_swizzled_layout,  # noqa: F401
     make_quarter_bank_swizzled_layout,  # noqa: F401
+    make_linear_layout,  # noqa: F401
 )
 from .gemm_sp import make_metadata_layout  # noqa: F401

tilelang/layout/swizzle.py

@@ -25,6 +25,7 @@ def make_wgmma_swizzled_layout(buffer: tvm.tir.Buffer,
     assert len(buffer.shape) == 2
     if continuity is None:
         continuity = int(buffer.shape[1])
+    print(f"make_wgmma_swizzled_layout: {buffer.shape[0]}, {buffer.shape[1]}, {continuity}, {tvm.DataType(buffer.dtype).bits}, {k_major}")
     return _ffi_api.make_wgmma_swizzled_layout(
         int(buffer.shape[0]),
         int(buffer.shape[1]),
@@ -107,3 +108,23 @@ def make_quarter_bank_swizzled_layout(*args):
         continuous,
         element_size,
     )
+
+def make_linear_layout(*args):
+    """
+    Args:
+        args: buffer or (stride, continuous)
+    Examples:
+        make_linear_layout(buffer)
+        make_linear_layout(stride, continuous)
+    """
+    if len(args) == 1:
+        buffer = args[0]
+        stride, continuous = int(buffer.shape[0]), int(buffer.shape[1])
+    elif len(args) == 2:
+        stride, continuous = args
+    else:
+        raise ValueError(f"Invalid arguments: {args}")
+    return _ffi_api.make_linear_layout(
+        stride,
+        continuous,
+    )

tilelang/tileop/gemm/gemm_wgmma.py

@@ -34,7 +34,8 @@ def infer_layout(self, target: Target, thread_nums: int):
 
         if self.is_gemm_ss():
             a_continuity = self.M if a_is_k_major else 4 * self.K // m_warp
-            b_continuity = self.N if b_is_k_major else 4 * self.K // n_warp
+            b_continuity = self.K if b_is_k_major else self.N // n_warp
+
             return {
                 # WGMMA does not support padding
                 self.A: