[REFACTOR] Phase out LegacyReprPrinter and improve CommonSubExprElim

include/tvm/node/repr_printer.h

@@ -52,32 +52,6 @@ class ReprPrinter {
   TVM_DLL static FType& vtable();
 };
 
-/*! \brief Legacy behavior of ReprPrinter. */
-class ReprLegacyPrinter {
- public:
-  /*! \brief The indentation level. */
-  int indent{0};
-
-  explicit ReprLegacyPrinter(std::ostream& stream)  // NOLINT(*)
-      : stream(stream) {}
-
-  /*! \brief The node to be printed. */
-  TVM_DLL void Print(const ObjectRef& node);
-  /*! \brief Print indent to the stream */
-  TVM_DLL void PrintIndent();
-  /*! \brief Could the LegacyPrinter dispatch the node */
-  TVM_DLL static bool CanDispatch(const ObjectRef& node);
-  /*! \brief Return the ostream it maintains */
-  TVM_DLL std::ostream& Stream() const;
-  // Allow registration to be printer.
-  using FType = NodeFunctor<void(const ObjectRef&, ReprLegacyPrinter*)>;
-  TVM_DLL static FType& vtable();
-
- private:
-  /*! \brief The output stream */
-  std::ostream& stream;
-};
-
 /*!
  * \brief Dump the node to stderr, used for debug purposes.
  * \param node The input node
@@ -113,12 +87,6 @@ inline std::ostream& operator<<(std::ostream& os, const Variant<V...>& n) {  //
   return os;
 }
 
-inline std::string AsLegacyRepr(const ObjectRef& n) {
-  std::ostringstream os;
-  ReprLegacyPrinter(os).Print(n);
-  return os.str();
-}
 }  // namespace ffi
-using ffi::AsLegacyRepr;
 }  // namespace tvm
 #endif  // TVM_NODE_REPR_PRINTER_H_

src/ir/analysis.cc

@@ -31,7 +31,7 @@ namespace ir {
 Map<GlobalVar, Array<GlobalVar>> CollectCallMap(const IRModule& mod) {
   struct CalleeCollectorImpl : CalleeCollector {
     void Mark(GlobalVar gvar) override { gvars.push_back(gvar); }
-    support::OrderedSet<GlobalVar> gvars;
+    support::OrderedSet<GlobalVar, ObjectPtrHash, ObjectPtrEqual> gvars;
   };
 
   Map<GlobalVar, Array<GlobalVar>> call_map;

src/node/repr_printer.cc

@@ -97,38 +97,6 @@ ReprPrinter::FType& ReprPrinter::vtable() {
   return inst;
 }
 
-void ReprLegacyPrinter::Print(const ObjectRef& node) {
-  static const FType& f = vtable();
-  if (!node.defined()) {
-    stream << "(nullptr)";
-  } else if (f.can_dispatch(node)) {
-    f(node, this);
-  } else {
-    try {
-      stream << node;  // Use ReprPrinter
-    } catch (const tvm::Error& e) {
-      LOG(WARNING) << "ReprPrinter fails";
-      stream << node->GetTypeKey() << '(' << node.get() << ')';
-    }
-  }
-}
-
-bool ReprLegacyPrinter::CanDispatch(const ObjectRef& node) {
-  static const FType& f = vtable();
-  return !node.defined() || f.can_dispatch(node);
-}
-
-void ReprLegacyPrinter::PrintIndent() {
-  for (int i = 0; i < indent; ++i) {
-    stream << ' ';
-  }
-}
-
-ReprLegacyPrinter::FType& ReprLegacyPrinter::vtable() {
-  static FType inst;
-  return inst;
-}
-
 void Dump(const runtime::ObjectRef& n) { std::cerr << n << "\n"; }
 
 void Dump(const runtime::Object* n) { Dump(runtime::GetRef<runtime::ObjectRef>(n)); }
@@ -138,7 +106,4 @@ TVM_FFI_REGISTER_GLOBAL("node.AsRepr").set_body_typed([](ffi::Any obj) {
   os << obj;
   return os.str();
 });
-
-TVM_FFI_REGISTER_GLOBAL("node.AsLegacyRepr").set_body_typed(ffi::AsLegacyRepr);
-
 }  // namespace tvm

src/node/script_printer.cc

@@ -32,7 +32,10 @@ TVMScriptPrinter::FType& TVMScriptPrinter::vtable() {
 
 std::string TVMScriptPrinter::Script(const ObjectRef& node, const Optional<PrinterConfig>& cfg) {
   if (!TVMScriptPrinter::vtable().can_dispatch(node)) {
-    return AsLegacyRepr(node);
+    std::ostringstream os;
+    ReprPrinter printer(os);
+    printer.Print(node);
+    return os.str();
   }
   return TVMScriptPrinter::vtable()(node, cfg.value_or(PrinterConfig()));
 }

src/relax/analysis/computable_at_compile_time.cc

@@ -83,7 +83,7 @@ class CompileTimeCollector : ExprVisitor {
     }
   }
 
-  support::OrderedSet<Var> known_relax_vars_;
+  support::OrderedSet<Var, ObjectPtrHash, ObjectPtrEqual> known_relax_vars_;
   std::unordered_set<tir::Var> known_tir_vars_;
 };
 }  // namespace

src/relax/analysis/udchain.cc

@@ -56,8 +56,8 @@ class UDChain : relax::ExprVisitor {
  private:
   Map<Var, Expr> bound_values;
   std::unordered_set<Var> forward_declarations;
-  std::unordered_map<Var, support::OrderedSet<Var>> usage_map;
-  support::OrderedSet<Var> outputs;
+  std::unordered_map<Var, support::OrderedSet<Var, ObjectPtrHash, ObjectPtrEqual>> usage_map;
+  support::OrderedSet<Var, ObjectPtrHash, ObjectPtrEqual> outputs;
 
   Optional<Var> cur_user_;
 

src/relax/ir/binding_rewrite.cc

@@ -321,7 +321,8 @@ Expr RemoveAllUnused(Expr expr) {
   auto var_usage = CollectVarUsage(expr);
 
   // For the purpose of
-  support::OrderedSet<Var> externally_exposed(var_usage.outputs.begin(), var_usage.outputs.end());
+  support::OrderedSet<Var, ObjectPtrHash, ObjectPtrEqual> externally_exposed(
+      var_usage.outputs.begin(), var_usage.outputs.end());
   for (const auto& [var, expr] : var_usage.bound_values) {
     if (ContainsImpureCall(expr)) {
       externally_exposed.insert(var);

src/relax/transform/inline_functions.cc

@@ -138,7 +138,7 @@ class FunctionInliner : public ExprMutator {
   }
 
   const Map<Variant<String, GlobalVar>, Function>& replacements_;
-  support::OrderedSet<GlobalVar> inline_stack_;
+  std::unordered_set<GlobalVar, ObjectPtrHash, ObjectPtrEqual> inline_stack_;
 };
 }  // namespace
 

src/relax/transform/run_codegen.cc

@@ -44,7 +44,7 @@ class CodeGenRunner : ExprMutator {
                Array<String> entry_function_names) {
     IRModule mod = builder_->GetContextIRModule();
 
-    support::OrderedSet<GlobalVar> entry_functions;
+    support::OrderedSet<GlobalVar, ObjectPtrHash, ObjectPtrEqual> entry_functions;
     // Any user-provided functions are treated as entry functions.
     for (const auto& name : entry_function_names) {
       entry_functions.insert(mod->GetGlobalVar(name));