triton.heuristics

triton.heuristics(values)

用于指定如何计算某些元参数(meta-parameters)值的装饰器。这在自动调优代价过高或不适用的情况下非常有用。

# smallest power-of-two >= x_size
@triton.heuristics(values={'BLOCK_SIZE': lambda args: triton.next_power_of_2(args['x_size'])})
@triton.jit
def kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr):
    ...
参数:

values (dict[str, Callable[[dict[str, Any]], Any]]) – 一个包含元参数名称及其计算函数的字典。每个此类函数都接受一个位置参数列表作为输入。