Implement simple no-cloning validation for flattened kernels (#698)

Simplified version of #607 until we have a proper solution. This is
enough to support gemini logical kernels for now.

Closes #695 .

Author: david-pl

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

@@ -0,0 +1,4 @@
+from .validation import (
+    FlatKernelNoCloningValidation as FlatKernelNoCloningValidation,
+    _FlatKernelNoCloningAnalysis as _FlatKernelNoCloningAnalysis,
+)

src/bloqade/analysis/validation/simple_nocloning/validation.py

@@ -0,0 +1,73 @@
+from typing import Any
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.lattice import EmptyLattice
+from kirin.analysis import Forward
+from kirin.validation import ValidationPass
+from kirin.analysis.forward import ForwardFrame
+
+from bloqade.analysis.address import Address, AddressAnalysis
+
+
+class _FlatKernelNoCloningAnalysis(Forward[EmptyLattice]):
+    """Simple no-cloning validation for kernels that have been aggressively inlined."""
+
+    keys = ("validate.nocloning.flatkernel",)
+    lattice = EmptyLattice
+    _address_frame: ForwardFrame[Address] | None = None
+
+    def eval_fallback(
+        self, frame: ForwardFrame[EmptyLattice], node: ir.Statement
+    ) -> None:
+        pass
+
+    def run(self, method: ir.Method, *args: EmptyLattice, **kwargs: EmptyLattice):
+        if self._address_frame is None:
+            address_analysis = AddressAnalysis(method.dialects)
+            address_frame, _ = address_analysis.run(method)
+            self._address_frame = address_frame
+        return super().run(method, *args, **kwargs)
+
+    def method_self(self, method: ir.Method) -> EmptyLattice:
+        return EmptyLattice.bottom()
+
+    def collect_errors(self, stmt: ir.Statement, addresses: list[int]):
+        seen = set()
+        duplicates = set()
+
+        for addr in addresses:
+            if addr in seen:
+                duplicates.add(addr)
+            else:
+                seen.add(addr)
+
+        self.add_validation_error(
+            stmt,
+            ir.ValidationError(
+                stmt,
+                f"Gate {stmt.name.upper()} applies to the qubits {duplicates} more than once.",
+            ),
+        )
+
+
+@dataclass
+class FlatKernelNoCloningValidation(ValidationPass):
+    _analysis: _FlatKernelNoCloningAnalysis = field(init=False)
+
+    def name(self) -> str:
+        return "No-Cloning Validation"
+
+    def get_required_analyses(self) -> list[type]:
+        return [AddressAnalysis]
+
+    def run(self, method: ir.Method) -> tuple[Any, list[ir.ValidationError]]:
+        analysis = _FlatKernelNoCloningAnalysis(method.dialects)
+        frame, _ = analysis.run(
+            method, *(EmptyLattice.bottom() for _ in range(len(method.args) - 1))
+        )
+
+        self._analysis = analysis
+        errors = analysis.get_validation_errors()
+
+        return frame, errors

src/bloqade/gemini/analysis/logical_validation/analysis.py

@@ -13,7 +13,7 @@
 
 @dataclass
 class _GeminiLogicalValidationAnalysis(Forward[EmptyLattice]):
-    keys = ["gemini.validate.logical"]
+    keys = ("gemini.validate.logical",)
 
     lattice = EmptyLattice
     addr_frame: ForwardFrame[Address]

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

@@ -10,6 +10,7 @@
 
 from bloqade.squin import gate, qubit
 from bloqade.rewrite.passes import AggressiveUnroll
+from bloqade.analysis.validation.simple_nocloning import FlatKernelNoCloningValidation
 
 from ._dialect import dialect
 
@@ -71,7 +72,11 @@ def run_pass(
             )
 
             validator = ValidationSuite(
-                [GeminiLogicalValidation, GeminiTerminalMeasurementValidation]
+                [
+                    GeminiLogicalValidation,
+                    GeminiTerminalMeasurementValidation,
+                    FlatKernelNoCloningValidation,
+                ]
             )
             validation_result = validator.validate(mt)
             validation_result.raise_if_invalid()

src/bloqade/squin/__init__.py

@@ -53,6 +53,9 @@
     single_qubit_pauli_channel as single_qubit_pauli_channel,
 )
 from .analysis.fidelity import impls as impls
+from .analysis.validation.simple_nocloning import (  # noqa: F401
+    impls as simple_nocloning_impls,
+)
 
 # NOTE: it's important to keep these imports here since they import squin.kernel
 # we skip isort here

src/bloqade/squin/analysis/validation/simple_nocloning/__init__.py

@@ -0,0 +1,149 @@
+from kirin import interp
+from kirin.lattice import EmptyLattice
+from kirin.analysis import ForwardFrame
+
+from bloqade.squin import gate, noise
+from bloqade.analysis.address.lattice import AddressReg, PartialIList
+from bloqade.analysis.validation.simple_nocloning import _FlatKernelNoCloningAnalysis
+
+
+@gate.dialect.register(key="validate.nocloning.flatkernel")
+class GateMethods(interp.MethodTable):
+    @interp.impl(gate.stmts.X)
+    @interp.impl(gate.stmts.Y)
+    @interp.impl(gate.stmts.Z)
+    @interp.impl(gate.stmts.H)
+    @interp.impl(gate.stmts.S)
+    @interp.impl(gate.stmts.T)
+    @interp.impl(gate.stmts.Rx)
+    @interp.impl(gate.stmts.Ry)
+    @interp.impl(gate.stmts.Rz)
+    def single_qubit_gate(
+        self,
+        interp_: _FlatKernelNoCloningAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: gate.stmts.SingleQubitGate,
+    ):
+        if interp_._address_frame is None:
+            return
+
+        qubit_addrs = interp_._address_frame.get(stmt.qubits)
+
+        if not isinstance(qubit_addrs, AddressReg):
+            return
+
+        unique_addrs = set(qubit_addrs.data)
+        if len(qubit_addrs.data) == len(unique_addrs):
+            return
+
+        interp_.collect_errors(stmt, list(qubit_addrs.data))
+
+    @interp.impl(gate.stmts.CX)
+    @interp.impl(gate.stmts.CY)
+    @interp.impl(gate.stmts.CZ)
+    def controlled_gate(
+        self,
+        interp_: _FlatKernelNoCloningAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: gate.stmts.ControlledGate,
+    ):
+
+        if interp_._address_frame is None:
+            return
+
+        control_addrs = interp_._address_frame.get(stmt.controls)
+        target_addrs = interp_._address_frame.get(stmt.targets)
+
+        if not isinstance(control_addrs, AddressReg) or not isinstance(
+            target_addrs, AddressReg
+        ):
+            return
+
+        all_addrs = list(control_addrs.data) + list(target_addrs.data)
+        unique_addrs = set(all_addrs)
+
+        if len(all_addrs) == len(unique_addrs):
+            return
+
+        interp_.collect_errors(stmt, all_addrs)
+
+
+@noise.dialect.register(key="validate.nocloning.flatkernel")
+class NoiseMethods(interp.MethodTable):
+    @interp.impl(noise.stmts.Depolarize)
+    @interp.impl(noise.stmts.SingleQubitPauliChannel)
+    @interp.impl(noise.stmts.QubitLoss)
+    def single_qubit_noise_channel(
+        self,
+        interp_: _FlatKernelNoCloningAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: (
+            noise.stmts.SingleQubitPauliChannel
+            | noise.stmts.Depolarize
+            | noise.stmts.QubitLoss
+        ),
+    ):
+        if interp_._address_frame is None:
+            return
+
+        qubit_addrs = interp_._address_frame.get(stmt.qubits)
+
+        if not isinstance(qubit_addrs, AddressReg):
+            return
+
+        if len(qubit_addrs.data) == len(set(qubit_addrs.data)):
+            return
+
+        interp_.collect_errors(stmt, list(qubit_addrs.data))
+
+    @interp.impl(noise.stmts.Depolarize2)
+    @interp.impl(noise.stmts.TwoQubitPauliChannel)
+    def two_qubit_noise_channel(
+        self,
+        interp_: _FlatKernelNoCloningAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: noise.stmts.Depolarize2 | noise.stmts.TwoQubitPauliChannel,
+    ):
+        if interp_._address_frame is None:
+            return
+
+        control_addrs = interp_._address_frame.get(stmt.controls)
+        target_addrs = interp_._address_frame.get(stmt.targets)
+
+        if not isinstance(control_addrs, AddressReg) or not isinstance(
+            target_addrs, AddressReg
+        ):
+            return
+
+        all_addrs = list(control_addrs.data) + list(target_addrs.data)
+
+        if len(all_addrs) == len(set(all_addrs)):
+            return
+
+        interp_.collect_errors(stmt, all_addrs)
+
+    @interp.impl(noise.stmts.CorrelatedQubitLoss)
+    def correlated_loss(
+        self,
+        interp_: _FlatKernelNoCloningAnalysis,
+        frame: ForwardFrame[EmptyLattice],
+        stmt: noise.stmts.CorrelatedQubitLoss,
+    ):
+        if interp_._address_frame is None:
+            return
+
+        qubit_addrs = interp_._address_frame.get(stmt.qubits)
+
+        if not isinstance(qubit_addrs, PartialIList):
+            return
+
+        all_addresses = []
+        for group_addrs in qubit_addrs.data:
+            if not isinstance(group_addrs, AddressReg):
+                continue
+            all_addresses.extend(group_addrs.data)
+
+        if len(all_addresses) == len(set(all_addresses)):
+            return
+
+        interp_.collect_errors(stmt, all_addresses)

src/bloqade/squin/analysis/validation/simple_nocloning/impls.py

@@ -205,3 +205,13 @@ def main():
         squin.broadcast.h(q[1:])
 
     main.print()
+
+
+def test_nocloning():
+    with pytest.raises(ValidationErrorGroup):
+
+        @gemini.logical.kernel
+        def main():
+            q = squin.qalloc(2)
+            squin.broadcast.ry(0.123, [q[0], q[0]])
+            squin.cx(q[1], q[1])

test/gemini/test_logical_validation.py

@@ -0,0 +1,55 @@
+import pytest
+from kirin.validation import ValidationSuite
+from kirin.ir.exception import ValidationErrorGroup
+
+from bloqade import squin
+from bloqade.rewrite.passes import AggressiveUnroll
+from bloqade.analysis.validation.simple_nocloning import FlatKernelNoCloningValidation
+
+
+def test_gates():
+
+    @squin.kernel
+    def bad_kernel():
+        q = squin.qalloc(3)
+        squin.broadcast.x([q[0], q[1], q[0]])
+        squin.broadcast.rx(0.123, [q[1], q[1]])
+        squin.cx(q[2], q[2])
+
+    AggressiveUnroll(bad_kernel.dialects).fixpoint(bad_kernel)
+
+    _, errors = FlatKernelNoCloningValidation().run(bad_kernel)
+    assert len(errors) == 3
+
+
+def test_noise():
+    @squin.kernel
+    def bad_kernel():
+        q = squin.qalloc(3)
+        squin.broadcast.depolarize(0.1, [q[0], q[0]])
+        squin.broadcast.depolarize2(0.1, [q[0]], [q[0]])
+
+    AggressiveUnroll(bad_kernel.dialects).fixpoint(bad_kernel)
+
+    validation_suite = ValidationSuite([FlatKernelNoCloningValidation])
+    result = validation_suite.validate(bad_kernel)
+    assert result.error_count() == 2
+
+    with pytest.raises(ValidationErrorGroup):
+        result.raise_if_invalid()
+
+
+def test_correlated_loss():
+    @squin.kernel
+    def bad_kernel():
+        q = squin.qalloc(3)
+        squin.broadcast.correlated_qubit_loss(0.1, [[q[0], q[1]], [q[1], q[2]]])
+
+    AggressiveUnroll(bad_kernel.dialects).fixpoint(bad_kernel)
+
+    validation_suite = ValidationSuite([FlatKernelNoCloningValidation])
+    result = validation_suite.validate(bad_kernel)
+    assert result.error_count() == 1
+
+    with pytest.raises(ValidationErrorGroup):
+        result.raise_if_invalid()