#include <cuda_bf16.h>
#include "../common.cuh"

// Naive GEMM: each thread computes one element of C.
// C[i][j] = sum_k A[i][k] * B[k][j]
// All matrices are row-major.
__global__ void gemm_naive_bf16(
    const __nv_bfloat16* A, const __nv_bfloat16* B, __nv_bfloat16* C,
    int M, int N, int K
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
        }
        C[row * N + col] = __float2bfloat16(sum);
    }
}

__global__ void gemm_naive_f32(
    const float* A, const float* B, float* C,
    int M, int N, int K
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

extern "C" {

void launch_gemm_naive_bf16(
    const void* A, const void* B, void* C,
    int M, int N, int K, void* stream
) {
    dim3 block(16, 16);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    gemm_naive_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const __nv_bfloat16*)A, (const __nv_bfloat16*)B, (__nv_bfloat16*)C, M, N, K
    );
    CUDA_CHECK_LAST_ERROR();
}

void launch_gemm_naive_f32(
    const void* A, const void* B, void* C,
    int M, int N, int K, void* stream
) {
    dim3 block(16, 16);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    gemm_naive_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
        (const float*)A, (const float*)B, (float*)C, M, N, K
    );
    CUDA_CHECK_LAST_ERROR();
}

} // extern "C"