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

# Allow TF32 tensor-core math for FP32 matmuls (explicitly permitted by the
# precision rule). Storage and compute type stay FP32.
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# ---------------------------------------------------------------------------
# Fused implementation of:  out = min(x @ W.T + bias, c) - c
#
# The op is completely dominated by the 16384x16384 FP32 GEMM (it streams the
# 1 GiB weight matrix from HBM). Strategy:
#   * Call cublasLt directly (TF32 compute + fused BIAS epilogue) instead of
#     going through the PyTorch dispatcher, which re-queries cublasLt
#     heuristics on every single call.
#   * One-time autotune (absorbed by warmup): benchmark every cublasLt
#     heuristic algorithm for this shape with CUDA events and cache the best
#     one per (M, N, K, device).
#   * torch.min(x, c) and x - c are fused into a single in-place vectorized
#     kernel, so the 8 MiB output is read/written exactly once after the GEMM.
# ---------------------------------------------------------------------------
_lt_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <algorithm>
#include <map>
#include <mutex>
#include <tuple>

// ---------------- fused epilogue: data = min(data, c) - c (in place) -------
__global__ void min_sub_vec4_kernel(float4* __restrict__ data, const float c,
                                    const long n4) {
    const long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = data[i];
        v.x = fminf(v.x, c) - c;
        v.y = fminf(v.y, c) - c;
        v.z = fminf(v.z, c) - c;
        v.w = fminf(v.w, c) - c;
        data[i] = v;
    }
}

__global__ void min_sub_kernel(float* __restrict__ data, const float c,
                               const long n) {
    const long i = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (i < n) {
        data[i] = fminf(data[i], c) - c;
    }
}

static void launch_min_sub(float* p, long n, float c, cudaStream_t stream) {
    const int block = 256;
    if ((n & 3) == 0) {
        const long n4 = n >> 2;
        const int grid = (int)((n4 + block - 1) / block);
        min_sub_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<float4*>(p), c, n4);
    } else {
        const int grid = (int)((n + block - 1) / block);
        min_sub_kernel<<<grid, block, 0, stream>>>(p, c, n);
    }
}

// ---------------- cublasLt state -------------------------------------------
struct LtEntry {
    cublasLtMatmulAlgo_t algo;
    size_t workspace = 0;
};

static cublasLtHandle_t g_handle = nullptr;
static std::map<std::tuple<long, long, long, int>, LtEntry> g_cache;
static torch::Tensor g_ws;
static std::mutex g_mu;

static bool ensure_handle() {
    if (g_handle != nullptr) return true;
    if (cublasLtCreate(&g_handle) != CUBLAS_STATUS_SUCCESS) {
        g_handle = nullptr;
        return false;
    }
    return true;
}

static void ensure_workspace(size_t bytes, const torch::TensorOptions& opt) {
    if (bytes == 0) return;
    if (!g_ws.defined() || (size_t)g_ws.numel() < bytes) {
        g_ws = at::empty({(long)bytes}, opt.dtype(at::kByte));
    }
}

// Create matmul descriptor (TF32 compute + fused bias epilogue) and the
// matrix layouts for out(MxN row-major) = x(MxK rm) @ w(NxK rm)^T.
// Column-major view: D(NxM, ld=N) = op_T(A=w, KxN ld=K) @ B(x^T, KxM ld=K).
static bool make_desc_and_layouts(long M, long N, long K,
                                  const float* bias_ptr,
                                  cublasLtMatmulDesc_t* desc,
                                  cublasLtMatrixLayout_t* Ad,
                                  cublasLtMatrixLayout_t* Bd,
                                  cublasLtMatrixLayout_t* Cd) {
    if (cublasLtMatmulDescCreate(desc, CUBLAS_COMPUTE_32F_FAST_TF32,
                                 CUDA_R_32F) != CUBLAS_STATUS_SUCCESS)
        return false;
    cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
    cublasLtMatmulDescSetAttribute(*desc, CUBLASLT_MATMUL_DESC_TRANSA,
                                   &opT, sizeof(opT));
    cublasLtMatmulDescSetAttribute(*desc, CUBLASLT_MATMUL_DESC_TRANSB,
                                   &opN, sizeof(opN));
    cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_BIAS;
    cublasLtMatmulDescSetAttribute(*desc, CUBLASLT_MATMUL_DESC_EPILOGUE,
                                   &epi, sizeof(epi));
    cublasLtMatmulDescSetAttribute(*desc, CUBLASLT_MATMUL_DESC_BIAS,
                                   &bias_ptr, sizeof(bias_ptr));
    if (cublasLtMatrixLayoutCreate(Ad, CUDA_R_32F, K, N, K) !=
            CUBLAS_STATUS_SUCCESS ||
        cublasLtMatrixLayoutCreate(Bd, CUDA_R_32F, K, M, K) !=
            CUBLAS_STATUS_SUCCESS ||
        cublasLtMatrixLayoutCreate(Cd, CUDA_R_32F, N, M, N) !=
            CUBLAS_STATUS_SUCCESS)
        return false;
    return true;
}

static void destroy_desc_and_layouts(cublasLtMatmulDesc_t desc,
                                     cublasLtMatrixLayout_t Ad,
                                     cublasLtMatrixLayout_t Bd,
                                     cublasLtMatrixLayout_t Cd) {
    if (desc) cublasLtMatmulDescDestroy(desc);
    if (Ad) cublasLtMatrixLayoutDestroy(Ad);
    if (Bd) cublasLtMatrixLayoutDestroy(Bd);
    if (Cd) cublasLtMatrixLayoutDestroy(Cd);
}

static bool run_algo(const cublasLtMatmulAlgo_t& algo, size_t ws_bytes,
                     const float* x, const float* w, float* out,
                     long M, long N, long K, cudaStream_t stream) {
    cublasLtMatmulDesc_t desc = nullptr;
    cublasLtMatrixLayout_t Ad = nullptr, Bd = nullptr, Cd = nullptr;
    // NOTE: bias pointer is baked into the descriptor at creation time; the
    // caller rebuilds descriptors when the bias pointer changes (it does not
    // for a given module).
    return false;  // placeholder, never used
}

torch::Tensor linear_min_sub_cuda(torch::Tensor x, torch::Tensor w,
                                  torch::Tensor b, double c) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(), "CUDA tensors expected");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat &&
                b.scalar_type() == at::kFloat, "float32 expected");

    const c10::cuda::CUDAGuard guard(x.device());
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();

    const long M = xc.size(0);
    const long K = xc.size(1);
    const long N = wc.size(0);
    auto out = at::empty({M, N}, xc.options());
    const float cf = static_cast<float>(c);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    std::lock_guard<std::mutex> lock(g_mu);

    if (!ensure_handle()) {
        auto y = at::linear(xc, wc, bc);
        launch_min_sub(y.data_ptr<float>(), y.numel(), cf, stream);
        return y;
    }

    int dev = 0;
    cudaGetDevice(&dev);
    const auto key = std::make_tuple(M, N, K, dev);

    cublasLtMatmulDesc_t desc = nullptr;
    cublasLtMatrixLayout_t Ad = nullptr, Bd = nullptr, Cd = nullptr;
    if (!make_desc_and_layouts(M, N, K, bc.data_ptr<float>(), &desc, &Ad, &Bd, &Cd)) {
        destroy_desc_and_layouts(desc, Ad, Bd, Cd);
        auto y = at::linear(xc, wc, bc);
        launch_min_sub(y.data_ptr<float>(), y.numel(), cf, stream);
        return y;
    }

    const float alpha = 1.0f, beta = 0.0f;
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    auto it = g_cache.find(key);
    if (it == g_cache.end()) {
        // First time we see this shape: query heuristics, benchmark every
        // candidate algorithm with CUDA events, and cache the fastest.
        cublasLtMatmulPreference_t pref = nullptr;
        cublasLtMatmulPreferenceCreate(&pref);
        size_t max_ws = 64ull << 20;
        cublasLtMatmulPreferenceSetAttribute(
            pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_ws, sizeof(max_ws));

        const int kMaxAlgo = 32;
        cublasLtMatmulHeuristicResult_t res[kMaxAlgo];
        int nres = 0;
        cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
            g_handle, desc, Ad, Bd, Cd, Cd, pref, kMaxAlgo, res, &nres);
        if (pref) cublasLtMatmulPreferenceDestroy(pref);

        if (st != CUBLAS_STATUS_SUCCESS || nres == 0) {
            destroy_desc_and_layouts(desc, Ad, Bd, Cd);
            auto y = at::linear(xc, wc, bc);
            launch_min_sub(y.data_ptr<float>(), y.numel(), cf, stream);
            return y;
        }

        size_t need_ws = 0;
        for (int i = 0; i < nres; ++i) need_ws = std::max(need_ws, res[i].workspaceSize);
        ensure_workspace(need_ws, xc.options());
        void* ws_ptr = need_ws ? (void*)g_ws.data_ptr<uint8_t>() : nullptr;

        const bool tune = (2.0 * M * N * K >= 1e9);  // only tune heavy shapes
        int best = 0;
        float best_ms = 1e30f;
        for (int i = 0; i < nres; ++i) {
            // Every candidate computes the full function, so `out` always
            // holds a valid result after this loop.
            st = cublasLtMatmul(g_handle, desc, &alpha, wp, Ad, xp, Bd, &beta,
                                op, Cd, op, Cd, &res[i].algo, ws_ptr,
                                res[i].workspaceSize, stream);
            if (st != CUBLAS_STATUS_SUCCESS) continue;
            launch_min_sub(op, M * N, cf, stream);
            if (!tune) { best = i; break; }

            cublasLtMatmul(g_handle, desc, &alpha, wp, Ad, xp, Bd, &beta, op,
                           Cd, op, Cd, &res[i].algo, ws_ptr,
                           res[i].workspaceSize, stream);  // warm up
            cudaEvent_t e0, e1;
            cudaEventCreate(&e0);
            cudaEventCreate(&e1);
            cudaEventRecord(e0, stream);
            const int reps = 3;
            for (int r = 0; r < reps; ++r) {
                cublasLtMatmul(g_handle, desc, &alpha, wp, Ad, xp, Bd, &beta,
                               op, Cd, op, Cd, &res[i].algo, ws_ptr,
                               res[i].workspaceSize, stream);
                launch_min_sub(op, M * N, cf, stream);
            }
            cudaEventRecord(e1, stream);
            cudaEventSynchronize(e1);
            float ms = 0.0f;
            cudaEventElapsedTime(&ms, e0, e1);
            cudaEventDestroy(e0);
            cudaEventDestroy(e1);
            if (ms < best_ms) { best_ms = ms; best = i; }
        }

        // Leave a valid result in `out` (the best algorithm), then cache it.
        cublasLtMatmul(g_handle, desc, &alpha, wp, Ad, xp, Bd, &beta, op, Cd,
                       op, Cd, &res[best].algo, ws_ptr, res[best].workspaceSize,
                       stream);
        launch_min_sub(op, M * N, cf, stream);

        LtEntry entry;
        entry.algo = res[best].algo;
        entry.workspace = res[best].workspaceSize;
        g_cache[key] = entry;
    } else {
        ensure_workspace(it->second.workspace, xc.options());
        void* ws_ptr = it->second.workspace ? (void*)g_ws.data_ptr<uint8_t>()
                                            : nullptr;
        cublasStatus_t st = cublasLtMatmul(
            g_handle, desc, &alpha, wp, Ad, xp, Bd, &beta, op, Cd, op, Cd,
            &it->second.algo, ws_ptr, it->second.workspace, stream);
        if (st != CUBLAS_STATUS_SUCCESS) {
            destroy_desc_and_layouts(desc, Ad, Bd, Cd);
            auto y = at::linear(xc, wc, bc);
            launch_min_sub(y.data_ptr<float>(), y.numel(), cf, stream);
            return y;
        }
        launch_min_sub(op, M * N, cf, stream);
    }

    destroy_desc_and_layouts(desc, Ad, Bd, Cd);
    return out;
}
"""

_lt_cpp_source = (
    "torch::Tensor linear_min_sub_cuda(torch::Tensor x, torch::Tensor w, "
    "torch::Tensor b, double c);"
)

_ext = None
try:
    _ext = load_inline(
        name="fused_linear_min_sub_cublaslt",
        cpp_sources=_lt_cpp_source,
        cuda_sources=_lt_source,
        functions=["linear_min_sub_cuda"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
        extra_ldflags=["-lcublasLt", "-lcublas"],
    )
except Exception:
    _ext = None


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

    - GEMM + bias runs through cublasLt directly with TF32 tensor cores
      (allowed by the precision rule; tensors stay FP32), the fused BIAS
      epilogue, and a cached best algorithm picked by a one-time CUDA-event
      autotune over all cublasLt heuristic candidates for this shape.
    - torch.min(x, c) and x - c are fused into a single vectorized kernel,
      so the GEMM output is read/written exactly once.
    - Falls back to the standard PyTorch path on any non-CUDA/CPU input or
      if the extension is unavailable.
    """
    def __init__(self, in_features, out_features, constant):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.constant = nn.Parameter(torch.tensor(constant))
        self._const_val = float(constant)

    def forward(self, x):
        w = self.linear.weight
        b = self.linear.bias
        if (_ext is not None and x.is_cuda and x.dtype == torch.float32
                and x.dim() == 2 and w.is_cuda and w.dtype == torch.float32
                and b is not None):
            return _ext.linear_min_sub_cuda(x, w, b, self._const_val)
        out = F.linear(x, w, b)
        return torch.min(out, self.constant) - self.constant
```