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

# Design: the whole pipeline (conv3x3 + subtract + Mish) is fused into one
# implicit-GEMM TF32 tensor-core kernel.
#
# Improvements over the previous (7.1x) revision of this kernel:
#   * 32-bit addressing everywhere (guarded host-side), cutting register
#     pressure and 64-bit IMAD work in the hot K loop.
#   * Occupancy raised from 2 to 3 resident blocks/SM (768 threads/SM) for
#     better memory-latency hiding.
#   * Weights are pre-rounded to TF32 once at setup, removing 72 cvt ops per
#     thread from the inner loop.
#   * Out-of-range pixel rows are address-clamped instead of predicate-masked,
#     removing per-load selects from the K loop (their outputs are simply not
#     stored).
#   * Software double-buffering of the A fragments: the next tap's loads are
#     issued before the current tap's MMAs; the first tile's loads overlap
#     the shared-memory weight staging.
#   * Output stores use the streaming cache hint (__stcs) so the 2.1 GB of
#     write-once output does not evict the input/weights from L2.
#
# A one-time runtime self-check validates the fast kernel against the cuDNN
# reference path; on any mismatch or unsupported configuration it permanently
# falls back to the proven channels_last cuDNN conv + fused subtract/Mish
# kernel (~2.5x).
fused_pipe_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define DEVINL __device__ __forceinline__

// mish(x) = x * v/(v+2) with v = e^x (e^x + 2); algebraically identical to
// x*tanh(softplus(x)); threshold matches PyTorch's softplus (x > 20 -> x).
DEVINL float mish_fast(float x) {
    float u = __expf(x);
    float v = u * (u + 2.0f);
    float y = x * __fdividef(v, v + 2.0f);
    return (x > 20.0f) ? x : y;
}

DEVINL float mish_stable(float x) {
    float sp = (x > 20.0f) ? x : log1pf(expf(x));
    return x * tanhf(sp);
}

// Round fp32 -> tf32 (round-to-nearest-even), result kept in a b32 register.
DEVINL float cvt_tf32(float x) {
    uint32_t r;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(r) : "f"(x));
    return __uint_as_float(r);
}

DEVINL void mma1688_tf32(float (&d)[4], const float (&a)[4],
                         const float (&b)[2]) {
    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"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
        : "r"(__float_as_uint(a[0])), "r"(__float_as_uint(a[1])),
          "r"(__float_as_uint(a[2])), "r"(__float_as_uint(a[3])),
          "r"(__float_as_uint(b[0])), "r"(__float_as_uint(b[1])));
}

// Gather one k8 A-fragment tile (this thread's 2 pixel rows x channels
// {t4, t4+4}) straight from the NCHW input and round to tf32.
DEVINL void load_a_tile(const float* __restrict__ in, const int (&rb)[2][2],
                        const int koff, const int coff0, const int coff1,
                        float (&a)[2][4]) {
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        const float* p0 = in + (rb[mt][0] + koff);
        const float* p1 = in + (rb[mt][1] + koff);
        a[mt][0] = cvt_tf32(__ldg(p0 + coff0));
        a[mt][1] = cvt_tf32(__ldg(p1 + coff0));
        a[mt][2] = cvt_tf32(__ldg(p0 + coff1));
        a[mt][3] = cvt_tf32(__ldg(p1 + coff1));
    }
}

// ---------------------------------------------------------------------------
// Fused conv3x3 (stride 1, pad 0, Cin=8, Cout=64) + bias + subtract + Mish.
// GEMM view: M = N*Ho*Wo pixels, N = 64 channels, K = 72 taps.
// Block tile: 128 x 64, 8 warps (4 in M x 2 in N), warp tile 32x32.
// All indexing is 32-bit (guarded by the host wrapper). B is pre-repacked to
// [72][64] and pre-rounded to tf32. Output written as NHWC (channels_last)
// with streaming stores.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256, 3)
conv3x3_c8_o64_mish_kernel(const float* __restrict__ in,
                           const float* __restrict__ wmat,
                           const float* __restrict__ bias,
                           float* __restrict__ out,
                           const int H, const int W,
                           const int Ho, const int Wo,
                           const int M, const float sub) {
    __shared__ float bs[72][72];  // padded stride 72 -> bank-conflict-free

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;   // groupID
    const int t4   = lane & 3;    // threadID_in_group
    const int wm   = warp >> 1;   // 0..3 (M direction)
    const int wn   = warp & 1;    // 0..1 (N direction)

    const int HW = H * W;

    // NCHW base offsets (c=0, kh=kw=0) for this thread's 4 pixels. Rows
    // beyond M are clamped to pixel 0 (in-bounds; their results are masked
    // in the epilogue), so the K loop needs no per-load predicates.
    const int m0 = blockIdx.x * 128 + wm * 32 + g;
    int rbase[2][2];
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int m   = m0 + mt * 16 + r * 8;
            const int mm  = (m < M) ? m : 0;
            const int wo  = mm % Wo;
            const int t   = mm / Wo;
            const int ho  = t % Ho;
            const int img = t / Ho;
            rbase[mt][r] = (img * 8 * H + ho) * W + wo;
        }
    }
    const int coff0 = t4 * HW;        // channel t4
    const int coff1 = coff0 + 4 * HW; // channel t4 + 4

    // Prefetch the first A tile; its latency overlaps the B staging below.
    float a[2][2][4];  // [buffer][mt][reg]
    load_a_tile(in, rbase, 0, coff0, coff1, a[0]);

    // Stage the tf32 pre-rounded B matrix (72x64) into shared memory.
    #pragma unroll
    for (int i = 0; i < 18; ++i) {
        const int idx = tid + i * 256;
        bs[idx >> 6][idx & 63] = __ldg(wmat + idx);
    }
    __syncthreads();

    float acc[2][4][4];
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt)
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt)
            #pragma unroll
            for (int e = 0; e < 4; ++e) acc[mt][nt][e] = 0.0f;

    // K loop: 9 steps of k=8 (one (kh,kw) tap = 8 channels per step),
    // software-pipelined: issue step s+1's global loads before step s's MMAs.
    #pragma unroll
    for (int s = 0; s < 9; ++s) {
        const int cur = s & 1;
        if (s < 8) {
            const int sn = s + 1;
            load_a_tile(in, rbase, (sn / 3) * W + (sn % 3),
                        coff0, coff1, a[cur ^ 1]);
        }
        const int ks = s * 8;
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const int n = wn * 32 + nt * 8 + g;
            float b[2];
            b[0] = bs[ks + t4][n];
            b[1] = bs[ks + t4 + 4][n];
            mma1688_tf32(acc[0][nt], a[cur][0], b);
            mma1688_tf32(acc[1][nt], a[cur][1], b);
        }
    }

    // Epilogue: bias + subtract + Mish, streaming FP32 stores into the NHWC
    // buffer (write-once data kept out of L2 as much as possible).
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const int ncol = wn * 32 + nt * 8 + t4 * 2;
            const float b0 = __ldg(bias + ncol);
            const float b1 = __ldg(bias + ncol + 1);
            #pragma unroll
            for (int r = 0; r < 2; ++r) {
                const int m = m0 + mt * 16 + r * 8;
                if (m < M) {
                    float2 v;
                    v.x = mish_fast(acc[mt][nt][r * 2 + 0] + b0 - sub);
                    v.y = mish_fast(acc[mt][nt][r * 2 + 1] + b1 - sub);
                    __stcs(reinterpret_cast<float2*>(out + m * 64 + ncol), v);
                }
            }
        }
    }
}

torch::Tensor conv3x3_c8_o64_mish_cuda(torch::Tensor x, torch::Tensor wmat,
                                       torch::Tensor bias, double sub) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4 &&
                x.scalar_type() == torch::kFloat32,
                "x must be a CUDA 4D float32 tensor");
    TORCH_CHECK(x.size(1) == 8, "fast conv requires Cin == 8");
    TORCH_CHECK(x.size(2) >= 3 && x.size(3) >= 3, "spatial dims too small");
    TORCH_CHECK(wmat.numel() == 72 * 64 && bias.numel() == 64,
                "bad weight/bias shapes");
    torch::Tensor xc = x.is_contiguous() ? x : x.contiguous();
    torch::Tensor wm = wmat.is_contiguous() ? wmat : wmat.contiguous();
    torch::Tensor bc = bias.is_contiguous() ? bias : bias.contiguous();

    const long N  = xc.size(0);
    const long H  = xc.size(2);
    const long W  = xc.size(3);
    const long Ho = H - 2;
    const long Wo = W - 2;
    const long M  = N * Ho * Wo;
    TORCH_CHECK(M > 0, "spatial dims too small");
    // 32-bit indexing guards.
    TORCH_CHECK(xc.numel() < (1L << 31), "input too large for fast conv");
    TORCH_CHECK(M * 64 + 64 < (1L << 31), "output too large for fast conv");

    // Logical NCHW tensor whose dense storage is NHWC.
    auto out = torch::empty({N, 64, Ho, Wo}, xc.options(),
                            torch::MemoryFormat::ChannelsLast);
    const long blocks = (M + 127) / 128;
    conv3x3_c8_o64_mish_kernel<<<(unsigned int)blocks, 256, 0,
                                 at::cuda::getCurrentCUDAStream()>>>(
        xc.data_ptr<float>(), wm.data_ptr<float>(), bc.data_ptr<float>(),
        out.data_ptr<float>(), (int)H, (int)W, (int)Ho, (int)Wo, (int)M,
        (float)sub);
    return out;
}

// --------------------- one-time tf32 weight pre-rounding --------------------
__global__ void round_tf32_kernel(const float* __restrict__ in,
                                  float* __restrict__ out, const long n) {
    const long i = (long)blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = cvt_tf32(in[i]);
    }
}

torch::Tensor round_tf32_cuda(torch::Tensor w) {
    TORCH_CHECK(w.is_cuda() && w.scalar_type() == torch::kFloat32,
                "w must be a CUDA float32 tensor");
    torch::Tensor wc = w.contiguous();
    auto out = torch::empty_like(wc);
    const long n = wc.numel();
    if (n > 0) {
        const int block = 256;
        const long blocks = (n + block - 1) / block;
        round_tf32_kernel<<<(unsigned int)blocks, block, 0,
                            at::cuda::getCurrentCUDAStream()>>>(
            wc.data_ptr<float>(), out.data_ptr<float>(), n);
    }
    return out;
}

// ------------------------- fused subtract + Mish ---------------------------
__global__ void sub_mish_vec4_kernel(const float4* __restrict__ in,
                                     float4* __restrict__ out,
                                     const float sub, const long n4) {
    const long idx = (long)blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx < n4) {
        float4 v = in[idx];
        v.x = mish_stable(v.x - sub);
        v.y = mish_stable(v.y - sub);
        v.z = mish_stable(v.z - sub);
        v.w = mish_stable(v.w - sub);
        out[idx] = v;
    }
}

__global__ void sub_mish_scalar_kernel(const float* __restrict__ in,
                                       float* __restrict__ out,
                                       const float sub,
                                       const long start, const long n) {
    const long idx = start + (long)blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx < n) {
        out[idx] = mish_stable(in[idx] - sub);
    }
}

torch::Tensor fused_sub_mish_cuda(torch::Tensor x, double sub) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");

    // Any dense layout works for an elementwise op; only genuinely strided
    // tensors need a copy (never happens for conv output).
    torch::Tensor xc = x;
    if (!x.is_contiguous(at::MemoryFormat::Contiguous) &&
        !x.is_contiguous(at::MemoryFormat::ChannelsLast)) {
        xc = x.contiguous();
    }

    auto out = torch::empty_like(xc);  // preserves xc's memory format
    const long n = xc.numel();
    if (n == 0) {
        return out;
    }

    const float* in_ptr = xc.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();
    const float sub_f = static_cast<float>(sub);
    const int block = 256;

    const bool aligned =
        ((reinterpret_cast<uintptr_t>(in_ptr) |
          reinterpret_cast<uintptr_t>(out_ptr)) & 0xF) == 0;

    const long n4 = aligned ? (n / 4) : 0;
    if (n4 > 0) {
        const long blocks = (n4 + block - 1) / block;
        sub_mish_vec4_kernel<<<(unsigned int)blocks, block>>>(
            reinterpret_cast<const float4*>(in_ptr),
            reinterpret_cast<float4*>(out_ptr), sub_f, n4);
    }

    const long rem = n - n4 * 4;
    if (rem > 0) {
        const long blocks = (rem + block - 1) / block;
        sub_mish_scalar_kernel<<<(unsigned int)blocks, block>>>(
            in_ptr, out_ptr, sub_f, n4 * 4, n);
    }

    return out;
}
"""

fused_pipe_cpp_source = (
    "torch::Tensor conv3x3_c8_o64_mish_cuda(torch::Tensor x, "
    "torch::Tensor wmat, torch::Tensor bias, double sub);\n"
    "torch::Tensor round_tf32_cuda(torch::Tensor w);\n"
    "torch::Tensor fused_sub_mish_cuda(torch::Tensor x, double sub);"
)

fused_pipe = load_inline(
    name="fused_conv_pipe_v6",
    cpp_sources=fused_pipe_cpp_source,
    cuda_sources=fused_pipe_cuda_source,
    functions=["conv3x3_c8_o64_mish_cuda", "round_tf32_cuda",
               "fused_sub_mish_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model.

    Fast path (CUDA, fp32, 3x3 stride-1 pad-0 conv with Cin=8, Cout=64):
      a single custom implicit-GEMM TF32 tensor-core conv kernel that reads
      the NCHW input directly and fuses bias + subtract + Mish into its
      epilogue, eliminating the separate 4.2 GB activation pass. Uses 32-bit
      addressing, software-pipelined A-fragment loads, tf32 pre-rounded
      weights, 3 resident blocks/SM, and streaming output stores. A one-time
      self-check against the cuDNN path validates the kernel at runtime; on
      any mismatch (or unsupported config, or pre-sm80 GPU) it permanently
      falls back.

    Fallback path: channels_last cuDNN conv + fused subtract/Mish kernel
    (the proven ~2.5x configuration).
    """
    def __init__(self, in_channels, out_channels, kernel_size,
                 subtract_value_1, subtract_value_2):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        # Fold the two scalar subtractions into one (plain attribute, so it
        # stays out of the state_dict and weight loading is unaffected).
        self.sub_total = float(subtract_value_1) + float(subtract_value_2)
        self._cl_ready = False
        self._checked = False
        self._use_fast = False
        self._wmat = None
        self._bias = None
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True

    def _fast_supported(self):
        c = self.conv
        k = c.kernel_size if isinstance(c.kernel_size, tuple) else (c.kernel_size,) * 2
        s = c.stride if isinstance(c.stride, tuple) else (c.stride,) * 2
        p = c.padding if isinstance(c.padding, tuple) else (c.padding,) * 2
        d = c.dilation if isinstance(c.dilation, tuple) else (c.dilation,) * 2
        if not (k == (3, 3) and s == (1, 1) and p == (0, 0) and d == (1, 1)):
            return False
        if not (c.in_channels == 8 and c.out_channels == 64 and c.groups == 1):
            return False
        if c.bias is None:
            return False
        if not (c.weight.is_cuda and c.weight.dtype == torch.float32):
            return False
        if c.bias.dtype != torch.float32:
            return False
        if torch.cuda.get_device_capability(c.weight.device)[0] < 8:
            return False
        return True

    @staticmethod
    def _shape_ok(x):
        if not (x.dim() == 4 and x.dtype == torch.float32
                and x.shape[2] >= 3 and x.shape[3] >= 3):
            return False
        if x.numel() >= 2**31:
            return False
        m64 = x.shape[0] * (x.shape[2] - 2) * (x.shape[3] - 2) * 64
        return m64 + 64 < 2**31

    def forward(self, x):
        if not x.is_cuda:
            # Safe fallback for CPU tensors.
            return F.mish(self.conv(x) - self.sub_total)

        if not self._cl_ready:
            # Convert weights to channels_last once (done lazily here so it
            # happens after any state_dict loading, not at __init__ time).
            with torch.no_grad():
                self.conv.weight.data = self.conv.weight.data.to(
                    memory_format=torch.channels_last)
            self._cl_ready = True

        if not self._checked:
            self._checked = True
            use_fast = False
            if self._shape_ok(x) and self._fast_supported():
                try:
                    with torch.no_grad():
                        # Repack [64, 8, 3, 3] -> [(kh*3+kw)*8 + c][o] = [72][64]
                        # and pre-round to tf32 (the kernel's MMA gets exact
                        # tf32 operands with no per-step conversion).
                        wmat = (self.conv.weight.detach()
                                .permute(2, 3, 1, 0).reshape(72, 64)
                                .contiguous())
                        wmat = fused_pipe.round_tf32_cuda(wmat)
                        bias = self.conv.bias.detach().contiguous()
                        out_fast = fused_pipe.conv3x3_c8_o64_mish_cuda(
                            x, wmat, bias, self.sub_total)
                        ref = fused_pipe.fused_sub_mish_cuda(
                            self.conv(x.to(memory_format=torch.channels_last)),
                            self.sub_total)
                        rel = ((out_fast - ref).abs()
                               / ref.abs().clamp_min(1.0)).max().item()
                        if rel < 2e-2 and bool(torch.isfinite(out_fast).all().item()):
                            self._wmat = wmat
                            self._bias = bias
                            use_fast = True
                except Exception:
                    use_fast = False
            self._use_fast = use_fast

        if self._use_fast and self._shape_ok(x):
            return fused_pipe.conv3x3_c8_o64_mish_cuda(
                x, self._wmat, self._bias, self.sub_total)

        # Proven fallback: NHWC cuDNN conv + fused subtract/Mish on the
        # dense NHWC buffer (no extra layout copies).
        x = x.to(memory_format=torch.channels_last)
        x = self.conv(x)
        return fused_pipe.fused_sub_mish_cuda(x, self.sub_total)
```