[cc] Compose ActionRecorder outputs into a single C file for Emscript…

…en (#1629)

* app

* misc

* tmp save

* emscripten

* TI_ACTION_RECORD=record.yml env var

* clean

* cli

* fix cli

* [skip ci] cleaner

* prepare_sandbox now take no arg for simplicity

* [skip ci] Update python/taichi/core/util.py

Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

* revert kernel.cpp since yuanming's insist

* Apply suggestions from code review

* rev part

* [skip ci] Update python/taichi/core/util.py

Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

python/taichi/cc_compose.py

@@ -0,0 +1,96 @@
+import yaml
+import warnings
+import sys
+
+
+class Composer:
+    def __init__(self, fout, entries, emscripten=False):
+        self.fout = fout
+        self.entries = entries
+        self.emscripten = emscripten
+        self.launches = []
+
+    def emit(self, line):
+        print(line, file=self.fout)
+
+    def run(self):
+        if self.emscripten:
+            self.emit('#include <emscripten.h>')
+
+        for e in self.entries:
+            action = e['action']
+            func = getattr(self, 'do_' + action, self.do_unknown)
+            func(e)
+
+    def do_unknown(self, e):
+        self.emit(f"Unknown action type: {e['action']}")
+
+    def do_compile_runtime(self, e):
+        header = e['runtime_header']
+        source = e['runtime_source']
+
+        self.emit(header)
+        self.emit(source)
+        self.emit('')
+
+    def do_compile_layout(self, e):
+        source = e['layout_source']
+
+        self.emit(source)
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+        self.emit('struct Ti_S0root Ti_root;')
+        self.emit('')
+
+    def do_allocate_buffer(self, e):
+        root_size = e['root_size']
+        gtmp_size = e['gtmp_size']
+        extr_size = 4 * 1024 * 1024  # pinpoint: 4 MB
+
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+            self.emit(f'Ti_i8 Ti_extr[{extr_size}];')
+
+        self.emit(f'Ti_i8 Ti_gtmp[{gtmp_size}];')
+
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+
+        self.emit('union Ti_BitCast Ti_args[8];')
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+        self.emit('Ti_i32 Ti_earg[8 * 8];')
+        self.emit('')
+
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+        self.emit('struct Ti_Context Ti_ctx = {')
+        self.emit('  &Ti_root, Ti_gtmp, Ti_args, Ti_earg,')
+        self.emit('};')
+        self.emit('')
+
+    def do_compile_kernel(self, e):
+        name = e['kernel_name']
+        source = e['kernel_source']
+
+        if self.emscripten:
+            self.emit('EMSCRIPTEN_KEEPALIVE')
+        self.emit(source)
+        self.emit('')
+
+    def do_launch_kernel(self, e):
+        self.launches.append(e)
+
+
+def main(fin_name, fout_name, emscripten=False):
+    with open(fin_name, 'r') as fin:
+        warnings.filterwarnings('ignore')
+        obj = yaml.load(fin)
+
+    with open(fout_name, 'w') as fout:
+        comp = Composer(fout, obj, emscripten)
+        comp.run()
+
+
+if __name__ == '__main__':
+    main(sys.argv[1], sys.argv[2], len(sys.argv) > 3)

python/taichi/core/util.py

@@ -251,17 +251,17 @@ def check_exists(src):
         )
 
 
-def prepare_sandbox(src):
-    global g_tmp_dir
-    check_exists(src)
+def prepare_sandbox():
+    '''
+    Returns a temporary directory, which will be automatically deleted on exit.
+    It may contain the taichi_core shared object or some misc. files.
+    '''
     import atexit
     import shutil
     from tempfile import mkdtemp
     tmp_dir = mkdtemp(prefix='taichi-')
     atexit.register(shutil.rmtree, tmp_dir)
     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/'))
     return tmp_dir
 
@@ -285,6 +285,8 @@ def get_unique_task_id():
     import_ti_core()
     if get_os_name() != 'win':
         dll = ctypes.CDLL(get_core_shared_object(), mode=ctypes.RTLD_LOCAL)
+        # The C backend needs a temporary directory for the generated .c and compiled .so files:
+        ti_core.set_tmp_dir(prepare_sandbox())  # TODO: always allocate a tmp_dir for all situations
 
     ti_core.set_python_package_dir(package_root())
     os.makedirs(ti_core.get_repo_dir(), exist_ok=True)
@@ -295,12 +297,15 @@ def get_unique_task_id():
         os.environ['DYLD_FALLBACK_LIBRARY_PATH'] = get_runtime_directory()
         lib_path = os.path.join(bin_dir, 'libtaichi_core.dylib')
         tmp_cwd = os.getcwd()
-        tmp_dir = prepare_sandbox(lib_path)
+        tmp_dir = prepare_sandbox()
+        check_exists(lib_path)
+        shutil.copy(lib_path, os.path.join(tmp_dir, 'taichi_core.so'))
         os.chdir(tmp_dir)
         sys.path.append(tmp_dir)
         import taichi_core as ti_core
         os.chdir(tmp_cwd)
 
+    # TODO: unify importing infrastructure:
     elif get_os_name() == 'linux':
         bin_dir = get_bin_directory()
         if 'LD_LIBRARY_PATH' in os.environ:
@@ -310,7 +315,9 @@ def get_unique_task_id():
         lib_path = os.path.join(bin_dir, 'libtaichi_core.so')
         check_exists(lib_path)
         tmp_cwd = os.getcwd()
-        tmp_dir = prepare_sandbox(lib_path)
+        tmp_dir = prepare_sandbox()
+        check_exists(lib_path)
+        shutil.copy(lib_path, os.path.join(tmp_dir, 'taichi_core.so'))
         os.chdir(tmp_dir)
         sys.path.append(tmp_dir)
         try:
@@ -423,6 +430,10 @@ def at_startup():
 
     ti_core.set_core_state_python_imported(True)
 
+    record_file = os.environ.get('TI_ACTION_RECORD')
+    if record_file:
+        ti_core.start_recording(record_file)
+
 
 def start_memory_monitoring(output_fn, pid=-1, interval=1):
     # removing dependency on psutil

python/taichi/main.py

@@ -1003,6 +1003,30 @@ def dist(self, arguments: list = sys.argv[2:]):
         sys.argv.append(args.mode)
         runpy.run_path('build.py')
 
+    @register
+    def cc_compose(self, arguments: list = sys.argv[2:]):
+        """Compose C backend action record into a complete C file"""
+        parser = argparse.ArgumentParser(prog='ti cc_compose',
+                                         description=f"{self.cc_compose.__doc__}")
+        parser.add_argument(
+            'fin_name',
+            help='Action record YAML file name from C backend, e.g. program.yml')
+        parser.add_argument(
+            'fout_name',
+            help='The output C source file name, e.g. program.c')
+        parser.add_argument(
+            '-e',
+            '--emscripten',
+            required=False,
+            default=False,
+            dest='emscripten',
+            action='store_true',
+            help='Generate output C file for Emscripten instead of raw C')
+        args = parser.parse_args(arguments)
+
+        from .cc_compose import main
+        main(args.fin_name, args.fout_name, args.emscripten)
+
 
 def main():
     cli = TaichiMain()

python/taichi/testing.py

@@ -2,7 +2,6 @@
 
 print('[Taichi] loading test module')
 
-
 ## Helper functions
 def approx(expected, **kwargs):
     '''Tweaked pytest.approx for OpenGL low percisions'''

taichi/backends/cc/cc_configuation.h

@@ -9,7 +9,7 @@ struct CCConfiguation {
   std::string compile_cmd, link_cmd;
 
   CCConfiguation()
-      : compile_cmd("gcc -Wc99-c11-compat -c -o '{}' '{}'"),
+      : compile_cmd("gcc -Wc99-c11-compat -c -o '{}' '{}' -O3"),
         link_cmd("gcc -shared -fPIC -o '{}' '{}'") {
   }
 };

taichi/backends/cc/context.h

@@ -23,6 +23,7 @@ union Ti_BitCast {
   Ti_u8 *ptr_u8;
   Ti_f32 *ptr_f32;
   Ti_f64 *ptr_f64;
+  void *ptr_void;
 };
 
 struct Ti_Context {

taichi/util/str.cpp

@@ -17,7 +17,6 @@ std::string c_quoted(std::string const &str) {
       REG_ESC('\n', "n");
       REG_ESC('\a', "a");
       REG_ESC('\b', "b");
-      REG_ESC('\?', "?");
       REG_ESC('\v', "v");
       REG_ESC('\t', "t");
       REG_ESC('\f', "f");