Skip to content

Commit 5fff8d7

Browse files
committed
add distributed training deps.
1 parent 4dbcb97 commit 5fff8d7

File tree

1 file changed

+43
-0
lines changed

1 file changed

+43
-0
lines changed

paddle/fluid/framework/ir/graph.cc

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,49 @@ Graph::Graph(const ProgramDesc &program) : program_(program) {
6161
var->inputs.push_back(node);
6262
}
6363
}
64+
65+
std::vector<ir::Node *> send_ops;
66+
ir::Node *send_bar = nullptr;
67+
std::vector<ir::Node *> recv_ops;
68+
ir::Node *fetch_bar = nullptr;
69+
for (ir::Node *node : Nodes()) {
70+
if (node->Name() == "send") {
71+
send_ops.push_back(node);
72+
} else if (node->Name() == "send_barrier") {
73+
PADDLE_ENFORCE(!send_bar, "only has one send barrier");
74+
send_bar = node;
75+
} else if (node->Name() == "recv") {
76+
recv_ops.push_back(node);
77+
} else if (node->Name() == "fetch_barrier") {
78+
PADDLE_ENFORCE(!fetch_bar, "only has one fetch barrier");
79+
fetch_bar = node;
80+
}
81+
}
82+
if (send_bar) {
83+
for (ir::Node *send : send_ops) {
84+
ir::Node *dep_var = CreateControlDepVar();
85+
send->outputs.push_back(dep_var);
86+
dep_var->inputs.push_back(send);
87+
send_bar->inputs.push_back(dep_var);
88+
dep_var->outputs.push_back(send_bar);
89+
}
90+
for (ir::Node *recv : recv_ops) {
91+
ir::Node *dep_var = CreateControlDepVar();
92+
recv->inputs.push_back(dep_var);
93+
dep_var->outputs.push_back(recv);
94+
send_bar->outputs.push_back(dep_var);
95+
dep_var->inputs.push_back(send_bar);
96+
}
97+
}
98+
if (fetch_bar) {
99+
for (ir::Node *recv : recv_ops) {
100+
ir::Node *dep_var = CreateControlDepVar();
101+
recv->outputs.push_back(dep_var);
102+
dep_var->inputs.push_back(recv);
103+
fetch_bar->inputs.push_back(dep_var);
104+
dep_var->outputs.push_back(fetch_bar);
105+
}
106+
}
64107
/**
65108
* We only handle write after read(WAR), since it should not have a write
66109
* after write in program. If there are write after write operators, we need

0 commit comments

Comments
 (0)