You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/01 15:53:38 UTC
(tvm) branch tqchen/printer-refinement created (now 6bb89c913d)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
at 6bb89c913d [REFACTOR][Script] Replace docsifier with native Doc translation
This branch includes the following new commits:
new 6bb89c913d [REFACTOR][Script] Replace docsifier with native Doc translation
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(tvm) 01/01: [REFACTOR][Script] Replace docsifier with native Doc translation
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6bb89c913dab053db36fe4cad3bbc65089acb839
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 15:53:22 2026 +0000
[REFACTOR][Script] Replace docsifier with native Doc translation
Use a single native hook contract for expression translation, bindings and
statement emission. Keep node and operation hooks in their owning printer
files, with variable scopes and Doc-to-IR provenance managed by the
translator.
Separate IR-to-Doc translation, private Doc-to-text rendering and script
entry/path mapping. Keep printer configuration read-only and derived
state local to each invocation.
Remove the old docsifier dispatch, signature adapters and legacy display
mode. Let native hooks own compact syntax and defaults, and separate
intrinsic result inference from validation while preserving Relax's
deferred inference.
---
CMakeLists.txt | 22 +-
docs/arch/tvmscript.rst | 23 +-
docs/reference/api/python/script/ir_builder.rst | 2 +-
docs/reference/api/python/script/printer.rst | 4 +-
include/tvm/ir/expr.h | 4 +-
include/tvm/ir/op.h | 6 +
include/tvm/s_tir/script/ir_builder/ir.h | 20 +-
include/tvm/script/printer/config.h | 26 +-
include/tvm/script/printer/doc.h | 158 ++-
include/tvm/script/printer/doc_translator.h | 422 +++++++
include/tvm/script/printer/ir_docsifier.h | 365 ------
include/tvm/script/printer/ir_docsifier_functor.h | 199 ---
include/tvm/script/printer/printer.h | 60 +-
include/tvm/tirx/op_attr_types.h | 2 +-
python/tvm/backend/cuda/ptx/engine.py | 14 +-
python/tvm/ir/base.py | 17 +-
python/tvm/ir/expr.py | 4 +-
python/tvm/relax/base_py_module.py | 2 -
python/tvm/relax/script/ir_builder/__init__.py | 45 +
python/tvm/relax/script/ir_builder/ir.py | 2 +
python/tvm/relax/type.py | 2 +-
python/tvm/runtime/script_printer.py | 35 +-
python/tvm/s_tir/script/ir_builder/ir.py | 33 +-
.../tvm/s_tir/script/ir_builder/parser_protocol.py | 21 +-
python/tvm/script/ir_builder/__init__.py | 36 +-
python/tvm/tirx/op.py | 8 +-
python/tvm/tirx/script/ir_builder/ir.py | 114 +-
python/tvm/tirx/script/ir_builder/op.py | 4 +-
.../tvm/tirx/script/ir_builder/parser_protocol.py | 4 +-
src/backend/cuda/codegen/codegen_cuda.cc | 4 +-
src/backend/cuda/codegen/intrin_rule_cuda.cc | 26 +-
src/backend/cuda/op/iket.cc | 26 +-
src/backend/cuda/op/target_builtin.cc | 94 +-
src/backend/metal/codegen/intrin_rule_metal.cc | 17 +-
src/backend/metal/op/target_builtin.cc | 8 +-
src/backend/trn/op/target_builtin.cc | 5 +-
src/backend/webgpu/codegen/intrin_rule_webgpu.cc | 16 +-
src/ir/expr.cc | 7 +-
src/ir/prim/builtin.cc | 1 +
src/ir/prim/expr.cc | 3 +-
src/ir/transform.cc | 4 +
src/relax/op/op.cc | 3 +-
src/relax/op/op_common.h | 14 +-
src/relax/op/tensor/binary.cc | 4 +-
src/relax/script/printer/binding.cc | 199 +--
src/relax/script/printer/call.cc | 535 ++++----
src/relax/script/printer/dependent_type.cc | 231 ++--
src/relax/script/printer/distributed.cc | 201 ++-
src/relax/script/printer/expr.cc | 179 +--
src/relax/script/printer/function.cc | 235 ++--
src/relax/script/printer/global_info.cc | 67 +-
src/relax/script/printer/region.cc | 161 ++-
src/relax/script/printer/tir.cc | 135 +--
src/relax/script/printer/type.cc | 51 +-
src/relax/script/printer/utils.h | 159 +--
src/s_tir/script/ir_builder/frame.cc | 49 +-
src/s_tir/script/ir_builder/ir.cc | 8 +-
src/script/printer/config.cc | 26 +-
src/script/printer/dialect_prefix.h | 7 +-
src/script/printer/doc.cc | 12 +-
.../python_doc_printer.cc => doc_printer.cc} | 875 ++++++++++++-
.../printer/{dialect_prefix.h => doc_printer.h} | 18 +-
src/script/printer/doc_printer/base_doc_printer.cc | 386 ------
src/script/printer/doc_printer/base_doc_printer.h | 318 -----
src/script/printer/doc_translator.cc | 279 +++++
src/script/printer/doc_translator.h | 57 +
src/script/printer/ir/ir.cc | 323 ++---
src/script/printer/ir/misc.cc | 245 ++--
src/script/printer/ir/prim_expr.cc | 432 ++++---
src/script/printer/ir/prim_type.cc | 137 ++-
src/script/printer/ir/utils.cc | 266 ++++
src/script/printer/ir/utils.h | 75 +-
src/script/printer/ir_docsifier.cc | 218 ----
src/script/printer/printer.cc | 302 +++++
src/script/printer/script_printer.cc | 158 ++-
src/script/printer/utils.h | 228 +---
src/target/llvm/llvm_module.cc | 6 +-
src/tirx/op/builtin.cc | 190 +--
src/tirx/op/op.cc | 82 +-
src/tirx/op/tirx.cc | 60 +-
src/tirx/script/ir_builder/ir.cc | 33 +-
src/tirx/script/printer/block.cc | 483 ++++----
src/tirx/script/printer/buffer.cc | 1281 ++++++--------------
src/tirx/script/printer/expr.cc | 985 ++++++++-------
src/tirx/script/printer/for_loop.cc | 256 ++--
src/tirx/script/printer/function.cc | 414 +++----
src/tirx/script/printer/ir.cc | 91 +-
src/tirx/script/printer/stmt.cc | 655 +++++-----
src/tirx/script/printer/utils.h | 332 +----
src/tirx/transform/lower_thread_allreduce.h | 2 +-
.../relax/script/test_relax_script_printer.py | 52 +-
.../s_tir/script/test_s_tir_script_ir_builder.py | 6 +-
.../s_tir/script/test_s_tir_script_printer.py | 157 +--
...form_merge_dynamic_shared_memory_allocations.py | 20 +-
.../transform/test_s_tir_transform_thread_sync.py | 11 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 12 +
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 8 +
.../tirx/script/test_tirx_script_basic_usage.py | 4 +-
.../python/tirx/script/test_tirx_script_buffers.py | 42 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 4 +-
.../tirx/script/test_tirx_script_ir_builder.py | 4 +-
.../python/tirx/script/test_tirx_script_printer.py | 162 +--
.../script/test_tirx_script_printer_underlining.py | 1 +
.../script/test_tirx_script_source_locations.py | 4 +-
104 files changed, 7023 insertions(+), 6792 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 719c66da25..58a296f469 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -303,21 +303,23 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/topi/*.cc
src/driver/*.cc
src/support/*.cc
- # TVMScript shared core (Doc IR + dispatch infrastructure + IR-layer
- # printer/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/doc_printer/*.cc
- src/script/printer/ir_docsifier.cc
- src/script/printer/ir/*.cc
+ src/script/printer/doc_translator.cc
+ src/script/printer/ir/utils.cc
+ src/script/printer/ir/ir.cc
+ src/script/printer/ir/misc.cc
+ src/script/printer/ir/prim_type.cc
+ src/script/printer/ir/prim_expr.cc
+ src/script/printer/printer.cc
+ src/script/printer/doc_printer.cc
+ src/script/printer/script_printer.cc
src/relax/ir/*.cc
src/relax/op/*.cc
src/relax/analysis/*.cc
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 51eb6b719e..d7d64fb3d2 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -29,7 +29,7 @@ printer converts IR back into readable, parseable source.
Source + definition context
→ Python AST → syntax transpiler → Python builder program → TVM IR
TVM IR
- → IRDocsifier → Doc tree → Python text
+ → DocTranslator → Doc tree + origins → diagnostic paths → Python text
Source and frontend
-------------------
@@ -139,9 +139,17 @@ protocol.
Printing and round trips
------------------------
-``IRDocsifier`` dispatches on IR types and context to produce a ``Doc`` tree of
-expressions and statements. It tracks scopes and names; ``DocToPythonScript`` formats
-the tree as Python text. This document tree is separate from the parser's Python AST.
+``DocTranslator`` invokes native type and operation hooks to produce a ``Doc`` tree
+of expressions and statements. The translation engine tracks scopes, names and each
+Doc's original IR object. The script entry point maps these origins to diagnostic
+paths; the private Doc printer formats the tree, annotations and underlines as Python
+text. Printer configuration stays read-only throughout. ``DocToPythonScript`` also
+formats an existing Doc directly. This tree is separate from the parser's Python AST.
+
+The public ``Script`` text entry points are declared in
+``tvm/script/printer/printer.h``. Their orchestration and diagnostic path mapping
+live in ``src/script/printer/printer.cc``; ``doc_translator.h`` exposes the
+IR-to-Doc translation protocol.
For example, a small function can be authored, printed and parsed again:
@@ -156,7 +164,7 @@ For example, a small function can be authored, printed and parsed again:
A[i] = A[i] + 1.0
text = increment.script()
- reparsed = tvm.script.from_source(text)
+ reparsed = tvm.script.from_source(text, extra_vars={"T": T})
tvm.ir.assert_structural_equal(increment, reparsed)
Printed source uses canonical forms rather than preserving the original spelling.
@@ -170,7 +178,10 @@ A language variant supplies construction hooks and a public script namespace, th
registers its aliases with ``register_namespace``. ``register_namespace_initializer``
supports lazy setup. ``tvm.script.register_dialect`` exposes a package through the
public script namespace. Syntax policies are registered beside the operations that
-need them, while IR printing is extended through docsifier registrations.
+need them. IR printing uses ``FDocTranslate`` hooks from
+``tvm/script/printer/doc_translator.h``, registered through the existing type or
+operation attributes. A hook returns an expression Doc or emits completed statements
+through its translator context.
This division keeps source acquisition, syntax translation and formatting shared.
Language-specific construction and validation remain with the namespace and its
diff --git a/docs/reference/api/python/script/ir_builder.rst b/docs/reference/api/python/script/ir_builder.rst
index b17ba6de31..7975118155 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,7 +23,7 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
- :exclude-members: DataTypeImm, GenericConst, Range, StringImm, StringType
+ :exclude-members: DataTypeImm, FuncType, GenericConst, PrimType, Range, StringImm, StringType, Type
tvm.relax.script.ir_builder.distributed
***************************************
diff --git a/docs/reference/api/python/script/printer.rst b/docs/reference/api/python/script/printer.rst
index b5dddb4396..cc16c55c3b 100644
--- a/docs/reference/api/python/script/printer.rst
+++ b/docs/reference/api/python/script/printer.rst
@@ -16,10 +16,10 @@
under the License.
tvm.script.printer
-------------------
+-------------------------
tvm.script.printer
-******************
+*************************
.. automodule:: tvm.script.printer
:members:
:imported-members:
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index d30833e8a6..a597cfe0a8 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -519,8 +519,8 @@ class Call : public Expr {
/*! \brief Recompute a result type from the Call's current explicit inputs.
*
* This ignores the Call's stored result type and does not mutate the Call.
- * A context-free inference hook must be registered for the operator. Missing
- * or invalid inputs are reported by the hook rather than screened here.
+ * The operator must register TFixedReturnType or a context-free FInferType hook.
+ * This does not invoke the operator's validator.
*/
TVM_DLL static Type ReinferType(const CallNode* call);
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index 403541a104..dcb985b882 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -43,6 +43,12 @@ namespace tvm {
template <typename>
class OpAttrMap;
+/*!
+ * \brief An operator's result type independent of operands, attributes, and type arguments.
+ * When present, this concrete type takes precedence over FInferType.
+ */
+using TFixedReturnType = Type;
+
/*! \brief Infer a Call's result type from its explicit inputs without builder state. */
using FInferType = ffi::reflection::NativeFunctionView<Type(const CallNode* call)>;
diff --git a/include/tvm/s_tir/script/ir_builder/ir.h b/include/tvm/s_tir/script/ir_builder/ir.h
index 8c13911544..4b7427654d 100644
--- a/include/tvm/s_tir/script/ir_builder/ir.h
+++ b/include/tvm/s_tir/script/ir_builder/ir.h
@@ -120,38 +120,42 @@ namespace axis {
/*!
* \brief The spatial block axis defining function.
* \param dom The domain of the iteration variable.
- * \param binding The binding value of the iteration variable.
+ * \param binding The binding value, omitted only for a no_realize block.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
-Var Spatial(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
+Var Spatial(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
+ PrimType dtype = PrimType::Int(32));
/*!
* \brief The reduced block axis defining function.
* \param dom The domain of the iteration variable.
- * \param binding The binding value of the iteration variable.
+ * \param binding The binding value, omitted only for a no_realize block.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
-Var Reduce(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
+Var Reduce(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
+ PrimType dtype = PrimType::Int(32));
/*!
* \brief The scanning block axis defining function.
* \param dom The domain of the iteration variable.
- * \param binding The binding value of the iteration variable.
+ * \param binding The binding value, omitted only for a no_realize block.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
-Var Scan(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
+Var Scan(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
+ PrimType dtype = PrimType::Int(32));
/*!
* \brief The opaque block axis defining function.
* \param dom The domain of the iteration variable.
- * \param binding The binding value of the iteration variable.
+ * \param binding The binding value, omitted only for a no_realize block.
* \param dtype The data type of the iteration variable.
* \return The iteration variable.
*/
-Var Opaque(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
+Var Opaque(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
+ PrimType dtype = PrimType::Int(32));
/*!
* \brief The block axis remapping function.
diff --git a/include/tvm/script/printer/config.h b/include/tvm/script/printer/config.h
index 02bd96bba7..e770de90f1 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -21,8 +21,8 @@
* \brief Configuration object for the TVMScript printer.
*
* Contains PrinterConfig / PrinterConfigNode, GetBuiltinKeywords, GetExtraConfig,
- * and RedirectedReprPrinterMethod. The entry-point free function tvm::Script()
- * and the dispatch vtable TVMScriptPrinter live in printer.h.
+ * and RedirectedReprPrinterMethod. Translation APIs live in doc_translator.h,
+ * and text entry points live in printer.h.
*/
#ifndef TVM_SCRIPT_PRINTER_CONFIG_H_
#define TVM_SCRIPT_PRINTER_CONFIG_H_
@@ -49,20 +49,11 @@ class PrinterConfigNode : public ffi::Object {
/*! \brief The prefix of IR nodes */
ffi::String ir_prefix = "I";
/*!
- * \brief The alias of the current module at cross-function call
- * \note Directly use module name if it's empty.
+ * \brief Compatibility option; global references use the actual module name.
*/
ffi::String module_alias = "cls";
- /*! \brief Default buffer dtype */
+ /*! \brief Compatibility option; buffer hooks print the actual dtype. */
DLDataType buffer_dtype = DLDataType{kDLFloat, 32, 1};
- /*! \brief Default data type of integer literals */
- DLDataType int_dtype = DLDataType{kDLInt, 32, 1};
- /*!
- * \brief Default data type of float literals. Right now we always print out the explicit type
- * of floating point values, so setting it to Void means we do not print without the
- * T.float32/T.float64 wrapper.
- */
- DLDataType float_dtype = DLDataType{kDLOpaqueHandle, 0, 0};
/*! \brief Whether or not to verbose print expressions. */
bool verbose_expr = false;
/*! \brief Number of spaces used for indentation*/
@@ -73,8 +64,6 @@ class PrinterConfigNode : public ffi::Object {
int num_context_lines = -1;
/*! \brief Whether to output with syntax sugar, set false for complete printing. */
bool syntax_sugar = true;
- /*! \brief Whether variable names should include the object's address */
- bool show_object_address = false;
/*! \brief Whether to render access-path context for invisible underlined paths. Defaults true. */
bool render_invisible_path_info = true;
@@ -121,14 +110,11 @@ class PrinterConfigNode : public ffi::Object {
.def_ro("ir_prefix", &PrinterConfigNode::ir_prefix)
.def_ro("module_alias", &PrinterConfigNode::module_alias)
.def_ro("buffer_dtype", &PrinterConfigNode::buffer_dtype)
- .def_ro("int_dtype", &PrinterConfigNode::int_dtype)
- .def_ro("float_dtype", &PrinterConfigNode::float_dtype)
.def_ro("verbose_expr", &PrinterConfigNode::verbose_expr)
.def_ro("indent_spaces", &PrinterConfigNode::indent_spaces)
.def_ro("print_line_numbers", &PrinterConfigNode::print_line_numbers)
.def_ro("num_context_lines", &PrinterConfigNode::num_context_lines)
.def_ro("syntax_sugar", &PrinterConfigNode::syntax_sugar)
- .def_ro("show_object_address", &PrinterConfigNode::show_object_address)
.def_ro("render_invisible_path_info", &PrinterConfigNode::render_invisible_path_info)
.def_ro("path_to_underline", &PrinterConfigNode::path_to_underline)
.def_ro("path_to_annotate", &PrinterConfigNode::path_to_annotate)
@@ -152,9 +138,9 @@ class TVM_DLL PrinterConfig : public ffi::ObjectRef {
};
/*!
- * \brief The fallback body used by TVMScriptPrinter::Register.
+ * \brief Render an IR object for its registered repr callback.
*
- * Tries to format \p obj via tvm::Script; on error falls back to a plain
+ * 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
* <tvm/runtime/logging.h> is not pulled into this public header.
*/
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index 84a62477cb..14a270b4bb 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -35,16 +35,6 @@ 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
- */
-ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg);
-
/*!
* \brief The base class of all Doc.
*
@@ -107,19 +97,19 @@ class ExprDocNode : public DocNode {
* \brief Create a doc representing attribute access on the current ExprDoc
* \param attr The attribute to access.
*/
- ExprDoc Attr(ffi::String attr) const;
+ TVM_DLL ExprDoc Attr(ffi::String attr) const;
/*!
* \brief Create a doc representing index access on the current ExprDoc
* \param indices The indices to access.
*/
- ExprDoc operator[](ffi::Array<Doc> indices) const;
+ TVM_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
/*!
* \brief Create a doc representing calling the current ExprDoc
* \param args The positional arguments of the function call.
*/
- ExprDoc Call(ffi::Array<ExprDoc, void> args) const;
+ TVM_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args) const;
/*!
* \brief Create a doc representing attribute access on the current ExprDoc
@@ -127,9 +117,9 @@ class ExprDocNode : public DocNode {
* \param kwargs_keys Keys of keywords arguments of the function call.
* \param kwargs_values Values of keywords arguments of the function call.
*/
- ExprDoc Call(ffi::Array<ExprDoc, void> args, //
- ffi::Array<ffi::String> kwargs_keys, //
- ffi::Array<ExprDoc, void> kwargs_values) const;
+ TVM_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args, //
+ ffi::Array<ffi::String> kwargs_keys, //
+ ffi::Array<ExprDoc, void> kwargs_values) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -152,7 +142,7 @@ class ExprDoc : public Doc {
* \brief Create a doc representing index access on the current ExprDoc
* \param indices The indices to access.
*/
- ExprDoc operator[](ffi::Array<Doc> indices) const;
+ TVM_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
explicit ExprDoc(ffi::ObjectPtr<ExprDocNode> data) : Doc(data) {
TVM_FFI_ICHECK(data != nullptr);
@@ -226,7 +216,7 @@ class StmtBlockDoc : public Doc {
* \brief Constructor of StmtBlockDoc.
* \param stmts The list of statements.
*/
- explicit StmtBlockDoc(ffi::Array<StmtDoc> stmts);
+ TVM_DLL explicit StmtBlockDoc(ffi::Array<StmtDoc> stmts);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtBlockDoc, Doc, StmtBlockDocNode);
};
@@ -262,7 +252,7 @@ class LiteralDocNode : public ExprDocNode {
*/
class LiteralDoc : public ExprDoc {
protected:
- explicit LiteralDoc(ffi::Any value, const ffi::Optional<AccessPath>& object_path);
+ TVM_DLL explicit LiteralDoc(ffi::Any value, const ffi::Optional<AccessPath>& object_path);
public:
/*!
@@ -281,7 +271,12 @@ class LiteralDoc : public ExprDoc {
return LiteralDoc(IntImm::Int64(v), p);
}
- /*! rief Create an integer literal without narrowing its typed payload. */
+ /*!
+ * \brief Create an integer literal preserving its typed payload.
+ * \param v The typed integer value.
+ * \param p The object path.
+ * \return The integer literal.
+ */
static LiteralDoc Int(IntImm v, const ffi::Optional<AccessPath>& p) {
return LiteralDoc(std::move(v), p);
}
@@ -362,7 +357,7 @@ class ExprStringDoc : public ExprDoc {
* \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_DLL explicit ExprStringDoc(ExprDoc value, const ffi::Optional<AccessPath>& object_path);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprStringDoc, ExprDoc, ExprStringDocNode);
};
@@ -395,7 +390,7 @@ class IdDoc : public ExprDoc {
* \brief Constructor of IdDoc.
* \param name The name of identifier.
*/
- explicit IdDoc(ffi::String name);
+ TVM_DLL explicit IdDoc(ffi::String name);
explicit IdDoc(std::nullptr_t) : ExprDoc(nullptr) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IdDoc, ExprDoc, IdDocNode);
};
@@ -433,7 +428,7 @@ class AttrAccessDoc : public ExprDoc {
* \param value The target expression of attribute access.
* \param name The name of attribute to access.
*/
- explicit AttrAccessDoc(ExprDoc value, ffi::String name);
+ TVM_DLL explicit AttrAccessDoc(ExprDoc value, ffi::String name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrAccessDoc, ExprDoc, AttrAccessDocNode);
};
@@ -476,7 +471,7 @@ class IndexDoc : public ExprDoc {
* \param value The target expression of index access.
* \param indices The indices to access.
*/
- explicit IndexDoc(ExprDoc value, ffi::Array<Doc> indices);
+ TVM_DLL explicit IndexDoc(ExprDoc value, ffi::Array<Doc> indices);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IndexDoc, ExprDoc, IndexDocNode);
};
@@ -526,8 +521,8 @@ class CallDoc : public ExprDoc {
* \param kwargs_keys Keys of keyword arguments.
* \param kwargs_values Values of keyword arguments, must have the same length as `kwargs_keys.
*/
- CallDoc(ExprDoc callee, ffi::Array<ExprDoc> args, ffi::Array<ffi::String> kwargs_keys,
- ffi::Array<ExprDoc> kwargs_values);
+ TVM_DLL CallDoc(ExprDoc callee, ffi::Array<ExprDoc> args, ffi::Array<ffi::String> kwargs_keys,
+ ffi::Array<ExprDoc> kwargs_values);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallDoc, ExprDoc, CallDocNode);
};
@@ -606,7 +601,7 @@ class OperationDoc : public ExprDoc {
* \param kind The kind of operation.
* \param operands Operands of this expression.
*/
- explicit OperationDoc(OperationDocNode::Kind kind, ffi::Array<ExprDoc> operands);
+ TVM_DLL explicit OperationDoc(OperationDocNode::Kind kind, ffi::Array<ExprDoc> operands);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OperationDoc, ExprDoc, OperationDocNode);
};
@@ -646,7 +641,7 @@ class LambdaDoc : public ExprDoc {
* \param args Arguments of this function.
* \param body Body expression of this function.
*/
- explicit LambdaDoc(ffi::Array<IdDoc> args, ExprDoc body);
+ TVM_DLL explicit LambdaDoc(ffi::Array<IdDoc> args, ExprDoc body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LambdaDoc, ExprDoc, LambdaDocNode);
};
@@ -682,7 +677,7 @@ class TupleDoc : public ExprDoc {
* \brief Constructor of TupleDoc
* \param elements Elements of tuple.
*/
- explicit TupleDoc(ffi::Array<ExprDoc> elements);
+ TVM_DLL explicit TupleDoc(ffi::Array<ExprDoc> elements);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleDoc, ExprDoc, TupleDocNode);
};
@@ -718,7 +713,7 @@ class ListDoc : public ExprDoc {
* \brief Constructor of ListDoc
* \param elements Elements of list.
*/
- explicit ListDoc(ffi::Array<ExprDoc> elements);
+ TVM_DLL explicit ListDoc(ffi::Array<ExprDoc> elements);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ListDoc, ExprDoc, ListDocNode);
};
@@ -764,7 +759,7 @@ class DictDoc : public ExprDoc {
* \param keys Keys of dictionary.
* \param values Values of dictionary, must have same length as `keys`.
*/
- explicit DictDoc(ffi::Array<ExprDoc> keys, ffi::Array<ExprDoc> values);
+ TVM_DLL explicit DictDoc(ffi::Array<ExprDoc> keys, ffi::Array<ExprDoc> values);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DictDoc, ExprDoc, DictDocNode);
};
@@ -807,8 +802,8 @@ class SliceDoc : public Doc {
* \param stop The exclusive end of slice.
* \param step The step of slice.
*/
- explicit SliceDoc(ffi::Optional<ExprDoc> start, ffi::Optional<ExprDoc> stop,
- ffi::Optional<ExprDoc> step);
+ TVM_DLL explicit SliceDoc(ffi::Optional<ExprDoc> start, ffi::Optional<ExprDoc> stop,
+ ffi::Optional<ExprDoc> step);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SliceDoc, Doc, SliceDocNode);
};
@@ -853,7 +848,8 @@ class AssignDoc : public StmtDoc {
* \param rhs The right hand side of the assignment.
* \param annotation The type annotation of this assignment.
*/
- explicit AssignDoc(ExprDoc lhs, ffi::Optional<ExprDoc> rhs, ffi::Optional<ExprDoc> annotation);
+ TVM_DLL explicit AssignDoc(ExprDoc lhs, ffi::Optional<ExprDoc> rhs,
+ ffi::Optional<ExprDoc> annotation);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssignDoc, StmtDoc, AssignDocNode);
};
@@ -894,8 +890,8 @@ class IfDoc : public StmtDoc {
* \param then_branch The then branch of the if-then-else statement.
* \param else_branch The else branch of the if-then-else statement.
*/
- explicit IfDoc(ExprDoc predicate, ffi::Array<StmtDoc> then_branch,
- ffi::Array<StmtDoc> else_branch);
+ TVM_DLL explicit IfDoc(ExprDoc predicate, ffi::Array<StmtDoc> then_branch,
+ ffi::Array<StmtDoc> else_branch);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfDoc, StmtDoc, IfDocNode);
};
@@ -932,7 +928,7 @@ class WhileDoc : public StmtDoc {
* \param predicate The predicate of the while statement.
* \param body The body of the while statement.
*/
- explicit WhileDoc(ExprDoc predicate, ffi::Array<StmtDoc> body);
+ TVM_DLL explicit WhileDoc(ExprDoc predicate, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WhileDoc, StmtDoc, WhileDocNode);
};
@@ -961,7 +957,7 @@ class BreakDoc : public StmtDoc {
/*!
* \brief Constructor of BreakDoc.
*/
- explicit BreakDoc();
+ TVM_DLL explicit BreakDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BreakDoc, StmtDoc, BreakDocNode);
};
@@ -990,7 +986,7 @@ class ContinueDoc : public StmtDoc {
/*!
* \brief Constructor of ContinueDoc.
*/
- explicit ContinueDoc();
+ TVM_DLL explicit ContinueDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ContinueDoc, StmtDoc, ContinueDocNode);
};
@@ -1035,7 +1031,7 @@ class ForDoc : public StmtDoc {
* \param rhs The right hand side of the assignment of iterating variable.
* \param body The body of the for statement.
*/
- explicit ForDoc(ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body);
+ TVM_DLL explicit ForDoc(ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ForDoc, StmtDoc, ForDocNode);
};
@@ -1057,13 +1053,16 @@ class ScopeDocNode : public StmtDocNode {
ExprDoc rhs{ffi::UnsafeInit()};
/*! \brief The body of the scope doc. */
ffi::Array<StmtDoc> body;
+ /*! \brief Whether a tail-position renderer may spell this scope concisely. */
+ bool allow_concise_scoping{false};
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScopeDocNode>()
.def_ro("lhs", &ScopeDocNode::lhs)
.def_ro("rhs", &ScopeDocNode::rhs)
- .def_ro("body", &ScopeDocNode::body);
+ .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);
};
@@ -1080,15 +1079,17 @@ class ScopeDoc : public StmtDoc {
* \param lhs The name of the scoped variable.
* \param rhs The value of the scoped variable.
* \param body The body of the scope doc.
+ * \param allow_concise_scoping Whether a tail-position renderer may use a concise form.
*/
- explicit ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body);
+ TVM_DLL explicit ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body,
+ bool allow_concise_scoping = false);
/*!
* \brief Constructor of ScopeDoc.
* \param rhs The value of the scoped variable.
* \param body The body of the scope doc.
*/
- explicit ScopeDoc(ExprDoc rhs, ffi::Array<StmtDoc> body);
+ TVM_DLL explicit ScopeDoc(ExprDoc rhs, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ScopeDoc, StmtDoc, ScopeDocNode);
};
@@ -1121,7 +1122,7 @@ class ExprStmtDoc : public StmtDoc {
* \brief Constructor of ExprStmtDoc.
* \param expr The expression represented by this doc.
*/
- explicit ExprStmtDoc(ExprDoc expr);
+ TVM_DLL explicit ExprStmtDoc(ExprDoc expr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprStmtDoc, StmtDoc, ExprStmtDocNode);
};
@@ -1158,7 +1159,7 @@ class AssertDoc : public StmtDoc {
* \param test The expression to test.
* \param msg The optional error message when assertion failed.
*/
- explicit AssertDoc(ExprDoc test, ffi::Optional<ExprDoc> msg = std::nullopt);
+ TVM_DLL explicit AssertDoc(ExprDoc test, ffi::Optional<ExprDoc> msg = std::nullopt);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertDoc, StmtDoc, AssertDocNode);
};
@@ -1190,7 +1191,7 @@ class ReturnDoc : public StmtDoc {
* \brief Constructor of ReturnDoc.
* \param value The value to return.
*/
- explicit ReturnDoc(ExprDoc value);
+ TVM_DLL explicit ReturnDoc(ExprDoc value);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ReturnDoc, StmtDoc, ReturnDocNode);
};
@@ -1255,9 +1256,9 @@ class FunctionDoc : public StmtDoc {
* \param body The body of function.
* \param type_params The PEP 695 type parameters of the function.
*/
- explicit 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 = {});
+ TVM_DLL explicit 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 = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionDoc, StmtDoc, FunctionDocNode);
};
@@ -1298,7 +1299,7 @@ class ClassDoc : public StmtDoc {
* \param decorators The decorator of class.
* \param body The body of class.
*/
- explicit ClassDoc(IdDoc name, ffi::Array<ExprDoc> decorators, ffi::Array<StmtDoc> body);
+ TVM_DLL explicit ClassDoc(IdDoc name, ffi::Array<ExprDoc> decorators, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ClassDoc, StmtDoc, ClassDocNode);
};
@@ -1323,7 +1324,7 @@ class CommentDocNode : public StmtDocNode {
*/
class CommentDoc : public StmtDoc {
public:
- explicit CommentDoc(ffi::String comment);
+ TVM_DLL explicit CommentDoc(ffi::String comment);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommentDoc, StmtDoc, CommentDocNode);
};
@@ -1348,7 +1349,7 @@ class DocStringDocNode : public StmtDocNode {
*/
class DocStringDoc : public StmtDoc {
public:
- explicit DocStringDoc(ffi::String docs);
+ TVM_DLL explicit DocStringDoc(ffi::String docs);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DocStringDoc, StmtDoc, DocStringDocNode);
};
@@ -1398,11 +1399,60 @@ class OpCallDoc : public StmtDoc {
* \param config The config of this op call.
* \param dispatch The optional dispatch variant name of this op call.
*/
- explicit OpCallDoc(ExprDoc callee, ffi::Array<Doc> args, ffi::Optional<DictDoc> workspace,
- ffi::Optional<DictDoc> config, ffi::Optional<ExprDoc> dispatch = std::nullopt);
+ TVM_DLL explicit OpCallDoc(ExprDoc callee, ffi::Array<Doc> args, ffi::Optional<DictDoc> workspace,
+ ffi::Optional<DictDoc> config,
+ ffi::Optional<ExprDoc> dispatch = std::nullopt);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpCallDoc, StmtDoc, OpCallDocNode);
};
+/*!
+ * \brief A reference to a canonical script namespace.
+ *
+ * The renderer resolves the configured alias without remapping ordinary identifiers.
+ */
+class NamespaceDocNode : public ExprDocNode {
+ public:
+ /*! \brief The namespace name before alias resolution. */
+ ffi::String canonical_name;
+
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<NamespaceDocNode>().def_ro("canonical_name",
+ &NamespaceDocNode::canonical_name);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.NamespaceDoc", NamespaceDocNode, ExprDocNode);
+};
+
+/*!
+ * \brief An owning reference to a deferred script namespace.
+ */
+class NamespaceDoc : public ExprDoc {
+ public:
+ /*!
+ * \brief Construct a deferred namespace reference.
+ * \param canonical_name The namespace name before alias resolution.
+ */
+ TVM_DLL explicit NamespaceDoc(ffi::String canonical_name);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NamespaceDoc, ExprDoc, NamespaceDocNode);
+};
+
+/*!
+ * \brief Render a completed Doc tree with the supplied configuration.
+ * \param doc The Doc tree to render.
+ * \param config The rendering configuration.
+ * \return The rendered Python source.
+ */
+TVM_DLL ffi::String DocToPythonScript(Doc doc, const PrinterConfig& config);
+
+/*!
+ * \brief Render explicit headers followed by a completed Doc tree.
+ * \param doc The Doc tree to render.
+ * \param config The rendering configuration.
+ * \param header Ordered String source chunks or CommentDocs.
+ * \return The rendered Python source.
+ */
+TVM_DLL ffi::String DocToPythonScriptWithHeader(Doc doc, const PrinterConfig& config,
+ const ffi::Array<ffi::Any>& header);
+
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/include/tvm/script/printer/doc_translator.h b/include/tvm/script/printer/doc_translator.h
new file mode 100644
index 0000000000..d3bdefc7c2
--- /dev/null
+++ b/include/tvm/script/printer/doc_translator.h
@@ -0,0 +1,422 @@
+/*
+ * 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_DOC_TRANSLATOR_H_
+#define TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_H_
+
+#include <tvm/ffi/container/dict.h>
+#include <tvm/ffi/container/list.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/expected.h>
+#include <tvm/ffi/reflection/accessor.h>
+#include <tvm/ir/expr.h>
+#include <tvm/script/printer/doc.h>
+
+#include <optional>
+#include <utility>
+
+namespace tvm {
+
+namespace script {
+namespace printer {
+
+class DocTranslatorObj;
+
+/*!
+ * \brief Type attribute containing a native or packed translation hook.
+ */
+inline constexpr const char* kDocTranslate = "__tvm_doc_translate__";
+/*!
+ * \brief Op-specific Call translation hook with the FDocTranslate contract.
+ */
+inline constexpr const char* kOpCallDocTranslate = "__tvm_doc_translate_op_call__";
+/*!
+ * \brief Source-Type hook receiving the original TensorLoad and optional binder.
+ */
+inline constexpr const char* kTensorLoadDocTranslate = "__tvm_doc_translate_tensor_load__";
+
+/*!
+ * \brief Engine callbacks borrowed for the translator's lifetime.
+ */
+struct DocTranslatorVTable {
+ /*!
+ * \brief Translate a value with an optional pre-mapped destination.
+ * \param translator The borrowed translation context.
+ * \param value The borrowed input value.
+ * \param destination A borrowed Var, or nullptr when no binding is requested.
+ * \return An owning Expected<Optional<ExprDoc>> encoded as TVMFFIAny. None means
+ * emission completed the request; errors use Expected's error encoding.
+ */
+ TVMFFIAny (*translate)(DocTranslatorObj* translator, ffi::AnyView value,
+ const ffi::Object* destination) noexcept;
+ /*!
+ * \brief Reserve an identifier without mapping a variable.
+ * \param translator The borrowed translation context.
+ * \param hint The borrowed string name hint.
+ * \return An owning Expected<IdDoc> encoded as TVMFFIAny, including errors.
+ */
+ TVMFFIAny (*alloc_id)(DocTranslatorObj* translator, ffi::AnyView hint) noexcept;
+ /*!
+ * \brief Find or allocate a variable's identifier.
+ * \param translator The borrowed translation context.
+ * \param var The borrowed variable.
+ * \param explicit_def Whether to remove the binding from pending definitions.
+ * \return An owning Expected<IdDoc> encoded as TVMFFIAny, including errors.
+ */
+ TVMFFIAny (*var_get_or_alloc_id)(DocTranslatorObj* translator, ffi::AnyView var,
+ bool explicit_def) noexcept;
+ /*!
+ * \brief Get the current variable scope's pending definitions.
+ * \param translator The borrowed translation context.
+ * \return An owning Expected<Dict<Var, IdDoc>> encoded as TVMFFIAny, including errors.
+ */
+ TVMFFIAny (*get_implicit_defs)(DocTranslatorObj* translator) noexcept;
+ /*!
+ * \brief Begin a lexical variable scope.
+ * \param translator The borrowed translation context.
+ * \return Expected<void> encoded as TVMFFIAny, including errors.
+ */
+ TVMFFIAny (*begin_var_scope)(DocTranslatorObj* translator) noexcept;
+ /*!
+ * \brief End the current lexical variable scope.
+ * \param translator The borrowed translation context.
+ * \return Expected<void> encoded as TVMFFIAny, including errors.
+ */
+ TVMFFIAny (*end_var_scope)(DocTranslatorObj* translator) noexcept;
+};
+
+/*!
+ * \brief Native translation hook with borrowed inputs and an owning result.
+ * \param translator Translation context borrowed for the duration of the call.
+ * \param value Input value borrowed for the duration of the call.
+ * \param destination Optional binding variable borrowed for the duration of the call.
+ * \return An owning expression Doc, or None after completing emission.
+ * Return an expression for the caller to consume, or None after completing emission
+ * and any requested binding. Packed hooks use the same contract and None for a null
+ * destination; native errors cross the ABI through Expected.
+ */
+using FDocTranslate = ffi::reflection::NativeFunctionView<ffi::Optional<ExprDoc>(
+ DocTranslatorObj* translator, ffi::AnyView value, const ffi::Object* destination)>;
+
+/*!
+ * \brief Eager translation context shared by registered hooks.
+ * Hooks construct Docs, use explicit lexical scopes, and attach origins when emitting.
+ *
+ * \code{.cpp}
+ * ffi::Optional<ExprDoc> Hook(DocTranslatorObj* d, ffi::AnyView value,
+ * const ffi::Object* destination) {
+ * ExprDoc rhs = LiteralDoc::Str("value", std::nullopt);
+ * if (!destination) return rhs;
+ * auto var = ffi::GetRef<Var>(static_cast<const VarNode*>(destination));
+ * IdDoc lhs(d->VarGetOrAllocId(var, true)->name);
+ * d->Emit(AssignDoc(lhs, rhs, std::nullopt), var);
+ * return std::nullopt;
+ * }
+ * // Register with TypeAttrDef<Node>().attr(kDocTranslate, FDocTranslate::FromNative<&Hook>()).
+ * \endcode
+ */
+class DocTranslatorObj : public ffi::Object {
+ public:
+ /*!
+ * \brief Translate a value with an optional pre-mapped binding variable.
+ * \param value The input value.
+ * \param bind_var The binding variable, or None when no binding is requested.
+ * \return An expression for the caller, or None when emission completed the request.
+ * Callers requiring an expression use value(); for a specific Doc subtype,
+ * use value().as_or_throw<T>().
+ */
+ ffi::Optional<ExprDoc> Translate(ffi::AnyView value,
+ const ffi::Optional<Var> bind_var = std::nullopt) {
+ auto result = ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>(
+ vtable_->translate(this, value,
+ bind_var.has_value() ? bind_var.value().get() : nullptr))
+ .value();
+ if (result.has_value()) {
+ if (const auto* origin = value.as<ffi::Object>()) {
+ RecordOrigin(result.value(), ffi::GetRef<ffi::ObjectRef>(origin));
+ }
+ }
+ return result;
+ }
+ /*!
+ * \brief Append a completed Doc and record its optional origin.
+ * \param doc The emitted Doc.
+ * \param origin The original IR object represented by this emission.
+ */
+ void Emit(const Doc& doc, ffi::Optional<ffi::ObjectRef> origin = std::nullopt) {
+ CurrentScopeDocs().push_back(doc);
+ if (origin.has_value()) RecordOrigin(doc, origin.value());
+ }
+
+ /*!
+ * \brief Reuse or allocate a variable's stable identifier.
+ * \param var The variable to name.
+ * \param explicit_def Whether this occurrence defines the variable.
+ * \return The binding's IdDoc, with explicit definitions removed from pending entries.
+ */
+ IdDoc VarGetOrAllocId(const Var& var, bool explicit_def) {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<IdDoc>(
+ vtable_->var_get_or_alloc_id(this, var, explicit_def))
+ .value();
+ }
+ /*!
+ * \brief Reserve a fresh identifier without mapping a variable.
+ * \param hint The preferred name.
+ * \return The allocated identifier.
+ */
+ IdDoc AllocId(const ffi::String& hint) {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<IdDoc>(vtable_->alloc_id(this, hint))
+ .value();
+ }
+ /*!
+ * \brief Get the current variable scope's pending definitions.
+ * \return The live mutable dictionary of variables and their identifiers.
+ */
+ ffi::Dict<Var, IdDoc> GetImplicitDefs() {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Dict<Var, IdDoc>>(
+ vtable_->get_implicit_defs(this))
+ .value();
+ }
+
+ /*!
+ * \brief Read a typed extension option.
+ * \tparam T The option type.
+ * \param name The option name.
+ * \param fallback The value to use when absent.
+ * \return The typed option or fallback, rejecting an existing wrong type.
+ */
+ template <typename T>
+ T GetExtraConfig(const ffi::String& name, T fallback) const {
+ auto it = extra_config_.find(name);
+ return it == extra_config_.end() ? fallback : (*it).second.as_or_throw<T>();
+ }
+ /*!
+ * \brief Get or initialize typed extension state.
+ * \tparam T The state type, default constructed when absent.
+ * \param name The state name.
+ * \return The stored value, rejecting an existing wrong type; mutable handles share storage.
+ */
+ template <typename T>
+ T GetOrCreateExtraState(const ffi::String& name) {
+ auto it = extra_state_.find(name);
+ if (it != extra_state_.end()) return (*it).second.as_or_throw<T>();
+ T value{};
+ extra_state_.Set(name, value);
+ return value;
+ }
+ /*!
+ * \brief Replace or remove extension state.
+ * \param name The state name.
+ * \param value The replacement value, or nullopt to remove the entry.
+ * \return The previous value, or nullopt when absent.
+ */
+ std::optional<ffi::Any> ExchangeExtraState(const ffi::String& name,
+ std::optional<ffi::Any> value) {
+ auto previous = extra_state_.Get(name);
+ if (value.has_value()) {
+ extra_state_.Set(name, std::move(value.value()));
+ } else {
+ extra_state_.erase(name);
+ }
+ return previous;
+ }
+
+ /*!
+ * \brief Begin a lexical variable scope independent of Doc collection.
+ */
+ void BeginVarScope() {
+ ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<void>(vtable_->begin_var_scope(this)).value();
+ }
+ /*!
+ * \brief End the latest naming region, retaining the root region.
+ * Explicit local bindings expire. Unresolved implicit references and their
+ * identifiers remain pending in the enclosing region, preserving captures
+ * shared by sibling functions.
+ */
+ void EndVarScope() {
+ ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<void>(vtable_->end_var_scope(this)).value();
+ }
+ /*!
+ * \brief Push an empty Doc collection.
+ */
+ void BeginDocScope() { doc_scopes_.push_back(ffi::List<Doc>()); }
+ /*!
+ * \brief Pop the latest Doc collection.
+ * \return The collected Docs.
+ */
+ ffi::List<Doc> EndDocScope() {
+ auto docs = CurrentScopeDocs();
+ doc_scopes_.pop_back();
+ return docs;
+ }
+ /*!
+ * \brief Get the current Doc collection without popping it.
+ * \return A mutable handle sharing the active collection's storage.
+ */
+ ffi::List<Doc> CurrentScopeDocs() const {
+ TVM_FFI_CHECK(!doc_scopes_.empty(), ValueError)
+ << "printer emission requires a collection scope";
+ return doc_scopes_.back();
+ }
+ /*!
+ * \brief Run a callback in a balanced Doc collection scope.
+ * \tparam Callback The callable type.
+ * \param callback The callback that emits Docs.
+ * \return The collected Docs.
+ */
+ template <typename Callback>
+ ffi::List<Doc> WithDocScope(Callback&& callback) {
+ size_t depth = doc_scopes_.size();
+ BeginDocScope();
+ try {
+ std::forward<Callback>(callback)();
+ TVM_FFI_CHECK(doc_scopes_.size() == depth + 1, ValueError)
+ << "printer unbalanced Doc collection scopes";
+ return EndDocScope();
+ } catch (...) {
+ while (doc_scopes_.size() > depth) doc_scopes_.pop_back();
+ throw;
+ }
+ }
+
+ /*!
+ * \brief Get the owned Doc-to-original-object associations.
+ * \return A mutable handle sharing the origin dictionary's storage.
+ */
+ ffi::Dict<Doc, ffi::ObjectRef> GetDocOrigins() const { return doc_origins_; }
+ /*!
+ * \brief Set a Doc's original IR object, replacing any previous association.
+ * \param doc The constructed Doc.
+ * \param origin The original IR object.
+ */
+ void RecordOrigin(const Doc& doc, const ffi::ObjectRef& origin) { doc_origins_.Set(doc, origin); }
+ /*!
+ * \brief Dispatch directly to the registered native or packed hook.
+ * \param value The input value.
+ * \param destination A borrowed Var, or nullptr when no binding is requested.
+ * \return An expression for the caller, or None after completed emission.
+ */
+ ffi::Optional<ExprDoc> DefaultTranslate(ffi::AnyView value,
+ const ffi::Object* destination = nullptr) {
+ static ffi::reflection::TypeAttrColumn column(kDocTranslate);
+ ffi::AnyView attr = column[value.type_index()];
+ if (attr.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>(
+ reinterpret_cast<decltype(DocTranslatorVTable::translate)>(attr.cast<void*>())(
+ this, value, destination))
+ .value();
+ }
+ if (attr.type_index() == ffi::TypeIndex::kTVMFFIFunction) {
+ ffi::Any destination_arg = nullptr;
+ if (destination) destination_arg = ffi::GetRef<ffi::ObjectRef>(destination);
+ return attr.cast<ffi::Function>()
+ .CallExpected<ffi::Optional<ExprDoc>>(this, value, destination_arg)
+ .value();
+ }
+ TVM_FFI_THROW(TypeError) << (attr.type_index() == ffi::TypeIndex::kTVMFFINone
+ ? std::string("printer has no Doc hook for ") +
+ value.GetTypeKey()
+ : "printer Doc hook must be a native pointer or ffi.Function");
+ }
+
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO("script.printer.DocTranslator", DocTranslatorObj, ffi::Object);
+
+ protected:
+ /*!
+ * \brief Initialize an engine's translation context.
+ * \param vtable The engine callbacks, valid for this object's lifetime.
+ * \param extra_config The owned configuration snapshot.
+ */
+ explicit DocTranslatorObj(const DocTranslatorVTable* vtable,
+ ffi::Map<ffi::String, ffi::Any> extra_config = {})
+ : extra_config_(std::move(extra_config)), vtable_(vtable) {}
+
+ /*!
+ * \brief Active Doc collections, independent of lexical variable scopes.
+ */
+ ffi::List<ffi::List<Doc>> doc_scopes_;
+
+ private:
+ /*!
+ * \brief Owned extension configuration.
+ * Options are read eagerly by translation hooks.
+ */
+ ffi::Map<ffi::String, ffi::Any> extra_config_;
+ /*!
+ * \brief Owned mutable extension state.
+ * Values retain their objects and mutable handles share storage.
+ */
+ ffi::Dict<ffi::String, ffi::Any> extra_state_;
+ /*!
+ * \brief Owned associations between Docs and their original IR objects.
+ */
+ ffi::Dict<Doc, ffi::ObjectRef> doc_origins_;
+ /*!
+ * \brief Borrowed engine callbacks valid for this object's lifetime.
+ */
+ const DocTranslatorVTable* vtable_;
+};
+
+/*!
+ * \brief Convert collected Docs to statements in their original order.
+ * \param docs The Docs to convert, without modifying the collection.
+ * \return An owning array containing each StmtBlockDoc's statements in place,
+ * and each other Doc converted with as_or_throw<StmtDoc>().
+ * \throws ffi::Error If a Doc is neither a StmtBlockDoc nor a StmtDoc.
+ */
+inline ffi::Array<StmtDoc> ToStmtDocArray(const ffi::List<Doc>& docs) {
+ ffi::Array<StmtDoc> statements;
+ for (const auto& doc : docs) {
+ if (auto block = doc.as<StmtBlockDoc>()) {
+ for (const auto& stmt : block.value()->stmts) statements.push_back(stmt);
+ } else {
+ statements.push_back(doc.as_or_throw<StmtDoc>());
+ }
+ }
+ return statements;
+}
+
+/*!
+ * \brief Owning reference to a hook translation context.
+ */
+class DocTranslator : public ffi::ObjectRef {
+ public:
+ /*!
+ * \brief Retain a translation context.
+ * \param n The context object.
+ */
+ explicit DocTranslator(ffi::ObjectPtr<DocTranslatorObj> n) : ObjectRef(std::move(n)) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DocTranslator, ffi::ObjectRef, DocTranslatorObj);
+};
+
+/*!
+ * \brief Translate IR into a complete Doc including declarations and metadata setup.
+ * \param ir The input IR value.
+ * \param doc_origins Optional output replaced with the owned origins only on success.
+ * \param extra_config The translation configuration snapshot.
+ * \return The translated Doc, with text headers and diagnostic recovery left to the text adapter.
+ */
+TVM_DLL Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_origins = nullptr,
+ ffi::Map<ffi::String, ffi::Any> extra_config = {});
+
+} // namespace printer
+} // namespace script
+} // namespace tvm
+
+#endif // TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_H_
diff --git a/include/tvm/script/printer/ir_docsifier.h b/include/tvm/script/printer/ir_docsifier.h
deleted file mode 100644
index 944e511457..0000000000
--- a/include/tvm/script/printer/ir_docsifier.h
+++ /dev/null
@@ -1,365 +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_IR_DOCSIFIER_H_
-#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
-
-#include <tvm/ffi/device.h>
-#include <tvm/ffi/reflection/access_path.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/module.h>
-#include <tvm/script/printer/config.h>
-#include <tvm/script/printer/doc.h>
-#include <tvm/script/printer/ir_docsifier_functor.h>
-
-#include <string>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace script {
-namespace printer {
-
-using AccessPath = ffi::reflection::AccessPath;
-
-//////////////////////// Frame ////////////////////////
-
-class IRDocsifier;
-class IRDocsifierNode;
-
-/*!
- * Frame is the core data structure for semantic information
- * when printing IR graph into TVMScript code.
- */
-class FrameNode : public ffi::Object {
- public:
- /*! The docs generated in the frame */
- ffi::Array<StmtDoc> stmts;
- /*! The corresponding IRDocsifier */
- IRDocsifierNode* d;
- /*! The callbacks that are going to be invoked when the frame exits */
- std::vector<std::function<void()>> callbacks;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<FrameNode>().def_ro("stmts", &FrameNode::stmts);
- }
-
- static constexpr const bool _type_mutable = true;
- TVM_FFI_DECLARE_OBJECT_INFO("script.printer.Frame", FrameNode, ffi::Object);
-
- public:
- virtual ~FrameNode() = default;
-
- /*!
- * \brief Add a callback function to be called when this frame exits.
- * \param cb The callback function. It should have signature void().
- */
- template <typename TCallback>
- void AddExitCallback(TCallback&& cb) {
- callbacks.emplace_back(std::forward<TCallback>(cb));
- }
- /*!
- * \brief Add a dispatch token to the docsifier, and a callback that pops the token when this
- * frame exits.
- * \param d The docsifier.
- * \param token The token to be added.
- */
- void AddDispatchToken(const IRDocsifier& d, const ffi::String& token);
- /*!
- * \brief Method that's called when Frame enters the scope.
- */
- virtual void EnterWithScope();
- /*!
- * \brief Method that's called when Frame exits the scope.
- */
- virtual void ExitWithScope();
-};
-
-/*!
- * \brief Reference type of FrameNode
- */
-class Frame : public ffi::ObjectRef {
- protected:
- Frame() = default;
-
- public:
- virtual ~Frame() = default;
-
- /*! \brief Method that's called when Frame enters the scope. */
- void EnterWithScope() { get()->EnterWithScope(); }
-
- /*! \brief Method that's called when Frame exits the scope. */
- void ExitWithScope() { get()->ExitWithScope(); }
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Frame, ffi::ObjectRef, FrameNode);
-};
-
-//////////////////////// IRDocsifier ////////////////////////
-
-/*!
- * \brief IRDocsifier is the top-level interface in the IR->Doc process.
- *
- * It provides methods to convert IR node object to Doc, operate on Frame
- * objects and change dispatch tokens.
- */
-class IRDocsifierNode : public ffi::Object {
- public:
- /*! \brief A function that creates the doc for a variable */
- using DocCreator = std::function<ExprDoc()>;
- /*! \brief Information about a variable, including its optional name and its doc creator */
- struct VariableInfo {
- /*! \brief The creator */
- DocCreator creator;
- /*! \brief The name of the variable */
- ffi::Optional<ffi::String> name;
- };
- /*! \brief The configuration of the printer */
- PrinterConfig cfg{ffi::UnsafeInit()};
- /*!
- * \brief The stack of frames.
- * \sa FrameNode
- */
- ffi::Array<Frame> frames;
- /*!
- * \brief The stack of dispatch tokens.
- *
- * The dispatch token on the top decides which dispatch function to use
- * when converting IR node object to Doc.
- */
- ffi::Array<ffi::String> dispatch_tokens;
- /*! \brief Mapping from a var to its info */
- std::unordered_map<ffi::ObjectRef, VariableInfo, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- obj2info;
- /*! \brief Metadata printing */
- std::unordered_map<ffi::String, ffi::Array<ffi::Any>> metadata;
- /*! \brief GlobalInfo printing */
- std::unordered_map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
- /*! \brief The variable names used already */
- std::unordered_set<ffi::String> defined_names;
- /*! \brief Common prefixes of variable usages */
- std::unordered_map<const ffi::Object*, std::vector<const ffi::Object*>> common_prefix;
- /*! \brief The IR usages for headers printing */
- std::unordered_set<std::string> ir_usage;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IRDocsifierNode>()
- .def_ro("frames", &IRDocsifierNode::frames)
- .def_ro("dispatch_tokens", &IRDocsifierNode::dispatch_tokens);
- }
-
- static constexpr const bool _type_mutable = true;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IRDocsifier", IRDocsifierNode, ffi::Object);
-
- public:
- /*!
- * \brief Define variable by name.
- * \param obj The variable object.
- * \param frame The frame that this variable is defined in.
- * \param name_hint The hint for variable name.
- *
- * \return The id doc for this variable.
- *
- * This function will rename the variable to avoid name conflict with other variables
- * in the table.
- */
- IdDoc Define(const ffi::ObjectRef& obj, const Frame& frame, const ffi::String& name_hint);
-
- /*!
- * \brief Define variable by doc factory.
- * \param obj The variable object.
- * \param frame The frame that this variable is defined in.
- * \param doc_factory The function to return an ExprDoc object for this variable.
- *
- * This function is a special form of `Define`. Variable is mapped to ExprDoc rather
- * than IdDoc. It's useful when a variable is implicitly defined without a name, like
- * the buf->data in TIR, which should be mapped to `AttrDoc(IdDoc("<buffer_name>"), "data")`.
- *
- * This function takes a DocFactory instead of Doc. It's because GetVarDoc needs to
- * return a new Doc object every time it's called, as the returned doc will have
- * different `source_path`. Currently there isn't a good way to deep copy a TVMObject
- * so VarTable needs to call a factory function to get a freshly-constructed Doc object
- * every time GetVarDoc is called.
- */
- void Define(const ffi::ObjectRef& obj, const Frame& frame, DocCreator doc_factory);
-
- /*!
- * \brief Get the doc for variable.
- * \param obj The variable object.
- *
- * \return The doc for variable, if it exists in the table. Otherwise it returns std::nullopt.
- */
- ffi::Optional<ExprDoc> GetVarDoc(const ffi::ObjectRef& obj) const;
- /*! \brief Add a TVM object to the metadata section*/
- ExprDoc AddMetadata(const ffi::Any& obj);
- /*! \brief Add a GlobalInfo to the global_infos map.
- * \param name The name of key of global_infos.
- * \param ginfo The GlobalInfo to be added.
- */
- void AddGlobalInfo(const ffi::String& name, const GlobalInfo& ginfo);
- /*!
- * \brief Check if a variable exists in the table.
- * \param obj The variable object.
- *
- * \return a boolean for whether variable exists.
- */
- bool IsVarDefined(const ffi::ObjectRef& obj) const;
- /*! \brief Remove the variable defined */
- void RemoveVar(const ffi::ObjectRef& obj);
- /*!
- * \brief Set the common prefix information of variable usage.
- * \param root The root of the AST.
- * \param is_var A function that returns true if the given object is considered a variable.
- */
- void SetCommonPrefix(const ffi::ObjectRef& root, ffi::TypedFunction<bool(ffi::ObjectRef)> is_var);
- /*!
- * \brief Transform the input object into TDoc.
- * \param obj The object to be transformed.
- * \param path The path to this object.
- *
- * \return The Doc object.
- */
- template <class TDoc = Doc>
- inline TDoc AsDoc(const Any& obj, const AccessPath& path) const;
-};
-
-/*!
- * \brief Reference type of IRDocsifierNode.
- */
-class IRDocsifier : public ffi::ObjectRef {
- public:
- using FType = IRDocsifierFunctor<printer::Doc, AccessPath, IRDocsifier>;
- /*! \brief Create a IRDocsifier. */
- explicit IRDocsifier(const PrinterConfig& cfg);
- /*! \brief The registration table for IRDocsifier. */
- TVM_DLL static FType& vtable();
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRDocsifier, ffi::ObjectRef, IRDocsifierNode);
-};
-
-//////////////////////// Implementation ////////////////////////
-
-inline void FrameNode::EnterWithScope() {
- if (d != nullptr) {
- d->frames.push_back(ffi::GetRef<Frame>(this));
- }
-}
-
-inline void FrameNode::ExitWithScope() {
- for (const std::function<void()>& callback : callbacks) {
- callback();
- }
- callbacks.clear();
- if (d != nullptr) {
- d->frames.pop_back();
- }
-}
-
-template <class TDoc>
-inline static void AddDocDecoration(const Doc& d, const ffi::ObjectRef& obj, const AccessPath& path,
- const PrinterConfig& cfg) {
- if (cfg->obj_to_annotate.count(obj)) {
- if (const auto* stmt = d.as<StmtDocNode>()) {
- if (stmt->comment.has_value()) {
- stmt->comment = stmt->comment.value() + "\n" + cfg->obj_to_annotate.at(obj);
- } else {
- stmt->comment = cfg->obj_to_annotate.at(obj);
- }
- } else {
- LOG(WARNING) << "Expect StmtDoc to be annotated for object " << obj << ", but got "
- << d.as_or_throw<TDoc>()->_type_key;
- }
- }
- for (const ffi::ObjectRef& o : cfg->obj_to_underline) {
- if (o.same_as(obj)) {
- cfg->path_to_underline.push_back(path);
- }
- }
- for (const auto& pair : cfg->path_to_annotate) {
- AccessPath p = pair.first;
- ffi::String attn = pair.second;
- if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
- if (const auto* stmt = d.as<StmtDocNode>()) {
- if (stmt->comment.has_value()) {
- stmt->comment = stmt->comment.value() + "\n" + attn;
- } else {
- stmt->comment = attn;
- }
- } else {
- LOG(WARNING) << "Expect StmtDoc to be annotated at object path " << p << ", but got "
- << d.as_or_throw<TDoc>()->_type_key;
- }
- }
- }
-}
-
-template <class TDoc>
-inline TDoc IRDocsifierNode::AsDoc(const Any& value, const AccessPath& path) const {
- switch (value.type_index()) {
- case ffi::TypeIndex::kTVMFFINone:
- return LiteralDoc::None(path).as_or_throw<TDoc>();
- case ffi::TypeIndex::kTVMFFIBool:
- return LiteralDoc::Boolean(value.as<bool>().value(), path).as_or_throw<TDoc>();
- case ffi::TypeIndex::kTVMFFIInt:
- return LiteralDoc::Int(value.as<int64_t>().value(), path).as_or_throw<TDoc>();
- case ffi::TypeIndex::kTVMFFIFloat:
- return LiteralDoc::Float(value.as<double>().value(), path).as_or_throw<TDoc>();
- case ffi::TypeIndex::kTVMFFISmallStr:
- case ffi::TypeIndex::kTVMFFIStr: {
- std::string string_value = value.cast<std::string>();
- bool has_multiple_lines = string_value.find_first_of('\n') != std::string::npos;
- if (has_multiple_lines) {
- Doc d = const_cast<IRDocsifierNode*>(this)->AddMetadata(string_value);
- // TODO(tqchen): cross check AddDocDecoration
- return d.as_or_throw<TDoc>();
- }
- return LiteralDoc::Str(string_value, path).as_or_throw<TDoc>();
- }
- case ffi::TypeIndex::kTVMFFIDataType:
- return LiteralDoc::DataType(value.as<DLDataType>().value(), path).as_or_throw<TDoc>();
- case ffi::TypeIndex::kTVMFFIDevice:
- return LiteralDoc::Device(value.as<DLDevice>().value(), path).as_or_throw<TDoc>();
- default: {
- if (auto opt_obj = value.as<ffi::ObjectRef>()) {
- ffi::ObjectRef obj = opt_obj.value();
- Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path,
- ffi::GetRef<IRDocsifier>(this));
- d->source_paths.push_back(path);
- AddDocDecoration<TDoc>(d, obj, path, cfg);
- return d.as_or_throw<TDoc>();
- } else {
- TVM_FFI_THROW(TypeError) << "Cannot handle Any type: `" << value.GetTypeKey() << "`";
- TVM_FFI_UNREACHABLE();
- }
- }
- }
-}
-
-inline void FrameNode::AddDispatchToken(const IRDocsifier& d, const ffi::String& token) {
- d->dispatch_tokens.push_back(token);
- this->AddExitCallback([doc = d.get()]() { doc->dispatch_tokens.pop_back(); });
-}
-
-} // namespace printer
-} // namespace script
-} // namespace tvm
-
-#endif // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
diff --git a/include/tvm/script/printer/ir_docsifier_functor.h b/include/tvm/script/printer/ir_docsifier_functor.h
deleted file mode 100644
index 6d4d598633..0000000000
--- a/include/tvm/script/printer/ir_docsifier_functor.h
+++ /dev/null
@@ -1,199 +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_IR_DOCSIFIER_FUNCTOR_H_
-#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
-
-#include <tvm/ffi/function.h>
-#include <tvm/runtime/logging.h>
-
-#include <optional>
-#include <string>
-#include <type_traits>
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace script {
-namespace printer {
-
-/*!
- * \brief Dynamic dispatch functor based on AccessPath.
- *
- * This functor dispatches based on the type of object and the input dispatch token.
- */
-template <typename R, typename... Args>
-class IRDocsifierFunctor {
- private:
- using TSelf = IRDocsifierFunctor<R, Args...>;
-
- template <class TObjectRef, class TCallable>
- using IsDispatchFunction =
- typename std::is_convertible<TCallable, std::function<R(TObjectRef, Args...)>>;
-
- public:
- /*!
- * \brief Call the dispatch function.
- * \param token The dispatch token.
- * \param obj The object.
- * \param args Other args.
- *
- * \return The return value of the dispatch function
- *
- * If the TObjectRef isn't registered with the token, it will try to find
- * dispatch function for TObjectRef with the default dispatch token (empty string).
- */
- template <class TObjectRef>
- R operator()(const ffi::String& token, TObjectRef obj, Args... args) const {
- uint32_t type_index = obj.defined() ? obj->type_index() : 0;
- const ffi::Function* pf = nullptr;
- if ((pf = LookupDispatchTable(token, type_index)) != nullptr) {
- return (*pf)(obj, args...).template cast<R>();
- }
- if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
- return (*pf)(obj, args...).template cast<R>();
- }
- if ((pf = LookupFallback()) != nullptr) {
- return (*pf)(obj, args...).template cast<R>();
- }
-
- LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
- << ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
- << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
- TVM_FFI_ICHECK(false) << "ObjectFunctor calls un-registered function on type: "
- << ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
- << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
- TVM_FFI_UNREACHABLE();
- }
-
- /*!
- * \brief Set the dispatch function
- * \param token The dispatch token.
- * \param type_index The TVM object type index for this dispatch function.
- * \param f The dispatch function.
- *
- * This takes a type-erased packed function as input. It should be used
- * through FFI boundary, for example, registering dispatch function from Python.
- */
- TSelf& set_dispatch(ffi::String token, uint32_t type_index, ffi::Function f) {
- std::vector<ffi::Function>* table = &dispatch_table_[token];
- if (table->size() <= type_index) {
- table->resize(type_index + 1, nullptr);
- }
- ffi::Function& slot = (*table)[type_index];
- if (slot != nullptr) {
- TVM_FFI_ICHECK(false) << "Dispatch for type is already registered: "
- << ffi::Object::TypeIndex2Key(type_index);
- }
- slot = f;
- return *this;
- }
-
- TSelf& set_fallback(ffi::Function f) {
- TVM_FFI_ICHECK(!dispatch_fallback_.has_value()) << "Fallback is already defined";
- dispatch_fallback_ = f;
- return *this;
- }
-
- void remove_fallback() { dispatch_fallback_ = std::nullopt; }
-
- /*!
- * \brief Set the dispatch function
- * \param token The dispatch token.
- * \param f The dispatch function.
- */
- template <typename TObjectRef, typename TCallable,
- typename = std::enable_if_t<IsDispatchFunction<TObjectRef, TCallable>::value>>
- TSelf& set_dispatch(ffi::String token, TCallable f) {
- return set_dispatch(token, TObjectRef::ContainerType::RuntimeTypeIndex(),
- ffi::TypedFunction<R(TObjectRef, Args...)>(f));
- }
-
- template <typename TCallable,
- typename = std::enable_if_t<IsDispatchFunction<ffi::ObjectRef, TCallable>::value>>
- TSelf& set_fallback(TCallable f) {
- ffi::Function func = ffi::TypedFunction<R(ffi::ObjectRef, Args...)>(f);
- return set_fallback(func);
- }
-
- /*!
- * \brief Remove dispatch function
- * \param token The dispatch token.
- * \param type_index The TVM object type index for the dispatch function to be removed.
- *
- * This is useful when dispatch function comes from other language's runtime, and
- * those function should be removed before that language runtime shuts down.
- */
- void remove_dispatch(ffi::String token, uint32_t type_index) {
- std::vector<ffi::Function>* table = &dispatch_table_[token];
- if (table->size() <= type_index) {
- return;
- }
- (*table)[type_index] = nullptr;
- }
-
- private:
- /*!
- * \brief Look up the dispatch table for the given token and type_index.
- * \param token The dispatch token.
- * \param type_index The TVM object type index.
- * \return Returns the functor if the lookup succeeds, nullptr otherwise.
- */
- const ffi::Function* LookupDispatchTable(const ffi::String& token, uint32_t type_index) const {
- auto it = dispatch_table_.find(token);
- if (it == dispatch_table_.end()) {
- return nullptr;
- }
- const std::vector<ffi::Function>& tab = it->second;
- if (type_index >= tab.size()) {
- return nullptr;
- }
- const ffi::Function* f = &tab[type_index];
- if (f->defined()) {
- return f;
- } else {
- return nullptr;
- }
- }
-
- /*!
- * \brief Look up the fallback to be used if no handler is registered
- */
- const ffi::Function* LookupFallback() const {
- if (dispatch_fallback_.has_value()) {
- return &*dispatch_fallback_;
- } else {
- return nullptr;
- }
- }
-
- /*
- * This type alias and the following free functions are created to reduce the binary bloat
- * from template and also hide implementation details from this header
- */
- using DispatchTable = std::unordered_map<std::string, std::vector<ffi::Function>>;
- /*! \brief The dispatch table. */
- DispatchTable dispatch_table_;
- std::optional<ffi::Function> dispatch_fallback_;
-};
-
-} // namespace printer
-} // namespace script
-} // namespace tvm
-#endif // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
diff --git a/include/tvm/script/printer/printer.h b/include/tvm/script/printer/printer.h
index 93d5a0d64e..83f313013b 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -18,57 +18,41 @@
*/
/*!
* \file tvm/script/printer/printer.h
- * \brief Entry-point header for TVMScript printing.
- *
- * Declares the free function `tvm::Script(node, optional_config)` and the
- * dispatch vtable `TVMScriptPrinter::vtable()` used by per-dialect printers.
- * `PrinterConfig` and its dataclass helpers live in config.h; this header is
- * what callers include to invoke printing.
+ * \brief TVMScript text entry points.
*/
#ifndef TVM_SCRIPT_PRINTER_PRINTER_H_
#define TVM_SCRIPT_PRINTER_PRINTER_H_
-#include <tvm/ir/object_functor.h>
+#include <tvm/ffi/optional.h>
#include <tvm/script/printer/config.h>
+#include <optional>
+#include <string>
+
namespace tvm {
-/*! \brief Print \p node as TVMScript with the given \p config.
- *
- * Falls back to ffi::ReprPrint for types not registered with TVMScriptPrinter.
+/*!
+ * \brief Print an IR object as TVMScript, using repr when no translation hook exists.
+ * \param node The input IR object.
+ * \param config Optional translation and rendering configuration.
+ * \return The rendered script or fallback representation.
*/
TVM_DLL std::string Script(const ffi::ObjectRef& node,
const ffi::Optional<PrinterConfig>& config = std::nullopt);
-/*! \brief Dispatch table for TVMScript printing and repr registration. */
-class TVMScriptPrinter {
- public:
- using FType = ObjectFunctor<std::string(const ffi::ObjectRef&, const PrinterConfig&)>;
- TVM_DLL static FType& vtable();
+namespace script {
+namespace printer {
- /*! \brief Register a printer method for script dispatch and FFI repr.
- * \tparam ObjectType Concrete object node type.
- * \tparam Method Callable printer method type.
- * \param method Printer dispatch method for that type.
- *
- * Example:
- * \code
- * TVM_FFI_STATIC_INIT_BLOCK() {
- * TVMScriptPrinter::Register<tirx::ForNode>(ReprPrintTIR);
- * TVMScriptPrinter::Register<tirx::WhileNode>(ReprPrintTIR);
- * }
- * \endcode
- */
- template <typename ObjectType, typename Method>
- static void Register(Method method) {
- namespace refl = tvm::ffi::reflection;
- refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,
- [](ffi::ObjectRef obj, ffi::Function) -> ffi::String {
- return RedirectedReprPrinterMethod(obj);
- });
- vtable().SetDispatch<ObjectType>(method);
- }
-};
+/*!
+ * \brief Translate IR, recover diagnostic paths, and render Python text.
+ * \param obj The input IR object.
+ * \param config The translation and rendering options.
+ * \return The rendered script.
+ */
+TVM_DLL ffi::String Script(const ffi::ObjectRef& obj, const PrinterConfig& config);
+} // namespace printer
+} // namespace script
} // namespace tvm
+
#endif // TVM_SCRIPT_PRINTER_PRINTER_H_
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index db185cdf47..8ca859ae72 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -55,7 +55,7 @@ using FLowerIntrinsic = ffi::TypedFunction<PrimExpr(PrimExpr)>;
using FLegalize = ffi::TypedFunction<PrimExpr(PrimExpr)>;
/*!
- * \brief The operator's name in TVMScript printer
+ * \brief The fully qualified TVMScript name, including its dialect namespace.
*/
using TScriptPrinterName = ffi::String;
diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py
index 3a6ee65bb0..587288f1e7 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -105,9 +105,10 @@ def register_table(table: dict[str, InstructionEntry]) -> None:
# `PTXNamespace.__getattr__`. The escape is the identity for every
# other family, and `gen_stubs` already spells the attribute this way.
family = escape_token(entry.family)
- register_op_attr(entry.op_name, "TScriptPrinterName", f"ptx.{family}")
+ register_op_attr(entry.op_name, "TScriptPrinterName", f"tirx.ptx.{family}")
register_op_attr(entry.op_name, "TIRxOpCategory", "device_intrin")
register_op_attr(entry.op_name, "TDeviceIntrinsicNamespace", "ptx")
+ register_op_attr(entry.op_name, "TFixedReturnType", PrimType("void"))
register_codegen(f"ptx.{entry.name}")(_make_codegen(entry))
@@ -117,9 +118,10 @@ def register_addr() -> None:
op = Op.get(_ADDR_OP_NAME)
op.add_arg("base", "The base pointer.")
op.add_arg("byte_offset", "The offset in bytes.")
- register_op_attr(_ADDR_OP_NAME, "TScriptPrinterName", "ptx.addr")
+ register_op_attr(_ADDR_OP_NAME, "TScriptPrinterName", "tirx.ptx.addr")
register_op_attr(_ADDR_OP_NAME, "TIRxOpCategory", "device_intrin")
register_op_attr(_ADDR_OP_NAME, "TDeviceIntrinsicNamespace", "ptx")
+ register_op_attr(_ADDR_OP_NAME, "FInferType", _infer_addr_type)
register_codegen("ptx.addr")(_unconsumed_addr_codegen)
@@ -130,6 +132,10 @@ def _unconsumed_addr_codegen(*_args):
)
+def _infer_addr_type(call):
+ return call.args[0].ty
+
+
# ---------------------------------------------------------------------------
# Codegen (compile time): table -> asm volatile helper
# ---------------------------------------------------------------------------
@@ -813,7 +819,9 @@ class _InstrChain:
# ascending order, to restore the operand positions) and then the
# register-class tags, whose indices are those same positions.
args = list(args)
- for flag in sorted(f for f in flags if f.startswith("s")):
+ for flag in sorted(
+ (f for f in flags if f.startswith("s")), key=lambda flag: int(flag[1:])
+ ):
args.insert(int(flag[1:]), SINK)
for flag in flags:
if flag.startswith("p") and flag != "pred":
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index e593783ac0..fa7bba369b 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -30,10 +30,13 @@ class Node(Object):
"""Base class of all IR Nodes."""
def __repr__(self) -> str:
- from tvm.runtime.script_printer import _script
+ from tvm.runtime.script_printer import PrinterConfig, _script
try:
- return _script(self, None)
+ return _script(
+ self,
+ PrinterConfig(extra_config={"script.comment_imports": True}),
+ )
except Exception:
return super().__repr__()
@@ -188,13 +191,15 @@ def assert_structural_equal(lhs, rhs, map_free_vars=False):
first_mismatch = tvm_ffi.get_first_structural_mismatch(lhs, rhs, map_free_vars)
if first_mismatch is not None:
from tvm.runtime.script_printer import ( # pylint: disable=import-outside-toplevel
- PrinterConfig,
- _script,
+ Scriptable,
)
lhs_path, rhs_path = first_mismatch
- lhs_script = _script(lhs, PrinterConfig(syntax_sugar=False, path_to_underline=[lhs_path]))
- rhs_script = _script(rhs, PrinterConfig(syntax_sugar=False, path_to_underline=[rhs_path]))
+ # Diagnostics use the same display policy as Object.script(), including
+ # dialect selection and commented imports. The unbound method also
+ # accepts IR objects that do not inherit the convenience mixin.
+ lhs_script = Scriptable.script(lhs, syntax_sugar=False, path_to_underline=[lhs_path])
+ rhs_script = Scriptable.script(rhs, syntax_sugar=False, path_to_underline=[rhs_path])
raise ValueError(
f"StructuralEqual check failed, caused by lhs at {lhs_path}:\n"
f"{lhs_script}\n"
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 6129ea8c50..9c261c38c3 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -596,8 +596,8 @@ class Call(_CallableExprWithOp):
def reinfer_type(call: Call) -> "tvm.ir.Type":
"""Derive a Call's result type from its current inputs without changing the Call.
- The operator must have a context-free inference rule. Invalid inputs and
- missing rules raise errors.
+ The operator must register a fixed return type or a context-free inference
+ rule. This does not invoke the operator's validator.
"""
return _ffi_api.reinfer_type(call)
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index f1acdeb169..0db6a57659 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -549,7 +549,6 @@ class BasePyModule:
print_line_numbers: bool = False,
num_context_lines: int = -1,
syntax_sugar: bool = True,
- show_object_address: bool = False,
show_all_ty: bool = True,
extra_config: dict | None = None,
) -> str:
@@ -571,7 +570,6 @@ class BasePyModule:
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
- show_object_address=show_object_address,
show_all_ty=show_all_ty,
extra_config=extra_config,
)
diff --git a/python/tvm/relax/script/ir_builder/__init__.py b/python/tvm/relax/script/ir_builder/__init__.py
index 91c87aca82..6c9a953db3 100644
--- a/python/tvm/relax/script/ir_builder/__init__.py
+++ b/python/tvm/relax/script/ir_builder/__init__.py
@@ -33,3 +33,48 @@ from .parser_protocol import (
)
__all__ = [*_native.__all__, *_op.__all__, *_protocol.__all__, "constexpr", "dist"]
+
+
+def _register_printer_names():
+ import types # pylint: disable=import-outside-toplevel
+
+ from tvm.ir import Op, register_op_attr # pylint: disable=import-outside-toplevel
+
+ active = set()
+
+ def register(value, path):
+ # These constructors use specialized out_ty/argument conventions.
+ if not callable(value) or path.rsplit(".", 1)[-1].startswith("call_tir"):
+ return
+ try:
+ op = Op.get(path)
+ except AttributeError:
+ return
+ if op.has_attr("TScriptPrinterName") and op.get_attr("TScriptPrinterName") == path:
+ return
+ register_op_attr(path, "TScriptPrinterName", path)
+
+ def visit(value, path):
+ if id(value) in active:
+ return
+ active.add(id(value))
+ register(value, path)
+ for name in getattr(value, "__all__", dir(value)):
+ if name.startswith("_"):
+ continue
+ try:
+ member = getattr(value, name)
+ except AttributeError:
+ continue
+ member_path = f"{path}.{name}"
+ if isinstance(member, types.ModuleType) and member.__name__.startswith("tvm.relax.op."):
+ visit(member, member_path)
+ elif callable(member):
+ register(member, member_path)
+ active.remove(id(value))
+
+ visit(_op, "relax")
+ visit(dist, "relax.dist")
+
+
+_register_printer_names()
diff --git a/python/tvm/relax/script/ir_builder/ir.py b/python/tvm/relax/script/ir_builder/ir.py
index b560176f34..195d0b79c3 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -561,6 +561,7 @@ def match_cast(value, ty, *, span=None):
Any = Object
+PackedFunc = _relax.PackedFuncType
Range = _ir.Range
__all__ = [
@@ -569,6 +570,7 @@ __all__ = [
"DTensor",
"ExternFunc",
"Object",
+ "PackedFunc",
"Range",
"Shape",
"ShapeExpr",
diff --git a/python/tvm/relax/type.py b/python/tvm/relax/type.py
index 4ed8e7a2ff..cff57560af 100644
--- a/python/tvm/relax/type.py
+++ b/python/tvm/relax/type.py
@@ -187,5 +187,5 @@ class FuncType(Type):
"""
if isinstance(derive_func, str):
- derive_func = EnvFunc.get("tvm.relax.type.infer_view_ty")
+ derive_func = EnvFunc.get(derive_func)
return _ffi_api.FuncTypeOpaqueFunc(ret, derive_func, purity, span) # type: ignore
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index bd6a9290bc..45874a8e49 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -30,21 +30,23 @@ from . import _ffi_node_api
@register_object("script.PrinterConfig")
class PrinterConfig(Object):
- """Configuration of TVMScript printer"""
+ """Configuration of TVMScript printer.
+
+ The ``int_dtype`` and ``float_dtype`` arguments populate the string-valued
+ ``ir.int_dtype`` and ``script.float_dtype`` entries in ``extra_config``.
+ Explicit entries in ``extra_config`` override those arguments.
+ """
binding_names: Sequence[str]
show_meta: bool
ir_prefix: str
module_alias: str
buffer_dtype: str
- int_dtype: str
- float_dtype: str
verbose_expr: bool
indent_spaces: int
print_line_numbers: bool
num_context_lines: int
syntax_sugar: bool
- show_object_address: bool
extra_config: dict
path_to_underline: list[AccessPath] | None
path_to_annotate: dict[AccessPath, str] | None
@@ -66,7 +68,6 @@ class PrinterConfig(Object):
print_line_numbers: bool = False,
num_context_lines: int | None = None,
syntax_sugar: bool = True,
- show_object_address: bool = False,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -81,14 +82,13 @@ class PrinterConfig(Object):
"ir_prefix": ir_prefix,
"module_alias": module_alias,
"buffer_dtype": buffer_dtype,
- "int_dtype": int_dtype,
- "float_dtype": float_dtype,
+ "ir.int_dtype": int_dtype,
+ "script.float_dtype": float_dtype,
"verbose_expr": verbose_expr,
"indent_spaces": indent_spaces,
"print_line_numbers": print_line_numbers,
"num_context_lines": num_context_lines,
"syntax_sugar": syntax_sugar,
- "show_object_address": show_object_address,
"path_to_underline": path_to_underline,
"path_to_annotate": path_to_annotate,
"obj_to_underline": obj_to_underline,
@@ -133,7 +133,6 @@ class Scriptable:
print_line_numbers: bool = False,
num_context_lines: int = -1,
syntax_sugar: bool = True,
- show_object_address: bool = False,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -152,8 +151,7 @@ class Scriptable:
ir_prefix : str = "I"
The prefix of AST nodes from tvm.ir
module_alias : str = "cls"
- The alias of the current module at cross-function call,
- Directly use module name if it's empty.
+ Retained for compatibility. Cross-function calls use the module name directly.
int_dtype : str = "int32"
The default data type of integer
float_dtype : str = "void"
@@ -168,8 +166,6 @@ class Scriptable:
The number of lines of context to print before and after the line to underline.
syntax_sugar: bool = True
Whether to output with syntax sugar, set false for complete printing.
- show_object_address: bool = False
- Whether to include the object's address as part of the TVMScript name
show_all_ty: bool = True
If True (default), annotate all variable bindings with the struct
info of that variable. If False, only add annotations where
@@ -200,6 +196,9 @@ class Scriptable:
merged_extra: dict = {}
if extra_config is not None:
merged_extra.update(extra_config)
+ # Keep the historical interactive display header. Direct printer calls
+ # retain executable imports for standalone parser round-trips.
+ merged_extra.setdefault("script.comment_imports", True)
if "script.use_pep695" not in merged_extra:
merged_extra["script.use_pep695"] = merged_extra.get(
"relax.use_pep695", sys.version_info >= (3, 12)
@@ -245,7 +244,6 @@ class Scriptable:
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
- show_object_address=show_object_address,
show_all_ty=show_all_ty,
extra_config=merged_extra if merged_extra else None,
path_to_underline=path_to_underline,
@@ -269,7 +267,6 @@ class Scriptable:
print_line_numbers: bool = False,
num_context_lines: int = -1,
syntax_sugar: bool = True,
- show_object_address: bool = False,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
path_to_annotate: dict[AccessPath, str] | None = None,
@@ -295,7 +292,6 @@ class Scriptable:
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
- show_object_address=show_object_address,
extra_config=merged_extra,
path_to_underline=path_to_underline,
path_to_annotate=path_to_annotate,
@@ -320,7 +316,6 @@ class Scriptable:
print_line_numbers: bool = False,
num_context_lines: int = -1,
syntax_sugar: bool = True,
- show_object_address: bool = False,
show_all_ty: bool = True,
extra_config: dict | None = None,
path_to_underline: list[AccessPath] | None = None,
@@ -362,8 +357,7 @@ class Scriptable:
ir_prefix : str = "I"
The prefix of AST nodes from tvm.ir
module_alias : str = "cls"
- The alias of the current module at cross-function call,
- Directly use module name if it's empty.
+ Retained for compatibility. Cross-function calls use the module name directly.
int_dtype : str = "int32"
The default data type of integer
float_dtype : str = "void"
@@ -378,8 +372,6 @@ class Scriptable:
The number of lines of context to print before and after the line to underline.
syntax_sugar: bool = True
Whether to output with syntax sugar, set false for complete printing.
- show_object_address: bool = False
- Whether to include the object's address as part of the TVMScript name
show_all_ty: bool = True
If True (default), annotate all variable bindings with the struct
info of that variable. If False, only add annotations where
@@ -415,7 +407,6 @@ class Scriptable:
print_line_numbers=print_line_numbers,
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
- show_object_address=show_object_address,
show_all_ty=show_all_ty,
extra_config=extra_config,
path_to_underline=path_to_underline,
diff --git a/python/tvm/s_tir/script/ir_builder/ir.py b/python/tvm/s_tir/script/ir_builder/ir.py
index e7159b8119..faba9bffee 100644
--- a/python/tvm/s_tir/script/ir_builder/ir.py
+++ b/python/tvm/s_tir/script/ir_builder/ir.py
@@ -24,6 +24,7 @@ from typing import Any
from tvm import ir
from tvm import tirx as tir
from tvm.ir import TensorLoad, TensorRegion, is_prim_expr
+from tvm.script.ir_builder.base import MISSING
from tvm.script.parser.protocol_registry import register_mutable_decl as _register_mutable_decl
from tvm.tirx import Buffer, Expr, IntImm, Var
from tvm.tirx.layout import Layout
@@ -43,7 +44,7 @@ def match_buffer(
scope: str = "global",
align: int = -1,
offset_factor: int = 0,
- layout: str | Layout | None = "default",
+ layout: str | Layout | None = MISSING,
allocated_addr: Expr | int | tuple[Expr | int, ...] | None = None,
) -> Buffer:
"""Bind a buffer subregion inside an S-TIR block.
@@ -191,7 +192,7 @@ def sblock_alloc_buffer(
scope: str = "global",
align: int = -1,
offset_factor: int = 0,
- layout: str | Layout | None = "default",
+ layout: str | Layout | None = MISSING,
allocated_addr: int | tuple[int, ...] | None = None,
) -> Buffer:
"""SBlock-level buffer allocation function.
@@ -284,7 +285,7 @@ class axis: # pylint: disable=invalid-name
@staticmethod
def spatial(
dom: ir.Range | list[Expr] | tuple[Expr],
- binding: Expr,
+ binding: Expr | None = None,
dtype: str = "int32",
) -> Var:
"""The spatial block axis defining function.
@@ -294,8 +295,9 @@ class axis: # pylint: disable=invalid-name
dom : Union[Range, List[Expr], Tuple[Expr]]
The domain of the iteration variable.
- binding : Expr
- The binding value of the iteration variable.
+ binding : Expr, optional
+ The binding value of the iteration variable. Omit only inside a
+ block with no_realize=True.
dtype : str
The data type of the iteration variable.
@@ -312,7 +314,7 @@ class axis: # pylint: disable=invalid-name
@staticmethod
def reduce(
dom: ir.Range | list[Expr] | tuple[Expr],
- binding: Expr,
+ binding: Expr | None = None,
dtype: str = "int32",
) -> Var:
"""The reduced block axis defining function.
@@ -322,8 +324,9 @@ class axis: # pylint: disable=invalid-name
dom : Union[Range, List[Expr], Tuple[Expr]]
The domain of the iteration variable.
- binding : Expr
- The binding value of the iteration variable.
+ binding : Expr, optional
+ The binding value of the iteration variable. Omit only inside a
+ block with no_realize=True.
dtype : str
The data type of the iteration variable.
@@ -340,7 +343,7 @@ class axis: # pylint: disable=invalid-name
@staticmethod
def scan(
dom: ir.Range | list[Expr] | tuple[Expr],
- binding: Expr,
+ binding: Expr | None = None,
dtype: str = "int32",
) -> Var:
"""The scanning block axis defining function.
@@ -350,8 +353,9 @@ class axis: # pylint: disable=invalid-name
dom : Union[Range, List[Expr], Tuple[Expr]]
The domain of the iteration variable.
- binding : Expr
- The binding value of the iteration variable.
+ binding : Expr, optional
+ The binding value of the iteration variable. Omit only inside a
+ block with no_realize=True.
dtype : str
The data type of the iteration variable.
@@ -368,7 +372,7 @@ class axis: # pylint: disable=invalid-name
@staticmethod
def opaque(
dom: ir.Range | list[Expr] | tuple[Expr],
- binding: Expr,
+ binding: Expr | None = None,
dtype: str = "int32",
) -> Var:
"""The opaque block axis defining function.
@@ -378,8 +382,9 @@ class axis: # pylint: disable=invalid-name
dom : Union[Range, List[Expr], Tuple[Expr]]
The domain of the iteration variable.
- binding : Expr
- The binding value of the iteration variable.
+ binding : Expr, optional
+ The binding value of the iteration variable. Omit only inside a
+ block with no_realize=True.
dtype : str
The data type of the iteration variable.
diff --git a/python/tvm/s_tir/script/ir_builder/parser_protocol.py b/python/tvm/s_tir/script/ir_builder/parser_protocol.py
index f86bf4009b..243d7f52a4 100644
--- a/python/tvm/s_tir/script/ir_builder/parser_protocol.py
+++ b/python/tvm/s_tir/script/ir_builder/parser_protocol.py
@@ -20,7 +20,6 @@
from tvm import ir as _ir
from tvm import tirx as _tir
from tvm.script.ir_builder import base as _base
-from tvm.tirx.script.ir_builder.ir import buffer
from tvm.tirx.script.ir_builder.parser_protocol import arg_ as _shared_arg
from tvm.tirx.script.ir_builder.parser_protocol import bind_ as _shared_bind
@@ -50,21 +49,11 @@ def function_(*, private=False, persistent=False, decl=False, span=None):
def arg_(name, annotation, *, span=None):
- """Normalize eagerly constructed parameter buffers for S-TIR."""
- if _tir.is_buffer_var(annotation) and annotation.ty.layout is not None:
- ty = annotation.ty
- annotation = buffer(
- ty.shape,
- ty.dtype,
- strides=ty.strides,
- elem_offset=ty.elem_offset,
- scope=ty.storage_scope,
- align=ty.data_alignment,
- offset_factor=ty.offset_factor,
- layout=None,
- allocated_addr=list(ty.allocated_addr),
- buffer_name=name,
- )
+ """Preserve parameter annotations, including explicitly constructed layouts.
+
+ Buffer annotations are evaluated inside the function frame, which already
+ supplies the S-TIR default layout when the constructor omits it.
+ """
return _shared_arg(name, annotation, span=span)
diff --git a/python/tvm/script/ir_builder/__init__.py b/python/tvm/script/ir_builder/__init__.py
index f1faadf35e..90743bde5f 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -16,7 +16,19 @@
# under the License.
"""Shared TVMScript construction APIs and lazy language variant builders."""
-from tvm.ir import DataTypeImm, GenericConst, Range, StringImm, StringType
+from tvm.ir import Call as _IRCall
+from tvm.ir import (
+ DataTypeImm,
+ FuncType,
+ GenericConst,
+ PrimType,
+ Range,
+ StringImm,
+ StringType,
+ Type,
+ make_node,
+ reinfer_type,
+)
from .base import (
MISSING,
@@ -46,17 +58,38 @@ from .parser_protocol import (
resolve_global_info_,
)
+
+def Call(op, args, attrs=None, ty_args=None, ty=None): # pylint: disable=invalid-name
+ """Construct a shared IR Call, inferring its result type when ``ty`` is omitted.
+
+ ``ty_args`` are independent explicit inputs to the inference hook.
+ An explicit ``ty`` preserves the supplied Call fields without invoking its
+ operator validator, including provisional calls with missing input types.
+ Omitting ``ty`` requests inference followed by ordinary Call validation.
+ """
+
+ if ty is None:
+ provisional = _IRCall.unchecked(op, args, attrs=attrs, ty_args=ty_args)
+ ty = reinfer_type(provisional)
+ return _IRCall(op, args, attrs=attrs, ty_args=ty_args, ret_ty=ty)
+ return _IRCall.unchecked(op, args, attrs=attrs, ty_args=ty_args, ret_ty=ty)
+
+
# Keep source namespaces independent of imported helper modules and lazy builders.
__all__ = [
"MISSING",
"AlreadyEmitted",
+ "Call",
"DataTypeImm",
+ "FuncType",
"GenericConst",
"IRBuilder",
"IRModuleFrame",
+ "PrimType",
"Range",
"StringImm",
"StringType",
+ "Type",
"at_",
"check_well_formed_",
"constexpr",
@@ -64,6 +97,7 @@ __all__ = [
"dtype",
"dynamic",
"ir_module",
+ "make_node",
"meta_var",
"module_attrs",
"module_get_attr",
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index f54910acc3..d583d8930f 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -508,7 +508,9 @@ def call_llvm_intrin(dtype, name, *args, span=None):
return call_intrin(
dtype,
Op.get("tirx.call_llvm_intrin"),
- tvm.tirx.const(llvm_id, "uint32"),
+ name
+ if isinstance(name, IntImm)
+ else tvm.tirx.const(llvm_id, "int32" if isinstance(name, str) else "uint32"),
*args,
span=span,
)
@@ -550,7 +552,9 @@ def call_llvm_pure_intrin(dtype, name, *args, span=None):
return call_intrin(
dtype,
Op.get("tirx.call_llvm_pure_intrin"),
- tvm.tirx.const(llvm_id, "uint32"),
+ name
+ if isinstance(name, IntImm)
+ else tvm.tirx.const(llvm_id, "int32" if isinstance(name, str) else "uint32"),
*args,
span=span,
)
diff --git a/python/tvm/tirx/script/ir_builder/ir.py b/python/tvm/tirx/script/ir_builder/ir.py
index f39a099787..3b7725991b 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -37,7 +37,7 @@ from typing import Literal
from tvm import DataType, ir
from tvm import tirx as tir
from tvm.ir import TensorLoad, Type, is_prim_expr
-from tvm.script.ir_builder.base import IRBuilder
+from tvm.script.ir_builder.base import MISSING, IRBuilder
from tvm.script.ir_builder.ir import meta_var
from tvm.script.parser.protocol_registry import (
register_mutable_decl as _register_mutable_decl,
@@ -74,12 +74,7 @@ from . import _ffi_api, frame
def _get_layout(layout: str | Layout | None, shape: list[Expr], scope: str) -> Layout | None:
- if layout is None:
- return None
- if isinstance(layout, Layout):
- return layout
- assert isinstance(layout, str)
- if layout == "default":
+ if layout is MISSING:
if IRBuilder.is_in_scope():
for function_frame in reversed(list(IRBuilder.current().frames)):
if isinstance(function_frame, frame.PrimFuncFrame):
@@ -87,6 +82,15 @@ def _get_layout(layout: str | Layout | None, shape: list[Expr], scope: str) -> L
if scope in ["trn.sbuf", "trn.psum"]:
return None
return TileLayout(S[tuple(shape)])
+ if layout is None:
+ return None
+ if isinstance(layout, Layout):
+ return layout
+ assert isinstance(layout, str)
+ if layout == "default":
+ # An explicit default denotes the same TileLayout in every dialect.
+ # Only an omitted argument consults the enclosing function's policy.
+ return TileLayout(S[tuple(shape)])
shape = tuple(shape)
if scope == "trn.sbuf":
layout = TileLayout.trainium(layout, shape)
@@ -206,6 +210,90 @@ def _record_meta_resource(value: Any, skip_frames: int = 2) -> None:
@_register_mutable_decl("tirx.Buffer", syntax="parameter")
@_annotation_constructor
+def _buffer_type(
+ shape: list[Expr] | tuple[Expr] | Expr | Integral,
+ dtype: str = "float32",
+ data: Var = None,
+ strides: list[Expr] | None = None,
+ elem_offset: Expr = None,
+ byte_offset: Expr = None,
+ scope: str = "global",
+ align: int = 0,
+ offset_factor: int = 0,
+ layout: str | Layout | None = MISSING,
+ allocated_addr: int | tuple[int, ...] | None = None,
+ *,
+ span=None,
+) -> tir.BufferType:
+ """Construct a buffer type for annotations and explicit type-valued fields.
+
+ Parameters
+ ----------
+ shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
+ The shape of the buffer prior to flattening.
+
+ dtype : str
+ The data type in the content of the buffer.
+
+ data : Var
+ An optional pointer whose storage scope determines the buffer type scope.
+
+ strides : List[Expr]
+ The strides of each dimension.
+
+ elem_offset : Expr
+ The offset in terms of number of dtype elements (including lanes).
+
+ byte_offset : Expr, optional
+ The offset in bytes, as an alternative to elem_offset.
+
+ scope : str
+ The optional storage scope of buffer data pointer.
+
+ align : int
+ The alignment requirement of data pointer in bytes.
+
+ offset_factor : int
+ The factor of elem_offset field.
+
+ layout : str or Layout, optional
+ The buffer layout. "default" constructs the shape's default TileLayout;
+ omission uses the enclosing dialect and scope default. None omits a layout.
+
+ allocated_addr : int or tuple of int, optional
+ Addresses assigned to the buffer allocation.
+
+ Returns
+ -------
+ res : BufferType
+ The buffer type. Function annotations introduce variables of this type;
+ allocation and declaration operations construct buffer variables.
+ """
+ shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
+ shape = tuple(shape)
+ if strides is None:
+ strides = []
+ if allocated_addr is None:
+ allocated_addr = []
+ if not isinstance(allocated_addr, list | tuple):
+ allocated_addr = [allocated_addr]
+ result = _ffi_api.BufferType( # type: ignore[attr-defined] # pylint: disable=no-member
+ shape,
+ dtype,
+ data,
+ strides,
+ _get_elem_offset(elem_offset, byte_offset, dtype),
+ scope,
+ align,
+ offset_factor,
+ _get_layout(layout, shape, scope),
+ allocated_addr,
+ )
+ return _at(span, result)
+
+
+@_register_mutable_decl("tirx.buffer", syntax="parameter")
+@_annotation_constructor
def buffer(
shape: list[Expr] | tuple[Expr] | Expr | Integral,
dtype: str = "float32",
@@ -216,7 +304,7 @@ def buffer(
scope: str = "global",
align: int = 0,
offset_factor: int = 0,
- layout: str | Layout | None = "default",
+ layout: str | Layout | None = MISSING,
allocated_addr: int | tuple[int, ...] | None = None,
buffer_name: str = "",
*,
@@ -254,7 +342,8 @@ def buffer(
The factor of elem_offset field.
layout : str or Layout, optional
- The buffer layout; "default" selects the layout for the buffer scope.
+ The buffer layout. "default" constructs the shape's default TileLayout;
+ omission uses the enclosing dialect and scope default. None omits a layout.
allocated_addr : int or tuple of int, optional
Addresses assigned to the buffer allocation.
@@ -470,7 +559,7 @@ def alloc_buffer(
scope: str = "global",
align: int = -1,
offset_factor: int = 0,
- layout: str | Layout | None = "default",
+ layout: str | Layout | None = MISSING,
allocated_addr: int | tuple[int, ...] | None = None,
annotations: dict[str, Any] | None = None,
) -> Buffer:
@@ -678,7 +767,7 @@ def decl_buffer(
scope="global",
align=0,
offset_factor=0,
- layout="default",
+ layout=MISSING,
allocated_addr=None,
) -> Buffer:
"""Create a buffer declaration node.
@@ -1614,8 +1703,7 @@ def Ptr(dtype, storage_scope="global", *, span=None):
return _at(span, ptr(dtype, storage_scope))
-Buffer = buffer
-_register_mutable_decl("tirx.buffer", syntax="parameter")(buffer)
+Buffer = _buffer_type
__all__ = [
"Buffer",
diff --git a/python/tvm/tirx/script/ir_builder/op.py b/python/tvm/tirx/script/ir_builder/op.py
index a8e794dfef..f2173f361f 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -330,13 +330,13 @@ def register_script_namespace(name: str, namespace: object, override: bool = Fal
if isinstance(module_all, list) and name not in module_all:
module_all.append(name)
- _register_script_namespace_printer_names(namespace, name, override)
+ _register_script_namespace_printer_names(namespace, f"tirx.{name}", override)
return namespace
def _register_tir_namespace_printer_names():
try:
- _register_script_namespace_printer_names(webgpu, "webgpu")
+ _register_script_namespace_printer_names(webgpu, "tirx.webgpu")
except Exception:
# Best-effort registration; avoid import-time hard failure
pass
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 268dabd861..1d0a9cc3ea 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -635,7 +635,9 @@ def while_(condition: Any, *, span: _Span = None) -> frame.WhileFrame:
def range_(*args: Any, annotations: dict[str, Any] | None = None) -> frame.ForFrame:
"""Implements :func:`tvm.script.ir_builder.parser_protocol.range_`."""
if len(args) == 1:
- args = (0, args[0], None)
+ # serial constructs the omitted zero in the stop expression's dtype.
+ # A Python zero would promote narrow typed bounds to int32.
+ return serial(args[0], annotations=annotations)
elif len(args) == 2:
args = (*args, None)
elif len(args) != 3:
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 8b365cdfa5..471cba6edc 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1295,14 +1295,14 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
std::string global_buffer =
this->PrintExpr(addr_buffer->source.as_or_throw<tvm::tirx::BufferVar>().data());
std::string local_addr = this->PrintExpr(op->args[3]);
+ std::string destination = op->args[0].as<TensorLoadNode>() ? reg : reg + "[" + local_addr + "]";
this->stream << "asm volatile (\n";
this->stream << "\"{.reg .pred p;\\n\"\n";
this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n";
this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n";
this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n";
// stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ;
- stream << ": \"=f\"(" << reg << "[" << local_addr << "]"
- << ")\n";
+ stream << ": \"=f\"(" << destination << ")\n";
stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)"
<< guard << ")\n";
stream << ");\n";
diff --git a/src/backend/cuda/codegen/intrin_rule_cuda.cc b/src/backend/cuda/codegen/intrin_rule_cuda.cc
index a8dad1d14d..bbbbcca4f7 100644
--- a/src/backend/cuda/codegen/intrin_rule_cuda.cc
+++ b/src/backend/cuda/codegen/intrin_rule_cuda.cc
@@ -34,6 +34,13 @@ namespace intrin {
// Add float suffix to the intrinsics, CUDA fast math.
using tirx::FLowerIntrinsic;
+template <size_t N>
+static Type InferTypeReturnArgType(const CallNode* call) {
+ TVM_FFI_CHECK_GT(call->args.size(), N, ValueError)
+ << "Return type inference requires argument " << N;
+ return call->args[N]->ty;
+}
+
struct CUDAMath {
std::string operator()(const PrimType& ty, std::string name) const {
if (ty.MatchesCode(DLDataTypeCode::kDLFloat)) {
@@ -278,9 +285,11 @@ void RegisterCudaIntrinRules() {
sig::arg("var", "The variable to sync."),
sig::arg("lane", "The source thread id."),
sig::arg("width", "The warp thread width, must be a power of 2."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_sync"))
+ .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.__shfl_sync"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "__shfl_sync")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<bool>("cuda.need_warp_shuffle", true);
@@ -290,9 +299,11 @@ void RegisterCudaIntrinRules() {
sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be added."),
sig::arg("width", "The warp thread width, must be a power of 2."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_up_sync"))
+ .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.__shfl_up_sync"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "__shfl_up_sync")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<bool>("cuda.need_warp_shuffle", true);
@@ -302,10 +313,11 @@ void RegisterCudaIntrinRules() {
sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be subtracted."),
sig::arg("width", "The warp thread width, must be a power of 2."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("cuda.__shfl_down_sync"))
+ ffi::String("tirx.cuda.__shfl_down_sync"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "__shfl_down_sync")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<bool>("cuda.need_warp_shuffle", true);
@@ -315,17 +327,21 @@ void RegisterCudaIntrinRules() {
sig::arg("var", "The variable to sync."),
sig::arg("lane_mask", "The lane mask."),
sig::arg("width", "The warp thread width, must be a power of 2."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__shfl_xor_sync"))
+ .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.__shfl_xor_sync"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "__shfl_xor_sync")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<bool>("cuda.need_warp_shuffle", true);
OpDef("tirx.cuda.__activemask")
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::UInt(32))
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.__activemask"))
+ .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.__activemask"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "__activemask")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
.set_attr<bool>("cuda.need_warp_shuffle", true);
diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc
index aaec8a8003..0186402fa4 100644
--- a/src/backend/cuda/op/iket.cc
+++ b/src/backend/cuda/op/iket.cc
@@ -30,49 +30,47 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.cuda.iket_mark")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "mark"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.iket.mark"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_start")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::UInt(32))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "range_start"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.iket.range_start"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_end")
.signature(sig::arg("token", "The token."), sig::var_args("args"))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "range_end"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.iket.range_end"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_push")
.signature(sig::arg("name", "The name."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "range_push"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.iket.range_push"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_range_pop")
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "range_pop"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.iket.range_pop"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_sentinel_token")
.signature(sig::arg("name", "The name."))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "sentinel_token"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.iket.sentinel_token"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.cuda.iket_official_event")
@@ -80,8 +78,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("source_code", "The source code."), sig::var_args("args"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("cuda"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.iket."
- "official_event"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.cuda.iket.official_event"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc
index 2011ab9481..660d113f94 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -39,8 +39,22 @@ namespace builtin {
namespace {
void RegisterDeviceIntrinsicAliases();
+
+Type InferTypeCudaLdg(const CallNode* call) {
+ TVM_FFI_CHECK_GE(call->args.size(), 2U, ValueError)
+ << "CUDA load type inference requires a dtype operand";
+ return PrimType(ffi::StringToDLDataType(call->args[1].as_or_throw<StringImm>()->value));
+}
+
+template <size_t N>
+static Type InferTypeReturnArgType(const CallNode* call) {
+ TVM_FFI_CHECK_GT(call->args.size(), N, ValueError)
+ << "Return type inference requires argument " << N;
+ return call->args[N]->ty;
}
+} // namespace
+
void RegisterCudaTargetBuiltins() {
static bool registered = false;
if (registered) return;
@@ -51,7 +65,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
sig::arg("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."),
sig::arg("stride", "The stride."), sig::arg("layout", "The layout."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_load_matrix_sync"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.tvm_load_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
@@ -62,7 +76,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The A fragment index."),
sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The B fragment index."),
sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The C fragment index."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_mma_sync"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.tvm_mma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -72,7 +86,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("fragment_a", "The A fragment."), sig::arg("index_a", "The A fragment index."),
sig::arg("fragment_b", "The B fragment."), sig::arg("index_b", "The B fragment index."),
sig::arg("fragment_c", "The C fragment."), sig::arg("index_c", "The C fragment index."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_bmma_sync"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.tvm_bmma_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -80,7 +94,7 @@ void RegisterCudaTargetBuiltins() {
.signature(sig::arg("fragment", "The matrix fragment."), sig::arg("m", "The M dimension."),
sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
sig::arg("index", "The index."), sig::arg("value", "The value to use."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_fill_fragment"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.tvm_fill_fragment"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -89,7 +103,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("n", "The N dimension."), sig::arg("k", "The K dimension."),
sig::arg("index", "The index."), sig::arg("buffer_ptr", "The buffer pointer."),
sig::arg("stride", "The stride."), sig::arg("layout", "The layout."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tvm_store_matrix_sync"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.tvm_store_matrix_sync"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -104,14 +118,14 @@ void RegisterCudaTargetBuiltins() {
sig::arg("src_ptr", "The source pointer."),
sig::arg("src_offset", "The source offset."),
sig::arg("dst_stride", "The destination stride."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_store_legacy"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.mma_store_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.mma_fill_legacy")
.signature(sig::arg("local_size", "The local allocation size."),
sig::arg("local_ptr", "The local pointer."), sig::arg("offset", "The offset."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_fill_legacy"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.mma_fill_legacy"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -119,7 +133,7 @@ void RegisterCudaTargetBuiltins() {
.signature(sig::arg("reg", "The register."), sig::arg("guard", "The guard predicate."),
sig::arg("addr", "The address."), sig::arg("local_addr", "The local address."))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("s_tir.ldg32"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.s_tir.ldg32"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("s_tir"));
@@ -135,7 +149,7 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("s_tir"))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("s_tir.cp_async_raw"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.s_tir.cp_async_raw"))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst));
@@ -145,7 +159,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("src_ptr", "The source pointer."),
sig::arg("src_offset", "The source offset."),
sig::arg("dst_stride", "The destination stride."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_store"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.mma_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
@@ -154,7 +168,7 @@ void RegisterCudaTargetBuiltins() {
OpDef("tirx.mma_fill")
.signature(sig::arg("local_size", "The local allocation size."),
sig::arg("local_ptr", "The local pointer."), sig::arg("offset", "The offset."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.mma_fill"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.mma_fill"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
@@ -166,7 +180,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("profiler_write_offset", "The profiler write offset."),
sig::arg("num_groups", "The number of groups."),
sig::arg("group_id", "The group identifier."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_init"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.timer_init"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -177,7 +191,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("profiler_write_offset", "The profiler write offset."),
sig::arg("profiler_write_stride", "The profiler write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_start"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.timer_start"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -188,7 +202,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("profiler_write_offset", "The profiler write offset."),
sig::arg("profiler_write_stride", "The profiler write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_end"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.timer_end"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -198,7 +212,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg("profiler_write_offset", "The profiler write offset."),
sig::arg("profiler_write_stride", "The profiler write stride."),
sig::arg("leader_cond", "The leader condition."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("cuda.timer_finalize"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.cuda.timer_finalize"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -225,14 +239,15 @@ TVM_FFI_NO_INLINE DeviceIntrinsicNames MakeDeviceIntrinsicNames(const char* op_n
std::string canonical = "tirx." + namespace_name + "." + suffix;
// Match the nested construction namespaces at the canonical registration site.
if (namespace_name == "cuda" &&
- (suffix.rfind("tcgen05_", 0) == 0 || suffix.rfind("wgmma_", 0) == 0)) {
+ (suffix.rfind("tcgen05_", 0) == 0 || suffix.rfind("wgmma_", 0) == 0 ||
+ suffix.rfind("iket_", 0) == 0)) {
suffix[suffix.find('_')] = '.';
} else if (namespace_name == "nvshmem" &&
((suffix.size() >= 6 && suffix.compare(suffix.size() - 6, 6, "_block") == 0) ||
(suffix.size() >= 5 && suffix.compare(suffix.size() - 5, 5, "_warp") == 0))) {
suffix[suffix.rfind('_')] = '.';
}
- return {std::move(canonical), namespace_name + "." + suffix};
+ return {std::move(canonical), "tirx." + namespace_name + "." + suffix};
}
TVM_FFI_NO_INLINE void RegisterDeviceIntrinsicAttrs(OpDef& def, const char* op_namespace,
@@ -408,6 +423,51 @@ void RegisterDeviceIntrinsicAliases() {
sig::arg("a_ptr"), sig::arg("a_offset"), sig::arg("b_ptr"),
sig::arg("b_offset"), sig::arg("acc_ptr"), sig::arg("c_offset"),
sig::arg("saturate"), sig::var_args("args"));
+
+ for (const char* name : {
+ "tirx.cuda.tcgen05_encode_matrix_descriptor",
+ "tirx.cuda.wgmma_encode_matrix_descriptor",
+ "tirx.cuda.tcgen05_encode_instr_descriptor",
+ "tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled",
+ "tirx.cuda.wgmma_noop_barrier",
+ "tirx.cuda.cluster_sync",
+ "tirx.cuda.cta_sync",
+ "tirx.cuda.grid_sync",
+ "tirx.cuda.warp_sync",
+ "tirx.cuda.thread_fence",
+ "tirx.cuda.mbarrier_wait",
+ "tirx.cuda.printf",
+ "tirx.cuda.nano_sleep",
+ "tirx.nvshmem.fence",
+ "tirx.nvshmem.quiet",
+ "tirx.nvshmem.barrier_all",
+ "tirx.nvshmem.getmem_nbi",
+ "tirx.nvshmem.getmem_nbi_warp",
+ "tirx.nvshmem.putmem_nbi",
+ "tirx.nvshmem.putmem_nbi_warp",
+ "tirx.nvshmem.putmem_nbi_block",
+ "tirx.nvshmem.signal_op",
+ "tirx.nvshmem.wait_until",
+ "tirx.nvshmem.putmem_signal_nbi",
+ "tirx.nvshmem.putmem_signal_nbi_warp",
+ "tirx.nvshmem.putmem_signal_nbi_block",
+ }) {
+ OpDef(name).set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
+ }
+ for (const char* name : {"tirx.cuda.syncthreads_and", "tirx.cuda.syncthreads_or"}) {
+ OpDef(name).set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(64));
+ }
+ for (const char* name : {"tirx.nvshmem.my_pe", "tirx.nvshmem.n_pes"}) {
+ OpDef(name).set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(32));
+ }
+ OpDef("tirx.cuda.elect_sync").set_attr<TFixedReturnType>("TFixedReturnType", PrimType::UInt(32));
+ OpDef("tirx.cuda.ldg")
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeCudaLdg>());
+ OpDef("tirx.s_tir.ldg32").set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Float(32));
+ OpDef("tirx.cuda.atomic_add")
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>());
+ OpDef("tirx.cuda.atomic_cas")
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<1>>());
}
} // namespace
diff --git a/src/backend/metal/codegen/intrin_rule_metal.cc b/src/backend/metal/codegen/intrin_rule_metal.cc
index 306e81a55f..2f3d559ff2 100644
--- a/src/backend/metal/codegen/intrin_rule_metal.cc
+++ b/src/backend/metal/codegen/intrin_rule_metal.cc
@@ -30,6 +30,13 @@ namespace codegen {
namespace intrin {
using tirx::FLowerIntrinsic;
+template <size_t N>
+static Type InferTypeReturnArgType(const CallNode* call) {
+ TVM_FFI_CHECK_GT(call->args.size(), N, ValueError)
+ << "Return type inference requires argument " << N;
+ return call->args[N]->ty;
+}
+
struct MetalWarpIntrinsic {
const Op operator()(PrimType t, const Op& orig_op) const {
if (orig_op.same_as(builtin::tvm_warp_shuffle())) {
@@ -149,29 +156,33 @@ void RegisterMetalIntrinRules() {
OpDef("tirx.metal.simd_shuffle")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("lane", "The source thread id."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
- .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simd_shuffle"))
+ .set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
+ ffi::String("tirx.metal.simd_shuffle"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "simd_shuffle")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.metal.simd_shuffle_up")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be added."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("metal.simd_shuffle_up"))
+ ffi::String("tirx.metal.simd_shuffle_up"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "simd_shuffle_up")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.metal.simd_shuffle_down")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be subtracted."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("metal"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("metal.simd_shuffle_down"))
+ ffi::String("tirx.metal.simd_shuffle_down"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "simd_shuffle_down")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
// clang-format on
diff --git a/src/backend/metal/op/target_builtin.cc b/src/backend/metal/op/target_builtin.cc
index 7f6c8bf60e..d9f3be63c0 100644
--- a/src/backend/metal/op/target_builtin.cc
+++ b/src/backend/metal/op/target_builtin.cc
@@ -41,7 +41,7 @@ void RegisterMetalTargetBuiltins() {
sig::arg("value", "The value to use."), sig::arg("col", "The column index."),
sig::arg("row", "The row index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
- ffi::String("metal.make_filled_simdgroup_matrix"))
+ ffi::String("tirx.metal.make_filled_simdgroup_matrix"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -50,7 +50,7 @@ void RegisterMetalTargetBuiltins() {
sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
sig::arg("col", "The column index."), sig::arg("row", "The row index."),
sig::arg("transpose_matrix", "Whether to transpose the matrix."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simdgroup_load"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.metal.simdgroup_load"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -59,7 +59,7 @@ void RegisterMetalTargetBuiltins() {
sig::arg("ptr", "The pointer."), sig::arg("stride", "The stride."),
sig::arg("col", "The column index."), sig::arg("row", "The row index."),
sig::arg("transpose_matrix", "Whether to transpose the matrix."))
- .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("metal.simdgroup_store"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String("tirx.metal.simdgroup_store"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
@@ -69,7 +69,7 @@ void RegisterMetalTargetBuiltins() {
sig::arg("b", "The B operand."), sig::arg("index_b", "The B fragment index."),
sig::arg("c", "The C operand."), sig::arg("index_c", "The C fragment index."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
- ffi::String("metal.simdgroup_multiply_accumulate"))
+ ffi::String("tirx.metal.simdgroup_multiply_accumulate"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
}
diff --git a/src/backend/trn/op/target_builtin.cc b/src/backend/trn/op/target_builtin.cc
index 2beb4dc45e..7d9e835e92 100644
--- a/src/backend/trn/op/target_builtin.cc
+++ b/src/backend/trn/op/target_builtin.cc
@@ -62,11 +62,12 @@ TVM_FFI_NO_INLINE NKIIntrinsicNames MakeNKIIntrinsicNames(const char* op_name) {
if (suffix.rfind(prefix, 0) == 0) {
suffix = suffix.substr(prefix.size());
}
- return {"tirx.nki." + suffix, "nki." + suffix};
+ return {"tirx.nki." + suffix, "tirx.nki." + suffix};
}
TVM_FFI_NO_INLINE void RegisterNKIIntrinsicAttrs(OpDef& def, const std::string& printer_name) {
- def.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
+ def.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("nki"))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName", ffi::String(printer_name));
diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
index 12d0ad8c0b..b55c517000 100644
--- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
+++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc
@@ -32,6 +32,13 @@ namespace intrin {
using tirx::FLowerIntrinsic;
+template <size_t N>
+static Type InferTypeReturnArgType(const CallNode* call) {
+ TVM_FFI_CHECK_GT(call->args.size(), N, ValueError)
+ << "Return type inference requires argument " << N;
+ return call->args[N]->ty;
+}
+
// warp-level primitives. Follows implementation in intrin_rule_metal.cc
struct WebGPUWarpIntrinsic {
const Op operator()(PrimType t, const Op& orig_op) const {
@@ -176,30 +183,33 @@ void RegisterWebGPUIntrinRules() {
OpDef("tirx.webgpu.subgroup_shuffle")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("lane", "The source thread id."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("webgpu.subgroup_shuffle"))
+ ffi::String("tirx.webgpu.subgroup_shuffle"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffle")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.webgpu.subgroup_shuffle_up")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be added."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("webgpu.subgroup_shuffle_up"))
+ ffi::String("tirx.webgpu.subgroup_shuffle_up"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffleUp")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
OpDef("tirx.webgpu.subgroup_shuffle_down")
.signature(sig::arg("var", "The variable to sync."),
sig::arg("delta", "The source lane id offset to be subtracted."))
+ .set_attr<FInferType>("FInferType", FInferType::FromNative<&InferTypeReturnArgType<0>>())
.set_attr<tirx::TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
.set_attr<tirx::TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace", ffi::String("webgpu"))
.set_attr<tirx::TScriptPrinterName>("TScriptPrinterName",
- ffi::String("webgpu.subgroup_shuffle_down"))
+ ffi::String("tirx.webgpu.subgroup_shuffle_down"))
.set_attr<TGlobalSymbol>("TGlobalSymbol", "subgroupShuffleDown")
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kOpaque));
// clang-format on
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 03710bb080..6b883640e7 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -1208,6 +1208,10 @@ Type Call::ReinferType(const CallNode* call) {
TVM_FFI_CHECK(call != nullptr, ValueError) << "Call::ReinferType expects a defined Call";
auto op = call->op.as<Op>();
TVM_FFI_CHECK(op.has_value(), ValueError) << "Call::ReinferType requires an Op callee";
+ if (Op::HasAttrMap("TFixedReturnType")) {
+ static auto fixed_return_type = Op::GetAttrMap<TFixedReturnType>("TFixedReturnType");
+ if (fixed_return_type.count(op.value())) return fixed_return_type[op.value()];
+ }
TVM_FFI_CHECK(Op::HasAttrMap("FInferType"), ValueError)
<< "No context-free FInferType hook is registered for " << op.value();
static auto infer_type = Op::GetAttrMap<FInferType>("FInferType");
@@ -1247,8 +1251,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ss << ref;
return ss.str();
});
- // Note: kRepr for GlobalVarNode is registered in script/printer/ir/ir.cc
- // via TVMScriptPrinter::Register<GlobalVarNode>(ReprPrintIR).
+ // Note: kRepr for GlobalVarNode is registered in script/printer/script_printer.cc.
}
} // namespace tvm
diff --git a/src/ir/prim/builtin.cc b/src/ir/prim/builtin.cc
index 76a1eff6f5..f4b7e43d57 100644
--- a/src/ir/prim/builtin.cc
+++ b/src/ir/prim/builtin.cc
@@ -52,6 +52,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("prim.vscale")
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(32))
.set_attr<TCallEffectKind>("TCallEffectKind", static_cast<int64_t>(CallEffectKind::kPure));
OpDef("prim.ceil")
diff --git a/src/ir/prim/expr.cc b/src/ir/prim/expr.cc
index 71f24e4d4a..b6e67282df 100644
--- a/src/ir/prim/expr.cc
+++ b/src/ir/prim/expr.cc
@@ -439,8 +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 via TVMScriptPrinter::Register in
- // src/script/printer/tirx/expr.cc (-> ReprPrintTIR which delegates to TVMScriptPrinter).
+ // Note: kRepr for VarNode is registered in script/printer/script_printer.cc.
}
// Cast
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index d33895c997..baf088a5a9 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -313,6 +313,10 @@ std::optional<std::string> RenderScriptWithUnderline(const ffi::ObjectRef& node,
if (!config_fn.has_value() || !script_fn.has_value()) return std::nullopt;
try {
ffi::Map<ffi::String, ffi::Any> config_dict;
+ // Error snippets use the same display conventions as script() and repr().
+ config_dict.Set("extra_config", ffi::Map<ffi::String, ffi::Any>{
+ {"script.comment_imports", true},
+ });
config_dict.Set("path_to_underline", ffi::Array<ffi::reflection::AccessPath>{path});
// Show enough context that a small function renders in full (no
// "(... N lines skipped ...)" marker), while still bounding a large
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 52487c5dbe..fb776eda1e 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -825,8 +825,7 @@ void ValidateCallPyFunc(const CallNode* call) {
TVM_FFI_ICHECK(arg_tuple.as<TupleNode>() || arg_tuple.as<VarNode>())
<< "Operation " << call->op << " must hold its arguments as an in-line tuple. "
- << "However, " << ffi::GetRef<Call>(call) << " has arguments " << arg_tuple
- << ", which is neither an in-line tuple, "
+ << "However, the argument tuple is " << arg_tuple << ", which is neither an in-line tuple, "
<< "nor a variable binding that may be normalized to an in-line tuple.";
}
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index 4ef7d1ea8f..7b30b84038 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -339,6 +339,8 @@ inline std::optional<PrimType> GetElementDType(const Type& ty) {
*/
inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, const Type& lhs_ty,
const Type& rhs_ty) {
+ // Formatting the full Call in an inference error would invoke the printer
+ // and reenter the same inference hook before the exception is thrown.
auto opt_lhs_dtype = GetElementDType(lhs_ty);
if (!opt_lhs_dtype) {
if (const auto* lhs_tensor = lhs_ty.as<TensorTypeNode>()) {
@@ -346,9 +348,8 @@ inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, cons
}
TVM_FFI_VISIT_THROW(TypeError, call)
<< "Binary operators must have the same datatype for both operands. "
- << "However, " << call << " has argument " << call->args[0] << " on the LHS, with type "
- << lhs_ty << ". This is of type " << lhs_ty->GetTypeKey()
- << ", which does not have a datatype.";
+ << "However, " << call->op << " has LHS type " << lhs_ty << ". This is of type "
+ << lhs_ty->GetTypeKey() << ", which does not have a datatype.";
}
auto lhs_dtype = opt_lhs_dtype.value();
@@ -359,9 +360,8 @@ inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, cons
}
TVM_FFI_VISIT_THROW(TypeError, call)
<< "Binary operators must have the same datatype for both operands. "
- << "However, " << call << " has argument " << call->args[1] << " on the RHS, with type "
- << rhs_ty << ". This is of type " << rhs_ty->GetTypeKey()
- << ", which does not have a datatype.";
+ << "However, " << call->op << " has RHS type " << rhs_ty << ". This is of type "
+ << rhs_ty->GetTypeKey() << ", which does not have a datatype.";
}
auto rhs_dtype = opt_rhs_dtype.value();
@@ -369,7 +369,7 @@ inline ffi::Optional<PrimType> InferBinaryArithOpOutDtype(const Call& call, cons
!rhs_dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
TVM_FFI_VISIT_THROW(TypeError, call)
<< "Binary operators must have the same datatype for both operands. "
- << "However, " << call << " uses datatype " << lhs_dtype << " on the LHS (Type of "
+ << "However, " << call->op << " uses datatype " << lhs_dtype << " on the LHS (Type of "
<< lhs_ty << "), and datatype " << rhs_dtype << " on the RHS (Type of " << rhs_ty << ").";
}
return lhs_dtype;
diff --git a/src/relax/op/tensor/binary.cc b/src/relax/op/tensor/binary.cc
index 2e1f0522e5..22603e0bec 100644
--- a/src/relax/op/tensor/binary.cc
+++ b/src/relax/op/tensor/binary.cc
@@ -42,10 +42,10 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu
TVM_FFI_CHECK(lhs_ty.as<PrimTypeNode>() || lhs_ty.as<TensorTypeNode>(), TypeError)
<< "Arguments to binary operators must be either R.Tensor or R.Prim types, "
- << "but expression " << call << " has LHS " << call->args[0] << ", which has Type " << lhs_ty;
+ << "but " << call->op << " has LHS type " << lhs_ty;
TVM_FFI_CHECK(rhs_ty.as<PrimTypeNode>() || rhs_ty.as<TensorTypeNode>(), TypeError)
<< "Arguments to binary operators must be either R.Tensor or R.Prim types, "
- << "but expression " << call << " has RHS " << call->args[1] << ", which has Type " << rhs_ty;
+ << "but " << call->op << " has RHS type " << rhs_ty;
// Dtype
ffi::Optional<PrimType> output_dtype = f_compute_out_dtype(call, ctx, lhs_ty, rhs_ty);
diff --git a/src/relax/script/printer/binding.cc b/src/relax/script/printer/binding.cc
index 0896d2e5a8..1eb978080f 100644
--- a/src/relax/script/printer/binding.cc
+++ b/src/relax/script/printer/binding.cc
@@ -16,95 +16,156 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ir/op.h>
+#include <tvm/relax/op_attr_types.h>
+#include <tvm/tirx/type.h>
-#include "./utils.h"
+#include <optional>
+
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-IfDoc PrintIfExpr(const relax::If& n, const AccessPath& n_p,
- const IRDocsifier& d, //
- const ffi::Optional<ExprDoc>& var, const ffi::Optional<ExprDoc>& ann) {
- using relax::SeqExpr;
- ExprDoc cond = d->AsDoc<ExprDoc>(n->cond, n_p->Attr("cond"));
- std::vector<ffi::Array<StmtDoc>> branches{
- PrintSeqExpr(n->true_branch, n_p->Attr("true_branch"), d, false),
- PrintSeqExpr(n->false_branch, n_p->Attr("false_branch"), d, false),
- };
- if (var.has_value()) {
- for (ffi::Array<StmtDoc>& stmts : branches) {
- ExprDoc ret = stmts.back().as_or_throw<ExprStmtDoc>()->expr;
- stmts.Set(stmts.size() - 1, AssignDoc(var.value(), ret, ann));
- }
- }
- return IfDoc(cond, branches[0], branches[1]);
+namespace {
+
+ffi::Optional<ExprDoc> MatchCastDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* binding =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::MatchCastNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ ExprDoc rhs =
+ NamespaceDoc("relax")
+ ->Attr("match_cast")
+ ->Call({d->Translate(binding->value).value(), d->Translate(binding->ty).value()});
+ IdDoc lhs = VarDoc(d, binding->var);
+ ffi::Optional<ExprDoc> annotation = std::nullopt;
+ if (!binding->var->ty.IsMissing()) annotation = d->Translate(binding->var->ty).value();
+ if (!d->GetExtraConfig<bool>("relax.show_all_ty", true)) annotation = std::nullopt;
+ d->Emit(AssignDoc(lhs, rhs, annotation), ffi::GetRef<ffi::ObjectRef>(binding));
+ return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::MatchCast>(
- "", [](relax::MatchCast n, AccessPath n_p, IRDocsifier d) -> Doc {
- using tvm::Type;
- using relax::MatchType;
- ffi::Optional<ExprDoc> ann = std::nullopt;
- if (d->cfg->GetExtraConfig<bool>("relax.show_all_ty", true)) {
- ann = TypeAsAnn(n->var, n_p->Attr("var"), d, n->value);
- }
- ExprDoc rhs = Relax(d, "match_cast")
- ->Call({d->AsDoc<ExprDoc>(n->value, n_p->Attr("value")),
- d->AsDoc<ExprDoc>(n->ty, n_p->Attr("ty"))});
- ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
- return AssignDoc(lhs, rhs, ann);
- });
+ ffi::reflection::TypeAttrDef<relax::MatchCastNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&MatchCastDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::VarBinding>( //
- "", [](relax::VarBinding n, AccessPath n_p, IRDocsifier d) -> Doc {
- if (const auto if_ = n->value.as<relax::IfNode>()) {
- ffi::Optional<ExprDoc> ann = TypeAsAnn(n->var, n_p->Attr("var"), d, n->value);
- if (!ann.has_value() && n->var->ty.as<PrimTypeNode>()) {
- ann = d->AsDoc<ExprDoc>(n->var->ty, n_p->Attr("var")->Attr("ty"));
- }
- ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
- return PrintIfExpr(ffi::GetRef<relax::If>(if_), n_p->Attr("value"), d, lhs, ann);
- } else if (n->value->IsInstance<tvm::BaseFuncNode>() &&
- !n->value->IsInstance<relax::ExternFuncNode>()) {
- IdDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
- d->cfg->binding_names.push_back(lhs->name);
- Doc ret = d->AsDoc(n->value, n_p->Attr("value"));
- d->cfg->binding_names.pop_back();
- return ret;
- } else if (d->cfg->syntax_sugar && relax::HasVoidType(n->value) &&
- relax::HasVoidType(n->var)) {
- ExprDoc rhs = d->AsDoc<ExprDoc>(n->value, n_p->Attr("value"));
- return ExprStmtDoc(rhs);
- } else {
- ExprDoc rhs = d->AsDoc<ExprDoc>(n->value, n_p->Attr("value"));
- ffi::Optional<ExprDoc> ann = TypeAsAnn(n->var, n_p->Attr("var"), d, n->value);
- if (!ann.has_value() && n->var->ty.as<PrimTypeNode>()) {
- ann = d->AsDoc<ExprDoc>(n->var->ty, n_p->Attr("var")->Attr("ty"));
- }
- ExprDoc lhs = DefineRelaxVar(n->var, d->frames.back(), d);
- return AssignDoc(lhs, rhs, ann);
+ffi::Optional<ExprDoc> VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* binding =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::VarBindingNode>(
+ input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ if (auto func = binding->value.as<relax::FunctionNode>()) {
+ d->Emit(CommentDoc("from tvm.script import relax as R"), ffi::GetRef<ffi::ObjectRef>(binding));
+ d->Emit(CommentDoc(""), ffi::GetRef<ffi::ObjectRef>(binding));
+ IdDoc lhs = VarDoc(d, binding->var);
+ d->Translate(binding->value);
+ FunctionDoc function = d->CurrentScopeDocs().back().as_or_throw<FunctionDoc>();
+ function->name = lhs;
+ ExprDoc decorator = NamespaceDoc("relax")->Attr("function");
+ if (!func->is_pure) {
+ decorator = decorator->Call({}, {"pure"}, {LiteralDoc::Boolean(false, std::nullopt)});
+ }
+ function->decorators = {decorator};
+ return std::nullopt;
+ }
+ if (relax::HasVoidType(binding->value) && relax::HasVoidType(binding->var)) {
+ VarDoc(d, binding->var);
+ ffi::Optional<ExprDoc> rhs = d->Translate(binding->value, binding->var);
+ if (rhs.has_value()) d->Emit(ExprStmtDoc(rhs.value()), ffi::GetRef<ffi::ObjectRef>(binding));
+ return std::nullopt;
+ }
+ IdDoc lhs = VarDoc(d, binding->var);
+ ffi::Optional<ExprDoc> rhs = d->Translate(binding->value, binding->var);
+ if (!rhs.has_value()) return std::nullopt;
+ ffi::Optional<ExprDoc> annotation = std::nullopt;
+ bool infer_vdevice = false;
+ if (auto tensor = binding->var->ty.as<relax::TensorTypeNode>();
+ tensor && tensor->vdevice.has_value()) {
+ if (auto call = binding->value.as<CallNode>()) {
+ if (auto op = call->op.as<Op>()) infer_vdevice = op.value()->name == "relax.to_vdevice";
+ }
+ }
+ bool show_all_ty = d->GetExtraConfig<bool>("relax.show_all_ty", true);
+ const auto* call = binding->value.as<CallNode>();
+ bool output_type_argument = call && call->ty_args.size() == 1 &&
+ ffi::StructuralEqual()(binding->var->ty, call->ty_args[0]);
+ ffi::Optional<Type> inferred = std::nullopt;
+ if (!show_all_ty || output_type_argument) {
+ if (!call) {
+ inferred = binding->value->ty;
+ } else if (auto op = call->op.as<Op>()) {
+ static const auto fixed = Op::GetAttrMap<TFixedReturnType>("TFixedReturnType");
+ static const auto context_free = Op::GetAttrMap<FInferType>("FInferType");
+ if (fixed.count(op.value()) || context_free.count(op.value())) {
+ try {
+ inferred = Call::ReinferType(call);
+ } catch (const ffi::Error&) {
+ // Contextual inference remains the parser's responsibility.
}
- });
+ }
+ }
+ }
+ bool inferable = inferred.has_value() && !inferred.value().IsMissing() &&
+ ffi::StructuralEqual()(binding->var->ty, inferred.value());
+ bool explicit_output_type = output_type_argument && inferable;
+ // Primitive aliases need their annotation. Without context-free inference,
+ // keep the binding type and let the parser's deferred inference use it.
+ bool elide_annotation = infer_vdevice || explicit_output_type ||
+ (!show_all_ty && !binding->var->ty.as<PrimTypeNode>() && inferable);
+ if (!binding->var->ty.IsMissing() && !elide_annotation) {
+ annotation = d->Translate(binding->var->ty).value();
+ }
+ d->Emit(AssignDoc(lhs, rhs.value(), annotation), ffi::GetRef<ffi::ObjectRef>(binding));
+ return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::If>(
- "", [](relax::If n, AccessPath n_p, IRDocsifier d) -> Doc {
- return PrintIfExpr(n, n_p, d, std::nullopt, std::nullopt);
- });
+ ffi::reflection::TypeAttrDef<relax::VarBindingNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&VarBindingDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* branch =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::IfNode>(input);
+ // A value context needs a usable expression after the conditional. A
+ // destination supplied by a binding is completed directly on both arms.
+ ffi::Optional<Var> temporary = std::nullopt;
+ if (!destination) {
+ temporary = Var("if_result", branch->ty);
+ VarDoc(d, temporary.value());
+ destination = temporary.value().get();
+ }
+ TVM_FFI_CHECK(destination->IsInstance<VarNode>(), TypeError)
+ << "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;
+ if (!var->ty.IsMissing()) annotation = d->Translate(var->ty).value();
+ ExprDoc condition = d->Translate(branch->cond).value();
+ d->Emit(IfDoc(condition, RelaxSeqBody(d, branch->true_branch.get(), lhs, annotation, destination),
+ RelaxSeqBody(d, branch->false_branch.get(), lhs, annotation, destination)),
+ ffi::GetRef<ffi::ObjectRef>(branch));
+ return temporary.has_value() ? ffi::Optional<ExprDoc>(lhs.value()) : std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::MatchCastNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::VarBindingNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::IfNode>(ReprPrintRelax);
+ ffi::reflection::TypeAttrDef<relax::IfNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&IfDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index 6410895c8f..4a7b514dcc 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -16,344 +16,309 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/reflection/accessor.h>
-#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/op.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/distributed/type.h>
+#include <tvm/tirx/op_attr_types.h>
+#include <tvm/tirx/type.h>
+#include <algorithm>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "../../../script/printer/ir/utils.h"
#include "../../../tirx/script/printer/utils.h"
-#include "./utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-class AttrPrinter {
- public:
- explicit AttrPrinter(AccessPath p, const IRDocsifier& d, ffi::Array<ffi::String>* keys,
- ffi::Array<ExprDoc>* values)
- : p(std::move(p)), d(d), keys(keys), values(values) {}
-
- void operator()(const Attrs& attrs) {
- if (const auto* dict_attrs = attrs.as<DictAttrsNode>()) {
- for (const auto& [key, value] : dict_attrs->dict) {
- keys->push_back(key);
- values->push_back(d->AsDoc<ExprDoc>(value, p->Attr(key)));
+// Shared Call takes Expr values rather than a named wrapper's Python sugar.
+ExprDoc MaterializeCallArgument(DocTranslatorObj* d, const Expr& arg, ExprDoc doc) {
+ if (const auto* tuple = arg.as<TupleNode>()) {
+ if (auto fields = doc.as<TupleDocNode>()) {
+ ffi::Array<ExprDoc> values;
+ for (size_t i = 0; i < tuple->fields.size(); ++i) {
+ values.push_back(MaterializeCallArgument(d, tuple->fields[i], fields->elements[i]));
}
- } else {
- const TVMFFITypeInfo* attrs_tinfo = TVMFFIGetTypeInfo(attrs->type_index());
- TVM_FFI_ICHECK(attrs_tinfo->metadata != nullptr)
- << "Object `" << attrs->GetTypeKey()
- << "` misses reflection registration and do not support serialization";
- // new printing mechanism using the new reflection
- ffi::reflection::ForEachFieldInfo(attrs_tinfo, [&](const TVMFFIFieldInfo* field_info) {
- ffi::String field_name = ffi::String(field_info->name);
- Any field_value = ffi::reflection::FieldGetter(field_info)(attrs);
- keys->push_back(field_name);
- values->push_back(d->AsDoc<ExprDoc>(field_value, p->Attr(field_name)));
- });
+ doc = NamespaceDoc("relax")->Attr("tuple")->Call(values);
}
}
+ if (const auto* region = arg.as<TensorRegionNode>()) {
+ doc = TensorRegionValue(d, region, true);
+ }
+ d->RecordOrigin(doc, arg);
+ return doc;
+}
- AccessPath p;
- const IRDocsifier& d;
- ffi::Array<ffi::String>* keys;
- ffi::Array<ExprDoc>* values;
-};
+namespace {
-ExprDoc PrintCallee(const relax::Expr& n, const AccessPath& n_p, const IRDocsifier& d) {
- // TODO(@junrushao): handle callee better
- if (const auto* ext = n.as<relax::ExternFuncNode>()) {
- return LiteralDoc::Str(ext->global_symbol, n_p);
- } else {
- return d->AsDoc<ExprDoc>(n, n_p);
- }
+// Relax's named constructors defer result typing until a binding supplies it.
+// A standalone named call therefore reconstructs only a missing stored result.
+bool HasRelaxCallResult(const CallNode* call, const ffi::Object* destination) {
+ if (call->ty.IsMissing()) return true;
+ const auto* var = destination && destination->IsInstance<VarNode>()
+ ? static_cast<const VarNode*>(destination)
+ : nullptr;
+ return var && ffi::StructuralEqual()(var->ty, call->ty);
}
-ffi::Optional<ExprDoc> PrintCallTIRDPSPacked(const Call& n, const AccessPath& n_p,
- const IRDocsifier& d) {
- static const Op call_tir_op = Op::Get("relax.call_tir");
- static const Op call_tir_inplace_op = Op::Get("relax.call_tir_inplace");
- static const Op call_dps_packed_op = Op::Get("relax.call_dps_packed");
- static const Op call_tir_with_grad_op = Op::Get("relax.call_tir_with_grad");
- static const Op call_tir_local_view = Op::Get("relax.dist.call_tir_local_view");
- if (!n->op.same_as(call_tir_op) && !n->op.same_as(call_dps_packed_op) &&
- !n->op.same_as(call_tir_with_grad_op) && !n->op.same_as(call_tir_local_view) &&
- !n->op.same_as(call_tir_inplace_op)) {
- return std::nullopt;
- }
- TVM_FFI_ICHECK_EQ(n->args.size(), 2);
- TVM_FFI_ICHECK(n->ty_args.size() == 1);
- ffi::Array<ExprDoc> args;
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- // Step 1. Print n->args[0], the callee
- args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayItem(0), d));
- // Step 2. Print n->args[1], the input arguments
- args.push_back(d->AsDoc<ExprDoc>(n->args[1], n_p->Attr("args")->ArrayItem(1)));
- // Step 3. Print n->ty_args, the output type
- Type out_ty = n->ty_args[0];
- AccessPath out_ty_p = n_p->Attr("ty_args")->ArrayItem(0);
- bool is_dtensor = false;
- kwargs_keys.push_back("out_ty");
- if (const auto* o = out_ty.as<relax::TupleTypeNode>()) {
- ffi::Array<ExprDoc> fields;
- AccessPath fields_p = out_ty_p->Attr("fields");
- for (int i = 0, l = o->fields.size(); i < l; ++i) {
- if (o->fields[i].as<relax::distributed::DTensorTypeNode>()) {
- is_dtensor = true;
- }
- fields.push_back(d->AsDoc<ExprDoc>(o->fields[i], fields_p->ArrayItem(i)));
+ffi::Optional<ExprDoc> RelaxCallDocTranslate(DocTranslatorObj* d, const CallNode* call,
+ const ffi::Array<ExprDoc>& args) {
+ const Op& op = call->op.as_or_throw<Op>();
+ auto relax_format = [&](size_t index) -> ExprDoc {
+ if (const auto* format = call->args[index].as<StringImmNode>()) {
+ return NamespaceDoc("relax")->Attr("str")->Call(
+ {LiteralDoc::Str(format->value, std::nullopt)});
}
- kwargs_values.push_back(ListDoc(fields));
- } else {
- if (out_ty.as<relax::distributed::DTensorTypeNode>()) {
- is_dtensor = true;
- }
- kwargs_values.push_back(d->AsDoc<ExprDoc>(out_ty, out_ty_p));
- }
+ return args[index];
+ };
- // for call_tir_inplace, we also need to include the inplace args
- if (n->op.same_as(call_tir_inplace_op)) {
- kwargs_keys.push_back("inplace_indices");
- ffi::Array<ExprDoc> index_fields;
- if (auto* call_tir_inplace_attrs = n->attrs.as<relax::CallTIRInplaceAttrs>()) {
- for (int64_t inplace_index : call_tir_inplace_attrs->inplace_indices) {
- index_fields.push_back(
- LiteralDoc::Int(inplace_index, n_p->Attr("attrs")->Attr("inplace_indices")));
- }
+ if (op->name == "relax.call_py_func" && call->args.size() == 2 && !call->attrs.defined() &&
+ call->ty_args.size() == 1) {
+ const auto* tuple = call->args[1].as<TupleNode>();
+ if (!tuple || !call->args[0].as<StringImmNode>()) return std::nullopt;
+ const Type& output = call->ty_args[0];
+ const auto* output_tuple = output.as<TupleTypeNode>();
+ if (output_tuple && output_tuple->fields.size() == 1) return std::nullopt;
+ ffi::Array<Type> types = output_tuple ? output_tuple->fields : ffi::Array<Type>{output};
+ ffi::Array<ExprDoc> outputs;
+ for (const Type& type : types) {
+ const auto* tensor = type.as<relax::TensorTypeNode>();
+ if (!tensor || !tensor->shape.as<relax::ShapeExprNode>()) return std::nullopt;
+ outputs.push_back(d->Translate(type).value());
}
- kwargs_values.push_back(ListDoc(index_fields));
+ ExprDoc output_doc = output_tuple ? ExprDoc(ListDoc(outputs)) : outputs[0];
+ d->RecordOrigin(output_doc, output);
+ ffi::Array<ExprDoc> packed_args = {args[0]};
+ for (const Expr& arg : tuple->fields) packed_args.push_back(d->Translate(arg).value());
+ return NamespaceDoc("relax")->Attr("call_py_func")->Call(packed_args, {"out_ty"}, {output_doc});
}
- // start of specially handling call_tir_with_grad
- if (const auto* call_tir_with_grad_attrs = n->attrs.as<relax::CallTIRWithGradAttrs>()) {
- kwargs_keys.push_back("te_grad_name");
- kwargs_values.push_back(LiteralDoc::Str(call_tir_with_grad_attrs->te_grad_name,
- n_p->Attr("attrs")->Attr("te_grad_name")));
- if (!call_tir_with_grad_attrs->te_grad_kwargs.empty()) {
- kwargs_keys.push_back("te_grad_kwargs");
- kwargs_values.push_back(d->AsDoc<ExprDoc>(call_tir_with_grad_attrs->te_grad_kwargs,
- n_p->Attr("attrs")->Attr("te_grad_kwargs")));
- }
- }
- if (n->op.same_as(call_tir_with_grad_op)) {
- return Relax(d, "call_tir_with_grad")->Call(args, kwargs_keys, kwargs_values);
+ if (op->name == "relax.print" && !call->args.empty() && !call->attrs.defined() &&
+ call->ty_args.empty()) {
+ ffi::Array<ExprDoc> positional;
+ for (size_t i = 1; i < args.size(); ++i) positional.push_back(args[i]);
+ return NamespaceDoc("relax")->Attr("print")->Call(positional, {"format"}, {relax_format(0)});
}
- // end of specially handling call_tir_with_grad
- if (n->op.same_as(call_dps_packed_op)) {
- return Relax(d, "call_dps_packed")->Call(args, kwargs_keys, kwargs_values);
+ if (op->name == "relax.assert_op" && call->args.size() >= 2 && !call->attrs.defined() &&
+ call->ty_args.empty()) {
+ ffi::Array<ExprDoc> positional = {args[0]};
+ for (size_t i = 2; i < args.size(); ++i) positional.push_back(args[i]);
+ return NamespaceDoc("relax")
+ ->Attr("assert_op")
+ ->Call(positional, {"format"}, {relax_format(1)});
}
- if (n->op.same_as(call_tir_local_view)) {
- return Relax(d, "dist.call_tir_local_view")->Call(args, kwargs_keys, kwargs_values);
- } else if (is_dtensor) {
- return Relax(d, "dist.call_tir")->Call(args, kwargs_keys, kwargs_values);
- } else if (n->op.same_as(call_tir_inplace_op)) {
- return Relax(d, "call_tir_inplace")->Call(args, kwargs_keys, kwargs_values);
- } else {
- return Relax(d, "call_tir")->Call(args, kwargs_keys, kwargs_values);
+
+ if (op->name == "relax.hint_on_device") {
+ const auto* attrs = call->attrs.as<relax::HintOnDeviceAttrs>();
+ if (!attrs || args.size() != 1 || !call->ty_args.empty()) return std::nullopt;
+ ExprDoc device = NamespaceDoc("relax")->Attr("device")->Call(
+ {LiteralDoc::Int(attrs->device_type, std::nullopt),
+ LiteralDoc::Int(attrs->index, std::nullopt)});
+ return NamespaceDoc("relax")
+ ->Attr("hint_on_device")
+ ->Call({args[0], device, LiteralDoc::Str(attrs->memory_scope, std::nullopt)});
}
-}
-ffi::Optional<ExprDoc> PrintAssertOp(const Call& n, const AccessPath& n_p, const IRDocsifier& d) {
- static const Op assert_op = Op::Get("relax.assert_op");
- if (!n->op.same_as(assert_op)) {
+ static const auto& names = Op::GetAttrMap<tirx::TScriptPrinterName>("TScriptPrinterName");
+ if (!names.count(op) || names[op].find("relax.") != 0 || !call->ty_args.empty()) {
return std::nullopt;
}
- TVM_FFI_ICHECK(n->args.size() >= 2);
- // special handling: it is important to indicate that the format string (second argument)
- // is the _format_ string, or else roundtripping will fail
- // (the format string will be interpreted as an argument and there will be a new default format
- // string given)
- ffi::Array<ExprDoc> args;
- args.push_back(d->AsDoc<ExprDoc>(n->args[0], n_p->Attr("args")->ArrayItem(0)));
- ExprDoc second_arg = d->AsDoc<ExprDoc>(n->args[1], n_p->Attr("args")->ArrayItem(1));
- for (size_t i = 2; i < n->args.size(); i++) {
- args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayItem(i)));
+ if (call->attrs.defined() ? op->attrs_type_key != call->attrs->GetTypeKey()
+ : !op->attrs_type_key.empty()) {
+ return std::nullopt;
}
- return Relax(d, "assert_op")->Call(args, {"format"}, {second_arg});
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (call->attrs.defined()) {
+ std::vector<std::pair<ffi::String, ffi::Any>> fields;
+ if (auto attrs = call->attrs.as<DictAttrsNode>()) {
+ for (const auto& [key, value] : attrs->dict) fields.emplace_back(key, value);
+ } else {
+ ffi::reflection::ForEachFieldInfo(
+ TVMFFIGetTypeInfo(call->attrs->type_index()), [&](const TVMFFIFieldInfo* field) {
+ fields.emplace_back(ffi::String(field->name),
+ ffi::reflection::FieldGetter(field)(call->attrs));
+ });
+ }
+ std::sort(fields.begin(), fields.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ for (const auto& [key, value] : fields) {
+ if (key == "ty" || key == "ty_args" || key == "attrs") return std::nullopt;
+ keys.push_back(key);
+ values.push_back(AnyValue(d, value));
+ }
+ }
+ ffi::Array<ExprDoc> positional = args;
+ if (op->name == "relax.call_builtin_with_ctx" && !call->args.empty() &&
+ call->args[0].as<StringImmNode>()) {
+ // A Python string would be converted to an ExternFunc by this constructor.
+ positional.Set(0, AnyValue(d, call->args[0]));
+ }
+ return NamedCallCallee(names[op])->Call(positional, keys, values);
}
-ffi::Optional<ExprDoc> PrintHintOnDevice(const Call& n, const AccessPath& n_p,
- const IRDocsifier& d) {
- static const Op hint_on_device_op = Op::Get("relax.hint_on_device");
- if (!n->op.same_as(hint_on_device_op)) {
- return std::nullopt;
+ffi::Optional<ExprDoc> CallTIRDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input);
+ if (!HasRelaxCallResult(call, destination) || call->args.size() != 2 ||
+ !call->args[1].as<TupleNode>() || call->ty_args.size() != 1) {
+ return RawCall(d, call, false);
}
- ffi::Array<ExprDoc> args;
-
- args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayItem(0), d));
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- TVM_FFI_ICHECK(n->attrs.defined());
- if (n->attrs.as<relax::HintOnDeviceAttrs>()) {
- AttrPrinter(n_p->Attr("attrs"), d, &kwargs_keys, &kwargs_values)(n->attrs);
- ExprDoc scope_val = kwargs_values.back();
- kwargs_keys.pop_back();
- kwargs_values.pop_back();
- args.push_back(Relax(d, "device")->Call({}, kwargs_keys, kwargs_values));
- args.push_back(scope_val);
+ const Op& op = call->op.as_or_throw<Op>();
+ const auto* inplace = call->attrs.as<relax::CallTIRInplaceAttrs>();
+ const auto* grad = call->attrs.as<relax::CallTIRWithGradAttrs>();
+ if (op->name == "relax.call_tir_inplace" ? inplace == nullptr
+ : op->name == "relax.call_tir_with_grad" ? grad == nullptr
+ : call->attrs.defined()) {
+ return RawCall(d, call, false);
}
- return Relax(d, "hint_on_device")->Call(args);
-}
-ffi::Optional<ExprDoc> PrintToVDevice(const Call& n, const AccessPath& n_p, const IRDocsifier& d) {
- static const Op to_vdevice_op = Op::Get("relax.to_vdevice");
- if (!n->op.same_as(to_vdevice_op)) {
- return std::nullopt;
+ const Type& output = call->ty_args[0];
+ const auto* tuple = output.as<TupleTypeNode>();
+ // The constructor unwraps a single-element output list into a tensor type.
+ if (tuple && tuple->fields.size() == 1) return RawCall(d, call, false);
+ ffi::Array<Type> output_types = tuple ? tuple->fields : ffi::Array<Type>{output};
+ bool distributed =
+ !output_types.empty() && output_types[0].as<relax::distributed::DTensorTypeNode>();
+ if (distributed && op->name != "relax.call_tir") return RawCall(d, call, false);
+ ffi::Array<ExprDoc> output_docs;
+ for (const Type& type : output_types) {
+ if (distributed) {
+ const auto* tensor = type.as<relax::distributed::DTensorTypeNode>();
+ if (!tensor || !tensor->tensor_ty->shape.as<relax::ShapeExprNode>()) {
+ return RawCall(d, call, false);
+ }
+ } else {
+ const auto* tensor = type.as<relax::TensorTypeNode>();
+ if (!tensor || !tensor->shape.as<relax::ShapeExprNode>()) return RawCall(d, call, false);
+ }
+ output_docs.push_back(d->Translate(type).value());
}
+ ExprDoc output_doc = tuple ? ExprDoc(ListDoc(output_docs)) : output_docs[0];
+ d->RecordOrigin(output_doc, output);
ffi::Array<ExprDoc> args;
-
- args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayItem(0), d));
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- TVM_FFI_ICHECK(n->attrs.defined());
- if (const auto* attrs = n->attrs.as<relax::ToVDeviceAttrs>()) {
- relax::VDevice vdev = attrs->dst_vdevice;
- std::string dev_kind = vdev->target->kind->name;
- int dev_index = FindVDeviceIndexByTargetKind(vdev, d);
- kwargs_keys.push_back("dst_vdevice");
- kwargs_values.push_back(
- LiteralDoc::Str(dev_kind + ":" + std::to_string(dev_index) + ":" + vdev->memory_scope,
- n_p->Attr("dst_vdevice")));
+ for (const Expr& arg : call->args) args.push_back(d->Translate(arg).value());
+ if (op->name == "relax.call_dps_packed" || distributed) {
+ if (const auto* packed = call->args[0].as<relax::ExternFuncNode>();
+ packed &&
+ ffi::StructuralEqual()(packed->ty, relax::ExternFunc(packed->global_symbol)->ty)) {
+ ExprDoc symbol = LiteralDoc::Str(packed->global_symbol, std::nullopt);
+ d->RecordOrigin(symbol, call->args[0]);
+ args.Set(0, symbol);
+ } else if (call->args[0].as<StringImmNode>()) {
+ args.Set(0, AnyValue(d, call->args[0]));
+ }
+ }
+ ffi::Array<ffi::String> keys = {"out_ty"};
+ ffi::Array<ExprDoc> values = {output_doc};
+ if (inplace) {
+ keys.push_back("inplace_indices");
+ values.push_back(AnyValue(d, inplace->inplace_indices));
+ } else if (grad) {
+ keys.push_back("te_grad_name");
+ values.push_back(LiteralDoc::Str(grad->te_grad_name, std::nullopt));
+ keys.push_back("te_grad_kwargs");
+ values.push_back(AnyValue(d, grad->te_grad_kwargs));
}
- return Relax(d, "to_vdevice")->Call(args, kwargs_keys, kwargs_values);
+ return NamedCallCallee(distributed ? ffi::String("relax.dist.call_tir") : op->name)
+ ->Call(args, keys, values);
}
-ffi::Optional<ExprDoc> PrintRelaxPrint(const Call& n, const AccessPath& n_p, const IRDocsifier& d) {
- static const Op print_op = Op::Get("relax.print");
- if (!n->op.same_as(print_op)) {
- return std::nullopt;
- }
- TVM_FFI_ICHECK(n->args.size() >= 1);
- // special handling: it is important to indicate that the format string (first argument)
- // is the _format_ string, or else roundtripping will fail
- // (the format string will be interpreted as an argument and there will be a new default format
- // string given)
- ExprDoc first_arg = d->AsDoc<ExprDoc>(n->args[0], n_p->Attr("args")->ArrayItem(0));
- ffi::Array<ExprDoc> args;
- for (size_t i = 1; i < n->args.size(); i++) {
- args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayItem(i)));
+TVM_FFI_STATIC_INIT_BLOCK() {
+ for (const char* name : {"relax.call_tir", "relax.call_tir_with_grad", "relax.call_tir_inplace",
+ "relax.call_dps_packed"}) {
+ OpDef(name).set_attr<FDocTranslate>(kOpCallDocTranslate,
+ FDocTranslate::FromNative<&CallTIRDocTranslate>());
}
- return Relax(d, "print")->Call(args, {"format"}, {first_arg});
}
-bool ShouldPrintAsTIR(const Call& call) {
- if (call->op->ty.as<relax::FuncTypeNode>() || call->op.as<tvm::VarNode>() ||
- call->op.as<relax::FunctionNode>() || call->op.as<relax::ExternFuncNode>()) {
- return false;
+ffi::Optional<ExprDoc> CallDefaultDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input);
+ if (auto op = call->op.as<Op>(); op && op.value()->name.find("relax.") == 0) {
+ if (HasRelaxCallResult(call, destination)) {
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) args.push_back(d->Translate(arg).value());
+ if (auto doc = RelaxCallDocTranslate(d, call, args)) return doc;
+ return RawCall(d, call, false, args);
+ }
+ return RawCall(d, call, false);
+ }
+ if ((call->op.as<VarNode>() || call->op.as<GlobalVarNode>() ||
+ call->op.as<relax::FunctionNode>()) &&
+ call->op->ty.as<relax::FuncTypeNode>() && HasRelaxCallResult(call, destination) &&
+ !call->attrs.defined() && call->ty_args.empty()) {
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) {
+ args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ }
+ return d->Translate(call->op).value()->Call(args);
}
- if (auto op = call->op.as<Op>()) {
- return op.value()->name.find("relax.") != 0;
+ ffi::Optional<Type> inferred = std::nullopt;
+ if (call->op.as<Op>() && std::all_of(call->args.begin(), call->args.end(),
+ [](const Expr& arg) { return !arg->ty.IsMissing(); })) {
+ try {
+ inferred = Call::ReinferType(call);
+ } catch (const ffi::Error&) {
+ // A registered hook may require information absent from this Call.
+ }
}
- if (!call->ty.as<PrimTypeNode>()) {
- return false;
+ if (!inferred || !ffi::StructuralEqual()(inferred.value(), call->ty)) {
+ return RawCall(d, call, false);
}
- return true;
+ const Type& result_type = inferred.value();
+ if (auto doc = TIRCallPrefixDocTranslate(d, call)) return doc;
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) args.push_back(d->Translate(arg).value());
+ if (auto doc = FFIKernelDocTranslate(d, call, result_type, args)) return doc;
+ if (auto doc = TIRCallDocTranslate(d, call, result_type, args)) return doc;
+ return RawCall(d, call, true, args);
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Call>( //
- "", [](Call call, AccessPath n_p, IRDocsifier d) -> Doc {
- if (ShouldPrintAsTIR(call)) {
- return PrintTIRCall(call, n_p, d);
- }
- Call n = call;
- // Special case: call_tir, call_dps_packed, call_tir_with_grad
- if (ffi::Optional<ExprDoc> doc = PrintCallTIRDPSPacked(n, n_p, d)) {
- return doc.value();
- }
- // Special case: assert_op
- if (ffi::Optional<ExprDoc> doc = PrintAssertOp(n, n_p, d)) {
- return doc.value();
- }
- // Special case: hint_on_device
- if (ffi::Optional<ExprDoc> doc = PrintHintOnDevice(n, n_p, d)) {
- return doc.value();
- }
- // Special case: to_vdevice
- if (ffi::Optional<ExprDoc> doc = PrintToVDevice(n, n_p, d)) {
- return doc.value();
- }
- // Special case: print
- if (ffi::Optional<ExprDoc> doc = PrintRelaxPrint(n, n_p, d)) {
- return doc.value();
- }
- ExprDoc prefix{ffi::UnsafeInit()};
- ffi::Array<ExprDoc> args;
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- // Step 1. Print op
- if (const auto* op = n->op.as<relax::ExternFuncNode>()) {
- prefix = Relax(d, "call_packed");
- args.push_back(LiteralDoc::Str(op->global_symbol, n_p->Attr("op")));
- } else if (const auto* op = n->op.as<tvm::OpNode>()) {
- std::string name = op->name;
- if (name.rfind("relax.", 0) == 0) {
- prefix = Relax(d, name.substr(6));
- } else {
- prefix = IdDoc(name);
- }
- prefix->source_paths.push_back(n_p->Attr("op"));
- } else if (n->op->IsInstance<tvm::VarNode>() || n->op->IsInstance<tvm::GlobalVarNode>()) {
- prefix = d->AsDoc<ExprDoc>(n->op, n_p->Attr("op"));
- } else {
- TVM_FFI_THROW(TypeError) << "Unsupported op: " << n->op->GetTypeKey();
- }
- // Step 2. Print args
- if (!n->args.empty()) {
- args.push_back(PrintCallee(n->args[0], n_p->Attr("args")->ArrayItem(0), d));
- }
- for (int i = 1, l = n->args.size(); i < l; ++i) {
- args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayItem(i)));
- }
- // Step 3. Print attrs
- if (n->attrs.defined()) {
- if (n->op->IsInstance<relax::ExternFuncNode>()) {
- kwargs_keys.push_back("attrs_type_key");
- kwargs_values.push_back(LiteralDoc::Str(n->attrs->GetTypeKey(), n_p->Attr("attrs")));
- }
- if (const auto* attrs = n->attrs.as<tvm::DictAttrsNode>()) {
- std::vector<std::pair<ffi::String, ffi::Any>> sorted;
- for (const auto& kv : attrs->dict) {
- sorted.push_back(kv);
- }
- std::sort(sorted.begin(), sorted.end(),
- [](const auto& a, const auto& b) { return a.first < b.first; });
- for (const auto& kv : sorted) {
- kwargs_keys.push_back(kv.first);
- kwargs_values.push_back(
- d->AsDoc<ExprDoc>(kv.second, n_p->Attr("attrs")->Attr(kv.first)));
- }
- } else {
- AttrPrinter(n_p->Attr("attrs"), d, &kwargs_keys, &kwargs_values)(n->attrs);
- }
- }
- // Step 4. Print type_args
- if (n->ty_args.size() > 0) {
- AccessPath ty_args_p = n_p->Attr("ty_args");
- ffi::Array<ExprDoc> ty_args;
- for (int i = 0, l = n->ty_args.size(); i < l; ++i) {
- ty_args.push_back(d->AsDoc<ExprDoc>(n->ty_args[i], ty_args_p->ArrayItem(i)));
- }
- kwargs_keys.push_back("ty_args");
- kwargs_values.push_back(TupleDoc(ty_args));
- }
- return prefix->Call(args, kwargs_keys, kwargs_values);
- });
+ffi::Optional<ExprDoc> CallDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const CallNode>(input);
+ if (auto op = call->op.as<Op>(); op && Op::HasAttrMap(kOpCallDocTranslate)) {
+ static const auto& overrides = Op::GetAttrMap<ffi::Any>(kOpCallDocTranslate);
+ if (overrides.count(op.value())) {
+ ffi::Any hook = overrides[op.value()];
+ Call object = ffi::GetRef<Call>(call);
+ if (hook.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>(
+ reinterpret_cast<decltype(DocTranslatorVTable::translate)>(hook.cast<void*>())(
+ d, object, destination))
+ .value();
+ }
+ ffi::Any destination_arg = nullptr;
+ if (destination) destination_arg = ffi::GetRef<ffi::ObjectRef>(destination);
+ return hook.cast<ffi::Function>()
+ .CallExpected<ffi::Optional<ExprDoc>>(d, object, destination_arg)
+ .value();
+ }
+ }
+ return CallDefaultDocTranslate(d, input, destination);
}
-std::string ReprPrintCall(const ffi::ObjectRef& obj, const PrinterConfig& cfg) {
- Call call = obj.as_or_throw<Call>();
- if (ShouldPrintAsTIR(call)) {
- return ReprPrintTIR(obj, cfg);
- }
- return ReprPrintRelax(obj, cfg);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<CallNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&CallDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() { TVMScriptPrinter::Register<CallNode>(ReprPrintCall); }
+} // namespace
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/dependent_type.cc b/src/relax/script/printer/dependent_type.cc
index ec8115897a..1dea84cb4e 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -16,148 +16,137 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/cast.h>
+
+#include <tvm/tirx/type.h>
+
+#include <optional>
#include "../../../script/printer/ir/utils.h"
-#include "./utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<AnyType>( //
- "", [](AnyType n, AccessPath n_p, IRDocsifier d) -> Doc { return Relax(d, "Any"); });
+namespace {
+
+ffi::Optional<ExprDoc> AnyTypeDocTranslate(DocTranslatorObj*, ffi::AnyView, const ffi::Object*) {
+ return NamespaceDoc("relax")->Attr("Any");
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::ShapeType>(
- "", [](relax::ShapeType n, AccessPath n_p, IRDocsifier d) -> Doc {
- if (n->values.has_value()) {
- ffi::Array<PrimExpr> shape = n->values.value();
- AccessPath shape_p = n_p->Attr("values");
- ffi::Array<ExprDoc> shape_docs;
- for (int i = 0, ndim = shape.size(); i < ndim; ++i) {
- shape_docs.push_back(d->AsDoc<ExprDoc>(shape[i], shape_p->ArrayItem(i)));
- }
- return Relax(d, "Shape")->Call({ListDoc(shape_docs)});
- }
- return Relax(d, "Shape")->Call({}, {"ndim"}, {LiteralDoc::Int(n->ndim, n_p->Attr("ndim"))});
- });
+ ffi::reflection::TypeAttrDef<AnyTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&AnyTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> ShapeTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::ShapeTypeNode>(input);
+ if (ty->values.has_value()) {
+ ffi::Array<ExprDoc> values;
+ for (const PrimExpr& value : ty->values.value()) values.push_back(RelaxShapeDim(d, value));
+ return NamespaceDoc("relax")->Attr("Shape")->Call({ListDoc(values)});
+ }
+ return NamespaceDoc("relax")->Attr("Shape")->Call({}, {"ndim"},
+ {LiteralDoc::Int(ty->ndim, std::nullopt)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::TensorType>( //
- "", [](relax::TensorType n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> args;
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- if (n->shape.has_value()) {
- // Need to dig into ShapeExpr to preserve the `R.shape` prefix
- if (const auto* shape = n->shape.value().as<relax::ShapeExprNode>()) {
- auto shape_expr = ffi::GetRef<relax::ShapeExpr>(shape);
- AccessPath shape_p = n_p->Attr("shape")->Attr("values");
- ffi::Array<ExprDoc> shape_docs;
- for (int i = 0, ndim = shape_expr->values.size(); i < ndim; ++i) {
- shape_docs.push_back(d->AsDoc<ExprDoc>(shape_expr->values[i], shape_p->ArrayItem(i)));
- }
- args.push_back(TupleDoc(shape_docs));
- } else {
- args.push_back(d->AsDoc<ExprDoc>(n->shape.value(), n_p->Attr("shape")));
- }
- }
- if (!n->IsUnknownDtype()) {
- kwargs_keys.push_back("dtype");
- kwargs_values.push_back(
- LiteralDoc::DataType(n->dtype.value()->dtype, n_p->Attr("dtype")));
- }
- if (!n->shape.has_value() && !n->IsUnknownNdim()) {
- kwargs_keys.push_back("ndim");
- kwargs_values.push_back(LiteralDoc::Int(n->ndim, n_p->Attr("ndim")));
- }
- if (n->vdevice.has_value() && n->vdevice.value()->target.defined()) {
- // Function annotations defer module-owned selectors until declaration.
- bool has_relax_frame = false;
- bool has_module_frame = false;
- for (const Frame& frame : d->frames) {
- if (const auto* relax_frame = frame.as<RelaxFrameNode>()) {
- has_relax_frame = true;
- if (relax_frame->func_vars != nullptr) {
- d->ir_usage.insert("future_annotations");
- }
- } else if (const auto* ir_frame = frame.as<IRFrameNode>()) {
- has_module_frame = ir_frame->global_infos != nullptr;
- }
- }
- kwargs_keys.push_back("vdevice");
- if (has_relax_frame) {
- std::string dev_kind = n->vdevice.value()->target->kind->name;
- int dev_index = FindVDeviceIndexByTargetKind(n->vdevice.value(), d);
- kwargs_values.push_back(LiteralDoc::Str(
- dev_kind + ":" + std::to_string(dev_index) + ":" + n->vdevice.value()->memory_scope,
- n_p->Attr("vdevice")));
- } else if (has_module_frame) {
- // Class assignments execute eagerly; retrieve their existing concrete metadata.
- kwargs_values.push_back(
- Relax(d, "lookup_vdevice")
- ->Call({LiteralDoc::Str(n->vdevice.value()->target->kind->name,
- n_p->Attr("vdevice")),
- LiteralDoc::Int(FindVDeviceIndexByTargetKind(n->vdevice.value(), d),
- n_p->Attr("vdevice"))}));
- } else {
- kwargs_values.push_back(d->AsDoc<ExprDoc>(n->vdevice.value(), n_p->Attr("vdevice")));
- }
- }
- if (args.empty() && kwargs_keys.empty()) {
- return Relax(d, "Tensor");
- }
- return Relax(d, "Tensor")->Call(args, kwargs_keys, kwargs_values);
- });
+ ffi::reflection::TypeAttrDef<relax::ShapeTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ShapeTypeDocTranslate>());
+}
+
+} // namespace
+
+ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
+ bool include_vdevice) {
+ ffi::Array<ExprDoc> args;
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (ty->shape.has_value()) {
+ if (auto shape = ty->shape.value().as<relax::ShapeExprNode>()) {
+ ffi::Array<ExprDoc> dimensions;
+ for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim));
+ args.push_back(TupleDoc(dimensions));
+ } else {
+ args.push_back(d->Translate(ty->shape.value()).value());
+ }
+ }
+ if (!ty->IsUnknownDtype()) {
+ keys.push_back("dtype");
+ values.push_back(LiteralDoc::DataType(ty->dtype.value()->dtype, std::nullopt));
+ }
+ if (!ty->shape.has_value() && !ty->IsUnknownNdim()) {
+ keys.push_back("ndim");
+ values.push_back(LiteralDoc::Int(ty->ndim, std::nullopt));
+ }
+ if (include_vdevice && ty->vdevice.has_value()) {
+ keys.push_back("vdevice");
+ if (auto selector = GlobalInfoSelector(d, ty->vdevice.value())) {
+ values.push_back(LiteralDoc::Str(selector.value(), std::nullopt));
+ } else {
+ values.push_back(AnyValue(d, ty->vdevice.value()));
+ }
+ }
+ if (args.empty() && keys.empty()) return NamespaceDoc("relax")->Attr("Tensor");
+ return NamespaceDoc("relax")->Attr("Tensor")->Call(args, keys, values);
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::TensorTypeNode>(
+ input);
+ return RelaxTensorTypeDoc(d, ty, true);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::FuncType>( //
- "", [](relax::FuncType n, AccessPath n_p, IRDocsifier d) -> Doc {
- auto ret_doc = d->AsDoc<ExprDoc>(n->ret, n_p->Attr("ret"));
- auto purity_doc = LiteralDoc::Boolean(n->purity, n_p->Attr("purity"));
-
- if (n->IsOpaque()) {
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc, void> values;
-
- if (!n->ret->IsInstance<AnyTypeNode>()) {
- keys.push_back("ret");
- values.push_back(ret_doc);
- }
- if (n->purity) {
- keys.push_back("purity");
- values.push_back(purity_doc);
- }
-
- if (keys.size()) {
- return Relax(d, "Callable")->Call({}, keys, values);
- } else {
- return Relax(d, "Callable");
- }
- }
- // TODO(@junrushao): track symbolic shape relation
- ffi::Array<ExprDoc> params_doc;
- ffi::Array<tvm::Type> params = n->params.value();
- AccessPath params_p = n_p->Attr("params");
- for (int i = 0, n_params = params.size(); i < n_params; ++i) {
- params_doc.push_back(d->AsDoc<ExprDoc>(params[i], params_p->ArrayItem(i)));
- }
- return Relax(d, "Callable")->Call({TupleDoc(params_doc), ret_doc, purity_doc});
- });
+ ffi::reflection::TypeAttrDef<relax::TensorTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TensorTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> RelaxFuncTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::FuncTypeNode>(input);
+ if (!ty->params.has_value()) {
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (!ty->ret.as<AnyTypeNode>()) {
+ keys.push_back("ret");
+ values.push_back(d->Translate(ty->ret).value());
+ }
+ if (ty->purity) {
+ keys.push_back("purity");
+ values.push_back(LiteralDoc::Boolean(true, std::nullopt));
+ }
+ if (ty->derive_func.has_value()) {
+ keys.push_back("derive_func");
+ values.push_back(LiteralDoc::Str(ty->derive_func.value()->name, std::nullopt));
+ }
+ return keys.empty() ? NamespaceDoc("relax")->Attr("Callable")
+ : NamespaceDoc("relax")->Attr("Callable")->Call({}, keys, values);
+ }
+ ffi::Array<ExprDoc> params;
+ for (const Type& param : ty->params.value()) params.push_back(d->Translate(param).value());
+ return NamespaceDoc("relax")
+ ->Attr("Callable")
+ ->Call({TupleDoc(params), d->Translate(ty->ret).value(),
+ LiteralDoc::Boolean(ty->purity, std::nullopt)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<AnyTypeNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::ShapeTypeNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::TensorTypeNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::FuncTypeNode>(ReprPrintRelax);
+ ffi::reflection::TypeAttrDef<relax::FuncTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&RelaxFuncTypeDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/distributed.cc b/src/relax/script/printer/distributed.cc
index 5cac35ef22..f76a3c2626 100644
--- a/src/relax/script/printer/distributed.cc
+++ b/src/relax/script/printer/distributed.cc
@@ -16,132 +16,117 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/cast.h>
-#include <tvm/ir/expr.h>
#include <tvm/relax/distributed/type.h>
+#include <optional>
+
#include "../../../script/printer/ir/utils.h"
-#include "./utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-// distributed::Placement
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::distributed::Placement>(
- "", [](relax::distributed::Placement n, AccessPath n_p, IRDocsifier d) -> Doc {
- return d->AsDoc<Doc>(n->ToString(), n_p);
- });
+namespace {
+
+ffi::Optional<ExprDoc> PlacementDocTranslate(DocTranslatorObj*, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* placement = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<
+ const relax::distributed::PlacementNode>(input);
+ return LiteralDoc::Str(placement->ToString(), std::nullopt);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::distributed::DTensorType>(
- "", [](relax::distributed::DTensorType n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> args;
- ffi::Array<ffi::String> kwargs_keys;
- ffi::Array<ExprDoc> kwargs_values;
- bool require_kwargs = false;
- if (n->tensor_ty->shape.has_value()) {
- // Need to dig into ShapeExpr to preserve the `R.shape` prefix
- if (const auto* shape = n->tensor_ty->shape.value().as<relax::ShapeExprNode>()) {
- auto shape_expr = ffi::GetRef<relax::ShapeExpr>(shape);
- AccessPath shape_p = n_p->Attr("shape")->Attr("values");
- ffi::Array<ExprDoc> shape_docs;
- for (int i = 0, ndim = shape_expr->values.size(); i < ndim; ++i) {
- shape_docs.push_back(d->AsDoc<ExprDoc>(shape_expr->values[i], shape_p->ArrayItem(i)));
- }
- args.push_back(TupleDoc(shape_docs));
- } else {
- args.push_back(d->AsDoc<ExprDoc>(n->tensor_ty->shape.value(), n_p->Attr("shape")));
- }
- } else {
- require_kwargs = true;
- }
- if (!n->tensor_ty->IsUnknownDtype()) {
- if (!require_kwargs) {
- args.push_back(
- LiteralDoc::DataType(n->tensor_ty->dtype.value()->dtype, n_p->Attr("dtype")));
- } else {
- kwargs_keys.push_back("dtype");
- kwargs_values.push_back(
- LiteralDoc::DataType(n->tensor_ty->dtype.value()->dtype, n_p->Attr("dtype")));
- }
- } else {
- require_kwargs = true;
- }
- if (!require_kwargs) {
- args.push_back(d->AsDoc<ExprDoc>(n->device_mesh, n_p->Attr("device_mesh")));
- } else {
- kwargs_keys.push_back("device_mesh");
- kwargs_values.push_back(d->AsDoc<ExprDoc>(n->device_mesh, n_p->Attr("device_mesh")));
- }
- if (!require_kwargs) {
- args.push_back(d->AsDoc<ExprDoc>(n->placement, n_p->Attr("placement")));
- } else {
- kwargs_keys.push_back("placement");
- kwargs_values.push_back(d->AsDoc<ExprDoc>(n->placement, n_p->Attr("placement")));
- }
- if (!n->tensor_ty->shape.has_value() && !n->tensor_ty->IsUnknownNdim()) {
- kwargs_keys.push_back("ndim");
- kwargs_values.push_back(LiteralDoc::Int(n->tensor_ty->ndim, n_p->Attr("ndim")));
- }
- return Relax(d, "DTensor")->Call(args, kwargs_keys, kwargs_values);
- });
+ ffi::reflection::TypeAttrDef<relax::distributed::PlacementNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&PlacementDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> DTensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<
+ const relax::distributed::DTensorTypeNode>(input);
+ ffi::Array<ExprDoc> args;
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ const relax::TensorType& tensor = ty->tensor_ty;
+ bool require_keywords = !tensor->shape.has_value();
+ if (tensor->shape.has_value()) {
+ if (auto shape = tensor->shape.value().as<relax::ShapeExprNode>()) {
+ ffi::Array<ExprDoc> dimensions;
+ for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim));
+ args.push_back(TupleDoc(dimensions));
+ } else {
+ args.push_back(d->Translate(tensor->shape.value()).value());
+ }
+ }
+ if (!tensor->IsUnknownDtype()) {
+ ExprDoc dtype = LiteralDoc::DataType(tensor->dtype.value()->dtype, std::nullopt);
+ if (require_keywords) {
+ keys.push_back("dtype");
+ values.push_back(dtype);
+ } else {
+ args.push_back(dtype);
+ }
+ } else {
+ require_keywords = true;
+ }
+ ExprDoc mesh = d->Translate(ty->device_mesh).value();
+ if (auto selector = GlobalInfoSelector(d, ty->device_mesh)) {
+ mesh = LiteralDoc::Str(selector.value(), std::nullopt);
+ d->ExchangeExtraState("script.future_annotations", true);
+ }
+ ExprDoc placement = d->Translate(ty->placement).value();
+ if (require_keywords) {
+ keys.push_back("device_mesh");
+ values.push_back(mesh);
+ keys.push_back("placement");
+ values.push_back(placement);
+ } else {
+ args.push_back(mesh);
+ args.push_back(placement);
+ }
+ if (!tensor->shape.has_value() && !tensor->IsUnknownNdim()) {
+ keys.push_back("ndim");
+ values.push_back(LiteralDoc::Int(tensor->ndim, std::nullopt));
+ }
+ return NamespaceDoc("relax")->Attr("DTensor")->Call(args, keys, values);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::distributed::DeviceMesh>(
- "", [](relax::distributed::DeviceMesh n, AccessPath n_p, IRDocsifier d) -> Doc {
- bool has_relax_frame = false;
- const IRFrameNode* f = nullptr;
- for (const Frame& frame : d->frames) {
- if (frame.as<RelaxFrameNode>()) {
- has_relax_frame = true;
- break;
- } else if (const auto* ir_frame = frame.as<IRFrameNode>()) {
- f = ir_frame;
- }
- }
- if (!has_relax_frame || !f) {
- ffi::Array<ExprDoc> args;
- args.push_back(d->AsDoc<ExprDoc>(n->shape, n_p->Attr("shape")));
- if (n->device_range.has_value()) {
- args.push_back(d->AsDoc<ExprDoc>(n->device_range, n_p->Attr("device_range")));
- } else {
- args.push_back(d->AsDoc<ExprDoc>(n->device_ids, n_p->Attr("device_ids")));
- }
- return Relax(d, "device_mesh")->Call(args);
- } else {
- for (const auto& kv : *f->global_infos) {
- for (int i = 0; i < static_cast<int>(kv.second.size()); i++) {
- if (kv.second[i].same_as(n)) {
- // Module-owned selectors must be evaluated inside the declaration frame.
- for (const Frame& frame : d->frames) {
- if (const auto* relax_frame = frame.as<RelaxFrameNode>()) {
- if (relax_frame->func_vars != nullptr) {
- d->ir_usage.insert("future_annotations");
- break;
- }
- }
- }
- std::stringstream ss;
- ss << kv.first << "[" << i << "]";
- return d->AsDoc<Doc>(ffi::String(ss.str()), n_p);
- }
- }
- }
- TVM_FFI_THROW(InternalError) << "Cannot find device mesh in global infos";
- TVM_FFI_UNREACHABLE();
- }
- });
+ ffi::reflection::TypeAttrDef<relax::distributed::DTensorTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DTensorTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> DeviceMeshDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* mesh = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<
+ const relax::distributed::DeviceMeshNode>(input);
+ ffi::Array<ExprDoc> dimensions;
+ for (int64_t value : mesh->shape) {
+ dimensions.push_back(LiteralDoc::Int(value, std::nullopt));
+ }
+ ExprDoc devices = LiteralDoc::None(std::nullopt);
+ if (mesh->device_range.has_value()) {
+ CallDoc range = d->Translate(mesh->device_range.value()).value().as_or_throw<CallDoc>();
+ range->callee = NamespaceDoc("relax")->Attr("Range");
+ devices = range;
+ } else {
+ ffi::Array<ExprDoc> ids;
+ for (int64_t value : mesh->device_ids) ids.push_back(LiteralDoc::Int(value, std::nullopt));
+ devices = ListDoc(ids);
+ }
+ return NamespaceDoc("relax")->Attr("device_mesh")->Call({TupleDoc(dimensions), devices});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::distributed::DeviceMeshNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::distributed::PlacementNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::distributed::DTensorTypeNode>(ReprPrintRelax);
+ ffi::reflection::TypeAttrDef<relax::distributed::DeviceMeshNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DeviceMeshDocTranslate>());
}
+
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index c03be24742..abd5180f4d 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -17,172 +17,73 @@
* under the License.
*/
-#include <tvm/relax/distributed/type.h>
+#include <optional>
-#include <cmath>
-#include <limits>
-
-#include "../../../script/printer/dialect_prefix.h"
+#include "../../../script/printer/ir/utils.h"
#include "../../../tirx/script/printer/utils.h"
-#include "./utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
-
-TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("relax.prefix", "R"); }
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<StringImm>(
- "relax", [](StringImm n, AccessPath n_p, IRDocsifier d) -> Doc {
- return Relax(d, "str")->Call({LiteralDoc::Str(n->value, n_p->Attr("value"))});
- });
+namespace details {
+
+namespace {
+
+ffi::Optional<ExprDoc> TupleDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* tuple =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleNode>(input);
+ if (tuple->fields.empty()) return NamespaceDoc("relax")->Attr("tuple")->Call({});
+ ffi::Array<ExprDoc> fields;
+ for (const Expr& field : tuple->fields) fields.push_back(d->Translate(field).value());
+ return TupleDoc(fields);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::Tuple>( //
- "", [](relax::Tuple n, AccessPath n_p, IRDocsifier d) -> Doc {
- // TODO(@junrushao): revisit tuple printing
- if (n->fields.empty()) {
- return Relax(d, "tuple")->Call({});
- }
- ffi::Array<ExprDoc> fields_doc;
- AccessPath fields_p = n_p->Attr("fields");
- for (int i = 0, l = n->fields.size(); i < l; ++i) {
- fields_doc.push_back(d->AsDoc<ExprDoc>(n->fields[i], fields_p->ArrayItem(i)));
- }
- return TupleDoc(fields_doc);
- });
+ ffi::reflection::TypeAttrDef<TupleNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&TupleDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::TupleGetItem>( //
- "", [](relax::TupleGetItem n, AccessPath n_p, IRDocsifier d) -> Doc {
- ExprDoc idx = LiteralDoc::Int(n->index, n_p->Attr("index"));
- return d->AsDoc<ExprDoc>(n->tuple, n_p->Attr("tuple"))[{idx}];
- });
+ffi::Optional<ExprDoc> TupleGetItemDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* item =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleGetItemNode>(input);
+ return d->Translate(item->tuple).value()[{LiteralDoc::Int(item->index, std::nullopt)}];
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::ShapeExpr>( //
- "", [](relax::ShapeExpr n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> values_doc;
- AccessPath values_p = n_p->Attr("values");
- for (int i = 0, l = n->values.size(); i < l; ++i) {
- values_doc.push_back(d->AsDoc<ExprDoc>(n->values[i], values_p->ArrayItem(i)));
- }
- return Relax(d, "shape")->Call({ListDoc(values_doc)});
- });
+ ffi::reflection::TypeAttrDef<TupleGetItemNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TupleGetItemDocTranslate>());
}
-ffi::Optional<ExprDoc> SpecialScalar(const runtime::Tensor& n, const AccessPath& p) {
- DLDataType dtype = n.DataType();
- const void* data = n->data;
- if (n->ndim != 0 || n->device.device_type != kDLCPU) {
- return std::nullopt;
- }
-
- if (dtype == DLDataType{kDLInt, 8, 1}) {
- return LiteralDoc::Int(*reinterpret_cast<const int8_t*>(data), p);
- } else if (dtype == DLDataType{kDLInt, 16, 1}) {
- return LiteralDoc::Int(*reinterpret_cast<const int16_t*>(data), p);
- } else if (dtype == DLDataType{kDLInt, 32, 1}) {
- return LiteralDoc::Int(*reinterpret_cast<const int32_t*>(data), p);
- } else if (dtype == DLDataType{kDLInt, 64, 1}) {
- return LiteralDoc::Int(*reinterpret_cast<const int64_t*>(data), p);
- } else if (dtype == DLDataType{kDLFloat, 16, 1}) {
- // From IEEE-754 float16 definition
- //
- // Ref: https://en.wikipedia.org/wiki/Half-precision_floating-point_format
- uint16_t bits = *reinterpret_cast<const uint16_t*>(data);
- uint16_t sign_bit = (bits & 0b1000'0000'0000'0000) >> 15;
- uint16_t exponent = (bits & 0b0111'1100'0000'0000) >> 10;
- uint16_t fraction = (bits & 0b0000'0011'1111'1111) >> 0;
-
- double value;
- if (exponent == 0b1'1111 && fraction == 0) {
- value = std::numeric_limits<double>::infinity();
- } else if (exponent == 0b1'1111) {
- value = std::numeric_limits<double>::quiet_NaN();
- } else if (exponent == 0 && fraction == 0) {
- value = 0.0;
- } else if (exponent == 0) {
- value = std::pow(2.0, -24) * static_cast<double>(fraction);
- } else {
- value = std::pow(2.0, static_cast<double>(exponent) - 25) *
- static_cast<double>(fraction | (1 << 10));
- }
- if (sign_bit) {
- value *= -1.0;
- }
-
- return LiteralDoc::Float(value, p);
- } else if (dtype == DLDataType{kDLFloat, 32, 1}) {
- return LiteralDoc::Float(*reinterpret_cast<const float*>(data), p);
- } else if (dtype == DLDataType{kDLFloat, 64, 1}) {
- return LiteralDoc::Float(*reinterpret_cast<const double*>(data), p);
- } else if (dtype == DLDataType{kDLBool, 8, 1}) {
- return LiteralDoc::Boolean(*reinterpret_cast<const uint8_t*>(data), p);
- } else {
- return std::nullopt;
- }
+ffi::Optional<ExprDoc> ShapeExprDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* shape =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::ShapeExprNode>(input);
+ ffi::Array<ExprDoc> dimensions;
+ for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim));
+ return NamespaceDoc("relax")->Attr("shape")->Call({ListDoc(dimensions)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<::tvm::GenericConst>( //
- "relax", [](::tvm::GenericConst n, AccessPath n_p, IRDocsifier d) -> Doc {
- if (n->value.as<DLDataType>()) {
- return IRDocsifier::vtable()("", n, n_p, d);
- }
- auto data = n->value.cast<runtime::Tensor>();
- if (ffi::Optional<ExprDoc> s = SpecialScalar(data, n_p->Attr("value"))) {
- if (n->ty.as<relax::distributed::DTensorTypeNode>()) {
- ExprDoc ann = d->AsDoc<ExprDoc>(n->ty, n_p->Attr("ty"));
- return Relax(d, "dist.const")->Call({s.value(), ann});
- }
- return Relax(d, "const")
- ->Call({
- s.value(),
- LiteralDoc::DataType(data.DataType(), n_p->Attr("value")->Attr("dtype")),
- });
- }
- return d->AddMetadata(n);
- });
+ ffi::reflection::TypeAttrDef<relax::ShapeExprNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ShapeExprDocTranslate>());
}
-Doc PrintRelaxVar(tvm::Var n, AccessPath p, IRDocsifier d) {
- if (!d->IsVarDefined(n)) {
- ExprDoc ann = d->AsDoc<ExprDoc>(n->ty, p->Attr("ty"));
- Frame f = d->frames.back();
- ExprDoc var = DefineRelaxVar(n, f, d);
- f->stmts.push_back(AssignDoc(var, std::nullopt, ann));
- }
- return d->GetVarDoc(n).value();
+ffi::Optional<ExprDoc> DataflowVarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ return VarDocTranslate(d, input, destination);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::DataflowVar>("relax", PrintRelaxVar);
-}
-
-std::string ReprPrintVar(const ffi::ObjectRef& obj, const PrinterConfig& cfg) {
- Var var = obj.as_or_throw<Var>();
- if (var->ty.as<PrimTypeNode>() || var->ty.as<PointerTypeNode>() ||
- var->ty.as<tirx::BufferTypeNode>()) {
- return ReprPrintTIR(obj, cfg);
- }
- return ReprPrintRelax(obj, cfg);
+ ffi::reflection::TypeAttrDef<relax::DataflowVarNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DataflowVarDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::TupleNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::TupleGetItemNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::ShapeExprNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<VarNode>(ReprPrintVar);
- TVMScriptPrinter::Register<relax::DataflowVarNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<::tvm::GenericConstNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<::tvm::DataTypeImmNode>(ReprPrintRelax);
-}
+} // namespace
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc
index 53037d5d75..8913af7eca 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -16,146 +16,141 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ir/function.h>
+#include <tvm/tirx/type.h>
+
#include <algorithm>
+#include <optional>
+#include <utility>
+#include <vector>
-#include "./utils.h"
+#include "../../../script/printer/dialect_prefix.h"
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-static bool HasDefaultExternFuncType(const relax::ExternFunc& n) {
- const auto* ty = n->ty.as<relax::FuncTypeNode>();
- if (ty == nullptr || ty->params.has_value() || ty->purity ||
- !ty->ret->IsInstance<AnyTypeNode>()) {
- return false;
- }
- return true;
-}
+namespace {
+
+ffi::Optional<ExprDoc> FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* func =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::FunctionNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ VarScope vars(d);
-bool AtTopLevelFunction(const IRDocsifier& d) {
- // fewer than 2 frames: not in a function at all
- if (d->frames.size() < 2) {
- return false;
+ ffi::String name = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).value_or("main");
+ ffi::Array<AssignDoc> args;
+ ffi::Array<IdDoc> param_ids;
+ // A prior annotation may depend on a later scalar parameter.
+ for (const Var& param : func->params) {
+ param_ids.push_back(VarDoc(d, param));
+ }
+ size_t param_index = 0;
+ for (const Var& param : func->params) {
+ IdDoc lhs = param_ids[param_index++];
+ ffi::Optional<ExprDoc> annotation = std::nullopt;
+ if (!param->ty.IsMissing()) annotation = d->Translate(param->ty).value();
+ AssignDoc argument(lhs, std::nullopt, annotation);
+ d->RecordOrigin(argument, param);
+ args.push_back(argument);
+ }
+ auto signature_candidates = CopyImplicitDefs(d);
+ ffi::Optional<ExprDoc> ret_type = std::nullopt;
+ if (!func->ret_ty.IsMissing()) {
+ if (auto tensor = func->ret_ty.as<relax::TensorTypeNode>()) {
+ ret_type = RelaxTensorTypeDoc(d, tensor, true);
+ } else {
+ ret_type = d->Translate(func->ret_ty).value();
+ }
+ }
+ ffi::Array<ffi::String> decorator_keys;
+ ffi::Array<ExprDoc> decorator_values;
+ if (!func->is_pure) {
+ decorator_keys.push_back("pure");
+ decorator_values.push_back(LiteralDoc::Boolean(false, std::nullopt));
}
- // if the first frame is a RelaxFrame, then this is not inside a module.
- // 2 frames => we are at a function (more than 2 => nested function)
- if (d->frames[0]->IsInstance<RelaxFrameNode>()) {
- return d->frames.size() == 2;
+ if (!func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
+ decorator_keys.push_back("private");
+ decorator_values.push_back(LiteralDoc::Boolean(true, std::nullopt));
}
- // otherwise the first two frames pertain to an IR module,
- // so 3 frames => we are at a top-level function (more than 3 => nested function)
- return d->frames.size() == 3;
+ ExprDoc decorator = NamespaceDoc("relax")->Attr("function");
+ if (!decorator_keys.empty()) decorator = decorator->Call({}, decorator_keys, decorator_values);
+ auto docs = d->WithDocScope([&]() {
+ std::vector<std::pair<ffi::String, ffi::Any>> attrs;
+ for (const auto& [key, value] : func->attrs->dict) {
+ if (key != tvm::attr::kGlobalSymbol) attrs.emplace_back(key, value);
+ }
+ if (!attrs.empty()) {
+ std::sort(attrs.begin(), attrs.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ ffi::Array<ExprDoc> keys;
+ ffi::Array<ExprDoc> values;
+ for (const auto& [key, value] : attrs) {
+ keys.push_back(LiteralDoc::Str(key, std::nullopt));
+ values.push_back(AnyValue(d, value));
+ }
+ d->Emit(ExprStmtDoc(NamespaceDoc("relax")->Attr("func_attr")->Call({DictDoc(keys, values)})),
+ ffi::GetRef<ffi::ObjectRef>(func));
+ }
+ for (const relax::BindingBlock& block : func->body->blocks) d->Translate(block);
+ ExprDoc result = d->Translate(func->body->body).value();
+ if (auto integer = func->body->body.as<IntImmNode>();
+ integer && ffi::StructuralEqual()(integer->ty, func->ret_ty)) {
+ // The declared primitive return type reconstructs this literal's dtype.
+ result = LiteralDoc::Int(ffi::GetRef<IntImm>(integer), std::nullopt);
+ }
+ d->Emit(ReturnDoc(result), ffi::GetRef<ffi::ObjectRef>(func));
+ });
+ auto body = ToStmtDocArray(docs);
+ FunctionDoc function(IdDoc(name), args, {decorator}, ret_type, body);
+ FinalizeFunctionDefinitions(d, signature_candidates, function);
+ vars.Close();
+ d->Emit(function, ffi::GetRef<ffi::ObjectRef>(func));
+ return std::nullopt;
}
-TVM_FFI_STATIC_INIT_BLOCK() { RelaxFrameNode::RegisterReflection(); }
-
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::Function>(
- "", [](relax::Function n, AccessPath n_p, IRDocsifier d) -> Doc {
- std::unordered_set<const VarNode*> func_vars;
- std::unordered_set<const VarNode*> type_vars;
- std::unordered_set<const VarNode*> prim_params;
- With<RelaxFrame> f(d);
-
- IdDoc func_name("");
- // if we are binding a local definition, then calling d->Define
- // will result in a repeated definition and an incorrect displayed name
- if (ffi::Optional<ffi::String> name = GetBindingName(d)) {
- func_name = IdDoc(name.value());
- } else {
- func_name = IdDoc(FindFunctionName(d, n).value_or("main"));
- }
- (*f)->AddDispatchToken(d, "relax");
- (*f)->is_func = true;
- (*f)->func_vars = &func_vars;
- (*f)->type_vars = &type_vars;
- (*f)->prim_params = &prim_params;
- for (const Var& param : n->params) {
- if (param->ty.as<PrimTypeNode>()) {
- prim_params.insert(param.get());
- }
- }
- if (!prim_params.empty()) {
- d->ir_usage.insert("future_annotations");
- }
- // Step 1. Print params
- ffi::Array<AssignDoc> params;
- {
- AccessPath params_p = n_p->Attr("params");
- for (int i = 0, l = n->params.size(); i < l; ++i) {
- params.push_back(AssignDoc(
- /*lhs=*/DefineRelaxVar(n->params[i], *f, d),
- /*rhs=*/std::nullopt,
- TypeAsAnn(n->params[i], params_p->ArrayItem(i), d, std::nullopt)));
- }
- }
- // Step 2. Print the return type
- ffi::Optional<ExprDoc> ret_type = d->AsDoc<ExprDoc>(n->ret_ty, n_p->Attr("ret_ty"));
- // Step 3. Clean up func variables
- (*f)->func_vars = nullptr;
- (*f)->type_vars = nullptr;
- (*f)->prim_params = nullptr;
- // Step 4. Print attributes
- ffi::Map<ffi::String, Any> printable_attrs;
- for (const auto& [key, value] : n->attrs->dict) {
- // A matching global symbol is implicit for a top-level function.
- if (key == tvm::attr::kGlobalSymbol && AtTopLevelFunction(d) &&
- value.as_or_throw<ffi::String>() == func_name->name) {
- continue;
- }
- printable_attrs.Set(key, value);
- }
- if (!printable_attrs.empty()) {
- (*f)->stmts.push_back(ExprStmtDoc(
- Relax(d, "func_attr") //
- ->Call({d->AsDoc<ExprDoc>(DictAttrs(printable_attrs), n_p->Attr("attrs"))})));
- }
- // Step 5. Prepare the decorator (include purity if it's impure)
- ExprDoc decorator = Relax(d, "function");
- ffi::Array<ExprDoc, void> pos_args = {};
- ffi::Array<ffi::String, void> dec_keys;
- ffi::Array<ExprDoc, void> dec_values;
- if (!n->is_pure) {
- dec_keys.push_back("pure");
- dec_values.push_back(LiteralDoc::Boolean(false, ffi::Optional<AccessPath>()));
- }
- // if the function is global or is not in a module and does not have a global symbol,
- // indicate that it's private
- if (AtTopLevelFunction(d) && !n->attrs->dict.count(tvm::attr::kGlobalSymbol)) {
- dec_keys.push_back("private");
- dec_values.push_back(LiteralDoc::Boolean(true, ffi::Optional<AccessPath>()));
- }
- if (dec_keys.size()) {
- decorator = decorator->Call(pos_args, dec_keys, dec_values);
- }
-
- // Step 6. Print body
- ffi::Array<StmtDoc> body = PrintSeqExpr(n->body, n_p->Attr("body"), d, /*use_ret=*/true);
- (*f)->stmts.insert((*f)->stmts.end(), body.begin(), body.end());
- auto type_var_docs = DefineTypeVarDocs(type_vars, d);
- return WrapFunctionDocWithTypeVars(
- d, FunctionDoc(func_name, params, {decorator}, ret_type, (*f)->stmts), type_var_docs);
- });
+ RegisterDialectPrefix("relax.prefix", "R");
+ ffi::reflection::TypeAttrDef<relax::FunctionNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&FunctionDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::ExternFunc>( //
- "", [](relax::ExternFunc n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> args;
- args.push_back(LiteralDoc::Str(n->global_symbol, n_p->Attr("global_symbol")));
- if (!HasDefaultExternFuncType(n)) {
- args.push_back(d->AsDoc<ExprDoc>(n->ty, n_p->Attr("ty")));
- }
- return Relax(d, "ExternFunc")->Call(args);
- });
+ffi::Optional<ExprDoc> ExternFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* func =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::ExternFuncNode>(
+ input);
+ if (!func->ty.as<relax::FuncTypeNode>()) {
+ // The ExternFunc constructor requires a FuncType; retain diagnostic IR in metadata.
+ return AddMetadata(d, ffi::GetRef<relax::ExternFunc>(func));
+ }
+ ffi::Array<ExprDoc> args = {LiteralDoc::Str(func->global_symbol, std::nullopt)};
+ if (!ffi::StructuralEqual()(func->ty, relax::ExternFunc(func->global_symbol)->ty)) {
+ ExprDoc type = d->Translate(func->ty).value();
+ if (auto opaque = func->ty.as<relax::FuncTypeNode>();
+ opaque && !opaque->params.has_value() && opaque->ret.as<AnyTypeNode>() && !opaque->purity &&
+ !opaque->derive_func.has_value()) {
+ type = type->Call({});
+ }
+ args.push_back(type);
+ }
+ return NamespaceDoc("relax")->Attr("ExternFunc")->Call(args);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::FunctionNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::ExternFuncNode>(ReprPrintRelax);
+ ffi::reflection::TypeAttrDef<relax::ExternFuncNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ExternFuncDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/global_info.cc b/src/relax/script/printer/global_info.cc
index e23ba3eaa2..9c06dcac71 100644
--- a/src/relax/script/printer/global_info.cc
+++ b/src/relax/script/printer/global_info.cc
@@ -16,33 +16,68 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/relax/global_info.h>
+#include <tvm/target/target.h>
-#include "./utils.h"
+#include <optional>
+#include <string>
+
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::DummyGlobalInfo>(
- "", [](GlobalInfo ginfo, AccessPath p, IRDocsifier d) -> Doc {
- return Relax(d, "dummy_global_info")->Call({});
- });
+using GlobalInfoMap = ffi::Dict<ffi::String, ffi::List<GlobalInfo>>;
+
+// Selectors are derived from the module's forward registry, never cached as
+// object-to-string state. Definitions and standalone objects keep constructors.
+ffi::Optional<ffi::String> GlobalInfoSelector(DocTranslatorObj* d, const GlobalInfo& info) {
+ auto scope = d->GetOrCreateExtraState<ffi::Optional<GlobalInfoMap>>("ir.global_info_map");
+ if (!scope.has_value()) return std::nullopt;
+ auto infos = scope.value();
+ if (auto query = info.as<relax::VDevice>()) {
+ if (auto devices = infos.Get("vdevice")) {
+ ffi::String kind = query.value()->target->kind->name;
+ size_t index = 0;
+ for (const GlobalInfo& entry : devices.value()) {
+ if (auto device = entry.as<relax::VDevice>();
+ device && device.value()->target->kind->name == kind) {
+ if (entry.same_as(info)) {
+ return ffi::String(std::string(kind) + ":" + std::to_string(index) + ":" +
+ std::string(device.value()->memory_scope));
+ }
+ ++index;
+ }
+ }
+ }
+ }
+ if (auto meshes = infos.Get("mesh")) {
+ size_t index = 0;
+ for (const GlobalInfo& entry : meshes.value()) {
+ if (entry.same_as(info)) return ffi::String("mesh[" + std::to_string(index) + "]");
+ ++index;
+ }
+ }
+ return std::nullopt;
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> DummyGlobalInfoDocTranslate(DocTranslatorObj*, ffi::AnyView,
+ const ffi::Object*) {
+ return NamespaceDoc("relax")->Attr("dummy_global_info")->Call({});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::VDevice>(
- "", [](relax::VDevice vdev, AccessPath p, IRDocsifier d) -> Doc {
- d->AddGlobalInfo("vdevice", vdev);
- ffi::Map<ffi::String, ffi::Any> config = vdev->target->ToConfig();
- return Relax(d, "vdevice")
- ->Call({d->AsDoc<ExprDoc>(config, p),
- LiteralDoc::Int(vdev->vdevice_id, p->Attr("vdevice_id")),
- LiteralDoc::Str(vdev->memory_scope, p->Attr("memory_scope"))});
- });
+ ffi::reflection::TypeAttrDef<relax::DummyGlobalInfoNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DummyGlobalInfoDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/region.cc b/src/relax/script/printer/region.cc
index 42fb43666b..95f8ab94bc 100644
--- a/src/relax/script/printer/region.cc
+++ b/src/relax/script/printer/region.cc
@@ -16,94 +16,127 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/cast.h>
-#include "./utils.h"
+#include <tvm/tirx/type.h>
+
+#include <optional>
+
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-ffi::Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const AccessPath& n_p,
- const IRDocsifier& d, bool use_ret) {
- With<RelaxFrame> f(d);
- const ffi::Array<relax::BindingBlock>& blocks = n->blocks;
- AccessPath blocks_p = n_p->Attr("blocks");
- ffi::Array<StmtDoc>* stmts = &(*f)->stmts;
- for (int i = 0, l = blocks.size(); i < l; ++i) {
- Doc block = d->AsDoc(blocks[i], blocks_p->ArrayItem(i));
- if (const auto* stmt_block = block.as<StmtBlockDocNode>()) {
- stmts->insert(stmts->end(), stmt_block->stmts.begin(), stmt_block->stmts.end());
- } else if (const auto* stmt = block.as<StmtDocNode>()) {
- stmts->push_back(ffi::GetRef<StmtDoc>(stmt));
- } else {
- TVM_FFI_THROW(TypeError) << "Unknown type: " << block->GetTypeKey();
- }
- }
- ExprDoc ret = d->AsDoc<ExprDoc>(n->body, n_p->Attr("body"));
- if (use_ret) {
- stmts->push_back(ReturnDoc(ret));
- } else {
- stmts->push_back(ExprStmtDoc(ret));
- }
- return *stmts;
-}
+namespace {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::SeqExpr>(
- "", [](relax::SeqExpr n, AccessPath n_p, IRDocsifier d) -> Doc {
- return StmtBlockDoc(PrintSeqExpr(n, n_p, d, false));
- });
+ffi::Optional<Var> BindingVar(const ffi::Object* destination) {
+ if (!destination) return std::nullopt;
+ TVM_FFI_CHECK(destination->IsInstance<VarNode>(), TypeError)
+ << "printer binding destination must be a Var";
+ return ffi::GetRef<Var>(static_cast<const VarNode*>(destination));
}
-ffi::Array<StmtDoc> PrintBindingBlock(const relax::BindingBlock& n, const AccessPath& n_p,
- const IRDocsifier& d,
- ffi::Array<ExprDoc>* non_dataflow_vars) {
- const ffi::Array<relax::Binding>& bindings = n->bindings;
- AccessPath bindings_p = n_p->Attr("bindings");
- ffi::Array<StmtDoc> stmts;
- for (int i = 0, l = bindings.size(); i < l; ++i) {
- const relax::Binding& binding = bindings[i];
- AccessPath binding_p = bindings_p->ArrayItem(i);
- TVM_FFI_ICHECK(binding->var.defined());
- Doc binding_doc = d->AsDoc(binding, binding_p);
- if (const auto* stmt = binding_doc.as<StmtDocNode>()) {
- stmts.push_back(ffi::GetRef<StmtDoc>(stmt));
- } else if (const auto* stmt_block = binding_doc.as<StmtBlockDocNode>()) {
- stmts.insert(stmts.end(), stmt_block->stmts.begin(), stmt_block->stmts.end());
+} // namespace
+
+ffi::Array<StmtDoc> RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq,
+ ffi::Optional<IdDoc> destination,
+ ffi::Optional<ExprDoc> annotation,
+ const ffi::Object* destination_object) {
+ auto docs = d->WithDocScope([&]() {
+ ffi::Optional<ExprDoc> value =
+ d->Translate(ffi::GetRef<relax::SeqExpr>(seq), BindingVar(destination_object));
+ if (destination.has_value()) {
+ if (value.has_value()) {
+ d->Emit(AssignDoc(destination.value(), value.value(), annotation),
+ ffi::GetRef<ffi::ObjectRef>(seq));
+ } else {
+ TVM_FFI_CHECK(destination_object != nullptr, ValueError)
+ << "printer Relax SeqExpr completed a binding without a destination";
+ }
} else {
- TVM_FFI_THROW(TypeError) << "Unknown type: " << binding_doc->GetTypeKey();
+ TVM_FFI_CHECK(value.has_value(), ValueError)
+ << "printer Relax SeqExpr needs a value without a destination";
+ d->Emit(ExprStmtDoc(value.value()), ffi::GetRef<ffi::ObjectRef>(seq));
}
- if (non_dataflow_vars != nullptr && !binding->var->IsInstance<relax::DataflowVarNode>()) {
- non_dataflow_vars->push_back(d->AsDoc<ExprDoc>(binding->var, binding_p->Attr("var")));
+ });
+ return ToStmtDocArray(docs);
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> SeqExprDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* seq =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::SeqExprNode>(input);
+ for (const relax::BindingBlock& block : seq->blocks) d->Translate(block);
+ if (seq->blocks.empty()) {
+ if (auto var = seq->body.as<Var>()) {
+ d->VarGetOrAllocId(var.value(), false);
+ if (d->GetImplicitDefs().count(var.value()) && !var.value()->ty.as<PrimTypeNode>()) {
+ IdDoc id = VarDoc(d, var.value());
+ d->Emit(AssignDoc(id, std::nullopt, d->Translate(var.value()->ty).value()),
+ ffi::GetRef<ffi::ObjectRef>(var.value().get()));
+ }
}
}
- return stmts;
+ return d->Translate(seq->body, BindingVar(destination));
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::BindingBlock>( //
- "", [](relax::BindingBlock n, AccessPath n_p, IRDocsifier d) -> Doc {
- return StmtBlockDoc(PrintBindingBlock(n, n_p, d, nullptr));
- });
+ ffi::reflection::TypeAttrDef<relax::SeqExprNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&SeqExprDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> BindingBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* block =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::BindingBlockNode>(
+ input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ for (const relax::Binding& binding : block->bindings) d->Translate(binding);
+ return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::DataflowBlock>( //
- "", [](relax::DataflowBlock n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> non_dataflow_vars;
- ffi::Array<StmtDoc> stmts = PrintBindingBlock(n, n_p, d, &non_dataflow_vars);
- stmts.push_back(ExprStmtDoc(Relax(d, "output")->Call(non_dataflow_vars)));
- return ScopeDoc(std::nullopt, Relax(d, "dataflow")->Call({}), stmts);
- });
+ ffi::reflection::TypeAttrDef<relax::BindingBlockNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BindingBlockDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> DataflowBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* block =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const relax::DataflowBlockNode>(
+ input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ auto docs = d->WithDocScope([&]() {
+ ffi::Array<ExprDoc> outputs;
+ for (const relax::Binding& binding : block->bindings) {
+ d->Translate(binding);
+ if (!binding->var.as<relax::DataflowVarNode>()) {
+ outputs.push_back(d->Translate(binding->var).value());
+ }
+ }
+ d->Emit(ExprStmtDoc(NamespaceDoc("relax")->Attr("output")->Call(outputs)),
+ ffi::GetRef<ffi::ObjectRef>(block));
+ });
+ d->Emit(ScopeDoc(std::nullopt, NamespaceDoc("relax")->Attr("dataflow")->Call({}),
+ ToStmtDocArray(docs)),
+ ffi::GetRef<ffi::ObjectRef>(block));
+ return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::SeqExprNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::BindingBlockNode>(ReprPrintRelax);
- TVMScriptPrinter::Register<relax::DataflowBlockNode>(ReprPrintRelax);
+ ffi::reflection::TypeAttrDef<relax::DataflowBlockNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DataflowBlockDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc
index 8ae3a49a69..d2ae3d0177 100644
--- a/src/relax/script/printer/tir.cc
+++ b/src/relax/script/printer/tir.cc
@@ -16,129 +16,36 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/cast.h>
-#include <tvm/ir/expr.h>
+#include <tvm/ir/prim/op.h>
-#include "../../../tirx/script/printer/utils.h"
-#include "./utils.h"
+#include <optional>
+
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
-
-/*! \brief Find the innermost Relax function frame, or the innermost Relax frame. */
-RelaxFrameNode* GetRelaxFrame(IRDocsifier d) {
- RelaxFrameNode* f = nullptr;
- for (auto it = d->frames.rbegin(); it != d->frames.rend(); ++it) {
- if (const auto* relax_frame = (*it).as<RelaxFrameNode>()) {
- if (relax_frame->is_func) {
- return const_cast<RelaxFrameNode*>(relax_frame);
- } else if (f == nullptr) {
- f = const_cast<RelaxFrameNode*>(relax_frame);
- }
- }
- }
- return f;
-}
-
-Doc PrintCanonicalVar(Var n, AccessPath n_p, IRDocsifier d) {
- if (!n->ty.as<PrimTypeNode>()) {
- return PrintRelaxVar(n, n_p, d);
- }
- PrimVar prim_var = n.as_or_throw<PrimVar>();
- if (!d->IsVarDefined(n)) {
- PrimType n_ty = n->ty.as_or_throw<PrimType>();
- TVM_FFI_CHECK(!n_ty.IsScalableVector() && !n_ty.IsFixedLengthVector(), TypeError)
- << "Relax only uses scalar TIR variables,"
- << "but received TIR variable " << n << " with dtype " << n_ty->dtype;
- RelaxFrameNode* f = GetRelaxFrame(d);
- // There should be at least one Relax frame
- if (f == nullptr) {
- TVM_FFI_THROW(IndexError) << "No relax environment is found when printing a TIR var under "
- "relax's dispatch token";
- }
- // If the Relax function frame is collecting func vars
- if (f->func_vars) {
- TVM_FFI_ICHECK(f->is_func);
- f->func_vars->insert(n.get());
- if (!f->prim_params->count(n.get())) {
- f->type_vars->insert(n.get());
- }
- }
- Frame frame = d->frames.front();
- IdDoc var = d->Define(n, frame, n->name.empty() ? "v" : n->name);
- var->source_paths.push_back(n_p);
- if (!f->func_vars || f->prim_params->count(n.get()) || !f->type_vars->count(n.get())) {
- frame->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d), std::nullopt));
- }
+namespace details {
+
+ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim) {
+ ExprDoc translated = d->Translate(dim).value();
+ if (auto mul = dim.as<prim::MulNode>()) {
+ ExprDoc compact = OperationDoc(OperationDocNode::Kind::kMult,
+ {RelaxShapeDim(d, mul->a), RelaxShapeDim(d, mul->b)});
+ d->RecordOrigin(compact, dim);
+ return compact;
}
- if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
- return doc.value();
+ if (auto integer = dim.as<IntImmNode>();
+ integer && integer->ty.as_or_throw<PrimType>()->dtype == (DLDataType{kDLInt, 64, 1})) {
+ ExprDoc compact = LiteralDoc::Int(ffi::GetRef<IntImm>(integer), std::nullopt);
+ d->RecordOrigin(compact, dim);
+ return compact;
}
- TVM_FFI_THROW(IndexError) << "Variable is not defined in the environment: " << n;
- TVM_FFI_UNREACHABLE();
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() { IRDocsifier::vtable().set_dispatch<Var>("relax", PrintCanonicalVar); }
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tvm::IntImm>( //
- "relax", [](tvm::IntImm n, AccessPath n_p, IRDocsifier d) -> Doc { //
- // TODO(@junrushao): support non-int64 cases
- if (n->ty.as_or_throw<PrimType>().MatchesElementType(DLDataTypeCode::kDLBool, 8)) {
- return LiteralDoc::Boolean(static_cast<bool>(n->value), n_p);
- } else {
- return LiteralDoc::Int(n, n_p);
- }
- });
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tvm::GlobalVar>( //
- "relax", [](tvm::GlobalVar n, AccessPath n_p, IRDocsifier d) -> Doc { //
- if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
- return doc.value();
- } else {
- IdDoc ret(n->name_hint);
- ret->source_paths.push_back(n_p);
- return ret;
- }
- });
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tvm::IRModule>( //
- "relax", [](tvm::IRModule mod, AccessPath n_p, IRDocsifier d) -> Doc { //
- ffi::Optional<ExprDoc> doc = d->GetVarDoc(mod);
- TVM_FFI_ICHECK(doc) << "Unable to print IRModule before definition in Relax.";
- if (d->cfg->module_alias.empty()) {
- // Use Module Name directly
- return doc.value();
- }
- RelaxFrameNode* f = GetRelaxFrame(d);
- TVM_FFI_CHECK(f != nullptr && f->is_func, IndexError)
- << "No relax environment is found when printing a module alias var "
- "under relax's dispatch token";
- if (!f->module_alias_printed) {
- // If the module_alias is not defined before, define it.
- f->stmts.push_back(AssignDoc(IdDoc(d->cfg->module_alias), doc.value(), std::nullopt));
- f->module_alias_printed = true;
- }
- return IdDoc(d->cfg->module_alias);
- });
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Range>(
- "relax", [](Range range, AccessPath p, IRDocsifier d) -> Doc {
- return Relax(d, "Range")
- ->Call({
- d->AsDoc<ExprDoc>(range->min, p->Attr("min")),
- d->AsDoc<ExprDoc>(range->extent + range->min, p->Attr("extent")),
- });
- });
+ return translated;
}
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/type.cc b/src/relax/script/printer/type.cc
index 65408d855f..50d6024be9 100644
--- a/src/relax/script/printer/type.cc
+++ b/src/relax/script/printer/type.cc
@@ -16,59 +16,30 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/reflection/registry.h>
-#include "./utils.h"
+#include "../../../script/printer/ir/utils.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<relax::PackedFuncType>( //
- "", [](relax::PackedFuncType n, AccessPath n_p, IRDocsifier d) -> Doc {
- return Relax(d, "PackedFunc"); // TODO(@junrushao): verify if this is correct
- });
-}
+namespace {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tvm::TupleType>( //
- "relax", [](tvm::TupleType n, AccessPath n_p, IRDocsifier d) -> Doc {
- if (n->fields.empty()) {
- return Relax(d, "Tuple");
- }
- ffi::Array<ExprDoc> fields_doc;
- AccessPath fields_p = n_p->Attr("fields");
- for (int i = 0, l = n->fields.size(); i < l; ++i) {
- fields_doc.push_back(d->AsDoc<ExprDoc>(n->fields[i], fields_p->ArrayItem(i)));
- }
- return Relax(d, "Tuple")->Call(fields_doc);
- });
+ffi::Optional<ExprDoc> PackedFuncTypeDocTranslate(DocTranslatorObj*, ffi::AnyView,
+ const ffi::Object*) {
+ return NamespaceDoc("relax")->Attr("PackedFunc");
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<tvm::FuncType>(
- "relax", [](tvm::FuncType n, AccessPath n_p, IRDocsifier d) -> Doc {
- ffi::Array<ExprDoc> arg_types_doc;
- ffi::Array<Type> arg_types = n->arg_types;
- AccessPath arg_types_p = n_p->Attr("arg_types");
- for (int i = 0, n_params = arg_types.size(); i < n_params; ++i) {
- arg_types_doc.push_back(d->AsDoc<ExprDoc>(arg_types[i], arg_types_p->ArrayItem(i)));
- }
- return Relax(d, "Callable")
- ->Call({TupleDoc(arg_types_doc), //
- d->AsDoc<ExprDoc>(n->ret_type, n_p->Attr("ret_type"))});
- });
+ ffi::reflection::TypeAttrDef<relax::PackedFuncTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&PackedFuncTypeDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<relax::PackedFuncTypeNode>(ReprPrintRelax);
-}
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("script.printer.ReprPrintRelax", ReprPrintRelax);
-}
+} // namespace
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index 9b42a43dc8..7cfba8b01a 100644
--- a/src/relax/script/printer/utils.h
+++ b/src/relax/script/printer/utils.h
@@ -16,155 +16,34 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_RELAX_SCRIPT_PRINTER_UTILS_H_
-#define TVM_RELAX_SCRIPT_PRINTER_UTILS_H_
+#ifndef SRC_RELAX_SCRIPT_PRINTER_UTILS_H_
+#define SRC_RELAX_SCRIPT_PRINTER_UTILS_H_
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/relax/analysis.h>
-#include <tvm/relax/op_attr_types.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/global_info.h>
#include <tvm/relax/type.h>
-#include <tvm/relax/utils.h>
-#include <tvm/script/printer/ir_docsifier.h>
-#include <string>
-#include <unordered_set>
-#include <utility>
-#include <vector>
+#include <optional>
-#include "../../../script/printer/utils.h"
+#include "../../../script/printer/ir/utils.h"
namespace tvm {
namespace script {
namespace printer {
-
-class RelaxFrameNode : public FrameNode {
- public:
- bool is_func = false;
- bool module_alias_printed = false;
- std::unordered_set<const VarNode*>* func_vars = nullptr;
- std::unordered_set<const VarNode*>* type_vars = nullptr;
- std::unordered_set<const VarNode*>* prim_params = nullptr;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<RelaxFrameNode>()
- .def_ro("is_func", &RelaxFrameNode::is_func)
- .def_ro("module_alias_printed", &RelaxFrameNode::module_alias_printed);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.RelaxFrame", RelaxFrameNode, FrameNode);
-};
-
-class RelaxFrame : public Frame {
- public:
- explicit RelaxFrame(const IRDocsifier& d) {
- ffi::ObjectPtr<RelaxFrameNode> n = ffi::make_object<RelaxFrameNode>();
- n->stmts.clear();
- n->d = d.get();
- n->is_func = false;
- n->func_vars = nullptr;
- n->type_vars = nullptr;
- n->prim_params = nullptr;
- data_ = std::move(n);
- }
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RelaxFrame, Frame, RelaxFrameNode);
-};
-
-/*! \brief Redirected method for the ffi repr hook */
-inline std::string ReprPrintRelax(const ffi::ObjectRef& obj, const PrinterConfig& cfg) {
- IRDocsifier d(cfg);
- With<RelaxFrame> f(d);
- (*f)->AddDispatchToken(d, "relax");
- return Docsify(obj, d, *f, cfg);
-}
-
-inline IdDoc DefineRelaxVar(const tvm::Var& var, const Frame& frame, const IRDocsifier& d) {
- return d->Define(var, frame, var->name.empty() ? "v" : var->name);
-}
-
-inline ffi::Optional<ExprDoc> TypeAsAnn(const tvm::Var& v, const AccessPath& v_p,
- const IRDocsifier& d,
- const ffi::Optional<relax::Expr>& rhs) {
- if (v->ty.IsMissing()) {
- return std::nullopt;
- }
- bool attempt_to_hide_ty = !d->cfg->GetExtraConfig<bool>("relax.show_all_ty", true);
-
- if (rhs.has_value()) {
- if (const auto* call = rhs.as<tvm::CallNode>()) {
- static const Op call_tir_op = Op::Get("relax.call_tir");
- static const Op call_dps_packed_op = Op::Get("relax.call_dps_packed");
- if (call->op.same_as(call_tir_op) || call->op.same_as(call_dps_packed_op)) {
- attempt_to_hide_ty = true;
- }
- }
- }
- if (attempt_to_hide_ty && rhs.has_value()) {
- ffi::Optional<tvm::Type> inferred_ty = std::nullopt;
- if (auto opt = rhs.as<tvm::Call>()) {
- auto call = opt.value();
- if (auto opt = call->op.as<Op>()) {
- auto op = opt.value();
-
- static auto op_map_context_free = Op::GetAttrMap<FInferType>("FInferType");
- if (op_map_context_free.count(op)) {
- inferred_ty = Call::ReinferType(call.get());
- } else {
- static auto op_map_infer_ty =
- Op::GetAttrMap<relax::FInferTypeWithBuilder>("relax.FInferTypeWithBuilder");
- if (op_map_infer_ty.count(op)) {
- auto temp_builder = relax::BlockBuilder::Create(std::nullopt);
- inferred_ty = op_map_infer_ty[op](call, temp_builder);
- }
- }
- } else if (auto opt = call->op.as<relax::FuncType>()) {
- auto temp_builder = relax::BlockBuilder::Create(std::nullopt);
- inferred_ty =
- DeriveCallRetType(opt.value(), call, temp_builder, temp_builder->GetAnalyzer());
- }
-
- } else if (const auto* tuple = rhs.as<relax::TupleNode>()) {
- inferred_ty = relax::TupleType(tuple->fields.Map(relax::GetType));
-
- } else if (const auto* get_item = rhs.as<relax::TupleGetItemNode>()) {
- if (auto ptr = get_item->tuple->ty.as<relax::TupleTypeNode>();
- ptr && get_item->index < static_cast<int>(ptr->fields.size())) {
- inferred_ty = ptr->fields[get_item->index];
- }
-
- } else if (const auto* trivial_binding = rhs.as<tvm::VarNode>()) {
- inferred_ty = trivial_binding->ty.as<tvm::Type>();
- }
-
- if (inferred_ty && ffi::StructuralEqual()(inferred_ty, v->ty)) {
- return std::nullopt;
- }
- }
- return d->AsDoc<ExprDoc>(v->ty, v_p->Attr("ty"));
-}
-
-ffi::Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const AccessPath& n_p,
- const IRDocsifier& d, bool use_ret);
-
-Doc PrintRelaxVar(tvm::Var n, AccessPath p, IRDocsifier d);
-
-inline int FindVDeviceIndexByTargetKind(const relax::VDevice& vdevice, const IRDocsifier& d) {
- ffi::Array<GlobalInfo> vdevices = d->global_infos["vdevice"];
- int kind_index = 0;
- for (size_t i = 0; i < vdevices.size(); ++i) {
- auto vdev = vdevices[i].as_or_throw<relax::VDevice>();
- if (vdev.same_as(vdevice)) {
- return kind_index;
- }
- if (vdev->target->kind->name == vdevice->target->kind->name) {
- kind_index++;
- }
- }
- return -1;
-}
-
+namespace details {
+
+ffi::Optional<ffi::String> GlobalInfoSelector(DocTranslatorObj* d, const GlobalInfo& info);
+ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim);
+ExprDoc RelaxTensorTypeDoc(DocTranslatorObj* d, const relax::TensorTypeNode* ty,
+ bool include_vdevice);
+ffi::Array<StmtDoc> RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq,
+ ffi::Optional<IdDoc> destination = std::nullopt,
+ ffi::Optional<ExprDoc> annotation = std::nullopt,
+ const ffi::Object* destination_object = nullptr);
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
-#endif // TVM_RELAX_SCRIPT_PRINTER_UTILS_H_
+#endif // SRC_RELAX_SCRIPT_PRINTER_UTILS_H_
diff --git a/src/s_tir/script/ir_builder/frame.cc b/src/s_tir/script/ir_builder/frame.cc
index b6fe374568..5cac209755 100644
--- a/src/s_tir/script/ir_builder/frame.cc
+++ b/src/s_tir/script/ir_builder/frame.cc
@@ -17,7 +17,6 @@
* under the License.
*/
#include <tvm/s_tir/script/ir_builder/frame.h>
-#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/function.h>
#include "./script_complete.h"
@@ -27,22 +26,6 @@ namespace tvm {
namespace script {
namespace ir_builder {
namespace s_tir {
-namespace {
-
-// Annotations may have been evaluated before entering this frame. Normalize
-// their buffer layouts and all matching body references during finalization.
-class STirBufferLayoutNormalizer : public tvm::s_tir::StmtExprMutator {
- public:
- using tvm::s_tir::StmtExprMutator::Mutate;
- using tvm::s_tir::StmtExprMutator::Mutate_;
- void Register(const tvm::tirx::BufferVar& old_buf, const tvm::tirx::BufferVar& new_buf) {
- VarRemapSet(old_buf, new_buf);
- }
- bool Empty() const { return var_remap_.empty(); }
-};
-
-} // namespace
-
TVM_FFI_STATIC_INIT_BLOCK() {
PrimFuncFrameNode::RegisterReflection();
SBlockFrameNode::RegisterReflection();
@@ -66,39 +49,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
tvm::tirx::PrimFunc PrimFuncFrameNode::FinalizeFunction(tvm::tirx::PrimFunc func) {
TVM_FFI_CHECK(!is_declaration || root_alloc_buffers.empty(), ValueError)
<< "A function declaration cannot allocate buffers";
- auto normalizer = ffi::make_object<STirBufferLayoutNormalizer>();
- auto normalize = [&](tvm::tirx::BufferVar buffer) {
- if (buffer->layout.has_value()) {
- auto type = tvm::tirx::CopyBufferType(buffer);
- type->layout = std::nullopt;
- auto replacement = tvm::tirx::RebuildBufferVar(buffer, std::move(type));
- normalizer->Register(buffer, replacement);
- return replacement;
- }
- return buffer;
- };
- ffi::Array<tvm::tirx::Var> params;
- for (const auto& param : func->params) {
- if (auto buffer = param.as<tvm::tirx::BufferVar>()) {
- params.push_back(normalize(buffer.value()).var());
- } else {
- params.push_back(param);
- }
- }
- ffi::Array<tvm::tirx::BufferVar> alloc_buffers;
- for (const auto& buffer : root_alloc_buffers) {
- alloc_buffers.push_back(normalize(buffer));
- }
- if (!normalizer->Empty()) {
- auto* n = func.CopyOnWrite();
- n->params = params;
- if (!is_declaration) {
- n->body = normalizer->Mutate(n->body, InplaceMode::kAllow).ValueOrUnchanged(n->body);
- }
- }
if (!is_declaration) {
func = WithAttr(std::move(func), tvm::attr::kSTir, true);
- func = tvm::s_tir::ScriptComplete(std::move(func), alloc_buffers);
+ func = tvm::s_tir::ScriptComplete(std::move(func), root_alloc_buffers);
}
return func;
}
diff --git a/src/s_tir/script/ir_builder/ir.cc b/src/s_tir/script/ir_builder/ir.cc
index 592a96f8f6..a54c9251a1 100644
--- a/src/s_tir/script/ir_builder/ir.cc
+++ b/src/s_tir/script/ir_builder/ir.cc
@@ -223,11 +223,13 @@ BufferVar SBlockAllocBuffer(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::Opt
}
namespace axis {
-IterVar PushBlockVar(IterVar iter_var, PrimExpr binding) {
+IterVar PushBlockVar(IterVar iter_var, ffi::Optional<PrimExpr> binding) {
if (ffi::Optional<SBlockFrame> opt_frame = IRBuilder::Current()->GetLastFrame<SBlockFrame>()) {
SBlockFrame frame = opt_frame.value();
+ TVM_FFI_CHECK(binding.has_value() || frame->no_realize, TypeError)
+ << "Block axes require a binding unless no_realize=True";
frame->iter_vars.push_back(iter_var);
- frame->iter_values.push_back(binding);
+ if (binding.has_value()) frame->iter_values.push_back(binding.value());
} else {
TVM_FFI_THROW(InternalError) << "TypeError: The last frame is not SBlockFrame";
}
@@ -235,7 +237,7 @@ IterVar PushBlockVar(IterVar iter_var, PrimExpr binding) {
}
#define TVM_S_TIR_IR_BUILDER_AXIS(Method, Kind, Name) \
- Var Method(Range dom, PrimExpr binding, PrimType dtype) { \
+ Var Method(Range dom, ffi::Optional<PrimExpr> binding, PrimType dtype) { \
TVM_FFI_ICHECK(dom.defined()) << Name << " axis must have a domain"; \
PrimType min_ty = dom->min.ty(); \
PrimType extent_ty = dom->extent.ty(); \
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 993e4b3f44..518f93bb91 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -23,6 +23,7 @@
#include <map>
#include <string>
#include <unordered_set>
+#include <utility>
#include "./dialect_prefix.h"
@@ -58,6 +59,8 @@ void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_pr
<< "Duplicate printer dialect prefix: " << key;
}
+const std::map<ffi::String, ffi::String>& GetDialectPrefixes() { return DialectPrefixes(); }
+
} // namespace printer
} // namespace script
@@ -84,10 +87,14 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
n->buffer_dtype = ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
}
if (auto v = get("int_dtype")) {
- n->int_dtype = ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
+ n->extra_config.Set(
+ "ir.int_dtype",
+ ffi::DLDataTypeToString(ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>())));
}
if (auto v = get("float_dtype")) {
- n->float_dtype = ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
+ n->extra_config.Set(
+ "script.float_dtype",
+ ffi::DLDataTypeToString(ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>())));
}
if (auto v = get("verbose_expr")) {
n->verbose_expr = v.value().cast<bool>();
@@ -125,9 +132,6 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
if (auto v = get("syntax_sugar")) {
n->syntax_sugar = v.value().cast<bool>();
}
- if (auto v = get("show_object_address")) {
- n->show_object_address = v.value().cast<bool>();
- }
if (auto v = get("render_invisible_path_info")) {
n->render_invisible_path_info = v.value().cast<bool>();
}
@@ -147,7 +151,17 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
}
}
- // Validate all registered prefixes before names can be assigned by a docsifier.
+ n->extra_config.Set("ir.prefix", n->ir_prefix);
+
+ // Keep legacy constructor arguments as adapters; explicit extra_config values
+ // take precedence. Store validated dtype strings with the existing defaults.
+ for (const auto& [key, fallback] :
+ {std::pair{"ir.int_dtype", "int32"}, std::pair{"script.float_dtype", "void"}}) {
+ ffi::String value = n->GetExtraConfig<ffi::String>(key, fallback);
+ n->extra_config.Set(key, ffi::DLDataTypeToString(ffi::StringToDLDataType(value)));
+ }
+
+ // Validate all registered prefixes before translation assigns names.
n->GetBuiltinKeywords();
this->data_ = std::move(n);
diff --git a/src/script/printer/dialect_prefix.h b/src/script/printer/dialect_prefix.h
index 5aaae11293..0eac91b5ea 100644
--- a/src/script/printer/dialect_prefix.h
+++ b/src/script/printer/dialect_prefix.h
@@ -21,14 +21,19 @@
#include <tvm/ffi/string.h>
+#include <map>
+
namespace tvm {
namespace script {
namespace printer {
// Register during dialect static initialization so configuration can validate
-// and reserve the prefix before any docsifier assigns variable names.
+// and reserve the prefix before translation assigns variable names.
void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix);
+// Read the existing registration table without introducing a second namespace registry.
+const std::map<ffi::String, ffi::String>& GetDialectPrefixes();
+
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/doc.cc b/src/script/printer/doc.cc
index 4727532128..71360b31df 100644
--- a/src/script/printer/doc.cc
+++ b/src/script/printer/doc.cc
@@ -225,11 +225,13 @@ ForDoc::ForDoc(ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body) {
this->data_ = std::move(n);
}
-ScopeDoc::ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body) {
+ScopeDoc::ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body,
+ bool allow_concise_scoping) {
ffi::ObjectPtr<ScopeDocNode> n = ffi::make_object<ScopeDocNode>();
n->lhs = lhs;
n->rhs = rhs;
n->body = body;
+ n->allow_concise_scoping = allow_concise_scoping;
this->data_ = std::move(n);
}
@@ -522,6 +524,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+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/doc_printer/python_doc_printer.cc b/src/script/printer/doc_printer.cc
similarity index 51%
rename from src/script/printer/doc_printer/python_doc_printer.cc
rename to src/script/printer/doc_printer.cc
index d10b95bd76..302ba6d069 100644
--- a/src/script/printer/doc_printer/python_doc_printer.cc
+++ b/src/script/printer/doc_printer.cc
@@ -16,28 +16,680 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include "doc_printer.h"
+
#include <tvm/ffi/error.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/script/printer/doc.h>
+#include <tvm/ir/object_functor.h>
+#include <tvm/ir/prim/expr.h>
#include <algorithm>
#include <cmath>
+#include <limits>
#include <map>
+#include <memory>
#include <optional>
+#include <ostream>
#include <sstream>
#include <streambuf>
#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
-#include "../../../support/str_escape.h"
-#include "../../../support/utils.h"
-#include "./base_doc_printer.h"
+#include "../../support/str_escape.h"
+#include "../../support/utils.h"
+#include "dialect_prefix.h"
namespace tvm {
namespace script {
namespace printer {
namespace {
+/*! \brief Range of byte offsets in a string */
+using ByteSpan = std::pair<size_t, size_t>;
+
+/*!
+ * \brief DocPrinter is responsible for printing Doc tree into text format
+ * \details This is the base class for translating Doc into string.
+ * Each target language needs to have its subclass of DocPrinter
+ * to define the actual logic of printing Doc.
+ *
+ * \sa Doc
+ */
+class DocPrinter {
+ public:
+ /*!
+ * \brief The constructor of DocPrinter
+ *
+ * \param options the option for printer
+ */
+ explicit DocPrinter(const PrinterConfig& options);
+
+ virtual ~DocPrinter() = default;
+
+ /*! \brief Append an ordered entry-owned header with line/span accounting.
+ * Elements are fixed String source chunks or ordinary CommentDocs, in the
+ * exact order to emit. Comments use normal Doc spans and newline handling.
+ * Any other element type raises TypeError. The printer retains no header.
+ */
+ void AppendHeader(const ffi::Array<ffi::Any>& header);
+
+ /*!
+ * \brief Append a doc to the final content
+ *
+ * \param doc Doc to be printed
+ * \param paths Source paths requested for underlining
+ *
+ * \sa GetString
+ */
+ void Append(const Doc& doc, const ffi::Array<AccessPath>& paths);
+
+ /*!
+ * \brief Get the printed string of all Doc appended
+ *
+ * The content of each Doc in the returned string will
+ * appear in the same order as they are appended.
+ *
+ * \sa Append
+ */
+ ffi::String GetString() const;
+
+ /*!
+ * \brief Return the visible source path selected for each requested
+ * underline path.
+ */
+ ffi::Array<ffi::Optional<AccessPath>> GetVisiblePaths() const;
+
+ protected:
+ /*!
+ * \brief Get the printed string
+ *
+ * It will dispatch to the PrintTypedDoc method based on
+ * the actual type of Doc.
+ *
+ * \sa PrintTypedDoc
+ */
+ void PrintDoc(const Doc& doc);
+
+ /*!
+ * \brief Virtual method to print a LiteralDoc
+ */
+ 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
+ */
+ virtual void PrintTypedDoc(const IdDoc& doc) = 0;
+
+ /*! \brief Print a canonical namespace using this invocation's alias. */
+ virtual void PrintTypedDoc(const NamespaceDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an AttrAccessDoc
+ */
+ virtual void PrintTypedDoc(const AttrAccessDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an IndexDoc
+ */
+ virtual void PrintTypedDoc(const IndexDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an OperationDoc
+ */
+ virtual void PrintTypedDoc(const OperationDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a CallDoc
+ */
+ virtual void PrintTypedDoc(const CallDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a LambdaDoc
+ */
+ virtual void PrintTypedDoc(const LambdaDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ListDoc
+ */
+ virtual void PrintTypedDoc(const ListDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a TupleDoc
+ */
+ virtual void PrintTypedDoc(const TupleDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a DictDoc
+ */
+ virtual void PrintTypedDoc(const DictDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a SliceDoc
+ */
+ virtual void PrintTypedDoc(const SliceDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a StmtBlockDoc
+ */
+ virtual void PrintTypedDoc(const StmtBlockDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an AssignDoc
+ */
+ virtual void PrintTypedDoc(const AssignDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an IfDoc
+ */
+ virtual void PrintTypedDoc(const IfDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a WhileDoc
+ */
+ virtual void PrintTypedDoc(const WhileDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a BreakDoc
+ */
+ virtual void PrintTypedDoc(const BreakDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ContinueDoc
+ */
+ virtual void PrintTypedDoc(const ContinueDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ForDoc
+ */
+ virtual void PrintTypedDoc(const ForDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ScopeDoc
+ */
+ virtual void PrintTypedDoc(const ScopeDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an ExprStmtDoc
+ */
+ virtual void PrintTypedDoc(const ExprStmtDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print an AssertDoc
+ */
+ virtual void PrintTypedDoc(const AssertDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ReturnDoc
+ */
+ virtual void PrintTypedDoc(const ReturnDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a FunctionDoc
+ */
+ virtual void PrintTypedDoc(const FunctionDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a ClassDoc
+ */
+ virtual void PrintTypedDoc(const ClassDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a CommentDoc
+ */
+ virtual void PrintTypedDoc(const CommentDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a DocStringDoc
+ */
+ virtual void PrintTypedDoc(const DocStringDoc& doc) = 0;
+
+ /*!
+ * \brief Virtual method to print a OpCallDoc
+ */
+ virtual void PrintTypedDoc(const OpCallDoc& doc) = 0;
+
+ /*!
+ * \brief Increase the indent level of any content to be
+ * printed after this call
+ */
+ void IncreaseIndent() { indent_ += options_->indent_spaces; }
+
+ /*!
+ * \brief Decrease the indent level of any content to be
+ * printed after this call
+ */
+ void DecreaseIndent() { indent_ -= options_->indent_spaces; }
+
+ /*!
+ * \brief Add a new line into the output stream
+ *
+ * \sa output_
+ */
+ std::ostream& NewLine() {
+ size_t start_pos = output_.tellp();
+ output_ << "\n";
+ line_starts_.push_back(output_.tellp());
+ output_ << std::string(indent_, ' ');
+ size_t end_pos = output_.tellp();
+ underlines_exempted_.push_back({start_pos, end_pos});
+ return output_;
+ }
+
+ /*!
+ * \brief The output stream of printer
+ *
+ * All printed content will be stored in this stream and returned
+ * when GetString is called.
+ *
+ * \sa GetString
+ */
+ std::ostringstream output_;
+
+ /*! \brief Configuration is read only while a Doc is rendered. */
+ const PrinterConfig& config() const { return options_; }
+
+ /*! \brief Spans that we have already committed to underline exemption. */
+ std::vector<ByteSpan> underlines_exempted_;
+
+ private:
+ void MarkSpan(const ByteSpan& span, const AccessPath& path);
+
+ /*! \brief Options to customize certain aspects of the output */
+ PrinterConfig options_;
+
+ /*! \brief the current level of indent */
+ int indent_ = 0;
+
+ /*! \brief For each line in the output_, byte offset of its first character */
+ std::vector<size_t> line_starts_;
+
+ /*! \brief Path of the object that we would like to underline */
+ ffi::Array<AccessPath> path_to_underline_;
+
+ /*!
+ * \brief Candidate spans to be underlined, until we find a better match.
+ * (A better match is an object with a longer path that is still a prefix of path_to_underline_.)
+ */
+ std::vector<std::vector<ByteSpan>> current_underline_candidates_;
+
+ /*! \brief Path length of the objects that are current candidates for underlining. */
+ std::vector<int> current_max_path_depth_;
+
+ /*! \brief Visible path selected by the current best underline candidate. */
+ std::vector<ffi::Optional<AccessPath>> current_visible_paths_;
+
+ /*! \brief Spans that we have already committed to underline. */
+ std::vector<ByteSpan> underlines_;
+};
+
+namespace {
+
+std::vector<ByteSpan> MergeAndExemptSpans(const std::vector<ByteSpan>& spans,
+ const std::vector<ByteSpan>& spans_exempted) {
+ // use prefix sum to merge and exempt spans
+ std::vector<ByteSpan> res;
+ std::vector<std::pair<size_t, int>> prefix_stamp;
+ for (ByteSpan span : spans) {
+ prefix_stamp.push_back({span.first, 1});
+ prefix_stamp.push_back({span.second, -1});
+ }
+ // at most spans.size() spans accumulated in prefix sum
+ // use spans.size() + 1 as stamp unit to exempt all positive spans
+ // with only one negative span
+ int max_n = spans.size() + 1;
+ for (ByteSpan span : spans_exempted) {
+ prefix_stamp.push_back({span.first, -max_n});
+ prefix_stamp.push_back({span.second, max_n});
+ }
+ std::sort(prefix_stamp.begin(), prefix_stamp.end());
+ int prefix_sum = 0;
+ int n = prefix_stamp.size();
+ for (int i = 0; i < n - 1; ++i) {
+ prefix_sum += prefix_stamp[i].second;
+ // positive prefix sum leads to spans without exemption
+ // different stamp positions guarantee the stamps in same position accumulated
+ if (prefix_sum > 0 && prefix_stamp[i].first < prefix_stamp[i + 1].first) {
+ if (res.size() && res.back().second == prefix_stamp[i].first) {
+ // merge to the last spans if it is successive
+ res.back().second = prefix_stamp[i + 1].first;
+ } else {
+ // add a new independent span
+ res.push_back({prefix_stamp[i].first, prefix_stamp[i + 1].first});
+ }
+ }
+ }
+ return res;
+}
+
+size_t GetTextWidth(const std::string& text, const ByteSpan& span) {
+ // FIXME: this only works for ASCII characters.
+ // To do this "correctly", we need to parse UTF-8 into codepoints
+ // and call wcwidth() or equivalent for every codepoint.
+ size_t ret = 0;
+ for (size_t i = span.first; i != span.second; ++i) {
+ if (isprint(text[i])) {
+ ret += 1;
+ }
+ }
+ return ret;
+}
+
+size_t MoveBack(size_t pos, size_t distance) { return distance > pos ? 0 : pos - distance; }
+
+size_t MoveForward(size_t pos, size_t distance, size_t max) {
+ return distance > max - pos ? max : pos + distance;
+}
+
+size_t GetLineIndex(size_t byte_pos, const std::vector<size_t>& line_starts) {
+ auto it = std::upper_bound(line_starts.begin(), line_starts.end(), byte_pos);
+ return (it - line_starts.begin()) - 1;
+}
+
+using UnderlineIter = typename std::vector<ByteSpan>::const_iterator;
+
+ByteSpan PopNextUnderline(UnderlineIter* next_underline, UnderlineIter end_underline) {
+ if (*next_underline == end_underline) {
+ return {std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max()};
+ } else {
+ return *(*next_underline)++;
+ }
+}
+
+void PrintChunk(const std::pair<size_t, size_t>& lines_range,
+ const std::pair<UnderlineIter, UnderlineIter>& underlines, const std::string& text,
+ const std::vector<size_t>& line_starts, const PrinterConfig& options,
+ size_t line_number_width, std::string* out) {
+ UnderlineIter next_underline = underlines.first;
+ ByteSpan current_underline = PopNextUnderline(&next_underline, underlines.second);
+
+ for (size_t line_idx = lines_range.first; line_idx < lines_range.second; ++line_idx) {
+ if (options->print_line_numbers) {
+ std::string line_num_str = std::to_string(line_idx + 1);
+ line_num_str.push_back(' ');
+ for (size_t i = line_num_str.size(); i < line_number_width; ++i) {
+ out->push_back(' ');
+ }
+ *out += line_num_str;
+ }
+
+ size_t line_start = line_starts.at(line_idx);
+ size_t line_end =
+ line_idx + 1 == line_starts.size() ? text.size() : line_starts.at(line_idx + 1);
+ out->append(text.begin() + line_start, text.begin() + line_end);
+
+ bool printed_underline = false;
+ size_t line_pos = line_start;
+ bool printed_extra_caret = 0;
+ while (current_underline.first < line_end) {
+ if (!printed_underline) {
+ *out += std::string(line_number_width, ' ');
+ printed_underline = true;
+ }
+
+ size_t underline_end_for_line = std::min(line_end, current_underline.second);
+ size_t num_spaces = GetTextWidth(text, {line_pos, current_underline.first});
+ if (num_spaces > 0 && printed_extra_caret) {
+ num_spaces -= 1;
+ printed_extra_caret = false;
+ }
+ *out += std::string(num_spaces, ' ');
+
+ size_t num_carets = GetTextWidth(text, {current_underline.first, underline_end_for_line});
+ if (num_carets == 0 && !printed_extra_caret) {
+ // Special case: when underlineing an empty or unprintable string, make sure to print
+ // at least one caret still.
+ num_carets = 1;
+ printed_extra_caret = true;
+ } else if (num_carets > 0 && printed_extra_caret) {
+ num_carets -= 1;
+ printed_extra_caret = false;
+ }
+ *out += std::string(num_carets, '^');
+
+ line_pos = current_underline.first = underline_end_for_line;
+ if (current_underline.first == current_underline.second) {
+ current_underline = PopNextUnderline(&next_underline, underlines.second);
+ }
+ }
+
+ if (printed_underline) {
+ out->push_back('\n');
+ }
+ }
+}
+
+void PrintCut(size_t num_lines_skipped, std::string* out) {
+ if (num_lines_skipped != 0) {
+ std::ostringstream s;
+ s << "(... " << num_lines_skipped << " lines skipped ...)\n";
+ *out += s.str();
+ }
+}
+
+std::pair<size_t, size_t> GetLinesForUnderline(const ByteSpan& underline,
+ const std::vector<size_t>& line_starts,
+ size_t num_lines, const PrinterConfig& options) {
+ int context_lines = options->num_context_lines < 0 ? std::numeric_limits<int32_t>::max()
+ : options->num_context_lines;
+ size_t first_line_of_underline = GetLineIndex(underline.first, line_starts);
+ size_t first_line_of_chunk = MoveBack(first_line_of_underline, context_lines);
+ size_t end_line_of_underline = GetLineIndex(underline.second - 1, line_starts) + 1;
+ size_t end_line_of_chunk = MoveForward(end_line_of_underline, context_lines, num_lines);
+
+ return {first_line_of_chunk, end_line_of_chunk};
+}
+
+// If there is only one line between the chunks, it is better to print it as is,
+// rather than something like "(... 1 line skipped ...)".
+constexpr const size_t kMinLinesToCutOut = 2;
+
+bool TryMergeChunks(std::pair<size_t, size_t>* cur_chunk,
+ const std::pair<size_t, size_t>& new_chunk) {
+ if (new_chunk.first < cur_chunk->second + kMinLinesToCutOut) {
+ cur_chunk->second = new_chunk.second;
+ return true;
+ } else {
+ return false;
+ }
+}
+
+size_t GetNumLines(const std::string& text, const std::vector<size_t>& line_starts) {
+ if (line_starts.back() == text.size()) {
+ // Final empty line doesn't count as a line
+ return line_starts.size() - 1;
+ } else {
+ return line_starts.size();
+ }
+}
+
+size_t GetLineNumberWidth(size_t num_lines, const PrinterConfig& options) {
+ if (options->print_line_numbers) {
+ return std::to_string(num_lines).size() + 1;
+ } else {
+ return 0;
+ }
+}
+
+std::string DecorateText(const std::string& text, const std::vector<size_t>& line_starts,
+ const PrinterConfig& options, const std::vector<ByteSpan>& underlines) {
+ size_t num_lines = GetNumLines(text, line_starts);
+ size_t line_number_width = GetLineNumberWidth(num_lines, options);
+
+ std::string ret;
+ if (underlines.empty()) {
+ PrintChunk({0, num_lines}, {underlines.begin(), underlines.begin()}, text, line_starts, options,
+ line_number_width, &ret);
+ return ret;
+ }
+
+ size_t last_end_line = 0;
+ std::pair<size_t, size_t> cur_chunk =
+ GetLinesForUnderline(underlines[0], line_starts, num_lines, options);
+ if (cur_chunk.first < kMinLinesToCutOut) {
+ cur_chunk.first = 0;
+ }
+
+ auto first_underline_in_cur_chunk = underlines.begin();
+ for (auto underline_it = underlines.begin() + 1; underline_it != underlines.end();
+ ++underline_it) {
+ std::pair<size_t, size_t> new_chunk =
+ GetLinesForUnderline(*underline_it, line_starts, num_lines, options);
+
+ if (!TryMergeChunks(&cur_chunk, new_chunk)) {
+ PrintCut(cur_chunk.first - last_end_line, &ret);
+ PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underline_it}, text, line_starts,
+ options, line_number_width, &ret);
+ last_end_line = cur_chunk.second;
+ cur_chunk = new_chunk;
+ first_underline_in_cur_chunk = underline_it;
+ }
+ }
+
+ PrintCut(cur_chunk.first - last_end_line, &ret);
+ if (num_lines - cur_chunk.second < kMinLinesToCutOut) {
+ cur_chunk.second = num_lines;
+ }
+ PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underlines.end()}, text, line_starts,
+ options, line_number_width, &ret);
+ PrintCut(num_lines - cur_chunk.second, &ret);
+ return ret;
+}
+
+} // namespace
+
+DocPrinter::DocPrinter(const PrinterConfig& options) : options_(options) {
+ line_starts_.push_back(0);
+}
+
+void DocPrinter::AppendHeader(const ffi::Array<ffi::Any>& header) {
+ for (const ffi::Any& item : header) {
+ if (auto comment = item.as<CommentDoc>()) {
+ PrintDoc(comment.value());
+ NewLine();
+ } else {
+ // Executable imports and separators are fixed entry-owned text. Comment
+ // headers use ordinary Docs and their normal span/exemption machinery.
+ ffi::String text = item.as_or_throw<ffi::String>();
+ for (size_t i = 0; i < text.size(); ++i) {
+ if (text.data()[i] == '\n')
+ NewLine();
+ else
+ output_ << text.data()[i];
+ }
+ }
+ }
+}
+
+void DocPrinter::Append(const Doc& doc, const ffi::Array<AccessPath>& paths) {
+ for (const AccessPath& p : paths) {
+ path_to_underline_.push_back(p);
+ current_max_path_depth_.push_back(0);
+ current_visible_paths_.push_back(std::nullopt);
+ current_underline_candidates_.push_back(std::vector<ByteSpan>());
+ }
+ PrintDoc(doc);
+ for (const auto& c : current_underline_candidates_) {
+ underlines_.insert(underlines_.end(), c.begin(), c.end());
+ }
+}
+
+ffi::String DocPrinter::GetString() const {
+ std::string text = output_.str();
+
+ // Remove any trailing indentation
+ while (!text.empty() && text.back() == ' ') {
+ text.pop_back();
+ }
+
+ if (!text.empty() && text.back() != '\n') {
+ text.push_back('\n');
+ }
+
+ return DecorateText(text, line_starts_, options_,
+ MergeAndExemptSpans(underlines_, underlines_exempted_));
+}
+
+ffi::Array<ffi::Optional<AccessPath>> DocPrinter::GetVisiblePaths() const {
+ return ffi::Array<ffi::Optional<AccessPath>>(current_visible_paths_);
+}
+
+void DocPrinter::PrintDoc(const Doc& doc) {
+ size_t start_pos = output_.tellp();
+
+ static const auto dispatch = []() {
+ ObjectFunctor<void(const ffi::ObjectRef&, DocPrinter*)> table;
+#define TVM_SCRIPT_PRINT_DISPATCH(DocType) \
+ table.SetDispatch<DocType##Node>([](const ffi::ObjectRef& obj, DocPrinter* self) { \
+ self->PrintTypedDoc(ffi::GetRef<DocType>(static_cast<const DocType##Node*>(obj.get()))); \
+ });
+ TVM_SCRIPT_PRINT_DISPATCH(LiteralDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ExprStringDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(IdDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(NamespaceDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(AttrAccessDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(IndexDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(OperationDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(CallDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(LambdaDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ListDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(TupleDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(DictDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(SliceDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(StmtBlockDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(AssignDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(IfDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(WhileDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(BreakDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ContinueDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ForDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ScopeDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ExprStmtDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(AssertDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ReturnDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(FunctionDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(ClassDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(CommentDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(DocStringDoc);
+ TVM_SCRIPT_PRINT_DISPATCH(OpCallDoc);
+#undef TVM_SCRIPT_PRINT_DISPATCH
+ table.Finalize();
+ return table;
+ }();
+ dispatch(doc, this);
+
+ size_t end_pos = output_.tellp();
+ for (const AccessPath& path : doc->source_paths) {
+ MarkSpan({start_pos, end_pos}, path);
+ }
+}
+
+void DocPrinter::MarkSpan(const ByteSpan& span, const AccessPath& path) {
+ int n = path_to_underline_.size();
+ for (int i = 0; i < n; ++i) {
+ AccessPath p = path_to_underline_[i];
+ if (path->depth >= current_max_path_depth_[i] && path->IsPrefixOf(p)) {
+ if (path->depth > current_max_path_depth_[i] || !current_visible_paths_[i].has_value()) {
+ current_max_path_depth_[i] = path->depth;
+ current_visible_paths_[i] = path;
+ current_underline_candidates_[i].clear();
+ }
+ current_underline_candidates_[i].push_back(span);
+ }
+ }
+}
+
+namespace {
+
/*!
* \brief Escape an ExprStringDoc child while preserving final output positions.
*
@@ -227,6 +879,7 @@ ExprPrecedence GetExprPrecedence(const ExprDoc& doc) {
{LiteralDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{ExprStringDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{IdDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
+ {NamespaceDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{AttrAccessDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{IndexDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
{CallDocNode::RuntimeTypeIndex(), ExprPrecedence::kIdentity},
@@ -262,6 +915,7 @@ class PythonDocPrinter : public DocPrinter {
void PrintTypedDoc(const LiteralDoc& doc) final;
void PrintTypedDoc(const ExprStringDoc& doc) final;
void PrintTypedDoc(const IdDoc& doc) final;
+ void PrintTypedDoc(const NamespaceDoc& doc) final;
void PrintTypedDoc(const AttrAccessDoc& doc) final;
void PrintTypedDoc(const IndexDoc& doc) final;
void PrintTypedDoc(const OperationDoc& doc) final;
@@ -288,7 +942,15 @@ class PythonDocPrinter : public DocPrinter {
void PrintTypedDoc(const DocStringDoc& doc) final;
void PrintTypedDoc(const OpCallDoc& doc) final;
+ void PrintAssignment(const ExprDoc& lhs, const ffi::Optional<ExprDoc>& rhs,
+ const ffi::Optional<ExprDoc>& annotation);
+
private:
+ /*! \brief Borrowed tuple whose parentheses are omitted during the active print call. */
+ const TupleDocNode* unparenthesized_tuple_{nullptr};
+ /*! \brief Whether that tuple's singleton comma is supplied by the surrounding syntax. */
+ bool omit_tuple_singleton_comma_{false};
+
void NewLineWithoutIndent() {
size_t start_pos = output_.tellp();
output_ << "\n";
@@ -311,10 +973,13 @@ class PythonDocPrinter : public DocPrinter {
void PrintIndentedBlock(const ffi::Array<StmtDoc>& docs) {
IncreaseIndent();
- for (const StmtDoc& d : docs) {
+ bool outer_is_last_stmt = is_last_stmt_;
+ for (size_t i = 0; i < docs.size(); ++i) {
NewLine();
- PrintDoc(d);
+ is_last_stmt_ = i + 1 == docs.size();
+ PrintDoc(docs[i]);
}
+ is_last_stmt_ = outer_is_last_stmt;
if (docs.empty()) {
NewLine();
output_ << "pass";
@@ -425,6 +1090,9 @@ class PythonDocPrinter : public DocPrinter {
PrintDocString(comment);
DecreaseIndent();
}
+
+ // A ScopeDoc may use concise spelling only with enclosing-list context.
+ bool is_last_stmt_{false};
};
void PythonDocPrinter::PrintTypedDoc(const LiteralDoc& doc) {
@@ -487,6 +1155,16 @@ void PythonDocPrinter::PrintTypedDoc(const ExprStringDoc& doc) {
void PythonDocPrinter::PrintTypedDoc(const IdDoc& doc) { output_ << doc->name; }
+void PythonDocPrinter::PrintTypedDoc(const NamespaceDoc& doc) {
+ const ffi::String& name = doc->canonical_name;
+ ffi::String key = std::string(name) + ".prefix";
+ const auto& prefixes = GetDialectPrefixes();
+ auto registered = prefixes.find(key);
+ ffi::String fallback = registered == prefixes.end() ? name : registered->second;
+ output_ << (name == "ir" ? config()->ir_prefix
+ : config()->GetExtraConfig<ffi::String>(key, fallback));
+}
+
void PythonDocPrinter::PrintTypedDoc(const AttrAccessDoc& doc) {
PrintChildExpr(doc->value, doc);
output_ << "." << doc->name;
@@ -636,14 +1314,15 @@ void PythonDocPrinter::PrintTypedDoc(const ListDoc& doc) {
}
void PythonDocPrinter::PrintTypedDoc(const TupleDoc& doc) {
- output_ << "(";
+ bool parentheses = doc.get() != unparenthesized_tuple_ || doc->elements.empty();
+ if (parentheses) output_ << "(";
if (doc->elements.size() == 1) {
PrintDoc(doc->elements[0]);
- output_ << ",";
+ if (parentheses || !omit_tuple_singleton_comma_) output_ << ",";
} else {
PrintJoinedDocs(doc->elements, ", ");
}
- output_ << ")";
+ if (parentheses) output_ << ")";
}
void PythonDocPrinter::PrintTypedDoc(const DictDoc& doc) {
@@ -678,37 +1357,43 @@ void PythonDocPrinter::PrintTypedDoc(const SliceDoc& doc) {
}
void PythonDocPrinter::PrintTypedDoc(const StmtBlockDoc& doc) {
- for (const StmtDoc& stmt : doc->stmts) {
- PrintDoc(stmt);
- if (stmt != doc->stmts.back()) {
+ bool outer_is_last_stmt = is_last_stmt_;
+ for (size_t i = 0; i < doc->stmts.size(); ++i) {
+ is_last_stmt_ = i + 1 == doc->stmts.size() && outer_is_last_stmt;
+ PrintDoc(doc->stmts[i]);
+ if (i + 1 != doc->stmts.size()) {
NewLine();
}
}
+ is_last_stmt_ = outer_is_last_stmt;
}
-void PythonDocPrinter::PrintTypedDoc(const AssignDoc& doc) {
- if (const auto* tuple_doc = doc->lhs.as<TupleDocNode>()) {
- PrintJoinedDocs(tuple_doc->elements, ", ");
- } else {
- PrintDoc(doc->lhs);
- }
-
- if (doc->annotation) {
+void PythonDocPrinter::PrintAssignment(const ExprDoc& lhs, const ffi::Optional<ExprDoc>& rhs,
+ const ffi::Optional<ExprDoc>& annotation) {
+ const TupleDocNode* outer_tuple = unparenthesized_tuple_;
+ bool outer_comma = omit_tuple_singleton_comma_;
+ unparenthesized_tuple_ = lhs.as<TupleDocNode>();
+ omit_tuple_singleton_comma_ = true;
+ PrintDoc(lhs);
+ unparenthesized_tuple_ = outer_tuple;
+ omit_tuple_singleton_comma_ = outer_comma;
+
+ if (annotation) {
output_ << ": ";
- PrintDoc(doc->annotation.value());
+ PrintDoc(annotation.value());
}
- if (doc->rhs) {
+ if (rhs) {
output_ << " = ";
- if (const auto* tuple_doc = doc->rhs.as<TupleDocNode>()) {
- if (tuple_doc->elements.size() > 1) {
- PrintJoinedDocs(tuple_doc->elements, ", ");
- } else {
- PrintDoc(doc->rhs.value());
- }
- } else {
- PrintDoc(doc->rhs.value());
+ if (const auto* tuple = rhs.as<TupleDocNode>(); tuple && tuple->elements.size() > 1) {
+ unparenthesized_tuple_ = tuple;
}
+ PrintDoc(rhs.value());
+ unparenthesized_tuple_ = outer_tuple;
}
+}
+
+void PythonDocPrinter::PrintTypedDoc(const AssignDoc& doc) {
+ PrintAssignment(doc->lhs, doc->rhs, doc->annotation);
MaybePrintCommentInline(doc);
}
@@ -743,16 +1428,10 @@ void PythonDocPrinter::PrintTypedDoc(const ContinueDoc& doc) { output_ << "conti
void PythonDocPrinter::PrintTypedDoc(const ForDoc& doc) {
MaybePrintCommenMultiLines(doc, true);
output_ << "for ";
- if (const auto* tuple = doc->lhs.as<TupleDocNode>()) {
- if (tuple->elements.size() == 1) {
- PrintDoc(tuple->elements[0]);
- output_ << ",";
- } else {
- PrintJoinedDocs(tuple->elements, ", ");
- }
- } else {
- PrintDoc(doc->lhs);
- }
+ const TupleDocNode* outer_tuple = unparenthesized_tuple_;
+ unparenthesized_tuple_ = doc->lhs.as<TupleDocNode>();
+ PrintDoc(doc->lhs);
+ unparenthesized_tuple_ = outer_tuple;
output_ << " in ";
PrintDoc(doc->rhs);
output_ << ":";
@@ -762,6 +1441,19 @@ void PythonDocPrinter::PrintTypedDoc(const ForDoc& doc) {
void PythonDocPrinter::PrintTypedDoc(const ScopeDoc& doc) {
MaybePrintCommenMultiLines(doc, true);
+ if (doc->allow_concise_scoping && is_last_stmt_ && !doc->comment.has_value()) {
+ if (doc->lhs.has_value()) {
+ PrintDoc(doc->lhs.value());
+ output_ << " = ";
+ }
+ PrintDoc(doc->rhs);
+ for (size_t i = 0; i < doc->body.size(); ++i) {
+ NewLine();
+ is_last_stmt_ = i + 1 == doc->body.size();
+ PrintDoc(doc->body[i]);
+ }
+ return;
+ }
output_ << "with ";
PrintDoc(doc->rhs);
if (doc->lhs != nullptr) {
@@ -917,12 +1609,94 @@ void PythonDocPrinter::PrintTypedDoc(const OpCallDoc& doc) {
output_ << ")";
}
-ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
- if (cfg->num_context_lines < 0) {
- cfg->num_context_lines = std::numeric_limits<int32_t>::max();
+// Associate each recovered path with its nearest statement container. Expression
+// children keep their more precise paths while annotations use the enclosing
+// statement's ordinary comment rendering. Transparent statement lists introduce
+// no separate comment position.
+void CollectAnnotationTargets(ffi::AnyView value, ffi::Optional<StmtDoc> enclosing,
+ std::vector<std::pair<AccessPath, StmtDoc>>* targets,
+ std::unordered_set<const ffi::Object*>* active,
+ bool function_parameter = false) {
+ const ffi::Object* object = value.as<ffi::Object>();
+ if (!object || !active->insert(object).second) return;
+ if (auto stmt = value.as<StmtDoc>();
+ stmt && !value.as<StmtBlockDocNode>() && !(function_parameter && value.as<AssignDocNode>())) {
+ enclosing = stmt.value();
+ }
+ if (auto doc = value.as<DocNode>(); doc && enclosing.has_value()) {
+ for (const AccessPath& path : doc->source_paths) {
+ targets->emplace_back(path, enclosing.value());
+ }
+ }
+ if (auto array = value.as<ffi::Array<ffi::Any>>()) {
+ for (const auto& item : array.value()) {
+ CollectAnnotationTargets(item, enclosing, targets, active, function_parameter);
+ }
+ } else if (auto map = value.as<ffi::Map<ffi::Any, ffi::Any>>()) {
+ for (const auto& [key, item] : map.value()) {
+ CollectAnnotationTargets(item, enclosing, targets, active, function_parameter);
+ }
+ } else {
+ ffi::reflection::ForEachFieldInfo(
+ TVMFFIGetTypeInfo(object->type_index()), [&](const TVMFFIFieldInfo* field) {
+ if (ffi::String(field->name) == "source_paths") return;
+ // Parameters are AssignDocs for signature layout, but cannot carry
+ // statement comments. Keep their expressions' precise source paths
+ // while placing requested annotations on the enclosing function.
+ bool parameter = function_parameter || (value.as<FunctionDoc>().has_value() &&
+ ffi::String(field->name) == "args");
+ CollectAnnotationTargets(ffi::reflection::FieldGetter(field)(object), enclosing, targets,
+ active, parameter);
+ });
+ }
+ active->erase(object);
+}
+
+void AttachAnnotations(const Doc& root, const ffi::Map<AccessPath, ffi::String>& annotations) {
+ if (annotations.empty()) return;
+ std::vector<std::pair<AccessPath, StmtDoc>> targets;
+ std::unordered_set<const ffi::Object*> active;
+ CollectAnnotationTargets(root, std::nullopt, &targets, &active);
+ std::unordered_map<const ffi::Object*, std::unordered_set<std::string>> attached;
+ for (const auto& [requested, message] : annotations) {
+ int deepest = -1;
+ for (const auto& [path, stmt] : targets) {
+ if (path->IsPrefixOf(requested)) deepest = std::max(deepest, path->depth);
+ }
+ for (const auto& [path, stmt] : targets) {
+ if (path->depth != deepest || !path->IsPrefixOf(requested)) continue;
+ if (!attached[stmt.get()].insert(std::string(message)).second) continue;
+ bool inline_comment = stmt.as<AssignDocNode>() || stmt.as<ExprStmtDocNode>() ||
+ stmt.as<AssertDocNode>() || stmt.as<ReturnDocNode>();
+ std::string text = message;
+ if (inline_comment) {
+ // Diagnostic requests may converge on one visible statement. Keep
+ // every message without constructing an invalid multiline inline Doc.
+ size_t newline = 0;
+ while ((newline = text.find('\n', newline)) != std::string::npos) {
+ text.replace(newline, 1, "; ");
+ newline += 2;
+ }
+ }
+ stmt->comment = stmt->comment.has_value() ? ffi::String(std::string(stmt->comment.value()) +
+ (inline_comment ? "; " : "\n") + text)
+ : ffi::String(text);
+ }
}
+}
+
+} // namespace
+
+namespace details {
+
+ffi::String RenderPythonScript(Doc doc, const PrinterConfig& cfg,
+ const ffi::Array<ffi::Any>& header,
+ const ffi::Array<AccessPath>& underline_paths,
+ const ffi::Map<AccessPath, ffi::String>& annotations) {
+ AttachAnnotations(doc, annotations);
PythonDocPrinter printer(cfg);
- printer.Append(doc, cfg);
+ printer.AppendHeader(header);
+ printer.Append(doc, underline_paths);
std::string script = printer.GetString();
// GetString terminates non-empty output with one newline. Preserve the
@@ -934,7 +1708,18 @@ ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
}
if (!cfg->render_invisible_path_info) return script;
- return RenderInvisiblePathInfo(script, cfg->path_to_underline, printer.GetVisiblePaths());
+ return RenderInvisiblePathInfo(script, underline_paths, printer.GetVisiblePaths());
+}
+
+} // namespace details
+
+ffi::String DocToPythonScriptWithHeader(Doc doc, const PrinterConfig& cfg,
+ const ffi::Array<ffi::Any>& header) {
+ return details::RenderPythonScript(std::move(doc), cfg, header, cfg->path_to_underline, {});
+}
+
+ffi::String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
+ return DocToPythonScriptWithHeader(std::move(doc), cfg, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/script/printer/dialect_prefix.h b/src/script/printer/doc_printer.h
similarity index 61%
copy from src/script/printer/dialect_prefix.h
copy to src/script/printer/doc_printer.h
index 5aaae11293..819e3467db 100644
--- a/src/script/printer/dialect_prefix.h
+++ b/src/script/printer/doc_printer.h
@@ -16,21 +16,25 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
-#define TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
+#ifndef SRC_SCRIPT_PRINTER_DOC_PRINTER_H_
+#define SRC_SCRIPT_PRINTER_DOC_PRINTER_H_
-#include <tvm/ffi/string.h>
+#include <tvm/script/printer/doc.h>
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-// Register during dialect static initialization so configuration can validate
-// and reserve the prefix before any docsifier assigns variable names.
-void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix);
+// Render a Doc with this invocation's recovered paths and statement annotations.
+ffi::String RenderPythonScript(Doc doc, const PrinterConfig& config,
+ const ffi::Array<ffi::Any>& header,
+ const ffi::Array<AccessPath>& underline_paths,
+ const ffi::Map<AccessPath, ffi::String>& annotations);
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
+#endif // SRC_SCRIPT_PRINTER_DOC_PRINTER_H_
diff --git a/src/script/printer/doc_printer/base_doc_printer.cc b/src/script/printer/doc_printer/base_doc_printer.cc
deleted file mode 100644
index 6c1a71f304..0000000000
--- a/src/script/printer/doc_printer/base_doc_printer.cc
+++ /dev/null
@@ -1,386 +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 "./base_doc_printer.h"
-
-namespace tvm {
-namespace script {
-namespace printer {
-
-namespace {
-
-std::vector<ByteSpan> MergeAndExemptSpans(const std::vector<ByteSpan>& spans,
- const std::vector<ByteSpan>& spans_exempted) {
- // use prefix sum to merge and exempt spans
- std::vector<ByteSpan> res;
- std::vector<std::pair<size_t, int>> prefix_stamp;
- for (ByteSpan span : spans) {
- prefix_stamp.push_back({span.first, 1});
- prefix_stamp.push_back({span.second, -1});
- }
- // at most spans.size() spans accumulated in prefix sum
- // use spans.size() + 1 as stamp unit to exempt all positive spans
- // with only one negative span
- int max_n = spans.size() + 1;
- for (ByteSpan span : spans_exempted) {
- prefix_stamp.push_back({span.first, -max_n});
- prefix_stamp.push_back({span.second, max_n});
- }
- std::sort(prefix_stamp.begin(), prefix_stamp.end());
- int prefix_sum = 0;
- int n = prefix_stamp.size();
- for (int i = 0; i < n - 1; ++i) {
- prefix_sum += prefix_stamp[i].second;
- // positive prefix sum leads to spans without exemption
- // different stamp positions guarantee the stamps in same position accumulated
- if (prefix_sum > 0 && prefix_stamp[i].first < prefix_stamp[i + 1].first) {
- if (res.size() && res.back().second == prefix_stamp[i].first) {
- // merge to the last spans if it is successive
- res.back().second = prefix_stamp[i + 1].first;
- } else {
- // add a new independent span
- res.push_back({prefix_stamp[i].first, prefix_stamp[i + 1].first});
- }
- }
- }
- return res;
-}
-
-size_t GetTextWidth(const std::string& text, const ByteSpan& span) {
- // FIXME: this only works for ASCII characters.
- // To do this "correctly", we need to parse UTF-8 into codepoints
- // and call wcwidth() or equivalent for every codepoint.
- size_t ret = 0;
- for (size_t i = span.first; i != span.second; ++i) {
- if (isprint(text[i])) {
- ret += 1;
- }
- }
- return ret;
-}
-
-size_t MoveBack(size_t pos, size_t distance) { return distance > pos ? 0 : pos - distance; }
-
-size_t MoveForward(size_t pos, size_t distance, size_t max) {
- return distance > max - pos ? max : pos + distance;
-}
-
-size_t GetLineIndex(size_t byte_pos, const std::vector<size_t>& line_starts) {
- auto it = std::upper_bound(line_starts.begin(), line_starts.end(), byte_pos);
- return (it - line_starts.begin()) - 1;
-}
-
-using UnderlineIter = typename std::vector<ByteSpan>::const_iterator;
-
-ByteSpan PopNextUnderline(UnderlineIter* next_underline, UnderlineIter end_underline) {
- if (*next_underline == end_underline) {
- return {std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max()};
- } else {
- return *(*next_underline)++;
- }
-}
-
-void PrintChunk(const std::pair<size_t, size_t>& lines_range,
- const std::pair<UnderlineIter, UnderlineIter>& underlines, const std::string& text,
- const std::vector<size_t>& line_starts, const PrinterConfig& options,
- size_t line_number_width, std::string* out) {
- UnderlineIter next_underline = underlines.first;
- ByteSpan current_underline = PopNextUnderline(&next_underline, underlines.second);
-
- for (size_t line_idx = lines_range.first; line_idx < lines_range.second; ++line_idx) {
- if (options->print_line_numbers) {
- std::string line_num_str = std::to_string(line_idx + 1);
- line_num_str.push_back(' ');
- for (size_t i = line_num_str.size(); i < line_number_width; ++i) {
- out->push_back(' ');
- }
- *out += line_num_str;
- }
-
- size_t line_start = line_starts.at(line_idx);
- size_t line_end =
- line_idx + 1 == line_starts.size() ? text.size() : line_starts.at(line_idx + 1);
- out->append(text.begin() + line_start, text.begin() + line_end);
-
- bool printed_underline = false;
- size_t line_pos = line_start;
- bool printed_extra_caret = 0;
- while (current_underline.first < line_end) {
- if (!printed_underline) {
- *out += std::string(line_number_width, ' ');
- printed_underline = true;
- }
-
- size_t underline_end_for_line = std::min(line_end, current_underline.second);
- size_t num_spaces = GetTextWidth(text, {line_pos, current_underline.first});
- if (num_spaces > 0 && printed_extra_caret) {
- num_spaces -= 1;
- printed_extra_caret = false;
- }
- *out += std::string(num_spaces, ' ');
-
- size_t num_carets = GetTextWidth(text, {current_underline.first, underline_end_for_line});
- if (num_carets == 0 && !printed_extra_caret) {
- // Special case: when underlineing an empty or unprintable string, make sure to print
- // at least one caret still.
- num_carets = 1;
- printed_extra_caret = true;
- } else if (num_carets > 0 && printed_extra_caret) {
- num_carets -= 1;
- printed_extra_caret = false;
- }
- *out += std::string(num_carets, '^');
-
- line_pos = current_underline.first = underline_end_for_line;
- if (current_underline.first == current_underline.second) {
- current_underline = PopNextUnderline(&next_underline, underlines.second);
- }
- }
-
- if (printed_underline) {
- out->push_back('\n');
- }
- }
-}
-
-void PrintCut(size_t num_lines_skipped, std::string* out) {
- if (num_lines_skipped != 0) {
- std::ostringstream s;
- s << "(... " << num_lines_skipped << " lines skipped ...)\n";
- *out += s.str();
- }
-}
-
-std::pair<size_t, size_t> GetLinesForUnderline(const ByteSpan& underline,
- const std::vector<size_t>& line_starts,
- size_t num_lines, const PrinterConfig& options) {
- size_t first_line_of_underline = GetLineIndex(underline.first, line_starts);
- size_t first_line_of_chunk = MoveBack(first_line_of_underline, options->num_context_lines);
- size_t end_line_of_underline = GetLineIndex(underline.second - 1, line_starts) + 1;
- size_t end_line_of_chunk =
- MoveForward(end_line_of_underline, options->num_context_lines, num_lines);
-
- return {first_line_of_chunk, end_line_of_chunk};
-}
-
-// If there is only one line between the chunks, it is better to print it as is,
-// rather than something like "(... 1 line skipped ...)".
-constexpr const size_t kMinLinesToCutOut = 2;
-
-bool TryMergeChunks(std::pair<size_t, size_t>* cur_chunk,
- const std::pair<size_t, size_t>& new_chunk) {
- if (new_chunk.first < cur_chunk->second + kMinLinesToCutOut) {
- cur_chunk->second = new_chunk.second;
- return true;
- } else {
- return false;
- }
-}
-
-size_t GetNumLines(const std::string& text, const std::vector<size_t>& line_starts) {
- if (line_starts.back() == text.size()) {
- // Final empty line doesn't count as a line
- return line_starts.size() - 1;
- } else {
- return line_starts.size();
- }
-}
-
-size_t GetLineNumberWidth(size_t num_lines, const PrinterConfig& options) {
- if (options->print_line_numbers) {
- return std::to_string(num_lines).size() + 1;
- } else {
- return 0;
- }
-}
-
-std::string DecorateText(const std::string& text, const std::vector<size_t>& line_starts,
- const PrinterConfig& options, const std::vector<ByteSpan>& underlines) {
- size_t num_lines = GetNumLines(text, line_starts);
- size_t line_number_width = GetLineNumberWidth(num_lines, options);
-
- std::string ret;
- if (underlines.empty()) {
- PrintChunk({0, num_lines}, {underlines.begin(), underlines.begin()}, text, line_starts, options,
- line_number_width, &ret);
- return ret;
- }
-
- size_t last_end_line = 0;
- std::pair<size_t, size_t> cur_chunk =
- GetLinesForUnderline(underlines[0], line_starts, num_lines, options);
- if (cur_chunk.first < kMinLinesToCutOut) {
- cur_chunk.first = 0;
- }
-
- auto first_underline_in_cur_chunk = underlines.begin();
- for (auto underline_it = underlines.begin() + 1; underline_it != underlines.end();
- ++underline_it) {
- std::pair<size_t, size_t> new_chunk =
- GetLinesForUnderline(*underline_it, line_starts, num_lines, options);
-
- if (!TryMergeChunks(&cur_chunk, new_chunk)) {
- PrintCut(cur_chunk.first - last_end_line, &ret);
- PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underline_it}, text, line_starts,
- options, line_number_width, &ret);
- last_end_line = cur_chunk.second;
- cur_chunk = new_chunk;
- first_underline_in_cur_chunk = underline_it;
- }
- }
-
- PrintCut(cur_chunk.first - last_end_line, &ret);
- if (num_lines - cur_chunk.second < kMinLinesToCutOut) {
- cur_chunk.second = num_lines;
- }
- PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underlines.end()}, text, line_starts,
- options, line_number_width, &ret);
- PrintCut(num_lines - cur_chunk.second, &ret);
- return ret;
-}
-
-} // namespace
-
-DocPrinter::DocPrinter(const PrinterConfig& options) : options_(options) {
- line_starts_.push_back(0);
-}
-
-void DocPrinter::Append(const Doc& doc) { Append(doc, PrinterConfig()); }
-
-void DocPrinter::Append(const Doc& doc, const PrinterConfig& cfg) {
- for (const AccessPath& p : cfg->path_to_underline) {
- path_to_underline_.push_back(p);
- current_max_path_depth_.push_back(0);
- current_visible_paths_.push_back(std::nullopt);
- current_underline_candidates_.push_back(std::vector<ByteSpan>());
- }
- PrintDoc(doc);
- for (const auto& c : current_underline_candidates_) {
- underlines_.insert(underlines_.end(), c.begin(), c.end());
- }
-}
-
-ffi::String DocPrinter::GetString() const {
- std::string text = output_.str();
-
- // Remove any trailing indentation
- while (!text.empty() && text.back() == ' ') {
- text.pop_back();
- }
-
- if (!text.empty() && text.back() != '\n') {
- text.push_back('\n');
- }
-
- return DecorateText(text, line_starts_, options_,
- MergeAndExemptSpans(underlines_, underlines_exempted_));
-}
-
-ffi::Array<ffi::Optional<AccessPath>> DocPrinter::GetVisiblePaths() const {
- return ffi::Array<ffi::Optional<AccessPath>>(current_visible_paths_);
-}
-
-void DocPrinter::PrintDoc(const Doc& doc) {
- size_t start_pos = output_.tellp();
-
- 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>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<IndexDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<OperationDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<CallDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<LambdaDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ListDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<TupleDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<DictDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<SliceDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<StmtBlockDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<AssignDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<IfDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<WhileDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<BreakDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ContinueDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ForDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ScopeDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ExprStmtDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<AssertDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ReturnDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<FunctionDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<ClassDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<CommentDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<DocStringDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else if (auto doc_node = doc.as<OpCallDoc>()) {
- PrintTypedDoc(doc_node.value());
- } else {
- TVM_FFI_THROW(InternalError) << "Do not know how to print " << doc->GetTypeKey();
- throw;
- }
-
- size_t end_pos = output_.tellp();
- for (const AccessPath& path : doc->source_paths) {
- MarkSpan({start_pos, end_pos}, path);
- }
-}
-
-void DocPrinter::MarkSpan(const ByteSpan& span, const AccessPath& path) {
- int n = path_to_underline_.size();
- for (int i = 0; i < n; ++i) {
- AccessPath p = path_to_underline_[i];
- if (path->depth >= current_max_path_depth_[i] && path->IsPrefixOf(p)) {
- if (path->depth > current_max_path_depth_[i] || !current_visible_paths_[i].has_value()) {
- current_max_path_depth_[i] = path->depth;
- current_visible_paths_[i] = path;
- current_underline_candidates_[i].clear();
- }
- current_underline_candidates_[i].push_back(span);
- }
- }
-}
-
-} // namespace printer
-} // namespace script
-} // namespace tvm
diff --git a/src/script/printer/doc_printer/base_doc_printer.h b/src/script/printer/doc_printer/base_doc_printer.h
deleted file mode 100644
index c8335f951a..0000000000
--- a/src/script/printer/doc_printer/base_doc_printer.h
+++ /dev/null
@@ -1,318 +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_DOC_PRINTER_BASE_DOC_PRINTER_H_
-#define TVM_SCRIPT_PRINTER_DOC_PRINTER_BASE_DOC_PRINTER_H_
-
-#include <tvm/script/printer/doc.h>
-
-#include <limits>
-#include <memory>
-#include <ostream>
-#include <string>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace script {
-namespace printer {
-
-/*! \brief Range of byte offsets in a string */
-using ByteSpan = std::pair<size_t, size_t>;
-
-/*!
- * \brief DocPrinter is responsible for printing Doc tree into text format
- * \details This is the base class for translating Doc into string.
- * Each target language needs to have its subclass of DocPrinter
- * to define the actual logic of printing Doc.
- *
- * \sa Doc
- */
-class DocPrinter {
- public:
- /*!
- * \brief The constructor of DocPrinter
- *
- * \param options the option for printer
- */
- explicit DocPrinter(const PrinterConfig& options);
-
- virtual ~DocPrinter() = default;
-
- /*!
- * \brief Append a doc into the final content
- *
- * \param doc the Doc to be printed
- *
- * \sa GetString
- */
- void Append(const Doc& doc);
-
- /*!
- * \brief Append a doc to the final content
- *
- * \param doc Doc to be printed
- * \param path_to_underline ffi::Object path to be underlined
- *
- * \sa GetString
- */
- void Append(const Doc& doc, const PrinterConfig& cfg);
-
- /*!
- * \brief Get the printed string of all Doc appended
- *
- * The content of each Doc in the returned string will
- * appear in the same order as they are appended.
- *
- * \sa Append
- */
- ffi::String GetString() const;
-
- /*!
- * \brief Return the visible source path selected for each requested
- * underline path.
- */
- ffi::Array<ffi::Optional<AccessPath>> GetVisiblePaths() const;
-
- protected:
- /*!
- * \brief Get the printed string
- *
- * It will dispatch to the PrintTypedDoc method based on
- * the actual type of Doc.
- *
- * \sa PrintTypedDoc
- */
- void PrintDoc(const Doc& doc);
-
- /*!
- * \brief Virtual method to print a LiteralDoc
- */
- 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
- */
- virtual void PrintTypedDoc(const IdDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an AttrAccessDoc
- */
- virtual void PrintTypedDoc(const AttrAccessDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an IndexDoc
- */
- virtual void PrintTypedDoc(const IndexDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an OperationDoc
- */
- virtual void PrintTypedDoc(const OperationDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a CallDoc
- */
- virtual void PrintTypedDoc(const CallDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a LambdaDoc
- */
- virtual void PrintTypedDoc(const LambdaDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ListDoc
- */
- virtual void PrintTypedDoc(const ListDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a TupleDoc
- */
- virtual void PrintTypedDoc(const TupleDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a DictDoc
- */
- virtual void PrintTypedDoc(const DictDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a SliceDoc
- */
- virtual void PrintTypedDoc(const SliceDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a StmtBlockDoc
- */
- virtual void PrintTypedDoc(const StmtBlockDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an AssignDoc
- */
- virtual void PrintTypedDoc(const AssignDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an IfDoc
- */
- virtual void PrintTypedDoc(const IfDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a WhileDoc
- */
- virtual void PrintTypedDoc(const WhileDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a BreakDoc
- */
- virtual void PrintTypedDoc(const BreakDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ContinueDoc
- */
- virtual void PrintTypedDoc(const ContinueDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ForDoc
- */
- virtual void PrintTypedDoc(const ForDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ScopeDoc
- */
- virtual void PrintTypedDoc(const ScopeDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an ExprStmtDoc
- */
- virtual void PrintTypedDoc(const ExprStmtDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print an AssertDoc
- */
- virtual void PrintTypedDoc(const AssertDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ReturnDoc
- */
- virtual void PrintTypedDoc(const ReturnDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a FunctionDoc
- */
- virtual void PrintTypedDoc(const FunctionDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a ClassDoc
- */
- virtual void PrintTypedDoc(const ClassDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a CommentDoc
- */
- virtual void PrintTypedDoc(const CommentDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a DocStringDoc
- */
- virtual void PrintTypedDoc(const DocStringDoc& doc) = 0;
-
- /*!
- * \brief Virtual method to print a OpCallDoc
- */
- virtual void PrintTypedDoc(const OpCallDoc& doc) = 0;
-
- /*!
- * \brief Increase the indent level of any content to be
- * printed after this call
- */
- void IncreaseIndent() { indent_ += options_->indent_spaces; }
-
- /*!
- * \brief Decrease the indent level of any content to be
- * printed after this call
- */
- void DecreaseIndent() { indent_ -= options_->indent_spaces; }
-
- /*!
- * \brief Add a new line into the output stream
- *
- * \sa output_
- */
- std::ostream& NewLine() {
- size_t start_pos = output_.tellp();
- output_ << "\n";
- line_starts_.push_back(output_.tellp());
- output_ << std::string(indent_, ' ');
- size_t end_pos = output_.tellp();
- underlines_exempted_.push_back({start_pos, end_pos});
- return output_;
- }
-
- /*!
- * \brief The output stream of printer
- *
- * All printed content will be stored in this stream and returned
- * when GetString is called.
- *
- * \sa GetString
- */
- std::ostringstream output_;
-
- /*! \brief Spans that we have already committed to underline exemption. */
- std::vector<ByteSpan> underlines_exempted_;
-
- private:
- void MarkSpan(const ByteSpan& span, const AccessPath& path);
-
- /*! \brief Options to customize certain aspects of the output */
- PrinterConfig options_;
-
- /*! \brief the current level of indent */
- int indent_ = 0;
-
- /*! \brief For each line in the output_, byte offset of its first character */
- std::vector<size_t> line_starts_;
-
- /*! \brief Path of the object that we would like to underline */
- ffi::Array<AccessPath> path_to_underline_;
-
- /*!
- * \brief Candidate spans to be underlined, until we find a better match.
- * (A better match is an object with a longer path that is still a prefix of path_to_underline_.)
- */
- std::vector<std::vector<ByteSpan>> current_underline_candidates_;
-
- /*! \brief Path length of the objects that are current candidates for underlining. */
- std::vector<int> current_max_path_depth_;
-
- /*! \brief Visible path selected by the current best underline candidate. */
- std::vector<ffi::Optional<AccessPath>> current_visible_paths_;
-
- /*! \brief Spans that we have already committed to underline. */
- std::vector<ByteSpan> underlines_;
-};
-
-} // namespace printer
-} // namespace script
-} // namespace tvm
-
-#endif // TVM_SCRIPT_PRINTER_DOC_PRINTER_BASE_DOC_PRINTER_H_
diff --git a/src/script/printer/doc_translator.cc b/src/script/printer/doc_translator.cc
new file mode 100644
index 0000000000..7cc3a0fb91
--- /dev/null
+++ b/src/script/printer/doc_translator.cc
@@ -0,0 +1,279 @@
+/*
+ * 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 "doc_translator.h"
+
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/serialization.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/base.h>
+#include <tvm/script/printer/doc.h>
+#include <tvm/script/printer/doc_translator.h>
+
+#include <algorithm>
+#include <cctype>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "dialect_prefix.h"
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+namespace {
+
+// The only mutable structural state. Identifier allocation lives with the
+// remapping environment; no statements, scopes, paths, or config live here.
+struct VariableEnvironment {
+ ffi::Dict<Var, IdDoc> mapped;
+ ffi::Dict<Var, IdDoc> implicit_defs;
+ std::unordered_set<std::string> allocated;
+};
+
+// Bindings and reservations stay stable within each function naming region.
+class TranslationEngine : public DocTranslatorObj {
+ public:
+ explicit TranslationEngine(ffi::Map<ffi::String, ffi::Any> config)
+ : DocTranslatorObj(VTable(), std::move(config)) {
+ for (const char* keyword :
+ {"range", "False", "None", "True", "and", "as", "assert", "async", "await",
+ "break", "class", "continue", "def", "del", "elif", "else", "except", "finally",
+ "for", "from", "global", "if", "import", "in", "is", "lambda", "nonlocal",
+ "not", "or", "pass", "raise", "return", "try", "while", "with", "yield"}) {
+ AllocId(keyword);
+ }
+ }
+
+ static IdDoc Allocate(DocTranslatorObj* d, ffi::AnyView value) {
+ auto* self = static_cast<TranslationEngine*>(d);
+ std::string base = value.cast<ffi::String>();
+ if (base.empty()) base = "v";
+ for (char& c : base)
+ if (!std::isalnum(static_cast<unsigned char>(c)) && c != '_') c = '_';
+ if (std::isdigit(static_cast<unsigned char>(base[0]))) base.insert(base.begin(), '_');
+ std::string name = base;
+ auto reserved = [&](const std::string& candidate) {
+ return std::any_of(self->var_scopes_.begin(), self->var_scopes_.end(),
+ [&](const auto& frame) { return frame.allocated.count(candidate); });
+ };
+ for (int i = 1; reserved(name); ++i) name = base + '_' + std::to_string(i);
+ self->var_scopes_.back().allocated.insert(name);
+ return IdDoc(name);
+ }
+
+ static IdDoc GetOrAllocate(DocTranslatorObj* d, ffi::AnyView value, bool explicit_def) {
+ auto* self = static_cast<TranslationEngine*>(d);
+ Var var = value.cast<Var>();
+ for (size_t i = self->var_scopes_.size(); i != 0; --i) {
+ auto& frame = self->var_scopes_[i - 1];
+ auto found = frame.mapped.find(var);
+ if (found == frame.mapped.end()) continue;
+ IdDoc id = (*found).second;
+ if (explicit_def) {
+ frame.implicit_defs.erase(var);
+ }
+ return id;
+ }
+ IdDoc id = Allocate(d, var->name);
+ auto& frame = self->var_scopes_.back();
+ frame.mapped.Set(var, id);
+ if (!explicit_def) frame.implicit_defs.Set(var, id);
+ return id;
+ }
+ static ffi::Dict<Var, IdDoc> ImplicitDefinitions(DocTranslatorObj* d) {
+ return static_cast<TranslationEngine*>(d)->var_scopes_.back().implicit_defs;
+ }
+ static void BeginScope(DocTranslatorObj* d) {
+ static_cast<TranslationEngine*>(d)->var_scopes_.emplace_back();
+ }
+ static void EndScope(DocTranslatorObj* d) {
+ auto* self = static_cast<TranslationEngine*>(d);
+ TVM_FFI_CHECK(self->var_scopes_.size() > 1, ValueError)
+ << "printer cannot pop the root variable scope";
+ auto& local = self->var_scopes_.back();
+ auto& parent = self->var_scopes_[self->var_scopes_.size() - 2];
+ // Unbound references remain captures of the enclosing naming region.
+ for (const auto& [var, id] : local.implicit_defs) {
+ parent.mapped.Set(var, id);
+ parent.implicit_defs.Set(var, id);
+ parent.allocated.insert(id->name);
+ }
+ self->var_scopes_.pop_back();
+ }
+
+ private:
+ template <typename Signature, auto Function>
+ static auto Adapter() {
+ using View = ffi::reflection::NativeFunctionView<Signature>;
+ return reinterpret_cast<typename View::ABIType>(
+ ffi::Any(View::template FromNative<Function>()).template cast<void*>());
+ }
+ static const DocTranslatorVTable* VTable() {
+ static const DocTranslatorVTable table{
+ Adapter<ffi::Optional<ExprDoc>(DocTranslatorObj*, ffi::AnyView, const ffi::Object*),
+ [](DocTranslatorObj* d, ffi::AnyView value, const ffi::Object* destination) {
+ return d->DefaultTranslate(value, destination);
+ }>(),
+ Adapter<IdDoc(DocTranslatorObj*, ffi::AnyView), &Allocate>(),
+ Adapter<IdDoc(DocTranslatorObj*, ffi::AnyView, bool), &GetOrAllocate>(),
+ Adapter<ffi::Dict<Var, IdDoc>(DocTranslatorObj*), &ImplicitDefinitions>(),
+ Adapter<void(DocTranslatorObj*), &BeginScope>(),
+ Adapter<void(DocTranslatorObj*), &EndScope>()};
+ return &table;
+ }
+ /*! \brief Owned variable environments, with a persistent root scope. */
+ std::vector<VariableEnvironment> var_scopes_{1};
+};
+
+details::TranslationResult TranslateComplete(ffi::AnyView obj, TranslationEngine* engine,
+ ffi::Map<ffi::String, ffi::Any> extra_config) {
+ auto binding_names = engine->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ // Reserve the same registered or explicitly configured aliases used by rendering.
+ for (const auto& [key, fallback] : GetDialectPrefixes()) {
+ if (!extra_config.count(key)) engine->AllocId(fallback);
+ }
+ for (const auto& [key, value] : extra_config) {
+ std::string option = key;
+ if (option.size() >= 7 && option.compare(option.size() - 7, 7, ".prefix") == 0)
+ engine->AllocId(value.cast<ffi::String>());
+ }
+ for (const ffi::String& name : binding_names) engine->AllocId(name);
+ if (binding_names.empty()) {
+ if (auto name = engine->GetExtraConfig<ffi::Optional<ffi::String>>("script.module_name",
+ std::nullopt)) {
+ engine->AllocId(name.value());
+ }
+ }
+ // Allocate generated names after all configured names. A dialect or module
+ // named tvm/metadata must not shadow the loader or metadata registry.
+ IdDoc metadata_id = engine->AllocId("metadata");
+ IdDoc metadata_loader = engine->AllocId("tvm");
+ engine->ExchangeExtraState("ir.metadata_id", metadata_id);
+ auto docs = engine->WithDocScope([&]() {
+ ffi::Optional<ExprDoc> expression = engine->Translate(obj);
+ if (expression.has_value()) {
+ if (!engine->GetExtraConfig<bool>("script.verbose_expr", true))
+ engine->CurrentScopeDocs().clear();
+ engine->Emit(ExprStmtDoc(expression.value()), obj.as<ffi::ObjectRef>());
+ }
+ if (!binding_names.empty() && !engine->CurrentScopeDocs().empty()) {
+ Doc root = engine->CurrentScopeDocs().back();
+ ffi::String name = binding_names.back();
+ if (auto function = root.as<FunctionDoc>()) {
+ function.value()->name = IdDoc(name);
+ } else if (auto module = root.as<ClassDoc>()) {
+ module.value()->name = IdDoc(name);
+ }
+ }
+ });
+ auto root_declarations = engine->WithDocScope([&]() {
+ auto pending = engine->GetImplicitDefs();
+ std::vector<std::pair<Var, IdDoc>> remaining(pending.begin(), pending.end());
+ std::sort(remaining.begin(), remaining.end(),
+ [](const auto& a, const auto& b) { return a.second->name < b.second->name; });
+ for (const auto& [var, _] : remaining) {
+ if (!pending.count(var)) continue;
+ if (auto rhs = engine->Translate(var, var.get())) {
+ engine->Emit(AssignDoc(engine->VarGetOrAllocId(var, true), rhs.value(), std::nullopt), var);
+ }
+ }
+ });
+ if (!root_declarations.empty()) {
+ for (const Doc& doc : docs) root_declarations.push_back(doc);
+ docs = root_declarations;
+ }
+ using MetadataMap = ffi::Dict<ffi::String, ffi::List<ffi::Any>>;
+ auto metadata = engine->GetOrCreateExtraState<MetadataMap>("ir.metadata_map");
+ ffi::Array<StmtDoc> body;
+ if (!metadata.empty() && engine->GetExtraConfig<bool>("script.show_meta", true)) {
+ // Serialize the whole registry as one graph so aliases shared across
+ // entries survive. Immutable maps/arrays preserve the loader's contract.
+ ffi::Map<ffi::String, ffi::Any> serialized;
+ std::vector<ffi::String> keys;
+ for (const auto& [key, entries] : metadata) keys.push_back(key);
+ std::sort(keys.begin(), keys.end());
+ for (const auto& key : keys) {
+ auto entries = metadata[key];
+ serialized.Set(key, ffi::Array<ffi::Any>(entries.begin(), entries.end()));
+ }
+ ffi::String json = ffi::json::Stringify(
+ ffi::ToJSONGraph(serialized, ffi::json::Object{{"tvm_version", TVM_VERSION}}), 2);
+ body.push_back(AssignDoc(IdDoc(metadata_id->name),
+ IdDoc(metadata_loader->name)
+ ->Attr("ir")
+ ->Attr("load_json")
+ ->Call({LiteralDoc::Str(json, std::nullopt)}),
+ std::nullopt));
+ }
+ for (const auto& stmt : ToStmtDocArray(docs)) body.push_back(stmt);
+ if (!metadata.empty() && !engine->GetExtraConfig<bool>("script.show_meta", true)) {
+ body.push_back(
+ CommentDoc("Metadata omitted. Use show_meta=True in script() method to show it."));
+ }
+ return {StmtBlockDoc(body), engine->GetDocOrigins(),
+ engine->GetOrCreateExtraState<bool>("script.future_annotations"), metadata_loader->name,
+ !metadata.empty() && engine->GetExtraConfig<bool>("script.show_meta", true)};
+}
+
+} // namespace
+
+namespace details {
+TranslationResult TranslateWithOptions(ffi::AnyView ir,
+ ffi::Map<ffi::String, ffi::Any> extra_config) {
+ auto engine = ffi::make_object<TranslationEngine>(extra_config);
+ return TranslateComplete(ir, engine.get(), std::move(extra_config));
+}
+} // namespace details
+
+Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_origins,
+ ffi::Map<ffi::String, ffi::Any> extra_config) {
+ auto result = details::TranslateWithOptions(ir, std::move(extra_config));
+ if (doc_origins) *doc_origins = std::move(result.origins);
+ return result.doc;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = ffi::reflection;
+ refl::ObjectDef<DocTranslatorObj>();
+ refl::EnsureTypeAttrColumn(kDocTranslate);
+ refl::EnsureTypeAttrColumn(kTensorLoadDocTranslate);
+ refl::GlobalDef()
+ .def("script.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.Translate", &DocTranslatorObj::Translate)
+ .def_method("script.printer.Emit", &DocTranslatorObj::Emit)
+ .def_method("script.printer.VarGetOrAllocId", &DocTranslatorObj::VarGetOrAllocId)
+ .def_method("script.printer.GetImplicitDefs", &DocTranslatorObj::GetImplicitDefs)
+ .def_method("script.printer.AllocId", &DocTranslatorObj::AllocId)
+ .def_method("script.printer.BeginVarScope", &DocTranslatorObj::BeginVarScope)
+ .def_method("script.printer.EndVarScope", &DocTranslatorObj::EndVarScope)
+ .def_method("script.printer.BeginDocScope", &DocTranslatorObj::BeginDocScope)
+ .def_method("script.printer.EndDocScope", &DocTranslatorObj::EndDocScope)
+ .def("script.printer.WithDocScope", [](const DocTranslator& d, const ffi::Function& fn) {
+ return d->WithDocScope([&]() { fn(); });
+ });
+}
+
+} // namespace printer
+} // namespace script
+} // namespace tvm
diff --git a/src/script/printer/doc_translator.h b/src/script/printer/doc_translator.h
new file mode 100644
index 0000000000..54a85f431a
--- /dev/null
+++ b/src/script/printer/doc_translator.h
@@ -0,0 +1,57 @@
+/*
+ * 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_DOC_TRANSLATOR_PRIVATE_H_
+#define TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_PRIVATE_H_
+
+#include <tvm/script/printer/doc_translator.h>
+
+namespace tvm {
+namespace script {
+namespace printer {
+namespace details {
+
+/*! \brief Complete translated body, origins and explicit header facts. */
+struct TranslationResult {
+ /*! \brief Complete body, including declarations and metadata. */
+ Doc doc;
+ /*! \brief Original IR objects associated with generated Docs. */
+ ffi::Dict<Doc, ffi::ObjectRef> origins;
+ /*! \brief Whether the script needs postponed annotation evaluation. */
+ bool future_annotations;
+ /*! \brief Allocated import name for the metadata loader. */
+ ffi::String metadata_loader;
+ /*! \brief Whether the body uses the metadata loader. */
+ bool imports_metadata;
+};
+
+/*!
+ * \brief Translate an input into its complete Doc representation.
+ * \param ir The input to translate.
+ * \param extra_config Read-only options for this translation.
+ * \return Complete body, origins and header facts.
+ */
+TranslationResult TranslateWithOptions(ffi::AnyView ir,
+ ffi::Map<ffi::String, ffi::Any> extra_config);
+
+} // namespace details
+} // namespace printer
+} // namespace script
+} // namespace tvm
+
+#endif // TVM_SCRIPT_PRINTER_DOC_TRANSLATOR_PRIVATE_H_
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index bc60d69694..24c758e3a2 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -16,196 +16,209 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ir/type.h>
+#include <tvm/ir/module.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/type.h>
+#include <tvm/tirx/function.h>
-#include "./utils.h"
+#include <algorithm>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "../dialect_prefix.h"
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-TVM_FFI_STATIC_INIT_BLOCK() { IRFrameNode::RegisterReflection(); }
-
-struct SortableFunction {
- int priority;
- GlobalVar gv;
- BaseFunc func;
-
- explicit SortableFunction(const std::pair<GlobalVar, BaseFunc>& obj)
- : priority(0), gv(obj.first), func(obj.second) {
- if (gv->name_hint == "main") {
- priority = 1000;
- } else if (func.defined()) {
- if (func->GetTypeKey() == "tirx.PrimFunc") {
- priority = 1;
- } else if (func->GetTypeKey() == "relax.expr.ExternFunc") {
- priority = 2;
- } else if (func->GetTypeKey() == "relax.expr.Function") {
- priority = 3;
- } else {
- TVM_FFI_THROW(TypeError) << "TVMScript cannot print functions of type: "
- << func->GetTypeKey();
- }
- } else {
- // PrimFuncPass may leave undefined GlobalVar slots when transforming
- // this function (see tirx/ir/transform.cc); this transient state may
- // be encountered during the internal call Dump(mod) executed in
- // PrimFuncPass during debugging.
- priority = 999;
- LOG(INFO) << "Function " << gv->name_hint << " is undefined";
- }
- }
+using GlobalInfoMap = ffi::Dict<ffi::String, ffi::List<GlobalInfo>>;
- bool operator<(const SortableFunction& other) const {
- if (this->priority != other.priority) {
- return this->priority < other.priority;
- }
- return this->gv->name_hint < other.gv->name_hint;
- }
-};
+namespace {
-ffi::Optional<ffi::String> GetDynamicDeclarationName(const StmtDoc& stmt) {
- const auto* assign = stmt.as<AssignDocNode>();
- if (assign == nullptr || !assign->rhs.has_value()) {
- return std::nullopt;
+class GlobalInfoScope {
+ public:
+ GlobalInfoScope(DocTranslatorObj* d, const IRModuleNode* mod) : d_(d) {
+ GlobalInfoMap infos;
+ for (const auto& [key, entries] : mod->global_infos) {
+ infos.Set(key, ffi::List<GlobalInfo>(entries.begin(), entries.end()));
+ }
+ saved_ = d_->ExchangeExtraState("ir.global_info_map", infos);
}
- const auto* lhs = assign->lhs.as<IdDocNode>();
- const auto* call = assign->rhs.value().as<CallDocNode>();
- if (lhs == nullptr || call == nullptr) {
- return std::nullopt;
+ ~GlobalInfoScope() noexcept {
+ if (d_) {
+ try {
+ Close();
+ } catch (...) { /* Preserve translation failure. */
+ }
+ }
}
- const auto* callee = call->callee.as<AttrAccessDocNode>();
- if (callee == nullptr || callee->name != "dynamic") {
- return std::nullopt;
+ void Close() {
+ DocTranslatorObj* d = std::exchange(d_, nullptr);
+ d->ExchangeExtraState("ir.global_info_map", std::move(saved_));
}
- return lhs->name;
-}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<IRModule>(
- "", [](IRModule mod, AccessPath p, IRDocsifier d) -> Doc {
- std::vector<SortableFunction> functions;
- ffi::Array<StmtDoc> dynamic_decls;
- std::unordered_set<ffi::String> declared_dynamic_names;
- for (const auto& kv : mod->functions) {
- functions.push_back(SortableFunction(kv));
- }
- std::sort(functions.begin(), functions.end());
- With<IRFrame> f(d);
- (*f)->AddDispatchToken(d, "ir");
- IdDoc module_doc = d->Define(mod, f(), GetBindingName(d).value_or("Module"));
- (*f)->global_infos = &mod->global_infos;
- if (!mod->attrs->dict.empty()) {
- (*f)->stmts.push_back(
- ExprStmtDoc(IR(d, "module_attrs") //
- ->Call({d->AsDoc<ExprDoc>(mod->attrs, p->Attr("attrs"))})));
- }
- if (mod->global_infos.defined() && !mod->global_infos.empty()) {
- (*f)->stmts.push_back(ExprStmtDoc(
- IR(d, "module_global_infos") //
- ->Call({d->AsDoc<ExprDoc>(mod->global_infos, p->Attr("global_infos"))})));
- }
- // Declare GlobalVars first
- IdDoc module_alias =
- d->cfg->module_alias.empty() ? module_doc : IdDoc(d->cfg->module_alias);
- for (const auto& entry : functions) {
- const GlobalVar& gv = entry.gv;
- d->Define(gv, f(), [=]() {
- return d->AsDoc<ExprDoc>(mod, p->Attr("global_vars"))->Attr(gv->name_hint);
- });
- }
- // Print functions
- for (const auto& entry : functions) {
- const GlobalVar& gv = entry.gv;
- const BaseFunc& base_func = entry.func;
- d->cfg->binding_names.push_back(gv->name_hint);
- Doc doc = d->AsDoc(base_func, p->Attr("functions")->MapItem(gv));
- d->cfg->binding_names.pop_back();
- if (const auto* stmt_block = doc.as<StmtBlockDocNode>()) {
- for (const StmtDoc& stmt : stmt_block->stmts) {
- if (ffi::Optional<ffi::String> name = GetDynamicDeclarationName(stmt)) {
- if (!declared_dynamic_names.count(name.value())) {
- declared_dynamic_names.insert(name.value());
- dynamic_decls.push_back(stmt);
- }
+ private:
+ DocTranslatorObj* d_;
+ std::optional<ffi::Any> saved_;
+};
+
+ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* mod =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IRModuleNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ auto names = d->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ ffi::String module_name =
+ names.empty() ? d->GetExtraConfig<ffi::String>("script.module_name", "Module") : names.back();
+ d->AllocId(module_name);
+ GlobalInfoScope global_infos(d, mod);
+ std::vector<std::pair<GlobalVar, BaseFunc>> functions(mod->functions.begin(),
+ mod->functions.end());
+ auto rank = [](const BaseFunc& func) {
+ if (func.as<relax::ExternFuncNode>()) return 0;
+ if (func.as<tirx::PrimFuncNode>()) return 1;
+ return 2;
+ };
+ std::sort(functions.begin(), functions.end(), [&](const auto& a, const auto& b) {
+ int left = rank(a.second);
+ int right = rank(b.second);
+ return left == right ? a.first->name_hint < b.first->name_hint : left < right;
+ });
+ auto body = d->WithDocScope([&]() {
+ if (!mod->attrs->dict.empty()) {
+ ffi::Array<ExprDoc> keys;
+ ffi::Array<ExprDoc> values;
+ std::vector<std::pair<ffi::String, ffi::Any>> attrs(mod->attrs->dict.begin(),
+ mod->attrs->dict.end());
+ std::sort(attrs.begin(), attrs.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ for (const auto& [key, value] : attrs) {
+ keys.push_back(LiteralDoc::Str(key, std::nullopt));
+ values.push_back(AnyValue(d, value));
+ }
+ d->Emit(ExprStmtDoc(NamespaceDoc("ir")->Attr("module_attrs")->Call({DictDoc(keys, values)})),
+ ffi::GetRef<ffi::ObjectRef>(mod));
+ }
+ if (!mod->global_infos.empty()) {
+ ffi::Array<ExprDoc> keys;
+ ffi::Array<ExprDoc> values;
+ std::vector<std::pair<ffi::String, ffi::Array<GlobalInfo>>> infos(mod->global_infos.begin(),
+ mod->global_infos.end());
+ std::sort(infos.begin(), infos.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ for (const auto& [key, entries] : infos) {
+ ffi::Array<ExprDoc> items;
+ for (const GlobalInfo& entry : entries) {
+ ExprDoc item = AnyValue(d, entry);
+ if (key == "mesh") {
+ if (auto mesh = item.as<CallDoc>(); mesh && mesh.value()->args.size() == 2) {
+ if (auto range = mesh.value()->args[1].as<CallDoc>()) {
+ range.value()->callee = NamespaceDoc("ir")->Attr("Range");
}
}
- (*f)->stmts.push_back(stmt_block->stmts.back());
- (*f)->stmts.back()->source_paths = std::move(doc->source_paths);
- } else if (auto stmt = doc.as<StmtDoc>()) {
- (*f)->stmts.push_back(stmt.value());
- } else if (auto func = doc.as<FunctionDoc>()) {
- (*f)->stmts.push_back(func.value());
- } else if (auto expr = doc.as<ExprDoc>()) {
- ExprDoc lhs = IdDoc(gv->name_hint);
- AssignDoc assignment(lhs, expr.value(), std::nullopt);
- (*f)->stmts.push_back(assignment);
- } else {
- TVM_FFI_THROW(TypeError)
- << "Expected IRModule to only contain functions, "
- << " but mod[" << gv->name_hint << "] with type " << base_func->GetTypeKey()
- << " produced Doc type of " << doc->GetTypeKey();
}
+ items.push_back(item);
}
- ClassDoc class_doc(module_doc, {IR(d, "ir_module")}, (*f)->stmts);
- if (dynamic_decls.empty()) {
- return HeaderWrapper(d, class_doc);
- }
- dynamic_decls.push_back(class_doc);
- return HeaderWrapper(d, StmtBlockDoc(dynamic_decls));
- });
+ keys.push_back(LiteralDoc::Str(key, std::nullopt));
+ values.push_back(ListDoc(items));
+ }
+ d->Emit(ExprStmtDoc(
+ NamespaceDoc("ir")->Attr("module_global_infos")->Call({DictDoc(keys, values)})),
+ ffi::GetRef<ffi::ObjectRef>(mod));
+ }
+ for (const auto& [gv, func] : functions) {
+ if (func.as<relax::ExternFuncNode>()) {
+ d->Emit(AssignDoc(IdDoc(gv->name_hint), d->Translate(func).value(), std::nullopt),
+ ffi::GetRef<ffi::ObjectRef>(func.get()));
+ continue;
+ }
+ TVM_FFI_CHECK(func.as<tirx::PrimFuncNode>() || func.as<relax::FunctionNode>(), TypeError)
+ << "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);
+ if (auto symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
+ symbol && symbol.value() != gv->name_hint) {
+ ExprDoc attr = func.as<tirx::PrimFuncNode>() ? NamespaceDoc("tirx")->Attr("func_attr")
+ : NamespaceDoc("relax")->Attr("func_attr");
+ doc->body.insert(doc->body.begin(),
+ ExprStmtDoc(attr->Call(
+ {DictDoc({LiteralDoc::Str(tvm::attr::kGlobalSymbol, std::nullopt)},
+ {LiteralDoc::Str(symbol.value(), std::nullopt)})})));
+ }
+ }
+ });
+ global_infos.Close();
+ d->Emit(
+ ClassDoc(IdDoc(module_name), {NamespaceDoc("ir")->Attr("ir_module")}, ToStmtDocArray(body)),
+ ffi::GetRef<ffi::ObjectRef>(mod));
+ return std::nullopt;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<DictAttrs>(
- "", [](DictAttrs attrs, AccessPath p, IRDocsifier d) -> Doc {
- return d->AsDoc(attrs->dict, p->Attr("dict"));
- });
+ RegisterDialectPrefix("ir.prefix", "I");
+ ffi::reflection::TypeAttrDef<IRModuleNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&IRModuleDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<GlobalVar>(
- "", [](GlobalVar gv, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "GlobalVar")->Call({LiteralDoc::Str(gv->name_hint, p->Attr("name_hint"))});
- });
+ffi::Optional<ExprDoc> DictAttrsDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* attrs =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DictAttrsNode>(input);
+ return AnyValue(d, attrs->dict);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Op>("", [](Op op, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "Op")->Call({LiteralDoc::Str(op->name, p->Attr("name"))});
- });
+ ffi::reflection::TypeAttrDef<DictAttrsNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DictAttrsDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> GlobalVarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* var =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const GlobalVarNode>(input);
+ return GlobalReference(d, var->name_hint);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<FuncType>(
- "", [](FuncType func_type, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "FuncType")
- ->Call({
- d->AsDoc<ExprDoc>(func_type->arg_types, p->Attr("arg_types")),
- d->AsDoc<ExprDoc>(func_type->ret_type, p->Attr("ret_type")),
- });
- });
+ ffi::reflection::TypeAttrDef<GlobalVarNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&GlobalVarDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> FuncTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FuncTypeNode>(input);
+ return TypeValue(d, ffi::GetRef<FuncType>(ty), false);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Range>(
- "ir", [](Range range, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "Range")
- ->Call({
- d->AsDoc<ExprDoc>(range->min, p->Attr("min")),
- d->AsDoc<ExprDoc>(range->extent + range->min, p->Attr("extent")),
- });
- });
+ ffi::reflection::TypeAttrDef<FuncTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&FuncTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> RangeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* range =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const RangeNode>(input);
+ ExprDoc min = d->Translate(range->min).value();
+ ExprDoc end = d->Translate(range->min + range->extent).value();
+ return NamespaceDoc("ir")->Attr("Range")->Call({min, end});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVMScriptPrinter::Register<GlobalVarNode>(ReprPrintIR);
- TVMScriptPrinter::Register<DictAttrsNode>(ReprPrintIR);
- TVMScriptPrinter::Register<FuncTypeNode>(ReprPrintIR);
- TVMScriptPrinter::Register<RangeNode>(ReprPrintIR);
- TVMScriptPrinter::Register<IRModuleNode>(ReprPrintIR);
+ ffi::reflection::TypeAttrDef<RangeNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&RangeDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index 400cb4fe5d..bc98604d75 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -17,102 +17,195 @@
* under the License.
*/
#include <tvm/ffi/container/shape.h>
+#include <tvm/ffi/extra/module.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/relax/distributed/type.h>
+#include <tvm/relax/global_info.h>
#include <tvm/runtime/tensor.h>
+#include <tvm/target/target.h>
-#include "./utils.h"
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <limits>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<DataTypeImm>(
- "", [](DataTypeImm n, AccessPath p, IRDocsifier d) -> Doc {
- return TIR(d, "dtype")->Call({LiteralDoc::DataType(n->value, p->Attr("value"))});
- });
- IRDocsifier::vtable().set_dispatch<DataTypeImm>(
- "ir", [](DataTypeImm n, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "dtype")->Call({LiteralDoc::DataType(n->value, p->Attr("value"))});
- });
- IRDocsifier::vtable().set_dispatch<GenericConst>(
- "", [](GenericConst n, AccessPath p, IRDocsifier d) -> Doc {
- if (auto dtype = n->value.as<DLDataType>()) {
- return TIR(d, "dtype")->Call({LiteralDoc::DataType(*dtype, p->Attr("value"))});
- }
- if (n->value.as<runtime::Tensor>()) {
- // Tensor-valued constants are Relax constants and keep the Relax spelling.
- return IRDocsifier::vtable()("relax", n, p, d);
- }
- return d->AddMetadata(n);
- });
+namespace {
+
+ffi::Optional<ExprDoc> DataTypeImmDocTranslate(DocTranslatorObj*, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* imm =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DataTypeImmNode>(input);
+ return NamespaceDoc("tirx")->Attr("dtype")->Call(
+ {LiteralDoc::DataType(imm->value, std::nullopt)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<GenericConst>(
- "ir", [](GenericConst n, AccessPath p, IRDocsifier d) -> Doc {
- if (auto dtype = n->value.as<DLDataType>()) {
- return IR(d, "dtype")->Call({LiteralDoc::DataType(*dtype, p->Attr("value"))});
- }
- return IRDocsifier::vtable()("", n, p, d);
- });
+ ffi::reflection::TypeAttrDef<DataTypeImmNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&DataTypeImmDocTranslate>());
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<ffi::Array<Any>>( //
- "", [](ffi::Array<Any> array, AccessPath p, IRDocsifier d) -> Doc {
- int n = array.size();
- ffi::Array<ExprDoc> results;
- results.reserve(n);
- for (int i = 0; i < n; ++i) {
- results.push_back(d->AsDoc<ExprDoc>(array[i], p->ArrayItem(i)));
- }
- return ListDoc(results);
- });
+ffi::Optional<ExprDoc> GenericConstDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* constant =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const GenericConstNode>(input);
+ if (auto dtype = constant->value.as<DLDataType>()) {
+ return LiteralDoc::DataType(dtype.value(), std::nullopt);
+ }
+ auto tensor = constant->value.as<runtime::Tensor>();
+ if (!tensor.has_value()) return AddMetadata(d, ffi::GetRef<GenericConst>(constant));
+ const runtime::Tensor& data = tensor.value();
+ DLDataType dtype = data.DataType();
+ auto print_const = [&](ExprDoc value) -> ExprDoc {
+ if (constant->ty.as<relax::distributed::DTensorTypeNode>()) {
+ return NamespaceDoc("relax")->Attr("dist")->Attr("const")->Call(
+ {value, d->Translate(constant->ty).value()});
+ }
+ return NamespaceDoc("relax")->Attr("const")->Call(
+ {value, LiteralDoc::DataType(dtype, std::nullopt)});
+ };
+ if (data->ndim != 0) return AddMetadata(d, ffi::GetRef<GenericConst>(constant));
+ if (data->device.device_type != kDLCPU)
+ return AddMetadata(d, ffi::GetRef<GenericConst>(constant));
+ if (dtype.lanes != 1 || dtype.bits == 0 || dtype.bits > 64 || dtype.bits % 8 != 0) {
+ return AddMetadata(d, ffi::GetRef<GenericConst>(constant));
+ }
+ alignas(double) uint8_t scalar_bytes[sizeof(double)]{};
+ data.CopyToBytes(scalar_bytes, dtype.bits / 8);
+ auto read_scalar = [&](auto value) {
+ std::memcpy(&value, scalar_bytes, sizeof(value));
+ return value;
+ };
+ ExprDoc scalar(ffi::UnsafeInit{});
+ if (dtype == (DLDataType{kDLInt, 8, 1})) {
+ scalar = LiteralDoc::Int(read_scalar(int8_t{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLInt, 16, 1})) {
+ scalar = LiteralDoc::Int(read_scalar(int16_t{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLInt, 32, 1})) {
+ scalar = LiteralDoc::Int(read_scalar(int32_t{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLInt, 64, 1})) {
+ scalar = LiteralDoc::Int(read_scalar(int64_t{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLFloat, 16, 1})) {
+ uint16_t bits = read_scalar(uint16_t{});
+ uint16_t exponent = (bits >> 10) & 31;
+ uint16_t fraction = bits & 1023;
+ double value = exponent == 31
+ ? (fraction ? std::numeric_limits<double>::quiet_NaN()
+ : std::numeric_limits<double>::infinity())
+ : (exponent == 0 ? std::ldexp(static_cast<double>(fraction), -24)
+ : std::ldexp(static_cast<double>(fraction | 1024),
+ static_cast<int>(exponent) - 25));
+ if (bits & 0x8000) value = -value;
+ scalar = LiteralDoc::Float(value, std::nullopt);
+ } else if (dtype == (DLDataType{kDLFloat, 32, 1})) {
+ scalar = LiteralDoc::Float(read_scalar(float{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLFloat, 64, 1})) {
+ scalar = LiteralDoc::Float(read_scalar(double{}), std::nullopt);
+ } else if (dtype == (DLDataType{kDLBool, 8, 1})) {
+ scalar = LiteralDoc::Boolean(read_scalar(uint8_t{}), std::nullopt);
+ } else {
+ return AddMetadata(d, ffi::GetRef<GenericConst>(constant));
+ }
+ return print_const(scalar);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<ffi::Map<Any, Any>>( //
- "", [](ffi::Map<Any, Any> dict, AccessPath p, IRDocsifier d) -> Doc {
- using POO = std::pair<Any, Any>;
- std::vector<POO> items{dict.begin(), dict.end()};
- bool is_str_map = true;
- for (const auto& kv : items) {
- if (!kv.first.as<ffi::String>()) {
- is_str_map = false;
- break;
- }
- }
- if (is_str_map) {
- std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
- return lhs.first.as_or_throw<ffi::String>() < rhs.first.as_or_throw<ffi::String>();
- });
- }
- int n = dict.size();
- ffi::Array<ExprDoc> ks;
- ffi::Array<ExprDoc> vs;
- ks.reserve(n);
- vs.reserve(n);
- for (int i = 0; i < n; ++i) {
- ks.push_back(d->AsDoc<ExprDoc>(items[i].first, p->MapItemMissing(items[i].first)));
- vs.push_back(d->AsDoc<ExprDoc>(items[i].second, p->MapItem(items[i].first)));
- }
- return DictDoc(ks, vs);
- });
+ ffi::reflection::TypeAttrDef<GenericConstNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&GenericConstDocTranslate>());
+}
+
+} // namespace
+
+ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value) {
+ if (value.type_index() == ffi::TypeIndex::kTVMFFINone) return LiteralDoc::None(std::nullopt);
+ if (value.as<ffi::Module>() || value.as<runtime::Tensor>()) return AddMetadata(d, value);
+ if (auto array = value.as<ffi::Array<ffi::Any>>()) {
+ ffi::Array<ExprDoc> elements;
+ for (const auto& element : array.value()) elements.push_back(AnyValue(d, element));
+ return ListDoc(elements);
+ }
+ if (auto map = value.as<ffi::Map<ffi::Any, ffi::Any>>()) {
+ std::vector<std::pair<ffi::String, ffi::Any>> entries;
+ for (const auto& [key, item] : map.value()) {
+ entries.emplace_back(key.as_or_throw<ffi::String>(), item);
+ }
+ std::sort(entries.begin(), entries.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ ffi::Array<ExprDoc> keys;
+ ffi::Array<ExprDoc> values;
+ for (const auto& [key, item] : entries) {
+ keys.push_back(LiteralDoc::Str(key, std::nullopt));
+ values.push_back(AnyValue(d, item));
+ }
+ return DictDoc(keys, values);
+ }
+ if (auto string = value.as<ffi::String>()) {
+ return LiteralDoc::Str(string.value(), std::nullopt);
+ }
+ if (auto boolean = value.as<bool>()) {
+ return LiteralDoc::Boolean(boolean.value(), std::nullopt);
+ }
+ if (auto integer = value.as<int64_t>()) {
+ return LiteralDoc::Int(integer.value(), std::nullopt);
+ }
+ if (auto floating = value.as<double>()) {
+ return LiteralDoc::Float(floating.value(), std::nullopt);
+ }
+ if (auto dtype = value.as<DLDataType>()) {
+ return LiteralDoc::DataType(dtype.value(), std::nullopt);
+ }
+ if (const auto* dtype = value.as<DataTypeImmNode>()) {
+ return NamespaceDoc("ir")->Attr("dtype")->Call(
+ {LiteralDoc::DataType(dtype->value, std::nullopt)});
+ }
+ if (auto vdevice = value.as<relax::VDevice>()) {
+ return NamespaceDoc("relax")->Attr("vdevice")->Call(
+ {NamespaceDoc("tirx")->Attr("target")->Call(
+ {AnyValue(d, vdevice.value()->target->ToConfig())})},
+ {"vdevice_id", "memory_scope"},
+ {LiteralDoc::Int(vdevice.value()->vdevice_id, std::nullopt),
+ LiteralDoc::Str(vdevice.value()->memory_scope, std::nullopt)});
+ }
+ if (const auto* string = value.as<StringImmNode>()) {
+ // Any-valued fields must retain the IR object; Expr-valued fields convert literals.
+ ExprDoc doc =
+ NamespaceDoc("ir")->Attr("StringImm")->Call({LiteralDoc::Str(string->value, std::nullopt)});
+ d->RecordOrigin(doc, ffi::GetRef<StringImm>(string));
+ return doc;
+ }
+ return d->Translate(value).value();
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> ShapeDocTranslate(DocTranslatorObj*, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* shape =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ffi::ShapeObj>(input);
+ ffi::Array<ExprDoc> dimensions;
+ for (int64_t value : ffi::GetRef<ffi::Shape>(shape)) {
+ dimensions.push_back(LiteralDoc::Int(value, std::nullopt));
+ }
+ return IdDoc("Shape")->Call(dimensions);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<ffi::Shape>(
- "", [](ffi::Shape n, AccessPath n_p, IRDocsifier d) -> Doc {
- int s = n.size();
- ffi::Array<ExprDoc> results;
- results.reserve(s);
- for (int i = 0; i < s; ++i) {
- results.push_back(d->AsDoc<ExprDoc>(IntImm::Int32(n[i]), n_p->ArrayItem(i)));
- }
- return TupleDoc(results);
- });
+ ffi::reflection::TypeAttrDef<ffi::ShapeObj>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ShapeDocTranslate>());
}
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/ir/prim_expr.cc b/src/script/printer/ir/prim_expr.cc
index 649c2c5826..9bf33fd3ed 100644
--- a/src/script/printer/ir/prim_expr.cc
+++ b/src/script/printer/ir/prim_expr.cc
@@ -17,185 +17,331 @@
* under the License.
*/
#include <tvm/ir/prim/op.h>
+#include <tvm/ir/prim/vector_expr.h>
-#include "./utils.h"
+#include <cstring>
+#include <optional>
+
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
+
+namespace {
+
+template <typename T, OperationDocNode::Kind kind, PrimExpr (*operation)(PrimExpr, PrimExpr, Span)>
+ffi::Optional<ExprDoc> BinaryOpDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const T>(input);
+ ExprDoc a = d->Translate(node->a).value();
+ ExprDoc b = d->Translate(node->b).value();
+ // Replay operator overloads before choosing sugar. Even one constant can
+ // simplify an explicit node, while two plain Python literals fold earlier.
+ if (!(a.as<LiteralDocNode>() && b.as<LiteralDocNode>())) {
+ try {
+ PrimExpr replay = operation(node->a, node->b, Span());
+ if (const auto* result = replay.template as<T>();
+ result && result->a.same_as(node->a) && result->b.same_as(node->b)) {
+ return OperationDoc(kind, {a, b});
+ }
+ } catch (const ffi::Error&) {
+ // A constructor can preserve operands rejected by a simplifying helper.
+ }
+ }
+ return NamespaceDoc("tirx")->Attr(std::strrchr(T::_type_key, '.') + 1)->Call({a, b});
+}
+
+template <typename T, OperationDocNode::Kind kind, PrimExpr (*operation)(PrimExpr, Span)>
+ExprDoc UnaryOpDocTranslate(DocTranslatorObj* d, const T* node) {
+ ExprDoc value = d->Translate(node->a).value();
+ if (!value.as<LiteralDocNode>()) {
+ try {
+ PrimExpr replay = operation(node->a, Span());
+ if (const auto* result = replay.template as<T>(); result && result->a.same_as(node->a)) {
+ return OperationDoc(kind, {value});
+ }
+ } catch (const ffi::Error&) {
+ // Preserve the explicit node even when the overloaded spelling rejects it.
+ }
+ }
+ return NamespaceDoc("tirx")->Attr(std::strrchr(T::_type_key, '.') + 1)->Call({value});
+}
+
+template <typename T, PrimExpr (*operation)(PrimExpr, PrimExpr, Span)>
+ExprDoc BinaryHelperDocTranslate(DocTranslatorObj* d, const T* node, const char* helper) {
+ ExprDoc a = d->Translate(node->a).value();
+ ExprDoc b = d->Translate(node->b).value();
+ try {
+ PrimExpr replay = operation(node->a, node->b, Span());
+ if (const auto* result = replay.template as<T>();
+ result && result->a.same_as(node->a) && result->b.same_as(node->b)) {
+ return NamespaceDoc("tirx")->Attr(helper)->Call({a, b});
+ }
+ } catch (const ffi::Error&) {
+ // Keep the explicit constructor when the helper cannot preserve its inputs.
+ }
+ return NamespaceDoc("tirx")->Attr(std::strrchr(T::_type_key, '.') + 1)->Call({a, b});
+}
+
+ffi::Optional<ExprDoc> BitwiseNotDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::BitwiseNotNode>(input);
+ return UnaryOpDocTranslate<prim::BitwiseNotNode, OperationDocNode::Kind::kInvert,
+ tvm::bitwise_neg>(d, node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<prim::BitwiseNotNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BitwiseNotDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> NotDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::NotNode>(input);
+ return UnaryOpDocTranslate<prim::NotNode, OperationDocNode::Kind::kNot, tvm::logical_not>(d,
+ node);
+}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::BitwiseNot>(
- "", [](prim::BitwiseNot node, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc a = d->AsDoc<ExprDoc>(node->a, p->Attr("a"));
- if (a->IsInstance<LiteralDocNode>()) {
- return TIR(d, "BitwiseNot")->Call({a});
- }
- return OperationDoc(OperationDocNode::Kind::kInvert, {a});
- });
+ ffi::reflection::TypeAttrDef<prim::NotNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&NotDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> StringImmDocTranslate(DocTranslatorObj*, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* imm =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const StringImmNode>(input);
+ return LiteralDoc::Str(imm->value, std::nullopt);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<StringImmNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&StringImmDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> CastDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* cast =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::CastNode>(input);
+ return NamespaceDoc("tirx")->Attr("Cast")->Call(
+ {LiteralDoc::DataType(cast->ty.as_or_throw<PrimType>()->dtype, std::nullopt),
+ d->Translate(cast->value).value()});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Not>(
- "", [](prim::Not node, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc a = d->AsDoc<ExprDoc>(node->a, p->Attr("a"));
- if (a->IsInstance<LiteralDocNode>()) {
- return TIR(d, "Not")->Call({a});
- }
- return OperationDoc(OperationDocNode::Kind::kNot, {a});
- });
+ ffi::reflection::TypeAttrDef<prim::CastNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&CastDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> SelectDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::SelectNode>(input);
+ return NamespaceDoc("tirx")->Attr("Select")->Call({d->Translate(node->condition).value(),
+ d->Translate(node->true_value).value(),
+ d->Translate(node->false_value).value()});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<StringImm>(
- "", [](StringImm s, AccessPath p, IRDocsifier d) -> Doc {
- if (HasMultipleLines(s->value)) {
- return d->AddMetadata(s);
- } else {
- return d->AsDoc<ExprDoc>(s->value, p->Attr("value"));
- }
- });
+ ffi::reflection::TypeAttrDef<prim::SelectNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&SelectDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> RampDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::RampNode>(input);
+ return NamespaceDoc("tirx")->Attr("Ramp")->Call({d->Translate(node->base).value(),
+ d->Translate(node->stride).value(),
+ d->Translate(node->lanes).value()});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Cast>(
- "", [](prim::Cast cast, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc dtype = LiteralDoc::DataType(cast.ty()->dtype, p->Attr("dtype"));
- ExprDoc value = d->AsDoc<ExprDoc>(cast->value, p->Attr("value"));
- return TIR(d, "Cast")->Call({dtype, value});
- });
+ ffi::reflection::TypeAttrDef<prim::RampNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&RampDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> BroadcastDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::BroadcastNode>(input);
+ return NamespaceDoc("tirx")
+ ->Attr("Broadcast")
+ ->Call({d->Translate(node->value).value(), d->Translate(node->lanes).value()});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Select>(
- "", [](prim::Select select, AccessPath p, IRDocsifier d) -> Doc {
- return TIR(d, "Select")
- ->Call({
- d->AsDoc<ExprDoc>(select->condition, p->Attr("condition")),
- d->AsDoc<ExprDoc>(select->true_value, p->Attr("true_value")),
- d->AsDoc<ExprDoc>(select->false_value, p->Attr("false_value")),
- });
- });
+ ffi::reflection::TypeAttrDef<prim::BroadcastNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BroadcastDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> ShuffleDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::ShuffleNode>(input);
+ ExprDoc vectors = AnyValue(d, node->vectors);
+ ExprDoc indices = AnyValue(d, node->indices);
+ d->RecordOrigin(vectors, node->vectors);
+ d->RecordOrigin(indices, node->indices);
+ return NamespaceDoc("tirx")->Attr("Shuffle")->Call({vectors, indices});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Ramp>(
- "", [](prim::Ramp ramp, AccessPath ramp_p, IRDocsifier d) -> Doc {
- return TIR(d, "Ramp")->Call({
- d->AsDoc<ExprDoc>(ramp->base, ramp_p->Attr("base")),
- d->AsDoc<ExprDoc>(ramp->stride, ramp_p->Attr("stride")),
- d->AsDoc<ExprDoc>(ramp->lanes, ramp_p->Attr("lanes")),
- });
- });
+ ffi::reflection::TypeAttrDef<prim::ShuffleNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ShuffleDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> LetDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::LetNode>(input);
+ ExprDoc value = d->Translate(node->value).value();
+ IdDoc var = VarDoc(d, node->var);
+ if (auto type = node->var->ty.as<PrimType>()) {
+ ExprDoc declaration = NamespaceDoc("ir")->Attr("dynamic")->Call(
+ {LiteralDoc::Str(node->var->name, std::nullopt)}, {"dtype"},
+ {LiteralDoc::DataType(type.value()->dtype, std::nullopt)});
+ d->Emit(AssignDoc(var, declaration, std::nullopt), node->var);
+ }
+ ExprDoc body = d->Translate(node->body).value();
+ return NamespaceDoc("tirx")->Attr("Let")->Call({body}, {"where"}, {DictDoc({var}, {value})});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Broadcast>(
- "", [](prim::Broadcast bc, AccessPath bc_p, IRDocsifier d) -> Doc {
- return TIR(d, "Broadcast")
- ->Call({
- d->AsDoc<ExprDoc>(bc->value, bc_p->Attr("value")),
- d->AsDoc<ExprDoc>(bc->lanes, bc_p->Attr("lanes")),
- });
- });
+ ffi::reflection::TypeAttrDef<prim::LetNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&LetDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> DivDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::DivNode>(input);
+ PrimType a_type = node->a.ty();
+ PrimType b_type = node->b.ty();
+ if (a_type.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
+ b_type.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
+ return NamespaceDoc("tirx")->Attr("Div")->Call(
+ {d->Translate(node->a).value(), d->Translate(node->b).value()});
+ }
+ return BinaryOpDocTranslate<prim::DivNode, OperationDocNode::Kind::kDiv, tvm::div>(d, input,
+ destination);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Shuffle>( //
- "", [](prim::Shuffle shuffle, AccessPath p, IRDocsifier d) -> Doc {
- return TIR(d, "Shuffle")
- ->Call({
- d->AsDoc<ExprDoc>(shuffle->vectors, p->Attr("vectors")),
- d->AsDoc<ExprDoc>(shuffle->indices, p->Attr("indices")),
- });
- });
+ ffi::reflection::TypeAttrDef<prim::DivNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&DivDocTranslate>());
+ ffi::reflection::TypeAttrDef<prim::AddNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::AddNode, OperationDocNode::Kind::kAdd, tvm::add>>());
+ ffi::reflection::TypeAttrDef<prim::SubNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::SubNode, OperationDocNode::Kind::kSub, tvm::sub>>());
+ ffi::reflection::TypeAttrDef<prim::MulNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::MulNode, OperationDocNode::Kind::kMult, tvm::mul>>());
+ ffi::reflection::TypeAttrDef<prim::FloorDivNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::FloorDivNode, OperationDocNode::Kind::kFloorDiv, tvm::floordiv>>());
+ ffi::reflection::TypeAttrDef<prim::FloorModNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::FloorModNode, OperationDocNode::Kind::kMod, tvm::floormod>>());
+ ffi::reflection::TypeAttrDef<prim::LShiftNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::LShiftNode, OperationDocNode::Kind::kLShift, tvm::left_shift>>());
+ ffi::reflection::TypeAttrDef<prim::RShiftNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::RShiftNode, OperationDocNode::Kind::kRShift, tvm::right_shift>>());
+ ffi::reflection::TypeAttrDef<prim::BitwiseAndNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::BitwiseAndNode, OperationDocNode::Kind::kBitAnd, tvm::bitwise_and>>());
+ ffi::reflection::TypeAttrDef<prim::BitwiseOrNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::BitwiseOrNode, OperationDocNode::Kind::kBitOr, tvm::bitwise_or>>());
+ ffi::reflection::TypeAttrDef<prim::BitwiseXorNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<&BinaryOpDocTranslate<
+ prim::BitwiseXorNode, OperationDocNode::Kind::kBitXor, tvm::bitwise_xor>>());
+ ffi::reflection::TypeAttrDef<prim::LTNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::LTNode, OperationDocNode::Kind::kLt, tvm::less>>());
+ ffi::reflection::TypeAttrDef<prim::LENode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::LENode, OperationDocNode::Kind::kLtE, tvm::less_equal>>());
+ ffi::reflection::TypeAttrDef<prim::EQNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::EQNode, OperationDocNode::Kind::kEq, tvm::equal>>());
+ ffi::reflection::TypeAttrDef<prim::NENode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::NENode, OperationDocNode::Kind::kNotEq, tvm::not_equal>>());
+ ffi::reflection::TypeAttrDef<prim::GTNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::GTNode, OperationDocNode::Kind::kGt, tvm::greater>>());
+ ffi::reflection::TypeAttrDef<prim::GENode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::GENode, OperationDocNode::Kind::kGtE, tvm::greater_equal>>());
+ ffi::reflection::TypeAttrDef<prim::AndNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::AndNode, OperationDocNode::Kind::kAnd, tvm::logical_and>>());
+ ffi::reflection::TypeAttrDef<prim::OrNode>().attr(
+ kDocTranslate,
+ FDocTranslate::FromNative<
+ &BinaryOpDocTranslate<prim::OrNode, OperationDocNode::Kind::kOr, tvm::logical_or>>());
+}
+
+ffi::Optional<ExprDoc> ModDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::ModNode>(input);
+ return BinaryHelperDocTranslate<prim::ModNode, tvm::truncmod>(d, node, "truncmod");
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Let>(
- "", [](prim::Let let, AccessPath p, IRDocsifier d) -> Doc {
- DictDoc where({d->AsDoc<ExprDoc>(let->var, p->Attr("var"))},
- {d->AsDoc<ExprDoc>(let->value, p->Attr("value"))});
- return TIR(d, "Let")->Call({d->AsDoc<ExprDoc>(let->body, p->Attr("body"))}, //
- {"where"}, {where});
- });
-}
-
-#define TVM_SCRIPT_PRINTER_DEF_BINARY(NodeType, OpString) \
- IRDocsifier::vtable().set_dispatch<prim::NodeType>( \
- "", [](prim::NodeType node, AccessPath p, IRDocsifier d) -> Doc { \
- ExprDoc a = d->AsDoc<ExprDoc>(node->a, p->Attr("a")); \
- ExprDoc b = d->AsDoc<ExprDoc>(node->b, p->Attr("b")); \
- return TIR(d, OpString)->Call({a, b}); \
- });
+ ffi::reflection::TypeAttrDef<prim::ModNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&ModDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> MinDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::MinNode>(input);
+ return BinaryHelperDocTranslate<prim::MinNode, tvm::min>(d, node, "min");
+}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<prim::Div>(
- "", [](prim::Div node, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc a = d->AsDoc<ExprDoc>(node->a, p->Attr("a"));
- ExprDoc b = d->AsDoc<ExprDoc>(node->b, p->Attr("b"));
- PrimExpr ret = tvm::div(node->a, node->b);
- if (!ret->IsInstance<prim::DivNode>()) {
- return TIR(d, "Div")->Call({a, b});
- }
- PrimType a_ty = node->a.ty();
- PrimType b_ty = node->b.ty();
- if (a_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
- b_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
- return TIR(d, "Div")->Call({a, b});
- }
- return OperationDoc(OperationDocNode::Kind::kDiv, {a, b});
- });
-}
-
-#define TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(NodeType, NodeObj, NodeFunc, OpString, OpKind) \
- IRDocsifier::vtable().set_dispatch<prim::NodeType>( \
- "", [](prim::NodeType node, AccessPath p, IRDocsifier d) -> Doc { \
- ExprDoc a = d->AsDoc<ExprDoc>(node->a, p->Attr("a")); \
- ExprDoc b = d->AsDoc<ExprDoc>(node->b, p->Attr("b")); \
- PrimExpr ret = tvm::NodeFunc(node->a, node->b); \
- if (const auto* ret_node = ret.as<tvm::NodeObj>()) { \
- if (ret_node->a.same_as(node->a) && ret_node->b.same_as(node->b)) { \
- return OperationDoc(OperationDocNode::Kind::OpKind, {a, b}); \
- } \
- } \
- return TIR(d, OpString)->Call({a, b}); \
- });
+ ffi::reflection::TypeAttrDef<prim::MinNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&MinDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> MaxDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* node =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const prim::MaxNode>(input);
+ return BinaryHelperDocTranslate<prim::MaxNode, tvm::max>(d, node, "max");
+}
TVM_FFI_STATIC_INIT_BLOCK() {
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(Add, prim::AddNode, add, "Add", kAdd);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(Sub, prim::SubNode, sub, "Sub", kSub);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(Mul, prim::MulNode, mul, "Mul", kMult);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(FloorDiv, prim::FloorDivNode, floordiv, "FloorDiv",
- kFloorDiv);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(FloorMod, prim::FloorModNode, floormod, "FloorMod",
- kMod);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(LShift, prim::LShiftNode, left_shift, "LShift", kLShift);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(RShift, prim::RShiftNode, right_shift, "RShift",
- kRShift);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(BitwiseAnd, prim::BitwiseAndNode, bitwise_and,
- "BitwiseAnd", kBitAnd);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(BitwiseOr, prim::BitwiseOrNode, bitwise_or, "BitwiseOr",
- kBitOr);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(BitwiseXor, prim::BitwiseXorNode, bitwise_xor,
- "BitwiseXor", kBitXor);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(LT, prim::LTNode, less, "LT", kLt);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(LE, prim::LENode, less_equal, "LE", kLtE);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(EQ, prim::EQNode, equal, "EQ", kEq);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(NE, prim::NENode, not_equal, "NE", kNotEq);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(GT, prim::GTNode, greater, "GT", kGt);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(GE, prim::GENode, greater_equal, "GE", kGtE);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(And, prim::AndNode, logical_and, "And", kAnd);
- TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR(Or, prim::OrNode, logical_or, "Or", kOr);
-
- TVM_SCRIPT_PRINTER_DEF_BINARY(Mod, "truncmod");
- TVM_SCRIPT_PRINTER_DEF_BINARY(Min, "min");
- TVM_SCRIPT_PRINTER_DEF_BINARY(Max, "max");
-}
-
-#undef TVM_SCRIPT_PRINTER_DEF_BINARY_WITH_SUGAR
-#undef TVM_SCRIPT_PRINTER_DEF_BINARY
+ ffi::reflection::TypeAttrDef<prim::MaxNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&MaxDocTranslate>());
+}
+
+} // namespace
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/ir/prim_type.cc b/src/script/printer/ir/prim_type.cc
index 3816143808..22debd5500 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -16,64 +16,125 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ir/type.h>
-#include "./utils.h"
+#include <tvm/tirx/type.h>
+
+#include <algorithm>
+#include <functional>
+#include <optional>
+
+#include "utils.h"
namespace tvm {
namespace script {
namespace printer {
+namespace details {
+
+namespace {
+
+ffi::Optional<ExprDoc> IntImmDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* imm =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const IntImmNode>(input);
+ DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
+ ExprDoc value = LiteralDoc::Int(ffi::GetRef<IntImm>(imm), std::nullopt);
+ DLDataType default_dtype =
+ ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.int_dtype", "int32"));
+ if (dtype == default_dtype) return value;
+ return NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(dtype))->Call({value});
+}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<IntImm>(
- "", [](IntImm imm, AccessPath imm_p, IRDocsifier d) -> Doc {
- DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
- if (dtype == d->cfg->int_dtype) {
- return LiteralDoc::Int(imm, imm_p->Attr("value"));
- } else if (dtype == DLDataType{kDLBool, 8, 1}) {
- return TIR(d, DType2Str(dtype))
- ->Call({LiteralDoc::Boolean(static_cast<bool>(imm->value), imm_p->Attr("value"))});
- } else {
- return TIR(d, DType2Str(dtype))->Call({LiteralDoc::Int(imm, imm_p->Attr("value"))});
- }
- });
+ ffi::reflection::TypeAttrDef<IntImmNode>().attr(kDocTranslate,
+ FDocTranslate::FromNative<&IntImmDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> FloatImmDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* imm =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const FloatImmNode>(input);
+ DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
+ DLDataType default_dtype =
+ ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("script.float_dtype", "void"));
+ ExprDoc value = LiteralDoc::Float(imm->value, std::nullopt);
+ if (dtype == default_dtype) return value;
+ return NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(dtype))->Call({value});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<FloatImm>(
- "", [](FloatImm imm, AccessPath imm_p, IRDocsifier d) -> Doc {
- DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
- if (dtype == d->cfg->float_dtype) {
- return LiteralDoc::Float(imm->value, imm_p->Attr("value"));
- } else {
- return TIR(d, DType2Str(dtype))
- ->Call({LiteralDoc::Float(imm->value, imm_p->Attr("value"))});
- }
- });
+ ffi::reflection::TypeAttrDef<FloatImmNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&FloatImmDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> PrimTypeDocTranslate(DocTranslatorObj*, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const PrimTypeNode>(input);
+ if (ty->dtype.lanes != 1) {
+ // TIRx exports vectors with 2/4/8/16/32/64 lanes for 8/16/32/64-bit integers,
+ // 16/32/64-bit floats, and the float8, float6, and float4 formats below.
+ // Other fixed or scalable vectors use the generic type constructor.
+ if (ty->dtype.lanes == 2 || ty->dtype.lanes == 4 || ty->dtype.lanes == 8 ||
+ ty->dtype.lanes == 16 || ty->dtype.lanes == 32 || ty->dtype.lanes == 64) {
+ PrimType type = ffi::GetRef<PrimType>(ty);
+ bool supported_element =
+ (type.MatchesCode(kDLInt, kDLUInt) &&
+ (type.bits() == 8 || type.bits() == 16 || type.bits() == 32 || type.bits() == 64)) ||
+ (type.MatchesCode(kDLFloat) &&
+ (type.bits() == 16 || type.bits() == 32 || type.bits() == 64)) ||
+ type.MatchesCode(kDLFloat8_e3m4, kDLFloat8_e4m3, kDLFloat8_e4m3b11fnuz, kDLFloat8_e4m3fn,
+ kDLFloat8_e4m3fnuz, kDLFloat8_e5m2, kDLFloat8_e5m2fnuz,
+ kDLFloat8_e8m0fnu, kDLFloat6_e2m3fn, kDLFloat6_e3m2fn, kDLFloat4_e2m1fn);
+ if (supported_element) return NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(ty->dtype));
+ }
+ return NamespaceDoc("ir")
+ ->Attr("PrimType")
+ ->Call({LiteralDoc::DataType(ty->dtype, std::nullopt)});
+ }
+ return NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(ty->dtype));
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<PrimType>(
- "",
- [](PrimType ty, AccessPath p, IRDocsifier d) -> Doc { return TIR(d, DType2Str(ty->dtype)); });
+ ffi::reflection::TypeAttrDef<PrimTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&PrimTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> StringTypeDocTranslate(DocTranslatorObj*, ffi::AnyView, const ffi::Object*) {
+ return NamespaceDoc("ir")->Attr("StringType")->Call({});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<StringType>(
- "", [](StringType ty, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "StringType")->Call({});
- });
+ ffi::reflection::TypeAttrDef<StringTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&StringTypeDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> TupleTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const TupleTypeNode>(input);
+ if (ty->fields.empty()) return NamespaceDoc("relax")->Attr("Tuple");
+ ffi::Array<ExprDoc> fields;
+ for (const Type& field : ty->fields) fields.push_back(d->Translate(field).value());
+ std::function<bool(const Type&)> is_primitive = [&](const Type& field) {
+ if (field.as<PrimType>() || field.as<PointerType>()) return true;
+ if (const auto* tuple = field.as<TupleTypeNode>()) {
+ return std::all_of(tuple->fields.begin(), tuple->fields.end(), is_primitive);
+ }
+ return false;
+ };
+ bool primitive = std::all_of(ty->fields.begin(), ty->fields.end(), is_primitive);
+ return (primitive ? NamespaceDoc("tirx")->Attr("Tuple") : NamespaceDoc("relax")->Attr("Tuple"))
+ ->Call(fields);
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<TupleType>(
- "", [](TupleType ty, AccessPath p, IRDocsifier d) -> Doc {
- if (ty->fields.empty()) {
- return LiteralDoc::None(p);
- }
- return TIR(d, "Tuple")->Call(d->AsDoc<ListDoc>(ty->fields, p->Attr("fields"))->elements);
- });
+ ffi::reflection::TypeAttrDef<TupleTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TupleTypeDocTranslate>());
}
+
+} // namespace
+
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
new file mode 100644
index 0000000000..b3f35c3d53
--- /dev/null
+++ b/src/script/printer/ir/utils.cc
@@ -0,0 +1,266 @@
+/*
+ * 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 "utils.h"
+
+#include <tvm/ffi/reflection/accessor.h>
+#include <tvm/ir/global_info.h>
+#include <tvm/ir/op.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/type.h>
+
+#include <algorithm>
+#include <optional>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace printer {
+namespace details {
+
+ExprDoc AddMetadata(DocTranslatorObj* d, ffi::Any value) {
+ TVM_FFI_CHECK(value != nullptr, TypeError) << "Metadata cannot contain None";
+ using MetadataMap = ffi::Dict<ffi::String, ffi::List<ffi::Any>>;
+ auto metadata = d->GetOrCreateExtraState<MetadataMap>("ir.metadata_map");
+ ffi::String key = value.GetTypeKey();
+ auto entries = metadata.Get(key).value_or(ffi::List<ffi::Any>{});
+ size_t index = 0;
+ for (; index < entries.size(); ++index) {
+ if (ffi::AnyEqual()(entries[index], value)) break;
+ }
+ if (index == entries.size()) {
+ entries.push_back(value);
+ metadata.Set(key, entries);
+ }
+ auto metadata_id = d->GetOrCreateExtraState<ffi::Optional<IdDoc>>("ir.metadata_id");
+ TVM_FFI_ICHECK(metadata_id.has_value());
+ ExprDoc entries_doc = metadata_id.value()[{LiteralDoc::Str(key, std::nullopt)}];
+ ExprDoc doc = entries_doc[{LiteralDoc::Int(index, std::nullopt)}];
+ if (auto origin = value.as<ffi::ObjectRef>()) d->RecordOrigin(doc, origin.value());
+ return doc;
+}
+
+IdDoc VarDoc(DocTranslatorObj* d, const Var& var, bool explicit_def) {
+ IdDoc doc = IdDoc(d->VarGetOrAllocId(var, explicit_def)->name);
+ d->RecordOrigin(doc, var);
+ return doc;
+}
+
+bool SyntaxSugar(DocTranslatorObj* d) {
+ return d->GetExtraConfig<bool>("script.syntax_sugar", true);
+}
+
+bool IsAnnotated(DocTranslatorObj* d, const ffi::ObjectRef& object) {
+ for (const auto& item :
+ d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.annotated_objects", {})) {
+ if (item.same_as(object)) return true;
+ }
+ return false;
+}
+
+bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef& object) {
+ for (const auto& item :
+ d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.underlined_objects", {})) {
+ if (item.same_as(object)) return true;
+ }
+ return false;
+}
+
+bool HasAnnotatedDescendant(DocTranslatorObj* d, const ffi::ObjectRef& object) {
+ for (const auto& item :
+ d->GetExtraConfig<ffi::Array<ffi::ObjectRef>>("script.annotation_ancestors", {})) {
+ if (item.same_as(object)) return true;
+ }
+ return false;
+}
+
+ExprDoc NamedCallCallee(const ffi::String& canonical_name) {
+ std::string name = canonical_name;
+ size_t end = name.find('.');
+ ExprDoc callee = NamespaceDoc(name.substr(0, end));
+ while (end != std::string::npos) {
+ size_t start = end + 1;
+ end = name.find('.', start);
+ callee = callee->Attr(name.substr(start, end - start));
+ }
+ return callee;
+}
+
+ExprDoc GlobalReference(DocTranslatorObj* d, const ffi::String& name) {
+ auto names = d->GetExtraConfig<ffi::Array<ffi::String>>("script.binding_names", {});
+ return IdDoc(names.empty() ? d->GetExtraConfig<ffi::String>("script.module_name", "Module")
+ : names.back())
+ ->Attr(name);
+}
+
+namespace {
+
+ExprDoc TypeValueImpl(DocTranslatorObj* d, const Type& type, bool dtype_literal) {
+ if (type.as<relax::PackedFuncTypeNode>())
+ return NamespaceDoc("relax")->Attr("PackedFunc")->Call({});
+ if (type.as<AnyTypeNode>()) return NamespaceDoc("relax")->Attr("Any")->Call({});
+ if (auto tuple = type.as<TupleTypeNode>(); tuple && tuple->fields.empty()) {
+ return NamespaceDoc("relax")->Attr("Tuple")->Call({});
+ }
+ if (auto function = type.as<relax::FuncTypeNode>(); function && !function->params.has_value()) {
+ ExprDoc doc = d->Translate(type).value();
+ return doc.as<AttrAccessDocNode>() ? doc->Call({}) : doc;
+ }
+ if (auto primitive = type.as<PrimType>()) {
+ ExprDoc dtype = LiteralDoc::DataType(primitive.value()->dtype, std::nullopt);
+ return dtype_literal ? dtype : NamespaceDoc("ir")->Attr("PrimType")->Call({dtype});
+ }
+ if (auto function = type.as<FuncTypeNode>()) {
+ ffi::Array<ExprDoc> args;
+ for (const Type& arg : function->arg_types) args.push_back(TypeValue(d, arg, false));
+ return NamespaceDoc("ir")
+ ->Attr("FuncType")
+ ->Call({ListDoc(args), TypeValue(d, function->ret_type, false)});
+ }
+ if (type.IsMissing()) {
+ return NamespaceDoc("ir")->Attr("Type")->Attr("missing")->Call({});
+ }
+ if (auto pointer = type.as<PointerTypeNode>()) {
+ if (auto primitive = pointer->element_type.as<PrimType>();
+ primitive.has_value() && primitive.value().IsVoid() && pointer->storage_scope == "global") {
+ return NamespaceDoc("tirx")->Attr("handle")->Call({})->Attr("ty");
+ }
+ return d->Translate(type).value()->Attr("ty");
+ }
+ ExprDoc doc = d->Translate(type).value();
+ if (type.as<relax::TensorTypeNode>() && doc.as<AttrAccessDocNode>()) return doc->Call({});
+ return doc;
+}
+
+} // namespace
+
+ExprDoc TypeValue(DocTranslatorObj* d, const Type& type, bool dtype_literal) {
+ ExprDoc doc = TypeValueImpl(d, type, dtype_literal);
+ d->RecordOrigin(doc, type);
+ return doc;
+}
+
+namespace {
+
+ExprDoc CallAttrsValue(DocTranslatorObj* d, const Attrs& attrs) {
+ if (attrs.as<DictAttrsNode>()) return d->Translate(attrs).value();
+ std::vector<std::pair<ffi::String, ffi::Any>> fields;
+ ffi::reflection::ForEachFieldInfo(
+ TVMFFIGetTypeInfo(attrs->type_index()), [&](const TVMFFIFieldInfo* field) {
+ fields.emplace_back(ffi::String(field->name), ffi::reflection::FieldGetter(field)(attrs));
+ });
+ std::sort(fields.begin(), fields.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ for (const auto& [key, value] : fields) {
+ if (key == "type_key") return AddMetadata(d, attrs);
+ keys.push_back(key);
+ values.push_back(AnyValue(d, value));
+ }
+ ExprDoc inline_doc =
+ NamespaceDoc("ir")
+ ->Attr("make_node")
+ ->Call({LiteralDoc::Str(attrs->GetTypeKey(), std::nullopt)}, keys, values);
+ d->RecordOrigin(inline_doc, attrs);
+ // These reflected fields have a complete inline representation. Keep them
+ // visible in diagnostics and usable when no metadata table is requested.
+ return inline_doc;
+}
+
+} // namespace
+
+// The explicit fallback retains every field, including typed attribute objects.
+ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call, bool infer_result,
+ ffi::Optional<ffi::Array<ExprDoc>> translated_args) {
+ ffi::Optional<Op> op = call->op.as<Op>();
+ ffi::Array<ExprDoc> args;
+ if (translated_args.has_value()) {
+ args = translated_args.value();
+ } else {
+ for (const Expr& arg : call->args) args.push_back(d->Translate(arg).value());
+ }
+ for (size_t i = 0; i < args.size(); ++i) {
+ args.Set(i, MaterializeCallArgument(d, call->args[i], args[i]));
+ }
+ ExprDoc callee = op.has_value() ? LiteralDoc::Str(op.value()->name, std::nullopt)
+ : call->op.as<StringImmNode>() ? AnyValue(d, call->op)
+ : d->Translate(call->op).value();
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (call->attrs.defined()) {
+ keys.push_back("attrs");
+ values.push_back(CallAttrsValue(d, call->attrs));
+ }
+ if (!call->ty_args.empty()) {
+ ffi::Array<ExprDoc> types;
+ for (const Type& type : call->ty_args) types.push_back(TypeValue(d, type, false));
+ keys.push_back("ty_args");
+ values.push_back(ListDoc(types));
+ }
+ if (!infer_result) {
+ keys.push_back("ty");
+ auto primitive = call->ty.as<PrimType>();
+ // Call interprets the string "handle" as a void pointer, not PrimType.
+ bool dtype_literal = !primitive ||
+ !primitive.value().MatchesCode(DLDataTypeCode::kDLOpaqueHandle) ||
+ primitive.value().IsVoid();
+ values.push_back(TypeValue(d, call->ty, dtype_literal));
+ }
+ return NamespaceDoc("ir")->Attr("Call")->Call({callee, ListDoc(args)}, keys, values);
+}
+
+// The query aliases the active frame; candidate classification must own a copy.
+ffi::Dict<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d) {
+ ffi::Dict<Var, IdDoc> candidates;
+ for (const auto& [var, id] : d->GetImplicitDefs()) candidates.Set(var, id);
+ return candidates;
+}
+
+void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc>& signature,
+ const FunctionDoc& function) {
+ // A module's existing scoped registry also identifies shared external captures.
+ auto module =
+ d->GetOrCreateExtraState<ffi::Optional<ffi::Dict<ffi::String, ffi::List<GlobalInfo>>>>(
+ "ir.global_info_map");
+ bool use_pep695 = !module.has_value() && d->GetExtraConfig<bool>("script.use_pep695", false);
+ auto pending = d->GetImplicitDefs();
+ std::vector<Var> parameters;
+ if (use_pep695) {
+ for (const auto& [var, id] : signature) {
+ if (pending.count(var) && var->ty.as<PrimTypeNode>()) parameters.push_back(var);
+ }
+ std::sort(parameters.begin(), parameters.end(),
+ [&](const Var& a, const Var& b) { return signature[a]->name < signature[b]->name; });
+ if (!parameters.empty()) d->ExchangeExtraState("script.future_annotations", true);
+ for (const Var& var : parameters) {
+ IdDoc id = VarDoc(d, var);
+ if (var->ty.as_or_throw<PrimType>()->dtype == (DLDataType{kDLInt, 64, 1})) {
+ function->type_params.push_back(id);
+ } else {
+ function->type_params.push_back(AssignDoc(id, std::nullopt, d->Translate(var->ty).value()));
+ }
+ }
+ }
+}
+
+} // namespace details
+} // namespace printer
+} // namespace script
+} // namespace tvm
diff --git a/src/script/printer/ir/utils.h b/src/script/printer/ir/utils.h
index d04517b79d..272aa1ad44 100644
--- a/src/script/printer/ir/utils.h
+++ b/src/script/printer/ir/utils.h
@@ -16,59 +16,64 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_PRINTER_IR_UTILS_H_
-#define TVM_SCRIPT_PRINTER_IR_UTILS_H_
+#ifndef SRC_SCRIPT_PRINTER_IR_UTILS_H_
+#define SRC_SCRIPT_PRINTER_IR_UTILS_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
-#include <tvm/ir/function.h>
-#include <tvm/ir/op.h>
-#include <tvm/ir/with_context.h>
-#include <tvm/script/printer/ir_docsifier.h>
+#include <tvm/ir/type.h>
+#include <tvm/script/printer/doc_translator.h>
-#include <string>
+#include <optional>
#include <utility>
-#include "../utils.h"
-
namespace tvm {
namespace script {
namespace printer {
-class IRFrameNode : public FrameNode {
- public:
- ffi::Map<ffi::String, ffi::Array<GlobalInfo>>* global_infos = nullptr;
+namespace details {
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- // global infos is not exposed
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IRFrame", IRFrameNode, FrameNode);
-};
+ExprDoc AddMetadata(DocTranslatorObj* d, ffi::Any value);
+IdDoc VarDoc(DocTranslatorObj* d, const Var& var, bool explicit_def = true);
+bool SyntaxSugar(DocTranslatorObj* d);
+bool IsAnnotated(DocTranslatorObj* d, const ffi::ObjectRef& object);
+bool IsUnderlined(DocTranslatorObj* d, const ffi::ObjectRef& object);
+bool HasAnnotatedDescendant(DocTranslatorObj* d, const ffi::ObjectRef& object);
+ExprDoc GlobalReference(DocTranslatorObj* d, const ffi::String& name);
+ExprDoc NamedCallCallee(const ffi::String& canonical_name);
-class IRFrame : public Frame {
+ExprDoc TypeValue(DocTranslatorObj* d, const Type& type, bool dtype_literal = true);
+ExprDoc MaterializeCallArgument(DocTranslatorObj* d, const Expr& arg, ExprDoc doc);
+ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call, bool infer_result,
+ ffi::Optional<ffi::Array<ExprDoc>> translated_args = std::nullopt);
+ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value);
+ffi::Dict<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d);
+void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc>& signature,
+ const FunctionDoc& function);
+
+class VarScope {
public:
- explicit IRFrame(const IRDocsifier& d) {
- ffi::ObjectPtr<IRFrameNode> n = ffi::make_object<IRFrameNode>();
- n->stmts.clear();
- n->d = d.get();
- n->global_infos = nullptr;
- data_ = std::move(n);
+ explicit VarScope(DocTranslatorObj* d) : d_(d) { d_->BeginVarScope(); }
+ ~VarScope() noexcept {
+ if (d_) {
+ try {
+ d_->EndVarScope();
+ } catch (...) { /* Preserve translation failure. */
+ }
+ }
+ }
+ void Close() {
+ DocTranslatorObj* d = std::exchange(d_, nullptr);
+ d->EndVarScope();
}
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRFrame, Frame, IRFrameNode);
+ private:
+ DocTranslatorObj* d_;
};
-/*! \brief Redirected method for the ffi repr hook */
-inline std::string ReprPrintIR(const ffi::ObjectRef& obj, const PrinterConfig& cfg) {
- IRDocsifier d(cfg);
- With<IRFrame> f(d);
- (*f)->AddDispatchToken(d, "ir");
- return Docsify(obj, d, *f, cfg);
-}
-
+} // namespace details
} // namespace printer
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_PRINTER_IR_UTILS_H_
+#endif // SRC_SCRIPT_PRINTER_IR_UTILS_H_
diff --git a/src/script/printer/ir_docsifier.cc b/src/script/printer/ir_docsifier.cc
... 8288 lines suppressed ...