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

#define TILE 16

// Optimized double-buffered tiled kernel for fused Linear (GEMM + bias + optional ReLU).
// Computes: C = A * W^T + bias
//  - A is (B, K) (input matrix)
//  - W is (N, K) (weight matrix, row-major; each row corresponds to the weights of an output neuron)
//  - bias is (N)
//  - C is (B, N) computed as: C[b,n] = sum_{k=0}^{K-1} A[b,k]*W[n,k] + bias[n]
// Double buffering overlapped global loads with computation.
__global__ void fused_linear_tiled_kernel(const float* __restrict__ A,
                                          const float* __restrict__ W,
                                          const float* __restrict__ bias,
                                          float* __restrict__ C,
                                          int B, int K, int N,
                                          int apply_relu) {
    // Calculate output element indices.
    int row = blockIdx.y * TILE + threadIdx.y;  // batch index
    int col = blockIdx.x * TILE + threadIdx.x;  // neuron index

    float sum = 0.0f;
    
    // Declare double-buffered shared memory for tiles from A and W.
    __shared__ float As_db[2][TILE][TILE];
    __shared__ float Bs_db[2][TILE][TILE];
    
    int num_tiles = (K + TILE - 1) / TILE;
    
    if (num_tiles > 0) {
        // Load the first tile into buffer 0.
        int t = 0;
        int a_index = t * TILE + threadIdx.x;
        if (row < B && a_index < K)
            As_db[0][threadIdx.y][threadIdx.x] = A[row * K + a_index];
        else
            As_db[0][threadIdx.y][threadIdx.x] = 0.0f;
            
        int w_index = t * TILE + threadIdx.y;
        if (col < N && w_index < K)
            Bs_db[0][threadIdx.y][threadIdx.x] = W[col * K + w_index];
        else
            Bs_db[0][threadIdx.y][threadIdx.x] = 0.0f;
        
        __syncthreads();
        
        int curr = 0;
        // Loop over remaining tiles with double buffering.
        for (t = 0; t < num_tiles - 1; t++) {
            int next = (t + 1) & 1;  // alternate buffer index (0 or 1)
            // Preload next tile into buffer "next".
            int a_index_next = (t + 1) * TILE + threadIdx.x;
            if (row < B && a_index_next < K)
                As_db[next][threadIdx.y][threadIdx.x] = A[row * K + a_index_next];
            else
                As_db[next][threadIdx.y][threadIdx.x] = 0.0f;
            
            int w_index_next = (t + 1) * TILE + threadIdx.y;
            if (col < N && w_index_next < K)
                Bs_db[next][threadIdx.y][threadIdx.x] = W[col * K + w_index_next];
            else
                Bs_db[next][threadIdx.y][threadIdx.x] = 0.0f;
            
            __syncthreads();
            // Compute partial dot-product using current buffer.
            #pragma unroll
            for (int i = 0; i < TILE; i++) {
                sum += As_db[curr][threadIdx.y][i] * Bs_db[curr][i][threadIdx.x];
            }
            curr = next;  // swap buffer
            __syncthreads();
        }
        // Process the final tile which is now in buffer "curr".
        #pragma unroll
        for (int i = 0; i < TILE; i++) {
            sum += As_db[curr][threadIdx.y][i] * Bs_db[curr][i][threadIdx.x];
        }
    }
    
    // Write the computed result, adding bias and applying ReLU if required.
    if (row < B && col < N) {
        sum += bias[col];
        if (apply_relu)
            sum = fmaxf(sum, 0.0f);
        C[row * N + col] = sum;
    }
}

torch::Tensor fused_linear(torch::Tensor x,
                           torch::Tensor weight,
                           torch::Tensor bias,
                           int apply_relu) {
    // Ensure we operate on the correct CUDA device.
    c10::cuda::CUDAGuard device_guard(x.device());
    
    int B = x.size(0);  // Batch size
    int K = x.size(1);  // Input dimension
    int N = weight.size(0);  // Number of output neurons
    
    auto out = torch::empty({B, N}, x.options());
    
    // Configure grid and block dimensions.
    dim3 block(TILE, TILE);
    dim3 grid((N + TILE - 1) / TILE, (B + TILE - 1) / TILE);
    
    fused_linear_tiled_kernel<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        B, K, N,
        apply_relu
    );
    
    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_linear", &fused_linear, "Fused Linear (GEMM + bias + optional ReLU) operation using double-buffered tiled shared-memory kernel");
}