[REFACTOR][SCRIPT] tvmscript streamline: lift printer.h, restore one-way dep, migrate dialect config to extra_config#19631
Merged
Conversation
…nter.h Add a new public header `include/tvm/script/printer/printer.h` that declares the free function `tvm::Script(node, optional_config)` plus the dispatch vtable used by per-dialect printers. `PrinterConfig` and related dataclass helpers stay in `config.h`; `printer.h` is the entry-point header that callers include to invoke printing. The `TVMScriptPrinter::Script(...)` static method body moves to the free function `tvm::Script(...)` in `src/script/printer/script_printer.cc`. `TVMScriptPrinter` class and `TVM_REGISTER_SCRIPT_AS_REPR` macro move from `config.h` to `printer.h`. Existing direct callers in s_tir/schedule/error.cc, script/printer/config.cc, and script/printer/utils.h are rewritten to the new entry point.
The macro injected a `Script()` member method on IR Object types and forced those headers to reverse-include the script printer's config. The IR layer should not depend on the script layer; script depends on ir. Remove the macro from include/tvm/script/printer/config.h and drop its 5 usages (PrimExprNode, IRModuleNode, tirx Buffer/PrimFunc/Stmt nodes). The single C++ caller of `mod->Script()` in src/s_tir/meta_schedule/database/json_database.cc switches to the explicit `tvm::Script(...)` free function from include/tvm/script/printer/printer.h. The reverse-include of script/printer/config.h is dropped from the five IR/tirx headers, restoring one-way dependency: script depends on ir, never the other way.
…a_config PrinterConfig currently has 5 dialect-specific fields (tir_prefix, tir_import_module, tirx_prefix, relax_prefix, buffer_dtype). The first four have zero in-tree readers and are removed entirely. The fifth, buffer_dtype, has a single reader in src/tirx/script/printer/buffer.cc which is rewritten to look up the value from PrinterConfig::extra_config under the key "tirx.buffer_dtype". After this commit, PrinterConfig's core schema is dialect-agnostic; per-dialect customizations live in the extra_config ffi::Map.
Following the C++ PrinterConfig cleanup, drop the four dialect-specific
keyword arguments from the Python script entry points (Scriptable.script,
_relax_script, show, BasePyModule.script): tir_prefix, tir_import_module,
tirx_prefix, relax_prefix, buffer_dtype. Add a generic
extra_config: dict | None = None kwarg that maps directly to the C++
PrinterConfig.extra_config.
The auto-switch logic in Scriptable.script that detects tirx PrimFunc /
IRModule and switches to "Tx" prefix is preserved; it now stamps
extra_config["tirx.prefix"] instead of the removed tir_prefix kwarg.
Callers can suppress auto-switch or override the prefix via
extra_config={"tirx.prefix": "..."}.
Update six call sites in tests/python/tirx/ that previously passed
dialect kwargs to use extra_config={"tirx.prefix": "Tx"} instead.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the TVMScript printer configuration by removing dialect-specific fields (such as tir_prefix, relax_prefix, and buffer_dtype) from the core PrinterConfig class and introducing a generic extra_config dictionary. It also replaces the static TVMScriptPrinter::Script method with a free function tvm::Script defined in a new entry-point header. One critical issue was identified in python/tvm/runtime/script_printer.py where passing "tirx.buffer_dtype" as a string inside extra_config will cause a runtime type mismatch in C++ when downcasting to DataType. A code suggestion is provided to intercept and promote this key to the top level of the configuration dictionary so that the C++ side can parse it correctly.
Removing #include <tvm/script/printer/config.h> from ir/expr.h and ir/module.h (in the macro-drop commit) broke three headers and one test that relied on those transitively pulling in PrinterConfig and NodeFunctor: - include/tvm/script/printer/doc.h: add config.h (uses PrinterConfig in DocToPythonScript declaration) - include/tvm/script/printer/ir_docsifier.h: add config.h (PrinterConfig member on IRDocsifierNode arrived via module.h → config.h) - include/tvm/script/ir_builder/base.h: add node_functor.h (NodeFunctor arrived via expr.h → config.h → node_functor.h) - tests/cpp/tir_scalable_datatype.cc: rewrite call->Script() to tvm::Script(call) since the Script() member was removed with TVM_OBJECT_ENABLE_SCRIPT_PRINTER
tqchen
force-pushed
the
tvm-tvmscript-streamline
branch
from
c1fd97f to
e42381b
Compare
tlopex
approved these changes
May 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
The
tvm::irlayer previously had a reverse dependency ontvm::script, injected via theTVM_OBJECT_ENABLE_SCRIPT_PRINTER()macro that added aScript()member method to IR node types (IRModule, PrimExpr, Buffer, PrimFunc, Stmt). This violated the intended one-way dependency:scriptshould depend onir, never the other way around.Additionally,
PrinterConfigNodeaccumulated dialect-specific fields (tir_prefix,tir_import_module,tirx_prefix,relax_prefix) that created leakage between the generic printer infrastructure and dialect internals.Changes
This PR restores the clean dependency direction and encapsulates dialect config properly, in 5 commits:
Lift TVMScript entry point into
script/printer/printer.h: New headerinclude/tvm/script/printer/printer.hintroduces:tvm::Script()free function replacingTVMScriptPrinter::Script()static methodTVMScriptPrinterclass with vtable (NodeFunctor<std::string(...)>)TVM_REGISTER_SCRIPT_AS_REPRmacro for registering per-type repr callbacksDrop
TVM_OBJECT_ENABLE_SCRIPT_PRINTERmacro: Remove the macro from all IR headers (ir/expr.h,ir/module.h,tirx/buffer.h,tirx/function.h,tirx/stmt.h), eliminating the reverseir→scriptdependency. All call sites of.Script()member methods updated to usetvm::Script().Move dialect-specific
PrinterConfigfields toextra_config: Removetir_prefix,tir_import_module,tirx_prefix,relax_prefixfromPrinterConfigNode. Dialect internals now read their config viaGetExtraConfig<T>(key, fallback)with dotted keys (e.g.,"tirx.prefix").buffer_dtypeis kept as a top-level field alongsideint_dtype/float_dtypesince it is a shared scalar-literal default, not a dialect-specific knob.Python: drop dialect kwargs, expose
extra_config: UpdatePrinterConfig,Scriptable.script(),Scriptable.show(),Scriptable._relax_script(), andBasePyModule.script()to useextra_config: dict | None = Noneinstead of individual dialect kwargs. The tirx auto-switch logic is preserved.Fix transitive include breakage: Explicitly add direct includes for
config.handnode_functor.hwhere headers previously relied on transitive paths throughexpr.h/module.h.Testing
tests/python/tirx/test_printer_tir_namespaces.py): 13/13 pass