dist_eval.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_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED
#define TILEDARRAY_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED
 
#include <TiledArray/config.h>
#include <TiledArray/perm_index.h>
#include <TiledArray/permutation.h>
#include <TiledArray/tensor_impl.h>
#include <TiledArray/type_traits.h>
#ifdef TILEDARRAY_HAS_CUDA
#include <TiledArray/cuda/cuda_task_fn.h>
#include <TiledArray/external/cuda.h>
#endif
 
namespace TiledArray {
namespace detail {
 
 
template <typename Tile, typename Policy>
class DistEvalImpl : public TensorImpl<Policy>,
                     public madness::CallbackInterface {
 public:
  typedef DistEvalImpl<Tile, Policy> DistEvalImpl_;  
  typedef TiledArray::detail::TensorImpl<Policy> TensorImpl_;
 
  typedef typename TensorImpl_::ordinal_type ordinal_type;  
  typedef typename TensorImpl_::trange_type
      trange_type;  
  typedef
      typename TensorImpl_::range_type range_type;  
  typedef typename TensorImpl_::shape_type shape_type;  
  typedef typename TensorImpl_::pmap_interface
      pmap_interface;       
  typedef Tile value_type;  
  typedef typename eval_trait<value_type>::type
      eval_type;  
 
 private:
  madness::uniqueidT id_;       
  PermIndex source_to_target_;  
  PermIndex target_to_source_;  
 
  // The following variables are used to track the total number of tasks run
  // on the local node, task_count_, and the number of tiles set on this
  // node, set_counter_. They are used to track the progress of work done by
  // this node, which allows us to wait for the completion of these tasks
  // without waiting for all tasks.
 
  volatile int task_count_;         
  madness::AtomicInt set_counter_;  
 
 protected:
 
  ordinal_type perm_index_to_target(ordinal_type index) const {
    TA_ASSERT(index < TensorImpl_::trange().tiles_range().volume());
    return (source_to_target_ ? source_to_target_(index) : index);
  }
 
 
  ordinal_type perm_index_to_source(ordinal_type index) const {
    TA_ASSERT(index < TensorImpl_::trange().tiles_range().volume());
    return (target_to_source_ ? target_to_source_(index) : index);
  }
 
 public:
 
  DistEvalImpl(World& world, const trange_type& trange, const shape_type& shape,
               const std::shared_ptr<pmap_interface>& pmap,
               const Permutation& perm)
      : TensorImpl_(world, trange, shape, pmap),
        id_(world.unique_obj_id()),
        source_to_target_(),
        target_to_source_(),
        task_count_(-1),
        set_counter_() {
    set_counter_ = 0;
 
    if (perm) {
      Permutation inv_perm(-perm);
      range_type source_range = inv_perm * trange.tiles_range();
      source_to_target_ = PermIndex(source_range, perm);
      target_to_source_ = PermIndex(trange.tiles_range(), inv_perm);
    }
  }
 
  virtual ~DistEvalImpl() {}
 
 
  const madness::uniqueidT& id() const { return id_; }
 
 
  virtual Future<value_type> get_tile(ordinal_type i) const = 0;
 
 
  virtual void discard_tile(ordinal_type i) const = 0;
 
 
  void set_tile(ordinal_type i, const value_type& value) {
    // Store value
    madness::DistributedID id(id_, i);
    TensorImpl_::world().gop.send(TensorImpl_::owner(i), id, value);
 
    // Record the assignment of a tile
    DistEvalImpl_::notify();
  }
 
 
  void set_tile(ordinal_type i, Future<value_type> f) {
    // Store value
    madness::DistributedID id(id_, i);
    TensorImpl_::world().gop.send(TensorImpl_::owner(i), id, f);
 
    // Record the assignment of a tile
    f.register_callback(this);
  }
 
  virtual void notify() { set_counter_++; }
 
  void wait() const {
    const int task_count = task_count_;
    if (task_count > 0) {
      auto report_and_abort = [&, this](const char* type,
                                        const char* what = nullptr) {
        std::stringstream ss;
        ss << "!! ERROR TiledArray: Aborting due to " << type << " exception.\n"
           << (what != nullptr ? "!! ERROR TiledArray: " : "")
           << (what != nullptr ? what : "") << (what != nullptr ? "\n" : "")
           << "!! ERROR TiledArray: rank=" << TensorImpl_::world().rank()
           << " id=" << id_ << " " << set_counter_ << " of " << task_count
           << " tiles set" << std::endl;
        std::cerr << ss.str().c_str();
        abort();
      };
      try {
        TensorImpl_::world().await(
            [this, task_count]() { return this->set_counter_ == task_count; });
      } catch (TiledArray::Exception& e) {
        report_and_abort("TiledArray", e.what());
      } catch (madness::MadnessException& e) {
        report_and_abort("MADNESS", e.what());
      } catch (SafeMPI::Exception& e) {
        report_and_abort("SafeMPI", e.what());
      } catch (std::exception& e) {
        report_and_abort("std", e.what());
      } catch (...) {
        report_and_abort("", nullptr);
      }
    }
  }
 
 private:
 
  virtual int internal_eval() = 0;
 
 public:
 
  void eval() {
    TA_ASSERT(task_count_ == -1);
    task_count_ = this->internal_eval();
    TA_ASSERT(task_count_ >= 0);
  }
 
};  // class DistEvalImpl
 
 
template <typename Tile, typename Policy>
class DistEval {
 public:
  typedef DistEval<Tile, Policy> DistEval_;  
  typedef DistEvalImpl<Tile, Policy>
      impl_type;  
  typedef typename impl_type::ordinal_type ordinal_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::pmap_interface
      pmap_interface;  
  typedef typename impl_type::value_type value_type;  
  typedef typename impl_type::eval_type eval_type;    
  typedef Future<value_type> future;                  
 
 private:
  std::shared_ptr<impl_type> pimpl_;  
 
 public:
 
  template <typename Impl>
  DistEval(const std::shared_ptr<Impl>& pimpl)
      : pimpl_(std::static_pointer_cast<impl_type>(pimpl)) {
    TA_ASSERT(pimpl_);
  }
 
 
  DistEval(const DistEval_& other) : pimpl_(other.pimpl_) {}
 
 
  DistEval_& operator=(const DistEval_& other) {
    pimpl_ = other.pimpl_;
    return *this;
  }
 
 
  void eval() { pimpl_->eval(); }
 
 
  const range_type& range() const { return pimpl_->tiles_range(); }
 
 
  ordinal_type size() const { return pimpl_->size(); }
 
 
  ProcessID owner(ordinal_type i) const { return pimpl_->owner(i); }
 
 
  bool is_local(ordinal_type i) const { return pimpl_->is_local(i); }
 
 
  bool is_zero(ordinal_type i) const { return pimpl_->is_zero(i); }
 
 
  const std::shared_ptr<pmap_interface>& pmap() const { return pimpl_->pmap(); }
 
 
  bool is_dense() const { return pimpl_->is_dense(); }
 
 
  const shape_type& shape() const { return pimpl_->shape(); }
 
 
  const trange_type& trange() const { return pimpl_->trange(); }
 
 
  future get(ordinal_type i) const { return pimpl_->get_tile(i); }
 
 
  virtual void discard(ordinal_type i) const { pimpl_->discard_tile(i); }
 
 
  World& world() const { return pimpl_->world(); }
 
 
  madness::uniqueidT id() const { return pimpl_->id(); }
 
  void wait() const { pimpl_->wait(); }
 
};  // class DistEval
 
}  // namespace detail
}  // namespace TiledArray
 
#endif  // TILEDARRAY_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED