import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 1D convolution operation with asymmetric input and a square kernel, potentially dilated and strided.

    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.
        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, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 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(x)

# Test code
batch_size = 64
in_channels = 64
out_channels = 128
kernel_size = 3
# longer signal
length = 524280
stride = 3
dilation = 4

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

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

static inline int64_t ceil_div_i64(int64_t a, int64_t b) {
    return (a + b - 1) / b;
}

// Specialized for the benchmark case: Cin=64, Cout=128, K=3, stride=3, dilation=4.
// Each block computes 256 output positions for 32 output channels for one batch item.
template<int OC_TILE>
__global__ __launch_bounds__(256, 2)
void conv1d_k3_s3_d4_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) {
    constexpr int LT = 256;
    constexpr int CIN = 64;
    constexpr int SPAN = (LT - 1) * 3 + 8 + 1; // 774 valid max offset + 1, rounded = 776

    extern __shared__ float smem[];
    float* sx = smem;          // SPAN floats of input window
    float* sw = smem + SPAN;   // OC_TILE*3 weights for current input channel

    const int tx = threadIdx.x;
    const int l0 = blockIdx.x * LT;
    const int b = blockIdx.y;
    const int oc0 = blockIdx.z * OC_TILE;
    const int l_rel = tx;
    const int l_out = l0 + l_rel;
    const bool valid_l = (l_out < Lout);

    float acc[OC_TILE];
#pragma unroll
    for (int oo = 0; oo < OC_TILE; ++oo) {
        acc[oo] = has_bias ? bias[oc0 + oo] : 0.0f;
    }

#pragma unroll 1
    for (int ic = 0; ic < CIN; ++ic) {
        const int64_t in_base = ((int64_t)b * CIN + ic) * (int64_t)L + (int64_t)l0 * 3;
        const int remain = L - (int)(l0 * 3);

        // Cooperatively load the contiguous input span.  This turns the stride-3
        // per-output access pattern into coalesced global loads and reuses it for
        // all 32 output channels in this tile.
#pragma unroll
        for (int idx = tx; idx < SPAN; idx += LT) {
            sx[idx] = (idx < remain) ? x[in_base + idx] : 0.0f;
        }

        if (tx < OC_TILE * 3) {
            const int oo = tx / 3;
            const int kk = tx - oo * 3;
            sw[tx] = w[((oc0 + oo) * CIN + ic) * 3 + kk];
        }
        __syncthreads();

        if (valid_l) {
            const int p = l_rel * 3;
            const float v0 = sx[p];
            const float v1 = sx[p + 4];
            const float v2 = sx[p + 8];
#pragma unroll
            for (int oo = 0; oo < OC_TILE; ++oo) {
                const float w0 = sw[oo * 3 + 0];
                const float w1 = sw[oo * 3 + 1];
                const float w2 = sw[oo * 3 + 2];
                float a = acc[oo];
                a = fmaf(v0, w0, a);
                a = fmaf(v1, w1, a);
                a = fmaf(v2, w2, a);
                acc[oo] = a;
            }
        }
        __syncthreads();
    }

    if (valid_l) {
#pragma unroll
        for (int oo = 0; oo < OC_TILE; ++oo) {
            y[((int64_t)b * 128 + (oc0 + oo)) * (int64_t)Lout + l_out] = acc[oo];
        }
    }
}

torch::Tensor conv1d_forward(torch::Tensor x,
                             torch::Tensor weight,
                             torch::Tensor bias,
                             int64_t stride,
                             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");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "only float32 weights supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");

    const int64_t B64 = x.size(0);
    const int64_t Cin64 = x.size(1);
    const int64_t L64 = x.size(2);
    const int64_t Cout64 = weight.size(0);
    const int64_t K64 = weight.size(2);
    const int64_t Lout64 = (L64 - dilation * (K64 - 1) - 1) / stride + 1;

    TORCH_CHECK(Cin64 == 64 && Cout64 == 128 && K64 == 3 && stride == 3 && dilation == 4,
                "this optimized kernel expects Cin=64, Cout=128, K=3, stride=3, dilation=4");
    TORCH_CHECK(B64 <= 65535 && L64 <= INT_MAX && Lout64 <= INT_MAX, "unsupported tensor size");
    if (bias.numel() != 0) {
        TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kFloat32 && bias.is_contiguous(),
                    "bias must be contiguous CUDA float32 when provided");
    }

    auto y = torch::empty({B64, Cout64, Lout64}, x.options());
    if (y.numel() == 0) return y;

    const c10::cuda::CUDAGuard device_guard(x.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    constexpr int OC_TILE = 32;
    constexpr int LT = 256;
    constexpr int SPAN = (LT - 1) * 3 + 8 + 1;
    dim3 block(LT);
    dim3 grid((unsigned)ceil_div_i64(Lout64, LT), (unsigned)B64, (unsigned)ceil_div_i64(Cout64, OC_TILE));
    size_t shmem = (SPAN + OC_TILE * 3) * sizeof(float);

    const float* bptr = (bias.numel() == 0) ? nullptr : bias.data_ptr<float>();
    conv1d_k3_s3_d4_kernel<OC_TILE><<<grid, block, shmem, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
        (int)B64, (int)L64, (int)Lout64, bias.numel() != 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='conv1d_k3s3d4_opt_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv1d_forward'],
    extra_cflags=['-O3'],
    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, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size,
                                stride=stride, dilation=dilation, bias=bias)
        self.stride = stride
        self.dilation = dilation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv1d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv1d_forward(x, self.conv1d.weight, b, self.stride, self.dilation)