Skip to content

GitLab

"...composable_kernel.git" did not exist on "d072790fe27a5915d12b9172020b6b74e485aadf"
Unverified Commit 042c60fb authored by Yang Chen's avatar Yang Chen Committed by GitHub
Browse files

[Enhancement] Output cache-file-related messages with verbose=True (#683) 

This is a minor enhancement to output verbose messages indicating where
cache files are saved and loaded. These messages are useful for
examining the relevant intermediate files.
parent eb026b79
Hide whitespace changes
Inline Side-by-side
Showing with 60 additions and 8 deletions
tilelang/autotuner/param.py
View file @ 042c60fb
......@@ -149,14 +149,14 @@ class AutotuneResult:
func: Optional[Callable] = None
kernel: Optional[Callable] = None
def _save_kernel_to_disk(self, cache_path: Path, kernel: JITKernel):
def _save_kernel_to_disk(self, cache_path: Path, kernel: JITKernel, verbose: bool = False):
"""
Persists a compiled kernel to disk cache.
Args:
key (str): The hash key identifying the kernel.
cache_path (Path): The root path for the cache files.
kernel (JITKernel): The compiled kernel to be saved.
func (Callable, optional): The original function.
verbose (bool): Enable verbose log messages.
Note:
Saves the following files:
......@@ -170,6 +170,8 @@ class AutotuneResult:
# Save kernel source code
try:
kernel_path = os.path.join(cache_path, KERNEL_PATH)
if verbose:
logger.debug(f"Saving kernel source code to file: {kernel_path}")
if kernel.artifact.kernel_source is not None:
with open(kernel_path, "w") as f:
f.write(kernel.artifact.kernel_source)
......@@ -179,6 +181,8 @@ class AutotuneResult:
# Save wrapped kernel source code
try:
wrapped_kernel_path = os.path.join(cache_path, WRAPPED_KERNEL_PATH)
if verbose:
logger.debug(f"Saving wrapped kernel source code to file: {wrapped_kernel_path}")
with open(wrapped_kernel_path, "w") as f:
f.write(kernel.get_kernel_source())
except Exception as e:
......@@ -188,6 +192,8 @@ class AutotuneResult:
try:
kernel_lib_path = os.path.join(cache_path, KERNEL_LIB_PATH)
src_lib_path = kernel.adapter.libpath
if verbose:
logger.debug(f"Saving kernel library to file: {kernel_lib_path}")
shutil.copy(src_lib_path, kernel_lib_path)
except Exception as e:
logger.error(f"Error saving kernel library to disk: {e}")
......@@ -195,6 +201,8 @@ class AutotuneResult:
# Save kernel parameters
try:
params_path = os.path.join(cache_path, PARAMS_PATH)
if verbose:
logger.debug(f"Saving kernel parameters to disk: {params_path}")
with open(params_path, "wb") as f:
cloudpickle.dump(kernel.params, f)
except Exception as e:
......@@ -209,6 +217,7 @@ class AutotuneResult:
execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython",
pass_configs: dict = None,
func: Callable = None,
verbose: bool = False,
) -> JITKernel:
"""
Loads a previously compiled kernel from disk cache.
......@@ -221,6 +230,7 @@ class AutotuneResult:
execution_backend (Literal): Backend type for execution. Defaults to "cython".
pass_configs (dict, optional): Configuration for compiler passes.
func (Callable, optional): The original function.
verbose (bool): Enable verbose log messages.
Returns:
JITKernel: The loaded kernel if found, None otherwise.
......@@ -234,6 +244,8 @@ class AutotuneResult:
try:
wrapped_kernel_path = os.path.join(cache_path, WRAPPED_KERNEL_PATH)
if verbose:
logger.debug(f"Loading wrapped kernel source code from file: {wrapped_kernel_path}")
with open(wrapped_kernel_path, "r") as f:
kernel_global_source = f.read()
except Exception as e:
......@@ -244,6 +256,8 @@ class AutotuneResult:
# Load kernel parameters
try:
params_path = os.path.join(cache_path, PARAMS_PATH)
if verbose:
logger.debug(f"Loading kernel parameters from file: {params_path}")
with open(params_path, "rb") as f:
kernel_params = cloudpickle.load(f)
except Exception as e:
......@@ -264,19 +278,25 @@ class AutotuneResult:
else:
return None
def save_to_disk(self, path: Path):
def save_to_disk(self, path: Path, verbose: bool = False):
if not os.path.exists(path):
os.makedirs(path)
# save best config
if verbose:
logger.debug(f"Saving best config to file: {path / BEST_CONFIG_PATH}")
with open(path / BEST_CONFIG_PATH, "w") as f:
json.dump(self.config, f)
# save function
if verbose:
logger.debug(f"Saving function to file: {path / FUNCTION_PATH}")
with open(path / FUNCTION_PATH, "wb") as f:
cloudpickle.dump(self.func, f)
# save ref latency
if verbose:
logger.debug(f"Saving latency to file: {path / LATENCY_PATH}")
with open(path / LATENCY_PATH, "w") as f:
json.dump({
"latency": self.latency,
......@@ -291,15 +311,22 @@ class AutotuneResult:
if not os.path.exists(path):
return None
verbose = compile_args.verbose
# load best config
if verbose:
logger.debug(f"Loading best config from file: {path / BEST_CONFIG_PATH}")
with open(path / BEST_CONFIG_PATH, "r") as f:
config = json.load(f)
# load function
if verbose:
logger.debug(f"Loading function from file: {path / FUNCTION_PATH}")
with open(path / FUNCTION_PATH, "rb") as f:
func = cloudpickle.load(f)
# load latency
if verbose:
logger.debug(f"Loading latency from file: {path / LATENCY_PATH}")
with open(path / LATENCY_PATH, "r") as f:
latency = json.load(f)
latency, ref_latency = latency["latency"], latency["ref_latency"]
......
tilelang/autotuner/tuner.py
View file @ 042c60fb
......@@ -257,7 +257,7 @@ class AutoTuner:
return hashlib.sha256(key_string.encode()).hexdigest()
def _save_result_to_disk(self, key, result: AutotuneResult):
result.save_to_disk(self.cache_dir / key)
result.save_to_disk(self.cache_dir / key, self.compile_args.verbose)
def _load_result_from_disk(self, key) -> AutotuneResult:
result = AutotuneResult.load_from_disk(self.cache_dir / key, self.compile_args)
......
tilelang/cache/kernel_cache.py
View file @ 042c60fb
......@@ -165,7 +165,7 @@ class KernelCache:
# Then check disk cache
kernel = self._load_kernel_from_disk(key, target, target_host, out_idx,
execution_backend, pass_configs, func)
execution_backend, pass_configs, func, verbose)
if kernel is not None:
if verbose:
self.logger.debug(
......@@ -174,6 +174,8 @@ class KernelCache:
self._memory_cache[key] = kernel
return kernel
if verbose:
self.logger.debug(f"No cached kernel for {func.attrs['global_symbol']}")
# Compile kernel if cache miss; leave critical section
kernel = JITKernel(
func,
......@@ -189,7 +191,7 @@ class KernelCache:
else:
with self._lock:
if is_cache_enabled():
self._save_kernel_to_disk(key, kernel, func)
self._save_kernel_to_disk(key, kernel, func, verbose)
# Store in memory cache after compilation
self._memory_cache[key] = kernel
......@@ -231,7 +233,11 @@ class KernelCache:
# Use atomic POSIX replace, so other processes cannot see a partial write
os.replace(temp_path, path)
def _save_kernel_to_disk(self, key: str, kernel: JITKernel, func: Callable = None):
def _save_kernel_to_disk(self,
key: str,
kernel: JITKernel,
func: Callable = None,
verbose: bool = False):
"""
Persists a compiled kernel to disk cache.
......@@ -239,6 +245,7 @@ class KernelCache:
key (str): The hash key identifying the kernel.
kernel (JITKernel): The compiled kernel to be saved.
func (Callable, optional): The original function.
verbose (bool): Enable verbose log messages.
Note:
Saves the following files:
......@@ -253,6 +260,8 @@ class KernelCache:
# Save kernel source code
try:
kernel_path = os.path.join(cache_path, KERNEL_PATH)
if verbose:
self.logger.debug(f"Saving kernel source code to file: {kernel_path}")
if kernel.artifact.kernel_source is not None:
KernelCache._safe_write_file(kernel_path, "w",
lambda file: file.write(kernel.artifact.kernel_source))
......@@ -262,6 +271,9 @@ class KernelCache:
# Save wrapped kernel source code
try:
wrapped_kernel_path = os.path.join(cache_path, WRAPPED_KERNEL_PATH)
if verbose:
self.logger.debug(
f"Saving wrapped kernel source code to file: {wrapped_kernel_path}")
KernelCache._safe_write_file(
wrapped_kernel_path, "w",
lambda file: file.write(kernel.adapter.get_kernel_source()))
......@@ -274,6 +286,8 @@ class KernelCache:
kernel_lib_path = KERNEL_CUBIN_PATH if self.execution_backend == "nvrtc" else KERNEL_LIB_PATH
kernel_lib_path = os.path.join(cache_path, kernel_lib_path)
src_lib_path = kernel.adapter.libpath
if verbose:
self.logger.debug(f"Saving kernel library to file: {kernel_lib_path}")
KernelCache._safe_write_file(
kernel_lib_path, "wb",
lambda file: file.write(KernelCache._load_binary(src_lib_path)))
......@@ -282,6 +296,8 @@ class KernelCache:
if self.execution_backend == "nvrtc":
kernel_py_path = os.path.join(cache_path, KERNEL_PY_PATH)
src_lib_path = src_lib_path.replace(".cubin", ".py")
if verbose:
self.logger.debug(f"Saving kernel nvrtc python code to file: {kernel_py_path}")
KernelCache._safe_write_file(
kernel_py_path, "wb",
lambda file: file.write(KernelCache._load_binary(src_lib_path)))
......@@ -291,6 +307,8 @@ class KernelCache:
# Save kernel parameters
try:
params_path = os.path.join(cache_path, PARAMS_PATH)
if verbose:
self.logger.debug(f"Saving kernel parameters to disk: {params_path}")
KernelCache._safe_write_file(params_path, "wb",
lambda file: cloudpickle.dump(kernel.params, file))
except Exception as e:
......@@ -305,6 +323,7 @@ class KernelCache:
execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
pass_configs: dict = None,
func: Callable = None,
verbose: bool = False,
) -> Optional[JITKernel]:
"""
Loads a previously compiled kernel from disk cache.
......@@ -317,6 +336,7 @@ class KernelCache:
execution_backend (Literal): Backend type for execution. Defaults to "cython".
pass_configs (dict, optional): Configuration for compiler passes.
func (Callable, optional): The original function.
verbose (bool): Enable verbose log messages.
Returns:
JITKernel: The loaded kernel if found, None otherwise.
......@@ -334,6 +354,9 @@ class KernelCache:
# Load the kernel source file (optional)
try:
if verbose:
self.logger.debug(
f"Loading wrapped kernel source code from file: {wrapped_kernel_path}")
with open(wrapped_kernel_path, "r") as f:
kernel_global_source = f.read()
except Exception as e:
......@@ -341,6 +364,8 @@ class KernelCache:
# Load kernel parameters
try:
if verbose:
self.logger.debug(f"Loading kernel parameters from file: {params_path}")
with open(params_path, "rb") as f:
kernel_params = cloudpickle.load(f)
except Exception as e:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment