```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 CUDA kernel: (add_value + Swish + Tanh + GELU(erf) + Hardtanh) in a
# single memory pass over the GEMM output. The GEMM itself (with the linear
# bias) is handled by cuBLAS via F.linear, so the whole forward is just
# 2 kernel launches instead of ~7.
fused_bias_act_cpp_source = (
    "torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor bias);"
)

fused_bias_act_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BLOCK_SIZE 256

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

template <bool POW2>
__global__ void fused_bias_act_vec4_kernel(const float4* __restrict__ x,
                                           const float4* __restrict__ bias,
                                           float4* __restrict__ out,
                                           unsigned int n4,
                                           unsigned int row4,
                                           unsigned int mask) {
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n4) return;
    // row is a multiple of 4, so a float4 never straddles two rows.
    unsigned int c = POW2 ? (i & mask) : (i % row4);
    float4 xv = x[i];
    float4 bv = bias[c];
    float4 ov;
    ov.x = act_chain(xv.x + bv.x);
    ov.y = act_chain(xv.y + bv.y);
    ov.z = act_chain(xv.z + bv.z);
    ov.w = act_chain(xv.w + bv.w);
    out[i] = ov;
}

__global__ void fused_bias_act_scalar_kernel(const float* __restrict__ x,
                                             const float* __restrict__ bias,
                                             float* __restrict__ out,
                                             long long total,
                                             long long row) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)gridDim.x * blockDim.x;
    for (; i < total; i += stride) {
        out[i] = act_chain(x[i] + bias[i % row]);
    }
}

torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor bias) {
    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");

    const c10::cuda::CUDAGuard device_guard(x.device());

    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto out = torch::empty_like(xc);

    long long total = xc.numel();
    long long row = xc.size(-1);
    TORCH_CHECK(bc.numel() == row, "bias must match the last dimension of x");
    if (total == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    bool aligned = (reinterpret_cast<uintptr_t>(xp) % 16 == 0) &&
                   (reinterpret_cast<uintptr_t>(bp) % 16 == 0) &&
                   (reinterpret_cast<uintptr_t>(op) % 16 == 0);

    if (aligned && (row % 4 == 0) && ((unsigned long long)(total / 4) < 4000000000ULL)) {
        unsigned int n4 = (unsigned int)(total / 4);
        unsigned int row4 = (unsigned int)(row / 4);
        unsigned int grid = (n4 + BLOCK_SIZE - 1) / BLOCK_SIZE;
        if ((row4 & (row4 - 1)) == 0) {
            fused_bias_act_vec4_kernel<true><<<grid, BLOCK_SIZE, 0, stream>>>(
                reinterpret_cast<const float4*>(xp),
                reinterpret_cast<const float4*>(bp),
                reinterpret_cast<float4*>(op), n4, row4, row4 - 1);
        } else {
            fused_bias_act_vec4_kernel<false><<<grid, BLOCK_SIZE, 0, stream>>>(
                reinterpret_cast<const float4*>(xp),
                reinterpret_cast<const float4*>(bp),
                reinterpret_cast<float4*>(op), n4, row4, 0);
        }
    } else {
        long long want = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;
        int grid = (int)(want < 65535 ? want : 65535);
        fused_bias_act_scalar_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
            xp, bp, op, total, row);
    }
    return out;
}
"""

fused_bias_act_ops = load_inline(
    name="fused_bias_act_ext",
    cpp_sources=fused_bias_act_cpp_source,
    cuda_sources=fused_bias_act_cuda_source,
    functions=["fused_bias_act"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      * GEMM + linear bias via cuBLAS (F.linear / addmm epilogue).
      * add_value + Swish + Tanh + GELU + Hardtanh fused into a single
        vectorized CUDA kernel (one read + one write of the activations
        instead of six separate elementwise passes).
    """
    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))

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            x = F.linear(x, self.matmul.weight, self.matmul.bias)
            return fused_bias_act_ops.fused_bias_act(x, self.add_value)
        # Fallback (CPU / non-float32): reference implementation.
        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, max_val=1)
        return x
```