73 lines
2.8 KiB
Plaintext
73 lines
2.8 KiB
Plaintext
#ifndef __MATMUL_KERNEL_CUH__
|
|
#define __MATMUL_KERNEL_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
#include "cutlass/cutlass.h"
|
|
#include "cutlass/gemm/device/gemm.h"
|
|
#include "/work/hdd/bche/hshah6/cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h"
|
|
|
|
//referenced: https://github.com/NVIDIA/cutlass/blob/main/examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu
|
|
|
|
__global__ void bias_add_filler_kernel(float* __restrict__ out, const float* __restrict__ bias, int BT, int OC) {
|
|
int x = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (x < BT * OC) out[x] += bias[x % OC];
|
|
}
|
|
|
|
inline void matmul_forward(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight,
|
|
const float* __restrict__ bias,int B,int T, int C, int OC) {
|
|
int BT =B*T;
|
|
float alpha = 1.0f;
|
|
float beta = 0.0f;
|
|
cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, BT, C, &alpha, weight, C, inp, C, &beta, out, OC);
|
|
if(bias != NULL) {
|
|
int total = BT * OC;
|
|
bias_add_filler_kernel<<<(total + 255) / 256, 256>>>(out, bias, BT, OC);
|
|
}
|
|
}
|
|
|
|
// Launch kernel here
|
|
inline void matmul_forward_fused_gelu(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;
|
|
|
|
int total = BT * OC;
|
|
bias_add_filler_kernel<<<(total + 255) / 256, 256>>>(out, bias, BT, OC);
|
|
|
|
using ElementOutput = float;
|
|
using ElementAccumulator = float;
|
|
using ElementCompute = float;
|
|
|
|
using EpilogueOp = cutlass::epilogue::thread::LinearCombinationGELU<ElementOutput, 128/ cutlass::sizeof_bits<ElementOutput>::value, ElementAccumulator, ElementCompute>;
|
|
|
|
using Gemm = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
ElementOutput, cutlass::layout::RowMajor,
|
|
ElementAccumulator,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<128, 128, 32>,
|
|
cutlass::gemm::GemmShape<64, 64, 32>,
|
|
cutlass::gemm::GemmShape<16, 8, 8>,
|
|
EpilogueOp,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
2 // Stages
|
|
>;
|
|
|
|
Gemm gemm_op;
|
|
|
|
typename Gemm::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{out, OC}, // C matrix
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 1.0f} // alpha, beta for the epilogue
|
|
};
|
|
|
|
cutlass::Status status = gemm_op(args);
|
|
}
|
|
|
|
#endif // __MATMUL_KERNEL_CUH__ |