distributed_storage.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_DISTRIBUTED_STORAGE_H__INCLUDED
#define TILEDARRAY_DISTRIBUTED_STORAGE_H__INCLUDED
 
#include <TiledArray/pmap/pmap.h>
 
namespace TiledArray {
namespace detail {
 
 
template <typename T>
class DistributedStorage : public madness::WorldObject<DistributedStorage<T> > {
 public:
  typedef DistributedStorage<T> DistributedStorage_;  
  typedef madness::WorldObject<DistributedStorage_>
      WorldObject_;  
 
  typedef std::size_t size_type;      
  typedef size_type key_type;         
  typedef T value_type;               
  typedef Future<value_type> future;  
  typedef Pmap pmap_interface;        
  typedef madness::ConcurrentHashMap<key_type, future>
      container_type;  
  typedef typename container_type::accessor
      accessor;  
  typedef typename container_type::const_accessor
      const_accessor;  
 
 private:
  const size_type max_size_;  
  std::shared_ptr<pmap_interface>
      pmap_;  
  mutable container_type data_;     
  madness::AtomicInt num_live_ds_;  
 
  // not allowed
  DistributedStorage(const DistributedStorage_&);
  DistributedStorage_& operator=(const DistributedStorage_&);
 
  void set_handler(const size_type i, const value_type& value) {
    future& f = get_local(i);
 
    // Check that the future has not been set already.
    TA_ASSERT(!f.probe() && "Tile has already been assigned.");
 
    f.set(value);
  }
 
  void get_handler(const size_type i,
                   const typename future::remote_refT& ref) const {
    const future& f = get_local(i);
    future remote_f(ref);
    remote_f.set(f);
  }
 
  void set_remote(const size_type i, const value_type& value) {
    WorldObject_::task(owner(i), &DistributedStorage_::set_handler, i, value,
                       madness::TaskAttributes::hipri());
  }
 
  struct DelayedSet : public madness::CallbackInterface {
   private:
    DistributedStorage_& ds_;  
    size_type index_;          
    future future_;            
 
   public:
    DelayedSet(DistributedStorage_& ds, size_type i, const future& f)
        : ds_(ds), index_(i), future_(f) {
      ++ds_.num_live_ds_;
    }
 
    virtual ~DelayedSet() { --ds_.num_live_ds_; }
 
    virtual void notify() {
      ds_.set_remote(index_, future_);
      delete this;
    }
  };  // struct DelayedSet
  friend struct DelayedSet;
 
 public:
 
  DistributedStorage(World& world, size_type max_size,
                     const std::shared_ptr<pmap_interface>& pmap)
      : WorldObject_(world),
        max_size_(max_size),
        pmap_(pmap),
        data_((max_size / world.size()) + 11) {
    // Check that the process map is appropriate for this storage object
    TA_ASSERT(pmap_);
    TA_ASSERT(pmap_->size() == max_size);
    TA_ASSERT(pmap_->rank() == pmap_interface::size_type(world.rank()));
    TA_ASSERT(pmap_->procs() == pmap_interface::size_type(world.size()));
    num_live_ds_ = 0;
    WorldObject_::process_pending();
  }
 
  virtual ~DistributedStorage() {
    if (num_live_ds_ != 0) {
      madness::print_error(
          "DistributedStorage (object id=\", id(), \") destroyed while "
          "outstanding tasks exist. Add a fence() to extend the lifetime of "
          "this object.");
      abort();
    }
  }
 
  using WorldObject_::get_world;
 
 
  const std::shared_ptr<pmap_interface>& pmap() const { return pmap_; }
 
 
  ProcessID owner(size_type i) const {
    TA_ASSERT(i < max_size_);
    TA_ASSERT(pmap_);
    return pmap_->owner(i);
  }
 
 
  bool is_local(size_type i) const {
    TA_ASSERT(i < max_size_);
    TA_ASSERT(pmap_);
    return pmap_->is_local(i);
  }
 
 
  size_type size() const { return data_.size(); }
 
 
  size_type max_size() const { return max_size_; }
 
 
  future get(size_type i) const {
    TA_ASSERT(i < max_size_);
    if (is_local(i)) {
      return get_local(i);
    } else {
      // Send a request to the owner of i for the element.
      future result;
      WorldObject_::task(owner(i), &DistributedStorage_::get_handler, i,
                         result.remote_ref(get_world()),
                         madness::TaskAttributes::hipri());
 
      return result;
    }
  }
 
 
  const future& get_local(const size_type i) const {
    TA_ASSERT(pmap_->is_local(i));
 
    // Return the local element.
    const_accessor acc;
    [[maybe_unused]] const bool inserted = data_.insert(acc, i);
    return acc->second;
  }
 
 
  future& get_local(const size_type i) {
    TA_ASSERT(pmap_->is_local(i));
 
    // Return the local element.
    accessor acc;
    [[maybe_unused]] const bool inserted = data_.insert(acc, i);
    return acc->second;
  }
 
 
  void set(size_type i, const value_type& value) {
    TA_ASSERT(i < max_size_);
    if (is_local(i))
      set_handler(i, value);
    else
      set_remote(i, value);
  }
 
 
  void set(size_type i, const future& f) {
    TA_ASSERT(i < max_size_);
    if (is_local(i)) {
      const_accessor acc;
      if (!data_.insert(acc, typename container_type::datumT(i, f))) {
        // The element was already in the container, so set it with f.
        future existing_f = acc->second;
        acc.release();
 
        // Check that the future has not been set already.
        TA_ASSERT(!existing_f.probe() && "Tile has already been assigned.");
        // Set the future
        existing_f.set(f);
      }
    } else {
      if (f.probe()) {
        set_remote(i, f);
      } else {
        DelayedSet* set_callback = new DelayedSet(*this, i, f);
        const_cast<future&>(f).register_callback(set_callback);
      }
    }
  }
 
};  // class DistributedStorage
 
}  // namespace detail
}  // namespace TiledArray
 
#endif  // TILEDARRAY_DISTRIBUTED_STORAGE_H__INCLUDED