Copied code from bloqade-circuit to bloqade-gemini (#288)

* refactor: copied code and tests for 'gemini' directories from bloqade-circuit

* fix: fixed pyright check errors

* refactor: added impls to gemini + moved gemini to bloqade/gemini from bloqade/lanes/gemini

* build: bumped bloqade-circuit to 0.14.0

* style: ran formatter

* build: used public pypi index

* fix: Fixing the pyright issues by changing import style

---------

Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

Author: jasonhan3

demo/msd.py

@@ -1,10 +1,10 @@
 import math
 
-from bloqade.gemini import logical as gemini_logical
-from bloqade.gemini.logical.stdlib import default_post_processing
 from kirin.dialects import ilist
 
 from bloqade import qubit, squin
+from bloqade.gemini import logical as gemini_logical
+from bloqade.gemini.logical.stdlib import default_post_processing
 from bloqade.lanes.logical_mvp import (
     compile_to_stim_program,
 )

demo/pipeline_demo.py

@@ -1,7 +1,7 @@
-from bloqade.gemini import logical as gemini_logical
 from kirin.dialects import ilist
 
 from bloqade import squin
+from bloqade.gemini import logical as gemini_logical
 from bloqade.lanes.logical_mvp import (
     compile_to_stim_program,
 )

demo/pipeline_details.py

@@ -9,10 +9,9 @@
 # to start we create a simple squin kernel that uses the gemini logical dialect to define a
 # terminal measurement on the circuit.
 
-from bloqade.gemini import logical as gemini_logical
-
 # %%
 from bloqade import squin
+from bloqade.gemini import logical as gemini_logical
 from bloqade.lanes.heuristics.logical_placement import LogicalPlacementStrategy
 
 

demo/simulator_device_demo.py

@@ -3,11 +3,11 @@
 
 import numpy as np
 from bloqade.decoders import BpLsdDecoder
-from bloqade.gemini import logical
-from bloqade.gemini.logical.stdlib import default_post_processing
 from kirin.dialects import ilist
 
 from bloqade import qubit, squin
+from bloqade.gemini import logical
+from bloqade.gemini.logical.stdlib import default_post_processing
 from bloqade.lanes import GeminiLogicalSimulator
 
 

pyproject.toml

@@ -11,7 +11,7 @@ authors = [
     { name = "Phillip Weinberg", email = "pweinberg@quera.com" }
 ]
 dependencies = [
-    "bloqade-circuit[cirq, qasm2, tsim]~=0.13.0",
+    "bloqade-circuit[cirq, qasm2, tsim]~=0.14.0",
     "bloqade-geometry~=0.5.0",
     "deprecated>=1.3.1",
     "kirin-toolchain~=0.22.6",

python/bloqade/gemini/__init__.py

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

python/bloqade/gemini/analysis/__init__.py

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

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

@@ -0,0 +1,64 @@
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+from bloqade.analysis.address import Address, AddressAnalysis, AddressReg
+from kirin import ir
+from kirin.analysis import Forward, ForwardFrame
+from kirin.interp import InterpreterError
+from kirin.lattice import EmptyLattice
+from kirin.validation import ValidationPass
+
+from bloqade import squin
+
+
+@dataclass
+class _GeminiLogicalValidationAnalysis(Forward[EmptyLattice]):
+    keys = ("gemini.validate.logical",)
+
+    lattice = EmptyLattice
+    addr_frame: ForwardFrame[Address]
+
+    first_gates: dict[int, bool] = field(init=False, default_factory=dict)
+
+    max_qubits: ClassVar[int] = 10
+
+    def eval_fallback(self, frame: ForwardFrame, node: ir.Statement):
+        if isinstance(node, squin.gate.stmts.Gate):
+            raise InterpreterError(f"Missing implementation for gate {node}")
+
+        return tuple(self.lattice.bottom() for _ in range(len(node.results)))
+
+    def check_first_gate(self, qubits: ir.SSAValue) -> bool:
+        address = self.addr_frame.get(qubits)
+
+        if not isinstance(address, AddressReg):
+            # NOTE: we should have a flat kernel with simple address analysis, so in case we don't
+            # get concrete addresses, we might as well error here since something's wrong
+            return False
+
+        is_first = True
+        for addr_int in address.data:
+            is_first = is_first and self.first_gates.get(addr_int, True)
+            self.first_gates[addr_int] = False
+
+        return is_first
+
+    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]]:
+        addr_frame, _ = AddressAnalysis(method.dialects).run(method)
+        analysis = _GeminiLogicalValidationAnalysis(
+            method.dialects, addr_frame=addr_frame
+        )
+        frame, _ = analysis.run(method)
+
+        return frame, analysis.get_validation_errors()

python/bloqade/gemini/analysis/logical_validation/analysis.py

@@ -0,0 +1,155 @@
+from bloqade.analysis import address
+from bloqade.squin import gate
+from kirin import interp as _interp, ir
+from kirin.analysis import ForwardFrame, const
+from kirin.dialects import func, scf
+
+from bloqade import qubit
+
+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: ForwardFrame,
+        stmt: scf.IfElse,
+    ):
+        interp.add_validation_error(
+            stmt,
+            ir.ValidationError(
+                stmt, "If statements are not supported in logical Gemini programs!"
+            ),
+        )
+        return (interp.lattice.bottom(),)
+
+    @_interp.impl(scf.For)
+    def for_loop(
+        self,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: scf.For,
+    ):
+        if not isinstance(stmt.iterable.hints.get("const"), const.Value):
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    "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: ForwardFrame,
+        stmt: func.Invoke,
+    ):
+        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(interp.lattice.bottom() for _ in stmt.results)
+
+
+@gate.dialect.register(key="gemini.validate.logical")
+class __GateGeminiLogicalValidation(_interp.MethodTable):
+    @_interp.impl(gate.stmts.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: ForwardFrame,
+        stmt: gate.stmts.SingleQubitGate | gate.stmts.RotationGate,
+    ):
+        if interp.check_first_gate(stmt.qubits):
+            return ()
+
+        interp.add_validation_error(
+            stmt,
+            ir.ValidationError(
+                stmt,
+                f"Non-clifford gate {stmt.name} can only be used for initial state preparation, i.e. as the first gate!",
+            ),
+        )
+        return ()
+
+    @_interp.impl(gate.stmts.X)
+    @_interp.impl(gate.stmts.Y)
+    @_interp.impl(gate.stmts.SqrtX)
+    @_interp.impl(gate.stmts.SqrtY)
+    @_interp.impl(gate.stmts.Z)
+    @_interp.impl(gate.stmts.H)
+    @_interp.impl(gate.stmts.S)
+    def clifford(
+        self,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: gate.stmts.SingleQubitGate,
+    ):
+        # NOTE: ignore result, but make sure the first gate flag is set to False
+        interp.check_first_gate(stmt.qubits)
+
+        return ()
+
+    @_interp.impl(gate.stmts.CX)
+    @_interp.impl(gate.stmts.CY)
+    @_interp.impl(gate.stmts.CZ)
+    def controlled_gate(
+        self,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: gate.stmts.ControlledGate,
+    ):
+        interp.check_first_gate(stmt.controls)
+        interp.check_first_gate(stmt.targets)
+        return ()
+
+
+@qubit.dialect.register(key="gemini.validate.logical")
+class QubitMethods(_interp.MethodTable):
+    @_interp.impl(qubit.stmts.New)
+    def check_qubit_allocation(
+        self,
+        interp: _GeminiLogicalValidationAnalysis,
+        frame: ForwardFrame,
+        stmt: qubit.stmts.New,
+    ):
+        qubit_val = interp.addr_frame.get(stmt.result)
+        if not isinstance(qubit_val, address.AddressQubit):
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    "Cannot determine qubit address location.",
+                ),
+            )
+            return (interp.lattice.bottom(),)
+
+        if qubit_val.data >= interp.max_qubits:
+            interp.add_validation_error(
+                stmt,
+                ir.ValidationError(
+                    stmt,
+                    f"Qubit allocations exceeded {interp.max_qubits}.",
+                ),
+            )
+
+        return (interp.lattice.bottom(),)