#ifndef ENCODER_FORWARD_KERNEL_CUH
#define ENCODER_FORWARD_KERNEL_CUH

#include <cuda_runtime.h>
#include "../utils/cuda_utils.cuh"

__global__ void encoder_forward_kernel(float* __restrict__ out, const int* __restrict__ inp, const float* __restrict__ wte, const float* __restrict__ wpe,
                                       int B, int T, int C) {
    // Implement this
    int row = blockIdx.x * blockDim.y + threadIdx.y;
    if (row >= B * T){
        return;
    }
    int b = row / T;
    int t = row % T;
    int token = inp[b * T + t];
    const float* __restrict__ wte_ptr = wte + token * C;
    const float* __restrict__ wpe_ptr = wpe + t * C;
    float* __restrict__ out_ptr = out + row * C;
    for (int i = threadIdx.x; i < C; i += 32) {
        out_ptr[i] = wte_ptr[i] + wpe_ptr[i];
    }
}

// Launch kernel here
void encoder_forward(float* out, const int* inp, const float* wte, const float* wpe, int B, int T, int C) {
    // Implement this
    dim3 block(32, 8, 1);
    int total_rows = B * T;
    dim3 grid((total_rows + 8 - 1) / 8, 1, 1);
    encoder_forward_kernel<<<grid, block>>>(out, inp, wte, wpe, B, T, C);
}


#endif // ENCODER_FORWARD_KERNEL_CUH