```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: fuses the entire post-linear epilogue
# (bias add via cuBLAS addmm, then subtract + multiply + ReLU in one
# vectorized pass) into a single kernel launch.
fused_epilogue_cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_sub_mul_relu_vec4_kernel(float4* __restrict__ data,
                                               float sub,
                                               float mul,
                                               int n4) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n4) {
        float4 v = data[idx];
        v.x = fmaxf((v.x - sub) * mul, 0.0f);
        v.y = fmaxf((v.y - sub) * mul, 0.0f);
        v.z = fmaxf((v.z - sub) * mul, 0.0f);
        v.w = fmaxf((v.w - sub) * mul, 0.0f);
        data[idx] = v;
    }
}

__global__ void fused_sub_mul_relu_scalar_kernel(float* __restrict__ data,
                                                 float sub,
                                                 float mul,
                                                 int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float v = data[idx];
        data[idx] = fmaxf((v - sub) * mul, 0.0f);
    }
}

torch::Tensor fused_linear_sub_mul_relu(torch::Tensor x,
                                        torch::Tensor weight,
                                        torch::Tensor bias,
                                        double sub,
                                        double mul) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                weight.scalar_type() == torch::kFloat32 &&
                bias.scalar_type() == torch::kFloat32,
                "Only float32 is supported");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");

    // GEMM + bias fused by cuBLAS (equivalent to nn.Linear).
    auto out = torch::addmm(bias, x, weight.t());

    const int64_t n = out.numel();
    const float s = static_cast<float>(sub);
    const float m = static_cast<float>(mul);
    const int block = 256;

    if (n % 4 == 0) {
        const int64_t n4 = n / 4;
        const int64_t grid = (n4 + block - 1) / block;
        fused_sub_mul_relu_vec4_kernel<<<static_cast<unsigned int>(grid), block>>>(
            reinterpret_cast<float4*>(out.data_ptr<float>()), s, m,
            static_cast<int>(n4));
    } else {
        const int64_t grid = (n + block - 1) / block;
        fused_sub_mul_relu_scalar_kernel<<<static_cast<unsigned int>(grid), block>>>(
            out.data_ptr<float>(), s, m, static_cast<int>(n));
    }

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_linear_sub_mul_relu(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double sub, double mul);"
)

fused_epilogue = load_inline(
    name="fused_linear_sub_mul_relu",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_linear_sub_mul_relu"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the matmul + bias is executed by cuBLAS (addmm) and the
    subtract / multiply / ReLU epilogue is fused into a single custom
    vectorized CUDA kernel, eliminating three separate elementwise kernel
    launches and their intermediate memory traffic.
    """
    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:
            return self.fused_epilogue.fused_linear_sub_mul_relu(
                x.contiguous(),
                self.linear.weight,
                self.linear.bias,
                self.subtract_value,
                self.multiply_value,
            )
        # Fallback for CPU / non-float32 inputs.
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```