```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 kernel: GroupNorm + Swish + (channel-wise multiply) + Swish
#
# The GEMM (1024x8192x8192) dominates FLOPs, so we leave it to cuBLAS
# (via F.linear). Everything after the GEMM is memory-bound, so we fuse
# all of it into a single kernel:
#   - One warp handles one (batch_row, group) pair.
#   - The group statistics (mean, rstd) are computed with a warp
#     butterfly reduction (no shared memory / block sync needed).
#   - Normalization, both Swish gates and the channel multiply are
#     applied in registers, reading x once and writing the output once.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float sigmoid_f(float x) {
    return 1.0f / (1.0f + expf(-x));
}

__global__ void fused_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int batch_size,
    int channels,
    int num_groups,
    int group_size,
    float eps)
{
    const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
    const int lane = threadIdx.x & 31;
    const int total_warps = (gridDim.x * blockDim.x) >> 5;
    const int total_groups = batch_size * num_groups;

    for (int g = warp_id; g < total_groups; g += total_warps) {
        const int b = g / num_groups;
        const int grp = g - b * num_groups;
        const int64_t row_off = (int64_t)b * channels;
        const int ch_off = grp * group_size;
        const float* xr = x + row_off + ch_off;
        float* orow = out + row_off + ch_off;

        // Group statistics: sum and sum of squares
        float sum = 0.0f;
        float sumsq = 0.0f;
        for (int i = lane; i < group_size; i += 32) {
            float v = xr[i];
            sum += v;
            sumsq += v * v;
        }
        // Butterfly reduction: every lane ends up with the totals
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            sum   += __shfl_xor_sync(0xffffffffu, sum, off);
            sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off);
        }
        const float inv_n = 1.0f / (float)group_size;
        const float mean = sum * inv_n;
        float var = sumsq * inv_n - mean * mean;
        var = var < 0.0f ? 0.0f : var;
        const float rstd = rsqrtf(var + eps);

        // Apply: GroupNorm -> Swish -> Multiply -> Swish
        for (int i = lane; i < group_size; i += 32) {
            const int c = ch_off + i;
            float h = (xr[i] - mean) * rstd * gn_w[c] + gn_b[c];
            h = h * sigmoid_f(h);   // first Swish
            h = h * mul_w[c];       // channel-wise multiply
            h = h * sigmoid_f(h);   // second Swish
            orow[i] = h;
        }
    }
}

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");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(gn_w.is_contiguous() && gn_b.is_contiguous() && mul_w.is_contiguous(),
                "weights must be contiguous");

    const int batch_size = (int)x.size(0);
    const int channels = (int)x.size(1);
    const int group_size = channels / (int)num_groups;

    auto out = at::empty_like(x);

    const int threads = 256;
    const int warps_per_block = threads / 32;
    const int total_warps = batch_size * (int)num_groups;
    const int blocks = (total_warps + 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>(),
        gn_w.data_ptr<float>(),
        gn_b.data_ptr<float>(),
        mul_w.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        channels,
        (int)num_groups,
        group_size,
        (float)eps);

    return out;
}
"""

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=cpp_source,
    cuda_sources=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 + Multiply + Swish in one pass over the data.
    """
    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

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            # GEMM (cuBLAS, includes bias epilogue)
            h = F.linear(x, self.gemm.weight, self.gemm.bias)
            # Fully fused: GroupNorm + Swish + Multiply + Swish
            return fused_ops.fused_gn_swish_mul_swish_cuda(
                h.contiguous(),
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.group_norm.eps,
            )
        # Reference fallback for non-CUDA / non-float32 inputs
        h = self.gemm(x)
        h = self.group_norm(h)
        h = h * torch.sigmoid(h)
        h = h * self.multiply_weight
        h = h * torch.sigmoid(h)
        return h
```