```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

_conv1d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <cstring>

#define DEVINL __device__ __forceinline__

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
#define CONV_HAS_TF32 1
#else
#define CONV_HAS_TF32 0
#endif

// ---------------------------------------------------------------------------
// Fast path: Conv1d, kernel_size=3, stride=1, pad=0, dil=1, groups=1, no bias
// Block tile: 128 OC x 128 L, 8 warps each computing 32 OC x 64 L with
// mma.m16n8k8.tf32 tensor core instructions (TF32 math, FP32 in/out).
// ---------------------------------------------------------------------------

DEVINL uint32_t f2tf32(float x) {
#if CONV_HAS_TF32
    uint32_t r;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
#else
    uint32_t r; std::memcpy(&r, &x, 4); return r;
#endif
}

DEVINL void mma_tf32(float &c0, float &c1, float &c2, float &c3,
                     uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                     uint32_t b0, uint32_t b1) {
#if CONV_HAS_TF32
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
#else
    float a[4], b[2];
    std::memcpy(a, &a0, 4); std::memcpy(a+1, &a1, 4);
    std::memcpy(a+2, &a2, 4); std::memcpy(a+3, &a3, 4);
    std::memcpy(b, &b0, 4); std::memcpy(b+1, &b1, 4);
    c0 += a[0]*b[0] + a[2]*b[0];  // scalar fallback (never used on sm80+)
    c1 += a[0]*b[1] + a[2]*b[1];
    c2 += a[1]*b[0] + a[3]*b[0];
    c3 += a[1]*b[1] + a[3]*b[1];
#endif
}

DEVINL void cpa16(void* dst, const void* src) {
#if CONV_HAS_TF32
    uint32_t d = (uint32_t)__cvta_generic_to_shared(dst);
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" :: "r"(d), "l"(src));
#else
    const float4* s = (const float4*)src;
    *(float4*)dst = *s;
#endif
}

DEVINL void cp_commit() {
#if CONV_HAS_TF32
    asm volatile("cp.async.commit_group;\n" ::);
#endif
}

template <int N>
DEVINL void cp_wait() {
#if CONV_HAS_TF32
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
#endif
}

constexpr int LT  = 128;   // output-length tile per block
constexpr int OCT = 128;   // output-channel tile per block
constexpr int XLD = 136;   // padded leading dim of input smem tile (conflict-free)
constexpr int ALD = 12;    // padded leading dim of weight smem slice

__global__ void __launch_bounds__(256, 2)
conv1d_k3_tc(const float* __restrict__ X,
             const float* __restrict__ Wt,   // [OC][k][ic] contiguous (kk = k*IC + ic)
             float* __restrict__ Y,
             int Bsz, int IC, int OC, int Lin, int Lout) {
    const int nlt   = (Lout + LT - 1) / LT;
    const int ltile = blockIdx.x % nlt;
    const int b     = blockIdx.x / nlt;
    const int l0    = ltile * LT;

    const int ksteps = IC >> 3;   // IC assumed multiple of 8
    const int S      = 3 * ksteps;

    __shared__ __align__(16) float Xs[64][XLD];   // input tile for this block
    __shared__ __align__(16) float As[2][OCT][ALD]; // double-buffered weight slice

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int mw   = warp & 3;                 // warp row: oc offset = mw*32
    const int nw   = warp >> 2;                // warp col: l  offset = nw*64
    const int gid  = lane >> 2;
    const int tig  = lane & 3;

    const long xbase = ((long)b * IC) * Lin;

    // ---- Load input tile into shared memory (as TF32 bit patterns) --------
    // Row ic covers input positions [l0, l0+131] (33 float4 chunks).
    #pragma unroll
    for (int i = 0; i < 9; ++i) {
        int e  = tid + i * 256;
        int r  = e / 33;
        int c  = e - r * 33;
        if (r < IC) {
            int p = l0 + c * 4;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (p + 3 < Lin) {
                v = *reinterpret_cast<const float4*>(X + xbase + (long)r * Lin + p);
            } else if (p < Lin) {
                const float* s = X + xbase + (long)r * Lin;
                v.x = s[p];
                if (p + 1 < Lin) v.y = s[p+1];
                if (p + 2 < Lin) v.z = s[p+2];
            }
            v.x = __uint_as_float(f2tf32(v.x));
            v.y = __uint_as_float(f2tf32(v.y));
            v.z = __uint_as_float(f2tf32(v.z));
            v.w = __uint_as_float(f2tf32(v.w));
            *reinterpret_cast<float4*>(&Xs[r][c * 4]) = v;
        }
    }

    // ---- Pipelined weight-slice loader ------------------------------------
    const long wbase = (long)blockIdx.y * OCT * (3 * IC);
    auto issue_w = [&](int s) {
        int buf = s & 1;
        int r   = tid >> 1;
        int h   = (tid & 1) * 4;
        int kk  = (s << 3) + h;    // == k*IC + ic0 + h
        cpa16(&As[buf][r][h], Wt + wbase + (long)r * (3 * IC) + kk);
        cp_commit();
    };
    issue_w(0);

    float acc[8][4];
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        #pragma unroll
        for (int q = 0; q < 4; ++q) acc[j][q] = 0.f;
    }

    // ---- Main K loop -------------------------------------------------------
    for (int s = 0; s < S; ++s) {
        if (s + 1 < S) { issue_w(s + 1); cp_wait<1>(); }
        else           { cp_wait<0>(); }
        __syncthreads();

        const int k   = s / ksteps;
        const int ic0 = (s - (s / ksteps) * ksteps) << 3;
        const int buf = s & 1;

        // A fragments (2 x 16x8), convert weights to TF32
        const float* arow = &As[buf][mw * 32][0];
        uint32_t a00 = f2tf32(arow[(gid    ) * ALD + tig]);
        uint32_t a10 = f2tf32(arow[(gid + 8) * ALD + tig]);
        uint32_t a20 = f2tf32(arow[(gid    ) * ALD + tig + 4]);
        uint32_t a30 = f2tf32(arow[(gid + 8) * ALD + tig + 4]);
        uint32_t a01 = f2tf32(arow[(gid + 16) * ALD + tig]);
        uint32_t a11 = f2tf32(arow[(gid + 24) * ALD + tig]);
        uint32_t a21 = f2tf32(arow[(gid + 16) * ALD + tig + 4]);
        uint32_t a31 = f2tf32(arow[(gid + 24) * ALD + tig + 4]);

        const int nbase = nw * 64 + k;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            int col = nbase + (j << 3) + gid;
            uint32_t b0 = __float_as_uint(Xs[ic0 + tig    ][col]);
            uint32_t b1 = __float_as_uint(Xs[ic0 + tig + 4][col]);
            mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                     a00, a10, a20, a30, b0, b1);
            mma_tf32(acc[j][4], acc[j][5], acc[j][6], acc[j][7],
                     a01, a11, a21, a31, b0, b1);
        }
        __syncthreads();
    }

    // ---- Epilogue: masked FP32 stores --------------------------------------
    const int lvalid = min(LT, Lout - l0);
    float* yp = Y + ((long)b * OC + blockIdx.y * OCT) * Lout + l0;
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        int l = nw * 64 + (j << 3) + tig * 2;
        if (l < lvalid) {
            const bool s1 = (l + 1) < lvalid;
            float* r0 = yp + (long)(mw * 32 + gid     ) * Lout + l;
            r0[0] = acc[j][0]; if (s1) r0[1] = acc[j][1];
            float* r1 = yp + (long)(mw * 32 + gid +  8) * Lout + l;
            r1[0] = acc[j][2]; if (s1) r1[1] = acc[j][3];
            float* r2 = yp + (long)(mw * 32 + gid + 16) * Lout + l;
            r2[0] = acc[j][4]; if (s1) r2[1] = acc[j][5];
            float* r3 = yp + (long)(mw * 32 + gid + 24) * Lout + l;
            r3[0] = acc[j][6]; if (s1) r3[1] = acc[j][7];
        }
    }
}

// ---------------------------------------------------------------------------
// Generic fallback path (groups=1, any other conv1d configuration).
// ---------------------------------------------------------------------------
__global__ void conv1d_generic(const float* __restrict__ x,
                               const float* __restrict__ w,
                               const float* __restrict__ bias,
                               float* __restrict__ y,
                               int Bsz, int IC, int OC, int Lin, int Lout,
                               int K, int stride, int pad, int dil) {
    int l  = blockIdx.x * blockDim.x + threadIdx.x;
    int oc = blockIdx.y;
    int b  = blockIdx.z;
    if (l >= Lout) return;
    float sum = bias ? bias[oc] : 0.f;
    const float* xp = x + (long)b * IC * Lin;
    const float* wp = w + (long)oc * IC * K;
    int p0 = l * stride - pad;
    for (int ic = 0; ic < IC; ++ic) {
        const float* xrow = xp + (long)ic * Lin;
        const float* wrow = wp + (long)ic * K;
        int p = p0;
        for (int k = 0; k < K; ++k, p += dil) {
            if (p >= 0 && p < Lin) sum += xrow[p] * wrow[k];
        }
    }
    y[((long)b * OC + oc) * Lout + l] = sum;
}

// ---------------------------------------------------------------------------
torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias,
                             int64_t stride, int64_t pad, int64_t dil,
                             int64_t groups) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "inputs must be float32");
    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int Bsz = xc.size(0), IC = xc.size(1), Lin = xc.size(2);
    const int OC  = wc.size(0), K = wc.size(2);
    const int Lout = (Lin + 2 * (int)pad - (int)dil * (K - 1) - 1) / (int)stride + 1;

    auto y = at::empty({Bsz, OC, Lout}, xc.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    const bool fast = (K == 3) && stride == 1 && pad == 0 && dil == 1 &&
                      groups == 1 && !bias.has_value() &&
                      IC == 64 && (OC % OCT) == 0 && (Lin & 3) == 0;

    if (fast) {
        // Rearrange weight once: [OC][ic][k] -> [OC][k][ic]
        auto wt = wc.permute({0, 2, 1}).contiguous();
        int nlt = (Lout + LT - 1) / LT;
        dim3 grid(nlt * Bsz, OC / OCT);
        conv1d_k3_tc<<<grid, 256, 0, stream>>>(
            xc.data_ptr<float>(), wt.data_ptr<float>(), y.data_ptr<float>(),
            Bsz, IC, OC, Lin, Lout);
    } else {
        TORCH_CHECK(groups == 1, "custom kernel supports groups==1 only");
        dim3 grid((Lout + 255) / 256, OC, Bsz);
        const float* bp = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
        conv1d_generic<<<grid, 256, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bp, y.data_ptr<float>(),
            Bsz, IC, OC, Lin, Lout, K, (int)stride, (int)pad, (int)dil);
    }
    return y;
}
"""

_conv1d_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias,
                             int64_t stride, int64_t pad, int64_t dil,
                             int64_t groups);
"""

_conv1d_mod = load_inline(
    name="conv1d_tf32_tc",
    cpp_sources=_conv1d_cpp_src,
    cuda_sources=_conv1d_cuda_src,
    functions=["conv1d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv1d accelerated with a custom TF32 tensor-core CUDA kernel.
    Computes in FP32 precision using TF32 tensor-core math (same class of
    numerics as the cuDNN reference). Falls back to a generic custom kernel
    for configurations outside the fast path, and to cuDNN for anything the
    custom kernels do not cover (CPU tensors, groups != 1, non-fp32).
    """

    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.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        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:
        if x.is_cuda and x.dtype == torch.float32 and self.groups == 1:
            bias = self.conv1d.bias
            bias_arg = bias if (bias is not None and bias.numel() > 0) else None
            return _conv1d_mod.conv1d_forward(
                x, self.conv1d.weight, bias_arg,
                self.stride, self.padding, self.dilation, self.groups)
        return self.conv1d(x)
```