wai-quantum 0.6.0

// wai-quantum fabric: the host side of the npu-cores member. The device
// runtime's interface for a neural processor (a hardware context over a
// registered program) is C++, so the member compiles this file with the
// system's C++ compiler against that runtime and loads it; this is the whole
// of what it calls. Every function returns 0 on success, or 1 with the
// runtime's message in `err`.
#include <xrt/xrt_bo.h>
#include <xrt/xrt_device.h>
#include <xrt/xrt_hw_context.h>
#include <xrt/xrt_kernel.h>
#include <xrt/experimental/xrt_xclbin.h>

#include <cstdint>
#include <cstring>
#include <exception>
#include <fstream>
#include <string>
#include <vector>

namespace {

struct Npu {
  xrt::device device;
  xrt::hw_context context;
  xrt::kernel kernel;
  xrt::bo instructions;
  uint32_t count = 0;
};

void say(char *err, size_t len, const std::string &what) {
  if (len == 0) return;
  size_t n = what.size() < len - 1 ? what.size() : len - 1;
  std::memcpy(err, what.data(), n);
  err[n] = 0;
}

std::vector<uint32_t> words_of(const char *path) {
  std::ifstream f(path, std::ios::binary);
  if (!f) throw std::runtime_error(std::string("cannot read ") + path);
  std::vector<char> bytes((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
  if (bytes.empty() || bytes.size() % 4 != 0) throw std::runtime_error(std::string(path) + " is not a whole number of words");
  std::vector<uint32_t> w(bytes.size() / 4);
  std::memcpy(w.data(), bytes.data(), bytes.size());
  return w;
}

}  // namespace

// The device's name as the runtime gives it.
extern "C" int wq_npu_name(char *name, size_t len, char *err, size_t errlen) {
  try {
    xrt::device d(0);
    say(name, len, d.get_info<xrt::info::device::name>());
    return 0;
  } catch (const std::exception &e) {
    say(err, errlen, e.what());
    return 1;
  }
}

// Opens device 0 with the program `xclbin` and its instruction stream
// `insts`, kernel `kernel` (the first whose name starts with it).
extern "C" int wq_npu_open(const char *xclbin, const char *insts, const char *kernel, void **out, char *err, size_t errlen) {
  try {
    auto *n = new Npu();
    n->device = xrt::device(0);
    xrt::xclbin x{std::string(xclbin)};
    std::string name;
    for (auto &k : x.get_kernels()) {
      if (k.get_name().rfind(kernel, 0) == 0) {
        name = k.get_name();
        break;
      }
    }
    if (name.empty()) {
      delete n;
      say(err, errlen, std::string("the program has no kernel named ") + kernel);
      return 1;
    }
    n->device.register_xclbin(x);
    n->context = xrt::hw_context(n->device, x.get_uuid());
    n->kernel = xrt::kernel(n->context, name);
    auto w = words_of(insts);
    n->count = static_cast<uint32_t>(w.size());
    n->instructions = xrt::bo(n->device, w.size() * sizeof(uint32_t), XCL_BO_FLAGS_CACHEABLE, n->kernel.group_id(1));
    std::memcpy(n->instructions.map<void *>(), w.data(), w.size() * sizeof(uint32_t));
    n->instructions.sync(XCL_BO_SYNC_BO_TO_DEVICE);
    *out = n;
    return 0;
  } catch (const std::exception &e) {
    say(err, errlen, e.what());
    return 1;
  }
}

// One dispatch: `in` (in_words) to the program's first buffer, its second
// buffer (out_words) back into `out`.
extern "C" int wq_npu_run(void *h, const int32_t *in, size_t in_words, int32_t *out, size_t out_words, char *err, size_t errlen) {
  try {
    auto *n = static_cast<Npu *>(h);
    xrt::bo a(n->device, in_words * 4, XRT_BO_FLAGS_HOST_ONLY, n->kernel.group_id(3));
    xrt::bo b(n->device, out_words * 4, XRT_BO_FLAGS_HOST_ONLY, n->kernel.group_id(4));
    std::memcpy(a.map<void *>(), in, in_words * 4);
    std::memset(b.map<void *>(), 0, out_words * 4);
    a.sync(XCL_BO_SYNC_BO_TO_DEVICE);
    b.sync(XCL_BO_SYNC_BO_TO_DEVICE);
    unsigned int opcode = 3;
    auto run = n->kernel(opcode, n->instructions, n->count, a, b);
    auto state = run.wait();
    if (state != ERT_CMD_STATE_COMPLETED) {
      say(err, errlen, "the dispatch ended in state " + std::to_string(static_cast<int>(state)));
      return 1;
    }
    b.sync(XCL_BO_SYNC_BO_FROM_DEVICE);
    std::memcpy(out, b.map<void *>(), out_words * 4);
    return 0;
  } catch (const std::exception &e) {
    say(err, errlen, e.what());
    return 1;
  }
}

extern "C" void wq_npu_close(void *h) {
  delete static_cast<Npu *>(h);
}