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
31namespace glow {
32
33/// This class implements \p NodeWalker interface to translate from Node to
34/// Instruction IR.
35class IRGenVisitor : public NodeWalker {
36private:
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
57public:
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