import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 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.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        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, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, 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 = 32
in_channels = 64
out_channels = 128
kernel_size = 3
length = 131072

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

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

#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")

static inline void check_cublas(cublasStatus_t s) {
    TORCH_CHECK(s == CUBLAS_STATUS_SUCCESS, "cuBLAS call failed");
}

__global__ void pack_k3_weight_kernel(const float* __restrict__ w,
                                      float* __restrict__ wp,
                                      int OC, int IC) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 3 * IC * OC;
    if (idx >= total) return;
    int kk = idx / (IC * OC);
    int r = idx - kk * IC * OC;
    int c = r % IC;
    int oc = r / IC;
    // wp[kk] is a column-major IC x OC matrix for cuBLAS.
    wp[idx] = w[(oc * IC + c) * 3 + kk];
}

torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor weight) {
    CHECK_CUDA(x);
    CHECK_CUDA(weight);
    CHECK_CONTIGUOUS(x);
    CHECK_CONTIGUOUS(weight);
    CHECK_FLOAT(x);
    CHECK_FLOAT(weight);

    TORCH_CHECK(x.dim() == 3, "x must have shape (N,C,L)");
    TORCH_CHECK(weight.dim() == 3, "weight must have shape (OC,IC,3)");

    const int B  = (int)x.size(0);
    const int IC = (int)x.size(1);
    const int L  = (int)x.size(2);
    const int OC = (int)weight.size(0);
    TORCH_CHECK((int)weight.size(1) == IC, "weight IC mismatch");
    TORCH_CHECK((int)weight.size(2) == 3, "this optimized kernel expects kernel_size == 3");
    TORCH_CHECK(L >= 3, "input length must be at least 3");
    const int Lout = L - 2;

    c10::cuda::CUDAGuard device_guard(x.device());
    auto y = torch::empty({B, OC, Lout}, x.options());
    auto wp = torch::empty({3, IC, OC}, weight.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int total = 3 * IC * OC;
    pack_k3_weight_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(), wp.data_ptr<float>(), OC, IC);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    check_cublas(cublasSetStream(handle, stream));

    cublasMath_t old_math;
    check_cublas(cublasGetMathMode(handle, &old_math));
    // Use true FP32 accumulation/multiply for tight 1e-4 correctness.
    check_cublas(cublasSetMathMode(handle, CUBLAS_PEDANTIC_MATH));

    const float alpha = 1.0f;
    const float beta0 = 0.0f;
    const float beta1 = 1.0f;

    const float* xptr = x.data_ptr<float>();
    const float* bptr = wp.data_ptr<float>();
    float* yptr = y.data_ptr<float>();

    // Interpret contiguous row-major tensors as column-major matrices:
    // A: (Lout x IC) column-major view of x[:, :, kk:kk+Lout], lda=L.
    // B: (IC x OC) packed weights for kernel position kk.
    // C: (Lout x OC) column-major view of y, which is contiguous (OC,Lout).
    const int m = Lout;
    const int n = OC;
    const int k = IC;
    const int lda = L;
    const int ldb = IC;
    const int ldc = Lout;
    const long long strideA = (long long)IC * (long long)L;
    const long long strideB = 0;
    const long long strideC = (long long)OC * (long long)Lout;

    for (int kk = 0; kk < 3; ++kk) {
        const float* A = xptr + kk;
        const float* Bmat = bptr + (long long)kk * IC * OC;
        const float* beta = (kk == 0) ? &beta0 : &beta1;
        check_cublas(cublasSgemmStridedBatched(
            handle,
            CUBLAS_OP_N, CUBLAS_OP_N,
            m, n, k,
            &alpha,
            A, lda, strideA,
            Bmat, ldb, strideB,
            beta,
            yptr, ldc, strideC,
            B));
    }

    check_cublas(cublasSetMathMode(handle, old_math));
    return y;
}
'''

_ext = load_inline(
    name='conv1d_k3_cublas_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv1d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    with_cuda=True,
    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, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self._stride = stride
        self._padding = padding
        self._dilation = dilation
        self._groups = groups
        self._kernel_size = kernel_size
        self._has_bias = bias

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Optimized path for the benchmark configuration: fp32, contiguous CUDA,
        # groups=1, stride=1, padding=0, dilation=1, kernel_size=3, no bias.
        if (self._stride != 1 or self._padding != 0 or self._dilation != 1 or
            self._groups != 1 or self._kernel_size != 3 or self.conv1d.bias is not None):
            raise NotImplementedError('ModelNew optimized kernel supports only kernel_size=3, stride=1, padding=0, dilation=1, groups=1, bias=False')
        return _ext.conv1d_forward(x, self.conv1d.weight)