mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-31 08:13:15 +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
93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
# Regression test for https://github.com/tile-ai/tilelang/issues/2817
|
|
# Cached kernel parameters used to be stored with cloudpickle, so loading a
|
|
# cache entry deserialized arbitrary pickle bytes. They are now stored as
|
|
# JSON produced by TVM's reflection (``tvm.ir.save_json``), which only
|
|
# rebuilds IR nodes and never executes code. Host-side only, no GPU needed.
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.cache.kernel_cache import KernelCache
|
|
from tilelang.engine.param import KernelParam, dump_kernel_params, load_kernel_params
|
|
from tilelang.env import env
|
|
from tvm import tirx
|
|
|
|
|
|
def _params():
|
|
n = tirx.Var("n", "int32")
|
|
return [
|
|
KernelParam(tvm.DataType("float16"), [n, 64]),
|
|
KernelParam(tvm.DataType("int8"), [n * 2, 8]),
|
|
KernelParam(tvm.DataType("float8_e4m3"), [4, 4]),
|
|
KernelParam(tvm.DataType("int32"), []),
|
|
]
|
|
|
|
|
|
def _check_round_trip(params, loaded):
|
|
assert [str(p.dtype) for p in loaded] == [str(p.dtype) for p in params]
|
|
assert [len(p.shape) for p in loaded] == [2, 2, 2, 0]
|
|
assert loaded[0].shape[1] == 64 and isinstance(loaded[0].shape[1], int)
|
|
assert isinstance(loaded[0].shape[0], tirx.Var) and loaded[0].shape[0].name == "n"
|
|
# A symbolic dim shared by two params is the same Var after loading.
|
|
assert loaded[1].shape[0].a.same_as(loaded[0].shape[0])
|
|
tvm.ir.assert_structural_equal(loaded[1].shape[0], loaded[0].shape[0] * 2)
|
|
assert loaded[2].is_float8()
|
|
assert loaded[3].is_scalar()
|
|
|
|
|
|
def test_kernel_params_round_trip_through_json():
|
|
params = _params()
|
|
text = dump_kernel_params(params)
|
|
# Plain JSON text, not a pickle stream.
|
|
assert isinstance(text, str)
|
|
json.loads(text)
|
|
_check_round_trip(params, load_kernel_params(text))
|
|
|
|
|
|
def test_load_kernel_params_rejects_non_json_input():
|
|
with pytest.raises(ValueError):
|
|
load_kernel_params("\x80\x04not a json document")
|
|
|
|
|
|
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, params):
|
|
self.adapter = _FakeAdapter(libpath)
|
|
self.kernel_source = "// device kernel"
|
|
self.params = params
|
|
|
|
|
|
def test_kernel_cache_stores_params_as_json(tmp_path, monkeypatch):
|
|
cache_dir = tmp_path / "cache"
|
|
cache_dir.mkdir()
|
|
monkeypatch.setattr(env, "TILELANG_CACHE_DIR", str(cache_dir))
|
|
cache = KernelCache()
|
|
lib_path = tmp_path / "kernel_lib.so"
|
|
lib_path.write_bytes(b"fake-so")
|
|
params = _params()
|
|
|
|
cache._save_kernel_to_disk("json-params", _FakeKernel(str(lib_path), params))
|
|
|
|
cache_path = Path(cache._get_cache_path("json-params"))
|
|
params_file = cache_path / cache.params_path
|
|
assert params_file.suffix == ".json"
|
|
assert not (cache_path / "params.pkl").exists()
|
|
text = params_file.read_text()
|
|
json.loads(text)
|
|
_check_round_trip(params, load_kernel_params(text))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|