tvm_tirx.tirx.script.ir_builder.triton.TritonKernel#

class tvm_tirx.tirx.script.ir_builder.triton.TritonKernel(func)#

A kernel from Triton JIT function.

This class bridges the Triton kernel with TVM runtime. The compilation includes the following steps:

Deduce the kernel signature and generate the Triton kernel
Embed the compiled kernel into the current IRModule as an external module
Generate a call to the Triton kernel following its calling convention via call_packed.

compile_to_device_module(launch_args, *args, **kwargs)#

Compile the kernel to a device module.

Parameters:

launch_args (List[int]) – The grid size of the kernel. A list of one to three expressions, representing the number of “blockIdx.x”, “blockIdx.y”, and “blockIdx.z” respectively.
args (List[Any]) – Arguments to the kernel function.
kwargs (Dict[str, Any]) – Additional options for the kernel compilation.

Return type:

tuple[str, Module, list[Any]]

tvm_tirx.tirx.script.ir_builder.triton.TritonKernel