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

# Fully fused Conv2d(3x3) + HardSwish + ReLU kernel.
#
# Key idea: a single kernel computes the convolution and applies the
# activation relu(hardswish(x)) in registers before writing out, so the
# output tensor is never round-tripped through HBM for the activations.
#
# relu(hardswish(x)) simplifies nicely:
#   - x <= 0          -> 0            (hardswish(x) <= 0 here, relu kills it)
#   - 0 < x < 3       -> x*(x+3)/6
#   - x >= 3          -> x
# We still evaluate it as relu((x * relu6(x+3)) / 6) to match PyTorch's
# floating point operation order as closely as possible.
#
# Convolution strategy (direct, tiled, shared-memory):
#   - Each thread block produces a 32 (wide) x 8 (high) tile of outputs for
#     all output channels of one batch element.
#   - The corresponding 34x10 input halo (per input channel) is staged in
#     shared memory once and reused for every output channel.
#   - Weights are staged in shared memory transposed to [tap][out_channel]
#     so the 16 per-channel weights of a tap are contiguous (vectorized
#     float4 broadcast loads).
#   - Each thread accumulates 16 output channels per pass in registers.

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

#define TILE_W 32
#define TILE_H 8
#define HALO_W (TILE_W + 2)
#define HALO_H (TILE_H + 2)
#define THREADS 256
#define OC_CHUNK 16

__global__ __launch_bounds__(THREADS)
void conv3x3_hardswish_relu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int Cin, const int H, const int W,
    const int Cout, const int Ho, const int Wo)
{
    extern __shared__ float smem[];
    float* tile = smem;                             // Cin * HALO_H * HALO_W
    float* wbuf = tile + Cin * HALO_H * HALO_W;     // (Cin*9) * Cout, [tap][oc]
    float* bbuf = wbuf + (Cin * 9) * Cout;          // Cout

    const int tid = threadIdx.x;
    const int tx  = tid % TILE_W;   // column in tile (0..31)
    const int ty  = tid / TILE_W;   // row in tile (0..7)

    const int ow0 = blockIdx.x * TILE_W;
    const int oh0 = blockIdx.y * TILE_H;
    const int n   = blockIdx.z;

    // ---- stage input halo into shared memory ----
    const float* xn = x + (int64_t)n * Cin * H * W;
    const int tileElems = Cin * (HALO_H * HALO_W);
    for (int idx = tid; idx < tileElems; idx += THREADS) {
        const int ic = idx / (HALO_H * HALO_W);
        const int r  = idx - ic * (HALO_H * HALO_W);
        const int yy = r / HALO_W;
        const int xx = r - yy * HALO_W;
        const int gy = oh0 + yy;
        const int gx = ow0 + xx;
        float v = 0.0f;
        if (gy < H && gx < W) {
            v = xn[((int64_t)ic * H + gy) * W + gx];
        }
        tile[idx] = v;
    }

    // ---- stage weights (transposed to [tap][oc]) ----
    const int wElems = (Cin * 9) * Cout;
    for (int idx = tid; idx < wElems; idx += THREADS) {
        const int tap = idx / Cout;
        const int oc  = idx - tap * Cout;
        wbuf[idx] = w[(int64_t)oc * (Cin * 9) + tap];
    }

    // ---- stage bias ----
    for (int idx = tid; idx < Cout; idx += THREADS) {
        bbuf[idx] = bias[idx];
    }

    __syncthreads();

    const int oh = oh0 + ty;
    const int ow = ow0 + tx;
    const bool valid = (oh < Ho) && (ow < Wo);

    float* yn = y + (int64_t)n * Cout * Ho * Wo;

    for (int oc0 = 0; oc0 < Cout; oc0 += OC_CHUNK) {
        float acc[OC_CHUNK];
        #pragma unroll
        for (int o = 0; o < OC_CHUNK; ++o) acc[o] = 0.0f;

        for (int ic = 0; ic < Cin; ++ic) {
            const float* t  = tile + ic * (HALO_H * HALO_W);
            const float* wt = wbuf + (ic * 9) * Cout + oc0;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    const float iv = t[(ty + kh) * HALO_W + (tx + kw)];
                    const float4* wp =
                        reinterpret_cast<const float4*>(wt + (kh * 3 + kw) * Cout);
                    const float4 wa = wp[0];
                    const float4 wb = wp[1];
                    const float4 wc = wp[2];
                    const float4 wd = wp[3];
                    acc[0]  = fmaf(iv, wa.x, acc[0]);
                    acc[1]  = fmaf(iv, wa.y, acc[1]);
                    acc[2]  = fmaf(iv, wa.z, acc[2]);
                    acc[3]  = fmaf(iv, wa.w, acc[3]);
                    acc[4]  = fmaf(iv, wb.x, acc[4]);
                    acc[5]  = fmaf(iv, wb.y, acc[5]);
                    acc[6]  = fmaf(iv, wb.z, acc[6]);
                    acc[7]  = fmaf(iv, wb.w, acc[7]);
                    acc[8]  = fmaf(iv, wc.x, acc[8]);
                    acc[9]  = fmaf(iv, wc.y, acc[9]);
                    acc[10] = fmaf(iv, wc.z, acc[10]);
                    acc[11] = fmaf(iv, wc.w, acc[11]);
                    acc[12] = fmaf(iv, wd.x, acc[12]);
                    acc[13] = fmaf(iv, wd.y, acc[13]);
                    acc[14] = fmaf(iv, wd.z, acc[14]);
                    acc[15] = fmaf(iv, wd.w, acc[15]);
                }
            }
        }

        if (valid) {
            #pragma unroll
            for (int o = 0; o < OC_CHUNK; ++o) {
                const int oc = oc0 + o;
                if (oc < Cout) {
                    float v = acc[o] + bbuf[oc];
                    // relu(hardswish(v)) = max(0, (v * relu6(v+3)) / 6)
                    const float r6 = fminf(fmaxf(v + 3.0f, 0.0f), 6.0f);
                    const float hs = (v * r6) / 6.0f;
                    v = fmaxf(hs, 0.0f);
                    yn[((int64_t)oc * Ho + oh) * Wo + ow] = v;
                }
            }
        }
    }
}

torch::Tensor conv3x3_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    // Reference fallback for configurations the fused kernel does not cover.
    if (!x.is_cuda() || x.scalar_type() != at::kFloat ||
        w.scalar_type() != at::kFloat ||
        !bias.defined() || bias.scalar_type() != at::kFloat ||
        w.dim() != 4 || w.size(2) != 3 || w.size(3) != 3 ||
        w.size(1) != x.size(1) || (w.size(0) % 4) != 0 ||
        x.size(2) < 3 || x.size(3) < 3) {
        auto out = at::conv2d(x, w, bias);
        out = at::hardswish(out);
        out = at::relu(out);
        return out;
    }

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = bias.contiguous();

    const int N    = (int)xc.size(0);
    const int Cin  = (int)xc.size(1);
    const int H    = (int)xc.size(2);
    const int W    = (int)xc.size(3);
    const int Cout = (int)wc.size(0);
    const int Ho   = H - 2;
    const int Wo   = W - 2;

    auto y = at::empty({N, Cout, Ho, Wo}, xc.options());

    const size_t smem =
        ((size_t)Cin * (HALO_H * HALO_W) + (size_t)(Cin * 9) * Cout + Cout) * sizeof(float);

    if (smem > 48 * 1024) {
        cudaFuncSetAttribute(conv3x3_hardswish_relu_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem);
    }

    dim3 block(THREADS);
    dim3 grid((Wo + TILE_W - 1) / TILE_W, (Ho + TILE_H - 1) / TILE_H, N);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3x3_hardswish_relu_kernel<<<grid, block, smem, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
        y.data_ptr<float>(),
        N, Cin, H, W, Cout, Ho, Wo);

    return y;
}
"""

cpp_source = "torch::Tensor conv3x3_hardswish_relu(torch::Tensor x, torch::Tensor w, torch::Tensor bias);"

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


class ModelNew(nn.Module):
    """
    Optimized model: Conv2d + HardSwish + ReLU executed as a single fused
    CUDA kernel (one launch, no intermediate global-memory traffic).
    """
    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):
        if x.is_cuda and x.dtype == torch.float32:
            return fused_conv_act.conv3x3_hardswish_relu(
                x, self.conv.weight, self.conv.bias)
        # Generic fallback (e.g. CPU tensors).
        x = self.conv(x)
        x = F.hardswish(x)
        x = torch.relu(x)
        return x
```