preliminary handling of Apply

Author: johnzl-777

src/bloqade/squin/rewrite/__init__.py

@@ -1,4 +1,5 @@
 from .stim import (
+    SquinToStim as SquinToStim,
     SitesAttribute as SitesAttribute,
     AddressAttribute as AddressAttribute,
     WrapSquinAnalysis as WrapSquinAnalysis,

src/bloqade/squin/rewrite/stim.py

@@ -115,7 +115,27 @@ def get_stim_1q_gate(self, squin_op: op.stmts.Operator):
     # constants, seems to be more for lowering from Python AST
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        pass
+
+        match node:
+            case wire.Apply() | qubit.Apply():
+                return self.rewrite_Apply(node)
+            case wire.Wrap():
+                return self.rewrite_Wrap(node)
+            case _:
+                return RewriteResult()
+
+        return RewriteResult()
+
+    def rewrite_Wrap(self, wrap_stmt: wire.Wrap) -> RewriteResult:
+
+        # get the wire going into the statement
+        wire_ssa = wrap_stmt.wire
+        # remove the wrap statement altogether, then the wire that went into it
+        wrap_stmt.delete()
+        wire_ssa.delete()
+
+        # do NOT want to delete the qubit SSA! Leave that alone!
+        return RewriteResult(has_done_something=True)
 
     def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
 
@@ -140,8 +160,8 @@ def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
 
             stim_1q_stmt = stim_1q_op(targets=tuple(qubit_idx_ssas))
 
-            apply_stmt.replace_by(stim_1q_stmt)
-            apply_stmt.delete()
+            # can't do any of this because of dependencies downstream
+            # apply_stmt.replace_by(stim_1q_stmt)
 
             return RewriteResult(has_done_something=True)
 
@@ -151,16 +171,21 @@ def rewrite_Apply(self, apply_stmt: qubit.Apply | wire.Apply) -> RewriteResult:
             for wire_ssa in wires_ssa:
                 address_attribute = self.get_address(wire_ssa)
                 # get parent qubit idx
-                wire_address = address_attribute.data
+                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)
 
             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
 
-            apply_stmt.replace_by(stim_1q_stmt)
-            apply_stmt.delete()
+            # apply's results become wires that go to other apply's/wrap stmts
+            # apply_stmt.replace_by(stim_1q_stmt)
 
             return RewriteResult(has_done_something=True)
 

test/squin/stim/stim.py

@@ -0,0 +1,110 @@
+from kirin import ir, types
+from kirin.passes import Fold
+from kirin.rewrite import Walk, Fixpoint, DeadCodeElimination
+from kirin.dialects import py, func, ilist
+
+from bloqade import qasm2, squin
+from bloqade.analysis import address
+from bloqade.squin.rewrite import SquinToStim, WrapSquinAnalysis
+from bloqade.squin.analysis import nsites
+
+
+def as_int(value: int):
+    return py.constant.Constant(value=value)
+
+
+def as_float(value: float):
+    return py.constant.Constant(value=value)
+
+
+def gen_func_from_stmts(stmts):
+
+    extended_dialect = squin.groups.wired.add(qasm2.core).add(ilist)
+
+    block = ir.Block(stmts)
+    block.args.append_from(types.MethodType[[], types.NoneType], "main_self")
+    func_wrapper = func.Function(
+        sym_name="main",
+        signature=func.Signature(inputs=(), output=types.NoneType),
+        body=ir.Region(blocks=block),
+    )
+
+    constructed_method = ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name="main",
+        dialects=extended_dialect,
+        code=func_wrapper,
+        arg_names=[],
+    )
+
+    fold_pass = Fold(extended_dialect)
+    fold_pass(constructed_method)
+
+    return constructed_method
+
+
+def test_1q():
+
+    stmts: list[ir.Statement] = [
+        # Create qubit register
+        (n_qubits := as_int(1)),
+        (qreg := qasm2.core.QRegNew(n_qubits=n_qubits.result)),
+        # Get qubit out
+        (idx0 := as_int(0)),
+        (q0 := qasm2.core.QRegGet(reg=qreg.result, idx=idx0.result)),
+        # Unwrap to get wires
+        (w0 := squin.wire.Unwrap(qubit=q0.result)),
+        # pass the wires through some 1 Qubit operators
+        (op1 := squin.op.stmts.S()),
+        (op2 := squin.op.stmts.H()),
+        (op3 := squin.op.stmts.X()),
+        (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])),
+        (
+            squin.wire.Wrap(v2.results[0], q0.result)
+        ),  # for wrap, just free a use for the result SSAval
+        (ret_none := func.ConstantNone()),
+        (func.Return(ret_none)),
+        # the fact I return a wire here means DCE will NOT go ahead and
+        # eliminate all the other wire.Apply stmts
+    ]
+
+    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)
+
+    # attempt to wrap analysis results
+
+    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)
+
+    # Get rid of the unused statements
+    dce = Fixpoint(Walk(DeadCodeElimination()))
+    rewrite_res = dce.rewrite(constructed_method.code)
+    print(rewrite_res)
+
+    constructed_method.print()
+
+
+test_1q()