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 | #ifndef GLOW_IR_IRGEN_H |
17 | #define GLOW_IR_IRGEN_H |
18 | |
19 | #include "glow/Backend/Backend.h" |
20 | #include "glow/Graph/Graph.h" |
21 | #include "glow/Graph/Nodes.h" |
22 | #include "glow/IR/IR.h" |
23 | #include "glow/IR/IRBuilder.h" |
24 | |
25 | #include <unordered_set> |
26 | |
27 | //===----------------------------------------------------------------------===// |
28 | // IRGen visitor - the code that generates the IR. |
29 | //===----------------------------------------------------------------------===// |
30 | |
31 | namespace glow { |
32 | |
33 | /// This class implements \p NodeWalker interface to translate from Node to |
34 | /// Instruction IR. |
35 | class IRGenVisitor : public NodeWalker { |
36 | private: |
37 | using NodeValueToDestTy = std::unordered_map<NodeValue, Value *>; |
38 | using NodeToInstrTy = std::unordered_map<Node *, Instruction *>; |
39 | |
40 | /// A set of visited nodes. |
41 | std::unordered_set<Node *> visited_; |
42 | /// Holds the mapping between graph nodes to the destination buffers. |
43 | NodeValueToDestTy generatedNodeDest_; |
44 | /// Holds the mapping between graph nodes and the lowered instructions. This |
45 | /// map is used by instructions that want to inspect the generated |
46 | /// instructions. For example, gradient instructions that look at operands |
47 | /// that do not exist at the graph level. Not all variables are representible. |
48 | NodeToInstrTy nodeToInstr_; |
49 | |
50 | /// The function that we are building. |
51 | IRFunction *F_; |
52 | /// The builder that adds instructions into the function. |
53 | IRBuilder builder_; |
54 | /// The Backend /p B is used for custom lowering of Node to Instruction IR. |
55 | const Backend &B_; |
56 | |
57 | public: |
58 | bool shouldVisit(Node *parent, Node *N) override; |
59 | |
60 | explicit IRGenVisitor(IRFunction *M, const Backend &B) |
61 | : F_(M), builder_(F_), B_(B) {} |
62 | |
63 | /// \returns the generated instruction for the node \p N. |
64 | Value *valueForNode(NodeValue N); |
65 | |
66 | /// Saves the generated IR in \p v for the node \p N. |
67 | void registerIR(NodeValue N, Value *v); |
68 | |
69 | /// Adds to Node \p N --> Instruction \p inst map. |
70 | void setNodeToIR(Node *N, Instruction *inst); |
71 | |
72 | /// Return Instruction that is mapped to Node \p N. |
73 | /// If mapping doesn't exists returns nullptr. |
74 | Instruction *getNodeToIR(Node *N); |
75 | |
76 | void post(Node *parent, Node *N) override; |
77 | |
78 | IRBuilder *getBuilder() { return &builder_; } |
79 | }; |
80 | |
81 | } // namespace glow |
82 | #endif // GLOW_IR_IRGEN_H |
83 | |