| 12 | |
| 13 | |
| 14 | class Autotuner(triton.KernelInterface): |
| 15 | def __init__(self, fn, arg_names, configs, key, reset_to_zero, prune_configs_by: Dict = None, nearest_power_of_two: bool = False): |
| 16 | ''' |
| 17 | :param prune_configs_by: a dict of functions that are used to prune configs, fields: |
| 18 | 'perf_model': performance model used to predicate running time with different configs, returns running time |
| 19 | 'top_k': number of configs to bench |
| 20 | 'prune_num_stages_by'(optional): a function used to prune num_stages. It take configs:List[Config] as its input, and returns pruned configs. |
| 21 | 'nearest_power_of_two'(optional): whether to round key arguments to the nearest power of two when caching tuning results |
| 22 | ''' |
| 23 | if not configs: |
| 24 | self.configs = [triton.Config({}, num_warps=4, num_stages=2)] |
| 25 | else: |
| 26 | self.configs = configs |
| 27 | self.key_idx = [arg_names.index(k) for k in key] |
| 28 | self.nearest_power_of_two = nearest_power_of_two |
| 29 | self.cache = {} |
| 30 | # hook to reset all required tensor to zeros before relaunching a kernel |
| 31 | self.hook = lambda args: 0 |
| 32 | if reset_to_zero is not None: |
| 33 | self.reset_idx = [arg_names.index(k) for k in reset_to_zero] |
| 34 | |
| 35 | def _hook(args): |
| 36 | for i in self.reset_idx: |
| 37 | args[i].zero_() |
| 38 | self.hook = _hook |
| 39 | self.arg_names = arg_names |
| 40 | # prune configs |
| 41 | if prune_configs_by: |
| 42 | perf_model, top_k = prune_configs_by['perf_model'], prune_configs_by['top_k'] |
| 43 | if 'early_config_prune' in prune_configs_by: |
| 44 | early_config_prune = prune_configs_by['early_config_prune'] |
| 45 | else: |
| 46 | perf_model, top_k, early_config_prune = None, None, None |
| 47 | self.perf_model, self.configs_top_k = perf_model, top_k |
| 48 | self.early_config_prune = early_config_prune |
| 49 | self.fn = fn |
| 50 | |
| 51 | def _bench(self, *args, config, **meta): |
| 52 | # check for conflicts, i.e. meta-parameters both provided |
| 53 | # as kwargs and by the autotuner |
| 54 | conflicts = meta.keys() & config.kwargs.keys() |
| 55 | if conflicts: |
| 56 | raise ValueError( |
| 57 | f"Conflicting meta-parameters: {', '.join(conflicts)}." |
| 58 | " Make sure that you don't re-define auto-tuned symbols." |
| 59 | ) |
| 60 | # augment meta-parameters with tunable ones |
| 61 | current = dict(meta, **config.kwargs) |
| 62 | |
| 63 | def kernel_call(): |
| 64 | if config.pre_hook: |
| 65 | config.pre_hook(self.nargs) |
| 66 | self.hook(args) |
| 67 | self.fn.run(*args, num_warps=config.num_warps, num_stages=config.num_stages, **current) |
| 68 | try: |
| 69 | # In testings using only 40 reps seems to be close enough and it appears to be what PyTorch uses |
| 70 | # PyTorch also sets fast_flush to True, but I didn't see any speedup so I'll leave the default |
| 71 | return triton.testing.do_bench(kernel_call, rep=40) |