Metal support: MTLBuffer access and current_work_stream for kDLMetal

[ved1beta]

Hi! I'd like to help land Metal support but I'm new to tvm-ffi internals
and want to ask before I start coding.

The downstream context: I'm working with huggingface/kernels, which has
Metal kernel-builder support staged on a WIP branch
and a loader-side tvm-ffi-metal variant just landed in PR #506.
@danieldk pointed at two upstream blockers:

1. current_work_stream returns nullptr for kDLMetal in
   _build_optional_torch_c_dlpack.py
   — so a kernel can't find the producer's MTLCommandQueue.

2. DLTensor.data for MPS tensors is the buffer's contents pointer, not the
   MTLBuffer object — but [encoder setBuffer:offset:atIndex:] needs the
   buffer object, not a raw pointer.

(1) feels straightforward — add a kDLMetal branch using
at::mps::getCurrentMPSStream(). Happy to send that as a small PR.

(2) is where I'd love guidance on architecture before writing code:

• Should the MTLBuffer accessor live on DLPackExchangeAPI as a new
  callback (additive, but a dmlc/dlpack spec change)?
• Or should it be a tvm-ffi-only side API that doesn't touch the DLPack
  standard?
• Or something else entirely — e.g., reinterpret DLTensor.data as
  id<MTLBuffer> when device_type == kDLMetal?

I don't have strong opinions; I'd rather match whatever direction you'd
prefer. Happy to iterate on a design here before opening any PRs. Thanks!

[tqchen]

Thanks for the note. This actually needs a bit more thoughts. (2) is acually not that complicated, since DLTensor.data should actualy be id<MTLBuffer> (i believe it already is since for opaque cases like metal data should be the buffer object)

The main issue is (1), because ideally we want some reasonably ABI agnostic way of passing the current command queue for MPS, and ideally we also want it to be reasonably generic (e.g. ideally if things can work for mlx if possible). Unlike cuda, metal do not have the concept of default device and defacto default API to quickly launch a kernel, so commandqueue handling is somewhat done per framework. There is also not a common interface that we can rely on on job submission. At the momentum, there are different approaches in exposing things (e.g. does things eager vs lazy, how much command encoder reuse across kernels). One thing that could work is some form of generic interface and have an addon tvm-ffi-metal that allows registration of such interface, so torch variant can be passed and have the kernels depend on the addon, and when the respective torch extension get imported the default interface set to use torch's variant(while allowing other variants e.g. if it is desirable to lazily reuse command encoder which i believe is what mlx is doing for host perf reasons).

class TVMFFIMetalStreamInterface {
 public:
  virtual id<MTLComputeCommandEncoder> AcquireEncoder() = 0;
  virtual void ReleaseEncoder() = 0;
  virtual id<MTLDevice> GetDevice() = 0;
  virtual dispatch_queue_t GetDispatchQueue() = 0;
  virtual ~TVMFFIMetalStream() = default;
};

The above code can be a sketch of how this can looks like (but likely need more thoughts to really be confident that it is the stable low-level interface that can work). Taking a step back, one can of course directly call at::mps::getCurrentMPSStream() in tvm ffi defined interface, but of course that brings a direct dependency on libtorch.

[ved1beta]

Quick follow-up on (2) @tqchen you're right, I confirmed PyTorch's MPS DLPack
export already puts id<MTLBuffer> into DLTensor.data:

• The MPS allocator constructs c10::DataPtr directly from the
  id<MTLBuffer> (MPSAllocator.mm ), so storage().data() is
  the buffer object cast to void*.
• toDLPackImpl then writes that base pointer into DLTensor.data and
  sets byte_offset = storage_offset * elemsize (DLConvertor.cpp).

The DLPack spec covers this data is allowed to be an opaque handle
("cl_mem in OpenCL, may be opaque on some device types"), but the spec
doesn't explicitly call out Metal. Worth a one-line clarification to the
spec to name Metal alongside OpenCL? Could make a quick PR to dlpack if
that'd help future producers.

One small wrinkle: toDLPackNonOwning (the path tvm-ffi uses for the fast
dltensor_from_py_object_no_sync exchange) doesn't have the MPS
special-case — it stuffs data_ptr() (= MTLBuffer + storage_offset*elemsize)
into data with byte_offset=0, which produces an invalid id<MTLBuffer>
for sliced MPS tensors. That's a PyTorch bug I'll file separately, but
flagging in case it changes anything on the tvm-ffi side.

So (2) drops from scope. Back to (1) — happy to dig into the
AcquireEncoder/ReleaseEncoder design questions whenever you have time.

[tqchen]

On the Pytorch note. For backend where data is opaque and not addressable(corresponding to a real address on gpu), placing things in byte offset is indeed the right approach

[ved1beta]

shipped the toDLPackNonOwning fix: pytorch/pytorch#182924

A few things on the addon:

Singleton or thread-local? What if torch and mlx are both loaded?
Does AcquireEncoder fit mlx's lazy model, or would mlx need its own handle?
ReleaseEncoder = release last one on this thread?
Who calls GetDispatchQueue — kernel or framework?

[tqchen]

have't throw a lot of thoughts into it, likely we could use it to be TLS context(with global fallback) we could make use of RAII if needed, but even if it is set init to the framework first run would be good. the interface was meant for kernels that proxies APIs to framework

[oraluben]

btw: tilelang's metal backend would also benefit from this, currently we're using a very ad-hoc way for torch-metal sync: tile-ai/tvm#41

[ved1beta]

waiting on up stream i have a draft ready we ca start testing as soon as pytorch/pytorch#182924 gets in