import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, adds an input tensor, and applies HardSwish activation.
    """
    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, add_input):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).
            add_input (torch.Tensor): Input tensor to be added after transposed convolution, of shape (batch_size, out_channels, D, H, W).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, D, H, W) after HardSwish activation.
        """
        x = self.conv_transpose(x)
        x = x + add_input
        x = x * torch.nn.functional.hardswish(x)
        return x


batch_size = 128
in_channels = 32
out_channels = 64
D, H, W = 16, 16, 16
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias_shape = (out_channels, 1, 1, 1, 1)

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W), torch.rand(batch_size, out_channels, D*stride, H*stride, W*stride)]

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_src = r"""
#include <torch/extension.h>

torch::Tensor fused_forward(torch::Tensor x,
                            torch::Tensor add,
                            torch::Tensor weight,
                            torch::Tensor conv_bias);
"""

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

// Specialized for the problem sizes:
// input  (N,32,16,16,16), weight (32,64,3,3,3), output/add (N,64,32,32,32)
__global__ void ct3d_fused_kernel(const float* __restrict__ x,
                                  const float* __restrict__ add,
                                  const float* __restrict__ w,
                                  const float* __restrict__ b,
                                  float* __restrict__ out,
                                  unsigned long long total) {
    unsigned long long idx = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ow = (int)(idx & 31ULL);
    const int oh = (int)((idx >> 5) & 31ULL);
    const int od = (int)((idx >> 10) & 31ULL);
    const int oc = (int)((idx >> 15) & 63ULL);
    const int n  = (int)(idx >> 21);

    int did[2], dkd[2], hid[2], hkd[2], wid[2], wkd[2];
    int nd = 0, nh = 0, nw = 0;

    if ((od & 1) == 0) {
        did[0] = od >> 1; dkd[0] = 1; nd = 1;
    } else {
        int t = (od + 1) >> 1;       // kernel index 0
        if (t < 16) { did[nd] = t; dkd[nd] = 0; ++nd; }
        did[nd] = (od - 1) >> 1; dkd[nd] = 2; ++nd;
    }

    if ((oh & 1) == 0) {
        hid[0] = oh >> 1; hkd[0] = 1; nh = 1;
    } else {
        int t = (oh + 1) >> 1;
        if (t < 16) { hid[nh] = t; hkd[nh] = 0; ++nh; }
        hid[nh] = (oh - 1) >> 1; hkd[nh] = 2; ++nh;
    }

    if ((ow & 1) == 0) {
        wid[0] = ow >> 1; wkd[0] = 1; nw = 1;
    } else {
        int t = (ow + 1) >> 1;
        if (t < 16) { wid[nw] = t; wkd[nw] = 0; ++nw; }
        wid[nw] = (ow - 1) >> 1; wkd[nw] = 2; ++nw;
    }

    float acc = __ldg(b + oc);
    const int n_base = n << 17;  // n * 32 * 16 * 16 * 16

    #pragma unroll
    for (int ic = 0; ic < 32; ++ic) {
        const int x_c_base = n_base + (ic << 12);
        const int w_c_base = ((ic << 6) + oc) * 27;

        #pragma unroll
        for (int a = 0; a < 2; ++a) {
            if (a >= nd) continue;
            const int x_d_base = x_c_base + (did[a] << 8);
            const int w_d_base = w_c_base + dkd[a] * 9;
            #pragma unroll
            for (int bb = 0; bb < 2; ++bb) {
                if (bb >= nh) continue;
                const int x_h_base = x_d_base + (hid[bb] << 4);
                const int w_h_base = w_d_base + hkd[bb] * 3;
                #pragma unroll
                for (int c = 0; c < 2; ++c) {
                    if (c >= nw) continue;
                    acc = fmaf(__ldg(x + x_h_base + wid[c]),
                               __ldg(w + w_h_base + wkd[c]), acc);
                }
            }
        }
    }

    float v = acc + __ldg(add + idx);
    float r;
    if (v <= -3.0f) {
        r = 0.0f;
    } else if (v >= 3.0f) {
        r = v * v;
    } else {
        r = (v * v) * (v + 3.0f) * 0.166666666666666657f;
    }
    out[idx] = r;
}

torch::Tensor fused_forward(torch::Tensor x,
                            torch::Tensor add,
                            torch::Tensor weight,
                            torch::Tensor conv_bias) {
    auto out = torch::empty_like(add);
    const unsigned long long total = (unsigned long long)add.numel();
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    ct3d_fused_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), add.data_ptr<float>(), weight.data_ptr<float>(),
        conv_bias.data_ptr<float>(), out.data_ptr<float>(), total);
    return out;
}
"""

_ext = load_inline(
    name="ct3d_fused_hardswish_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Fused replacement for ConvTranspose3d + add_input + x * HardSwish(x).
    """
    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
        )
        # Present to match the reference module state_dict; intentionally unused, as in the reference forward.
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x, add_input):
        return _ext.fused_forward(x, add_input, self.conv_transpose.weight, self.conv_transpose.bias)