consolidate circuits with the same number of shots

qiskit/primitives/backend_sampler_v2.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import warnings
+from collections import defaultdict
 from dataclasses import dataclass
 from typing import Iterable
 
@@ -142,36 +143,74 @@ def _validate_pubs(self, pubs: list[SamplerPub]):
                 )
 
     def _run(self, pubs: Iterable[SamplerPub]) -> PrimitiveResult[PubResult]:
-        results = [self._run_pub(pub) for pub in pubs]
+        pub_dict = defaultdict(list)
+        # consolidate pubs with the same number of shots
+        for i, pub in enumerate(pubs):
+            pub_dict[pub.shots].append(i)
+
+        results = [None] * len(pubs)
+        for shots, lst in pub_dict.items():
+            # run pubs with the same number of shots at once
+            pub_results = self._run_pubs([pubs[i] for i in lst], shots)
+            # reconstruct the result of pubs
+            for i, pub_result in zip(lst, pub_results):
+                results[i] = pub_result
         return PrimitiveResult(results)
 
-    def _run_pub(self, pub: SamplerPub) -> PubResult:
-        meas_info, max_num_bytes = _analyze_circuit(pub.circuit)
-        bound_circuits = pub.parameter_values.bind_all(pub.circuit)
-        arrays = {
-            item.creg_name: np.zeros(
-                bound_circuits.shape + (pub.shots, item.num_bytes), dtype=np.uint8
-            )
-            for item in meas_info
-        }
-        flatten_circuits = np.ravel(bound_circuits).tolist()
-        result_memory, _ = _run_circuits(
+    def _run_pubs(self, pubs: list[SamplerPub], shots: int) -> list[PubResult]:
+        # prepare circuits
+        bound_circuits = [pub.parameter_values.bind_all(pub.circuit) for pub in pubs]
+        flatten_circuits = []
+        for circuits in bound_circuits:
+            flatten_circuits.extend(np.ravel(circuits).tolist())
+
+        # run circuits
+        results, _ = _run_circuits(
             flatten_circuits,
             self._backend,
             memory=True,
-            shots=pub.shots,
+            shots=shots,
             seed_simulator=self._options.seed_simulator,
         )
-        memory_list = _prepare_memory(result_memory, max_num_bytes)
+        result_memory = _prepare_memory(results)
+
+        # pack memory to an ndarray of uint8
+        results = []
+        start = 0
+        for pub, bound in zip(pubs, bound_circuits):
+            meas_info, max_num_bytes = _analyze_circuit(pub.circuit)
+            end = start + bound.size
+            results.append(
+                self._postprocessing_pub(
+                    result_memory[start:end], shots, bound.shape, meas_info, max_num_bytes
+                )
+            )
+            start = end
+
+        return results
 
-        for samples, index in zip(memory_list, np.ndindex(*bound_circuits.shape)):
+    def _postprocessing_pub(
+        self,
+        result_memory: list[list[str]],
+        shots: int,
+        shape: tuple[int, ...],
+        meas_info: list[_MeasureInfo],
+        max_num_bytes: int,
+    ) -> PubResult:
+        arrays = {
+            item.creg_name: np.zeros(shape + (shots, item.num_bytes), dtype=np.uint8)
+            for item in meas_info
+        }
+        memory_array = _memory_array(result_memory, max_num_bytes)
+
+        for samples, index in zip(memory_array, np.ndindex(*shape)):
             for item in meas_info:
                 ary = _samples_to_packed_array(samples, item.num_bits, item.start)
                 arrays[item.creg_name][index] = ary
 
         data_bin_cls = make_data_bin(
             [(item.creg_name, BitArray) for item in meas_info],
-            shape=bound_circuits.shape,
+            shape=shape,
         )
         meas = {
             item.creg_name: BitArray(arrays[item.creg_name], item.num_bits) for item in meas_info
@@ -202,17 +241,29 @@ def _analyze_circuit(circuit: QuantumCircuit) -> tuple[list[_MeasureInfo], int]:
     return meas_info, _min_num_bytes(max_num_bits)
 
 
-def _prepare_memory(results: list[Result], num_bytes: int) -> NDArray[np.uint8]:
+def _prepare_memory(results: list[Result]) -> list[list[str]]:
+    """Join splitted results if exceeding max_experiments"""
     lst = []
     for res in results:
         for exp in res.results:
             if hasattr(exp.data, "memory") and exp.data.memory:
-                data = b"".join(int(i, 16).to_bytes(num_bytes, "big") for i in exp.data.memory)
-                data = np.frombuffer(data, dtype=np.uint8).reshape(-1, num_bytes)
+                lst.append(exp.data.memory)
             else:
                 # no measure in a circuit
-                data = np.zeros((exp.shots, num_bytes), dtype=np.uint8)
-            lst.append(data)
+                lst.append(["0x0"] * exp.shots)
+    return lst
+
+
+def _memory_array(results: list[list[str]], num_bytes: int) -> NDArray[np.uint8]:
+    lst = []
+    for memory in results:
+        if num_bytes > 0:
+            data = b"".join(int(i, 16).to_bytes(num_bytes, "big") for i in memory)
+            data = np.frombuffer(data, dtype=np.uint8).reshape(-1, num_bytes)
+        else:
+            # no measure in a circuit
+            data = np.zeros((len(memory), num_bytes), dtype=np.uint8)
+        lst.append(data)
     ary = np.array(lst, copy=False)
     return np.unpackbits(ary, axis=-1, bitorder="big")
 

test/python/primitives/test_backend_sampler_v2.py

@@ -671,6 +671,24 @@ def test_empty_creg(self, backend):
         result = sampler.run([qc], shots=self._shots).result()
         self.assertEqual(result[0].data.c1.array.shape, (self._shots, 0))
 
+    @combine(backend=BACKENDS)
+    def test_diff_shots(self, backend):
+        """Test of pubs with different shots"""
+        pm = generate_preset_pass_manager(optimization_level=0, backend=backend)
+
+        bell, _, target = self._cases[1]
+        bell = pm.run(bell)
+        sampler = BackendSamplerV2(backend=backend, options=self._options)
+        shots2 = self._shots + 2
+        target2 = {k: v + 1 for k, v in target.items()}
+        job = sampler.run([(bell, None, self._shots), (bell, None, shots2)])
+        result = job.result()
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0].data.meas.num_shots, self._shots)
+        self._assert_allclose(result[0].data.meas, np.array(target))
+        self.assertEqual(result[1].data.meas.num_shots, shots2)
+        self._assert_allclose(result[1].data.meas, np.array(target2))
+
 
 if __name__ == "__main__":
     unittest.main()