Runtime for phase and shift operators

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -4,11 +4,13 @@
 from bloqade.pyqrack.base import PyQrackInterpreter
 
 from .runtime import (
+    RotRuntime,
     KronRuntime,
     MultRuntime,
     ScaleRuntime,
     AdjointRuntime,
     ControlRuntime,
+    PhaseOpRuntime,
     IdentityRuntime,
     OperatorRuntime,
     ProjectorRuntime,
@@ -61,47 +63,29 @@ def control(
         )
         return (rt,)
 
-    # @interp.impl(op.stmts.Rot)
-    # def rot(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Rot):
-    #     axis: ir.SSAValue = info.argument(OpType)
-    #     angle: ir.SSAValue = info.argument(types.Float)
-    #     result: ir.ResultValue = info.result(OpType)
+    @interp.impl(op.stmts.Rot)
+    def rot(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Rot):
+        axis = frame.get(stmt.axis)
+        angle = frame.get(stmt.angle)
+        return (RotRuntime(axis, angle),)
 
     @interp.impl(op.stmts.Identity)
     def identity(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Identity
     ):
         return (IdentityRuntime(sites=stmt.sites),)
 
-    # @interp.impl(op.stmts.PhaseOp)
-    # def phaseop(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.PhaseOp
-    # ):
-    #     """
-    #     A phase operator.
-
-    #     $$
-    #     PhaseOp(theta) = e^{i \theta} I
-    #     $$
-    #     """
-
-    #     theta: ir.SSAValue = info.argument(types.Float)
-    #     result: ir.ResultValue = info.result(OpType)
-
-    # @interp.impl(op.stmts.ShiftOp)
-    # def shiftop(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.ShiftOp
-    # ):
-    #     """
-    #     A phase shift operator.
-
-    #     $$
-    #     Shift(theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
-    #     $$
-    #     """
-
-    #     theta: ir.SSAValue = info.argument(types.Float)
-    #     result: ir.ResultValue = info.result(OpType)
+    @interp.impl(op.stmts.PhaseOp)
+    @interp.impl(op.stmts.ShiftOp)
+    def phaseop(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.PhaseOp | op.stmts.ShiftOp,
+    ):
+        theta = frame.get(stmt.theta)
+        global_ = isinstance(stmt, op.stmts.PhaseOp)
+        return (PhaseOpRuntime(theta, global_=global_),)
 
     @interp.impl(op.stmts.X)
     @interp.impl(op.stmts.Y)

src/bloqade/pyqrack/squin/runtime.py

@@ -5,7 +5,7 @@
 from bloqade.pyqrack import PyQrackQubit
 
 
-@dataclass
+@dataclass(frozen=True)
 class OperatorRuntimeABC:
     def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         raise NotImplementedError(
@@ -16,7 +16,7 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         raise NotImplementedError(f"Can't apply controlled version of {self}")
 
 
-@dataclass
+@dataclass(frozen=True)
 class OperatorRuntime(OperatorRuntimeABC):
     method_name: str
 
@@ -36,7 +36,7 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         getattr(target.sim_reg, method_name)(target.addr, ctrls)
 
 
-@dataclass
+@dataclass(frozen=True)
 class ControlRuntime(OperatorRuntimeABC):
     op: OperatorRuntimeABC
     n_controls: int
@@ -48,7 +48,7 @@ def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         self.op.control_apply(target, *ctrls, adjoint=adjoint)
 
 
-@dataclass
+@dataclass(frozen=True)
 class ProjectorRuntime(OperatorRuntimeABC):
     to_state: bool
 
@@ -62,7 +62,7 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
-@dataclass
+@dataclass(frozen=True)
 class IdentityRuntime(OperatorRuntimeABC):
     # TODO: do we even need sites? The apply never does anything
     sites: int
@@ -74,7 +74,7 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         pass
 
 
-@dataclass
+@dataclass(frozen=True)
 class MultRuntime(OperatorRuntimeABC):
     lhs: OperatorRuntimeABC
     rhs: OperatorRuntimeABC
@@ -97,7 +97,7 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
             self.lhs.control_apply(*qubits, adjoint=adjoint)
 
 
-@dataclass
+@dataclass(frozen=True)
 class KronRuntime(OperatorRuntimeABC):
     lhs: OperatorRuntimeABC
     rhs: OperatorRuntimeABC
@@ -107,7 +107,7 @@ def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         self.rhs.apply(qubits[1], adjoint=adjoint)
 
 
-@dataclass
+@dataclass(frozen=True)
 class ScaleRuntime(OperatorRuntimeABC):
     op: OperatorRuntimeABC
     factor: complex
@@ -140,7 +140,37 @@ def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
-@dataclass
+@dataclass(frozen=True)
+class PhaseOpRuntime(OperatorRuntimeABC):
+    theta: float
+    global_: bool
+
+    def mat(self, adjoint: bool):
+        sign = (-1) ** (not adjoint)
+        phase = np.exp(sign * 1j * self.theta)
+        return [self.global_ * phase, 0, 0, phase]
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        target = qubits[-1]
+        target.sim_reg.mtrx(self.mat(adjoint=adjoint), target.addr)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        target = qubits[-1]
+        ctrls = [qbit.addr for qbit in qubits[:-1]]
+
+        m = self.mat(adjoint=adjoint)
+
+        target.sim_reg.mcmtrx(ctrls, m, target.addr)
+
+
+@dataclass(frozen=True)
+class RotRuntime(OperatorRuntimeABC):
+    axis: OperatorRuntimeABC
+    angle: float
+    # TODO: how does this work?
+
+
+@dataclass(frozen=True)
 class AdjointRuntime(OperatorRuntimeABC):
     op: OperatorRuntimeABC
 

test/pyqrack/test_squin.py

@@ -150,11 +150,32 @@ def main():
     assert result == [1]
 
 
+def test_phase():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        h = squin.op.h()
+        squin.qubit.apply(h, q)
+
+        # rotate local phase by pi/2
+        p = squin.op.shift(math.pi / 2)
+        squin.qubit.apply(p, q)
+
+        # the next hadamard should rotate it back to 0
+        squin.qubit.apply(h, q)
+        return squin.qubit.measure(q)
+
+    target = PyQrack(1)
+    result = target.run(main)
+    assert result == [0]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
 # test_basic_ops("x")
 # test_cx()
 # test_mult()
 # test_kron()
-# test_scale()
+test_scale()
+test_phase()