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 9059d9d251 [REFACTOR][Script] Simplify region and call printing
(#20640)
9059d9d251 is described below
commit 9059d9d25149c9ca7a015717b5023d1e2e1185fd
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 10 11:39:30 2026 -0400
[REFACTOR][Script] Simplify region and call printing (#20640)
---
include/tvm/ir/expr.h | 8 ++
include/tvm/script/printer/doc.h | 52 ---------
python/tvm/backend/trn/tensor_instructions.py | 8 +-
python/tvm/ir/op.py | 17 +++
python/tvm/relax/op/base.py | 55 +++++++--
.../tvm/tirx/script/ir_builder/parser_protocol.py | 33 ++++--
python/tvm/tirx/tensor_instruction.py | 14 +--
src/ir/side_effect.cc | 16 +--
src/relax/script/printer/call.cc | 44 +++++++-
src/script/printer/doc.cc | 20 ----
src/script/printer/doc_printer.cc | 71 ------------
src/script/printer/ir/call.cc | 12 +-
src/script/printer/ir/prim_expr.cc | 97 +++++++++++-----
src/script/printer/ir/utils.cc | 45 +-------
src/script/printer/ir/utils.h | 7 +-
src/tirx/ir/buffer.cc | 10 ++
src/tirx/op/region.cc | 10 ++
src/tirx/script/ir_builder/utils.h | 6 +-
src/tirx/script/printer/expr.cc | 15 ++-
src/tirx/script/printer/stmt.cc | 123 +++++++--------------
src/tirx/script/printer/utils.h | 1 -
.../relax/script/test_relax_script_printer.py | 30 +++--
.../relax/test_transform_error_enrichment.py | 7 +-
.../tirx/script/test_tirx_script_basic_usage.py | 8 ++
24 files changed, 320 insertions(+), 389 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 0c0891725f..7eb3755c96 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -1098,6 +1098,14 @@ TVM_DLL CallEffectKind SideEffect(const Expr& expr);
namespace ffi {
+template <>
+inline constexpr bool use_default_type_traits_v<TensorRegion> = false;
+
+template <>
+struct TypeTraits<TensorRegion> : public
ObjectRefWithFallbackTraitsBase<TensorRegion, TensorLoad> {
+ TVM_DLL static TensorRegion ConvertFallbackValue(TensorLoad value);
+};
+
template <>
inline constexpr bool use_default_type_traits_v<PrimVar> = false;
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index 4c08bfded7..a7a8ebf186 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -1352,58 +1352,6 @@ class DocStringDoc : public StmtDoc {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DocStringDoc, StmtDoc,
DocStringDocNode);
};
-/*!
- * \brief Doc that represents call to an TIRX operator
- *
- * \sa OpCallDoc
- */
-class OpCallDocNode : public StmtDocNode {
- public:
- /*! \brief The callee of this function call */
- ExprDoc callee{ffi::UnsafeInit()};
- /*! \brief The positional arguments */
- ffi::Array<Doc> args;
- /*! \brief The workspace of this op call */
- ffi::Optional<DictDoc> workspace{std::nullopt};
- /*! \brief The config of this op call */
- ffi::Optional<DictDoc> config{std::nullopt};
- /*! \brief The optional dispatch variant of this op call */
- ffi::Optional<ExprDoc> dispatch{std::nullopt};
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<OpCallDocNode>()
- .def_ro("callee", &OpCallDocNode::callee)
- .def_ro("args", &OpCallDocNode::args)
- .def_ro("workspace", &OpCallDocNode::workspace)
- .def_ro("config", &OpCallDocNode::config)
- .def_ro("dispatch", &OpCallDocNode::dispatch);
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.OpCallDoc", OpCallDocNode,
StmtDocNode);
-};
-
-/*!
- * \brief Reference type of OpCallDocNode.
- *
- * \sa OpCallDocNode
- */
-class OpCallDoc : public StmtDoc {
- public:
- /*!
- * \brief Constructor of OpCallDoc
- * \param callee The callee of this function call.
- * \param args The positional arguments.
- * \param workspace The workspace of this op call.
- * \param config The config of this op call.
- * \param dispatch The optional dispatch variant name of this op call.
- */
- TVM_DLL explicit OpCallDoc(ExprDoc callee, ffi::Array<Doc> args,
ffi::Optional<DictDoc> workspace,
- ffi::Optional<DictDoc> config,
- ffi::Optional<ExprDoc> dispatch = std::nullopt);
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpCallDoc, StmtDoc,
OpCallDocNode);
-};
-
/*!
* \brief A reference to a canonical script namespace.
*
diff --git a/python/tvm/backend/trn/tensor_instructions.py
b/python/tvm/backend/trn/tensor_instructions.py
index cfc66f5467..1d02167bae 100644
--- a/python/tvm/backend/trn/tensor_instructions.py
+++ b/python/tvm/backend/trn/tensor_instructions.py
@@ -16,7 +16,7 @@
# under the License.
"""Trainium tensor instructions with explicit NKI instruction identity."""
-from tvm.ir import LambdaExpr, PrimType, Tuple, const
+from tvm.ir import Expr, LambdaExpr, PrimType, Tuple, const
from tvm.tirx.tensor_instruction import Instruction, Operand, namespace
@@ -296,10 +296,10 @@ def make_namespace():
original = ns.affine_select
def affine_select(dst, true_value, false_value, pred, **kwargs):
- from tvm.tirx.tensor_instruction import _region
+ from tvm.tirx.tensor_instruction import _operand
- dst = _region(dst)
- if not isinstance(pred, LambdaExpr):
+ dst = _operand(dst, "region")
+ if not isinstance(pred, Expr):
pred = LambdaExpr([PrimType("int32")] * len(dst.region), pred)
return original(dst, true_value, false_value, pred, **kwargs)
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index e6a42a153d..4269b0e262 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -45,6 +45,8 @@ def _make_op_api(op, module_name):
):
# Bind named operands using the registered signature, without inventing
# defaults or interpreting arbitrary Python wrapper signatures.
+ is_region = op.get_attr("FRegionGetBodyParams") is not None
+ body_params = kwargs.pop("body_params", None) if is_region else None
operands = list(args)
for index, info in enumerate(op.args_info):
if index < len(args):
@@ -60,6 +62,19 @@ def _make_op_api(op, module_name):
if attrs is not None:
raise TypeError(f"{op.name}: cannot mix attrs with attribute
keywords")
attrs = make_node(op.attrs_type_key, **kwargs)
+ if is_region:
+ from tvm.script.ir_builder.stmt import region
+
+ if (
+ ty is not None
+ or ty_args is not None
+ or not isinstance(loc, Location)
+ or not loc.same_as(UNKNOWN_LOC)
+ ):
+ raise TypeError(
+ f"{op.name}: region builders do not accept ty, ty_args or
an explicit loc"
+ )
+ return region(op, operands, body_params=body_params, attrs=attrs)
return Call(op, operands, attrs=attrs, ty_args=ty_args, loc=loc, ty=ty)
call.__name__ = op.name.rsplit(".", 1)[-1]
@@ -80,6 +95,8 @@ def _init_op_api(namespace, target_module_name=None):
Generated functions accept registered positional/named operands plus
``attrs``, ``ty_args``, ``loc`` and ``ty``.
+ Region operations construct a RegionFrame and accept ``body_params``
instead
+ of Call result/type arguments.
Omitting the result invokes an available Op inference hook; without one,
Call retains a missing type.
Explicit results and inference errors are preserved; Call.validate checks
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index aa83837075..c3991930bb 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -27,7 +27,7 @@ from tvm.ir import Attrs, Call, Op
from tvm.ir.attrs import make_node as _make_attrs
from tvm.ir.location import UNKNOWN_LOC, Location
from tvm.ir.op import _make_op_api
-from tvm.runtime import Object, ObjectConvertible
+from tvm.runtime import DataTypeCode, Object, ObjectConvertible
from ..expr import Expr, ExternFunc, GlobalVar, Var
from ..type import Type
@@ -86,7 +86,7 @@ def _wrap_inline_arg_tuple(args) -> Expr:
"""
if isinstance(args, tuple | list):
- return tvm.relax.Tuple([convert_to_expr(a) for a in args])
+ return tvm.ir.Tuple([convert_to_expr(arg) for arg in args])
elif (
isinstance(args, Expr)
and not isinstance(args, tvm.relax.Tuple)
@@ -97,6 +97,26 @@ def _wrap_inline_arg_tuple(args) -> Expr:
return args
+def _wrap_tir_arg_tuple(func, args) -> Expr:
+ """Bind native scalar literals without changing explicit IR arguments."""
+ if not isinstance(args, tuple | list):
+ return _wrap_inline_arg_tuple(args)
+ signature = getattr(func, "ty", None)
+ arg_types = signature.arg_types if isinstance(signature, tvm.ir.FuncType)
else ()
+ values = []
+ for i, value in enumerate(args):
+ ty = arg_types[i] if i < len(arg_types) else None
+ if isinstance(ty, tvm.ir.PrimType) and ty.dtype.lanes == 1:
+ if type(value) in (int, bool) and ty.matches_code(
+ DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL
+ ):
+ value = tvm.tirx.IntImm(ty, value)
+ elif type(value) is float and ty.matches_code(DataTypeCode.FLOAT,
DataTypeCode.BFLOAT):
+ value = tvm.tirx.FloatImm(ty, value)
+ values.append(convert_to_expr(value))
+ return tvm.ir.Tuple(values)
+
+
_call_tir = _make_op_api(Op.get("relax.call_tir"), __name__)
@@ -111,7 +131,13 @@ def call_tir(
and locs are forwarded unchanged.
"""
return _call_tir(
- func, _wrap_inline_arg_tuple(args), ty_args=ty_args, attrs=attrs,
ty=ty, loc=loc, **kwargs
+ func,
+ _wrap_tir_arg_tuple(func, args),
+ ty_args=ty_args,
+ attrs=attrs,
+ ty=ty,
+ loc=loc,
+ **kwargs,
)
@@ -187,7 +213,7 @@ def call_tir_packed(gvar: GlobalVar, args: Expr, *,
ty=None, loc: Location = UNK
relax.call_tir_packed(copy, (source, destination))
"""
- args = _wrap_inline_arg_tuple(args)
+ args = _wrap_tir_arg_tuple(gvar, args)
return Call("relax.call_tir_packed", [gvar, args], ty=ty, loc=loc)
@@ -212,7 +238,13 @@ def call_tir_with_grad(
if attrs is None and kwargs.get("te_grad_kwargs") is None:
kwargs["te_grad_kwargs"] = {}
return _call_tir_with_grad(
- func, _wrap_inline_arg_tuple(args), ty_args=ty_args, attrs=attrs,
ty=ty, loc=loc, **kwargs
+ func,
+ _wrap_tir_arg_tuple(func, args),
+ ty_args=ty_args,
+ attrs=attrs,
+ ty=ty,
+ loc=loc,
+ **kwargs,
)
@@ -242,7 +274,13 @@ def call_tir_inplace(
if isinstance(kwargs.get("inplace_indices"), int):
kwargs["inplace_indices"] = [kwargs["inplace_indices"]]
return _call_tir_inplace(
- func, _wrap_inline_arg_tuple(args), ty_args=ty_args, attrs=attrs,
ty=ty, loc=loc, **kwargs
+ func,
+ _wrap_tir_arg_tuple(func, args),
+ ty_args=ty_args,
+ attrs=attrs,
+ ty=ty,
+ loc=loc,
+ **kwargs,
)
@@ -263,9 +301,8 @@ def call_dps_packed(
"""
if isinstance(func, str):
func = ExternFunc(func)
- return _call_dps_packed(
- func, _wrap_inline_arg_tuple(args), ty_args=ty_args, attrs=attrs,
ty=ty, loc=loc, **kwargs
- )
+ args = _wrap_inline_arg_tuple(args)
+ return _call_dps_packed(func, args, ty_args=ty_args, attrs=attrs, ty=ty,
loc=loc, **kwargs)
def call_py_func(
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 3c086a226c..b259e3e7a0 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -34,7 +34,6 @@ import tvm
# isort: on
from tvm import ir as _ir
from tvm import tirx as _tir
-from tvm.ir import StringImm as _StringImm
from tvm.ir import is_prim_expr
from tvm.script.ir_builder import base as _base
from tvm.script.ir_builder.base import AlreadyEmitted
@@ -755,18 +754,23 @@ def thread_binding(
)
-def device_entry(*, launch=None, kernel_attrs=None) -> frame.RegionFrame:
+def device_entry(
+ *launch_values, launch=None, kernel_attrs=None, attrs=None,
body_params=None
+) -> frame.RegionFrame:
"""Enter a device kernel with an independent CUDA launch configuration.
CUDA entries use ``LaunchConfig(grid=..., block=...)`` and optional
``KernelAttributes``. Configuration values are ordinary region operands, so
host expressions remain visible to substitution and free-variable analysis.
+ Canonical operands and attributes are also accepted for IR reconstruction.
Other backends may use the argument-free device entry.
"""
if launch is None:
if kernel_attrs is not None:
raise ValueError("device_entry kernel_attrs require a launch
configuration")
- return region("tirx.device_entry", [])
+ return region("tirx.device_entry", launch_values, attrs=attrs,
body_params=body_params)
+ if launch_values or attrs is not None:
+ raise ValueError("device_entry launch cannot be combined with
canonical operands or attrs")
from tvm.backend.cuda.launch._impl import pack_kernel_attrs, pack_launch
names, values = pack_launch(launch)
@@ -777,28 +781,35 @@ def device_entry(*, launch=None, kernel_attrs=None) ->
frame.RegionFrame:
"cuda.launch_fields": names,
"cuda.kernel_attrs": pack_kernel_attrs(kernel_attrs, launch),
},
+ body_params=body_params,
)
-def device_context(device_type: Expr, device_id: Expr) -> frame.RegionFrame:
+def device_context(
+ device_type: Expr, device_id: Expr, *, attrs=None, body_params=None
+) -> frame.RegionFrame:
"""Supply lexical device context for allocation and packed-call lowering.
This region does not change the active runtime device.
"""
- return region("tirx.device_context", [device_type, device_id])
+ return region(
+ "tirx.device_context", [device_type, device_id], attrs=attrs,
body_params=body_params
+ )
-def compute_scope(name: str) -> frame.RegionFrame:
+def compute_scope(name: str, *, attrs=None, body_params=None) ->
frame.RegionFrame:
"""Outline the body as a named CPU compute helper."""
- return region("tirx.compute_scope", [_StringImm(name)])
+ return region("tirx.compute_scope", [name], attrs=attrs,
body_params=body_params)
-def parallel_launch() -> frame.RegionFrame:
+def parallel_launch(*, attrs=None, body_params=None) -> frame.RegionFrame:
"""Launch a CPU worker team around parallel loops and team barriers."""
- return region("tirx.parallel_launch", [])
+ return region("tirx.parallel_launch", [], attrs=attrs,
body_params=body_params)
-def launch_thread(thread_tag: str, extent: Expr) -> frame.RegionFrame:
+def launch_thread(
+ thread_tag: str, extent: Expr, *, attrs=None, body_params=None
+) -> frame.RegionFrame:
"""Launch a hardware or virtual thread with a fresh lexical variable.
The extent determines the variable's scalar integer type. Tags starting
with
@@ -812,7 +823,7 @@ def launch_thread(thread_tag: str, extent: Expr) ->
frame.RegionFrame:
with T.launch_thread("threadIdx.x", 32) as tx:
T.evaluate(tx)
"""
- return region("tirx.launch_thread", [_StringImm(thread_tag), extent])
+ return region("tirx.launch_thread", [thread_tag, extent], attrs=attrs,
body_params=body_params)
# --------------------------------------
diff --git a/python/tvm/tirx/tensor_instruction.py
b/python/tvm/tirx/tensor_instruction.py
index 890e6cb96a..8b732a7a92 100644
--- a/python/tvm/tirx/tensor_instruction.py
+++ b/python/tvm/tirx/tensor_instruction.py
@@ -35,6 +35,7 @@ from tvm.ir import (
PointerType,
PrimType,
StringImm,
+ TensorLoad,
TensorRegion,
Tuple,
const,
@@ -86,11 +87,6 @@ class Instruction:
op.set_signature(args=[arg.name for arg in self.operands])
op.set_attrs_type_key("tirx.tensor." + self.schema)
tvm_ffi.get_global_func("tirx.ConfigureTensorInstruction")(op,
self.validate)
- register_op_attr(
- self.name,
- "__tvm_doc_translate_op_call__",
- tvm_ffi.get_global_func("script.printer.TensorCallDocTranslate")(),
- )
SPECS[self.name] = self
def construct(*args, **qualifiers):
@@ -279,12 +275,16 @@ def _is_absent(value):
def _operand(value, role="expr"):
from tvm.tirx import is_tensor_var
- if value is None:
+ if value is None or _is_absent(value) or (isinstance(value, tuple | list)
and not value):
return Tuple([])
if role == "region":
+ if isinstance(value, TensorLoad):
+ value = tvm_ffi.get_global_func("tirx.AsTensorRegion")(value)
return _region(value)
if role == "selectors":
- return Tuple([Tuple([_expr(cond), _region(candidate)]) for cond,
candidate in value])
+ return Tuple(
+ [Tuple([_expr(cond), _operand(candidate, "region")]) for cond,
candidate in value]
+ )
if role == "tensor":
return value
if is_tensor_var(value):
diff --git a/src/ir/side_effect.cc b/src/ir/side_effect.cc
index 6a3bd54418..7173f29543 100644
--- a/src/ir/side_effect.cc
+++ b/src/ir/side_effect.cc
@@ -25,21 +25,12 @@
#include <tvm/ir/op.h>
#include <algorithm>
-#include <unordered_set>
namespace tvm {
CallEffectKind SideEffect(const Expr& expr) {
static auto effects =
Op::GetAttrMap<TCallEffectKind>(tvm::op_attr::kCallEffectKind);
CallEffectKind kind = CallEffectKind::kPure;
- std::unordered_set<const ExprNode*> visited;
- auto visit_expr =
- [&](const ExprNode* node,
- ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
- // Effects are context-independent; inspect each shared expression only
once.
- if (!visited.insert(node).second) return std::nullopt;
- return visitor->DefaultVisitExpected(node);
- };
ffi::StructuralVisit(
expr,
[&](const CallNode* node,
@@ -48,19 +39,18 @@ CallEffectKind SideEffect(const Expr& expr) {
effects.get(node->op,
static_cast<TCallEffectKind>(CallEffectKind::kOpaque)));
kind = std::max(kind, std::min(effect, CallEffectKind::kUpdateState));
if (kind == CallEffectKind::kUpdateState) return ffi::VisitInterrupt();
- return visit_expr(node, visitor);
+ return visitor->DefaultVisitExpected(node);
},
[&](const TensorLoadNode* node,
ffi::StructuralVisitorObj* visitor) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
kind = std::max(kind, CallEffectKind::kReadState);
- return visit_expr(node, visitor);
+ return visitor->DefaultVisitExpected(node);
},
[](const TypeNode*,
ffi::StructuralVisitorObj*) ->
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> {
// Effects describe evaluated expressions, not expressions embedded in
types.
return std::nullopt;
- },
- visit_expr);
+ });
return kind;
}
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index eaaae8e204..5c0469eef3 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -32,12 +32,50 @@ namespace printer {
namespace details {
namespace {
+// These semantic wrappers convert bare Python scalar tuple fields to 64-bit
+// values. Use raw construction when ordinary tuple syntax would change the IR.
+bool InlineTupleNeedsRawCall(DocTranslatorObj* d, const Expr& expr) {
+ if (const auto* tuple = expr.as<TupleNode>()) {
+ for (const Expr& field : tuple->fields) {
+ if (InlineTupleNeedsRawCall(d, field)) return true;
+ }
+ } else if (const auto* imm = expr.as<IntImmNode>()) {
+ DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
+ auto implicit =
+ ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.int_dtype",
"int32"));
+ return dtype == implicit && dtype != ffi::StringToDLDataType("int64");
+ } else if (const auto* imm = expr.as<FloatImmNode>()) {
+ DLDataType dtype = imm->ty.as_or_throw<PrimType>()->dtype;
+ auto implicit =
+
ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.float_dtype",
"void"));
+ return dtype == implicit && dtype != ffi::StringToDLDataType("float64");
+ }
+ return false;
+}
+
+ffi::Optional<ExprDoc> InlineTupleCallDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
+ const ffi::Object*) {
+ const auto* call =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
+ if (call->args.size() > 1 && InlineTupleNeedsRawCall(d, call->args[1]))
return RawCall(d, call);
+ if (auto doc = StandardCallDocTranslate(d, call)) return doc;
+ return RawCall(d, call);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ for (const char* name : {"relax.call_builtin_with_ctx", "relax.make_closure",
+ "relax.invoke_closure",
"relax.invoke_pure_closure"}) {
+
OpDef(name).set_attr<FDocTranslate>(tvm::script::printer::op_attr::kOpCallDocTranslate,
+
FDocTranslate::FromNative<&InlineTupleCallDocTranslate>());
+ }
+}
+
ffi::Optional<ExprDoc> CallDPSPackedDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object*) {
const auto* call =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
if (call->attrs.defined() || call->args.size() != 2 ||
!call->args[1].as<TupleNode>() ||
- call->ty_args.size() != 1) {
+ call->ty_args.size() != 1 || InlineTupleNeedsRawCall(d, call->args[1])) {
return RawCall(d, call);
}
ExprDoc callee = d->Translate(call->args[0]).value();
@@ -49,11 +87,9 @@ ffi::Optional<ExprDoc>
CallDPSPackedDocTranslate(DocTranslatorObj* d, ffi::AnyVi
} else if (call->args[0].as<StringImmNode>()) {
// A bare Python string now constructs ExternFunc, so retain this explicit
IR value.
callee = AnyValue(d, call->args[0]);
- } else {
- callee = MaterializeCallArgument(d, call->args[0], callee);
}
d->RecordOrigin(callee, call->args[0]);
- ExprDoc args = MaterializeCallArgument(d, call->args[1],
d->Translate(call->args[1]).value());
+ ExprDoc args = d->Translate(call->args[1]).value();
ffi::Array<ffi::String> keys = {"ty_args"};
ffi::Array<ExprDoc> values = {ListDoc({TypeValue(d, call->ty_args[0],
false)})};
bool omit_result = false;
diff --git a/src/script/printer/doc.cc b/src/script/printer/doc.cc
index abe00a5883..df34776bf4 100644
--- a/src/script/printer/doc.cc
+++ b/src/script/printer/doc.cc
@@ -293,17 +293,6 @@ DocStringDoc::DocStringDoc(ffi::String docs) {
this->data_ = std::move(n);
}
-OpCallDoc::OpCallDoc(ExprDoc callee, ffi::Array<Doc> args,
ffi::Optional<DictDoc> workspace,
- ffi::Optional<DictDoc> config, ffi::Optional<ExprDoc>
dispatch) {
- ffi::ObjectPtr<OpCallDocNode> n = ffi::make_object<OpCallDocNode>();
- n->callee = callee;
- n->args = args;
- n->workspace = workspace;
- n->config = config;
- n->dispatch = dispatch;
- this->data_ = std::move(n);
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
@@ -513,15 +502,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](ffi::String docs) { return DocStringDoc(docs); });
}
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("script.printer.OpCallDoc",
- [](ExprDoc callee, ffi::Array<Doc> args, DictDoc
workspace, DictDoc config,
- ffi::Optional<ExprDoc> dispatch) {
- return OpCallDoc(callee, args, workspace, config,
dispatch);
- });
-}
-
NamespaceDoc::NamespaceDoc(ffi::String canonical_name) {
auto node = ffi::make_object<NamespaceDocNode>();
node->canonical_name = std::move(canonical_name);
diff --git a/src/script/printer/doc_printer.cc
b/src/script/printer/doc_printer.cc
index 502fa7c2d7..dd131c0d41 100644
--- a/src/script/printer/doc_printer.cc
+++ b/src/script/printer/doc_printer.cc
@@ -253,11 +253,6 @@ class DocPrinter {
*/
virtual void PrintTypedDoc(const DocStringDoc& doc) = 0;
- /*!
- * \brief Virtual method to print a OpCallDoc
- */
- virtual void PrintTypedDoc(const OpCallDoc& doc) = 0;
-
/*!
* \brief Increase the indent level of any content to be
* printed after this call
@@ -660,7 +655,6 @@ void DocPrinter::PrintDoc(const Doc& doc) {
TVM_SCRIPT_PRINT_DISPATCH(ClassDoc);
TVM_SCRIPT_PRINT_DISPATCH(CommentDoc);
TVM_SCRIPT_PRINT_DISPATCH(DocStringDoc);
- TVM_SCRIPT_PRINT_DISPATCH(OpCallDoc);
#undef TVM_SCRIPT_PRINT_DISPATCH
table.Finalize();
return table;
@@ -940,7 +934,6 @@ class PythonDocPrinter : public DocPrinter {
void PrintTypedDoc(const ClassDoc& doc) final;
void PrintTypedDoc(const CommentDoc& doc) final;
void PrintTypedDoc(const DocStringDoc& doc) final;
- void PrintTypedDoc(const OpCallDoc& doc) final;
void PrintAssignment(const ExprDoc& lhs, const ffi::Optional<ExprDoc>& rhs,
const ffi::Optional<ExprDoc>& annotation);
@@ -1543,70 +1536,6 @@ void PythonDocPrinter::PrintTypedDoc(const DocStringDoc&
doc) {
}
}
-void PythonDocPrinter::PrintTypedDoc(const OpCallDoc& doc) {
- PrintDoc(doc->callee);
-
- output_ << "(";
-
- // Print positional args
- bool wrote_any = false;
- for (const Doc& arg : doc->args) {
- if (wrote_any) {
- output_ << ", ";
- }
- wrote_any = true;
- PrintDoc(arg);
- }
- // workspace first (if present and non-empty)
- if (doc->workspace.has_value() && !doc->workspace.value()->keys.empty()) {
- if (wrote_any) output_ << ", ";
- wrote_any = true;
- output_ << "workspace=";
- PrintDoc(doc->workspace.value());
- }
- // dispatch next (if present)
- if (doc->dispatch.has_value()) {
- if (wrote_any) output_ << ", ";
- wrote_any = true;
- output_ << "dispatch=";
- PrintDoc(doc->dispatch.value());
- }
- // Flatten config as keyword args: key=value
- if (doc->config.has_value() && !doc->config.value()->keys.empty()) {
- const auto* dict = doc->config.value().as<DictDocNode>();
- // Only flatten if all keys are literal strings; otherwise, fallback to
config={...}
- bool all_str_keys = true;
- for (const ExprDoc& k : dict->keys) {
- if (!k.as<LiteralDocNode>()) {
- all_str_keys = false;
- break;
- }
- const auto* lit = k.as<LiteralDocNode>();
- if (!lit->value.as<ffi::String>()) {
- all_str_keys = false;
- break;
- }
- }
- if (all_str_keys) {
- int n = dict->keys.size();
- for (int i = 0; i < n; ++i) {
- const auto* lit = dict->keys[i].as<LiteralDocNode>();
- std::string key = lit->value.as_or_throw<ffi::String>();
- if (wrote_any) output_ << ", ";
- wrote_any = true;
- output_ << key << "=";
- PrintDoc(dict->values[i]);
- }
- } else {
- if (wrote_any) output_ << ", ";
- wrote_any = true;
- output_ << "config=";
- PrintDoc(doc->config.value());
- }
- }
- output_ << ")";
-}
-
// Associate each recovered path with its nearest statement container.
Expression
// children keep their more precise paths while annotations use the enclosing
// statement's ordinary comment rendering. Transparent statement lists
introduce
diff --git a/src/script/printer/ir/call.cc b/src/script/printer/ir/call.cc
index f2ba0e3b35..fe331c7dbc 100644
--- a/src/script/printer/ir/call.cc
+++ b/src/script/printer/ir/call.cc
@@ -32,7 +32,6 @@ namespace tvm {
namespace script {
namespace printer {
namespace details {
-namespace {
// The name is published only for a callable implementing the shared Op
contract.
// Semantic wrappers register their own Op hooks.
@@ -55,7 +54,7 @@ ffi::Optional<ExprDoc>
StandardCallDocTranslate(DocTranslatorObj* d, const CallN
}
ffi::Array<ExprDoc> args;
for (const Expr& arg : call->args) {
- args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ args.push_back(d->Translate(arg).value());
}
ffi::Array<ffi::String> keys;
ffi::Array<ExprDoc> values;
@@ -65,6 +64,7 @@ ffi::Optional<ExprDoc>
StandardCallDocTranslate(DocTranslatorObj* d, const CallN
keys.push_back("ty_args");
values.push_back(ListDoc(types));
}
+ bool default_attrs = false;
if (call->attrs.defined()) {
std::vector<std::pair<ffi::String, ffi::Any>> fields;
ffi::reflection::ForEachFieldInfo(
@@ -80,9 +80,8 @@ ffi::Optional<ExprDoc>
StandardCallDocTranslate(DocTranslatorObj* d, const CallN
}
fields.emplace_back(ffi::String(field->name), std::move(value));
});
- // Keep an explicit schema when no attribute fields remain in the spelling,
- // including empty and entirely default-valued attrs.
- if (fields.empty()) return RawCall(d, call);
+ // Named constructors supply the registered attribute schema and defaults.
+ default_attrs = fields.empty();
std::sort(fields.begin(), fields.end(),
[](const auto& a, const auto& b) { return a.first < b.first; });
for (const auto& [key, value] : fields) {
@@ -104,12 +103,15 @@ ffi::Optional<ExprDoc>
StandardCallDocTranslate(DocTranslatorObj* d, const CallN
// Missing information, unavailable inference, or errors require the exact
stored type.
}
if (!omit_result) {
+ if (default_attrs) return RawCall(d, call);
keys.push_back("ty");
values.push_back(TypeValue(d, call->ty));
}
return NamedCallCallee(names[op])->Call(args, keys, values);
}
+namespace {
+
ffi::Optional<ExprDoc> CallDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object* destination) {
const auto* call =
diff --git a/src/script/printer/ir/prim_expr.cc
b/src/script/printer/ir/prim_expr.cc
index 2836c7c516..f254ade234 100644
--- a/src/script/printer/ir/prim_expr.cc
+++ b/src/script/printer/ir/prim_expr.cc
@@ -16,11 +16,13 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/prim/vector_expr.h>
#include <tvm/script/printer/doc_translator.h>
#include <cstring>
+#include <limits>
#include <optional>
#include "utils.h"
@@ -71,40 +73,79 @@ ffi::Array<Doc> TensorIndices(DocTranslatorObj* d, const
ffi::Array<PrimExpr>& i
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);
+ bool indexable = region->ty.as<TensorRegionTypeNode>() &&
!region->region.empty();
+ using Slice =
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>;
+ ffi::Array<ffi::Variant<Slice, PrimExpr>> indices;
+ ffi::Array<Doc> slices;
+ bool has_slice = false;
+ if (indexable) {
+ for (size_t i = 0; i < region->region.size(); ++i) {
+ const Range& range = region->region[i];
+ if (range->min.ty()->dtype.lanes != 1 || range->extent.ty()->dtype.lanes
!= 1) {
+ indexable = false;
+ break;
+ }
+ ExprDoc start = d->Translate(range->min).value();
+ const auto* extent = range->extent.as<IntImmNode>();
+ bool point =
+ extent && extent->value == 1 && ffi::StructuralEqual()(extent->ty,
range->min.ty());
+ if (point && (has_slice || i + 1 < region->region.size())) {
+ slices.push_back(start);
+ indices.push_back(range->min);
+ } else {
+ try {
+ PrimExpr stop = range->min + range->extent;
+ // Scalar constructors accept signed 64-bit Python literals. A
computed
+ // endpoint outside that range cannot reconstruct through indexing.
+ if (const auto* imm = stop.as<IntImmNode>();
+ imm && (imm->value > std::numeric_limits<int64_t>::max() ||
+ imm->value < std::numeric_limits<int64_t>::min())) {
+ indexable = false;
+ break;
+ }
+ slices.push_back(SliceDoc(start, d->Translate(stop).value(),
std::nullopt));
+ indices.push_back(Slice(range->min, stop, std::nullopt));
+ has_slice = true;
+ } catch (const ffi::Error&) {
+ indexable = false;
+ break;
+ }
+ }
+ }
+ }
+ if (indexable) {
+ try {
+ // Ask the existing type-directed subscription owner whether this syntax
+ // reconstructs the exact region, without depending on a dialect type.
+ static const auto realize =
ffi::Function::GetGlobalRequired("ir.SubscriptExprRealize");
+ auto restored = realize(region->source, indices,
region->loc).cast<TensorRegion>();
+ if (ffi::StructuralEqual()(restored, ffi::GetRef<TensorRegion>(region)))
{
+ return d->Translate(region->source).value()[slices];
+ }
+ } catch (const ffi::Error&) {
+ // Unsupported subscription or a different expression kind needs
explicit construction.
+ }
+ }
+
+ ffi::Array<ExprDoc> ranges;
+ for (const Range& range : region->region) {
+ ranges.push_back(
+ NamespaceDoc("ir")
+ ->Attr("Range")
+ ->Attr("from_min_extent")
+ ->Call({d->Translate(range->min).value(),
d->Translate(range->extent).value()}));
+ }
+ return NamespaceDoc("ir")
+ ->Attr("TensorRegion")
+ ->Call({d->Translate(region->source).value(), ListDoc(ranges)}, {"ty"},
+ {TypeValue(d, region->ty, false)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index 126b01eaa9..1469dac7cc 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -160,33 +160,6 @@ ExprDoc CallAttrsValue(DocTranslatorObj* d, const Attrs&
attrs) {
} // namespace
-// Translate operands as IR values when a dialect offers compact expression
sugar.
-ExprDoc MaterializeCallArgument(DocTranslatorObj* d, const Expr& arg, ExprDoc
doc) {
- if (const auto* tuple = arg.as<TupleNode>()) {
- ffi::Array<ExprDoc> fields;
- for (const Expr& field : tuple->fields) {
- fields.push_back(MaterializeCallArgument(d, field,
d->Translate(field).value()));
- }
- doc = NamespaceDoc("ir")->Attr("Tuple")->Call({ListDoc(fields)});
- }
- if (const auto* region = arg.as<TensorRegionNode>()) {
- ffi::Array<ExprDoc> ranges;
- for (const Range& range : region->region) {
- ranges.push_back(
- NamespaceDoc("ir")
- ->Attr("Range")
- ->Attr("from_min_extent")
- ->Call({d->Translate(range->min).value(),
d->Translate(range->extent).value()}));
- }
- doc = NamespaceDoc("ir")
- ->Attr("TensorRegion")
- ->Call({d->Translate(region->source).value(), ListDoc(ranges)},
{"ty"},
- {TypeValue(d, region->ty, false)});
- }
- d->RecordOrigin(doc, arg);
- return doc;
-}
-
// The explicit fallback retains every field, including typed attribute
objects.
ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call,
ffi::Optional<ffi::Array<ExprDoc>> translated_args) {
@@ -197,9 +170,6 @@ ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call,
} else {
for (const Expr& arg : call->args)
args.push_back(d->Translate(arg).value());
}
- for (size_t i = 0; i < args.size(); ++i) {
- args.Set(i, MaterializeCallArgument(d, call->args[i], args[i]));
- }
ExprDoc callee = op.has_value() ?
LiteralDoc::Str(op.value()->name, std::nullopt)
: call->op.as<StringImmNode>() ? AnyValue(d, call->op)
:
d->Translate(call->op).value();
@@ -221,24 +191,19 @@ ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call,
}
// The query aliases the active frame; candidate classification must own a
copy.
-ffi::Dict<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d) {
- ffi::Dict<Var, IdDoc> candidates;
+ffi::Map<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d) {
+ ffi::Map<Var, IdDoc> candidates;
for (const auto& [var, id] : d->GetImplicitDefs()) candidates.Set(var, id);
return candidates;
}
-void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var,
IdDoc>& signature,
+void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Map<Var,
IdDoc>& signature,
const FunctionDoc& function) {
auto module = d->GetOrCreateExtraState<ffi::Optional<IRModule>>("ir.module");
auto pending = d->GetImplicitDefs();
- std::vector<Var> parameters;
if (!module.has_value()) {
- for (const auto& [var, id] : signature) {
- if (pending.count(var) && var->ty.as<PrimTypeNode>())
parameters.push_back(var);
- }
- std::sort(parameters.begin(), parameters.end(),
- [&](const Var& a, const Var& b) { return signature[a]->name <
signature[b]->name; });
- for (const Var& var : parameters) {
+ for (const auto& [var, signature_id] : signature) {
+ if (!pending.count(var) || !var->ty.as<PrimTypeNode>()) continue;
IdDoc id = VarDoc(d, var);
if (var->ty.as_or_throw<PrimType>()->dtype == (DLDataType{kDLInt, 64,
1})) {
function->type_params.push_back(id);
diff --git a/src/script/printer/ir/utils.h b/src/script/printer/ir/utils.h
index 24a72e9e9f..16c69b4151 100644
--- a/src/script/printer/ir/utils.h
+++ b/src/script/printer/ir/utils.h
@@ -37,7 +37,6 @@ 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);
@@ -45,12 +44,12 @@ ExprDoc GlobalReference(DocTranslatorObj* d, const
ffi::String& name);
ExprDoc NamedCallCallee(const ffi::String& canonical_name);
ExprDoc TypeValue(DocTranslatorObj* d, const Type& type, bool dtype_literal =
true);
-ExprDoc MaterializeCallArgument(DocTranslatorObj* d, const Expr& arg, ExprDoc
doc);
+ffi::Optional<ExprDoc> StandardCallDocTranslate(DocTranslatorObj* d, const
CallNode* call);
ExprDoc RawCall(DocTranslatorObj* d, const CallNode* call,
ffi::Optional<ffi::Array<ExprDoc>> translated_args =
std::nullopt);
ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value);
-ffi::Dict<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d);
-void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var,
IdDoc>& signature,
+ffi::Map<Var, IdDoc> CopyImplicitDefs(DocTranslatorObj* d);
+void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Map<Var,
IdDoc>& signature,
const FunctionDoc& function);
class VarScope {
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 325bb01dd5..9d3d58936c 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -39,6 +39,15 @@
#include "../../sym/pattern_match.h"
namespace tvm {
+
+TensorRegion ffi::TypeTraits<TensorRegion>::ConvertFallbackValue(TensorLoad
value) {
+ ffi::Array<Range> ranges;
+ for (const PrimExpr& index : value->indices) {
+ ranges.push_back(Range::FromMinExtent(index, IntImm(index.ty(), 1)));
+ }
+ return tirx::BufferRegion(value->source.as_or_throw<tirx::TensorVar>(),
ranges, value->loc);
+}
+
namespace tirx {
using namespace tvm::prim;
@@ -170,6 +179,7 @@ TensorRegion BufferRegion(TensorVar buffer,
ffi::Array<Range> region, Location l
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("tirx.AsTensorRegion", [](TensorRegion region) {
return region; });
refl::GlobalDef().def("tirx.BufferRegion", [](TensorVar buffer,
ffi::Array<Range> region) {
return BufferRegion(buffer, region);
});
diff --git a/src/tirx/op/region.cc b/src/tirx/op/region.cc
index 72bda4a532..cbe8be5bc0 100644
--- a/src/tirx/op/region.cc
+++ b/src/tirx/op/region.cc
@@ -25,6 +25,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/stmt.h>
+#include <tvm/script/printer/doc_translator.h>
#include <tvm/tirx/op/region.h>
#include <tvm/tirx/op_attr_types.h>
@@ -55,6 +56,8 @@ const Op& launch_thread_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.launch_thread", "Bind a thread index within a body with a launch
extent.")
+ .set_attr<ffi::String>(tvm::script::printer::op_attr::kScriptPrinterName,
+ "tirx.launch_thread")
.signature(sig::arg<StringImm>("tag"), sig::arg<IntExpr>("extent"))
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&LaunchThreadBodyParams>())
@@ -83,6 +86,7 @@ void ValidateDeviceEntry(const RegionStmtNode* region) {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_entry", "Enter a device kernel with independent launch
configuration.")
+
.set_attr<ffi::String>(tvm::script::printer::op_attr::kScriptPrinterName,
"tirx.device_entry")
.signature(sig::var_args("launch_values"),
sig::call_attrs<DictAttrsNode>())
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
@@ -98,6 +102,8 @@ const Op& device_context_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.device_context", "Supply the device type and ID within a
region.")
+ .set_attr<ffi::String>(tvm::script::printer::op_attr::kScriptPrinterName,
+ "tirx.device_context")
.signature(sig::arg<IntExpr>("device_type"),
sig::arg<IntExpr>("device_id"))
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
@@ -113,6 +119,8 @@ const Op& compute_scope_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.compute_scope", "Outline a named CPU compute region.")
+ .set_attr<ffi::String>(tvm::script::printer::op_attr::kScriptPrinterName,
+ "tirx.compute_scope")
.signature(sig::arg<StringImm>("name"))
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
@@ -128,6 +136,8 @@ const Op& parallel_launch_op() {
TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.parallel_launch", "Launch a CPU worker team around a region.")
+ .set_attr<ffi::String>(tvm::script::printer::op_attr::kScriptPrinterName,
+ "tirx.parallel_launch")
.signature()
.set_attr<FRegionGetBodyParams>(tvm::op_attr::kRegionGetBodyParams,
FRegionGetBodyParams::FromNative<&RegionNoBodyParams>())
diff --git a/src/tirx/script/ir_builder/utils.h
b/src/tirx/script/ir_builder/utils.h
index 00b451360e..06286ccbcf 100644
--- a/src/tirx/script/ir_builder/utils.h
+++ b/src/tirx/script/ir_builder/utils.h
@@ -64,11 +64,7 @@ inline FunctionFrame FindFunctionFrame(const ffi::String&
method) {
* \return The converted TensorRegion.
*/
inline tvm::TensorRegion TensorRegionFromLoad(tvm::TensorLoad buffer_load) {
- ffi::Array<Range> ranges;
- for (const PrimExpr& index : buffer_load->indices) {
- ranges.push_back(Range::FromMinExtent(index, IntImm(index.ty(), 1)));
- }
- return
tvm::tirx::BufferRegion(buffer_load->source.as_or_throw<tvm::tirx::TensorVar>(),
ranges);
+ return ffi::TypeTraits<tvm::TensorRegion>::ConvertFallbackValue(buffer_load);
}
} // namespace tirx
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 4da38c5c06..0338e8ed3d 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -94,13 +94,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FDocTranslate::FromNative<&VarDocTranslate>());
}
-bool CanTranslateExplicitResultCall(const CallNode* call) {
- return !call->attrs.defined() && call->ty_args.empty() &&
call->ty.as<PrimType>() &&
- std::all_of(call->args.begin(), call->args.end(), [](const Expr& arg)
{
- return !arg->ty.as<MissingType>().has_value() &&
!arg.as<TensorRegionNode>();
- });
-}
-
namespace {
ffi::Optional<ExprDoc> IndexMapDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
@@ -139,6 +132,12 @@ ffi::Optional<ExprDoc>
LLVMIntrinsicDocTranslate(DocTranslatorObj* d, ffi::AnyVi
!call->ty.as<PrimType>()) {
return RawCall(d, call);
}
+ // The LLVM convenience builders reject direct region operands. Keep their
+ // raw Call representation while ordinary tuple operands recurse normally.
+ if (std::any_of(call->args.begin(), call->args.end(),
+ [](const Expr& arg) { return arg.as<TensorRegionNode>() !=
nullptr; })) {
+ return RawCall(d, call);
+ }
const auto* id = call->args[0].as<IntImmNode>();
// The named constructor uses an int32 intrinsic identifier. Other stored
// representations retain their exact operand type through the full Call.
@@ -159,7 +158,7 @@ ffi::Optional<ExprDoc>
LLVMIntrinsicDocTranslate(DocTranslatorObj* d, ffi::AnyVi
d->RecordOrigin(name_doc, call->args[0]);
ffi::Array<ExprDoc> args = {name_doc};
for (size_t i = 1; i < call->args.size(); ++i) {
- args.push_back(MaterializeCallArgument(d, call->args[i],
d->Translate(call->args[i]).value()));
+ args.push_back(d->Translate(call->args[i]).value());
}
return NamespaceDoc("tirx")
->Attr(call->op.same_as(tirx::call_llvm_intrin_op()) ? "call_llvm_intrin"
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index f7330c1fef..cac5c52da2 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -42,52 +42,6 @@ namespace details {
namespace {
-ExprDoc TensorOperandDoc(DocTranslatorObj* d, const Expr& value) {
- if (const auto* region = value.as<TensorRegionNode>()) {
- ExprDoc doc = TensorRegionValue(d, region, true);
- d->RecordOrigin(doc, value);
- return doc;
- }
- if (const auto* tuple = value.as<TupleNode>()) {
- if (tuple->fields.empty()) return LiteralDoc::None(std::nullopt);
- ffi::Array<ExprDoc> fields;
- for (const auto& field : tuple->fields)
fields.push_back(TensorOperandDoc(d, field));
- return TupleDoc(fields);
- }
- return AnyValue(d, value);
-}
-
-ffi::Optional<ExprDoc> TensorCallDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
- const ffi::Object* destination) {
- const auto* call =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
- const auto op = call->op.as_or_throw<Op>();
- op.Validate(call);
- ffi::Array<ExprDoc> args;
- for (const auto& arg : call->args) args.push_back(TensorOperandDoc(d, arg));
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- ffi::reflection::ForEachFieldInfo(
- TVMFFIGetTypeInfo(call->attrs->type_index()), [&](const TVMFFIFieldInfo*
field) {
- ffi::Any value = ffi::reflection::FieldGetter(field)(call->attrs);
- if ((field->flags & kTVMFFIFieldFlagBitMaskHasDefault) &&
- ffi::StructuralEqual()(
- value,
ffi::AnyView::CopyFromTVMFFIAny(field->default_value_or_factory)))
- return;
- keys.push_back(ffi::String(field->name));
- values.push_back(AnyValue(d, value));
- });
- static const auto& names =
-
Op::GetAttrMap<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName);
- return NamedCallCallee(names[op])->Call(args, keys, values);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::GlobalDef().def("script.printer.TensorCallDocTranslate",
[]() {
- return FDocTranslate::FromNative<&TensorCallDocTranslate>();
- });
-}
-
ffi::Optional<ExprDoc> EvaluateDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object* destination) {
const auto* stmt =
@@ -172,48 +126,45 @@ ffi::Optional<ExprDoc>
RegionStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView
// Inputs, attributes, and parameter types are evaluated before the body
// parameters enter scope. Explicit Var constructors preserve their exact
types.
- ExprDoc rhs(ffi::UnsafeInit{});
- if (stmt->op.same_as(tirx::device_entry_op()) && stmt->attrs->dict.empty()) {
- rhs = NamespaceDoc("tirx")->Attr("device_entry")->Call({});
- } else if (stmt->op.same_as(tirx::launch_thread_op())) {
- rhs = NamespaceDoc("tirx")
- ->Attr("launch_thread")
-
->Call({LiteralDoc::Str(stmt->args[0].as_or_throw<StringImm>()->value,
std::nullopt),
-
d->Translate(stmt->args[1].as_or_throw<PrimExpr>()).value()});
- } else if (stmt->op.same_as(tirx::device_context_op())) {
- rhs = NamespaceDoc("tirx")
- ->Attr("device_context")
- ->Call({d->Translate(stmt->args[0]).value(),
d->Translate(stmt->args[1]).value()});
- } else if (stmt->op.same_as(tirx::compute_scope_op())) {
- rhs =
- NamespaceDoc("tirx")
- ->Attr("compute_scope")
-
->Call({LiteralDoc::Str(stmt->args[0].as_or_throw<StringImm>()->value,
std::nullopt)});
- } else if (stmt->op.same_as(tirx::parallel_launch_op())) {
- rhs = NamespaceDoc("tirx")->Attr("parallel_launch")->Call({});
- } else {
- ffi::Array<ExprDoc> args;
- for (const Expr& arg : stmt->args)
args.push_back(d->Translate(arg).value());
- for (size_t i = 0; i < args.size(); ++i) {
- args.Set(i, MaterializeCallArgument(d, stmt->args[i], args[i]));
- }
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- if (!stmt->body_params.empty()) {
- ffi::Array<ExprDoc> params;
- for (const Var& param : stmt->body_params) {
- ExprDoc value = NamespaceDoc("tirx")->Attr("Var")->Call(
- {LiteralDoc::Str(param->name, std::nullopt), TypeValue(d,
param->ty)});
- d->RecordOrigin(value, param);
- params.push_back(value);
- }
- keys.push_back("body_params");
- values.push_back(ListDoc(params));
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : stmt->args) args.push_back(d->Translate(arg).value());
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ bool infer_params = false;
+ try {
+ auto inferred = GetRegionBodyParams(stmt->op, stmt->args, stmt->attrs);
+ infer_params = inferred.size() == stmt->body_params.size();
+ for (size_t i = 0; infer_params && i < inferred.size(); ++i) {
+ infer_params = ffi::StructuralEqual()(inferred[i]->ty,
stmt->body_params[i]->ty);
}
- if (!stmt->attrs->dict.empty()) {
- keys.push_back("attrs");
- values.push_back(AnyValue(d, stmt->attrs));
+ } catch (const ffi::Error&) {
+ // Explicit parameters preserve regions whose inference is unavailable.
+ }
+ if (!infer_params) {
+ ffi::Array<ExprDoc> params;
+ for (const Var& param : stmt->body_params) {
+ ExprDoc value = NamespaceDoc("tirx")->Attr("Var")->Call(
+ {LiteralDoc::Str("", std::nullopt), TypeValue(d, param->ty)});
+ d->RecordOrigin(value, param);
+ params.push_back(value);
}
+ keys.push_back("body_params");
+ values.push_back(ListDoc(params));
+ }
+ if (!stmt->attrs->dict.empty()) {
+ keys.push_back("attrs");
+ values.push_back(AnyValue(d, stmt->attrs));
+ }
+ ExprDoc rhs(ffi::UnsafeInit{});
+ ffi::Optional<ffi::String> name;
+ if (Op::HasAttrMap(tvm::script::printer::op_attr::kScriptPrinterName)) {
+ auto names =
+
Op::GetAttrMap<TScriptPrinterName>(tvm::script::printer::op_attr::kScriptPrinterName);
+ if (names.count(stmt->op)) name = names[stmt->op];
+ }
+ if (name.has_value() && !name.value().empty()) {
+ rhs = NamedCallCallee(name.value())->Call(args, keys, values);
+ } else {
rhs = NamespaceDoc("tirx")->Attr("region")->Call(
{LiteralDoc::Str(stmt->op->name, std::nullopt), ListDoc(args)}, keys,
values);
}
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index 97c6059923..d6011b497f 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -34,7 +34,6 @@ namespace details {
void PrintFunction(DocTranslatorObj* d, const tirx::FunctionNode* func,
ExprDoc decorator,
const ffi::String& dialect_attr);
-bool CanTranslateExplicitResultCall(const CallNode* call);
bool IsScalarBuffer(DocTranslatorObj* d, const Expr& source);
ffi::Optional<ExprDoc> VarDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object* destination);
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index 8cb18f2253..dbf8a9f15f 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -166,7 +166,7 @@ from __future__ import annotations
@R.function
def foo(x: R.DTensor((128, 128), "float32", R.device_mesh((2, 2), I.Range(0,
4)), "S[0], R")) -> R.DTensor((128, 128), "float32", R.device_mesh((2, 2),
I.Range(0, 4)), "S[0], R"):
- gv0 = R.call_tir(Module.tir_func, I.Tuple([x]), ty_args=[R.DTensor((128,
128), "float32", R.device_mesh((2, 2), I.Range(0, 4)), "S[0], R")])
+ gv0 = R.call_tir(Module.tir_func, (x,), ty_args=[R.DTensor((128, 128),
"float32", R.device_mesh((2, 2), I.Range(0, 4)), "S[0], R")])
return gv0
""",
)
@@ -198,13 +198,13 @@ class Module:
with Ts.sblock(""):
v = Ts.axis.spatial(T.int64(128), i, dtype="int64")
v_1 = Ts.axis.spatial(T.int64(128), j, dtype="int64")
- Ts.reads(x[v, v_1])
- Ts.writes(y[v, v_1])
+ Ts.reads(x[v, v_1:v_1 + T.int64(1)])
+ Ts.writes(y[v, v_1:v_1 + T.int64(1)])
y[v, v_1] = x[v, v_1] + T.float32(1.0)
@R.function
def foo(x: R.DTensor((128, 128), "float32", "mesh[0]", "S[0], R")) ->
R.DTensor((128, 128), "float32", "mesh[0]", "S[0], R"):
- gv0 = R.call_tir(Module.tir_func, I.Tuple([x]),
ty_args=[R.DTensor((128, 128), "float32", "mesh[0]", "S[0], R")])
+ gv0 = R.call_tir(Module.tir_func, (x,), ty_args=[R.DTensor((128, 128),
"float32", "mesh[0]", "S[0], R")])
return gv0
""",
)
@@ -720,7 +720,7 @@ def test_call():
"""
x = I.dynamic("x", dtype="int64")
a: R.Tensor((1, x, 3), dtype="float32")
-I.Call("relax.call_tir", [Module.tir_func, I.Tuple([a, x])],
ty_args=[R.Tensor((1, x, 3), dtype="float32")], ty=R.Tensor((1, x, 3),
dtype="float32"))
+I.Call("relax.call_tir", [Module.tir_func, (a, x)], ty_args=[R.Tensor((1, x,
3), dtype="float32")], ty=R.Tensor((1, x, 3), dtype="float32"))
""",
)
_assert_print_lines(
@@ -728,7 +728,7 @@ I.Call("relax.call_tir", [Module.tir_func, I.Tuple([a,
x])], ty_args=[R.Tensor((
"""
x = I.dynamic("x", dtype="int64")
a: R.Tensor((1, x, 3), dtype="float32")
-R.call_dps_packed("my_dps_func", I.Tuple([a]), ty_args=[R.Tensor((1, x, 3),
dtype="float32")])
+R.call_dps_packed("my_dps_func", (a,), ty_args=[R.Tensor((1, x, 3),
dtype="float32")])
""",
)
@@ -749,7 +749,7 @@ def test_call_tir_with_grad():
"""
v0: R.Tensor((54, 96), dtype="float32")
x = I.dynamic("x", dtype="int64")
-I.Call("relax.call_tir_with_grad", [Module.tir_func, I.Tuple([v0])],
attrs=I.make_node("relax.attrs.CallTIRWithGradAttrs", te_grad_kwargs={"k": 1.0,
"x": x}, te_grad_name="grad_func"), ty_args=[R.Tensor((54, 96),
dtype="float32")], ty=R.Tensor((54, 96), dtype="float32"))
+I.Call("relax.call_tir_with_grad", [Module.tir_func, (v0,)],
attrs=I.make_node("relax.attrs.CallTIRWithGradAttrs", te_grad_kwargs={"k": 1.0,
"x": x}, te_grad_name="grad_func"), ty_args=[R.Tensor((54, 96),
dtype="float32")], ty=R.Tensor((54, 96), dtype="float32"))
""",
)
@@ -776,7 +776,7 @@ def test_call_tir_inplace():
x: R.Tensor((32, 32), dtype="int32")
y: R.Tensor((32, 32), dtype="int32")
t = I.dynamic("t", dtype="int64")
-I.Call("relax.call_tir_inplace", [Module.tir_func, I.Tuple([x, y, t])],
attrs=I.make_node("relax.attrs.CallTIRInplaceAttrs", inplace_indices=[-1, 0]),
ty_args=[R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32),
dtype="int32"))], ty=R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32,
32), dtype="int32")))
+I.Call("relax.call_tir_inplace", [Module.tir_func, (x, y, t)],
attrs=I.make_node("relax.attrs.CallTIRInplaceAttrs", inplace_indices=[-1, 0]),
ty_args=[R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32),
dtype="int32"))], ty=R.Tuple(R.Tensor((32, 32), dtype="int32"), R.Tensor((32,
32), dtype="int32")))
""",
)
@@ -965,7 +965,7 @@ class Module:
@R.function
def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,),
dtype="float32"):
- gv0 = R.call_tir(Module.tir_func, I.Tuple([x]),
ty_args=[R.Tensor((128,), dtype="float32")])
+ gv0 = R.call_tir(Module.tir_func, (x,), ty_args=[R.Tensor((128,),
dtype="float32")])
return gv0
""",
)
@@ -990,7 +990,7 @@ class Module:
@R.function
def foo(x: R.Tensor((128,), dtype="float32")) -> R.Tensor((128,),
dtype="float32"):
- gv0 = R.call_tir(Module.tir_func, I.Tuple([x]),
ty_args=[R.Tensor((128,), dtype="float32")])
+ gv0 = R.call_tir(Module.tir_func, (x,), ty_args=[R.Tensor((128,),
dtype="float32")])
return gv0
""",
)
@@ -1163,14 +1163,13 @@ def test_reused_extern_func():
"""
from __future__ import annotations
-# from tvm.script import ir as I
# from tvm.script import relax as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128),
dtype="float32"):
extern_func: R.Callable(derive_func="tvm.relax.type.infer_by_ty_args") =
R.ExternFunc("extern_func")
- y = R.call_dps_packed(extern_func, I.Tuple([x]), ty_args=[R.Tensor((128,
128), dtype="float32")])
- z = R.call_dps_packed(extern_func, I.Tuple([y]), ty_args=[R.Tensor((128,
128), dtype="float32")])
+ y = R.call_dps_packed(extern_func, (x,), ty_args=[R.Tensor((128, 128),
dtype="float32")])
+ z = R.call_dps_packed(extern_func, (y,), ty_args=[R.Tensor((128, 128),
dtype="float32")])
return z
""",
)
@@ -1194,13 +1193,12 @@ def test_inline_extern_func():
"""
from __future__ import annotations
-# from tvm.script import ir as I
# from tvm.script import relax as R
@R.function
def func(x: R.Tensor((128, 128), dtype="float32")) -> R.Tensor((128, 128),
dtype="float32"):
- y = R.call_dps_packed("extern_func", I.Tuple([x]), ty_args=[R.Tensor((128,
128), dtype="float32")])
- z = R.call_dps_packed("extern_func", I.Tuple([y]), ty_args=[R.Tensor((128,
128), dtype="float32")])
+ y = R.call_dps_packed("extern_func", (x,), ty_args=[R.Tensor((128, 128),
dtype="float32")])
+ z = R.call_dps_packed("extern_func", (y,), ty_args=[R.Tensor((128, 128),
dtype="float32")])
return z
""",
)
diff --git a/tests/python/relax/test_transform_error_enrichment.py
b/tests/python/relax/test_transform_error_enrichment.py
index 53e5d39f41..97a954cfcd 100644
--- a/tests/python/relax/test_transform_error_enrichment.py
+++ b/tests/python/relax/test_transform_error_enrichment.py
@@ -59,15 +59,12 @@ def test_pass_error_renders_underlined_tvmscript():
"Location (TVMScript):\n"
"Access path: <root>.body.blocks[0].bindings[0].value\n\n"
"from __future__ import annotations\n\n"
- "# from tvm.script import ir as I\n"
"# from tvm.script import relax as R\n\n"
"@R.function\n"
'def main(x: R.Tensor((3, 4), dtype="float32"), '
'y: R.Tensor((5, 6), dtype="float32")) -> R.Tensor((3, 6),
dtype="float32"):\n'
- ' lv = I.Call("relax.matmul", [x, y],
attrs=I.make_node("relax.attrs.M'
- 'atmulAttrs", out_dtype=None), ty=I.MissingType())\n'
- "
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^"
- "^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n"
+ " lv = R.matmul(x, y)\n"
+ " ^^^^^^^^^^^^^^\n"
" return lv"
)
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py
b/tests/python/tirx/script/test_tirx_script_basic_usage.py
index d0802f9e69..0f9a3694d2 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -334,6 +334,14 @@ def func() -> None:
assert_structural_equal(func, from_source(func.script()))
+def test_tensor_region_roundtrip():
+ @T.function
+ def func(A: T.Tensor((4, 4), "float32")):
+ T.cuda.tile.sqrt(A[0, 1:2], A[1, 2:3])
+
+ assert_structural_equal(func, from_source(func.script()))
+
+
def test_roundtrip_unary_inplace():
"""Single-arg unary ops (in-place) should round-trip."""
