[REFACTOR][SCRIPT] Move script_printer.h → script/printer/config.h; rename TVM_SCRIPT_REPR; replace dialect fields with extra_config#19478
Conversation
…ename TVM_SCRIPT_REPR; replace dialect fields with extra_config include/tvm/ir/script_printer.h misplaced PrinterConfig and TVMScriptPrinter under tvm/ir/ — these are printer-layer concepts, not IR-specific. Move to include/tvm/script/printer/config.h next to the rest of the printer's public surface (doc.h, ir_docsifier.h, ir_docsifier_functor.h). Rename TVM_SCRIPT_REPR → TVM_REGISTER_SCRIPT_AS_REPR for clarity: the macro registers Script as the kRepr callback for an object type plus a per-type TVMScriptPrinter::vtable() dispatch entry. The new name is explicit about the registration semantic and aligns with the TVM_REGISTER_* family. Drop hard-coded dialect-specific fields from PrinterConfig (tir_prefix, relax_prefix, show_all_struct_info, buffer_dtype) in favor of a generic ffi::Map<String, Any> extra_config keyed by "<dialect>.<knob>" (e.g. "tirx.prefix", "relax.show_all_struct_info"). Each dialect reads its own keys via GetExtraConfig<T>(key, fallback) with defaults at the call site; the shared core never names a specific dialect through hardcoded fields. ir_prefix and module_alias also flip from std::string to ffi::String for consistency with the rest of the config. RedirectedReprPrinterMethod is de-inlined into a new src/script/printer/config.cc to keep <tvm/runtime/logging.h> out of the public header. ~6 header includers + ~80 macro use sites updated; Python config wrappers updated to write into extra_config with dotted keys.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the PrinterConfig by moving it to a new header and replacing dialect-specific fields with a generic extra_config map to improve decoupling. It introduces a GetExtraConfig helper for accessing these settings and updates the script representation registration macro. Feedback highlights that GetExtraConfig incorrectly uses Downcast for POD types and that the configuration constructor still contains hardcoded dialect-specific logic, which should be generalized to fully achieve the architectural goals.
| T GetExtraConfig(const ffi::String& key, T fallback) const { | ||
| auto it = extra_config.find(key); | ||
| if (it == extra_config.end()) return fallback; | ||
| return Downcast<T>((*it).second); |
There was a problem hiding this comment.
The use of Downcast<T> here is incorrect for non-ObjectRef types. Since extra_config stores ffi::Any values, and this method is used with types like bool (in binding.cc) and DataType (in buffer.cc), Downcast will fail to compile or error at runtime for these POD-like types. Using ffi::Any::cast<T>() is the correct way to retrieve values from an Any container generically.
return (*it).second.cast<T>();| for (const char* key : {"tirx.prefix", "relax.prefix"}) { | ||
| if (auto v = config_dict.Get(key)) { | ||
| n->extra_config.Set(ffi::String(key), v.value()); | ||
| } | ||
| } | ||
| // "tirx.buffer_dtype" is passed as a DLDataType string from Python; convert to DataType. | ||
| if (auto v = config_dict.Get("tirx.buffer_dtype")) { | ||
| DataType dt(ffi::StringToDLDataType(Downcast<ffi::String>(v.value()))); | ||
| n->extra_config.Set(ffi::String("tirx.buffer_dtype"), ffi::Any(dt)); | ||
| } | ||
| // Boolean dialect keys. | ||
| if (auto v = config_dict.Get("relax.show_all_struct_info")) { | ||
| n->extra_config.Set(ffi::String("relax.show_all_struct_info"), v.value()); | ||
| } |
There was a problem hiding this comment.
The implementation here contradicts the architectural goal stated in the PR description: "shared core never names a specific dialect through hard-coded fields." This constructor still contains hardcoded logic for tirx and relax dialects, including specific key mapping and DataType conversion for tirx.buffer_dtype. To achieve true decoupling, this logic should be generalized (e.g., by iterating over all non-core keys in config_dict) or moved to a registration-based system where dialects can define their own configuration conversion and validation logic.
|
Folded into #19479 — single coherent TVMScript dialect-friendly refactor PR. Commits absorbed; this PR's branch can be deleted. |
…estructure + dialect registry (#19479) ## Summary Restructure TVMScript to be dialect-agnostic at the script-core layer while letting each extension dialect (TIRX, Relax) own its own per-dialect script subtree. IR is below script in the dependency stack and is NOT a peer dialect — its script handlers stay in the shared core. This PR folds together two coupled refactors that were initially opened as separate PRs (#19478 and the original #19479); they share rename / relocation surface so they ship as one cohesive change. ## What this PR does ### Per-dialect script subtree (originally #19479) - Moves per-dialect printer + builder from `src/script/{printer,ir_builder}/{tirx,relax}/` to `src/{tirx,relax}/script/{printer,builder}/`. - Tightens `src/script/*.cc` CMake glob to the dialect-free core. - Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry (`script.ir_builder.decl_function.<type-key>`); removes cross-dialect includes from the shared core. - Adds `tvm.script.register_dialect` API + `__getattr__` + a `sys.meta_path` finder for Python-side dialect discovery. In-tree dialects (tirx, relax) registered centrally in `python/tvm/__init__.py`. - Drops the obsolete static re-export shims at `python/tvm/script/{parser,ir_builder}/{tirx,relax}/`. ### Dialect-agnostic printer config (originally #19478) - Relocates `include/tvm/ir/script_printer.h` → `include/tvm/script/printer/config.h` next to the rest of the printer's public surface. The header is not IR-specific. - Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for clarity (the macro registers Script as the kRepr callback + per-type vtable dispatch). Aligns with the `TVM_REGISTER_*` family. - Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`, `relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor of a generic `ffi::Map<String, Any> extra_config` keyed by `"<dialect>.<knob>"`. Each call site reads via the templated accessor `config->GetExtraConfig<T>("...", default)`. - Promotes `std::string` config fields to `ffi::String`. After this lands, the script-printer core knows nothing specific about any dialect — new dialects plug in via the registry pattern with zero core edits. Public Python API surface unchanged.
1 participant
Summary
include/tvm/ir/script_printer.hplacedPrinterConfigandTVMScriptPrinterundertvm/ir/— these are printer-layer concepts, not IR-specific. Move toinclude/tvm/script/printer/config.halongside the printer's other public headers (doc.h,ir_docsifier.h,ir_docsifier_functor.h).Rename
TVM_SCRIPT_REPR→TVM_REGISTER_SCRIPT_AS_REPR: the macro registers Script as thekReprcallback for an object type and installs the per-typeTVMScriptPrinter::vtable()dispatch entry. The new name is explicit about the registration semantic and aligns with theTVM_REGISTER_*family.Drop dialect-specific
PrinterConfigfields (tir_prefix,relax_prefix,show_all_struct_info,buffer_dtype) in favor of a genericffi::Map<String, Any> extra_configkeyed by"<dialect>.<knob>"(e.g."tirx.prefix","relax.show_all_struct_info"). Each dialect reads its own keys viaGetExtraConfig<T>(key, fallback)with defaults at the call site; the shared core never names a specific dialect through hard-coded fields.ir_prefixandmodule_aliasalso flip fromstd::stringtoffi::Stringfor consistency.RedirectedReprPrinterMethodis de-inlined intosrc/script/printer/config.ccto keep<tvm/runtime/logging.h>out of the public header.~6 header includers and ~80 macro use sites updated; Python config wrappers updated to write into
extra_configwith dotted keys.