From 9783858bf7c5e04d0044098e570bec7777e83c20 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 29 Oct 2024 14:56:57 -0500 Subject: [PATCH] modularize fetching the number of candidates --- pyop2/transforms/auto_tiling.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/pyop2/transforms/auto_tiling.py b/pyop2/transforms/auto_tiling.py index fdbafde85..96201ce27 100644 --- a/pyop2/transforms/auto_tiling.py +++ b/pyop2/transforms/auto_tiling.py @@ -1766,7 +1766,10 @@ def estimated_exec_time(self, tiling_config): ) return (total_time, nsync) - def __call__(self) -> Tuple[ParametricTiling, ...]: + def _get_all_configurations(self) -> FrozenSet[ParametricTiling]: + """ + Returns a :class:`frozenset` of all configurations in our search space. + """ from itertools import product threads_to_cells = {} @@ -1845,9 +1848,11 @@ def get_eta_shared_mem_alias(tiles): False, ) ) + return frozenset(params) - # sort the parameters with highest occupancy. - params.sort(key=lambda P: self.estimated_exec_time(P)) + def __call__(self) -> Tuple[ParametricTiling, ...]: + params = sorted(self._get_all_configurations(), + key=self.estimated_exec_time) return tuple(params[: self.num_param_tiling_candidates])