mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Yichen Yan
c025b8d08e
[BugFix][Metal] Resolve auto backend to torch and skip disk cache for torch (#2722) execution_backend="auto" previously resolved to tvm_ffi for metal targets because it picks the first registered spec and tvm_ffi was registered before torch, contradicting the documented default (metal->torch). Register torch first so bare @tilelang.jit runs on Metal via torch.mps.compile_shader; tvm_ffi remains available as an explicit opt-in. KernelCache also crashed on the torch backend: MetalKernelAdapter has no libpath artifact (the MSL source is compiled in-process via torch.mps.compile_shader), but _save_so_cubin_to_disk unconditionally read kernel.adapter.libpath, logging "Error during atomic cache save" with an AttributeError on every run and never persisting a usable entry (the database-load path has no torch branch either). The torch backend now disables disk caching when it compiles, so the disk cache is never written for it, and libpath access is guarded so lib-less adapters can never crash the save path again. Adds regression tests in testing/python/metal/test_metal_cache.py.
100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
"""Regression tests for issue #2722: TileLang on Metal.
|
|
|
|
Covers two behaviors that previously made Metal support opaque:
|
|
1. `@tilelang.jit` with no explicit `execution_backend` must resolve to the
|
|
torch backend (Metal adapter) instead of `tvm_ffi`.
|
|
2. The torch backend must not attempt to persist a compiled library to the disk
|
|
cache (it has no `libpath` artifact; `torch.mps.compile_shader` compiles the
|
|
MSL source in-process). Before the fix this raised
|
|
`AttributeError: 'MetalKernelAdapter' object has no attribute 'libpath'`
|
|
and logged "Error during atomic cache save" on every run.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
from tilelang.cache.kernel_cache import KernelCache
|
|
|
|
|
|
@tilelang.jit
|
|
def _matmul_gemm_auto(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):
|
|
@T.prim_func
|
|
def gemm_kernel(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
C: T.Tensor((M, N), accum_dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_K), dtype, scope="shared")
|
|
B_shared = T.alloc_shared((block_K, block_N), dtype, scope="shared")
|
|
C_local = T.alloc_shared((block_M, block_N), accum_dtype, scope="shared")
|
|
|
|
T.clear(C_local)
|
|
|
|
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=0):
|
|
T.copy(A[by * block_M, ko * block_K], A_shared)
|
|
T.copy(B[ko * block_K, bx * block_N], B_shared)
|
|
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return gemm_kernel
|
|
|
|
|
|
def _run_gemm(jit_kernel, M, N, K):
|
|
a = torch.randn(M, K, dtype=torch.float16, device="mps")
|
|
b = torch.randn(K, N, dtype=torch.float16, device="mps")
|
|
c = torch.zeros(M, N, dtype=torch.float32, device="mps")
|
|
jit_kernel(a, b, c)
|
|
ref = a.to(torch.float32) @ b.to(torch.float32)
|
|
assert torch.allclose(ref, c, atol=1e-2), f"max diff: {(ref - c).abs().max().item()}"
|
|
|
|
|
|
@tilelang.testing.requires_metal
|
|
def test_auto_backend_resolves_to_torch():
|
|
"""Bare @tilelang.jit must resolve to the torch (Metal) backend."""
|
|
M, N, K = 64, 64, 64
|
|
kernel = _matmul_gemm_auto(M, N, K, 16, 16, 8)
|
|
assert kernel.execution_backend == "torch", kernel.execution_backend
|
|
_run_gemm(kernel, M, N, K)
|
|
|
|
|
|
@tilelang.testing.requires_metal
|
|
def test_torch_backend_no_disk_cache_write():
|
|
"""The torch backend must not write disk cache entries or log save errors."""
|
|
cache_logger = logging.getLogger("tilelang.cache.kernel_cache")
|
|
records = []
|
|
handler = logging.Handler()
|
|
handler.emit = lambda record: records.append(record)
|
|
cache_logger.addHandler(handler)
|
|
cache_logger.setLevel(logging.ERROR)
|
|
|
|
cache_root = KernelCache._get_cache_root()
|
|
before = set(os.listdir(cache_root)) if os.path.isdir(cache_root) else set()
|
|
|
|
try:
|
|
M, N, K = 64, 64, 64
|
|
kernel = _matmul_gemm_auto(M, N, K, 16, 16, 8)
|
|
_run_gemm(kernel, M, N, K)
|
|
# Second run in the same process: memory-cache hit, still no disk I/O.
|
|
_run_gemm(kernel, M, N, K)
|
|
finally:
|
|
cache_logger.removeHandler(handler)
|
|
|
|
error_messages = [r.getMessage() for r in records if r.levelno >= logging.ERROR]
|
|
assert not any("Error during atomic cache save" in msg for msg in error_messages), error_messages
|
|
|
|
after = set(os.listdir(cache_root)) if os.path.isdir(cache_root) else set()
|
|
assert after == before, f"torch backend wrote disk cache entries: {after - before}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if torch.mps.is_available():
|
|
tilelang.testing.main()
|