move qasm2 from kirin-circuit (#1)

This PR moves the QASM2 implementation into this new namespace package
with some minor cleanups:

- qrack is no longer living with dialects, it is in independent module
so we can install it conditionally with extra `qrack`
- implement qrack interpreter for QASM2

Co-authored-by: Roger-luo <[email protected]>
Co-authored-by: kaihsin <[email protected]>
Co-authored-by: weinbe58 <[email protected]>
Co-authored-by: johnzl-777 <[email protected]>

Author: 4 people

.github/workflows/ci.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Set up Python ${{ matrix.python-version }}
         run: uv python install ${{ matrix.python-version }}
       - name: Install the project
-        run: uv sync --all-extras --dev
+        run: uv sync --all-extras --dev --index="https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/kirin/simple/ https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/quera-pypi-algo/simple/"
       - name: Run tests
         # For example, using `pytest`
         run: uv run just coverage

.github/workflows/devdoc.yml

@@ -24,7 +24,7 @@ jobs:
           enable-cache: true
           cache-dependency-glob: "uv.lock"
       - name: Install Documentation dependencies
-        run: uv sync --group doc
+        run: uv sync --group doc --index="https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/kirin/simple/ https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/quera-pypi-algo/simple/"
       - name: Set up build cache
         uses: actions/cache@v4
         id: cache

.github/workflows/doc.yml

@@ -24,7 +24,7 @@ jobs:
           enable-cache: true
           cache-dependency-glob: "uv.lock"
       - name: Install Documentation dependencies
-        run: uv sync --group doc
+        run: uv sync --group doc --index="https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/kirin/simple/ https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/quera-pypi-algo/simple/"
       - name: Set up build cache
         uses: actions/cache@v4
         id: cache

.github/workflows/pub_doc.yml

@@ -24,7 +24,7 @@ jobs:
           enable-cache: true
           cache-dependency-glob: "uv.lock"
       - name: Install Documentation dependencies
-        run: uv sync --group doc
+        run: uv sync --group doc --index="https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/kirin/simple/ https://${{ secrets.JFROG_USER}}:${{ secrets.JFROG_TOKEN }}@quera.jfrog.io/artifactory/api/pypi/quera-pypi-algo/simple/"
       - name: Set up build cache
         uses: actions/cache@v4
         id: cache

.gitignore

@@ -185,3 +185,4 @@ main.py
 .ruff_cache
 .python-version
 !package.json
+!src/**/**/main.py

.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v5.0.0
+  hooks:
+    - id: check-yaml
+      args: ['--unsafe']
+    - id: end-of-file-fixer
+    - id: trailing-whitespace
+- repo: https://github.com/pycqa/isort
+  rev: 5.13.2
+  hooks:
+    - id: isort
+      name: isort (python)
+- repo: https://github.com/psf/black
+  rev: 24.10.0
+  hooks:
+    - id: black
+- repo: https://github.com/charliermarsh/ruff-pre-commit
+  # Ruff version.
+  rev: "v0.9.2"
+  hooks:
+    - id: ruff

justfile

@@ -0,0 +1,22 @@
+coverage-run:
+    coverage run -m pytest test
+
+coverage-xml:
+    coverage xml
+
+coverage-html:
+    coverage html
+
+coverage-report:
+    coverage report
+
+coverage-open:
+    open htmlcov/index.html
+
+coverage: coverage-run coverage-xml coverage-report
+
+doc:
+    mkdocs serve
+
+doc-build:
+    mkdocs build

pyproject.toml

@@ -11,7 +11,15 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "kirin-toolchain>=0.9.0",
+    "kirin-toolchain>=0.9.1",
+]
+
+[project.optional-dependencies]
+qasm2 = [
+    "lark>=1.2.2",
+]
+pyqrack-cpu = [
+    "pyqrack-cpu>=1.34.9",
 ]
 
 [build-system]
@@ -89,4 +97,4 @@ exclude = [
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
 [tool.coverage.run]
-include = ["src/kirin/*"]
+include = ["src/bloqade/*"]

src/bloqade/qasm2/analysis/__init__.py renamed to src/bloqade/analysis/__init__.py

@@ -0,0 +1,10 @@
+from . import impls as impls
+from .lattice import (
+    Address as Address,
+    NotQubit as NotQubit,
+    AddressReg as AddressReg,
+    AnyAddress as AnyAddress,
+    AddressQubit as AddressQubit,
+    AddressTuple as AddressTuple,
+)
+from .analysis import AddressAnalysis as AddressAnalysis

src/bloqade/analysis/address/__init__.py

@@ -0,0 +1,96 @@
+from typing import TypeVar, Iterable
+
+from kirin import ir, types, interp
+from bloqade.types import QubitType
+from kirin.analysis import Forward, const
+from kirin.dialects import cf, py, func, ilist
+from kirin.exceptions import InterpreterError
+from kirin.analysis.forward import ForwardFrame
+
+from .lattice import Address
+
+
+class AddressAnalysis(Forward[Address]):
+    keys = ["qubit.address"]
+    lattice = Address
+
+    def __init__(
+        self,
+        dialects: ir.DialectGroup | Iterable[ir.Dialect],
+        *,
+        fuel: int | None = None,
+        save_all_ssa: bool = False,
+        max_depth: int = 128,
+        max_python_recursion_depth: int = 8192,
+    ):
+        super().__init__(
+            dialects,
+            fuel=fuel,
+            save_all_ssa=save_all_ssa,
+            max_depth=max_depth,
+            max_python_recursion_depth=max_python_recursion_depth,
+        )
+        self.next_address: int = 0
+        self.constprop_results: dict[ir.SSAValue, const.JointResult] = {}
+
+    def clear(self):
+        self.next_address = 0
+        self.constprop_results.clear()
+
+    @property
+    def qubit_count(self) -> int:
+        return self.next_address
+
+    T = TypeVar("T")
+
+    def get_const_value(self, typ: type[T], value: ir.SSAValue) -> T:
+        if isinstance(value.type, types.Hinted) and isinstance(
+            value.type.data, const.Value
+        ):
+            data = value.type.data.data
+            if isinstance(data, typ):
+                return value.type.data.data
+            raise InterpreterError(
+                f"Expected constant value <type = {typ}>, got {data}"
+            )
+        raise InterpreterError(f"Expected constant value <type = {typ}>, got {value}")
+
+    def run_stmt_fallback(
+        self, frame: ForwardFrame[Address, None], stmt: ir.Statement
+    ) -> tuple[Address, ...] | interp.SpecialResult[Address]:
+        return tuple(
+            (
+                self.lattice.top()
+                if result.type.is_subseteq(QubitType)
+                else self.lattice.bottom()
+            )
+            for result in stmt.results
+        )
+
+    def should_exec_stmt(self, stmt: ir.Statement):
+        return (
+            stmt.has_trait(ir.ConstantLike)
+            or stmt.dialect in self.dialects.data
+            or isinstance(
+                stmt,
+                (
+                    func.Return,
+                    func.Invoke,
+                    py.tuple.New,
+                    ilist.New,
+                    py.GetItem,
+                    py.Alias,
+                    py.Add,
+                    cf.Branch,
+                    cf.ConditionalBranch,
+                ),
+            )
+        )
+
+    def run_method(
+        self, method: ir.Method, args: tuple[Address, ...]
+    ) -> Address | interp.Err[Address]:
+        if len(self.state.frames) >= self.max_depth:
+            raise InterpreterError("maximum recursion depth exceeded")
+        # 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)

src/bloqade/analysis/address/analysis.py

@@ -0,0 +1,94 @@
+"""
+qubit.address method table for a few builtin dialects.
+"""
+
+from kirin import interp
+from kirin.dialects import cf, py, func, ilist
+
+from .lattice import NotQubit, AddressReg, AddressQubit, AddressTuple
+from .analysis import AddressAnalysis
+
+
+@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):
+        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(),)
+
+
+@py.tuple.dialect.register(key="qubit.address")
+class PyTuple(interp.MethodTable):
+    @interp.impl(py.tuple.New)
+    def new_tuple(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: py.tuple.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.args)),)
+
+
+@ilist.dialect.register(key="qubit.address")
+class IList(interp.MethodTable):
+    @interp.impl(ilist.New)
+    def new_ilist(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: ilist.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.args)),)
+
+
+@py.list.dialect.register(key="qubit.address")
+class PyList(interp.MethodTable):
+    @interp.impl(py.list.New)
+    def new_ilist(
+        self,
+        interp: AddressAnalysis,
+        frame: interp.Frame,
+        stmt: py.list.New,
+    ):
+        return (AddressTuple(frame.get_values(stmt.args)),)
+
+
+@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):
+        idx = interp.get_const_value(int, stmt.index)
+        obj = frame.get(stmt.obj)
+        if isinstance(obj, AddressTuple):
+            return (obj.data[idx],)
+        elif isinstance(obj, AddressReg):
+            return (AddressQubit(obj.data[idx]),)
+        else:
+            return (NotQubit(),)
+
+
+@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):
+        return (frame.get(stmt.value),)
+
+
+@func.dialect.register(key="qubit.address")
+class Func(interp.MethodTable):
+    @interp.impl(func.Return)
+    def return_(self, _: AddressAnalysis, frame: interp.Frame, stmt: func.Return):
+        return interp.ReturnValue(frame.get(stmt.value))
+
+
+@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

src/bloqade/analysis/address/impls.py

@@ -0,0 +1,74 @@
+from typing import Sequence, final
+from dataclasses import dataclass
+
+from kirin.lattice import (
+    SingletonMeta,
+    BoundedLattice,
+    SimpleJoinMixin,
+    SimpleMeetMixin,
+)
+
+
+@dataclass
+class Address(
+    SimpleJoinMixin["Address"],
+    SimpleMeetMixin["Address"],
+    BoundedLattice["Address"],
+):
+
+    @classmethod
+    def bottom(cls) -> "Address":
+        return NotQubit()
+
+    @classmethod
+    def top(cls) -> "Address":
+        return AnyAddress()
+
+
+@final
+@dataclass
+class NotQubit(Address, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Address) -> bool:
+        return True
+
+
+@final
+@dataclass
+class AnyAddress(Address, metaclass=SingletonMeta):
+
+    def is_subseteq(self, other: Address) -> bool:
+        return isinstance(other, AnyAddress)
+
+
+@final
+@dataclass
+class AddressTuple(Address):
+    data: tuple[Address, ...]
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressTuple):
+            return all(a.is_subseteq(b) for a, b in zip(self.data, other.data))
+        return False
+
+
+@final
+@dataclass
+class AddressReg(Address):
+    data: Sequence[int]
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressReg):
+            return self.data == other.data
+        return False
+
+
+@final
+@dataclass
+class AddressQubit(Address):
+    data: int
+
+    def is_subseteq(self, other: Address) -> bool:
+        if isinstance(other, AddressQubit):
+            return self.data == other.data
+        return False