[TIR, TVMScript] Add TIR - Triton integration

[vinx13]

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]

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.

• Let us think of a base class tir.kernel.BaseKernel which implements the base methods needed
  • compile_to_device_module(
    • calls get_meta_data to get the meta_data
  • have a registry from class name to Kernel constructor(can be in next PR)
    • So the construct and import of related kernel can be done only when we see the related class name
    • call_kernel will lookup registry and construct related kernels if needed
  • TritonKernel subclasses that, call_kernel will lookup the constructor, construct it and call related class method to generate the necessary downstream classes.
• Let us also add a tir.kernel.CUDAKernel as an example extension pt
  • CUDAKernel takes in cuda C source device code, and ensure call_kernel mechanism works the same

Overall Mechanism

• A new DSL/kernel defines a subclass of tir.kernel.Kernel
  • the subclass describes the mechanism on how to compile the source to the device module
  • it also registers into the mapping a class name => Kernel type
• User can either manually construct a instance of tir.kernel.Kernel, or pass in the original data structure triton.JITFunction and call_kernel leverages the registered mapping to do automatic conversion.
• call_kernel can contain possible specialization hints(constants etc)

[tqchen]

add a testcase of rewrite usecase

[tqchen]

as a followup, add a function to check if the module is_device_module, this should include cuda, rocm, webgpu, vulkan, opencl