From 5e49f423afb95da856ca6a2a55c0c5981a427535 Mon Sep 17 00:00:00 2001
From: Ye Kuang <yekuang.ky@gmail.com>
Date: Sun, 8 Nov 2020 21:45:14 +0900
Subject: [PATCH] [async] More fixes of constant folding in async mode

---
 taichi/transforms/constant_fold.cpp | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/taichi/transforms/constant_fold.cpp b/taichi/transforms/constant_fold.cpp
index 2e7829d8b2962..7b469e7ca2f2e 100644
--- a/taichi/transforms/constant_fold.cpp
+++ b/taichi/transforms/constant_fold.cpp
@@ -96,9 +96,12 @@ class ConstantFold : public BasicStmtVisitor {
     auto launch_ctx = ker->make_launch_context();
     launch_ctx.set_arg_raw(0, lhs.val_u64);
     launch_ctx.set_arg_raw(1, rhs.val_u64);
-    (*ker)(launch_ctx);
     auto &current_program = stmt->get_kernel()->program;
-    ret.val_i64 = current_program.fetch_result<int64_t>(0);
+    {
+      std::lock_guard<std::mutex> _(current_program.jit_evaluator_cache_mut);
+      (*ker)(launch_ctx);
+      ret.val_i64 = current_program.fetch_result<int64_t>(0);
+    }
     return true;
   }
 
@@ -116,9 +119,12 @@ class ConstantFold : public BasicStmtVisitor {
     auto *ker = get_jit_evaluator_kernel(id);
     auto launch_ctx = ker->make_launch_context();
     launch_ctx.set_arg_raw(0, operand.val_u64);
-    (*ker)(launch_ctx);
     auto &current_program = stmt->get_kernel()->program;
-    ret.val_i64 = current_program.fetch_result<int64_t>(0);
+    {
+      std::lock_guard<std::mutex> _(current_program.jit_evaluator_cache_mut);
+      (*ker)(launch_ctx);
+      ret.val_i64 = current_program.fetch_result<int64_t>(0);
+    }
     return true;
   }