more testing, verification implemented

Author: johnzl-777

src/bloqade/squin/rewrite/stim.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, cast
 from dataclasses import dataclass
 
 from kirin import ir
@@ -9,7 +9,7 @@
 from bloqade import stim
 from bloqade.squin import op, wire, qubit
 from bloqade.analysis.address import Address, AddressWire, AddressQubit, AddressTuple
-from bloqade.squin.analysis.nsites import Sites
+from bloqade.squin.analysis.nsites import Sites, NumberSites
 
 # Probably best to move these attributes to a
 # separate file? Keep here for now
@@ -93,7 +93,9 @@ def get_address_attr(self, value: ir.SSAValue) -> AddressAttribute:
 
     def get_sites_attr(self, value: ir.SSAValue):
         try:
-            return value.hints["sites"]
+            sites_attr = value.hints["sites"]
+            assert isinstance(sites_attr, SitesAttribute)
+            return sites_attr
         except KeyError:
             raise KeyError(f"The sites analysis hint for {value} does not exist")
 
@@ -111,7 +113,7 @@ 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
+            case op.stmts.Identity():
                 return stim.gate.Identity
             case _:
                 raise NotImplementedError(
@@ -193,6 +195,64 @@ def insert_qubit_idx_after_apply(
                 "unsupported statement detected, only wire.Apply and qubit.Apply statements are supported by this method"
             )
 
+    def verify_num_site_Apply(self, apply_stmt: wire.Apply | qubit.Apply):
+
+        # get the number of wires/qubits that went into the statement
+        if isinstance(apply_stmt, wire.Apply):
+            num_sites_targeted = len(apply_stmt.inputs)
+        elif isinstance(apply_stmt, qubit.Apply):
+            address_attr = self.get_address_attr(apply_stmt.qubits)
+            # ilist has AddressTuple type,
+            # should be the case that the types INSIDE the AddressTuple
+            # are all AddressQubit
+            address_tuple = address_attr.address
+            assert isinstance(address_tuple, AddressTuple)
+            num_sites_targeted = len(address_tuple.data)
+        else:
+            raise TypeError(
+                "Number of sites verification cannot occur on statements other than wire.Apply and qubit.Apply"
+            )
+
+        # The only single qubit operator that can have its size customized is the Identity gate.
+        # There are two possible valid uses for size.
+        # Either:
+        ## Apply(Identity(size=n), wire0, ..., wire_n)
+        # Or:
+        ## Apply(Identity(size=1), wire0, ..., wire_n)
+        # both should have the same effect, and can naturally be represented in Stim as:
+        # 1QGate q1 q2 q3 q4
+
+        op_ssa = apply_stmt.operator
+        op_stmt = op_ssa.owner
+        cast(ir.Statement, op_stmt)
+
+        sites_attr = self.get_sites_attr(op_ssa)
+        sites_type = sites_attr.sites
+        assert isinstance(sites_type, NumberSites)
+        num_sites_supported = sites_type.sites
+
+        if isinstance(op_stmt, op.stmts.Identity):
+            if num_sites_supported != 1 or num_sites_supported != num_sites_targeted:
+                raise ValueError(
+                    "squin.op.Identity must either have sites = 1 or sites = the number of qubits/wires it is being applied on"
+                )
+        elif isinstance(op_stmt, op.stmts.Control):
+            # in Stim control gates have the following supported syntax
+            ## CX 1 2
+            ## CX 1 2 3 4 (equivalent to CX 1 2, then CX 3 4)
+
+            if (
+                num_sites_targeted < num_sites_supported
+                or num_sites_targeted % num_sites_supported != 0
+            ):
+                raise ValueError(
+                    "Mismatch found between Control gate supported number of qubits/wires and number of qubits/wires being supplied."
+                )
+        else:
+            return None
+
+        return None
+
     # 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
@@ -206,6 +266,7 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         match node:
             case wire.Apply() | qubit.Apply():
+                self.verify_num_site_Apply(node)
                 return self.rewrite_Apply(node)
             case wire.Wrap():
                 return self.rewrite_Wrap(node)
@@ -244,11 +305,6 @@ def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
         # but we can handle X, Y, Z, H, and S here just fine
         stim_1q_op = self.get_stim_1q_gate(applied_op)
 
-        # wire.Apply -> tuple of SSA -> AddressTuple
-        # qubit.Apply -> list of qubits ->  AddressTuple
-        ## Both cases the statements follow the Stim semantics of
-        ## 1QGate a b c d ....
-
         if isinstance(apply_stmt, qubit.Apply):
             address_attr = self.get_address_attr(apply_stmt.qubits)
             qubit_idx_ssas = self.insert_qubit_idx_from_address(
@@ -266,6 +322,12 @@ def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
         stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
         stim_1q_stmt.insert_before(apply_stmt)
 
+        # Could I safely delete the apply statements?
+        # If it's a a qubit.Apply yes, because it doesn't return anything
+        # If it's a wire.Apply no, because the `results` of that Apply get used later on
+        if isinstance(apply_stmt, qubit.Apply):
+            apply_stmt.delete()
+
         return RewriteResult(has_done_something=True)
 
     def rewrite_Control(
@@ -278,12 +340,9 @@ def rewrite_Control(
         ctrl_op = apply_stmt_ctrl.operator.owner
         assert isinstance(ctrl_op, op.stmts.Control)
 
-        # enforce that n_controls is 1
-
         ctrl_op_target_gate = ctrl_op.op.owner
         assert isinstance(ctrl_op_target_gate, op.stmts.Operator)
 
-        # should enforce that this is some multiple of 2
         qubit_idx_ssas = self.insert_qubit_idx_after_apply(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)
@@ -314,6 +373,9 @@ def rewrite_Control(
 
         stim_stmt.insert_before(apply_stmt_ctrl)
 
+        if isinstance(apply_stmt_ctrl, qubit.Apply):
+            apply_stmt_ctrl.delete()
+
         return RewriteResult(has_done_something=True)
 
     def rewrite_Measure(

test/squin/stim/stim.py

@@ -78,6 +78,135 @@ def test_wire_1q():
     constructed_method.print()
 
 
+def test_parallel_wire_1q_application():
+
+    stmts: list[ir.Statement] = [
+        # Create qubit register
+        (n_qubits := as_int(4)),
+        (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
+        # Get qubits 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)),
+        (idx2 := as_int(2)),
+        (q2 := qasm2.core.QRegGet(reg=qreg.result, idx=idx2.result)),
+        (idx3 := as_int(3)),
+        (q3 := qasm2.core.QRegGet(reg=qreg.result, idx=idx3.result)),
+        # Unwrap to get wires
+        (w0 := squin.wire.Unwrap(qubit=q0.result)),
+        (w1 := squin.wire.Unwrap(qubit=q1.result)),
+        (w2 := squin.wire.Unwrap(qubit=q2.result)),
+        (w3 := squin.wire.Unwrap(qubit=q3.result)),
+        # Apply with stim semantics
+        (h_op := squin.op.stmts.H()),
+        (
+            app_res := squin.wire.Apply(
+                h_op.result, w0.result, w1.result, w2.result, w3.result
+            )
+        ),
+        # Wrap everything back
+        (squin.wire.Wrap(app_res.results[0], q0.result)),
+        (squin.wire.Wrap(app_res.results[1], q1.result)),
+        (squin.wire.Wrap(app_res.results[2], q2.result)),
+        (squin.wire.Wrap(app_res.results[3], q3.result)),
+        (ret_none := func.ConstantNone()),
+        (func.Return(ret_none)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    constructed_method.print()
+
+    squin_to_stim = squin_passes.SquinToStim(constructed_method.dialects)
+    squin_to_stim(constructed_method)
+
+    constructed_method.print()
+
+
+def test_parallel_qubit_1q_application():
+
+    stmts: list[ir.Statement] = [
+        # Create qubit register
+        (n_qubits := as_int(4)),
+        (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
+        # Get qubits 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)),
+        (idx2 := as_int(2)),
+        (q2 := qasm2.core.QRegGet(reg=qreg.result, idx=idx2.result)),
+        (idx3 := as_int(3)),
+        (q3 := qasm2.core.QRegGet(reg=qreg.result, idx=idx3.result)),
+        # create ilist of qubits
+        (q_list := ilist.New(values=(q0.result, q1.result, q2.result, q3.result))),
+        # Apply with stim semantics
+        (h_op := squin.op.stmts.H()),
+        (app_res := squin.qubit.Apply(h_op.result, q_list.result)),  # noqa: F841
+        # Measure everything out
+        (meas_res := squin.qubit.Measure(q_list.result)),  # noqa: F841
+        (ret_none := func.ConstantNone()),
+        (func.Return(ret_none)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    constructed_method.print()
+
+    squin_to_stim = squin_passes.SquinToStim(constructed_method.dialects)
+    squin_to_stim(constructed_method)
+
+    constructed_method.print()
+
+
+def test_parallel_control_gate_wire_application():
+
+    stmts: list[ir.Statement] = [
+        # Create qubit register
+        (n_qubits := as_int(4)),
+        (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
+        # Get qubits 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)),
+        (idx2 := as_int(2)),
+        (q2 := qasm2.core.QRegGet(reg=qreg.result, idx=idx2.result)),
+        (idx3 := as_int(3)),
+        (q3 := qasm2.core.QRegGet(reg=qreg.result, idx=idx3.result)),
+        # Unwrap to get wires
+        (w0 := squin.wire.Unwrap(qubit=q0.result)),
+        (w1 := squin.wire.Unwrap(qubit=q1.result)),
+        (w2 := squin.wire.Unwrap(qubit=q2.result)),
+        (w3 := squin.wire.Unwrap(qubit=q3.result)),
+        # Create and apply CX gate
+        (x_op := squin.op.stmts.X()),
+        (ctrl_x_op := squin.op.stmts.Control(x_op.result, n_controls=1)),
+        (
+            app_res := squin.wire.Apply(
+                ctrl_x_op.result, w0.result, w1.result, w2.result, w3.result
+            )
+        ),
+        # measure it all out
+        (meas_res_0 := squin.wire.Measure(app_res.results[0])),  # noqa: F841
+        (meas_res_1 := squin.wire.Measure(app_res.results[1])),  # noqa: F841
+        (meas_res_2 := squin.wire.Measure(app_res.results[2])),  # noqa: F841
+        (meas_res_3 := squin.wire.Measure(app_res.results[3])),  # noqa: F841
+        (ret_none := func.ConstantNone()),
+        (func.Return(ret_none)),
+    ]
+
+    constructed_method = gen_func_from_stmts(stmts)
+
+    constructed_method.print()
+
+    squin_to_stim = squin_passes.SquinToStim(constructed_method.dialects)
+    squin_to_stim(constructed_method)
+
+    constructed_method.print()
+
+
 def test_wire_control():
 
     stmts: list[ir.Statement] = [
@@ -255,7 +384,10 @@ def test_wire_measure_and_reset():
     constructed_method.print()
 
 
-test_wire_measure_and_reset()
+# test_wire_measure_and_reset()
 # test_qubit_measure_and_reset()
 # test_wire_reset()
-# test_qubit_reset()
+
+# test_parallel_qubit_1q_application()
+# test_parallel_wire_1q_application()
+test_parallel_control_gate_wire_application()