microsoft/qdk
Publicmirrored from https://github.com/microsoft/qdkAvailable
source/pip/qsharp/_device/_atom/__init__.py
299lines · modecode
| 1 | # Copyright (c) Microsoft Corporation. |
| 2 | # Licensed under the MIT License. |
| 3 | |
| 4 | from .._device import Device, Zone, ZoneType |
| 5 | from ..._simulation import NoiseConfig, run_qir_clifford, run_qir_cpu, run_qir_gpu |
| 6 | from ..._native import try_create_gpu_adapter |
| 7 | from ..._qsharp import QirInputData |
| 8 | from ... import telemetry_events |
| 9 | |
| 10 | from typing import List, Literal, Optional |
| 11 | import time |
| 12 | |
| 13 | |
| 14 | class NeutralAtomDevice(Device): |
| 15 | """ |
| 16 | Representation of a neutral atom device quantum computer. |
| 17 | """ |
| 18 | |
| 19 | def __init__( |
| 20 | self, |
| 21 | column_count: int = 40, |
| 22 | register_zone_row_count: int = 25, |
| 23 | interaction_zone_row_count: int = 2, |
| 24 | measurement_zone_row_count: int = 2, |
| 25 | ): |
| 26 | default_layout = ( |
| 27 | column_count == 40 |
| 28 | and register_zone_row_count == 25 |
| 29 | and interaction_zone_row_count == 2 |
| 30 | and measurement_zone_row_count == 2 |
| 31 | ) |
| 32 | telemetry_events.on_neutral_atom_init(default_layout) |
| 33 | |
| 34 | super().__init__( |
| 35 | column_count, |
| 36 | [ |
| 37 | Zone("Register 1", register_zone_row_count, ZoneType.REG), |
| 38 | Zone("Interaction Zone", interaction_zone_row_count, ZoneType.INTER), |
| 39 | Zone("Measurement Zone", measurement_zone_row_count, ZoneType.MEAS), |
| 40 | ], |
| 41 | ) |
| 42 | |
| 43 | def _init_home_locs(self): |
| 44 | # Set up the home locations for qubits in the NeutralAtomDevice layout. |
| 45 | assert len(self.zones) == 3 |
| 46 | assert ( |
| 47 | self.zones[0].type == ZoneType.REG |
| 48 | and self.zones[1].type == ZoneType.INTER |
| 49 | and self.zones[2].type == ZoneType.MEAS |
| 50 | ) |
| 51 | rz1_rows = range(self.zones[0].row_count - 1, -1, -1) |
| 52 | self.home_locs = [] |
| 53 | for row in range(self.zones[0].row_count): |
| 54 | for col in range(self.column_count): |
| 55 | self.home_locs.append((rz1_rows[row], col)) |
| 56 | |
| 57 | def compile( |
| 58 | self, |
| 59 | program: str | QirInputData, |
| 60 | verbose: bool = False, |
| 61 | ) -> QirInputData: |
| 62 | """ |
| 63 | Compile a QIR program for the NeutralAtomDevice device. This includes decomposing gates to the native gate set, |
| 64 | optimizing sequences of single qubit gates, pruning unused functions, and reordering instructions to |
| 65 | enable better scheduling during execution. |
| 66 | |
| 67 | :param program: The QIR program to compile, either as a string or as QirInputData. |
| 68 | :param verbose: If true, print detailed information about each compilation step. |
| 69 | :returns QirInputData: The compiled QIR program. |
| 70 | """ |
| 71 | |
| 72 | from ._optimize import ( |
| 73 | OptimizeSingleQubitGates, |
| 74 | PruneUnusedFunctions, |
| 75 | ) |
| 76 | from ._decomp import ( |
| 77 | DecomposeMultiQubitToCZ, |
| 78 | DecomposeSingleRotationToRz, |
| 79 | DecomposeSingleQubitToRzSX, |
| 80 | ReplaceResetWithMResetZ, |
| 81 | ) |
| 82 | from ._reorder import Reorder |
| 83 | from pyqir import Module, Context |
| 84 | |
| 85 | start_time = time.monotonic() |
| 86 | all_start_time = start_time |
| 87 | telemetry_events.on_neutral_atom_compile() |
| 88 | |
| 89 | name = "" |
| 90 | if isinstance(program, QirInputData): |
| 91 | name = program._name |
| 92 | |
| 93 | if verbose: |
| 94 | print(f"Compiling program {name} for NeutralAtomDevice device...") |
| 95 | |
| 96 | module = Module.from_ir(Context(), str(program)) |
| 97 | if verbose: |
| 98 | end_time = time.monotonic() |
| 99 | print(f" Loaded module in {end_time - start_time:.2f} seconds") |
| 100 | start_time = end_time |
| 101 | |
| 102 | OptimizeSingleQubitGates().run(module) |
| 103 | if verbose: |
| 104 | end_time = time.monotonic() |
| 105 | print( |
| 106 | f" Optimized single qubit gates in {end_time - start_time:.2f} seconds" |
| 107 | ) |
| 108 | start_time = end_time |
| 109 | |
| 110 | DecomposeMultiQubitToCZ().run(module) |
| 111 | if verbose: |
| 112 | end_time = time.monotonic() |
| 113 | print( |
| 114 | f" Decomposed multi-qubit gates to CZ in {end_time - start_time:.2f} seconds" |
| 115 | ) |
| 116 | start_time = end_time |
| 117 | |
| 118 | OptimizeSingleQubitGates().run(module) |
| 119 | if verbose: |
| 120 | end_time = time.monotonic() |
| 121 | print( |
| 122 | f" Optimized single qubit gates in {end_time - start_time:.2f} seconds" |
| 123 | ) |
| 124 | start_time = end_time |
| 125 | |
| 126 | DecomposeSingleRotationToRz().run(module) |
| 127 | if verbose: |
| 128 | end_time = time.monotonic() |
| 129 | print( |
| 130 | f" Decomposed single rotations to Rz in {end_time - start_time:.2f} seconds" |
| 131 | ) |
| 132 | start_time = end_time |
| 133 | |
| 134 | OptimizeSingleQubitGates().run(module) |
| 135 | if verbose: |
| 136 | end_time = time.monotonic() |
| 137 | print( |
| 138 | f" Optimized single qubit gates in {end_time - start_time:.2f} seconds" |
| 139 | ) |
| 140 | start_time = end_time |
| 141 | |
| 142 | DecomposeSingleQubitToRzSX().run(module) |
| 143 | if verbose: |
| 144 | end_time = time.monotonic() |
| 145 | print( |
| 146 | f" Decomposed single qubit gates to Rz and SX in {end_time - start_time:.2f} seconds" |
| 147 | ) |
| 148 | start_time = end_time |
| 149 | |
| 150 | OptimizeSingleQubitGates().run(module) |
| 151 | if verbose: |
| 152 | end_time = time.monotonic() |
| 153 | print( |
| 154 | f" Optimized single qubit gates in {end_time - start_time:.2f} seconds" |
| 155 | ) |
| 156 | start_time = end_time |
| 157 | |
| 158 | ReplaceResetWithMResetZ().run(module) |
| 159 | if verbose: |
| 160 | end_time = time.monotonic() |
| 161 | print( |
| 162 | f" Replaced resets with mresetz in {end_time - start_time:.2f} seconds" |
| 163 | ) |
| 164 | start_time = end_time |
| 165 | |
| 166 | PruneUnusedFunctions().run(module) |
| 167 | if verbose: |
| 168 | end_time = time.monotonic() |
| 169 | print(f" Pruned unused functions in {end_time - start_time:.2f} seconds") |
| 170 | start_time = end_time |
| 171 | |
| 172 | Reorder(self).run(module) |
| 173 | if verbose: |
| 174 | end_time = time.monotonic() |
| 175 | print(f" Reordered instructions in {end_time - start_time:.2f} seconds") |
| 176 | start_time = end_time |
| 177 | |
| 178 | end_time = time.monotonic() |
| 179 | telemetry_events.on_neutral_atom_compile_end((end_time - all_start_time) * 1000) |
| 180 | if verbose: |
| 181 | print( |
| 182 | f"Finished compiling program {name} in {end_time - all_start_time:.2f} seconds" |
| 183 | ) |
| 184 | |
| 185 | return QirInputData(name, str(module)) |
| 186 | |
| 187 | def show_trace(self, qir: str | QirInputData): |
| 188 | """ |
| 189 | Visualize the execution trace of a QIR program on the NeutralAtomDevice device using the Atoms widget. |
| 190 | This includes approximate layout and scheduling of the program to show the parallelism of gates and |
| 191 | movement of qubits during execution. |
| 192 | |
| 193 | :param qir: The QIR program to visualize, either as a string or as QirInputData. |
| 194 | """ |
| 195 | |
| 196 | try: |
| 197 | from qsharp_widgets import Atoms |
| 198 | except ImportError: |
| 199 | raise ImportError( |
| 200 | "The qsharp-widgets package is required for showing atom trace visualization. " |
| 201 | "Please install it via 'pip install \"qdk[jupyter]\"' or 'pip install qsharp-widgets'." |
| 202 | ) |
| 203 | from ._trace import Trace |
| 204 | from ._validate import ValidateNoConditionalBranches |
| 205 | from ._scheduler import Schedule |
| 206 | from pyqir import Module, Context |
| 207 | from IPython.display import display |
| 208 | |
| 209 | start_time = time.monotonic() |
| 210 | telemetry_events.on_neutral_atom_trace() |
| 211 | |
| 212 | # Compile and visualize the trace in one step. |
| 213 | compiled = self.compile(qir) |
| 214 | module = Module.from_ir(Context(), str(compiled)) |
| 215 | ValidateNoConditionalBranches().run(module) |
| 216 | Schedule(self).run(module) |
| 217 | tracer = Trace(self) |
| 218 | tracer.run(module) |
| 219 | display(Atoms(machine_layout=self.get_layout(), trace_data=tracer.trace)) |
| 220 | |
| 221 | end_time = time.monotonic() |
| 222 | telemetry_events.on_neutral_atom_trace_end((end_time - start_time) * 1000) |
| 223 | |
| 224 | def simulate( |
| 225 | self, |
| 226 | qir: str | QirInputData, |
| 227 | shots=1, |
| 228 | noise: NoiseConfig | None = None, |
| 229 | type: Optional[Literal["clifford", "cpu", "gpu"]] = None, |
| 230 | seed: Optional[int] = None, |
| 231 | ) -> List: |
| 232 | """ |
| 233 | Simulate a QIR program on the NeutralAtomDevice device. This includes approximate layout and scheduling of the program |
| 234 | to model the parallelism of gates and movement of qubits during execution. The simulation can optionally |
| 235 | include noise based on a provided noise configuration. |
| 236 | |
| 237 | :param qir: The QIR program to simulate, either as a string or as QirInputData. |
| 238 | :param shots: The number of shots to simulate. Defaults to 1. |
| 239 | :param noise: An optional NoiseConfig to include noise in the simulation. |
| 240 | :param type: The type of simulator to use: |
| 241 | Use `"clifford"` if your QIR only contains Clifford gates and measurements. |
| 242 | Use `"gpu"` if you have a GPU available in your system. |
| 243 | Use `"cpu"` as a fallback option if you don't have a GPU in your system. |
| 244 | If `None` (default), the GPU simulator will be tried first, falling back to |
| 245 | CPU if a suitable GPU device could not be located. |
| 246 | :param seed: An optional random seed for reproducibility. |
| 247 | :returns: The results of each shot of the simulation as a list. |
| 248 | """ |
| 249 | |
| 250 | from ._validate import ValidateNoConditionalBranches |
| 251 | from ._scheduler import Schedule |
| 252 | from ._decomp import DecomposeRzAnglesToCliffordGates |
| 253 | from pyqir import Module, Context |
| 254 | |
| 255 | start_time = time.monotonic() |
| 256 | |
| 257 | using_noise = noise is not None |
| 258 | if noise is None: |
| 259 | noise = NoiseConfig() |
| 260 | |
| 261 | compiled = self.compile(qir) |
| 262 | module = Module.from_ir(Context(), str(compiled)) |
| 263 | ValidateNoConditionalBranches().run(module) |
| 264 | Schedule(self).run(module) |
| 265 | |
| 266 | if type is None: |
| 267 | try: |
| 268 | try_create_gpu_adapter() |
| 269 | type = "gpu" |
| 270 | except OSError: |
| 271 | telemetry_events.on_neutral_atom_cpu_fallback() |
| 272 | type = "cpu" |
| 273 | |
| 274 | telemetry_events.on_neutral_atom_simulate(shots, using_noise, type) |
| 275 | |
| 276 | match type: |
| 277 | case "clifford": |
| 278 | DecomposeRzAnglesToCliffordGates().run(module) |
| 279 | result = run_qir_clifford( |
| 280 | str(module), |
| 281 | shots, |
| 282 | noise, |
| 283 | seed, |
| 284 | ) |
| 285 | case "cpu": |
| 286 | result = run_qir_cpu(str(module), shots, noise, seed) |
| 287 | case "gpu": |
| 288 | result = run_qir_gpu(str(module), shots, noise, seed) |
| 289 | case _: |
| 290 | raise ValueError(f"Simulation type {type} is not supported") |
| 291 | |
| 292 | end_time = time.monotonic() |
| 293 | telemetry_events.on_neutral_atom_simulate_end( |
| 294 | (end_time - start_time) * 1000, shots, using_noise, type |
| 295 | ) |
| 296 | return result |
| 297 | |
| 298 | |
| 299 | __all__ = ["NeutralAtomDevice"] |
| 300 | |