#define _POSIX_C_SOURCE 200809L

#include <dlfcn.h>
#include <math.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#if defined(__has_include)
#if __has_include(<cublasLt.h>)
#include <cublasLt.h>
#define HAVE_CUBLASLT_HEADERS 1
#else
#define HAVE_CUBLASLT_HEADERS 0
#endif
#else
#define HAVE_CUBLASLT_HEADERS 0
#endif

typedef int CUdevice;
typedef uint64_t CUdeviceptr;
typedef int CUresult;
typedef void *CUcontext;
typedef void *CUmodule;
typedef void *CUfunction;
typedef void *CUstream;

#define CU_SUCCESS 0
#define CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT 16
#define CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR 75
#define CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR 76
#define MAX_STRESS_STREAMS 16
#define MIN_PROFILE_BUDGET_BYTES ((size_t)4u * 1024u * 1024u)
#define MIN_STREAM_BUDGET_BYTES ((size_t)64u * 1024u * 1024u)
#define MAX_SINGLE_PRECISION_STREAMS 4
#define MAX_SINGLE_PRECISION_PROFILE_BUDGET_BYTES ((size_t)2u * 1024u * 1024u * 1024u)

static const char *ptx_source =
    ".version 6.0\n"
    ".target sm_30\n"
    ".address_size 64\n"
    "\n"
    ".visible .entry burn(\n"
    "    .param .u64 data,\n"
    "    .param .u32 words,\n"
    "    .param .u32 rounds\n"
    ")\n"
    "{\n"
    "    .reg .pred %p<2>;\n"
    "    .reg .b32 %r<8>;\n"
    "    .reg .b64 %rd<5>;\n"
    "\n"
    "    ld.param.u64 %rd1, [data];\n"
    "    ld.param.u32 %r1, [words];\n"
    "    ld.param.u32 %r2, [rounds];\n"
    "    mov.u32 %r3, %ctaid.x;\n"
    "    mov.u32 %r4, %ntid.x;\n"
    "    mov.u32 %r5, %tid.x;\n"
    "    mad.lo.s32 %r0, %r3, %r4, %r5;\n"
    "    setp.ge.u32 %p0, %r0, %r1;\n"
    "    @%p0 bra DONE;\n"
    "    mul.wide.u32 %rd2, %r0, 4;\n"
    "    add.s64 %rd3, %rd1, %rd2;\n"
    "    ld.global.u32 %r6, [%rd3];\n"
    "LOOP:\n"
    "    setp.eq.u32 %p1, %r2, 0;\n"
    "    @%p1 bra STORE;\n"
    "    mad.lo.u32 %r6, %r6, 1664525, 1013904223;\n"
    "    sub.u32 %r2, %r2, 1;\n"
    "    bra LOOP;\n"
    "STORE:\n"
    "    st.global.u32 [%rd3], %r6;\n"
    "DONE:\n"
    "    ret;\n"
    "}\n";

typedef CUresult (*cuInit_fn)(unsigned int);
typedef CUresult (*cuDeviceGetCount_fn)(int *);
typedef CUresult (*cuDeviceGet_fn)(CUdevice *, int);
typedef CUresult (*cuDeviceGetName_fn)(char *, int, CUdevice);
typedef CUresult (*cuDeviceGetAttribute_fn)(int *, int, CUdevice);
typedef CUresult (*cuCtxCreate_fn)(CUcontext *, unsigned int, CUdevice);
typedef CUresult (*cuCtxDestroy_fn)(CUcontext);
typedef CUresult (*cuCtxSynchronize_fn)(void);
typedef CUresult (*cuMemAlloc_fn)(CUdeviceptr *, size_t);
typedef CUresult (*cuMemFree_fn)(CUdeviceptr);
typedef CUresult (*cuMemsetD8_fn)(CUdeviceptr, unsigned char, size_t);
typedef CUresult (*cuMemcpyHtoD_fn)(CUdeviceptr, const void *, size_t);
typedef CUresult (*cuMemcpyDtoH_fn)(void *, CUdeviceptr, size_t);
typedef CUresult (*cuModuleLoadDataEx_fn)(CUmodule *, const void *, unsigned int, void *, void *);
typedef CUresult (*cuModuleGetFunction_fn)(CUfunction *, CUmodule, const char *);
typedef CUresult (*cuLaunchKernel_fn)(CUfunction,
                                      unsigned int,
                                      unsigned int,
                                      unsigned int,
                                      unsigned int,
                                      unsigned int,
                                      unsigned int,
                                      unsigned int,
                                      CUstream,
                                      void **,
                                      void **);
typedef CUresult (*cuMemGetInfo_fn)(size_t *, size_t *);
typedef CUresult (*cuStreamCreate_fn)(CUstream *, unsigned int);
typedef CUresult (*cuStreamDestroy_fn)(CUstream);
typedef CUresult (*cuGetErrorName_fn)(CUresult, const char **);
typedef CUresult (*cuGetErrorString_fn)(CUresult, const char **);

struct cuda_api {
    void *lib;
    cuInit_fn cuInit;
    cuDeviceGetCount_fn cuDeviceGetCount;
    cuDeviceGet_fn cuDeviceGet;
    cuDeviceGetName_fn cuDeviceGetName;
    cuDeviceGetAttribute_fn cuDeviceGetAttribute;
    cuCtxCreate_fn cuCtxCreate;
    cuCtxDestroy_fn cuCtxDestroy;
    cuCtxSynchronize_fn cuCtxSynchronize;
    cuMemAlloc_fn cuMemAlloc;
    cuMemFree_fn cuMemFree;
    cuMemsetD8_fn cuMemsetD8;
    cuMemcpyHtoD_fn cuMemcpyHtoD;
    cuMemcpyDtoH_fn cuMemcpyDtoH;
    cuModuleLoadDataEx_fn cuModuleLoadDataEx;
    cuModuleGetFunction_fn cuModuleGetFunction;
    cuLaunchKernel_fn cuLaunchKernel;
    cuMemGetInfo_fn cuMemGetInfo;
    cuStreamCreate_fn cuStreamCreate;
    cuStreamDestroy_fn cuStreamDestroy;
    cuGetErrorName_fn cuGetErrorName;
    cuGetErrorString_fn cuGetErrorString;
};

struct stress_report {
    char backend[32];
    char device[128];
    int cc_major;
    int cc_minor;
    int buffer_mb;
    int stream_count;
    unsigned long iterations;
    uint64_t checksum;
    char details[16384];
};

/* Kept as one translation unit so the dynamically loaded CUDA ABI and static
 * helpers remain private, while each backend can be reviewed independently. */
#include "bee-gpu-stress-cuda.inc"
#include "bee-gpu-stress-cublaslt.inc"
#include "bee-gpu-stress-main.inc"