Refactor gemini logical validation to use kirin's infrastructure (#629)

Co-authored-by: kaihsin <kaihsinwu@gmail.com>

Author: david-pl

pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.22.0",
     "scipy>=1.13.1",
-    "kirin-toolchain>=0.21.0,<0.23.0",
+    "kirin-toolchain~=0.22.2",
     "rich>=13.9.4",
     "pydantic>=1.3.0,<2.11.0",
     "pandas>=2.2.3",

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

@@ -1,2 +1,2 @@
-from . import impls as impls  # NOTE: register methods
-from .analysis import GeminiLogicalValidationAnalysis as GeminiLogicalValidationAnalysis
+from . import impls as impls, analysis as analysis  # NOTE: register methods
+from .analysis import GeminiLogicalValidation as GeminiLogicalValidation

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

@@ -1,17 +1,40 @@
+from typing import Any
+from dataclasses import dataclass
+
 from kirin import ir
+from kirin.lattice import EmptyLattice
+from kirin.analysis import Forward, ForwardFrame
+from kirin.validation import ValidationPass
 
 from bloqade import squin
-from bloqade.validation.analysis import ValidationFrame, ValidationAnalysis
 
 
-class GeminiLogicalValidationAnalysis(ValidationAnalysis):
+class _GeminiLogicalValidationAnalysis(Forward[EmptyLattice]):
     keys = ["gemini.validate.logical"]
 
     first_gate = True
+    lattice = EmptyLattice
 
-    def eval_fallback(self, frame: ValidationFrame, node: ir.Statement):
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
         if isinstance(node, 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 super().eval_fallback(frame, node)
+        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 GeminiLogicalValidation(ValidationPass):
+    """Validates a logical gemini program"""
+
+    def name(self) -> str:
+        return "Gemini Logical Validation"
+
+    def run(self, method: ir.Method) -> tuple[Any, list[ir.ValidationError]]:
+        analysis = _GeminiLogicalValidationAnalysis(method.dialects)
+        frame, _ = analysis.run(method)
+
+        return frame, analysis.get_validation_errors()

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

@@ -1,12 +1,10 @@
 from kirin import ir, interp as _interp
-from kirin.analysis import const
+from kirin.analysis import ForwardFrame, 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
 
-from .analysis import GeminiLogicalValidationAnalysis
+from .analysis import _GeminiLogicalValidationAnalysis
 
 
 @scf.dialect.register(key="gemini.validate.logical")
@@ -15,87 +13,82 @@ class __ScfGeminiLogicalValidation(_interp.MethodTable):
     @_interp.impl(scf.IfElse)
     def if_else(
         self,
-        interp: GeminiLogicalValidationAnalysis,
-        frame: ValidationFrame,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
         stmt: scf.IfElse,
     ):
-        frame.errors.append(
+        interp.add_validation_error(
+            stmt,
             ir.ValidationError(
                 stmt, "If statements are not supported in logical Gemini programs!"
-            )
-        )
-        return (
-            Error(
-                message="If statements are not supported in logical Gemini programs!"
             ),
         )
+        return (interp.lattice.bottom(),)
 
     @_interp.impl(scf.For)
     def for_loop(
         self,
-        interp: GeminiLogicalValidationAnalysis,
-        frame: ValidationFrame,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
         stmt: scf.For,
     ):
-        if isinstance(stmt.iterable.hints.get("const"), const.Value):
-            return (interp.lattice.top(),)
+        if not isinstance(stmt.iterable.hints.get("const"), const.Value):
 
-        frame.errors.append(
-            ir.ValidationError(
+            interp.add_validation_error(
                 stmt,
-                "Non-constant iterable in for loop is not supported in Gemini logical programs!",
+                ir.ValidationError(
+                    stmt,
+                    "Non-constant iterable in for loop is not supported in Gemini logical programs!",
+                ),
             )
-        )
 
-        return (
-            Error(
-                message="Non-constant iterable in for loop is not supported in Gemini logical programs!"
-            ),
-        )
+        return (interp.lattice.bottom(),)
 
 
 @func.dialect.register(key="gemini.validate.logical")
 class __FuncGeminiLogicalValidation(_interp.MethodTable):
     @_interp.impl(func.Invoke)
     def invoke(
         self,
-        interp: GeminiLogicalValidationAnalysis,
-        frame: ValidationFrame,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
         stmt: func.Invoke,
     ):
-        frame.errors.append(
+        interp.add_validation_error(
+            stmt,
             ir.ValidationError(
                 stmt,
                 "Function invocations not supported in logical Gemini program!",
                 help="Make sure to decorate your function with `@logical(inline = True)` or `@logical(aggressive_unroll = True)` to inline function calls",
-            )
+            ),
         )
 
-        return tuple(
-            Error(
-                message="Function invocations not supported in logical Gemini program!"
-            )
-            for _ in stmt.results
-        )
+        return tuple(interp.lattice.bottom() for _ in stmt.results)
 
 
 @gate.dialect.register(key="gemini.validate.logical")
 class __GateGeminiLogicalValidation(_interp.MethodTable):
+
     @_interp.impl(gate.stmts.U3)
-    def u3(
+    @_interp.impl(gate.stmts.T)
+    @_interp.impl(gate.stmts.Rx)
+    @_interp.impl(gate.stmts.Ry)
+    @_interp.impl(gate.stmts.Rz)
+    def non_clifford(
         self,
-        interp: GeminiLogicalValidationAnalysis,
-        frame: ValidationFrame,
-        stmt: gate.stmts.U3,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: gate.stmts.SingleQubitGate | gate.stmts.RotationGate,
     ):
         if interp.first_gate:
             interp.first_gate = False
             return ()
 
-        frame.errors.append(
+        interp.add_validation_error(
+            stmt,
             ir.ValidationError(
                 stmt,
-                "U3 gate can only be used for initial state preparation, i.e. as the first gate!",
-            )
+                f"Non-clifford gate {stmt.name} can only be used for initial state preparation, i.e. as the first gate!",
+            ),
         )
         return ()

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

@@ -4,13 +4,13 @@
 from kirin.passes import Default
 from kirin.prelude import structural_no_opt
 from kirin.dialects import py, func, ilist
+from kirin.validation import ValidationSuite
 from typing_extensions import Doc
 from kirin.passes.inline import InlinePass
 
 from bloqade.squin import gate, qubit
-from bloqade.validation import KernelValidation
 from bloqade.rewrite.passes import AggressiveUnroll
-from bloqade.gemini.analysis.logical_validation import GeminiLogicalValidationAnalysis
+from bloqade.gemini.analysis.logical_validation import GeminiLogicalValidation
 
 from ._dialect import dialect
 
@@ -63,8 +63,9 @@ def run_pass(
             default_pass.fixpoint(mt)
 
         if verify:
-            validator = KernelValidation(GeminiLogicalValidationAnalysis)
-            validator.run(mt, no_raise=no_raise)
+            validator = ValidationSuite([GeminiLogicalValidation])
+            validation_result = validator.validate(mt)
+            validation_result.raise_if_invalid()
             mt.verify()
 
     return run_pass