import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, multiplies by a learnable scalar, applies LeakyReLU, and then GELU.
    """
    def __init__(self, in_channels, out_channels, kernel_size, multiplier_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape)) 
        self.leaky_relu = nn.LeakyReLU()

    def forward(self, x):
        x = self.conv(x)
        x = x * self.multiplier
        x = self.leaky_relu(x)
        x = torch.nn.functional.gelu(x)
        return x

batch_size = 64
in_channels = 64
out_channels = 64
height, width = 256, 256
kernel_size = 3
multiplier_shape = (out_channels, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, multiplier_shape]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r"""
#include <torch/extension.h>

torch::Tensor fused_postop(torch::Tensor x, torch::Tensor multiplier);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ __forceinline__ float erf_approx(float x) {
    float s = x < 0.0f ? -1.0f : 1.0f;
    float ax = fabsf(x);
    float t = 1.0f / (1.0f + 0.3275911f * ax);
    float p = (((((1.061405429f * t - 1.453152027f) * t) + 1.421413741f) * t - 0.284496736f) * t + 0.254829592f) * t;
    return s * (1.0f - p * expf(-ax * ax));
}

__device__ __forceinline__ float gelu_after_leaky(float v) {
    v = (v >= 0.0f) ? v : (0.01f * v);
    return 0.5f * v * (1.0f + erf_approx(v * 0.70710678118654752440f));
}

__global__ void fused_postop_vec4x4_kernel(float4* __restrict__ x4, const float* __restrict__ m,
                                           int HW4, int C) {
    int nc = blockIdx.y;
    int c = nc % C;
    float mult = m[c];
    int base_hw4 = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    long plane = (long)nc * HW4;
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        int hw4 = base_hw4 + i;
        if (hw4 < HW4) {
            float4 r = x4[plane + hw4];
            r.x = gelu_after_leaky(r.x * mult);
            r.y = gelu_after_leaky(r.y * mult);
            r.z = gelu_after_leaky(r.z * mult);
            r.w = gelu_after_leaky(r.w * mult);
            x4[plane + hw4] = r;
        }
    }
}

__global__ void fused_postop_scalar_kernel(float* __restrict__ x, const float* __restrict__ m,
                                           long total, int C, int HW) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (idx / HW) % C;
    x[idx] = gelu_after_leaky(x[idx] * m[c]);
}

torch::Tensor fused_postop(torch::Tensor x, torch::Tensor multiplier) {
    const long total = x.numel();
    const int C = x.size(1);
    const int HW = x.size(2) * x.size(3);
    const int threads = 256;
    if ((HW & 3) == 0) {
        const int HW4 = HW >> 2;
        dim3 grid((HW4 + threads * 4 - 1) / (threads * 4), x.size(0) * C);
        fused_postop_vec4x4_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            reinterpret_cast<float4*>(x.data_ptr<float>()), multiplier.data_ptr<float>(), HW4, C);
    } else {
        const int blocks = (total + threads - 1) / threads;
        fused_postop_scalar_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), multiplier.data_ptr<float>(), total, C, HW);
    }
    return x;
}
"""

_ext = load_inline(name="fused_conv_postop_ext_final", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["fused_postop"], verbose=False, extra_cuda_cflags=["-O3", "--use_fast_math"])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, multiplier_shape):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape))
    def forward(self, x):
        y = self.conv(x)
        return _ext.fused_postop(y, self.multiplier)