Implement no-cloning validation

Author: zhenrongliew

src/bloqade/analysis/validation/nocloning/__init__.py

@@ -0,0 +1,2 @@
+from . import impls as impls
+from .analysis import NoCloningValidation as NoCloningValidation

src/bloqade/analysis/validation/nocloning/analysis.py

@@ -0,0 +1,122 @@
+from dataclasses import field
+
+from kirin import ir
+from kirin.analysis import Forward, TypeInference
+from kirin.dialects import func
+from kirin.analysis.forward import ForwardFrame
+
+from bloqade.analysis.address import (
+    Address,
+    AddressReg,
+    AddressQubit,
+    AddressAnalysis,
+)
+from bloqade.analysis.address.lattice import QubitLike
+
+from .lattice import QubitValidation
+
+
+class NoCloningValidation(Forward[QubitValidation]):
+    """
+    Validates the no-cloning theorem by tracking qubit addresses.
+
+    Built on top of AddressAnalysis to get qubit address information.
+    """
+
+    keys = ["validate.nocloning"]
+    lattice = QubitValidation
+    _address_frame: ForwardFrame[Address] = field(init=False)
+    _type_frame: ForwardFrame = field(init=False)
+    method: ir.Method
+    violations: int = field(default=0, init=False)
+
+    def __init__(self, mtd: ir.Method):
+        """
+        Input:
+          - an ir.Method / kernel function
+            infer dialects from it and remember method.
+        """
+        self.method = mtd
+        super().__init__(mtd.dialects)
+
+    def initialize(self):
+        super().initialize()
+
+        address_analysis = AddressAnalysis(self.dialects)
+        address_analysis.initialize()
+        self._address_frame, _ = address_analysis.run_analysis(self.method)
+
+        type_inference = TypeInference(self.dialects)
+        type_inference.initialize()
+        self._type_frame, _ = type_inference.run_analysis(self.method)
+
+        return self
+
+    def method_self(self, method: ir.Method) -> QubitValidation:
+        return self.lattice.bottom()
+
+    def get_qubit_addresses(self, addr: Address) -> frozenset[int]:
+        """Extract concrete qubit addresses from an Address lattice element."""
+        match addr:
+            case AddressQubit(data=qubit_addr):
+                return frozenset([qubit_addr])
+            case AddressReg(data=addrs):
+                return frozenset(addrs)
+            case _:
+                return frozenset()
+
+    def get_stmt_info(self, stmt: ir.Statement) -> str:
+        """String Report about the statement for violation messages."""
+        if isinstance(stmt, func.Invoke) and hasattr(stmt, "callee"):
+            gate_name = stmt.callee.sym_name.upper()
+            return f"{gate_name} Gate"
+
+        return f"{stmt.__class__.__name__}@{stmt}"
+
+    def eval_stmt_fallback(
+        self, frame: ForwardFrame[QubitValidation], stmt: ir.Statement
+    ) -> tuple[QubitValidation, ...]:
+        """
+        Default statement evaluation: check for qubit usage violations.
+        """
+
+        if not isinstance(stmt, func.Invoke):
+            return tuple(QubitValidation.bottom() for _ in stmt.results)
+
+        address_frame = self._address_frame
+        if address_frame is None:
+            return tuple(QubitValidation.top() for _ in stmt.results)
+
+        has_qubit_args = any(
+            isinstance(address_frame.get(arg), QubitLike) for arg in stmt.args
+        )
+
+        if not has_qubit_args:
+            return tuple(QubitValidation.bottom() for _ in stmt.results)
+
+        used_addrs: list[int] = []
+        for arg in stmt.args:
+            addr = address_frame.get(arg)
+            qubit_addrs = self.get_qubit_addresses(addr)
+            used_addrs.extend(qubit_addrs)
+
+        seen: set[int] = set()
+        violations: list[str] = []
+        stmt_info = self.get_stmt_info(stmt)
+
+        for qubit_addr in used_addrs:
+            if qubit_addr in seen:
+                violations.append(f"Qubit[{qubit_addr}] at {stmt_info}")
+            seen.add(qubit_addr)
+
+        if not violations:
+            return tuple(QubitValidation(violations=frozenset()) for _ in stmt.results)
+
+        usage = QubitValidation(violations=frozenset(violations))
+        return tuple(usage for _ in stmt.results) if stmt.results else (usage,)
+
+    def run_method(
+        self, method: ir.Method, args: tuple[QubitValidation, ...]
+    ) -> tuple[ForwardFrame[QubitValidation], QubitValidation]:
+        self_mt = self.method_self(method)
+        return self.run_callable(method.code, (self_mt,) + args)

src/bloqade/analysis/validation/nocloning/impls.py

@@ -0,0 +1,33 @@
+from kirin import interp
+from kirin.analysis import ForwardFrame
+from kirin.dialects import scf
+
+from .lattice import QubitValidation
+from .analysis import NoCloningValidation
+
+
+@scf.dialect.register(key="validate.nocloning")
+class Scf(interp.MethodTable):
+    @interp.impl(scf.IfElse)
+    def if_else(
+        self,
+        interp_: NoCloningValidation,
+        frame: ForwardFrame[QubitValidation],
+        stmt: scf.IfElse,
+    ):
+        cond_validation = frame.get(stmt.cond)
+
+        then_results = interp_.run_callable_region(
+            frame, stmt, stmt.then_body, (cond_validation,)
+        )
+
+        if stmt.else_body:
+            else_results = interp_.run_callable_region(
+                frame, stmt, stmt.else_body, (cond_validation,)
+            )
+
+            merged = tuple(then_results.join(else_results) for _ in stmt.results)
+        else:
+            merged = tuple(then_results for _ in stmt.results)
+
+        return merged if merged else (QubitValidation.bottom(),)

src/bloqade/analysis/validation/nocloning/lattice.py

@@ -0,0 +1,77 @@
+from typing import FrozenSet, final
+from dataclasses import field, dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+
+@dataclass
+class QubitValidation(
+    SimpleJoinMixin["QubitValidation"],
+    SimpleMeetMixin["QubitValidation"],
+    BoundedLattice["QubitValidation"],
+):
+    """Tracks cloning violations detected during analysis."""
+
+    violations: FrozenSet[str] = field(default_factory=frozenset)
+
+    @classmethod
+    def bottom(cls) -> "QubitValidation":
+        """No violations detected"""
+        return Bottom()
+
+    @classmethod
+    def top(cls) -> "QubitValidation":
+        """Unknown state - assume potential violations"""
+        return Top()
+
+    def is_subseteq(self, other: "QubitValidation") -> bool:
+        """Check if this state is a subset of another.
+
+        Lattice ordering:
+        Bottom ⊑ {{'Qubit[1] at CX Gate'}} ⊑ {{'Qubit[0] at CX Gate'},{'Qubit[1] at CX Gate'}} ⊑ Top
+        """
+        if isinstance(other, Top):
+            return True
+        if isinstance(self, Bottom):
+            return True
+        if isinstance(other, Bottom):
+            return False
+
+        return self.violations.issubset(other.violations)
+
+    def __repr__(self) -> str:
+        """Custom repr to show violations clearly."""
+        if not self.violations:
+            return "QubitValidation()"
+        return f"QubitValidation(violations={self.violations})"
+
+
+@final
+class Bottom(QubitValidation, metaclass=SingletonMeta):
+    """Bottom element representing no violations."""
+
+    def is_subseteq(self, other: QubitValidation) -> bool:
+        """Bottom is subset of everything."""
+        return True
+
+    def __repr__(self) -> str:
+        """Cleaner printing."""
+        return "⊥ (Bottom)"
+
+
+@final
+class Top(QubitValidation, metaclass=SingletonMeta):
+    """Top element representing unknown state with potential violations."""
+
+    def is_subseteq(self, other: QubitValidation) -> bool:
+        """Top is only subset of Top."""
+        return isinstance(other, Top)
+
+    def __repr__(self) -> str:
+        """Cleaner printing."""
+        return "⊤ (Top)"

test/analysis/validation/test_no_cloning.py

@@ -0,0 +1,80 @@
+from typing import Any
+
+import pytest
+from util import collect_validation_errors
+from kirin import ir
+from kirin.dialects.ilist.runtime import IList
+
+from bloqade import squin
+from bloqade.types import Qubit
+from bloqade.analysis.validation.nocloning.lattice import QubitValidation
+from bloqade.analysis.validation.nocloning.analysis import NoCloningValidation
+
+
+@pytest.mark.parametrize("control_gate", [squin.cx, squin.cy, squin.cz])
+def test_control_gate_fail(control_gate: ir.Method[[Qubit, Qubit], Any]):
+    @squin.kernel
+    def bad_control():
+        q = squin.qalloc(1)
+        control_gate(q[0], q[0])
+
+    validation = NoCloningValidation(bad_control)
+    validation.initialize()
+    frame, _ = validation.run_analysis(bad_control)
+    print()
+    bad_control.print(analysis=frame.entries)
+    validation_errors = collect_validation_errors(frame, QubitValidation)
+    assert len(validation_errors) == 1
+
+
+@pytest.mark.parametrize("control_gate", [squin.cx, squin.cy, squin.cz])
+def test_control_gate_conditionals_fail(control_gate: ir.Method[[Qubit, Qubit], Any]):
+    @squin.kernel
+    def bad_control(cond: bool):
+        q = squin.qalloc(10)
+        if cond:
+            control_gate(q[0], q[0])
+        else:
+            control_gate(q[0], q[1])
+        squin.cx(q[1], q[1])
+
+    validation = NoCloningValidation(bad_control)
+    validation.initialize()
+    frame, _ = validation.run_analysis(bad_control)
+    print()
+    bad_control.print(analysis=frame.entries)
+    validation_errors = collect_validation_errors(frame, QubitValidation)
+    # print("Violations:", validation_errors)
+    assert len(validation_errors) == 2
+
+
+@pytest.mark.parametrize("control_gate", [squin.cx, squin.cy, squin.cz])
+def test_control_gate_parallel_fail(control_gate: ir.Method[[Qubit, Qubit], Any]):
+    @squin.kernel
+    def bad_control():
+        q = squin.qalloc(2)
+        control_gate(q[0], q[1])
+
+    validation = NoCloningValidation(bad_control)
+    validation.initialize()
+    frame, _ = validation.run_analysis(bad_control)
+    print()
+    bad_control.print(analysis=frame.entries)
+    validation_errors = collect_validation_errors(frame, QubitValidation)
+    assert len(validation_errors) == 0
+
+
+def test_control_gate_parallel_pass():
+    @squin.kernel
+    def good_kernel():
+        q = squin.qalloc(2)
+        squin.cx(q[0], q[1])
+        squin.cy(q[1], q[1])
+
+    validation = NoCloningValidation(good_kernel)
+    validation.initialize()
+    frame, _ = validation.run_analysis(good_kernel)
+    print()
+    good_kernel.print(analysis=frame.entries)
+    validation_errors = collect_validation_errors(frame, QubitValidation)
+    assert len(validation_errors) == 1

test/analysis/validation/util.py

@@ -0,0 +1,17 @@
+from typing import TypeVar
+
+from kirin.analysis import ForwardFrame
+
+from bloqade.analysis.validation.nocloning.lattice import QubitValidation
+
+T = TypeVar("T", bound=QubitValidation)
+
+
+def collect_validation_errors(
+    frame: ForwardFrame[QubitValidation], typ: type[T]
+) -> list[T]:
+    return [
+        validation_errors
+        for validation_errors in frame.entries.values()
+        if isinstance(validation_errors, typ) and len(validation_errors.violations) > 0
+    ]