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
34namespace 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 */
88class 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