import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 1D convolution operation.

    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.
        output_padding (int, optional): Additional size added to one side of the output shape. Defaults to 0.
        groups (int, optional): Number of blocked connections from input channels to output channels. 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, output_padding: int = 0, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d_transpose = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, 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 = 64
in_channels = 128
out_channels = 128
kernel_size = 3
# much larger signal length for heavier workload
length = 65536

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

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.h>
#include <cuda_runtime.h>

constexpr int T_TILE = 64;
constexpr int TX = 32;
constexpr int OC_TILE = 16;
constexpr int OC_HALF = 8;
constexpr int KSIZE = 3;

__global__ void convt1d_k3_kernel(const float* __restrict__ x,
                                  const float* __restrict__ w,
                                  float* __restrict__ y,
                                  int IC, int L, int OC) {
    extern __shared__ float sh[];
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    const int n = blockIdx.z;
    const int tbase = blockIdx.x * T_TILE;
    const int ocbase = blockIdx.y * OC_TILE;
    const int Lout = L + 2;
    const int cols = T_TILE + 2;
    const int total = IC * cols;
    const float* __restrict__ xn = x + (long long)n * IC * L;

    // Cache the input window needed for this time tile.  The same window is
    // reused by all 16 output channels computed by the CTA.
    for (int idx = tid; idx < total; idx += blockDim.x * blockDim.y) {
        int ic = idx / cols;
        int c = idx - ic * cols;
        int pos = tbase + c - 2;
        float v = 0.0f;
        if ((unsigned)pos < (unsigned)L) v = xn[(long long)ic * L + pos];
        sh[idx] = v;
    }
    __syncthreads();

    const int t0 = tbase + tx;
    const int t1 = t0 + TX;
    const int oc0 = ocbase + ty;
    const int oc1 = oc0 + OC_HALF;
    const bool vt0 = t0 < Lout;
    const bool vt1 = t1 < Lout;
    const bool voc1 = oc1 < OC;

    float a00 = 0.0f, a01 = 0.0f;
    float a10 = 0.0f, a11 = 0.0f;

    if (oc0 < OC && (vt0 || vt1)) {
        #pragma unroll 4
        for (int ic = 0; ic < IC; ++ic) {
            const float* __restrict__ s = sh + ic * cols + tx;
            float x02 = s[0],      x01 = s[1],      x00 = s[2];
            float x12 = s[TX],     x11 = s[TX + 1], x10 = s[TX + 2];

            const float* __restrict__ wp0 = w + ((long long)ic * OC + oc0) * KSIZE;
            float w00 = wp0[0], w01 = wp0[1], w02 = wp0[2];
            a00 = fmaf(x00, w00, a00);
            a00 = fmaf(x01, w01, a00);
            a00 = fmaf(x02, w02, a00);
            a10 = fmaf(x10, w00, a10);
            a10 = fmaf(x11, w01, a10);
            a10 = fmaf(x12, w02, a10);

            if (voc1) {
                const float* __restrict__ wp1 = w + ((long long)ic * OC + oc1) * KSIZE;
                float w10 = wp1[0], w11 = wp1[1], w12 = wp1[2];
                a01 = fmaf(x00, w10, a01);
                a01 = fmaf(x01, w11, a01);
                a01 = fmaf(x02, w12, a01);
                a11 = fmaf(x10, w10, a11);
                a11 = fmaf(x11, w11, a11);
                a11 = fmaf(x12, w12, a11);
            }
        }

        long long base0 = ((long long)n * OC + oc0) * Lout;
        if (vt0) y[base0 + t0] = a00;
        if (vt1) y[base0 + t1] = a10;
        if (voc1) {
            long long base1 = ((long long)n * OC + oc1) * Lout;
            if (vt0) y[base1 + t0] = a01;
            if (vt1) y[base1 + t1] = a11;
        }
    }
}

torch::Tensor conv_transpose1d_forward(torch::Tensor x, torch::Tensor weight) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(), "inputs must be contiguous");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat, "float32 only");
    TORCH_CHECK(x.dim() == 3 && weight.dim() == 3, "bad tensor rank");
    TORCH_CHECK(weight.size(2) == 3, "this kernel is specialized for kernel_size=3");

    int64_t N = x.size(0), IC = x.size(1), L = x.size(2), OC = weight.size(1);
    auto y = torch::empty({N, OC, L + 2}, x.options());

    dim3 block(TX, OC_HALF);
    dim3 grid((L + 2 + T_TILE - 1) / T_TILE, (OC + OC_TILE - 1) / OC_TILE, N);
    size_t shmem = (size_t)IC * (T_TILE + 2) * sizeof(float);
    convt1d_k3_kernel<<<grid, block, shmem, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), y.data_ptr<float>(),
        (int)IC, (int)L, (int)OC);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='convt1d_k3_ext_submit',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_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, output_padding: int = 0,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv1d_transpose = nn.ConvTranspose1d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.conv_transpose1d_forward(
            x.contiguous(), self.conv1d_transpose.weight.contiguous())