Squin Dialect Address Analysis (#4)

* Move from bloqade to bloqade-circuit WIP on analysis

* Saving comments/work from last sync with Roger

* completed unit tests

* fix isort

* fix unit test, remove leftover print statements

* remove commented out code for wire dialect wrapping, not possible to handle multiple return results

Author: johnzl-777

src/bloqade/analysis/address/__init__.py

@@ -4,6 +4,7 @@
     NotQubit as NotQubit,
     AddressReg as AddressReg,
     AnyAddress as AnyAddress,
+    AddressWire as AddressWire,
     AddressQubit as AddressQubit,
     AddressTuple as AddressTuple,
 )

src/bloqade/analysis/address/analysis.py

@@ -11,7 +11,9 @@
 
 
 class AddressAnalysis(Forward[Address]):
-    """This analysis pass can be used to track the global addresses of qubits."""
+    """
+    This analysis pass can be used to track the global addresses of qubits and wires.
+    """
 
     keys = ["qubit.address"]
     lattice = Address

src/bloqade/analysis/address/impls.py

@@ -6,7 +6,16 @@
 from kirin.analysis import ForwardFrame, const
 from kirin.dialects import cf, py, scf, func, ilist
 
-from .lattice import Address, NotQubit, AddressReg, AddressQubit, AddressTuple
+from bloqade import squin
+
+from .lattice import (
+    Address,
+    NotQubit,
+    AddressReg,
+    AddressWire,
+    AddressQubit,
+    AddressTuple,
+)
 from .analysis import AddressAnalysis
 
 
@@ -64,10 +73,16 @@ def new_ilist(
 class PyIndexing(interp.MethodTable):
     @interp.impl(py.GetItem)
     def getitem(self, interp: AddressAnalysis, frame: interp.Frame, stmt: py.GetItem):
+        # Integer index into the thing being indexed
         idx = interp.get_const_value(int, stmt.index)
+        # 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],)
+        # an AddressReg is guaranteed to just have some sequence
+        # of integers which is directly pluggable to AddressQubit
         elif isinstance(obj, AddressReg):
             return (AddressQubit(obj.data[idx]),)
         else:
@@ -147,3 +162,67 @@ def for_loop(
             return  # if terminate is Return, there is no result
 
         return loop_vars
+
+
+# Address lattice elements we can work with:
+## NotQubit (bottom), AnyAddress (top)
+
+## AddressTuple -> data: tuple[Address, ...]
+### Recursive type, could contain itself or other variants
+### This pops up in cases where you can have an IList/Tuple
+### That contains elements that could be other Address types
+
+## AddressReg -> data: Sequence[int]
+### specific to creation of a register of qubits
+
+## AddressQubit -> data: int
+### Base qubit address type
+
+
+@squin.wire.dialect.register(key="qubit.address")
+class SquinWireMethodTable(interp.MethodTable):
+
+    @interp.impl(squin.wire.Unwrap)
+    def unwrap(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: squin.wire.Unwrap,
+    ):
+
+        origin_qubit = frame.get(stmt.qubit)
+
+        return (AddressWire(origin_qubit=origin_qubit),)
+
+    @interp.impl(squin.wire.Apply)
+    def apply(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: squin.wire.Apply,
+    ):
+
+        origin_qubits = tuple(
+            [frame.get(input_elem).origin_qubit for input_elem in stmt.inputs]
+        )
+        new_address_wires = tuple(
+            [AddressWire(origin_qubit=origin_qubit) for origin_qubit in origin_qubits]
+        )
+        return new_address_wires
+
+
+@squin.qubit.dialect.register(key="qubit.address")
+class SquinQubitMethodTable(interp.MethodTable):
+
+    # This can be treated like a QRegNew impl
+    @interp.impl(squin.qubit.New)
+    def new(
+        self,
+        interp_: AddressAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: squin.qubit.New,
+    ):
+        n_qubits = interp_.get_const_value(int, stmt.n_qubits)
+        addr = AddressReg(range(interp_.next_address, interp_.next_address + n_qubits))
+        interp_.next_address += n_qubits
+        return (addr,)

src/bloqade/analysis/address/lattice.py

@@ -72,3 +72,14 @@ def is_subseteq(self, other: Address) -> bool:
         if isinstance(other, AddressQubit):
             return self.data == other.data
         return False
+
+
+@final
+@dataclass
+class AddressWire(Address):
+    origin_qubit: AddressQubit
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressWire):
+            return self.origin_qubit == self.origin_qubit
+        return False

src/bloqade/qasm2/dialects/glob.py

@@ -5,7 +5,7 @@
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QRegType
 from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
-from bloqade.analysis.schedule import DagScheduleAnalysis
+from bloqade.squin.analysis.schedule import DagScheduleAnalysis
 
 dialect = ir.Dialect("qasm2.glob")
 

src/bloqade/qasm2/dialects/parallel.py

@@ -8,7 +8,7 @@
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QubitType
 from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
-from bloqade.analysis.schedule import DagScheduleAnalysis
+from bloqade.squin.analysis.schedule import DagScheduleAnalysis
 
 dialect = ir.Dialect("qasm2.parallel")
 

src/bloqade/qasm2/dialects/uop/schedule.py

@@ -1,7 +1,7 @@
 from kirin import interp
 from kirin.analysis import ForwardFrame
 
-from bloqade.analysis.schedule import DagScheduleAnalysis
+from bloqade.squin.analysis.schedule import DagScheduleAnalysis
 
 from . import stmts
 from ._dialect import dialect

src/bloqade/qasm2/passes/parallel.py

@@ -20,14 +20,15 @@
 )
 from kirin.analysis import const
 
-from bloqade.analysis import address, schedule
+from bloqade.analysis import address
 from bloqade.qasm2.rewrite import (
     MergePolicyABC,
     ParallelToUOpRule,
     RaiseRegisterRule,
     UOpToParallelRule,
     SimpleOptimalMergePolicy,
 )
+from bloqade.squin.analysis import schedule
 
 
 @dataclass

src/bloqade/qasm2/rewrite/uop_to_parallel.py

@@ -9,7 +9,7 @@
 
 from bloqade.analysis import address
 from bloqade.qasm2.dialects import uop, core, parallel
-from bloqade.analysis.schedule import StmtDag
+from bloqade.squin.analysis.schedule import StmtDag
 
 
 class MergePolicyABC(abc.ABC):