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/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/registry.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 {
 
using AccessPath = ffi::reflection::AccessPath;
 
 
class IRDocsifier;
class IRDocsifierNode;
 
class FrameNode : public Object {
 public:
  ffi::Array<StmtDoc> stmts;
  IRDocsifierNode* d;
  std::vector<std::function<void()>> callbacks;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FrameNode>().def_ro("stmts", &FrameNode::stmts);
  }
 
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO("script.printer.Frame", 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 ffi::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_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Frame, ObjectRef, FrameNode);
};
 
 
class IRDocsifierNode : public Object {
 public:
  using DocCreator = std::function<ExprDoc()>;
  struct VariableInfo {
    DocCreator creator;
    ffi::Optional<ffi::String> name;
  };
  PrinterConfig cfg{ffi::UnsafeInit()};
  ffi::Array<Frame> frames;
  ffi::Array<ffi::String> dispatch_tokens;
  std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual> obj2info;
  std::unordered_map<ffi::String, ffi::Array<ffi::Any>> metadata;
  std::unordered_map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
  std::unordered_set<ffi::String> defined_names;
  std::unordered_map<const Object*, std::vector<const Object*>> common_prefix;
  std::unordered_set<std::string> ir_usage;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IRDocsifierNode>()
        .def_ro("frames", &IRDocsifierNode::frames)
        .def_ro("dispatch_tokens", &IRDocsifierNode::dispatch_tokens);
  }
 
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IRDocsifier", IRDocsifierNode, Object);
 
 public:
  IdDoc Define(const ObjectRef& obj, const Frame& frame, const ffi::String& name_hint);
 
  void Define(const ObjectRef& obj, const Frame& frame, DocCreator doc_factory);
 
  ffi::Optional<ExprDoc> GetVarDoc(const ObjectRef& obj) const;
  ExprDoc AddMetadata(const ffi::Any& obj);
  void AddGlobalInfo(const ffi::String& name, const GlobalInfo& ginfo);
  bool IsVarDefined(const ObjectRef& obj) const;
  void RemoveVar(const ObjectRef& obj);
  void SetCommonPrefix(const ObjectRef& root, ffi::TypedFunction<bool(ObjectRef)> is_var);
  template <class TDoc = Doc>
  inline TDoc AsDoc(const Any& obj, const AccessPath& path) const;
};
 
class IRDocsifier : public ObjectRef {
 public:
  using FType = IRDocsifierFunctor<printer::Doc, AccessPath, IRDocsifier>;
  explicit IRDocsifier(const PrinterConfig& cfg);
  TVM_DLL static FType& vtable();
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRDocsifier, ObjectRef, IRDocsifierNode);
};
 
 
inline void FrameNode::EnterWithScope() {
  if (d != nullptr) {
    d->frames.push_back(ffi::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 AccessPath& path,
                                    const PrinterConfig& cfg) {
  if (cfg->obj_to_annotate.count(obj)) {
    if (const auto* stmt = d.as<StmtDocNode>()) {
      if (stmt->comment.has_value()) {
        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) {
    AccessPath p = pair.first;
    ffi::String attn = pair.second;
    if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
      if (const auto* stmt = d.as<StmtDocNode>()) {
        if (stmt->comment.has_value()) {
          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 Any& value, const AccessPath& path) const {
  switch (value.type_index()) {
    case ffi::TypeIndex::kTVMFFINone:
      return Downcast<TDoc>(LiteralDoc::None(path));
    case ffi::TypeIndex::kTVMFFIBool:
      return Downcast<TDoc>(LiteralDoc::Boolean(value.as<bool>().value(), path));
    case ffi::TypeIndex::kTVMFFIInt:
      return Downcast<TDoc>(LiteralDoc::Int(value.as<int64_t>().value(), path));
    case ffi::TypeIndex::kTVMFFIFloat:
      return Downcast<TDoc>(LiteralDoc::Float(value.as<double>().value(), path));
    case ffi::TypeIndex::kTVMFFISmallStr:
    case ffi::TypeIndex::kTVMFFIStr: {
      std::string string_value = value.cast<std::string>();
      bool has_multiple_lines = string_value.find_first_of('\n') != std::string::npos;
      if (has_multiple_lines) {
        Doc d = const_cast<IRDocsifierNode*>(this)->AddMetadata(string_value);
        // TODO(tqchen): cross check AddDocDecoration
        return Downcast<TDoc>(d);
      }
      return Downcast<TDoc>(LiteralDoc::Str(string_value, path));
    }
    case ffi::TypeIndex::kTVMFFIDataType:
      return Downcast<TDoc>(LiteralDoc::DataType(value.as<runtime::DataType>().value(), path));
    case ffi::TypeIndex::kTVMFFIDevice:
      return Downcast<TDoc>(LiteralDoc::Device(value.as<DLDevice>().value(), path));
    default: {
      if (auto opt_obj = value.as<ObjectRef>()) {
        ObjectRef obj = opt_obj.value();
        Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path,
                                      ffi::GetRef<IRDocsifier>(this));
        d->source_paths.push_back(path);
        AddDocDecoration<TDoc>(d, obj, path, cfg);
        return Downcast<TDoc>(d);
      } else {
        LOG(FATAL) << "TypeError: Cannot handle Any type: `" << value.GetTypeKey() << "`";
        TVM_FFI_UNREACHABLE();
      }
    }
  }
}
 
inline void FrameNode::AddDispatchToken(const IRDocsifier& d, const ffi::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_