Implement methods for qubit statements

Author: david-pl

src/bloqade/pyqrack/__init__.py

@@ -14,4 +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, glob as glob, parallel as parallel
+from .squin import qubit as qubit
 from .target import PyQrack as PyQrack

src/bloqade/pyqrack/squin/qubit.py

@@ -0,0 +1,74 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.squin import qubit
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+
+
+@qubit.dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
+    @interp.impl(qubit.New)
+    def new(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New):
+        n_qubits: int = frame.get(stmt.n_qubits)
+        qreg = ilist.IList(
+            [
+                PyQrackQubit(i, interp.memory.sim_reg, QubitState.Active)
+                for i in interp.memory.allocate(n_qubits=n_qubits)
+            ]
+        )
+        return (qreg,)
+
+    @interp.impl(qubit.Apply)
+    def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
+        # TODO
+        # operator: ir.SSAValue = info.argument(OpType)
+        # qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
+        pass
+
+    @interp.impl(qubit.Measure)
+    def measure(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Measure
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        result = [qbit.sim_reg.m(qbit.addr) for qbit in qubits]
+        return (result,)
+
+    @interp.impl(qubit.MeasureAndReset)
+    def measure_and_reset(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: qubit.MeasureAndReset,
+    ):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        result = [qbit.sim_reg.m(qbit.addr) for qbit in qubits]
+        for qbit in qubits:
+            qbit.sim_reg.force_m(qbit.addr, 0)
+
+        return (result,)
+
+    @interp.impl(qubit.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Reset):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            qbit.sim_reg.force_m(qbit.addr, 0)
+
+    # @interp.impl(glob.UGate)
+    # def ugate(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: glob.UGate):
+    #     registers: ilist.IList[ilist.IList[PyQrackQubit, Any], 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 ()

test/pyqrack/test_squin.py

@@ -0,0 +1,65 @@
+from kirin.dialects import ilist
+
+from bloqade import squin
+from bloqade.pyqrack import PyQrack, PyQrackQubit
+
+
+def test_qubit():
+    @squin.kernel
+    def new():
+        return squin.qubit.new(3)
+
+    new.print()
+
+    target = PyQrack(3)
+    result = target.run(new)
+    assert isinstance(result, ilist.IList)
+    assert isinstance(qubit := result[0], PyQrackQubit)
+
+    out = qubit.sim_reg.out_ket()
+    assert out == [1.0] + [0.0] * (2**3 - 1)
+
+    @squin.kernel
+    def measure():
+        q = squin.qubit.new(3)
+        m = squin.qubit.measure(q)
+        squin.qubit.reset(q)
+        return m
+
+    target = PyQrack(3)
+    result = target.run(measure)
+    assert isinstance(result, list)
+    assert result == [0, 0, 0]
+
+    @squin.kernel
+    def measure_and_reset():
+        q = squin.qubit.new(3)
+        m = squin.qubit.measure_and_reset(q)
+        return m
+
+    target = PyQrack(3)
+    result = target.run(measure_and_reset)
+    assert isinstance(result, list)
+    assert result == [0, 0, 0]
+
+
+# @squin.kernel
+# def main():
+#     q = squin.qubit.new(3)
+#     x = squin.op.x()
+#     id = squin.op.identity(sites=2)
+
+#     # FIXME? Should we have a method apply(x, q, idx)?
+#     squin.qubit.apply(squin.op.kron(x, id), q)
+
+#     return squin.qubit.measure(q)
+
+
+# main.print()
+
+# target = PyQrack(2)
+# result = target.run(main)
+
+
+if __name__ == "main":
+    test_qubit()