Verify that U3 can only occur at first position

Author: david-pl

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

@@ -1,17 +1,18 @@
 from kirin import ir
 
+from bloqade import squin
 from bloqade.validation.analysis import ValidationFrame, ValidationAnalysis
-from bloqade.validation.analysis.lattice import ErrorType
 
 
 class GeminiLogicalValidationAnalysis(ValidationAnalysis):
     keys = ["gemini.validate.logical"]
-    lattice = ErrorType
 
-    has_allocated_qubits: bool = False
-
-    def run_method(self, method: ir.Method, args: tuple[ErrorType, ...]):
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+    first_gate = True
 
     def eval_stmt_fallback(self, frame: ValidationFrame, stmt: ir.Statement):
+
+        if isinstance(stmt, squin.gate.stmts.Gate):
+            # NOTE: to validate that only the first encountered gate can be non-Clifford, we need to track this here
+            self.first_gate = False
+
         return (self.lattice.top(),)

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

@@ -2,6 +2,7 @@
 from kirin.analysis import const
 from kirin.dialects import scf, func
 
+from bloqade.squin import gate
 from bloqade.validation.analysis import ValidationFrame
 from bloqade.validation.analysis.lattice import Error
 
@@ -58,3 +59,24 @@ def invoke(
                 help="Make sure to decorate your function with `@logical(inline = True)` or `@logical(aggressive_unroll = True)` to inline function calls",
             ),
         )
+
+
+@gate.dialect.register(key="gemini.validate.logical")
+class __GateGeminiLogicalValidation(_interp.MethodTable):
+    @_interp.impl(gate.stmts.U3)
+    def u3(
+        self,
+        interp: GeminiLogicalValidationAnalysis,
+        frame: ValidationFrame,
+        stmt: gate.stmts.U3,
+    ):
+        if interp.first_gate:
+            interp.first_gate = False
+            return (interp.lattice.top(),)
+
+        return (
+            Error(
+                stmt,
+                "U3 gate can only be used for initial state preparation, i.e. as the first gate!",
+            ),
+        )

src/bloqade/gemini/groups.py

@@ -42,6 +42,9 @@ def run_pass(
         no_raise: Annotated[bool, Doc("do not raise exception during analysis")] = True,
     ) -> None:
 
+        if inline and not aggressive_unroll:
+            InlinePass(mt.dialects, no_raise=no_raise).fixpoint(mt)
+
         if aggressive_unroll:
             AggressiveUnroll(mt.dialects, no_raise=no_raise).fixpoint(mt)
         else:
@@ -56,9 +59,6 @@ def run_pass(
 
             default_pass.fixpoint(mt)
 
-        if inline and not aggressive_unroll:
-            InlinePass(mt.dialects, no_raise=no_raise).fixpoint(mt)
-
         if verify:
             validator = KernelValidation(GeminiLogicalValidationAnalysis)
             validator.run(mt, no_raise=no_raise)

src/bloqade/squin/gate/stmts.py

@@ -8,7 +8,13 @@
 
 
 @statement
-class SingleQubitGate(ir.Statement):
+class Gate(ir.Statement):
+    # NOTE: just for easier isinstance checks elsewhere, all gates inherit from this class
+    pass
+
+
+@statement
+class SingleQubitGate(Gate):
     traits = frozenset({lowering.FromPythonCall()})
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
 
@@ -59,7 +65,7 @@ class SqrtY(SingleQubitNonHermitianGate):
 
 
 @statement
-class RotationGate(ir.Statement):
+class RotationGate(Gate):
     # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
     traits = frozenset({lowering.FromPythonCall()})
     angle: ir.SSAValue = info.argument(types.Float)
@@ -85,7 +91,7 @@ class Rz(RotationGate):
 
 
 @statement
-class ControlledGate(ir.Statement):
+class ControlledGate(Gate):
     traits = frozenset({lowering.FromPythonCall()})
     controls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
     targets: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
@@ -110,7 +116,7 @@ class CZ(ControlledGate):
 
 
 @statement(dialect=dialect)
-class U3(ir.Statement):
+class U3(Gate):
     # NOTE: don't inherit from SingleQubitGate here so the wrapper doesn't have qubits as first arg
     traits = frozenset({lowering.FromPythonCall()})
     theta: ir.SSAValue = info.argument(types.Float)

src/bloqade/validation/analysis/analysis.py

@@ -1,13 +1,17 @@
 from abc import ABC
+from typing import Iterable
+from dataclasses import field, dataclass
 
 from kirin import ir
+from kirin.interp import AbstractFrame
 from kirin.analysis import Forward, ForwardFrame
 
 from .lattice import ErrorType
 
 ValidationFrame = ForwardFrame[ErrorType]
 
 
+@dataclass
 class ValidationAnalysis(Forward[ErrorType], ABC):
     """Analysis pass that indicates errors in the IR according to the respective method tables.
 
@@ -17,9 +21,43 @@ class ValidationAnalysis(Forward[ErrorType], ABC):
 
     lattice = ErrorType
 
+    additional_errors: list[ErrorType] = field(default_factory=list)
+    """List to store return values that are not associated with an SSA Value (e.g. when the statement has no ResultValue)"""
+
     def run_method(self, method: ir.Method, args: tuple[ErrorType, ...]):
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+        return self.run_callable(method.code, (self.lattice.top(),) + args)
 
     def eval_stmt_fallback(self, frame: ValidationFrame, stmt: ir.Statement):
         # NOTE: default to no errors
         return (self.lattice.top(),)
+
+    def set_values(
+        self,
+        frame: AbstractFrame[ErrorType],
+        ssa: Iterable[ir.SSAValue],
+        results: Iterable[ErrorType],
+    ):
+        """Set the abstract values for the given SSA values in the frame.
+
+        This method is overridden to account for additional errors we may
+        encounter when they are not associated to an SSA Value.
+        """
+
+        number_of_ssa_values = 0
+        for ssa_value, result in zip(ssa, results):
+            number_of_ssa_values += 1
+            if ssa_value in frame.entries:
+                frame.entries[ssa_value] = frame.entries[ssa_value].join(result)
+            else:
+                frame.entries[ssa_value] = result
+
+        if isinstance(results, tuple):
+            # NOTE: usually what we have
+            self.additional_errors.extend(results[number_of_ssa_values:])
+
+        for i, result in enumerate(results):
+            # NOTE: only sure-fire way I found to get remaining values from an Iterable
+            if i < number_of_ssa_values:
+                continue
+
+            self.additional_errors.append(result)

src/bloqade/validation/kernel_validation.py

@@ -1,9 +1,10 @@
+import itertools
 from dataclasses import dataclass
 
 from kirin import ir
 
 from .analysis import ValidationFrame, ValidationAnalysis
-from .analysis.lattice import Error
+from .analysis.lattice import Error, ErrorType
 
 
 @dataclass
@@ -14,7 +15,9 @@ def run(self, mt: ir.Method, **kwargs) -> None:
         validation_analysis = self.validation_analysis_cls(mt.dialects)
         validation_frame, _ = validation_analysis.run_analysis(mt, **kwargs)
 
-        errors = self.get_exceptions(mt, validation_frame)
+        errors = self.get_exceptions(
+            mt, validation_frame, validation_analysis.additional_errors
+        )
 
         if len(errors) == 0:
             # Valid program
@@ -23,9 +26,16 @@ def run(self, mt: ir.Method, **kwargs) -> None:
         # TODO: Make something similar to an ExceptionGroup that pretty-prints ValidationErrors
         raise errors[0]
 
-    def get_exceptions(self, mt: ir.Method, validation_frame: ValidationFrame):
+    def get_exceptions(
+        self,
+        mt: ir.Method,
+        validation_frame: ValidationFrame,
+        additional_errors: list[ErrorType],
+    ):
         errors = []
-        for value in validation_frame.entries.values():
+        for value in itertools.chain(
+            validation_frame.entries.values(), additional_errors
+        ):
             if not isinstance(value, Error):
                 continue
 

test/gemini/test_logical.py

@@ -84,4 +84,27 @@ def invalid():
             sub_kernel(q[0])
 
 
-test_func()
+def test_clifford_gates():
+    @gemini.logical
+    def main():
+        q = squin.qalloc(2)
+        squin.u3(0.123, 0.253, 1.2, q[0])
+
+        squin.h(q[0])
+        squin.cx(q[0], q[1])
+
+    with pytest.raises(ir.ValidationError):
+
+        @gemini.logical(no_raise=False)
+        def invalid():
+            q = squin.qalloc(2)
+
+            squin.h(q[0])
+            squin.cx(q[0], q[1])
+            squin.u3(0.123, 0.253, 1.2, q[0])
+
+        frame, _ = GeminiLogicalValidationAnalysis(invalid.dialects).run_analysis(
+            invalid, no_raise=False
+        )
+
+        invalid.print(analysis=frame.entries)