import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, adds a value, takes the minimum, applies GELU, and multiplies by a value.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, add_value, multiply_value):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)
        self.add_value = add_value
        self.multiply_value = multiply_value

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x + self.add_value
        x = torch.min(x, torch.tensor(0.0, device=x.device))
        x = torch.nn.functional.gelu(x)
        x = x * self.multiply_value
        return x

batch_size = 128
in_channels = 64
out_channels = 128
height, width = 64, 64
kernel_size = 4
stride = 2
add_value = 0.5
multiply_value = 2.0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, add_value, multiply_value]

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

cpp_src = r"""
#include <torch/extension.h>
void activation_inplace(torch::Tensor y, double add_value, double multiply_value);
"""

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

__device__ __forceinline__ float erfc_approx_pos(float z) {
    float t = 1.0f / (1.0f + 0.3275911f * z);
    float p = (((((1.061405429f * t - 1.453152027f) * t) + 1.421413741f) * t - 0.284496736f) * t + 0.254829592f) * t;
    return p * __expf(-z * z);
}

__global__ void act_kernel(float* __restrict__ y, long long n, float addv, float mulv) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float inv_sqrt2 = 0.70710678118654752440f;
    for (long long i = idx; i < n; i += stride) {
        float v = y[i] + addv;
        if (v >= 0.0f) {
            y[i] = 0.0f;
        } else {
            y[i] = (0.5f * v * erfc_approx_pos(-v * inv_sqrt2)) * mulv;
        }
    }
}

void activation_inplace(torch::Tensor y, double add_value, double multiply_value) {
    long long n = y.numel();
    if (n == 0) return;
    int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    act_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(y.data_ptr<float>(), n, (float)add_value, (float)multiply_value);
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, add_value, multiply_value):
        super().__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride)
        self.conv_transpose.to(memory_format=torch.channels_last)
        self.add_value = add_value
        self.multiply_value = multiply_value
    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last)
        y = self.conv_transpose(x)
        _ext.activation_inplace(y, float(self.add_value), float(self.multiply_value))
        return y