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_GRAPH_HOOK_H |
17 | #define GLOW_GRAPH_HOOK_H |
18 | |
19 | #include "llvm/ADT/StringRef.h" |
20 | #include <list> |
21 | |
22 | namespace glow { |
23 | |
24 | class Function; |
25 | class Node; |
26 | class Placeholder; |
27 | class SaveNode; |
28 | |
29 | /// This struct is used to keep information returned from |
30 | /// hooking a function. |
31 | struct HookedFunction { |
32 | /// Saves the function hookNode() creates and inserts hooks in. |
33 | Function *function; |
34 | /// List of Save nodes hookNode() inserts at the outputs of the layer. |
35 | std::list<SaveNode *> outputSaves; |
36 | /// List of the Placeholders associated with Save nodes in outputSaves. |
37 | std::list<Placeholder *> outputs; |
38 | /// List of Save nodes hookNode() inserts at the input of the layer. |
39 | std::list<SaveNode *> inputSaves; |
40 | /// List of the Placeholders associated with Save nodes in inputSaves. |
41 | std::list<Placeholder *> inputs; |
42 | }; |
43 | |
44 | /// Given a function \p F and a node \p node it creates a function |
45 | /// in the same module and populates it with a recursive clone of |
46 | /// the node \p node then inserts Save nodes at the output of to capture it. |
47 | /// If \p hookInputs is set it also inserts Save nodes to capture the inputs. |
48 | /// \returns the list of inserted nodes and associated placeholder in |
49 | /// \p HookedFunction object. |
50 | HookedFunction hookNode(Function *F, Node *node, bool hookInputs = false); |
51 | |
52 | /// Given a function \p F and a layer name \p nodeName it finds the |
53 | /// corresponding nodes in the function \p F then calls above override. |
54 | HookedFunction hookNode(Function *F, llvm::StringRef nodeName, |
55 | bool hookInputs = false); |
56 | |
57 | } // namespace glow |
58 | |
59 | #endif // GLOW_GRAPH_HOOK_H |
60 | |