spack solve: add a --profile option

We often want to know what the solver is thinking about in order to optimize it. We can
generate a profile of which atoms are most frequently propagated and/or undone by the
solver using a custom propagator class.

This adds a `spack.solver.profiler` module with a simple propagator that counts `undo`
and `propagate` operations and associates them with atoms in the solve. To simplify
things, rather than enumerating *every* atom in the logic program, we coarsen to show
only function names for everything but `attr()`, and we show the first argument of any
`attr()` atom (e.g. `attr("node")`).

`spack solve --profile` will print the top 40 most propagated atoms in the program,
along with how many times they are undone (these are usually but not always similar).
The list can give people working with the concretizer a rough idea of where the solver
is spending its time, and hopefully point to areas of the program to optimize.

This is based on some code that Tom Scogland wrote a year or so ago. I did the work of
integrating it as an option so that I do not have to keep looking it up and running this
from the CLI :).

Co-authored-by: Tom Scogland <scogland1@llnl.gov>
Signed-off-by: Todd Gamblin <tgamblin@llnl.gov>

lib/spack/spack/cmd/solve.py

@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import argparse
 import re
 import sys
 
@@ -48,6 +49,13 @@ def setup_parser(subparser):
     subparser.add_argument(
         "--stats", action="store_true", default=False, help="print out statistics from clingo"
     )
+    subparser.add_argument(
+        "--profile",
+        action="store_true",
+        default=False,
+        help="profile the solve phase and print out statistics on atoms",
+    )
+    subparser.add_argument("specs", nargs=argparse.REMAINDER, help="specs of packages")
 
     spack.cmd.spec.setup_parser(subparser)
 
@@ -147,6 +155,7 @@ def solve(parser, args):
             stats=args.stats,
             setup_only=setup_only,
             allow_deprecated=allow_deprecated,
+            profile=args.profile,
         )
         if not setup_only:
             _process_result(result, show, required_format, kwargs)

lib/spack/spack/solver/asp.py

@@ -64,6 +64,7 @@
     parse_term,
 )
 from .counter import FullDuplicatesCounter, MinimalDuplicatesCounter, NoDuplicatesCounter
+from .profiler import ProfilePropagator
 from .version_order import concretization_version_order
 
 GitOrStandardVersion = Union[spack.version.GitVersion, spack.version.StandardVersion]
@@ -800,7 +801,16 @@ def __init__(self, cores=True):
         # This attribute will be reset at each call to solve
         self.control = None
 
-    def solve(self, setup, specs, reuse=None, output=None, control=None, allow_deprecated=False):
+    def solve(
+        self,
+        setup,
+        specs,
+        reuse=None,
+        output=None,
+        control=None,
+        allow_deprecated=False,
+        profile=False,
+    ):
         """Set up the input and solve for dependencies of ``specs``.
 
         Arguments:
@@ -826,6 +836,11 @@ def solve(self, setup, specs, reuse=None, output=None, control=None, allow_depre
         # Initialize the control object for the solver
         self.control = control or default_clingo_control()
 
+        # if profiling is enabled, register a profiling propagator
+        if profile:
+            propagator = ProfilePropagator()
+            self.control.register_propagator(propagator)
+
         # ensure core deps are present on Windows
         # needs to modify active config scope, so cannot be run within
         # bootstrap config scope
@@ -925,13 +940,18 @@ def on_model(model):
             result.cores.extend(cores)
 
         if output.timers:
+            tty.msg("Timers:")
             timer.write_tty()
             print()
 
         if output.stats:
-            print("Statistics:")
+            tty.msg("Statistics:")
             pprint.pprint(self.control.statistics)
 
+        if profile:
+            tty.msg("Profile:")
+            propagator.print_profile(40)
+
         result.raise_if_unsat()
 
         if result.satisfiable and result.unsolved_specs and setup.concretize_everything:
@@ -4361,6 +4381,7 @@ def solve_with_stats(
         tests=False,
         setup_only=False,
         allow_deprecated=False,
+        profile=False,
     ):
         """
         Concretize a set of specs and track the timing and statistics for the solve
@@ -4382,7 +4403,12 @@ def solve_with_stats(
         setup = SpackSolverSetup(tests=tests)
         output = OutputConfiguration(timers=timers, stats=stats, out=out, setup_only=setup_only)
         return self.driver.solve(
-            setup, specs, reuse=reusable_specs, output=output, allow_deprecated=allow_deprecated
+            setup,
+            specs,
+            reuse=reusable_specs,
+            output=output,
+            allow_deprecated=allow_deprecated,
+            profile=profile,
         )
 
     def solve(self, specs, **kwargs):

lib/spack/spack/solver/profiler.py

@@ -0,0 +1,111 @@
+# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+"""Profiling propagator for clingo solves."""
+
+import re
+from typing import Any, Dict, List
+
+import llnl.util.tty as tty
+
+
+class Data:
+    """Counters for propagations and undos of an atom."""
+
+    __slots__ = ["atom", "literal", "prop", "undo"]
+
+    # currently we use Any for clingo types because clingo has a bunch of import
+    # wrappers around it that make typing difficult (see spack.solver.core for details)
+    def __init__(self, atom: Any, literal: int, prop: int, undo: int):
+        self.atom = atom
+        self.literal = literal
+        self.prop = prop
+        self.undo = undo
+
+
+class AggregatedData:
+    """Aggregated data for a profile, constructed from ``Data``.
+
+    We coarsen from atom granularity to string keys when aggregating.
+    """
+
+    __slots__ = ["name", "prop", "undo"]
+
+    def __init__(self, name: str, prop: int, undo: int):
+        self.name = name
+        self.prop = prop
+        self.undo = undo
+
+
+class ProfilePropagator:
+    """Profiling propagator for `spack solve --profile`.
+
+    Register this with the ``clingo.Control`` object to profile a solve.
+    """
+
+    _literal_to_atom: Dict
+    _profile: Dict[int, Data]
+
+    def init(self, init) -> None:
+        self._literal_to_atom = {}
+        self._profile = {}
+        for atom in init.symbolic_atoms:
+            solver_literal = init.solver_literal(atom.literal)
+            self._profile[solver_literal] = Data(atom, solver_literal, 0, 0)
+            init.add_watch(solver_literal)
+
+    def propagate(self, ctl, changes: List[int]) -> bool:
+        """Record a propagation in the solve."""
+        for literal in changes:
+            data = self._profile[literal]
+            data.prop += 1
+        return True
+
+    def undo(self, solver_id: int, assign, undo: List[int]) -> None:
+        """Record an undo in the solve."""
+        for literal in undo:
+            data = self._profile[literal]
+            data.undo += 1
+
+    def color_sym(self, string: str) -> str:
+        """Colorize a symbol for profile output"""
+        string = re.sub(r"^(\w+)", r"@C{\1}", string)
+        string = re.sub(r'("[^"]*")', r"@G{\1}", string)
+        string = re.sub(r"([\(\)])", r"@b{\1}", string)
+        return tty.color.colorize(string)
+
+    def key(self, atom) -> str:
+        """Convert an atom into an aggregate key for our profile.
+
+        Currently this compresses most things to their function name, and expands
+        ``attr("name", ...)`` to ``attr("name")`` so we can see which attributes affect
+        the solve most.
+
+        """
+        sym = atom.symbol
+        return f"attr({sym.arguments[0]})" if sym.name == "attr" else sym.name
+
+    def print_profile(self, n_atoms: int) -> None:
+        """Aggregate and print nicely formatted profile data."""
+        aggregated = {}
+        for data in self._profile.values():
+            name = self.key(data.atom)
+            if name not in aggregated:
+                aggregated[name] = AggregatedData(name, data.prop, data.undo)
+            else:
+                agg = aggregated[name]
+                agg.prop += data.prop
+                agg.undo += data.undo
+
+        values = sorted(
+            (x for x in aggregated.values() if x.prop), key=lambda x: x.prop, reverse=True
+        )
+
+        # format the output nicely
+        w = 10  # width for number fields
+        print(tty.color.colorize(f"  @*{{{'Prop':<{w}}{'Undo':<{w}}{'Symbol'}}}"))
+        for a in values[:n_atoms]:
+            print(f"  {a.prop:<{w}}{a.undo:<{w}}{self.color_sym(a.name)}")
+        if len(values) > n_atoms:
+            print("  ...")

lib/spack/spack/test/concretization/core.py

@@ -2871,6 +2871,18 @@ def test_specifying_different_versions_build_deps(self):
         assert any(x.satisfies(hdf5_str) for x in result.specs)
         assert any(x.satisfies(pinned_str) for x in result.specs)
 
+    def test_solve_with_profile(self, capsys):
+        """For now, just ensure that the profiler runs."""
+        solver = spack.solver.asp.Solver()
+        solver.solve([Spec("hdf5")], profile=True)
+
+        out, _ = capsys.readouterr()
+        assert "Profile:" in out
+        assert "Symbol" in out
+        assert "Prop" in out
+        assert "Undo" in out
+        assert "internal_error" in out  # symbol is always in small solves
+
 
 @pytest.mark.parametrize(
     "v_str,v_opts,checksummed",

share/spack/spack-completion.bash

@@ -1849,7 +1849,7 @@ _spack_restage() {
 _spack_solve() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --show --timers --stats -l --long -L --very-long -N --namespaces -I --install-status --no-install-status -y --yaml -j --json --format -c --cover -t --types -U --fresh --reuse --fresh-roots --reuse-deps --deprecated"
+        SPACK_COMPREPLY="-h --help --show --timers --stats --profile -l --long -L --very-long -N --namespaces -I --install-status --no-install-status -y --yaml -j --json --format -c --cover -t --types -U --fresh --reuse --fresh-roots --reuse-deps --deprecated"
     else
         _all_packages
     fi

share/spack/spack-completion.fish

@@ -2799,8 +2799,9 @@ complete -c spack -n '__fish_spack_using_command restage' -s h -l help -f -a hel
 complete -c spack -n '__fish_spack_using_command restage' -s h -l help -d 'show this help message and exit'
 
 # spack solve
-set -g __fish_spack_optspecs_spack_solve h/help show= timers stats l/long L/very-long N/namespaces I/install-status no-install-status y/yaml j/json format= c/cover= t/types U/fresh reuse fresh-roots deprecated
+set -g __fish_spack_optspecs_spack_solve h/help show= timers stats profile l/long L/very-long N/namespaces I/install-status no-install-status y/yaml j/json format= c/cover= t/types U/fresh reuse fresh-roots deprecated
 complete -c spack -n '__fish_spack_using_command_pos_remainder 0 solve' -f -k -a '(__fish_spack_specs_or_id)'
+complete -c spack -n '__fish_spack_using_command_pos_remainder 1 solve' -f -k -a '(__fish_spack_specs_or_id)'
 complete -c spack -n '__fish_spack_using_command solve' -s h -l help -f -a help
 complete -c spack -n '__fish_spack_using_command solve' -s h -l help -d 'show this help message and exit'
 complete -c spack -n '__fish_spack_using_command solve' -l show -r -f -a show
@@ -2809,6 +2810,8 @@ complete -c spack -n '__fish_spack_using_command solve' -l timers -f -a timers
 complete -c spack -n '__fish_spack_using_command solve' -l timers -d 'print out timers for different solve phases'
 complete -c spack -n '__fish_spack_using_command solve' -l stats -f -a stats
 complete -c spack -n '__fish_spack_using_command solve' -l stats -d 'print out statistics from clingo'
+complete -c spack -n '__fish_spack_using_command solve' -l profile -f -a profile
+complete -c spack -n '__fish_spack_using_command solve' -l profile -d 'profile the solve phase and print out statistics on atoms'
 complete -c spack -n '__fish_spack_using_command solve' -s l -l long -f -a long
 complete -c spack -n '__fish_spack_using_command solve' -s l -l long -d 'show dependency hashes as well as versions'
 complete -c spack -n '__fish_spack_using_command solve' -s L -l very-long -f -a very_long