import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 2D transposed convolution operation with asymmetric input and kernel, with optional padding.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Size of the convolution kernel (height, width).
        stride (tuple, optional): Stride of the convolution (height, width). Defaults to (1, 1).
        padding (tuple, optional): Padding applied to the input (height, width). Defaults to (0, 0).
        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: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 2D transposed 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 = 32
out_channels = 32
kernel_size = (3, 7)
height = 512
width = 1024
stride = (1, 1)
padding = (1, 3)

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, stride, padding]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

#ifndef __CUDA_ARCH__
#define __CUDA_ARCH__ 0
#endif

static __global__ __launch_bounds__(128, 2)
void deconv32_3x7_s1_kernel(const float* __restrict__ x,
                            const float* __restrict__ w,
                            const float* __restrict__ b,
                            float* __restrict__ yout,
                            int N, int H, int W, int has_bias) {
    // Specialized for IC=OC=32, kernel=(3,7), stride=(1,1), padding=(1,3).
    // Output spatial size equals input spatial size for the problem dimensions.
    constexpr int IC = 32;
    constexpr int OC = 32;
    constexpr int KH = 3;
    constexpr int KW = 7;
    constexpr int OCP = 8;

    const long long HW = (long long)H * W;
    const long long total_sp = (long long)N * HW;
    long long p = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= total_sp) return;

    const int og = blockIdx.y;              // 0..3, group of 8 output channels
    const int oc0 = og * OCP;
    const int xx = (int)(p % W);
    long long t = p / W;
    const int yy = (int)(t % H);
    const int n  = (int)(t / H);

    float a0 = has_bias ? __ldg(b + oc0 + 0) : 0.0f;
    float a1 = has_bias ? __ldg(b + oc0 + 1) : 0.0f;
    float a2 = has_bias ? __ldg(b + oc0 + 2) : 0.0f;
    float a3 = has_bias ? __ldg(b + oc0 + 3) : 0.0f;
    float a4 = has_bias ? __ldg(b + oc0 + 4) : 0.0f;
    float a5 = has_bias ? __ldg(b + oc0 + 5) : 0.0f;
    float a6 = has_bias ? __ldg(b + oc0 + 6) : 0.0f;
    float a7 = has_bias ? __ldg(b + oc0 + 7) : 0.0f;

    const long long in_n_base = (long long)n * IC * HW;

    if (yy >= 1 && yy < H - 1 && xx >= 3 && xx < W - 3) {
        // Interior: all 21 taps are valid, no bounds checks in the hot path.
        #pragma unroll
        for (int ic = 0; ic < IC; ++ic) {
            const float* __restrict__ xp = x + in_n_base + (long long)ic * HW;
            const int w_ic_base = ic * OC * KH * KW + oc0 * KH * KW;
            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                const int iy = yy + 1 - kh;
                const long long row = (long long)iy * W;
                #pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const float v = __ldg(xp + row + (xx + 3 - kw));
                    const int wb = w_ic_base + kh * KW + kw;
                    a0 = fmaf(v, __ldg(w + wb + 0 * KH * KW), a0);
                    a1 = fmaf(v, __ldg(w + wb + 1 * KH * KW), a1);
                    a2 = fmaf(v, __ldg(w + wb + 2 * KH * KW), a2);
                    a3 = fmaf(v, __ldg(w + wb + 3 * KH * KW), a3);
                    a4 = fmaf(v, __ldg(w + wb + 4 * KH * KW), a4);
                    a5 = fmaf(v, __ldg(w + wb + 5 * KH * KW), a5);
                    a6 = fmaf(v, __ldg(w + wb + 6 * KH * KW), a6);
                    a7 = fmaf(v, __ldg(w + wb + 7 * KH * KW), a7);
                }
            }
        }
    } else {
        // Borders: apply transposed-convolution bounds induced by padding.
        #pragma unroll
        for (int ic = 0; ic < IC; ++ic) {
            const float* __restrict__ xp = x + in_n_base + (long long)ic * HW;
            const int w_ic_base = ic * OC * KH * KW + oc0 * KH * KW;
            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                const int iy = yy + 1 - kh;
                if ((unsigned)iy >= (unsigned)H) continue;
                const long long row = (long long)iy * W;
                #pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const int ix = xx + 3 - kw;
                    if ((unsigned)ix >= (unsigned)W) continue;
                    const float v = __ldg(xp + row + ix);
                    const int wb = w_ic_base + kh * KW + kw;
                    a0 = fmaf(v, __ldg(w + wb + 0 * KH * KW), a0);
                    a1 = fmaf(v, __ldg(w + wb + 1 * KH * KW), a1);
                    a2 = fmaf(v, __ldg(w + wb + 2 * KH * KW), a2);
                    a3 = fmaf(v, __ldg(w + wb + 3 * KH * KW), a3);
                    a4 = fmaf(v, __ldg(w + wb + 4 * KH * KW), a4);
                    a5 = fmaf(v, __ldg(w + wb + 5 * KH * KW), a5);
                    a6 = fmaf(v, __ldg(w + wb + 6 * KH * KW), a6);
                    a7 = fmaf(v, __ldg(w + wb + 7 * KH * KW), a7);
                }
            }
        }
    }

    const long long out_n_base = (long long)n * OC * HW;
    const long long sp = (long long)yy * W + xx;
    yout[out_n_base + (long long)(oc0 + 0) * HW + sp] = a0;
    yout[out_n_base + (long long)(oc0 + 1) * HW + sp] = a1;
    yout[out_n_base + (long long)(oc0 + 2) * HW + sp] = a2;
    yout[out_n_base + (long long)(oc0 + 3) * HW + sp] = a3;
    yout[out_n_base + (long long)(oc0 + 4) * HW + sp] = a4;
    yout[out_n_base + (long long)(oc0 + 5) * HW + sp] = a5;
    yout[out_n_base + (long long)(oc0 + 6) * HW + sp] = a6;
    yout[out_n_base + (long long)(oc0 + 7) * HW + sp] = a7;
}

torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "only float32 supported");

    x = x.contiguous();
    weight = weight.contiguous();
    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);
    TORCH_CHECK(IC == 32 && weight.size(0) == 32 && weight.size(1) == 32 &&
                weight.size(2) == 3 && weight.size(3) == 7,
                "specialized kernel expects IC=OC=32, kernel=(3,7)");

    auto out = torch::empty({N, 32, H, W}, x.options());
    const int threads = 128;
    const long long total_sp = (long long)N * H * W;
    dim3 grid((unsigned int)((total_sp + threads - 1) / threads), 4, 1);
    int has_bias = (bias.defined() && bias.numel() >= 32) ? 1 : 0;
    const float* bptr = has_bias ? bias.contiguous().data_ptr<float>() : nullptr;
    deconv32_3x7_s1_kernel<<<grid, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bptr, out.data_ptr<float>(), N, H, W, has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        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)