ECE 408 sem project

kernels_op_16/matmul.cuh

@@ -0,0 +1,73 @@
+#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__