```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 kernel: GroupNorm + Swish + Multiply + Swish in a single pass.
# The GEMM is left to cuBLAS (nn.Linear) since it dominates FLOPs and is
# already near-optimal. Everything after the GEMM is memory-bound, so fusing
# it into one kernel avoids 4+ extra global memory round-trips.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

// Specialized kernel for the common case group_size == 32 (one warp per group,
// each lane holds exactly one element in registers).
__global__ void fused_gn_swish_mul_swish_gs32_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int channels,
    int num_groups,
    int n_warps,
    float eps)
{
    int warp_id = blockIdx.x * blockDim.y + threadIdx.y;
    if (warp_id >= n_warps) return;
    int sample = warp_id / num_groups;
    int group  = warp_id % num_groups;
    int lane   = threadIdx.x;
    int c      = group * 32 + lane;
    int64_t idx = (int64_t)sample * channels + c;

    float v = x[idx];

    // mean
    float sum = v;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_xor_sync(0xffffffffu, sum, offset);
    float mean = sum * (1.0f / 32.0f);

    // biased variance (two-pass, matches GroupNorm)
    float d = v - mean;
    float sq = d * d;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sq += __shfl_xor_sync(0xffffffffu, sq, offset);
    float inv_std = rsqrtf(sq * (1.0f / 32.0f) + eps);

    // group norm affine
    float y = fmaf((v - mean) * inv_std, gamma[c], beta[c]);
    // swish 1
    float s = 1.0f / (1.0f + expf(-y));
    y = y * s;
    // multiply weight
    y = y * mul_w[c];
    // swish 2
    s = 1.0f / (1.0f + expf(-y));
    y = y * s;

    out[idx] = y;
}

// Generic fallback kernel: one warp per group, loop over group elements.
__global__ void fused_gn_swish_mul_swish_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int channels,
    int num_groups,
    int group_size,
    int n_warps,
    float eps)
{
    int warp_id = blockIdx.x * blockDim.y + threadIdx.y;
    if (warp_id >= n_warps) return;
    int sample = warp_id / num_groups;
    int group  = warp_id % num_groups;
    int lane   = threadIdx.x;
    int64_t base = (int64_t)sample * channels + group * group_size;

    float sum = 0.0f;
    for (int i = lane; i < group_size; i += 32) sum += x[base + i];
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_xor_sync(0xffffffffu, sum, offset);
    float mean = sum / (float)group_size;

    float sq = 0.0f;
    for (int i = lane; i < group_size; i += 32) {
        float d = x[base + i] - mean;
        sq += d * d;
    }
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sq += __shfl_xor_sync(0xffffffffu, sq, offset);
    float inv_std = rsqrtf(sq / (float)group_size + eps);

    for (int i = lane; i < group_size; i += 32) {
        int c = group * group_size + i;
        float y = fmaf((x[base + i] - mean) * inv_std, gamma[c], beta[c]);
        float s = 1.0f / (1.0f + expf(-y));
        y = y * s;
        y = y * mul_w[c];
        s = 1.0f / (1.0f + expf(-y));
        y = y * s;
        out[base + i] = y;
    }
}

torch::Tensor fused_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 CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    auto x_c     = x.contiguous();
    auto gamma_c = gamma.contiguous();
    auto beta_c  = beta.contiguous();
    auto mulw_c  = mul_w.contiguous();
    auto out = torch::empty_like(x_c);

    int64_t batch      = x_c.size(0);
    int64_t channels   = x_c.size(1);
    TORCH_CHECK(channels % num_groups == 0, "channels must be divisible by num_groups");
    int64_t group_size = channels / num_groups;
    int64_t n_warps    = batch * num_groups;

    dim3 block(32, 8); // 8 warps per block, one warp per (sample, group)
    int64_t grid = (n_warps + block.y - 1) / block.y;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (group_size == 32) {
        fused_gn_swish_mul_swish_gs32_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            x_c.data_ptr<float>(), gamma_c.data_ptr<float>(), beta_c.data_ptr<float>(),
            mulw_c.data_ptr<float>(), out.data_ptr<float>(),
            (int)channels, (int)num_groups, (int)n_warps, (float)eps);
    } else {
        fused_gn_swish_mul_swish_generic_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            x_c.data_ptr<float>(), gamma_c.data_ptr<float>(), beta_c.data_ptr<float>(),
            mulw_c.data_ptr<float>(), out.data_ptr<float>(),
            (int)channels, (int)num_groups, (int)group_size, (int)n_warps, (float)eps);
    }
    return out;
}
"""

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

fused_ops = load_inline(
    name="fused_gn_swish_mul_swish",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_gn_swish_mul_swish"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM via cuBLAS (nn.Linear), then a single fused CUDA
    kernel performing GroupNorm + Swish + Multiply + Swish in one pass.
    """
    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 = 1e-5

    def forward(self, x):
        x = self.gemm(x)
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            return fused_ops.fused_gn_swish_mul_swish(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )
        # Fallback path (non-CUDA / non-fp32 inputs)
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```