Improvements to squin-related infrastructure (#310)

These were changes I originally planned to make in separate PRs but
considering how useful they are for other work along with the fact it
would get a bit messy having so many very small PRs I thought it best to
consolidate them here.

Changes include:
- support for site analysis to work with kirin scf dialect
- enabling the squin to stim control gate mapping to be externally
accessible

---------

Co-authored-by: Phillip Weinberg <[email protected]>

Author: 2 people

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

@@ -1,4 +1,6 @@
 from kirin import interp
+from kirin.dialects import scf
+from kirin.dialects.scf.typeinfer import TypeInfer as ScfTypeInfer
 
 from bloqade.squin import op, wire
 
@@ -78,3 +80,8 @@ def rot(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Rot):
     def scale(self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.Scale):
         op_sites = frame.get(stmt.op)
         return (op_sites,)
+
+
+@scf.dialect.register(key="op.nsites")
+class ScfSquinOp(ScfTypeInfer):
+    pass

src/bloqade/squin/passes/__init__.py

@@ -1,11 +1,7 @@
-from .wire_to_stim import SquinWireToStim as SquinWireToStim
-from .qubit_to_stim import SquinQubitToStim as SquinQubitToStim
-from .squin_measure import SquinMeasureToStim as SquinMeasureToStim
 from .wrap_analysis import (
     SitesAttribute as SitesAttribute,
     AddressAttribute as AddressAttribute,
-    WrapSquinAnalysis as WrapSquinAnalysis,
-)
-from .wire_identity_elimination import (
-    SquinWireIdentityElimination as SquinWireIdentityElimination,
+    WrapOpSiteAnalysis as WrapOpSiteAnalysis,
+    WrapAddressAnalysis as WrapAddressAnalysis,
 )
+from .remove_dangling_qubits import RemoveDeadRegister as RemoveDeadRegister

src/bloqade/squin/passes/stim.py

@@ -0,0 +1,19 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import qubit
+
+
+class RemoveDeadRegister(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, qubit.New):
+            return RewriteResult()
+
+        if bool(node.result.uses):
+            return RewriteResult()
+        else:
+            node.delete()
+
+        return RewriteResult(has_done_something=True)

src/bloqade/squin/rewrite/__init__.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from dataclasses import dataclass
 
 from kirin import ir
@@ -40,33 +41,47 @@ def print_impl(self, printer: Printer) -> None:
 
 
 @dataclass
-class WrapSquinAnalysis(RewriteRule):
+class WrapAnalysis(RewriteRule):
 
+    @abstractmethod
+    def wrap(self, value: ir.SSAValue) -> bool:
+        pass
+
+    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
+        has_done_something = any(self.wrap(arg) for arg in node.args)
+        return RewriteResult(has_done_something=has_done_something)
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        has_done_something = any(self.wrap(result) for result in node.results)
+        return RewriteResult(has_done_something=has_done_something)
+
+
+@dataclass
+class WrapAddressAnalysis(WrapAnalysis):
     address_analysis: dict[ir.SSAValue, Address]
-    op_site_analysis: dict[ir.SSAValue, Sites]
 
     def wrap(self, value: ir.SSAValue) -> bool:
         address_analysis_result = self.address_analysis[value]
-        op_site_analysis_result = self.op_site_analysis[value]
 
-        if value.hints.get("address") and value.hints.get("sites"):
+        if value.hints.get("address") is not None:
             return False
-        else:
-            value.hints["address"] = AddressAttribute(address_analysis_result)
-            value.hints["sites"] = SitesAttribute(op_site_analysis_result)
+
+        value.hints["address"] = AddressAttribute(address_analysis_result)
 
         return True
 
-    def rewrite_Block(self, node: ir.Block) -> RewriteResult:
-        has_done_something = False
-        for arg in node.args:
-            if self.wrap(arg):
-                has_done_something = True
-        return RewriteResult(has_done_something=has_done_something)
 
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        has_done_something = False
-        for result in node.results:
-            if self.wrap(result):
-                has_done_something = True
-        return RewriteResult(has_done_something=has_done_something)
+@dataclass
+class WrapOpSiteAnalysis(WrapAnalysis):
+
+    op_site_analysis: dict[ir.SSAValue, Sites]
+
+    def wrap(self, value: ir.SSAValue) -> bool:
+        op_site_analysis_result = self.op_site_analysis[value]
+
+        if value.hints.get("sites") is not None:
+            return False
+
+        value.hints["sites"] = SitesAttribute(op_site_analysis_result)
+
+        return True

src/bloqade/squin/rewrite/remove_dangling_qubits.py

@@ -1,4 +1,4 @@
-from . import emit as emit, parse as parse, dialects as dialects
+from . import emit as emit, parse as parse, passes as passes, dialects as dialects
 from .groups import main as main
 from ._wrappers import (
     h as h,

src/bloqade/squin/rewrite/wrap_analysis.py

@@ -47,7 +47,7 @@ def print_impl(self, printer: Printer) -> None:
 
 @statement(dialect=dialect)
 class ConstBool(ir.Statement):
-    """IR Statement representing a constant float value."""
+    """IR Statement representing a constant boolean value."""
 
     name = "constant.bool"
     traits = frozenset({ir.Pure(), ir.ConstantLike(), lowering.FromPythonCall()})

src/bloqade/stim/__init__.py

@@ -0,0 +1 @@
+from .squin_to_stim import SquinToStim as SquinToStim

src/bloqade/stim/dialects/auxiliary/stmts/const.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+
+from kirin.passes import Fold
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    DeadCodeElimination,
+    CommonSubexpressionElimination,
+)
+from kirin.ir.method import Method
+from kirin.passes.abc import Pass
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.stim.groups import main as stim_main_group
+from bloqade.stim.rewrite import (
+    SquinWireToStim,
+    PyConstantToStim,
+    SquinQubitToStim,
+    SquinMeasureToStim,
+    SquinWireIdentityElimination,
+)
+from bloqade.squin.rewrite import RemoveDeadRegister
+
+
+@dataclass
+class SquinToStim(Pass):
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        fold_pass = Fold(mt.dialects)
+        # propagate constants
+        rewrite_result = fold_pass(mt)
+
+        # Assume that address analysis and
+        # wrapping has been done before this pass!
+
+        # Wrap Rewrite + SquinToStim can happen w/ standard walk
+        rewrite_result = (
+            Walk(
+                Chain(
+                    SquinQubitToStim(),
+                    SquinWireToStim(),
+                    SquinMeasureToStim(),  # reduce duplicated logic, can split out even more rules later
+                    SquinWireIdentityElimination(),
+                )
+            )
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        # Convert all PyConsts to Stim Constants
+        rewrite_result = (
+            Walk(Chain(PyConstantToStim())).rewrite(mt.code).join(rewrite_result)
+        )
+
+        # remove any squin.qubit.new that's left around
+        ## Not considered pure so DCE won't touch it but
+        ## it isn't being used anymore considering everything is a
+        ## stim dialect statement
+        rewrite_result = (
+            Fixpoint(
+                Walk(
+                    Chain(
+                        DeadCodeElimination(),
+                        CommonSubexpressionElimination(),
+                        RemoveDeadRegister(),
+                    )
+                )
+            )
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        # do program verification here,
+        # ideally use built-in .verify() to catch any
+        # incompatible statements as the full rewrite process should not
+        # leave statements from any other dialects (other than the stim main group)
+        mt_verification_clone = mt.similar(stim_main_group)
+
+        # suggested by Kai, will work for now
+        for stmt in mt_verification_clone.code.walk():
+            assert (
+                stmt.dialect in stim_main_group
+            ), "Statements detected that are not part of the stim dialect, please verify the original code is valid for rewrite!"
+
+        return rewrite_result