(tvm) branch main updated: [QoL][Relax] Return well-formed IR from relax::Function::CreateEmpty (#16861)

[wuwei]

This is an automated email from the ASF dual-hosted git repository.

wuwei 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 94a44d7d62 [QoL][Relax] Return well-formed IR from 
relax::Function::CreateEmpty (#16861)
94a44d7d62 is described below

commit 94a44d7d62206849b891c1c262843d88bfb54c3b
Author: Eric Lunderberg <lunderb...@users.noreply.github.com>
AuthorDate: Tue Apr 16 16:26:54 2024 -0500

    [QoL][Relax] Return well-formed IR from relax::Function::CreateEmpty 
(#16861)
    
    Prior to this commit, the static method `relax::Function::CreateEmpty`
    returned a function with a nullptr as the body.  While only intended for 
use in
    bookkeeping for TVMScript, allowing nullptr in this location can cause
    unexpected segfaults while debugging.  For example, adding a print
    statement
    
    This commit updates the `relax::Function::CreateEmpty` function to
    contain a placeholder body, consistent with the `ret_struct_info`
    argument provided.
---
 include/tvm/relax/expr.h |  2 ++
 src/relax/ir/expr.cc     | 24 +++++++++++++++++++-----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index 40707675fe..e2176cf720 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -1045,6 +1045,8 @@ class ExternFuncNode : public BaseFuncNode {
 class ExternFunc : public BaseFunc {
  public:
   TVM_DLL ExternFunc(String global_symbol, Span span = Span());
+  TVM_DLL ExternFunc(String global_symbol, StructInfo struct_info, Span span = 
Span());
+
   TVM_DEFINE_OBJECT_REF_METHODS(ExternFunc, BaseFunc, ExternFuncNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncNode);
 };
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 1b5551e509..0530bb770b 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -559,10 +559,18 @@ Function Function::CreateEmpty(Array<Var> params, 
StructInfo ret_struct_info, bo
 
   FuncStructInfo finfo(param_sinfo, ret_struct_info, is_pure);
 
+  // A dummy body, to ensure that the empty function is still well-formed.
+  Expr body = [&]() -> Expr {
+    Var output("output", ret_struct_info);
+    Call expr(ExternFunc("_dummy_function", FuncStructInfo({}, 
ret_struct_info)), {});
+
+    return SeqExpr({BindingBlock({VarBinding(output, expr)})}, output);
+  }();
+
   // set the fields
   ObjectPtr<FunctionNode> n = make_object<FunctionNode>();
   n->params = std::move(params);
-  n->body = Expr();
+  n->body = std::move(body);
   n->is_pure = is_pure;
   n->checked_type_ = GetStaticType(finfo);
   n->struct_info_ = std::move(finfo);
@@ -602,13 +610,19 @@ FuncStructInfo GetExternFuncStructInfo() {
 
 TVM_REGISTER_NODE_TYPE(ExternFuncNode);
 
-ExternFunc::ExternFunc(String global_symbol, Span span) {
+ExternFunc::ExternFunc(String global_symbol, Span span)
+    : ExternFunc(global_symbol, GetExternFuncStructInfo(), span) {}
+
+ExternFunc::ExternFunc(String global_symbol, StructInfo struct_info, Span 
span) {
+  CHECK(struct_info.as<FuncStructInfoNode>())
+      << "ExternFunc must have FuncStructInfo, "
+      << "but declaration of '" << global_symbol << "' received " << 
struct_info;
+
   ObjectPtr<ExternFuncNode> n = make_object<ExternFuncNode>();
   n->global_symbol = std::move(global_symbol);
   n->span = span;
-  static auto sinfo = GetExternFuncStructInfo();
-  n->struct_info_ = sinfo;
-  n->checked_type_ = GetStaticType(sinfo);
+  n->struct_info_ = struct_info;
+  n->checked_type_ = GetStaticType(struct_info);
   data_ = std::move(n);
 }