From c453f2f3c0817f8617913d97584fcfafaf6426bd Mon Sep 17 00:00:00 2001
From: archibate <17721388340@163.com>
Date: Fri, 8 May 2020 13:57:19 +0800
Subject: [PATCH] use KernelReturnStmt from #917

---
 taichi/common/testing.h             | 5 ++---
 taichi/transforms/constant_fold.cpp | 6 +++---
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/taichi/common/testing.h b/taichi/common/testing.h
index 4753e07b77ef84..d4944e50783b58 100644
--- 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
index 7f443f2eaea09e..a7663e9f40715f 100644
--- 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;
   }