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

# Custom fused CUDA operator:
#   out = relu( mul * (x @ weight^T + bias - sub) )
#       = relu( mul * (x @ weight^T) + mul * (bias - sub) )
# computed with a single cublasLt GEMM using the RELU_BIAS epilogue
# (alpha = mul, bias vector = mul * (bias - sub) built by a tiny CUDA kernel).
# A vectorized fused (sub, mul, relu) epilogue kernel is also provided and
# used as a fallback on top of the regular cuBLAS linear.
_cpp_src = """
torch::Tensor linear_sub_mul_relu(torch::Tensor x, torch::Tensor weight,
                                  torch::Tensor bias, double sub, double mul,
                                  bool tf32);
torch::Tensor sub_mul_relu(torch::Tensor y, double sub, double mul);
"""

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

#define CUBLASLT_CHECK(expr)                                                   \
  do {                                                                         \
    cublasStatus_t _status = (expr);                                           \
    TORCH_CHECK(_status == CUBLAS_STATUS_SUCCESS, "cublasLt call failed: ",    \
                #expr, " status=", (int)_status);                              \
  } while (0)

namespace {

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

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

__global__ void sub_mul_relu_vec4_kernel(const float4* __restrict__ in,
                                         float4* __restrict__ out, long n4,
                                         float sub, float mul) {
  const long stride = (long)gridDim.x * blockDim.x;
  for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < n4;
       i += stride) {
    float4 v = in[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 sub_mul_relu_kernel(const float* __restrict__ in,
                                    float* __restrict__ out, long n, float sub,
                                    float mul) {
  const long stride = (long)gridDim.x * blockDim.x;
  for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < n;
       i += stride) {
    out[i] = fmaxf((in[i] - sub) * mul, 0.0f);
  }
}

}  // namespace

torch::Tensor linear_sub_mul_relu(torch::Tensor x, torch::Tensor weight,
                                  torch::Tensor bias, double sub, double mul,
                                  bool tf32) {
  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,
              "FP32 tensors required");
  TORCH_CHECK(x.dim() >= 2, "x must have at least 2 dims");
  TORCH_CHECK(weight.dim() == 2, "weight must be 2D");

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

  const int64_t K = xc.size(-1);
  const int64_t N = wc.size(0);
  TORCH_CHECK(wc.size(1) == K, "weight shape mismatch");
  TORCH_CHECK(bc.numel() == N, "bias shape mismatch");

  auto x2 = xc.reshape({-1, K});
  const int64_t M = x2.size(0);

  auto out_sizes = xc.sizes().vec();
  out_sizes.back() = N;
  auto out = at::empty({M, N}, xc.options());
  if (M == 0 || N == 0) {
    return out.reshape(out_sizes);
  }
  TORCH_CHECK(K > 0 && K <= INT32_MAX && M <= INT32_MAX && N <= INT32_MAX,
              "unsupported dims");

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

  // bias_eff = (bias - sub) * mul   (length N, one per output feature)
  auto bias_eff = at::empty({N}, xc.options());
  {
    const int threads = 256;
    const int blocks = (int)((N + threads - 1) / threads);
    bias_transform_kernel<<<blocks, threads, 0, stream>>>(
        bc.data_ptr<float>(), bias_eff.data_ptr<float>(), (int)N, (float)sub,
        (float)mul);
  }

  const cublasComputeType_t compute_type =
      tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;

  cublasLtMatmulDesc_t op_desc = nullptr;
  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&op_desc, compute_type, CUDA_R_32F));

  // Row-major C(M,N) = X(M,K) * W^T(K,N) is computed as the equivalent
  // column-major problem C^T(N,M) = W(N,K) * X^T(K,M):
  //   A = weight buffer viewed as col-major (K,N), op(A) = A^T = W (N,K)
  //   B = x buffer      viewed as col-major (K,M), op(B) = X^T (K,M)
  cublasOperation_t transa = CUBLAS_OP_T;
  cublasOperation_t transb = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)));

  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
  float* bias_ptr = bias_eff.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));

  cublasLtMatrixLayout_t a_desc = nullptr, b_desc = nullptr, c_desc = nullptr;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_desc, CUDA_R_32F, K, N, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_32F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&c_desc, CUDA_R_32F, N, M, N));

  using Key = std::tuple<int64_t, int64_t, int64_t, bool>;
  static std::map<Key, cublasLtMatmulHeuristicResult_t> algo_cache;
  const Key key = std::make_tuple(M, N, K, tf32);

  cublasLtMatmulHeuristicResult_t heuristic;
  auto it = algo_cache.find(key);
  if (it != algo_cache.end()) {
    heuristic = it->second;
  } else {
    cublasLtMatmulPreference_t pref = nullptr;
    CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
    size_t max_workspace = 32ull * 1024ull * 1024ull;
    CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
        pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace,
        sizeof(max_workspace)));
    int returned = 0;
    CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
        getLtHandle(), op_desc, a_desc, b_desc, c_desc, c_desc, pref, 1,
        &heuristic, &returned));
    cublasLtMatmulPreferenceDestroy(pref);
    TORCH_CHECK(returned > 0, "cublasLt found no suitable algorithm");
    algo_cache.emplace(key, heuristic);
  }

  auto workspace = at::empty({(int64_t)heuristic.workspaceSize},
                             xc.options().dtype(at::kByte));

  const float alpha = (float)mul;
  const float beta = 0.0f;
  CUBLASLT_CHECK(cublasLtMatmul(
      getLtHandle(), op_desc, &alpha,
      wc.data_ptr<float>(), a_desc,
      x2.data_ptr<float>(), b_desc,
      &beta,
      out.data_ptr<float>(), c_desc,
      out.data_ptr<float>(), c_desc,
      &heuristic.algo,
      workspace.data_ptr(), heuristic.workspaceSize,
      stream));

  cublasLtMatrixLayoutDestroy(a_desc);
  cublasLtMatrixLayoutDestroy(b_desc);
  cublasLtMatrixLayoutDestroy(c_desc);
  cublasLtMatmulDescDestroy(op_desc);

  return out.reshape(out_sizes);
}

torch::Tensor sub_mul_relu(torch::Tensor y, double sub, double mul) {
  TORCH_CHECK(y.is_cuda(), "CUDA tensor required");
  TORCH_CHECK(y.scalar_type() == at::kFloat, "FP32 tensor required");
  auto yc = y.contiguous();
  auto out = at::empty_like(yc);
  const long n = (long)yc.numel();
  if (n == 0) {
    return out;
  }
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const float* in_ptr = yc.data_ptr<float>();
  float* out_ptr = out.data_ptr<float>();
  const int threads = 256;
  if ((n & 3) == 0) {
    const long n4 = n >> 2;
    long blocks = (n4 + threads - 1) / threads;
    if (blocks > 16384) blocks = 16384;
    sub_mul_relu_vec4_kernel<<<(int)blocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(in_ptr),
        reinterpret_cast<float4*>(out_ptr), n4, (float)sub, (float)mul);
  } else {
    long blocks = (n + threads - 1) / threads;
    if (blocks > 65536) blocks = 65536;
    sub_mul_relu_kernel<<<(int)blocks, threads, 0, stream>>>(
        in_ptr, out_ptr, n, (float)sub, (float)mul);
  }
  return out;
}
"""

_ext = load_inline(
    name="fused_linear_sub_mul_relu",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["linear_sub_mul_relu", "sub_mul_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas", "-lcublasLt"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The whole forward (matmul + bias + subtract + multiply + ReLU) is fused
    into a single custom cublasLt-based CUDA operator by folding the subtract
    into the bias vector and the multiply into the GEMM alpha, with a fused
    ReLU+bias epilogue. If the fused path is unavailable, it falls back to
    cuBLAS linear followed by a custom fused subtract-multiply-ReLU kernel.
    """

    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_fused = True

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            weight = self.linear.weight
            bias = self.linear.bias
            if self._use_fused:
                try:
                    return _ext.linear_sub_mul_relu(
                        x, weight, bias, self.subtract_value,
                        self.multiply_value,
                        bool(torch.backends.cuda.matmul.allow_tf32))
                except RuntimeError:
                    self._use_fused = False
            y = F.linear(x, weight, bias)
            return _ext.sub_mul_relu(y, self.subtract_value,
                                     self.multiply_value)
        # CPU fallback: reference computation.
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```