```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

# ---------------------------------------------------------------------------
# Custom fused HardSwish + ReLU CUDA kernel.
#
# The reference model runs:
#     x = conv(x)                       (cuDNN kernel)
#     x = F.hardswish(x)                (elementwise kernel, full read+write)
#     x = F.relu(x)                     (elementwise kernel, full read+write)
#
# For the given shapes the conv output is 128 x 64 x 126 x 126 (~520 MB), so
# the two separate elementwise ops cost ~2 GB of extra DRAM traffic. This
# kernel fuses both activations into a single read+write pass and uses
# vectorized 128-bit (float4) accesses.
#
# The math exactly matches ATen's hardswish (x * clamp(x+3, 0, 6) / 6, with
# the multiply performed before the divide) followed by relu (max(h, 0)), so
# results are bit-exact with the reference implementation.
# ---------------------------------------------------------------------------

fused_act_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ float hardswish_relu(float x) {
    // hardswish: x * clamp(x + 3, 0, 6) / 6  (same operation order as ATen)
    float h = x * fminf(fmaxf(x + 3.0f, 0.0f), 6.0f) / 6.0f;
    // relu
    return fmaxf(h, 0.0f);
}

__global__ void hardswish_relu_vec4_kernel(const float4* __restrict__ in,
                                           float4* __restrict__ out,
                                           int64_t n4) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = in[i];
        v.x = hardswish_relu(v.x);
        v.y = hardswish_relu(v.y);
        v.z = hardswish_relu(v.z);
        v.w = hardswish_relu(v.w);
        out[i] = v;
    }
}

__global__ void hardswish_relu_scalar_kernel(const float* __restrict__ in,
                                             float* __restrict__ out,
                                             int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = hardswish_relu(in[i]);
    }
}

torch::Tensor hardswish_relu_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");

    at::Tensor xc = x;
    if (!xc.is_contiguous() &&
        !xc.is_contiguous(at::MemoryFormat::ChannelsLast)) {
        xc = xc.contiguous();
    }
    auto out = torch::empty_like(xc);

    const int64_t n = xc.numel();
    if (n == 0) {
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;

    if ((n & 3) == 0) {
        // Vectorized path: one 128-bit load + one 128-bit store per thread.
        const int64_t n4 = n >> 2;
        const unsigned int grid =
            (unsigned int)((n4 + block - 1) / block);
        hardswish_relu_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            n4);
    } else {
        const unsigned int grid =
            (unsigned int)((n + block - 1) / block);
        hardswish_relu_scalar_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(),
            out.data_ptr<float>(),
            n);
    }
    return out;
}
"""

fused_act_cpp_src = "torch::Tensor hardswish_relu_cuda(torch::Tensor x);"

# Compile the fused activation kernel.
fused_act = load_inline(
    name="fused_hardswish_relu",
    cpp_sources=fused_act_cpp_src,
    cuda_sources=fused_act_cuda_src,
    functions=["hardswish_relu_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Conv2d (unchanged cuDNN call) + custom fused HardSwish+ReLU CUDA kernel.

    Fusing the two back-to-back elementwise operators into one kernel removes
    an entire read+write pass over the ~520 MB activation tensor, cutting
    total DRAM traffic by roughly a third and eliminating one kernel launch.
    """
    def __init__(self, in_channels, out_channels, kernel_size):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)

    def forward(self, x):
        x = self.conv(x)
        if x.is_cuda and x.dtype == torch.float32:
            return fused_act.hardswish_relu_cuda(x)
        # Fallback for non-CUDA / non-fp32 inputs.
        return torch.relu(F.hardswish(x))
```