From 327eef3f22d9bb0954a0091e49d46cb2b5a1fa95 Mon Sep 17 00:00:00 2001
From: Andrej Karpathy <andrej.karpathy@gmail.com>
Date: Sun, 28 Apr 2024 20:10:33 +0000
Subject: [PATCH] incorporate faster encoder_forward kernel to fp32 CUDA
 version

---
 dev/cuda/encoder_forward.cu | 13 +++++++------
 train_gpt2_fp32.cu          | 36 +++++++++++++++++++-----------------
 2 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu
index 5ab32bc..f56a9c8 100644
--- a/dev/cuda/encoder_forward.cu
+++ b/dev/cuda/encoder_forward.cu
@@ -9,6 +9,9 @@ version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops ov
 
 version 2 is more optimized, parallelizes over all of B,T,C
 ./encoder_forward 2
+
+version 3 is like version 2 but uses float4 reads/writes
+./encoder_forward 3
 */
 
 #include <stdio.h>
@@ -82,27 +85,25 @@ __global__ void encoder_forward_kernel2(float* out,
     }
 }
 
-// use float4 for memory coalescing
-__device__ inline float4 operator+(const float4& a, const float4& b) {
+__device__ inline float4 add_float4(const float4& a, const float4& b) {
     return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
 }
 
+// use of float4 leads to using 128-bit LDG / STG instructions in SASS,
+// very helpful in memory-bound kernels like encoder_forward
 __global__ void encoder_forward_kernel3(float4* out,
                                const int* inp, const float4* wte, const float4* wpe,
                                int B, int T, int C) {
     int C4 = C / 4;
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     int N = B * T * C4;
-
     if (idx < N) {
         int bt = idx / C4;
         int b = bt / T;
         int t = bt % T;
         int c4 = idx % C4;
-
         int ix = inp[b * T + t];
-
-        out[b * T * C4 + t * C4 + c4] = wte[ix * C4 + c4] + wpe[t * C4 + c4];
+        out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]);
     }
 }
 
diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu
index 19bc574..8b90e3c 100644
--- a/train_gpt2_fp32.cu
+++ b/train_gpt2_fp32.cu
@@ -153,24 +153,25 @@ __device__ float warpReduceSum(float val) {
     return val;
 }
 
-__global__ void encoder_forward_kernel2(float* out,
-                               int* inp, float* wte, float* wpe,
-                               int B, int T, int C) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int N = B * T * C;
+__device__ inline float4 add_float4(const float4& a, const float4& b) {
+    return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
+}
 
+// use of float4 leads to using 128-bit LDG / STG instructions in SASS,
+// very helpful in memory-bound kernels like encoder_forward
+__global__ void encoder_forward_kernel3(float4* out,
+                               const int* inp, const float4* wte, const float4* wpe,
+                               int B, int T, int C) {
+    int C4 = C / 4;
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int N = B * T * C4;
     if (idx < N) {
-        int bt = idx / C;
+        int bt = idx / C4;
         int b = bt / T;
         int t = bt % T;
-        int c = idx % C;
-
+        int c4 = idx % C4;
         int ix = inp[b * T + t];
-
-        float* out_btc = out + b * T * C + t * C + c;
-        float* wte_ix = wte + ix * C + c;
-        float* wpe_tc = wpe + t * C + c;
-        *out_btc = *wte_ix + *wpe_tc;
+        out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]);
     }
 }
 
@@ -812,12 +813,13 @@ __global__ void fused_classifier_kernel3(float* logits, float* losses, float* pr
 // kernel launchers
 
 void encoder_forward(float* out,
-                     int* inp, float* wte, float* wpe,
+                     const int* inp, const float* wte, const float* wpe,
                      int B, int T, int C) {
+    assert(C % 4 == 0);
+    const int block_size = 512;
     const int N = B * T * C;
-    const int block_size = 256;
-    const int grid_size = CEIL_DIV(N, block_size);
-    encoder_forward_kernel2<<<grid_size, block_size>>>(out, inp, wte, wpe, B, T, C);
+    const int grid_size = CEIL_DIV(N / 4, block_size);
+    encoder_forward_kernel3<<<grid_size, block_size>>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C);
     cudaCheck(cudaGetLastError());
 }