[GitHub] [tvm] tqchen commented on a change in pull request #7084: [TIR] Support Return in TIR

[GitBox <gi...@apache.org>]

tqchen commented on a change in pull request #7084:
URL: https://github.com/apache/tvm/pull/7084#discussion_r554276298



##########
File path: src/tir/transforms/make_packed_api.cc
##########
@@ -41,6 +41,56 @@
 namespace tvm {
 namespace tir {
 
+class ReturnRewriter : public StmtMutator {
+ public:
+  explicit ReturnRewriter(Var ret_var, Var ret_tcode) : ret_var_(ret_var), ret_tcode_(ret_tcode) {}
+
+  Stmt VisitStmt_(const EvaluateNode* node) override {
+    Stmt ret = StmtMutator::VisitStmt_(node);
+    const EvaluateNode* eval = ret.as<EvaluateNode>();
+    ICHECK(eval);
+    if (const CallNode* call = eval->value.as<CallNode>()) {
+      if (call->op.same_as(builtin::ret())) {
+        ICHECK_EQ(call->args.size(), 1) << "tir.ret expect a single argument.";
+        ret = WriteToOut(call->args[0], ret_var_, ret_tcode_);
+      }
+    }
+    return ret;
+  }
+
+ private:
+  std::pair<int, PrimExpr> ConvertForFFI(PrimExpr val) {
+    DataType dtype = val.dtype();
+    if (dtype.is_int() || dtype.is_uint()) {
+      return {kTVMArgInt, Cast(DataType::Int(64), val)};
+    } else if (dtype.is_float()) {
+      return {kTVMArgFloat, Cast(DataType::Float(64), val)};
+    } else if (dtype.is_void()) {
+      return {kTVMNullptr, val};
+    } else {
+      LOG(FATAL) << "data type " << dtype << " not supported yet";
+    }
+    return {kTVMNullptr, val};
+  }
+  // convert val's data type to FFI data type, return type code
+  Stmt WriteToOut(PrimExpr val, Var ret_var, Var ret_tcode) {
+    auto p = ConvertForFFI(val);
+    int tcode = p.first;
+    val = p.second;
+    Stmt store_val = Store(ret_var_, val, 0, const_true());
+    Stmt store_tcode = Store(ret_tcode_, tcode, 0, const_true());
+    return SeqStmt({store_val, store_tcode});

Review comment:
       We might need a low level ret intrinsic to terminate the control flow, instead of continue running things.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org