Add optimizer to IR tree #1

mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/dataset_node.cc

@@ -20,6 +20,7 @@
 #include <memory>
 #include <set>
 
+#include "minddata/dataset/engine/opt/pass.h"
 #include "minddata/dataset/util/random.h"
 
 namespace mindspore {
@@ -227,6 +228,7 @@ std::shared_ptr<DatasetNode> DatasetNode::SetNumWorkers(int32_t num_workers) {
   num_workers_ = num_workers;
   return shared_from_this();
 }
+
 DatasetNode::DatasetNode() {
   // Fetch some default value from config manager
   std::shared_ptr<ConfigManager> cfg = GlobalContext::config_manager();
@@ -236,5 +238,17 @@ DatasetNode::DatasetNode() {
   worker_connector_size_ = cfg->worker_connector_size();
 }
 
+// In DFS tree traversal, each node is visited twice. Accept is called on the first visit.
+Status DatasetNode::Accept(NodePass *p, bool *modified) {
+  // This method will only be called if its derived class does not implement one.
+  return p->Visit(shared_from_this(), modified);
+}
+
+// In DFS tree traversal, each node is visited twice. AcceptAfter is called on the second visit
+// after all child nodes are visited.
+Status DatasetNode::AcceptAfter(NodePass *p, bool *modified) {
+  // This method will only be called if its derived class does not implement one.
+  return p->VisitAfter(shared_from_this(), modified);
+}
 }  // namespace dataset
 }  // namespace mindspore

mindspore/ccsrc/minddata/dataset/engine/ir/datasetops/dataset_node.h

@@ -31,6 +31,7 @@ namespace dataset {
 
 class Dataset;
 class SamplerObj;
+class NodePass;
 
 #define RETURN_EMPTY_IF_ERROR(_s) \
   do {                            \
@@ -107,6 +108,24 @@ class DatasetNode : public std::enable_shared_from_this<DatasetNode> {
   /// \return Shared pointer to the original object
   std::shared_ptr<DatasetNode> SetNumWorkers(int32_t num_workers);
 
+  /// \brief Base method for NodePass visit. A tree walk consists of walking down the tree and also walking back up
+  ///     in a depth-first order. Accept is the node visit on the way down, whereas AcceptAfter is the node
+  ///     visit on the way back up the tree after its descendants are visited.
+  /// \notes Subclass needs to override this if it requires special node visit access.
+  ///     Check "dataset/engine/opt/pass.h" for more details.
+  /// \param[in] p The node to visit
+  /// \param[out] modified Indicator if the node was modified
+  /// \return Status of the node visit
+  virtual Status Accept(NodePass *p, bool *modified);
+
+  /// \brief Base method for NodePass visit on the way back up the tree after its descendants are visited.
+  /// \notes Subclass needs to override this if it requires special node visit access.
+  ///     Check "dataset/engine/opt/pass.h" for more details.
+  /// \param[in] p The node to visit
+  /// \param[out] modified Indicator if the node was modified
+  /// \return Status of the node visit
+  virtual Status AcceptAfter(NodePass *p, bool *modified);
+
  protected:
   std::vector<std::shared_ptr<DatasetNode>> children;
   std::shared_ptr<DatasetNode> parent;