// Cublas.upload: The first n floats of an array copied to the device; answers the buffer id. The array is consumed.

// bend-blas: the platform CBLAS and LAPACK, opened once per process
// (Accelerate on macOS; OpenBLAS, else libcblas + libblas + liblapack,
// on Linux). Bend's clang command is fixed, so nothing is linked: every
// symbol comes through dlsym.
#ifndef BEND_BLAS_LOADER
#define BEND_BLAS_LOADER 1
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dlfcn.h>
#include <math.h>

static void* bend_lib_open(const char** names) {
  void* lib = NULL;
  for (int i = 0; names[i] != NULL && lib == NULL; i += 1) {
    lib = dlopen(names[i], RTLD_NOW);
  }
  return lib;
}

// a symbol from the BLAS, the CBLAS wrapper beside a reference BLAS, or
// the LAPACK library, in that order; NULL when none holds it
static void* bend_blas_sym(const char* name) {
  static void* libs[3] = { NULL, NULL, NULL };
  static int tried = 0;
  if (!tried) {
    const char* blas[] = {
      "/System/Library/Frameworks/Accelerate.framework/Accelerate",
      "libopenblas.so.0", "libopenblas.so", "libblas.so.3", "libblas.so", NULL };
    const char* cblas[] = { "libcblas.so.3", "libcblas.so", NULL };
    const char* lapack[] = { "liblapack.so.3", "liblapack.so", NULL };
    tried = 1;
    libs[0] = bend_lib_open(blas);
    libs[1] = bend_lib_open(cblas);
    libs[2] = bend_lib_open(lapack);
  }
  for (int i = 0; i < 3; i += 1) {
    void* sym = libs[i] == NULL ? NULL : dlsym(libs[i], name);
    if (sym != NULL) {
      return sym;
    }
  }
  return NULL;
}

// an F32 argument is the low 32 bits of its word; an F32 answer the same
static float bend_f32_of(Term t) {
  u32 bits = (u32)t;
  float x;
  memcpy(&x, &bits, 4);
  return x;
}

static Term bend_f32_term(float x) {
  u32 bits;
  memcpy(&bits, &x, 4);
  return (Term)bits;
}

// a fresh block of the smallest power-of-two size holding n floats, the
// tail zero. Answers 0 past 2^31 floats, the largest block class; the
// heap itself never answers, a full heap is a fail-stop of the runtime.
static u64 bend_block_floats(Env e, u64 n, u32* cls_out) {
  u32 cls = 0;
  while ((1ull << cls) < n) {
    cls += 1;
  }
  if (cls > 31) {
    return 0;
  }
  u64 loc = heap_alloc(e, buf_wcls(cls));
  float* c = (float*)blk_ptr(e.mem, loc, 0);
  for (u64 i = n; i < (1ull << cls); i += 1) {
    c[i] = 0.0f;
  }
  *cls_out = cls;
  return loc;
}

// a fresh block holding the first n floats of an operand
static u64 bend_block_copy(Env e, const float* src, u64 n, u32* cls_out) {
  u64 loc = bend_block_floats(e, n, cls_out);
  if (loc != 0) {
    memcpy(blk_ptr(e.mem, loc, 0), src, (size_t)n * sizeof(float));
  }
  return loc;
}

// a row-major m x n operand copied to a column-major buffer (malloc)
static float* bend_colmajor(const float* a, u32 m, u32 n) {
  float* c = (float*)malloc((size_t)m * n * sizeof(float) + 4);
  for (u32 i = 0; i < m; i += 1) {
    for (u32 j = 0; j < n; j += 1) {
      c[(size_t)j * m + i] = a[(size_t)i * n + j];
    }
  }
  return c;
}

// a symmetric operand copied as it is: its column-major view is itself
static float* bend_symcopy(const float* a, u32 n) {
  float* c = (float*)malloc((size_t)n * n * sizeof(float) + 4);
  memcpy(c, a, (size_t)n * n * sizeof(float));
  return c;
}

// a column-major m x n buffer (leading dimension ld) into a fresh
// row-major block
static u64 bend_block_rowmajor(Env e, const float* c, u32 m, u32 n, u32 ld, u32* cls_out) {
  u64 loc = bend_block_floats(e, (u64)m * n, cls_out);
  if (loc != 0) {
    float* r = (float*)blk_ptr(e.mem, loc, 0);
    for (u32 i = 0; i < m; i += 1) {
      for (u32 j = 0; j < n; j += 1) {
        r[(size_t)i * n + j] = c[(size_t)j * ld + i];
      }
    }
  }
  return loc;
}

// the capacity of an Array<F32> in floats: its block class is in the term
#define BEND_CAP(t) (1ull << (u32)term_aux(t))
#define BEND_ARR(t) ((const float*)blk_ptr(e.mem, term_loc(t), 0))
// an operand must hold the floats the sizes claim, and a size product
// must fit a block: both are checked before anything is read
#define BEND_NEED(t, n) \
  if ((u64)(n) > 0x7fffffffull) { return io_fail(e, 2, "bend-blas: a size past 2^31 floats"); } \
  if (BEND_CAP(t) < (u64)(n)) { return io_fail(e, 2, "bend-blas: an operand smaller than its sizes claim"); }
#define BEND_SYM(var, type, name) \
  type var = (type)bend_blas_sym(name); \
  if (var == NULL) { return io_fail(e, 2, "bend-blas: no BLAS or LAPACK found (Accelerate, OpenBLAS, libcblas, liblapack): " name); }
#define BEND_OUT(loc, cls, n) \
  u32 cls = 0; \
  u64 loc = bend_block_floats(e, (u64)(n), &cls); \
  if (loc == 0) { return io_fail(e, 2, "bend-blas: a result past 2^31 floats"); }
#define BEND_OUT_COPY(loc, cls, src, n) \
  u32 cls = 0; \
  u64 loc = bend_block_copy(e, (src), (u64)(n), &cls); \
  if (loc == 0) { return io_fail(e, 2, "bend-blas: a result past 2^31 floats"); }
#define BEND_DONE(loc, cls) return io_done(e, term_blk(false, cls, loc))
#endif

// bend-blas: cuBLAS, opened at runtime with dlopen (libcudart, libcublas),
// no CUDA headers needed to build; on a machine without them every
// routine answers Fail{(1, ..)}. cuBLAS is column-major, so a row-major
// product is computed as C^T = op(B)^T op(A)^T: the operands and their
// flags swap. Device buffers live in a table by id.
#ifndef BEND_CUBLAS_LOADER
#define BEND_CUBLAS_LOADER 1

typedef int (*bend_cuda_malloc_fn)(void** p, size_t n);
typedef int (*bend_cuda_free_fn)(void* p);
typedef int (*bend_cuda_memcpy_fn)(void* dst, const void* src, size_t n, int kind);
typedef int (*bend_cublas_create_fn)(void** h);
typedef int (*bend_cublas_sgemm_fn)(void* h, int ta, int tb, int m, int n, int k,
  const float* alpha, const float* a, int lda, const float* b, int ldb,
  const float* beta, float* c, int ldc);

static bend_cuda_malloc_fn  bend_cuda_malloc   = NULL;
static bend_cuda_free_fn    bend_cuda_free     = NULL;
static bend_cuda_memcpy_fn  bend_cuda_memcpy   = NULL;
static bend_cublas_sgemm_fn bend_cublas_sgemm  = NULL;
static void*                bend_cublas_handle = NULL;
static int                  bend_cublas_state  = 0;   // 0 untried, 1 open, 2 absent

static int bend_cublas_open(void) {
  if (bend_cublas_state != 0) {
    return bend_cublas_state == 1 ? 0 : 1;
  }
  bend_cublas_state = 2;
  const char* rts[] = { "libcudart.so", "libcudart.so.12", "libcudart.so.13", NULL };
  const char* bls[] = { "libcublas.so", "libcublas.so.12", "libcublas.so.13", NULL };
  void* rt = bend_lib_open(rts);
  void* bl = bend_lib_open(bls);
  if (rt == NULL || bl == NULL) {
    return 1;
  }
  bend_cuda_malloc = (bend_cuda_malloc_fn)dlsym(rt, "cudaMalloc");
  bend_cuda_free   = (bend_cuda_free_fn)dlsym(rt, "cudaFree");
  bend_cuda_memcpy = (bend_cuda_memcpy_fn)dlsym(rt, "cudaMemcpy");
  bend_cublas_create_fn create = (bend_cublas_create_fn)dlsym(bl, "cublasCreate_v2");
  bend_cublas_sgemm = (bend_cublas_sgemm_fn)dlsym(bl, "cublasSgemm_v2");
  if (bend_cuda_malloc == NULL || bend_cuda_free == NULL || bend_cuda_memcpy == NULL
      || create == NULL || bend_cublas_sgemm == NULL || create(&bend_cublas_handle) != 0) {
    bend_cublas_sgemm = NULL;
    return 1;
  }
  bend_cublas_state = 1;
  return 0;
}

// a device copy of n floats, or NULL; memcpy kinds: host to device 1,
// device to host 2
static void* bend_cuda_upload(const float* p, u64 n) {
  void* d = NULL;
  size_t bytes = (size_t)n * sizeof(float) + 4;
  if (bend_cuda_malloc(&d, bytes) != 0) {
    return NULL;
  }
  if (bend_cuda_memcpy(d, p, (size_t)n * sizeof(float), 1) != 0) {
    bend_cuda_free(d);
    return NULL;
  }
  return d;
}

// the device buffer table
typedef struct BendDev { void* p; u64 n; } BendDev;
static BendDev* bend_devs     = NULL;
static u32      bend_devs_len = 0;
static u32      bend_devs_cap = 0;

static u32 bend_dev_add(void* p, u64 n) {
  if (bend_devs_len == bend_devs_cap) {
    bend_devs_cap = bend_devs_cap == 0 ? 16 : bend_devs_cap * 2;
    bend_devs = (BendDev*)realloc(bend_devs, bend_devs_cap * sizeof(BendDev));
  }
  bend_devs[bend_devs_len].p = p;
  bend_devs[bend_devs_len].n = n;
  bend_devs_len += 1;
  return bend_devs_len - 1;
}

static BendDev* bend_dev_get(u32 id) {
  return id < bend_devs_len && bend_devs[id].p != NULL ? &bend_devs[id] : NULL;
}

// C (m x n, row-major) = alpha op(A) op(B) + beta C with da, db device
// pointers and c a host buffer holding C in and out; answers 0 on
// success. cublasOperation_t: N 0, T 1.
static int bend_cublas_run(u32 ta, u32 tb, u32 m, u32 n, u32 k, float alpha, float beta, const void* da, const void* db, float* c) {
  void* dc = NULL;
  size_t fc = (size_t)m * n * sizeof(float);
  if (bend_cuda_malloc(&dc, fc + 4) != 0) {
    return 3;
  }
  int ok = beta == 0.0f || bend_cuda_memcpy(dc, c, fc, 1) == 0;
  ok = ok && bend_cublas_sgemm(bend_cublas_handle, tb ? 1 : 0, ta ? 1 : 0, (int)n, (int)m, (int)k, &alpha,
    (const float*)db, tb ? (int)k : (int)n, (const float*)da, ta ? (int)m : (int)k, &beta, (float*)dc, (int)n) == 0;
  ok = ok && bend_cuda_memcpy(c, dc, fc, 2) == 0;
  bend_cuda_free(dc);
  return ok ? 0 : 3;
}

#define BEND_CUDA_OPEN() \
  if (bend_cublas_open() != 0) { return io_fail(e, 1, "bend-blas: cuBLAS not found (libcudart, libcublas)"); }
#endif

static Term cublas_upload_body(Env e, Term* f, IoWork* w) {
  u32 n = (u32)f[0];
  BEND_NEED(f[1], n);
  (void)n;
  BEND_CUDA_OPEN();
  void* d = bend_cuda_upload(BEND_ARR(f[1]), n);
  if (d == NULL) { return io_fail(e, 3, "bend-blas: cudaMalloc or cudaMemcpy failed"); }
  return io_done(e, (Term)bend_dev_add(d, n));
}

// the runtime does not drop an effect's arguments: every array operand
// is dropped here, after the body has read it
Term cublas_upload_run(Env e, Term* f, IoWork* w) {
  Term r = cublas_upload_body(e, f, w);
  term_drop(e, f[1]);
  return r;
}

static void __attribute__((constructor)) cublas_upload_use(void) {
  io_eff(CID(Cublas.upload), cublas_upload_run, 0);
}