```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: fuse the WHOLE pipeline into a single convolution kernel.
#
# Profiling in earlier turns showed the pipeline was dominated by memory
# traffic: cuDNN conv writes 2.1 GB, then the fused subtract+Mish pass reads
# 2.1 GB and writes 2.1 GB again. Since the activation is elementwise it can
# be applied in the convolution epilogue, eliminating 4.2 GB of traffic.
#
# Fast path (CUDA, fp32, 3x3 stride-1 pad-0 conv with Cin=8, Cout=64, sm80+):
#   A custom implicit-GEMM conv that reads the NCHW input directly (fully
#   coalesced: consecutive lanes cover consecutive pixels per channel plane),
#   uses TF32 tensor-core mma (m16n8k8) with inputs rounded via
#   cvt.rna.tf32.f32 to match cuDNN TF32 numerics, stages the pre-repacked
#   [72][64] weight matrix in padded shared memory (bank-conflict-free), and
#   fuses bias + (subtract_value_1 + subtract_value_2) + Mish into the
#   epilogue, writing the output directly as a channels_last tensor.
#
# A one-time runtime self-check validates the custom kernel against the
# cuDNN path on the first forward; on any mismatch or unsupported config it
# permanently falls back to the proven path (channels_last cuDNN conv +
# layout-agnostic fused subtract+Mish kernel, ~2.48x speedup).
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])));
}

// ---------------------------------------------------------------------------
// 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.
// A fragments are gathered directly from the NCHW input (no im2col, no
// layout conversion); B is pre-repacked to [72][64] and staged in shared
// memory once per block. Output written as NHWC (channels_last) buffer.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256, 2)
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 long 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)

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

    const long HW = (long)H * (long)W;
    const long m0 = (long)blockIdx.x * 128 + wm * 32 + g;

    // NCHW base offsets (c=0, kh=kw=0) for this thread's 4 pixels.
    long rbase[2][2];
    bool pv[2][2];
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            long m = m0 + mt * 16 + r * 8;
            pv[mt][r] = (m < M);
            long mm = pv[mt][r] ? m : 0;
            int wo = (int)(mm % Wo);
            long t = mm / Wo;
            int ho = (int)(t % Ho);
            long img = t / Ho;
            rbase[mt][r] = (img * 8 * (long)H + ho) * (long)W + wo;
        }
    }

    // The 9 (kh,kw) tap offsets in an NCHW plane.
    int koff[9];
    #pragma unroll
    for (int s = 0; s < 9; ++s) {
        koff[s] = (s / 3) * W + (s % 3);
    }

    const long coff0 = (long)t4 * HW;        // channel t4
    const long coff1 = coff0 + 4 * HW;       // channel t4 + 4

    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).
    #pragma unroll
    for (int s = 0; s < 9; ++s) {
        const int ks = s * 8;
        float a[2][4];
        #pragma unroll
        for (int mt = 0; mt < 2; ++mt) {
            const float* p0 = in + rbase[mt][0] + koff[s];
            const float* p1 = in + rbase[mt][1] + koff[s];
            a[mt][0] = cvt_tf32(pv[mt][0] ? __ldg(p0 + coff0) : 0.0f);
            a[mt][1] = cvt_tf32(pv[mt][1] ? __ldg(p1 + coff0) : 0.0f);
            a[mt][2] = cvt_tf32(pv[mt][0] ? __ldg(p0 + coff1) : 0.0f);
            a[mt][3] = cvt_tf32(pv[mt][1] ? __ldg(p1 + coff1) : 0.0f);
        }
        #pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
            const int n = wn * 32 + nt * 8 + g;
            float b[2];
            b[0] = cvt_tf32(bs[ks + t4][n]);
            b[1] = cvt_tf32(bs[ks + t4 + 4][n]);
            mma1688_tf32(acc[0][nt], a[0], b);
            mma1688_tf32(acc[1][nt], a[1], b);
        }
    }

    // Epilogue: bias + subtract + Mish, write FP32 into the NHWC buffer.
    #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) {
                if (pv[mt][r]) {
                    const long m = m0 + mt * 16 + r * 8;
                    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);
                    *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(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");

    // 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, M,
        (float)sub);
    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 fused_sub_mish_cuda(torch::Tensor x, double sub);"
)

fused_pipe = load_inline(
    name="fused_conv_pipe_v5",
    cpp_sources=fused_pipe_cpp_source,
    cuda_sources=fused_pipe_cuda_source,
    functions=["conv3x3_c8_o64_mish_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. 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.48x 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 torch.cuda.get_device_capability(c.weight.device)[0] < 8:
            return False
        return True

    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 (x.dim() == 4 and x.dtype == torch.float32
                    and self._fast_supported()
                    and x.shape[2] >= 3 and x.shape[3] >= 3):
                try:
                    with torch.no_grad():
                        # Repack [64, 8, 3, 3] -> [(kh*3+kw)*8 + c][o] = [72][64].
                        wmat = (self.conv.weight.detach()
                                .permute(2, 3, 1, 0).reshape(72, 64)
                                .contiguous())
                        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 x.dim() == 4 and x.dtype == torch.float32
                and x.shape[2] >= 3 and x.shape[3] >= 3):
            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)
```