[skip ci] suces

python/taichi/lang/kernel.py

@@ -387,13 +387,17 @@ def call_back():
                 self.runtime.target_tape.insert(self, args)
 
             t_kernel()
+            # todo union_cast<int> from uint64
+            ret = taichi_lang_core.fetch_return_result()
 
             if callbacks:
                 import taichi as ti
                 ti.sync()
                 for c in callbacks:
                     c()
 
+            return ret
+
         return func__
 
     def match_ext_arr(self, v, needed):
@@ -412,7 +416,8 @@ def __call__(self, *args, **kwargs):
         instance_id, arg_features = self.mapper.lookup(args)
         key = (self.func, instance_id)
         self.materialize(key=key, args=args, arg_features=arg_features)
-        return self.compiled_functions[key](*args)
+        ret = self.compiled_functions[key](*args)
+        return ret
 
 
 # For a Taichi class definition like below:
@@ -481,7 +486,7 @@ def wrapped(*args, **kwargs):
 
         @wraps(func)
         def wrapped(*args, **kwargs):
-            primal(*args, **kwargs)
+            return primal(*args, **kwargs)
 
         wrapped.grad = adjoint
 

taichi/python/export_lang.cpp

@@ -309,6 +309,10 @@ void export_lang(py::module &m) {
     current_ast_builder().insert(Stmt::make<FrontendKernelReturnStmt>(value));
   });
 
+  m.def("fetch_return_result", []() {
+    return get_current_program().context.get_arg_as_uint64(0);
+  });
+
   m.def("insert_continue_stmt", [&]() {
     current_ast_builder().insert(Stmt::make<FrontendContinueStmt>());
   });