From 19bfd2de69c6991cbe21495532e8010b5ac123ee Mon Sep 17 00:00:00 2001
From: yujianfeng <yujianfeng5@huawei.com>
Date: Mon, 18 May 2020 19:03:19 +0800
Subject: [PATCH] Add reduce_scatter communication op fusion pass

---
 .../pre_activate/ascend/ascend_backend_optimization.cc     | 1 +
 .../ccsrc/pre_activate/pass/communication_op_fusion.h      | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc
index 8f3e5a5737e..220b309200d 100644
--- a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc
+++ b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc
@@ -287,6 +287,7 @@ void AscendBackendOptimization(const std::shared_ptr<session::KernelGraph> &kern
   auto other_pm = std::make_shared<PassManager>("other_pm");
   other_pm->AddPass(std::make_shared<AllReduceFusion>());
   other_pm->AddPass(std::make_shared<AllGatherFusion>());
+  other_pm->AddPass(std::make_shared<ReduceScatterFusion>());
   other_pm->AddPass(std::make_shared<BroadcastFusion>());
   other_pm->AddPass(std::make_shared<ParameterTransOpFusion>());
   other_pm->AddPass(std::make_shared<RefreshParameterFormat>());
diff --git a/mindspore/ccsrc/pre_activate/pass/communication_op_fusion.h b/mindspore/ccsrc/pre_activate/pass/communication_op_fusion.h
index e01d1816164..d00180f97f0 100644
--- a/mindspore/ccsrc/pre_activate/pass/communication_op_fusion.h
+++ b/mindspore/ccsrc/pre_activate/pass/communication_op_fusion.h
@@ -68,6 +68,13 @@ class BroadcastFusion : public CommunicationOpFusion {
   explicit BroadcastFusion(size_t groups = 1) : CommunicationOpFusion("broadcast_fusion", kBroadcastOpName, groups) {}
   ~BroadcastFusion() override = default;
 };
+
+class ReduceScatterFusion : public CommunicationOpFusion {
+ public:
+  explicit ReduceScatterFusion(size_t groups = 1)
+      : CommunicationOpFusion("reduce_scatter_fusion", kReduceScatterOpName, groups) {}
+  ~ReduceScatterFusion() override = default;
+};
 }  // namespace opt
 }  // namespace mindspore
 #endif  // MINDSPORE_CCSRC_PRE_ACTIVATE_PASS_COMMUNICATION_OP_FUSION_H_