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

# ---------------------------------------------------------------------------
# Custom CUDA kernel: fused epilogue
#   z = gemm_out + linear_bias + add_value
#   z = z * sigmoid(z)        (Swish)
#   z = tanh(z)               (Tanh)
#   z = gelu(z)               (GELU, exact erf formulation)
#   z = clamp(z, -1, 1)       (Hardtanh)
# All in a single pass over the data (vectorized with float4).
# ---------------------------------------------------------------------------
_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float act_chain(float z) {
    // Swish: z * sigmoid(z)
    float sig = 1.0f / (1.0f + expf(-z));
    float s = z * sig;
    // Tanh
    float t = tanhf(s);
    // GELU (erf formulation, matches torch.nn.functional.gelu default)
    float g = 0.5f * t * (1.0f + erff(t * 0.70710678118654752440f));
    // Hardtanh on [-1, 1]
    return fminf(fmaxf(g, -1.0f), 1.0f);
}

__global__ void fused_bias_act_vec4_kernel(float4* __restrict__ x,
                                           const float4* __restrict__ b,
                                           const float4* __restrict__ a,
                                           int ncols4,
                                           long long n4) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n4) return;
    int c = (int)(i % ncols4);
    float4 v  = x[i];
    float4 bb = b[c];
    float4 aa = a[c];
    v.x = act_chain(v.x + bb.x + aa.x);
    v.y = act_chain(v.y + bb.y + aa.y);
    v.z = act_chain(v.z + bb.z + aa.z);
    v.w = act_chain(v.w + bb.w + aa.w);
    x[i] = v;
}

__global__ void fused_bias_act_kernel(float* __restrict__ x,
                                      const float* __restrict__ b,
                                      const float* __restrict__ a,
                                      int ncols,
                                      long long n) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    int c = (int)(i % ncols);
    x[i] = act_chain(x[i] + b[c] + a[c]);
}

torch::Tensor fused_bias_activation(torch::Tensor x, torch::Tensor bias, torch::Tensor add_val) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto x_c = x.contiguous();
    auto b_c = bias.contiguous();
    auto a_c = add_val.contiguous();

    const long long n  = (long long)x_c.numel();
    const int ncols    = (int)a_c.numel();
    const int threads  = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    if ((n % 4 == 0) && (ncols % 4 == 0)) {
        const long long n4     = n / 4;
        const long long blocks = (n4 + threads - 1) / threads;
        fused_bias_act_vec4_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            reinterpret_cast<float4*>(x_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(b_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(a_c.data_ptr<float>()),
            ncols / 4, n4);
    } else {
        const long long blocks = (n + threads - 1) / threads;
        fused_bias_act_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            x_c.data_ptr<float>(), b_c.data_ptr<float>(), a_c.data_ptr<float>(),
            ncols, n);
    }
    return x_c;
}
"""

_fused_cpp_src = "torch::Tensor fused_bias_activation(torch::Tensor x, torch::Tensor bias, torch::Tensor add_val);"

_fused_ops = load_inline(
    name="fused_bias_swish_tanh_gelu_hardtanh_v1",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_bias_activation"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Two optimizations:
      1. The whole pointwise chain (bias add + add_value + Swish + Tanh + GELU
         + Hardtanh) is fused into a single custom CUDA kernel, eliminating
         5 intermediate kernel launches and their global-memory traffic.
      2. On Ampere+ GPUs the FP32 GEMM is evaluated with the accurate
         "3xTF32" tensor-core algorithm: A = Ah + Al, B = Bh + Bl with the low
         13 mantissa bits split off, then
             A@B ~= Ah@Bh + Ah@Bl + Al@Bh
         using TF32 tensor cores with FP32 accumulation. This is ~2.5x faster
         than a plain FP32 SGEMM while retaining FP32-level accuracy
         (the dropped Al@Bl term is ~2^-22 relative).
         The weight split is cached and invalidated if the weight changes.
    """

    def __init__(self, in_features, out_features, add_value_shape):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.add_value = nn.Parameter(torch.randn(add_value_shape))

        self._use_tf32 = (
            torch.cuda.is_available()
            and torch.cuda.get_device_capability(0)[0] >= 8
        )
        self._w_hi = None
        self._w_lo = None
        self._w_key = None

    def _reference_path(self, x):
        x = self.matmul(x)
        x = x + self.add_value
        x = torch.sigmoid(x) * x
        x = torch.tanh(x)
        x = F.gelu(x)
        x = F.hardtanh(x, min_val=-1.0, max_val=1.0)
        return x

    def _get_split_weight(self):
        w = self.matmul.weight.detach()
        key = (w.data_ptr(), w._version)
        if self._w_key != key:
            # Split off the low 13 mantissa bits so w_hi is exact in TF32.
            w_hi = (w.view(torch.int32) & -8192).view(torch.float32)
            w_lo = w - w_hi
            self._w_hi = w_hi
            self._w_lo = w_lo
            self._w_key = key
        return self._w_hi, self._w_lo

    def forward(self, x):
        # Exact reference fallback for CPU tensors / non-float32 inputs.
        if (not x.is_cuda) or x.dtype != torch.float32:
            return self._reference_path(x)

        x = x.contiguous()

        if self._use_tf32 and x.dim() == 2:
            w_hi, w_lo = self._get_split_weight()
            x_hi = (x.view(torch.int32) & -8192).view(torch.float32)
            x_lo = x - x_hi

            prev_tf32 = torch.backends.cuda.matmul.allow_tf32
            torch.backends.cuda.matmul.allow_tf32 = True
            try:
                out = torch.mm(x_hi, w_hi.t())
                out.addmm_(x_hi, w_lo.t())
                out.addmm_(x_lo, w_hi.t())
            finally:
                torch.backends.cuda.matmul.allow_tf32 = prev_tf32
        else:
            out = F.linear(x, self.matmul.weight)

        # Single fused kernel: + linear bias, + add_value, Swish, Tanh, GELU,
        # Hardtanh.
        return _fused_ops.fused_bias_activation(
            out, self.matmul.bias, self.add_value
        )
```