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

// ---------------------------------------------------------------------------
// Fast path: Conv1d kernel_size=3, stride=1, pad=0, dil=1, groups=1, no bias.
//   Y[b][oc][l] = sum_{k=0..2} sum_{ic} W[oc][ic][k] * X[b][ic][l+k]
// Implicit GEMM: C[OC, L] = sum_k W_k[OC, IC] * X_k[IC, L], K-dim = 3*IC.
// FP32 inputs/outputs/accumulation with TF32 tensor-core math
// (mma.m16n8k8), the same numerics class as the TF32 cuDNN reference.
//
//  * Weights are pre-transposed to [K*IC][OC] AND pre-converted to TF32 bit
//    patterns by a tiny kernel -> the mainloop has no conversions.
//  * Input tile is staged in smem as TF32 bit patterns (converted once).
//  * 3-stage cp.async pipeline for the weight tiles -> one __syncthreads per
//    k-step and deep latency hiding.
//  * Fully-coalesced 8-byte (float2) stores in the epilogue so the output
//    (2/3 of all DRAM traffic) is written with 100% sector utilization.
// ---------------------------------------------------------------------------

DEVINL uint32_t f2tf32(float x) {
    uint32_t r;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

// mma.m16n8k8.row.col  D(16x8) += A(16x8) * B(8x8), tf32 -> f32
// A frag: a0=(g,t) a1=(g+8,t) a2=(g,t+4) a3=(g+8,t+4)
// B frag: b0=(t,g) b1=(t+4,g)
// C frag: c0=(g,2t) c1=(g,2t+1) c2=(g+8,2t) c3=(g+8,2t+1)
//   g = lane>>2, t = lane&3
DEVINL void mma16808(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) {
    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));
}

DEVINL void cpa16(void* dst, const void* src) {
    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));
}

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

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

constexpr int LT  = 128;   // output-length tile per block
constexpr int OCT = 128;   // output-channel tile per block
constexpr int XLD = 136;   // padded ld of the input smem tile  (== 8 mod 32)
constexpr int ALD = 136;   // padded ld of the weight smem tile (== 8 mod 32)

// Preprocess: W[OC][IC][K] (fp32) -> Wt[K*IC][OC] (TF32 bit patterns)
__global__ void conv1d_w_pre(const float* __restrict__ w,
                             float* __restrict__ wt,
                             int IC, int OC, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int tot = OC * IC * K;
    if (idx >= tot) return;
    int oc  = idx / (IC * K);
    int rem = idx - oc * (IC * K);
    int ic  = rem / K;
    int k   = rem - ic * K;
    wt[(long)(k * IC + ic) * OC + oc] = __uint_as_float(f2tf32(w[idx]));
}

__global__ void __launch_bounds__(256, 2)
conv1d_k3_tc(const float* __restrict__ X,
             const float* __restrict__ Wt,   // [K*IC][OC], TF32 bits
             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 is a multiple of 8
    const int S      = 3 * ksteps;       // total k-tiles (K=3 taps)

    __shared__ __align__(16) float Xs[64][XLD];        // TF32 bits
    __shared__ __align__(16) float As[3][8][ALD];      // TF32 bits, 3 stages

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

    // ---- pipelined weight-tile loader: tile s = Wt[s*8 .. s*8+7][0..OC) ---
    const long wbase = (long)blockIdx.y * OCT;
    auto issue_w = [&](int s) {
        int buf = s % 3;
        int r   = tid >> 5;              // 0..7   (k row)
        int h   = (tid & 31) * 4;        // 0..124 (oc)
        cpa16(&As[buf][r][h], Wt + ((long)(s << 3) + r) * OC + wbase + h);
        cp_commit();
    };

    // Kick off the first two weight tiles before staging X (overlap).
    issue_w(0);
    issue_w(1);

    // ---- input tile -> smem (converted once to TF32 bit patterns) --------
    const long xbase = ((long)b * IC) * Lin;
    #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;
        }
    }

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

    // ---- main K loop: one __syncthreads per step --------------------------
    for (int s = 0; s < S; ++s) {
        if (s + 1 < S) cp_wait<1>(); else cp_wait<0>();
        __syncthreads();
        if (s + 2 < S) issue_w(s + 2);

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

        // A fragments: rows mw*32 + g, g+8 (m-tile 0) and +16, +24 (m-tile 1)
        // As stored as [k][oc] with ld=136 (8 mod 32) -> bank-conflict-free.
        const float* ap = &As[buf][0][0];
        uint32_t a00 = __float_as_uint(ap[(tig    ) * ALD + mw * 32 + gid     ]);
        uint32_t a10 = __float_as_uint(ap[(tig    ) * ALD + mw * 32 + gid +  8]);
        uint32_t a20 = __float_as_uint(ap[(tig + 4) * ALD + mw * 32 + gid     ]);
        uint32_t a30 = __float_as_uint(ap[(tig + 4) * ALD + mw * 32 + gid +  8]);
        uint32_t a01 = __float_as_uint(ap[(tig    ) * ALD + mw * 32 + gid + 16]);
        uint32_t a11 = __float_as_uint(ap[(tig    ) * ALD + mw * 32 + gid + 24]);
        uint32_t a21 = __float_as_uint(ap[(tig + 4) * ALD + mw * 32 + gid + 16]);
        uint32_t a31 = __float_as_uint(ap[(tig + 4) * ALD + mw * 32 + gid + 24]);

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

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

// ---------------------------------------------------------------------------
// Generic FP32 fallback (groups=1, any 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();

    static const int dev_major = at::cuda::getCurrentDeviceProperties()->major;

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

    if (fast) {
        // Pre-transpose + TF32-convert the small weight tensor.
        auto wt = at::empty({(long)K * IC * OC}, xc.options());
        int tot = K * IC * OC;
        conv1d_w_pre<<<(tot + 255) / 256, 256, 0, stream>>>(
            wc.data_ptr<float>(), wt.data_ptr<float>(), IC, OC, K);
        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_v5",
    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.

    Fast path (kernel_size=3, stride=1, padding=0, dilation=1, groups=1, no
    bias, in_channels=64, out_channels % 128 == 0): a tiled implicit-GEMM
    kernel using mma.m16n8k8 TF32 tensor-core instructions with FP32
    inputs/outputs/accumulation (same numerics class as the TF32 cuDNN
    reference). Weights are pre-transposed/pre-converted by a tiny kernel so
    the mainloop has no conversions, a 3-stage cp.async pipeline hides weight
    latency with one __syncthreads per k-step, and the epilogue writes
    fully-coalesced float2 stores. Other configurations fall back to a
    generic custom FP32 kernel, and anything unsupported (CPU tensors,
    groups != 1, non-fp32) falls back to cuDNN.
    """

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