```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 torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernel for 1D convolution (stride=1, padding=0, dilation=1, groups=1).
#
# Strategy:
#   * Each thread block handles a tile of (8 output channels) x (512 positions)
#     of one batch element.
#   * Each thread computes 4 consecutive output positions for all 8 output
#     channels (32 independent FMA chains -> high ILP).
#   * The kernel weights are repacked once per forward into a
#     [C_in, C_out, 4] zero-padded layout so every thread can fetch all 3 taps
#     of one (channel, output-channel) pair with a single 16-byte vector load
#     (broadcast across the warp, served from L1).
#   * Input values are read directly through __ldg (coalesced across the warp,
#     overlap between neighboring threads captured by L1), keeping the inner
#     loop free of __syncthreads.
#
# This yields ~96 FMA per ~14 load instructions per channel step, i.e. the
# kernel is firmly FMA-issue bound rather than memory bound.

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

constexpr int TILE_O  = 8;    // output channels per block
constexpr int POS     = 4;    // output positions per thread
constexpr int THREADS = 128;  // threads per block (block covers THREADS*POS positions)

template <int K>
__global__ __launch_bounds__(THREADS)
void conv1d_kernel(const float* __restrict__ x,
                   const float* __restrict__ w4,    // [C_in, C_out, 4] zero-padded weights
                   const float* __restrict__ bias,  // nullptr when no bias
                   float* __restrict__ out,
                   int C_in, int C_out, int L, int L_out) {
    const int tx = threadIdx.x;
    const int i  = blockIdx.x * (THREADS * POS) + tx * POS;   // first position of this thread
    const int o0 = blockIdx.y * TILE_O;                        // first output channel of this block
    const int b  = blockIdx.z;                                 // batch index

    // Clamped output-channel indices for weight loads (avoids OOB on the edge block).
    int oc[TILE_O];
    #pragma unroll
    for (int o = 0; o < TILE_O; ++o) {
        int t = o0 + o;
        oc[o] = (t < C_out) ? t : 0;
    }

    float acc[TILE_O][POS];
    #pragma unroll
    for (int o = 0; o < TILE_O; ++o)
        #pragma unroll
        for (int p = 0; p < POS; ++p)
            acc[o][p] = 0.0f;

    const long xb = (long)b * (long)C_in * (long)L;

    for (int c = 0; c < C_in; ++c) {
        const float* xp = x + xb + (long)c * (long)L;

        // POS + K - 1 input values feed this thread's POS outputs for channel c.
        float xv[POS + K - 1];
        #pragma unroll
        for (int t = 0; t < POS + K - 1; ++t) {
            const int idx = i + t;
            xv[t] = (idx < L) ? __ldg(xp + idx) : 0.0f;
        }

        // All K taps of one (c, o) pair in a single 16B broadcast load.
        const float4* wrow = reinterpret_cast<const float4*>(w4) + (long)c * (long)C_out;

        #pragma unroll
        for (int o = 0; o < TILE_O; ++o) {
            const float4 wv = __ldg(wrow + oc[o]);
            const float* wf = reinterpret_cast<const float*>(&wv);
            #pragma unroll
            for (int p = 0; p < POS; ++p) {
                float a = acc[o][p];
                #pragma unroll
                for (int k = 0; k < K; ++k)
                    a = fmaf(xv[p + k], wf[k], a);
                acc[o][p] = a;
            }
        }
    }

    #pragma unroll
    for (int o = 0; o < TILE_O; ++o) {
        const int och = o0 + o;
        if (och < C_out) {
            const float bv = bias ? __ldg(bias + och) : 0.0f;
            float* op = out + ((long)b * (long)C_out + (long)och) * (long)L_out;
            #pragma unroll
            for (int p = 0; p < POS; ++p) {
                const int idx = i + p;
                if (idx < L_out) op[idx] = acc[o][p] + bv;
            }
        }
    }
}

torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w4, torch::Tensor bias, int64_t k) {
    const int B     = x.size(0);
    const int C_in  = x.size(1);
    const int L     = x.size(2);
    const int C_out = w4.size(1);
    const int L_out = L - (int)k + 1;

    auto out = torch::empty({B, C_out, L_out > 0 ? (long)L_out : 0L}, x.options());
    if (B == 0 || C_out == 0 || L_out <= 0) return out;

    dim3 block(THREADS);
    dim3 grid((L_out + THREADS * POS - 1) / (THREADS * POS),
              (C_out + TILE_O - 1) / TILE_O,
              B);
    auto stream = at::cuda::getCurrentCUDAStream();

    const float* xp = x.data_ptr<float>();
    const float* wp = w4.data_ptr<float>();
    const float* bp = (bias.numel() > 0) ? bias.data_ptr<float>() : nullptr;
    float* op = out.data_ptr<float>();

    switch (k) {
        case 1: conv1d_kernel<1><<<grid, block, 0, stream>>>(xp, wp, bp, op, C_in, C_out, L, L_out); break;
        case 2: conv1d_kernel<2><<<grid, block, 0, stream>>>(xp, wp, bp, op, C_in, C_out, L, L_out); break;
        case 3: conv1d_kernel<3><<<grid, block, 0, stream>>>(xp, wp, bp, op, C_in, C_out, L, L_out); break;
        case 4: conv1d_kernel<4><<<grid, block, 0, stream>>>(xp, wp, bp, op, C_in, C_out, L, L_out); break;
        default: TORCH_CHECK(false, "conv1d_forward: unsupported kernel size ", k);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = (
    "torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w4, torch::Tensor bias, int64_t k);"
)

_conv1d_ext = load_inline(
    name="conv1d_custom",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv1d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv1d accelerated with a custom CUDA kernel.

    The custom kernel handles the common fast path (stride=1, padding=0,
    dilation=1, groups=1, fp32 CUDA tensors, kernel_size <= 4). Any other
    configuration transparently falls back to the native PyTorch operator, so
    the module remains fully functional for all valid inputs.
    """
    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._w4_cache = None
        self._w4_key = None

    def _packed_weight(self, w: torch.Tensor, k: int) -> torch.Tensor:
        # Repack [C_out, C_in, K] -> zero-padded [C_in, C_out, 4] for float4 loads.
        # Cached on (data_ptr, version, device) so repeated inference is free while
        # any in-place weight update invalidates the cache.
        key = (w.data_ptr(), w._version, w.device)
        if self._w4_key != key or self._w4_cache is None:
            self._w4_cache = F.pad(w.permute(1, 0, 2), (0, 4 - k)).contiguous()
            self._w4_key = key
        return self._w4_cache

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv1d
        k = conv.kernel_size[0]
        w = conv.weight
        if (x.is_cuda and x.dtype == torch.float32 and w.dtype == torch.float32
                and conv.groups == 1 and conv.stride[0] == 1 and conv.padding[0] == 0
                and conv.dilation[0] == 1 and 1 <= k <= 4
                and x.size(0) <= 65535):
            w4 = self._packed_weight(w, k)
            bias = conv.bias if conv.bias is not None else x.new_empty(0)
            return _conv1d_ext.conv1d_forward(x.contiguous(), w4, bias, k)
        return conv(x)
```