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

# Autotune conv algorithms (helps the generic fallback path) and allow TF32,
# matching the PyTorch reference configuration.
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.benchmark_limit = 0
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

# ---------------------------------------------------------------------------
# Fully fused Conv2d(3x3, stride 1, pad 0, Cin=8) + bias + HardSwish + ReLU.
#
# Strategy: direct convolution formulated as 9 tiny TF32 tensor-core GEMMs
# (one per kernel tap) using nvcuda::wmma 16x16x8 fragments. Each thread block
# computes an 8x16 tile of output pixels for all output channels:
#   - the input tile (10x18x8, zero-padded halo) is staged in shared memory
#     directly from NCHW global memory (no separate NHWC transform pass);
#   - weights are staged transposed as [kh][kw][ic][oc] in shared memory;
#   - each warp owns one output row of 16 pixels, accumulates 16x16 fragments
#     over the 9 taps and all 8 input channels (one K-step per tap);
#   - the bias + HardSwish + ReLU epilogue is applied in registers/shared and
#     written straight to NCHW global memory.
#
# This eliminates every intermediate memory pass of the reference pipeline
# (cudnn conv + bias add + hardswish + relu + layout transforms). Only the
# input (67MB) is read and the final output (520MB) is written, once.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <mma.h>

namespace wm = nvcuda::wmma;

#define TILE_W 16
#define TILE_H 8
#define CIN 8
#define PITCH 12            // padded channel pitch (mult of 4 for tf32 wmma)
#define MAX_NT 4            // supports up to 64 output channels
#define NTHREADS 256

__device__ __forceinline__ float bias_hs_relu(float x, float b) {
    x += b;
    float t = fminf(fmaxf(x + 3.0f, 0.0f), 6.0f);
    return fmaxf(x * t * 0.16666667163372039794921875f, 0.0f);
}

template <int NT>
__global__ void __launch_bounds__(NTHREADS) conv3x3_bias_hs_relu_kernel(
        const float* __restrict__ x,     // NCHW input
        const float* __restrict__ w,     // [C_out][8][3][3]
        const float* __restrict__ bias,  // [C_out]
        float* __restrict__ out,         // NCHW output
        int N, int H, int W, int C_out, int H_out, int W_out) {
    __shared__ float sm_in[(TILE_H + 2) * (TILE_W + 2) * PITCH];
    __shared__ float sm_w[3 * 3 * CIN * MAX_NT * 16];
    __shared__ float stage[TILE_H][16 * 16];

    const int n  = blockIdx.z;
    const int h0 = blockIdx.y * TILE_H;
    const int w0 = blockIdx.x * TILE_W;
    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;

    // ---- stage input tile [hh][ww][c] with zero-padded halo ----
    const int HWt = (TILE_H + 2) * (TILE_W + 2);   // 180
    const int tile_elems = CIN * HWt;              // 1440
    const long xbase = (long)n * CIN * H * W;
    for (int idx = tid; idx < tile_elems; idx += NTHREADS) {
        int c  = idx / HWt;
        int r  = idx - c * HWt;
        int hh = r / (TILE_W + 2);
        int ww = r - hh * (TILE_W + 2);
        int h = h0 - 1 + hh;
        int w = w0 - 1 + ww;
        float v = 0.0f;
        if ((unsigned)h < (unsigned)H && (unsigned)w < (unsigned)W) {
            v = x[xbase + ((long)c * H + h) * W + w];
        }
        sm_in[(hh * (TILE_W + 2) + ww) * PITCH + c] = v;
    }

    // ---- stage weights transposed to [kh][kw][ic][oc] ----
    const int w_elems = 9 * CIN * C_out;
    for (int idx = tid; idx < w_elems; idx += NTHREADS) {
        int oc = idx % C_out;
        int t  = idx / C_out;
        int ic = t % CIN;
        t /= CIN;
        int kw = t % 3;
        int kh = t / 3;
        sm_w[idx] = w[((long)(oc * CIN + ic) * 3 + kh) * 3 + kw];
    }
    __syncthreads();

    const int h = h0 + warp;   // output row owned by this warp
    if (h >= H_out) return;

    wm::fragment<wm::accumulator, 16, 16, 8, float> acc[NT];
    #pragma unroll
    for (int j = 0; j < NT; ++j) wm::fill_fragment(acc[j], 0.0f);

    // direct conv as 9 tensor-core GEMMs (M=16 pixels, N=16ch tiles, K=8)
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
            wm::fragment<wm::matrix_a, 16, 16, 8, wm::precision::tf32, wm::row_major> a;
            wm::load_matrix_sync(a, &sm_in[((warp + kh) * (TILE_W + 2) + kw) * PITCH], PITCH);
            #pragma unroll
            for (int j = 0; j < NT; ++j) {
                wm::fragment<wm::matrix_b, 16, 16, 8, wm::precision::tf32, wm::row_major> b;
                wm::load_matrix_sync(b, &sm_w[((kh * 3 + kw) * CIN) * C_out + j * 16], C_out);
                wm::mma_sync(acc[j], a, b, acc[j]);
            }
        }
    }

    // ---- fused bias + HardSwish + ReLU epilogue, store NCHW ----
    const long HWo = (long)H_out * W_out;
    const long obase = (long)n * C_out * HWo + (long)h * W_out + w0;
    const int r  = lane & 15;    // pixel column within the 16-wide row
    const int ch = lane >> 4;    // 0 or 1 -> two channels per pass
    const bool valid = (w0 + r) < W_out;

    #pragma unroll
    for (int j = 0; j < NT; ++j) {
        wm::store_matrix_sync(stage[warp], acc[j], 16, wm::mem_row_major);
        __syncwarp();
        #pragma unroll
        for (int cc = 0; cc < 8; ++cc) {
            int cl = cc * 2 + ch;
            int c  = j * 16 + cl;
            float v = stage[warp][r * 16 + cl];
            v = bias_hs_relu(v, __ldg(bias + c));
            if (valid) out[obase + (long)c * HWo + r] = v;
        }
        __syncwarp();
    }
}

torch::Tensor conv3x3_bias_hs_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "float32 only");
    TORCH_CHECK(x.dim() == 4 && w.dim() == 4, "bad dims");

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

    const int N     = (int)xc.size(0);
    const int C_in  = (int)xc.size(1);
    const int H     = (int)xc.size(2);
    const int W     = (int)xc.size(3);
    const int C_out = (int)wc.size(0);
    const int H_out = H - 2;
    const int W_out = W - 2;

    TORCH_CHECK(C_in == CIN, "fused kernel needs in_channels == 8");
    TORCH_CHECK((C_out % 16) == 0 && C_out >= 16 && C_out <= MAX_NT * 16,
                "fused kernel needs out_channels in {16,32,48,64}");
    TORCH_CHECK(H_out > 0 && W_out > 0, "bad spatial size");

    auto out = torch::empty({N, C_out, H_out, W_out}, xc.options());
    if (out.numel() == 0) return out;

    dim3 grid((unsigned)((W_out + TILE_W - 1) / TILE_W),
              (unsigned)((H_out + TILE_H - 1) / TILE_H),
              (unsigned)N);
    dim3 block(NTHREADS);

    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    switch (C_out / 16) {
        case 1:
            conv3x3_bias_hs_relu_kernel<1><<<grid, block>>>(xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 2:
            conv3x3_bias_hs_relu_kernel<2><<<grid, block>>>(xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 3:
            conv3x3_bias_hs_relu_kernel<3><<<grid, block>>>(xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 4:
            conv3x3_bias_hs_relu_kernel<4><<<grid, block>>>(xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        default:
            TORCH_CHECK(false, "unsupported out_channels");
    }
    return out;
}

// ---------------------------------------------------------------------------
// Generic fallback epilogue: bias + HardSwish + ReLU in a single memory pass
// (used when the fully fused conv kernel's shape constraints are not met).
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(256) bias_hs_relu_vec_kernel(
        const float4* __restrict__ in,
        float4* __restrict__ out,
        const float* __restrict__ bias,
        int hw4, int C) {
    int plane = blockIdx.y;
    float b = __ldg(bias + (plane % C));
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < hw4) {
        long off = (long)plane * hw4 + idx;
        float4 v = in[off];
        v.x = bias_hs_relu(v.x, b);
        v.y = bias_hs_relu(v.y, b);
        v.z = bias_hs_relu(v.z, b);
        v.w = bias_hs_relu(v.w, b);
        out[off] = v;
    }
}

__global__ void __launch_bounds__(256) bias_hs_relu_scalar_kernel(
        const float* __restrict__ in,
        float* __restrict__ out,
        const float* __restrict__ bias,
        int hw, int C) {
    int plane = blockIdx.y;
    float b = __ldg(bias + (plane % C));
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < hw) {
        long off = (long)plane * hw + idx;
        out[off] = bias_hs_relu(in[off], b);
    }
}

torch::Tensor bias_hardswish_relu_cuda(torch::Tensor x, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "input must be float32");
    TORCH_CHECK(x.dim() == 4, "expected NCHW 4D input");

    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    if (xc.numel() == 0) return out;

    const long C  = xc.size(1);
    const long HW = xc.size(2) * xc.size(3);
    const long planes = xc.size(0) * C;
    TORCH_CHECK(planes <= 65535, "too many planes for grid.y");

    const int block = 256;
    const float* in_ptr = xc.data_ptr<float>();
    float* out_ptr      = out.data_ptr<float>();
    const float* b_ptr  = bias.data_ptr<float>();

    dim3 grid;
    grid.y = (unsigned)planes;
    if ((HW & 3) == 0) {
        int hw4 = (int)(HW >> 2);
        grid.x = (unsigned)((hw4 + block - 1) / block);
        bias_hs_relu_vec_kernel<<<grid, block>>>(
            reinterpret_cast<const float4*>(in_ptr),
            reinterpret_cast<float4*>(out_ptr),
            b_ptr, hw4, (int)C);
    } else {
        grid.x = (unsigned)((HW + block - 1) / block);
        bias_hs_relu_scalar_kernel<<<grid, block>>>(
            in_ptr, out_ptr, b_ptr, (int)HW, (int)C);
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor conv3x3_bias_hs_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b);\n"
    "torch::Tensor bias_hardswish_relu_cuda(torch::Tensor x, torch::Tensor bias);"
)

fused_ops = load_inline(
    name="fused_conv3x3_bias_hardswish_relu",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["conv3x3_bias_hs_relu_cuda", "bias_hardswish_relu_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      Fast path (Cin=8, 3x3 stride-1 pad-0 conv, Cout in {16,32,48,64}):
        a single fully fused CUDA kernel performing the convolution with TF32
        tensor cores plus the bias + HardSwish + ReLU epilogue, reading the
        input and writing the final output exactly once (no intermediate
        activations, no layout-transform passes).
      Fallback path: cudnn conv (no bias) + fused epilogue 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):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height, width).
        """
        c = self.conv
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and c.bias is not None
            and c.groups == 1
            and c.stride == (1, 1)
            and c.padding == (0, 0)
            and c.dilation == (1, 1)
            and c.kernel_size == (3, 3)
            and c.in_channels == 8
            and (c.out_channels % 16) == 0
            and 16 <= c.out_channels <= 64
            and x.size(0) <= 65535
            and x.size(2) >= 3
            and x.size(3) >= 3
        ):
            return fused_ops.conv3x3_bias_hs_relu_cuda(x, c.weight, c.bias)

        y = F.conv2d(x, c.weight, None, c.stride, c.padding, c.dilation, c.groups)
        if y.is_cuda and y.dtype == torch.float32 and c.bias is not None:
            return fused_ops.bias_hardswish_relu_cuda(y, c.bias)
        if c.bias is not None:
            y = y + c.bias.view(1, -1, 1, 1)
        return F.relu(F.hardswish(y))


batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128
kernel_size = 3


def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size]
```