[JIT] Track angle-bracket includes in source hash (#28273)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2026-06-15 02:58:14 -07:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent 2a33724c9b
commit 3b419f66da
+8 -2
View File
@@ -67,7 +67,8 @@ def _make_wrapper(tup: Tuple[str, str]) -> str:
return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
_LOCAL_INCLUDE_RE = re.compile(r'^\s*#\s*include\s+"([^"]+)"', re.MULTILINE)
_QUOTED_INCLUDE_RE = re.compile(r'^\s*#\s*include\s*"([^"]+)"', re.MULTILINE)
_ANGLE_INCLUDE_RE = re.compile(r"^\s*#\s*include\s*<(sgl_kernel/[^>]+)>", re.MULTILINE)
def _local_jit_source_hash(source_files: List[str]) -> str:
@@ -75,6 +76,7 @@ def _local_jit_source_hash(source_files: List[str]) -> str:
digest = hashlib.sha256()
seen: set[pathlib.Path] = set()
stack = [pathlib.Path(path).resolve() for path in source_files]
include_dir = KERNEL_PATH / "include"
while stack:
path = stack.pop()
@@ -95,10 +97,14 @@ def _local_jit_source_hash(source_files: List[str]) -> str:
digest.update(b"\0")
text = data.decode("utf-8", errors="ignore")
for include in _LOCAL_INCLUDE_RE.findall(text):
for include in _QUOTED_INCLUDE_RE.findall(text):
include_path = (path.parent / include).resolve()
if include_path.is_file():
stack.append(include_path)
for include in _ANGLE_INCLUDE_RE.findall(text):
include_path = (include_dir / include).resolve()
if include_path.is_file():
stack.append(include_path)
return digest.hexdigest()[:16]