```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 (stride=1, direct) + HardSwish + ReLU in a single pass.
# Key ideas:
#  - Each thread computes OC_TILE output channels for one spatial location, so each
#    loaded input value is reused OC_TILE times (higher arithmetic intensity).
#  - The (tiny) filter weights are staged in shared memory and broadcast to all threads.
#  - HardSwish (x * clamp(x+3,0,6) / 6) and the trailing ReLU are fused into the
#    epilogue, eliminating two full element-wise passes over the output tensor.
_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>

#define OC_TILE 8

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

// Specialized kernel: compile-time channel count / kernel size (fully unrolled).
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]);
        }
    }
}

// Generic fallback kernel: runtime channel count / kernel size.
__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");

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

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

    const dim3 block(32, 8);
    const int tilesX = (OW + 31) / 32;
    const int tilesY = (OH + 7) / 8;
    const dim3 grid(tilesX * tilesY, (OC + OC_TILE - 1) / OC_TILE, N);

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

    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>(), N, H, W, OC, OH, OW, (int)stride, (int)pad);
    } else {
        const size_t shmem = (size_t)OC_TILE * C * KH * KW * sizeof(float);
        TORCH_CHECK(shmem <= 48 * 1024, "kernel too large for fused conv");
        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>(), N, C, H, W, OC, KH, KW, OH, OW, (int)stride, (int)pad);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

fused_conv_act_mod = load_inline(
    name="fused_conv_hardswish_relu_v1",
    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.
    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
        # Fused kernel stages OC_TILE(=8) filters in shared memory.
        self._fused_ok = (in_channels * kh * kw * 8 * 4) <= 48 * 1024

    def forward(self, x):
        if self._fused_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
```