File size: 6,518 Bytes
d934615 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
"""Persistent .o cache for CuTe DSL compiled kernels.
Compiled kernels are exported as object files (.o) via export_to_c.
On subsequent runs the .o is loaded via tvm_ffi (~1ms) instead of
re-generating IR + re-JIT'ing (~100ms per kernel).
Controls:
QUACK_CACHE_ENABLED=0 — disable persistent .o cache (default: enabled)
QUACK_CACHE_DIR=path — override default cache directory
"""
import fcntl
import functools
import hashlib
import os
import pickle
import sys
import tempfile
import time
from collections import namedtuple
from getpass import getuser
from pathlib import Path
import cutlass
import cutlass.cute as cute
import tvm_ffi
CACHE_ENABLED: bool = os.getenv("QUACK_CACHE_ENABLED", "1") == "1"
CACHE_DIR: str | None = os.getenv("QUACK_CACHE_DIR", None)
COMPILE_ONLY: bool = False
# Downstream projects can append directories here to include their sources
# in the cache fingerprint. Must be set before the first jit_cache call.
EXTRA_SOURCE_DIRS: list[Path] = []
EXPORT_FUNC_NAME = "func"
LOCK_TIMEOUT = 60
CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
def _noop_kernel(*args, **kwargs):
pass
def get_cache_path() -> Path:
if CACHE_DIR is not None:
cache_dir = Path(CACHE_DIR)
else:
cache_dir = Path(tempfile.gettempdir()) / getuser() / "quack_cache"
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir
def _hash_source_dir(h, root: Path) -> None:
"""Hash all Python sources under *root* into *h*."""
for src in sorted(root.rglob("*.py")):
if not src.is_file():
continue
h.update(src.relative_to(root).as_posix().encode())
content = src.read_bytes()
h.update(len(content).to_bytes(8, "little"))
h.update(content)
@functools.lru_cache(maxsize=1)
def _compute_source_fingerprint() -> str:
"""Hash quack + extra source dirs plus runtime ABI stamps into a fingerprint."""
h = hashlib.sha256()
h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
h.update(f"cutlass={cutlass.__version__}".encode())
h.update(f"tvm_ffi={tvm_ffi.__version__}".encode())
_hash_source_dir(h, Path(__file__).resolve().parent)
for extra_dir in EXTRA_SOURCE_DIRS:
_hash_source_dir(h, Path(extra_dir).resolve())
return h.hexdigest()
def _key_to_hash(key: tuple) -> str:
return hashlib.sha256(pickle.dumps(key)).hexdigest()
# ---------------------------------------------------------------------------
# File locking
# ---------------------------------------------------------------------------
class FileLock:
"""Advisory file lock using fcntl.flock with timeout."""
def __init__(self, lock_path: Path, exclusive: bool, timeout: float = 15):
self.lock_path = lock_path
self.exclusive = exclusive
self.timeout = timeout
self._fd: int = -1
def __enter__(self) -> "FileLock":
flags = os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT
lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH
self._fd = os.open(str(self.lock_path), flags)
deadline = time.monotonic() + self.timeout
while time.monotonic() < deadline:
try:
fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB)
return self
except OSError:
time.sleep(0.1)
os.close(self._fd)
self._fd = -1
raise RuntimeError(f"Timed out waiting for lock: {self.lock_path}")
def __exit__(self, *exc) -> None:
if self._fd >= 0:
fcntl.flock(self._fd, fcntl.LOCK_UN)
os.close(self._fd)
self._fd = -1
# ---------------------------------------------------------------------------
# JIT cache decorator
# ---------------------------------------------------------------------------
def jit_cache(fn):
"""Decorator that caches compiled CuTe DSL kernels in-memory and on disk.
The decorated function should return a compiled kernel (i.e. call cute.compile).
The disk cache key is (fn.__qualname__, *args, **sorted_kwargs).
"""
cache = {}
hits = 0
misses = 0
@functools.wraps(fn)
def wrapper(*args, **kwargs):
nonlocal hits, misses
cache_key = args + tuple(sorted(kwargs.items())) if kwargs else args
# 1. In-memory hit
if cache_key in cache:
hits += 1
return _noop_kernel if COMPILE_ONLY else cache[cache_key]
# 2. Disk hit
disk_key = (fn.__qualname__,) + cache_key
if CACHE_ENABLED:
sha = _key_to_hash(disk_key)
cache_path = get_cache_path() / _compute_source_fingerprint()
cache_path.mkdir(parents=True, exist_ok=True)
o_path = cache_path / f"{sha}.o"
lock_path = cache_path / f"{sha}.lock"
try:
with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT):
if o_path.exists():
m = cute.runtime.load_module(str(o_path), enable_tvm_ffi=True)
loaded = m[EXPORT_FUNC_NAME]
cache[cache_key] = loaded
hits += 1
return _noop_kernel if COMPILE_ONLY else loaded
except RuntimeError:
pass
# 3. Compile
misses += 1
compiled_fn = fn(*args, **kwargs)
# 4. Store
cache[cache_key] = compiled_fn
if CACHE_ENABLED:
try:
with FileLock(lock_path, exclusive=True, timeout=LOCK_TIMEOUT):
if not o_path.exists():
o_path.parent.mkdir(parents=True, exist_ok=True)
compiled_fn.export_to_c(
object_file_path=str(o_path),
function_name=EXPORT_FUNC_NAME,
)
except Exception as e:
print(f"quack cache: export failed for key {sha}: {e}")
return _noop_kernel if COMPILE_ONLY else compiled_fn
def cache_clear():
nonlocal hits, misses
cache.clear()
hits = 0
misses = 0
def cache_info():
return CacheInfo(hits=hits, misses=misses, maxsize=None, currsize=len(cache))
wrapper.cache = cache
wrapper.cache_clear = cache_clear
wrapper.cache_info = cache_info
return wrapper
|