[REFACTOR] Phase out LegacyReprPrinter and improve CommonSubExprElim

This PR phases out LegacyReprPrinter. Previously common subexpr elim
relies on sorting on legacy repr for determism, which is hacky.

This PR introduces an ordered_map impl in support to ensure determinism
and migrates the CSE pass to use that instead.

Author: tqchen

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,7 @@ TVMScriptPrinter::FType& TVMScriptPrinter::vtable() {
 
 std::string TVMScriptPrinter::Script(const ObjectRef& node, const Optional<PrinterConfig>& cfg) {
   if (!TVMScriptPrinter::vtable().can_dispatch(node)) {
-    return AsLegacyRepr(node);
+    TVM_FFI_THROW(RuntimeError) << "TVMScriptPrinter cannot dispatch node: " << node.GetTypeKey();
   }
   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));