```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 kernel that fuses the entire element-wise chain:
#   (+ add_value) -> Swish -> Tanh -> GELU (erf) -> Hardtanh
# This eliminates 5 separate global-memory round-trips over the
# (batch_size, out_features) activation tensor and replaces them
# with a single fused, vectorized (float4) pass.
fused_act_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

namespace {

constexpr float kGeluAlpha = 0.70710678118654752440f;  // 1/sqrt(2)

__device__ __forceinline__ float act_chain(float v) {
    // Swish: x * sigmoid(x)
    float sig = 1.0f / (1.0f + expf(-v));
    float s = sig * v;
    // Tanh
    float t = tanhf(s);
    // GELU (exact, erf-based), matching F.gelu default
    float g = (t * 0.5f) * (1.0f + erff(t * kGeluAlpha));
    // Hardtanh in [-1, 1]
    g = fminf(fmaxf(g, -1.0f), 1.0f);
    return g;
}

__global__ void fused_act_kernel(const float* __restrict__ x,
                                 const float* __restrict__ bias,
                                 float* __restrict__ out,
                                 long total,
                                 int ncols) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int col = (int)(idx % ncols);
        out[idx] = act_chain(x[idx] + bias[col]);
    }
}

__global__ void fused_act_vec4_kernel(const float4* __restrict__ x,
                                      const float4* __restrict__ bias,
                                      float4* __restrict__ out,
                                      long total4,
                                      int ncols4) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total4) {
        int col4 = (int)(idx % ncols4);
        float4 xv = x[idx];
        float4 bv = bias[col4];
        float4 r;
        r.x = act_chain(xv.x + bv.x);
        r.y = act_chain(xv.y + bv.y);
        r.z = act_chain(xv.z + bv.z);
        r.w = act_chain(xv.w + bv.w);
        out[idx] = r;
    }
}

}  // namespace

torch::Tensor fused_act_cuda(torch::Tensor x, torch::Tensor bias) {
    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));
    TORCH_CHECK(x.is_cuda() && bias.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");

    auto x_c = x.contiguous();
    auto b_c = bias.contiguous();
    auto out = torch::empty_like(x_c);

    long total = x_c.numel();
    int ncols = (int)x_c.size(-1);
    TORCH_CHECK(b_c.numel() == ncols, "bias must match last dim of x");

    if (total == 0) {
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;

    if (ncols % 4 == 0) {
        long total4 = total / 4;
        int ncols4 = ncols / 4;
        long blocks = (total4 + block - 1) / block;
        fused_act_vec4_kernel<<<blocks, block, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(b_c.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            total4, ncols4);
    } else {
        long blocks = (total + block - 1) / block;
        fused_act_kernel<<<blocks, block, 0, stream>>>(
            x_c.data_ptr<float>(), b_c.data_ptr<float>(),
            out.data_ptr<float>(), total, ncols);
    }
    return out;
}
"""

fused_act_cpp_source = (
    "torch::Tensor fused_act_cuda(torch::Tensor x, torch::Tensor bias);"
)

fused_act_ext = load_inline(
    name="fused_act_ext",
    cpp_sources=fused_act_cpp_source,
    cuda_sources=fused_act_cuda_source,
    functions=["fused_act_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      - matmul + Linear bias is dispatched to cuBLAS via torch.addmm
        (fused beta epilogue, no extra memory pass).
      - The remaining chain (+ add_value, Swish, Tanh, GELU, Hardtanh)
        is executed by a single fused custom CUDA kernel.
    """
    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_act = fused_act_ext

    def forward(self, x):
        in_shape = x.shape
        x2 = x.reshape(-1, in_shape[-1])
        # cuBLAS GEMM with fused bias epilogue
        y = torch.addmm(self.matmul.bias, x2, self.matmul.weight.t())
        # One fused kernel for: + add_value, Swish, Tanh, GELU, Hardtanh
        y = self.fused_act.fused_act_cuda(y, self.add_value)
        return y.reshape(*in_shape[:-1], y.shape[-1])
```