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 | |
20 | /*! |
21 | * \file tvm/printer/meta_data.h |
22 | * \brief Meta data context for printers. |
23 | */ |
24 | #ifndef TVM_PRINTER_META_DATA_H_ |
25 | #define TVM_PRINTER_META_DATA_H_ |
26 | |
27 | #include <tvm/node/serialization.h> |
28 | |
29 | #include <string> |
30 | #include <unordered_map> |
31 | |
32 | #include "doc.h" |
33 | |
34 | namespace tvm { |
35 | /*! |
36 | * \brief Meta data context for Printers |
37 | * |
38 | * This is an important part to enable bi-directional serializability. |
39 | * We use tvm's Node system to build the current IR. |
40 | * It can be hard to design a text format for all the possible nodes |
41 | * as the set of nodes can grow when we do more extensions. |
42 | * |
43 | * Instead of trying to design readable text format for every node, |
44 | * we support a meta data section in the text format. |
45 | * We allow the text format to refer to a node in the meta data section. |
46 | * |
47 | * The meta data section is a json serialized string of an Map<string, Array<NodeRef>>. |
48 | * Each element in the meta data section can be referenced by the text format. |
49 | * Each meta data node is printed in the following format. |
50 | * |
51 | * meta[type-key-of-node>][<index-in-meta-section>] |
52 | * |
53 | * Specifically, consider the following IR(constructed by python). |
54 | * |
55 | * \code |
56 | * |
57 | * n = tvm.var("n") |
58 | * x = tvm.relay.var("x", shape=(n, 1)) |
59 | * f = tvm.relay.Function([x], x) |
60 | * print(f.astext()) |
61 | * |
62 | * \endcode |
63 | * |
64 | * The corresponding text format is shown in the following code block. |
65 | * |
66 | * \code |
67 | * |
68 | * fn (%x: Tensor[(meta[Variable][0],), float32]) { |
69 | * %x |
70 | * } |
71 | * # Meta data section is a json-serialized string |
72 | * # of the following array. |
73 | * # [tvm.var("n")] |
74 | * |
75 | * \endcode |
76 | * |
77 | * Note that we store tvm.var("n") in the meta data section. |
78 | * Since it is stored in the index-0 in the meta data section, |
79 | * we print it as meta[Variable][0]. |
80 | * |
81 | * The text parser can recover this object by loading from the corresponding |
82 | * location in the meta data section. |
83 | * |
84 | * This is a design trade-off. |
85 | * It allows us to embedded any meta data in the text format, |
86 | * while still being able to tweak the text part of the printed IR easily. |
87 | */ |
88 | class TextMetaDataContext { |
89 | public: |
90 | /*! |
91 | * \brief Get text representation of meta node. |
92 | * \param node The node to be converted to meta node. |
93 | * \return A string representation of the meta node. |
94 | */ |
95 | Doc GetMetaNode(const ObjectRef& node) { |
96 | auto it = meta_repr_.find(node); |
97 | if (it != meta_repr_.end()) { |
98 | return it->second; |
99 | } |
100 | std::string type_key = node->GetTypeKey(); |
101 | ICHECK(!type_key.empty()); |
102 | Array<ObjectRef>& mvector = meta_data_[type_key]; |
103 | int64_t index = static_cast<int64_t>(mvector.size()); |
104 | mvector.push_back(node); |
105 | Doc doc; |
106 | doc << "meta[" << type_key << "][" << index << "]" ; |
107 | meta_repr_[node] = doc; |
108 | return meta_repr_[node]; |
109 | } |
110 | |
111 | /*! |
112 | * \brief Test whether a node has been put in meta |
113 | * \param node The query node |
114 | * \return whether the node has been put in meta |
115 | */ |
116 | bool InMeta(const ObjectRef& node) { return meta_repr_.find(node) != meta_repr_.end(); } |
117 | |
118 | /*! |
119 | * \brief Print a key value pair |
120 | */ |
121 | Doc PrintKeyValue(const std::string& str, const Doc& v) const { |
122 | return Doc() << "\"" << str << "\": " << v; |
123 | } |
124 | |
125 | /*! |
126 | * \brief Get the metadata section in json format. |
127 | * \return the meta data string. |
128 | */ |
129 | Doc GetMetaSection() const { |
130 | if (meta_data_.size() == 0) return Doc(); |
131 | return Doc::RawText(SaveJSON(Map<String, ObjectRef>(meta_data_.begin(), meta_data_.end()))); |
132 | } |
133 | |
134 | /*! \return whether the meta data context is empty. */ |
135 | bool empty() const { return meta_data_.empty(); } |
136 | |
137 | private: |
138 | /*! \brief additional metadata stored in TVM json format */ |
139 | std::unordered_map<String, Array<ObjectRef>> meta_data_; |
140 | /*! \brief map from meta data into its string representation */ |
141 | std::unordered_map<ObjectRef, Doc, ObjectPtrHash, ObjectPtrEqual> meta_repr_; |
142 | }; |
143 | } // namespace tvm |
144 | #endif // TVM_PRINTER_META_DATA_H_ |
145 | |