Quantum runtime analysis (#68)

* adding analysis to check for quantum runtime instructions

* refactor analysis

* WIP: runtime analysis

* adding statements to frame

* removing unused imporst

* Removing runtime impl this is not a quantum runtime

* adding support for scf, and func dialects

* adding run_method + importing method tables

* Adding tests + fixing some minor bugs

* Adding simplified interface

* testing methods of other dialects + bug fixes

* switching to using interpreters lattice property

Author: weinbe58

src/bloqade/shuttle/analysis/runtime.py

@@ -0,0 +1,136 @@
+from dataclasses import dataclass, field
+
+from kirin import interp, ir
+from kirin.analysis import ForwardExtra, ForwardFrame, const
+from kirin.dialects import func, scf
+from kirin.lattice import EmptyLattice
+
+
+@dataclass
+class RuntimeFrame(ForwardFrame[EmptyLattice]):
+    """Frame for quantum runtime analysis.
+    This frame is used to track the state of quantum operations within a method.
+    """
+
+    quantum_stmts: set[ir.Statement] = field(default_factory=set)
+    """Set of quantum statements in the frame."""
+    is_quantum: bool = False
+    """Whether the frame contains quantum operations."""
+
+
+class RuntimeAnalysis(ForwardExtra[RuntimeFrame, EmptyLattice]):
+    """Forward dataflow analysis to check if a method has quantum runtime.
+
+    This analysis checks if a method contains any quantum runtime operations.
+    It is used to determine if the method can be executed on a quantum device.
+    """
+
+    keys = ["runtime"]
+    lattice = EmptyLattice
+
+    def eval_stmt_fallback(self, frame: RuntimeFrame, stmt: ir.Statement):
+        return tuple(self.lattice.top() for _ in stmt.results)
+
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> RuntimeFrame:
+        return RuntimeFrame(code, has_parent_access=has_parent_access)
+
+    def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
+        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+
+    def has_quantum_runtime(self, method: ir.Method) -> bool:
+        """Return True if the method has quantum runtime operations, False otherwise."""
+        frame, _ = self.run_analysis(method)
+        return frame.is_quantum
+
+
+@scf.dialect.register(key="runtime")
+class Scf(interp.MethodTable):
+
+    @interp.impl(scf.IfElse)
+    def ifelse(self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: scf.IfElse):
+        # If either branch is quantum, the whole ifelse is quantum
+        with _interp.new_frame(stmt, has_parent_access=True) as then_frame:
+            then_result = _interp.run_ssacfg_region(
+                then_frame, stmt.then_body, (_interp.lattice.top(),)
+            )
+
+        with _interp.new_frame(stmt, has_parent_access=True) as else_frame:
+            else_result = _interp.run_ssacfg_region(
+                else_frame, stmt.else_body, (_interp.lattice.top(),)
+            )
+
+        frame.is_quantum = (
+            frame.is_quantum or then_frame.is_quantum or else_frame.is_quantum
+        )
+        frame.quantum_stmts.update(then_frame.quantum_stmts, else_frame.quantum_stmts)
+        match (then_result, else_result):
+            case (interp.ReturnValue(), tuple()):
+                return else_result
+            case (tuple(), interp.ReturnValue()):
+                return then_result
+            case (tuple(), tuple()):
+                return tuple(
+                    then_result.join(else_result)
+                    for then_result, else_result in zip(then_result, else_result)
+                )
+            case _:
+                return tuple(_interp.lattice.top() for _ in stmt.results)
+
+    @interp.impl(scf.For)
+    def for_loop(self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: scf.For):
+        args = (_interp.lattice.top(),) * (len(stmt.initializers) + 1)
+        with _interp.new_frame(stmt, has_parent_access=True) as body_frame:
+            result = _interp.run_ssacfg_region(
+                body_frame, stmt.body, (_interp.lattice.bottom(),)
+            )
+
+        frame.is_quantum = frame.is_quantum or body_frame.is_quantum
+        frame.quantum_stmts.update(body_frame.quantum_stmts)
+        if isinstance(result, interp.ReturnValue) or result is None:
+            return args[1:]
+        else:
+            return tuple(arg.join(res) for arg, res in zip(args[1:], result))
+
+    @interp.impl(scf.Yield)
+    def yield_stmt(
+        self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: scf.Yield
+    ):
+        return interp.YieldValue(frame.get_values(stmt.args))
+
+
+@func.dialect.register(key="runtime")
+class Func(interp.MethodTable):
+
+    @interp.impl(func.Invoke)
+    def invoke(self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: func.Invoke):
+        args = (_interp.lattice.top(),) * len(stmt.inputs)
+        callee_frame, result = _interp.run_method(stmt.callee, args)
+        frame.is_quantum = frame.is_quantum or callee_frame.is_quantum
+        return (result,)
+
+    @interp.impl(func.Call)
+    def call(self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: func.Call):
+        # Check if the called method is quantum
+        callee_result = stmt.callee.hints.get("const")
+        args = (_interp.lattice.top(),) * len(stmt.inputs)
+        if (
+            isinstance(callee_result, const.PartialLambda)
+            and (trait := callee_result.code.get_trait(ir.CallableStmtInterface))
+            is not None
+        ):
+            body = trait.get_callable_region(callee_result.code)
+            with _interp.new_frame(stmt) as callee_frame:
+                result = _interp.run_ssacfg_region(callee_frame, body, args)
+        else:
+            raise InterruptedError("Dynamic method calls are not supported")
+
+        frame.is_quantum = frame.is_quantum or callee_frame.is_quantum
+        return (result,)
+
+    @interp.impl(func.Return)
+    def return_stmt(
+        self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: func.Return
+    ):
+        return interp.ReturnValue(frame.get_values(stmt.results))

src/bloqade/shuttle/dialects/gate/__init__.py

@@ -1,3 +1,4 @@
+from . import runtime as runtime
 from ._dialect import dialect as dialect
 from ._interface import (
     global_r as global_r,

src/bloqade/shuttle/dialects/gate/runtime.py

@@ -0,0 +1,29 @@
+from kirin import interp
+
+from bloqade.shuttle.analysis.runtime import (
+    RuntimeAnalysis,
+    RuntimeFrame,
+)
+
+from ._dialect import dialect
+from .stmts import GlobalR, GlobalRz, LocalR, LocalRz, TopHatCZ
+
+
+@dialect.register(key="runtime")
+class HasQuantumRuntimeMethodTable(interp.MethodTable):
+
+    @interp.impl(TopHatCZ)
+    @interp.impl(LocalRz)
+    @interp.impl(LocalR)
+    @interp.impl(GlobalR)
+    @interp.impl(GlobalRz)
+    def gate(
+        self,
+        interp: RuntimeAnalysis,
+        frame: RuntimeFrame,
+        stmt: TopHatCZ | LocalRz | LocalR | GlobalR | GlobalRz,
+    ) -> interp.StatementResult[RuntimeFrame]:
+        """Handle gate statements and mark the frame as quantum."""
+        frame.is_quantum = True
+        frame.quantum_stmts.add(stmt)
+        return ()

src/bloqade/shuttle/dialects/init/__init__.py

@@ -1,3 +1,4 @@
+from . import runtime as runtime
 from ._dialect import dialect as dialect
 from ._interface import fill as fill
 from .stmts import Fill as Fill

src/bloqade/shuttle/dialects/init/runtime.py

@@ -0,0 +1,20 @@
+from kirin import interp
+
+from bloqade.shuttle.analysis.runtime import (
+    RuntimeAnalysis,
+    RuntimeFrame,
+)
+
+from ._dialect import dialect
+from .stmts import Fill
+
+
+@dialect.register(key="runtime")
+class HasQuantumRuntimeMethodTable(interp.MethodTable):
+
+    @interp.impl(Fill)
+    def gate(self, interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: Fill):
+        """Handle gate statements and mark the frame as quantum."""
+        frame.is_quantum = True
+        frame.quantum_stmts.add(stmt)
+        return ()

src/bloqade/shuttle/dialects/init/stmts.py

@@ -10,11 +10,7 @@
 class Fill(ir.Statement):
     name = "fill"
 
-    traits = frozenset(
-        {
-            lowering.FromPythonCall(),
-        }
-    )
+    traits = frozenset({lowering.FromPythonCall()})
     locations: ir.SSAValue = info.argument(
         ilist.IListType[grid.GridType[types.Any, types.Any], types.Any]
     )

src/bloqade/shuttle/dialects/measure/__init__.py

@@ -1,3 +1,4 @@
+from . import runtime as runtime
 from ._dialect import dialect as dialect
 from ._interface import measure as measure
 from .stmts import Measure as Measure, New as New

src/bloqade/shuttle/dialects/measure/runtime.py

@@ -0,0 +1,20 @@
+from kirin import interp
+
+from bloqade.shuttle.analysis.runtime import (
+    RuntimeAnalysis,
+    RuntimeFrame,
+)
+
+from ._dialect import dialect
+from .stmts import Measure
+
+
+@dialect.register(key="runtime")
+class HasQuantumRuntimeMethodTable(interp.MethodTable):
+
+    @interp.impl(Measure)
+    def gate(self, _interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: Measure):
+        """Handle gate statements and mark the frame as quantum."""
+        frame.is_quantum = True
+        frame.quantum_stmts.add(stmt)
+        return (_interp.lattice.top(),)

src/bloqade/shuttle/dialects/path/__init__.py

@@ -1,3 +1,4 @@
+from . import runtime as runtime
 from ._dialect import dialect as dialect
 from .concrete import PathInterpreter as PathInterpreter
 from .constprop import ConstProp as ConstProp

src/bloqade/shuttle/dialects/path/runtime.py

@@ -0,0 +1,22 @@
+from kirin import interp
+
+from bloqade.shuttle.analysis.runtime import (
+    RuntimeAnalysis,
+    RuntimeFrame,
+)
+
+from ._dialect import dialect
+from .stmts import Play
+
+
+@dialect.register(key="runtime")
+class HasQuantumRuntimeMethodTable(interp.MethodTable):
+
+    @interp.impl(Play)
+    def gate(
+        self, interp: RuntimeAnalysis, frame: RuntimeFrame, stmt: Play
+    ) -> interp.StatementResult[RuntimeFrame]:
+        """Handle gate statements and mark the frame as quantum."""
+        frame.is_quantum = True
+        frame.quantum_stmts.add(stmt)
+        return ()