From a910db149813e1c8f7f90cbc3bc0e5f481bcac20 Mon Sep 17 00:00:00 2001
From: tronzhang <6517937+tronzhang@user.noreply.gitee.com>
Date: Wed, 20 Jan 2021 11:32:56 +0800
Subject: [PATCH] atomic add should not enable for some case

---
 .../backend/optimizer/graph_kernel/add_atomic_clean_gpu.cc    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean_gpu.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean_gpu.cc
index 2b73ff8840..a7f3df0a04 100644
--- a/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean_gpu.cc
+++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/add_atomic_clean_gpu.cc
@@ -133,8 +133,10 @@ bool AtomicCleanInsertter::CanActivateAtomicAdd(const AnfNodePtr &anf_node) {
     if (reduce_cnt != 1) {
       return false;
     }
+    real_output_num_ = inputs.size() - 1;
   } else if (IsPrimitiveCNode(real_return_node, prim::kPrimReduceSum)) {
     atomic_add_node_ = real_return_node->cast<CNodePtr>();
+    real_output_num_ = 1;
   } else {
     return false;
   }
@@ -200,7 +202,6 @@ void AtomicCleanInsertter::CreateInplaceAssignNodeAndCorrectReturn(const FuncGra
   auto retrun_node = sub_graph->get_return()->input(kFirstDataInputIndex);
   if (IsPrimitiveCNode(retrun_node, prim::kPrimMakeTuple)) {
     const auto &outs = retrun_node->cast<CNodePtr>()->inputs();
-    real_output_num_ = outs.size() - 1;
     for (size_t i = 1; i < outs.size(); ++i) {
       if (i != reduce_real_output_index_ + 1) {
         out_node = outs[i];
@@ -209,7 +210,6 @@ void AtomicCleanInsertter::CreateInplaceAssignNodeAndCorrectReturn(const FuncGra
       }
     }
   } else {
-    real_output_num_ = 1;
     out_node = atomic_add_node_;  // Use result data itself, and set attr "fake_out" true.
     fake_out = true;
   }