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

# Fused CUDA kernel: GroupNorm + Swish + channel-wise Multiply + Swish
# Each warp processes one (batch_row, group) pair, computing the group
# statistics with warp shuffle reductions and applying the full fused
# element-wise chain in registers.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float warp_allreduce_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffffu, v, offset);
    }
    return v;
}

__global__ void gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int N, int C, int num_groups, int group_size, float eps)
{
    int warps_per_block = blockDim.x >> 5;
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int group_idx = blockIdx.x * warps_per_block + warp;
    if (group_idx >= N * num_groups) return;

    int n = group_idx / num_groups;
    int g = group_idx - n * num_groups;

    const long row_off = (long)n * C + (long)g * group_size;
    const long ch_off  = (long)g * group_size;
    const float* xg = x + row_off;
    float* og = out + row_off;
    const float* gg = gamma + ch_off;
    const float* bg = beta + ch_off;
    const float* wg = mul_w + ch_off;

    // GroupNorm mean
    float sum = 0.0f;
    for (int i = lane; i < group_size; i += 32) sum += xg[i];
    sum = warp_allreduce_sum(sum);
    sum = __shfl_sync(0xffffffffu, sum, 0);
    float mean = sum / (float)group_size;

    // GroupNorm biased variance
    float dsum = 0.0f;
    for (int i = lane; i < group_size; i += 32) {
        float d = xg[i] - mean;
        dsum += d * d;
    }
    dsum = warp_allreduce_sum(dsum);
    dsum = __shfl_sync(0xffffffffu, dsum, 0);
    float rstd = rsqrtf(dsum / (float)group_size + eps);

    // Fused: normalize -> affine -> swish -> multiply -> swish
    for (int i = lane; i < group_size; i += 32) {
        float v = xg[i];
        float h = (v - mean) * rstd * gg[i] + bg[i];
        float s1 = h / (1.0f + expf(-h));   // swish
        float m = s1 * wg[i];               // channel-wise multiply
        float s2 = m / (1.0f + expf(-m));   // swish
        og[i] = s2;
    }
}

torch::Tensor gn_swish_mul_swish(
    torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor mul_w, int64_t num_groups, double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto wc = mul_w.contiguous();

    int64_t N = xc.size(0);
    int64_t C = xc.size(1);
    auto out = torch::empty_like(xc);

    int group_size = (int)(C / num_groups);
    int total_groups = (int)(N * num_groups);
    int threads = 256;
    int warps_per_block = threads / 32;
    int blocks = (total_groups + warps_per_block - 1) / warps_per_block;

    auto stream = at::cuda::getCurrentCUDAStream();
    gn_swish_mul_swish_kernel<<<blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(),
        wc.data_ptr<float>(), out.data_ptr<float>(),
        (int)N, (int)C, (int)num_groups, group_size, (float)eps);
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor gn_swish_mul_swish(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor mul_w, int64_t num_groups, double eps);"
)

fused_op = load_inline(
    name="fused_gn_swish_mul_swish",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["gn_swish_mul_swish"],
    verbose=True,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM is kept as the cuBLAS-backed nn.Linear, while
    GroupNorm + Swish + Multiply + Swish are fused into a single custom
    CUDA kernel (one warp per (row, group), shuffle-based statistics).
    """
    def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
        self.num_groups = num_groups
        self.eps = self.group_norm.eps
        self.fused_op = fused_op

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features)
        x = self.gemm(x)
        # fused GroupNorm + Swish + Multiply + Swish
        x = self.fused_op.gn_swish_mul_swish(
            x,
            self.group_norm.weight,
            self.group_norm.bias,
            self.multiply_weight,
            self.num_groups,
            self.eps,
        )
        return x
```