```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

# Fused CUDA kernel that applies (x - subtract_value) * multiply_value followed by ReLU
# in a single memory pass (in-place), instead of three separate elementwise kernels.
# The matmul + bias itself is left to the highly optimized cuBLAS path via nn.Linear,
# which is compute-bound and hard to beat; the elementwise chain is memory-bound and
# benefits greatly from fusion + vectorized (float4) memory access.
fused_epilogue_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__global__ void sub_mul_relu_scalar_kernel(float* __restrict__ data,
                                           const float sub,
                                           const float mul,
                                           const long n) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)gridDim.x * blockDim.x;
    for (; idx < n; idx += stride) {
        float v = (data[idx] - sub) * mul;
        data[idx] = v > 0.0f ? v : 0.0f;
    }
}

__global__ void sub_mul_relu_vec4_kernel(float4* __restrict__ data,
                                         const float sub,
                                         const float mul,
                                         const long n4) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long stride = (long)gridDim.x * blockDim.x;
    for (; idx < n4; idx += stride) {
        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;
    }
}

torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul) {
    TORCH_CHECK(x.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Input must be float32");
    TORCH_CHECK(x.is_contiguous(), "Input must be contiguous");

    const long n = x.numel();
    if (n == 0) {
        return x;
    }

    const float sub_f = static_cast<float>(sub);
    const float mul_f = static_cast<float>(mul);
    auto stream = at::cuda::getCurrentCUDAStream();
    const int block_size = 256;

    // 16-byte aligned fast path using float4 vectorized loads/stores.
    if ((n % 4) == 0 && (reinterpret_cast<uintptr_t>(x.data_ptr<float>()) % 16) == 0) {
        const long n4 = n / 4;
        long grid = (n4 + block_size - 1) / block_size;
        if (grid > 65535L) grid = 65535L;
        sub_mul_relu_vec4_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
            reinterpret_cast<float4*>(x.data_ptr<float>()), sub_f, mul_f, n4);
    } else {
        long grid = (n + block_size - 1) / block_size;
        if (grid > 65535L) grid = 65535L;
        sub_mul_relu_scalar_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
            x.data_ptr<float>(), sub_f, mul_f, n);
    }

    return x;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul);"
)

fused_epilogue = load_inline(
    name="fused_sub_mul_relu",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_source,
    functions=["sub_mul_relu_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS linear followed by a single fused CUDA kernel that
    performs subtraction, multiplication and ReLU in one in-place pass.
    """
    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 = subtract_value
        self.multiply_value = multiply_value
        self.fused_epilogue = fused_epilogue

    def forward(self, x):
        x = self.linear(x)
        x = self.fused_epilogue.sub_mul_relu_cuda(
            x, self.subtract_value, self.multiply_value
        )
        return x
```