Change operator runtime implementation

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -5,7 +5,7 @@
 # from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
-from .runtime import IdentityRuntime, OperatorRuntime, ProjectorRuntime
+from .runtime import ControlRuntime, IdentityRuntime, OperatorRuntime, ProjectorRuntime
 
 # from kirin.dialects import ilist
 
@@ -49,10 +49,9 @@ def control(
         op = frame.get(stmt.op)
         n_controls = stmt.n_controls
         # FIXME: the method name here is dirty
-        rt = OperatorRuntime(
+        rt = ControlRuntime(
             method_name="mc" + op.method_name,
-            target_index=n_controls,
-            ctrl_index=list(range(n_controls)),
+            n_controls=n_controls,
         )
         return (rt,)
 
@@ -66,7 +65,7 @@ def control(
     def identity(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Identity
     ):
-        return (IdentityRuntime(target_index=0, sites=stmt.sites),)
+        return (IdentityRuntime(sites=stmt.sites),)
 
     # @interp.impl(op.stmts.PhaseOp)
     # def phaseop(
@@ -112,7 +111,7 @@ def operator(
             op.stmts.X | op.stmts.Y | op.stmts.Z | op.stmts.H | op.stmts.S | op.stmts.T
         ),
     ):
-        return (OperatorRuntime(method_name=stmt.name.lower(), target_index=0),)
+        return (OperatorRuntime(method_name=stmt.name.lower()),)
 
     @interp.impl(op.stmts.P0)
     @interp.impl(op.stmts.P1)
@@ -123,7 +122,7 @@ def projector(
         stmt: op.stmts.P0 | op.stmts.P1,
     ):
         state = isinstance(stmt, op.stmts.P1)
-        return (ProjectorRuntime(to_state=state, target_index=0),)
+        return (ProjectorRuntime(to_state=state),)
 
     @interp.impl(op.stmts.Sn)
     def sn(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Sn):

src/bloqade/pyqrack/squin/qubit.py

@@ -27,7 +27,7 @@ def new(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New):
     def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
         operator: OperatorRuntimeABC = frame.get(stmt.operator)
         qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
-        operator.apply(qubits=qubits)
+        operator.apply(*qubits)
 
     @interp.impl(qubit.Measure)
     def measure(

src/bloqade/pyqrack/squin/runtime.py

@@ -1,16 +1,11 @@
-from typing import Any, Optional
 from dataclasses import dataclass
 
-from kirin.dialects import ilist
-
 from bloqade.pyqrack import PyQrackQubit
 
 
 @dataclass
 class OperatorRuntimeABC:
-    target_index: int
-
-    def apply(self, qubits: ilist.IList[PyQrackQubit, Any]) -> None:
+    def apply(self, *qubits: PyQrackQubit) -> None:
         raise NotImplementedError(
             "Operator runtime base class should not be called directly, override the method"
         )
@@ -19,18 +14,27 @@ def apply(self, qubits: ilist.IList[PyQrackQubit, Any]) -> None:
 @dataclass
 class OperatorRuntime(OperatorRuntimeABC):
     method_name: str
-    ctrl_index: Optional[list[int]] = None
 
     def apply(
         self,
-        qubits: ilist.IList[PyQrackQubit, Any],
-    ):
-        target_qubit = qubits[self.target_index]
-        if self.ctrl_index is not None:
-            ctrls = [qubits[i].addr for i in self.ctrl_index]
-            getattr(target_qubit.sim_reg, self.method_name)(ctrls, target_qubit.addr)
-        else:
-            getattr(target_qubit.sim_reg, self.method_name)(target_qubit.addr)
+        *qubits: PyQrackQubit,
+    ) -> None:
+        getattr(qubits[-1].sim_reg, self.method_name)(qubits[-1].addr)
+
+
+@dataclass
+class ControlRuntime(OperatorRuntimeABC):
+    method_name: str
+    n_controls: int
+
+    def apply(
+        self,
+        *qubits: PyQrackQubit,
+    ) -> 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]]
+        target = qubits[-1]
+        getattr(target.sim_reg, self.method_name)(ctrls, target.addr)
 
 
 @dataclass
@@ -39,16 +43,15 @@ class ProjectorRuntime(OperatorRuntimeABC):
 
     def apply(
         self,
-        qubits: ilist.IList[PyQrackQubit, Any],
-    ):
-        target_qubit = qubits[self.target_index]
-        target_qubit.sim_reg.force_m(target_qubit.addr, self.to_state)
+        *qubits: PyQrackQubit,
+    ) -> None:
+        qubits[-1].sim_reg.force_m(qubits[-1].addr, self.to_state)
 
 
 @dataclass
 class IdentityRuntime(OperatorRuntimeABC):
     # TODO: do we even need sites? The apply never does anything
     sites: int
 
-    def apply(self, qubits: ilist.IList[PyQrackQubit, Any]):
+    def apply(self, *qubits: PyQrackQubit) -> None:
         pass