```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.
#
# Direct convolution formulated as 9 tiny TF32 tensor-core GEMMs (one per
# kernel tap) using nvcuda::wmma 16x16x8 fragments. Each thread block computes
# an 8x32 tile of output pixels for all output channels:
#   - the input tile (10x34x8, zero-padded halo) is staged in dynamic shared
#     memory directly from NCHW global memory (channel pitch 12 keeps the
#     fragment loads conflict-light);
#   - weights are staged as [kh][kw][oc-tile][ic][16] with a padded pitch of
#     20 floats so weight-fragment loads are conflict-free;
#   - each warp owns one output row of 32 pixels (two 16-pixel column tiles)
#     and accumulates against all 16-channel tiles, so every loaded weight
#     fragment feeds 2 mmas and every input fragment feeds NT mmas
#     (0.75 shared-fragment loads per mma vs 1.25 for a 16-pixel tile);
#   - the bias + HardSwish + ReLU epilogue is applied through a padded
#     column-major shared stage and written straight to NCHW global memory
#     with fully coalesced 128-byte store transactions, one store per
#     (channel, 32-pixel row).
#
# The block needs ~57KB of shared memory, so it is placed in dynamic shared
# memory with an opt-in raised limit (sm_80 supports up to 163KB per block).
#
# Only the input (67MB) is read and the final output (520MB) is written, once.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <mma.h>

namespace wm = nvcuda::wmma;

#define TILE_W 32
#define TILE_H 8
#define CIN 8
#define PITCH 12            // input channel pitch (mult of 4 for tf32 wmma)
#define WPITCH 20           // padded weight-tile pitch (conflict-free)
#define SLD 36              // epilogue stage pitch (mult of 4, conflict-free)
#define MAX_NT 4            // supports up to 64 output channels
#define NTHREADS 256
#define SM_IN_FLOATS ((TILE_H + 2) * (TILE_W + 2) * PITCH)
#define STAGE_FLOATS (TILE_H * 16 * SLD)

__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) {
    extern __shared__ __align__(16) float dyn_smem[];
    float* sm_in = dyn_smem;
    float* sm_w  = dyn_smem + SM_IN_FLOATS;
    float* stage = dyn_smem + SM_IN_FLOATS + 9 * NT * CIN * WPITCH;

    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 WWt = TILE_W + 2;                    // 34
    const int HWt = (TILE_H + 2) * WWt;            // 340
    const int tile_elems = CIN * HWt;              // 2720
    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 / WWt;
        int ww = r - hh * WWt;
        int h = h0 + hh;
        int w = w0 + 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 * WWt + ww) * PITCH + c] = v;
    }

    // ---- stage weights into padded [kh][kw][oc16][ic][16] tiles ----
    // Iterate in the natural global order so global reads stay coalesced.
    const int w_elems = C_out * 72;
    for (int idx = tid; idx < w_elems; idx += NTHREADS) {
        int kw = idx % 3;
        int t  = idx / 3;
        int kh = t % 3;
        t /= 3;
        int ic = t & 7;
        int oc = t >> 3;
        sm_w[(((kh * 3 + kw) * NT + (oc >> 4)) * CIN + ic) * WPITCH + (oc & 15)] = w[idx];
    }
    __syncthreads();   // last block-wide barrier; warps may exit freely below

    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[2][NT];
    #pragma unroll
    for (int t2 = 0; t2 < 2; ++t2)
        #pragma unroll
        for (int j = 0; j < NT; ++j) wm::fill_fragment(acc[t2][j], 0.0f);

    // direct conv as 9 tensor-core GEMMs (M=32 pixels, N=16ch tiles, K=8)
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
            const float* a_ptr = &sm_in[((warp + kh) * WWt + kw) * PITCH];
            wm::fragment<wm::matrix_a, 16, 16, 8, wm::precision::tf32, wm::row_major> a0, a1;
            wm::load_matrix_sync(a0, a_ptr, PITCH);
            wm::load_matrix_sync(a1, a_ptr + 16 * 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) * NT + j) * CIN * WPITCH], WPITCH);
                wm::mma_sync(acc[0][j], a0, b, acc[0][j]);
                wm::mma_sync(acc[1][j], a1, b, acc[1][j]);
            }
        }
    }

    // ---- fused bias + HardSwish + ReLU epilogue, store NCHW ----
    const long HWo = (long)H_out * W_out;
    const long obase = ((long)n * C_out * H_out + h) * W_out + w0;
    float* stg = stage + warp * (16 * SLD);   // per-warp 32px x 16ch stage
    const int rem = W_out - w0;

    if (rem >= TILE_W) {
        // fast path: all 32 pixels of this warp's row are valid
        #pragma unroll
        for (int j = 0; j < NT; ++j) {
            wm::store_matrix_sync(stg,      acc[0][j], SLD, wm::mem_col_major);
            wm::store_matrix_sync(stg + 16, acc[1][j], SLD, wm::mem_col_major);
            __syncwarp();
            #pragma unroll
            for (int i = 0; i < 16; ++i) {
                const int c = j * 16 + i;
                const float v = bias_hs_relu(stg[i * SLD + lane], __ldg(bias + c));
                out[obase + (long)c * HWo + lane] = v;   // 128B coalesced
            }
            __syncwarp();
        }
    } else {
        // edge path: predicate every pixel
        #pragma unroll
        for (int j = 0; j < NT; ++j) {
            wm::store_matrix_sync(stg,      acc[0][j], SLD, wm::mem_col_major);
            wm::store_matrix_sync(stg + 16, acc[1][j], SLD, wm::mem_col_major);
            __syncwarp();
            #pragma unroll
            for (int i = 0; i < 16; ++i) {
                const int c = j * 16 + i;
                if (lane < rem) {
                    const float v = bias_hs_relu(stg[i * SLD + lane], __ldg(bias + c));
                    out[obase + (long)c * HWo + lane] = v;
                }
            }
            __syncwarp();
        }
    }
}

template <int NT>
static void launch_fused(dim3 grid, dim3 block,
                         const float* xp, const float* wp, const float* bp, float* op,
                         int N, int H, int W, int C_out, int H_out, int W_out) {
    constexpr int kSmemBytes =
        (SM_IN_FLOATS + 9 * NT * CIN * WPITCH + STAGE_FLOATS) * (int)sizeof(float);
    if constexpr (kSmemBytes > 48 * 1024) {
        static bool configured = [] {
            cudaError_t err = cudaFuncSetAttribute(
                reinterpret_cast<const void*>(conv3x3_bias_hs_relu_kernel<NT>),
                cudaFuncAttributeMaxDynamicSharedMemorySize,
                kSmemBytes);
            TORCH_CHECK(err == cudaSuccess,
                        "cudaFuncSetAttribute failed: ", cudaGetErrorString(err));
            return true;
        }();
        (void)configured;
    }
    conv3x3_bias_hs_relu_kernel<NT><<<grid, block, kSmemBytes>>>(
        xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
}

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:
            launch_fused<1>(grid, block, xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 2:
            launch_fused<2>(grid, block, xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 3:
            launch_fused<3>(grid, block, xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        case 4:
            launch_fused<4>(grid, block, xp, wp, bp, op, N, H, W, C_out, H_out, W_out);
            break;
        default:
            TORCH_CHECK(false, "unsupported out_channels");
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    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);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    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_v4",
    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). Each warp computes a full
        32-pixel output row against all output channels for a high
        mma-to-fragment-load ratio and fully coalesced 128B stores. The >48KB
        shared-memory footprint is enabled via dynamic shared memory opt-in.
      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 3 <= x.size(2) <= 8 * 65535
            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]
```