triton.autotune¶

triton.autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None, cache_results=False)¶

用于自动调优triton.jit装饰函数的装饰器。

@triton.autotune(configs=[
    triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
    triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
  ],
  key=['x_size'] # the two above configs will be evaluated anytime
                 # the value of x_size changes
)
@triton.jit
def kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr):
    ...

Note:: 当所有配置都被评估后，内核将多次运行。这意味着内核更新的任何值都会被多次更新。为了避免这种不期望的行为，您可以使用reset_to_zero参数，该参数会在运行任何配置之前将提供的张量值重置为zero。

如果环境变量 TRITON_PRINT_AUTOTUNING 设置为 "1"，Triton 将在自动调优每个内核后向标准输出打印一条消息，包括自动调优所花费的时间和最佳配置。

Parameters:

configs (list[triton.Config]) – 一个包含triton.Config对象的列表
key (list[str]) - 一个参数名称列表，这些参数值的变化将触发所有提供配置的重新评估。
prune_configs_by – 用于剪枝配置的字典函数，包含以下字段： ‘perf_model’: 用于预测不同配置运行时间的性能模型，返回运行时间 ‘top_k’: 需要进行基准测试的配置数量 ‘early_config_prune’(可选): 用于早期剪枝的函数（例如num_stages）。它以configs:List[Config]作为输入，并返回剪枝后的配置。
reset_to_zero (list[str]) - 一个参数名称列表，在评估任何配置之前，这些参数的值将被重置为零。
restore_value (list[str]) - 一个参数名称列表，这些参数的值在评估任何配置后将被恢复。
pre_hook (lambda args, reset_only) – 一个在调用内核前会被调用的函数。这会覆盖用于'reset_to_zero'和'restore_value'的默认pre_hook。 'kwargs': 传递给内核的所有参数的字典。 'reset_only': 一个布尔值，表示pre_hook是否仅被调用来重置值，而没有对应的post_hook。
post_hook (lambda args, exception) - 一个在内核调用后将被调用的函数。这会覆盖用于'restore_value'的默认post_hook。 'kwargs': 传递给内核的所有参数的字典。 'exception': 在编译或运行时错误情况下由内核引发的异常。
warmup (int) – 传递给基准测试的预热时间(以毫秒计)(已弃用)。
rep (int) – 传递给基准测试的重复时间(以毫秒为单位)(已弃用)。
do_bench (lambda fn, quantiles) - 一个用于测量每次运行时间的基准测试函数。
cache_results – 是否将自动调优时间缓存到磁盘。默认为False。

“type cache_results: bool