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
25namespace glow {
26
27using FXModule = folly::dynamic;
28using FXNode = folly::dynamic;
29using FXNodeList = folly::dynamic;
30using 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.
34class 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
51public:
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