microsoft/qdk
Publicmirrored from https://github.com/microsoft/qdkAvailable
source/pip/qsharp/_device/_atom/_optimize.py
315lines · modecode
| 1 | # Copyright (c) Microsoft Corporation. |
| 2 | # Licensed under the MIT License. |
| 3 | |
| 4 | from pyqir import ( |
| 5 | Type, |
| 6 | Function, |
| 7 | FunctionType, |
| 8 | FloatConstant, |
| 9 | Linkage, |
| 10 | PointerType, |
| 11 | const, |
| 12 | ptr_id, |
| 13 | is_entry_point, |
| 14 | QirModuleVisitor, |
| 15 | ) |
| 16 | from math import pi |
| 17 | |
| 18 | from ._utils import TOLERANCE |
| 19 | |
| 20 | |
| 21 | class OptimizeSingleQubitGates(QirModuleVisitor): |
| 22 | """ |
| 23 | Optimizes single qubit gates by looking for sequences of a gate and its adjoint on a given qubit. |
| 24 | Will also try to replace certain patterns with simpler gates. |
| 25 | """ |
| 26 | |
| 27 | sx_func: Function |
| 28 | mresetz_func: Function |
| 29 | |
| 30 | def _on_module(self, module): |
| 31 | void = Type.void(module.context) |
| 32 | qubit_ty = PointerType(Type.void(module.context)) |
| 33 | result_ty = qubit_ty |
| 34 | self.double_ty = Type.double(module.context) |
| 35 | self.used_qubits = set() |
| 36 | # Find or create the intrinsic gate functions |
| 37 | for func in module.functions: |
| 38 | match func.name: |
| 39 | case "__quantum__qis__mresetz__body": |
| 40 | self.mresetz_func = func |
| 41 | case "__quantum__qis__sx__body": |
| 42 | self.sx_func = func |
| 43 | if not hasattr(self, "sx_fun"): |
| 44 | self.sx_func = Function( |
| 45 | FunctionType(void, [qubit_ty]), |
| 46 | Linkage.EXTERNAL, |
| 47 | "__quantum__qis__sx__body", |
| 48 | module, |
| 49 | ) |
| 50 | if not hasattr(self, "mresetz_func"): |
| 51 | self.mresetz_func = Function( |
| 52 | FunctionType(void, [qubit_ty, result_ty]), |
| 53 | Linkage.EXTERNAL, |
| 54 | "__quantum__qis__mresetz__body", |
| 55 | module, |
| 56 | ) |
| 57 | super()._on_module(module) |
| 58 | |
| 59 | def _drop_ops(self, qubits): |
| 60 | # Since instructions are only removed when they are canceled out by their adjoint or folded with another |
| 61 | # instruction, we can just pop the entries for these qubits so they start fresh with the next gates. |
| 62 | for qubit in qubits: |
| 63 | q = ptr_id(qubit) |
| 64 | self.qubit_ops.pop(q, None) |
| 65 | self.last_meas.pop(q, None) |
| 66 | self.used_qubits.add(q) |
| 67 | |
| 68 | def _schedule_gate(self, instr, key, name, adj): |
| 69 | if key in self.qubit_ops: |
| 70 | # There are previous operations on this qubit, so check if the last one was the adjoint of this one. |
| 71 | if self.qubit_ops[key][-1][1] == adj: |
| 72 | (other_instr, _) = self.qubit_ops[key].pop() |
| 73 | # Erase the adjoint instruction and the current instruction since they cancel out. |
| 74 | other_instr.erase() |
| 75 | instr.erase() |
| 76 | elif ( |
| 77 | len(self.qubit_ops[key]) > 1 |
| 78 | and name == "h" |
| 79 | and self.qubit_ops[key][-1][1] == "s" |
| 80 | and self.qubit_ops[key][-2][1] == "h" |
| 81 | ): |
| 82 | # We have a sequence of h s h, which can be replaced with a single sx. |
| 83 | self.builder.insert_before(instr) |
| 84 | self.builder.call(self.sx_func, [instr.args[0]]) |
| 85 | instr.erase() |
| 86 | (other_instr, _) = self.qubit_ops[key].pop() |
| 87 | other_instr.erase() |
| 88 | (other_instr, _) = self.qubit_ops[key].pop() |
| 89 | other_instr.erase() |
| 90 | else: |
| 91 | # The last operation was not the adjoint of this one, so add this instruction to the list. |
| 92 | self.qubit_ops[key].append((instr, name)) |
| 93 | self.used_qubits.add(key) |
| 94 | self.last_meas.pop(key, None) |
| 95 | |
| 96 | if len(self.qubit_ops[key]) == 0: |
| 97 | # There are no more operations on this qubit, so pop it's entry to avoid having empty lists in the dict. |
| 98 | self.qubit_ops.pop(key) |
| 99 | |
| 100 | else: |
| 101 | # No previous operations on this qubit, so create a new list from this instruction. |
| 102 | self.qubit_ops[key] = [(instr, name)] |
| 103 | self.used_qubits.add(key) |
| 104 | self.last_meas.pop(key, None) |
| 105 | |
| 106 | def _schedule_rotation(self, instr, key, name): |
| 107 | if isinstance(instr.args[0], FloatConstant): |
| 108 | # The angle is constant, so we can try to fold this rotation with other instances of the same rotation |
| 109 | # tht are constant. |
| 110 | if key in self.qubit_ops: |
| 111 | if self.qubit_ops[key][-1][1] == name and isinstance( |
| 112 | self.qubit_ops[key][-1][0].args[0], FloatConstant |
| 113 | ): |
| 114 | # The last operation on this qubit was also a rotation of the same type by a constant angle. |
| 115 | (other_instr, _) = self.qubit_ops[key].pop() |
| 116 | new_angle = instr.args[0].value + other_instr.args[0].value |
| 117 | sign = -1 if new_angle < 0 else 1 |
| 118 | abs_new_angle = abs(new_angle) |
| 119 | # Normalize the angle to be within 0 to 2*pi |
| 120 | while abs_new_angle > 2 * pi: |
| 121 | abs_new_angle -= 2 * pi |
| 122 | new_angle = sign * abs_new_angle |
| 123 | if ( |
| 124 | abs(new_angle) > TOLERANCE |
| 125 | and abs(abs(new_angle) - (2 * pi)) > TOLERANCE |
| 126 | ): |
| 127 | # Create a new rotation instruction with the sum of the angles, |
| 128 | # and insert it, but only if the angle is above our threshold. |
| 129 | self.builder.insert_before(instr) |
| 130 | new_instr = self.builder.call( |
| 131 | instr.callee, |
| 132 | [const(self.double_ty, new_angle), instr.args[1]], |
| 133 | ) |
| 134 | self.qubit_ops[key].append((new_instr, name)) |
| 135 | self.used_qubits.add(key) |
| 136 | self.last_meas.pop(key, None) |
| 137 | # Erase the old instructions the new rotation replaces. |
| 138 | other_instr.erase() |
| 139 | instr.erase() |
| 140 | else: |
| 141 | # Can't fold this rotation with the previous one, so just add it to the list. |
| 142 | self.qubit_ops[key].append((instr, name)) |
| 143 | self.used_qubits.add(key) |
| 144 | self.last_meas.pop(key, None) |
| 145 | |
| 146 | if len(self.qubit_ops[key]) == 0: |
| 147 | # There are no more operations on this qubit, so pop it's entry to avoid having empty lists in the dict. |
| 148 | self.qubit_ops.pop(key) |
| 149 | |
| 150 | else: |
| 151 | # No previous operations on this qubit, so create a new list from this instruction. |
| 152 | self.qubit_ops[key] = [(instr, name)] |
| 153 | self.used_qubits.add(key) |
| 154 | self.last_meas.pop(key, None) |
| 155 | else: |
| 156 | # This angle is not constant, so append it to the list of operations on this qubit. |
| 157 | if key in self.qubit_ops: |
| 158 | self.qubit_ops[key].append((instr, name)) |
| 159 | else: |
| 160 | self.qubit_ops[key] = [(instr, name)] |
| 161 | self.used_qubits.add(key) |
| 162 | self.last_meas.pop(key, None) |
| 163 | |
| 164 | def _on_function(self, function): |
| 165 | self.last_meas = {} |
| 166 | self.qubit_ops = {} |
| 167 | super()._on_function(function) |
| 168 | # At the end of a function, if there are any remaining entries in self.last_meas, it means |
| 169 | # that there were measurements on qubits that were never reset. Convert those into mresetz. |
| 170 | for key, (instr, target, result) in self.last_meas.items(): |
| 171 | self.builder.insert_before(instr) |
| 172 | self.builder.call( |
| 173 | self.mresetz_func, |
| 174 | [target, result], |
| 175 | ) |
| 176 | instr.erase() |
| 177 | for key in self.qubit_ops: |
| 178 | if self.qubit_ops[key][-1][1] == "reset": |
| 179 | # The last operation on this qubit was a reset, so we can drop it. |
| 180 | (instr, _) = self.qubit_ops[key].pop() |
| 181 | instr.erase() |
| 182 | |
| 183 | def _on_block(self, block): |
| 184 | # Each block is independent, so start from an empty list of operations per qubit. |
| 185 | self.qubit_ops = {} |
| 186 | self.last_meas = {} |
| 187 | super()._on_block(block) |
| 188 | |
| 189 | def _on_call_instr(self, call): |
| 190 | if call.callee.name == "__quantum__qis__sx__body": |
| 191 | self._drop_ops([call.args[0]]) |
| 192 | elif call.callee.name == "__quantum__qis__move__body": |
| 193 | self._drop_ops([call.args[0]]) |
| 194 | elif call.callee.name == "__quantum__qis__barrier__body": |
| 195 | # Don't optimize across barrier calls. Treat this as a drop of all tracked gates, |
| 196 | # which effectively flushes all scheduled operations. |
| 197 | self.qubit_ops = {} |
| 198 | self.last_meas = {} |
| 199 | else: |
| 200 | super()._on_call_instr(call) |
| 201 | |
| 202 | def _on_qis_h(self, call, target): |
| 203 | self._schedule_gate(call, ptr_id(target), "h", "h") |
| 204 | |
| 205 | def _on_qis_s(self, call, target): |
| 206 | self._schedule_gate(call, ptr_id(target), "s", "s_adj") |
| 207 | |
| 208 | def _on_qis_s_adj(self, call, target): |
| 209 | self._schedule_gate(call, ptr_id(target), "s_adj", "s") |
| 210 | |
| 211 | def _on_qis_t(self, call, target): |
| 212 | self._schedule_gate(call, ptr_id(target), "t", "t_adj") |
| 213 | |
| 214 | def _on_qis_t_adj(self, call, target): |
| 215 | self._schedule_gate(call, ptr_id(target), "t_adj", "t") |
| 216 | |
| 217 | def _on_qis_x(self, call, target): |
| 218 | self._schedule_gate(call, ptr_id(target), "x", "x") |
| 219 | |
| 220 | def _on_qis_y(self, call, target): |
| 221 | self._schedule_gate(call, ptr_id(target), "y", "y") |
| 222 | |
| 223 | def _on_qis_z(self, call, target): |
| 224 | self._schedule_gate(call, ptr_id(target), "z", "z") |
| 225 | |
| 226 | def _on_qis_rx(self, call, angle, target): |
| 227 | self._schedule_rotation(call, ptr_id(target), "rx") |
| 228 | |
| 229 | def _on_qis_rxx(self, call, angle, target1, target2): |
| 230 | self._drop_ops([target1, target2]) |
| 231 | |
| 232 | def _on_qis_ry(self, call, angle, target): |
| 233 | self._schedule_rotation(call, ptr_id(target), "ry") |
| 234 | |
| 235 | def _on_qis_ryy(self, call, angle, target1, target2): |
| 236 | self._drop_ops([target1, target2]) |
| 237 | |
| 238 | def _on_qis_rz(self, call, angle, target): |
| 239 | self._schedule_rotation(call, ptr_id(target), "rz") |
| 240 | |
| 241 | def _on_qis_rzz(self, call, angle, target1, target2): |
| 242 | self._drop_ops([target1, target2]) |
| 243 | |
| 244 | def _on_qis_ccx(self, call, ctrl1, ctrl2, target): |
| 245 | self._drop_ops([ctrl1, ctrl2, target]) |
| 246 | |
| 247 | def _on_qis_cx(self, call, target1, target2): |
| 248 | self._drop_ops([target1, target2]) |
| 249 | |
| 250 | def _on_qis_cy(self, call, target1, target2): |
| 251 | self._drop_ops([target1, target2]) |
| 252 | |
| 253 | def _on_qis_cz(self, call, target1, target2): |
| 254 | self._drop_ops([target1, target2]) |
| 255 | |
| 256 | def _on_qis_swap(self, call, target1, target2): |
| 257 | self._drop_ops([target1, target2]) |
| 258 | |
| 259 | def _on_qis_m(self, call, target, result): |
| 260 | self._drop_ops([target]) |
| 261 | self.last_meas[ptr_id(target)] = (call, target, result) |
| 262 | |
| 263 | def _on_qis_mz(self, call, target, result): |
| 264 | self._on_qis_m(call, target, result) |
| 265 | |
| 266 | def _on_qis_mresetz(self, call, target, result): |
| 267 | self._on_qis_m(call, target, result) |
| 268 | |
| 269 | def _on_qis_reset(self, call, target): |
| 270 | id = ptr_id(target) |
| 271 | if id in self.last_meas: |
| 272 | # Since the last operation on this qubit was a measurement, |
| 273 | # we can combine that measurement with the reset. |
| 274 | (instr, target, result) = self.last_meas.pop(id) |
| 275 | instr.erase() |
| 276 | self.builder.insert_before(call) |
| 277 | new_call = self.builder.call( |
| 278 | self.mresetz_func, |
| 279 | [target, result], |
| 280 | ) |
| 281 | call.erase() |
| 282 | self.last_meas[ptr_id(target)] = (new_call, target, result) |
| 283 | elif not id in self.used_qubits: |
| 284 | # This qubit was never used, so we can just erase the reset instruction. |
| 285 | call.erase() |
| 286 | elif id in self.qubit_ops and self.qubit_ops[id][-1][1] == "reset": |
| 287 | # The last operation on this qubit was also a reset, so we drop the current, |
| 288 | # extra one. |
| 289 | call.erase() |
| 290 | else: |
| 291 | self._drop_ops([target]) |
| 292 | self._schedule_gate(call, id, "reset", "") |
| 293 | |
| 294 | |
| 295 | class PruneUnusedFunctions(QirModuleVisitor): |
| 296 | def _on_module(self, module): |
| 297 | # Assume every non-entry point function is unused. |
| 298 | self.funcs_to_drop = [f for f in module.functions if not is_entry_point(f)] |
| 299 | super()._on_module(module) |
| 300 | # Delete all unused functions. |
| 301 | for func in self.funcs_to_drop: |
| 302 | func.delete() |
| 303 | |
| 304 | def _on_call_instr(self, call): |
| 305 | # Remove calls to initialization and barrier functions, since they aren't handled |
| 306 | # by most of the stack. |
| 307 | if call.callee.name == "__quantum__rt__initialize": |
| 308 | call.erase() |
| 309 | elif call.callee.name == "__quantum__qis__barrier__body": |
| 310 | call.erase() |
| 311 | elif call.callee in self.funcs_to_drop: |
| 312 | # This function is used in a call, so remove it from the list of |
| 313 | # functions to drop. |
| 314 | assert isinstance(call.callee, Function) |
| 315 | self.funcs_to_drop.remove(call.callee) |
| 316 | |