[Question] Kindly ask for metal example.

[oraluben]

This is related to #255.

I'm trying to migrate a torch extension for mps backend to tvm-ffi and have met various issue, mainly the internal metal context management not exposed by torch.

I saw many code in tvm-ffi involves metal logic so I'd like to know if you have some working example of writing (or loading) metal kernel with tvm-ffi.

Here's a very simple torch extension that I'm trying to get started with: https://gist.github.com/oraluben/9f4834ba13b32090967d86babbc72389

[tqchen]

Indeed one challenge here is to get the metal context management, unlike cuda, metal do not have a predefined global context (e.g. in cuda we can simply cudaMalloc device=0) and it would work. So runtimes like torch, mlx do not necessarily share global context.

Previously we also have a tvm-runtime in metal as part of tvm compiler that solves the problem of kernel launch. But the main remaining issue is the ability to share, or query current global context, if the goal is interop. Maybe we need some extra extension to get such interop working, but i haven't think very deeply into it

[oraluben]

Based on my experiment, mlx's dlpack support isn't very well (it cannot load from a dlpack closure, and memory synchronization seems to have issue when modifying a dlpack tensor exported by mlx from torch).

So I think a reasonable short-term goal is to make tvm-ffi work with torch mps backend, i.e. compile metal kernel w.o. torch (and minimal efforts to migrating from an existing extension).

[junrushao]

I didn’t very much touch TVM’s Metal backend before, and I’m trying to understand the challenges here.

the ability to share, or query current global context, if the goal is interop

@tqchen Could you kindly point out what specifically is missing in forms of TVM code? I’m reading the runtime code here (https://github.com/apache/tvm/blob/main/src/runtime/metal/metal_module.mm#L85) and am trying to figure out which extra APIs are needed to enable tvm-ffi for Metals

[oraluben]

I also tried to simply enable tvm-ffi backend in tilelang, it works without error but the kernel seems not working and gives empty result: tile-ai/tilelang#1289

tilelang with tvm-ffi backend should be using the tvm runtime already, so it's supposed to work?

[junrushao]

Yes. TileLang is using tvm runtime, which is perfectly fine and work already out of box. The main discussion here is more relevant to cases where tvm runtime is not available.

[oraluben]

Yes. TileLang is using tvm runtime, which is perfectly fine and work already out of box.

One thing that I'm curious about is that it's not working for metal backend. The kernel seems to run but the tensor is unchanged. Do you have some suggestion to debug this?

[tqchen]

To be clear, if we use metal through tvm runtime, aka allocate tensor as tvm.nd.Tensor with device tvm.metal() it should work of box. This has been our approach on enabling iOS LLM support for example. For enabling execution through tvm runtime, perhaps we can start from an upstream tvm metal example, which should work, and dig into a bit difference here.

The main challenge comes in we want to exchange Tensors, e.g. when we pass in torch.mps Tensor or mlx.Tensor and treat them as native tvm runtime Tensor. The complication comes from the metal context management issue i noted above

[junrushao]

it's not working for metal backend. The kernel seems to run but the tensor is unchanged

I suppose this repo is tested with TVM’s metal backend, and unittests inside tvm main repo should already work out of the box. Could you test and share a case that doesn’t work as expected?

[oraluben]

I suppose this repo is tested with TVM’s metal backend, and unittests inside tvm main repo should already work out of the box. Could you test and share a case that doesn’t work as expected?

Sure. If you checkout code from tile-ai/tilelang#1289 and run the code below, you'll see the tensor not updated. To see expected behavior, change execution_backend='tvm_ffi' to execution_backend='torch'.

import tilelang
from tilelang import tvm as tvm
import tilelang.language as T
import torch


@tilelang.jit(execution_backend='tvm_ffi')
def add(M, block_M, dtype="float32"):

    @T.prim_func
    def add_kernel(
            A: T.Tensor((M, ), dtype),
    ):
        with T.Kernel(T.ceildiv(M, block_M), threads=128) as bx:
            for i in T.Parallel(block_M):
                A[bx * block_M + i] = 1

    return add_kernel


def assert_gemm(
    M,
    block_M,
    dtype="int",
    atol=1e-8,
):
    jit_kernel = add(M, block_M, dtype=dtype)

    torch_dtype = getattr(torch, dtype)
    a = torch.randint(100, (M, ), dtype=torch_dtype, device='mps')

    jit_kernel(a)

    print(a)

    assert jit_kernel.kernel_source is not None
    print(jit_kernel.kernel_source)


if __name__ == "__main__":
    assert_gemm(1024, 16)

But it makes sense that this is a torch tensor, not tvm.nd.Tensor so the dlpack / ctx management have some issue.