```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
from torch.utils.cpp_extension import load_inline

# Custom CUDA implementation:
#  1) The Linear (GEMM + bias) is dispatched to cuBLAS via at::addmm, which is
#     already near-peak for this 1024x8192x8192 FP32 GEMM.
#  2) All five subsequent pointwise operators (add bias vector, Swish, Tanh,
#     exact GELU, Hardtanh) are fused into a single memory-bound CUDA kernel,
#     eliminating ~10 extra global-memory round trips of the 32MB activation
#     tensor that PyTorch would otherwise perform.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define INV_SQRT2 0.707106781186547524400844362104849039f

__device__ __forceinline__ float act_chain(float v, float addv) {
    // add broadcast bias vector
    v = v + addv;
    // Swish: v * sigmoid(v)
    float sig = 1.0f / (1.0f + expf(-v));
    v = v * sig;
    // Tanh
    v = tanhf(v);
    // Exact GELU (matches torch.nn.functional.gelu default): 0.5*v*(1+erf(v/sqrt(2)))
    v = 0.5f * v * (1.0f + erff(v * INV_SQRT2));
    // Hardtanh clamp to [-1, 1]
    v = fminf(fmaxf(v, -1.0f), 1.0f);
    return v;
}

// Vectorized (float4) path, used when the feature dim is divisible by 4.
__global__ void fused_epilogue_vec4_kernel(float* __restrict__ data,
                                           const float* __restrict__ addv,
                                           const int n4,
                                           const int tail,
                                           const int features) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 val = reinterpret_cast<float4*>(data)[i];
        int j = (i << 2) % features;
        const float4 av = *reinterpret_cast<const float4*>(addv + j);
        val.x = act_chain(val.x, av.x);
        val.y = act_chain(val.y, av.y);
        val.z = act_chain(val.z, av.z);
        val.w = act_chain(val.w, av.w);
        reinterpret_cast<float4*>(data)[i] = val;
    } else if (i - n4 < tail) {
        int idx = (n4 << 2) + (i - n4);
        int j = idx % features;
        data[idx] = act_chain(data[idx], addv[j]);
    }
}

// Scalar fallback path for arbitrary feature dims.
__global__ void fused_epilogue_scalar_kernel(float* __restrict__ data,
                                             const float* __restrict__ addv,
                                             const int n,
                                             const int features) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        int j = idx % features;
        data[idx] = act_chain(data[idx], addv[j]);
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor bias, torch::Tensor add_value) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() && add_value.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Only float32 is supported");

    auto xc = x.contiguous();
    // GEMM + linear bias via cuBLAS (weight.t() is handled natively, no copy).
    auto out = at::addmm(bias, xc, weight.t());
    auto addc = add_value.contiguous();

    const int n = static_cast<int>(out.numel());
    const int features = static_cast<int>(out.size(-1));

    if (n > 0) {
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        const int block = 256;
        if ((features & 3) == 0) {
            const int n4 = n >> 2;
            const int tail = n & 3;
            const int threads_needed = n4 + tail;
            const int blocks = (threads_needed + block - 1) / block;
            fused_epilogue_vec4_kernel<<<blocks, block, 0, stream>>>(
                out.data_ptr<float>(), addc.data_ptr<float>(), n4, tail, features);
        } else {
            const int blocks = (n + block - 1) / block;
            fused_epilogue_scalar_kernel<<<blocks, block, 0, stream>>>(
                out.data_ptr<float>(), addc.data_ptr<float>(), n, features);
        }
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, torch::Tensor add_value);"
)

fused_op = load_inline(
    name="fused_linear_swish_tanh_gelu_hardtanh",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM + a single fused CUDA kernel implementing
    (bias add -> Swish -> Tanh -> GELU -> Hardtanh) in one memory pass.
    Keeps identical parameter names/shapes to the reference Model so
    state_dict loading matches exactly.
    """
    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._fused_op = fused_op

    def _reference_forward(self, x):
        x = self.matmul(x)
        x = x + self.add_value
        x = torch.sigmoid(x) * x
        x = torch.tanh(x)
        x = torch.nn.functional.gelu(x)
        x = torch.nn.functional.hardtanh(x, min_val=-1, max_val=1)
        return x

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            return self._fused_op.fused_forward(
                x, self.matmul.weight, self.matmul.bias, self.add_value
            )
        # Safe fallback (CPU tensors / non-fp32 / higher-rank inputs).
        return self._reference_forward(x)
```