1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | #ifndef TVM_RELAY_PRINTER_META_DATA_H_ |
20 | #define TVM_RELAY_PRINTER_META_DATA_H_ |
21 | |
22 | #include <tvm/node/serialization.h> |
23 | |
24 | #include <string> |
25 | #include <unordered_map> |
26 | |
27 | #include "doc.h" |
28 | |
29 | namespace tvm { |
30 | namespace relay { |
31 | /*! |
32 | * \brief Meta data context for Printers |
33 | * |
34 | * This is an important part to enable bi-directional serializability. |
35 | * We use tvm's Node system to build the current IR. |
36 | * It can be hard to design a text format for all the possible nodes |
37 | * as the set of nodes can grow when we do more extensions. |
38 | * |
39 | * Instead of trying to design readable text format for every node, |
40 | * we support a meta data section in the text format. |
41 | * We allow the text format to refer to a node in the meta data section. |
42 | * |
43 | * The meta data section is a json serialized string of an Map<string, Array<NodeRef>>. |
44 | * Each element in the meta data section can be referenced by the text format. |
45 | * Each meta data node is printed in the following format. |
46 | * |
47 | * meta[type-key-of-node>][<index-in-meta-section>] |
48 | * |
49 | * Specifically, consider the following IR(constructed by python). |
50 | * |
51 | * \code |
52 | * |
53 | * n = tvm.var("n") |
54 | * x = tvm.relay.var("x", shape=(n, 1)) |
55 | * f = tvm.relay.Function([x], x) |
56 | * print(f.astext()) |
57 | * |
58 | * \endcode |
59 | * |
60 | * The corresponding text format is shown in the following code block. |
61 | * |
62 | * \code |
63 | * |
64 | * fn (%x: Tensor[(meta[Variable][0],), float32]) { |
65 | * %x |
66 | * } |
67 | * # Meta data section is a json-serialized string |
68 | * # of the following array. |
69 | * # [tvm.var("n")] |
70 | * |
71 | * \endcode |
72 | * |
73 | * Note that we store tvm.var("n") in the meta data section. |
74 | * Since it is stored in the index-0 in the meta data section, |
75 | * we print it as meta[Variable][0]. |
76 | * |
77 | * The text parser can recover this object by loading from the corresponding |
78 | * location in the meta data section. |
79 | * |
80 | * This is a design trade-off. |
81 | * It allows us to embedded any meta data in the text format, |
82 | * while still being able to tweak the text part of the printed IR easily. |
83 | */ |
84 | class TextMetaDataContext { |
85 | public: |
86 | /*! |
87 | * \brief Get text representation of meta node. |
88 | * \param node The node to be converted to meta node. |
89 | * \return A string representation of the meta node. |
90 | */ |
91 | Doc GetMetaNode(const ObjectRef& node) { |
92 | auto it = meta_repr_.find(node); |
93 | if (it != meta_repr_.end()) { |
94 | return it->second; |
95 | } |
96 | std::string type_key = node->GetTypeKey(); |
97 | ICHECK(!type_key.empty()); |
98 | Array<ObjectRef>& mvector = meta_data_[type_key]; |
99 | int64_t index = static_cast<int64_t>(mvector.size()); |
100 | mvector.push_back(node); |
101 | Doc doc; |
102 | doc << "meta[" << type_key << "][" << index << "]" ; |
103 | meta_repr_[node] = doc; |
104 | return meta_repr_[node]; |
105 | } |
106 | |
107 | /*! |
108 | * \brief Test whether a node has been put in meta |
109 | * \param node The query node |
110 | * \return whether the node has been put in meta |
111 | */ |
112 | bool InMeta(const ObjectRef& node) { return meta_repr_.find(node) != meta_repr_.end(); } |
113 | |
114 | /*! |
115 | * \brief Print a key value pair |
116 | */ |
117 | Doc PrintKeyValue(const std::string& str, const Doc& v) const { |
118 | return Doc() << "\"" << str << "\": " << v; |
119 | } |
120 | |
121 | /*! |
122 | * \brief Get the metadata section in json format. |
123 | * \return the meta data string. |
124 | */ |
125 | Doc GetMetaSection() const { |
126 | if (meta_data_.size() == 0) return Doc(); |
127 | return Doc::RawText(SaveJSON(Map<String, ObjectRef>(meta_data_.begin(), meta_data_.end()))); |
128 | } |
129 | |
130 | /*! \return whether the meta data context is empty. */ |
131 | bool empty() const { return meta_data_.empty(); } |
132 | |
133 | private: |
134 | /*! \brief additional metadata stored in TVM json format */ |
135 | std::unordered_map<String, Array<ObjectRef>> meta_data_; |
136 | /*! \brief map from meta data into its string representation */ |
137 | std::unordered_map<ObjectRef, Doc, ObjectPtrHash, ObjectPtrEqual> meta_repr_; |
138 | }; |
139 | } // namespace relay |
140 | } // namespace tvm |
141 | #endif // TVM_RELAY_PRINTER_META_DATA_H_ |
142 | |