import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, followed by a sum, 
    a residual add, a multiplication, and another residual add.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        x = self.conv_transpose(x)
        original_x = x.clone().detach()
        x = x + self.bias
        x = x + original_x
        x = x * original_x
        x = x + original_x
        return x

batch_size = 16
in_channels = 32
out_channels = 64
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias_shape = (out_channels, 1, 1, 1)

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape]

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>

torch::Tensor fused_residual_bias_forward(torch::Tensor y, torch::Tensor bias);
"""

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

__global__ void fused_residual_bias_kernel(float* __restrict__ y,
                                           const float* __restrict__ bias,
                                           long long total,
                                           int C,
                                           long long spatial) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        long long c = (idx / spatial) % C;
        float v = y[idx];
        // Reference after transposed convolution y:
        // ((y + custom_bias) + y) * y + y = 2*y*y + (custom_bias + 1)*y
        y[idx] = fmaf(2.0f * v, v, (bias[c] + 1.0f) * v);
    }
}

torch::Tensor fused_residual_bias_forward(torch::Tensor y, torch::Tensor bias) {
    const c10::cuda::CUDAGuard device_guard(y.device());
    y = y.contiguous();
    bias = bias.contiguous();

    const long long total = y.numel();
    if (total == 0) return y;

    const int C = (int)y.size(1);
    const long long spatial = (long long)y.size(2) * y.size(3) * y.size(4);

    const int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;

    fused_residual_bias_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), total, C, spatial);
    return y;
}
"""

_fused_ext = load_inline(
    name="fused_convtranspose3d_epilogue_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["fused_residual_bias_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Same module interface as the reference.  The expensive ConvTranspose3d is
    executed by PyTorch/cuDNN, while the clone/add/add/mul/add tail is replaced
    by one custom CUDA epilogue kernel executed in-place on the convolution
    output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
        )
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        y = F.conv_transpose3d(
            x,
            self.conv_transpose.weight,
            self.conv_transpose.bias,
            stride=self.conv_transpose.stride,
            padding=self.conv_transpose.padding,
            output_padding=self.conv_transpose.output_padding,
            groups=self.conv_transpose.groups,
            dilation=self.conv_transpose.dilation,
        )
        return _fused_ext.fused_residual_bias_forward(y, self.bias)