```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: bias-add + GroupNorm + Swish + Multiply + Swish
# The GEMM itself is left to cuBLAS (torch.matmul without bias); everything
# after the GEMM is fused into a single kernel where each warp processes one
# (batch-row, group) pair using warp-shuffle reductions.
fused_cpp_source = """
#include <torch/extension.h>
torch::Tensor fused_bias_gn_swish_mul_swish(torch::Tensor x, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mul_w, int64_t num_groups, double eps);
"""

fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define WARP_SIZE 32

__global__ void fused_bias_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,      // GEMM output (no bias yet), (batch, channels)
    const float* __restrict__ bias,   // (channels)
    const float* __restrict__ gamma,  // (channels)
    const float* __restrict__ beta,   // (channels)
    const float* __restrict__ mul_w,  // (channels)
    float* __restrict__ out,          // (batch, channels)
    const int batch_size,
    const int channels,
    const int num_groups,
    const float eps)
{
    const int group_size = channels / num_groups;
    const int lane = threadIdx.x & (WARP_SIZE - 1);
    const int warps_per_block = blockDim.x >> 5;
    const long global_warp = (long)blockIdx.x * warps_per_block + (threadIdx.x >> 5);
    const long total_warps = (long)batch_size * num_groups;
    if (global_warp >= total_warps) return;

    const int b = (int)(global_warp / num_groups);
    const int g = (int)(global_warp % num_groups);
    const int cbase = g * group_size;

    const float* xrow = x + (long)b * channels + cbase;
    float* orow = out + (long)b * channels + cbase;

    // Fast path: one element per lane, keep it in a register across passes.
    if (group_size == WARP_SIZE) {
        const int c = cbase + lane;
        const float v = xrow[lane] + bias[c];

        float sum = v;
        #pragma unroll
        for (int off = WARP_SIZE / 2; off > 0; off >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        }
        const float mean = __shfl_sync(0xffffffffu, sum, 0) * (1.0f / (float)WARP_SIZE);

        const float d = v - mean;
        float vsum = d * d;
        #pragma unroll
        for (int off = WARP_SIZE / 2; off > 0; off >>= 1) {
            vsum += __shfl_down_sync(0xffffffffu, vsum, off);
        }
        const float var = __shfl_sync(0xffffffffu, vsum, 0) * (1.0f / (float)WARP_SIZE);
        const float rstd = rsqrtf(var + eps);

        float y = d * rstd * gamma[c] + beta[c];
        y = y / (1.0f + expf(-y));   // swish: y * sigmoid(y)
        y = y * mul_w[c];
        y = y / (1.0f + expf(-y));   // swish again
        orow[lane] = y;
        return;
    }

    // Generic path for arbitrary group sizes.
    float sum = 0.0f;
    for (int i = lane; i < group_size; i += WARP_SIZE) {
        sum += xrow[i] + bias[cbase + i];
    }
    #pragma unroll
    for (int off = WARP_SIZE / 2; off > 0; off >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    }
    const float mean = __shfl_sync(0xffffffffu, sum, 0) / (float)group_size;

    float vsum = 0.0f;
    for (int i = lane; i < group_size; i += WARP_SIZE) {
        const float d = xrow[i] + bias[cbase + i] - mean;
        vsum += d * d;
    }
    #pragma unroll
    for (int off = WARP_SIZE / 2; off > 0; off >>= 1) {
        vsum += __shfl_down_sync(0xffffffffu, vsum, off);
    }
    const float var = __shfl_sync(0xffffffffu, vsum, 0) / (float)group_size;
    const float rstd = rsqrtf(var + eps);

    for (int i = lane; i < group_size; i += WARP_SIZE) {
        const int c = cbase + i;
        float y = (xrow[i] + bias[c] - mean) * rstd * gamma[c] + beta[c];
        y = y / (1.0f + expf(-y));   // swish
        y = y * mul_w[c];
        y = y / (1.0f + expf(-y));   // swish
        orow[i] = y;
    }
}

torch::Tensor fused_bias_gn_swish_mul_swish(
    torch::Tensor x,
    torch::Tensor bias,
    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 bias_c = bias.contiguous();
    auto gamma_c = gamma.contiguous();
    auto beta_c = beta.contiguous();
    auto w_c = mul_w.contiguous();

    auto out = torch::empty_like(xc);
    const int batch_size = xc.size(0);
    const int channels = xc.size(1);

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_bias_gn_swish_mul_swish_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(),
        bias_c.data_ptr<float>(),
        gamma_c.data_ptr<float>(),
        beta_c.data_ptr<float>(),
        w_c.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        channels,
        (int)num_groups,
        (float)eps);

    return out;
}
"""

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


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM (bias handled in the epilogue) followed by a
    single fused CUDA kernel implementing
    bias-add + GroupNorm + Swish + Multiply + Swish.
    """
    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 = float(self.group_norm.eps)

    def forward(self, x):
        # Fallback to the reference implementation for non-CUDA / non-fp32 inputs.
        if (not x.is_cuda) or x.dtype != torch.float32:
            y = F.linear(x, self.gemm.weight, self.gemm.bias)
            y = F.group_norm(y, self.num_groups, self.group_norm.weight,
                             self.group_norm.bias, self.eps)
            y = y * torch.sigmoid(y)
            y = y * self.multiply_weight
            y = y * torch.sigmoid(y)
            return y

        in_shape = x.shape
        if x.dim() != 2:
            x = x.reshape(-1, in_shape[-1])

        # GEMM via cuBLAS; the linear bias is fused into the epilogue kernel.
        y = torch.matmul(x, self.gemm.weight.t())

        # Fused: + bias, GroupNorm, Swish, * multiply_weight, Swish.
        y = fused_ops.fused_bias_gn_swish_mul_swish(
            y,
            self.gemm.bias,
            self.group_norm.weight,
            self.group_norm.bias,
            self.multiply_weight,
            self.num_groups,
            self.eps,
        )

        if len(in_shape) != 2:
            y = y.reshape(tuple(in_shape[:-1]) + (y.shape[-1],))
        return y
```