1 /**
2 * This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
3 *
4 * Copyright 2019-2022 Huawei Technologies Co., Ltd
5 *
6 * Licensed under the Apache License, Version 2.0 (the "License");
7 * you may not use this file except in compliance with the License.
8 * You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing, software
13 * distributed under the License is distributed on an "AS IS" BASIS,
14 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 * See the License for the specific language governing permissions and
16 * limitations under the License.
17 */
18
19 #ifndef MINDSPORE_CORE_IR_GRAPH_UTILS_H_
20 #define MINDSPORE_CORE_IR_GRAPH_UTILS_H_
21
22 #include <functional>
23 #include <string>
24 #include <vector>
25
26 #include "utils/hash_map.h"
27 #include "utils/hash_set.h"
28 #include "ir/anf.h"
29 #include "ir/primitive.h"
30 #include "ir/scalar.h"
31 #include "ir/tensor.h"
32 #include "utils/label.h"
33 #include "mindapi/base/macros.h"
34
35 namespace mindspore {
36 enum IncludeType { FOLLOW, NOFOLLOW, EXCLUDE };
37
38 using IncludeFunc = std::function<IncludeType(const AnfNodePtr &)>;
39 using FilterFunc = std::function<bool(const AnfNodePtr &)>;
40 using GraphFilterFunc = std::function<bool(const FuncGraphPtr &)>;
41 using SuccFunc = std::function<AnfNodeWeakPtrList(AnfNodePtr)>;
42 using DeprecatedSuccFunc = std::function<AnfNodePtrList(AnfNodePtr)>;
43 using MatchFunc = std::function<bool(const CNodePtr &)>;
44 using NodeVisitFunc = std::function<void(const AnfNodePtr &)>;
45
46 MS_CORE_API AnfNodeWeakPtrList SuccDeeperSimple(const AnfNodePtr &node);
47 MS_CORE_API AnfNodeWeakPtrList SuccIncoming(const AnfNodePtr &node);
48 AnfNodeWeakPtrList SuccIncludeFV(const FuncGraphPtr &fg, const AnfNodePtr &node);
49 MS_CORE_API AnfNodeWeakPtrList SuccWithFilter(const GraphFilterFunc &graph_filter, const AnfNodePtr &node);
50
51 MS_CORE_API const AnfNodePtrList GetInputs(const AnfNodePtr &node);
52 MS_CORE_API const AnfNodeWeakPtrList &GetWeakInputs(const AnfNodePtr &node);
53
AlwaysInclude(const AnfNodePtr &)54 inline IncludeType AlwaysInclude(const AnfNodePtr &) { return FOLLOW; }
55 MS_CORE_API IncludeType IncludeBelongGraph(const FuncGraphPtr &fg, const AnfNodePtr &node);
56
57 MS_CORE_API AnfNodePtrList DeepScopedGraphSearch(const AnfNodePtr &root, const IncludeFunc &include = AlwaysInclude);
58
59 MS_CORE_API AnfNodePtrList DeepLinkedGraphSearch(const AnfNodePtr &root, const IncludeFunc &include = AlwaysInclude);
60
61 MS_CORE_API AnfNodePtrList DeepScopedGraphSearchWithFilter(const AnfNodePtr &root, const IncludeFunc &include,
62 const FilterFunc &filter);
63
64 MS_CORE_API AnfNodePtrList TopoSort(const AnfNodePtr &root, const SuccFunc &succ = SuccIncoming,
65 const IncludeFunc &include = AlwaysInclude, bool exclude_circle_node = false);
66
67 // @deprecated
68 // To use 'AnfNodePtrList TopoSort(const AnfNodePtr &, const SuccFunc &, const IncludeFunc &)' instead.
69 MS_CORE_API AnfNodePtrList TopoSort(const AnfNodePtr &root, const DeprecatedSuccFunc &deprecated_succ,
70 const IncludeFunc &include = AlwaysInclude, bool exclude_circle_node = false);
71
72 MS_CORE_API std::vector<CNodePtr> BroadFirstSearchGraphCNodes(const CNodePtr &root);
73 std::vector<FuncGraphPtr> BroadFirstSearchGraphUsed(const FuncGraphPtr &root,
74 const GraphFilterFunc &filter = GraphFilterFunc());
75
76 MS_CORE_API CNodePtr BroadFirstSearchFirstOf(const std::vector<CNodePtr> &roots, const MatchFunc &match_predicate);
77
78 using DumpIRPrividerFunction = void (*)(std::ostringstream &, const FuncGraphPtr &, bool, int, bool);
79 MS_CORE_API void SetDumpIRPrivider(const DumpIRPrividerFunction &func);
80 using DumpIRStorageFunction = void (*)(const std::string &, const std::string &, const std::string &);
81 MS_CORE_API void SetDumpIRStorage(const DumpIRStorageFunction &func);
82 constexpr auto kTopoSortCircle = "TopoSortCircle";
83 } // namespace mindspore
84
85 #endif // MINDSPORE_CORE_IR_GRAPH_UTILS_H_
86