enable tvm-ffi for metal

Author: oraluben

testing/python/metal/test_metal_codegen.py

@@ -5,7 +5,7 @@
 import torch
 
 
-@tilelang.jit(execution_backend='torch')
+@tilelang.jit(execution_backend='tvm_ffi')
 def matmul(M, N, K, block_M, block_N, block_K, dtype="float32", accum_dtype="float"):
 
     @T.prim_func

tilelang/engine/lower.py

@@ -161,6 +161,8 @@ def device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule:
         device_mod = tvm.ffi.get_global_func("target.build.tilelang_cuda")(device_mod, target)
     elif target.kind.name == "hip":
         device_mod = tvm.ffi.get_global_func("target.build.tilelang_hip")(device_mod, target)
+    elif target.kind.name == "metal":
+        device_mod = tvm.ffi.get_global_func("target.build.metal")(device_mod, target)
     else:
         raise ValueError(f"Target {target.kind.name} is not supported")
 

tilelang/jit/execution_backend.py

@@ -35,7 +35,7 @@ def allowed_backends_for_target(target: Target, *, include_unavailable: bool = T
     elif kind == "hip":
         allowed = ["tvm_ffi", "cython", "ctypes"]
     elif kind == "metal":
-        allowed = ["torch"]
+        allowed = ["tvm_ffi", "torch"]
     elif kind == "c":  # CPU C backend
         allowed = ["cython", "ctypes", "tvm_ffi"]
     else: