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 68d7a73b96 [REFACTOR][Script] Relocate shared and dialect printer
registrations (#20638)
68d7a73b96 is described below
commit 68d7a73b9685e56356891f67ccb85e53a0c1c9f9
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 9 20:41:06 2026 -0400
[REFACTOR][Script] Relocate shared and dialect printer registrations
(#20638)
Move dialect repr registrations and Relax constant/VDevice printing into
their existing owner printers. Relocate shared tuple and tensor
expression printers, AnyType printing, and simple statement translators
into shared Script, retaining the existing Doc and TensorLoad dispatch
and printed forms.
Keep existing IRModule and type-value handling, mixed TIRx statement
policies, Var/PointerType printing, and the direct S-TIR function path
unchanged. This reduces dependencies in the shared repr and
miscellaneous printers without adding policy keys or context protocols;
full shared-printer independence remains outside this change.
---
CMakeLists.txt | 1 +
src/relax/script/printer/dependent_type.cc | 9 --
src/relax/script/printer/expr.cc | 119 ++++++++++++++++----
src/relax/script/printer/global_info.cc | 15 +++
src/s_tir/script/printer/expr.cc | 5 +-
src/script/printer/ir/misc.cc | 82 --------------
src/script/printer/ir/prim_expr.cc | 114 +++++++++++++++++++
src/script/printer/ir/prim_type.cc | 11 +-
src/script/printer/ir/stmt.cc | 171 +++++++++++++++++++++++++++++
src/script/printer/ir/utils.h | 6 +
src/script/printer/script_printer.cc | 49 +--------
src/te/script/printer/expr.cc | 6 +-
src/tirx/script/printer/buffer.cc | 92 ----------------
src/tirx/script/printer/expr.cc | 19 ++++
src/tirx/script/printer/stmt.cc | 137 -----------------------
src/tirx/script/printer/utils.h | 2 -
16 files changed, 447 insertions(+), 391 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index ee90e23774..d0b3de5dbf 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -318,6 +318,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/script/printer/ir/misc.cc
src/script/printer/ir/prim_type.cc
src/script/printer/ir/prim_expr.cc
+ src/script/printer/ir/stmt.cc
src/script/printer/printer.cc
src/script/printer/doc_printer.cc
src/script/printer/script_printer.cc
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index e4a81d7a7c..e223e4d6b9 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -31,15 +31,6 @@ namespace details {
namespace {
-ffi::Optional<ExprDoc> AnyTypeDocTranslate(DocTranslatorObj*, ffi::AnyView,
const ffi::Object*) {
- return NamespaceDoc("relax")->Attr("Any");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<AnyTypeNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&AnyTypeDocTranslate>());
-}
-
ffi::Optional<ExprDoc> ShapeTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
const auto* ty =
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index 6984e63590..e8be93f65f 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -17,9 +17,19 @@
* under the License.
*/
+#include <tvm/relax/distributed/type.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/type.h>
+#include <tvm/runtime/tensor.h>
+#include <tvm/script/printer/printer.h>
+
+#include <cmath>
+#include <cstring>
+#include <limits>
#include <optional>
#include "../../../script/printer/ir/utils.h"
+#include "../../../script/printer/utils.h"
#include "../../../tirx/script/printer/utils.h"
#include "utils.h"
@@ -30,30 +40,73 @@ 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);
- 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() {
- ffi::reflection::TypeAttrDef<TupleNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&TupleDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> TupleGetItemDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ffi::Optional<ExprDoc> GenericConstDocTranslate(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)}];
+ 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() {
- ffi::reflection::TypeAttrDef<TupleGetItemNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&TupleGetItemDocTranslate>());
+ ffi::reflection::TypeAttrDef<GenericConstNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&GenericConstDocTranslate>());
}
ffi::Optional<ExprDoc> ShapeExprDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
@@ -80,4 +133,30 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} // namespace details
} // namespace printer
} // namespace script
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ using script::printer::details::RegisterScriptRepr;
+ ffi::reflection::GlobalDef().def("script.printer.ReprPrintRelax",
+ [](const ffi::ObjectRef& obj, const
PrinterConfig& config) {
+ return tvm::Script(obj, config);
+ });
+ RegisterScriptRepr<relax::BindingBlockNode>();
+ RegisterScriptRepr<relax::DataflowBlockNode>();
+ RegisterScriptRepr<relax::DataflowVarNode>();
+ RegisterScriptRepr<relax::ExternFuncNode>();
+ RegisterScriptRepr<relax::FuncTypeNode>();
+ RegisterScriptRepr<relax::FunctionNode>();
+ RegisterScriptRepr<relax::IfExprNode>();
+ RegisterScriptRepr<relax::MatchCastNode>();
+ RegisterScriptRepr<relax::PackedFuncTypeNode>();
+ RegisterScriptRepr<relax::SeqExprNode>();
+ RegisterScriptRepr<relax::ShapeExprNode>();
+ RegisterScriptRepr<relax::ShapeTypeNode>();
+ RegisterScriptRepr<relax::TensorTypeNode>();
+ RegisterScriptRepr<relax::VarBindingNode>();
+ RegisterScriptRepr<relax::distributed::DTensorTypeNode>();
+ RegisterScriptRepr<relax::distributed::DeviceMeshNode>();
+ RegisterScriptRepr<relax::distributed::PlacementNode>();
+}
+
} // namespace tvm
diff --git a/src/relax/script/printer/global_info.cc
b/src/relax/script/printer/global_info.cc
index af00725cd3..9c7096213b 100644
--- a/src/relax/script/printer/global_info.cc
+++ b/src/relax/script/printer/global_info.cc
@@ -63,6 +63,21 @@ ffi::Optional<ffi::String>
GlobalInfoSelector(DocTranslatorObj* d, const GlobalI
namespace {
+ffi::Optional<ExprDoc> VDeviceDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object*) {
+ const auto* device =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
relax::VDeviceNode>(input);
+ return NamespaceDoc("relax")->Attr("vdevice")->Call(
+ {d->Translate(device->target).value()}, {"vdevice_id", "memory_scope"},
+ {LiteralDoc::Int(device->vdevice_id, std::nullopt),
+ LiteralDoc::Str(device->memory_scope, std::nullopt)});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<relax::VDeviceNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&VDeviceDocTranslate>());
+}
+
ffi::Optional<ExprDoc> DummyGlobalInfoDocTranslate(DocTranslatorObj*,
ffi::AnyView,
const ffi::Object*) {
return NamespaceDoc("relax")->Attr("dummy_global_info")->Call({});
diff --git a/src/s_tir/script/printer/expr.cc b/src/s_tir/script/printer/expr.cc
index ec9f80f376..736f111a82 100644
--- a/src/s_tir/script/printer/expr.cc
+++ b/src/s_tir/script/printer/expr.cc
@@ -20,12 +20,15 @@
#include <tvm/s_tir/iter_var.h>
#include <tvm/s_tir/stmt.h>
-#include "../../../tirx/script/printer/utils.h"
+#include "../../../script/printer/ir/utils.h"
+#include "../../../script/printer/utils.h"
namespace tvm {
namespace script {
namespace printer {
namespace details {
+
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterScriptRepr<s_tir::IterVarNode>(); }
namespace {
ffi::Optional<ExprDoc> IterVarDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index e89c9341ea..c169969848 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -19,15 +19,10 @@
#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 <algorithm>
#include <cmath>
-#include <cstring>
-#include <limits>
#include <optional>
#include <utility>
#include <vector>
@@ -54,75 +49,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&DataTypeImmDocTranslate>());
}
-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() {
- ffi::reflection::TypeAttrDef<GenericConstNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&GenericConstDocTranslate>());
-}
-
} // namespace
ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value) {
@@ -170,14 +96,6 @@ ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value) {
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(
diff --git a/src/script/printer/ir/prim_expr.cc
b/src/script/printer/ir/prim_expr.cc
index 9e0ff08992..5c9a5cd320 100644
--- a/src/script/printer/ir/prim_expr.cc
+++ b/src/script/printer/ir/prim_expr.cc
@@ -49,8 +49,122 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&LambdaExprDocTranslate>());
}
+ffi::Array<Doc> TensorIndices(DocTranslatorObj* d, const ffi::Array<PrimExpr>&
indices,
+ bool store) {
+ ffi::Array<Doc> docs;
+ for (const PrimExpr& index : indices) {
+ if (const auto* ramp = index.as<prim::RampNode>(); store && ramp) {
+ if (const auto* stride = ramp->stride.as<IntImmNode>()) {
+ ffi::Optional<ExprDoc> step = std::nullopt;
+ if (stride->value != 1) step = d->Translate(ramp->stride).value();
+ SliceDoc slice(d->Translate(ramp->base).value(),
+ d->Translate(ramp->base + ramp->lanes *
ramp->stride).value(), step);
+ d->RecordOrigin(slice, index);
+ docs.push_back(slice);
+ continue;
+ }
+ }
+ docs.push_back(d->Translate(index).value());
+ }
+ return docs;
+}
+
+ExprDoc TensorRegionValue(DocTranslatorObj* d, const TensorRegionNode* region,
+ bool require_region) {
+ ffi::Array<Doc> slices;
+ for (const Range& range : region->region) {
+ ExprDoc start = d->Translate(range->min).value();
+ if (const auto* extent = range->extent.as<IntImmNode>();
+ !require_region && extent && extent->value == 1) {
+ slices.push_back(start);
+ } else {
+ ExprDoc end = [&]() -> ExprDoc {
+ auto lower = range->min.as<IntImmNode>();
+ auto length = range->extent.as<IntImmNode>();
+ if (lower && length) {
+ auto last = (lower->value + length->value).as<int64_t>();
+ if (last) {
+ return LiteralDoc::Int(*last, std::nullopt);
+ }
+ }
+ return OperationDoc(OperationDocNode::Kind::kAdd,
+ {start, d->Translate(range->extent).value()});
+ }();
+ slices.push_back(SliceDoc(start, end, std::nullopt));
+ }
+ }
+ return d->Translate(region->source).value()[slices];
+}
+
namespace {
+ffi::Optional<ExprDoc> TensorRegionDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* region =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorRegionNode>(input);
+ return TensorRegionValue(d, region, false);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<TensorRegionNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TensorRegionDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> TensorLoadDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* load =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorLoadNode>(input);
+ static ffi::reflection::TypeAttrColumn column(kTensorLoadDocTranslate);
+ ffi::AnyView hook = column[load->source->ty->type_index()];
+ ffi::Any value = ffi::GetRef<TensorLoad>(load);
+ if (hook.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) {
+ return
ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>(
+
reinterpret_cast<decltype(DocTranslatorVTable::translate)>(hook.cast<void*>())(
+ d, value, destination))
+ .value();
+ }
+ if (hook.type_index() == ffi::TypeIndex::kTVMFFIFunction) {
+ ffi::Any binder = destination ?
ffi::Any(ffi::GetRef<ffi::ObjectRef>(destination)) : nullptr;
+ return hook.cast<ffi::Function>()
+ .CallExpected<ffi::Optional<ExprDoc>>(d, value, binder)
+ .value();
+ }
+ TVM_FFI_CHECK(hook.type_index() == ffi::TypeIndex::kTVMFFINone, TypeError)
+ << "TensorLoad type hook must be a native pointer or ffi.Function";
+ return IndexDoc(d->Translate(load->source).value(), TensorIndices(d,
load->indices));
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<TensorLoadNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TensorLoadDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> TupleDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object*) {
+ const auto* tuple =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TupleNode>(input);
+ 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() {
+ ffi::reflection::TypeAttrDef<TupleNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&TupleDocTranslate>());
+}
+
+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() {
+ ffi::reflection::TypeAttrDef<TupleGetItemNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&TupleGetItemDocTranslate>());
+}
+
template <typename T, OperationDocNode::Kind kind, PrimExpr
(*operation)(PrimExpr, PrimExpr, Span)>
ffi::Optional<ExprDoc> BinaryOpDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
diff --git a/src/script/printer/ir/prim_type.cc
b/src/script/printer/ir/prim_type.cc
index dc74aadde6..4a50a88f05 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -17,8 +17,6 @@
* under the License.
*/
-#include <tvm/tirx/type.h>
-
#include <algorithm>
#include <functional>
#include <optional>
@@ -32,6 +30,15 @@ namespace details {
namespace {
+ffi::Optional<ExprDoc> AnyTypeDocTranslate(DocTranslatorObj*, ffi::AnyView,
const ffi::Object*) {
+ return NamespaceDoc("relax")->Attr("Any");
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<AnyTypeNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&AnyTypeDocTranslate>());
+}
+
ffi::Optional<ExprDoc> MissingTypeDocTranslate(DocTranslatorObj*, ffi::AnyView,
const ffi::Object*) {
return NamespaceDoc("ir")->Attr("MissingType")->Call({});
diff --git a/src/script/printer/ir/stmt.cc b/src/script/printer/ir/stmt.cc
new file mode 100644
index 0000000000..5639fde416
--- /dev/null
+++ b/src/script/printer/ir/stmt.cc
@@ -0,0 +1,171 @@
+/*
+ * 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/stmt.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace script {
+namespace printer {
+namespace details {
+
+ffi::Array<StmtDoc> Body(const Stmt& stmt, DocTranslatorObj* d) {
+ return ToStmtDocArray(d->WithDocScope([&]() { d->Translate(stmt); }));
+}
+
+namespace {
+
+ffi::Optional<ExprDoc> BindDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
BindNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ d->VarGetOrAllocId(stmt->var, false);
+ bool existing = !d->GetImplicitDefs().count(stmt->var);
+ IdDoc lhs = VarDoc(d, stmt->var);
+ auto rhs = d->Translate(stmt->value, stmt->var);
+ // None means the child completed emission. A returned expression is still
+ // an RHS, so this binding owns the one remaining assignment.
+ if (!rhs.has_value()) return std::nullopt;
+ ffi::Optional<ExprDoc> annotation = std::nullopt;
+ if (!existing) {
+ ExprDoc type = d->Translate(stmt->var->ty).value();
+ if (auto primitive = stmt->var->ty.as<PrimType>();
+ primitive && primitive.value().IsScalableVector()) {
+ type = TypeValue(d, stmt->var->ty);
+ }
+ annotation = NamespaceDoc("tirx")->Attr("let")[{type}];
+ }
+ d->Emit(AssignDoc(lhs, rhs.value(), annotation),
ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<BindNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&BindDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> ReturnDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ReturnNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ d->Emit(ReturnDoc(d->Translate(stmt->value).value()),
ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<ReturnNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&ReturnDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> AssertStmtDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
AssertStmtNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ ffi::Array<ExprDoc> parts;
+ for (const StringImm& part : stmt->message_parts) {
+ parts.push_back(LiteralDoc::Str(part->value, std::nullopt));
+ }
+ d->Emit(
+ AssertDoc(d->Translate(stmt->condition).value(),
+ TupleDoc({LiteralDoc::Str(stmt->error_kind->value,
std::nullopt), ListDoc(parts)})),
+ ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<AssertStmtNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&AssertStmtDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> WhileDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
WhileNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ d->Emit(WhileDoc(d->Translate(stmt->condition).value(), Body(stmt->body, d)),
+ ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<WhileNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&WhileDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> BreakDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
BreakNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ d->Emit(BreakDoc(), ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<BreakNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&BreakDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> ContinueDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ContinueNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ d->Emit(ContinueDoc(), ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<ContinueNode>().attr(
+ kDocTranslate, FDocTranslate::FromNative<&ContinueDocTranslate>());
+}
+
+ffi::Optional<ExprDoc> IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
+ const ffi::Object* destination) {
+ const auto* stmt =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IfNode>(input);
+ TVM_FFI_CHECK(destination == nullptr, TypeError)
+ << "printer statement-only node cannot fulfill a destination";
+ ExprDoc condition = d->Translate(stmt->condition).value();
+ ffi::Array<StmtDoc> then_body = Body(stmt->then_case, d);
+ ffi::Optional<ffi::Array<StmtDoc>> else_body;
+ if (stmt->else_case.has_value()) else_body = Body(stmt->else_case.value(),
d);
+ d->Emit(IfDoc(condition, then_body, else_body),
ffi::GetRef<ffi::ObjectRef>(stmt));
+ return std::nullopt;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<IfNode>().attr(kDocTranslate,
+
FDocTranslate::FromNative<&IfDocTranslate>());
+}
+
+} // namespace
+} // 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 d7ed373928..24a72e9e9f 100644
--- a/src/script/printer/ir/utils.h
+++ b/src/script/printer/ir/utils.h
@@ -21,6 +21,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
+#include <tvm/ir/stmt.h>
#include <tvm/ir/type.h>
#include <tvm/script/printer/doc_translator.h>
@@ -33,6 +34,11 @@ namespace printer {
namespace details {
+ffi::Array<StmtDoc> Body(const Stmt& stmt, DocTranslatorObj* d);
+ffi::Array<Doc> TensorIndices(DocTranslatorObj* d, const ffi::Array<PrimExpr>&
indices,
+ bool store = false);
+ExprDoc TensorRegionValue(DocTranslatorObj* d, const TensorRegionNode* region,
bool require_region);
+
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);
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index ad148bcc0b..51d311d8cc 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -24,16 +24,8 @@
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/vector_expr.h>
-#include <tvm/relax/distributed/type.h>
-#include <tvm/relax/expr.h>
-#include <tvm/relax/type.h>
+#include <tvm/ir/stmt.h>
#include <tvm/script/printer/printer.h>
-#include <tvm/te/operation.h>
-#include <tvm/tirx/exec_scope.h>
-#include <tvm/tirx/function.h>
-#include <tvm/tirx/index_map.h>
-#include <tvm/tirx/layout.h>
-#include <tvm/tirx/stmt.h>
#include <utility>
@@ -46,11 +38,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
using script::printer::details::RegisterScriptRepr;
refl::GlobalDef()
.def("node.TVMScriptPrinterScript", tvm::Script)
- .def("script.printer.Script", tvm::Script)
- .def("script.printer.ReprPrintRelax",
- [](const ffi::ObjectRef& obj, const PrinterConfig& config) {
- return tvm::Script(obj, config);
- });
+ .def("script.printer.Script", tvm::Script);
RegisterScriptRepr<DataTypeImmNode>();
RegisterScriptRepr<GenericConstNode>();
@@ -69,6 +57,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<StringImmNode>();
RegisterScriptRepr<TensorLoadNode>();
RegisterScriptRepr<TensorRegionTypeNode>();
+ RegisterScriptRepr<TupleNode>();
+ RegisterScriptRepr<TupleGetItemNode>();
RegisterScriptRepr<TupleTypeNode>();
RegisterScriptRepr<VarNode>();
RegisterScriptRepr<prim::AddNode>();
@@ -101,49 +91,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<prim::SelectNode>();
RegisterScriptRepr<prim::ShuffleNode>();
RegisterScriptRepr<prim::SubNode>();
- RegisterScriptRepr<relax::BindingBlockNode>();
- RegisterScriptRepr<relax::DataflowBlockNode>();
- RegisterScriptRepr<relax::DataflowVarNode>();
- RegisterScriptRepr<relax::ExternFuncNode>();
- RegisterScriptRepr<relax::FuncTypeNode>();
- RegisterScriptRepr<relax::FunctionNode>();
- RegisterScriptRepr<relax::IfExprNode>();
- RegisterScriptRepr<relax::MatchCastNode>();
- RegisterScriptRepr<relax::PackedFuncTypeNode>();
- RegisterScriptRepr<relax::SeqExprNode>();
- RegisterScriptRepr<relax::ShapeExprNode>();
- RegisterScriptRepr<relax::ShapeTypeNode>();
- RegisterScriptRepr<relax::TensorTypeNode>();
- RegisterScriptRepr<relax::TupleGetItemNode>();
- RegisterScriptRepr<relax::TupleNode>();
- RegisterScriptRepr<relax::VarBindingNode>();
- RegisterScriptRepr<relax::distributed::DTensorTypeNode>();
- RegisterScriptRepr<relax::distributed::DeviceMeshNode>();
- RegisterScriptRepr<relax::distributed::PlacementNode>();
- RegisterScriptRepr<te::CommReducerNode>();
- RegisterScriptRepr<te::ReduceNode>();
RegisterScriptRepr<AssertStmtNode>();
RegisterScriptRepr<RegionStmtNode>();
RegisterScriptRepr<BindNode>();
RegisterScriptRepr<BreakNode>();
RegisterScriptRepr<TensorStoreNode>();
- RegisterScriptRepr<tirx::TensorTypeNode>();
- RegisterScriptRepr<tirx::ComposeLayoutNode>();
RegisterScriptRepr<ContinueNode>();
RegisterScriptRepr<EvaluateNode>();
- RegisterScriptRepr<tirx::ExecScopeNode>();
RegisterScriptRepr<ForNode>();
RegisterScriptRepr<IfNode>();
- RegisterScriptRepr<tirx::IndexMapNode>();
- RegisterScriptRepr<tirx::IterNode>();
- RegisterScriptRepr<s_tir::IterVarNode>();
RegisterScriptRepr<LambdaExprNode>();
- RegisterScriptRepr<tirx::FunctionNode>();
RegisterScriptRepr<ReturnNode>();
- RegisterScriptRepr<tirx::ScopeIdDefNode>();
- RegisterScriptRepr<tirx::ScopeIdDefStmtNode>();
RegisterScriptRepr<SeqStmtNode>();
- RegisterScriptRepr<tirx::TileLayoutNode>();
RegisterScriptRepr<WhileNode>();
RegisterScriptRepr<TensorRegionNode>();
}
diff --git a/src/te/script/printer/expr.cc b/src/te/script/printer/expr.cc
index ca6c171804..2c27fd2374 100644
--- a/src/te/script/printer/expr.cc
+++ b/src/te/script/printer/expr.cc
@@ -20,9 +20,13 @@
#include <tvm/te/operation.h>
#include "../../../script/printer/ir/utils.h"
-#include "../../../tirx/script/printer/utils.h"
+#include "../../../script/printer/utils.h"
namespace tvm::script::printer::details {
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterScriptRepr<te::CommReducerNode>();
+ RegisterScriptRepr<te::ReduceNode>();
+}
namespace {
ffi::Optional<ExprDoc> CommReducerDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object*) {
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 62e223409f..019ce73ec1 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -223,69 +223,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace {
-ffi::Array<Doc> TensorIndices(DocTranslatorObj* d, const ffi::Array<PrimExpr>&
indices,
- bool store = false) {
- ffi::Array<Doc> docs;
- for (const PrimExpr& index : indices) {
- if (const auto* ramp = index.as<prim::RampNode>(); store && ramp) {
- if (const auto* stride = ramp->stride.as<IntImmNode>()) {
- ffi::Optional<ExprDoc> step = std::nullopt;
- if (stride->value != 1) step = d->Translate(ramp->stride).value();
- SliceDoc slice(d->Translate(ramp->base).value(),
- d->Translate(ramp->base + ramp->lanes *
ramp->stride).value(), step);
- d->RecordOrigin(slice, index);
- docs.push_back(slice);
- continue;
- }
- }
- docs.push_back(d->Translate(index).value());
- }
- return docs;
-}
-
-} // namespace
-
-ExprDoc TensorRegionValue(DocTranslatorObj* d, const TensorRegionNode* region,
- bool require_region) {
- ffi::Array<Doc> slices;
- for (const Range& range : region->region) {
- ExprDoc start = d->Translate(range->min).value();
- if (const auto* extent = range->extent.as<IntImmNode>();
- !require_region && extent && extent->value == 1) {
- slices.push_back(start);
- } else {
- ExprDoc end = [&]() -> ExprDoc {
- auto lower = range->min.as<IntImmNode>();
- auto length = range->extent.as<IntImmNode>();
- if (lower && length) {
- auto last = (lower->value + length->value).as<int64_t>();
- if (last) {
- return LiteralDoc::Int(*last, std::nullopt);
- }
- }
- return OperationDoc(OperationDocNode::Kind::kAdd,
- {start, d->Translate(range->extent).value()});
- }();
- slices.push_back(SliceDoc(start, end, std::nullopt));
- }
- }
- return d->Translate(region->source).value()[slices];
-}
-
-namespace {
-
-ffi::Optional<ExprDoc> TensorRegionDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object*) {
- const auto* region =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorRegionNode>(input);
- return TensorRegionValue(d, region, false);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<TensorRegionNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&TensorRegionDocTranslate>());
-}
-
ffi::Optional<ExprDoc> TensorStoreDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object* destination)
{
const auto* store =
@@ -323,35 +260,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kTensorLoadDocTranslate,
FDocTranslate::FromNative<&TIRxTensorLoadDocTranslate>());
}
-ffi::Optional<ExprDoc> TensorLoadDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object* destination) {
- const auto* load =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorLoadNode>(input);
- static ffi::reflection::TypeAttrColumn column(kTensorLoadDocTranslate);
- ffi::AnyView hook = column[load->source->ty->type_index()];
- ffi::Any value = ffi::GetRef<TensorLoad>(load);
- if (hook.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) {
- return
ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>(
-
reinterpret_cast<decltype(DocTranslatorVTable::translate)>(hook.cast<void*>())(
- d, value, destination))
- .value();
- }
- if (hook.type_index() == ffi::TypeIndex::kTVMFFIFunction) {
- ffi::Any binder = destination ?
ffi::Any(ffi::GetRef<ffi::ObjectRef>(destination)) : nullptr;
- return hook.cast<ffi::Function>()
- .CallExpected<ffi::Optional<ExprDoc>>(d, value, binder)
- .value();
- }
- TVM_FFI_CHECK(hook.type_index() == ffi::TypeIndex::kTVMFFINone, TypeError)
- << "TensorLoad type hook must be a native pointer or ffi.Function";
- return IndexDoc(d->Translate(load->source).value(), TensorIndices(d,
load->indices));
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<TensorLoadNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&TensorLoadDocTranslate>());
-}
-
ffi::Optional<ExprDoc> IterDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
const auto* iter =
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 25499493d4..ac08b8283e 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -21,9 +21,13 @@
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/exec_scope.h>
+#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
+#include <tvm/tirx/layout.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op_attr_types.h>
+#include <tvm/tirx/stmt.h>
#include <algorithm>
#include <limits>
@@ -33,6 +37,7 @@
#include <vector>
#include "../../../script/printer/ir/utils.h"
+#include "../../../script/printer/utils.h"
#include "utils.h"
namespace tvm {
@@ -283,4 +288,18 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} // namespace details
} // namespace printer
} // namespace script
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ using script::printer::details::RegisterScriptRepr;
+ RegisterScriptRepr<tirx::TensorTypeNode>();
+ RegisterScriptRepr<tirx::ComposeLayoutNode>();
+ RegisterScriptRepr<tirx::ExecScopeNode>();
+ RegisterScriptRepr<tirx::IndexMapNode>();
+ RegisterScriptRepr<tirx::IterNode>();
+ RegisterScriptRepr<tirx::FunctionNode>();
+ RegisterScriptRepr<tirx::ScopeIdDefNode>();
+ RegisterScriptRepr<tirx::ScopeIdDefStmtNode>();
+ RegisterScriptRepr<tirx::TileLayoutNode>();
+}
+
} // namespace tvm
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 4e2ce7bdfc..842a4ab431 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -39,10 +39,6 @@ namespace script {
namespace printer {
namespace details {
-ffi::Array<StmtDoc> Body(const Stmt& stmt, DocTranslatorObj* d) {
- return ToStmtDocArray(d->WithDocScope([&]() { d->Translate(stmt); }));
-}
-
namespace {
ExprDoc TensorOperandDoc(DocTranslatorObj* d, const Expr& value) {
@@ -112,139 +108,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&EvaluateDocTranslate>());
}
-ffi::Optional<ExprDoc> ReturnDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ReturnNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->Emit(ReturnDoc(d->Translate(stmt->value).value()),
ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<ReturnNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&ReturnDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> BindDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
BindNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->VarGetOrAllocId(stmt->var, false);
- bool existing = !d->GetImplicitDefs().count(stmt->var);
- IdDoc lhs = VarDoc(d, stmt->var);
- auto rhs = d->Translate(stmt->value, stmt->var);
- // None means the child completed emission. A returned expression is still
- // an RHS, so this binding owns the one remaining assignment.
- if (!rhs.has_value()) return std::nullopt;
- ffi::Optional<ExprDoc> annotation = std::nullopt;
- if (!existing) {
- ExprDoc type = d->Translate(stmt->var->ty).value();
- if (auto primitive = stmt->var->ty.as<PrimType>();
- primitive && primitive.value().IsScalableVector()) {
- type = TypeValue(d, stmt->var->ty);
- }
- annotation = NamespaceDoc("tirx")->Attr("let")[{type}];
- }
- d->Emit(AssignDoc(lhs, rhs.value(), annotation),
ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<BindNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&BindDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> AssertStmtDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
AssertStmtNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- ffi::Array<ExprDoc> parts;
- for (const StringImm& part : stmt->message_parts) {
- parts.push_back(LiteralDoc::Str(part->value, std::nullopt));
- }
- d->Emit(
- AssertDoc(d->Translate(stmt->condition).value(),
- TupleDoc({LiteralDoc::Str(stmt->error_kind->value,
std::nullopt), ListDoc(parts)})),
- ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<AssertStmtNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&AssertStmtDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> WhileDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
WhileNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->Emit(WhileDoc(d->Translate(stmt->condition).value(), Body(stmt->body, d)),
- ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<WhileNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&WhileDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> BreakDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
BreakNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->Emit(BreakDoc(), ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<BreakNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&BreakDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> ContinueDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ContinueNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- d->Emit(ContinueDoc(), ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<ContinueNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&ContinueDocTranslate>());
-}
-
-ffi::Optional<ExprDoc> IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IfNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- ExprDoc condition = d->Translate(stmt->condition).value();
- ffi::Array<StmtDoc> then_body = Body(stmt->then_case, d);
- ffi::Optional<ffi::Array<StmtDoc>> else_body;
- if (stmt->else_case.has_value()) else_body = Body(stmt->else_case.value(),
d);
- d->Emit(IfDoc(condition, then_body, else_body),
ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<IfNode>().attr(kDocTranslate,
-
FDocTranslate::FromNative<&IfDocTranslate>());
-}
-
ffi::Optional<ExprDoc> SeqStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object* destination) {
const auto* stmt =
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index c111719afa..97c6059923 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -36,8 +36,6 @@ void PrintFunction(DocTranslatorObj* d, const
tirx::FunctionNode* func, ExprDoc
const ffi::String& dialect_attr);
bool CanTranslateExplicitResultCall(const CallNode* call);
bool IsScalarBuffer(DocTranslatorObj* d, const Expr& source);
-ffi::Array<StmtDoc> Body(const Stmt& stmt, DocTranslatorObj* d);
-ExprDoc TensorRegionValue(DocTranslatorObj* d, const TensorRegionNode* region,
bool require_region);
ffi::Optional<ExprDoc> VarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object* destination);
