[adlantz]

Issue #336 asks for improvements to the tfq.convert_to_tensor speed function. The majorirty of the run time occurs in the cirq serialization logic:

============================================================
Where does convert_to_tensor spend its time?
============================================================
  tensorflow_quantum 0.7.7 | cirq 1.5.0 | python 3.10.14
  serialize_circuit: deepcopy (baseline)  [git HEAD]
  best of 3 runs per phase

100 x 10q depth 20
  --------------------------------------------------------
  convert_to_tensor (total)                  397.4 ms   100.0%
  --------------------------------------------------------
  1. serialize_circuit  (cirq, Python)       392.4 ms    98.7%
  2. SerializeToString  (protobuf)             3.3 ms     0.8%
  3. remainder          (tensor packing)       1.7 ms     0.4%

100 x 20q depth 50
  --------------------------------------------------------
  convert_to_tensor (total)                 1913.2 ms   100.0%
  --------------------------------------------------------
  1. serialize_circuit  (cirq, Python)      1941.8 ms   101.5%
  2. SerializeToString  (protobuf)            18.9 ms     1.0%
  3. remainder          (tensor packing)     -47.6 ms    -2.5%

50 x 30q depth 100
  --------------------------------------------------------
  convert_to_tensor (total)                 2944.9 ms   100.0%
  --------------------------------------------------------
  1. serialize_circuit  (cirq, Python)      2938.0 ms    99.8%
  2. SerializeToString  (protobuf)            29.5 ms     1.0%
  3. remainder          (tensor packing)     -22.6 ms    -0.8%

100 x 20q depth 50 (+ctrl)
  --------------------------------------------------------
  convert_to_tensor (total)                 1960.2 ms   100.0%
  --------------------------------------------------------
  1. serialize_circuit  (cirq, Python)      1947.9 ms    99.4%
  2. SerializeToString  (protobuf)            19.7 ms     1.0%
  3. remainder          (tensor packing)      -7.4 ms    -0.4%

Phases are timed independently, so the remainder is a difference of minima. Small negatives are noise, read as ~0.

Inspecting serialize_circuit I noticed a quick win would be to not do a full deep copy of the input circuit, instead create a shallow copy with circuit = circuit_inp.copy(). The risk with this is that we may modify the input circuit (change in behavior). We can get away with this because moments are replaces rather than edited in place by the serializer, i.e

circuit[i] = cirq.Moment(
            new_ops[op.qubits] if op.qubits in new_ops else op for op in moment)

That being said, in the controlled operations demotion loop, we reference and mutate op.sub_operation which points to an attribute on the input circuit. Therefore, we make a copy there of each sub operation.

for op in controlled_ops:
    # Copy before tagging: `sub_operation` is a live reference into
    # the caller's circuit, and these attributes would leak onto it.
    tfq_compatible = copy.copy(op.sub_operation)
    tfq_compatible._tfq_control_qubits = op.controls
    tfq_compatible._tfq_control_values = op.control_values
    new_ops[op.qubits] = tfq_compatible

Overall the speedup is significant:

==================================================================
serialize_circuit: before vs after
==================================================================
  baseline  (HEAD):  deepcopy
  candidate (working tree): shallow copy
  cirq 1.5.0 | python 3.10.14 | Darwin arm64
  best of 5 alternating rounds

  correctness
    output identical      yes   (2b2bc6dd05e14f75...)
    input circuits clean  yes

  100 x 10q depth 20                378.5 ->    277.1 ms    1.37x
  500 x 10q depth 20               1895.1 ->   1363.3 ms    1.39x
  100 x 20q depth 50               1888.7 ->   1402.4 ms    1.35x
  50 x 30q depth 100               2841.5 ->   2121.3 ms    1.34x
  100 x 20q depth 50 (+ctrl)       1928.5 ->   1436.3 ms    1.34x

The downside is less peace of mind that the input circuit is not mutated, and any changes here in the future need to keep that in mind.

I added a test to ensure serialization does not mutate the input circuit.

[mhucka]

Thanks for this improvement and for the careful before/after numbers!

Here's a preliminary review. I still need to do some more testing, and will come back to this when I can.

[mhucka]

@adlantz Thank you for the benchmarks in the PR description. Could you also attach the benchmark script?

[adlantz]

@mhucka Here are the two scripts I used. First is to profile convert_to_tensor and see where it spends its time.

Profile convert_to_tensor()

#!/usr/bin/env python3
"""Break tfq.convert_to_tensor into serialize / protobuf / packing phases.

Usage
-----
    ./.venv/bin/python profile_convert_to_tensor.py --serializer-ref master

Pass --serializer-ref to choose which serialize_circuit is profiled. Without
it, the script uses whatever version is in the installed wheel, which may not
match the code you think you are measuring.
"""

import argparse
import cProfile
import contextlib
import importlib.util
import io
import os
import pstats
import re
import subprocess
import sys
import tempfile
import time

# A checkout on sys.path[0] shadows the installed package; drop it.
_REPO = os.path.dirname(os.path.abspath(__file__))
sys.path[:] = [p for p in sys.path if os.path.abspath(p or ".") != _REPO]

import cirq
import numpy as np
import tensorflow_quantum as tfq
from tensorflow_quantum.core.serialize import serializer

_PATH = "tensorflow_quantum/core/serialize/serializer.py"

GATE_DOMAIN = {
    cirq.X: 1,
    cirq.Y: 1,
    cirq.Z: 1,
    cirq.H: 1,
    cirq.CNOT: 2,
    cirq.CZ: 2,
}

WORKLOADS = [
    # label, n_circuits, n_qubits, depth, controlled
    ("100 x 10q depth 20", 100, 10, 20, False),
    ("100 x 20q depth 50", 100, 20, 50, False),
    ("50 x 30q depth 100", 50, 30, 100, False),
    # Adds one controlled op per circuit, so control demotion (and the copy of
    # its sub_operation) is exercised at least once.
    ("100 x 20q depth 50 (+ctrl)", 100, 20, 50, True),
]


def serializer_variant():
    """Which copy strategy is serialize_circuit currently using?"""
    import inspect

    source = inspect.getsource(serializer.serialize_circuit)
    if "deepcopy(circuit_inp)" in source:
        return "deepcopy (baseline)"
    if "circuit_inp.copy()" in source:
        return "shallow copy (patched)"
    return "unrecognised"


def load_ref(ref, tmpdir):
    """Load serializer.py at a git ref as a standalone module.

    serialize_circuit is pure Python with absolute imports, so an older copy
    runs against the installed wheel without a rebuild -- everything else
    (protos, op_serializer, the compiled ops) still comes from site-packages.
    """
    try:
        source = subprocess.run(["git", "show", f"{ref}:{_PATH}"],
                                cwd=_REPO,
                                capture_output=True,
                                check=True).stdout
    except subprocess.CalledProcessError as exc:
        sys.exit(f"Could not read {_PATH} at {ref}:\n"
                 f"{exc.stderr.decode(errors='replace').strip()}")
    path = os.path.join(tmpdir, "_serializer_ref.py")
    with open(path, "wb") as handle:
        handle.write(source)

    spec = importlib.util.spec_from_file_location("_serializer_ref", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules["_serializer_ref"] = module
    spec.loader.exec_module(module)
    return module


def make_batch(n_circuits, n_qubits, depth, controlled=False, seed=11):
    """A reproducible batch of random circuits on GridQubits."""
    rand = np.random.RandomState(seed)
    qubits = cirq.GridQubit.rect(1, n_qubits)
    batch = []
    for _ in range(n_circuits):
        circuit = cirq.testing.random_circuit(qubits,
                                              depth,
                                              0.9,
                                              GATE_DOMAIN,
                                              random_state=rand)
        if controlled:
            circuit.append(
                cirq.ControlledOperation([qubits[0]],
                                         cirq.ISWAP(qubits[1], qubits[2])**0.5))
        batch.append(circuit)
    return batch


def best_of(fn, reps):
    """Minimum wall time over `reps` runs, in ms."""
    best = float("inf")
    for _ in range(reps):
        start = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - start)
    return best * 1000.0


def report(label, circuits, reps):
    """Time each phase separately and print it as a share of the total."""
    serialize_ms = best_of(
        lambda: [serializer.serialize_circuit(c) for c in circuits], reps)

    protos = [serializer.serialize_circuit(c) for c in circuits]
    proto_ms = best_of(lambda: [p.SerializeToString() for p in protos], reps)

    total_ms = best_of(lambda: tfq.convert_to_tensor(circuits), reps)
    remainder = total_ms - serialize_ms - proto_ms

    print(f"\n{label}")
    print("  " + "-" * 56)
    print(f"  {'convert_to_tensor (total)':<38} {total_ms:9.1f} ms   100.0%")
    print("  " + "-" * 56)
    for name, value in (("1. serialize_circuit  (cirq, Python)", serialize_ms),
                        ("2. SerializeToString  (protobuf)", proto_ms),
                        ("3. remainder          (tensor packing)", remainder)):
        print(f"  {name:<38} {value:9.1f} ms   "
              f"{100.0 * value / total_ms:5.1f}%")


def profile_hotspots(circuits, top, sort):
    """cProfile serialize_circuit and print the heaviest entries."""
    profiler = cProfile.Profile()
    profiler.enable()
    for circuit in circuits:
        serializer.serialize_circuit(circuit)
    profiler.disable()

    stream = io.StringIO()
    pstats.Stats(profiler, stream=stream).sort_stats(sort).print_stats(top)

    print(f"\ncProfile: serialize_circuit hotspots (by {sort})")
    print("-" * 60)
    for line in stream.getvalue().splitlines():
        # Strip directory components so the output stays readable, keeping the
        # ncalls/tottime/cumtime columns intact.
        print(re.sub(r"/\S*/", "", line))


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--reps",
                        type=int,
                        default=3,
                        help="timing repetitions per phase (default 3)")
    parser.add_argument("--profile",
                        action="store_true",
                        help="also print a cProfile breakdown")
    parser.add_argument("--profile-top",
                        type=int,
                        default=12,
                        help="cProfile rows to show (default 12)")
    parser.add_argument("--profile-sort",
                        default="cumulative",
                        choices=["cumulative", "tottime", "ncalls"],
                        help="cProfile sort key (default cumulative); tottime "
                        "surfaces leaves that cumulative buries")
    parser.add_argument("--serializer-ref",
                        default=None,
                        metavar="REF",
                        help="profile serialize_circuit as of this git ref "
                        "instead of the installed one, e.g. master for "
                        "pre-change numbers (no rebuild needed)")
    args = parser.parse_args()

    with contextlib.ExitStack() as stack:
        source = "installed wheel"
        if args.serializer_ref:
            tmpdir = stack.enter_context(tempfile.TemporaryDirectory())
            # util.py looks this attribute up at call time, so the swap
            # redirects convert_to_tensor too, not just direct calls.
            serializer.serialize_circuit = load_ref(args.serializer_ref,
                                                    tmpdir).serialize_circuit
            source = f"git {args.serializer_ref}"

        print("=" * 60)
        print("Where does convert_to_tensor spend its time?")
        print("=" * 60)
        print(f"  tensorflow_quantum {getattr(tfq, '__version__', 'unknown')} | "
              f"cirq {cirq.__version__} | python {sys.version.split()[0]}")
        print(f"  serialize_circuit: {serializer_variant()}  [{source}]")
        print(f"  best of {args.reps} runs per phase")

        for label, n_circuits, n_qubits, depth, controlled in WORKLOADS:
            report(label, make_batch(n_circuits, n_qubits, depth, controlled),
                   args.reps)

        print("\n  Phases are timed independently, so the remainder is a\n"
              "  difference of minima. Small negatives are noise, read as ~0.")

        if args.profile:
            profile_hotspots(make_batch(100, 20, 50), args.profile_top,
                             args.profile_sort)


if __name__ == "__main__":
    main()

And here is the script that benchmarks the before and after of replacing the deep copy.

Benchmark before and after changes

#!/usr/bin/env python3
"""Before/after benchmark for the `serialize_circuit` deepcopy change.

Loads two versions of `serializer.py` side by side in one process -- the
baseline from git, the candidate from your working tree -- then checks they
produce identical bytes and times them against each other.

Nothing is stashed, rebuilt, or reinstalled, and your working tree is never
modified. Both versions are loaded the same way, so neither gets an unfair
advantage from import caching.

Timing method
-------------
Rounds alternate baseline/candidate rather than running one block then the
other, so machine drift (thermal throttling, background load) hits both sides
roughly equally instead of landing entirely on whichever ran second. Each
side's reported figure is its minimum across rounds.

Usage
-----
    ./.venv/bin/python ab_benchmark_serializer.py

    --rounds N          alternating measurement rounds (default 5)
    --baseline-ref REF  git ref to take the baseline from (default master)
    --markdown          emit a markdown table for pasting into GitHub
    --quick             smaller workloads, for a fast sanity check
"""

import argparse
import hashlib
import importlib.util
import os
import platform
import subprocess
import sys
import tempfile
import time

_REPO = os.path.dirname(os.path.abspath(__file__))
_SERIALIZER_PATH = "tensorflow_quantum/core/serialize/serializer.py"

# This script lives in the source checkout, and Python puts a script's own
# directory first on sys.path. That shadows the installed tensorflow_quantum
# with the un-built source tree. Drop it so the real package wins.
if os.path.isdir(os.path.join(_REPO, "tensorflow_quantum", "core")):
    sys.path[:] = [p for p in sys.path if os.path.abspath(p or ".") != _REPO]

try:
    import cirq
    import numpy as np
except ImportError as exc:
    sys.exit(f"Could not import a dependency: {exc}")

GATE_DOMAIN = {
    cirq.X: 1,
    cirq.Y: 1,
    cirq.Z: 1,
    cirq.H: 1,
    cirq.CNOT: 2,
    cirq.CZ: 2,
}

WORKLOADS = [
    # label, n_circuits, n_qubits, depth, controlled
    ("100 x 10q depth 20", 100, 10, 20, False),
    ("500 x 10q depth 20", 500, 10, 20, False),
    ("100 x 20q depth 50", 100, 20, 50, False),
    ("50 x 30q depth 100", 50, 30, 100, False),
    ("100 x 20q depth 50 (+ctrl)", 100, 20, 50, True),
]

QUICK_WORKLOADS = [
    ("50 x 10q depth 20", 50, 10, 20, False),
    ("50 x 20q depth 50 (+ctrl)", 50, 20, 50, True),
]


# ---------------------------------------------------------------------------
# Loading two serializer versions at once
# ---------------------------------------------------------------------------


def git_show(ref, path):
    """Return a file's contents at a given git ref."""
    try:
        return subprocess.run(["git", "show", f"{ref}:{path}"],
                              cwd=_REPO,
                              capture_output=True,
                              check=True).stdout
    except subprocess.CalledProcessError as exc:
        sys.exit(f"Could not read {path} at {ref}:\n"
                 f"{exc.stderr.decode(errors='replace').strip()}")


def load_module(source_bytes, name, tmpdir):
    """Import `source_bytes` as a standalone module called `name`.

    `serializer.py` uses only absolute imports, so it can be loaded outside
    its package as long as `tensorflow_quantum` itself is importable.
    """
    path = os.path.join(tmpdir, f"{name}.py")
    with open(path, "wb") as handle:
        handle.write(source_bytes)

    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    try:
        spec.loader.exec_module(module)
    except ImportError as exc:
        sys.exit(
            f"Could not load {name}: {exc}\n\n"
            "tensorflow_quantum must be importable. Install a built wheel:\n"
            "  bazel build -c opt //release:build_pip_package\n"
            "  PYTHON_BIN_PATH=\"$PWD/.venv/bin/python\" \\\n"
            "      ./bazel-bin/release/build_pip_package /tmp/wheel\n"
            "  pip install /tmp/wheel/*.whl")
    return module


def describe(module):
    """Which copy strategy does this version use?"""
    import inspect

    source = inspect.getsource(module.serialize_circuit)
    if "deepcopy(circuit_inp)" in source:
        return "deepcopy"
    if "circuit_inp.copy()" in source:
        return "shallow copy"
    return "unrecognised"


# ---------------------------------------------------------------------------
# Workloads and measurement
# ---------------------------------------------------------------------------


def make_batch(n_circuits, n_qubits, depth, controlled, seed=11):
    rand = np.random.RandomState(seed)
    qubits = cirq.GridQubit.rect(1, n_qubits)
    batch = []
    for _ in range(n_circuits):
        circuit = cirq.testing.random_circuit(qubits,
                                              depth,
                                              0.9,
                                              GATE_DOMAIN,
                                              random_state=rand)
        if controlled:
            circuit.append(
                cirq.ControlledOperation([qubits[0]],
                                         cirq.ISWAP(qubits[1], qubits[2])**0.5))
        batch.append(circuit)
    return batch


def time_once(module, circuits):
    """Wall time in ms for one pass over the batch."""
    start = time.perf_counter()
    for circuit in circuits:
        module.serialize_circuit(circuit)
    return (time.perf_counter() - start) * 1000.0


def digest(module, circuits):
    """sha256 of the deterministic proto bytes for the whole batch."""
    blob = b"".join(
        module.serialize_circuit(c).SerializeToString(deterministic=True)
        for c in circuits)
    return hashlib.sha256(blob).hexdigest()


def tags_leaked(module, circuits):
    """Did serialization leave control tags on the caller's circuits?"""
    for circuit in circuits:
        for moment in circuit:
            for op in moment:
                sub_op = getattr(op, "sub_operation", op)
                if hasattr(sub_op, "_tfq_control_qubits"):
                    return True
    return False


def measure(baseline, candidate, circuits, rounds):
    """Alternate the two versions across rounds; return the minimum for each."""
    # Warm up both so neither pays first-call costs inside a measured round.
    time_once(baseline, circuits[:2])
    time_once(candidate, circuits[:2])

    base_times, cand_times = [], []
    for index in range(rounds):
        # Swap which side goes first each round, so any residual ordering
        # effect (cache state, GC timing) does not always favour the same one.
        if index % 2 == 0:
            base_times.append(time_once(baseline, circuits))
            cand_times.append(time_once(candidate, circuits))
        else:
            cand_times.append(time_once(candidate, circuits))
            base_times.append(time_once(baseline, circuits))

    return min(base_times), min(cand_times)


# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--rounds", type=int, default=5,
                        help="alternating measurement rounds (default 5)")
    parser.add_argument("--baseline-ref", default="master",
                        help="git ref for the baseline (default master)")
    parser.add_argument("--markdown", action="store_true",
                        help="emit a markdown table")
    parser.add_argument("--quick", action="store_true",
                        help="smaller workloads for a fast sanity check")
    args = parser.parse_args()

    working_path = os.path.join(_REPO, _SERIALIZER_PATH)
    if not os.path.exists(working_path):
        sys.exit(f"Not a tensorflow_quantum checkout: {working_path} missing.")

    with tempfile.TemporaryDirectory() as tmpdir:
        baseline = load_module(git_show(args.baseline_ref, _SERIALIZER_PATH),
                               "_serializer_baseline", tmpdir)
        with open(working_path, "rb") as handle:
            candidate = load_module(handle.read(), "_serializer_candidate",
                                    tmpdir)

        base_kind = describe(baseline)
        cand_kind = describe(candidate)

        if base_kind == cand_kind:
            sys.exit(
                f"Both versions use '{base_kind}' -- nothing to compare.\n"
                f"The baseline comes from {args.baseline_ref}; pass a ref "
                "that predates the change\nwith --baseline-ref.")

        workloads = QUICK_WORKLOADS if args.quick else WORKLOADS

        # --- correctness first: a speedup that changes output is worthless ---
        check = make_batch(40, 8, 15, controlled=True)
        base_digest = digest(baseline, check)
        cand_digest = digest(candidate, check)
        identical = base_digest == cand_digest

        leak_check = make_batch(20, 8, 15, controlled=True)
        digest(candidate, leak_check)
        leaked = tags_leaked(candidate, leak_check)

        if not args.markdown:
            print("=" * 66)
            print("serialize_circuit: before vs after")
            print("=" * 66)
            print(f"  baseline  ({args.baseline_ref}):  {base_kind}")
            print(f"  candidate (working tree): {cand_kind}")
            print(f"  cirq {cirq.__version__} | python "
                  f"{platform.python_version()} | {platform.system()} "
                  f"{platform.machine()}")
            print(f"  best of {args.rounds} alternating rounds")
            print()
            print("  correctness")
            print(f"    output identical      {'yes' if identical else 'NO'}"
                  f"   ({base_digest[:16]}...)")
            print(f"    input circuits clean  "
                  f"{'yes' if not leaked else 'NO -- TAGS LEAKED'}")
            print()
        else:
            print(f"<!-- baseline {args.baseline_ref} ({base_kind}) vs working "
                  f"tree ({cand_kind}) | cirq {cirq.__version__} | "
                  f"{platform.system()} {platform.machine()} -->\n")
            print(f"Output identical: **{'yes' if identical else 'NO'}** "
                  f"(`sha256={base_digest[:16]}...`)  ")
            print(f"Input circuits unmodified: "
                  f"**{'yes' if not leaked else 'NO'}**\n")
            print("| workload | before | after | speedup |")
            print("|---|---:|---:|---:|")

        speedups = []
        for label, n_circuits, n_qubits, depth, controlled in workloads:
            circuits = make_batch(n_circuits, n_qubits, depth, controlled)
            base_ms, cand_ms = measure(baseline, candidate, circuits,
                                       args.rounds)
            speedup = base_ms / cand_ms if cand_ms else float("nan")
            speedups.append(speedup)

            if args.markdown:
                print(f"| {label} | {base_ms:.0f} ms | {cand_ms:.0f} ms | "
                      f"{speedup:.2f}x |")
            else:
                print(f"  {label:<30} {base_ms:8.1f} -> {cand_ms:8.1f} ms   "
                      f"{speedup:5.2f}x")

        summary = (f"\nRange {min(speedups):.2f}x - {max(speedups):.2f}x "
                   f"across {len(speedups)} workloads.")
        print(summary)

        if not identical:
            print("\nWARNING: the two versions produced different bytes. "
                  "Investigate before trusting the timings.")
        if leaked:
            print("\nWARNING: the candidate left control tags on its input "
                  "circuits. This is the failure mode a bare shallow copy "
                  "causes -- see copy.copy(op.sub_operation).")


if __name__ == "__main__":
    main()

And here is a little guide on actually running these scripts:

Running the benchmark scripts

Both scripts are attached above. They aren't part of this PR's diff.

1. Check out this branch

git clone https://github.com/tensorflow/quantum.git && cd quantum
git fetch origin pull/1114/head:serialize-shallow-copy
git checkout serialize-shallow-copy

2. Install TFQ into a virtualenv. The scripts import the installed tensorflow_quantum for the compiled ops and protos, so a plain source checkout isn't enough. I built the wheel from this repo (see docs/install.md):

python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
./configure.sh
bazel build -c opt release:build_pip_package
./bazel-bin/release/build_pip_package /tmp/tfquantum/
pip install /tmp/tfquantum/*.whl

3. Save both scripts in the repo root. They find the repo from their own location and read serializer.py from git, so they have to sit inside the checkout.

4. Run them

# A/B: master's serialize_circuit vs. this branch, with an output-equality check
python ab_benchmark_serializer.py              # --markdown for a table, --quick for a fast check

# Where convert_to_tensor spends its time (pre-change code)
python profile_convert_to_tensor.py --serializer-ref master   # add --profile for a cProfile breakdown

Notes:

• Neither script changes the working tree. They load serializer.py from a git ref, or from the working tree, as a separate module.
• Always pass --serializer-ref to the profiler. Without it, the profiler measures whatever version is in your installed wheel.
• Expect some variation. Timings vary by a few percent from run to run. What should stay stable is the ratio between before and after.
• My numbers: cirq 1.5.0, Python 3.10, macOS arm64.