```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: direct Conv1d implemented as a tiled implicit-GEMM.
# C[oc, n] = sum_{ic, k} W[oc, ic*K+k] * X[b, ic, f(n)+k*dilation]
# where n indexes the output length and f(n) = n*stride - padding.
# Each thread block computes a BM x BN output tile with an 8x4 register
# micro-tile per thread, staging weight/input tiles in shared memory.
_conv1d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 64
#define BKMAX 24
#define TM 8
#define TN 4
#define THREADS 256

__global__ void __launch_bounds__(THREADS)
conv1d_fp32_kernel(const float* __restrict__ x,
                   const float* __restrict__ w,
                   const float* __restrict__ bias,
                   float* __restrict__ out,
                   const int C_in, const int C_out, const int L, const int K,
                   const int stride, const int padding, const int dilation,
                   const int L_out)
{
    const int b  = blockIdx.z;
    const int m0 = blockIdx.y * BM;   // output-channel tile origin
    const int n0 = blockIdx.x * BN;   // output-length tile origin

    __shared__ alignas(16) float As[BKMAX][BM + 4];
    __shared__ alignas(16) float Bs[BKMAX][BN + 4];

    const int tid = threadIdx.x;
    const int tm = tid >> 4;          // 0..15
    const int tn = tid & 15;          // 0..15

    float acc[TM][TN];
    #pragma unroll
    for (int i = 0; i < TM; ++i)
        #pragma unroll
        for (int j = 0; j < TN; ++j)
            acc[i][j] = 0.0f;

    const float* __restrict__ x_b = x + (long)b * (long)C_in * (long)L;
    const int Kdim = C_in * K;
    int BIC = BKMAX / K;
    if (BIC < 1) BIC = 1;
    const int BKeff = BIC * K;

    // Static (per-thread) load coordinates for the As tile (BKeff x BM),
    // advanced incrementally to avoid integer division in the loop.
    int a_m  = tid / BKeff;
    int a_kk = tid - a_m * BKeff;
    const int a_dm = THREADS / BKeff;
    const int a_dr = THREADS - a_dm * BKeff;

    // Static (per-thread) load coordinates for the Bs tile (BKeff x BN).
    const int b_n   = tid % BN;
    int       b_kk  = tid / BN;
    int       b_icl = b_kk / K;
    int       b_k   = b_kk - b_icl * K;
    const int gn      = n0 + b_n;
    const bool n_ok   = (gn < L_out);
    const int pos_base = gn * stride - padding;

    for (int ic0 = 0; ic0 < C_in; ic0 += BIC) {
        // ---- stage weights tile As (BKeff x BM), coalesced global reads ----
        {
            int am = a_m, akk = a_kk;
            #pragma unroll
            for (int i = 0; i < (BKMAX * BM) / THREADS; ++i) {
                if (am < BM) {
                    float v = 0.0f;
                    const int gm = m0 + am;
                    const int gk = ic0 * K + akk;
                    if (gm < C_out && gk < Kdim)
                        v = w[(long)gm * Kdim + gk];
                    As[akk][am] = v;
                }
                am += a_dm; akk += a_dr;
                if (akk >= BKeff) { akk -= BKeff; ++am; }
            }
        }
        // ---- stage input tile Bs (BKeff x BN), coalesced global reads ----
        {
            int kk = b_kk, icl = b_icl, k = b_k;
            #pragma unroll
            for (int i = 0; i < (BKMAX * BN) / THREADS; ++i) {
                if (kk < BKeff) {
                    float v = 0.0f;
                    const int ic = ic0 + icl;
                    if (n_ok && ic < C_in) {
                        const int pos = pos_base + k * dilation;
                        if (pos >= 0 && pos < L)
                            v = x_b[(long)ic * L + pos];
                    }
                    Bs[kk][b_n] = v;
                }
                kk += 4;
                k += 4;
                while (k >= K) { k -= K; ++icl; }
            }
        }
        __syncthreads();

        if (BKeff == BKMAX) {
            float ra[TM], rb[TN];
            #pragma unroll
            for (int kk = 0; kk < BKMAX; ++kk) {
                *reinterpret_cast<float4*>(&ra[0]) =
                    *reinterpret_cast<const float4*>(&As[kk][tm * TM]);
                *reinterpret_cast<float4*>(&ra[4]) =
                    *reinterpret_cast<const float4*>(&As[kk][tm * TM + 4]);
                *reinterpret_cast<float4*>(&rb[0]) =
                    *reinterpret_cast<const float4*>(&Bs[kk][tn * TN]);
                #pragma unroll
                for (int ii = 0; ii < TM; ++ii)
                    #pragma unroll
                    for (int jj = 0; jj < TN; ++jj)
                        acc[ii][jj] = fmaf(ra[ii], rb[jj], acc[ii][jj]);
            }
        } else {
            for (int kk = 0; kk < BKeff; ++kk) {
                #pragma unroll
                for (int ii = 0; ii < TM; ++ii) {
                    const float av = As[kk][tm * TM + ii];
                    #pragma unroll
                    for (int jj = 0; jj < TN; ++jj)
                        acc[ii][jj] = fmaf(av, Bs[kk][tn * TN + jj], acc[ii][jj]);
                }
            }
        }
        __syncthreads();
    }

    // ---- epilogue (with optional fused bias) ----
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int gm = m0 + tm * TM + i;
        if (gm >= C_out) continue;
        const float bv = (bias != nullptr) ? bias[gm] : 0.0f;
        float* __restrict__ optr = out + ((long)b * C_out + gm) * (long)L_out;
        #pragma unroll
        for (int j = 0; j < TN; ++j) {
            const int gnn = n0 + tn * TN + j;
            if (gnn < L_out)
                optr[gnn] = acc[i][j] + bv;
        }
    }
}

torch::Tensor conv1d_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                          int64_t stride, int64_t padding, int64_t dilation) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "conv1d_cuda: inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "conv1d_cuda: only float32 is supported");
    TORCH_CHECK(x.dim() == 3 && w.dim() == 3, "conv1d_cuda: bad ranks");

    const c10::cuda::CUDAGuard device_guard(x.device());

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

    const int B     = (int)xc.size(0);
    const int C_in  = (int)xc.size(1);
    const int L     = (int)xc.size(2);
    const int C_out = (int)wc.size(0);
    const int K     = (int)wc.size(2);
    TORCH_CHECK(wc.size(1) == C_in, "conv1d_cuda: groups != 1 not supported");
    TORCH_CHECK(K <= BKMAX, "conv1d_cuda: kernel_size too large");
    TORCH_CHECK(B <= 65535, "conv1d_cuda: batch too large");

    const long L_out = ((long)L + 2L * padding - dilation * (K - 1) - 1L) / stride + 1L;
    TORCH_CHECK(L_out > 0, "conv1d_cuda: non-positive output length");

    auto out = at::empty({B, C_out, L_out}, xc.options());

    const float* bias_ptr = nullptr;
    torch::Tensor bc;
    if (bias.defined() && bias.numel() == C_out) {
        bc = bias.contiguous();
        bias_ptr = bc.data_ptr<float>();
    }

    dim3 grid((unsigned)((L_out + BN - 1) / BN),
              (unsigned)((C_out + BM - 1) / BM),
              (unsigned)B);
    dim3 block(THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();

    conv1d_fp32_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(),
        C_in, C_out, L, K, (int)stride, (int)padding, (int)dilation,
        (int)L_out);

    return out;
}
"""

_conv1d_cpp_src = (
    "torch::Tensor conv1d_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "int64_t stride, int64_t padding, int64_t dilation);"
)

_conv1d_ext = load_inline(
    name="fused_conv1d_fp32",
    cpp_sources=_conv1d_cpp_src,
    cuda_sources=_conv1d_cuda_src,
    functions=["conv1d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv1d accelerated with a custom tiled implicit-GEMM CUDA kernel.
    Falls back to the standard PyTorch op for unsupported configurations.
    """
    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._custom_ok = (groups == 1 and kernel_size <= 24)
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv1d
        if (self._custom_ok and x.is_cuda and x.dtype == torch.float32
                and x.size(0) <= 65535):
            bias = c.bias if c.bias is not None else self._empty_bias
            return _conv1d_ext.conv1d_cuda(
                x, c.weight, bias,
                c.stride[0], c.padding[0], c.dilation[0])
        return c(x)
```