flopscope 0.2.0__py3-none-any.whl

flopscope/_opt_einsum/_path_random.py

@@ -0,0 +1,462 @@
+"""Support for random optimizers, including the random-greedy path."""
+
+import functools
+import heapq
+import math
+import time
+from collections import deque
+from collections.abc import Generator, Iterable
+from decimal import Decimal
+from random import choices as random_choices
+from random import seed as random_seed
+from typing import Any
+
+from . import _paths as paths
+from ._helpers import compute_size_by_dict
+from ._subgraph_symmetry import SubgraphSymmetryOracle
+from ._typing import ArrayIndexType, ArrayType, PathType
+
+__all__ = ["RandomGreedy", "random_greedy", "random_greedy_128"]
+
+
+class RandomOptimizer(paths.PathOptimizer):
+    """Base class for running any random path finder that benefits
+    from repeated calling, possibly in a parallel fashion. Custom random
+    optimizers should subclass this, and the `setup` method should be
+    implemented with the following signature:
+
+    ```python
+    def setup(self, inputs, output, size_dict):
+        # custom preparation here ...
+        return trial_fn, trial_args
+    ```
+
+    Where `trial_fn` itself should have the signature::
+
+    ```python
+    def trial_fn(r, *trial_args):
+        # custom computation of path here
+        return ssa_path, cost, size
+    ```
+
+    Where `r` is the run number and could for example be used to seed a
+    random number generator. See `RandomGreedy` for an example.
+
+
+    Parameters:
+        max_repeats: The maximum number of repeat trials to have.
+        max_time: The maximum amount of time to run the algorithm for.
+        minimize:  Whether to favour paths that minimize the total estimated flop-count or
+            the size of the largest intermediate created.
+        parallel: Whether to parallelize the random trials, by default `False`. If
+            `True`, use a `concurrent.futures.ProcessPoolExecutor` with the same
+            number of processes as cores. If an integer is specified, use that many
+            processes instead. Finally, you can supply a custom executor-pool which
+            should have an API matching that of the python 3 standard library
+            module `concurrent.futures`. Namely, a `submit` method that returns
+            `Future` objects, themselves with `result` and `cancel` methods.
+        pre_dispatch: If running in parallel, how many jobs to pre-dispatch so as to avoid
+            submitting all jobs at once. Should also be more than twice the number
+            of workers to avoid under-subscription. Default: 128.
+
+    Attributes:
+        path: The best path found so far.
+        costs: The list of each trial's costs found so far.
+        sizes: The list of each trial's largest intermediate size so far.
+    """
+
+    def __init__(
+        self,
+        max_repeats: int = 32,
+        max_time: float | None = None,
+        minimize: str = "flops",
+        parallel: bool | Decimal | int = False,
+        pre_dispatch: int = 128,
+    ):
+        if minimize not in ("flops", "size"):
+            raise ValueError("`minimize` should be one of {'flops', 'size'}.")
+
+        self.max_repeats = max_repeats
+        self.max_time = max_time
+        self.minimize = minimize
+        self.better = paths.get_better_fn(minimize)
+        self._parallel: bool | Decimal | int = False
+        self.parallel = parallel
+        self.pre_dispatch = pre_dispatch
+
+        self.costs: list[int] = []
+        self.sizes: list[int] = []
+        self.best: dict[str, Any] = {"flops": float("inf"), "size": float("inf")}
+
+        self._repeats_start = 0
+        self._executor: Any
+        self._futures: Any
+
+    @property
+    def path(self) -> PathType:
+        """The best path found so far."""
+        return paths.ssa_to_linear(self.best["ssa_path"])
+
+    @property
+    def parallel(self) -> bool | Decimal | int:
+        return self._parallel
+
+    @parallel.setter
+    def parallel(self, parallel: bool | Decimal | int) -> None:
+        # shutdown any previous executor if we are managing it
+        if getattr(self, "_managing_executor", False):
+            self._executor.shutdown()
+
+        self._parallel = parallel
+        self._managing_executor = False
+
+        if parallel is False:
+            self._executor = None
+            return
+
+        if parallel is True:
+            from concurrent.futures import ProcessPoolExecutor
+
+            self._executor = ProcessPoolExecutor()
+            self._managing_executor = True
+            return
+
+        if isinstance(parallel, (int, Decimal)):
+            from concurrent.futures import ProcessPoolExecutor
+
+            self._executor = ProcessPoolExecutor(int(parallel))
+            self._managing_executor = True
+            return
+
+        # assume a pool-executor has been supplied
+        self._executor = parallel
+
+    def _gen_results_parallel(
+        self, repeats: Iterable[int], trial_fn: Any, args: Any
+    ) -> Generator[Any, None, None]:
+        """Lazily generate results from an executor without submitting all jobs at once."""
+        self._futures = deque()
+
+        # the idea here is to submit at least ``pre_dispatch`` jobs *before* we
+        # yield any results, then do both in tandem, before draining the queue
+        for r in repeats:
+            if len(self._futures) < self.pre_dispatch:
+                self._futures.append(self._executor.submit(trial_fn, r, *args))
+                continue
+            yield self._futures.popleft().result()
+
+        while self._futures:
+            yield self._futures.popleft().result()
+
+    def _cancel_futures(self) -> None:
+        if self._executor is not None:
+            for f in self._futures:
+                f.cancel()
+
+    def setup(
+        self,
+        inputs: list[ArrayIndexType],
+        output: ArrayIndexType,
+        size_dict: dict[str, int],
+        **kwargs: Any,
+    ) -> tuple[Any, Any]:
+        raise NotImplementedError
+
+    def __call__(
+        self,
+        inputs: list[ArrayIndexType],
+        output: ArrayIndexType,
+        size_dict: dict[str, int],
+        memory_limit: int | None = None,
+        symmetry_oracle: "SubgraphSymmetryOracle | None" = None,
+        **kwargs: Any,
+    ) -> PathType:
+        self._check_args_against_first_call(inputs, output, size_dict)
+
+        self._symmetry_oracle = symmetry_oracle
+
+        # start a timer?
+        t0 = 0.0  # assigned below if max_time is set
+        if self.max_time is not None:
+            t0 = time.time()
+
+        trial_fn, trial_args = self.setup(
+            inputs, output, size_dict, symmetry_oracle=symmetry_oracle, **kwargs
+        )
+
+        r_start = self._repeats_start + len(self.costs)
+        r_stop = r_start + self.max_repeats
+        repeats = range(r_start, r_stop)
+
+        # create the trials lazily
+        if self._executor is not None:
+            trials = self._gen_results_parallel(repeats, trial_fn, trial_args)
+        else:
+            trials = (trial_fn(r, *trial_args) for r in repeats)
+
+        # assess the trials
+        for ssa_path, cost, size in trials:
+            # keep track of all costs and sizes
+            self.costs.append(cost)
+            self.sizes.append(size)
+
+            # check if we have found a new best
+            found_new_best = self.better(
+                cost, size, self.best["flops"], self.best["size"]
+            )
+
+            if found_new_best:
+                self.best["flops"] = cost
+                self.best["size"] = size
+                self.best["ssa_path"] = ssa_path
+
+            # check if we have run out of time
+            if (self.max_time is not None) and (time.time() > t0 + self.max_time):
+                break
+
+        self._cancel_futures()
+        return self.path
+
+    def __del__(self):
+        # if we created the parallel pool-executor, shut it down
+        if getattr(self, "_managing_executor", False):
+            self._executor.shutdown()
+
+
+def thermal_chooser(queue, remaining, nbranch=8, temperature=1, rel_temperature=True):
+    """A contraction 'chooser' that weights possible contractions using a
+    Boltzmann distribution. Explicitly, given costs `c_i` (with `c_0` the
+    smallest), the relative weights, `w_i`, are computed as:
+
+        $$w_i = exp( -(c_i - c_0) / temperature)$$
+
+    Additionally, if `rel_temperature` is set, scale `temperature` by
+    `abs(c_0)` to account for likely fluctuating cost magnitudes during the
+    course of a contraction.
+
+    Parameters:
+        queue: The heapified list of candidate contractions.
+        remaining: Mapping of remaining inputs' indices to the ssa id.
+        temperature: When choosing a possible contraction, its relative probability will be
+            proportional to `exp(-cost / temperature)`. Thus the larger
+            `temperature` is, the further random paths will stray from the normal
+            'greedy' path. Conversely, if set to zero, only paths with exactly the
+            same cost as the best at each step will be explored.
+        rel_temperature: Whether to normalize the `temperature` at each step to the scale of
+            the best cost. This is generally beneficial as the magnitude of costs
+            can vary significantly throughout a contraction.
+        nbranch: How many potential paths to calculate probability for and choose from at each step.
+
+    Returns:
+        cost
+        k1
+        k2
+        k3
+    """
+    n = 0
+    choices = []
+    while queue and n < nbranch:
+        cost, k1, k2, k12 = heapq.heappop(queue)
+        if k1 not in remaining or k2 not in remaining:
+            continue  # candidate is obsolete
+        choices.append((cost, k1, k2, k12))
+        n += 1
+
+    if n == 0:
+        return None
+    if n == 1:
+        return choices[0]
+
+    costs = [choice[0][0] for choice in choices]
+    cmin = costs[0]
+
+    # adjust by the overall scale to account for fluctuating absolute costs
+    if rel_temperature:
+        temperature *= max(1, abs(cmin))
+
+    # compute relative probability for each potential contraction
+    if temperature == 0.0:
+        energies = [1 if c == cmin else 0 for c in costs]
+    else:
+        # shift by cmin for numerical reasons
+        energies = [math.exp(-(c - cmin) / temperature) for c in costs]
+
+    # randomly choose a contraction based on energies
+    (chosen,) = random_choices(range(n), weights=energies)
+    cost, k1, k2, k12 = choices.pop(chosen)
+
+    # put the other choice back in the heap
+    for other in choices:
+        heapq.heappush(queue, other)
+
+    return cost, k1, k2, k12
+
+
+def ssa_path_compute_cost(
+    ssa_path: PathType,
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    symmetry_oracle: "SubgraphSymmetryOracle | None" = None,
+) -> tuple[int, int]:
+    """Compute the flops and max size of an ssa path, using the symmetry oracle
+    if provided."""
+    inputs = list(map(frozenset, inputs))
+    output = frozenset(output)
+    remaining = set(range(len(inputs)))
+    num_operands = len(inputs)
+    ssa_to_subset: dict[int, frozenset[int]] = {
+        k: frozenset({k}) for k in range(num_operands)
+    }
+    total_cost = 0
+    max_size = 0
+
+    for i, j in ssa_path:
+        k12, flops12, _sym12 = paths.calc_k12_flops(
+            inputs,  # type: ignore[arg-type]
+            output,
+            remaining,  # type: ignore[arg-type]
+            i,
+            j,
+            size_dict,
+            oracle=symmetry_oracle,
+            ssa_to_subset=ssa_to_subset,
+        )
+        # Update ssa_to_subset for the new intermediate
+        new_ssa_id = len(inputs)
+        ssa_to_subset[new_ssa_id] = ssa_to_subset[i] | ssa_to_subset[j]
+
+        remaining.discard(i)
+        remaining.discard(j)
+        remaining.add(new_ssa_id)
+        inputs.append(k12)
+        total_cost += flops12
+        max_size = max(max_size, compute_size_by_dict(k12, size_dict))
+
+    return total_cost, max_size
+
+
+def _trial_greedy_ssa_path_and_cost(
+    r: int,
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    size_dict: dict[str, int],
+    choose_fn: Any,
+    cost_fn: Any,
+    symmetry_oracle: "SubgraphSymmetryOracle | None" = None,
+) -> tuple[PathType, int, int]:
+    """A single, repeatable, greedy trial run. **Returns:** ``ssa_path`` and cost."""
+    if r == 0:
+        # always start with the standard greedy approach
+        choose_fn = None
+
+    random_seed(r)
+
+    num_operands = len(inputs)
+    ssa_to_subset = {k: frozenset({k}) for k in range(num_operands)}
+
+    ssa_path = paths.ssa_greedy_optimize(
+        inputs,
+        output,
+        size_dict,
+        choose_fn,
+        cost_fn,
+        symmetry_oracle=symmetry_oracle,
+        ssa_to_subset=ssa_to_subset,
+    )
+    cost, size = ssa_path_compute_cost(
+        ssa_path, inputs, output, size_dict, symmetry_oracle=symmetry_oracle
+    )
+
+    return ssa_path, cost, size
+
+
+class RandomGreedy(RandomOptimizer):
+    def __init__(
+        self,
+        cost_fn: str = "memory-removed-jitter",
+        temperature: float = 1.0,
+        rel_temperature: bool = True,
+        nbranch: int = 8,
+        **kwargs: Any,
+    ):
+        """Parameters:
+        cost_fn: A function that returns a heuristic 'cost' of a potential contraction
+                with which to sort candidates. Should have signature
+                `cost_fn(size12, size1, size2, k12, k1, k2)`.
+        temperature: When choosing a possible contraction, its relative probability will be
+                proportional to `exp(-cost / temperature)`. Thus the larger
+                `temperature` is, the further random paths will stray from the normal
+                'greedy' path. Conversely, if set to zero, only paths with exactly the
+                same cost as the best at each step will be explored.
+        rel_temperature: Whether to normalize the ``temperature`` at each step to the scale of
+                the best cost. This is generally beneficial as the magnitude of costs
+                can vary significantly throughout a contraction. If False, the
+                algorithm will end up branching when the absolute cost is low, but
+                stick to the 'greedy' path when the cost is high - this can also be
+                beneficial.
+        nbranch: How many potential paths to calculate probability for and choose from at each step.
+        kwargs: Supplied to RandomOptimizer.
+        """
+        self.cost_fn = cost_fn
+        self.temperature = temperature
+        self.rel_temperature = rel_temperature
+        self.nbranch = nbranch
+        super().__init__(**kwargs)
+
+    @property
+    def choose_fn(self) -> Any:
+        """The function that chooses which contraction to take - make this a
+        property so that ``temperature`` and ``nbranch`` etc. can be updated
+        between runs.
+        """
+        if self.nbranch == 1:
+            return None
+
+        return functools.partial(
+            thermal_chooser,
+            temperature=self.temperature,  # type: ignore[arg-type]
+            nbranch=self.nbranch,
+            rel_temperature=self.rel_temperature,
+        )
+
+    def setup(
+        self,
+        inputs: list[ArrayIndexType],
+        output: ArrayIndexType,
+        size_dict: dict[str, int],
+        symmetry_oracle: "SubgraphSymmetryOracle | None" = None,
+        **kwargs: Any,
+    ) -> tuple[Any, Any]:
+        fn = _trial_greedy_ssa_path_and_cost
+        args = (
+            inputs,
+            output,
+            size_dict,
+            self.choose_fn,
+            self.cost_fn,
+            symmetry_oracle,
+        )
+        return fn, args
+
+
+def random_greedy(
+    inputs: list[ArrayIndexType],
+    output: ArrayIndexType,
+    idx_dict: dict[str, int],
+    memory_limit: int | None = None,
+    symmetry_oracle: "SubgraphSymmetryOracle | None" = None,
+    **optimizer_kwargs: Any,
+) -> ArrayType:
+    """A simple wrapper around the RandomGreedy optimizer."""
+    optimizer = RandomGreedy(**optimizer_kwargs)
+    return optimizer(
+        inputs,
+        output,
+        idx_dict,
+        memory_limit,
+        symmetry_oracle=symmetry_oracle,
+    )
+
+
+random_greedy_128 = functools.partial(random_greedy, max_repeats=128)