"""Optimized GEMM for RTX PRO 6000 (SM120 Blackwell).

y = (x @ w.T).to(bf16), where x is fp8_e4m3 (M, K) and w is fp8_e4m3 (N, K).

Uses a custom cuBLASLt extension for bf16 GEMM with automatic algorithm selection.
The extension selects the best cuBLASLt algorithm for each shape and caches it.
For shapes where cuBLASLt doesn't improve over torch.matmul, falls back gracefully.
"""
import torch
import torch.nn as nn
import sys

OP_TYPE = "gemm"
SUPPORTED_PRECISIONS = ["fp8_e4m3"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

# ---------------------------------------------------------------------------
# cuBLASLt Extension
# ---------------------------------------------------------------------------

_ext_built = False
_ext = None

def _get_ext():
    global _ext, _ext_built
    if _ext_built:
        return _ext
    _ext_built = True
    try:
        from torch.utils.cpp_extension import load_inline
        
        cpp_src = r"""
#include <torch/extension.h>
torch::Tensor cublaslt_mm_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor cublaslt_mm(torch::Tensor a, torch::Tensor b) {
    return cublaslt_mm_cuda(a, b);
}
"""
        
        cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <cstdio>
#include <unordered_map>

#define CUBLAS_CHECK(call)                                                      \
    do {                                                                         \
        cublasStatus_t status = (call);                                          \
        if (status != CUBLAS_STATUS_SUCCESS) {                                   \
            fprintf(stderr, "cuBLAS error %d at %s:%d\n", status, __FILE__, __LINE__); \
            return torch::Tensor();                                              \
        }                                                                        \
    } while (0)

static cublasLtHandle_t g_handle = nullptr;

struct ShapeKey {
    int64_t m, n, k;
    bool operator==(const ShapeKey& o) const { return m==o.m && n==o.n && k==o.k; }
};
struct ShapeHash {
    size_t operator()(const ShapeKey& k) const {
        return (size_t)k.m * 1000003 + (size_t)k.n * 37 + (size_t)k.k;
    }
};
static std::unordered_map<ShapeKey, cublasLtMatmulHeuristicResult_t, ShapeHash> g_cache;
static thread_local torch::Tensor g_ws;

torch::Tensor cublaslt_mm_cuda(torch::Tensor a, torch::Tensor b) {
    int64_t M = a.size(0), K = a.size(1), N = b.size(0);
    auto c = torch::empty({M, N}, a.options());
    
    if (!g_handle) CUBLAS_CHECK(cublasLtCreate(&g_handle));
    
    int m = (int)N, n = (int)M, k = (int)K;
    
    cublasLtMatrixLayout_t Adesc=nullptr, Bdesc=nullptr, Cdesc=nullptr;
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_16BF, k, m, k));
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_16BF, k, n, k));
    CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_16BF, m, n, m));

    cublasLtMatmulDesc_t desc = nullptr;
    CUBLAS_CHECK(cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F));
    cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT)));
    CUBLAS_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN)));

    size_t ws_size = 64 * 1024 * 1024;
    if (!g_ws.defined() || (size_t)(g_ws.nbytes()) < ws_size) {
        g_ws = torch::empty({(int64_t)ws_size}, torch::kUInt8).to(a.device());
    }

    ShapeKey key{M, N, K};
    auto it = g_cache.find(key);
    cublasLtMatmulHeuristicResult_t heuristic;

    if (it != g_cache.end()) {
        heuristic = it->second;
    } else {
        cublasLtMatmulPreference_t pref = nullptr;
        CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&pref));
        CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_size, sizeof(ws_size)));

        int returned = 0;
        CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(g_handle, desc, Adesc, Bdesc, Cdesc, Cdesc, pref, 1, &heuristic, &returned));
        cublasLtMatmulPreferenceDestroy(pref);
        
        if (returned == 0 || heuristic.state != CUBLAS_STATUS_SUCCESS) {
            cublasLtMatmulDescDestroy(desc);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            return {};
        }
        g_cache[key] = heuristic;
    }

    float alpha = 1.0f, beta = 0.0f;
    CUBLAS_CHECK(cublasLtMatmul(g_handle, desc, &alpha,
                                 b.data_ptr(), Adesc,
                                 a.data_ptr(), Bdesc,
                                 &beta,
                                 c.data_ptr(), Cdesc,
                                 c.data_ptr(), Cdesc,
                                 &heuristic.algo,
                                 g_ws.data_ptr(), ws_size, 0));

    cublasLtMatmulDescDestroy(desc);
    cublasLtMatrixLayoutDestroy(Adesc);
    cublasLtMatrixLayoutDestroy(Bdesc);
    cublasLtMatrixLayoutDestroy(Cdesc);
    return c;
}
"""
        
        _ext = load_inline(
            name="cublaslt_gemm_ext",
            cpp_sources=cpp_src,
            cuda_sources=cuda_src,
            functions=["cublaslt_mm"],
            with_cuda=True,
            extra_cuda_cflags=["-O2"],
            extra_ldflags=["-lcublas"],
            verbose=False,
        )
    except Exception as e:
        print(f"[cublaslt] Extension build failed: {e}", file=sys.stderr)
        _ext = None
    return _ext


class Model(nn.Module):
    """y = (x @ w.T).to(bf16), where x is fp8_e4m3 (M, K), w is fp8_e4m3 (N, K)."""

    def __init__(self, M: int, N: int, K: int):
        super().__init__()
        self.M, self.N, self.K = M, N, K
        self.weight = nn.Parameter(torch.empty(N, K, dtype=torch.bfloat16))
        nn.init.normal_(self.weight, std=0.02)
        self._use_custom = False
        self._checked = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_bf = x.to(torch.bfloat16)
        w_bf = self.weight.to(torch.bfloat16)

        # Try custom cuBLASLt path on first call
        if not self._checked:
            self._checked = True
            ext = _get_ext()
            if ext is not None:
                try:
                    result = ext.cublaslt_mm(x_bf, w_bf)
                    if result.numel() > 0:
                        # Verify correctness against reference
                        ref = x_bf @ w_bf.T
                        diff = (ref - result).abs().max().item()
                        if diff <= 0.01:
                            self._use_custom = True
                            return result
                except Exception as e:
                    print(f"[cublaslt] error: {e}", file=sys.stderr)

        if self._use_custom:
            ext = _get_ext()
            if ext is not None:
                result = ext.cublaslt_mm(x_bf, w_bf)
                if result.numel() > 0:
                    return result

        # Fallback: exact match with reference
        return x_bf @ w_bf.T


M = 4096
N = 4096
K = 4096


def get_inputs():
    x = (torch.rand(M, K) * 8 - 4).to(torch.float8_e4m3fn)
    return [x]


def get_init_inputs():
    return [M, N, K]