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>

torch::Tensor conv_transpose_act_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, 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>
#include <math_constants.h>

#define TILE_M 16
#define TILE_N 32
#define TILE_K 32

__device__ __forceinline__ float gelu_exact_f(float x) {
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f));
}

// Specialized for N=128, Cin=64, Cout=128, H=W=64, kernel=4, stride=2, OH=OW=130.
// Computes one output-parity sub-problem per grid.z as a tiled GEMM with fused epilogue.
__global__ void deconv4s2_gelu_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ b,
                                      float* __restrict__ y,
                                      float addv, float mulv) {
    __shared__ float As[TILE_M][TILE_K];
    __shared__ float Bs[TILE_K][TILE_N];

    const int tid = threadIdx.x; // 0..255
    const int lr = tid >> 4;     // 0..15
    const int lc = tid & 15;     // 0..15

    const int parity = blockIdx.z;
    const int ph = parity >> 1;
    const int pw = parity & 1;

    const int M = 128 * 65 * 65;
    const int row = blockIdx.x * TILE_M + lr;
    const int co0 = blockIdx.y * TILE_N + lc;
    const int co1 = co0 + 16;

    float acc0 = 0.0f;
    float acc1 = 0.0f;

    int n = 0, rr = 0, cc = 0;
    if (row < M) {
        int tmp = row;
        cc = tmp % 65;
        tmp /= 65;
        rr = tmp % 65;
        n = tmp / 65;
    }

    // K = Cin * 4 spatial contributors for a fixed parity.
    #pragma unroll
    for (int kt = 0; kt < 256; kt += TILE_K) {
        // Load A: 16 x 32, two elements per thread.
        int aidx = tid;
        #pragma unroll
        for (int rep = 0; rep < 2; ++rep) {
            int idx = aidx + rep * 256;
            int ar = idx >> 5;       // /32
            int ak = idx & 31;
            int grow = blockIdx.x * TILE_M + ar;
            int gk = kt + ak;
            float aval = 0.0f;
            if (grow < M) {
                int tmp = grow;
                int c = tmp % 65;
                tmp /= 65;
                int r = tmp % 65;
                int bn = tmp / 65;
                int ci = gk >> 2;
                int term = gk & 3;
                int dh = term >> 1;
                int dw = term & 1;
                int ih = r - dh;
                int iw = c - dw;
                if ((unsigned)ih < 64u && (unsigned)iw < 64u) {
                    aval = x[((bn * 64 + ci) * 64 + ih) * 64 + iw];
                }
            }
            As[ar][ak] = aval;
        }

        // Load B: 32 x 32, four elements per thread. Weight layout [Cin, Cout, 4, 4].
        #pragma unroll
        for (int rep = 0; rep < 4; ++rep) {
            int idx = tid + rep * 256;
            int bk = idx >> 5;       // /32
            int bc = idx & 31;
            int gk = kt + bk;
            int co = blockIdx.y * TILE_N + bc;
            int ci = gk >> 2;
            int term = gk & 3;
            int dh = term >> 1;
            int dw = term & 1;
            int kh = ph + (dh << 1);
            int kw = pw + (dw << 1);
            Bs[bk][bc] = w[(((ci * 128 + co) * 4 + kh) * 4 + kw)];
        }
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE_K; ++k) {
            float av = As[lr][k];
            acc0 = fmaf(av, Bs[k][lc], acc0);
            acc1 = fmaf(av, Bs[k][lc + 16], acc1);
        }
        __syncthreads();
    }

    if (row < M) {
        const int oh = (rr << 1) + ph;
        const int ow = (cc << 1) + pw;
        const int base = ((n * 128) * 130 + oh) * 130 + ow;
        float z0 = acc0 + b[co0] + addv;
        z0 = fminf(z0, 0.0f);
        y[base + co0 * 130 * 130] = gelu_exact_f(z0) * mulv;

        float z1 = acc1 + b[co1] + addv;
        z1 = fminf(z1, 0.0f);
        y[base + co1 * 130 * 130] = gelu_exact_f(z1) * mulv;
    }
}

torch::Tensor conv_transpose_act_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double add_value, double multiply_value) {
    auto y = torch::empty({x.size(0), weight.size(1), (x.size(2) - 1) * 2 + 4, (x.size(3) - 1) * 2 + 4}, x.options());
    const int M = (int)(x.size(0) * 65 * 65);
    dim3 block(256);
    dim3 grid((M + TILE_M - 1) / TILE_M, 4, 4);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    deconv4s2_gelu_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(),
        (float)add_value, (float)multiply_value);
    return y;
}
'''

_ext = load_inline(
    name='deconv4s2_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv_transpose_act_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

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):
        return _ext.conv_transpose_act_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, float(self.add_value), float(self.multiply_value))