optimised version of the cool new fused classifier + hopefully bugfixes?

dev/cuda/classifier_fused.cu

@@ -10,6 +10,7 @@ nvcc -O3 --use_fast_math classifier_fused.cu -o classifier_fused
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <float.h>
 #include <cuda_runtime.h>
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
@@ -72,7 +73,7 @@ void crossentropy_softmax_backward_cpu(float* dlogits,
             for (int i = 0; i < V; i++) {
                 float p = probs_bt[i];
                 float indicator = i == ix ? 1.0f : 0.0f;
-                dlogits_bt[i] += (p - indicator) * dloss;
+                dlogits_bt[i] = (p - indicator) * dloss;
             }
         }
     }
@@ -111,8 +112,7 @@ __device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp,
     return SoftmaxParams{norm, global_maxval};
 }
 
-
-__global__ void fused_classifier_kernel(float* dlogits, float* losses,
+__global__ void fused_classifier_kernel1(float* dlogits, float* losses,
                              const float* logits, const float* dlosses, const int* targets,
                              int B, int T, int V) {
     namespace cg = cooperative_groups;
@@ -139,15 +139,96 @@ __global__ void fused_classifier_kernel(float* dlogits, float* losses,
 
     // calculate all the gradients
     for (int i = warp.thread_rank(); i < V; i += warp.size()) {
-        float prob = expf(logits[i] - sp.Offset) * sp.Scale;
+        float prob = expf(logits[idx * V + i] - sp.Offset) * sp.Scale;
         float* dlogits_bt = dlogits + b * T * V + t * V;
         float dloss = dlosses[b * T + t];
         int ix = targets[b * T + t];
         float p = prob;
         float indicator = i == ix ? 1.0f : 0.0f;
-        dlogits_bt[i] += (p - indicator) * dloss;
+        dlogits_bt[i] = (p - indicator) * dloss;
+    }
+}
+
+
+__device__ float vec_at(const float4& vec, int index) {
+    return reinterpret_cast<const float*>(&vec)[index];
+}
+
+__device__ SoftmaxParams prepare_softmax2(cg::thread_block_tile<32>& warp,
+                                         int idx, const float* inp, int V) {
+    // one row of inp, i.e. inp[idx, :] of shape (V,)
+    // float4 to get 128-bit loads and memory level parallelism
+    // this is only possible if V is a multiple of 4
+    const float4* x_vec4 = reinterpret_cast<const float4*>(inp + idx * V);
+
+    float thread_maxval = -INFINITY;
+    float thread_sumval = 0.0f;
+    // do the loop in reverse to maximise probability of L2 cache hits
+    // so even small L2s get some hits on the 2nd read of the same thread
+    for (int i = (V/4) + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) {
+        float4 v = x_vec4[i];
+        #pragma unroll
+        for(int k = 0; k < 4; ++k) {
+            float old_maxval = thread_maxval;
+            thread_maxval = fmaxf(thread_maxval, vec_at(v, k));
+            thread_sumval *= expf((old_maxval - thread_maxval));
+            thread_sumval += expf(vec_at(v, k) - thread_maxval);
+        }
     }
 
+    // reduction in 2 stages: 1) inside warp, 2) between warps
+    // this results in much cleaner code than using a multi-warp cg::reduce
+    // todo benchmark to make sure it's faster, possibly too many reductions
+    // we could do the 2nd set of reductions per-block rather than per-warp
+    // but that would require an extra __syncthreads() unfortunately
+    __shared__ float shared_maxval[32];
+    __shared__ float shared_sumval[32];
+    int num_warps = blockDim.x / 32;
+    int warp_id = threadIdx.x / 32;
+    int lane_id = threadIdx.x % 32;
+
+    // reduce maxval
+    shared_maxval[warp_id] = cg::reduce(warp, thread_maxval, cg::greater<float>{});
+    __syncthreads();
+    float warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;
+    float block_maxval = cg::reduce(warp, warp_maxval, cg::greater<float>{});
+    // use maxval to scale sumval to avoid numerical instability / overflow
+    thread_sumval *= expf(thread_maxval - block_maxval);
+    // reduce sumval
+    shared_sumval[warp_id] = cg::reduce(warp, thread_sumval, cg::plus<float>{});
+    __syncthreads();
+    float warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.f;
+    float block_sumval = cg::reduce(warp, warp_sumval, cg::plus<float>{});
+    
+    return SoftmaxParams{1.f / block_sumval, block_maxval};
+}
+
+__global__ void fused_classifier_kernel2(float* dlogits, float* losses,
+                                         const float* logits, const float* dlosses, const int* targets,
+                                         int B, int T, int V) {
+    namespace cg = cooperative_groups;
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+    
+    int idx = blockIdx.x;
+    int ix = targets[idx];
+    float dloss = dlosses[idx];
+    
+    auto sp = prepare_softmax2(warp, idx, logits, V);
+
+    // calculate the probability needed for the loss and update.
+    // single-threaded
+    if(threadIdx.x == 0) {
+        float prob = expf(logits[idx * V + ix] - sp.Offset) * sp.Scale;
+        losses[idx] = -logf(prob);
+    }
+
+    // calculate all the gradients
+    for (int i = threadIdx.x; i < V; i += blockDim.x) {
+        float prob = expf(__ldcs(&logits[idx * V + i]) - sp.Offset) * sp.Scale;
+        float indicator = i == ix ? 1.0f : 0.0f;
+        dlogits[idx * V + i] = (prob - indicator) * dloss;
+    }
 }
 
 // ----------------------------------------------------------------------------
@@ -157,8 +238,17 @@ void fused_classifier1(float* dlogits, float* losses,
                       const float* logits, const float* dlosses, const int* targets,
                       int B, int T, int V, int block_size) {
     const int N = B * T;
+    const int grid_size = N / (block_size / 32);
+    fused_classifier_kernel1<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V);
+    cudaCheck(cudaGetLastError());
+}
+
+void fused_classifier2(float* dlogits, float* losses,
+                      const float* logits, const float* dlosses, const int* targets,
+                      int B, int T, int V, int block_size) {
+    const int N = B * T;
     const int grid_size = N;
-    fused_classifier_kernel<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V);
+    fused_classifier_kernel2<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V);
     cudaCheck(cudaGetLastError());
 }
 
@@ -169,6 +259,13 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses,
         case 1:
             fused_classifier1(dlogits, losses, logits, dlosses, targets, B, T, V, block_size);
             break;
+        case 2:
+            if((V % 4) != 0) {
+                printf("V needs to be a multiple of 4 for this kernel to work!\n");
+                exit(EXIT_FAILURE);
+            }
+            fused_classifier2(dlogits, losses, logits, dlosses, targets, B, T, V, block_size);
+            break;
         default:
             printf("Invalid kernel number\n");
             exit(1);
@@ -182,7 +279,7 @@ int main(int argc, char **argv) {
 
     int B = 8;
     int T = 1024;
-    int V = 50257;
+    int V = 50272; // rounded up from 50257 for alignment purposes
 
     int deviceIdx = 0;
     cudaCheck(cudaSetDevice(deviceIdx));