From dbab352861ffad834f77cf7a3847f6a368a84e3b Mon Sep 17 00:00:00 2001
From: sheng <sheng.yang1@huawei.com>
Date: Fri, 18 Dec 2020 00:51:53 +0100
Subject: [PATCH] support user defined strategy

---
 .../auto_parallel/rec_core/rec_generate_strategy.cc      | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc b/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc
index f00da76532..6716acd3b3 100644
--- a/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc
+++ b/mindspore/ccsrc/frontend/parallel/auto_parallel/rec_core/rec_generate_strategy.cc
@@ -25,6 +25,7 @@
 #include "frontend/parallel/auto_parallel/rec_core/rec_partition.h"
 #include "frontend/parallel/ops_info/operator_info.h"
 #include "frontend/parallel/strategy.h"
+#include "frontend/parallel/step_parallel.h"
 
 namespace mindspore {
 namespace parallel {
@@ -43,6 +44,14 @@ void GenerateStrategy(const std::shared_ptr<Graph> &graph, const std::vector<std
   GenerateEliminatedOperatorStrategyForward(graph, ops, input_tensor_names, index_list, no_stra_op_list);
   GenerateEliminatedOperatorStrategyBackward(ops, input_tensor_names, no_stra_op_list);
   GenerateRemainingOperatorStrategy(graph, ops, input_tensor_names, index_list, no_stra_op_list);
+
+  for (auto &op : ops) {
+    auto attrs = op->attrs();
+    if (StrategyFound(attrs)) {
+      StrategyPtr user_defined_stra = parallel::ExtractStrategy(attrs);
+      op->SetSelectedStrategyAndCost(user_defined_stra, op->selected_cost());
+    }
+  }
 }
 
 Strategys PrepareMatMul(const std::shared_ptr<Graph> &graph, const std::vector<std::shared_ptr<OperatorInfo>> &ops,