Benchmark compilation time for cancel_inverses and merge_rotations, catalyst-vs-PL

frontend/catalyst/compiler.py

@@ -189,6 +189,14 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
     ],
 )
 
+PEEPHOLE_BENCHMARK_PASS = (
+    "PeepholeBenchmarkPass",
+    [
+        "remove-chained-self-inverse{func-name=circuit}",
+        "merge-rotations{func-name=circuit}",
+    ],
+)
+
 HLO_LOWERING_PASS = (
     "HLOLoweringPass",
     [
@@ -306,6 +314,7 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
 
 
 DEFAULT_PIPELINES = [
+    PEEPHOLE_BENCHMARK_PASS,
     ENFORCE_RUNTIME_INVARIANTS_PASS,
     HLO_LOWERING_PASS,
     QUANTUM_COMPILATION_PASS,
@@ -323,6 +332,7 @@ def run_writing_command(command: List[str], compile_options: Optional[CompileOpt
 
 DEFAULT_ASYNC_PIPELINES = [
     ENFORCE_RUNTIME_INVARIANTS_PASS,
+    PEEPHOLE_BENCHMARK_PASS,
     HLO_LOWERING_PASS,
     QUANTUM_COMPILATION_PASS,
     BUFFERIZATION_PASS,

mlir/lib/Quantum/Transforms/merge_rotation.cpp

@@ -14,6 +14,8 @@
 
 #define DEBUG_TYPE "merge-rotation"
 
+#include <chrono>
+
 #include "Catalyst/IR/CatalystDialect.h"
 #include "Quantum/IR/QuantumOps.h"
 #include "Quantum/Transforms/Patterns.h"
@@ -40,6 +42,7 @@ struct MergeRotationsPass : impl::MergeRotationsPassBase<MergeRotationsPass> {
     {
         LLVM_DEBUG(dbgs() << "merge rotation pass"
                           << "\n");
+        auto start = std::chrono::high_resolution_clock::now();
 
         Operation *module = getOperation();
         Operation *targetfunc;
@@ -67,6 +70,10 @@ struct MergeRotationsPass : impl::MergeRotationsPassBase<MergeRotationsPass> {
         if (failed(applyPatternsAndFoldGreedily(targetfunc, std::move(patterns)))) {
             return signalPassFailure();
         }
+
+        auto stop = std::chrono::high_resolution_clock::now();
+        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
+        llvm::errs() << "merge rotation pass runtime: " << duration.count() << " microseconds\n";
     }
 };
 

mlir/lib/Quantum/Transforms/remove_chained_self_inverse.cpp

@@ -14,6 +14,7 @@
 
 #define DEBUG_TYPE "remove-chained-self-inverse"
 
+#include <chrono>
 #include <memory>
 #include <vector>
 
@@ -51,6 +52,7 @@ struct RemoveChainedSelfInversePass
     {
         LLVM_DEBUG(dbgs() << "remove chained self inverse pass"
                           << "\n");
+        auto start = std::chrono::high_resolution_clock::now();
 
         // Run cse pass before running remove-chained-self-inverse,
         // to aid identifying equivalent SSA values when verifying
@@ -88,6 +90,10 @@ struct RemoveChainedSelfInversePass
         if (failed(applyPatternsAndFoldGreedily(targetfunc, std::move(patterns)))) {
             return signalPassFailure();
         }
+
+        auto stop = std::chrono::high_resolution_clock::now();
+        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(stop - start);
+        llvm::errs() << "cancel inverse pass runtime: " << duration.count() << " microseconds\n";
     }
 };
 

my_toy_circuit.py

@@ -0,0 +1,112 @@
+import os
+import sys
+
+import jax
+import numpy as np
+import pennylane as qml
+
+import catalyst
+from catalyst import qjit
+from catalyst.debug import instrumentation
+
+from timeit import default_timer as timer
+
+dev = qml.device("lightning.qubit", wires=2)
+
+peephole_pipeline = {"cancel_inverses": {}, "merge_rotations": {}}
+
+
+# @qjit(seed=37, keep_intermediate=True)
+# @qjit(keep_intermediate=True)
+"""
+@qjit(pipelines = [
+    ("0_canonicalize", ["canonicalize"]),
+    ("peephole",
+    [
+    "builtin.module(remove-chained-self-inverse{func-name=circuit})",
+    "builtin.module(merge-rotations{func-name=circuit})"
+    ]
+    ), # peephole
+    ("inline_nested_module", ["inline-nested-module"]),
+    ], # pipelines =
+    #circuit_transform_pipeline = peephole_pipeline,
+    autograph=True,
+    keep_intermediate=True)
+"""
+
+
+@qjit(
+    autograph=True,
+    keep_intermediate=False,
+    # circuit_transform_pipeline = peephole_pipeline,
+)
+@qml.qnode(dev)
+def circuit(theta, loop_size):
+    for i in range(loop_size):
+        # for j in range(loop_size):
+        qml.Hadamard(0)
+        qml.Hadamard(0)
+        qml.RX(theta, wires=1)
+        qml.RX(-theta, wires=1)
+    return qml.probs()
+
+
+num_of_iters = int(sys.argv[1:][0])
+os.remove("my_toy_circuit.yml")
+with instrumentation("my_toy_circuit", filename="my_toy_circuit.yml", detailed=True):
+    res = circuit(12.3, num_of_iters)
+    # print(res)
+
+# with open('my_toy_circuit.yml', 'r') as f:
+#    print(f.read())
+
+
+####################### core PL #######################
+
+
+@qml.qnode(dev)
+def circuit_corePL(theta, loop_size):
+    for i in range(loop_size):
+        # for j in range(loop_size):
+        qml.Hadamard(0)
+        qml.Hadamard(0)
+        qml.RX(theta, wires=1)
+        qml.RX(-theta, wires=1)
+    return qml.probs()
+
+
+from matplotlib.patches import Rectangle
+
+draw = False
+if draw:
+    plt = qml.draw_mpl(circuit_corePL, show_all_wires=True, style="pennylane", label_options={"color": "white"})(1.23, 1)[0]
+    plt.text(0.5, 0.06, "repeated $N$ times", fontsize=12, ha="center")
+    square = Rectangle(
+        (-0.52, -0.5), 1, 2, linestyle="--", linewidth=1, edgecolor="r", facecolor="none"
+    )
+    plt.gca().add_patch(square)
+    plt.savefig("circuit_optimized")
+
+
+(tape,), _ = qml.workflow.construct_batch(circuit_corePL, level=0)(12.3, num_of_iters)
+start1 = timer()
+_ = qml.transforms.cancel_inverses(tape)
+end1 = timer()
+elapsed_seconds = end1 - start1
+elapsed_ms1 = elapsed_seconds * 1e3
+
+
+tape = _[0][0]  # <QuantumScript>
+start2 = timer()
+_ = qml.transforms.merge_rotations(tape)
+end2 = timer()
+elapsed_seconds = end2 - start2
+elapsed_ms2 = elapsed_seconds * 1e3
+
+total_elapsed_ms = elapsed_ms1 + elapsed_ms2
+with open("core_peephole_time.txt", "w") as f:
+    print(total_elapsed_ms, file=f)
+
+
+# res = circuit_corePL(12.3, num_of_iters)
+# print(res)

peephole_benchmark_driver.py

@@ -0,0 +1,132 @@
+import subprocess
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def do(command):
+    subprocess.call(command, shell=True)
+
+def stderr(list_of_data):
+    return np.std(list_of_data, ddof=1) / np.sqrt(len(list_of_data))
+
+def run_one_circuit(timings, core_PL_timings, num_of_iters):
+    do(f"python3 my_toy_circuit.py {num_of_iters}")
+
+    with open("my_toy_circuit.yml", "r") as f:
+        lines = f.readlines()
+        for i, line in enumerate(lines):
+            if "PeepholeBenchmarkPass" in line:
+                walltime = lines[i + 1]
+                walltime = float(walltime[walltime.find(":") + 2 :].strip("\n"))
+                cputime = lines[i + 2]
+                cputime = float(cputime[cputime.find(":") + 2 :].strip("\n"))
+                programsize = lines[i + 3]
+                programsize = int(programsize[programsize.find(":") + 2 :].strip("\n"))
+
+                timings.append(
+                    {"walltime": walltime, "cputime": cputime, "programsize": programsize}
+                )
+
+    with open("core_peephole_time.txt", "r") as f:
+        core_PL_time = float(f.readlines()[0].strip("\n"))
+        # print(core_PL_time)
+        core_PL_timings.append(core_PL_time)
+
+    return timings, core_PL_timings
+
+
+def collect_mean(raw_data, core_PL_timings):
+    """
+    raw_data is something like
+    [{'walltime': 0.307265, 'cputime': 0.306, 'programsize': 48},
+     {'walltime': 0.301666, 'cputime': 0.299, 'programsize': 48},
+     {'walltime': 0.342119, 'cputime': 0.341, 'programsize': 48},
+     {'walltime': 0.313134, 'cputime': 0.312, 'programsize': 48}]
+
+    core_PL_timings is just a list of numbers
+    """
+    mean_data = {"walltime": [], "cputime": [], "programsize": []}
+    for i, d in enumerate(raw_data):
+        mean_data["walltime"].append(d["walltime"])
+        mean_data["cputime"].append(d["cputime"])
+        mean_data["programsize"].append(d["programsize"])
+
+    _ = mean_data["walltime"]
+    mean_data["walltime"] = (np.mean(_), stderr(_))
+    _ = mean_data["cputime"]
+    mean_data["cputime"] = (np.mean(_), stderr(_))
+    _ = mean_data["programsize"]
+    mean_data["programsize"] = (np.mean(_), stderr(_))
+
+    return mean_data, (np.mean(core_PL_timings), stderr(core_PL_timings))
+
+
+def run_one_loopsize(loopsize):
+    timings = []
+    core_PL_timings = []
+    for i in range(3):
+        _ = run_one_circuit(timings, core_PL_timings, loopsize)
+        timings = _[0]
+        core_PL_timings = _[1]
+
+    # print(timings)
+    # print(collect_mean(timings))
+    return collect_mean(timings, core_PL_timings)
+
+
+############# main ##################
+# loopsizes = [10, 50, 100, 500, 1000, 5000, 10000]
+# loopsizes = [10, 20, 30, 40, 100, 150, 200]
+loopsizes = np.geomspace(10, 50000, 25, dtype=int)
+walltimes = []
+cputimes = []
+programsizes = []
+core_PL_times = []
+
+walltime_errs = []
+cputime_errs = []
+programsize_errs = []
+core_PL_time_errs = []
+
+for loopsize in loopsizes:
+    _ = run_one_loopsize(loopsize)
+    #breakpoint()
+    core_PL_times.append(_[1][0])
+    core_PL_time_errs.append(_[1][1])
+
+    catalyst_times = _[0]
+    walltimes.append(catalyst_times["walltime"][0])
+    cputimes.append(catalyst_times["cputime"][0])
+    programsizes.append(catalyst_times["programsize"][0])
+    walltime_errs.append(catalyst_times["walltime"][0])
+    cputime_errs.append(catalyst_times["cputime"][0])
+    programsize_errs.append(catalyst_times["programsize"][0])
+
+
+print(loopsizes, walltimes, cputimes, programsizes, core_PL_times)
+
+loopsizes = np.array(loopsizes)
+walltimes = np.array(walltimes)
+cputimes = np.array(cputimes)
+programsizes = np.array(programsizes)
+core_PL_times = np.array(core_PL_times)
+walltime_errs = np.array(walltime_errs)
+cputime_errs = np.array(cputime_errs)
+programsize_errs = np.array(programsize_errs)
+core_PL_time_errs = np.array(core_PL_time_errs)
+
+np.savez(
+    "timeit_peephole_benchmark_data_geom25_err",
+    loopsizes=loopsizes,
+    walltimes=walltimes,
+    cputimes=cputimes,
+    programsizes=programsizes,
+    core_PL_times=core_PL_times,
+    walltime_errs=walltime_errs,
+    cputime_errs=cputime_errs,
+    programsize_errs=programsize_errs,
+    core_PL_time_errs=core_PL_time_errs
+)
+
+do("rm -rf core_peephole_time.txt my_toy_circuit.yml")