| 102 | return self.fn.run(*args, num_warps=config.num_warps, num_stages=config.num_stages, **kwargs, **config.kwargs) |
| 103 | |
| 104 | def prune_configs(self, kwargs): |
| 105 | pruned_configs = self.configs |
| 106 | if self.early_config_prune: |
| 107 | pruned_configs = self.early_config_prune(self.configs, self.nargs) |
| 108 | if self.perf_model: |
| 109 | top_k = self.configs_top_k |
| 110 | if isinstance(top_k, float) and top_k <= 1.0: |
| 111 | top_k = int(len(self.configs) * top_k) |
| 112 | if len(pruned_configs) > top_k: |
| 113 | est_timing = { |
| 114 | config: self.perf_model(**self.nargs, **kwargs, **config.kwargs, num_stages=config.num_stages, |
| 115 | num_warps=config.num_warps) |
| 116 | for config in pruned_configs |
| 117 | } |
| 118 | pruned_configs = sorted(est_timing.keys(), key=lambda x: est_timing[x])[:top_k] |
| 119 | return pruned_configs |
| 120 | |
| 121 | def warmup(self, *args, **kwargs): |
| 122 | self.nargs = dict(zip(self.arg_names, args)) |