[TVMScript] Represent tir::builtin::ret() using python "return"

python/tvm/script/parser/tir/parser.py

@@ -520,7 +520,8 @@ def visit_return(self: Parser, node: doc.Return) -> None:
     node : doc.Return
         The doc AST return node.
     """
-    self.report_error(node, "Return is not allowed.")
+    value = self.eval_expr(node.value)
+    T.evaluate(tvm.tir.ret(value))
 
 
 @dispatch.register(token="tir", type_name="tvm_declare_function")

src/script/printer/tir/stmt.cc

@@ -64,8 +64,30 @@ bool IsAncestorOfAllVarUse(const tir::Stmt& node, const ObjectRef& var, const IR
   return false;
 }
 
+Optional<PrimExpr> FindReturnValue(const tir::Stmt& node) {
+  auto eval = node.as<tir::EvaluateNode>();
+  if (!eval) return NullOpt;
+
+  auto call = eval->value.as<tir::CallNode>();
+  if (!call) return NullOpt;
+
+  if (!call->op.same_as(tir::builtin::ret())) return NullOpt;
+
+  if (call->args.size() != 1) return NullOpt;
+
+  return call->args[0];
+}
+
 TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
     .set_dispatch<tir::Evaluate>("", [](tir::Evaluate eval, ObjectPath 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")->ArrayIndex(0));
+          return ReturnDoc(value);
+        }
+      }
+
       ExprDoc value = d->AsDoc<ExprDoc>(eval->value, p->Attr("value"));
       if (eval->value->IsInstance<tir::CallNode>()) {
         return ExprStmtDoc(value);

tests/python/tvmscript/test_tvmscript_printer_tir.py

@@ -900,5 +900,22 @@ def func(a_name: T.handle):
     assert re.match(expected_regex, script)
 
 
+def test_return_statement():
+    from tvm.script import tir as T
+
+    @T.prim_func
+    def func():
+        T.evaluate(T.ret(5))
+
+    expected_output = """
+# from tvm.script import tir as T
+
+@T.prim_func
+def func():
+    return 5
+    """
+    _assert_print(func, expected_output)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/tvmscript/test_tvmscript_syntax_sugar.py

@@ -492,5 +492,19 @@ def implicit():
     assert_structural_equal_ignore_global_symbol(implicit, explicit)
 
 
+def test_return_statement():
+    """A python `return` statement uses `T.ret`"""
+
+    @T.prim_func
+    def explicit():
+        T.evaluate(T.ret(5))
+
+    @T.prim_func
+    def implicit():
+        return 5
+
+    assert_structural_equal_ignore_global_symbol(implicit, explicit)
+
+
 if __name__ == "__main__":
     tvm.testing.main()