19 #ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
20 #define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
22 #include <tvm/ffi/reflection/access_path.h>
23 #include <tvm/ffi/reflection/registry.h>
30 #include <unordered_map>
31 #include <unordered_set>
44 class IRDocsifierNode;
64 static constexpr
const char*
_type_key =
"script.printer.Frame";
75 template <
typename TCallback>
77 callbacks.emplace_back(std::forward<TCallback>(cb));
99 class Frame :
public ObjectRef {
149 std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual>
obj2info;
151 std::unordered_map<String, Array<ffi::Any>>
metadata;
157 std::unordered_map<const Object*, std::vector<const Object*>>
common_prefix;
163 refl::ObjectDef<IRDocsifierNode>()
168 static constexpr
const char*
_type_key =
"script.printer.IRDocsifier";
210 Optional<ExprDoc>
GetVarDoc(
const ObjectRef& obj)
const;
232 void SetCommonPrefix(
const ObjectRef& root, ffi::TypedFunction<
bool(ObjectRef)> is_var);
240 template <
class TDoc = Doc>
262 d->
frames.push_back(GetRef<Frame>(
this));
267 for (
const std::function<
void()>& callback :
callbacks) {
276 template <
class TDoc>
277 inline static void AddDocDecoration(
const Doc& d,
const ObjectRef& obj,
const AccessPath& path,
279 if (cfg->obj_to_annotate.count(obj)) {
281 if (stmt->comment.has_value()) {
282 stmt->comment = stmt->comment.value() +
"\n" + cfg->obj_to_annotate.at(obj);
284 stmt->comment = cfg->obj_to_annotate.at(obj);
287 LOG(WARNING) <<
"Expect StmtDoc to be annotated for object " << obj <<
", but got "
288 << Downcast<TDoc>(d)->_type_key;
291 for (
const ObjectRef& o : cfg->obj_to_underline) {
292 if (o.same_as(obj)) {
293 cfg->path_to_underline.push_back(path);
296 for (
const auto& pair : cfg->path_to_annotate) {
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;
304 stmt->comment = attn;
307 LOG(WARNING) <<
"Expect StmtDoc to be annotated at object path " << p <<
", but got "
308 << Downcast<TDoc>(d)->_type_key;
314 template <
class TDoc>
316 switch (value.type_index()) {
317 case ffi::TypeIndex::kTVMFFINone:
319 case ffi::TypeIndex::kTVMFFIBool:
321 case ffi::TypeIndex::kTVMFFIInt:
322 return Downcast<TDoc>(
LiteralDoc::Int(value.as<int64_t>().value(), path));
323 case ffi::TypeIndex::kTVMFFIFloat:
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) {
332 return Downcast<TDoc>(d);
336 case ffi::TypeIndex::kTVMFFIDataType:
338 case ffi::TypeIndex::kTVMFFIDevice:
341 if (
auto opt_obj = value.as<ObjectRef>()) {
342 ObjectRef obj = opt_obj.value();
344 d->source_paths.push_back(path);
345 AddDocDecoration<TDoc>(d, obj, path,
cfg);
346 return Downcast<TDoc>(d);
348 LOG(FATAL) <<
"TypeError: Cannot handle Any type: `" << value.GetTypeKey() <<
"`";
349 TVM_FFI_UNREACHABLE();
357 this->
AddExitCallback([doc =
d.get()]() { doc->dispatch_tokens.pop_back(); });
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
virtual ~FrameNode()=default
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
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