support for control gates confirmed

Author: johnzl-777

src/bloqade/squin/analysis/nsites/impls.py

@@ -2,7 +2,7 @@
 
 from kirin import ir, interp
 
-from bloqade.squin import op
+from bloqade.squin import op, wire
 
 from .lattice import (
     NoSites,
@@ -11,6 +11,15 @@
 from .analysis import NSitesAnalysis
 
 
+@wire.dialect.register(key="op.nsites")
+class SquinWire(interp.MethodTable):
+
+    @interp.impl(wire.Apply)
+    def apply(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: wire.Apply):
+
+        return tuple([frame.get(input) for input in stmt.inputs])
+
+
 @op.dialect.register(key="op.nsites")
 class SquinOp(interp.MethodTable):
 

src/bloqade/squin/rewrite/stim.py

@@ -102,9 +102,46 @@ def get_stim_1q_gate(self, squin_op: op.stmts.Operator):
                 return stim.gate.H
             case op.stmts.S():
                 return stim.gate.S
+            case op.stmts.Identity():  # enforce sites defined = num wires in
+                return stim.gate.Identity
             case _:
                 return None
 
+    # get the qubit indices from the Apply statement argument
+    # wires/qubits
+    def insert_qubit_idx_ssa(
+        self, apply_stmt: wire.Apply | qubit.Apply
+    ) -> tuple[ir.SSAValue, ...]:
+
+        if isinstance(apply_stmt, qubit.Apply):
+            qubits = apply_stmt.qubits
+            address_attribute: AddressAttribute = self.get_address(qubits)
+            # Should get an AddressTuple out of the address stored in attribute
+            address_tuple = address_attribute.address
+            qubit_idx_ssas: list[ir.SSAValue] = []
+            for address_qubit in address_tuple.data:
+                qubit_idx = address_qubit.data
+                qubit_idx_stmt = py.Constant(qubit_idx)
+                qubit_idx_stmt.insert_before(apply_stmt)
+                qubit_idx_ssas.append(qubit_idx_stmt.result)
+
+            return tuple(qubit_idx_ssas)
+
+        elif isinstance(apply_stmt, wire.Apply):
+            wire_ssas = apply_stmt.inputs
+            qubit_idx_ssas: list[ir.SSAValue] = []
+            for wire_ssa in wire_ssas:
+                address_attribute = self.get_address(wire_ssa)
+                # get parent qubit idx
+                wire_address = address_attribute.address
+                qubit_idx = wire_address.origin_qubit.data
+                qubit_idx_stmt = py.Constant(qubit_idx)
+                # accumulate all qubit idx SSA to instantiate stim gate stmt
+                qubit_idx_ssas.append(qubit_idx_stmt.result)
+                qubit_idx_stmt.insert_before(apply_stmt)
+
+            return tuple(qubit_idx_ssas)
+
     # might be worth attempting multiple dispatch like qasm2 rewrites
     # for Glob and Parallel to UOp
     # The problem is I'd have to introduce names for all the statements
@@ -142,51 +179,60 @@ def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
         # this is an SSAValue, need it to be the actual operator
         applied_op = apply_stmt.operator.owner
 
-        # need to handle Identity and Control through separate means
-        # but we can handle X, Y, Z, and H here just fine
-        stim_1q_op = self.get_stim_1q_gate(applied_op)
-
-        if isinstance(apply_stmt, qubit.Apply):
-            qubits = apply_stmt.qubits
-            address_attribute: AddressAttribute = self.get_address(qubits)
-            # Should get an AddressTuple out of the address stored in attribute
-            address_tuple = address_attribute.address
-            qubit_idx_ssas: list[ir.SSAValue] = []
-            for address_qubit in address_tuple.data:
-                qubit_idx = address_qubit.data
-                qubit_idx_stmt = py.Constant(qubit_idx)
-                qubit_idx_ssas.append(qubit_idx_stmt.result)
-                qubit_idx_stmt.insert_before(apply_stmt)
-
-            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
-
-            # can't do any of this because of dependencies downstream
-            # apply_stmt.replace_by(stim_1q_stmt)
-
-            return RewriteResult(has_done_something=True)
+        if isinstance(applied_op, op.stmts.Control):
+            return self.rewrite_Control(apply_stmt)
 
-        elif isinstance(apply_stmt, wire.Apply):
-            wires_ssa = apply_stmt.inputs
-            qubit_idx_ssas: list[ir.SSAValue] = []
-            for wire_ssa in wires_ssa:
-                address_attribute = self.get_address(wire_ssa)
-                # get parent qubit idx
-                wire_address = address_attribute.address
-                qubit_idx = wire_address.origin_qubit.data
-                qubit_idx_stmt = py.Constant(qubit_idx)
-                # accumulate all qubit idx SSA to instantiate stim gate stmt
-                qubit_idx_ssas.append(qubit_idx_stmt.result)
-                qubit_idx_stmt.insert_before(apply_stmt)
+        # need to handle Control through separate means
+        # but we can handle X, Y, Z, H, and S here just fine
+        stim_1q_op = self.get_stim_1q_gate(applied_op)
 
-            stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
-            stim_1q_stmt.insert_before(apply_stmt)
+        qubit_idx_ssas = self.insert_qubit_idx_ssa(apply_stmt=apply_stmt)
+        stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
+        stim_1q_stmt.insert_before(apply_stmt)
 
-            # There is something depending on the results of the statement,
-            # need to handle that so replacement/deletion can occur without problems
+        return RewriteResult(has_done_something=True)
 
-            # apply's results become wires that go to other apply's/wrap stmts
-            # apply_stmt.replace_by(stim_1q_stmt)
+    def rewrite_Control(
+        self, apply_stmt_ctrl: qubit.Apply | wire.Apply
+    ) -> RewriteResult:
+        # stim only supports CX, CY, CZ so we have to check the
+        # operator of Apply is a Control gate, enforce it's only asking for 1 control qubit,
+        # and that the target of the control is X, Y, Z in squin
+
+        ctrl_op: op.stmts.Control = apply_stmt_ctrl.operator.owner
+        # enforce that n_controls is 1
+
+        ctrl_op_target_gate = ctrl_op.op.owner
+
+        # should enforce that this is some multiple of 2
+        qubit_idx_ssas = self.insert_qubit_idx_ssa(apply_stmt=apply_stmt_ctrl)
+        # according to stim, final result can be:
+        # CX 1 2 3 4 -> CX(1, targ=2), CX(3, targ=4)
+        target_qubits = []
+        ctrl_qubits = []
+        # definitely a better way to do this but
+        # can't think of it right now
+        for i in range(len(qubit_idx_ssas)):
+            if (i % 2) == 0:
+                ctrl_qubits.append(qubit_idx_ssas[i])
+            else:
+                target_qubits.append(qubit_idx_ssas[i])
+
+        target_qubits = tuple(target_qubits)
+        ctrl_qubits = tuple(ctrl_qubits)
+
+        match ctrl_op_target_gate:
+            case op.stmts.X():
+                stim_stmt = stim.CX(controls=ctrl_qubits, targets=target_qubits)
+            case op.stmts.Y():
+                stim_stmt = stim.CY(controls=ctrl_qubits, targets=target_qubits)
+            case op.stmts.Z():
+                stim_stmt = stim.CZ(controls=ctrl_qubits, targets=target_qubits)
+            case _:
+                raise NotImplementedError(
+                    "Control gates beyond CX, CY, and CZ are not supported"
+                )
 
-            return RewriteResult(has_done_something=True)
+        stim_stmt.insert_before(apply_stmt_ctrl)
 
-        return RewriteResult()
+        return RewriteResult(has_done_something=True)

test/squin/stim/stim.py

@@ -58,12 +58,14 @@ def test_1q():
         # pass the wires through some 1 Qubit operators
         (op1 := squin.op.stmts.S()),
         (op2 := squin.op.stmts.H()),
-        (op3 := squin.op.stmts.X()),
+        (op3 := squin.op.stmts.Identity(sites=1)),
+        (op4 := squin.op.stmts.Identity(sites=1)),
         (v0 := squin.wire.Apply(op1.result, w0.result)),
         (v1 := squin.wire.Apply(op2.result, v0.results[0])),
         (v2 := squin.wire.Apply(op3.result, v1.results[0])),
+        (v3 := squin.wire.Apply(op4.result, v2.results[0])),
         (
-            squin.wire.Wrap(v2.results[0], q0.result)
+            squin.wire.Wrap(v3.results[0], q0.result)
         ),  # for wrap, just free a use for the result SSAval
         (ret_none := func.ConstantNone()),
         (func.Return(ret_none)),
@@ -107,4 +109,64 @@ def test_1q():
     constructed_method.print()
 
 
-test_1q()
+def test_control():
+
+    stmts: list[ir.Statement] = [
+        # Create qubit register
+        (n_qubits := as_int(2)),
+        (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
+        # Get qubis out
+        (idx0 := as_int(0)),
+        (q0 := qasm2.core.QRegGet(reg=qreg.result, idx=idx0.result)),
+        (idx1 := as_int(1)),
+        (q1 := qasm2.core.QRegGet(reg=qreg.result, idx=idx1.result)),
+        # Unwrap to get wires
+        (w0 := squin.wire.Unwrap(qubit=q0.result)),
+        (w1 := squin.wire.Unwrap(qubit=q1.result)),
+        # set up control gate
+        (op1 := squin.op.stmts.X()),
+        (cx := squin.op.stmts.Control(op1.result, n_controls=1)),
+        (app := squin.wire.Apply(cx.result, w0.result, w1.result)),
+        # wrap things back
+        (squin.wire.Wrap(wire=app.results[0], qubit=q0.result)),
+        (squin.wire.Wrap(wire=app.results[1], qubit=q1.result)),
+        (ret_none := func.ConstantNone()),
+        (func.Return(ret_none)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+    constructed_method.print()
+
+    address_frame, _ = address.AddressAnalysis(
+        constructed_method.dialects
+    ).run_analysis(constructed_method, no_raise=False)
+
+    nsites_frame, _ = nsites.NSitesAnalysis(constructed_method.dialects).run_analysis(
+        constructed_method, no_raise=False
+    )
+
+    constructed_method.print(analysis=address_frame.entries)
+    constructed_method.print(analysis=nsites_frame.entries)
+
+    wrap_squin_analysis = WrapSquinAnalysis(
+        address_analysis=address_frame.entries, op_site_analysis=nsites_frame.entries
+    )
+    fix_walk_squin_analysis = Fixpoint(Walk(wrap_squin_analysis))
+    rewrite_res = fix_walk_squin_analysis.rewrite(constructed_method.code)
+
+    # attempt rewrite to Stim
+    # Be careful with Fixpoint, can go to infinity until reaches defined threshold
+    squin_to_stim = Walk(SquinToStim())
+    rewrite_res = squin_to_stim.rewrite(constructed_method.code)
+
+    constructed_method.print()
+
+    # Get rid of the unused statements
+    dce = Fixpoint(Walk(DeadCodeElimination()))
+    rewrite_res = dce.rewrite(constructed_method.code)
+    print(rewrite_res)
+
+    constructed_method.print()
+
+
+test_control()