broke everything up into smaller rulesgst

Author: johnzl-777

src/bloqade/qasm2/dialects/uop/stmts.py

@@ -28,7 +28,7 @@ class TwoQubitCtrlGate(ir.Statement):
 
 @statement(dialect=dialect)
 class CX(TwoQubitCtrlGate):
-    """Alias for the CNOT or CH gate operations."""
+    """Alias for the CNOT or CX gate operations."""
 
     name = "CX"  # Note this is capitalized
 

src/bloqade/squin/passes/qasm2_gate_func_to_squin.py

@@ -40,9 +40,12 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
             Chain(
                 QASM2ToPyRule(),
                 qasm2_rule.QASM2CoreToSquin(),
-                qasm2_rule.QASM2UOPToSquin(),
-                qasm2_rule.QASM2NoiseToSquin(),
                 qasm2_rule.QASM2GlobParallelToSquin(),
+                qasm2_rule.QASM2NoiseToSquin(),
+                qasm2_rule.QASM2IdToSquin(),
+                qasm2_rule.QASM2UOp1QToSquin(),
+                qasm2_rule.QASM2ParametrizedUOp1QToSquin(),
+                qasm2_rule.QASM2UOp2QToSquin(),
             )
         )
 

src/bloqade/squin/passes/qasm2_to_squin.py

@@ -8,10 +8,13 @@
 
 from bloqade import squin
 from bloqade.squin.rewrite.qasm2 import (
-    QASM2UOPToSquin,
+    QASM2IdToSquin,
     QASM2CoreToSquin,
     QASM2NoiseToSquin,
+    QASM2UOp1QToSquin,
+    QASM2UOp2QToSquin,
     QASM2GlobParallelToSquin,
+    QASM2ParametrizedUOp1QToSquin,
 )
 
 # There's a QASM2Py pass that only applies an _QASM2Py rewrite rule,
@@ -31,9 +34,12 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
             Chain(
                 QASM2ToPyRule(),
                 QASM2CoreToSquin(),
-                QASM2UOPToSquin(),
                 QASM2GlobParallelToSquin(),
                 QASM2NoiseToSquin(),
+                QASM2IdToSquin(),
+                QASM2UOp1QToSquin(),
+                QASM2ParametrizedUOp1QToSquin(),
+                QASM2UOp2QToSquin(),
             )
         ).rewrite(mt.code)
 

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

@@ -1,4 +1,9 @@
-from .uop_to_squin import QASM2UOPToSquin as QASM2UOPToSquin
+from .id_to_squin import QASM2IdToSquin as QASM2IdToSquin
 from .core_to_squin import QASM2CoreToSquin as QASM2CoreToSquin
 from .noise_to_squin import QASM2NoiseToSquin as QASM2NoiseToSquin
+from .uop_1q_to_squin import QASM2UOp1QToSquin as QASM2UOp1QToSquin
+from .uop_2q_to_squin import QASM2UOp2QToSquin as QASM2UOp2QToSquin
 from .glob_parallel_to_squin import QASM2GlobParallelToSquin as QASM2GlobParallelToSquin
+from .parametrized_uop_1q_to_squin import (
+    QASM2ParametrizedUOp1QToSquin as QASM2ParametrizedUOp1QToSquin,
+)

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

@@ -5,29 +5,34 @@
 from bloqade import squin
 from bloqade.qasm2.dialects.core import stmts as core_stmts
 
+CORE_TO_SQUIN_MAP = {
+    core_stmts.QRegNew: squin.qubit.qalloc,
+    core_stmts.Reset: squin.qubit.reset,
+}
+
 
 class QASM2CoreToSquin(RewriteRule):
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
-        match 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()
+        if isinstance(node, core_stmts.QRegGet):
+            py_get_item = py.GetItem(
+                obj=node.reg,
+                index=node.idx,
+            )
+            node.replace_by(py_get_item)
+            return RewriteResult(has_done_something=True)
+
+        if isinstance(node, core_stmts.QRegNew):
+            args = (node.n_qubits,)
+        elif isinstance(node, core_stmts.Reset):
+            args = (node.qarg,)
+        else:
+            return RewriteResult()
 
+        new_stmt = func.Invoke(
+            callee=CORE_TO_SQUIN_MAP[type(node)],
+            inputs=args,
+        )
+        node.replace_by(new_stmt)
         return RewriteResult(has_done_something=True)

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

@@ -5,47 +5,30 @@
 from bloqade import squin
 from bloqade.qasm2.dialects import glob, parallel
 
+GLOBAL_PARALLEL_TO_SQUIN_MAP = {
+    glob.UGate: squin.broadcast.u3,
+    parallel.UGate: squin.broadcast.u3,
+    parallel.RZ: squin.broadcast.rz,
+}
+
 
 class QASM2GlobParallelToSquin(RewriteRule):
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
-        match node:
-            case glob.UGate() | parallel.UGate() | parallel.RZ():
-                return self.rewrite_1q_gates(node)
-            case _:
-                return RewriteResult()
-
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_1q_gates(
-        self, stmt: glob.UGate | parallel.UGate | parallel.RZ
-    ) -> RewriteResult:
-
-        match stmt:
-            case glob.UGate(theta=theta, phi=phi, lam=lam) | parallel.UGate(
-                theta=theta, phi=phi, lam=lam
-            ):
-                # ever so slight naming difference,
-                # exists because intended semantics are different
-                match stmt:
-                    case glob.UGate():
-                        qargs = stmt.registers
-                    case parallel.UGate():
-                        qargs = stmt.qargs
-
-                invoke_u_broadcast_stmt = func.Invoke(
-                    callee=squin.broadcast.u3,
-                    inputs=(theta, phi, lam, qargs),
-                )
-                stmt.replace_by(invoke_u_broadcast_stmt)
-            case parallel.RZ(theta=theta, qargs=qargs):
-                invoke_rz_broadcast_stmt = func.Invoke(
-                    callee=squin.broadcast.rz,
-                    inputs=(theta, qargs),
-                )
-                stmt.replace_by(invoke_rz_broadcast_stmt)
-            case _:
-                return RewriteResult()
-
+        if isinstance(node, glob.UGate):
+            args = (node.theta, node.phi, node.lam, node.registers)
+        elif isinstance(node, parallel.UGate):
+            args = (node.theta, node.phi, node.lam, node.qargs)
+        elif isinstance(node, parallel.RZ):
+            args = (node.theta, node.qargs)
+        else:
+            return RewriteResult()
+
+        squin_equivalent_stmt = GLOBAL_PARALLEL_TO_SQUIN_MAP[type(node)]
+        invoke_stmt = func.Invoke(
+            callee=squin_equivalent_stmt,
+            inputs=args,
+        )
+        node.replace_by(invoke_stmt)
         return RewriteResult(has_done_something=True)

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

@@ -0,0 +1,15 @@
+from kirin import ir
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+import bloqade.qasm2.dialects.uop.stmts as uop_stmts
+
+
+class QASM2IdToSquin(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, uop_stmts.Id):
+            return RewriteResult()
+
+        node.delete()
+        return RewriteResult(has_done_something=True)

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

@@ -1,67 +1,43 @@
 from kirin import ir
-from kirin.dialects import py, func
+from kirin.dialects import func
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade import squin
 from bloqade.qasm2.dialects.noise import stmts as noise_stmts
 
+from .util import num_to_py_constant
 
-class QASM2NoiseToSquin(RewriteRule):
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-
-        match node:
-            case noise_stmts.AtomLossChannel():
-                return self.rewrite_AtomLossChannel(node)
-            case noise_stmts.PauliChannel():
-                return self.rewrite_PauliChannel(node)
-            case noise_stmts.CZPauliChannel():
-                return self.rewrite_CZPauliChannel(node)
-            case _:
-                return RewriteResult()
-
-        return RewriteResult()
+NOISE_TO_SQUIN_MAP = {
+    noise_stmts.AtomLossChannel: squin.broadcast.qubit_loss,
+    noise_stmts.PauliChannel: squin.broadcast.single_qubit_pauli_channel,
+}
 
-    def rewrite_AtomLossChannel(
-        self, stmt: noise_stmts.AtomLossChannel
-    ) -> RewriteResult:
 
-        qargs = stmt.qargs
-        # this is a raw float, not in SSA form yet!
-        prob = stmt.prob
-        prob_stmt = py.Constant(value=prob)
-        prob_stmt.insert_before(stmt)
-
-        invoke_loss_stmt = func.Invoke(
-            callee=squin.broadcast.qubit_loss,
-            inputs=(prob_stmt.result, qargs),
-        )
-
-        stmt.replace_by(invoke_loss_stmt)
-
-        return RewriteResult(has_done_something=True)
-
-    def rewrite_PauliChannel(self, stmt: noise_stmts.PauliChannel) -> RewriteResult:
-
-        qargs = stmt.qargs
-        p_x = stmt.px
-        p_y = stmt.py
-        p_z = stmt.pz
-
-        probs = [p_x, p_y, p_z]
-        probs_ssas = []
+class QASM2NoiseToSquin(RewriteRule):
 
-        for prob in probs:
-            prob_stmt = py.Constant(value=prob)
-            prob_stmt.insert_before(stmt)
-            probs_ssas.append(prob_stmt.result)
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
-        invoke_pauli_channel_stmt = func.Invoke(
-            callee=squin.broadcast.single_qubit_pauli_channel,
-            inputs=(*probs_ssas, qargs),
+        if isinstance(node, noise_stmts.AtomLossChannel):
+            qargs = node.qargs
+            prob = node.prob
+            prob_ssas = num_to_py_constant([prob], stmt_to_insert_before=node)
+        elif isinstance(node, noise_stmts.PauliChannel):
+            qargs = node.qargs
+            p_x = node.px
+            p_y = node.py
+            p_z = node.pz
+            prob_ssas = num_to_py_constant([p_x, p_y, p_z], stmt_to_insert_before=node)
+        elif isinstance(node, noise_stmts.CZPauliChannel):
+            return self.rewrite_CZPauliChannel(node)
+        else:
+            return RewriteResult()
+
+        squin_noise_stmt = NOISE_TO_SQUIN_MAP[type(node)]
+        invoke_stmt = func.Invoke(
+            callee=squin_noise_stmt,
+            inputs=(*prob_ssas, qargs),
         )
-
-        stmt.replace_by(invoke_pauli_channel_stmt)
+        node.replace_by(invoke_stmt)
         return RewriteResult(has_done_something=True)
 
     def rewrite_CZPauliChannel(self, stmt: noise_stmts.CZPauliChannel) -> RewriteResult:
@@ -78,11 +54,8 @@ def rewrite_CZPauliChannel(self, stmt: noise_stmts.CZPauliChannel) -> RewriteRes
 
         error_probs = [px_ctrl, py_ctrl, pz_ctrl, px_qarg, py_qarg, pz_qarg]
         # first half of entries for control qubits, other half for targets
-        error_prob_ssas = []
-        for error_prob in error_probs:
-            error_prob_stmt = py.Constant(value=error_prob)
-            error_prob_stmt.insert_before(stmt)
-            error_prob_ssas.append(error_prob_stmt.result)
+
+        error_prob_ssas = num_to_py_constant(error_probs, stmt_to_insert_before=stmt)
 
         ctrl_pauli_channel_invoke = func.Invoke(
             callee=squin.broadcast.single_qubit_pauli_channel,

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

@@ -0,0 +1,46 @@
+from math import pi
+
+from kirin import ir
+from kirin.dialects import py, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade import squin
+from bloqade.qasm2.dialects.uop import stmts as uop_stmts
+
+PARAMETRIZED_1Q_GATES_TO_SQUIN_MAP = {
+    uop_stmts.UGate: squin.u3,
+    uop_stmts.U1: squin.u3,
+    uop_stmts.U2: squin.u3,
+    uop_stmts.RZ: squin.rz,
+    uop_stmts.RX: squin.rx,
+    uop_stmts.RY: squin.ry,
+}
+
+
+class QASM2ParametrizedUOp1QToSquin(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if isinstance(node, (uop_stmts.RX, uop_stmts.RY, uop_stmts.RZ)):
+            args = (node.theta, node.qarg)
+        elif isinstance(node, (uop_stmts.UGate)):
+            args = (node.theta, node.phi, node.lam, node.qarg)
+        elif isinstance(node, (uop_stmts.U1)):
+            zero_stmt = py.Constant(value=0.0)
+            zero_stmt.insert_before(node)
+            args = (zero_stmt.result, zero_stmt.result, node.lam, node.qarg)
+        elif isinstance(node, (uop_stmts.U2)):
+            half_pi_stmt = py.Constant(value=pi / 2)
+            half_pi_stmt.insert_before(node)
+            args = (half_pi_stmt.result, node.phi, node.lam, node.qarg)
+        else:
+            return RewriteResult()
+
+        squin_equivalent_stmt = PARAMETRIZED_1Q_GATES_TO_SQUIN_MAP[type(node)]
+        invoke_stmt = func.Invoke(
+            callee=squin_equivalent_stmt,
+            inputs=args,
+        )
+        node.replace_by(invoke_stmt)
+
+        return RewriteResult(has_done_something=True)

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

@@ -0,0 +1,33 @@
+from kirin import ir
+from kirin.dialects import func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade import squin
+from bloqade.qasm2.dialects.uop import stmts as uop_stmts
+
+ONE_Q_GATES_TO_SQUIN_MAP = {
+    uop_stmts.X: squin.x,
+    uop_stmts.Y: squin.y,
+    uop_stmts.Z: squin.z,
+    uop_stmts.H: squin.h,
+    uop_stmts.S: squin.s,
+    uop_stmts.T: squin.t,
+    uop_stmts.SX: squin.sqrt_x,
+}
+
+
+class QASM2UOp1QToSquin(RewriteRule):
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        squin_1q_gate = ONE_Q_GATES_TO_SQUIN_MAP.get(type(node))
+        if squin_1q_gate is None:
+            return RewriteResult()
+
+        invoke_stmt = func.Invoke(
+            callee=squin_1q_gate,
+            inputs=(node.qarg,),
+        )
+        node.replace_by(invoke_stmt)
+
+        return RewriteResult(has_done_something=True)