This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-doc-translation in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 5363be9991d9e2ad243533e344e00c7d80140e1d Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 30 20:36:42 2026 +0000 [REFACTOR][Printer] Move shared printing into tvm printer --- CMakeLists.txt | 18 ++-- docs/reference/api/python/index.rst | 2 +- docs/reference/api/python/{script => }/printer.rst | 6 +- include/tvm/{script => }/printer/config.h | 10 +- include/tvm/{script => }/printer/doc.h | 89 +++++++-------- include/tvm/{script => }/printer/doc_translator.h | 20 ++-- include/tvm/{script => }/printer/printer.h | 10 +- include/tvm/script/printer/v2/doc.h | 23 ---- include/tvm/script/printer/v2/doc_translator.h | 23 ---- python/tvm/{script => }/printer/__init__.py | 0 python/tvm/{script => }/printer/_ffi_api.py | 4 +- python/tvm/{script => }/printer/doc.py | 56 +++++----- python/tvm/{script => }/printer/doc_printer.py | 0 python/tvm/{script => }/printer/highlight.py | 0 python/tvm/relax/base_py_module.py | 2 +- python/tvm/runtime/script_printer.py | 8 +- python/tvm/s_tir/schedule/trace.py | 4 +- src/ir/expr.cc | 2 +- src/ir/prim/expr.cc | 2 +- src/{script => }/printer/config.cc | 4 +- src/{script => }/printer/dialect_prefix.h | 8 +- src/{script => }/printer/doc.cc | 120 +++++++++++---------- .../printer/v2 => printer}/doc_translator.cc | 80 +++++++------- src/{script/printer/v2 => printer}/entry.cc | 18 ++-- .../printer/v2 => printer}/python_doc_printer.cc | 13 +-- src/{script => }/printer/script_printer.cc | 32 ++---- src/{script => }/printer/utils.h | 8 +- src/s_tir/meta_schedule/database/json_database.cc | 2 +- src/s_tir/schedule/error.cc | 2 +- src/s_tir/script/printer/config.cc | 4 +- src/script/printer/v2/doc.cc | 36 ------- tests/cpp/tir_scalable_datatype.cc | 6 +- tests/python/script/test_script_printer_doc.py | 4 +- 33 files changed, 242 insertions(+), 374 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 2ebe9fade2..6d1f599254 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -303,18 +303,18 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS src/topi/*.cc src/driver/*.cc src/support/*.cc - # TVMScript shared core (Doc IR, translation, rendering and IR builder). Per-dialect (relax, tirx, ...) printer/builder pieces - # live under src/<dialect>/script/. The list below is intentionally explicit - # (not src/script/*.cc) so that any new file accidentally added under - # src/script/{printer,ir_builder}/<dialect>/ for a non-IR dialect is - # rejected at link time rather than silently included. + # Shared printer and TVMScript IR builder. Per-dialect printer/builder pieces + # live under src/<dialect>/script/. List shared files explicitly so a dialect + # implementation cannot be silently included through a recursive glob. src/script/ir_builder/base.cc src/script/ir_builder/ir.cc src/script/ir_builder/frame.cc - src/script/printer/config.cc - src/script/printer/script_printer.cc - src/script/printer/doc.cc - src/script/printer/v2/*.cc + src/printer/config.cc + src/printer/doc.cc + src/printer/doc_translator.cc + src/printer/entry.cc + src/printer/python_doc_printer.cc + src/printer/script_printer.cc src/relax/ir/*.cc src/relax/op/*.cc src/relax/analysis/*.cc diff --git a/docs/reference/api/python/index.rst b/docs/reference/api/python/index.rst index f81de243f9..234916315d 100644 --- a/docs/reference/api/python/index.rst +++ b/docs/reference/api/python/index.rst @@ -75,7 +75,7 @@ Python API script/script script/parser script/ir_builder - script/printer + printer .. toctree:: :maxdepth: 1 diff --git a/docs/reference/api/python/script/printer.rst b/docs/reference/api/python/printer.rst similarity index 92% rename from docs/reference/api/python/script/printer.rst rename to docs/reference/api/python/printer.rst index b5dddb4396..0194e3f676 100644 --- a/docs/reference/api/python/script/printer.rst +++ b/docs/reference/api/python/printer.rst @@ -15,11 +15,11 @@ specific language governing permissions and limitations under the License. -tvm.script.printer +tvm.printer ------------------ -tvm.script.printer +tvm.printer ****************** -.. automodule:: tvm.script.printer +.. automodule:: tvm.printer :members: :imported-members: diff --git a/include/tvm/script/printer/config.h b/include/tvm/printer/config.h similarity index 96% rename from include/tvm/script/printer/config.h rename to include/tvm/printer/config.h index 6057e23f14..fc2d1dad15 100644 --- a/include/tvm/script/printer/config.h +++ b/include/tvm/printer/config.h @@ -17,15 +17,15 @@ * under the License. */ /*! - * \file tvm/script/printer/config.h + * \file tvm/printer/config.h * \brief Configuration object for the TVMScript printer. * * Contains PrinterConfig / PrinterConfigNode, GetBuiltinKeywords, GetExtraConfig, * and RedirectedReprPrinterMethod. The translation entry points live in * doc_translator.h. */ -#ifndef TVM_SCRIPT_PRINTER_CONFIG_H_ -#define TVM_SCRIPT_PRINTER_CONFIG_H_ +#ifndef TVM_PRINTER_CONFIG_H_ +#define TVM_PRINTER_CONFIG_H_ #include <tvm/ffi/any.h> #include <tvm/ffi/container/array.h> @@ -145,10 +145,10 @@ class TVM_DLL PrinterConfig : public ffi::ObjectRef { * \brief Render an IR object for its registered repr callback. * * Tries to translate \p obj directly; on error falls back to a plain - * address string. Defined in src/script/printer/script_printer.cc so that + * address string. Defined in src/printer/script_printer.cc so that * <tvm/runtime/logging.h> is not pulled into this public header. */ TVM_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj); } // namespace tvm -#endif // TVM_SCRIPT_PRINTER_CONFIG_H_ +#endif // TVM_PRINTER_CONFIG_H_ diff --git a/include/tvm/script/printer/doc.h b/include/tvm/printer/doc.h similarity index 91% rename from include/tvm/script/printer/doc.h rename to include/tvm/printer/doc.h index 3a444ed753..fb6d2a31c4 100644 --- a/include/tvm/script/printer/doc.h +++ b/include/tvm/printer/doc.h @@ -16,35 +16,24 @@ * specific language governing permissions and limitations * under the License. */ -#ifndef TVM_SCRIPT_PRINTER_DOC_H_ -#define TVM_SCRIPT_PRINTER_DOC_H_ +#ifndef TVM_PRINTER_DOC_H_ +#define TVM_PRINTER_DOC_H_ #include <tvm/ffi/dtype.h> #include <tvm/ffi/reflection/access_path.h> #include <tvm/ffi/reflection/registry.h> #include <tvm/ir/expr.h> #include <tvm/ir/type.h> +#include <tvm/printer/config.h> #include <tvm/runtime/device_api.h> -#include <tvm/script/printer/config.h> #include <string> namespace tvm { -namespace script { namespace printer { using AccessPath = ffi::reflection::AccessPath; -// Forward declaration -class Doc; - -/*! - * \brief Convert Doc into Python script. - * \param doc Doc to be converted - * \param cfg The configuration of the printer - */ -TVM_DLL ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg); - /*! * \brief The base class of all Doc. * @@ -74,7 +63,7 @@ class DocNode : public ffi::Object { static constexpr bool _type_mutable = true; - TVM_FFI_DECLARE_OBJECT_INFO("script.printer.Doc", DocNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("printer.Doc", DocNode, ffi::Object); public: virtual ~DocNode() = default; @@ -135,7 +124,7 @@ class ExprDocNode : public DocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ExprDocNode>(); } - TVM_FFI_DECLARE_OBJECT_INFO("script.printer.ExprDoc", ExprDocNode, DocNode); + TVM_FFI_DECLARE_OBJECT_INFO("printer.ExprDoc", ExprDocNode, DocNode); }; /*! @@ -182,7 +171,7 @@ class StmtDocNode : public DocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<StmtDocNode>().def_rw("comment", &StmtDocNode::comment); } - TVM_FFI_DECLARE_OBJECT_INFO("script.printer.StmtDoc", StmtDocNode, DocNode); + TVM_FFI_DECLARE_OBJECT_INFO("printer.StmtDoc", StmtDocNode, DocNode); }; /*! @@ -213,7 +202,7 @@ class StmtBlockDocNode : public DocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<StmtBlockDocNode>().def_ro("stmts", &StmtBlockDocNode::stmts); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.StmtBlockDoc", StmtBlockDocNode, DocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.StmtBlockDoc", StmtBlockDocNode, DocNode); }; /*! @@ -252,7 +241,7 @@ class LiteralDocNode : public ExprDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<LiteralDocNode>().def_ro("value", &LiteralDocNode::value); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.LiteralDoc", LiteralDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.LiteralDoc", LiteralDocNode, ExprDocNode); }; /*! @@ -352,7 +341,7 @@ class ExprStringDocNode : public ExprDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ExprStringDocNode>().def_ro("value", &ExprStringDocNode::value); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ExprStringDoc", ExprStringDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ExprStringDoc", ExprStringDocNode, ExprDocNode); }; /*! @@ -386,7 +375,7 @@ class IdDocNode : public ExprDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<IdDocNode>().def_ro("name", &IdDocNode::name); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IdDoc", IdDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.IdDoc", IdDocNode, ExprDocNode); }; /*! @@ -423,7 +412,7 @@ class AttrAccessDocNode : public ExprDocNode { .def_ro("value", &AttrAccessDocNode::value) .def_ro("name", &AttrAccessDocNode::name); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.AttrAccessDoc", AttrAccessDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.AttrAccessDoc", AttrAccessDocNode, ExprDocNode); }; /*! @@ -466,7 +455,7 @@ class IndexDocNode : public ExprDocNode { .def_ro("value", &IndexDocNode::value) .def_ro("indices", &IndexDocNode::indices); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IndexDoc", IndexDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.IndexDoc", IndexDocNode, ExprDocNode); }; /*! @@ -514,7 +503,7 @@ class CallDocNode : public ExprDocNode { .def_ro("kwargs_keys", &CallDocNode::kwargs_keys) .def_ro("kwargs_values", &CallDocNode::kwargs_values); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.CallDoc", CallDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.CallDoc", CallDocNode, ExprDocNode); }; /*! @@ -596,7 +585,7 @@ class OperationDocNode : public ExprDocNode { .def_ro("kind", &OperationDocNode::kind) .def_ro("operands", &OperationDocNode::operands); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.OperationDoc", OperationDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.OperationDoc", OperationDocNode, ExprDocNode); }; /*! @@ -636,7 +625,7 @@ class LambdaDocNode : public ExprDocNode { .def_ro("args", &LambdaDocNode::args) .def_ro("body", &LambdaDocNode::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.LambdaDoc", LambdaDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.LambdaDoc", LambdaDocNode, ExprDocNode); }; /*! @@ -669,7 +658,7 @@ class TupleDocNode : public ExprDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<TupleDocNode>().def_ro("elements", &TupleDocNode::elements); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.TupleDoc", TupleDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.TupleDoc", TupleDocNode, ExprDocNode); }; /*! @@ -705,7 +694,7 @@ class ListDocNode : public ExprDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ListDocNode>().def_ro("elements", &ListDocNode::elements); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ListDoc", ListDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ListDoc", ListDocNode, ExprDocNode); }; /*! @@ -750,7 +739,7 @@ class DictDocNode : public ExprDocNode { .def_ro("keys", &DictDocNode::keys) .def_ro("values", &DictDocNode::values); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.DictDoc", DictDocNode, ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.DictDoc", DictDocNode, ExprDocNode); }; /*! @@ -796,7 +785,7 @@ class SliceDocNode : public DocNode { .def_ro("stop", &SliceDocNode::stop) .def_ro("step", &SliceDocNode::step); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.SliceDoc", SliceDocNode, DocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.SliceDoc", SliceDocNode, DocNode); }; /*! @@ -842,7 +831,7 @@ class AssignDocNode : public StmtDocNode { .def_ro("rhs", &AssignDocNode::rhs) .def_ro("annotation", &AssignDocNode::annotation); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.AssignDoc", AssignDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.AssignDoc", AssignDocNode, StmtDocNode); }; /*! @@ -884,7 +873,7 @@ class IfDocNode : public StmtDocNode { .def_ro("then_branch", &IfDocNode::then_branch) .def_ro("else_branch", &IfDocNode::else_branch); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IfDoc", IfDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.IfDoc", IfDocNode, StmtDocNode); }; /*! @@ -923,7 +912,7 @@ class WhileDocNode : public StmtDocNode { .def_ro("predicate", &WhileDocNode::predicate) .def_ro("body", &WhileDocNode::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.WhileDoc", WhileDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.WhileDoc", WhileDocNode, StmtDocNode); }; /*! @@ -954,7 +943,7 @@ class BreakDocNode : public StmtDocNode { refl::ObjectDef<BreakDocNode>(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.BreakDoc", BreakDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.BreakDoc", BreakDocNode, StmtDocNode); }; /*! @@ -983,7 +972,7 @@ class ContinueDocNode : public StmtDocNode { refl::ObjectDef<ContinueDocNode>(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ContinueDoc", ContinueDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ContinueDoc", ContinueDocNode, StmtDocNode); }; /*! @@ -1025,7 +1014,7 @@ class ForDocNode : public StmtDocNode { .def_ro("rhs", &ForDocNode::rhs) .def_ro("body", &ForDocNode::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ForDoc", ForDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ForDoc", ForDocNode, StmtDocNode); }; /*! @@ -1074,7 +1063,7 @@ class ScopeDocNode : public StmtDocNode { .def_ro("body", &ScopeDocNode::body) .def_ro("allow_concise_scoping", &ScopeDocNode::allow_concise_scoping); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ScopeDoc", ScopeDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ScopeDoc", ScopeDocNode, StmtDocNode); }; /*! @@ -1118,7 +1107,7 @@ class ExprStmtDocNode : public StmtDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ExprStmtDocNode>().def_ro("expr", &ExprStmtDocNode::expr); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ExprStmtDoc", ExprStmtDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ExprStmtDoc", ExprStmtDocNode, StmtDocNode); }; /*! @@ -1154,7 +1143,7 @@ class AssertDocNode : public StmtDocNode { .def_ro("test", &AssertDocNode::test) .def_ro("msg", &AssertDocNode::msg); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.AssertDoc", AssertDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.AssertDoc", AssertDocNode, StmtDocNode); }; /*! @@ -1187,7 +1176,7 @@ class ReturnDocNode : public StmtDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ReturnDocNode>().def_ro("value", &ReturnDocNode::value); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ReturnDoc", ReturnDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ReturnDoc", ReturnDocNode, StmtDocNode); }; /*! @@ -1247,7 +1236,7 @@ class FunctionDocNode : public StmtDocNode { .def_ro("body", &FunctionDocNode::body) .def_ro("type_params", &FunctionDocNode::type_params); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.FunctionDoc", FunctionDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.FunctionDoc", FunctionDocNode, StmtDocNode); }; /*! @@ -1293,7 +1282,7 @@ class ClassDocNode : public StmtDocNode { .def_ro("decorators", &ClassDocNode::decorators) .def_ro("body", &ClassDocNode::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ClassDoc", ClassDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.ClassDoc", ClassDocNode, StmtDocNode); }; /*! @@ -1324,7 +1313,7 @@ class CommentDocNode : public StmtDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<CommentDocNode>(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.CommentDoc", CommentDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.CommentDoc", CommentDocNode, StmtDocNode); }; /*! @@ -1349,7 +1338,7 @@ class DocStringDocNode : public StmtDocNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<DocStringDocNode>(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.DocStringDoc", DocStringDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.DocStringDoc", DocStringDocNode, StmtDocNode); }; /*! @@ -1391,7 +1380,7 @@ class OpCallDocNode : public StmtDocNode { .def_ro("dispatch", &OpCallDocNode::dispatch); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.OpCallDoc", OpCallDocNode, StmtDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.OpCallDoc", OpCallDocNode, StmtDocNode); }; /*! @@ -1415,8 +1404,6 @@ class OpCallDoc : public StmtDoc { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpCallDoc, StmtDoc, OpCallDocNode); }; -namespace v2 { - /*! * \brief A reference to a canonical script namespace. * @@ -1431,8 +1418,7 @@ class NamespaceDocNode : public ExprDocNode { ffi::reflection::ObjectDef<NamespaceDocNode>().def_ro("canonical_name", &NamespaceDocNode::canonical_name); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.v2.NamespaceDoc", NamespaceDocNode, - ExprDocNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("printer.NamespaceDoc", NamespaceDocNode, ExprDocNode); }; /*! @@ -1466,10 +1452,7 @@ TVM_DLL ffi::String DocToPythonScript(Doc doc, const PrinterConfig& config); TVM_DLL ffi::String DocToPythonScriptWithHeader(Doc doc, const PrinterConfig& config, const ffi::Array<ffi::Any>& header); -} // namespace v2 - } // namespace printer -} // namespace script } // namespace tvm -#endif // TVM_SCRIPT_PRINTER_DOC_H_ +#endif // TVM_PRINTER_DOC_H_ diff --git a/include/tvm/script/printer/doc_translator.h b/include/tvm/printer/doc_translator.h similarity index 97% rename from include/tvm/script/printer/doc_translator.h rename to include/tvm/printer/doc_translator.h index 7e3c6701f6..2bef66bf12 100644 --- a/include/tvm/script/printer/doc_translator.h +++ b/include/tvm/printer/doc_translator.h @@ -15,8 +15,8 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#ifndef TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_H_ -#define TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_H_ +#ifndef TVM_PRINTER_DOC_TRANSLATOR_H_ +#define TVM_PRINTER_DOC_TRANSLATOR_H_ #include <tvm/ffi/container/dict.h> #include <tvm/ffi/container/list.h> @@ -24,7 +24,7 @@ #include <tvm/ffi/expected.h> #include <tvm/ffi/reflection/accessor.h> #include <tvm/ir/expr.h> -#include <tvm/script/printer/doc.h> +#include <tvm/printer/doc.h> #include <optional> #include <utility> @@ -40,9 +40,7 @@ namespace tvm { TVM_DLL std::string Script(const ffi::ObjectRef& node, const ffi::Optional<PrinterConfig>& config = std::nullopt); -namespace script { namespace printer { -namespace v2 { class DocTranslatorObj; @@ -53,7 +51,7 @@ inline constexpr const char* kDocTranslate = "__tvm_doc_translate__"; /*! * \brief Op-specific Call translation hook with the FDocTranslate contract. */ -inline constexpr const char* kOpCallTranslate = "script_printer.call_translate_v2"; +inline constexpr const char* kOpCallTranslate = "script_printer.call_translate"; /*! * \brief Source-Type hook receiving the original TensorLoad and optional binder. */ @@ -167,7 +165,7 @@ class DocTranslatorObj : public ffi::Object { T TranslateAs(ffi::AnyView value, const ffi::Optional<Var> bind_var = std::nullopt) { auto result = Translate(value, bind_var); TVM_FFI_CHECK(result.has_value(), ValueError) - << "v2 expression context requires a returned Doc"; + << "printer expression context requires a returned Doc"; return result.value().as_or_throw<T>(); } /*! @@ -289,7 +287,7 @@ class DocTranslatorObj : public ffi::Object { try { std::forward<Callback>(callback)(); TVM_FFI_CHECK(doc_scopes_.size() == depth + 1, ValueError) - << "v2 unbalanced Doc collection scopes"; + << "printer unbalanced Doc collection scopes"; return EndDocScope(); } catch (...) { while (doc_scopes_.size() > depth) doc_scopes_.pop_back(); @@ -324,7 +322,7 @@ class DocTranslatorObj : public ffi::Object { const ffi::Object* destination = nullptr); static constexpr bool _type_mutable = true; - TVM_FFI_DECLARE_OBJECT_INFO("script.printer.v2.DocTranslator", DocTranslatorObj, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("printer.DocTranslator", DocTranslatorObj, ffi::Object); protected: /*! @@ -406,9 +404,7 @@ TVM_DLL Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_or */ TVM_DLL ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config); -} // namespace v2 } // namespace printer -} // namespace script } // namespace tvm -#endif // TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_H_ +#endif // TVM_PRINTER_DOC_TRANSLATOR_H_ diff --git a/include/tvm/script/printer/printer.h b/include/tvm/printer/printer.h similarity index 81% rename from include/tvm/script/printer/printer.h rename to include/tvm/printer/printer.h index bab8adca96..1a6d02b08e 100644 --- a/include/tvm/script/printer/printer.h +++ b/include/tvm/printer/printer.h @@ -17,12 +17,12 @@ * under the License. */ /*! - * \file tvm/script/printer/printer.h + * \file tvm/printer/printer.h * \brief Compatibility include for the TVMScript translation entry points. */ -#ifndef TVM_SCRIPT_PRINTER_PRINTER_H_ -#define TVM_SCRIPT_PRINTER_PRINTER_H_ +#ifndef TVM_PRINTER_PRINTER_H_ +#define TVM_PRINTER_PRINTER_H_ -#include <tvm/script/printer/doc_translator.h> +#include <tvm/printer/doc_translator.h> -#endif // TVM_SCRIPT_PRINTER_PRINTER_H_ +#endif // TVM_PRINTER_PRINTER_H_ diff --git a/include/tvm/script/printer/v2/doc.h b/include/tvm/script/printer/v2/doc.h deleted file mode 100644 index c3fb658a35..0000000000 --- a/include/tvm/script/printer/v2/doc.h +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#ifndef TVM_SCRIPT_PRINTER_V2_DOC_H_ -#define TVM_SCRIPT_PRINTER_V2_DOC_H_ - -#include <tvm/script/printer/doc.h> - -#endif // TVM_SCRIPT_PRINTER_V2_DOC_H_ diff --git a/include/tvm/script/printer/v2/doc_translator.h b/include/tvm/script/printer/v2/doc_translator.h deleted file mode 100644 index c2a4c1a161..0000000000 --- a/include/tvm/script/printer/v2/doc_translator.h +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#ifndef TVM_SCRIPT_PRINTER_V2_DOC_TRANSLATOR_H_ -#define TVM_SCRIPT_PRINTER_V2_DOC_TRANSLATOR_H_ - -#include <tvm/script/printer/doc_translator.h> - -#endif // TVM_SCRIPT_PRINTER_V2_DOC_TRANSLATOR_H_ diff --git a/python/tvm/script/printer/__init__.py b/python/tvm/printer/__init__.py similarity index 100% rename from python/tvm/script/printer/__init__.py rename to python/tvm/printer/__init__.py diff --git a/python/tvm/script/printer/_ffi_api.py b/python/tvm/printer/_ffi_api.py similarity index 86% rename from python/tvm/script/printer/_ffi_api.py rename to python/tvm/printer/_ffi_api.py index b366fd823b..89213734be 100644 --- a/python/tvm/script/printer/_ffi_api.py +++ b/python/tvm/printer/_ffi_api.py @@ -14,8 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""FFI APIs for tvm.script.printer""" +"""FFI APIs for tvm.printer""" import tvm_ffi -tvm_ffi.init_ffi_api("script.printer", __name__) # pylint: disable=protected-access +tvm_ffi.init_ffi_api("printer", __name__) # pylint: disable=protected-access diff --git a/python/tvm/script/printer/doc.py b/python/tvm/printer/doc.py similarity index 91% rename from python/tvm/script/printer/doc.py rename to python/tvm/printer/doc.py index 2e64649b8d..57542494c9 100644 --- a/python/tvm/script/printer/doc.py +++ b/python/tvm/printer/doc.py @@ -29,12 +29,12 @@ from tvm.runtime import Object from . import _ffi_api -@register_object("script.printer.Doc") +@register_object("printer.Doc") class Doc(Object): """Base class of all Docs""" -@register_object("script.printer.ExprDoc") +@register_object("printer.ExprDoc") class ExprDoc(Doc): """Base class of all expression Docs""" @@ -104,12 +104,12 @@ class ExprDoc(Doc): raise RuntimeError(f"{self.__class__} cannot be used as iterable.") -@register_object("script.printer.StmtDoc") +@register_object("printer.StmtDoc") class StmtDoc(Doc): """Base class of statement doc""" -@register_object("script.printer.StmtBlockDoc") +@register_object("printer.StmtBlockDoc") class StmtBlockDoc(Doc): """The container doc that holds a list of StmtDoc. @@ -123,7 +123,7 @@ class StmtBlockDoc(Doc): self.__init_handle_by_constructor__(_ffi_api.StmtBlockDoc, stmts) # type: ignore # pylint: disable=no-member -@register_object("script.printer.LiteralDoc") +@register_object("printer.LiteralDoc") class LiteralDoc(ExprDoc): """Doc that represents literal value""" @@ -164,7 +164,7 @@ class LiteralDoc(ExprDoc): raise TypeError(f"Unsupported type {type(value)} for LiteralDoc") -@register_object("script.printer.IdDoc") +@register_object("printer.IdDoc") class IdDoc(ExprDoc): """Doc that represents identifier""" @@ -174,7 +174,7 @@ class IdDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.IdDoc, name) # type: ignore # pylint: disable=no-member -@register_object("script.printer.AttrAccessDoc") +@register_object("printer.AttrAccessDoc") class AttrAccessDoc(ExprDoc): """Doc that represents attribute access on an expression""" @@ -185,7 +185,7 @@ class AttrAccessDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.AttrAccessDoc, value, name) # type: ignore # pylint: disable=no-member -@register_object("script.printer.IndexDoc") +@register_object("printer.IndexDoc") class IndexDoc(ExprDoc): """Doc that represents index access on an expression""" @@ -196,7 +196,7 @@ class IndexDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.IndexDoc, value, indices) # type: ignore # pylint: disable=no-member -@register_object("script.printer.CallDoc") +@register_object("printer.CallDoc") class CallDoc(ExprDoc): """Doc that represents function call""" @@ -222,7 +222,7 @@ class OperationKind(IntEnum): """ This enum represents the kind of operation (operator) in OperationDoc - It's mirrored from OperationDocNode::Kind at include/tvm/script/printer/doc.h + It's mirrored from OperationDocNode::Kind at include/tvm/printer/doc.h """ # The name convention follows https://docs.python.org/3/library/ast.html @@ -265,7 +265,7 @@ class OperationKind(IntEnum): # pylint: enable=invalid-name -@register_object("script.printer.OperationDoc") +@register_object("printer.OperationDoc") class OperationDoc(ExprDoc): """ Doc that represents operation @@ -281,7 +281,7 @@ class OperationDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.OperationDoc, kind, operands) # type: ignore # pylint: disable=no-member -@register_object("script.printer.LambdaDoc") +@register_object("printer.LambdaDoc") class LambdaDoc(ExprDoc): """Doc that represents lambda function""" @@ -292,7 +292,7 @@ class LambdaDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.LambdaDoc, args, body) # type: ignore # pylint: disable=no-member -@register_object("script.printer.TupleDoc") +@register_object("printer.TupleDoc") class TupleDoc(ExprDoc): """Doc that represents tuple literal""" @@ -302,7 +302,7 @@ class TupleDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.TupleDoc, elements) # type: ignore # pylint: disable=no-member -@register_object("script.printer.ListDoc") +@register_object("printer.ListDoc") class ListDoc(ExprDoc): """Doc that represents list literal""" @@ -312,7 +312,7 @@ class ListDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.ListDoc, elements) # type: ignore # pylint: disable=no-member -@register_object("script.printer.DictDoc") +@register_object("printer.DictDoc") class DictDoc(ExprDoc): """Doc that represents dict literal""" @@ -325,7 +325,7 @@ class DictDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.DictDoc, keys, values) # type: ignore # pylint: disable=no-member -@register_object("script.printer.SliceDoc") +@register_object("printer.SliceDoc") class SliceDoc(ExprDoc): """ Doc that represents slice in Index expression @@ -346,7 +346,7 @@ class SliceDoc(ExprDoc): self.__init_handle_by_constructor__(_ffi_api.SliceDoc, start, stop, step) # type: ignore # pylint: disable=no-member -@register_object("script.printer.AssignDoc") +@register_object("printer.AssignDoc") class AssignDoc(StmtDoc): """Doc that represents assign statement.""" @@ -363,7 +363,7 @@ class AssignDoc(StmtDoc): ) -@register_object("script.printer.IfDoc") +@register_object("printer.IfDoc") class IfDoc(StmtDoc): """Doc that represent if-then-else statement.""" @@ -380,7 +380,7 @@ class IfDoc(StmtDoc): ) -@register_object("script.printer.WhileDoc") +@register_object("printer.WhileDoc") class WhileDoc(StmtDoc): """Doc that represents while statement.""" @@ -391,7 +391,7 @@ class WhileDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.WhileDoc, predicate, body) # type: ignore # pylint: disable=no-member -@register_object("script.printer.ForDoc") +@register_object("printer.ForDoc") class ForDoc(StmtDoc): """Doc that represents for statement.""" @@ -403,7 +403,7 @@ class ForDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.ForDoc, lhs, rhs, body) # type: ignore # pylint: disable=no-member -@register_object("script.printer.ScopeDoc") +@register_object("printer.ScopeDoc") class ScopeDoc(StmtDoc): """ Doc that represents special scopes. @@ -422,7 +422,7 @@ class ScopeDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.ScopeDoc, lhs, rhs, body) # type: ignore # pylint: disable=no-member -@register_object("script.printer.ExprStmtDoc") +@register_object("printer.ExprStmtDoc") class ExprStmtDoc(StmtDoc): """Doc that represents an expression as statement.""" @@ -432,7 +432,7 @@ class ExprStmtDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.ExprStmtDoc, expr) # type: ignore # pylint: disable=no-member -@register_object("script.printer.AssertDoc") +@register_object("printer.AssertDoc") class AssertDoc(StmtDoc): """Doc that represents assert statement.""" @@ -443,7 +443,7 @@ class AssertDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.AssertDoc, test, msg) # type: ignore # pylint: disable=no-member -@register_object("script.printer.ReturnDoc") +@register_object("printer.ReturnDoc") class ReturnDoc(StmtDoc): """Doc that represents return statement.""" @@ -453,7 +453,7 @@ class ReturnDoc(StmtDoc): self.__init_handle_by_constructor__(_ffi_api.ReturnDoc, value) # type: ignore # pylint: disable=no-member -@register_object("script.printer.FunctionDoc") +@register_object("printer.FunctionDoc") class FunctionDoc(StmtDoc): """Doc that represents function definition.""" @@ -484,7 +484,7 @@ class FunctionDoc(StmtDoc): ) -@register_object("script.printer.ClassDoc") +@register_object("printer.ClassDoc") class ClassDoc(StmtDoc): """Doc that represents class definition.""" @@ -501,7 +501,7 @@ class ClassDoc(StmtDoc): ) -@register_object("script.printer.CommentDoc") +@register_object("printer.CommentDoc") class CommentDoc(StmtDoc): """Doc that represents comment.""" @@ -512,7 +512,7 @@ class CommentDoc(StmtDoc): ) -@register_object("script.printer.DocStringDoc") +@register_object("printer.DocStringDoc") class DocStringDoc(StmtDoc): """Doc that represents docstring.""" diff --git a/python/tvm/script/printer/doc_printer.py b/python/tvm/printer/doc_printer.py similarity index 100% rename from python/tvm/script/printer/doc_printer.py rename to python/tvm/printer/doc_printer.py diff --git a/python/tvm/script/printer/highlight.py b/python/tvm/printer/highlight.py similarity index 100% rename from python/tvm/script/printer/highlight.py rename to python/tvm/printer/highlight.py diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py index f1acdeb169..4e08656d04 100644 --- a/python/tvm/relax/base_py_module.py +++ b/python/tvm/relax/base_py_module.py @@ -646,7 +646,7 @@ class BasePyModule: This method extends the standard IRModule show() method to handle Python functions stored in the IRModule's pyfuncs attribute. """ - from tvm.script.printer.highlight import cprint # pylint: disable=import-outside-toplevel + from tvm.printer.highlight import cprint # pylint: disable=import-outside-toplevel if black_format is None: env = os.environ.get("TVM_BLACK_FORMAT") diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py index dfed7325f2..40b3575fd7 100644 --- a/python/tvm/runtime/script_printer.py +++ b/python/tvm/runtime/script_printer.py @@ -115,7 +115,7 @@ def _script(obj: Object, config: PrinterConfig) -> str: def _relax_script(obj: Object, config: PrinterConfig) -> str: - func = get_global_func("script.printer.ReprPrintRelax") + func = get_global_func("printer.ReprPrintRelax") return func(obj, config) @@ -203,7 +203,7 @@ class Scriptable: merged_extra: dict = {} if extra_config is not None: merged_extra.update(extra_config) - # Keep the historical interactive display header. Direct v2 calls + # Keep the historical interactive display header. Direct printer calls # retain executable imports for standalone parser round-trips. merged_extra.setdefault("script.comment_imports", True) merged_extra.setdefault("script.legacy_display", True) @@ -341,7 +341,7 @@ class Scriptable: ---------- style : str, optional Pygmentize printing style, auto-detected if None. See - `tvm.script.printer.highlight.cprint` for more details. + `tvm.printer.highlight.cprint` for more details. black_format: Optional[bool] @@ -403,7 +403,7 @@ class Scriptable: Object to be annotated """ - from tvm.script.printer.highlight import cprint # pylint: disable=import-outside-toplevel + from tvm.printer.highlight import cprint # pylint: disable=import-outside-toplevel if black_format is None: env = os.environ.get("TVM_BLACK_FORMAT") diff --git a/python/tvm/s_tir/schedule/trace.py b/python/tvm/s_tir/schedule/trace.py index 3316d67912..22c253cc76 100644 --- a/python/tvm/s_tir/schedule/trace.py +++ b/python/tvm/s_tir/schedule/trace.py @@ -278,7 +278,7 @@ class Trace(Object): style : str, optional Pygmentize printing style, auto-detected if None. See - `tvm.script.printer.highlight.cprint` for more details. + `tvm.printer.highlight.cprint` for more details. black_format: bool @@ -286,7 +286,7 @@ class Trace(Object): If None, determine based on the "TVM_BLACK_FORMAT" environment variable. """ - from tvm.script.printer.highlight import ( # pylint: disable=import-outside-toplevel + from tvm.printer.highlight import ( # pylint: disable=import-outside-toplevel cprint, ) diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 4c1d9da2fc..58dedf97e5 100644 --- a/src/ir/expr.cc +++ b/src/ir/expr.cc @@ -1247,7 +1247,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ss << ref; return ss.str(); }); - // Note: kRepr for GlobalVarNode is registered in script/printer/script_printer.cc. + // Note: kRepr for GlobalVarNode is registered in printer/script_printer.cc. } } // namespace tvm diff --git a/src/ir/prim/expr.cc b/src/ir/prim/expr.cc index b6e67282df..a24a99304d 100644 --- a/src/ir/prim/expr.cc +++ b/src/ir/prim/expr.cc @@ -439,7 +439,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("prim.convert", [](ffi::Variant<PrimExpr, ffi::Array<PrimExpr>> expr) { return expr; }); - // Note: kRepr for VarNode is registered in script/printer/script_printer.cc. + // Note: kRepr for VarNode is registered in printer/script_printer.cc. } // Cast diff --git a/src/script/printer/config.cc b/src/printer/config.cc similarity index 98% rename from src/script/printer/config.cc rename to src/printer/config.cc index 993f22dd4f..28a5e8f5fe 100644 --- a/src/script/printer/config.cc +++ b/src/printer/config.cc @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -#include <tvm/script/printer/config.h> +#include <tvm/printer/config.h> #include <algorithm> #include <cctype> @@ -51,7 +51,6 @@ bool IsIdentifier(const std::string& name) { } // namespace -namespace script { namespace printer { void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix) { @@ -60,7 +59,6 @@ void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_pr } } // namespace printer -} // namespace script PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) { ffi::ObjectPtr<PrinterConfigNode> n = ffi::make_object<PrinterConfigNode>(); diff --git a/src/script/printer/dialect_prefix.h b/src/printer/dialect_prefix.h similarity index 86% rename from src/script/printer/dialect_prefix.h rename to src/printer/dialect_prefix.h index 5aaae11293..9be6410218 100644 --- a/src/script/printer/dialect_prefix.h +++ b/src/printer/dialect_prefix.h @@ -16,13 +16,12 @@ * specific language governing permissions and limitations * under the License. */ -#ifndef TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ -#define TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ +#ifndef TVM_PRINTER_DIALECT_PREFIX_H_ +#define TVM_PRINTER_DIALECT_PREFIX_H_ #include <tvm/ffi/string.h> namespace tvm { -namespace script { namespace printer { // Register during dialect static initialization so configuration can validate @@ -30,7 +29,6 @@ namespace printer { void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix); } // namespace printer -} // namespace script } // namespace tvm -#endif // TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_ +#endif // TVM_PRINTER_DIALECT_PREFIX_H_ diff --git a/src/script/printer/doc.cc b/src/printer/doc.cc similarity index 80% rename from src/script/printer/doc.cc rename to src/printer/doc.cc index 941cbd7b38..f1e1a6deec 100644 --- a/src/script/printer/doc.cc +++ b/src/printer/doc.cc @@ -21,10 +21,9 @@ #include <tvm/ffi/error.h> #include <tvm/ffi/function.h> #include <tvm/ffi/reflection/registry.h> -#include <tvm/script/printer/doc.h> +#include <tvm/printer/doc.h> namespace tvm { -namespace script { namespace printer { TVM_FFI_STATIC_INIT_BLOCK() { @@ -309,16 +308,16 @@ OpCallDoc::OpCallDoc(ExprDoc callee, ffi::Array<Doc> args, ffi::Optional<DictDoc TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "script.printer.DocSetSourcePaths", + "printer.DocSetSourcePaths", [](Doc doc, ffi::Array<AccessPath> source_paths) { doc->source_paths = source_paths; }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def_method("script.printer.ExprDocAttr", &ExprDocNode::Attr) - .def_method("script.printer.ExprDocIndex", &ExprDocNode::operator[]) - .def_method("script.printer.ExprDocCall", + .def_method("printer.ExprDocAttr", &ExprDocNode::Attr) + .def_method("printer.ExprDocIndex", &ExprDocNode::operator[]) + .def_method("printer.ExprDocCall", [](ExprDoc doc, ffi::Array<ExprDoc> args, ffi::Array<ffi::String> kwargs_keys, ffi::Array<ExprDoc> kwargs_values) { return doc->Call(args, kwargs_keys, kwargs_values); @@ -328,139 +327,137 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "script.printer.StmtDocSetComment", + "printer.StmtDocSetComment", [](StmtDoc doc, ffi::Optional<ffi::String> comment) { doc->comment = comment; }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.StmtBlockDoc", + refl::GlobalDef().def("printer.StmtBlockDoc", [](ffi::Array<StmtDoc> stmts) { return StmtBlockDoc(stmts); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def("script.printer.LiteralDocNone", LiteralDoc::None) - .def("script.printer.LiteralDocInt", + .def("printer.LiteralDocNone", LiteralDoc::None) + .def("printer.LiteralDocInt", static_cast<LiteralDoc (*)(int64_t, const ffi::Optional<AccessPath>&)>(LiteralDoc::Int)) - .def("script.printer.LiteralDocBoolean", LiteralDoc::Boolean) - .def("script.printer.LiteralDocFloat", LiteralDoc::Float) - .def("script.printer.LiteralDocStr", LiteralDoc::Str); + .def("printer.LiteralDocBoolean", LiteralDoc::Boolean) + .def("printer.LiteralDocFloat", LiteralDoc::Float) + .def("printer.LiteralDocStr", LiteralDoc::Str); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.IdDoc", [](ffi::String name) { return IdDoc(name); }); + refl::GlobalDef().def("printer.IdDoc", [](ffi::String name) { return IdDoc(name); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.AttrAccessDoc", + refl::GlobalDef().def("printer.AttrAccessDoc", [](ExprDoc value, ffi::String attr) { return AttrAccessDoc(value, attr); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.IndexDoc", [](ExprDoc value, ffi::Array<Doc> indices) { + refl::GlobalDef().def("printer.IndexDoc", [](ExprDoc value, ffi::Array<Doc> indices) { return IndexDoc(value, indices); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.CallDoc", [](ExprDoc callee, // - ffi::Array<ExprDoc> args, // - ffi::Array<ffi::String> kwargs_keys, // - ffi::Array<ExprDoc> kwargs_values) { + refl::GlobalDef().def("printer.CallDoc", [](ExprDoc callee, // + ffi::Array<ExprDoc> args, // + ffi::Array<ffi::String> kwargs_keys, // + ffi::Array<ExprDoc> kwargs_values) { return CallDoc(callee, args, kwargs_keys, kwargs_values); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.OperationDoc", - [](int32_t kind, ffi::Array<ExprDoc> operands) { - return OperationDoc(OperationDocNode::Kind(kind), operands); - }); + refl::GlobalDef().def("printer.OperationDoc", [](int32_t kind, ffi::Array<ExprDoc> operands) { + return OperationDoc(OperationDocNode::Kind(kind), operands); + }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.LambdaDoc", + refl::GlobalDef().def("printer.LambdaDoc", [](ffi::Array<IdDoc> args, ExprDoc body) { return LambdaDoc(args, body); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.TupleDoc", + refl::GlobalDef().def("printer.TupleDoc", [](ffi::Array<ExprDoc> elements) { return TupleDoc(elements); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ListDoc", + refl::GlobalDef().def("printer.ListDoc", [](ffi::Array<ExprDoc> elements) { return ListDoc(elements); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "script.printer.DictDoc", + "printer.DictDoc", [](ffi::Array<ExprDoc> keys, ffi::Array<ExprDoc> values) { return DictDoc(keys, values); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.SliceDoc", + refl::GlobalDef().def("printer.SliceDoc", [](ffi::Optional<ExprDoc> start, ffi::Optional<ExprDoc> stop, ffi::Optional<ExprDoc> step) { return SliceDoc(start, stop, step); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.AssignDoc", [](ExprDoc lhs, ffi::Optional<ExprDoc> rhs, - ffi::Optional<ExprDoc> annotation) { + refl::GlobalDef().def("printer.AssignDoc", [](ExprDoc lhs, ffi::Optional<ExprDoc> rhs, + ffi::Optional<ExprDoc> annotation) { return AssignDoc(lhs, rhs, annotation); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "script.printer.IfDoc", - [](ExprDoc predicate, ffi::Array<StmtDoc> then_branch, ffi::Array<StmtDoc> else_branch) { - return IfDoc(predicate, then_branch, else_branch); - }); + refl::GlobalDef().def("printer.IfDoc", [](ExprDoc predicate, ffi::Array<StmtDoc> then_branch, + ffi::Array<StmtDoc> else_branch) { + return IfDoc(predicate, then_branch, else_branch); + }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.WhileDoc", [](ExprDoc predicate, ffi::Array<StmtDoc> body) { + refl::GlobalDef().def("printer.WhileDoc", [](ExprDoc predicate, ffi::Array<StmtDoc> body) { return WhileDoc(predicate, body); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.BreakDoc", []() { return BreakDoc(); }); + refl::GlobalDef().def("printer.BreakDoc", []() { return BreakDoc(); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ContinueDoc", []() { return ContinueDoc(); }); + refl::GlobalDef().def("printer.ContinueDoc", []() { return ContinueDoc(); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "script.printer.ForDoc", - [](ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body) { return ForDoc(lhs, rhs, body); }); + refl::GlobalDef().def("printer.ForDoc", [](ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body) { + return ForDoc(lhs, rhs, body); + }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ScopeDoc", + refl::GlobalDef().def("printer.ScopeDoc", [](ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body) { return ScopeDoc(lhs, rhs, body); }); @@ -468,36 +465,34 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ExprStmtDoc", - [](ExprDoc expr) { return ExprStmtDoc(expr); }); + refl::GlobalDef().def("printer.ExprStmtDoc", [](ExprDoc expr) { return ExprStmtDoc(expr); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "script.printer.AssertDoc", + "printer.AssertDoc", [](ExprDoc test, ffi::Optional<ExprDoc> msg = std::nullopt) { return AssertDoc(test, msg); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ReturnDoc", [](ExprDoc value) { return ReturnDoc(value); }); + refl::GlobalDef().def("printer.ReturnDoc", [](ExprDoc value) { return ReturnDoc(value); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.FunctionDoc", - [](IdDoc name, ffi::Array<AssignDoc> args, ffi::Array<ExprDoc> decorators, - ffi::Optional<ExprDoc> return_type, ffi::Array<StmtDoc> body, - ffi::Array<Doc> type_params = {}) { - return FunctionDoc(name, args, decorators, return_type, body, - type_params); - }); + refl::GlobalDef().def( + "printer.FunctionDoc", [](IdDoc name, ffi::Array<AssignDoc> args, + ffi::Array<ExprDoc> decorators, ffi::Optional<ExprDoc> return_type, + ffi::Array<StmtDoc> body, ffi::Array<Doc> type_params = {}) { + return FunctionDoc(name, args, decorators, return_type, body, type_params); + }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.ClassDoc", + refl::GlobalDef().def("printer.ClassDoc", [](IdDoc name, ffi::Array<ExprDoc> decorators, ffi::Array<StmtDoc> body) { return ClassDoc(name, decorators, body); }); @@ -505,25 +500,32 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.CommentDoc", + refl::GlobalDef().def("printer.CommentDoc", [](ffi::String comment) { return CommentDoc(comment); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.DocStringDoc", + refl::GlobalDef().def("printer.DocStringDoc", [](ffi::String docs) { return DocStringDoc(docs); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.OpCallDoc", + refl::GlobalDef().def("printer.OpCallDoc", [](ExprDoc callee, ffi::Array<Doc> args, DictDoc workspace, DictDoc config, ffi::Optional<ExprDoc> dispatch) { return OpCallDoc(callee, args, workspace, config, dispatch); }); } +NamespaceDoc::NamespaceDoc(ffi::String canonical_name) { + auto node = ffi::make_object<NamespaceDocNode>(); + node->canonical_name = std::move(canonical_name); + data_ = std::move(node); +} + +TVM_FFI_STATIC_INIT_BLOCK() { NamespaceDocNode::RegisterReflection(); } + } // namespace printer -} // namespace script } // namespace tvm diff --git a/src/script/printer/v2/doc_translator.cc b/src/printer/doc_translator.cc similarity index 98% rename from src/script/printer/v2/doc_translator.cc rename to src/printer/doc_translator.cc index 0e927cd8eb..b4e7e52a1d 100644 --- a/src/script/printer/v2/doc_translator.cc +++ b/src/printer/doc_translator.cc @@ -28,6 +28,8 @@ #include <tvm/ir/prim/op.h> #include <tvm/ir/prim/vector_expr.h> #include <tvm/ir/type.h> +#include <tvm/printer/doc.h> +#include <tvm/printer/doc_translator.h> #include <tvm/relax/attrs/op.h> #include <tvm/relax/block_builder.h> #include <tvm/relax/distributed/type.h> @@ -37,8 +39,6 @@ #include <tvm/relax/type.h> #include <tvm/runtime/tensor.h> #include <tvm/s_tir/stmt.h> -#include <tvm/script/printer/doc.h> -#include <tvm/script/printer/doc_translator.h> #include <tvm/sym/analyzer.h> #include <tvm/target/target.h> #include <tvm/te/operation.h> @@ -69,9 +69,7 @@ #include <vector> namespace tvm { -namespace script { namespace printer { -namespace v2 { ffi::Optional<ExprDoc> DocTranslatorObj::DefaultTranslate(ffi::AnyView value, const ffi::Object* destination) { @@ -91,13 +89,14 @@ ffi::Optional<ExprDoc> DocTranslatorObj::DefaultTranslate(ffi::AnyView value, .value(); } TVM_FFI_THROW(TypeError) << (attr.type_index() == ffi::TypeIndex::kTVMFFINone - ? std::string("v2 has no Doc hook for ") + value.GetTypeKey() - : "v2 Doc hook must be a native pointer or ffi.Function"); + ? std::string("printer has no Doc hook for ") + + value.GetTypeKey() + : "printer Doc hook must be a native pointer or ffi.Function"); } void DocTranslatorObj::BeginDocScope() { doc_scopes_.push_back(ffi::List<Doc>()); } ffi::List<Doc> DocTranslatorObj::CurrentScopeDocs() const { - TVM_FFI_CHECK(!doc_scopes_.empty(), ValueError) << "v2 emission requires a collection scope"; + TVM_FFI_CHECK(!doc_scopes_.empty(), ValueError) << "printer emission requires a collection scope"; return doc_scopes_.back(); } void DocTranslatorObj::Emit(const Doc& doc, ffi::Optional<ffi::ObjectRef> origin) { @@ -339,7 +338,8 @@ ExprDoc TensorElementDoc(const uint8_t* bytes, DLDataType dtype) { read(&value); return LiteralDoc::Boolean(value != 0, std::nullopt); } - TVM_FFI_THROW(TypeError) << "v2 inline tensor constant has unsupported dtype " << DType(dtype); + TVM_FFI_THROW(TypeError) << "printer inline tensor constant has unsupported dtype " + << DType(dtype); } ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* constant) { if (auto dtype = constant->value.as<DLDataType>()) { @@ -375,13 +375,13 @@ ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* const data->shape[i] >= 0 && (count == 0 || static_cast<size_t>(data->shape[i]) <= std::numeric_limits<size_t>::max() / count), TypeError) - << "v2 tensor constant shape overflows host size"; + << "printer tensor constant shape overflows host size"; count *= static_cast<size_t>(data->shape[i]); } const size_t element_size = dtype.bits / 8; - TVM_FFI_CHECK(element_size != 0, TypeError) << "v2 tensor constant element size is zero"; + TVM_FFI_CHECK(element_size != 0, TypeError) << "printer tensor constant element size is zero"; TVM_FFI_CHECK(count <= std::numeric_limits<size_t>::max() / element_size, TypeError) - << "v2 tensor constant byte size overflows host size"; + << "printer tensor constant byte size overflows host size"; std::vector<uint8_t> bytes(count * element_size); if (!bytes.empty()) data.CopyToBytes(bytes.data(), bytes.size()); size_t offset = 0; @@ -1677,7 +1677,7 @@ ExprDoc TranslateLet(DocTranslatorObj* d, const prim::LetNode* node) { } ExprDoc TranslateCommReducer(DocTranslatorObj* d, const te::CommReducerNode* reducer) { TVM_FFI_CHECK(reducer->lhs.size() == reducer->rhs.size(), TypeError) - << "v2 comm reducer arguments must be paired"; + << "printer comm reducer arguments must be paired"; ffi::Array<IdDoc> args; ffi::Array<ExprDoc> result; { @@ -1832,7 +1832,7 @@ void EmitFor(DocTranslatorObj* d, const tirx::ForNode* loop) { method = "thread_binding"; break; default: - TVM_FFI_THROW(TypeError) << "v2 unknown loop kind"; + TVM_FFI_THROW(TypeError) << "printer unknown loop kind"; } ExprDoc min = d->TranslateAs(loop->min); ExprDoc extent = d->TranslateAs(loop->extent); @@ -1848,7 +1848,7 @@ void EmitFor(DocTranslatorObj* d, const tirx::ForNode* loop) { ffi::Array<ExprDoc> values; if (loop->kind == tirx::ForKind::kThreadBinding) { TVM_FFI_CHECK(loop->thread_binding.has_value(), TypeError) - << "v2 thread-binding loop lacks thread tag"; + << "printer thread-binding loop lacks thread tag"; keys.push_back("thread"); values.push_back(LiteralDoc::Str(loop->thread_binding.value()->thread_tag, std::nullopt)); } @@ -1884,7 +1884,7 @@ void EmitFor(DocTranslatorObj* d, const tirx::ForNode* loop) { } if (loop->step.has_value()) { TVM_FFI_CHECK(loop->kind != tirx::ForKind::kThreadBinding, TypeError) - << "v2 thread-binding loop cannot have a step"; + << "printer thread-binding loop cannot have a step"; keys.push_back("step"); values.push_back(d->TranslateAs(loop->step.value())); } @@ -2007,7 +2007,7 @@ ffi::String ScopeIdApi(tirx::ScopeBinding scope) { if (parent == "warp" && current == "thread") return "lane_id"; if (parent == "cta" && current == "thread") return "thread_id"; if (parent == "warpgroup" && current == "thread") return "thread_id_in_wg"; - TVM_FFI_THROW(ValueError) << "v2 unknown scope-id binding " << parent << "/" << current; + TVM_FFI_THROW(ValueError) << "printer unknown scope-id binding " << parent << "/" << current; } ExprDoc ScopeExtents(DocTranslatorObj* d, const ffi::Array<PrimExpr>& values) { ffi::Array<ExprDoc> items; @@ -2031,7 +2031,7 @@ void EmitScopeIdDef(DocTranslatorObj* d, const tirx::ScopeIdDefStmtNode* stmt) { if (!def->def_ids.empty()) { PrimType dtype = def->def_ids[0].ty(); for (const PrimVar& var : def->def_ids) { - TVM_FFI_CHECK(var.ty() == dtype, TypeError) << "v2 mixed scope-id dtypes"; + TVM_FFI_CHECK(var.ty() == dtype, TypeError) << "printer mixed scope-id dtypes"; } if (dtype != PrimType::Int(32)) { keys.push_back("dtype"); @@ -2103,7 +2103,7 @@ void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc {LiteralDoc::DataType(primitive.value()->dtype, std::nullopt)}); } else { TVM_FFI_CHECK(!var->ty.IsMissing(), TypeError) - << "v2 free non-scalar variable requires a type annotation"; + << "printer free non-scalar variable requires a type annotation"; annotation = d->TranslateAs(var->ty); } ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>( @@ -2319,10 +2319,10 @@ void EmitTilePrimitiveCall(DocTranslatorObj* d, const tirx::TilePrimitiveCallNod static const OpAttrMap<tirx::TScriptPrinterName>& names = Op::GetAttrMap<tirx::TScriptPrinterName>("TScriptPrinterName"); TVM_FFI_CHECK(names.count(stmt->op), TypeError) - << "v2 tile primitive has no canonical script name: " << stmt->op->name; + << "printer tile primitive has no canonical script name: " << stmt->op->name; std::string name = names[stmt->op]; TVM_FFI_CHECK(name.find("tirx.tile.") == 0, TypeError) - << "v2 tile primitive name must be in tirx.tile namespace: " << name; + << "printer tile primitive name must be in tirx.tile namespace: " << name; name.erase(0, 10); ffi::String scope; switch (stmt->scope->kind) { @@ -2426,7 +2426,7 @@ ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode* blo kind = "opaque"; break; default: - TVM_FFI_THROW(TypeError) << "v2 unsupported SBlock iter var kind"; + TVM_FFI_THROW(TypeError) << "printer unsupported SBlock iter var kind"; } ExprDoc domain = prim::is_zero(iter->dom->min) && iter->dom->min.ty() == iter->dom->extent.ty() @@ -2639,7 +2639,7 @@ void EmitFunction(DocTranslatorObj* d, const tirx::PrimFuncNode* func) { ffi::Optional<Var> BindingVar(const ffi::Object* destination) { if (!destination) return std::nullopt; TVM_FFI_CHECK(destination->IsInstance<VarNode>(), TypeError) - << "v2 binding destination must be a Var"; + << "printer binding destination must be a Var"; return ffi::GetRef<Var>(static_cast<const VarNode*>(destination)); } ffi::Array<StmtDoc> RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq, @@ -2761,11 +2761,11 @@ ffi::Array<StmtDoc> RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* ffi::GetRef<ffi::ObjectRef>(seq)); } else { TVM_FFI_CHECK(destination_object != nullptr, ValueError) - << "v2 Relax SeqExpr completed a binding without a destination"; + << "printer Relax SeqExpr completed a binding without a destination"; } } else { TVM_FFI_CHECK(value.has_value(), ValueError) - << "v2 Relax SeqExpr needs a value without a destination"; + << "printer Relax SeqExpr needs a value without a destination"; d->Emit(ExprStmtDoc(value.value()), ffi::GetRef<ffi::ObjectRef>(seq)); } }); @@ -2797,7 +2797,7 @@ ffi::Optional<ExprDoc> EmitRelaxIf(DocTranslatorObj* d, const relax::IfNode* bra destination = temporary.value().get(); } TVM_FFI_CHECK(destination->IsInstance<VarNode>(), TypeError) - << "v2 Relax If destination must be a Var"; + << "printer Relax If destination must be a Var"; Var var = ffi::GetRef<Var>(static_cast<const VarNode*>(destination)); ffi::Optional<IdDoc> lhs = VarDoc(d, var); ffi::Optional<ExprDoc> annotation = std::nullopt; @@ -2942,7 +2942,7 @@ void EmitModule(DocTranslatorObj* d, const IRModuleNode* mod) { continue; } TVM_FFI_CHECK(func.as<tirx::PrimFuncNode>() || func.as<relax::FunctionNode>(), TypeError) - << "v2 IRModule needs a registered TIRx, Relax, or extern function hook"; + << "printer IRModule needs a registered TIRx, Relax, or extern function hook"; d->Translate(func); FunctionDoc doc = d->CurrentScopeDocs().back().as_or_throw<FunctionDoc>(); doc->name = IdDoc(gv->name_hint); @@ -2978,7 +2978,7 @@ ffi::Optional<ExprDoc> InvokeNative(DocTranslatorObj* d, ffi::AnyView value, } } else if constexpr (std::is_void_v<decltype(Hook(d, node))>) { TVM_FFI_CHECK(destination == nullptr, TypeError) - << "v2 statement-only node cannot fulfill a destination"; + << "printer statement-only node cannot fulfill a destination"; Hook(d, node); return std::nullopt; } else { @@ -3155,27 +3155,25 @@ TVM_FFI_STATIC_INIT_BLOCK() { RegisterHook<relax::FunctionNode, EmitRelaxFunction>(); RegisterHook<IRModuleNode, EmitModule>(); refl::GlobalDef() - .def("script.printer.v2.DocTranslate", + .def("printer.DocTranslate", [](ffi::AnyView ir, ffi::Map<ffi::String, ffi::Any> extra_config) { return DocTranslate(ir, nullptr, std::move(extra_config)); }) - .def_method("script.printer.v2.Translate", &DocTranslatorObj::Translate) - .def_method("script.printer.v2.Emit", &DocTranslatorObj::Emit) - .def_method("script.printer.v2.VarGetOrAllocId", &DocTranslatorObj::VarGetOrAllocId) - .def_method("script.printer.v2.GetImplicitDefs", &DocTranslatorObj::GetImplicitDefs) - .def_method("script.printer.v2.AllocId", &DocTranslatorObj::AllocId) - .def_method("script.printer.v2.BeginVarScope", &DocTranslatorObj::BeginVarScope) - .def_method("script.printer.v2.EndVarScope", &DocTranslatorObj::EndVarScope) - .def_method("script.printer.v2.BeginDocScope", &DocTranslatorObj::BeginDocScope) - .def_method("script.printer.v2.EndDocScope", &DocTranslatorObj::EndDocScope) - .def("script.printer.v2.WithDocScope", + .def_method("printer.Translate", &DocTranslatorObj::Translate) + .def_method("printer.Emit", &DocTranslatorObj::Emit) + .def_method("printer.VarGetOrAllocId", &DocTranslatorObj::VarGetOrAllocId) + .def_method("printer.GetImplicitDefs", &DocTranslatorObj::GetImplicitDefs) + .def_method("printer.AllocId", &DocTranslatorObj::AllocId) + .def_method("printer.BeginVarScope", &DocTranslatorObj::BeginVarScope) + .def_method("printer.EndVarScope", &DocTranslatorObj::EndVarScope) + .def_method("printer.BeginDocScope", &DocTranslatorObj::BeginDocScope) + .def_method("printer.EndDocScope", &DocTranslatorObj::EndDocScope) + .def("printer.WithDocScope", [](const DocTranslator& d, const ffi::Function& fn) { return d->WithDocScope([&]() { fn(); }); }) - .def("script.printer.v2.Script", Script); + .def("printer.Script", Script); } -} // namespace v2 } // namespace printer -} // namespace script } // namespace tvm diff --git a/src/script/printer/v2/entry.cc b/src/printer/entry.cc similarity index 98% rename from src/script/printer/v2/entry.cc rename to src/printer/entry.cc index f98d2fac0a..e93b0a4d6f 100644 --- a/src/script/printer/v2/entry.cc +++ b/src/printer/entry.cc @@ -21,11 +21,11 @@ #include <tvm/ffi/reflection/accessor.h> #include <tvm/ir/module.h> #include <tvm/ir/prim/expr.h> +#include <tvm/printer/doc.h> +#include <tvm/printer/doc_translator.h> #include <tvm/relax/expr.h> #include <tvm/runtime/base.h> #include <tvm/s_tir/stmt.h> -#include <tvm/script/printer/doc.h> -#include <tvm/script/printer/doc_translator.h> #include <tvm/te/operation.h> #include <tvm/tirx/function.h> #include <tvm/tirx/index_map.h> @@ -44,9 +44,7 @@ #include <vector> namespace tvm { -namespace script { namespace printer { -namespace v2 { namespace { // The only mutable structural state. Identifier allocation lives with the @@ -171,7 +169,7 @@ class TranslationEngine : public DocTranslatorObj { static void EndScope(DocTranslatorObj* d) { auto* self = static_cast<TranslationEngine*>(d); TVM_FFI_CHECK(self->var_scopes_.size() > 1, ValueError) - << "v2 cannot pop the root variable scope"; + << "printer cannot pop the root variable scope"; // Scalar state is keyed by stable IdDoc. Remove only bindings owned here; // captured outer bindings and whole-output accumulators remain untouched. auto scalars = @@ -684,7 +682,7 @@ class EntryEngine : public TranslationEngine { Emit(AssignDoc(id, creation, std::nullopt), var); } else { TVM_FFI_CHECK(!var->ty.IsMissing(), TypeError) - << "v2 free non-scalar variable requires a type annotation"; + << "printer free non-scalar variable requires a type annotation"; Emit(AssignDoc(id, std::nullopt, TranslateAs(var->ty)), var); } } @@ -705,7 +703,7 @@ class EntryEngine : public TranslationEngine { {LiteralDoc::DataType(primitive.value()->dtype, std::nullopt)}); } else { TVM_FFI_CHECK(!var->ty.IsMissing(), TypeError) - << "v2 free non-scalar variable requires a type annotation"; + << "printer free non-scalar variable requires a type annotation"; annotation = TranslateAs(var->ty); } ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>( @@ -857,7 +855,7 @@ CompleteTranslation TranslateComplete(ffi::AnyView obj, EntryEngine* engine, if (auto func = obj.as<tirx::PrimFunc>()) { if (auto symbol = func.value()->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) { TVM_FFI_CHECK(symbol.value() == name, TypeError) - << "v2 PrimFunc global_symbol must match its definition name"; + << "printer PrimFunc global_symbol must match its definition name"; } } function.value()->name = IdDoc(name); @@ -967,10 +965,8 @@ ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config) { } if (!header.empty()) header.push_back(ffi::String("\n")); } - return v2::DocToPythonScriptWithHeader(result.doc, local, header); + return DocToPythonScriptWithHeader(result.doc, local, header); } -} // namespace v2 } // namespace printer -} // namespace script } // namespace tvm diff --git a/src/script/printer/v2/python_doc_printer.cc b/src/printer/python_doc_printer.cc similarity index 99% rename from src/script/printer/v2/python_doc_printer.cc rename to src/printer/python_doc_printer.cc index e1ff060c9b..3226f1e675 100644 --- a/src/script/printer/v2/python_doc_printer.cc +++ b/src/printer/python_doc_printer.cc @@ -20,8 +20,7 @@ #include <tvm/ffi/function.h> #include <tvm/ffi/reflection/registry.h> #include <tvm/ir/prim/expr.h> -#include <tvm/script/printer/doc.h> -#include <tvm/script/printer/v2/doc.h> +#include <tvm/printer/doc.h> #include <algorithm> #include <cmath> @@ -36,13 +35,11 @@ #include <utility> #include <vector> -#include "../../../support/str_escape.h" -#include "../../../support/utils.h" +#include "../support/str_escape.h" +#include "../support/utils.h" namespace tvm { -namespace script { namespace printer { -namespace v2 { namespace { /*! \brief Range of byte offsets in a string */ @@ -1661,10 +1658,8 @@ ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("script.printer.v2.DocToPythonScript", DocToPythonScript); + refl::GlobalDef().def("printer.DocToPythonScript", DocToPythonScript); } -} // namespace v2 } // namespace printer -} // namespace script } // namespace tvm diff --git a/src/script/printer/script_printer.cc b/src/printer/script_printer.cc similarity index 89% rename from src/script/printer/script_printer.cc rename to src/printer/script_printer.cc index 0d9f23e303..dd0cfde4bd 100644 --- a/src/script/printer/script_printer.cc +++ b/src/printer/script_printer.cc @@ -24,12 +24,12 @@ #include <tvm/ir/module.h> #include <tvm/ir/prim/expr.h> #include <tvm/ir/prim/vector_expr.h> +#include <tvm/printer/doc_translator.h> #include <tvm/relax/distributed/type.h> #include <tvm/relax/expr.h> #include <tvm/relax/type.h> #include <tvm/runtime/logging.h> #include <tvm/s_tir/stmt.h> -#include <tvm/script/printer/doc_translator.h> #include <tvm/te/operation.h> #include <tvm/tirx/exec_scope.h> #include <tvm/tirx/function.h> @@ -75,13 +75,13 @@ void RegisterScriptRepr() { 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::v2::kDocTranslate); + static ffi::reflection::TypeAttrColumn translate(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::v2::Script(node, config)); + return std::string(printer::Script(node, config)); } std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) { @@ -89,11 +89,11 @@ std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) { PrinterConfig config; // Python repr/str is an interactive display surface. In particular, // backend helpers parse the quoted repr of StringImm values; the direct - // v2 Script entry keeps its explicit IR constructor for round-trips. + // Script entry keeps its explicit IR constructor for round-trips. config->extra_config.Set("script.legacy_display", true); config->extra_config.Set("script.comment_imports", true); // Call translation directly so an unsupported type cannot recurse through ffi repr. - return std::string(script::printer::v2::Script(obj, config)); + return std::string(printer::Script(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(); @@ -103,27 +103,15 @@ std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) { } } -namespace script { -namespace printer { - -ffi::String DocToPythonScript(Doc doc, const PrinterConfig& config) { - return v2::DocToPythonScript(std::move(doc), config); -} - -} // namespace printer -} // namespace script - TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = ffi::reflection; refl::GlobalDef() .def("node.TVMScriptPrinterScript", tvm::Script) - .def("script.printer.DocToPythonScript", script::printer::DocToPythonScript) - .def("script.printer.ReprPrintRelax", - [](const ffi::ObjectRef& obj, const PrinterConfig& config) { - return script::printer::v2::Script(obj, config); - }); - script::printer::RegisterDialectPrefix("tirx.prefix", "T"); - script::printer::RegisterDialectPrefix("relax.prefix", "R"); + .def("printer.ReprPrintRelax", [](const ffi::ObjectRef& obj, const PrinterConfig& config) { + return printer::Script(obj, config); + }); + printer::RegisterDialectPrefix("tirx.prefix", "T"); + printer::RegisterDialectPrefix("relax.prefix", "R"); RegisterScriptRepr<DataTypeImmNode>(); RegisterScriptRepr<GenericConstNode>(); diff --git a/src/script/printer/utils.h b/src/printer/utils.h similarity index 89% rename from src/script/printer/utils.h rename to src/printer/utils.h index ebb9581b8f..beb579c360 100644 --- a/src/script/printer/utils.h +++ b/src/printer/utils.h @@ -16,15 +16,14 @@ * specific language governing permissions and limitations * under the License. */ -#ifndef TVM_SCRIPT_PRINTER_UTILS_H_ -#define TVM_SCRIPT_PRINTER_UTILS_H_ +#ifndef TVM_PRINTER_UTILS_H_ +#define TVM_PRINTER_UTILS_H_ #include <tvm/ffi/dtype.h> #include <string> namespace tvm { -namespace script { namespace printer { /*! @@ -39,7 +38,6 @@ inline std::string DType2Str(DLDataType dtype) { } } // namespace printer -} // namespace script } // namespace tvm -#endif // TVM_SCRIPT_PRINTER_UTILS_H_ +#endif // TVM_PRINTER_UTILS_H_ diff --git a/src/s_tir/meta_schedule/database/json_database.cc b/src/s_tir/meta_schedule/database/json_database.cc index 5db35c3776..bc927afaa0 100644 --- a/src/s_tir/meta_schedule/database/json_database.cc +++ b/src/s_tir/meta_schedule/database/json_database.cc @@ -17,7 +17,7 @@ * under the License. */ #include <tvm/ffi/reflection/registry.h> -#include <tvm/script/printer/printer.h> +#include <tvm/printer/printer.h> #include <set> #include <thread> diff --git a/src/s_tir/schedule/error.cc b/src/s_tir/schedule/error.cc index b6aed81da9..bfd8f436d2 100644 --- a/src/s_tir/schedule/error.cc +++ b/src/s_tir/schedule/error.cc @@ -18,7 +18,7 @@ */ #include <tvm/ffi/extra/visit_error_context.h> #include <tvm/ffi/reflection/registry.h> -#include <tvm/script/printer/printer.h> +#include <tvm/printer/printer.h> #include "./utils.h" diff --git a/src/s_tir/script/printer/config.cc b/src/s_tir/script/printer/config.cc index 461858edad..e1abb2cee9 100644 --- a/src/s_tir/script/printer/config.cc +++ b/src/s_tir/script/printer/config.cc @@ -18,14 +18,12 @@ */ #include <tvm/ffi/reflection/registry.h> -#include "../../../script/printer/dialect_prefix.h" +#include "../../../printer/dialect_prefix.h" namespace tvm { -namespace script { namespace printer { TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("s_tir.prefix", "Ts"); } } // namespace printer -} // namespace script } // namespace tvm diff --git a/src/script/printer/v2/doc.cc b/src/script/printer/v2/doc.cc deleted file mode 100644 index 55beb45215..0000000000 --- a/src/script/printer/v2/doc.cc +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include <tvm/script/printer/v2/doc.h> - -namespace tvm { -namespace script { -namespace printer { -namespace v2 { - -NamespaceDoc::NamespaceDoc(ffi::String canonical_name) { - auto node = ffi::make_object<NamespaceDocNode>(); - node->canonical_name = std::move(canonical_name); - data_ = std::move(node); -} - -TVM_FFI_STATIC_INIT_BLOCK() { NamespaceDocNode::RegisterReflection(); } - -} // namespace v2 -} // namespace printer -} // namespace script -} // namespace tvm diff --git a/tests/cpp/tir_scalable_datatype.cc b/tests/cpp/tir_scalable_datatype.cc index 8023344a96..5eb4b1366f 100644 --- a/tests/cpp/tir_scalable_datatype.cc +++ b/tests/cpp/tir_scalable_datatype.cc @@ -22,14 +22,14 @@ #include <tvm/ffi/dtype.h> #include <tvm/ir/prim/builtin.h> #include <tvm/ir/prim/expr.h> -#include <tvm/script/printer/printer.h> +#include <tvm/printer/printer.h> #include <tvm/tirx/builtin.h> #ifdef TVM_LLVM_VERSION #include <llvm/IR/Intrinsics.h> #endif -#include "../../src/script/printer/utils.h" +#include "../../src/printer/utils.h" using ::testing::HasSubstr; @@ -209,5 +209,5 @@ TEST(ScalablePrimType, TestScalableIntrinCall) { TEST(ScalablePrimType, TestTIRScriptScalableDtype2Str) { tvm::PrimType scalable_type = tvm::PrimType::ScalableVector(kDLInt, 32, 4); - ASSERT_EQ(tvm::script::printer::DType2Str(scalable_type->dtype), "int32xvscalex4"); + ASSERT_EQ(tvm::printer::DType2Str(scalable_type->dtype), "int32xvscalex4"); } diff --git a/tests/python/script/test_script_printer_doc.py b/tests/python/script/test_script_printer_doc.py index 77f9383e18..0d9d609a94 100644 --- a/tests/python/script/test_script_printer_doc.py +++ b/tests/python/script/test_script_printer_doc.py @@ -24,7 +24,7 @@ import pytest from tvm_ffi.access_path import AccessPath import tvm -from tvm.script.printer.doc import ( +from tvm.printer.doc import ( AssertDoc, AssignDoc, CallDoc, @@ -49,7 +49,7 @@ from tvm.script.printer.doc import ( TupleDoc, WhileDoc, ) -from tvm.script.printer.doc_printer import to_python_script +from tvm.printer.doc_printer import to_python_script def format_script(s: str) -> str:
