diff --git a/pyquil/api/_qpu.py b/pyquil/api/_qpu.py index d7aab078e..676791e20 100644 --- a/pyquil/api/_qpu.py +++ b/pyquil/api/_qpu.py @@ -58,11 +58,7 @@ def _extract_memory_regions( ro_sources: Dict[MemoryReference, str], buffers: Dict[str, np.ndarray], ) -> Dict[str, np.ndarray]: - # hack to extract num_shots indirectly from the shape of the returned data - first, *rest = buffers.values() - num_shots = first.shape[0] - - def alloc(spec: ParameterSpec) -> np.ndarray: + def alloc(spec: ParameterSpec, num_shots: int) -> np.ndarray: dtype = { "BIT": np.int64, "INTEGER": np.int64, @@ -76,31 +72,45 @@ def alloc(spec: ParameterSpec) -> np.ndarray: regions: Dict[str, np.ndarray] = {} + # filter out the ro_sources to only get the ones that are used and group + # mrefs by their name + sources = {mem_desc: defaultdict(list) for mem_desc in memory_descriptors} for mref, key in ro_sources.items(): - # Translation sometimes introduces ro_sources that the user didn't ask for. - # That's fine, we just ignore them. - if mref.name not in memory_descriptors: + # Translation sometimes introduces ro_sources that the user didn't + # ask for. That's fine, we just ignore them. + if mref.name not in sources: continue - elif mref.name not in regions: - regions[mref.name] = alloc(memory_descriptors[mref.name]) - - buf = buffers[key] - if buf.ndim == 1: - buf = buf.reshape((num_shots, 1)) - - if np.iscomplexobj(buf): - buf = np.column_stack((buf.real, buf.imag)) - _, width = buf.shape - - end = mref.offset + width - region_width = memory_descriptors[mref.name].length - if end > region_width: - raise ValueError( - f"Attempted to fill {mref.name}[{mref.offset}, {end})" - f"but the declared region has width {region_width}." - ) - - regions[mref.name][:, mref.offset : end] = buf + sources[mref.name][key].append(mref) + + # hack to extract the number of shots + # the number of shots is the length of the buffer + # divided by how many sources are in that buffer + # it should be the same for all buffers + mref_name = next(iter(sources)) + buf_key = next(iter(sources[mref_name])) + num_shots = buffers[buf_key].shape[0] // len(sources[mref_name][buf_key]) + + for mref_name, mref_dict in sources.items(): + regions[mref_name] = alloc(memory_descriptors[mref_name], num_shots) + region_width = memory_descriptors[mref_name].length + for key, mref_list in mref_dict.items(): + num_lanes = len(mref_list) + # sort the items in the mref_list by their offset + sorted_mref = [mref for _, mref in sorted(zip((m.offset for m in mref_list), mref_list))] + buf = buffers[key] + if buf.ndim == 1: + buf = buf[:, None] + if np.iscomplexobj(buf): + buf = np.column_stack((buf.real, buf.imag)) + _, width = buf.shape + for i, mref in enumerate(sorted_mref): + end = mref.offset + width + if end > region_width: + raise ValueError( + f"Attempted to fill {mref_name}[{mref.offset}, {end})" + f"but the declared region has width {region_width}." + ) + regions[mref_name][:, mref.offset : end] = buf[i::num_lanes, :] return regions