This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 9bfefb7e4b [TIRx] Introduce first-class Return statement (#20018)
9bfefb7e4b is described below
commit 9bfefb7e4b2f13f92b59ed4755bc83d856369c40
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Jul 17 05:34:50 2026 +0800
[TIRx] Introduce first-class Return statement (#20018)
Return is control flow, but TIRx currently represents it as an
Evaluate-wrapped intrinsic call. This prevents return values from
participating naturally in statement traversal and requires special-case
handling across the pipeline.
This change introduces a reflected tirx.Return statement carrying an
Expr, wires it through TVMScript, statement visitors and mutators,
lowering, storage planning, and C/LLVM code generation, and removes the
legacy tirx.ret and T.ret surfaces.
---
include/tvm/tirx/builtin.h | 4 --
include/tvm/tirx/op.h | 10 -----
include/tvm/tirx/script/builder/ir.h | 6 +++
include/tvm/tirx/stmt.h | 28 +++++++++++++
include/tvm/tirx/stmt_functor.h | 4 ++
python/tvm/tirx/__init__.py | 4 +-
python/tvm/tirx/op.py | 21 ----------
python/tvm/tirx/script/builder/ir.py | 8 +++-
python/tvm/tirx/script/parser/parser.py | 2 +-
python/tvm/tirx/stmt.py | 20 +++++++++
python/tvm/tirx/stmt_functor.py | 18 ++++++++
src/relax/op/tensor/inspect.cc | 4 +-
src/relax/transform/compute_prim_value.cc | 3 +-
src/target/llvm/codegen_llvm.cc | 31 +++++++-------
src/target/llvm/codegen_llvm.h | 1 +
src/target/source/codegen_c.cc | 12 ++++--
src/target/source/codegen_c.h | 1 +
src/tirx/ir/stmt.cc | 16 ++++++++
src/tirx/ir/stmt_functor.cc | 13 ++++++
src/tirx/ir/tir_visitor_with_path.cc | 4 ++
src/tirx/ir/tir_visitor_with_path.h | 1 +
src/tirx/op/builtin.cc | 5 ---
src/tirx/op/op.cc | 14 -------
src/tirx/script/builder/ir.cc | 3 ++
src/tirx/script/printer/stmt.cc | 29 ++++---------
src/tirx/transform/make_packed_api.cc | 19 +++------
src/tirx/transform/split_host_device.cc | 27 +++---------
src/tirx/transform/storage_rewrite.cc | 11 +++++
tests/python/codegen/test_target_codegen_cuda.py | 2 +-
tests/python/codegen/test_target_codegen_llvm.py | 6 +--
.../relax/test_transform_compute_prim_value.py | 4 +-
.../test_s_tir_analysis_is_pure_function.py | 10 ++---
tests/python/tirx-base/test_tir_base.py | 48 ++++++++++++++++++----
tests/python/tirx-base/test_tir_imm_values.py | 46 ++++++++++-----------
tests/python/tirx-base/test_tir_specialize.py | 4 +-
.../test_tir_inline_private_functions.py | 2 +-
.../test_tir_transform_make_packed_api.py | 13 ++++++
.../test_tir_transform_split_host_device.py | 23 ++++++++++-
tests/python/tirx/transform/test_stmt_functor.py | 25 +++++++++++
.../python/tvmscript/test_tvmscript_parser_tir.py | 6 +--
.../python/tvmscript/test_tvmscript_printer_tir.py | 3 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 8 ++--
.../tvmscript/test_tvmscript_syntax_sugar.py | 4 +-
43 files changed, 328 insertions(+), 195 deletions(-)
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index cbff6680b6..ee9d698ed2 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -41,10 +41,6 @@ namespace tirx {
/*! \brief Collection of builtin intrinsics as ops */
namespace builtin {
-/*!
- * \brief Return value.
- */
-TVM_DLL const Op& ret();
/*!
* \brief Return from a GPU thread.
*/
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index edf6477d79..3be4fb739b 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -78,16 +78,6 @@ TVM_DLL Type GetType(const PrimExpr& expr);
*/
TVM_DLL Type GetTypeFromRuntimeDataType(DLDataType dtype);
-/*!
- * \brief Return the value.
- *
- * \param value The returned value.
- * \param span The location of this operation in the source.
- * \return The return expression.
- */
-TVM_DLL PrimExpr ret(PrimExpr value, Span span = Span());
-TVM_DLL Expr ret(Expr value, Span span = Span());
-
/*!
* \brief Return from a thread.
*
diff --git a/include/tvm/tirx/script/builder/ir.h
b/include/tvm/tirx/script/builder/ir.h
index 2a8a313370..9997e3ec52 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -367,6 +367,12 @@ AttrFrame DeviceEntry();
*/
WhileFrame While(PrimExpr condition);
+/*!
+ * \brief Create a return statement.
+ * \param value The value to return.
+ */
+void Return(Expr value);
+
/*!
* \brief Create a break statement.
*/
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 1d06f1568c..af7ef30ae5 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -687,6 +687,34 @@ class While : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(WhileNode);
};
+/*!
+ * \brief A return from the current function.
+ */
+class ReturnNode : public StmtNode {
+ public:
+ /*! \brief The value to return. */
+ Expr value;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<ReturnNode>().def_ro("value", &ReturnNode::value);
+ }
+
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Return", ReturnNode, StmtNode);
+};
+
+/*!
+ * \brief Managed reference to ReturnNode.
+ * \sa ReturnNode
+ */
+class Return : public Stmt {
+ public:
+ TVM_DLL explicit Return(Expr value, Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Return, Stmt, ReturnNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(ReturnNode);
+};
+
/*!
* \brief A Break in control flow.
*/
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index d4cf5a7f9c..0df163d5f9 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -90,6 +90,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
virtual R VisitStmt_(const IfThenElseNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
+ virtual R VisitStmt_(const ReturnNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const BreakNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const ContinueNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const AllocBufferNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
@@ -116,6 +117,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode);
IR_STMT_FUNCTOR_DISPATCH(ForNode);
IR_STMT_FUNCTOR_DISPATCH(WhileNode);
+ IR_STMT_FUNCTOR_DISPATCH(ReturnNode);
IR_STMT_FUNCTOR_DISPATCH(BreakNode);
IR_STMT_FUNCTOR_DISPATCH(ContinueNode);
IR_STMT_FUNCTOR_DISPATCH(AllocBufferNode);
@@ -173,6 +175,7 @@ class TVM_DLL StmtVisitor : protected
StmtFunctor<void(const Stmt&)> {
void VisitStmt_(const IfThenElseNode* op) override;
void VisitStmt_(const ForNode* op) override;
void VisitStmt_(const WhileNode* op) override;
+ void VisitStmt_(const ReturnNode* op) override;
void VisitStmt_(const BreakNode* op) override;
void VisitStmt_(const ContinueNode* op) override;
void VisitStmt_(const AllocBufferNode* op) override;
@@ -293,6 +296,7 @@ class TVM_DLL StmtMutator : protected
StmtFunctor<Stmt(const Stmt&)> {
Stmt VisitStmt_(const IfThenElseNode* op) override;
Stmt VisitStmt_(const ForNode* op) override;
Stmt VisitStmt_(const WhileNode* op) override;
+ Stmt VisitStmt_(const ReturnNode* op) override;
Stmt VisitStmt_(const BreakNode* op) override;
Stmt VisitStmt_(const ContinueNode* op) override;
Stmt VisitStmt_(const AllocBufferNode* op) override;
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index efec3e35d6..8d530465e5 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -34,7 +34,7 @@ from .expr import Min, Max, EQ, NE, LT, LE, GT, GE, And, Or,
Not
from .expr import Select, BufferLoad, ProducerLoad, Ramp, Broadcast, Shuffle
from .expr import CallEffectKind, Let, IterVar, CommReducer
-from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While
+from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While, Return, Break,
Continue
# Legacy alias: LetStmt was folded into Bind (which now accepts an optional
body)
LetStmt = Bind
@@ -50,7 +50,7 @@ from .function import PrimFunc, TensorIntrin, IndexMap
from .op import call_packed_lowered, call_cpacked_lowered, call_tir
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
-from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any,
min_value, max_value, trace
+from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value, trace
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get,
tvm_struct_set
from .op import address_of, lookup_param, assume, undef
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index f15b6f554f..782fbcfe58 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -1182,27 +1182,6 @@ def dp4a(vec1, vec2, acc=0):
return call_intrin("int32", "tirx.dp4a", vec1, vec2, acc)
-def ret(val, span=None):
- """Create a tir return expression
-
- Parameters
- ----------
- val : Expr
- The returned tir expression, whose data type is int, float or void
pointer.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- ret : Expr
- The return expression
- """
- if not isinstance(val, Expr):
- val = tirx.convert(val)
- return Call(Op.get("tirx.ret"), [val], span=span, ret_ty=val.ty)
-
-
def any(*args, span=None):
"""Create a new experssion of the union of all conditions in the arguments
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index 4e839565a8..ccd091ae63 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -1700,6 +1700,11 @@ def While(condition: Expr) -> frame.WhileFrame: #
pylint: disable=invalid-name
return _ffi_api.While(condition) # type: ignore[attr-defined] # pylint:
disable=no-member
+def Return(value: Expr) -> None: # pylint: disable=invalid-name
+ """Create a return node."""
+ return _ffi_api.Return(value) # type: ignore[attr-defined] # pylint:
disable=no-member
+
+
def Break() -> None: # pylint: disable=invalid-name
"""Create a break node."""
return _ffi_api.Break() # type: ignore[attr-defined] # pylint:
disable=no-member
@@ -3118,7 +3123,6 @@ popcount = _op_wrapper(_tir_op.popcount)
pow = _op_wrapper(_tir_op.pow) # pylint: disable=redefined-builtin
q_multiply_shift = _op_wrapper(_tir_op.q_multiply_shift)
q_multiply_shift_per_axis = _op_wrapper(_tir_op.q_multiply_shift_per_axis)
-ret = _op_wrapper(_tir_op.ret)
continue_loop = _op_wrapper(_tir_op.continue_loop)
break_loop = _op_wrapper(_tir_op.break_loop)
round = _op_wrapper(_tir_op.round) # pylint: disable=redefined-builtin
@@ -3357,6 +3361,7 @@ __all__ = [
"attr",
"hint",
"While",
+ "Return",
"Break",
"Continue",
"If",
@@ -3432,7 +3437,6 @@ __all__ = [
"pow",
"q_multiply_shift",
"q_multiply_shift_per_axis",
- "ret",
"continue_loop",
"break_loop",
"reinterpret",
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index ed262825e2..076569a9f3 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -874,7 +874,7 @@ def visit_return(self: Parser, node: doc.Return) -> None:
value = self.eval_expr(node.value)
if value is None:
self.report_error(node, "Expression to be returned must be a Expr")
- T.evaluate(tvm.tirx.ret(value))
+ T.Return(value)
@dispatch.register(token="tirx", type_name="tvm_declare_function")
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index ba645985ea..1c46f4e011 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -860,6 +860,26 @@ class Break(Stmt):
self.__init_handle_by_constructor__(_ffi_api.Break, span) # type:
ignore
+@tvm_ffi.register_object("tirx.Return")
+class Return(Stmt):
+ """Return node.
+
+ Parameters
+ ----------
+ value : Expr
+ The value to return.
+
+ span : Optional[Span]
+ The location of this statement in the source code.
+ """
+
+ value: Expr
+ span: Span | None
+
+ def __init__(self, value: Expr, span: Span | None = None) -> None:
+ self.__init_handle_by_constructor__(_ffi_api.Return, value, span) #
type: ignore
+
+
@tvm_ffi.register_object("tirx.Continue")
class Continue(Stmt):
"""Continue node.
diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py
index 743518a196..2f85a0d7ae 100644
--- a/python/tvm/tirx/stmt_functor.py
+++ b/python/tvm/tirx/stmt_functor.py
@@ -38,6 +38,7 @@ class StmtFunctor:
"tirx.IfThenElse": self.visit_if_then_else_,
"tirx.For": self.visit_for_,
"tirx.While": self.visit_while_,
+ "tirx.Return": self.visit_return_,
"tirx.Break": self.visit_break_,
"tirx.Continue": self.visit_continue_,
"tirx.Allocate": self.visit_allocate_,
@@ -112,6 +113,10 @@ class StmtFunctor:
"""Visitor for While nodes."""
return self.visit_stmt_default_(op)
+ def visit_return_(self, op):
+ """Visitor for Return nodes."""
+ return self.visit_stmt_default_(op)
+
def visit_break_(self, op):
"""Visitor for Break nodes."""
return self.visit_stmt_default_(op)
@@ -253,6 +258,10 @@ class StmtVisitor(StmtFunctor):
self.visit_expr(op.condition)
self.visit_stmt(op.body)
+ def visit_return_(self, op):
+ """Visitor implementation for Return."""
+ self.visit_expr(op.value)
+
def visit_break_(self, op):
"""Visitor implementation for Break."""
pass
@@ -470,6 +479,15 @@ class StmtMutator(StmtFunctor):
return tvm.tirx.While(condition, body, op.span)
+ def visit_return_(self, op):
+ """Mutator implementation for Return."""
+ value = self.visit_expr(op.value)
+
+ if value is op.value:
+ return op
+
+ return tvm.tirx.Return(value, op.span)
+
def visit_break_(self, op):
"""Mutator implementation for Break."""
return op
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index 8248aa11aa..b6abbfa54b 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -97,7 +97,7 @@ tirx::PrimFunc
GetDLTensorField(tirx::builtin::TVMStructFieldKind field, PrimTyp
{tirx::Bind(value, tvm::Call(field_ty, tirx::builtin::tvm_struct_get(),
{dlpack_handle, IntImm::Int32(0),
IntImm::Int32(field)})
.as_or_throw<PrimExpr>()),
- tirx::Evaluate(tvm::ret(value.as_or_throw<PrimExpr>()))});
+ tirx::Return(value)});
DictAttrs attrs({{"tirx.is_scheduled", true}, {"tirx.is_host_func", true}});
@@ -282,7 +282,7 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const
Call& call) {
IntImm::Int32(tirx::builtin::TVMStructFieldKind::kDLTensorShape)})),
tirx::DeclBuffer(shape_buffer),
tirx::Bind(extent, tirx::BufferLoad(shape_buffer,
{axis.as_or_throw<PrimExpr>()})),
- tirx::Evaluate(tvm::ret(extent.as_or_throw<PrimExpr>()))});
+ tirx::Return(extent)});
DictAttrs attrs({{"tirx.is_scheduled", true}, {"tirx.is_host_func",
true}});
diff --git a/src/relax/transform/compute_prim_value.cc
b/src/relax/transform/compute_prim_value.cc
index 0d73745124..1e582f13a8 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -102,8 +102,7 @@ class PrimExprComputeInjector : public ExprMutator {
tvm::PrimType ret_ty = node.ty();
auto param_vars = tirx::UndefinedVars(node);
- tirx::Stmt body =
- tirx::Evaluate(tvm::Call(node.ty(), tirx::builtin::ret(),
{node}).as_or_throw<PrimExpr>());
+ tirx::Stmt body = tirx::Return(node);
tirx::PrimFunc func(param_vars, body, ret_ty, {},
DictAttrs({{tirx::attr::kIsHostFunc, true},
{tvm::attr::kSTir, true}}));
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 4e8c3beda5..97bd1b0f26 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -353,8 +353,7 @@ void CodeGenLLVM::AddFunctionInternal(const GlobalVar&
gvar, const PrimFunc& f)
EmitDebugLocation(f->span);
if (IsVoidType(f->ret_type)) {
- // All other return types are handled when encountering
- // builtin::ret().
+ // All other return types are handled when encountering Return.
builder_->CreateRetVoid();
} else {
builder_->CreateRet(ConstInt32(0));
@@ -1463,19 +1462,6 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const
CallNode* op) {
value->addIncoming(then_value, then_value_block);
value->addIncoming(else_value, else_value_block);
return value;
- } else if (op->op.same_as(builtin::ret())) {
- auto const* val = args[0].as<IntImmNode>();
- TVM_FFI_ICHECK(val) << "the tirx.ret should be transformed to return zero "
- << "before the llvm code generation.";
- TVM_FFI_ICHECK_EQ(val->value, 0) << "the tirx.ret should be transformed to
"
- << "return zero before the llvm code
generation.";
- builder_->CreateRet(ConstInt32(0));
- // LLVM allows exactly one terminator in a single basic block
- // append a new dummy basic block to avoid error.
- llvm::BasicBlock* ret_dummy =
- llvm::BasicBlock::Create(*llvm_target_->GetContext(), "ret_dummy",
function_);
- builder_->SetInsertPoint(ret_dummy);
- return ret_dummy;
} else if (op->op.same_as(builtin::continue_loop())) {
TVM_FFI_ICHECK(!loop_frame_jump_tgts_.empty())
<< "the tirx.continue_loop should be inserted under at least one For
or While stmts.";
@@ -2070,6 +2056,21 @@ void CodeGenLLVM::VisitStmt_(const WhileNode* op) {
builder_->SetInsertPoint(while_merge);
}
+void CodeGenLLVM::VisitStmt_(const ReturnNode* op) {
+ EmitDebugLocation(op);
+ auto const* val = op->value.as<IntImmNode>();
+ TVM_FFI_ICHECK(val) << "Return should be transformed to return zero "
+ << "before LLVM code generation.";
+ TVM_FFI_ICHECK_EQ(val->value, 0)
+ << "Return should be transformed to return zero before LLVM code
generation.";
+ builder_->CreateRet(ConstInt32(0));
+ // LLVM allows exactly one terminator in a basic block. Append a dummy block
+ // so code generation can continue after the return statement.
+ llvm::BasicBlock* ret_dummy =
+ llvm::BasicBlock::Create(*llvm_target_->GetContext(), "ret_dummy",
function_);
+ builder_->SetInsertPoint(ret_dummy);
+}
+
void CodeGenLLVM::VisitStmt_(const IfThenElseNode* op) {
EmitDebugLocation(op);
llvm::Value* cond = MakeValue(op->condition);
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 8e8e81a0d6..698e2571e9 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -230,6 +230,7 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const
Expr&)>,
void VisitStmt_(const BufferStoreNode* op) override;
void VisitStmt_(const ForNode* op) override;
void VisitStmt_(const WhileNode* op) override;
+ void VisitStmt_(const ReturnNode* op) override;
void VisitStmt_(const IfThenElseNode* op) override;
void VisitStmt_(const AllocBufferNode* op) override;
void VisitStmt_(const AttrStmtNode* op) override;
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index e422295f84..6f9c25c592 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -673,10 +673,7 @@ void CodeGenC::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLINT(*)
if (auto opt_call_op = op->op.as<Op>()) {
auto call_op = opt_call_op.value();
- if (op->op.same_as(builtin::ret())) {
- os << "return ";
- PrintExpr(op->args[0], os);
- } else if (op->op.same_as(builtin::continue_loop())) {
+ if (op->op.same_as(builtin::continue_loop())) {
os << "continue;";
} else if (op->op.same_as(builtin::break_loop())) {
os << "break;";
@@ -1348,6 +1345,13 @@ void CodeGenC::VisitStmt_(const WhileNode* op) {
stream << "}\n";
}
+void CodeGenC::VisitStmt_(const ReturnNode* op) {
+ PrintIndent();
+ stream << "return ";
+ PrintExpr(op->value, stream);
+ stream << ";\n";
+}
+
void CodeGenC::VisitStmt_(const BreakNode* op) {
PrintIndent();
stream << "break;\n";
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index a937a96a99..02fdc56391 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -197,6 +197,7 @@ class CodeGenC : public ExprFunctor<void(const Expr&,
std::ostream&)>,
void VisitStmt_(const BufferStoreNode* op) override;
void VisitStmt_(const ForNode* op) override;
void VisitStmt_(const WhileNode* op) override;
+ void VisitStmt_(const ReturnNode* op) override;
void VisitStmt_(const BreakNode* op) override;
void VisitStmt_(const ContinueNode* op) override;
void VisitStmt_(const IfThenElseNode* op) override;
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 30e002f4f0..e5fd278ece 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -48,6 +48,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IfThenElseNode::RegisterReflection();
ForNode::RegisterReflection();
WhileNode::RegisterReflection();
+ ReturnNode::RegisterReflection();
BreakNode::RegisterReflection();
ContinueNode::RegisterReflection();
BufferRegionNode::RegisterReflection();
@@ -241,6 +242,21 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+// Return
+Return::Return(Expr value, Span span) {
+ TVM_FFI_ICHECK(value.defined());
+
+ ffi::ObjectPtr<ReturnNode> node = ffi::make_object<ReturnNode>();
+ node->value = std::move(value);
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("tirx.Return", [](Expr value, Span span) { return
Return(value, span); });
+}
+
// Break
Break::Break(Span span) {
ffi::ObjectPtr<BreakNode> node = ffi::make_object<BreakNode>();
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index def6059c95..7ae935601f 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -59,6 +59,8 @@ void StmtVisitor::VisitStmt_(const WhileNode* op) {
this->VisitStmt(op->body);
}
+void StmtVisitor::VisitStmt_(const ReturnNode* op) {
this->VisitExpr(op->value); }
+
void StmtVisitor::VisitStmt_(const BreakNode* op) {}
void StmtVisitor::VisitStmt_(const ContinueNode* op) {}
@@ -344,6 +346,17 @@ Stmt StmtMutator::VisitStmt_(const WhileNode* op) {
}
}
+Stmt StmtMutator::VisitStmt_(const ReturnNode* op) {
+ Expr value = this->VisitExpr(op->value);
+ if (value.same_as(op->value)) {
+ return ffi::GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->value = std::move(value);
+ return Stmt(n);
+ }
+}
+
Stmt StmtMutator::VisitStmt_(const BreakNode* op) { return
ffi::GetRef<Stmt>(op); }
Stmt StmtMutator::VisitStmt_(const ContinueNode* op) { return
ffi::GetRef<Stmt>(op); }
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index e891027596..0d0df45f2b 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -212,6 +212,10 @@ void TIRVisitorWithPath::VisitStmt_(const WhileNode* op,
AccessPath path) {
bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
}
+void TIRVisitorWithPath::VisitStmt_(const ReturnNode* op, AccessPath path) {
+ Visit(op->value, path->Attr("value"));
+}
+
void TIRVisitorWithPath::VisitStmt_(const BreakNode* op, AccessPath path) {}
void TIRVisitorWithPath::VisitStmt_(const ContinueNode* op, AccessPath path) {}
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index 5d3f85b2f7..eb3974e4ee 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -129,6 +129,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
void VisitStmt_(const IfThenElseNode* op, ffi::reflection::AccessPath path)
override;
void VisitStmt_(const ForNode* op, ffi::reflection::AccessPath path)
override;
void VisitStmt_(const WhileNode* op, ffi::reflection::AccessPath path)
override;
+ void VisitStmt_(const ReturnNode* op, ffi::reflection::AccessPath path)
override;
void VisitStmt_(const BreakNode* op, ffi::reflection::AccessPath path)
override;
void VisitStmt_(const ContinueNode* op, ffi::reflection::AccessPath path)
override;
void VisitStmt_(const AllocBufferNode* op, ffi::reflection::AccessPath path)
override;
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 4a16e11139..a5d7b342ce 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -44,11 +44,6 @@ TIR_DEFINE_BUILTIN_FUNC(reinterpret)
static_cast<int64_t>(ScriptDtypePrintLocation::kFirst))
.set_num_inputs(1);
-TIR_DEFINE_BUILTIN_FUNC(ret)
- .set_attr<TCallEffectKind>("TCallEffectKind",
-
static_cast<int64_t>(CallEffectKind::kControlJump))
- .set_num_inputs(1);
-
TIR_DEFINE_BUILTIN_FUNC(thread_return)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kControlJump))
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index e23ea8e1bd..416a0085b8 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -288,19 +288,6 @@ void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Span
span) { // NOLINT(*)
}
}
-PrimExpr ret(PrimExpr value, Span span) {
- TVM_FFI_ICHECK(value.defined());
- return Call(value.ty(), tirx::builtin::ret(), {value}, {}, {},
span).as_or_throw<PrimExpr>();
-}
-
-Expr ret(Expr value, Span span) {
- TVM_FFI_ICHECK(value.defined());
- if (auto prim_value = value.as<PrimExpr>()) {
- return ret(prim_value.value(), span);
- }
- return Call(value->ty, tirx::builtin::ret(), {value}, {}, {}, span);
-}
-
PrimExpr thread_return(Span span) {
return Call(PrimType::Void(), tirx::builtin::thread_return(), {}, {}, {},
span)
.as_or_throw<PrimExpr>();
@@ -319,7 +306,6 @@ PrimExpr break_loop(Span span) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("tirx.ret", [](Expr value, Span span) { return ret(value, span); })
.def("tirx.thread_return", thread_return)
.def("tirx.continue_loop", continue_loop)
.def("tirx.break_loop", break_loop);
diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc
index a3af8da702..e9dc5c14da 100644
--- a/src/tirx/script/builder/ir.cc
+++ b/src/tirx/script/builder/ir.cc
@@ -698,6 +698,8 @@ WhileFrame While(PrimExpr condition) {
return WhileFrame(n);
}
+void Return(Expr value) { AddToParent(tvm::tirx::Return(std::move(value),
Span())); }
+
void Break() { AddToParent(tvm::tirx::Break(Span())); }
void Continue() { AddToParent(tvm::tirx::Continue(Span())); }
@@ -990,6 +992,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.Attr", Attr)
.def("script.ir_builder.tirx.DeviceEntry", DeviceEntry)
.def("script.ir_builder.tirx.While", While)
+ .def("script.ir_builder.tirx.Return", Return)
.def("script.ir_builder.tirx.Break", Break)
.def("script.ir_builder.tirx.Continue", Continue)
.def("script.ir_builder.tirx.If", If)
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 22e86d44e6..ffafa41494 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -66,20 +66,6 @@ bool IsAncestorOfAllVarUse(const tirx::Stmt& node, const
ffi::ObjectRef& var,
return false;
}
-ffi::Optional<Expr> FindReturnValue(const tirx::Stmt& node) {
- auto eval = node.as<tirx::EvaluateNode>();
- if (!eval) return std::nullopt;
-
- auto call = eval->value.as<CallNode>();
- if (!call) return std::nullopt;
-
- if (!call->op.same_as(tirx::builtin::ret())) return std::nullopt;
-
- if (call->args.size() != 1) return std::nullopt;
-
- return call->args[0];
-}
-
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tirx::TilePrimitiveCall>(
"", [](tirx::TilePrimitiveCall op_call, AccessPath p, IRDocsifier d)
-> Doc {
@@ -189,14 +175,6 @@ TVM_SCRIPT_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR);
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tirx::Evaluate>("", [](tirx::Evaluate eval, AccessPath p,
IRDocsifier d) -> Doc {
- if (d->cfg->syntax_sugar) {
- if (auto return_value = FindReturnValue(eval)) {
- ExprDoc value =
- d->AsDoc<ExprDoc>(return_value.value(),
p->Attr("value")->Attr("args")->ArrayItem(0));
- return ReturnDoc(value);
- }
- }
-
ExprDoc value = d->AsDoc<ExprDoc>(eval->value, p->Attr("value"));
if (eval->value->IsInstance<CallNode>()) {
return ExprStmtDoc(value);
@@ -204,6 +182,12 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return ExprStmtDoc(TIR(d, "evaluate")->Call({value}));
});
+TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
+ .set_dispatch<tirx::Return>("", [](tirx::Return stmt, AccessPath p,
IRDocsifier d) -> Doc {
+ ExprDoc value = d->AsDoc<ExprDoc>(stmt->value, p->Attr("value"));
+ return ReturnDoc(value);
+ });
+
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tirx::Bind>("", [](tirx::Bind stmt, AccessPath p,
IRDocsifier d) -> Doc {
// Step 1. Type annotation
@@ -894,6 +878,7 @@ TVM_SCRIPT_REPR(tirx::AttrStmtNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::AssertStmtNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::WhileNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::AllocBufferNode, ReprPrintTIR);
+TVM_SCRIPT_REPR(tirx::ReturnNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::BreakNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::ContinueNode, ReprPrintTIR);
TVM_SCRIPT_REPR(tirx::DeclBufferNode, ReprPrintTIR);
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 32dddb90d9..162ae78bab 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -56,18 +56,9 @@ class ReturnRewriter : public StmtMutator {
return ret;
}
- Stmt VisitStmt_(const EvaluateNode* node) override {
- Stmt ret = StmtMutator::VisitStmt_(node);
- const EvaluateNode* eval = ret.as<EvaluateNode>();
- TVM_FFI_ICHECK(eval);
- if (const CallNode* call = eval->value.as<CallNode>()) {
- if (call->op.same_as(builtin::ret())) {
- TVM_FFI_ICHECK_EQ(in_parallel_, 0) << "tirx.ret cannot be used in
parallel scope.";
- TVM_FFI_ICHECK_EQ(call->args.size(), 1) << "tirx.ret expect a single
argument.";
- ret = WriteToOut(call->args[0]);
- }
- }
- return ret;
+ Stmt VisitStmt_(const ReturnNode* node) override {
+ TVM_FFI_ICHECK_EQ(in_parallel_, 0) << "Return cannot be used in parallel
scope.";
+ return WriteToOut(this->VisitExpr(node->value));
}
private:
@@ -124,7 +115,7 @@ class ReturnRewriter : public StmtMutator {
{ret_var_, IntImm::Int32(0),
IntImm::Int32(tirx::builtin::kTVMFFIAnyUnionValue), info.expr})
.as_or_throw<PrimExpr>());
- Stmt ret_zero = Evaluate(tvm::ret(0));
+ Stmt ret_zero = Return(IntImm::Int32(0));
return SeqStmt({store_tindex, store_zero_padding, store_val, ret_zero});
}
@@ -275,7 +266,7 @@ PrimFunc MakePackedAPI(PrimFunc func) {
}
// Return error code of zero on success
- body = SeqStmt({body, Evaluate(ret(IntImm::Int32(0)))});
+ body = SeqStmt({body, Return(IntImm::Int32(0))});
body = MergeNest({std::move(result.init_nest), seq_check,
std::move(result.asserts),
std::move(result.decl_buffers)},
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index a54bc3e0fd..5bf8c77584 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -172,7 +172,7 @@ class HostDeviceSplitter : public StmtMutator {
Type kernel_ret_type = Type::Missing();
if (can_propagate_errors) {
kernel_ret_type = PrimType::Int(32);
- body = SeqStmt::Flatten(body, Evaluate(ret(success)));
+ body = SeqStmt::Flatten(body, Return(success));
} else {
kernel_ret_type = VoidType();
}
@@ -393,26 +393,11 @@ class ReturnRemover : public StmtExprMutator {
}
private:
- using Parent = StmtExprMutator;
- Stmt VisitStmt_(const EvaluateNode* op) override {
- if (auto* call = op->value.as<CallNode>()) {
- if (call->op.same_as(builtin::ret())) {
- TVM_FFI_ICHECK_EQ(call->args.size(), 1);
- auto as_int = call->args[0].as<IntImmNode>();
- TVM_FFI_ICHECK(as_int && as_int->value == 0)
- << "Device kernel may only contain successful return, T.ret(0)";
- return Evaluate(0);
- }
- }
- return Parent::VisitStmt_(op);
- }
-
- Expr VisitExpr_(const CallNode* op) override {
- if (op->op.same_as(builtin::ret())) {
- TVM_FFI_THROW(InternalError)
- << "Call to builtin::ret() should only appear within an Evaluate
node";
- }
- return Parent::VisitExpr_(op);
+ Stmt VisitStmt_(const ReturnNode* op) override {
+ auto as_int = op->value.as<IntImmNode>();
+ TVM_FFI_ICHECK(as_int && as_int->value == 0)
+ << "Device kernel may only contain a successful return, return 0";
+ return Evaluate(0);
}
};
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 0b0697446b..650021d125 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -179,6 +179,17 @@ class LinearAccessPatternFinder final : public
StmtExprVisitor {
}
}
+ void VisitStmt_(const ReturnNode* op) final {
+ scope_.push_back(StmtEntry());
+ StmtExprVisitor::VisitStmt_(op);
+ StmtEntry e = scope_.back();
+ scope_.pop_back();
+ if (e.touched.size() != 0) {
+ e.stmt = op;
+ linear_seq_.push_back(e);
+ }
+ }
+
void VisitExpr_(const VarNode* buf) final {
// Directly reference to the variable count as a read.
auto it = alloc_info_.find(buf);
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index 8bb32f68a4..7fbcb8935a 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -1096,7 +1096,7 @@ def test_device_host_call_same_func():
C[bx, tx] = Module.add(A[bx, tx], B[bx, tx]) # Call from
device
# 1. If we set host to llvm, it will raise an error of
- # "the tirx.ret should be transformed to return zero before the llvm
code generation."
+ # "Return should be transformed to return zero before LLVM code
generation."
# Need to revisit this.
# 2. We set a dummy mcpu value for testing purpose,
# in order to avoid checking a function is host or device based on the
"cpu" substring.
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index 9de6b73bb1..64f77c4acc 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -941,15 +941,15 @@ def test_llvm_order_functions():
class Module:
@T.prim_func(s_tir=True)
def Danny(v: T.float32) -> T.float32:
- T.ret(T.call_extern("float32", "Dave", v))
+ return T.call_extern("float32", "Dave", v)
@T.prim_func(s_tir=True)
def Sammy(v: T.float32) -> T.float32:
- T.ret(T.call_extern("float32", "Eve", v))
+ return T.call_extern("float32", "Eve", v)
@T.prim_func(s_tir=True)
def Kirby(v: T.float32) -> T.float32:
- T.ret(T.call_extern("float32", "Fred", v))
+ return T.call_extern("float32", "Fred", v)
ir_text = tvm.tirx.build(Module, target="llvm").inspect_source("ll")
# Skip functions whose names start with _.
diff --git a/tests/python/relax/test_transform_compute_prim_value.py
b/tests/python/relax/test_transform_compute_prim_value.py
index 37a4fafcb1..db5e1e7f77 100644
--- a/tests/python/relax/test_transform_compute_prim_value.py
+++ b/tests/python/relax/test_transform_compute_prim_value.py
@@ -43,7 +43,7 @@ def test_prim_value_in_assert_condition():
@T.prim_func(private=True, s_tir=True)
def compute_symbolic_expr(N: T.int64) -> T.bool:
T.func_attr({"tirx.is_host_func": True})
- T.ret(N % 16 == 0)
+ return N % 16 == 0
After = tvm.relax.transform.ComputePrimValue()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -76,7 +76,7 @@ def test_prim_value_in_branch_condition():
@T.prim_func(private=True, s_tir=True)
def compute_symbolic_expr(N: T.int64) -> T.bool:
T.func_attr({"tirx.is_host_func": True})
- T.ret(N % 16 == 0)
+ return N % 16 == 0
After = tvm.relax.transform.ComputePrimValue()(Before)
tvm.ir.assert_structural_equal(After, Expected)
diff --git
a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
index a10ad96750..23f624c653 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
@@ -48,19 +48,19 @@ class TestNoOp(CheckPureFunction):
class TestReturnValue(CheckPureFunction):
@T.prim_func(s_tir=True)
def func() -> T.int32:
- T.ret(42)
+ return 42
class TestComputeValueAndReturn(CheckPureFunction):
@T.prim_func(s_tir=True)
def func(N: T.int32, M: T.int32) -> T.int32:
- T.ret(N * M)
+ return N * M
class TestReadBufferArgument(CheckPureFunction):
@T.prim_func(s_tir=True)
def func(A: T.Buffer(16, "float32")) -> T.float32:
- T.ret(A[0])
+ return A[0]
class TestWriteToBufferArgument(CheckImpureFunction):
@@ -78,13 +78,13 @@ class TestWriteToInternalAllocation(CheckPureFunction):
for i, j in T.grid(16, 16):
Sum[()] = Sum[()] + A[i, j]
- T.ret(Sum[()])
+ return Sum[()]
class TestCallPureBuiltin(CheckPureFunction):
@T.prim_func(s_tir=True)
def func(x: T.float32) -> T.float32:
- T.ret(T.cos(x))
+ return T.cos(x)
class TestCallPureExtern(CheckPureFunction):
diff --git a/tests/python/tirx-base/test_tir_base.py
b/tests/python/tirx-base/test_tir_base.py
index 477062c497..62af878e00 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx-base/test_tir_base.py
@@ -19,6 +19,7 @@ import itertools
import numpy as np
import pytest
+import tvm_ffi
import tvm
from tvm import tirx
@@ -50,8 +51,7 @@ def test_scalar_add():
lhs = tirx.Cast(lhs_type, lhs_input)
rhs = tirx.Cast(rhs_type, rhs_input)
output = lhs + rhs
- output = tirx.ret(output)
- output = tirx.Evaluate(output)
+ output = tirx.Return(output)
func = tirx.PrimFunc([lhs_input, rhs_input], output)
func = build_tir_func(func)
out = func(1.0, 2.0)
@@ -92,22 +92,54 @@ def test_cast_between_types():
assignment_helper(store_dtype, value_dtype)
-def test_ret_const():
+def test_return_const():
a = tirx.const(0)
- b = tirx.ret(a)
- b = tirx.Evaluate(b)
+ b = tirx.Return(a)
func = tirx.PrimFunc([], b)
func = build_tir_func(func)
out = func()
assert out == 0
+def test_return_accepts_expr_and_roundtrips():
+ value = tvm.relax.ShapeExpr([2, 3])
+ span = tvm.ir.Span(tvm.ir.SourceName("return_test"), 1, 1, 1, 9)
+ stmt = tirx.Return(value, span)
+
+ assert stmt.value.same_as(value)
+ assert stmt.span.same_as(span)
+ assert not tvm.ir.is_prim_expr(stmt.value)
+
+ restored = tvm.ir.load_json(tvm.ir.save_json(stmt))
+ tvm.ir.assert_structural_equal(restored, stmt)
+ assert tvm_ffi.structural_hash(restored) == tvm_ffi.structural_hash(stmt)
+
+ with pytest.raises(tvm.error.InternalError):
+ tirx.Return(None)
+
+
+def test_return_stmt_functor_traversal_and_mutation():
+ x = tirx.Var("x", "int32")
+ span = tvm.ir.Span(tvm.ir.SourceName("return_test"), 1, 1, 1, 9)
+ stmt = tirx.Return(x + 1, span)
+ visited = []
+
+ tirx.stmt_functor.post_order_visit(stmt, visited.append)
+ assert any(node.same_as(x) for node in visited)
+ assert any(isinstance(node, tirx.Return) for node in visited)
+
+ rewritten = tirx.stmt_functor.substitute(stmt, {x: tirx.IntImm("int32",
4)})
+ expected = tirx.Return(tirx.Add(tirx.IntImm("int32", 4),
tirx.IntImm("int32", 1)), span)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+ assert rewritten.span.same_as(span)
+
+
def test_control_flow_jump():
@T.prim_func(s_tir=True)
def func(a: T.float32, b: T.float32):
if True:
- T.evaluate(T.ret(a))
- T.evaluate(T.ret(b))
+ return a
+ return b
func = build_tir_func(func)
out = func(1.0, 2.0)
@@ -191,7 +223,7 @@ def test_eq_ops():
if __name__ == "__main__":
test_scalar_add()
- test_ret_const()
+ test_return_const()
test_control_flow_jump()
test_exception()
test_eq_ops()
diff --git a/tests/python/tirx-base/test_tir_imm_values.py
b/tests/python/tirx-base/test_tir_imm_values.py
index d2bd44f811..2d9048f7d1 100644
--- a/tests/python/tirx-base/test_tir_imm_values.py
+++ b/tests/python/tirx-base/test_tir_imm_values.py
@@ -153,7 +153,7 @@ def test_tir_too_large_literal_f64():
# object is still constructed, and eval to infinity.
@T.prim_func(s_tir=True)
def imm_overflow_fp64() -> T.float64:
- T.evaluate(T.ret(T.float64(1.7976e309), dtype="float64"))
+ return T.float64(1.7976e309)
f = tvm.compile(imm_overflow_fp64, target="llvm")
assert math.isinf(f())
@@ -321,23 +321,23 @@ def test_tir_int8_const_fold():
@T.prim_func(s_tir=True)
def imm_multiply(x: T.int8, y: T.int8) -> T.int8:
- T.evaluate(T.ret(x * y, dtype="int8"))
+ return x * y
@T.prim_func(s_tir=True)
def imm_add(x: T.int8, y: T.int8) -> T.int8:
- T.evaluate(T.ret(x + y, dtype="int8"))
+ return x + y
@T.prim_func(s_tir=True)
def imm_sub(x: T.int8, y: T.int8) -> T.int8:
- T.evaluate(T.ret(x - y, dtype="int8"))
+ return x - y
@T.prim_func(s_tir=True)
def imm_truncdiv(x: T.int8, y: T.int8) -> T.int8:
- T.evaluate(T.ret(T.truncdiv(x, y), dtype="int8"))
+ return T.truncdiv(x, y)
@T.prim_func(s_tir=True)
def imm_floordiv(x: T.int8, y: T.int8) -> T.int8:
- T.evaluate(T.ret(T.floordiv(x, y), dtype="int8"))
+ return T.floordiv(x, y)
fmul = tvm.compile(imm_multiply, target="llvm")
fadd = tvm.compile(imm_add, target="llvm")
@@ -377,23 +377,23 @@ def test_tir_uint8_const_fold():
@T.prim_func(s_tir=True)
def imm_multiply(x: T.uint8, y: T.uint8) -> T.uint8:
- T.evaluate(T.ret(x * y, dtype="uint8"))
+ return x * y
@T.prim_func(s_tir=True)
def imm_add(x: T.uint8, y: T.uint8) -> T.uint8:
- T.evaluate(T.ret(x + y, dtype="uint8"))
+ return x + y
@T.prim_func(s_tir=True)
def imm_sub(x: T.uint8, y: T.uint8) -> T.uint8:
- T.evaluate(T.ret(x - y, dtype="uint8"))
+ return x - y
@T.prim_func(s_tir=True)
def imm_truncdiv(x: T.uint8, y: T.uint8) -> T.uint8:
- T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint8"))
+ return T.truncdiv(x, y)
@T.prim_func(s_tir=True)
def imm_floordiv(x: T.uint8, y: T.uint8) -> T.uint8:
- T.evaluate(T.ret(T.floordiv(x, y), dtype="uint8"))
+ return T.floordiv(x, y)
fmul = tvm.compile(imm_multiply, target="llvm")
fadd = tvm.compile(imm_add, target="llvm")
@@ -442,31 +442,31 @@ def test_tir_int32_const_fold():
@T.prim_func(s_tir=True)
def imm_multiply(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(x * y, dtype="int32"))
+ return x * y
@T.prim_func(s_tir=True)
def imm_add(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(x + y, dtype="int32"))
+ return x + y
@T.prim_func(s_tir=True)
def imm_sub(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(x - y, dtype="int32"))
+ return x - y
@T.prim_func(s_tir=True)
def imm_truncdiv(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(T.truncdiv(x, y), dtype="int32"))
+ return T.truncdiv(x, y)
@T.prim_func(s_tir=True)
def imm_truncmod(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(T.truncmod(x, y), dtype="int32"))
+ return T.truncmod(x, y)
@T.prim_func(s_tir=True)
def imm_floordiv(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(T.floordiv(x, y), dtype="int32"))
+ return T.floordiv(x, y)
@T.prim_func(s_tir=True)
def imm_floormod(x: T.int32, y: T.int32) -> T.int32:
- T.evaluate(T.ret(T.floormod(x, y), dtype="int32"))
+ return T.floormod(x, y)
fmul = tvm.compile(imm_multiply, target="llvm")
fadd = tvm.compile(imm_add, target="llvm")
@@ -530,23 +530,23 @@ def test_tir_uint32_const_fold():
@T.prim_func(s_tir=True)
def imm_multiply(x: T.uint32, y: T.uint32) -> T.uint32:
- T.evaluate(T.ret(x * y, dtype="uint32"))
+ return x * y
@T.prim_func(s_tir=True)
def imm_add(x: T.uint32, y: T.uint32) -> T.uint32:
- T.evaluate(T.ret(x + y, dtype="uint32"))
+ return x + y
@T.prim_func(s_tir=True)
def imm_sub(x: T.uint32, y: T.uint32) -> T.uint32:
- T.evaluate(T.ret(x - y, dtype="uint32"))
+ return x - y
@T.prim_func(s_tir=True)
def imm_truncdiv(x: T.uint32, y: T.uint32) -> T.uint32:
- T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint32"))
+ return T.truncdiv(x, y)
@T.prim_func(s_tir=True)
def imm_floordiv(x: T.uint32, y: T.uint32) -> T.uint32:
- T.evaluate(T.ret(T.floordiv(x, y), dtype="uint32"))
+ return T.floordiv(x, y)
fmul = tvm.compile(imm_multiply, target="llvm")
fadd = tvm.compile(imm_add, target="llvm")
diff --git a/tests/python/tirx-base/test_tir_specialize.py
b/tests/python/tirx-base/test_tir_specialize.py
index c29f26d389..a1428348d2 100644
--- a/tests/python/tirx-base/test_tir_specialize.py
+++ b/tests/python/tirx-base/test_tir_specialize.py
@@ -341,11 +341,11 @@ def test_specialization_updates_ty():
@T.prim_func(private=True, s_tir=True)
def before(n: T.int32) -> T.int32:
- T.ret(n * 10)
+ return n * 10
@T.prim_func(private=True, s_tir=True)
def expected() -> T.int32:
- T.ret(50)
+ return 50
ty_before = tvm.relax.FuncType([tvm.ir.PrimType("int32")],
tvm.ir.PrimType("int32"))
tvm.ir.assert_structural_equal(before.ty, ty_before)
diff --git a/tests/python/tirx-transform/test_tir_inline_private_functions.py
b/tests/python/tirx-transform/test_tir_inline_private_functions.py
index 13a4a05216..fc6c060e9b 100644
--- a/tests/python/tirx-transform/test_tir_inline_private_functions.py
+++ b/tests/python/tirx-transform/test_tir_inline_private_functions.py
@@ -195,7 +195,7 @@ class TestInlineCallOccurringInExpression(BaseTestCase):
cos = T.cos(T.cast(i, "float32"))
sin = T.sin(T.cast(i, "float32"))
retval = cos * cos + sin * sin
- T.ret(retval)
+ return retval
@I.ir_module(s_tir=True)
class Expected:
diff --git a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
index 7ac26d37ac..4b69f93335 100644
--- a/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
+++ b/tests/python/tirx-transform/test_tir_transform_make_packed_api.py
@@ -219,6 +219,19 @@ def test_pointer_return():
assert 4 in return_type_indices # ffi::TypeIndex::kTVMFFIOpaquePtr
+def test_return_from_parallel_scope_is_rejected():
+ """A parallel loop cannot return from its enclosing function."""
+
+ i = tirx.Var("i", "int32")
+ body = tirx.For(i, 0, 1, tirx.ForKind.PARALLEL, tirx.Return(i))
+ func = tirx.PrimFunc([], body, tvm.ir.PrimType("int32"))
+ func = func.with_attr("global_symbol", "main")
+ func = func.with_attr("target", tvm.target.Target("llvm", host="llvm"))
+
+ with pytest.raises(tvm.error.InternalError, match="Return cannot be used
in parallel scope"):
+ tvm.tirx.transform.MakePackedAPI()(tvm.IRModule({"main": func}))
+
+
def test_int_parameter():
"""Int parameter emits type check accepting int or bool."""
diff --git
a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
index 79c0829177..8a026c7a8a 100644
--- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
+++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import pytest
+
import tvm
import tvm.testing
from tvm.script import ir as I
@@ -116,12 +118,31 @@ def test_split_host_device_on_cpu():
}
)
T.evaluate(n)
- T.ret(0)
+ return 0
After = tvm.tirx.transform.SplitHostDevice()(Before)
tvm.ir.assert_structural_equal(After, Expected)
+def test_device_kernel_nonzero_return_is_rejected():
+ """A device kernel may only return the zero success code."""
+
+ device_target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+ target = tvm.target.Target(device_target, host="llvm")
+ body = tvm.tirx.AttrStmt(
+ device_target,
+ "target",
+ 0,
+ tvm.tirx.Return(tvm.tirx.IntImm("int32", 1)),
+ )
+ func = tvm.tirx.PrimFunc([], body)
+ func = func.with_attr("global_symbol", "main")
+ func = func.with_attr("target", target)
+
+ with pytest.raises(tvm.error.InternalError, match="successful return"):
+ tvm.tirx.transform.SplitHostDevice()(tvm.IRModule({"main": func}))
+
+
def test_split_host_device_without_func_host_attribute():
"""Like test_split_host_device, but no host specified in the host's target
diff --git a/tests/python/tirx/transform/test_stmt_functor.py
b/tests/python/tirx/transform/test_stmt_functor.py
index aee9a78254..393cfa47a0 100644
--- a/tests/python/tirx/transform/test_stmt_functor.py
+++ b/tests/python/tirx/transform/test_stmt_functor.py
@@ -98,6 +98,12 @@ class ASTPrinter(StmtVisitor):
self.visit_stmt(op.body)
self.log.pop_scope()
+ def visit_return_(self, op):
+ self.log.add("Return")
+ self.log.push_scope()
+ self.visit_expr(op.value)
+ self.log.pop_scope()
+
def visit_buffer_store_(self, op):
self.log.add("BufferStore")
self.log.push_scope()
@@ -256,6 +262,11 @@ class ASTPrinterMutator(StmtMutator):
self.log.add("While")
return result
+ def visit_return_(self, op):
+ result = super().visit_return_(op)
+ self.log.add("Return")
+ return result
+
def visit_buffer_store_(self, op):
result = super().visit_buffer_store_(op)
self.log.add("BufferStore")
@@ -641,6 +652,9 @@ def create_test_statements():
# While loop
while_loop = tir.While(tir.LT(x, int_imm), evaluate_stmt)
+ # Return
+ return_stmt = tir.Return(add_expr)
+
# Buffer operations
buffer_var = tir.Var("buf", "handle")
buffer = tir.decl_buffer((10,), "int32", buffer_var.name)
@@ -685,6 +699,7 @@ def create_test_statements():
"let": let_stmt,
"for": for_loop,
"while": while_loop,
+ "return": return_stmt,
"buffer_store": buffer_store,
"seq_stmt": seq_stmt,
"block_realize": block_realize,
@@ -758,6 +773,16 @@ def test_while():
)
+def test_return():
+ """Test return statement."""
+ return_stmt = create_test_statements()["return"]
+ basic_check(
+ return_stmt,
+ "\n".join(["Return", "\tAdd", "\t\tVar", "\t\tIntImm"]),
+ "\n".join(["Var", "IntImm", "Add", "Return"]),
+ )
+
+
def test_buffer_store():
"""Test buffer store statement."""
buffer_store = create_test_statements()["buffer_store"]
diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py
b/tests/python/tvmscript/test_tvmscript_parser_tir.py
index fd1f6a7901..157e59f59e 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py
@@ -407,7 +407,7 @@ def test_inferred_ty_with_prim_args():
@T.prim_func(s_tir=True)
def func(M: T.int32, N: T.int32) -> T.int32:
- T.ret(M * N)
+ return M * N
expected = tvm.relax.FuncType(
[
@@ -425,7 +425,7 @@ def test_inferred_ty_with_buffer_args():
@T.prim_func(s_tir=True)
def func(A: T.Buffer([16, 16], "float32"), B: T.Buffer([256], "int32")) ->
T.float32:
- T.ret(T.float32(42.0))
+ return T.float32(42.0)
expected = tvm.relax.FuncType(
[
@@ -452,7 +452,7 @@ def test_inferred_ty_with_internal_allocation():
for i, j in T.grid(16, 16):
Sum[()] = Sum[()] + A[i, j]
- T.ret(Sum[()])
+ return Sum[()]
expected = tvm.relax.FuncType(
[
diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py
b/tests/python/tvmscript/test_tvmscript_printer_tir.py
index a6c822eedd..c29b99eaa3 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py
@@ -910,7 +910,7 @@ def test_return_statement():
@T.prim_func(s_tir=True)
def func():
- T.evaluate(T.ret(5))
+ return T.int32(5)
expected_output = """
# from tvm.script import tirx as T
@@ -921,6 +921,7 @@ def func():
return 5
"""
_assert_print(func, expected_output)
+ assert func.script(verbose_expr=True, syntax_sugar=False).strip() ==
expected_output.strip()
CUSTOM_FLOAT_DTYPES = [
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py
b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index 08bf42decf..71aedaecbe 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -3013,7 +3013,7 @@ def subroutine_call_returning_int():
@T.prim_func(s_tir=True)
def subroutine(x: T.float32) -> T.float32:
- T.ret(x * x)
+ return x * x
return mod
@@ -3091,7 +3091,7 @@ def subroutine_call_without_arguments():
def return_zero():
@T.prim_func(s_tir=True)
def func() -> T.int32:
- T.ret(0)
+ return 0
return func
@@ -3099,7 +3099,7 @@ def return_zero():
def return_zero_private():
@T.prim_func(private=True, s_tir=True)
def func() -> T.int32:
- T.ret(0)
+ return 0
return func
@@ -3108,7 +3108,7 @@ def return_zero_private_with_attr():
@T.prim_func(private=True, s_tir=True)
def func() -> T.int32:
T.func_attr({"greeting": "hello"})
- T.ret(0)
+ return 0
return func
diff --git a/tests/python/tvmscript/test_tvmscript_syntax_sugar.py
b/tests/python/tvmscript/test_tvmscript_syntax_sugar.py
index 6150439190..57354f72db 100644
--- a/tests/python/tvmscript/test_tvmscript_syntax_sugar.py
+++ b/tests/python/tvmscript/test_tvmscript_syntax_sugar.py
@@ -506,11 +506,11 @@ def test_foldable_boolean_in_assert():
def test_return_statement():
- """A python `return` statement uses `T.ret`"""
+ """A Python `return` statement creates a first-class Return node."""
@T.prim_func(s_tir=True)
def explicit():
- T.evaluate(T.ret(5))
+ T.Return(T.int32(5))
@T.prim_func(s_tir=True)
def implicit():
