From 5eaf812ed03b506c082c890ca029869210630c69 Mon Sep 17 00:00:00 2001
From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com>
Date: Mon, 13 Jul 2026 23:12:35 +0800
Subject: [PATCH] fix(jit): invalidate kernel cache on #included source changes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The JIT content hash covered only a kernel's top-level .hip plus include/mori.
Ops kernel .hip units #include src/ops/dispatch_combine/*.cpp and ep_common.hip,
so editing an included file left the hash unchanged and a stale .hsaco was
reused — the edit silently had no effect.

Hash the whole subsystem source tree instead of the single .hip, and include
.hip files when hashing directories so #included translation units are covered.
---
 python/mori/jit/cache.py | 5 +++--
 python/mori/jit/core.py  | 8 ++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/python/mori/jit/cache.py b/python/mori/jit/cache.py
index 92f9f0761..1d325bfac 100644
--- a/python/mori/jit/cache.py
+++ b/python/mori/jit/cache.py
@@ -42,7 +42,8 @@ def get_cache_root() -> Path:
 def _hash_tree(paths: list[Path]) -> str:
     """Compute a short content hash over files and directories.
 
-    For directories, all ``.hpp``, ``.h``, and ``.cpp`` files are included.
+    For directories, all ``.hpp``, ``.h``, ``.cpp``, and ``.hip`` files are
+    included, so ``#include``d translation units invalidate the cache too.
     """
     h = hashlib.sha256()
     for p in sorted(paths):
@@ -50,7 +51,7 @@ def _hash_tree(paths: list[Path]) -> str:
             h.update(p.read_bytes())
         elif p.is_dir():
             for f in sorted(p.rglob("*")):
-                if f.suffix in (".hpp", ".h", ".cpp"):
+                if f.suffix in (".hpp", ".h", ".cpp", ".hip"):
                     h.update(f.read_bytes())
     return h.hexdigest()[:12]
 
diff --git a/python/mori/jit/core.py b/python/mori/jit/core.py
index aa0028524..3d5c4cebd 100644
--- a/python/mori/jit/core.py
+++ b/python/mori/jit/core.py
@@ -495,7 +495,7 @@ def compile_genco(
     sub_kernels = _PARALLEL_KERNEL_GROUPS.get(kernel_name)
     if sub_kernels:
         source_paths = [
-            mori_root / "src" / "ops" / "kernels",
+            mori_root / "src" / "ops",
             mori_root / "include" / "mori",
         ]
         cache_dir = get_cache_dir(
@@ -544,7 +544,11 @@ def compile_genco(
     if not source.is_file():
         raise FileNotFoundError(f"Kernel source not found: {source}")
 
-    source_paths = [source, mori_root / "include" / "mori"]
+    # The .hip translation unit #includes sibling sources from its subsystem
+    # (e.g. ops kernels pull in src/ops/dispatch_combine/*), so hash the whole
+    # subsystem source tree; hashing only the top-level .hip reuses a stale
+    # .hsaco when an included file changes.
+    source_paths = [(mori_root / source_dir).parent, mori_root / "include" / "mori"]
     cache_dir = get_cache_dir(cfg.arch, source_paths, nic, profiler=profiler, ccqe=ccqe)
     hsaco_path = cache_dir / f"{kernel_name}.hsaco"