tvm
ir_docsifier.h
Go to the documentation of this file.
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_SCRIPT_PRINTER_IR_DOCSIFIER_H_
20 #define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
21 
22 #include <tvm/ffi/reflection/access_path.h>
23 #include <tvm/ffi/reflection/registry.h>
24 #include <tvm/ir/module.h>
25 #include <tvm/node/node.h>
26 #include <tvm/script/printer/doc.h>
28 
29 #include <string>
30 #include <unordered_map>
31 #include <unordered_set>
32 #include <utility>
33 #include <vector>
34 
35 namespace tvm {
36 namespace script {
37 namespace printer {
38 
40 
42 
43 class IRDocsifier;
44 class IRDocsifierNode;
45 
50 class FrameNode : public Object {
51  public:
53  Array<StmtDoc> stmts;
57  std::vector<std::function<void()>> callbacks;
58 
59  static void RegisterReflection() {
60  namespace refl = tvm::ffi::reflection;
61  refl::ObjectDef<FrameNode>().def_ro("stmts", &FrameNode::stmts);
62  }
63 
64  static constexpr const char* _type_key = "script.printer.Frame";
65 
67 
68  public:
69  virtual ~FrameNode() = default;
70 
75  template <typename TCallback>
76  void AddExitCallback(TCallback&& cb) {
77  callbacks.emplace_back(std::forward<TCallback>(cb));
78  }
85  void AddDispatchToken(const IRDocsifier& d, const String& token);
89  virtual void EnterWithScope();
93  virtual void ExitWithScope();
94 };
95 
99 class Frame : public ObjectRef {
100  protected:
101  Frame() = default;
102 
103  public:
104  virtual ~Frame() = default;
105 
107  void EnterWithScope() { get()->EnterWithScope(); }
108 
110  void ExitWithScope() { get()->ExitWithScope(); }
111 
113 };
114 
116 
123 class IRDocsifierNode : public Object {
124  public:
126  using DocCreator = std::function<ExprDoc()>;
128  struct VariableInfo {
132  Optional<String> name;
133  };
135  PrinterConfig cfg{nullptr};
140  Array<Frame> frames;
147  Array<String> dispatch_tokens;
149  std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual> obj2info;
151  std::unordered_map<String, Array<ffi::Any>> metadata;
153  std::unordered_map<String, Array<GlobalInfo>> global_infos;
155  std::unordered_set<String> defined_names;
157  std::unordered_map<const Object*, std::vector<const Object*>> common_prefix;
159  std::unordered_set<std::string> ir_usage;
160 
161  static void RegisterReflection() {
162  namespace refl = tvm::ffi::reflection;
163  refl::ObjectDef<IRDocsifierNode>()
164  .def_ro("frames", &IRDocsifierNode::frames)
165  .def_ro("dispatch_tokens", &IRDocsifierNode::dispatch_tokens);
166  }
167 
168  static constexpr const char* _type_key = "script.printer.IRDocsifier";
169 
171 
172  public:
184  IdDoc Define(const ObjectRef& obj, const Frame& frame, const String& name_hint);
185 
202  void Define(const ObjectRef& obj, const Frame& frame, DocCreator doc_factory);
203 
210  Optional<ExprDoc> GetVarDoc(const ObjectRef& obj) const;
212  ExprDoc AddMetadata(const ffi::Any& obj);
217  void AddGlobalInfo(const String& name, const GlobalInfo& ginfo);
224  bool IsVarDefined(const ObjectRef& obj) const;
226  void RemoveVar(const ObjectRef& obj);
232  void SetCommonPrefix(const ObjectRef& root, ffi::TypedFunction<bool(ObjectRef)> is_var);
240  template <class TDoc = Doc>
241  inline TDoc AsDoc(const Any& obj, const AccessPath& path) const;
242 };
243 
247 class IRDocsifier : public ObjectRef {
248  public:
251  explicit IRDocsifier(const PrinterConfig& cfg);
253  TVM_DLL static FType& vtable();
254 
256 };
257 
259 
261  if (d != nullptr) {
262  d->frames.push_back(GetRef<Frame>(this));
263  }
264 }
265 
267  for (const std::function<void()>& callback : callbacks) {
268  callback();
269  }
270  callbacks.clear();
271  if (d != nullptr) {
272  d->frames.pop_back();
273  }
274 }
275 
276 template <class TDoc>
277 inline static void AddDocDecoration(const Doc& d, const ObjectRef& obj, const AccessPath& path,
278  const PrinterConfig& cfg) {
279  if (cfg->obj_to_annotate.count(obj)) {
280  if (const auto* stmt = d.as<StmtDocNode>()) {
281  if (stmt->comment.has_value()) {
282  stmt->comment = stmt->comment.value() + "\n" + cfg->obj_to_annotate.at(obj);
283  } else {
284  stmt->comment = cfg->obj_to_annotate.at(obj);
285  }
286  } else {
287  LOG(WARNING) << "Expect StmtDoc to be annotated for object " << obj << ", but got "
288  << Downcast<TDoc>(d)->_type_key;
289  }
290  }
291  for (const ObjectRef& o : cfg->obj_to_underline) {
292  if (o.same_as(obj)) {
293  cfg->path_to_underline.push_back(path);
294  }
295  }
296  for (const auto& pair : cfg->path_to_annotate) {
297  AccessPath p = pair.first;
298  String attn = pair.second;
299  if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
300  if (const auto* stmt = d.as<StmtDocNode>()) {
301  if (stmt->comment.has_value()) {
302  stmt->comment = stmt->comment.value() + "\n" + attn;
303  } else {
304  stmt->comment = attn;
305  }
306  } else {
307  LOG(WARNING) << "Expect StmtDoc to be annotated at object path " << p << ", but got "
308  << Downcast<TDoc>(d)->_type_key;
309  }
310  }
311  }
312 }
313 
314 template <class TDoc>
315 inline TDoc IRDocsifierNode::AsDoc(const Any& value, const AccessPath& path) const {
316  switch (value.type_index()) {
317  case ffi::TypeIndex::kTVMFFINone:
318  return Downcast<TDoc>(LiteralDoc::None(path));
319  case ffi::TypeIndex::kTVMFFIBool:
320  return Downcast<TDoc>(LiteralDoc::Boolean(value.as<bool>().value(), path));
321  case ffi::TypeIndex::kTVMFFIInt:
322  return Downcast<TDoc>(LiteralDoc::Int(value.as<int64_t>().value(), path));
323  case ffi::TypeIndex::kTVMFFIFloat:
324  return Downcast<TDoc>(LiteralDoc::Float(value.as<double>().value(), path));
325  case ffi::TypeIndex::kTVMFFISmallStr:
326  case ffi::TypeIndex::kTVMFFIStr: {
327  std::string string_value = value.cast<std::string>();
328  bool has_multiple_lines = string_value.find_first_of('\n') != std::string::npos;
329  if (has_multiple_lines) {
330  Doc d = const_cast<IRDocsifierNode*>(this)->AddMetadata(string_value);
331  // TODO(tqchen): cross check AddDocDecoration
332  return Downcast<TDoc>(d);
333  }
334  return Downcast<TDoc>(LiteralDoc::Str(string_value, path));
335  }
336  case ffi::TypeIndex::kTVMFFIDataType:
337  return Downcast<TDoc>(LiteralDoc::DataType(value.as<runtime::DataType>().value(), path));
338  case ffi::TypeIndex::kTVMFFIDevice:
339  return Downcast<TDoc>(LiteralDoc::Device(value.as<DLDevice>().value(), path));
340  default: {
341  if (auto opt_obj = value.as<ObjectRef>()) {
342  ObjectRef obj = opt_obj.value();
343  Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path, GetRef<IRDocsifier>(this));
344  d->source_paths.push_back(path);
345  AddDocDecoration<TDoc>(d, obj, path, cfg);
346  return Downcast<TDoc>(d);
347  } else {
348  LOG(FATAL) << "TypeError: Cannot handle Any type: `" << value.GetTypeKey() << "`";
349  TVM_FFI_UNREACHABLE();
350  }
351  }
352  }
353 }
354 
355 inline void FrameNode::AddDispatchToken(const IRDocsifier& d, const String& token) {
356  d->dispatch_tokens.push_back(token);
357  this->AddExitCallback([doc = d.get()]() { doc->dispatch_tokens.pop_back(); });
358 }
359 
360 } // namespace printer
361 } // namespace script
362 } // namespace tvm
363 
364 #endif // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
Managed reference to GlobalInfoNode.
Definition: global_info.h:56
Definition: script_printer.h:155
Runtime primitive data type.
Definition: data_type.h:47
Reference type of DocNode.
Definition: doc.h:88
Reference type of ExprDocNode.
Definition: doc.h:148
Definition: ir_docsifier.h:50
static constexpr const char * _type_key
Definition: ir_docsifier.h:64
virtual void EnterWithScope()
Method that's called when Frame enters the scope.
Definition: ir_docsifier.h:260
static void RegisterReflection()
Definition: ir_docsifier.h:59
IRDocsifierNode * d
Definition: ir_docsifier.h:55
TVM_DECLARE_BASE_OBJECT_INFO(FrameNode, Object)
virtual void ExitWithScope()
Method that's called when Frame exits the scope.
Definition: ir_docsifier.h:266
void AddExitCallback(TCallback &&cb)
Add a callback function to be called when this frame exits.
Definition: ir_docsifier.h:76
Array< StmtDoc > stmts
Definition: ir_docsifier.h:53
std::vector< std::function< void()> > callbacks
Definition: ir_docsifier.h:57
void AddDispatchToken(const IRDocsifier &d, const String &token)
Add a dispatch token to the docsifier, and a callback that pops the token when this frame exits.
Definition: ir_docsifier.h:355
Reference type of FrameNode.
Definition: ir_docsifier.h:99
void EnterWithScope()
Method that's called when Frame enters the scope.
Definition: ir_docsifier.h:107
virtual ~Frame()=default
void ExitWithScope()
Method that's called when Frame exits the scope.
Definition: ir_docsifier.h:110
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Frame, ObjectRef, FrameNode)
Dynamic dispatch functor based on AccessPath.
Definition: ir_docsifier_functor.h:43
IRDocsifier is the top-level interface in the IR->Doc process.
Definition: ir_docsifier.h:123
void Define(const ObjectRef &obj, const Frame &frame, DocCreator doc_factory)
Define variable by doc factory.
std::unordered_set< std::string > ir_usage
The IR usages for headers printing.
Definition: ir_docsifier.h:159
std::function< ExprDoc()> DocCreator
A function that creates the doc for a variable.
Definition: ir_docsifier.h:126
TVM_DECLARE_FINAL_OBJECT_INFO(IRDocsifierNode, Object)
std::unordered_map< String, Array< ffi::Any > > metadata
Metadata printing.
Definition: ir_docsifier.h:151
std::unordered_map< String, Array< GlobalInfo > > global_infos
GlobalInfo printing.
Definition: ir_docsifier.h:153
Optional< ExprDoc > GetVarDoc(const ObjectRef &obj) const
Get the doc for variable.
ExprDoc AddMetadata(const ffi::Any &obj)
Add a TVM object to the metadata section.
void SetCommonPrefix(const ObjectRef &root, ffi::TypedFunction< bool(ObjectRef)> is_var)
Set the common prefix information of variable usage.
Array< String > dispatch_tokens
The stack of dispatch tokens.
Definition: ir_docsifier.h:147
void AddGlobalInfo(const String &name, const GlobalInfo &ginfo)
Add a GlobalInfo to the global_infos map.
std::unordered_set< String > defined_names
The variable names used already.
Definition: ir_docsifier.h:155
std::unordered_map< ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual > obj2info
Mapping from a var to its info.
Definition: ir_docsifier.h:149
TDoc AsDoc(const Any &obj, const AccessPath &path) const
Transform the input object into TDoc.
Definition: ir_docsifier.h:315
IdDoc Define(const ObjectRef &obj, const Frame &frame, const String &name_hint)
Define variable by name.
PrinterConfig cfg
The configuration of the printer.
Definition: ir_docsifier.h:135
bool IsVarDefined(const ObjectRef &obj) const
Check if a variable exists in the table.
Array< Frame > frames
The stack of frames.
Definition: ir_docsifier.h:140
static constexpr const char * _type_key
Definition: ir_docsifier.h:168
std::unordered_map< const Object *, std::vector< const Object * > > common_prefix
Common prefixes of variable usages.
Definition: ir_docsifier.h:157
static void RegisterReflection()
Definition: ir_docsifier.h:161
void RemoveVar(const ObjectRef &obj)
Remove the variable defined.
Reference type of IRDocsifierNode.
Definition: ir_docsifier.h:247
static FType & vtable()
The registration table for IRDocsifier.
IRDocsifier(const PrinterConfig &cfg)
Create a IRDocsifier.
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRDocsifier, ObjectRef, IRDocsifierNode)
Reference type of IdDocNode.
Definition: doc.h:358
static LiteralDoc DataType(const runtime::DataType &v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent string.
Definition: doc.h:315
static LiteralDoc Boolean(bool v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent boolean.
Definition: doc.h:293
static LiteralDoc Int(int64_t v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent integer.
Definition: doc.h:285
static LiteralDoc None(const Optional< AccessPath > &p)
Create a LiteralDoc to represent None/null/empty value.
Definition: doc.h:279
static LiteralDoc Float(double v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent float.
Definition: doc.h:301
static LiteralDoc Device(const DLDevice &v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent device.
Definition: doc.h:324
static LiteralDoc Str(const String &v, const Optional< AccessPath > &p)
Create a LiteralDoc to represent string.
Definition: doc.h:309
The base class of statement doc.
Definition: doc.h:167
IRModule that holds the functions and type definitions.
Definition: repr_printer.h:91
ffi::reflection::AccessPath AccessPath
Definition: doc.h:35
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Definitions and helper macros for IR/AST nodes.
Information about a variable, including its optional name and its doc creator.
Definition: ir_docsifier.h:128
DocCreator creator
The creator.
Definition: ir_docsifier.h:130
Optional< String > name
The name of the variable.
Definition: ir_docsifier.h:132