```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: (add_value broadcast) + Swish + Tanh + GELU(erf) + Hardtanh
# applied in-place on the GEMM+bias result produced by cuBLAS (at::addmm).
# This replaces 6 separate memory-bound elementwise kernels with a single
# vectorized (float4) pass, while keeping the matmul on the same cuBLAS
# path (and identical numerics) as nn.Linear.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float swish_tanh_gelu_hardtanh(float v) {
    // Swish: x * sigmoid(x)
    float s = v / (1.0f + expf(-v));
    // Tanh
    float t = tanhf(s);
    // GELU, exact erf formulation (matches PyTorch default approximate='none')
    float g = 0.5f * t * (1.0f + erff(t * 0.70710678118654752440f));
    // Hardtanh on [-1, 1] (same clamp semantics as PyTorch)
    g = g < -1.0f ? -1.0f : (g > 1.0f ? 1.0f : g);
    return g;
}

// Vectorized path: out_features % 4 == 0. Processes one float4 per thread.
__global__ void fused_bias_act_vec4_kernel(
    float* __restrict__ out,
    const float* __restrict__ addv,
    const long total4,
    const int cols4) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total4; idx += stride) {
        const int col4 = (int)(idx % cols4);
        const float4 a = reinterpret_cast<const float4*>(addv)[col4];
        float4 v = reinterpret_cast<float4*>(out)[idx];
        v.x = swish_tanh_gelu_hardtanh(v.x + a.x);
        v.y = swish_tanh_gelu_hardtanh(v.y + a.y);
        v.z = swish_tanh_gelu_hardtanh(v.z + a.z);
        v.w = swish_tanh_gelu_hardtanh(v.w + a.w);
        reinterpret_cast<float4*>(out)[idx] = v;
    }
}

// Generic scalar fallback for any shape.
__global__ void fused_bias_act_scalar_kernel(
    float* __restrict__ out,
    const float* __restrict__ addv,
    const long total,
    const int cols) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total; idx += stride) {
        const int col = (int)(idx % cols);
        out[idx] = swish_tanh_gelu_hardtanh(out[idx] + addv[col]);
    }
}

torch::Tensor fused_linear_act_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor add_value) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(weight.scalar_type() == at::kFloat);
    TORCH_CHECK(bias.scalar_type() == at::kFloat);
    TORCH_CHECK(add_value.scalar_type() == at::kFloat);

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

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b_c = bias.contiguous();
    auto a_c = add_value.contiguous();

    // GEMM + bias via cuBLAS, identical call path/numerics to nn.Linear.
    auto out = at::addmm(b_c, x_c, w_c.t());

    const long rows = out.size(0);
    const long cols = out.size(1);
    const long total = rows * cols;

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

    if ((cols % 4) == 0 && total > 0) {
        const long cols4 = cols / 4;
        const long total4 = total / 4;
        long blocks = (total4 + threads - 1) / threads;
        if (blocks > 65535L * 32L) blocks = 65535L * 32L;
        fused_bias_act_vec4_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            out.data_ptr<float>(), a_c.data_ptr<float>(), total4, (int)cols4);
    } else if (total > 0) {
        long blocks = (total + threads - 1) / threads;
        if (blocks > 65535L * 32L) blocks = 65535L * 32L;
        fused_bias_act_scalar_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            out.data_ptr<float>(), a_c.data_ptr<float>(), total, (int)cols);
    }

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_linear_act_forward("
    "torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor add_value);"
)

fused_linear_act = load_inline(
    name="fused_linear_act",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_linear_act_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM+bias (same as nn.Linear) followed by a single
    fused CUDA kernel performing (add_value) + Swish + Tanh + GELU + Hardtanh.
    """

    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:
            return fused_linear_act.fused_linear_act_forward(
                x, self.matmul.weight, self.matmul.bias, self.add_value
            )
        # Fallback (e.g. CPU tensors): original PyTorch ops.
        x = self.matmul(x)
        x = x + self.add_value
        x = torch.sigmoid(x) * x  # Swish
        x = torch.tanh(x)
        x = torch.nn.functional.gelu(x)  # GELU
        x = torch.nn.functional.hardtanh(x, min_val=-1, max_val=1)  # Hardtanh
        return x
```