mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
93 lines
2.3 KiB
Python
93 lines
2.3 KiB
Python
import threading
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
from tilelang.jit.adapter.tvm_ffi import TVMFFIKernelAdapter
|
|
|
|
|
|
class _FakeKernelParam:
|
|
shape = [1]
|
|
dtype = SimpleNamespace(bits=32, lanes=1)
|
|
|
|
@staticmethod
|
|
def torch_dtype():
|
|
return torch.float32
|
|
|
|
|
|
class _TestAdapter(TVMFFIKernelAdapter):
|
|
@property
|
|
def prim_func(self):
|
|
return self._test_prim_func
|
|
|
|
|
|
def _make_adapter():
|
|
adapter = _TestAdapter.__new__(_TestAdapter)
|
|
tir_param = object()
|
|
adapter.params = [_FakeKernelParam()]
|
|
adapter.result_idx = []
|
|
adapter._test_prim_func = SimpleNamespace(
|
|
params=[tir_param],
|
|
buffer_map={tir_param: SimpleNamespace(dtype="float32")},
|
|
)
|
|
adapter._process_dynamic_symbolic = lambda: {}
|
|
adapter.executable = None
|
|
adapter._executable_lock = threading.Lock()
|
|
|
|
created = []
|
|
|
|
def make_executable():
|
|
def executable(*args):
|
|
return None
|
|
|
|
created.append(executable)
|
|
return executable
|
|
|
|
adapter._make_executable = make_executable
|
|
return adapter, created
|
|
|
|
|
|
def test_cold_compiled_dispatch_does_not_probe_cuda(monkeypatch):
|
|
adapter, created = _make_adapter()
|
|
func = adapter._convert_torch_func()
|
|
tensor = torch.empty(1)
|
|
cuda_probe_count = 0
|
|
|
|
def counted_is_available():
|
|
nonlocal cuda_probe_count
|
|
cuda_probe_count += 1
|
|
return False
|
|
|
|
monkeypatch.setattr(torch.cuda, "is_available", counted_is_available)
|
|
|
|
for _ in range(3):
|
|
func(tensor)
|
|
|
|
assert cuda_probe_count == 0
|
|
assert len(created) == 1
|
|
assert adapter.executable is created[0]
|
|
|
|
|
|
def test_executable_is_initialized_once_and_reused():
|
|
adapter, created = _make_adapter()
|
|
|
|
executable = adapter._get_executable()
|
|
|
|
assert adapter._get_executable() is executable
|
|
assert adapter.get_exportable_executable() is executable
|
|
assert adapter.executable is executable
|
|
assert len(created) == 1
|
|
|
|
|
|
def test_preloaded_executable_is_reused():
|
|
adapter, created = _make_adapter()
|
|
|
|
def preloaded_executable(*args):
|
|
return None
|
|
|
|
adapter.executable = preloaded_executable
|
|
|
|
assert adapter._get_executable() is preloaded_executable
|
|
assert adapter.get_exportable_executable() is preloaded_executable
|
|
assert created == []
|