fix: Working tests and benchmarks

benchmark/benchmark.py

@@ -6,6 +6,7 @@
 # always the same for repeatability
 np.random.seed(0x1C2C6D66)
 
+batch_size = (10, 100)
 n_qubits = range(3, 21)
 exact_shots_results = (
     "state_vector",
@@ -22,7 +23,7 @@
 )
 
 
-def generate_ghz(nq: int, result_type: str):
+def ghz(nq: int, result_type: str):
     source = f"OPENQASM 3.0;\nqubit[{nq}] q;\nh q[0];\n"
     for q in range(1, nq - 1):
         source += f"cnot q[0], q[{q}];\n"
@@ -31,32 +32,75 @@ def generate_ghz(nq: int, result_type: str):
     return source
 
 
+def qft(nq: int, result_type: str):
+    source = f"OPENQASM 3.0;\nqubit[{nq}] q;\n"
+    for q in range(nq - 1):
+        angle = np.pi / 2.0
+        source += f"h q[{q}];\n"
+        for ctrl_q in range(q + 1, nq - 1):
+            source += f"cphaseshift({angle}) q[{ctrl_q}], q[{q}];\n"
+            angle /= 2.0
+
+    source += f"#pragma braket result {result_type}\n"
+    return source
+
+
 def run_sim(oq3_prog, sim, shots):
     sim.run(oq3_prog, shots=shots)
     return
 
 
+def run_sim_batch(oq3_prog, sim, shots):
+    sim.run_batch(oq3_prog, shots=shots)
+    return
+
+
 device_ids = ("braket_sv", "braket_sv_v2", "braket_dm", "braket_dm_v2")
 
+generators = (ghz, qft)
+
 
 @pytest.mark.parametrize("device_id", device_ids)
 @pytest.mark.parametrize("nq", n_qubits)
 @pytest.mark.parametrize("exact_results", exact_shots_results)
-def test_exact_shots(benchmark, device_id, nq, exact_results):
+@pytest.mark.parametrize("circuit", generators)
+def test_exact_shots(benchmark, device_id, nq, exact_results, circuit):
     if device_id in ("braket_dm_v2", "braket_dm") and (
         exact_results in ("state_vector",) or nq > 10
     ):
         pytest.skip()
     if (
         device_id in ("braket_sv",)
         and exact_results in ("density_matrix q[0], q[1]",)
-        and nq > 17
+        and nq >= 17
     ):
         pytest.skip()
     result_type = exact_results
-    oq3_prog = Program(source=generate_ghz(nq, result_type))
+    oq3_prog = Program(source=circuit(nq, result_type))
     sim = LocalSimulator(device_id)
-    benchmark.pedantic(run_sim, args=(oq3_prog, sim, 0))
+    benchmark.pedantic(run_sim, args=(oq3_prog, sim, 0), iterations=5, warmup_rounds=1)
+
+
+@pytest.mark.parametrize("device_id", device_ids)
+@pytest.mark.parametrize("nq", n_qubits)
+@pytest.mark.parametrize("batch_size", batch_size)
+@pytest.mark.parametrize("exact_results", exact_shots_results)
+@pytest.mark.parametrize("circuit", generators)
+def test_exact_shots_batched(
+    benchmark, device_id, nq, batch_size, exact_results, circuit
+):
+    if device_id in ("braket_dm_v2", "braket_dm") and (exact_results in ("state_vector,") or nq >= 5):
+        pytest.skip()
+    if nq >= 10:
+        pytest.skip()
+    # skip all for now as this is very expensive
+    pytest.skip()
+    result_type = exact_results
+    oq3_prog = [Program(source=circuit(nq, result_type)) for _ in range(batch_size)]
+    sim = LocalSimulator(device_id)
+    benchmark.pedantic(
+        run_sim_batch, args=(oq3_prog, sim, 0), iterations=5, warmup_rounds=1
+    )
 
 
 shots = (100,)
@@ -66,10 +110,40 @@ def test_exact_shots(benchmark, device_id, nq, exact_results):
 @pytest.mark.parametrize("nq", n_qubits)
 @pytest.mark.parametrize("shots", shots)
 @pytest.mark.parametrize("nonzero_shots_results", nonzero_shots_results)
-def test_nonzero_shots(benchmark, device_id, nq, shots, nonzero_shots_results):
+@pytest.mark.parametrize("circuit", generators)
+def test_nonzero_shots(benchmark, device_id, nq, shots, nonzero_shots_results, circuit):
     if device_id in ("braket_dm_v2", "braket_dm") and nq > 10:
         pytest.skip()
     result_type = nonzero_shots_results
-    oq3_prog = Program(source=generate_ghz(nq, result_type))
+    oq3_prog = Program(source=circuit(nq, result_type))
+    sim = LocalSimulator(device_id)
+    benchmark.pedantic(
+        run_sim, args=(oq3_prog, sim, shots), iterations=5, warmup_rounds=1
+    )
+    del sim
+
+
+@pytest.mark.parametrize("device_id", device_ids)
+@pytest.mark.parametrize("nq", n_qubits)
+@pytest.mark.parametrize("batch_size", batch_size)
+@pytest.mark.parametrize("shots", shots)
+@pytest.mark.parametrize("nonzero_shots_results", nonzero_shots_results)
+@pytest.mark.parametrize("circuit", generators)
+def test_nonzero_shots_batched(
+    benchmark, device_id, nq, batch_size, shots, nonzero_shots_results, circuit
+):
+    if device_id in ("braket_dm_v2", "braket_dm") and nq >= 5:
+        pytest.skip()
+    if nq >= 10:
+        pytest.skip()
+
+    # skip all for now as this is very expensive
+    pytest.skip()
+
+    result_type = nonzero_shots_results
+    oq3_prog = [Program(source=circuit(nq, result_type)) for _ in range(batch_size)]
     sim = LocalSimulator(device_id)
-    benchmark.pedantic(run_sim, args=(oq3_prog, sim, shots))
+    benchmark.pedantic(
+        run_sim_batch, args=(oq3_prog, sim, shots), iterations=5, warmup_rounds=1
+    )
+    del sim