```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 os
import sys
import glob
import site

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


# ---------------------------------------------------------------------------
# Fused epilogue: + combined bias (linear bias + add_value), Swish, Tanh, GELU.
# Note: Tanh maps into (-1, 1) and GELU on (-1, 1) stays within [-0.17, 0.85],
# so the final Hardtanh([-1, 1]) is mathematically an identity and is omitted.
# Fast-but-accurate math: __expf-based sigmoid/tanh (inf-safe), and a Taylor
# polynomial for erf valid on |z| <= 0.7072 (max abs err ~1.2e-7).
# ---------------------------------------------------------------------------
epilogue_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float act_chain(float x) {
    // Swish: x * sigmoid(x)
    float sig = 1.0f / (1.0f + __expf(-x));
    float s = x * sig;
    // Tanh: 1 - 2 / (exp(2s) + 1)  (safe for exp overflow -> inf)
    float u = __expf(2.0f * s);
    float t = 1.0f - __fdividef(2.0f, u + 1.0f);
    // GELU: 0.5 * t * (1 + erf(t / sqrt(2))), |t| <= 1
    float z = t * 0.70710678118654752f;
    float z2 = z * z;
    float p = 0.000120553394f;
    p = fmaf(p, z2, -0.000854832702f);
    p = fmaf(p, z2, 0.00522397763f);
    p = fmaf(p, z2, -0.02686617065f);
    p = fmaf(p, z2, 0.11283791671f);
    p = fmaf(p, z2, -0.37612638903f);
    p = fmaf(p, z2, 1.12837916710f);
    float e = z * p;
    return 0.5f * t * (1.0f + e);
}

template<int UNROLL>
__global__ void fused_bias_act_vec4(const float4* __restrict__ x,
                                    const float4* __restrict__ b,
                                    float4* __restrict__ out,
                                    int n4, int row4) {
    int base = blockIdx.x * (blockDim.x * UNROLL) + threadIdx.x;
    bool pow2 = (row4 & (row4 - 1)) == 0;
    #pragma unroll
    for (int u = 0; u < UNROLL; ++u) {
        int i = base + u * blockDim.x;
        if (i < n4) {
            int j = pow2 ? (i & (row4 - 1)) : (i % row4);
            float4 v = __ldcs(&x[i]);
            float4 q = __ldg(&b[j]);
            v.x = act_chain(v.x + q.x);
            v.y = act_chain(v.y + q.y);
            v.z = act_chain(v.z + q.z);
            v.w = act_chain(v.w + q.w);
            __stcs(&out[i], v);
        }
    }
}

__global__ void fused_bias_act_scalar(const float* __restrict__ x,
                                      const float* __restrict__ b,
                                      float* __restrict__ out,
                                      int n, int row) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = act_chain(__ldcs(&x[i]) + __ldg(&b[i % row]));
    }
}

torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor b) {
    TORCH_CHECK(x.is_cuda() && b.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Input must be float32");
    auto xc = x.contiguous();
    auto bc = b.contiguous();
    auto out = torch::empty_like(xc);

    int64_t n = xc.numel();
    int64_t row = bc.numel();
    if (n == 0) return out;

    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if ((row % 4 == 0) && (n % 4 == 0)) {
        int64_t n4 = n / 4;
        int64_t row4 = row / 4;
        constexpr int UNROLL = 4;
        int64_t per_block = (int64_t)threads * UNROLL;
        int blocks = (int)((n4 + per_block - 1) / per_block);
        fused_bias_act_vec4<UNROLL><<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<const float4*>(bc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            (int)n4, (int)row4);
    } else {
        int blocks = (int)((n + threads - 1) / threads);
        fused_bias_act_scalar<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(),
            bc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)n, (int)row);
    }
    return out;
}
"""

epilogue_cpp_source = "torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor b);"

epilogue_ext = None
try:
    epilogue_ext = load_inline(
        name="fused_bias_act_v6",
        cpp_sources=epilogue_cpp_source,
        cuda_sources=epilogue_source,
        functions=["fused_bias_act"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    epilogue_ext = None


# ---------------------------------------------------------------------------
# cuBLASLt GEMM wrapper with per-shape autotuning over heuristic algorithms,
# workspace sizes (so split-K variants become available) and weight layouts
# (TN via W, NN via a pre-transposed W). FP32 compute with TF32 tensor cores
# (CUBLAS_COMPUTE_32F_FAST_TF32), as permitted.
# ---------------------------------------------------------------------------
lt_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <vector>
#include <map>
#include <mutex>
#include <memory>
#include <string>

#define CULT_CHECK(x) do { cublasStatus_t _s = (x); TORCH_CHECK(_s == CUBLAS_STATUS_SUCCESS, "cublasLt error ", (int)_s); } while (0)

namespace {

struct LtPlan {
    cublasLtMatmulDesc_t opDesc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
    std::vector<cublasLtMatmulHeuristicResult_t> heur;
    void* workspace = nullptr;
    size_t ws_bytes = 0;
    ~LtPlan() {
        if (opDesc) cublasLtMatmulDescDestroy(opDesc);
        if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
        if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
        if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
        if (workspace) cudaFree(workspace);
    }
};

std::string make_key(int64_t M, int64_t N, int64_t K, bool trans_w, size_t ws) {
    return std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K)
           + (trans_w ? "_t_" : "_n_") + std::to_string(ws);
}

cublasLtHandle_t lt_handle() {
    static cublasLtHandle_t h = nullptr;
    if (!h) CULT_CHECK(cublasLtCreate(&h));
    return h;
}

std::map<std::string, std::shared_ptr<LtPlan>>& plan_cache() {
    static std::map<std::string, std::shared_ptr<LtPlan>> c;
    return c;
}
std::mutex& cache_mutex() { static std::mutex m; return m; }

// Row-major C[M,N] = X[M,K] * W^T, computed as col-major:
// C_cm[N,M] = A[N,K] * B[K,M], where B is x viewed col-major and A is the
// weight viewed col-major (trans_w=false: W is [N,K] row-major, use OP_T;
// trans_w=true: Wt is [K,N] row-major, use OP_N).
std::shared_ptr<LtPlan> get_plan(int64_t M, int64_t N, int64_t K, bool trans_w, size_t ws_bytes) {
    std::string key = make_key(M, N, K, trans_w, ws_bytes);
    std::lock_guard<std::mutex> lock(cache_mutex());
    auto it = plan_cache().find(key);
    if (it != plan_cache().end()) return it->second;

    auto plan = std::make_shared<LtPlan>();
    CULT_CHECK(cublasLtMatmulDescCreate(&plan->opDesc, CUBLAS_COMPUTE_32F_FAST_TF32, CUDA_R_32F));
    cublasOperation_t opA = trans_w ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasOperation_t opB = CUBLAS_OP_N;
    CULT_CHECK(cublasLtMatmulDescSetAttribute(plan->opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opA, sizeof(opA)));
    CULT_CHECK(cublasLtMatmulDescSetAttribute(plan->opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opB, sizeof(opB)));

    if (trans_w) {
        CULT_CHECK(cublasLtMatrixLayoutCreate(&plan->Adesc, CUDA_R_32F, K, N, K));
    } else {
        CULT_CHECK(cublasLtMatrixLayoutCreate(&plan->Adesc, CUDA_R_32F, N, K, N));
    }
    CULT_CHECK(cublasLtMatrixLayoutCreate(&plan->Bdesc, CUDA_R_32F, K, M, K));
    CULT_CHECK(cublasLtMatrixLayoutCreate(&plan->Cdesc, CUDA_R_32F, N, M, N));

    plan->ws_bytes = ws_bytes;
    if (ws_bytes > 0) {
        TORCH_CHECK(cudaMalloc(&plan->workspace, ws_bytes) == cudaSuccess, "cublasLt workspace alloc failed");
    }

    cublasLtMatmulPreference_t pref = nullptr;
    CULT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
    CULT_CHECK(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                                                    &ws_bytes, sizeof(ws_bytes)));

    const int max_results = 64;
    plan->heur.resize(max_results);
    int returned = 0;
    CULT_CHECK(cublasLtMatmulAlgoGetHeuristic(lt_handle(), plan->opDesc,
                                              plan->Adesc, plan->Bdesc, plan->Cdesc, plan->Cdesc,
                                              pref, max_results, plan->heur.data(), &returned));
    plan->heur.resize(returned);
    cublasLtMatmulPreferenceDestroy(pref);

    std::vector<cublasLtMatmulHeuristicResult_t> good;
    good.reserve(plan->heur.size());
    for (auto& h : plan->heur) {
        if (h.state == CUBLAS_STATUS_SUCCESS && h.workspaceSize <= ws_bytes) good.push_back(h);
    }
    plan->heur.swap(good);
    TORCH_CHECK(!plan->heur.empty(), "cublasLt returned no algorithms");

    plan_cache()[key] = plan;
    return plan;
}

} // namespace

int64_t lt_num_algos(int64_t M, int64_t N, int64_t K, bool trans_w, int64_t ws_bytes) {
    auto plan = get_plan(M, N, K, trans_w, (size_t)ws_bytes);
    return (int64_t)plan->heur.size();
}

void lt_gemm(torch::Tensor x, torch::Tensor w, torch::Tensor out, bool trans_w, int64_t algo_idx, int64_t ws_bytes) {
    int64_t M = x.size(0);
    int64_t K = x.size(1);
    int64_t N = out.size(1);
    auto plan = get_plan(M, N, K, trans_w, (size_t)ws_bytes);
    TORCH_CHECK(algo_idx >= 0 && algo_idx < (int64_t)plan->heur.size(), "invalid algo index");

    float alpha = 1.0f;
    float beta = 0.0f;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    CULT_CHECK(cublasLtMatmul(lt_handle(), plan->opDesc, &alpha,
                              w.data_ptr<float>(), plan->Adesc,
                              x.data_ptr<float>(), plan->Bdesc,
                              &beta,
                              out.data_ptr<float>(), plan->Cdesc,
                              out.data_ptr<float>(), plan->Cdesc,
                              &plan->heur[algo_idx].algo,
                              plan->workspace, plan->ws_bytes,
                              stream));
}

void lt_keep_only(int64_t M, int64_t N, int64_t K, bool trans_w, int64_t ws_bytes) {
    std::string prefix = std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K) + "_";
    std::string keep = make_key(M, N, K, trans_w, (size_t)ws_bytes);
    std::lock_guard<std::mutex> lock(cache_mutex());
    auto& c = plan_cache();
    for (auto it = c.begin(); it != c.end();) {
        if (it->first.compare(0, prefix.size(), prefix) == 0 && it->first != keep) {
            it = c.erase(it);
        } else {
            ++it;
        }
    }
}
"""

lt_cpp_source = (
    "int64_t lt_num_algos(int64_t M, int64_t N, int64_t K, bool trans_w, int64_t ws_bytes);\n"
    "void lt_gemm(torch::Tensor x, torch::Tensor w, torch::Tensor out, bool trans_w, int64_t algo_idx, int64_t ws_bytes);\n"
    "void lt_keep_only(int64_t M, int64_t N, int64_t K, bool trans_w, int64_t ws_bytes);"
)


def _cublaslt_paths():
    inc_dirs, lib_dirs = [], []
    roots = []
    try:
        roots += site.getsitepackages()
    except Exception:
        pass
    try:
        roots.append(os.path.dirname(os.path.dirname(os.path.abspath(torch.__file__))))
    except Exception:
        pass
    roots.append(sys.prefix)
    roots.append(os.environ.get("CUDA_HOME", "/usr/local/cuda"))
    roots.append("/usr/local/cuda")
    for r in dict.fromkeys(roots):
        for sub in ("lib64", "lib",
                    os.path.join("nvidia", "cublas", "lib"),
                    os.path.join("torch", "lib")):
            p = os.path.join(r, sub)
            if os.path.isdir(p):
                lib_dirs.append(p)
        for sub in ("include", os.path.join("nvidia", "cublas", "include")):
            p = os.path.join(r, sub)
            if os.path.isdir(p):
                inc_dirs.append(p)
    lib_file = None
    for d in lib_dirs:
        for pat in ("libcublasLt.so", "libcublasLt.so.*"):
            hits = sorted(glob.glob(os.path.join(d, pat)))
            hits = [h for h in hits if os.path.isfile(h)]
            if hits:
                lib_file = hits[0]
                break
        if lib_file:
            break
    if lib_file:
        ldflags = [lib_file, "-Wl,-rpath," + os.path.dirname(lib_file)]
    else:
        ldflags = ["-lcublasLt"]
    return inc_dirs, ldflags


lt_ext = None
try:
    _inc, _ld = _cublaslt_paths()
    lt_ext = load_inline(
        name="cublaslt_gemm_autotune_v3",
        cpp_sources=lt_cpp_source,
        cuda_sources=lt_source,
        functions=["lt_num_algos", "lt_gemm", "lt_keep_only"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        extra_include_paths=_inc,
        extra_ldflags=_ld,
    )
except Exception:
    lt_ext = None


class ModelNew(nn.Module):
    """
    Optimized model:
      * GEMM (FP32 with TF32 tensor cores) via cuBLASLt, autotuned per shape over
        heuristic algorithms, workspace sizes (unlocking split-K) and weight
        layouts (TN / pre-transposed NN). Falls back to the PyTorch cuBLAS path
        when cublasLt is unavailable.
      * Single fused CUDA epilogue kernel: combined bias (linear bias + add_value)
        + Swish + Tanh + GELU. The trailing Hardtanh([-1, 1]) is provably an
        identity after Tanh -> GELU and is dropped.
      * Whole forward is captured into a CUDA graph (keyed by input/parameter
        buffer signatures) and replayed on subsequent calls to remove host and
        launch overhead; any mismatch or failure falls back to the eager path.
    """
    def __init__(self, in_features, out_features, add_value_shape):
        super(ModelNew, self).__init__()
        # TF32 tensor-core math is explicitly permitted for FP32 workloads.
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

        self.matmul = nn.Linear(in_features, out_features)
        self.add_value = nn.Parameter(torch.randn(add_value_shape))

        self._ws_list = [0, 32 << 20, 64 << 20, 128 << 20, 256 << 20]
        self._topk = 6
        self._gemm_choice = {}
        self._wt = None
        self._wt_ver = None
        self._comb = None
        self._comb_ver = None

        self._graphs = {}
        self._counts = {}
        self._bad = set()
        self._max_graphs = 16
        self._graphs_ok = hasattr(torch.cuda, "CUDAGraph") and hasattr(torch.cuda, "graph")

    # ----- cached derived buffers -------------------------------------------
    def _get_wt(self):
        w = self.matmul.weight
        ver = (w.data_ptr(), w._version)
        if self._wt is None or self._wt_ver != ver:
            with torch.no_grad():
                self._wt = w.t().contiguous()
            self._wt_ver = ver
        return self._wt

    def _get_comb(self):
        b = self.matmul.bias
        av = self.add_value
        ver = (av.data_ptr(), av._version,
               0 if b is None else b.data_ptr(), -1 if b is None else b._version)
        if self._comb is None or self._comb_ver != ver:
            with torch.no_grad():
                self._comb = av.detach().clone() if b is None else (av + b).contiguous()
            self._comb_ver = ver
        return self._comb

    # ----- reference / eager paths ------------------------------------------
    def _reference(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        y = y + self.add_value
        y = torch.sigmoid(y) * y
        y = torch.tanh(y)
        y = F.gelu(y)
        y = F.hardtanh(y, min_val=-1, max_val=1)
        return y

    @torch.no_grad()
    def _autotune(self, x):
        M, K = x.shape
        N = self.matmul.weight.shape[0]
        w = self.matmul.weight
        wt = self._get_wt()
        out = torch.empty(M, N, device=x.device, dtype=torch.float32)

        cands = []
        cands.append((("torch_tn",), lambda: torch.matmul(x, w.t(), out=out)))
        cands.append((("torch_nn",), lambda: torch.mm(x, wt, out=out)))

        if lt_ext is not None:
            xc = x if x.is_contiguous() else x.contiguous()
            for trans_w, wbuf in ((False, w), (True, wt)):
                for ws in self._ws_list:
                    try:
                        n_alg = int(lt_ext.lt_num_algos(M, N, K, trans_w, ws))
                    except Exception:
                        continue
                    for i in range(min(n_alg, self._topk)):
                        cands.append((("lt", trans_w, i, ws),
                                      (lambda wb=wbuf, tw=trans_w, ii=i, ww=ws:
                                       lt_ext.lt_gemm(xc, wb, out, tw, ii, ww))))

        best, best_t = None, float("inf")
        for tag, fn in cands:
            try:
                for _ in range(3):
                    fn()
                torch.cuda.synchronize(x.device)
                s = torch.cuda.Event(enable_timing=True)
                e = torch.cuda.Event(enable_timing=True)
                s.record()
                for _ in range(6):
                    fn()
                e.record()
                torch.cuda.synchronize(x.device)
                t = s.elapsed_time(e)
                if t < best_t:
                    best_t, best = t, tag
            except Exception:
                continue
        if best is None:
            best = ("torch_tn",)

        # Free losing cublasLt plans (and their workspaces) for this shape.
        if lt_ext is not None:
            try:
                if best[0] == "lt":
                    lt_ext.lt_keep_only(M, N, K, best[1], best[3])
                else:
                    lt_ext.lt_keep_only(M, N, K, False, -1)
            except Exception:
                pass
        return best

    def _run_gemm(self, x, choice):
        w = self.matmul.weight
        N = w.shape[0]
        M = x.shape[0]
        tag = choice[0]
        if tag == "lt":
            _, trans_w, idx, ws = choice
            if not x.is_contiguous():
                x = x.contiguous()
            wbuf = self._get_wt() if trans_w else w
            out = torch.empty(M, N, device=x.device, dtype=torch.float32)
            lt_ext.lt_gemm(x, wbuf, out, trans_w, idx, ws)
            return out
        if tag == "torch_nn":
            return torch.mm(x, self._get_wt())
        return torch.matmul(x, w.t())

    def _eager(self, x):
        key = (tuple(x.shape), x.device.index)
        choice = self._gemm_choice.get(key)
        if choice is None:
            choice = self._autotune(x)
            self._gemm_choice[key] = choice
        y = self._run_gemm(x, choice)
        if epilogue_ext is not None:
            return epilogue_ext.fused_bias_act(y, self._get_comb())
        y = y + self._get_comb()
        y = torch.sigmoid(y) * y
        y = torch.tanh(y)
        y = F.gelu(y)
        return y

    def _capture(self, x, key):
        # Warm up on a side stream (materializes workspaces, caches, etc.).
        s = torch.cuda.Stream(x.device)
        s.wait_stream(torch.cuda.current_stream(x.device))
        with torch.cuda.stream(s):
            for _ in range(3):
                self._eager(x)
        torch.cuda.current_stream(x.device).wait_stream(s)
        torch.cuda.synchronize(x.device)

        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            static_out = self._eager(x)
        self._graphs[key] = (g, static_out)

    def forward(self, x):
        if (self._graphs_ok and x.is_cuda and x.dim() == 2
                and x.dtype == torch.float32 and self.matmul.weight.is_cuda):
            w = self.matmul.weight
            b = self.matmul.bias
            key = (x.data_ptr(), tuple(x.shape), tuple(x.stride()), x.device.index,
                   w.data_ptr(), w._version,
                   0 if b is None else b.data_ptr(), -1 if b is None else b._version,
                   self.add_value.data_ptr(), self.add_value._version)
            ent = self._graphs.get(key)
            if ent is not None:
                try:
                    ent[0].replay()
                    return ent[1]
                except Exception:
                    self._graphs.pop(key, None)
                    self._bad.add(key)
            elif key not in self._bad:
                n = self._counts.get(key, 0) + 1
                self._counts[key] = n
                if n >= 3 and len(self._graphs) < self._max_graphs:
                    try:
                        self._capture(x, key)
                        ent = self._graphs.get(key)
                        if ent is not None:
                            ent[0].replay()
                            return ent[1]
                    except Exception:
                        self._graphs.pop(key, None)
                        self._bad.add(key)
            return self._eager(x)
        # Non-CUDA fallback: replicate the reference computation exactly.
        return self._reference(x)
```