```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA operator for Conv1d (stride=1, padding=0, dilation=1, groups=1).
#
# Key idea: a 1D convolution with kernel size K and stride 1 can be written as
# a sum of K batched GEMMs, where the "im2col" matrices are just shifted views
# of the input (zero-copy, no im2col buffer materialization):
#
#   out[b, o, i] = sum_k sum_c w[o, c, k] * x[b, c, i + k]
#                = sum_k (W_k @ X_k)[b, o, i]
#
# Each term is evaluated with a strided-batched GEMM (beta=1 accumulates the
# taps). A small custom CUDA kernel repacks the (Cout, Cin, K) weight into
# per-tap contiguous (Cout, Cin) slices so they map onto cuBLAS column-major
# matrices with a simple leading dimension.
# ---------------------------------------------------------------------------

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

// Repack conv1d weight (Cout, Cin, K) -> (K, Cout, Cin) so that each per-tap
// (Cout, Cin) slice is contiguous. A contiguous row-major (Cout, Cin) matrix
// is exactly a column-major (Cin, Cout) matrix with leading dimension Cin,
// which is what the GEMM below consumes.
__global__ void repack_weight_kernel(const float* __restrict__ w,
                                     float* __restrict__ wrep,
                                     const int Cout, const int Cin, const int K) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = Cout * Cin * K;
    if (idx < total) {
        const int o = idx / (Cin * K);
        const int rem = idx - o * Cin * K;
        const int c = rem / K;
        const int k = rem - c * K;
        wrep[(long long)k * Cout * Cin + (long long)o * Cin + c] = w[idx];
    }
}

static inline void cublas_check(cublasStatus_t st, const char* what) {
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cuBLAS call failed (", what, "): ", (int)st);
}

torch::Tensor conv1d_forward_cuda(torch::Tensor x, torch::Tensor w, bool use_tf32) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "x and w must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 is supported by the custom conv1d kernel");
    TORCH_CHECK(x.dim() == 3 && w.dim() == 3, "x must be (B,Cin,L), w must be (Cout,Cin,K)");

    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int64_t B    = xc.size(0);
    const int64_t Cin  = xc.size(1);
    const int64_t L    = xc.size(2);
    const int64_t Cout = wc.size(0);
    const int64_t K    = wc.size(2);
    TORCH_CHECK(wc.size(1) == Cin, "weight/input channel mismatch");
    const int64_t Lout = L - K + 1;
    TORCH_CHECK(Lout > 0, "input length too small for kernel size");

    auto out  = at::empty({B, Cout, Lout}, xc.options());
    auto wrep = at::empty({K, Cout, Cin}, wc.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // 1) Custom kernel: repack weights into per-tap contiguous slices.
    const int total = (int)(Cout * Cin * K);
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    repack_weight_kernel<<<blocks, threads, 0, stream>>>(
        wc.data_ptr<float>(), wrep.data_ptr<float>(), (int)Cout, (int)Cin, (int)K);

    // 2) Decide compute type: mirror PyTorch's cudnn.allow_tf32 policy.
    static int tf32_supported = -1;
    if (tf32_supported < 0) {
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        tf32_supported = (prop.major >= 8) ? 1 : 0;
    }
    const cublasComputeType_t ctype =
        (use_tf32 && tf32_supported) ? CUBLAS_COMPUTE_32F_FAST_TF32
                                     : CUBLAS_COMPUTE_32F_PEDANTIC;

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublas_check(cublasSetStream(handle, stream), "set stream");

    const float alpha = 1.0f, beta0 = 0.0f, beta1 = 1.0f;
    const float* xptr = xc.data_ptr<float>();
    const float* wptr = wrep.data_ptr<float>();
    float* optr = out.data_ptr<float>();

    // 3) Accumulate K strided-batched GEMMs (one per kernel tap).
    //    Column-major view:  C(Lout, Cout) = A(Lout, Cin) * B(Cin, Cout)
    //    where A is a shifted view of x (ld = L) and B is a repacked weight
    //    slice (ld = Cin). C accumulates with beta=1 after the first tap.
    for (int64_t k = 0; k < K; ++k) {
        cublas_check(cublasGemmStridedBatchedEx(
            handle, CUBLAS_OP_N, CUBLAS_OP_N,
            (int)Lout, (int)Cout, (int)Cin,
            &alpha,
            (const void*)(xptr + k), CUDA_R_32F, (int)L, (long long)(Cin * L),
            (const void*)(wptr + k * Cout * Cin), CUDA_R_32F, (int)Cin, (long long)0,
            (k == 0) ? &beta0 : &beta1,
            (void*)optr, CUDA_R_32F, (int)Lout, (long long)(Cout * Lout),
            (int)B,
            ctype, CUBLAS_GEMM_DEFAULT), "gemm");
    }
    return out;
}
"""

cpp_src = "torch::Tensor conv1d_forward_cuda(torch::Tensor x, torch::Tensor w, bool use_tf32);"

# Make sure cuBLAS can be found and linked both in full-toolkit installs and in
# pip-only (nvidia-*-cu12 wheel) environments.
_extra_ldflags = ["-lcublas"]
try:
    _nvidia_cublas_lib = os.path.abspath(
        os.path.join(os.path.dirname(torch.__file__), "..", "nvidia", "cublas", "lib")
    )
    if os.path.isdir(_nvidia_cublas_lib):
        _extra_ldflags.insert(0, "-L" + _nvidia_cublas_lib)
except Exception:
    pass

_ext = load_inline(
    name="fast_conv1d_gemm_ext",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["conv1d_forward_cuda"],
    extra_ldflags=_extra_ldflags,
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized 1D convolution. Uses the custom CUDA/GEMM operator for the
    standard configuration (stride=1, padding=0, dilation=1, groups=1) and
    falls back to the native nn.Conv1d implementation otherwise.

    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(ModelNew, self).__init__()
        # Keep an identical parameter layout/state_dict as the reference model.
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self._fast_path = (stride == 1 and padding == 0 and dilation == 1 and groups == 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast_path and x.is_cuda and x.dtype == torch.float32:
            w = self.conv1d.weight
            if w.size(1) == x.size(1) and x.size(2) >= w.size(2):
                out = _ext.conv1d_forward_cuda(x, w, bool(torch.backends.cudnn.allow_tf32))
                b = self.conv1d.bias
                if b is not None:
                    out.add_(b.view(1, -1, 1))
                return out
        return self.conv1d(x)
```