ir_docsifier.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_

#include <tvm/ir/module.h>
#include <tvm/node/node.h>
#include <tvm/script/printer/doc.h>
#include <tvm/script/printer/ir_docsifier_functor.h>

#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

namespace tvm {
namespace script {
namespace printer {


class IRDocsifier;
class IRDocsifierNode;

class FrameNode : public Object {
 public:
  Array<StmtDoc> stmts;
  IRDocsifierNode* d;
  std::vector<std::function<void()>> callbacks;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("stmts", &stmts);
    // `d` is not visited
    // `callbacks` is not visited
  }

  static constexpr const char* _type_key = "script.printer.Frame";
  TVM_DECLARE_BASE_OBJECT_INFO(FrameNode, Object);

 public:
  virtual ~FrameNode() = default;

  template <typename TCallback>
  void AddExitCallback(TCallback&& cb) {
    callbacks.emplace_back(std::forward<TCallback>(cb));
  }
  void AddDispatchToken(const IRDocsifier& d, const String& token);
  virtual void EnterWithScope();
  virtual void ExitWithScope();
};

class Frame : public ObjectRef {
 protected:
  Frame() = default;

 public:
  virtual ~Frame() = default;

  void EnterWithScope() { get()->EnterWithScope(); }

  void ExitWithScope() { get()->ExitWithScope(); }

  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Frame, ObjectRef, FrameNode);
};


class IRDocsifierNode : public Object {
 public:
  using DocCreator = std::function<ExprDoc()>;
  struct VariableInfo {
    DocCreator creator;
    Optional<String> name;
  };
  PrinterConfig cfg{nullptr};
  Array<Frame> frames;
  Array<String> dispatch_tokens;
  std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual> obj2info;
  std::unordered_map<String, Array<ObjectRef>> metadata;
  std::unordered_set<String> defined_names;
  std::unordered_map<const Object*, std::vector<const Object*>> common_prefix;
  std::unordered_set<std::string> ir_usage;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("frames", &frames);
    v->Visit("dispatch_tokens", &dispatch_tokens);
    // `obj2info` is not visited
    // `metadata` is not visited
    // `defined_names` is not visited
    // `common_prefix` is not visited
    // `ir_usage` is not visited
  }

  static constexpr const char* _type_key = "script.printer.IRDocsifier";
  TVM_DECLARE_FINAL_OBJECT_INFO(IRDocsifierNode, Object);

 public:
  IdDoc Define(const ObjectRef& obj, const Frame& frame, const String& name_hint);

  void Define(const ObjectRef& obj, const Frame& frame, DocCreator doc_factory);

  Optional<ExprDoc> GetVarDoc(const ObjectRef& obj) const;
  ExprDoc AddMetadata(const ObjectRef& obj);
  bool IsVarDefined(const ObjectRef& obj) const;
  void RemoveVar(const ObjectRef& obj);
  void SetCommonPrefix(const ObjectRef& root, runtime::TypedPackedFunc<bool(ObjectRef)> is_var);
  template <class TDoc = Doc>
  inline TDoc AsDoc(const ObjectRef& obj, const ObjectPath& path) const;
};

class IRDocsifier : public ObjectRef {
 public:
  using FType = IRDocsifierFunctor<printer::Doc, ObjectPath, IRDocsifier>;
  explicit IRDocsifier(const PrinterConfig& cfg);
  TVM_DLL static FType& vtable();

  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRDocsifier, ObjectRef, IRDocsifierNode);
};


inline void FrameNode::EnterWithScope() {
  if (d != nullptr) {
    d->frames.push_back(GetRef<Frame>(this));
  }
}

inline void FrameNode::ExitWithScope() {
  for (const std::function<void()>& callback : callbacks) {
    callback();
  }
  callbacks.clear();
  if (d != nullptr) {
    d->frames.pop_back();
  }
}

template <class TDoc>
inline static void AddDocDecoration(const Doc& d, const ObjectRef& obj, const ObjectPath& path,
                                    const PrinterConfig& cfg) {
  if (cfg->obj_to_annotate.count(obj)) {
    if (const auto* stmt = d.as<StmtDocNode>()) {
      if (stmt->comment.defined()) {
        stmt->comment = stmt->comment.value() + "\n" + cfg->obj_to_annotate.at(obj);
      } else {
        stmt->comment = cfg->obj_to_annotate.at(obj);
      }
    } else {
      LOG(WARNING) << "Expect StmtDoc to be annotated for object " << obj << ", but got "
                   << Downcast<TDoc>(d)->_type_key;
    }
  }
  for (const ObjectRef& o : cfg->obj_to_underline) {
    if (o.same_as(obj)) {
      cfg->path_to_underline.push_back(path);
    }
  }
  for (const auto& pair : cfg->path_to_annotate) {
    ObjectPath p = pair.first;
    String attn = pair.second;
    if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
      if (const auto* stmt = d.as<StmtDocNode>()) {
        if (stmt->comment.defined()) {
          stmt->comment = stmt->comment.value() + "\n" + attn;
        } else {
          stmt->comment = attn;
        }
      } else {
        LOG(WARNING) << "Expect StmtDoc to be annotated at object path " << p << ", but got "
                     << Downcast<TDoc>(d)->_type_key;
      }
    }
  }
}

template <class TDoc>
inline TDoc IRDocsifierNode::AsDoc(const ObjectRef& obj, const ObjectPath& path) const {
  if (obj.defined()) {
    Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path, GetRef<IRDocsifier>(this));
    d->source_paths.push_back(path);
    AddDocDecoration<TDoc>(d, obj, path, cfg);
    return Downcast<TDoc>(d);
  }
  return Downcast<TDoc>(LiteralDoc::None(path));
}

inline void FrameNode::AddDispatchToken(const IRDocsifier& d, const String& token) {
  d->dispatch_tokens.push_back(token);
  this->AddExitCallback([doc = d.get()]() { doc->dispatch_tokens.pop_back(); });
}

}  // namespace printer
}  // namespace script
}  // namespace tvm

#endif  // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_