From f26cf00a615e4f643a3660071e4ef86010045db7 Mon Sep 17 00:00:00 2001
From: Andrej Karpathy <andrej.karpathy@gmail.com>
Date: Thu, 11 Apr 2024 15:49:01 +0000
Subject: [PATCH] make gelu constant be a #define, speeds up the kernel by 1%
 or so

---
 dev/cuda/gelu_forward.cu | 8 ++++----
 train_gpt2.c             | 9 ++++-----
 train_gpt2.cu            | 4 ++--
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu
index fad7be8..436b5af 100644
--- a/dev/cuda/gelu_forward.cu
+++ b/dev/cuda/gelu_forward.cu
@@ -38,12 +38,13 @@ void cudaCheck(cudaError_t error, const char *file, int line) {
 // ----------------------------------------------------------------------------
 // CPU code reference
 
+#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
+
 void gelu_forward_cpu(float* out, float* inp, int N) {
-    float s = sqrtf(2.0f / M_PI);
     for (int i = 0; i < N; i++) {
         float x = inp[i];
         float cube = 0.044715f * x * x * x;
-        out[i] = 0.5f * x * (1.0f + tanhf(s * (x + cube)));
+        out[i] = 0.5f * x * (1.0f + tanhf(GELU_SCALING_FACTOR * (x + cube)));
     }
 }
 
@@ -53,11 +54,10 @@ void gelu_forward_cpu(float* out, float* inp, int N) {
 // elementwise ops are nice and ez
 __global__ void gelu_kernel(float* out, const float* inp, int N) {
     int i = blockIdx.x * blockDim.x + threadIdx.x;
-    float s = sqrtf(2.0f / M_PI);
     if (i < N) {
         float xi = inp[i];
         float cube = 0.044715f * xi * xi * xi;
-        out[i] = 0.5f * xi * (1.0f + tanhf(s * (xi + cube)));
+        out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
     }
 }
 
diff --git a/train_gpt2.c b/train_gpt2.c
index 830085c..a54a7b9 100644
--- a/train_gpt2.c
+++ b/train_gpt2.c
@@ -335,25 +335,24 @@ void attention_backward(float* dinp, float* dpreatt, float* datt,
     }
 }
 
+#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
 void gelu_forward(float* out, float* inp, int N) {
-    float s = sqrtf(2.0f / M_PI);
     for (int i = 0; i < N; i++) {
         float x = inp[i];
         float cube = 0.044715f * x * x * x;
-        out[i] = 0.5f * x * (1.0f + tanhf(s * (x + cube)));
+        out[i] = 0.5f * x * (1.0f + tanhf(GELU_SCALING_FACTOR * (x + cube)));
     }
 }
 
 void gelu_backward(float* dinp, float* inp, float* dout, int N) {
-    float s = sqrtf(2.0f / M_PI);
     for (int i = 0; i < N; i++) {
         float x = inp[i];
         float cube = 0.044715f * x * x * x;
-        float tanh_arg = s * (x + cube);
+        float tanh_arg = GELU_SCALING_FACTOR * (x + cube);
         float tanh_out = tanhf(tanh_arg);
         float coshf_out = coshf(tanh_arg);
         float sech_out = 1.0f / (coshf_out * coshf_out);
-        float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * s * (1.0f + 3.0f * 0.044715f * x * x);
+        float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x);
         dinp[i] += local_grad * dout[i];
     }
 }
diff --git a/train_gpt2.cu b/train_gpt2.cu
index c7d0992..0a2cb46 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -305,13 +305,13 @@ __global__ void residual_forward_kernel(float* out, float* inp1, float* inp2, in
     }
 }
 
+#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
 __global__ void gelu_kernel(float* out, const float* inp, int N) {
     int i = blockIdx.x * blockDim.x + threadIdx.x;
-    float s = sqrtf(2.0f / M_PI);
     if (i < N) {
         float xi = inp[i];
         float cube = 0.044715f * xi * xi * xi;
-        out[i] = 0.5f * xi * (1.0f + tanhf(s * (xi + cube)));
+        out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
     }
 }