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
from torch.utils.cpp_extension import load_inline

_cpp_source = r'''
#include <torch/extension.h>
#include <vector>

void fused_poly_kernel_launcher(torch::Tensor y, torch::Tensor bias);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor bias,
                           int64_t stride, int64_t padding, int64_t output_padding) {
    auto y = torch::conv_transpose3d(x, weight, conv_bias,
                                     {stride, stride, stride},
                                     {padding, padding, padding},
                                     {output_padding, output_padding, output_padding},
                                     1,
                                     {1, 1, 1});
    fused_poly_kernel_launcher(y, bias);
    return y;
}
'''

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

__global__ void fused_poly_kernel(float* __restrict__ y, const float* __restrict__ bias,
                                  int C, long long spatial, int tiles) {
    int tile = blockIdx.x;
    int nc = blockIdx.y;
    int c = nc % C;
    long long start = (spatial * (long long)tile) / tiles;
    long long end = (spatial * (long long)(tile + 1)) / tiles;
    long long base = (long long)nc * spatial;
    float b = bias[c] + 1.0f;
    for (long long s = start + threadIdx.x; s < end; s += blockDim.x) {
        float z = y[base + s];
        y[base + s] = (2.0f * z + b) * z;
    }
}

void fused_poly_kernel_launcher(torch::Tensor y, torch::Tensor bias) {
    const int N = (int)y.size(0);
    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 tiles = 8;
    if (spatial < 32768) tiles = 1;
    dim3 blocks(tiles, N * C);
    fused_poly_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), bias.data_ptr<float>(), C, spatial, tiles);
}
'''

_ext = load_inline(
    name='fused_convtranspose3d_poly_ext_v3',
    cpp_sources=[_cpp_source],
    cuda_sources=[_cuda_source],
    functions=['fused_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    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))
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding

    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias,
                                  int(self.stride), int(self.padding), int(self.output_padding))