import torch
import torch.nn as nn

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

    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 convolution kernel (square, e.g., 3 for a 3x3 kernel).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. 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, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, 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_in, width_in). 

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


# Test code
batch_size = 16
in_channels = 32
out_channels = 64
kernel_size = 3
height_in = 64
width_in = 128
stride = 5
padding = 1
dilation = 2

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

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

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_transpose_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int64_t stride, int64_t padding, int64_t dilation);
"""

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

__global__ void fill_bias_kernel(float* __restrict__ out, const float* __restrict__ bias,
                                 long long total, int Cout, int HWout) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long step = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += step) {
        int oc = (idx / HWout) % Cout;
        out[idx] = bias[oc];
    }
}

// Specialized for the benchmark shape: Cin=32, Cout=64, K=3, stride=5, pad=1, dilation=2.
// One block computes all 64*3*3 non-zero values generated by one input pixel.
__global__ void deconv2d_s5p1d2_k3_c32_c64_kernel(const float* __restrict__ x,
                                                   const float* __restrict__ w,
                                                   float* __restrict__ out,
                                                   int N, int H, int W,
                                                   int Hout, int Wout,
                                                   bool has_bias) {
    int sp = blockIdx.x;
    int iw = sp % W;
    int t = sp / W;
    int ih = t % H;
    int n  = t / H;

    __shared__ float sx[32];
    if (threadIdx.x < 32) {
        sx[threadIdx.x] = x[((n * 32 + threadIdx.x) * H + ih) * W + iw];
    }
    __syncthreads();

    // idx = oc*9 + kh*3 + kw, matching contiguous weight layout for each input channel.
    for (int idx = threadIdx.x; idx < 576; idx += blockDim.x) {
        int k = idx - (idx / 9) * 9;
        int oc = idx / 9;
        int kh = k / 3;
        int kw = k - kh * 3;
        int oh = ih * 5 - 1 + kh * 2;
        int ow = iw * 5 - 1 + kw * 2;
        if ((unsigned)oh < (unsigned)Hout && (unsigned)ow < (unsigned)Wout) {
            float sum = 0.0f;
            #pragma unroll
            for (int c = 0; c < 32; ++c) {
                sum = fmaf(sx[c], w[c * 576 + idx], sum);
            }
            long long oidx = (((long long)n * 64 + oc) * Hout + oh) * Wout + ow;
            if (has_bias) out[oidx] += sum;
            else out[oidx] = sum;
        }
    }
}

__global__ void deconv2d_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        float* __restrict__ out,
                                        int N, int Cin, int H, int W,
                                        int Cout, int K, int stride, int pad, int dil,
                                        int Hout, int Wout,
                                        bool has_bias, bool use_atomic) {
    int sp = blockIdx.x;
    int iw = sp % W;
    int tmp = sp / W;
    int ih = tmp % H;
    int n = tmp / H;

    extern __shared__ float sx[];
    for (int c = threadIdx.x; c < Cin; c += blockDim.x) {
        sx[c] = x[((n * Cin + c) * H + ih) * W + iw];
    }
    __syncthreads();

    int KK = K * K;
    int total = Cout * KK;
    for (int idx = threadIdx.x; idx < total; idx += blockDim.x) {
        int kk = idx % KK;
        int oc = idx / KK;
        int kh = kk / K;
        int kw = kk - kh * K;
        int oh = ih * stride - pad + kh * dil;
        int ow = iw * stride - pad + kw * dil;
        if ((unsigned)oh < (unsigned)Hout && (unsigned)ow < (unsigned)Wout) {
            float sum = 0.0f;
            for (int c = 0; c < Cin; ++c) {
                sum = fmaf(sx[c], w[((c * Cout + oc) * K + kh) * K + kw], sum);
            }
            long long oidx = (((long long)n * Cout + oc) * Hout + oh) * Wout + ow;
            if (use_atomic) {
                atomicAdd(out + oidx, sum);
            } else {
                if (has_bias) out[oidx] += sum;
                else out[oidx] = sum;
            }
        }
    }
}

torch::Tensor conv_transpose_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                     int64_t stride64, int64_t padding64, int64_t dilation64) {
    const int N = (int)x.size(0);
    const int Cin = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int Cout = (int)weight.size(1);
    const int K = (int)weight.size(2);
    const int stride = (int)stride64;
    const int pad = (int)padding64;
    const int dil = (int)dilation64;
    const int Hout = (H - 1) * stride - 2 * pad + dil * (K - 1) + 1;
    const int Wout = (W - 1) * stride - 2 * pad + dil * (K - 1) + 1;

    auto out = torch::empty({N, Cout, Hout, Wout}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const bool has_bias = (bias.numel() != 0);
    long long total_out = (long long)N * Cout * Hout * Wout;

    if (has_bias) {
        int threads = 256;
        int blocks = (int)min(65535LL, (total_out + threads - 1) / threads);
        fill_bias_kernel<<<blocks, threads, 0, stream>>>(out.data_ptr<float>(), bias.data_ptr<float>(), total_out, Cout, Hout * Wout);
    } else {
        cudaMemsetAsync(out.data_ptr<float>(), 0, total_out * sizeof(float), stream);
    }

    int blocks_sp = N * H * W;
    int threads = 256;
    if (Cin == 32 && Cout == 64 && K == 3 && stride == 5 && pad == 1 && dil == 2) {
        deconv2d_s5p1d2_k3_c32_c64_kernel<<<blocks_sp, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), out.data_ptr<float>(), N, H, W, Hout, Wout, has_bias);
    } else {
        // Spatial overlap occurs when adjacent input pixels can write the same output pixel.
        bool use_atomic = (stride <= dil * (K - 1));
        deconv2d_generic_kernel<<<blocks_sp, threads, (size_t)Cin * sizeof(float), stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), out.data_ptr<float>(),
            N, Cin, H, W, Cout, K, stride, pad, dil, Hout, Wout, has_bias, use_atomic);
    }
    return out;
}
"""

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

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