use KernelReturnStmt from taichi-dev#917

archibate · May 8, 2020 · c453f2f · c453f2f
1 parent 6bed95f
commit c453f2f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/taichi/common/testing.h b/taichi/common/testing.h
@@ -31,9 +31,8 @@ TI_NAMESPACE_BEGIN
 #define TI_TEST(x) TEST_CASE(x, ("[" x "]"))
 #define TI_CHECK(x) CHECK(x)
 #define TI_TEST_PROGRAM \
-  auto prog_ = std::make_unique<Program>(Arch::x64); \
-  prog_->materialize_layout(); \
-  TI_ASSERT(prog_->llvm_runtime);
+  auto prog_ = std::make_unique<Program>(); \
+  prog_->materialize_layout();
 
 int run_tests(std::vector<std::string> argv);
 

diff --git a/taichi/transforms/constant_fold.cpp b/taichi/transforms/constant_fold.cpp
@@ -92,7 +92,7 @@ class ConstantFold : public BasicStmtVisitor {
           oper->cast<UnaryOpStmt>()->cast_type = id.rhs;
         }
       }
-      auto ret = Stmt::make<ArgStoreStmt>(0, oper.get());
+      auto ret = Stmt::make<KernelReturnStmt>(oper.get());
       current_ast_builder().insert(std::move(lhstmt));
       if (id.is_binary)
         current_ast_builder().insert(std::move(rhstmt));
@@ -126,7 +126,7 @@ class ConstantFold : public BasicStmtVisitor {
     ctx.set_arg<int64_t>(2, rhs.val_i64);
     irpass::print(ker->ir);
     (*ker)();
-    ret.val_i64 = ctx.get_arg<int64_t>(0); // TODO: replace by fetch_reader_result, gen
+    ret.val_i64 = get_current_program().fetch_result<int64_t>(0);
     //TI_INFO("JITEVALf = {}", ret.val_f32);
     TI_INFO("JITEVALi = {}", ret.val_i32);
     return true;
@@ -142,7 +142,7 @@ class ConstantFold : public BasicStmtVisitor {
     ctx.set_arg<int64_t>(1, lhs.val_i64);
     irpass::print(ker->ir);
     (*ker)();
-    ret.val_i64 = ctx.get_arg<int64_t>(0);
+    ret.val_i64 = get_current_program().fetch_result<int64_t>(0);
     return true;
   }