Repository: QuEraComputing/bloqade-pyqrack
Pull request #23: Feat: Dynamic memory allocator. 
Author: login: weinbe58
Created: 2025-03-02T05:01:23Z | Merged: 2025-03-05T05:15:28Z

This PR refactors the `Memory` interface to make it more general and extends the interface to use pyqrack's dynamic memory allocation of qubits. For now there are not deallocate statements inside the qasm2, so no deallocate method is required for the interface. 


@jon-wurtz here you go!


Diff:
diff --git a/src/bloqade/pyqrack/__init__.py b/src/bloqade/pyqrack/__init__.py
index 7cbbc98..66a6ca7 100644
--- a/src/bloqade/pyqrack/__init__.py
+++ b/src/bloqade/pyqrack/__init__.py
@@ -6,7 +6,11 @@
     Measurement as Measurement,
     PyQrackQubit as PyQrackQubit,
 )
-from .base import Memory as Memory, PyQrackInterpreter as PyQrackInterpreter
+from .base import (
+    StackMemory as StackMemory,
+    DynamicMemory as DynamicMemory,
+    PyQrackInterpreter as PyQrackInterpreter,
+)
 from .noise import native as native
 
 # NOTE: The following import is for registering the method tables
diff --git a/src/bloqade/pyqrack/base.py b/src/bloqade/pyqrack/base.py
index b6d2926..48090a0 100644
--- a/src/bloqade/pyqrack/base.py
+++ b/src/bloqade/pyqrack/base.py
@@ -1,3 +1,4 @@
+import abc
 import typing
 from dataclasses import field, dataclass
 
@@ -11,8 +12,20 @@
     from pyqrack import QrackSimulator
 
 
+class MemoryABC(abc.ABC):
+    @abc.abstractmethod
+    def allocate(self, n_qubits: int) -> tuple[int, ...]:
+        """Allocate `n_qubits` qubits and return their ids."""
+        ...
+
+    @abc.abstractmethod
+    def reset(self):
+        """Reset the memory, releasing all qubits."""
+        ...
+
+
 @dataclass
-class Memory:
+class StackMemory(MemoryABC):
     total: int
     allocated: int
     sim_reg: "QrackSimulator"
@@ -30,11 +43,37 @@ def allocate(self, n_qubits: int):
 
         return tuple(range(curr_allocated, self.allocated))
 
+    def reset(self):
+        self.sim_reg.reset_all()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    sim_reg: "QrackSimulator"
+
+    def __post_init__(self):
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+        self.reset()
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+    def reset(self):
+        for qid in self.sim_reg.dump_ids():
+            self.sim_reg.release(qid)
+
 
 @dataclass
 class PyQrackInterpreter(Interpreter):
     keys = ["pyqrack", "main"]
-    memory: Memory = field(kw_only=True)
+    memory: MemoryABC = field(kw_only=True)
     rng_state: np.random.Generator = field(
         default_factory=np.random.default_rng, kw_only=True
     )
@@ -43,5 +82,5 @@ class PyQrackInterpreter(Interpreter):
 
     def initialize(self) -> Self:
         super().initialize()
-        self.memory.allocated = 0  # reset allocated qubits
+        self.memory.reset()  # reset allocated qubits
         return self
diff --git a/src/bloqade/pyqrack/qasm2/core.py b/src/bloqade/pyqrack/qasm2/core.py
index 969a0aa..add6a01 100644
--- a/src/bloqade/pyqrack/qasm2/core.py
+++ b/src/bloqade/pyqrack/qasm2/core.py
@@ -54,12 +54,18 @@ def measure(
         qarg: PyQrackQubit = frame.get(stmt.qarg)
         carg: CBitRef = frame.get(stmt.carg)
         if qarg.is_active():
-            carg.set_value(Measurement(qarg.ref.sim_reg.m(qarg.addr)))
+            carg.set_value(Measurement(qarg.sim_reg.m(qarg.addr)))
         else:
             carg.set_value(interp.loss_m_result)
 
         return ()
 
+    @interp.impl(core.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.Reset):
+        qarg: PyQrackQubit = frame.get(stmt.qarg)
+        qarg.sim_reg.force_m(qarg.addr, 0)
+        return ()
+
     @interp.impl(core.CRegEq)
     def creg_eq(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: core.CRegEq
diff --git a/src/bloqade/pyqrack/target.py b/src/bloqade/pyqrack/target.py
index fb27b51..19f834a 100644
--- a/src/bloqade/pyqrack/target.py
+++ b/src/bloqade/pyqrack/target.py
@@ -1,23 +1,36 @@
-from typing import Any, Dict, List, TypeVar, ParamSpec
+from typing import List, TypeVar, ParamSpec, TypedDict
 from dataclasses import field, dataclass
 
 from kirin import ir
 from pyqrack import QrackSimulator
 from kirin.passes import Fold
-from bloqade.pyqrack.base import Memory, PyQrackInterpreter
+from bloqade.pyqrack.base import StackMemory, DynamicMemory, PyQrackInterpreter
 from bloqade.analysis.address import AnyAddress, AddressAnalysis
 
-Params = ParamSpec("Params")
-RetType = TypeVar("RetType")
+
+class PyQrackOptions(TypedDict):
+    isTensorNetwork: bool
+    isSchmidtDecomposeMulti: bool
+    isSchmidtDecompose: bool
+    isStabilizerHybrid: bool
+    isBinaryDecisionTree: bool
+    isPaged: bool
+    isCpuGpuHybrid: bool
+    isOpenCL: bool
+    isHostPointer: bool
 
 
-def _default_pyqrack_args():
+def _default_pyqrack_args() -> PyQrackOptions:
     return {
         "isTensorNetwork": False,
         "isOpenCL": False,
     }
 
 
+Params = ParamSpec("Params")
+RetType = TypeVar("RetType")
+
+
 @dataclass
 class PyQrack:
     """PyQrack target runtime for Bloqade."""
@@ -26,10 +39,39 @@ class PyQrack:
     """Minimum number of qubits required for the PyQrack simulator.
     Useful when address analysis fails to determine the number of qubits.
     """
-    pyqrack_options: Dict[str, Any] = field(default_factory=_default_pyqrack_args)
+    dynamic_qubits: bool = False
+    """Whether to use dynamic qubit allocation. Cannot use with tensor network simulations."""
+
+    pyqrack_options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
     """Options to pass to the QrackSimulator object, node `qubitCount` will be overwritten."""
 
-    memory: Memory | None = field(init=False, default=None)
+    def _get_interp(self, mt: ir.Method[Params, RetType]):
+        if self.dynamic_qubits:
+            return PyQrackInterpreter(
+                mt.dialects,
+                memory=DynamicMemory(
+                    QrackSimulator(qubitCount=0, **self.pyqrack_options)
+                ),
+            )
+        else:
+            address_analysis = AddressAnalysis(mt.dialects)
+            frame, _ = address_analysis.run_analysis(mt)
+            if self.min_qubits == 0 and any(
+                isinstance(a, AnyAddress) for a in frame.entries.values()
+            ):
+                raise ValueError(
+                    "All addresses must be resolved. Or set min_qubits to a positive integer."
+                )
+
+            num_qubits = max(address_analysis.qubit_count, self.min_qubits)
+            self.pyqrack_options.pop("qubitCount", None)
+            memory = StackMemory(
+                num_qubits,
+                allocated=0,
+                sim_reg=QrackSimulator(qubitCount=num_qubits, **self.pyqrack_options),
+            )
+
+            return PyQrackInterpreter(mt.dialects, memory=memory)
 
     def run(
         self,
@@ -49,24 +91,7 @@ def run(
         """
         fold = Fold(mt.dialects)
         fold(mt)
-        address_analysis = AddressAnalysis(mt.dialects)
-        frame, ret = address_analysis.run_analysis(mt)
-        if self.min_qubits == 0 and any(
-            isinstance(a, AnyAddress) for a in frame.entries.values()
-        ):
-            raise ValueError(
-                "All addresses must be resolved. Or set min_qubits to a positive integer."
-            )
-
-        num_qubits = max(address_analysis.qubit_count, self.min_qubits)
-        self.pyqrack_options.pop("qubitCount", None)
-        self.memory = Memory(
-            num_qubits,
-            allocated=0,
-            sim_reg=QrackSimulator(qubitCount=num_qubits, **self.pyqrack_options),
-        )
-        interpreter = PyQrackInterpreter(mt.dialects, memory=self.memory)
-        return interpreter.run(mt, args, kwargs).expect()
+        return self._get_interp(mt).run(mt, args, kwargs).expect()
 
     def multi_run(
         self,
@@ -87,26 +112,12 @@ def multi_run(
             List of results of the kernel method, one for each shot.
 
         """
+        fold = Fold(mt.dialects)
+        fold(mt)
 
-        address_analysis = AddressAnalysis(mt.dialects)
-        frame, ret = address_analysis.run_analysis(mt)
-        if any(isinstance(a, AnyAddress) for a in frame.entries.values()):
-            raise ValueError("All addresses must be resolved.")
-
-        memory = Memory(
-            address_analysis.next_address,
-            allocated=0,
-            sim_reg=QrackSimulator(
-                qubitCount=address_analysis.next_address,
-                **self.pyqrack_options,
-            ),
-        )
-
+        interpreter = self._get_interp(mt)
         batched_results = []
         for _ in range(_shots):
-            memory.allocated = 0
-            memory.sim_reg.reset_all()
-            interpreter = PyQrackInterpreter(mt.dialects, memory=memory)
             batched_results.append(interpreter.run(mt, args, kwargs).expect())
 
         return batched_results
diff --git a/test/runtime/noise/native/test_loss.py b/test/runtime/noise/native/test_loss.py
index a97e3c3..fb53786 100644
--- a/test/runtime/noise/native/test_loss.py
+++ b/test/runtime/noise/native/test_loss.py
@@ -2,7 +2,7 @@
 
 from bloqade import qasm2
 from bloqade.noise import native
-from bloqade.pyqrack import Memory, PyQrackInterpreter, reg
+from bloqade.pyqrack import StackMemory, PyQrackInterpreter, reg
 
 simulation = qasm2.extended.add(native)
 
@@ -21,7 +21,7 @@ def test_atom_loss(c: qasm2.CReg):
     rng_state = Mock()
     rng_state.uniform.return_value = 0.7
     input = reg.CRegister(1)
-    memory = Memory(total=2, allocated=0, sim_reg=Mock())
+    memory = StackMemory(total=2, allocated=0, sim_reg=Mock())
 
     result: reg.PyQrackReg[Mock] = (
         PyQrackInterpreter(simulation, memory=memory, rng_state=rng_state)
diff --git a/test/runtime/noise/native/test_pauli.py b/test/runtime/noise/native/test_pauli.py
index 0f4754c..9616fb0 100644
--- a/test/runtime/noise/native/test_pauli.py
+++ b/test/runtime/noise/native/test_pauli.py
@@ -3,13 +3,13 @@
 from kirin import ir
 from bloqade import qasm2
 from bloqade.noise import native
-from bloqade.pyqrack import Memory, PyQrackInterpreter
+from bloqade.pyqrack import StackMemory, PyQrackInterpreter
 
 simulation = qasm2.extended.add(native)
 
 
 def run_mock(size, program: ir.Method, rng_state: Mock) -> Mock:
-    memory = Memory(total=2, allocated=0, sim_reg=Mock())
+    memory = StackMemory(total=2, allocated=0, sim_reg=Mock())
 
     PyQrackInterpreter(program.dialects, memory=memory, rng_state=rng_state).run(
         program, ()
diff --git a/test/runtime/test_dyn_memory.py b/test/runtime/test_dyn_memory.py
new file mode 100644
index 0000000..9296743
--- /dev/null
+++ b/test/runtime/test_dyn_memory.py
@@ -0,0 +1,32 @@
+from collections import Counter
+
+from bloqade import qasm2
+from bloqade.pyqrack import PyQrack
+
+
+def test():
+
+    @qasm2.extended
+    def ghz(n: int):
+        q = qasm2.qreg(n)
+        c = qasm2.creg(n)
+
+        qasm2.h(q[0])
+        for i in range(1, n):
+            qasm2.cx(q[0], q[i])
+
+        for i in range(n):
+            qasm2.measure(q[i], c[i])
+
+        return c
+
+    target = PyQrack(
+        pyqrack_options={"isTensorNetwork": False, "isStabilizerHybrid": True},
+        dynamic_qubits=True,
+    )
+
+    N = 50
+
+    result = target.multi_run(ghz, 100, N)
+    result = Counter("".join(str(int(bit)) for bit in bits) for bits in result)
+    assert result.keys() == {"0" * N, "1" * N}
diff --git a/test/runtime/test_qrack.py b/test/runtime/test_qrack.py
index c6c98b9..b3a5078 100644
--- a/test/runtime/test_qrack.py
+++ b/test/runtime/test_qrack.py
@@ -7,7 +7,7 @@
 
 
 def run_mock(size: int, program: ir.Method) -> Mock:
-    memory = pyqrack.Memory(size, 0, Mock())
+    memory = pyqrack.StackMemory(size, 0, Mock())
     interp = pyqrack.PyQrackInterpreter(qasm2.main, memory=memory)
     interp.run(program, ())
 


Reviews:
user: login: Roger-luo
body: cool feature. I think this is good to me in general. Just minor comment on that type parameter. I think you can use a parameteric class there instead of annotating with `MemoryABC`
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-03T18:31:33Z
user: login: weinbe58
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-05T01:59:44Z
user: login: Roger-luo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T03:44:59Z

Review comments:
diff_hunk: @@ -30,11 +43,37 @@ def allocate(self, n_qubits: int):
 
         return tuple(range(curr_allocated, self.allocated))
 
+    def reset(self):
+        self.sim_reg.reset_all()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    sim_reg: "QrackSimulator"
+
+    def __post_init__(self):
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+        self.reset()
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+    def reset(self):
+        for qid in self.sim_reg.dump_ids():
+            self.sim_reg.release(qid)
+
 
 @dataclass
 class PyQrackInterpreter(Interpreter):
     keys = ["pyqrack", "main"]
-    memory: Memory = field(kw_only=True)
+    memory: MemoryABC = field(kw_only=True)
path: src/bloqade/pyqrack/base.py
user: login: Roger-luo
body: should this be a type parameter then?
created_at: 2025-03-03T18:29:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -30,11 +43,37 @@ def allocate(self, n_qubits: int):
 
         return tuple(range(curr_allocated, self.allocated))
 
+    def reset(self):
+        self.sim_reg.reset_all()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    sim_reg: "QrackSimulator"
+
+    def __post_init__(self):
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+        self.reset()
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+    def reset(self):
+        for qid in self.sim_reg.dump_ids():
+            self.sim_reg.release(qid)
+
 
 @dataclass
 class PyQrackInterpreter(Interpreter):
     keys = ["pyqrack", "main"]
-    memory: Memory = field(kw_only=True)
+    memory: MemoryABC = field(kw_only=True)
path: src/bloqade/pyqrack/base.py
user: login: Roger-luo
body: otherwise users of this interpreter won't get the type right.
created_at: 2025-03-03T18:29:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -30,11 +43,37 @@ def allocate(self, n_qubits: int):
 
         return tuple(range(curr_allocated, self.allocated))
 
+    def reset(self):
+        self.sim_reg.reset_all()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    sim_reg: "QrackSimulator"
+
+    def __post_init__(self):
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+        self.reset()
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+    def reset(self):
+        for qid in self.sim_reg.dump_ids():
+            self.sim_reg.release(qid)
+
 
 @dataclass
 class PyQrackInterpreter(Interpreter):
     keys = ["pyqrack", "main"]
-    memory: Memory = field(kw_only=True)
+    memory: MemoryABC = field(kw_only=True)
path: src/bloqade/pyqrack/base.py
user: login: weinbe58
body: I can add it but I don't think it will do much here. The only difference is runtime behavior when the address analysis fails and the program overflows the stack. otherwise the linter can't infer anything from the type hint. 
created_at: 2025-03-05T01:59:44Z
author_association: MEMBER
diff_hunk: @@ -30,11 +43,37 @@ def allocate(self, n_qubits: int):
 
         return tuple(range(curr_allocated, self.allocated))
 
+    def reset(self):
+        self.sim_reg.reset_all()
+        self.allocated = 0
+
+
+@dataclass
+class DynamicMemory(MemoryABC):
+    sim_reg: "QrackSimulator"
+
+    def __post_init__(self):
+        if self.sim_reg.is_tensor_network:
+            raise ValueError("DynamicMemory does not support tensor networks")
+
+        self.reset()
+
+    def allocate(self, n_qubits: int):
+        start = self.sim_reg.num_qubits()
+        for i in range(start, start + n_qubits):
+            self.sim_reg.allocate_qubit(i)
+
+        return tuple(range(start, start + n_qubits))
+
+    def reset(self):
+        for qid in self.sim_reg.dump_ids():
+            self.sim_reg.release(qid)
+
 
 @dataclass
 class PyQrackInterpreter(Interpreter):
     keys = ["pyqrack", "main"]
-    memory: Memory = field(kw_only=True)
+    memory: MemoryABC = field(kw_only=True)
path: src/bloqade/pyqrack/base.py
user: login: Roger-luo
body: fair. I think let's keep as it is then.
created_at: 2025-03-05T03:44:59Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-02T05:02:08Z
body: # ☂️ Python Coverage
> current status: ✅
## Overall Coverage
| Lines | Covered | Coverage | Threshold | Status |
| :---: | :-----: | :------: | :-------: | :----: |
|  451  |   393   |    87%   |     0%    |   🟢   |
## New Files
No new covered files...
## Modified Files
| File                              | Coverage | Status |
| :-------------------------------- | :------: | :----: |
| src/bloqade/pyqrack/\__init__.py  |   100%   |   🟢   |
| src/bloqade/pyqrack/base.py       |    91%   |   🟢   |
| src/bloqade/pyqrack/qasm2/core.py |    80%   |   🟢   |
| src/bloqade/pyqrack/target.py     |    98%   |   🟢   |
| **TOTAL**                         |  **92%** |   🟢   |


> **updated for commit: `c028e54` by [action](https://github.com/marketplace/actions/python-coverage)🐍**
author_association: NONE
user: login: codecov[bot]
created_at: 2025-03-02T05:02:17Z
body: ## [Codecov](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing) Report
Attention: Patch coverage is `89.85507%` with `7 lines` in your changes missing coverage. Please review.
| [Files with missing lines](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing) | Patch % | Lines |
|---|---|---|
| [src/bloqade/pyqrack/base.py](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing#diff-c3JjL2Jsb3FhZGUvcHlxcmFjay9iYXNlLnB5) | 89.65% | [3 Missing :warning: ](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing) |
| [src/bloqade/pyqrack/qasm2/core.py](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing#diff-c3JjL2Jsb3FhZGUvcHlxcmFjay9xYXNtMi9jb3JlLnB5) | 50.00% | [3 Missing :warning: ](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing) |
| [src/bloqade/pyqrack/target.py](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing#diff-c3JjL2Jsb3FhZGUvcHlxcmFjay90YXJnZXQucHk=) | 96.96% | [1 Missing :warning: ](https://app.codecov.io/gh/QuEraComputing/bloqade-pyqrack/pull/23?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=QuEraComputing) |

:loudspeaker: Thoughts on this report? [Let us know!](https://github.com/codecov/feedback/issues/255)
author_association: NONE