compiler pipeline

src/bloqade/builder2/backend/__init__.py

@@ -1,28 +1,9 @@
-from .quera import SubmitQuEra, FlattenedQuEra
-from .braket import SubmitBraket, FlattenedBraket
-from .bloqade import SubmitBloqade, FlattenedBloqade
+from .quera import QuEraService
+from .braket import BraketService
+from .bloqade import BloqadeService
 
 
-class SubmitBackendRoute(SubmitQuEra, SubmitBraket, SubmitBloqade):
-    def device(self, name: str, *args, **kwargs):
-        if name == "quera.aquila":
-            dev = self.quera.aquila
-        elif name == "quera.gemini":
-            dev = self.quera.gemini
-        elif name == "braket.aquila":
-            dev = self.braket.aquila
-        elif name == "braket.simu":
-            dev = self.braket.simu
-        elif name == "bloqade.python":
-            dev = self.bloqade.python
-        elif name == "bloqade.julia":
-            dev = self.bloqade.julia
-        else:
-            raise ValueError(f"Unknown device: {name}")
-        return dev(*args, **kwargs)
-
-
-class FlattenedBackendRoute(FlattenedQuEra, FlattenedBraket, FlattenedBloqade):
+class BackendRoute(QuEraService, BraketService, BloqadeService):
     def device(self, name: str, *args, **kwargs):
         if name == "quera.aquila":
             dev = self.quera.aquila

src/bloqade/builder2/backend/base.py

@@ -1,72 +1,112 @@
-from typing import List
-from ..base import Builder
-
-# from ... import ir
+from itertools import repeat
+from typing import Dict
+import numbers
 
+from ... import ir
+from ..base import Builder
 from ..compile.trait import Compile
-from ..compile.stream import BuilderStream
+from ..compile.stream import BuilderStream, BuilderNode
+from ...task2.base import Batch
 
 
 class Backend(Builder, Compile):
-    def __init__(self, parent: Builder | None = None) -> None:
+    def __init__(self, cache: bool = False, parent: Builder | None = None) -> None:
         super().__init__(parent)
         self._ir_cache = None
+        self._cluster_spacing = None
+        self._static_params = {}
+        self._batch_params = {}
+        self._orders = ()
+        self._cache = cache
+
+    def _parse_args(self, *args: numbers.Real) -> Dict[str, float]:
+        if len(args) != len(self._orders):
+            raise ValueError(
+                f"Number of arguments ({len(args)}) does not match "
+                f"number of flattened variables ({len(self._orders)})."
+            )
+
+        return dict(zip(self._orders, map(float, args)))
+
+    def _compile_builder(self) -> None:
+        from ..assign import Assign, BatchAssign
+        from ..flatten import Flatten
+        from ..parallelize import Parallelize, ParallelizeFlatten
 
-    def __compile_ir__(self):
-        from ..assign import Assign
+        stream = BuilderStream(self)
 
+        while stream.curr is not None:
+            node = stream.read_next(
+                [Assign, BatchAssign, Flatten, ParallelizeFlatten, Parallelize]
+            )
+            match node:
+                case BuilderNode(node=Assign(static_params)):
+                    self._static_params = static_params
+                case BuilderNode(node=BatchAssign(batch_params)):
+                    self._batch_params = batch_params
+                case BuilderNode(node=Parallelize(cluster_spacing)) | BuilderNode(
+                    node=ParallelizeFlatten(cluster_spacing)
+                ):
+                    self._cluster_spacing = cluster_spacing
+                case BuilderNode(node=Flatten(orders)):
+                    self._orders = orders
+
+    def _compile_cache(self) -> ir.Program:
         if self._ir_cache:
             return self._ir_cache
 
-        stream = BuilderStream(self)
-        assign_options = stream.eat([Assign])
-
-        if assign_options is not None:
-            assignments = assign_options.node._assignments
-        else:
-            assignments = {}
+        self._ir_cache = self.compile_bloqade_ir(**self._static_params)
 
-        self._ir_cache = self.compile_bloqade_ir(**assignments)
         return self._ir_cache
 
-    def compile(self):
+    def _compile_ir(self, *args):
+        from ...codegen.common.static_assign import StaticAssignProgram
+
+        self._compile_builder()
+        ir_cache = self._compile_cache()
+
+        input_mapping = self._parse_args(*args)
+
+        if self._batch_params:
+            batch_iterators = [
+                zip(repeat(name), values) for name, values in self._batch_params.items()
+            ]
+
+            mappings = (
+                {**input_mapping, **dict(tuples)} for tuples in zip(*batch_iterators)
+            )
+
+            for mapping in mappings:
+                yield mapping, StaticAssignProgram(mapping).emit(ir_cache)
+        else:
+            yield input_mapping, ir_cache
+
+    def _compile_task(self, bloqade_ir: ir.Program, shots: int, **metadata):
         raise NotImplementedError(
-            "Compilation backend not implemented for "
+            "Compilation not implemented for "
             f"'{self.__service_name__}.{self.__device_name__}'."
         )
 
 
 class LocalBackend(Backend):
-    def __call__(self):
-        raise NotImplementedError(
-            "__call__ backend not implemented for "
-            f"'{self.__service_name__}.{self.__device_name__}'."
-        )
+    def __call__(self, *args, shots: int = 1):
+        tasks = [
+            self._compile_task(program, shots, **metadata)
+            for metadata, program in self._compile_ir(*args)
+        ]
+        batch = Batch(tasks)
 
+        return batch
 
-class RemoteBackend(Backend):
-    def __init__(self, parent: Builder | None = None) -> None:
-        super().__init__(parent)
 
-    def __call__(self):
+class RemoteBackend(Backend):
+    def __call__(self, *args, shots: int = 1):
         batch = self.submit()
         batch.fetch()
         return batch
 
-    def submit(self):
+    def submit(self, *args, shots: int = 1):
         raise NotImplementedError(
             "Submission backend not implemented for "
             f"'{self.__service_name__}.{self.__device_name__}'."
         )
-
-
-class FlattenedBackend(Backend):
-    def __init__(self, arg_list: List[str], parent: Builder | None = None) -> None:
-        super().__init__(parent)
-        self._arg_list = arg_list
-
-    def __call__(self, *args, nshots=1):
-        raise NotImplementedError(
-            "__call__ backend not implemented for "
-            f"'{self.__service_name__}.{self.__device_name__}'."
-        )

src/bloqade/builder2/backend/bloqade.py

@@ -1,8 +1,8 @@
 from ..base import Builder
-from .base import LocalBackend, FlattenedBackend
+from .base import LocalBackend
 
 
-class SubmitBloqade(Builder):
+class BloqadeService(Builder):
     @property
     def bloqade(self):
         return BloqadeDeviceRoute(self)
@@ -36,40 +36,3 @@ def __init__(
     ) -> None:
         super().__init__(solver, parent)
         self._nthreads = nthreads
-
-
-class FlattenedBloqade(Builder):
-    @property
-    def bloqade(self):
-        return FlattenedBloqadeDeviceRoute(self)
-
-
-class FlattenedBloqadeDeviceRoute(Builder):
-    def python(self, solver: str):
-        return FlattenedBloqadePython(solver, self)
-
-    def julia(self, solver: str, nthreads: int = 1):
-        return FlattenedBloqadeJulia(solver, nthreads, self)
-
-
-class FlattenedBloqadeBackend(FlattenedBackend):
-    __service_name__ = "bloqade"
-
-
-class FlattenedBloqadePython(FlattenedBloqadeBackend):
-    ___device_name__ = "python"
-
-    def __init__(self, solver: str, parent: Builder | None = None) -> None:
-        super().__init__(parent)
-        self._solver = solver
-
-
-class FlattenedBloqadeJulia(FlattenedBloqadeBackend):
-    ___device_name__ = "julia"
-
-    def __init__(
-        self, solver: str, nthreads: int = 1, parent: Builder | None = None
-    ) -> None:
-        super().__init__(parent)
-        self._solver = solver
-        self._nthreads = nthreads

src/bloqade/builder2/backend/braket.py

@@ -1,58 +1,81 @@
 from bloqade.builder2.base import Builder
-from .base import LocalBackend, RemoteBackend, FlattenedBackend
+from .base import LocalBackend, RemoteBackend
+from pydantic.dataclasses import dataclass
+from braket.analog_hamiltonian_simulator.rydberg.constants import (
+    RYDBERG_INTERACTION_COEF,
+)
 
 
-class SubmitBraket(Builder):
+class BraketService(Builder):
     @property
     def braket(self):
         return BraketDeviceRoute(self)
 
 
 class BraketDeviceRoute(Builder):
-    def aquila(self, nshots: int) -> "Aquila":
-        return Aquila(nshots, self)
+    def aquila(self) -> "Aquila":
+        return Aquila(self)
 
-    def simu(self, nshots: int) -> "Simu":
-        return Simu(nshots, self)
+    def simu(self) -> "Simu":
+        return Simu(self)
 
 
 class Aquila(RemoteBackend):
     __service_name__ = "braket"
     __device_name__ = "aquila"
 
-    def __init__(self, nshots: int, parent: Builder | None = None) -> None:
+    def __init__(self, parent: Builder | None = None) -> None:
         super().__init__(parent)
-        self._nshots = nshots
+
+
+@dataclass
+class SimuOptions:
+    steps: int = (1000,)
+    rydberg_interaction_coef: float = (RYDBERG_INTERACTION_COEF,)
+    blockade_radius: float = (0.0,)
+    progress_bar: bool = (False,)
+    atol: float = (1e-8,)
+    rtol: float = (1e-6,)
+    solver_method: str = ("adams",)
+    order: int = (12,)
+    nsteps: int = (1000,)
+    first_step: int = (0,)
+    max_step: int = (0,)
+    min_step: int = (0,)
 
 
 class Simu(LocalBackend):
     __service_name__ = "braket"
     __device_name__ = "simu"
 
-    def __init__(self, nshots: int, parent: Builder | None = None) -> None:
-        super().__init__(parent)
-        self._nshots = nshots
-
-
-class FlattenedBraket(Builder):
-    @property
-    def braket(self):
-        return FlattenedBraketDevice(self)
-
-
-class FlattenedBraketDevice(Builder):
-    def aquila(self, nshots: int) -> "Aquila":
-        return FlattenedAquila(nshots, self)
-
-    def simu(self, nshots: int) -> "Simu":
-        return FlattenedSimu(nshots, self)
-
-
-class FlattenedAquila(FlattenedBackend):
-    def __init__(self, parent: Builder | None = None) -> None:
-        super().__init__(parent)
-
-
-class FlattenedSimu(FlattenedBackend):
-    def __init__(self, parent: Builder | None = None) -> None:
+    def __init__(
+        self,
+        steps: int = 1000,
+        rydberg_interaction_coef: float = RYDBERG_INTERACTION_COEF,
+        blockade_radius: float = 0.0,
+        progress_bar: bool = False,
+        atol: float = 1e-8,
+        rtol: float = 1e-6,
+        solver_method: str = "adams",
+        order: int = 12,
+        nsteps: int = 1000,
+        first_step: int = 0,
+        max_step: int = 0,
+        min_step: int = 0,
+        parent: Builder | None = None,
+    ) -> None:
         super().__init__(parent)
+        self._local_emulator_options = SimuOptions(
+            steps,
+            rydberg_interaction_coef,
+            blockade_radius,
+            progress_bar,
+            atol,
+            rtol,
+            solver_method,
+            order,
+            nsteps,
+            first_step,
+            max_step,
+            min_step,
+        )