import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 1D convolution operation with asymmetric input and square kernel.
    Supports padding, striding, and 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 square 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 = 16
in_channels = 32
out_channels = 64
kernel_size = 3
# long sequence
length = 131072
stride = 2
padding = 1
dilation = 2

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 conv_transpose1d_forward(torch::Tensor x, torch::Tensor weight);
"""

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

__device__ __forceinline__ float to_tf32_rn(float v) {
    unsigned int x = __float_as_uint(v);
    x += 0x00000fffU + ((x >> 13) & 1U);
    x &= 0xffffe000U;
    return __uint_as_float(x);
}

__global__ void convt1d_kernel(const float* __restrict__ x,
                               const float* __restrict__ w,
                               float* __restrict__ y,
                               int L) {
    int lane = threadIdx.x;
    int coy  = threadIdx.y;
    int t0 = blockIdx.x * 32 + lane;
    int co = blockIdx.y * 8 + coy;
    int b  = blockIdx.z;
    if (t0 >= L) return;

    int t = t0 + 1;
    float acc = 0.0f;

    #pragma unroll
    for (int ci = 0; ci < 32; ++ci) {
        const float* xx = x + ((b * 32 + ci) * L);
        const float* ww = w + (ci * 64 + co) * 3;
        if (t < L) {
            acc = fmaf(to_tf32_rn(xx[t]), to_tf32_rn(ww[0]), acc);
        }
        acc = fmaf(to_tf32_rn(xx[t - 1]), to_tf32_rn(ww[1]), acc);
        if (t >= 2) {
            acc = fmaf(to_tf32_rn(xx[t - 2]), to_tf32_rn(ww[2]), acc);
        }
    }

    int Lout = 2 * L + 1;
    float* yb = y + ((b * 64 + co) * Lout);
    int pos = 2 * t0;
    yb[pos] = 0.0f;
    yb[pos + 1] = acc;
    if (t0 == L - 1) yb[2 * L] = 0.0f;
}

torch::Tensor conv_transpose1d_forward(torch::Tensor x, torch::Tensor weight) {
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    int L = (int)xc.size(2);
    auto y = torch::empty({xc.size(0), 64, (int64_t)2 * L + 1}, xc.options());
    dim3 block(32, 8, 1);
    dim3 grid((L + 31) / 32, 8, (unsigned)xc.size(0));
    convt1d_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), y.data_ptr<float>(), L);
    return y;
}
"""

_ext = load_inline(name="convt1d_s2p1d2_tf32_final", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["conv_transpose1d_forward"],
                   extra_cuda_cflags=["-O3"], 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.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:
        return _ext.conv_transpose1d_forward(x, self.conv1d_transpose.weight)