```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 CUDA kernel: Conv2d (3x3, stride 1, pad 0) + HardSwish + ReLU in a single pass.
# The composition relu(hardswish(x)) is applied on-the-fly before writing to global
# memory, eliminating two full elementwise kernel launches and their memory traffic.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TILE 32

__device__ __forceinline__ float hardswish_relu(float x) {
    // hardswish(x) = x * relu6(x + 3) / 6, then relu()
    float h = x * fminf(fmaxf(x + 3.0f, 0.0f), 6.0f) / 6.0f;
    return fmaxf(h, 0.0f);
}

template <int K>
__global__ void conv2d_hardswish_relu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ out,
    const int N, const int C_in, const int H, const int W,
    const int C_out, const int H_out, const int W_out)
{
    constexpr int IN_TILE = TILE + K - 1;   // input tile rows/cols needed
    constexpr int LD = IN_TILE + 1;         // padded leading dim to reduce bank conflicts

    extern __shared__ float smem[];
    float* s_in = smem;                          // C_in * IN_TILE * LD
    float* s_w  = smem + C_in * IN_TILE * LD;    // C_in * K * K

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    const int nthreads = blockDim.x * blockDim.y;

    const int oc = blockIdx.z % C_out;
    const int n  = blockIdx.z / C_out;

    // Load this block's output-channel weights into shared memory.
    const float* wp = w + (long long)oc * C_in * K * K;
    for (int i = tid; i < C_in * K * K; i += nthreads) {
        s_w[i] = wp[i];
    }

    const int oh0 = blockIdx.y * TILE;
    const int ow0 = blockIdx.x * TILE;

    // Cooperatively load the input tile (zero-filled outside the image).
    const long long xBase = (long long)n * C_in * H * W;
    const int total = C_in * IN_TILE * IN_TILE;
    for (int idx = tid; idx < total; idx += nthreads) {
        int c   = idx / (IN_TILE * IN_TILE);
        int rem = idx - c * IN_TILE * IN_TILE;
        int iy  = rem / IN_TILE;
        int ix  = rem - iy * IN_TILE;
        int gy  = oh0 + iy;
        int gx  = ow0 + ix;
        float v = 0.0f;
        if (gy < H && gx < W) {
            v = x[xBase + ((long long)c * H + gy) * W + gx];
        }
        s_in[(c * IN_TILE + iy) * LD + ix] = v;
    }
    __syncthreads();

    // Each thread computes a 2x2 block of outputs (tile coords (2ty,2tx)..(2ty+1,2tx+1)).
    const float bias = b[oc];
    float acc00 = bias, acc01 = bias, acc10 = bias, acc11 = bias;

    const int ry = 2 * ty;
    const int rx = 2 * tx;

    for (int c = 0; c < C_in; ++c) {
        const float* s = s_in + c * IN_TILE * LD;
        // Stage the (K+1)x(K+1) patch for the 2x2 outputs into registers.
        float v[K + 1][K + 1];
        #pragma unroll
        for (int i = 0; i < K + 1; ++i) {
            #pragma unroll
            for (int j = 0; j < K + 1; ++j) {
                v[i][j] = s[(ry + i) * LD + (rx + j)];
            }
        }
        const float* wc = s_w + c * K * K;
        #pragma unroll
        for (int ky = 0; ky < K; ++ky) {
            #pragma unroll
            for (int kx = 0; kx < K; ++kx) {
                float wv = wc[ky * K + kx];
                acc00 += v[ky][kx]     * wv;
                acc01 += v[ky][kx + 1] * wv;
                acc10 += v[ky + 1][kx] * wv;
                acc11 += v[ky + 1][kx + 1] * wv;
            }
        }
    }

    // Fused activation + guarded store.
    const int oh = oh0 + ry;
    const int ow = ow0 + rx;
    const long long oBase = ((long long)n * C_out + oc) * H_out * W_out;
    if (oh < H_out) {
        if (ow < W_out)     out[oBase + (long long)oh * W_out + ow]             = hardswish_relu(acc00);
        if (ow + 1 < W_out) out[oBase + (long long)oh * W_out + ow + 1]         = hardswish_relu(acc01);
    }
    if (oh + 1 < H_out) {
        if (ow < W_out)     out[oBase + (long long)(oh + 1) * W_out + ow]       = hardswish_relu(acc10);
        if (ow + 1 < W_out) out[oBase + (long long)(oh + 1) * W_out + ow + 1]   = hardswish_relu(acc11);
    }
}

torch::Tensor conv2d_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();

    const int N     = xc.size(0);
    const int C_in  = xc.size(1);
    const int H     = xc.size(2);
    const int W     = xc.size(3);
    const int C_out = wc.size(0);
    const int K     = wc.size(2);
    const int H_out = H - K + 1;
    const int W_out = W - K + 1;

    auto out = torch::empty({N, C_out, H_out, W_out}, xc.options());

    constexpr int KK = 3;
    constexpr int IN_TILE = TILE + KK - 1;
    constexpr int LD = IN_TILE + 1;
    const size_t smem = (size_t)(C_in * IN_TILE * LD + C_in * KK * KK) * sizeof(float);

    if (K == KK && H_out > 0 && W_out > 0 && smem <= 48 * 1024 && (long long)N * C_out <= 65535) {
        auto stream = at::cuda::getCurrentCUDAStream();
        dim3 block(16, 16, 1);
        dim3 grid((W_out + TILE - 1) / TILE, (H_out + TILE - 1) / TILE, N * C_out);
        conv2d_hardswish_relu_kernel<KK><<<grid, block, smem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            out.data_ptr<float>(), N, C_in, H, W, C_out, H_out, W_out);
        return out;
    }

    // Generic fallback for unexpected configurations.
    out = at::conv2d(xc, wc, bc);
    out = at::hardswish(out);
    out = at::relu(out);
    return out;
}
"""

cpp_source = "torch::Tensor conv2d_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b);"

fused_op = load_inline(
    name="conv2d_hardswish_relu_fused",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv2d_hardswish_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d + HardSwish + ReLU fused into a single custom CUDA kernel.
    Keeps the same nn.Conv2d parameters so the state dict matches the original 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.fused = fused_op

    def forward(self, x):
        w, b = self.conv.weight, self.conv.bias
        if (x.is_cuda and x.dtype == torch.float32
                and self.conv.stride == (1, 1)
                and self.conv.padding == (0, 0)
                and self.conv.dilation == (1, 1)
                and self.conv.groups == 1):
            return self.fused.conv2d_hardswish_relu(x, w, b)
        # Safe fallback matching the original semantics.
        x = F.conv2d(x, w, b, self.conv.stride, self.conv.padding,
                     self.conv.dilation, self.conv.groups)
        x = F.hardswish(x)
        x = torch.relu(x)
        return x
```