Implement pyqrack interpreter method for qasm2 glob.u gate (#152)

* Reorder UGate arguments to match wrapper function arg order

The wrapper was already in use in a bunch of places with this order,
but using named arguments so it never errored. Still this should be
consistent and I stuck with the order of the wrapper as it is already used.

* Make PyQrackReg iterable by throwing IndexError in getitem method

* Implement pyqrack interpreter method for glob.u gate

* Fix pyright issue in core & reg

* Fix order of glob.u decorator yet again (to match original UGate class args)

* Move new method to separate method table

* Add TODO to clean up PyQrackReg

Author: david-pl

src/bloqade/pyqrack/__init__.py

@@ -14,5 +14,5 @@
 
 # NOTE: The following import is for registering the method tables
 from .noise import native as native
-from .qasm2 import uop as uop, core as core, parallel as parallel
+from .qasm2 import uop as uop, core as core, glob as glob, parallel as parallel
 from .target import PyQrack as PyQrack

src/bloqade/pyqrack/qasm2/core.py

@@ -14,7 +14,6 @@
 
 @core.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
-
     @interp.impl(core.QRegNew)
     def qreg_new(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.QRegNew

src/bloqade/pyqrack/qasm2/glob.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.pyqrack.reg import PyQrackReg
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.qasm2.dialects import glob
+
+
+@glob.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    @interp.impl(glob.UGate)
+    def ugate(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: glob.UGate):
+        registers: ilist.IList[PyQrackReg, Any] = frame.get(stmt.registers)
+        theta, phi, lam = (
+            frame.get(stmt.theta),
+            frame.get(stmt.phi),
+            frame.get(stmt.lam),
+        )
+
+        for qreg in registers:
+            for qarg in qreg:
+                if qarg.is_active():
+                    interp.memory.sim_reg.u(qarg.addr, theta, phi, lam)
+        return ()

src/bloqade/pyqrack/reg.py

@@ -33,7 +33,7 @@ class CBitRef:
     pos: int
     """The position of this bit in the classical register."""
 
-    def set_value(self, value: bool):
+    def set_value(self, value: Measurement):
         self.ref[self.pos] = value
 
     def get_value(self):
@@ -46,7 +46,7 @@ class QubitState(enum.Enum):
 
 
 @dataclass(frozen=True)
-class PyQrackReg(QReg):
+class PyQrackReg(QReg):  # TODO: clean up implementation with list base class
     """Simulation runtime value of a quantum register."""
 
     size: int
@@ -72,6 +72,8 @@ def drop(self, pos: int):
         self.qubit_state[pos] = QubitState.Lost
 
     def __getitem__(self, pos: int):
+        if not 0 <= pos < self.size:
+            raise IndexError("Qubit index out of bounds of register.")
         return PyQrackQubit(self, pos)
 
 

src/bloqade/qasm2/glob.py

@@ -11,7 +11,7 @@
 
 @wraps(glob.UGate)
 def u(
-    theta: float, phi: float, lam: float, registers: ilist.IList[QReg, Any] | list
+    registers: ilist.IList[QReg, Any] | list, theta: float, phi: float, lam: float
 ) -> None:
     """Apply a U gate to all qubits in the input registers.
 

test/pyqrack/test_target.py

@@ -38,3 +38,77 @@ def ghz():
     assert math.isclose(out[0].real, val, abs_tol=abs_tol)
     assert math.isclose(out[-1].real, val, abs_tol=abs_tol)
     assert all(math.isclose(ele.real, 0.0, abs_tol=abs_tol) for ele in out[1:-1])
+
+
+def test_target_glob():
+    @qasm2.extended
+    def global_h():
+        q = qasm2.qreg(3)
+
+        # rotate around Y by pi/2, i.e. perform a hadamard
+        qasm2.glob.u([q], math.pi / 2.0, 0, 0)
+
+        return q
+
+    target = PyQrack(3)
+    q = target.run(global_h)
+
+    assert isinstance(q, reg.PyQrackReg)
+
+    out = q.sim_reg.out_ket()
+
+    # remove global phase introduced by pyqrack
+    phase = out[0] / abs(out[0])
+    out = [ele / phase for ele in out]
+
+    for element in out:
+        assert math.isclose(element.real, 1 / math.sqrt(8), abs_tol=2.2e-7)
+        assert math.isclose(element.imag, 0, abs_tol=2.2e-7)
+
+    @qasm2.extended
+    def multiple_registers():
+        q1 = qasm2.qreg(2)
+        q2 = qasm2.qreg(2)
+        q3 = qasm2.qreg(2)
+
+        # hadamard on first register
+        qasm2.glob.u(
+            [q1],
+            math.pi / 2.0,
+            0,
+            0,
+        )
+
+        # apply hadamard to the other two
+        qasm2.glob.u(
+            [q2, q3],
+            math.pi / 2.0,
+            0,
+            0,
+        )
+
+        # rotate all of them back down
+        qasm2.glob.u(
+            [q1, q2, q3],
+            -math.pi / 2.0,
+            0,
+            0,
+        )
+
+        return q1
+
+    target = PyQrack(6)
+    q1 = target.run(multiple_registers)
+
+    assert isinstance(q1, reg.PyQrackReg)
+
+    out = q1.sim_reg.out_ket()
+
+    assert out[0] == 1
+    for i in range(1, len(out)):
+        assert out[i] == 0
+
+    assert True
+
+
+test_target_glob()

test/qasm2/passes/test_heuristic_noise.py

@@ -265,7 +265,7 @@ def test_global_noise():
     def test_method():
         q0 = qasm2.qreg(1)
         q1 = qasm2.qreg(1)
-        glob.UGate([q0, q1], 0.1, 0.2, 0.3)
+        qasm2.glob.u([q0, q1], 0.1, 0.2, 0.3)
 
     px = 0.01
     py = 0.01