```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-multiply -> Swish
# The GEMM is left to cuBLAS (nn.Linear); everything after it is fused into
# a single memory-bound kernel using one warp per (sample, group).
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define FULL_MASK 0xffffffffu

// Specialization for group_size == 32 (exactly one element per lane).
__global__ void fused_gn_swish_mul_swish_gs32_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    const int batch_size,
    const int num_groups,
    const int out_features,
    const float eps)
{
    const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
    const int lane    = threadIdx.x & 31;
    if (warp_id >= batch_size * num_groups) return;

    const int b = warp_id / num_groups;
    const int g = warp_id % num_groups;
    const int c = g * 32 + lane;

    const long idx = (long)b * out_features + c;
    const float v  = x[idx];

    // group mean (32 lanes, butterfly reduction -> all lanes get the sum)
    float sum = v;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(FULL_MASK, sum, o);
    const float mean = sum * (1.0f / 32.0f);

    // group variance (biased, as in nn.GroupNorm)
    const float d = v - mean;
    float vs = d * d;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) vs += __shfl_xor_sync(FULL_MASK, vs, o);
    const float rstd = rsqrtf(vs * (1.0f / 32.0f) + eps);

    // normalize + affine
    float y = fmaf(d * rstd, gn_w[c], gn_b[c]);
    // swish: y * sigmoid(y)
    y = y / (1.0f + expf(-y));
    // per-channel multiply
    y = y * mul_w[c];
    // swish again
    y = y / (1.0f + expf(-y));

    out[idx] = y;
}

// General fallback for arbitrary group sizes (warp per (sample, group)).
__global__ void fused_gn_swish_mul_swish_general_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    const int batch_size,
    const int num_groups,
    const int group_size,
    const int out_features,
    const float eps)
{
    const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
    const int lane    = threadIdx.x & 31;
    if (warp_id >= batch_size * num_groups) return;

    const int b = warp_id / num_groups;
    const int g = warp_id % num_groups;
    const long base = (long)b * out_features + (long)g * group_size;

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

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

    for (int i = lane; i < group_size; i += 32) {
        const int c = g * group_size + i;
        float y = fmaf((x[base + i] - mean) * rstd, gn_w[c], gn_b[c]);
        y = y / (1.0f + expf(-y));
        y = y * mul_w[c];
        y = y / (1.0f + expf(-y));
        out[(long)b * out_features + c] = y;
    }
}

torch::Tensor fused_gn_swish_mul_swish_cuda(
    torch::Tensor x,
    torch::Tensor gn_w,
    torch::Tensor gn_b,
    torch::Tensor mul_w,
    int64_t num_groups,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");

    auto x_c     = x.contiguous();
    auto gn_w_c  = gn_w.contiguous();
    auto gn_b_c  = gn_b.contiguous();
    auto mul_w_c = mul_w.contiguous();

    const int batch_size   = (int)x_c.size(0);
    const int out_features = (int)x_c.size(1);
    const int group_size   = out_features / (int)num_groups;
    TORCH_CHECK(out_features % (int)num_groups == 0, "invalid num_groups");

    auto out = torch::empty_like(x_c);

    const long total_warps = (long)batch_size * (long)num_groups;
    const int threads = 256;
    const long blocks = (total_warps * 32 + threads - 1) / threads;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (group_size == 32) {
        fused_gn_swish_mul_swish_gs32_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            x_c.data_ptr<float>(), gn_w_c.data_ptr<float>(), gn_b_c.data_ptr<float>(),
            mul_w_c.data_ptr<float>(), out.data_ptr<float>(),
            batch_size, (int)num_groups, out_features, (float)eps);
    } else {
        fused_gn_swish_mul_swish_general_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            x_c.data_ptr<float>(), gn_w_c.data_ptr<float>(), gn_b_c.data_ptr<float>(),
            mul_w_c.data_ptr<float>(), out.data_ptr<float>(),
            batch_size, (int)num_groups, group_size, out_features, (float)eps);
    }

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_gn_swish_mul_swish_cuda("
    "torch::Tensor x, torch::Tensor gn_w, torch::Tensor gn_b, torch::Tensor mul_w, "
    "int64_t num_groups, double eps);"
)

fused_gn_act = load_inline(
    name="fused_gn_swish_mul_swish",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_gn_swish_mul_swish_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM followed by a single fused CUDA kernel that
    performs GroupNorm + Swish + per-channel 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 = self.group_norm.eps

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features)
        x = self.gemm(x)

        if (x.is_cuda and x.dtype == torch.float32 and x.dim() == 2
                and self.group_norm.affine
                and x.size(1) % self.num_groups == 0):
            # Fused: GroupNorm -> Swish -> Multiply -> Swish
            return fused_gn_act.fused_gn_swish_mul_swish_cuda(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )

        # Fallback to the original PyTorch ops for unsupported configurations.
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```