import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 1D convolution operation with square input and asymmetric kernel, optionally with dilation.

    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.
        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.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 1D convolution.

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

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, length_out).
        """
        return self.conv1d_transpose(x)

# Test code
batch_size = 32
in_channels = 32
out_channels = 64
kernel_size = 5
length = 131072
stride = 1
padding = 0
dilation = 3

def get_inputs():
    x = torch.rand(batch_size, in_channels, length)
    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 deconv1d_forward(torch::Tensor x, torch::Tensor w, int64_t stride, int64_t padding, int64_t dilation, bool has_bias, torch::Tensor bias);
'''

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

__global__ void deconv1d_co4_kernel(const float* __restrict__ x, const float* __restrict__ w, float* __restrict__ y, int B, int L, int Lout) {
    long long u = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * 16LL * (long long)Lout;
    if (u >= total) return;
    int t = (int)(u % Lout);
    int cg = (int)((u / Lout) & 15LL);
    int co = cg * 4;
    int b = (int)(u / ((long long)16 * Lout));
    const float* xb = x + ((long long)b * 32LL * L);
    float s0=0.f, s1=0.f, s2=0.f, s3=0.f;
    if (t >= 12 && t < L) {
        #pragma unroll
        for (int ci=0; ci<32; ++ci) {
            const float* xp = xb + (long long)ci * L + t;
            const float* wp = w + (ci * 64 + co) * 5;
            float xv;
            xv=xp[0];   s0=fmaf(xv,wp[0],s0);  s1=fmaf(xv,wp[5],s1);  s2=fmaf(xv,wp[10],s2); s3=fmaf(xv,wp[15],s3);
            xv=xp[-3];  s0=fmaf(xv,wp[1],s0);  s1=fmaf(xv,wp[6],s1);  s2=fmaf(xv,wp[11],s2); s3=fmaf(xv,wp[16],s3);
            xv=xp[-6];  s0=fmaf(xv,wp[2],s0);  s1=fmaf(xv,wp[7],s1);  s2=fmaf(xv,wp[12],s2); s3=fmaf(xv,wp[17],s3);
            xv=xp[-9];  s0=fmaf(xv,wp[3],s0);  s1=fmaf(xv,wp[8],s1);  s2=fmaf(xv,wp[13],s2); s3=fmaf(xv,wp[18],s3);
            xv=xp[-12]; s0=fmaf(xv,wp[4],s0);  s1=fmaf(xv,wp[9],s1);  s2=fmaf(xv,wp[14],s2); s3=fmaf(xv,wp[19],s3);
        }
    } else {
        #pragma unroll
        for (int ci=0; ci<32; ++ci) {
            const float* xpbase = xb + (long long)ci * L;
            const float* wp = w + (ci * 64 + co) * 5;
            int p; float xv;
            p=t;    if ((unsigned)p < (unsigned)L) { xv=xpbase[p]; s0=fmaf(xv,wp[0],s0); s1=fmaf(xv,wp[5],s1); s2=fmaf(xv,wp[10],s2); s3=fmaf(xv,wp[15],s3); }
            p=t-3;  if ((unsigned)p < (unsigned)L) { xv=xpbase[p]; s0=fmaf(xv,wp[1],s0); s1=fmaf(xv,wp[6],s1); s2=fmaf(xv,wp[11],s2); s3=fmaf(xv,wp[16],s3); }
            p=t-6;  if ((unsigned)p < (unsigned)L) { xv=xpbase[p]; s0=fmaf(xv,wp[2],s0); s1=fmaf(xv,wp[7],s1); s2=fmaf(xv,wp[12],s2); s3=fmaf(xv,wp[17],s3); }
            p=t-9;  if ((unsigned)p < (unsigned)L) { xv=xpbase[p]; s0=fmaf(xv,wp[3],s0); s1=fmaf(xv,wp[8],s1); s2=fmaf(xv,wp[13],s2); s3=fmaf(xv,wp[18],s3); }
            p=t-12; if ((unsigned)p < (unsigned)L) { xv=xpbase[p]; s0=fmaf(xv,wp[4],s0); s1=fmaf(xv,wp[9],s1); s2=fmaf(xv,wp[14],s2); s3=fmaf(xv,wp[19],s3); }
        }
    }
    float* yp = y + ((long long)b * 64 + co) * Lout + t;
    yp[0] = s0; yp[Lout] = s1; yp[(long long)2 * Lout] = s2; yp[(long long)3 * Lout] = s3;
}

__global__ void deconv1d_generic_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias,
                                        float* __restrict__ y, int B, int Cin, int L, int Cout, int K,
                                        int stride, int padding, int dilation, int Lout, bool has_bias) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)B * Cout * Lout;
    if (idx >= total) return;
    int t = (int)(idx % Lout);
    int co = (int)((idx / Lout) % Cout);
    int b = (int)(idx / ((long long)Cout * Lout));
    float sum = has_bias ? bias[co] : 0.0f;
    for (int ci=0; ci<Cin; ++ci) {
        const float* xp = x + ((long long)b * Cin + ci) * L;
        const float* wp = w + ((long long)ci * Cout + co) * K;
        for (int k=0; k<K; ++k) {
            int pos = t + padding - k * dilation;
            if (pos % stride == 0) {
                int ip = pos / stride;
                if ((unsigned)ip < (unsigned)L) sum = fmaf(xp[ip], wp[k], sum);
            }
        }
    }
    y[idx] = sum;
}

torch::Tensor deconv1d_forward(torch::Tensor x, torch::Tensor w, int64_t stride, int64_t padding, int64_t dilation, bool has_bias, torch::Tensor bias) {
    int B = x.size(0), Cin = x.size(1), L = x.size(2), Cout = w.size(1), K = w.size(2);
    int Lout = (L - 1) * (int)stride - 2 * (int)padding + (int)dilation * (K - 1) + 1;
    auto y = torch::empty({B, Cout, Lout}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    if (x.scalar_type() == at::kFloat && Cin == 32 && Cout == 64 && K == 5 && stride == 1 && padding == 0 && dilation == 3 && !has_bias) {
        long long total = (long long)B * 16LL * Lout;
        int blocks = (int)((total + threads - 1) / threads);
        deconv1d_co4_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), y.data_ptr<float>(), B, L, Lout);
    } else {
        long long total = (long long)B * Cout * Lout;
        int blocks = (int)((total + threads - 1) / threads);
        deconv1d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), w.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), B, Cin, L, Cout, K, (int)stride, (int)padding, (int)dilation, Lout, has_bias);
    }
    return y;
}
'''

_ext = load_inline(name='deconv1d_ext_final_co4', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['deconv1d_forward'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])

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.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv1d_transpose.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.deconv1d_forward(x.contiguous(), self.conv1d_transpose.weight.contiguous(), self.stride, self.padding, self.dilation, self.conv1d_transpose.bias is not None, b)