comp files

kernels/softmax.cuh

@@ -6,30 +6,93 @@
 #include <float.h>
 #include "../utils/cuda_utils.cuh"
 
+#define BLOCK_DIM 256
+#define WARP_SIZE 32
+
+// lecture 15 slide 32
+__device__ __inline__ float warpReduceSum(float partialSum) {
+    for (unsigned int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
+        partialSum += __shfl_down_sync(0xffffffff, partialSum, stride);
+    }
+    return partialSum;
+}
+
+// same pattern for max
+__device__ __inline__ float warpReduceMax(float partialMax) {
+    for (unsigned int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
+        partialMax = fmaxf(partialMax, __shfl_down_sync(0xffffffff, partialMax, stride));
+    }
+    return partialMax;
+}
+
 __global__ void softmax_forward_kernel(float* out, float inv_temperature, const float* inp, int N, int T) {
     // Implement this
-    int x = blockDim.x * blockIdx.x + threadIdx.x;
+    int row = blockIdx.x;
+    if(row >= N*T){
+        return;
+    }
 
-    if(x < N*T){
-        int own_pos = x % T;
-        const float* a = inp + x * T;
+    int own_pos = row % T;
+    const float* a = inp + row * T;
+    float* o = out + row * T;
 
-        float maxval = -FLT_MAX;
-        for (int i = 0; i <= own_pos; i++) {
-            maxval = fmaxf(maxval,a[i]);
-        }
+    //find max
+    //thread coarsening
+    float partial_max = -FLT_MAX;
+    for(int i = threadIdx.x; i <= own_pos; i+=BLOCK_DIM){
+        partial_max = fmaxf(partial_max,a[i]);
+    }
 
-        float sumval = 0.0f;
-        for (int i = 0; i <= own_pos; ++i) {
-            float ev = expf(inv_temperature * (a[i] - maxval));
-            sumval += ev;
-            out[x * T + i] = ev;
-        }
+    //lecture 15 slide 38
+    partial_max = warpReduceMax(partial_max);
 
-        float norm = 1.0f / sumval;
-        for (int i = 0; i <= own_pos; ++i) {
-            out[x * T + i] *= norm;
-        }
+    __shared__ float partial_max_s[BLOCK_DIM / WARP_SIZE];
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_max_s[threadIdx.x / WARP_SIZE] = partial_max;
+    }
+    __syncthreads();
+    
+    if (threadIdx.x < WARP_SIZE) {
+        partial_max = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_max_s[threadIdx.x] : -FLT_MAX;
+        partial_max = warpReduceMax(partial_max);
+    }
+
+    __shared__ float s_max;
+    if (threadIdx.x == 0) {
+        s_max = partial_max;
+    }
+    __syncthreads();
+    float maxval = s_max;
+
+    float partial_sum = 0.0f;
+    for (int i = threadIdx.x; i <= own_pos; i += BLOCK_DIM) {
+        float ev = expf(inv_temperature * (a[i] - maxval));
+        o[i] = ev;
+        partial_sum += ev;
+    }
+
+    partial_sum = warpReduceSum(partial_sum);
+
+    __shared__ float partial_sums_s[BLOCK_DIM / WARP_SIZE];
+    if (threadIdx.x % WARP_SIZE == 0) {
+        partial_sums_s[threadIdx.x / WARP_SIZE] = partial_sum;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < WARP_SIZE) {
+        partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sums_s[threadIdx.x] : 0.0f;
+        partial_sum = warpReduceSum(partial_sum);
+    }
+
+    __shared__ float s_sum;
+    if (threadIdx.x == 0) {
+        s_sum = partial_sum;
+    }
+    __syncthreads();
+    float norm = 1.0f / s_sum;
+
+    for (int i = threadIdx.x; i <= own_pos; i += BLOCK_DIM) {
+        o[i] *= norm;
     }
 }