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_post_inplace_cuda(torch::Tensor x, torch::Tensor multiplier);

torch::Tensor fused_post_inplace(torch::Tensor x, torch::Tensor multiplier) {
    return fused_post_inplace_cuda(x, multiplier);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_post_kernel(float* __restrict__ x,
                                  const float* __restrict__ multiplier,
                                  int C, int HW) {
    const int nc = (int)blockIdx.y;          // combined N,C index
    const int c = nc % C;
    const float m = multiplier[c];
    const long long base = (long long)nc * (long long)HW;

    for (int hw = (int)blockIdx.x * blockDim.x + threadIdx.x;
         hw < HW;
         hw += (int)gridDim.x * blockDim.x) {
        const long long idx = base + (long long)hw;
        float v = x[idx] * m;
        v = (v >= 0.0f) ? v : (v * 0.01f);
        // Exact GELU: 0.5*x*(1+erf(x/sqrt(2))).
        float e = erff(v * 0.70710678118654752440f);
        x[idx] = 0.5f * v * (1.0f + e);
    }
}

torch::Tensor fused_post_inplace_cuda(torch::Tensor x, torch::Tensor multiplier) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int HW = H * W;

    const int threads = 256;
    int bx = (HW + threads - 1) / threads;
    if (bx > 64) bx = 64;   // grid-stride over spatial dimension; fewer CTAs, no per-element channel divides
    if (bx < 1) bx = 1;
    const dim3 blocks((unsigned)bx, (unsigned)(N * C));

    fused_post_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), multiplier.data_ptr<float>(), C, HW);
    return x;
}
"""

_ext = load_inline(
    name="fused_conv_post_gelu_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_post_inplace"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Conv2d followed by per-output-channel multiplier, LeakyReLU, and exact GELU.
    The convolution is kept as cuDNN/PyTorch Conv2d while the three pointwise
    operations are fused into one custom CUDA kernel writing the temporary
    convolution output in-place.
    """
    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_post_inplace(y, self.multiplier)