This is an automated email from the ASF dual-hosted git repository.
csullivan 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 0c09547c76 [TIR][CodeGen] Define PackedFunc error code in
MakePackedAPI (#15076)
0c09547c76 is described below
commit 0c09547c7610bd1bfde3e53d1ac87baad11b6026
Author: Eric Lunderberg <[email protected]>
AuthorDate: Thu Jun 15 14:50:51 2023 -0400
[TIR][CodeGen] Define PackedFunc error code in MakePackedAPI (#15076)
* [TIR][CodeGen] Define PackedFunc error code in MakePackedAPI
Previously, the return value of a PackedFunc was hard-coded as the
string `"return 0;"` in `CodeGenCHost`, which could cause compilation
errors for `PrimFunc` returning `DataType::Void()`. This PR removes
this explicit return statement from `CodeGenCHost`, replacing it with
`tir::ret(Integer(0))` in the `MakePackedAPI` and `MakeUnpackedAPI`
transforms.
This is related to https://github.com/apache/tvm/pull/15073, which
performs an analogous change for the function signature.
* Handle builtin::ret() in CodeGenC
* Place T.ret(0) inside asserts, rather than outside
This causes fewer unit tests to break, and has more readable
TVMScript.
* Update unit tests to look inside SeqStmt
* Handle T.ret(0) in CodeGenStackVM
* Update MakeUnpackedAPI tests to expect T.ret
---
src/target/source/codegen_c.cc | 6 +++---
src/target/source/codegen_c.h | 4 ----
src/target/source/codegen_c_host.cc | 5 -----
src/target/source/codegen_c_host.h | 1 -
src/target/stackvm/codegen_stackvm.cc | 6 ++++++
src/tir/transforms/make_packed_api.cc | 9 +++++++--
src/tir/transforms/make_unpacked_api.cc | 4 +++-
.../unittest/test_tir_transform_lower_tvm_builtin.py | 9 +++++++--
.../python/unittest/test_tir_transform_make_packed_api.py | 15 ++++++++++++---
.../unittest/test_tir_transform_make_unpacked_api.py | 6 ++++++
10 files changed, 44 insertions(+), 21 deletions(-)
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index bcdd0bfea0..a7cc320562 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -123,7 +123,6 @@ void CodeGenC::AddFunction(const PrimFunc& f) {
this->PreFunctionBody(f);
int func_scope = this->BeginScope();
this->PrintStmt(f->body);
- this->PrintFinalReturn();
this->EndScope(func_scope);
this->PrintIndent();
this->stream << "}\n\n";
@@ -133,8 +132,6 @@ void CodeGenC::PrintFuncPrefix(std::ostream& os) {}
void CodeGenC::PrintExtraAttrs(const PrimFunc& f) {}
-void CodeGenC::PrintFinalReturn() {}
-
std::string CodeGenC::Finish() { return decl_stream.str() + stream.str(); }
void CodeGenC::PrintExpr(const PrimExpr& n, std::ostream& os) { // NOLINT(*)
@@ -538,6 +535,9 @@ void CodeGenC::VisitExpr_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
PrintExpr(op->args[0], os);
os << " ) return ";
PrintExpr(op->args[1], os);
+ } else if (op->op.same_as(builtin::ret())) {
+ os << "return ";
+ PrintExpr(op->args[0], os);
} else if (op->op.same_as(builtin_call_extern_) ||
op->op.same_as(builtin_call_pure_extern_)) {
ICHECK_GE(op->args.size(), 1U);
auto func = Downcast<StringImm>(op->args[0]);
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index de9c2f1745..93f9ea519c 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -110,10 +110,6 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&,
std::ostream&)>,
* Example: __launch_bounds__(256) for CUDA functions
*/
virtual void PrintExtraAttrs(const PrimFunc& f);
- /*!
- * \brief Print the final return at the end the function.
- */
- virtual void PrintFinalReturn(); // NOLINT(*)
/*!
* \brief Insert statement before function body.
* \param f The function to be compiled.
diff --git a/src/target/source/codegen_c_host.cc
b/src/target/source/codegen_c_host.cc
index e98852c270..3255e11c5d 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -128,11 +128,6 @@ void CodeGenCHost::PrintFuncPrefix(std::ostream& os) { //
NOLINT(*)
<< "TVM_DLL ";
}
-void CodeGenCHost::PrintFinalReturn() { // NOLINT(*)
- this->PrintIndent();
- stream << "return 0;\n";
-}
-
std::string CodeGenCHost::Finish() { // NOLINT(*)
std::string ret = decl_stream.str();
if (emit_fwd_func_decl_) {
diff --git a/src/target/source/codegen_c_host.h
b/src/target/source/codegen_c_host.h
index 9c71f197f0..c4db634705 100644
--- a/src/target/source/codegen_c_host.h
+++ b/src/target/source/codegen_c_host.h
@@ -58,7 +58,6 @@ class CodeGenCHost : public CodeGenC {
using CodeGenC::PrintType;
void PrintType(DataType t, std::ostream& os) final; // NOLINT(*)
void PrintFuncPrefix(std::ostream& os) final; // NOLINT(*)
- void PrintFinalReturn() final; // NOLINT(*)
// overload visitor functions
void VisitExpr_(const BroadcastNode* op, std::ostream& os) final; //
NOLINT(*)
diff --git a/src/target/stackvm/codegen_stackvm.cc
b/src/target/stackvm/codegen_stackvm.cc
index db6e32d65f..fa2cd6b09d 100644
--- a/src/target/stackvm/codegen_stackvm.cc
+++ b/src/target/stackvm/codegen_stackvm.cc
@@ -284,6 +284,12 @@ void CodeGenStackVM::VisitExpr_(const CallNode* op) {
this->Push(op->args[0]);
this->PushOp(StackVM::PUSH_I64, 0);
this->PushOp(StackVM::EQ_HANDLE);
+ } else if (op->op.same_as(builtin::ret())) {
+ CHECK(op->args.size() == 1 && op->args[0]->IsInstance<IntImmNode>() &&
+ op->args[0].as<IntImmNode>()->value == 0)
+ << "StackVM does not support return values, "
+ << "and the return value " << op->args
+ << " is not special case of returning an error code of zero.";
} else {
LOG(FATAL) << "unknown function call " << op->op;
}
diff --git a/src/tir/transforms/make_packed_api.cc
b/src/tir/transforms/make_packed_api.cc
index a6673a19ad..e387204045 100644
--- a/src/tir/transforms/make_packed_api.cc
+++ b/src/tir/transforms/make_packed_api.cc
@@ -353,11 +353,16 @@ PrimFunc MakePackedAPI(PrimFunc func) {
}
}
+ // Return error code of zero on success
+ body = SeqStmt({body, Evaluate(ret(Integer(0)))});
+
+ // Apply all argument assertions
std::ostringstream num_args_error;
num_args_error << name_hint << ": num_args should be " << num_args;
std::vector<Stmt> arg_assert = {MakeAssertEQ(v_num_packed_args, num_args,
num_args_error.str())};
- func_ptr->body =
- MergeNest({arg_assert, seq_init, binder.init_nest(), seq_check,
binder.asserts()}, body);
+ body = MergeNest({arg_assert, seq_init, binder.init_nest(), seq_check,
binder.asserts()}, body);
+
+ func_ptr->body = body;
func_ptr->params = args;
Array<Var> undefined = UndefinedVars(func_ptr->body, func_ptr->params);
diff --git a/src/tir/transforms/make_unpacked_api.cc
b/src/tir/transforms/make_unpacked_api.cc
index 4b1b3bf517..2646b5baea 100644
--- a/src/tir/transforms/make_unpacked_api.cc
+++ b/src/tir/transforms/make_unpacked_api.cc
@@ -147,7 +147,9 @@ PrimFunc MakeUnpackedAPI(PrimFunc func) {
device_init.push_back(AttrStmt(node, attr::device_type, device_type, nop));
}
- func_ptr->body = MergeNest(device_init, func_ptr->body);
+ Stmt body = MergeNest(device_init, SeqStmt({func_ptr->body,
Evaluate(ret(Integer(0)))}));
+
+ func_ptr->body = body;
func_ptr->params = args;
func_ptr->ret_type = PrimType(DataType::Int(32));
func_ptr->buffer_map = Map<Var, Buffer>();
diff --git a/tests/python/unittest/test_tir_transform_lower_tvm_builtin.py
b/tests/python/unittest/test_tir_transform_lower_tvm_builtin.py
index 2e0784cc31..6eac5e90b5 100644
--- a/tests/python/unittest/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/unittest/test_tir_transform_lower_tvm_builtin.py
@@ -70,8 +70,13 @@ def check_packed_func(target="llvm"):
node = prim_func.body
# Recursively visit PrimFunc until we meet the for-loop:
- while isinstance(node, (tvm.tir.AssertStmt, tvm.tir.LetStmt,
tvm.tir.AttrStmt)):
- node = node.body
+ while True:
+ if isinstance(node, (tvm.tir.AssertStmt, tvm.tir.LetStmt,
tvm.tir.AttrStmt)):
+ node = node.body
+ elif isinstance(node, tvm.tir.SeqStmt):
+ node = node[0]
+ else:
+ break
# For-loop:
assert isinstance(node, tvm.tir.stmt.For)
diff --git a/tests/python/unittest/test_tir_transform_make_packed_api.py
b/tests/python/unittest/test_tir_transform_make_packed_api.py
index 34adcbb9ae..6f84b6f6d4 100644
--- a/tests/python/unittest/test_tir_transform_make_packed_api.py
+++ b/tests/python/unittest/test_tir_transform_make_packed_api.py
@@ -60,9 +60,18 @@ def _find_assignment(stmt, var_name):
def _find_next(stmt, type):
- while not isinstance(stmt, type):
- stmt = stmt.body
- return stmt
+ search_stack = [stmt]
+
+ while search_stack:
+ stmt = search_stack.pop()
+ if isinstance(stmt, type):
+ return stmt
+ elif isinstance(stmt, tvm.tir.SeqStmt):
+ search_stack.extend(reversed(stmt))
+ else:
+ search_stack.append(stmt.body)
+
+ return None
def _find_compute_scope(func):
diff --git a/tests/python/unittest/test_tir_transform_make_unpacked_api.py
b/tests/python/unittest/test_tir_transform_make_unpacked_api.py
index 1931f7aef3..868d30db36 100644
--- a/tests/python/unittest/test_tir_transform_make_unpacked_api.py
+++ b/tests/python/unittest/test_tir_transform_make_unpacked_api.py
@@ -172,6 +172,7 @@ class TestTargetHostRemoved(tvm.testing.CompareBeforeAfter):
T.attr("default", "device_id", 0)
T.attr("default", "device_type", 2)
mod.subroutine(A_data)
+ T.ret(T.int32(0))
@T.prim_func
def subroutine(A_data: T.handle("float32")):
@@ -215,6 +216,7 @@ class
TestInternalSubroutineCall(tvm.testing.CompareBeforeAfter):
T.attr("default", "device_id", 0)
T.attr("default", "device_type", 1)
mod.subroutine(A_data)
+ T.ret(T.int32(0))
@T.prim_func
def subroutine(A_data: T.handle("float32")):
@@ -259,11 +261,13 @@ class
TestSubroutineCallToExternallyVisibleSubroutine(tvm.testing.CompareBeforeA
T.attr("default", "device_id", 0)
T.attr("default", "device_type", 1)
mod.subroutine(A_data)
+ T.ret(T.int32(0))
@T.prim_func
def subroutine(A_data: T.handle("float32")) -> T.int32:
T.func_attr({"global_symbol": "subroutine", "target":
T.target("llvm")})
T.evaluate(A_data)
+ T.ret(T.int32(0))
return mod
@@ -316,6 +320,7 @@ class
TestCallExternallyVisibleSubroutineWithDLTensor(tvm.testing.CompareBeforeA
T.attr("default", "device_id", 0)
T.attr("default", "device_type", 1)
mod.subroutine(A_data)
+ T.ret(T.int32(0))
@T.prim_func
def subroutine(A_data: T.handle("float32")) -> T.int32:
@@ -323,6 +328,7 @@ class
TestCallExternallyVisibleSubroutineWithDLTensor(tvm.testing.CompareBeforeA
T.attr("default", "device_id", 0)
T.attr("default", "device_type", 1)
T.evaluate(A_data)
+ T.ret(T.int32(0))
return mod
