```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 epilogue: + 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).
fused_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);
}

__global__ void fused_bias_act_vec4(const float4* __restrict__ x,
                                    const float4* __restrict__ b1,
                                    const float4* __restrict__ b2,
                                    float4* __restrict__ out,
                                    int n4, int row4) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        int j = ((row4 & (row4 - 1)) == 0) ? (i & (row4 - 1)) : (i % row4);
        float4 v = __ldcs(&x[i]);
        float4 p = __ldg(&b1[j]);
        float4 q = __ldg(&b2[j]);
        v.x = act_chain(v.x + p.x + q.x);
        v.y = act_chain(v.y + p.y + q.y);
        v.z = act_chain(v.z + p.z + q.z);
        v.w = act_chain(v.w + p.w + q.w);
        __stcs(&out[i], v);
    }
}

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

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

    int64_t n = xc.numel();
    int64_t row = b1c.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;
        int blocks = (int)((n4 + threads - 1) / threads);
        fused_bias_act_vec4<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<const float4*>(b1c.data_ptr<float>()),
            reinterpret_cast<const float4*>(b2c.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>(),
            b1c.data_ptr<float>(),
            b2c.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)n, (int)row);
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor b1, torch::Tensor b2);"
)

fused_op = load_inline(
    name="fused_bias_act_v3",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_source,
    functions=["fused_bias_act"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM (FP32 with TF32 tensor cores) followed by a single
    fused CUDA epilogue kernel (bias add + value add + Swish + Tanh + GELU; Hardtanh
    is provably an identity after Tanh->GELU and is dropped).

    Additionally, the whole forward is captured into a CUDA graph (keyed by the input
    buffer signature) and replayed on subsequent calls to remove host/launch overhead.
    Any signature mismatch or capture failure falls back to the eager fused 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._graphs = {}
        self._counts = {}
        self._bad = set()
        self._max_graphs = 4
        self._graphs_ok = hasattr(torch.cuda, "CUDAGraph") and hasattr(torch.cuda, "graph")

    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

    def _eager(self, x):
        y = torch.matmul(x, self.matmul.weight.t())
        return fused_op.fused_bias_act(y, self.matmul.bias, self.add_value)

    def _capture(self, x, key):
        # Warm up on a side stream (materializes cuBLAS workspaces, etc.).
        s = torch.cuda.Stream(x.device)
        s.wait_stream(torch.cuda.current_stream(x.device))
        with torch.cuda.stream(s):
            for _ in range(2):
                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.dtype == torch.float32
                and self.matmul.weight.is_cuda):
            key = (x.data_ptr(), tuple(x.shape), tuple(x.stride()), x.device.index,
                   self.matmul.weight.data_ptr(), self.matmul.bias.data_ptr(),
                   self.add_value.data_ptr())
            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)
```