```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(3x3, stride=1, pad=0) + HardSwish + ReLU CUDA kernel.
#
# Strategy: direct convolution with shared-memory tiling.
#  - Each block computes a 63x63 output tile for 2 adjacent output channels.
#  - The required 65x65 input slice for each input channel is staged in shared
#    memory (double buffered across input channels).
#  - Each of the 21x21=441 threads computes a 3x3 spatial patch x 2 output
#    channels (18 accumulators), so each input value loaded from shared memory
#    feeds 18 FMAs, making the kernel compute-bound.
#  - HardSwish followed by ReLU is applied in the epilogue for free:
#    relu(hardswish(x)) = max( x * min(max(x+3,0),6) / 6, 0 )
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <algorithm>

#define TILE 63
#define TS 65
#define NTHREADS 441
#define LOAD_ITERS 10

template <bool INTERIOR>
__global__ __launch_bounds__(NTHREADS, 2)
void conv3x3_hardswish_relu_kernel(
    const float* __restrict__ X,
    const float* __restrict__ Wt,
    const float* __restrict__ Bias,
    float* __restrict__ Y,
    int C_in, int H, int W,
    int C_out, int H_out, int W_out, int OC2)
{
    extern __shared__ float smem[];
    const int C9 = C_in * 9;
    float* s_w = smem;              // [2*C9] weights for the 2 output channels
    float* b0  = s_w + 2 * C9;      // [TS*TS] input slice buffer 0
    float* b1  = b0 + TS * TS;      // [TS*TS] input slice buffer 1

    const int tx  = threadIdx.x;    // 0..20
    const int ty  = threadIdx.y;    // 0..20
    const int tid = ty * 21 + tx;

    const int z   = blockIdx.z;
    const int n   = z / OC2;
    const int ocp = z - n * OC2;
    const int oc0 = ocp * 2;

    const int oy = blockIdx.y * TILE;
    const int ox = blockIdx.x * TILE;

    // Weights: two adjacent output channels are contiguous in memory.
    const float* wsrc = Wt + (size_t)oc0 * C9;
    for (int t = tid; t < 2 * C9; t += NTHREADS) s_w[t] = wsrc[t];

    const int r00 = 3 * ty;
    const int c00 = 3 * tx;

    // Cooperative-loader coordinates (linear walk over the TS x TS slice).
    const int lr = tid / TS;
    const int lc = tid - lr * TS;

    const float* xn = X + (size_t)n * C_in * H * W;

    // Load slice for ic = 0 into b0.
    {
        const float* src = xn + (size_t)oy * W + ox;
        int r = lr, c = lc;
        #pragma unroll
        for (int it = 0; it < LOAD_ITERS; ++it) {
            if (r < TS) {
                float v;
                if (INTERIOR) {
                    v = src[(size_t)r * W + c];
                } else {
                    int gy = oy + r, gx = ox + c;
                    v = (gy < H && gx < W) ? src[(size_t)r * W + c] : 0.f;
                }
                b0[r * TS + c] = v;
            }
            r += 6; c += 51; if (c >= TS) { c -= TS; r += 1; }
        }
    }
    __syncthreads();

    const float bias0 = Bias[oc0];
    const float bias1 = Bias[oc0 + 1];
    float a0[3][3], a1[3][3];
    #pragma unroll
    for (int i = 0; i < 3; ++i)
        #pragma unroll
        for (int j = 0; j < 3; ++j) { a0[i][j] = bias0; a1[i][j] = bias1; }

    for (int ic = 0; ic < C_in; ++ic) {
        const float* cur = (ic & 1) ? b1 : b0;
        float*       nxt = (ic & 1) ? b0 : b1;

        // Preload next input-channel slice into the other buffer.
        if (ic + 1 < C_in) {
            const float* src = xn + (size_t)(ic + 1) * H * W + (size_t)oy * W + ox;
            int r = lr, c = lc;
            #pragma unroll
            for (int it = 0; it < LOAD_ITERS; ++it) {
                if (r < TS) {
                    float v;
                    if (INTERIOR) {
                        v = src[(size_t)r * W + c];
                    } else {
                        int gy = oy + r, gx = ox + c;
                        v = (gy < H && gx < W) ? src[(size_t)r * W + c] : 0.f;
                    }
                    nxt[r * TS + c] = v;
                }
                r += 6; c += 51; if (c >= TS) { c -= TS; r += 1; }
            }
        }

        // Compute 3x3 (spatial) x 2 (channels) outputs from this slice.
        float v[5][5];
        const float* curT = cur + r00 * TS + c00;
        #pragma unroll
        for (int i = 0; i < 5; ++i)
            #pragma unroll
            for (int j = 0; j < 5; ++j)
                v[i][j] = curT[i * TS + j];

        const float* w0p = s_w + ic * 9;
        const float* w1p = s_w + C9 + ic * 9;
        #pragma unroll
        for (int dy = 0; dy < 3; ++dy) {
            #pragma unroll
            for (int dx = 0; dx < 3; ++dx) {
                const float w0 = w0p[dy * 3 + dx];
                const float w1 = w1p[dy * 3 + dx];
                #pragma unroll
                for (int py = 0; py < 3; ++py) {
                    #pragma unroll
                    for (int px = 0; px < 3; ++px) {
                        const float vv = v[py + dy][px + dx];
                        a0[py][px] = fmaf(vv, w0, a0[py][px]);
                        a1[py][px] = fmaf(vv, w1, a1[py][px]);
                    }
                }
            }
        }
        __syncthreads();
    }

    // Epilogue: fused HardSwish + ReLU.
    float* ybase = Y + ((size_t)n * C_out + oc0) * H_out * W_out;
    const size_t plane = (size_t)H_out * W_out;
    #pragma unroll
    for (int py = 0; py < 3; ++py) {
        const int gy = oy + r00 + py;
        #pragma unroll
        for (int px = 0; px < 3; ++px) {
            const int gx = ox + c00 + px;
            if (INTERIOR || (gy < H_out && gx < W_out)) {
                const float t0 = a0[py][px];
                const float h0 = (t0 * fminf(fmaxf(t0 + 3.f, 0.f), 6.f)) / 6.f;
                const float t1 = a1[py][px];
                const float h1 = (t1 * fminf(fmaxf(t1 + 3.f, 0.f), 6.f)) / 6.f;
                const size_t off = (size_t)gy * W_out + gx;
                ybase[off]         = fmaxf(h0, 0.f);
                ybase[plane + off] = fmaxf(h1, 0.f);
            }
        }
    }
}

torch::Tensor fused_conv3x3_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.dim() == 4 && w.dim() == 4, "bad dims");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "expect float32");

    const int N     = (int)x.size(0);
    const int C_in  = (int)x.size(1);
    const int H     = (int)x.size(2);
    const int W     = (int)x.size(3);
    const int C_out = (int)w.size(0);
    TORCH_CHECK((C_out & 1) == 0, "out_channels must be even");
    TORCH_CHECK((int)w.size(1) == C_in && w.size(2) == 3 && w.size(3) == 3, "bad weight shape");
    TORCH_CHECK(H >= 3 && W >= 3, "input too small");

    const int H_out = H - 2;
    const int W_out = W - 2;
    const int OC2   = C_out / 2;

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

    dim3 block(21, 21);
    dim3 grid((W_out + TILE - 1) / TILE, (H_out + TILE - 1) / TILE, N * OC2);
    const size_t smem = (size_t)(2 * C_in * 9 + 2 * TS * TS) * sizeof(float);

    const bool interior =
        ((int)grid.x - 1) * TILE + TS <= W && ((int)grid.y - 1) * TILE + TS <= H;

    auto stream = at::cuda::getCurrentCUDAStream();
    if (interior) {
        conv3x3_hardswish_relu_kernel<true><<<grid, block, smem, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(),
            out.data_ptr<float>(), C_in, H, W, C_out, H_out, W_out, OC2);
    } else {
        conv3x3_hardswish_relu_kernel<false><<<grid, block, smem, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(),
            out.data_ptr<float>(), C_in, H, W, C_out, H_out, W_out, OC2);
    }
    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
torch::Tensor fused_conv3x3_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b);
"""

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


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d + HardSwish + ReLU fused into a single CUDA kernel.
    Falls back to the reference PyTorch implementation for unsupported configs.
    """

    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_ok = (
            self.conv.kernel_size == (3, 3)
            and self.conv.stride == (1, 1)
            and self.conv.padding == (0, 0)
            and self.conv.dilation == (1, 1)
            and self.conv.groups == 1
            and self.conv.bias is not None
            and out_channels % 2 == 0
            and in_channels <= 1024
        )

    def forward(self, x):
        if (
            self._fused_ok
            and x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 4
            and x.size(2) >= 3
            and x.size(3) >= 3
        ):
            return fused_ext.fused_conv3x3_hardswish_relu(
                x.contiguous(),
                self.conv.weight.contiguous(),
                self.conv.bias.contiguous(),
            )
        x = self.conv(x)
        x = F.hardswish(x)
        return torch.relu(x)
```