[TVMScript] Use default fallback for un-registered type

include/tvm/script/printer/ir_docsifier_functor.h

@@ -23,6 +23,7 @@
 #include <tvm/runtime/logging.h>
 #include <tvm/runtime/packed_func.h>
 
+#include <optional>
 #include <string>
 #include <type_traits>
 #include <unordered_map>
@@ -69,6 +70,10 @@ class IRDocsifierFunctor {
     if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
       return (*pf)(obj, args...);
     }
+    if ((pf = LookupFallback()) != nullptr) {
+      return (*pf)(obj, args...);
+    }
+
     LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
                  << runtime::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
                  << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
@@ -100,6 +105,14 @@ class IRDocsifierFunctor {
     return *this;
   }
 
+  TSelf& set_fallback(runtime::PackedFunc f) {
+    ICHECK(!dispatch_fallback_.has_value()) << "Fallback is already defined";
+    dispatch_fallback_ = f;
+    return *this;
+  }
+
+  void remove_fallback() { dispatch_fallback_ = std::nullopt; }
+
   /*!
    * \brief Set the dispatch function
    * \param token The dispatch token.
@@ -112,6 +125,13 @@ class IRDocsifierFunctor {
                         runtime::TypedPackedFunc<R(TObjectRef, Args...)>(f));
   }
 
+  template <typename TCallable,
+            typename = std::enable_if_t<IsDispatchFunction<ObjectRef, TCallable>::value>>
+  TSelf& set_fallback(TCallable f) {
+    runtime::PackedFunc func = runtime::TypedPackedFunc<R(ObjectRef, Args...)>(f);
+    return set_fallback(func);
+  }
+
   /*!
    * \brief Remove dispatch function
    * \param token The dispatch token.
@@ -151,13 +171,26 @@ class IRDocsifierFunctor {
       return nullptr;
     }
   }
+
+  /*!
+   * \brief Look up the fallback to be used if no handler is registered
+   */
+  const runtime::PackedFunc* LookupFallback() const {
+    if (dispatch_fallback_.has_value()) {
+      return &*dispatch_fallback_;
+    } else {
+      return nullptr;
+    }
+  }
+
   /*
    * This type alias and the following free functions are created to reduce the binary bloat
    * from template and also hide implementation details from this header
    */
   using DispatchTable = std::unordered_map<std::string, std::vector<runtime::PackedFunc>>;
   /*! \brief The dispatch table. */
   DispatchTable dispatch_table_;
+  std::optional<runtime::PackedFunc> dispatch_fallback_;
 };
 
 }  // namespace printer

src/script/printer/ir_docsifier.cc

@@ -160,6 +160,11 @@ IRDocsifier::FType& IRDocsifier::vtable() {
 TVM_REGISTER_NODE_TYPE(FrameNode);
 TVM_REGISTER_NODE_TYPE(IRDocsifierNode);
 
+TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
+    .set_fallback([](ObjectRef obj, ObjectPath p, IRDocsifier d) -> Doc {
+      return d->AddMetadata(obj);
+    });
+
 }  // namespace printer
 }  // namespace script
 }  // namespace tvm