#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define TILE_DIM 16

// Kernel to compute y = relu(x @ weight^T + bias)
// Let A = x (shape: N x K), weight (shape: M x K), bias (shape: M)
// Then, C[i, j] = dot(A[i, :], weight[j, :]) + bias[j] for i in [0, N) and j in [0, M)
// This kernel performs a tiled GEMM using shared memory tiling. The weight matrix is stored 
// in row-major order (M x K) and is implicitly transposed during multiplication.
__global__ void gemm_bias_relu_kernel(const float* __restrict__ A,
                                      const float* __restrict__ weight,
                                      const float* __restrict__ bias,
                                      float* __restrict__ C,
                                      int N, // batch size
                                      int K, // in_features
                                      int M) // out_features
{
    // Compute output row and column.
    int row = blockIdx.y * TILE_DIM + threadIdx.y;  // in A, row index
    int col = blockIdx.x * TILE_DIM + threadIdx.x;    // corresponds to weight row index & output column

    float sum = 0.0f;
    
    // Loop over tiles.
    for (int t = 0; t < (K + TILE_DIM - 1) / TILE_DIM; t++) {
        __shared__ float sA[TILE_DIM][TILE_DIM];
        __shared__ float sB[TILE_DIM][TILE_DIM];
        
        // Load tile from A.
        int aCol = t * TILE_DIM + threadIdx.x;
        if (row < N && aCol < K) {
            sA[threadIdx.y][threadIdx.x] = A[row * K + aCol];
        } else {
            sA[threadIdx.y][threadIdx.x] = 0.0f;
        }
        
        // Load tile from weight.
        // Weight is stored as (M x K) in row-major order.
        int wCol = t * TILE_DIM + threadIdx.y;
        if (col < M && wCol < K) {
            sB[threadIdx.y][threadIdx.x] = weight[col * K + wCol];
        } else {
            sB[threadIdx.y][threadIdx.x] = 0.0f;
        }
        
        __syncthreads();
        
        // Multiply tiles.
        #pragma unroll
        for (int i = 0; i < TILE_DIM; i++) {
            sum += sA[threadIdx.y][i] * sB[i][threadIdx.x];
        }
        __syncthreads();
    }
    
    // Write the output.
    if (row < N && col < M) {
        // Use __ldg to load bias from read-only cache.
        float result = sum + __ldg(&bias[col]);
        C[row * M + col] = (result > 0.0f) ? result : 0.0f;
    }
}

torch::Tensor forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    // Ensure we're using the correct device.
    c10::cuda::CUDAGuard guard(x.device());
    
    int N = x.size(0);       // batch size
    int K = x.size(1);       // in_features
    int M = weight.size(0);  // out_features
    
    auto output = torch::empty({N, M}, x.options());
    
    dim3 block(TILE_DIM, TILE_DIM);
    dim3 grid((M + TILE_DIM - 1) / TILE_DIM, (N + TILE_DIM - 1) / TILE_DIM);
    
    gemm_bias_relu_kernel<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        N, K, M
    );
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused GEMM, bias addition, and ReLU using tiled GEMM with unrolling and __ldg optimization");
}