[FIX][Script] Preserve device identity and unify rendering configuration

docs/arch/tvmscript.rst

@@ -158,11 +158,12 @@ Printing and round trips
 of expressions and statements. The translation engine tracks scopes, names and each
 Doc's original IR object. The script entry point maps these origins to diagnostic
 paths; the private Doc printer formats the tree, annotations and underlines as Python
-text. Printer configuration stays read-only throughout. ``DocToPythonScript`` also
-formats an existing Doc directly. This tree is separate from the parser's Python AST.
+text. Printer configuration stays read-only throughout. The Python document printer
+helper also formats an existing Doc directly for document-level tests. This tree is
+separate from the parser's Python AST.
 
-The public ``Script`` text entry points are declared in
-``tvm/script/printer/printer.h``. Their orchestration and diagnostic path mapping
+The public ``tvm::Script`` text entry point is declared in
+``tvm/script/printer/printer.h``. Its orchestration and diagnostic path mapping
 live in ``src/script/printer/printer.cc``; ``doc_translator.h`` exposes the
 IR-to-Doc translation protocol.
 

include/tvm/script/printer/doc.h

@@ -25,7 +25,6 @@
 #include <tvm/ir/expr.h>
 #include <tvm/ir/type.h>
 #include <tvm/runtime/device_api.h>
-#include <tvm/script/printer/config.h>
 
 #include <string>
 
@@ -1435,24 +1434,6 @@ class NamespaceDoc : public ExprDoc {
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NamespaceDoc, ExprDoc, NamespaceDocNode);
 };
 
-/*!
- * \brief Render a completed Doc tree with the supplied configuration.
- * \param doc The Doc tree to render.
- * \param config The rendering configuration.
- * \return The rendered Python source.
- */
-TVM_DLL ffi::String DocToPythonScript(Doc doc, const PrinterConfig& config);
-
-/*!
- * \brief Render explicit headers followed by a completed Doc tree.
- * \param doc The Doc tree to render.
- * \param config The rendering configuration.
- * \param header Ordered String source chunks or CommentDocs.
- * \return The rendered Python source.
- */
-TVM_DLL ffi::String DocToPythonScriptWithHeader(Doc doc, const PrinterConfig& config,
-                                                const ffi::Array<ffi::Any>& header);
-
 }  // namespace printer
 }  // namespace script
 }  // namespace tvm

include/tvm/script/printer/printer.h

@@ -54,14 +54,6 @@ TVM_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& d
 /*! \brief Read the registered namespace aliases. */
 TVM_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases();
 
-/*!
- * \brief Translate IR, recover diagnostic paths, and render Python text.
- * \param obj The input IR object.
- * \param config The translation and rendering options.
- * \return The rendered script.
- */
-TVM_DLL ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config);
-
 }  // namespace printer
 }  // namespace script
 }  // namespace tvm

src/relax/script/printer/binding.cc

@@ -64,8 +64,6 @@ ffi::Optional<ExprDoc> VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView
   TVM_FFI_CHECK(destination == nullptr, TypeError)
       << "printer statement-only node cannot fulfill a destination";
   if (auto func = binding->value.as<relax::FunctionNode>()) {
-    d->Emit(CommentDoc("from tvm.script import relax as R"), ffi::GetRef<ffi::ObjectRef>(binding));
-    d->Emit(CommentDoc(""), ffi::GetRef<ffi::ObjectRef>(binding));
     IdDoc lhs = VarDoc(d, binding->var);
     d->Translate(binding->value);
     FunctionDoc function = d->CurrentScopeDocs().back().as_or_throw<FunctionDoc>();

src/relax/script/printer/dependent_type.cc

@@ -58,10 +58,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
       kDocTranslate, FDocTranslate::FromNative<&ShapeTypeDocTranslate>());
 }
 
-}  // namespace
-
-ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
-                           bool include_vdevice) {
+ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+                                              const ffi::Object*) {
+  const auto* ty =
+      ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::TensorTypeNode>(
+          input);
   ffi::Array<ExprDoc> args;
   ffi::Array<ffi::String> keys;
   ffi::Array<ExprDoc> values;
@@ -82,7 +83,7 @@ ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
     keys.push_back("ndim");
     values.push_back(LiteralDoc::Int(ty->ndim, std::nullopt));
   }
-  if (include_vdevice && ty->vdevice.has_value()) {
+  if (ty->vdevice.has_value()) {
     keys.push_back("vdevice");
     if (auto selector = GlobalInfoSelector(d, ty->vdevice.value())) {
       values.push_back(LiteralDoc::Str(selector.value(), std::nullopt));
@@ -94,16 +95,6 @@ ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
   return NamespaceDoc("relax")->Attr("Tensor")->Call(args, keys, values);
 }
 
-namespace {
-
-ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
-                                              const ffi::Object*) {
-  const auto* ty =
-      ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::TensorTypeNode>(
-          input);
-  return RelaxTensorTypeDoc(d, ty, true);
-}
-
 TVM_FFI_STATIC_INIT_BLOCK() {
   ffi::reflection::TypeAttrDef<relax::TensorTypeNode>().attr(
       kDocTranslate, FDocTranslate::FromNative<&TensorTypeDocTranslate>());

src/relax/script/printer/distributed.cc

@@ -107,9 +107,7 @@ ffi::Optional<ExprDoc> DeviceMeshDocTranslate(DocTranslatorObj* d, ffi::AnyView
   }
   ExprDoc devices = LiteralDoc::None(std::nullopt);
   if (mesh->device_range.has_value()) {
-    CallDoc range = d->Translate(mesh->device_range.value()).value().as_or_throw<CallDoc>();
-    range->callee = NamespaceDoc("relax")->Attr("Range");
-    devices = range;
+    devices = d->Translate(mesh->device_range.value()).value();
   } else {
     ffi::Array<ExprDoc> ids;
     for (int64_t value : mesh->device_ids) ids.push_back(LiteralDoc::Int(value, std::nullopt));

src/relax/script/printer/expr.cc

@@ -71,14 +71,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
       kDocTranslate, FDocTranslate::FromNative<&ShapeExprDocTranslate>());
 }
 
-ffi::Optional<ExprDoc> DataflowVarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
-                                               const ffi::Object* destination) {
-  return VarDocTranslate(d, input, destination);
-}
-
 TVM_FFI_STATIC_INIT_BLOCK() {
   ffi::reflection::TypeAttrDef<relax::DataflowVarNode>().attr(
-      kDocTranslate, FDocTranslate::FromNative<&DataflowVarDocTranslate>());
+      kDocTranslate, FDocTranslate::FromNative<&VarDocTranslate>());
 }
 
 }  // namespace

src/relax/script/printer/function.cc

@@ -63,11 +63,7 @@ ffi::Optional<ExprDoc> FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in
   auto signature_candidates = CopyImplicitDefs(d);
   ffi::Optional<ExprDoc> ret_type = std::nullopt;
   if (!func->ret_ty.IsMissing()) {
-    if (auto tensor = func->ret_ty.as<relax::TensorTypeNode>()) {
-      ret_type = RelaxTensorTypeDoc(d, tensor, true);
-    } else {
-      ret_type = d->Translate(func->ret_ty).value();
-    }
+    ret_type = d->Translate(func->ret_ty).value();
   }
   ffi::Array<ffi::String> decorator_keys;
   ffi::Array<ExprDoc> decorator_values;

src/relax/script/printer/global_info.cc

@@ -44,8 +44,7 @@ ffi::Optional<ffi::String> GlobalInfoSelector(DocTranslatorObj* d, const GlobalI
         if (auto device = entry.as<relax::VDevice>();
             device && device.value()->target->kind->name == kind) {
           if (entry.same_as(info)) {
-            return ffi::String(std::string(kind) + ":" + std::to_string(index) + ":" +
-                               std::string(device.value()->memory_scope));
+            return ffi::String(std::string(kind) + ":" + std::to_string(index));
           }
           ++index;
         }

src/relax/script/printer/utils.h

@@ -34,8 +34,6 @@ namespace details {
 
 ffi::Optional<ffi::String> GlobalInfoSelector(DocTranslatorObj* d, const GlobalInfo& info);
 ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim);
-ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
-                           bool include_vdevice);
 ffi::Array<StmtDoc> RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq,
                                  ffi::Optional<IdDoc> destination = std::nullopt,
                                  ffi::Optional<ExprDoc> annotation = std::nullopt,

src/script/printer/doc_printer.cc

@@ -1159,8 +1159,8 @@ void PythonDocPrinter::PrintTypedDoc(const NamespaceDoc& doc) {
   const ffi::String& name = doc->canonical_name;
   ffi::String key = std::string(name) + ".prefix";
   ffi::String fallback = GetNamespaceAliases().Get(key).value_or(name);
-  output_ << (name == "ir" ? config()->ir_prefix
-                           : config()->GetExtraConfig<ffi::String>(key, fallback));
+  output_ << config()->GetExtraConfig<ffi::String>(key,
+                                                   name == "ir" ? config()->ir_prefix : fallback);
 }
 
 void PythonDocPrinter::PrintTypedDoc(const AttrAccessDoc& doc) {
@@ -1698,7 +1698,7 @@ ffi::String RenderPythonScript(Doc doc, const PrinterConfig& cfg,
   std::string script = printer.GetString();
 
   // GetString terminates non-empty output with one newline.  Preserve the
-  // established DocToPythonScript result without normalizing any other
+  // established rendering result without normalizing any other
   // trailing whitespace.
   if (!script.empty()) {
     TVM_FFI_ICHECK_EQ(script.back(), '\n');
@@ -1711,18 +1711,12 @@ ffi::String RenderPythonScript(Doc doc, const PrinterConfig& cfg,
 
 }  // namespace details
 
-ffi::String DocToPythonScriptWithHeader(Doc doc, const PrinterConfig& cfg,
-                                        const ffi::Array<ffi::Any>& header) {
-  return details::RenderPythonScript(std::move(doc), cfg, header, cfg->path_to_underline, {});
-}
-
-ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
-  return DocToPythonScriptWithHeader(std::move(doc), cfg, {});
-}
-
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
-  refl::GlobalDef().def("script.printer.DocToPythonScript", DocToPythonScript);
+  refl::GlobalDef().def("script.printer.DocToPythonScript", [](Doc doc,
+                                                               const PrinterConfig& config) {
+    return details::RenderPythonScript(std::move(doc), config, {}, config->path_to_underline, {});
+  });
 }
 
 }  // namespace printer

src/script/printer/doc_printer.h

@@ -19,6 +19,7 @@
 #ifndef SRC_SCRIPT_PRINTER_DOC_PRINTER_H_
 #define SRC_SCRIPT_PRINTER_DOC_PRINTER_H_
 
+#include <tvm/script/printer/config.h>
 #include <tvm/script/printer/doc.h>
 
 namespace tvm {

src/script/printer/ir/ir.cc

@@ -84,15 +84,7 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView in
       for (const auto& [key, entries] : infos) {
         ffi::Array<ExprDoc> items;
         for (const GlobalInfo& entry : entries) {
-          ExprDoc item = AnyValue(d, entry);
-          if (key == "mesh") {
-            if (auto mesh = item.as<CallDoc>(); mesh && mesh.value()->args.size() == 2) {
-              if (auto range = mesh.value()->args[1].as<CallDoc>()) {
-                range.value()->callee = NamespaceDoc("ir")->Attr("Range");
-              }
-            }
-          }
-          items.push_back(item);
+          items.push_back(AnyValue(d, entry));
         }
         keys.push_back(LiteralDoc::Str(key, std::nullopt));
         values.push_back(ListDoc(items));

src/script/printer/printer.cc

@@ -20,12 +20,14 @@
 #include <tvm/ffi/extra/serialization.h>
 #include <tvm/ffi/reflection/accessor.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/logging.h>
 #include <tvm/script/printer/doc.h>
 #include <tvm/script/printer/doc_translator.h>
 #include <tvm/script/printer/printer.h>
 
 #include <algorithm>
 #include <functional>
+#include <sstream>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -108,8 +110,8 @@ ffi::Map<ffi::String, ffi::Array<ffi::Any>> CollectMetadata(
 ffi::Array<ffi::String> DisplayAliases(const PrinterConfig& config) {
   ffi::Array<ffi::String> aliases;
   for (const auto& [key, fallback] : GetNamespaceAliases()) {
-    aliases.push_back(key == "ir.prefix" ? config->ir_prefix
-                                         : config->GetExtraConfig<ffi::String>(key, fallback));
+    aliases.push_back(config->GetExtraConfig<ffi::String>(
+        key, key == "ir.prefix" ? config->ir_prefix : fallback));
   }
   return aliases;
 }
@@ -330,9 +332,7 @@ class MapDocPaths {
   std::vector<std::pair<ffi::ObjectRef, AccessPath>> occurrences_;
 };
 
-}  // namespace
-
-ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
+ffi::String RenderScript(const ffi::ObjectRef& obj, const PrinterConfig& config) {
   ffi::Dict<Doc, ffi::ObjectRef> origins;
   Doc doc = DocTranslate(obj, &origins, config->extra_config);
   auto block = doc.as<StmtBlockDoc>();
@@ -410,8 +410,8 @@ ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
       std::string canonical = key;
       canonical.resize(canonical.size() - std::string(".prefix").size());
       if (!namespaces.count(canonical)) continue;
-      ffi::String alias = canonical == "ir" ? config->ir_prefix
-                                            : config->GetExtraConfig<ffi::String>(key, fallback);
+      ffi::String alias = config->GetExtraConfig<ffi::String>(
+          key, canonical == "ir" ? config->ir_prefix : fallback);
       ffi::String import = "from tvm.script import " + canonical + " as " + std::string(alias);
       if (config->GetExtraConfig<bool>("ir.comment_imports", false) && !config->show_meta) {
         header.push_back(CommentDoc(import));
@@ -421,13 +421,64 @@ ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) {
     }
     if (!header.empty()) header.push_back(ffi::String("\n"));
   }
+  if (config->path_to_underline.empty() && config->obj_to_underline.empty() &&
+      config->path_to_annotate.empty() && config->obj_to_annotate.empty() &&
+      !config->render_invisible_path_info) {
+    return details::RenderPythonScript(doc, config, header, {}, {});
+  }
   MapDocPaths paths(obj, config);
   auto [underline_paths, annotations] = paths.Map(doc, origins);
   return details::RenderPythonScript(doc, config, header, underline_paths, annotations);
 }
 
-TVM_FFI_STATIC_INIT_BLOCK() { ffi::reflection::GlobalDef().def("script.printer.Script", Script); }
+}  // namespace
 
 }  // namespace printer
 }  // namespace script
+namespace {
+
+std::string RenderFallbackWithInvisiblePathInfo(const ffi::String& script,
+                                                const PrinterConfig& config) {
+  if (!config->render_invisible_path_info || config->path_to_underline.empty()) {
+    return std::string(script);
+  }
+
+  std::ostringstream os;
+  for (size_t i = 0; i < config->path_to_underline.size(); ++i) {
+    if (i != 0) os << "\n";
+    os << "Access path: " << config->path_to_underline[i]
+       << "\nNote: No visible object for this path is rendered in TVMScript.";
+  }
+  os << "\n\n" << script;
+  return os.str();
+}
+
+}  // namespace
+
+std::string Script(const ffi::ObjectRef& node, const ffi::Optional<PrinterConfig>& cfg) {
+  PrinterConfig config = cfg.value_or(PrinterConfig());
+  static ffi::reflection::TypeAttrColumn translate(script::printer::kDocTranslate);
+  // Builtin runtime roots keep their native repr; hooks still translate them within IR.
+  if (!node.defined() || node->type_index() < ffi::TypeIndex::kTVMFFIDynObjectBegin ||
+      translate[node->type_index()].type_index() == ffi::TypeIndex::kTVMFFINone) {
+    return RenderFallbackWithInvisiblePathInfo(ffi::ReprPrint(ffi::Any(node)), config);
+  }
+  return std::string(script::printer::RenderScript(node, config));
+}
+
+std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
+  try {
+    PrinterConfig config;
+    config->extra_config.Set("ir.comment_imports", true);
+    // Call translation directly so an unsupported type cannot recurse through ffi repr.
+    return std::string(script::printer::RenderScript(obj, config));
+  } catch (const tvm::ffi::Error& e) {
+    LOG(WARNING) << "TVMScript printer falls back to the basic address printer with the error:\n"
+                 << e.what();
+    std::ostringstream os;
+    os << obj->GetTypeKey() << '(' << obj.get() << ')';
+    return os.str();
+  }
+}
+
 }  // namespace tvm