apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.67k stars 3.45k forks source link

[TIR, TVMScript] Add TIR - Triton integration #17395

Closed vinx13 closed 1 week ago

vinx13 commented 2 weeks ago

Added a macro T.call_triton in TIR script parser, which expands to AOT compilation of the kernel and the host TIR code to launch the kernel.

cc @tqchen @cyx-6

tqchen commented 2 weeks ago

This is a great mechanism to integrate kernel generators. Some notes on design that might help generalize it abit

Would be great to change the intrinsic to T.call_kernel, which checks the first parameter for kernel types and dispatch accordingly.

Overall Mechanism