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

_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ __launch_bounds__(256, 2)
void deconv1d_fixed_vec8_kernel(const float* __restrict__ x,
                                const float* __restrict__ w,
                                const float* __restrict__ bias,
                                float* __restrict__ y,
                                int B, int L, int Lout, int has_bias) {
    // Specialized for Cin=32, Cout=64, K=5, stride=1, padding=0, dilation=3.
    // One thread computes 8 output channels for one (batch, position).  The
    // linearization keeps positions contiguous, so each of the 8 stores is
    // coalesced across the warp.
    const int groups = 8; // 64 / 8
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    const long long total = (long long)B * groups * Lout;
    if (idx >= total) return;

    int l = (int)(idx % Lout);
    long long t = idx / Lout;
    int g = (int)(t % groups);
    int b = (int)(t / groups);
    int oc = g * 8;

    float a0, a1, a2, a3, a4, a5, a6, a7;
    if (has_bias) {
        a0 = __ldg(bias + oc + 0); a1 = __ldg(bias + oc + 1);
        a2 = __ldg(bias + oc + 2); a3 = __ldg(bias + oc + 3);
        a4 = __ldg(bias + oc + 4); a5 = __ldg(bias + oc + 5);
        a6 = __ldg(bias + oc + 6); a7 = __ldg(bias + oc + 7);
    } else {
        a0 = a1 = a2 = a3 = a4 = a5 = a6 = a7 = 0.0f;
    }

    #pragma unroll
    for (int ic = 0; ic < 32; ++ic) {
        const float* __restrict__ xp = x + ((b * 32 + ic) * L);
        const float* __restrict__ wp0 = w + (ic * 64 + oc) * 5;

        int il = l;
        if ((unsigned)il < (unsigned)L) {
            float xv = __ldg(xp + il);
            const float* wp = wp0 + 0;
            a0 = fmaf(xv, __ldg(wp + 0*5), a0);
            a1 = fmaf(xv, __ldg(wp + 1*5), a1);
            a2 = fmaf(xv, __ldg(wp + 2*5), a2);
            a3 = fmaf(xv, __ldg(wp + 3*5), a3);
            a4 = fmaf(xv, __ldg(wp + 4*5), a4);
            a5 = fmaf(xv, __ldg(wp + 5*5), a5);
            a6 = fmaf(xv, __ldg(wp + 6*5), a6);
            a7 = fmaf(xv, __ldg(wp + 7*5), a7);
        }
        il = l - 3;
        if ((unsigned)il < (unsigned)L) {
            float xv = __ldg(xp + il);
            const float* wp = wp0 + 1;
            a0 = fmaf(xv, __ldg(wp + 0*5), a0);
            a1 = fmaf(xv, __ldg(wp + 1*5), a1);
            a2 = fmaf(xv, __ldg(wp + 2*5), a2);
            a3 = fmaf(xv, __ldg(wp + 3*5), a3);
            a4 = fmaf(xv, __ldg(wp + 4*5), a4);
            a5 = fmaf(xv, __ldg(wp + 5*5), a5);
            a6 = fmaf(xv, __ldg(wp + 6*5), a6);
            a7 = fmaf(xv, __ldg(wp + 7*5), a7);
        }
        il = l - 6;
        if ((unsigned)il < (unsigned)L) {
            float xv = __ldg(xp + il);
            const float* wp = wp0 + 2;
            a0 = fmaf(xv, __ldg(wp + 0*5), a0);
            a1 = fmaf(xv, __ldg(wp + 1*5), a1);
            a2 = fmaf(xv, __ldg(wp + 2*5), a2);
            a3 = fmaf(xv, __ldg(wp + 3*5), a3);
            a4 = fmaf(xv, __ldg(wp + 4*5), a4);
            a5 = fmaf(xv, __ldg(wp + 5*5), a5);
            a6 = fmaf(xv, __ldg(wp + 6*5), a6);
            a7 = fmaf(xv, __ldg(wp + 7*5), a7);
        }
        il = l - 9;
        if ((unsigned)il < (unsigned)L) {
            float xv = __ldg(xp + il);
            const float* wp = wp0 + 3;
            a0 = fmaf(xv, __ldg(wp + 0*5), a0);
            a1 = fmaf(xv, __ldg(wp + 1*5), a1);
            a2 = fmaf(xv, __ldg(wp + 2*5), a2);
            a3 = fmaf(xv, __ldg(wp + 3*5), a3);
            a4 = fmaf(xv, __ldg(wp + 4*5), a4);
            a5 = fmaf(xv, __ldg(wp + 5*5), a5);
            a6 = fmaf(xv, __ldg(wp + 6*5), a6);
            a7 = fmaf(xv, __ldg(wp + 7*5), a7);
        }
        il = l - 12;
        if ((unsigned)il < (unsigned)L) {
            float xv = __ldg(xp + il);
            const float* wp = wp0 + 4;
            a0 = fmaf(xv, __ldg(wp + 0*5), a0);
            a1 = fmaf(xv, __ldg(wp + 1*5), a1);
            a2 = fmaf(xv, __ldg(wp + 2*5), a2);
            a3 = fmaf(xv, __ldg(wp + 3*5), a3);
            a4 = fmaf(xv, __ldg(wp + 4*5), a4);
            a5 = fmaf(xv, __ldg(wp + 5*5), a5);
            a6 = fmaf(xv, __ldg(wp + 6*5), a6);
            a7 = fmaf(xv, __ldg(wp + 7*5), a7);
        }
    }

    float* yp = y + (b * 64 + oc) * Lout + l;
    yp[0 * Lout] = a0; yp[1 * Lout] = a1;
    yp[2 * Lout] = a2; yp[3 * Lout] = a3;
    yp[4 * Lout] = a4; yp[5 * Lout] = a5;
    yp[6 * Lout] = a6; yp[7 * Lout] = a7;
}

__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 Lout,
                             int stride, int padding, int dilation, int 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 l = (int)(idx % Lout);
    long long t = idx / Lout;
    int oc = (int)(t % Cout);
    int b = (int)(t / Cout);
    float acc = has_bias ? __ldg(bias + oc) : 0.0f;
    for (int ic = 0; ic < Cin; ++ic) {
        for (int k = 0; k < K; ++k) {
            int num = l + padding - k * dilation;
            if (num % stride == 0) {
                int il = num / stride;
                if ((unsigned)il < (unsigned)L) {
                    acc = fmaf(__ldg(x + (b * Cin + ic) * L + il),
                               __ldg(w + (ic * Cout + oc) * K + k), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor conv_transpose1d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                       int64_t stride, int64_t padding, int64_t dilation) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    int B = (int)x.size(0);
    int Cin = (int)x.size(1);
    int L = (int)x.size(2);
    int Cout = (int)weight.size(1);
    int K = (int)weight.size(2);
    int Lout = (int)((L - 1) * stride - 2 * padding + dilation * (K - 1) + 1);
    auto y = torch::empty({B, Cout, Lout}, x.options());
    int has_bias = (bias.numel() != 0);
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;

    if (Cin == 32 && Cout == 64 && K == 5 && stride == 1 && padding == 0 && dilation == 3) {
        const int threads = 256;
        long long total = (long long)B * 8 * Lout;
        int blocks = (int)((total + threads - 1) / threads);
        deconv1d_fixed_vec8_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp,
                                                        y.data_ptr<float>(), B, L, Lout, has_bias);
    } else {
        const int threads = 256;
        long long total = (long long)B * Cout * Lout;
        int blocks = (int)((total + threads - 1) / threads);
        deconv1d_generic_kernel<<<blocks, threads>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bp,
                                                     y.data_ptr<float>(), B, Cin, L, Cout, K, Lout,
                                                     (int)stride, (int)padding, (int)dilation, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='conv_transpose1d_dil3_ext',
    cpp_sources='',
    cuda_sources=_src,
    functions=['conv_transpose1d_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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.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.conv_transpose1d_forward(x, self.conv1d_transpose.weight, b, self.stride, self.padding, self.dilation)