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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>
std::vector<torch::Tensor> pack_wb(torch::Tensor w, torch::Tensor b);
torch::Tensor post_interleave(torch::Tensor y, double addv, double mulv);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <vector>

__global__ void pack_kernel(const float* __restrict__ w, const float* __restrict__ b,
                            float* __restrict__ p, float* __restrict__ pb,
                            int I, int O) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 4 * O * I * 4;
    if (idx < total) {
        int t = idx;
        int kk = t & 3; t >>= 2;
        int ci = t % I; t /= I;
        int oc = t;
        int co = oc % O;
        int par = oc / O;
        int ph = par >> 1;
        int pw = par & 1;
        int kh, kw;
        if (kk == 0) { kh = ph + 2; kw = pw + 2; }
        else if (kk == 1) { kh = ph + 2; kw = pw; }
        else if (kk == 2) { kh = ph; kw = pw + 2; }
        else { kh = ph; kw = pw; }
        p[idx] = w[(((ci * O + co) * 4 + kh) * 4 + kw)];
    }
    if (idx < 4 * O) pb[idx] = b[idx % O];
}

std::vector<torch::Tensor> pack_wb(torch::Tensor w, torch::Tensor b) {
    int I = (int)w.size(0);
    int O = (int)w.size(1);
    auto p = torch::empty({4 * O, I, 2, 2}, w.options());
    auto pb = torch::empty({4 * O}, b.options());
    int total = 4 * O * I * 4;
    pack_kernel<<<(total + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        w.data_ptr<float>(), b.data_ptr<float>(), p.data_ptr<float>(), pb.data_ptr<float>(), I, O);
    return {p, pb};
}

__device__ __forceinline__ float post_fun(float v, float addv, float mulv) {
    float x = v + addv;
    if (x >= 0.0f) return 0.0f;
    const float inv_sqrt2 = 0.7071067811865475244f;
    return (0.5f * x * (1.0f + erff(x * inv_sqrt2))) * mulv;
}

__global__ void inter_kernel(const float* __restrict__ y, float* __restrict__ out,
                             long long total, int C, int H, int W,
                             float addv, float mulv) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    int Hp = H >> 1;
    int Wp = W >> 1;
    for (long long linear = idx; linear < total; linear += stride) {
        long long t = linear;
        int wo = (int)(t % W); t /= W;
        int ho = (int)(t % H); t /= H;
        int co = (int)(t % C); t /= C;
        int n = (int)t;
        int par = ((ho & 1) << 1) | (wo & 1);
        int r = ho >> 1;
        int c = wo >> 1;
        long long src = (((long long)n * (4 * C) + par * C + co) * Hp + r) * Wp + c;
        out[linear] = post_fun(y[src], addv, mulv);
    }
}

torch::Tensor post_interleave(torch::Tensor y, double addv, double mulv) {
    int N = (int)y.size(0);
    int C4 = (int)y.size(1);
    int Hp = (int)y.size(2);
    int Wp = (int)y.size(3);
    int C = C4 / 4;
    int H = Hp * 2;
    int W = Wp * 2;
    auto out = torch::empty({N, C, H, W}, y.options());
    long long total = out.numel();
    int blocks = (int)((total + 255) / 256);
    if (blocks > 65535) blocks = 65535;
    inter_kernel<<<blocks, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), out.data_ptr<float>(), total, C, H, W, (float)addv, (float)mulv);
    return out;
}
'''

_ext = load_inline(
    name='deconv_stride2_k4_post_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['pack_wb', 'post_interleave'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, add_value, multiply_value):
        super(ModelNew, 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):
        p, pb = _ext.pack_wb(self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous())
        y = F.conv2d(x, p, pb, padding=1)
        return _ext.post_interleave(y, float(self.add_value), float(self.multiply_value))