```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
from torch.utils.cpp_extension import load_inline

# Custom fused CUDA op:
#   1) GEMM (delegated to cuBLAS via at::linear, same numerics as nn.Linear)
#   2) GroupNorm + Swish + per-channel Multiply + Swish fused into ONE kernel.
#      Each warp handles one (batch-row, group) pair, computing mean/var with
#      warp-shuffle reductions and applying the whole elementwise chain from
#      registers, so the tensor is only read/written once for the entire epilogue.
fused_ops_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__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,
    const int C,
    const int cpg,
    const int num_groups,
    const float eps)
{
    const int warp_id = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int groups_per_block = blockDim.x >> 5;
    const int group = blockIdx.x * groups_per_block + warp_id;
    if (group >= num_groups) return;

    const int row = blockIdx.y;
    const long base = (long)row * (long)C + (long)group * (long)cpg;

    // ---- group mean ----
    float sum = 0.0f;
    for (int i = lane; i < cpg; i += 32) sum += x[base + i];
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    const float mean = __shfl_sync(0xffffffffu, sum, 0) / (float)cpg;

    // ---- group variance (biased) ----
    float sq = 0.0f;
    for (int i = lane; i < cpg; i += 32) {
        const float d = x[base + i] - mean;
        sq += d * d;
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        sq += __shfl_down_sync(0xffffffffu, sq, off);
    const float rstd = rsqrtf(__shfl_sync(0xffffffffu, sq, 0) / (float)cpg + eps);

    // ---- normalize -> swish -> multiply -> swish ----
    for (int i = lane; i < cpg; i += 32) {
        const int c = group * cpg + i;
        const float xv = x[base + i];
        const float norm = (xv - mean) * rstd * gn_w[c] + gn_b[c];
        float v = norm / (1.0f + expf(-norm));  // swish: x * sigmoid(x)
        v = v * mul_w[c];                       // per-channel multiply
        v = v / (1.0f + expf(-v));              // swish again
        out[base + i] = v;
    }
}

torch::Tensor fused_forward(
    torch::Tensor x,
    torch::Tensor gemm_w,
    torch::Tensor gemm_b,
    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(num_groups > 0, "num_groups must be positive");

    // GEMM through cuBLAS (bit-identical path to nn.Linear).
    auto h = at::linear(x.contiguous(), gemm_w.contiguous(), gemm_b.contiguous());

    const int64_t C = h.size(-1);
    TORCH_CHECK(C % num_groups == 0, "channels must be divisible by num_groups");
    const int cpg = (int)(C / num_groups);

    auto h2 = h.reshape({-1, C});
    const int64_t B = h2.size(0);

    auto gn_wc = gn_w.contiguous();
    auto gn_bc = gn_b.contiguous();
    auto mul_wc = mul_w.contiguous();

    auto out = torch::empty_like(h2);

    const int threads = 256;                 // 8 warps -> 8 groups per block
    const int groups_per_block = threads / 32;
    dim3 grid((unsigned int)((num_groups + groups_per_block - 1) / groups_per_block),
              (unsigned int)B);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_gn_swish_mul_swish_kernel<<<grid, threads, 0, stream>>>(
        h2.data_ptr<float>(),
        gn_wc.data_ptr<float>(),
        gn_bc.data_ptr<float>(),
        mul_wc.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)C, cpg, (int)num_groups, (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out.reshape(h.sizes());
}
"""

fused_ops_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor gemm_w, "
    "torch::Tensor gemm_b, torch::Tensor gn_w, torch::Tensor gn_b, "
    "torch::Tensor mul_w, int64_t num_groups, double eps);"
)

fused_ops = load_inline(
    name="fused_gemm_gn_swish_mul_swish",
    cpp_sources=fused_ops_cpp_source,
    cuda_sources=fused_ops_cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM runs on cuBLAS and the entire
    GroupNorm + Swish + Multiply + Swish chain is fused into a single
    custom CUDA kernel (one pass over the data instead of 4+ kernel launches).
    """
    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):
        return fused_ops.fused_forward(
            x,
            self.gemm.weight,
            self.gemm.bias,
            self.group_norm.weight,
            self.group_norm.bias,
            self.multiply_weight,
            self.num_groups,
            self.eps,
        )
```