!29946 Clear device tensor in front value node.

Merge pull request !29946 from gaoyong10/r1.6_2
This commit is contained in:
i-robot 2022-02-11 12:13:14 +00:00 committed by Gitee
commit f8da97b31f
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
3 changed files with 16 additions and 3 deletions

View File

@ -343,7 +343,7 @@ void UpdateRefCountForGraphOutput(const std::vector<KernelWithIndex> &output_wit
} // namespace
GraphCompilerInfo::~GraphCompilerInfo() {
GraphScheduler::GetInstance().Clear(name_, graphs_, origin_parameters_order_);
GraphScheduler::GetInstance().Clear(name_, graphs_, origin_parameters_order_, control_node_parser_);
}
GraphId GraphCompiler::CompileGraph(const GraphSegmentPtr &segment, const AnfNodePtrList &outputs,

View File

@ -182,7 +182,8 @@ void IntHandler(int, siginfo_t *, void *) {
} // namespace
void GraphScheduler::Clear(const ActorInfo &actor_info, const std::vector<KernelGraphPtr> &graphs,
const std::vector<AnfNodePtr> &root_graph_parameters) noexcept {
const std::vector<AnfNodePtr> &root_graph_parameters,
const ControlNodeParserPtr &parser) noexcept {
// Terminate the actors of actor info.
if (actors_.count(actor_info) > 0) {
auto actor_manager = ActorMgr::GetActorMgrRef();
@ -204,6 +205,17 @@ void GraphScheduler::Clear(const ActorInfo &actor_info, const std::vector<Kernel
ClearNodeInfo(graph);
}
if (parser != nullptr && parser->IsInited()) {
const auto &front_value_nodes = parser->front_value_nodes();
for (const auto &front_value_node : front_value_nodes) {
const auto &node = front_value_node.first.first;
size_t index = front_value_node.first.second;
if (AnfAlgo::OutputAddrExist(node, index)) {
AnfAlgo::SetOutputAddr(nullptr, index, node.get());
}
}
}
// Clear the member of DeviceTensorStore.
for (auto &root_graph_parameter : root_graph_parameters) {
DeviceTensorStore::GetInstance().Remove(root_graph_parameter.get());

View File

@ -57,7 +57,8 @@ class GraphScheduler {
// Clear the members.
void Clear();
void Clear(const ActorInfo &actor_info, const std::vector<KernelGraphPtr> &graphs,
const std::vector<AnfNodePtr> &root_graph_parameters) noexcept;
const std::vector<AnfNodePtr> &root_graph_parameters,
const ControlNodeParserPtr &parser = nullptr) noexcept;
// The control flow actors will generate some data in the loop body execution, so need clear on the end of execution.
void ClearActorData(const ActorSet *actor_set);