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

// Scalar fused activation kernel (in-place).
// For each element at index idx in the GEMM output tensor (shape: [batch, out_features]):
//   y = gemm[idx] + bias[col]  (where col = idx % out_features),
//   then y *= scaling_factor,
//   then y is clamped: y = clamp(y, hardtanh_min, hardtanh_max),
//   and finally y is transformed using approximate GELU.
__global__ void fused_activation_kernel(float* data,
                                          const float* __restrict__ bias,
                                          int out_features,
                                          int total,
                                          float scaling_factor,
                                          float hardtanh_min,
                                          float hardtanh_max) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int col = idx % out_features;
        float y = data[idx] + bias[col];
        y *= scaling_factor;
        y = fminf(fmaxf(y, hardtanh_min), hardtanh_max);
        const float sqrt_2_over_pi = 0.7978845608028654f;
        float y_cube = y * y * y;
        float gelu = 0.5f * y * (1.0f + tanhf(sqrt_2_over_pi * (y + 0.044715f * y_cube)));
        data[idx] = gelu;
    }
}

// Vectorized fused activation kernel (in-place) using float4.
// Assumptions:
//   - total number of elements (batch * out_features) is divisible by 4,
//   - out_features is divisible by 4 so that the bias can be processed as float4.
__global__ void fused_activation_kernel_vec(float4* data,
                                              const float4* __restrict__ bias_vec,
                                              int out_features,
                                              int vec_total,
                                              float scaling_factor,
                                              float hardtanh_min,
                                              float hardtanh_max) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < vec_total) {
        int vec_out_features = out_features / 4;
        // Determine the vector-column index in the bias.
        int row = idx / vec_out_features;
        int col = idx % vec_out_features;  // index within one row (vectorized)
        float4 in_val = data[idx];
        float4 bias_val = bias_vec[col];    // bias broadcast along batch
        
        float4 out_val;
        const float sqrt_2_over_pi = 0.7978845608028654f;
        
        // Process x component.
        float y = in_val.x + bias_val.x;
        y *= scaling_factor;
        y = fminf(fmaxf(y, hardtanh_min), hardtanh_max);
        float y_cube = y * y * y;
        out_val.x = 0.5f * y * (1.0f + tanhf(sqrt_2_over_pi * (y + 0.044715f * y_cube)));
        
        // Process y component.
        y = in_val.y + bias_val.y;
        y *= scaling_factor;
        y = fminf(fmaxf(y, hardtanh_min), hardtanh_max);
        y_cube = y * y * y;
        out_val.y = 0.5f * y * (1.0f + tanhf(sqrt_2_over_pi * (y + 0.044715f * y_cube)));
        
        // Process z component.
        y = in_val.z + bias_val.z;
        y *= scaling_factor;
        y = fminf(fmaxf(y, hardtanh_min), hardtanh_max);
        y_cube = y * y * y;
        out_val.z = 0.5f * y * (1.0f + tanhf(sqrt_2_over_pi * (y + 0.044715f * y_cube)));
        
        // Process w component.
        y = in_val.w + bias_val.w;
        y *= scaling_factor;
        y = fminf(fmaxf(y, hardtanh_min), hardtanh_max);
        y_cube = y * y * y;
        out_val.w = 0.5f * y * (1.0f + tanhf(sqrt_2_over_pi * (y + 0.044715f * y_cube)));
        
        data[idx] = out_val;
    }
}

torch::Tensor forward(torch::Tensor x,
                      torch::Tensor weight,
                      torch::Tensor bias,
                      float scaling_factor,
                      float hardtanh_min,
                      float hardtanh_max) {
    // Ensure operations execute on the correct device.
    c10::cuda::CUDAGuard device_guard(x.device());
    
    // Compute GEMM using at::mm: gemm = x * weight^T.
    // weight shape: [out_features, in_features]
    auto gemm = at::mm(x, weight.t()); // gemm shape: [batch, out_features]
    int total = gemm.numel();
    int out_features = gemm.size(1);
    
    // Get current CUDA stream.
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    int threads = 256;
    
    // Use vectorized kernel if total elements and out_features are divisible by 4.
    if ((total % 4 == 0) && (out_features % 4 == 0)) {
        int vec_total = total / 4;
        int blocks = (vec_total + threads - 1) / threads;
        const float4* bias_vec = reinterpret_cast<const float4*>(bias.data_ptr<float>());
        float4* data_vec = reinterpret_cast<float4*>(gemm.data_ptr<float>());
        fused_activation_kernel_vec<<<blocks, threads, 0, stream>>>(
            data_vec,
            bias_vec,
            out_features,
            vec_total,
            scaling_factor,
            hardtanh_min,
            hardtanh_max
        );
    } else {
        int blocks = (total + threads - 1) / threads;
        fused_activation_kernel<<<blocks, threads, 0, stream>>>(
            gemm.data_ptr<float>(),
            bias.data_ptr<float>(),
            out_features,
            total,
            scaling_factor,
            hardtanh_min,
            hardtanh_max
        );
    }
    
    return gemm;  // gemm is updated in-place.
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused GEMM with in-place bias addition, scaling, hardtanh, and GELU activation (vectorized when possible)");
}