xma.autotuner.tuner

class AutotunedFunction(function: Callable, configs: list[AutotuneConfig], triggers: set[str], warmup_iterations: int, benchmark_iterations: int, functional_triggers: dict[str, Callable] = {}, reset_to_zero: dict = {})

Bases: object

property exposed_signature: Signature

autotune(configs: list[AutotuneConfig], triggers: set[str] = {}, functional_triggers: dict[str, Callable] = {}, warmup_iterations: int = 5, benchmark_iterations: int = 10, reset_to_zero: dict = {}) → AutotunedFunction

autotuner for any function or kernel

Parameters:

• configs (list[AutotuneConfig]) – list of configs to autotune over

• triggers (set[str]) – change in these parameters will trigger autotuning

• functional_triggers (dict[str, Callable]) – key, function mapping. change in the function outputs will trigger autotuning.

• warmup_iterations (int) – iterations for warmup. Defaults to 5.

• benchmark_iterations (int) – iterations for benchmarking. Defaults to 10.

• reset_to_zero (dict) – A dictionary mapping tensor argument names to an optional callable condition. The specified tensors will be zeroed out after each benchmark iteration if the condition (if provided) returns True.

Returns:

autotuned version of the function

Return type:

_Autotune