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/tensor_impl.h>
#include <TiledArray/permutation.h>
#include <TiledArray/perm_index.h>
#include <TiledArray/type_traits.h>

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_::size_type size_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:



      size_type perm_index_to_target(size_type index) const {
        TA_ASSERT(index < TensorImpl_::trange().tiles_range().volume());
        return (source_to_target_ ? source_to_target_(index) : index);
      }


      size_type perm_index_to_source(size_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(size_type i) const = 0;


      virtual void discard_tile(size_type i) const = 0;


      void set_tile(size_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(size_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::size_type size_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() { return pimpl_->eval(); }



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


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


      ProcessID owner(size_type i) const { return pimpl_->owner(i); }


      bool is_local(size_type i) const { return pimpl_->is_local(i); }


      bool is_zero(size_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(size_type i) const { return pimpl_->get_tile(i); }


      virtual void discard(size_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