import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 2D convolution with square input and square kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        output_padding (int, optional): Additional size added to one side of the output shape. Defaults to 0.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)

# Test code
batch_size = 8
in_channels = 64  # double channels for heavier compute
out_channels = 64
kernel_size = 3
# larger square input
height = 1024
width = 1024

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

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_transpose2d_forward(torch::Tensor x,
                                      torch::Tensor weight,
                                      torch::Tensor bias,
                                      int64_t stride,
                                      int64_t padding,
                                      int64_t output_padding,
                                      int64_t groups);
'''

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

#define ACCUM8(v, wp) \
    acc0 = fmaf((v), __ldg((wp) + 0*9), acc0); \
    acc1 = fmaf((v), __ldg((wp) + 1*9), acc1); \
    acc2 = fmaf((v), __ldg((wp) + 2*9), acc2); \
    acc3 = fmaf((v), __ldg((wp) + 3*9), acc3); \
    acc4 = fmaf((v), __ldg((wp) + 4*9), acc4); \
    acc5 = fmaf((v), __ldg((wp) + 5*9), acc5); \
    acc6 = fmaf((v), __ldg((wp) + 6*9), acc6); \
    acc7 = fmaf((v), __ldg((wp) + 7*9), acc7)

__global__ void deconv3x3_s1p0_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ bias,
                                      float* __restrict__ y,
                                      int N, int IC, int OC,
                                      int H, int W, int OH, int OW,
                                      int has_bias) {
    const int tx = threadIdx.x;                 // 0..31: neighbouring output pixels
    const int ty = threadIdx.y;                 // several output-channel groups per block
    const int pix = blockIdx.x * blockDim.x + tx; // flattened N*OH*OW
    const int ocg = (blockIdx.y * blockDim.y + ty) * 8;

    const int total_pix = N * OH * OW;
    if (pix >= total_pix || ocg >= OC) return;

    const int OHW = OH * OW;
    const int HW = H * W;
    const int n = pix / OHW;
    const int sp = pix - n * OHW;
    const int oh = sp / OW;
    const int ow = sp - oh * OW;

    float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f;
    float acc4 = 0.0f, acc5 = 0.0f, acc6 = 0.0f, acc7 = 0.0f;

    // Fast path used by essentially all pixels for the benchmark shape.
    if (ocg + 7 < OC && oh >= 2 && oh < H && ow >= 2 && ow < W) {
        const float* x_n = x + n * IC * HW;
        #pragma unroll 1
        for (int ic = 0; ic < IC; ++ic) {
            const float* x_c = x_n + ic * HW;
            const float* w_c = w + (ic * OC + ocg) * 9;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int ih = oh - kh;
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    const int iw = ow - kw;
                    const float v = __ldg(x_c + ih * W + iw);
                    const float* wp = w_c + kh * 3 + kw;
                    ACCUM8(v, wp);
                }
            }
        }
        if (has_bias) {
            acc0 += __ldg(bias + ocg + 0); acc1 += __ldg(bias + ocg + 1);
            acc2 += __ldg(bias + ocg + 2); acc3 += __ldg(bias + ocg + 3);
            acc4 += __ldg(bias + ocg + 4); acc5 += __ldg(bias + ocg + 5);
            acc6 += __ldg(bias + ocg + 6); acc7 += __ldg(bias + ocg + 7);
        }
        float* y_n = y + n * OC * OHW + sp;
        y_n[(ocg + 0) * OHW] = acc0;
        y_n[(ocg + 1) * OHW] = acc1;
        y_n[(ocg + 2) * OHW] = acc2;
        y_n[(ocg + 3) * OHW] = acc3;
        y_n[(ocg + 4) * OHW] = acc4;
        y_n[(ocg + 5) * OHW] = acc5;
        y_n[(ocg + 6) * OHW] = acc6;
        y_n[(ocg + 7) * OHW] = acc7;
        return;
    }

    // Generic edge/tail path for the one-pixel-wide transposed-convolution border.
    const float* x_n = x + n * IC * HW;
    #pragma unroll 1
    for (int ic = 0; ic < IC; ++ic) {
        const float* x_c = x_n + ic * HW;
        const float* w_c = w + (ic * OC + ocg) * 9;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            const int ih = oh - kh;
            if ((unsigned)ih >= (unsigned)H) continue;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                const int iw = ow - kw;
                if ((unsigned)iw >= (unsigned)W) continue;
                const float v = __ldg(x_c + ih * W + iw);
                const float* wp = w_c + kh * 3 + kw;
                if (ocg + 7 < OC) {
                    ACCUM8(v, wp);
                } else {
                    if (ocg + 0 < OC) acc0 = fmaf(v, __ldg(wp + 0*9), acc0);
                    if (ocg + 1 < OC) acc1 = fmaf(v, __ldg(wp + 1*9), acc1);
                    if (ocg + 2 < OC) acc2 = fmaf(v, __ldg(wp + 2*9), acc2);
                    if (ocg + 3 < OC) acc3 = fmaf(v, __ldg(wp + 3*9), acc3);
                    if (ocg + 4 < OC) acc4 = fmaf(v, __ldg(wp + 4*9), acc4);
                    if (ocg + 5 < OC) acc5 = fmaf(v, __ldg(wp + 5*9), acc5);
                    if (ocg + 6 < OC) acc6 = fmaf(v, __ldg(wp + 6*9), acc6);
                    if (ocg + 7 < OC) acc7 = fmaf(v, __ldg(wp + 7*9), acc7);
                }
            }
        }
    }
    float* y_n = y + n * OC * OHW + sp;
    if (has_bias) {
        if (ocg + 0 < OC) acc0 += __ldg(bias + ocg + 0);
        if (ocg + 1 < OC) acc1 += __ldg(bias + ocg + 1);
        if (ocg + 2 < OC) acc2 += __ldg(bias + ocg + 2);
        if (ocg + 3 < OC) acc3 += __ldg(bias + ocg + 3);
        if (ocg + 4 < OC) acc4 += __ldg(bias + ocg + 4);
        if (ocg + 5 < OC) acc5 += __ldg(bias + ocg + 5);
        if (ocg + 6 < OC) acc6 += __ldg(bias + ocg + 6);
        if (ocg + 7 < OC) acc7 += __ldg(bias + ocg + 7);
    }
    if (ocg + 0 < OC) y_n[(ocg + 0) * OHW] = acc0;
    if (ocg + 1 < OC) y_n[(ocg + 1) * OHW] = acc1;
    if (ocg + 2 < OC) y_n[(ocg + 2) * OHW] = acc2;
    if (ocg + 3 < OC) y_n[(ocg + 3) * OHW] = acc3;
    if (ocg + 4 < OC) y_n[(ocg + 4) * OHW] = acc4;
    if (ocg + 5 < OC) y_n[(ocg + 5) * OHW] = acc5;
    if (ocg + 6 < OC) y_n[(ocg + 6) * OHW] = acc6;
    if (ocg + 7 < OC) y_n[(ocg + 7) * OHW] = acc7;
}

// Simple fully-generic fp32 fallback for uncommon stride/padding/groups settings.
__global__ void deconv_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ bias,
                                      float* __restrict__ y,
                                      int N, int IC, int OC, int H, int W, int K,
                                      int stride, int padding, int groups,
                                      int OH, int OW, int has_bias) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * OC * OH * OW;
    if (idx >= total) return;
    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int oc = t % OC;
    int n = t / OC;
    int ocpg = OC / groups;
    int icpg = IC / groups;
    int g = oc / ocpg;
    float acc = 0.0f;
    for (int ic0 = 0; ic0 < icpg; ++ic0) {
        int ic = g * icpg + ic0;
        for (int kh = 0; kh < K; ++kh) {
            int num_h = oh + padding - kh;
            if (num_h % stride) continue;
            int ih = num_h / stride;
            if ((unsigned)ih >= (unsigned)H) continue;
            for (int kw = 0; kw < K; ++kw) {
                int num_w = ow + padding - kw;
                if (num_w % stride) continue;
                int iw = num_w / stride;
                if ((unsigned)iw >= (unsigned)W) continue;
                float xv = __ldg(x + ((n * IC + ic) * H + ih) * W + iw);
                float wv = __ldg(w + (((ic * ocpg + (oc - g * ocpg)) * K + kh) * K + kw));
                acc = fmaf(xv, wv, acc);
            }
        }
    }
    if (has_bias) acc += __ldg(bias + oc);
    y[idx] = acc;
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x,
                                      torch::Tensor weight,
                                      torch::Tensor bias,
                                      int64_t stride,
                                      int64_t padding,
                                      int64_t output_padding,
                                      int64_t groups) {
    const at::cuda::CUDAGuard device_guard(x.device());
    const int N = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int K = (int)weight.size(2);
    const int OC = (int)weight.size(1) * (int)groups;
    const int OH = (H - 1) * (int)stride - 2 * (int)padding + K + (int)output_padding;
    const int OW = (W - 1) * (int)stride - 2 * (int)padding + K + (int)output_padding;
    auto y = torch::empty({N, OC, OH, OW}, x.options());
    const int has_bias = (bias.numel() > 0);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (K == 3 && stride == 1 && padding == 0 && output_padding == 0 && groups == 1) {
        dim3 block(32, 4, 1);
        dim3 grid((N * OH * OW + block.x - 1) / block.x, (OC + 8 * block.y - 1) / (8 * block.y), 1);
        deconv3x3_s1p0_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), N, IC, OC, H, W, OH, OW, has_bias);
    } else {
        int total = N * OC * OH * OW;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        deconv_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), N, IC, OC, H, W, K, (int)stride, (int)padding, (int)groups, OH, OW, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='deconv2d_3x3_s1_ext',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['conv_transpose2d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, bias=bias)
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding
        self.groups = groups
        self.register_buffer('_empty_bias', torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv_transpose2d.bias
        if b is None:
            b = self._empty_bias
        return _ext.conv_transpose2d_forward(x, self.conv_transpose2d.weight, b, self.stride, self.padding, self.output_padding, self.groups)