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

# ---------------------------------------------------------------------------
# Fused Conv1d (stride=1, padding=0, dilation=1, groups=1) on tensor cores.
#
#   out[b, o, i] = bias[o] + sum_{t,c} w[o, c, t] * x[b, c, i + t]
#
# computed as a direct GEMM per (batch, tile) using sm80+ mma.sync.m16n8k8
# with TF32 operands and FP32 accumulation.
#
# Key optimizations over a naive tiling:
#   * the input patch is staged in shared memory with cp.async (16-byte
#     copies with hardware zero-fill for the length tail), so staging issues
#     ~8 fire-and-forget instructions per thread instead of ~33 dependent
#     scalar load/convert/store chains;
#   * channels are processed in two 32-channel chunks with a 2-stage
#     cp.async pipeline, overlapping the next chunk's global loads with the
#     current chunk's tensor-core math;
#   * weights are repacked once per call (~3us) into per-tap contiguous,
#     TF32-rounded (Cout, Cin) slices, zero-padded so channel / out-channel
#     tails need no special casing in the inner loop;
#   * shared row stride 136 (% 32 == 8) keeps the B-fragment scalar loads
#     bank-conflict free;
#   * bias is fused into the epilogue and output tails use predicated float2
#     stores, so global traffic is exactly: input read once + fp32 output
#     written once.
# ---------------------------------------------------------------------------

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

#define CONV_BN 128     // length tile per block
#define CONV_BO 128     // out-channel tile per block
#define CONV_BC 32      // in-channels per pipelined chunk
#define CONV_STRIDE 136 // shared row stride; % 32 == 8 -> conflict-free

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

__device__ __forceinline__ void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                                         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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

__device__ __forceinline__ void cp_async16(unsigned smem, const void* gmem, int src_bytes) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;\n"
                 :: "r"(smem), "l"(gmem), "r"(src_bytes));
}
__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n" ::);
}
template <int N>
__device__ __forceinline__ void cp_async_wait() {
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}

// w (Cout, Cin, K) fp32 -> wrep (K, CoutPad, CinPad) tf32, zero padded
__global__ void repack_w_kernel(const float* __restrict__ w, float* __restrict__ wrep,
                                int Cout, int Cin, int K, int CoutPad, int CinPad) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = Cout * Cin * K;
    if (idx < total) {
        const int o = idx / (Cin * K);
        const int rem = idx - o * Cin * K;
        const int c = rem / K;
        const int t = rem - c * K;
        wrep[((long long)t * CoutPad + o) * CinPad + c] =
            __uint_as_float(f32_to_tf32(w[idx]));
    }
}

__global__ __launch_bounds__(256, 2) void conv1d_tf32_kernel(
        const float* __restrict__ x, const float* __restrict__ wrep,
        const float* __restrict__ bias, float* __restrict__ out,
        int Cin, int L, int Cout, int K, int Lout, int CoutPad, int CinPad) {
    __shared__ __align__(16) float sh[2][CONV_BC][CONV_STRIDE];

    const int b  = blockIdx.z;
    const int i0 = blockIdx.x * CONV_BN;   // tile start along output length
    const int o0 = blockIdx.y * CONV_BO;   // tile start in out-channels

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int tig  = lane & 3;             // thread id in group of 4
    const int gid  = lane >> 2;            // group id 0..7
    const int wo   = (warp & 3) * 32;      // warp out-channel offset
    const int wn   = (warp >> 2) * 64;     // warp length offset

    const float* xb = x + (long long)b * Cin * L;

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

    const int nstage = CONV_BN + K - 1;      // staged positions per row
    const int nchunk = (nstage + 3) >> 2;    // 16-byte chunks per row
    const int cpjobs = CONV_BC * nchunk;

    // Stage one 32-channel chunk of the input patch into shared buffer `buf`
    // using 16B cp.async with zero-fill for out-of-range positions/channels.
    auto stage = [&](int cc, int buf) {
        const int rows = min(CONV_BC, Cin - cc);
        const unsigned base =
            static_cast<unsigned>(__cvta_generic_to_shared(&sh[buf][0][0]));
        for (int p = tid; p < cpjobs; p += 256) {
            const int c = p / nchunk;
            const int k = p - c * nchunk;
            const int g0 = i0 + (k << 2);
            int bytes = 0;
            const float* src = xb;
            if (c < rows) {
                const int avail = L - g0;
                bytes = (avail >= 4) ? 16 : ((avail > 0) ? (avail << 2) : 0);
                if (bytes > 0) src = xb + (long long)(cc + c) * L + g0;
            }
            const unsigned dst = base + (unsigned)(c * CONV_STRIDE + (k << 2)) * 4u;
            cp_async16(dst, src, bytes);
        }
    };

    const int ncchunk = (Cin + CONV_BC - 1) / CONV_BC;

    stage(0, 0);
    cp_async_commit();

    for (int j = 0; j < ncchunk; ++j) {
        const int cc  = j * CONV_BC;
        const int buf = j & 1;
        if (j + 1 < ncchunk) {
            stage(cc + CONV_BC, buf ^ 1);
            cp_async_commit();
            cp_async_wait<1>();
        } else {
            cp_async_wait<0>();
        }
        __syncthreads();

        const float* wc = wrep + (long long)o0 * CinPad + cc;
        for (int t = 0; t < K; ++t) {
            const float* wt = wc + (long long)t * CoutPad * CinPad;
            #pragma unroll
            for (int c0 = 0; c0 < CONV_BC; c0 += 8) {
                uint32_t a[2][4];
                #pragma unroll
                for (int f = 0; f < 2; ++f) {
                    const float* ap = wt + (long long)(wo + f * 16) * CinPad + c0;
                    a[f][0] = __float_as_uint(ap[(long long)gid * CinPad + tig]);
                    a[f][1] = __float_as_uint(ap[(long long)(gid + 8) * CinPad + tig]);
                    a[f][2] = __float_as_uint(ap[(long long)gid * CinPad + tig + 4]);
                    a[f][3] = __float_as_uint(ap[(long long)(gid + 8) * CinPad + tig + 4]);
                }
                const int nb = wn + t + gid;
                #pragma unroll
                for (int f = 0; f < 8; ++f) {
                    const int n = nb + f * 8;
                    const uint32_t b0 = f32_to_tf32(sh[buf][c0 + tig][n]);
                    const uint32_t b1 = f32_to_tf32(sh[buf][c0 + tig + 4][n]);
                    mma_tf32(acc[0][f][0], acc[0][f][1], acc[0][f][2], acc[0][f][3],
                             a[0][0], a[0][1], a[0][2], a[0][3], b0, b1);
                    mma_tf32(acc[1][f][0], acc[1][f][1], acc[1][f][2], acc[1][f][3],
                             a[1][0], a[1][1], a[1][2], a[1][3], b0, b1);
                }
            }
        }
        __syncthreads();
    }

    // Epilogue: fuse bias, predicated float2 stores for length/o tails.
    float* outb = out + (long long)b * Cout * Lout;
    const bool vec_ok = (Lout & 1) == 0;
    #pragma unroll
    for (int f = 0; f < 2; ++f) {
        const int oA = o0 + wo + f * 16 + gid;
        const bool v0 = (oA < Cout);
        const bool v1 = (oA + 8 < Cout);
        float bs0 = 0.f, bs1 = 0.f;
        if (bias != nullptr) {
            if (v0) bs0 = bias[oA];
            if (v1) bs1 = bias[oA + 8];
        }
        #pragma unroll
        for (int g = 0; g < 8; ++g) {
            const int iA = i0 + wn + g * 8 + tig * 2;
            if (vec_ok && iA + 1 < Lout) {
                if (v0) *reinterpret_cast<float2*>(outb + (long long)oA * Lout + iA) =
                            make_float2(acc[f][g][0] + bs0, acc[f][g][1] + bs0);
                if (v1) *reinterpret_cast<float2*>(outb + (long long)(oA + 8) * Lout + iA) =
                            make_float2(acc[f][g][2] + bs1, acc[f][g][3] + bs1);
            } else {
                if (v0 && iA < Lout) outb[(long long)oA * Lout + iA] = acc[f][g][0] + bs0;
                if (v0 && iA + 1 < Lout) outb[(long long)oA * Lout + iA + 1] = acc[f][g][1] + bs0;
                if (v1 && iA < Lout) outb[(long long)(oA + 8) * Lout + iA] = acc[f][g][2] + bs1;
                if (v1 && iA + 1 < Lout) outb[(long long)(oA + 8) * Lout + iA + 1] = acc[f][g][3] + bs1;
            }
        }
    }
}

torch::Tensor conv1d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "x and w must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 supported");
    TORCH_CHECK(x.dim() == 3 && w.dim() == 3, "x (B,Cin,L), w (Cout,Cin,K)");

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

    const int64_t B    = xc.size(0);
    const int64_t Cin  = xc.size(1);
    const int64_t L    = xc.size(2);
    const int64_t Cout = wc.size(0);
    const int64_t K    = wc.size(2);
    const int64_t Lout = L - K + 1;

    TORCH_CHECK(Lout > 0, "kernel larger than length");
    TORCH_CHECK(K <= 9, "kernel_size too large for fast path");
    TORCH_CHECK((L % 4) == 0, "length must be a multiple of 4 for fast path");
    TORCH_CHECK((reinterpret_cast<uintptr_t>(xc.data_ptr<float>()) & 15) == 0,
                "input must be 16-byte aligned");

    const int64_t CoutPad = (Cout + CONV_BO - 1) / CONV_BO * CONV_BO;
    const int64_t CinPad  = (Cin + CONV_BC - 1) / CONV_BC * CONV_BC;

    auto out  = at::empty({B, Cout, Lout}, xc.options());
    auto wrep = at::zeros({K, CoutPad, CinPad}, wc.options());

    const float* bias_ptr = nullptr;
    if (bias.defined() && bias.numel() > 0) {
        bias_ptr = bias.contiguous().data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int wtotal = (int)(Cout * Cin * K);
    repack_w_kernel<<<(wtotal + 255) / 256, 256, 0, stream>>>(
        wc.data_ptr<float>(), wrep.data_ptr<float>(),
        (int)Cout, (int)Cin, (int)K, (int)CoutPad, (int)CinPad);

    dim3 grid((unsigned)((Lout + CONV_BN - 1) / CONV_BN),
              (unsigned)(CoutPad / CONV_BO),
              (unsigned)B);
    conv1d_tf32_kernel<<<grid, 256, 0, stream>>>(
        xc.data_ptr<float>(), wrep.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
        (int)Cin, (int)L, (int)Cout, (int)K, (int)Lout, (int)CoutPad, (int)CinPad);
    return out;
}
"""

cpp_src = "torch::Tensor conv1d_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias);"

_ext = load_inline(
    name="fast_conv1d_tf32_cpasync_ext",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["conv1d_forward_cuda"],
    with_cuda=True,
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized 1D convolution.

    Fast path: custom fused TF32 tensor-core (mma.sync m16n8k8, fp32
    accumulate) CUDA kernel with cp.async pipelined staging, for the standard
    configuration (stride=1, padding=0, dilation=1, groups=1) on sm80+ GPUs
    with fp32 inputs, kernel_size <= 9 and length % 4 == 0.  Bias is fused in
    the kernel epilogue.  Anything else falls back to the native nn.Conv1d.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel.
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    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__()
        # Keep identical parameters/state_dict layout as the reference model.
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self._shape_ok = (
            stride == 1 and padding == 0 and dilation == 1 and groups == 1
            and kernel_size <= 9
        )
        self._sm_ok = False
        if torch.cuda.is_available():
            try:
                self._sm_ok = torch.cuda.get_device_capability(0)[0] >= 8
            except Exception:
                self._sm_ok = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        K = self.conv1d.kernel_size[0]
        if (self._shape_ok and self._sm_ok and x.is_cuda
                and x.dtype == torch.float32 and x.dim() == 3
                and x.size(1) == self.conv1d.in_channels
                and x.size(2) >= K and (x.size(2) % 4) == 0
                and (x.storage_offset() % 4) == 0):
            w = self.conv1d.weight
            b = self.conv1d.bias
            if b is None:
                b = x.new_empty(0)
            return _ext.conv1d_forward_cuda(x, w, b)
        return self.conv1d(x)
```