mindspore2022/mindspore/ccsrc/dataset/engine/execution_tree.h

/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef DATASET_ENGINE_EXECUTION_TREE_H_
#define DATASET_ENGINE_EXECUTION_TREE_H_

#include <functional>
#include <memory>
#include <stack>
#include <vector>
#include "dataset/engine/datasetops/dataset_op.h"
#include "dataset/util/status.h"

namespace mindspore {
namespace dataset {
// Forward declares
class TaskGroup;
class DatasetOp;

class ExecutionTree {
 public:
  // Prepare flags used during tree prepare phase
  enum PrepareFlags {
    kDePrepNone = 0,
    kDePrepRepeat = 1  //  Processing a repeat operation
  };

  // State flags for the lifecycle of the tree
  enum TreeState {
    kDeTStateInit = 0,  // The freshly initialized state after construction
    kDeTStateBuilding,  // The tree is being built, nodes are being added
    kDeTStatePrepare,   // The tree has been assigned a root node and is pending prepare
    kDeTStateReady,     // The tree has been prepared and is ready to be launched
    kDeTStateExecuting  // The tree has been launched and is executing
  };

  class Iterator {
   public:
    // Constructor
    // @param root The root node to start iterating from
    explicit Iterator(const std::shared_ptr<DatasetOp> &root = nullptr);

    // Destructor
    ~Iterator() {}

    Iterator &operator++() {
      ++ind_;
      return *this;
    }  // prefix ++ overload
    Iterator operator++(int) {
      Iterator it = *this;
      it.ind_ = ind_;
      ind_++;
      return it;
    }  // post-fix ++ overload
    Iterator &operator--() {
      --ind_;
      return *this;
    }  // prefix -- overload
    Iterator operator--(int) {
      Iterator it = *this;
      it.ind_ = ind_;
      ind_--;
      return it;
    }                                                 // post-fix -- overload
    DatasetOp &operator*() { return *nodes_[ind_]; }  // dereference operator
    std::shared_ptr<DatasetOp> operator->() { return nodes_[ind_]; }

    // getter function
    // @return Shared pointer to the current operator
    std::shared_ptr<DatasetOp> get() { return nodes_[ind_]; }

    bool operator!=(const Iterator &rhs) { return nodes_[ind_] != rhs.nodes_[rhs.ind_]; }

   private:
    int ind_;                                        // the cur node our Iterator points to
    std::vector<std::shared_ptr<DatasetOp>> nodes_;  // store the nodes in post order
    void PostOrderTraverse(const std::shared_ptr<DatasetOp> &);
  };

  // Constructor
  ExecutionTree();

  // Destructor
  ~ExecutionTree();

  // Associates a DatasetOp with this tree. This assigns a valid node id to the operator and
  // provides it with a link to the tree. A node cannot form any relationships (parent/child) with
  // other nodes unless they are associated with the same tree.
  // @param op - The operator to associate
  // @return Status - The error code return
  Status AssociateNode(const std::shared_ptr<DatasetOp> &op);

  // Sets the root node of the tree
  // @param op - The operator to assign as root
  // @return Status - The error code return
  Status AssignRoot(const std::shared_ptr<DatasetOp> &op);

  // Start the execution of the tree
  // @return Status - The error code return
  Status Launch();

  // A print method typically used for debugging
  // @param out - The output stream to write output to
  // @param show_all - A bool to control if you want to show all info or just a summary
  void Print(std::ostream &out, bool show_all) const;

  // Returns an iterator positioned at the start
  // @return Iterator - The iterator
  ExecutionTree::Iterator begin(const std::shared_ptr<DatasetOp> &root = nullptr) const {
    return Iterator((root == nullptr) ? root_ : root);
  }

  // Returns an iterator positioned at the end
  // @return Iterator - The iterator
  ExecutionTree::Iterator end() const { return Iterator(nullptr); }

  // << Stream output operator overload
  // @notes This allows you to write the debug print info using stream operators
  // @param out - reference to the output stream being overloaded
  // @param exe_tree - reference to the execution tree to display
  // @return - the output stream must be returned
  friend std::ostream &operator<<(std::ostream &out, ExecutionTree &exe_tree) {
    exe_tree.Print(out, false);
    return out;
  }

  // Given the number of workers, launches the worker entry function for each. Essentially a
  // wrapper for the TaskGroup handling that is stored inside the execution tree.
  // @param num_workers - The number of workers to launch
  // @param func - The function entry point that workers will execute
  // @return Status - The error code return
  Status LaunchWorkers(int32_t num_workers, std::function<Status(uint32_t)> func);

  // Getter method
  // @return shared_ptr to the root operator
  std::shared_ptr<DatasetOp> root() const { return root_; }

  // Getter method
  // @return the prepare flags
  uint32_t PrepareFlags() const { return prepare_flags_; }

  // The driver of the prepare phase of the execution tree. The prepare phase will recursively
  // walk the tree to perform modifications to the tree or specific nodes within the tree to get
  // it ready for execution.
  // @return Status - The error code return
  Status Prepare();

  // Recursive function used during prepare phase to visit a node and drive any pre- and post-
  // node actions during a tree walk.
  // @param op - The dataset op to work on
  // @return Status - The error code return
  Status PrepareNode(const std::shared_ptr<DatasetOp> &dataset_op);

  // Adds an operator to the repeat stack during prepare phase.
  // @param op - The dataset op to work add to repeat stack
  // @return Status - The error code return
  void AddToRepeatStack(std::shared_ptr<DatasetOp> dataset_op);

  // Pops an operator from the repeat stack during prepare phase.
  // @return shared_ptr to the popped operator
  std::shared_ptr<DatasetOp> PopFromRepeatStack();

  // Return the pointer to the TaskGroup
  // @return raw pointer to the TaskGroup
  TaskGroup *AllTasks() const { return tg_.get(); }

 private:
  std::unique_ptr<TaskGroup> tg_;                        // Class for worker management
  std::shared_ptr<DatasetOp> root_;                      // The root node of the tree
  int32_t id_count_;                                     // Counter for generating operator id's
  uint32_t prepare_flags_;                               // Flags used during tree prepare
  TreeState tree_state_;                                 // Tracking the current tree state
  std::stack<std::shared_ptr<DatasetOp>> repeat_stack_;  // A stack used during prepare phase
};
}  // namespace dataset
}  // namespace mindspore

#endif  // DATASET_ENGINE_EXECUTION_TREE_H_