use callgraph pass instead of nesting rewrite rule applications

Author: johnzl-777

src/bloqade/squin/passes/qasm2_gate_func_to_squin.py

@@ -0,0 +1,52 @@
+from kirin import ir, passes
+from kirin.rewrite import Walk, Chain
+from kirin.dialects import func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.rewrite.passes import CallGraphPass
+
+from ..rewrite import qasm2 as qasm2_rule
+
+
+class QASM2GateFuncToKirinFunc(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        from bloqade.qasm2.dialects.expr.stmts import GateFunction
+
+        if not isinstance(node, GateFunction):
+            return RewriteResult()
+
+        kirin_func = func.Function(
+            sym_name=node.sym_name,
+            signature=node.signature,
+            body=node.body,
+        )
+        node.replace_by(kirin_func)
+
+        return RewriteResult(has_done_something=True)
+
+
+class QASM2GateFuncToSquinPass(passes.Pass):
+
+    def unsafe_run(self, mt: ir.Method) -> RewriteResult:
+        convert_to_kirin_func = CallGraphPass(
+            dialects=mt.dialects, rule=Walk(QASM2GateFuncToKirinFunc())
+        )
+        rewrite_result = convert_to_kirin_func(mt)
+
+        combined_qasm2_rules = Walk(
+            Chain(
+                qasm2_rule.QASM2ExprToSquin(),
+                qasm2_rule.QASM2CoreToSquin(),
+                qasm2_rule.QASM2UOPToSquin(),
+                qasm2_rule.QASM2NoiseToSquin(),
+                qasm2_rule.QASM2GlobParallelToSquin(),
+            )
+        )
+
+        body_conversion_pass = CallGraphPass(
+            dialects=mt.dialects, rule=combined_qasm2_rules
+        )
+        rewrite_result = body_conversion_pass(mt).join(rewrite_result)
+
+        return rewrite_result

src/bloqade/squin/passes/qasm2_to_squin.py

@@ -11,11 +11,12 @@
     QASM2UOPToSquin,
     QASM2CoreToSquin,
     QASM2ExprToSquin,
-    QASM2FuncToSquin,
     QASM2NoiseToSquin,
     QASM2GlobParallelToSquin,
 )
 
+from .qasm2_gate_func_to_squin import QASM2GateFuncToSquinPass
+
 
 @dataclass
 class QASM2ToSquin(Pass):
@@ -25,7 +26,6 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
         # rewrite all QASM2 to squin first
         rewrite_result = Walk(
             Chain(
-                QASM2FuncToSquin(),
                 QASM2ExprToSquin(),
                 QASM2CoreToSquin(),
                 QASM2UOPToSquin(),
@@ -34,6 +34,13 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
             )
         ).rewrite(mt.code)
 
+        # go into subkernels
+        rewrite_result = (
+            QASM2GateFuncToSquinPass(dialects=mt.dialects)
+            .unsafe_run(mt)
+            .join(rewrite_result)
+        )
+
         # kernel should be entirely in squin dialect now
         mt.dialects = squin.kernel
 

src/bloqade/squin/rewrite/qasm2/__init__.py

@@ -1,6 +1,5 @@
 from .uop_to_squin import QASM2UOPToSquin as QASM2UOPToSquin
 from .core_to_squin import QASM2CoreToSquin as QASM2CoreToSquin
 from .expr_to_squin import QASM2ExprToSquin as QASM2ExprToSquin
-from .func_to_squin import QASM2FuncToSquin as QASM2FuncToSquin
 from .noise_to_squin import QASM2NoiseToSquin as QASM2NoiseToSquin
 from .glob_parallel_to_squin import QASM2GlobParallelToSquin as QASM2GlobParallelToSquin