dist_array.h

Go to the documentation of this file.

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2013  Virginia Tech
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifndef TILEDARRAY_ARRAY_H__INCLUDED
#define TILEDARRAY_ARRAY_H__INCLUDED

#include <cstdlib>

#include <TiledArray/replicator.h>
#include <TiledArray/pmap/replicated_pmap.h>
//#include <TiledArray/tensor.h>
#include <TiledArray/policies/dense_policy.h>
#include <TiledArray/array_impl.h>
#include <TiledArray/conversions/truncate.h>
#include <TiledArray/conversions/clone.h>
#include <TiledArray/tile_interface/cast.h>

namespace TiledArray {

  // Forward declarations
  template <typename, typename> class Tensor;
  namespace expressions {
    template <typename, bool> class TsrExpr;
  } // namespace expressions



  template <typename Tile = Tensor<double, Eigen::aligned_allocator<double> >,
      typename Policy = DensePolicy>
  class DistArray {
  public:
    typedef DistArray<Tile, Policy> DistArray_; 
    typedef TiledArray::detail::ArrayImpl<Tile, Policy> impl_type;
    typedef typename impl_type::policy_type policy_type; 
    typedef typename detail::numeric_type<Tile>::type element_type; 
    typedef typename detail::scalar_type<Tile>::type scalar_type; 
    typedef typename impl_type::trange_type trange_type; 
    typedef typename impl_type::range_type range_type; 
    typedef typename impl_type::shape_type shape_type; 
    typedef typename impl_type::range_type::index index; 
    typedef typename impl_type::size_type size_type; 
    typedef typename impl_type::value_type value_type; 
    typedef typename impl_type::eval_type eval_type; 
    typedef typename impl_type::reference future; 
    typedef typename impl_type::reference reference; 
    typedef typename impl_type::const_reference const_reference; 
    typedef typename impl_type::iterator iterator; 
    typedef typename impl_type::const_iterator const_iterator; 
    typedef typename impl_type::pmap_interface pmap_interface; 

  private:

    std::shared_ptr<impl_type> pimpl_; 

    static madness::AtomicInt cleanup_counter_;


    static void lazy_deleter(const impl_type* const pimpl) {
      if(pimpl) {
        if(madness::initialized()) {
          World& world = pimpl->world();
          const madness::uniqueidT id = pimpl->id();
          cleanup_counter_++;

          try {
            world.gop.lazy_sync(id, [pimpl]() {
              delete pimpl;
              DistArray_::cleanup_counter_--;
            });
          }
          catch(madness::MadnessException& e) {
            fprintf(stderr, "!! ERROR TiledArray: madness::MadnessException thrown in Array::lazy_deleter().\n"
                            "%s\n"
                            "!! ERROR TiledArray: The exception has been absorbed.\n"
                            "!! ERROR TiledArray: rank=%i\n", e.what(), world.rank());

            cleanup_counter_--;
            delete pimpl;
          }
          catch(std::exception& e) {
            fprintf(stderr, "!! ERROR TiledArray: std::exception thrown in Array::lazy_deleter().\n"
                            "%s\n"
                            "!! ERROR TiledArray: The exception has been absorbed.\n"
                            "!! ERROR TiledArray: rank=%i\n", e.what(), world.rank());

            cleanup_counter_--;
            delete pimpl;
          }
          catch(...) {
            fprintf(stderr, "!! ERROR TiledArray: An unknown exception was thrown in Array::lazy_deleter().\n"
                            "!! ERROR TiledArray: The exception has been absorbed.\n"
                            "!! ERROR TiledArray: rank=%i\n", world.rank());

            cleanup_counter_--;
            delete pimpl;
          }
        } else {
          delete pimpl;
        }
      }
    }


    static std::shared_ptr<impl_type>
    init(World& world, const trange_type& trange, const shape_type& shape,
        std::shared_ptr<pmap_interface> pmap)
    {
      // User level validation of input

      if(! pmap) {
        // Construct a default process map
        pmap = Policy::default_pmap(world, trange.tiles_range().volume());
      } else {
        // Validate the process map
        TA_USER_ASSERT(pmap->size() == trange.tiles_range().volume(),
            "Array::Array() -- The size of the process map is not equal to the number of tiles in the TiledRange object.");
        TA_USER_ASSERT(pmap->rank() == typename pmap_interface::size_type(world.rank()),
            "Array::Array() -- The rank of the process map is not equal to that of the world object.");
        TA_USER_ASSERT(pmap->procs() == typename pmap_interface::size_type(world.size()),
            "Array::Array() -- The number of processes in the process map is not equal to that of the world object.");
      }

      // Validate the shape
      TA_USER_ASSERT(! shape.empty(),
          "Array::Array() -- The shape is not initialized.");
      TA_USER_ASSERT(shape.validate(trange.tiles_range()),
          "Array::Array() -- The range of the shape is not equal to the tiles range.");

      return std::shared_ptr<impl_type>(new impl_type(world, trange, shape, pmap), lazy_deleter);
    }

  public:


    DistArray() : pimpl_() { }


    DistArray(const DistArray_& other) : pimpl_(other.pimpl_) { }


    DistArray(World& world, const trange_type& trange,
        const std::shared_ptr<pmap_interface>& pmap = std::shared_ptr<pmap_interface>()) :
      pimpl_(init(world, trange, shape_type(), pmap))
    { }


    DistArray(World& world, const trange_type& trange, const shape_type& shape,
        const std::shared_ptr<pmap_interface>& pmap = std::shared_ptr<pmap_interface>()) :
      pimpl_(init(world, trange, shape, pmap))
    { }


    template <typename OtherTile, typename = std::enable_if_t<not std::is_same<DistArray_,DistArray<OtherTile,Policy>>::value>>
    explicit DistArray(const DistArray<OtherTile,Policy>& other) :
        pimpl_()
    {
      *this = foreach<Tile,OtherTile>(other, [](Tile& result, const OtherTile& source) { result = TiledArray::Cast<Tile,OtherTile>{}(source); });
    }


    template <typename OtherTile, typename Op>
    DistArray(const DistArray<OtherTile,Policy>& other, Op&& op) :
      pimpl_()
    {
      *this = foreach<Tile,OtherTile,Op>(other, std::forward<Op>(op));
    }


    ~DistArray() { }


    DistArray_ clone() const {
      return TiledArray::clone(*this);
    }


    static void wait_for_lazy_cleanup(World& world,
                                      const double = 60.0)
    {
      try {
        world.await([&]() { return (cleanup_counter_ == 0); }, true);
      } catch(...) {
        printf("%i: Array lazy cleanup timeout with %i pending cleanup(s)\n",
            world.rank(), int(cleanup_counter_));
        throw;
      }
    }


    DistArray_& operator=(const DistArray_& other) {
      pimpl_ = other.pimpl_;

      return *this;
    }


    madness::uniqueidT id() const { return pimpl_->id(); }


    iterator begin() {
      check_pimpl();
      return pimpl_->begin();
    }


    const_iterator begin() const {
      check_pimpl();
      return pimpl_->cbegin();
    }


    iterator end() {
      check_pimpl();
      return pimpl_->end();
    }


    const_iterator end() const {
      check_pimpl();
      return pimpl_->cend();
    }


    template <typename Index>
    Future<value_type> find(const Index& i) const {
      check_index(i);
      return pimpl_->get(i);
    }


    template <typename Integer>
    Future<value_type> find(const std::initializer_list<Integer>& i) const {
      return find<std::initializer_list<Integer>>(i);
    }


    template <typename Index, typename InIter>
    typename std::enable_if<detail::is_input_iterator<InIter>::value>::type
    set(const Index& i, InIter first) {
      check_index(i);
      pimpl_->set(i, value_type(pimpl_->trange().make_tile_range(i), first));
    }


    template <typename Integer, typename InIter>
    typename std::enable_if<detail::is_input_iterator<InIter>::value>::type
    set(const std::initializer_list<Integer>& i, InIter first) {
      set<std::initializer_list<Integer>>(i, first);
    }


    template <typename Index>
    void set(const Index& i, const element_type& value = element_type()) {
      check_index(i);
      pimpl_->set(i, value_type(pimpl_->trange().make_tile_range(i), value));
    }


    template <typename Integer>
    void set(const std::initializer_list<Integer>& i,
             const element_type& value = element_type()) {
      set<std::initializer_list<Integer>>(i, value);
    }


    template <typename Index>
    void set(const Index& i, const Future<value_type>& f) {
      check_index(i);
      pimpl_->set(i, f);
    }


    template <typename Integer>
    void set(const std::initializer_list<Integer>& i,
             const Future<value_type>& f) {
      set<std::initializer_list<Integer>>(i, f);
    }


    template <typename Index>
    void set(const Index& i, const value_type& v) {
      check_index(i);
      pimpl_->set(i, v);
    }


    template <typename Integer>
    void set(const std::initializer_list<Integer>& i, const value_type& v) {
      set<std::initializer_list<Integer>>(i, v);
    }


    void fill_local(const element_type& value = element_type(), bool skip_set = false) {
      init_tiles([=] (const range_type& range)
          { return value_type(range, value); }, skip_set);
    }


    void fill(const element_type& value = element_type(), bool skip_set = false) {
      fill_local(value, skip_set);
    }

    void fill_random(bool skip_set = false) {
      init_elements([](const auto &) {
        return (element_type)std::rand() / RAND_MAX;
      });
    }


    template <typename Op>
    void init_tiles(Op&& op, bool skip_set = false) {
      check_pimpl();

      auto it = pimpl_->pmap()->begin();
      const auto end = pimpl_->pmap()->end();
      for(; it != end; ++it) {
        const auto index = *it;
        if(! pimpl_->is_zero(index)) {
          if (skip_set) {
            auto fut = find(index);
            if (fut.probe())
              continue;
          }
          Future<value_type> tile = pimpl_->world().taskq.add(
              [] (DistArray_* array, const size_type index, const Op& op) -> value_type
              { return op(array->trange().make_tile_range(index)); },
              this, index, op);
          set(index, tile);
        }
      }
    }


    template <typename Op>
    void init_elements(Op&& op, bool skip_set = false) {
      init_tiles([op] (const TiledArray::Range& range) -> value_type
      {
        // Initialize the tile with the given range object
        TiledArray::Tensor<double> tile(range);

        // Initialize tile elements
        for(auto& idx: range)
          tile[idx] = op(idx);

        return tile;
       });
    }


    const trange_type& trange() const {
      check_pimpl();
      return pimpl_->trange();
    }


    const range_type& range() const {
      check_pimpl();
      return pimpl_->tiles_range();
    }

    DEPRECATED const typename trange_type::range_type& elements() const {
      return elements_range();
    }


    const typename trange_type::range_type& elements_range() const {
      check_pimpl();
      return pimpl_->trange().elements_range();
    }

    size_type size() const {
      check_pimpl();
      return pimpl_->size();
    }


    TiledArray::expressions::TsrExpr<const DistArray_, true>
    operator ()(const std::string& vars) const {
#ifndef NDEBUG
      const unsigned int n = 1u + std::count_if(vars.begin(), vars.end(),
          [](const char c) { return c == ','; });
      if(bool(pimpl_) && n != pimpl_->trange().tiles_range().rank()) {
        if(TiledArray::get_default_world().rank() == 0) {
          TA_USER_ERROR_MESSAGE( \
              "The number of array annotation variables is not equal to the array dimension:" \
              << "\n    number of variables  = " << n \
              << "\n    array dimension      = " << pimpl_->trange().tiles_range().rank() );
        }

        TA_EXCEPTION("The number of array annotation variables is not equal to the array dimension.");
      }
#endif // NDEBUG
      return TiledArray::expressions::TsrExpr<const DistArray_, true>(*this, vars);
    }


    TiledArray::expressions::TsrExpr<DistArray_, true>
    operator ()(const std::string& vars) {
#ifndef NDEBUG
      const unsigned int n = 1u + std::count_if(vars.begin(), vars.end(),
          [](const char c) { return c == ','; });
      if(bool(pimpl_) && n != pimpl_->trange().tiles_range().rank()) {
        if(TiledArray::get_default_world().rank() == 0) {
          TA_USER_ERROR_MESSAGE( \
              "The number of array annotation variables is not equal to the array dimension:" \
              << "\n    number of variables  = " << n \
              << "\n    array dimension      = " << pimpl_->trange().tiles_range().rank() );
        }

        TA_EXCEPTION("The number of array annotation variables is not equal to the array dimension.");
      }
#endif // NDEBUG
      return TiledArray::expressions::TsrExpr<DistArray_, true>(*this, vars);
    }

    DEPRECATED World& get_world() const {
      check_pimpl();
      return pimpl_->world();
    }


    World& world() const {
      check_pimpl();
      return pimpl_->world();
    }

    DEPRECATED const std::shared_ptr<pmap_interface>& get_pmap() const {
      check_pimpl();
      return pimpl_->pmap();
    }


    const std::shared_ptr<pmap_interface>& pmap() const {
      check_pimpl();
      return pimpl_->pmap();
    }


    bool is_dense() const {
      check_pimpl();
      return pimpl_->is_dense();
    }

    DEPRECATED const shape_type& get_shape() const {  return pimpl_->shape(); }


    inline const shape_type& shape() const {  return pimpl_->shape(); }


    template <typename Index>
    ProcessID owner(const Index& i) const {
      check_index(i);
      return pimpl_->owner(i);
    }


    template <typename Index1>
    ProcessID owner(const std::initializer_list<Index1>& i) const {
      return owner<std::initializer_list<Index1>>(i);
    }


    template <typename Index>
    bool is_local(const Index& i) const {
      check_index(i);
      return pimpl_->is_local(i);
    }


    template <typename Index1>
    bool is_local(const std::initializer_list<Index1>& i) const {
      return is_local<std::initializer_list<Index1>>(i);
    }


    template <typename Index>
    bool is_zero(const Index& i) const {
      check_index(i);
      return pimpl_->is_zero(i);
    }


    template <typename Index1>
    bool is_zero(const std::initializer_list<Index1>& i) const {
      return is_zero<std::initializer_list<Index1>>(i);
    }


    void swap(DistArray_& other) { std::swap(pimpl_, other.pimpl_); }

    void make_replicated() {
      check_pimpl();
      if((! pimpl_->pmap()->is_replicated()) && (world().size() > 1)) {
        // Construct a replicated array
        auto pmap = std::make_shared<detail::ReplicatedPmap>(world(), size());
        DistArray_ result = DistArray_(world(), trange(), shape(), pmap);

        // Create the replicator object that will do an all-to-all broadcast of
        // the local tile data.
        auto replicator =
            std::make_shared<detail::Replicator<DistArray_>>(*this, result);

        // Put the replicator pointer in the deferred cleanup object so it will
        // be deleted at the end of the next fence.
        TA_ASSERT(replicator.unique()); // Required for deferred_cleanup
        madness::detail::deferred_cleanup(world(), replicator);

        DistArray_::operator=(result);
      }
    }


    void truncate() { TiledArray::truncate(*this); }


    bool is_initialized() const { return static_cast<bool>(pimpl_); }

  private:

    template <typename Index>
    typename std::enable_if<std::is_integral<Index>::value>::type
    check_index(const Index i) const {
      check_pimpl();
      TA_USER_ASSERT(pimpl_->tiles_range().includes(i),
          "The ordinal index used to access an array tile is out of range.");
    }

    template <typename Index>
    typename std::enable_if<! std::is_integral<Index>::value>::type
    check_index(const Index& i) const {
      check_pimpl();
      TA_USER_ASSERT(pimpl_->tiles_range().includes(i),
          "The coordinate index used to access an array tile is out of range.");
      TA_USER_ASSERT(i.size() == pimpl_->trange().tiles_range().rank(),
          "The number of elements in the coordinate index does not match the dimension of the array.");
    }

    template <typename Index1>
    void check_index(const std::initializer_list<Index1>& i) const {
      check_index<std::initializer_list<Index1>>(i);
    }

    void check_pimpl() const {
      TA_USER_ASSERT(pimpl_,
          "The Array has not been initialized, likely reason: it was default constructed and used.");
    }

  }; // class Array


  template <typename Tile, typename Policy>
  madness::AtomicInt DistArray<Tile, Policy>::cleanup_counter_;

#ifndef TILEDARRAY_HEADER_ONLY

  extern template
  class DistArray<Tensor<double, Eigen::aligned_allocator<double> >, DensePolicy>;
  extern template
  class DistArray<Tensor<float, Eigen::aligned_allocator<float> >, DensePolicy>;
  extern template
  class DistArray<Tensor<int, Eigen::aligned_allocator<int> >, DensePolicy>;
  extern template
  class DistArray<Tensor<long, Eigen::aligned_allocator<long> >, DensePolicy>;
//  extern template
//  class DistArray<Tensor<std::complex<double>, Eigen::aligned_allocator<std::complex<double> > >, DensePolicy>;
//  extern template
//  class DistArray<Tensor<std::complex<float>, Eigen::aligned_allocator<std::complex<float> > >, DensePolicy>

  extern template
  class DistArray<Tensor<double, Eigen::aligned_allocator<double> >, SparsePolicy>;
  extern template
  class DistArray<Tensor<float, Eigen::aligned_allocator<float> >, SparsePolicy>;
  extern template
  class DistArray<Tensor<int, Eigen::aligned_allocator<int> >, SparsePolicy>;
  extern template
  class DistArray<Tensor<long, Eigen::aligned_allocator<long> >, SparsePolicy>;
//  extern template
//  class DistArray<Tensor<std::complex<double>, Eigen::aligned_allocator<std::complex<double> > >, SparsePolicy>;
//  extern template
//  class DistArray<Tensor<std::complex<float>, Eigen::aligned_allocator<std::complex<float> > >, SparsePolicy>;

#endif // TILEDARRAY_HEADER_ONLY


  template <typename Tile, typename Policy>
  inline std::ostream& operator<<(std::ostream& os, const DistArray<Tile, Policy>& a) {
    if(a.world().rank() == 0) {
      for(std::size_t i = 0; i < a.size(); ++i)
        if(! a.is_zero(i)) {
          const typename DistArray<Tile, Policy>::value_type tile = a.find(i).get();
          os << i << ": " << tile  << "\n";
        }
    }
    a.world().gop.fence();
    return os;
  }

} // namespace TiledArray

#endif // TILEDARRAY_ARRAY_H__INCLUDED