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 * \file source_map.h
21 * \brief A map from source names to source code.
22 */
23#ifndef TVM_IR_SOURCE_MAP_H_
24#define TVM_IR_SOURCE_MAP_H_
25
26#include <tvm/node/node.h>
27#include <tvm/runtime/object.h>
28#include <tvm/runtime/packed_func.h>
29#include <tvm/runtime/registry.h>
30
31#include <fstream>
32#include <string>
33#include <utility>
34#include <vector>
35
36namespace tvm {
37
38/*!
39 * \brief The source name in the Span
40 * \sa SourceNameNode, Span
41 */
42class SourceName;
43/*!
44 * \brief The name of a source fragment.
45 */
46class SourceNameNode : public Object {
47 public:
48 /*! \brief The source name. */
49 String name;
50 // override attr visitor
51 void VisitAttrs(AttrVisitor* v) { v->Visit("name", &name); }
52
53 static constexpr bool _type_has_method_sequal_reduce = true;
54
55 bool SEqualReduce(const SourceNameNode* other, SEqualReducer equal) const {
56 return equal(name, other->name);
57 }
58
59 static constexpr const char* _type_key = "SourceName";
60 TVM_DECLARE_FINAL_OBJECT_INFO(SourceNameNode, Object);
61};
62
63/*!
64 * \brief The source name of a file span.
65 * \sa SourceNameNode, Span
66 */
67class SourceName : public ObjectRef {
68 public:
69 /*!
70 * \brief Get an SourceName for a given operator name.
71 * Will raise an error if the source name has not been registered.
72 * \param name Name of the operator.
73 * \return SourceName valid throughout program lifetime.
74 */
75 TVM_DLL static SourceName Get(const String& name);
76
77 TVM_DEFINE_OBJECT_REF_METHODS(SourceName, ObjectRef, SourceNameNode);
78};
79
80/*!
81 * \brief Span information for debugging purposes
82 */
83class Span;
84/*!
85 * \brief Stores locations in frontend source that generated a node.
86 */
87class SpanNode : public Object {
88 public:
89 /*! \brief The source name. */
90 SourceName source_name;
91 /*! \brief The line number. */
92 int line;
93 /*! \brief The column offset. */
94 int column;
95 /*! \brief The end line number. */
96 int end_line;
97 /*! \brief The end column number. */
98 int end_column;
99
100 // override attr visitor
101 void VisitAttrs(AttrVisitor* v) {
102 v->Visit("source_name", &source_name);
103 v->Visit("line", &line);
104 v->Visit("column", &column);
105 v->Visit("end_line", &end_line);
106 v->Visit("end_column", &end_column);
107 }
108 static constexpr bool _type_has_method_sequal_reduce = true;
109
110 bool SEqualReduce(const SpanNode* other, SEqualReducer equal) const {
111 return equal(source_name, other->source_name) && equal(line, other->line) &&
112 equal(column, other->column) && equal(end_line, other->end_line) &&
113 equal(end_column, other->end_column);
114 }
115
116 static constexpr const char* _type_key = "Span";
117 TVM_DECLARE_FINAL_OBJECT_INFO(SpanNode, Object);
118};
119
120class Span : public ObjectRef {
121 public:
122 TVM_DLL Span(SourceName source_name, int line, int end_line, int column, int end_column);
123
124 /*! \brief Merge two spans into one which captures the combined regions. */
125 TVM_DLL Span Merge(const Span& other) const;
126
127 TVM_DEFINE_OBJECT_REF_METHODS(Span, ObjectRef, SpanNode);
128};
129
130/*! \brief A program source in any language.
131 *
132 * Could represent the source from an ML framework or a source
133 * representing a tvm::IRModule.
134 */
135class Source;
136
137class SourceNode : public Object {
138 public:
139 /*! \brief The source name. */
140 SourceName source_name;
141
142 /*! \brief The raw source. */
143 String source;
144
145 /*! \brief A mapping of line breaks into the raw source. */
146 std::vector<std::pair<int, int>> line_map;
147
148 // override attr visitor
149 void VisitAttrs(AttrVisitor* v) {
150 v->Visit("source_name", &source_name);
151 v->Visit("source", &source);
152 }
153
154 static constexpr const char* _type_key = "Source";
155 TVM_DECLARE_FINAL_OBJECT_INFO(SourceNode, Object);
156};
157
158class Source : public ObjectRef {
159 public:
160 TVM_DLL Source(SourceName src_name, std::string source);
161 TVM_DLL tvm::String GetLine(int line);
162
163 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Source, ObjectRef, SourceNode);
164};
165
166/*!
167 * \brief A mapping from a unique source name to source fragment.
168 */
169class SourceMap;
170/*!
171 * \brief Stores locations in frontend source that generated a node.
172 */
173class SourceMapNode : public Object {
174 public:
175 /*! \brief The source mapping. */
176 Map<SourceName, Source> source_map;
177
178 // override attr visitor
179 void VisitAttrs(AttrVisitor* v) { v->Visit("source_map", &source_map); }
180
181 bool SEqualReduce(const SourceMapNode* other, SEqualReducer equal) const {
182 return equal(source_map, other->source_map);
183 }
184
185 static constexpr const char* _type_key = "SourceMap";
186 TVM_DECLARE_FINAL_OBJECT_INFO(SourceMapNode, Object);
187};
188
189class SourceMap : public ObjectRef {
190 public:
191 explicit SourceMap(Map<SourceName, Source> source_map);
192
193 explicit SourceMap(std::initializer_list<std::pair<SourceName, Source>> source_map)
194 : SourceMap(Map<SourceName, Source>(source_map)) {}
195
196 SourceMap() : SourceMap(Map<SourceName, Source>()) {}
197
198 void Add(const Source& source);
199
200 SourceMapNode* operator->() {
201 ICHECK(get() != nullptr);
202 return static_cast<SourceMapNode*>(get_mutable());
203 }
204
205 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SourceMap, ObjectRef, SourceMapNode);
206};
207
208} // namespace tvm
209
210#endif // TVM_IR_SOURCE_MAP_H_
211