Simplyfing qubit.new (#518)

Closes #508 however I am not sure where to reexport `new` since I Can't
reexport it into the `qubit` module. For now I have re-exported it into
`squin.new`.

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: weinbe58

pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
     "scipy>=1.13.1",
-    "kirin-toolchain~=0.17.26",
+    "kirin-toolchain~=0.17.30",
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",

src/bloqade/cirq_utils/emit/base.py

@@ -6,11 +6,11 @@
 from kirin import ir, types, interp
 from kirin.emit import EmitABC, EmitError, EmitFrame
 from kirin.interp import MethodTable, impl
-from kirin.passes import inline
-from kirin.dialects import func
+from kirin.dialects import py, func
 from typing_extensions import Self
 
 from bloqade.squin import kernel
+from bloqade.rewrite.passes import AggressiveUnroll
 
 
 def emit_circuit(
@@ -28,7 +28,7 @@ def emit_circuit(
     Keyword Args:
         circuit_qubits (Sequence[cirq.Qid] | None):
             A list of qubits to use as the qubits in the circuit. Defaults to None.
-            If this is None, then `cirq.LineQubit`s are inserted for every `squin.qubit.new`
+            If this is None, then `cirq.LineQubit`s are inserted for every `squin.qalloc`
             statement in the order they appear inside the kernel.
             **Note**: If a list of qubits is provided, make sure that there is a sufficient
             number of qubits for the resulting circuit.
@@ -48,7 +48,7 @@ def emit_circuit(
 
     @squin.kernel
     def main():
-        q = squin.qubit.new(2)
+        q = squin.qalloc(2)
         squin.h(q[0])
         squin.cx(q[0], q[1])
 
@@ -74,8 +74,10 @@ def entangle(q: ilist.IList[squin.qubit.Qubit, Literal[2]]):
 
     @squin.kernel
     def main():
-        q = squin.qubit.new(2)
-        entangle(q)
+        q = squin.qalloc(2)
+        q2 = squin.qalloc(3)
+        squin.cx(q[1], q2[2])
+
 
     # custom list of qubits on grid
     qubits = [cirq.GridQubit(i, i+1) for i in range(5)]
@@ -112,10 +114,43 @@ def main():
 
     emitter = EmitCirq(qubits=circuit_qubits)
 
-    mt_ = mt.similar(mt.dialects)
-    inline.InlinePass(mt_.dialects).fixpoint(mt_)
+    symbol_op_trait = mt.code.get_trait(ir.SymbolOpInterface)
+    if (symbol_op_trait := mt.code.get_trait(ir.SymbolOpInterface)) is None:
+        raise EmitError("The method is not a symbol, cannot emit circuit!")
+
+    sym_name = symbol_op_trait.get_sym_name(mt.code).unwrap()
+
+    if (signature_trait := mt.code.get_trait(ir.HasSignature)) is None:
+        raise EmitError(
+            f"The method {sym_name} does not have a signature, cannot emit circuit!"
+        )
+
+    signature = signature_trait.get_signature(mt.code)
+    new_signature = func.Signature(inputs=(), output=signature.output)
+
+    callable_region = mt.callable_region.clone()
+    entry_block = callable_region.blocks[0]
+    args_ssa = list(entry_block.args)
+    first_stmt = entry_block.first_stmt
+
+    assert first_stmt is not None, "Method has no statements!"
+    if len(args_ssa) - 1 != len(args):
+        raise EmitError(
+            f"The method {sym_name} takes {len(args_ssa) - 1} arguments, but you passed in {len(args)} via the `args` keyword!"
+        )
+
+    for arg, arg_ssa in zip(args, args_ssa[1:], strict=True):
+        (value := py.Constant(arg)).insert_before(first_stmt)
+        arg_ssa.replace_by(value.result)
+        entry_block.args.delete(arg_ssa)
+
+    new_func = func.Function(
+        sym_name=sym_name, body=callable_region, signature=new_signature
+    )
+    mt_ = ir.Method(None, None, sym_name, [], mt.dialects, new_func)
 
-    return emitter.run(mt_, args=args)
+    AggressiveUnroll(mt_.dialects).fixpoint(mt_)
+    return emitter.run(mt_, args=())
 
 
 @dataclass

src/bloqade/cirq_utils/emit/qubit.py

@@ -10,19 +10,13 @@
 class EmitCirqQubitMethods(MethodTable):
     @impl(qubit.New)
     def new(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: qubit.New):
-        n_qubits = frame.get(stmt.n_qubits)
-
         if frame.qubits is not None:
-            cirq_qubits = tuple(
-                frame.qubits[i + frame.qubit_index] for i in range(n_qubits)
-            )
+            cirq_qubit = frame.qubits[frame.qubit_index]
         else:
-            cirq_qubits = tuple(
-                cirq.LineQubit(i + frame.qubit_index) for i in range(n_qubits)
-            )
+            cirq_qubit = cirq.LineQubit(frame.qubit_index)
 
-        frame.qubit_index += n_qubits
-        return (cirq_qubits,)
+        frame.qubit_index += 1
+        return (cirq_qubit,)
 
     @impl(qubit.MeasureQubit)
     def measure_qubit(

src/bloqade/cirq_utils/lowering.py

@@ -6,7 +6,7 @@
 from kirin.rewrite import Walk, CFGCompactify
 from kirin.dialects import py, scf, func, ilist
 
-from bloqade.squin import gate, noise, qubit, kernel
+from bloqade.squin import gate, noise, qubit, kernel, qalloc
 
 
 def load_circuit(
@@ -92,7 +92,7 @@ def load_circuit(
     @squin.kernel
     def main():
         qreg = get_entangled_qubits()
-        qreg2 = squin.qubit.new(1)
+        qreg2 = squin.qalloc(1)
         entangle_qubits([qreg[1], qreg2[0]])
         return squin.qubit.measure(qreg2)
     ```
@@ -142,7 +142,7 @@ def main():
         body=body,
     )
 
-    return ir.Method(
+    mt = ir.Method(
         mod=None,
         py_func=None,
         sym_name=kernel_name,
@@ -151,6 +151,11 @@ def main():
         code=code,
     )
 
+    assert (run_pass := kernel.run_pass) is not None
+    run_pass(mt, typeinfer=True)
+
+    return mt
+
 
 CirqNode = (
     cirq.Circuit
@@ -254,7 +259,9 @@ def run(
                 # NOTE: create a new register of appropriate size
                 n_qubits = len(self.qreg_index)
                 n = frame.push(py.Constant(n_qubits))
-                self.qreg = frame.push(qubit.New(n_qubits=n.result)).result
+                self.qreg = frame.push(
+                    func.Invoke((n.result,), callee=qalloc, kwargs=())
+                ).result
 
             self.visit(state, stmt)
 
@@ -382,8 +389,16 @@ def bool_op_or(x: bool, y: bool) -> bool:
         # NOTE: remove stmt from parent block
         then_stmt.detach()
         then_body = ir.Block((then_stmt,))
+        then_body.args.append_from(types.Bool, name="cond")
+        then_body.stmts.append(scf.Yield())
 
-        return state.current_frame.push(scf.IfElse(condition, then_body=then_body))
+        else_body = ir.Block(())
+        else_body.args.append_from(types.Bool, name="cond")
+        else_body.stmts.append(scf.Yield())
+
+        return state.current_frame.push(
+            scf.IfElse(condition, then_body=then_body, else_body=else_body)
+        )
 
     def visit_MeasurementGate(
         self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation

src/bloqade/pyqrack/squin/qubit.py

@@ -11,15 +11,12 @@
 @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,)
+    def new_qubit(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New
+    ):
+        (addr,) = interp.memory.allocate(1)
+        qb = PyQrackQubit(addr, interp.memory.sim_reg, QubitState.Active)
+        return (qb,)
 
     def _measure_qubit(self, qbit: PyQrackQubit, interp: PyQrackInterpreter):
         if qbit.is_active():

src/bloqade/rewrite/passes/aggressive_unroll.py

@@ -38,6 +38,7 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             InlineGetField(),
             InlineGetItem(),
             ilist.rewrite.InlineGetItem(),
+            ilist.rewrite.FlattenAdd(),
             ilist.rewrite.HintLen(),
         )
         result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
@@ -68,7 +69,7 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             .rewrite(mt.code)
             .join(result)
         )
-        result = self.typeinfer.unsafe_run(mt).join(result)
+        self.typeinfer.unsafe_run(mt)
         result = self.fold.unsafe_run(mt).join(result)
         result = Walk(Inline(self.inline_heuristic)).rewrite(mt.code).join(result)
         result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)

src/bloqade/squin/__init__.py

@@ -3,9 +3,9 @@
     noise as noise,
     qubit as qubit,
     analysis as analysis,
-    _typeinfer as _typeinfer,
 )
 from .groups import kernel as kernel
+from .stdlib.qubit import qalloc as qalloc
 from .stdlib.simple import (
     h as h,
     s as s,

src/bloqade/squin/_typeinfer.py

@@ -3,7 +3,7 @@
 
 from bloqade.analysis.address.lattice import (
     Address,
-    AddressReg,
+    AddressQubit,
 )
 from bloqade.analysis.address.analysis import AddressAnalysis
 
@@ -27,15 +27,13 @@
 @qubit.dialect.register(key="qubit.address")
 class SquinQubitMethodTable(interp.MethodTable):
 
-    # This can be treated like a QRegNew impl
     @interp.impl(qubit.New)
-    def new(
+    def new_qubit(
         self,
         interp_: AddressAnalysis,
         frame: ForwardFrame[Address],
         stmt: qubit.New,
     ):
-        n_qubits = interp_.get_const_value(int, stmt.n_qubits)
-        addr = AddressReg(range(interp_.next_address, interp_.next_address + n_qubits))
-        interp_.next_address += n_qubits
+        addr = AddressQubit(interp_.next_address)
+        interp_.next_address += 1
         return (addr,)

src/bloqade/squin/analysis/address_impl.py

@@ -20,7 +20,7 @@ def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
             fold_pass.fixpoint(method)
 
         if typeinfer:
-            typeinfer_pass(method)
+            typeinfer_pass(method)  # infer types before desugaring
             desugar_pass.rewrite(method.code)
 
         ilist_desugar_pass(method)