Implement gemini.logical and draft kernel validation (#558)

This implements the simple dialect group as discussed in #531.

However, it also contains a draft (RFC) for validating kernels, using
the above as an example.

I thought a bit about it and in the end decided to split the validation
into two parts:

1. A `ValidationAnalysis` pass with an appropriate lattice (could be
extended to warnings, etc.) that checks for errors in the IR.
2. A simple `KernelValidation` class, that gathers up the errors and
throws them as `kirin.ir.ValidationError` with the appropriate source
info pointing to the line of the error. I went back and forth as to how
to implement this, since it is a bit similar to interpretation. However,
in the end I just decided to go with a very simple dataclass, since it
is sufficiently different from an interpreter or a rewrite.

One thing that's not great and I'm not sure how to improve: we probably
don't want the `ValidationAnalysis` to implement a lot of methods via
method tables. For example, the `gemini.logical` dialect doesn't support
`scf.IfElse`. However, if we want to re-use the `ValidationAnalysis`
elsewhere, then registering a `MethodTable` for `scf` with a method for
`scf.IfElse` might fail validation in kernels where this should be fine.
At the same time, we might want to share validation between different
kernels.

The way I "solved" this for now is to define a new analysis pass just
for the logical dialect that inherits from the generic
`ValidationAnalysis`. If we go this route, then the `ValidationAnalysis`
basically just defines the lattice and a fallback method.

There's also some TODOs left here, e.g. how to display multiple errors,
but I thought I'd get comments as soon as possible.

---------

Co-authored-by: Phillip Weinberg <[email protected]>

Author: david-pl

src/bloqade/gemini/__init__.py

@@ -0,0 +1 @@
+from .groups import logical as logical

src/bloqade/gemini/analysis/__init__.py

@@ -0,0 +1,3 @@
+from .logical_validation.analysis import (
+    GeminiLogicalValidationAnalysis as GeminiLogicalValidationAnalysis,
+)

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

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

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

@@ -0,0 +1,17 @@
+from kirin import ir
+
+from bloqade import squin
+from bloqade.validation.analysis import ValidationFrame, ValidationAnalysis
+
+
+class GeminiLogicalValidationAnalysis(ValidationAnalysis):
+    keys = ["gemini.validate.logical"]
+
+    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 super().eval_stmt_fallback(frame, stmt)

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

@@ -0,0 +1,101 @@
+from kirin import ir, interp as _interp
+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
+
+from .analysis import GeminiLogicalValidationAnalysis
+
+
+@scf.dialect.register(key="gemini.validate.logical")
+class __ScfGeminiLogicalValidation(_interp.MethodTable):
+
+    @_interp.impl(scf.IfElse)
+    def if_else(
+        self,
+        interp: GeminiLogicalValidationAnalysis,
+        frame: ValidationFrame,
+        stmt: scf.IfElse,
+    ):
+        frame.errors.append(
+            ir.ValidationError(
+                stmt, "If statements are not supported in logical Gemini programs!"
+            )
+        )
+        return (
+            Error(
+                message="If statements are not supported in logical Gemini programs!"
+            ),
+        )
+
+    @_interp.impl(scf.For)
+    def for_loop(
+        self,
+        interp: GeminiLogicalValidationAnalysis,
+        frame: ValidationFrame,
+        stmt: scf.For,
+    ):
+        if isinstance(stmt.iterable.hints.get("const"), const.Value):
+            return (interp.lattice.top(),)
+
+        frame.errors.append(
+            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!"
+            ),
+        )
+
+
+@func.dialect.register(key="gemini.validate.logical")
+class __FuncGeminiLogicalValidation(_interp.MethodTable):
+    @_interp.impl(func.Invoke)
+    def invoke(
+        self,
+        interp: GeminiLogicalValidationAnalysis,
+        frame: ValidationFrame,
+        stmt: func.Invoke,
+    ):
+        frame.errors.append(
+            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
+        )
+
+
+@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 ()
+
+        frame.errors.append(
+            ir.ValidationError(
+                stmt,
+                "U3 gate can only be used for initial state preparation, i.e. as the first gate!",
+            )
+        )
+        return ()

src/bloqade/gemini/groups.py

@@ -0,0 +1,67 @@
+from typing import Annotated
+
+from kirin import ir
+from kirin.passes import Default
+from kirin.prelude import structural_no_opt
+from kirin.dialects import py, func, ilist
+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 .analysis import GeminiLogicalValidationAnalysis
+
+
+@ir.dialect_group(structural_no_opt.union([gate, py.constant, qubit, func, ilist]))
+def logical(self):
+    """Compile a function to a Gemini logical kernel."""
+
+    def run_pass(
+        mt,
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc("run type inference and apply the inferred type to IR, default `True`"),
+        ] = True,
+        fold: Annotated[bool, Doc("run folding passes")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+        inline: Annotated[bool, Doc("inline function calls, default `True`")] = True,
+        aggressive_unroll: Annotated[
+            bool,
+            Doc(
+                "Run aggressive inlining and unrolling pass on the IR, default `False`"
+            ),
+        ] = False,
+        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:
+            default_pass = Default(
+                self,
+                verify=verify,
+                fold=fold,
+                aggressive=aggressive,
+                typeinfer=typeinfer,
+                no_raise=no_raise,
+            )
+
+            default_pass.fixpoint(mt)
+
+        if verify:
+            validator = KernelValidation(GeminiLogicalValidationAnalysis)
+            validator.run(mt, no_raise=no_raise)
+            mt.verify()
+
+    return run_pass

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/__init__.py

@@ -0,0 +1,2 @@
+from . import analysis as analysis
+from .kernel_validation import KernelValidation as KernelValidation

src/bloqade/validation/analysis/__init__.py

@@ -0,0 +1,5 @@
+from . import lattice as lattice
+from .analysis import (
+    ValidationFrame as ValidationFrame,
+    ValidationAnalysis as ValidationAnalysis,
+)

src/bloqade/validation/analysis/analysis.py

@@ -0,0 +1,41 @@
+from abc import ABC
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.analysis import ForwardExtra, ForwardFrame
+
+from .lattice import ErrorType
+
+
+@dataclass
+class ValidationFrame(ForwardFrame[ErrorType]):
+    # NOTE: cannot be set[Error] since that's not hashable
+    errors: list[ir.ValidationError] = field(default_factory=list)
+    """List of all ecnountered errors.
+
+    Append a `kirin.ir.ValidationError` to this list in the method implementation
+    in order for it to get picked up by the `KernelValidation` run.
+    """
+
+
+@dataclass
+class ValidationAnalysis(ForwardExtra[ValidationFrame, ErrorType], ABC):
+    """Analysis pass that indicates errors in the IR according to the respective method tables.
+
+    If you need to implement validation for a dialect shared by many groups (for example, if you need to ascertain if statements have a specific form)
+    you'll need to inherit from this class.
+    """
+
+    lattice = ErrorType
+
+    def run_method(self, method: ir.Method, args: tuple[ErrorType, ...]):
+        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 tuple(self.lattice.top() for _ in stmt.results)
+
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> ValidationFrame:
+        return ValidationFrame(code, has_parent_access=has_parent_access)