From 535219ba43406465753b69d8d24aeaea2668181c Mon Sep 17 00:00:00 2001
From: Margaret_wangrui <wangrui178@huawei.com>
Date: Wed, 6 Apr 2022 10:11:59 +0800
Subject: [PATCH] Avoid adding duplicate umonad to output of Bprop funcgraph.

---
 .../ccsrc/frontend/optimizer/ad/kprim.cc      | 24 ++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc b/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc
index b5e92de35ad..e7a023c0c10 100644
--- a/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc
+++ b/mindspore/ccsrc/frontend/optimizer/ad/kprim.cc
@@ -471,6 +471,24 @@ MetaFuncGraphPtr KPrim::KMetaFuncGraph(const PrimitivePtr &prim) {
   MS_LOG(EXCEPTION) << "Fail to find bprop function for " << prim->name() << ".";
 }
 
+bool HasUMonadInput(const CNodePtr &node) {
+  for (size_t index = node->inputs().size() - 1; index > 0; --index) {
+    if (HasAbstractUMonad(node->input(index))) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool HasIOMonadInput(const CNodePtr &node) {
+  for (size_t index = node->inputs().size() - 1; index > 0; --index) {
+    if (HasAbstractIOMonad(node->input(index))) {
+      return true;
+    }
+  }
+  return false;
+}
+
 static void AppendMonadOutput(const FuncGraphPtr &bprop_fg, const AnfNodePtr &monad) {
   const auto &output = bprop_fg->output();
   MS_EXCEPTION_IF_NULL(output);
@@ -484,7 +502,11 @@ static void AppendMonadOutput(const FuncGraphPtr &bprop_fg, const AnfNodePtr &mo
     }
   }
   if (output_cnode != nullptr) {
-    output_cnode->add_input(monad);
+    bool need_umonad = HasAbstractUMonad(monad) && !HasUMonadInput(output_cnode);
+    bool need_iomonad = HasAbstractIOMonad(monad) && !HasIOMonadInput(output_cnode);
+    if (need_umonad || need_iomonad) {
+      output_cnode->add_input(monad);
+    }
     return;
   }
   // If output is an empty tuple, create a (make_tuple, monad) as the new output.