Rework impl of control to work with all supported operators

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -7,6 +7,7 @@
     KronRuntime,
     MultRuntime,
     ScaleRuntime,
+    AdjointRuntime,
     ControlRuntime,
     IdentityRuntime,
     OperatorRuntime,
@@ -33,13 +34,12 @@ def mult(
         rhs = frame.get(stmt.rhs)
         return (MultRuntime(lhs, rhs),)
 
-    # @interp.impl(op.stmts.Adjoint)
-    # def adjoint(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Adjoint
-    # ):
-    #     is_unitary: bool = info.attribute(default=False)
-    #     op: ir.SSAValue = info.argument(OpType)
-    #     result: ir.ResultValue = info.result(OpType)
+    @interp.impl(op.stmts.Adjoint)
+    def adjoint(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Adjoint
+    ):
+        op = frame.get(stmt.op)
+        return (AdjointRuntime(op),)
 
     @interp.impl(op.stmts.Scale)
     def scale(
@@ -55,9 +55,8 @@ def control(
     ):
         op = frame.get(stmt.op)
         n_controls = stmt.n_controls
-        # FIXME: the method name here is dirty
         rt = ControlRuntime(
-            method_name="mc" + op.method_name,
+            op=op,
             n_controls=n_controls,
         )
         return (rt,)

src/bloqade/pyqrack/squin/runtime.py

@@ -1,50 +1,76 @@
 from dataclasses import dataclass
 
+import numpy as np
+
 from bloqade.pyqrack import PyQrackQubit
 
 
 @dataclass
 class OperatorRuntimeABC:
-    def apply(self, *qubits: PyQrackQubit) -> None:
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         raise NotImplementedError(
             "Operator runtime base class should not be called directly, override the method"
         )
 
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        raise NotImplementedError(f"Can't apply controlled version of {self}")
+
 
 @dataclass
 class OperatorRuntime(OperatorRuntimeABC):
     method_name: str
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
-        getattr(qubits[0].sim_reg, self.method_name)(qubits[0].addr)
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        method_name = self.method_name
+        if adjoint:
+            method_name = "adj" + method_name
+        getattr(qubits[0].sim_reg, method_name)(qubits[0].addr)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        target = qubits[-1]
+        method_name = "mc"
+        if adjoint:
+            method_name += "adj"
+        method_name += self.method_name
+        getattr(target.sim_reg, method_name)(target.addr, ctrls)
 
 
 @dataclass
 class ControlRuntime(OperatorRuntimeABC):
-    method_name: str
+    op: OperatorRuntimeABC
     n_controls: int
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         # NOTE: this is a bit odd, since you can "skip" qubits by making n_controls < len(qubits)
-        ctrls = [qbit.addr for qbit in qubits[: self.n_controls]]
+        ctrls = qubits[: self.n_controls]
         target = qubits[-1]
-        getattr(target.sim_reg, self.method_name)(ctrls, target.addr)
+        self.op.control_apply(target, *ctrls, adjoint=adjoint)
 
 
 @dataclass
 class ProjectorRuntime(OperatorRuntimeABC):
     to_state: bool
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
-        qubits[-1].sim_reg.force_m(qubits[-1].addr, self.to_state)
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        qubits[0].sim_reg.force_m(qubits[0].addr, self.to_state)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        m = [not self.to_state, 0, 0, self.to_state]
+        target = qubits[-1]
+        ctrls = [qbit.addr for qbit in qubits[:-1]]
+        target.sim_reg.mcmtrx(ctrls, m, target.addr)
 
 
 @dataclass
 class IdentityRuntime(OperatorRuntimeABC):
     # TODO: do we even need sites? The apply never does anything
     sites: int
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        pass
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
         pass
 
 
@@ -53,32 +79,73 @@ class MultRuntime(OperatorRuntimeABC):
     lhs: OperatorRuntimeABC
     rhs: OperatorRuntimeABC
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
-        self.rhs.apply(*qubits)
-        self.lhs.apply(*qubits)
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        if adjoint:
+            # NOTE: inverted order
+            self.lhs.apply(*qubits, adjoint=adjoint)
+            self.rhs.apply(*qubits, adjoint=adjoint)
+        else:
+            self.rhs.apply(*qubits)
+            self.lhs.apply(*qubits)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        if adjoint:
+            self.lhs.control_apply(*qubits, adjoint=adjoint)
+            self.rhs.control_apply(*qubits, adjoint=adjoint)
+        else:
+            self.rhs.control_apply(*qubits, adjoint=adjoint)
+            self.lhs.control_apply(*qubits, adjoint=adjoint)
 
 
 @dataclass
 class KronRuntime(OperatorRuntimeABC):
     lhs: OperatorRuntimeABC
     rhs: OperatorRuntimeABC
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
-        self.lhs.apply(qubits[0])
-        self.rhs.apply(qubits[1])
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        self.lhs.apply(qubits[0], adjoint=adjoint)
+        self.rhs.apply(qubits[1], adjoint=adjoint)
 
 
 @dataclass
 class ScaleRuntime(OperatorRuntimeABC):
     op: OperatorRuntimeABC
     factor: complex
 
-    def apply(self, *qubits: PyQrackQubit) -> None:
-        target = qubits[0]
-        self.op.apply(target)
+    def mat(self, adjoint: bool):
+        if adjoint:
+            return [np.conj(self.factor), 0, 0, self.factor]
+        else:
+            return [self.factor, 0, 0, self.factor]
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        self.op.apply(*qubits, adjoint=adjoint)
+
+        target = qubits[-1]
 
         # NOTE: just factor * eye(2)
-        mat = [self.factor, 0, 0, self.factor]
+        m = self.mat(adjoint)
 
         # TODO: output seems to always be normalized -- no-op?
-        target.sim_reg.mtrx(mat, target.addr)
+        target.sim_reg.mtrx(m, target.addr)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
+        self.op.control_apply(*qubits, adjoint=adjoint)
+
+        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
+class AdjointRuntime(OperatorRuntimeABC):
+    op: OperatorRuntimeABC
+
+    def apply(self, *qubits: PyQrackQubit, adjoint: bool = True) -> None:
+        self.op.apply(*qubits, adjoint=adjoint)
+
+    def control_apply(self, *qubits: PyQrackQubit, adjoint: bool = True) -> None:
+        self.op.control_apply(*qubits, adjoint=adjoint)