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/external/madness.h>
 
#ifdef TILEDARRAY_HAS_CUDA
#include <TiledArray/cuda/cuda_task_fn.h>
#include <TiledArray/external/cuda.h>
#include <TiledArray/tensor/type_traits.h>
#include <TiledArray/util/time.h>
#endif
 
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) {
        TA_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
 
#ifdef TILEDARRAY_HAS_CUDA
 
    static void CUDART_CB cuda_reduceobject_delete_callback(void* userData) {
      const auto t0 = TiledArray::now();
 
      std::vector<void*>* objects = static_cast<std::vector<void*>*>(userData);
 
      madness::World* world = static_cast<madness::World*>((*objects)[0]);
 
      auto destroy_vector = [](std::vector<void*>* objects) {
        std::size_t n_objects = objects->size();
        for (std::size_t i = 1; i < n_objects; i++) {
          // convert void* to ReduceObject*
          ReduceObject* reduce_object =
              static_cast<ReduceObject*>((*objects)[i]);
          // delete ReduceObject
          ReduceObject::destroy(reduce_object);
        }
        delete objects;
        //            std::cout << std::to_string(
        //                             TiledArray::get_default_world().rank()) +
        //                             " call 1\n";
      };
 
      world->taskq.add(destroy_vector, objects, TaskAttributes::hipri());
 
      const auto t1 = TiledArray::now();
      TiledArray::detail::cuda_callback_duration_ns<0>() +=
          TiledArray::duration_in_ns(t0, t1);
    }
 
    static void CUDART_CB cuda_dependency_dec_callback(void* userData) {
      const auto t0 = TiledArray::now();
 
      std::vector<void*>* objects = static_cast<std::vector<void*>*>(userData);
 
      for (auto& item : *objects) {
        // convert void* to DependencyInterface
        ReduceTaskImpl* dep = static_cast<ReduceTaskImpl*>(item);
        // call dec
        dep->dec();
      }
      delete objects;
      //          std::cout <<
      //          std::to_string(TiledArray::get_default_world().rank()) +
      //                           " call 2\n";
 
      const auto t1 = TiledArray::now();
      TiledArray::detail::cuda_callback_duration_ns<1>() +=
          TiledArray::duration_in_ns(t0, t1);
    }
 
    static void CUDART_CB
    cuda_dependency_dec_reduceobject_delete_callback(void* userData) {
      const auto t0 = TiledArray::now();
 
      std::vector<void*>* objects = static_cast<std::vector<void*>*>(userData);
 
      assert(objects->size() == 3);
 
      madness::World* world = static_cast<madness::World*>(objects->at(0));
 
      // convert void* to DependencyInterface
      ReduceTaskImpl* dep = static_cast<ReduceTaskImpl*>(objects->at(1));
      // call dec
      dep->dec();
 
      // convert void* to ReduceObject*
      ReduceObject* reduce_object = static_cast<ReduceObject*>(objects->at(2));
 
      auto destroy = [](ReduceObject* object) {
        ReduceObject::destroy(object);
        //            std::cout << std::to_string(
        //                             TiledArray::get_default_world().rank()) +
        //                             " call 3\n";
      };
 
      // delete ReduceObject
      world->taskq.add(destroy, reduce_object, TaskAttributes::hipri());
 
      delete objects;
 
      const auto t1 = TiledArray::now();
      TiledArray::detail::cuda_callback_duration_ns<2>() +=
          TiledArray::duration_in_ns(t0, t1);
    }
 
    static void CUDART_CB cuda_readyresult_reset_callback(void* userData) {
      const auto t0 = TiledArray::now();
 
      std::vector<void*>* objects = static_cast<std::vector<void*>*>(userData);
 
      madness::World* world = static_cast<madness::World*>((*objects)[0]);
 
      auto reset = [](std::vector<void*>* objects) {
        // skip the first one since it always be madness::World*
        // convert void* to the correct type
        std::shared_ptr<result_type>* result =
            static_cast<std::shared_ptr<result_type>*>((*objects)[1]);
        // call reset on shared_ptr
        result->reset();
        delete objects;
        //            std::cout <<
        //            std::to_string(TiledArray::get_default_world().rank()) +
        //                         " call 4\n";
      };
 
      world->taskq.add(reset, objects, TaskAttributes::hipri());
 
      const auto t1 = TiledArray::now();
      TiledArray::detail::cuda_callback_duration_ns<3>() +=
          TiledArray::duration_in_ns(t0, t1);
    }
 
#endif
    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
#ifdef TILEDARRAY_HAS_CUDA
          auto stream_ptr = tls_cudastream_accessor();
 
          if (stream_ptr == nullptr) {
            ReduceObject::destroy(ready_object);
            this->dec();
          } else {
            auto callback_object = new std::vector<void*>(3);
            (*callback_object)[0] = &world_;
            (*callback_object)[1] = this;
            (*callback_object)[2] = ready_object;
            CudaSafeCall(
                cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
            CudaSafeCall(cudaLaunchHostFunc(
                *stream_ptr, cuda_dependency_dec_reduceobject_delete_callback,
                callback_object));
            synchronize_stream(nullptr);
            //                std::cout << std::to_string(world().rank()) + "
            //                add 3\n";
          }
#else
          ReduceObject::destroy(ready_object);
          this->dec();
#endif
        } 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
#ifdef TILEDARRAY_HAS_CUDA
          auto stream_ptr = tls_cudastream_accessor();
          if (stream_ptr == nullptr) {
            ready_result.reset();
          } else {
            auto ready_result_heap =
                new std::shared_ptr<result_type>(ready_result);
            auto callback_object = new std::vector<void*>(2);
            (*callback_object)[0] = &world_;
            (*callback_object)[1] = ready_result_heap;
            CudaSafeCall(
                cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
            CudaSafeCall(cudaLaunchHostFunc(
                *stream_ptr, cuda_readyresult_reset_callback, callback_object));
            synchronize_stream(nullptr);
            //                std::cout << std::to_string(world().rank()) + "
            //                add 4\n";
          }
#else
          ready_result.reset();
#endif
        } 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
#ifdef TILEDARRAY_HAS_CUDA
      auto stream_ptr = tls_cudastream_accessor();
      if (stream_ptr == nullptr) {
        ReduceObject::destroy(object);
      } else {
        auto callback_object = new std::vector<void*>(2);
        (*callback_object)[0] = &world_;
        (*callback_object)[1] = const_cast<ReduceObject*>(object);
        CudaSafeCall(
            cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
        CudaSafeCall(cudaLaunchHostFunc(
            *stream_ptr, cuda_reduceobject_delete_callback, callback_object));
        synchronize_stream(nullptr);
        //            std::cout << std::to_string(world().rank()) + " add 1\n";
      }
#else
      ReduceObject::destroy(object);
#endif
      // 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.
#ifdef TILEDARRAY_HAS_CUDA
      if (stream_ptr == nullptr) {
        this->dec();
      } else {
        auto callback_object2 = new std::vector<void*>(1);
        (*callback_object2)[0] = this;
        CudaSafeCall(
            cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
        CudaSafeCall(cudaLaunchHostFunc(
            *stream_ptr, cuda_dependency_dec_callback, callback_object2));
        //            std::cout << std::to_string(world().rank()) + " add 2\n";
      }
#else
      this->dec();
#endif
    }
 
    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
#ifdef TILEDARRAY_HAS_CUDA
      auto stream_ptr = tls_cudastream_accessor();
      if (stream_ptr == nullptr) {
        ReduceObject::destroy(object1);
        ReduceObject::destroy(object2);
      } else {
        auto callback_object1 = new std::vector<void*>(3);
        (*callback_object1)[0] = &world_;
        (*callback_object1)[1] = const_cast<ReduceObject*>(object1);
        (*callback_object1)[2] = const_cast<ReduceObject*>(object2);
        CudaSafeCall(
            cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
        CudaSafeCall(cudaLaunchHostFunc(
            *stream_ptr, cuda_reduceobject_delete_callback, callback_object1));
        synchronize_stream(nullptr);
        //            std::cout << std::to_string(world().rank()) + " add 1\n";
      }
#else
      ReduceObject::destroy(object1);
      ReduceObject::destroy(object2);
#endif
 
      // 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.
#ifdef TILEDARRAY_HAS_CUDA
      if (stream_ptr == nullptr) {
        this->dec();
        this->dec();
      } else {
        auto callback_object2 = new std::vector<void*>(2);
        (*callback_object2)[0] = this;
        (*callback_object2)[1] = this;
        CudaSafeCall(
            cudaSetDevice(cudaEnv::instance()->current_cuda_device_id()));
        CudaSafeCall(cudaLaunchHostFunc(
            *stream_ptr, cuda_dependency_dec_callback, callback_object2));
        //            std::cout << std::to_string(world().rank()) + " add 2\n";
      }
 
#else
      this->dec();
      this->dec();
#endif
    }
 
#ifdef TILEDARRAY_HAS_CUDA
    template <typename Result = result_type>
    std::enable_if_t<detail::is_cuda_tile_v<Result>, void> internal_run(
        const madness::TaskThreadEnv&) {
      TA_ASSERT(ready_result_);
 
      auto post_result = madness::add_cuda_task(world_, op_, *ready_result_);
      result_.set(post_result);
 
      if (callback_) {
        result_.register_callback(callback_);
      }
    }
 
    template <typename Result = result_type>
    std::enable_if_t<!detail::is_cuda_tile_v<Result>, void>
#else
    void
#endif
    internal_run(const madness::TaskThreadEnv&) {
      TA_ASSERT(ready_result_);
      result_.set(op_(*ready_result_));
 
      if (callback_) callback_->notify();
    }
 
    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& threadEnv) {
      internal_run(threadEnv);
    }
 
 
    void ready(ReduceObject* object) {
      TA_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
        TA_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
        TA_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) {
    TA_ASSERT(pimpl_);
    pimpl_->inc();
    new typename ReduceTaskImpl::ReduceObject(pimpl_, arg, callback);
    return ++count_;
  }
 
 
  int count() const { return count_; }
 
 
  Future<result_type> submit() {
    TA_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