98 lines
2.9 KiB
Plaintext
98 lines
2.9 KiB
Plaintext
#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__ |