```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

# Fallback matmuls should at least use TF32 tensor cores.
torch.backends.cuda.matmul.allow_tf32 = True

# ---------------------------------------------------------------------------
# Fused epilogue: bias-add + GroupNorm + Swish + per-channel multiply + Swish
# ---------------------------------------------------------------------------
epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define FULL_MASK 0xffffffffu

// Vectorized kernel for group_size == 32 and num_groups % 4 == 0.
// Each warp processes 4 consecutive groups (128 floats) with float4 I/O;
// each group lives entirely within an 8-lane segment of the warp.
__global__ void fused_bias_gn_swish_mul_swish_vec4_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    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_quads)
{
    long warp_id = (long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    if (warp_id >= total_quads) return;
    int lane = threadIdx.x & 31;

    long base = warp_id << 7;  // 128 floats per warp
    int ch0 = (int)(((warp_id << 2) % num_groups) << 5);
    int off = ch0 + (lane << 2);

    float4 v = *reinterpret_cast<const float4*>(x + base + (lane << 2));
    float4 b = *reinterpret_cast<const float4*>(bias + off);
    v.x += b.x; v.y += b.y; v.z += b.z; v.w += b.w;

    float s = (v.x + v.y) + (v.z + v.w);
    #pragma unroll
    for (int o = 4; o > 0; o >>= 1) s += __shfl_xor_sync(FULL_MASK, s, o, 8);
    float mean = s * (1.0f / 32.0f);

    float d0 = v.x - mean, d1 = v.y - mean, d2 = v.z - mean, d3 = v.w - mean;
    float ss = fmaf(d0, d0, fmaf(d1, d1, fmaf(d2, d2, d3 * d3)));
    #pragma unroll
    for (int o = 4; o > 0; o >>= 1) ss += __shfl_xor_sync(FULL_MASK, ss, o, 8);
    float rstd = rsqrtf(ss * (1.0f / 32.0f) + eps);

    float4 w = *reinterpret_cast<const float4*>(gn_w + off);
    float4 g = *reinterpret_cast<const float4*>(gn_b + off);
    float4 m = *reinterpret_cast<const float4*>(mul_w + off);

    #define EPILOGUE(val, ww, gg, mm)                      \
        {                                                  \
            float sc = rstd * (ww);                        \
            float y = fmaf(val, sc, (gg) - mean * sc);     \
            y = y / (1.0f + __expf(-y));                   \
            y = y * (mm);                                  \
            y = y / (1.0f + __expf(-y));                   \
            val = y;                                       \
        }
    EPILOGUE(v.x, w.x, g.x, m.x);
    EPILOGUE(v.y, w.y, g.y, m.y);
    EPILOGUE(v.z, w.z, g.z, m.z);
    EPILOGUE(v.w, w.w, g.w, m.w);
    #undef EPILOGUE

    *reinterpret_cast<float4*>(out + base + (lane << 2)) = v;
}

// Generic fallback: one warp per group, arbitrary group_size.
__global__ void fused_bias_gn_swish_mul_swish_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    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] + bias[ch0 + i];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_xor_sync(FULL_MASK, sum, o);
        sumsq += __shfl_xor_sync(FULL_MASK, 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) {
        int c = ch0 + i;
        float v = x[base + i] + bias[c];
        float sc = rstd * gn_w[c];
        float y = fmaf(v, sc, gn_b[c] - mean * sc);
        y = y / (1.0f + __expf(-y));
        y = y * mul_w[c];
        y = y / (1.0f + __expf(-y));
        out[base + i] = y;
    }
}

torch::Tensor fused_bias_gn_swish_mul_swish_cuda(
    torch::Tensor x,
    torch::Tensor bias,
    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 bc = bias.contiguous();
    auto wc = gn_w.contiguous();
    auto gc = gn_b.contiguous();
    auto mc = 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;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (group_size == 32 && (num_groups % 4) == 0) {
        int64_t quads = total_groups / 4;
        int64_t blocks = (quads + warps_per_block - 1) / warps_per_block;
        fused_bias_gn_swish_mul_swish_vec4_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            wc.data_ptr<float>(), gc.data_ptr<float>(), mc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups, (float)eps, quads);
    } else {
        int64_t blocks = (total_groups + warps_per_block - 1) / warps_per_block;
        fused_bias_gn_swish_mul_swish_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            wc.data_ptr<float>(), gc.data_ptr<float>(), mc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups, (int)group_size, (float)eps, total_groups);
    }

    return out;
}
"""

epilogue_cpp_source = (
    "torch::Tensor fused_bias_gn_swish_mul_swish_cuda("
    "torch::Tensor x, torch::Tensor bias, torch::Tensor gn_w, "
    "torch::Tensor gn_b, torch::Tensor mul_w, int64_t num_groups, double eps);"
)

fused_epilogue = load_inline(
    name="fused_bias_gn_swish_mul_swish_v2",
    cpp_sources=epilogue_cpp_source,
    cuda_sources=epilogue_cuda_source,
    functions=["fused_bias_gn_swish_mul_swish_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)

# ---------------------------------------------------------------------------
# Fast GEMM: FP16 inputs on tensor cores, FP32 accumulate, FP32 output.
# (FP16 has the same 11-bit significand as TF32, but 2x the throughput.)
# Bias is intentionally left out; it is folded into the fused epilogue.
# ---------------------------------------------------------------------------
gemm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>

torch::Tensor gemm_fp16acc32_cuda(torch::Tensor x, torch::Tensor w) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(w.scalar_type() == at::kHalf || w.scalar_type() == at::kFloat,
                "w must be float16 or float32");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    int64_t K = xc.size(-1);
    int64_t M = xc.numel() / K;
    int64_t N = wc.size(0);
    TORCH_CHECK(wc.size(1) == K, "weight shape mismatch");

    auto out = at::empty({M, N}, xc.options());
    auto xh = xc.to(at::kHalf);
    auto wh = (wc.scalar_type() == at::kHalf) ? wc : wc.to(at::kHalf);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    const float alpha = 1.0f;
    const float beta = 0.0f;
    // Row-major trick: out^T (N x M, col-major) = W^T * X^T.
    cublasStatus_t status = cublasGemmEx(
        handle, CUBLAS_OP_T, CUBLAS_OP_N,
        (int)N, (int)M, (int)K,
        &alpha,
        wh.data_ptr(), CUDA_R_16F, (int)K,
        xh.data_ptr(), CUDA_R_16F, (int)K,
        &beta,
        out.data_ptr<float>(), CUDA_R_32F, (int)N,
        CUBLAS_COMPUTE_32F,
        CUBLAS_GEMM_DEFAULT);
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "cublasGemmEx failed: ", (int)status);
    return out;
}
"""

gemm_cpp_source = (
    "torch::Tensor gemm_fp16acc32_cuda(torch::Tensor x, torch::Tensor w);"
)

try:
    fast_gemm = load_inline(
        name="fast_gemm_fp16acc32_v1",
        cpp_sources=gemm_cpp_source,
        cuda_sources=gemm_cuda_source,
        functions=["gemm_fp16acc32_cuda"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        extra_ldflags=["-lcublas"],
    )
except Exception:
    fast_gemm = None


class ModelNew(nn.Module):
    """
    Optimized model:
      * GEMM on FP16 tensor cores (FP32 accumulate/output, cached half weight)
      * Fused bias + GroupNorm + Swish + Multiply + Swish epilogue 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._w_half = None
        self._w_half_key = None

    def _half_weight(self):
        w = self.gemm.weight
        key = (w._version, w.data_ptr(), w.device)
        if self._w_half is None or self._w_half_key != key:
            self._w_half = w.detach().half()
            self._w_half_key = key
        return self._w_half

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            x2 = x.reshape(-1, x.shape[-1])
            if fast_gemm is not None:
                h = fast_gemm.gemm_fp16acc32_cuda(x2, self._half_weight())
            else:
                torch.backends.cuda.matmul.allow_tf32 = True
                h = F.linear(x2, self.gemm.weight)
            y = fused_epilogue.fused_bias_gn_swish_mul_swish_cuda(
                h,
                self.gemm.bias,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.group_norm.eps,
            )
            return y.view(*x.shape[:-1], y.shape[-1])
        # 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
```