
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

__global__ void sidrc_kernel_batched_fp16(
    const at::Half* __restrict__ A,
    const at::Half* __restrict__ B,
    at::Half* __restrict__ C,
    int M_blocks, int K_blocks, int N_blocks) 
{
    int bx = blockIdx.x; 
    int by = blockIdx.y; 
    int bz = blockIdx.z; 
    int tx = threadIdx.x; 

    if (bx >= N_blocks || by >= M_blocks) return;

    int batch_stride_A = M_blocks * K_blocks * 1024;
    int batch_stride_B = K_blocks * N_blocks * 1024;
    int batch_stride_C = M_blocks * N_blocks * 1024;

    const at::Half* batch_A = A + (bz * batch_stride_A);
    const at::Half* batch_B = B + (bz * batch_stride_B);
    at::Half* batch_C = C + (bz * batch_stride_C);

    float c_accum[32] = {0.0f};

    for (int k = 0; k < K_blocks; ++k) {
        const at::Half* tile_A = batch_A + (by * K_blocks + k) * 1024;
        const at::Half* tile_B = batch_B + (k * N_blocks + bx) * 1024;

        float rA[32];
        float rB[32];

        #pragma unroll
        for(int i=0; i<32; i+=4) {
            rA[i]   = __half2float(tile_A[(i)   * 32 + tx]);
            rA[i+1] = __half2float(tile_A[(i+1) * 32 + tx]);
            rA[i+2] = __half2float(tile_A[(i+2) * 32 + tx]);
            rA[i+3] = __half2float(tile_A[(i+3) * 32 + tx]);

            rB[i]   = __half2float(tile_B[(i)   * 32 + tx]);
            rB[i+1] = __half2float(tile_B[(i+1) * 32 + tx]);
            rB[i+2] = __half2float(tile_B[(i+2) * 32 + tx]);
            rB[i+3] = __half2float(tile_B[(i+3) * 32 + tx]);
        }

        #pragma unroll
        for (int d = 0; d < 32; ++d) {
            float local_sum = 0.0f;
            #pragma unroll
            for (int s = 0; s < 32; ++s) {
                int q = (d - s + 32) % 32;
                float b_val = rB[q];
                int src_lane = (tx + s) % 32;
                float shifted_b = __shfl_sync(0xffffffff, b_val, src_lane);
                local_sum += rA[s] * shifted_b;
            }
            c_accum[d] += local_sum;
        }
    }

    at::Half* tile_C = batch_C + (by * N_blocks + bx) * 1024;
    #pragma unroll
    for (int d = 0; d < 32; ++d) {
        tile_C[d * 32 + tx] = __float2half(c_accum[d]);
    }
}

torch::Tensor sidrc_forward_cuda(torch::Tensor A, torch::Tensor B) {
    int Batch    = A.size(0);
    int M_blocks = A.size(1); 
    int K_blocks = A.size(2);
    int N_blocks = B.size(2);
    
    auto C = torch::zeros({Batch, M_blocks, N_blocks, 32, 32}, A.options());

    dim3 grid(N_blocks, M_blocks, Batch);
    dim3 block(32);

    sidrc_kernel_batched_fp16<<<grid, block>>>(
        A.data_ptr<at::Half>(),
        B.data_ptr<at::Half>(),
        C.data_ptr<at::Half>(),
        M_blocks, K_blocks, N_blocks
    );
    return C;
}
