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 de2551db8b [REFACTOR][TVMScript] Align C++ builder and printer
ownership (#20443)
de2551db8b is described below
commit de2551db8b4f111e52e42b68c115737aa6fae936
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Sep 25 07:12:32 2026 -0400
[REFACTOR][TVMScript] Align C++ builder and printer ownership (#20443)
Place TIRx, S-TIR, and Relax C++ builders under their
`script/ir_builder` directories and flatten the shared IR builder
directory, keeping parser completion beside its S-TIR builder consumer.
Move shared primitive-expression and literal/type handlers into the IR
printer layer while preserving dialect repr/frame behavior. Keep
PrinterConfig parsing dialect-neutral by forwarding extension keys and
registering dialect-owned prefixes before name allocation. Consolidate
shared shape printing and remove obsolete printer aliases and forwarding
wrappers.
---
CMakeLists.txt | 3 +-
.../relax/script/{builder => ir_builder}/frame.h | 10 +-
.../tvm/relax/script/{builder => ir_builder}/ir.h | 8 +-
.../s_tir/script/{builder => ir_builder}/frame.h | 8 +-
.../tvm/s_tir/script/{builder => ir_builder}/ir.h | 10 +-
include/tvm/script/ir_builder/{ir => }/frame.h | 6 +-
include/tvm/script/ir_builder/{ir => }/ir.h | 8 +-
include/tvm/script/printer/config.h | 2 +-
.../tirx/script/{builder => ir_builder}/frame.h | 8 +-
.../tvm/tirx/script/{builder => ir_builder}/ir.h | 8 +-
.../script/{builder => ir_builder}/distributed.cc | 2 +-
src/relax/script/{builder => ir_builder}/frame.cc | 4 +-
src/relax/script/{builder => ir_builder}/ir.cc | 4 +-
src/relax/script/{builder => ir_builder}/utils.h | 10 +-
src/relax/script/printer/expr.cc | 3 +
src/s_tir/script/{builder => ir_builder}/frame.cc | 2 +-
src/s_tir/script/{builder => ir_builder}/ir.cc | 2 +-
.../{builder => ir_builder}/script_complete.cc | 2 +-
.../{builder => ir_builder}/script_complete.h | 8 +-
src/s_tir/script/{builder => ir_builder}/utils.h | 10 +-
.../script/printer/config.cc} | 19 +-
src/script/ir_builder/{ir => }/frame.cc | 2 +-
src/script/ir_builder/{ir => }/ir.cc | 2 +-
src/script/ir_builder/{ir => }/utils.h | 8 +-
src/script/printer/config.cc | 165 +++++++++++++-
.../{ir/distributed.cc => dialect_prefix.h} | 24 +-
src/script/printer/ir/ir.cc | 6 +-
src/script/printer/ir/misc.cc | 15 ++
src/script/printer/ir/prim_expr.cc | 201 +++++++++++++++++
.../ir.cc => script/printer/ir/prim_type.cc} | 63 +-----
src/script/printer/script_printer.cc | 146 +-----------
src/script/printer/utils.h | 9 -
src/tirx/script/{builder => ir_builder}/frame.cc | 4 +-
src/tirx/script/{builder => ir_builder}/ir.cc | 2 +-
src/tirx/script/{builder => ir_builder}/utils.h | 10 +-
src/tirx/script/printer/block.cc | 6 +-
src/tirx/script/printer/buffer.cc | 26 +--
src/tirx/script/printer/expr.cc | 244 +++------------------
src/tirx/script/printer/ir.cc | 52 +----
src/tirx/script/printer/stmt.cc | 30 +--
src/tirx/script/printer/utils.h | 4 -
.../python/tvmscript/test_tvmscript_printer_ir.py | 55 ++++-
42 files changed, 594 insertions(+), 617 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 482d356230..719c66da25 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -310,7 +310,8 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
# src/script/{printer,ir_builder}/<dialect>/ for a non-IR dialect is
# rejected at link time rather than silently included.
src/script/ir_builder/base.cc
- src/script/ir_builder/ir/*.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
diff --git a/include/tvm/relax/script/builder/frame.h
b/include/tvm/relax/script/ir_builder/frame.h
similarity index 98%
rename from include/tvm/relax/script/builder/frame.h
rename to include/tvm/relax/script/ir_builder/frame.h
index 2d3cd6cf8f..8ddba3659d 100644
--- a/include/tvm/relax/script/builder/frame.h
+++ b/include/tvm/relax/script/ir_builder/frame.h
@@ -16,15 +16,15 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
-#define TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
+#ifndef TVM_RELAX_SCRIPT_IR_BUILDER_FRAME_H_
+#define TVM_RELAX_SCRIPT_IR_BUILDER_FRAME_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/script/ir_builder/base.h>
-#include <tvm/script/ir_builder/ir/frame.h>
-#include <tvm/script/ir_builder/ir/ir.h>
+#include <tvm/script/ir_builder/frame.h>
+#include <tvm/script/ir_builder/ir.h>
#include <utility>
@@ -344,4 +344,4 @@ class ElseFrame : public SeqExprFrame {
} // namespace script
} // namespace tvm
-#endif // TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
+#endif // TVM_RELAX_SCRIPT_IR_BUILDER_FRAME_H_
diff --git a/include/tvm/relax/script/builder/ir.h
b/include/tvm/relax/script/ir_builder/ir.h
similarity index 96%
rename from include/tvm/relax/script/builder/ir.h
rename to include/tvm/relax/script/ir_builder/ir.h
index be008474c2..85091f540d 100644
--- a/include/tvm/relax/script/builder/ir.h
+++ b/include/tvm/relax/script/ir_builder/ir.h
@@ -16,11 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_RELAX_SCRIPT_BUILDER_IR_H_
-#define TVM_RELAX_SCRIPT_BUILDER_IR_H_
+#ifndef TVM_RELAX_SCRIPT_IR_BUILDER_IR_H_
+#define TVM_RELAX_SCRIPT_IR_BUILDER_IR_H_
#include <tvm/relax/expr.h>
-#include <tvm/relax/script/builder/frame.h>
+#include <tvm/relax/script/ir_builder/frame.h>
#include <tvm/relax/type.h>
#include <tvm/script/ir_builder/base.h>
@@ -164,4 +164,4 @@ ElseFrame Else();
} // namespace script
} // namespace tvm
-#endif // TVM_RELAX_SCRIPT_BUILDER_IR_H_
+#endif // TVM_RELAX_SCRIPT_IR_BUILDER_IR_H_
diff --git a/include/tvm/s_tir/script/builder/frame.h
b/include/tvm/s_tir/script/ir_builder/frame.h
similarity index 97%
rename from include/tvm/s_tir/script/builder/frame.h
rename to include/tvm/s_tir/script/ir_builder/frame.h
index 9f59569911..3630dd85c4 100644
--- a/include/tvm/s_tir/script/builder/frame.h
+++ b/include/tvm/s_tir/script/ir_builder/frame.h
@@ -16,11 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_S_TIR_SCRIPT_BUILDER_FRAME_H_
-#define TVM_S_TIR_SCRIPT_BUILDER_FRAME_H_
+#ifndef TVM_S_TIR_SCRIPT_IR_BUILDER_FRAME_H_
+#define TVM_S_TIR_SCRIPT_IR_BUILDER_FRAME_H_
#include <tvm/s_tir/stmt.h>
-#include <tvm/tirx/script/builder/frame.h>
+#include <tvm/tirx/script/ir_builder/frame.h>
namespace tvm {
namespace script {
@@ -175,4 +175,4 @@ class BlockInitFrame : public TIRFrame {
} // namespace script
} // namespace tvm
-#endif // TVM_S_TIR_SCRIPT_BUILDER_FRAME_H_
+#endif // TVM_S_TIR_SCRIPT_IR_BUILDER_FRAME_H_
diff --git a/include/tvm/s_tir/script/builder/ir.h
b/include/tvm/s_tir/script/ir_builder/ir.h
similarity index 96%
rename from include/tvm/s_tir/script/builder/ir.h
rename to include/tvm/s_tir/script/ir_builder/ir.h
index a14d2ac489..fac2b0fcd5 100644
--- a/include/tvm/s_tir/script/builder/ir.h
+++ b/include/tvm/s_tir/script/ir_builder/ir.h
@@ -16,11 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_S_TIR_SCRIPT_BUILDER_IR_H_
-#define TVM_S_TIR_SCRIPT_BUILDER_IR_H_
+#ifndef TVM_S_TIR_SCRIPT_IR_BUILDER_IR_H_
+#define TVM_S_TIR_SCRIPT_IR_BUILDER_IR_H_
-#include <tvm/s_tir/script/builder/frame.h>
-#include <tvm/tirx/script/builder/ir.h>
+#include <tvm/s_tir/script/ir_builder/frame.h>
+#include <tvm/tirx/script/ir_builder/ir.h>
namespace tvm {
namespace script {
@@ -148,4 +148,4 @@ ffi::Array<Var> Remap(ffi::String kinds,
ffi::Array<PrimExpr> bindings,
} // namespace script
} // namespace tvm
-#endif // TVM_S_TIR_SCRIPT_BUILDER_IR_H_
+#endif // TVM_S_TIR_SCRIPT_IR_BUILDER_IR_H_
diff --git a/include/tvm/script/ir_builder/ir/frame.h
b/include/tvm/script/ir_builder/frame.h
similarity index 95%
rename from include/tvm/script/ir_builder/ir/frame.h
rename to include/tvm/script/ir_builder/frame.h
index 797ececd1e..ce8086d5dc 100644
--- a/include/tvm/script/ir_builder/ir/frame.h
+++ b/include/tvm/script/ir_builder/frame.h
@@ -16,8 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
-#define TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
+#ifndef TVM_SCRIPT_IR_BUILDER_FRAME_H_
+#define TVM_SCRIPT_IR_BUILDER_FRAME_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
@@ -84,4 +84,4 @@ class IRModuleFrame : public IRBuilderFrame {
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
+#endif // TVM_SCRIPT_IR_BUILDER_FRAME_H_
diff --git a/include/tvm/script/ir_builder/ir/ir.h
b/include/tvm/script/ir_builder/ir.h
similarity index 92%
rename from include/tvm/script/ir_builder/ir/ir.h
rename to include/tvm/script/ir_builder/ir.h
index 761e2a995f..1a6e4c83f0 100644
--- a/include/tvm/script/ir_builder/ir/ir.h
+++ b/include/tvm/script/ir_builder/ir.h
@@ -16,12 +16,12 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_IR_BUILDER_IR_IR_H_
-#define TVM_SCRIPT_IR_BUILDER_IR_IR_H_
+#ifndef TVM_SCRIPT_IR_BUILDER_IR_H_
+#define TVM_SCRIPT_IR_BUILDER_IR_H_
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
-#include <tvm/script/ir_builder/ir/frame.h>
+#include <tvm/script/ir_builder/frame.h>
#include <vector>
@@ -58,4 +58,4 @@ TVM_DLL void DefFunction(const ffi::String& func_name, const
BaseFunc& func);
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_IR_BUILDER_IR_IR_H_
+#endif // TVM_SCRIPT_IR_BUILDER_IR_H_
diff --git a/include/tvm/script/printer/config.h
b/include/tvm/script/printer/config.h
index f8cb7cba8a..7f45b34cd5 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -155,7 +155,7 @@ class TVM_DLL PrinterConfig : public ffi::ObjectRef {
* \brief The fallback body used by TVM_REGISTER_SCRIPT_AS_REPR (defined in
printer.h).
*
* Tries to format \p obj via tvm::Script; on error falls back to a plain
- * address string. Defined in src/script/printer/config.cc so that
+ * address string. Defined in src/script/printer/script_printer.cc so that
* <tvm/runtime/logging.h> is not pulled into this public header.
*/
TVM_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj);
diff --git a/include/tvm/tirx/script/builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
similarity index 99%
rename from include/tvm/tirx/script/builder/frame.h
rename to include/tvm/tirx/script/ir_builder/frame.h
index d045209907..7836276b56 100644
--- a/include/tvm/tirx/script/builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -16,11 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
-#define TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
+#ifndef TVM_TIRX_SCRIPT_IR_BUILDER_FRAME_H_
+#define TVM_TIRX_SCRIPT_IR_BUILDER_FRAME_H_
#include <tvm/script/ir_builder/base.h>
-#include <tvm/script/ir_builder/ir/frame.h>
+#include <tvm/script/ir_builder/frame.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
@@ -581,4 +581,4 @@ class HintFrame : public TIRFrame {
} // namespace script
} // namespace tvm
-#endif // TVM_TIRX_SCRIPT_BUILDER_FRAME_H_
+#endif // TVM_TIRX_SCRIPT_IR_BUILDER_FRAME_H_
diff --git a/include/tvm/tirx/script/builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
similarity index 99%
rename from include/tvm/tirx/script/builder/ir.h
rename to include/tvm/tirx/script/ir_builder/ir.h
index 246891abb1..91533fe0d2 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -16,8 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
-#define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
+#ifndef TVM_TIRX_SCRIPT_IR_BUILDER_IR_H_
+#define TVM_TIRX_SCRIPT_IR_BUILDER_IR_H_
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/container/variant.h>
@@ -26,7 +26,7 @@
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/script/builder/frame.h>
+#include <tvm/tirx/script/ir_builder/frame.h>
#include <tvm/tirx/tile_primitive.h>
#include <tvm/tirx/type.h>
@@ -485,4 +485,4 @@ TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Void, PrimType::Void());
} // namespace script
} // namespace tvm
-#endif // TVM_TIRX_SCRIPT_BUILDER_IR_H_
+#endif // TVM_TIRX_SCRIPT_IR_BUILDER_IR_H_
diff --git a/src/relax/script/builder/distributed.cc
b/src/relax/script/ir_builder/distributed.cc
similarity index 97%
rename from src/relax/script/builder/distributed.cc
rename to src/relax/script/ir_builder/distributed.cc
index dc9440bdee..77e9b6f652 100644
--- a/src/relax/script/builder/distributed.cc
+++ b/src/relax/script/ir_builder/distributed.cc
@@ -20,7 +20,7 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/distributed/type.h>
-#include <tvm/relax/script/builder/ir.h>
+#include <tvm/relax/script/ir_builder/ir.h>
#include <tvm/relax/type.h>
#include <tvm/tirx/op.h>
diff --git a/src/relax/script/builder/frame.cc
b/src/relax/script/ir_builder/frame.cc
similarity index 99%
rename from src/relax/script/builder/frame.cc
rename to src/relax/script/ir_builder/frame.cc
index 78399a0c91..0c5ec90534 100644
--- a/src/relax/script/builder/frame.cc
+++ b/src/relax/script/ir_builder/frame.cc
@@ -21,8 +21,8 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
-#include <tvm/relax/script/builder/frame.h>
-#include <tvm/relax/script/builder/ir.h>
+#include <tvm/relax/script/ir_builder/frame.h>
+#include <tvm/relax/script/ir_builder/ir.h>
#include "./utils.h"
diff --git a/src/relax/script/builder/ir.cc b/src/relax/script/ir_builder/ir.cc
similarity index 99%
rename from src/relax/script/builder/ir.cc
rename to src/relax/script/ir_builder/ir.cc
index eec633205e..4104523ac9 100644
--- a/src/relax/script/builder/ir.cc
+++ b/src/relax/script/ir_builder/ir.cc
@@ -18,12 +18,12 @@
*/
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
-#include <tvm/relax/script/builder/ir.h>
+#include <tvm/relax/script/ir_builder/ir.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/op.h>
-#include "../../../script/ir_builder/ir/utils.h"
+#include "../../../script/ir_builder/utils.h"
#include "./utils.h"
namespace tvm {
diff --git a/src/relax/script/builder/utils.h
b/src/relax/script/ir_builder/utils.h
similarity index 96%
rename from src/relax/script/builder/utils.h
rename to src/relax/script/ir_builder/utils.h
index 247dd91aae..dd4349b02a 100644
--- a/src/relax/script/builder/utils.h
+++ b/src/relax/script/ir_builder/utils.h
@@ -16,11 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_RELAX_SCRIPT_BUILDER_UTILS_H_
-#define TVM_RELAX_SCRIPT_BUILDER_UTILS_H_
+#ifndef TVM_RELAX_SCRIPT_IR_BUILDER_UTILS_H_
+#define TVM_RELAX_SCRIPT_IR_BUILDER_UTILS_H_
-#include <tvm/relax/script/builder/frame.h>
-#include <tvm/relax/script/builder/ir.h>
+#include <tvm/relax/script/ir_builder/frame.h>
+#include <tvm/relax/script/ir_builder/ir.h>
#include <tvm/relax/type_functor.h>
#include <tvm/relax/utils.h>
@@ -139,4 +139,4 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const
SeqExprFrame& frame, ffi::S
} // namespace script
} // namespace tvm
-#endif // TVM_RELAX_SCRIPT_BUILDER_UTILS_H_
+#endif // TVM_RELAX_SCRIPT_IR_BUILDER_UTILS_H_
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index d30e932234..c1901459e5 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -22,6 +22,7 @@
#include <cmath>
#include <limits>
+#include "../../../script/printer/dialect_prefix.h"
#include "../../../tirx/script/printer/utils.h"
#include "./utils.h"
@@ -29,6 +30,8 @@ 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 {
diff --git a/src/s_tir/script/builder/frame.cc
b/src/s_tir/script/ir_builder/frame.cc
similarity index 99%
rename from src/s_tir/script/builder/frame.cc
rename to src/s_tir/script/ir_builder/frame.cc
index 76b05b7eaf..b6fe374568 100644
--- a/src/s_tir/script/builder/frame.cc
+++ b/src/s_tir/script/ir_builder/frame.cc
@@ -16,7 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/s_tir/script/builder/frame.h>
+#include <tvm/s_tir/script/ir_builder/frame.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/function.h>
diff --git a/src/s_tir/script/builder/ir.cc b/src/s_tir/script/ir_builder/ir.cc
similarity index 99%
rename from src/s_tir/script/builder/ir.cc
rename to src/s_tir/script/ir_builder/ir.cc
index 8579774b30..a6698f4e2e 100644
--- a/src/s_tir/script/builder/ir.cc
+++ b/src/s_tir/script/ir_builder/ir.cc
@@ -18,7 +18,7 @@
*/
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/s_tir/script/builder/ir.h>
+#include <tvm/s_tir/script/ir_builder/ir.h>
#include <tvm/sym/analyzer.h>
#include "./utils.h"
diff --git a/src/s_tir/script/builder/script_complete.cc
b/src/s_tir/script/ir_builder/script_complete.cc
similarity index 99%
rename from src/s_tir/script/builder/script_complete.cc
rename to src/s_tir/script/ir_builder/script_complete.cc
index 7f9ea20eb4..d3675fd7e1 100644
--- a/src/s_tir/script/builder/script_complete.cc
+++ b/src/s_tir/script/ir_builder/script_complete.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file s_tir/script/builder/script_complete.cc
+ * \file s_tir/script/ir_builder/script_complete.cc
* \brief Used by TVM Script parser to expand incomplete TIR input
*/
diff --git a/src/s_tir/script/builder/script_complete.h
b/src/s_tir/script/ir_builder/script_complete.h
similarity index 84%
rename from src/s_tir/script/builder/script_complete.h
rename to src/s_tir/script/ir_builder/script_complete.h
index 12191a750d..f752262b09 100644
--- a/src/s_tir/script/builder/script_complete.h
+++ b/src/s_tir/script/ir_builder/script_complete.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file s_tir/script/builder/script_complete.h
+ * \file s_tir/script/ir_builder/script_complete.h
* \brief Used by TVM Script parser to expand incomplete TIR input
*/
-#ifndef TVM_S_TIR_SCRIPT_BUILDER_SCRIPT_COMPLETE_H_
-#define TVM_S_TIR_SCRIPT_BUILDER_SCRIPT_COMPLETE_H_
+#ifndef TVM_S_TIR_SCRIPT_IR_BUILDER_SCRIPT_COMPLETE_H_
+#define TVM_S_TIR_SCRIPT_IR_BUILDER_SCRIPT_COMPLETE_H_
#include <tvm/ffi/function.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
@@ -35,4 +35,4 @@ tvm::tirx::PrimFunc ScriptComplete(tvm::tirx::PrimFunc func,
} // namespace s_tir
} // namespace tvm
-#endif // TVM_S_TIR_SCRIPT_BUILDER_SCRIPT_COMPLETE_H_
+#endif // TVM_S_TIR_SCRIPT_IR_BUILDER_SCRIPT_COMPLETE_H_
diff --git a/src/s_tir/script/builder/utils.h
b/src/s_tir/script/ir_builder/utils.h
similarity index 89%
rename from src/s_tir/script/builder/utils.h
rename to src/s_tir/script/ir_builder/utils.h
index 40fbc591a8..f91c118782 100644
--- a/src/s_tir/script/builder/utils.h
+++ b/src/s_tir/script/ir_builder/utils.h
@@ -16,12 +16,12 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_S_TIR_SCRIPT_BUILDER_UTILS_H_
-#define TVM_S_TIR_SCRIPT_BUILDER_UTILS_H_
+#ifndef TVM_S_TIR_SCRIPT_IR_BUILDER_UTILS_H_
+#define TVM_S_TIR_SCRIPT_IR_BUILDER_UTILS_H_
-#include <tvm/s_tir/script/builder/frame.h>
+#include <tvm/s_tir/script/ir_builder/frame.h>
-#include "../../../tirx/script/builder/utils.h"
+#include "../../../tirx/script/ir_builder/utils.h"
namespace tvm {
namespace script {
@@ -58,4 +58,4 @@ inline SBlockFrame FindSBlockFrame(const ffi::String& method)
{
} // namespace script
} // namespace tvm
-#endif // TVM_S_TIR_SCRIPT_BUILDER_UTILS_H_
+#endif // TVM_S_TIR_SCRIPT_IR_BUILDER_UTILS_H_
diff --git a/src/script/printer/ir/distributed.cc
b/src/s_tir/script/printer/config.cc
similarity index 64%
copy from src/script/printer/ir/distributed.cc
copy to src/s_tir/script/printer/config.cc
index 89db9a70d7..461858edad 100644
--- a/src/script/printer/ir/distributed.cc
+++ b/src/s_tir/script/printer/config.cc
@@ -16,26 +16,15 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/container/shape.h>
-#include <tvm/ir/expr.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include "../../../script/printer/dialect_prefix.h"
-#include "./utils.h"
namespace tvm {
namespace script {
namespace printer {
-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);
- });
-}
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("s_tir.prefix", "Ts"); }
} // namespace printer
} // namespace script
diff --git a/src/script/ir_builder/ir/frame.cc b/src/script/ir_builder/frame.cc
similarity index 98%
rename from src/script/ir_builder/ir/frame.cc
rename to src/script/ir_builder/frame.cc
index 284cd7e8d9..0684a32be5 100644
--- a/src/script/ir_builder/ir/frame.cc
+++ b/src/script/ir_builder/frame.cc
@@ -19,7 +19,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
-#include <tvm/script/ir_builder/ir/frame.h>
+#include <tvm/script/ir_builder/frame.h>
namespace tvm {
namespace script {
diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir.cc
similarity index 99%
rename from src/script/ir_builder/ir/ir.cc
rename to src/script/ir_builder/ir.cc
index e351dde55a..e7523211d3 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir.cc
@@ -21,7 +21,7 @@
#include <tvm/ir/global_info.h>
#include <tvm/ir/module.h>
#include <tvm/runtime/logging.h>
-#include <tvm/script/ir_builder/ir/ir.h>
+#include <tvm/script/ir_builder/ir.h>
#include "./utils.h"
diff --git a/src/script/ir_builder/ir/utils.h b/src/script/ir_builder/utils.h
similarity index 92%
rename from src/script/ir_builder/ir/utils.h
rename to src/script/ir_builder/utils.h
index 397b54a822..d0278a51ea 100644
--- a/src/script/ir_builder/ir/utils.h
+++ b/src/script/ir_builder/utils.h
@@ -16,10 +16,10 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
-#define TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
+#ifndef TVM_SCRIPT_IR_BUILDER_UTILS_H_
+#define TVM_SCRIPT_IR_BUILDER_UTILS_H_
-#include <tvm/script/ir_builder/ir/frame.h>
+#include <tvm/script/ir_builder/frame.h>
namespace tvm {
namespace script {
@@ -57,4 +57,4 @@ inline IRModuleFrame FindModuleFrame() {
} // namespace script
} // namespace tvm
-#endif // TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
+#endif // TVM_SCRIPT_IR_BUILDER_UTILS_H_
diff --git a/src/script/printer/config.cc b/src/script/printer/config.cc
index 87ca87979f..993e4b3f44 100644
--- a/src/script/printer/config.cc
+++ b/src/script/printer/config.cc
@@ -16,23 +16,164 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/runtime/logging.h>
-#include <tvm/script/printer/printer.h>
+#include <tvm/script/printer/config.h>
-#include <sstream>
+#include <algorithm>
+#include <cctype>
+#include <map>
+#include <string>
+#include <unordered_set>
+
+#include "./dialect_prefix.h"
namespace tvm {
+namespace {
+
+std::map<ffi::String, ffi::String>& DialectPrefixes() {
+ static std::map<ffi::String, ffi::String> prefixes;
+ return prefixes;
+}
+
+bool IsIdentifier(const std::string& name) {
+ // Python identifiers follow the regex: "^[a-zA-Z_][a-zA-Z0-9_]*$"
+ // `std::regex` would cause a symbol conflict with PyTorch, we avoids to use
it in the codebase.
+ //
+ // We convert the regex into following conditions:
+ // 1. The name is not empty.
+ // 2. The first character is either an alphabet or an underscore.
+ // 3. The rest of the characters are either an alphabet, a digit or an
underscore.
+ return name.size() > 0 && //
+ (std::isalpha(name[0]) || name[0] == '_') && //
+ std::all_of(name.begin() + 1, name.end(),
+ [](char c) { return std::isalnum(c) || c == '_'; });
+}
+
+} // namespace
+
+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;
+}
+
+} // namespace printer
+} // namespace script
-std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
- try {
- return tvm::Script(obj, std::nullopt);
- } catch (const tvm::ffi::Error& e) {
- LOG(WARNING) << "TVMScript printer falls back to the basic address printer
with the error:\n"
- << e.what();
- std::ostringstream os;
- os << obj->GetTypeKey() << '(' << obj.get() << ')';
- return os.str();
+PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
+ ffi::ObjectPtr<PrinterConfigNode> n = ffi::make_object<PrinterConfigNode>();
+ std::unordered_set<ffi::String> core_keys;
+ auto get = [&](const char* key) {
+ core_keys.insert(key);
+ return config_dict.Get(key);
+ };
+ if (auto v = get("name")) {
+ n->binding_names.push_back(v.value().as_or_throw<ffi::String>());
+ }
+ if (auto v = get("show_meta")) {
+ n->show_meta = v.value().cast<bool>();
+ }
+ if (auto v = get("ir_prefix")) {
+ n->ir_prefix = v.value().as_or_throw<ffi::String>();
+ }
+ if (auto v = get("module_alias")) {
+ n->module_alias = v.value().as_or_throw<ffi::String>();
+ }
+ if (auto v = get("buffer_dtype")) {
+ 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>());
+ }
+ if (auto v = get("float_dtype")) {
+ n->float_dtype =
ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
+ }
+ if (auto v = get("verbose_expr")) {
+ n->verbose_expr = v.value().cast<bool>();
+ }
+ if (auto v = get("indent_spaces")) {
+ n->indent_spaces = v.value().cast<int>();
+ }
+ if (auto v = get("print_line_numbers")) {
+ n->print_line_numbers = v.value().cast<bool>();
}
+ if (auto v = get("num_context_lines")) {
+ n->num_context_lines = v.value().cast<int>();
+ }
+ if (auto v = get("path_to_underline")) {
+ n->path_to_underline =
+
v.value().as_or_throw<ffi::Optional<ffi::Array<ffi::reflection::AccessPath>>>().value_or(
+ ffi::Array<ffi::reflection::AccessPath>());
+ }
+ if (auto v = get("path_to_annotate")) {
+ n->path_to_annotate =
+ v.value()
+ .as_or_throw<ffi::Optional<ffi::Map<ffi::reflection::AccessPath,
ffi::String>>>()
+ .value_or(ffi::Map<ffi::reflection::AccessPath, ffi::String>());
+ }
+ if (auto v = get("obj_to_underline")) {
+ n->obj_to_underline =
+
v.value().as_or_throw<ffi::Optional<ffi::Array<ffi::ObjectRef>>>().value_or(
+ ffi::Array<ffi::ObjectRef>());
+ }
+ if (auto v = get("obj_to_annotate")) {
+ n->obj_to_annotate =
+ 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>();
+ }
+ auto extra_config = get("extra_config");
+ for (const auto& kv : config_dict) {
+ if (!core_keys.count(kv.first)) {
+ n->extra_config.Set(kv.first, kv.second);
+ }
+ }
+ if (extra_config) {
+ auto extra = extra_config.value().as_or_throw<ffi::Map<ffi::String,
ffi::Any>>();
+ for (auto kv : extra) {
+ n->extra_config.Set(kv.first, kv.second);
+ }
+ if (auto render = extra.Get("render_invisible_path_info")) {
+ n->render_invisible_path_info = render.value().cast<bool>();
+ }
+ }
+
+ // Validate all registered prefixes before names can be assigned by a
docsifier.
+ n->GetBuiltinKeywords();
+
+ this->data_ = std::move(n);
+}
+
+ffi::Array<ffi::String> PrinterConfigNode::GetBuiltinKeywords() {
+ TVM_FFI_ICHECK(IsIdentifier(std::string(ir_prefix))) << "Invalid
`ir_prefix`: " << ir_prefix;
+ 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()) {
+ ffi::String prefix = GetExtraConfig<ffi::String>(key, default_prefix);
+ TVM_FFI_ICHECK(IsIdentifier(std::string(prefix))) << "Invalid `" << key <<
"`: " << prefix;
+ result.push_back(prefix);
+ }
+ if (!module_alias.empty()) {
+ result.push_back(module_alias);
+ }
+ return result;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ PrinterConfigNode::RegisterReflection();
+ ffi::reflection::GlobalDef().def(
+ "node.PrinterConfig",
+ [](ffi::Map<ffi::String, Any> config_dict) { return
PrinterConfig(config_dict); });
}
} // namespace tvm
diff --git a/src/script/printer/ir/distributed.cc
b/src/script/printer/dialect_prefix.h
similarity index 64%
rename from src/script/printer/ir/distributed.cc
rename to src/script/printer/dialect_prefix.h
index 89db9a70d7..5aaae11293 100644
--- a/src/script/printer/ir/distributed.cc
+++ b/src/script/printer/dialect_prefix.h
@@ -16,27 +16,21 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ffi/container/shape.h>
-#include <tvm/ir/expr.h>
+#ifndef TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
+#define TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
+
+#include <tvm/ffi/string.h>
-#include "./utils.h"
namespace tvm {
namespace script {
namespace printer {
-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);
- });
-}
+// 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);
} // namespace printer
} // namespace script
} // namespace tvm
+
+#endif // TVM_SCRIPT_PRINTER_DIALECT_PREFIX_H_
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index f697bd9aaf..d4d41342b6 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -198,15 +198,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-std::string ReprPrintIRModule(const ffi::ObjectRef& mod, const PrinterConfig&
cfg) {
- return ReprPrintIR(mod, cfg);
-}
-
TVM_REGISTER_SCRIPT_AS_REPR(GlobalVarNode, ReprPrintIR);
TVM_REGISTER_SCRIPT_AS_REPR(DictAttrsNode, ReprPrintIR);
TVM_REGISTER_SCRIPT_AS_REPR(FuncTypeNode, ReprPrintIR);
TVM_REGISTER_SCRIPT_AS_REPR(RangeNode, ReprPrintIR);
-TVM_REGISTER_SCRIPT_AS_REPR(IRModuleNode, ReprPrintIRModule);
+TVM_REGISTER_SCRIPT_AS_REPR(IRModuleNode, ReprPrintIR);
} // namespace printer
} // namespace script
diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index 8792225a17..53e57f8afb 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/container/shape.h>
+
#include "./utils.h"
namespace tvm {
@@ -65,6 +67,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+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);
+ });
+}
+
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/ir/prim_expr.cc
b/src/script/printer/ir/prim_expr.cc
new file mode 100644
index 0000000000..649c2c5826
--- /dev/null
+++ b/src/script/printer/ir/prim_expr.cc
@@ -0,0 +1,201 @@
+/*
+ * 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/ir/prim/op.h>
+
+#include "./utils.h"
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+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});
+ });
+}
+
+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});
+ });
+}
+
+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"));
+ }
+ });
+}
+
+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});
+ });
+}
+
+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")),
+ });
+ });
+}
+
+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")),
+ });
+ });
+}
+
+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")),
+ });
+ });
+}
+
+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")),
+ });
+ });
+}
+
+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}); \
+ });
+
+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});
\
+ });
+
+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
+
+} // namespace printer
+} // namespace script
+} // namespace tvm
diff --git a/src/tirx/script/printer/ir.cc b/src/script/printer/ir/prim_type.cc
similarity index 52%
copy from src/tirx/script/printer/ir.cc
copy to src/script/printer/ir/prim_type.cc
index f9086d941b..3816143808 100644
--- a/src/tirx/script/printer/ir.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -16,8 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/target/target.h>
-#include <tvm/tirx/type.h>
+#include <tvm/ir/type.h>
#include "./utils.h"
@@ -25,8 +24,6 @@ namespace tvm {
namespace script {
namespace printer {
-TVM_FFI_STATIC_INIT_BLOCK() { TIRFrameNode::RegisterReflection(); }
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<IntImm>(
"", [](IntImm imm, AccessPath imm_p, IRDocsifier d) -> Doc {
@@ -55,17 +52,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Range>(
- "tirx", [](Range range, AccessPath p, IRDocsifier d) -> Doc {
- return TIR(d, "Range")
- ->Call({
- d->AsDoc<ExprDoc>(range->min, p->Attr("min")),
- d->AsDoc<ExprDoc>(range->extent + range->min,
p->Attr("extent")),
- });
- });
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<PrimType>(
"",
@@ -79,38 +65,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<PointerType>(
- "", [](PointerType ty, AccessPath ty_p, IRDocsifier d) -> Doc {
- ExprDoc element_type{ffi::UnsafeInit()};
- TVM_FFI_ICHECK(!ty->element_type.IsMissing())
- << "InternalError: PointerType.element_type is missing";
- if (const auto* prim_type = ty->element_type.as<PrimTypeNode>()) {
- if (ffi::GetRef<PrimType>(prim_type).IsVoid()) {
- if (ty->storage_scope == "global") {
- return TIR(d, "handle");
- }
- return TIR(d, "handle")
- ->Call({}, {"storage_scope"},
- {LiteralDoc::Str(ty->storage_scope,
ty_p->Attr("storage_scope"))});
- }
- element_type = LiteralDoc::DataType(prim_type->dtype, //
-
ty_p->Attr("element_type")->Attr("dtype"));
- } else if (ty->element_type.as<tirx::TensorMapTypeNode>()) {
- return TIR(d, "TensorMap")->Call({});
- } else {
- element_type = d->AsDoc<ExprDoc>(ty->element_type,
ty_p->Attr("element_type"));
- }
- if (ty->storage_scope == "") {
- return TIR(d, "handle")->Call({element_type});
- } else {
- return TIR(d, "handle")
- ->Call(
- {element_type, LiteralDoc::Str(ty->storage_scope,
ty_p->Attr("storage_scope"))});
- }
- });
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<TupleType>(
"", [](TupleType ty, AccessPath p, IRDocsifier d) -> Doc {
@@ -120,21 +74,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return TIR(d, "Tuple")->Call(d->AsDoc<ListDoc>(ty->fields,
p->Attr("fields"))->elements);
});
}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<Target>(
- "", [](Target target, AccessPath p, IRDocsifier d) -> Doc {
- ffi::Map<ffi::String, ffi::Any> config = target->ToConfig();
- return TIR(d, "target")->Call({d->AsDoc<ExprDoc>(config, p)});
- });
-}
-
-TVM_REGISTER_SCRIPT_AS_REPR(IntImmNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(FloatImmNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(PrimTypeNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(PointerTypeNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(TupleTypeNode, ReprPrintTIR);
-
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index b0ea44ef9d..a58d52c439 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -20,9 +20,9 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
+#include <tvm/runtime/logging.h>
#include <tvm/script/printer/printer.h>
-#include <algorithm>
#include <sstream>
namespace tvm {
@@ -47,8 +47,6 @@ std::string RenderFallbackWithInvisiblePathInfo(const
ffi::String& script,
} // namespace
-TVM_FFI_STATIC_INIT_BLOCK() { PrinterConfigNode::RegisterReflection(); }
-
TVMScriptPrinter::FType& TVMScriptPrinter::vtable() {
static FType inst;
return inst;
@@ -63,140 +61,20 @@ std::string Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig
return TVMScriptPrinter::vtable()(node, config);
}
-bool IsIdentifier(const std::string& name) {
- // Python identifiers follow the regex: "^[a-zA-Z_][a-zA-Z0-9_]*$"
- // `std::regex` would cause a symbol conflict with PyTorch, we avoids to use
it in the codebase.
- //
- // We convert the regex into following conditions:
- // 1. The name is not empty.
- // 2. The first character is either an alphabet or an underscore.
- // 3. The rest of the characters are either an alphabet, a digit or an
underscore.
- return name.size() > 0 && //
- (std::isalpha(name[0]) || name[0] == '_') && //
- std::all_of(name.begin() + 1, name.end(),
- [](char c) { return std::isalnum(c) || c == '_'; });
-}
-
-PrinterConfig::PrinterConfig(ffi::Map<ffi::String, Any> config_dict) {
- ffi::ObjectPtr<PrinterConfigNode> n = ffi::make_object<PrinterConfigNode>();
- if (auto v = config_dict.Get("name")) {
- n->binding_names.push_back(v.value().as_or_throw<ffi::String>());
- }
- if (auto v = config_dict.Get("show_meta")) {
- n->show_meta = v.value().cast<bool>();
- }
- if (auto v = config_dict.Get("ir_prefix")) {
- n->ir_prefix = v.value().as_or_throw<ffi::String>();
- }
- if (auto v = config_dict.Get("module_alias")) {
- n->module_alias = v.value().as_or_throw<ffi::String>();
- }
- if (auto v = config_dict.Get("buffer_dtype")) {
- n->buffer_dtype =
ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
- }
- if (auto v = config_dict.Get("int_dtype")) {
- n->int_dtype =
ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
- }
- if (auto v = config_dict.Get("float_dtype")) {
- n->float_dtype =
ffi::StringToDLDataType(v.value().as_or_throw<ffi::String>());
- }
- if (auto v = config_dict.Get("verbose_expr")) {
- n->verbose_expr = v.value().cast<bool>();
- }
- if (auto v = config_dict.Get("indent_spaces")) {
- n->indent_spaces = v.value().cast<int>();
- }
- if (auto v = config_dict.Get("print_line_numbers")) {
- n->print_line_numbers = v.value().cast<bool>();
- }
- if (auto v = config_dict.Get("num_context_lines")) {
- n->num_context_lines = v.value().cast<int>();
- }
- if (auto v = config_dict.Get("path_to_underline")) {
- n->path_to_underline =
-
v.value().as_or_throw<ffi::Optional<ffi::Array<ffi::reflection::AccessPath>>>().value_or(
- ffi::Array<ffi::reflection::AccessPath>());
- }
- if (auto v = config_dict.Get("path_to_annotate")) {
- n->path_to_annotate =
- v.value()
- .as_or_throw<ffi::Optional<ffi::Map<ffi::reflection::AccessPath,
ffi::String>>>()
- .value_or(ffi::Map<ffi::reflection::AccessPath, ffi::String>());
- }
- if (auto v = config_dict.Get("obj_to_underline")) {
- n->obj_to_underline =
-
v.value().as_or_throw<ffi::Optional<ffi::Array<ffi::ObjectRef>>>().value_or(
- ffi::Array<ffi::ObjectRef>());
- }
- if (auto v = config_dict.Get("obj_to_annotate")) {
- n->obj_to_annotate =
- v.value().as_or_throw<ffi::Optional<ffi::Map<ffi::ObjectRef,
ffi::String>>>().value_or(
- ffi::Map<ffi::ObjectRef, ffi::String>());
- }
- if (auto v = config_dict.Get("syntax_sugar")) {
- n->syntax_sugar = v.value().cast<bool>();
- }
- if (auto v = config_dict.Get("show_object_address")) {
- n->show_object_address = v.value().cast<bool>();
- }
- if (auto v = config_dict.Get("render_invisible_path_info")) {
- n->render_invisible_path_info = v.value().cast<bool>();
- }
- // Dialect-specific keys are stored in extra_config with dotted-name keys.
- // String-typed dialect keys passed through directly.
- for (const char* key : {"tirx.prefix", "s_tir.prefix", "relax.prefix"}) {
- if (auto v = config_dict.Get(key)) {
- n->extra_config.Set(ffi::String(key), v.value());
- }
- }
- // Boolean dialect keys.
- if (auto v = config_dict.Get("relax.show_all_ty")) {
- n->extra_config.Set(ffi::String("relax.show_all_ty"), v.value());
- }
- if (auto v = config_dict.Get("extra_config")) {
- auto extra = v.value().as_or_throw<ffi::Map<ffi::String, ffi::Any>>();
- for (auto kv : extra) {
- n->extra_config.Set(kv.first, kv.second);
- }
- if (auto render = extra.Get("render_invisible_path_info")) {
- n->render_invisible_path_info = render.value().cast<bool>();
- }
- }
-
- // Checking prefixes if they are valid Python identifiers.
- TVM_FFI_ICHECK(IsIdentifier(std::string(n->ir_prefix)))
- << "Invalid `ir_prefix`: " << n->ir_prefix;
- ffi::String tir_prefix = n->GetExtraConfig<ffi::String>("tirx.prefix", "T");
- ffi::String stir_prefix = n->GetExtraConfig<ffi::String>("s_tir.prefix",
"Ts");
- TVM_FFI_ICHECK(IsIdentifier(std::string(stir_prefix)))
- << "Invalid `s_tir.prefix`: " << stir_prefix;
- ffi::String relax_prefix = n->GetExtraConfig<ffi::String>("relax.prefix",
"R");
- TVM_FFI_ICHECK(IsIdentifier(std::string(tir_prefix))) << "Invalid
`tirx.prefix`: " << tir_prefix;
- TVM_FFI_ICHECK(IsIdentifier(std::string(relax_prefix)))
- << "Invalid `relax.prefix`: " << relax_prefix;
- TVM_FFI_ICHECK(n->module_alias.empty() ||
IsIdentifier(std::string(n->module_alias)))
- << "Invalid `module_alias`: " << n->module_alias;
-
- this->data_ = std::move(n);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("node.TVMScriptPrinterScript", tvm::Script);
}
-ffi::Array<ffi::String> PrinterConfigNode::GetBuiltinKeywords() {
- ffi::String tir_prefix = GetExtraConfig<ffi::String>("tirx.prefix", "T");
- ffi::String relax_prefix = GetExtraConfig<ffi::String>("relax.prefix", "R");
- ffi::Array<ffi::String> result{this->ir_prefix, tir_prefix, relax_prefix,
- GetExtraConfig<ffi::String>("s_tir.prefix",
"Ts")};
- if (!this->module_alias.empty()) {
- result.push_back(this->module_alias);
+std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
+ try {
+ return tvm::Script(obj, std::nullopt);
+ } catch (const tvm::ffi::Error& e) {
+ LOG(WARNING) << "TVMScript printer falls back to the basic address printer
with the error:\n"
+ << e.what();
+ std::ostringstream os;
+ os << obj->GetTypeKey() << '(' << obj.get() << ')';
+ return os.str();
}
- return result;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("node.PrinterConfig",
- [](ffi::Map<ffi::String, Any> config_dict) { return
PrinterConfig(config_dict); })
- .def("node.TVMScriptPrinterScript", tvm::Script);
}
} // namespace tvm
diff --git a/src/script/printer/utils.h b/src/script/printer/utils.h
index 9e9c1bcab3..7441906ee3 100644
--- a/src/script/printer/utils.h
+++ b/src/script/printer/utils.h
@@ -42,12 +42,6 @@ namespace tvm {
namespace script {
namespace printer {
-// Note: the `TVM_SCRIPT_REPR` macro is intentionally duplicated in each
-// dialect-local `src/<dialect>/script/printer/utils.h`. Keeping a single
-// definition here would force the dialect headers to depend on this shared
-// header, which the per-dialect restructure aims to avoid for cross-directory
-// references. See each `<dialect>/script/printer/utils.h` for the macro.
-
inline std::string Docsify(const ffi::ObjectRef& obj, const IRDocsifier& d,
const Frame& f,
const PrinterConfig& cfg) {
Doc doc = d->AsDoc(obj, AccessPath::Root());
@@ -115,9 +109,6 @@ inline ExprDoc STIR(const IRDocsifier& d, const
ffi::String& attr) {
return IdDoc(d->cfg->GetExtraConfig<ffi::String>("s_tir.prefix",
"Ts"))->Attr(attr);
}
-/*! \brief Alias for TIR — historical TIRx name used by tirx printer code */
-inline ExprDoc TIRx(const IRDocsifier& d, const ffi::String& attr) { return
TIR(d, attr); }
-
/*! \brief Creates the Relax common prefix, which is by default `R` */
inline ExprDoc Relax(const IRDocsifier& d, const ffi::String& attr) {
d->ir_usage.insert("relax");
diff --git a/src/tirx/script/builder/frame.cc
b/src/tirx/script/ir_builder/frame.cc
similarity index 99%
rename from src/tirx/script/builder/frame.cc
rename to src/tirx/script/ir_builder/frame.cc
index 25e8faa996..cc0acfc322 100644
--- a/src/tirx/script/builder/frame.cc
+++ b/src/tirx/script/ir_builder/frame.cc
@@ -20,12 +20,12 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/runtime/logging.h>
-#include <tvm/script/ir_builder/ir/ir.h>
+#include <tvm/script/ir_builder/ir.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/script/builder/frame.h>
+#include <tvm/tirx/script/ir_builder/frame.h>
#include <map>
diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
similarity index 99%
rename from src/tirx/script/builder/ir.cc
rename to src/tirx/script/ir_builder/ir.cc
index 14a23e1698..e569af59f4 100644
--- a/src/tirx/script/builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -30,7 +30,7 @@
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
-#include <tvm/tirx/script/builder/ir.h>
+#include <tvm/tirx/script/ir_builder/ir.h>
#include <tvm/tirx/tile_primitive.h>
#include "./utils.h"
diff --git a/src/tirx/script/builder/utils.h
b/src/tirx/script/ir_builder/utils.h
similarity index 95%
rename from src/tirx/script/builder/utils.h
rename to src/tirx/script/ir_builder/utils.h
index b3ed8d224e..f4b9425311 100644
--- a/src/tirx/script/builder/utils.h
+++ b/src/tirx/script/ir_builder/utils.h
@@ -16,14 +16,14 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_TIRX_SCRIPT_BUILDER_UTILS_H_
-#define TVM_TIRX_SCRIPT_BUILDER_UTILS_H_
+#ifndef TVM_TIRX_SCRIPT_IR_BUILDER_UTILS_H_
+#define TVM_TIRX_SCRIPT_IR_BUILDER_UTILS_H_
#include <tvm/ffi/cast.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/script/builder/frame.h>
-#include <tvm/tirx/script/builder/ir.h>
+#include <tvm/tirx/script/ir_builder/frame.h>
+#include <tvm/tirx/script/ir_builder/ir.h>
#include <tvm/tirx/stmt.h>
namespace tvm {
@@ -130,4 +130,4 @@ inline tvm::TensorRegion
BufferRegionFromLoad(tvm::TensorLoad buffer_load) {
} // namespace script
} // namespace tvm
-#endif // TVM_TIRX_SCRIPT_BUILDER_UTILS_H_
+#endif // TVM_TIRX_SCRIPT_IR_BUILDER_UTILS_H_
diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc
index eecef93677..e6b392f3cd 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/tirx/script/printer/block.cc
@@ -283,7 +283,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_SCRIPT_REPR(tirx::ScopeIdDefStmtNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ScopeIdDefStmtNode, ReprPrintTIR);
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::ExecScope>(
@@ -292,7 +292,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return doc;
});
}
-TVM_SCRIPT_REPR(tirx::ExecScopeNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ExecScopeNode, ReprPrintTIR);
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::ScopeIdDef>(
@@ -308,7 +308,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return doc;
});
}
-TVM_SCRIPT_REPR(tirx::ScopeIdDefNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ScopeIdDefNode, ReprPrintTIR);
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 6f75ab0829..9218bd4603 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -470,7 +470,7 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier d,
AccessPath p) {
keys.push_back("offset");
values.push_back(DictDoc(offset_keys, offset_values));
}
- return TIRx(d, "TileLayout")->Attr("from_iters")->Call({}, keys, values);
+ return TIR(d, "TileLayout")->Attr("from_iters")->Call({}, keys, values);
}
// Compose `Tx.S[..] [+ Tx.R[..]] [+ offset_expr]`.
@@ -484,10 +484,10 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier
d, AccessPath p) {
ffi::Optional<ExprDoc> spec;
if (layout->shard.size() > 0) {
- add_term(spec, iters_to_index(TIRx(d, "S"), layout->shard));
+ add_term(spec, iters_to_index(TIR(d, "S"), layout->shard));
}
if (layout->replica.size() > 0) {
- add_term(spec, iters_to_index(TIRx(d, "R"), layout->replica));
+ add_term(spec, iters_to_index(TIR(d, "R"), layout->replica));
}
if (layout->offset.size() > 0) {
// Sort by axis name so the printed text is deterministic across builds
@@ -515,7 +515,7 @@ Doc PrintTileLayout(tirx::TileLayout layout, IRDocsifier d,
AccessPath p) {
add_term(spec, off_doc.value());
}
- return TIRx(d, "TileLayout")->Call({spec.value()}, {}, {});
+ return TIR(d, "TileLayout")->Call({spec.value()}, {}, {});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -541,7 +541,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kwargs_values.push_back(
LiteralDoc::Boolean(layout->swizzle_inner,
p->Attr("swizzle_inner")));
}
- return TIRx(d, "ComposeLayout")
+ return TIR(d, "ComposeLayout")
->Call({per_element, swizzle_len, atom_len, tile_doc},
kwargs_keys, kwargs_values);
});
}
@@ -558,14 +558,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_SCRIPT_REPR(tvm::TensorRegionNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(TensorLoadNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::BufferStoreNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::BufferTypeNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::IterNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::TileLayoutNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::ComposeLayoutNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(s_tir::MatchBufferRegionNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tvm::TensorRegionNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(TensorLoadNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::BufferStoreNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::BufferTypeNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::IterNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::TileLayoutNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ComposeLayoutNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(s_tir::MatchBufferRegionNode, ReprPrintTIR);
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index ffc6145dd9..7bfb7fd50e 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -153,93 +153,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-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});
- });
-}
-
-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});
- });
-}
-
-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"));
- }
- });
-}
-
-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});
- });
-}
-
-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")),
- });
- });
-}
-
-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")),
- });
- });
-}
-
-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")),
- });
- });
-}
-
-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")),
- });
- });
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<te::CommReducer>( //
"", [](te::CommReducer r, AccessPath p, IRDocsifier d) -> Doc {
@@ -325,16 +238,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-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});
- });
-}
-
Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier d) {
if (call->op.same_as(tirx::builtin::buffer_data())) {
TVM_FFI_ICHECK_EQ(call->args.size(), 1);
@@ -519,117 +422,42 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-#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}); \
- });
-
-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});
\
- });
-
-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
-
-TVM_SCRIPT_REPR(tirx::IterVarNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(StringImmNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::CastNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::AddNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::SubNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::MulNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::DivNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::ModNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::FloorDivNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::FloorModNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::LShiftNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::RShiftNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::BitwiseAndNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::BitwiseOrNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::BitwiseXorNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::BitwiseNotNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::MinNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::MaxNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::LTNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::LENode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::EQNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::NENode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::GTNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::GENode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::AndNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::OrNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::NotNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::SelectNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::RampNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::BroadcastNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::LetNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(prim::ShuffleNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(te::CommReducerNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::IndexMapNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(te::ReduceNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::LambdaExprNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::IterVarNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(StringImmNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::CastNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::AddNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::SubNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::MulNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::DivNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::ModNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::FloorDivNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::FloorModNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::LShiftNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::RShiftNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseAndNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseOrNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseXorNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseNotNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::MinNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::MaxNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::LTNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::LENode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::EQNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::NENode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::GTNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::GENode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::AndNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::OrNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::NotNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::SelectNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::RampNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::BroadcastNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::LetNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(prim::ShuffleNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(te::CommReducerNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::IndexMapNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(te::ReduceNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::LambdaExprNode, ReprPrintTIR);
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/ir.cc b/src/tirx/script/printer/ir.cc
index f9086d941b..3b5044b463 100644
--- a/src/tirx/script/printer/ir.cc
+++ b/src/tirx/script/printer/ir.cc
@@ -19,6 +19,7 @@
#include <tvm/target/target.h>
#include <tvm/tirx/type.h>
+#include "../../../script/printer/dialect_prefix.h"
#include "./utils.h"
namespace tvm {
@@ -27,33 +28,7 @@ namespace printer {
TVM_FFI_STATIC_INIT_BLOCK() { TIRFrameNode::RegisterReflection(); }
-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"))});
- }
- });
-}
-
-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"))});
- }
- });
-}
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterDialectPrefix("tirx.prefix", "T"); }
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<Range>(
@@ -66,19 +41,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<PrimType>(
- "",
- [](PrimType ty, AccessPath p, IRDocsifier d) -> Doc { return TIR(d,
DType2Str(ty->dtype)); });
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<StringType>(
- "", [](StringType ty, AccessPath p, IRDocsifier d) -> Doc {
- return IR(d, "StringType")->Call({});
- });
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<PointerType>(
"", [](PointerType ty, AccessPath ty_p, IRDocsifier d) -> Doc {
@@ -111,16 +73,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-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);
- });
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<Target>(
"", [](Target target, AccessPath p, IRDocsifier d) -> Doc {
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index af9ec33da0..b0d6e6159c 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -108,9 +108,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
auto scoped_callee = [&](const ffi::String& op_name) -> ExprDoc {
ffi::Optional<ffi::String> ns = scope_ns(op_call->scope->kind);
if (ns.has_value()) {
- return TIRx(d, ns.value())->Attr(op_name);
+ return TIR(d, ns.value())->Attr(op_name);
}
- return TIRx(d, "tile")->Attr(op_name);
+ return TIR(d, "tile")->Attr(op_name);
};
// Trim trailing None args (e.g. optional bias=None, scale=None)
size_t n_args = op_call->args.size();
@@ -144,7 +144,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
d->AsDoc<DictDoc>(op_call->config,
p->Attr("config")), disp);
});
}
-TVM_SCRIPT_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR);
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::Evaluate>(
@@ -840,18 +840,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_SCRIPT_REPR(tirx::BindNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::AttrStmtNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::AssertStmtNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::WhileNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::AllocBufferNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::ReturnNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::BreakNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::ContinueNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::DeclBufferNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::SeqStmtNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::IfThenElseNode, ReprPrintTIR);
-TVM_SCRIPT_REPR(tirx::EvaluateNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::BindNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::AttrStmtNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::AssertStmtNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::WhileNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::AllocBufferNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ReturnNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::BreakNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::ContinueNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::DeclBufferNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::SeqStmtNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::IfThenElseNode, ReprPrintTIR);
+TVM_REGISTER_SCRIPT_AS_REPR(tirx::EvaluateNode, ReprPrintTIR);
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index d6ae5a222e..12bf30b6b9 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -344,10 +344,6 @@ LambdaDoc PrintLambda(const ffi::ObjectRef& pred, const
ffi::Array<tirx::Var>& v
const AccessPath& vs_p, const PrimExpr& p, const
AccessPath& p_p,
const IRDocsifier& d);
-#ifndef TVM_SCRIPT_REPR
-#define TVM_SCRIPT_REPR(ObjectType, Method)
TVM_REGISTER_SCRIPT_AS_REPR(ObjectType, Method)
-#endif
-
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/tests/python/tvmscript/test_tvmscript_printer_ir.py
b/tests/python/tvmscript/test_tvmscript_printer_ir.py
index b7213a4777..0cefc4a319 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_ir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_ir.py
@@ -17,8 +17,10 @@
# pylint: disable=missing-docstring
import pytest
+from tvm_ffi import get_global_func
-from tvm import IRModule
+from tvm import IRModule, tirx
+from tvm.runtime.script_printer import _script
from tvm.s_tir.script.ir_builder import prim_func as build_prim_func
from tvm.script import ir_builder as I
from tvm.script.ir_builder import IRBuilder
@@ -65,6 +67,57 @@ def test_failed_invalid_prefix():
mod.script(ir_prefix="2I")
+def test_config_extension_passthrough():
+ make_config = get_global_func("node.PrinterConfig")
+ cfg = make_config(
+ {
+ "extension.option": 7,
+ "custom_key": "value",
+ "syntax_sugar": False,
+ "render_invisible_path_info": True,
+ "tirx.prefix": "invalid-prefix",
+ "extra_config": {
+ "extension.option": 9,
+ "render_invisible_path_info": False,
+ "tirx.prefix": "Custom",
+ },
+ }
+ )
+ assert cfg.extra_config["extension.option"] == 9
+ assert cfg.extra_config["custom_key"] == "value"
+ assert "syntax_sugar" not in cfg.extra_config
+ assert "extra_config" not in cfg.extra_config
+ assert not cfg.syntax_sugar
+ assert not cfg.render_invisible_path_info
+ assert make_config({}).syntax_sugar
+ assert _script(tirx.Var("Custom", "int32"), cfg) == "Custom_1"
+
+
[email protected]("key", ["tirx.prefix", "relax.prefix",
"s_tir.prefix"])
[email protected]("value", ["2prefix", 17])
[email protected]("nested", [False, True])
+def test_config_validates_dialect_prefixes(key, value, nested):
+ config = {key: value}
+ if nested:
+ config = {"extra_config": config}
+ with pytest.raises((RuntimeError, TypeError)):
+ get_global_func("node.PrinterConfig")(config)
+
+
[email protected](
+ "prefixes",
+ [{}, {"tirx.prefix": "CustomT", "relax.prefix": "CustomR", "s_tir.prefix":
"CustomTs"}],
+)
+def test_config_reserves_dialect_prefixes_before_variable_definition(prefixes):
+ tir_prefix = prefixes.get("tirx.prefix", "T")
+ relax_prefix = prefixes.get("relax.prefix", "R")
+ for name in [tir_prefix, relax_prefix, prefixes.get("s_tir.prefix", "Ts")]:
+ var = tirx.Var(name, "int32")
+ assert var.script(verbose_expr=True, extra_config=prefixes).strip() ==
(
+ f"{name}_1 = {tir_prefix}.int32()\n{name}_1"
+ )
+
+
if __name__ == "__main__":
test_ir_module()
test_failed_invalid_prefix()
