tn脚本更新

bench_profile.py

@@ -0,0 +1,460 @@
+"""Benchmark: qibotn/quimb generic TN — single-process torch profiling version."""
+
+import os
+import pickle
+import time
+import argparse
+import numpy as np
+import cotengra as ctg
+import qibo
+from qibo import Circuit, gates
+
+
+def make_circuit(circuit_type, nqubits, nlayers=1):
+    c = Circuit(nqubits)
+
+    if circuit_type == "qft":
+        from qibo.models import QFT
+        return QFT(nqubits)
+
+    elif circuit_type == "variational":
+        for layer in range(nlayers):
+            for q in range(nqubits):
+                c.add(gates.RY(q, theta=np.random.uniform(0, 2 * np.pi)))
+
+            offset = layer % 2
+            for q in range(offset, nqubits - 1, 2):
+                c.add(gates.CZ(q, q + 1))
+
+    elif circuit_type == "ghz":
+        c.add(gates.H(0))
+        for q in range(nqubits - 1):
+            c.add(gates.CNOT(q, q + 1))
+
+    elif circuit_type == "brickwork":
+        for q in range(nqubits):
+            c.add(gates.H(q))
+
+        for layer in range(nlayers):
+            offset = layer % 2
+            for q in range(offset, nqubits - 1, 2):
+                c.add(gates.CNOT(q, q + 1))
+                c.add(gates.RZ(q, theta=np.random.uniform(0, 2 * np.pi)))
+                c.add(gates.RZ(q + 1, theta=np.random.uniform(0, 2 * np.pi)))
+
+    else:
+        raise ValueError(f"Unknown circuit: {circuit_type}")
+
+    return c
+
+
+def make_z_observable(nqubits):
+    """Z on qubit 0 only — single contraction for benchmarking."""
+    return ["z"], [(0,)], [1.0]
+
+
+def export_profiler_outputs(prof, trace_path):
+    """Export Chrome trace and text table."""
+    prof.export_chrome_trace(trace_path)
+
+    table_path = trace_path.replace(".json", ".txt")
+    with open(table_path, "w") as f:
+        f.write(
+            prof.key_averages().table(
+                sort_by="self_cpu_time_total",
+                row_limit=200,
+            )
+        )
+
+    print(f"  [torch profiler trace] {trace_path}")
+    print(f"  [torch profiler table] {table_path}")
+
+
+def run_quimb_tn(
+    circuit,
+    nqubits,
+    num_slices,
+    load_path=None,
+    save_path=None,
+):
+    """Mode: expval — compute <Z_0> via local_expectation."""
+    qibo.set_backend("qibotn", platform="quimb")
+    b = qibo.get_backend()
+    b.configure_tn_simulation(ansatz="tn")
+
+    operators, sites, coeffs = make_z_observable(nqubits)
+    ops = b._string_to_quimb_operator(operators[0])
+
+    qc = b._qibo_circuit_to_quimb(
+        circuit,
+        quimb_circuit_type=b.circuit_ansatz,
+        gate_opts={"max_bond": None, "cutoff": 1e-10},
+    )
+
+    if load_path:
+        with open(load_path, "rb") as f:
+            saved = pickle.load(f)
+        tree = saved["tree"]
+        t_search = 0.0
+        print(f"  [path loaded]  {load_path}")
+    else:
+        opt = ctg.HyperOptimizer(
+            methods=["kahypar", "random-greedy", "spinglass"],
+            max_repeats=16,
+            parallel=True,
+            max_time=60,
+            slicing_opts={"target_slices": num_slices},
+            progbar=True,
+        )
+
+        t0 = time.time()
+        rehearsal = qc.local_expectation(
+            ops,
+            where=sites[0],
+            optimize=opt,
+            simplify_sequence="R",
+            rehearse=True,
+        )
+        t_search = time.time() - t0
+        tree = rehearsal["tree"]
+
+        print(
+            f"  [path search]  {t_search:.3f}s  "
+            f"flops~2^{tree.contraction_cost():.2f}  "
+            f"size~2^{tree.contraction_width():.2f}  "
+            f"slices={tree.multiplicity}"
+        )
+
+        if save_path:
+            with open(save_path, "wb") as f:
+                pickle.dump({"tree": tree}, f)
+            print(f"  [path saved]  {save_path}")
+
+    t0 = time.time()
+    expval = qc.local_expectation(
+        ops,
+        where=sites[0],
+        optimize=tree,
+        simplify_sequence="R",
+    )
+    t_contract = time.time() - t0
+
+    print(f"  [contraction]  {t_contract:.3f}s")
+
+    return float(expval.real), t_search + t_contract
+
+
+def run_quimb_tn_statevector(
+    circuit,
+    nqubits,
+    num_slices,
+    load_path=None,
+    save_path=None,
+    profile=False,
+    profile_dir="profiles",
+):
+    """Mode: statevector — contract full TN to dense vector, single process."""
+    qibo.set_backend("qibotn", platform="quimb")
+    b = qibo.get_backend()
+    b.configure_tn_simulation(ansatz="tn")
+
+    import torch
+
+    qc = b._qibo_circuit_to_quimb(
+        circuit,
+        quimb_circuit_type=b.circuit_ansatz,
+        gate_opts={"max_bond": None, "cutoff": 1e-10},
+    )
+
+    # 让 quimb 生成 torch tensor，这样 torch.profiler 能抓到 aten op。
+    qc.to_backend = torch.from_numpy
+
+    if load_path:
+        with open(load_path, "rb") as f:
+            saved = pickle.load(f)
+        tree = saved["tree"]
+        t_search = 0.0
+        print(f"  [path loaded]  {load_path}")
+    else:
+        opt = ctg.HyperOptimizer(
+            methods=["kahypar", "random-greedy", "spinglass"],
+            max_repeats=500,
+            parallel=48,
+            max_time=100,
+            minimize="size",
+            slicing_opts={"target_slices": num_slices},
+            progbar=True,
+        )
+
+        t0 = time.time()
+        rehearsal = qc.to_dense(optimize=opt, rehearse=True)
+        t_search = time.time() - t0
+        tree = rehearsal["tree"]
+
+        print(
+            f"  [path search]  {t_search:.3f}s  "
+            f"flops~2^{tree.contraction_cost():.2f}  "
+            f"size~2^{tree.contraction_width():.2f}  "
+            f"slices={tree.multiplicity}"
+        )
+
+        if save_path:
+            with open(save_path, "wb") as f:
+                pickle.dump({"tree": tree}, f)
+            print(f"  [path saved]  {save_path}")
+
+    os.makedirs(profile_dir, exist_ok=True)
+
+    if profile:
+        from torch.profiler import profile as torch_profile
+        from torch.profiler import ProfilerActivity, record_function
+
+        trace_path = os.path.join(
+            profile_dir,
+            (
+                f"trace_statevector_"
+                f"{circuit.nqubits}q_"
+                f"slices{tree.multiplicity}_"
+                f"{int(time.time())}.json"
+            ),
+        )
+
+        t0 = time.time()
+
+        with torch_profile(
+            activities=[ProfilerActivity.CPU],
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=True,
+        ) as prof:
+            with record_function("qibotn_to_dense_contraction"):
+                sv = qc.to_dense(optimize=tree).reshape(-1)
+
+            with record_function("torch_to_numpy_view_or_copy"):
+                if type(sv).__module__.startswith("torch"):
+                    sv_tn = sv.detach().cpu().numpy()
+                else:
+                    sv_tn = np.asarray(sv)
+
+        t_contract = time.time() - t0
+
+        export_profiler_outputs(prof, trace_path)
+
+    else:
+        t0 = time.time()
+        sv = qc.to_dense(optimize=tree).reshape(-1)
+        t_contract = time.time() - t0
+
+        if type(sv).__module__.startswith("torch"):
+            sv_tn = sv.detach().cpu().numpy()
+        else:
+            sv_tn = np.asarray(sv)
+
+    print(f"  [contraction]  {t_contract:.3f}s")
+
+    return sv_tn, t_search + t_contract
+
+
+def run_quimb_tn_samples(circuit, nshots=1024):
+    """Mode: samples — sample from circuit output distribution."""
+    qibo.set_backend("qibotn", platform="quimb")
+    b = qibo.get_backend()
+    b.configure_tn_simulation(ansatz="tn")
+
+    t0 = time.time()
+    result = b.execute_circuit(circuit, nshots=nshots)
+    t_total = time.time() - t0
+
+    print(f"  [sampling]     {t_total:.3f}s  nshots={nshots}")
+
+    try:
+        freqs = result.frequencies()
+        print(f"  top states: {dict(list(freqs.items())[:5])}")
+    except Exception:
+        pass
+
+    return result, t_total
+
+
+def qibojit_expval(circuit, nqubits):
+    """Compute <Z_0> via qibojit statevector."""
+    qibo.set_backend("qibojit", platform="numba")
+
+    t0 = time.time()
+    result = circuit()
+    elapsed = time.time() - t0
+
+    sv = np.array(result.state(), dtype=complex).flatten()
+    probs = np.abs(sv) ** 2
+
+    bits = (np.arange(len(probs)) >> (nqubits - 1)) & 1
+    expval = float(np.dot(probs, 1 - 2 * bits))
+
+    return expval, elapsed
+
+
+def run_qibojit(circuit):
+    """Compute full statevector via qibojit."""
+    qibo.set_backend("qibojit", platform="numba")
+
+    t0 = time.time()
+    result = circuit()
+    elapsed = time.time() - t0
+
+    sv = np.array(result.state(), dtype=complex).flatten()
+
+    return sv, elapsed
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--nqubits", type=int, default=10)
+    parser.add_argument(
+        "--circuit",
+        type=str,
+        default="qft",
+        choices=["qft", "variational", "ghz", "brickwork"],
+    )
+    parser.add_argument("--nlayers", type=int, default=3)
+    parser.add_argument("--num-slices", type=int, default=1)
+    parser.add_argument("--nshots", type=int, default=1024)
+
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="statevector",
+        choices=["expval", "statevector", "samples"],
+        help="expval: local_expectation; statevector: to_dense; samples: sampling",
+    )
+
+    parser.add_argument(
+        "--no-compare",
+        action="store_true",
+        help="Skip qibojit reference run",
+    )
+
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default=None,
+        help="Save contraction tree to a pickle file",
+    )
+
+    parser.add_argument(
+        "--load-path",
+        type=str,
+        default=None,
+        help="Load contraction tree from a pickle file and skip path search",
+    )
+
+    parser.add_argument(
+        "--profile",
+        action="store_true",
+        help="Enable torch profiler for statevector contraction stage",
+    )
+
+    parser.add_argument(
+        "--profile-dir",
+        type=str,
+        default="profiles",
+        help="Directory to save torch profiler traces",
+    )
+
+    parser.add_argument(
+        "--save-statevector",
+        action="store_true",
+        help="Save TN statevector to data/sv_tn_*.npy",
+    )
+
+    args = parser.parse_args()
+
+    print(
+        f"Circuit: {args.circuit}, "
+        f"nqubits={args.nqubits}, "
+        f"nlayers={args.nlayers}, "
+        f"mode={args.mode}, "
+        f"profile={args.profile}"
+    )
+
+    np.random.seed(42)
+    circuit_tn = make_circuit(args.circuit, args.nqubits, args.nlayers)
+
+    try:
+        if args.mode == "expval":
+            expval_tn, t_tn = run_quimb_tn(
+                circuit_tn,
+                args.nqubits,
+                args.num_slices,
+                load_path=args.load_path,
+                save_path=args.save_path,
+            )
+
+            print(f"\n[quimb TN]  time={t_tn:.4f}s  <Z_0>={expval_tn:.8f}")
+
+        elif args.mode == "statevector":
+            sv_tn, t_tn = run_quimb_tn_statevector(
+                circuit_tn,
+                args.nqubits,
+                args.num_slices,
+                load_path=args.load_path,
+                save_path=args.save_path,
+                profile=args.profile,
+                profile_dir=args.profile_dir,
+            )
+
+            print(
+                f"\n[quimb TN]  time={t_tn:.4f}s  "
+                f"statevector shape={sv_tn.shape}"
+            )
+
+            if args.save_statevector:
+                os.makedirs("data", exist_ok=True)
+                out_path = f"data/sv_tn_{args.circuit}{args.nqubits}.npy"
+                np.save(out_path, sv_tn)
+                print(f"[saved statevector] {out_path}")
+
+        else:
+            _, t_tn = run_quimb_tn_samples(
+                circuit_tn,
+                nshots=args.nshots,
+            )
+
+            print(f"\n[quimb TN]  time={t_tn:.4f}s")
+            args.no_compare = True
+
+    except Exception as e:
+        print(f"[quimb TN] FAILED: {e}")
+        raise
+
+    if not args.no_compare and args.mode != "statevector":
+        np.random.seed(42)
+        circuit_ref = make_circuit(args.circuit, args.nqubits, args.nlayers)
+
+        expval_ref, t_ref = qibojit_expval(circuit_ref, args.nqubits)
+
+        print(f"[qibojit]   time={t_ref:.4f}s  <Z_0>={expval_ref:.8f}")
+        print(f"\nDiff     : {abs(expval_tn - expval_ref):.2e}")
+
+        if t_tn > 0:
+            print(f"Speedup  : {t_ref / t_tn:.2f}x")
+
+    elif not args.no_compare and args.mode == "statevector" and sv_tn is not None:
+        np.random.seed(42)
+        circuit_ref = make_circuit(args.circuit, args.nqubits, args.nlayers)
+
+        sv_ref, t_ref = run_qibojit(circuit_ref)
+
+        fid = abs(np.dot(sv_ref.conj(), sv_tn)) ** 2
+        l2_err = np.linalg.norm(sv_ref - sv_tn)
+
+        print(f"[qibojit]   time={t_ref:.4f}s")
+        print(f"Fidelity : {fid:.8f}  (1=perfect)")
+        print(f"L2 error : {l2_err:.2e}")
+
+        if t_tn > 0:
+            print(f"Speedup  : {t_ref / t_tn:.2f}x")
+
+
+if __name__ == "__main__":
+    main()