bloqade-circuit 0.6.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/analysis/address/__init__.py

@@ -1,11 +1,15 @@
 from . import impls as impls
 from .lattice import (
+    Bottom as Bottom,
     Address as Address,
-    NotQubit as NotQubit,
+    Unknown as Unknown,
     AddressReg as AddressReg,
-    AnyAddress as AnyAddress,
-    AddressWire as AddressWire,
+    UnknownReg as UnknownReg,
+    ConstResult as ConstResult,
     AddressQubit as AddressQubit,
-    AddressTuple as AddressTuple,
+    PartialIList as PartialIList,
+    PartialTuple as PartialTuple,
+    UnknownQubit as UnknownQubit,
+    PartialLambda as PartialLambda,
 )
 from .analysis import AddressAnalysis as AddressAnalysis

bloqade/analysis/address/analysis.py

@@ -1,13 +1,13 @@
-from typing import TypeVar
+from typing import Any, Type, TypeVar
 from dataclasses import field
 
-from kirin import ir, interp
+from kirin import ir, types, interp
 from kirin.analysis import Forward, const
+from kirin.dialects.ilist import IList
 from kirin.analysis.forward import ForwardFrame
+from kirin.analysis.const.lattice import PartialLambda
 
-from bloqade.types import QubitType
-
-from .lattice import Address
+from .lattice import Address, AddressReg, ConstResult, PartialIList, PartialTuple
 
 
 class AddressAnalysis(Forward[Address]):
@@ -15,13 +15,16 @@ class AddressAnalysis(Forward[Address]):
     This analysis pass can be used to track the global addresses of qubits and wires.
     """
 
-    keys = ["qubit.address"]
+    keys = ("qubit.address",)
+    _const_prop: const.Propagate
     lattice = Address
     next_address: int = field(init=False)
 
     def initialize(self):
         super().initialize()
         self.next_address: int = 0
+        self._const_prop = const.Propagate(self.dialects)
+        self._const_prop.initialize()
         return self
 
     @property
@@ -31,30 +34,117 @@ class AddressAnalysis(Forward[Address]):
 
     T = TypeVar("T")
 
-    def get_const_value(self, typ: type[T], value: ir.SSAValue) -> T:
-        if isinstance(hint := value.hints.get("const"), const.Value):
-            data = hint.data
-            if isinstance(data, typ):
-                return hint.data
-            raise interp.InterpreterError(
-                f"Expected constant value <type = {typ}>, got {data}"
-            )
-        raise interp.InterpreterError(
-            f"Expected constant value <type = {typ}>, got {value}"
-        )
-
-    def eval_stmt_fallback(
-        self, frame: ForwardFrame[Address], stmt: ir.Statement
-    ) -> tuple[Address, ...] | interp.SpecialValue[Address]:
-        return tuple(
-            (
-                self.lattice.top()
-                if result.type.is_subseteq(QubitType)
-                else self.lattice.bottom()
-            )
-            for result in stmt.results
-        )
-
-    def run_method(self, method: ir.Method, args: tuple[Address, ...]):
-        # NOTE: we do not support dynamic calls here, thus no need to propagate method object
-        return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+    def to_address(self, result: const.Result):
+        return ConstResult(result)
+
+    def try_eval_const_prop(
+        self,
+        frame: ForwardFrame[Address],
+        stmt: ir.Statement,
+        args: tuple[ConstResult, ...],
+    ) -> interp.StatementResult[Address]:
+        _frame = self._const_prop.initialize_frame(frame.code)
+        _frame.set_values(stmt.args, tuple(x.result for x in args))
+        result = self._const_prop.frame_eval(_frame, stmt)
+
+        match result:
+            case interp.ReturnValue(constant_ret):
+                return interp.ReturnValue(self.to_address(constant_ret))
+            case interp.YieldValue(constant_values):
+                return interp.YieldValue(tuple(map(self.to_address, constant_values)))
+            case interp.Successor(block, block_args):
+                return interp.Successor(block, *map(self.to_address, block_args))
+            case tuple():
+                return tuple(map(self.to_address, result))
+            case _:
+                return result
+
+    def unpack_iterable(self, iterable: Address):
+        """Extract the values of a container lattice element.
+
+        Args:
+            iterable: The lattice element representing a container.
+
+        Returns:
+            A tuple of the container type and the contained values.
+
+        """
+
+        def from_constant(constant: const.Result) -> Address:
+            return ConstResult(constant)
+
+        def from_literal(literal: Any) -> Address:
+            return ConstResult(const.Value(literal))
+
+        match iterable:
+            case PartialIList(data):
+                return PartialIList, data
+            case PartialTuple(data):
+                return PartialTuple, data
+            case AddressReg():
+                return PartialIList, iterable.qubits
+            case ConstResult(const.Value(IList() as data)):
+                return PartialIList, tuple(map(from_literal, data))
+            case ConstResult(const.Value(tuple() as data)):
+                return PartialTuple, tuple(map(from_literal, data))
+            case ConstResult(const.PartialTuple(data)):
+                return PartialTuple, tuple(map(from_constant, data))
+            case _:
+                return None, ()
+
+    def run_lattice(
+        self,
+        callee: Address,
+        inputs: tuple[Address, ...],
+        keys: tuple[str, ...],
+        kwargs: tuple[Address, ...],
+    ) -> Address:
+        """Run a callable lattice element with the given inputs and keyword arguments.
+
+        Args:
+            callee (Address): The lattice element representing the callable.
+            inputs (tuple[Address, ...]): The input lattice elements.
+            kwargs (tuple[str, ...]): The keyword argument names.
+
+        Returns:
+            Address: The resulting lattice element after invoking the callable.
+
+        """
+
+        match callee:
+            case PartialLambda(code=code):
+                _, ret = self.call(
+                    code, callee, *inputs, **{k: v for k, v in zip(keys, kwargs)}
+                )
+            case ConstResult(const.Value(ir.Method() as method)):
+                _, ret = self.call(
+                    method.code,
+                    self.method_self(method),
+                    *inputs,
+                    **{k: v for k, v in zip(keys, kwargs)},
+                )
+                return ret
+            case _:
+                return Address.top()
+
+    def get_const_value(self, addr: Address, typ: Type[T]) -> T | None:
+        if not isinstance(addr, ConstResult):
+            return None
+
+        if not isinstance(result := addr.result, const.Value):
+            return None
+
+        if not isinstance(value := result.data, typ):
+            return None
+
+        return value
+
+    def eval_fallback(self, frame: ForwardFrame[Address], node: ir.Statement):
+        args = frame.get_values(node.args)
+        if types.is_tuple_of(args, ConstResult):
+            return self.try_eval_const_prop(frame, node, args)
+
+        return tuple(Address.from_type(result.type) for result in node.results)
+
+    def method_self(self, method: ir.Method) -> Address:
+        return ConstResult(const.Value(method))

bloqade/analysis/address/impls.py

@@ -2,32 +2,56 @@
 qubit.address method table for a few builtin dialects.
 """
 
-from kirin import interp
+from itertools import chain
+
+from kirin import ir, interp
 from kirin.analysis import ForwardFrame, const
 from kirin.dialects import cf, py, scf, func, ilist
 
 from .lattice import (
     Address,
-    NotQubit,
-    AddressReg,
-    AddressQubit,
-    AddressTuple,
+    ConstResult,
+    PartialIList,
+    PartialTuple,
+    PartialLambda,
 )
 from .analysis import AddressAnalysis
 
 
+@py.constant.dialect.register(key="qubit.address")
+class PyConstant(interp.MethodTable):
+    @interp.impl(py.Constant)
+    def constant(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: py.Constant,
+    ):
+        return (ConstResult(const.Value(stmt.value.unwrap())),)
+
+
 @py.binop.dialect.register(key="qubit.address")
 class PyBinOp(interp.MethodTable):
-
     @interp.impl(py.Add)
-    def add(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.Add):
+    def add(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: py.Add,
+    ):
         lhs = frame.get(stmt.lhs)
         rhs = frame.get(stmt.rhs)
 
-        if isinstance(lhs, AddressTuple) and isinstance(rhs, AddressTuple):
-            return (AddressTuple(data=lhs.data + rhs.data),)
-        else:
-            return (NotQubit(),)
+        lhs_type, lhs_values = interp_.unpack_iterable(lhs)
+        rhs_type, rhs_values = interp_.unpack_iterable(rhs)
+
+        if lhs_type is None or rhs_type is None:
+            return (Address.top(),)
+
+        if lhs_type is not rhs_type:
+            return (Address.bottom(),)
+
+        return (lhs_type(tuple(chain(lhs_values, rhs_values))),)
 
 
 @py.tuple.dialect.register(key="qubit.address")
@@ -35,110 +59,299 @@ class PyTuple(interp.MethodTable):
     @interp.impl(py.tuple.New)
     def new_tuple(
         self,
-        interp: AddressAnalysis,
-        frame: interp.Frame,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
         stmt: py.tuple.New,
     ):
-        return (AddressTuple(frame.get_values(stmt.args)),)
+        return (PartialTuple(frame.get_values(stmt.args)),)
 
 
 @ilist.dialect.register(key="qubit.address")
-class IList(interp.MethodTable):
+class IListMethods(interp.MethodTable):
     @interp.impl(ilist.New)
     def new_ilist(
         self,
-        interp: AddressAnalysis,
-        frame: interp.Frame,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
         stmt: ilist.New,
     ):
-        return (AddressTuple(frame.get_values(stmt.values)),)
-
+        return (PartialIList(frame.get_values(stmt.args)),)
 
-@py.list.dialect.register(key="qubit.address")
-class PyList(interp.MethodTable):
-    @interp.impl(py.list.New)
-    def new_ilist(
+    @interp.impl(ilist.ForEach)
+    @interp.impl(ilist.Map)
+    def map_(
         self,
-        interp: AddressAnalysis,
-        frame: interp.Frame,
-        stmt: py.list.New,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: ilist.Map | ilist.ForEach,
     ):
-        return (AddressTuple(frame.get_values(stmt.args)),)
+        fn = frame.get(stmt.fn)
+        collection = frame.get(stmt.collection)
+        collection_type, values = interp_.unpack_iterable(collection)
+
+        if collection_type is None:
+            return (Address.top(),)
+
+        if collection_type is not PartialIList:
+            return (Address.bottom(),)
+
+        results = []
+        for ele in values:
+            ret = interp_.run_lattice(fn, (ele,), (), ())
+            results.append(ret)
+
+        if isinstance(stmt, ilist.Map):
+            return (PartialIList(tuple(results)),)
+
+
+@py.len.dialect.register(key="qubit.address")
+class PyLen(interp.MethodTable):
+    @interp.impl(py.Len)
+    def len_(
+        self, interp_: AddressAnalysis, frame: ForwardFrame[Address], stmt: py.Len
+    ):
+        obj = frame.get(stmt.value)
+        _, values = interp_.unpack_iterable(obj)
+
+        if values is None:
+            return (Address.top(),)
+
+        return (ConstResult(const.Value(len(values))),)
 
 
 @py.indexing.dialect.register(key="qubit.address")
 class PyIndexing(interp.MethodTable):
     @interp.impl(py.GetItem)
-    def getitem(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.GetItem):
-
+    def getitem(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: py.GetItem,
+    ):
         # determine if the index is an int constant
         # or a slice
-        hint = stmt.index.hints.get("const")
-        if hint is None:
-            return (NotQubit(),)
-
-        if isinstance(hint, const.Value):
-            idx = hint.data
-        elif isinstance(hint, slice):
-            idx = hint
-        else:
-            return (NotQubit(),)
-
-        # The object being indexed into
         obj = frame.get(stmt.obj)
-        # The `data` attributes holds onto other Address types
-        # so we just extract that here
-        if isinstance(obj, AddressTuple):
-            return (obj.data[idx],)
-        # If idx is an integer index into an AddressReg,
-        # then it's safe to assume a single qubit is being accessed.
-        # On the other hand, if it's a slice, we return
-        # a new AddressReg to preserve the new sequence.
-        elif isinstance(obj, AddressReg):
-            if isinstance(idx, slice):
-                return (AddressReg(data=obj.data[idx]),)
-            if isinstance(idx, int):
-                return (AddressQubit(obj.data[idx]),)
-        else:
-            return (NotQubit(),)
+        index = frame.get(stmt.index)
+
+        typ, values = interp_.unpack_iterable(obj)
+        if typ is None:
+            return (Address.top(),)
+
+        int_index = interp_.get_const_value(index, int)
+        if int_index is not None:
+            return (values[int_index],)
+
+        slice_index = interp_.get_const_value(index, slice)
+        if slice_index is not None:
+            return (typ(values[slice_index]),)
+
+        return (Address.top(),)
 
 
 @py.assign.dialect.register(key="qubit.address")
 class PyAssign(interp.MethodTable):
     @interp.impl(py.Alias)
-    def alias(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.Alias):
+    def alias(
+        self,
+        interp: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: py.Alias,
+    ):
         return (frame.get(stmt.value),)
 
 
+# TODO: look for abstract method table for func.
 @func.dialect.register(key="qubit.address")
 class Func(interp.MethodTable):
     @interp.impl(func.Return)
-    def return_(self, _: AddressAnalysis, frame: interp.Frame, stmt: func.Return):
+    def return_(
+        self,
+        _: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: func.Return,
+    ):
         return interp.ReturnValue(frame.get(stmt.value))
 
     # TODO: replace with the generic implementation
     @interp.impl(func.Invoke)
-    def invoke(self, interp_: AddressAnalysis, frame: interp.Frame, stmt: func.Invoke):
-        _, ret = interp_.run_method(
-            stmt.callee,
-            interp_.permute_values(
-                stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
-            ),
+    def invoke(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: func.Invoke,
+    ):
+        _, ret = interp_.call(
+            stmt.callee.code,
+            interp_.method_self(stmt.callee),
+            *frame.get_values(stmt.inputs),
         )
+
         return (ret,)
 
-    # TODO: support lambda?
+    @interp.impl(func.Lambda)
+    def lambda_(
+        self,
+        inter_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: func.Lambda,
+    ):
+        arg_names = [
+            arg.name or str(idx) for idx, arg in enumerate(stmt.body.blocks[0].args)
+        ]
+        return (
+            PartialLambda(
+                arg_names,
+                stmt,
+                frame.get_values(stmt.captured),
+            ),
+        )
+
+    @interp.impl(func.Call)
+    def call(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: func.Call,
+    ):
+        result = interp_.run_lattice(
+            frame.get(stmt.callee),
+            frame.get_values(stmt.inputs),
+            stmt.keys,
+            frame.get_values(stmt.kwargs),
+        )
+        return (result,)
+
+    @interp.impl(func.GetField)
+    def get_field(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: func.GetField,
+    ):
+        self_mt = frame.get(stmt.obj)
+        match self_mt:
+            case PartialLambda(captured=captured):
+                return (captured[stmt.field],)
+            case ConstResult(const.Value(ir.Method() as mt)):
+                return (ConstResult(const.Value(mt.fields[stmt.field])),)
+
+        return (Address.top(),)
 
 
 @cf.dialect.register(key="qubit.address")
-class Cf(cf.typeinfer.TypeInfer):
-    # NOTE: cf just re-use the type infer method table
-    # it's the same process as type infer.
-    pass
+class Cf(interp.MethodTable):
+
+    @interp.impl(cf.Branch)
+    def branch(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: cf.Branch,
+    ):
+        frame.worklist.append(
+            interp.Successor(stmt.successor, *frame.get_values(stmt.arguments))
+        )
+        return ()
+
+    @interp.impl(cf.ConditionalBranch)
+    def conditional_branch(
+        self,
+        interp_: const.Propagate,
+        frame: ForwardFrame[Address],
+        stmt: cf.ConditionalBranch,
+    ):
+        address_cond = frame.get(stmt.cond)
+
+        if isinstance(address_cond, ConstResult) and isinstance(
+            cond := address_cond.result, const.Value
+        ):
+            else_successor = interp.Successor(
+                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+            )
+            then_successor = interp.Successor(
+                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+            )
+            if cond.data:
+                frame.worklist.append(then_successor)
+            else:
+                frame.worklist.append(else_successor)
+        else:
+            frame.entries[stmt.cond] = ConstResult(const.Value(True))
+            then_successor = interp.Successor(
+                stmt.then_successor, *frame.get_values(stmt.then_arguments)
+            )
+            frame.worklist.append(then_successor)
+
+            frame.entries[stmt.cond] = ConstResult(const.Value(False))
+            else_successor = interp.Successor(
+                stmt.else_successor, *frame.get_values(stmt.else_arguments)
+            )
+            frame.worklist.append(else_successor)
+
+            frame.entries[stmt.cond] = address_cond
+        return ()
 
 
 @scf.dialect.register(key="qubit.address")
-class Scf(scf.absint.Methods):
+class Scf(interp.MethodTable):
+    @interp.impl(scf.Yield)
+    def yield_(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: scf.Yield,
+    ):
+        return interp.YieldValue(frame.get_values(stmt.values))
+
+    @interp.impl(scf.IfElse)
+    def ifelse(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: scf.IfElse,
+    ):
+        address_cond = frame.get(stmt.cond)
+        # run specific branch
+        if isinstance(address_cond, ConstResult) and isinstance(
+            const_cond := address_cond.result, const.Value
+        ):
+            body = stmt.then_body if const_cond.data else stmt.else_body
+            with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+                ret = interp_.frame_call_region(body_frame, stmt, body, address_cond)
+                # interp_.set_values(frame, body_frame.entries.keys(), body_frame.entries.values())
+                return ret
+        else:
+            # run both branches
+            with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
+                then_results = interp_.frame_call_region(
+                    then_frame,
+                    stmt,
+                    stmt.then_body,
+                    address_cond,
+                )
+                frame.set_values(then_frame.entries.keys(), then_frame.entries.values())
+
+            with interp_.new_frame(stmt, has_parent_access=True) as else_frame:
+                else_results = interp_.frame_call_region(
+                    else_frame,
+                    stmt,
+                    stmt.else_body,
+                    address_cond,
+                )
+                frame.set_values(else_frame.entries.keys(), else_frame.entries.values())
+            # TODO: pick the non-return value
+            if isinstance(then_results, interp.ReturnValue) and isinstance(
+                else_results, interp.ReturnValue
+            ):
+                return interp.ReturnValue(then_results.value.join(else_results.value))
+            elif isinstance(then_results, interp.ReturnValue):
+                ret = else_results
+            elif isinstance(else_results, interp.ReturnValue):
+                ret = then_results
+            else:
+                ret = interp_.join_results(then_results, else_results)
+
+            return ret
 
     @interp.impl(scf.For)
     def for_loop(
@@ -147,32 +360,22 @@ class Scf(scf.absint.Methods):
         frame: ForwardFrame[Address],
         stmt: scf.For,
     ):
-        if not isinstance(hint := stmt.iterable.hints.get("const"), const.Value):
-            return interp_.eval_stmt_fallback(frame, stmt)
-
-        iterable = hint.data
         loop_vars = frame.get_values(stmt.initializers)
-        body_block = stmt.body.blocks[0]
-        block_args = body_block.args
-
-        # NOTE: we need to actually run iteration in case there are
-        # new allocations/re-assign in the loop body.
-        for _ in iterable:
-            with interp_.new_frame(stmt) as body_frame:
-                body_frame.entries.update(frame.entries)
-                body_frame.set_values(
-                    block_args,
-                    (NotQubit(),) + loop_vars,
+        iter_type, iterable = interp_.unpack_iterable(frame.get(stmt.iterable))
+
+        if iter_type is None:
+            return interp_.eval_fallback(frame, stmt)
+
+        for value in iterable:
+            with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+                loop_vars = interp_.frame_call_region(
+                    body_frame, stmt, stmt.body, value, *loop_vars
                 )
-                loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body, ())
 
             if loop_vars is None:
                 loop_vars = ()
+
             elif isinstance(loop_vars, interp.ReturnValue):
                 return loop_vars
 
-        if isinstance(body_block.last_stmt, func.Return):
-            frame.worklist.append(interp.Successor(body_block, NotQubit(), *loop_vars))
-            return  # if terminate is Return, there is no result
-
         return loop_vars