[PIR] Fix exe bug of hang (#61491)

zhangbo9674 · web-flow · commit ff98389b3962 · 2024-02-04T11:00:42.000+08:00
* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/framework/new_executor/interpreter/dependency_builder.cc b/paddle/fluid/framework/new_executor/interpreter/dependency_builder.cc
@@ -565,6 +565,40 @@ PirDependencyBuilder::PirDependencyBuilder() {
   op_happens_before_ = std::make_shared<std::vector<std::vector<bool>>>();
 }
 
+void PirDependencyBuilder::AddDependencyForCommunicationOp() {
+  size_t dependence_op_idx = ULLONG_MAX;
+  for (size_t op_idx = 0; op_idx < op_num_; ++op_idx) {
+    if (instructions_.at(op_idx)->Operation() &&
+        IsCommunicationOp(instructions_.at(op_idx)->Operation())) {
+      if (dependence_op_idx != ULLONG_MAX) {
+        AddDownstreamOp(dependence_op_idx, op_idx);
+      }
+      dependence_op_idx = op_idx;
+    }
+  }
+
+  // TODO(zhiqiu): there still some cases not handled
+  // add dependency for c_sync_comm_stream
+
+  // in program, we can add only one c_sync_comm_stream to sync all
+  // communication ops.
+  // c_allreduce_sum(a)
+  // c_allreduce_sum(b)
+  // c_allreduce_sum(c)
+  // c_sync_comm_stream(a)
+  const std::string kSyncComm = dialect::CSyncCommStreamOp::name();
+  dependence_op_idx = ULLONG_MAX;
+  for (size_t op_idx = 0; op_idx < op_num_; ++op_idx) {
+    if (instructions_.at(op_idx)->Name() == kSyncComm) {
+      dependence_op_idx = op_idx;
+    } else {
+      if (dependence_op_idx != ULLONG_MAX) {
+        AddDownstreamOp(dependence_op_idx, op_idx);
+      }
+    }
+  }
+}
+
 void PirDependencyBuilder::AddDependencyForRandomOp() {
   const std::set<std::string> random_op_set = {
       dialect::BernoulliOp::name(),
@@ -616,6 +650,11 @@ const std::map<size_t, std::set<size_t>>& PirDependencyBuilder::Build(
     AddDependencyForSequentialRun();
   }
 
+  if (FLAGS_add_dependency_for_communication_op) {
+    AddDependencyForCommunicationOp();
+    VLOG(6) << "Finish AddDependencyForSequentialRun";
+  }
+
   // TODO(zhangbo): Add dependency for special op ？
   // Note(lvyongkang): necessary for reproducibility
   AddDependencyForRandomOp();
diff --git a/paddle/fluid/framework/new_executor/interpreter/dependency_builder.h b/paddle/fluid/framework/new_executor/interpreter/dependency_builder.h
@@ -59,7 +59,7 @@ class DependencyBuilder {
 
  protected:
   void AddDependencyForCoalesceTensorOp();
-  void AddDependencyForCommunicationOp();
+  virtual void AddDependencyForCommunicationOp();
   virtual void AddDependencyForRandomOp();
   void AddDependencyForReadOp();
   void AddDependencyForSequentialRun();
@@ -117,6 +117,8 @@ class PirDependencyBuilder : public DependencyBuilder {
   void ShareDependencyFrom(const PirDependencyBuilder& src);
 
  private:
+  void AddDependencyForCommunicationOp() override;
+
   void AddDependencyForRandomOp() override;
 
   std::vector<paddle::framework::InstructionBase*> instructions_;  // not_owned
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -1321,7 +1321,6 @@ paddle::framework::FetchList PirInterpreter::Run(
     PreAnalysis();
     VLOG(4) << "Done PreAnalysis";
 
-    // Run
     if (FLAGS_enable_pir_in_executor_trace_run || onednn_op_num_ ||
         execution_config_.used_for_inference ||
         ((execution_config_.used_for_jit || execution_config_.used_for_cinn) &&
