[Enhancement] Include compile flags into the hash key of cached kernels (#911)

* [Cache] Add compile_flags parameter to KernelCache hash keys

* [Cache] Update compile_flags parameter to accept both List[str] and str types

* lint

* [Refactor] Update compile_flags parameter to accept Union[List[str], str] type

tilelang/cache/__init__.py

@@ -20,7 +20,7 @@ def cached(
     execution_backend: Optional[Literal["dlpack", "ctypes", "cython", "nvrtc"]] = "cython",
     verbose: Optional[bool] = False,
     pass_configs: Optional[dict] = None,
-    compile_flags: Optional[List[str]] = None,
+    compile_flags: Optional[Union[List[str], str]] = None,
 ) -> JITKernel:
     """
     Caches and reuses compiled kernels (using KernelCache class).

tilelang/cache/kernel_cache.py

@@ -73,6 +73,7 @@ class KernelCache:
         target: Union[str, Target] = "auto",
         target_host: Union[str, Target] = None,
         pass_configs: dict = None,
+        compile_flags: Optional[Union[List[str], str]] = None,
     ) -> str:
         """
         Generates a unique hash key for caching compiled kernels.
@@ -101,6 +102,7 @@ class KernelCache:
             "target_host": str(target_host) if target_host else None,
             "execution_backend": execution_backend,
             "pass_configs": pass_configs,
+            "compile_flags": compile_flags,
         }
         # Sort keys to ensure consistency
         key_string = json.dumps(key_data, sort_keys=True)
@@ -117,7 +119,7 @@ class KernelCache:
         execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
         verbose: bool = False,
         pass_configs: dict = None,
-        compile_flags: Optional[List[str]] = None,
+        compile_flags: Optional[Union[List[str], str]] = None,
     ) -> JITKernel:
         """
         Caches and reuses compiled kernels to avoid redundant compilation.
@@ -152,6 +154,7 @@ class KernelCache:
             target=target,
             target_host=target_host,
             pass_configs=pass_configs,
+            compile_flags=compile_flags,
         )
         with self._lock:
             # First check in-memory cache
@@ -165,7 +168,8 @@ class KernelCache:
 
             # Then check disk cache
             kernel = self._load_kernel_from_disk(key, target, target_host, out_idx,
-                                                 execution_backend, pass_configs, func, verbose)
+                                                 execution_backend, pass_configs, compile_flags,
+                                                 func, verbose)
             if kernel is not None:
                 if verbose:
                     self.logger.debug(
@@ -185,6 +189,7 @@ class KernelCache:
             target_host=target_host,
             verbose=verbose,
             pass_configs=pass_configs,
+            compile_flags=compile_flags,
         )
         if execution_backend == "dlpack":
             self.logger.warning("DLPack backend does not support cache saving to disk.")
@@ -322,6 +327,7 @@ class KernelCache:
         out_idx: List[int] = None,
         execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
         pass_configs: dict = None,
+        compile_flags: Optional[Union[List[str], str]] = None,
         func: Callable = None,
         verbose: bool = False,
     ) -> Optional[JITKernel]:
@@ -382,6 +388,7 @@ class KernelCache:
                 out_idx=out_idx,
                 execution_backend=execution_backend,
                 pass_configs=pass_configs,
+                compile_flags=compile_flags,
             )
         else:
             return None

tilelang/jit/__init__.py

@@ -95,7 +95,7 @@ class _JitImplementation:
     verbose: bool
     pass_configs: Optional[Dict[str, Any]]
     debug_root_path: Optional[str]
-    compile_flags: Optional[List[str]]
+    compile_flags: Optional[Union[List[str], str]]
 
     def __init__(self,
                  out_idx: Any = None,
@@ -105,7 +105,7 @@ class _JitImplementation:
                  verbose: bool = False,
                  pass_configs: Optional[Dict[str, Any]] = None,
                  debug_root_path: Optional[str] = None,
-                 compile_flags: Optional[List[str]] = None):
+                 compile_flags: Optional[Union[List[str], str]] = None):
         """
         Initializes the JIT compiler decorator.
 
@@ -137,6 +137,9 @@ class _JitImplementation:
             If None, no debug information is saved (default: None).
             If a relative path is given, it's made absolute relative to the project root
             or current working directory.
+        compile_flags : Optional[Union[List[str], str]], optional
+            Additional compilation flags to pass to the compiler.
+            If None, no additional compilation flags are passed (default: None).
         """
         self.out_idx = out_idx
         self.execution_backend = execution_backend