schedule.h

Go to the documentation of this file.

#ifndef __SCHEDULE_H__
#define __SCHEDULE_H__
#include "common.h"
#include <queue>
#include "idx_tensor.h"

namespace CTF {

  enum TensorOperationTypes {
    TENSOR_OP_NONE,
    TENSOR_OP_SET,
    TENSOR_OP_SUM,
    TENSOR_OP_SUBTRACT,
    TENSOR_OP_MULTIPLY };

  class TensorOperationBase {
  public:
    virtual ~TensorOperationBase() {}
  };

  class TensorOperation : public TensorOperationBase {
  public:
    TensorOperation(TensorOperationTypes op,
        Idx_Tensor* lhs,
        const CTF_int::Term* rhs) :
          dependency_count(0),
          op(op),
          lhs(lhs),
          rhs(rhs),
          cached_estimated_cost(0) {}

    void get_outputs(std::set<Idx_Tensor*, CTF_int::tensor_name_less >* outputs_set) const;

    void get_inputs(std::set<Idx_Tensor*, CTF_int::tensor_name_less >* inputs_set) const;

    void execute(std::map<CTF_int::tensor*, CTF_int::tensor*>* remap = NULL);

    double  estimate_time();

    bool is_dummy() {
      return op == TENSOR_OP_NONE;
    }

    // Number of dependencies I have
    int dependency_count;
    // List of all successors - operations that depend on me
    std::vector<TensorOperation* > successors;
    std::vector<TensorOperation* > reads;

    int dependency_left;

    const char* name() {
      return lhs->parent->name;
    }

  protected:
    TensorOperationTypes op;
    Idx_Tensor* lhs;
    const CTF_int::Term* rhs;

    double  cached_estimated_cost;
  };

  // untemplatized scheduler abstract base class to assist in global operations
  class ScheduleBase {
  public:
    virtual void add_operation(TensorOperationBase* op) = 0;
  };

  extern ScheduleBase* global_schedule;

  struct ScheduleTimer {
    double comm_down_time;
    double exec_time;
    double imbalance_wall_time;
    double imbalance_acuum_time;
    double comm_up_time;
    double total_time;

    ScheduleTimer():
      comm_down_time(0),
      exec_time(0),
      imbalance_wall_time(0),
      imbalance_acuum_time(0),
      comm_up_time(0),
      total_time(0) {}

    void operator+=(ScheduleTimer const & B) {
      comm_down_time += B.comm_down_time;
      exec_time += B.exec_time;
      imbalance_wall_time += B.imbalance_wall_time;
      imbalance_acuum_time += B.imbalance_acuum_time;
      comm_up_time += B.comm_up_time;
      total_time += B.total_time;
    }
  };

  
  class Schedule : public ScheduleBase {
  public:
    Schedule(World* world = NULL) :
      world(world),
      partitions(0) {}

    void record();

    ScheduleTimer execute();

    inline ScheduleTimer partition_and_execute();

    inline void schedule_op_successors(TensorOperation* op);

    void add_operation_typed(TensorOperation* op);
    void add_operation(TensorOperationBase* op);

    void set_max_partitions(int in_partitions) {
      partitions = in_partitions;
    }

  protected:
    World* world;

    // Tasks with no dependencies, which can be executed at the start
    std::deque<TensorOperation*> root_tasks;

    // For debugging purposes - the steps in the original input order
    std::deque<TensorOperation*> steps_original;

    // Last operation writing to the key tensor
    std::map<CTF_int::tensor*, TensorOperation*> latest_write;

    // Ready queue of tasks with all dependencies satisfied
    std::deque<TensorOperation*> ready_tasks;

    int partitions;

  };

}
#endif