#include <stdio.h>
#include <assert.h>
#include "ATen/ATen.h"

template <typename F>
__global__ void kernel_forward(const int B, const int T, const int C, const int H,
                               float *__restrict__ _state, const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b,
                               F *__restrict__ const _y)
{
    const int bbb = blockIdx.x / H;  // batch index
    const int h = blockIdx.x % H;    // head index
    const int i = threadIdx.x;
    _state += bbb*C*_N_ + h*_N_*_N_ + i*_N_;  // now supports B > 1

    float state[_N_];
    #pragma unroll
    for (int j = 0; j < _N_; j++)
        state[j] = _state[j];

    __shared__ float r[_N_], k[_N_], w[_N_], a[_N_], b[_N_];

    for (int _t = 0; _t < T; _t++)
    {
        const int t = bbb*T*C + h*_N_ + i + _t * C;  // use bbb instead of e
        __syncthreads();
        r[i] = float(_r[t]);
        w[i] = __expf(-__expf(float(_w[t])));
        k[i] = float(_k[t]);
        a[i] = float(_a[t]);
        b[i] = float(_b[t]);
        __syncthreads();

        float sa = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            sa += a[j] * state[j];
        }

        float vv = float(_v[t]);
        float y = 0;
        #pragma unroll
        for (int j = 0; j < _N_; j++)
        {
            float& s = state[j];
            s = s * w[j] + k[j] * vv + sa * b[j];
            y += s * r[j];
        }
        _y[t] = F(y);
    }
    #pragma unroll
    for (int j = 0; j < _N_; j++)
        _state[j] = state[j];    
}

template <typename scalar_t>
void cuda_forward_template(int B, int T, int C, int H, float *state, scalar_t *r, scalar_t *w, scalar_t *k, scalar_t *v, scalar_t *a, scalar_t *b, scalar_t *y)
{
    assert(H*_N_ == C);
    // B > 1 is now supported
    kernel_forward<scalar_t><<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
}

void cuda_forward_fp16(int B, int T, int C, int H, float *state, at::Half *r, at::Half* w, at::Half *k, at::Half *v, at::Half *a, at::Half *b, at::Half *y)
{
    cuda_forward_template<at::Half>(B, T, C, H, state, r, w, k, v, a, b, y);
}

void cuda_forward_bf16(int B, int T, int C, int H, float *state, at::BFloat16 *r, at::BFloat16* w, at::BFloat16 *k, at::BFloat16 *v, at::BFloat16 *a, at::BFloat16 *b, at::BFloat16 *y)
{
    cuda_forward_template<at::BFloat16>(B, T, C, H, state, r, w, k, v, a, b, y);
}

void cuda_forward_fp32(int B, int T, int C, int H, float *state, float *r, float* w, float *k, float *v, float *a, float *b, float *y)
{
    cuda_forward_template<float>(B, T, C, H, state, r, w, k, v, a, b, y);
}
