pass registration

Author: panyx0718

paddle/fluid/framework/details/multi_devices_graph_builder.cc

@@ -34,30 +34,16 @@ namespace paddle {
 namespace framework {
 namespace details {
 
+void MultiDevSSAGraphBuilder::Init() const {
+  loss_var_name_ = Get<std::string>("loss_var_name");
+  places_ = Get<std::vector<platform::Place>>("places");
+  local_scopes_ = Get<std::vector<Scope *>>("local_scopes");
+  strategy_ = Get<BuildStrategy>("strategy");
 #ifdef PADDLE_WITH_CUDA
-MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
-    const std::vector<platform::Place> &places,
-    const std::string &loss_var_name,
-    const std::unordered_set<std::string> &params,
-    const std::vector<Scope *> &local_scopes,
-    platform::NCCLContextMap *nccl_ctxs, const BuildStrategy &strategy)
-    : loss_var_name_(loss_var_name),
-      places_(places),
-      local_scopes_(local_scopes),
-      nccl_ctxs_(nccl_ctxs),
-      strategy_(strategy) {
-#else
-MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
-    const std::vector<platform::Place> &places,
-    const std::string &loss_var_name,
-    const std::unordered_set<std::string> &params,
-    const std::vector<Scope *> &local_scopes, const BuildStrategy &strategy)
-    : loss_var_name_(loss_var_name),
-      places_(places),
-      local_scopes_(local_scopes),
-      strategy_(strategy) {
+  nccl_ctxs_ = &Get<platform::NCCLContextMap>("nccl_ctxs");
 #endif
-  for (auto &p : params) {
+
+  for (auto &p : Get<std::unordered_set<std::string>>("params")) {
     grad_names_.insert(GradVarName(p));
   }
   balance_vars_.resize(places_.size(), 0);
@@ -241,6 +227,7 @@ std::vector<ir::Node *> SortOpsAndDelayOptimizeOp(const ir::Graph &graph) {
 
 std::unique_ptr<ir::Graph> MultiDevSSAGraphBuilder::Apply(
     std::unique_ptr<ir::Graph> graph) const {
+  Init();
   // Give the topology sort order and rebuild the graph structure.
   std::vector<ir::Node *> sorted_ops = SortOpsAndDelayOptimizeOp(*graph);
   auto nodes = graph->ReleaseNodes();

paddle/fluid/framework/details/multi_devices_graph_builder.h

@@ -32,36 +32,23 @@ namespace details {
 
 class MultiDevSSAGraphBuilder : public SSAGraphBuilder {
  public:
-#ifdef PADDLE_WITH_CUDA
-  MultiDevSSAGraphBuilder(const std::vector<platform::Place> &places,
-                          const std::string &loss_var_name,
-                          const std::unordered_set<std::string> &params,
-                          const std::vector<Scope *> &local_scopes,
-                          platform::NCCLContextMap *nccl_ctxs,
-                          const BuildStrategy &strategy);
-#else
-  MultiDevSSAGraphBuilder(const std::vector<platform::Place> &places,
-                          const std::string &loss_var_name,
-                          const std::unordered_set<std::string> &params,
-                          const std::vector<Scope *> &local_scopes,
-                          const BuildStrategy &strategy);
-#endif
   std::unique_ptr<ir::Graph> Apply(
       std::unique_ptr<ir::Graph> graph) const override;
   int GetVarDeviceID(const std::string &varname) const override;
 
  private:
   void CreateOpHandleIOs(ir::Graph *result, ir::Node *node,
                          size_t device_id) const;
+  void Init() const;
 
  private:
-  std::string loss_var_name_;
-  const std::vector<platform::Place> &places_;
-  const std::vector<Scope *> &local_scopes_;
-  std::unordered_set<std::string> grad_names_;
+  mutable std::string loss_var_name_;
+  mutable std::vector<platform::Place> places_;
+  mutable std::vector<Scope *> local_scopes_;
+  mutable std::unordered_set<std::string> grad_names_;
 
 #ifdef PADDLE_WITH_CUDA
-  platform::NCCLContextMap *nccl_ctxs_;
+  mutable platform::NCCLContextMap *nccl_ctxs_;
 #endif
 
   bool IsScaleLossOp(ir::Node *node) const;
@@ -113,7 +100,7 @@ class MultiDevSSAGraphBuilder : public SSAGraphBuilder {
       const std::vector<std::string> &var_names) const;
 
  private:
-  BuildStrategy strategy_;
+  mutable BuildStrategy strategy_;
   mutable std::unordered_map<std::string, VarDesc *> all_vars_;
   mutable std::unordered_map<std::string, int> var_name_on_devices_;
   mutable std::vector<int64_t> balance_vars_;

paddle/fluid/framework/details/ssa_graph_builder_factory.cc

@@ -22,26 +22,29 @@ namespace paddle {
 namespace framework {
 namespace details {
 std::unique_ptr<SSAGraphBuilder> SSAGraphBuilderFactory::Create() {
-  std::unique_ptr<SSAGraphBuilder> res(
+  std::unique_ptr<SSAGraphBuilder> res(new MultiDevSSAGraphBuilder);
+  res->SetNotOwned<std::vector<platform::Place>>("places", &places_);
+  res->SetNotOwned<std::string>("loss_var_name", &loss_var_name_);
+  res->SetNotOwned<std::unordered_set<std::string>>("params", &param_names_);
+  res->SetNotOwned<std::vector<Scope *>>("local_scopes", &local_scopes_);
+  res->SetNotOwned<BuildStrategy>("strategy", &strategy_);
 #ifdef PADDLE_WITH_CUDA
-      new MultiDevSSAGraphBuilder(places_, loss_var_name_, param_names_,
-                                  local_scopes_, nccl_ctxs_, strategy_)
-#else
-      new MultiDevSSAGraphBuilder(places_, loss_var_name_, param_names_,
-                                  local_scopes_, strategy_)
+  res->SetNotOwned<platform::NCCLContextMap>("nccl_ctxs", nccl_ctxs_);
 #endif
-          );  // NOLINT
 
   if (!strategy_.debug_graphviz_path_.empty()) {
-    std::unique_ptr<std::ostream> fout(
-        new std::ofstream(strategy_.debug_graphviz_path_));
-    PADDLE_ENFORCE(fout->good());
-    std::unique_ptr<GraphvizSSAGraphPrinter> graphviz_printer(
-        new GraphvizSSAGraphPrinter());
-    res.reset(new SSAGraghBuilderWithPrinter(
-        std::move(fout), std::move(graphviz_printer), std::move(res)));
+    SSAGraphBuilder *previous_pass = res.release();
+    res.reset(new SSAGraghBuilderWithPrinter);
+    res->Set<SSAGraphBuilder>("previous_pass", previous_pass);
+    res->SetNotOwned<std::string>("debug_graphviz_path",
+                                  &strategy_.debug_graphviz_path_);
+    res->Set<GraphvizSSAGraphPrinter>("graph_printer",
+                                      new GraphvizSSAGraphPrinter);
   }
-  res.reset(new SSAGraghBuilderWithChecker(std::move(res)));
+
+  SSAGraphBuilder *previous_pass = res.release();
+  res.reset(new SSAGraghBuilderWithChecker);
+  res->Set<SSAGraphBuilder>("previous_pass", previous_pass);
 
   return res;
 }

paddle/fluid/framework/details/ssa_graph_checker.h

@@ -24,25 +24,19 @@ namespace details {
 
 class SSAGraghBuilderWithChecker : public SSAGraphBuilder {
  public:
-  explicit SSAGraghBuilderWithChecker(
-      std::unique_ptr<SSAGraphBuilder>&& builder)
-      : builder_(std::move(builder)) {}
-
   std::unique_ptr<ir::Graph> Apply(
       std::unique_ptr<ir::Graph> graph) const override {
-    auto new_graph = builder_->Apply(std::move(graph));
+    auto new_graph =
+        Get<SSAGraphBuilder>("previous_pass").Apply(std::move(graph));
     PADDLE_ENFORCE(IsValidGraph(new_graph.get()));
     return new_graph;
   }
 
   int GetVarDeviceID(const std::string& var_name) const override {
-    return builder_->GetVarDeviceID(var_name);
+    return Get<SSAGraphBuilder>("previous_pass").GetVarDeviceID(var_name);
   }
 
   bool IsValidGraph(const ir::Graph* graph) const;
-
- private:
-  std::unique_ptr<SSAGraphBuilder> builder_;
 };
 
 }  // namespace details

paddle/fluid/framework/details/ssa_graph_printer.h

@@ -14,7 +14,9 @@
 
 #pragma once
 
+#include <fstream>
 #include <iosfwd>
+#include <ostream>
 #include <string>
 #include "paddle/fluid/framework/details/ssa_graph_builder.h"
 
@@ -35,37 +37,21 @@ class GraphvizSSAGraphPrinter : public SSAGraphPrinter {
 
 class SSAGraghBuilderWithPrinter : public SSAGraphBuilder {
  public:
-  SSAGraghBuilderWithPrinter(std::ostream& sout,
-                             std::unique_ptr<SSAGraphPrinter>&& printer,
-                             std::unique_ptr<SSAGraphBuilder>&& builder)
-      : printer_(std::move(printer)),
-        builder_(std::move(builder)),
-        stream_ref_(sout) {}
-
-  SSAGraghBuilderWithPrinter(std::unique_ptr<std::ostream>&& sout,
-                             std::unique_ptr<SSAGraphPrinter>&& printer,
-                             std::unique_ptr<SSAGraphBuilder>&& builder)
-      : printer_(std::move(printer)),
-        builder_(std::move(builder)),
-        stream_ptr_(std::move(sout)),
-        stream_ref_(*stream_ptr_) {}
-
   std::unique_ptr<ir::Graph> Apply(
       std::unique_ptr<ir::Graph> graph) const override {
-    auto new_graph = builder_->Apply(std::move(graph));
-    printer_->Print(*new_graph, stream_ref_);
+    auto new_graph =
+        Get<SSAGraphBuilder>("previous_pass").Apply(std::move(graph));
+
+    std::unique_ptr<std::ostream> fout(
+        new std::ofstream(Get<std::string>("debug_graphviz_path")));
+    PADDLE_ENFORCE(fout->good());
+    Get<GraphvizSSAGraphPrinter>("graph_printer").Print(*new_graph, *fout);
     return new_graph;
   }
 
   int GetVarDeviceID(const std::string& var_name) const override {
-    return builder_->GetVarDeviceID(var_name);
+    return Get<SSAGraphBuilder>("previous_pass").GetVarDeviceID(var_name);
   }
-
- private:
-  std::unique_ptr<SSAGraphPrinter> printer_;
-  std::unique_ptr<SSAGraphBuilder> builder_;
-  std::unique_ptr<std::ostream> stream_ptr_;
-  std::ostream& stream_ref_;
 };
 
 }  // namespace details

paddle/fluid/framework/ir/graph_viz_pass.cc

@@ -23,7 +23,8 @@ namespace ir {
 
 std::unique_ptr<ir::Graph> GraphVizPass::Apply(
     std::unique_ptr<ir::Graph> graph) const {
-  std::unique_ptr<std::ostream> fout(new std::ofstream(graph_viz_path_));
+  const std::string graph_viz_path = Get<std::string>("graph_viz_path");
+  std::unique_ptr<std::ostream> fout(new std::ofstream(graph_viz_path));
   PADDLE_ENFORCE(fout->good());
   std::ostream& sout = *fout;
 
@@ -61,6 +62,9 @@ std::unique_ptr<ir::Graph> GraphVizPass::Apply(
   sout << "}\n";
   return graph;
 }
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle
+
+REGISTER_PASS(graph_viz_pass, paddle::framework::ir::GraphVizPass);

paddle/fluid/framework/ir/graph_viz_pass.h

@@ -29,14 +29,8 @@ namespace ir {
 
 class GraphVizPass : public Pass {
  public:
-  explicit GraphVizPass(const std::string& graph_viz_path)
-      : graph_viz_path_(graph_viz_path) {}
-
   std::unique_ptr<ir::Graph> Apply(
       std::unique_ptr<ir::Graph> graph) const override;
-
- private:
-  const std::string graph_viz_path_;
 };
 
 }  // namespace ir

paddle/fluid/framework/ir/pass.cc

@@ -15,5 +15,12 @@ limitations under the License. */
 #include "paddle/fluid/framework/ir/pass.h"
 
 namespace paddle {
-namespace framework {}  // namespace framework
+namespace framework {
+namespace ir {
+PassRegistry& PassRegistry::Instance() {
+  static PassRegistry g_pass_info_map;
+  return g_pass_info_map;
+}
+}  // namespace ir
+}  // namespace framework
 }  // namespace paddle