tilelang.cache.kernel_cache¶
The cache utils with class and database persistence - KernelCache Class
Attributes¶
Classes¶
Caches compiled kernels using a class and database persistence to avoid redundant compilation. |
Module Contents¶
- tilelang.cache.kernel_cache.KERNEL_PATH = 'kernel.cu'¶
- tilelang.cache.kernel_cache.WRAPPED_KERNEL_PATH = 'wrapped_kernel.cu'¶
- tilelang.cache.kernel_cache.KERNEL_LIB_PATH = 'kernel_lib.so'¶
- tilelang.cache.kernel_cache.KERNEL_CUBIN_PATH = 'kernel.cubin'¶
- tilelang.cache.kernel_cache.KERNEL_PY_PATH = 'kernel.py'¶
- tilelang.cache.kernel_cache.PARAMS_PATH = 'params.pkl'¶
- class tilelang.cache.kernel_cache.KernelCache¶
Caches compiled kernels using a class and database persistence to avoid redundant compilation. Cache files:
kernel.cu: The compiled kernel source code wrapped_kernel.cu: The compiled wrapped kernel source code kernel_lib.so: The compiled kernel library params.pkl: The compiled kernel parameters
- execution_backend: Literal['dlpack', 'ctypes', 'cython', 'nvrtc'] = 'cython'¶
- cached(func=None, out_idx=None, *args, target='auto', target_host=None, execution_backend='cython', verbose=False, pass_configs=None, compile_flags=None)¶
Caches and reuses compiled kernels to avoid redundant compilation.
- Parameters:
func (tvm.tir.PrimFunc) – Function to be compiled or a prepared PrimFunc
out_idx (List[int]) – Indices specifying which outputs to return
target (Union[str, tvm.target.Target]) – Compilation target platform
target_host (Union[str, tvm.target.Target]) – Host target platform
*args – Arguments passed to func
execution_backend (Literal['dlpack', 'ctypes', 'cython', 'nvrtc'])
verbose (bool)
pass_configs (dict)
compile_flags (Optional[List[str]])
- Returns:
The compiled kernel, either freshly compiled or from cache
- Return type:
- clear_cache()¶
Clears the entire kernel cache, including both in-memory and disk cache.