Notebook CI Job -- report times and fix regression (#1746)

- fixes performance regression introduced in #1732 
- After some experimentation, run on 8 processes with
maxtasksperworker=1
 - Report 10 slowest notebooks

Author: mpharrigan

.github/workflows/pr.yaml

@@ -76,7 +76,7 @@ jobs:
           pip install -r dev_tools/requirements/envs/pytest.env.txt
           pip install --no-deps -e .
       - run: |
-          python dev_tools/execute-notebooks.py
+          python dev_tools/execute-notebooks.py --n-workers=8
         env:
           NUMBA_NUM_THREADS: 4
 

dev_tools/execute-notebooks.py

@@ -22,9 +22,13 @@ def parse_args():
     p.add_argument('--output-nbs', action=argparse.BooleanOptionalAction, default=True)
     p.add_argument('--output-md', action=argparse.BooleanOptionalAction, default=False)
     p.add_argument('--only-out-of-date', action=argparse.BooleanOptionalAction, default=True)
+    p.add_argument('--n-workers', type=int, default=None)
     args = p.parse_args()
     execute_and_export_notebooks(
-        output_nbs=args.output_nbs, output_md=args.output_md, only_out_of_date=args.only_out_of_date
+        output_nbs=args.output_nbs,
+        output_md=args.output_md,
+        only_out_of_date=args.only_out_of_date,
+        n_workers=args.n_workers,
     )
 
 

dev_tools/qualtran_dev_tools/notebook_execution.py

@@ -13,11 +13,14 @@
 #  limitations under the License.
 import asyncio
 import multiprocessing
+import random
 import subprocess
 import sys
+import time
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import attrs
 import filelock
 import nbconvert
 import nbformat
@@ -180,6 +183,13 @@ def execute_and_export_notebook(paths: _NBInOutPaths) -> Optional[Exception]:
     return None
 
 
+@attrs.frozen
+class _NotebookRunResult:
+    nb_in: Path
+    err: Optional[Exception]
+    duration_s: float
+
+
 class _NotebookRunClosure:
     """Used to run notebook execution logic in subprocesses."""
 
@@ -189,7 +199,7 @@ def __init__(self, reporoot: Path, output_nbs: bool, output_md: bool, only_out_o
         self.output_md = output_md
         self.only_out_of_date = only_out_of_date
 
-    def __call__(self, nb_rel_path: Path, sourceroot: Path) -> Tuple[Path, Optional[Exception]]:
+    def __call__(self, nb_rel_path: Path, sourceroot: Path) -> _NotebookRunResult:
         paths = _NBInOutPaths.from_nb_rel_path(
             nb_rel_path,
             reporoot=self.reporoot,
@@ -200,15 +210,21 @@ def __call__(self, nb_rel_path: Path, sourceroot: Path) -> Tuple[Path, Optional[
 
         if self.only_out_of_date and not paths.needs_reexport():
             print(f'{nb_rel_path} up to date')
-            return paths.nb_in, None
+            return _NotebookRunResult(paths.nb_in, None, 0.0)
 
+        start = time.time()
         err = execute_and_export_notebook(paths)
-        print(f"Exported {nb_rel_path}")
-        return paths.nb_in, err
+        end = time.time()
+        print(f"Exported {nb_rel_path} in {end-start:.2f} seconds.")
+        return _NotebookRunResult(paths.nb_in, err, duration_s=end - start)
 
 
 def execute_and_export_notebooks(
-    *, output_nbs: bool, output_md: bool, only_out_of_date: bool = True
+    *,
+    output_nbs: bool,
+    output_md: bool,
+    only_out_of_date: bool = True,
+    n_workers: Optional[int] = None,
 ):
     """Find, execute, and export all checked-in ipynbs.
 
@@ -217,23 +233,41 @@ def execute_and_export_notebooks(
         output_md: Whether to save the executed notebooks as markdown
         only_out_of_date: Only re-execute and re-export notebooks whose output files
             are out of date.
+        n_workers: If set to 1, do not use parallelization. If set to `None` (the detault),
+            `multiprocessing.Pool()` will be used, which uses the number of processors as
+            a default. Otherwise, this argument is passed to
+            `multiprocessing.Pool(n_workers)` to execute notebooks in parallel on this many
+            worker processes..
     """
     reporoot = get_git_root()
     nb_rel_paths = get_nb_rel_paths(sourceroot=reporoot / 'qualtran')
     nb_rel_paths += get_nb_rel_paths(sourceroot=reporoot / 'tutorials')
+    random.shuffle(nb_rel_paths)
+    print(f"Found {len(nb_rel_paths)} notebooks.")
     func = _NotebookRunClosure(
         reporoot=reporoot,
         output_nbs=output_nbs,
         output_md=output_md,
         only_out_of_date=only_out_of_date,
     )
-    with multiprocessing.Pool() as pool:
-        results = pool.starmap(func, nb_rel_paths)
-    bad_nbs = [nbname for nbname, err in results if err is not None]
+    if n_workers == 1:
+        print("(Not using multiprocessing, n_workers=1)")
+        results = [func(nb_rel_path, sourceroot) for nb_rel_path, sourceroot in nb_rel_paths]
+    else:
+        print(f"Multiprocessing with {n_workers=}")
+        with multiprocessing.Pool(n_workers, maxtasksperchild=1) as pool:
+            results = pool.starmap(func, nb_rel_paths)
+        assert results
+    bad_nbs = [result.nb_in for result in results if result.err is not None]
 
     if len(bad_nbs) > 0:
         print()
         print("Errors in notebooks:")
         for nb in bad_nbs:
             print(' ', nb)
         sys.exit(1)
+
+    duration_nbs = sorted(results, key=lambda r: r.duration_s, reverse=True)
+    print("Slowest 10 notebooks:")
+    for result in duration_nbs[:10]:
+        print(f'{result.duration_s:5.2f}s {result.nb_in}')

qualtran/bloqs/block_encoding/lcu_block_encoding_test.py

@@ -48,7 +48,7 @@ def test_ctrl_lcu_be_cost():
     assert bloq.controlled(CtrlSpec(cvs=0)) == attrs.evolve(bloq, control_val=0)
 
     assert get_cost_value(bloq.controlled(), QECGatesCost()) == GateCounts(
-        cswap=28, and_bloq=77, clifford=438, rotation=16, measurement=77
+        cswap=28, and_bloq=77, clifford=294, rotation=16, measurement=77
     )
 
 

qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard.py

@@ -39,7 +39,7 @@
 """
 
 from functools import cached_property
-from typing import Dict, Iterator, Optional, Tuple
+from typing import Dict, Iterator, Optional, Set, Tuple, Union
 
 import attrs
 import cirq
@@ -62,13 +62,15 @@
     Signature,
     SoquetT,
 )
-from qualtran.bloqs.basic_gates import CSwap, TwoBitCSwap
+from qualtran.bloqs.basic_gates import CNOT, CSwap, CZ, TwoBitCSwap
 from qualtran.bloqs.bookkeeping import Join2, Split2
+from qualtran.bloqs.mcmt import And
 from qualtran.bloqs.multiplexers.apply_gate_to_lth_target import ApplyGateToLthQubit
 from qualtran.bloqs.multiplexers.select_base import SelectOracle
 from qualtran.bloqs.multiplexers.selected_majorana_fermion import SelectedMajoranaFermion
 from qualtran.cirq_interop import decompose_from_cirq_style_method
 from qualtran.drawing import Circle, TextBox, WireSymbol
+from qualtran.resource_counting import BloqCountDictT, BloqCountT, SympySymbolAllocator
 from qualtran.symbolics import ceil, is_symbolic, log2, SymbolicInt
 
 
@@ -275,6 +277,10 @@ def log_m(self) -> SymbolicInt:
             raise NotImplementedError("Currently only supports the case where x_dim=y_dim.")
         return ceil(log2(self.x_dim))
 
+    @cached_property
+    def N(self):
+        return 2 * self.x_dim * self.y_dim
+
     @cached_property
     def control_registers(self) -> Tuple[Register, ...]:
         return () if self.control_val is None else (Register('control', QBit()),)
@@ -289,7 +295,7 @@ def selection_registers(self) -> Tuple[Register, ...]:
 
     @cached_property
     def target_registers(self) -> Tuple[Register, ...]:
-        return (Register('target', QAny(self.x_dim * self.y_dim * 2)),)
+        return (Register('target', QAny(self.N)),)
 
     @cached_property
     def signature(self) -> Signature:
@@ -330,6 +336,15 @@ def build_composite_bloq(
             x, y, spin, target = bb.add_from(smf, x=x, y=y, spin=spin, target=target)
             return {'x': x, 'y': y, 'spin': spin, 'target': target}
 
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
+
+        count = self.N - 1
+        if self.control_val is None:
+            count -= 1
+        return {And(): count, And().adjoint(): count}
+
     def wire_symbol(
         self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple()
     ) -> 'WireSymbol':
@@ -473,6 +488,15 @@ def build_composite_bloq(
 
         return ret | {'x': x, 'y': y, 'target': target}
 
+    def build_call_graph(
+        self, ssa: 'SympySymbolAllocator'
+    ) -> Union['BloqCountDictT', Set['BloqCountT']]:
+        half_N = self.x_dim * self.y_dim
+        count = half_N
+        if self.control_val is None:
+            count -= 1
+        return {And(): count, And().adjoint(): count, CNOT(): half_N - 1, CZ(): half_N}
+
     def wire_symbol(
         self, reg: Optional['Register'], idx: Tuple[int, ...] = tuple()
     ) -> 'WireSymbol':

qualtran/bloqs/chemistry/hubbard_model/qubitization/select_hubbard_test.py

@@ -16,8 +16,13 @@
 import numpy as np
 import pytest
 
+import qualtran.testing as qlt_testing
 from qualtran import QAny, QUInt
-from qualtran.bloqs.chemistry.hubbard_model.qubitization import HubbardSpinUpZ, SelectHubbard
+from qualtran.bloqs.chemistry.hubbard_model.qubitization import (
+    HubbardMajorannaOperator,
+    HubbardSpinUpZ,
+    SelectHubbard,
+)
 from qualtran.bloqs.chemistry.hubbard_model.qubitization.select_hubbard import (
     _hubb_majoranna,
     _hubb_majoranna_small,
@@ -26,6 +31,12 @@
     _sel_hubb,
 )
 from qualtran.resource_counting import GateCounts, get_cost_value, QECGatesCost
+from qualtran.resource_counting.generalizers import (
+    generalize_cvs,
+    ignore_alloc_free,
+    ignore_cliffords,
+    ignore_split_join,
+)
 from qualtran.simulation.classical_sim import do_phased_classical_simulation
 
 
@@ -62,6 +73,22 @@ def test_hubbard_majoranna_small(bloq_autotester):
     bloq_autotester(_hubb_majoranna_small)
 
 
+def test_hubb_majoranna_bloq_counts():
+    for n in [2, 3, 10, 11]:
+        bloq = HubbardMajorannaOperator(x_dim=n, y_dim=n, control_val=None)
+        qlt_testing.assert_equivalent_bloq_counts(
+            bloq,
+            generalizer=[ignore_cliffords, ignore_alloc_free, ignore_split_join, generalize_cvs],
+        )
+
+    for n in [2, 3, 10, 11]:
+        bloq = HubbardMajorannaOperator(x_dim=n, y_dim=n, control_val=1)
+        qlt_testing.assert_equivalent_bloq_counts(
+            bloq,
+            generalizer=[ignore_cliffords, ignore_alloc_free, ignore_split_join, generalize_cvs],
+        )
+
+
 def test_hubbard_spin_up_z_symb(bloq_autotester):
     bloq_autotester(_hubb_spin_up_z)
 
@@ -70,6 +97,20 @@ def test_hubbard_spin_up_z(bloq_autotester):
     bloq_autotester(_hubb_spin_up_z_small)
 
 
+def test_hubbard_spin_up_z_bloq_counts():
+    for n in [2, 3, 10, 11]:
+        bloq = HubbardSpinUpZ(x_dim=n, y_dim=n, control_val=None)
+        qlt_testing.assert_equivalent_bloq_counts(
+            bloq, generalizer=[ignore_alloc_free, ignore_split_join, generalize_cvs]
+        )
+
+    for n in [2, 3, 10, 11]:
+        bloq = HubbardSpinUpZ(x_dim=n, y_dim=n, control_val=1)
+        qlt_testing.assert_equivalent_bloq_counts(
+            bloq, generalizer=[ignore_alloc_free, ignore_split_join, generalize_cvs]
+        )
+
+
 def test_hubbard_spin_up_z_classical():
     rng = np.random.default_rng(52)
 

qualtran/resource_counting/generalizers.py

@@ -43,9 +43,17 @@ def _ignore_wrapper(f: Callable[[Bloq], Optional[Bloq]], b: Bloq) -> Optional[Bl
 
 def ignore_split_join(b: Bloq) -> Optional[Bloq]:
     """A generalizer that ignores split and join operations."""
-    from qualtran.bloqs.bookkeeping import AutoPartition, Cast, Join, Partition, Split
-
-    if isinstance(b, (Split, Join, Partition, Cast)):
+    from qualtran.bloqs.bookkeeping import (
+        AutoPartition,
+        Cast,
+        Join,
+        Join2,
+        Partition,
+        Split,
+        Split2,
+    )
+
+    if isinstance(b, (Split, Split2, Join, Join2, Partition, Cast)):
         return None
     if isinstance(b, AutoPartition):
         return ignore_split_join(b.bloq)