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 <madness/world/parallel_archive.h>
 
#include "TiledArray/array_impl.h"
#include "TiledArray/conversions/clone.h"
#include "TiledArray/conversions/truncate.h"
#include "TiledArray/pmap/replicated_pmap.h"
#include "TiledArray/policies/dense_policy.h"
#include "TiledArray/replicator.h"
#include "TiledArray/tile_interface/cast.h"
#include "TiledArray/util/annotation.h"
#include "TiledArray/util/initializer_list.h"
#include "TiledArray/util/random.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 madness::archive::ParallelSerializableObject {
 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 numeric_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::index1_type
      index1_type;  
  typedef typename impl_type::range_type::index
      index;  
  typedef typename impl_type::ordinal_type ordinal_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;  
 
  typedef typename value_type::value_type element_type;
 
  template <typename OtherTile>
  using is_my_type = std::is_same<DistArray_, DistArray<OtherTile, Policy>>;
 
  template <typename OtherTile>
  using enable_if_not_my_type =
      std::enable_if_t<not is_my_type<OtherTile>::value>;
 
  template <typename Index>
  static constexpr bool is_integral_or_integral_range_v =
      std::is_integral_v<Index> || detail::is_integral_range_v<Index>;
 
  template <typename Index>
  using enable_if_is_integral_or_integral_range =
      std::enable_if_t<is_integral_or_integral_range_v<Index>>;
 
  template <typename Index>
  using enable_if_is_integral = std::enable_if_t<std::is_integral_v<Index>>;
 
  template <typename Value>
  static constexpr bool is_value_or_future_to_value_v =
      std::is_same_v<std::decay_t<Value>, Future<value_type>> ||
      std::is_same_v<std::decay_t<Value>, value_type>;
 
 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_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_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_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_ASSERT(!shape.empty() &&
              "Array::Array() -- The shape is not initialized.");
    TA_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(1, trange), 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 T>
  DistArray(World& world, detail::vector_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
  template <typename T>
  DistArray(World& world, detail::matrix_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
  template <typename T>
  DistArray(World& world, detail::tensor3_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
  template <typename T>
  DistArray(World& world, detail::tensor4_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
  template <typename T>
  DistArray(World& world, detail::tensor5_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
  template <typename T>
  DistArray(World& world, detail::tensor6_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, il)) {}
 
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::vector_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::matrix_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::tensor3_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::tensor4_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::tensor5_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
  template <typename T>
  DistArray(World& world, const trange_type& trange, detail::tensor6_il<T> il)
      : DistArray(array_from_il<DistArray_>(world, trange, il)) {}
 
 
  template <typename OtherTile, typename = enable_if_not_my_type<OtherTile>>
  explicit DistArray(const DistArray<OtherTile, Policy>& other) : pimpl_() {
    *this = foreach<Tile>(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>(other, std::forward<Op>(op));
  }
 
 
  ~DistArray() {}
 
 
  DistArray_ clone() const { return TiledArray::clone(*this); }
 
 
  std::shared_ptr<const impl_type> pimpl() const { return pimpl_; }
 
 
  std::shared_ptr<impl_type> pimpl() { return pimpl_; }
 
 
  std::weak_ptr<const impl_type> weak_pimpl() const { return pimpl_; }
 
 
  std::weak_ptr<impl_type> weak_pimpl() { return pimpl_; }
 
 
  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;
    }
  }
 
 
  void wait_for_lazy_cleanup() const {
    DistArray::wait_for_lazy_cleanup(world());
  }
 
 
  DistArray_& operator=(const DistArray_& other) {
    pimpl_ = other.pimpl_;
 
    return *this;
  }
 
 
  madness::uniqueidT id() const { return impl_ref().id(); }
 
 
  iterator begin() { return impl_ref().begin(); }
 
 
  const_iterator begin() const { return impl_ref().cbegin(); }
 
 
  iterator end() { return impl_ref().end(); }
 
 
  const_iterator end() const { return impl_ref().cend(); }
 
 
  template <typename Index,
            typename = enable_if_is_integral_or_integral_range<Index>>
  Future<value_type> find(const Index& i) const {
    check_index(i);
    return pimpl_->get(i);
  }
 
 
  template <typename Integer, typename = enable_if_is_integral<Integer>>
  Future<value_type> find(const std::initializer_list<Integer>& i) const {
    return find<std::initializer_list<Integer>>(i);
  }
 
 
  template <typename Index,
            typename = std::enable_if_t<std::is_integral_v<Index> ||
                                        detail::is_integral_range_v<Index>>>
  const Future<value_type>& find_local(const Index& i) const {
    check_local_index(i);
    return pimpl_->get_local(i);
  }
 
 
  template <typename Integer,
            typename = std::enable_if_t<(std::is_integral_v<Integer>)>>
  const Future<value_type>& find_local(
      const std::initializer_list<Integer>& i) const {
    return find_local<std::initializer_list<Integer>>(i);
  }
 
 
  template <typename Index,
            typename = std::enable_if_t<std::is_integral_v<Index> ||
                                        detail::is_integral_range_v<Index>>>
  Future<value_type>& find_local(const Index& i) {
    check_local_index(i);
    return pimpl_->get_local(i);
  }
 
 
  template <typename Integer,
            typename = std::enable_if_t<(std::is_integral_v<Integer>)>>
  Future<value_type>& find_local(const std::initializer_list<Integer>& i) {
    return find_local<std::initializer_list<Integer>>(i);
  }
 
  template <
      typename Index, typename InIter,
      typename = std::enable_if_t<is_integral_or_integral_range_v<Index> &&
                                  detail::is_input_iterator<InIter>::value>>
  void 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_t<(std::is_integral_v<Integer>)&&detail::
                                            is_input_iterator<InIter>::value>>
  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,
            typename = enable_if_is_integral_or_integral_range<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, typename = enable_if_is_integral<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, typename Value,
      typename = std::enable_if_t<(is_integral_or_integral_range_v<Index> &&
                                   is_value_or_future_to_value_v<Value>)>>
  void set(const Index& i, Value&& v) {
    check_index(i);
    pimpl_->set(i, std::forward<Value>(v));
  }
 
 
  template <
      typename Index, typename Value,
      typename = std::enable_if_t<
          (std::is_integral_v<Index>)&&is_value_or_future_to_value_v<Value>>>
  void set(const std::initializer_list<Index>& i, Value&& v) {
    set<std::initializer_list<Index>>(i, std::forward<Value>(v));
  }
 
 
  void fill_local(const element_type& value = element_type(),
                  bool skip_set = false) {
    init_tiles(
        [value](const range_type& range) { return value_type(range, value); },
        skip_set);
  }
 
 
  void fill(const element_type& value = numeric_type(), bool skip_set = false) {
    fill_local(value, skip_set);
  }
 
  template <typename T = element_type,
            typename = detail::enable_if_can_make_random_t<T>>
  void fill_random(bool skip_set = false) {
    init_elements(
        [](const auto&) { return detail::MakeRandom<T>::generate_value(); });
  }
 
 
  template <typename Op>
  void init_tiles(Op&& op, bool skip_set = false) {
    // lifetime management of op depends on whether it is a lvalue ref (i.e. has
    // an external owner) or an rvalue ref
    // - if op is an lvalue ref: pass op to tasks
    // - if op is an rvalue ref pass make_shared_function(op) to tasks
    auto op_shared_handle = make_op_shared_handle(std::forward<Op>(op));
 
    auto it = impl_ref().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(
            [pimpl = this->weak_pimpl(), index = ordinal_type(index),
             op_shared_handle]() -> value_type {
              auto pimpl_ptr = pimpl.lock();
              if (pimpl_ptr)
                return op_shared_handle(
                    pimpl_ptr->trange().make_tile_range(index));
              else
                return {};
            });
        set(index, tile);
      }
    }
  }
 
 
  template <typename Op>
  void init_elements(Op&& op, bool skip_set = false) {
    auto op_shared_handle = make_op_shared_handle(std::forward<Op>(op));
    init_tiles(
        [op = std::move(op_shared_handle)](
            const TiledArray::Range& range) -> value_type {
          // Initialize the tile with the given range object
          Tile tile(range);
 
          // Initialize tile elements
          for (auto& idx : range) tile[idx] = op(idx);
 
          return tile;
        },
        skip_set);
  }
 
 
  const trange_type& trange() const { return impl_ref().trange(); }
 
 
  // TODO: uncomment [[deprecated("use DistArray::tiles_range()")]]
  const range_type& range() const { return impl_ref().tiles_range(); }
 
 
  const range_type& tiles_range() const { return impl_ref().tiles_range(); }
 
  [[deprecated("use DistArray::elements_range()")]] const typename trange_type::
      range_type&
      elements() const {
    return elements_range();
  }
 
 
  const typename trange_type::range_type& elements_range() const {
    return impl_ref().trange().elements_range();
  }
 
 
  auto size() const { return impl_ref().size(); }
 
 
  TiledArray::expressions::TsrExpr<const DistArray_, true> operator()(
      const std::string& vars) const {
    check_str_index(vars);
    return TiledArray::expressions::TsrExpr<const DistArray_, true>(*this,
                                                                    vars);
  }
 
 
  TiledArray::expressions::TsrExpr<DistArray_, true> operator()(
      const std::string& vars) {
    check_str_index(vars);
    return TiledArray::expressions::TsrExpr<DistArray_, true>(*this, vars);
  }
 
  [[deprecated]] World& get_world() const { return world(); }
 
 
  World& world() const { return impl_ref().world(); }
 
  [[deprecated]] const std::shared_ptr<pmap_interface>& get_pmap() const {
    return pmap();
  }
 
 
  const std::shared_ptr<pmap_interface>& pmap() const {
    return impl_ref().pmap();
  }
 
 
  bool is_dense() const { return impl_ref().is_dense(); }
 
  [[deprecated]] const shape_type& get_shape() const { return shape(); }
 
 
  inline const shape_type& shape() const { return impl_ref().shape(); }
 
 
  template <typename Index,
            typename = enable_if_is_integral_or_integral_range<Index>>
  ProcessID owner(const Index& i) const {
    check_index(i);
    return pimpl_->owner(i);
  }
 
 
  template <typename Index, typename = enable_if_is_integral<Index>>
  ProcessID owner(const std::initializer_list<Index>& i) const {
    return owner<std::initializer_list<Index>>(i);
  }
 
 
  template <typename Index,
            typename = enable_if_is_integral_or_integral_range<Index>>
  bool is_local(const Index& i) const {
    check_index(i);
    return pimpl_->is_local(i);
  }
 
 
  template <typename Index, typename = enable_if_is_integral<Index>>
  bool is_local(const std::initializer_list<Index>& i) const {
    return is_local<std::initializer_list<Index>>(i);
  }
 
 
  template <typename Index,
            typename = enable_if_is_integral_or_integral_range<Index>>
  bool is_zero(const Index& i) const {
    check_index(i);
    return pimpl_->is_zero(i);
  }
 
 
  template <typename Index, typename = enable_if_is_integral<Index>>
  bool is_zero(const std::initializer_list<Index>& i) const {
    return is_zero<std::initializer_list<Index>>(i);
  }
 
 
  void swap(DistArray_& other) { std::swap(pimpl_, other.pimpl_); }
 
  void make_replicated() {
    if ((!impl_ref().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(
      typename shape_type::value_type thresh = shape_type::threshold()) {
    TiledArray::truncate(*this, thresh);
  }
 
 
  bool is_initialized() const { return static_cast<bool>(pimpl_); }
 
 
  explicit operator bool() const { return is_initialized(); }
 
 
  template <typename Archive,
            typename = std::enable_if_t<
                !Archive::is_parallel_archive &&
                madness::archive::is_output_archive<Archive>::value>>
  void serialize(const Archive& ar) const {
    // serialize array type, world size, rank, and pmap type to be able
    // to ensure same data type and same data distribution expected
    ar& typeid(*this).hash_code() & world().size() & world().rank() & trange() &
        shape() & typeid(pmap().get()).hash_code();
    int64_t count = 0;
    for (auto it = begin(); it != end(); ++it) ++count;
    ar& count;
    for (auto it = begin(); it != end(); ++it) ar & it->get();
  }
 
 
  template <typename Archive,
            typename = std::enable_if_t<
                !Archive::is_parallel_archive &&
                madness::archive::is_input_archive<Archive>::value>>
  void serialize(const Archive& ar) {
    auto& world = TiledArray::get_default_world();
 
    std::size_t typeid_hash = 0l;
    ar& typeid_hash;
    if (typeid_hash != typeid(*this).hash_code())
      TA_EXCEPTION(
          "DistArray::serialize: source DistArray type != this DistArray type");
 
    ProcessID world_size = -1;
    ProcessID world_rank = -1;
    ar& world_size& world_rank;
    if (world_size != world.size() || world_rank != world.rank())
      TA_EXCEPTION(
          "DistArray::serialize: source DistArray world != this DistArray "
          "world");
 
    trange_type trange;
    shape_type shape;
    ar& trange& shape;
 
    // use default pmap, ensure it's the same pmap used to serialize
    auto volume = trange.tiles_range().volume();
    auto pmap = detail::policy_t<DistArray_>::default_pmap(world, volume);
    size_t pmap_hash_code = 0;
    ar& pmap_hash_code;
    if (pmap_hash_code != typeid(pmap.get()).hash_code())
      TA_EXCEPTION(
          "DistArray::serialize: source DistArray pmap != this DistArray pmap");
    pimpl_.reset(
        new impl_type(world, std::move(trange), std::move(shape), pmap));
 
    int64_t count = 0;
    ar& count;
    for (auto it = begin(); it != end(); ++it, --count) {
      Tile tile;
      ar& tile;
      this->set(it.ordinal(), std::move(tile));
    }
    if (count != 0)
      TA_EXCEPTION(
          "DistArray::serialize: # of tiles in archive != # of tiles expected");
  }
 
 
  template <typename Archive>
  void load(World& world, Archive& ar) {
    auto me = world.rank();
    const Tag tag = world.mpi.unique_tag();  // for broadcasting metadata
 
    if (ar.dofence()) world.gop.fence();
 
    if (ar.is_io_node()) {  // on each io node ...
 
      auto& localar = ar.local_archive();
 
      // make sure source data matches the expected type
      // TODO would be nice to be able to convert the data upon reading
      std::size_t typeid_hash = 0l;
      localar& typeid_hash;
      if (typeid_hash != typeid(*this).hash_code())
        TA_EXCEPTION(
            "DistArray::load: source DistArray type != this DistArray type");
 
      // make sure same number of clients for every I/O node
      int num_io_clients = 0;
      localar& num_io_clients;
      if (num_io_clients != ar.num_io_clients())
        TA_EXCEPTION("DistArray::load: invalid parallel archive");
 
      trange_type trange;
      shape_type shape;
      localar& trange& shape;
 
      // send trange and shape to every client
      for (ProcessID p = 0; p < world.size(); ++p) {
        if (p != me && ar.io_node(p) == me) {
          world.mpi.Send(int(1), p, tag);  // Tell client to expect the data
          madness::archive::MPIOutputArchive dest(world, p);
          dest& trange& shape;
          dest.flush();
        }
      }
 
      // use default pmap
      auto volume = trange.tiles_range().volume();
      auto pmap = detail::policy_t<DistArray_>::default_pmap(world, volume);
      pimpl_.reset(
          new impl_type(world, std::move(trange), std::move(shape), pmap));
 
      int64_t count = 0;
      localar& count;
      for (size_t ord = 0; ord != volume; ++ord) {
        if (!is_zero(ord)) {
          auto owner_rank = pmap->owner(ord);
          if (ar.io_node(owner_rank) == me) {
            Tile tile;
            localar& tile;
            this->set(ord, std::move(tile));
            --count;
          }
        }
      }
      if (count != 0)
        TA_EXCEPTION(
            "DistArray::load: # of tiles in archive != # of tiles expected");
    } else {  // non-I/O node still needs to initialize metadata
 
      trange_type trange;
      shape_type shape;
 
      ProcessID p = ar.my_io_node();
      int flag;
      world.mpi.Recv(flag, p, tag);
      TA_ASSERT(flag == 1);
      madness::archive::MPIInputArchive source(world, p);
      source& trange& shape;
 
      // use default pmap
      auto volume = trange.tiles_range().volume();
      auto pmap = detail::policy_t<DistArray_>::default_pmap(world, volume);
      pimpl_.reset(
          new impl_type(world, std::move(trange), std::move(shape), pmap));
    }
 
    if (ar.dofence()) world.gop.fence();
  }
 
 
  template <typename Archive>
  void store(Archive& ar) const {
    auto me = world().rank();
 
    if (ar.dofence()) world().gop.fence();
 
    if (ar.is_io_node()) {  // on each io node ...
      auto& localar = ar.local_archive();
      // ... store metadata first ...
      localar& typeid(*this).hash_code() & ar.num_io_clients() & trange() &
          shape();
      // ... then loop over tiles and dump the data from ranks
      // assigned to this I/O node in order ...
      // for sanity check dump tile count assigned to this I/O node
      const auto volume = trange().tiles_range().volume();
      int64_t count = 0;
      for (size_t ord = 0; ord != volume; ++ord) {
        if (!is_zero(ord)) {
          const auto owner_rank = pmap()->owner(ord);
          if (ar.io_node(owner_rank) == me) {
            ++count;
          }
        }
      }
      localar& count;
      for (size_t ord = 0; ord != volume; ++ord) {
        if (!is_zero(ord)) {
          auto owner_rank = pmap()->owner(ord);
          if (ar.io_node(owner_rank) == me) localar& find(ord).get();
        }
      }
    }  // am I an I/O node?
    if (ar.dofence()) world().gop.fence();
  }
 
 
  static void register_set_notifier(
      std::function<void(const impl_type&, int64_t)> notifier = {}) {
    impl_type::set_notifier_accessor() = notifier;
  }
 
 private:
  template <typename Index>
  std::enable_if_t<std::is_integral_v<Index>, void> check_index(
      const Index i) const {
    TA_ASSERT(
        impl_ref().tiles_range().includes(i) &&
        "The ordinal index used to access an array tile is out of range.");
  }
 
  template <typename Index>
  std::enable_if_t<detail::is_integral_range_v<Index>, void> check_index(
      const Index& i) const {
    TA_ASSERT(
        impl_ref().tiles_range().includes(i) &&
        "The coordinate index used to access an array tile is out of range.");
  }
 
  template <typename Index1>
  void check_index(const std::initializer_list<Index1>& i) const {
    check_index<std::initializer_list<Index1>>(i);
  }
 
  template <typename Index>
  std::enable_if_t<std::is_integral_v<Index>, void> check_local_index(
      const Index i) const {
    check_index(i);
    TA_ASSERT(pimpl_->is_local(i)  // pimpl_ already checked
              &&
              "The ordinal index used to access an array tile is not local.");
  }
 
  template <typename Index>
  std::enable_if_t<detail::is_integral_range_v<Index>, void> check_local_index(
      const Index& i) const {
    check_index(i);
    TA_ASSERT(
        pimpl_->is_local(i)  // pimpl_ already checked
        && "The coordinate index used to access an array tile is not local.");
  }
 
  template <typename Index1>
  void check_local_index(const std::initializer_list<Index1>& i) const {
    check_local_index<std::initializer_list<Index1>>(i);
  }
 
  void check_str_index(const std::string& vars) const {
#if (TA_ASSERT_POLICY != TA_ASSERT_IGNORE)
    // Only check indices if the PIMPL is initialized (okay to not initialize
    // the RHS of an equation)
    if (!is_initialized()) return;
 
    constexpr bool is_tot = detail::is_tensor_of_tensor_v<value_type>;
    const auto rank = range().rank();
    // TODO: Make constexpr and use structured bindings when CUDA supports C++17
    if (is_tot) {
      // Make sure the index is capable of being interpreted as a ToT index
      TA_ASSERT(detail::is_tot_index(vars));
 
      // Rank of outer tiles must match number of outer indices
      // is_tot_index(vars) implies vars.find(';') < vars.size()
      TA_ASSERT(std::count(vars.begin(), vars.begin() + vars.find(';'), ',') +
                    1ul ==
                rank);
 
      // Check inner index rank?
    } else {
      // Better not be a ToT index
      TA_ASSERT(!detail::is_tot_index(vars));
 
      // Number of indices must match rank
      TA_ASSERT(std::count(vars.begin(), vars.end(), ',') + 1ul == rank);
    }
#endif  // NDEBUG
  }
 
  void assert_pimpl() const {
    TA_ASSERT(pimpl_ &&
              "The Array has not been initialized, likely reason: it was "
              "default constructed and used.");
  }
 
  auto& impl_ref() const {
    assert_pimpl();
    return *pimpl_;
  }
 
  auto& impl_ref() {
    assert_pimpl();
    return *pimpl_;
  }
 
};  // class DistArray
 
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;
}
 
template <typename Tile, typename Policy>
auto rank(const DistArray<Tile, Policy>& a) {
  return a.trange().tiles_range().rank();
}
 
template <typename Tile, typename Policy>
size_t volume(const DistArray<Tile, Policy>& a) {
  // this is the number of tiles
  if (a.size() > 0)  // assuming dense shape
    return a.trange().elements_range().volume();
  return 0;
}
 
template <typename Tile, typename Policy>
auto abs_min(const DistArray<Tile, Policy>& a) {
  return a(detail::dummy_annotation(rank(a))).abs_min();
}
 
template <typename Tile, typename Policy>
auto abs_max(const DistArray<Tile, Policy>& a) {
  return a(detail::dummy_annotation(rank(a))).abs_max();
}
 
template <typename Tile, typename Policy>
auto dot(const DistArray<Tile, Policy>& a, const DistArray<Tile, Policy>& b) {
  return (a(detail::dummy_annotation(rank(a)))
              .dot(b(detail::dummy_annotation(rank(b)))))
      .get();
}
 
template <typename Tile, typename Policy>
auto inner_product(const DistArray<Tile, Policy>& a,
                   const DistArray<Tile, Policy>& b) {
  return (a(detail::dummy_annotation(rank(a)))
              .inner_product(b(detail::dummy_annotation(rank(b)))))
      .get();
}
 
template <typename Tile, typename Policy>
auto squared_norm(const DistArray<Tile, Policy>& a) {
  return a(detail::dummy_annotation(rank(a))).squared_norm();
}
 
template <typename Tile, typename Policy>
auto norm2(const DistArray<Tile, Policy>& a) {
  return std::sqrt(squared_norm(a));
}
 
}  // namespace TiledArray
 
// serialization
namespace madness {
namespace archive {
template <class Tile, class Policy>
struct ArchiveLoadImpl<ParallelInputArchive,
                       TiledArray::DistArray<Tile, Policy>> {
  static inline void load(const ParallelInputArchive& ar,
                          TiledArray::DistArray<Tile, Policy>& x) {
    x.load(*ar.get_world(), ar);
  }
};
 
template <class Tile, class Policy>
struct ArchiveStoreImpl<ParallelOutputArchive,
                        TiledArray::DistArray<Tile, Policy>> {
  static inline void store(const ParallelOutputArchive& ar,
                           const TiledArray::DistArray<Tile, Policy>& x) {
    x.store(ar);
  }
};
}  // namespace archive
 
template <class Tile, class Policy>
void save(const TiledArray::DistArray<Tile, Policy>& x,
          const std::string name) {
  archive::ParallelOutputArchive ar2(x.world(), name.c_str(), 1);
  ar2& x;
}
 
template <class Tile, class Policy>
void load(TiledArray::DistArray<Tile, Policy>& x, const std::string name) {
  archive::ParallelInputArchive ar2(x.world(), name.c_str(), 1);
  ar2& x;
}
 
}  // namespace madness
 
#endif  // TILEDARRAY_ARRAY_H__INCLUDED