```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: bias-add + subtract + multiply + ReLU
# into a single memory-bound epilogue pass over the GEMM output.
# This replaces 4 separate PyTorch kernel launches (addmm bias epilogue /
# sub / mul / relu) with one fused pass, minimizing global memory traffic.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

// Vectorized (float4) path: each thread handles 4 consecutive columns.
__global__ void fused_bias_sub_mul_relu_vec4_kernel(
    const float4* __restrict__ x,
    const float4* __restrict__ bias,
    float4* __restrict__ out,
    long rows,
    int cols4,
    float sub,
    float mul)
{
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= cols4) return;

    float4 bv = bias[c];
    for (long r = blockIdx.y; r < rows; r += gridDim.y) {
        long idx = r * (long)cols4 + c;
        float4 xv = x[idx];
        float4 res;
        res.x = fmaxf((xv.x + bv.x - sub) * mul, 0.0f);
        res.y = fmaxf((xv.y + bv.y - sub) * mul, 0.0f);
        res.z = fmaxf((xv.z + bv.z - sub) * mul, 0.0f);
        res.w = fmaxf((xv.w + bv.w - sub) * mul, 0.0f);
        out[idx] = res;
    }
}

// Scalar fallback path for column counts not divisible by 4.
__global__ void fused_bias_sub_mul_relu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long rows,
    int cols,
    float sub,
    float mul)
{
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= cols) return;

    float b = bias[c];
    for (long r = blockIdx.y; r < rows; r += gridDim.y) {
        long idx = r * (long)cols + c;
        float v = (x[idx] + b - sub) * mul;
        out[idx] = v > 0.0f ? v : 0.0f;
    }
}

torch::Tensor fused_bias_sub_mul_relu_cuda(
    torch::Tensor x,
    torch::Tensor bias,
    double subtract_value,
    double multiply_value)
{
    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(x.dim() == 2, "x must be 2D");

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

    long rows = x_c.size(0);
    long cols = x_c.size(1);
    if (rows == 0 || cols == 0) return out;

    const int block = 256;
    float sub = (float)subtract_value;
    float mul = (float)multiply_value;
    unsigned int grid_y = (unsigned int)(rows < 65535 ? rows : 65535);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (cols % 4 == 0) {
        int cols4 = (int)(cols / 4);
        unsigned int grid_x = (unsigned int)((cols4 + block - 1) / block);
        dim3 grid(grid_x, grid_y);
        fused_bias_sub_mul_relu_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias_c.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            rows, cols4, sub, mul);
    } else {
        unsigned int grid_x = (unsigned int)((cols + block - 1) / block);
        dim3 grid(grid_x, grid_y);
        fused_bias_sub_mul_relu_kernel<<<grid, block, 0, stream>>>(
            x_c.data_ptr<float>(),
            bias_c.data_ptr<float>(),
            out.data_ptr<float>(),
            rows, (int)cols, sub, mul);
    }

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_bias_sub_mul_relu_cuda(torch::Tensor x, torch::Tensor bias, "
    "double subtract_value, double multiply_value);"
)

# Compile the inline CUDA fused epilogue
fused_epilogue = load_inline(
    name="fused_bias_sub_mul_relu",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_bias_sub_mul_relu_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      - The GEMM (x @ W^T) is delegated to cuBLAS via torch.matmul (highly tuned).
      - bias-add + subtract + multiply + ReLU are fused into a single custom
        CUDA epilogue kernel, eliminating 3 extra kernel launches and 3 extra
        full read/write passes over the (batch, out_features) activation.
    """
    def __init__(self, in_features, out_features, subtract_value, multiply_value):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.subtract_value = float(subtract_value)
        self.multiply_value = float(multiply_value)
        self.fused_epilogue = fused_epilogue

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            # GEMM without bias; bias is fused into the custom epilogue.
            mm = torch.matmul(x, self.linear.weight.t())
            return self.fused_epilogue.fused_bias_sub_mul_relu_cuda(
                mm,
                self.linear.bias,
                self.subtract_value,
                self.multiply_value,
            )
        # Fallback (CPU / non-float32) path: original PyTorch ops.
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```