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/02 11:16:30 UTC
(tvm) branch main updated: [REFACTOR][SCRIPT] Introduce NativeFunction backed Printer Mechanism (#20516)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b38e214658 [REFACTOR][SCRIPT] Introduce NativeFunction backed Printer Mechanism (#20516)
b38e214658 is described below
commit b38e214658c11cc6d188c425f20f4f1da51568ea
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Fri Oct 2 07:16:23 2026 -0400
[REFACTOR][SCRIPT] Introduce NativeFunction backed Printer Mechanism (#20516)
---
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 | 29 +-
include/tvm/script/printer/doc.h | 158 ++-
include/tvm/script/printer/doc_translator.h | 418 +++++++
include/tvm/script/printer/ir_docsifier.h | 365 ------
include/tvm/script/printer/ir_docsifier_functor.h | 199 ---
include/tvm/script/printer/printer.h | 71 +-
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/__init__.py | 3 +
python/tvm/relax/base_py_module.py | 4 -
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 | 61 +-
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 | 29 +-
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/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 | 200 ++-
src/relax/script/printer/expr.cc | 179 +--
src/relax/script/printer/function.cc | 235 ++--
src/relax/script/printer/global_info.cc | 66 +-
src/relax/script/printer/region.cc | 161 ++-
src/relax/script/printer/tir.cc | 139 +--
src/relax/script/printer/type.cc | 51 +-
src/relax/script/printer/utils.h | 159 +--
src/s_tir/schedule/error.cc | 1 -
src/s_tir/script/ir_builder/frame.cc | 49 +-
src/s_tir/script/ir_builder/ir.cc | 8 +-
src/s_tir/script/printer/config.cc | 5 +-
src/script/printer/config.cc | 49 +-
src/script/printer/doc.cc | 12 +-
.../python_doc_printer.cc => doc_printer.cc} | 873 ++++++++++++-
.../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 | 223 ++++
src/script/printer/ir/ir.cc | 301 +++--
src/script/printer/ir/misc.cc | 272 +++--
src/script/printer/ir/prim_expr.cc | 432 ++++---
src/script/printer/ir/prim_type.cc | 136 ++-
src/script/printer/ir/utils.cc | 230 ++++
src/script/printer/ir/utils.h | 71 +-
src/script/printer/ir_docsifier.cc | 218 ----
src/script/printer/printer.cc | 433 +++++++
src/script/printer/script_printer.cc | 156 ++-
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 | 1283 ++++++--------------
src/tirx/script/printer/expr.cc | 1022 +++++++++-------
src/tirx/script/printer/for_loop.cc | 256 ++--
src/tirx/script/printer/function.cc | 353 ++----
src/tirx/script/printer/ir.cc | 91 +-
src/tirx/script/printer/stmt.cc | 649 +++++-----
src/tirx/script/printer/utils.h | 332 +----
src/tirx/transform/lower_thread_allreduce.h | 2 +-
.../relax/script/test_relax_script_basic_usage.py | 6 +
.../script/test_relax_script_dynamic_shape.py | 52 +-
.../script/test_relax_script_meta_programming.py | 2 +
.../relax/script/test_relax_script_printer.py | 108 +-
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../relax/test_transform_error_enrichment.py | 1 +
.../script/test_s_tir_script_dynamic_shape.py | 3 +
.../s_tir/script/test_s_tir_script_ir_builder.py | 6 +-
.../s_tir/script/test_s_tir_script_namespace.py | 1 +
.../s_tir/script/test_s_tir_script_printer.py | 262 ++--
...form_merge_dynamic_shared_memory_allocations.py | 20 +-
.../transform/test_s_tir_transform_thread_sync.py | 6 +-
tests/python/tirx/codegen/test_ptx_addr.py | 2 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 14 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 8 +
.../tirx/script/test_tirx_script_basic_usage.py | 10 +-
.../python/tirx/script/test_tirx_script_buffers.py | 58 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 52 +-
.../tirx/script/test_tirx_script_ir_builder.py | 4 +-
.../python/tirx/script/test_tirx_script_printer.py | 173 +--
.../test_tirx_script_printer_structural_equal.py | 4 +-
.../script/test_tirx_script_printer_underlining.py | 8 +-
.../script/test_tirx_script_source_locations.py | 30 +-
tests/python/tirx/test_op_namespace_cleanup.py | 4 +-
114 files changed, 7198 insertions(+), 6998 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..9d0a562864 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*/
@@ -71,10 +62,6 @@ class PrinterConfigNode : public ffi::Object {
bool print_line_numbers = false;
/*! \brief Number of context lines to print around the underlined text */
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 +108,10 @@ 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 +135,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..7288899682
--- /dev/null
+++ b/include/tvm/script/printer/doc_translator.h
@@ -0,0 +1,418 @@
+/*
+ * 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.
+ */
+ void SetExtraState(const ffi::String& name, std::optional<ffi::Any> value) {
+ if (value.has_value()) {
+ extra_state_.Set(name, std::move(value.value()));
+ } else {
+ extra_state_.erase(name);
+ }
+ }
+
+ /*!
+ * \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..22d4c52e56 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -18,57 +18,52 @@
*/
/*!
* \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/container/map.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 Register a namespace alias during dialect static initialization.
+ * \param key The existing prefix configuration key, such as "tirx.prefix".
+ * \param default_alias The alias reserved before translation assigns variable names.
+ */
+TVM_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias);
+/*! \brief Read the registered namespace aliases. */
+TVM_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases();
+
+/*!
+ * \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..f572b02e37 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={"ir.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, path_to_underline=[lhs_path])
+ rhs_script = Scriptable.script(rhs, 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/__init__.py b/python/tvm/relax/__init__.py
index a5620ce9cc..e19a31acac 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -134,3 +134,6 @@ def _initialize_script_namespace() -> None:
from tvm.script.parser import register_namespace_initializer as _register_namespace_initializer
_register_namespace_initializer(_initialize_script_namespace, aliases=("relax",))
+
+# Register the exported operator printer names for compiler diagnostics as well.
+from .script import ir_builder as _script_ir_builder
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index f1acdeb169..6495a77e19 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -548,8 +548,6 @@ class BasePyModule:
indent_spaces: int = 4,
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:
@@ -570,8 +568,6 @@ class BasePyModule:
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,
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..483e02e3b3 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -17,7 +17,6 @@
"""Configuration of TVMScript printer"""
import os
-import sys
from collections.abc import Sequence
from tvm_ffi import get_global_func, register_object
@@ -30,21 +29,26 @@ 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 ``ir.float_dtype`` entries in ``extra_config``.
+ Explicit entries in ``extra_config`` override those arguments.
+
+ ``tirx.scalar_buffer_as_mutable_var`` in ``extra_config`` defaults to true.
+ It prints eligible local scalar buffers as typed mutable variables. Set it
+ to false to retain explicit buffer allocations and indexed loads/stores.
+ """
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
@@ -65,8 +69,6 @@ class PrinterConfig(Object):
indent_spaces: int = 4,
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 +83,12 @@ 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,
+ "ir.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,
@@ -132,8 +132,6 @@ class Scriptable:
indent_spaces: int = 4,
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 +150,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"
@@ -166,10 +163,6 @@ class Scriptable:
Whether to print line numbers
num_context_lines : int = -1
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,10 +193,9 @@ class Scriptable:
merged_extra: dict = {}
if extra_config is not None:
merged_extra.update(extra_config)
- if "script.use_pep695" not in merged_extra:
- merged_extra["script.use_pep695"] = merged_extra.get(
- "relax.use_pep695", sys.version_info >= (3, 12)
- )
+ # Keep the historical interactive display header. Direct printer calls
+ # retain executable imports for standalone parser round-trips.
+ merged_extra.setdefault("ir.comment_imports", True)
# Only auto-switch if the caller has not already set a tirx.prefix override.
if "tirx.prefix" not in merged_extra:
@@ -244,8 +236,6 @@ class Scriptable:
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,
show_all_ty=show_all_ty,
extra_config=merged_extra if merged_extra else None,
path_to_underline=path_to_underline,
@@ -268,8 +258,6 @@ class Scriptable:
indent_spaces: int = 4,
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,
@@ -277,10 +265,6 @@ class Scriptable:
obj_to_annotate: dict[Object, str] | None = None,
) -> str:
merged_extra = dict(extra_config or {})
- if "script.use_pep695" not in merged_extra:
- merged_extra["script.use_pep695"] = merged_extra.get(
- "relax.use_pep695", sys.version_info >= (3, 12)
- )
return _relax_script(
self,
PrinterConfig(
@@ -294,8 +278,6 @@ class Scriptable:
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,
extra_config=merged_extra,
path_to_underline=path_to_underline,
path_to_annotate=path_to_annotate,
@@ -319,8 +301,6 @@ class Scriptable:
indent_spaces: int = 4,
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 +342,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"
@@ -376,10 +355,6 @@ class Scriptable:
Whether to print line numbers
num_context_lines : int = -1
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
@@ -414,8 +389,6 @@ class Scriptable:
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,
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..4eff21210a 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,
)
@@ -829,7 +833,7 @@ def tvm_thread_invariant(cond):
return call_intrin(_primexpr_ty(cond), "tirx.tvm_thread_invariant", cond)
-def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
+def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1, *, dtype="void"):
"""Perform synchronization in specified scope.
Parameters
@@ -837,18 +841,31 @@ def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
storage_scope : str
The storage scope to perform synchronization.
- is_load : bool
+ is_load : bool or Expr or None
Whether to perform load synchronization. (for global sync only)
+ Set both ``is_load`` and ``num_blocks`` to None to omit these operands.
- num_blocks : int
+ num_blocks : int or Expr or None
The number of blocks to synchronize. (for global sync only)
+ Set to None to omit this operand.
+
+ dtype : str or tvm.ir.Type
+ The stored result type. Defaults to void.
Returns
-------
call : Expr
The call expression.
"""
- return call_intrin("void", "tirx.tvm_storage_sync", storage_scope, is_load, num_blocks)
+ args = [storage_scope]
+ if is_load is None:
+ if num_blocks is not None:
+ raise ValueError("num_blocks must be None when is_load is omitted")
+ else:
+ args.append(is_load)
+ if num_blocks is not None:
+ args.append(num_blocks)
+ return call_intrin(dtype, "tirx.tvm_storage_sync", *args)
def tvm_kernel_replace_point():
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/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..d1fe44f357 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>{
+ {"ir.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..e0386ef8f7 100644
--- a/src/relax/script/printer/distributed.cc
+++ b/src/relax/script/printer/distributed.cc
@@ -16,132 +16,116 @@
* 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);
+ }
+ 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..42cba66e0a 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/script/printer/printer.h>
+#include <tvm/tirx/type.h>
+
#include <algorithm>
+#include <optional>
+#include <utility>
+#include <vector>
-#include "./utils.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);
- });
+ RegisterNamespaceAlias("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..d6ad01bf25 100644
--- a/src/relax/script/printer/global_info.cc
+++ b/src/relax/script/printer/global_info.cc
@@ -16,33 +16,67 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/relax/global_info.h>
+#include <tvm/ir/module.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({});
- });
+// 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 module = d->GetOrCreateExtraState<ffi::Optional<IRModule>>("ir.module");
+ if (!module.has_value()) return std::nullopt;
+ const auto& infos = module.value()->global_infos;
+ 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..9c15d81ee6 100644
--- a/src/relax/script/printer/tir.cc
+++ b/src/relax/script/printer/tir.cc
@@ -16,129 +16,46 @@
* 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 {
+
+namespace {
+
+ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim, const ExprDoc& translated) {
+ auto operation = translated.as<OperationDoc>();
+ if (auto mul = dim.as<prim::MulNode>();
+ mul && operation && operation.value()->kind == OperationDocNode::Kind::kMult) {
+ ExprDoc compact = OperationDoc(OperationDocNode::Kind::kMult,
+ {RelaxShapeDim(d, mul->a, operation.value()->operands[0]),
+ RelaxShapeDim(d, mul->b, operation.value()->operands[1])});
+ 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();
+ return translated;
}
-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);
- });
-}
+} // namespace
-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")),
- });
- });
+ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim) {
+ return RelaxShapeDim(d, dim, d->Translate(dim).value());
}
+} // 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/schedule/error.cc b/src/s_tir/schedule/error.cc
index b6aed81da9..b233e91061 100644
--- a/src/s_tir/schedule/error.cc
+++ b/src/s_tir/schedule/error.cc
@@ -51,7 +51,6 @@ ffi::String ScheduleErrorContextObj::RenderReport(const ffi::String& primitive)
int n_locs = locs.size();
std::string msg = DetailRenderTemplate();
PrinterConfig cfg;
- cfg->syntax_sugar = false;
if (n_locs > 0) {
for (int i = 0; i < n_locs; ++i) {
std::string name = locs[i]->GetTypeKey() + '#' + std::to_string(i);
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/s_tir/script/printer/config.cc b/src/s_tir/script/printer/config.cc
index 461858edad..c0e135d7de 100644
--- a/src/s_tir/script/printer/config.cc
+++ b/src/s_tir/script/printer/config.cc
@@ -17,14 +17,13 @@
* under the License.
*/
#include <tvm/ffi/reflection/registry.h>
-
-#include "../../../script/printer/dialect_prefix.h"
+#include <tvm/script/printer/printer.h>
namespace tvm {
namespace script {
namespace printer {
-TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("s_tir.prefix", "Ts"); }
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("s_tir.prefix", "Ts"); }
} // namespace printer
} // namespace script
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 993e4b3f44..3e384a2628 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -17,21 +17,21 @@
* under the License.
*/
#include <tvm/script/printer/config.h>
+#include <tvm/script/printer/printer.h>
#include <algorithm>
#include <cctype>
-#include <map>
#include <string>
#include <unordered_set>
-
-#include "./dialect_prefix.h"
+#include <utility>
+#include <vector>
namespace tvm {
namespace {
-std::map<ffi::String, ffi::String>& DialectPrefixes() {
- static std::map<ffi::String, ffi::String> prefixes;
- return prefixes;
+ffi::Map<ffi::String, ffi::String>& NamespaceAliases() {
+ static ffi::Map<ffi::String, ffi::String> aliases;
+ return aliases;
}
bool IsIdentifier(const std::string& name) {
@@ -53,11 +53,13 @@ bool IsIdentifier(const std::string& name) {
namespace script {
namespace printer {
-void RegisterDialectPrefix(const ffi::String& key, const ffi::String& default_prefix) {
- TVM_FFI_ICHECK(DialectPrefixes().emplace(key, default_prefix).second)
- << "Duplicate printer dialect prefix: " << key;
+void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias) {
+ TVM_FFI_ICHECK(!NamespaceAliases().count(key)) << "Duplicate printer dialect prefix: " << key;
+ NamespaceAliases().Set(key, default_alias);
}
+const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases() { return NamespaceAliases(); }
+
} // namespace printer
} // namespace script
@@ -84,10 +86,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(
+ "ir.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>();
@@ -122,12 +128,6 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
v.value().as_or_throw<ffi::Optional<ffi::Map<ffi::ObjectRef, ffi::String>>>().value_or(
ffi::Map<ffi::ObjectRef, ffi::String>());
}
- 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 +147,15 @@ PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
}
}
- // Validate all registered prefixes before names can be assigned by a docsifier.
+ // 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{"ir.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);
@@ -158,7 +166,10 @@ ffi::Array<ffi::String> PrinterConfigNode::GetBuiltinKeywords() {
TVM_FFI_ICHECK(module_alias.empty() || IsIdentifier(std::string(module_alias)))
<< "Invalid `module_alias`: " << module_alias;
ffi::Array<ffi::String> result{ir_prefix};
- for (const auto& [key, default_prefix] : DialectPrefixes()) {
+ const auto& aliases = NamespaceAliases();
+ std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end());
+ std::sort(ordered.begin(), ordered.end());
+ for (const auto& [key, default_prefix] : ordered) {
ffi::String prefix = GetExtraConfig<ffi::String>(key, default_prefix);
TVM_FFI_ICHECK(IsIdentifier(std::string(prefix))) << "Invalid `" << key << "`: " << prefix;
result.push_back(prefix);
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..37080b6c4d 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 <tvm/script/printer/printer.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"
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,14 @@ 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";
+ ffi::String fallback = GetNamespaceAliases().Get(key).value_or(name);
+ 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 +1312,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 +1355,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 +1426,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 +1439,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 +1607,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 +1706,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%
rename from src/script/printer/dialect_prefix.h
rename 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..8db9bda4b7
--- /dev/null
+++ b/src/script/printer/doc_translator.cc
@@ -0,0 +1,223 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/base.h>
+#include <tvm/script/printer/doc.h>
+#include <tvm/script/printer/doc_translator.h>
+#include <tvm/script/printer/printer.h>
+
+#include <algorithm>
+#include <cctype>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+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();
+ self->RecordOrigin(id, var);
+ 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 ffi::Optional<ExprDoc> Dispatch(DocTranslatorObj* d, ffi::AnyView value,
+ const ffi::Object* destination) {
+ return d->DefaultTranslate(value, destination);
+ }
+ static const DocTranslatorVTable* VTable() {
+ static const DocTranslatorVTable table{
+ Adapter<ffi::Optional<ExprDoc>(DocTranslatorObj*, ffi::AnyView, const ffi::Object*),
+ &Dispatch>(),
+ 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};
+};
+
+Doc TranslateComplete(ffi::AnyView obj, TranslationEngine* engine,
+ const ffi::Map<ffi::String, ffi::Any>& extra_config) {
+ // Reserve the same registered or explicitly configured aliases used by rendering.
+ // Keep key order so overlapping aliases reserve the same generated names.
+ const auto& aliases = GetNamespaceAliases();
+ std::vector<std::pair<ffi::String, ffi::String>> ordered(aliases.begin(), aliases.end());
+ std::sort(ordered.begin(), ordered.end());
+ for (const auto& [key, fallback] : ordered) {
+ 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>());
+ }
+ // Metadata references share one allocated identifier across the completed Doc.
+ IdDoc metadata_id = engine->AllocId("metadata");
+ engine->SetExtraState("ir.metadata_id", metadata_id);
+ auto docs = engine->WithDocScope([&]() {
+ ffi::Optional<ExprDoc> expression = engine->Translate(obj);
+ if (expression.has_value()) {
+ engine->Emit(ExprStmtDoc(expression.value()), obj.as<ffi::ObjectRef>());
+ }
+ });
+ auto root_declarations = engine->WithDocScope([&]() {
+ auto pending = engine->GetImplicitDefs();
+ std::vector<std::pair<Var, IdDoc>> remaining(pending.begin(), pending.end());
+ for (const auto& [var, _] : remaining) {
+ if (!pending.count(var)) continue;
+ if (auto rhs = engine->Translate(var, var)) {
+ 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;
+ }
+ ffi::Array<StmtDoc> body;
+ for (const auto& stmt : ToStmtDocArray(docs)) body.push_back(stmt);
+ return StmtBlockDoc(body);
+}
+
+} // namespace
+
+Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_origins,
+ ffi::Map<ffi::String, ffi::Any> extra_config) {
+ auto engine = ffi::make_object<TranslationEngine>(extra_config);
+ Doc doc = TranslateComplete(ir, engine.get(), extra_config);
+ if (doc_origins) *doc_origins = engine->GetDocOrigins();
+ return 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/ir/ir.cc b/src/script/printer/ir/ir.cc
index bc60d69694..43ec200062 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -16,196 +16,179 @@
* 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/script/printer/printer.h>
+#include <tvm/tirx/function.h>
-#include "./utils.h"
+#include <algorithm>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#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";
- }
- }
+namespace {
- bool operator<(const SortableFunction& other) const {
- if (this->priority != other.priority) {
- return this->priority < other.priority;
+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";
+ ffi::String module_name = d->GetExtraConfig<ffi::String>("ir.module_name", "Module");
+ IdDoc module_id = d->AllocId(module_name);
+ d->SetExtraState("ir.module_id", module_id);
+ d->SetExtraState("ir.module", ffi::GetRef<IRModule>(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));
}
- return this->gv->name_hint < other.gv->name_hint;
- }
-};
-
-ffi::Optional<ffi::String> GetDynamicDeclarationName(const StmtDoc& stmt) {
- const auto* assign = stmt.as<AssignDocNode>();
- if (assign == nullptr || !assign->rhs.has_value()) {
- return std::nullopt;
- }
- const auto* lhs = assign->lhs.as<IdDocNode>();
- const auto* call = assign->rhs.value().as<CallDocNode>();
- if (lhs == nullptr || call == nullptr) {
- return std::nullopt;
- }
- const auto* callee = call->callee.as<AttrAccessDocNode>();
- if (callee == nullptr || callee->name != "dynamic") {
- return std::nullopt;
- }
- 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);
- }
+ 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)})})));
+ }
+ }
+ });
+ d->SetExtraState("ir.module", std::nullopt);
+ d->SetExtraState("ir.module_id", std::nullopt);
+ d->Emit(ClassDoc(module_id, {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"));
- });
+ RegisterNamespaceAlias("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..e89c9341ea 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -17,102 +17,224 @@
* 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::Any, ffi::Any>> entries(map->begin(), map->end());
+ if (std::all_of(entries.begin(), entries.end(), [](const auto& entry) {
+ return entry.first.template as<ffi::String>().has_value();
+ })) {
+ std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& b) {
+ return a.first.template as_or_throw<ffi::String>() <
+ b.first.template as_or_throw<ffi::String>();
});
+ }
+ ffi::Array<ExprDoc> keys;
+ ffi::Array<ExprDoc> values;
+ for (const auto& [key, item] : entries) {
+ keys.push_back(AnyValue(d, key));
+ 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)});
+ }
+ // Any-valued fields must retain IR objects; Expr-valued fields convert literals.
+ if (const auto* integer = value.as<IntImmNode>()) {
+ ExprDoc doc = NamespaceDoc("tirx")->Attr("IntImm")->Call(
+ {LiteralDoc::DataType(integer->ty.as_or_throw<PrimType>()->dtype, std::nullopt),
+ LiteralDoc::Int(ffi::GetRef<IntImm>(integer), std::nullopt)});
+ d->RecordOrigin(doc, ffi::GetRef<IntImm>(integer));
+ return doc;
+ }
+ if (const auto* floating = value.as<FloatImmNode>()) {
+ ExprDoc literal = LiteralDoc::Float(floating->value, std::nullopt);
+ if (floating->ty.as_or_throw<PrimType>()->dtype == (DLDataType{kDLFloat, 32, 1})) {
+ ExprDoc doc = NamespaceDoc("tirx")->Attr("float32")->Call({literal});
+ d->RecordOrigin(doc, ffi::GetRef<FloatImm>(floating));
+ return doc;
+ }
+ if (!std::isfinite(floating->value)) {
+ // Nonfinite literals render as strings; the scalar helper converts them to a number.
+ literal = NamespaceDoc("tirx")->Attr("float64")->Call({literal})->Attr("value");
+ }
+ ExprDoc doc =
+ NamespaceDoc("tirx")
+ ->Attr("FloatImm")
+ ->Call({LiteralDoc::DataType(floating->ty.as_or_throw<PrimType>()->dtype, std::nullopt),
+ literal});
+ d->RecordOrigin(doc, ffi::GetRef<FloatImm>(floating));
+ return doc;
+ }
+ if (const auto* string = value.as<StringImmNode>()) {
+ 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..bf64ce7965 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -16,64 +16,124 @@
* 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>("ir.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);
+ 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);
+ bool supported_lanes = ty->dtype.lanes == 1 || ty->dtype.lanes == 2 || ty->dtype.lanes == 4 ||
+ ty->dtype.lanes == 8 || ty->dtype.lanes == 16 || ty->dtype.lanes == 32 ||
+ ty->dtype.lanes == 64;
+ bool supported_scalar = type.IsScalar() && (type.MatchesElementType(kDLBool, 8) ||
+ type.MatchesElementType(kDLBfloat, 16));
+ // Other widths and fixed or scalable vectors have no exported TIRx shorthand.
+ if ((supported_element && supported_lanes) || supported_scalar) {
+ return NamespaceDoc("tirx")->Attr(ffi::DLDataTypeToString(ty->dtype));
+ }
+ return NamespaceDoc("ir")
+ ->Attr("PrimType")
+ ->Call({LiteralDoc::DataType(ty->dtype, std::nullopt)});
}
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..50b8ed2fb9
--- /dev/null
+++ b/src/script/printer/ir/utils.cc
@@ -0,0 +1,230 @@
+/*
+ * 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/module.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;
+}
+
+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 module = d->GetOrCreateExtraState<ffi::Optional<IdDoc>>("ir.module_id");
+ if (!module.has_value()) {
+ module = IdDoc(d->GetExtraConfig<ffi::String>("ir.module_name", "Module"));
+ d->SetExtraState("ir.module_id", module);
+ }
+ return module.value()->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");
+ values.push_back(TypeValue(d, call->ty));
+ }
+ 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) {
+ auto module = d->GetOrCreateExtraState<ffi::Optional<IRModule>>("ir.module");
+ auto pending = d->GetImplicitDefs();
+ std::vector<Var> parameters;
+ if (!module.has_value()) {
+ 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; });
+ 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..2dd96e46ec 100644
--- a/src/script/printer/ir/utils.h
+++ b/src/script/printer/ir/utils.h
@@ -16,59 +16,60 @@
* 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);
+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
deleted file mode 100644
index 885908fe71..0000000000
--- a/src/script/printer/ir_docsifier.cc
+++ /dev/null
@@ -1,218 +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.
- */
... 9350 lines suppressed ...