apache / apache/tvm-ffi

Metal support: MTLBuffer access and current_work_stream for kDLMetal

Open
#578 7 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
C++
Stars
462
Forks
99
Avg merge
8h 37m
Merged PRs (30d)
73

Description

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](https://github.com/huggingface/kernels/tree/tvm-ffi-metal)
and a loader-side `tvm-ffi-metal` variant just landed in [PR #506](https://github.com/huggingface/kernels/pull/506).
@danieldk [pointed at two upstream blockers](https://github.com/huggingface/kernels/pull/506#issuecomment-4322456165):

1. `current_work_stream` returns `nullptr` for `kDLMetal` in
[`_build_optional_torch_c_dlpack.py`](https://github.com/apache/tvm-ffi/blob/main/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py#L508)
— 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` 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!

Contributor guide

Open the contributing guide

Research direction

Start with python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py around the kDLMetal current_work_stream handling, then review the DLPackExchangeAPI options described here. Done means the maintainers agree on an architecture for exposing MTLBuffer access and define a scoped implementation path alongside the Metal stream behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend-api-design
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.