HaloMatrixWrapper.h

#ifndef DASH__HALO_HALOMATRIXWRAPPER_H
#define DASH__HALO_HALOMATRIXWRAPPER_H

#include <dash/dart/if/dart.h>

#include <dash/Matrix.h>
#include <dash/Pattern.h>
#include <dash/halo/StencilOperator.h>
#include <dash/halo/CoordinateAccess.h>
#include <dash/halo/Types.h>

#include <type_traits>
#include <vector>



namespace dash {

namespace halo {

template <typename MatrixT, SignalReady SigReady = SignalReady::OFF>
class HaloMatrixWrapper {
private:
  using Pattern_t       = typename MatrixT::pattern_type;
  using pattern_index_t = typename Pattern_t::index_type;

  static constexpr auto NumDimensions = Pattern_t::ndim();
  using GlobMem_t = typename MatrixT::GlobMem_t;

public:
  using Element_t = typename MatrixT::value_type;

  using ViewSpec_t      = ViewSpec<NumDimensions, pattern_index_t>;
  using GlobBoundSpec_t = GlobalBoundarySpec<NumDimensions>;
  using HaloBlock_t     = HaloBlock<Element_t, Pattern_t, GlobMem_t>;
  using HaloMemory_t    = HaloMemory<HaloBlock_t>;
  using HaloUpdateEnv_t = HaloUpdateEnv<HaloBlock_t, SigReady>;
  using ElementCoords_t = std::array<pattern_index_t, NumDimensions>;
  using region_index_t  = internal::region_index_t;
  using stencil_dist_t  = internal::spoint_value_t;

private:
  static constexpr auto MemoryArrange = Pattern_t::memory_order();

  using pattern_size_t        = typename Pattern_t::size_type;
  using signed_pattern_size_t = typename std::make_signed<pattern_size_t>::type;
  using HaloSpec_t            = HaloSpec<NumDimensions>;
  using Region_t              = Region<Element_t, Pattern_t, typename MatrixT::GlobMem_t>;

public:
  template <size_t NumStencilPointsFirst, typename... StencilSpecRestT>
  HaloMatrixWrapper(MatrixT& matrix, const GlobBoundSpec_t& glob_bnd_spec,
                    const StencilSpec<StencilPoint<NumDimensions>, NumStencilPointsFirst>& stencil_spec_first,
                    const StencilSpecRestT&... stencil_spec)
  : _matrix(matrix), _glob_bnd_spec(glob_bnd_spec),
    _halo_spec(stencil_spec_first, stencil_spec...),
    _view_global(matrix.local.offsets(), matrix.local.extents()),
    _haloblock(matrix.begin().globmem(), matrix.pattern(), _view_global,
               _halo_spec, glob_bnd_spec),
    _view_local(_haloblock.view_local()),
    //_halomemory(_haloblock),
    _halo_env(_haloblock, matrix.lbegin(), matrix.team(), matrix.pattern().teamspec()) {
  }

  template <typename StencilPointT = StencilPoint<NumDimensions>>
  HaloMatrixWrapper(MatrixT& matrix, const GlobBoundSpec_t& glob_bnd_spec,
                    stencil_dist_t dist, std::enable_if_t<std::is_integral<stencil_dist_t>::value, std::nullptr_t> = nullptr )
  : HaloMatrixWrapper(matrix, glob_bnd_spec, StencilSpecFactory<StencilPointT>::full_stencil_spec(dist)) {
  }

  template <typename StencilPointT = StencilPoint<NumDimensions>>
  HaloMatrixWrapper(MatrixT& matrix, stencil_dist_t dist, std::enable_if_t<std::is_integral<stencil_dist_t>::value, std::nullptr_t> = nullptr )
  : HaloMatrixWrapper(matrix, GlobBoundSpec_t(), StencilSpecFactory<StencilPointT>::full_stencil_spec(dist)) {
  }

  template <typename StencilPointT, std::size_t NumStencilPoints>
  HaloMatrixWrapper(MatrixT& matrix, const StencilSpec<StencilPointT,NumStencilPoints> stencil_spec)
  : HaloMatrixWrapper(matrix, GlobBoundSpec_t(), stencil_spec) {}

  template <typename... StencilSpecT, typename std::enable_if_t<sizeof...(StencilSpecT) >= 2, bool>>
  HaloMatrixWrapper(MatrixT& matrix, const StencilSpecT&... stencil_spec)
  : HaloMatrixWrapper(matrix, GlobBoundSpec_t(), stencil_spec...) {}

  HaloMatrixWrapper() = delete;

  const HaloBlock_t& halo_block() { return _haloblock; }

  void update() {
    _halo_env.update();
  }

  void update_at(region_index_t index) {
    _halo_env.update_at(index);
  }

  void update_async() {
    _halo_env.update_async();
  }

  void update_async_at(region_index_t index) {
    _halo_env.update_async_at(index);
  }

  void wait() {
    _halo_env.wait();
  }

  void wait(region_index_t index) {
    _halo_env.wait(index);
  }

  const ViewSpec_t& view_local() const { return _view_local; }

  HaloUpdateEnv_t& halo_env() { return _halo_env; }

  const HaloUpdateEnv_t& halo_env() const { return _halo_env; }

  MatrixT& matrix() { return _matrix; }

  const MatrixT& matrix() const { return _matrix; }

  template <typename FunctionT>
  void set_custom_halos(FunctionT f) {
    using signed_extent_t = typename std::make_signed<pattern_size_t>::type;
    for(const auto& region : _haloblock.boundary_regions()) {
      if(!region.is_custom_region()) {
        continue;
      }

      const auto& spec    = region.spec();
      std::array<signed_extent_t, NumDimensions> coords_offset{};
      const auto& reg_ext = region.view().extents();
      for(auto d = 0; d < NumDimensions; ++d) {
        if(spec[d] == 0) {
          coords_offset[d] -= reg_ext[d];
          continue;
        }
        if(spec[d] == 2)
          coords_offset[d] = reg_ext[d];
      }

      auto range_mem = _halo_env.halo_memory().range_at(region.index());
      auto it_mem = range_mem.first;
      auto it_reg_end  = region.end();
      DASH_ASSERT_MSG(
          std::distance(range_mem.first, range_mem.second) == region.size(),
          "Range distance of the HaloMemory is unequal region size");
      const auto& pattern = _matrix.pattern();
      for(auto it = region.begin(); it != it_reg_end; ++it, ++it_mem) {
        auto coords = pattern.coords(it.rpos(), region.view());
        for(auto d = 0; d < NumDimensions; ++d) {
          coords[d] += coords_offset[d];
        }

        *it_mem = f(coords);
      }
    }
  }

  Element_t* halo_element_at_global(ElementCoords_t coords) {
    const auto& offsets = _view_global.offsets();
    for(auto d = 0; d < NumDimensions; ++d) {
      coords[d] -= offsets[d];
    }

    return halo_element_at(coords);
  }

  Element_t* halo_element_at_local(ElementCoords_t coords) {
    return halo_element_at(coords);
  }

  template <typename StencilSpecT>
  StencilOperator<HaloBlock_t, StencilSpecT> stencil_operator(
    const StencilSpecT& stencil_spec) {
    for(const auto& stencil : stencil_spec.specs()) {
      DASH_ASSERT_MSG(
        stencil.max()
          <= _halo_spec.extent(RegionSpec<NumDimensions>::index(stencil)),
        "Stencil point extent higher than halo region extent.");
    }

    return StencilOperator<HaloBlock_t, StencilSpecT>(
      &_haloblock, _matrix.lbegin(), &_halo_env.halo_memory(), stencil_spec);
  }

  CoordinateAccess<HaloBlock_t> coordinate_access() {
    return CoordinateAccess<HaloBlock_t>(&_haloblock, _matrix.lbegin(),&_halo_env.halo_memory());
  }

private:

  Element_t* halo_element_at(ElementCoords_t& coords) {
    auto        index     = _haloblock.index_at(_view_local, coords);
    const auto& spec      = _halo_spec.spec(index);
    auto& halo_memory     = _halo_env.halo_memory();
    auto        range_mem = halo_memory.range_at(index);
    if(spec.level() == 0 || range_mem.first == range_mem.second)
      return nullptr;

    if(!halo_memory.to_halo_mem_coords_check(index, coords))
      return nullptr;

    return &*(range_mem.first + halo_memory.offset(index, coords));
  }

private:
  MatrixT&                       _matrix;
  const GlobBoundSpec_t          _glob_bnd_spec;
  const HaloSpec_t               _halo_spec;
  const ViewSpec_t               _view_global;
  const HaloBlock_t              _haloblock;
  const ViewSpec_t&              _view_local;
  HaloUpdateEnv_t                _halo_env;
};

}  // namespace halo

}  // namespace dash

#endif  // DASH__HALO_HALOMATRIXWRAPPER_H