writer_adios2.hxx

#pragma once

#include <kg/io.h>

#include "fields3d.inl"

#include <cstdio>

// #define PSC_USE_IO_THREADS

#ifdef PSC_USE_IO_THREADS

#include <thread>
#include <mutex>
#include <condition_variable>
#include <queue>

static std::mutex writer_mutex;

#endif

class WriterADIOS2
{
#ifdef PSC_USE_IO_THREADS
  using task_type = std::function<void(void)>;
#endif

public:
  WriterADIOS2()
  {
    MPI_Comm_dup(MPI_COMM_WORLD, &comm_);
#ifdef PSC_USE_IO_THREADS
    writer_thread_ = std::thread(&WriterADIOS2::thread_func, this);
#endif
  }

  WriterADIOS2(const WriterADIOS2&) = delete;
  WriterADIOS2& operator=(const WriterADIOS2&) = delete;

  ~WriterADIOS2()
  {
#ifdef PSC_USE_IO_THREADS
    std::unique_lock<std::mutex> lock(queue_lock_);
    exit_ = true;
    lock.unlock();
    cv_.notify_all();

    if (writer_thread_.joinable()) {
      writer_thread_.join();
    }
#endif
    MPI_Comm_free(&comm_);
  }

  explicit operator bool() const { return pfx_.size() != 0; }

  void open(const std::string& pfx, const std::string& dir = ".")
  {
    assert(pfx_.size() == 0);
    pfx_ = pfx;
    dir_ = dir;
  }

  void close()
  {
    assert(pfx_.size() != 0);
    pfx_.clear();
    dir_.clear();
  }

  void begin_step(const Grid_t& grid)
  {
    _begin_step(file_, grid.timestep(), grid.time(), grid.domain.length,
                grid.domain.corner, grid.domain.gdims);
  }

  void end_step() { _end_step(file_); }

  template <typename E>
  void write(const E& expr, const Grid_t& grid, const std::string& name,
             const std::vector<std::string>& comp_names)
  {
    static int pr_write, pr_eval;
    if (!pr_write) {
      pr_write = prof_register("adios2_write", 1., 0, 0);
      pr_eval = prof_register("adios2_eval", 1., 0, 0);
    }

    prof_start(pr_eval);
    // FIXME, constness hack
    auto&& evaluated = gt::eval(const_cast<E&>(expr));
    auto&& h_expr = gt::host_mirror(evaluated);
    gt::copy(evaluated, h_expr);
    prof_stop(pr_eval);

    prof_start(pr_write);
    std::vector<Int3> patch_off(grid.n_patches());
    for (int p = 0; p < grid.n_patches(); p++) {
      patch_off[p] = grid.patches[p].off;
    }
    write(file_, grid.ldims, grid.domain.gdims, patch_off, h_expr, name,
          comp_names);
    prof_stop(pr_write);
  }

  template <typename E>
  static void write(kg::io::Engine& file, const Int3& ldims, const Int3& gdims,
                    const std::vector<Int3>& patch_off, const E& h_expr,
                    const std::string& name,
                    const std::vector<std::string>& comp_names)
  {
#if 0
    file.prefixes_.push_back(name);
    write_4d(file, ldims, gdims, patch_off, h_expr);
    file.prefixes_.pop_back();
#else
    int n_comps = h_expr.shape(3);
    assert(n_comps == comp_names.size());

    for (int m = 0; m < n_comps; m++) {
      write_3d(file, ldims, gdims, patch_off,
               h_expr.view(_all, _all, _all, m, _all), comp_names[m]);
    }
#endif
  }

  template <typename E>
  void write_step(const Grid_t& grid, const Int3& rn, const Int3& rx,
                  const E& expr, const std::string& name,
                  const std::vector<std::string>& comp_names)
  {

    static int pr, pr_copy, pr_write, pr_thread, pr_lock;
    if (!pr) {
      pr = prof_register("write_step", 1., 0, 0);
      pr_copy = prof_register("ws copy", 1., 0, 0);
      pr_write = prof_register("ws write", 1., 0, 0);
      pr_thread = prof_register("ws thread", 1., 0, 0);
      pr_lock = prof_register("ws lock", 1., 0, 0);
    }

    prof_start(pr);

    // FIXME? rn, rx ignored
    prof_start(pr_copy);
    // FIXME, constness hack
    auto&& evaluated = gt::eval(const_cast<E&>(expr));
    auto&& h_expr = gt::host_mirror(evaluated);
    gt::copy(evaluated, h_expr);
    prof_stop(pr_copy);

#ifdef PSC_USE_IO_THREADS
    prof_start(pr_write);
    int step = grid.timestep();
    double time = grid.time();

    assert(grid.n_patches() == h_expr.shape(4));
    Int3 ldims = grid.ldims;
    Int3 gdims = grid.domain.gdims;
    int n_patches = grid.n_patches();
    std::vector<Int3> patch_off(n_patches);
    for (int p = 0; p < n_patches; p++) {
      patch_off[p] = grid.patches[p].off;
    }
    Double3 length = grid.domain.length;
    Double3 corner = grid.domain.corner;

    auto write_func = [this, step, time, h_expr = std::move(h_expr), name,
                       comp_names, ldims, gdims, length, corner,
                       patch_off = std::move(patch_off)]() {
      // std::this_thread::sleep_for(std::chrono::milliseconds(1000));

      prof_start(pr_thread);

      {
        // FIXME not sure how necessary this lock really is, it certainly could
        // spin for a long time if another thread is writing another file
        prof_start(pr_lock);
        // std::lock_guard<std::mutex> guard(writer_mutex);
        prof_stop(pr_lock);

        prof_start(pr_write);
        kg::io::Engine file;
        _begin_step(file, step, time, length, corner, gdims);
        write(file, ldims, gdims, patch_off, h_expr, name, comp_names);
        _end_step(file);
        prof_stop(pr_write);
      }

      prof_stop(pr_thread);
    };

    std::unique_lock<std::mutex> lock(queue_lock_);
    queue_.push(write_func);
    lock.unlock();
    cv_.notify_one();
#else
    prof_start(pr_write);
    begin_step(grid);
    write(expr, grid, name, comp_names);
    end_step();
    prof_stop(pr_write);
#endif
    prof_stop(pr);
  }

private:
  void _begin_step(kg::io::Engine& file, int step, double time,
                   const Double3& length, const Double3& corner,
                   const Int3& gdims)
  {
    int len = dir_.size() + pfx_.size() + 20;
    char filename[len];
    snprintf(filename, len, "%s/%s.%09d.bp", dir_.c_str(), pfx_.c_str(), step);
    file = io_.open(filename, kg::io::Mode::Write, comm_, pfx_);

    file.beginStep(kg::io::StepMode::Append);
    file.put("psc_output_version", std::string("1.0.0"));
    file.put("step", step);
    file.put("time", time);
    // additionally, write time as a variable (above writes it as an attribute)
    // so that xarray-adios2 can use it as a dimension coordinate
    file.file_.putVariable("time", time);
    file.put("length", length);
    file.put("corner", corner);
    file.put("step_dimension", std::string("time"));

    if (file.mpi_rank_ == 0) {
      Double3 dx = length / Double3(gdims);
      for (int d = 0; d < 3; d++) {
        auto crd = gt::zeros<double>({gdims[d]});
        for (int i = 0; i < gdims[d]; i++) {
          crd(i) = corner[d] + (i + .5) * dx[d];
        }
        std::string name = d == 0 ? "x" : d == 1 ? "y" : "z";
        file.prefixes_.push_back(name);
        file.putVariable(crd.data(), kg::io::Mode::Blocking,
                         kg::io::Dims({size_t(gdims[d])}), {}, {});
        file.prefixes_.pop_back();
        file.put(std::string(name) + "/dimensions", name);
      }
    }

    file_.performPuts();
  }

  static void _end_step(kg::io::Engine& file)
  {
    file.endStep();
    file.close();
  }

#ifdef PSC_USE_IO_THREADS
  void thread_func()
  {
    std::unique_lock<std::mutex> lock(queue_lock_);

    do {
      cv_.wait(lock, [this] { return queue_.size() || exit_; });

      if (!exit_ && queue_.size()) {
        auto task = std::move(queue_.front());
        queue_.pop();

        lock.unlock();
        task();
        lock.lock();
      }
    } while (!exit_);
  }
#endif

  // Our writer thread may be writing one file via adios2 at the same time that
  // another thread is writing another file, and adios2 isn't thread safe at
  // all, so both threads may be making MPI calls that interfere with each
  // other, which we prevent by using a unique communicator.
  MPI_Comm comm_;
  kg::io::IOAdios2 io_;
  kg::io::Engine file_;
  std::string pfx_;
  std::string dir_;
#ifdef PSC_USE_IO_THREADS
  std::thread writer_thread_;
  std::queue<task_type> queue_;
  std::mutex queue_lock_;
  std::condition_variable cv_;
  bool exit_ = false;
#endif
};