triton.heuristics

triton.heuristics(values)

用于指定如何计算某些元参数值的装饰器。这对于自动调优成本过高或根本不适用的情况很有用。

# 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]]) – 包含元参数名称和计算元参数值函数的字典。每个这样的函数都接受一个位置参数列表作为输入。