move runtime file inside the sandbox as well suggested by @KLozes #501 (

#517)

* move runtime file inside the sandbox as well (as @KLozes suggested in #501)

* Update util.py

Co-authored-by: Yuanming Hu <[email protected]>

python/taichi/core/util.py

@@ -19,7 +19,7 @@ def in_docker():
     return True
 
 
-def import_tc_core():
+def import_tc_core(tmp_dir=None):
   global tc_core
   if get_os_name() != 'win':
     old_flags = sys.getdlopenflags()
@@ -38,6 +38,8 @@ def import_tc_core():
     sys.setdlopenflags(old_flags)
   lib_dir = os.path.join(package_root(), 'lib')
   core.set_lib_dir(locale_encode(lib_dir))
+  if tmp_dir is not None:
+    core.set_tmp_dir(locale_encode(tmp_dir))
 
 
 def locale_encode(s):
@@ -160,6 +162,7 @@ def build():
   os.chdir(tmp_cwd)
 
 def prepare_sandbox(src):
+  global g_tmp_dir
   assert os.path.exists(src)
   import atexit
   import shutil
@@ -169,6 +172,8 @@ def prepare_sandbox(src):
   print(f'[Taichi] preparing sandbox at {tmp_dir}')
   dest = os.path.join(tmp_dir, 'taichi_core.so')
   shutil.copy(src, dest)
+  os.mkdir(os.path.join(tmp_dir, 'runtime/'))
+  print(f'[Taichi] sandbox prepared')
   return tmp_dir
 
 
@@ -213,7 +218,7 @@ def prepare_sandbox(src):
     os.chdir(tmp_dir)
     sys.path.append(tmp_dir)
     try:
-      import_tc_core()
+      import_tc_core(tmp_dir)
     except Exception as e:
       from colorama import Fore, Back, Style
       print_red_bold("Taichi core import failed: ", end='')

python/taichi/misc/settings.py

@@ -32,7 +32,7 @@ def get_project_directory(project=None):
 
 
 def get_runtime_directory():
-  bin_rel_path = ['extneral', 'lib']
+  bin_rel_path = ['external', 'lib']
   return os.environ.get('TAICHI_BIN_DIR',
                         os.path.join(get_repo_directory(), *bin_rel_path))
 

taichi/python_bindings.cpp

@@ -18,6 +18,7 @@ TI_NAMESPACE_END
 TLANG_NAMESPACE_BEGIN
 
 std::string compiled_lib_dir;
+std::string runtime_tmp_dir;
 
 Expr expr_index(const Expr &expr, const Expr &index) {
   return expr[index];
@@ -445,6 +446,7 @@ void export_lang(py::module &m) {
   m.def("libdevice_path", libdevice_path);
 
   m.def("set_lib_dir", [&](const std::string &dir) { compiled_lib_dir = dir; });
+  m.def("set_tmp_dir", [&](const std::string &dir) { runtime_tmp_dir = dir; });
 
   m.def("get_commit_hash", get_commit_hash);
   m.def("get_version_string", get_version_string);

taichi/taichi_llvm_context.cpp

@@ -97,8 +97,15 @@ std::string get_runtime_fn(Arch arch) {
   return fmt::format("runtime_{}.bc", arch_name(arch));
 }
 
+std::string get_runtime_src_dir() {
+    return get_repo_dir() + "/taichi/runtime/";
+}
+
 std::string get_runtime_dir() {
-  return get_repo_dir() + "/taichi/runtime/";
+  if (runtime_tmp_dir.size() == 0)
+    return get_runtime_src_dir();
+  else
+    return runtime_tmp_dir + "/runtime/";
 }
 
 void compile_runtime_bitcode(Arch arch) {
@@ -110,6 +117,7 @@ void compile_runtime_bitcode(Arch arch) {
     auto clang = find_existing_command({"clang-7", "clang-8", "clang-9", "clang"});
     TI_ASSERT(command_exist("llvm-as"));
     TI_TRACE("Compiling runtime module bitcode...");
+    auto runtime_src_folder = get_runtime_src_dir();
     auto runtime_folder = get_runtime_dir();
     std::string macro = fmt::format(" -D ARCH_{} ", arch_name(arch));
 #if defined(TI_ARCH_ARM)
@@ -118,7 +126,7 @@ void compile_runtime_bitcode(Arch arch) {
     int ret = std::system(
         fmt::format(
             "{} -S {}runtime.cpp -o {}runtime.ll -emit-llvm -std=c++17 {}",
-            clang, runtime_folder, runtime_folder, macro)
+            clang, runtime_src_folder, runtime_folder, macro)
             .c_str());
     if (ret) {
       TI_ERROR("Runtime compilation failed.");

taichi/tlang_util.h

@@ -342,6 +342,7 @@ struct CompileConfig {
 
 extern CompileConfig default_compile_config;
 extern std::string compiled_lib_dir;
+extern std::string runtime_tmp_dir;
 
 bool command_exist(const std::string &command);