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
29namespace tvm {
30namespace 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 */
84class 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