update

cmake/external/xpu.cmake

@@ -35,7 +35,7 @@ ELSE ()
 ENDIF()
 
 SET(XPU_BASE_URL_WITHOUT_DATE "https://baidu-kunlun-product.cdn.bcebos.com/KL-SDK/klsdk-dev")
-SET(XPU_BASE_URL "${XPU_BASE_URL_WITHOUT_DATE}/20210701")
+SET(XPU_BASE_URL "${XPU_BASE_URL_WITHOUT_DATE}/20210729")
 SET(XPU_XRE_URL  "${XPU_BASE_URL}/${XPU_XRE_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)
 SET(XPU_XDNN_URL "${XPU_BASE_URL}/${XPU_XDNN_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)
 SET(XPU_XCCL_URL "${XPU_BASE_URL_WITHOUT_DATE}/20210623/${XPU_XCCL_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)

paddle/fluid/framework/CMakeLists.txt

@@ -188,8 +188,13 @@ cc_library(op_kernel_type SRCS op_kernel_type.cc DEPS device_context place)
 
 cc_library(unused_var_check SRCS unused_var_check.cc DEPS glog no_need_buffer_vars_inference)
 
+IF(WITH_XPU)
+cc_library(operator SRCS operator.cc DEPS xpu_op_list op_info device_context tensor scope glog trainer_desc_proto data_feed_proto
+    shape_inference data_transform lod_tensor profiler transfer_scope_cache op_kernel_type op_call_stack unused_var_check nan_inf_utils)
+ELSE()
 cc_library(operator SRCS operator.cc DEPS op_info device_context tensor scope glog trainer_desc_proto data_feed_proto
     shape_inference data_transform lod_tensor profiler transfer_scope_cache op_kernel_type op_call_stack unused_var_check nan_inf_utils)
+ENDIF()
 
 cc_test(operator_test SRCS operator_test.cc DEPS operator op_registry device_context)
 cc_test(operator_exception_test SRCS operator_exception_test.cc DEPS operator op_registry device_context)
@@ -405,7 +410,7 @@ configure_file(commit.h.in commit.h)
 # Adapt to custom op mechanism: Include the header files related to the data type
 # to avoid exposing the path of the underlying file
 include_directories(${PADDLE_SOURCE_DIR}/paddle/fluid/platform)
-include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../extension/include)
+include_directories(${PADDLE_SOURCE_DIR}/paddle/fluid/extension/include)
 
 if(WITH_ROCM)
   hip_library(custom_tensor SRCS ../extension/src/ext_tensor.cc DEPS lod_tensor memory enforce)

paddle/fluid/framework/details/multi_devices_helper.h

@@ -77,10 +77,6 @@ typedef std::vector<std::pair<std::string, std::string>> ParamsAndGrads;
 constexpr char kParamsAndDenseGrads[] = "params_and_dense_grads";
 constexpr char kParamsAndSparseGrads[] = "params_and_sparse_grads";
 
-typedef std::vector<ProgramDesc> ProgramDescs;
-constexpr char kProgramDescs[] = "program_descs";
-constexpr char kStartupProgramDescs[] = "startup_program_descs";
-
 typedef std::unordered_set<std::string> PinnedVars;
 constexpr char kPinnedVars[] = "pinned_vars";
 

paddle/fluid/framework/ir/graph_helper.cc

@@ -15,7 +15,9 @@ limitations under the License. */
 #include "paddle/fluid/framework/ir/graph_helper.h"
 #include <queue>
 #include <stack>
+#include "paddle/fluid/framework/op_proto_maker.h"
 
+DECLARE_bool(convert_all_blocks);
 DEFINE_string(print_sub_graph_dir, "",
               "FLAGS_print_sub_graph_dir is used "
               "to print the nodes of sub_graphs.");
@@ -431,6 +433,117 @@ std::vector<ir::Node *> TopologySortGraphByDescOrder(const Graph &graph) {
   return ret;
 }
 
+static OpDesc *ReplaceScaleLossGradOp(const Node &node, OpDesc *desc) {
+  desc->SetType("fill_constant");
+  desc->SetAttr(
+      OpProtoAndCheckerMaker::OpRoleAttrName(),
+      (static_cast<int>(OpRole::kBackward) | static_cast<int>(OpRole::kLoss)));
+  desc->SetAttr("value", 1.0f);
+  std::vector<std::string> output_names;
+  for (auto out : node.outputs) {
+    output_names.emplace_back(out->Name());
+  }
+  desc->SetOutput("Out", output_names);
+  return desc;
+}
+
+static void GetGraphOpDesc(const std::vector<Node *> &nodes,
+                           std::vector<OpDesc> *ops) {
+  for (Node *n : nodes) {
+    // if node is not Op, skip
+    if (!n->IsOp()) continue;
+
+    // create fill_constant op
+    if (n->Name() == "scale_loss_grad") {
+      ops->emplace_back();
+      auto &desc = ops->back();
+      ReplaceScaleLossGradOp(*n, &desc);
+    } else if (n->Op()) {
+      ops->emplace_back(*n->Op());
+    }
+    // delete no OpDesc op
+  }
+}
+
+static void GraphToBlock(const Graph &graph, proto::BlockDesc *block,
+                         const SortKind *sort_kind) {
+  // Remove the unneeded variables after memory optimization.
+  std::unordered_set<std::string> vars2remove;
+  if (graph.Has(kGraphToProgramVarsToRemove)) {
+    vars2remove =
+        graph.Get<std::unordered_set<std::string>>(kGraphToProgramVarsToRemove);
+    VLOG(2) << "graph (id: " << block->idx() << ") to program remove "
+            << vars2remove.size() << " nodes";
+  }
+
+  block->clear_vars();
+  std::unordered_set<std::string> visited_vars;
+  for (Node *n : graph.Nodes()) {
+    if (n->IsVar()) {
+      if (n->Var() && visited_vars.count(n->Var()->Name()) == 0 &&
+          !vars2remove.count(n->Var()->Name()) &&
+          n->GetVarNodeBlockId() == graph.GetBlockId()) {
+        visited_vars.insert(n->Var()->Name());
+        block->add_vars()->MergeFrom(*n->Var()->Proto());
+      }
+    }
+  }
+  block->clear_ops();
+
+  std::vector<Node *> nodes;
+  if (sort_kind != nullptr) {
+    // Inference Memory Optimize relays on this branch.
+    nodes = TopologyVarientSort(graph, *sort_kind);
+  } else {
+    if (FLAGS_convert_all_blocks) {
+      nodes = TopologySortGraphByDescOrder(graph);
+    } else {
+      nodes = TopologySortOperations(graph);
+    }
+  }
+
+  std::vector<OpDesc> ops;
+  GetGraphOpDesc(nodes, &ops);
+  for (auto &op : ops) {
+    block->add_ops()->MergeFrom(*op.Proto());
+  }
+}
+
+void GraphToProgram(const Graph &graph, ProgramDesc *program,
+                    const SortKind *sort_kind) {
+  PADDLE_ENFORCE_EQ(graph.IsMainGraph(), true,
+                    platform::errors::InvalidArgument(
+                        "This graph is a sub_graph, "
+                        "and can't convert to program individually"));
+  PADDLE_ENFORCE_NOT_NULL(
+      program,
+      platform::errors::InvalidArgument(
+          "program must not be nullptr when converting graph to program"));
+
+  proto::ProgramDesc program_pb(*(program->Proto()));
+  auto block = program_pb.mutable_blocks(kRootBlockIndex);
+  block->set_idx(kRootBlockIndex);
+
+  if (FLAGS_convert_all_blocks) {
+    GraphToBlock(*graph.GetSubGraph(kRootBlockIndex), block, sort_kind);
+
+    VLOG(3) << "Graph to program need convert " << graph.SubGraphsSize()
+            << " sub graph";
+    for (size_t idx = 0; idx < graph.SubGraphsSize(); ++idx) {
+      // avoid kRootBlockIndex not 0
+      if (idx == kRootBlockIndex) continue;
+
+      block = program_pb.add_blocks();
+      block->set_idx(idx);
+      GraphToBlock(*graph.GetSubGraph(idx), block, sort_kind);
+    }
+  } else {
+    GraphToBlock(graph, block, sort_kind);
+  }
+
+  program->CopyFrom(program_pb);
+}
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/graph_helper.h

@@ -27,6 +27,10 @@ namespace paddle {
 namespace framework {
 namespace ir {
 
+constexpr char kGraphToProgramVarsToRemove[] =
+    "__graph_to_program_vars_to_remove__";
+constexpr char kGraphToProgramSortKind[] = "__graph_to_program_sort_kind__";
+
 // Compare nodes via node id.
 class Graph;
 
@@ -117,6 +121,9 @@ std::vector<T *> FilterByNodeWrapper(const Graph &graph) {
 
 std::vector<ir::Node *> TopologySortGraphByDescOrder(const Graph &graph);
 
+void GraphToProgram(const Graph &graph, ProgramDesc *p_program,
+                    const SortKind *sort_kind = nullptr);
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/graph_to_program_pass.cc

@@ -17,11 +17,8 @@ limitations under the License. */
 #include <gflags/gflags.h>
 #include <algorithm>
 
-#include "paddle/fluid/framework/ir/graph_helper.h"
 #include "paddle/fluid/framework/op_proto_maker.h"
 
-DECLARE_bool(convert_all_blocks);
-
 namespace paddle {
 namespace framework {
 class ProgramDesc;
@@ -33,116 +30,12 @@ namespace framework {
 namespace ir {
 
 void GraphToProgramPass::ApplyImpl(ir::Graph* graph) const {
-  PADDLE_ENFORCE_EQ(graph->IsMainGraph(), true,
-                    platform::errors::InvalidArgument(
-                        "This graph is a sub_graph, "
-                        "and can't convert to program individually"));
-
-  ProgramDesc& program = Get<ProgramDesc>("program");
-
-  std::unique_ptr<proto::ProgramDesc> program_pb(
-      new proto::ProgramDesc(*program.Proto()));
-
-  auto block = program_pb->mutable_blocks(kRootBlockIndex);
-  block->set_idx(kRootBlockIndex);
-
-  if (FLAGS_convert_all_blocks) {
-    GraphToBlock(graph->GetSubGraph(kRootBlockIndex), block);
-
-    VLOG(3) << "Graph to program need convert " << graph->SubGraphsSize()
-            << " sub graph";
-    for (size_t idx = 0; idx < graph->SubGraphsSize(); ++idx) {
-      // avoid kRootBlockIndex not 0
-      if (idx == kRootBlockIndex) continue;
-
-      block = program_pb->add_blocks();
-      block->set_idx(idx);
-      GraphToBlock(graph->GetSubGraph(idx), block);
-    }
-  } else {
-    GraphToBlock(graph, block);
-  }
-
-  program.CopyFrom(*program_pb);
-}
-
-OpDesc* ReplaceScaleLossGradOp(ir::Node* node, OpDesc* desc) {
-  desc->SetType("fill_constant");
-  desc->SetAttr(
-      OpProtoAndCheckerMaker::OpRoleAttrName(),
-      (static_cast<int>(OpRole::kBackward) | static_cast<int>(OpRole::kLoss)));
-  desc->SetAttr("value", 1.0f);
-  std::vector<std::string> output_names;
-  for (auto out : node->outputs) {
-    output_names.emplace_back(out->Name());
-  }
-  desc->SetOutput("Out", output_names);
-  return desc;
-}
-
-std::vector<OpDesc>* GetGraphOpDesc(const std::vector<ir::Node*>& nodes,
-                                    std::vector<OpDesc>* ops) {
-  for (ir::Node* n : nodes) {
-    // if node is not Op, skip
-    if (!n->IsOp()) continue;
-
-    // create fill_constant op
-    if (n->Name() == "scale_loss_grad") {
-      ops->emplace_back();
-      auto& desc = ops->back();
-      ReplaceScaleLossGradOp(n, &desc);
-    } else if (n->Op()) {
-      ops->emplace_back(*n->Op());
-    } else {
-      // delete no OpDesc op
-    }
-  }
-  return ops;
-}
-
-void GraphToProgramPass::GraphToBlock(const Graph* graph,
-                                      proto::BlockDesc* block) const {
-  // Remove the unneeded variables after memory optimization.
-  std::unordered_set<std::string> vars2remove;
-  if (graph->Has(kGraphToProgramVarsToRemove)) {
-    vars2remove = graph->Get<std::unordered_set<std::string>>(
-        kGraphToProgramVarsToRemove);
-    VLOG(2) << "graph (id: " << block->idx() << ") to program remove "
-            << vars2remove.size() << " nodes";
-  }
-
-  block->clear_vars();
-  std::unordered_set<std::string> visited_vars;
-  for (ir::Node* n : graph->Nodes()) {
-    if (n->IsVar()) {
-      if (n->Var() && visited_vars.count(n->Var()->Name()) == 0 &&
-          !vars2remove.count(n->Var()->Name()) &&
-          n->GetVarNodeBlockId() == graph->GetBlockId()) {
-        visited_vars.insert(n->Var()->Name());
-        block->add_vars()->MergeFrom(*n->Var()->Proto());
-      }
-    }
-  }
-  block->clear_ops();
-
-  std::vector<ir::Node*> nodes;
+  auto& program = Get<ProgramDesc>("program");
   if (Has(kGraphToProgramSortKind)) {
-    // Inference Memory Optimize relays on this branch.
-    int sort_kind = Get<int>(kGraphToProgramSortKind);
-    nodes = TopologyVarientSort(
-        *graph, static_cast<framework::ir::SortKind>(sort_kind));
+    auto sort_kind = static_cast<SortKind>(Get<int>(kGraphToProgramSortKind));
+    GraphToProgram(*graph, &program, &sort_kind);
   } else {
-    if (FLAGS_convert_all_blocks) {
-      nodes = TopologySortGraphByDescOrder(*graph);
-    } else {
-      nodes = TopologySortOperations(*graph);
-    }
-  }
-
-  std::vector<OpDesc> ops;
-  GetGraphOpDesc(nodes, &ops);
-  for (auto& op : ops) {
-    block->add_ops()->MergeFrom(*op.Proto());
+    GraphToProgram(*graph, &program, nullptr);
   }
 }
 

paddle/fluid/framework/ir/graph_to_program_pass.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include "paddle/fluid/framework/ir/graph_helper.h"
 #include "paddle/fluid/framework/ir/pass.h"
 
 namespace paddle {
@@ -22,16 +23,9 @@ namespace ir {
 
 class Graph;
 
-const char kGraphToProgramVarsToRemove[] =
-    "__graph_to_program_vars_to_remove__";
-const char kGraphToProgramSortKind[] = "__graph_to_program_sort_kind__";
-
 class GraphToProgramPass : public Pass {
  protected:
   void ApplyImpl(ir::Graph* graph) const override;
-
- private:
-  void GraphToBlock(const Graph* graph, proto::BlockDesc* block) const;
 };
 
 }  // namespace ir

paddle/fluid/framework/ir/pass.cc

@@ -69,6 +69,26 @@ Graph* Pass::Apply(Graph* graph) const {
   return graph;
 }
 
+void Pass::Apply(ProgramDesc* main_program,
+                 ProgramDesc* startup_program) const {
+  PADDLE_ENFORCE_NOT_NULL(main_program, platform::errors::InvalidArgument(
+                                            "main program must be provided"));
+  PADDLE_ENFORCE_NOT_NULL(
+      startup_program,
+      platform::errors::InvalidArgument("startup program must be provided"));
+
+  Graph graph(*main_program);
+  Apply(&graph);
+
+  // TODO(zjl): support details::kStartupProgramDescs and details::kProgramDescs
+  ProgramDesc new_main_program;
+  GraphToProgram(graph, &new_main_program);
+  main_program->CopyFrom(*new_main_program.Proto());
+
+  startup_program->Flush();
+  main_program->Flush();
+}
+
 PassRegistry& PassRegistry::Instance() {
   static PassRegistry g_pass_info_map;
   return g_pass_info_map;