wip on converting invoke/functions

src/bloqade/squin/passes/qasm2_to_squin.py

@@ -11,6 +11,7 @@
     QASM2UOPToSquin,
     QASM2CoreToSquin,
     QASM2ExprToSquin,
+    QASM2FuncToSquin,
     QASM2NoiseToSquin,
     QASM2GlobParallelToSquin,
 )
@@ -24,6 +25,7 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
         # rewrite all QASM2 to squin first
         rewrite_result = Walk(
             Chain(
+                QASM2FuncToSquin(),
                 QASM2ExprToSquin(),
                 QASM2CoreToSquin(),
                 QASM2UOPToSquin(),
@@ -43,4 +45,6 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
         )
         TypeInfer(dialects=squin.kernel).unsafe_run(mt).join(rewrite_result)
 
+        mt.dialects = squin.kernel
+
         return rewrite_result

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

@@ -1,5 +1,6 @@
 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

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

@@ -1,9 +1,8 @@
 from kirin import ir
-from kirin.dialects import py, func, ilist
+from kirin.dialects import py, func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade import squin
-from bloqade.types import MeasurementResultType
 from bloqade.qasm2.dialects.core import stmts as core_stmts
 
 
@@ -12,42 +11,23 @@ class QASM2CoreToSquin(RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         match node:
-            case core_stmts.QRegNew():
-                return self.rewrite_QRegNew(node)
-            case core_stmts.CRegNew():
-                return self.rewrite_CRegNew(node)
-            case core_stmts.Reset():
-                return self.rewrite_Reset(node)
-            case core_stmts.QRegGet():
-                return self.rewrite_Get(node)
+            case core_stmts.QRegNew(n_qubits=n_qubits):
+                qalloc_invoke_stmt = func.Invoke(
+                    callee=squin.qubit.qalloc, inputs=(n_qubits,)
+                )
+                node.replace_by(qalloc_invoke_stmt)
+            case core_stmts.Reset(qarg=qarg):
+                reset_invoke_stmt = func.Invoke(
+                    callee=squin.qubit.reset, inputs=(qarg,)
+                )
+                node.replace_by(reset_invoke_stmt)
+            case core_stmts.QRegGet(reg=reg, idx=idx):
+                get_item_stmt = py.GetItem(
+                    obj=reg,
+                    index=idx,
+                )
+                node.replace_by(get_item_stmt)
             case _:
                 return RewriteResult()
 
-    def rewrite_QRegNew(self, stmt: core_stmts.QRegNew) -> RewriteResult:
-        qalloc_invoke_stmt = func.Invoke(
-            callee=squin.qubit.qalloc, inputs=(stmt.n_qubits,)
-        )
-        stmt.replace_by(qalloc_invoke_stmt)
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_CRegNew(self, stmt: core_stmts.CRegNew) -> RewriteResult:
-
-        measurement_list = ilist.New(values=(), elem_type=MeasurementResultType)
-        stmt.replace_by(measurement_list)
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_Reset(self, stmt: core_stmts.Reset) -> RewriteResult:
-
-        squin_reset_stmt = squin.qubit.stmts.Reset(qubits=stmt.qarg)
-        stmt.replace_by(squin_reset_stmt)
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_Get(self, stmt: core_stmts.QRegGet) -> RewriteResult:
-
-        get_item_stmt = py.GetItem(
-            obj=stmt.reg,
-            index=stmt.idx,
-        )
-
-        stmt.replace_by(get_item_stmt)
         return RewriteResult(has_done_something=True)

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

@@ -0,0 +1,25 @@
+from kirin import ir
+from kirin.dialects import func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+class QASM2FuncToSquin(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, func.Invoke):
+            return RewriteResult()
+
+        callee = node.callee
+
+        return self.rewrite_Region(callee.callable_region)
+
+    def rewrite_Region(self, region: ir.Region) -> RewriteResult:
+
+        rewrite_result = RewriteResult()
+
+        for stmt in list(region.walk()):
+            result = self.rewrite_Statement(stmt)
+            rewrite_result = rewrite_result.join(result)
+
+        return rewrite_result

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

@@ -30,6 +30,8 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
                 return self.rewrite_SingleQubit_with_parameters(node)
             case uop_stmts.UGate() | uop_stmts.U1() | uop_stmts.U2():
                 return self.rewrite_u_gates(node)
+            case uop_stmts.Id():
+                return self.rewrite_Id(node)
             case _:
                 return RewriteResult()
 
@@ -147,9 +149,18 @@ def rewrite_SingleQubit_with_parameters(
 
         invoke_stmt = func.Invoke(
             callee=squin_1q_stdlib,
-            inputs=(qarg, theta),
+            inputs=(theta, qarg),
         )
 
         stmt.replace_by(invoke_stmt)
 
         return RewriteResult(has_done_something=True)
+
+    def rewrite_Id(self, stmt: uop_stmts.Id) -> RewriteResult:
+
+        # Identity does not exist in squin,
+        # we can just remove it from the program
+
+        stmt.delete()
+
+        return RewriteResult(has_done_something=True)