57 lines
1.9 KiB
Plaintext
57 lines
1.9 KiB
Plaintext
#ifndef __MATMUL_KERNEL_CUH__
|
|
#define __MATMUL_KERNEL_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
#define TILE_SIZE 32
|
|
#define PADDING 1
|
|
|
|
//op12
|
|
// Adding +1 padding to the inner dimension shifts each row by one bank, distributing the 32 column reads across all 32 banks.
|
|
|
|
__global__ void matmul_kernel(const float* __restrict__ inp, float* __restrict__ out, const float* __restrict__ weights,
|
|
const float* __restrict__ biases, int B, int T, int OC, int C ) {
|
|
// Implement this
|
|
// reshape (BT X C) * (OC X C)^T -> (BT X OC)
|
|
|
|
int bt = blockIdx.y * TILE_SIZE + threadIdx.y;
|
|
int oc = blockIdx.x * TILE_SIZE + threadIdx.x;
|
|
int BT = B*T;
|
|
|
|
__shared__ float A_smem_tile[TILE_SIZE][TILE_SIZE + PADDING];
|
|
__shared__ float B_smem_tile[TILE_SIZE][TILE_SIZE + PADDING];
|
|
|
|
float acc = 0.0f;
|
|
int num_tiles = (C + TILE_SIZE - 1) / TILE_SIZE;
|
|
|
|
for(int t = 0; t < num_tiles; t++){
|
|
int c = t * TILE_SIZE + threadIdx.x;
|
|
int c_ = t * TILE_SIZE + threadIdx.y;
|
|
A_smem_tile[threadIdx.y][threadIdx.x] = (bt < BT && c < C) ? inp[bt * C + c] : 0.0f;
|
|
B_smem_tile[threadIdx.y][threadIdx.x] = (oc < OC && c_ < C) ? weights[oc * C + c_] : 0.0f;
|
|
|
|
__syncthreads();
|
|
|
|
#pragma unroll
|
|
for(int i = 0; i < TILE_SIZE; i++){
|
|
acc += A_smem_tile[threadIdx.y][i] * B_smem_tile[i][threadIdx.x];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if(bt < BT && oc < OC){
|
|
if(biases != nullptr){
|
|
acc += biases[oc];
|
|
}
|
|
out[bt * OC + oc] = acc;
|
|
}
|
|
}
|
|
|
|
void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
|
|
int B, int T, int C, int OC){
|
|
dim3 block(TILE_SIZE, TILE_SIZE);
|
|
dim3 grid((OC + TILE_SIZE - 1) / TILE_SIZE, (B * T + TILE_SIZE - 1) / TILE_SIZE);
|
|
matmul_kernel<<<grid, block>>>(inp, out, weight, bias, B, T, OC, C);
|
|
}
|
|
|
|
#endif |