Skip to content

Commit

Permalink
modularize fetching the number of candidates
Browse files Browse the repository at this point in the history
  • Loading branch information
kaushikcfd committed Oct 29, 2024
1 parent 7b2d30d commit 9783858
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions pyop2/transforms/auto_tiling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down Expand Up @@ -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])

Expand Down

0 comments on commit 9783858

Please sign in to comment.