diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a5b888b2..ebbfb3413 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -199,6 +199,12 @@ if(NOT TILELANG_BACKEND_USER_SELECTED) endif() if(USE_METAL) + if(NOT APPLE) + # On non-Apple platforms USE_METAL=ON enables only codegen (Metal source + # generation) without requiring the Metal/Foundation frameworks. + message(STATUS "Metal backend on non-Apple: enabling codegen-only mode (no Metal runtime)") + set(USE_METAL OFF) + endif() file(GLOB TILE_LANG_METAL_SRCS src/target/rt_mod_metal.cc ) diff --git a/testing/python/metal/test_metal_codegen_linux.py b/testing/python/metal/test_metal_codegen_linux.py new file mode 100644 index 000000000..690a0a3a2 --- /dev/null +++ b/testing/python/metal/test_metal_codegen_linux.py @@ -0,0 +1,72 @@ +"""Test Metal code generation on any platform (including Linux). + +These tests verify that TileLang can compile kernels down to Metal shader +source code without requiring a Metal runtime or macOS. +""" + +import tilelang +from tilelang import tvm as tvm +import tilelang.testing +import tilelang.language as T + + +def matmul(M, N, K, block_M, block_N, block_K, dtype=T.float32, accum_dtype=T.float32): + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), 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_fragment((block_M, block_N), accum_dtype) + + 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, coalesced_width=2) + T.copy(B[ko * block_K, bx * block_N], B_shared, coalesced_width=2) + + for i, j in T.Parallel(block_M, block_N): + for k in T.Serial(block_K): + C_local[i, j] += A_shared[i, k] * B_shared[k, j] + + T.copy(C_local, C[by * block_M, bx * block_N], coalesced_width=2) + + return main + + +def assert_metal_codegen( + M, + N, + K, + block_M, + block_N, + block_K, + dtype=T.float32, + accum_dtype=T.float32, +): + func = matmul(M, N, K, block_M, block_N, block_K, dtype=dtype, accum_dtype=accum_dtype) + with tvm.transform.PassContext(), tvm.target.Target("metal"): + artifact = tilelang.lower(func, target="metal") + + src_code = artifact.kernel_source + assert src_code is not None + assert "metal" in src_code or "kernel void" in src_code + + +def test_metal_codegen_float32(): + assert_metal_codegen(1024, 1024, 1024, 16, 16, 16) + + +def test_metal_codegen_float16(): + assert_metal_codegen(1024, 1024, 1024, 16, 16, 16, dtype=T.float16) + + +def test_metal_codegen_int32(): + assert_metal_codegen(1024, 1024, 1024, 16, 16, 16, dtype=T.int32) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/tilelang/engine/lower.py b/tilelang/engine/lower.py index 2aa18806c..40a757c7a 100644 --- a/tilelang/engine/lower.py +++ b/tilelang/engine/lower.py @@ -3,7 +3,6 @@ from __future__ import annotations from typing import Callable -import sys import tilelang.transform from tilelang import tvm as tvm from tvm import tir @@ -145,7 +144,20 @@ def canon_target_host(target: str | Target, target_host: str | Target | None): return target_host -def host_codegen(host_mod: tvm.IRModule, target_host: Target) -> tvm.IRModule: +def host_codegen(host_mod: tvm.IRModule, target_host: Target, target: Target | None = None) -> tvm.IRModule: + """Generate host-side code from the lowered IR module. + + Parameters + ---------- + host_mod : tvm.IRModule + The host-side IR module to compile. + target_host : Target + The host compilation target (e.g. "llvm" or "c"). + target : Target, optional + The device target. When the device target is Metal, the pass + MarkHostMetalContext is applied so that the generated host code + contains the Metal/MPS synchronisation logic. + """ host_mod = tir.transform.BindTarget(target_host)(host_mod) host_mod = tir.transform.FP8StorageLegalize()(host_mod) host_mod = tir.transform.BF16StorageLegalize()(host_mod) @@ -154,7 +166,7 @@ def host_codegen(host_mod: tvm.IRModule, target_host: Target) -> tvm.IRModule: host_mod = tilelang.transform.LowerIntrin()(host_mod) host_mod = tilelang.transform.LowerDeviceStorageAccessInfo()(host_mod) host_mod = tir.transform.CombineContextCall()(host_mod) - if sys.platform == "darwin": + if target is not None and target.kind.name == "metal": host_mod = MarkHostMetalContext()(host_mod) if target_host.kind.name == "llvm": host_mod = tvm.ffi.get_global_func("target.build.llvm")(host_mod, target_host) @@ -257,7 +269,7 @@ def lower( codegen_mod = device_codegen(device_mod, target) if enable_device_compile else device_codegen_without_compile(device_mod, target) if enable_host_codegen: - host_mod = host_codegen(host_mod, target_host) + host_mod = host_codegen(host_mod, target_host, target=target) host_mod.import_module(codegen_mod) return CompiledArtifact(host_mod, device_mod, params, codegen_mod.inspect_source(), rt_mod=host_mod)