Merge branch 'hp/func-base-refactorization' into hp/func-base-refactorization-factorize-cache

Author: hughperkins

gstaichi/runtime/cuda/jit_cuda.cpp

@@ -78,7 +78,7 @@ JITSessionCUDA::JITSessionCUDA(GsTaichiLLVMContext *tlctx,
     : JITSession(tlctx, config),
       data_layout(data_layout),
       program_impl_(program_impl),
-      config(config) {
+      config_(config) {
   PtxCache::Config ptx_cache_config;
   ptx_cache_config.offline_cache_path = config.offline_cache_file_path;
   ptx_cache_ = std::make_unique<PtxCache>(ptx_cache_config, config);
@@ -91,7 +91,7 @@ JITModule *JITSessionCUDA::add_module(std::unique_ptr<llvm::Module> M,
                                       int max_reg) {
   const char *dump_ir_env = std::getenv(DUMP_IR_ENV.data());
   if (dump_ir_env != nullptr && std::string(dump_ir_env) == "1") {
-    std::filesystem::path ir_dump_dir = config.debug_dump_path;
+    std::filesystem::path ir_dump_dir = config_.debug_dump_path;
     std::filesystem::create_directories(ir_dump_dir);
     std::string dumpName = moduleToDumpName(M.get());
     std::filesystem::path filename =

gstaichi/runtime/cuda/jit_cuda.h

@@ -85,7 +85,7 @@ class JITSessionCUDA : public JITSession {
   std::unique_ptr<PtxCache> ptx_cache_;
   ProgramImpl *program_impl_;
   std::unique_ptr<Finalizer> finalizer_;
-  const CompileConfig &config;
+  const CompileConfig &config_;
 };
 
 #endif

python/gstaichi/lang/_func_base.py

@@ -260,9 +260,11 @@ def get_tree_and_ctx(
         return tree, ctx
 
     def process_args(self, is_pyfunc: bool, is_func: bool, args: tuple[Any, ...], kwargs) -> tuple[Any, ...]:
+        """
+        - expand dataclass args
+        - fuse args and kwargs into a single list of args
+        """
         if is_func and not is_pyfunc:
-            # if typing.TYPE_CHECKING:
-            #     assert isinstance(self, Func)
             current_kernel = self.current_kernel
             if typing.TYPE_CHECKING:
                 assert current_kernel is not None
@@ -315,7 +317,6 @@ def process_args(self, is_pyfunc: bool, is_func: bool, args: tuple[Any, ...], kw
                         raise GsTaichiSyntaxError(f"Unexpected argument '{key}'.")
         elif num_missing_args:
             for i in range(num_args, num_arg_metas):
-                arg = fused_args[i]
                 if fused_args[i] is _ARG_EMPTY:
                     arg_meta = self.arg_metas_expanded[i]
                     raise GsTaichiSyntaxError(f"Missing argument '{arg_meta.name}'.")

python/gstaichi/lang/_template_mapper.py

@@ -6,7 +6,8 @@
 from gstaichi.lang.impl import Program
 from gstaichi.lang.kernel_arguments import ArgMetadata
 
-from ._template_mapper_hotpath import _extract_arg
+from .._test_tools import warnings_helper
+from ._template_mapper_hotpath import _extract_arg, _primitive_types
 
 ArgsHash: TypeAlias = tuple[int, ...]
 Key: TypeAlias = tuple[Any, ...]
@@ -38,7 +39,7 @@ def __init__(self, arguments: list[ArgMetadata], template_slot_locations: list[i
         self.template_slot_locations: list[int] = template_slot_locations
         self.mapping: dict[Key, int] = {}
         self._mapping_cache: dict[ArgsHash, tuple[int, Key]] = {}
-        self._mapping_cache_tracker: dict[ArgsHash, list[ReferenceType]] = {}
+        self._mapping_cache_tracker: dict[ArgsHash, list[ReferenceType | None]] = {}
         self._prog_weakref: ReferenceType[Program] | None = None
 
     def extract(self, raise_on_templated_floats: bool, args: tuple[Any, ...]) -> Key:
@@ -64,8 +65,12 @@ def lookup(self, raise_on_templated_floats: bool, args: tuple[Any, ...]) -> tupl
             prog = self._prog_weakref()
         assert prog is not None
 
-        mapping_cache_tracker: list[ReferenceType] | None = None
-        args_hash: ArgsHash = tuple([id(arg) for arg in args])
+        # Note that it is necessary to handle primitive types separately. First, using their address as cache key must
+        # be avoided, because even though it is theoretically possible, it is overly restrictive. Second, it does not
+        # make sense to use these arguments to track the lifetime of the corresponding cache entry and taking weakref
+        # of primitive types if forbidden anyway.
+        mapping_cache_tracker: list[ReferenceType | None] | None = None
+        args_hash: ArgsHash = tuple([id(arg) if type(arg) not in _primitive_types else arg for arg in args])
         try:
             mapping_cache_tracker = self._mapping_cache_tracker[args_hash]
         except KeyError:
@@ -79,13 +84,17 @@ def lookup(self, raise_on_templated_floats: bool, args: tuple[Any, ...]) -> tupl
         except KeyError:
             count = self.mapping[key] = len(self.mapping)
 
-        mapping_cache_tracker_: list[ReferenceType] = []
+        # Note that it is important to prepend the cache tracker with 'None' to avoid misclassifying no argument with
+        # expired cache entry caused by deallocated argument.
+        mapping_cache_tracker_: list[ReferenceType | None] = [None]
         clear_callback = lambda ref: mapping_cache_tracker_.clear()
         try:
-            mapping_cache_tracker_ += [ReferenceType(arg, clear_callback) for arg in args]
+            mapping_cache_tracker_ += [
+                ReferenceType(arg, clear_callback) for arg in args if type(arg) not in _primitive_types
+            ]
             self._mapping_cache_tracker[args_hash] = mapping_cache_tracker_
             self._mapping_cache[args_hash] = (count, key)
-        except TypeError:
-            pass
+        except TypeError as e:
+            warnings_helper.warn_once(f"{e}. Template mapper caching disabled.")
 
         return (count, key)

python/gstaichi/lang/_template_mapper_hotpath.py

@@ -89,10 +89,10 @@ def _extract_arg(raise_on_templated_floats: bool, arg: Any, annotation: Annotati
             # TODO(zhanlue): replacing "tuple(args)" with "hash of argument values"
             # This can resolve the following issues:
             # 1. Invalid weak-ref will leave a dead(dangling) entry in both caches: "self.mapping" and "self.compiled_functions"
-            # 2. Different argument instances with same type and same value, will get templatized into seperate kernels.
+            # 2. Different argument instances with same type and same value, will get templatized into separate kernels.
             return weakref.ref(arg)
 
-        # [Primitive arguments] Return the value
+        # Return value directly for other types, i.e. primitive types and all ti.Field-derived classes
         if raise_on_templated_floats and arg_type is float:
             raise ValueError("Floats not allowed as templated types.")
         return arg

python/gstaichi/lang/impl.py

@@ -519,10 +519,10 @@ def materialize(self):
         self._check_gradient_field_not_placed("dual")
         self._check_matrix_field_member_shape()
         self._calc_matrix_field_dynamic_index_stride()
-        self.global_vars = []
-        self.grad_vars = []
-        self.dual_vars = []
-        self.matrix_fields = []
+        self.global_vars.clear()
+        self.grad_vars.clear()
+        self.dual_vars.clear()
+        self.matrix_fields.clear()
 
     def _register_signal_handlers(self):
         if self._signal_handler_registry is None:

python/gstaichi/lang/kernel.py

@@ -95,7 +95,7 @@ def __init__(self) -> None:
         #   the launch context being stored in cache.
         # See 'launch_kernel' for details regarding the intended use of caching.
         self._launch_ctx_cache: dict[ArgsHash, KernelLaunchContext] = {}
-        self._launch_ctx_cache_tracker: dict[ArgsHash, list[ReferenceType]] = {}
+        self._launch_ctx_cache_tracker: dict[ArgsHash, list[ReferenceType | None]] = {}
         self._prog_weakref: ReferenceType[Program] | None = None
 
     def _destroy_callback(self, ref: ReferenceType):
@@ -113,11 +113,12 @@ def cache(
         self._launch_ctx_cache[args_hash] = cached_launch_ctx
 
         # Note that the clearing callback will only be called once despite being registered for each tracked
-        # objects, because all the weakrefs get deallocated right away, and their respective callback
-        # vanishes with them, without even getting a chance to get called. This means that registring the
-        # clearing callback systematically does not incur any cumulative runtime penalty yet ensures full
-        # memory safety.
-        launch_ctx_cache_tracker_: list[ReferenceType] = []
+        # objects, because all the weakrefs get deallocated right away, and their respective callback vanishes
+        # with them, without even getting a chance to get called. This means that registring the clearing
+        # callback systematically does not incur any cumulative runtime penalty yet ensures full memory safety.
+        # Note that it is important to prepend the cache tracker with 'None' to avoid misclassifying no argument
+        # with expired cache entry caused by deallocated argument.
+        launch_ctx_cache_tracker_: list[ReferenceType | None] = []
         clear_callback = lambda ref: launch_ctx_cache_tracker_.clear()
         if launch_ctx_args := launch_ctx_buffer.get(_TI_ARRAY):
             _, arrs = zip(*launch_ctx_args)
@@ -400,7 +401,7 @@ def launch_kernel(self, t_kernel: KernelCxx, compiled_kernel_data: CompiledKerne
         assert len(args) == len(self.arg_metas), f"{len(self.arg_metas)} arguments needed but {len(args)} provided"
 
         callbacks: list[Callable[[], None]] = []
-        args_hash: ArgsHash = tuple(map(id, args))
+        args_hash: ArgsHash = (id(t_kernel), *[id(arg) for arg in args if type(arg) is not template])
         launch_ctx = t_kernel.make_launch_context()
         prog, _populated_launch_ctx = self.launch_context_buffer_cache.populate_launch_ctx_from_cache(
             launch_ctx, args_hash

tests/python/test_cache.py

@@ -0,0 +1,78 @@
+import gstaichi as ti
+from gstaichi.lang.misc import get_host_arch_list
+
+from tests import test_utils
+
+
+@test_utils.test(arch=get_host_arch_list())
+def test_cache_primitive_args():
+    @ti.data_oriented
+    class StructStaticConfig:
+        flag_1: bool = True
+
+    @ti.kernel
+    def fun(static_args: ti.template(), flag_2: ti.template(), value: ti.types.ndarray()):
+        if ti.static(static_args.flag_1):
+            if ti.static(flag_2):
+                value[None] = value[None] + 1
+            else:
+                assert "Invalid 'flag_2' branch"
+        else:
+            assert "Invalid 'static_args.flag_1' branch"
+
+    assert len(fun._primal.mapper._mapping_cache) == 0
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 0
+    assert len(fun._primal._launch_ctx_cache) == 0
+    assert len(fun._primal._launch_ctx_cache_tracker) == 0
+
+    static_args = StructStaticConfig()
+    flag_2 = True
+    value = ti.ndarray(ti.i32, shape=())
+    value[None] = 1
+
+    fun(static_args, flag_2, value)
+    assert value[None] == 2
+    assert len(fun._primal.mapper._mapping_cache) == 1
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 1
+    assert len(fun._primal._launch_ctx_cache) == 1
+    assert len(fun._primal._launch_ctx_cache_tracker) == 1
+
+    fun(static_args, flag_2, value)
+    assert value[None] == 3
+    assert len(fun._primal.mapper._mapping_cache) == 1
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 1
+    assert len(fun._primal._launch_ctx_cache) == 1
+    assert len(fun._primal._launch_ctx_cache_tracker) == 1
+
+
+@test_utils.test(arch=get_host_arch_list())
+def test_cache_fields_only():
+    @ti.kernel
+    def fun(flag: ti.template(), value: ti.template()):
+        if ti.static(flag):
+            value[None] = value[None] + 1
+        else:
+            assert "Invalid 'static_args.flag_1' branch"
+
+    assert len(fun._primal.mapper._mapping_cache) == 0
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 0
+    assert len(fun._primal._launch_ctx_cache) == 0
+    assert len(fun._primal._launch_ctx_cache_tracker) == 0
+
+    flag = True
+    value = ti.field(ti.i32, shape=())
+    value[None] = 1
+
+    fun(flag, value)
+    assert value[None] == 2
+    assert len(fun._primal.mapper._mapping_cache) == 1
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 1
+    assert len(fun._primal._launch_ctx_cache) == 1
+    assert len(fun._primal._launch_ctx_cache_tracker) == 1
+
+    fun(flag, value)
+    assert value[None] == 3
+    assert len(fun._primal.mapper._mapping_cache) == 1
+    assert len(fun._primal.mapper._mapping_cache_tracker) == 1
+    assert len(fun._primal._launch_ctx_cache) == 1
+    assert len(fun._primal._launch_ctx_cache_tracker) == 1