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

#define DEVINL __device__ __forceinline__

// ---------------------------------------------------------------------------
// Fast path: Conv1d k=3, s=1, p=0, d=1, g=1, no bias.
//   Y[b][oc][l] = sum_{k=0..2} sum_{ic} W[oc][ic][k] * X[b][ic][l+k]
// Implicit GEMM with TF32 tensor cores (mma.m16n8k8), FP32 in/out/accum.
//
// v6 design (fixes the v5 stall sources):
//   * weights pre-transposed to [3*IC][OC] AND pre-converted to TF32 bits,
//     then double-buffered in smem with cp.async (2 stages, one sync/k-step).
//   * input tile staged once into smem as TF32 bits ([k][l], padded ld=136).
//   * inner j-loop unrolled by 4 so 4 independent (load,load,mma,mma) groups
//     are in flight -> hides smem->register latency without extra registers.
//   * grid rasterization is l-major within a batch so concurrently-running
//     blocks share the same batch -> better DRAM page locality, and their
//     cp.async weight streams are in-phase -> bursty L2-friendly traffic.
//   * float2 fully-coalesced epilogue stores.
// ---------------------------------------------------------------------------

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

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 input smem tile  (8 mod 32)
constexpr int ALD = 136;   // padded ld of weight smem tile (8 mod 32)

// W[OC][IC][K] fp32 -> Wt[K*IC][OC] with 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;          // l-major rasterization
    const int b     = blockIdx.x / nlt;
    const int l0    = ltile * LT;

    const int ksteps = IC >> 3;                  // IC % 8 == 0
    const int S      = 3 * ksteps;

    __shared__ __align__(16) float Xs[64][XLD];   // TF32 bits
    __shared__ __align__(16) float As[2][8][ALD]; // TF32 bits, 2 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;

    // ---- 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 & 1;
        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();
    };
    issue_w(0);

    // ---- stage input tile -> smem (converted once to TF32 bits) -----------
    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;

    // precomputed smem source pointers (immediate-offset addressing)
    const float* abase = &As[0][0][mw * 32 + gid];
    const int   aoff0  = tig * ALD;
    const int   aoff1  = (tig + 4) * ALD;
    const int   ktap   = 0;  // placeholder to keep structure clear
    (void)ktap;

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

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

        const float* ap = abase + buf * (8 * ALD);
        uint32_t a00 = __float_as_uint(ap[aoff0     ]);
        uint32_t a10 = __float_as_uint(ap[aoff0 +  8]);
        uint32_t a20 = __float_as_uint(ap[aoff1     ]);
        uint32_t a30 = __float_as_uint(ap[aoff1 +  8]);
        uint32_t a01 = __float_as_uint(ap[aoff0 + 16]);
        uint32_t a11 = __float_as_uint(ap[aoff0 + 24]);
        uint32_t a21 = __float_as_uint(ap[aoff1 + 16]);
        uint32_t a31 = __float_as_uint(ap[aoff1 + 24]);

        const int nbase = nw * 64 + k + gid;
        const float* x0 = &Xs[ic0 + tig    ][nbase];
        const float* x1 = &Xs[ic0 + tig + 4][nbase];

        // 2 groups of 4: deep (LDS,LDS,MMA,MMA) pipelines, 16 B-regs live.
        #pragma unroll
        for (int j = 0; j < 8; j += 4) {
            uint32_t b00 = __float_as_uint(x0[(j + 0) << 3]);
            uint32_t b01 = __float_as_uint(x1[(j + 0) << 3]);
            uint32_t b10 = __float_as_uint(x0[(j + 1) << 3]);
            uint32_t b11 = __float_as_uint(x1[(j + 1) << 3]);
            uint32_t b20 = __float_as_uint(x0[(j + 2) << 3]);
            uint32_t b21 = __float_as_uint(x1[(j + 2) << 3]);
            uint32_t b30 = __float_as_uint(x0[(j + 3) << 3]);
            uint32_t b31 = __float_as_uint(x1[(j + 3) << 3]);
            mma16808(acc[j+0][0], acc[j+0][1], acc[j+0][2], acc[j+0][3],
                     a00, a10, a20, a30, b00, b01);
            mma16808(acc[j+0][4], acc[j+0][5], acc[j+0][6], acc[j+0][7],
                     a01, a11, a21, a31, b00, b01);
            mma16808(acc[j+1][0], acc[j+1][1], acc[j+1][2], acc[j+1][3],
                     a00, a10, a20, a30, b10, b11);
            mma16808(acc[j+1][4], acc[j+1][5], acc[j+1][6], acc[j+1][7],
                     a01, a11, a21, a31, b10, b11);
            mma16808(acc[j+2][0], acc[j+2][1], acc[j+2][2], acc[j+2][3],
                     a00, a10, a20, a30, b20, b21);
            mma16808(acc[j+2][4], acc[j+2][5], acc[j+2][6], acc[j+2][7],
                     a01, a11, a21, a31, b20, b21);
            mma16808(acc[j+3][0], acc[j+3][1], acc[j+3][2], acc[j+3][3],
                     a00, a10, a20, a30, b30, b31);
            mma16808(acc[j+3][4], acc[j+3][5], acc[j+3][6], acc[j+3][7],
                     a01, a11, a21, a31, b30, b31);
        }
    }

    // ---- 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) {
        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_v6",
    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): tiled implicit GEMM with
    mma.m16n8k8 TF32 tensor-core math and FP32 inputs/outputs/accumulation
    (same numerics class as the TF32 cuDNN reference). Improvements over the
    previous version: double-buffered cp.async weight pipeline (in-phase
    across neighboring blocks due to l-major rasterization), immediate-offset
    shared-memory addressing, and a 4x-unrolled inner loop that keeps four
    independent load/mma groups in flight to hide smem->register latency.
    Other configurations use a generic custom FP32 kernel; CPU tensors,
    groups != 1 or non-fp32 inputs fall 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)
```