Recompile stale `BASE_DIR` by kennykos · Pull Request #415 · kokkos/pykokkos

kennykos · 2026-04-29T21:31:00Z

Features

1. AST-hashed compilation directories
pk_cpp output directories now include an AST_<hash> subdirectory that encodes a structural hash of the compiled workunit. On each invocation, PyKokkos computes the current AST hash and compares it against the on-disk directory name; if the AST hashes differ, recompilation is triggered. This replaces the previous approach which did not check if the compiled AST matched the current AST.

2. Incremental recompilation on AST change
When a workunit's AST hash changes, initialize_directory renames the old AST_<hash> directory to the new hash rather than discarding it. Only CMakeCache.txt and cmake.check_cache are deleted (as they contain hardcoded absolute paths that CMake rejects after a rename), while all object files and link artifacts in build/ are preserved. This allows CMake to perform an incremental rebuild rather than compiling from scratch.

3. Recompilation unit test
Adds test_recompilation which verifies end-to-end that modifying a @pk.workunit body triggers a correct recompile. The test writes a buggy kernel (arr[i] += 2), runs it and asserts the output, then overwrites it with the fixed kernel (arr[i] += 1) and asserts the recompiled output, confirming that stale builds are invalidated and that the new kernel produces correct results.

We add a hashing to the AST directly to check if the workunit needs to be recompiled. NOTE: Ideally, we could edit the functor.hpp directly so we do not need to bring in all of Kokkos again, but this approach works for now.

gliga · 2026-04-30T01:38:26Z


+        # get parser signature
+        signature = ast.dump(self.tree)
+        self.signature = hashlib.md5(signature.encode()).hexdigest()


how expensive is this? I assume our ast's for kernels are small that this should not be an issue?

I didn't do a performance test. This block is only done when the workunit is first seen by pykokkos, so it is a startup cost that will not add overhead during repeated execution.

In most cases the ast's should be relatively small; e.g., in my Ewald code, the slowest time for this block of code is 0.013 seconds.

gliga · 2026-04-30T01:44:12Z

+    except AssertionError as e:
+        raise AssertionError(
+            f"kernel is incorrect\nactual:  {arr_correct}\ndesired: {expected}"
+        ) from e


this test does not check that values of signature on disk changed or something was renamed

I kept it abstract in case we want to change in implementation details down the road (i.e., if we want to move away from renaming or change the directory structure).

gliga

Overall, good change. See comments.

gliga · 2026-04-30T01:44:56Z

I do not like the renaming part, but we can keep it and see when/if it breaks.

gliga · 2026-04-30T01:45:57Z

Are we ever going to change the name of that directory? It has been a while that we wanted to move to .pykokkos or we can use .pk

kennykos · 2026-04-30T02:36:33Z

Are we ever going to change the name of that directory? It has been a while that we wanted to move to .pykokkos or we can use .pk

Changed to .pykokkos in #416.

JBludau · 2026-04-30T17:23:42Z

+        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.


maybe we want to reverse the order of the comments?

Changed in ee7e739.

JBludau · 2026-04-30T17:24:47Z

+        # make the parent directory if necessary
+        try:
+            os.makedirs(name.parent, exist_ok=True)
+        except FileExistsError:
+            pass
+


why would we want to catch a FileExistsError if we pass exists_ok=true?

Removed in ee7e739, this was a vestigial code pattern.

JBludau · 2026-04-30T17:25:32Z

+        old_dir = [
+            f for f in name.parent.parent.iterdir() if f.is_dir() and f != name.parent
+        ]
+        if len(old_dir) == 1:


why a list if we just expect 0 or 1 entries?

I chose list comprehension as it seemed the most pythonic thing to do, added a comment to specify that we have 0 or 1 entries.

JBludau · 2026-04-30T17:31:09Z

+        if ast_signature is not None:
+            out_dir = out_dir / f"AST_{ast_signature}"


hmm so we have a hierarchy which has the ast signature as the highest priority ... but is there cases where the ast signature would not trigger but the others would?

Yes, consider the scenario:

workunit v0 is compiled with pk.float arrays;

an edit is made and we now have ATS has for workunit v1;

workunit v1 is compiled with pk.double arrays;

workunit v1 is called again with pk.float arrays. In this case, the type hash would trigger, but the ats hash would not as it is out of date.

JBludau · 2026-04-30T17:32:51Z

+    # assert add2 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


maybe not call it buggy but rather v1 and v2?

Addressed in 19b56c0.

* Change ordering of comments for clarity; * remove vestigial try/except blocks

kennykos · 2026-04-30T19:36:46Z

@JBludau I want to make sure you're aware of the new commit a810722.

kennykos added 5 commits April 29, 2026 11:26

Add recompilation test

e223c92

Ensure both workunits have the same type string

c979fb9

Recompile test passes

5c2e4e6

We add a hashing to the AST directly to check if the workunit needs to be recompiled. NOTE: Ideally, we could edit the functor.hpp directly so we do not need to bring in all of Kokkos again, but this approach works for now.

AST signature for fused workunits

9f79823

Copy over old AST directory for incremental recomp

d87b2c6

kennykos requested review from JBludau and gliga April 29, 2026 21:31

Format

4acad26

gliga reviewed Apr 30, 2026

View reviewed changes

Comment thread pykokkos/core/cpp_setup.py Outdated

gliga reviewed Apr 30, 2026

View reviewed changes

Comment thread tests/test_recompile.py Outdated

gliga reviewed Apr 30, 2026

View reviewed changes

gliga approved these changes Apr 30, 2026

View reviewed changes

Update comment to reflect AST hash check

7548ed5

kennykos changed the title ~~Recompile stale pk_cpp~~ Recompile stale BASE_DIR Apr 30, 2026

Fix comment to reflect correct assertion for add2

5c5e734

JBludau reviewed Apr 30, 2026

View reviewed changes

kennykos added 3 commits April 30, 2026 14:27

Address code review comments

ee7e739

* Change ordering of comments for clarity; * remove vestigial try/except blocks

Descriptive workunit names

19b56c0

Remove unnecessary rmtree in initialize directory.

a810722

kennykos requested a review from JBludau April 30, 2026 19:36

		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.

		if ast_signature is not None:
		out_dir = out_dir / f"AST_{ast_signature}"

Conversation

kennykos commented Apr 29, 2026

Features

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gliga left a comment

Choose a reason for hiding this comment

Uh oh!

gliga commented Apr 30, 2026

Uh oh!

gliga commented Apr 30, 2026

Uh oh!

kennykos commented Apr 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kennykos commented Apr 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants