Restructure and move qubit dialect (#557)

As discussed in #549, this restructures the qubit dialect. I also moved
it to the top level of bloqade, since it's used in multiple places. It's
still re-exported by squin, though.

I've got it down to a single failing test, where the `AggressiveUnroll`
fails because the constant propagation doesn't hint the loop iterable as
constant for some reason.

This breaks the syntax such that you have to change the function calls
using lists such as `squin.qubit.measure(q: IList[Qubit])` to
`squin.broadcast.measure(q: IList[Qubit])`.

**Edit**: turns out I was just too eager to remove the type inference
method @weinbe58 added in #508. That resolves the failing test.

**Edit2**: I've also changed `QubitId` and `MeasurementId` to be
"broadcasted" (i.e., they now operate on lists of qubits / measurements
and return lists). This means that now **all statements except
`qubit.new` are broadcasted**, which is as consistent as we can have it,
I think.

---------

Co-authored-by: John Long <[email protected]>

Author: david-pl

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):