Recompile stale BASE_DIR

pykokkos/core/compiler.py

@@ -275,8 +275,9 @@ def compile_entity(
             main,
             module_setup.metadata,
             space,
-            module_setup.types_signature,
-            module_setup.restrict_signature,
+            module_setup.ast_signature,
+            types_signature=module_setup.types_signature,
+            restrict_signature=module_setup.restrict_signature,
         )
         c_start: float = time.perf_counter()
         cpp_setup.compile(

pykokkos/core/cpp_setup.py

@@ -117,7 +117,8 @@ def compile(
 
     def initialize_directory(self, name: Path) -> None:
         """
-        Creates an output directory, overwriting an existing directory with the same name
+        Creates an output directory, overwriting an existing directory with the same name.
+        Checks if an older AST directory exists, and if so copies over the relevant information.
 
         :param name: the name of the directory
         """
@@ -127,6 +128,31 @@ def initialize_directory(self, name: Path) -> None:
         except OSError:
             pass
 
+        # make the parent directory if necessary
+        try:
+            os.makedirs(name.parent, exist_ok=True)
+        except FileExistsError:
+            pass
+
+        # check if an older AST library exists
+        old_dir = [
+            f for f in name.parent.parent.iterdir() if f.is_dir() and f != name.parent
+        ]
+        if len(old_dir) == 1:
+            # if an older AST directory exists, change the name to the new AST directory
+            old_dir = old_dir[0]
+            old_dir.rename(name.parent)
+            # CMakeCache.txt has hardcoded absolute paths to the old AST directory.
+            # Delete only the cache files so CMake reconfigures with correct paths
+            # while keeping object files in build/ for incremental recompilation.
+            stale_cache = name / "build" / "CMakeCache.txt"
+            if stale_cache.is_file():
+                stale_cache.unlink()
+            stale_cache_check = name / "build" / "CMakeFiles" / "cmake.check_cache"
+            if stale_cache_check.is_file():
+                stale_cache_check.unlink()
+
+        # make the name directory if necessary
         try:
             os.makedirs(name, exist_ok=True)
         except FileExistsError:

pykokkos/core/module_setup.py

@@ -82,13 +82,16 @@ def __init__(
         self,
         entity: Union[Callable[..., None], type, List[Callable[..., None]]],
         space: ExecutionSpace,
+        ast_signature: str,
+        *,
         types_signature: Optional[str] = None,
         restricted_views: Optional[Set[str]] = None,
     ):
         """
         ModuleSetup constructor
 
         :param entity: the functor/workunit/workload or list of workunits for fusion
+        :param ast_signature: hash/string to identify workunit signature against AST
         :param types_signature: hash/string to identify workunit signature against types
         :param restricted_views: a set of view names that do not alias any other views
         """
@@ -103,6 +106,7 @@ def __init__(
             self.metadata = [get_metadata(entity)]
 
         self.space: ExecutionSpace = space
+        self.ast_signature = ast_signature
         self.types_signature = types_signature
         self.restrict_signature: Optional[str] = None
         if restricted_views is not None:
@@ -126,7 +130,12 @@ def __init__(
             self.main: Path = self.get_main_path()
 
         self.output_dir: Optional[Path] = self.get_output_dir(
-            self.main, self.metadata, space, types_signature, self.restrict_signature
+            self.main,
+            self.metadata,
+            space,
+            ast_signature,
+            types_signature=types_signature,
+            restrict_signature=self.restrict_signature,
         )
         self.gpu_module_files: List[str] = []
         if km.is_multi_gpu_enabled():
@@ -154,6 +163,8 @@ def get_output_dir(
         main: Path,
         metadata: List[EntityMetadata],
         space: ExecutionSpace,
+        ast_signature,
+        *,
         types_signature: Optional[str] = None,
         restrict_signature: Optional[str] = None,
     ) -> Optional[Path]:
@@ -163,6 +174,7 @@ def get_output_dir(
         :param main: the path to the main file in the current PyKokkos application
         :param metadata: the metadata of the entity or fused entities being compiled
         :param space: the execution space to compile for
+        :param ast_signature: hash/string to identify workunit signature against AST
         :param types_signature: optional identifier/hash string for types of parameters
         :param restrict_signature: optional identifier/hash string from the views that do not alias any other views
         :returns: the path to the output directory for a specific execution space
@@ -180,6 +192,8 @@ def get_output_dir(
             out_dir = out_dir / f"types_{types_signature}"
         if restrict_signature is not None:
             out_dir = out_dir / f"restrict_{restrict_signature}"
+        if ast_signature is not None:
+            out_dir = out_dir / f"AST_{ast_signature}"
 
         out_dir = out_dir / space.value
 
@@ -258,7 +272,8 @@ def is_compiled(self) -> bool:
                 self.main,
                 self.metadata,
                 self.space,
-                self.types_signature,
-                self.restrict_signature,
+                self.ast_signature,
+                types_signature=self.types_signature,
+                restrict_signature=self.restrict_signature,
             )
         )

pykokkos/core/parsers/parser.py

@@ -1,4 +1,5 @@
 import ast
+import hashlib
 from dataclasses import dataclass
 from enum import Enum, auto
 from typing import Callable, Dict, List, Optional, Tuple, Union
@@ -79,6 +80,10 @@ def __init__(self, path: Optional[str], pk_import: Optional[str] = None):
             self.functors: Dict[str, PyKokkosEntity] = {}
             self.workunits: Dict[str, PyKokkosEntity] = {}
 
+        # get parser signature
+        signature = ast.dump(self.tree)
+        self.signature = hashlib.md5(signature.encode()).hexdigest()
+
     def get_import(self) -> str:
         """
         Get the pykokkos import identifier

pykokkos/core/runtime.py

@@ -5,6 +5,7 @@
 import sys
 from typing import Any, Callable, Dict, Optional, Set, Tuple, Type, Union, List
 import sysconfig
+import hashlib
 
 import numpy as np
 
@@ -174,6 +175,7 @@ def run_workload(self, space: ExecutionSpace, workload: object) -> None:
     def precompile_workunit(
         self,
         workunit: Callable[..., None],
+        ast_signature: str,
         space: ExecutionSpace,
         updated_decorator: Optional[UpdatedDecorator],
         updated_types: Optional[UpdatedTypes],
@@ -186,15 +188,21 @@ def precompile_workunit(
         precompile the workunit
 
         :param workunit: the workunit function object
+        :param ast_signature: Hash/identifer string for workunit module against AST
         :param space: the ExecutionSpace for which the bindings are generated
         :param updated_decorator: Object for decorator specifier
         :param updated_types: Object with type inference information
+        :param types_signature: Hash/identifer string for workunit module against data types
         :param restrict_views: a set of view names that do not alias any other views
         :returns: the members the functor is containing
         """
 
         module_setup: ModuleSetup = self.get_module_setup(
-            workunit, space, types_signature, restrict_signature
+            workunit,
+            space,
+            ast_signature,
+            types_signature=types_signature,
+            restrict_signature=restrict_signature,
         )
         members: PyKokkosMembers = self.compiler.compile_object(
             module_setup,
@@ -342,9 +350,17 @@ def execute_workunit(
             }
             restrict_views, restrict_signature = get_restrict_views(view_dict)
 
+        # Set ast signature
+        if isinstance(parser, list):
+            ast_signature = "".join([p.signature for p in parser])
+            ast_signature = hashlib.md5(ast_signature.encode()).hexdigest()
+        else:
+            ast_signature = parser.signature
+
         execution_space: ExecutionSpace = policy.space.space
         members: PyKokkosMembers = self.precompile_workunit(
             workunit,
+            ast_signature,
             execution_space,
             updated_decorator,
             updated_types,
@@ -355,7 +371,11 @@ def execute_workunit(
         )
 
         module_setup: ModuleSetup = self.get_module_setup(
-            workunit, execution_space, types_signature, restrict_signature
+            workunit,
+            execution_space,
+            ast_signature,
+            types_signature=types_signature,
+            restrict_signature=restrict_signature,
         )
         return self.execute(
             workunit,
@@ -812,6 +832,8 @@ def get_module_setup(
         self,
         entity: Union[object, Callable[..., None]],
         space: ExecutionSpace,
+        ast_signature: str,
+        *,
         types_signature: Optional[str] = None,
         restrict_signature: Optional[str] = None,
     ) -> ModuleSetup:
@@ -820,6 +842,7 @@ def get_module_setup(
 
         :param entity: the workload or workunit object
         :param space: the execution space
+        :param ast_signature: Hash/identifer string for workunit module against AST
         :param types_signature: Hash/identifer string for workunit module against data types
         :param restrict_signature: Hash/identifer string for views that do not alias any other views
         :returns: the ModuleSetup object
@@ -830,13 +853,23 @@ def get_module_setup(
         )
 
         module_setup_id = self.get_module_setup_id(
-            entity, space, types_signature, restrict_signature
+            entity,
+            space,
+            ast_signature,
+            types_signature=types_signature,
+            restrict_signature=restrict_signature,
         )
 
         if module_setup_id in self.module_setups:
             return self.module_setups[module_setup_id]
 
-        module_setup = ModuleSetup(entity, space, types_signature, restrict_signature)
+        module_setup = ModuleSetup(
+            entity,
+            space,
+            ast_signature,
+            types_signature=types_signature,
+            restricted_views=restrict_signature,
+        )
         self.module_setups[module_setup_id] = module_setup
 
         return module_setup
@@ -845,6 +878,8 @@ def get_module_setup_id(
         self,
         entity: Union[object, Callable[..., None]],
         space: ExecutionSpace,
+        ast_signature: str,
+        *,
         types_signature: Optional[str] = None,
         restrict_signature: Optional[str] = None,
     ) -> Tuple:
@@ -856,6 +891,7 @@ def get_module_setup_id(
 
         :param entity: the workload or workunit object
         :param space: the execution space
+        :param ast_signature: Hash/identifer string for workunit module against AST
         :param types_signature: optional identifier/hash string for
             types of parameters against workunit module
         :param restrict_signature: Hash/identifer string for views
@@ -890,6 +926,7 @@ def get_module_setup_id(
                 module_setup_id_list.append(types_signature)
             if restrict_signature is not None:
                 module_setup_id_list.append(restrict_signature)
+            module_setup_id_list.append(ast_signature)
 
             module_setup_id = tuple(module_setup_id_list)
 

tests/test_recompile.py

@@ -0,0 +1,98 @@
+"""
+pytest: PyKokkos JIT kernel recompilation test.
+
+Tests that modifying a @pk.workunit source file triggers recompilation
+and that both versions of the kernel produce correct results.
+
+Workflow:
+  1. Write the "buggy" kernel (arr[i] += 2) to a temp module file.
+  2. Import it and run it — verify each element increased by 2.
+  3. Overwrite the module file with the "fixed" kernel (arr[i] += 1).
+  4. Invalidate Python's module cache so the new source is picked up.
+  5. Re-import and run — verify each element increased by 1.
+"""
+
+import sys
+import textwrap
+import importlib
+from pathlib import Path
+
+import numpy as np
+import pykokkos as pk
+
+BUGGY_SOURCE = textwrap.dedent(
+    """\
+    import pykokkos as pk
+
+    @pk.workunit
+    def add1(i: int, arr: pk.View1D[int]):
+        arr[i] += 2
+"""
+)
+
+CORRECT_SOURCE = textwrap.dedent(
+    """\
+    import pykokkos as pk
+
+    @pk.workunit
+    def add1(i: int, arr: pk.View1D[int]):
+        arr[i] += 1
+"""
+)
+
+MODULE_NAME = "_test_jit_kernel_add1"
+
+
+def _load_fresh(path: Path):
+    """Force a fresh import of the module at *path*, bypassing sys.modules."""
+    # Remove any previously cached version.
+    sys.modules.pop(MODULE_NAME, None)
+
+    spec = importlib.util.spec_from_file_location(MODULE_NAME, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[MODULE_NAME] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+# ---------- Test JIT Recompilation ----------
+
+
+def test_recompilation(tmp_path):
+    kernel_file = tmp_path / "_test_jit_recompile.py"
+
+    # ---- Stage 1: buggy kernel
+    # write buggy source
+    kernel_file.write_text(BUGGY_SOURCE, encoding="utf-8")
+
+    # load buggy source
+    mod_buggy = _load_fresh(kernel_file)
+
+    # run the buggy kernel
+    n = 5
+    arr_buggy = np.zeros(n, dtype=np.int32)
+    pk.parallel_for(n, mod_buggy.add1, arr=arr_buggy)
+
+    # assert buggy array is correct
+    try:
+        np.testing.assert_equal(arr_buggy, np.zeros(n, dtype=np.int32) + 2)
+    except AssertionError as e:
+        raise AssertionError("buggy kernel is incorrect") from e
+
+    # ---- Stage 2: correct kernel
+
+    # reload pykokkos to clear cache
+    importlib.reload(sys.modules["pykokkos"])
+
+    kernel_file.write_text(CORRECT_SOURCE, encoding="utf-8")
+    mod_correct = _load_fresh(kernel_file)
+
+    arr_correct = np.zeros(n, dtype=np.int32)
+    pk.parallel_for(n, mod_correct.add1, arr=arr_correct)
+    expected = np.zeros(n, dtype=np.int32) + 1
+    try:
+        np.testing.assert_equal(arr_correct, expected)
+    except AssertionError as e:
+        raise AssertionError(
+            f"kernel is incorrect\nactual:  {arr_correct}\ndesired: {expected}"
+        ) from e