```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 operator:
#   1) linear (matmul + bias) is dispatched to highly-optimized cuBLAS via at::addmm
#   2) the subtract -> multiply -> ReLU chain is fused into a single custom
#      vectorized (float4) in-place CUDA kernel, eliminating 2 extra kernel
#      launches and 2 full read/write passes over the activation tensor.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void fused_sub_mul_relu_kernel(float* __restrict__ data,
                                          const float sub,
                                          const float mul,
                                          const long n) {
    const long n4 = n >> 2;  // number of float4 chunks
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n4) {
        float4* d4 = reinterpret_cast<float4*>(data);
        float4 v = d4[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);
        d4[idx] = v;
    } else if (idx == n4) {
        // handle the tail (n % 4 elements) with a single thread
        for (long i = (n4 << 2); i < n; ++i) {
            float v = data[i];
            data[i] = 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() == at::kFloat &&
                weight.scalar_type() == at::kFloat &&
                bias.scalar_type() == at::kFloat,
                "All tensors must be float32");
    TORCH_CHECK(x.size(-1) == weight.size(1), "x and weight dims mismatch");
    TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias dims mismatch");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    // Flatten leading dims so any batched input works with addmm.
    auto x2 = (xc.dim() == 2) ? xc : xc.reshape({-1, xc.size(-1)});

    // cuBLAS SGEMM with fused bias: out = x2 @ W^T + b
    auto out = at::addmm(bc, x2, wc.t());

    // Fused epilogue: out = relu((out - sub) * mul), in-place, one kernel.
    const long n = out.numel();
    if (n > 0) {
        const long n4 = n >> 2;
        const long tail_thread = (n & 3) ? 1 : 0;
        const long threads_needed = n4 + tail_thread;
        const int block = 256;
        const long blocks = (threads_needed + block - 1) / block;
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        fused_sub_mul_relu_kernel<<<(unsigned int)blocks, block, 0, stream>>>(
            out.data_ptr<float>(),
            static_cast<float>(sub),
            static_cast<float>(mul),
            n);
    }

    if (xc.dim() != 2) {
        auto sizes = xc.sizes().vec();
        sizes.back() = out.size(-1);
        out = out.reshape(sizes);
    }
    return out;
}
"""

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

fused_linear_op = load_inline(
    name="fused_linear_sub_mul_relu",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_linear_sub_mul_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS linear + single fused custom CUDA kernel for
    subtract / multiply / ReLU.
    """
    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_op = fused_linear_op

    def forward(self, x):
        return self.fused_op.fused_linear_sub_mul_relu(
            x,
            self.linear.weight,
            self.linear.bias,
            self.subtract_value,
            self.multiply_value,
        )
```