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

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>

#define CUBLASLT_CHECK(expr)                                                 \
  do {                                                                       \
    cublasStatus_t _st = (expr);                                             \
    TORCH_CHECK(_st == CUBLAS_STATUS_SUCCESS, "cublasLt error: ", (int)_st); \
  } while (0)

// bias_prime = (bias - sub) * mul
__global__ void bias_prime_kernel(const float* __restrict__ bias,
                                  float* __restrict__ out,
                                  float sub, float mul, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) out[i] = (bias[i] - sub) * mul;
}

// out = relu((y - sub) * mul), vectorized
__global__ void fused_epilogue_vec4_kernel(const float4* __restrict__ y,
                                           float4* __restrict__ out,
                                           float sub, float mul, long n4) {
  long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
  long stride = (long)gridDim.x * blockDim.x;
  for (; i < n4; i += stride) {
    float4 v = y[i];
    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);
    out[i] = v;
  }
}

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

static cublasLtHandle_t get_lt_handle() {
  static cublasLtHandle_t handle = nullptr;
  if (handle == nullptr) {
    CUBLASLT_CHECK(cublasLtCreate(&handle));
  }
  return handle;
}

// Fully fused: out = relu(mul * (x @ W^T) + (bias - sub) * mul)
// The whole linear + subtract + multiply + relu collapses into a single
// cuBLASLt GEMM with a fused RELU_BIAS epilogue (alpha = mul does the
// scaling for free inside the GEMM).
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(),
              "CUDA tensors required");
  TORCH_CHECK(x.scalar_type() == at::kFloat &&
              weight.scalar_type() == at::kFloat &&
              bias.scalar_type() == at::kFloat, "float32 required");

  auto x_c = x.contiguous();
  auto w_c = weight.contiguous();
  auto b_c = bias.contiguous();

  const int64_t M = x_c.size(0);
  const int64_t K = x_c.size(1);
  const int64_t N = w_c.size(0);
  TORCH_CHECK(w_c.size(1) == K && b_c.numel() == N, "shape mismatch");

  auto out = torch::empty({M, N}, x_c.options());
  auto bias_prime = torch::empty({N}, x_c.options());

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

  {
    const int threads = 256;
    const int blocks = (int)((N + threads - 1) / threads);
    bias_prime_kernel<<<blocks, threads, 0, stream>>>(
        b_c.data_ptr<float>(), bias_prime.data_ptr<float>(),
        (float)sub, (float)mul, (int)N);
  }

  cublasLtHandle_t lt = get_lt_handle();

  cublasLtMatmulDesc_t opDesc;
  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F));
  // Row-major torch buffers reinterpreted as column-major:
  //   W[N,K] row-major == (W^T)[K,N] col-major -> opA = T gives W (N x K)
  //   X[M,K] row-major == (X^T)[K,M] col-major -> opB = N
  //   OUT[M,N] row-major == (OUT^T)[N,M] col-major
  //   OUT^T = W * X^T  <=>  OUT = X * W^T
  cublasOperation_t opA = CUBLAS_OP_T;
  cublasOperation_t opB = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opA, sizeof(opA)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opB, sizeof(opB)));
  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      opDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
  const float* bias_ptr = bias_prime.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      opDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));

  cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_32F, K, N, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_32F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N));

  const size_t workspace_size = 32ull * 1024ull * 1024ull;
  auto workspace = torch::empty({(int64_t)workspace_size},
                                x_c.options().dtype(torch::kUInt8));

  cublasLtMatmulPreference_t pref;
  CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
  CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
      pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
      &workspace_size, sizeof(workspace_size)));

  cublasLtMatmulHeuristicResult_t heuristic;
  int returned = 0;
  CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
      lt, opDesc, Adesc, Bdesc, Cdesc, Cdesc, pref, 1, &heuristic, &returned));
  TORCH_CHECK(returned > 0, "cublasLt: no algorithm found");

  float alpha = (float)mul;
  float beta = 0.0f;

  cublasStatus_t st = cublasLtMatmul(
      lt, opDesc, &alpha,
      w_c.data_ptr<float>(), Adesc,
      x_c.data_ptr<float>(), Bdesc,
      &beta,
      out.data_ptr<float>(), Cdesc,
      out.data_ptr<float>(), Cdesc,
      &heuristic.algo,
      workspace.data_ptr(), workspace_size,
      stream);

  cublasLtMatmulPreferenceDestroy(pref);
  cublasLtMatrixLayoutDestroy(Adesc);
  cublasLtMatrixLayoutDestroy(Bdesc);
  cublasLtMatrixLayoutDestroy(Cdesc);
  cublasLtMatmulDescDestroy(opDesc);

  CUBLASLT_CHECK(st);
  return out;
}

// Fallback fused elementwise op: out = relu((y - sub) * mul)
torch::Tensor fused_sub_mul_relu(torch::Tensor y, double sub, double mul) {
  TORCH_CHECK(y.is_cuda() && y.scalar_type() == at::kFloat);
  auto y_c = y.contiguous();
  auto out = torch::empty_like(y_c);
  const long n = (long)y_c.numel();
  if (n == 0) return out;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int threads = 256;
  if (n % 4 == 0) {
    const long n4 = n / 4;
    long blocks = (n4 + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    fused_epilogue_vec4_kernel<<<(int)blocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(y_c.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()),
        (float)sub, (float)mul, n4);
  } else {
    long blocks = (n + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    fused_epilogue_scalar_kernel<<<(int)blocks, threads, 0, stream>>>(
        y_c.data_ptr<float>(), out.data_ptr<float>(), (float)sub, (float)mul, n);
  }
  return out;
}
"""

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

_ext = load_inline(
    name="fused_linear_sub_mul_relu_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_linear_sub_mul_relu", "fused_sub_mul_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublasLt", "-lcublas"],
)


class ModelNew(nn.Module):
    """
    Optimized model: the matrix multiply, bias add, subtraction,
    multiplication and ReLU are fused into a single custom CUDA operator
    (cuBLASLt GEMM with a fused RELU_BIAS epilogue, where the scalar
    multiply is folded into the GEMM alpha and the subtraction is folded
    into the epilogue bias). A fused elementwise CUDA kernel is kept as a
    fallback path.
    """
    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._use_lt = True

    def forward(self, x):
        weight = self.linear.weight
        bias = self.linear.bias
        if x.is_cuda:
            if self._use_lt:
                try:
                    return _ext.fused_linear_sub_mul_relu(
                        x, weight, bias,
                        self.subtract_value, self.multiply_value)
                except RuntimeError:
                    self._use_lt = False
            y = F.linear(x, weight, bias)
            return _ext.fused_sub_mul_relu(
                y, self.subtract_value, self.multiply_value)
        y = F.linear(x, weight, bias)
        return torch.relu((y - self.subtract_value) * self.multiply_value)
```