This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5745c209f4 [REFACTOR][SCRIPT] Keep dependent shape recursion in
docsifier (#19940)
5745c209f4 is described below
commit 5745c209f482179f7edff965b185eb1f6c83e968
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Jul 5 06:29:56 2026 +0800
[REFACTOR][SCRIPT] Keep dependent shape recursion in docsifier (#19940)
Replace the nested `DocToPythonScript` invocation in Relax
dependent-shape printing with an expression-string Doc rendered during
the active `PythonDocPrinter` traversal.
This keeps recursive IR-to-doc conversion inside the active docsifier,
preserves naming, precedence, escaping, source paths, and printer
configuration, and avoids a nested top-level renderer or a new public
rendering entry point.
Expression-string escaping is streamed directly into the final output so
wrapper and nested source spans retain exact escaped byte offsets.
---
include/tvm/script/printer/doc.h | 34 +++++++
src/relax/script/printer/dependent_type.cc | 5 +-
src/script/printer/doc.cc | 10 ++
src/script/printer/doc_printer/base_doc_printer.cc | 2 +
src/script/printer/doc_printer/base_doc_printer.h | 5 +
.../printer/doc_printer/python_doc_printer.cc | 109 +++++++++++++++++++++
tests/python/relax/test_tvmscript_printer_relax.py | 51 ++++++++++
7 files changed, 212 insertions(+), 4 deletions(-)
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index bc90e53657..0442852692 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -328,6 +328,40 @@ class LiteralDoc : public ExprDoc {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LiteralDoc, ExprDoc,
LiteralDocNode);
};
+/*!
+ * \brief Doc that renders an expression as a Python string literal.
+ *
+ * \sa ExprStringDoc
+ */
+class ExprStringDocNode : public ExprDocNode {
+ public:
+ /*! \brief The expression to render as a string. */
+ ExprDoc value{ffi::UnsafeInit()};
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<ExprStringDocNode>().def_ro("value",
&ExprStringDocNode::value);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ExprStringDoc",
ExprStringDocNode, ExprDocNode);
+};
+
+/*!
+ * \brief Reference type of ExprStringDocNode.
+ *
+ * \sa ExprStringDocNode
+ */
+class ExprStringDoc : public ExprDoc {
+ public:
+ /*!
+ * \brief Constructor of ExprStringDoc.
+ * \param value The expression to render as a string.
+ * \param object_path The object path.
+ */
+ explicit ExprStringDoc(ExprDoc value, const ffi::Optional<AccessPath>&
object_path);
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprStringDoc, ExprDoc,
ExprStringDocNode);
+};
+
/*!
* \brief Doc that represents identifier.
*
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index c3b77223b5..2c6ea7e9c7 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -54,10 +54,7 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath&
e_p, const IRDocsifie
}
// Step 3. Stringify the PrimExpr if func var exists
if (func_var_mode) {
- // This nested render is only converting a shape expression into one token.
- // Give it an independent configuration so it cannot consume or emit the
- // enclosing invocation's access-path diagnostics.
- return LiteralDoc::Str(DocToPythonScript(expr_doc, PrinterConfig()), e_p);
+ return ExprStringDoc(expr_doc, e_p);
}
return expr_doc;
}
diff --git a/src/script/printer/doc.cc b/src/script/printer/doc.cc
index ffdb081a48..97ebabaaa4 100644
--- a/src/script/printer/doc.cc
+++ b/src/script/printer/doc.cc
@@ -33,6 +33,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
StmtDocNode::RegisterReflection();
StmtBlockDocNode::RegisterReflection();
LiteralDocNode::RegisterReflection();
+ ExprStringDocNode::RegisterReflection();
IdDocNode::RegisterReflection();
AttrAccessDocNode::RegisterReflection();
IndexDocNode::RegisterReflection();
@@ -94,6 +95,15 @@ LiteralDoc::LiteralDoc(ffi::Any value, const
ffi::Optional<AccessPath>& object_p
this->data_ = std::move(n);
}
+ExprStringDoc::ExprStringDoc(ExprDoc value, const ffi::Optional<AccessPath>&
object_path) {
+ ffi::ObjectPtr<ExprStringDocNode> n = ffi::make_object<ExprStringDocNode>();
+ n->value = value;
+ if (object_path.defined()) {
+ n->source_paths.push_back(object_path.value());
+ }
+ this->data_ = std::move(n);
+}
+
IdDoc::IdDoc(ffi::String name) {
ffi::ObjectPtr<IdDocNode> n = ffi::make_object<IdDocNode>();
n->name = name;
diff --git a/src/script/printer/doc_printer/base_doc_printer.cc
b/src/script/printer/doc_printer/base_doc_printer.cc
index 93e0233d38..4b4d799cd5 100644
--- a/src/script/printer/doc_printer/base_doc_printer.cc
+++ b/src/script/printer/doc_printer/base_doc_printer.cc
@@ -301,6 +301,8 @@ void DocPrinter::PrintDoc(const Doc& doc) {
if (auto doc_node = doc.as<LiteralDoc>()) {
PrintTypedDoc(doc_node.value());
+ } else if (auto doc_node = doc.as<ExprStringDoc>()) {
+ PrintTypedDoc(doc_node.value());
} else if (auto doc_node = doc.as<IdDoc>()) {
PrintTypedDoc(doc_node.value());
} else if (auto doc_node = doc.as<AttrAccessDoc>()) {
diff --git a/src/script/printer/doc_printer/base_doc_printer.h
b/src/script/printer/doc_printer/base_doc_printer.h
index 1dee2322f3..c8335f951a 100644
--- a/src/script/printer/doc_printer/base_doc_printer.h
+++ b/src/script/printer/doc_printer/base_doc_printer.h
@@ -105,6 +105,11 @@ class DocPrinter {
*/
virtual void PrintTypedDoc(const LiteralDoc& doc) = 0;
+ /*!
+ * \brief Virtual method to print an ExprStringDoc
+ */
+ virtual void PrintTypedDoc(const ExprStringDoc& doc) = 0;
+
/*!
* \brief Virtual method to print an IdDoc
*/
diff --git a/src/script/printer/doc_printer/python_doc_printer.cc
b/src/script/printer/doc_printer/python_doc_printer.cc
index c8f32eef0c..0a82430e8e 100644
--- a/src/script/printer/doc_printer/python_doc_printer.cc
+++ b/src/script/printer/doc_printer/python_doc_printer.cc
@@ -27,6 +27,7 @@
#include <map>
#include <optional>
#include <sstream>
+#include <streambuf>
#include <string>
#include "../../../support/str_escape.h"
@@ -38,6 +39,100 @@ namespace script {
namespace printer {
namespace {
+/*!
+ * \brief Escape an ExprStringDoc child while preserving final output
positions.
+ *
+ * ExprStringDoc must print its child through the active PythonDocPrinter so
that the child's
+ * precedence, printer configuration, and source-path bookkeeping stay in the
current traversal.
+ * However, escaping after printing into a temporary buffer would make
PrintDoc record positions
+ * in that unescaped buffer, rather than the final Python source. This scoped
adapter instead
+ * replaces the output stream's buffer only while the child is printed. It
escapes each write and
+ * forwards it immediately to the original final buffer; the caller emits the
surrounding quotes
+ * outside the scope so that only the child is transformed.
+ *
+ * The adapter owns neither the output stream nor its original buffer. Both
outlive this stack
+ * object. The constructor saves the original buffer in destination_ and
installs this adapter;
+ * the destructor restores the original buffer before the caller continues
writing to the stream.
+ *
+ * PrintDoc uses tellp() before and after every child to record source spans.
tellp() reaches
+ * seekoff(0, cur, out), which is the only seek this adapter accepts and
delegates directly to the
+ * original buffer. Because that buffer already contains the preceding output
and advances by the
+ * escaped byte count, child spans use absolute positions in the final source,
including any
+ * expansion introduced by StrEscape.
+ *
+ * A raw newline violates the one-line expression-string contract. xsputn
records its presence
+ * while still escaping and forwarding the write, and the caller checks
saw_newline() together
+ * with the stream state before completing the literal. A destination short
write is reported as
+ * an input write failure so that the output stream records the error. On
every exit path,
+ * including exception unwinding, the noexcept destructor restores the
original buffer and then
+ * reapplies the child traversal's stream state, which rdbuf() would otherwise
clear.
+ */
+class ScopedExprStringEscapeBuf : public std::streambuf {
+ public:
+ explicit ScopedExprStringEscapeBuf(std::ostream* output)
+ : output_(output), destination_(output->rdbuf()) {
+ TVM_FFI_ICHECK(output_->good()) << "Cannot escape into a failed output
stream";
+ output_->rdbuf(this);
+ }
+
+ ~ScopedExprStringEscapeBuf() noexcept {
+ // Swapping rdbuf resets rdstate, so retain the child's state across
restoration. setstate may
+ // throw under the stream's exception mask; suppress that throw so an
in-flight exception is
+ // never replaced, while setstate still records the bits before throwing.
+ std::ios_base::iostate state = output_->rdstate();
+ output_->rdbuf(destination_);
+ try {
+ output_->setstate(state);
+ } catch (...) {
+ // Preserve the stream state without replacing an exception already in
flight.
+ }
+ }
+
+ ScopedExprStringEscapeBuf(const ScopedExprStringEscapeBuf&) = delete;
+ ScopedExprStringEscapeBuf& operator=(const ScopedExprStringEscapeBuf&) =
delete;
+
+ bool saw_newline() const { return saw_newline_; }
+
+ protected:
+ std::streamsize xsputn(const char* data, std::streamsize count) final {
+ if (count <= 0) return 0;
+ saw_newline_ = saw_newline_ || std::find(data, data + count, '\n') != data
+ count;
+ // StrEscape is byte-wise, so each write can be transformed independently
and forwarded
+ // without retaining or copying the complete output. Report consumed input
bytes, not the
+ // potentially larger number of escaped bytes written to the destination.
+ std::string escaped = support::StrEscape(data, static_cast<size_t>(count));
+ return destination_->sputn(escaped.data(), escaped.size()) ==
+ static_cast<std::streamsize>(escaped.size())
+ ? count
+ : 0;
+ }
+
+ int_type overflow(int_type ch) final {
+ if (traits_type::eq_int_type(ch, traits_type::eof())) {
+ return traits_type::not_eof(ch);
+ }
+ char value = traits_type::to_char_type(ch);
+ return xsputn(&value, 1) == 1 ? ch : traits_type::eof();
+ }
+
+ int sync() final { return destination_->pubsync(); }
+
+ pos_type seekoff(off_type offset, std::ios_base::seekdir direction,
+ std::ios_base::openmode mode) final {
+ // Support the current-position query used by tellp(), preserving the
destination's absolute
+ // escaped offset. ExprStringDoc rendering never needs to reposition the
output sequence.
+ if (offset == 0 && direction == std::ios_base::cur && (mode &
std::ios_base::out)) {
+ return destination_->pubseekoff(offset, direction, mode);
+ }
+ return pos_type(off_type(-1));
+ }
+
+ private:
+ std::ostream* output_;
+ std::streambuf* destination_;
+ bool saw_newline_{false};
+};
+
ffi::String RenderInvisiblePathInfo(const ffi::String& script,
const ffi::Array<AccessPath>&
requested_paths,
const
ffi::Array<ffi::Optional<AccessPath>>& visible_paths) {
@@ -147,6 +242,7 @@ ExprPrecedence GetExprPrecedence(const ExprDoc& doc) {
// Key is the type index of Doc
static const std::unordered_map<uint32_t, ExprPrecedence>
doc_type_precedence = {
{LiteralDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
+ {ExprStringDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{IdDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{AttrAccessDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{IndexDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
@@ -181,6 +277,7 @@ class PythonDocPrinter : public DocPrinter {
using DocPrinter::PrintDoc;
void PrintTypedDoc(const LiteralDoc& doc) final;
+ void PrintTypedDoc(const ExprStringDoc& doc) final;
void PrintTypedDoc(const IdDoc& doc) final;
void PrintTypedDoc(const AttrAccessDoc& doc) final;
void PrintTypedDoc(const IndexDoc& doc) final;
@@ -394,6 +491,18 @@ void PythonDocPrinter::PrintTypedDoc(const LiteralDoc&
doc) {
}
}
+void PythonDocPrinter::PrintTypedDoc(const ExprStringDoc& doc) {
+ this->output_ << '"';
+ {
+ ScopedExprStringEscapeBuf escaping_scope(&this->output_);
+ this->PrintDoc(doc->value);
+ TVM_FFI_ICHECK(this->output_.good()) << "Failed to render an expression
string literal";
+ TVM_FFI_ICHECK(!escaping_scope.saw_newline())
+ << "An expression rendered inside a Python string literal must be one
line";
+ }
+ this->output_ << '"';
+}
+
void PythonDocPrinter::PrintTypedDoc(const IdDoc& doc) { output_ << doc->name;
}
void PythonDocPrinter::PrintTypedDoc(const AttrAccessDoc& doc) {
diff --git a/tests/python/relax/test_tvmscript_printer_relax.py
b/tests/python/relax/test_tvmscript_printer_relax.py
index a1a7caeab4..b98fb33c41 100644
--- a/tests/python/relax/test_tvmscript_printer_relax.py
+++ b/tests/python/relax/test_tvmscript_printer_relax.py
@@ -17,10 +17,12 @@
# pylint: disable=missing-docstring
# ruff: noqa: E501, F841
+from tvm_ffi.access_path import AccessPath
import tvm
import tvm.testing
from tvm import IRModule, relax, tirx
+from tvm.runtime.script_printer import PrinterConfig, _script
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
@@ -53,6 +55,55 @@ def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
)
+def test_function_dependent_shape_escaped_source_spans():
+ n = tirx.Var("n", "int64")
+ cast = tirx.Cast("int64", n)
+ x = relax.Var("x", relax.TensorType([cast], "float32"))
+ ret_ty = relax.TensorType(dtype="float32", ndim=1)
+ func = relax.Function([x], x, ret_ty=ret_ty).with_attr("global_symbol",
"main")
+ cast_path = (
+ AccessPath.root()
+ .attr("params")
+ .array_item(0)
+ .attr("ty")
+ .attr("shape")
+ .attr("values")
+ .array_item(0)
+ )
+
+ def render(path):
+ config = PrinterConfig(
+ verbose_expr=True,
+ num_context_lines=0,
+ path_to_underline=[path],
+ extra_config={"render_invisible_path_info": True},
+ )
+ first = _script(func, config)
+ assert _script(func, config) == first
+ assert first.count("Access path:") == 1
+ lines = first.splitlines()
+ definition_index = next(i for i, line in enumerate(lines) if "def
main" in line)
+ assert "Access path:" not in lines[definition_index]
+ return lines[definition_index], lines[definition_index + 1]
+
+ expression = r'"T.Cast(\"int64\", n)"'
+ definition, underline = render(cast_path)
+ expression_start = definition.index(expression)
+ assert underline[expression_start : expression_start + len(expression)] ==
"^" * len(expression)
+ assert underline.strip() == "^" * len(expression)
+
+ escaped_dtype = r"\"int64\""
+ definition, underline = render(cast_path.attr("dtype"))
+ dtype_start = definition.index(escaped_dtype)
+ assert underline[dtype_start : dtype_start + len(escaped_dtype)] == "^" *
len(escaped_dtype)
+ assert underline.strip() == "^" * len(escaped_dtype)
+
+ definition, underline = render(cast_path.attr("value"))
+ variable_start = definition.index(expression) + expression.rindex("n")
+ assert underline[variable_start] == "^"
+ assert underline.strip() == "^"
+
+
def test_lone_private_function():
@R.function(private=True)
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)): # type: ignore
