import torch
from torch.utils.cpp_extension import load_inline

# CUDA source for fused FMA + Activation
fma_activation_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void fma_relu_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    const float* __restrict__ c,
    float* __restrict__ output,
    int total_elements
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < total_elements) {
        // Fused multiply-add with ReLU activation
        float fma_result = a[idx] * b[idx] + c[idx];
        output[idx] = fmaxf(0.0f, fma_result);  // ReLU
    }
}

// Kernel for fused FMA + Sigmoid
__global__ void fma_sigmoid_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    const float* __restrict__ c,
    float* __restrict__ output,
    int total_elements
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < total_elements) {
        // Fused multiply-add with Sigmoid activation
        float fma_result = a[idx] * b[idx] + c[idx];
        // Fast sigmoid approximation: 1 / (1 + exp(-x))
        output[idx] = 1.0f / (1.0f + expf(-fma_result));
    }
}

torch::Tensor fma_activation_cuda(
    torch::Tensor a,
    torch::Tensor b,
    torch::Tensor c,
    std::string activation = "relu"
) {
    TORCH_CHECK(a.is_cuda() && a.is_contiguous());
    TORCH_CHECK(b.is_cuda() && b.is_contiguous());
    TORCH_CHECK(c.is_cuda() && c.is_contiguous());
    TORCH_CHECK(a.sizes() == b.sizes() && a.sizes() == c.sizes());
    TORCH_CHECK(a.dim() == 2);

    int N = a.size(0);
    int M = a.size(1);
    int total_elements = N * M;
    
    auto output = torch::empty_like(a);

    // Use 1D grid for element-wise operations
    int block_size = 256;
    int grid_size = (total_elements + block_size - 1) / block_size;

    if (activation == "relu") {
        fma_relu_kernel<<<grid_size, block_size>>>(
            a.data_ptr<float>(),
            b.data_ptr<float>(),
            c.data_ptr<float>(),
            output.data_ptr<float>(),
            total_elements
        );
    } else if (activation == "sigmoid") {
        fma_sigmoid_kernel<<<grid_size, block_size>>>(
            a.data_ptr<float>(),
            b.data_ptr<float>(),
            c.data_ptr<float>(),
            output.data_ptr<float>(),
            total_elements
        );
    } else {
        TORCH_CHECK(false, "Unsupported activation. Use 'relu' or 'sigmoid'");
    }
    
    return output;
}
"""

fma_activation_cpp_source = """
torch::Tensor fma_activation_cuda(
    torch::Tensor a,
    torch::Tensor b,
    torch::Tensor c,
    std::string activation = "relu"
);
"""

# Compile module
fma_activation_module = load_inline(
    name="fma_activation",
    cpp_sources=fma_activation_cpp_source,
    cuda_sources=fma_activation_source,
    functions=["fma_activation_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.fma_activation = fma_activation_module

    def forward(self, a, b, c, activation='relu'):
        if not a.is_contiguous(): a = a.contiguous()
        if not b.is_contiguous(): b = b.contiguous()
        if not c.is_contiguous(): c = b.contiguous()
        return self.fma_activation.fma_activation_cuda(a, b, c, activation)
