215 lines
7.7 KiB
Plaintext
215 lines
7.7 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
|
|
|
|
//removing bias overhead and fusing with cutlass
|
|
// __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];
|
|
// }
|
|
|
|
//using big cutlass for kernels for BT>=64
|
|
using GemmDefault = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<128,128,32>,
|
|
cutlass::gemm::GemmShape<64,64,32>,
|
|
cutlass::gemm::GemmShape<16,8,8>,
|
|
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
3 // Stages
|
|
>;
|
|
|
|
//small cutlass for kernels for BT<64, we can use smaller threadblocks and warps to increase occupancy
|
|
using GemmsmallDefault = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<32,128,32>,
|
|
cutlass::gemm::GemmShape<32,64,32>,
|
|
cutlass::gemm::GemmShape<16,8,8>,
|
|
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
2
|
|
>;
|
|
|
|
using GemmMediumDefault = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<64,128,32>,
|
|
cutlass::gemm::GemmShape<32,64,32>,
|
|
cutlass::gemm::GemmShape<16,8,8>,
|
|
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
3
|
|
>;
|
|
|
|
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;
|
|
|
|
if(bias != NULL){
|
|
if(BT <= 32){
|
|
GemmsmallDefault gemm_op;
|
|
typename GemmsmallDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{bias, 0}, // C matrix
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 1.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
else if(BT <= 128){
|
|
GemmMediumDefault gemm_op;
|
|
typename GemmMediumDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{bias, 0}, // C matrix
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 1.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
else {
|
|
GemmDefault gemm_op;
|
|
typename GemmDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{bias, 0}, // C matrix
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 1.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
}
|
|
else{
|
|
if(BT <= 32){
|
|
GemmsmallDefault gemm_op;
|
|
typename GemmsmallDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{out, OC}, // C matrix (output)
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 0.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
else if(BT <= 128){
|
|
GemmDefault gemm_op;
|
|
typename GemmDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{out, OC}, // C matrix (output)
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 0.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
else{
|
|
GemmDefault gemm_op;
|
|
typename GemmDefault::Arguments args{
|
|
{BT, OC, C}, // problem size
|
|
{inp, C}, // A matrix
|
|
{weight, C}, // B matrix
|
|
{out, OC}, // C matrix (output)
|
|
{out, OC}, // D matrix (output)
|
|
{1.0f, 0.0f} // alpha, beta for the epilogue
|
|
};
|
|
gemm_op(args);
|
|
}
|
|
}
|
|
|
|
|
|
}
|
|
|
|
using Gemmgelu = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<128, 128, 32>,
|
|
cutlass::gemm::GemmShape<64, 64, 32>,
|
|
cutlass::gemm::GemmShape<16, 8, 8>,
|
|
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
3>;
|
|
|
|
using GemmgeluSmall = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<32, 128, 32>,
|
|
cutlass::gemm::GemmShape<32, 64, 32>,
|
|
cutlass::gemm::GemmShape<16, 8, 8>,
|
|
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
2>;
|
|
|
|
using GemmgeluMedium = cutlass::gemm::device::Gemm<
|
|
float, cutlass::layout::RowMajor,
|
|
float, cutlass::layout::ColumnMajor,
|
|
float, cutlass::layout::RowMajor,
|
|
float,
|
|
cutlass::arch::OpClassTensorOp,
|
|
cutlass::arch::Sm80,
|
|
cutlass::gemm::GemmShape<64, 128, 32>,
|
|
cutlass::gemm::GemmShape<32, 64, 32>,
|
|
cutlass::gemm::GemmShape<16, 8, 8>,
|
|
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
|
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
|
3>;
|
|
|
|
// 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;
|
|
|
|
if (BT <= 32) {
|
|
GemmgeluSmall op;
|
|
typename GemmgeluSmall::Arguments args{
|
|
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
|
op(args);
|
|
} else if (BT <= 128) {
|
|
GemmgeluMedium op;
|
|
typename GemmgeluMedium::Arguments args{
|
|
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
|
op(args);
|
|
} else {
|
|
Gemmgelu op;
|
|
typename Gemmgelu::Arguments args{
|
|
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
|
op(args);
|
|
}
|
|
}
|
|
|
|
#endif // __MATMUL_KERNEL_CUH__ |