Fixing compilation route (#630)

* refactor compilation pipeline so that builder parser directly returns `Routine` instead of tuple * renaming `parse_source` to `parse`. * fixing trait. * fixing parser trait of routine to evaluate object values.
QuEraComputing · Sep 25, 2023 · e27a19c · e27a19c
1 parent 057db7e
commit e27a19c
Show file tree

Hide file tree

Showing 10 changed files with 51 additions and 56 deletions.
diff --git a/src/bloqade/builder/backend/bloqade.py b/src/bloqade/builder/backend/bloqade.py
@@ -26,9 +26,7 @@ def python(self):
             -> `...python().run(shots, ...)`
                 :: Run the current program using bloqade python backend
         """
-        from bloqade.ir.routine.bloqade import BloqadeServiceOptions
-
-        return BloqadeServiceOptions(self).python()
+        return self.parse().bloqade.python()
 
     def julia(self, solver: str, nthreads: int = 1):
         raise NotImplementedError
diff --git a/src/bloqade/builder/backend/braket.py b/src/bloqade/builder/backend/braket.py
@@ -39,9 +39,7 @@ def aquila(self):
                 and results returned
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).braket.aquila()
+        return self.parse().braket.aquila()
 
     def local_emulator(self):
         """
@@ -57,6 +55,4 @@ def local_emulator(self):
                 :: run on local emulator
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).braket.local_emulator()
+        return self.parse().braket.local_emulator()
diff --git a/src/bloqade/builder/backend/quera.py b/src/bloqade/builder/backend/quera.py
@@ -48,9 +48,7 @@ def device(self, config_file: Optional[str] = None, **api_config):
 
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).quera.device(config_file, **api_config)
+        return self.parse().quera.device(config_file, **api_config)
 
     def aquila(self):
         """
@@ -75,9 +73,7 @@ def aquila(self):
 
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).quera.aquila()
+        return self.parse().quera.aquila()
 
     def cloud_mock(self):
         """
@@ -102,9 +98,7 @@ def cloud_mock(self):
 
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).quera.cloud_mock()
+        return self.parse().quera.cloud_mock()
 
     def mock(self, state_file: str = ".mock_state.txt"):
         """
@@ -129,6 +123,4 @@ def mock(self, state_file: str = ".mock_state.txt"):
 
 
         """
-        from bloqade.ir.routine.base import Routine
-
-        return Routine(self).quera.mock(state_file)
+        return self.parse().quera.mock(state_file)
diff --git a/src/bloqade/builder/parse/builder.py b/src/bloqade/builder/parse/builder.py
@@ -14,7 +14,8 @@
 from typing import TYPE_CHECKING, Tuple, Union, Dict, List, Optional, Set
 
 if TYPE_CHECKING:
-    from bloqade.ir.routine.params import Params, ParamType
+    from bloqade.ir.routine.params import ParamType
+    from bloqade.ir.routine.base import Routine
     from bloqade.ir.analog_circuit import AnalogCircuit
 
 
@@ -260,9 +261,10 @@ def parse_circuit(self, builder: Builder) -> "AnalogCircuit":
 
         return circuit
 
-    def parse_source(self, builder: Builder) -> Tuple["AnalogCircuit", "Params"]:
-        from bloqade.ir.routine.params import Params
+    def parse(self, builder: Builder) -> "Routine":
         from bloqade.ir.analog_circuit import AnalogCircuit
+        from bloqade.ir.routine.params import Params
+        from bloqade.ir.routine.base import Routine
 
         self.reset(builder)
         self.read_register()
@@ -276,4 +278,4 @@ def parse_source(self, builder: Builder) -> Tuple["AnalogCircuit", "Params"]:
         )
         circuit = AnalogCircuit(self.register, self.sequence)
 
-        return circuit, params
+        return Routine(builder, circuit, params)
diff --git a/src/bloqade/builder/parse/trait.py b/src/bloqade/builder/parse/trait.py
@@ -1,4 +1,4 @@
-from typing import Tuple, Union, TYPE_CHECKING
+from beartype.typing import Union, TYPE_CHECKING
 import json
 import bloqade.ir as ir
 from bloqade.visualization import display_builder
@@ -7,7 +7,7 @@
 if TYPE_CHECKING:
     from bloqade.ir import AtomArrangement, ParallelRegister, Sequence
     from bloqade.ir.analog_circuit import AnalogCircuit
-    from bloqade.ir.routine.params import Params
+    from bloqade.ir.routine.base import Routine
     from bloqade.builder.base import Builder
 
 
@@ -47,16 +47,11 @@ def parse_circuit(self: "Builder") -> "AnalogCircuit":
 
 
 class ParseRoutine:
-    """Parse the program to return an AnalogCircuit as well as the parameters
-    for the circuit.
-
-
-    """
-
-    def parse_source(self: "Builder") -> Tuple["AnalogCircuit", "Params"]:
+    def parse(self: "Builder") -> "Routine":
+        """Parse the program to return a Routine object."""
         from bloqade.builder.parse.builder import Parser
 
-        return Parser().parse_source(self)
+        return Parser().parse(self)
 
 
 class Parse(ParseRegister, ParseSequence, ParseCircuit, ParseRoutine):
@@ -75,7 +70,7 @@ def n_atoms(self: "Builder"):
     def __repr__(self: "Builder"):
         from .builder import Parser
 
-        analog_circ, metas = Parser().parse_source(self)
+        analog_circ, metas = Parser().parse(self)
 
         return repr(analog_circ) + "\n" + repr(metas)
 

diff --git a/src/bloqade/ir/routine/base.py b/src/bloqade/ir/routine/base.py
@@ -7,30 +7,33 @@
 
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Union, Tuple
+from typing import TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
     from bloqade.ir.routine.braket import BraketServiceOptions
     from bloqade.ir.routine.quera import QuEraServiceOptions
+    from bloqade.ir.routine.bloqade import BloqadeServiceOptions
 
 
 class RoutineParse:
     def parse_register(self: "RoutineBase") -> Union[AtomArrangement, ParallelRegister]:
-        return self.source.parse_register()
+        return self.circuit.register
 
     def parse_sequence(self: "RoutineBase") -> Sequence:
-        return self.source.parse_sequence()
+        return self.circuit.sequence
 
     def parse_circuit(self: "RoutineBase") -> AnalogCircuit:
-        return self.source.parse_circuit()
+        return self.circuit
 
-    def parse_source(self: "RoutineBase") -> Tuple[AnalogCircuit, Params]:
-        return self.source.parse_source()
+    def parse(self: "RoutineBase") -> "Routine":
+        return self
 
 
 @dataclass(frozen=True)
 class RoutineBase(RoutineParse):
     source: Builder
+    circuit: AnalogCircuit
+    params: Params
 
 
 @dataclass(frozen=True)
@@ -41,10 +44,16 @@ class Routine(RoutineBase):
     def braket(self) -> "BraketServiceOptions":
         from .braket import BraketServiceOptions
 
-        return BraketServiceOptions(source=self.source)
+        return BraketServiceOptions(self.source, self.circuit, self.params)
 
     @property
     def quera(self) -> "QuEraServiceOptions":
         from .quera import QuEraServiceOptions
 
-        return QuEraServiceOptions(source=self.source)
+        return QuEraServiceOptions(self.source, self.circuit, self.params)
+
+    @property
+    def bloqade(self) -> "BloqadeServiceOptions":
+        from .bloqade import BloqadeServiceOptions
+
+        return BloqadeServiceOptions(self.source, self.circuit, self.params)
diff --git a/src/bloqade/ir/routine/bloqade.py b/src/bloqade/ir/routine/bloqade.py
@@ -10,7 +10,7 @@
 @dataclass(frozen=True)
 class BloqadeServiceOptions(RoutineBase):
     def python(self):
-        return BloqadePythonRoutine(source=self.source)
+        return BloqadePythonRoutine(self.source, self.circuit, self.params)
 
 
 @dataclass(frozen=True)
@@ -29,7 +29,7 @@ def _compile(
         from bloqade.emulate.codegen.hamiltonian import CompileCache
         from bloqade.task.bloqade import BloqadeTask
 
-        circuit, params = self.source.parse_source()
+        circuit, params = self.circuit, self.params
 
         circuit = AssignAnalogCircuit(params.static_params).visit(circuit)
 

diff --git a/src/bloqade/ir/routine/braket.py b/src/bloqade/ir/routine/braket.py
@@ -17,10 +17,10 @@ def aquila(self) -> "BraketHardwareRoutine":
         backend = BraketBackend(
             device_arn="arn:aws:braket:us-east-1::device/qpu/quera/Aquila"
         )
-        return BraketHardwareRoutine(source=self.source, backend=backend)
+        return BraketHardwareRoutine(self.source, self.circuit, self.params, backend)
 
     def local_emulator(self):
-        return BraketLocalEmulatorRoutine(source=self.source)
+        return BraketLocalEmulatorRoutine(self.source, self.circuit, self.params)
 
 
 @dataclass(frozen=True)
@@ -40,7 +40,7 @@ def _compile(
 
         capabilities = self.backend.get_capabilities()
 
-        circuit, params = self.parse_source()
+        circuit, params = self.circuit, self.params
         circuit = AssignAnalogCircuit(params.static_params).visit(circuit)
 
         tasks = OrderedDict()
@@ -175,7 +175,7 @@ def _compile(
         from bloqade.ir.analysis.assignment_scan import AssignmentScan
         from bloqade.submission.ir.braket import to_braket_task_ir
 
-        circuit, params = self.parse_source()
+        circuit, params = self.circuit, self.params
         circuit = AssignAnalogCircuit(params.static_params).visit(circuit)
 
         if isinstance(circuit.register, ParallelRegister):

diff --git a/src/bloqade/ir/routine/quera.py b/src/bloqade/ir/routine/quera.py
@@ -24,20 +24,20 @@ def device(self, config_file: Optional[str], **api_config):
 
         backend = QuEraBackend(**api_config)
 
-        return QuEraHardwareRoutine(source=self.source, backend=backend)
+        return QuEraHardwareRoutine(self.source, self.circuit, self.params, backend)
 
     def aquila(self) -> "QuEraHardwareRoutine":
         backend = QuEraBackend(**load_config("Aquila"))
-        return QuEraHardwareRoutine(source=self.source, backend=backend)
+        return QuEraHardwareRoutine(self.source, self.circuit, self.params, backend)
 
     def cloud_mock(self) -> "QuEraHardwareRoutine":
         backend = QuEraBackend(**load_config("Mock"))
-        return QuEraHardwareRoutine(source=self.source, backend=backend)
+        return QuEraHardwareRoutine(self.source, self.circuit, self.params, backend)
 
     @beartype
     def mock(self, state_file: str = ".mock_state.txt") -> "QuEraHardwareRoutine":
         backend = MockBackend(state_file=state_file)
-        return QuEraHardwareRoutine(source=self.source, backend=backend)
+        return QuEraHardwareRoutine(self.source, self.circuit, self.params, backend)
 
 
 @dataclass(frozen=True)
@@ -55,7 +55,7 @@ def _compile(
         from bloqade.ir.analysis.assignment_scan import AssignmentScan
         from bloqade.codegen.hardware.quera import QuEraCodeGen
 
-        circuit, params = self.parse_source()
+        circuit, params = self.circuit, self.params
         capabilities = self.backend.get_capabilities()
         circuit = AssignAnalogCircuit(params.static_params).visit(circuit)
 

diff --git a/src/bloqade/visualization/display.py b/src/bloqade/visualization/display.py
@@ -69,7 +69,10 @@ def liner(txt):
 def builder_figure(builder, batch_id):
     from bloqade.builder.parse.builder import Parser
 
-    analog_circ, metas = Parser().parse_source(builder)
+    routine = Parser().parse(builder)
+
+    analog_circ = routine.circuit
+    metas = routine.params
 
     kwargs = metas.static_params
     kwargs.update(metas.batch_params[batch_id])