reduce_task.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_REDUCE_TASK_H__INCLUDED
#define TILEDARRAY_REDUCE_TASK_H__INCLUDED

#include <TiledArray/config.h>
#include <TiledArray/error.h>
#include <TiledArray/madness.h>

namespace TiledArray {
  namespace detail {

    template <typename T>
    struct ArgumentHelper {
      typedef Future<T> type;
    }; // struct ArgumentHelper

    template <typename T>
    struct ArgumentHelper<Future<T> > {
      typedef Future<T> type;
    }; // struct ArgumentHelper

    template <typename T, typename U>
    struct ArgumentHelper<std::pair<Future<T>, Future<U> > > {
      typedef std::pair<Future<T>, Future<U> > type;
    }; // struct ArgumentHelper


    template <typename opT>
    class ReducePairOpWrapper {
    public:
      typedef typename opT::result_type result_type;
      typedef typename std::remove_cv<typename std::remove_reference<
          typename opT::first_argument_type>::type>::type first_argument_type;
      typedef typename std::remove_cv<typename std::remove_reference<
          typename opT::second_argument_type>::type>::type second_argument_type;
      typedef std::pair<Future<first_argument_type>,
          Future<second_argument_type> > argument_type;

    private:
      opT op_; 

    public:
      ReducePairOpWrapper() : op_() { }


      ReducePairOpWrapper(const opT& op) : op_(op) { }


      ReducePairOpWrapper(const ReducePairOpWrapper<opT>& other) :
        op_(other.op_)
      { }

      ~ReducePairOpWrapper() { }


      ReducePairOpWrapper<opT>& operator=(const ReducePairOpWrapper<opT>& other) {
        op_ = other.op_;
        return *this;
      }

      result_type operator()() const { return op_(); }

      result_type operator()(result_type& temp) const { return op_(temp); }


      void operator()(result_type& result, const result_type& arg) {
        op_(result, arg);
      }


      void operator()(result_type& result, const argument_type& arg) const {
        op_(result, arg.first, arg.second);
      }

    }; // class ReducePairOpWrapper



    template <typename opT>
    class ReduceTask {
    private:
      typedef typename opT::result_type result_type;
      typedef typename std::remove_const<typename std::remove_reference<
          typename opT::argument_type>::type>::type argument_type;


      class ReduceTaskImpl : public madness::TaskInterface {
      public:


        class ReduceObject : public madness::CallbackInterface {
        private:

          ReduceTaskImpl* parent_; 
          typename ArgumentHelper<argument_type>::type arg_; 
          madness::CallbackInterface* callback_; 
          madness::AtomicInt count_; 


          template <typename T>
          void register_callbacks(Future<T>& f) {
            if(f.probe()) {
              parent_->ready(this);
            } else {
              count_ = 1;
              f.register_callback(this);
            }
          }


          template <typename T, typename U>
          void register_callbacks(std::pair<Future<T>, Future<U> >& p) {
            if(p.first.probe() && p.second.probe()) {
              parent_->ready(this);
            } else {
              count_ = 2;
              p.first.register_callback(this);
              p.second.register_callback(this);
            }
          }

        public:


          template <typename Arg>
          ReduceObject(ReduceTaskImpl* parent, const Arg& arg, madness::CallbackInterface* callback) :
          parent_(parent), arg_(arg), callback_(callback)
          {
            MADNESS_ASSERT(parent_);
            register_callbacks(arg_);
          }

          virtual ~ReduceObject() { }

          virtual void notify() { if((--count_) == 0) parent_->ready(this); }


          const argument_type& arg() const { return arg_; }


          static void destroy(const ReduceObject* object) {
            static constexpr const bool trace_tasks =
#ifdef TILEDARRAY_ENABLE_TASK_DEBUG_TRACE
                true
#else
                false
#endif
            ;
            if(object->callback_) {
              if (trace_tasks)
                object->callback_->notify_debug("destroy(*ReduceObject)");
              else
                object->callback_->notify();
            }
            delete object;
          }

        }; // class ReduceObject

        virtual void get_id(std::pair<void*,unsigned short>& id) const {
          return PoolTaskInterface::make_id(id, *this);
        }


        void reduce(std::shared_ptr<result_type>& result) {
          while(result) {
            lock_.lock(); // <<< Begin critical section
            if(ready_object_) {
              // Get the ready argument
              ReduceObject* ready_object = const_cast<ReduceObject*>(ready_object_);
              ready_object_ = nullptr;
              lock_.unlock(); // <<< End critical section

              // Reduce the argument that was held by ready_object_
              op_(*result, ready_object->arg());

              // cleanup the argument
              ReduceObject::destroy(ready_object);
              this->dec();
            } else if(ready_result_) {
              // Get the ready result
              std::shared_ptr<result_type> ready_result = ready_result_;
              ready_result_.reset();
              lock_.unlock(); // <<< End critical section

              // Reduce the result that was held by ready_result_
              op_(*result, *ready_result);

              // cleanup the result
              ready_result.reset();
            } else {
              // Nothing is ready, so place result in the ready state.
              ready_result_ = result;
              result.reset();
              lock_.unlock(); // <<< End critical section
            }
          }
        }


        void reduce_result_object(std::shared_ptr<result_type> result, const ReduceObject* object) {
          // Reduce the argument
          op_(*result, object->arg());

          // Cleanup the argument
          ReduceObject::destroy(object);

          // Check for more reductions
          reduce(result);

          // Decrement the dependency counter for the argument. This must
          // be done after the reduce call to avoid a race condition.
          this->dec();
        }

        void reduce_object_object(const ReduceObject* object1, const ReduceObject* object2) {
          // Construct an empty result object
          auto result = std::make_shared<result_type>(op_());

          // Reduce the two arguments
          op_(*result, object1->arg());
          op_(*result, object2->arg());

          // Cleanup arguments
          ReduceObject::destroy(object1);
          ReduceObject::destroy(object2);

          // Check for more reductions
          reduce(result);

          // Decrement the dependency counter for the two arguments. This
          // must be done after the reduce call to avoid a race condition.
          this->dec();
          this->dec();
        }

        World& world_; 
        opT op_; 
        std::shared_ptr<result_type> ready_result_; 
        volatile ReduceObject* ready_object_; 
        Future<result_type> result_; 
        madness::Spinlock lock_; 
        madness::CallbackInterface* callback_; 

      public:


        ReduceTaskImpl(World& world, opT op, madness::CallbackInterface* callback) :
          madness::TaskInterface(1, TaskAttributes::hipri()),
          world_(world), op_(op), ready_result_(std::make_shared<result_type>(op())),
          ready_object_(nullptr), result_(), lock_(), callback_(callback)
        { }

        virtual ~ReduceTaskImpl() { }

        virtual void run(const madness::TaskThreadEnv&) {
          MADNESS_ASSERT(ready_result_);
          result_.set(op_(*ready_result_));
          if(callback_)
            callback_->notify();
        }


        void ready(ReduceObject* object) {
          MADNESS_ASSERT(object);
          lock_.lock(); // <<< Begin critical section
          if(ready_result_) {
            std::shared_ptr<result_type> ready_result = ready_result_;
            ready_result_.reset();
            lock_.unlock(); // <<< End critical section
            MADNESS_ASSERT(ready_result);
            world_.taskq.add(this, & ReduceTaskImpl::reduce_result_object,
                ready_result, object, TaskAttributes::hipri());
          } else if(ready_object_) {
            ReduceObject* ready_object = const_cast<ReduceObject*>(ready_object_);
            ready_object_ = nullptr;
            lock_.unlock(); // <<< End critical section
            MADNESS_ASSERT(ready_object);
            world_.taskq.add(this, & ReduceTaskImpl::reduce_object_object,
                object, ready_object, TaskAttributes::hipri());
          } else {
            ready_object_ = object;
            lock_.unlock(); // <<< End critical section
          }
        }


        const Future<result_type>& result() const { return result_; }


        World& world() const { return world_; }

      }; // class ReduceTaskImpl


      ReduceTaskImpl* pimpl_; 
      std::size_t count_; 

    public:

      ReduceTask() : pimpl_(nullptr), count_(0ul) { }



      ReduceTask(World& world, const opT& op = opT(),
          madness::CallbackInterface* callback = nullptr) :
        pimpl_(new ReduceTaskImpl(world, op, callback)), count_(0ul)
      { }


      ReduceTask(ReduceTask<opT>&& other) noexcept :
        pimpl_(other.pimpl_), count_(other.count_)
      {
        other.pimpl_ = nullptr;
        other.count_ = 0ul;
      }


      ~ReduceTask() { delete pimpl_; }


      ReduceTask<opT>& operator=(ReduceTask<opT>&& other) noexcept {
        pimpl_ = other.pimpl_;
        count_ = other.count_;
        other.pimpl_ = nullptr;
        other.count_ = 0;
        return *this;
      }

      // Non-copyable
      ReduceTask(const ReduceTask<opT>&) = delete;
      ReduceTask<opT>& operator=(const ReduceTask<opT>&) = delete;


      template <typename Arg>
      int add(const Arg& arg, madness::CallbackInterface* callback = nullptr) {
        MADNESS_ASSERT(pimpl_);
        pimpl_->inc();
        new typename ReduceTaskImpl::ReduceObject(pimpl_, arg, callback);
        return ++count_;
      }


      int count() const { return count_; }


      Future<result_type> submit() {
        MADNESS_ASSERT(pimpl_);

        // Get the result before submitting/running the task, otherwise the
        // task could run and be deleted before we are done here.
        Future<result_type> result = pimpl_->result();

        pimpl_->dec();
        World& world = pimpl_->world();
        world.taskq.add(pimpl_);

        pimpl_ = nullptr;
        return result;
      }


      operator bool() const { return pimpl_ != nullptr; }

    }; // class ReduceTask




    template <typename opT>
    class ReducePairTask : public ReduceTask<ReducePairOpWrapper<opT> > {
    private:

      typedef ReducePairOpWrapper<opT> op_type; 
      typedef typename op_type::first_argument_type first_argument_type; 
      typedef typename op_type::second_argument_type second_argument_type; 
      typedef typename op_type::argument_type argument_type; 
      typedef ReduceTask<op_type> ReduceTask_; 

    public:

      ReducePairTask() : ReduceTask_() { }


      ReducePairTask(World& world, const opT& op = opT(), madness::CallbackInterface* callback = nullptr) :
        ReduceTask_(world, op_type(op), callback)
      { }


      ReducePairTask(ReducePairTask<opT>&& other) noexcept :
        ReduceTask_(std::move(other))
      { }


      ReducePairTask<opT>& operator=(ReducePairTask<opT>&& other) noexcept {
        ReduceTask_::operator=(std::move(other));
        return *this;
      }

      ReducePairTask(const ReducePairTask<opT>&) = delete;
      ReducePairTask<opT> operator=(const ReducePairTask<opT>&) = delete;


      template <typename L, typename R>
      void add(const L& left, const R& right, madness::CallbackInterface* callback = nullptr) {
        ReduceTask_::add(argument_type(Future<first_argument_type>(left),
            Future<second_argument_type>(right)), callback);
      }

    }; // class ReducePairTask

  } // namespace detail
} // namespace TiledArray

#endif // TILEDARRAY_REDUCE_TASK_H__INCLUDED