ECE 408 sem project

kernels_req_1/matmul.cuh

@@ -0,0 +1,98 @@
+#ifndef __MATMUL_KERNEL_CUH__
+#define __MATMUL_KERNEL_CUH__
+
+#include <cuda_runtime.h>
+#include <mma.h>
+#include "../utils/cuda_utils.cuh"
+
+using namespace nvcuda;
+
+#define WMMA_M 16
+#define WMMA_N 16
+#define WMMA_K 8
+
+#define WARPS_M 4
+#define WARPS_N 2
+#define BM 64
+#define BN 32
+#define BK WMMA_K
+
+__global__ void matmul_forward_kernel(float* out, const float* inp, const float* weight, 
+                                      const float* bias,int B,int T,  int C, int OC) {
+    // Implement this
+    int BT = B*T;
+
+    // block output
+    int start_row = blockIdx.x * BM;
+    int start_col = blockIdx.y * BN;
+
+    // which warp
+    int warp_id = threadIdx.x / 32;
+    int warp_row = warp_id / WARPS_N;
+    int warp_col = warp_id % WARPS_N;
+
+    int warp_out_row = start_row + warp_row * WMMA_M;
+    int warp_out_col = start_col + warp_col * WMMA_N;
+
+    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
+    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::col_major> b_frag;
+    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;
+
+    wmma::fill_fragment(c_frag, 0.0f);
+
+    for(int k = 0; k < C; k += WMMA_K){
+        if(warp_out_row < BT && k < C){
+            wmma::load_matrix_sync(a_frag, inp + warp_out_row * C + k, C);
+
+            #pragma unroll
+            for(int i = 0; i< a_frag.num_elements; i++){
+                a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
+            }
+        }
+
+        if (warp_out_col < OC && k < C) {
+            wmma::load_matrix_sync(b_frag, weight + warp_out_col * C + k, C);
+            #pragma unroll
+            for (int i = 0; i < b_frag.num_elements; i++) {
+                b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
+            }
+        }
+
+        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
+    }
+
+    //store  in s mem, add bais durinthe writw out  for performance
+    __shared__ float out_tile[BM * BN];
+
+    int row_smem = warp_row * WMMA_M;
+    int col_smem = warp_col * WMMA_N;
+    wmma::store_matrix_sync(out_tile + row_smem * BN + col_smem, c_frag, BN, wmma::mem_row_major);
+    __syncthreads();
+
+    for(int i = threadIdx.x; i < BM*BN; i += blockDim.x){
+        int r_l = i / BN;
+        int c_l = i % BN;
+        int r_g = start_row + r_l;
+        int c_g = start_col + c_l;
+
+        if(r_g < BT && c_g < OC){
+            if(bias != NULL){
+                out_tile[i] += bias[c_g];
+            }
+            out[r_g * OC + c_g] = out_tile[i];
+        }
+    }
+}
+
+// Launch kernel here
+void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
+                    int B, int T, int C, int OC) {
+    // Implement this
+    int BT = B*T;
+    dim3 block(256);
+    dim3 grid((BT + BM - 1) / BM, (OC + BN - 1) / BN);
+    matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
+
+}
+
+#endif // __MATMUL_KERNEL_CUH__