[refactor] [autodiff] Clean redundant compiled functions and refactor kernel key

python/taichi/ad/_ad.py

@@ -7,7 +7,6 @@
 
 from taichi._snode.fields_builder import FieldsBuilder
 from taichi.lang import impl
-from taichi.lang.enums import AutodiffMode
 from taichi.lang.field import ScalarField
 from taichi.lang.snode import SNode
 
@@ -299,14 +298,13 @@ def __exit__(self, _type, value, tb):
         self.runtime.fwd_mode_manager = None
         self.recover_kernels()
 
-    def insert(self, func):
-        self.calls.append(func)
+    def insert(self, func, mode_original):
+        self.calls.append((func, mode_original))
 
     def recover_kernels(self):
         assert self.entered, "Before recover the kernels, fwd mode manager must be entered."
-        for f in self.calls:
-            f.autodiff_mode = AutodiffMode.NONE
-            f.compiled_functions = f.runtime.compiled_functions
+        for f, mode_original in self.calls:
+            f.autodiff_mode = mode_original
         self.kernels_recovered = True
 
 

python/taichi/lang/impl.py

@@ -219,8 +219,6 @@ def __init__(self, kernels=None):
         self.materialized = False
         self.prog = None
         self.compiled_functions = {}
-        self.compiled_fwd_mode_grad_functions = {}
-        self.compiled_grad_functions = {}
         self.src_info_stack = []
         self.inside_kernel = False
         self.current_kernel = None
@@ -235,9 +233,7 @@ def __init__(self, kernels=None):
         self._signal_handler_registry = None
 
     def get_num_compiled_functions(self):
-        return len(self.compiled_functions) + len(
-            self.compiled_grad_functions) + len(
-                self.compiled_fwd_mode_grad_functions)
+        return len(self.compiled_functions)
 
     def src_info_guard(self, info):
         return SrcInfoGuard(self.src_info_stack, info)

python/taichi/lang/kernel_impl.py

@@ -405,6 +405,8 @@ def __init__(self, _func, autodiff_mode, _classkernel=False):
         self.func = _func
         self.kernel_counter = Kernel.counter
         Kernel.counter += 1
+        assert autodiff_mode in (AutodiffMode.NONE, AutodiffMode.FORWARD,
+                                 AutodiffMode.REVERSE)
         self.autodiff_mode = autodiff_mode
         self.grad = None
         self.arguments = []
@@ -426,14 +428,6 @@ def __init__(self, _func, autodiff_mode, _classkernel=False):
 
     def reset(self):
         self.runtime = impl.get_runtime()
-        if self.autodiff_mode == AutodiffMode.NONE:
-            self.compiled_functions = self.runtime.compiled_functions
-        elif self.autodiff_mode == AutodiffMode.REVERSE:
-            self.compiled_functions = self.runtime.compiled_grad_functions
-        elif self.autodiff_mode == AutodiffMode.FORWARD:
-            self.compiled_functions = self.runtime.compiled_fwd_mode_grad_functions
-        else:
-            raise NotImplementedError("Unknown autodiff mode.")
 
     def extract_arguments(self):
         sig = inspect.signature(self.func)
@@ -488,17 +482,10 @@ def extract_arguments(self):
 
     def materialize(self, key=None, args=None, arg_features=None):
         if key is None:
-            key = (self.func, 0)
+            key = (self.func, 0, self.autodiff_mode)
         self.runtime.materialize()
 
-        # Transform the primal kernel to forward mode grad kernel
-        # then recover to primal when exiting the forward mode manager
-        if self.runtime.fwd_mode_manager:
-            self.autodiff_mode = AutodiffMode.FORWARD
-            self.runtime.fwd_mode_manager.insert(self)
-            self.compiled_functions = self.runtime.compiled_fwd_mode_grad_functions
-
-        if key in self.compiled_functions:
+        if key in self.runtime.compiled_functions:
             return
 
         grad_suffix = ""
@@ -549,8 +536,9 @@ def taichi_ast_generator(kernel_cxx):
 
         self.kernel_cpp = taichi_kernel
 
-        assert key not in self.compiled_functions
-        self.compiled_functions[key] = self.get_function_body(taichi_kernel)
+        assert key not in self.runtime.compiled_functions
+        self.runtime.compiled_functions[key] = self.get_function_body(
+            taichi_kernel)
         self.compiled_kernels[key] = taichi_kernel
 
     def get_torch_callbacks(self, v, has_torch, is_ndarray=True):
@@ -808,7 +796,7 @@ def match_ext_arr(v):
 
     def ensure_compiled(self, *args):
         instance_id, arg_features = self.mapper.lookup(args)
-        key = (self.func, instance_id)
+        key = (self.func, instance_id, self.autodiff_mode)
         self.materialize(key=key, args=args, arg_features=arg_features)
         return key
 
@@ -817,14 +805,25 @@ def ensure_compiled(self, *args):
     @_shell_pop_print
     def __call__(self, *args, **kwargs):
         args = _process_args(self, args, kwargs)
+
+        # Transform the primal kernel to forward mode grad kernel
+        # then recover to primal when exiting the forward mode manager
+        if self.runtime.fwd_mode_manager:
+            # TODO: if we would like to compute 2nd-order derivatives by forward-on-reverse in a nested context manager fashion,
+            # i.e., a `Tape` nested in the `FwdMode`, we can transform the kernels with `mode_original == AutodiffMode.REVERSE` only,
+            # to avoid duplicate computation for 1st-order derivatives
+            mode_original = self.autodiff_mode
+            self.autodiff_mode = AutodiffMode.FORWARD
+            self.runtime.fwd_mode_manager.insert(self, mode_original)
+
         if self.autodiff_mode != AutodiffMode.NONE and impl.current_cfg(
         ).opt_level == 0:
             _logging.warn(
                 """opt_level = 1 is enforced to enable gradient computation."""
             )
             impl.current_cfg().opt_level = 1
         key = self.ensure_compiled(*args)
-        return self.compiled_functions[key](*args)
+        return self.runtime.compiled_functions[key](*args)
 
 
 # For a Taichi class definition like below: