```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:
#   out = Hardtanh(GELU(Tanh(Swish(matmul_out + linear_bias + add_value))))
# The GEMM itself is left to cuBLAS (torch.matmul), which is optimal for the
# 1024x8192 @ 8192x8192 problem. Everything after the GEMM (bias add, residual
# parameter add, Swish, Tanh, GELU, Hardtanh) is fused into a single
# memory-bound kernel, eliminating ~6 full passes over the 32MB activation
# tensor and all intermediate kernel launches.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__device__ __forceinline__ float fused_act(float z) {
    // Swish: sigmoid(z) * z
    float sig = 1.0f / (1.0f + expf(-z));
    float sw = sig * z;
    // Tanh
    float t = tanhf(sw);
    // GELU (exact erf formulation, matches PyTorch's 'none' approximation)
    float g = t * 0.5f * (1.0f + erff(t * 0.7071067811865476f));
    // Hardtanh on [-1, 1]
    return fminf(fmaxf(g, -1.0f), 1.0f);
}

// Vectorized path: one float4 per thread, 2D grid (row, col-chunk) so no
// modulo is needed to find the bias column.
__global__ void fused_bias_act_vec4_kernel(const float4* __restrict__ y,
                                           const float4* __restrict__ bias,
                                           const float4* __restrict__ addv,
                                           float4* __restrict__ out,
                                           const int cols4) {
    int col4 = blockIdx.x * blockDim.x + threadIdx.x;
    if (col4 >= cols4) return;
    long long idx = (long long)blockIdx.y * cols4 + col4;
    float4 yv = y[idx];
    float4 bv = __ldg(&bias[col4]);
    float4 av = __ldg(&addv[col4]);
    float4 o;
    o.x = fused_act(yv.x + bv.x + av.x);
    o.y = fused_act(yv.y + bv.y + av.y);
    o.z = fused_act(yv.z + bv.z + av.z);
    o.w = fused_act(yv.w + bv.w + av.w);
    out[idx] = o;
}

// Generic scalar fallback path (any shape / alignment).
__global__ void fused_bias_act_scalar_kernel(const float* __restrict__ y,
                                             const float* __restrict__ bias,
                                             const float* __restrict__ addv,
                                             float* __restrict__ out,
                                             const int cols,
                                             const long long n) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    int col = (int)(idx % (long long)cols);
    out[idx] = fused_act(y[idx] + __ldg(&bias[col]) + __ldg(&addv[col]));
}

torch::Tensor fused_forward(torch::Tensor y, torch::Tensor bias, torch::Tensor addv) {
    TORCH_CHECK(y.is_cuda() && bias.is_cuda() && addv.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "fused_forward expects float32");

    auto yc = y.contiguous();
    auto bc = bias.contiguous();
    auto ac = addv.contiguous();
    auto out = torch::empty_like(yc);

    const long long n = yc.numel();
    if (n == 0) return out;
    const int cols = (int)yc.size(-1);
    const int block = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    bool aligned = ((reinterpret_cast<uintptr_t>(yc.data_ptr<float>()) & 15) == 0) &&
                   ((reinterpret_cast<uintptr_t>(bc.data_ptr<float>()) & 15) == 0) &&
                   ((reinterpret_cast<uintptr_t>(ac.data_ptr<float>()) & 15) == 0) &&
                   ((reinterpret_cast<uintptr_t>(out.data_ptr<float>()) & 15) == 0);

    const long long rows = n / cols;
    if (cols % 4 == 0 && aligned && rows <= 65535) {
        const int cols4 = cols / 4;
        dim3 grid((cols4 + block - 1) / block, (unsigned int)rows);
        fused_bias_act_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(yc.data_ptr<float>()),
            reinterpret_cast<const float4*>(bc.data_ptr<float>()),
            reinterpret_cast<const float4*>(ac.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            cols4);
    } else {
        const long long grid = (n + block - 1) / block;
        fused_bias_act_scalar_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            yc.data_ptr<float>(),
            bc.data_ptr<float>(),
            ac.data_ptr<float>(),
            out.data_ptr<float>(),
            cols, n);
    }
    return out;
}
"""

cpp_source = "torch::Tensor fused_forward(torch::Tensor y, torch::Tensor bias, torch::Tensor addv);"

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


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM (without bias) followed by a single fused CUDA
    kernel that performs bias add + parameter add + Swish + Tanh + GELU +
    Hardtanh in one pass over the data.
    """
    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:
            # GEMM via cuBLAS; bias is folded into the fused epilogue kernel.
            y = torch.matmul(x, self.matmul.weight.t())
            return _fused_epilogue.fused_forward(y, self.matmul.bias, self.add_value)
        # Fallback to the reference implementation for non-CUDA / non-fp32.
        x = self.matmul(x)
        x = x + self.add_value
        x = torch.sigmoid(x) * x  # Swish
        x = torch.tanh(x)
        x = F.gelu(x)  # GELU
        x = F.hardtanh(x, min_val=-1, max_val=1)  # Hardtanh
        return x
```