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

# Fully fused operator:
#   out = relu( mul * (x @ weight^T + bias - sub) )
#       = relu( mul * (x @ weight^T) + mul * (bias - sub) )
# One tiny kernel builds bias_eff = mul * (bias - sub); a single cublasLt
# GEMM with TF32 tensor-core math (explicitly requested, as allowed) and the
# fused RELU_BIAS epilogue computes everything else (alpha = mul).
# On the first call for a given (M, N, K) the top cublasLt heuristics are
# timed and the fastest algorithm is cached.
_cpp_src = """
torch::Tensor linear_sub_mul_relu(torch::Tensor x, torch::Tensor weight,
                                  torch::Tensor bias, double sub, double mul);
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 <climits>
#include <map>
#include <tuple>
#include <vector>
#include <algorithm>

#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;
}

bool tf32Supported() {
  static bool supported = []() {
    int dev = 0;
    if (cudaGetDevice(&dev) != cudaSuccess) return false;
    int major = 0;
    if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor,
                               dev) != cudaSuccess)
      return false;
    return major >= 8;
  }();
  return supported;
}

__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);
  }
}

using Key = std::tuple<int64_t, int64_t, int64_t>;

std::map<Key, cublasLtMatmulHeuristicResult_t>& planCache() {
  static std::map<Key, cublasLtMatmulHeuristicResult_t> cache;
  return cache;
}

}  // namespace

torch::Tensor 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(), "CUDA tensors required");
  TORCH_CHECK(bias.defined() && bias.is_cuda(), "CUDA bias 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();
  cublasLtHandle_t handle = getLtHandle();

  // bias_eff = (bias - sub) * mul   (one entry 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);
  }

  // TF32 tensor-core math is explicitly allowed for this task.
  const cublasComputeType_t compute_type =
      tf32Supported() ? 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) 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));

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

  const Key key = std::make_tuple(M, N, K);
  auto& cache = planCache();
  auto it = cache.find(key);

  cublasLtMatmulHeuristicResult_t heuristic;
  if (it != 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)));

    const int max_results = 4;
    std::vector<cublasLtMatmulHeuristicResult_t> results(max_results);
    int returned = 0;
    CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
        handle, op_desc, a_desc, b_desc, c_desc, c_desc, pref, max_results,
        results.data(), &returned));
    cublasLtMatmulPreferenceDestroy(pref);
    TORCH_CHECK(returned > 0, "cublasLt found no suitable algorithm");
    results.resize(returned);

    size_t ws = 1;
    for (const auto& r : results) ws = std::max(ws, r.workspaceSize);
    auto workspace =
        at::empty({(int64_t)ws}, xc.options().dtype(at::kByte));

    // One-time micro autotune over the candidate algorithms.
    cudaEvent_t ev0, ev1;
    cudaEventCreate(&ev0);
    cudaEventCreate(&ev1);
    float best_ms = 1e30f;
    int best = -1;
    for (int i = 0; i < returned; ++i) {
      cublasStatus_t st = cublasLtMatmul(
          handle, 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,
          &results[i].algo, workspace.data_ptr(), results[i].workspaceSize,
          stream);
      if (st != CUBLAS_STATUS_SUCCESS) continue;
      cudaEventRecord(ev0, stream);
      for (int r = 0; r < 3; ++r) {
        cublasLtMatmul(handle, 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,
                       &results[i].algo, workspace.data_ptr(),
                       results[i].workspaceSize, stream);
      }
      cudaEventRecord(ev1, stream);
      cudaEventSynchronize(ev1);
      float ms = 0.0f;
      cudaEventElapsedTime(&ms, ev0, ev1);
      if (ms < best_ms) {
        best_ms = ms;
        best = i;
      }
    }
    cudaEventDestroy(ev0);
    cudaEventDestroy(ev1);
    TORCH_CHECK(best >= 0, "no runnable cublasLt algorithm found");
    heuristic = results[best];
    cache.emplace(key, heuristic);
  }

  auto workspace = at::empty(
      {(int64_t)std::max<size_t>(heuristic.workspaceSize, 1)},
      xc.options().dtype(at::kByte));

  CUBLASLT_CHECK(cublasLtMatmul(
      handle, 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_v2",
    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
    TF32 tensor-core GEMM and a 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
        # TF32 tensor-core math is explicitly permitted for this task; make
        # sure the ATen fallback path also benefits from it.
        if torch.cuda.is_available():
            try:
                torch.backends.cuda.matmul.allow_tf32 = True
            except Exception:
                pass

    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 and bias is not None:
                try:
                    return _ext.linear_sub_mul_relu(
                        x, weight, bias, self.subtract_value,
                        self.multiply_value)
                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)
```