Merge branch 'main' into john/repeat-support

Author: johnzl-777

src/bloqade/analysis/measure_id/impls.py

@@ -4,7 +4,7 @@
 from kirin.dialects import py, scf, func, ilist
 from kirin.ir.attrs.py import PyAttr
 
-from bloqade import qubit, annotate
+from bloqade import qubit, gemini, annotate
 
 from .lattice import (
     Predicate,
@@ -18,6 +18,8 @@
 )
 from .analysis import MeasureIDFrame, MeasurementIDAnalysis
 
+# from bloqade.gemini.dialects.logical import stmts as gemini_stmts, dialect as logical_dialect
+
 
 @qubit.dialect.register(key="measure_id")
 class SquinQubit(interp.MethodTable):
@@ -80,6 +82,31 @@ def measurement_predicate(
         return (MeasureIdTuple(data=tuple(predicate_measure_ids)),)
 
 
+@gemini.logical.dialect.register(key="measure_id")
+class LogicalQubit(interp.MethodTable):
+    @interp.impl(gemini.logical.stmts.TerminalLogicalMeasurement)
+    def terminal_measurement(
+        self,
+        interp: MeasurementIDAnalysis,
+        frame: interp.Frame,
+        stmt: gemini.logical.stmts.TerminalLogicalMeasurement,
+    ):
+        # try to get the length of the list
+        qubits_type = stmt.qubits.type
+        # vars[0] is just the type of the elements in the ilist,
+        # vars[1] can contain a literal with length information
+        num_qubits = qubits_type.vars[1]
+        if not isinstance(num_qubits, kirin_types.Literal):
+            return (AnyMeasureId(),)
+
+        measure_id_bools = []
+        for _ in range(num_qubits.data):
+            interp.measure_count += 1
+            measure_id_bools.append(RawMeasureId(interp.measure_count))
+
+        return (MeasureIdTuple(data=tuple(measure_id_bools)),)
+
+
 @annotate.dialect.register(key="measure_id")
 class Annotate(interp.MethodTable):
     @interp.impl(annotate.stmts.SetObservable)

src/bloqade/gemini/analysis/measurement_validation/__init__.py

@@ -0,0 +1,4 @@
+from . import impls as impls, analysis as analysis  # NOTE: register methods
+from .analysis import (
+    GeminiTerminalMeasurementValidation as GeminiTerminalMeasurementValidation,
+)

src/bloqade/gemini/analysis/measurement_validation/analysis.py

@@ -0,0 +1,66 @@
+from typing import Any
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.lattice import EmptyLattice
+from kirin.analysis import Forward, ForwardFrame
+from kirin.validation import ValidationPass
+
+from bloqade.analysis import address, measure_id
+
+
+@dataclass
+class _GeminiTerminalMeasurementValidationAnalysis(Forward[EmptyLattice]):
+    keys = ("gemini.validate.terminal_measurement",)
+
+    measurement_analysis_results: ForwardFrame
+    unique_qubits_allocated: int = 0
+    terminal_measurement_encountered: bool = False
+    lattice = EmptyLattice
+
+    # boilerplate, not really worried about these right now
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
+        return tuple(self.lattice.bottom() for _ in range(len(node.results)))
+
+    def method_self(self, method: ir.Method) -> EmptyLattice:
+        return self.lattice.bottom()
+
+
+@dataclass
+class GeminiTerminalMeasurementValidation(ValidationPass):
+
+    analysis_cache: dict = field(default_factory=dict)
+
+    def name(self) -> str:
+        return "Gemini Terminal Measurement Validation"
+
+    def get_required_analyses(self) -> list[type]:
+        return [measure_id.MeasurementIDAnalysis]
+
+    def set_analysis_cache(self, cache: dict[type, Any]) -> None:
+        self.analysis_cache.update(cache)
+
+    def run(self, method: ir.Method) -> tuple[Any, list[ir.ValidationError]]:
+
+        # get the data out of the cache and forward it to the underlying analysis
+        measurement_analysis_results = self.analysis_cache.get(
+            measure_id.MeasurementIDAnalysis
+        )
+
+        address_analysis = address.AddressAnalysis(dialects=method.dialects)
+        address_analysis.run(method)
+        unique_qubits_allocated = address_analysis.qubit_count
+
+        assert (
+            measurement_analysis_results is not None
+        ), "Measurement ID analysis results not found in cache"
+
+        analysis = _GeminiTerminalMeasurementValidationAnalysis(
+            method.dialects,
+            measurement_analysis_results,
+            unique_qubits_allocated=unique_qubits_allocated,
+        )
+
+        frame, _ = analysis.run(method)
+
+        return frame, analysis.get_validation_errors()

src/bloqade/gemini/analysis/measurement_validation/impls.py

@@ -0,0 +1,89 @@
+from kirin import ir, interp as _interp
+from kirin.analysis import ForwardFrame
+from kirin.dialects import func
+
+from bloqade import qubit, gemini
+from bloqade.analysis.address.impls import Func as AddressFuncMethodTable
+from bloqade.analysis.measure_id.lattice import MeasureIdTuple
+
+from .analysis import _GeminiTerminalMeasurementValidationAnalysis
+
+
+@qubit.dialect.register(key="gemini.validate.terminal_measurement")
+class __QubitGeminiMeasurementValidation(_interp.MethodTable):
+
+    # This is a non-logical measurement, can safely flag as invalid
+    @_interp.impl(qubit.stmts.Measure)
+    def measure(
+        self,
+        interp: _GeminiTerminalMeasurementValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: qubit.stmts.Measure,
+    ):
+
+        interp.add_validation_error(
+            stmt,
+            ir.ValidationError(
+                stmt,
+                "Non-terminal measurements are not allowed in Gemini programs!",
+            ),
+        )
+
+        return (interp.lattice.bottom(),)
+
+
+@gemini.logical.dialect.register(key="gemini.validate.terminal_measurement")
+class __GeminiLogicalMeasurementValidation(_interp.MethodTable):
+
+    @_interp.impl(gemini.logical.stmts.TerminalLogicalMeasurement)
+    def terminal_measure(
+        self,
+        interp: _GeminiTerminalMeasurementValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: gemini.logical.stmts.TerminalLogicalMeasurement,
+    ):
+
+        # should only be one terminal measurement EVER
+        if not interp.terminal_measurement_encountered:
+            interp.terminal_measurement_encountered = True
+        else:
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    "Multiple terminal measurements are not allowed in Gemini logical programs!",
+                ),
+            )
+            return (interp.lattice.bottom(),)
+
+        measurement_analysis_results = interp.measurement_analysis_results
+        total_qubits_allocated = interp.unique_qubits_allocated
+
+        measure_lattice_element = measurement_analysis_results.get(stmt.result)
+        if not isinstance(measure_lattice_element, MeasureIdTuple):
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    "Measurement ID Analysis failed to produce the necessary results needed for validation.",
+                ),
+            )
+            return (interp.lattice.bottom(),)
+
+        if len(measure_lattice_element.data) != total_qubits_allocated:
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    "The number of qubits in the terminal measurement does not match the number of total qubits allocated! "
+                    + f"{total_qubits_allocated} qubits were allocated but only {len(measure_lattice_element.data)} were measured.",
+                ),
+            )
+            return (interp.lattice.bottom(),)
+
+        return (interp.lattice.bottom(),)
+
+
+@func.dialect.register(key="gemini.validate.terminal_measurement")
+class Func(AddressFuncMethodTable):
+    pass

src/bloqade/gemini/dialects/logical/_interface.py

@@ -7,8 +7,8 @@
 
 from .stmts import TerminalLogicalMeasurement
 
-Len = TypeVar("Len", bound=int)
-CodeN = TypeVar("CodeN", bound=int)
+Len = TypeVar("Len")
+CodeN = TypeVar("CodeN")
 
 
 @lowering.wraps(TerminalLogicalMeasurement)

src/bloqade/gemini/dialects/logical/groups.py

@@ -10,7 +10,6 @@
 
 from bloqade.squin import gate, qubit
 from bloqade.rewrite.passes import AggressiveUnroll
-from bloqade.gemini.analysis.logical_validation import GeminiLogicalValidation
 
 from ._dialect import dialect
 
@@ -63,7 +62,17 @@ def run_pass(
             default_pass.fixpoint(mt)
 
         if verify:
-            validator = ValidationSuite([GeminiLogicalValidation])
+            # stop circular import problems
+            from bloqade.gemini.analysis.logical_validation import (
+                GeminiLogicalValidation,
+            )
+            from bloqade.gemini.analysis.measurement_validation import (
+                GeminiTerminalMeasurementValidation,
+            )
+
+            validator = ValidationSuite(
+                [GeminiLogicalValidation, GeminiTerminalMeasurementValidation]
+            )
             validation_result = validator.validate(mt)
             validation_result.raise_if_invalid()
             mt.verify()

src/bloqade/gemini/dialects/logical/stmts.py

@@ -6,8 +6,8 @@
 
 from ._dialect import dialect
 
-Len = types.TypeVar("Len", bound=types.Int)
-CodeN = types.TypeVar("CodeN", bound=types.Int)
+Len = types.TypeVar("Len")
+CodeN = types.TypeVar("CodeN")
 
 
 @statement(dialect=dialect)

src/bloqade/squin/rewrite/non_clifford_to_U3.py

@@ -0,0 +1,102 @@
+from kirin import ir
+from kirin.rewrite import abc as rewrite_abc
+from kirin.dialects import py
+
+from bloqade.squin.gate import stmts as gate_stmts
+
+
+class RewriteNonCliffordToU3(rewrite_abc.RewriteRule):
+    """Rewrite non-Clifford gates to U3 gates.
+
+    This rewrite rule transforms specific non-Clifford single-qubit gates
+    into equivalent U3 gate representations. The following transformations are applied:
+    - T gate (with adjoint attribute) to U3 gate with parameters (0, 0, ±π/4)
+    - Rx gate to U3 gate with parameters (angle, -π/2, π/2)
+    - Ry gate to U3 gate with parameters (angle, 0, 0)
+    - Rz gate is U3 gate with parameters (0, 0, angle)
+
+    This rewrite should be paired with `U3ToClifford` to canonicalize the circuit.
+
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
+        if not isinstance(
+            node,
+            (
+                gate_stmts.T,
+                gate_stmts.Rx,
+                gate_stmts.Ry,
+                gate_stmts.Rz,
+            ),
+        ):
+            return rewrite_abc.RewriteResult()
+
+        rule = getattr(self, f"rewrite_{type(node).__name__}")
+
+        return rule(node)
+
+    def rewrite_T(self, node: gate_stmts.T) -> rewrite_abc.RewriteResult:
+        if node.adjoint:
+            lam_value = -1.0 / 8.0
+        else:
+            lam_value = 1.0 / 8.0
+
+        (theta_stmt := py.Constant(0.0)).insert_before(node)
+        (phi_stmt := py.Constant(0.0)).insert_before(node)
+        (lam_stmt := py.Constant(lam_value)).insert_before(node)
+
+        node.replace_by(
+            gate_stmts.U3(
+                qubits=node.qubits,
+                theta=theta_stmt.result,
+                phi=phi_stmt.result,
+                lam=lam_stmt.result,
+            )
+        )
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_Rx(self, node: gate_stmts.Rx) -> rewrite_abc.RewriteResult:
+        (phi_stmt := py.Constant(-0.25)).insert_before(node)
+        (lam_stmt := py.Constant(0.25)).insert_before(node)
+
+        node.replace_by(
+            gate_stmts.U3(
+                qubits=node.qubits,
+                theta=node.angle,
+                phi=phi_stmt.result,
+                lam=lam_stmt.result,
+            )
+        )
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_Ry(self, node: gate_stmts.Ry) -> rewrite_abc.RewriteResult:
+        (phi_stmt := py.Constant(0.0)).insert_before(node)
+        (lam_stmt := py.Constant(0.0)).insert_before(node)
+
+        node.replace_by(
+            gate_stmts.U3(
+                qubits=node.qubits,
+                theta=node.angle,
+                phi=phi_stmt.result,
+                lam=lam_stmt.result,
+            )
+        )
+
+        return rewrite_abc.RewriteResult(has_done_something=True)
+
+    def rewrite_Rz(self, node: gate_stmts.Rz) -> rewrite_abc.RewriteResult:
+        (theta_stmt := py.Constant(0.0)).insert_before(node)
+        (phi_stmt := py.Constant(0.0)).insert_before(node)
+
+        node.replace_by(
+            gate_stmts.U3(
+                qubits=node.qubits,
+                theta=theta_stmt.result,
+                phi=phi_stmt.result,
+                lam=node.angle,
+            )
+        )
+
+        return rewrite_abc.RewriteResult(has_done_something=True)

test/analysis/measure_id/test_measure_id.py

@@ -1,7 +1,7 @@
 from kirin.dialects import scf
 from kirin.passes.inline import InlinePass
 
-from bloqade import squin
+from bloqade import squin, gemini
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
 from bloqade.stim.passes.flatten import Flatten
 from bloqade.analysis.measure_id.lattice import (
@@ -339,3 +339,23 @@ def test():
     assert frame.get(results["is_zero_bools"]) == expected_is_zero_bools
     assert frame.get(results["is_one_bools"]) == expected_is_one_bools
     assert frame.get(results["is_lost_bools"]) == expected_is_lost_bools
+
+
+def test_terminal_logical_measurement():
+
+    @gemini.logical.kernel(no_raise=False, typeinfer=True, aggressive_unroll=True)
+    def tm_logical_kernel():
+        q = squin.qalloc(3)
+        tm = gemini.logical.terminal_measure(q)
+        return tm
+
+    frame, _ = MeasurementIDAnalysis(tm_logical_kernel.dialects).run(tm_logical_kernel)
+    # will have a MeasureIdTuple that's not from the terminal measurement,
+    # basically a container of InvalidMeasureIds from the qubits that get allocated
+    analysis_results = [
+        val for val in frame.entries.values() if isinstance(val, MeasureIdTuple)
+    ]
+    expected_result = MeasureIdTuple(
+        data=tuple([RawMeasureId(idx=i) for i in range(1, 4)])
+    )
+    assert expected_result in analysis_results