get rid of redundant rewrite of QASM2 const/expr to py const/expr - already exists

Author: johnzl-777

src/bloqade/squin/passes/qasm2_gate_func_to_squin.py

@@ -4,6 +4,7 @@
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.rewrite.passes import CallGraphPass
+from bloqade.qasm2.passes.qasm2py import _QASM2Py as QASM2ToPyRule
 
 from ..rewrite import qasm2 as qasm2_rule
 
@@ -37,7 +38,7 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
 
         combined_qasm2_rules = Walk(
             Chain(
-                qasm2_rule.QASM2ExprToSquin(),
+                QASM2ToPyRule(),
                 qasm2_rule.QASM2CoreToSquin(),
                 qasm2_rule.QASM2UOPToSquin(),
                 qasm2_rule.QASM2NoiseToSquin(),

src/bloqade/squin/passes/qasm2_to_squin.py

@@ -10,11 +10,14 @@
 from bloqade.squin.rewrite.qasm2 import (
     QASM2UOPToSquin,
     QASM2CoreToSquin,
-    QASM2ExprToSquin,
     QASM2NoiseToSquin,
     QASM2GlobParallelToSquin,
 )
 
+# There's a QASM2Py pass that only applies an _QASM2Py rewrite rule,
+# I just want the rule here.
+from bloqade.qasm2.passes.qasm2py import _QASM2Py as QASM2ToPyRule
+
 from .qasm2_gate_func_to_squin import QASM2GateFuncToSquinPass
 
 
@@ -26,7 +29,7 @@ def unsafe_run(self, mt: ir.Method) -> RewriteResult:
         # rewrite all QASM2 to squin first
         rewrite_result = Walk(
             Chain(
-                QASM2ExprToSquin(),
+                QASM2ToPyRule(),
                 QASM2CoreToSquin(),
                 QASM2UOPToSquin(),
                 QASM2GlobParallelToSquin(),

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

@@ -1,5 +1,4 @@
 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 .noise_to_squin import QASM2NoiseToSquin as QASM2NoiseToSquin
 from .glob_parallel_to_squin import QASM2GlobParallelToSquin as QASM2GlobParallelToSquin

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

@@ -0,0 +1,32 @@
+import io
+
+from kirin import ir
+
+from bloqade import stim, squin
+from bloqade.stim.emit import EmitStimMain
+from bloqade.stim.passes import SquinToStimPass
+
+
+def codegen(mt: ir.Method):
+    # method should not have any arguments!
+    buf = io.StringIO()
+    emit = EmitStimMain(dialects=stim.main, io=buf)
+    emit.initialize()
+    emit.run(mt)
+    return buf.getvalue().strip()
+
+
+@squin.kernel
+def test_simple_linear():
+
+    qs = squin.qalloc(4)
+    m0 = squin.broadcast.measure(qs)
+    squin.set_detector([m0[0], m0[1]], coordinates=[0, 0])
+    m1 = squin.broadcast.measure(qs)
+    squin.set_detector([m1[0], m1[1]], coordinates=[1, 1])
+
+
+test_simple_linear.print()
+SquinToStimPass(dialects=test_simple_linear.dialects)(test_simple_linear)
+test_simple_linear.print()
+print(codegen(test_simple_linear))

test/analysis/measure_id/test_refactored_measure_id.py

@@ -1,8 +1,7 @@
 import math
 
 from kirin import types as kirin_types
-from kirin.rewrite import Walk
-from kirin.dialects import py, func, math as kirin_math
+from kirin.dialects import py
 from kirin.dialects.ilist import IListType
 
 from bloqade import qasm2, squin
@@ -11,79 +10,9 @@
 from bloqade.squin.passes import QASM2ToSquin
 from bloqade.rewrite.passes import AggressiveUnroll
 from bloqade.analysis.address import AddressAnalysis
-from bloqade.squin.rewrite.qasm2 import (
-    QASM2ExprToSquin,
-)
 from bloqade.analysis.address.lattice import AddressReg
 
 
-def test_expr_rewrite():
-
-    @qasm2.main
-    def expr_program():
-        # constants
-        x = 0  # noqa: F841
-        # ConstPI only added from lowering
-        y = qasm2.dialects.expr.stmts.ConstPI()  # noqa: F841
-        z = -1.75  # noqa: F841
-        # binary ops
-        a = 1 + 1  # noqa: F841
-        b = 2 * 2  # noqa: F841
-        c = 3 - 3  # noqa: F841
-        d = 4 / 4  # noqa: F841
-        e = 5**2  # noqa: F841
-
-        # math
-        a = 0.2
-        qasm2.sin(a)
-        qasm2.cos(a)
-        qasm2.tan(a)
-        qasm2.exp(a)
-        qasm2.ln(a)
-        qasm2.sqrt(a)
-        return
-
-    expr_program.print()
-
-    Walk(QASM2ExprToSquin()).rewrite(expr_program.code)
-
-    expr_program.print()
-
-    actual_stmt_sequence = list(expr_program.callable_region.walk())
-
-    def is_pi_const(stmt: py.Constant):
-        return isinstance(stmt, py.Constant) and math.isclose(stmt.value.data, math.pi)
-
-    assert any(is_pi_const(stmt) for stmt in actual_stmt_sequence)
-
-    assert qasm2.expr.ConstFloat not in actual_stmt_sequence
-    assert qasm2.expr.ConstInt not in actual_stmt_sequence
-    assert qasm2.expr.ConstPI not in actual_stmt_sequence
-
-    no_const_actual_sequence = [
-        type(stmt)
-        for stmt in actual_stmt_sequence
-        if not isinstance(stmt, (py.Constant, func.ConstantNone, func.Return))
-    ]
-
-    expected_stmt_sequence = [
-        py.unary.stmts.USub,
-        py.binop.Add,
-        py.binop.Mult,
-        py.binop.Sub,
-        py.binop.Div,
-        py.binop.Pow,
-        kirin_math.stmts.sin,
-        kirin_math.stmts.cos,
-        kirin_math.stmts.tan,
-        kirin_math.stmts.exp,
-        kirin_math.stmts.log2,
-        kirin_math.stmts.sqrt,
-    ]
-
-    assert no_const_actual_sequence == expected_stmt_sequence
-
-
 def test_qasm2_core():
 
     @qasm2.main(fold=False)