1 | /** |
2 | * Copyright (c) Glow Contributors. See CONTRIBUTORS file. |
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 | #ifndef GLOW_GRAPH_FXIRWRAPPER_H |
18 | #define GLOW_GRAPH_FXIRWRAPPER_H |
19 | |
20 | #include "glow/Graph/Graph.h" |
21 | #include "llvm/ADT/MapVector.h" |
22 | |
23 | #include <folly/dynamic.h> |
24 | |
25 | namespace glow { |
26 | |
27 | using FXModule = folly::dynamic; |
28 | using FXNode = folly::dynamic; |
29 | using FXNodeList = folly::dynamic; |
30 | using FXWeight = folly::dynamic; |
31 | |
32 | /// A thin wrapper around FX IR to provides common APIs. This is not mean to be |
33 | /// yet another IR, e.g., Glow high level graph IR. |
34 | class FXIRWrapper final : public IRContainer { |
35 | /// Mapping from constant name to a void pointer points to where the constant |
36 | /// is actually stored. |
37 | const llvm::StringMap<const void *> &constants_; |
38 | |
39 | /// A reference to the owner of the function. |
40 | Module *parent_; |
41 | |
42 | /// Mapping from getattrs to the underlying name of the constants they alias. |
43 | llvm::StringMap<const std::string> getattrs_; |
44 | |
45 | /// A reference to the FXIR that this FXIRWrapper wraps. |
46 | const FXModule *fx_mod_; |
47 | |
48 | /// Map nodes to their names. |
49 | llvm::StringMap<const FXNode &> namedNodes_; |
50 | |
51 | public: |
52 | FXIRWrapper(const FXModule &FXIR, const std::string &submod, |
53 | const llvm::StringMap<const void *> &constants, |
54 | Module *glowModule, llvm::StringRef name = {}) |
55 | : IRContainer(name), constants_(constants), parent_(glowModule) { |
56 | fx_mod_ = submod == "" ? &FXIR : &(FXIR["modules" ][submod]); |
57 | // Create mapping from getattrs to the underlying name of the constants they |
58 | // alias. |
59 | for (const auto &node : fx_mod_->at("nodes" )) { |
60 | const auto &opCode = node["op_code" ].getString(); |
61 | if (opCode == "get_attr" ) { |
62 | const auto &nodeName = node["name" ].getString(); |
63 | bool inserted = |
64 | getattrs_.try_emplace(nodeName, node["target" ].getString()).second; |
65 | CHECK(inserted) << "Already mapped a getattr by name " << nodeName |
66 | << " to its underlying Constant" ; |
67 | } |
68 | const auto &nodeName = node["name" ].getString(); |
69 | namedNodes_.try_emplace(nodeName, node); |
70 | } |
71 | } |
72 | |
73 | ~FXIRWrapper() = default; |
74 | |
75 | /// A map to store (key) node name to (value) placeholder/constant. |
76 | llvm::StringMap<const Storage *> mapNodeNameToStorage_ = {}; |
77 | |
78 | /// A map to store (key) Glow constant(storage) node name to FX IR weight node |
79 | /// name. This is in spirit the reverse of the above map. |
80 | /// Example use of this is to map constant node name from the symbol table to |
81 | /// a FX weight node name. |
82 | llvm::StringMap<std::string> mapGlowConstNodeToFXNodeName = {}; |
83 | |
84 | IRKind getIRKind() const override { return IRKind::GlowFXIRKind; }; |
85 | |
86 | static bool classof(const IRContainer *I) { |
87 | return I->getIRKind() == IRKind::GlowFXIRKind; |
88 | } |
89 | |
90 | static bool classof(const FXIRWrapper *F) { return true; } |
91 | |
92 | /// \returns the name of input \p node. Fatals if \p node is not |
93 | /// specified as is_node. If \p optional then \returns an empty string if \p |
94 | /// node is null. |
95 | const std::string &getInputNodeName(const folly::dynamic &node, |
96 | bool optional = false) const; |
97 | |
98 | /// \returns the underlying name of a Constant given provided \p name. If \p |
99 | /// name is already the name of a Constant it is returned, else looks for |
100 | /// getattr aliases to return the name of the actual underlying Constant. |
101 | const char *getConstantName(llvm::StringRef name) const; |
102 | |
103 | /// \returns the underlying constant pointer given provided \p name. If \p |
104 | /// name is already the name of a Constant, then the name is directly used, |
105 | /// else looks for getattr aliases to find the name of the actual underlying |
106 | /// Constant. |
107 | const void *getConstant(llvm::StringRef name); |
108 | |
109 | /// \returns true if a constant with given \p name exists. |
110 | bool constantExists(llvm::StringRef name) const; |
111 | |
112 | /// \returns the weights of the graph. |
113 | const FXWeight &getWeight() const; |
114 | |
115 | /// \returns the nodes of the graph. |
116 | const FXNodeList &getNodes() const; |
117 | |
118 | /// \returns parent module that owns this graph. |
119 | Module *getParent() override { return parent_; } |
120 | const Module *getParent() const override { return parent_; } |
121 | |
122 | /// \returns fx module. |
123 | const FXModule &getFXModule() const { return *fx_mod_; } |
124 | |
125 | /// \returns the name of the node. |
126 | const FXNode &getFXNodeByName(llvm::StringRef nodeName) const; |
127 | |
128 | const llvm::StringMap<const Storage *> &getMapNodeNameToStorage() const { |
129 | return mapNodeNameToStorage_; |
130 | } |
131 | |
132 | const llvm::StringMap<const void *> &getConstantsStringMap() const { |
133 | return constants_; |
134 | } |
135 | |
136 | /// For a given weights node, get the underlying Storage. |
137 | const Storage *getStorageFromNodeName(llvm::StringRef name) const; |
138 | |
139 | /// Given the name of Glow Constant Node, return the FX Weight node name. |
140 | const std::string |
141 | getFXWeightNameFromGlowConstNodeName(llvm::StringRef name) const; |
142 | |
143 | /// When FXIR has the notion of memory/buffers. This function returns |
144 | /// the memory buffer a node(operator) writes to. |
145 | const FXNode &getDestinationBufferForNode(const FXNode &node) const; |
146 | }; |
147 | |
148 | } // namespace glow |
149 | |
150 | #endif // GLOW_GRAPH_FXIRWRAPPER_H |
151 | |