tn脚本更新

benchmarks/benchmark_quimb.py

@@ -0,0 +1,519 @@
+"""Benchmark and profile the qibotn/quimb CPU backend.
+
+This script is intended to be the stable baseline for quimb backend
+optimization work. It supports:
+
+- multiple circuit families
+- MPS or generic TN execution
+- statevector, sampling, conversion, and local expectation workloads
+- warmup/repeat timing
+- optional correctness checks against qibojit
+- optional cProfile output
+"""
+
+from __future__ import annotations
+
+import argparse
+import cProfile
+import json
+import math
+import os
+import pstats
+import time
+from pathlib import Path
+from statistics import mean, pstdev
+
+import numpy as np
+import qibo
+from qibo import Circuit, gates
+
+
+def configure_runtime_env(quimb_num_procs: int | None, blas_threads: int | None):
+    """Pin process-level thread settings before heavy work starts."""
+    if quimb_num_procs is not None:
+        os.environ["QUIMB_NUM_PROCS"] = str(quimb_num_procs)
+    if blas_threads is not None:
+        value = str(blas_threads)
+        os.environ["OMP_NUM_THREADS"] = value
+        os.environ["OPENBLAS_NUM_THREADS"] = value
+        os.environ["MKL_NUM_THREADS"] = value
+        os.environ["NUMEXPR_NUM_THREADS"] = value
+
+
+def make_circuit(
+    circuit_type: str,
+    nqubits: int,
+    nlayers: int,
+    seed: int,
+    add_measurements: bool = False,
+) -> Circuit:
+    """Construct repeatable workloads covering low/high entanglement cases."""
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    if circuit_type == "qft":
+        from qibo.models import QFT
+
+        circuit = QFT(nqubits)
+    elif circuit_type == "variational":
+        for layer in range(nlayers):
+            for qubit in range(nqubits):
+                circuit.add(gates.RY(qubit, theta=rng.uniform(0.0, 2.0 * np.pi)))
+            offset = layer % 2
+            for qubit in range(offset, nqubits - 1, 2):
+                circuit.add(gates.CZ(qubit, qubit + 1))
+    elif circuit_type == "ghz":
+        circuit.add(gates.H(0))
+        for qubit in range(nqubits - 1):
+            circuit.add(gates.CNOT(qubit, qubit + 1))
+    elif circuit_type == "qaoa":
+        for _ in range(nlayers):
+            for qubit in range(nqubits):
+                circuit.add(gates.RZ(qubit, theta=rng.uniform(0.0, 2.0 * np.pi)))
+            for qubit in range(0, nqubits - 1, 2):
+                circuit.add(gates.CZ(qubit, qubit + 1))
+            for qubit in range(nqubits):
+                circuit.add(gates.RX(qubit, theta=rng.uniform(0.0, 2.0 * np.pi)))
+    elif circuit_type == "ising1d":
+        for _ in range(nlayers):
+            for qubit in range(nqubits):
+                circuit.add(gates.RX(qubit, theta=rng.uniform(0.0, 2.0 * np.pi)))
+            for qubit in range(0, nqubits - 1, 2):
+                circuit.add(gates.CZ(qubit, qubit + 1))
+            for qubit in range(1, nqubits - 1, 2):
+                circuit.add(gates.CZ(qubit, qubit + 1))
+    elif circuit_type == "rcs":
+        cols = math.ceil(math.sqrt(nqubits))
+        rows = math.ceil(nqubits / cols)
+        single_qubit_gates = [gates.T, gates.X, gates.Y]
+        for layer in range(nlayers):
+            for qubit in range(nqubits):
+                gate_cls = single_qubit_gates[rng.integers(0, len(single_qubit_gates))]
+                circuit.add(gate_cls(qubit))
+            if layer % 2 == 0:
+                for row in range(rows):
+                    for col in range(0, cols - 1, 2):
+                        q1, q2 = row * cols + col, row * cols + col + 1
+                        if q2 < nqubits:
+                            circuit.add(gates.CZ(q1, q2))
+            else:
+                for row in range(0, rows - 1, 2):
+                    for col in range(cols):
+                        q1, q2 = row * cols + col, (row + 1) * cols + col
+                        if q2 < nqubits:
+                            circuit.add(gates.CZ(q1, q2))
+    else:
+        raise ValueError(f"Unknown circuit type: {circuit_type}")
+
+    if add_measurements:
+        circuit.add(gates.M(*range(nqubits)))
+    return circuit
+
+
+def prepare_quimb_backend(
+    ansatz: str,
+    max_bond: int | None,
+    svd_cutoff: float,
+    optimizer: str,
+    n_most_frequent_states: int,
+):
+    """Create and configure the qibotn/quimb backend once."""
+    qibo.set_backend("qibotn", platform="quimb")
+    backend = qibo.get_backend()
+    backend.configure_tn_simulation(
+        ansatz=ansatz,
+        max_bond_dimension=max_bond,
+        svd_cutoff=svd_cutoff,
+        n_most_frequent_states=n_most_frequent_states,
+    )
+    backend.contractions_optimizer = optimizer
+    return backend
+
+
+def run_qibojit_state(circuit: Circuit):
+    qibo.set_backend("qibojit", platform="numba")
+    t0 = time.perf_counter()
+    result = circuit()
+    elapsed = time.perf_counter() - t0
+    state = np.asarray(result.state(), dtype=complex).reshape(-1)
+    return state, elapsed
+
+
+def run_qibojit_shots(circuit: Circuit, nshots: int):
+    qibo.set_backend("qibojit", platform="numba")
+    t0 = time.perf_counter()
+    result = circuit(nshots=nshots)
+    elapsed = time.perf_counter() - t0
+    return dict(result.frequencies()), elapsed
+
+
+def z_expectation_from_statevector(statevector: np.ndarray, nqubits: int, qubit: int):
+    probs = np.abs(np.asarray(statevector).reshape(-1)) ** 2
+    bit_index = nqubits - qubit - 1
+    bits = (np.arange(len(probs)) >> bit_index) & 1
+    return float(np.dot(probs, 1.0 - 2.0 * bits))
+
+
+def fidelity_and_l2(reference: np.ndarray, candidate: np.ndarray):
+    ref = np.asarray(reference, dtype=complex).reshape(-1)
+    cand = np.asarray(candidate, dtype=complex).reshape(-1)
+    fidelity = abs(np.vdot(ref, cand)) ** 2
+    l2_error = np.linalg.norm(ref - cand)
+    return float(fidelity), float(l2_error)
+
+
+def total_variation_distance(reference: dict[str, int], candidate: dict[str, int], nshots: int):
+    keys = set(reference) | set(candidate)
+    return 0.5 * sum(abs(reference.get(key, 0) - candidate.get(key, 0)) for key in keys) / nshots
+
+
+def profile_callable(func, output_path: Path, sort_by: str):
+    """Profile a single invocation and dump textual stats."""
+    profiler = cProfile.Profile()
+    profiler.enable()
+    result = func()
+    profiler.disable()
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", encoding="utf-8") as stream:
+        stats = pstats.Stats(profiler, stream=stream)
+        stats.strip_dirs().sort_stats(sort_by).print_stats(80)
+        stats.print_callers(30)
+    return result
+
+
+def time_callable(func, repeats: int, warmup: int, profile_output: Path | None, profile_sort: str):
+    for _ in range(warmup):
+        func()
+
+    profiled_payload = None
+    if profile_output is not None:
+        profiled_payload = profile_callable(func, profile_output, profile_sort)
+
+    samples = []
+    payloads = []
+
+    for _ in range(repeats):
+        t0 = time.perf_counter()
+        payload = func()
+        elapsed = time.perf_counter() - t0
+        samples.append(elapsed)
+        payloads.append(payload)
+
+    final_payload = payloads[-1] if payloads else profiled_payload
+    return samples, final_payload
+
+
+def summarize_samples(samples: list[float]):
+    return {
+        "min_s": min(samples),
+        "mean_s": mean(samples),
+        "max_s": max(samples),
+        "std_s": pstdev(samples) if len(samples) > 1 else 0.0,
+        "repeats": len(samples),
+    }
+
+
+def workload_state(args):
+    circuit = make_circuit(args.circuit, args.nqubits, args.nlayers, args.seed)
+    backend = prepare_quimb_backend(
+        ansatz=args.ansatz,
+        max_bond=args.max_bond,
+        svd_cutoff=args.svd_cutoff,
+        optimizer=args.optimizer,
+        n_most_frequent_states=args.topk,
+    )
+
+    def run_once():
+        result = backend.execute_circuit(circuit, return_array=True)
+        return np.asarray(result.statevector).reshape(-1)
+
+    samples, statevector = time_callable(
+        run_once, args.repeats, args.warmup, args.profile_output, args.profile_sort
+    )
+    summary = summarize_samples(samples)
+
+    correctness = None
+    if not args.no_compare:
+        ref_state, ref_time = run_qibojit_state(circuit)
+        fidelity, l2_error = fidelity_and_l2(ref_state, statevector)
+        correctness = {
+            "qibojit_time_s": ref_time,
+            "fidelity": fidelity,
+            "l2_error": l2_error,
+        }
+
+    return summary, correctness
+
+
+def workload_shots(args):
+    circuit = make_circuit(
+        args.circuit, args.nqubits, args.nlayers, args.seed, add_measurements=True
+    )
+    backend = prepare_quimb_backend(
+        ansatz=args.ansatz,
+        max_bond=args.max_bond,
+        svd_cutoff=args.svd_cutoff,
+        optimizer=args.optimizer,
+        n_most_frequent_states=args.topk,
+    )
+
+    def run_once():
+        result = backend.execute_circuit(circuit, nshots=args.nshots)
+        return dict(result.frequencies())
+
+    samples, frequencies = time_callable(
+        run_once, args.repeats, args.warmup, args.profile_output, args.profile_sort
+    )
+    summary = summarize_samples(samples)
+
+    correctness = None
+    if not args.no_compare:
+        ref_freq, ref_time = run_qibojit_shots(circuit, args.nshots)
+        correctness = {
+            "qibojit_time_s": ref_time,
+            "tvd": total_variation_distance(ref_freq, frequencies, args.nshots),
+        }
+
+    return summary, correctness
+
+
+def workload_convert(args):
+    circuit = make_circuit(args.circuit, args.nqubits, args.nlayers, args.seed)
+    backend = prepare_quimb_backend(
+        ansatz=args.ansatz,
+        max_bond=args.max_bond,
+        svd_cutoff=args.svd_cutoff,
+        optimizer=args.optimizer,
+        n_most_frequent_states=args.topk,
+    )
+
+    def run_once():
+        quimb_circuit = backend._qibo_circuit_to_quimb(  # pylint: disable=protected-access
+            circuit,
+            quimb_circuit_type=backend.circuit_ansatz,
+            gate_opts={"max_bond": backend.max_bond_dimension, "cutoff": backend.svd_cutoff},
+        )
+        return len(quimb_circuit.gates)
+
+    samples, gate_count = time_callable(
+        run_once, args.repeats, args.warmup, args.profile_output, args.profile_sort
+    )
+    summary = summarize_samples(samples)
+    summary["gate_count"] = gate_count
+    return summary, None
+
+
+def workload_expectation(args):
+    circuit = make_circuit(args.circuit, args.nqubits, args.nlayers, args.seed)
+    backend = prepare_quimb_backend(
+        ansatz=args.ansatz,
+        max_bond=args.max_bond,
+        svd_cutoff=args.svd_cutoff,
+        optimizer=args.optimizer,
+        n_most_frequent_states=args.topk,
+    )
+    operators = ["z"]
+    sites = [(args.observable_qubit,)]
+    coeffs = [1.0]
+
+    def run_once():
+        return float(
+            backend.exp_value_observable_symbolic(
+                circuit, operators, sites, coeffs, args.nqubits
+            )
+        )
+
+    samples, expval = time_callable(
+        run_once, args.repeats, args.warmup, args.profile_output, args.profile_sort
+    )
+    summary = summarize_samples(samples)
+
+    correctness = None
+    if not args.no_compare:
+        ref_state, ref_time = run_qibojit_state(circuit)
+        correctness = {
+            "qibojit_time_s": ref_time,
+            "reference_expval": z_expectation_from_statevector(
+                ref_state, args.nqubits, args.observable_qubit
+            ),
+            "abs_error": abs(
+                z_expectation_from_statevector(ref_state, args.nqubits, args.observable_qubit)
+                - expval
+            ),
+        }
+
+    return summary, correctness
+
+
+def workload_raw_local_exp(args):
+    circuit = make_circuit(args.circuit, args.nqubits, args.nlayers, args.seed)
+    backend = prepare_quimb_backend(
+        ansatz=args.ansatz,
+        max_bond=args.max_bond,
+        svd_cutoff=args.svd_cutoff,
+        optimizer=args.optimizer,
+        n_most_frequent_states=args.topk,
+    )
+
+    def run_once():
+        metrics = {}
+        t0 = time.perf_counter()
+        quimb_circuit = backend._qibo_circuit_to_quimb(  # pylint: disable=protected-access
+            circuit,
+            quimb_circuit_type=backend.circuit_ansatz,
+            gate_opts={"max_bond": backend.max_bond_dimension, "cutoff": backend.svd_cutoff},
+        )
+        metrics["convert_s"] = time.perf_counter() - t0
+
+        operator = backend._string_to_quimb_operator("z")  # pylint: disable=protected-access
+        if args.rehearse:
+            t1 = time.perf_counter()
+            rehearsal = quimb_circuit.local_expectation(
+                operator,
+                where=(args.observable_qubit,),
+                backend=backend.backend,
+                optimize=backend.contractions_optimizer,
+                simplify_sequence="R",
+                rehearse=True,
+            )
+            metrics["rehearse_s"] = time.perf_counter() - t1
+            optimize = rehearsal["tree"]
+        else:
+            metrics["rehearse_s"] = 0.0
+            optimize = backend.contractions_optimizer
+
+        t2 = time.perf_counter()
+        expval = quimb_circuit.local_expectation(
+            operator,
+            where=(args.observable_qubit,),
+            backend=backend.backend,
+            optimize=optimize,
+            simplify_sequence="R",
+        )
+        metrics["contract_s"] = time.perf_counter() - t2
+        metrics["total_inner_s"] = (
+            metrics["convert_s"] + metrics["rehearse_s"] + metrics["contract_s"]
+        )
+        metrics["expval"] = float(np.real(expval))
+        return metrics
+
+    samples, metrics = time_callable(
+        run_once, args.repeats, args.warmup, args.profile_output, args.profile_sort
+    )
+    summary = summarize_samples(samples)
+    summary.update(
+        {
+            "convert_s": metrics["convert_s"],
+            "rehearse_s": metrics["rehearse_s"],
+            "contract_s": metrics["contract_s"],
+            "total_inner_s": metrics["total_inner_s"],
+        }
+    )
+
+    correctness = None
+    if not args.no_compare:
+        ref_state, ref_time = run_qibojit_state(circuit)
+        ref_expval = z_expectation_from_statevector(
+            ref_state, args.nqubits, args.observable_qubit
+        )
+        correctness = {
+            "qibojit_time_s": ref_time,
+            "reference_expval": ref_expval,
+            "abs_error": abs(ref_expval - metrics["expval"]),
+        }
+
+    return summary, correctness
+
+
+WORKLOADS = {
+    "state": workload_state,
+    "shots": workload_shots,
+    "convert": workload_convert,
+    "expectation": workload_expectation,
+    "raw-local-exp": workload_raw_local_exp,
+}
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--mode",
+        choices=sorted(WORKLOADS),
+        default="raw-local-exp",
+        help="Workload to benchmark.",
+    )
+    parser.add_argument(
+        "--circuit",
+        choices=["ghz", "ising1d", "qaoa", "qft", "rcs", "variational"],
+        default="variational",
+    )
+    parser.add_argument("--nqubits", type=int, default=10)
+    parser.add_argument("--nlayers", type=int, default=3)
+    parser.add_argument("--ansatz", choices=["mps", "tn"], default="tn")
+    parser.add_argument("--max-bond", type=int, default=None)
+    parser.add_argument("--svd-cutoff", type=float, default=1e-10)
+    parser.add_argument("--optimizer", type=str, default="auto-hq")
+    parser.add_argument("--observable-qubit", type=int, default=0)
+    parser.add_argument("--nshots", type=int, default=1024)
+    parser.add_argument("--topk", type=int, default=100)
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--repeats", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--quimb-num-procs", type=int, default=None)
+    parser.add_argument("--blas-threads", type=int, default=None)
+    parser.add_argument("--rehearse", action="store_true")
+    parser.add_argument("--no-compare", action="store_true")
+    parser.add_argument("--profile-output", type=Path, default=None)
+    parser.add_argument("--profile-sort", type=str, default="cumulative")
+    parser.add_argument("--json-output", type=Path, default=None)
+    return parser
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+
+    configure_runtime_env(args.quimb_num_procs, args.blas_threads)
+
+    print(
+        f"mode={args.mode} circuit={args.circuit} nqubits={args.nqubits} "
+        f"nlayers={args.nlayers} ansatz={args.ansatz} optimizer={args.optimizer}"
+    )
+    if args.quimb_num_procs is not None or args.blas_threads is not None:
+        print(
+            "threads:"
+            f" QUIMB_NUM_PROCS={os.environ.get('QUIMB_NUM_PROCS')}"
+            f" OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS')}"
+        )
+
+    workload = WORKLOADS[args.mode]
+    summary, correctness = workload(args)
+
+    print("\nTiming")
+    for key, value in summary.items():
+        if isinstance(value, float):
+            print(f"{key:>16}: {value:.6f}")
+        else:
+            print(f"{key:>16}: {value}")
+
+    if correctness is not None:
+        print("\nCorrectness")
+        for key, value in correctness.items():
+            if isinstance(value, float):
+                print(f"{key:>16}: {value:.6e}")
+            else:
+                print(f"{key:>16}: {value}")
+
+    if args.profile_output is not None:
+        print(f"\nProfile written to: {args.profile_output}")
+
+    if args.json_output is not None:
+        payload = {"timing": summary, "correctness": correctness, "args": vars(args)}
+        args.json_output.parent.mkdir(parents=True, exist_ok=True)
+        args.json_output.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
+        print(f"JSON written to: {args.json_output}")
+
+
+if __name__ == "__main__":
+    main()