[export] [cc] Support ti.RecordGroupHint to group launched kernels

docs/export_kernels.rst

@@ -21,7 +21,7 @@ a binary format, by compiling and linking exported C code to their project.
 The workflow of exporting
 -------------------------
 
-Use ``ti.core.start_recording`` in the Taichi program you want to export.
+Use ``ti.start_recording`` in the Taichi program you want to export.
 
 Suppose you want to export `examples/mpm88.py <https://github.com/taichi-dev/taichi/blob/master/examples/mpm88.py>`_, here is the workflow:
 
@@ -34,7 +34,7 @@ First, modify ``mpm88.py`` as shown below:
 
     import taichi as ti
 
-    ti.core.start_recording('mpm88.yml')
+    ti.start_recording('mpm88.yml')
     ti.init(arch=ti.cc)
 
     ... # your program
@@ -108,6 +108,8 @@ source file ``mpm88.c``:
    It should also be functional when compiled using a C++ compiler.
 
 
+
+
 Calling the exported kernels
 ----------------------------
 
@@ -203,3 +205,69 @@ shared object. Then it can be called from other langurages that provides a C
 interface, including but not limited to Julia, Matlab, Mathematica, Java, etc.
 
 TODO: WIP.
+
+Record kernel group hints
++++++++++++++++++++++++++
+
+Suppose you have a program with lots of kernel.
+
+To run this program in C or Javascript, you have to rewrite their names in
+the exact same order as they were in Python.
+
+Not to say implicitly generated meta kernels like ``fill_tensor`` and
+``clear_gradients`` which is invisible to end-users.
+
+So you may find it hard and error-prone to figure out the correct launch order
+and mangled names.
+
+No worry, we provide a handy tool for such situation: you may guard the desired
+kernels with ``ti.RecordKernelGroup``. For example:
+
+.. code-block:: python
+
+    import taichi as ti
+
+    ti.start_recording('record.yml')
+    ti.init(arch=ti.cc)
+
+    loss = ti.field(float, (), needs_grad=True)
+    x = ti.field(float, 233, needs_grad=True)
+
+
+    @ti.kernel
+    def compute_loss():
+        for i in x:
+            loss[None] += x[i]**2
+
+
+    @ti.kernel
+    def do_some_works():
+        for i in x:
+            x[i] -= x.grad[i]
+
+
+    with ti.RecordKernelGroup('my_substep'):  # HERE!
+        x.fill(0)
+        with ti.Tape(loss):
+            compute_loss()
+        do_some_works()
+
+
+Then the ``ti cc_compose`` command will add a comment at the end of ``record.c``
+as a hint of launch order:
+
+.. code-block:: python
+
+    // group my_substep: ['fill_tensor_c8_0', 'clear_gradients_c24_0', 'clear_gradients_c24_1', 'snode_writer_2', 'snode_writer_4', 'compute_loss_c4_0', 'compute_loss_c5_0_grad_grad', 'do_some_works_c6_0']
+
+
+This is the name of all the kernels launched within the ``ti.RecordGroupHint`` scope,
+sorted by launch order. Copy the list and somehow iterate them in C or Javascript
+and the launch order is exactly same as we had in Python, e.g.:
+
+.. code-block:: c
+
+    Tk_fill_tensor_c8_0(&Ti_ctx);
+    Tk_clear_gradients_c24_0(&Ti_ctx);
+    Tk_clear_gradients_c24_1(&Ti_ctx);
+    ...

misc/demo_record_kernel_group.py

@@ -0,0 +1,26 @@
+import taichi as ti
+
+ti.start_recording('record.yml')
+ti.init(arch=ti.cc)
+
+loss = ti.field(float, (), needs_grad=True)
+x = ti.field(float, 233, needs_grad=True)
+
+
+@ti.kernel
+def compute_loss():
+    for i in x:
+        loss[None] += x[i]**2
+
+
+@ti.kernel
+def do_some_works():
+    for i in x:
+        x[i] -= x.grad[i]
+
+
+with ti.RecordKernelGroup('my_substep'):
+    x.fill(0)
+    with ti.Tape(loss):
+        compute_loss()
+    do_some_works()

misc/prtags.json

@@ -32,5 +32,6 @@
   "cc"              : "C source backend",
   "error"           : "Error messages",
   "blender"         : "Blender intergration",
+  "export"          : "Exporting kernels",
   "release"         : "Release"
 }

python/taichi/cc_compose.py

@@ -3,13 +3,45 @@
 import sys
 
 
-class Composer:
+class ComposerBase:
+    def __init__(self, entries):
+        self.entries = entries
+        self.current_group = None
+        self.groups = {}
+        self.launches = []
+
+    def run(self):
+        for e in self.entries:
+            action = e['action']
+            func = getattr(self, 'do_' + action, self.do_unknown)
+            func(e)
+
+    def do_unknown(self, e):
+        pass
+
+    def do_group_begin(self, e):
+        name = e['content']
+        self.current_group = name
+        self.launches = []
+
+    def do_group_end(self, e):
+        name = e['content']
+        self.groups[self.current_group] = list(self.launches)
+        self.current_group = None
+        self.launches = []
+
+    def do_launch_kernel(self, e):
+        name = e['kernel_name']
+        self.launches.append(name)
+
+
+class ComposerCC(ComposerBase):
     def __init__(self, entries, fout, hdrout, emscripten=False):
+        super().__init__(entries)
+
         self.fout = fout
         self.hdrout = hdrout
-        self.entries = entries
         self.emscripten = emscripten
-        self.launches = []
 
     def emit(self, line):
         print(line, file=self.fout)
@@ -21,13 +53,10 @@ 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)
+        super().run()
 
-    def do_unknown(self, e):
-        pass
+        for key, launches in self.groups.items():
+            self.emit(f'// group {key}: {launches}')
 
     def do_compile_runtime(self, e):
         header = e['runtime_header']
@@ -88,9 +117,6 @@ def do_compile_kernel(self, e):
         declaration = source.split('{', 1)[0].strip()
         self.emit_header(f'extern {declaration};')
 
-    def do_launch_kernel(self, e):
-        self.launches.append(e)
-
 
 def main(fin_name, fout_name, hdrout_name, emscripten=False):
     with open(fin_name, 'r') as fin:
@@ -99,7 +125,7 @@ def main(fin_name, fout_name, hdrout_name, emscripten=False):
 
     with open(hdrout_name, 'w') as hdrout:
         with open(fout_name, 'w') as fout:
-            comp = Composer(obj, fout, hdrout, emscripten)
+            comp = ComposerCC(obj, fout, hdrout, emscripten)
             comp.run()
 
 

python/taichi/core/__init__.py

@@ -1,5 +1,6 @@
 from .util import *
 from .settings import *
+from .record import *
 from .unit import unit
 
 ti_core.build = build

python/taichi/core/record.py

@@ -0,0 +1,61 @@
+from .util import ti_core
+import os
+
+
+def record_action_hint(name, content=None):
+    if content is None:
+        name, content = 'hint', name
+    ti_core.record_action_hint(name, content)
+
+
+def start_recording(filename):
+    ti_core.start_recording(filename)
+
+
+def stop_recording():
+    ti_core.stop_recording()
+
+
+class RecordAction:
+    def __init__(self, filename):
+        self.filename = filename
+
+    def __enter__(self):
+        start_recording(self.filename)
+        return self
+
+    def __exit__(self, *args):
+        stop_recording()
+
+
+class RecordKernelGroup:
+    def __init__(self, name):
+        if name in RecordKernelGroup.recorded:
+            self.name = None
+        else:
+            RecordKernelGroup.recorded.add(name)
+            self.name = name
+
+    def __enter__(self):
+        if self.name is not None:
+            record_action_hint('group_begin', self.name)
+        return self
+
+    def __exit__(self, *args):
+        if self.name is not None:
+            record_action_hint('group_end', self.name)
+
+    recorded = set()
+
+
+record_file = os.environ.get('TI_ACTION_RECORD')
+if record_file:
+    start_recording(record_file)
+
+__all__ = [
+    'start_recording',
+    'stop_recording',
+    'record_action_hint',
+    'RecordAction',
+    'RecordKernelGroup',
+]

python/taichi/core/util.py

@@ -315,10 +315,6 @@ 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/lang/__init__.py

@@ -9,19 +9,6 @@
 
 core = taichi_lang_core
 
-
-def record_action_hint(s):
-    core.record_action_hint(s)
-
-
-def begin_recording(fn):
-    core.begin_recording(fn)
-
-
-def stop_recording():
-    core.stop_recording()
-
-
 runtime = get_runtime()
 
 i = indices(0)

taichi/python/export_lang.cpp

@@ -632,8 +632,8 @@ void export_lang(py::module &m) {
     return result;
   });
 
-  m.def("record_action_hint", [](std::string content) {
-    ActionRecorder::get_instance().record("hint",
+  m.def("record_action_hint", [](std::string name, std::string content) {
+    ActionRecorder::get_instance().record(name,
                                           {ActionArg("content", content)});
   });