```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 (stride=1, pad=0) + HardSwish + ReLU CUDA kernel.
# Key idea: ReLU(HardSwish(x)) == (x > 0) ? x * min(x + 3, 6) / 6 : 0
# so the whole activation collapses into a couple of FMA-friendly ops
# applied directly on the convolution accumulator (no extra memory passes).
_fused_conv_act_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TILE 16
#define OC_TILE 8

// Direct convolution kernel:
//  - each block computes a TILE x TILE output tile for OC_TILE output channels
//    of one batch element.
//  - the (TILE+K-1)^2 x Cin input patch and the OC_TILE filter banks are staged
//    in shared memory.
//  - fused HardSwish+ReLU epilogue.
__global__ void conv2d_hardswish_relu_kernel(
    const float* __restrict__ input,   // N x Cin x H x W
    const float* __restrict__ weight,  // Cout x Cin x K x K
    const float* __restrict__ bias,    // Cout (or nullptr)
    float* __restrict__ output,        // N x Cout x Hout x Wout
    const int N, const int Cin, const int H, const int W,
    const int Cout, const int K, const int Hout, const int Wout,
    const int has_bias)
{
    extern __shared__ float smem[];
    const int TS = TILE + K - 1;
    const int in_tile_elems = Cin * TS * TS;
    float* s_in = smem;
    float* s_w  = smem + in_tile_elems;

    const int n   = blockIdx.z;
    const int oc0 = blockIdx.y * OC_TILE;
    const int tiles_w = (Wout + TILE - 1) / TILE;
    const int tile_y  = (blockIdx.x / tiles_w) * TILE;
    const int tile_x  = (blockIdx.x % tiles_w) * TILE;

    const int tid = threadIdx.y * TILE + threadIdx.x;   // block is TILE x TILE
    const int nthreads = TILE * TILE;

    // Stage the input patch (zero-filled outside the image; only threads that
    // produce valid outputs ever read the valid region).
    for (int idx = tid; idx < in_tile_elems; idx += nthreads) {
        const int c   = idx / (TS * TS);
        const int rem = idx - c * TS * TS;
        const int r   = rem / TS;
        const int col = rem - r * TS;
        const int ih  = tile_y + r;
        const int iw  = tile_x + col;
        float v = 0.0f;
        if (ih < H && iw < W) {
            v = input[((n * Cin + c) * H + ih) * W + iw];
        }
        s_in[idx] = v;
    }

    // Stage the OC_TILE filter banks.
    const int w_elems = OC_TILE * Cin * K * K;
    for (int idx = tid; idx < w_elems; idx += nthreads) {
        const int i  = idx / (Cin * K * K);
        const int oc = oc0 + i;
        s_w[idx] = (oc < Cout) ? weight[oc * Cin * K * K + (idx - i * Cin * K * K)] : 0.0f;
    }
    __syncthreads();

    const int oy = tile_y + threadIdx.y;
    const int ox = tile_x + threadIdx.x;
    if (oy >= Hout || ox >= Wout) return;

    float acc[OC_TILE];
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) {
        const int oc = oc0 + i;
        acc[i] = (has_bias && oc < Cout) ? bias[oc] : 0.0f;
    }

    const int ty = threadIdx.y;
    const int tx = threadIdx.x;
    for (int c = 0; c < Cin; ++c) {
        const float* tile_c = s_in + c * TS * TS;
        const float* w_c    = s_w  + c * K * K;
        for (int ky = 0; ky < K; ++ky) {
            const int row = (ty + ky) * TS + tx;
            for (int kx = 0; kx < K; ++kx) {
                const float iv = tile_c[row + kx];
                #pragma unroll
                for (int i = 0; i < OC_TILE; ++i) {
                    acc[i] += iv * w_c[i * Cin * K * K + ky * K + kx];
                }
            }
        }
    }

    // Fused epilogue: ReLU(HardSwish(x)) = (x > 0) ? x * min(x + 3, 6) / 6 : 0
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i) {
        const int oc = oc0 + i;
        if (oc < Cout) {
            const float x = acc[i];
            const float y = (x > 0.0f) ? (x * fminf(x + 3.0f, 6.0f) / 6.0f) : 0.0f;
            output[((n * Cout + oc) * Hout + oy) * Wout + ox] = y;
        }
    }
}

torch::Tensor conv2d_hardswish_relu_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias) {
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "input/weight must be CUDA tensors");
    TORCH_CHECK(input.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(input.dim() == 4 && weight.dim() == 4, "input/weight must be 4-D");

    auto x = input.contiguous();
    auto w = weight.contiguous();

    const int N    = (int)x.size(0);
    const int Cin  = (int)x.size(1);
    const int H    = (int)x.size(2);
    const int W    = (int)x.size(3);
    const int Cout = (int)w.size(0);
    const int K    = (int)w.size(2);
    TORCH_CHECK((int)w.size(1) == Cin && (int)w.size(3) == K, "weight shape mismatch");

    const int Hout = H - K + 1;
    const int Wout = W - K + 1;
    TORCH_CHECK(Hout > 0 && Wout > 0, "kernel larger than input");

    auto out = torch::empty({N, Cout, Hout, Wout}, x.options());

    const int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;
    auto b = has_bias ? bias.contiguous() : bias;
    const float* bias_ptr = has_bias ? b.data_ptr<float>() : nullptr;

    dim3 block(TILE, TILE);
    const int tiles_w = (Wout + TILE - 1) / TILE;
    const int tiles_h = (Hout + TILE - 1) / TILE;
    dim3 grid(tiles_w * tiles_h, (Cout + OC_TILE - 1) / OC_TILE, N);

    const int TS = TILE + K - 1;
    const size_t smem = (size_t)(Cin * TS * TS + OC_TILE * Cin * K * K) * sizeof(float);

    auto stream = at::cuda::getCurrentCUDAStream();
    conv2d_hardswish_relu_kernel<<<grid, block, smem, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
        N, Cin, H, W, Cout, K, Hout, Wout, has_bias);

    return out;
}
"""

_fused_conv_act_cpp_source = (
    "torch::Tensor conv2d_hardswish_relu_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);"
)

_fused_conv_act = load_inline(
    name="fused_conv2d_hardswish_relu",
    cpp_sources=_fused_conv_act_cpp_source,
    cuda_sources=_fused_conv_act_cuda_source,
    functions=["conv2d_hardswish_relu_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d + HardSwish + ReLU fused into a single CUDA kernel.

    The custom kernel performs a tiled direct convolution and applies the
    algebraically-simplified activation
        ReLU(HardSwish(x)) = (x > 0) ? x * min(x + 3, 6) / 6 : 0
    directly on the accumulator, eliminating two full read/write passes over
    the feature map. Falls back to the reference implementation for
    unsupported configurations.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)

    def _can_fuse(self, x):
        c = self.conv
        if not (x.is_cuda and x.dtype == torch.float32):
            return False
        if (c.stride != (1, 1) or c.padding != (0, 0)
                or c.dilation != (1, 1) or c.groups != 1):
            return False
        if c.kernel_size[0] != c.kernel_size[1]:
            return False
        k = c.kernel_size[0]
        ts = 16 + k - 1
        smem_bytes = (x.size(1) * ts * ts + 8 * x.size(1) * k * k) * 4
        return smem_bytes <= 46 * 1024

    def forward(self, x):
        if self._can_fuse(x):
            bias = self.conv.bias
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=x.dtype)
            return _fused_conv_act.conv2d_hardswish_relu_cuda(x, self.conv.weight, bias)
        out = self.conv(x)
        out = F.hardswish(out)
        return torch.relu(out)
```