#ifdef ENABLE_CUDA

// cuBLASLt-based GEMM peak benchmark.
//
// Uses NVIDIA's high-level cuBLASLt matmul API rather than hand-written WMMA
// to reach close to the vendor-advertised TFLOPS / TOPS.  H2D / D2H transfers
// are excluded from timing; we measure only the kernel launch + execution
// window via cuEvents.
//
// Note on consumer Blackwell (RTX 50 series, sm_120):
// NVIDIA caps several tensor-core dtypes on consumer parts as product
// segmentation vs the datacenter dies (B100 / B200).  Measured on RTX 5060:
//   * tf32 ~= fp32 (no tensor-core boost; TF32 runs on the regular ALUs)
//   * bf16 = 1/2 of fp16
//   * fp16 = full tensor-core rate
//   * fp8  = 2x fp16
// This is intentional, not a kernel-selection bug -- the numbers are at peak
// once the dtype/algo combo is correct.  Don't try to "fix" bf16 ~ 1/2 fp16
// on these cards.

#include <cuda/cuda_peak.h>
#include <cublasLt.h>
#include <common/common.h>
#include <common/dynlib.h>
#include <chrono>
#include <sstream>
#include <vector>

// ---------------------------------------------------------------------------
// Optional cuBLASLt loader.  cuBLASLt is not part of the NVIDIA driver, so it
// is resolved at run time; if absent the GEMM benchmark is skipped and the rest
// of the CUDA backend still runs with just the driver.  Function-pointer types
// are taken from the header via decltype (exact signatures), and macros redirect
// the existing call sites to the loaded pointers -- so the body below is
// unchanged.
// ---------------------------------------------------------------------------
namespace {
struct CublasLtApi
{
  void *lib = nullptr;
  decltype(&::cublasLtCreate)                       Create = nullptr;
  decltype(&::cublasLtDestroy)                      Destroy = nullptr;
  decltype(&::cublasLtMatmul)                       Matmul = nullptr;
  decltype(&::cublasLtMatmulAlgoGetHeuristic)       AlgoGetHeuristic = nullptr;
  decltype(&::cublasLtMatmulDescCreate)             DescCreate = nullptr;
  decltype(&::cublasLtMatmulDescDestroy)            DescDestroy = nullptr;
  decltype(&::cublasLtMatmulDescSetAttribute)       DescSetAttribute = nullptr;
  decltype(&::cublasLtMatmulPreferenceCreate)       PrefCreate = nullptr;
  decltype(&::cublasLtMatmulPreferenceDestroy)      PrefDestroy = nullptr;
  decltype(&::cublasLtMatmulPreferenceSetAttribute) PrefSetAttribute = nullptr;
  decltype(&::cublasLtMatrixLayoutCreate)           LayoutCreate = nullptr;
  decltype(&::cublasLtMatrixLayoutDestroy)          LayoutDestroy = nullptr;
  bool load();
};
CublasLtApi g_lt;
bool CublasLtApi::load()
{
  if (lib)
    return true;
  lib = clpeak::dynOpen({"libcublasLt.so", "libcublasLt.so.13",
                         "libcublasLt.so.12", "libcublasLt.so.11",
                         "cublasLt64_13.dll", "cublasLt64_12.dll",
                         "cublasLt64_11.dll"});
  if (!lib)
    return false;
  bool ok = true;
#define CLPEAK_LT_SYM(member, name)                                      \
  member = reinterpret_cast<decltype(member)>(clpeak::dynSym(lib, name)); \
  ok = ok && (member != nullptr)
  CLPEAK_LT_SYM(Create,           "cublasLtCreate");
  CLPEAK_LT_SYM(Destroy,          "cublasLtDestroy");
  CLPEAK_LT_SYM(Matmul,           "cublasLtMatmul");
  CLPEAK_LT_SYM(AlgoGetHeuristic, "cublasLtMatmulAlgoGetHeuristic");
  CLPEAK_LT_SYM(DescCreate,       "cublasLtMatmulDescCreate");
  CLPEAK_LT_SYM(DescDestroy,      "cublasLtMatmulDescDestroy");
  CLPEAK_LT_SYM(DescSetAttribute, "cublasLtMatmulDescSetAttribute");
  CLPEAK_LT_SYM(PrefCreate,       "cublasLtMatmulPreferenceCreate");
  CLPEAK_LT_SYM(PrefDestroy,      "cublasLtMatmulPreferenceDestroy");
  CLPEAK_LT_SYM(PrefSetAttribute, "cublasLtMatmulPreferenceSetAttribute");
  CLPEAK_LT_SYM(LayoutCreate,     "cublasLtMatrixLayoutCreate");
  CLPEAK_LT_SYM(LayoutDestroy,    "cublasLtMatrixLayoutDestroy");
#undef CLPEAK_LT_SYM
  if (!ok)
  {
    clpeak::dynClose(lib);
    lib = nullptr;
  }
  return ok;
}
} // namespace

#define cublasLtCreate                       g_lt.Create
#define cublasLtDestroy                      g_lt.Destroy
#define cublasLtMatmul                       g_lt.Matmul
#define cublasLtMatmulAlgoGetHeuristic       g_lt.AlgoGetHeuristic
#define cublasLtMatmulDescCreate             g_lt.DescCreate
#define cublasLtMatmulDescDestroy            g_lt.DescDestroy
#define cublasLtMatmulDescSetAttribute       g_lt.DescSetAttribute
#define cublasLtMatmulPreferenceCreate       g_lt.PrefCreate
#define cublasLtMatmulPreferenceDestroy      g_lt.PrefDestroy
#define cublasLtMatmulPreferenceSetAttribute g_lt.PrefSetAttribute
#define cublasLtMatrixLayoutCreate           g_lt.LayoutCreate
#define cublasLtMatrixLayoutDestroy          g_lt.LayoutDestroy

namespace {

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const char *cuErrStrLocal(CUresult r)
{
    const char *s = nullptr;
    cuGetErrorString(r, &s);
    return s ? s : "unknown CUDA error";
}

#define CU_TRY(call) do { CUresult _r = (call); \
    if (_r != CUDA_SUCCESS) { \
        CLPEAK_VLOG("CUDA error at %s:%d: %s\n", __FILE__, __LINE__, cuErrStrLocal(_r)); \
        return -1; } } while (0)

#define LT_TRY(call) do { cublasStatus_t _s = (call); \
    if (_s != CUBLAS_STATUS_SUCCESS) { \
        CLPEAK_VLOG("cuBLASLt error at %s:%d: status=%d\n", __FILE__, __LINE__, (int)_s); \
        return -1; } } while (0)

// Pick a square GEMM dim that scales with the device's compute budget so
// small / large GPUs land in similar wall-clock windows.  Same shape as the
// Metal version: 2048 + numSMs*256, clamped to [2048, 16384], 256-aligned,
// then capped to 25% of VRAM (worst-case fp32 = 4 bytes / element).
uint32_t pickGemmDim(const cuda_device_info_t &info)
{
    uint32_t sms = (uint32_t)(info.numSMs > 0 ? info.numSMs : 16);
    uint64_t D = 2048 + (uint64_t)sms * 256;
    D = (D + 255) & ~uint64_t(255);
    if (D < 2048)  D = 2048;
    if (D > 16384) D = 16384;

    uint64_t budget = info.totalGlobalMem ? info.totalGlobalMem / 4 : ((uint64_t)4 << 30);
    // Worst-case element size is fp64 = 8 bytes (DGEMM); shrink D if the
    // 3*D*D*8 footprint won't fit in 1/4 VRAM.
    while (D > 1024 && 3ULL * D * D * 8 > budget)
        D /= 2;
    return (uint32_t)D;
}

// One-shot timing helper: run `n` cublasLtMatmul calls between an event pair.
// Returns mean per-iter time in microseconds (or <0 on error).
double timeCublasLt(cublasLtHandle_t lt, CUstream stream,
                    cublasLtMatmulDesc_t opDesc,
                    const void *alpha, const void *A, cublasLtMatrixLayout_t Adesc,
                                       const void *B, cublasLtMatrixLayout_t Bdesc,
                    const void *beta,  const void *C, cublasLtMatrixLayout_t Cdesc,
                                             void *D, cublasLtMatrixLayout_t Ddesc,
                    const cublasLtMatmulAlgo_t *algo,
                    void *workspace, size_t workspaceSize,
                    unsigned int n)
{
    CUevent start = nullptr, stop = nullptr;
    cuEventCreate(&start, CU_EVENT_DEFAULT);
    cuEventCreate(&stop,  CU_EVENT_DEFAULT);

    cuStreamSynchronize(stream);
    cuEventRecord(start, stream);
    for (unsigned int i = 0; i < n; i++)
    {
        cublasStatus_t s = cublasLtMatmul(lt, opDesc,
                                          alpha, A, Adesc, B, Bdesc,
                                          beta,  C, Cdesc, D, Ddesc,
                                          algo, workspace, workspaceSize, stream);
        if (s != CUBLAS_STATUS_SUCCESS)
        {
            cuEventDestroy(start); cuEventDestroy(stop);
            CLPEAK_VLOG("cublasLtMatmul failed: status=%d\n", (int)s);
            return -1.0;
        }
    }
    cuEventRecord(stop, stream);
    cuEventSynchronize(stop);

    float ms = 0;
    cuEventElapsedTime(&ms, start, stop);
    cuEventDestroy(start);
    cuEventDestroy(stop);
    return (double)ms * 1000.0 / (double)n; // -> microseconds / iter
}

} // namespace

int CudaPeak::runCublas(CudaDevice &dev, benchmark_config_t &cfg, Category category)
{
    (void)cfg;

    // cuBLASLt is an optional runtime dependency (not part of the driver).
    if (!g_lt.load())
    {
        const bool isInt = category == Category::IntCompute;
        auto t = currentDeviceScope->beginTest(
            {isInt ? "cublas-int" : "cublas-fp", "cuBLASLt GEMM peak",
             isInt ? "tops" : "tflops", Category::Unknown,
             "Matrix-multiply speed through NVIDIA's own tuned library, on a "
             "large square problem.  Where the WMMA rows show what the tensor "
             "cores can do in principle, this shows what shipping code reaches "
             "on the operation most AI work is built from."});
        t.skip(isInt ? "int8" : "fp32", ResultStatus::Unsupported,
               "cuBLASLt library not found; GEMM skipped");
        return 0;
    }

    const uint32_t D = pickGemmDim(dev.info);
    const uint32_t M = D, N = D, K = D;
    const double  flops_per_iter = 2.0 * (double)M * (double)N * (double)K;

    // Layout per variant: fp32 (CUDA cores) uses NN; tensor-core dtypes use TN
    // because that keeps the K axis contiguous for both A and B, matching the
    // MMA load pattern.
    std::stringstream dimStr; dimStr << M << "x" << N << "x" << K;

    // -----------------------------------------------------------------------
    // Allocate persistent buffers (A, B, C, workspace).  Worst-case input
    // size is fp64 = 8 bytes/elem (DGEMM); smaller dtypes alias the same
    // buffers.  All sizes are bytes for fp64.
    // -----------------------------------------------------------------------
    const size_t aBytes = (size_t)M * K * sizeof(double);
    const size_t bBytes = (size_t)K * N * sizeof(double);
    const size_t cBytes = (size_t)M * N * sizeof(double);
    const size_t wsBytes = 256ULL * 1024 * 1024; // 256 MB workspace

    CUdeviceptr dA = 0, dB = 0, dC = 0, dWS = 0;
    if (cuMemAlloc(&dA, aBytes) != CUDA_SUCCESS ||
        cuMemAlloc(&dB, bBytes) != CUDA_SUCCESS ||
        cuMemAlloc(&dC, cBytes) != CUDA_SUCCESS ||
        cuMemAlloc(&dWS, wsBytes) != CUDA_SUCCESS)
    {
        log->note("  Failed to allocate device buffers\n");
        if (dA)  cuMemFree(dA);
        if (dB)  cuMemFree(dB);
        if (dC)  cuMemFree(dC);
        if (dWS) cuMemFree(dWS);
        return -1;
    }

    // Fill A, B with non-zero garbage once (excluded from timing).  cuMemsetD32
    // is the fastest device-side fill that sticks with reasonable bit pattern;
    // 0x3f800000 == fp32 1.0f, so the GEMM produces deterministic finite output.
    cuMemsetD32(dA, 0x3f800000, aBytes / 4);
    cuMemsetD32(dB, 0x3f800000, bBytes / 4);
    cuMemsetD8(dC, 0, cBytes);

    // -----------------------------------------------------------------------
    // cuBLASLt handle
    // -----------------------------------------------------------------------
    cublasLtHandle_t lt = nullptr;
    if (cublasLtCreate(&lt) != CUBLAS_STATUS_SUCCESS)
    {
        log->note("  cublasLtCreate failed\n");
        cuMemFree(dA); cuMemFree(dB); cuMemFree(dC); cuMemFree(dWS);
        return -1;
    }

    logger::TestScope *blasTest = nullptr;

    auto runVariantAB = [&](const char *label, const char *note,
                            cudaDataType_t aType, cudaDataType_t bType,
                            cudaDataType_t cType,
                            cublasComputeType_t computeType,
                            cudaDataType_t scaleType,
                            const void *alphaPtr, const void *betaPtr,
                            bool useTN) -> int
    {
        cublasLtMatmulDesc_t opDesc = nullptr;
        cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;

        // TN (TRANSA=T, TRANSB=N) puts the K axis on the contiguous side of
        // both stored A and stored B -- the MMA load pattern tensor-core
        // kernels are tuned around.  Stored A is K x M col-major, B is K x N
        // col-major.  NN is kept for fp32 because cuBLASLt's CUDA-core fp32
        // kernels are better tuned for the NN layout (TN measured ~50% slower
        // for fp32 on RTX 5060).
        const uint64_t Arows = useTN ? K : M;
        const uint64_t Acols = useTN ? M : K;
        const uint64_t Ald   = Arows;
        if (cublasLtMatmulDescCreate(&opDesc, computeType, scaleType) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Adesc, aType, Arows, Acols, Ald) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Bdesc, bType, K, N, K) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Cdesc, cType, M, N, M) != CUBLAS_STATUS_SUCCESS)
        {
            blasTest->skip(label, ResultStatus::Error, "descriptor create failed", note);
            if (opDesc) cublasLtMatmulDescDestroy(opDesc);
            if (Adesc)  cublasLtMatrixLayoutDestroy(Adesc);
            if (Bdesc)  cublasLtMatrixLayoutDestroy(Bdesc);
            if (Cdesc)  cublasLtMatrixLayoutDestroy(Cdesc);
            return -1;
        }

        cublasOperation_t opA = useTN ? CUBLAS_OP_T : CUBLAS_OP_N;
        cublasOperation_t opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opA, sizeof(opA));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));

        // Heuristic: ask cuBLASLt for its top-K algo candidates.  The top-1
        // result is an estimator pick, not a measured pick; for some dtypes
        // (notably bf16 on RTX 5060 sm_120) the estimator's top algo is ~50%
        // slower than the measured-best.  We request K candidates, run a
        // short timing pass on each, and use the measured fastest for the
        // full timing run.
        const int kCandidates = 8;
        cublasLtMatmulPreference_t pref = nullptr;
        cublasLtMatmulPreferenceCreate(&pref);
        cublasLtMatmulPreferenceSetAttribute(pref,
            CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &wsBytes, sizeof(wsBytes));

        cublasLtMatmulHeuristicResult_t heurs[kCandidates] = {};
        int returnedResults = 0;
        cublasStatus_t hs = cublasLtMatmulAlgoGetHeuristic(
            lt, opDesc, Adesc, Bdesc, Cdesc, Cdesc,
            pref, kCandidates, heurs, &returnedResults);
        cublasLtMatmulPreferenceDestroy(pref);

        if (hs != CUBLAS_STATUS_SUCCESS || returnedResults == 0)
        {
            blasTest->skip(label, ResultStatus::Unsupported,
                          std::string("unsupported on ") + dev.info.archName, note);
            cublasLtMatmulDescDestroy(opDesc);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            return 0;
        }

        // Measured algo selection.  4 iters per candidate is enough to filter
        // out estimator outliers without burning a noticeable wall-clock budget
        // (top-K total is ~4 * K * per_iter_us; on a fast dtype that's well
        // under a second).  Skip candidates that fail to launch (CUBLAS_STATUS
        // not SUCCESS within timeCublasLt -- it returns a negative time).
        const unsigned int probeIters = 4;
        int bestIdx = -1;
        double bestProbeUs = 1e30;
        for (int i = 0; i < returnedResults; i++)
        {
            double t = timeCublasLt(lt, dev.stream, opDesc,
                alphaPtr, (void*)dA, Adesc, (void*)dB, Bdesc,
                betaPtr,  (void*)dC, Cdesc, (void*)dC, Cdesc,
                &heurs[i].algo, (void*)dWS, wsBytes, probeIters);
            if (t > 0.0 && t < bestProbeUs) { bestProbeUs = t; bestIdx = i; }
        }
        if (bestIdx < 0)
        {
            blasTest->skip(label, ResultStatus::Error, "all candidate algos failed", note);
            cublasLtMatmulDescDestroy(opDesc);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            return -1;
        }

        // alpha / beta scalar bit-widths must match `scaleType`; the caller
        // supplies them by pointer so we don't have to fan out a switch here.
        unsigned int warmup = warmupCount > 0 ? warmupCount : 2;
        double per_iter_us = timeCublasLt(lt, dev.stream, opDesc,
            alphaPtr, (void*)dA, Adesc, (void*)dB, Bdesc,
            betaPtr,  (void*)dC, Cdesc, (void*)dC, Cdesc,
            &heurs[bestIdx].algo, (void*)dWS, wsBytes, warmup);

        if (per_iter_us <= 0.0)
        {
            blasTest->skip(label, ResultStatus::Error, "timing probe failed", note);
            cublasLtMatmulDescDestroy(opDesc);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            return -1;
        }

        // BLAS uses a 5 s budget (vs 250 ms for shader-path tests) -- cuBLAS
        // GEMMs aren't subject to the watchdog rules that drive the shader
        // budget, and the longer window stabilises tensor-core throughput.
        unsigned int iters = pickIters(per_iter_us, 5000000u,
                                       forceIters ? specifiedIters : 0);
        double mean_us = timeCublasLt(lt, dev.stream, opDesc,
            alphaPtr, (void*)dA, Adesc, (void*)dB, Bdesc,
            betaPtr,  (void*)dC, Cdesc, (void*)dC, Cdesc,
            &heurs[bestIdx].algo, (void*)dWS, wsBytes, iters);

        double tops = flops_per_iter * 1.0e6 / mean_us / 1.0e12;
        blasTest->emit(label, (float)tops, {false, note});

        cublasLtMatmulDescDestroy(opDesc);
        cublasLtMatrixLayoutDestroy(Adesc);
        cublasLtMatrixLayoutDestroy(Bdesc);
        cublasLtMatrixLayoutDestroy(Cdesc);
        return 0;
    };

    // -----------------------------------------------------------------------
    // Floating-point variants.  All share col-major NN layout; cuBLASLt's
    // heuristic picks the best algo per dtype.  Where a dtype isn't supported
    // by the device's compute capability the heuristic returns 0 results and
    // the runner prints "unsupported on sm_<cc>".
    // -----------------------------------------------------------------------

    // Convenience wrapper: most variants use the same dtype for A and B.
    auto runVariant = [&](const char *label, const char *note,
                          cudaDataType_t abType, cudaDataType_t cType,
                          cublasComputeType_t computeType,
                          cudaDataType_t scaleType,
                          const void *alphaPtr, const void *betaPtr,
                          bool useTN) -> int
    {
        return runVariantAB(label, note, abType, abType, cType, computeType,
                            scaleType, alphaPtr, betaPtr, useTN);
    };

    // Scale-type scalars.  Bit widths MUST match the cublasLtMatmul scaleType:
    //   * R_32F  -> float
    //   * R_16F  -> half (16-bit; we use uint16_t with the IEEE half bit
    //               pattern: 0x3c00 = 1.0, 0x0000 = 0.0)
    //   * R_32I  -> int32
    const float    alpha32   = 1.0f,  beta32   = 0.0f;
    const double   alpha64   = 1.0,   beta64   = 0.0;
    const uint16_t alpha16   = 0x3c00, beta16  = 0x0000;
    const int32_t  alpha32i  = 1,     beta32i  = 0;

#if defined(CLPEAK_CUBLASLT_HAS_FP4)
    // Block-scaled FP4 GEMM (Blackwell, CUDA 12.8+).  Both operands are
    // CUDA_R_4F_E2M1 (4-bit e2m1) and carry a companion per-block scale tensor:
    //   * mxfp4 -> 1 UE8M0 scale per 32 elements (OCP MX)
    //   * nvfp4 -> 1 UE4M3 scale per 16 elements (NVIDIA NVFP4)
    // Output is bf16, compute + scale are fp32.  We bind neutral (==1.0) scale
    // tensors filled once -- their values affect numerics, not throughput.  The
    // scale buffers are over-allocated to the operand element count (always >=
    // the padded scale-tensor size cuBLASLt expects) to stay safe regardless of
    // the exact block-scale layout.  TN layout, like the fp8 path.
    auto runVariantFp4 = [&](const char *label, const char *note,
                             cublasLtMatmulMatrixScale_t scaleMode,
                             uint8_t neutralScaleByte) -> int
    {
        const size_t scaleABytes = (size_t)M * K; // upper bound, 1 byte/scale
        const size_t scaleBBytes = (size_t)K * N;
        CUdeviceptr dSA = 0, dSB = 0;
        if (cuMemAlloc(&dSA, scaleABytes) != CUDA_SUCCESS ||
            cuMemAlloc(&dSB, scaleBBytes) != CUDA_SUCCESS)
        {
            blasTest->skip(label, ResultStatus::Error, "scale buffer alloc failed", note);
            if (dSA) cuMemFree(dSA);
            if (dSB) cuMemFree(dSB);
            return -1;
        }
        cuMemsetD8(dSA, neutralScaleByte, scaleABytes);
        cuMemsetD8(dSB, neutralScaleByte, scaleBBytes);

        cublasLtMatmulDesc_t opDesc = nullptr;
        cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
        // TN: stored A is K x M col-major, B is K x N col-major, C/D are M x N.
        if (cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_4F_E2M1, K, M, K) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_4F_E2M1, K, N, K) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_16BF, M, N, M) != CUBLAS_STATUS_SUCCESS)
        {
            blasTest->skip(label, ResultStatus::Error, "descriptor create failed", note);
            if (opDesc) cublasLtMatmulDescDestroy(opDesc);
            if (Adesc)  cublasLtMatrixLayoutDestroy(Adesc);
            if (Bdesc)  cublasLtMatrixLayoutDestroy(Bdesc);
            if (Cdesc)  cublasLtMatrixLayoutDestroy(Cdesc);
            cuMemFree(dSA); cuMemFree(dSB);
            return -1;
        }

        cublasOperation_t opA = CUBLAS_OP_T, opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opA, sizeof(opA));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));

        // Bind the per-block scale tensors + their mode for A and B.
        CUdeviceptr saPtr = dSA, sbPtr = dSB;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &saPtr, sizeof(saPtr));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &sbPtr, sizeof(sbPtr));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scaleMode, sizeof(scaleMode));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &scaleMode, sizeof(scaleMode));

        cublasLtMatmulPreference_t pref = nullptr;
        cublasLtMatmulPreferenceCreate(&pref);
        cublasLtMatmulPreferenceSetAttribute(pref,
            CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &wsBytes, sizeof(wsBytes));

        const int kCandidates = 8;
        cublasLtMatmulHeuristicResult_t heurs[kCandidates] = {};
        int returnedResults = 0;
        cublasStatus_t hs = cublasLtMatmulAlgoGetHeuristic(
            lt, opDesc, Adesc, Bdesc, Cdesc, Cdesc,
            pref, kCandidates, heurs, &returnedResults);
        cublasLtMatmulPreferenceDestroy(pref);

        auto cleanup = [&]() {
            cublasLtMatmulDescDestroy(opDesc);
            cublasLtMatrixLayoutDestroy(Adesc);
            cublasLtMatrixLayoutDestroy(Bdesc);
            cublasLtMatrixLayoutDestroy(Cdesc);
            cuMemFree(dSA); cuMemFree(dSB);
        };

        if (hs != CUBLAS_STATUS_SUCCESS || returnedResults == 0)
        {
            blasTest->skip(label, ResultStatus::Unsupported,
                          std::string("unsupported on ") + dev.info.archName, note);
            cleanup();
            return 0;
        }

        const unsigned int probeIters = 4;
        int bestIdx = -1;
        double bestProbeUs = 1e30;
        for (int i = 0; i < returnedResults; i++)
        {
            double t = timeCublasLt(lt, dev.stream, opDesc,
                &alpha32, (void*)dA, Adesc, (void*)dB, Bdesc,
                &beta32,  (void*)dC, Cdesc, (void*)dC, Cdesc,
                &heurs[i].algo, (void*)dWS, wsBytes, probeIters);
            if (t > 0.0 && t < bestProbeUs) { bestProbeUs = t; bestIdx = i; }
        }
        if (bestIdx < 0)
        {
            blasTest->skip(label, ResultStatus::Error, "all candidate algos failed", note);
            cleanup();
            return -1;
        }

        unsigned int warmup = warmupCount > 0 ? warmupCount : 2;
        double per_iter_us = timeCublasLt(lt, dev.stream, opDesc,
            &alpha32, (void*)dA, Adesc, (void*)dB, Bdesc,
            &beta32,  (void*)dC, Cdesc, (void*)dC, Cdesc,
            &heurs[bestIdx].algo, (void*)dWS, wsBytes, warmup);
        if (per_iter_us <= 0.0)
        {
            blasTest->skip(label, ResultStatus::Error, "timing probe failed", note);
            cleanup();
            return -1;
        }

        unsigned int iters = pickIters(per_iter_us, 5000000u,
                                       forceIters ? specifiedIters : 0);
        double mean_us = timeCublasLt(lt, dev.stream, opDesc,
            &alpha32, (void*)dA, Adesc, (void*)dB, Bdesc,
            &beta32,  (void*)dC, Cdesc, (void*)dC, Cdesc,
            &heurs[bestIdx].algo, (void*)dWS, wsBytes, iters);

        double tops = flops_per_iter * 1.0e6 / mean_us / 1.0e12;
        blasTest->emit(label, (float)tops, {false, note});
        cleanup();
        return 0;
    };
#endif // CLPEAK_CUBLASLT_HAS_FP4

    if (category == Category::FpCompute)
    {
        // -----------------------------------------------------------------------
        // Floating-point variants -- reported in TFLOPS.
        // -----------------------------------------------------------------------
        auto testFp = currentDeviceScope->beginTest(
            {"cublas-fp", "cuBLASLt GEMM peak", "tflops", Category::Unknown,
             "Matrix-multiply speed through NVIDIA's own tuned library, on a large "
             "square problem.  Where the WMMA rows show what the tensor cores can "
             "do in principle, this shows what shipping code reaches on the "
             "operation most AI work is built from."});
        blasTest = &testFp;

        // fp32: full-precision GEMM on CUDA cores.  NN layout -- TN measured ~50%
        // slower for fp32 on RTX 5060 because cuBLASLt's heuristic falls off the
        // tuned kernel set for fp32 + TN.
        runVariant("fp32", "Full 32-bit precision, on the ordinary shader cores "
                   "rather than the tensor cores.",
                   CUDA_R_32F,  CUDA_R_32F, CUBLAS_COMPUTE_32F,
                   CUDA_R_32F, &alpha32, &beta32, /*useTN=*/false);

        // fp64: DGEMM on CUDA cores (tensor-core fp64 acceleration kicks in
        // automatically on sm_80+; the heuristic picks the right algo).  NN
        // layout matches fp32 -- DGEMM kernels are tuned for that shape.
        runVariant("fp64", "Full 64-bit precision, for scientific computing.  "
                   "Consumer cards run this far slower than the datacenter parts.",
                   CUDA_R_64F,  CUDA_R_64F, CUBLAS_COMPUTE_64F,
                   CUDA_R_64F, &alpha64, &beta64, /*useTN=*/false);

        // tf32: fp32 inputs, but cuBLASLt internally rounds to TF32 and runs on
        // tensor cores (Ampere+).  Inputs/outputs stay fp32.
        if (dev.info.tf32GemmSupported)
            runVariant("tf32", "32-bit numbers in and out, but rounded internally "
                       "to a shorter form so the tensor cores can take them.",
                       CUDA_R_32F,  CUDA_R_32F, CUBLAS_COMPUTE_32F_FAST_TF32,
                       CUDA_R_32F, &alpha32, &beta32, /*useTN=*/true);

        // fp16: half inputs + half output + half compute.  scaleType is R_16F so
        // alpha/beta must be 16-bit half values, not float.
        if (dev.info.fp16Supported)
            runVariant("fp16", "16-bit inputs and totals -- the everyday precision "
                       "of AI inference, and usually the fastest row here.",
                       CUDA_R_16F,  CUDA_R_16F, CUBLAS_COMPUTE_16F,
                       CUDA_R_16F, &alpha16, &beta16, /*useTN=*/true);

        // bf16: bf16 inputs + fp32 output + fp32 compute (mixed-precision).
        // bf16 has no native accumulator format -- compute is always fp32.
        // R_32F output matches the dtype combo NVIDIA's published bf16 GEMM
        // peaks are quoted against.
        if (dev.info.bf16Supported)
            runVariant("bf16", "bfloat16 inputs with 32-bit totals -- 16 bits "
                       "arranged for AI work, trading digits of accuracy for the "
                       "number range of a full float.",
                       CUDA_R_16BF, CUDA_R_32F, CUBLAS_COMPUTE_32F,
                       CUDA_R_32F, &alpha32, &beta32, /*useTN=*/true);

        // fp8: cuBLASLt's fp8 matmul requires TN layout; inputs are 1-byte fp8,
        // accumulator + output are bf16, scale is fp32.  Gated on cc >= 8.9
        // (Ada / Hopper / Blackwell).
        //
        // cuBLASLt does NOT support both A and B being E5M2 -- the heuristic
        // returns 0 results for that combo on every CUDA version we've tested
        // (incl. 13.2 on sm_120).  The supported pairs are:
        //     E4M3 x E4M3   -- canonical "training" fp8 path
        //     E4M3 x E5M2   -- mixed
        //     E5M2 x E4M3   -- mixed (typical inference: e5m2 activations,
        //                              e4m3 weights)
        // We label the second variant "fp8_e5m2" and run E5M2 x E4M3 to measure
        // the throughput of fp8 paths that include an E5M2 input.  Hardware
        // throughput is the same as E4M3 x E4M3; the asymmetry comes from
        // cuBLASLt's algorithm coverage, not the tensor cores.
        if (dev.info.fp8MmaSupported)
        {
            runVariant("fp8_e4m3", "8-bit inputs, in the variant that spends its "
                       "bits on accuracy rather than range.",
                       CUDA_R_8F_E4M3, CUDA_R_16BF, CUBLAS_COMPUTE_32F,
                       CUDA_R_32F, &alpha32, &beta32, /*useTN=*/true);
            runVariantAB("fp8_e5m2", "8-bit inputs including the variant that "
                         "spends its bits on range rather than accuracy -- the "
                         "usual choice for inference.",
                         CUDA_R_8F_E5M2, CUDA_R_8F_E4M3, CUDA_R_16BF,
                         CUBLAS_COMPUTE_32F, CUDA_R_32F,
                         &alpha32, &beta32, /*useTN=*/true);
        }

        // mxfp4 / nvfp4: block-scaled 4-bit GEMM on Blackwell tensor cores.
        // Same dtypes the wmma microbench measures (mxf4_e2m1 / nvf4_e2m1), so
        // the cublas-fp rows line up against the wmma rows.  Gated on
        // fp4GemmSupported (all Blackwell) rather than fp4MmaSupported (the
        // consumer-only raw-mma microbench): the library FP4 path runs on
        // datacenter sm_100/103 as well, and self-skips where it can't.
        // Declared outside the #if: both branches below use them.
        const char *mxf4Note = "4-bit inputs with a shared scale factor per block "
                               "of 32, the open MX format.  The scale is what "
                               "makes 4 bits usable for real models.";
        const char *nvf4Note = "NVIDIA's own 4-bit block format, with a finer "
                               "scale shared by every 16 values instead of every "
                               "32 -- more accurate than MX at the same width.";
#if defined(CLPEAK_CUBLASLT_HAS_FP4)
        if (dev.info.fp4GemmSupported)
        {
            // 0x7F = 1.0 in UE8M0 (exponent-only, bias 127);
            // 0x38 = 1.0 in UE4M3 (e4m3, bias 7).
            runVariantFp4("mxf4_e2m1", mxf4Note,
                          CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0, 0x7F);
            runVariantFp4("nvf4_e2m1", nvf4Note,
                          CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3, 0x38);
        }
        else
        {
            blasTest->skip("mxf4_e2m1", ResultStatus::Unsupported,
                           std::string("FP4 tensor cores require Blackwell -- unsupported on ") + dev.info.archName,
                           mxf4Note);
            blasTest->skip("nvf4_e2m1", ResultStatus::Unsupported,
                           std::string("FP4 tensor cores require Blackwell -- unsupported on ") + dev.info.archName,
                           nvf4Note);
        }
#else
        blasTest->skip("mxf4_e2m1", ResultStatus::Unsupported,
                       "block-scaled FP4 GEMM API not in this cuBLASLt (needs CUDA 12.8+)",
                       mxf4Note);
        blasTest->skip("nvf4_e2m1", ResultStatus::Unsupported,
                       "block-scaled FP4 GEMM API not in this cuBLASLt (needs CUDA 12.8+)",
                       nvf4Note);
#endif
    }

    if (category != Category::IntCompute)
    {
        cublasLtDestroy(lt);
        cuMemFree(dA); cuMemFree(dB); cuMemFree(dC); cuMemFree(dWS);
        return 0;
    }

    // -----------------------------------------------------------------------
    // Integer variants -- reported in TOPS
    // -----------------------------------------------------------------------
    auto testInt = currentDeviceScope->beginTest(
        {"cublas-int", "cuBLASLt GEMM peak", "tops", Category::Unknown,
         "The same tuned-library matrix multiply on whole-number formats, which "
         "is how a quantized (compressed) model actually runs."});
    blasTest = &testInt;

    // int8: 1-byte signed inputs, int32 accumulator + output, int32 compute
    // and scale.  cc >= 7.5 (Turing) for IMMA tensor cores.
    if (dev.info.int8GemmSupported)
        runVariant("int8", "8-bit whole numbers with 32-bit totals -- the format "
                   "quantized neural networks use.",
                   CUDA_R_8I, CUDA_R_32I, CUBLAS_COMPUTE_32I,
                   CUDA_R_32I, &alpha32i, &beta32i, /*useTN=*/true);

    // int4: packed 4-bit signed inputs, int32 accumulator + output.  cc >= 9.0
    // (Hopper) for the IMMA int4 path.  cuBLASLt may refuse this even when the
    // device's compute capability claims support -- the algo heuristic
    // returning 0 results is the canonical "unsupported" signal and we report
    // that line cleanly.
    if (dev.info.int4GemmSupported)
        runVariant("int4", "4-bit whole numbers.  The library may refuse this "
                   "even where the hardware claims it.",
                   CUDA_R_4I, CUDA_R_32I, CUBLAS_COMPUTE_32I,
                   CUDA_R_32I, &alpha32i, &beta32i, /*useTN=*/true);

    cublasLtDestroy(lt);
    cuMemFree(dA); cuMemFree(dB); cuMemFree(dC); cuMemFree(dWS);
    return 0;
}

#endif // ENABLE_CUDA
