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

# ---------------------------------------------------------------------------
# One merged extension:
#   * vectorized FP32 -> FP16 cast
#   * GEMM with one-time autotuning across:
#       - cublasLt FP16 (FP32 accumulate) heuristic algos
#       - cublasGemmEx FP16 (FP32 accumulate)
#       - cublasGemmEx FP32 with CUBLAS_COMPUTE_32F_FAST_16F (native fp16 TC)
#   * fused bias + GroupNorm + Swish + per-channel multiply + Swish epilogue
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <functional>
#include <algorithm>

#define FULL_MASK 0xffffffffu

// ------------------------- vectorized FP32 -> FP16 -------------------------
__global__ void cast_f32_f16_vec4_kernel(const float4* __restrict__ in,
                                         float2* __restrict__ out,
                                         long n4) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n4) return;
    float4 v = in[i];
    __half2 h[2];
    h[0] = __floats2half2_rn(v.x, v.y);
    h[1] = __floats2half2_rn(v.z, v.w);
    out[i] = *reinterpret_cast<float2*>(h);
}

static void cast_to_half(const float* in, void* out, long numel, cudaStream_t stream) {
    long n4 = numel >> 2;
    const int threads = 256;
    long blocks = (n4 + threads - 1) / threads;
    cast_f32_f16_vec4_kernel<<<blocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(in),
        reinterpret_cast<float2*>(out), n4);
}

// ------------------------------- GEMM plan ---------------------------------
namespace {

struct GemmPlan {
    long M = -1, N = -1, K = -1;
    int mode = 2;                 // 0: Lt fp16, 1: gemmEx FAST_16F, 2: gemmEx fp16
    cublasLtMatmulAlgo_t algo{};
    size_t ws = 0;
    bool valid = false;
};

GemmPlan g_plan;
cublasLtHandle_t g_lt = nullptr;
torch::Tensor* g_ws = nullptr;    // leaked on purpose (avoid exit-time teardown)

float time_op(const std::function<cublasStatus_t()>& fn, int warmup, int iters,
              cudaStream_t stream, bool* ok) {
    for (int i = 0; i < warmup; i++) {
        if (fn() != CUBLAS_STATUS_SUCCESS) { *ok = false; return 1e30f; }
    }
    cudaEvent_t e0, e1;
    cudaEventCreate(&e0);
    cudaEventCreate(&e1);
    cudaEventRecord(e0, stream);
    cublasStatus_t st = CUBLAS_STATUS_SUCCESS;
    for (int i = 0; i < iters; i++) {
        st = fn();
        if (st != CUBLAS_STATUS_SUCCESS) break;
    }
    cudaEventRecord(e1, stream);
    cudaEventSynchronize(e1);
    float ms = 0.f;
    cudaEventElapsedTime(&ms, e0, e1);
    cudaEventDestroy(e0);
    cudaEventDestroy(e1);
    *ok = (st == CUBLAS_STATUS_SUCCESS);
    return *ok ? ms / iters : 1e30f;
}

}  // namespace

torch::Tensor gemm_autotune_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor w_half) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto whc = w_half.contiguous();
    long K = xc.size(-1);
    long M = xc.numel() / K;
    long N = wc.size(0);
    TORCH_CHECK(wc.size(1) == K, "weight shape mismatch");
    TORCH_CHECK((xc.numel() % 4) == 0, "x numel must be divisible by 4");

    auto out = at::empty({M, N}, xc.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasHandle_t blas = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(blas, stream);

    const float alpha = 1.0f;
    const float beta0 = 0.0f;

    torch::Tensor xh;
    bool need_plan = !(g_plan.valid && g_plan.M == M && g_plan.N == N && g_plan.K == K);
    bool need_xh = need_plan || g_plan.mode != 1;
    if (need_xh) {
        xh = at::empty({M, K}, xc.options().dtype(at::kHalf));
        cast_to_half(xc.data_ptr<float>(), xh.data_ptr(), xc.numel(), stream);
    }

    auto run_fast16f = [&]() -> cublasStatus_t {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)N, (int)M, (int)K, &alpha,
                            wc.data_ptr<float>(), CUDA_R_32F, (int)K,
                            xc.data_ptr<float>(), CUDA_R_32F, (int)K,
                            &beta0,
                            out.data_ptr<float>(), CUDA_R_32F, (int)N,
                            CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT);
    };
    auto run_fp16 = [&]() -> cublasStatus_t {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)N, (int)M, (int)K, &alpha,
                            whc.data_ptr(), CUDA_R_16F, (int)K,
                            xh.data_ptr(), CUDA_R_16F, (int)K,
                            &beta0,
                            out.data_ptr<float>(), CUDA_R_32F, (int)N,
                            CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    };

    if (!need_plan) {
        cublasStatus_t st = CUBLAS_STATUS_SUCCESS;
        if (g_plan.mode == 1) {
            st = run_fast16f();
        } else if (g_plan.mode == 2) {
            st = run_fp16();
        } else {
            void* wsp = (g_plan.ws > 0) ? (*g_ws).data_ptr() : nullptr;
            cublasLtMatmulDesc_t desc;
            cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
            cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
            cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
            cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
            cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc;
            cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_16F, K, N, K);
            cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_16F, K, M, K);
            cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N);
            st = cublasLtMatmul(g_lt, desc, &alpha,
                                whc.data_ptr(), Adesc, xh.data_ptr(), Bdesc, &beta0,
                                out.data_ptr<float>(), Cdesc, out.data_ptr<float>(), Cdesc,
                                &g_plan.algo, wsp, g_plan.ws, stream);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            cublasLtMatmulDescDestroy(desc);
        }
        TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "GEMM failed: ", (int)st);
        return out;
    }

    // ------------------------------ autotune ------------------------------
    GemmPlan best;
    best.mode = 2;
    best.ws = 0;
    float best_t = 1e30f;
    bool ok = false;

    {
        float t = time_op(run_fp16, 1, 3, stream, &ok);
        if (ok && t < best_t) { best_t = t; best.mode = 2; }
    }
    {
        float t = time_op(run_fast16f, 1, 3, stream, &ok);
        if (ok && t < best_t) { best_t = t; best.mode = 1; }
    }

    // cublasLt heuristic search (fp16 in, fp32 out, fp32 accumulate)
    if (g_lt == nullptr) {
        if (cublasLtCreate(&g_lt) != CUBLAS_STATUS_SUCCESS) g_lt = nullptr;
    }
    if (g_lt != nullptr) {
        cublasLtMatmulDesc_t desc;
        cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
        cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
        cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc;
        cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_16F, K, N, K);
        cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_16F, K, M, K);
        cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N);
        cublasLtMatmulPreference_t pref;
        cublasLtMatmulPreferenceCreate(&pref);
        size_t max_ws = 64ull << 20;
        cublasLtMatmulPreferenceSetAttribute(
            pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_ws, sizeof(max_ws));

        const int max_results = 10;
        cublasLtMatmulHeuristicResult_t results[max_results];
        int count = 0;
        cublasLtMatmulAlgoGetHeuristic(g_lt, desc, Adesc, Bdesc, Cdesc, Cdesc,
                                       pref, max_results, results, &count);

        size_t need_ws = 0;
        for (int i = 0; i < count; i++) {
            if (results[i].state == CUBLAS_STATUS_SUCCESS)
                need_ws = std::max(need_ws, results[i].workspaceSize);
        }
        if (need_ws > 0) {
            if (g_ws == nullptr) g_ws = new torch::Tensor();
            if (!(*g_ws).defined() || (size_t)(*g_ws).numel() < need_ws) {
                *g_ws = at::empty({(long)need_ws}, xc.options().dtype(at::kByte));
            }
        }
        void* wsp = need_ws > 0 ? (*g_ws).data_ptr() : nullptr;

        for (int i = 0; i < count; i++) {
            if (results[i].state != CUBLAS_STATUS_SUCCESS) continue;
            cublasLtMatmulAlgo_t algo = results[i].algo;
            size_t a_ws = results[i].workspaceSize;
            auto fn = [&]() -> cublasStatus_t {
                return cublasLtMatmul(g_lt, desc, &alpha,
                                      whc.data_ptr(), Adesc, xh.data_ptr(), Bdesc, &beta0,
                                      out.data_ptr<float>(), Cdesc,
                                      out.data_ptr<float>(), Cdesc,
                                      &algo, wsp, a_ws, stream);
            };
            float t = time_op(fn, 1, 3, stream, &ok);
            if (ok && t < best_t) {
                best_t = t;
                best.mode = 0;
                best.algo = algo;
                best.ws = a_ws;
            }
        }
        cublasLtMatmulPreferenceDestroy(pref);
        cublasLtMatrixLayoutDestroy(Adesc);
        cublasLtMatrixLayoutDestroy(Bdesc);
        cublasLtMatrixLayoutDestroy(Cdesc);
        cublasLtMatmulDescDestroy(desc);
    }

    g_plan = best;
    g_plan.M = M; g_plan.N = N; g_plan.K = K;
    g_plan.valid = true;

    // Final run with the winning algorithm so `out` is deterministic.
    cublasStatus_t st = CUBLAS_STATUS_SUCCESS;
    if (g_plan.mode == 1) {
        st = run_fast16f();
    } else if (g_plan.mode == 2) {
        st = run_fp16();
    } else {
        void* wsp = (g_plan.ws > 0) ? (*g_ws).data_ptr() : nullptr;
        cublasLtMatmulDesc_t desc;
        cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
        cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
        cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc;
        cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_16F, K, N, K);
        cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_16F, K, M, K);
        cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N);
        st = cublasLtMatmul(g_lt, desc, &alpha,
                            whc.data_ptr(), Adesc, xh.data_ptr(), Bdesc, &beta0,
                            out.data_ptr<float>(), Cdesc, out.data_ptr<float>(), Cdesc,
                            &g_plan.algo, wsp, g_plan.ws, stream);
        cublasLtMatrixLayoutDestroy(Adesc);
        cublasLtMatrixLayoutDestroy(Bdesc);
        cublasLtMatrixLayoutDestroy(Cdesc);
        cublasLtMatmulDescDestroy(desc);
    }
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "GEMM failed: ", (int)st);
    return out;
}

// -------------------- fused epilogue: bias + GN + swish --------------------
// Vectorized kernel for group_size == 32 and num_groups % 4 == 0.
__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;
}
"""

fused_cpp_source = (
    "torch::Tensor gemm_autotune_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor w_half);\n"
    "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_ops = load_inline(
    name="fused_gemm_gn_swish_v3",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["gemm_autotune_cuda", "fused_bias_gn_swish_mul_swish_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas", "-lcublasLt"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * GEMM via one-time autotuned cuBLASLt/cuBLAS FP16 tensor-core path
        (FP32 accumulate / FP32 output), with an optional native FAST_16F
        mode that skips the explicit cast when it is faster.
      * Cached half-precision 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.eps = self.group_norm.eps
        self._w_half = None
        self._w_half_key = None
        self._gemm_fn = fused_ops.gemm_autotune_cuda
        self._epilogue_fn = fused_ops.fused_bias_gn_swish_mul_swish_cuda

    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().contiguous()
            self._w_half_key = key
        return self._w_half

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            two_d = x.dim() == 2
            x2 = x if two_d else x.reshape(-1, x.shape[-1])
            h = self._gemm_fn(x2, self.gemm.weight, self._half_weight())
            y = self._epilogue_fn(
                h,
                self.gemm.bias,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )
            return y if two_d else 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
```