Add disk cache for bytecode to cubin

Signed-off-by: Jay Gu <jagu@nvidia.com>

Author: haijieg

changelog.d/disk-cache.md

@@ -0,0 +1,5 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Add bytecode-to-cubin disk cache to avoid recompilation of unchanged kernels.
+  Controlled by ``CUDA_TILE_CACHE_DIR`` and ``CUDA_TILE_CACHE_SIZE``.

docs/source/debugging.rst

@@ -39,3 +39,15 @@ debugging :class:`TileTypeError`.
 
 Set ``CUDA_TILE_TEMP_DIR`` to configure the directory
 for storing temporary files.
+
+Set ``CUDA_TILE_CACHE_DIR`` to configure the directory
+for the bytecode-to-cubin disk cache. Compiled cubins
+are cached here to avoid recompilation of unchanged
+kernels. Set to ``0``, ``off``, ``none``, or an empty
+string to disable caching. Defaults to
+``~/.cache/cutile-python``.
+
+Set ``CUDA_TILE_CACHE_SIZE`` to configure the maximum
+disk cache size in bytes. Oldest entries are evicted
+when the cache exceeds this limit. Defaults to
+2 GB (2147483648).

src/cuda/tile/_cache.py

@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import hashlib
+import logging
+import os
+import sqlite3
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+_CREATE_TABLE_SQL = """
+CREATE TABLE IF NOT EXISTS cache (
+    key       TEXT PRIMARY KEY,
+    blob      BLOB NOT NULL,
+    blob_size INTEGER NOT NULL,
+    atime     REAL NOT NULL
+)
+"""
+
+_CREATE_ATIME_INDEX_SQL = """
+CREATE INDEX IF NOT EXISTS idx_cache_atime ON cache(atime)
+"""
+
+
+_CACHE_FILENAME = "cache.db"
+
+
+def _close(conn):
+    if conn:
+        try:
+            conn.close()
+        except sqlite3.Error:
+            pass
+
+
+def _open_db(db_path: str) -> sqlite3.Connection:
+    conn = sqlite3.connect(db_path, timeout=5.0)
+    conn.execute(_CREATE_TABLE_SQL)
+    conn.execute(_CREATE_ATIME_INDEX_SQL)
+    return conn
+
+
+def _connect(cache_dir: str) -> sqlite3.Connection:
+    os.makedirs(cache_dir, exist_ok=True)
+    db_path = os.path.join(cache_dir, _CACHE_FILENAME)
+    try:
+        return _open_db(db_path)
+    except sqlite3.Error:
+        logger.debug("cache db corrupt, recreating %s", db_path,
+                     exc_info=True)
+        try:
+            os.unlink(db_path)
+        except OSError:
+            pass
+        return _open_db(db_path)
+
+
+_CACHE_VERSION = b''
+
+
+def cache_key(compiler_version: str, sm_arch: str, opt_level: int,
+              bytecode: bytes) -> str:
+
+    def encode_uint(x: int):
+        return int.to_bytes(x, 4, byteorder='big', signed=False)
+
+    version = compiler_version.encode()
+    arch = sm_arch.encode()
+
+    h = hashlib.sha256()
+    h.update(_CACHE_VERSION)
+    h.update(encode_uint(len(version)))
+    h.update(version)
+    h.update(encode_uint(len(arch)))
+    h.update(arch)
+    h.update(encode_uint(opt_level))
+    h.update(encode_uint(len(bytecode)))
+    h.update(bytecode)
+    return h.hexdigest()
+
+
+def cache_lookup(cache_dir: str, key: str,
+                 temp_dir: str) -> Optional[Path]:
+    conn = None
+    try:
+        conn = _connect(cache_dir)
+        row = conn.execute(
+            "SELECT blob FROM cache WHERE key = ?", (key,)
+        ).fetchone()
+        if row is None:
+            return None
+        conn.execute(
+            "UPDATE cache SET atime = ? WHERE key = ?",
+            (time.time(), key)
+        )
+        conn.commit()
+        blob = row[0]
+    except (sqlite3.Error, OSError):
+        logger.debug("cache lookup failed for %s", key, exc_info=True)
+        return None
+    finally:
+        _close(conn)
+
+    try:
+        with tempfile.NamedTemporaryFile(dir=temp_dir, suffix=".cubin", delete=False) as f:
+            f.write(blob)
+            return Path(f.name)
+    except OSError:
+        logger.debug("cache lookup failed for %s", key, exc_info=True)
+        return None
+
+
+def cache_store(cache_dir: str, key: str, cubin_path) -> None:
+    conn = None
+    try:
+        blob = Path(cubin_path).read_bytes()
+        conn = _connect(cache_dir)
+        conn.execute(
+            "INSERT OR IGNORE INTO cache"
+            " (key, blob, blob_size, atime) VALUES (?, ?, ?, ?)",
+            (key, blob, len(blob), time.time())
+        )
+        conn.commit()
+    except (sqlite3.Error, OSError):
+        logger.debug("cache store failed for %s", key, exc_info=True)
+    finally:
+        _close(conn)
+
+
+def evict_lru(cache_dir: str, size_limit: int) -> None:
+    conn = None
+    try:
+        conn = _connect(cache_dir)
+        row_limit = 100
+        while True:
+            res = conn.execute("""
+            DELETE FROM cache WHERE key IN (SELECT key FROM
+                (SELECT key, SUM(blob_size) OVER (ORDER BY atime, key) as cumul_size
+                    FROM cache ORDER BY atime, key limit ?)
+                WHERE cumul_size <= (SELECT SUM(blob_size) - ? FROM cache)
+            )
+            """, (row_limit, size_limit))
+            if res.rowcount < row_limit:
+                break
+            row_limit *= 10
+        conn.commit()
+    except sqlite3.Error:
+        logger.debug("cache evict failed", exc_info=True)
+    finally:
+        _close(conn)

src/cuda/tile/_compile.py

@@ -51,6 +51,7 @@
 from cuda.tile._passes.check_ampere_fp8 import check_ampere_fp8
 from cuda.tile._passes.dce import dead_code_elimination_pass
 from cuda.tile._passes.token_order import token_order_pass
+from cuda.tile._cache import cache_key, cache_lookup, cache_store, evict_lru
 from cuda.tile._ir2bytecode import generate_bytecode_for_kernel
 from cuda.tile._version import __version__ as cutile_version
 import cuda.tile._bytecode as bc
@@ -248,6 +249,21 @@ def compile_tile(pyfunc,
             print("Can't print MLIR because the internal extension is missing. "
                   "This is currently not a public feature.", file=sys.stderr)
 
+    # Check disk cache before invoking tileiras
+    cache_dir = context.config.cache_dir
+    compiler_ver = _get_compiler_version_string()
+    key = None
+    if cache_dir is None:
+        logger.debug("disk cache disabled: context.config.cache_dir is not set")
+    elif compiler_ver is None:
+        logger.warning("disk cache disabled: compiler version is unknown")
+    else:
+        opt_level = compiler_options.specialize_for_target(sm_arch).opt_level
+        key = cache_key(compiler_ver, sm_arch, opt_level, bytecode_buf)
+        cubin_path = cache_lookup(cache_dir, key, context.config.temp_dir)
+        if cubin_path is not None:
+            return TileLibrary(func_ir.name, cubin_path, bytecode_buf, func_ir.body)
+
     # Compile MLIR module and generate cubin
     with tempfile.NamedTemporaryFile(suffix='.bytecode', prefix=func_ir.name,
                                      dir=context.config.temp_dir, delete=False) as f:
@@ -265,6 +281,10 @@ def compile_tile(pyfunc,
 
             raise e
 
+    if cache_dir is not None and key is not None:
+        cache_store(cache_dir, key, cubin_file)
+        evict_lru(cache_dir, context.config.cache_size_limit)
+
     return TileLibrary(func_ir.name, cubin_file, bytecode_buf, func_ir.body)
 
 
@@ -436,6 +456,13 @@ def _try_get_compiler_version(compiler_bin) -> Optional[str]:
         return None
 
 
+@cache
+def _get_compiler_version_string() -> str | None:
+    binary = _find_compiler_bin()
+    version = _try_get_compiler_version(binary.path)
+    return version
+
+
 @cache
 def get_sm_arch() -> str:
     major, minor = get_compute_capability()

src/cuda/tile/_context.py

@@ -5,6 +5,7 @@
 import atexit
 import os
 import shutil
+import sys
 import tempfile
 from dataclasses import dataclass
 from typing import Optional
@@ -16,14 +17,18 @@ class TileContextConfig:
     log_keys: list[str]
     compiler_timeout_sec: Optional[int]
     enable_crash_dump: bool
+    cache_dir: Optional[str]
+    cache_size_limit: int
 
 
 def init_context_config_from_env():
     config = TileContextConfig(
             temp_dir=get_temp_dir_from_env(),
             log_keys=get_log_keys_from_env(),
             compiler_timeout_sec=get_compile_timeout_from_env(),
-            enable_crash_dump=get_enable_crash_dump_from_env()
+            enable_crash_dump=get_enable_crash_dump_from_env(),
+            cache_dir=get_cache_dir_from_env(),
+            cache_size_limit=get_cache_size_limit_from_env()
             )
     return config
 
@@ -71,3 +76,20 @@ def get_enable_crash_dump_from_env() -> bool:
     key = "CUDA_TILE_ENABLE_CRASH_DUMP"
     env = os.environ.get(key, "0").lower()
     return env in ("1", "true", "yes", "on")
+
+
+def get_cache_dir_from_env() -> Optional[str]:
+    home_cache = os.path.join(os.path.expanduser("~"), ".cache")
+    if sys.platform == "win32":
+        base = os.environ.get("LOCALAPPDATA", home_cache)
+    else:
+        base = os.environ.get("XDG_CACHE_HOME", home_cache)
+    default = os.path.join(base, "cutile-python")
+    env = os.environ.get("CUDA_TILE_CACHE_DIR", default)
+    if env.strip().lower() in ("0", "off", "none", ""):
+        return None
+    return env
+
+
+def get_cache_size_limit_from_env() -> int:
+    return int(os.environ.get("CUDA_TILE_CACHE_SIZE", 1 << 31))  # 2GB

test/test_cache.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sqlite3
+import time
+
+import pytest
+
+from cuda.tile._cache import cache_key, cache_lookup, cache_store, evict_lru
+
+
+def test_cache_key_equal():
+    k1 = cache_key("v1", "sm_90", 3, b"data")
+    k2 = cache_key("v1", "sm_90", 3, b"data")
+    assert k1 == k2
+
+
+def test_cache_key_differs():
+    base = cache_key("v1", "sm_90", 3, b"data")
+    assert cache_key("v2", "sm_90", 3, b"data") != base
+    assert cache_key("v1", "sm_80", 3, b"data") != base
+    assert cache_key("v1", "sm_90", 2, b"data") != base
+    assert cache_key("v1", "sm_90", 3, b"other") != base
+
+
+def make_cubin(tmp_path, name, content):
+    src = tmp_path / name
+    src.write_bytes(content)
+    return src
+
+
+@pytest.fixture
+def cache_env(tmp_path):
+    cache_dir = str(tmp_path / "cache")
+    return cache_dir, tmp_path
+
+
+def test_store_then_lookup(cache_env):
+    cache_dir, tmp_path = cache_env
+    key = cache_key("v1", "sm_90", 3, b"data")
+    content = b"\x7fELF_fake_cubin_data"
+
+    cache_store(cache_dir, key, make_cubin(tmp_path, "kernel.cubin", content))
+
+    result = cache_lookup(cache_dir, key, str(tmp_path))
+    assert result is not None
+    assert result.read_bytes() == content
+
+
+def test_lookup_updates_atime(cache_env):
+    cache_dir, tmp_path = cache_env
+    key = cache_key("v1", "sm_90", 3, b"data")
+
+    cache_store(cache_dir, key, make_cubin(tmp_path, "kernel.cubin", b"data"))
+
+    # Manually set old atime in DB
+    import os
+    db_path = os.path.join(cache_dir, "cache.db")
+    old_time = time.time() - 1000
+    conn = sqlite3.connect(db_path)
+    conn.execute("UPDATE cache SET atime = ? WHERE key = ?", (old_time, key))
+    conn.commit()
+    conn.close()
+
+    cache_lookup(cache_dir, key, str(tmp_path))
+
+    conn = sqlite3.connect(db_path)
+    atime = conn.execute(
+        "SELECT atime FROM cache WHERE key = ?", (key,)
+    ).fetchone()[0]
+    conn.close()
+    assert atime > old_time
+
+
+def test_lookup_miss(cache_env):
+    cache_dir, _ = cache_env
+
+    result = cache_lookup(cache_dir, "a" * 64, str(cache_dir))
+    assert result is None
+
+
+def test_evict_lru(cache_env):
+    cache_dir, tmp_path = cache_env
+    import os
+    db_path = os.path.join(cache_dir, "cache.db")
+
+    # Populate 5 entries (1000 bytes each, 5000 total)
+    keys = []
+    for i in range(5):
+        key = cache_key(str(i), "sm_90", 3, b"data")
+        keys.append(key)
+        cache_store(cache_dir, key,
+                    make_cubin(tmp_path, f"k{i}.cubin", b"x" * 1000))
+
+    # Set controlled atimes so eviction order is deterministic
+    conn = sqlite3.connect(db_path)
+    for i, key in enumerate(keys):
+        conn.execute(
+            "UPDATE cache SET atime = ? WHERE key = ?",
+            (float(i), key)
+        )
+    conn.commit()
+    conn.close()
+
+    # Evict to keep 3000 bytes; newest 3 survive (indices 2, 3, 4)
+    evict_lru(cache_dir, 3000)
+
+    remaining = [k for k in keys
+                 if cache_lookup(cache_dir, k, str(tmp_path)) is not None]
+    assert remaining == keys[2:]