1 /**
2 * Copyright 2020 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <regex>
18 #include "frontend/parallel/graph_util/graph_info.h"
19 #include "debug/anf_ir_dump.h"
20 #include "debug/anf_ir_utils.h"
21 #include "debug/draw.h"
22 #include "utils/ms_context.h"
23 #include "ir/graph_utils.h"
24 #include "pipeline/jit/pipeline.h"
25
26 namespace mindspore {
27 namespace parallel {
FindPrimtive(const FuncGraphPtr & graph,const std::string & name)28 std::vector<PrimitivePtr> FindPrimtive(const FuncGraphPtr &graph, const std::string &name) {
29 AnfNodePtr ret = graph->get_return();
30 MS_EXCEPTION_IF_NULL(ret);
31 std::vector<AnfNodePtr> all_nodes = DeepScopedGraphSearch(ret);
32 std::vector<PrimitivePtr> prim_list;
33 for (auto &node : all_nodes) {
34 if (!IsValueNode<Primitive>(node)) {
35 continue;
36 }
37 ValueNodePtr prim_node_anf = node->cast<ValueNodePtr>();
38 MS_EXCEPTION_IF_NULL(prim_node_anf);
39 PrimitivePtr node_prim = prim_node_anf->value()->cast<PrimitivePtr>();
40 MS_EXCEPTION_IF_NULL(node_prim);
41 if (node_prim->name() == name) {
42 prim_list.emplace_back(node_prim);
43 }
44 }
45 return prim_list;
46 }
47
DumpGraph(const FuncGraphPtr & root,const std::string & name)48 void DumpGraph(const FuncGraphPtr &root, const std::string &name) {
49 #ifdef ENABLE_DUMP_IR
50 if (MsContext::GetInstance()->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG)) {
51 draw::Draw(name + ".dot", root);
52 DumpIR(name + ".ir", root);
53 ExportIR(name + ".dat", root);
54 }
55 #endif
56 }
57
58 // Return true if the cnode is in a for-loop and loop_index indicates the i-th loop;
59 // otherwise return false
GetLoopIndexFromCNode(const CNodePtr & cnode,size_t * loop_index)60 bool GetLoopIndexFromCNode(const CNodePtr &cnode, size_t *loop_index) {
61 std::regex pattern(CELLLIST_KEYWORD_PATTERN);
62 std::smatch result;
63 const auto &cnode_fullname = cnode->fullname_with_scope();
64 if (std::regex_search(cnode_fullname, result, pattern)) {
65 if (result.length() < 2) {
66 MS_LOG(EXCEPTION) << "Wrong format of fullname_with_scope: " << cnode_fullname;
67 }
68 *loop_index = IntToSize(std::stoi(result[1]));
69 return true;
70 }
71 return false;
72 }
73
SetOpsNumToExecutor(size_t num_ops)74 void SetOpsNumToExecutor(size_t num_ops) {
75 auto executor = pipeline::GraphExecutorPy::GetInstance();
76 executor->SetNumOpsInfo(num_ops);
77 }
78 } // namespace parallel
79 } // namespace mindspore
80