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

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():
+    '''
+    return a runtime temporary directory name, which will be automatically
+    deleted on exit.
+    '''
     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)
+        # TODO: add tmp dir for win too, so that C backend can be functional:
+        ti_core.set_tmp_dir(prepare_sandbox())
 
     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

@@ -1,7 +1,5 @@
 import taichi as ti
 
-print('[Taichi] loading test module')
-
 
 ## Helper functions
 def approx(expected, **kwargs):

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/backends/opengl/shaders/random.glsl.h

@@ -5,11 +5,11 @@ STR(
 uvec4 _rand_;
 
 void _init_rand() {
-  uint i = (54321 + gl_GlobalInvocationID.x) * (12345 + _rand_state_);
-  _rand_.x = 123456789 * i * 1000000007;
-  _rand_.y = 362436069;
-  _rand_.z = 521288629;
-  _rand_.w = 88675123;
+  uint i = (54321u + gl_GlobalInvocationID.x) * (12345u + uint(_rand_state_));
+  _rand_.x = 123456789u * i * 1000000007u;
+  _rand_.y = 362436069u;
+  _rand_.z = 521288629u;
+  _rand_.w = 88675123u;
 
   // Yes, this is not an atomic operation, but just fine since no matter
   // how `_rand_state_` changes, `gl_GlobalInvocationID.x` can still help
@@ -22,7 +22,7 @@ uint _rand_u32() {
   uint t = _rand_.x ^ (_rand_.x << 11);
   _rand_.xyz = _rand_.yzw;
   _rand_.w = (_rand_.w ^ (_rand_.w >> 19)) ^ (t ^ (t >> 8));
-  return _rand_.w * 1000000007;
+  return _rand_.w * 1000000007u;
 }
 
 float _rand_f32() {

taichi/program/kernel.cpp

@@ -119,12 +119,8 @@ void Kernel::set_arg_float(int i, float64 d) {
   TI_ASSERT_INFO(
       !args[i].is_nparray,
       "Assigning a scalar value to a numpy array argument is not allowed");
-
-  ActionRecorder::get_instance().record(
-      "set_kernel_arg_float64", {ActionArg("kernel_name", name),
-                                 ActionArg("arg_id", i), ActionArg("val", d)});
-
   auto dt = args[i].dt;
+
   if (dt == DataType::f32) {
     program.context.set_arg(i, (float32)d);
   } else if (dt == DataType::f64) {
@@ -154,12 +150,8 @@ void Kernel::set_arg_int(int i, int64 d) {
   TI_ASSERT_INFO(
       !args[i].is_nparray,
       "Assigning scalar value to numpy array argument is not allowed");
-
-  ActionRecorder::get_instance().record(
-      "set_kernel_arg_int64", {ActionArg("kernel_name", name),
-                               ActionArg("arg_id", i), ActionArg("val", d)});
-
   auto dt = args[i].dt;
+
   if (dt == DataType::i32) {
     program.context.set_arg(i, (int32)d);
   } else if (dt == DataType::i64) {
@@ -247,12 +239,6 @@ void Kernel::set_arg_nparray(int i, uint64 ptr, uint64 size) {
   TI_ASSERT_INFO(args[i].is_nparray,
                  "Assigning numpy array to scalar argument is not allowed");
 
-  ActionRecorder::get_instance().record(
-      "set_kernel_arg_ext_ptr",
-      {ActionArg("kernel_name", name), ActionArg("arg_id", i),
-       ActionArg("address", fmt::format("0x{:x}", ptr)),
-       ActionArg("array_size_in_bytes", (int64)size)});
-
   args[i].size = size;
   program.context.set_arg(i, ptr);
 }

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");