```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) )
# A tiny kernel builds bias_eff = mul * (bias - sub); a single cublasLt GEMM
# with TF32 tensor-core math (explicitly allowed for this task) and the fused
# RELU_BIAS epilogue computes everything else (alpha = mul).
#
# Compared to the previous version:
#   * One-time plan cache: cublasLt descriptors, the chosen algorithm and the
#     workspace are created once per (M, N, K, device) and reused every call.
#   * Exhaustive one-time autotune: besides the cublasLt heuristics, all
#     enumerated algorithms are validated and timed, for BOTH the TN layout
#     (weight as-is) and the NN layout (pre-transposed weight, cached with
#     version checking). The fastest overall is kept.
#   * bias_eff is cached and only recomputed when the bias tensor actually
#     changes (pointer/version check), into a persistent buffer so the cached
#     epilogue bias pointer stays valid.
#   * No per-call reshapes/views; the output is allocated in its final shape.
_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 <c10/cuda/CUDAGuard.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <cstring>
#include <map>
#include <tuple>
#include <vector>

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

// TF32 tensor-core math is explicitly allowed for this task (sm_80+).
bool tf32Supported() {
  static int supported = -1;
  if (supported < 0) {
    int dev = 0;
    supported = 0;
    if (cudaGetDevice(&dev) == cudaSuccess) {
      int major = 0;
      if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor,
                                 dev) == cudaSuccess &&
          major >= 8) {
        supported = 1;
      }
    }
  }
  return supported == 1;
}

inline int64_t ver_of(const at::Tensor& t) {
  return (int64_t)t.unsafeGetTensorImpl()->version_counter().current_version();
}

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

struct GemmPlan {
  cublasLtMatmulDesc_t op_desc = nullptr;
  cublasLtMatrixLayout_t a_desc = nullptr;
  cublasLtMatrixLayout_t b_desc = nullptr;
  cublasLtMatrixLayout_t c_desc = nullptr;
  cublasLtMatmulHeuristicResult_t heuristic{};
  bool use_wt = false;
  at::Tensor workspace;
  at::Tensor bias_eff;
  at::Tensor weight_t;
  at::Tensor bias_src;
  at::Tensor weight_src;
  const void* bias_src_ptr = nullptr;
  const void* weight_src_ptr = nullptr;
  int64_t bias_src_version = -1;
  int64_t weight_src_version = -1;
  double sub = 0.0;
  double mul = 0.0;
};

cublasLtMatmulDesc_t make_op_desc(cublasComputeType_t ct,
                                  cublasOperation_t transa, float* bias_ptr) {
  cublasLtMatmulDesc_t d = nullptr;
  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&d, ct, CUDA_R_32F));
  cublasOperation_t tb = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      d, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(d, CUBLASLT_MATMUL_DESC_TRANSB,
                                                &tb, sizeof(tb)));
  cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      d, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      d, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));
  return d;
}

struct TimedAlgo {
  cublasLtMatmulHeuristicResult_t hr;
  float ms;
};

// Gather candidate algorithms (cublasLt heuristics + full enumeration) for one
// matrix-layout configuration and micro-benchmark each of them.
void collect_and_time(std::vector<TimedAlgo>& sink, cublasComputeType_t ct,
                      cublasLtMatmulDesc_t op_desc,
                      cublasLtMatrixLayout_t a_desc,
                      cublasLtMatrixLayout_t b_desc,
                      cublasLtMatrixLayout_t c_desc, const float* A,
                      const float* B, float* C, float alpha, float beta,
                      void* ws_ptr, size_t ws_cap, int max_heur, int max_enum,
                      cudaStream_t stream) {
  cublasLtHandle_t handle = getLtHandle();
  std::vector<cublasLtMatmulHeuristicResult_t> cands;

  {
    cublasLtMatmulPreference_t pref = nullptr;
    if (cublasLtMatmulPreferenceCreate(&pref) == CUBLAS_STATUS_SUCCESS) {
      cublasLtMatmulPreferenceSetAttribute(
          pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_cap,
          sizeof(ws_cap));
      std::vector<cublasLtMatmulHeuristicResult_t> hres((size_t)max_heur);
      int nret = 0;
      if (cublasLtMatmulAlgoGetHeuristic(handle, op_desc, a_desc, b_desc,
                                         c_desc, c_desc, pref, max_heur,
                                         hres.data(),
                                         &nret) == CUBLAS_STATUS_SUCCESS) {
        for (int i = 0; i < nret; ++i) {
          if (hres[i].state == CUBLAS_STATUS_SUCCESS &&
              hres[i].workspaceSize <= ws_cap) {
            cands.push_back(hres[i]);
          }
        }
      }
      cublasLtMatmulPreferenceDestroy(pref);
    }
  }

  {
    int ids[64];
    int nids = 0;
    if (cublasLtMatmulAlgoGetIds(handle, ct, CUDA_R_32F, CUDA_R_32F,
                                 CUDA_R_32F, CUDA_R_32F, 64, ids,
                                 &nids) == CUBLAS_STATUS_SUCCESS) {
      int lim = nids < max_enum ? nids : max_enum;
      for (int i = 0; i < lim; ++i) {
        cublasLtMatmulAlgo_t algo;
        std::memset(&algo, 0, sizeof(algo));
        if (cublasLtMatmulAlgoInit(handle, ct, CUDA_R_32F, CUDA_R_32F,
                                   CUDA_R_32F, CUDA_R_32F, ids[i],
                                   &algo) != CUBLAS_STATUS_SUCCESS) {
          continue;
        }
        bool dup = false;
        for (size_t j = 0; j < cands.size(); ++j) {
          if (std::memcmp(&cands[j].algo, &algo, sizeof(algo)) == 0) {
            dup = true;
            break;
          }
        }
        if (dup) continue;
        cublasLtMatmulHeuristicResult_t chk;
        std::memset(&chk, 0, sizeof(chk));
        if (cublasLtMatmulAlgoCheck(handle, op_desc, a_desc, b_desc, c_desc,
                                    c_desc, &algo,
                                    &chk) != CUBLAS_STATUS_SUCCESS) {
          continue;
        }
        if (chk.state != CUBLAS_STATUS_SUCCESS || chk.workspaceSize > ws_cap) {
          continue;
        }
        chk.algo = algo;
        cands.push_back(chk);
      }
    }
  }

  if (cands.empty()) return;

  cudaEvent_t e0, e1;
  cudaEventCreate(&e0);
  cudaEventCreate(&e1);
  for (size_t i = 0; i < cands.size(); ++i) {
    cublasLtMatmulHeuristicResult_t& hr = cands[i];
    cublasStatus_t st =
        cublasLtMatmul(handle, op_desc, &alpha, A, a_desc, B, b_desc, &beta, C,
                       c_desc, C, c_desc, &hr.algo, ws_ptr, hr.workspaceSize,
                       stream);
    if (st != CUBLAS_STATUS_SUCCESS) continue;
    if (cudaStreamSynchronize(stream) != cudaSuccess) continue;
    cudaEventRecord(e0, stream);
    for (int r = 0; r < 3; ++r) {
      cublasLtMatmul(handle, op_desc, &alpha, A, a_desc, B, b_desc, &beta, C,
                     c_desc, C, c_desc, &hr.algo, ws_ptr, hr.workspaceSize,
                     stream);
    }
    cudaEventRecord(e1, stream);
    cudaEventSynchronize(e1);
    float ms = 0.0f;
    cudaEventElapsedTime(&ms, e0, e1);
    TimedAlgo ta;
    ta.hr = hr;
    ta.ms = ms / 3.0f;
    sink.push_back(ta);
  }
  cudaEventDestroy(e0);
  cudaEventDestroy(e1);
}

const TimedAlgo* best_of(const std::vector<TimedAlgo>& v) {
  const TimedAlgo* b = nullptr;
  for (size_t i = 0; i < v.size(); ++i) {
    if (b == nullptr || v[i].ms < b->ms) b = &v[i];
  }
  return b;
}

// Build a plan for a fixed (M, N, K): autotune both the TN layout (weight
// as-is) and the NN layout (pre-transposed weight) and keep the fastest.
GemmPlan create_plan(int64_t M, int64_t N, int64_t K, int use_tf32,
                     const float* x_ptr, const at::Tensor& w, float* out_ptr,
                     float* bias_eff_ptr, float alpha, cudaStream_t stream) {
  const cublasComputeType_t ct =
      use_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;
  const float beta = 0.0f;
  const size_t ws_cap = (size_t)256 << 20;
  auto scratch = at::empty({(int64_t)ws_cap}, w.options().dtype(at::kByte));
  void* ws_ptr = scratch.data_ptr();

  // Row-major C(M,N) = X(M,K) * W^T(K,N) as the equivalent column-major
  // problem C^T(N,M) = W(N,K) * X^T(K,M).
  cublasLtMatrixLayout_t b_desc = nullptr, c_desc = nullptr;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_32F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&c_desc, CUDA_R_32F, N, M, N));

  // Config 1 (TN): A = weight buffer viewed col-major (K,N), op(A) = A^T.
  cublasLtMatmulDesc_t op_tn = make_op_desc(ct, CUBLAS_OP_T, bias_eff_ptr);
  cublasLtMatrixLayout_t a_tn = nullptr;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_tn, CUDA_R_32F, K, N, K));
  std::vector<TimedAlgo> res_tn;
  collect_and_time(res_tn, ct, op_tn, a_tn, b_desc, c_desc,
                   w.data_ptr<float>(), x_ptr, out_ptr, alpha, beta, ws_ptr,
                   ws_cap, 16, 24, stream);

  // Config 2 (NN): A = pre-transposed weight viewed col-major (N,K), op(A)=A.
  at::Tensor wt = w.t().contiguous();
  cublasLtMatmulDesc_t op_nn = make_op_desc(ct, CUBLAS_OP_N, bias_eff_ptr);
  cublasLtMatrixLayout_t a_nn = nullptr;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_nn, CUDA_R_32F, N, K, N));
  std::vector<TimedAlgo> res_nn;
  collect_and_time(res_nn, ct, op_nn, a_nn, b_desc, c_desc,
                   wt.data_ptr<float>(), x_ptr, out_ptr, alpha, beta, ws_ptr,
                   ws_cap, 16, 24, stream);

  const TimedAlgo* bt = best_of(res_tn);
  const TimedAlgo* bn = best_of(res_nn);
  TORCH_CHECK(bt != nullptr || bn != nullptr,
              "cublasLt found no runnable algorithm");
  const bool choose_nn = (bn != nullptr) && (bt == nullptr || bn->ms < bt->ms);

  GemmPlan plan;
  plan.use_wt = choose_nn;
  if (choose_nn) {
    plan.op_desc = op_nn;
    plan.a_desc = a_nn;
    plan.heuristic = bn->hr;
    plan.weight_t = wt;
    cublasLtMatmulDescDestroy(op_tn);
    cublasLtMatrixLayoutDestroy(a_tn);
  } else {
    plan.op_desc = op_tn;
    plan.a_desc = a_tn;
    plan.heuristic = bt->hr;
    cublasLtMatmulDescDestroy(op_nn);
    cublasLtMatrixLayoutDestroy(a_nn);
  }
  plan.b_desc = b_desc;
  plan.c_desc = c_desc;
  const size_t ws_need =
      plan.heuristic.workspaceSize > 0 ? plan.heuristic.workspaceSize : 1;
  plan.workspace =
      at::empty({(int64_t)ws_need}, w.options().dtype(at::kByte));
  return plan;
}

}  // 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() && 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");

  const c10::cuda::OptionalCUDAGuard guard(at::device_of(x));

  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");
  TORCH_CHECK(K > 0, "in_features must be positive");
  TORCH_CHECK(K <= INT32_MAX && N <= INT32_MAX, "unsupported dims");
  const int64_t M = xc.numel() / K;
  TORCH_CHECK(M <= INT32_MAX, "unsupported dims");

  auto out_sizes = xc.sizes().vec();
  out_sizes.back() = N;
  auto out = at::empty(out_sizes, xc.options());
  if (M == 0 || N == 0) {
    return out;
  }

  // The cached algorithm assumes 16-byte aligned pointers (always true for
  // freshly allocated CUDA tensors; fix up exotic views just in case).
  if ((reinterpret_cast<uintptr_t>(xc.data_ptr()) & 15) != 0) xc = xc.clone();
  if ((reinterpret_cast<uintptr_t>(wc.data_ptr()) & 15) != 0) wc = wc.clone();

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int use_tf32 = tf32Supported() ? 1 : 0;
  const int dev = (int)xc.get_device();
  const float alpha = (float)mul;
  const float beta = 0.0f;

  typedef std::tuple<int64_t, int64_t, int64_t, int, int> Key;
  static std::map<Key, GemmPlan>* plans = new std::map<Key, GemmPlan>();
  const Key key = std::make_tuple(M, N, K, use_tf32, dev);

  std::map<Key, GemmPlan>::iterator it = plans->find(key);
  if (it == plans->end()) {
    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);
    }
    GemmPlan p = create_plan(M, N, K, use_tf32, xc.data_ptr<float>(), wc,
                             out.data_ptr<float>(),
                             bias_eff.data_ptr<float>(), alpha, stream);
    p.bias_eff = bias_eff;
    p.bias_src = bc;  // keep the storage alive so its pointer cannot be reused
    p.bias_src_ptr = bc.data_ptr();
    p.bias_src_version = ver_of(bc);
    p.sub = sub;
    p.mul = mul;
    if (p.use_wt) {
      p.weight_src = wc;
      p.weight_src_ptr = wc.data_ptr();
      p.weight_src_version = ver_of(wc);
    }
    it = plans->emplace(key, std::move(p)).first;
  }
  GemmPlan& plan = it->second;

  // Refresh bias_eff only if the bias (or scalars) actually changed.
  const bool bias_ok = plan.bias_src_ptr == bc.data_ptr() &&
                       plan.bias_src_version == ver_of(bc) &&
                       plan.sub == sub && plan.mul == mul;
  if (!bias_ok) {
    const int threads = 256;
    const int blocks = (int)((N + threads - 1) / threads);
    bias_transform_kernel<<<blocks, threads, 0, stream>>>(
        bc.data_ptr<float>(), plan.bias_eff.data_ptr<float>(), (int)N,
        (float)sub, (float)mul);
    plan.bias_src = bc;
    plan.bias_src_ptr = bc.data_ptr();
    plan.bias_src_version = ver_of(bc);
    plan.sub = sub;
    plan.mul = mul;
  }

  if (plan.use_wt) {
    const bool w_ok = plan.weight_src_ptr == wc.data_ptr() &&
                      plan.weight_src_version == ver_of(wc);
    if (!w_ok) {
      plan.weight_t.copy_(wc.t());
      plan.weight_src = wc;
      plan.weight_src_ptr = wc.data_ptr();
      plan.weight_src_version = ver_of(wc);
    }
  }

  const float* A = plan.use_wt ? plan.weight_t.data_ptr<float>()
                               : wc.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmul(
      getLtHandle(), plan.op_desc, &alpha, A, plan.a_desc,
      xc.data_ptr<float>(), plan.b_desc, &beta, out.data_ptr<float>(),
      plan.c_desc, out.data_ptr<float>(), plan.c_desc, &plan.heuristic.algo,
      plan.workspace.data_ptr(), plan.heuristic.workspaceSize, stream));

  return out;
}

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_v4",
    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"],
)

_fused = _ext.linear_sub_mul_relu
_ew = _ext.sub_mul_relu


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, using a TF32
    tensor-core GEMM (explicitly allowed) with a fused ReLU+bias epilogue.
    Descriptors, the autotuned algorithm (searched exhaustively over both the
    TN and pre-transposed NN layouts), the workspace and the transformed bias
    are all cached across calls. 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:
            lin = self.linear
            if self._use_fused and lin.bias is not None:
                try:
                    return _fused(x, lin.weight, lin.bias,
                                  self.subtract_value, self.multiply_value)
                except RuntimeError:
                    self._use_fused = False
            return _ew(F.linear(x, lin.weight, lin.bias),
                       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)
```