```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 + per-channel Multiply + Swish
# The GEMM (nn.Linear) is kept as-is since cuBLAS/cublasLt is already
# near-optimal; everything after the GEMM is fused into a single
# memory-bound kernel (one warp per GroupNorm group).
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

// Specialized kernel for group_size == 32 (one element per lane).
// Each warp handles exactly one (batch, group) pair.
__global__ void fused_gn_swish_mul_swish_g32_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int num_groups,
    float eps,
    long total_groups)
{
    long warp_id = (long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    if (warp_id >= total_groups) return;
    int lane = threadIdx.x & 31;

    long idx = warp_id * 32 + lane;
    int channel = (int)(warp_id % num_groups) * 32 + lane;

    float v = x[idx];

    // Two-pass mean/variance for numerical parity with torch GroupNorm.
    float sum = v;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1)
        sum += __shfl_xor_sync(0xffffffff, sum, o);
    float mean = sum * (1.0f / 32.0f);

    float d = v - mean;
    float ss = d * d;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1)
        ss += __shfl_xor_sync(0xffffffff, ss, o);
    float rstd = rsqrtf(ss * (1.0f / 32.0f) + eps);

    // GroupNorm affine folded into scale/shift
    float scale = rstd * gn_w[channel];
    float shift = gn_b[channel] - mean * scale;

    float y = fmaf(v, scale, shift);
    y = y / (1.0f + expf(-y));   // Swish
    y = y * mul_w[channel];      // per-channel multiply
    y = y / (1.0f + expf(-y));   // Swish
    out[idx] = y;
}

// Generic fallback: one warp per group, arbitrary group_size.
__global__ void fused_gn_swish_mul_swish_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int num_groups,
    int group_size,
    float eps,
    long total_groups)
{
    long warp_id = (long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    if (warp_id >= total_groups) return;
    int lane = threadIdx.x & 31;

    long base = warp_id * (long)group_size;
    int ch0 = (int)(warp_id % num_groups) * group_size;

    float sum = 0.f, sumsq = 0.f;
    for (int i = lane; i < group_size; i += 32) {
        float v = x[base + i];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_xor_sync(0xffffffff, sum, o);
        sumsq += __shfl_xor_sync(0xffffffff, sumsq, o);
    }
    float inv_n = 1.0f / (float)group_size;
    float mean = sum * inv_n;
    float var = fmaf(-mean, mean, sumsq * inv_n);
    float rstd = rsqrtf(var + eps);

    for (int i = lane; i < group_size; i += 32) {
        float v = x[base + i];
        int c = ch0 + i;
        float y = fmaf(v - 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[base + i] = 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() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto gnwc = gn_w.contiguous();
    auto gnbc = gn_b.contiguous();
    auto mwc = mul_w.contiguous();

    auto out = at::empty_like(xc);

    int64_t channels = xc.size(-1);
    int64_t batch = xc.numel() / channels;
    int64_t group_size = channels / num_groups;
    int64_t total_groups = batch * num_groups;

    const int threads = 256;
    const int warps_per_block = threads / 32;
    int64_t blocks = (total_groups + warps_per_block - 1) / warps_per_block;

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

    if (group_size == 32) {
        fused_gn_swish_mul_swish_g32_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(),
            gnwc.data_ptr<float>(),
            gnbc.data_ptr<float>(),
            mwc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups,
            (float)eps,
            total_groups);
    } else {
        fused_gn_swish_mul_swish_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(),
            gnwc.data_ptr<float>(),
            gnbc.data_ptr<float>(),
            mwc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups,
            (int)group_size,
            (float)eps,
            total_groups);
    }

    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_ops = 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_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM + fused
    (GroupNorm + Swish + Multiply + Swish) CUDA kernel.
    """
    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.fused_ops = fused_ops

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            # GEMM via cuBLAS/cublasLt (bias fused in epilogue)
            x = F.linear(x, self.gemm.weight, self.gemm.bias)
            # Fused GroupNorm + Swish + Multiply + Swish
            return self.fused_ops.fused_gn_swish_mul_swish_cuda(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.group_norm.eps,
            )
        # Fallback for CPU / non-float32 inputs
        x = self.gemm(x)
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```