```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: implicit-GEMM tensor-core (TF32 WMMA) kernel specialized for
# C=8, K=3x3. Each block computes a 128(ow) x 64(oc) output tile:
#   - A staged in shared as col-major [BK=24][BM=128] (fully unrolled, compile-time
#     im2col index decomposition, coalesced global loads, conflict-free shared stores).
#   - B (filters) staged as col-major [BN=64][BK=24].
#   - 4 warps each compute a 32x64 sub-tile via 16x16x8 tf32 wmma fragments.
#   - Epilogue stages C through padded shared memory (ldm=68), applies
#     bias + HardSwish + ReLU, then writes coalesced.
# Fallback: previous direct FP32 SIMT fused kernel (any shape/arch), then ATen.
_fused_cpp_src = "torch::Tensor fused_conv_act(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int64_t stride, int64_t pad);"

_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 CLD (BN + 4)

__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);
}

__device__ __forceinline__ float to_tf32(float x) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    float r;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=f"(r) : "f"(x));
    return r;
#else
    return x;
#endif
}

// ---------------- Tensor-core (TF32) fused conv + HardSwish + ReLU --------------
template <int C, int KH, int KW>
__global__ void __launch_bounds__(BM) 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;
    extern __shared__ __align__(32) float smem[];
    float* As_t = smem;              // [BK][BM]  col-major A tile
    float* Bs_t = smem + BK * BM;    // [BN][BK]  col-major B tile
    float* Cs   = smem;              // [BM][CLD] reused after the k-loop

    const int tid  = threadIdx.x;    // 128 threads = 4 warps
    const int warp = tid >> 5;

    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;

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

    const float* in_n = input + (long)n * (C * H * W);

    #pragma unroll
    for (int k0 = 0; k0 < K; k0 += BK) {
        __syncthreads();
        // Stage A tile (col-major): fully unrolled -> k is compile-time.
        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            const int k  = k0 + kk;
            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 = oh * stride + kh - pad;
            const int iw = ow * stride + kw - pad;
            float v = 0.0f;
            if (ih >= 0 && ih < H && iw >= 0 && iw < W)
                v = in_n[(c * H + ih) * W + iw];
            As_t[kk * BM + tid] = to_tf32(v);
        }
        // Stage B tile (col-major): [BN][BK].
        #pragma unroll
        for (int i = 0; i < (BN * BK) / BM; ++i) {
            const int idx = tid + i * BM;
            const int ocl = idx / BK;
            const int kk  = idx - ocl * BK;
            const int k   = k0 + kk;
            const int oc  = oc0 + ocl;
            float v = 0.0f;
            if (oc < OC) v = weight[(long)oc * K + k];
            Bs_t[idx] = to_tf32(v);
        }
        __syncthreads();
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::col_major> a0, a1;
            wmma::load_matrix_sync(a0, As_t + kk * BM + warp * 32, BM);
            wmma::load_matrix_sync(a1, As_t + kk * BM + warp * 32 + 16, BM);
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> b;
                wmma::load_matrix_sync(b, Bs_t + j * 16 * BK + kk, BK);
                wmma::mma_sync(acc[0][j], a0, b, acc[0][j]);
                wmma::mma_sync(acc[1][j], a1, b, acc[1][j]);
            }
        }
    }
    __syncthreads();
    // Stage C tile in shared (padded leading dim avoids bank conflicts).
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::store_matrix_sync(Cs + (warp * 32 + i * 16) * CLD + j * 16,
                                    acc[i][j], CLD, wmma::mem_row_major);
    __syncthreads();
    // Fused epilogue: bias + HardSwish + ReLU, coalesced global writes.
    if (ow < OW) {
        float* out_p = output + (long)n * OC * OH * OW + (long)oh * OW + ow;
        const float* cs_p = Cs + tid * CLD;
        #pragma unroll 4
        for (int j = 0; j < BN; ++j) {
            const int oc = oc0 + j;
            if (oc < OC) {
                out_p[(long)oc * OH * OW] = relu_hardswish(cs_p[j] + bias[oc]);
            }
        }
    }
#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)
{
    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 fast path.
    const long octiles = (OC + BN - 1) / BN;
    const long tilesX  = (OW + BM - 1) / BM;
    if (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));
        constexpr size_t shmem = (size_t)BM * CLD * sizeof(float);
        conv_act_tc_kernel<8, 3, 3><<<grid, BM, shmem, 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 (keeps output buffer 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_v2",
    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 TF32 tensor cores (implicit GEMM); falls back to a direct FP32
    fused kernel (or ATen) 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)
        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])

    def forward(self, x):
        if self._ok and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            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],
            )
        x = self.conv(x)
        x = F.hardswish(x)
        x = torch.relu(x)
        return x
```