Metal support: MTLBuffer access and current_work_stream for kDLMetal
- 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
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