mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Rishabh Sinha
11ec2397fe
The kernel cache and the autotuner result cache loaded params.pkl with cloudpickle.load(), so a writable cache directory allowed arbitrary code execution on the next cache hit. Kernel params are plain data, so store them as JSON via TVM's reflection (tvm.ir.save_json / load_json), which rebuilds IR nodes without executing code. Entries without params.json are treated as incomplete and rebuilt. Fixes #2817
402 lines
14 KiB
Python
402 lines
14 KiB
Python
import builtins
|
|
import errno
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from tilelang import tvm
|
|
|
|
import tilelang.cache.kernel_cache as kernel_cache_mod
|
|
from tilelang.backend import create_backend_context
|
|
from tilelang.cache.kernel_cache import KernelCache
|
|
from tilelang.engine.param import KernelParam, dump_kernel_params
|
|
from tilelang.env import env
|
|
from tilelang.jit.adapter.base import CachedTextSource
|
|
from tilelang.jit.adapter.kernel_cache import TVMFFIKernelCache
|
|
from tilelang.jit.adapter.nvrtc.kernel_cache import NVRTCKernelCache
|
|
|
|
|
|
class _FakeAdapter:
|
|
def __init__(self, libpath: str):
|
|
self.libpath = libpath
|
|
|
|
def get_kernel_source(self):
|
|
return "// host kernel"
|
|
|
|
|
|
class _FakeKernel:
|
|
def __init__(self, libpath: str):
|
|
self.adapter = _FakeAdapter(libpath)
|
|
self.kernel_source = "// device kernel"
|
|
self.params = [KernelParam(tvm.DataType("float32"), [4])]
|
|
|
|
|
|
@pytest.fixture
|
|
def cache_dirs(tmp_path, monkeypatch):
|
|
cache_dir = tmp_path / "cache"
|
|
cache_dir.mkdir()
|
|
monkeypatch.setattr(env, "TILELANG_CACHE_DIR", str(cache_dir))
|
|
return cache_dir
|
|
|
|
|
|
def _make_fake_kernel(tmp_path):
|
|
lib_path = tmp_path / "kernel_lib.so"
|
|
lib_path.write_bytes(b"fake-so")
|
|
return _FakeKernel(str(lib_path))
|
|
|
|
|
|
def _make_fake_nvrtc_kernel(tmp_path):
|
|
lib_path = tmp_path / "kernel.cubin"
|
|
lib_path.write_bytes(b"fake-cubin")
|
|
lib_path.with_suffix(".py").write_text("# fake launcher")
|
|
return _FakeKernel(str(lib_path))
|
|
|
|
|
|
def _write_complete_kernel_cache_entry(
|
|
cache: KernelCache,
|
|
key: str,
|
|
device_source: str = "// device kernel",
|
|
host_source: str = "// host kernel",
|
|
) -> Path:
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
cache_path.mkdir(parents=True)
|
|
(cache_path / cache.device_kernel_path).write_text(device_source)
|
|
(cache_path / cache.host_kernel_path).write_text(host_source)
|
|
(cache_path / cache.kernel_lib_path).write_bytes(b"fake-so")
|
|
(cache_path / cache.params_path).write_text(dump_kernel_params([KernelParam(tvm.DataType("float32"), [4])]))
|
|
cache._write_manifest(str(cache_path))
|
|
return cache_path
|
|
|
|
|
|
def test_kernel_cache_disk_hit_defers_source_loading(cache_dirs, monkeypatch):
|
|
cache = KernelCache()
|
|
key = "lazy-source-load"
|
|
cache_path = _write_complete_kernel_cache_entry(cache, key)
|
|
|
|
sentinel = object()
|
|
captured = {}
|
|
|
|
def fail_source_load(*args, **kwargs):
|
|
raise AssertionError("disk cache hit should pass source paths through for lazy loading")
|
|
|
|
def fake_from_database(cls, **kwargs):
|
|
captured.update(kwargs)
|
|
return sentinel
|
|
|
|
monkeypatch.setattr(cache, "_load_kernel_source", fail_source_load)
|
|
monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fake_from_database))
|
|
|
|
backend_context = create_backend_context("cuda", execution_backend="tvm_ffi")
|
|
loaded = cache._load_kernel_from_disk(
|
|
key,
|
|
backend_context=backend_context,
|
|
out_idx=[0],
|
|
pass_configs=None,
|
|
compile_flags=None,
|
|
func=None,
|
|
)
|
|
|
|
assert loaded is sentinel
|
|
assert captured["host_kernel_source"] == CachedTextSource(path=str(cache_path / cache.host_kernel_path))
|
|
assert captured["device_kernel_source"] == CachedTextSource(path=str(cache_path / cache.device_kernel_path))
|
|
assert captured["kernel_lib_path"] == str(cache_path / cache.kernel_lib_path)
|
|
assert captured["params"] == [KernelParam(tvm.DataType("float32"), [4])]
|
|
assert captured["backend_context"] is backend_context
|
|
|
|
|
|
def test_kernel_cache_disk_hit_perf_skips_large_source_file_reads(cache_dirs, monkeypatch):
|
|
cache = KernelCache()
|
|
key = "lazy-source-load-perf"
|
|
large_source = "// source\n" + ("x" * (2 * 1024 * 1024))
|
|
cache_path = _write_complete_kernel_cache_entry(
|
|
cache,
|
|
key,
|
|
device_source=large_source,
|
|
host_source=large_source,
|
|
)
|
|
source_paths = {
|
|
(cache_path / cache.device_kernel_path).resolve(),
|
|
(cache_path / cache.host_kernel_path).resolve(),
|
|
}
|
|
source_read_count = 0
|
|
sentinel = object()
|
|
|
|
real_open = builtins.open
|
|
|
|
def tracking_open(file, *args, **kwargs):
|
|
nonlocal source_read_count
|
|
mode = args[0] if args else kwargs.get("mode", "r")
|
|
try:
|
|
path = Path(file).resolve()
|
|
except TypeError:
|
|
return real_open(file, *args, **kwargs)
|
|
if "r" in mode and path in source_paths:
|
|
source_read_count += 1
|
|
raise AssertionError("cache perf regression: source file read during disk cache hit")
|
|
return real_open(file, *args, **kwargs)
|
|
|
|
def fake_from_database(cls, **kwargs):
|
|
return sentinel
|
|
|
|
monkeypatch.setattr(builtins, "open", tracking_open)
|
|
monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fake_from_database))
|
|
|
|
loaded = cache._load_kernel_from_disk(
|
|
key,
|
|
backend_context=create_backend_context("cuda", execution_backend="tvm_ffi"),
|
|
out_idx=[0],
|
|
pass_configs=None,
|
|
compile_flags=None,
|
|
func=None,
|
|
)
|
|
|
|
assert loaded is sentinel
|
|
assert source_read_count == 0
|
|
|
|
|
|
def test_lazy_jit_closure_specialization_re_elaborates_before_kernel_cache(cache_dirs):
|
|
import tilelang
|
|
import tilelang.language as T
|
|
from tilelang.cache import _dispatch_map
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
pytest.skip("PyTorch is required for CUDA availability checks")
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA is required to compile the test kernel")
|
|
|
|
for cache in _dispatch_map.values():
|
|
cache._memory_cache.clear()
|
|
|
|
def build(n: int):
|
|
@tilelang.jit
|
|
def make_kernel():
|
|
m = T.dynamic("M")
|
|
|
|
@T.prim_func
|
|
def copy_kernel(
|
|
x: T.Tensor[(m, n), T.float32],
|
|
y: T.Tensor[(m, n), T.float32],
|
|
):
|
|
with T.Kernel(m, threads=128) as pid:
|
|
for j in T.Parallel(n):
|
|
y[pid, j] = x[pid, j]
|
|
|
|
return copy_kernel
|
|
|
|
return make_kernel()
|
|
|
|
kernel_128 = build(128)
|
|
kernel_256 = build(256)
|
|
|
|
assert kernel_128 is not kernel_256
|
|
assert [tuple(map(str, param.shape)) for param in kernel_128.adapter.params] == [("M", "128"), ("M", "128")]
|
|
assert [tuple(map(str, param.shape)) for param in kernel_256.adapter.params] == [("M", "256"), ("M", "256")]
|
|
|
|
|
|
def test_kernel_cache_disk_hit_rejects_entries_missing_sources(cache_dirs, monkeypatch):
|
|
cache = KernelCache()
|
|
key = "missing-source-entry"
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
cache_path.mkdir(parents=True)
|
|
(cache_path / cache.kernel_lib_path).write_bytes(b"fake-so")
|
|
(cache_path / cache.params_path).write_text(dump_kernel_params([KernelParam(tvm.DataType("float32"), [4])]))
|
|
|
|
def fail_from_database(cls, **kwargs):
|
|
raise AssertionError("incomplete cache entries should miss before rebuilding from database")
|
|
|
|
monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fail_from_database))
|
|
|
|
loaded = cache._load_kernel_from_disk(
|
|
key,
|
|
backend_context=create_backend_context("cuda", execution_backend="tvm_ffi"),
|
|
out_idx=[0],
|
|
pass_configs=None,
|
|
compile_flags=None,
|
|
func=None,
|
|
)
|
|
|
|
assert loaded is None
|
|
|
|
|
|
def test_nvrtc_adapter_host_source_lazy_loads(tmp_path):
|
|
pytest.importorskip("cuda.bindings.driver", reason="NVRTC adapter requires cuda-python")
|
|
from tilelang.jit.adapter.nvrtc.adapter import NVRTCKernelAdapter
|
|
|
|
host_source_path = tmp_path / "host_kernel.cu"
|
|
host_source_path.write_text("// nvrtc host source")
|
|
adapter = NVRTCKernelAdapter.__new__(NVRTCKernelAdapter)
|
|
adapter.host_func = None
|
|
adapter._host_kernel_source_path = str(host_source_path)
|
|
|
|
assert adapter.get_host_source() == "// nvrtc host source"
|
|
assert adapter.host_func == "// nvrtc host source"
|
|
|
|
|
|
def test_cutedsl_adapter_host_source_lazy_loads(tmp_path):
|
|
from tilelang.jit.adapter.cutedsl.adapter import CuTeDSLKernelAdapter
|
|
|
|
host_source_path = tmp_path / "kernel.py"
|
|
host_source_path.write_text("# cutedsl host source")
|
|
adapter = CuTeDSLKernelAdapter.__new__(CuTeDSLKernelAdapter)
|
|
adapter.host_kernel_source = None
|
|
adapter.host_func = None
|
|
adapter._host_kernel_source_path = str(host_source_path)
|
|
|
|
assert adapter.get_host_source() == "# cutedsl host source"
|
|
assert adapter.host_kernel_source == "# cutedsl host source"
|
|
|
|
|
|
def test_tvm_ffi_source_fallback_handles_missing_runtime_module():
|
|
from tilelang.jit.adapter.tvm_ffi import TVMFFIKernelAdapter
|
|
|
|
adapter = TVMFFIKernelAdapter.__new__(TVMFFIKernelAdapter)
|
|
adapter.host_kernel_source = None
|
|
adapter.device_kernel_source = None
|
|
adapter._host_kernel_source_path = None
|
|
adapter._device_kernel_source_path = None
|
|
adapter.rt_mod = None
|
|
|
|
assert adapter.get_host_source() is None
|
|
assert adapter.get_device_source() is None
|
|
assert adapter.get_kernel_source(kernel_only=True) == ""
|
|
assert adapter.get_kernel_source(kernel_only=False) == ""
|
|
|
|
|
|
def test_kernel_cache_rewrites_incomplete_cache_dir(cache_dirs, tmp_path):
|
|
cache = KernelCache()
|
|
key = "atomic-repair"
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
cache_path.mkdir(parents=True)
|
|
(cache_path / "stale.txt").write_text("partial")
|
|
|
|
cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
|
|
|
|
assert (cache_path / cache.device_kernel_path).exists()
|
|
assert (cache_path / cache.host_kernel_path).exists()
|
|
assert (cache_path / cache.kernel_lib_path).exists()
|
|
assert (cache_path / cache.params_path).exists()
|
|
assert not (cache_path / "stale.txt").exists()
|
|
|
|
|
|
def test_kernel_cache_logs_write_oserror_instead_of_treating_it_as_race(cache_dirs, tmp_path, monkeypatch):
|
|
cache = KernelCache()
|
|
key = "atomic-write-error"
|
|
logged = []
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
staging_root = Path(cache._get_staging_root())
|
|
|
|
def raise_write_error(*args, **kwargs):
|
|
raise OSError(errno.ENOSPC, "No space left on device")
|
|
|
|
def record_exception(message, *args, **kwargs):
|
|
logged.append(message)
|
|
|
|
monkeypatch.setattr(cache, "_save_so_cubin_to_disk", raise_write_error)
|
|
monkeypatch.setattr(cache.logger, "exception", record_exception)
|
|
|
|
cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
|
|
|
|
assert not cache_path.exists()
|
|
assert "Error during atomic cache save" in logged
|
|
assert not staging_root.exists() or not any(staging_root.iterdir())
|
|
|
|
|
|
def test_kernel_cache_does_not_publish_incomplete_dir_when_device_source_is_missing(cache_dirs, tmp_path, monkeypatch):
|
|
cache = KernelCache()
|
|
key = "atomic-missing-device-source"
|
|
kernel = _make_fake_kernel(tmp_path)
|
|
kernel.kernel_source = None
|
|
logged = []
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
staging_root = Path(cache._get_staging_root())
|
|
|
|
def record_exception(message, *args, **kwargs):
|
|
logged.append(message)
|
|
|
|
monkeypatch.setattr(cache.logger, "exception", record_exception)
|
|
|
|
cache._save_kernel_to_disk(key, kernel)
|
|
|
|
assert not cache_path.exists()
|
|
assert "Error during atomic cache save" in logged
|
|
assert not staging_root.exists() or not any(staging_root.iterdir())
|
|
|
|
|
|
def test_nvrtc_kernel_cache_rewrites_dir_missing_launcher(cache_dirs, tmp_path):
|
|
cache = NVRTCKernelCache()
|
|
key = "nvrtc-atomic-repair"
|
|
cache_path = Path(cache._get_cache_path(key))
|
|
cache_path.mkdir(parents=True)
|
|
(cache_path / cache.device_kernel_path).write_text("// device kernel")
|
|
(cache_path / cache.host_kernel_path).write_text("// host kernel")
|
|
(cache_path / cache.kernel_lib_path).write_bytes(b"old-cubin")
|
|
(cache_path / cache.params_path).write_bytes(b"old-params")
|
|
(cache_path / "legacy.txt").write_text("stale")
|
|
|
|
cache._save_kernel_to_disk(key, _make_fake_nvrtc_kernel(tmp_path))
|
|
|
|
assert (cache_path / cache.kernel_py_path).exists()
|
|
assert not (cache_path / "legacy.txt").exists()
|
|
|
|
|
|
def test_safe_write_executable_uses_explicit_export_kwargs(cache_dirs, tmp_path):
|
|
captured = []
|
|
|
|
class FakeExecutable:
|
|
def export_library(self, path, **kwargs):
|
|
captured.append(dict(kwargs))
|
|
with open(path, "wb") as f:
|
|
f.write(b"fake-so")
|
|
|
|
export_kwargs = {"options": ["-Wl,-dead_strip"], "custom": "keep"}
|
|
|
|
KernelCache._safe_write_executable(
|
|
FakeExecutable(),
|
|
str(tmp_path / "explicit.so"),
|
|
export_kwargs=export_kwargs,
|
|
)
|
|
|
|
assert captured == [export_kwargs]
|
|
|
|
|
|
def test_tvm_ffi_kernel_cache_selects_export_kwargs_from_host_target(cache_dirs, tmp_path, monkeypatch):
|
|
"""Regression: device target may be cuda while the exported host module is LLVM.
|
|
Source-compilation options must not be passed to that LLVM object-link step.
|
|
"""
|
|
|
|
source_compile_args = {"options": ["-x", "objective-c++"], "source": "keep"}
|
|
export_link_args = {"link": "keep"}
|
|
captured = []
|
|
|
|
monkeypatch.setattr(KernelCache, "_get_source_compile_args", staticmethod(lambda: source_compile_args))
|
|
monkeypatch.setattr(KernelCache, "_get_export_link_args", staticmethod(lambda: export_link_args))
|
|
|
|
class FakeExecutable:
|
|
def export_library(self, path, **kwargs):
|
|
captured.append(dict(kwargs))
|
|
with open(path, "wb") as f:
|
|
f.write(b"fake-so")
|
|
|
|
class FakeAdapter:
|
|
def get_exportable_executable(self):
|
|
return FakeExecutable()
|
|
|
|
def make_kernel(target_host_kind: str):
|
|
return SimpleNamespace(
|
|
adapter=FakeAdapter(),
|
|
target=SimpleNamespace(kind=SimpleNamespace(name="cuda")),
|
|
artifact=SimpleNamespace(target_host=SimpleNamespace(kind=SimpleNamespace(name=target_host_kind))),
|
|
)
|
|
|
|
cache = TVMFFIKernelCache()
|
|
|
|
cache._save_so_cubin_to_disk(make_kernel("llvm"), str(tmp_path))
|
|
assert captured[-1] == export_link_args
|
|
assert "options" not in captured[-1]
|
|
|
|
cache._save_so_cubin_to_disk(make_kernel("c"), str(tmp_path))
|
|
assert captured[-1] == source_compile_args
|