[tvm-ffi] Enable tvm-ffi for metal backend

[oraluben]

This PR enables tvm-ffi backend for metal device, and is still under development.

It explicitly passes command buffer maintained by torch to tvm and reuse that buffer.
To run example, checkout tvm: tile-ai/tvm#21 and run TILELANG_DISABLE_CACHE=1 python test.py

Current example hardcode necessary transformation by register_c_postproc_callback. To archive our goal, the host runtime passes torch::mps::get_command_buffer() to tvm and force it to use it in the module. The runtime module also needs extra torch dependency.

Ideally, those transformation should be done by a tvm pass, and the changes to metal_module.mm doesn't look very elegant. It would be great if someone familiar with tvm could give some advice.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[oraluben]

Some update:

tvm's compiled device code into binary and embed them to runtime mod (inside f"{pack_lib_prefix}devc.c" is c array containing compiled device code):

https://github.com/tile-ai/tvm/blob/03ad7ccd20058ad9576571375771129a98e72756/python/tvm/runtime/module.py#L288-L304

What we need is to fork the current metal codegen from tvm, make sure it's aware of torch's mps context, and see if it works well.

[LeiWang1999]

@oraluben Thanks and feel free to connect if you need any help :)