#ifndef __SOFTMAX_KERNEL_CUH__
#define __SOFTMAX_KERNEL_CUH__

#include <cuda_runtime.h>
#include <math.h>
#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 row = blockIdx.x;
    if(row >= N*T){
        return;
    }

    int own_pos = row % T;
    const float* a = inp + row * T;
    float* o = out + row * T;

    //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]);
    }

    //lecture 15 slide 38
    partial_max = warpReduceMax(partial_max);

    __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;
    }
}

#endif // __SOFTMAX_KERNEL_CUH__