```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
from torch.utils.cpp_extension import load_inline

# Fused Conv2d(3x3, stride=1, pad=0) + HardSwish + ReLU CUDA kernel.
# Note: ReLU(HardSwish(v)) == max(0, v * clamp(v + 3, 0, 6) / 6), which lets us
# fuse both activations into a single elementwise epilogue inside the conv kernel.
_cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>

#define TILE 16                 // spatial output tile per block (TILE x TILE)
#define OC_TILE 64              // output channels handled per block
#define OC_GROUP 16             // output channels per thread group
#define OC_CHUNK 4              // output channels per register-tile step
#define KF 3                    // fast-path kernel size
#define TS (TILE + KF - 1)      // shared input tile edge (with halo)

__device__ __forceinline__ float hs_relu(float v) {
    // hardswish: v * min(max(v + 3, 0), 6) / 6, then relu: max(., 0)
    float t = fminf(fmaxf(v + 3.0f, 0.0f), 6.0f);
    float h = v * t / 6.0f;
    return fmaxf(h, 0.0f);
}

// Fast path: 3x3 convolution, stride 1, padding 0, fused hardswish+relu.
// Block: 256 threads = (8 sx) * (8 sy) * (4 oc groups).
// Each thread owns a 2x2 spatial patch and iterates over its 16 output
// channels in chunks of 4, keeping 16 accumulators in registers.
__global__ void conv3x3_hs_relu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    int N, int C_in, int H, int W,
    int C_out, int H_out, int W_out)
{
    extern __shared__ float smem[];
    float* s_in = smem;                     // C_in * TS * TS
    float* s_w  = smem + C_in * TS * TS;    // OC_TILE * C_in * 9

    const int tid = threadIdx.x;
    const int n   = blockIdx.z % N;
    const int oct = blockIdx.z / N;
    const int oh0 = blockIdx.y * TILE;
    const int ow0 = blockIdx.x * TILE;
    const int oc0 = oct * OC_TILE;

    // ---- stage input tile (with halo) into shared memory ----
    {
        const float* xn = x + (long)n * C_in * H * W;
        const int total = C_in * TS * TS;
        for (int idx = tid; idx < total; idx += 256) {
            int c = idx / (TS * TS);
            int rem = idx - c * TS * TS;
            int r = rem / TS;
            int q = rem - r * TS;
            int ih = oh0 + r;
            int iw = ow0 + q;
            float v = 0.0f;
            if (ih < H && iw < W)
                v = xn[((long)c * H + ih) * W + iw];
            s_in[idx] = v;
        }
    }
    // ---- stage this block's weight slice into shared memory ----
    {
        const int ckk = C_in * KF * KF;
        const int total = OC_TILE * ckk;
        const int oc_lim = min(OC_TILE, C_out - oc0);
        for (int idx = tid; idx < total; idx += 256) {
            int oc = idx / ckk;
            float v = 0.0f;
            if (oc < oc_lim)
                v = w[(long)(oc0 + oc) * ckk + (idx - oc * ckk)];
            s_w[idx] = v;
        }
    }
    __syncthreads();

    const int sx = tid & 7;
    const int sy = (tid >> 3) & 7;
    const int og = tid >> 6;

    const int ty = sy * 2;
    const int tx = sx * 2;
    const int oh = oh0 + ty;
    const int ow = ow0 + tx;
    const int oc_base = oc0 + og * OC_GROUP;
    const int ckk = C_in * KF * KF;

    #pragma unroll 1
    for (int ocg = 0; ocg < OC_GROUP; ocg += OC_CHUNK) {
        const int oc = oc_base + ocg;
        float acc[OC_CHUNK][2][2];
        #pragma unroll
        for (int j = 0; j < OC_CHUNK; ++j) {
            float bv = (b != nullptr && (oc + j) < C_out) ? b[oc + j] : 0.0f;
            acc[j][0][0] = bv;
            acc[j][0][1] = bv;
            acc[j][1][0] = bv;
            acc[j][1][1] = bv;
        }
        const float* wp0 = s_w + (og * OC_GROUP + ocg) * ckk;
        #pragma unroll 1
        for (int c = 0; c < C_in; ++c) {
            const float* in_c = s_in + c * TS * TS + ty * TS + tx;
            const float* wp = wp0 + c * KF * KF;
            #pragma unroll
            for (int ky = 0; ky < KF; ++ky) {
                #pragma unroll
                for (int kx = 0; kx < KF; ++kx) {
                    const int wi = ky * KF + kx;
                    float w0 = wp[wi];
                    float w1 = wp[wi + ckk];
                    float w2 = wp[wi + 2 * ckk];
                    float w3 = wp[wi + 3 * ckk];
                    float i00 = in_c[ky * TS + kx];
                    float i01 = in_c[ky * TS + kx + 1];
                    float i10 = in_c[(ky + 1) * TS + kx];
                    float i11 = in_c[(ky + 1) * TS + kx + 1];
                    acc[0][0][0] = fmaf(w0, i00, acc[0][0][0]);
                    acc[0][0][1] = fmaf(w0, i01, acc[0][0][1]);
                    acc[0][1][0] = fmaf(w0, i10, acc[0][1][0]);
                    acc[0][1][1] = fmaf(w0, i11, acc[0][1][1]);
                    acc[1][0][0] = fmaf(w1, i00, acc[1][0][0]);
                    acc[1][0][1] = fmaf(w1, i01, acc[1][0][1]);
                    acc[1][1][0] = fmaf(w1, i10, acc[1][1][0]);
                    acc[1][1][1] = fmaf(w1, i11, acc[1][1][1]);
                    acc[2][0][0] = fmaf(w2, i00, acc[2][0][0]);
                    acc[2][0][1] = fmaf(w2, i01, acc[2][0][1]);
                    acc[2][1][0] = fmaf(w2, i10, acc[2][1][0]);
                    acc[2][1][1] = fmaf(w2, i11, acc[2][1][1]);
                    acc[3][0][0] = fmaf(w3, i00, acc[3][0][0]);
                    acc[3][0][1] = fmaf(w3, i01, acc[3][0][1]);
                    acc[3][1][0] = fmaf(w3, i10, acc[3][1][0]);
                    acc[3][1][1] = fmaf(w3, i11, acc[3][1][1]);
                }
            }
        }
        // fused hardswish + relu epilogue, guarded stores
        #pragma unroll
        for (int j = 0; j < OC_CHUNK; ++j) {
            if ((oc + j) < C_out) {
                float* yp = y + ((long)n * C_out + oc + j) * H_out * W_out;
                if (oh < H_out) {
                    if (ow < W_out)     yp[(long)oh * W_out + ow]         = hs_relu(acc[j][0][0]);
                    if (ow + 1 < W_out) yp[(long)oh * W_out + ow + 1]     = hs_relu(acc[j][0][1]);
                }
                if (oh + 1 < H_out) {
                    if (ow < W_out)     yp[(long)(oh + 1) * W_out + ow]     = hs_relu(acc[j][1][0]);
                    if (ow + 1 < W_out) yp[(long)(oh + 1) * W_out + ow + 1] = hs_relu(acc[j][1][1]);
                }
            }
        }
    }
}

// Generic fallback: any kernel size / channel count, stride 1, padding 0.
__global__ void conv_hs_relu_naive_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    int N, int C_in, int H, int W, int C_out, int KH, int KW,
    int H_out, int W_out, long total)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % W_out;
    long t = idx / W_out;
    int oh = t % H_out;
    t /= H_out;
    int oc = t % C_out;
    int n  = t / C_out;
    float acc = (b != nullptr) ? b[oc] : 0.0f;
    for (int c = 0; c < C_in; ++c) {
        for (int ky = 0; ky < KH; ++ky) {
            int ih = oh + ky;
            for (int kx = 0; kx < KW; ++kx) {
                int iw = ow + kx;
                acc = fmaf(x[(((long)n * C_in + c) * H + ih) * W + iw],
                           w[(((long)oc * C_in + c) * KH + ky) * KW + kx], acc);
            }
        }
    }
    y[idx] = hs_relu(acc);
}

torch::Tensor conv_hswish_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    torch::Tensor bc;
    const float* bptr = nullptr;
    if (b.defined() && b.numel() > 0) {
        bc = b.contiguous();
        bptr = bc.data_ptr<float>();
    }

    const int N = xc.size(0), C_in = xc.size(1), H = xc.size(2), W = xc.size(3);
    const int C_out = wc.size(0), KH = wc.size(2), KW = wc.size(3);
    const int H_out = H - KH + 1;
    const int W_out = W - KW + 1;
    TORCH_CHECK(H_out > 0 && W_out > 0, "kernel size larger than input");

    auto y = at::empty({N, C_out, H_out, W_out}, xc.options());
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* yp = y.data_ptr<float>();

    bool fast = (KH == KF && KW == KF);
    size_t smem = 0;
    if (fast) {
        smem = ((size_t)C_in * TS * TS + (size_t)OC_TILE * C_in * KF * KF) * sizeof(float);
        if (smem > 160 * 1024) fast = false;
    }

    if (fast) {
        const int oc_tiles = (C_out + OC_TILE - 1) / OC_TILE;
        dim3 grid((W_out + TILE - 1) / TILE, (H_out + TILE - 1) / TILE, N * oc_tiles);
        if (smem > 48 * 1024) {
            cudaFuncSetAttribute(conv3x3_hs_relu_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 (int)smem);
        }
        conv3x3_hs_relu_kernel<<<grid, 256, smem>>>(
            xp, wp, bptr, yp, N, C_in, H, W, C_out, H_out, W_out);
    } else {
        const long total = (long)N * C_out * H_out * W_out;
        const int threads = 256;
        const long blocks = (total + threads - 1) / threads;
        conv_hs_relu_naive_kernel<<<(unsigned int)blocks, threads>>>(
            xp, wp, bptr, yp, N, C_in, H, W, C_out, KH, KW, H_out, W_out, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_cpp_src = "torch::Tensor conv_hswish_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b);"

_fused_conv = load_inline(
    name="conv_hswish_relu_fused",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv_hswish_relu_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Model that performs a convolution, applies HardSwish, and then ReLU,
    using a single fused custom CUDA kernel.
    """
    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):
        weight = self.conv.weight
        bias = self.conv.bias
        if bias is None:
            bias = x.new_zeros(weight.size(0))
        return _fused_conv.conv_hswish_relu_cuda(x, weight, bias)
```