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

# Fused Conv2d + HardSwish + ReLU.
#
# Fast path (v5 tensor-core kernel): implicit-GEMM using 16x16x8 TF32 wmma
# fragments with the "3xTF32" trick (split every fp32 into hi+lo TF32 and issue
# hi*hi, hi*lo, lo*hi MMAs). This gives tensor-core throughput with ~fp32
# accuracy (error ~2^-22, far below plain TF32's ~2^-11), so tight correctness
# checks still pass. Shared memory is kept under the static 48KB limit by
# aliasing the epilogue bounce buffer (Cs) onto the A/B staging buffers, which
# are dead once the main loop finishes.
#
# Fallbacks: proven direct FP32 fused SIMT kernel (any shape/arch), then ATen.
# The TC path is enabled only after a one-time runtime self-check against a
# strict-fp32 reference; on any mismatch it silently stays on the FP32 kernel.
_fused_cpp_src = (
    "torch::Tensor fused_conv_act(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride, int64_t pad, bool use_tc);"
)

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <mma.h>

#define OC_TILE 8
#define BM 128
#define BN 64
#define BK 24
#define LDB 32
#define LDC 65
#define NTHREADS 256

__device__ __forceinline__ float relu_hardswish(float x) {
    float t = fminf(fmaxf(x + 3.0f, 0.0f), 6.0f);
    float hs = x * t / 6.0f;
    return fmaxf(hs, 0.0f);
}

// Split fp32 into (hi, lo) tf32 bit patterns with x ~= hi + lo.
__device__ __forceinline__ void split_tf32(float x, float& hi, float& lo) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    unsigned int h, l;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(h) : "f"(x));
    float xh = __uint_as_float(h);
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(l) : "f"(x - xh));
    hi = xh;
    lo = __uint_as_float(l);
#else
    hi = x;
    lo = 0.0f;
#endif
}

// ---------------- Tensor-core (3xTF32) fused conv + HardSwish + ReLU -----------
template <int C, int KH, int KW>
__global__ void __launch_bounds__(NTHREADS) conv_act_tc_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int H, int W, int OC,
    int OH, int OW, int stride, int pad, int octiles)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    constexpr int K = C * KH * KW;
    constexpr int KK = KH * KW;
    // Aliased shared memory: A/B staging live during the main loop; the C
    // bounce buffer (BM*LDC floats <= SM_A+SM_B) reuses the same storage
    // afterwards. s_bias sits past the staging region so it survives.
    constexpr int SM_A = 2 * BK * BM;       // 6144 floats
    constexpr int SM_B = 2 * BN * LDB;      // 4096 floats
    constexpr int SM_MAIN = SM_A + SM_B;    // 10240 floats (>= BM*LDC = 8320)
    __shared__ __align__(32) float smem[SM_MAIN + BN];
    float* As_hi  = smem;
    float* As_lo  = smem + BK * BM;
    float* Bs_hi  = smem + SM_A;
    float* Bs_lo  = smem + SM_A + BN * LDB;
    float* s_bias = smem + SM_MAIN;
    float* Cs     = smem;  // alias; A/B are dead after the main loop

    const int tid  = threadIdx.x;             // 256 threads = 8 warps
    const int warp = tid >> 5;
    const int wm = warp & 3;                  // 4 warps along M (rows)
    const int wn = warp >> 2;                 // 2 warps along N (cols)

    const int n   = blockIdx.z / octiles;
    const int oct = blockIdx.z - n * octiles;
    const int oc0 = oct * BN;
    const int oh  = blockIdx.y;
    const int ow0 = blockIdx.x * BM;
    const int ow  = ow0 + (tid & (BM - 1));
    const int ih0 = oh * stride - pad;
    const int iw0 = ow * stride - pad;

    namespace wmma = nvcuda::wmma;
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_hh[2][2];
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_md[2][2];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::fill_fragment(acc_hh[i][j], 0.0f);
            wmma::fill_fragment(acc_md[i][j], 0.0f);
        }

    if (tid < BN) {
        const int oc = oc0 + tid;
        s_bias[tid] = (oc < OC) ? bias[oc] : 0.0f;
    }

    const float* in_n = input + (long)n * (C * H * W);
    const int m  = tid & (BM - 1);            // column (output-position) index
    const int khalf = tid >> 7;               // 0/1 -> two k-rows per pass
    const int bn = tid & (BN - 1);            // B row (output channel)
    const int kq = tid >> 6;                  // 0..3 -> four k per pass

    #pragma unroll
    for (int k0 = 0; k0 < K; k0 += BK) {
        // ---- Stage A chunk: As[k][m], col-major (ld = BM) ----
        #pragma unroll
        for (int p = 0; p < BK / 2; ++p) {
            const int k  = k0 + p * 2 + khalf;
            const int c  = k / KK;
            const int r  = k - c * KK;
            const int kh = r / KW;
            const int kw = r - kh * KW;
            const int ih = ih0 + kh;
            const int iw = iw0 + kw;
            float v = 0.0f;
            if (ih >= 0 && ih < H && iw >= 0 && iw < W && ow < OW)
                v = in_n[(c * H + ih) * W + iw];
            float vh, vl;
            split_tf32(v, vh, vl);
            const int off = (p * 2 + khalf) * BM + m;
            As_hi[off] = vh;
            As_lo[off] = vl;
        }
        // ---- Stage B chunk: Bs[n][k], col-major (ld = LDB) ----
        #pragma unroll
        for (int p = 0; p < BK / 4; ++p) {
            const int k  = k0 + p * 4 + kq;
            const int oc = oc0 + bn;
            float v = 0.0f;
            if (oc < OC) v = weight[(long)oc * K + k];
            float vh, vl;
            split_tf32(v, vh, vl);
            const int off = bn * LDB + p * 4 + kq;
            Bs_hi[off] = vh;
            Bs_lo[off] = vl;
        }
        __syncthreads();

        // ---- Compute: BK/8 ksteps ----
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::col_major> ah[2], al[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> bh[2], bl[2];
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(ah[i], As_hi + kk * BM + wm * 32 + i * 16, BM);
                wmma::load_matrix_sync(al[i], As_lo + kk * BM + wm * 32 + i * 16, BM);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(bh[j], Bs_hi + (wn * 32 + j * 16) * LDB + kk, LDB);
                wmma::load_matrix_sync(bl[j], Bs_lo + (wn * 32 + j * 16) * LDB + kk, LDB);
            }
            #pragma unroll
            for (int i = 0; i < 2; ++i)
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc_hh[i][j], ah[i], bh[j], acc_hh[i][j]);
                    wmma::mma_sync(acc_md[i][j], ah[i], bl[j], acc_md[i][j]);
                    wmma::mma_sync(acc_md[i][j], al[i], bh[j], acc_md[i][j]);
                }
        }
        __syncthreads();
    }

    // ---- Epilogue: merge hi/lo partial sums, bounce through smem ----
    // (A/B staging is dead; the final __syncthreads() above makes the alias safe.)
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int e = 0; e < acc_hh[i][j].num_elements; ++e)
                acc_hh[i][j].x[e] += acc_md[i][j].x[e];
            wmma::store_matrix_sync(
                Cs + (wm * 32 + i * 16) * LDC + (wn * 32 + j * 16),
                acc_hh[i][j], LDC, wmma::mem_row_major);
        }
    __syncthreads();

    // ---- Fused bias + HardSwish + ReLU, coalesced writes ----
    float* out_p = output + (long)n * OC * OH * OW + (long)oh * OW + ow0;
    #pragma unroll
    for (int i = 0; i < (BM * BN) / NTHREADS; ++i) {
        const int idx = tid + i * NTHREADS;
        const int j = idx / BM;               // output channel (local)
        const int mm = idx & (BM - 1);        // output position
        const int oc = oc0 + j;
        if (oc < OC && (ow0 + mm) < OW) {
            out_p[(long)oc * OH * OW + mm] =
                relu_hardswish(Cs[mm * LDC + j] + s_bias[j]);
        }
    }
#endif
}

// ---------------- Generic direct FP32 fallback (any shape/arch) -----------------
template <int C, int K>
__global__ void conv_hswish_relu_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int H, int W, int OC,
    int OH, int OW, int stride, int pad)
{
    constexpr int KSIZE = C * K * K;
    __shared__ float s_weight[OC_TILE * KSIZE];

    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
    const int nthreads = blockDim.x * blockDim.y;
    const int oc0 = blockIdx.y * OC_TILE;

    for (int i = tid; i < OC_TILE * KSIZE; i += nthreads) {
        int oc = oc0 + i / KSIZE;
        int w = i % KSIZE;
        s_weight[i] = (oc < OC) ? weight[oc * KSIZE + w] : 0.0f;
    }
    __syncthreads();

    const int tilesX = (OW + (int)blockDim.x - 1) / (int)blockDim.x;
    const int tileY = blockIdx.x / tilesX;
    const int tileX = blockIdx.x % tilesX;
    const int oh = tileY * (int)blockDim.y + (int)threadIdx.y;
    const int ow = tileX * (int)blockDim.x + (int)threadIdx.x;
    if (oh >= OH || ow >= OW) return;

    const int n = blockIdx.z;
    const float* in_n = input + (long)n * (C * H * W);
    const int ih0 = oh * stride - pad;
    const int iw0 = ow * stride - pad;

    float acc[OC_TILE];
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) acc[i] = 0.0f;

    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* in_c = in_n + c * (H * W);
        #pragma unroll
        for (int kh = 0; kh < K; ++kh) {
            const int ih = ih0 + kh;
            const bool hok = (ih >= 0) && (ih < H);
            #pragma unroll
            for (int kw = 0; kw < K; ++kw) {
                const int iw = iw0 + kw;
                float v = (hok && iw >= 0 && iw < W) ? in_c[ih * W + iw] : 0.0f;
                const int woff = (c * K + kh) * K + kw;
                #pragma unroll
                for (int i = 0; i < OC_TILE; ++i) {
                    acc[i] = fmaf(v, s_weight[i * KSIZE + woff], acc[i]);
                }
            }
        }
    }

    float* out_n = output + (long)n * OC * OH * OW + oh * OW + ow;
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) {
        const int oc = oc0 + i;
        if (oc < OC) {
            out_n[(long)oc * OH * OW] = relu_hardswish(acc[i] + bias[oc]);
        }
    }
}

__global__ void conv_hswish_relu_generic_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int N, int C, int H, int W, int OC,
    int KH, int KW, int OH, int OW, int stride, int pad)
{
    extern __shared__ float s_weight[];
    const int ksize = C * KH * KW;
    const int tid = threadIdx.y * blockDim.x + threadIdx.x;
    const int nthreads = blockDim.x * blockDim.y;
    const int oc0 = blockIdx.y * OC_TILE;

    for (int i = tid; i < OC_TILE * ksize; i += nthreads) {
        int oc = oc0 + i / ksize;
        int w = i % ksize;
        s_weight[i] = (oc < OC) ? weight[oc * ksize + w] : 0.0f;
    }
    __syncthreads();

    const int tilesX = (OW + (int)blockDim.x - 1) / (int)blockDim.x;
    const int tileY = blockIdx.x / tilesX;
    const int tileX = blockIdx.x % tilesX;
    const int oh = tileY * (int)blockDim.y + (int)threadIdx.y;
    const int ow = tileX * (int)blockDim.x + (int)threadIdx.x;
    if (oh >= OH || ow >= OW) return;

    const int n = blockIdx.z;
    const float* in_n = input + (long)n * C * H * W;
    const int ih0 = oh * stride - pad;
    const int iw0 = ow * stride - pad;

    float acc[OC_TILE];
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) acc[i] = 0.0f;

    for (int c = 0; c < C; ++c) {
        const float* in_c = in_n + (long)c * H * W;
        for (int kh = 0; kh < KH; ++kh) {
            const int ih = ih0 + kh;
            const bool hok = (ih >= 0) && (ih < H);
            for (int kw = 0; kw < KW; ++kw) {
                const int iw = iw0 + kw;
                float v = (hok && iw >= 0 && iw < W) ? in_c[ih * W + iw] : 0.0f;
                const int woff = (c * KH + kh) * KW + kw;
                #pragma unroll
                for (int i = 0; i < OC_TILE; ++i) {
                    acc[i] = fmaf(v, s_weight[i * ksize + woff], acc[i]);
                }
            }
        }
    }

    float* out_n = output + (long)n * OC * OH * OW + oh * OW + ow;
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) {
        const int oc = oc0 + i;
        if (oc < OC) {
            out_n[(long)oc * OH * OW] = relu_hardswish(acc[i] + bias[oc]);
        }
    }
}

torch::Tensor fused_conv_act(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride, int64_t pad, bool use_tc)
{
    TORCH_CHECK(input.is_cuda() && weight.is_cuda() && bias.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "bad ranks");

    const long N  = input.size(0);
    const long C  = input.size(1);
    const long H  = input.size(2);
    const long W  = input.size(3);
    const long OC = weight.size(0);
    const long KH = weight.size(2);
    const long KW = weight.size(3);
    const long OH = (H + 2 * pad - KH) / stride + 1;
    const long OW = (W + 2 * pad - KW) / stride + 1;

    auto output = at::empty({N, OC, OH, OW}, input.options());
    if (output.numel() == 0 || N == 0) return output;

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

    // Tensor-core 3xTF32 fast path.
    const long octiles = (OC + BN - 1) / BN;
    const long tilesX  = (OW + BM - 1) / BM;
    if (use_tc && C == 8 && KH == 3 && KW == 3 &&
        at::cuda::getCurrentDeviceProperties()->major >= 8 &&
        OH <= 65535 && N * octiles <= 65535) {
        const dim3 grid((unsigned)tilesX, (unsigned)OH, (unsigned)(N * octiles));
        conv_act_tc_kernel<8, 3, 3><<<grid, NTHREADS, 0, stream>>>(
            input.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            output.data_ptr<float>(), (int)N, (int)H, (int)W, (int)OC,
            (int)OH, (int)OW, (int)stride, (int)pad, (int)octiles);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return output;
    }

    // Direct FP32 fallback.
    const dim3 block(32, 8);
    const long tx = (OW + 31) / 32;
    const long ty = (OH + 7) / 8;
    if (tx * ty <= 65535 && N <= 65535 &&
        (size_t)OC_TILE * C * KH * KW * sizeof(float) <= 48 * 1024) {
        const dim3 grid((unsigned)(tx * ty), (unsigned)((OC + OC_TILE - 1) / OC_TILE), (unsigned)N);
        if (KH == 3 && KW == 3 && C == 8) {
            conv_hswish_relu_kernel<8, 3><<<grid, block, 0, stream>>>(
                input.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
                output.data_ptr<float>(), (int)N, (int)H, (int)W, (int)OC,
                (int)OH, (int)OW, (int)stride, (int)pad);
        } else {
            const size_t shmem = (size_t)OC_TILE * C * KH * KW * sizeof(float);
            conv_hswish_relu_generic_kernel<<<grid, block, shmem, stream>>>(
                input.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
                output.data_ptr<float>(), (int)N, (int)C, (int)H, (int)W, (int)OC,
                (int)KH, (int)KW, (int)OH, (int)OW, (int)stride, (int)pad);
        }
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return output;
    }

    // Last-resort ATen fallback (identical semantics).
    auto y = at::conv2d(input, weight, bias, {stride, stride}, {pad, pad});
    y = at::hardswish(y);
    y = at::relu(y);
    output.copy_(y);
    return output;
}
"""

fused_conv_act_mod = load_inline(
    name="fused_conv_hardswish_relu_v5",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_conv_act"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d + HardSwish + ReLU fused into a single custom CUDA kernel.
    Fast path uses tensor cores with 3xTF32 (near-fp32 accuracy, enabled only after a
    one-time runtime self-check against a strict fp32 reference). Falls back to a
    direct FP32 fused kernel (or the original PyTorch ops) for other shapes/devices.
    Keeps the same nn.Conv2d submodule so state_dict / initialization is identical
    to the reference Model.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        kh, kw = self.conv.kernel_size
        self._ok = (self.conv.groups == 1 and
                    self.conv.dilation == (1, 1) and
                    self.conv.stride[0] == self.conv.stride[1] and
                    self.conv.padding[0] == self.conv.padding[1] and
                    self.conv.bias is not None and
                    (in_channels * kh * kw * 8 * 4) <= 48 * 1024)
        self._use_tc = False
        self._checked = False
        if torch.cuda.is_available():
            self._try_check()

    @torch.no_grad()
    def _try_check(self):
        self._checked = True
        self._use_tc = False
        try:
            ic = self.conv.in_channels
            oc = self.conv.out_channels
            kh, kw = self.conv.kernel_size
            if not (ic == 8 and kh == 3 and kw == 3):
                return
            if not torch.cuda.is_available():
                return
            if torch.cuda.get_device_capability(torch.cuda.current_device())[0] < 8:
                return
            dev = torch.device('cuda')
            tmp = nn.Conv2d(ic, oc, kh,
                            stride=self.conv.stride,
                            padding=self.conv.padding).to(dev)
            x = torch.rand(2, ic, 34, 34, device=dev)
            prev = torch.backends.cudnn.allow_tf32
            try:
                torch.backends.cudnn.allow_tf32 = False
                ref = F.relu(F.hardswish(F.conv2d(x, tmp.weight, tmp.bias,
                                                  tmp.stride, tmp.padding)))
            finally:
                torch.backends.cudnn.allow_tf32 = prev
            tc = fused_conv_act_mod.fused_conv_act(
                x, tmp.weight, tmp.bias, tmp.stride[0], tmp.padding[0], True)
            # Accept the TC path only if it is numerically ~fp32-exact.
            self._use_tc = bool(torch.allclose(tc, ref, rtol=1e-4, atol=1e-5))
        except Exception:
            self._use_tc = False

    def forward(self, x):
        if self._ok and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            if not self._checked:
                self._try_check()
            return fused_conv_act_mod.fused_conv_act(
                x.contiguous(),
                self.conv.weight.contiguous(),
                self.conv.bias.contiguous(),
                self.conv.stride[0],
                self.conv.padding[0],
                self._use_tc,
            )
        x = self.conv(x)
        x = F.hardswish(x)
        x = torch.relu(x)
        return x
```