sgemm_tcore: Implement A transpose for coalesced smem access

tests/regression/sgemm_tcore/kernel.cpp

@@ -7,7 +7,7 @@
 #include "common.h"
 
 #define USE_TENSOR_CORE 1
-#define TC_SINGLE_WARP 0
+#define TC_SINGLE_WARP 1
 
 #define NUM_LANES 8
 
@@ -23,9 +23,9 @@
 //   (BM*BN) / (TM*TN) == threadblock size >= NT * CORES_PER_CLUSTER
 // * Combining BM * BK >= (BM*BN) / (TM*TN) == threadblock yields
 //   BM <= BK*TM*TN
-#define BM 16
+#define BM 8
-#define BN 16
+#define BN 8
-#define BK 32
+#define BK 8
 #define TCM 8
 #define TCN 8
 #define TCK 8
@@ -34,9 +34,13 @@
 #define WMITER (WM / TCM)
 #define WNITER (WN / TCN)
 #define TM 1
-#define TN ((TCM * TCN) / NUM_LANES / TM)
+// #define TN ((TCM * TCN) / NUM_LANES / TM)
-// #define TN 1
+#define TN 1
 
+// number of loop around the inner 0..TCK..BK loop to simulate perfect-DRAM
+// scenario
+#define BK_LOOP 1
+#define TRANSPOSE_AS 1
 
 inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) {
   const int tg = tid / 4;
@@ -130,7 +134,7 @@ inline void vx_wmma() {
 }
 
 // `local_k` is assumed to be multiple of TCK
-void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int local_k,
+inline void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int local_k,
                   const int warp_col, const int warp_row, const int wn_iter,
                   const int wm_iter, const int thread_in_warp) {
   int tid = thread_in_warp;
@@ -145,6 +149,7 @@ void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int loca
   int smem_B_rows = BK;
   int smem_B_cols = BN;
 
+  if constexpr (!TRANSPOSE_AS) {
     int A_offset = (row + WM * warp_row + TCM * wm_iter) * smem_A_cols;
 
     // @perf: bank conflicts
@@ -156,9 +161,20 @@ void vx_wmma_load(volatile float *smem_A, volatile float *smem_B, const int loca
     asm volatile("flw f5, %0" ::"m"(smem_A[A_offset + (local_k + 5)]));
     asm volatile("flw f6, %0" ::"m"(smem_A[A_offset + (local_k + 6)]));
     asm volatile("flw f7, %0" ::"m"(smem_A[A_offset + (local_k + 7)]));
+  } else {
+    // transposed A
+    asm volatile("flw  f0, %0" ::"m"(smem_A[((local_k + 0) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f1, %0" ::"m"(smem_A[((local_k + 1) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f2, %0" ::"m"(smem_A[((local_k + 2) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f3, %0" ::"m"(smem_A[((local_k + 3) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f4, %0" ::"m"(smem_A[((local_k + 4) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f5, %0" ::"m"(smem_A[((local_k + 5) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f6, %0" ::"m"(smem_A[((local_k + 6) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+    asm volatile("flw  f7, %0" ::"m"(smem_A[((local_k + 7) * smem_A_rows) + (WM * warp_row + TCM * wm_iter) + row]));
+  }
 
-  asm volatile("flw f8 , %0" ::"m"(smem_B[((local_k + 0) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
+  asm volatile("flw  f8, %0" ::"m"(smem_B[((local_k + 0) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
-  asm volatile("flw f9 , %0" ::"m"(smem_B[((local_k + 1) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
+  asm volatile("flw  f9, %0" ::"m"(smem_B[((local_k + 1) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
   asm volatile("flw f10, %0" ::"m"(smem_B[((local_k + 2) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
   asm volatile("flw f11, %0" ::"m"(smem_B[((local_k + 3) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
   asm volatile("flw f12, %0" ::"m"(smem_B[((local_k + 4) * smem_B_cols) + (WN * warp_col + TCN * wn_iter) + col]));
@@ -233,10 +249,10 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
 
   const uint32_t local_a_row = tid_in_threadblock / BK;
   const uint32_t local_a_col = tid_in_threadblock % BK;
+  const uint32_t local_as_row = tid_in_threadblock / BM;
+  const uint32_t local_as_col = tid_in_threadblock % BM;
   const uint32_t local_b_row = tid_in_threadblock / BN;
   const uint32_t local_b_col = tid_in_threadblock % BN;
-  const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
-  const uint32_t global_b_col = BN * threadblock_id_x + local_b_col;
 
   const uint32_t local_c_row = tid_in_threadblock / (BN / TN);
   const uint32_t local_c_col = tid_in_threadblock % (BN / TN);
@@ -259,10 +275,6 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
   volatile float *local_warp_results =
       local_b + local_b_elems + (warp_in_threadblock * TCM * TCN);
 
-  // number of rows a full TB can read at a time
-  constexpr uint32_t row_stride_a = (BM * BN) / BK / (TM * TN);
-  constexpr uint32_t row_stride_b = (BM * BN) / BN / (TM * TN);
-
   // clear out C
   initialize_C();
 
@@ -273,16 +285,34 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
     // neighboring threads to ensure GMEM coalescing.
     //
     // TODO: Sharedmem swizzling is important here
-#pragma GCC unroll 2
+    if constexpr (!TRANSPOSE_AS) {
+      const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
+      // number of rows a full TB can read at a time
+      constexpr uint32_t row_stride_a = (BM * BN) / BK / (TM * TN);
+#pragma GCC unroll 1
       for (uint32_t load_offset = 0; load_offset < BM; load_offset += row_stride_a) {
         const uint32_t global_a_offset =
             dim_k * (global_a_row + load_offset) + (k + local_a_col);
-      // FIXME: all threads in TB (BM*BN) will do this load, even if this is
+        // NOTE: all threads in TB  will do this load; make sure this is not
         // out-of-bounds of BM*BK
         local_a[BK * (local_a_row + load_offset) + local_a_col] =
             A[global_a_offset];
       }
-#pragma GCC unroll 2
+    } else {
+      const uint32_t global_a_row = BM * threadblock_id_y + local_as_col;
+      constexpr uint32_t row_stride_a = (BM * BN) / BM / (TM * TN);
+#pragma GCC unroll 1
+      for (uint32_t load_offset = 0; load_offset < BK; load_offset += row_stride_a) {
+        const uint32_t global_a_offset =
+            dim_k * (global_a_row + load_offset) + (k + local_as_row);
+        local_a[BM * (local_as_row + load_offset) + local_as_col] =
+            A[global_a_offset];
+      }
+    }
+
+  constexpr uint32_t row_stride_b = (BM * BN) / BN / (TM * TN);
+  const uint32_t global_b_col = BN * threadblock_id_x + local_b_col;
+#pragma GCC unroll 1
     for (uint32_t load_offset = 0; load_offset < BK; load_offset += row_stride_b) {
       const uint32_t global_b_offset =
           dim_n * (k + local_b_row + load_offset) + global_b_col;
@@ -294,6 +324,10 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
                         threadblock_dim_y);
 
 #if USE_TENSOR_CORE
+// #pragma GCC unroll 1
+    for (int i = 0; i < BK_LOOP; i++) {
+#pragma GCC unroll 1
+      // @perf: this loop spills to stack a lot because of all the flws in vx_wmma_load
       for (uint32_t local_k = 0; local_k < BK; local_k += TCK) {
         // perform wmma
         // vx_wmma_load(local_a, local_b, warp_x, warp_y, tid_in_warp);
@@ -306,8 +340,10 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
 #if TC_SINGLE_WARP
             if (warp_in_threadblock == 0) {
 #endif
+              // SMEM -> RF
               vx_wmma_load(local_a, local_b, local_k, warp_col, warp_row, wn_iter,
                   wm_iter, tid_in_warp);
+              // compute
               vx_wmma();
 #if TC_SINGLE_WARP
             }
@@ -315,6 +351,7 @@ void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
           }
         }
       }
+    }
 
 #else