```
# 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 -> per-channel Multiply -> Swish
# The GEMM is left to cuBLAS (nn.Linear), which is already near-optimal.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

// x: (N, C) contiguous fp32; gamma/beta/mw: (C,)
// One warp handles one (batch, group) pair.
__global__ void fused_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mw,
    float* __restrict__ out,
    const int num_groups,
    const int group_size,
    const long total_tasks,
    const float eps)
{
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int warps_per_block = blockDim.x >> 5;
    const long task = (long)blockIdx.x * warps_per_block + warp;
    if (task >= total_tasks) return;

    const int b = (int)(task / num_groups);
    const int g = (int)(task % num_groups);
    const long C = (long)num_groups * (long)group_size;

    const float* row = x + (long)b * C + (long)g * group_size;
    float* orow = out + (long)b * C + (long)g * group_size;
    const int base = g * group_size;

    if (group_size <= 32) {
        // Fast path: one element per lane (this is the case for C=8192, G=256).
        const float v = (lane < group_size) ? row[lane] : 0.f;

        float sum = v;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
        const float mean = sum / (float)group_size;

        const float d = (lane < group_size) ? (v - mean) : 0.f;
        float vsum = d * d;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) vsum += __shfl_xor_sync(0xffffffffu, vsum, o);
        const float rstd = rsqrtf(vsum / (float)group_size + eps);

        if (lane < group_size) {
            const float nrm = (v - mean) * rstd * gamma[base + lane] + beta[base + lane];
            const float s = nrm / (1.f + expf(-nrm));   // Swish
            const float m = s * mw[base + lane];        // per-channel multiply
            orow[lane] = m / (1.f + expf(-m));          // Swish
        }
    } else {
        // General path: group_size <= 512, values cached in registers.
        float vals[16];
        int cnt = 0;
        float sum = 0.f;
        for (int i = lane; i < group_size; i += 32) {
            const float v = row[i];
            vals[cnt++] = v;
            sum += v;
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
        const float mean = sum / (float)group_size;

        float vsum = 0.f;
        for (int k = 0; k < cnt; ++k) {
            const float d = vals[k] - mean;
            vsum += d * d;
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) vsum += __shfl_xor_sync(0xffffffffu, vsum, o);
        const float rstd = rsqrtf(vsum / (float)group_size + eps);

        int idx = lane;
        for (int k = 0; k < cnt; ++k, idx += 32) {
            const float nrm = (vals[k] - mean) * rstd * gamma[base + idx] + beta[base + idx];
            const float s = nrm / (1.f + expf(-nrm));
            const float m = s * mw[base + idx];
            orow[idx] = m / (1.f + expf(-m));
        }
    }
}

torch::Tensor fused_gn_swish_mul_swish(
    torch::Tensor x,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor mw,
    int64_t num_groups,
    double eps)
{
    const int batch = x.size(0);
    const int C = x.size(1);
    const int group_size = (int)(C / num_groups);

    auto out = torch::empty_like(x);
    const long total_tasks = (long)batch * (long)num_groups;
    if (total_tasks == 0) return out;

    const int warps_per_block = 8;
    const int threads = warps_per_block * 32;
    const long blocks = (total_tasks + warps_per_block - 1) / warps_per_block;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_gn_swish_mul_swish_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        mw.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)num_groups,
        group_size,
        total_tasks,
        (float)eps);
    return out;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: GEMM via cuBLAS (nn.Linear), followed by a single fused
    CUDA kernel performing 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.out_features = out_features
        self.eps = float(self.group_norm.eps)

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features)
        x = self.gemm(x)
        group_size = self.out_features // self.num_groups
        if (x.is_cuda and x.dtype == torch.float32
                and self.out_features % self.num_groups == 0
                and group_size <= 512):
            # Fused: GroupNorm -> Swish -> Multiply -> Swish
            return fused_ops.fused_gn_swish_mul_swish(
                x.contiguous(),
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )
        # Fallback to the reference PyTorch implementation
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```