ninetoothed.jit

ninetoothed.jit(func=None, *, caller='torch', kernel_name=None, num_warps=None, num_stages=None, max_num_configs=None, _prettify=False)

A decorator for generating compute kernels.

Parameters:

• func – The function to be compiled.

• caller – Who will call the compute kernel.

• kernel_name – The name for the generated kernel.

• num_warps – The number of warps to use.

• num_stages – The number of pipeline stages.

• max_num_configs – The maximum number of auto-tuning configurations to use.

• _prettify – Whether to prettify the generated code.

Returns:

A handle to the compute kernel.

Note

The _prettify parameter is experimental, which might break the generated code.