refactor builder (#364)

* refactor builder

* keep refactoring

* flatten works

* cleanup

* add a sketch of compile pipeline

* support device fn

* setup a basic lexer like codegen for builder

* waveform and address is working

* reexporting objects to `bloqade.ir` and `bloqade`

* adding `BoundedBravais` and `ParallelRegister` to imports.

* adding `create`/`copy` interface to `BuilderStream`

* wrap up whatever I have

* docstring.

* renambing fields of pulse and sequence

* implmenting compiler for pulse sequence.

* fix routing.

* adding better error messages for not implemented backends.

* stashing changes.

* adding `static_assign` functionality to scalars

* moving from just kargs to storing parameters as Scalar variables.

* stash

* rename

* fixing typo

* adding unit tests for

* rename class.

* moving visitor methods for primatives.

* update dev

* new compiler for sequence.

* removing main2.py

* rename modules.

* creating visitor for scalar AST, removing `static_assign` implement assign visitor for bloqade AST.

* adding canonicalization

* reorder imports.

* adding match arguments and serializers

* adding JSON serialization for Scalar IR.

* moving ir compiler into compiler subpackage

* creating stream submodule and json submodule.

* adding Sample node for builder.

* rename `checkpoint` to `coeff`

* adding json compiler for waveforms.

* adding json encoder for builder.

* task sketch.

* adding `apply` builder method for sequence objects.

* adding explicit check for `Waveform` type.

* adding serialization for builder + BloqadeIR

* fixing parallelize option for submit/flatten routes.

* commenting out old compiler. Adding register compiler

* setting up compiler trait.

* fixing compiler after refactoring circular imports

* fixing imports

* adding JSON compiler to builder.

* adding compiler traits.

* compiler pipeline

* adding ir caching option.

* updating caching to happen at all levels of bloqade-IR

* adding deserialization for bloqade-IR

* refactor base backend.

* finished draft implementation of serializer/deserializer

* Khwu/refactor task (#391)

* update, finished hardware tasl

* relocate Geometry, isolcate deps

* fix bug on Report

* update, change status behavior to request call. fix submit order does not shuffle in OrderDict

* add repr for Batch

* rename submit as resubmit

* finished restructure of Task/Batch

* update, finished bridge for task<->builder v2

* update, add simple docs

* fix namespace conflict in backend

* fix bug in RemoteBackend

* adding option for caching compiled programs in backends.

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

* removing shots from options for backends.

* keyword arg for `parent`

* missing impl for static assignment on register IR objects.

* sketching bug fix.

* fix static assign, and visit_register issue

* fixing builder compiler to blqoade IR

* tm

* fixing another bug in builder compiler.

* remove redundant print

* fix bugs for batches

* update, fix bugs for report()

* fix bugs in waveform builder

* fix bugs with PythonFn waveform default kwargs. fix most of the testing

* Khwu/rename (#399)

* rename

* remove factory

* temporary add assert to make test_quera_internal_api fail

* fixing bugs in test.

* fix bugs in incomplete LocalTask/ LocalBatch. Fix report for LocalBatch

* change constrain of scale to Union[Real,str,Scalar].
2. compile visiting of record node, replace Variable(name) with cast(name)
3. fix tests cases for test_program_visitor.
4. add location method for scale

* more testing case

* mask coverage for deprecated codes builder_old/ task_old/

* add get_tasks/remove_tasks with status codes

* merging on github.  (#396)

* updating program definition. refactor compiler

* going back to old arg parser.

* adding error messages for flatten.

* refactor builder.

* moving builder to builder_old

* moving builder2 to builder

* moving task to task_old

* moving task2 to task

* fixing directory

* tmp save

* commented cache_compiled_program

* finished refactor on bridge using schemaCompiler

* starting serialization of task

* finished serialization

* fixing bugs.

* fixing braket local emulator.

* removing test call

* fixing some tests

* fixing interface for backends.

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>

* adding serialization for batch. (#400)

* adding serialization for batch.

* adding missing object ot serialization.

* update, fix bugs in naming _compile_taskdata, recover more tests

* add load_batch + doc for tree_depth()

* fix bugs in DefaultVariable should be Scalar not real, and member should be Literal instead of Decimal.
2) Fix PythonFn kwargs does not properly processed

* fix backend api_config need to pass in with dereferences

* Phil/save load batch (#404)

* adding serialization for batch.

* adding missing object ot serialization.

* adding save and load functions for batch objects.

* removing dict dispatching.

* fix bug in static_assign visit Negative waveform feed in wrong ast

* fixing tests

* fixing issue with pattern matching with `Poly`

* Khwu/testing v3 (#408)

* test +6

* 1) add checking for batch_assign lengths. 2) more tests recover,

* update, add trait BatchAssignable

* test +1, fix bug in load_batch

* adding fix to parser to account for `Sample` builder node.

* fix test case for mod/ DefaultVariable

* fix covergerc wrong format

* ocmply ruff

* fixing last tests.

---------

CI is now passing for builder refactor

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>

* Khwu/testing v4 (#409)

* test +6

* 1) add checking for batch_assign lengths. 2) more tests recover,

* update, add trait BatchAssignable

* test +1, fix bug in load_batch

* adding fix to parser to account for `Sample` builder node.

* fix test case for mod/ DefaultVariable

* fix covergerc wrong format

* ocmply ruff

* add more tests

* add missing __init__.py for builder/compile to comply with doc builder

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

---------

Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>
Co-authored-by: Kai-Hsin Wu <kaihsinwu@gmail.com>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>

.coveragerc

@@ -1,5 +1,5 @@
 [run]
-omit = */tests/*
+omit = */tests/*, */src/bloqade/task_old/*, */src/bloqade/builder_old/*
 
 [report]
-omit = */tests/*
+omit = */tests/*, */src/bloqade/task_old/*, */src/bloqade/builder_old/*

src/bloqade/__init__.py

@@ -1,15 +1,28 @@
-from bloqade.ir.location import start
-from bloqade.ir import var, cast, Variable, Literal
+from bloqade.ir import var, cast, Variable, Literal, start
+
+
 from bloqade.builder.factory import (
     piecewise_linear,
     piecewise_constant,
     linear,
     constant,
 )
 import bloqade.ir as _ir
+from bloqade.task.json import save_batch, load_batch
 
 
 def tree_depth(depth: int = None):
+    """Setting globally maximum depth for tree printing
+
+    If `depth=None`, return current depth.
+    If `depth` is provided, setting current depth to `depth`
+
+    Args:
+        depth (int, optional): the user specified depth. Defaults to None.
+
+    Returns:
+        int: current updated depth
+    """
     if depth is not None:
         _ir.tree_print.max_tree_depth = depth
     return _ir.tree_print.max_tree_depth
@@ -26,4 +39,6 @@ def tree_depth(depth: int = None):
     "linear",
     "constant",
     "set_print_depth",
+    "save_batch",
+    "load_batch",
 ]

src/bloqade/builder/assign.py

@@ -0,0 +1,30 @@
+from typing import Optional
+from .base import Builder
+from .pragmas import Parallelizable, Flattenable, BatchAssignable
+from .backend import BackendRoute
+from .compile.trait import Parse
+import numpy as np
+
+
+class AssignBase(Builder):
+    __match_args__ = ("_assignments", "__parent__")
+
+    def __init__(self, parent: Optional[Builder] = None, **assignments) -> None:
+        super().__init__(parent)
+        # TODO: implement checks for assignments
+        self._assignments = assignments
+
+
+class Assign(
+    AssignBase, BatchAssignable, Flattenable, Parallelizable, BackendRoute, Parse
+):
+    pass
+
+
+class BatchAssign(AssignBase, Parallelizable, BackendRoute, Parse):
+    def __init__(self, parent: Optional[Builder] = None, **assignments) -> None:
+        super().__init__(parent, **assignments)
+        if not len(np.unique(list(map(len, assignments.values())))) == 1:
+            raise ValueError(
+                "all the assignment variables need to have same number of elements."
+            )

src/bloqade/builder/backend/__init__.py

@@ -0,0 +1,22 @@
+from .quera import QuEraService
+from .braket import BraketService
+from .bloqade import BloqadeService
+
+
+class BackendRoute(QuEraService, BraketService, BloqadeService):
+    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.local_emulator
+        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)

src/bloqade/builder/backend/base.py

@@ -0,0 +1,82 @@
+# from ... import ir
+from ..base import Builder
+
+from bloqade.task.batch import RemoteBatch, LocalBatch
+from typing import Any, Tuple
+from numbers import Real
+
+# from ..compile.quera import QuEraSchemaCompiler
+
+
+class Backend(Builder):
+    pass
+
+
+class LocalBackend(Backend):
+    def run(
+        self,
+        shots: int,
+        args: Tuple[Real, ...] = (),
+        name: str | None = None,
+        **kwargs,
+    ):
+        tasks = self.compile_tasks(shots, *args)
+
+        batch = LocalBatch(dict(zip(range(len(tasks)), tasks)), name)
+
+        # kwargs is the tuning params for integrators
+        batch._run(**kwargs)
+
+        return batch
+
+    def __call__(
+        self,
+        *args: Tuple[Real, ...],
+        shots: int = 1,
+        name: str | None = None,
+        **kwargs,
+    ) -> Any:
+        return self.run(shots, args, **kwargs)
+
+    def compile_tasks(self, shots, *args):
+        raise NotImplementedError
+
+
+class RemoteBackend(Backend):
+    def run(
+        self,
+        shots: int,
+        args: Tuple[Real, ...] = (),
+        name: str | None = None,
+        shuffle: bool = False,
+    ):
+        batch = self.submit(shots, args, name, shuffle)
+        batch.pull()  # blocking
+        return batch
+
+    def __call__(
+        self,
+        *args: Tuple[Real, ...],
+        shots: int = 1,
+        name: str | None = None,
+        shuffle: bool = False,
+    ) -> Any:
+        return self.run(shots, args, name, shuffle)
+
+    def compile_tasks(self, shots, *args):
+        raise NotImplementedError
+
+    def submit(
+        self,
+        shots: int,
+        args: Tuple[Real, ...] = (),
+        name: str = None,
+        shuffle: bool = False,
+    ):
+        tasks = self.compile_tasks(shots, *args)
+
+        batch = RemoteBatch(dict(zip(range(len(tasks)), tasks)), name)
+
+        batch._submit(shuffle)
+
+        return batch

src/bloqade/builder/backend/bloqade.py

@@ -0,0 +1,49 @@
+from ..base import Builder
+from .base import LocalBackend
+
+
+class BloqadeService(Builder):
+    @property
+    def bloqade(self):
+        return BloqadeDeviceRoute(self)
+
+
+class BloqadeDeviceRoute(Builder):
+    def python(self, solver: str):
+        return BloqadePython(solver, parent=self)
+
+    def julia(self, solver: str, nthreads: int = 1):
+        return BloqadeJulia(solver, nthreads, parent=self)
+
+
+class SubmitBloqadeBackend(LocalBackend):
+    __service_name__ = "bloqade"
+
+    def __init__(
+        self,
+        solver: str,
+        cache_compiled_program: bool = False,
+        parent: Builder | None = None,
+    ) -> None:
+        # super().__init__(cache_compiled_program, parent=parent)
+        super().__init__(parent=parent)
+        self._solver = solver
+
+
+class BloqadePython(SubmitBloqadeBackend):
+    __device_name__ = "python"
+
+
+class BloqadeJulia(SubmitBloqadeBackend):
+    __device_name__ = "julia"
+
+    def __init__(
+        self,
+        solver: str,
+        nthreads: int = 1,
+        cache_compiled_program: bool = False,
+        parent: Builder | None = None,
+    ) -> None:
+        # super().__init__(solver, cache_compiled_program, parent)
+        super().__init__(solver, parent)
+        self._nthreads = nthreads

src/bloqade/builder/backend/braket.py

@@ -0,0 +1,107 @@
+from typing import Optional
+from bloqade.builder.base import Builder
+from .base import LocalBackend, RemoteBackend
+
+# import bloqade.ir as ir
+from bloqade.task.braket import BraketTask
+from bloqade.task.braket_simulator import BraketEmulatorTask
+import bloqade.submission.braket as braket_submit
+
+# from bloqade.submission.ir.braket import to_braket_task_ir
+
+
+class BraketService(Builder):
+    @property
+    def braket(self):
+        return BraketDeviceRoute(self)
+
+
+class BraketDeviceRoute(Builder):
+    def aquila(self) -> "Aquila":
+        return Aquila(parent=self)
+
+    def local_emulator(self) -> "BraketEmulator":
+        return BraketEmulator(parent=self)
+
+
+class Aquila(RemoteBackend):
+    __service_name__ = "braket"
+    __device_name__ = "aquila"
+
+    def __init__(
+        self,
+        # cache_compiled_programs: bool = False,
+        parent: Builder | None = None,
+    ) -> None:
+        # super().__init__(cache_compiled_programs, parent=parent)
+        super().__init__(parent=parent)
+
+    """
+    def _compile_task(self, bloqade_ir: ir.Program, shots: int, **metadata):
+        from bloqade.codegen.hardware.quera import SchemaCodeGen
+
+        backend = braket_submit.BraketBackend()
+
+        capabilities = backend.get_capabilities()
+        schema_compiler = SchemaCodeGen({}, capabilities=capabilities)
+        task_ir = schema_compiler.emit(shots, bloqade_ir)
+        task_ir = task_ir.discretize(capabilities)
+        return BraketTask(
+            task_ir=task_ir,
+            backend=backend,
+            parallel_decoder=schema_compiler.parallel_decoder,
+        )
+    """
+
+    def compile_taskdata(self, shots, *args):
+        backend = braket_submit.BraketBackend()
+        return self._compile_tasks(shots, backend, *args)
+
+    def _compile_taskdata(self, shots, backend, *args):
+        from ..compile.quera import QuEraSchemaCompiler
+
+        capabilities = backend.get_capabilities()
+
+        quera_task_data_list = QuEraSchemaCompiler(self, capabilities).compile(
+            shots, *args
+        )
+        return quera_task_data_list
+
+    def compile_tasks(self, shots, *args):
+        backend = braket_submit.BraketBackend()
+        task_data = self._compile_taskdata(shots, backend, *args)
+
+        return [BraketTask(task_data=dat, backend=backend) for dat in task_data]
+
+
+class BraketEmulator(LocalBackend):
+    __service_name__ = "braket"
+    __device_name__ = "local_emulator"
+
+    def __init__(
+        self,
+        # cache_compiled_programs: bool = False,
+        parent: Optional[Builder] = None,
+    ) -> None:
+        # super().__init__(cache_compiled_programs, parent=parent)
+        super().__init__(parent=parent)
+
+    """
+    def _compile_task(self, bloqade_ir: ir.Program, shots: int, **metadata):
+        from bloqade.codegen.hardware.quera import SchemaCodeGen
+
+        schema_compiler = SchemaCodeGen({})
+        task_ir = schema_compiler.emit(shots, bloqade_ir)
+        return BraketEmulatorTask(task_ir=to_braket_task_ir(task_ir))
+    """
+
+    def compile_taskdata(self, shots, *args):
+        from ..compile.braket_simulator import BraketEimulatorCompiler
+
+        braketemu_task_data_list = BraketEimulatorCompiler(self).compile(shots, *args)
+        return braketemu_task_data_list
+
+    def compile_tasks(self, shots, *args):
+        task_data = self.compile_taskdata(shots, *args)
+
+        return [BraketEmulatorTask(task_data=dat) for dat in task_data]