ArrayHandler.h

#ifndef LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_ARRAYHANDLER_H
#define LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_ARRAYHANDLER_H
#include <algorithm>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <numeric>
#include <set>
#include <stdexcept>
#include <string>
#include <vector>
 
#include <molpro/linalg/array/type_traits.h>
#include <molpro/linalg/itsolv/subspace/Matrix.h>
#include <molpro/linalg/scalar_traits.h>
#include <molpro/linalg/itsolv/wrap_util.h>
 
namespace molpro::linalg::array {
using molpro::linalg::itsolv::CVecRef;
using molpro::linalg::itsolv::VecRef;
using molpro::linalg::itsolv::subspace::Matrix;
 
namespace util {
 
struct ArrayHandlerError : public std::logic_error {
  using std::logic_error::logic_error;
};
 
template <typename... Args>
struct OperationRegister {
  using OP = std::tuple<Args...>;
  std::list<std::tuple<Args...>> m_register; 
 
  template <int N, class ArgEqual>
  void push(const Args &...args, ArgEqual equal) {
    auto &&new_op = OP{args...};
    auto &ref = std::get<N>(new_op);
    auto rend = std::find_if(m_register.rbegin(), m_register.rend(), [&ref, &equal](const auto &el) {
      auto xx = std::get<N>(el);
      return equal(ref, xx);
    });
    auto end_of_group = m_register.end();
    if (rend != m_register.rend())
      end_of_group = rend.base();
    m_register.insert(end_of_group, std::forward<OP>(new_op));
  }
 
  void push(const Args &...args) { m_register.push_back({args...}); }
 
  bool empty() { return m_register.empty(); }
  void clear() { m_register.clear(); }
};
 
template <typename X, typename Y, typename Z, class EqualX, class EqualY, class EqualZ>
std::tuple<std::vector<std::tuple<size_t, size_t, size_t>>, std::vector<X>, std::vector<Y>, std::vector<Z>>
remove_duplicates(const std::list<std::tuple<X, Y, Z>> &reg, EqualX equal_x, EqualY equal_y, EqualZ equal_z) {
  auto n_op = reg.size();
  std::vector<std::tuple<size_t, size_t, size_t>> op_register;
  op_register.reserve(n_op);
  std::vector<X> xx;
  std::vector<Y> yy;
  std::vector<Z> zz;
  for (const auto &op : reg) {
    auto x = std::get<0>(op);
    auto y = std::get<1>(op);
    auto z = std::get<2>(op);
    auto it_x = std::find_if(cbegin(xx), cend(xx), [&x, &equal_x](const auto &el) { return equal_x(x, el); });
    auto it_y = std::find_if(cbegin(yy), cend(yy), [&y, &equal_y](const auto &el) { return equal_y(y, el); });
    auto it_z = std::find_if(cbegin(zz), cend(zz), [&z, &equal_z](const auto &el) { return equal_z(z, el); });
    auto ix = distance(cbegin(xx), it_x);
    auto iy = distance(cbegin(yy), it_y);
    auto iz = distance(cbegin(zz), it_z);
    if (it_x == cend(xx))
      xx.push_back(x);
    if (it_y == cend(yy))
      yy.push_back(y);
    if (it_z == cend(zz))
      zz.push_back(z);
    op_register.emplace_back(ix, iy, iz);
  }
  return {op_register, xx, yy, zz};
}
 
template <typename T = int>
struct RefEqual {
  bool operator()(const std::reference_wrapper<T> &l, const std::reference_wrapper<T> &r) {
    return std::addressof(l.get()) == std::addressof(r.get());
  }
};
} // namespace util
 
template <class AL, class AR = AL>
class ArrayHandler {
protected:
  ArrayHandler() : m_counter(std::make_unique<Counter>()){};
  ArrayHandler(const ArrayHandler &) = default;
 
  struct Counter {
    int scal = 0;
    int dot = 0;
    int axpy = 0;
    int copy = 0;
    int gemm_inner = 0;
    int gemm_outer = 0;
  };
 
  std::unique_ptr<Counter> m_counter;
 
public:
  using value_type_L = typename array::mapped_or_value_type_t<AL>;
  using value_type_R = typename array::mapped_or_value_type_t<AR>;
  using value_type = decltype(value_type_L{} * value_type_R{});
  using value_type_abs = decltype(check_abs<value_type>());
 
  virtual AL copy(const AR &source) = 0;
  virtual void copy(AL &x, const AR &y) = 0;
  virtual void scal(value_type alpha, AL &x) = 0;
  virtual void fill(value_type alpha, AL &x) = 0;
  virtual void axpy(value_type alpha, const AR &x, AL &y) = 0;
  virtual value_type dot(const AL &x, const AR &y) = 0;
 
  virtual void gemm_outer(const Matrix<value_type> alphas, const CVecRef<AR> &xx, const VecRef<AL> &yy) = 0;
 
  virtual Matrix<value_type> gemm_inner(const CVecRef<AL> &xx, const CVecRef<AR> &yy) = 0;
 
  virtual std::map<size_t, value_type_abs> select_max_dot(size_t n, const AL &x, const AR &y) = 0;
 
  virtual std::map<size_t, value_type> select(size_t n, const AL &x, bool max = false, bool ignore_sign = false) = 0;
 
  const Counter &counter() const { return *m_counter; }
 
  std::string counter_to_string(std::string L, std::string R) {
    std::string output = "";
    if (m_counter->scal > 0)
      output.append(std::to_string(m_counter->scal) + " scaling operations of the " + L + " vectors, ");
    if (m_counter->copy > 0)
      output.append(std::to_string(m_counter->copy) + " " + L + "<-" + R + " copy operations, ");
    if (m_counter->dot > 0)
      output.append(std::to_string(m_counter->dot) + " dot product operations between the " + L + " and " + R +
                    " vectors, ");
    if (m_counter->axpy > 0)
      output.append(std::to_string(m_counter->axpy) + " axpy (" + L + " = a*" + R + " + " + L + ") operations, ");
    if (m_counter->gemm_inner > 0)
      output.append(std::to_string(m_counter->gemm_inner) + " gemm_inner operations between the " + L + " and " + R +
                    " vectors, ");
    if (m_counter->gemm_outer > 0)
      output.append(std::to_string(m_counter->gemm_outer) + " gemm_outer operations between the " + L + " and " + R +
                    " vectors, ");
    return output;
  };
 
  void clear_counter() {
    m_counter->scal = 0;
    m_counter->copy = 0;
    m_counter->dot = 0;
    m_counter->axpy = 0;
    m_counter->gemm_inner = 0;
    m_counter->gemm_outer = 0;
  }
 
  virtual ~ArrayHandler() {
    std::for_each(m_lazy_handles.begin(), m_lazy_handles.end(), [](auto &el) {
      if (auto handle = el.lock())
        handle->invalidate();
    });
  }
 
protected:
  virtual void error(const std::string &message) { throw util::ArrayHandlerError{message}; };
 
  virtual void fused_axpy(const std::vector<std::tuple<size_t, size_t, size_t>> &reg,
                          const std::vector<value_type> &alphas,
                          const std::vector<std::reference_wrapper<const AR>> &xx,
                          std::vector<std::reference_wrapper<AL>> &yy) {
    for (const auto &i : reg) {
      size_t ai, xi, yi;
      std::tie(ai, xi, yi) = i;
      axpy(alphas[ai], xx[xi].get(), yy[yi].get());
    }
  }
 
  virtual void fused_dot(const std::vector<std::tuple<size_t, size_t, size_t>> &reg,
                         const std::vector<std::reference_wrapper<const AL>> &xx,
                         const std::vector<std::reference_wrapper<const AR>> &yy,
                         std::vector<std::reference_wrapper<value_type>> &out) {
    for (const auto &i : reg) {
      size_t xi, yi, zi;
      std::tie(xi, yi, zi) = i;
      out[zi].get() = dot(xx[xi].get(), yy[yi].get());
    }
  }
 
  class LazyHandle {
  public:
    using value_type = ArrayHandler<AL, AR>::value_type;
    template <typename T>
    using ref_wrap = std::reference_wrapper<T>;
 
  protected:
    std::set<std::string> m_op_types;
    util::OperationRegister<value_type, ref_wrap<const AR>, ref_wrap<AL>> m_axpy;
    util::OperationRegister<ref_wrap<const AL>, ref_wrap<const AR>, ref_wrap<value_type>> m_dot;
 
    void error(std::string message) { m_handler.error(message); };
 
    virtual bool register_op_type(const std::string &type) {
      if (m_op_types.count(type) == 0 && !m_op_types.empty())
        return false;
      m_op_types.insert(type);
      return true;
    }
 
    void clear() {
      m_op_types.clear();
      m_axpy.clear();
      m_dot.clear();
    }
 
  public:
    explicit LazyHandle(ArrayHandler<AL, AR> &handler) : m_handler{handler} {}
    virtual ~LazyHandle() { LazyHandle::eval(); }
 
    virtual void axpy(value_type alpha, const AR &x, AL &y) {
      if (register_op_type("axpy"))
        m_axpy.push(alpha, std::cref(x), std::ref(y));
      else
        error("Failed to register operation type axpy with the current state of the LazyHandle");
    }
    virtual void dot(const AL &x, const AR &y, value_type &out) {
      if (register_op_type("dotLR"))
        m_dot.push(std::cref(x), std::cref(y), std::ref(out));
      else
        error("Failed to register operation type dot with the current state of the LazyHandle");
    }
 
    virtual void eval() {
      if (m_invalid)
        return;
      if (!m_axpy.empty()) {
        auto reg = util::remove_duplicates<value_type, ref_wrap<const AR>, ref_wrap<AL>, std::equal_to<value_type>,
                                           util::RefEqual<const AR>, util::RefEqual<AL>>(m_axpy.m_register, {}, {}, {});
        m_handler.fused_axpy(std::get<0>(reg), std::get<1>(reg), std::get<2>(reg), std::get<3>(reg));
      }
      if (!m_dot.empty()) {
        auto reg =
            util::remove_duplicates<ref_wrap<const AL>, ref_wrap<const AR>, ref_wrap<value_type>,
                                    util::RefEqual<const AL>, util::RefEqual<const AR>, util::RefEqual<value_type>>(
                m_dot.m_register, {}, {}, {});
        m_handler.fused_dot(std::get<0>(reg), std::get<1>(reg), std::get<2>(reg), std::get<3>(reg));
      }
      clear();
    }
 
    void invalidate() { m_invalid = true; }
    bool invalid() { return m_invalid; }
 
  protected:
    ArrayHandler<AL, AR> &m_handler; 
    bool m_invalid = false;          
  };
 
  class ProxyHandle {
  public:
    ProxyHandle(std::shared_ptr<LazyHandle> handle) : m_lazy_handle{std::move(handle)} {}
 
    template <typename... Args>
    void axpy(Args &&...args) {
      m_lazy_handle->axpy(std::forward<Args>(args)...);
      if (m_off)
        eval();
    }
    template <typename... Args>
    void dot(Args &&...args) {
      m_lazy_handle->dot(std::forward<Args>(args)...);
      if (m_off)
        eval();
    }
    void eval() { m_lazy_handle->eval(); }
    void invalidate() { m_lazy_handle->invalidate(); }
    bool invalid() { return m_lazy_handle->invalid(); }
 
    void off() { m_off = true; };
    void on() { m_off = false; };
    bool is_off() { return m_off; }
 
  protected:
    std::shared_ptr<LazyHandle> m_lazy_handle;
    bool m_off = false; 
  };
 
  std::vector<std::weak_ptr<LazyHandle>> m_lazy_handles; 
 
  void save_handle(const std::shared_ptr<LazyHandle> &handle) {
    auto empty_handle =
        std::find_if(m_lazy_handles.begin(), m_lazy_handles.end(), [](const auto &el) { return el.expired(); });
    if (empty_handle == m_lazy_handles.end())
      m_lazy_handles.push_back(handle);
    else
      *empty_handle = handle;
  }
 
  ProxyHandle lazy_handle(ArrayHandler<AL, AR> &handler) {
    auto handle = std::make_shared<typename ArrayHandler<AL, AR>::LazyHandle>(handler);
    save_handle(handle);
    return handle;
  };
 
public:
  virtual ProxyHandle lazy_handle() = 0;
};
 
} // namespace molpro::linalg::array
 
#endif // LINEARALGEBRA_SRC_MOLPRO_LINALG_ARRAY_ARRAYHANDLER_H