Merge branch 'main' into khwu/stim_annotate

Author: kaihsin

src/bloqade/analysis/measure_id/impls.py

@@ -2,7 +2,7 @@
 from kirin.analysis import const
 from kirin.dialects import py, scf, func, ilist
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
 from .lattice import (
     AnyMeasureId,
@@ -21,22 +21,12 @@
 @qubit.dialect.register(key="measure_id")
 class SquinQubit(interp.MethodTable):
 
-    @interp.impl(qubit.MeasureQubit)
-    def measure_qubit(
-        self,
-        interp: MeasurementIDAnalysis,
-        frame: interp.Frame,
-        stmt: qubit.MeasureQubit,
-    ):
-        interp.measure_count += 1
-        return (MeasureIdBool(interp.measure_count),)
-
-    @interp.impl(qubit.MeasureQubitList)
+    @interp.impl(qubit.stmts.Measure)
     def measure_qubit_list(
         self,
         interp: MeasurementIDAnalysis,
         frame: interp.Frame,
-        stmt: qubit.MeasureQubitList,
+        stmt: qubit.stmts.Measure,
     ):
 
         # try to get the length of the list

src/bloqade/cirq_utils/emit/base.py

@@ -230,3 +230,17 @@ def emit_invoke(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Invoke):
             "Function invokes should need to be inlined! "
             "If you called the emit_circuit method, that should have happened, please report this issue."
         )
+
+    @impl(func.Return)
+    def return_(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: func.Return):
+        # NOTE: should only be hit if ignore_returns == True
+        return ()
+
+
+@py.indexing.dialect.register(key="emit.cirq")
+class __Concrete(interp.MethodTable):
+
+    @interp.impl(py.indexing.GetItem)
+    def getindex(self, interp, frame: interp.Frame, stmt: py.indexing.GetItem):
+        # NOTE: no support for indexing into single statements in cirq
+        return ()

src/bloqade/cirq_utils/emit/qubit.py

@@ -1,7 +1,7 @@
 import cirq
 from kirin.interp import MethodTable, impl
 
-from bloqade.squin import qubit
+from bloqade.qubit import stmts as qubit
 
 from .base import EmitCirq, EmitCirqFrame
 
@@ -18,17 +18,9 @@ def new(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.New):
         frame.qubit_index += 1
         return (cirq_qubit,)
 
-    @impl(qubit.MeasureQubit)
-    def measure_qubit(
-        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.MeasureQubit
-    ):
-        qbit = frame.get(stmt.qubit)
-        frame.circuit.append(cirq.measure(qbit))
-        return (emit.void,)
-
-    @impl(qubit.MeasureQubitList)
+    @impl(qubit.Measure)
     def measure_qubit_list(
-        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.MeasureQubitList
+        self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.Measure
     ):
         qbits = frame.get(stmt.qubits)
         frame.circuit.append(cirq.measure(qbits))

src/bloqade/cirq_utils/lowering.py

@@ -6,7 +6,8 @@
 from kirin.rewrite import Walk, CFGCompactify
 from kirin.dialects import py, scf, func, ilist
 
-from bloqade.squin import gate, noise, qubit, kernel, qalloc
+from bloqade import qubit
+from bloqade.squin import gate, noise, kernel, qalloc
 
 
 def load_circuit(
@@ -403,13 +404,8 @@ def bool_op_or(x: bool, y: bool) -> bool:
     def visit_MeasurementGate(
         self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
     ):
-        cirq_qubits = node.qubits
-        if len(cirq_qubits) == 1:
-            qbit = self.lower_qubit_getindex(state, node.qubits[0])
-            stmt = state.current_frame.push(qubit.MeasureQubit(qbit))
-        else:
-            qubits = self.lower_qubit_getindices(state, node.qubits)
-            stmt = state.current_frame.push(qubit.MeasureQubitList(qubits))
+        qubits = self.lower_qubit_getindices(state, node.qubits)
+        stmt = state.current_frame.push(qubit.stmts.Measure(qubits))
 
         # NOTE: add for classically controlled lowering
         key = node.gate.key

src/bloqade/native/_prelude.py

@@ -5,7 +5,7 @@
 from kirin.prelude import structural_no_opt
 from typing_extensions import Doc
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
 from .dialects import gates
 

src/bloqade/native/dialects/gates/_interface.py

@@ -3,7 +3,7 @@
 from kirin import lowering
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
 from .stmts import CZ, R, Rz
 

src/bloqade/native/dialects/gates/stmts.py

@@ -2,7 +2,7 @@
 from kirin.decl import info, statement
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
+from bloqade.types import QubitType
 
 from ._dialect import dialect
 
@@ -12,20 +12,20 @@
 @statement(dialect=dialect)
 class CZ(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
-    ctrls: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, N])
-    qargs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, N])
+    ctrls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
 
 
 @statement(dialect=dialect)
 class R(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
-    inputs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, types.Any])
+    inputs: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
     axis_angle: ir.SSAValue = info.argument(types.Float)
     rotation_angle: ir.SSAValue = info.argument(types.Float)
 
 
 @statement(dialect=dialect)
 class Rz(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
-    inputs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, types.Any])
+    inputs: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
     rotation_angle: ir.SSAValue = info.argument(types.Float)

src/bloqade/native/stdlib/broadcast.py

@@ -3,7 +3,7 @@
 
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
+from bloqade import qubit
 from bloqade.native._prelude import kernel
 from bloqade.native.dialects.gates import _interface as native
 

src/bloqade/native/stdlib/simple.py

@@ -1,6 +1,6 @@
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
 from . import broadcast
 from .._prelude import kernel

src/bloqade/pyqrack/squin/qubit.py

@@ -3,7 +3,7 @@
 from kirin import interp
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
+from bloqade.qubit import stmts as qubit
 from bloqade.pyqrack.reg import QubitState, Measurement, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
@@ -27,38 +27,32 @@ def _measure_qubit(self, qbit: PyQrackQubit, interp: PyQrackInterpreter):
         interp.set_global_measurement_id(m)
         return m
 
-    @interp.impl(qubit.MeasureQubitList)
+    @interp.impl(qubit.Measure)
     def measure_qubit_list(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: qubit.MeasureQubitList,
+        stmt: qubit.Measure,
     ):
         qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
         result = ilist.IList([self._measure_qubit(qbit, interp) for qbit in qubits])
         return (result,)
 
-    @interp.impl(qubit.MeasureQubit)
-    def measure_qubit(
-        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasureQubit
-    ):
-        qbit: PyQrackQubit = frame.get(stmt.qubit)
-        result = self._measure_qubit(qbit, interp)
-        return (result,)
-
     @interp.impl(qubit.QubitId)
     def qubit_id(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.QubitId
     ):
-        qbit: PyQrackQubit = frame.get(stmt.qubit)
-        return (qbit.addr,)
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        ids = ilist.IList([qbit.addr for qbit in qubits])
+        return (ids,)
 
     @interp.impl(qubit.MeasurementId)
     def measurement_id(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasurementId
     ):
-        measurement: Measurement = frame.get(stmt.measurement)
-        return (measurement.measurement_id,)
+        measurements: ilist.IList[Measurement, Any] = frame.get(stmt.measurements)
+        ids = ilist.IList([measurement.measurement_id for measurement in measurements])
+        return (ids,)
 
     @interp.impl(qubit.Reset)
     def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Reset):