From 01df41fbba703f951269dec3aefade5633f53c8f Mon Sep 17 00:00:00 2001
From: Andrej Karpathy <andrej.karpathy@gmail.com>
Date: Thu, 2 May 2024 19:24:04 +0000
Subject: [PATCH] careful to guard the overloaded warpReduceSum and small fixes

---
 train_gpt2.cu | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/train_gpt2.cu b/train_gpt2.cu
index f293b2b..00b7669 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -213,13 +213,6 @@ __device__ void atomicAddX(float* addr, float val) {
     atomicAdd(addr, val);
 }
 
-__device__ floatX warpReduceSum(floatX val) {
-    for (int offset = 16; offset > 0; offset /= 2) {
-        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
-    }
-    return val;
-}
-
 // warp-level reduction for summing values
 __device__ float warpReduceSum(float val) {
     for (int offset = 16; offset > 0; offset /= 2) {
@@ -236,6 +229,15 @@ __device__ float warpReduceMax(float val) {
     return val;
 }
 
+#if defined(ENABLE_BF16) || defined(ENABLE_FP16)
+__device__ floatX warpReduceSum(floatX val) {
+    for (int offset = 16; offset > 0; offset /= 2) {
+        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
+    }
+    return val;
+}
+#endif
+
 // ----------------------------------------------------------------------------
 // Packed128 data structure, which forces the compiler to use 128-bit loads/stores
 // in GPUs that support (the LDG.128 and STS.128 instructions)