From 3b7e01c6987f6a2b1905d4d26404cee52662abda Mon Sep 17 00:00:00 2001
From: wilfChen <chenweifeng720@huawei.com>
Date: Wed, 28 Oct 2020 21:02:33 +0800
Subject: [PATCH] fix momentum fusion pass

---
 .../gpu/apply_momentum_scale_fusion.cc        |  3 +--
 .../gpu/apply_momentum_scale_fusion.h         | 23 ++--------------
 .../gpu/apply_momentum_weight_scale_fusion.cc | 22 +++++++++++++++
 .../gpu/apply_momentum_weight_scale_fusion.h  | 27 ++-----------------
 4 files changed, 27 insertions(+), 48 deletions(-)

diff --git a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc
index c21f107805..1c139c3660 100644
--- a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc
+++ b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.cc
@@ -26,7 +26,6 @@
 
 namespace mindspore {
 namespace opt {
-namespace {
 bool ApplyMomentumScaleFusion::IsScalar(const BaseRef &n) {
   if (utils::isa<AnfNodePtr>(n)) {
     AnfNodePtr in = utils::cast<AnfNodePtr>(n);
@@ -86,5 +85,5 @@ const AnfNodePtr ApplyMomentumScaleFusion::Process(const FuncGraphPtr &graph, co
   replace_node->set_scope(node->scope());
   return replace_node;
 }
-}  // namespace
 }  // namespace opt
+}  // namespace mindspore
diff --git a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h
index 349c6be338..8888f40c7b 100644
--- a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h
+++ b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_scale_fusion.h
@@ -37,27 +37,8 @@ class ApplyMomentumScaleFusion : public PatternProcessPass {
   const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
 
  private:
-  static bool IsScalar(const BaseRef &n) {
-    if (utils::isa<AnfNodePtr>(n)) {
-      AnfNodePtr in = utils::cast<AnfNodePtr>(n);
-      MS_EXCEPTION_IF_NULL(in);
-      auto shape = in->Shape()->cast<abstract::ShapePtr>();
-      MS_EXCEPTION_IF_NULL(shape);
-      if (shape->shape().size() != 0) {
-        return false;
-      }
-      auto dtype = in->Type();
-      if (dtype->type_id() != kObjectTypeTensorType) {
-        return false;
-      }
-      auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
-      if (element_type != kNumberTypeFloat32) {
-        return false;
-      }
-      return true;
-    }
-    return false;
-  }
+  static bool IsScalar(const BaseRef &n);
+
   VarPtr scale_;
   VarPtr variable_;
   VarPtr accumulation_;
diff --git a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc
index 9e235a756f..743015c50c 100644
--- a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc
+++ b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.cc
@@ -26,6 +26,28 @@
 
 namespace mindspore {
 namespace opt {
+bool ApplyMomentumWeightDecayScaleFusion::IsScalar(const BaseRef &n) {
+  if (utils::isa<AnfNodePtr>(n)) {
+    AnfNodePtr in = utils::cast<AnfNodePtr>(n);
+    MS_EXCEPTION_IF_NULL(in);
+    auto shape = in->Shape()->cast<abstract::ShapePtr>();
+    MS_EXCEPTION_IF_NULL(shape);
+    if (shape->shape().size() != 0) {
+      return false;
+    }
+    auto dtype = in->Type();
+    if (dtype->type_id() != kObjectTypeTensorType) {
+      return false;
+    }
+    auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
+    if (element_type != kNumberTypeFloat32) {
+      return false;
+    }
+    return true;
+  }
+  return false;
+}
+
 const BaseRef ApplyMomentumWeightDecayScaleFusion::DefinePattern() const {
   VectorRef weight = VectorRef(
     {prim::kPrimAddN, VectorRef({prim::kPrimMul, variable_, weight_decay_}), VectorRef({prim::kPrimCast, gradient_})});
diff --git a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h
index d8fb16e7b9..c1b92c8242 100644
--- a/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h
+++ b/mindspore/ccsrc/backend/optimizer/gpu/apply_momentum_weight_scale_fusion.h
@@ -21,30 +21,6 @@
 
 namespace mindspore {
 namespace opt {
-namespace {
-bool IsScalar(const BaseRef &n) {
-  if (utils::isa<AnfNodePtr>(n)) {
-    AnfNodePtr in = utils::cast<AnfNodePtr>(n);
-    MS_EXCEPTION_IF_NULL(in);
-    auto shape = in->Shape()->cast<abstract::ShapePtr>();
-    MS_EXCEPTION_IF_NULL(shape);
-    if (shape->shape().size() != 0) {
-      return false;
-    }
-    auto dtype = in->Type();
-    if (dtype->type_id() != kObjectTypeTensorType) {
-      return false;
-    }
-    auto element_type = dyn_cast<TensorType>(dtype)->element()->type_id();
-    if (element_type != kNumberTypeFloat32) {
-      return false;
-    }
-    return true;
-  }
-  return false;
-}
-}  // namespace
-
 class ApplyMomentumWeightDecayScaleFusion : public PatternProcessPass {
  public:
   explicit ApplyMomentumWeightDecayScaleFusion(bool multigraph = true)
@@ -62,9 +38,10 @@ class ApplyMomentumWeightDecayScaleFusion : public PatternProcessPass {
   const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override;
 
  private:
+  static bool IsScalar(const BaseRef &n);
+
   VarPtr weight_decay_;
   VarPtr scale_;
-
   VarPtr variable_;
   VarPtr accumulation_;
   VarPtr learning_rate_;